基于梯度外推的策略优化

# 基于梯度外推的策略优化 (GXPO)

来源: https://arxiv.org/html/2605.06755

**作者:** Ismam Nur Swapnil$^1$, Aranya Saha$^2$, Tanvir Ahmed Khan$^3$, Mohammad Ariful Haque$^1$, Ser-Nam Lim$^4$

**机构:**
1. Bangladesh University of Engineering and Technology
2. University of Maryland, College Park
3. Illinois Institute of Technology
4. University of Central Florida

###### 摘要

强化学习被广泛用于提升大语言模型（LLM）的推理能力，尤其是在答案可以自动验证的场景中。标准的 GRPO 风格训练仅利用当前步骤更新模型，而完整的多步前瞻（multi-step lookahead）虽然能提供更好的更新方向，但由于需要大量的反向传播，计算成本过高。我们提出了**基于梯度外推的策略优化（GXPO, Gradient Extrapolation-Based Policy Optimization）**，这是一种适用于 GRPO 风格推理 RL 的即插即用式策略更新规则。GXPO 在活跃阶段仅使用三次反向传播来近似更长的局部前瞻。它重用相同的 rollout 批次、奖励、优势（advantages）和 GRPO 损失，因此在前瞻点不需要新的 rollout 或奖励计算。GXPO 执行两次快速的优化器步骤，测量梯度的变化，预测一个虚拟的 $K$ 步前瞻点，将策略部分移动向该点，然后在新的位置使用真实梯度应用纠正性更新。当前瞻信号变得不稳定时，GXPO 会自动切换回标准的单次 GRPO。我们还提供了一个基于普通梯度下降的代理分析，解释了外推何时精确以及其局部误差的来源。在 Qwen2.5 和 Llama 数学推理实验中，GXPO 相较于 GRPO 将平均采样 pass@1 提升了 $\mathbf{+1.65}$ 到 $\mathbf{+5.00}$ 个点，相较于最强的 SFPO 设置提升了 $\mathbf{+0.14}$ 到 $\mathbf{+1.28}$ 个点，同时保持活跃阶段的成本固定在三次反向传播。它还实现了高达 $\mathbf{4.00\times}$ 的步骤加速、$\mathbf{2.33\times}$ 的挂钟时间加速，以及在达到 GRPO 峰值精度时 $\mathbf{1.33\times}$ 的反向传播加速。

## 1. 引言

策略梯度强化学习构成了现代语言模型对齐和推理的基础 (Sutton et al., 1999; Williams, 1992; Schulman et al., 2017)。在可验证奖励的强化学习（RLVR）中，模型生成候选解，接收可验证的奖励，并使用一阶梯度更新策略 (Shao et al., 2024; Yu et al., 2025)。这一设置在 MATH 和 GSM8K 风格的推理基准测试中至关重要，因为长文本生成后的正确性可以自动检查 (Hendrycks et al., 2021; Cobbe et al., 2021)。在 LLM 规模下，每增加一次反向传播都会倍增训练时间和内存消耗，使得更新质量与重复后训练迭代中的每步成本之间产生张力。

> **图 1 说明：** GXPO 整体训练框架。每个活跃步骤执行三次反向传播：两次在梯度外推期间探测梯度，一次在重新定位点 $\tilde{\theta}^t$ 处计算纠正梯度 $g_{\text{slow}}$。缓慢纠正始终在当前步骤应用。更新后，$z$-score 门控检查 $\|g_{\text{slow}}\|$；如果 $Z > \tau$，所有后续步骤永久回退到单次 GRPO。

像 PPO 和 GRPO 这样的单步方法避免了这种成本，但仅使用当前策略的梯度 (Schulman et al., 2017; Shao et al., 2024)。显式前瞻通过付出代价来恢复轨迹信息：$h$ 步策略镜像下降改善了正则化策略迭代 (Protopapas and Barakat, 2024)，树扩展策略降低了策略梯度方差 (Dalal et al., 2025)，在线或自适应规划使用学习模型的 rollout 进行动作选择 (Sikchi et al., 2021; Rosenberg et al., 2023)。这些方法依赖于规划、树扩展或辅助价值/模型估计，而不是手头已有的 rollout 批次，这使得它们难以在不改变数据路径的情况下直接嵌入标准的 GRPO 风格推理流水线。

最近的 LLM 推理 RL 工作通过目标函数、价值估计、熵控制、过滤、奖励设计、配方或离策略重用提高了稳定性和效率 (Liu et al., 2025; Dai et al., 2025; Yue et al., 2025; Xiong et al., 2025; Cui et al., 2025; Wen et al., 2025; Fatemi et al., 2025; Shen et al., 2025; Mroueh, 2025; Mroueh et al., 2025)，以及通过样本选择、降采样、选择性 rollout、少样本训练、测试时缩放或 RLHF/生成系统加速 (Chen et al., 2024; Xia et al., 2024; Ye et al., 2025; Li et al., 2025; Xu et al., 2026; Wang et al., 2025; Zheng et al., 2025; Muennighoff et al., 2025; Sheng et al., 2025; Kwon et al., 2023; He et al., 2025)。

