TreeRL: LLM Reinforcement Learning with On-Policy Tree Search¶

会议: ACL 2025
arXiv: 2506.11902
代码: https://github.com/THUDM/TreeRL
领域: LLM推理
关键词: 树搜索, 强化学习, 过程监督, 熵引导分叉, LLM推理

一句话总结¶

提出 TreeRL，将基于熵引导的树搜索（EPTree）直接集成到 LLM 的 on-policy 强化学习训练中，通过在高不确定性 token 处分叉来扩展推理路径多样性，并利用树结构提供的全局+局部优势作为过程监督信号，在数学和代码推理任务上超过传统的多链采样 RL。

研究背景与动机¶

领域现状：当前 LLM 推理能力的强化学习（RL）方法主要采用独立多链采样（i.i.d multi-chain sampling），对每个 prompt 独立生成 K 个完整回答，基于最终答案的正确性进行奖励反馈和策略更新（如 GRPO、RLOO）。
现有痛点：
独立多链采样生成的回答多样性有限，难以充分探索推理空间
outcome-level 的奖励信号太稀疏，只在回答末尾给一次反馈，中间步骤缺乏梯度信号
虽然有 Process Reward Model (PRM) 提供中间步骤监督，但离线训练的 PRM 面临分布偏移和 reward hacking 问题
核心矛盾：树搜索在 AlphaZero 等传统领域已被验证有效，但在 LLM RL 中仍未被充分利用。现有尝试要么只在推理阶段使用树搜索提升表现，要么用树搜索生成离线训练数据（SFT/DPO），而非直接在 on-policy RL 训练中使用。经典 MCTS 逐步生成的方式需要大量迭代，对 VLLM 等现代推理引擎不友好，效率低下。
本文要解决什么？
如何设计一种高效的树搜索算法，在相同推理预算下比多链采样和 MCTS 产生更多样的正确回答
如何将树搜索的结构信息转化为可靠的 on-policy 过程监督信号用于 RL 训练
切入角度：模型在生成过程中某些 token 的预测熵（不确定性）特别高——这些位置是推理路径分叉的关键点。在这些高不确定性的 token 处分叉生成新分支，既高效又能最大化探索的多样性。
核心 idea 一句话：用 token 级别的熵来引导树搜索分叉，配合树结构推导的 on-policy 过程监督奖励，直接将树搜索集成到 LLM 强化学习训练中。

方法详解¶

整体框架¶

TreeRL 的 pipeline：给定 prompt → 用 EPTree 生成推理树（包含多个分支和叶节点）→ 对每个叶节点验证答案正确性 → 基于树结构计算每个中间步骤的过程监督信号 → 用 policy gradient 更新策略模型。整个过程是 on-policy 的：每轮 RL 迭代都重新生成树，奖励信号直接来自当前策略的生成结果。

关键设计¶

EPTree (Entropy-guided Parallel Tree search):
做什么：一种高效并行的树搜索算法，在相同推理预算下生成比 MCTS 和多链采样更多样的回答
核心思路：
- 初始化：并行生成 \(M\) 个独立链作为初始树
- 分叉点选择：计算每个 token 的交叉熵 \(H(\mathbf{y}_t) = -\log \pi_\theta(\mathbf{y}_t | \mathbf{x}, \mathbf{y}_{<t})\)，在整棵树中选择熵最高的 top-\(N\) 个 token 作为分叉点（同时 mask 掉接近序列末尾的 token）
- 扩展：从每个分叉点的前缀继续生成 \(T\) 个新分支直到完成
- 重复分叉和扩展 \(L\) 次，最终得到 \(M \times (T \times N \times L + 1)\) 个叶节点
设计动机：MCTS 需要逐步展开非常多的迭代，效率低；EPTree 只需约 2 轮迭代就能构建多样化的搜索树。高熵 token 是模型最"纠结"的位置（如数学运算符、逻辑连词、"wait" 等自我反思词），从这些位置分叉能最大化路径多样性
关键参数记为 \((M, N, L, T)\)-tree，例如 \((6,2,1,2)\) 表示 6 棵并行树，每轮选 2 个分叉点，1 轮迭代，每个分叉点扩展 2 个分支
On-policy 过程监督信号:
做什么：利用树结构为每个中间推理步骤计算奖励信号
核心思路：
- 每个节点 \(s_n\) 的 value：其所有叶子后代中正确答案的比例 \(V(s_n) = \frac{1}{|L(s_n)|} \sum_{l \in L(s_n)} \mathbf{1}(l \text{ correct})\)
- 全局优势 (Global Advantage)：\(G_A(s_n) = V(s_n) - V(\text{root})\)，衡量该步骤相对于整体平均正确率的优势
- 局部优势 (Local Advantage)：\(L_A(s_n) = V(s_n) - V(p(s_n))\)，衡量该步骤相对于父节点的改进
- 最终奖励：\(R(s_n) = (G_A(s_n) + L_A(s_n)) / \sqrt{|L(s_n)|}\)
设计动机：全局优势提供对步骤质量的宏观评估，局部优势捕捉每一步的边际贡献。除以 \(\sqrt{|L(s_n)|}\) 是因为共享前缀的非叶节点会在多条序列中重复出现，需要降权防止过拟合。这种设计可以看作 GAE (Generalized Advantage Estimation) 的特殊形式
与 PRM 的关键区别：完全 on-policy，不需要额外训练 reward model，天然抗 reward hacking
树到序列的训练:
做什么：将树结构展开为多条序列用于 RL 训练
核心思路：每个叶节点对应一条从根到叶的完整序列，树中共有 \(M \times (N \times L \times T + 1)\) 条序列。每条序列中每个 token 都带有对应步骤的过程奖励 \(R(s_n)\)，用 policy gradient 更新
设计动机：相比多链采样 16 条独立链，TreeRL 的 \((6,2,1,2)\) 配置在相同推理预算下产生约 30 条序列，提供更多训练数据（虽然训练计算量更大）

