首先，我们来回顾o1模型的整体技术架构。基于目前已公开的论文与代码，我们提炼出o1可能的实现路径，如下图所示：

关于这张架构图的具体含义，这里不再赘述——在之前的《框架篇》中已详细讲解。简单来说，它高度概括了o1技术报告中指出的核心思想：将更多计算资源分配至推理阶段，从而显著提升模型的逻辑推理能力。从本文开始，我们将用具体算法逐步填充这张架构图的细节。

今天我们将重点展开框架图中的Inference部分（黄色区块）。从图中可以看出，Inference部分通常承担两项职责：

• 职责1：直接优化推理过程。具体方法包括：
• PRM + 搜索方法。PRM是一种专门评估模型中间步骤质量的奖励模型，而非仅评估最终答案。前面《框架篇》已给出示例，此处不再重复。
• MCTS（蒙特卡洛树搜索）。该方法曾在AlphaGo中应用，通过自我对弈（self-play）寻找一条从原始问题、中间步骤到最终答案的最优路径。广义上讲，PRM+搜索本质上也是一种MCTS风格的搜索——只不过MCTS依靠“探索”来估计节点的奖励，而训练好的PRM直接替代了这一探索与评估过程。如果觉得抽象，别担心，MCTS正是本文要详细阐述的内容，即将展示其实现细节。
• 职责2：在训练后（post-training）阶段筛选高质量数据。根据现有开源工作的总结，提升模型推理能力有一个关键原则：尽量少用人标注，多借助已有模型（base generator）自身的能力自动化生成训练数据，然后利用这些数据通过SFT或强化学习等post-training方法进一步改善。为了保证自动生成数据的质量，可以引入Inference模块来辅助搜索出高质量数据。

因此，Inference模块可以视为o1实现中的一块积木。理解这块积木的目的与可能的实现方式后，你可以按需灵活地将其组装到o1的任何环节。网上关于o1的资料中经常出现“MCTS、self-play”等关键词，它们正是这块黄色积木的一种实现。需要指出的是，o1很可能不纯粹采用推理优化路线（即图中的framework1），更可能走post-training + inference路线（framework3），因为技术报告中提到o1把计算资源也投入了强化学习阶段。但无论如何，了解这块积木的实现都是必要的。

今天这篇文章，我们将以微软今年开源的rStar为例，从源码出发，详细解析MCTS技术如何在NLP逻辑推理任务中落地（毕竟我们对MCTS的主要印象来自AlphaGo，自然好奇它在自然语言场景下如何运作，尤其是搜索空间如何定义）。阅读本文不需要任何MCTS先验知识，我们将逐步讲解。

一、为什么选择rStar

rStar的目标同样是增强模型的逻辑推理能力，但它采用了framework1——纯靠推理阶段的搜索优化实现目标，具体方法是MCTS，而非PRM+搜索。rStar如此选择的原因很现实：

• 基座生成器（Base generator）是小模型（SLM）。rStar针对小模型场景——小模型本身能力有限，无法依靠预训练阶段的能力生成高质量训练数据，因此post-training自产自销的路线在小模型上走不通。而且在多数业务场景下，我们也往往缺乏充足的训练资源。
• PRM的训练成本高昂。如果必须人工标注，那么这笔费用大概率会投入PRM训练（《框架篇》中介绍过OpenAI训练PRM的方式）。对于资源有限的团队和背负落地绩效压力的管理者来说，时间和金钱能省则省。

正因为rStar走纯推理路线，它特别适合从“一块积木”的视角来理解框架图中的黄色区块。此外，它利好小模型的场景也更贴近资源有限的我们。最关键的是，rStar的代码完全开源，让我们可以深入探究所有细节，减少凭空想象的成分（rStar的论文写得较为精简，缺失许多细节，阅读代码具有一定难度）。

二、按照人的思考方式构造一棵搜索树

