s1: 简单的测试时扩展

ArXiv ID: 2501.19393
作者: Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, Tatsunori Hashimoto
机构: Stanford University, University of Washington, Hugging Face
发布日期: 2025-01-31
模型: s1-32B (基于 Qwen2.5-32B-Instruct)

---

摘要

OpenAI o1 等推理模型展示了测试时计算扩展的巨大潜力，但其训练方法（大规模强化学习）成本高昂且不透明。本文证明，仅用 1000 个精选问题进行监督微调，配合 budget forcing 技术控制思考过程，就能实现强大的测试时扩展能力。s1-32B 在竞赛数学题上超越 o1-preview 多达 27%，在 AIME24 上通过扩展从 50% 提升至 57% 准确率。

---

核心洞察

主要发现

传统观点：
- 推理能力需要大规模 RL 训练
- 测试时扩展是复杂系统的专利
- 需要大量推理轨迹数据

s1 的发现：
- ✅ 仅用 1000 个问题就能学会推理
- ✅ 简单 SFT + budget forcing = 强大推理
- ✅ 测试时计算可被简单控制
- ✅ 32B 模型可超越 o1-preview

关键创新：Budget Forcing

什么是 Budget Forcing？

通过操控模型的思考过程来控制测试时计算量：

标准推理：
[思考...（不确定多长）] → [答案]

Budget Forcing 推理：
[思考...（恰好 N 步）] → [答案]

实现方式：

# 在训练数据中添加预算标记
training_example = {
    "prompt": "数学问题...",
    "response": """<budget>1000</budget>
<thinking>
步骤 1: 分析问题...
步骤 2: 制定计划...
步骤 3: 执行计算...
...（直到达到预算）
</thinking>
<answer>最终答案</answer>"""
}

---

方法详解

数据构建

数据来源：

数据集	问题类型	数量
AIME	数学竞赛	200
MATH	高中数学	300
AMC	数学竞赛	150
Codeforces	编程竞赛	200
GSM8K-Hard	算术推理	150
总计		1000

数据特点：

• 精选高质量问题
• 覆盖多个难度级别
• 包含详细解题步骤
• 答案格式统一

数据示例

问题：
Find the number of positive integers less than 1000 such that
x + ⌊x⌋ = 1000.

答案：
<budget>2000</budget>
<thinking>
让我分析这个问题...

首先，我需要理解 x + ⌊x⌋ = 1000 的含义。

设 x = n + f，其中 n 是整数部分，f 是小数部分（0 ≤ f < 1）

那么：
x + ⌊x⌋ = (n + f) + n = 2n + f = 1000

因为 2n 是整数，f 是小数，所以：
2n + f = 1000 意味着 2n = 999 或 2n = 1000

如果 2n = 999，则 n = 499.5，不是整数，舍去。
如果 2n = 1000，则 n = 500，f = 0。

所以 x = 500 是一个解。

等等，让我重新检查...

[继续推理直到预算用尽]
</thinking>
<answer>1</answer>

模型训练

训练配置：

参数	值
基座模型	Qwen2.5-32B-Instruct
训练数据	1000 个问题
学习率	2e-5
Batch Size	32
训练步数	3 epochs
序列长度	32768
训练时间	~8 小时 (8×H100)

训练成本对比：

方法	数据量	训练成本	性能
o1	百万级	$百万+	SOTA
s1	1000	<$1000	超越 o1-preview

---

Budget Forcing 技术

工作原理

思考预算控制：

┌─────────────────────────────────────────┐
│  模型学习在不同预算下调整推理深度       │
│                                         │
│  低预算 (100 tokens):                   │
│  [快速直觉 → 答案]                      │
│                                         │
│  中预算 (500 tokens):                   │
│  [分析 → 计划 → 执行 → 答案]            │
│                                         │
│  高预算 (2000 tokens):                  │
│  [分析 → 多方案比较 → 验证 → 答案]      │
│                                         │
└─────────────────────────────────────────┘

预算与性能关系

准确率
  │
  │         ● s1 (高预算)
  │      ●
  │   ●
  │ ●            ● o1-preview (固定)
  │
  └───────────────────────── 预算
    低    中    高

关键发现：

• 增加预算可持续提升性能
• 存在边际递减效应
• 最优预算因任务而异

---

实验结果

主要基准

数学竞赛

模型	AIME24	MATH500	AMC23	GSM8K
o1-preview	50%	82%	78%	96%
s1-32B (默认)	50%	79%	72%	94%
s1-32B (高预算)	57%	85%	81%	96%

关键结果：

• s1-32B 在 AIME24 上通过扩展超越 o1-preview 7 个百分点
• 在 MATH500 上超越 o1-preview 3 个百分点
• 在 AMC23 上超越 o1-preview 3 个百分点

综合对比

