DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models¶
会议: NeurIPS 2026
arXiv: 2605.16342
OpenReview: https://openreview.net/forum?id=vjbVsKvLXq
论文: https://arxiv.org/abs/2605.16342
领域: LLM 其他
关键词: 扩散语言模型, 强化学习, 信用分配, GRPO, 去噪进度评分
一句话总结¶
针对离散扩散语言模型强化学习中去噪时间步信用分配缺失及均值场似然估计偏置严重两大痛点,本文提出即插即用的 DACA-GRPO 框架,通过重用中间步骤 logits 构建去噪进度评分进行逐 token 优势调制,并利用分层掩码似然估计消除均值场偏置,在多基线与跨任务上实现显著增益。
研究背景与动机¶
离散掩码扩散语言模型(Diffusion LLMs,如 LLaDA、MDLM)通过在离散空间内对全掩码序列执行并行迭代去噪,展现出高吞吐以及任意顺序生成文本的独特优势。随着组相对策略优化(GRPO)等强化学习后训练算法在大语言模型推理任务上取得巨大成功,近期研究如 Diffu-GRPO、wd1 以及 GDPO 尝试将强化学习迁移至扩散语言模型。然而,现有方法在策略优化中普遍沿用黑盒范式,将整个复杂的去噪轨迹压缩为一个标量奖励或粗粒度优势,完全忽视了去噪时间维度上的精细结构。
这种机制存在两项根本性缺陷。首先是时间信用分配(Temporal Credit Assignment)的完全缺失:扩散生成过程在不同去噪步上的重要性差异极大,某些早期或关键过渡步决定了算术运算符、代码控制流等核心骨架的选择,而大量后续步骤仅负责空白填充或常规补齐;现有算法对所有去噪步一视同仁分配均等信用,既低估了关键决策步的贡献,又在模型发生自我怀疑或纠错时施加过度惩罚。其次是均值场似然估计的系统性偏置与高方差:扩散模型无法直接计算序列对数似然,主流方法在计算重要性权重时采用完全掩码的单样本独立分解估计,强制模型在零上下文先验下预测所有 token,造成极大的估计偏差和梯度噪声。
本文的切入点在于:扩散模型在每一去噪步本就会计算全部掩码位置的词表分布,但仅挑选置信度最高的少量 token 解掩,剩余中间预测包含着模型信念演化的丰富信息却被直接抛弃;同时,似然估计的均值场偏置可以通过在位置维度上引入部分上下文条件来系统性纠偏。核心 idea:利用生成过程中被丢弃的中间预测构建零额外开销的去噪进度评分(DPS)进行逐 token 优势调制,并结合位置分层的分层掩码似然(SML)降低均值场偏置,构建首个感知去噪过程的扩散强化学习框架 DACA-GRPO。
方法详解¶
整体框架¶
DACA-GRPO 由两大互补机制构成:去噪进度评分(Denoising Progress Scores, DPS) 与 分层掩码似然估计(Stratified Masking Likelihood, SML)。整体训练流程在生成阶段与损失计算阶段实现无缝解耦。在 rollout 阶段,扩散策略并行生成候选补全,并在每隔特定步长记录未解掩位置上的对数概率分布,计算去噪置信度增益并完成批次标准化,将其映射为每个 token 的优势调制权重;在反向传播与损失计算阶段,根据基础 GRPO 方法类型(比率裁剪型或无比率加权型),通过分层掩码重构序列似然估计,形成去噪感知的梯度更新。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入 Prompt 批次"] --> B["并行迭代去噪采样<br/>计算所有掩码位置 Logits"]
B --> C["去噪进度评分 (DPS)<br/>追踪未解掩位置信念演化与步增量"]
C --> D["去噪步 z-score 标准化与 Token 优势调制<br/>按生成出生步长加权 A~_i"]
B --> E["分层掩码似然 (SML)<br/>K 个互斥子集位置分层评估"]
D --> F["DACA-GRPO 联合损失优化<br/>比率嵌入 (Diffu-GRPO/GDPO) 或 似然正则 (wd1)"]
E --> F
F --> G["策略参数更新 θ"]
关键设计¶
1. 去噪进度评分(DPS):零开销挖掘中间状态的信念增量与逐 token 信用分配
传统扩散 RL 将整个生成的标量优势标定到所有 token 上,无法区分决定性决策与机械性填空。DPS 的核心发现是扩散模型在反向去噪的第 \(k\) 步已经针对所有仍在掩码集合 \(\mathcal{M}_k\) 中的位置计算了完整输出分布。定义第 \(k\) 步模型在后续掩码集合 \(\mathcal{M}_j\)(其中 \(j \ge k\))上的平均对数概率为:
$\(S(k,j) = \frac{1}{|\mathcal{M}_j|} \sum_{i \in \mathcal{M}_j} \log p_\theta(x_i^* \mid \mathbf{x}_k)\)$
其中 \(x_i^*\) 为模型在最终生成中自主选定的 token。从步长 \(k\) 到 \(k+1\) 的去噪进度差值定义为 \(\delta_k := S(k+1, k+1) - S(k, k+1)\)。由于评估固定在同一位置集合 \(\mathcal{M}_{k+1}\) 上,\(\delta_k\) 精确度量了因第 \(k+1\) 步新揭示 token 所带来的整体理解飞跃。若 \(\delta_k > 0\) 显著,说明该步完成关键结构承诺;若 \(\delta_k \le 0\) 则代表无效推进或迷茫。对于最终步无法计算后续差值的情形,算法采用一阶外推 \(\delta_T = \delta_{T-1}\),保持量纲一致且无需额外前向传播。
2. 步级标准化与出生步优势调制:消除非平稳时间尺度与正负向纠错反思
去噪初期上下文匮乏导致 \(\delta_k\) 幅度普遍偏小,而去噪后期上下文充沛导致差值剧烈,直接使用原始差值会使晚期步骤压倒早期决策。本文设计了沿 batch 维度的逐步 z-score 标准化:
$\(\bar{\delta}_k = \frac{\delta_k - \mu_k}{\sigma_k + \epsilon}\)$
其中 \(\mu_k\) 与 \(\sigma_k\) 为当前步在批次内的均值与标准差。每个 token 根据其首次被解掩的“出生步” \(\text{birth}(i)\) 继承对应的标准差分,调制后的逐 token 优势为:
$\(\tilde{A}_i^{(g)} = A^{(g)} \cdot \left(1 + \lambda \cdot \bar{\delta}_{\text{birth}(i)}\right)\)$
当样本优势 \(A^{(g)} > 0\) 时,\(\bar{\delta} > 0\) 放大对关键正确决策的奖励;当 \(A^{(g)} < 0\) 时,若模型在错误路径上表现出自信推进(\(\bar{\delta} > 0\)),则承受更重惩罚,而若模型出现犹豫或纠错态势(\(\bar{\delta} < 0\)),则有效减免惩罚,从而保护探索潜力。
3. 分层掩码似然估计(SML):打破均值场独立性假设以消除偏置
现有方法在估计生成序列对数似然 \(\log p_\theta(y)\) 时,采用全掩码均值场近似,迫使模型在零上下文下对每个 token 独立打分,带来了极高的估计偏置与方差。SML 将长度为 \(N\) 的补全序列均匀划分为 \(K\) 个互斥分层 \(\mathcal{S}_1, \dots, \mathcal{S}_K\)。在第 \(k\) 次前向传播中,仅将 \(\mathcal{S}_k\) 设为 [MASK],其余 \((K-1)/K\) 比例的上下文均显式可见:
$\(\log p_\theta^{\text{SML}}(o \mid q) = \sum_{k=1}^K \sum_{n \in \mathcal{S}_k} \log p_\theta(o_n \mid o_{\setminus \mathcal{S}_k}, q)\)$
\(K=1\) 即退化为均值场基线,而 \(K=4\)(默认设置)使每个 token 能在 75% 的上下文支持下进行预测,极大捕获了局部与长程依赖。在集成策略上,对于无重要性比率的 wd1,SML 作为似然正则项直接相加:\(\mathcal{L}_{\text{DACA-wd1}} = \mathcal{L}_{\text{wd1}}(\tilde{A}_i) - \frac{\eta}{G} \sum_g \log p_\theta^{\text{SML}}(y^{(g)})\);对于依赖 PPO 裁剪比率的 Diffu-GRPO 与 GDPO,将全掩码视角与分层掩码视角在对数概率空间对等平均,构成富集对数概率 \(\hat{\ell}_{\theta,i} = \frac{1}{2}[\log p_\theta(o_i \mid \mathbf{m}, q) + \log p_\theta(o_i \mid o_{\setminus \mathcal{S}_{k(i)}}, q)]\),从内嵌层面稳定信赖域。
损失函数 / 训练策略¶
DACA-GRPO 适用于不同基线优化目标。对于基于比率裁剪的 Diffu-GRPO,损失函数采用富集比率与 DPS 调制优势的联合形式: $\(\mathcal{L}_{\text{DACA}} = -\mathbb{E} \left[ \min\left(\rho_i^{\text{SML}} \cdot \tilde{A}_i^{(g)}, \;\text{clip}(\rho_i^{\text{SML}}, 1-\epsilon, 1+\epsilon) \cdot \tilde{A}_i^{(g)}\right) \right]\)$ 在训练参数上,骨干模型采用 LLaDA-8B-Instruct,以 bfloat16 精度配合 4-bit QLoRA(Rank 128, Alpha 64)进行微调。优化器采用 AdamW,学习率 \(3\times 10^{-6}\),梯度累积步长 2,每组候选数 \(G=6\)。DPS 的记录步长 \(s\) 支持自适应间隔抽样(推荐 \(s \in \{4, 8, 16, 32\}\)),在 \(s=4\) 时训练开销增加小于 1%;SML 在 \(K=4\) 时的前向开销约为 11.7%,在保证精度的同时维持极高的吞吐效率。
实验关键数据¶
主实验¶
在数学推理(MATH-500、GSM8K)与代码生成(MBPP、HumanEval)上,DACA-GRPO 在 Diffu-GRPO、wd1 及 GDPO 三大主流基线之上均展现出稳健增益。
| 基准模型 / 优化方法 | MATH-500 (L=512) | GSM8K (L=512) | MBPP (L=512) | HumanEval (L=512) | 约束 Sudoku (L=512) |
|---|---|---|---|---|---|
| LLaDA-8B-Instruct 基座 | 36.2% | 78.2% | 45.2% | 34.2% | 5.5% |
| Diffu-GRPO 基线 | 39.2% | 79.8% | 44.2% | 35.4% | 26.3% |
| + DPS | 40.2% | 80.9% | 45.6% | 36.5% | 37.5% |
| + SML | 40.5% | 81.5% | 48.6% | 40.6% | — |
| + DACA-GRPO (完整) | 40.5% | 81.5% | 50.5% | 40.1% | — |
| wd1 基线 | 39.6% | 80.7% | 45.8% | 37.5% | 34.6% |
| + DPS | 40.9% | 81.6% | 47.5% | 40.0% | 70.9% |
| + SML | 41.5% | 82.0% | 47.2% | 39.1% | — |
| + DACA-GRPO (完整) | 41.3% | 82.1% | 50.7% | 40.1% | — |
| GDPO 基线 | 40.7% | 80.0% | 46.1% | 37.0% | 25.8% |
| + DPS | 40.7% | 82.8% | 48.2% | 40.0% | 37.5% |
| + SML | 40.2% | 80.6% | 50.0% | 40.6% | — |
| + DACA-GRPO (完整) | 41.7% | 82.0% | 49.5% | 41.1% | — |
消融实验¶
针对 DPS 的末步外推模式以及步间标准化轴在 MATH-500 上的表现进行深入消融(以 wd1 为基准):
| 实验配置分类 | 具体配置设定 | MATH-500 (L=256) | MATH-500 (L=512) | 显著特征与结论说明 |
|---|---|---|---|---|
| wd1 原生基线 | Baseline (无 DPS/SML) | 36.7% | 39.6% | 均匀信用分配,缺乏时序感知 |
| 末步处理策略 (Last-step) | Raw Score (\(\delta_T = S(T,T)\)) | 38.3% | 35.6% | 绝对分数与差值量纲冲突,长序列严重劣化 |
| 末步处理策略 (Last-step) | Neutral (\(\delta_T = 0\)) | 38.2% | 36.3% | 丢弃最关键的不确定性消解步,表现受损 |
| 末步处理策略 (Last-step) | Mean (均值填充) | 37.1% | 39.2% | 抹杀去噪非平稳特性 |
| 末步处理策略 (Last-step) | Measured (额外前向传播) | 37.4% | 39.7% | 测量真实差值但带来额外显存与时间损耗 |
| 末步处理策略 (Last-step) | Extrapolate (一阶外推 \(\delta_{T-1}\)) | 38.6% | 40.2% | 零计算成本且在所有 stride 下超越基线 (6/6) |
| 标准化轴模式 (Norm Axis) | Trajectory 轨迹内标准化 | 36.9% | 39.3% | 丢失样本间对比,长文本弱于基线 |
| 标准化轴模式 (Norm Axis) | Group 全局标量标准化 | 37.6% | 39.9% | 晚期大步长主导全局统计量,压制早期关键步 |
| 标准化轴模式 (Norm Axis) | None 无标准化 | 37.5% | 39.4% | 对对数尺度极度敏感,超参调优脆弱 |
| 标准化轴模式 (Norm Axis) | Per-step 步级批次标准化 | 38.6% | 40.2% | 有效解耦去噪动态演变并保留跨样本辨识度 |
关键发现¶
- 两模块高度互补且分工明确:DPS 在短序列以及高度依赖严密时序承诺的任务(如 Sudoku 中 wd1 准确率由 34.6% 跃升至 70.9%,Countdown 增益达 10.4pp)中贡献巨大;而 SML 则在长序列复杂依赖任务(如 MBPP \(L=512\) 提升 4.9pp,HumanEval \(L=512\) 提升 4.7pp)中主导增益,两者结合实现近乎线性的效果叠加。
- SML 的适用边界受长度驱动而非结构驱动:在极短输出的数独与倒数任务中,75% 的掩码上下文会造成约束直接穿透退化为平凡补词,引发训练崩溃;但在长达 512 token 的 JSON 结构化输出(github-medium)中,SML 显著提升遵循率(wd1 从 72.57% 提升至 78.47%),证明只要序列具有足够的语义与结构延展度,部分上下文即可提供高质量的非偏似然监督。
- 训练收敛速度显著加快且稳定:在训练动态分析中,wd1 + DPS 在前几百步内即与原生 baseline 拉开差距,全程展现出平滑且严格单调上升的奖励曲线,证明基于固有优势调制的机制完全不会破坏信赖域。
亮点与洞察¶
- 废弃信号变废为宝:巧妙利用离散扩散模型在去噪过程中每一轮前向传播所计算但随后丢弃的大量掩码位置 logits,以零额外前向计算代价提取出模型信念演化的导数(\(\delta_k\)),极具工程智慧与理论美感。
- 方向感知型的惩罚缩放:DPS 对负优势样本的处理赋予了符号敏感性,对在错误道路上“盲目自信推进”的决策加大惩罚,而对自我纠错和犹豫行为给予豁免,极大缓解了 RL 后训练中的过度拟合。
- 正交分解扩散训练误差源:清晰界定了时间轴方差缩减(SDMC、耦合适样)与位置空间均值场偏置(SML)的物理差异,为后续连续/离散混合生成模型的强化学习提供了通用的解构框架。
局限与展望¶
- SML 对极短上下文存在分布偏移失真:由于生成过程始于 100% 掩码,而 SML 在 \(K=4\) 时依赖 75% 的可见上下文,对于依赖纯发散搜索的超短硬约束任务易诱发平凡填空诱导的模式崩溃,未来需要研究动态衰减的退火分层掩码方案。
- 静态 Stride 采样局限:当前采用固定步长 \(s\) 记录快照,未来若能根据 \(\delta_k\) 的即时二阶导数自适应调整探测频率,有望进一步提升细粒度时间信用的信噪比。
相关工作与启发¶
- vs Diffu-GRPO / wd1 / GDPO: 既有扩散 RL 方法均视去噪轨迹为黑盒,给全部 token 施加均等优势标量,并忍受全掩码似然估计的高方差和偏置;DACA-GRPO 作为上层插件,同时重构了优势输入与似然评估,全面提升既有算法上限。
- vs DyLLM: DyLLM 利用注意力上下文相似度挖掘去噪显著位置以在推理期跳步加速;DPS 则聚焦于未揭示位置的输出置信度动态,在训练期用于优化时间信用分配,形成推理加速与训练增强的互补视角。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次揭示并利用扩散去噪轨迹的固有 logits 实现无开销时序信用分配,构思新颖。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 7 大跨领域基准,全面评测 3 种基线模型在 3 种序列长度下的表现,消融严谨。
- 写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻,公式表述清晰节制,图表与实验对比分析极具说服力。
- 价值: ⭐⭐⭐⭐⭐ 具备极高的通用性与即插即用特性,有望成为扩散语言模型 RL 训练的标准范式。