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

AI summary

Overview

Research area: Machine learning, specifically white-box knowledge distillation (KD) of large language models (LLMs).

Technical level: Advanced. The paper assumes familiarity with transformer internals (residual streams, attention, MLP blocks), KL and Jensen-Shannon divergence, and the language of mechanistic interpretability.

Scope: The paper proposes KDFP ("Knowledge Distillation from First Principles"), a token-level white-box distillation methodology for decoder-only LLMs, and compares it against traditional KD, layer-wise distillation (LWD), and TED on 9 zero-shot benchmarks using Pythia teacher and student models.

What This Paper Is About

Most recent LLM distillation research targets abilities learned during post-training, such as instruction following, chain-of-thought reasoning, and tool usage, leaving a gap in general knowledge distillation. The authors argue general distillation matters for building efficient, private, deployable systems, and they set out to design a distillation method for modern LLMs by revisiting lessons from prior work (TED, TinyBERT) alongside findings from mechanistic interpretability. They identify four design parameters of a general KD method — student initialization, activation projection, knowledge gap measure, and objective function — and explore each empirically to arrive at KDFP.

Key Contributions

1. KDFP methodology. A white-box general distillation recipe combining pretrained student initialization, an "upward-online" single projection matrix, a parametrized divergence with β = 0.25 (a 75% teacher / 25% student target mixing ratio), and an objective that distills attention, MLP, and residual stream outputs at every student block plus embeddings, output distributions, and cross-entropy loss, with all six components weighted equally (α = 1/6).
2. Demonstration that previously ignored activations matter. The paper reports that distilling attention and MLP outputs — not just residuals, as in LWD and TED — is essential for peak student performance.
3. Evidence on teacher–student alignment. It reports that finetuning the teacher and scaling the teacher from 2.8B to 6.9B both failed to help students, suggesting more closely aligned representations are preferable.
4. Evidence that more data can hurt. Increasing training data 5x (to 500M tokens) lowered peak average score.
5. A proposed explanation for TED's failure. The authors hypothesize TED's task-aware filters disrupt information processing, and that removing filters at inference introduces further disruption, since LLMs appear to store knowledge across layers.

Main Findings

• Peak average score. KDFP reached an average of 0.4358 across the 9 benchmarks with a 2.8B teacher and 100M tokens, versus Baseline 0.4236, LWD 0.4291, KD 0.4232, FT 0.4227, and TED 0.4154. The abstract states KDFP outperforms existing methods by 1.6% – 4.9% across 9 benchmarks; the main text describes the gain versus TED as nearly 5% and versus Baseline as 2.9%.
• TED degrades below baseline. TED scored 0.4154, below the Baseline pretrained student at 0.4236, which the authors present as evidence that TED does not scale to modern models and evaluations.
• Efficiency. KDFP trains only 5.2M temporary parameters (the projector) versus nearly 600M for TED's filters with Pythia 2.8B and 1.4B, which the abstract reports as up to a 99.1% increase in training efficiency through ephemeral parameter reduction.
• Finetuned teacher (Table 3a). KDFP led again at 0.4316 (LWD 0.4303, KD 0.4232, TED 0.4140), but the peak average dropped by approximately 1% relative to the main results.
• Larger teacher (Table 3b). With Pythia 6.9B as teacher, KDFP scored 0.4303 (LWD 0.4284, KD 0.4232, TED 0.4165), a 1.3% drop from the peak. Pythia 6.9B's embedding dimension (4096) is double the student's (2048), while 2.8B's is only 25% larger.
• 5x data (Table 3c). KDFP scored 0.4264 (LWD 0.4266, FT 0.4263, KD 0.4207, TED 0.4154), a 2.1% drop from peak; the main-results KDFP configuration still outperformed these with an 80% reduction in training cost.
• Projection. A single upward-online projector was found best; the authors report they found no alternative projection formulation in their literature review. Projection direction made no difference for offline and randomized projectors and only a very small difference for online.
• Shallow students failed. Shallow (reduced-layer) student initializations performed very poorly under the given training budgets, so no shallow results appear in the main results section.
• Attention heads were excluded. Attention head activations were not distilled, because head specialization makes meaningful teacher-to-student head mapping difficult for pretrained students.

Methodology in Plain English

The researchers frame distillation as a parametrized system with four knobs — how the student is initialized, how student and teacher activations are put into a comparable space, how the gap between them is measured, and which signals go into the loss — and then search over those knobs.

