面向大模型强化学习的预测性散度掩码 Predictive Divergence Masks for LLM RL
用散度方向导数重定义信任域方向判据，改进LLM强化学习稳定性。
前置知识
重要性比 (importance ratio)
off-policy RL 中，当前训练策略 $\pi$ 与采样所用的行为策略 $\mu$ 在采样 token $y_t$ 上的概率之比 $r_t = \pi(y_t|s_t)/\mu(y_t|s_t)$。它用于把行为策略采集的数据“重加权”到当前策略下，是 PPO 裁剪与遮蔽的核心量。
整篇论文的批判对象：基于比值的判据只是单样本代理，无法反映全词表分布的漂移，理解它才能理解为何要换成散度。
信任域掩码 (trust-region mask)
在带裁剪/遮蔽的代理目标中，掩码由两部分组成：接近判据（$|r_t-1|>\epsilon$ 或 $D_t>\delta$）判断是否越界，方向判据（$\text{sign}(\hat{A}_t(r_t-1))>0$）判断更新是否把策略进一步推离行为策略。越界且向外时才遮蔽该 token 的梯度。
论文的修改点正是掩码中的“方向判据”，理解掩码的二元结构是读懂方法的前提。
DPPO 与 top-K 散度
DPPO 把 PPO 的比值接近判据换成基于 top-K 的前向 KL 散度 $D_t = D_{KL}^{TopK}(\mu\|\pi)$，支撑集取 $\mathcal{K}=\text{TopK}(\mu,K)\cup\{y_t\}$，词表剩余部分塌缩成一个聚合尾部桶，top-K 通常捕获 99% 以上概率质量，是对真实 KL 的极紧下界。
DPPO-TopK-KL 是本文的直接基线与改进起点，方法在其基础上只替换方向判据。
方向导数 (directional derivative)
函数沿某方向的变化率。本文计算 KL 散度沿一个策略梯度步方向 $v_i=\mathbf{1}[i=k]-\pi_i$ 的方向导数 $\dot{D}=\frac{d}{d\eta}D_{KL}(\mu\|\pi_\eta)|_{\eta=0}$，其一阶符号判断散度下一步会增还是减。
这是“预测性散度掩码”的核心数学工具，是方法创新的技术载体。
softmax 归一化耦合
softmax 策略 $\pi_i=e^{z_i}/\sum_j e^{z_j}$ 中，改变任一 token 的概率必然通过归一化分母连带改变词表其余 token 的概率。方向导数中的全局项 $\sum_i\pi_i(\mu_i-\pi_i)$ 正是这种分布级耦合的体现。
全文的关键洞察：比值判据丢失的正是这个全局项，本文把它补回来以对齐散度变化。
研究动机
在大语言模型强化学习（LLM RL）中，由于训练栈与推理引擎的数值行为不一致（training-inference mismatch），以及策略陈旧（同一批 rollout 上多次 minibatch 更新），采样数据并非来自当前正在更新的策略，实际训练本质上是 off-policy 的。主流方法（PPO、GRPO、DAPO）在行为策略 $\mu$ 采样的 rollout 上优化带掩码或裁剪的代理目标来稳定更新。PPO 用采样 token 的重要性比 $r_t = \pi(y_t|s_t)/\mu(y_t|s_t)$ 同时决定接近判据 $|r_t-1|>\epsilon$ 与方向判据 $\text{sign}(\hat{A}_t(r_t-1))>0$。由于 $|r_t-1|/2$ 仅是总变差距离的单样本蒙特卡洛估计，比值是噪声很大的分布漂移代理。近期的 DPPO 把接近判据升级为基于 top-K 的前向 KL 散度 $D_t = D_{KL}^{TopK}(\mu\|\pi)$，但方向判据仍沿用 PPO 的比值形式，造成内部不一致：散度本身是全词表分布级的，而方向却只看单个采样 token，其符号可能与 $D_t$ 的真实变化相反。
本文的目标是本文目标是消除这种“接近判据是分布级、方向判据却是单样本”的不一致。作者主张：既然信任域用散度 $D_t$ 定义，那么方向判据也应该直接来自散度本身，而非采样比值。具体而言，给定沿代理梯度迈出有效步长 $\eta$ 后得到的训练策略 $\pi_\eta$，作者要预测这一步会让 $D_t$ 增大还是减小，从而只在“更新会进一步把策略推到信任域之外”时才遮蔽该 token。最终目标是给出一个既与基于散度的接近判据一致、又能用现有 rollout 工具（仅暴露 top-K 概率）实现、且不引入新超参的替换式方向判据，以提升 RL 训练在不同模型规模（4B/8B/30B）与精度设置（FP8）下的稳定性与有效性。
与已有工作不同的是，DPPO 的工作只动了接近判据这一半，把方向判据原封不动地继承自 PPO。本文的独特切入点是把方向判据也“分布化”，做法是用散度自身的方向导数来定义它。关键洞察在于：对于自回归 LLM 所用的 softmax 策略，散度沿一个策略梯度步的方向导数有闭式解，且可精确分解为一个“局部项”$(\pi_k-\mu_k)$（其符号恰好等于基于比值的判据，因为 $\text{sign}(\pi_k-\mu_k)=\text{sign}(r-1)$）加上一个“全局项”$\sum_i\pi_i(\mu_i-\pi_i)$（捕捉 softmax 归一化在全词表上的耦合，单样本比值无法看到）。这一分解既解释了比值判据为何会失效（它丢弃了全局项），又给出了一个几乎零额外开销的修复：把全局项加回来即可。
核心方法
方法整体思路是“用散度的方向导数预判更新趋势”。直觉上，DPPO 的接近判据问“现在是否已超出信任域”，本文的方向判据问“下一步会不会让散度 $D_t$ 继续增大”。技术上，把迈一步后的训练策略写成 $\pi_\eta$，沿单位方向 $v_i = \mathbf{1}[i=k] - \pi_i$（$k$ 为采样 token）对 $D_{KL}(\mu\|\pi_\eta)$ 求方向导数，得到闭式 $\dot{D} = (\pi_k - \mu_k) + \sum_i \pi_i(\mu_i - \pi_i)$。该系数只与训练已有的量（top-K 概率、散度、优势 $\hat{A}_t$）相关，且遮蔽只用符号 $\text{sign}(\hat{A}_t \cdot \dot{D}_t)$，因此对有效步长 $\eta$ 与比值 $r$ 的量级不敏感。针对生产 rollout 仅暴露 top-K（$K=20$）的截断，作者设计聚合尾部和均匀尾部两种轻量估计器，开销可忽略。最终目标为 $\nabla_\theta L = \mathbb{E}_{y\sim\mu}\sum_t M_t^{pred} r_t \hat{A}_t \nabla_\theta\log\pi(y_t|s_t)$。
核心创新是方向判据不再读采样比值，而是读散度自身的一阶变化 $\dot{D}_t$。与已有方法的本质区别在于：比值判据等价于只保留 $\dot{D}$ 分解中的局部项 $(\pi_k - \mu_k)$，而完全忽略全局项 $\sum_i \pi_i(\mu_i - \pi_i)$。全局项来自 softmax 归一化：改变采样 token 的概率必然连带改变词表其余部分的分布，这种分布级耦合是单样本比值无法观测的。本文把两项都保留，使符号能追踪散度的真实变化。一个漂亮的理论结果是：在 DPPO 的二值 KL 近似下，聚合尾部估计恰好退化为基于比值的方向判据（附录 C 证明），说明本文方法是它的严格推广而非另起炉灶。另一个关键性质是“步长无关”：符号 $\text{sign}(\hat{A}\cdot\dot{D})$ 不需要估计 $\eta$ 或 $r$ 的量级，只依赖一阶主导项，而预测更新后真实值 $D_t(\pi_\eta)$ 才需要这些量和高阶 $O(\eta^2)$ 项。
方法步骤详情
步骤分四步。第一步计算基于 top-K 的前向 KL 散度 $D_t = D_{KL}^{TopK}(\mu\|\pi)$：支撑集 $\mathcal{K} = \text{TopK}(\mu,K) \cup \{y_t\}$（$K=20$），词表剩余部分塌缩成聚合尾部桶 $\mu_{tail}=1-\sum_{i\in\mathcal{K}}\mu_i$。第二步接近判据：仅当 $D_t > \delta$（主实验 $\delta=0.15$）时该 token 进入遮蔽候选集。第三步计算方向导数系数——聚合尾部用 $\dot{D}_{agg} = (\pi_k-\mu_k) + \sum_{i\in\mathcal{K}}\pi_i(\mu_i-\pi_i) + \pi_{tail}(\mu_{tail}-\pi_{tail})$；均匀尾部把尾项除以 $n-m$ 得 $\dot{D}_{uni}$（$n$ 约 $10^5$，$m\le 20$）。第四步方向判据：当 $\text{sign}(\hat{A}_t \cdot \dot{D}_t) > 0$ 且 $D_t>\delta$ 时掩码置 0 遮蔽该 token 梯度，否则置 1 保留。整个过程不引入新超参，只复用已有的 $\hat{A}_t$、top-K 概率与散度。
技术新颖性
新颖性体现在三个层面。其一，把方向判据从“读比值”变为“读散度的方向导数”，并在 softmax 策略下给出闭式分解，明确指出比值判据丢掉了全局项，理论上是比值判据的严格推广（附录 C 给出退化等价关系）。其二，理论上的“步长无关性”非常实用：$\text{sign}(\hat{A}\cdot\dot{D})$ 无需估计有效步长 $\eta$ 或比值 $r$ 的量级，只依赖一阶主导项，这与需要高阶项才能预测 $D_t(\pi_\eta)$ 真实值的做法形成鲜明对比。其三，针对 top-K 截断设计了两个轻量估计器，并证明两者之差 $\dot{D}_{agg}-\dot{D}_{uni} \approx \pi_{tail}(\mu_{tail}-\pi_{tail})$ 在大词表下可忽略，因此方法对尾部建模假设不敏感；附录还给出与比值判据的退化等价证明、以及把接近判据换为 top-K TV 时（附录 D）的方向判据对应形式。
实验结果
主实验（图 1）在推荐阈值 $\delta=0.15$ 下于四种设置评测 avg@16（AIME24+AIME25 每题采样 16 次）：Qwen3-4B-Base、Qwen3-8B-Base、Qwen3-30B-A3B-Base 的 FP8 E2E 与 FP8 Rollout。GRPO clip-higher（$\epsilon_{low}=0.2,\epsilon_{high}=0.28$）在两个 30B 设置下崩溃，而基于散度的方法保持稳定；两种预测散度掩码均稳定优于 DPPO-TopK-KL——由于二者共享相同的 top-K KL 接近判据、差异仅来自方向判据，这一提升直接验证了核心假设。聚合尾部与均匀尾部表现几乎一致。图 2 显示把阈值收紧到 $\delta=0.05$ 会损害所有方法（信任域过紧），但预测掩码仍优于 DPPO-TopK-KL，对超参更鲁棒。token 级分析（61 个种子，每种子约 1435 个越界 token，仅约 82 个产生分歧）显示：不安全保留率从 36.9% 降到 34.2%（-2.7 pp），61 个种子中 38 个上散度判据更优；保留更新中收缩散度的占比从 62.9% 升到 65.1%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME24+AIME25 数学推理 | avg@16 准确率（训练曲线） | 预测散度掩码（聚合/均匀尾部）在 4B/8B/30B-FP8 四种设置上一致稳定提升 | DPPO-TopK-KL（同接近判据，仅方向判据不同） | 训练更稳、收敛更高；GRPO clip-higher 在 30B 崩溃而本方法稳定 |
| 方向判据机制验证（token 级） | 不安全保留率（越界 token 中保留后 $\Delta D>0$ 的比例） | 基于散度的方向判据 34.2% | 基于比值的方向判据 36.9% | -2.7 pp，且在 61 个种子中 38 个种子更优 |
| 方向判据机制验证（token 级，互补口径） | 保留更新中收缩散度的占比（$\Delta D<0$） | 65.1% | 62.9% | 保留更多修正性更新 |
局限与改进
作者明确承认方法是一阶局部近似：真实参数更新会聚合 batch 内所有 token 的梯度，某 token 的实际散度变化还会受其他 token 更新的影响，因此预测的符号只是散度趋势的局部估计，而非对更新后 $D_t$ 的精确预测（文末 Remark 3）。其次，效果虽然方向一致但“温和”——分歧 token 集很小（每种子约 82/1435），token 级证据强度有限，作者自己也指出“effect is modest because the disagreement set is small”。另外补充我自己的观察：评估只覆盖数学推理（AIME24/25），未验证代码生成、多轮对话或通用能力；高阶项被完全忽略，在有效步长较大时符号可能失效；TV 版接近判据（附录 D）只给了形式而无主结果实证；阈值过紧（$\delta=0.05$）时所有方法都变差，说明性能对 $\delta$ 仍有依赖。
独立分析的弱点
第一，一阶近似的脆弱性：当有效步长 $\eta$ 较大（如激进学习率、单 rollout 多次更新）时，二阶及更高阶项可能主导，使 $\text{sign}(\hat{A}\cdot\dot{D})$ 失真；改进方向是引入二阶修正或自适应阈值，或直接用蒙特卡洛估计真实 $\Delta D$ 校准符号。第二，分歧集过小导致增益温和：可考虑把方向判据用于更宽松的接近阈值，或与 token 级裁剪/动态采样联动以放大收益。第三，任务覆盖单一（仅数学推理、仅 AIME 两个基准）：应在代码生成、多轮对话、安全对齐等更广任务上验证泛化性，并报告最终峰值分数而非只画曲线。第四，两种尾部估计在极端长尾（尾部质量大）下差异会放大，可结合真实尾部分布学习一个更稳的估计器，而非均匀/二值假设。
未来方向
作者方向包括：把“分布化方向判据”思想推广到其它散度（如总变差 TV，附录 D 已给形式）、探索二阶预测项以缓解大步长失效、以及研究 token 间梯度交互如何影响整体散度变化。基于成果可延伸的方向有：将方向导数思路迁移到其它带信任域的优化器（如 TRPO 的自然梯度、DAPO 的动态采样与组相对优势）、把全局项 $\sum_i\pi_i(\mu_i-\pi_i)$ 作为可解释的诊断量来定位训练崩溃点、结合 importance sampling 修正 stale policy，或在多轮 RLHF / 层级信任域控制场景下验证方向导数的有效性。
复现评估
复现性整体较好：训练基于开源 VeRL 框架，数据用公开的 DAPO-Math-17k（过滤后约 13k 样本），评测用公开 AIME24/AIME25，模型为公开 Qwen3-4B/8B/30B-A3B-Base。核心超参已在正文给出：$K=20$，$\delta=0.15$（敏感度实验 $\delta=0.05$），clip-higher 用 $\epsilon_{low}=0.2,\epsilon_{high}=0.28$；详细超参与硬件在附录 E 的表 1（4 节点共 32 张 NVIDIA H20 GPU）。主要算力门槛高（30B 模型 × 多种精度，外加 token 级分析跑 61 个种子），普通实验室难以完整复现主结果，但 token 级机制分析方法本身可在较小模型上低成本验证。代码是否开源、学习率/优化器等细节未在正文明确给出，建议作者补充代码仓库以保证完全复现。
论文图表
包含学习率、优化器、batch 大小、rollout 配置、$K=20$、$\delta$ 取值、clip-higher 的 $\epsilon_{low}/\epsilon_{high}$，以及 4 节点共 32 张 NVIDIA H20 GPU 等训练与评测超参。
复现实验所需的关键配置清单，对判断方法的算力门槛与可复现性至关重要，建议读者结合此表评估复现可行性。