Title: Bringing Stability to Diffusion: Decomposing and Reducing Variance of Training Masked Diffusion Models
Primary area: foundation or frontier models, including LLMs
Keywords: Masked Diffusion Models, Training Variance, Training Stability, Mask Schedule, Mask Sampling

Abstract: Recently, masked diffusion models (MDMs) have emerged as a promising direction beyond auto-regressive models (ARMs). However, MDM training consistently exhibits high variance, leading to less stable optimization and lower final performance. Despite several attempts at mitigation, there has been no theoretical explanation or systematic solution to this instability. In this paper, we derive the first decomposition of MDM training variance into (A) masking pattern noise, (B) masking ratio noise, and (C) data noise, while ARMs are only affected by (C). Building on this foundation, we design six variance-reduction methods, including two core methods: 1) P-POTS, a Pareto-optimal t-distribution that samples harder t values more often with smaller update steps to reduce (A), (B) and (C) simultaneously, and 2) MIRROR, which constructs negatively correlated samples to reduce (A). Experiments show that, compared to standard MDM training, our methods improve accuracy by 7–8% on complex reasoning tasks (such as GSM8K and HiTab), while simultaneously reducing run-to-run variability from 5.07% to 1.35%, thereby substantially narrowing the gap with strong ARM baselines.

TL;DR: We stabilize MDM training by deriving a variance decomposition and introducing two core methods: P-POTS, which is Pareto-optimal among all unbiased t-samplers, and MIRROR, which complements it. Experiments yield clear gains on final performance.