字节：多掩码扩散实现少步生成

标题Multi-Mask Diffusion Language Models for Few-Step Generation来源arXiv, 2607.19686v1️文章简介研究问题如何解决掩码扩散模型因终端状态熵为零而难以进行一致性蒸馏加速的问题主要贡献论文提出MultiMDM通过引入多掩码机制保留掩码结构并解决终端退化实现了高质量少步生成。重点思路设计多掩码前向过程将单一吸收态扩展为多个掩码集合每个干净Token映射到指定掩码使前向轨迹在混合前偏向指定掩码从而在反向过程中保留部分信息以支持草稿生成。推导闭式ELBO训练目标包含干净Token重建项和掩码内识别项该目标支持从预训练的单掩码扩散模型进行持续微调无需从头训练。提出纯离散状态的一致性蒸馏方案利用共享Gumbel噪声耦合构建低熵路径将一致性目标定义为沿耦合路径的未来滤波后验从而减少路径随机性并实现少步采样。分析总结在LM1B和OpenWebText数据集上的预训练实验显示MultiMDM在困惑度与样本熵的权衡上优于现有的均匀态扩散和单掩码扩散基线且从预训练模型持续微调效果更佳。经过一致性蒸馏后MultiMDM在极少步数如4步生成下显著降低了生成困惑度证明了多掩码教师模型能为蒸馏提供更丰富的中间状态信息。消融实验表明增加掩码数量能提升性能但在M50左右达到饱和将该方法适配到8B参数规模的LLaDA模型上在数学和代码基准测试中也取得了改进验证了其可扩展性。个人观点论文在离散扩散中引入了“多掩码”这一中间态既避免了均匀态扩散中噪声与语义难以区分的问题又解决了传统掩码扩散终端无信息的困境。