GXPO 是正交的：它保持 rollout 批次、奖励、优势和 GRPO 损失不变，仅改变策略更新。这使得更新规则成为一种窄干预，而不是新的推理 RL 训练配方或奖励流水线组件。

优化器侧的前瞻更接近我们的设置。Lookahead 优化器将快速的内部步骤与缓慢的更新交错 (Zhang et al., 2019; Zhou et al., 2021)，而 SFPO 将这一思想移植到策略优化中，在一次缓慢纠正之前执行 $K$ 个快速内部步骤 (Wang et al., 2026)——每次更新需要 $K+1$ 次反向传播。

我们要解决的问题是：策略更新是否可以在不增加反向传播次数的情况下，获得类似于显式前瞻方法的梯度轨迹信息？我们介绍了**基于梯度外推的策略优化（GXPO）**，这是一种与 GRPO 兼容的更新方法，用三次反向传播近似 $K$ 个局部策略梯度步骤，且独立于虚拟深度 $K$（见图 1）。GXPO 重用 rollout 批次、奖励、优势和正则化项；从两个探测梯度估计每个坐标的保留率 $r_i = g_{1,i}/g_{0,i}$；向几何位移方向移动；并在重新定位的策略上应用纠正梯度，因此最后一步仍然锚定在真实目标上，而不是外推预测上。滚动的 $z$-score 门控在训练期间纠正梯度范数变得不稳定时回退到单次更新。

我们的贡献包括：
- 一种与 GRPO 兼容的更新方法，对于任意虚拟前瞻深度 $K$，在活跃阶段仅使用三次反向传播，利用两个探测梯度、几何外推和一个纠正梯度；
- 一种固定批次的实现，重用 rollout、奖励、优势、损失和正则化项，并包含一个门控机制，当局部轨迹信号不可靠时回退到基础单次更新；
- 局部二次代理分析，以及跨两个模型系列的基准、预算、消融和诊断证据。

## 2. 方法：基于梯度外推的策略优化

GXPO 用三步更新替换单次 GRPO 更新，同时重用相同的 rollout、奖励、优势和目标，因此不需要额外的数据或奖励计算。在训练期间，它首先使用基础 actor 优化器（在我们的实验中为 AdamW (Loshchilov and Hutter, 2019)）执行两次快速优化步骤，并观察参数如何变化。然后，它使用这种变化来估计更新方向，部分地朝该方向移动，并应用最终的纠正步骤。理论部分研究了一个简化版本以阐明这种行为，而实验部分则在所选优化器下评估该方法。

### 设置与符号

令 $\theta \in \mathbb{R}^d$ 参数化策略，令 $\mathcal{L}(\theta)$ 为 GRPO 损失。记 $g(\theta) = \nabla_\theta \mathcal{L}(\theta)$, $H(\theta) = \nabla_\theta^2 \mathcal{L}(\theta)$，以及 $g_n = g(\theta_n)$。令 $\eta > 0$ 为学习率，$K$ 为虚拟深度。

### 2.1 从泰勒展开到几何缩放

GXPO 需要一种廉价的方法来估计梯度在附近几个步骤中会如何变化。局部直觉很简单：如果损失曲率在 $\theta_0$ 附近大致稳定，那么梯度演化可以通过可预测的局部递推来近似。在逐坐标代理下，这意味着每个梯度坐标保留其前一个值的可测量部分。泰勒展开使这种直觉变得精确。

在 $\theta_0$ 周围，
$$
g(\theta_0 + \Delta) = g(\theta_0) + H(\theta_0)\Delta + R_2(\Delta), \quad \|R_2(\Delta)\| \leq \frac{M_3}{2} \|\Delta\|^2, \quad (1)
$$
其中 $M_3 = \sup_\xi \|\nabla^3 \mathcal{L}(\xi)\|$。忽略余项得到局部二次模型 $g(\theta_0 + \Delta) \approx g_0 + H_0 \Delta$。

###### 假设 1（局部二次模型）

Hessian 在局部外推区域内固定为 $H_0$。此假设仅用于外推：较小的学习率保持探测的局部性，并且最终更新仍然使用重新定位点的真实梯度。

在普通 GD 代理中，对于一步梯度下降 $\theta_1 = \theta_0 - \eta g_0$，模型给出：
$$
g_1 = g_0 - \eta H_0 g_0 = (I - \eta H_0) g_0. \quad (2)
$$
重复此递推得到：

###### 定理 1（局部二次模型下的梯度演化）

