Knowledge Distillation: Teacher-Student Training for LLMs

Modern language models keep growing. GPT-3 has 175 billion parameters. GPT-4 is estimated to be considerably larger. These models perform remarkably well, but deploying them at scale is expensive: each inference call touches hundreds of gigabytes of weights, requires high-end GPUs, and adds latency that users notice. You cannot run a 175B parameter model on a smartphone. You cannot serve it cheaply for millions of requests per day without a datacenter budget that most teams simply do not have.

The instinctive response to this problem is compression: make the model smaller. But naive size reduction, like simply removing layers or pruning weights, destroys the model's learned representations. The accuracy collapse is severe. Something important is being lost, and it is not immediately obvious what. The problem is not just that the model has fewer weights. The problem is that those weights were arranged in a specific way that encoded a rich, interlocking structure of knowledge. Removing pieces of that structure without care collapses the whole.

Knowledge distillation answers a different question: rather than compressing a model, can we train a small model to behave like a large one? Instead of shrinking GPT-3, can we train a much smaller model that has absorbed what GPT-3 knows? The large model becomes a teacher, and the small model becomes a student. The student learns from labeled data and from the teacher's outputs, specifically from the soft probability distributions the teacher assigns to each possible output.

The key insight, due to Geoffrey Hinton, Oriol Vinyals, and Jeff Dean in their 2015 paper "Distilling the Knowledge in a Neural Network," is that a model's output distribution carries far more information than a hard label. When a classifier says a handwritten "7" is 73% likely to be a 7, 18% likely to be a 1, and 5% likely to be a 4, those secondary probabilities encode structural knowledge about which concepts are similar. A student trained on these soft targets learns richer representations than one trained only on the correct label "7." The teacher's uncertainty is itself instructive, because it reveals the geometry of the problem space.

In the context of language models, the teacher's output distribution over a vocabulary of 50,000 tokens at each position communicates which words are semantically interchangeable, which completions are plausible, and what syntactic patterns are expected. A student trained to match this distribution inherits a compressed form of the teacher's understanding of language, without needing to replicate all 175 billion of the teacher's parameters.

This chapter develops the distillation framework from first principles. We start with the mathematical objective, work through temperature scaling in detail, examine how to choose and evaluate teacher models, and build a working implementation. We then scale these ideas to the LLM setting, examining DistilBERT, TinyBERT, and instruction-following distillation as concrete case studies. The chapter closes with an honest assessment of where distillation succeeds and where it falls short.

To understand why distillation works, we need to examine what information is encoded in a model's output probabilities and how the student is trained to capture it. The mechanism is elegant, but it takes some care to see why it provides more signal than ordinary supervised training.

A standard language model training objective uses cross-entropy loss against the true one-hot label: the correct next token gets probability 1.0, all others get 0.0. This is called training with hard labels. Hard labels are efficient and mathematically clean, but they discard an enormous amount of information.

Consider what happens when a language model processes the sentence "The cat sat on the ___." The model assigns high probability to "mat" but also substantial probability to "floor," "rug," "sofa," and "chair." These secondary probabilities are not noise. They reflect semantic knowledge: all of these completions are plausible, and the differences between their probabilities reflect subtle contextual cues about which is most natural in this exact phrasing. Hard-label training treats this entire distribution as equivalent to a one-hot vector pointing at "mat." All that relational information evaporates.

The same problem appears at finer scales. If the teacher assigns 60% probability to "cat" and 35% probability to "feline," that relationship between near-synonyms is discarded by hard-label training. The student only sees "cat: 1, everything else: 0." But the teacher's soft assignment reveals that "cat" and "feline" occupy the same semantic neighborhood in the model's representation space. A student that learns to match this distribution will develop a similar geometric structure in its own embeddings.

Distillation trains the student on soft labels: the teacher's full probability distribution over all possible outputs. Let $p^T_i$ be the teacher's probability for token $i$, and $p^S_i$ be the student's predicted probability for token $i$. The distillation loss is the cross-entropy between these two distributions:

where:

• $V$: the vocabulary size
• $p^T_i$: the teacher's probability assigned to token $i$ (a soft label)
• $p^S_i$: the student's predicted probability for token $i$
• The sum runs over all $V$ tokens in the vocabulary

This is Kullback-Leibler divergence between the teacher and student distributions up to a constant, so minimizing the cross-entropy is equivalent to minimizing the KL divergence:

where $H(p^T)$ is the teacher's entropy, which is constant with respect to the student's parameters. Minimizing the distillation loss therefore pushes the student's distribution toward the teacher's, which is exactly what we want.

Why does this help? The teacher has been trained on vastly more computation than the student could afford. Its output distribution encodes the result of fitting billions of parameters to massive datasets. When the student is trained to match that distribution, it gets a rich supervisory signal that goes far beyond binary correct-or-wrong feedback. Every training example now communicates a full probability distribution rather than a single spike, and the student can learn the relative relationships between outputs from that distribution.

In practice, most distillation setups combine the distillation loss with the standard supervised loss against the ground-truth labels. This ensures the student learns from labeled data directly while also absorbing the teacher's knowledge:

where:

• $\mathcal{L}_{\text{distill}}$: the distillation loss (cross-entropy against teacher soft labels)
• $\mathcal{L}_{\text{CE}}$: the standard cross-entropy loss against ground-truth labels
• $\alpha \in [0, 1]$: a mixing coefficient that controls the balance between the two objectives

When $\alpha = 1$, the student learns purely from the teacher's soft labels. When $\alpha = 0$, it learns only from ground truth. Typical values in practice lie in the range $[0.5, 0.9]$, favoring the distillation signal while retaining some grounding in labeled data. The appropriate value depends on how well the teacher's distributions match the true underlying distribution and how much labeled data is available.

This combination is useful for a subtle reason. The teacher's soft labels encode structural knowledge about relationships between classes, but they are still the teacher's beliefs, not ground truth. The teacher makes errors. If the teacher systematically assigns too much probability to certain tokens due to biases in its training data, the student will inherit those biases. Mixing in the ground-truth loss acts as a corrective force, anchoring the student to the actual labels even while it absorbs the teacher's relational knowledge.

The most practically important concept in distillation is temperature scaling, and it is worth understanding in detail because it fundamentally changes what information the student receives. Without temperature scaling, distillation often provides little improvement over hard-label training, because the soft labels are too similar to hard labels to provide additional signal.

Language model outputs tend to be sharply peaked. Given a sentence like "The sky is ___", the model assigns very high probability to "blue" and near-zero probability to essentially everything else. When these near-zero probabilities are used as soft labels, they provide almost no additional information compared to hard labels: the distribution is essentially one-hot.

The teacher's knowledge about relationships between less likely tokens is compressed into these tiny probability values. Words that are semantically related to "blue" (like "clear," "gray," or "overcast") receive probabilities that are only slightly higher than completely unrelated words. At floating-point precision, these differences exist, but they are numerically tiny. A student trying to match these near-zero values receives a training signal that is essentially the same as trying to match zeros. The relational structure is there in principle, but it is buried beneath the noise floor of floating-point arithmetic.

This is the fundamental problem that temperature scaling solves. The information is present in the teacher's logits, but it is compressed in a way that makes it invisible to the training objective. We need to uncompress it.

Temperature scaling addresses this by flattening the distribution before using it as a soft label. The teacher's logits $z^T_i$ are divided by a temperature parameter $\tau$ before the softmax is applied:

where:

• $z^T_i$: the teacher's logit (pre-softmax score) for token $i$
• $\tau$: the temperature parameter ($\tau > 0$)
• $p^T_i(\tau)$: the temperature-scaled probability for token $i$

When $\tau = 1$, this is the standard softmax. When $\tau > 1$, the distribution becomes flatter: smaller probabilities grow relative to the largest ones. When $\tau < 1$, the distribution becomes sharper, concentrating even more mass on the peak.

To see why temperature matters concretely, consider logits $[10, 2, 1, 0]$. At $\tau = 1$, the softmax gives approximately $[0.9997, 0.0003, 0.0001, 0.00004]$: essentially a one-hot distribution. At $\tau = 4$, the same logits give roughly $[0.76, 0.13, 0.08, 0.05]$: a much softer distribution where the relative magnitudes of all four logits are visible in the probability values. The smaller logits now carry meaningful probability mass and can guide the student's learning.

The mechanics of why this works come from the exponential function. When logits are divided by a large $\tau$, the differences between them shrink. The exponential function maps small differences to values close to 1, so the ratio between the largest and smallest probabilities decreases. What was a thousand-fold difference at $\tau = 1$ becomes a moderate difference at $\tau = 4$, moving probabilities that were numerically invisible into ranges where the gradient signal is strong.