正式进入MCTS之前，我们先抛开术语，只思考一个问题：解决推理问题时，人类通常会采用哪些思考方式？

假设我们有一个简单问题：

user_question:
If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?

针对此问题，我们可能采用以下几种思考模式（所有模式以字母A开头，代表Action）：

2.1 A1（提出一步推理）

我们会将推理过程分解，每次只提出一个步骤，直至得出最终答案。我们将这种思考方式称为A1。例如：

A1（propose a one-step-thought）

### Instruction:
If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?

### Response:
Let's think step by step.
Step 1: Start with the number of cars that are already in the parking lot, which is 3 cars.
Step 2: Add the number of cars that arrive, which is 2 cars.
Step 3: Add the numbers together. there are 3 cars + 2 cars = 5 cars in the parking lot.
Step 4: The answer is 5.

注意上述格式：

• 总是以“Let's think step by step.”开头
• 每个step包含“推理文字 + 该步得到的数值答案”。例如step1中，在推理文字后可提取出“3”这个数字。
• 最后一步以“The answer is”开头，表示给出最终答案。

2.2 A2（提出剩余推理步骤一步到位）

对于简单问题，我们可能不逐步拆解，而是直接通过简洁推理给出完整答案：

### Instruction:
If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?

### Response:
Let's think step by step. There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5. The answer is: 5.

2.3 A3（提出下一个子问题并回答）

有时我们会将原始问题拆解为若干子问题，依次回答每个子问题，最终得到答案。例如：

Question 1: If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?
Question 1.1: How many cars are there in the park before?
Answer 1.1: There are 3 cars in the park before.
Question 1.2: How many cars arrive then?
Answer 1.2: 2 more cars arrive.
Question 1.3: Now we can answer the question: how many cars are in the parking lot?
Answer 1.3: There are 3 + 2 = 5 cars in the parking lot now. The answer is 5.

其中Question1.3属于终结子问题，因为回答它等同于直接回答最终答案。这种拆解方式更适合解决困难问题，此处用简单例子仅展示形式。

2.4 A4（重新回答子问题）

A4通常与A3搭配使用。例如对A3中的Question1.1，如果对Answer1.1不确定，可以重新思考一次其答案。此时可能采用A2的方式做些简单推理，重新得出Answer1.1——相当于将原答案替换为A2例子的输出结果。此处不再举例。

2.5 A5（重述问题/子问题）

有时我们会在原始问题的大量描述中先提取关键信息，例如“condition1... condition2...”等，然后基于这些关键信息回答。这种方式相当于先重述问题，再做回答。A5即执行此操作，比较直观，也不再举例。

2.6 整合：构造一棵搜索树

总结起来，根据人类思维方式，我们归纳出5种动作：

• A1（提出一步推理）：逐步推理，每步有中间答案，最后一步得到最终答案
• A2（提出剩余推理步骤）：一次性推理，直接得到最终答案
• A3（提出下一个子问题并回答）：将问题拆解为若干子问题分别回答，最后一个子问题的答案即最终答案（与A1类似，但采用子问-子答形式）
• A4（重新回答子问题）：对A3中某个子问题的答案不放心，重新用A2模板回答
• A5（重述问题/子问题）：将原始问题或子问题重新表述，例如提取关键条件，去除冗长描述

在代码中，我们会按照2.1~2.5的示例构造相应的prompt，指示模型执行不同动作。下面给出A1的prompt示例，更多细节可查看源码中的 rStar/prompts 部分：

当人类解决问题时，会根据难度决定采用哪种模式。但模型在搜索时很难预知难度，因此我们希望模型尽可能将所有可能的动作都探索一遍。

接下来，我们配合rStar的源码，一起看看这棵搜索树的具体形态（这里不用论文中的图，因为缺失太多细节，我们直接根据源码重新绘制）：

基本信息：

• 根结点是原始问题。
• 方形node表示叶结点（leaf node），例如图中的cot结点（A2）。注意，A1的最后一个step、A3的最后一个子问题-子回答也可成为leaf node。
• 虚线表示可选探索（取决于脚本配置），实线表示必选探索。

