Teaching language models to reason consistently

At this year’s ACL, Amazon researchers won an outstanding-paper award for showing that knowledge distillation using contrastive decoding in the teacher model and counterfactual reasoning in the student model improves the consistency of “chain of thought” reasoning.

Teaching large language models (LLMs) to reason is an active topic of research in natural-language processing, and a popular approach to that problem is the so-called chain-of-thought paradigm, in which a model is prompted not just to answer questions but to provide rationales for its answers.

In-context learning.png
The structure of the type of prompt used to induce chain-of-thought reasoning in a large language model.

However, given LLMs’ tendency to hallucinate (that is, make spurious factual assertions), the generated rationales may be inconsistent with the predicted answers, making them untrustworthy.

In a paper we presented at this year’s meeting of the Association for Computational Linguistics (ACL), we show how to improve the consistency of chain-of-thought reasoning through knowledge distillation: given pairs of questions and answers from a training set, an LLM — the “teacher” — generates rationales for a smaller “student” model, which learns to both answer questions and provide rationales for its answers. Our paper received one of the conference’s outstanding-paper awards, reserved for 39 of the 1,074 papers accepted to the main conference.

Sample rationale.png
Example of a student model outputting a rationale together with the answer to a question.

With knowledge distillation (KD), we still have to contend with the possibility that the rationales generated by the teacher are spurious or vacuous. On the student side, the risk is that while the model may learn to produce rationales, and it may learn to deliver answers, it won’t learn the crucial logical relationships between the two; it might, for instance, learn inferential short cuts between questions and answers that bypass the whole reasoning process.

False rationales.png
In a study involving a leading LLM, we found that 42% of generated rationales were vacuous (top), and 37% were irrelevant (bottom).

To curb hallucination, on the teacher side, we use contrastive decoding, which ensures that the rationales generated for true assertions differ as much as possible from the rationales generated for false assertions.

To train the student model, we use counterfactual reasoning, in which the model is trained on both true and false rationales and must learn to provide the answer that corresponds to the rationale, even if it’s wrong. To ensure that this doesn’t compromise model performance, during training, we label true rationales “factual” and false rationales “counterfactual”.

Counterfactual training.png
Counterfactual training eliminates reasoning short cuts, in which the student model uses incidental features of the input question to leap to an answer, without performing the intervening inferential steps.

To evaluate our model, we compared it to a chain-of-thought model built using ordinary knowledge distillation, on datasets for four different reasoning tasks. We asked human reviewers to evaluate the rationales generated by the teacher models. To evaluate the student models, we used the leakage-adjusted simulatability (LAS) metric, which measures the ability of a simulator (an external model) to predict the student’s outputs from the generated rationales. Across the board, our models outperformed the baselines, while preserving accuracy on the reasoning tasks.

Contrastive decoding

As our teacher model, we use a trained LLM whose parameters are frozen. To generate training examples for the student model, we use in-context learning, in which we provide the teacher with a handful of examples of questions, answers, and human-annotated rationales, then supply a final question-answer pair. The model generates the rationale for the final pair.

Related content
Methods for controlling the outputs of large generative models and integrating symbolic reasoning with machine learning are among the conference’s hot topics.

During training, LLMs learn the probabilities of sequences of words. At generation time, they either select the single most probable word to continue a sequence or sample from the top-ranked words. This is the standard decoding step, which doesn’t guarantee that the generated rationales justify the model’s answers.

We can control the decoding process without making any adjustments to the LLM parameters. With contrastive decoding, we perform the same in-context rationale generation twice, once with the true answer in the final question-answer pair and once with a perturbed answer.

Then, when we’re decoding the true question-answer pair, we select words that are not only probable given the true pair but relatively improbable given the false pair. In other words, we force the rationale for the true pair to diverge from the rationale for the false pair. In this way, we ensure that the output skews toward rationales particularized to the answers in the question-answer pairs.

In our experiments, we considered two types of perturbation to the true answers: null answers, where no answer at all was supplied, and false answers. We found that contrastive decoding using false answers consistently yielded better rationales than contrastive decoding using null answers.

Counterfactual reasoning

Past research has shown that question-answering models will often exploit short cuts in their training data to improve performance. For instance, answering “who?” questions with the first proper name encountered in a source document will yield the right answer with surprising frequency.

Similarly, a chain-of-thought model might learn to use shortcuts in answering questions and generate rationales as a parallel task, without learning the crucial connection between the two. The goal of training our model on a counterfactual-reasoning objective is to break that short cut.

Related content
Amazon’s Dan Roth on a hot new research topic — that he’s been studying for more than 25 years.

