季逸超 Peak：迈向复现 OpenAI o1 的一小步之 Steiner 开源模型阶段性进展报告｜Z Talk

文｜季逸超 Peak，Peak Labs 创始人，真格基金 EIR（Entrepreneur In Residence，入驻创业者）

从 OpenAI o1 发布以来，我就把复现 o1 作为业余时间的兴趣项目，其过程中获得了很多有趣的知识与洞察 (与失败)。考虑到个人能力、精力、财力有限，我很可能不是走到最后的选手，因此我觉得很有必要将这些价值数十张 H100 的经验及时分享出来。

本文的英文版同步发表于：

https://medium.com/@peakji/a-small-step-towards-reproducing-openai-o1-b9a756a00855

01

TL;DR

Steiner 是一个 reasoning 模型，能在推理时以自回归的形式探索多种路径，并在必要时自主进行验证或回溯。训练过程分为三个步骤：首先，通过随机截断合成的 reasoning path 并进行引导式再补全，将每个样本转化为从问题到答案的有向无环图（DAG）。接着，从 10K 个 DAG 中采样出 50K 个具有回溯节点的 reasoning path 样本。最后，使用 DAG 中节点的入边和出边数量及其距离作为 heuristic rewards，通过强化学习让模型学会平衡探索的广度和深度。该模型在 GPQA-Diamond 数据集上实现了 +5.56 的提升，但未能复现 inference-time scaling。

模型下载地址：

https://huggingface.co/collections/peakji/steiner-preview-6712c6987110ce932a44e9a6

02

简介

Steiner 是一系列使用强化学习在合成数据上训练得到的 reasoning 模型，能够在推理时以自回归的形式尝试多种推理路径，并在必要时自主进行验证和回溯，从而实现在同一上下文中对隐含的搜索树的线性遍历。

Steiner 能够在推理时以自回归的形式尝试多种推理路径，并在必要时自主进行验证和回溯，从而实现在同一上下文中对隐含的搜索树的线性遍历。

Steiner 是 季逸超 (Yichao 'Peak' Ji) 的个人兴趣项目，受 OpenAI o1 启发，最终目标是复现 o1 并验证 inference-time scaling 曲线。Steiner-preview 模型是一个 work-in-progress。开源的目的是因为我观察到以选择题为主要构成的自动化评估手段较难充分体现 reasoning 模型的进展：事实上，「正确答案一定在选项中」这一前提就与真实的 reasoning 场景并不对齐，它会鼓励模型进行代入验证而非开放式探索。因此，我选择将阶段性成果开源出来，并在精力允许的情况下 build in public，从而在分享知识的同时获得更多来自真实人类的评估和反馈。

⚠️Disclaimer: 到目前为止，Steiner 可以在无需 Chain of Thought (CoT) prompting 和 agent framework 的情况下获得较高质量的 zero-shot 结果，但是仍未能复现 o1 展示出的 inference-time scaling 能力：在使用特殊的 logits processor（https://gist.github.com/peakji/f81c032b6c24b358054ed763c426a46f）对 reasoning tokens 进行干预的实验中，增加额外的 reasoning steps 没能提升性能，反而在 MMLU-Pro 和 GPQA 等 benchmark 上有所下降。因此 Steiner 暂时还不能被视为 OpenAI o1 的成功复现，可能在训练方法和数据质量上都存在不足，请谨慎参考。

03

背景

与传统 LLM 相比，OpenAI o1 最显著的变化是在推理时引入了 reasoning tokens，进而实现了 inference-time scaling：即通过增加推理时的 compute budget 来提升模型的表现。当我们讨论 inference-time scaling 时，最直觉的方法可能是引入 tree search 或 agentic framework。但是，我在阅读 o1 相关的 (有限的) 官方信息时，注意到其 report 的 benchmark 多是 pass@1 和 majority voting，而且 OpenAI 团队也提到 o1 是 single model 而非 system。这让我对其实现的方法非常好奇。

OpenAI API 中至今仍未将 reasoning tokens 的具体内容开放给开发者，但好在 tokens 用量的统计数据中会包含 reasoning tokens 的数量 (毕竟要根据这个向开发者收钱)。借此我设计了一个简单的实验，尝试通过 o1 的 API 来统计 completion (包含 reasoning) tokens 的数量与总请求时长的关系。我们知道，若使用 tree search，为了复用 cache 和提升 GPU 利用率会尽可能的并行推理，那么得到曲线应是 sub-linear 的。然而，实验的结果却是几条漂亮的直线，o1-mini 的波动甚至比 GPT-4o-mini 还小：

OpenAI o1 可能仍是一个进行线性自回归解码的模型，但这不代表它在推理阶段没有进行 「搜索」。

