S1：简单高效的测试时推理能力扩展方法

引言：测试时计算扩展的兴起与挑战

近年来，大型语言模型（LLMs）在自然语言处理领域取得了举世瞩目的成就，其性能的飞跃在很大程度上得益于大规模自监督预训练和不断增长的模型规模。然而，模型规模的不断扩大也带来了训练成本和计算资源需求的急剧增加，这使得进一步提升模型性能变得越来越困难。为了在不进一步增加模型规模的前提下提升性能，一种新的范式——测试时计算扩展（Test-Time Computation Extension）——应运而生。测试时计算扩展的核心思想是在测试阶段（即模型的推理阶段）投入额外的计算资源，从而增强模型的表现。这种方法避免了对模型进行重新训练或大规模修改，提供了一种灵活且高效的性能提升途径。

OpenAI 的 o1 模型（其具体实现细节并未公开）展示了测试时扩展的巨大潜力，这引发了大量的复现和改进工作。许多研究者试图理解并复制 o1 的成功，但由于缺乏公开信息，这些努力面临着巨大的挑战。为了探索一种既简单又有效的测试时扩展方法，斯坦福研究人员的论文《s1: Simple test-time scaling》提出了 S1 方法。S1 通过在一个精心构建的小规模推理数据集上进行微调，并结合易于实现的预算强制技术，实现了强大的推理性能。S1 的目标是提供一种可复现、易理解且高效的测试时扩展方案，为 LLM 社区提供一个有价值的参考。

S1 方法的核心：简约而不简单

S1 的整体思路非常简洁，可以概括为两个主要步骤：首先，在一个名为 S1K 的、仅包含 1000 个高质量推理样本的数据集上，对预训练语言模型进行监督微调（Supervised Fine-Tuning, SFT）。这一步使得模型能够更好地适应推理任务的特定需求。然后，在测试阶段，采用预算强制（Budget Forcing）技术，精确控制模型的推理过程，以此来提升模型的性能。预算强制通过限制模型的推理步数或推理时间，迫使模型在有限的资源内进行更高效的思考。

S1 的简单性体现在多个方面。首先，S1K 数据集规模很小，但质量非常高，其中的样本都经过了精心挑选，确保了训练数据的高效性。其次，模型的训练过程非常简洁，只需要标准的监督微调即可完成，无需复杂的训练策略或算法设计。最后，预算强制技术本身也易于理解和实现，不需要复杂的数学推导或编程技巧。这种简单性使得 S1 方法易于复现和应用，降低了研究门槛。

S1K 数据集：小而精的推理基石

S1K 数据集的构建是 S1 方法成功的关键。该数据集遵循三个核心原则：质量、难度和多样性。首先，样本必须具有高质量，这意味着样本必须避免格式错误、内容不准确或逻辑混乱等问题。只有高质量的样本才能为模型提供有效的学习信号。其次，样本应具有一定的挑战性，能够有效考察模型的推理能力，避免过于简单的问题导致模型无法得到充分的训练。最后，样本应覆盖多个领域和主题，以确保模型的泛化能力，避免模型在特定领域过拟合。

S1K 数据集的构建过程包括以下几个步骤：首先，从 16 个不同的数据源（涵盖数学、物理、生物等多个领域）收集了 59,000 个初始样本。这些样本包含了推理问题，并使用 Google Gemini Flash Thinking API 生成了相应的推理轨迹和答案。然后，进行质量过滤，移除包含 API 错误、格式问题、重复样本等低质量样本，确保数据集的基本质量。接下来，利用 Qwen2.5-7B-Instruct 和 Qwen2.5-32B-Instruct 模型的推理表现和推理轨迹长度作为指标，筛选出具有一定难度的样本，这保证了数据集能够有效评估模型的推理能力。最后，根据数学主题分类系统（MSC）对样本进行分类，并从每个类别中挑选具有代表性的样本，确保了数据集在不同领域和主题上的覆盖度。最终形成的 S1K 数据集具有小规模、高质量和多样性的特点，为 S1 方法的成功奠定了坚实的基础。

预算强制：S1 测试时扩展的核心机制

测试时扩展方法主要分为两类：顺序型和并行型。顺序型方法指的是模型的推理过程是逐步进行的，每一步的推理都依赖于前一步的结果。而并行型方法则允许模型同时进行多个推理过程，例如同时探索多个推理路径。S1 采用的是顺序型扩展方法，其核心技术是预算强制。预算强制通过两种方式来控制模型的推理过程。一方面，当模型生成的思考 Token 数量达到预设上限时，强制结束思考过程，并引导模型开始生成答案，这避免了模型无限制地进行推理，导致资源浪费。另一方面，当模型过早地尝试结束思考过程时，抑制其生成结束思考的 Token，并通过添加 "Wait" 等提示词，鼓励模型进行更深入的思考，这防止了模型在推理不充分的情况下给出答案。

