Combining knowledge graphs, quickly and accurately

Novel cross-graph-attention and self-attention mechanisms enable state-of-the-art performance.

Knowledge graphs are a way of representing information that can capture complex relationships more easily than conventional databases. At Amazon, we use knowledge graphs to represent the hierarchical relationships between product types on amazon.com; the relationships between creators and content on Amazon Music and Prime Video; and general information for Alexa’s question-answering service — among other things.

Expanding a knowledge graph often involves integrating it with another knowledge graph. But different graphs may use different terms for the same entities, which can lead to errors and inconsistencies during integration. Hence the need for automated techniques of entity alignment, or determining which elements of different graphs refer to the same entities.

In a paper accepted to the Web Conference, my colleagues and I describe a new entity alignment technique that factors in information about the graph in the vicinity of the entity name. In tests involving the integration of two movie databases, our system improved upon the best-performing of ten baseline systems by 10% on a metric called area under the precision-recall curve (PRAUC), which evaluates the trade-off between true-positive and true-negative rates.

Despite our system’s improved performance, it remains highly computationally efficient. One of the baseline systems we used for comparison is a neural-network-based system called DeepMatcher, which was specifically designed with scalability in mind. On two tasks, involving movie databases and music databases, our system reduced training time by 95% compared to DeepMatcher, while offering enormous improvements in PRAUC.

To implement our model, we used a new open-source tool called DGL (Deep Graph Library), which was developed by researchers in Amazon Web Services.

A graph is a mathematical object that consists of nodes, usually depicted as circles, and edges, usually depicted as line segments connecting the circles. Network diagrams, organizational charts, and flow charts are familiar examples of graphs.

Our work specifically addresses the problem of merging multi-type knowledge graphs, or knowledge graphs whose nodes represent more than one type of entity. For instance, in the movie data sets we worked with, a node might represent an actor, a director, a film, a film genre, and so on. Edges represented the relationships between entities — acted in, directed, wrote, and so on.

Entity alignment.png
This example illustrates the challenge of entity alignment. IMDB lists the writer of the movie Don’t Stop Dreaming as Aditya Raj, but the (now defunct) Freebase database lists him as Aditya Raj Kapoor. Are they the same person?

Our system is an example of a graph neural network, a type of neural network that has recently become popular for graph-related tasks. To get a sense for how it works, consider the Freebase example above, which includes what we call the “neighborhood” of the node representing Aditya Raj Kapoor. This is a two-hop local graph, meaning that it contains the nodes connected to Kapoor (one hop) and the nodes connected to them (two hops), but it doesn’t fan out any farther through the knowledge graph. The neighborhood thus consists of six nodes.

With a standard graph neural network (GNN), the first step — known as the level-0 step — is to embed each of the nodes, or convert it to a fixed-length vector representation. That representation is intended to capture information about node attributes useful for the network’s task — in this case, entity alignment — and it’s learned during the network’s training.

Next, in the level-1 step, the network considers the central node (here, Aditya Raj Kapoor) and the nodes one hop away from it (Don’t Stop Dreaming and Sambar Salsa). For each of these nodes, it produces a new embedding, which consists of the node's level-0 embedding concatenated with the sum of its immediate neighbors' level-0 embeddings.

At the level-2 step — the final step in a two-hop network — the network produces a new embedding for the central node, which consists of that node’s level-1 embedding concatenated with the summation of the level-1 embeddings of its immediate neighbors.

Graph neural network
A demonstration of how graph neural networks use recursive embedding to condense all the information in a two-hop graph into a single vector. Relationships between entities — such as "produce" and "write" in a movie database (red and yellow arrows, respectively) — are encoded in the level-0 embeddings of the entities themselves (red and orange blocks).
Stacy Reilly

In our example, this process compresses the entire six-node neighborhood graph from the Freebase database into a single vector. It would do the same with the ten-node neighborhood graph from IMDB, and comparing the vectors is the basis for the network’s decision about whether or not the entities at the centers of the graphs — Aditya Raj and Aditya Raj Kapoor — are the same.