上述实验让我产生了一种猜想：OpenAI o1 可能仍是一个进行线性自回归解码的模型，但这不代表它在推理阶段没有进行「搜索」。想象有一颗搜索树，我们在对它进行遍历的时候，产生的轨迹其实是线性的。如果我们能够训练一个自回归语言模型，让它不仅能生成推理路径，还能在必要时进行验证、回溯、或切换思路，那么在最理想的情况下它实际上完成了在同一上下文中对隐含的搜索树的线性遍历。这种线性遍历听起来很浪费，但与并行搜索相比，它具有三个极其优秀的性质：

1. 前置的尝试无论对错都在 context memory 中，每一步决策均基于过去的完整信息；
2. 隐含的 backtracking 无需保证目标节点已经在搜索树内，探索更加自由；
3. 工程层面可以复用一切现有的经过高度优化的推理 infrastructure。

想象有一颗搜索树，我们在对它进行遍历的时候，产生的轨迹其实是线性的。

04

方法

要获得具有线性搜索能力的模型并不容易，无论是数据合成还是模型训练都面临许多挑战。

首先，目前可以取得的 reasoning 数据集大多仅有合成的 CoT 或 reasoning steps，它们一般是通过将「问题-答案」tuple 输入给强大的 LLM，然后要求模型拆解思考过程获得的。这种方法决定了这些数据集不会包含合理的 backtracking 节点，故而利用这种数据训练得到模型其实只学到了 shortcut，或者说是将 CoT internalized 了。

针对这个问题，我设计了两种数据合成与扩增方法：

1. 对上述 shortcut 数据集进行随机截断，并隐藏正确答案，让强大的 LLM 基于被截断后的前缀去尝试正向推理一定的 steps，然后再提供正确答案，以获得 backtracking 样本；
2. 对上一步产生的 steps 进行聚类后赋予唯一 ID，将同一个问题下的所有 steps 构建为有向无环图 (DAG)，并对 DAG 进行随机采样以获得多项式量级的 reasoning path 样本。
3. 通过上述方法 (以及大量人工和奇技淫巧)，我最终获得了 10K 张有向无环图，并在此基础上采样了 50K 条带有 backtracking 的 reasoning path 样本。其中每个样本的 reasoning tokens 的平均数量约 1600 个，刚好与先前对 o1/o1-mini 进行测试时的统计数据非常接近！考虑到训练成本，我仅保留了 reasoning tokens 数量在 4096 以下，且 prompt + reasoning + completion 总 tokens 数量在 8192 以下的样本。

接着，我将 Steiner 系列模型的训练分为三个阶段：

1. Continual Pre-Training (CPT)：混合普通文本语料和 reasoning path 进行训练，目的是让模型习惯 long reasoning output，并且初步训练新增的 14 种特殊 tokens（https://huggingface.co/peakji/steiner-32b-preview/blob/278989ea17f74b14e2b32d9544eb53a17b4ad087/special_tokens_map.json#L16-L29）的 embeddings。需要指出的是，在一些小参数量模型的测试中表明，这一步可能是冗余的，直接在 SFT 阶段用大量 reasoning 数据训练似乎也能获得不错的表征，但 32B 的 CPT 很早就做完了，我就沿用到现在了；
2. Supervised Fine-Tuning (SFT)：使用 chat template 进行训练，目的是让模型学会模仿 reasoning 的格式：先给每个 step 想一个 name，再输出完整的 thought，然后将 thought 总结为 summary，接着对到此为止的推理进行 reflection，最后决定下一步是 proceed、backtrack 或是结束 reasoning 并开始正式回答问题。你可能会好奇，开源模型既然无需隐藏 thoughts 为什么也要像 o1 那样生成 summary 呢？这是因为我在为后续的具有多轮对话能力的 Steiner 模型做准备。理论上来说，经过训练后可以选择将先前对话中的完整 thought 用 summary 取代，来减少无法命中 prefix cache 时的 pre-fill 开销。目前 Steiner 还未针对多轮对话做优化，仅保留 summary 可能会产生负向的 few-shot 效果；
3. Reinforcement Learning with Step-Level Reward (RL)：经过前两阶段，模型已经学会生成和补完 reasoning path，但它还不知道哪些是正确且高效的选择。如果我们盲目 reward 更短的 reasoning path，则会降级为 shortcut 学习来 internalize CoT。这里我设计了一种启发式的 reward 机制：基于 DAG 中每个节点的入边数量 e_i、出边数量 e_o、距离原始问题的距离 d_s、距离正确答案的距离 d_e 来加权，为每个 step 以及整个 reasoning path 赋予 reward，以引导模型学会平衡探索的广度和深度。

