PaperlayerView exact revision on arXiv

Source-linked AI summary

KDFP: A first-principles approach to knowledge distillation in large language models

Ryan Swift, Konstantinos Psounis

arXiv:2610.10854v1cs.LG

TL;DR

General knowledge distillation for LLMs has received less attention than distilling post-training abilities, despite its relevance to efficient, private on-device systems. The paper introduces KDFP, a first-principles white-box method for general LLM distillation. It reports better performance across nine benchmarks and higher training efficiency than existing methods.

  • Problem

    Recent LLM distillation has focused more on post-training abilities than general knowledge distillation, which is relevant to efficient, privacy-preserving on-device systems.

    Source: §1 · p. 1

  • Method

    KDFP is a first-principles white-box methodology for general knowledge distillation in LLMs.

    Source: §1 · p. 1, §9 · p. 10

  • Results

    KDFP outperforms existing methods by 1.6%–4.9% across 9 benchmarks and increases training efficiency by up to 99.1%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    KDFP’s peak performance required only 100M training tokens and a small, unmodified teacher.

    Source: §9 · p. 10

    The evaluation covers continual pretraining, not full pretraining, which the authors could not test because of resource constraints.

    Source: §8.3 · p. 9

Abstract

from arXiv · show

Knowledge distillation is an established technique for improving the capabilities of small, efficient student models by training them with the representations of larger, more capable teacher models. Much of the recent work in the distillation of large language models (LLMs) has focused on distilling abilities learned during post-training, such as instruction following, chain-of-thought reasoning, and tool usage. This has left a large research gap in general knowledge distillation for LLMs, which is essential for developing efficient and private systems suitable for deployment on edge devices. We take a first-principles approach, evaluating previous lessons from prior works and conducting new explorations to develop a distillation methodology suitable for modern LLMs. We present KDFP, a novel methodology for white-box general knowledge distillation in LLMs. We demonstrate that KDFP outperforms existing methods by 1.6% $-$ 4.9% across 9 benchmarks while increasing training efficiency by up to 99.1% through ephemeral parameter reduction.

1 Introduction

General knowledge distillation for LLMs has received less attention than distilling post-training abilities, despite its relevance to capable, private on-device systems. The paper introduces KDFP and reports gains across benchmarks and several findings about activations, teacher alignment, data, and divergence mixing.

1 Introduction

  • General knowledge distillation has received comparatively less attention than distilling instruction-following, reasoning, and tool-use abilities.The paper argues that distillation can support capable, privacy-preserving on-device systems for resource-constrained platforms.Source: §1 · p. 1
  • The paper introduces KDFP, a first-principles methodology for general knowledge distillation in transformers.Its approach draws on prior work and mechanistic interpretability; Figure 1 depicts distillation of attention, MLP, and residual outputs, embeddings, and model outputs, using one learned projection matrix.Source: §1 · p. 1, Figure 1 · p. 2
  • In Section 7.1, KDFP outperforms existing methods by 1.6%–4.9% across 9 benchmarks, while Section 8.1 reports an efficiency increase of up to 99.1%.Source: §1 · p. 2
  • Previously ignored activations are essential for peak student performance.This finding is reported in Section 7.1.Source: §1 · p. 2
  • A 75% teacher and 25% student target-distribution mix outperforms existing divergence formulations.The paper also reports that teacher–student alignment may benefit students and that increasing training data may decrease performance.Source: §1 · p. 2

2 Preliminaries

These preliminaries introduce knowledge distillation, distinguish black-box from white-box approaches, and explain divergence measures and transformer components. Together, they establish the paper’s focus on using model internals and comparing student and teacher distributions.

2.1 Knowledge Distillation

  • Knowledge distillation compresses a teacher’s knowledge into a smaller student, first through synthetic-data training and later through output-distribution matching.The original formulation generated synthetic data for classification; Hinton et al. (2015) used forward KL divergence on teacher outputs.Source: §2.1 · p. 2

2.1.1 White-Box vs Black-Box

  • Black-box KD relies on model inputs and outputs, while white-box KD can also distill information-rich internal activations.Black-box synthetic-data methods include Vicuna (Zheng et al., 2023) and Alpaca (Taori et al., 2023); transformer methods have distilled residual-stream activations (Jiao et al., 2020; Liang et al., 2023b).Source: §2.1.1 · pp. 2–3

2.2 Divergence Functions

  • Divergence functions are nonnegative measures of statistical distance between a true distribution and its approximation, with zero indicating a perfect match.These properties motivate their use over cross-entropy in distillation.Source: §2.2 · p. 3

2.2.1 Kullback-Leibler Divergence

  • Forward KL can yield mean-seeking students, whereas reverse KL yields mode-seeking students that may lack diversity but can be specific and accurate.Forward KL’s mean-seeking behavior can produce diverse but potentially nonsensical responses; reverse KL’s behavior follows from minimizing where student probability is low despite teacher probability being high.Source: §2.2.1 · p. 3