This is the standard implementation of the GNN for the entity alignment problem. Unfortunately, in our experiments, it performed terribly. So we made two significant modifications.

The first was a cross-graph attention mechanism. During the level-1 and level-2 aggregation stages, when the network sums the embeddings of each node’s neighbors, it weights those sums based on a comparison with the other graph.

In our example, that means that during the level-1 and level-2 aggregations, the nodes Don’t Stop Dreaming and Sambar Salsa, which show up in both the IMDB and Freebase graphs, will get greater weight than Gawaahi and Shamaal, which show up only in IMDB.

Cross-graph attention.png
In this example, our cross-graph attention mechanism (blue lines) gives added weight (dotted red lines) to the embeddings of entities shared between neighborhood graphs.

The cross-graph attention mechanism thus emphasizes correspondences between the graphs and de-emphasizes differences. After all, the differences between the graphs is why it’s useful to combine their information in the first place.

Radioactive.png
The original version of “Radioactive” and the remix are distinct tracks, but they share so many attributes that a naïve entity alignment system might misclassify them as identical.

This approach has one major problem, however: sometimes the differences between graphs matter more than their correspondences. Consider the example at left, which compares two different versions of Imagine Dragons’ hit “Radioactive”, the original album cut and the remix featuring Kendrick Lamar.

Here, the cross-graph attention mechanism might overweight the many similarities between the two tracks and underweight the key difference: the main performer. So our network also includes a self-attention mechanism.

Self-attention.png
The application of our self-attention mechanism in our running example involving Aditya Raj.

During training, the self-attention mechanism learns which attributes of an entity are most important for distinguishing it from entities that look similar. In this case, it would learn that many distinct recordings may share the same songwriter or songwriters; what distinguishes them is the performer.

These two modifications are chiefly responsible for the improved performance of our model versus the ten baselines we compared it with.

Finally, a quick remark about one of the several techniques we used to increase our model’s computational efficiency. Although, for purposes of entity alignment, we compare two-hop neighborhoods, we don’t necessarily include a given entity’s entire two-hop neighborhood. We impose a cap on the number of nodes included in the neighborhood, and to select nodes for inclusion, we use weighted sampling.

The sample weights have an inverse relationship to the number of neighbor nodes that share the same relationship to the node of interest. So, for instance, a movie might have dozens of actors but only one director. In that case, our method would have a much higher chance of including the director node in our sampled neighborhood than it would of including any given actor node. Restricting the neighborhood size in this way prevents our method’s computational complexity from getting out of hand.

Related content