预算强制的优势在于其简单性和有效性。它易于理解和实现，无需复杂的算法设计。同时，它能够有效控制模型的推理过程，显著提升模型在多个任务上的性能。与其他测试时扩展方法相比，例如条件长度控制、拒绝采样和强化学习，预算强制在简单性和有效性之间取得了较好的平衡。条件长度控制的局限性在于模型可能难以准确理解和遵循 Prompt 中的长度限制。拒绝采样的效率较低，并且可能导致模型过度依赖较短的推理路径。强化学习方法虽然理论上能够实现更精细的资源控制，但训练过程较为复杂。

S1-32B 模型与实验结果

S1-32B 模型是 S1 方法的具体实现，它基于 Qwen2.5-32B-Instruct 模型。S1-32B 的训练过程非常简洁，仅仅是在 S1K 数据集上进行监督微调。训练过程使用 16 块 H100 GPU，总训练时间仅需 26 分钟。实验结果表明，S1-32B 在多个推理任务（如 MATH500、AIME24 和 GPQA Diamond）上都取得了优异的性能。而且，随着预算强制计算量的增加，模型性能也随之提升，这有力地验证了预算强制的有效性。

为了深入理解 S1 的成功因素，研究团队进行了一系列消融实验。数据消融实验旨在探究 S1K 数据集构建的三个原则（质量、难度和多样性）对模型性能的影响。实验结果表明，这三个原则的结合对于模型性能的提升至关重要，缺一不可。测试时扩展方法消融实验则旨在比较不同测试时扩展方法对模型性能的影响。实验结果表明，预算强制在控制性、扩展性和最终性能方面均优于其他方法。

S1 的优势、局限性与启示

S1 方法的优势在于其简单性、高效性和有效性。它易于理解、实现和复现，研究者和开发者可以快速上手，进行实验和应用。同时，使用少量高质量数据即可快速完成模型训练，显著降低了训练成本。更重要的是，S1 在多个推理任务上取得了与 o1-preview 相当的性能，证明了其在提升 LLM 推理能力方面的潜力。 并且S1的开源特性肯定会进一步促进了LLM社区的协作发展。

然而，S1 方法也存在一些局限性。一是S1 的性能一般来说可能是存在上限的，低于更复杂的方法（如 DeepSeek-R1）。这可能是因为 S1 仅依赖于简单的监督微调和预算强制，而没有引入更复杂的训练策略或模型架构。二是预算强制本身也存在局限性。随着预算强制计算量的增加，模型性能的提升会逐渐趋于平缓，甚至可能出现下降。这意味着预算强制存在一个最优的计算量范围，超过这个范围可能导致资源浪费或性能下降。另外，预算强制的有效性受到模型上下文窗口长度的限制，对于需要更长推理路径的任务，预算强制可能无法充分发挥作用。预算强制还缺乏自适应性，它是一种静态的控制方法，无法根据不同的输入样本动态调整计算资源。对于不同难度的样本，使用相同的计算预算可能并不是最优策略。 S1还对顺序推理有依赖, 这限制了它利用并行计算加速推理的能力。 并且预算强制机制的内部工作原理仍然不够透明, 增加了理解模型决策的难度.

当然，S1 的成功也带来了一些重要的启示。首先，它强调了数据质量的重要性，即使是规模较小的数据集，只要质量足够高，也能显著提升模型性能。其次，它表明简单的策略在特定场景下也能取得与复杂方法相媲美的效果，这启示我们不应一味追求复杂性，而应根据具体问题选择合适的解决方案。最后，S1 的成功再次证明了测试时扩展的巨大潜力，为 LLM 性能提升提供了一条新的途径。

S1的总体评价

S1方法为测试时计算扩展提供了一个新视角，证明了即使是简单的方法，也可以在提升大型语言模型推理能力上取得显著效果。 它的主要优点是简单, 高效, 并且开源。这使得研究者和开发者更容易复现和应用S1, 促进了整个LLM社区的合作与发展。

但是，S1 的局限性也是显而易见的。 它的性能上限可能低于更复杂的方法。预算强制技术本身也存在局限, 比如性能提升会趋于平缓, 并且受限于模型的上下文窗口大小。 S1高度依赖于S1K数据集的质量, 并且对顺序推理的依赖限制了它利用并行计算的能力。 此外, 预算强制的内部机制不够透明, 使得模型的可解释性较差。这些局限性也为未来的研究提供了方向。

什么是测试时扩展？

这个术语太不友好，我花点篇幅解释一下吧。

