学习蒸馏什么：用于大型语言模型自蒸馏的双层 Top-K Token 选择

arXiv:2610.07247v1 公告类型：新 摘要：大型语言模型展现出了强大的推理能力，但其高昂的推理成本使知识蒸馏成为在资源受限场景下将这些能力迁移到紧凑型模型的重要方法。基于策略的数据自蒸馏进一步减少了对外部大型教师模型的依赖，同时提升了紧凑型语言模型的推理能力。然而，现有方法通常要么对所有 token 位置进行统一蒸馏，要么使用固定的启发式标准选择 token，并对选定位置赋予相同的蒸馏强度，而不是自适应地学习哪些 token 最有利于蒸馏。为解决这些局限，我们提出 BiToK-SD（Bilevel Top-K Token Selection for Self-Distillation），一种基于双层优化的 token 选择方法，用于学习在基于策略的自蒸馏过程中应在哪些位置进行蒸馏。具体而言，BiToK-SD 被表述为一个双层优化问题：下层问题将 Top-K token 选择建模为基于可微阈值的松弛，使选定位置能够随着学生策略的演化而自适应调整；上层问题则在选定位置执行知识蒸馏。在数学推理基准测试上的实验表明，在所有对比方法中，BiToK-SD 取得了最佳平均性能，同时仅需额外增加少量计算开销。