Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs
Authors: name: Anshumann; name: Mohd Abbas Zaidi; name: Akhil Kedia; name: Jinwoo Ahn; name: Taehwak Kwon; name: Kangwook Lee; name: Haejun Lee; name: Joohyung Lee
Published: 2025-03-21 05:58:18

Abstract: 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.