在机器学习领域，模型的生命周期通常划分为两个主要阶段：训练阶段和测试阶段（也称为推理阶段）。训练阶段是模型学习的过程，通过大量的数据和计算，模型不断调整其内部参数，以尽可能准确地完成预定的任务，如图像识别、文本生成或语言翻译等。一旦模型训练完成，其内部参数就被固定下来。测试阶段则是对训练好的模型进行性能评估的环节，在这个阶段，模型会接收之前未见过的新数据，并根据这些新数据进行预测或生成输出。关键在于，在传统的测试流程中，模型不会再更新其参数，而是直接利用训练阶段学习到的知识来进行推理。

测试时扩展（Test-Time Computation Extension，简称 Test-Time Extension）是一种旨在提升模型在测试阶段性能的技术范式。它的核心思想非常直观：在测试阶段，也就是模型进行推理的时候，为模型提供额外的计算资源或者计算步骤，以此来提高模型输出结果的准确率或整体性能，而无需对模型本身进行任何修改（即不改变模型在训练阶段学习到的参数）。这种方法与传统的“训练-测试”二阶段流程有所不同，它在测试阶段引入了额外的计算干预。

采用测试时扩展的主要动机有以下几点：首先，它允许在不增加模型规模的前提下提升性能。训练更大、更复杂的模型通常需要耗费更多的计算资源和时间，而测试时扩展提供了一条在资源受限的情况下提高性能的途径。其次，测试时扩展具有很强的灵活性。由于它不涉及模型的重新训练，因此可以根据实际需要在测试时动态地调整计算量，从而在模型性能和计算成本之间取得平衡。此外，对于某些特定类型的任务，尤其是那些需要复杂推理的任务，标准的推理过程可能不足以让模型充分发挥其潜力，测试时扩展允许模型在这些任务上进行更深入、更细致的“思考”。最后，测试时扩展还可以在一定程度上弥补训练阶段的不足，例如当训练数据不够充分或者训练方法存在局限时，通过在测试时增加计算，模型有可能更好地处理那些在训练时没有充分学习到的复杂模式。

测试时扩展与一些常见的机器学习技术有着明显的区别。例如，它不同于训练时的数据增强（Data Augmentation）。数据增强是在训练阶段对训练数据进行各种变换，以增加数据的多样性，从而提高模型的泛化能力。而测试时扩展则完全作用于测试阶段，与训练数据无关。它也不同于集成学习（Ensemble Learning）。集成学习通常需要训练多个不同的模型，然后在测试时将这些模型的预测结果进行整合，例如通过投票或平均的方式。测试时扩展通常只针对单个模型，通过增加该模型在测试时的计算量来提升其性能。测试时扩展也不等同于直接增加模型的深度或复杂度，后者是在模型设计阶段就确定下来的，并且会改变模型本身的结构，而测试时扩展不会改变模型的固有结构。

常见的测试时扩展方法包括多次推理（让模型对同一输入进行多次独立推理并选择最一致的答案）、思考链（通过提示词引导模型逐步推理）、自修正（让模型对自己的输出进行迭代改进）以及搜索算法（在测试时探索多个可能的输出）等。每种方法都有其适用的场景和优缺点。 S1方法，通过在小规模数据集上进行微调，再结合推理时的预算强制，实际上也是一种测试时扩展策略，它通过限制推理步骤来达到提升效果的目的。

总之，测试时扩展是一种在不改变模型本身的情况下，通过在测试阶段增加计算资源或计算步骤来提升模型性能的技术。它提供了一种灵活且有效的方法，尤其适用于那些需要复杂推理的任务，为大型语言模型在实际应用中的性能提升开辟了新的可能性。

结论与未来展望

S1 的主要贡献在于提出了一种简单、高效且有效的提升 LLM 推理能力的方法，并构建和开源了高质量的 S1K 推理数据集。同时，S1 验证了预算强制技术在测试时扩展中的有效性，并强调了 S1 方法的开源特性。未来的研究方向包括探索更有效的测试时扩展方法，例如改进预算强制技术，或者将其与强化学习等更高级的技术结合起来。此外，可以将 S1 的思想应用于其他领域和任务，如代码生成、文本摘要、对话系统等，以验证其通用性。另一个重要的研究方向是突破上下文窗口的限制，例如采用并行扩展方法或分块推理等，以处理更复杂的推理任务。增强模型的可解释性, 深入理解预算强制机制的内部工作原理也是一个重要的研究方向。继续扩大和完善 S1K 数据集，使其涵盖更广泛的领域和任务, 以及探索将人类反馈融入S1, 都有助于提高模型的泛化能力和可靠性。

总而言之，S1 方法为 LLM 测试时计算扩展提供了一个简单而有效的解决方案，并为未来的研究提供了宝贵的启示。随着 LLM 技术的不断发展，相信 S1 方法及其衍生技术将在更广泛的领域得到应用和推广，推动 LLM 推理能力的持续提升。

参考论文： arXiv:2501.19393v2 [cs.CL] 3 Feb 2025

本文转载自​​上堵吟​​​，作者：​​上堵吟​​