针对大模型压缩中知识蒸馏的粗粒度问题,提出按词元动态调整损失的方法。
ToDi: Token-wise Distillation via Fine-Grained Divergence Control
- 按词元自适应融合正向与反向KL散度,提升分布对齐精度
- 在多个指令跟随任务上超越现有基线模型,性能更优
- 适合资源受限场景下的轻量化大模型部署
大语言模型虽表现优异,但因高延迟和高能耗难以在资源受限环境下部署。知识蒸馏通过将大模型的知识迁移至小模型缓解此问题。然而,传统方法如前向KL(FKL)和反向KL(RKL)对整个词表施加统一的差异损失,忽略词元级预测差异。通过梯度分析发现,FKL增强被低估的词元,而RKL抑制被高估的词元,二者具有互补性。基于此,我们提出词元级蒸馏(ToDi),采用基于教师-学生概率对数比的Sigmoid加权函数,动态组合每词元的FKL与RKL。ToDi能为每个词元精准选择合适偏差,实现精细化分布对齐。实验表明,ToDi在指令遵循基准测试中持续优于使用统一或粗粒度策略的近期蒸馏方法。大量消融研究与效率分析进一步验证其有效性与实用性。
原文摘要 · Abstract (English)
Large language models (LLMs) offer impressive performance but are impractical for resource-constrained deployment due to high latency and energy consumption. Knowledge distillation (KD) addresses this by transferring knowledge from a large teacher to a smaller student model. However, conventional KD, notably approaches like Forward KL (FKL) and Reverse KL (RKL), apply uniform divergence loss across the entire vocabulary, neglecting token-level prediction discrepancies. By investigating these representative divergences via gradient analysis, we reveal that FKL boosts underestimated tokens, while RKL suppresses overestimated ones, showing their complementary roles. Based on this observation, we propose Token-wise Distillation (ToDi), a novel method that adaptively combines FKL and RKL per token using a sigmoid-based weighting function derived from the teacher-student probability log-ratio. ToDi dynamically emphasizes the appropriate divergence for each token, enabling precise distribution alignment. We demonstrate that ToDi consistently outperforms recent distillation baselines using uniform or less granular strategies across instruction-following benchmarks. Extensive ablation studies and efficiency analysis further validate ToDi's effectiveness and practicality.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。