US, WA, Seattle
Economists in this role partner with business stakeholders to distill complex problems into testable economic questions and generate actionable insights. They collaborate with engineers and scientists to estimate models on large-scale data, design pilots, measure impact, and scale successful prototypes into improved policies and programs. They leverage AI tools to scale economic study for broader business impact. They communicate findings to business leaders, incorporate feedback, and deliver customer-centric solutions at scale.
US, NY, New York
Are you passionate about solving big problems from ground-up? Do you enjoy building new state-of-the-art products at internet scale? Come lead the innovation in this startup team, vertical ad products. This is a green field problem without a known answer or a pattern to follow. We have ambitious vision to simplify full funnel advertising solutions, at scale, with specialized agentic AI-powered models and diversify the demand to strategic verticals including finserv, autos, locals.. etc. We are seeking an experienced Applied Scientist to drive innovation in our Ads Foundational Model. In this individual contributor role, you will apply advanced machine learning techniques to improve advertiser performance and customer experience. Key job responsibilities As an Applied Scientist on this team, you will: 1. Develop and drive the science strategy for Ads Foundational Model (Ads-FM), aligning it with the program's objectives and overall business goals. 2. Identify high-impact opportunities within Ads-FM program and lead the ideation, planning, and execution of science initiatives to address them. 3. Build and deploy machine learning models using computer vision, natural language processing, and deep learning to evaluate and enhance ad effectiveness. 4. Develop algorithms that extract meaningful signals from image, video, and audio content to predict and improve customer engagement 5. Leverage Amazon's extensive data repository to create predictive models that generate actionable recommendations for more compelling ad creative 6. Collaborate with business leaders and cross-functional teams to implement ML-powered solutions 7. Contribute to the ML roadmap for the Ads-FM program through innovation and research.
US, WA, Seattle
This role will contribute to developing the Economics and Science products and services in the Fee domain, with specialization in supply chain systems and fees. Through the lens of economics, you will develop causal links for how Amazon, Sellers and Customers interact. You will be a key and senior scientist, advising Amazon leaders how to price our services. You will work on developing frameworks and scaleable, repeatable models supporting optimal pricing and policy in the two-sided marketplace that is central to Amazon's business. The pricing for Amazon services is complex. You will partner with science and technology teams across Amazon including Advertising, Supply Chain, Operations, Prime, Consumer Pricing, and Finance. We are looking for an experienced Principal Economist to improve our understanding of seller Economics, enhance our ability to estimate the causal impact of fees, and work with partner teams to design pricing policy changes. In this role, you will provide guidance to scientists to develop econometric models to influence our fee pricing worldwide. You will lead the development of causal models to help isolate the impact of fee and policy changes from other business actions, using experiments when possible, or observational data when not. Key job responsibilities The ideal candidate will have extensive Economics knowledge, demonstrated strength in practical and policy relevant structural econometrics, strong collaboration skills, proven ability to lead highly ambiguous and large projects, and a drive to deliver results. They will work closely with Economists, Data / Applied Scientists, Strategy Analysts, Data Engineers, and Product leads to integrate economic insights into policy and systems production. Familiarity with systems and services that constitute seller supply chains is a plus but not required. About the team The Stores Economics and Sciences team is a central science team that supports Amazon's Retail and Supply Chain leadership. We tackle some of Amazon's most challenging economics and machine learning problems, where our mandate is to impact the business on massive scale.
US, CA, San Diego
The Private Brands team is looking for a Research Scientist to join the team in building science solutions at scale. Our team applies Optimization, Machine Learning, Statistics, Causal Inference, and Econometrics/Economics to derive actionable insights about the complex economy of Amazon’s retail business and develop Statistical Models and Algorithms to drive strategic business decisions and improve operations. We are an interdisciplinary team of Scientists, Engineers, and Economists. Key job responsibilities You will work with business leaders, scientists, and economists to translate business and functional requirements into concrete deliverables, including the design, development, testing, and deployment of highly scalable optimization solutions and ML models. This is a unique, high visibility opportunity for someone who wants to have business impact, dive deep into large-scale problems, enable measurable actions on the consumer economy, and work closely with scientists and economists. As a Research Scientist, you bring business and industry context to science and technology decisions. You set the standard for scientific excellence and make decisions that affect the way we build and integrate algorithms. Your solutions are exemplary in terms of algorithm design, clarity, model structure, efficiency, and extensibility. You tackle intrinsically hard problems, acquiring expertise as needed. You decompose complex problems into straightforward solutions. We are particularly interested in candidates with experience in Operations Research and predictive models and working with distributed systems. Academic and/or practical background in Operations Research, Machine Learning and Reinforcement Learning are particularly relevant for this position. To know more about Amazon science, Please visit https://www.amazon.science
US, CA, Palo Alto
Alexa for Shopping (previously Rufus) is seeking a Senior Manager, Applied Science to lead multidisciplinary teams of Applied Scientists and Machine Learning Engineers building next-generation conversational AI and multi-agent systems powering customer-facing experiences at scale. This leader will drive both scientific innovation and execution across large language models (LLMs), agent orchestration, retrieval and grounding systems, evaluation frameworks, and scalable AI infrastructure. The role requires a combination of deep technical judgment, organizational leadership, product and engineering partnership, and operational excellence. The ideal candidate has a strong track record of building high-performing science and engineering teams, translating ambiguous business problems into scalable AI solutions, and delivering measurable customer impact through applied machine learning and generative AI technologies. Key job responsibilities - Lead and grow teams of Applied Scientists and Machine Learning Engineers working on conversational AI and multi-agent orchestration systems. - Define and drive technical strategy for large-scale generative AI systems, including LLM routing, prompting, grounding, memory, tool use, personalization, and response optimization. - Partner closely with Product, Engineering, and Tech leadership to align AI investments with long-term business and customer goals. - Drive end-to-end delivery of production AI systems balancing quality, latency, scalability, safety, and operational reliability. - Establish scientific and engineering best practices across experimentation, evaluation, model iteration, and production deployment. - Lead roadmap prioritization and execution across research innovation and product delivery timelines. - Build scalable evaluation methodologies and quality frameworks for multilingual and global customer experiences. - Mentor and develop technical leaders across both science and engineering disciplines. - Foster a high-performance culture centered on customer obsession, innovation, operational excellence, and strong cross-functional collaboration.
US, NY, New York
We are seeking a Human-Robot Interaction (HRI) Applied Scientist to develop cutting-edge interactions that make robots feel alive, personal, and fun. In this role, you will focus on verbal and non-verbal conversational systems, social dynamics, memory, and long-term relationship formation between robots, their environments, and the people they interact with. Your contributions will be essential in advancing robotics by enabling expressive, socially intelligent, and trustworthy interactions between robots and humans. Key job responsibilities - Develop interactive systems that leverage large language models, multimodal inputs and outputs, reinforcement learning from human feedback, or other advanced techniques to achieve fluid, engaging, and socially appropriate robot behavior - Design and implement intelligent conversational systems that handle turn-taking, grounding, interruption, and incorporates context drawn from a robot's physical environment and shared history with a user - Integrate perceptual sensor streams including gaze, facial expression, gesture, posture, and more to understand social context and produce coherent, lifelike interactions. - Develop memory and personalization systems that allow robots to form lasting relationships with individual users, learn their environments, and adapt their behavior over weeks and months - Stay updated on advancements in HRI, NLP, multimodal AI, and cognitive and social science to apply cutting-edge techniques to robot interaction challenges - Lead technical projects from conception through production deployment - Mentor junior scientists and engineers - Bridge research initiatives with practical engineering implementation
IN, KA, Bengaluru
Do you want to join an innovative team of scientists applying machine learning and advanced statistical techniques to protect Amazon customers and enable a trusted eCommerce experience? Are you excited about modeling terabytes of data and building state-of-the-art algorithms to solve complex, real-world fraud and risk challenges? Do you enjoy owning end-to-end machine learning problems, directly influencing customer experience and company profitability, while collaborating in a diverse, high-performing team? If so, the Amazon Buyer Risk Prevention (BRP) Machine Learning team may be the right fit for you. We are seeking an Applied Scientist to design, develop, and deploy advanced algorithmic systems that safeguard millions of transactions every day. In this role, you will independently drive model development from problem formulation to production deployment, build scalable ML solutions, and leverage emerging technologies—including Generative AI and LLMs—to enhance fraud detection and next-generation risk prevention systems. Key job responsibilities Own end-to-end development of machine learning models for large-scale risk management systems Analyze large volumes of historical and real-time data to identify fraud patterns and emerging risk trends Design, develop, validate, and deploy innovative models to production environments Apply GenAI/LLM technologies to automate risk evaluation and improve operational efficiency Collaborate closely with software engineering teams to implement scalable, real-time model solutions Partner with operations and business stakeholders to translate risk insights into measurable impact Establish scalable and automated processes for data analysis, model experimentation, validation, and monitoring Track model performance and business metrics; communicate insights clearly to technical and non-technical stakeholders Research and implement novel machine learning and statistical methodologies
IN, KA, Bengaluru
Do you want to join an innovative team applying machine learning and advanced statistical techniques to protect Amazon customers and enable a trusted eCommerce experience? Are you excited about working with large-scale datasets and developing models that solve real-world fraud and risk challenges? If so, the Amazon Buyer Risk Prevention (BRP) Machine Learning team may be the right fit for you. We are seeking an Applied Scientist to help develop scalable machine learning solutions that safeguard millions of transactions every day. In this role, you will partner with senior scientists and engineers to translate business problems into data-driven solutions, build and evaluate models, and contribute to next-generation risk prevention systems, including applications of Generative AI and LLM technologies. Key job responsibilities Apply machine learning and statistical techniques to build and improve risk management models Analyze large-scale historical data to identify risk patterns and emerging trends Develop, validate, and deploy innovative models under the guidance of senior scientists Experiment with emerging technologies, including GenAI/LLMs, to enhance automation and risk evaluation Collaborate closely with software engineers to implement models in real-time production systems Partner with operations and business teams to improve risk policies and operational efficiency Build scalable, automated pipelines for data analysis, model training, and validation Monitor model performance and provide clear reporting on key risk and business metrics Research and prototype new modeling approaches to improve system performance
IN, KA, Bengaluru
Do you want to join an innovative team of scientists applying machine learning and advanced statistical techniques to protect Amazon customers and enable a trusted eCommerce experience? Are you excited about modeling terabytes of data and building state-of-the-art algorithms to solve complex, real-world fraud and risk challenges? Do you enjoy owning end-to-end machine learning problems, directly influencing customer experience and company profitability, while collaborating in a diverse, high-performing team? If so, the Amazon Buyer Risk Prevention (BRP) Machine Learning team may be the right fit for you. We are seeking an Applied Scientist to design, develop, and deploy advanced algorithmic systems that safeguard millions of transactions every day. In this role, you will independently drive model development from problem formulation to production deployment, build scalable ML solutions, and leverage emerging technologies—including Generative AI and LLMs—to enhance fraud detection and next-generation risk prevention systems. Key job responsibilities Own end-to-end development of machine learning models for large-scale risk management systems Analyze large volumes of historical and real-time data to identify fraud patterns and emerging risk trends Design, develop, validate, and deploy innovative models to production environments Apply GenAI/LLM technologies to automate risk evaluation and improve operational efficiency Collaborate closely with software engineering teams to implement scalable, real-time model solutions Partner with operations and business stakeholders to translate risk insights into measurable impact Establish scalable and automated processes for data analysis, model experimentation, validation, and monitoring Track model performance and business metrics; communicate insights clearly to technical and non-technical stakeholders Research and implement novel machine learning and statistical methodologies
IN, KA, Bengaluru
Do you want to lead the development of advanced machine learning systems that protect millions of customers and power a trusted global eCommerce experience? Are you passionate about modeling terabytes of data, solving highly ambiguous fraud and risk challenges, and driving step-change improvements through scientific innovation? If so, the Amazon Buyer Risk Prevention (BRP) Machine Learning team may be the right place for you. We are seeking a Senior Applied Scientist to define and drive the scientific direction of large-scale risk management systems that safeguard millions of transactions every day. In this role, you will lead the design and deployment of advanced machine learning solutions, influence cross-team technical strategy, and leverage emerging technologies—including Generative AI and LLMs—to build next-generation risk prevention platforms. Key job responsibilities Lead the end-to-end scientific strategy for large-scale fraud and risk modeling initiatives Define problem statements, success metrics, and long-term modeling roadmaps in partnership with business and engineering leaders Design, develop, and deploy highly scalable machine learning systems in real-time production environments Drive innovation using advanced ML, deep learning, and GenAI/LLM technologies to automate and transform risk evaluation Influence system architecture and partner with engineering teams to ensure robust, scalable implementations Establish best practices for experimentation, model validation, monitoring, and lifecycle management Mentor and raise the technical bar for junior scientists through reviews, technical guidance, and thought leadership Communicate complex scientific insights clearly to senior leadership and cross-functional stakeholders Identify emerging scientific trends and translate them into impactful production solutions