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

cs.LGOct 27, 2026

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 ·

Abstract

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

Panel of Figure 1 — 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

P_{\mathrm{low}}

, 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

P_{\mathrm{low}}

. Hidden-state MSE over the cumulative teacher-committed set

C_{t+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.

In-domain mathematics reasoning
	GSM8K			MATH500
Method	Acc. $\uparrow$	RTL $\downarrow$	Time $\downarrow$	Acc. $\uparrow$	RTL $\downarrow$	Time $\downarrow$
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	$L_{\mathrm{T}}=T_{\mathrm{T}}=1024$
Student training window	$L_{\mathrm{S}}=512$
Teacher response-prefix cap	$c=L_{\mathrm{S}}$
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. $\uparrow$	RTL $\downarrow$	Time $\downarrow$	Acc. $\uparrow$	RTL $\downarrow$	Time $\downarrow$
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

Diffusion language model decoding generates text through iterative parallel refinement, balancing fewer denoising steps with accuracy, coherence, and control.

See all 250 papers

LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models
Mar 24, 2026Shaik AmanMasked Diffusion Models Diffusion Language Model Decoding
Independent Researcher, Nellore, Andhra Pradesh, India
Mask-Aware Policy Gradients for Diffusion Language Models
Jul 16, 2026Haran Raajesh, Kulin Shah, Adam Klivans +1Masked Diffusion Models RL for Language Model Reasoning
Department of Computer Science The University of Texas at Austin
Reliable Parallel Decoding in Masked Diffusion Language Models
Sep 29, 2026Zhenghao He, Bohan Liu, Guangzhi Xiong +1Masked Diffusion Models Diffusion Language Model Decoding
Department of Computer Science, University of Virginia

Mar 24, 2026J/K move · Enter open · S save

LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models

Shaik Aman

Independent Researcher, Nellore, Andhra Pradesh, India

Masked Diffusion Models Diffusion Language Model Decoding LLM Reasoning

Open paper AskPDF arXiv

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.

Abstract

Diffusion Language Model Decoding

Explore similar work

LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models

Mask-Aware Policy Gradients for Diffusion Language Models

Reliable Parallel Decoding in Masked Diffusion Language Models