# Learning Unmasking Policies for Diffusion Language Models

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

## Key facts

- Cited by 30 papers in the map (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 72 of its 87 references that are in the map.
- First posted on arXiv 9 Dec 2025.

## Placed in

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

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [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)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding](https://mlsanity.com/paper/2505.22618/) (2025)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [SkipNet: Learning Dynamic Routing in Convolutional Networks](https://mlsanity.com/paper/1711.09485/) (2017)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation](https://mlsanity.com/paper/2303.00848/) (2023)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Training Language Models to Reason Efficiently](https://mlsanity.com/paper/2502.04463/) (2025)
- [Continuous diffusion for categorical data](https://mlsanity.com/paper/2211.15089/) (2022)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Remasking Discrete Diffusion Models with Inference-Time Scaling](https://mlsanity.com/paper/2503.00307/) (2025)

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

- [dUltra: Ultra-Fast Diffusion Language Models via Reinforcement Learning](https://mlsanity.com/paper/2512.21446/) (2025)
- [Stop Training for the Worst: Progressive Unmasking Accelerates Masked Diffusion Training](https://mlsanity.com/paper/2602.10314/) (2026)
- [Adaptation to Intrinsic Dependence in Diffusion Language Models](https://mlsanity.com/paper/2602.20126/) (2026)
- [Activation Steering for Masked Diffusion Language Models](https://mlsanity.com/paper/2512.24143/) (2025)
- [DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking](https://mlsanity.com/paper/2603.01367/) (2026)
- [Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models](https://mlsanity.com/paper/2604.02560/) (2026)
- [The Design Space of Tri-Modal Masked Diffusion Models](https://mlsanity.com/paper/2602.21472/) (2026)
- [A Tale of Two Temperatures: Simple, Efficient, and Diverse Sampling from Diffusion Language Models](https://mlsanity.com/paper/2604.09921/) (2026)
- [Fixed-Point Masked Generative Modeling](https://mlsanity.com/paper/2605.31215/) (2026)
- [ILRR: Inference-Time Steering Method for Masked Diffusion Language Models](https://mlsanity.com/paper/2601.21647/) (2026)
- [TABES: Trajectory-Aware Backward-on-Entropy Steering for Masked Diffusion Models](https://mlsanity.com/paper/2602.00250/) (2026)
- [DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention](https://mlsanity.com/paper/2603.08026/) (2026)
- [DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Discrete Diffusion Models](https://mlsanity.com/paper/2604.24357/) (2026)
- [Mask-Aware Policy Gradients for Diffusion Language Models](https://mlsanity.com/paper/2607.15200/) (2026)
- [Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models](https://mlsanity.com/paper/2606.10829/) (2026)
- [Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from k-Parity](https://mlsanity.com/paper/2601.22450/) (2026)
- [Visual-Redundancy-Controlled Parallel Decoding for Diffusion-Based Multimodal Large Language Models](https://mlsanity.com/paper/2605.25820/) (2026)
- [Re-evaluating Confidence Remasking in Masked Diffusion Language Models](https://mlsanity.com/paper/2606.12232/) (2026)
- [Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models](https://mlsanity.com/paper/2606.23567/) (2026)
- [Accelerating Masked Diffusion Large Language Models: A Survey of Efficient Inference Techniques](https://mlsanity.com/paper/2607.12829/) (2026)
- [From Interface to Inference: Eliciting Any-Order Inference from Any-Order Models](https://mlsanity.com/paper/2607.26504/) (2026)
- [Temporal Self-Distillation: Faster Inference in Discrete Diffusion Language Models](https://mlsanity.com/paper/2609.15177/) (2026)
- [Reconsidering Positional Supervision in Masked Diffusion Language Model Training](https://mlsanity.com/paper/2601.22947/) (2026)
- [Progressive Refinement Regulation for Accelerating Diffusion Language Model Decoding](https://mlsanity.com/paper/2603.04514/) (2026)
- [Adaptive Order Policies for Masked Diffusion](https://mlsanity.com/paper/2606.00295/) (2026)
- [One-step lowest-variance selection in a Gaussian random-field model motivated by masked diffusion: Total correlation and a square root collision threshold](https://mlsanity.com/paper/2607.17522/) (2026)
- [Prefix-Denoising Consistency: Test-Time Verification for Diffusion Language Models](https://mlsanity.com/paper/2608.25311/) (2026)
- [Fork-dLLM: Avoiding the Flexibility Trap in Diffusion Language Models](https://mlsanity.com/paper/2609.39859/) (2026)
- [ITC-MoE: Importance-guided Token-aware Compression for MoE Diffusion Language Models](https://mlsanity.com/paper/2610.01296/) (2026)
- [Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models](https://mlsanity.com/paper/2610.05125/) (2026)

## Abstract (arXiv)

Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. One critical design aspect of dLLMs is the sampling procedure that selects which tokens to unmask at each diffusion step. Indeed, recent work has found that heuristic strategies such as confidence thresholding improve both sample quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger block sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive (block) generation, while outperforming them in the full-diffusion setting.