New method improves knowledge-graph-based question answering

Replacing hand annotation with a machine learning component reduces labor, while an intersection operation enables multiple-entity queries.

Question answering is a popular task in natural-language processing, where models are given questions such as “What city is the Mona Lisa in?” and trained to predict correct answers, such as “Paris”. 

One way to train question-answering models is to use a knowledge graph, which stores facts about the world in a structured format. Historically, knowledge-graph-based question-answering systems required separate semantic-parsing and entity resolution models, which were costly to train and maintain. But since 2020, the field has been moving toward differentiable knowledge graphs, which allow a single end-to-end model to take questions as inputs and directly output answers.

At this year’s Conference on Empirical Methods in Natural Language Processing (EMNLP), we presented two extensions to end-to-end question answering using differentiable knowledge graphs.

In “End-to-end entity resolution and question answering using differentiable knowledge graphs”, we explain how to integrate entity resolution into a knowledge-graph-based question-answering model, so that it will be performed automatically. This is a potentially labor-saving innovation, as some existing approaches require hand annotation of entities.

And in “Expanding end-to-end question answering on differentiable knowledge graphs with intersection”, we explain how to handle queries that involve multiple entities. In experiments involving two different datasets, our approach improved performance on multi-entity queries by 14% and 19%.

Traditional approaches

In a typical knowledge graph, the nodes represent entities, and the edges between nodes represent relationships between entities. For example, a knowledge graph could contain the fact “Mona Lisa | exhibited at | Louvre Museum”, which links the entities “Mona Lisa” and “Louvre Museum” with the relationship “exhibited at”. Similarly, the graph could link the entities “Louvre Museum” and “Paris” with the relationship “located in”. A model can learn to use the facts in the knowledge graph to get to answer questions.

GraphTraversal_Cropped.gif
Question answering through traversal of links in a knowledge graph.

Traditional approaches to knowledge-graph-based question answering use a pipeline of models. First, the question goes into a semantic-parsing model, which is trained to predict queries. Queries can be thought of as instructions of what to do in the knowledge graph — for example, “find out where the ‘Mona Lisa’ is exhibited, then look up what city it’s in.” 

Next, the question goes into an entity resolution model, which links parts of the sentence, like “Mona Lisa”, to IDs in the knowledge graph, like Q12418.

PipelinedApproach.gif
The pipelined approach to knowledge-graph-based question answering.

While this pipeline approach works, it does have some flaws. Each model is trained independently, so it has to be evaluated and updated separately, and each model requires annotations for training, which are time-consuming and expensive to collect. 

End-to-end question answering

End-to-end question answering is a way to rectify these flaws. During training, an end-to-end question answering model is given the question and the answer but no instructions about what to do in the knowledge graph. Instead, the model learns the instructions based on the correct answer. 

In 2020, Cohen et al. proposed a way to perform end-to-end question answering using differentiable knowledge graphs, which represent knowledge graphs as tensors and queries as differentiable mathematical operations. This allows for fully differentiable training, so that the answer alone provides a training signal to update all parts of the model.

DifferentiableApproach.gif
End-to-end question answering with a differentiable knowledge graph.

“End-to-end entity resolution and question answering using differentiable knowledge graphs”

In our first paper, we extend end-to-end models to include the training of an entity resolution component. Previous work left the task of entity resolution (linking “Mona Lisa” to Q12418) out of the scope of the model, relying instead on a separate entity resolution model or hand annotation of entities. 

We propose a way to train entity resolution as part of the question-answering model. To do this, we start with a baseline model similar to the implementation by Cohen et al., which has an encoder-decoder structure and an attention mechanism that takes in a question and returns predicted answers with probabilities.

We add entity resolution to this baseline model by first introducing a span detection component. This identifies all the possible parts of the sentence (spans) that could refer to an entity. For example, in the question “Who wrote films starring Tom Hanks?”, there are multiple spans, such as “films” or “Tom”, and we want the model to learn to give a higher score to the correct span, “Tom Hanks”. 

Entity resolution.png
An end-to-end model that jointly learns knowledge-graph-based question answering and entity resolution.

Then for each of the identified spans, our model ranks all the possible entities in the knowledge graph that the span could refer to. For example, “Tom Hanks” is also the name of a seismologist and a theologian, but we want the model to learn to give the actor a higher score.

