Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs
Knowledge distillation is much cheaper if you can pre-compute and cache the teacher's output logits once, rather than running the teacher live during every student training step. The obvious way to keep that cache small — storing only the teacher's top-K probabilities per token — turns out to be biased: it systematically distorts the target distribution the student learns from, hurting both accuracy and calibration. We prove this formally, then propose an importance-sampling scheme that gives the student an unbiased estimate of the full teacher distribution while storing as few as 12 logits per token. It preserves the gradient in expectation and matches full-distribution distillation quality at a fraction of the storage and compute cost, making distillation practical at pretraining scale.
