[Arxiv 2024] rStar: Mutual Reasoning Makes Smaller LLMs Stronger Problem-Solvers

Introduction

• 作者提出 Self-play muTuAl Reasoning (rStar) 用于增强 SLM (Small Language Model) 的推理能力，不需要进行额外的模型训练或 reward model，仅依靠推理时的搜索就能提升模型的推理能力；作者的贡献点主要包括：(1) 定义了不同的推理 action 有效扩大了搜索空间；(2) 提出使用一个额外的 SLM 作为 discriminator 帮助模型从搜索树中选择合适的推理路径 (不需要依赖额外的 reward model)
• 个人觉得这篇文章里特别是定义 reasoning action 构建搜索树还是很值得学习的，我们甚至可以将其扩展到 post-training 上并尝试自定义地添加 “reflection” action 来帮助模型获得类似 o1 的反思能力，也可以自定义地添加更多推理动作，例如 Verification (“Let me check my answer”)、Backtracking (“Let’s try a different approach, what if we…”)、Backward Chaining (“Working backwards, 24 is 8 times 3”) 等 (from Cognitive Behaviors that Enable Self-Improving Reasoners, or, Four Habits of Highly Effective STaRs)；此外，无需 PRM 评估推理步骤 reward 的方法极大降低了训练成本，也可以用在自动化生成 reward model 训练数据上，整篇论文感觉还是比较有价值的
  
• Q. (1) rStar 每次推理需要做若干次 MCTS rollout + discriminator 推理，作者没有给出推理成本的分析；(2) 作者对于 discriminator 的选择没有探索得很清楚，例如 table 5 里 llama-3-8b-instruct 使用 3.8B 的 Phi3-mini-4k 作为 discriminator 效果比较好，但用自身作为 discriminator 效果却明显更差，这是为什么？

Method

A Rich Set of Human-like Reasoning Actions

• 问题回复的过程可以被划分为若干推理步，作者按照人类的思维方式，给推理步设计了 5 种不同的 reasoning actions，模型在推理时可以采用其中一种 action 进行推理，能够帮助模型更有效地构建搜索树；这些回复都可以用 prompt 来达成，可以参见源码中 rStar/prompts 部分
• A1: Propose an 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.

• A2: Propose the remaining thought steps. 一次性推理完毕，直接得出最终答案 (standard CoT enabling “fast thinking”)

### 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.

• A3: Propose next sub-question along with its answer. 将原始问题拆解成若干子问题并做相关回答。最后一个子问题的答案即是最终答案

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.

• A4: Answer the sub-question again. 有时 A3 中某个子问题的回答不一定可信 (A3 使用的是 least-to-most problem decomposition prompt)，我们尝试重新回答它。这时我们会采用 A2 的模版 (few-shot CoT)，重新回答这个子问题；例如，对于 A3 的 Question1.1，你可能并不确定 Answer1.1 是否正确，这时你想重新再思考一次 Answer1.1 的答案。由于此时你只是对某一个子答案做修正，因此你可能采用 A2 (propose the remaining thought steps) 的方式，做一些简单的推理，重新取得 Answer1.1。此时相当于把 Answer1.1 用 A2 例子中的输出结果进行替代
• A5: Rephrase the question/sub-question. 重新复述一个原始问题/子问题。例如去掉大段文字表述信息，只把关键部分提取成 condition1…, condition2 之类的形式，用这个形式当作新的问题

Constructing Action Space

• 通过对上述 5 种 reasoning actions 进行采样，我们可以得到一个巨大的搜索树，使得模型在每一个推理步都可以选择一种最合适的推理方式进行推理，搜索树如下所示：

---

