1 分钟阅读

简介

逻辑推理与决策规划:LLM+MCTS

  1. 为什么要将 LLM 与 MCTS 结合起来?
  2. 为什么 LLM 可以与 MCTS 结合起来?
  3. LLM 要如何与 MCTS 有效结合起来?
    1. 在训练过程中,MCTS 可以构造出更高质量的数据(比如<问题,推理轨迹COT,答案>)以供 LLM 训练;PS:引导让模型生成多个标签,这些标签对应于搜索所需的具体推理步骤。。在训练过程中,首先使用收集到的提示通过由预训练值模型指导的蒙特卡罗树搜索找到答案。随后,使用生成的问题-答案对来同时训练policy/actor模型和reward模型,迭代地改进该过程。这种方法的失败在于next-token的维度爆炸问题非常严重,在优先探索时间下只能采样一部分路径,这些路径可能是不良的,或者是局部最优的,而相关的细粒度reward模型也很难训练,最终导致policy模型难以迭代改进。
    2. 在推理过程中,LLM 通过与 MCTS 的多步交互与迭代,以时间换正确率。PS:MCTS是解码策略的一种,跟Beam search、Top-k sampling、top-p sampling等策略发挥的作用是一样的,只是后几种都是固定规则,而MCTS 有回溯的能力。具有回溯能力的前提是要有一个很好的PRM。每一个query的推理,都是一个全新的MCTS 树。
    3. 无论是推理还是训练,难点是以一个多大的粒度作为 树节点,一个token肯定不能是一个。此外就是如何得到一个好的PRM。
  4. LLM + MCTS 是终极方案么,有何局限性?

原理

OpenAI明确表明 o1的训练借鉴了AlphaGo的强化学习思路,而AlphaGo主要使用了Self-Play和蒙特卡洛搜索(MCTS)。根据围棋的规则,棋子可以在19 X 19的棋盘上选择落点。如果使用暴搜法,那么将有 361!种可能,即使去除其中不合法的情况,可能性仍比宇宙中的原子个数$10^80$高出20个数量级。此时,可以使用MCTS来近似暴搜的结果,MCTS是一种用于决策过程的启发式搜索算法,它包含四个步骤:

  1. 选择(Selection): 从根节点开始,根据某种策略(如上置信区间 UCB),选择一个最优的子节点,直到到达一个尚未完全展开或终止的节点。
  2. 扩展(Expansion): 如果所选节点不是终端节点,从该节点中随机选择一个未被访问过的子节点,添加到搜索树中。
  3. 模拟(Simulation): 从新扩展的节点开始,进行随机模拟(也称为“Playout”,通常使用随机策略),直到达到游戏的终局状态。根据终局状态的结果(如胜负或得分),评估该模拟。
  4. 反向传播(Backpropagation): 将模拟的结果反向传播到所有经过的节点,更新各个节点的统计信息(如胜率、访问次数),以便未来的决策可以基于更精确的估计。PS:类似于前人某条路走了xx次,胜率xx,后人可以参考

简单通俗地讲,AlphaGo包括两部分:一个CNN卷积神经网络和MCTS算法,CNN相当于人类的”棋感“系统,它负责根据当前棋盘上的形式给出下一步可能的落点的集合;MCTS相当于理性思考,可以在CNN给出的候选落点集合里找到胜率最高的位置。提升的CNN可以找到更好的落点,MCTS的性能也会提升,MCTS性能的提升的结果反馈给CNN,则CNN的性能会再次提升,整个过程形成正反馈循环。AlphaGo首先在人类棋手的棋谱上进行监督学习,然后自我博弈(Self-Play),即在每一局自我对弈中,AlphaGo 的不同版本与自己对战,并根据对局结果更新模型的策略网络和价值网络,通过自我对弈,AlphaGo不再依赖人类棋谱,能够探索新的战术和策略,超越人类棋手的水平。

OpenAI o1模型的前世今生MCTS多在Action为闭集的场景下使用。笔者认为o1没有使用MCTS,甚至没有使用搜索算法,而是使用类似STaR方法中的。其实MCTS在除了下棋以外的场景上基本没有成功的案例,尤其是在自然语言组成的几乎无限空间的场景上。PS: MCTS主要解决局部最优问题(llm推理时的token选择是固定规则,比如sampling,Best-of-N),从作用看,mcts与sampling和Best-of-N是并列的,认为MCTS是Tree Search的一种。

