AI在线 AI在线

用更一致的轨迹、更少的解码步数「驯服」掩码扩散语言模型,扩散语言模型的推理性能和效率大幅提升

扩散大语言模型得到了突飞猛进的发展,早在 25 年 2 月 Inception Labs 推出 Mercury—— 第一个商业级扩散大型语言模型,同期人民大学发布第一个开源 8B 扩散大语言模型 LLaDA,5 月份 Gemini Diffusion 也接踵而至。 种种迹象表明,扩散大语言模型很可能是下一代大语言模型基础范式的有力竞争者。 但是针对于扩散大语言模型的解码策略和强化学习算法仍然是欠探索的。

图片

扩散大语言模型得到了突飞猛进的发展,早在 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)的高效解码策略 + 强化学习训练组合,显著提升了掩码扩散大语言模型的推理性能与效率,为扩散大语言模型的发展开辟了新路径。

图片
  • 代码仓库:https://github.com/yjyddq/EOSER-ASS-RL

  • 论文链接:https://arxiv.org/pdf/2509.23924

解决了什么问题?

掩码扩散大语言模型(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 在规划类任务中表现更佳,而数学类任务则更适合半自回归 / 分块解码,这比较符合人类的直观感受。

未来可探索混合推理模式,结合扩散与自回归优势,适应多样化任务需求。

📩欢迎引用、关注与合作交流。

相关资讯

黄学东等多名华人当选,2023美国工程院院士名单出炉

今日,美国工程院公布了新增院士名单,包括 106 名院士和 18 名外籍院士。
2/8/2023 10:29:00 AM
机器之心

一篇推文看一年!Jim Fan力荐2025必读清单:50篇论文,扫盲「全领域AI实战」

别再反复看Transformer之类的古董级论文了,该更新论文清单了! 现行的AI从工程技术角度可以分为十个领域:前沿大模型、基准评估、提示思维链、检索增强生成、智能体、代码生成、视觉、声音、图像/视频扩散、微调,每个领域选出5篇代表作和相关工作,看完 实践=AI全栈大神! 前沿大模型OpenAI发布的ChatGPT把大模型带入普通消费者的日常生活中,旗下的GPT系列模型也一直是行业标杆,其中GPT1/2/3, Codex, InstructGPT, GPT4都有论文,GPT3.5, 4o, o1和o3只有相关的宣发活动。
1/10/2025 9:20:00 AM
新智元

细粒度对齐无需仔细标注了!淘天提出视觉锚定奖励,自我校准实现多模态对齐

AIxiv专栏是AI在线发布学术、技术内容的栏目。 过去数年,AI在线AIxiv专栏接收报道了2000多篇内容,覆盖全球各大高校与企业的顶级实验室,有效促进了学术交流与传播。 如果您有优秀的工作想要分享,欢迎投稿或者联系报道。
1/19/2025 12:07:00 PM
机器之心
  • 1