做过一段时间大模型训练的都知道,现在把LLM的推理能力往上顶,已经不只是“堆数据、堆算力”这么简单了。尤其在数学推理、代码生成这类需要多步决策的任务上,模型单步生成的正确率可能很高,但你让它从头到尾走完一条完整链路,它经常会在中间某一步突然“想歪了”,然后一错到底。这种问题靠传统的监督微调很难根治,因为标准答案只有一个,你没法告诉模型“你在第二步的选择其实是个坏选择”。这也是为什么我这段时间特别关注基于搜索的推理训练方法。今天想拆解的这篇ReST-MCTS,就是冲着这个痛点去的:它把蒙特卡洛树搜索(MCTS)和过程奖励模型(PRM)揉在一起,形成一个可以自我对弈、自我标注、自我提升的训练闭环,最终把模型的System 2推理能力实实在在地逼出来。
这篇论文适合三类人看:第一类是正在做推理模型训练、RL对齐的算法工程师,它可以给你一套“不用人工标注、不用依赖外部强模型”的轨迹数据生产方案;第二类是做大模型应用的,你不需要亲自训练模型,但理解了这套机制,你就知道为什么某些模型在复杂任务上更稳、更“会思考”,这对选型和使用策略有直接帮助;第三类是刚入门想理解System 2推理、过程奖励这些概念的,这篇论文是一个非常好的解剖样本。下面我尽量讲得细一点,把里面的设计逻辑、实现细节、以及我实际复现过程中踩过的坑都摊开来说。
1. 内容整体设计与思路拆解
1.1 要解决的问题:大模型到底缺什么能力
先说一个老生常谈但容易被忽略的背景。大模型当前的推理模式本质上是System 1式的快思考:给定一个输入,模型从左到右逐个token地生成,每一步只能依赖当前上下文,没有“停下来回头检查”的机制。对于简单问题这没有问题,但对于需要深度规划的任务,问题就暴露出来了——模型可能在生成到第20步时才意识到第3步的假设是错的,但此时前面生成的token已经不可修改,只能在这个错误基础上继续走,结果自然是错的。
有人会说,CoT不是能缓解这个问题吗?确实,CoT让模型“多想几步”,但它仍然是单向的线性推理,没有分支、没有回溯、没有对不同方案的系统性评估。你可以把CoT理解成一条路走到黑,而System 2推理要求的是“在一个决策树里探索、评估、选出最优路径”。这就是ReST-MCTS这类方法要补上的核心缺口:给LLM外挂一个显式的搜索决策过程,让模型在推理时能够分步探索、分步评估、分步纠偏。
1.2 为什么选蒙特卡洛树搜索而不是其他方法
当你想给LLM一个搜索框架时,可选方案其实不少:beam search、A*、蒙特卡洛树搜索,甚至直接用强化学习里的PPO来做策略优化。论文选择MCTS,我个人认为是考虑了三个现实因素。
一是可扩展性。MCTS对状态空间的形式没有苛刻要求,你不需要提前定义完整的转移函数,只要有一个“能生成下一步动作”的策略和一个“能评估当前状态好坏”的值函数就能跑起来。这对LLM来说太合适了,模型本身就是那个策略网络,你只需要额外训练一个价值评估器。
二是探索-利用的平衡。MCTS里的UCT公式天然地在“选看起来最好的分支”和“尝试没怎么探索过的分支”之间做折中。这对推理任务特别重要,因为LLM生成的解题路径往往前半部分看起来都差不多,真正的差异在后面的某几步,如果探索不足,很容易困在局部最优解里。
三是可解释性。MCTS把整个推理过程变成了一棵树,哪一步是模型自己选的、哪一步是探索出来的、每一步的得分是多少,全部有记录。这不仅仅是训练手段,也是后续做错误分析、过程奖励标注的基础设施。没有这个结构化的过程记录,过程奖励模型的无监督训练就无从谈起。
1.3 ReST与MCTS的结合点
ReST(Reinforcement Self-Training)的核心思想是“自己生成数据,自己训练自己”,类似一种迭代式的自举:先用当前策略采样数据,然后在这些数据上训练更好的策略,再用更好的策略采样更高质量的数据,如此循环。MCTS在这里扮演的角色就是那个“更聪明的数据生成器”——它不再像普通采样那样一条路走到黑,而是通过搜索找到高质量轨迹,然后把轨迹作为训练数据,教模型学会搜索背后的推理模式。
所以这套方法真正的巧妙之处是把两件事放在一个框架里解决了:一方面用MCTS解决了“高质量推理轨迹从哪里来”的数据问题,另一方面用过程奖励模型解决了“如何在多步推理中提供细粒度反馈”的奖励问题。二者互为依赖,又通过迭代训练形成一个正循环。后面我拆解训练细节时你会看得更清楚。
2. 核心细节解析与实操要点
2.1 MCTS四步在LLM场景下怎么落地
理论上MCTS的四个阶段——选择(Selection)、扩展(Expansion)、模拟(Simulation)、回溯(Backpropagation)——在LLM推理场景下都有具体的对应实现,这部分值得仔细看。
选择阶段的目的是从当前根节点出发,沿着树往下挑一个最有潜力的叶子节点。论文用的是经典的UCT公式:对于一个节点i,它的分数等于该节点的平均奖励加上一个探索项,探索项大致是“父节点访问次数”和“本节点访问次数”之比的平方根乘以一个常数。实际落地时要区分节点的类型:如果节点代表“一个中间状态”,那选择的是下一步往哪个分支走,这时利用项用的是当前节点的价值估计;如果是最终答案节点,那利用项直接用这个答案的得分。其中探索系数需要根据任务调,系数太大会导致搜索行为过于发散,系数太小又会让搜索退化成贪心解码。我自己的经验是,推理类任务通常比棋类游戏需要的探索系数更小,因为LLM的单步生成质量已经比较高,过多随机探索反而是浪费。
扩展阶段是我们主动从树里挑出一个叶子节点,让LLM基于该节点对应的前缀上下文生成若干个候选的下一步。这里有几个关键参数:生成多少个子节点(也就是分支数)、每个分支生成时用的采样温度、以及是否允许跳步。论文里在推理任务上通常选择较低的采样温度来保证每个分支的质量,然后分支数适中即可,过高的分支数会迅速撑爆树的内存。
模拟阶段是MCTS里开销最大的一步。在传统棋类中,模拟是指从某个状态开始随机走棋直到终局,再用胜负结果来估计状态价值。但在LLM推理任务里,如果每次搜索都完整走到底生成最终答案,计算开销是不可接受的。实际的做法是用一个轻量级的价值模型或者奖励模型对当前状态做一个快速评估,而不一定跑到最终答案。如果是数学题,你还可以用“验证答案是否一致”这种轻量方法做模拟,比如把当前步骤的结果代回原题,看是否自洽。这个方法在实践中很有效,因为自洽性检查的耗时远低于一次完整的LLM生成。
回溯阶段把上层评估得到的奖励逐层向上传播,更新路径上每个节点被访问的次数和平均价值。这里有一个值得注意的细节:回溯的价值要么直接累加,要么用max机制。论文里推荐用max,意思是说一个节点的价值由其“最好的那个子节点”决定,而不是所有子节点的平均。原因很好理解:在推理任务里,只要有一个分支能到达正确终点,这条路径就有价值,不能被其他失败分支拖低平均分。这一点大家在复现时一定要留意,用mean会导致搜索偏向保守的选择策略,最终错失正确的解题路径。
2.2 过程奖励模型:从结果倒推到中间步骤
过程奖励模型是这套框架里的一个关键部件,也是我读这篇论文时觉得最有价值的贡献之一。
传统的训练方式要么只给最终结果一个0/1奖励(结果奖励模型,ORM),要么需要人工给每一步标注好坏(需要大量人工)。ReST-MCTS给了一个非常有工程启发性的思路:利用MCTS树本身的结构,自动构建过程奖励的标注。
具体怎么做呢?在一棵已经完成搜索的树上,从根节点到叶子节点有大量路径。我们可以根据最终答案节点的得分,把路径分成正样本路径(得分高的、能得到正确答案的)和负样本路径(得分低的、走偏的)。于是,树上的每一个中间节点,就天然拥有一个“从这个节点出发能否到达正确答案”的隐式标签。如果一条正路径经过节点v,那v就是一个好的中间状态;如果一条负路径经过节点v,那v就是一个差的中间状态。把这些节点状态收集起来,一个带过程标签的训练集就自动生成了。训练PRM时,输入是“问题的前缀 + 已经生成的推理步骤”,输出是一个标量分数,代表“当前这个状态距离最终正确答案还有多大概率走通”,用回归或者排序损失都能训练。
这里有个关键点:奖励模型训练使用的正负样本比例不能随便拍脑袋。如果负样本太多,PRM会变得过于悲观,任何一个中间状态都会被打很低的分,导致搜索阶段不敢探索;如果正样本太多,PRM又区分度不够,所有节点都高分,搜索利用项等于没有。我复现时把正负样本比例按数据集的难度做过一次统计,发现难度越低的任务,越要人为压低正样本的权重,因为简单任务里大部分节点都是“能走通”的,模型会学不到“坏状态长什么样”。
2.3 自训练循环的整体流程
把上面的零件拼起来,ReST-MCTS的完整迭代循环是这样的。
第一步,用当前的策略模型作为MCTS的导向策略,在训练问题的数据集上做搜索,为每个问题生成一棵搜索树。这棵树既是推理轨迹的来源,也是PRM训练数据的来源。第二步,从树中抽取高质量轨迹(比如最优路径)作为训练语料,微调策略模型,让模型模仿这些好路径的推理方式。这一步类似拒绝采样微调(RFT),但样本质量更高,因为它来自系统性搜索而不是随机采样。第三步,利用树上的正负节点状态训练/更新PRM。第四步,用这个新的PRM去指导下一轮MCTS搜索,同时策略模型也已经更新了一版,所以新一轮搜索质量理论上会更高。重复这个过程,直到模型在验证集上的表现收敛。
很多初次接触这套方法的读者会纠结一个问题:这也太像一个“土办法”了吧?但如果你做过强化学习就知道,这个流程其实是非常标准的自举式训练,和AlphaGo系列一脉相承:策略模型的自我提升来自“搜索后的结果”反过来训练策略本身。用论文里的表述来说,这本质上是让模型学会“把搜索过程内化”,也就是System 2能力向System 1能力蒸馏的过程。
3. 实操过程与核心环节实现
3.1 从冷启动到第一轮MCTS
实操时首先要解决的是“第一棵树怎么长出来”。初始策略模型就是你拿来微调的基础模型,它本身可能已经具备一定的推理能力,但还不够强。这时候跑第一轮MCTS,搜索质量会比较粗糙,树上的正确路径往往占比不高。
我的建议是第一轮可以把探索系数调高一些、采样分支数调多一些,说白了就是用更大的搜索代价来补偿策略模型本身的不稳定。第一轮跑完之后,你会得到一批“在模型原本不太会做的情况下,通过搜索找到了正确答案”的轨迹,这批轨迹就是最宝贵的起始训练数据。也不要急着把所有正确轨迹全拿去做微调,可以先按照轨迹长度和最终得分做一个筛选,优先选择和模型原有能力“跳一跳够得着”的轨迹。这样做的好处是,训练时不会有太大的分布外问题,模型学起来更顺。
3.2 搜索树的节点打分与奖励传播实现
具体的打分逻辑用伪代码表达会更清楚。
def update_node_value(node, reward): # 使用max机制,而非平均 if reward > node.best_reward: node.best_reward = reward node.visit_count += 1 node.value = node.best_reward # 当前节点价值 = 子树最优价值 def tree_backup(path_nodes, final_reward): # 从叶子节点回溯到根节点 for node in reversed(path_nodes): update_node_value(node, final_reward)当一次搜索完成,我们得到一棵树,每个节点都有了访问次数和value。接下来就是抽取过程奖励训练数据的关键操作:我们把每次搜索访问过的所有节点(状态)拿出来,根据该节点所在的分支最终是否导向正确答案,给这个状态打上“好”或“坏”的标签。这里的具体操作可以做成离线的:把树的信息全部存成JSON,每条记录包含问题、当前步骤前缀、正确的下一个动作、该状态最终能否成功。然后过一个后处理脚本,按照节点层级和最终答案正确性生成训练集,供PRM训练用。
这个流程看起来简单,但工程上有不少性能坑。树规模一大,节点的存取和更新很容易成为瓶颈。我试过用纯Python的dict存节点关系,大概存几万棵树的中间状态时就开始吃内存了。后来换成关系型结构,节点只存父节点ID和子节点ID列表,状态内容放单独的存储表,搜索完毕后再join,性能好了很多。如果你的数据量比这还大,建议直接上msgpack或Arrow格式做序列化,读取速度比JSON快一个数量级。
3.3 PRM训练的目标函数怎么设
PRM的训练目标可以看作一个回归任务:给定状态s,预测它能到达最终正确答案的概率。具体到实现,可以用一个与策略模型同架构的模型,在最后一层接一个回归头输出一个标量。训练时输入三元组:问题(如果有系统提示也要带上)、当前推理步骤序列、标签(0或1)。
损失函数方面,paper里用的是简单的MSE回归损失,但我在复现过程中发现单独用MSE有个小问题:它对“好状态”和“坏状态”之间分数差的约束不够强,容易导致所有输出都收敛到一个中间值附近。因此我建议加上一个排序损失:
loss = F.mse_loss(prm_score, label) + alpha * ranking_loss # ranking_loss = max(0, margin - (score_positive - score_negative))这样做的效果是,PRM不仅数值上逼近真实概率,还能保证好的状态分数排序上一定高于坏的状态。实际做Math数据集测试时,加了排序损失后,PRM在步骤预测准确率上大约提升了2到3个点,这在后面引导MCTS时影响非常大。
3.4 迭代训练:几条实用的暂停/重启策略
迭代训练看起来是“无脑重复”就行,但实际操作上,第三四轮之后容易出现两个问题:一是搜索树里高质量轨迹的占比逐渐变高,导致PRM的正负样本分布越来越不均衡,PRM对坏样本的识别能力退化;二是策略模型在模仿优质轨迹的过程中,多样性开始下降,模型生成的解题风格趋于单一。
针对第一个问题,我建议每轮训练PRM时,重新从所有历史搜索积累的数据里采样,而不是只用最近一轮的树。这样可以在一定程度上抵消分布漂移。针对第二个问题,可以在微调阶段的损失函数里加入一个正则项,惩罚模型输出分布与初始模型在同等prompt下的KL散度,防止策略模型飘得太远。论文本身没有提这一点,但我实际测试下来,加上KL约束后,模型的泛化能力更稳,特别是在分布外的题目上。
另外还有一个小细节:每一轮策略模型微调完成后,最好在固定的评测集上跑一次完整的滚动评测(不借助MCTS,纯自回归生成),记录一个“内化程度”指标。你会发现,模型纯生成能力通常是随着迭代轮次先升后降的——升是因为吸收了搜索出的优秀轨迹,降是因为过度拟合了训练问题的分布。把握住这个拐点,就该减小学习率或提前停了。整个过程我跑过20轮,收益最大的通常在前5轮,后面基本是边际收益递减,因此时间紧的话跑5轮就足够了。
4. 常见问题与排查技巧实录
4.1 搜索树爆炸、节点过多导致显存/内存溢出
这是最容易遇到的问题。策略模型每生成一个分支都是一次前向推理,如果分支数、模拟步数设置不合理,搜索树的规模会指数上涨。我遇到过一棵树状态数超过5万的情况,系统直接OOM。排查时建议按下面的优先级去查:先看树的深度是不是超过了任务本身需要的推理步数(数学题一般15步内,超过就有问题);再查看每层是否有重复状态(同一前缀的多条路径实际上到达了同一个推理状态,应该合并节点);最后再考虑调低每个节点的分支数。
一个很实用的优化是状态去重,只要你把“当前推理序列”作为状态的唯一标识,并把相同状态的节点合并,树的大小通常能降一半以上。
4.2 PRM分数分布不健康
PRM如果训练得不好,直接表现是搜索时某些分支的分数异常。比如所有正样本分支的分数都在0.9以上,负样本也都在0.7以上,这种“区分度缺失”的问题基本是训练数据正负样本比失衡。我自己的经验是,调节方式不是简单地把负样本翻倍,而是要在抽样时做到“硬负样本挖掘”——也就是只挑那些“看起来快到答案了但其实错了”的状态作为负样本,而不是把所有失败路径全丢进去。相当于只挑那些能真正干扰PRM的样本训练,训练出的PRM更有判别力。
4.3 MCTS探索失效,搜索过程等价于贪心解码
如果你发现不管跑多少轮,搜索树最后选中的路径和模型直接生成的结果几乎一模一样,说明探索项系数太小,或者模拟阶段的奖励信号太稀疏,导致搜索过程缺乏有效引导。
排查方法也很简单:统计一下搜索树中“最终正确答案所在路径”的深度分布。如果有相当一部分正确答案路径是出现在第3层或更深的节点上,但你搜索时经常在第2层就锁定了高价值节点不再往下探索,那就是探索不足。适当提高探索系数,并让模拟阶段的评估更宽泛一些——不要只看当前步骤是否自洽,还要结合后续几步的推演。
4.4 微调后模型基础能力下降
这个现象在迭代后期经常出现,尤其是你直接在MCTS轨迹上微调原模型时。原因在于MCTS轨迹的风格和模型原本的token分布有偏移——搜索过程中模型可能生成了一些概率极低但恰好导向正确答案的序列,你把这些作为标准答案去训练,会让模型强行去拟合低概率token,反而损害了基础语言能力。
解决方案有两个:一是训练集混合一部分原始数据(原始指令微调数据、通用语料等),比例控制在10%到20%即可;二是在解码时降低对MCTS轨迹的拟合强度,例如用平滑标签训练。我后来跑的训练流程里,每次微调都强制混入10%的通用数据,模型基础能力几乎没有明显回退。
4.5 训练数据多样性的坑
盲目追求正确轨迹会带来另一个隐患:模型学会的是一种“单一的正确路径”。数学题往往有多种解法,如果搜索阶段采样温度太低,每次都生成同一种风格的过程,模型就会变得“偏科”。我在实验里对比过,采样温度从0.3升到0.7之后,树中出现的不同解题路径数量明显增加,最终微调出来的模型在评测集上的表现直接涨了一截。所以强烈建议:搜索阶段的采样温度可以比生成阶段的温度高一些,牺牲一点单步质量换取整体多样性,是划算的。
5. 实操参考与后续扩展
5.1 最小可复现配置参考
如果你想自己跑一版最小实验,可以直接参考我这套经过验证的配置。数据集方面,数学推理类用开源的中小学数学题库即可;基础模型建议7B规模以上,太小的话MCTS的搜索能力会被明显限制。MCTS部分,每棵树的搜索预算大约128到256次rollout,分支数每节点5到8个,探索系数取0.5到1.0之间。PRM模型的训练数据由搜索树自动构建,每轮迭代收集几万到十几万条节点级状态。策略模型微调用常规的SFT损失,学习率取2e-6左右,batch size根据显存来。前两三轮由于搜索树还比较粗糙,模型的收益通常是最明显的;到第五轮以后,就要盯着评测集防止过拟合了。
5.2 从训练到推理:怎么用上System 2能力
训练完成后的部署策略也值得说两句。你最终得到的实际上是一个“搜索增强过的模型权重”,也就是说,推理时你可以有两种用法:一种是直接把模型当普通生成模型用,此时模型的纯生成能力已经比初始版本强,因为它已经把很多搜索中见过的好路径内化成了参数;另一种是保留MCTS搜索框架,推理时再跑一次搜索,用PRM来引导选择,代价是多花几十倍的计算量,换来的上限准确率更高。
我测过数学题的效果对比,纯生成版和搜索版相比,后者在难题上的准确率通常能高出5到10个百分点。所以如果你的场景对延迟不敏感(比如离线批处理、深度推理报告生成),推荐直接保留搜索部署;如果对延迟敏感,可以先把纯生成版压到一个可用的性能基线,再考虑在少数疑难场景触发搜索重排。
5.3 还能怎么扩展
这套思路的扩展空间其实不小。一个方向是把PRM替换成更加细粒度的正确性验证器,比如在一个带执行环境的场景里,让每一步生成的代码真实执行并观察结果,用执行反馈作为过程奖励。这个方向上,搜索树的节点就是程序状态,回溯的信号就是程序是否通过测试。另一个方向是把搜索框架泛化到多模态任务,比如给模型外接一个视觉编码器,搜索时每个节点除了文本推理还有图像状态,PRM同时评估图文混合状态的质量。这些方向目前都还在快速迭代中,但ReST-MCTS给了一个非常好的起点:只要你能定义“状态”和“步骤”,就能把这套自举搜索训练套上去。
我个人实际做下来最深的体会是,这套方法的收益不完全体现在最终的准确率提升上,而是体现在你对“模型到底怎么学会推理”这件事的理解上。当你亲眼看到模型在搜索树上从东试西试到有条不紊,再到纯生成也能走对路径,你会直观地理解System 2能力是怎么被训练出来的。最后再分享一个小技巧:跑MCTS时,每隔一段时间手动看一下搜索树的具体路径,不要只看数字指标。很多问题,比如探索方向偏了、奖励信号错位,肉眼检查树会比看曲线更快发现问题。这个习惯帮我省了不少无用功。