PA1-OPD: One-Step Prefix-Augmented On-Policy Distillation for Language Model Pretraining
作者: Zhenhe Wu, Jian Yang, Yaping Jin, Xinlai Kang, Jie Hu, Hanting Chen
分类: cs.CL
发布日期: 2026-10-14
💡 一句话要点
提出PA1-OPD以增强语言模型预训练中的知识蒸馏效果
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 知识蒸馏 语言模型 预训练 教师反馈 模型优化 自然语言处理 深度学习
📋 核心要点
- 现有的知识蒸馏方法在学生模型的预测上下文中缺乏直接的监督，导致效果受限。
- PA1-OPD通过在每个选定前缀后引入教师反馈，增强了语料路径蒸馏的监督效果。
- 实验结果表明，PA1-OPD在多个数据集上显著提高了交叉熵性能，尤其是在选择高惊讶度根节点时。
📝 摘要（中文）
知识蒸馏（KD）通过匹配教师分布来改善语言模型预训练，但在学生预测引导的上下文中缺乏直接监督。本文提出PA1-OPD，它在每个选定前缀后通过教师反馈增强语料路径蒸馏。这种方法扩展了监督状态，同时保留了下一个标记预测和语料路径KD。数据锚定根支持并行采样，且可微分树注意力在不改变推理的情况下共享前缀计算。在两个继续预训练的设置中，全面覆盖在每个评估检查点上改善了C4、PG-19和Pile的持出交叉熵。第二阶段中，选择学生惊讶度最高的10%根节点使KD的宏观交叉熵增益增加了48.85%。当分支监督排除已在语料路径上的后继状态时，增益依然存在，支持反馈在观察上下文之外的价值。额外的深度或广度在不改变每项监督权重的情况下改善了整体结果，且收益因语料而异。这些发现支持局部学生引导的监督作为预训练蒸馏的补充。
🔬 方法详解
问题定义：本文旨在解决现有知识蒸馏方法在学生模型预测上下文中缺乏直接监督的问题，这限制了模型的学习效果。
核心思路：PA1-OPD通过在每个选定前缀后引入教师反馈，扩展了监督状态，同时保留了下一个标记预测和语料路径KD，从而提高了模型的学习效率。
技术框架：该方法的整体架构包括数据锚定根的并行采样和可微分树注意力机制，确保在不改变推理的情况下共享前缀计算。
关键创新：PA1-OPD的核心创新在于通过教师反馈增强了语料路径蒸馏的监督效果，尤其是在学生模型的预测上下文中，这与传统的KD方法形成了鲜明对比。
关键设计：在参数设置上，选择学生惊讶度最高的10%根节点以最大化增益，同时在损失函数和网络结构上保持了每项监督权重不变，从而实现了更好的性能提升。
🖼️ 关键图片
📊 实验亮点
实验结果显示，PA1-OPD在C4、PG-19和Pile数据集上显著提高了交叉熵性能，特别是在选择高惊讶度根节点时，KD的宏观交叉熵增益增加了48.85%。这些结果表明该方法在知识蒸馏领域的有效性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理中的文本生成、对话系统和机器翻译等。通过增强语言模型的预训练效果，PA1-OPD能够提升模型在实际应用中的表现，具有重要的实际价值和未来影响。
📄 摘要（原文）
Knowledge distillation (KD) improves language-model pretraining by matching teacher distributions at corpus prefixes, but provides no direct supervision at contexts induced by student predictions. We introduce PA1-OPD, which augments corpus-path distillation with teacher feedback after one student-sampled transition per selected prefix. This expands the supervised states while retaining next-token prediction and corpus-path KD. Data-anchored roots support parallel sampling, and differentiable tree attention shares prefix computation without changing inference. In two continued-pretraining settings, full coverage improves held-out cross-entropy over KD across C4, PG-19, and Pile at every evaluated checkpoint. In Stage2, selecting the 10% of roots with highest student surprisal adds 48.85% to KD's Macro cross-entropy gain over next-token training. Gains remain when branch supervision excludes successor states already on the corpus path, supporting the value of feedback beyond observed contexts. Additional depth or breadth improves aggregate results at unchanged per-term supervision weight, with benefits varying by corpus. These findings support local student-induced supervision as a complement to pretraining distillation.