从根结点（第0层）出发：

• A1（一步推理）、A5（重述）为可选；A2（remaining steps，源码中称为cot）和A3（next subqs and subas）为必选。经过A5后相当于从新问题出发，因此A5之后的分支规律与根结点一致，图中不再重复画出。下面重点看第1层。
• 第1层的A1部分（蓝色块）：对原始问题并行采样若干个step1（代码默认3个），每个step1都是一个A1类型结点。并行采样用于拓宽搜索空间——因为对同一问题，不同人产出的第一步可能不同。
• 第1层的A2部分（绿色块）：对原始问题做简单推理后一次性生成最终结果。同样并行采样（默认32个）。由于已生成最终答案，A2结点均为leaf node，不再向下扩展。
• 第1层的A3部分（红色块）：对原始问题产出下一个子问题及其对应回答。同样并行采样（默认3个），例如对原始问题并行采样三个Question1.1和Answer1.1。
• 第1层的A5块（灰色块）：对原始问题进行重述，前面已提及。
• 根结点下为何没有A4？因为A4只出现在A3之后，因此根结点下不会出现。

以第1层为例，看不同类型结点可产生哪些子节点，以及可能的leaf node类型。理解了第1层，其余层可类推。

• A1（一步推理）类型结点：
  • 只能产生A1（可选）和A2（必选）两种子节点。基于当前step，必须产出一个一步到位的推理结果，选择性产出一个step by step的结果。
  • A1结点的leaf node可以是A2（直接得出最终结果）或A1自身（最后一个step的结果）。
  • A1结点的值包含“该步推理文本 + 该步阶段性答案”。
• A2（剩余推理）类型结点：如前所述，一步到位产出答案，是leaf node，不再有子节点。
• A3（子问题与回答）类型结点：
  • 可以产生A1（可选）、A2（必选）、A3（必选）、A4（必选）类型子节点。
  • 图中第1层只画出了并行采样出的第1个A3结点的子节点情况，其余类推，用省略号表示。
  • A3结点的leaf node可以是A1（最后一个step）、A2（一步到位）、或A3自身（最后一个子问题-子回答）。
• A4（重新回答子问题）类型结点：
  • 可以产生A1（可选）、A2（必选）、A3（必选）类型子节点。
  • A4本质上是重新修改父结点的子回答，其leaf node类型同上。

到此，我们已经解决了：

• 按照人类思维模式，定义了搜索动作空间（Action）。
• 根据动作之间的前后依赖关系，确定了完整的搜索树结构。

但还有两个关键问题未解决：

• 这棵树是我们站在上帝视角穷尽所有可能后构造出来的，模型该如何按照我们的想法实际构造出这棵树？
• 有了这棵树后，如何从根结点出发，选择一条最佳路径并产出最终答案？

接下来，MCTS算法正好能回答这两个问题。

三、使用MCTS搜索最佳推理路径

3.1 用rollout构造搜索树

现在，模型从原始问题出发开始构造搜索树。先看第一步：

针对根结点：

• 执行select步骤：选中根结点（select的细节稍后讲解，这里先关注它选中了根结点）。
• 执行expand步骤：按照第二节中各结点间的依赖规则，为选中的结点（这里是根结点）创建所有可能的子节点。为简便，图中略去了“并行采样”过程，但实际操作中依然并行采样。
• 执行simulate步骤：随机选择一个子节点，重复“expand-随机”步骤，直到遇到leaf node或达到最大搜索深度为止。注意，只有两种node可成为leaf node（这与第二节中理想情况略有差异）：
  • Terminal A3 node：当某个子问题是最后一个子问题（子问题中包含原始问题，或以“Now we can answer the question”开头，见2.3示例。这是通过prompt控制模型生成实现的）。
  • Terminal A2 node：本身就是一步推理产出最终结果的node。
