重现 OpenAI o1 的技术路线

OpenAI o1发布后，其强大的推理能力远超早期的大语言模型（LLM），达到了媲美博士级专业知识的性能水平。

目前，有两种复现o1的范式：

基于知识蒸馏：这是一种捷径方法，可以提取o1的数据并微调LLM（如Llama 3.2、Qwen2等）以模仿o1的推理风格。但是，这些方法受到教师模型的能力上限的限制。
基于强化学习：这是一种更正统的方法，循序渐进，专注于策略初始化、奖励设计、搜索和学习四个关键组成部分，以达到复现o1的目标，如图1所示。

虽然已经有一些关于复制o1的信息，但本文主要基于一个严肃的调查，介绍了基于强化学习的o1复制路线图。路线图如图2所示。

图2：使用强化学习重现o1的路线图和实现

生动的描述

想象一下，你正在准备一支探险队来寻找传说中的宝藏。此旅程遵循四个关键步骤：

策略初始化 —— 装备和培训团队

出发前，您必须装备您的探险家并对其进行训练。您从经验丰富的导游（预先训练的LLM）开始，教他们如何阅读地图和使用指南针（语言理解和推理）。然后，通过向他们展示如何解释复杂的线索并遵循指示（指令微调）来提高他们的技能，确保他们能够应对未来的挑战。

reward —— 设置激励和反馈

为了让团队走上正轨，您引入了一个奖励系统：

结果奖励：只有当他们找到宝藏时才会获得奖励（正确的最终答案）。
流程奖励：他们发现的每条线索或纠正他们采取的步骤（逐步改进）即可获得奖励。

这确保他们不仅以目标为目标，而且在此过程中获得指导。

搜索 —— 探索未知

团队不能仅仅依靠猜测，他们必须进行战略性探索。有些提前规划多条路线（Tree Search），而另一些则逐步调整（Sequential Revisions）。他们的决策通过自我评估（内部反馈）和环境线索（外部反馈）得到改善，帮助他们更有效地驾驭。

学习 —— 从经验中成长

每次探险后，团队都会回顾他们的旅程，了解哪些有效，哪些无效。通过完善他们的策略，他们成为更好的探索者。同样，该模型使用搜索生成的数据来优化其策略，从而随着时间的推移提高性能。

简而言之，这个路线图就像训练一个智能探险队——每个步骤都协同工作，帮助他们找到最佳解决方案，就像它使o1这样的模型能够发展出强大的推理能力一样。

策略初始化

图3：策略初始化的过程，包括预训练、指令微调和类似人类推理行为的注入。这些行为可以通过SFT学习，也可以通过提示触发

如图3所示，策略初始化涉及三个关键阶段：

预训练：该模型通过在大规模文本语料库上进行自我监督学习来学习基础语言理解。
指令微调：此步骤将模型从简单的标记预测转变为生成符合人类期望的响应。
开发类似人类的推理行为：虽然指令调整模型展示了广泛的任务能力和对用户意图的理解，但像o1这样的模型需要更复杂的推理技能才能充分发挥其潜力。如图4所示，确定了6种类似人类的推理行为，以帮助o1更有效地探索解决方案空间。

图4：类人推理行为的示例

类似o1的模型的策略初始化面临两个主要挑战：

平衡效率和探索：在高效采样和探索各种可能性之间取得适当的平衡是很棘手的。虽然人工演示提供了有价值的指导，但它们有时会限制发现更好的策略，如AlphaGo和AlphaGo Zero之间的区别。
跨领域泛化：模型不仅在数学和编码等特定领域表现出色，还需要适用于各种任务（包括安全检查）的多功能推理技能。构建这种广泛适用的功能既必要又复杂。
奖励设计
接下来，我们进入奖励设计阶段，旨在为搜索和学习提供指导信号。

LLM的奖励分为两大类：

结果奖励：根据预定义的目标对输出进行评分，例如检查数学解决方案是否正确。它们易于使用，但会忽略中间步骤中的错误，这可能会强化有缺陷的推理并阻碍逐步学习。尽管存在这些限制，但它们的简单性使其得到广泛使用。
流程奖励：另一方面，流程奖励为最终和中间步骤提供反馈，提供更详细的指导。但是，它们更难设计和实现。

