EPIG-Tree：梯度高效强化学习的计算最优分支

# EPIG-Tree：面向梯度高效强化学习的计算最优分支
来源：https://arxiv.org/html/2609.20004
###### 摘要

以分组相对策略优化（GRPO）为例的语言模型奖励强化学习方法，将整个随机轨迹坍缩为单个标量奖励。这种方法简洁且可扩展，但在探索和奖励分配上效率低下：一条轨迹可能包含许多因果决策、恢复尝试和环境随机事件，但每个词元或动作都继承了轨迹级别的优势值。我们将树结构的展开构建视为策略梯度估计的计算分配问题。核心主张是：树分支不应放置在策略仅存在不确定性的位置，而应放置在考虑预期计算成本后，额外分支能最大程度减少策略梯度不确定性的位置。基于局部策略梯度随机变量的全方差分解，我们推导出两条分配法则：新分支降低决策不确定性，而后缀重复采样降低持续不确定性。由此产生的 EPIG-Tree 评分利用已计算的展开来分配分支。它估计占据率和加权得分的不确定性，以及后缀法则：
$$n_e \propto \frac{w_e \|\nabla_\theta \log \pi(a_e|h_e)\| \sigma_e}{\sqrt{c_e}}$$
实证研究表明，EPIG 在克隆状态控制中降低了梯度均方误差，在十三个环境中的九个稠密连续控制环境中均取得优势，并近乎完美地恢复了参考梯度方向；同时，相对于更简单的熵分支法，它改善了冻结大语言模型的梯度校准效果。在在线单轮数学任务中，树局部信用分配优于平铺式 GRPO，但分支放置的重要性次于词元级信用分配。在在线多轮 Wordle 游戏中，随着训练进行，EPIG 达到了最高的最终胜率（0.850），超越了平铺式 GRPO（早期即饱和于 0.790）和熵分支法，证实了梯度估计优势可迁移到有状态、大动作空间的设置中。

## 1 引言

基于可验证奖励的强化学习近期成功，使得 GRPO 等策略内算法成为大语言模型推理训练的核心[1]。GRPO 为同一提示采样一组响应，在组内归一化终端奖励，并采用 PPO 风格的裁剪目标，无需训练单独的评论家。

然而，在长时域或多轮设置中，GRPO 存在三种结构性失效模式。

**恢复信用分配不佳**。假设一条轨迹在执行正确动作时出现错误并失败。基于奖励的 GRPO 会为整条轨迹分配一个负的组相对奖励，导致正确动作也受到惩罚。反之，如果一条轨迹尽管包含多个错误动作但最终成功，错误动作会被强化。因此，恢复行为并非长时域信用分配的特例，而是为整条路径分配单个标量的直接结果。

**对非确定性的处理效率低下**。在随机环境中，终端奖励是 $Q(h,a) = \mathbb{E}[R|h,a]$ 的带噪声观测。单条轨迹将策略选择与环境随机性混为一谈。如果环境随机破坏了一个好动作，GRPO 会惩罚该动作；如果环境拯救了一个坏动作，GRPO 会强化它。要估计条件持续价值，需要从相同状态重复进行反事实后缀展开。

**随展开次数增加的扩展性差**。增加独立全轨迹的展开次数可以探索更多，但许多展开仅在低重要性词元、格式或环境噪声上有所不同，这些并不能为模型学习提供有效信号。

树展开是一种自然的补救方法。我们不再采样 $K$ 条独立的全轨迹，而是重用前缀，在中间状态分支，并估计后代价值。这产生了局部的、类似过程的优势值，而无需训练单独的过程奖励模型。近期的大语言模型树方法沿袭了这一方向：TreeRL 的 EPTree 从高熵中间词元分叉[2]；TreePO 使用局部不确定性和分段树建模来分摊公共前缀的成本[3]；相关的树结构展开方法则在多轮智能体任务中推导相对优势。尚未解决的问题是：*树应该在何处分支？*

常见的答案是熵。熵是有用的：如果策略在某个前缀处是确定性的，其分支将重合。但熵本身并非策略梯度估计器所关注的量。在奖励等价的措辞间进行高熵选择不值得分支；而在高价值和低价值动作间进行中等熵的选择可能至关重要。因此，本文的组织原则是：
> 树是用于估计策略梯度的实验装置。最优分支将计算分配在单位成本下，额外分支能最大程度减少梯度估计器误差的位置——而非策略仅存在不确定性的位置。

具体而言：GRPO 将一条完整轨迹坍缩为一个标量优势值，而树则估计局部条件价值；熵仅衡量分支是否可能不同，而 EPIG 衡量的是它们对梯度是否重要。

