摘要：大模型推理性能的提升，真的只能靠堆数据、加算力吗？李飞飞等用仅1000个样本微调模型，并提出预算强制（Budget forcing）技术，成功让推理能力随测试计算量增加而提升。他们的s1-32B模型在多个基准测试中超越闭源模型OpenAI o1-preview
【新智元导读】大模型推理性能的提升，真的只能靠堆数据、加算力吗？李飞飞等用仅1000个样本微调模型，并提出预算强制（Budget forcing）技术，成功让推理能力随测试计算量增加而提升。他们的s1-32B模型在多个基准测试中超越闭源模型OpenAI o1-preview，成为目前最具样本效率的推理模型。
OpenAI o系列模型为何性能如此强大？
OpenAI将他们的方法描述为使用大规模强化学习（RL），暗示使用了大量的数据。
最近大火的DeepSeek-R1模型也通过使用数百万个样本和多个训练阶段使用强化学习的方式，成功地达到了o1级别的性能。
然而，至今为止没有人公开成功复现清晰的测试时扩展行为。
那么问题来了，实现测试时扩展和强推理性能的最简单方法是什么？
近日，来自斯坦福大学、华盛顿大学、Ai2等机构的研究人员发表了一篇题为「s1: Simple test-time scaling」的论文，回答了上述问题。
论文地址：https://arxiv.org/pdf/2501.19393
团队证明，仅使用1000个样本进行下一个token的预测训练，并在测试时通过一种简单的预算强制（budget forcing）技术来控制思维持续时间，就能获得一个强大的推理模型，其性能随着测试计算量的增加而提升。
预算强制（budget forcing）可以简单理解为通过强制提前结束模型的思考过程，或通过重复添加「Wait」来延长思考时间，从而影响模型的推理深度和最终答案。
这种方法可以引导模型进行自我检查，并修正推理过程中的错误，从而提高推理性能。
具体来说，他们构建了一个叫做「s1K」的数据集，由1000个精心筛选的问题组成，每个问题都配有推理轨迹（reasoning traces）和从Gemini Thinking Experimental蒸馏而来的答案。
接着团队在一个预训练模型上进行监督微调（SFT），仅使用16张H100 GPU训练26分钟。
训练完成后，使用预算强制（budget forcing）方法来控制模型在测试时的计算量：
若模型生成的推理token超过设定的上限，则强制结束推理过程，并附加思维结束（end-of-thinking）token，促使模型进入答案生成阶段。若希望模型在问题上投入更多测试时计算资源，则抑制思维结束token的生成，并在推理轨迹中追加 「Wait」，鼓励模型进行更深入的推理探索。基于这个简单的方法，并在1000个样本上进行SFT训练 + 测试时的预算强制（budget forcing）后，团队提出的s1-32B展现出了测试时扩展（test-time scaling）的能力。
此外，s1-32B也是目前最具样本效率（sample-efficient）的推理模型，在推理能力上超越了OpenAI的o1-preview等闭源模型。
如何创建s1K数据集
s1K数据集是一个包含1000个高质量推理问题的精选数据集。
团队创建过程主要分为两个阶段。
初始阶段，研究人员从16个不同的来源收集了59029个问题，并遵循三个指导原则：质量、难度和多样性。
这些来源包括现有的数学问题数据集（如 NuminaMATH、AIME、OmniMath 和 AGIEval），以及研究人员自己创建的概率问题集 (s1-prob) 和脑筋急转弯问题集 (s1-teasers)。
为了确保质量，研究人员检查了所有样本，并忽略了格式不佳的数据集。为了增加难度，他们选择需要大量推理努力的问题。为了确保多样性，他们涵盖了不同的领域和推理任务.
第二阶段，最终筛选1K样本。研究人员通过三个阶段的过滤，从59K样本中筛选出1000个样本，并继续依赖质量、难度和多样性这三个原则。
通过这些过程，研究人员创建了s1K数据集，该数据集包含50个不同领域的1000个高质量、多样化和高难度的问题，并附带推理过程。
这个数据集对于训练s1-32B模型至关重要。
测试时扩展方法
其核心思想是通过在测试时增加计算量来提高语言模型的性能。
论文将测试时扩展方法分为两类：顺序（Sequential）和并行（Parallel）。
顺序扩展是指后面的计算依赖于前面的计算，例如长的推理过程；并行扩展是指计算是独立运行的，例如多数投票。
论文主要关注顺序扩展，因为作者认为它可以更好地利用中间结果进行更深入的推理和迭代改进。
预算强制（Budget Forcing）：通过限制模型在测试时使用的最大和/或最小思考token数量来控制计算量。
论文通过实验证明了，这种简单的方法能够引导模型修正答案。
下图这个例子中，模型最初在回答一个关于「raspberry」中 「r」的数量的问题时给出了错误的答案 「2」。
然而，通过抑制结束思考的token生成，并追加「Wait」来强制模型继续推理，模型最终意识到自己快速阅读导致了错误，并最终给出了正确的答案「3」。
测试时扩展方法
如下图所示，s1-32B模型在使用预算强制技术后，其性能会随着测试时计算量的增加而提高。
具体来说，通过增加模型思考的token数量（例如，通过追加「Wait」），模型在 AIME24 基准测试上的表现得到了提升。
然而，这种提升最终会趋于平缓，过度抑制结束思考的token会导致模型进入重复循环。
结果表明，s1-32B模型是目前样本效率最高的开源推理模型。尽管只使用了1000个样本进行微调，s1-32B的性能仍明显优于其基础模型Qwen2.5-32B-Instruct。
同时，虽然DeepSeek r1-32B模型性能更强，但其使用了800倍的训练样本。
此外，s1-32B模型在AIME24上的表现几乎与Gemini 2.0 Thinking API持平，表明其蒸馏过程是有效的。
总之，实验结果证明了s1-32B模型在测试时扩展、样本效率和推理能力方面的优势，并验证了预算强制技术的有效性。
消融实验
数据消融实验：研究人员通过以下对比实验，验证了高质量、多样性、和难度这三个数据选择标准的重要性：
仅质量 (1K-random)：随机选取1000个高质量样本，性能明显低于s1K，表明难度和多样性过滤的重要性。仅多样性 (1K-diverse)：均匀选取各个领域样本，性能也远不如s1K，表明只关注多样性是不够的。仅难度 (1K-longest)：选择推理轨迹最长的1000个样本，在GPQA上有提升，但整体不如s1K，表明难度只是一个方面。最大化数据量 (59K-full)：使用所有59K样本训练，虽然性能略有提升，但训练资源消耗巨大，且提升幅度有限，说明精心挑选的少量数据比大量数据更高效。结果表明，将质量、难度和多样性相结合是实现样本高效推理训练的关键。
测试时扩展方法消融实验：研究人员通过比较不同的测试时扩展方法，验证了预算强制的优越性：
Token/步骤/类别 条件控制 (TCC/SCC/CCC)：这些方法都无法有效控制计算量或获得良好的扩展效果，表明仅在提示中告知模型计算量或步骤是不足的。拒绝采样 (RS)：使用拒绝采样会导致性能随着计算量的增加而下降，因为更短的生成往往是模型一开始就走在正确轨道上的结果。预算强制 (BF)：实验表明，预算强制在控制性、扩展性和性能方面都优于其他方法。在预算强制中，追加「Wait」能够鼓励模型进行额外的思考，从而提高性能。下图展示了在AIME24数据集上使用s1-32B模型进行拒绝采样（rejection sampling）的实验结果，结果表明，随着平均思考时间（以token数量衡量）的增加，模型的准确率反而下降，呈现出反向扩展趋势。
更长的推理过程并不一定意味着更好的性能。 这一结果也反衬了论文提出的预算强制方法（budget forcing）的优越性，因为预算强制可以更有效地控制测试时的计算量，并促使模型进行更有目的性的思考。
实验结果表明，预算强制是测试时扩展的最佳方法。
总结
尽管很多模型，例如DeepSeek-r1和k1.5，通过强化学习或使用数万个蒸馏样本来构建强大的推理模型。
但该研究表明，仅需在1000个样本上进行监督微调 (SFT) 就足以构建一个具有竞争力的推理模型，该模型可以与 OpenAI的o1-preview相媲美。
李飞飞团队研究人员认为，预训练阶段模型已经接触了大量的推理数据，因此微调阶段只需少量样本就能激活模型的推理能力，这与LIMA论文中提出的表面「对齐假说」类似。
论文提出的预算强制 (budget forcing) 是一种简单有效的顺序扩展方法，它通过控制模型思考的token数量来提高性能，并首次复现了OpenAI的测试时扩展曲线。
尽管预算强制有其局限性，例如最终会趋于平缓和受上下文窗口的限制，但它证明了测试时扩展的潜力，并为未来的研究提供了明确的指标：可控性、扩展性和性能。
为了克服顺序扩展的局限性，论文还探讨了并行扩展方法，如多数投票和基于REBASE的树搜索（下图）。实验结果表明，将顺序扩展与并行扩展方法相结合，还可以进一步扩展测试时的计算量。
未来方向
论文强调，未来的研究可以探索如何进一步改进预算强制，例如轮换使用不同的字符串或结合频率惩罚。
一个有前景的方向是将预算强制应用于通过强化学习训练的推理模型，并研究新的测试时扩展方法。
此外，可以研究如何进一步扩展测试时的计算量，以克服现有语言模型上下文窗口的限制。
来源：新智元