During distillation, the same temperature $\tau$ is applied to both the teacher's distribution (to produce soft labels) and the student's distribution (to compute the loss). The student is trained to match the teacher's soft output. After training is complete, the temperature is reset to $\tau = 1$ for inference. The student's weights have been shaped by exposure to the temperature-scaled distributions, but it produces predictions at the standard temperature.

Hinton et al. introduced the term "dark knowledge" to describe the information encoded in the small probabilities that standard training ignores. A teacher that assigns probability 0.0001 to "feline" when predicting the word "cat" is revealing a structural relationship between these two words. This relationship is invisible under hard-label training but becomes visible when the distribution is temperature-scaled.

The key insight is that the ratio of probabilities at high temperature reveals structure that absolute probabilities hide. If the teacher assigns logit 10 to "cat," logit 2 to "feline," and logit $-5$ to "airplane," the ratio $p(\text{feline}) / p(\text{airplane})$ encodes the semantic proximity of cat and feline relative to airplane. At low temperature, both probabilities round toward zero and the ratio becomes numerically unreliable. At high temperature, both probabilities are substantial and the ratio is preserved, giving the student a meaningful gradient signal that "feline" and "cat" are semantically related while "airplane" is not.

This is the distinction between distillation and a minor regularization variant. The teacher's logits contain a rich relational structure built up over training. Temperature scaling is the tool that makes this structure accessible to the student.

The right temperature depends on how peaked the teacher's distribution is and how much dark knowledge you want to expose. For tasks with very confident outputs (like clearly labeled image classification), moderate temperatures of $\tau = 2$ to $\tau = 4$ are usually sufficient to make the secondary probabilities visible. For tasks with ambiguous targets (like next-token prediction in language models), higher temperatures up to $\tau = 10$ can reveal richer structure, but there is a point of diminishing returns.

At very high temperatures, the distribution approaches a uniform distribution, and the distinctions between tokens blur. The teacher's guidance becomes noise rather than signal. At very low temperatures, the distribution collapses toward a one-hot vector and the soft labels carry no more information than hard labels. The sweet spot is in between: high enough that secondary probabilities are substantial, low enough that the teacher's ranking of outputs is still preserved.

A practical approach is to treat the temperature as a hyperparameter and tune it on a held-out validation set. In practice, values in the range $[2, 8]$ work well for most classification and language modeling tasks. For well-trained, high-confidence teachers, lean toward lower temperatures; for less confident or more uncertain teachers, lean toward higher temperatures.

When using temperature scaling during distillation, the gradients of the distillation loss are scaled by $1/\tau^2$ compared to the gradients of the standard supervised loss. This happens because the softmax function with temperature $\tau$ produces logit gradients that are $1/\tau$ times smaller than those of the standard softmax, and this factor appears twice in the computation: once for the teacher's distribution and once for the student's. The net effect is that increasing $\tau$ reduces the magnitude of the distillation gradient by a factor of $\tau^2$.

To keep the distillation loss contribution stable relative to the supervised loss when varying $\tau$, Hinton et al. recommend multiplying the distillation loss term by $\tau^2$:

where $\mathcal{L}_{\text{distill}}(\tau)$ denotes the distillation loss computed with temperature $\tau$. The $\tau^2$ factor compensates for the gradient shrinkage caused by the temperature scaling, making sure that the magnitude of the distillation gradient remains roughly constant as $\tau$ varies. This makes it easier to tune $\alpha$ independently of $\tau$: you can increase the temperature to reveal more dark knowledge without accidentally reducing the effective learning rate on the distillation objective.

This $\tau^2$ correction is often omitted in implementations without dramatic consequences, especially when $\alpha$ is tuned empirically, but it simplifies hyperparameter tuning and is good practice when sweeping over temperature values.

The success of distillation depends critically on choosing the right teacher model. A bad teacher teaches bad habits. A teacher that is too different from what the student can learn from provides targets the student cannot approximate. Getting teacher selection right is often more important than fine-tuning the temperature or the mixing coefficient.

Several properties characterize a good distillation teacher:

