Infinite Mask Diffusion for Few-Step Distillation

📄 arXiv: 2605.10518v1 📥 PDF

作者: Jaehoon Yoo, Wonjung Kim, Chanhyuk Lee, Seunghoon Hong

分类: cs.CL, cs.AI

发布日期: 2026-05-11

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

提出无限掩码扩散模型（IMDM），通过引入随机无限状态掩码突破掩码扩散模型的采样步数限制。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion)

关键词: 掩码扩散模型 文本生成 模型蒸馏 少步采样 分解误差 概率建模

📋 核心要点

  1. 现有掩码扩散模型（MDM）受限于确定性单状态掩码，导致采样过程中存在无法消除的分解误差，限制了其在少步生成任务中的表现。
  2. 论文提出无限掩码扩散模型（IMDM），通过引入随机无限状态掩码机制，从理论上突破了分解误差的下界，提升了模型生成效率。
  3. 实验证明IMDM在合成任务中表现优异，并在LM1B和OpenWebText数据集上，配合蒸馏技术实现了比现有方法更高效的少步文本生成效果。

📝 摘要（中文）

掩码扩散模型（MDMs）作为自回归模型的替代方案，在语言建模中展现出并行解码和双向上下文处理的优势。然而，由于同时更新标记导致的分解误差，MDMs通常需要多次采样迭代。研究发现，标准MDMs因使用确定性的单状态掩码，存在一个无法逾越的理论分解误差下界。为此，本文提出了无限掩码扩散模型（IMDM），引入随机无限状态掩码来降低该理论界限，同时保留了MDMs兼容预训练权重的特性。实验表明，MDM在简单合成任务中因分解误差界限无法实现少步生成，而IMDM能有效解决该问题。在结合适当的蒸馏方法后，IMDM在LM1B和OpenWebText数据集的少步生成任务中优于现有的蒸馏方法。

🔬 方法详解

问题定义：MDM在处理离散数据时，由于在每一步采样中同时更新多个标记，导致联合分布的分解产生误差。现有方法使用确定性的单状态掩码，导致该误差存在一个理论下界，限制了模型在极少步数下的生成质量。

核心思路：引入随机性以平滑采样过程。通过将确定性的掩码状态扩展为无限状态的随机过程，IMDM能够更灵活地控制标记的更新路径，从而在保持模型架构简洁性的同时，有效降低分解误差。

技术框架：IMDM架构基于标准的MDM框架，但在掩码机制上进行了重构。它不再将掩码视为简单的二元状态，而是通过引入随机变量来模拟无限状态空间，使得模型在推理阶段能够通过更平滑的概率转移实现少步采样。

关键创新：核心创新在于“无限状态掩码”的概念，它打破了传统MDM中掩码状态的离散限制，通过随机化处理消除了确定性掩码带来的分解误差下界，使模型在极少步数下仍能保持高生成质量。

关键设计：该方法通过与现有的蒸馏技术（如一致性蒸馏或概率蒸馏）无缝集成，利用预训练权重进行初始化，并在训练过程中优化随机掩码的分布参数，以适应少步推理的需求。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，在简单的合成任务中，传统MDM因分解误差界限无法收敛，而IMDM成功实现了高效生成。在LM1B和OpenWebText基准测试中，IMDM结合蒸馏方法在极少步数（Few-step）设置下，生成质量显著优于现有的蒸馏基线模型，证明了其在提升采样效率方面的显著优势。

🎯 应用场景

该研究主要应用于大规模语言模型的快速推理场景，特别是在需要低延迟、高吞吐的实时文本生成任务中。其技术价值在于通过改进扩散过程，使离散数据生成能够摆脱长采样链的束缚，未来有望在移动端设备、实时对话系统及高效内容创作工具中发挥重要作用。

📄 摘要（原文）

Masked Diffusion Models (MDMs) have emerged as a promising alternative to autoregressive models in language modeling, offering the advantages of parallel decoding and bidirectional context processing within a simple yet effective framework. Specifically, their explicit distinction between masked tokens and data underlies their simple framework and effective conditional generation. However, MDMs typically require many sampling iterations due to factorization errors stemming from simultaneous token updates. We observe that a theoretical lower bound of the factorization error exists, which standard MDMs cannot reduce due to their use of a deterministic single-state mask. In this paper, we propose the Infinite Mask Diffusion Model (IMDM), which introduces a stochastic infinite-state mask to mitigate the theoretical bound while directly inheriting the benefits of MDMs, including the compatibility with pre-trained weights. We empirically demonstrate that MDM fails to perform few-step generation even in a simple synthetic task due to the factorization error bound, whereas IMDM can find an efficient solution for the same task. Finally, when equipped with appropriate distillation methods, IMDM surpasses existing few-step distillation methods at small step counts on LM1B and OpenWebText. Code is available at https://Ugness.github.io/official_imdm.