以上思路看起来很简单，但在过去的一个月多里，每个周末（和没加班的夜晚）我都在与 OOM 和 reward hacking 做斗争。终于，在 OpenAI o1 发布后的第 38 天，训练出了不算太丢人的阶段性成果。

05

评估

不同训练阶段的 Steiner 模型在 GPQA-Diamond 数据集上的表现。Steiner 系列模型均不使用 CoT prompting。

图中展示的是不同训练阶段的 Steiner 模型在 GPQA-Diamond 数据集上的表现。可以看到引入强化学习阶段让模型有了 +3.53 的提升。若配合用于约束推理步数的 logits processor（https://gist.github.com/peakji/f81c032b6c24b358054ed763c426a46f），在最优配置下可以获得 +5.56 的提升。

使用特殊 logits processor 的 Steiner-32B 模型在 GPQA-Diamond 的各 subdomain 上的表现。

选择展示该 benchmark 首先是因为 o1/o1-mini 在该数据集上有较大提升，同时该数据集的 contamination 情况较为理想。其次，也是因为我发现 Steiner 在 MMLU 等数据集上与 baseline 相比没有明显的差异，这与 OpenAI 在关于 o1-mini 的博客（https://openai.com/index/openai-o1-mini-advancing-cost-efficient-reasoning/）中的观察类似，可能是受限于 32B 的模型在 pre-training 阶段所获得的 world knowledge。

无论从 benchmark 还是实际使用来看，都必须承认当前模型的能力较 o1-mini 和 o1 还有很大差距。但另一方面，主要由选择题构成的自动化评估 benchmark 可能无法充分体现出 reasoning 模型的能力：在训练阶段，reasoning 模型被鼓励对问题进行开放性的探索，而选择题存在「正确答案一定在选项中」这一前提，很明显将选项逐一代入验证才是更高效的方法。事实上，现有的大语言模型均有意或无意地掌握了这种技巧，无论是否使用特殊的 prompt。当然，正是因为自动化评估与真实 reasoning 需求的不对齐，让我认为更有必要将模型开源出来供真实的人类进行评估和反馈。

更大的问题在于 inference-time scaling 实验：很不幸，目前 Steiner 在使用 logits processor（https://gist.github.com/peakji/f81c032b6c24b358054ed763c426a46f）增加额外的推理步数后，在几乎所有测试的 benchmark 上都没有提升，甚至出现了下降。目前猜测的原因非常多，包括但不限于：

• CPT/SFT 并未让模型真正掌握超长的 output tokens；
• RL 阶段启发式的 reward 设计未能让模型在推理的广度和深度之间取得平衡；
• backtrack 时受 context 影响容易确认错误答案或否认原本正确的答案；
• 强行修改 logits 来增加推理步数的方法过于粗暴；
• … …

总之，这将是我接下来工作的重点。

06

限制

目前 Steiner 的 post-training 数据中不包含多轮对话的样本。当前性能最好的 Steiner 模型版本 (基于 Qwen2.5-32B) 完全不具备多轮对话能力。开源的 Steiner-preivew 模型 (基于 Qwen2.5-32B-Instruct) 可以兼容 chat 格式，但依然不建议用于多轮对话。

与 OpenAI o1-2024-09-12 类似，Steiner 也不建议使用自定义的 system prompt 或对 temperature 等 sampling 参数进行修改。目前 Steiner 还未在多样化的 system prompt 上进行训练，而修改其他参数可能导致 reasoning tokens 的格式出现错误。

Steiner 的 post-training 数据中语言的占比约为 90% 英文与 10% 的中文，但在 reasoning path 数据扩增过程中几乎完全使用英文。因此虽然模型的最终回答具备一定程度的语言跟随能力，但 reasoning tokens 部分可能均为英文。

07

彩蛋

关于 Steiner 这个名字的来源，不妨问问模型吧：

你的名字的来源是什么？

我的名字叫做 Steiner，来源于《命运石之门》中主角的能力「Reading Steiner」。这种能力使得主角能够保留多个世界线的记忆，这非常契合我作为一款具有推理能力模型的特点：在进行推理时，我可以尝试多种路径，并在必要时进行跳转或回溯，同时保留所有探索过的路径的记忆。

08

引用

如果你觉得我的工作有帮助，请考虑在你的研究或项目中引用它。你的认可将不胜感激！

@misc{ji2024steiner,title = {A Small Step Towards Reproducing OpenAI o1: Progress Report on the Steiner Open Source Models},url = {https://medium.com/@peakji/b9a756a00855},author = {Yichao Ji},month = {October},year = {2024}}