PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Selective Knowledge Distillation

Almog Tavor, Itay Ebenspanger, Neil Cnaan, Mor Geva

arXiv:2602.01395v1cs.CL

TL;DR

Selective KD for LLMs lacks clear evidence about which importance signals and selection policies best allocate teacher supervision. The paper systematically compares these choices, introduces student-entropy-guided SE-KD and three-axis SE-KD3X, and reports accuracy and efficiency gains while identifying scope limitations.

  • Problem

    It remains unclear which token-importance signals and position-selection policies most reliably benefit autoregressive LLM distillation and how they interact with class and sample sparsification.

  • Method

    The paper builds a framework for selective KD and introduces SE-KD, which ranks positions by student entropy, plus SE-KD3X, which selects samples, positions, and classes.

  • Results

    Across benchmarks, student-entropy top-20% selection improves average accuracy to 64.8 vs. 64.4 for Full KD and perplexity to 6.9 vs. 7.3, while SE-KD3X supports substantial efficiency gains.

  • Takeaways & Limitations

    Student-entropy-guided position selection is the most reliable overall strategy, and combining position, class, and sample sparsification yields favorable accuracy–efficiency trade-offs.

  • Takeaways & Limitations

    The evaluation uses one teacher–student pair and a fixed supervision budget, so validation across model families, scales, and longer contexts remains necessary.

Abstract

from arXiv · show

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.

1. Introduction

The paper revisits selective knowledge distillation for autoregressive LLMs, asking which importance signals and selection policies best determine where teacher supervision is applied. It introduces student-entropy-guided selection and combines position, class, and sample selection for efficiency.

  • Uniform supervision at every token position is often suboptimal, motivating selective KD based on position importance.
  • The study disentangles selective KD across positions, classes, and samples, comparing importance signals, selection policies, and their interactions.
  • SE-KD selects positions using student entropy, while SE-KD3X jointly selects samples, positions, and classes.
  • 64.8 vs. 64.4 for Full KD average evaluation accuracy and 6.9 vs. 7.3 perplexity result from selecting the top-20% positions by student entropy.Calibration changes from 0.273 to 0.276, while fewer logits require computation.
  • SE-KD and SE-KD3X remain competitive in on-policy and task-specific math-reasoning distillation settings.
  • 70% lower total wall time is achieved by student-entropy-based sample selection on 80M-token distillation while maintaining performance.Class-sampled offline caching cuts storage by 99.96%.

2. Related Work

Prior work reduces KD cost or adapts supervision through position, class, and sample selection, but generally studies individual heuristics or axes. The paper situates its unified comparison within these lines of research.

  • Position-selective KD has used student cross-entropy, teacher uncertainty, discrepancy, and curricula to focus supervision on selected sequence positions.
  • Unlike prior approaches that use uncertainty for loss reweighting, the paper uses uncertainty solely as a ranking signal for explicit position selection.
  • Class-sampling methods reduce output-distribution cost, with RS-KD using importance sampling to avoid biased gradient estimates and improve calibration.
  • Sample-selection methods reduce teacher queries by targeting informative or difficult examples using uncertainty or teacher–student cross-entropy discrepancy.

3. A Framework for Selective Knowledge Distillation

The framework organizes selective KD by how teacher information is aligned and sparsified across classes, positions, samples, and features. It formalizes selective objectives while identifying multi-axis selection and student entropy as underexplored directions.

  • The framework proposes a general formulation that encapsulates existing selective KD methods and exposes implementation design choices.
  • Selective KD can select classes at positions, positions within samples, and samples within the training set.
  • Five design axes are alignment criterion, position, class, sample, and feature selection.Feature-axis selection is explicitly left unexplored for future work.
  • Position selection is defined by an importance metric u(t) and a policy mapping scores to supervised-position indicators.
  • No prior work in the reviewed comparison selects across more than one axis, while student entropy remains underexplored as a distillation signal.
  • Table 1 records whether each selective KD method acts on the position, class, or sample axes.

4. Student Entropy Guided Selective KD

This section introduces student-entropy-guided selective KD and extends it across position, class, and sample axes. It also presents selection-aware optimizations that reduce logit-related computation and memory.

  • Student Entropy-based Position Selection (SE-KD): SE-KD uses student entropy as a position-importance signal for selective knowledge distillation.The method selects positions according to the student’s uncertainty.
  • Student Entropy-based Position Selection (SE-KD): SE-KD selects the top-k% most uncertain positions in each sequence for distillation.The threshold is chosen so exactly ⌈k(L −1)⌉ positions satisfy the selection mask.
  • Student Entropy-based Position Selection (SE-KD): Per-sequence loss normalization maintains a fixed supervision budget during selective distillation.
  • Cross-Axis Selection: SE-KD3X combines class selection, position selection, and student-entropy-based sample selection into a unified multi-axis method.Sample ranking uses average student entropy from a single forward-pass preprocessing step with a frozen student.
  • Selective LM Head and Chunked Entropy Computation: Selective LM-head and chunked entropy computation reduce logit-related memory by computing only needed logits or entropy scalars.Chunked computation avoids materializing the full [B, L, V] logits tensor, while selective projection uses Nselect positions.