阿里Marco-o1推理大模型技术报告解读由论文内容推理MTCS和LLM的算法流程,具体来说:

  1. 在选择阶段,从根节点开始,算法根据特定的策略(如UCT策略)导航到最有潜力的子节点,直到到达叶节点。
    1. 节点的选择基于其累计奖励(置信度得分)和访问次数,优先选择奖励较高的路径。具体选择标准是从当前节点出发的推理步骤中,挑选出奖励分数 较高的节点,即具有较高累计置信度的路径。
  2. 在扩展阶段,在叶节点处,除非它代表游戏的终止状态,否则会添加一个或多个新的子节点来表示未来的可能移动。
    1. 在选择阶段到达一个尚未完全展开的节点后,将其子节点扩展到搜索树中。从当前节点输入 LLM,生成潜在的下一步推理输出,这些输出被作为新节点添加到搜索树。生成的子节点代表不同的推理方向(例如下一步的逻辑步骤或不同的解法路径)。扩展过程通常生成多个候选输出,例如基于 LLM 前 5 个置信度最高的候选 token,从而捕获可能的推理分支。
  3. 在模拟阶段,从新添加的节点开始,算法进行随机模拟(通常称为“rollouts”),选择随机移动直到游戏结束,从而评估节点的潜力。
    1. 模拟由 LLM 完成,执行一个完整的推理链展开。在每一步中,LLM生成的 token 置信度通过 softmax 计算得出。整个模拟路径的总体奖励v是所有 token 置信度得分的平均值。模拟结束的条件可以是推理链完成、达到预设长度或模型预测的终止标记。
  4. 在回溯更新阶段,模拟结束后,结果(赢、输或平局)被反向传播回根节点,更新每个遍历节点的统计数据(如赢、输次数),以指导未来的决策。
    1. 将模拟结果(奖励v)沿路径从当前节点反向传播,更新父节点及祖先节点的统计信息。每个节点更新其累计奖励和访问次数:奖励:$w_i <- w_i + v$ ,访问次数:$n_i <- n_i +1$ 。奖励值v是 roll-out 路径的置信度得分的平均值,表示路径推理质量。回溯更新确保整个搜索树能够动态调整,未来的搜索倾向于奖励较高的路径。PS:有一个推理llm (带上cot?)生成一个推理路径,有一个对推理路径评分的模型,进而优化推理llm的参数。 MTCS的用处是扩展了搜索空间,呈现类似TOT的效果 通过反复迭代这些阶段,MCTS逐步构建决策树,优化在状态空间庞大且难以直接计算最佳策略的场景中的决策。

推理动作策略。

  1. 动作选择。研究者观察到,以动作作为蒙特卡洛树搜索(MCTS)的粒度较为粗糙,这种方式往往导致模型忽视解决复杂问题所需的关键推理路径。为了解决这一问题,研究者尝试了不同粒度的搜索单元。起初,研究者使用完整的“步骤”作为搜索的基本单位。随后,为了进一步扩展模型的搜索空间并提升其解决问题的能力,研究者尝试将步骤细化为更小的单元,即每32或64个Token组成的“微步骤(mini-step)”。这种更细粒度的划分使模型能够以更高的精度探索推理路径。尽管在理论上,以Token级别为单位的搜索能够提供最大的灵活性和精细化,但由于计算资源的高昂需求以及构建有效奖励模型的难度,目前在实践中尚不可行。在实验中,研究者在MCTS框架下实现了以下策略:
    1. 步骤(Step)作为动作(Action):允许模型生成完整的推理步骤作为动作,每个MCTS节点代表一个完整的思维过程或行动标签。这种方法在探索效率上具有优势,但可能忽略解决复杂问题时需要的更细粒度的推理路径。
    2. 微步骤(Mini-step)作为动作(Action):将32或64个Token作为一个微步骤的动作单元。这种更细的粒度扩大了问题解决的空间,通过在搜索过程中引入更多细微步骤,增强了模型应对复杂推理任务的能力。在这一粒度水平上探索解决方案空间,能够帮助模型发现以较大动作单元可能忽略的正确答案。 研究结果表明,采用更细粒度的MCTS搜索策略可以显著提升模型的推理能力,从而更有效地解决复杂问题。
  2. 反思机制。研究者引入了一种反思机制,通过在每次推理过程末尾添加短语“等等!也许我犯了一些错误!我需要从头开始重新思考。”来促使模型进行自我反思并重新评估推理步骤。这种机制显著提高了解题的准确性,特别是在原始模型初始错误解决的复杂问题上表现尤为突出。加入反思机制后,大约有一半此类困难问题能够被正确解决。从自我批评的角度来看,这一方法使模型能够充当自己的批评者,从而识别推理中的潜在错误。通过明确提示模型质疑其初始结论,这一机制鼓励模型重新表达并优化其思维过程。自我批评机制充分利用了模型检测自身输出中不一致性或错误的能力,从而提升了解题的准确性和可靠性。反思步骤作为一个内部反馈循环,有效增强了模型在没有外部干预情况下的自我纠错能力。这一机制不仅显著提升了模型解决复杂问题的能力,还强化了其解题的准确性和鲁棒性。

