反向 KL 策略蒸馏：统计保证与局限

arXiv:2610.30844v1 公告类型：新投稿 摘要：反向 Kullback-Leibler（KL）蒸馏在语言模型训练中展现出了良好的实证效果，但其统计保证与局限仍缺乏深入理解。尤其是，两项设计选择的影响尚不明确：获取教师的完整动作分布，是否比仅观察教师赋予采样动作的概率更有优势？训练轨迹是否必须来自学生（在线蒸馏），还是教师生成的轨迹（离线蒸馏）也足够？我们的工作为理解这些问题背后的理论迈出了第一步。我们研究了两种反馈模型下在线反向 KL 蒸馏的样本复杂度：教师提供其完整动作分布，或仅提供学生所生成动作的概率。我们为前者建立了极小极大界，并表明后者可能需要多得多的数据。相比之下，无论教师提供完整分布，还是仅提供教师所生成动作的概率，离线反向 KL 蒸馏都能达到相同的极小极大速率。关键区别在于，教师生成的动作本身就能揭示教师的偏好，而对学生所选动作进行查询则可能无法观测到这些偏好。这凸显了在线蒸馏相较于离线蒸馏所面临的额外统计挑战。这些统计保证并不能仅靠优化标准蒸馏目标来获得：我们的下界表明，当仅基于过去由学生生成的轨迹来评估候选策略时，常见的在线目标可能失效。我们提出了轨迹损失截断和重要性比率截断技术，并结合策略空间学习，实现了最优的在线蒸馏保证。