Masked Diffusion Models as Energy Minimization

📄 arXiv: 2509.13866v2 📥 PDF

作者: Sitong Chen, Shen Nie, Jiacheng Sun, Zijin Feng, Zhenguo Li, Ji-Rong Wen, Chongxuan Li

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-09-17 (更新: 2025-11-27)

期刊: Published at NeurIPS 2025

💡 一句话要点

将掩码扩散模型解释为能量最小化问题，并提出能量驱动的采样策略。

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 掩码扩散模型 能量最小化 最优传输 采样策略 Beta分布

📋 核心要点

  1. 现有掩码扩散模型缺乏统一的理论解释，手工设计的采样策略并非最优。
  2. 论文将掩码扩散模型视为能量最小化问题，推导出三种等价的能量公式。
  3. 提出基于Beta分布的参数化调度方法，通过二维搜索优化采样策略，提升低步采样性能。

📝 摘要（中文）

本文提出了一个系统的理论框架，将掩码扩散模型（MDM）解释为离散最优传输中能量最小化问题的解。具体而言，我们证明了在MDM的结构下，三种不同的能量公式——动能、条件动能和测地能量——在数学上是等价的，并且当掩码调度满足闭式最优条件时，MDM可以最小化所有这三种能量。这种统一不仅阐明了MDM的理论基础，而且激发了采样方面的实际改进。通过使用Beta分布参数化插值调度，我们将调度设计空间简化为可处理的二维搜索，从而无需修改模型即可实现高效的后训练调整。在合成和真实世界基准上的实验表明，我们受能量启发的调度优于手工设计的基线，尤其是在低步采样设置中。

🔬 方法详解

问题定义：论文旨在解决掩码扩散模型（MDM）缺乏统一理论基础，以及现有手工设计的采样策略并非最优的问题。现有方法通常依赖启发式规则设计掩码调度，缺乏理论指导，导致采样效率和生成质量受限。

核心思路：论文的核心思路是将MDM解释为离散最优传输中的能量最小化问题。通过建立MDM与能量最小化之间的数学联系，可以从能量角度理解MDM的运作机制，并推导出最优的掩码调度策略。

技术框架：论文的技术框架主要包含以下几个部分：1) 建立MDM与三种能量公式（动能、条件动能和测地能量）之间的数学等价性；2) 推导在特定条件下，MDM最小化这三种能量的闭式最优条件；3) 提出基于Beta分布的参数化插值调度方法，将掩码调度设计问题转化为二维搜索问题；4) 通过实验验证所提出的能量驱动的调度策略的有效性。

关键创新：论文最重要的技术创新在于建立了MDM与能量最小化之间的理论联系，并证明了三种能量公式的等价性。此外，提出的基于Beta分布的参数化调度方法，将复杂的掩码调度设计问题简化为可处理的二维搜索问题，为优化采样策略提供了新的思路。

关键设计：论文的关键设计包括：1) 使用Beta分布参数化插值调度，将调度函数表示为Beta分布的参数形式；2) 通过实验搜索Beta分布的参数，找到最优的掩码调度策略；3) 在低步采样设置下，验证所提出的能量驱动的调度策略的性能提升。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，论文提出的能量驱动的调度策略在合成和真实世界基准上均优于手工设计的基线。尤其是在低步采样设置中，性能提升显著，表明该方法能够有效提高采样效率，并生成高质量的样本。

🎯 应用场景

该研究成果可应用于图像修复、图像生成、文本生成等领域。通过优化掩码调度策略，可以提高生成模型的采样效率和生成质量，降低计算成本，并有望推动相关技术在实际场景中的应用，例如智能设计、内容创作等。

📄 摘要（原文）

We present a systematic theoretical framework that interprets masked diffusion models (MDMs) as solutions to energy minimization problems in discrete optimal transport. Specifically, we prove that three distinct energy formulations--kinetic, conditional kinetic, and geodesic energy--are mathematically equivalent under the structure of MDMs, and that MDMs minimize all three when the mask schedule satisfies a closed-form optimality condition. This unification not only clarifies the theoretical foundations of MDMs, but also motivates practical improvements in sampling. By parameterizing interpolation schedules via Beta distributions, we reduce the schedule design space to a tractable 2D search, enabling efficient post-training tuning without model modification. Experiments on synthetic and real-world benchmarks demonstrate that our energy-inspired schedules outperform hand-crafted baselines, particularly in low-step sampling settings.