无限掩码扩散：用于少步蒸馏
计算与语言
2026-05-12 v1 人工智能
摘要
掩码扩散模型 (MDM) 已成为语言建模中自回归模型的有前景的替代方案，具有并行解码和在简单且有效的框架内实现双向上下文处理的优势。 Specifically, their explicit distinction between masked tokens and data underlies their simple framework and effective conditional generation. 然而，MDM 通常需要许多采样迭代 due to factorization errors from simultaneous token updates. 我们观察到，存在一种理论上的因子误差下界，标准 MDM 无法通过使用确定性单状态掩码来降低。本文提出了无限掩码扩散模型 (IMDM)，它引入了随机无限状态掩码以缓解理论下界，同时直接继承了 MDM 的优势，包括与预训练权重的兼容性。我们在简单合成任务上经验性地展示了 MDM 由于因子误差下界而无法进行少步生成，而 IMDM 能够为同一任务找到高效解决方案。最后，在合适的蒸馏方法下，IMDM 在 LM1B 和 OpenWebText 上的小步数计数中超越了现有的少步蒸馏方法。代码可在 https://Ugness.github.io/official_imdm 查看。
引用
@article{arxiv.2605.10518,
title = {Infinite Mask Diffusion for Few-Step Distillation},
author = {Jaehoon Yoo and Wonjung Kim and Chanhyuk Lee and Seunghoon Hong},
journal= {arXiv preprint arXiv:2605.10518},
year = {2026}
}