arXiv:2605.10518cs.CLcs.AI2026-05

提出无限掩码扩散模型,实现少步生成并突破传统方法瓶颈。

Infinite Mask Diffusion for Few-Step Distillation

  • 用随机无限状态掩码替代确定性单状态掩码,缓解因子化误差
  • 在小步数下超越现有方法,在LM1B和OpenWebText上表现优异
  • 兼容预训练权重,适合需要高效生成的场景

掩码扩散模型(MDMs)作为自回归模型的替代方案,在语言建模中展现出并行解码和双向上下文处理的优势,其核心在于显式区分掩码词元与数据。然而,由于同时更新多个词元导致的因子化误差,传统MDMs通常需要大量采样步骤。我们发现该误差存在理论下界,而标准MDMs因采用确定性单状态掩码无法降低。本文提出无限掩码扩散模型(IMDM),引入随机无限状态掩码以突破该理论下界,同时保持原有优势,包括与预训练权重的兼容性。实验表明,即使在简单合成任务中,标准MDM也无法实现少步生成,而IMDM可找到高效解法。结合适当蒸馏方法后,IMDM在小步数下优于现有方法,在LM1B和OpenWebText上表现更优。代码已公开于https://Ugness.github.io/official_imdm。

原文摘要 · Abstract (English)

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.

扩散模型少步生成掩码机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。