论文

知识蒸馏 在训练中期有利于推理而不是事实回忆

Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

摘要

基于 logits 的 知识蒸馏 (KD) 用于在更强教师的监督下训练较小的语言模型 (LM),但其好处在整个训练阶段是否一致仍不清楚。通过对照实验,我们发现前向 Kullback-Leibler (KL) 蒸馏(标准 KD 公式)与经过训练的教师在训练中期(在精选语料库上自我监督学习的中间阶段)的行为有根本不同。令人惊讶的是,虽然相对于标准的下一个标记预测 (NTP),前向 KD 在 预训练 期间同时提高了推理和事实回忆,但它反而减慢了训练中期的事实回忆获取速度,尽管推理持续获得收益。我们将这一阶段的依赖归因于教师跨数据域的信心和学生不断发展的知识状态的不对称:教师对程序性数据比对知识密集型数据更有信心,而学生在训练早期就获得了低熵事实知识。为了缓解这种不平衡,我们提出了 Switch 蒸馏,这是一个简单的中期训练目标,它在教师有信心的情况下提取词元,使用教师预测熵作为轻量级路由信号,否则回退到交叉熵。 Switch 蒸馏 在不同教师规模下始终优于现有的 蒸馏 目标。相对于标准 NTP,它的推理性能提高了 1.61-1.71 倍,知识和常识性能提高了 1.13-1.19 倍,同时保留了 96.7-96.8% 的事实召回率。至关重要的是,这些好处在 后训练 之后仍然存在:切换 蒸馏 缩小了事实回忆差距,同时在推理、知识和常识方面分别保持 1.25-1.32 倍和 1.13-1.20 倍的增益。