一文通透OpenAI o1：从CoT、Self-Correct/STaR、Self-play RL、MCST等技术细节到工程复现

前言

注意，本文自10.12日起，正在每天更新的过程中..

可能是去年写或讲的关于ChatGPT原理的文章和课程，影响力太大了

导致自从OpenAI o1出来后，每过两天，就有朋友问，“校长 o1啥时候出解读或课程”，实在是各个大模型项目上的事情太多，加之近期一直在抠机器人(比如本博客内连发了五篇机器人相关的文章，再不刻意强插一下比如o1，真的很难停下来)

但，今天又有朋友来问，故，准备并行开写o1了

一开始，我主要是看的这几个资料

OpenAI o1的技术blog
张俊林老师的Reverse-o1:OpenAI o1原理逆向工程图解
GitHub上的o1论文列表，比如下面这几篇
1) Let's Verify Step by Step，发布于23年5月
2) Chain-of-Thought Prompting Elicits Reasoning in Large Language Models，思维链开山之作
3) Training Large Language Models for Reasoning through Reverse Curriculum Reinforcement Learning发布于24年2月，阐述R3方法——通过结果监督达到过程监督的效果
4) Training Language Models to Self-Correct via Reinforcement Learning，通过RL训练语言模型学会自我纠正
5) Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters，在加强模型的推理能力方面，延长思考时间比单纯增加模型参数更有效
6) Chain of Thought Empowers Transformers to Solve Inherently Serial Problems，CoT助力模型解决复杂数学问题
公号上的一些o1文章，我正在看
上交团队的一个o1复现

但随着本文越写越深，挖出来的相关资料也越来越多，我会竭尽所能，确保本文具备以下两大特征

尽可能在本文中把o1的所有重要细节都一次性说清楚
如果实在有些细节因为篇幅而没法面面俱到的，我也会尽可能提供相关的论文供相关的读者进一步深入
尽可能在行文过程中，对绝大部分结论不做主观猜测、推测
由于OpenAI实在是不够开放，一个技术报告完全是啥都没说，所以很容易引发各种猜测、推测，虽然我没法完全百分百知晓o1的所有技术细节，但本文会尽可能做到一切结论都有科学而严谨的权威论文做佐证

且在实在必须做一定的推测时，也可能确保是合理、有信服力的猜测，而非主观想象

至于，为何要写OpenAI o1呢，简言之，就是能力强大

其在物理、化学和生物学的具有挑战性的基准任务上的表现与博士生相似
且它在数学和编码方面表现出色。在国际数学奥林匹克 (IMO) 资格考试中，GPT-4o 仅正确解决了 13% 的问题，而推理模型得分为 83%，并在 Codeforces 比赛中达到了第 89 个百分位

第一部分如何让模型学会思考并自我纠正

1.1 如何让模型学会思考

1.1.1 具备问题拆解、自我纠正、多方尝试等综合能力的CoT

说到o1的原理，根据OpenAI o1的技术报告显示

在于大规模强化学习算法教会o1如何在高质量高效数据(data-efficient)的训练过程中利用其CoT(chain of thought)的能力进行有效思考。且他们发现，随着强化学习的更多赋能(训练时间计算)和思考时间的增加（测试时间计算），o1 的性能会持续提高
进一步，类似于人类在回答困难问题之前可能会思考很长时间，o1 在尝试解决问题时使用了“思维链”（chain of thought）策略
Similar to how a human may think for a long time before responding to a difficult question, o1 uses a chain of thought when attempting to solve a problem.
相当于在面对一个不是很轻松的问题时，模型会像人类一样，尽可能思考全面，以想出一个万全之策
那如何思考出来万全之策呢，OpenAI的策略是通过强化学习，让o1 学会优化其思维链并改进所使用的策略
Through reinforcement learning, o1 learns to hone its chain of thought and refine the strategies it uses.

接下来关键点来了
1) 它学会了识别并纠正自己的错误，即It learns to recognize and correct its mistakes.
2) 将复杂的步骤分解成更简单的步骤，即It learns to break down tricky steps into simpler ones.
3) 并在当前方法无效时尝试不同的解决方式，即It learns to try a different approach when the current one isn’t working.

每一点都是职场精英人士的优秀特征啊，总之，整个过程显著提升了模型的推理能力