生成的模型通常称为结果奖励模型（ORM）和过程奖励模型（PRM）。

图5：结果奖励（左）与流程奖励（右）。由于倒数第二行，整个解决方案不正确。因此，对于结果奖励，整个解决方案被标记为不正确，而对于过程奖励，前三个步骤被标记为正确，最后两个步骤被标记为错误

从本质上讲，结果奖励可以看作是过程奖励的简化版本，其中没有中间反馈。

模型的奖励设计方法可以使用来自真实或模拟环境的直接信号、AI判断或数据驱动方法。真实环境提供直接的回报，而模拟环境降低成本，但风险分配发生变化。RLHF等基于偏好的方法使行为与人类价值观保持一致，而逆向强化学习（IRL）等专家数据方法从专家轨迹中恢复奖励函数，但更复杂。奖励塑造可以优化弱信号以改善学习，但需要仔细设计以避免负面影响。

o1的奖励设计

o1的推理能力表明，其奖励模型包含多种奖励设计方法。

对于数学和编程任务，它可能会使用流程奖励来跟踪中间步骤，奖励塑造将结果奖励转换为流程奖励。
当环境中没有可用的奖励信号时，o1可能依赖于首选项数据和专家数据。
其强大的奖励模型跨不同领域进行了微调，并使用Ground Truth和解决方案对进行调整。
此外，它更有可能通过使用LLM而不是价值头来预测奖励。

复制o1的奖励设计面临三个主要挑战。

首先，当模型探索训练数据之外时，会发生分布变化，虽然扩展数据和迭代训练有帮助，但它们并不能完全解决问题。
其次，由于作空间巨大，并且需要平衡Token级、Step级或解决方案级的反馈，因此很难为语言任务设计细粒度的奖励。
第三，为复杂任务选择正确的数据很棘手，因为基于偏好的反馈有时会降低策略模型的性能，并且随着任务复杂性的增加，评估奖励的有效性变得更加困难。
搜索
搜索是指根据某些指导（例如奖励或启发式规则）通过多次尝试或战略探索来找到正确解决方案的过程。

对于像o1这样的模型，搜索对于在训练和推理期间解决复杂的推理任务至关重要。

内部指导依赖于模型自身的状态或评估功能，例如在贪婪或光束搜索中使用标记概率。另一方面，外部指导使用与环境或任务相关的信号来指导搜索过程。

图6：搜索的两个关键方面是解决方案选择的指导信号和获取候选解决方案的搜索策略。搜索策略用于获取候选解决方案或行动，而指导信号用于对候选解决方案或行动进行选择。S1表示步骤1

搜索策略可以分为两类：

Tree search：同时探索多个解决方案，提供全局视图和更广泛的覆盖范围。
Sequential revisions：根据之前的尝试优化每次尝试，作为一种更有效的局部搜索。
o1的搜索
在训练期间，o1可能使用树搜索技术，如BoN和其他树搜索算法，同时主要依赖外部指导。这种方法允许模型逐步提高其推理能力，因为树搜索可以并行生成许多候选解决方案，有效地创建丰富的高质量训练数据。
在测试时，o1可能依赖于带有内部指导的顺序修订，使用反思来优化和纠正其搜索。

当然，复现o1的搜索存在一些挑战：

克服反向缩放：减少测试时的搜索可以限制反向缩放，但也会缩小搜索范围。更好的方法是通过增加奖励模型的规模和训练数据来提高其泛化能力。
避免对简单任务过度思考：像“1+1=？”这样的简单问题不需要复杂的推理。通过奖励塑造对思维链（CoT）添加长度惩罚可以限制不必要的步骤，节省资源并减少延迟。
平衡树搜索和顺序修订：结合树搜索和顺序修订可以提高性能，但在预算固定的情况下，资源分配具有挑战性。经验缩放定律可以指导平衡。
提高搜索效率：搜索效率受到内存速度和某些树搜索算法（如MCTS）缺乏并行性的限制。解决方案包括工程优化，如KV缓存共享，以及算法改进，如KV缓存压缩和推测采样。
学习
策略初始化使用人类专家数据，但强化学习因其超越人类表现的能力而不可或缺。与专家数据不同，它通过试错学习，就像AlphaGo发现著名的“第37手”等开创性棋步一样。

