Multi-Mask Diffusion Language Models for Few-Step Generation
Sijin Chen, Yinuo Ren, Heyang Zhao, Ziheng Cheng, Quanquan Gu, Lexing Ying
COLM 2026
cs.CL, cs.LG
2026-07-22
ByteDance Seed提出MultiMDM:每个干净token先进入指定掩码再在掩码子空间混合。170M上从MDLM续训后,4步对齐熵的生成困惑度从721.1降到558.8。
掩码扩散语言模型(MDM)训练目标干净,重建任务也跟语言建模对齐,但少步生成一直别扭。正向过程最后会塌成「整句全是同一个掩码」,终点熵为零,一致性蒸馏没有随机种子可走。均匀状态扩散(USDM)把 token 打散到词表里,终点有熵,可干净 token 和噪声看起来一样,建模质量和训练效率通常更差。
ByteDance Seed 联合 Princeton、Stanford、UCLA、Berkeley 提出 MultiMDM:保留掩码结构,但用一组掩码代替一个吸收态,让少步生成有地方起步。
每个干净 token 被指派一个指定掩码。正向过程分两段:先按 αt 把可见 token 推进指定掩码,再按 βt 在掩码子空间里混匀。终点是掩码集合上的均匀分布,熵为 L log M,不再是零。反推时,模型可以先猜指定掩码当草稿,再细化成干净 token。掩码数 M 远小于词表;经验上 M 在 5 到 100 有用,主实验取 50。
训练目标有闭式 ELBO,两项相加:一项是 MDM 同款的干净 token 重建,一项是掩码子空间里的身份识别。可见位置不贡献损失。从预训练 MDM 续训很直接:干净 token 输出头不动,把原来的单掩码嵌入复制到所有新掩码上,再学着把它们分开。课程是先重建、再打开掩码识别项。
少步蒸馏仍在离散状态上做。坐标上对每个状态抽一次 Gumbel 噪声并在时间上共享,Gumbel-max 保证边缘分布不变,路径在给定噪声后是确定的。一致性目标对齐这条耦合路径上的干净 token 后验。实现上用 EMA 教师,把未来滤子后验压成只看当前序列的 Markov 预测。
主干是 170M DiT,两个未公开的英语语料(合规原因隐去名称)。生成质量用 GPT-2 测的生成困惑度(GenPPL,越低越好),并报告对齐语料熵和温度 τ=1 两套。预训练 20 万步;续训设定是先 15 万步 MDM,再 5 万步 MultiMDM。
Corpus A(语料熵 5.44)对齐熵的 GenPPL:
| 方法 | 4 步 | 8 步 | 16 步 | 64 步 |
| DUO | 995.3 | 509.4 | 193.6 | 70.9 |
| MDLM | 721.1 | 280.1 | 117.0 | 65.9 |
| CANDI | 666.5 | 237.5 | 126.8 | 69.6 |
| MultiMDM-cont | 558.8 | 219.1 | 102.8 | 59.7 |
4 步相对 MDLM 从 721.1 降到 558.8。Corpus B 上续训在 2 步把 730.3 降到 400.0,对 DUO 的 679.1 也有优势。从随机初始化的 MultiMDM-rand 多数设置只是持平或略好,增益主要来自「先 MDM、再多掩码」的续训。M=50 最好,M=100 在多步上反而变差(64 步 96.7,对 M=50 的 59.7)。
把 LLaDA-8B-Base 用 LoRA 续成 M=50 的 MultiLLaDA,MATH500 在每步 1 个 token 时从 27.0 到 30.2。GSM8K、HumanEval、MBPP 多数设置小幅上升,HumanEval 在每步 4 token 时从 14.3 掉到 9.5。8B 实验没有做共享 Gumbel 蒸馏,也不是全参微调。
想给已有 MDM 加速少步采样,这条路不用改成连续嵌入,也不用换成均匀噪声。指定掩码等于给每个 token 留一张可解码的草稿纸,终点有熵,噪声位置仍然显式可标。170M 无条件生成上数字清楚;8B 只是初步迁移证据,还不能当产品结论。
两个英语语料的真名被隐去,外部无法复现数据配方。主指标是 GPT-2 生成困惑度,和样本熵强相关,作者用对齐熵来补,但仍不是下游任务。8B 只用 LoRA、没有共享 Gumbel 蒸馏,HumanEval 在更粗的采样步上还退步。指定映射是列上独热,没有按语义设计。论文没有独立的 Limitations 节,这些边界来自实验设置和讨论。