What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models

1 Introduction

• The study treats supervised unmasking as distillation: a teacher labels committable positions from reference trajectories, and a policy imitates those labels.It varies the teacher’s labels, the states used for imitation, and policy inputs independently.Source: §1 · pp. 1–2
• Trajectory contraction accepts a commitment only when it directly reaches a later state on the reference trajectory, addressing the limits of final-output matching.The study compares this target jointly with whether training uses reference trajectories or states induced by the labels.Source: §1 · p. 2
• The study compares confidence, probability margins, hidden representations, and confidence combined with hidden representations as policy inputs.It measures NFE and end-to-end throughput while profiling policy overhead.Source: §1 · p. 2
• On LLaDA-8B-Instruct, the 856K-parameter policy achieves the best accuracy–NFE frontier and highest throughput on GSM8K, with competitive moderate-target efficiency on MATH-500 and HumanEval.Its advantage largely transfers to Dream-7B-Instruct and to unseen generation lengths for mathematical reasoning.Source: §1 · p. 2
• 28.7% of final-output-matching positives leave the reference trajectory when committed early, while hidden representations are weak alone but can complement confidence.Policies up to 7.5M parameters pay off in throughput in tested settings once they reduce NFE by more than about 10%.Source: §1 · p. 2

2 Related Work

• Training-free methods use confidence, decoding dynamics, or dependencies to choose parallel commitments, while learned policies adapt where or when to commit tokens.Fast-dLLM uses confidence-thresholded decoding while keeping the backbone frozen.Source: §2 · p. 2
• Learned unmasking work spans reinforcement learning, supervised oracle signals, ranking objectives, and learned dependency models.Learn2PD uses token-correctness signals, while TraceLock derives labels from completed traces and uses hidden-state representations.Source: §2 · p. 2
• Table 1 compares learned-policy design choices, and its caption notes that existing methods differ simultaneously in inputs, capacity, commitment targets, and training states.The table also marks candidates introduced in this work and estimates some configurations from released settings.Source: Table 1 · p. 3

3.1 Problem Formulation

• A masked dLLM commits selected greedy predictions at each step, trading fewer backbone evaluations against possible quality loss from premature commitments.Committed tokens are not revised later.Source: §3.1 · p. 3

3.2 Design Space and Protocol

• Supervised policies distill a fixed decoding procedure, and the study varies commitment targets, training states, and inputs while controlling other design choices.Existing methods vary across several axes, making individual effects difficult to isolate; candidates share architecture, optimizer, and training schedule.Source: §3.2 · p. 3

3.3 How Should Supervision Be Constructed?

• Trajectory contraction improves the accuracy–NFE frontier mainly at low NFE; on MBPP, reference-trajectory supervision reaches the 98% target at 48.6 NFE versus 64.8 for teacher-induced trajectories.At the highest accuracy targets, terminal agreement can be more accurate.Source: Figure 2 · p. 4, §3.3 · pp. 4–5

3.4 What Information Should the Policy Use?

• Confidence plus hidden states reaches the 98% target on GSM8K with 42.1 NFE versus 52.9 for confidence alone, though input benefits vary across tasks and targets.Hidden states alone require substantially more NFE; adding the margin yields no consistent gain.Source: Figure 3 · p. 5, §3.4 · p. 5

3.3 How Should Supervision Be Constructed?

• Terminal agreement labels 28.7% of positive predictions outside the contraction set, while teacher-induced trajectories retain only 34% of contraction states and 25% of terminal-agreement states.The reduced state coverage may explain why reference trajectories transfer better to MBPP, though it has little effect on GSM8K.Source: §3.3 · p. 5

3.5 Do NFE Gains Translate to Speedups?

• A roughly 10% NFE reduction offsets the overhead of policies up to 7.5M parameters across prompt lengths and caching settings.Even a 300M-parameter policy needs about 28% fewer NFEs to break even.Source: §3.5 · pp. 5–6, Table 2 · p. 6

4 Putting the Findings Together

• OnTrack combines the study’s selected supervision and input choices while keeping policy overhead small.The findings determine trajectory-contraction supervision, confidence plus hidden-state inputs, and a compact policy.Source: §4 · p. 6

4.1 Policy Architecture

• The policy projects confidence and hidden representations into a shared 128-dimensional space, combines them, and produces commitment logits conditioned on mask state and decoding step.The policy has 856K parameters, about 0.01% of the backbone.Source: §4.1 · p. 6

4.2 Policy Training

