HiLo-RD: Cross-Budget Self-Distillation for Masked Diffusion Language Models

📄 arXiv: 2610.36901v1 📥 PDF

作者: Xiang Li, Chengting Yu, Jungang Li, Yilun Zheng, Junzhuo Li, Xingyao Zhang, Wenbo Su, Xuming Hu

分类: cs.LG, cs.CL

发布日期: 2026-10-27

备注: 20 pages, 9 figures, 3 tables; Under review


💡 一句话要点

提出HiLo-RD以解决有限生成预算下的MDLMs问题

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 掩蔽扩散语言模型 自我蒸馏 生成预算 推理效率 数学基准测试 代码生成 模型优化

📋 核心要点

  1. 现有的掩蔽扩散语言模型在有限生成预算下面临预测内容和固定答案片段的双重挑战,导致推理效率低下。
  2. HiLo-RD通过将大窗口教师状态与小窗口视图配对,并在已承诺位置匹配隐藏表示,来优化模型的内容生成和承诺顺序。
  3. 在GSM8K基准测试中,HiLo-RD的准确率提高了9.4个百分点,同时平均可读长度减少约25%,显示出显著的性能提升。

📝 摘要(中文)

我们提出了HiLo-RD(高到低预算表示蒸馏),旨在改善在有限生成预算下的掩蔽扩散语言模型(MDLMs)。与自回归模型不同,MDLMs需要同时选择生成内容和确定何时固定答案。这带来了两个相互关联的挑战:在推理完成之前预测有用内容,以及决定何时固定影响后续预测的答案片段。我们识别出预算引起的过早承诺问题:即使正确的解决方案适合,模型也可能在解决中间步骤之前就固定靠近窗口边界的答案片段。这些片段成为后续预测的上下文,并重定向后续推理。现有的后训练方法无法直接转移同一模型在更大窗口下遵循的更强轨迹。HiLo-RD通过将大窗口教师状态与裁剪的小窗口视图配对,并在已承诺位置匹配其隐藏表示来解决这一问题。实验表明,HiLo-RD在数学基准测试中提高了低预算准确率,并缩短了响应时间;这些提升也扩展到代码生成任务上。

🔬 方法详解

问题定义:论文要解决的问题是掩蔽扩散语言模型在有限生成预算下的推理效率低下,现有方法无法有效处理预测内容和固定答案片段的挑战。

核心思路:论文的核心思路是通过高到低预算的表示蒸馏,利用大窗口的教师状态来指导小窗口的生成过程,从而提高模型的推理能力和准确性。

技术框架:整体架构包括两个主要模块:大窗口教师模型和小窗口学生模型。教师模型生成完整的上下文,而学生模型则在此基础上进行裁剪和学习。

关键创新:最重要的技术创新点在于通过匹配大窗口和小窗口的隐藏表示,解决了现有方法无法有效转移强轨迹的问题,实现了无外部教师的自我蒸馏。

关键设计:关键设计包括损失函数的设置,确保在已承诺位置的表示匹配,以及网络结构的优化,以适应不同窗口大小的输入。具体参数设置和训练策略在实验部分进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,在GSM8K基准测试中,HiLo-RD的准确率提高了9.4个百分点,同时平均可读长度减少约25%。这些结果表明,HiLo-RD在低预算条件下显著提升了模型的性能,且这一提升同样适用于代码生成任务。

🎯 应用场景

该研究的潜在应用场景包括自然语言处理中的文本生成、对话系统以及代码生成等领域。通过提高模型在有限预算下的推理能力,HiLo-RD能够在实际应用中提供更高效的响应,具有重要的实际价值和未来影响。

📄 摘要(原文)

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.