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, NY, New York
The Sponsored Products and Brands team at Amazon Ads is re-imagining the advertising landscape through generative AI technologies, revolutionizing how millions of customers discover products and engage with brands across Amazon.com and beyond. We are at the forefront of re-inventing advertising experiences, bridging human creativity with artificial intelligence to transform every aspect of the advertising lifecycle from ad creation and optimization to performance analysis and customer insights. We are a passionate group of innovators dedicated to developing responsible and intelligent AI technologies that balance the needs of advertisers, enhance the shopping experience, and strengthen the marketplace. If you're energized by solving complex challenges and pushing the boundaries of what's possible with AI, join us in shaping the future of advertising. About the team SPB Agent team's vision is to build a highly personalized and context-aware agentic advertiser guidance system that seamlessly integrates Large Language Models (LLMs) with sophisticated tooling, operating across all experiences. The SPB-Agent is the central agent that interfaces with advertisers across Ads Console, Selling Partner portals (Seller Central, KDP, Vendor Central), and internal Sales systems. We identify high-impact opportunities spanning from strategic product guidance to granular optimization and deliver them through personalized, scalable experiences grounded in state-of-the-art agent architectures, reasoning frameworks, sophisticated tool integration, and model customization approaches including fine-tuning, MCP, and preference optimization. This presents an exceptional opportunity to shape the future of e-commerce advertising through advanced AI technology at unprecedented scale, creating solutions that directly impact millions of advertisers.
DE, BE, Berlin
At Audible, we believe stories have the power to transform lives. It’s why we work with some of the world’s leading creators to produce and share audio storytelling with our millions of global listeners. We are dreamers and inventors who come from a wide range of backgrounds and experiences to empower and inspire each other. Imagine your future with us. ABOUT THIS ROLE As an Applied Scientist, you will solve large complex real-world problems at scale, draw inspiration from the latest science and technology to empower undefined/untapped business use cases, delve into customer requirements, collaborate with tech and product teams on design, and create production-ready models that span various domains, including Machine Learning (ML), Artificial Intelligence (AI) and Generative AI, Natural Language Processing (NLP), Reinforcement Learning (RL), real-time and distributed systems. ABOUT YOU Your work will focus on inventing or adapting scientific approaches, models, and algorithms driven by customer needs at the project level. You will develop components and/or end-to-end solutions that are deployed into production or directly support production systems, delivering consistently high-quality work that meets both scientific and engineering best practices. You will develop reusable science components and services that resolve architecture deficiencies and customers’ pain points, while making technical trade-offs for long-term/short- term. You will work semi-autonomously to deliver solutions, contribute to research papers at peer-reviewed venues when appropriate, and document your work thoroughly to enable others to understand and reproduce it. Your decision-making will consistently incorporate robust, data-driven business and technical judgment. You will collaborate with other scientists to raise the bar of both scientific and engineering complexity for the team and to foster valuable scientific partnership opportunities to help/guide science decisions. We work in a highly collaborative, fast-paced environment where scientists, engineers, and product managers work to test and build scalable foundational capabilities, as well as customer facing experiences. You will have the opportunity to innovate and think big within your projects scope, implement optimization services and algorithms, and influence the experiences of millions of customers. We are looking for a results-oriented Applied Scientist with deep knowledge in ML, NLP, Deep Learning, GenAI, and/or large-scale distributed computation. As an Applied Scientist, you will... - Understand use cases across the business and adopt/extend/design/invent solutions/models that are scalable, efficient, and automated for difficult problems that are not well defined - Work closely with fellow scientists and software engineers (at Audible and Amazon) to build and productionize models, deliver novel and highly impactful features - Review models of peers for the purpose of reducing and managing risk to the business, while improving customer experience - Design, develop, and deploy modeling techniques and solutions for Content Understanding, Recommendations, GenAI-based product features, by employing a wide range of methodologies, working from simple to complex - Contribute to initiatives that employ the most recent advances in ML/AI in a fast-paced, experimental environment - Push the boundary of innovation ABOUT AUDIBLE Audible is the leading producer and provider of audio storytelling. We spark listeners’ imaginations, offering immersive, cinematic experiences full of inspiration and insight to enrich our customers daily lives. We are a global company with an entrepreneurial spirit. We are dreamers and inventors who are passionate about the positive impact Audible can make for our customers and our neighbors. This spirit courses throughout Audible, supporting a culture of creativity and inclusion built on our People Principles and our mission to build more equitable communities in the cities we call home.
IN, KA, Bengaluru
This position is based in Bangalore, India The Last Mile team helps get packages from delivery stations to a customer’s doorstep. To provide new innovations for customers, awe are inventing the next-generation smart delivery operation. We are combining innovative mobile and IoT technologies, data streams (video, vehicle telematics, location, and presence), together with machine learning models and algorithms – all to create solutions that allow us to deliver faster, and with more confidence. Playing a key role in the Last Mile Driver Experience team, as a Applied Scientist you will be responsible for building machine learning models and algorithms in areas including mapping and location, pattern detection in sensor data, and computer vision. Using your research, you will work with your engineering and product management peers to drive designs from ideation through development and into production. You will bring your experience of research for similar products and solutions, preferably in consumer or industrial verticals. This role requires autonomy and an ability to deliver results, often within the ambiguity of building a v1 product. You will need to work efficiently to build the right things with limited guidance, raising the bar to create an amazing experience for our customers.
ES, M, Madrid
Amazon's EU International Technology (EU INTech) organisation is creating new ways for customers to discover products through innovative customer experiences. We are a science-only team within EU INTech, responsible for designing and developing AI/ML science solutions that support business needs across Amazon's global search and discovery experiences. Our mission is to make Amazon navigation easier for customers worldwide. We achieve this through two strategic pillars: making Amazon navigation more visual and improving Amazon navigation with more inspiring discovery tools and narrowing navigation. To support this vision, we build and deploy AI/ML models that surface the most relevant content to hundreds of millions of Amazon customers worldwide. Our team comprises Applied Scientists and we partner with other teams, collaborating with ML Engineers, Software Developers, Product Managers, Technical Product Managers, and UX Designers. We are located in the Madrid Technical Hub. We are looking for Applied Scientists who are passionate about solving highly ambiguous and challenging problems at global scale. This is a hands-on, end-to-end applied science role where you will own the full lifecycle of science solutions — from business problem analysis and science plan design, through development and experimentation, to production deployment. We are looking for AI/ML experts with knowledge on ranking, computer vision, recommendation systems, search, and customer experience design. What makes this role unique: • End-to-end ownership – You will analyse business problems, map them to science plans, and design and develop solutions from ideation to production. We are owners of the full science lifecycle. • Applied science with a research edge – While our focus is on delivering applied science solutions that drive measurable business impact, our team actively pushes the state of the art in areas such as computer vision and Generative AI. • Hands-on execution – We need scientists who thrive in building, experimenting, and shipping. What are we looking for? • A scientist who can independently analyse any business problem and design a rigorous science approach to solve it • Strong hands-on engineering skills — you build and ship, not just theorise • Deep expertise in one or more of: computer vision, generative AI, recommendation systems, ranking, or NLP • Experience taking ML models from research to production at scale • Comfort with ambiguity and the ability to structure complex, undefined problems • A passion for customer-centric innovation and measurable impact • A strong communicator capable to adapt the message from a science audience, to engineering or leadership Key job responsibilities • Analyse complex business problems and translate them into well-defined science plans with clear milestones and success criteria • Design, develop, and deliver ML/AI models end-to-end — from research and prototyping through to production systems at Amazon scale and extending solutions going beyond the state of the art • Work with state-of-the-art models in computer vision, ranking and generative AI to power new customer experiences globally • Own major science challenges for the team, driving solutions from ideation through experimentation to production deployment • Collaborate with a variety of roles and partner teams around the world to deliver integrated solutions • Influence scientific direction and best practices across the team • Maintain high quality standards on team deliverables • Contribute to expanding the state of the art in computer vision, ranking and GenAI through publications and internal knowledge sharing
ES, M, Madrid
Amazon's EU International Technology (EU INTech) organisation is creating new ways for customers to discover products through innovative customer experiences. We are a science-only team within EU INTech, responsible for designing and developing AI/ML science solutions that support business needs across Amazon's global search and discovery experiences. Our mission is to make Amazon navigation easier for customers worldwide. We achieve this through two strategic pillars: making Amazon navigation more visual and improving Amazon navigation with more inspiring discovery tools and narrowing navigation. To support this vision, we build and deploy AI/ML models that surface the most relevant content to hundreds of millions of Amazon customers worldwide. Our team comprises Applied Scientists and we partner with other teams, collaborating with ML Engineers, Software Developers, Product Managers, Technical Product Managers, and UX Designers. We are located in the Madrid Technical Hub. We are looking for Applied Scientists who are passionate about solving highly ambiguous and challenging problems at global scale. This is a hands-on, end-to-end applied science role where you will own the full lifecycle of science solutions — from business problem analysis and science plan design, through development and experimentation, to production deployment. We are looking for AI/ML experts with knowledge on ranking, computer vision, recommendation systems, search, and customer experience design. What makes this role unique: • End-to-end ownership – You will analyse business problems, map them to science plans, and design and develop solutions from ideation to production. We are owners of the full science lifecycle. • Applied science with a research edge – While our focus is on delivering applied science solutions that drive measurable business impact, our team actively pushes the state of the art in areas such as computer vision and Generative AI. • Hands-on execution – We need scientists who thrive in building, experimenting, and shipping. What are we looking for? • A scientist who can independently analyse any business problem and design a rigorous science approach to solve it • Strong hands-on engineering skills — you build and ship, not just theorise • Deep expertise in one or more of: computer vision, generative AI, recommendation systems, ranking, or NLP • Experience taking ML models from research to production at scale • Comfort with ambiguity and the ability to structure complex, undefined problems • A passion for customer-centric innovation and measurable impact • A strong communicator capable to adapt the message from a science audience, to engineering or leadership Key job responsibilities • Analyse complex business problems and translate them into well-defined science plans with clear milestones and success criteria • Design, develop, and deliver ML/AI models end-to-end — from research and prototyping through to production systems at Amazon scale and extending solutions going beyond the state of the art • Work with state-of-the-art models in computer vision, ranking and generative AI to power new customer experiences globally • Own major science challenges for the team, driving solutions from ideation through experimentation to production deployment • Collaborate with a variety of roles and partner teams around the world to deliver integrated solutions • Influence scientific direction and best practices across the team • Maintain high quality standards on team deliverables • Contribute to expanding the state of the art in computer vision, ranking and GenAI through publications and internal knowledge sharing
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) Customization Team is seeking a highly skilled and experienced Applied Scientist to support adoption and enable customization of Amazon Nova. The role focuses on developing state-of-the-art services and tools for model customization, including supervised fine-tuning, reinforcement learning, and knowledge distillation across large language models. As an Applied Scientist, you will play a important role in developing advanced customization capabilities that enable enterprises to build highly performant application-specific models without the need for training models from scratch. Your work will directly impact how companies leverage Amazon Nova models for their specific use cases. Key job responsibilities - Contribute to the development of novel customization techniques including extended post-training, continued pre-training, and advanced knowledge distillation - Collaborate with cross-functional teams to design and implement enterprise-ready tooling for various training techniques on Amazon SageMaker - Design and execute experiments to optimize model accuracy, latency, and cost across different customization approaches (SFT, DPO, PPO) - Develop and enhance preference learning algorithms and training curricula for customer-specific applications - Create robust evaluation frameworks for assessing model performance across different domains and use cases - Contribute to the development of the Responsible AI toolkit, including creating training and evaluation datasets for model alignment - Design and implement secure access mechanisms for early model checkpoints and weights - Communicate technical insights and results to both technical and non-technical stakeholders through presentations and documentation
IN, KA, Bengaluru
Amazon is seeking a passionate and inventive Applied Scientist II with a strong machine learning background to build industry-leading Speech and Language technology. Our mission is to deliver delightful customer experiences by advancing Automatic Speech Recognition (ASR), Natural Language Understanding (NLU), Machine Learning (ML), and Computer Vision (CV). You will work alongside internationally recognized experts to develop novel algorithms and modeling techniques that advance the state-of-the-art in human language technology. Your work will directly impact millions of customers through products and services powered by speech and language technology. You will gain hands-on experience with Amazon's heterogeneous speech, text, and structured data sources, and leverage large-scale computing resources to accelerate advances in spoken language understanding. We are hiring across all areas of human language technology: ASR, Machine Translation (MT), NLU, Text-to-Speech (TTS), Dialog Management, and Computer Vision. We also seek talent experienced in building large-scale, high-performing systems. Key job responsibilities Basic Qualifications PhD or M.Tech in Computer Science, Electrical Engineering, Mathematics, or Physics with specialization in one or more of: speech recognition, natural language processing, machine translation, time series analysis, signal processing, or machine learning 1-2 years of industry or research experience (including internships, co-ops, or post-doctoral work) in applied ML or related areas Proficiency in programming languages such as Python, C/C++, or Java Strong foundation in machine learning fundamentals and statistical modeling Preferred Qualifications Experience building speech recognition, machine translation, or natural language processing systems (e.g., commercial products, government projects, or published research with working prototypes) Hands-on experience with deep learning frameworks (e.g., PyTorch, TensorFlow) Track record of publications in top-tier conferences (e.g., NeurIPS, ICML, ACL, Interspeech, CVPR) Scientific thinking with demonstrated ability to innovate and contribute to advancing the field Solid software development practices and experience shipping production-quality code Strong written and verbal communication skills A day in the life 0
US, CA, San Jose
Are you excited about making business decisions using science and data? Are you interested in supporting consumer device concepts from idea inception to launch? Do you want to work on a Science Product team focused on scaling statistics and econometrics with custom tools? If so, this may be the role for you! Amazon.com strives to be Earth's most customer-centric company. The Amazon Devices and Services team focuses on delighting customer by enabling seamless functionality in supplying, entertaining, and managing the home -- and beyond. We seek and hire the world's brightest minds, offering them a fast-paced, technologically-sophisticated, and friendly work environment, where economic theory meets real-world industry. The Decision Science team in Devices owns demand estimates and pricing recommendations of concept devices before customers know they exist. We support devices and services ranging from Echo Frames to Kindle Paperwhite to Blink Video Camera …all prior to launch. We are a cross-functional Product team working to scale Econometrics through Amazon and beyond by incorporating Science into internal facing tools and making it easier for others to do so as well. In this role, you will have input in decision meetings with Amazon senior leadership, which include go/no-go decisions for brand new devices and services and build volume decisions for manufacture prior to receiving any customer signal. You will have direct input to pricing decisions. You will leverage Science and Tools produced by the Decision Science team such as conjoint demand models to produce these recommendations. You will work with Scientists, Economists, Product Managers, and Software Developers to provide meaningful feedback about stakeholder problems to inform business solutions and increase the velocity, quality, and scope behind our recommendations. You will also have the opportunity to work on special projects to both guide the business and advance your own knowledge and understanding of specific topics. Key job responsibilities Applies expertise to develop econometric/machine learning models to measure the demand of devices and the business; Reviews models and results for other scientists, mentors junior scientists; Generates economic insights for the Devices and Services business and work with stakeholders to run the business for effectively; Describes strategic importance of vision inside and outside of team; and, Identifies business opportunities, defines the problem and how to solve it; Engages with senior scientists, business leadership outside Devices and Services to understand interplay between different business units.
AU, VIC, Melbourne
Are you excited about leveraging and extending state-of-the-art Deep Learning, Information Retrieval, Natural Language Processing, Computer Vision algorithms to solve customer problems at the scale of Amazon? As an Applied Scientist Intern, you will be working in the Melbourne office in a fast-paced, cross-disciplinary team of experienced R&D scientists. You will take on complex problems, work on solutions that leverage existing academic and industrial research, and utilize your own out-of-the-box pragmatic thinking. In addition to coming up with novel solutions and prototypes, you may even deliver these to production in customer facing products. Key job responsibilities - Develop novel solutions and build prototypes - Work on complex problems in Deep Learning and Generative AI - Contribute to research that could significantly impact Amazon operations - Collaborate with a diverse team of experts in a fast-paced environment - Present your research findings to both technical and non-technical audiences - Collaborate with scientists on writing and submitting papers to top ML conferences, e.g. NeurIPS, ICML, ICLR, AISTATS, ACL ICCV, CVPR, KDD. Key Opportunities: - Work in a team of ML scientists to solve applied science problems at the scale of Amazon - Access to Amazon services and hardware - Potentially deliver solutions to production in customer-facing applications - Opportunities to be hired full-time after the internship Join us in shaping the future of AI at Amazon. Apply now and turn your research into real-world solutions!
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 subscriptions such as Apple TV+, HBO Max, Peacock, 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 team member, 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! Key job responsibilities - Lead research and development of speech and audio generation technology and end-to-end speech-to-speech architecture - Develop audio processing solutions for production environments, including source separation, enhancement, and mixing - Define the research roadmap for your area, identify high-impact problems, and communicate technical direction to senior leadership - Publish research, contribute to the broader scientific community, and bring external advances into production systems - Hire, mentor, and develop applied scientists. Grow the team's capabilities to meet evolving customer and business needs About the team This team's mission is to deeply understand all content and empower all customers with relevant language options, innovative accessibility assists, and rich title-information across all their content-experiences on Prime Video. We create and publish content on-time that's meaningful, accurate, and accessible to every customer globally. We delight our customers by pushing the boundaries of content understanding and enrichment. Through inclusion and innovation, we do the most fulfilling work of our career.