• Training retains reference trajectories with correct final answers and applies trajectory-contraction targets at their states, optimizing binary cross-entropy over eligible positions.Targets mark positions in the furthest valid contraction.Source: §4.2 · p. 6

5.1 Experimental Setup

• Evaluation covers GSM8K and MATH-500 reasoning plus HumanEval and MBPP code, comparing Fast-dLLM and EB-Sampler with learned Learn2PD.The study reports accuracy or Pass@1, NFE, and TPS, measuring throughput on the same hardware and configuration.Source: §5.1 · p. 6

5.2 Main Results on LLaDA

• On LLaDA, the policy has the best GSM8K accuracy–NFE frontier and highest TPS across targets, with highest TPS also at moderate targets on MATH-500 and HumanEval.It reaches 98% of reference accuracy on GSM8K with the fewest NFE; Figure 4 shows accuracy versus NFE and TPS.Source: §5.2 · p. 7, Figure 4 · p. 7

5.3 Compatibility with Prefix Caching

• With four-shot prefix caching, the policy remains fastest at 98% reference accuracy; caching reduces its accuracy by up to about two points, and zero-shot Fast-dLLM wins at 98%.Table 3 reports highest measured TPS at 98% and 95% of uncached reference accuracy, with and without caching.Source: §5.3 · pp. 7–8, Table 3 · p. 7

5.4 Generalization Across Backbones

• On Dream-7B-Instruct, the policy has the best frontier on HumanEval and reaches 98% of reference accuracy on MATH-500 with the fewest NFE, while MBPP methods perform comparably.On GSM8K it reaches the 95% target with the fewest NFE but saturates slightly below reference accuracy.Source: §5.4 · p. 8, Figure 5 · p. 8

5.5 Generalization Across Sequence Lengths

• A policy trained at L = 256 retains best or comparable frontiers through 95% of reference accuracy at L = 128 and 512, but combined length and domain shifts leave it behind on code.At L = 128 it reaches 98% on MATH-500 with the fewest NFE; at L = 512 it saturates below 98% there.Source: §5.5 · p. 8, Figure 6 · p. 8

• Trajectory-contraction targets improve the accuracy–NFE frontier mainly at low NFE, with an additional out-of-domain gain from training on reference trajectories.Source: §6 · p. 8
• Hidden representations are weak alone but can complement confidence at high accuracy targets; policies up to 7.5M parameters pay off in throughput after reducing NFE by more than about 10% in tested settings.Source: §6 · p. 8
• The 856K-parameter policy reaches best or comparable efficiency at moderate targets across two backbones, transfers from mathematical reasoning to code, and converts NFE reductions into throughput gains.Source: §6 · p. 8

• Training uses mathematical reasoning data at L = 256 with a fixed block length; transfer works to code at that length and to other lengths on math, but not both shifts together.Training across multiple domains and lengths is left to future work.Source: §7 · p. 9
• On some tasks, accuracy saturates below reference even with conservative thresholds, making the policy most beneficial at moderate targets; prefix caching reduces its accuracy by up to about two points.The caching drop persists when the policy unmasks one token per step, and the authors suggest perturbed backbone hidden representations as a plausible cause.Source: §7 · p. 9

• Efficiency may lower energy use, but deployment scale could increase overall environmental impact; faster inference may also facilitate large-scale misleading or harmful generation.Source: §8 · p. 9
• The method leaves backbone parameters unchanged and inherits backbone biases and failure modes; effects of decoding changes on output safety are not explicitly evaluated.Source: §8 · p. 9
• The work uses publicly available datasets and models under their licenses, collects no new human-subject data, and reports author review of AI-assisted outputs.Source: §8 · p. 9

A.1 Policy Architecture

• The policy projects confidence and hidden representations into a shared space, adds the hidden stream to confidence, and processes the sequence with a DiT-style transformer block.The backbone remains frozen, and commitment is restricted to masked positions in the current block.Source: §A.1 · p. 12, Figure 7 · p. 12

A.2 Trajectory-Contraction Teacher

• For each reference state, the teacher selects the furthest future state reachable by a valid contraction and uses its commitment set as the binary target.Validity is monotone, so target construction requires no backbone evaluations beyond those on the reference trajectory.Source: §A.2 · pp. 12–13

A.3 Decoding

• At each step, the decoder commits eligible positions whose policy scores exceed γ, falling back to the highest-scoring eligible position when none exceed it.Eligibility is limited to masked positions in the current block, and selected positions receive the frozen backbone’s top-1 predictions.Source: §A.3 · p. 13

B.1 Experimental Setup