• High accuracy on the target task: The teacher should be substantially better than a student trained without distillation. If the teacher and student have similar accuracy, distillation provides little benefit because the soft labels are not much more informative than what the student could generate on its own.
• Well-calibrated probabilities: The teacher's soft labels should encode meaningful uncertainty. An overconfident teacher that assigns 99.99% probability to all its predictions produces soft labels that are barely softer than hard labels, even with temperature scaling. Calibration can be measured by comparing the teacher's confidence to its actual accuracy on a held-out set.
• Same task and domain: The teacher and student should operate on the same input distribution. A teacher trained on formal news text distills poorly to a student meant for conversational text because the probability distributions over tokens reflect very different stylistic patterns.
• Compatible output space: Teacher and student must share the same vocabulary and output format for token-level distillation to work directly. Distilling across vocabularies requires additional alignment steps.

Beyond these requirements, a good teacher has also "learned the right things" in the sense that its internal representations are organized in a way that reflects linguistic or semantic structure, not artifacts of training noise. This is hard to measure directly, but models trained on more data and with more compute tend to have more structured internal representations.

A single model's soft labels reflect its particular training trajectory and may encode idiosyncratic biases. One teacher might assign systematically high probability to certain tokens due to artifacts in its pretraining data. Another might have learned slightly different probability assignments for ambiguous cases based on the random initialization it started from. These idiosyncrasies add noise to the distillation signal.

Using an ensemble of models as the teacher, averaging their output distributions, produces smoother and more reliable soft labels:

where:

• $M$: the number of teacher models in the ensemble
• $p^{T_m}_i$: the probability assigned to token $i$ by the $m$-th teacher

Ensemble distillation is especially effective when the individual teachers make different types of errors. Their averaged distribution is more informative because regions of uncertainty in one model are compensated by confidence in another. The student learns from a "collective wisdom" that no single model could provide. This is one of the original motivations in Hinton et al.'s paper: they noted that the ensemble provides a more stable and less biased training target than any individual model.

In practice, maintaining and querying an ensemble of large models is expensive, but even a small ensemble of two or three models substantially improves the quality of the soft labels. When the teachers are smaller models that fit easily in GPU memory, ensemble distillation is worth the additional compute.

A subtle but important issue is the capacity gap: when the student is much smaller than the teacher, it may not be able to approximate the teacher's distribution even in principle. A student with 10 layers trying to match a 96-layer teacher faces a learning problem that may be fundamentally unsolvable within its parameter budget.

The student can only represent distributions that are consistent with its own architecture. A very small model has limited expressivity. Its softmax output is constrained to distributions that its particular arrangement of weights can produce. If the teacher's distribution requires a more complex functional form than the student can represent, the student cannot match it regardless of how long it trains.

Empirically, very large capacity gaps lead to poor distillation. The student can learn that certain tokens are more likely than others, but it cannot reproduce the fine-grained probability differences that make the teacher's soft labels informative. Worse, trying to match an unachievable target can hurt the student's training, leading to worse performance than training without distillation. The student wastes optimization effort pursuing targets it cannot reach, rather than learning the best possible solution within its own capacity.

The practical solution is to use an intermediate "assistant" or "mediator" model that bridges the capacity gap. The original large teacher first distills into a medium-sized model, which then distills into the target small model. This progressive distillation produces students that outperform those trained directly from the large teacher when the gap is extreme. The intermediate student learns to represent a simplified version of the teacher's distribution, and this simplified version is much easier for the final small student to approximate.

The capacity gap problem also explains why some well-publicized distillation attempts produce disappointing results. When researchers report that distillation "doesn't work" for their use case, the capacity gap is often the culprit. If the student is too small relative to the teacher, and no progressive distillation is used, the student is set up to fail. The right solution is to either increase the student's size, reduce the teacher's size, or introduce an intermediate model.

Let's make these concepts concrete with a worked example. We'll simulate the distillation process on a sentiment classification task: binary classification of movie reviews into positive and negative. This is not an LLM example, but it clearly illustrates the mechanics before we scale up to the language model setting. The core insight, that soft labels encode more information than hard labels, is identical in both cases.

Suppose the teacher produces these logits for three reviews:

The first two reviews are clear-cut. The teacher is confident in both directions. Its soft labels at temperature $\tau = 2$ will be similar to hard labels: most of the probability mass will sit on the correct class, with a small residual on the wrong class. The third review is ambiguous. "It was fine, I guess" has a slightly positive teacher logit of 0.2 versus 0.1 for negative, but the difference is tiny. The teacher is essentially saying "I cannot tell."

Now consider what happens at training time. Hard-label training penalizes the student for assigning any probability to the negative class on review 3. The student learns that "fine, I guess" is unambiguously positive. This is technically correct based on the label, but it is overfitting to a noisy label. The true sentiment of that phrase really is ambiguous. A student trained on the teacher's soft labels at $\tau = 2$ learns that assigning nearly equal probabilities to both classes is acceptable for that phrase. It will generalize better to similar hedged expressions.