至于什么是CoT，特援引此文《ChatGPT技术原理解析：从RL之PPO算法、RLHF到GPT4、instructGPT》中“2.4.3 基于思维链(Chain-of-thought)技术下的prompt，是后来OpenAI o1的关键技术之一的”这段内容

为让大语言模型进一步具备解决数学推理问题的能力，22年1月，谷歌大脑团队的Jason Wei、Xuezhi Wang等人提出了最新的Prompting机制——Chain of Thought(简称CoT)

简言之就是给模型推理步骤的prompt，让其学习人类如何一步步思考/推理，从而让模型具备基本的推理能力，最终可以求解一些简单甚至相对复杂的数学推理能力
其对应的论文为，通过此篇论文《Chain-of-Thought Prompting Elicits Reasoning in Large Language Models》
该论文全部作者包括Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, Denny Zhou
PS，上面的最后一位作者Denny zhou是OpenAI推理团队首席科学家，其后来于24年

以下是一个示例(下图左侧为standard prompting，下图右侧为基于Cot的prompt，高亮部分为chain-of-thought)，模型在引入基于Cot技术的prompt的引导下，一步一步算出了正确答案，有没有一种眼前一亮的感觉？相当于模型具备了逻辑推理能力

1.2 如何让模型学会一步一步的自我验证：结合结果监督与过程监督

1.2.1 让我们一步一步验证：Let's Verify Step by Step

首先，我们来回顾下语言生成的整个过程

在每个时间步t，策略语言模型 $\pi_{\theta}^{R L}$ 由参数 $\theta$ 参数化的策略接收一个状态 $s_t$ ，该状态由输入prompt和到目前为止生成的文本组成

然后，策略的行动 $a_{t+1}$ 是在该状态 $s_t$ 的条件下生成下一个token，其概率为 $\pi_{\theta}\left(a_{t+1} \mid s_{t}\right)$
之后，环境返回一个奖励 $r(s_t, a_{t+1})$ ，状态以转移概率 $p(s_{t+1}|s_t, a_{t+1})$ 转移到 $s_{t+1}$
强化学习的目标是找到一个最优策略，以最大化轨迹 $\tau=\left\{s_{0}, a_{1}, \ldots, s_{T}, a_{T}\right\}$ 上的累积奖励（即回报），其中 $s_0$ 是初始状态（即提示）， $T$ 是动作的长度

策略梯度的一般形式给出如下
$\mathbb{E}_{\tau \sim \pi_{\theta}^{\mathrm{RL}}}\left[\sum_{t=1}^{T} \nabla_{\theta} \log \pi_{\theta}^{\mathrm{RL}}\left(a_{t} \mid s_{t-1}\right) R\left(s_{t-1}, a_{t}\right)\right]$
其中 $E_{\tau \sim \pi^{RL}_\theta}$ 表示在从策略 $\pi^{RL}_\theta$ 采样的轨迹分布下的期望
且其中回报 $R\left(s_{t-1}, a_{t}\right)=\sum_{t^{\prime}=t}^{T} \gamma^{t^{\prime}-t+1} r\left(s_{t^{\prime}-1}, a_{t^{\prime}}\right)$ 是从时间步骤 $t$ 开始的奖励的折现和，折现因子为 $\gamma \in[0,1)$
通过这个梯度，可以进行梯度上升来优化模型。如果回报是有利的，动作的选择概率会通过增加来“增强”
如果给定一个数据集 $\mathcal{D}=\left\{\left(s_{0}^{i}, \mathbf{a}^{i}\right)\right\}_{i=1}^{N}$ ，其中包含 N 对输入 $s_0$ 和人类生成的输出序列 $a$ ，其中a= (a1, a2, ..., aT)，整个轨迹为 $\tau=\left\{s_{0}, a_{1}, \ldots, s_{T-1}, a_{T}\right\}$
策略梯度则变为
$\begin{array}{l} \mathbb{E}_{s_{0} \sim \mathcal{D}}\left[\mathbb{E}_{\tau \sim \pi_{\theta}^{\mathrm{RL}}\left(\cdot \mid s_{0}\right)}[ \right. \\ \left.\left.\sum_{t=1}^{T} \nabla_{\theta} \log \pi_{\theta}^{\mathrm{RL}}\left(a_{t} \mid s_{t-1}\right) R\left(s_{t-1}, a_{t}\right)\right]\right] \end{array}$

当然，如果你对上面的这些还不太明白，则可以看下此文《强化学习极简入门：通俗理解MDP、DP MC TC和Q学习、策略梯度、PPO》

