跳转至

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 产出的高质量代码去微调基座模型)。

关键设计

  1. 知识库作为过程奖励模型:

    • 功能:用检索信号替代训练的 PRM
    • 核心思路:将当前扩展的算法步骤与知识库中的正确步骤做相似度匹配,高相似度 = 步骤质量高 → MCTS 优先探索该分支
    • 设计动机:同类算法的正确实现高度相似(如所有正确的快速排序都会有 partition 步骤)——利用这种同质性
  2. 相似度过滤去冗余:

    • 功能:在 MCTS 扩展阶段去除语义重复的候选节点
    • 核心思路:新生成的扩展节点与已有节点做相似度检查,超过阈值则视为冗余并过滤
    • 设计动机:LLM 的多次采样常产生语义相同但写法不同的代码——保留冗余浪费搜索预算
  3. 沙箱执行+错误反射:

    • 功能:在模拟阶段执行代码发现错误并反馈
    • 核心思路:对 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研究有双重贡献