Muestreo de Logits Dispersos: Aceleración de la Distilación de Conocimiento en Modelos de Lenguaje de Gran Escala

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

Resumen

La destilación de conocimiento puede ser una técnica rentable para transferir conocimiento en Modelos de Lenguaje de Gran Escala, si los logits de salida del profesor pueden precomputarse y almacenarse en caché. Sin embargo, la aplicación exitosa de esto durante el preentrenamiento sigue siendo en gran parte inexplorada. En este trabajo, demostramos que los enfoques ingenuos para la destilación de conocimiento disperso, como almacenar en caché las probabilidades Top-K, aunque intuitivos, proporcionan estimaciones sesgadas de la distribución de probabilidad del profesor al estudiante, lo que resulta en un rendimiento y calibración subóptimos. Proponemos un método basado en muestreo por importancia llamado `Destilación de Conocimiento por Muestreo Aleatorio', que proporciona estimaciones no sesgadas, preserva el gradiente en expectativa y requiere almacenar logits significativamente más dispersos. Nuestro método permite un entrenamiento más rápido de los modelos estudiantes con un sobrecosto marginal (<10%) en comparación con el entrenamiento basado en entropía cruzada, manteniendo un rendimiento competitivo en comparación con la destilación completa, en un rango de tamaños de modelo desde 300M hasta 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.