其次，近年来，大型语言模型在执行复杂的多步推理方面有了很大提升。然而，即使是最先进的模型仍然经常会产生逻辑错误，很容易出现一步错、则步步错的情况

所以，为了训练更可靠的模型，可以

选择结果监督的奖励模型ORMs
它仅为采样序列的最终结果提供反馈

即奖励模型只给最后一个结果分配奖励分数，而其他token的分数为0
$r_{o}\left(s_{t-1}, a_{t}\right)=\left\{\begin{array}{ll} r f_{o}\left(s_{t-1}, a_{t}\right), & t=T \\ 0, & t \neq T \end{array}\right.$
所以其可以在没有人类参与的情况下进行，针对所有待解决的问题自动得到答案
过程监督的奖励模型PRMs
它为每个中间推理步骤提供反馈，即奖励模型 $r_{mp}(·)$ 被训练用于为每个中间推理步骤分配一个奖励分数
$r_{p}\left(s_{t-1}, a_{t}\right)=\left\{\begin{array}{ll} r m_{p}\left(s_{t-1}, a_{t}\right), & t \in \mathcal{T}^{\text {Delimiter }} \\ 0, & t \notin \mathcal{T}^{\text {Delimiter }} \end{array}\right.$
如此则能够准确指出错误发生的位置
但该PRM则需要依赖于人工标注以提供过程监督，比如告诉模型每一步的正确与否，比如下图

ORM 在最后一个标记上的预测作为解决方案的整体得分，这点大家可能并不陌生了，因为instructGPT便是这么做的(详见这篇ChatGPT原理笔记)

而PRM好像很少见，有人可能会说，在这篇ChatGPT原理笔记中，在讲到DSC的实现时，不讲到了虽然奖励模型RM的输出结果reward侧重结果监督奖励，但critic输出的价值估计value则侧重衡量过程么

举个例⼦，假定在经验数据池里，有seq=[ i am learning machine learning with julyedu ]，其

KL系数，即β
kl_ctl: 0.02
旧策略下生成句子中不同词的对数概率
log_probs: [-12.0938, -20.1250, -1.3477, -1.8945, -0.9966, -0.5142, -1.4902]
SFT策略下生成句子中不同词的对数概率
ref_log_probs: [-12.1172, -20.1406, -1.3447, -2.1699, -1.2002, -0.7905, -1.6611]
KL = - kl_ctl * (log_probs - ref_log_probs)
kl_divergence_estimate: [-4.6873e-04, -3.1257e-04, 5.8591e-05, -5.5084e-03, -4.0741e-03, -5.5275e-03, -3.4180e-03]
rewards = [0, 0, 0, 0, 0, 0, reward_score]
token_id为7的token作为该对话的最后1个有效token，其对应的实际经验奖励reward_score将被⽤于表示整个对话的reward
即其奖励reward_score只会是1个标量
reward_clip(reward_score)
reward_clip: -0.0826
KL[最末有效token的idx] + reward_clip → KL_rewards
$r_{KL} = r(x,y) - \beta \log \frac{\pi_{old}^{RL}(y|x)}{\pi^{SFT}(y|x)}$
KL_rewards = [-4.6873e-04, -3.1257e-04, 5.8591e-05, -5.5084e-03, -4.0741e-03, -5.5275e-03, -8.5999e-02]

然后才有后来的优势值计算结果

$\begin{array}{c} \hat{A}_{t}=\delta_{t}+(\gamma \lambda) \delta_{t+1}+(\gamma \lambda)^{2} \delta_{t+2}+\cdots+(\gamma \lambda)^{T-t+1} \delta_{T-1} \\ \text { where } \delta_{t}=r_{K L, t}+\gamma \cdot V_{\text {old }}\left(s_{t+1}\right)-V_{\text {old }}\left(s_{t}\right) \end{array}$

以及

$R_{t}=\hat{A}_{t}+V_{t}$

最后，再有critic输出的价值估计Value对returns的拟合

$V_{clip} = clip(V_{new}, V_{old}-\phi, V_{old}+\phi)$

$vf\_loss = \frac{1}{2} \cdot E_{\tau \sim \pi_{old}^{RL}} E_{s_t \sim {\tau}} [\max((V_{new}(s_t)-R_t)^2, (V_{clip}(s_t)-R_t)^2)]$

其实不然，原因在于

prm orm本质都是奖励模型给出实际得分评判，他们本身如果需要都可以引入“critic”
prm更倾向于表达“走到这一步是好还是坏”，只看这时候如何
critic本质还是价值估计给出的都是价值预测
而且更多时候critic会更倾向于表达“给定状态和动作，最后的奖励可能是多少”，侧重关注对未来的影响

进一步，23年5月，OpenAI的这篇论文《Let's Verify Step by Step》发现过程监督在训练模型解决具有挑战性的MATH数据集的问题时，明显优于结果监督，比如过程监督模型解决了MATH测试集代表性子集中的78%问题

那如何训练PRM呢，训练模型的第一步当然是收集数据了

为了收集过程监督数据，首先向人工数据标注员提供由模型回答数学问题的逐步解决方案

在每次迭代中，他们为每个问题生成N个解决方案，，且尝试在问题级别（每个问题K个解决方案）或在整个数据集上（总共K个解决方案，不均匀分布在各个问题中）应用这种前K过滤
最终，将收集的步骤级标签的整个数据集称为PRM800K。PRM800K训练集包含针对12K个问题的75K个解决方案的800K个步骤级标签

他们的任务是为解决方案中的每一步分配一个正面、负面或中性的标签，如下图所示

正面标签表示该步骤是正确且合理的
负面标签表示该步骤不正确或不合理
中性标签表示存在歧义。在实践中，如果某个步骤存在微妙的误导性，或者是一个技术上仍然有效但不佳的建议，它可能会被标记为中性

1.2.2 R3：通过结果监督的手段达到过程监督的效果

考虑到结果监督为最终结果提供稀疏的奖励，但不识别错误位置(没法明确哪些动作导致了成功或失败)，而过程监督则提供逐步的奖励，但需要大量的人工标注

24年5月，来自复旦的研究者们(黄萱菁团队)，通过此篇论文《Training Large Language Models for Reasoning through Reverse Curriculum Reinforcement Learning》提出了一种名为R3的创新方法：通过Curriculum RL进行逆向推理学习

该方法仅使用结果监督来实现大型语言模型的过程监督优势。在将RL应用于复杂推理时，核心挑战在于识别出能够获得正向奖励的动作序列，并提供适当的监督以进行优化

对于一个多跳推理问题，可以通过不同的推理路径得出一个标准答案

假设可以访问至少一个示例，即一个通往标准答案的正确推理路径，就像在监督微调中一样。当模型从初始状态 $s_0$ 开始探索时，可能会遇到难以获得正奖励的问题

受以往在带示范的强化学习领域研究的启发，R3作者将给定示范的中间状态集定义为 $S_{\text{Inter}}\subset S$ ，并让策略 $\pi_{\theta}$ 从接近目标状态的中间状态 $s_k \in S_{\text{Inter}}$ 开始探索： $\pi_{\theta}(a_{k+1:T}|s_k)$ ，其中 $a_{k+1:T} = (a_{k+1}, ..., a_T)$

基于结果的奖励函数为最终结果提供反馈，作为在 $s_k$ 之后采取的动作的监督信号
在这种策略中，示范中 $s_k$ 之前的轨迹（即 $\{s_0, a_1, s_1, ..., a_k\}$ 可以作为一种指导，帮助模型更容易获得正向奖励，并避免陷入无方向、低效的探索过程，如下图所示
一旦策略学会了从接近目标的选定状态开始实现目标，它就可以将训练扩展到更远的状态（例如， $s_{k-1}$ ），从而利用已经获得的知识

在每个阶段，模型面临一个容易的探索问题，因为它已经学会解决大部分剩余的部分，因此很可能成功。通过这种方式，可以形成一个逐步增加探索难度的curriculum

这使得能够为模型提供大致的逐步监督信号，现在，策略梯度可以表示为：
$\begin{array}{l} \mathbb{E}_{s_{k} \sim \mathcal{S}^{\text {Inter }}} {\left[\mathbb{E}_{\tau \sim \pi_{\theta}^{\mathrm{RL}}\left(\cdot \mid s_{k}\right)}[ \right.} \\ \left.\left.\sum_{t=k+1}^{T} \nabla_{\theta} \log \pi_{\theta}^{\mathrm{RL}}\left(a_{t} \mid s_{t-1}\right) R_{o}\left(s_{t-1}, a_{t}\right)\right]\right] \end{array}$
其中， $S_{\text{Inter}}$ 指的是从数据集 $D$ 中采样的一个演示的中间状态集； $k$ 从 $T-1$ 开始并逐步回退到 0。在最后一步，模型开始从初始状态 $s_0$ 展开，这相当于原始的结果监督强化学习
接下来，便可以通过PPO算法迭代策略，且设计奖励函数
$\begin{array}{l} r_{\text {final }}\left(s_{t-1}, a_{t}\right)=r_{o}\left(s_{t-1}, a_{t}\right) \\ \quad-\beta \mathrm{KL}\left(\pi_{\theta}^{R L}\left(\cdot \mid s_{t-1}\right), \pi_{\theta}^{I n i t}\left(\cdot \mid s_{t-1}\right)\right) \end{array}$
且和此文中介绍的DSC一样，通过广义优势估计GAE来计算优势

1.3 验证出问题之后，如何让模型学会自我纠正

1.1节开头说过，OpenAI o1学会了识别并纠正自己的错误，那它怎么做到的呢

之前对自我纠正LLM的尝试要么依赖于提示工程，要么依赖于专门为自我纠正而微调模型

前者通常无法进行有意义的内在自我纠正
而微调方法在推理时需要运行多个模型，例如一个独立的改进模型，或依赖“教师”监督来指导自我纠正的过程
使用独立的教师监督模型时，自我纠正不一定优于并行的独立尝试

Google开发了一种不需要这些要求的有效自我纠正方法，称为通过强化学习进行自我纠正（Self-Correction via Reinforcement Learning，简称SCoRe），其对应的论文为24年9月下旬发布的：《Training Language Models to Self-Correct via Reinforcement Learning》

其仅训练一个模型，该模型既可以对问题产生响应，也可以在没有任何预测反馈的情况下纠正错误简言之，SCoRe是一种多轮次强化学习方法，用于教大型语言模型如何纠正自己的错误
且与其最接近的可能是「其对应的论文为Recursive introspection: Teaching language model agents how to self-improve」，他们利用一种类似STaR的迭代方法进行自我纠正。尽管该研究主要使用了预言式教师监督，但其仅通过五轮自我纠正训练获得的初步结果显示出微小的改进，这与我们在STaR中看到的结果一致

SCoRe采用两阶段训练：

在第一阶段，通过训练纠正第二次尝试的响应，同时将第一次尝试的分布约束在接近基础模型，以产生不易行为崩溃的初始化
然后在第二阶段对两次尝试进行训练以最大化奖励。关键是，多轮RL的第二阶段采用奖励塑造项，奖励“进步”以实现自我纠正，而不是最终响应的正确性

在效果上，相对于基础的Gemini模型，SCoRe的方法在MATH（Hendrycks等人，2021）的推理问题上实现了绝对15.6%的自我纠正提升，并且实现了绝对9.1%的提升

具体来说，如下图所示，表示的是自我纠正的问题设定「SCoRe不仅训练模型生成最佳可能的响应，还旨在训练模型在最后一次尝试中生成最佳最终响应。在第二轮中，可能会以指令的形式提供额外输入，要求模型自我纠正或由模型生成」

给定一个包含问题 $\boldsymbol{x}_{i}$ 和响应 $y_{i}^{*}$ 的数据集 $\mathcal{D}=\left\{\left(\boldsymbol{x}_{i}, \boldsymbol{y}_{i}^{*}\right)\right\}_{i=1}^{N}$
然后训练一个LLM策略 $\pi_{\theta}\left(\cdot \mid\left[x, \hat{y}_{1: l}, p_{1: l}\right]\right)$ ，该策略在给定问题 $x$ 、之前模型对该问题的 $l$ 次尝试 $\hat{\boldsymbol{y}}_{1: l}$ 以及辅助指令 $p_{1: l}$ （例如，寻找错误并改进响应的指令）的情况下，尽可能正确地解决问题 $x$ 。这种形式化类似于Qu等人（2024）中的多轮MDP
我们还假设可以访问一个oracle奖励̂ $\hat{r}\left(y, y^{*}\right)$ ，例如一个答案检查器（Uesato等人，2022），通过将响应 $y$ 与oracle响应 $y^{*}$ 进行比较来评估其正确性
至关重要的是，不假设在测试时能访问这个oracle；相反，模型必须推断是否存在错误，并在必要时进行纠正，这在例如数学推理问题中经常发生