mlsanityData to 7 Oct 2026

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

arXiv 2602.10314 ↗First posted 10 Feb 2026cs.LGcited by 8

On the map: Diffusion models → Discrete and text diffusion

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.

Key facts. Cited by 8 papers in the map (data through 7 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.

Builds on

41 of 51 references.

Closest ancestors · 22

Same direction, newest first

  1. 2025LLaDA2.0: Scaling Up Diffusion Language Models to 100B10 Dec
  2. 2025Learning Unmasking Policies for Diffusion Language Models9 Dec
  3. 2025Decoding Large Language Diffusion Models with Foreseeing Movement3 Dec
  4. 2025Encoder-Decoder Diffusion Language Models for Efficient Training and Inference26 Oct
  5. 2025Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies7 Oct
  6. 2025Demystifying MaskGIT Sampler and Beyond: Adaptive Order Selection in Masked Diffusion6 Oct
  7. 2025Planner Aware Path Learning in Diffusion Language Models Training27 Sep
  8. 2025Dream-Coder 7B: An Open Diffusion Language Model for Code1 Sep
  9. 2025Any-Order Flexible Length Masked Diffusion31 Aug
  10. 2025Reinforced Context Order Recovery for Adaptive Reasoning and Planning18 Aug
  11. 2025Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference4 Aug
  12. 2025DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation25 Jun
10 more
  1. 2025Large Language Diffusion Models14 Feb
  2. 2025Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions10 Feb
  3. 2025Path Planning for Masked Diffusion Model Sampling5 Feb
  4. 2024Scaling up Masked Diffusion Models on Text24 Oct
  5. 2024Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling4 Sep
  6. 2024Simple and Effective Masked Diffusion Language Models11 Jun
  7. 2024Simplified and Generalized Masked Diffusion for Discrete Data6 Jun
  8. 2024Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data6 Jun
  9. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution25 Oct
  10. 2023A Reparameterized Discrete Diffusion Model for Text Generation11 Feb

Other references · 19

  1. 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,084
  2. 2021Training Verifiers to Solve Math Word Problemscited by 9,743
  3. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  4. 2021Improved Denoising Diffusion Probabilistic Modelscited by 3,736
  5. 2021Program Synthesis with Large Language Modelscited by 3,509
  6. 2024Qwen2 Technical Reportcited by 2,076
  7. 2022MaskGIT: Masked Generative Image Transformercited by 1,082
  8. 2024Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Thinkcited by 666
  9. 2025Self Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusioncited by 569
  10. 2022A Continuous Time Framework for Discrete Denoising Modelscited by 437
  11. 2023Analyzing and Improving the Training Dynamics of Diffusion Modelscited by 396
  12. 2024Discrete Flow Matchingcited by 301
7 more
  1. 2025Mercury: Ultra-Fast Language Models Based on Diffusioncited by 225
  2. 2025KodCode: A Diverse, Challenging, and Verifiable Synthetic Dataset for Codingcited by 115
  3. 2023TinyGSM: achieving >80% on GSM8k with small language modelscited by 68
  4. 2025OpenCodeInstruct: A Large-scale Instruction Tuning Dataset for Code LLMscited by 57
  5. 2024Causal Language Modeling Can Elicit Search and Reasoning Capabilities on Logic Puzzlescited by 35
  6. 2026Stable-DiffCoder: Pushing the Frontier of Code Diffusion Large Language Modelcited by 12
  7. 2025Demystifying Diffusion Objectives: Reweighted Losses are Better Variational Boundscited by 7

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 8.

8 papers.

  1. 2026Fixed-Point Masked Generative Modelingcited by 2
  2. 2026Set Diffusion: Interpolating Token Orderings Between Autoregression and Diffusion for Fast and Flexible Decodingcited by 2
  3. 2026DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Discrete Diffusion Modelscited by 1
  4. 2026Learned Relay Representations for Forward-Thinking Discrete Diffusion Modelscited by 1
  5. 2026The Confidence Shortcut: A Reasoning Failure Mode of Masked Diffusion Modelscited by 1
  6. 2026CaRE Compute-aware Remasking Evaluation Protocol for Masked Diffusion Language Modelscited by 1
  7. 2026Looped Diffusion Language Modelscited by 0
  8. 2026From Interface to Inference: Eliciting Any-Order Inference from Any-Order Modelscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
  2. 2026Learn from Your Mistakes: Self-Correcting Masked Diffusion Models
  3. 2025Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models
  4. 2026Understanding and Accelerating the Training of Masked Diffusion Language Models
  5. 2026Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from k-Parity
  6. 2025Masked Diffusion Models are Secretly Learned-Order Autoregressive Models
  7. 2025Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models
  8. 2026On Trajectory-Aware Training for Masked Diffusion Language Models