**贡献**：
1. 我们推导出全方差分解，表明树构建承担两项不同任务：新分支降低决策不确定性，而后缀重采样降低持续不确定性。
2. 我们证明了计算最优的后缀分配法则 $n_e^* \propto \frac{w_e \|\nabla \log \pi(a_e|h_e)\| \sigma_e}{\sqrt{c_e}}$ 以及基于占据率和加权得分不确定性的边际分支法则。
3. 我们引入 EPIG-Tree 算法，该算法使用熵作为提议机制，但根据对梯度的预期预测信息增益来分配分支。
4. 我们在克隆状态连续控制、冻结大语言模型梯度校准、单轮大语言模型数学和在线多轮 Wordle 中验证了该理论，确立了 EPIG 的优势场景以及与更便宜基线持平的边界条件。

## 2 相关工作

**策略梯度与 PPO**。我们的推导始于标准策略梯度定理[4] 和 PPO 风格的裁剪优化[5]。新颖之处不在于新的策略梯度恒等式，而在于利用该恒等式将树构建表述为梯度的最优采样问题。

**GRPO 与 RLVR**。DeepSeekMath 引入了 GRPO，作为一种用于数学推理的无需评论家的 PPO 变体[1]，用组相对归一化替代了学习得到的价值基线。我们的批评是，组归一化的终端奖励在长时域设置中是局部优势值的粗糙估计器。

**大语言模型强化学习的树搜索**。TreeRL 提出了 EPTree，一种熵引导的树搜索，从高不确定性词元分叉，并从后代正确性推导过程监督[2]。TreePO 类似地使用局部不确定性和分段树建模及前缀分摊[3]。这些方法启发了我们的研究场景；我们用梯度估计目标替代了熵/不确定性启发式方法。

**大语言模型强化学习中的熵**。熵机制分析认为，推理模型的强化学习以策略熵换取下游性能，从 logit 更新与动作概率/优势值的协方差推导出熵变化[6]。互补工作发现，少数高熵词元充当推理分叉，并承载了大部分 RLVR 更新信号[7]。我们同意熵能识别可分支或对更新敏感的点。我们的结果更窄更精确：熵不是目标；它是计算分配问题的输入之一。

**奖励过度优化**。Gao, Schulman 和 Hilton 研究了最优奖励作为优化距离 $d = \sqrt{D_{\mathrm{KL}}(\pi\|\pi_0)}$ 的函数，发现了 best-of-n 和 RL 优化的不同缩放形式，以及代理奖励下的古德哈特式退化[8]。我们包含了一个针对代理奖励设置的古德哈特修正版 EPIG-Full 变体，而我们最强的证据涉及精确或克隆状态的梯度估计。

## 3 从结果 GRPO 到局部梯度实验

考虑一个历史或前缀 $h$，一个动作 $a \sim \pi_\theta(\cdot|h)$，以及一个效用
$$Y = U\bigl(G(h) + R_{\mathrm{future}}\bigr),$$
其中 $G(h)$ 是已累积的奖励，$U$ 是训练效用（线性回报、成功/失败、偏好效用、验证器分数或特定任务奖励）。定义
$$Q_h(a) = \mathbb{E}[Y|h,a], \quad \sigma_h^2(a) = \mathrm{Var}(Y|h,a), \quad \psi_h(a) = \nabla_\theta \log \pi_\theta(a|h).$$
对策略梯度的局部贡献为
$$g_h = \mu_h \mathbb{E}_{a \sim \pi_h} \bigl[\psi_h(a) Q_h(a)\bigr],$$
其中 $\mu_h$ 是前缀的占据率或训练权重。对于大语言模型词元前缀，$\mu_h$ 捕获当前策略到达该前缀的频率，或我们分配给它的训练质量；对于克隆的 Gym/MuJoCo 状态，$\mu_h$ 可能在采样状态上是均匀的。

仅基于结果的 GRPO 用单条轨迹标量估计许多局部项。对于一个提示的 $K$ 个采样完成 $\tau_i$，组优势值为
$$\widehat{A}_{i}^{\mathrm{GRPO}} = \frac{R_{i} - \bar{R}}{s_{R} + \varepsilon},$$
轨迹 $i$ 上的每个动作都接收相同的符号和比例。树估计器则通过来自共享前缀的后代展开来估计 $Q_h(a)$：
$$\widehat{Q}_h(a) = \frac{1}{n_{h,a}} \sum_{j=1}^{n_{h,a}} Y_{h,a,j}, \qquad \widehat{A}(h,a) = \widehat{Q}_h(a) - \widehat{V}(h).$$
现在的数学问题是：给定计算预算，我们应将下一个后缀展开或分支花费在何处？

## 4 决定树结构的方差分解

令单样本局部梯度贡献的随机向量为
$$Z_h = \mu_h \psi_h(A) Y, \qquad A \sim \pi_h.$$
其总方差分解如下。

