Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning

大型语言模型（LLMs）的最新进展突显了在测试时扩展计算能力以在复杂任务（如数学推理和代码生成）中取得优异表现的强大作用。这引发了一个关键问题：应该如何修改模型训练，以便在随后的测试时间计算策略和预算下优化性能？为了探讨这一点，我们专注于pass@N，这是一种简单的测试时策略，在$N$个独立样本中搜索正确答案。令人惊讶的是，我们发现使用交叉熵（CE）损失进行训练可能会与pass@N不一致，即随着训练时间的延长，pass@N准确率反而下降。我们解释了这种不一致的原因在于由CE引起的模型过度自信，并通过实验验证了过度自信是通过pass@N扩展测试时计算的障碍。此外，我们提出了一种有原则的、经过修改的训练损失函数，通过限制模型的置信度来更好地与pass@N对齐，从而恢复pass@N的测试性能。我们的算法在多个场景下展示了改进的数学推理能力，包括在MATH和MiniF2F基准测试中的表现：（1）回答数学问题；以及（2）通过搜索不同形状的证明树来证明定理。总体而言，我们的工作强调了共同设计LLM开发两个传统上独立的阶段的重要性：训练时协议和测试时的搜索及推理策略。

论文试图解决的问题是，在大规模语言模型（LLM）中，如何优化训练过程以适应特定的测试时计算策略和预算。具体来说，研究发现传统的交叉熵损失函数（CE loss）在训练过程中可能会导致与pass@N测试策略的不一致，即随着训练时间的增长，pass@N准确率反而下降。这并不是一个全新的问题，但在LLM背景下对这一现象的研究和解决方案是新颖的。

关键思路在于识别并解释了由交叉熵损失函数引起的模型过自信现象，并提出了一种新的训练损失函数来限制模型的自信度，从而改善pass@N测试性能。相比现有研究，该论文提出了通过调整训练损失函数来直接针对测试时策略进行优化的方法，这是较为新颖的尝试。

论文值得关注的地方包括：1. 通过实验验证了过自信是影响pass@N性能的主要因素；2. 提出了一个修改后的训练损失函数，该函数在数学推理和代码生成任务上表现出色；3. 在MATH和MiniF2F基准测试中展示了改进效果；4. 强调了训练时间和测试时搜索及推理策略共同设计的重要性。此外，作者还开源了相关代码，便于后续研究。

最近在这个领域内还有其他相关研究，例如《Improving Neural Net Robustness via Regularization of Confidence》探讨了如何通过正则化方法减少模型的过自信；《On Calibration of Modern Neural Networks》讨论了现代神经网络校准问题；以及《Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》研究了统一文本到文本转换器在不同任务间的迁移学习极限。