提出稀疏对数采样，加速大模型蒸馏训练且不损失精度。
Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs
- 用重要性采样替代贪心选Top-K，避免教师概率分布偏差。
- 仅存稀疏对数即可实现与全量蒸馏相当的性能，训练提速显著。
- 适合追求高效训练的大模型学生模型，尤其300M-3B规模。
知识蒸馏在大型语言模型中具有成本效益，若教师输出对数可预计算并缓存。然而，将其应用于预训练仍鲜有探索。本文证明，朴素的稀疏蒸馏方法（如缓存Top-K概率）会引入教师概率分布的偏差，导致学生模型性能与校准度下降。为此，我们提出基于重要性采样的随机采样知识蒸馏方法，能提供无偏估计，保持梯度期望一致，且只需存储极稀疏的对数。该方法在300M至3B模型规模上，相比交叉熵训练增加不到10%开销，实现显著更快的训练速度，同时性能媲美完整蒸馏。
原文摘要 · Abstract (English)
Knowledge distillation can be a cost-effective technique to distill knowledge in Large Language Models, if the teacher output logits can be pre-computed and cached. However, successfully applying this to pre-training remains largely unexplored. In this work, we prove that naive approaches for sparse knowledge distillation such as caching Top-K probabilities, while intuitive, provide biased estimates of teacher probability distribution to the student, resulting in suboptimal performance and calibration. We propose an importance-sampling-based method `Random Sampling Knowledge Distillation', which provides unbiased estimates, preserves the gradient in expectation, and requires storing significantly sparser logits. Our method enables faster training of student models with marginal overhead (<10%) compared to cross-entropy based training, while maintaining competitive performance compared to full distillation, across a range of model sizes from 300M to 3B.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。