在假设 1 下，第 $n$ 次梯度下降迭代 $\theta_n = \theta_{n-1} - \eta g_{n-1}$ 处的梯度满足：
$$
g_n = (I - \eta H_0)^n g_0. \quad (3)
$$
证明见附录 A.1。

### 2.2 每参数保留率

完整的局部递推涉及 Hessian，这在 LLM 规模下无法构建或相乘。因此，GXPO 直接从两个附近的梯度测量梯度演化。对于每个充分活跃的坐标，比率 $g_{1,i}/g_{0,i}$ 估计在经过一个局部步骤后该坐标梯度的保留程度。这个比率在实现的优化器更新中不被视为精确的 Hessian 量；它是沿实际快速优化器轨迹测量的经验信号。

形式上，对于活跃坐标，GXPO 定义：
$$
r_i \equiv \frac{g_{1,i}}{g_{0,i}}. \quad (4)
$$
在普通 GD 代理中，其中 $\theta_1 = \theta_0 - \eta g_0$，局部二次模型给出 $r_i = 1 - \eta \frac{[H_0 g_0]_i}{g_{0,i}}$。

在有限精度下，算法 1 仅在活跃集 $\mathcal{A}_t = \{i : |g_i^{t,0}| > \delta\}$ 上评估此比率。对于 $i \notin \mathcal{A}_t$，不形成比率，保留观察到的双探测位移。在活跃坐标上，$r_i$ 测量局部梯度保留：$r_i \approx 1$ 表示近乎平坦，$r_i \approx 0$ 表示快速衰减。

算法 1 GXPO: 基于梯度外推的策略优化
1: 参数 $\theta^0 \in \mathbb{R}^d$; 学习率 $\eta$; 虚拟步数 $K$; 混合系数 $\alpha$; 稳定性阈值 $\delta$; 触发阈值 $\tau$;
2: 初始化: 滚动缓冲区 $\mathcal{B} \leftarrow \emptyset$, $s^\star \leftarrow +\infty$
3: for each training step $t = 0, 1, 2, \dots$ do
4:     $g^{t,0} \leftarrow \nabla_\theta \mathcal{L}(\theta^t) \in \mathbb{R}^d$ $\triangleright$ 反向传播 1
5:     if $t < s^\star$ then
6:         $\theta^{t,1} \leftarrow \text{OptimStep}(\theta^t, g^{t,0})$
7:         $g^{t,1} \leftarrow \nabla_\theta \mathcal{L}(\theta^{t,1}) \in \mathbb{R}^d$ $\triangleright$ 反向传播 2
8:         $\theta^{t,2} \leftarrow \text{OptimStep}(\theta^{t,1}, g^{t,1})$
9:         $\mathcal{A}_t \leftarrow \{i \in [d] : |g_i^{t,0}| > \delta\}$ $\triangleright$ 活跃坐标
10:        $r_i^t \leftarrow g_i^{t,1} / g_i^{t,0}, \quad \forall i \in \mathcal{A}_t$ $\triangleright$ 保留率
11:        $S_i^t(n) \leftarrow \frac{1-(r_i^t)^n}{1-r_i^t}, \quad \forall i \in \mathcal{A}_t$ $\triangleright$ 几何和
12:        $\text{scale}_i^t \leftarrow \begin{cases} \frac{S_i^t(K)}{S_i^t(2)}, & i \in \mathcal{A}_t, \\ 1, & i \notin \mathcal{A}_t, \end{cases}$ $\triangleright$ 外推因子
13:        $\theta^{t,K} \leftarrow \theta^t + (\theta^{t,2} - \theta^t) \odot \text{scale}^t$ $\triangleright$ 外推
14:        $\tilde{\theta}^t \leftarrow \theta^t + \alpha(\theta^{t,K} - \theta^t)$ $\triangleright$ 重新定位
15:        $g_{\text{slow}}^t \leftarrow \nabla_\theta \mathcal{L}(\tilde{\theta}^t) \in \mathbb{R}^d$ $\triangleright$ 反向传播 3
16:        $\theta^{t+1} \leftarrow \text{OptimStep}(\tilde{\theta}^t, g_{\text{slow}}^t)$ $\triangleright$ 缓慢纠正
17:    end if
18:    if $\text{len}(\mathcal{B}) > 1$ then
19:        从当前缓冲区 $\mathcal{B}$ 计算滚动统计量 $\mu_t$ 和 $\sigma_t$
20:        $Z_t \leftarrow \frac{\|g_{\text{slow}}^t\|_2 - \mu_t}{\sigma_t + \epsilon}$ $\triangleright$ §2.3 中的自适应规则
21:        if $Z_t \geq \tau$ then
22:            $s^\star \leftarrow t + 1$ $\triangleright$ 永久禁用外推