5. Experiments

The experiments isolate selective KD design choices across position metrics, selection policies, baselines, and evaluation settings. They cover general-purpose pretraining-style distillation and task-specific reasoning distillation.

  • Experimental Protocol: The evaluation fixes all but one design axis at a time to isolate the effects of selective KD choices.This controlled protocol addresses the infeasibility of exhaustively evaluating the many configurations in the design space.
  • Position Selection: The position study compares importance metrics and policies including GLS, stochastic importance sampling, importance correction, and curriculum selection.The curriculum policy supervises a fixed k=20% of positions while shifting from low- to high-entropy tokens.
  • Baselines and Ablations: The baseline suite includes dense Full KD, AT-KD, RS-KD, random position and sample selection, TopSmp, an undistilled student, and the teacher upper bound.
  • Experimental Protocol: All methods share fixed KD hyperparameters selected once on validation data, with supervision budgets searched on validation splits.The setup uses temperature T = 1.0 and loss weighting λ = 1.0, yielding a KL-only objective.
  • Evaluation: General-purpose distillation trains models on 80 million FineWeb-Edu tokens and evaluates zero-shot performance across reasoning and long-range prediction benchmarks.The benchmark suite includes HellaSwag, PIQA, Arc-E, GSM8K, and LAMBADA.
  • Evaluation: Task-specific distillation applies KD directly to GSM8K and reports exact-match test accuracy under off-policy and on-policy settings.SE-KD3X is excluded because class-level sampling requires an offline teacher cache incompatible with dynamic student generation.

6. Results

Across position, sample, and multi-axis experiments, student-entropy-guided selection performs strongly under selective distillation, with benefits depending on the axis, budget, and task setting.

  • Position-importance metrics: 64.8 average accuracy and 6.9 perplexity make top-20% student-entropy selection stronger than Full KD at 64.4 accuracy and 7.3 perplexity.Top-20% KL, reverse-KL, and CE-ratio remain competitive, while teacher-entropy and teacher-CE rankings underperform.
  • Position-selection policies: 64.8 accuracy, 6.9 perplexity, and 21.4 instruction-following outperform Full KD at 64.4, 7.3, and 20.5 under top-20% student-entropy selection.SE-KD also outperforms random selection, GLS, curriculum scheduling, and AT-KD in accuracy and IFEval, although AT-KD calibrates best.
  • Distillation budgets: 20% is the best reported position budget for SE-KD and reverse-KL, while budgets near 1% can match or exceed Full KD and approximately 0.25% remains closer to the undistilled baseline.Both methods are robust across a wide range of budgets, with a shallow optimum at intermediate values.
  • Distillation budgets: Sample-selection accuracy changes little as the budget varies, while compute scales roughly linearly; the authors therefore use ℓ=20% in multi-axis experiments.Samples are ranked by average student entropy, and only the top-ℓ% are distilled.
  • Selection across axes: SE-KD3X combines position, class, and sample selection with competitive 64.4 accuracy, 20.7 IFEval, and 7.3 PPL while substantially reducing runtime, memory, and storage.Position selection is the dominant performance contributor, while the combined method has slightly worse calibration.
  • Task-specific distillation: On GSM8K, Full KD reaches 71.6 off-policy accuracy, whereas entropy-based top-20% position selection reaches 69.5; SE-KD plus TopSmp remains close at 70.9.In on-policy distillation, combining entropy-guided position selection with sample filtering yields the strongest results, but selective gains are less consistent off-policy.

7. Distillation Efficiency

Selective KD improves efficiency through complementary savings across sample, class, and position selection. SE-KD3X combines these axes to reduce storage, runtime, and peak memory while maintaining competitive performance.

  • Runtime and storage efficiency: Sample-level selection yields prominent runtime gains by reducing the number of sequences requiring teacher supervision.Reusing an offline cache of selected indices removes the upfront scoring step.
  • Runtime and storage efficiency: Class-level sampling enables orders-of-magnitude storage reductions for offline teacher supervision.For N=100B tokens, caching sampled teacher classes requires substantially fewer bytes than storing full teacher logits.
  • Runtime and storage efficiency: 99.96% storage reduction versus Full KD and 80% versus RS-KD reduce SE-KD3X storage to 3.84 TB.RS-KD reduces storage from 10,000 TB to 19.2 TB, while SE-KD3X further benefits from sample selection.
  • Runtime and storage efficiency: Position selection adds further speedups of up to ∼30% with a selective LM head and chunked entropy.Sample selection provides the main efficiency gains, while position selection contributes additional savings.
  • Memory efficiency: 18.3% lower sum of per-GPU peak memory allocations results from avoiding full logit materialization and restricting KD computation to selected positions.In the reported B=2, L=512 setting, memory falls from 33.18 GB to 27.10 GB.

