# Stop Training for the Worst: Progressive Unmasking Accelerates Masked Diffusion Training

arXiv: https://arxiv.org/abs/2602.10314 · HTML: https://mlsanity.com/paper/2602.10314/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 41 of its 51 references that are in the map.
- First posted on arXiv 10 Feb 2026.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 30 of 41)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think](https://mlsanity.com/paper/2410.06940/) (2024)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Self Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusion](https://mlsanity.com/paper/2506.08009/) (2025)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [Mercury: Ultra-Fast Language Models Based on Diffusion](https://mlsanity.com/paper/2506.17298/) (2025)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling](https://mlsanity.com/paper/2409.02908/) (2024)
- [Scaling up Masked Diffusion Models on Text](https://mlsanity.com/paper/2410.18514/) (2024)
- [DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation](https://mlsanity.com/paper/2506.20639/) (2025)
- [Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference](https://mlsanity.com/paper/2508.02193/) (2025)
- [LLaDA2.0: Scaling Up Diffusion Language Models to 100B](https://mlsanity.com/paper/2512.15745/) (2025)
- [A Reparameterized Discrete Diffusion Model for Text Generation](https://mlsanity.com/paper/2302.05737/) (2023)
- [KodCode: A Diverse, Challenging, and Verifiable Synthetic Dataset for Coding](https://mlsanity.com/paper/2503.02951/) (2025)
- [Dream-Coder 7B: An Open Diffusion Language Model for Code](https://mlsanity.com/paper/2509.01142/) (2025)
- [Path Planning for Masked Diffusion Model Sampling](https://mlsanity.com/paper/2502.03540/) (2025)
- [TinyGSM: achieving >80% on GSM8k with small language models](https://mlsanity.com/paper/2312.09241/) (2023)
- [OpenCodeInstruct: A Large-scale Instruction Tuning Dataset for Code LLMs](https://mlsanity.com/paper/2504.04030/) (2025)

## Built on by (most cited first, 10 of 10)

- [Fixed-Point Masked Generative Modeling](https://mlsanity.com/paper/2605.31215/) (2026)
- [DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Discrete Diffusion Models](https://mlsanity.com/paper/2604.24357/) (2026)
- [Set Diffusion: Interpolating Token Orderings Between Autoregression and Diffusion for Fast and Flexible Decoding](https://mlsanity.com/paper/2607.01775/) (2026)
- [Learned Relay Representations for Forward-Thinking Discrete Diffusion Models](https://mlsanity.com/paper/2605.22967/) (2026)
- [Looped Diffusion Language Models](https://mlsanity.com/paper/2605.26106/) (2026)
- [The Confidence Shortcut: A Reasoning Failure Mode of Masked Diffusion Models](https://mlsanity.com/paper/2605.29123/) (2026)
- [CaRE Compute-aware Remasking Evaluation Protocol for Masked Diffusion Language Models](https://mlsanity.com/paper/2607.24763/) (2026)
- [From Interface to Inference: Eliciting Any-Order Inference from Any-Order Models](https://mlsanity.com/paper/2607.26504/) (2026)
- [On the Design Space of Continuous Diffusion Language Models](https://mlsanity.com/paper/2610.13834/) (2026)
- [OmniPoser: Flexible Human Motion Recovery in the Wild with Masked Flow Matching](https://mlsanity.com/paper/2610.19857/) (2026)

## Abstract (arXiv)

Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces. By generating sequences in any order and allowing for parallel decoding, they enable fast inference and strong performance on non-causal tasks. However, this flexibility comes with a training complexity trade-off: MDMs train on an exponentially large set of masking patterns, which is not only computationally expensive, but also creates a train--test mismatch between the random masks used in training and the highly structured masks induced by inference-time unmasking. In this work, we propose Progressive UnMAsking (PUMA), a simple modification of the forward masking process that aligns training-time and inference-time masking patterns, thereby focusing optimization on inference-aligned masks and speeding up training. Empirically, PUMA speeds up pretraining at the 125M scale by ≈ 2.5× and offers complementary advantages on top of common recipes like autoregressive initialization. We open-source our codebase at https://github.com/JaeyeonKim01/PUMA.