提出新方法让扩散语言模型用更少步骤高效生成文本。
Taming Masked Diffusion Language Models via Consistency Trajectory Reinforcement Learning with Fewer Decoding Step
- 设计新解码策略,使模型在非自回归条件下更稳定
- 实验显示在数学与规划任务上仅用较少步数即达竞争力表现
- 适合关注推理效率与并行生成的AI研究者
掩码扩散语言模型(MDLMs)作为自回归模型的潜在替代方案,具备并行解码、灵活生成顺序及减少推理步数的潜力。然而,针对MDLM的解码策略与强化学习算法仍缺乏深入探索。直接套用自回归模型的方法存在训练-推理不一致问题,因MDLM解码是非因果的。为此,本文提出EOS早期拒绝(EOSER)与递增步长(ASS)解码调度器,充分释放全扩散式解码潜力,在少量解码步数下实现优异性能。同时提出一致性轨迹组相对策略优化(CJ-GRPO),强化推演轨迹与优化轨迹的一致性,降低跳步优化带来的误差。在LLaDA-8B-Instruct上进行大量推理任务实验,涵盖数学与规划基准,结果表明所提方法能有效且高效地驾驭MDLM。代码已开源。
原文摘要 · Abstract (English)
Masked diffusion language models (MDLMs) have recently emerged as a promising alternative to autoregressive (AR) language models, offering properties such as parallel decoding, flexible generation orders, and the potential for fewer inference steps. Despite these advantages, decoding strategies and reinforcement learning (RL) algorithms tailored for MDLMs remain underexplored. A naive approach is to directly transfer techniques well-established for AR models to MDLMs. However, this raises an immediate question: Is such a naive transfer truly optimal? For example, 1) Block-wise and semi-AR decoding strategies are not employed during the training of MDLMs, so why do they outperform full diffusion-style decoding during inference? 2) Applying RL algorithms designed for AR models directly to MDLMs exhibits a training-inference inconsistency, since MDLM decoding are non-causal (parallel). This results in inconsistencies between the rollout trajectory and the optimization trajectory. To address these challenges, we propose EOS Early Rejection (EOSER) and Ascending Step-Size (ASS) decoding scheduler, which unlock the potential of MDLMs to perform full diffusion-style decoding, achieving competitive performance with fewer decoding steps. Additionally, we introduce Consistency Trajectory Group Relative Policy Optimization (CJ-GRPO) for taming MDLMs, which emphasizes the consistency between rollout trajectory and optimization trajectory, and reduces the optimization errors caused by skip-step optimization. We conduct extensive experiments on reasoning tasks, such as mathematical and planning benchmarks, using LLaDA-8B-Instruct. The results demonstrate that the proposed EOSER and ASS mechanisms, together with CJ-GRPO, hold significant promise for effectively and efficiently taming MDLMs. Code: https://github.com/yjyddq/EOSER-ASS-RL.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。