To generate counterfactual training data, we randomly vary the answers in question-answer pairs and generate the corresponding rationales, just as we did for contrastive decoding. Then we train the student model using the questions and rationales as input, and it must generate the corresponding answers.

This means that the student model may very well see the same question multiple times during training, but with different answers (and rationales). The “factual” and “counterfactual” tags prevent it from getting confused about its task.

In our experiments, we compared our approach to one that also uses in-context learning but uses greedy decoding to produce rationales — that is, a decoding method that always selects the highest-probability word. We also used two other baselines: an LLM that directly generates rationales from in-context learning and a model trained on human-annotated rationales.

Our study with human evaluators showed that in-context learning with contrastive decoding generated more persuasive rationales than in-context learning with greedy decoding:

Teacher Model

Grammaticality

New Info

Supports Answer

Greedy

0.99

0.65

0.48

Contrast.-Empty

0.97

0.77

0.58

Contrast.-Wrong

0.97

0.82

0.63

Table: Human evaluation of data generated with greedy decoding, contrastive decoding using empty answers, and contrastive decoding using incorrect answers.

In the experiments using the LAS metric, knowledge distillation using contrastive decoding alone consistently outperformed all three baselines, and knowledge distillation with counterfactual reasoning and contrastive decoding consistently outperformed knowledge distillation with contrastive decoding alone. The model trained on the human-annotated dataset yielded the most-accurate results on downstream tasks, but its rationales fared badly. On average, our model was slightly more accurate than the one trained using greedy decoding.

LAS results.png
Experimental results, measured according to leakage-adjusted simulatability (LAS) and question-answering accuracy.

Research areas

Related content