**命题 1（决策与持续不确定性）**。对于固定的前缀 $h$，
$$\begin{aligned}
\mathrm{Var}(Z_h \mid h) &= \mu_h^2 \mathrm{Var}_{A \sim \pi_h} \bigl[\psi_h(A) Q_h(A)\bigr] \\
&\quad + \mu_h^2 \mathbb{E}_{A \sim \pi_h} \bigl[\psi_h(A) \psi_h(A)^\top \sigma_h^2(A)\bigr].
\end{aligned}$$
取迹可得一个标量梯度均方误差的代理项。

**证明**。应用全方差公式，$\mathrm{Var}(Z_h \mid h) = \mathrm{Var}(\mathbb{E}[Z_h \mid A,h] \mid h) + \mathbb{E}[\mathrm{Var}(Z_h \mid A,h) \mid h]$。由于 $\mathbb{E}[Z_h \mid A=a,h] = \mu_h \psi_h(a) Q_h(a)$ 且 $\mathrm{Var}(Y \mid h,a) = \sigma_h^2(a)$，结论成立。$\square$

这个分解是本文的核心。第一项是*决策不确定性*：关于哪个动作分支重要的不确定性。第二项是*持续不确定性*：分支动作已选定后后缀的噪声。因此，树构建包含两种不同的计算操作：
- • 添加从 $h$ 出发的新分支，以减少 $\mathrm{Var}_a [\psi_h(a) Q_h(a)]$；
- • 在现有边 $(h,a)$ 下添加后缀样本，以减少 $\sigma_h^2(a)$。

### 4.1 为什么熵是不够的

熵是 $H_h = H(\pi_\theta(\cdot|h))$：它衡量策略能产生多少种替代方案。但方差分解中没有独立的熵项。一个具有几乎恒定 $Q_h(a)$ 的高熵前缀，其决策不确定性对梯度影响很小；而一个 $Q_h$ 分布较广的中等熵前缀可能主导梯度均方误差。因此，熵是一个有用的提议机制——它告诉我们哪些地方的分支可能不同。EPIG 则追问这些差异是否重要。

## 5 最优分配法则

### 5.1 后缀分配

设边 $e=(h,a)$ 的训练权重为 $w_e$，得分向量为 $\psi_e = \nabla \log \pi(a|h)$，持续标准差为 $\sigma_e$，后缀成本为 $c_e$。如果在边 $e$ 下分配 $n_e$ 个独立后缀，其迹方差贡献近似为 $A_e/n_e$，其中
$$A_e = w_e^2 \|\psi_e\|^2 \sigma_e^2.$$
**定理 1（成本敏感的后缀分配）**。对于固定的候选边和预算 $B$，求解
$$\min_{n_e > 0} \sum_e \frac{A_e}{n_e} \qquad \text{s.t.} \qquad \sum_e c_e n_e \leq B$$
的解为
$$n_e^* \propto \sqrt{\frac{A_e}{c_e}} = \frac{w_e \|\psi_e\| \sigma_e}{\sqrt{c_e}}.$$
**证明**。拉格朗日函数为 $\mathcal{L}(n, \rho) = \sum_e A_e/n_e + \rho(\sum_e c_e n_e - B)$。令导数为零得 $-A_e/n_e^2 + \rho c_e = 0$，故 $n_e = \sqrt{A_e/(\rho c_e)}$；归一化到预算后即得比例关系。$\square$

该法则是策略梯度的奈曼分配类似物：在梯度杠杆高、后缀噪声大且成本低的地方进行更多采样。

### 5.2 分支分配

令 $m_h$ 为节点 $h$ 处已采样的动作分支数，并定义决策不确定性系数
$$B_h = \mu_h^2 \mathrm{tr}\!\left(\mathrm{Var}_{a \sim \pi_h}[\psi_h(a) Q_h(a)]\right).$$
对于 $m_h$ 个独立分支样本，动作蒙特卡洛项按 $B_h/m_h$ 缩放，因此添加一个分支的离散增益为
$$\Delta_h^{\mathrm{branch}} \approx \frac{B_h}{m_h(m_h + 1)} - \lambda c_h,$$
其中 $c_h$ 是预期的分支/后缀成本，$\lambda$ 是计算价格。实践中，$Q_h$ 和 $\psi_h$ 从试验分支中估计，EPIG 使用经验评分
$$S_{\mathrm{EPIG}}(h) = \frac{\widehat{\mu}_h^{\,2}\,\mathrm{tr}\,\widehat{\mathrm{Var}}_{a \in \mathcal{C}(h)}\!\left[\widehat{\psi}_h(a) \widehat{Q}_h(a)\right]}{\widehat{c}_h(m_h + 1)^2 + \varepsilon}.$$