For the two clearly labeled reviews, hard and soft losses are similar: the teacher is confident, and the soft label closely resembles the hard label. For the ambiguous third review, the hard loss penalizes the student's uncertainty, while the soft loss accepts it. This is exactly the desired behavior: distillation avoids forcing overconfidence on uncertain examples, and the student learns to be appropriately uncertain where the teacher is appropriately uncertain.

The teacher's soft label for review 3 is approximately $[0.47, 0.53]$: nearly equal probabilities for both classes. This is the teacher's honest assessment of that example. A student trained on this soft label learns that equal uncertainty is acceptable for hedged language. A student trained on the hard label $[0, 1]$ is pushed toward false confidence.

We'll now implement the distillation training loop from scratch to make the mechanics fully transparent. We use a simple two-class classification problem where the teacher is a larger neural network and the student is a smaller one. The goal is to see the training dynamics and accuracy differences clearly, without the complexity of a full language model.

The dataset has 20 features, 10 of which are informative and 5 of which are correlated noise. The class separation of 0.8 makes this a moderately difficult problem: a large model can learn the informative features and achieve high accuracy, while a small model without guidance will struggle to separate signal from noise.

We implement a simple multilayer perceptron without a deep learning framework to keep the code transparent. The teacher has three hidden layers, giving it enough capacity to learn complex decision boundaries.

The MLP class stores all intermediate activations during the forward pass. This is necessary for the gradient computation in the backward pass. The distillation_loss function implements the core distillation objective: it temperature-scales both the teacher and student distributions, then computes the cross-entropy of the student's distribution against the teacher's.

The train_distillation function implements the combined objective. Notice the $\tau^2$ scaling factor on the distillation gradient, which compensates for the reduced gradient magnitude that temperature scaling introduces. The alpha parameter controls how much weight the distillation signal receives relative to the supervised signal.

The key hyperparameters for the distillation training loop are:

• temperature ($\tau$): Controls how much to soften the teacher's distributions. Values in $[2, 8]$ are typical. Higher values reveal more dark knowledge but may make the targets too noisy. Start at $\tau = 4$ and tune from there.
• alpha ($\alpha$): The weight on the distillation loss versus the supervised loss. Values in $[0.5, 0.9]$ favor distillation. When labeled data is plentiful, lower alpha is acceptable; when it is scarce, higher alpha extracts more value from the teacher.
• teacher architecture: The teacher should be substantially larger than the student. A 2-4x parameter ratio is a reasonable starting point; very large ratios may require progressive distillation.

One practical note: when tuning these hyperparameters, tune them jointly rather than independently. The optimal temperature depends on how calibrated the teacher is, and the optimal alpha depends on the temperature (because temperature affects the information content of the distillation signal, and alpha controls how much weight that signal receives).

Scaling distillation to large language models introduces both new opportunities and new challenges. The core objective remains the same, but the scale, the nature of the outputs, and the relationship between teacher and student all change significantly.

Distilling a 70B-parameter model into a 7B model involves training the student while repeatedly querying the teacher for its output distribution over a vocabulary of tens of thousands of tokens. The teacher must be kept in memory alongside the student during training, which roughly doubles the GPU memory requirement. For very large teachers (100B+ parameters), this becomes practically infeasible unless the teacher is served on separate hardware and queried via API, with logits streamed back to the student's training process.

An alternative is offline distillation: pre-compute the teacher's soft labels on the entire training corpus and store them as training data. The student then trains on these pre-computed distributions without needing the teacher loaded during training. This trades memory for storage. Storing soft labels over a vocabulary of 50,000 tokens for billions of training tokens requires enormous disk space.

In practice, the teacher's distribution is often truncated to store only the top-$k$ probabilities (e.g., $k = 64$), discarding near-zero mass. This significantly reduces storage requirements while retaining the most informative part of the distribution. The bottom $49,936$ probabilities in a 50,000-token vocabulary are typically negligible, and storing them wastes space without providing meaningful signal.

There are two levels at which distillation can operate for language models.

Token-level distillation applies the standard distillation objective at every token position in the sequence. At position $t$, the student is trained to match the teacher's distribution over the next token given the prefix. This is the most direct application of the Hinton et al. framework and is used in models like DistilBERT and TinyBERT. It provides the richest training signal because the student receives a full distribution at every step, but it requires access to the teacher's logits, which may not be available for closed-source models.

