Weight Ensembling Improves Reasoning in Language Models

本文发现监督微调导致推理模型多样性坍塌损害 Pass@K，并提出通过插值早期与后期 SFT 检查点（WiSE-FT）的方法，有效提升模型多样性，同时提高 Pass@1 和 Pass@K，进而改善测试时缩放和强化学习效果。

Supervised Learning, Large Language Model, Reasoning, Fine-tuning, Efficiency

Xingyu Dang, Christina Baek, Kaiyue Wen, Zico Kolter, Aditi Raghunathan

Tsinghua University, Carnegie Mellon University, Stanford University

Generated by gemini-2.5-flash-preview-04-17

Background Problem

大型语言模型（LLMs）在复杂推理任务中表现出色，但这通常依赖于测试时的多样本策略（如采样多个推理轨迹并投票或使用奖励模型选择）。衡量这种能力的关键指标是 Pass@K，即 K 次独立采样中至少有一次正确的概率。然而，标准的监督微调（SFT）训练流程虽然能提高单次尝试的准确率（Pass@1），却常常导致“多样性坍塌”（diversity collapse），即模型生成的推理轨迹变得非常相似。这使得即使采样多次，也可能都犯同样的错误，从而损害了 Pass@K 性能。这种现象在数学推理等任务中尤为明显，形成了 Pass@1 提高与 Pass@K 下降之间的矛盾。

Method

• 核心问题: 解决 SFT 导致的推理模型多样性坍塌问题，从而同时提升 Pass@1 和 Pass@K。
• 核心思想: 利用 WiSE-FT（Weight-Space Ensembling for Fine-Tuning）技术，通过插值早期 SFT 检查点和后期 SFT 检查点的权重，结合两者的优点。
• 具体方法:
  • WiSE-FT 变体: 不同于原始 WiSE-FT 插值预训练模型和微调模型，本文插值的是早期 SFT 检查点（例如训练 1 epoch 后的模型 $w^0$）和后期 SFT 检查点（例如训练到时间步 t 的模型 $w^t$）。
  • 插值公式: 使用线性插值 $w_{\text{WiSE}(t)} = \delta \cdot w^0 + (1 - \delta) \cdot w^t$。论文发现 $\delta=1/2$ 在不同数据集上表现良好，并固定使用此值。
  • 工作原理: 早期 SFT 检查点通常具有较高的生成多样性（尽管 Pass@1 可能较低），而后期 SFT 检查点具有较高的 Pass@1（但多样性较低）。通过权重插值，WiSE-FT 模型旨在继承早期模型的探索能力和后期模型的准确性，从而在不牺牲 Pass@1 的前提下恢复甚至提高多样性，进而提升 Pass@K。
  • 理论分析: 论文从理论上分析了 SFT 和标准 RL（REINFORCE, GRPO）在离散 bandit 设置下可能导致多样性坍塌的原因，并形式化了 Pass@K 与 Pass@1 偏差和方差之间的关系，解释了多样性坍塌如何通过增加方差损害 Pass@K。

Experiment

• 模型与数据集: 主要实验在 Gemma-2-2B 和 Qwen-2.5-0.5B 模型上进行，使用 MetaMath40k 数据集（包含 GSM8k 和 MATH 的增强版本）进行 SFT 和后续的强化学习（GRPO）。评估数据集包括 GSM8K、MATH500 和 AIME24。
• 核心发现:
  • SFT 的影响: 实验证实，随着 SFT 步数增加，Pass@1 持续提高，但 Pass@K（特别是对于较大的 K）在训练早期达到峰值后迅速下降，表现出多样性坍塌。
  • WiSE-FT 的效果: 将最终 SFT 检查点与早期 SFT 检查点进行权重插值（WiSE-FT，固定插值系数 $\delta=1/2$）后，模型同时实现了比最终 SFT 模型更高的 Pass@1 和比早期 SFT 模型更高的 Pass@K。Pass@1 和 Pass@K 都能随着 SFT 步数（插值后的后期检查点）单调增长。
  • 测试时缩放: WiSE-FT 模型在使用多数投票或 ORM 进行 Best@K 评估时，相比过拟合的最终 SFT 模型表现出显著优势，性能提升 5-7%。
  • 强化学习: 从 WiSE-FT 检查点开始进行后续的 RL 训练，相比从早期或最终 SFT 检查点开始，能以更少的数据实现更高的 Pass@1 性能，表明 WiSE-FT 提供了更好的 RL 初始化点。
  • 通用推理: 在通用推理数据集 OpenThoughts-114k 上训练 Qwen2.5-7B-Instruct 并在 OOD 数据集 AIME24 上评估，WiSE-FT 同样能提升 Pass@1 和 Pass@K。
  • 多样性度量: 通过答案多样性、操作多样性和语义多样性指标，实验量化了 SFT 过程中模型生成多样性的显著下降。
  • 偏差-方差分析: 经验性地测量了 1-Pass@1 在测试集上的分布，发现 SFT 降低了偏差（平均错误率），但增加了方差（错误率分布变得双峰化）。温度缩放降低了方差但增加了偏差。WiSE-FT 独特地同时降低了偏差和方差。
• 实验设置合理性: 实验设置全面，涵盖了 SFT、RL、测试时缩放和 OOD 泛化等多个方面，使用了不同的模型和数据集，并通过多种指标（Pass@K, Best@K, 多样性指标, 偏差-方差）验证了方法的有效性。理论分析与实验结果相互印证。

Further Thoughts

本文的核心发现——监督微调（SFT）虽然提升了单次尝试的准确率（Pass@1），却牺牲了生成多样性，导致多样本推理（Pass@K）性能下降，这揭示了推理模型训练中的一个关键矛盾。WiSE-FT通过插值早期和后期SFT检查点来解决这个问题，其有效性令人惊讶，因为它没有引入复杂的训练目标或推理策略，仅仅是模型权重的简单组合。这暗示了SFT训练过程中，模型可能在不同的训练阶段学习到了互补的能力：早期阶段保留了探索性和多样性，而后期阶段则专注于提高特定模式下的准确性。将这种思想推广，或许可以在其他需要多样性输出的任务（如创意写作、代码生成）中探索类似的权重组合或模型融合策略。此外，论文对Pass@K的偏差-方差分解提供了一个清晰的理论视角，解释了多样性坍塌如何增加方差并损害Pass@K。理论分析指出标准RL算法在没有KL正则化时也会导致多样性坍塌，并且KL正则化并不能增加多样性，这为理解RL在推理任务中的局限性提供了深刻见解。最后，论文强调了当前解码策略与理论最优Pass@K之间的巨大差距，这指明了未来研究的重要方向：如何在训练或推理阶段更有效地引导模型生成多样且高质量的推理路径，而不仅仅依赖于后处理的解码技巧。