Learning to Reason in LLMs by Expectation Maximization
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915766026108928 |
|---|---|
| author | Lee, Junghyun Kveton, Branislav Rao, Anup Mukherjee, Subhojyoti Rossi, Ryan A. Choudhary, Sunav Siu, Alexa |
| author_facet | Lee, Junghyun Kveton, Branislav Rao, Anup Mukherjee, Subhojyoti Rossi, Ryan A. Choudhary, Sunav Siu, Alexa |
| contents | 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 a reward-based filtered expectation-maximization (FEM) 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 of rationales that justify correct answers. We instantiate and compare three sampling schemes: rejection sampling with a budget, self-taught reasoner (STaR), and prompt posterior sampling (PPS), which only keeps the rationalization stage of STaR that conditions on the correct answer in the prompt. We experiment with LLM-as-a-judge calibration and summarization from feedback tasks, where conditioning on the correct answer provides a strong guidance for generating rationales. Our experiments show the efficacy of PPS over other sampling schemes, and that the sampling scheme can have a significant impact on performance. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2512_20169 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Learning to Reason in LLMs by Expectation Maximization Lee, Junghyun Kveton, Branislav Rao, Anup Mukherjee, Subhojyoti Rossi, Ryan A. Choudhary, Sunav Siu, Alexa Machine Learning Computation and Language 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 a reward-based filtered expectation-maximization (FEM) 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 of rationales that justify correct answers. We instantiate and compare three sampling schemes: rejection sampling with a budget, self-taught reasoner (STaR), and prompt posterior sampling (PPS), which only keeps the rationalization stage of STaR that conditions on the correct answer in the prompt. We experiment with LLM-as-a-judge calibration and summarization from feedback tasks, where conditioning on the correct answer provides a strong guidance for generating rationales. Our experiments show the efficacy of PPS over other sampling schemes, and that the sampling scheme can have a significant impact on performance. |
| title | Learning to Reason in LLMs by Expectation Maximization |
| topic | Machine Learning Computation and Language |
| url | https://arxiv.org/abs/2512.20169 |