将掩码扩散模型蒸馏为单步生成器，实现高效图像生成。
Di$\mathtt{[M]}$O: Distilling Masked Diffusion Models into One-step Generator
- 通过令牌级分布匹配优化输出逻辑，解决单步生成信息瓶颈。
- 在图文生成任务中达到与多步教师模型相当的性能。
- 首次实现掩码扩散模型单步蒸馏，适合高效生成场景。
掩码扩散模型（MDMs）已成为强大的生成建模技术，但通常存在推理缓慢的问题。本文提出Di[M]O，一种将掩码扩散模型蒸馏为单步生成器的新方法。针对两个关键挑战：（1）中间步骤信息难以用于单步生成，通过基于辅助模型的‘在线策略’框架实现令牌级分布匹配，优化模型输出逻辑；（2）初始分布熵不足，采用令牌初始化策略注入随机性，同时保持与教师训练分布相似。实验表明，Di[M]O在类别条件与文本条件图像生成上均表现优异，性能媲美多步教师模型，且推理时间显著降低。据我们所知，这是首个成功实现掩码扩散模型单步蒸馏的工作，也是首个将离散蒸馏应用于文生图生成的研究，为高效生成建模开辟新路径。
原文摘要 · Abstract (English)
Masked Diffusion Models (MDMs) have emerged as a powerful generative modeling technique. Despite their remarkable results, they typically suffer from slow inference with several steps. In this paper, we propose Di$\mathtt{[M]}$O, a novel approach that distills masked diffusion models into a one-step generator. Di$\mathtt{[M]}$O addresses two key challenges: (1) the intractability of using intermediate-step information for one-step generation, which we solve through token-level distribution matching that optimizes model output logits by an 'on-policy framework' with the help of an auxiliary model; and (2) the lack of entropy in the initial distribution, which we address through a token initialization strategy that injects randomness while maintaining similarity to teacher training distribution. We show Di$\mathtt{[M]}$O's effectiveness on both class-conditional and text-conditional image generation, impressively achieving performance competitive to multi-step teacher outputs while drastically reducing inference time. To our knowledge, we are the first to successfully achieve one-step distillation of masked diffusion models and the first to apply discrete distillation to text-to-image generation, opening new paths for efficient generative modeling.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。