# Where-to-Unmask: Ground-Truth-Guided Unmasking Order Learning for Masked Diffusion Language Models

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 32 of its 47 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 32)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (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)
- [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)
- [Did Aristotle Use a Laptop? A Question Answering Benchmark with Implicit Reasoning Strategies](https://mlsanity.com/paper/2101.02235/) (2021)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (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)
- [Remasking Discrete Diffusion Models with Inference-Time Scaling](https://mlsanity.com/paper/2503.00307/) (2025)
- [DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation](https://mlsanity.com/paper/2506.20639/) (2025)
- [d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning](https://mlsanity.com/paper/2504.12216/) (2025)
- [LLaDA2.0: Scaling Up Diffusion Language Models to 100B](https://mlsanity.com/paper/2512.15745/) (2025)
- [Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning](https://mlsanity.com/paper/2410.14157/) (2024)
- [Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking](https://mlsanity.com/paper/2505.24857/) (2025)
- [Insertion Transformer: Flexible Sequence Generation via Insertion Operations](https://mlsanity.com/paper/1902.03249/) (2019)
- [Path Planning for Masked Diffusion Model Sampling](https://mlsanity.com/paper/2502.03540/) (2025)
- [A Survey on Diffusion Language Models](https://mlsanity.com/paper/2508.10875/) (2025)
- [Any-Order Flexible Length Masked Diffusion](https://mlsanity.com/paper/2509.01025/) (2025)
- [Diffusion Language Models are Super Data Learners](https://mlsanity.com/paper/2511.03276/) (2025)
- [Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective](https://mlsanity.com/paper/2512.03759/) (2025)
- [PiRank: Scalable Learning To Rank via Differentiable Sorting](https://mlsanity.com/paper/2012.06731/) (2020)
- [Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies](https://mlsanity.com/paper/2510.05725/) (2025)
- [Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models](https://mlsanity.com/paper/2511.05563/) (2025)
- [Flexible-length Text Infilling for Discrete Diffusion Models](https://mlsanity.com/paper/2506.13579/) (2025)
- [On the Reasoning Abilities of Masked Diffusion Language Models](https://mlsanity.com/paper/2510.13117/) (2025)

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

- [Chatterbox-Flash: Prior-Calibrated Block Diffusion for Streaming Zero-Shot TTS](https://mlsanity.com/paper/2605.30748/) (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)
- [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)
- [Decoding in Order-Agnostic Language Models: Chain-Rule Deviation and Uniform Spreading](https://mlsanity.com/paper/2606.00997/) (2026)
- [Unmask the State: When Does State Adaptation Matter for Masked Diffusion Language Models](https://mlsanity.com/paper/2609.33355/) (2026)
- [Fine-Tuning Diffusion Language Models with Context Selection and Target Weighting](https://mlsanity.com/paper/2609.38385/) (2026)
- [The Lattice of Transition Laws](https://mlsanity.com/paper/2610.11216/) (2026)

## Abstract (arXiv)

Masked Diffusion Language Models (MDLMs) generate text by iteratively filling masked tokens, requiring two coupled decisions at each step: which positions to unmask (where-to-unmask) and which tokens to place (what-to-unmask). While standard MDLM training directly optimizes token prediction (what-to-unmask), inference-time unmasking orders (where-to-unmask) are typically determined by heuristic confidence measures or trained through reinforcement learning with costly on-policy rollouts. To address this, we introduce Gt-Margin, a position-wise score derived from ground-truth tokens, defined as the probability margin between the correct token and its strongest alternative. Gt-Margin yields an oracle unmasking order that prioritizes easier positions first under each partially masked state. We demonstrate that leveraging this oracle unmasking order significantly enhances final generation quality, particularly on logical reasoning benchmarks. Building on this insight, we train a supervised unmasking planner via learning-to-rank to imitate the oracle ordering from masked contexts. The resulting planner integrates into standard MDLM sampling to select where-to-unmask, improving reasoning accuracy without modifying the token prediction model.