Source-linked AI summary
HiLo-RD: Cross-Budget Self-Distillation for Masked Diffusion Language Models
Xiang Li, Chengting Yu, Jungang Li, Yilun Zheng, Junzhuo Li, Xingyao Zhang, Wenbo Su, Xuming Hu
TL;DR
Masked diffusion models must choose both what to predict and when to commit it, and limited budgets can lead to premature answer fragments that redirect later reasoning. HiLo-RD distills the same model’s large-window representations into low-budget decoding. It improves low-budget accuracy while shortening responses, with reported gains extending to code generation without code-domain training.
Problem
A low-budget masked diffusion model can commit answer fragments before resolving intermediate reasoning, and existing post-training methods do not directly transfer the same model’s more effective large-window computation into small-window decisions.
Source: §1 · pp. 1–2Method
HiLo-RD pairs large-window teacher states with cropped small-window views and matches hidden representations at committed positions using the model’s own trajectories.
Source: §1 · p. 2Results
At the same small inference budget, HiLo-RD raises GSM8K accuracy from 68.84% to 78.24% while reducing mean readable length from 362.8 to 271.3 tokens.
Source: §1 · p. 2Takeaways & Limitations
The results show that the same model’s larger-window computation can be transferred into more accurate, more concise small-window responses.
Source: §1 · p. 2
Abstract
from arXiv · showhide
We propose HiLo-RD (High-to-Low Budget Representation Distillation) to improve masked diffusion language models (MDLMs) under a limited generation budget. Unlike autoregressive models, MDLMs choose both what to generate and where to commit it. This creates two coupled challenges: predicting useful content before reasoning is complete, and deciding when to fix answer tokens that shape later predictions. We identify budget-induced premature commitment: even when a correct solution fits, the model can fix answer fragments near the window boundary before resolving intermediate steps. These fragments become context for later predictions and redirect subsequent reasoning. Existing post-training methods do not directly transfer the stronger trajectories that the same model follows with a larger window. HiLo-RD addresses this gap by (1) pairing large-window teacher states with cropped small-window views and (2) matching their hidden representations at committed positions. This supervision guides both token prediction and commitment order without an external teacher. Across mathematics benchmarks, HiLo-RD improves low-budget accuracy and shortens responses; these gains extend to code generation without code-domain training. On GSM8K, accuracy rises by 9.4 percentage points as mean readable length falls by about 25%. Attention diagnostics and fixed-budget order interventions support the proposed mechanism. Together, these results show that large-window reasoning can improve small-window decoding.
1 INTRODUCTION
Masked diffusion models must choose both what to predict and when to commit it, so limited response windows can prompt premature answer fragments that redirect later reasoning. HiLo-RD transfers the same model’s large-window representations to low-budget decoding, improving answer quality while reducing response length.
1 INTRODUCTION
- A small response window can prompt the model to commit answer fragments before resolving intervening reasoning, even when a correct solution fits.The paper calls this budget-induced premature commitment.Source: §1 · pp. 1–2
- The paper reports that larger windows are associated with higher accuracy and lower mean readable length across the aggregate pattern.Figure 1 compares response-length shares, accuracy, and mean readable length for 512- and 1024-token windows.Source: Figure 1 · p. 1, §1 · p. 2
- Committed answer tokens become context for later predictions, so a premature fragment can shape the reasoning subsequently filled around it.The model must choose both what to predict and when to fix answer fragments.Source: §1 · p. 2
- HiLo-RD pairs large-window teacher states with cropped small-window views and matches hidden representations at committed positions, using the model’s own trajectories.This targets low-budget token prediction and commitment decisions without an external teacher.Source: §1 · p. 2
- At the same small inference budget, HiLo-RD raises GSM8K accuracy from 68.84% to 78.24% while reducing mean readable length from 362.8 to 271.3 tokens.The paper also reports gains in code generation without code-domain training.Source: §1 · p. 2
2 PRELIMINARIES AND RELATED WORK
The preliminaries describe masked diffusion decoding as iterative prediction and commitment across a bidirectional response window. They then explain how window length and update count shape computation, and situate HiLo-RD among diffusion-model post-training approaches.
2.1 MASKED DIFFUSION DECODING
- Masked diffusion decoding predicts tokens at masked positions and commits selected predictions, with a confidence rule determining which positions are filled at each update.Commitment need not follow response order; the window is processed afresh as decoding proceeds.Source: §2.1 · pp. 2–3
- A decoding budget specifies response-window length and planned updates; when window length equals update count, one position is committed per update, while fewer updates require multi-position commitments.Enlarging the window changes both visible context and the positions competing for commitment.Source: §2.1 · p. 3
2.2 DECODING BUDGETS AND REALIZED COMPUTATION
- The decoder recomputes the full prompt and response window at each update, and bidirectional states prevent reuse of a KV cache across updates.Each avoided update saves one full-window forward pass.Source: §2.2 · p. 3
2.3 RELATED WORK
- Prior work includes masked SFT, reward-based post-training, and self-distillation, while HiLo-RD pairs teacher states with cropped small-window views and matches hidden states.The paper’s single-block setting keeps all masked response positions eligible at each update.Source: §2.3 · p. 3
3 BUDGET-INDUCED PREMATURE COMMITMENT AND ITS CONSEQUENCES
In selected paired cases, low-budget decoding can commit a response-tail fragment before completing intermediate reasoning, while high-budget decoding reaches a correct answer. The section links boundary-focused attention and later attention to committed tail content with improved recovery when tail commitment is delayed.
3 BUDGET-INDUCED PREMATURE COMMITMENT AND ITS CONSEQUENCES
- The analysis uses 155 cases selected for incorrect native 512-token answers and correct native 1024-token answers, and traces premature tail commitment before intervening reasoning is complete.Tail Delay excludes original response-tail positions during the first 10% of remaining updates; the matched-order control perturbs selection at the same update indices without a tail-specific rule.Source: §3 · p. 3, §3.3 · p. 4
3.1 BOUNDARY ATTENTION PRECEDES PREMATURE COMMITMENT
- Before the first tail commitment, low-budget attention forms a band near the response-window boundary; adding visible masked positions weakens it, while native high-budget decoding attends more to earlier reasoning.The within-state comparison isolates visible response space, and the native trajectory shows computation with that space available throughout decoding.Source: §3.1 · pp. 3–4, Figure 2 · p. 4
3.2 COMMITTED TAIL CONTENT ATTRACTS SUBSEQUENT ATTENTION
- Across 145 valid traces, attention to the tail rises by 1.91 percentage points immediately after commitment, while attention to the established prefix falls by 0.41 points.At +32 updates, the changes reach +4.19 points to the tail and −1.87 points to the prefix.Source: §3.2 · p. 4, Figure 3 · p. 4
3.3 TAIL DELAY RESTORES CORRECT ANSWERS
- Tail Delay reaches the correct answer in 67 of 155 selected cases (43.2%), versus 47 (30.3%) for the matched order control, a paired gain of 12.9 percentage points (95% bootstrap interval: 7.7–18.7).Both interventions retain the 512-token window, predictor, and decoding schedule.Source: Figure 4 · p. 4, §3.3 · p. 4
4 HiLo-RD: CROSS-BUDGET REPRESENTATION SELF-DISTILLATION
HiLo-RD transfers high-budget trajectory representations into a cropped low-budget student, supervising teacher-committed positions. Completion-aware stopping reduces unnecessary decoding after a response prefix is complete.
4.1 CROSS-BUDGET REPRESENTATION DISTILLATION
- HiLo-RD aligns high-budget teacher representations with a student’s cropped low-budget view, using the model’s own trajectories to guide decoding.The teacher is frozen, and the student receives the crop ending at the smaller response-window boundary.Source: Figure 5 · p. 5, §4.1 · p. 5
- The student and teacher representations are paired at corresponding response positions preserved by cropping, with teacher states computed before cropping.The overlapping positions retain their token values and position indices.Source: §4.1 · p. 5
- The hidden-state loss supervises cumulative teacher-committed positions, including previously committed tokens and newly selected positions still masked in the teacher input.The representation objective transfers trajectory information rather than training on isolated tokens or separately specified commitment labels.Source: §4.1 · pp. 5–6
4.2 EFFICIENT TRAJECTORY ACQUISITION
- Completion-aware stopping avoids unnecessary rollout after response completion, while teacher acquisition can stop once the student-sized supervision prefix is filled.Clean-EOT stopping preserves the extracted response under matched decoding and first-EOT extraction; the teacher-prefix cap may truncate a longer solution.Source: §4.2 · p. 6
5 EXPERIMENTS
The experiments evaluate low-budget reasoning, cross-domain transfer, rollout efficiency, and successive self-distillation. HiLo-RD improves low-budget benchmark results while shortening responses, and completion-aware stopping and iterative distillation further address efficiency and refinement.
5.1 EXPERIMENTAL SETUP
- The evaluation covers GSM8K and MATH500 reasoning, HumanEval and MBPP code generation without code-domain training, and compares HiLo-RD with vanilla and four post-training methods.Training uses the corresponding mathematics training sets; comparisons use pure diffusion decoding, with fast rollout also assessed for vanilla and HiLo-RD.Source: §5.1 · pp. 6–7, §5 · p. 6
5.2 LOW-BUDGET PERFORMANCE AND CROSS-DOMAIN TRANSFER (RQ1 & RQ2)
- HiLo-RD improves over vanilla on all four benchmarks while producing shorter readable responses in the shared low-budget pure-diffusion fast-rollout setting.Table 1 compares accuracy or Pass@1, readable token length, and time per request across mathematics and code tasks.Source: §5.2 · p. 7, Table 1 · p. 7
- On GSM8K and MATH500, HiLo-RD achieves the highest accuracy and shortest readable responses among compared methods.The passage reports that stronger low-budget reasoning need not require longer readable solutions, with compression varying by task.Source: §5.2 · p. 7
- On MBPP, HiLo-RD combines the strongest Pass@k with the shortest responses; on HumanEval, it improves over vanilla but trails dOPSD slightly in Pass@k.The readable-length reduction on HumanEval is marginal, unlike the stronger compression reported for MBPP.Source: §5.2 · pp. 7–8
5.3 EFFICIENT TRAJECTORY ACQUISITION (RQ3)
- Fast rollout reduces recorded latency most clearly on GSM8K and MBPP, with little change on HumanEval; GSM8K HiLo-RD retains identical accuracy and readable length.The figure pairs Full/Fast results by dataset and model; reported timing differences remain subject to cross-server normalization.Source: §5.3 · p. 8, Figure 6 · p. 8
5.4 TOWARD SELF-EVOLVING QUALITY–BUDGET REFINEMENT (RQ4)
- Across C0–C3, low-budget accuracy rises as readable length and latency fall, while high-budget decoding remains above C0 accuracy and becomes more efficient.Each cycle reuses the same fixed subset of 1,000 GSM8K training examples and acquires new high-budget supervision from the preceding model.Source: §5.4 · p. 8, Figure 7 · p. 9
6 CONCLUSION AND FUTURE WORK
The paper identifies budget-induced premature commitment as a way early answer fragments can redirect masked-diffusion reasoning, and uses same-model high-budget representations to guide low-budget decoding. It concludes that accuracy gains with shorter responses extend across mathematics and code, while richer intermediate supervision may support longer self-distillation sequences.
6 CONCLUSION AND FUTURE WORK
- The paper links premature answer-tail commitments to redirected subsequent reasoning and distills the same model’s high-budget trajectories into low-budget decoding.Decoding traces, attention diagnostics, and commitment-order interventions support the proposed mechanism.Source: §6 · p. 9
- Successive self-distillation further improves the quality–budget frontier; richer intermediate supervision may help sustain refinement across longer cycles.The proposed refinement uses the model’s own trajectories.Source: §6 · p. 9
- Generative AI tools were used only to polish the manuscript, and at least two authors independently reviewed every AI-assisted passage.The authors state that the tools did not generate or analyze experimental results, research artifacts, or scientific claims.Source: AI USE STATEMENT · p. 10
A.1 MAIN DIAGNOSTIC COHORT AND CONTROLLED COMPARISONS
The diagnostics use paired low- and high-budget failures and successes to test how visible space and commitment order affect decoding at a fixed budget.
A.1 MAIN DIAGNOSTIC COHORT AND CONTROLLED COMPARISONS
- The 155 paired cases share incorrect native L=T=512 and correct native L=T=1024 outcomes, and probes replay the state before the first ordinary tail commitment.Window expansion appends 512 visible MASK positions while preserving the prompt, committed contents, query, and token targets; native high-budget trajectories are decoded separately.Source: §A.1 · p. 12
- Tail attention rises 1.908 percentage points immediately after commitment in 140/145 cases, while prefix attention changes by −0.413 points.At +32 updates, tail and prefix attention changes are +4.194 and −1.875 points, respectively; later offsets include additional natural commitments.Source: §A.1 · p. 13
- At fixed L=T=512, Tail Delay yields 12.903 percentage points more correct outcomes than matched-order control, with a 95% interval of [7.742,18.710].Tail Delay postpones selection of original response-tail positions for the first 10% of remaining updates; the control excludes its own current top-scoring position at matched offsets.Source: §A.1 · p. 13
A.2 COMPLEMENTARY CHECKS ON VISIBLE SPACE AND TAIL READING
Complementary checks distinguish visible generation space from added tensor length, characterize irreversible commitments and paired traces, and specify how training trajectories are sampled and weighted.
A.2 COMPLEMENTARY CHECKS ON VISIBLE SPACE AND TAIL READING
- Visible expansion reduces mean tail-target probability from 46.30% to 0.25%, while invisible padding gives 46.24%; the next position moves away from the tail in 19/20 cases.The front target changes from 47.65% to 52.12% with visible expansion and to 47.79% with invisible padding; blocking attention to committed tail keys repairs one case and breaks one among 20, leaving aggregate accuracy unchanged.Source: §A.2 · pp. 13–14
A.3 IRREVERSIBLE COMMITMENT AND REACHABLE COMPLETIONS
- Each irreversible commitment shrinks the set of compatible correct completions, and a wrong fixed fragment can remove successful continuations while changing later predictions and position scores.Budget-Induced Premature Commitment is boundary-near commitment while intervening reasoning remains unresolved despite sufficient response space.Source: §A.3 · p. 14
A.4 PAIRED TRACE EXAMPLES AND VISUALIZATION PROTOCOL
- In the selected 16 GSM8K pairs, low-budget answers are incorrect and committed early, while high-budget answers are correct and committed near completion.The figures illustrate outcome and ordering contrasts, not the prevalence of Budget-Induced Premature Commitment; readable lengths are 510 for each low-budget case and 162–355 for high-budget cases.Source: §A.4 · p. 14, Figure 8 · p. 15, Figure 9 · p. 16
B.1 TRAJECTORY SAMPLING AND OBJECTIVE NORMALIZATION
- Training pairs each full-window teacher state with its cropped small-window student view, supervising committed positions and normalizing each state’s loss by the number of supervised positions.The positive recipe accepts teacher responses whose extracted answers pass verification; this certifies the extracted answer, not every intermediate reasoning step.Source: §B.1 · p. 17
B.2 COMPLETED-PREFIX STOPPING
The appendix states conditions for completed-prefix preservation, bounds local token-distribution differences from representation matching, and clarifies what smaller-window students can retain and what optional repulsion does not guarantee.
B.2 COMPLETED-PREFIX STOPPING
- With identical initialization, update schedule, tie-breaking, and coupled randomness, fast decoding stopped after R updates matches the full decoder’s trajectory prefix, and completed prefixes remain unchanged.Under clean EOT, extraction strictly before the first EOT returns the same token sequence after either run.Source: §B.2 · p. 17
B.3 REPRESENTATION MATCHING AND TOKEN DISTRIBUTIONS
- Representation matching bounds local token-distribution differences under a shared fixed vocabulary head, but does not guarantee position-selection agreement or matching autonomous student states.The bound concerns finite raw logits, not hard logit suppression or discrete selection directly.Source: §B.3 · p. 18
B.4 WHAT THE SMALLER-WINDOW REPRESENTATION CAN RETAIN
- The unrestricted optimal smaller-window predictor is the conditional mean of the teacher representation given the prompt, cropped state, and queried position; it need not reconstruct details tied to omitted positions.A parameterized student may also incur approximation and optimization error, and the regression identity does not guarantee reasoning quality on autonomous rollouts.Source: §B.4 · p. 18
B.5 OPTIONAL OUTCOME-CONDITIONED REPULSION
- Optional negative representation repulsion stops increasing separation beyond its margin, but does not guarantee lower probability for a specified erroneous trajectory.The supplied negative teacher retains the long window while reducing the step count, for example (Lneg, Tneg)=(1024,128), not (128,128).Source: §B.5 · p. 19
C.1 TRAINING CONFIGURATION AND DEPLOYMENT
The appendix records the GSM8K reference setup and deployment, qualifies how latency should be interpreted, and explains readable-length and pass@k measurement conventions.
C.1 TRAINING CONFIGURATION AND DEPLOYMENT
- Main experiments use separate mathematics-trained checkpoints for GSM8K and MATH500, while code evaluation uses mathematics-trained checkpoints without code-domain adaptation.The documented GSM8K configuration is not a complete record for every dataset and baseline; the student adapter can be merged for autonomous decoding.Source: §C.1 · p. 19
C.2 COMPUTATION AND LATENCY
- Full/Fast latency ratios are not controlled speedup estimates because measurements include different servers and depend on hardware, kernels, batching, and prompt length.Teacher-acquisition savings require stopping the decoding loop, rather than retaining fewer states after a full rollout.Source: §C.2 · p. 20
C.3 READABLE LENGTH, EXTRACTION, AND PASS@𝑘
- Readable token length differs from occupied response positions and forward-pass counts, while pass@k interpretation depends on sample count, decoding distribution, and code-test protocol.Per-completion latency is not the total cost of obtaining a multi-sample pass@k score.Source: §C.3 · p. 20
D SUCCESSIVE DISTILLATION: PROCEDURE AND MEASUREMENTS
Successive distillation merges each cycle’s adapter into the next teacher–student initialization and reuses a fixed GSM8K subset. Across cycles, low-budget accuracy rises as readable length and latency fall.
D SUCCESSIVE DISTILLATION: PROCEDURE AND MEASUREMENTS
- Each cycle merges the preceding adapter into the backbone before initializing a new frozen teacher and trainable student; the teacher remains fixed within that cycle.The resulting merged model supplies supervision for the next cycle.Source: D SUCCESSIVE DISTILLATION: PROCEDURE… · p. 20
- Low-budget accuracy improves monotonically from C0 to C3 as readable length and latency decrease, while high-budget decoding retains an accuracy advantage over C0.The cycles reuse the same 1,000 GSM8K training examples; Table 3 reports accuracy, mean readable tokens, and seconds per request.Source: D SUCCESSIVE DISTILLATION: PROCEDURE… · p. 20, Table 3 · p. 20