Unmask the State: When Does State Adaptation Matter for Masked Diffusion Language Models
Injin Kong, Sunghwan Choi, Yohan Jo
cs.AI
2026-09-27
Splits MDM decoding into five axes, finds adaptation opportunity highly uneven, and gains up to 5.6 points by deviating from fixed strategies only on high-opportunity states.
Masked diffusion language models (MDMs) do not follow the left-to-right pass of autoregressive decoding. They start from a partially masked sequence and reveal tokens step by step, which makes generation order a free variable. Which positions to reveal, how many per step, and whether a revealed prediction stays locked are all inference-time decisions. The past two years produced a pile of proposals: confidence or KL-stability ordering, block-wise reveals, remasking (masking revealed tokens back for re-prediction), lookahead. Each reports gains on its own benchmarks.
All of them share one assumption: a single rule runs from start to finish. Early and late denoising states look very different, so the preferred rule has no reason to stay constant. The paper calls state-dependent switching "state adaptation" and asks whether the opportunity exists, where it concentrates, and whether it can be recognized before acting.
The strategy space factors into five axes: score (which positions to prioritize), cardinality (how many to reveal per step), region (where selection is allowed), commitment (whether revealed predictions can be revised), and planning (whether the decision accounts for future denoising). Experiments vary the first four and hold planning fixed. LLaDA, Fast-dLLM, KLASS, and ReMDM all map onto this grid, which separates choices those methods couple.
Measurement rests on a one-step state-action utility: change only the current decision, keep the rest of decoding fixed, and read off the expected terminal utility. The fixed action is the best candidate chosen on validation prompts, and adaptation opportunity is how much the best state-dependent action beats it, computed against an oracle ceiling. A short theorem splits total opportunity into reversal frequency times the average margin at reversals. Equal means can hide different shapes: small gains spread over many states, or large gains packed into a few. Only the second shape rewards detection. A two-fold rollout cross-fitting protocol (select actions with rollouts 0 and 1, score them with 2 and 3) controls maximization bias.
The deployed policy separates two problems. Action selection: each axis gets a training-free diagnostic. The score diagnostic looks back at which ranking rule proved trustworthy, the region diagnostic measures spatial drift of the priority landscape, the cardinality diagnostic weighs context supplied to downstream tokens against context gained by waiting, and the commitment diagnostic compares the context value of locking a prediction against the pressure of a posterior flip. All statistics come from the decoding trajectory, with no extra forward pass. Opportunity detection: diagnostic values are binned into quantiles on validation prompts, each bin carries the mean observed opportunity, and at inference the policy deviates from the fixed action only when predicted opportunity clears a threshold.
Three models (LLaDA-8B-Instruct, LLaDA-1.5, Dream-7B) across ten tasks. First, does opportunity exist?
| Model | Oracle gap, strongest axis |
| LLaDA-8B | 0.0025 to 0.0454 |
| LLaDA-1.5 | 0.0025 to 0.0450 |
| Dream-7B | 0 to 0.1325 |
The spread is huge. On LLaDA-1.5 CSV Missing Cells, the region axis has a 0.0% positive-opportunity rate: the fixed action is never beaten. Structured tasks (constrained JSON, HTML tag closing, carry circuits) expose the largest gaps, and the strongest axis changes with the model.
Existence is not exploitability. Mean AUROC across task-axis settings is 0.545, 0.541, and 0.571 for the three models, barely above chance, and only 3, 1, and 5 settings show a positive lift at 10% coverage. The bright spots: on LLaDA-8B constrained JSON filling, the region detector reaches 0.854 AUROC, and adapting the top 10% of states captures 56.9% of the oracle opportunity (84.3% at 20% coverage). On Dream Carry RTL, the cardinality detector captures 64.3% at 0.802 AUROC. The instructive failure: Dream JSON Mode Eval with the score detector hits 0.891 AUROC, yet its top 10% of states contain zero oracle mass and yield zero lift. Binary detection and utility-relevant ranking are different skills.
End to end, selective adaptation lifts task-macro utility from 0.392 to 0.424 on LLaDA-8B (+3.2 points), 0.344 to 0.384 on LLaDA-1.5 (+4.0), and 0.411 to 0.467 on Dream (+5.6). Following the diagnostic on every state does worse (0.370, 0.334, 0.447), and in one displayed setting it drops LLaDA-8B from 0.365 to 0.160. The value sits in judging when to intervene, not in adapting itself.
For MDM inference work, this is a map plus a warning. The map is the five-axis taxonomy, which places the scattered unmasking tricks in one coordinate system so future comparisons can state which axis a method actually moves. The warning: average detectability sits near a coin flip, unconditional adaptation loses utility, and the profitable shape is fixed-by-default with deviation on a small fraction of states, concentrated on structured tasks. Teams eyeing adaptive decoding should run this kind of opportunity audit before investing. The measurement protocol, one-step utility with rollout cross-fitting plus the separation of action selection from opportunity detection, transfers directly to other inference-time adaptation research.
The authors state the main one plainly: predictability is limited overall, so the conclusion is conditional. Selective adaptation pays only where the detector is accurate and the diagnostic action is useful, and most settings fail one or both conditions. The action-selection viability screen used held-out outcomes, which the authors flag as exploratory rather than confirmatory.
Open questions remain. Most of the ten tasks lean structured or synthetic (CSV filling, HTML tag closing, carry circuits, sparse-mask recovery, multi-span cloze); HumanEval and GSM8K are the only widely known benchmarks, and open-ended generation is untested. The planning axis is untouched, which is exactly where search-based methods live. The diagnostics rest on idealized assumptions, each works in only a subset of settings, and none of the seven combinations passing the screen replicates across all three models.