Sparse Logit Sampling: Versnelling van Kennisdistillatie in LLM's

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

Samenvatting

Kennisdistillatie kan een kosteneffectieve techniek zijn om kennis te destilleren in grote taalmodelen, indien de uitvoerlogits van de leraar vooraf kunnen worden berekend en opgeslagen. Het succesvol toepassen hiervan tijdens voorafgaande training blijft echter grotendeels onontgonnen. In dit werk tonen we aan dat naïeve benaderingen voor sparse kennisdistillatie, zoals het opslaan van Top-K waarschijnlijkheden, hoewel intuïtief, bevooroordeelde schattingen van de waarschijnlijkheidsverdeling van de leraar aan de student geven, wat resulteert in suboptimale prestaties en calibratie. We stellen een op importance sampling gebaseerde methode voor, genaamd `Random Sampling Knowledge Distillation', die onbevooroordeelde schattingen biedt, de gradiënt in verwachting behoudt en aanzienlijk sparser logits vereist om op te slaan. Onze methode maakt snellere training van studentmodellen mogelijk met een marginale overhead (<10%) in vergelijking met training op basis van kruisingsentropie, terwijl competitieve prestaties worden behouden in vergelijking met volledige distillatie, over een reeks modelgroottes van 300M tot 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.