s1:Simple tets-time scaling略微解析
-
- 论文步骤
- 《s1: Simple test-time scaling》的代码和资源仓库。这个仓库提供了论文中提到的所有资源,包括模型、数据集、训练脚本、评估脚本等。
- **`SamplingParams` 参数解析**
论文步骤
-
数据集构建
- 初始数据收集:从16个不同的来源收集了59,029个问题,遵循质量、难度和多样性三个指导原则。
- 质量:确保数据集的格式和内容质量高,忽略格式不佳的数据。
- 难度:选择需要显著推理努力的问题。
- 多样性:涵盖不同领域的推理任务。
- 数据筛选:通过三阶段过滤(质量→难度→多样性)从59K初始数据集中筛选出1,000个高价值样本,形成s1K数据集。
- 质量过滤:删除低质量样本,如格式错误、包含ASCII艺术图、不存在的图像引用等。
- 难度过滤:使用两个模型(Qwen2.5-7B-Instruct和Qwen2.5-32B-Instruct)评估问题难度,移除简单问题。
- 多样性过滤:从过滤后的数据中均匀采样,确保样本多样性。
- 初始数据收集:从16个不同的来源收集了59,029个问题,遵循质量、难度和多样性三个指导原则。
-
模型训练
- 使用s1K数据集对预训练模型(Qwen2.532B-Instruct)进行监督微调(SFT),仅需16张H100 GPU训练26分钟。
-
测试时扩展方法(预算强制)
- 确定token上限:
- 基于模型最大生成能力:模型的最大生成token数为32,768。因此,token上限需要确保不超过模型的最大生成能力,以避免超出模型的处理范围[12]。
- 实验验证与性能平衡:通过实验验证不同token上限对模型性能的影响。例如,研究人员将token上限设置为32,000,这个值在确保模型有足够空间进行推理的同时,也避免了过长的推理过程导致的性能下降[12]。
- 终止生成:如果模型生成的推理token接近或达到设定的上限(例如32,000),则强制结束推理过程,通过附加思维结束token分隔符使模型进入答案生成阶段[12]。
- 延长窗口:如果希望模型在问题上投入更多测试时计算资源,则抑制思维结束token的生成,并在推理轨迹中追加“Wait”,鼓励模型进行更深入的推理探索[12]。
- 确定token上限:
-
模型评估
- 在竞赛数学题(MATH和AIME24)等任务上评估模型性能。s1-32B模型在AIME24任务上通过预算强制技术,准确率从50%提升到57%[12]。
核心思想
-
测试时扩展(Test-time Scaling)
- 利用额外的测试时计算资源来提升语言模型的推理性能。通过预算强制技术,动态调整模型在测试时的计算量。
-
样本高效(Sample Efficiency)
- 仅使用1,000个高质量样本进行微调,通过精心筛选数据,而不是依赖大量数据,来提升模型性能。研究表明,59K全量数据训练仅提升3%性能,但算力消耗远高于1K样本[12]。
-
数据质量优于数量
- 通过质量、难度和多样性三个标准筛选数据,验证了数据筛选的重要性。随机采样或仅选择长序列样本的性能远低于通过三阶段过滤得到的1K样本[12]。
-
顺序扩展优于并行采样
- 预算强制引导的链式推理性能增益显著高于多数投票等并行采样方法。通过控制生成的令牌数,逼近蒙特卡洛树搜索效果,但无需奖励模型[12]。
-
开源与可扩展性
- 论文开源了完整的训练管线、模型和数据,并展示了模型在不同任务上的可扩展性,例如在GPQA博士级科学难题上,测试时计算量增加可提升准确率[12]。
通过这些步骤和核心思想,论文证明了小模型在高效微调下可以匹敌顶级闭源模型,为轻量化高精度推理模型提供了新的范式。
《s1: Simple test-time scaling》的代码和资源仓库。这个仓库提供了论文中提到的所有资源,包括模型、数据集、训练脚本、评估脚本等。
1. 项目概述
s1: Simple test-time scaling
是一个关于如何通过简单的测试时扩展(test-time scaling)技术提升语言模型推理性能的研究项目。它展示了如何仅用1,000个样本和预算强制(budget forcing)技术,就能让模型在推理任务上达到与大型闭源模型相当的性能。
2. 主要资源
- 论文:s1: Simple test-time scaling
- 模型:s1-32B
- 数据集:s1K
- s1-prob: s1-prob
- s1-teasers: s1-teasers
3. 项目结构
eval/
:评估脚本,用于测试模型性能。data/
:合成数据创建脚本,用于生成训练和测试数据。train/
:训练脚本,用于对模型进行微调。
4. 推理代码示例
使用 vLLM
进行推理
vLLM
是一个高效的推理库,支持并行化推理。以下是使用 vLLM
进行推理的代码示例:
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
# 初始化模型和分词器
model = LLM(
"simplescaling/s1-32B",
tensor_parallel_size=2,
)
tok = AutoTokenizer.from_pretrained("simplescaling/s1-32B")
# 设置停止token
stop_token_ids = tok("<|im_end|>")["input_ids"]
# 设置采样参数
sampling_params = SamplingParams(
max_tokens=32768,
min_tokens=0,
stop_token_ids=stop_token_ids,
)
# 构造提示
prompt = "How many r in raspberry"
prompt = "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n"
# 生成推理结果
o = model.generate(prompt, sampling_params=sampling_params)
print(o[0].outputs[0].text)
使用 vLLM
进行预算强制推理
预算强制技术允许模型在测试时动态调整生成的token数量,以优化推理性能。以下是代码示例:
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
# 设置推理时的token上限
MAX_TOKENS_THINKING = 32000
NUM_IGNORE = 1 # 忽略结束token的次数
model = LLM(
"simplescaling/s1-32B",
tensor_parallel_size=2,
)
tok = AutoTokenizer.from_pretrained("simplescaling/s1-32B")
# 设置停止token
stop_token_ids = tok("<|im_end|>")["input_ids"]
sampling_params = SamplingParams(
max_tokens=32768,
min_tokens=0,
stop_token_ids=stop_token_ids,
skip_special_tokens=False,
temperature=0.0,
)
# 构造提示
prompts = ["How many r in raspberry"]
for i, p in enumerate(prompts):
prompt