HiLo-RD: Cross-Budget Self-Distillation for Masked Diffusion Language Models
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.
Authors
- Chengting Yu (ORCID: https://orcid.org/0009-0007-7210-879X)
- Xuming Hu (ORCID: https://orcid.org/0000-0001-6075-4224)
- Wenbo Su (ORCID: https://orcid.org/0009-0009-3800-7543)
- Yilun Zheng
- Jungang Li
- Junzhuo Li
- Xingyao Zhang
- Xiang Li
Publication Details
- Journal
- arXiv (Cornell University)
- Published
- 2026-10-27
- DOI
- https://doi.org/10.48550/arxiv.2610.36901
- Primary Topic
- Topic Modeling
- Type
- preprint