US, CA, Sunnyvale
Amazon is on a mission to redefine the future of automation — and we're looking for exceptional talent to help lead the way. We are building the next generation of advanced robotic systems that seamlessly blend cutting-edge AI, sophisticated control systems, and novel mechanical design to create adaptable, intelligent automation solutions capable of operating safely alongside humans in dynamic, real-world environments. At Amazon, we leverage the power of machine learning, artificial intelligence, and advanced robotics to solve some of the most complex operational challenges at a scale unlike anywhere else in the world. Our fleet of robots spans hundreds of facilities globally, working in sophisticated coordination to deliver on our promise of customer excellence — and we're just getting started. As a Sr. Scientist in Robot Navigation, you will be at the forefront of this transformation — architecting and delivering navigation systems that are intelligent, safe, and scalable. You will bring deep expertise in learning-based planning and control, a strong understanding of foundation models and their application to embodied agents, and as well as have in-depth understanding of control-theoretic approaches such as model predictive control (MPC)-based trajectory planning. You will develop navigation solutions that seamlessly blend data-driven intelligence with principled control-theoretic guarantees. Our vision is bold: to build navigation systems that allow robots to move fluidly and safely through dynamic environments — understanding context, anticipating change, and adapting in real time. You will lead research that bridges the gap between cutting-edge academic advances and production grade deployment, collaborating with world-class teams pushing the boundaries of robotic autonomy, manipulation, and human-robot interaction. Join us in building the next generation of intelligent navigation systems that will define the future of autonomous robotics at scale. Key job responsibilities - Design, develop, and deploy perception algorithms for robotics systems, including object detection, segmentation, tracking, depth estimation, and scene understanding - Lead research initiatives in computer vision, sensor fusion and 3D perception - Collaborate with cross-functional teams including robotics engineers, software engineers, and product managers to define and deliver perception capabilities - Drive end-to-end ownership of ML models — from data collection and labeling strategy to training, evaluation, and deployment - Mentor junior scientists and engineers; contribute to a culture of technical excellence - Define and track key metrics to measure perception system performance in real-world environments - Publish research findings in top-tier venues (CVPR, ICCV, ECCV, ICRA, NeurIPS, etc.) and contribute to patents A day in the life - Train ML models for deployment in simulation and real-world robots, identify and document their limitations post-deployment - Drive technical discussions within your team and with key stakeholders to develop innovative solutions to address identified limitations - Actively contribute to brainstorming sessions on adjacent topics, bringing fresh perspectives that help peers grow and succeed — and in doing so, build lasting trust across the team - Mentor team members while maintaining significant hands-on contribution to technical solutions About the team Our team is a group is a diverse group of scientists and engineers passionate about building intelligent machines. We value curiosity, rigor, and a bias for action. We believe in learning from failure and iterating quickly toward solutions that matter.
IN, TS, Hyderabad
Have you ever wondered how Amazon launches and maintains a consistent customer experience across hundreds of countries and languages it serves its customers? Are you passionate about data and mathematics, and hope to impact the experience of millions of customers? Are you obsessed with designing simple algorithmic solutions to very challenging problems? If so, we look forward to hearing from you! At Amazon, we strive to be Earth's most customer-centric company, where both internal and external customers can find and discover anything they want in their own language of preference. Our Translations Services (TS) team plays a pivotal role in expanding the reach of our marketplace worldwide and enables thousands of developers and other stakeholders (Product Managers, Program Managers, Linguists) in developing locale specific solutions. Amazon Translations Services (TS) is seeking an Applied Scientist to be based in our Hyderabad office. As a key member of the Science and Engineering team of TS, this person will be responsible for designing algorithmic solutions based on data and mathematics for translating billions of words annually across 130+ and expanding set of locales. The successful applicant will ensure that there is minimal human touch involved in any language translation and accurate translated text is available to our worldwide customers in a streamlined and optimized manner. With access to vast amounts of data, technology, and a diverse community of talented individuals, you will have the opportunity to make a meaningful impact on the way customers and stakeholders engage with Amazon and our platform worldwide. Together, we will drive innovation, solve complex problems, and shape the future of e-commerce. Key job responsibilities * Apply your expertise in LLM models to design, develop, and implement scalable machine learning solutions that address complex language translation-related challenges in the eCommerce space. * Collaborate with cross-functional teams, including software engineers, data scientists, and product managers, to define project requirements, establish success metrics, and deliver high-quality solutions. * Conduct thorough data analysis to gain insights, identify patterns, and drive actionable recommendations that enhance seller performance and customer experiences across various international marketplaces. * Continuously explore and evaluate state-of-the-art modeling techniques and methodologies to improve the accuracy and efficiency of language translation-related systems. * Communicate complex technical concepts effectively to both technical and non-technical stakeholders, providing clear explanations and guidance on proposed solutions and their potential impact. About the team We are a start-up mindset team. As the long-term technical strategy is still taking shape, there is a lot of opportunity for this fresh Science team to innovate by leveraging Gen AI technoligies to build scalable solutions from scratch. Our Vision: Language will not stand in the way of anyone on earth using Amazon products and services. Our Mission: We are the enablers and guardians of translation for Amazon's customers. We do this by offering hands-off-the-wheel service to all Amazon teams, optimizing translation quality and speed at the lowest cost possible.
US, CA, San Diego
Do you want to join an innovative team of scientists and engineers who use terabytes of data and create state-of-the-art Generative AI algorithms to push the boundaries of AI creativity? We are building foundational behavioral models for Amazon Stores using Generative AI, LLMs and Large Model training techniques that fuses general world knowledge, customer shopping behavior and Amazon e-commerce domain knowledge. We are looking for scientists who are passionate about technology, innovation, and customer experience, and are ready to make a lasting impact on the industry using intelligent and transformative AI applications. Working closely with cross-functional teams, you will be an essential part of every stage of AI development, from ideation and design to rigorous testing and successful deployment, ensuring our AI projects drive innovation and provide value for our customers. 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 In this role you will leverage your background and expertise to lead developing foundational behavioral model for Amazon Stores using Generative AI, LLM and Large Model training techniques. On a day-to-day basis, you will: - Research and implement new algorithms and architectures for generative AI applications. - Optimize model performance and scalability for inference and deployment. - Collaborate with other talented applied scientists and engineers to gather and preprocess large datasets and develop an improved training infrastructure that accelerates innovation. - Experiment with SOTA methods to improve generative AI model quality. - Provide technical expertise and guidance to support the integration of generative AI solutions into various products and services.
IN, KA, Bengaluru
As a member of the CMT team, you'll play a key role in the evolution of our Competitive Monitoring systems to solve significantly complex and interesting technical challenges in machine learning, large language models in production, and recommender systems to name a few. The team's work directly impacts customer experience at a worldwide scale. Key job responsibilities Thought leader on the team and help set team directions Research multiple problem domains, suggest various approaches to try and be as hands-on as needed while providing more junior scientists with critical mentorship Collaborate with engineers to come up with the right LLD and HLD to solve key business problems Strong emphasis on communication via writing, internal and external talks, and being able to align with multiple stakeholders A day in the life As an Applied scientist II, a typical day will involve aligning with key product, engineering and business stakeholders ; advising junior scientists on the work they are doing ; reading current research papers and staying up-to-date on AI research ; diving deep as needed to improve CMT models and addressing stakeholders from the science perspective ; writing python code
IN, KA, Bengaluru
As a member of the CMT team, you'll play a key role in the evolution of our Competitive Monitoring systems to solve significantly complex and interesting technical challenges in machine learning, large language models in production, and recommender systems to name a few. The team's work directly impacts customer experience at a worldwide scale. Key job responsibilities 1. Research the problem domain and come up with various approaches to solve the problem. 2. Be willing to experiment quickly and fail fast. 3. Collaborate with engineers to come up with the right end to end solution to the business problems. 4. Ideate on future roadmap for science in CMT 5. Be willing to roll up your sleeves and learn core topics outside applied science, for example ML engineering A day in the life A typical day might involve (a) working on ideas for improving models around product similarity or price recommendations, (b) working closely with other scientists and our ML engineers to ensure that the best models are in production, (c) writing good maintainable code that can be reused and reproduced, (d) sharing your work across CMT and beyond via technical writings and presentations
US, CA, Santa Clara
We are looking for passionate, talented, and inventive Applied Scientists with a strong machine learning background to help build industry-leading Conversational AI Systems. Our mission is to provide a delightful experience to Amazon’s customers by pushing the envelope in Natural Language Understanding (NLU), Dialog Systems including Generative AI with Large Language Models (LLMs) and Applied Machine Learning (ML). You will work alongside internationally recognized experts to develop novel algorithms and modeling techniques to advance the state-of-the-art in human language technology. Your work will directly impact millions of our customers in the form of products and services that make use language technology. You will gain hands on experience with Amazon’s heterogeneous text, structured data sources, and large-scale computing resources to accelerate advances in language understanding. We are hiring in all areas of human language technology: NLU, Dialog Management, Conversational AI, LLMs and Generative AI. A day in the life The team uses generative AI and foundation models to reimagine the experience of all customers on AWS. We explore new technologies and find creative solutions. Curiosity and an explorative mindset can find a place here to impact the life of engineers around the world. If you are excited about this space and want to enlighten your peers with new capabilities, this is the team for you. We are open to hiring candidates to work out of one of the following locations: Santa Clara, CA, USA About the team AWS Utility Computing (UC) provides product innovations — from foundational services such as Amazon’s Simple Storage Service (S3) and Amazon Elastic Compute Cloud (EC2), to consistently released new product innovations that continue to set AWS’s services and features apart in the industry. As a member of the UC organization, you’ll support the development and management of Compute, Database, Storage, Internet of Things (Iot), Platform, and Productivity Apps services in AWS, including support for customers who require specialized security solutions for their cloud services. Diverse Experiences AWS values diverse experiences. Even if you do not meet all of the qualifications and skills listed in the job description, we encourage candidates to apply. If your career is just starting, hasn’t followed a traditional path, or includes alternative experiences, don’t let it stop you from applying. Why AWS? Amazon Web Services (AWS) is the world’s most comprehensive and broadly adopted cloud platform. We pioneered cloud computing and never stopped innovating — that’s why customers from the most successful startups to Global 500 companies trust our robust suite of products and services to power their businesses. Inclusive Team Culture Here at AWS, it’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be proud of our differences. Ongoing events and learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (gender diversity) conferences, inspire us to never stop embracing our uniqueness. Mentorship & Career Growth We’re continuously raising our performance bar as we strive to become Earth’s Best Employer. That’s why you’ll find endless knowledge-sharing, mentorship and other career-advancing resources here to help you develop into a better-rounded professional. Work/Life Balance We value work-life harmony. Achieving success at work should never come at the expense of sacrifices at home, which is why we strive for flexibility as part of our working culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve in the cloud. Hybrid Work We value innovation and recognize this sometimes requires uninterrupted time to focus on a build. We also value in-person collaboration and time spent face-to-face. Our team affords employees options to work in the office every day or in a flexible, hybrid work model near one of our U.S. Amazon offices.
US, WA, Seattle
The Automated Reasoning Group in the Amazon Neuron team is looking for an Applied Scientist to work on the intersection of Artificial Intelligence and program analysis to raise the code quality bar in our state-of-the-art deep learning compiler stack. This stack is designed to optimize application models across diverse domains, including Large Language and Vision, originating from leading frameworks such as PyTorch and JAX. Your role will involve working closely with our custom-built Machine Learning accelerator, Trainium, which represents the forefront of innovation for advanced ML capabilities, and is the underpinning of Generative AI. In this role as an Applied Scientist, you'll be instrumental in designing, developing, and deploying analyzers for ML compiler stages and compiler IRs. You will architect and implement business-critical tooling, publish research, and mentor a brilliant team of experienced scientists and engineers. You will need to be technically capable, credible, and curious in your own right as a trusted AWS Neuron engineer, innovating on behalf of our customers. Your responsibilities will involve tackling crucial challenges alongside a talented engineering team, contributing to leading-edge design and research in compiler technology and deep-learning systems software. Strong experience in programming languages, compilers, program analyzers, theorem provers, and program synthesis engines will be a benefit in this role. A background in machine learning and AI accelerators is preferred but not required.
US, CA, Santa Clara
We are looking for passionate, talented, and inventive Principal Applied Scientist with a strong machine learning background to help build industry-leading Conversational AI Systems. Our mission is to provide a delightful experience to Amazon’s customers by pushing the envelope in Natural Language Understanding (NLU), Dialog Systems including Generative AI with Large Language Models (LLMs) and Applied Machine Learning (ML). As part of our team, you will work alongside internationally recognized experts to develop novel algorithms and modeling techniques to advance the state-of-the-art in human language technology. Your work will directly impact millions of our customers in the form of products and services that make use language technology. You will gain hands on experience with Amazon’s heterogeneous text, structured data sources, and large-scale computing resources to accelerate advances in language understanding. We are hiring in all areas of human language technology: NLU, Dialog Management, Conversational AI, LLMs and Generative AI. A day in the life The team uses generative AI and foundation models to reimagine the experience of all customers on AWS. We explore new technologies and find creative solutions. Curiosity and an explorative mindset can find a place here to impact the life of engineers around the world. If you are excited about this space and want to enlighten your peers with new capabilities, this is the team for you.
US, CA, Palo Alto
The Demand Utilization team within Amazon Advertising is responsible for determining which ads to serve when hundreds of millions of shoppers search for products on Amazon. We sit at the intersection of customer intent understanding and advertiser value, solving one of the most complex matching problems in the industry, identifying the right ad, for the right shopper, at the right moment, across one of the world's largest product catalogs. Our systems deliver billions of ad impressions and millions of clicks daily under strict relevance and latency constraints. We are looking for a Principal Applied Scientist to set the technical vision and drive the science strategy for our ad retrieval and ranking systems. This is a high-impact leadership role where you will tackle unsolved problems at the frontier of large-scale information retrieval, natural language understanding, and multi-objective optimization, all operating in real time at Amazon scale. You will work on challenges such as: - Modeling shopper intent from sparse, ambiguous, and multi-modal signals - Designing retrieval architectures that balance relevance, advertiser and shopper experience across billions of candidate ads - Advancing personalization and cold-start strategies for new advertisers and emerging product categories This is a role for a scientist who wants to shape the future of performance advertising through rigorous research applied to real-world systems that directly impact Amazon's customers, sellers, and business. Key job responsibilities Key Responsibilities: - Own the science roadmap for ad retrieval and ranking within Demand Utilization, defining multi-year research priorities aligned with business goals - Lead the design and development of novel machine learning models and algorithms for relevance, intent understanding, and ad selection at scale - Drive end-to-end execution from problem formulation and experimentation through production deployment, measuring impact on shopper and advertiser outcomes - Mentor and elevate a team of applied scientists and research engineers, raising the technical bar and fostering a culture of scientific rigor - Collaborate cross-functionally with product, engineering, and business leaders to translate science capabilities into product strategy - Represent Amazon externally through publications at top-tier venues, patents, and participation in the broader ML/IR research community
US, WA, Seattle
As a Principal Applied Scientist at Prime Video, you will be a technical and strategic leader responsible for inventing, developing, and deploying groundbreaking AI solutions that power personalized, relevant, and delightful experiences for millions of global customers. You will help shape the vision and direction of key ML systems that support Prime Video’s mission to deliver AI-powered customer experiences. This role demands a unique blend of deep technical expertise in machine learning and recommendation systems, industry leadership, and strong collaboration skills. You will guide the development of high-impact systems end-to-end - leading innovation from foundational research through production deployment - while mentoring scientists and influencing product and engineering roadmaps. We are looking for a thought leader who brings a strong track record of delivering ML innovations at scale, along with the curiosity and drive to push boundaries. This is a rare opportunity to drive meaningful impact at one of the largest streaming services in the world. Key job responsibilities - Invent, prototype, and productionize large-scale AI solutions across Prime Video’s personalization and discovery ecosystem using deep learning, generative AI, reinforcement learning, and optimization techniques; - Provide technical leadership and influence product vision by collaborating closely with engineers, product managers, and senior stakeholders; - Design and lead high-impact A/B tests and data analyses to validate hypotheses and guide product direction; - Drive technical bar-raising across science and engineering teams through mentorship, design reviews, and collaboration; - Stay ahead of industry trends and emerging research; leverage them to evolve long-term strategy and architecture; - Publish impactful research internally and externally (e.g. top-tier conferences and journals).