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, CA, San Francisco
Amazon has launched a new research lab in San Francisco to develop foundational capabilities for useful AI agents. We’re enabling practical AI to make our customers more productive, empowered, and fulfilled. Our work leverages large vision language models (VLMs) with reinforcement learning (RL) and world modeling to solve perception, reasoning, and planning to build useful enterprise agents. Our lab is a small, talent-dense team with the resources and scale of Amazon. Each team in the lab has the autonomy to move fast and the long-term commitment to pursue high-risk, high-payoff research. We’re entering an exciting new era where agents can redefine what AI makes possible. Key job responsibilities You will contribute directly to AI agent development in an applied research role to improve the multi-model perception and visual-reasoning abilities of our agent. Daily responsibilities including model training, dataset design, and pre- and post-training optimization. You will be hired as a Member of Technical Staff.
IN, TN, Chennai
Are you excited about the digital media revolution and passionate about designing and delivering advanced analytics that directly influence the product decisions of Amazon's digital businesses. Do you see yourself as a champion of innovating on behalf of the customer by turning data insights into action? The Amazon Digital Acceleration Analytics team is looking for an analytical and technically skilled individual to join our team. In this role, you will invent, build and deploy state of the art machine-learning models and systems to enable and enhance the team's mission This role offers wide scope, autonomy, and ownership. You will work closely with software engineers & data engineers to put algorithms into practice. You should have strong business judgement, excellent written and verbal communication skills. The candidate should be willing to take on challenging initiatives and be capable of working both independently and with others as a team. Key job responsibilities We are looking for an experienced data scientist with strong foundations in mathematics, statistics & machine learning with exceptional communication and leadership skills, and a proven track record of delivery. In this role, You will Define a long-term science vision and roadmap for the team, driven fundamentally from our customers' needs, translating those directions into specific plans for engineering teams. Design and execute machine learning projects/products end-to-end: from ideation, analysis, prototyping, development, metrics, and monitoring. Drive end-to-end statistical analysis that have a high degree of ambiguity, scale, and complexity. Research and develop advanced Generative AI based solutions to solve diverse customer problems. About the team The MIDAS team operates within Amazon's Digital Analytics (DA) engineering organization, building analytics and data engineering solutions that support cross-digital teams. Our platform delivers a wide range of capabilities, including metadata discovery, data lineage, customer segmentation, compliance automation, AI-driven data access through generative AI and LLMs, and advanced data quality monitoring. Today, more than 100 Amazon business and technology teams rely on MIDAS, with over 20,000 monthly active users leveraging our mission-critical tools to drive data-driven decisions at Amazon scale.
US, WA, Seattle
Prime Video is a first-stop entertainment destination offering customers a vast collection of premium programming in one app available across thousands of devices. Prime members can customize their viewing experience and find their favorite movies, series, documentaries, and live sports – including Amazon MGM Studios-produced series and movies; licensed fan favorites; and programming from Prime Video add-on subscriptions such as Apple TV+, Max, Crunchyroll and MGM+. All customers, regardless of whether they have a Prime membership or not, can rent or buy titles via the Prime Video Store, and can enjoy even more content for free with ads. Are you interested in shaping the future of entertainment? Prime Video's technology teams are creating best-in-class digital video experience. As a Prime Video technologist, you’ll have end-to-end ownership of the product, user experience, design, and technology required to deliver state-of-the-art experiences for our customers. You’ll get to work on projects that are fast-paced, challenging, and varied. You’ll also be able to experiment with new possibilities, take risks, and collaborate with remarkable people. We’ll look for you to bring your diverse perspectives, ideas, and skill-sets to make Prime Video even better for our customers. With global opportunities for talented technologists, you can decide where a career Prime Video Tech takes you! We are forming a new organization within Prime Video to redefine our operational landscape through the power of artificial intelligence. As a Applied Scientist within this initiative, you will be a technical leader helping to design and build the intelligent systems that power our vision. You will tackle complex and ambiguous problems, designing and delivering scalable and resilient agentic AI and ML solutions from the ground up. You will not only write high-quality, maintainable software and models, but also mentor other scientists, influence our technical strategy, and drive engineering best practices across the team. Your work will directly contribute to making Prime Video's operations more efficient and will set the technical foundation for years to come. We're seeking candidates with strong experience in computer vision and generative AI technologies. In this role, you'll apply cutting-edge techniques in image and video understanding, visual content generation, and multimodal AI systems to transform how Prime Video operates at scale. Key job responsibilities • Lead the design and architecture of highly scalable, available, and resilient services for our AI automation platform. • Write high-quality, maintainable, and robust code to solve complex business problems, building flexible systems without over-engineering. • Act as a technical leader and mentor for other engineers on the team, assisting with career growth and encouraging excellence. • Work through ambiguous requirements, cut through complexity, and translate business needs into scalable technical solutions. • Take ownership of the full software development lifecycle, including design, testing, deployment, and operations. • Work closely with product managers, scientists, and other engineers to build and launch new features and systems. About the team This role offers a unique opportunity to shape the future of one of Amazon's most exciting businesses through the application of AI technologies. If you're passionate about leveraging AI to drive real-world impact at massive scale, we want to hear from you.
US, CA, San Francisco
Join the next revolution in robotics at Amazon's Frontier AI & Robotics team, where you'll work alongside world-renowned AI pioneers to push the boundaries of what's possible in robotic intelligence. As an Applied Scientist, you'll be at the forefront of developing breakthrough foundation models that enable robots to perceive, understand, and interact with the world in unprecedented ways. You'll drive independent research initiatives in areas such as perception, manipulation, science understanding, locomotion, manipulation, sim2real transfer, multi-modal foundation models and multi-task robot learning, designing novel frameworks that bridge the gap between state-of-the-art research and real-world deployment at Amazon scale. In this role, you'll balance innovative technical exploration with practical implementation, collaborating with platform teams to ensure your models and algorithms perform robustly in dynamic real-world environments. You'll have access to Amazon's vast computational resources, enabling you to tackle ambitious problems in areas like very large multi-modal robotic foundation models and efficient, promptable model architectures that can scale across diverse robotic applications. Key job responsibilities - Drive independent research initiatives across the robotics stack, including robotics foundation models, focusing on breakthrough approaches in perception, and manipulation, for example open-vocabulary panoptic scene understanding, scaling up multi-modal LLMs, sim2real/real2sim techniques, end-to-end vision-language-action models, efficient model inference, video tokenization - Design and implement novel deep learning architectures that push the boundaries of what robots can understand and accomplish - Lead full-stack robotics projects from conceptualization through deployment, taking a system-level approach that integrates hardware considerations with algorithmic development, ensuring robust performance in production environments - Collaborate with platform and hardware teams to ensure seamless integration across the entire robotics stack, optimizing and scaling models for real-world applications - Contribute to the team's technical strategy and help shape our approach to next-generation robotics challenges A day in the life - Design and implement novel foundation model architectures and innovative systems and algorithms, leveraging our extensive infrastructure to prototype and evaluate at scale - Collaborate with our world-class research team to solve complex technical challenges - Lead technical initiatives from conception to deployment, working closely with robotics engineers to integrate your solutions into production systems - Participate in technical discussions and brainstorming sessions with team leaders and fellow scientists - Leverage our massive compute cluster and extensive robotics infrastructure to rapidly prototype and validate new ideas - Transform theoretical insights into practical solutions that can handle the complexities of real-world robotics applications About the team At Frontier AI & Robotics, we're not just advancing robotics – we're reimagining it from the ground up. Our team is building the future of intelligent robotics through innovative foundation models and end-to-end learned systems. We tackle some of the most challenging problems in AI and robotics, from developing sophisticated perception systems to creating adaptive manipulation strategies that work in complex, real-world scenarios. What sets us apart is our unique combination of ambitious research vision and practical impact. We leverage Amazon's massive computational infrastructure and rich real-world datasets to train and deploy state-of-the-art foundation models. Our work spans the full spectrum of robotics intelligence – from multimodal perception using images, videos, and sensor data, to sophisticated manipulation strategies that can handle diverse real-world scenarios. We're building systems that don't just work in the lab, but scale to meet the demands of Amazon's global operations. Join us if you're excited about pushing the boundaries of what's possible in robotics, working with world-class researchers, and seeing your innovations deployed at unprecedented scale.
US, CA, San Francisco
Join the next revolution in robotics at Amazon's Frontier AI & Robotics team, where you'll work alongside world-renowned AI pioneers to push the boundaries of what's possible in robotic intelligence. As an Applied Scientist, you'll be at the forefront of developing breakthrough foundation models that enable robots to perceive, understand, and interact with the world in unprecedented ways. You'll drive independent research initiatives in areas such as perception, manipulation, science understanding, locomotion, manipulation, sim2real transfer, multi-modal foundation models and multi-task robot learning, designing novel frameworks that bridge the gap between state-of-the-art research and real-world deployment at Amazon scale. In this role, you'll balance innovative technical exploration with practical implementation, collaborating with platform teams to ensure your models and algorithms perform robustly in dynamic real-world environments. You'll have access to Amazon's vast computational resources, enabling you to tackle ambitious problems in areas like very large multi-modal robotic foundation models and efficient, promptable model architectures that can scale across diverse robotic applications. Key job responsibilities - Drive independent research initiatives across the robotics stack, including robotics foundation models, focusing on breakthrough approaches in perception, and manipulation, for example open-vocabulary panoptic scene understanding, scaling up multi-modal LLMs, sim2real/real2sim techniques, end-to-end vision-language-action models, efficient model inference, video tokenization - Design and implement novel deep learning architectures that push the boundaries of what robots can understand and accomplish - Lead full-stack robotics projects from conceptualization through deployment, taking a system-level approach that integrates hardware considerations with algorithmic development, ensuring robust performance in production environments - Collaborate with platform and hardware teams to ensure seamless integration across the entire robotics stack, optimizing and scaling models for real-world applications - Contribute to the team's technical strategy and help shape our approach to next-generation robotics challenges A day in the life - Design and implement novel foundation model architectures and innovative systems and algorithms, leveraging our extensive infrastructure to prototype and evaluate at scale - Collaborate with our world-class research team to solve complex technical challenges - Lead technical initiatives from conception to deployment, working closely with robotics engineers to integrate your solutions into production systems - Participate in technical discussions and brainstorming sessions with team leaders and fellow scientists - Leverage our massive compute cluster and extensive robotics infrastructure to rapidly prototype and validate new ideas - Transform theoretical insights into practical solutions that can handle the complexities of real-world robotics applications About the team At Frontier AI & Robotics, we're not just advancing robotics – we're reimagining it from the ground up. Our team is building the future of intelligent robotics through innovative foundation models and end-to-end learned systems. We tackle some of the most challenging problems in AI and robotics, from developing sophisticated perception systems to creating adaptive manipulation strategies that work in complex, real-world scenarios. What sets us apart is our unique combination of ambitious research vision and practical impact. We leverage Amazon's massive computational infrastructure and rich real-world datasets to train and deploy state-of-the-art foundation models. Our work spans the full spectrum of robotics intelligence – from multimodal perception using images, videos, and sensor data, to sophisticated manipulation strategies that can handle diverse real-world scenarios. We're building systems that don't just work in the lab, but scale to meet the demands of Amazon's global operations. Join us if you're excited about pushing the boundaries of what's possible in robotics, working with world-class researchers, and seeing your innovations deployed at unprecedented scale.
US, WA, Seattle
Are you excited to help customers discover the hottest and best reviewed products? The Discovery Tech team helps customers discover and engage with new, popular and relevant products across Amazon worldwide. We do this by combining technology, science, and innovation to build new customer-facing features and experiences alongside advanced tools for marketers. You will be responsible for creating and building critical services that automatically generate, target, and optimize Amazon’s cross-category marketing and merchandising. Through the enablement of intelligent marketing campaigns that leverage machine-learning models, you will help to deliver the best possible shopping experience for Amazon’s customers all over the globe. We are looking for analytical problem solvers who enjoy diving into data, excited about data science and statistics, can multi-task, and can credibly interface between engineering teams and business stakeholders. Your analytical abilities, business understanding, and technical savvy will be used to identify specific and actionable opportunities to solve existing business problems and look around corners for future opportunities. Your domain spans the design, development, testing, and deployment of data-driven and highly scalable machine learning solutions in product recommendation. As an Applied 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. To know more about Amazon science, please visit https://www.amazon.science
ES, B, Barcelona
Are you a scientist passionate about advancing the frontiers of computer vision, machine learning, or large language models? Do you want to work on innovative research projects that lead to innovative products and scientific publications? Would you value access to extensive datasets? If you answer yes to any of these questions, you'll find a great fit at Amazon. We're seeking a hands-on researcher eager to derive, implement, and test the next generation of Generative AI, computer vision, ML, and NLP algorithms. Our research is innovative, multidisciplinary, and far-reaching. We aim to define, deploy, and publish pioneering research that pushes the boundaries of what's possible. To achieve our vision, we think big and tackle complex technological challenges at the forefront of our field. Where technology doesn't exist, we create it. Where it does, we adapt it to function at Amazon's scale. We need team members who are passionate, curious, and willing to learn continuously. Key job responsibilities * Derive novel computer vision and ML models or LLMs/VLMs. * Design and develop scalable ML models. * Create and work with large datasets * Work with large GPU clusters. * Work closely with software engineering teams to deploy your innovations. * Publish your work at major conferences/journals. * Mentor team members in the use of your AI models. A day in the life As a Senior Applied Scientist at Amazon, your typical day might look like this: * Dive into coding, deriving new ML models for computer vision or NLP * Experiment with massive datasets on our GPU clusters * Brainstorm with your team to solve complex AI challenges * Collaborate with engineers to turn your research into real products * Write up your findings for publication in top journals or conferences * Mentor junior team members on AI concepts and implementation About the team DiscoVision, a science unit within Amazon's UPMT, focuses on advancing visual content capabilities through state-of-the-art AI technology. Our team specializes in developing state-of-the-art technologies in text-to-image/video Generative AI, 3D modeling, and multimodal Large Language Models (LLMs).
US, WA, Seattle
We are seeking a Principal Applied Scientist to lead research and development in automated reasoning, formal verification, and program analysis. You will drive innovation in making formal methods practical and accessible for real-world systems at cloud scale. Key job responsibilities - Lead research initiatives in automated reasoning, formal verification, SMT solving, model checking, or program analysis - Design and implement novel algorithms and techniques that advance the state of the art - Mentor and guide applied scientists, research scientists, and engineers - Collaborate with product teams to transition research into production systems - Define technical vision and strategy for automated reasoning initiatives - Represent AWS in the academic and research community - Drive cross-organizational impact through technical leadership About the team The Automated Reasoning Group at AWS develops and applies cutting-edge formal methods and automated reasoning techniques to ensure the security, reliability, and correctness of AWS services and customer applications. Our work innovates tools and services to perform verification at scale and apply them to build safe and secure systems at AWS. We are also pioneering the use of formal verification and automated reasoning to develop agentic systems, ensuring AI agents operate within defined safety boundaries.
IN, TS, Hyderabad
Do you want to join an innovative team of scientists who leverage machine learning and statistical techniques to revolutionize how businesses discover and purchase products on Amazon? Are you passionate about building intelligent systems that understand and predict complex B2B customer needs? The Amazon Business team is looking for exceptional Applied Science to help shape the future of B2B commerce. Amazon Business is one of Amazon's fastest-growing initiatives focused on serving business customers, from individual professionals to large institutions, with unique and complex purchasing needs. Our customers require sophisticated solutions that go beyond traditional B2C experiences, including bulk purchasing, approval workflows, and business-grade service support. The AB-MSET Applied Science team focuses on building intelligent systems for delivering personalized, contextual service experiences throughout the customer lifecycle. We apply advanced machine learning techniques to develop sophisticated intent detection models for business customer service needs, create intelligent matching algorithms for optimal service routing based on multiple variables including customer value, maturity, effort, and issue complexity, build predictive models to enable proactive service interventions, design recommendation systems for self-service solutions, and develop ML models for automated service resolution. As an Applied Scientist on the team, you will design and develop state-of-the-art ML models for service intent classification, routing optimization, and customer experience personalization. You will analyze large-scale business customer interaction data to identify patterns and opportunities for automation, create scalable solutions for complex B2B service scenarios using advanced ML techniques, and work closely with engineering teams to implement and deploy models in production. You will collaborate with business stakeholders to identify opportunities for ML applications, establish automated processes for model development, validation, and maintenance, lead research initiatives to advance the state-of-the-art in B2B service science, and mentor other scientists and engineers in applying ML techniques to business problems.
US, WA, Bellevue
Amazon Leo is an initiative to increase global broadband access through a constellation of 3,236 satellites in low Earth orbit (LEO). Its mission is to bring fast, affordable broadband to unserved and underserved communities around the world. Amazon Leo will help close the digital divide by delivering fast, affordable broadband to a wide range of customers, including consumers, businesses, government agencies, and other organizations operating in places without reliable connectivity. Do you get excited by aerospace, space exploration, and/or satellites? Do you want to help build solutions at Amazon Leo to transform the space industry? If so, then we would love to talk! Key job responsibilities Work cross-functionally with product, business development, and various technical teams (engineering, science, simulations, etc.) to execute on the long-term vision, strategy, and architecture for the science-based global demand forecast. Design and deliver modern, flexible, scalable solutions to integrate data from a variety of sources and systems (both internal and external) and develop Bandwidth Usage models at granular temporal and geographic grains, deployable to Leo traffic management systems. Work closely with the capacity planning science team to ensure that demand forecasts feed seamlessly into their systems to deliver continuous optimization of resources. Lead short and long terms technical roadmap definition efforts to deliver solutions that meet business needs in pre-launch, early-launch, and mature business phases. Synthesize and communicate insights and recommendations to audiences of varying levels of technical sophistication to drive change across Amazon Leo. Export Control Requirement: Due to applicable export control laws and regulations, candidates must be a U.S. citizen or national, U.S. permanent resident (i.e., current Green Card holder), or lawfully admitted into the U.S. as a refugee or granted asylum. About the team The Amazon Leo Global Demand Planning team's mission is to map customer demand across space and time. We enable Amazon Leo's long-term success by delivering actionable insights and scientific forecasts across geographies and customer segments to empower long range planning, capacity simulations, business strategy, and hardware manufacturing recommendations through scalable tools and durable mechanisms.