代码

MCT Self-Refine (MCTSr)的算法(包含代码理解) 代码层级的理解看这里。

AlphaZero

def alphazero_search(state, neural_net, n_simulations=800):
    root = MCTSNode(state)
    for _ in range(n_simulations):
        # 1. Selection: 按 PUCT 选最优子节点
        node = root
        while not node.is_leaf():
            node = node.select_child()

        # 2. Expansion: 神经网络评估叶子
        policy, value = neural_net(node.state)
        node.expand(policy)

        # 3. Backup: 把 value 反向传播到根
        node.backup(value)

    # 返回根节点的访问次数分布作为动作概率
    return root.compute_action_distribution()

在真正走一步棋之前,先做 800 次模拟(simulation),每次模拟分三步:选择(Selection)、扩张/评估(Expansion/Evaluation)、回溯(Backup)。传统 MCTS 要 rollout 到游戏结束才能拿到价值,AlphaZero 相对于传统 MCTS 的两大创新,用两个网络分别解决了搜索的两个痛点:

  1. 策略网络 $p_\theta(a\mid s)$:输出“在这个状态下,每个动作看起来有多好”的先验概率。它的作用是缩小搜索宽度——传统 MCTS 在每个节点要平等考虑所有合法动作,但有了策略网络,我们只需要重点搜索那些先验概率高的动作。
  2. 价值网络 $v_\theta(s)$:输出“这个状态最终能赢的概率”的估计。它的作用是缩短搜索深度——传统 MCTS 必须 rollout 到游戏结束(终局)才能拿到奖励,但有了价值网络,到了叶子节点直接用 $v_\theta$ 估值,不需要继续往下搜了。

就这么两个改变,让搜索效率提升了几个数量级。传统 MCTS 在围棋上需要百万次模拟才能下出好棋,AlphaZero 只需要800 次。

搜索的核心是:在树的每个节点,我们应该优先扩展哪个子节点?AlphaZero 用的是 PUCT(Predictor + Upper Confidence Bound)公式来给每个候选动作打分: \(\operatorname{PUCT}(a) = Q(s,a) + c_{\text{prior}}\cdot p_\theta(a\mid s)\cdot \frac{\sqrt{N(s)}}{1+N(s,a)}\)

