s1: Simple Test-time Scaling

背景知识

近年来，大型语言模型（LLMs）的性能提升主要依赖于训练时计算量的增加，通过大规模的自监督预训练实现（Kaplan et al., 2020; Hoffmann et al., 2022）。这些强大模型的出现为一种新的扩展范式奠定了基础：测试时扩展（test-time scaling）。这种方法的目的是在测试时增加计算量以获得更好的结果。

最近，OpenAI 的 o1 模型展示了这种能力，但没有公开分享其方法，导致了许多人尝试复现。测试时扩展是一种有前途的新方法，它通过在测试时增加额外计算来提高语言模型的性能。

研究概述

本文旨在寻找实现测试时扩展和强推理性能的最简单方法。

首先，作者们策划了一个名为 s1K 的小型数据集，其中包含 1000 个问题，这些问题配有推理过程，这些推理过程基于三个标准：难度、多样性和质量。

其次，开发了一种名为“预算强制”（budget forcing）的技术，通过强制终止模型的思考过程或通过在模型试图结束时多次附加“Wait”来延长思考过程，从而控制测试时的计算量。这可以引导模型复核其答案，经常修复不正确的推理步骤。

在 s1K 数据集上对 Qwen2.5-32B-Instruct 语言模型进行监督微调，并结合预算强制技术后，得到的模型 s1-32B 在竞赛数学问题上的表现比 o1-preview 高出 27%（在 MATH 和 AIME24 数据集上）。此外，通过预算强制对 s1-32B 进行扩展，使其性能超越了没有测试时干预的性能：在 AIME24 数据集上从 50% 提升到 57%。

方法

数据集

s1K 数据集包含 1000 个精心策划的问题，配有推理轨迹和答案。这些问题和答案来自于 Gemini Thinking Experimental (Google, 2024)。作者通过消融实验验证了难度、多样性和质量这三个标准。

质量：首先删除任何存在 API 错误的 questions，将数据集减少到 54,116 个样本。然后，通过检查是否存在格式问题的字符串模式（如 ASCII 艺术图、不存在的图像引用或不一致的问题编号）来过滤低质量样本，将数据集减少到 51,581 个样本。
难度：使用两个指标：模型性能和推理轨迹长度。在每个问题上评估两个模型：Qwen2.5-7B-Instruct 和 Qwen2.5-32B-Instruct (Qwen et al., 2024)，并由 Claude 3.5 Sonnet 评估其正确性。测量每个推理轨迹的 token 长度，以指示问题的难度。
多样性：为了量化多样性，作者使用 Claude 3.5 Sonnet 根据美国数学学会的数学学科分类（MSC）系统将每个问题分类到特定领域。

算法：Budget Forcing

预算强制是一种简单的解码时干预方法，通过在测试时强制执行最大和/或最小数量的思维 token 来实现。具体来说，通过简单地附加思维结束 token 分隔符和“Final Answer:”来强制执行最大 token 数，以提前退出思维阶段，并使模型提供其当前的最佳答案。为了强制执行最小值，将抑制思维结束 token 分隔符的生成，并可选地将字符串“Wait”附加到模型的当前推理轨迹中，以鼓励模型反思其当前生成。

训练过程

在 Qwen2.5-32B-Instruct 模型上进行了监督微调（SFT），使用了 s1K 数据集，获得了 s1-32B 模型。微调过程在 16 个 NVIDIA H100 GPU 上进行，使用 PyTorch FSDP，耗时 26 分钟。

实验

实验设计

为了验证他们的方法，作者在三个具有代表性的推理基准上对s1-32B进行了测试：

AIME24：包含2024年美国邀请数学考试（AIME）中的30个问题，测试内容包括算术、代数、计数、几何、数论、概率和其他中学数学知识。
MATH500：包含难度各异的竞赛数学题。
GPQA Diamond：包含生物、化学和物理领域的198个博士级科学问题。

作者将s1-32B与其他几个模型进行了比较，包括：

OpenAI o1系列：这些模型推广了测试时扩展的理念。
DeepSeek r1系列：具有与o1相当性能的推理模型。
Qwen的QwQ-32B-preview：一个32B的reasoning模型。
Google Gemini 2.0 Flash Thinking Experimental：用于构建s1K数据集的模型。

实验结果

测试时扩展：图1显示，s1-32B的性能随着测试时计算的增加而提高。
样本效率：图2（右）和表1表明，s1-32B是最具样本效率的reasoning模型。尽管只对额外的1000个样本进行了训练，s1-32B的表现明显优于其基础模型（Qwen2.5-32B-Instruct）。
消融实验：通过消融实验证明数据数量、多样性和难度三个标准的重要性。

总结与评价

总的来说，这篇文章展示了一种简单而有效的方法来实现测试时扩展和强大的推理性能。通过精心策划的小数据集和简单的测试时技术，s1-32B 模型在多个推理基准测试中取得了优异的性能，甚至可以与更大、更复杂的模型相媲美。

文章的优点在于：

方法简单有效：仅需 1000 个样本的监督微调，结合预算强制技术，就实现了强大的推理性能和测试时扩展能力。
开源：模型、数据和代码都是开源的，这有助于推动社区对测试时扩展的研究。
实验充分：文章进行了大量的消融实验，验证了数据选择标准和测试时扩展方法的有效性。

文章的不足之处在于：

依赖 Gemini：s1K 数据集是通过 Gemini Thinking Experimental API 生成的，这可能限制了数据集的多样性和独立性。
预算强制的局限性：虽然预算强制可以有效地控制测试时计算量，但它可能无法充分利用模型的潜力，尤其是在处理更复杂的问题时。

总之，这是一篇很有意思的文章，它提出了一种简单而有效的方法来增强语言模型的推理能力。它为未来的研究提供了新的思路，并为开发更强大的语言模型提供了有益的参考。

Links

paper link: https://arxiv.org/abs/2501.19393

github link: https://github.com/simplescaling/s1

[Arxiv] s1: Simple test-time scaling