They pick pretrained models rather than shallow copies of the teacher, since pretraining a shallow student from scratch is infeasible at modern scale. Because a pretrained student's activations live in a different embedding space than the teacher's, they learn a single linear projection matrix during training that maps student activations up into the teacher's space ("upward-online"), instead of the per-layer task-aware filters TED uses. Rather than choosing strictly forward KL (mean-seeking) or reverse KL (mode-seeking), they use a generalized Jensen-Shannon divergence with a β parameter and land on β = 0.25, meaning the target distribution is 75% teacher and 25% student. The objective combines loss terms on the attention output, MLP output, and residual stream at every student block, plus the embedding output, the output distribution, and the ground-truth cross-entropy loss, each weighted equally.

They map N linearly spaced teacher blocks onto the N student blocks. Experiments use Pythia 2.8B deduped as teacher and Pythia 1.4B deduped as student, trained on 100M tokens of FineWeb Edu with a 2048-token sequence length (Pythia's full context) in a continual pretraining setting with causal language modeling. TED's filter warmup stage was trained on 25M tokens to match the 1:4 filter-to-distillation ratio from Liang et al. (2023b). Evaluation uses EleutherAI's LM Evaluation Harness, all zero-shot: accuracy on BoolQ, MMLU, and WinoGrande; normalized accuracy on ARC-Easy, ARC-Challenge, PIQA, and SciQ; F1 on DROP; and exact match on TriviaQA. ARC is reported as the sum of ARC-Challenge and ARC-Easy.

An exploration phase preceded the final method, using HellaSwag accuracy with seed 3, Pythia 2.8B teacher, Pythia 1.4B student, 10M tokens to train offline projectors, and 100M tokens to train students, across 5 stages: basic exploration, offline projector alignment search, divergence parameter search, objective function search, and final validation. HellaSwag results are deliberately excluded from the main body to avoid bias.

Why This Matters

Impact on research. The paper reframes general LLM distillation as an explicitly parametrized design space and reports that several intuitions from earlier distillation work — that redundancy filtering helps, that more data helps, that a stronger or better-finetuned teacher helps — did not hold in their modern decoder-only setting. It also connects distillation design to mechanistic interpretability findings about cross-layer knowledge storage, offering a testable explanation for why TED's filters may disrupt processing.

Real-world applications:

• Compressing LLMs into small, capable models for smartphones and single-board computers, where the paper notes distilled models small enough to run on a workstation have surpassed models with 15x as many parameters.
• Privacy-preserving on-device inference, since distillation enables local deployment without sending data to a hosted API.
• Continual pretraining: cheaply improving an existing model using a subset of a new, high-quality corpus rather than fully pretraining a new model.
• Edge assistants and on-device natural-language features, in the mold of the 1.8B-parameter, 4-bit-quantized Gemini Nano-1 the authors cite as commercially deployed on edge devices (noting their Pythia 1.4B students are not directly comparable).

Industry relevance. The reported efficiency profile matters for practitioners: KDFP's only extra cost is storing and computing losses on activations already produced by the forward pass, and it avoids TED's long multi-stage filter training schedule. Reduced ephemeral parameter counts mean less memory and compute during distillation itself, not just at inference.

Future Directions

1. Full pretraining evaluation. The authors could only test KDFP in the continual pretraining setting because pretraining a model of even under 2B parameters was infeasible with their resources, and they explicitly leave full pretraining to future work.
2. Multi-GPU and broader model support. The implementation relies on TransformerLens, which standardizes activation extraction but does not fully support multi-GPU setups, limiting experiments to single-GPU-trainable models. Extending extraction via HuggingFace transformers is possible but not standardized.
3. Distilling attention heads. Crosscoders are suggested as a possible tool for building meaningful mappings between specialized teacher and student attention heads, which KDFP currently does not attempt.
4. Testing the limits of efficiency. The authors state that KDFP "appears to thrive with efficiency" and leave it to future work to probe how far that goes, and to further investigate their hypothesis about TED's filter-induced disruption.

Target Audience

Researchers and engineers working on LLM compression, small-model training, or on-device deployment who need to understand white-box distillation losses and where existing recipes break down. It is also relevant to mechanistic interpretability researchers interested in how internal-representation findings translate into training objectives, and to practitioners who need a distillation method that works without finetuning the teacher, without massive data budgets, or without long multi-stage filter training. Readers wanting a first introduction to distillation will find the preliminaries useful, but the results sections assume comfort with transformer activation hooks and divergence measures.

Authors’ abstract

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.