Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models

最近的研究表明，有效利用推理时的计算资源对于从大型语言模型（LLMs）中获得更好的性能至关重要。在这项工作中，我们提出了一种新的推理感知微调范式，其中模型以直接优化推理时策略性能的方式进行微调。我们使用简单而有效的最佳N选一（Best-of-N, BoN）推理策略来研究这一范式，在该策略中，验证器从一组由LLM生成的响应中选择最佳答案。我们设计了首个用于BoN感知微调的模仿学习和强化学习（RL）方法，克服了BoN内部具有挑战性的非可微分argmax操作符问题。我们通过实证展示了我们的BoN感知模型隐式地学习了一种元策略，这种策略将最佳响应与可能更适合测试时输入的更多样化的响应交织在一起——这一过程让人联想到RL中的探索与利用权衡。我们的实验展示了BoN感知微调在改进性能和推理时计算效率方面的有效性。特别是，我们展示了我们的方法将Gemma 2B在Hendrycks MATH上的Bo32性能从26.8%提高到30.8%，将pass@32从60.0%提高到67.0%，以及在HumanEval上的pass@16从61.6%提高到67.1%。

该论文旨在解决如何通过优化推理时间策略来提高大型语言模型（LLM）的性能。具体来说，它探讨了在推理阶段使用Best-of-N (BoN) 策略时，如何通过细调模型以更好地适应这一策略，从而提升模型性能和推理效率。这是一个相对较新的问题，特别是在直接针对推理时间进行优化方面。

关键思路在于提出了一种新的推理感知微调范式，该范式直接优化了推理时间策略的性能。研究者引入了模仿学习和强化学习方法来进行BoN-aware的微调，解决了BoN中非可微argmax操作的挑战。这种新方法使模型能够学习一种元策略，在推理时平衡最佳响应与多样化响应的选择，类似于强化学习中的探索-利用权衡。

亮点包括：1. 首次提出了针对BoN策略的模仿学习和强化学习微调方法；2. 实验证明，BoN-aware微调显著提高了Gemma 2B在Hendrycks MATH和HumanEval数据集上的性能；3. 模型学会了在推理过程中动态选择最佳和多样化响应的组合；4. 论文提供了详细的实验设计和结果分析，表明了该方法的有效性。关于代码开源情况未提及，但值得进一步关注。

近期相关研究包括：1. 探索不同的推理优化技术，如《Efficient Fine-Tuning of Language Models with Sparse Updates》；2. 强化学习应用于文本生成，《Reinforcement Learning for Text Generation: A Comprehensive Survey》；3. 大规模预训练模型的微调策略，《Adaptive Fine-Tuning of Pre-trained Language Models》。这些研究共同推动了大模型在特定任务上的性能提升。