公众号ID
｜
计算机视觉研究院
学习群
｜
扫码在主页获取加入方式
开源地址：
https://github.com/simplescaling/s1
Column of Computer Vision Institute
Test-time scaling是一种很有前景的语言建模新方法，它利用额外的测试时计算资源来提升性能。最近，OpenAI 的 o1 模型展现出了这种能力，但并未公开其方法，这引发了众多复现尝试。我们致力于探寻实现测试时缩放以及强大推理性能的最简方法。
Test-time scaling是一种很有前景的语言建模新方法，它利用额外的测试时计算资源来提升性能。最近，OpenAI的o1模型展现出了这种能力，但并未公开其方法，这引发了众多复现尝试。我们致力于探寻实现测试时缩放以及强大推理性能的最简方法。
首先，我们精心整理了一个名为s1K的小型数据集，该数据集包含1000道问题，并配有推理过程，这些内容依据我们通过对比分析验证的三个标准选取：难度、多样性和质量。其次，我们开发了 “预算强制” 方法，通过强行终止模型的思考过程，或者在模型试图结束生成时多次追加 “等待” 以延长其思考过程，从而控制测试时的计算量。这能促使模型对答案进行二次检查，常常可以修正错误的推理步骤。在s1K数据集上对Qwen2.5-32B-Instruct语言模型进行有监督微调，并为其配备 “预算强制” 功能后，我们的模型s1-32B在竞赛数学问题（MATH 和 AIME24）上的表现比o1-preview高出27%。此外，通过 “预算强制” 对s1-32B进行扩展，能够使其在无需测试时干预的情况下提升性能：在 AIME24 上的成绩从50%提高到57%。
我们发现，仅通过对1000个样本进行下一个词预测训练，并借助一种我们称之为 “预算强制” 的简单测试阶段技术来控制思考时长，就能够打造出一个强大的推理模型，该模型的性能会随着测试阶段计算资源的增加而提升。具体而言，我们构建了s1K数据集，它由1000个经过精心挑选的问题组成，这些问题配有从Gemini思维实验提炼而来的推理过程和答案。我们在这个小型数据集上对一个现成的预训练模型进行有监督微调（SFT），在16块H100 GPU上仅需26分钟的训练时间。
训练完成后，我们使用 “预算强制” 来控制模型在测试阶段所花费的计算量：（I）如果模型生成的思考词元数量超过了预期限制，我们会通过追加一个思考结束词元分隔符来强行终止思考过程。以这种方式结束思考会使模型转而生成答案。（II）如果我们希望模型在某个问题上花费更多测试阶段的计算资源，我们会抑制思考结束词元分隔符的生成，而是在模型当前的推理过程中追加 “等待”，以鼓励模型进行更多探索。
凭借这个简单的方法在1000个样本上进行有监督微调以及测试阶段的预算强制，我们的s1-32B模型展现出了测试阶段缩放的特性（下图）。
此外，s1-32B是样本利用效率最高的推理模型，其性能超过了诸如 OpenAI 的 o1 - preview 等闭源模型（下图）。
我们针对以下两个方面开展了广泛的消融实验：（a）我们选取的1000个（1K）推理样本；（b）我们的测试时缩放方法。
对于（a），我们发现将难度、多样性和质量指标共同纳入样本选择算法至关重要。随机选择样本、选择推理过程最长的样本，或者仅选择多样性最大的样本，都会导致性能显著下降（在AIME24测试中平均下降约30%）。在包含59000个示例的完整数据池（s1K的超集）上进行训练，相较于我们选取的1000个样本，并没有带来显著的性能提升。这凸显了精心选择数据的重要性，也与之前关于指令微调的研究结果相呼应。对于（b），我们为测试时缩放方法定义了评估标准，以便比较不同的方法。“预算强制” 方法实现了最佳的缩放效果，因为它具有完美的可控性，且呈现出明显的正斜率，从而带来了出色的性能表现。
初步收集 59000 个样本
我们依据三条指导原则，从16个不同来源初步收集了59029个问题。
-
质量：数据集应具有较高质量；我们会始终检查样本，对于格式不佳等问题的数据集予以忽略。
-
难度：数据集应具有挑战性，需要大量的推理工作。
-
多样性：数据集应来自不同领域，以涵盖不同的推理任务。
我们收集的数据集分为两类：
现有数据集的整理
我们最大的数据源是 NuminaMATH，它包含从在线网站收集的 30660 道数学问题。我们还纳入了美国数学邀请赛（AIME）的历史试题（1983-2021年）。为了增加多样性，我们添加了 OlympicArena，其中包含来自各类奥林匹克竞赛的4250道涵盖天文学、生物学、化学、计算机科学、地理学、数学和物理学的问题。OmniMath补充了4238道竞赛级别的数学问题。我们还纳入了来自AGIEval的2385道问题，该数据集的特点是包含如SAT和LSAT等标准化考试的题目，涵盖英语、法律和逻辑等领域。其他数据源详见下表。
定量推理新数据集
为了补充这些现有数据集，我们创建了两个原创数据集。
-
s1-prob：包含来自斯坦福大学统计系博士资格考试概率部分（https://statistics.stanford.edu）的182道问题，并配有涵盖复杂证明的手写解答。概率资格考试每年举行一次，需要专业水平的数学解题能力。
-
s1-teasers：包含23道具有挑战性的脑筋急转弯，这些题目常用于量化交易岗位的面试问题。每个样本包含一个问题和一个解答，均取自PuzzledQuant（https://www.puzzledquant.com/）。我们只选取了难度最高（“困难”）的示例。
对于每个问题，我们使用谷歌 Gemini Flash Thinking API生成推理过程和解答，并提取其推理轨迹和响应。这样就得到了59000个由问题、生成的推理过程和生成的解答组成的三元组。数据集的示例详如下。
我们使用 8-gram 方法对所有样本与评估问题（MATH500、GPQA Diamond、AIME24）进行去污染处理，并对数据进行去重。
最终选取1000个样本
我们本可以直接在包含59000个问题的数据池中进行训练，但我们的目标是用最少的资源找到最简单的方法。因此，我们依据质量、难度和多样性这三条数据指导原则，经过三个阶段的筛选，最终得到一个仅包含1000个样本的最小集合。
质量筛选
首先，我们剔除了所有调用API时出现错误的问题，将数据集规模缩减至54116个样本。接着，我们通过检查样本是否包含存在格式问题的字符串模式（如ASCII艺术图、不存在的图像引用或不一致的问题编号）来过滤掉低质量的示例，将数据集进一步缩减至51581个示例。从这个数据池中，我们从自认为高质量且无需进一步筛选的数据集中确定了384个样本，作为最终1000个样本的一部分。
难度筛选
对于难度筛选，我们使用两个指标：模型性能和推理过程长度。我们在每个问题上对两个模型进行评估：Qwen2.5-7B-Instruct和Qwen2.5-32B-Instruct，并由Claude 3.5 Sonnet将每次尝试的结果与参考答案进行对比，以此来评估解答的正确性。我们使用Qwen2.5分词器测量每个推理过程的词元长度，以此来衡量问题的难度，其依据的假设是，更难的问题需要更多的思考词元。根据评分结果，我们剔除了Qwen2.5-7B-Instruct或Qwen2.5-32B-Instruct能够正确解答的问题，因为这些问题可能太简单了。通过使用两个模型进行筛选，我们降低了因某个模型在简单问题上偶尔出错而使简单样本通过筛选的可能性。经过这一步筛选，样本总数降至24496个，为基于多样性的下一轮子采样做好了准备。虽然使用这两个模型进行筛选可能是针对我们的特定设置进行优化的（因为我们后续也会使用Qwen2.5-32B-Instruct作为要微调的模型），但基于模型进行筛选的思路可以推广到其他设置中。
多样性筛选
为了量化多样性，我们基于美国数学学会的《数学学科分类》（MSC）系统，使用 Claude 3.5 Sonnet 将每个问题分类到特定的领域（例如，几何、动态系统、实分析等）。该分类法主要聚焦于数学主题，但也涵盖了生物学、物理学和经济学等其他学科。
为了从24496个问题的数据池中选出最终的样本，我们首先从所有领域中随机均匀地选择一个领域。然后，根据一种倾向于选择推理过程更长问题的分布，从该领域中抽取一个问题，这一做法的动机源于前面 “难度筛选” 部分提到的观点。我们重复这个过程，直到选出总共1000个样本。
这一分为三个阶段的筛选过程产生了一个涵盖50个不同领域的数据集。
我们将测试时缩放方法分为两类：1）顺序式缩放，即后续的计算依赖于先前的计算（例如，一个较长的推理过程）；2）并行式缩放，即计算相互独立地进行（例如，多数投票法）。我们专注于顺序式缩放，因为直觉上我们认为它的缩放效果会更好，因为后续的计算可以基于中间结果进行，从而实现更深入的推理和迭代优化。我们提出了新的顺序式缩放方法以及对其进行基准测试的方式。
我们提出了一种简单的解码时干预方法，即在测试时强制规定思考词元的最大和 / 或最小数量。具体来说，我们通过简单地追加思考结束词元分隔符和 “最终答案：” 来强制实施最大词元计数，从而提前退出思考阶段，让模型给出当前的最佳答案。为了强制实施最小词元计数，我们会抑制思考结束词元分隔符的生成，并可选择在模型当前的推理过程中追加字符串 “等待”，以鼓励模型反思其当前的生成结果。图 3 展示了这种简单方法如何能让模型得出更好答案的一个示例。
我们使用以下方法对 “预算强制” 方法进行基准测试：
（I）条件长度控制方法：这类方法依赖于在提示信息中告知模型应该生成多长的内容。我们根据粒度将其分为以下几类：
（a）词元条件控制：我们在提示信息中指定思考词元的上限；
（b）步骤条件控制：我们指定思考步骤的上限，其中每个步骤大约包含100个词元；
（c）类别条件控制：我们编写两条通用提示信息，分别告知模型进行短时间或长时间的思考。
（II）拒绝采样法：该方法会不断进行采样，直到生成的内容符合预先设定的计算预算。这种理想化的方法能捕捉到基于内容长度的响应后验分布。
在上图（左图）中，我们对下图（中图）的曲线进行了拓展。结果表明，虽然我们可以利用 “预算强制” 技术和增加测试时的计算资源来提高模型在AIME24测试中的性能，但当计算资源增加到六倍时，性能提升最终会趋于平缓。如果过于频繁地抑制思考结束词元分隔符的生成，可能会导致模型陷入重复循环，而非持续进行推理。
在上图（右图）中，我们展示了在使用1000个样本对Qwen2.5-32B-Instruct模型进行训练得到s1-32B模型，并为其配备简单的 “预算强制” 技术后，该模型呈现出不同的性能缩放模式。通过多数投票法增加基础模型的测试时计算资源，其性能无法赶上s1-32B模型，这证实了我们在前面的直觉判断，即顺序式缩放比并行式缩放更有效。下图给出了s1-32B模型的示例生成结果。
