Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models

📄 arXiv: 2610.03665v1 📥 PDF

作者: Seo Hyun Kim, Sunwoo Hong, Younwoo Choi, Chen-Hao Chao, Se-Young Yun, Rahul G. Krishnan

分类: cs.LG, cs.CL

发布日期: 2026-10-02

备注: EMNLP 2026 Main (Oral)

💡 一句话要点

提出Pivot-SD以解决掩蔽扩散语言模型的信用分配问题

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

关键词: 掩蔽扩散语言模型 自蒸馏 信用分配 信息增益 深度学习

📋 核心要点

  1. 现有的掩蔽扩散语言模型在去噪过程中面临信用分配挑战，导致训练效率低下。
  2. Pivot-SD通过选择高影响力的承诺（枢轴）进行自蒸馏，优化了训练过程。
  3. 使用仅200个问题和四次回放，Pivot-SD在数学和代码基准测试中超越了全序列SFT和预算匹配的扩散RL基线。

📝 摘要（中文）

掩蔽扩散语言模型（dLMs）为复杂推理提供了一种有前景的并行替代方案，但在去噪过程中面临信用分配挑战。现有的后训练方法通常未能有效利用去噪信号，导致训练效率低下。本文提出了Pivot-SD，一个高效的离线自蒸馏框架，仅监督高影响力的承诺（枢轴），通过信息增益度量选择枢轴，并采用交叉熵和目标不可能性损失进行训练。实验表明，Pivot-SD在数学和代码基准测试中显著提升了LLaDA-8B-Instruct的性能。

🔬 方法详解

问题定义：本文旨在解决掩蔽扩散语言模型在去噪过程中面临的信用分配问题。现有方法通常未能有效利用去噪信号，导致训练效率低下。

核心思路：Pivot-SD的核心思路是通过选择高影响力的承诺（枢轴）进行自蒸馏，利用信息增益度量来选择这些枢轴，从而优化训练过程。

技术框架：Pivot-SD的整体架构包括信息增益度量模块、枢轴选择模块和训练模块。信息增益度量用于评估去噪过程中各个承诺对剩余掩蔽位置的不确定性减少的贡献。枢轴选择模块根据该度量选择高影响力的承诺进行训练。

关键创新：Pivot-SD的主要创新在于仅监督高影响力的承诺，而不是整个去噪步骤。这种方法使得训练更加高效，能够更好地聚焦于对最终输出影响最大的部分。

关键设计：在训练过程中，成功轨迹的枢轴使用交叉熵损失进行训练，而失败轨迹的枢轴则使用目标不可能性损失，未触及失败轨迹的其余部分。这种设计确保了训练的高效性和针对性。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果显示，Pivot-SD在仅使用200个问题和四次回放的情况下，显著提升了LLaDA-8B-Instruct的性能，超越了全序列SFT和预算匹配的扩散RL基线，展示了其在数学和代码基准测试中的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和代码生成等。通过提高掩蔽扩散语言模型的训练效率，Pivot-SD可以在实际应用中提升模型的推理能力和响应质量，具有广泛的实际价值和未来影响。

📄 摘要（原文）

Masked diffusion language models (dLMs) offer a promising parallel alternative to autoregressive models for complex reasoning. However, they face a distinct credit-assignment challenge, since a few commitments during denoising sharply reduce the uncertainty over the remaining masked positions and shape much of the response. Most post-training recipes for dLMs do not use this signal to decide which tokens to train on: they typically train on the final text or assign rewards to whole denoising steps, rather than selecting the individual commitments that shape the response. We introduce Pivot-SD, an efficient offline self-distillation framework that supervises only these high-impact commitments (pivots). Pivot-SD selects pivots using an information-gain metric measuring uncertainty reduction over the remaining masked positions. Pivots from successful trajectories are trained with cross-entropy, and pivots from failed trajectories with targeted unlikelihood, leaving the rest of the failed trajectory untouched. Using only 200 questions and four rollouts each, Pivot-SD improves LLaDA-8B-Instruct over full-sequence SFT and budget-matched diffusion RL baselines across math and code benchmarks.