• 执行backprop步骤：计算leaf node的reward，同时将本次搜索路径上所有node.reward += leaf_node.reward，node.freq += 1（freq表示被访问次数）。leaf node的reward如何计算？
  • Terminal A3 node reward：对最后一个子问题并行采样若干个子回答。假设采样n个回答，其中指向答案x、y、z的条数分别是a、b、c（n=a+b+c），那么答案x的占比为a/n，以此类推。选择占比最大的答案作为最终答案，并将该占比作为reward。
  • Terminal A2 node reward：同样，在所有并行采样结果中计算答案占比，方法同上。

这样一轮 select + expand + simulate + backprop 就称为1次rollout。不难发现，经过1次rollout后，我们构造出了一部分搜索树（这里先只谈构造，不谈搜索）。

接下来执行第2轮rollout：

• select步骤：从第1轮backprop后构造的搜索树开始，再次从根结点向下选择。走到第1层，发现有3个子节点都没被探索过，随机选择一个（比如A5），这个子节点将被用作expand。这里需要更系统地解释select步骤：
• 每次都从根结点出发，向下逐层探索，直到找到一个未被探索过的结点为止。
• 如果某一层所有结点都已经被探索过，那么计算每个结点的UCT值（3.2节会详细解释，UCT值用于衡量结点的探索价值，由node的reward、freq和探索权重超参C决定）。选择UCT值最大的结点继续向下搜索。
• 因此，select步骤的目的是找到一条未被探索或具有最高探索价值的路径，以便后续沿着它继续扩展，生成更好的搜索树。
• 第2轮的expand、simulate、backprop步骤与第一轮类似，不再赘述。

需要额外注意：生成搜索树的每一层时，都需要将前面所有层的推理结果作为上文，传递给模型。具体构造上文的细节可以阅读源码，这里不再展开。

好，到现在我们已经理清了单轮rollout的概念。循环多次rollout（代码默认16次）后，就能得到一棵相对完整的搜索树。接下来就可以基于这棵树找到一条最佳的推理路径了。但在介绍具体搜索方法之前，先看看如何计算一个结点的UCT值（UCT越大，该结点越值得探索）。

3.2 计算结点的UCT值

一个结点的UCT计算方法如下：

• Q：到本轮rollout为止，该结点的累积reward
• N：到本轮rollout为止，该结点的累积被访问次数
• N_parent：到本轮rollout为止，该结点的父结点累积被访问次数
• c：探索权重。c越大，越侧重探索（explore）而非利用（exploit）

什么样的结点更值得被访问？直觉上，平均reward（Q/N）越大的结点表现越好，应优先选择——这部分叫利用（exploit）。但如果一个结点被访问次数很少（例如它的父结点被访问了几百次，它才被访问几次），说明这条路径可能还有更多“宝藏”没被发现，应该给它更多机会——这部分叫探索（explore），由 c*sqrt(N_parent/N) 来体现。

人为设置的探索权重c控制着explore和exploit的平衡。实践中常对c采用退火策略：

• 前期rollouts用较大的c，侧重探索
• 后期rollouts用较小的c，侧重利用，因为此时树已经有了比较丰富的先验价值评估

理解了这一点，阅读代码中的相关部分就不难了。

3.3 搜索最佳路径

（1）直接从树中搜索

在若干轮rollouts后，我们终于有了一棵相对完整的搜索树。现在对于一个原始问题（根结点），该如何选择一条最佳路径呢？