Span detection and entity resolution happen jointly in a new entity resolution component, which returns possible entities with scores. Finding the entities in the knowledge graph and following the paths from the inference component yields the predicted answers. We believe that this joint modeling is responsible for our method’s increase in efficiency.

In our experiments, we use two English question-answering datasets. We find that although our ER and E2E models perform slightly worse than the baseline, they do come very close, with differences of about 7% and 5% on the two datasets. 

This is an impressive result, since the baseline model uses hand-annotated entities from the datasets, while our model is learning entity resolution jointly with question answering. With these findings, we demonstrate the feasibility of learning to perform entity resolution and multihop inference in a single end-to-end model.

“Expanding end-to-end question answering on differentiable knowledge graphs with intersection”

In our second paper, we extend end-to-end models to handle more complex questions with multiple entities. Take for example the question “Who did Natalie Portman play in Star Wars?”. This question has two entities, Natalie Portman and Star Wars. 

Previous end-to-end models were trained to follow paths originating with one entity in a knowledge graph. However, this is not enough to answer questions with multiple entities. If we started at “Natalie Portman” and found all the roles she played, we would get roles that were not in Star Wars. If we started at Star Wars and found all the characters, we would get characters that Natalie Portman didn’t play. What we need is an intersection of the characters played by Natalie Portman and the characters in Star Wars. 

To handle questions with multiple entities, we expand an end-to-end model with a new operation: intersection. For each entity in the question, the model follows paths from the entity independently and arrives at an intermediate answer. Then the model performs intersection, which we implemented as the element-wise minimum of two vectors, to identify which entities the intermediate answers have in common. Only entities that appear in all intermediate answers are returned in the final answer. 

Multiple entities.png
An end-to-end knowledge-graph-based question-answering model for queries with multiple entities.

In our experiments, we use two English question-answering datasets. Our results show that introducing intersection improves performance over the baseline by 3.7% on one and 8.9% on the other. 

More importantly, we see that improved performance comes from better handling of questions with multiple entities, where the intersection model surpasses the baseline by over 14% on one dataset and by 19% on the other.

In future work, we plan to continue developing end-to-end models by improving entity resolution, so it’s competitive with hand annotation of entities; integrating entity resolution with intersection; and learning to handle more-complex operations, such as maximums/minimums and counts.

Related content

