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的结构下，动能、条件动能和测地能量三种不同的能量形式是数学上等价的。这意味着优化MDM可以等价于最小化这些能量，从而为MDM的设计和优化提供了统一的理论基础。

技术框架：论文的技术框架主要包括以下几个部分：1) 建立MDM与离散最优传输之间的联系；2) 证明三种能量形式在MDM结构下的等价性；3) 推导最优掩码策略的闭式解；4) 使用Beta分布参数化掩码策略，并进行高效搜索。整体流程是从理论分析到算法设计，最终通过实验验证算法的有效性。

关键创新：论文最重要的技术创新点在于将MDM与能量最小化问题联系起来，并证明了不同能量形式的等价性。这一发现为MDM的理论分析和算法设计提供了新的视角。此外，通过Beta分布参数化掩码策略，并进行高效搜索，也为MDM的实际应用提供了便利。与现有方法相比，该方法不再依赖于手工调整掩码策略，而是通过优化能量函数来自动学习最优策略。

关键设计：论文的关键设计包括：1) 使用离散最优传输框架来描述MDM；2) 定义动能、条件动能和测地能量三种能量形式；3) 推导最优掩码策略的闭式解，该解依赖于扩散过程的噪声水平；4) 使用Beta分布来参数化掩码策略，Beta分布的两个参数控制了掩码策略的形状；5) 设计高效的搜索算法来优化Beta分布的参数，从而找到最优的掩码策略。

📊 实验亮点

实验结果表明，通过Beta分布参数化并优化后的掩码策略，在合成数据和真实数据集上均优于手工设计的基线方法。尤其是在低步采样设置下，性能提升更为显著，验证了能量最小化框架的有效性。

🎯 应用场景

该研究成果可应用于图像修复、文本补全、三维重建等领域，通过优化掩码策略，提高生成质量和采样效率。此外，该理论框架为理解和改进其他扩散模型提供了新的思路，具有重要的学术价值和潜在的工业应用前景。

📄 摘要（原文）

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.