rStar论文精读

论文简介

论文标题：《Mutual reasoning makes smaller LLMs stronger problem-solvers》
论文地址：https://arxiv.org/abs/2408.06195
录用会议：ICLR2025

背景与挑战

挑战1：在SLM中平衡exploration与exploitation。一些方法有很大的exploitation但限制任务多样性泛化性不好；一些方法有很大的exploration但产生很多较差的结果，需要很多次尝试才能获得理想结果，消耗推理资源大。

方法

为了解决挑战1，设置了5个模拟人类的推理动作以进行彻底的探索；为了解决挑战2，使用了针对SLM设计奖励函数去评估中间步骤，而不是使用不可靠的自我评估。此外，使用另一个SLM作为鉴别器来增强MCTS过程，与生成器SLM相互验证每个轨迹的正确性。

动作空间

MCTS生成的核心是动作空间，它定义了树探索的范围。大多数基于MCTS的方法使用单一的动作类型来构建树。例如，在RAP中，动作是提出下一个子问题，而在AlphaMath（Chen et al., 2024a）和MindStar（Kang et al., 2024）中，动作是生成下一个推理步骤。然而，依靠单一的动作类型很容易导致空间探索无效。因此为了解决以上问题，使用了人类思考的方法进行推理.

• A1: Propose an one-step thought.
• A2: Propose the remaining thought steps.
• A3: Propose next sub-question along with its answer.
• A4: Answer the sub-question again.
• A5: Rephrase the question/sub-question.

动作A1和A3每层最多5个节点，其他动作的节点数默认为1。

蒙特卡洛模拟

蒙特卡洛搜索过程参考了Alpha Go的推理过程。

在某个Node处，有5个Action可供选择，Q都被初始化为0。
每次进行rollout时对UCT最大的Action进行选择，使用SLM进行生成。
这样，收集了一些较为随机的rollout

奖励函数

没用使用特定的奖励函数模型，而是使用了模拟打分的方式作为奖励函数。引入了另一个SLM’，屏蔽后i步，通过前i-1步直接推理到结尾。比较SLM’完成的答案是否与原始轨迹t匹配。如果它们一致，将t视为最终选择的有效轨迹。

考虑学生在没有老师反馈的情况下解决问题。一个对自己的解不确定的学生（SLM1）可能会请同伴（SLM2）审查他们的推理过程。如果同伴在给定相同的初始步骤后得到相同的答案，学生就会对自己的解更有信心。这种同伴验证过程反映了我们旨在实现的互推理一致性。 由SLM1进行最终轨迹选择。

轨迹获得的奖励 = rollout过程中给出的奖励 * SLM2给出的答案奖励

我们通过将每个轨迹的奖励乘以从 rollout 中获得的终端节点的置信度得分来计算每个轨迹的最终得分。最终得分最高的轨迹被选为解。

总结

UCT：Rollout期间，Action的选择
SLM：Rollout期间，使用选择的Action作为提示进行生成
SLM’：对Mask的trajectory进行评估，是否符合自一致性