HiLo-RD: Cross-Budget Self-Distillation for Masked Diffusion Language Models
Authors: Xiang Li, Chengting Yu, Jungang Li, Yilun Zheng, Junzhuo Li, Xingyao Zhang, Wenbo Su, Xuming Hu
Organizations: The Hong Kong University of Science and Technology (Guangzhou) · Zhejiang University · Nanyang Technological University · Alibaba Group · The Hong Kong University of Science and Technology
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.
Figures & tables
Figure 1: More room, shorter and more accurate answers. At L=T=512 (blue) and L=T=1024 (orange), (a) shows response-length shares in 32-token bins on common axes and (b) compares accuracy and mean readable length. Values are approximate readings from the archived aggregate plot.
Figure 2: Boundary-focused attention precedes tail commitment. Visible-window expansion relaxes this focus at the same state and query (left vs. middle), while native high-budget decoding attends more to earlier response positions at the matched update (right). Panels share an attention scale; red dashed lines mark the original low-budget boundary.
Figure 3: Attention shifts toward the committed tail. Changes from the pre-commitment update, averaged across layers and 145 valid traces. The dashed line marks commitment; 0 is immediately after it.
Figure 4: Tail delay improves low-budget recovery. Correct continuations in 155 cases selected for an incorrect native 512-token answer and a correct native 1024-token answer. Both interventions use L=T=512 .
Figure 5: HiLo-RD overview. Left: Budget-Induced Premature Commitment. For the same prompt, low-budget decoding concentrates attention near its response boundary Plow , fixes an incorrect answer early, and produces a response that fills its allocated window. With a larger budget, attention remains on earlier reasoning positions; the same model commits the correct answer at the final decoding stage and produces a much shorter response. Colors indicate normalized commitment time. Right: Cross-budget representation self-distillation. A frozen teacher processes a high-budget state, while the trainable student receives the crop ending at Plow . Hidden-state MSE over the cumulative teacher-committed set Ct+1 , including newly and previously committed positions, distills the reliable high-budget trajectory encoded in the teacher representations into the low-budget student. The trained student follows a coherent low-budget path and reaches the correct answer.
GSM8K
MATH500
Method
Acc. ↑
RTL ↓
Time ↓
Acc. ↑
RTL ↓
Time ↓
In-domain mathematics reasoning
vanilla
68.84
362.8
73.35
28.40
479.6
31.09
Masked SFT
70.13
366.8
72.10
30.60
478.3
31.03
Diffu-GRPO
72.93
325.4
63.53
29.20
482.1
31.19
OPSD
69.90
360.7
73.25
30.20
478.5
30.96
Table 1: Low-budget reasoning and transfer. All methods use the same low-budget pure diffusion setting and fast rollout. Acc. and Pass@1 are percentages; RTL is mean readable tokens, and Time is seconds per request. Domain separates mathematics and code; HiLo-RD identifies our method; Best marks the highest task score or shortest response per benchmark. Bold and underlined values mark the best and second-best distinct values in each column, including ties.
Figure 6: Quality–latency trade-off. Color denotes dataset; shape and fill identify model and rollout. Linear marker size represents RTL. Lines pair Full/Fast results. Timings await normalization.
Figure 7: Successive cross-budget self-distillation on GSM8K. Each cycle reuses the same 1,000 training examples and acquires new high-budget supervision from the preceding model. Across C0–C3, low-budget accuracy rises while readable length and latency fall; high-budget decoding remains stronger than C0 and becomes progressively more efficient.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Paired premature-commitment traces, Cases 1–8. Left: low budget; right: high budget. Headers give RTL. Ribbons and excerpts encode commitment timing in final response order: purple is early, yellow late, and gray denotes special tokens. Timing is normalized per response, and panel widths are rescaled. Cases are selected for contrasting outcomes and commitment order.
Figure 9: Paired premature-commitment traces, Cases 9–16. The encoding follows Figure 8 . Low-budget incorrect answers are fixed before much of the intervening reasoning; the paired high-budget answers are correct and committed near completion. RTL provides the absolute length comparison.
Component
Setting
Backbone
LLaDA-8B-Instruct
Teacher budget
LT=TT=1024
Student training window
LS=512
Teacher response-prefix cap
c=LS
Retained states
Up to K=8 , including the initial state
Representation
Final RMSNorm, before vocabulary projection
Appendix
Table 2: GSM8K reference training configuration. Teacher and Student distinguish the two budget views. The remaining settings document trajectory selection and adapter training for the full-training-set run. The 1,000-example successive-distillation study uses a separate training subset.
Low budget
High budget
Cycle
Acc. ↑
RTL ↓
Time ↓
Acc. ↑
RTL ↓
Time ↓
C0 (vanilla)
68.84
362.8
73.35
79.61
278.4
66.05
C1
74.37
309.3
60.16
81.05
250.7
59.18
C2
76.19
275.5
53.90
81.05
235.7
55.63
C3
78.09
254.2
49.71
80.74
222.9
52.47
Appendix
Table 3: Successive distillation improves the low-budget quality–length trade-off. C0–C3 reuse the same 1,000 GSM8K training examples. Accuracy is in percent, RTL is mean readable tokens, and Time is seconds per request. C3 identifies the final cycle; Best marks the highest accuracy or shortest response within each budget. Bold indicates the best value in each column, including ties.
Diffusion language model decoding generates text through iterative parallel refinement, balancing fewer denoising steps with accuracy, coherence, and control.
Masked diffusion language models (MDLMs) generate text by iteratively unmasking tokens from a fully masked sequence. Their standard confidence-based unmasking strategy systematically defers high-entropy logical connective tokens, degrading reasoning performance. We introduce LogicDiff, an inference-time method that replaces confidence-based unmasking with logic-role-guided unmasking. A lightweight classification head (4.2M parameters, 0.05% of the base model) predicts the logical role of each masked position (premise, connective, derived step, conclusion, or filler) from the base model's hidden states with 98.4% accuracy, and a dependency-ordered scheduler unmasks tokens in logical order. In zero-shot settings, LogicDiff improves LLaDA-8B-Instruct accuracy from 22.0% to 60.7% on GSM8K (+38.7 percentage points) and from 23.6% to 29.2% on MATH-500 (+5.6 pp), with less than 6% speed overhead. However, with 8-shot chain-of-thought prompting, the baseline reaches approximately 70% and LogicDiff provides no additional improvement. Analysis reveals that few-shot prompting implicitly resolves the same ordering problem that LogicDiff explicitly addresses, and that fixed role-based ordering can cause premature commitment to numerical values before sufficient context is available. Our results characterize the Flexibility Trap as primarily a zero-shot phenomenon and identify context-adaptive ordering as a key direction for future work.
Shaik Aman
Independent Researcher, Nellore, Andhra Pradesh, India
Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.
Haran Raajesh, Kulin Shah, Adam Klivans +1
Department of Computer Science The University of Texas at Austin
Masked diffusion language models (MDLMs) can generate text efficiently by predicting multiple masked tokens in parallel, but predictions from the same forward pass are not necessarily reliable when committed together. We study when parallel commitment is reliable. Our diagnostics show that confidence alone does not determine a reliable commitment order: confident predictions near the end of the sequence can fix an answer before its supporting computations are established, and downstream predictions become less reliable as the uncertainty of their upstream context grows. At the same time, a single forward pass can already resolve several masked tokens, and predictions that remain stable across the final layers are more likely to be correct. Based on these findings, we propose Reliable Parallel Decoding (RPD), a training-free method that selects candidates by layerwise prediction stability and final confidence, and commits them under a cumulative entropy budget over their preceding masked positions. RPD defers predictions with uncertain upstream context while committing the remaining candidates in parallel, without relying on a fixed block schedule. Across mathematical reasoning and code generation benchmarks on LLaDA and Dream, RPD achieves the highest decoding throughput among the evaluated methods while maintaining or improving accuracy.
Zhenghao He, Bohan Liu, Guangzhi Xiong +1
Department of Computer Science, University of Virginia