What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
Masked diffusion language models generate tokens in parallel, but their efficiency depends critically on which predictions are committed at each step. This paper disentangles the design space of learned unmasking policies by treating them as distilled teachers and independently varying supervision targets, training states, and policy inputs. The authors find that trajectory-contraction labels and a confidence-plus-hidden-state input yield a lightweight policy that improves the accuracy–NFE frontier and end-to-end throughput once NFE drops by about 10%.
The paper presents a carefully controlled empirical study that isolates design choices for learned unmasking policies in masked diffusion language models. It introduces trajectory contraction, a stricter supervision target that improves the accuracy–NFE frontier mainly at low NFE, and combines it with confidence and hidden-state inputs in an 856K-parameter policy called OnTrack. While the controlled comparisons are methodologically sound, the gains are concentrated on moderate accuracy targets, and the policy shows sensitivity to approximate inference such as prefix caching.
The controlled ablation design is the paper’s main strength. By varying one axis at a time—supervision target, training states, and policy inputs—while holding architecture and capacity fixed, the authors cleanly attribute effects. The trajectory-contraction teacher is well-motivated by the exact statistic that 28.7% of terminal-agreement positives leave the reference trajectory when committed early. The latency profiling across policy sizes, prompt lengths, and caching regimes is thorough, and the code is publicly released.
The primary limitation is that the policy is trained only on mathematical reasoning data at a single generation length ($L=256$), and its advantage vanishes when both domain and length shift simultaneously. On some benchmarks, accuracy saturates below the reference: the authors note that "On MATH-500 on LLaDA and GSM8K on Dream, the accuracy of our policy saturates slightly below the reference, even with conservative thresholds." Additionally, prefix caching reduces accuracy by up to about two percentage points, while confidence-based baselines are unaffected; the authors hypothesize this is because hidden representations are perturbed by caching, but they do not validate or mitigate it. The study also excludes strong contemporary learned baselines such as LUP and TraceLock from the main experiments due to unavailable checkpoints, which weakens the empirical comparison to the full prior-art landscape.
The evidence generally supports the central claims about supervision and inputs. Trajectory contraction is justified both theoretically and with empirical statistics—"28.7% of terminal-agreement positives violate trajectory contraction"—and the accuracy–NFE frontiers in Figures 2 and 3 support the rankings of design choices. Comparisons to related work are fair where checkpoints exist (Fast-dLLM, EB-Sampler, Learn2PD), but the absence of LUP and TraceLock from the main evaluation—explained by the fact that "no trained checkpoints are publicly available for either method"—means the paper does not establish state-of-the-art gains against all learned baselines. Throughput experiments on the same hardware strengthen practical claims, though statistical significance intervals are only mentioned for a subset of results.
The authors release code at https://github.com/jinya1013/ontrack and report detailed hyperparameters, including training configuration, optimizer settings, and parameter grids. The backbone models and evaluation datasets are publicly available. Key details such as the AdamW learning rate of $3\times 10^{-5}$, batch size 16, and Fast-dLLM reference threshold of $0.9$ are provided. Potential blockers include the need for substantial GPU resources—"Training takes approximately 3 hours and 15 minutes on a single NVIDIA H200 GPU"—and the reliance on specific hardware for the exact latency profiling. Reproducing the reported throughputs will depend on matching the GPU and software stack, but the methodology is documented sufficiently for an independent implementation.
Masked diffusion language models generate tokens in parallel, but their efficiency depends on which predictions are committed. Learned unmasking policies improve the speed--accuracy trade-off, yet differ in supervision, inputs, and capacity, obscuring which choices matter. We view supervised policies as distilling a decoding procedure through a teacher and examine independently what the teacher labels as committable, on which states the policy imitates it, what the policy observes, and how much the number of function evaluations (NFE) must decrease to offset policy overhead. We find that labeling predictions matching the final output is overly permissive at low NFE, and hidden states are weak alone but can complement confidence. About a 10% NFE reduction suffices to offset policy overhead for policies up to 7.5M parameters across tested settings. We therefore train a policy on trajectory-contraction targets, which accept only commitments staying on a reference trajectory, using confidence and hidden states. On LLaDA-8B-Instruct, the resulting 856K-parameter policy achieves the best accuracy--NFE frontier and highest throughput on GSM8K, is competitive at moderate targets on MATH-500 and HumanEval, and largely transfers to Dream-7B-Instruct.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.