Learning to Reason in LLMs by Expectation Maximization
作者: Junghyun Lee, Branislav Kveton, Anup Rao, Subhojyoti Mukherjee, Ryan A. Rossi, Sunav Choudhary, Alexa Siu
分类: cs.LG, cs.CL, stat.ML
发布日期: 2025-12-23 (更新: 2026-01-31)
备注: 27 pages, 15 figures, 5 tables (ver2: major revision, including new experiments, reorganization, etc)
💡 一句话要点
提出基于期望最大化的LLM推理学习框架，优化生成合理化解释。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 推理学习 期望最大化 隐变量模型 奖励学习
📋 核心要点
- 现有LLM推理方法缺乏有效学习框架，难以生成既正确又合理的推理过程。
- 提出基于期望最大化的学习框架，将推理视为隐变量模型，通过优化采样分布提升推理能力。
- 实验表明，提示后验采样（PPS）优于其他采样方法，显著提升了LLM在推理任务中的性能。
📝 摘要（中文）
大型语言模型（LLMs）通过先生成推理过程，然后给出答案来解决推理问题。本文将推理形式化为一个隐变量模型，并推导出一个基于奖励的过滤期望最大化（FEM）目标，用于学习推理。这种观点将EM算法与现代基于奖励的优化联系起来，并表明主要挑战在于设计一个能够证明正确答案合理的推理采样分布。本文实例化并比较了三种采样方案：带预算的拒绝采样、自学推理器（STaR）和提示后验采样（PPS），后者仅保留STaR中以提示中的正确答案为条件的合理化阶段。本文通过LLM作为评判者的校准和来自反馈任务的摘要进行了实验，其中以正确答案为条件为生成理由提供了强大的指导。实验表明，PPS优于其他采样方案，并且采样方案可能对性能产生重大影响。
🔬 方法详解
问题定义：论文旨在解决大型语言模型（LLMs）在推理任务中生成合理化解释（rationales）的问题。现有的方法要么难以生成与正确答案一致的解释，要么效率低下。核心痛点在于如何有效地学习生成既能支持正确答案，又具有内在一致性的推理过程。
核心思路：论文的核心思路是将LLM的推理过程建模为一个隐变量模型，其中答案是可观测变量，而推理过程是隐变量。通过期望最大化（EM）算法，迭代地优化推理过程的生成分布，使得生成的推理过程能够更好地支持正确的答案。关键在于设计一个有效的采样策略，从推理过程的后验分布中采样。
技术框架：整体框架包含以下几个主要阶段： 1. 前向传播（E-step）：利用LLM生成候选的推理过程。 2. 奖励计算：根据生成的推理过程是否支持正确的答案，计算奖励信号。 3. 后验采样：基于奖励信号，对推理过程进行采样，得到高质量的推理过程样本。 4. 模型更新（M-step）：利用采样的推理过程样本，更新LLM的参数，提高生成高质量推理过程的能力。论文比较了三种采样方案：拒绝采样、自学推理器（STaR）和提示后验采样（PPS）。
关键创新：最重要的技术创新点在于提出了提示后验采样（PPS）方法。PPS仅保留STaR中以提示中的正确答案为条件的合理化阶段，直接利用正确答案的信息来指导推理过程的生成，从而更有效地生成高质量的推理过程。与现有方法的本质区别在于，PPS更加注重利用正确答案的监督信息，从而避免了生成大量无效的推理过程。
关键设计： 1. 奖励函数：奖励函数的设计至关重要，用于衡量生成的推理过程是否支持正确的答案。论文中使用了LLM作为评判者来评估推理过程的质量。 2. 采样策略：不同的采样策略会影响学习的效率和最终的性能。PPS通过直接以正确答案为条件，提高了采样效率。 3. 损失函数：论文推导了基于奖励的过滤期望最大化（FEM）目标，用于优化LLM的参数。
🖼️ 关键图片
📊 实验亮点
实验结果表明，提示后验采样（PPS）方法在LLM推理任务中表现优异，显著优于其他采样方法，包括拒绝采样和自学推理器（STaR）。具体而言，PPS在LLM作为评判者的校准和来自反馈任务的摘要等任务中，展现出更强的生成合理化解释的能力，证明了采样方案对性能的显著影响。
🎯 应用场景
该研究成果可应用于各种需要LLM进行推理的任务，例如问答系统、代码生成、数学问题求解等。通过提升LLM的推理能力，可以提高这些应用系统的准确性和可靠性。此外，该研究提出的学习框架也可以推广到其他需要生成合理化解释的场景，例如医疗诊断、金融风险评估等。
📄 摘要（原文）
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.