Research map / Image generation / Fast diffusion sampling
Masked Diffusion Models as Energy Minimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Frames masked diffusion models as energy minimizers in discrete optimal transport, proves equivalence of kinetic, conditional kinetic and geodesic energies, derives an optimal mask schedule condition, and proposes Beta-CDF schedule tuning for LLaDA 8B sampling.
Problem
Understanding why mask schedules in masked diffusion models affect few-step sampling efficiency, and how to choose them principledly rather than using hand-crafted linear or sine schedules.
Approach
Shows three energy functionals (kinetic, conditional kinetic, geodesic) are proportional under MDM structure, derives the closed-form optimal condition alpha_t = sin^2(pi/2 gamma_t) linking mask schedule to a geodesic interpolation schedule, and parameterizes gamma_t by a Beta CDF so schedule tuning reduces to a 2D search post-training without retraining.
Key results
- On MBPP and HumanEval, beta-parameterized schedules match the linear baseline quality with 2x fewer sampling steps.
- On Hendrycks Math, beta schedules reach linear-schedule performance with 4x fewer steps.
- On BBH and GSM8K, beta schedules show comparable rather than better performance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Beta-CDF-tuned schedules instead of the linear schedule | LLaDA 8B · MBPP | generation quality: 2× step reduction | No clear effect |
| Use Beta-CDF-tuned schedules instead of the linear schedule | LLaDA 8B · HumanEval | generation quality: 2× step reduction | No clear effect |
| Use Beta-CDF-tuned schedules instead of the linear schedule | LLaDA 8B · Hendrycks Math | performance: 4× fewer steps | No clear effect |
| Use Beta-CDF-tuned schedules instead of the linear schedule | LLaDA 8B · BBH | performance | No clear effect |
| Use Beta-CDF-tuned schedules instead of the linear schedule | LLaDA 8B · GSM8K | performance | No clear effect |
Surprising findings
- On BBH and GSM8K, beta-parameterized schedules perform comparably to, not better than, the baselines.
Limitations
- Relationship between high-dimensional real-world tasks and their optimal schedules is not well understood.
- Theoretical optimality is shown for the continuous case; finite-step discretization is treated only briefly.
- Gains are task-dependent and not uniform across benchmarks.
Experimental setup
Datasets: MBPP HumanEval BBH GSM8K Hendrycks Math Minerva Math
Benchmarks: MBPP HumanEval BBH GSM8K Hendrycks Math Minerva Math
Metrics: task performance vs sampling steps
Baselines: linear schedule sine schedule squared sine schedule
Builds on: discrete flow matching kinetic energy formulation LLaDA
Closest papers by meaning
- On scalable and efficient training of diffusion samplers (NeurIPS 2025 · similarity 0.88)
- Why Masking Diffusion Works: Condition on the Jump Schedule for Improved Discrete Diffusion (NeurIPS 2025 · similarity 0.88)
- Align Your Steps: Optimizing Sampling Schedules in Diffusion Models (ICML 2024 · similarity 0.88)
- Simplified and Generalized Masked Diffusion for Discrete Data (NeurIPS 2024 · similarity 0.88)
- Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics (ICLR 2025 · similarity 0.87)
- Hierarchical Variational Policies for Reward-Guided Diffusion (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.