• 首先，找出所有有效的solution nodes。满足以下任一条件的node就是solution node：
• node是terminal A3 node（定义见3.1）
• node是A2 node（天然terminal）
• node是terminal A1 node（最后一个step包含“The answer is...”字符串，这是prompt控制的结果）。可见，solution nodes = leaf nodes + terminal A1 nodes。
• 这些solution nodes及其路径就是所有备选。接下来需要量化“好不好”。
• 每个solution node下有一个最终答案。以答案key对solution nodes做group by，得到字典 {ans1: [solution nodes], ans2: [solution nodes], ...}。
• 用“占比计算法”统计每个答案的投票得分：假设共有n个solution nodes，答案ans1、ans2、ans3下分别有a、b、c个node (n=a+b+c)，则ans1的得分为a/n，以此类推。
• 计算每个solution node的prior weights，用于衡量路径的整体质量，展开如下：
• 对于一个solution node，从它的父结点开始往上逐层遍历。
• 在每一层，将“遍历到的结点 + 它之前所有祖先结点”的推理结果拼接起来，然后用A2（propose remaining thoughts）的prompt基于这个拼接结果并行采样若干个一步到位的答案。
• 计算这些答案中有多少个与solution node的答案一致，该比例记为depth_score——即该层推理上的得分。
• 最终solution node的prior weight = 路径上所有depth_score的乘积。
• 直觉上，如果一条路径的每一步推理都稳定地指向solution node的答案，说明推理过程高度自洽，就给一个较高的prior weight。这样可以避免“答案对但过程可能是蒙对的”情况。论文中没提这一点，代码也写得较难读，这里特作说明。
• 如果只按之前simulate步骤中的reward方法，直接选投票得分最高的答案即可。但有时一个答案投票分虽高，推理过程质量却不一定好。我们假设：
• 现在，对于一个solution node，我们有投票得分（评估答案）和prior weight（评估推理过程质量），最终得分就是两者的乘积。有两种方法选择最佳solution node：
• 方法1：选择得分最高的那个solution node及其路径（当前代码的默认做法）。
• 方法2：计算每个solution node的prob = 它的score / 所有solution nodes的score总和，然后以概率prob随机选中一个（代码中有但并未使用）。

（2）使用discriminator

除了直接从树中搜索，rStar还提供了另一种巧妙思路：引入一个discriminator。这样，构造的搜索树相当于generator，我们再用discriminator从generator的结果中挑出最可信的——与GAN非常相似。generator和discriminator都是小模型，但属于不同的小模型。

具体过程：

• 在搜索树中，取出每一次rollout找到的leaf node以及最后一次rollout找到的solution node作为备选项。discriminator会从中挑出最好的。
• 对每条备选路径，将“用户问题 → 推理步骤 → 答案”拼接成完整文本。
• 随机选择一个推理步骤，mask掉该步骤之后的所有结果（包括答案），如下图。Question + SLM1是一条备选路径，Masked Solution SLM2则是把SLM1中的部分结果mask掉后输入给discriminator（SLM2）。

• 将mask后的文本喂给discriminator（SLM2）。如果SLM2生成的答案和SLM1的一致，则这条推理路径可信（consistent，绿色示例），否则不可信（红色示例）。
• 利用这种一致性，筛选出generator和discriminator一致的结果，然后再用前面的各种得分计算方法找出最可信的一条路径。更多细节可自行阅读源码。

四、总结

在本文中，我们以rStar为例，从代码级细节给出了o1（可能的）实现框架中Inference这块积木的一种具体实现。原本想放一些源码和注释，但考虑到在公众号里占据篇幅过多、可读性不高，就没有呈现。不过前面已讲清了源码中最精华的部分，希望能大大降低大家阅读源码的难度。

有了对MCTS如何在NLP推理任务上落地的初步理解，接下来就可以按兴趣广泛探索这块黄色积木的各种实现方式了（其实本质上都大同小异）。在后续系列中，我们将继续对框架进行拆解，加入更多积木。

你是一名 AI 行业编辑，请围绕下面这条热点输出一份资讯解读：
热点：OpenAI o1增强推理能力：MCTS代码实践解读要求：
1. 先用一句话解释这条热点在讲什么
2. 再总结它为什么重要
3. 说明会影响哪些 AI 产品或内容方向
4. 最后给出 3 个适合资讯站使用的标题