LED: Improving Masked Diffusion Language Models by Lookahead Editing
作者: Yuyan Zhou, Zeyuan Yang, Yuxuan Song, Zhicheng Cai, Hongli Yu, Weiying Ma, Hao Zhou, James Kwok
分类: cs.LG, cs.CL
发布日期: 2026-10-16
备注: 10 pages, 5 figures
💡 一句话要点
提出Lookahead Editing以优化掩码扩散语言模型
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱四:生成式动作 (Generative Motion) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散模型 自然语言处理 文本生成 前瞻更新 模型优化 深度学习
📋 核心要点
- 现有的掩码扩散模型在表达能力上存在不足,均匀扩散模型虽然表达能力强,但训练复杂度高。
- 本文提出Lookahead Editing,通过引入可重写的提议状态,结合前瞻更新机制,优化掩码扩散解码过程。
- 在八个基准测试中,LED模型在性能上平均提升了11.4和7.7分,且LED-8B超越了经过多种训练方式的LLaDA2.1-16B Q模式。
📝 摘要(中文)
扩散大语言模型(dLLMs)作为自回归语言模型的有力替代,支持通过并行解码进行非自回归文本生成。现有的掩码扩散模型(MDMs)在优化上较为简单,但计算表达能力较低;而均匀扩散模型(UDMs)则具备更高的表达能力,但训练难度较大。为此,本文提出Lookahead Editing(LED),通过可重写的提议状态增强掩码扩散解码,在每一步解码中进行多次前瞻更新,以迭代优化未来预测,从而在保持MDMs简单优化调度的同时提升表达能力。实验结果表明,LED在多个基准测试中显著优于标准MDMs。
🔬 方法详解
问题定义:本文旨在解决掩码扩散模型(MDMs)在表达能力不足与均匀扩散模型(UDMs)训练复杂度高之间的权衡问题。现有方法在优化上存在简单性与表达能力之间的矛盾。
核心思路:提出Lookahead Editing(LED),通过引入可重写的提议状态,在每一步解码中进行前瞻更新,以迭代优化未来的预测,从而提升模型的表达能力,同时保持MDMs的简单优化调度。
技术框架:LED的整体架构包括掩码扩散解码模块和前瞻更新模块。解码过程中,模型首先生成提议状态,然后通过多次前瞻更新迭代优化该状态,最后将优化后的状态应用于掩码序列中。
关键创新:LED的核心创新在于引入了可重写的提议状态和前瞻更新机制,使得模型在保持简单优化的同时,显著提升了表达能力。这一设计与传统的掩码扩散模型形成了鲜明对比。
关键设计:在模型训练中,LED使用了特定的损失函数以平衡提议状态的优化与掩码序列的生成。此外,模型的参数设置经过精心调整,以确保在不同基准测试中表现出色。具体的网络结构和训练细节在实验部分进行了详细描述。
🖼️ 关键图片
📊 实验亮点
LED在八个基准测试中表现出色,平均提升了11.4和7.7分,尤其是LED-8B模型在使用单一的监督微调时超越了经过多种训练方式的LLaDA2.1-16B Q模式,显示出其优越的性能。
🎯 应用场景
该研究的潜在应用领域包括自然语言生成、对话系统和文本摘要等。通过提升掩码扩散模型的表达能力,LED能够在多种文本生成任务中提供更高质量的输出,具有广泛的实际价值和未来影响。
📄 摘要(原文)
Diffusion large language models (dLLMs) have emerged as a promising alternative to autoregressive language models, enabling non-autoregressive text generation through parallel decoding. Existing dLLMs can be broadly categorized into two families: masked diffusion models (MDMs) and uniform diffusion models (UDMs) In this work, we theoretically show that MDMs are easier to optimize due to their mask-based transition schedule but have lower computational expressivity, whereas UDMs enable iterative state revision and greater expressivity but are substantially harder to train. This trade-off raises a natural question: can we obtain the expressivity of iterative revision without giving up the simple, easy-to-optimize transition schedule of MDMs? Motivated by this analysis, we propose Lookahead Editing (\ours{}), which augments masked diffusion decoding with a rewritable proposal state. At each decoding step, the model performs several lookahead updates to iteratively refine future predictions in this proposal state before committing tokens to the masked sequence. This design introduces a rewritable proposal that uses additional compute to improve expressivity while keeping the simple mask-based transition schedule of MDMs. We fine-tune Qwen2.5-7B and Qwen3-8B to obtain LED-\texttt{7B} and LED-\texttt{8B}, respectively. Across eight benchmarks, LED SFT yields average gains of $+11.4$ and $+7.7$ points over standard MDM SFT. Moreover, using SFT alone, LED-\texttt{8B} outperforms LLaDA2.1-16B Q-mode, which is trained with continued pre-training, supervised fine-tuning, and reinforcement learning.