论文

BiToK-SD:学习自蒸馏时应选择的 Top-K 词元

Learning What to Distill: Bilevel Top-K Token Selection for Self-Distillation in Large Language Models

摘要

大语言模型表现出了强大的推理能力,但其较高的推理成本使得知识蒸馏成为资源受限场景下将这种能力转移到紧凑模型的重要方法。 同策略自蒸馏进一步减少了对外部大型教师模型的依赖,同时提高了紧凑语言模型的推理能力。然而,现有方法通常要么统一蒸馏所有 token 位置,要么使用固定启发式标准选择 token,为所选位置分配相同的蒸馏强度,而不是自适应地学习哪个 token 对蒸馏最有利。为了解决这些限制,我们提出了 BiToK-SD(用于自蒸馏的双层 Top-K token 选择),这是一种基于双层优化的 token 选择方法,可学习在策略自蒸馏期间应在何处应用蒸馏。具体来说,BiToK-SD 被表述为双层优化问题,其中较低级别的问题将 Top-K token 选择建模为可微分的基于阈值的松弛,允许所选位置随着 学生模型策略的发展而适应,而上层问题则对所选位置执行知识蒸馏。数学推理基准测试表明,BiToK-SD 在所有比较方法中实现了最佳的平均性能,同时只需要轻量级的额外计算。

BiToK-SD:学习自蒸馏时应选择的 Top-K 词元:论文原图
图 1:训练期间的探索动态。 (a) 训练步骤上的 学生模型 熵。 (b) 所有 token 上的 学生模型–教师模型 熵差。 (c) 所选 token 上的 学生模型–教师模型 熵间隙。 BiToK-SD 保持比基线更高的 学生模型 熵和更大的 学生模型–教师模型 熵差距,表明策略探索更强。