这个公式看起来有点复杂,我们把它拆成三项,每一项都有明确的含义:

  1. 第一项:$Q(s,a)$——这是利用项。动作 $a$ 的当前价值估计,也就是到目前为止,这条路径在模拟中平均拿到了多少价值。$Q$ 越高,说明这个动作历史表现越好,我们越应该继续考虑它。
  2. 第二项:$c_{\text{prior}}\cdot p_\theta(a\mid s)$——这是先验项。$p_\theta(a\mid s)$ 是策略网络给出的“这个动作应该被考虑的概率”。策略网络认为好的动作,即使还没怎么被探索过,我们也会给它初始的优先级。$c_{\text{prior}}$ 是一个常数,控制先验的权重。
  3. 第三项:$\sqrt{N(s)}/(1+N(s,a))$——这是探索项,来自经典的 UCB(Upper Confidence Bound)算法。$N(s)$ 是当前节点被访问的总次数,$N(s,a)$ 是动作 $a$ 被选中的次数。这个项的意思是:
    • 总访问次数越多,探索奖励越大(因为我们已经有足够多数据利用了,可以多探索);
    • 某个动作被选得越少,它的探索奖励越大(鼓励探索那些还没怎么试过的动作)。 这就是”利用-探索”平衡的精妙之处:既要选目前看起来好的,也要给那些”有潜力但还没试过”的机会。

你可能会问:”策略网络和价值网络是怎么训练出来的?一开始它们是随机的啊,搜索不也是乱搜吗?”答案是自我对弈(self-play)——AI 自己和自己下棋,用搜索的结果作为标签来训练自己,然后用更强的网络再做搜索,再对弈,循环往复,越来越强。这是一个”左右互搏”的过程。具体来说,每一局自我对弈是这样的:

def self_play_training(network, n_games=10000):
    for game in range(n_games):
        # 1. 自我对弈
        trajectory = []
        state = initial_state()
        while not state.is_terminal():
            policy = alphazero_search(state, network)
            action = sample_from(policy)
            trajectory.append((state, policy, action))
            state = state.next(action)

        # 2. 标注胜负
        winner = state.winner()
        for s, p, a in trajectory:
            value = +1 if winner == s.current_player else -1
            train_network(s, p, value)

我们一步步看:

  1. 自我对弈阶段:从初始状态开始,每一步都用当前网络做 MCTS 搜索,得到一个“更优的动作分布”(MCTS 算出来的访问次数),然后从这个分布里采样一个动作(加点探索),把(状态,MCTS 策略分布,动作)存下来。
  2. 标注阶段:一局棋下完,有了赢家。对于每一步,如果走棋的玩家最终赢了,这一步的价值标签就是 +1;如果输了,就是 -1。
  3. 训练阶段:用这些数据监督训练两个网络:
    • 策略网络:让它的输出 $p_\theta(a\mid s)$ 尽量接近 MCTS 搜索出来的动作分布——因为 MCTS 搜索比直接用网络选动作更强,这相当于“用搜索结果作为更好的策略标签”。
    • 价值网络:让它的输出 $v_\theta(s)$ 尽量接近最终的胜负结果(+1/-1)。

这个过程的美妙之处在于:标签不需要人给,是搜索自己生成的。一开始网络是随机的,搜索也很弱;但哪怕是弱搜索,也比纯随机网络强一点;用弱搜索的结果训练网络,网络变强一点;变强的网络又能让搜索变强一点……这样正反馈循环下去,网络和搜索一起进化,最终达到超人类水平。最惊人的是:整个过程不需要任何人类棋谱。AlphaZero 从零开始,不知道任何定式、任何开局、任何人类总结的围棋知识——它只知道规则,然后自己和自己下,72 小时后就超越了所有人类围棋程序。

AlphaZero 非常强大,但它有一个前提条件:你必须知道游戏规则——也就是知道状态转移是什么、合法动作有哪些、终局条件是什么( model-based RL)。因为 MCTS 需要知道”做了动作 a 之后,下一个状态是什么”。但如果我们不知道规则呢?如果环境的动力学是未知的呢?比如 Atari 游戏从像素输入,你根本不知道游戏代码;或者真实世界,物理规则太复杂没法写出来。

不太行?

