Summary
Masked diffusion language models (MDLMs) have to decide both what word to generate and where in the sentence to place it, which is tricky when they have limited space to work with. The researchers found that when these models have a small window to generate text, they might fix parts of the answer too early, even before figuring out the full reasoning. To help models do better with small windows, the authors developed HiLo-RD, a method that teaches the small-window model by comparing it with a bigger-window version of itself and guiding it on when and where to place words. This approach helped the model get better at math problem-solving and code generation, producing more accurate and shorter answers.
What this means in practice
- •For math tutoring platforms: Improve automatic problem-solving accuracy and response length for math question answering using efficient small-window language models guided by larger-window reasoning.
- •For software development teams: Enhance code generation quality and brevity by training code generation models to better use limited generation contexts informed by longer context reasoning.
Authors
Xiang Li, Chengting Yu, Jungang Li, Yilun Zheng, Junzhuo Li, Xingyao Zhang, Wenbo Su, Xuming Hu
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.