Look Before You Select: Rethinking Vocabulary Sparsification in On-Policy Distillation

📄 arXiv: 2609.34386v1 📥 PDF

作者: Yongliang Miao, Shuang Liu, Yanguang Liu, Yandong Bai, Mengnan Du

分类: cs.CL, cs.LG

发布日期: 2026-09-28

💡 一句话要点

提出SparseOPD以解决长序列蒸馏中的内存瓶颈问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 策略蒸馏 全词汇纠正 内存优化 长序列处理 机器学习 自然语言处理 模型训练

📋 核心要点

  1. 现有的策略蒸馏方法在处理长序列时，因全词汇反向传播导致内存消耗过大，影响模型性能。
  2. SparseOPD通过在选择反向传播的token logits之前，先进行全词汇的教师纠正，从而优化内存使用。
  3. 在六个不同任务中，SparseOPD在任务平均准确率上超越了现有方法，并显著降低了反向传播的内存需求。

📝 摘要（中文）

本论文提出了一种新的方法SparseOPD，用于在策略蒸馏中进行全词汇的教师纠正。现有方法在处理长序列时，因需要对所有token的logits进行反向传播而导致内存消耗过大。SparseOPD通过在选择需要反向传播的token logits之前，先构建全词汇的纠正信息，从而有效减少内存使用。该方法在数学、化学问答和多模态推理等六个任务中表现优异，任务平均准确率超越了现有的采样token和TopK方法，并在某些情况下与全词汇方法相当或更优。

🔬 方法详解

问题定义：现有的策略蒸馏方法在长序列处理时，因需要对所有token的logits进行反向传播，导致内存消耗过大，影响训练效率和模型性能。

核心思路：SparseOPD的核心思想是先进行全词汇的教师纠正，然后根据纠正的重要性选择需要反向传播的token logits，这样可以有效减少内存使用并保持性能。

技术框架：SparseOPD的整体流程包括：首先构建全词汇的纠正信息，但不保留其反向图；然后根据纠正的大小选择token；最后，仅对选定的token logits进行反向传播。

关键创新：SparseOPD的主要创新在于通过纠正的重要性来选择token，而不是依赖于学生模型的概率分布，这一设计显著减少了内存消耗并提高了训练效率。

关键设计：在具体实现中，采用了带符号的残差补偿来保持纠正的总量，并通过纠正感知的预算分配在不同位置分配稀疏支持，确保了模型的稳定性和性能。

📊 实验亮点

在实验中，SparseOPD在六个任务的平均准确率上超越了采样token和TopK方法，且在4B数学任务中，梯度余弦相似度达到99%。此外，8K全参数分析显示其反向传播内存需求降低了70.5%。

🎯 应用场景

SparseOPD的研究成果在多个领域具有广泛的应用潜力，尤其是在需要处理长序列的自然语言处理任务中，如机器翻译、对话系统和文本生成等。通过优化内存使用，该方法能够提升模型的训练效率，降低硬件成本，推动更复杂模型的应用与发展。

📄 摘要（原文）

On-policy distillation (OPD) uses teacher correction on student-generated responses. Full-vocabulary correction can provide important corrections even for tokens that the student assigns low probability, but backpropagating through all token logits becomes memory-intensive for long sequences. Existing memory-saving approaches estimate corrections from sampled tokens or restrict supervision to the student's TopK tokens, introducing sampling noise or changing the full-vocabulary correction. We introduce \textbf{SparseOPD}, which uses full-vocabulary teacher correction to determine which corrections matter before selecting the token logits to differentiate. SparseOPD first constructs the full-vocabulary correction without retaining its backward graph, then selects tokens by correction magnitude rather than student probability. Signed residual compensation preserves the total promoting and suppressing correction mass, while correction-aware budget allocation distributes the sparse support across positions. Finally, the update backpropagates only through the selected token logits. Across six task--scale settings spanning mathematics, chemistry QA, and multimodal reasoning, SparseOPD outperforms Sampled Token and TopK in task-average accuracy and matches or exceeds Full Vocabulary. Gradient cosine similarity reaches 99\% on 4B mathematics, while 8K full-parameter profiling shows 70.5\% lower backward memory.