MCTS为什么在LLM任务上不好用?最近半年仍然有很多文章在用MCTS来做LLM的decoding,或者用MCTS造数据来改进ReST以self-improvement范式迭代训练LLM和PRM。先说我的结论:广泛被探讨的1)节省token,2)造出质量更高的数据,我都没做出来(期待大佬做出来打脸)

  1. MCTS作为LLM的decoding策略是否有优势(MCTS相比BoN能提升token-efficiency)?
    1. 首先MCTS需要一个质量很高的PRM,如果PRM质量都不行的话何谈guide搜索。
    2. 认为MCTS相比BoN能提升token-efficiency的论据是:“BoN每条路径不管好坏都roll到底,但如果路径上已经出现明显错误,后续就无法roll出正确答案”,PRM-guided的MCTS可以把它剪枝掉。
    3. 但实际想做到却很困难。分两方面归因:
      1. 工程上,MCTS中每次expand若干个节点,其实很多节点后续没有被使用,造成浪费;并且在每个节点都需要额外用llm-as-judge判断一下是否已经得到了答案并判定为叶子结点
      2. “BoN的每条路径是不管好坏都roll到底,但显然如果路径上已经出现了明显的错误,后续就无法roll出正确答案了”,这个论断真的正确么?LLM任务相比围棋任务的一个显著特征就是【并没有严格的状态转移】。围棋中,一步烂棋就会导致后续很难挽回,所谓一步错,步步皆错;但LLM,可以通过’but’,’wait’补救回来,前面再多错误都没关系,只要cot足够长总有扭转乾坤的希望。
  2. 造出更高质量的数据。
    1. 无论是ReST-MCTS, 还是AlphaLLM, MCTS都只是起到一个”启发式搜索”的作用,相当于只是通过PRM剪枝掉了一些明显错误的路径,相比beam search没有本质区别。这些工作里的“伪”MCTS,相当于仅以root节点为起点做了一次MCTS,而不是alphago zero中在棋局的每个中间状态都进行。MCTS搜索了若干个iter后,会得到一棵搜索树,这棵树上有很多输出了answer的叶子结点,所有从根节点到叶子结点的路径都是一条reasoning trace,而只要叶子结点被verifier验证为答案正确,这条trace就可以被收集,作为SFT的训练数据,用来训练LLM。
    2. 但是,这种做法是否有意义?本质上这里每一步扩展节点都仍旧基于LLM本身的policy(也即alphago zero中的prior probability $p$,而不是improved probability $\pi$),roll出的数据其实没有本质提升,那么我就用BoN收集数据是不是就够了(退化为ReST或者说reject sampling)。

封闭域的围棋和开放域的LLM任务有什么区别?

  1. 树的宽度(每一步的动作空间规模):361 vs 词表规模(qwen是150000),后者的规模显著高于前者。
  2. 树的深度:围棋的树的最大深度是确定的(大多数职业围棋对局通常在200-300手之间结束),到达这个深度后一定能获得明确的胜利(+1)or失败(-1)的反馈信号。因此alphago中有如下论断:随着MCTS迭代次数趋向于无穷,一定能获得更多的信息,MCTS搜索后的improved policy一定能逼近最优策略。
  3. 但LLM任务呢,不确定!理论上对于一道题,reasoning trace可以无限制地roll下去。这就导致,你也不知道什么时候抵达了叶子结点。如前文所述,ReST-MCTS中额外引入了一个模块,让LLM自己判断当前节点是否已经输出了问题的答案,可以终止了。但这在工程上也导致token efficiency的下降。
  4. 语言具有显著的goto性质,前面说的都是错的没关系,一个but、wait就可以全部推翻重来,这就导致MCTS时的一个直观弊端:举个例子,你当前的推理步骤出了个很低级的错误,但后面靠着模型的随机性but一下又roll到正确解了,这步低级错误的back propagate后的value反而还挺高。但围棋是落子无悔的,这步下的很差后面就很难挽回,状态与动作之间具有强因果性。所以lookahead search和alphago zero的“真”MCTS尽管能跑通,但以上前两点导致它们计算开销巨大,得砸一堆卡和时间;第三点导致back propagate后的value bias很大(更何况PRM本身就有很大bias了),性能没法保障。相信也有其他从业者在小规模尝试发现不work后就放弃了scale。 MCTS在哪些任务上有潜力?
  5. 如果还是math、code这种token-level的任务,可能需要人为定义macro action,降低动作数量,例如微软rstar-math的前身r-star。
  6. web navigation等multi-turn的agent任务,天然具有强因果性的状态转移,适合被建模为MDP;且动作空间规模和trajectory长度都是几十这种可枚举数量级的,树的宽度深度都得到保障,那么MCTS很合适。

留下评论