Sequence-level distillation treats the problem differently. Rather than matching the teacher's token-by-token distribution, the student is trained on text generated by the teacher. The teacher generates sequences by sampling from its distribution (possibly with temperature), and the student uses those generated sequences as training data for standard language modeling. This approach, studied by Kim and Rush (2016), avoids the need to store or stream the teacher's logits and works well when high-quality teacher-generated data is available.

The choice between these approaches involves real tradeoffs. Token-level distillation provides a richer training signal but requires accessing the teacher's internal distributions. Sequence-level distillation can be performed from any model that produces text, making it applicable even when distilling from closed-source models like GPT-4. However, sequence-level distillation loses the detailed probability information that temperature-scaled soft labels provide. The student learns what the teacher says, but not how uncertain the teacher was when it said it.

For practical LLM compression, token-level distillation typically produces better results when the teacher's logits are accessible. Sequence-level distillation is the only option when they are not, and in those cases it still provides substantial benefit over training without teacher guidance.

DistilBERT, introduced by Sanh et al. (2019) at Hugging Face, demonstrates token-level distillation at scale. The goal was to compress BERT-base (110M parameters, 12 layers) into a model with 40% fewer parameters (66M parameters, 6 layers) while retaining 97% of BERT's performance on the GLUE benchmark.

The distillation approach combined three components:

where:

• $\mathcal{L}_{\text{distill}}$: the standard distillation loss matching the teacher's masked language model predictions
• $\mathcal{L}_{\text{CE}}$: supervised cross-entropy loss against ground-truth masked tokens
• $\mathcal{L}_{\text{cos}}$: cosine embedding loss between teacher and student hidden state vectors, encouraging alignment of internal representations

The third term, $\mathcal{L}_{\text{cos}}$, represents a form of feature distillation: rather than only matching the teacher's output probabilities, the student is also trained to produce similar hidden representations layer by layer. We will explore this further in the next chapter on distillation variants. For now, note that it provides an additional training signal that helps narrow the representation gap between teacher and student beyond what output-level supervision alone can achieve.

DistilBERT's training also used weight initialization from the teacher: rather than initializing the 6-layer student randomly, the authors initialized its layers by taking every other layer from the 12-layer BERT teacher. This dramatically speeds up convergence because the student starts with weights already in a good region of parameter space. The distillation training then refines these weights to account for the fact that the student has fewer layers between them.

The reported results were substantial. DistilBERT achieved 97% of BERT's performance on GLUE, ran 60% faster at inference, and reduced memory consumption by 40%. It became one of the most widely used compressed language models and demonstrated that careful distillation can retain nearly all of a model's capability at a substantial cost reduction.

DistilBERT's success established several practical lessons for LLM distillation:

• Weight initialization from the teacher substantially accelerates convergence and should be used whenever architectures are compatible.
• The distillation objective is more effective than standard training alone even when the same data is used, meaning the soft labels add information beyond what the raw text provides.
• A 40% parameter reduction is achievable with minimal performance loss; larger compressions require more sophisticated techniques.
• The cosine representation loss is a useful complement to the output distillation loss, helping align internal representations not just final outputs.

TinyBERT, introduced by Jiao et al. (2020), pushed the compression ratio further by distilling both intermediate representations and the attention matrices, in addition to the final output. The student is trained to match the teacher's attention patterns at corresponding layers in addition to matching the output distribution.

This multi-layer distillation approach requires the student to have the same number of attention heads as the teacher (or a compatible mapping between them), but it transfers much richer information. The teacher's attention patterns encode how it distributes focus across the input sequence when processing each token, and matching these patterns constrains the student to develop similar internal processing strategies.

