Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss
作者: Bakbergen Ryskulov, Iker García-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R. Garcia, Antonio Tiene, Román Orús
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-08-04
备注: Patent Application Pending. EP26382987.1
🔗 代码/项目: GITHUB
💡 一句话要点
提出高效知识蒸馏方法以提升小型语言模型训练效率
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 知识蒸馏 小型语言模型 训练效率 内存优化 深度学习
📋 核心要点
- 现有的知识蒸馏方法在训练效率和内存使用上存在显著瓶颈，限制了小型语言模型的应用。
- 本文提出的离线KD方法通过缓存教师模型的logits，显著提高了训练速度和吞吐量，同时降低了内存消耗。
- 实验结果表明，新的分块KL损失方法使得模型能够处理更长的上下文，提升了训练的可扩展性和效率。
📝 摘要（中文）
小型语言模型在延迟、成本和本地部署等限制下常常是唯一的选择，但通常并非从头开始训练，而是通过知识蒸馏（KD）恢复压缩模型。本文提出了两种系统贡献以提高蒸馏训练的效率。首先，离线KD通过缓存教师模型的前K个logits进行训练，速度比在线蒸馏快约29%，且在单个H200 GPU上吞吐量提高41%。其次，提出了一种融合的分块KL损失，避免了全词汇大小的logit张量生成，使得峰值内存与序列长度线性相关，从而在单个GPU上实现四倍的上下文训练（32,768个token）。
🔬 方法详解
问题定义：本文旨在解决小型语言模型在知识蒸馏过程中训练效率低和内存消耗高的问题。现有的在线蒸馏方法在训练速度和内存使用上存在显著不足，限制了其在实际应用中的可行性。
核心思路：提出离线KD方法，通过缓存教师模型的前K个logits进行训练，避免了实时计算logits的开销。同时，设计了一种融合的分块KL损失，减少内存峰值需求，使得训练过程更加高效。
技术框架：整体架构包括两个主要模块：离线KD模块和分块KL损失模块。离线KD模块负责缓存教师模型的logits，而分块KL损失模块则在训练过程中动态计算损失，确保内存使用最优化。
关键创新：最重要的技术创新在于提出了离线KD和分块KL损失的结合，前者提高了训练速度，后者降低了内存消耗。这一组合显著提升了小型语言模型的训练效率。
关键设计：在损失函数设计上，采用了融合的分块KL损失，避免了全词汇大小的logit张量生成，确保内存使用与序列长度线性相关。此外，实验中还进行了多种损失设计和序列打包的支持性消融实验。
🖼️ 关键图片
📊 实验亮点
实验结果显示，离线KD方法在单个H200 GPU上训练速度提高约29%，吞吐量提升高达41%。同时，分块KL损失使得模型能够处理长达32,768个token的上下文，显著提升了训练的可扩展性和效率。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、对话系统和实时翻译等场景，尤其是在资源受限的设备上。通过提高小型语言模型的训练效率，能够更好地满足实际应用中的延迟和成本要求，推动智能助手和其他AI应用的发展。
📄 摘要（原文）
Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-$K$ logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29\% faster per iteration, and reaching up to 41\% higher throughput on a single H200 GPU. Second, we introduce a \emph{fused, chunked KL loss} that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32{,}768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.