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

US, CA, Santa Clara
Join the next science and engineering revolution at Amazon's Delivery Foundation Model team, where you'll work alongside world-class scientists and engineers to pioneer the next frontier of logistics through advanced AI and foundation models. We are seeking an exceptional Senior Applied Scientist to help develop innovative foundation models that enable delivery of billions of packages worldwide. In this role, you'll combine highly technical work with scientific leadership, ensuring the team delivers robust solutions for dynamic real-world environments. Your team will leverage Amazon's vast data and computational resources to tackle ambitious problems across a diverse set of Amazon delivery use cases. Key job responsibilities - Design and implement novel deep learning architectures combining a multitude of modalities, including image, video, and geospatial data. - Solve computational problems to train foundation models on vast amounts of Amazon data and infer at Amazon scale, taking advantage of latest developments in hardware and deep learning libraries. - As a foundation model developer, collaborate with multiple science and engineering teams to help build adaptations that power use cases across Amazon Last Mile deliveries, improving experience and safety of a delivery driver, an Amazon customer, and improving efficiency of Amazon delivery network. - Guide technical direction for specific research initiatives, ensuring robust performance in production environments. - Mentor fellow scientists while maintaining strong individual technical contributions. A day in the life As a member of the Delivery Foundation Model team, you’ll spend your day on the following: - Develop and implement novel foundation model architectures, working hands-on with data and our extensive training and evaluation infrastructure - Guide and support fellow scientists in solving complex technical challenges, from trajectory planning to efficient multi-task learning - Guide and support fellow engineers in building scalable and reusable infra to support model training, evaluation, and inference - Lead focused technical initiatives from conception through deployment, ensuring successful integration with production systems- Drive technical discussions within the team and and key stakeholders - Conduct experiments and prototype new ideas - Mentor team members while maintaining significant hands-on contribution to technical solutions About the team The Delivery Foundation Model team combines ambitious research vision with real-world impact. Our foundation models provide generative reasoning capabilities required to meet the demands of Amazon's global Last Mile delivery network. We leverage Amazon's unparalleled computational infrastructure and extensive datasets to deploy state-of-the-art foundation models to improve the safety, quality, and efficiency of Amazon deliveries. Our work spans the full spectrum of foundation model development, from multimodal training using images, videos, and sensor data, to sophisticated modeling strategies that can handle diverse real-world scenarios. We build everything end to end, from data preparation to model training and evaluation to inference, along with all the tooling needed to understand and analyze model performance. Join us if you're excited about pushing the boundaries of what's possible in logistics, working with world-class scientists and engineers, and seeing your innovations deployed at unprecedented scale.
US, NY, New York
Join the next science and engineering revolution at Amazon's Delivery Foundation Model team, where you'll work alongside world-class scientists and engineers to pioneer the next frontier of logistics through advanced AI and foundation models. We are seeking an exceptional Senior Applied Scientist to help develop innovative foundation models that enable delivery of billions of packages worldwide. In this role, you'll combine highly technical work with scientific leadership, ensuring the team delivers robust solutions for dynamic real-world environments. Your team will leverage Amazon's vast data and computational resources to tackle ambitious problems across a diverse set of Amazon delivery use cases. Key job responsibilities - Design and implement novel deep learning architectures combining a multitude of modalities, including image, video, and geospatial data. - Solve computational problems to train foundation models on vast amounts of Amazon data and infer at Amazon scale, taking advantage of latest developments in hardware and deep learning libraries. - As a foundation model developer, collaborate with multiple science and engineering teams to help build adaptations that power use cases across Amazon Last Mile deliveries, improving experience and safety of a delivery driver, an Amazon customer, and improving efficiency of Amazon delivery network. - Guide technical direction for specific research initiatives, ensuring robust performance in production environments. - Mentor fellow scientists while maintaining strong individual technical contributions. A day in the life As a member of the Delivery Foundation Model team, you’ll spend your day on the following: - Develop and implement novel foundation model architectures, working hands-on with data and our extensive training and evaluation infrastructure - Guide and support fellow scientists in solving complex technical challenges, from trajectory planning to efficient multi-task learning - Guide and support fellow engineers in building scalable and reusable infra to support model training, evaluation, and inference - Lead focused technical initiatives from conception through deployment, ensuring successful integration with production systems- Drive technical discussions within the team and and key stakeholders - Conduct experiments and prototype new ideas - Mentor team members while maintaining significant hands-on contribution to technical solutions About the team The Delivery Foundation Model team combines ambitious research vision with real-world impact. Our foundation models provide generative reasoning capabilities required to meet the demands of Amazon's global Last Mile delivery network. We leverage Amazon's unparalleled computational infrastructure and extensive datasets to deploy state-of-the-art foundation models to improve the safety, quality, and efficiency of Amazon deliveries. Our work spans the full spectrum of foundation model development, from multimodal training using images, videos, and sensor data, to sophisticated modeling strategies that can handle diverse real-world scenarios. We build everything end to end, from data preparation to model training and evaluation to inference, along with all the tooling needed to understand and analyze model performance. Join us if you're excited about pushing the boundaries of what's possible in logistics, working with world-class scientists and engineers, and seeing your innovations deployed at unprecedented scale.
US, NY, New York
Are you a passionate Applied Scientist (AS) ready to shape the future of digital content creation? At Amazon, we're building Earth's most desired destination for creators to monetize their unique skills, inspire the next generation of customers, and help brands expand their reach. We build innovative products and experiences that drive growth for creators across Amazon's ecosystem. Our team owns the entire Creator product suite, ensuring a cohesive experience, optimizing compensation structures, and launching features that help creators achieve both monetary and non-monetary goals. Key job responsibilities As an AS on our team, you will: - Handle challenging problems that directly impact millions of creators and customers - Independently collect and analyze data - Develop and deliver scalable predictive models, using any necessary programming, machine learning, and statistical analysis software - Collaborate with other scientists, engineers, product managers, and business teams to creatively solve problems, measure and estimate risks, and constructively critique peer research - Consult with engineering teams to design data and modeling pipelines which successfully interface with new and existing software - Participate in design and implementation across teams to contribute to initiatives and develop optimal solutions that benefit the creators organization The successful candidate is a self-starter, comfortable with a dynamic, fast-paced environment, and able to think big while paying careful attention to detail. You have deep knowledge of an area/multiple areas of science, with a track record of applying this knowledge to deliver science solutions in a business setting and a demonstrated ability to operate at scale. You excel in a culture of invention and collaboration.
CA, ON, Toronto
About Sponsored Products and Brands 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 our team The Targeting and Recommendations team within Sponsored Products and Brands empowers advertisers with intelligent targeting controls and one-click campaign recommendations that automatically populate optimal settings based on ASIN data. This comprehensive suite provides advanced targeting capabilities through AI-generated keyword and ASIN suggestions, sophisticated targeting controls including Negative Targeting, Manual Targeting with Product Attribute Targeting (PAT) and Keyword Targeting (KWT), and Automated Targeting (ATv2). Our vision is to build a revolutionary, highly personalized and context-aware agentic advertiser guidance system that seamlessly integrates Large Language Models (LLMs) with sophisticated tooling, operating across both conversational and traditional ad console experiences while scaling from natural language queries to proactive, intelligent guidance delivery based on deep advertiser understanding, ultimately enhancing both targeting precision and one-click campaign optimization. Through strategic partnerships across Ad Console, Sales, and Marketing teams, we identify high-impact opportunities spanning from strategic product guidance to granular keyword 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 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. Key job responsibilities * Design and build targeting and 1 click recommendation agents to guide advertisers in conversational and non-conversational experience. * Design and implement advanced model and agent optimization techniques, including supervised fine-tuning, instruction tuning and preference optimization (e.g., DPO/IPO). * Collaborate with peers across engineering and product to bring scientific innovations into production. * Stay current with the latest research in LLMs, RL, and agent-based AI, and translate findings into practical applications. * Develop agentic architectures that integrate planning, tool use, and long-horizon reasoning. A day in the life As an Applied Scientist on our team, your days will be immersed in collaborative problem-solving and strategic innovation. You'll partner closely with expert applied scientists, software engineers, and product managers to tackle complex advertising challenges through creative, data-driven solutions. Your work will center on developing sophisticated machine learning and AI models, leveraging state-of-the-art techniques in natural language processing, recommendation systems, and agentic AI frameworks. From designing novel targeting algorithms to building personalized guidance systems, you'll contribute to breakthrough innovations
US, WA, Seattle
The AWS Supply Chain organization is looking for a Sr. Manager of Applied Science to lead science and data teams working on innovative AI-powered supply chain solutions. As part of the AWS Solutions organization, we have a vision to provide business applications, leveraging Amazon’s unique experience and expertise, that are used by millions of companies worldwide to manage day-to-day operations. We will accomplish this by accelerating our customers’ businesses through delivery of intuitive and differentiated technology solutions that solve enduring business challenges. We blend vision with curiosity and Amazon’s real-world experience to build opinionated, turnkey solutions. Where customers prefer to buy over build, we become their trusted partner with solutions that are no-brainers to buy and easy to use. Are you excited about developing state-of-the-art GenAI/Agentic AI based solutions for enterprise applications? As a Sr. Manager of Applied Scientist at AWS Supply Chain, you will bring AI advancements to customer facing enterprise applications. In this role, you will drive the technical vision and strategy for your team while fostering a culture of innovation and scientific excellence. You will be leading a fast-paced, cross-disciplinary team of researchers who are leaders in the field. You will take on challenging problems, distill real requirements, and then deliver solutions that either leverage existing academic and industrial research, or utilize your own out-of-the-box pragmatic thinking. In addition to coming up with novel solutions and prototypes, you may even need to deliver these to production in customer facing products. Key job responsibilities Building and mentoring teams of Applied Scientists, ML Engineers, and Data Scientists. Setting technical direction and research strategy aligned with business goals. Driving innovation in Supply Chains systems using AI/ML models and AI Agents. Collaborating with cross-functional teams to translate research into production. Managing project portfolios and resource allocation.
GB, MLN, Edinburgh
Do you want a role with deep meaning and the ability to make a major impact? As part of Intelligent Talent Acquisition (ITA), you'll have the opportunity to reinvent the hiring process and deliver unprecedented scale, sophistication, and accuracy for Amazon Talent Acquisition operations. ITA is an industry-leading people science and technology organization made up of scientists, engineers, analysts, product professionals and more, all with the shared goal of connecting the right people to the right jobs in a way that is fair and precise. Last year we delivered over 6 million online candidate assessments, and helped Amazon deliver billions of packages around the world by making it possible to hire hundreds of thousands of workers in the right quantity, at the right location and at exactly the right time. You’ll work on state-of-the-art research, advanced software tools, new AI systems, and machine learning algorithms, leveraging Amazon's in-house tech stack to bring innovative solutions to life. Join ITA in using technologies to transform the hiring landscape and make a meaningful difference in people's lives. Together, we can solve the world's toughest hiring problems. Key job responsibilities As an Applied Scientist, you will own the design and development of end-to-end systems. You’ll have the opportunity to write technical white papers, create technical roadmaps and drive production level projects that will support Amazon Science. You will have the opportunity to design new algorithms, models, or other technical solutions whilst experiencing Amazon’s customer focused culture. The ideal scientist must have the ability to work with diverse groups of people and cross-functional teams to solve complex business problems. About the team The Automated Performance Evaluation (APE) team is a hybrid team of Applied Scientists and Software Development Engineers who develop, deploy and own end-to-end machine learning services for use in the HR and Recruiting functions at Amazon.
US, CA, Sunnyvale
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 As an Applied Scientist at Prime Video, you will have end-to-end ownership of the product, related research and experimentation, applying advanced machine learning techniques in computer vision (CV), Generative AI, multimedia understanding and so on. You’ll work on diverse projects that enhance Prime Video’s content localization, image/video understanding, and content personalization, driving impactful innovations for our global audience. Other responsibilities include: - Research and develop generative models for controllable synthesis across images, video, vector graphics, and multimedia - Innovate in advanced diffusion and flow-based methods (e.g., inverse flow matching, parameter efficient training, guided sampling, test-time adaptation) to improve efficiency, controllability, and scalability. - Advance visual grounding, depth and 3D estimation, segmentation, and matting for integration into pre-visualization, compositing, VFX, and post-production pipelines. - Design multimodal GenAI workflows including visual-language model tooling, structured prompt orchestration, agentic pipelines. A day in the life Prime Video is pioneering the use of Generative AI to empower the next generation of creatives. Our mission is to make world-class media creation accessible, scalable, and efficient. We are seeking an Applied Scientist to advance the state of the art in Generative AI and to deliver these innovations as production-ready systems at Amazon scale. Your work will give creators unprecedented freedom and control while driving new efficiencies across Prime Video’s global content and marketing pipelines. This is a newly formed team within Prime Video Science!
US, NY, New York
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Applied Scientist to work on pre-training methodologies for Generative Artificial Intelligence (GenAI) models. You will interact closely with our customers and with the academic and research communities. Key job responsibilities Join us to work as an integral part of a team that has experience with GenAI models in this space. We work on these areas: - Scaling laws - Hardware-informed efficient model architecture, low-precision training - Optimization methods, learning objectives, curriculum design - Deep learning theories on efficient hyperparameter search and self-supervised learning - Learning objectives and reinforcement learning methods - Distributed training methods and solutions - AI-assisted research About the team The AGI team has a mission to push the envelope in GenAI with Large Language Models (LLMs) and multimodal systems, in order to provide the best-possible experience for our customers.
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Applied Scientist 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 an Applied Scientist with the AGI team, you will work with talented peers to support the development of algorithms and modeling techniques, to advance the state of the art with LLMs. 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.
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 early career research 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 Key Job Responsibilities: • Understand and contribute to model compression techniques (quantization, pruning, distillation, etc.) while developing theoretical understanding of Information Theory and Deep Learning fundamentals • Work with senior researchers to optimize Gen AI models for edge platforms using Amazon's Neural Edge Engine • Study and apply first principles of Information Theory, Scientific Computing, and Non-Equilibrium Thermodynamics to model optimization problems • Assist in research projects involving custom Gen AI model development, aiming to improve SOTA under mentorship • Co-author research papers for top-tier conferences (NeurIPS, ICLR, MLSys) and present at internal research meetings • Collaborate with compiler engineers, Applied Scientists, and Hardware Architects while learning about production ML systems • Participate in reading groups and research discussions to build expertise in efficient AI and edge computing