2.2.2 Jensen-Shannon Divergence

  • JS divergence averages forward and reverse KL, producing more diversity than reverse KL and more faithful responses than forward KL.Source: §2.2.2 · p. 4

2.2.3 Generalizing Divergence

  • A parameterized divergence can weight forward and reverse KL to tune the balance between response diversity and faithfulness.Source: §2.2.3 · p. 4

2.3 Mechanistic Interpretability

  • Mechanistic interpretability offers component-level intuitions: the residual stream accumulates model knowledge, attention mixes information across positions, and MLPs process each position independently.The residual stream is described as a checkpoint or “best guess”; the MLP is difficult to interpret cleanly because of superposition.Source: §2.3 · p. 4

3 Related Work

Prior distillation work varies both its objective functions and the internal activations it targets. The cited methods illustrate trade-offs in objective design and constraints on activation distillation.

3.1 Distillation Objective Functions

  • Distillation objectives range from forward KL combined with language-modeling loss to cosine-alignment losses, alternative divergences, and sequence-level reverse KL.DistilBERT adds cosine embedding loss (Sanh et al., 2020); f-Distill compares several divergences (Wen et al., 2023); MiniLLM uses reverse KL for sequence-based, on-policy KD (Gu et al., 2024).Source: §3.1 · p. 4

3.2 Distilling Internal Activations

  • TinyBERT and TED distill internal activations, but TinyBERT’s QK attention-matrix distillation requires matching teacher and student head counts, and TED’s redundancy premise may not transfer to modern LLMs.TED succeeded on GPT-2 Small, but its scaling to modern 1B+ models and benchmarks remained uncertain; research on cross-layer superposition questions simple redundancy notions.Source: §3.2 · pp. 4–5

4 Distillation as a Parametrized System

The paper frames distillation as a system with four design parameters and focuses on token-level matching of representations. Its explorations address student initialization, projection, divergence, and objective design.

4 Distillation as a Parametrized System

  • The paper treats internal activations and output distributions as core knowledge units and frames general KD around four parameters: student initialization, activation projection, knowledge-gap measure, and objective function.It focuses on token-level rather than sequence-level KD because sequence generation is inefficient for full and continual pretraining in modern settings.Source: §4 · p. 5

4.1 Student Initialization

  • Shallow students performed poorly under the available training budgets, so the paper excludes their results and compares pretrained students against baseline and finetuned students.Pretrained students sacrifice full compatibility for a stronger initialization; their activations may require projection, and their attention-head activations are incompatible with teacher activations.Source: §4.1 · p. 5

4.2 Activation Projection

  • When embedding dimensions differ, a trainable upward-online projection maps student activations into the teacher’s embedding space; the paper selected this formulation after exploring alternatives.Source: §4.2 · pp. 5–6

4.3 Knowledge Gap Measure

  • Rather than fixing forward or reverse KL, the paper explores target-distribution mixing ratios to identify a suitable divergence measure.The exploration varies the weighting between forward and reverse KL.Source: §4.3 · p. 6

4.4 Objective Function

  • The objective combines teacher output distributions with internal activations, which act as checkpoints in model processing and are intended to provide useful signals beyond outputs alone.Mechanistic-interpretability findings motivate the exploration of new objective formulations.Source: §4.4 · p. 6

5 Methodology

KDFP combines a pretrained student, an upward-online projector, and distillation across multiple internal activations and outputs. The comparison methods differ in their projection and loss choices, with TED using a staged filter process.

5 Methodology

  • KDFP distills attention, MLP, and residual-stream outputs at each student block, alongside the output distribution and ground-truth cross-entropy loss.It maps N linearly spaced teacher blocks to a student with N blocks and weights objective components equally.Source: §5 · p. 6
  • KDFP uses a pretrained student initialization, an upward-online projector, and a target distribution mixed from 75% teacher and 25% student.Source: §5 · p. 6

5.1 Comparisons to Other Distillation Methods

  • Traditional KD, LWD, and TED use forward KL; traditional KD adds no internal-activation loss, while LWD adds residual-activation MSE.TED trains teacher and student filters with cross-entropy in a warmup stage, then distills their task-aware-filter outputs.Source: §5.1 · pp. 6–7
  • KDFP and LWD use a single upward-online projector, but LWD distills only the residual; this LWD variant uses one projector rather than one per layer.The authors note that this makes their variant less expressive than the LWD variant previously beaten by TED.Source: §5.1 · p. 6

6 Experiments

The experiments evaluate KDFP through continual pretraining with causal language modeling, using Pythia teacher and student models and FineWeb Edu data. Students are assessed zero-shot on nine benchmarks.

