Разреженная выборка логитов: ускорение дистилляции знаний в больших языковых моделях

Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs

March 21, 2025

Авторы: Anshumann, Mohd Abbas Zaidi, Akhil Kedia, Jinwoo Ahn, Taehwak Kwon, Kangwook Lee, Haejun Lee, Joohyung Lee

cs.AI

Аннотация

Дистилляция знаний может стать экономически эффективной техникой для передачи знаний в больших языковых моделях, если выходные логиты учителя могут быть предварительно вычислены и кэшированы. Однако успешное применение этого метода на этапе предварительного обучения остается в значительной степени неисследованным. В данной работе мы доказываем, что наивные подходы к разреженной дистилляции знаний, такие как кэширование топ-K вероятностей, хотя и интуитивно понятны, предоставляют смещенные оценки распределения вероятностей учителя для ученика, что приводит к неоптимальной производительности и калибровке. Мы предлагаем метод, основанный на важностном сэмплировании, под названием `Random Sampling Knowledge Distillation`, который предоставляет несмещенные оценки, сохраняет градиент в ожидании и требует хранения значительно более разреженных логитов. Наш метод позволяет ускорить обучение моделей-учеников с минимальными накладными расходами (<10%) по сравнению с обучением на основе перекрестной энтропии, при этом сохраняя конкурентоспособную производительность по сравнению с полной дистилляцией для моделей различных размеров — от 300 миллионов до 3 миллиардов параметров.

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.