也来猜猜 o1 实现方法

作者：朱小霖

链接：https://zhuanlan.zhihu.com/p/839732117

最近几周都在开心地猜 o1 可能是怎么做的，目前思路上相对收敛了，所以来记录一下，一方面是可以等答案揭晓的时候拿来对比看看，另一方面是自己手里没有真的能去验证想法的资源，所以也可以把心收回去去补落下的工作，等社区的各种资源（例如数据集）更完备再说了。

目前来看 o1 这里唯一明确的信息就是用了 RL。也就是用某种 reward model 指引模型，来提升模型的效果。那么一个最重要的问题是如何定义 reward。我认为在考虑 reward model 的时候，有 2 个点是比较重要的：

一个是如何定义一种通用的 reward。我认为 openai 训练的 reward model 不应该是专门用来做数学题或者代码题的，而是应该用来判断模型的回答多大程度上有了某种更通用的特性，因为只有这样才能让这个 reward model 带来的模型能力可以泛化到别的领域去。举例来说，一个比较朴素的通用 reward 可以是，模型的输出越长，多样性越丰富，分数越高；

另一个是如何收集训练数据。虽然上面说了我认为 o1 应该不止是训了个做题家（这里我们排除 o1-mini，因为依照公开信息感觉 o1-mini 可能是一个从预训练数据开始就集中于做题的模型），但是在采集数据的时候应该是充分利用了数学题难解好判（hard to generate, easy to verify）的特点，再通过 RL 的步骤将生成模型的解题能力逐渐拉近至 reward model 的判题能力。所以我们可以先去考虑如何去生成基于数学题的 reward model 数据。

为了上一步的通用性，我不太相信数据只会被整理成 ORM（outcome reward model）需要的数据（即，每条数据只标记最后的结果是否是正确的），因为那样太容易让 reward model exploit 特定领域的一些小 trick。那么如果要给模型的中间结果打分，我觉得还是 https://arxiv.org/abs/2305.20050 中提到的 PRM（process-supervised reward model）更合理一些（即，收集数据的时候，记录每一个推理步骤的分数；训练的时候，让 reward model 去预测每一步的分数）。

随之而来的又有 2 个问题，如何高效标注中间步骤的分数以及如何生成足够长的 CoT 数据，以防止 reward model OOD。

对于前者，虽然 openai 财大气粗，可以先用大量人力 bootstrap，我认为还是需要一个自动标注的方式。这里我看到有 https://arxiv.org/abs/2312.08935 和 https://arxiv.org/abs/2406.06592 这样的工作，可以按某一步有多大概率最后能得到正确的结果，来给每一步的重要性赋值，感觉很合理。

对于后者，我觉得让单一模型去生成数据的难度比较大。https://arxiv.org/abs/2408.07055 中提到，模型输出长文的能力和 sft 数据集的长度分布关系比较大，那么我感觉很难强行扩长（例如使用 prompt engineering）现有模型的输出，来提供包含足够有效信息的训练数据。一个也许可行的思路，是用多模型（同一个模型，但是用不同的 system message），来造出足够长的数据，例如说一些模型负责做计划，一些模型负责细化，一些模型负责批评和验证。目前看到的论文中，https://arxiv.org/abs/2409.12917 比较像是这个思路，不过他的轮数还是很少。而且我没太想明白，假如要用多 agent 构造数据，如何能防止在生成数据这个阶段引入太多的先验，比如说要求模型每一步都反思一下。有这方面的担心，是因为在 openai 的访谈中，貌似像 backtracking 或者 “发现自己前面做错了” 这样的能力都是 “涌现” 的，而不是主动引导的。当然也有可能是力大砖飞，只要收集的数据的多样性足够好，或者迭代式更新模型，就不需要考虑这种问题了。

然后是训练方式。我不太相信模型是通过某种 AlphaZero 一样的纯 self play 进行训练的，因为 Noam Brown 大大在各种场合下强调，一个需要与人互动的模型，需要使用 "human data"。这是 llm 和下围棋的 AI 最大的不同，下围棋这样的 2 人零和博弈，只需要能搜到 minimax 解对应的 nash equilibrium 就可以保证胜率了，不需要考虑对手的思维，而 llm 要保证他说的人能听懂。这也是为啥上一步中我更倾向于 PRM 这种会指导模型中间步骤的方案。我能够想到的最直接地使用 "human data" 的方法就是先用上一步采到的数据sft 一下（也算是 reject sampling 了），之后再进行 RL 训练。

如果采用的是上文这样用单模型生成长输出的方案，我认为 RL 的流程可能一个简单改动版的 PPO 就够了。往哪个方向改我没啥思路，有可能是某种提升采样效率的方案，也有可能是前两天 https://arxiv.org/abs/2410.01679 这样暴增采样计算量的方案。除此之外可能还需要有一些能控制输出长度（即思考长度）的惩罚项或正则，这方面就更是需要具体做实验尝试了。

最后是推理形式。从上文可以看出，我比较相信 o1 是一个模型输出了很长的 CoT，而不是多个 agent 做推理，抑或是用了 MCTS 之类的搜索。我认为大家聊 MCTS 聊得比较多，只是因为 AlphaGo/AlphaZero 太过出名了。如果真的要用 MCTS，我更愿意相信是在造数据的阶段，甚至是 RL 的过程中，是一种为了提高效率而引入的优化手段。

以上的想法来源于各种 blog，视频，以及和各位朋友的讨论（就不一一 at 大家了，各种感谢~）。公开的资源里我比较相信 Noam Brown 的一些访谈，感觉他大有变成 openai 的下一个 Ilya 的势头，所以贴一些他最近的访谈或 talk 在下面，都非常推荐：

• https://www.youtube.com/watch?v=eaAonE58sLU
• https://www.youtube.com/watch?v=EY9iHSe82Hc
• https://www.youtube.com/watch?v=jPluSXJpdrA

以上，期待未来的打脸以及和大家的讨论~

P.S. 基于上述的思路，我还是想办法摘了摘 low hanging fruit，列在下面，有兴趣的朋友可以瞅瞅：

• vllm 之前不怎么支持 serving reward model，我加了一下 Qwen2.5-math-72B-rm 作为样例，已经 merge 了：https://github.com/vllm-project/vllm/pull/8896
• PPO 那边，我在跟着 owner 大大一起优化 OpenRLHF 里的 PPO 性能，OpenRLHF 真的是目前开源界设计得最干净的 llm + rl 的训练框架了~ 有兴趣的朋友可以试试：

GitHub - OpenRLHF/OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework (70B+ PPO Full Tuning & Iterative DPO & LoRA & Mixtral)* 开源社区没有能训 PRM 的框架，所以我在 OpenRLHF 里实现了一下，可以作为参考：https://github.com/OpenRLHF/OpenRLHF/pull/442