模型	平均	数学	代码	推理
GPT-4o	68%	65%	70%	69%
o1-mini	78%	82%	75%	77%
o1-preview	82%	85%	79%	82%
s1-32B	84%	88%	81%	83%

扩展曲线

AIME24 准确率 vs 预算：

60% │                          ●
    │                      ●
55% │                  ●
    │              ●
50% │          ● s1-32B
    │      ●
    │  ●
45% │●
    └────────────────────────────
      0   500  1000  2000  4000
            思考预算 (tokens)

观察：

• 预算增加，性能持续提升
• 在 2000-4000 tokens 趋于饱和
• 存在最优预算区间

效率分析

模型	准确率	平均 Token	每正确 1 题成本
o1-preview	50%	~5000	$0.10
s1-32B	57%	~3000	$0.02

s1 优势：

• 准确率更高
• 使用 token 更少
• 成本低 5 倍

---

案例研究

成功案例

问题（AIME24）：

Find the number of solutions to the equation sin(2x) = cos(3x) in the interval [0, 2π].

s1 的思考过程（简化）：

<budget>1500</budget>
<thinking>
我需要解 sin(2x) = cos(3x)...

方法 1：使用 sin(θ) = cos(π/2 - θ)
sin(2x) = cos(π/2 - 2x)
所以 cos(π/2 - 2x) = cos(3x)
π/2 - 2x = 3x + 2πn 或 π/2 - 2x = -3x + 2πn

情况 1: π/2 - 2x = 3x + 2πn
π/2 = 5x + 2πn
x = π/10 - 2πn/5

情况 2: π/2 - 2x = -3x + 2πn
π/2 = -x + 2πn
x = 2πn - π/2

现在找出 [0, 2π] 内的所有解...
[详细计算]

让我验证这些解...
[验证过程]

所以共有 5 个解。
</thinking>
<answer>5</answer>

失败分析

失败模式分类：

类型	比例	描述
计算错误	35%	算术或代数错误
推理跳跃	25%	步骤不完整
理解偏差	20%	误解问题
预算不足	15%	思考被截断
其他	5%	-

---

与相关方法对比

测试时扩展方法

方法	核心思想	数据需求	训练成本	性能
Repeated Sampling	多样本采样	无	无	中
Chain of Thought	思维链	少	低	中
Self-Refine	迭代改进	中	中	中高
o1 RL	大规模强化学习	百万级	高	SOTA
s1	Budget Forcing	1000	低	超 SOTA

核心优势对比

o1 方法：
优点：强大推理能力，SOTA 性能
缺点：黑盒，成本高，不可控

s1 方法：
优点：透明，低成本，可控，开源
缺点：需要精心设计预算机制

---

实践指南

部署 s1

from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载模型
tokenizer = AutoTokenizer.from_pretrained("s1-32B")
model = AutoModelForCausalLM.from_pretrained(
    "s1-32B",
    device_map="auto",
    torch_dtype=torch.bfloat16
)

def solve_with_budget(prompt, budget=1000):
    """使用指定预算解决问题"""
    messages = [
        {"role": "user", "content": prompt},
        {"role": "budget", "content": str(budget)}
    ]

    input_ids = tokenizer.apply_chat_template(
        messages,
        return_tensors="pt",
        max_length=budget + 1000
    ).to(model.device)

    output = model.generate(
        input_ids,
        max_new_tokens=budget + 500,
        do_sample=True,
        temperature=0.7
    )

    return tokenizer.decode(output[0])

# 使用示例
result = solve_with_budget(
    "Solve: x² + 5x + 6 = 0",
    budget=500
)

预算选择建议

任务类型	推荐预算	说明
简单算术	200-500	直接计算即可
高中数学	500-1000	需要多步推理
竞赛数学	1000-2000	复杂证明
研究级问题	2000-4000	深度探索

---

局限性与未来方向

当前局限

1. 领域限制：主要在数学/代码任务上验证
2. 预算选择：需要手动指定预算
3. 语言覆盖：主要支持英语

未来方向

1. 自适应预算：模型自动选择最优预算
2. 多模态扩展：视觉推理任务
3. 领域扩展：科学、法律、医疗

---

总结

s1 证明了简单的测试时扩展方法可以匹敌甚至超越复杂的推理系统：

核心贡献：

1. 仅用 1000 个问题实现强大推理能力
2. Budget Forcing 技术控制测试时计算
3. 在数学竞赛上超越 o1-preview 27%
4. 开源、透明、可复现

实际意义：

• 降低了推理模型的门槛
• 提供了测试时扩展的新思路
• 为资源受限团队提供可行方案

---

资源

• arXiv 论文
• GitHub 代码
• Hugging Face 模型

---

评分: 4.5/5.0 ⭐⭐⭐⭐⭐

推荐度: 强烈推荐。测试时扩展的简洁方案，开源可复现。