当AI自我纠错：一个简单的“Wait“提示如何让模型思考更深、推理更强

原论文：s1: Simple test-time scaling
作者：Niklas Muennighoff, Zitong Yang, Weijia Shi等（斯坦福大学、华盛顿大学、Allen AI研究所、Contextual AI）
论文链接：arXiv:2501.19393
代码仓库：GitHub - simplescaling/s1

1. 研究背景与问题

近年来，大语言模型的性能提升主要依赖于扩大预训练计算规模。而最近OpenAI的o1模型展示了一种新的范式——测试时间扩展（Test-time Scaling），即通过增加测试阶段的计算量来提升模型性能。然而，OpenAI并未公开其方法，这导致了许多复制尝试。

核心问题：如何以最简单的方法实现测试时间扩展并获得强大的推理性能？

2. 核心思路与方法

论文提出了两个关键创新：

2.1 高效数据策划（s1K数据集）

作者精心策划了一个仅包含1,000个问题的高效训练数据集，基于三个核心原则：

1. 质量：确保数据集中的问题和解答格式优良
2. 难度：选择足够具有挑战性的问题，需要复杂推理
3. 多样性：涵盖不同领域的问题（数学、物理、生物等）

数据筛选流程：

• 从16个来源收集初始的59,029个问题
• 使用API质量过滤和格式检查，减少到51,581个样本
• 使用模型性能（Qwen2.5-7B和32B均无法解决）作为难度指标
• 在50个不同领域中均衡采样，确保多样性

2.2 预算强制（Budget Forcing）技术

这是论文的核心创新，一种简单但高效的测试时间控制方法：

两种操作模式：

1. 最大限制强制：当模型生成的思考标记超过预设上限时，强制附加思考结束标记分隔符，迫使模型提供当前最佳答案
2. 最小限制强制：当模型试图过早结束思考时，抑制结束标记生成，并插入"Wait"提示词，鼓励模型继续思考和自我纠正

预算强制不需要额外训练，纯粹是解码时的干预技术，实现简单但效果显著。

3. 实验设计

3.1 模型训练

• 基础模型：Qwen2.5-32B-Instruct
• 训练数据：s1K数据集（1,000个推理样本）
• 训练效率：16个H100 GPU，仅需26分钟
• 训练参数：5个epoch，batch size=16，学习率1e-5，余弦学习率衰减

3.2 评估基准

选择了三个广泛使用的推理密集型基准：

1. AIME24：2024年美国邀请数学考试，30个高难度数学问题
2. MATH500：500个竞赛级数学问题
3. GPQA Diamond：198个博士级科学问题（生物、化学、物理）

3.3 测试时间扩展方法对比

研究比较了多种测试时间扩展方法：

1. 预算强制（提出的方法）
2. 标记条件控制：在提示中指定思考标记上限
3. 步骤条件控制：在提示中指定思考步骤上限
4. 类别条件控制：使用泛化提示（短/长思考）
5. 拒绝采样：采样直至生成符合预定长度的回答

评估指标包括：

• 控制性（Control）：方法能够精确控制测试时间计算的程度
• 扩展性（Scaling）：性能随计算增加的提升斜率
• 性能（Performance）：方法能达到的最高准确率

4. 实验结果

4.1 测试时间扩展效果

s1-32B模型在三个基准上都展现了明显的测试时间扩展能力：

• MATH500：随思考时间从512增至2048标记，准确率从约75%增至93%
• AIME24：随思考时间从512增至8192标记，准确率从约20%增至近60%
• GPQA Diamond：随思考时间从1024增至4096标记，准确率从约40%增至60%

4.2 不同测试时间扩展方法比较

方法	控制性(%)	扩展斜率	性能(AIME24)	评估次数
预算强制	100%	15	56.7%	5
标记条件控制	40%	-24	40.0%	5
标记条件+预算强制	100%	13	40.0%	5
步骤条件控制	60%	3	36.7%	5
步骤条件+预算强制	100%	6	36.7%	5
类别条件控制	50%	25	36.7%	2
拒绝采样	100%	-35	40.0%	5

