稳定扩散语言模型的强化学习训练 Stabilizing Reinforcement Learning for Diffusion Language Models

StableDRL的完整方法包含以下步骤：首先，对于从行为策略 $\pi_{\theta_{old}}$ 采样的rollout组 $\{x_1, \dots, x_G\}$，通过蒙特卡洛采样估计ELBO来计算重要性比率 $\hat{\rho}_j = \exp(\hat{\mathcal{L}}_\theta(x_j) - \hat{\mathcal{L}}_{\theta_{old}}(x_j))$。其次，对每个重要性比率进行无条件裁剪，定义裁剪后的权重 $w_{i,j} := \text{clip}_\epsilon(\hat{\rho}_{i,j})$，确保权重始终在 $[1-\epsilon, 1+\epsilon]$ 范围内。然后，使用自归一化计算梯度更新：$\nabla_\theta J_{\text{Ours}} = \mathbb{E}\left[\frac{1}{\sum_{i=1}^{G} \text{clip}_\epsilon(\hat{\rho}_i)} \sum_{j=1}^{G} \text{clip}_\epsilon(\hat{\rho}_j) A_j g_j\right]$。在实际实现中，为保证数值稳定性，所有计算在对数空间进行：定义裁剪后的对数比率 $\tilde{\ell}_j = \text{clip}(\hat{\mathcal{L}}_\theta(x_j) - \hat{\mathcal{L}}_{\theta_{old}}(x_j), \log(1-\epsilon), \log(1+\epsilon))$，然后通过稳定的softmax $\exp(\tilde{\ell}_j - \text{LSE}(\{\tilde{\ell}_k\}_{k=1}^{G}))$ 计算归一化系数，其中LSE是Log-Sum-Exp函数。对于块扩散模型，还引入阶梯注意力机制，使用双流输入（冻结的干净上下文和损坏的目标）配合块下三角掩码，实现 $O(1)$ 的无泄漏代理估计。

StableDRL的技术新颖性体现在多个层面。首先，从问题定义的角度，本文首次系统性地将dLLM训练不稳定性建模为一个自强化循环，而非单纯的估计误差问题，这一视角转变引导了全新的解决方案方向。其次，无条件裁剪虽然概念简洁，但其理论贡献在于证明了GRPO的条件裁剪在噪声比率下具有固有的缺陷——条件裁剪允许的“返回信任域加速”机制会被噪声利用，转化为梯度尖峰。第三，自归一化的引入不仅是经验性的改进，更有深刻的几何意义：将更新限制在样本梯度的凸包内，从结构上消除了组级随机性，这是对标准归一化方式的根本性改进。第四，阶梯注意力机制解决了块扩散模型在RL训练中的效率-泄漏两难问题，通过结构化掩码实现 $O(1)$ 的单次前向传播ELBO估计。最后，对数空间的数值稳定实现确保了在混合精度训练中的可靠性。

从独立分析的角度来看，StableDRL存在以下几个值得关注的弱点。首先，无条件裁剪阈值 $\epsilon$ 的选择需要仔细调优——实验表明 $\epsilon = 5$ 提供了最佳的探索-稳定性权衡，但当 $\epsilon$ 增大到100或1000时，“陷阱门”失败模式会重新出现，说明方法对超参数仍然敏感。改进方向可以是设计自适应的裁剪阈值，根据训练动态自动调整。其次，自归一化虽然在理论上将更新限制在凸包内，但当所有样本的重要性比率都被裁剪到接近边界时，凸包退化为一个点，学习信号可能被过度压缩。可以考虑引入温度参数来软化这种退化。第三，阶梯注意力机制虽然实现了 $O(1)$ 的单次前向传播，但其块下三角掩码的设计限制了块内token之间的信息流动，可能在需要精细块内交互的任务上存在性能瓶颈。第四，论文的方法论聚焦于重要性比率的噪声，但未深入分析奖励函数设计对dLLM训练稳定性的影响，奖励稀疏性或奖励黑客问题可能在更复杂的任务中成为新的瓶颈。第五，StableDRL的数值稳定实现依赖于对数空间计算和softmax技巧，这增加了实现复杂度，在某些硬件平台上可能引入额外的计算开销。

基于StableDRL的成果，未来研究可以在多个方向展开。首先，作者的工作为dLLM的RL训练奠定了稳定基础，一个自然的扩展是将StableDRL与更好的重要性比率估计方法结合——当前方法让优化器对噪声免疫，而改进的估计器可以减少噪声本身，两者可能是正交且互补的。其次，StableDRL的自归一化思想可以推广到其他策略梯度方法，如PPO、A3C等，探索更广泛的自归一化框架。第三，阶梯注意力机制为块扩散模型的RL训练打开了大门，未来可以探索更复杂的注意力模式来支持更长的块大小或更灵活的块间依赖。第四，当前实验集中在数学推理任务上，将StableDRL扩展到代码生成、工具使用、多轮对话等更广泛的任务类型是重要的未来方向。第五，随着dLLM规模的增长，研究StableDRL在更大模型（如70B+）上的scaling行为具有重要的实际意义。第六，论文的理论分析建立了不稳定性循环的数学框架，未来可以进一步发展更精细的收敛性分析，给出StableDRL的显式收敛速率。

在复现性方面，StableDRL具有较高的可复现性。作者提供了完整的代码仓库（https://github.com/JianyuanZhong/StableDRL），包含了所有实验的实现细节。论文在附录中详细描述了训练基础设施、模型配置和超参数，包括AdamW优化器、学习率 $1.0 \times 10^{-6}$、裁剪阈值 $\epsilon = 5$ 等。实验基于公开的LLaDA-8B-Instruct和SDAR-8B-Chat模型，数据集使用标准的GSM8K、MATH500、Countdown和Sudoku基准，训练和测试分割均公开。主要的复现挑战在于计算资源需求——全参数RL训练8B模型需要显著的GPU资源，且论文提到使用了自定义的JetEngine推理内核来实现计算可行的训练。此外，蒙特卡洛估计的随机性可能引入结果方差，需要多次运行取平均。总体而言，对于拥有足够计算资源的研究者，StableDRL的复现难度为中等。