Échantillonnage parcimonieux des logits : Accélération de la distillation des connaissances dans les LLM

Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs

March 21, 2025

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

cs.AI

Résumé

La distillation des connaissances peut s'avérer une technique rentable pour extraire le savoir des grands modèles de langage, à condition que les logits de sortie de l'enseignant puissent être pré-calculés et mis en cache. Cependant, l'application réussie de cette méthode au pré-entraînement reste largement inexplorée. Dans ce travail, nous démontrons que les approches naïves pour la distillation parcimonieuse des connaissances, comme la mise en cache des probabilités Top-K, bien qu'intuitives, fournissent des estimations biaisées de la distribution de probabilité de l'enseignant à l'élève, entraînant des performances et une calibration sous-optimales. Nous proposons une méthode basée sur l'échantillonnage d'importance, appelée `Random Sampling Knowledge Distillation', qui fournit des estimations non biaisées, préserve le gradient en espérance, et nécessite le stockage de logits significativement plus parcimonieux. Notre méthode permet un entraînement plus rapide des modèles élèves avec une surcharge marginale (<10%) par rapport à l'entraînement basé sur l'entropie croisée, tout en maintenant des performances compétitives par rapport à la distillation complète, pour une gamme de tailles de modèles allant de 300M à 3B.

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.