RPM-MCTS: Knowledge-Retrieval as Process Reward Model with Monte Carlo Tree Search for Code Generation¶
会议: AAAI 2026
arXiv: 2511.19895
代码: 有
领域: LLM推理 / 代码生成
关键词: 过程奖励模型, MCTS, 知识库检索, 代码生成, 算法步骤评估
一句话总结¶
提出 RPM-MCTS——用知识库检索替代训练的过程奖励模型(PRM)来指导代码生成的 MCTS 搜索。利用同类算法实现的同质性,从知识库中检索正确算法步骤作为评估信号,配合相似度过滤去除冗余扩展节点和沙箱执行定位错误,实现 ~15% token 减少同时超越 SOTA。
研究背景与动机¶
领域现状:树搜索(MCTS)用于代码生成已证明有效——通过多步搜索生成更好的算法代码。但核心瓶颈在于如何评估中间算法步骤的质量。
现有痛点: - 训练 PRM 需要大量标注的步骤级数据——收集和标注成本极高 - 现有方法无法有效定位和及时修正错误步骤——发现错误时可能已经浪费了大量搜索预算 - 搜索过程中大量冗余节点——不同扩展产生语义重复的候选
核心矛盾:需要评估中间算法步骤的质量但训练 PRM 太贵——有没有不需要训练的替代方案?
本文目标 用知识库检索替代训练的 PRM,同时解决搜索冗余和错误定位问题。
切入角度:同类算法的实现具有"同质性"——排序算法的正确步骤在不同问题间是相似的。因此可以从正确实现的知识库中检索,作为评估当前步骤质量的参考。
核心 idea:知识库中的正确算法步骤 = 免训练的过程奖励信号 → MCTS 搜索更高效。
方法详解¶
整体框架¶
三组件:(1) 知识库构建(收集各类算法的正确实现);(2) RPM-MCTS 搜索(在每步扩展时检索知识库评估步骤质量 + 相似度过滤去冗余 + 沙箱执行定位错误);(3) 可选的数据蒸馏(用 RPM-MCTS 产出的高质量代码去微调基座模型)。
关键设计¶
-
知识库作为过程奖励模型:
- 功能:用检索信号替代训练的 PRM
- 核心思路:将当前扩展的算法步骤与知识库中的正确步骤做相似度匹配,高相似度 = 步骤质量高 → MCTS 优先探索该分支
- 设计动机:同类算法的正确实现高度相似(如所有正确的快速排序都会有 partition 步骤)——利用这种同质性
-
相似度过滤去冗余:
- 功能:在 MCTS 扩展阶段去除语义重复的候选节点
- 核心思路:新生成的扩展节点与已有节点做相似度检查,超过阈值则视为冗余并过滤
- 设计动机:LLM 的多次采样常产生语义相同但写法不同的代码——保留冗余浪费搜索预算
-
沙箱执行+错误反射:
- 功能:在模拟阶段执行代码发现错误并反馈
- 核心思路:对 MCTS 模拟阶段的候选代码在沙箱中执行测试用例,失败时提取错误信息反馈给 LLM 进行定向修正
- 设计动机:代码生成的优势——可以实际执行来验证,不仅查分而且定位错误
损失函数 / 训练策略¶
- 搜索阶段无需训练——纯推理时方法
- 可选蒸馏:用 RPM-MCTS 产出的代码做 SFT 提升基座模型
实验关键数据¶
主实验(Pass@1)¶
| 方法 | APPS-Intro | APPS-Interv | APPS-Comp | CodeContests | HumanEval+ | MBPP+ | 平均 |
|---|---|---|---|---|---|---|---|
| Qwen3-8B 基线 | 56.7 | 35.3 | 29.3 | 10.7 | 75.6 | 72.2 | 52.1 |
| Qwen3-8B + ToT | 69.3 | 54.0 | 41.3 | 17.3 | 82.3 | 70.4 | 59.0 |
| Qwen3-8B + RPM-MCTS | 77.3 | 60.0 | 43.6 | 23.0 | 83.5 | 76.2 | 64.0 |
| Qwen3-235B 基线 | 78.0 | 54.7 | 42.7 | 24.0 | 86.0 | 78.8 | 64.6 |
| Qwen3-235B + RPM-MCTS | 86.7 | 67.3 | 59.3 | 36.7 | 87.8 | 81.2 | 72.3 |
- RPM-MCTS使Qwen3-8B平均提升+11.9个点(52.1→64.0),达到Qwen3-235B基线水平
- Qwen3-235B + RPM-MCTS在竞赛级任务上提升最大(CodeContests +12.7)
消融¶
| 配置 | 平均Pass@1 | 说明 |
|---|---|---|
| w/o 知识库评估 | 63.5 | 搜索方向不精确,退化为普通MCTS |
| w/o 相似度过滤 | 略低 | 大量冗余扩展浪费搜索预算 |
| w/o 沙箱执行反馈 | 较低 | 无法定位具体错误步骤 |
| 完整 RPM-MCTS | 64.0 | 三组件协同最优 |
| + SFT蒸馏到基座模型 | 基座显著提升 | RPM-MCTS也是高质量SFT数据生成方法 |
Token消耗对比¶
- RPM-MCTS比标准MCTS减少约15% token消耗——归功于相似度过滤和错误步骤截断
关键发现¶
- 知识库检索有效替代训练PRM——零训练成本获得过程级评估信号
- 相似度过滤减少~15% token——去除冗余节点后搜索更聚焦高价值路径
- 错误反射使修正精准化——不只知道"错了"还知道"错在第几步",截断后重新搜索
- 蒸馏可进一步提升基座模型——RPM-MCTS不仅是推理时方法也是数据生成方法
- 8B模型+RPM-MCTS(64.0)达到235B基线水平(64.6)——搜索弥补了30倍参数差距
亮点与洞察¶
- "算法实现的同质性"是核心洞察——利用同类算法的结构相似性作为免费的奖励信号
- 知识库检索 + MCTS 的组合可推广到其他有"正确参考"可检索的领域(数学证明、逻辑推理)
- 三个优化(知识库评估+冗余过滤+错误反射)各解决不同的搜索效率问题
局限与展望¶
- 知识库的覆盖范围决定了方法的适用范围——对全新算法类型无参考可检索
- 相似度阈值需要调优
- 仅在竞赛编程任务上验证,工程代码生成效果未知
- 知识库的维护和更新是额外开销
相关工作与启发¶
- vs MCTS-SQL:同样用 MCTS 做代码搜索。RPM-MCTS 加入了知识库评估维度
- vs PRM (Lightman et al.):需要大量训练数据。RPM-MCTS 零训练
- 知识库即奖励模型的思路对数学推理 PRM 也有启发
评分¶
- 新颖性: ⭐⭐⭐⭐ 知识库检索作为免训练PRM的思路简洁有效
- 实验充分度: ⭐⭐⭐⭐ MCTS对比+消融+蒸馏验证
- 写作质量: ⭐⭐⭐⭐ 三组件设计清晰
- 价值: ⭐⭐⭐⭐ 对代码生成和PRM研究有双重贡献