• The controlled comparisons evaluate accuracy–NFE frontiers against one-by-one confidence-based reference accuracy, varying one factor at a time.Evaluations use GSM8K, MATH-500, HumanEval, and MBPP, with accuracy targets at 98%, 95%, and 90% of the reference.Source: §B.1 · p. 13, Table 5 · p. 13

B.2 Commitment Target Construction

• Trajectory-contraction labels reject 28.7% of terminal-agreement positives that would leave the reference trajectory when committed early.Such violations occur in 56.0% of reference states; rates are 28.5% on GSM8K and 29.1% on MATH.Source: §B.2 · p. 13

B.3 Training-Trajectory Construction

• Teacher-induced trajectories retain 34.0% of reference states under trajectory contraction and 25.4% under terminal agreement.With trajectory contraction, induced states are a subsequence of the reference trajectory and inherit its targets; terminal-agreement states generally leave it and require recomputed predictions and targets.Source: §B.3 · pp. 13–14

B.4 Policy-Input Study

• Confidence and hidden-state inputs are each projected to 128 dimensions, and differences in policy parameter counts come only from input projections.The hidden representation is taken from the backbone’s final layer; its projection adds 524,416 parameters.Source: §B.4 · p. 14, Table 7 · p. 14

B.5 Per-Step Latency Profiling

• Per-step latency profiling includes backbone and policy forwards plus bookkeeping, with and without prefix caching; larger randomly initialized policies measure size-related overhead.Profiling uses LLaDA-8B-Instruct on one H200 GPU, batch size 1, generation length 256, and block length 32.Source: §B.5 · p. 14, Table 8 · p. 14

B.6 Results on MATH-500 and HumanEval

• On HumanEval, trajectory contraction with reference-state training has the best frontier, reaches 90% with the fewest NFE, and is the only configuration to reach 95%.On MATH-500, this combination has the best frontier through 95% but falls just short of 98%; terminal agreement on reference states exceeds reference accuracy.Source: §B.6 · p. 15
• On MATH-500, confidence plus hidden states has the best frontier through 95%, while adding margins reaches 98% with the fewest NFE and highest accuracy, exceeding the reference.On HumanEval, confidence-based inputs perform comparably through 90%; only margin-and-hidden-state combinations reach 95%, and none reaches 98%.Source: §B.6 · p. 15

C.1 Datasets and Prompt Templates

• Training uses GSM8K and MATH examples, while evaluation covers GSM8K, MATH-500, HumanEval, and MBPP with task-specific prompts and answer verification.Evaluation examples are not used for training; training prompts longer than 256 tokens are left-truncated.Source: §C.1 · pp. 15–16

C.2 Policy Optimization

• The policy is optimized jointly with AdamW, and training takes approximately 3 hours and 15 minutes on one NVIDIA H200 GPU.The backbone remains frozen during policy training.Source: §C.2 · p. 16, §C.3 · p. 16

C.3 Training and Decoding Configuration

• Policies are trained on 9,169 correct-answer reference trajectories, with no code-generation examples, making HumanEval and MBPP out-of-domain evaluations.The main configuration uses generation length 256 and block length 32; length-512 evaluation also applies End-of-Text Prediction.Source: §C.3 · p. 16, Table 11 · p. 16

C.4 Parameter Grids

• Commitment thresholds are swept over a fixed grid for OnTrack, Learn2PD, and Fast-dLLM, while EB-Sampler sweeps an entropy-budget parameter.The evaluated parameter grids are summarized in Table 12.Source: §C.4 · p. 16, Table 12 · p. 16

C.5 Throughput Measurement

• Throughput is measured on one H200 GPU as generated tokens divided by sampler time, including backbone and policy evaluation but excluding tokenization and scoring.Figure 10 reports GSM8K accuracy versus throughput under zero-shot and four-shot prompting; Figure 11 reports HumanEval and MBPP accuracy versus NFE at lengths 128 and 512.Source: §C.5 · pp. 16–17, Figure 10 · p. 17, Figure 11 · p. 17

C.6 Additional Results

• On GSM8K, the policy is more accuracy-sensitive to prefix caching than the baselines, while four-shot prompting brings substantial throughput improvements with caching.Source: §C.6 · p. 17
• HumanEval and MBPP frontiers are reported at generation lengths 128 and 512 for a policy trained at length 256.Source: §C.6 · p. 17

D Numerical Results

• Tables 13–26 list operating points underlying the supervision and input comparisons, backbone-transfer results, prefix-caching results, and sequence-length evaluations.Source: D Numerical Results · p. 17