近期，复旦大学等联合团队提出了一种针对掩码扩散大语言模型（MDLM）的高效解码策略与强化学习训练组合。该方法解决了MDLM全扩散式解码易陷入“EOS陷阱”导致性能下降的问题，通过“EOS早期拒绝机制”和“幂次递增解码步长调度器”，显著加速了推理过程，将步数从O(L)降至O(logL)。同时，创新性地提出“一致性轨迹分组优化（CJ-GRPO）”算法，有效缓解了强化学习训练中的轨迹不一致问题。实验表明，该组合策略在数学和规划任务上均能取得优异表现，尤其在少步数推理下，性能媲美甚至超越传统方法，为扩散大模型的进一步发展开辟了新路径。
💡 **解决“EOS陷阱”，提升全扩散解码性能：** 针对掩码扩散大语言模型（MDLM）在全扩散式解码时容易过早生成`` token导致生成中断的痛点，研究团队提出了“EOS早期拒绝机制”。该机制在解码早期主动抑制``的置信度，避免生成过早终止，并在后期逐步恢复其置信度以确保句子正常结束，从而显著提升了全扩散式解码的性能，并在部分任务上超越了分块解码。
🚀 **加速推理，引入幂次递增解码步长：** 基于对解码过程中token置信度变化趋势的观察（前期平缓，后期急剧上升），研究团队设计了“幂次递增解码步长调度器”。此举使得模型能够前期谨慎解码，后期大胆解码，将推理步数从线性复杂度O(L)大幅降低至对数复杂度O(logL)，极大地加速了推理过程，充分释放了扩散大语言模型的推理速度优势。
🔄 **优化训练稳定性，提出一致性轨迹分组优化：** 为解决将自回归模型的强化学习算法迁移至MDLM时可能遇到的“rollout轨迹与优化轨迹不一致”问题，团队提出了“一致性轨迹分组优化（CJ-GRPO）”算法。该算法在rollout过程中存储中间状态，逐步优化相邻状态间的转变，有效缓解了跳步优化带来的误差，提升了训练的稳定性和效果。结合递增步长调度器，还大幅削减了训练时的存储开销。
📊 **实验验证，展现多任务优势：** 在数学推理（GSM8K, MATH500）和规划任务（Countdown, Sudoku）上的广泛实验表明，所提出的方法在不同任务上均表现出色。CJ-GRPO结合EOSER和ASS策略在规划任务上性能提升高达2-4倍，而CJ-GRPO与半自回归（Semi-AR）结合在数学任务上表现突出。在仅使用log(L)步数的情况下，新方法性能仍优于基线方法，实现了“又快又好”的推理效果。
2025-11-05 12:13 四川
一套对于掩码扩散大语言模型的高效解码策略+强化学习训练组合，显著提升了掩码扩散大语言模型的推理性能与效率。
扩散大语言模型得到了突飞猛进的发展，早在 25 年 2 月 Inception Labs 推出 Mercury—— 第一个商业级扩散大型语言模型，同期人民大学发布第一个开源 8B 扩散大语言模型 LLaDA，5 月份 Gemini Diffusion 也接踵而至。种种迹象表明，扩散大语言模型很可能是下一代大语言模型基础范式的有力竞争者。但是针对于扩散大语言模型的解码策略和强化学习算法仍然是欠探索的。近期，复旦大学、上海人工智能实验室、上海交通大学联合研究团队发布最新论文《Taming Masked Diffusion Language Models via Consistency Trajectory Reinforcement Learning with Fewer Decoding Step》。他们提出了一套对于掩码扩散大语言模型（Masked Diffusion Large Language Model，MDLM）的高效解码策略 + 强化学习训练组合，显著提升了掩码扩散大语言模型的推理性能与效率，为扩散大语言模型的发展开辟了新路径。
解决了什么问题？掩码扩散大语言模型（MDLM）如 LLaDA 展现出与自回归模型相媲美的能力，并具备并行解码、灵活生成顺序、潜在少步推理等优势。然而，完全扩散式（Full Diffusion-Style）解码策略并未被广泛使用，取而代之的是分块解码（Block-wise）。因为目前的完全扩散式解码存在一大痛点 —— 性能大幅度逊色于分块解码。
但令人疑惑的是，掩码扩散大语言模型在预训练和指令微调阶段并未针对分块解码这种方式进行适配微调，所以这一现象背后的原因仍不为人所知。该团队基于这点发现刨根问底，最终定位到 MDLM 的全扩散式解码的三个特点：解码过程中的 token 置信度变化趋势：由平缓到陡升；解码过程中 <EOS> token 的置信度一直显著大于其他非 <EOS> token；以上两点导致在早期解码时都会有很大概率解码出 <EOS> token，模型像是陷入了 <EOS > 的陷阱无法自拔。而分块解码由于块的限制存在，则不会深陷其中。
此外，在将为 AR LLMs 定制的强化学习算法迁移到 MDLM 时，可能会遇到 rollout 轨迹和优化轨迹不一致的问题，因为 AR 模型具有因果性掩码，获得完整轨迹后计算得到的 token 概率与 rollout 时保持一致。而 MDLM 采用的是双向注意力，获得完整轨迹再计算得到的 token 概率与 rollout 时不一致。而现有的方法是（1）使用 prompt masking 近似的一步优化，或者（2）从 fully masked response 进行一步去噪优化。但是这二者都面临 rollout 轨迹和实际优化轨迹不一致的问题，可能会引起较大的优化误差。基于此，该团队将以上问题汇总为三个关键问题：解码陷阱：全扩散解码容易在早期步骤中过早生成过多的 <EOS>token，导致生成中断，性能下降。等间隔解码步长：根据 token 置信度的变化曲线，等间隔的解码步长可能不是最优的，因为解码前期对解码置信度并不高，而解码后期很确定。rollout 轨迹和优化轨迹不一致：直接将自回归模型的强化学习算法迁移至 MDLM 可能会面临 rollout 轨迹和优化轨迹不一致问题，影响优化效果。
提出了什么方法？
该团队提出三大核心贡献，致力于解决了上述问题：
针对于 <EOS > 陷阱！该团队提出 <EOS > 早期拒绝机制在解码早期主动抑制 <EOS> 的置信度，避免生成过早终止。后期逐步恢复 <EOS > 的置信度，确保句子正常结束。显著提升全扩散式解码性能，在部分任务上显著超越半自回归解码。
基于「早期置信度低、后期急剧上升」的观察，前期谨慎解码，后期大胆解码，设计幂次递增解码步长调度器。将推理步数从 O (L) 降至 O (logL)，大幅加速推理过程。
针对 rollout 轨迹和优化轨迹不一致问题，该团队提出一致性轨迹分组策略优化（CJ-GRPO），在 rollout 过程中存储每一步解码的中间状态，逐步优化相邻两个状态之间的转变，从而有效缓解跳步优化（或者说是不一致轨迹优化）带来的误差，提升训练稳定性与效果。
但是，由于过长的步数会引入较大的存储开销，而结合递增步长调度器后，既能保证轨迹的一致性，同时能大大缩减中间状态存储开销。于是该团队结合 <EOS > 早期拒绝机制、递增步长调度器和 CJ-GRPO 算法，削减训练时 CJ-GRPO 的中间状态存储开销，同时使得训练后的模型在少解码步数推理下甚至能达到 Baseline 方法多解码步数时的性能 —— 一石三鸟。
模型在少解码步数 (logL) 时仍能保持可观的性能（与 L/2 解码步数性能可比），真正激发扩散语言模型的推理速度优势的潜能。 训练时的解码时间 / 空间复杂度从 O (L) 降至 O (logL)，大幅加速
训练过程。
实验结果在数学推理（GSM8K、MATH500）和规划任务（Countdown、Sudoku）上，使用 LLaDA-8B-Instruct 模型进行了广泛的实验，结果显示：
一致性轨迹优化在数学和规划任务上实现了对 baseline 的全面超越。CJ-GRPO + Semi-AR 在数学类任务上表现尤为突出。CJ-GRPO + EOSER + ASS 在规划类任务上表现尤为突出，某些 benchmark 的性能提升高达 2–4 倍。发现规划任务适合并行推理，数学问题适合顺序推理。
在都仅使用 log (L) 步数的情况下，EOSER + ASS 性能仍优于分块解码和全扩散式解码策略。一致性轨迹优化结合递增步长调度器和 EOSER 解码策略，在少步数推理场景下保持良好的竞争力，真正实现「又快又好」。
研究意义与未来展望本工作探究了扩散语言模型的全扩散式解码策略，并对其使用更一致的轨迹、更少的解码步数进行优化，用更少的步数，越快越好地完成更复杂的推理任务，推动扩散语言模型的全扩散式解码、少步数解码、强化学习算法的发展。并行推理 v.s. 顺序推理：装配了并行解码的 MDLM 在规划类任务中表现更佳，而数学类任务则更适合半自回归 / 分块解码，这比较符合人类的直观感受。
未来可探索混合推理模式，结合扩散与自回归优势，适应多样化任务需求。
📩欢迎引用、关注与合作交流。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
阅读原文
跳转微信打开