6 Experiments

  • Experiments use continual pretraining for causal language modeling, with Pythia 2.8B as teacher and Pythia 1.4B as student.The Pythia models are deduplicated variants, and the teacher is not finetuned on the dataset for the Section 7.1 setup.Source: §6 · p. 7
  • Training uses FineWeb Edu, a pretraining corpus, in a continual-pretraining setup intended to produce a more capable model at lower cost than full pretraining on the dataset.Source: §6 · p. 7
  • Teachers and students train on 100M tokens at sequence length 2048; TED filters receive a 25M-token warmup.The filter warmup preserves the 1:4 filter-to-distillation data ratio.Source: §6 · p. 7
  • Students are evaluated zero-shot on nine benchmarks, using accuracy, normalized accuracy, F1, or exact match according to the task.The benchmarks include BoolQ, MMLU, WinoGrande, ARC, PIQA, SciQ, DROP, and TriviaQA; ARC combines Easy and Challenge scores.Source: §6 · p. 7
  • The main-results table covers 100M training tokens and a 2.8B-parameter teacher without finetuning.Source: Table 2 · p. 8

7 Results

In the main setup, KDFP achieves the strongest reported average across nine benchmarks, while the analyses find no consistent benefit from teacher finetuning, a larger teacher, or more training data.

7.1 Main Results

  • KDFP reaches a peak average score of 0.4358, nearly 5% above TED and 2.9% above the Baseline across nine benchmarks.The gain over TED required an additional 0.03% of the pretraining token budget; TED degraded performance below the Baseline student.Source: §7.1 · p. 8

7.2 Teacher Finetuning

  • With a finetuned teacher, KDFP remains the top performer, but the peak average score drops by approximately 1%.LWD gains notably but remains 1.3% below the peak; the proposed distribution-drift explanation is a hypothesis.Source: §7.2 · p. 8

7.3 Teacher Scaling

  • Scaling the teacher from 2.8B to 6.9B parameters largely does not help students, and the peak average score falls 1.3%.KDFP again performs best, while TED shows a very small improvement; the 6.9B teacher's 4096 embedding dimension is double the student's 2048.Source: §7.3 · p. 8

7.4 Data Scaling

  • In Table 3c, KDFP with a 5x increase in data (500M tokens) falls 2.1% behind the peak average score; Section 7.1 KDFP outperforms these variants with an 80% reduction in training cost.Among the compared variants, LWD performs best and FT improves.Source: §7.4 · p. 8

8 Discussion

KDFP reduces temporary trainable parameters while distilling additional activations already computed in the forward pass. The discussion also presents a hypothesis for TED's results and notes evaluation and implementation boundaries.

8.1 Training Efficiency

  • KDFP trains 5.2M temporary projector parameters versus nearly 600M temporary filter parameters for TED, reducing training requirements by over half a billion parameters.Its additional activations are already computed in the forward pass; the added compute is storing and calculating losses on them.Source: §8.1 · pp. 8–9
  • KDFP uses information available from small, unmodified teachers and does not require long, expensive training schedules; the authors leave the limits of this efficiency for future work.Source: §8.1 · p. 9

8.2 Explaining TED Performance

  • The authors hypothesize that TED's filters may disrupt information processing, including through their removal at inference, and point to TED's below-Baseline results as an indicator.They suggest future work investigate this explanation.Source: §8.2 · p. 9

8.3 Limitations

  • KDFP is evaluated only in continual pretraining because the authors lacked resources to test full pretraining, even for models below 2B parameters.Source: §8.3 · p. 9
  • TransformerLens's incomplete multi-GPU support limited experiments to models trainable on a single GPU.The authors say activation extraction with HuggingFace Transformers is feasible, but not standardized.Source: §8.3 · p. 9
  • Pythia 1.4B student performance is not directly comparable to Nano-1 because Nano-1 is 4-bit quantized.Source: §8.3 · p. 9

9 Conclusion

KDFP is presented as a general method for distilling LLMs, with better performance and efficiency than comparable methods. The conclusion highlights a 75% teacher/25% student target mix and peak performance with 100M tokens and a small, unmodified teacher.

9 Conclusion

  • KDFP is a new LLM distillation method that demonstrates superior performance and efficiency relative to comparable methods.The conclusion also says a 75% teacher and 25% student target distribution outperforms existing methods.Source: §9 · p. 10
  • KDFP reaches peak performance with 100M training tokens and a small, unmodified teacher; the conclusion also identifies less data and teacher–student alignment as potentially beneficial.The paper presents KDFP as an example of how understanding LLM science can lead to practical improvements.Source: §9 · p. 10

A Disclosure of LLM Use

  • The authors used LLMs to help generate table and TikZ figure code, developer API citations, and Python code for the divergence comparison plot.Source: A Disclosure of LLM Use · p. 22

