Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning to Reason in LLMs by Expectation Maximization

Junghyun Lee, Branislav Kveton, Sunav Choudhary, Subhojyoti Mukherjee, Anup Rao, Ryan A. Rossi, Alexa Siu

TL;DR

This work reframes reasoning in LLMs as a latent-variable problem with a rationale $z$ bridging $x$ to the correct answer $y^\star$ and derives an EM objective for learning to reason. It shows that exact EM steps are intractable in LLMs and introduces a practical filtered EM (FEM) update, where a reward $r(\hat{y}, y^\star)$ gates the gradient update and sampling of $(\hat{z}, \hat{y})$ from a proposal distribution $q(\cdot \mid x, y^\star; \theta)$. The authors compare three sampling schemes—RS-$M$, STaR, and PPS—on ARC, MMLU, and OpenBookQA using Llama and Qwen models, showing that PPS consistently yields higher accuracy and more concise, evidence-backed rationales. Their results highlight the critical role of the rationale sampling distribution in learning to reason, offering a simple, effective approach (PPS) that outperforms prior schemes and connects EM principles with reward-based fine-tuning.

Abstract

Large language models (LLMs) solve reasoning problems by first generating a rationale and then answering. We formalize reasoning as a latent variable model and derive an expectation-maximization (EM) objective for learning to reason. This view connects EM and modern reward-based optimization, and shows that the main challenge lies in designing a sampling distribution that generates rationales that justify correct answers. We instantiate and compare several sampling schemes: rejection sampling with a budget, self-taught reasoner (STaR), and prompt posterior sampling (PPS), which only keeps the rationalization stage of STaR. Our experiments on the ARC, MMLU, and OpenBookQA datasets with the Llama and Qwen models show that the sampling scheme can significantly affect the accuracy of learned reasoning models. Despite its simplicity, we observe that PPS outperforms the other sampling schemes.

Learning to Reason in LLMs by Expectation Maximization

TL;DR

This work reframes reasoning in LLMs as a latent-variable problem with a rationale bridging to the correct answer and derives an EM objective for learning to reason. It shows that exact EM steps are intractable in LLMs and introduces a practical filtered EM (FEM) update, where a reward gates the gradient update and sampling of from a proposal distribution . The authors compare three sampling schemes—RS-, STaR, and PPS—on ARC, MMLU, and OpenBookQA using Llama and Qwen models, showing that PPS consistently yields higher accuracy and more concise, evidence-backed rationales. Their results highlight the critical role of the rationale sampling distribution in learning to reason, offering a simple, effective approach (PPS) that outperforms prior schemes and connects EM principles with reward-based fine-tuning.

Abstract

Large language models (LLMs) solve reasoning problems by first generating a rationale and then answering. We formalize reasoning as a latent variable model and derive an expectation-maximization (EM) objective for learning to reason. This view connects EM and modern reward-based optimization, and shows that the main challenge lies in designing a sampling distribution that generates rationales that justify correct answers. We instantiate and compare several sampling schemes: rejection sampling with a budget, self-taught reasoner (STaR), and prompt posterior sampling (PPS), which only keeps the rationalization stage of STaR. Our experiments on the ARC, MMLU, and OpenBookQA datasets with the Llama and Qwen models show that the sampling scheme can significantly affect the accuracy of learned reasoning models. Despite its simplicity, we observe that PPS outperforms the other sampling schemes.
Paper Structure (8 sections, 1 theorem, 6 equations, 3 figures, 1 table)

This paper contains 8 sections, 1 theorem, 6 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. EM and Learning to Reason
  3. Different Rationale Sampling Schemes
  4. Experiments
  5. Setting
  6. Results
  7. Additional Related Work
  8. Conclusions

Key Result

Lemma 1

Let $(x, y^\star)$ be a question-answer pair and $r(\hat{y}, y^\star) = \mathds{1}[\hat{y} = y^\star]$. Then for any parameter vector $\theta$ and iteration $k \geq 1$, The bound is tight when $q(\cdot \mid x, y^\star; \theta^{(k - 1)}) \equiv \pi(\cdot \mid x; \theta)$.

Figures (3)

  • Figure 1: The rationale proposal prompt in all experiments. We set {question} to the asked question. We set {response} to four candidate answers, indexed by integers from $1$ to $4$. In prompt posterior sampling, we add the blue text and set {conditioning} to the index of the correct answer.
  • Figure 2: Test accuracy, data usage, and reasoning length in characters for four sampling schemes applied to Llama3.2-3B-Instruct.
  • Figure 3: Test accuracy, data usage, and reasoning length in characters for four sampling schemes applied to Qwen2.5-3B-Instruct.

Theorems & Definitions (1)

  • Lemma 1: Adapted from Lemma 1 of mukherjee2025multiturn