TinyBERT achieved competitive performance with only 14.5M parameters (about 13% of BERT-base's 110M) on GLUE tasks. The tradeoff is complexity: the distillation process requires careful layer mapping between teacher and student, and the multi-objective loss is harder to tune. But the compression ratio is dramatically better than DistilBERT, showing that when you extract richer signal from the teacher, you can get away with a much smaller student.

A different flavor of distillation has become prominent with instruction-following LLMs. Models like Alpaca, Vicuna, and Orca train a smaller model to mimic the instruction-following behavior of a larger model like GPT-4. The "teacher" in these cases is GPT-4, but its internal logits are not accessible. The distillation is purely sequence-level: GPT-4 generates responses to prompts, and the student is fine-tuned to produce similar responses on the same prompts.

This approach raises an important distinction between two uses of distillation:

• Capability distillation: transferring the teacher's factual knowledge, reasoning ability, and language model quality to the student
• Behavioral distillation: teaching the student to follow instructions, adhere to formats, and exhibit particular conversational styles

Both are valuable, but they require different teacher signals. Capability distillation benefits from token-level soft labels with temperature scaling, because the goal is to transfer the precise probabilistic structure of the teacher's language model. Behavioral distillation can work with generated sequences because the target behavior is the output format and conversational style, not the precise token probability distribution.

Alpaca and Vicuna demonstrated that even relatively modest fine-tuning on high-quality teacher-generated responses could substantially improve a smaller model's instruction-following capabilities. A LLaMA-7B model fine-tuned on 52,000 GPT-4 completions behaves far more helpfully and coherently than the base model. This is sequence-level distillation working as intended: the student learns the distribution of helpful responses by imitating a larger model's outputs.

Orca (Microsoft, 2023) took this further by using "explanation traces" from GPT-4 rather than final responses alone. The teacher was prompted to explain its reasoning step by step before giving an answer, and the student was trained on both the reasoning trace and the final answer. This richer supervision signal helped the student learn what to say and how to reason, resulting in substantially better performance on reasoning benchmarks than models trained on final answers alone.

A significant practical challenge in LLM distillation is distribution shift: the student generates different text than the teacher would, and the student must learn on text from the training corpus or generated by the teacher, rather than on its own outputs. This mismatch creates a situation where the student encounters teacher logits in contexts it would never reach when generating on its own.

To understand why this matters, consider what happens at inference time. The student generates token by token, and at each step its generation depends on all previous tokens it has produced. During training, the student has been trained on contexts from the training corpus or from the teacher's completions. If the student's own completions drift away from these contexts (which they do, because the student is different from both the corpus and the teacher), the student encounters contexts that are "off the training distribution." Its behavior in these off-distribution contexts is unpredictable.

On-policy distillation addresses this by having the student generate sequences during training and computing the teacher's logit distribution at each position of the student's own outputs. The student is trained to improve its distribution on the contexts it generates, rather than on fixed training data contexts. This is more expensive, requiring forward passes through both models on student-generated text, but it produces better alignment between the student's behavior and the teacher's guidance, especially for generation-heavy tasks.

On-policy distillation is related to techniques like RLHF (reinforcement learning from human feedback) where the model is also trained on its own outputs. The key difference is the reward signal: RLHF uses a scalar reward from a reward model, while on-policy distillation uses the teacher's full distribution as a dense, token-level reward. The dense signal is much more informative than a scalar, because it tells the student whether the output was good or bad and what a better output would have looked like at each token position.

When storing teacher distributions for offline distillation, the standard approach is to truncate to the top-$k$ logits. This raises an important question: how large should $k$ be?

If $k$ is too small, the stored distribution loses dark knowledge. With $k = 10$, only the ten most probable tokens are represented, and the relational structure among less likely tokens is discarded. With $k = 64$, a much richer slice of the distribution is preserved. With $k = 1000$, the storage cost grows but the distribution is nearly complete for most tokens.

In practice, $k = 64$ is a common choice that balances storage cost against information quality. For a 50,000-token vocabulary, the top 64 tokens typically account for more than 99% of the probability mass for most positions in natural language sequences. The remaining 49,936 tokens split a negligible fraction of probability between them. Their relative ordering is information the student could learn, but it is so faint that it rarely makes a measurable difference.

The choice of $k$ should also be matched to the temperature. At high temperature, more probability mass is spread across non-top tokens, making a larger $k$ more important. At low temperature, the distribution is sharp and a small $k$ captures almost everything.

Distillation is powerful, but it is not a universal solution to the model compression problem. Understanding its limitations is essential for applying it correctly and for knowing when to reach for a different technique.

The capacity gap problem is the most fundamental limitation. When the teacher is much larger than the student, the student cannot represent the teacher's distribution within its parameter budget. In these regimes, distillation can hurt: the student wastes capacity trying to approximate distributions it cannot represent, rather than learning the best possible solution within its own capacity. Empirical evidence suggests that the distillation benefit diminishes sharply when the parameter ratio exceeds roughly 10:1. This does not mean distillation fails entirely at large ratios, but it does mean the gains are smaller and less predictable, and progressive distillation is often necessary to realize the full potential.

Distillation also assumes the teacher is reliably better than the student. If the teacher makes confident errors, the student learns those errors as well, and the soft labels propagate misinformation. This is particularly problematic in domains where the teacher has limited training data, where it extrapolates far beyond its training distribution, or where it has developed systematic biases. In these cases, the distillation signal is noisier than the ground-truth labels, and a lower value of $\alpha$ is appropriate. There is also no mechanism in standard distillation for the student to recognize when the teacher is wrong: the student simply follows the teacher's distribution, errors and all.

The computational cost of distillation is often underestimated. Training the student requires the teacher's logits or soft labels at every training step, which means either keeping the teacher loaded in memory alongside the student (doubling the memory requirement), pre-computing and storing the teacher's distributions over the entire training corpus (requiring enormous storage), or making online API calls to a hosted teacher (adding latency and cost). For very large training datasets, none of these approaches is trivially cheap. The decision to use distillation should include an honest accounting of these costs.

The data distribution assumption is another limitation that practitioners frequently overlook. Distillation works best when teacher and student operate on the same data distribution. If the deployment distribution differs significantly from the training distribution, both teacher and student may fail, and the student inherits the teacher's failure modes rather than learning to handle distribution shift. Out-of-distribution generalization is not automatically transferred through distillation. It must be built into the training data or addressed through other techniques like data augmentation or distribution-aware training objectives.

Despite these limitations, distillation remains one of the most effective techniques for creating compact, deployable models. DistilBERT retains 97% of BERT's accuracy at 40% fewer parameters and 60% faster inference. TinyBERT achieves competitive performance with roughly 7x fewer parameters. Smaller instruction-following models like Mistral-7B and Phi-2 demonstrate that thoughtful distillation and data curation can produce models that punch far above their weight class. The Phi-2 model from Microsoft, with only 2.7B parameters, achieved performance on reasoning benchmarks comparable to models with five to ten times as many parameters, largely because of careful distillation-informed training data curation.

The practical takeaway is this: distillation works best when the capacity gap is moderate (2-10x), the teacher is well-calibrated, the data distribution is consistent, and the training pipeline can accommodate the teacher's memory or compute requirements. When these conditions are met, distillation reliably produces small models that outperform any alternative training strategy at the same parameter count.

Knowledge distillation transfers a large teacher model's learned knowledge to a smaller student model through soft label supervision. The core mechanism is simple: train the student to match the teacher's output probability distribution rather than just the correct class label. This soft target carries structural information about which outputs are similar and which are different, giving a richer training signal than hard labels alone.

Temperature scaling is the key technique for amplifying this information. By dividing the teacher's logits by a temperature $\tau > 1$ before applying softmax, the near-zero probabilities for non-predicted outputs become large enough to compare, revealing the teacher's implicit knowledge about relationships between outputs. The same temperature is applied to the student during training and removed at inference. Higher temperatures expose more dark knowledge but can introduce noise if taken too far; the range $[2, 8]$ works well for most tasks.

The combined distillation objective mixes the distillation loss (cross-entropy against the teacher's soft labels) with the standard supervised loss. The mixing coefficient $\alpha$ and the temperature $\tau$ are the two primary hyperparameters. The $\tau^2$ scaling factor compensates for the gradient magnitude reduction that temperature scaling introduces, keeping the distillation gradient stable as temperature varies.

Teacher selection matters significantly. The teacher should substantially outperform the student-capacity baseline, be well-calibrated, and operate on the same data distribution. Very large capacity gaps benefit from progressive distillation through intermediate models. Ensemble teachers produce smoother soft labels by averaging the outputs of multiple models, reducing the idiosyncratic biases of any single training run.

For large language models, distillation operates in two regimes. Token-level distillation matches the teacher's per-token probability distributions, as in DistilBERT and TinyBERT, giving the richest signal but requiring access to the teacher's logits. Sequence-level distillation trains on teacher-generated text, as in Alpaca and Vicuna, and works even with closed-source teachers but loses the detailed probability information. On-policy distillation addresses the distribution shift problem by computing teacher guidance on the student's own generated sequences, at higher cost but with better alignment for generation tasks.

The next chapter on distillation variants explores feature distillation, attention transfer, and other techniques that extend beyond output-level supervision to transfer richer internal structure from teacher to student.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about knowledge distillation.