Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Rethinking Selective Knowledge Distillation

Almog Tavor, Itay Ebenspanger, Neil Cnaan, Mor Geva

TL;DR

Across a suite of benchmarks, SE-KD often improves accuracy, downstream task adherence, and memory efficiency over dense distillation over autoregressive LLMs, and extending this approach across the class and sample axes yields complementary efficiency gains that make offline teacher caching feasible.

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.

Rethinking Selective Knowledge Distillation

TL;DR

Across a suite of benchmarks, SE-KD often improves accuracy, downstream task adherence, and memory efficiency over dense distillation over autoregressive LLMs, and extending this approach across the class and sample axes yields complementary efficiency gains that make offline teacher caching feasible.

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.
Paper Structure (41 sections, 10 equations, 8 figures, 16 tables)

This paper contains 41 sections, 10 equations, 8 figures, 16 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. KD with Position Selection
  4. Curriculum Learning
  5. Uncertainty-Guided Position-Weighting KD
  6. KD with Class Sampling
  7. KD with Sample Selection & Weighting
  8. A Framework for Selective Knowledge Distillation
  9. Problem Setup
  10. Key Choices for Selective Distillation
  11. Student Entropy Guided Selective KD
  12. Student Entropy-based Position Selection (SE-KD)
  13. Cross-Axis Selection
  14. Selective LM Head and Chunked Entropy Computation
  15. Experiments
  16. ...and 26 more sections

Figures (8)

  • Figure 1: Illustration of three key selection axes for knowledge distillation: (A) sample selection, (B) class sampling (RS-KD), (C) position-selective KD, and (D) our combined approach, SE-KD$_{\text{3X}}$, which integrates sample, class, and position selection. Blue cells denote active (selected) supervision, light gray indicates inactive but included elements, and dark gray denotes filtered samples.
  • Figure 2: Position-axis budget sweep. Average validation accuracy after distilling on 80M FineWeb-Edu tokens as a function of the supervised position budget $k\%$. We compare Top-$k\%$ student-entropy (SE-KD) and Top-$k\%$ reverse-KL, with Full KD and RandomPos as reference. The teacher accuracy is 77.0.
  • Figure 3: Sample-axis budget sweep. Average validation accuracy after distilling on 80M FineWeb-Edu tokens as a function of the sample-selection budget $\ell\%$. Only the top-$\ell\%$ samples ranked by average student entropy are distilled; Full KD and RandomSmp are shown for reference. The teacher accuracy is 77.0.
  • Figure 4: Temperature ablation for Full KD. We compare $T{=}2.0$ vs. $T{=}1.0$ and report average accuracy over five benchmarks (ArcEasy, GSM8K, HellaSwag, PIQA, and LAMBADA OpenAI).
  • Figure 5: Cross-entropy mixing ablation for Full KD. We compare $\alpha_{\mathrm{CE}}=1-\lambda{=}0.1$ vs. $\alpha_{\mathrm{CE}}{=}0.0$ and report the same average accuracy metric. This study uses a smaller 10M-token run and is included as a sanity check rather than a fully converged comparison.
  • ...and 3 more figures