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, MA, Boston
Our team is involved with pre-silicon design verification for custom IP. A critical requirement of the verification flow is the requirement of legal and realistic stimulus of a custom Machine Learning Accelerator Chip. Content creation is built using formal methods that model legal behavior of the design and then solving the problem to create the specific assembly tests. The entire frame work for creating these custom tests is developed using a SMT solver and custom software code to guide the solution space into templated scenarios. This highly visible and innovative role requires the design of this solving framework and collaborating with design verification engineers, hardware architects and designers to ensure that interesting content can be created for the projects needs. Key job responsibilities Develop an understanding for a custom machine learning instruction set architecture. Model correctness of instruction streams using first order logic. Create custom API's to allow control over scheduling and randomness. Deploy algorithms to ensure concurrent code is safely constructed. Create coverage metrics to ensure solution space coverage. Use novel methods like machine learning to automate content creation. About the team Utility Computing (UC) AWS Utility Computing (UC) provides product innovations — from foundational services such as Amazon’s Simple Storage Service (S3) and Amazon Elastic Compute Cloud (EC2), to consistently released new product innovations that continue to set AWS’s services and features apart in the industry. As a member of the UC organization, you’ll support the development and management of Compute, Database, Storage, Internet of Things (Iot), Platform, and Productivity Apps services in AWS, including support for customers who require specialized security solutions for customers who require specialized security solutions for their cloud services. Annapurna Labs (our organization within AWS UC) designs silicon and software that accelerates innovation. Customers choose us to create cloud solutions that solve challenges that were unimaginable a short time ago—even yesterday. Our custom chips, accelerators, and software stacks enable us to take on technical challenges that have never been seen before, and deliver results that help our customers change the world.
IN, KA, Bengaluru
Amazon Devices is an inventive research and development company that designs and engineer high-profile devices like the Kindle family of products, Fire Tablets, Fire TV, Health Wellness, Amazon Echo & Astro products. This is an exciting opportunity to join Amazon in developing state-of-the-art techniques that bring Gen AI on edge for our consumer products. We are looking for exceptional scientists to join our Applied Science team and help develop the next generation of edge models, and optimize them while doing co-designed with custom ML HW based on a revolutionary architecture. Work hard. Have Fun. Make History. Key job responsibilities - Quantize, prune, distill, finetune Gen AI models to optimize for edge platforms - Fundamentally understand Amazon’s underlying Neural Edge Engine to invent optimization techniques - Analyze deep learning workloads and provide guidance to map them to Amazon’s Neural Edge Engine - Use first principles of Information Theory, Scientific Computing, Deep Learning Theory, Non Equilibrium Thermodynamics - Train custom Gen AI models that beat SOTA and paves path for developing production models - Collaborate closely with compiler engineers, fellow Applied Scientists, Hardware Architects and product teams to build the best ML-centric solutions for our devices - Publish in open source and present on Amazon's behalf at key ML conferences - NeurIPS, ICLR, MLSys.
US, CA, Culver City
MULTIPLE POSITIONS AVAILABLE Employer: AMAZON.COM SERVICES LLC Offered Position: Applied Scientist III Job Location: Culver City, California Job Number: AMZ10564141 Position Responsibilities: Participate in the design, development, evaluation, deployment and updating of data-driven models and analytical solutions for machine learning (ML) and/or natural language (NL) applications. Develop and/or apply statistical modeling techniques (e.g. Bayesian models and deep neural networks), optimization methods, and other ML techniques to different applications in business and engineering. Routinely build and deploy ML models on available data, and run and analyze experiments in a production environment. Identify new opportunities for research in order to meet business goals. Research and implement novel ML and statistical approaches to add value to the business. Mentor junior engineers and scientists. 40 hours / week, 8:00am-5:00pm, Salary Range: $167,100/year to $226,100/year. Amazon is a total compensation company. Dependent on the position offered, equity, sign-on payments, and other forms of compensation may be provided as part of a total compensation package, in addition to a full range of medical, financial, and/or other benefits. For more information, visit: https://www.aboutamazon.com/workplace/employee-benefits. Amazon.com is an Equal Opportunity-Affirmative Action Employer – Minority / Female / Disability / Veteran / Gender Identity / Sexual Orientation.#0000
US, CA, San Francisco
If you are interested in this position, please apply on Twitch's Career site https://www.twitch.tv/jobs/en/ About Us: Twitch is the world’s biggest live streaming service, with global communities built around gaming, entertainment, music, sports, cooking, and more. It is where thousands of communities come together for whatever, every day. We’re about community, inside and out. You’ll find coworkers who are eager to team up, collaborate, and smash (or elegantly solve) problems together. We’re on a quest to empower live communities, so if this sounds good to you, see what we’re up to on LinkedIn and X, and discover the projects we’re solving on our Blog. Be sure to explore our Interviewing Guide to learn how to ace our interview process. About the Role We are looking for an Applied Scientist to solve challenging and open-ended problems in the domain of Creators and content contextual signals. As an Applied Scientist on Twitch's GTM Science team, you will use applied machine learning to tackle ambiguous business problems across problem areas such as Ads, Brand Safety, Creator Sponsorship and Content Understanding. You will lead high profile projects where you will use a wide toolbox of ML tools to handle multiple types of data, including user behavior, metadata, and user generated content such as text and video. You will collaborate with a team of passionate scientists, engineers, and product stakeholders to develop these models, prototype and put them into production, where they can help Twitch's creators and viewers succeed and build communities. You will report to our Senior Data Science Manager for GTM. This position is located in San Francisco, CA. You Will - Build machine learning products to enrich Twitch’s content understanding and help creators and viewers build and discover their communities. - Propose, design, and oversee execution for high-impact projects that work backwards from customer problems to develop the right solution for the job. Know when to use a classical ML model versus a state-of-the-art one. - Collaborate with cross-functional engineering and product teams to prototype and deploy your models into flexible data pipelines and ML-based services. - Bring innovation and help drive science culture on the team by staying up to date and experimenting with new techniques in LLM, computer vision, inference and model evaluation. Perks * Medical, Dental, Vision & Disability Insurance * 401(k) * Maternity & Parental Leave * Flexible PTO * Amazon Employee Discount
IN, KA, Bengaluru
Interested to build the next generation Financial systems that can handle billions of dollars in transactions? Interested to build highly scalable next generation systems that could utilize Amazon Cloud? Massive data volume + complex business rules in a highly distributed and service oriented architecture, a world class information collection and delivery challenge. Our challenge is to deliver the software systems which accurately capture, process, and report on the huge volume of financial transactions that are generated each day as millions of customers make purchases, as thousands of Vendors and Partners are paid, as inventory moves in and out of warehouses, as commissions are calculated, and as taxes are collected in hundreds of jurisdictions worldwide. Key job responsibilities • Understand the business and discover actionable insights from large volumes of data through application of machine learning, statistics or causal inference. • Analyse and extract relevant information from large amounts of Amazon’s historical transactions data to help automate and optimize key processes • Research, develop and implement novel machine learning and statistical approaches for anomaly, theft, fraud, abusive and wasteful transactions detection. • Use machine learning and analytical techniques to create scalable solutions for business problems. • Identify new areas where machine learning can be applied for solving business problems. • Partner with developers and business teams to put your models in production. • Mentor other scientists and engineers in the use of ML techniques. A day in the life • Understand the business and discover actionable insights from large volumes of data through application of machine learning, statistics or causal inference. • Analyse and extract relevant information from large amounts of Amazon’s historical transactions data to help automate and optimize key processes • Research, develop and implement novel machine learning and statistical approaches for anomaly, theft, fraud, abusive and wasteful transactions detection. • Use machine learning and analytical techniques to create scalable solutions for business problems. • Identify new areas where machine learning can be applied for solving business problems. • Partner with developers and business teams to put your models in production. • Mentor other scientists and engineers in the use of ML techniques. About the team The FinAuto TFAW(theft, fraud, abuse, waste) team is part of FGBS Org and focuses on building applications utilizing machine learning models to identify and prevent theft, fraud, abusive and wasteful(TFAW) financial transactions across Amazon. Our mission is to prevent every single TFAW transaction. As a Machine Learning Scientist in the team, you will be driving the TFAW Sciences roadmap, conduct research to develop state-of-the-art solutions through a combination of data mining, statistical and machine learning techniques, and coordinate with Engineering team to put these models into production. You will need to collaborate effectively with internal stakeholders, cross-functional teams to solve problems, create operational efficiencies, and deliver successfully against high organizational standards.
US, NY, New York
Fauna Robotics is building capable, safe, and delightful robots for everyday life, and voice is one of the most natural ways people will interact with them. Cloud speech and language models are good and getting better, but they can only work with the audio they receive, and a robot is a hard place to listen. Its microphones sit beside motors, fans, and moving joints. It speaks through its own loudspeaker while people talk over it. It moves, turns, and shares a room with several people at once. We are hiring a Principal Audio Scientist to be Fauna's technical authority on how our robots hear. You will design, prototype, and ship the hardest algorithms in the robot's audio system. You will set the audio architecture that other engineers build on, shape hardware decisions across robot generations, mentor the engineers and scientists working on audio, and be the person teams come to when the robot can't hear. Key job responsibilities - Set the long-range science roadmap and technical architecture for the robot's audio system, and serve as Fauna's primary technical authority on robot hearing - Design and implement suppression of the robot's own noise from motors, fans, and moving joints - Design and implement echo cancellation for the robot's own voice, so people can interrupt it naturally - Develop multi-microphone processing that holds up as the robot and the people around it move, including locating who is speaking so the robot can turn toward them - Make on-robot listening decisions robust to internal and external noise sources: wake word, voice activity, and whether speech is directed at the robot - Drive microphone and speaker placement, enclosure acoustics, and vibration isolation decisions with mechanical, electrical, and industrial design, backed by your own measurements - Design the robot-specific data collection and evaluation methods to validate the performance of our audio design - Present audio science and its tradeoffs to senior leadership and partner teams - Mentor scientists and engineers, raising the scientific bar for audio across the organization through design reviews, code reviews, and hiring
US, NY, New York
We are seeking a Robotics/AI Motor Control Scientist to develop cutting-edge machine learning algorithms for motor control systems in robots. In this role, you will focus on creating and optimizing intelligent motor control strategies to enable robots to perform complex, whole-body tasks. Your contributions will be essential in advancing robotics by enabling fluid, reliable, and safe interactions between robots and their environments. Key job responsibilities - Develop controllers that leverage reinforcement learning, imitation learning, or other advanced AI techniques to achieve natural, robust, and adaptive motor behaviors - Collaborate with multi-disciplinary teams to integrate motor control systems with robotic hardware, ensuring alignment with real-world constraints such as actuator dynamics and energy efficiency - Use simulation and real-world testing to refine and validate control algorithms - Stay updated on advancements in robotics, AI, and control systems to apply advanced techniques to robotic motion challenges - Lead technical projects from conception through production deployment - Mentor junior scientists and engineers - Bridge research initiatives with practical engineering implementation About the team Fauna Robotics, an Amazon company, is building capable, safe, and genuinely delightful robots for everyday life. Our goal is simple: make robots people actually want to live and interact with in everyday human spaces. We believe that future won’t arrive until building for robotics becomes far more accessible. Today, too much effort is spent reinventing the fundamentals. We’re changing that by developing tightly integrated hardware and software systems that make it faster, safer, and more intuitive to create real-world robotic products. Our work spans the full stack: mechanical design, control systems, dynamic modeling, and intelligent software. The focus is not just functionality, but experience. We’re building robots that feel responsive, expressive, and genuinely useful. At Fauna, you’ll work at the frontier of this space, helping define how robots move, manipulate, and interact with people in natural environments. It’s an opportunity to solve hard problems across hardware and software with a team focused on making robotics accessible and joyful to build. If you care about making robotics real for everyone and building systems that are as delightful as they are capable, we’re interested in hearing from you. an opportunity to solve hard problems across hardware and software with a team focused on making robotics accessible and joyful to build. If you care about making robotics real for everyone and building systems that are as delightful as they are capable, we’re interested in hearing from you.
US, CA, Sunnyvale
Amazon's Artificial General Intelligence (AGI) organization is seeking an Applied Scientist III to advance the science of Responsible AI evaluation for large language models and generative AI. In this role, you will lead the design and development of rigorous evaluation methods, benchmarks, and metrics that measure the safety, fairness, robustness, and trustworthiness of frontier models. You will work with large-scale datasets, modern deep learning frameworks, and world-class scientists and engineers to turn research into evaluation systems that shape model launch decisions at Amazon scale. Key job responsibilities - Lead the design and implementation of evaluation frameworks, benchmarks, and metrics for responsible AI, including safety, fairness, robustness, and harmful content. - Build scalable automated evaluation pipelines for large language models, including model-based and human-in-the-loop evaluation. - Partner with pretraining, post-training, and product teams to translate evaluation results into model improvements and launch decisions. - Conduct rigorous experimentation and statistical analysis, and publish research at top venues. - Mentor junior scientists and help raise the scientific bar of the team. - Champion responsible AI practices across the model development lifecycle. About the team The AGI Responsible AI (RAI) team builds the science and systems that make Amazon's large language models safe, fair, and trustworthy. We work on problems spanning safety evaluation, content moderation, watermarking, bias mitigation, and alignment. Our team values scientific rigor, customer obsession, and rapid iteration, and we collaborate closely with pretraining, post-training, and product teams across AGI.
IN, KA, Bengaluru
Amazon Prime Video is looking for passionate, talented, and inventive Applied Scientists with a strong machine learning background to help build industry-leading Audio, Video and Language technology using GenAI. As part of the AI team in Amazon Prime Video, you will work on developing and fine-tuning LLMs and developing new large multimodal models to advance state-of-the-art in the translation, transcription and generation of multingual cinematic content. You will use the latest GenAI techniques to create and assess cinematic content like videos, audio, and captions, and more. You will design novel approaches to integrate speech, vision, and text for holistic understanding and generation of content, and Collaborate cross-functionally with product managers, engineers, and domain experts to define and deliver high-impact solutions. Your work will directly impact millions of our customers in the form of products and services that make use of speech and language technology on Prime Video content. Key job responsibilities We are looking to hire a top notch scientist with strong machine learning background, sharp research acumen and a hunger for solving ambiguous and challenging science problems. This person will work on developing and fine-tuning LLMs and developing new large multimodal models to advance state-of-the-art in the translation, transcription and generation of multingual cinematic content. They will be required to use the latest GenAI advancements to create and assess cinematic content like videos, audio, and captions, and more. Flavor of problems we are trying to solve: a) generating high quality subtitles in many languages meeting high customer satisfaction bar, b) developing foundation model capable of enabling many down stream tasks in localization and accessibility space, c) generating captions with plot pertinency, character identification etc. d) developing solution for assessing quality of subtitles, captions and metadata unified across different content types, linguistic vs non-linguistic aspects, different languages, and which mimics end user satisfaction while at the same time is explainable. About the team By 2027, Prime Video (PV) aspires to be the primary destination for customers to watch and the primary channel for partners to distribute premium cinematic content. The Content Reasoning, Enrichment & Localization (CoREL) science team’s mission is to develop AI to deeply understand the diverse facets of Prime Video's content across multiple languages, and to create content that power immersive, cinematic experiences. To drive this mission and delight our customers, the CoREL Science Bangalore team addresses the localization and accessibility aspects by a) understanding fine-grained information contained in content to enable generation of localized assets in different languages, b) generating subtitles and accessibility assets like captions across different content types and languages, c) assessing quality of generated (1P) and partner submitted (3P) assets and metadata to ensure defect free customer experience on PV.
CH, Zurich
RIVR, an Amazon company, is building Physical AI by deploying autonomous robots for real-world doorstep delivery. Operating daily in diverse urban environments, RIVR's robots continuously learn from and navigate the millions of scenarios encountered during deliveries. By owning the full stack from software to hardware, RIVR is purpose-built for safety, reliability, and the customer from day one. Reinforcement learning is transforming our robotic intelligence, enabling autonomous behavior without human guidance. We are seeking a Senior AI Engineer with deep expertise in reinforcement learning and deep learning, including supervised and self-supervised learning with a focus on dexterous manipulation. Your role will involve leveraging both simulated and real-world data to address practical challenges in dynamic grasping, contact-rich manipulation, and object interaction. If you are passionate about advancing AI and developing innovative solutions, join us in shaping the future of intelligent robotics. Key job responsibilities Develop cutting-edge reinforcement learning algorithms to enable robust, contact-rich dexterous manipulation, translating vision, depth, tactile, and proprioceptive sensor input into precise end-effector and joint-level motor commands. Design, test, and refine algorithms to solve complex real-world manipulation challenges, such as handling diverse package form factors, dynamic hand-offs, and operating door handles or latches. Collaborate with the foundation model team to innovate methods that leverage both simulated and real-world data.