Amostragem de Logits Esparsa: Acelerando a Distilação de Conhecimento em LLMs

Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs

March 21, 2025

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

cs.AI

Resumo

A destilação de conhecimento pode ser uma técnica econômica para transferir conhecimento em Modelos de Linguagem de Grande Escala, se os logits de saída do professor puderem ser pré-computados e armazenados em cache. No entanto, a aplicação bem-sucedida desse método durante o pré-treinamento permanece amplamente inexplorada. Neste trabalho, demonstramos que abordagens ingênuas para destilação esparsa de conhecimento, como o armazenamento em cache das probabilidades Top-K, embora intuitivas, fornecem estimativas tendenciosas da distribuição de probabilidade do professor para o aluno, resultando em desempenho e calibração subótimos. Propomos um método baseado em amostragem por importância, chamado `Random Sampling Knowledge Distillation`, que fornece estimativas não tendenciosas, preserva o gradiente em expectativa e requer o armazenamento de logits significativamente mais esparsos. Nosso método permite um treinamento mais rápido dos modelos alunos com uma sobrecarga marginal (<10%) em comparação com o treinamento baseado em entropia cruzada, mantendo um desempenho competitivo em relação à destilação completa, em uma variedade de tamanhos de modelo, de 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.