2025/03/21 by Anshumann Anshumann, Anshumann, Mohd Abbas Zaidi +12 · 1 citation
Computer Science · #68T50 #Artificial Intelligence (cs.AI) #Computation and Language (cs.CL) #Data Mining Algorithms and Applications #Data Stream Mining Techniques #FOS: Computer and information sciences #I.2.7 #Imbalanced Data Classification Techniques #Machine Learning (cs.LG)
paper · pdf · doi:10.48550/arxiv.2503.16870
openalex publication_date 2025/03/21 · openalex created_date 2025/10/16 · openalex updated_date 2026/07/28
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.