策略可以通过策略梯度方法或行为克隆（监督学习）来改进。与策略初始化不同，这里的行为克隆是从搜索过程中学习，而不是从人类专家那里学习。

图7：4种学习方法在数据使用方面的差异。策略梯度使用包含所有搜索解决方案的Dsearch，而行为克隆仅使用奖励最高的解决方案Dexpert。奖励最高的解决方案以绿色突出显示，其他解决方案以红色标记

策略梯度方法有多种形式，包括REINFORCE、Actor-Critic、近端策略优化算法（PPO）和直接偏好优化（DPO）。图8比较了REINFORCE、PPO、DPO和行为克隆。

图8：可以基于梯度方差、内存成本和数据利用率三个因素对REINFORCE、PPO、DPO和行为克隆这4种学习方法进行比较。对于内存成本，✓表示需要奖励/价值/参考模型，意味着高内存成本。至于数据利用率，✓表示使用重放缓冲区/负解决方案进行学习，这表明数据利用率高

在内存成本方面，PPO内存密集，需要奖励函数、价值函数和参考策略。DPO更简单高效，但需要偏好数据。行为克隆不需要参考策略，内存效率最高。

在数据利用率方面，PPO和DPO使用搜索中的所有状态 - 行动对（Dsearch），包括负奖励。行为克隆仅使用搜索中具有高奖励的状态 - 行动对子集（Dexpert），限制了其范围。因此，PPO和DPO比行为克隆更好地利用数据，因为它们可以从正面和负面行动中学习以改进策略。

o1的学习

总体而言，o1可能会在高效热身阶段从行为克隆开始，当改进趋于平稳时过渡到PPO或DPO。这种组合利用了行为克隆早期的效率，同时利用PPO或DPO从更广泛的数据中学习以进行进一步优化的能力。这种方法与Llama2和Llama3中使用的训练后策略一致。

复现o1的学习也存在挑战：

训练效率：受MCTS等缓慢的搜索过程阻碍，可以通过优化算法或重用先前迭代的数据来提高。
学习强大的问题生成器：具有挑战性，因为问题必须相关、可解并能够促进学习。LLM可以缓解这个问题。
离策略学习：受分布转移困扰，可以通过限制搜索规模、从当前策略采样或使用重要性采样等技术来减少。另一种方法是应用行为克隆将离策略学习转换为在策略学习。
开源o1项目
虽然o1没有发表论文，但有几个开源实现可用。这些可以看作是这里讨论的框架的组件或特定案例。

图9：不同开源o1项目的比较

其中，我认为以下三个复现特别具有代表性，值得深入研究：

o1 Replication Journey — Part 1: Expert Iteration：通过波束搜索生成的树数据被遍历，特定节点由GPT-4优化，然后用于监督微调。这个被称为专家迭代的过程，通过GPT-4的优化突出了模型的自我反思能力。
o1 Replication Journey — Part 2: Distillation and CoT Recovery：它采用了不同的方法，专注于提炼o1-mini。虽然o1-mini仅输出思维链（CoT）摘要，但第2部分试图通过提示o1-mini来扩充摘要，以恢复隐藏的CoT。
Marco-o1：就像一支登山探险队，旨在找到登顶的最佳路径（解决复杂问题）。通过路线规划器（思维链（CoT）），探险队将大目标分解为更小的、可管理的步骤，逐步实现它们。在无人机（蒙特卡洛树搜索（MCTS））的帮助下，他们模拟并探索多条路线，使用置信度分数选择最佳路径。此外，团队依靠指南针（反思机制）来检查和纠正方向。

重现 OpenAI o1 的技术路线