B Extended Analysis Results

The extended analysis tables cover teacher finetuning, teacher size, and training-data scaling. The accompanying sections describe attention-head incompatibility, JS-loss computation, and projection formulations.

B Extended Analysis Results

  • Tables 4–6 cover teacher finetuning, larger Pythia teachers, and data scaling across benchmark columns; italics mark gains over non-finetuned, 2.8B-teacher, and 100M-token baselines, respectively.Source: Table 4 · p. 22, Table 5 · p. 22, Table 6 · p. 22

C Incompatibility of Pretrained Student Attention Heads

  • Specialized attention heads make pretrained-student head matching difficult, so the method omits their activations and leaves crosscoder mappings to future work.The text says crosscoders may be capable of creating such a mapping.Source: C Incompatibility of Pretrained… · p. 22

D JS Divergence Implementation

  • The JS implementation softens logits with temperature τ, computes log probabilities via log_softmax, excludes masked tokens, and averages JS over valid tokens with τ^2 scaling.Source: D JS Divergence Implementation · pp. 22–23

E Projection Matrices

  • Traditional trainable linear projections map student activations into teacher space; KDFP compares alternatives varying construction, training, and projection direction.The paper hypothesizes that an online learned projector may absorb knowledge intended to transfer to the student.Source: E Projection Matrices · p. 23

E.1 Online

  • An online projector learns its weights during distillation, but may capture the alignment knowledge instead of propagating it into the student.Source: §E.1 · p. 23

E.2 Offline

The offline-projector design maps student activations into teacher space using a separately trained linear autoencoder, while staged experiments narrowed the projection, divergence, and objective choices.

E.2 Offline

  • Offline projectors are trained separately as fully linear autoencoders, then inserted and frozen during student training.Their latent representations align to teacher activations, while the output is reconstructed; the alignment-loss weight β is tunable.Source: §E.2 · pp. 23, 25, Figure 4 · p. 24

E.3 Randomized

  • Upward randomized projections were evaluated despite lacking theoretical motivation, because distance-preservation guarantees apply only to downward projection.The randomized matrix uses Achlioptas's binary coin formulation and is frozen during training.Source: §E.3 · p. 25

F.1 Stage 1: Basic Exploration

  • Shallow students performed far worse than pretrained students, so they were excluded from later exploration stages.The paper notes that shallow students require full pretraining, making them infeasible in most modern settings.Source: §F.1 · p. 26, Figure 5 · p. 26

F.2 Stage 2: Search for Offline Projector Alignment

  • Every offline-projector alignment setting scored 0.4151 on 0-shot HellaSwag, so alignment was fixed at 0 and upward projection selected.The sweep covered five alignment values in both projection directions; the upward direction preserves full-fidelity teacher activations.Source: §F.2 · p. 26

F.3 Stage 3: Search for Divergence Parameter

  • β = 0.25 was selected because it ranked second-highest in both 0-shot and 5-shot evaluations, whose top and bottom parameters were inconsistent.β = 0 is forward KL, β = 0.5 is traditional JS divergence, and β = 1 is reverse KL.Source: §F.3 · p. 27, Table 8 · p. 27

F.4 Stage 4: Search for Objective Function

  • Objective-function exploration constrained each objective to a convex combination of included components with evenly distributed weights.The search used a two-part approach to reduce the number of component combinations examined.Source: §F.4 · p. 27

F.4.1 Stage 4A: All or Nothing

  • The Balanced objective with an online projector had a distinct advantage over other configurations, while projector variants performed identically for Only Outputs.The authors therefore continued Only Outputs experiments with the randomized projector, which required no training.Source: §F.4.1 · p. 28, Table 9 · p. 28

F.4.2 Stage 4B: Marginal Modifications

  • The Balanced objective with an online projector performed best, and mlp_out activations appeared to be the most important included component.The subtractive comparisons removed hook groups from the original Balanced objective.Source: §F.4.2 · p. 28, Table 10 · p. 28

F.5 Stage 5: Final Validation

Final validation favored β = 0.25 and upward projection for the pretrained student with the online projector and Balanced objective. Hyperparameters were tuned in Stage 1 and then held consistent across later exploration results.

F.5 Stage 5: Final Validation

  • β = 0.25 remained the best divergence parameter with the pretrained student, online projector, and Balanced objective fixed.The final validation revisited the divergence sweep.Source: §F.5 · p. 28, Table 11 · p. 29
  • Upward projection performed better than downward projection, despite training losses an order of magnitude higher.The authors attribute the difference in losses to the curse of dimensionality, since teacher embeddings are twice the student size.Source: §F.5 · p. 28, Table 12 · p. 29

G Hyperparameters

Loading 2610.10854v1…