Masked diffusion models lock answers too early: HiLo-RD distills large-window trajectories, lifting GSM8K accuracy by 9.4 points
Synopsis
Masked diffusion language models can fix answer fragments before the supporting reasoning is resolved when the response window is small, and those fragments then become context for later predictions. HiLo-RD uses the same model's large-window trajectory as a frozen teacher and aligns its hidden representations to cropped small-window views at committed positions. At the same small inference budget, GSM8K accuracy rises from 68.84% to 78.24% while mean readable length falls from 362.8 to 271.3 tokens, and the gains extend to code generation without code-domain training.
Interpretation
Under a small response window, a masked diffusion language model can commit answer-tail fragments before filling in the intermediate reasoning, even when a correct solution would fit. With a larger window the same model defers the answer and reaches a shorter, correct solution, and existing post-training methods do not directly transfer that large-window trajectory into small-window decisions. In 155 paired cases where low-budget decoding is wrong and high-budget decoding is correct, delaying tail commitments at the original low budget recovers correct answers in 43.2% of cases versus 30.3% for a matched order control, a paired gain of 12.9 percentage points with a 95% bootstrap interval of 7.7 to 18.7.
HiLo-RD transfers hidden representations from the same model's large-window trajectory into small-window decoding, shaping both token prediction and commitment order. It uses no external teacher and does not apply a loss to isolated tokens or separately specified commitment labels; it aligns cross-budget trajectories at the representation level. Teacher and student start from the same checkpoint, the teacher generates trajectories at a large budget, the student receives cropped small-window states, and their final normalized hidden states are matched at positions the teacher has cumulatively committed after an update; training adapts LoRA parameters while the backbone and vocabulary head stay frozen.
At the same small inference budget, HiLo-RD raises accuracy and shortens readable responses at once. Against vanilla LLaDA-8B-Instruct and four post-training methods, Masked SFT, Diffu-GRPO, OPSD and dOPSD, it reaches the highest accuracy and shortest readable responses on both GSM8K and MATH500. On GSM8K accuracy rises from 68.84% to 78.24% as mean readable length falls from 362.8 to 271.3 tokens; on MATH500 accuracy rises from 28.40% to 32.40%; the mathematics-trained checkpoint, without code-domain training, lifts MBPP pass@1 from 38.60% to 41.40%.
Perspective
The demonstrated setting is pure diffusion decoding with LLaDA-8B-Instruct, trained separately on the full GSM8K training set and the original MATH training set and evaluated on the GSM8K test set and MATH500 respectively; code results come from mathematics-trained checkpoints without additional code-domain training. The successive-distillation experiment reuses a fixed subset of 1,000 GSM8K training examples and is a separate sequence from the full-training-set main comparison.
The attainable compression is task-dependent: on HumanEval pass@1 is slightly below dOPSD and the reduction in readable length is marginal. Reported timing differences come from different servers and remain subject to normalization, and the completed-prefix preservation property depends on matched decoding and first-EOT extraction. Representation alignment transfers only the part of higher-budget information predictable from the shared prompt and cropped state, and a parameterized student may still incur approximation and optimization error.