IN, HR, Gurugram
Work on ML teams building large-scale forecasting and optimization systems that power Amazon’s global transportation network and directly impact customer experience and cost. As an Applied Scientist II, you will set scientific direction, mentor applied scientists, and partner with engineering and product leaders to deliver production-grade ML solutions at massive scale. Key job responsibilities 1. Lead and grow a high-performing team of Applied Scientists, providing technical guidance, mentorship, and career development. 2. Define and own the scientific vision and roadmap for ML solutions powering large-scale transportation planning and execution. 3. Guide model and system design across a range of techniques, including tree-based models, deep learning (LSTMs, transformers), LLMs, and reinforcement learning. 4. Ensure models are production-ready, scalable, and robust through close partnership with stakeholders. Partner with Product, Operations, and Engineering leaders to enable proactive decision-making and corrective actions. 5. Own end-to-end business metrics, directly influencing customer experience, cost optimization, and network reliability. 6. Help contribute to the broader ML community through publications, conference submissions, and internal knowledge sharing. A day in the life Your day includes reviewing model performance and business metrics, guiding technical design and experimentation, mentoring scientists, and driving roadmap execution. You’ll balance near-term delivery with long-term innovation while ensuring solutions are robust, interpretable, and scalable. Ultimately, your work helps improve delivery reliability, reduce costs, and enhance the customer experience at massive scale.
DE, BE, Berlin
Are you excited about developing agentic AI, LLM and computer vision models that revolutionize Amazon's Fulfillment network? Are you looking for opportunities to apply state-of-the-art AI on real-world problems at truly vast scale? At Amazon Fulfillment Technologies and Robotics, we are on a mission to build high-performance autonomous systems that perceive and act to further improve our world-class customer experience — at Amazon scale. To this end, we are looking for an Applied Scientist who will build and deploy models that make smarter decisions on a wide array of multi-modal signals. Together, we will be pushing beyond the state of the art in optimizing one of the most complex systems in the world: Amazon's Fulfillment Network. Key job responsibilities In this role, you will build agentic AI solutions and multi-modal deep learning models that understand how products and packages flowing through Amazon’s fulfillment network. You will build models that solve challenging problems like understanding warehouse operations systems, or visual defect detection on Amazon's entire retail catalog (billions of different items, thousands of new items every day). You will work with a diverse set of very large multi-modal real-world datasets, including imagery, natural language and structured data. You will face a high level of research ambiguity and problems that require creative, ambitious, and inventive solutions. A day in the life AFT AI delivers the AI solutions that empower Amazon’s fulfillment network to make smarter decisions. You will work on an interdisciplinary project involving scientists and engineers with deep expertise in developing state-of-the-art AI solutions at scale. You will work with images, videos, natural language, and sequences of events from existing or new hardware. You will adapt state-of-the-art agentic AI, deep learning, language understanding and computer vision techniques to develop solutions for business problems in the Amazon Fulfillment Network. About the team Amazon Fulfillment Technologies (AFT) powers Amazon’s global fulfillment network. We invent and deliver software, hardware, and science solutions that orchestrate processes, robots, machines, and people. We harmonize the physical and virtual world so Amazon customers can get what they want, when they want it. AFT AI is spread across NA (Bellevue, WA) and Europe (Berlin, Germany). We are hiring candidates to work out of the Berlin location. Publicly available articles showcasing some of our work: - Visual Defect Detection: https://www.amazon.science/blog/novel-kaputt-dataset-sets-new-benchmark-for-large-scale-visual-defect-detection - Eluna: https://www.aboutamazon.com/news/operations/new-robots-amazon-fulfillment-agentic-ai
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Member of Technical Staff with a strong deep learning background, to build industry-leading Generative Artificial Intelligence (GenAI) technology with Large Language Models (LLMs) and multimodal systems. Key job responsibilities As a Member of Technical Staff with the AGI team, you will support the development of algorithms and modeling techniques, to advance the state of the art with LLMs. You will support the foundational model development in an applied research role, including model training, dataset design, and pre- and post-training optimization. Your work will directly impact our customers in the form of products and services that make use of GenAI technology. You will leverage Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in LLMs. About the team The AGI team has a mission to push the envelope in GenAI with LLMs and multimodal systems, in order to provide the best-possible experience for our customers.
US, WA, Seattle
The Sponsored Products and Brands (SPB) 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. This position will be part of the Conversational Ad Experiences team within the Amazon Advertising organization. Our cross-functional team focuses on designing, developing and launching innovative ad experiences delivered to shoppers in conversational contexts. We utilize leading-edge engineering and science technologies in generative AI to help shoppers discover new products and brands through intuitive, conversational, multi-turn interfaces. We also empower advertisers to reach shoppers, using their own voice to explain and demonstrate how their products meet shoppers' needs. We collaborate with various teams across multiple Amazon organizations to push the boundary of what's possible in these fields. We are seeking a science leader for our team within the Sponsored Products & Brands organization. You'll be working with talented scientists, engineers, and product managers to innovate on behalf of our customers. An ideal candidate is able to navigate through ambiguous requirements, working with various partner teams, and has experience in generative AI, large language models (LLMs), information retrieval, and ads recommendation systems. Using a combination of generative AI and online experimentation, our scientists develop insights and optimizations that enable the monetization of Amazon properties while enhancing the experience of hundreds of millions of Amazon shoppers worldwide. If you're fired up about being part of a dynamic, driven team, then this is your moment to join us on this exciting journey! Key job responsibilities - Serve as a tech lead for defining the science roadmap for multiple projects in the conversational ad experiences space powered by LLMs. - Build POCs, optimize and deploy models into production, run experiments, perform deep dives on experiment data to gather actionable learnings and communicate them to senior leadership - Work closely with software engineers on detailed requirements, technical designs and implementation of end-to-end solutions in production. - Work closely with product managers to contribute to our mission, and proactively identify opportunities where science can help improve customer experience - Research new machine learning approaches to drive continued scientific innovation - Be a member of the Amazon-wide machine learning community, participating in internal and external meetups, hackathons and conferences - Help attract and recruit technical talent, mentor scientists and engineers in the team
US, WA, Seattle
Amazon Customer Service (CS) Data Intelligence builds the data and Artificial Intelligence (AI) foundations for CS to ensure Amazon delivers the best customer service possible. CS Economics sits within CS DI and contributes to the CS knowledge base and decision frameworks. CS Economics seeks economists to apply economic methods to solve business problems. The ideal candidate will work with engineers and applied scientists to design models that leverage large scale and unstructured data, design scalable agents for non-tech CS partners to understand the impact of their actions, and propose mechanism designs to robustly match customers to our services. CS Economics is looking for optimistic critical-thinkers who combine a strong technical economic toolbox with a desire to learn from other disciplines, and who know how to execute and deliver on big ideas as part of an interdisciplinary technical team. Ideal candidates enjoy working in a team setting with individuals from diverse disciplines and backgrounds. They will work with teammates to develop scientific models and conduct data analysis, modeling, and experimentation that is necessary for estimating and validating models. They will work closely with engineering teams to develop scalable data resources to support rapid insights, and take successful models and findings into production as new products and services. They will be customer-centric and will communicate scientific approaches and findings to business leaders, listening to and incorporate their feedback, and delivering successful scientific solutions. Key job responsibilities - Design and conduct rigorous evaluations of CS actions - Develop experiments to evaluate product launches - Communicate complex findings to business stakeholders in clear, actionable terms - Work with engineering teams to develop scalable tools that automate and streamline evaluation processes A day in the life Work with teammates to apply economic methods to business problems, e.g., identify the appropriate research question and identification strategy, write code to estimate heterogeneous treatment effects or conduct experiment analysis, write and present a document with findings to business leaders. We collaborate with partner teams within and outside of CS throughout the process, from understanding their challenges, to developing a research agenda that will address those challenges, to help them implement solutions. About the team Amazon Customer Service (CS) Economics provides estimates and measures of the causal impact of CS actions on costs and benefits. We build agents and guide leadership to establish processes to scale valid experimentation, causal inference, and mechanism design.
US, WA, Seattle
Innovators wanted! Are you an entrepreneur? A builder? A dreamer? This role is part of an Amazon Special Projects team that takes the company’s Think Big leadership principle to the limits. If you’re interested in innovating at scale to address big challenges in the world, this is the team for you. As a Senior Applied Scientist on our team, you will focus on building state-of-the-art ML models for healthcare. Our team rewards curiosity while maintaining a laser-focus in bringing products to market. Competitive candidates are responsive, flexible, and able to succeed within an open, collaborative, entrepreneurial, startup-like environment. At the forefront of both academic and applied research in this product area, you have the opportunity to work together with a diverse and talented team of scientists, engineers, and product managers and collaborate with other teams. This role offers a unique opportunity to work on projects that could fundamentally transform healthcare outcomes. Key job responsibilities In this role, you will: • Design and implement novel AI/ML solutions for complex healthcare challenges • Drive advancements in machine learning and data science • Balance theoretical knowledge with practical implementation • Work closely with customers and partners to understand their requirements • Navigate ambiguity and create clarity in early-stage product development • Collaborate with cross-functional teams while fostering innovation in a collaborative work environment to deliver impactful solutions • Establish best practices for ML experimentation, evaluation, development and deployment • Partner with leadership to define roadmap and strategic initiatives You’ll need a strong background in AI/ML, proven leadership skills, and the ability to translate complex concepts into actionable plans. You’ll also need to effectively translate research findings into practical solutions. A day in the life You will solve real-world problems by getting and analyzing large amounts of data, generate insights and opportunities, design simulations and experiments, and develop statistical and ML models. The team is driven by business needs, which requires collaboration with other Scientists, Engineers, and Product Managers across the Special Projects organization. You will prepare written and verbal presentations to share insights to audiences of varying levels of technical sophistication. About the team We represent Amazon's ambitious vision to solve the world's most pressing challenges. We are exploring new approaches to enhance research practices in the healthcare space, leveraging Amazon's scale and technological expertise. We operate with the agility of a startup while backed by Amazon's resources and operational excellence. We're looking for builders who are excited about working on ambitious, undefined problems and are comfortable with ambiguity.
US, CA, Palo Alto
About Sponsored Products and Brands The Sponsored Products and Brands (SPB) 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 our team SPB Ad Response Prediction team is your choice, if you want to join a highly motivated, collaborative, and fun-loving team with a strong entrepreneurial spirit and bias for action. We are seeking an experienced and motivated Applied Scientist with machine learning engineering background who loves to innovate at the intersection of customer experience, deep learning, and high-scale machine learning systems. We are looking for a talented Applied Scientist with a strong background in machine learning engineering to join our team and help us grow the business. In this role, you will partner with a team of engineers and scientists to build advanced machine learning models and infrastructure, from training to inference, including emerging LLM-based systems, that deliver highly relevant ads to shoppers across all Amazon platforms and surfaces worldwide. Key job responsibilities As an Applied Scientist, you will: * Develop scalable and effective machine learning models and optimization strategies to solve business problems. * Conduct research on new machine learning modeling to optimize all aspects of Sponsored Products business. * Enhance the scalability, automation, and efficiency of large-scale training and real-time inference systems. * Pioneer the development of LLM inference infrastructure to support next-generation GenAI workloads at Amazon Ads scale.
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
We are seeking an Applied Scientist II to work on development of an AI-based data intelligence and classification platform that will redefine how security and privacy assessments and enforcement are conducted at scale. This mission-critical platform will leverage AI-driven autonomous agents to conduct proactive, intelligent security operations across the company. The platform will integrate deeply with internal security, privacy, engineering, and cloud-native tools to provide self-serve, automated insights, verifications, and enforcement mechanisms. This role requires strong technical expertise in AI/ML, LLMs, and distributed cloud infrastructure, as well as thought leadership to drive alignment across multiple teams, customers, and business units. This is an opportunity to shape the future of AI-driven security and privacy assurance at an enterprise scale, defining standards, influencing company-wide security posture, and leading technical innovation at the highest level. Key job responsibilities * Architect and define the next-generation data classification and search matching platform, leading the technical strategy for AI-driven security automation across applied science and engineering teams * Build on multi-agent LLM framework, and influence your organizations in adopting the promising approaches * Develop a highly scalable, traditional ML-based as well as LLM-based intelligent security agent framework that enables internal teams to automate processing of structured and unstructured data * Combine depth and breadth of domain expertise and provide technical leadership to the entire team while also doing hands-on work by diving deep into details to diagnose complex system performance problems. About the team The Data Categorization team helps Amazonians understand their data and govern it at scale and ensures experiences delivered by Amazon to our customers uphold our high security and privacy standards. The science team harnesses AI to strengthen Amazon’s privacy and security posture more efficiently and effectively.
US, CA, Sunnyvale
Amazon Music is an immersive audio entertainment service that deepens connections between fans, artists, and creators. From personalized music playlists to exclusive podcasts, concert livestreams to artist merch, Amazon Music is innovating at some of the most exciting intersections of music and culture. We offer experiences that serve all listeners with our different tiers of service: Prime members get access to all the music in shuffle mode, and top ad-free podcasts, included with their membership; customers can upgrade to Amazon Music Unlimited for unlimited, on-demand access to 100 million songs, including millions in HD, Ultra HD, and spatial audio; and anyone can listen for free by downloading the Amazon Music app or via Alexa-enabled devices. Join us for the opportunity to influence how Amazon Music engages fans, artists, and creators on a global scale. Amazon Music - Search Science team is seeking an experienced Applied Scientist who will join a team of experts in the field of machine learning, and work together to break new ground in the world of understanding and classifying different forms of music, and creating interactive experiences to help users find the music they are in the mood for. We work on machine learning problems for music classification, recommender systems, dialogue systems, NLP, and music information retrieval. You'll work in a collaborative environment where you can pursue applied research, with many peta-bytes of data, work on problems that haven’t been solved before, quickly implement and deploy your algorithmic ideas at scale, understand whether they succeed via statistically relevant experiments across millions of customers, and publish your research. You'll see the work you do directly improve the experience of Amazon Music customers on Alexa/Echo, mobile, and web. Key job responsibilities - Use machine learning, deep learning, LLMs and Agentic AI techniques to create scalable solutions for business problems - Analyze and extract relevant information from large amounts of Amazon's data to help automate and optimize key processes - Design, development and evaluation of AI models for predictive learning - Work closely with software engineering teams to drive model implementations and new feature creations - Establish scalable, efficient, automated processes for large scale data analyses, model development, model validation and model implementation - Research and implement novel machine learning and statistical approaches About the team Everyone on our team has a meaningful impact on product features, new directions in music streaming, and customer engagement. We are looking for new team members across a variety of job functions including software engineering/development, marketing, design, ops and more. Come join us as we make history by launching exciting new projects in the coming year.Our team is focused on building a personalized, curated, and seamless music experience. We want to help our customers discover up-and-coming artists, while also having access to their favorite established musicians. We build systems that are distributed on a large scale, spanning our music apps, web player, and voice-forward audio engagement on mobile and Amazon Echo devices, powered by Alexa to support our customer base. Amazon Music offerings are available in countries around the world, and our applications support our mission of delivering music to customers in new and exciting ways that enhance their day-to-day lives.