Sparse Logit Sampling: Beschleunigung der Wissensdistillation in LLMs

Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs

March 21, 2025

Autoren: Anshumann, Mohd Abbas Zaidi, Akhil Kedia, Jinwoo Ahn, Taehwak Kwon, Kangwook Lee, Haejun Lee, Joohyung Lee

cs.AI

Zusammenfassung

Wissensdestillation kann eine kosteneffiziente Technik sein, um Wissen in großen Sprachmodellen zu destillieren, wenn die Ausgabe-Logits des Lehrers vorab berechnet und zwischengespeichert werden können. Die erfolgreiche Anwendung dieser Methode auf das Vortraining bleibt jedoch weitgehend unerforscht. In dieser Arbeit zeigen wir, dass naive Ansätze für spärliche Wissensdestillation, wie das Zwischenspeichern von Top-K-Wahrscheinlichkeiten, obwohl intuitiv, verzerrte Schätzungen der Wahrscheinlichkeitsverteilung des Lehrers für den Schüler liefern, was zu suboptimaler Leistung und Kalibrierung führt. Wir schlagen eine auf Importance Sampling basierende Methode vor, die `Random Sampling Knowledge Distillation`, die unverzerrte Schätzungen liefert, den Gradienten im Erwartungswert erhält und deutlich spärlichere Logits speichert. Unsere Methode ermöglicht eine schnellere Ausbildung von Schülermodellen mit minimalem Overhead (<10%) im Vergleich zum Training auf Basis der Kreuzentropie, während sie eine wettbewerbsfähige Leistung im Vergleich zur vollständigen Destillation über eine Reihe von Modellgrößen von 300M bis 3B beibehält.

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.