Rethinking Selective Knowledge Distillation

TL;DR

This paper systematically analyzes selective knowledge distillation in large language models, introducing a new method (SE-KD) that improves efficiency and performance by focusing on important tokens, classes, and samples.

Contribution

It disentangles importance signals and selection policies in selective KD, proposing SE-KD and its extensions to enhance efficiency and accuracy in LLM training.

Findings

SE-KD improves accuracy and efficiency over dense distillation.

Extending SE-KD across axes yields further efficiency gains.

Reduces wall time by 70%, memory by 18%, and storage by 80% without performance loss.

Abstract

Growing efforts to improve knowledge distillation (KD) in large language models (LLMs) replace dense teacher supervision with selective distillation, which uses a subset of token positions, vocabulary classes, or training samples for supervision. However, it remains unclear which importance signals, selection policies, and their interplay are most effective. In this work, we revisit where and how to distill in autoregressive LLMs. We disentangle selective KD along the position, class, and sample axes and systematically compare importance signals and selection policies. Then, guided by this analysis, we identify underexplored opportunities and introduce student-entropy-guided position selection (SE-KD). Across a suite of benchmarks, SE-KD often improves accuracy, downstream task adherence, and memory efficiency over dense distillation. Extending this approach across the class and sample axes (SE-KD 3X) yields complementary efficiency gains that make offline teacher caching feasible. In practice, this reduces wall time by 70% and peak memory by 18%, while cutting storage usage by 80% over prior methods without sacrificing performance.