• 基本信息. 搜索树的根结点是原始问题。方形 node 表示终止结点 (leaf node)，例如图中的 COT 结点 (A2)。但注意，不是只有 COT 结点才是 leaf node。例如 A1 中的最后一个 step (定义是最后一个 step 中含 “The answer is…” 字符串，能做到这一点是因为我们通过相关 prompt 来指示模型生成结果)，A3 中最后一个子问题的子回答都可以成为 leaf node (“最后一个的含义” 是，子问题中包含原始问题，或者子问题以 “Now we can answer the question" 开头，能做到这一点是因为我们通过相关 prompt 来指示模型生成结果)。虚线表示选择性探索 (根据你的脚本配置决定)，实线表示必须性探索
• 接下来我们来看图中的更多细节：我们从根结点 (第 0 层) 出发，根结点是用户的原始问题，对于根结点来说，A1 (a step)，A5 (rephrase) 是选做的，A2 (remaining steps/COT)，A3 (next subqs and subas) 是必做的。其中，经过 A5 后，相当于从一个全新的用户问题出发，所以 A5 之后创建分支的规律和根结点一致，因此图中不再画出。现在我们观察图中的第 1 层，也就是根结点所有的子节点。先看第 1 层中的 A1 部分 (蓝色块)，它表示对原始问题并行采样若干个 step1 (代码中默认采样 3 个)。每个 step1 都是一个 A1 类型结点，并行采样的目的也是为了拓宽搜索空间。A1 类型结点的值包含 “该 step 推理文本 + 该 step 阶段性答案”；再看第 1 层中的 A2 部分 (绿色块)，它表示对原始问题做简单推理后，一步生成最终结果。这里同样采用了并行采样的方式 (默认值为 32)。因为已经生成最终答案，所以 A2 结点都是 leaf node，它不会再往下做任何扩展。再看第 1 层中的 A3 部分 (红色块)，它表示对原始问题产出 next subquestion 及其对应的 subanswer。这里同样采用并行采样的方式 (默认值为3)，例如对原始问题，我们并行采样三个 Question1.1 和 Answer1.1。接下来我们从第 1 层出发，对于 A1 类型结点，它只能产生 A1 (选做) 和 A2 (必做) 两种类型的子节点。也就是说，基于当前 step，我们必须产出一个一步到位的推理结果，选择性产出一个 step by step 的结果；对于 A3 类型结点，它可以产生 A1 (选做)，A2 (必做)，A3 (必做)，A4 (必做) 类型子节点，对于 A4 类型结点，它可以产生 A1 (选做)，A2 (必做)，A3 (必做) 类型子节点，当我们执行 A4 时，你可以理解成只是重新修改了它的 parent 层的 sub answer

Solution Generation with MCTS Rollout

• 作者在搜索树的基础上利用 MCTS 高效采样，搜索出最好的推理步骤

---

• rollout: 执行 select 步骤，从根结点出发，如果发现某一层所有的结点都被探索过了，那么我们就选择 UCT 值最大的结点，向下层继续搜寻，直到找到一个未被探索过的结点为止，随机选择一个未被探索的结点；执行 expand 步骤，为被选中的结点创建所有可能的子节点；执行 simulate 步骤，随机选择一个子节点，重复执行 “expand-随机” 步骤，直到遇到 leaf node 或者达到设定的最大搜索深度为止，注意，这里的 leaf node 只有 Terminal A2/A3 node 两种；执行 backprop 步骤。这一步我们将计算 leaf node 的 reward，同时将本次搜索路径上所有 node.reward += leaf_node.reward，node.freq += 1
  • Terminal A3 node reward. 对于 A3 类型的 leaf node，我们对这最后一个子问题，并行采样若干个子回答。假设我们采样 $n$ 个子回答，这些回答中指向答案 $x,y,z$ 的条数分别是 $a,b,c$ ($n=a+b+c$)，那么 $x$ 答案的占比就是 $a/n$，以此类推，我们选择占比最大的那个答案作为最终答案，并将这个占比作为 reward
  • Terminal A2 node reward. 对于 A2 类型的 leaf node，我们则直接在它的所有并行采样结果中计算答案占比，计算方式同上
• 以此循环往复，在执行若干轮 rollouts (代码默认值为 16) 后，我们就有一棵相对完整的搜索树了，接下来我们就可以基于这棵树去找到一条最佳的推理路径了

结点的 UCT 值计算方式如下：
$Q$：截止到本轮 rollout 为止，该结点的累积 reward
$N$：截止到本轮 rollout 为止，该结点的累积被访问的次数
$N_{parent}$：截止到本轮 rollout 为止，该结点的父结点累积被访问的次数
$c$：探索权重，$c$ 值越大，更侧重 explore 而不是 exploit；一般来说，我们会对 $c$ 采用一些 "退火策略“ 即：对于初期的那些 rollouts，我们用较大的 $c$，侧重探索，对于后期的那些 rollouts，我们用较小的 $c$，侧重利用

Reasoning Trajectory Selection with Mutual Consistency

• 给定 MCTS 得到的搜索树，对于这棵搜索树，我们可以找到所有有效的 solution nodes (terminal A1 node / A2 node / terminal A3 node)，这些被找出的 solution nodes 及其相关的路径就是我们所有的备选项。我们想从这些备选项中选择出最佳推理路径
• 一种方式是直接从树中搜索：每一个 solution node 下都涵盖一个最终答案。我们以这个答案为 key 对所有 solution node 做 group by，即含有相同答案的 solution node 为一组。这时我们会得到一个字典，形式如 {ans1: [solution nodes], ans2: [solution nodes], ...}；接下来我们同样通过前文所说的 “占比计算法”，统计每个答案的投票得分。例如一共有 $n$ 个 solution nodes，答案 ans1, ans2, ans3 下的 solution nodes 数量分别为 $a,b,c$ ($n=a+b+c$)，那么 ans1 的投票得分为 $a/n$，以此类推。我们可以通过计算每个 solution node 的 prior weights 来衡量其所在推理路径的整体质量：如果按照之前构造搜索树时 simulate 步骤找到最终答案和计算 reward 的方法，我们只需要选择投票得分最高的答案即可，但是有时候一个答案的投票得分虽然高，整体的推理过程质量却不一定好。那么什么样的推理过程才算好呢？我们可以做一个简单的假设：对于一个 solution node，我们从它的父结点开始往上逐层遍历，在每一层中，我们将 “遍历到的结点 + 它之前所有的祖先结点” 产出的推理结果拼在一起，再使用一个 A2 (propose remaining thoughts) 的 prompt，基于这个拼接结果，并行采样出若干个一步到位的答案。我们计算这若干个一步到位的答案中，有多少个 solution node 的答案一致，我们记这个比例为 depth_score，也就是每一层推理上的得分。那么最终一个 solution node 的 prior weight = prod(depth_score)，也就是它的路径上每一层 depth score 的乘积。这个 prior weight 直觉上的理解就是，如果一条路径上，每一步推理过程都稳定地指向 solution node 的答案，说明整个推理过程是高度自洽的，这时我们就给这个 solution node 一个比较高的 prior weight。这样就可以避免 “答案是对的，但是过程可能是懵的” 的情况；现在，对于一个 solution node (及其相关路径)，我们有一个用于评估其答案的投票得分，和一个用于评估其推理过程质量的 prior weight，那么这个 solution node 的最终得分就是两者的乘积，选择得分最高的那个 solution node 及其 path 即可
• Mutual Reasoning Consistency by Discriminator ${\text{SLM}}_2$. rStar 还提供了另一种巧思：借助一个 discriminator (作者用 3.8B 的 Phi3-mini-4k 作为 discriminator)。也就是我们构造的搜索树相当于一个 generator，我们使用 discriminator 从 generator 的结果中找到最可信的那个。我们来看详细的过程：在构造好的搜索树中，我们取出每一次 rollout 找到的 leaf node，以及最后一次 rollout 找到的 solution node 作为备选项。对于每一条备选路径，我们把整条路径上的 “user question $\to$ 推理步骤 $\to$ 答案” 拼接起来，形成一个完整的文本。对于这个文本，我们随机选择一个推理步骤，然后 mask 掉这个推理步骤之后的所有结果 (包括答案)，如下图所示，Question + ${\text{SLM}}_1$ 是我们的一条备选路径。Masked Solution ${\text{SLM}}_2$ 则是我们把 ${\text{SLM}}_1$ 中的一些结果 mask 掉后，输入给我们的 discriminator ${\text{SLM}}_2$。如果 ${\text{SLM}}_2$ 生成的答案和 ${\text{SLM}}_1$ 的一致，则我们认为这个推理路径是可信的，因为两个不同的模型都基于同样的推理文本给出了相同的结果 (consistent，图中绿色示例)。否则就是不可信的 (图中红色示例)，我们使用这种一致性，从备选路径中筛选出一波 generator 和 discriminator 一致的结果。然后再用之前所说的计算各种得分的方式，找到最可信的 1 条路径
  

Experiments

• Main Results.
  
• Ablation Study.
  

References

• Qi, Zhenting, et al. “Mutual reasoning makes smaller llms stronger problem-solvers.” arXiv preprint arXiv:2408.06195 (2024).
• code: https://github.com/zhentingqi/rStar
• OpenAI o1 技术初探2：使用MCTS增强推理能力（基于源码实践的解读）