Campionamento Sparso dei Logit: Accelerazione della Distillazione della Conoscenza nei Modelli Linguistici di Grande Dimensione

Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs

March 21, 2025

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

cs.AI

Abstract

La distillazione della conoscenza può rappresentare una tecnica economicamente vantaggiosa per trasferire la conoscenza nei Large Language Models, qualora i logit di output del modello insegnante possano essere pre-calcolati e memorizzati. Tuttavia, l'applicazione efficace di questo approccio durante la fase di pre-training rimane in gran parte inesplorata. In questo lavoro, dimostriamo che approcci ingenui per la distillazione della conoscenza sparsa, come la memorizzazione delle probabilità Top-K, sebbene intuitivi, forniscono stime distorte della distribuzione di probabilità del modello insegnante allo studente, portando a prestazioni e calibrazione subottimali. Proponiamo un metodo basato sul campionamento per importanza, denominato `Random Sampling Knowledge Distillation', che fornisce stime non distorte, preserva il gradiente in attesa e richiede la memorizzazione di logit significativamente più sparsi. Il nostro metodo consente un addestramento più rapido dei modelli studente con un sovraccarico marginale (<10%) rispetto all'addestramento basato sull'entropia incrociata, mantenendo prestazioni competitive rispetto alla distillazione completa, su una gamma di dimensioni del modello che va da 300M a 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.