损失函数 / 训练策略¶

使用标准 policy gradient，KL 系数 \(\beta = 10^{-4}\)，学习率 \(1.5 \times 10^{-6}\)
采样温度 1.2，top-p = 0.95，最大序列长度 8192
规则型奖励：答案正确 +1，错误 0
训练数据来自 MATH-train 和 NuminaMath

实验关键数据¶

主实验¶

模型	MATH500	Omni-MATH-500	AIME2024	AMC	OlympiadBench	LiveCodeBench	Avg
Qwen-2.5-14B SFT	76.6	29.5	10.6	48.0	36.9	14.5	36.0
ChainRL (Qwen-14B)	81.6	32.7	22.2	53.9	41.1	18.2	41.6
TreeRL (Qwen-14B)	81.7	36.7	28.0	55.9	44.6	20.8	44.5
R1-Distilled-7B SFT	94.0	47.8	55.9	85.5	54.4	43.9	63.6
ChainRL (R1-7B)	93.6	48.1	59.7	85.5	54.5	46.1	64.5
TreeRL (R1-7B)	94.4	49.8	60.8	85.0	57.1	47.4	65.8

TreeRL 在 Qwen-14B 上平均比 ChainRL 高 2.9%，在 R1-Distilled-7B 上高 1.3%。

消融实验¶

过程奖励设计	#Response	MATH500	Omni-MATH	AIME2024	Avg Gain
\((G_A + L_A)/\sqrt{n}\)（完整）	30	81.7	36.7	28.0	—
\(G_A + L_A\)（无降权）	30	81.5	32.0	24.1	-1.9 ↓
\(G_A/\sqrt{n}\)（无局部优势）	30	80.1	35.1	24.7	-1.3 ↓
\((G_A + L_A)/\sqrt{n}\)（16条）	16	80.1	32.5	24.5	-3.2 ↓

分叉策略 (M,N,L,T)	熵引导	PassRate ↑	Token数 ↓
(6,2,1,2) EPTree	✓	56.9	22268
(6,2,1,2) 随机分叉	✗	54.8	24213
(16,0,0,0) 多链采样	-	52.4	19858

关键发现¶

EPTree 在相同推理预算下 PassRate 持续优于多链采样约 3%，同时生成约 2 倍多样的不同回答
过程监督三要素缺一不可：去掉降权因子（\(1/\sqrt{n}\)）掉 1.9%、去掉局部优势掉 1.3%、减少训练序列数掉 3.2%
熵引导分叉比随机分叉在更少 token 下取得更高 PassRate（56.9 vs 54.8），说明高熵位置确实是关键分叉点
高频分叉 token 包括数学运算符 (、逻辑连词、转折词 "Since"/"But" 以及自我反思词 "wait"，符合直觉
TreeRL 在训练约 100 步后开始优于 ChainRL 并持续拉大差距

亮点与洞察¶

熵引导分叉是最巧妙的设计：避免了 MCTS 逐步展开的低效问题，只需 2 轮迭代建树。利用 token 熵找到模型最"纠结"的位置进行分叉，既高效又有针对性，是一种很自然的不确定性引导探索策略
on-policy 过程监督无需训练 PRM：直接从树结构的正确率统计推导出梯度信号，避免了离线 PRM 的分布偏移问题。这种方法可以迁移到任何有 outcome reward 的 RL 场景
全局+局部优势的组合可视为 GAE 的特殊形式，提供了理论基础，且 \(1/\sqrt{n}\) 降权是一个值得关注的技术细节——共享前缀在多条序列中重复出现的过拟合问题是树结构 RL 的独有挑战

局限性 / 可改进方向¶

当前推理引擎（如 VLLM）未对树搜索做特殊优化，EPTree 仍需 2+ 轮迭代，实际速度约为多链采样的 1/2
实验仅在数学和代码推理上验证，未扩展到 NLP 生成、对话等任务
\((G_A + L_A)/\sqrt{n}\) 的奖励设计较为启发式，步骤重要性权重 \(\lambda_j\) 的最优分配尚未充分探索
R1-Distilled 模型上 RL 增益较小，可能因为训练数据对该模型太简单——训练数据难度匹配是 RL 的关键问题
可以探索将 EPTree 与 MCTS 的 UCB 策略结合，或与 reward model 联合使用

评分¶

新颖性: ⭐⭐⭐⭐ 树搜索+RL 的 idea 很多人提过，但 EPTree 的熵引导分叉和 on-policy 过程监督的具体设计有创新
实验充分度: ⭐⭐⭐⭐ 多个基准、多个基座模型、详细的消融，但缺少非推理任务验证
写作质量: ⭐⭐⭐⭐ 结构清晰，图表丰富，算法描述规范
价值: ⭐⭐⭐⭐ 为 LLM RL 训练提供了一种实用的树搜索方案，开源代码降低了复现门槛