Stabilizing Reinforcement Learning for Diffusion Language Models

该论文针对群相对策略优化(GRPO)直接应用于扩散语言模型时因概率比估计噪声导致的奖励崩溃问题,提出了通过无条件截断和自归一化机制来打破不稳定性循环的 StableDRL 方法,并进一步将其扩展至块级扩散模型。

Jianyuan Zhong, Kaibo Wang, Ding Ding, Zijin Feng, Haoli Bai, Yang Xiang, Jiacheng Sun, Qiang Xu

发布于 2026-03-10
📖 1 分钟阅读☕ 轻松阅读

Each language version is independently generated for its own context, not a direct translation.

这篇论文主要解决了一个大模型训练中的“翻车”问题,并提出了一个名为 StableDRL 的“急救包”。

为了让你轻松理解,我们可以把训练一个扩散语言模型(dLLM)想象成教一个盲人画家(AI)通过不断修改画作来画出一幅完美的风景画

1. 背景:为什么之前的方法会“翻车”?

传统的做法(GRPO):
想象一下,老师(算法)给盲人画家看了一组画(比如 8 张),然后说:“这组画里,哪张最好?哪张最差?我们要让画得好的那张‘多画几次’,画得差的‘少画几次’。”
在传统的自回归模型(像 AR 模型,一个字一个字写)中,老师能非常精准地算出每张画的“好分数”(概率),所以这个教学很顺利。

扩散模型的问题(dLLM):
但是,扩散模型(dLLM)是像“去噪”一样,从一团乱麻中慢慢变出画。老师很难直接算出那张画到底“有多好”(概率计算太难了,数学上叫“不可解”)。
于是,老师只能(用蒙特卡洛采样去估算)。

翻车现场(Reward Collapse):
因为是在“猜”,所以这个“好分数”经常猜错,而且错得离谱(比如把一张乱涂的画猜成满分,或者把满分猜成负分)。
这就导致了两个致命问题:

  1. 条件性“刹车”失灵(Conditional Clipping Failure):

    • 比喻: 老师给画家设了一个“安全区”。如果画家画得比预期好太多,老师会踩刹车,防止他太飘;如果画得差,老师会鼓励他“大改”。
    • 问题: 因为分数是“猜”的,有时候画家明明画得很烂(负分),但因为猜错了,分数突然变得超级高。这时候,老师以为画家表现太好了,不仅不踩刹车,反而猛踩油门(梯度爆炸),让画家朝着错误的方向狂奔。
    • 后果: 画家彻底跑偏,越改越烂,最后奖励(Reward)直接归零,训练崩溃。
  2. 群体“平均数”失效(Group Normalization):

    • 比喻: 老师通常看 8 张画的平均分来调整。如果其中一张画因为“猜错”分数变得巨大,它会瞬间拉高整个群体的平均值。
    • 问题: 这导致老师给所有人的指令都变得忽大忽小,像坐过山车一样,画家根本学不到东西。

恶性循环:
一旦画家跑偏了,他下次画的画就更不像原来的风格,导致老师下次“猜分数”时错得更离谱。错得更离谱 -> 指令更疯狂 -> 跑得更偏。这是一个死循环


2. 解决方案:StableDRL(稳定强化学习)

作者提出了 StableDRL,就像给这个混乱的教学过程装上了两个“稳定器”:

第一招:无条件“硬刹车”(Unconditional Clipping)

  • 以前的做法: 只有在分数特别高且是“正分”时才刹车;如果是“负分”但分数很高(因为猜错了),就不刹车。
  • StableDRL 的做法: 不管分数是正还是负,只要超过安全范围,一律强制刹车!
  • 比喻: 就像给画家的油门装了一个物理限位器。不管老师怎么喊“冲啊”,只要速度超过限速,物理结构就会强制把速度降下来。这样,哪怕老师猜错了分数,画家也不会因为一次错误的指令而飞出去。

第二招:自我“归一化”(Self-Normalization)

  • 以前的做法: 用固定的数字(比如 8 张画)来平均。如果有一张画分数异常高,平均值就被带偏了。
  • StableDRL 的做法: 不看固定的数字,而是看所有画家的“有效努力”总和
  • 比喻: 老师不再说“我们要把 8 个人的平均分提高”,而是说“我们要把这 8 个人实际贡献的总力量重新分配”。如果某个人因为猜错分数突然变得“力大无穷”,老师会把他多余的力量“稀释”掉,确保他的指令不会压倒其他人。
  • 效果: 这样,无论个别样本怎么“发疯”,整体的更新方向始终被限制在一个安全的、合理的范围内(凸包内),不会剧烈震荡。

3. 额外大招:楼梯式注意力(Staircase Attention)

对于一种更高级的“分块扩散模型”(Block Diffusion),还有一个难题:怎么在计算时不“作弊”?

  • 比喻: 就像做数学题,如果你能一眼看到答案(未来的信息),你就不是在学习,而是在作弊。
  • StableDRL 的做法: 设计了一种**“楼梯式”的遮光板**。
    • 画家在画第 2 层楼梯时,只能看到第 1 层(过去的信息),绝对看不到第 3 层(未来的答案)。
    • 这样既保证了计算速度快(一次算完),又保证了没有作弊(信息不泄露),让训练过程既快又稳。

4. 结果如何?

  • 以前: 训练到 300 步左右,奖励直接崩盘,模型变傻。
  • 现在(StableDRL): 模型可以稳稳地训练 1000 步以上,奖励持续上升。
  • 成绩: 在数学推理(MATH500)、逻辑谜题(Sudoku)等任务上,StableDRL 训练出来的模型比之前的所有方法都强,甚至超过了某些传统的自回归大模型。

总结

这篇论文的核心思想就是:在 AI 训练中,当“评分系统”本身不可靠(充满噪音)时,不能依赖灵活的“条件判断”,而必须依靠“死板的物理限制”(无条件截断)和“动态的平衡机制”(自我归一化)来防止模型因为一次错误的判断而彻底跑偏。

这就好比在暴风雨中开船,如果指南针(评分)经常乱指,你不能听它指挥猛转舵,而应该把舵锁死在安全角度,并依靠船身的平衡设计(归一化)来平稳前行。