Masked Diffusion Models as Energy Minimization

Ji-Rong Wen (Renmin University of China) · Sitong Chen (Renmin University of China) · Shen Nie (Renmin University of China) · Jiacheng Sun (Huawei Technologies Co., Ltd) · Zijin Feng (Chinese University of Hong Kong) · Zhenguo Li (Huawei Noah's Ark Lab, Hong Kong) · Chongxuan LI (Renmin University of China)

beta distributionsconditional kinetic energydiscrete optimal transportenergy minimizationgeodesic energyinterpolation scheduleskinetic energylow-step samplingmasked diffusion modelsoptimality conditionpost-training tuningreal-world benchmarks sampling efficiencyschedule design spacesynthetic benchmarks

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.