预算强制展示了最佳的综合表现：

• 完美的控制性（100%）
• 良好的扩展斜率（15）
• 最高的性能（56.7%）

4.3 与现有模型对比

模型	样本数量	AIME 2024	MATH 500	GPQA Diamond
API 模型
o1-preview	N.A.	44.6%	85.5%	73.3%
o1-mini	N.A.	70.0%	90.0%	60.0%
o1	N.A.	74.4%	94.8%	77.3%
Gemini 2.0 Flash	N.A.	60.0%	N.A.	N.A.
开源权重模型
Qwen2.5-32B-Instruct	N.A.	26.7%	84.0%	49.0%
QwQ-32B	N.A.	50.0%	90.6%	54.5%
r1	>800K	79.8%	97.3%	71.5%
r1-distill	800K	72.6%	94.3%	62.1%
开源权重和数据
Sky-T1	17K	43.3%	82.4%	56.8%
Bespoke-32B	17K	63.3%	93.0%	58.1%
s1-32B（本文）	1K	56.7%	93.0%	59.6%

s1-32B是最具样本效率的开源推理模型，性能超过o1-preview，并且数据集大小比其他模型少1-800倍。

4.4 数据策划方法的重要性

数据集	AIME 2024	MATH 500	GPQA Diamond
1K-random（仅质量）	36.7%	90.6%	52.0%
1K-diverse（仅多样性）	26.7%	91.2%	54.6%
1K-longest（仅难度）	33.3%	90.4%	59.6%
59K-full（全量数据）	53.3%	92.8%	58.1%
s1K（质量+难度+多样性）	50.0%	93.0%	57.6%

结果证明，三个数据筛选标准（质量、难度、多样性）的结合对于高效推理训练至关重要。仅依赖单一标准会导致明显的性能下降。

5. 关键发现与创新

5.1 预算强制的有效性

论文展示了一个简单但强大的预算强制示例（图3）：

• 当模型错误计算"raspberry"中字母’r’的数量（误认为是2个）
• 预算强制通过插入"Wait"提示模型重新检查
• 模型进行自我修正，得出正确答案（3个’r’）

这种简单的干预能提高模型思考质量，无需额外训练成本。

5.2 测试时间扩展的局限性

研究也发现了测试时间扩展的两个主要限制：

1. 性能平缓化：增加思考时间到一定程度后，性能提升趋于平缓
2. 上下文窗口限制：模型的上下文窗口限制了能使用的最大思考标记数

5.3 拒绝采样的反向扩展现象

有趣的是，研究发现使用拒绝采样时出现了反向扩展现象：随着允许的思考标记数增加，准确率反而下降。这表明：

• 较短的生成往往来自模型一开始就走在正确轨道的情况
• 较长的生成往往来自模型犯了错误需要回溯或自我质疑的情况

6. 结论与启示

s1论文提供了关于测试时间扩展的重要启示：

1. 简单有效：简单的方法（1K样本+预算强制）可以实现与复杂方法相媲美的性能
2. 资源高效：训练仅需7个H100 GPU小时，远低于竞争方法
3. 开源透明：全部代码、数据和模型完全开源，促进社区研究

论文表明，语言模型在预训练中已经获得了推理能力，只需要适当的监督微调和测试时干预就能激活这些能力。这为构建高性能、低成本的推理系统提供了新的可能性。

7. 未来研究方向

研究指出了几个有前景的未来方向：

1. 并行扩展补充：结合并行方法（如REBASE）可以突破顺序扩展的上下文窗口限制
2. 预算强制改进：尝试轮换不同提示词或结合频率惩罚来避免重复循环
3. 与强化学习结合：探索预算强制是否能与RL训练模型结合获得更好的外推性能

这项工作为推理模型研究建立了新的基准，并为测试时间计算优化提供了实用框架。