8. Conclusion and Discussion

The paper finds that dense, uniform logit supervision is often unnecessary for autoregressive LLM distillation. Student-entropy-guided position selection and its multi-axis extension provide favorable performance–efficiency trade-offs, within a limited evaluation scope.

  • Concentrating supervision on a small subset of high-uncertainty positions consistently matches or outperforms Full KD for general-purpose distillation.
  • Student-entropy-guided Top-20% selection is the most reliable overall strategy, with curriculum learning, CE-ratio ranking, and teacher–student KL as promising alternatives.
  • SE-KD3X integrates position, class, and sample sparsification to provide favorable accuracy–efficiency trade-offs through sample filtering, offline caching, and reduced peak memory.
  • The study evaluates one widely used teacher–student pair under a fixed supervision budget, leaving validation across model families, scales, and longer contexts for future work.

Impact Statement

The paper reports no societal impacts specific to this work beyond general considerations associated with training and deploying language models.

  • The authors identify no societal impacts specific to this work beyond general considerations for training and deploying language models.

A. Proof: Positional Random Sampling Selection is an Unbiased Estimator of Weighted KD

Positional random sampling selects token positions according to non-negative importance weights and estimates a weighted KD objective in expectation. Importance correction can instead recover Full KD in expectation.

  • Design choice: The analysis deliberately matches weighted KD that emphasizes tokens according to their entropy rather than transforming the estimator to match Full KD.
  • Sampling formulation: Positional random sampling draws K token indices independently from a distribution determined by non-negative per-token importance weights.The sequence contains N indexed positions, each with per-token distillation loss L_t and importance weight w(t).
  • Sampling formulation: The sampling probability for token t is proportional to its contribution in the weighted KD objective.
  • Unbiasedness: The empirical loss estimator is unbiased for the weighted KD objective in expectation.The proof represents c_t as the number of times token t is sampled and takes expectation over the sampling process.
  • Importance correction: Importance-corrected positional random sampling can be reweighted by inverse sampling probability to estimate the full unweighted KD objective.

C. Additional Results

Auxiliary ablations motivate the paper’s default hyperparameters: temperature T=1.0 generally outperforms higher temperatures, while cross-entropy mixing adds little benefit and limits efficiency optimizations.

  • Temperature T=1.0 mostly outperforms higher temperatures in the Full KD ablations.These experiments motivate the temperature setting used throughout the paper.
  • The cross-entropy component provides negligible benefit in the Full KD baseline.
  • Including cross-entropy would prevent some selection-based efficiency optimizations, including restricting gradient-carrying logits to selected positions.

D. Positional Random Sampling Underperformance

At a fixed position budget, entropy-proportional positional random sampling underperforms deterministic Top-k in the general-purpose setting, likely because concentrated entropy reduces coverage and increases selection variance. The section also connects selective-position choices to measurable memory and speed savings, while noting unresolved sampling and caching trade-offs.

  • Positional Random Sampling Underperformance: Positional random sampling underperforms deterministic Top-k in both accuracy and calibration at the same position budget.The comparison uses a fixed budget of k=20%.
  • Positional Random Sampling Underperformance: Entropy-mass concentration can allocate too much budget to a few extreme-entropy positions, reducing coverage of other informative positions and increasing update variance.
  • Positional Random Sampling Underperformance: Future Pos RS-KD variants include temperature smoothing, entropy clipping, position exclusions, and hybrid deterministic coverage, but their systematic study remains open.
  • Offline Cache Tradeoff: Offline caching retains a trade-off: position-level caching saves more storage but removes the evolving student-entropy curriculum, whereas sample-level caching preserves curriculum effects at higher cache cost.
  • Selection Choices: Student entropy was selected as the position-importance metric because it achieved top validation accuracy and requires no teacher-side information.This also enables selective teacher-side logit computation at non-selected positions.
  • Selection Choices: Top 20% was selected as the main policy because its small validation disadvantage relative to GLS and Curriculum favored simplicity and fewer hyperparameters.The reported validation differences were 0.1–0.2 points.
  • Efficiency Effects: Chunked-streaming at k=100% reduces student peak memory from 15.88 GB to 14.15 GB and yields a 9% speedup.
  • Efficiency Effects: At k=20%, selective computation lowers peak memory and can eliminate transient logit spikes when a selective LM head computes logits only at selected positions.The selective LM head computes logits at approximately 20% of positions rather than materializing full sequence logits.
Loading 2602.01395v1…