Distillation Variants: Feature, Attention, Progressive

Knowledge distillation, introduced in the previous chapter, compresses a large teacher model into a smaller student by matching softened probability outputs. That approach transfers what the teacher believes about each class. But the teacher knows far more than its final predictions. Inside its intermediate layers are rich geometric structures: feature maps that capture syntax, attention patterns that signal which words matter, and gradient dynamics that shape learning itself. Distillation variants extract these deeper signals, giving students access to knowledge that the output distribution alone cannot convey.

This chapter covers four extensions that push distillation beyond soft labels. Feature distillation transfers intermediate representations, aligning the hidden-layer geometry of student and teacher. Attention transfer teaches the student to replicate the teacher's attention patterns, guiding where it looks rather than just what it concludes. Progressive distillation takes a different angle, treating the teacher-student relationship as a staged refinement process that iteratively compresses models while preserving quality. On-policy distillation addresses a subtle but important failure mode of standard distillation: the mismatch between the distribution of inputs the student is trained on and the distribution it encounters at inference time.

Understanding these variants matters because each solves a different gap in vanilla output-level distillation. You will see that the choice of which variant to use depends on your task, your architecture, and how much computational budget you have for training. In practice, production pipelines rarely use a single method in isolation. The most capable compressed models combine feature distillation, attention transfer, and on-policy generation into unified training objectives, extracting every available signal from the teacher.

Standard knowledge distillation operates only at the output layer, comparing the student's softmax probabilities to the teacher's. This ignores everything the teacher computed internally. Feature distillation says: the intermediate representations themselves are worth learning from.

The core idea is that a well-trained deep network builds up a hierarchy of increasingly abstract representations. A language model's lower layers capture token-level patterns like morphology; middle layers encode syntactic structure; upper layers represent semantic roles and discourse-level relationships. A student that replicates only the final output could, in principle, build completely different internal representations and still produce correct answers. But it is likely to generalize worse, because those different representations have not been shaped by the same structural pressures that made the teacher generalize.

Think about what it means for two networks to converge to the same output distribution through different internal paths. The student that finds its own internal geometry might overfit to superficial statistics of the training set. The student whose internal geometry resembles the teacher's has been guided toward the teacher's more principled inductive bias. The teacher is a learned function that has organized knowledge across its layers, rather than merely a label generator. Feature distillation says: let the student borrow that organizational structure, along with the final answers.

The earliest formalization of feature distillation appeared in FitNets (Romero et al., 2015). The approach adds an auxiliary objective that minimizes the distance between corresponding intermediate feature maps in the teacher and student.

Let $\mathbf{h}^T_l \in \mathbb{R}^{d_T}$ denote the teacher's activations at layer $l$, and $\mathbf{h}^S_l \in \mathbb{R}^{d_S}$ denote the student's activations at the corresponding layer. Because the student is typically smaller, $d_S < d_T$. A learned linear projection $W_r \in \mathbb{R}^{d_T \times d_S}$ maps the student into the teacher's space:

where:

• $\mathbf{h}^T_l$: teacher hidden activations at layer $l$, treated as fixed targets
• $\mathbf{h}^S_l$: student hidden activations at the corresponding layer
• $W_r$: a trainable projection matrix (called the "regressor") that maps student activations to teacher dimensionality
• $\| \cdot \|_2^2$: squared Euclidean distance, penalizing deviations in every dimension equally

The total training loss combines this feature loss with the standard soft-label distillation loss:

where:

• $\mathcal{L}_{\text{CE}}$: standard cross-entropy loss against ground-truth labels
• $\mathcal{L}_{\text{KD}}$: the soft-label Kullback-Leibler divergence loss from the previous chapter
• $\mathcal{L}_{\text{feat}}$: the feature alignment loss above
• $\alpha, \beta, \gamma$: scalar coefficients controlling the relative contribution of each objective

The projection $W_r$ is discarded after training. It exists only to bridge the dimensionality gap during distillation. This is a subtle but important point: you do not expand the student's architecture permanently. The regressor is a training scaffold, not a structural change. Once training is done, the student runs inference without it at its native, smaller dimensionality.

Why does this work? The squared distance in $\mathcal{L}_{\text{feat}}$ penalizes the student's projected hidden state for every dimension in which it differs from the teacher's. This creates a dense gradient signal. Whereas the output KD loss gives the student one signal per token position in its vocabulary, the feature loss gives it one signal per hidden dimension per token position per layer. For a 128-dimensional hidden state on a 16-token sequence, that is 2,048 gradient components per example compared to the vocabulary-size components from the output loss. The student receives far more directional guidance about how to organize its internal representations.

A practical challenge in feature distillation is deciding which teacher layers to match to which student layers. A teacher with 24 layers and a student with 6 layers cannot align layer-by-layer. Common strategies include:

• Matching the $k$-th student layer to the $\lfloor k \cdot N_T / N_S \rfloor$-th teacher layer, distributing correspondences evenly across the network depth
• Matching only the final few layers, where representations are most semantically rich and directly predictive of the output
• Skipping early layers, whose low-level features are easier for the student to learn independently from the training data

The choice matters because different layers encode different kinds of knowledge. Early layers in language models tend to encode local token-level information: subword statistics, character n-gram patterns, simple morphological cues. Middle layers encode syntactic structure: verb-argument relationships, constituency boundaries, named entity spans. Late layers encode semantic and pragmatic content: discourse roles, sentiment orientation, world-knowledge associations. If you align only late layers, you give the student strong semantic guidance but leave it to independently discover syntax. If you align early layers, you give it a structural foundation but may under-constrain the higher-level representations.

Activation magnitudes also differ between teacher and student. A normalized version of the loss mitigates this:

This cosine-based formulation ignores magnitude and focuses on directional alignment, removing sensitivity to the scale differences that naturally arise between networks of different sizes. A teacher with 128-dimensional hidden states may have activation norms substantially larger or smaller than a student with 64-dimensional hidden states, especially early in training when the student's parameters are still poorly calibrated. The normalized form decouples directional structure from magnitude, which is typically the more informative signal during early training phases.

For transformer models, feature distillation typically targets the hidden states output by each transformer block. Patient Knowledge Distillation (PKD, Sun et al., 2019) aligns the student's $[\mathrm{CLS}]$ token representation at every layer to its teacher counterpart, encouraging the student to develop similarly structured hidden states at each depth.

More recently, DistilBERT (Sanh et al., 2019) applied feature distillation by initializing the student from every other teacher layer and matching hidden states with a cosine loss. The initialization from teacher layers is itself a powerful form of knowledge transfer: rather than having the student start from random parameters and slowly learn to approximate the teacher, it begins the distillation training already producing reasonable representations. The remaining distillation training then fine-tunes these borrowed representations within the student's smaller capacity. TinyBERT (Jiao et al., 2020) took this further, aligning hidden states, embedding layers, and attention matrices simultaneously, producing a multi-objective distillation pipeline demonstrated for language models at the time.

The appeal of feature distillation is clear: it provides a much richer training signal than soft labels alone. The risk is over-constraining the student. If you force the student to mirror the teacher's internal geometry exactly, you prevent it from finding more efficient representations suited to its own smaller capacity. The projection matrix $W_r$ helps by providing a linear adapter rather than requiring exact dimensionality-matched alignment, but the tension between fidelity and flexibility remains.

A related concern is representational collapse. If the feature loss coefficient $\gamma$ is too large relative to the task loss $\alpha$ and KD loss $\beta$, the student may learn to match teacher hidden states accurately while losing the ability to produce good output predictions. The student becomes a teacher-representation-replicator rather than a task-solving model. Balancing the three objectives requires careful hyperparameter tuning. In practice, researchers often anneal $\gamma$ over training: high early to shape internal geometry, lower later to let the student optimize task performance.

Attention mechanisms, covered in depth in the self-attention chapters of this book, produce one of the most interpretable outputs in a transformer: a distribution over positions indicating which tokens the model attends to when processing each token. These attention patterns contain structural knowledge that goes beyond the final logit distribution.

Attention transfer teaches the student to replicate the teacher's attention distributions, along with its output probabilities. The mechanism is direct: if the teacher has learned that the verb "is" should attend strongly to its grammatical subject, that syntactic relationship is implicit in the attention weights. Transferring it gives the student a scaffold for building correct representations.

This is distinct from feature distillation in an important way. Feature distillation constrains the student's hidden representations, which are the outputs of entire transformer blocks after multi-head attention, residual connections, and feed-forward layers. Attention transfer constrains the attention weights themselves, which are computed inside the multi-head attention operation before the subsequent processing stages. Attention weights reflect a more focused kind of structural knowledge: specifically, the token-to-token relevance relationships that the model has learned to track.

For a given layer $l$ and attention head $h$, let $A^T_{l,h} \in \mathbb{R}^{n \times n}$ be the teacher's attention weight matrix over a sequence of length $n$, and $A^S_{l,h} \in \mathbb{R}^{n \times n}$ be the student's counterpart. The standard formulation minimizes the mean squared error:

where:

• $L$: number of layers with attention transfer (usually all layers or a subset)
• $H$: number of attention heads per layer
• $A^T_{l,h}, A^S_{l,h}$: teacher and student attention matrices at layer $l$, head $h$
• $\| \cdot \|_F^2$: squared Frobenius norm, summing squared differences over all $n^2$ position pairs

The Frobenius norm sums across all $n \times n$ entries of the attention matrix. This means the loss penalizes discrepancies at every (query, key) position pair, including positions outside the attention maximum. A student that broadly approximates the teacher's attention pattern but misaligns on a few critical structural relationships will receive gradient pressure to correct those misalignments at every training step.

When the teacher and student have different numbers of heads, you have two options. You can aggregate across heads first by averaging into a single per-layer attention matrix:

and then align $\bar{A}^T_l$ to $\bar{A}^S_l$. Alternatively, you can match a subset of student heads to specific teacher heads, which requires a head correspondence mapping. Averaging is simpler and more stable in practice: individual heads often specialize differently across runs, so one-to-one head correspondence can be brittle. The averaged matrix captures the collective structural intent of all heads, which is the more stable transfer target.

A practical note on normalization: attention weights are already normalized to sum to 1.0 within each row (they are output of softmax over query-key dot products). So unlike hidden state features, there is no need for an explicit normalization step in the attention transfer loss. The MSE directly compares probability distributions over keys for each query position.

Consider a 4-token sequence: "The dog barks." Suppose the teacher at layer 0 has learned the following averaged attention pattern (each row sums to 1.0):

The teacher has learned that "barks" attends strongly to "dog" (0.65), capturing the grammatical subject-verb relationship. The student, early in training, might produce a more uniform pattern:

The Frobenius norm between these two matrices captures the structural gap. The attention transfer loss provides a gradient that pushes the student toward concentrating attention on "dog" when processing "barks," effectively teaching it the subject-verb relationship that the teacher has already discovered. This is knowledge that might take the student many more training steps to independently acquire from the task loss alone.

Attention transfer provides the most benefit for tasks that rely on precise structural relationships: parsing, coreference resolution, question answering requiring multi-hop reasoning, and translation. For these tasks, the teacher has developed clear attention patterns that encode syntactic and semantic structure. In contrast, for tasks like sentence classification where the answer is largely determined by a few key tokens and the overall theme, attention transfer adds less over standard output-level distillation.

One notable finding from TinyBERT is that attention transfer is particularly useful in the lower layers, where the teacher's heads specialize for syntactic dependencies. Upper-layer attention patterns are more diffuse and task-specific, making them harder for the student to replicate exactly. This aligns with the general observation about layer specialization in transformers: lower layers handle structural parsing, upper layers handle semantic aggregation. The structural parsing patterns are more consistent and transferable; the semantic aggregation patterns depend heavily on the specific task and are harder to match precisely.

Attention transfer incurs substantial computational cost. Computing $\mathcal{L}_{\text{attn}}$ requires storing and comparing all attention matrices at all layers, in addition to the final output. For a sequence of length 512, each attention matrix is $512 \times 512 = 262{,}144$ entries. For 12 layers and 12 heads, that is roughly 37 million entries per example. Compared to the vocabulary-size soft-label comparison (which involves maybe 50,000 entries per position), the memory and compute cost of attention transfer is substantially higher. This reinforces the practice of applying attention transfer selectively to a subset of layers rather than all of them.

An advanced variant uses not raw per-layer attention maps but attention rollout, which propagates attention through the network to estimate the effective attention from input tokens to each layer's output. This captures how information flows across multiple layers rather than within a single layer, giving a transfer target that integrates information across layers.

The motivation for rollout is that raw per-layer attention weights can be misleading. A head in layer 5 might attend strongly to position 3, but position 3's representation in layer 5 is already a mixture of all input positions due to the attention layers below it. The effective influence of token 3 at layer 5 is determined by more than the direct attention weight to position 3 at that layer; it is the product of the attention paths that connect the original input token 3 through all intermediate layers. Attention rollout computes this propagated influence by taking the matrix product of attention matrices across layers, accounting for the residual connections that also pass information forward directly.

Transferring attention rollout rather than raw attention weights gives the student a more accurate picture of the teacher's true information routing, but at higher computational cost: you must compute the matrix products across layers for every example in every training batch.

Feature distillation and attention transfer still require you to have both a large teacher and a final small student. Progressive distillation takes a different architectural approach: rather than a single student learning from a single teacher, you build a chain of successively smaller models, each distilling from its predecessor.

The canonical form works as follows. Start with a large pretrained model $M_0$. Train $M_1$ to distill from $M_0$, producing a slightly smaller model. Train $M_2$ to distill from $M_1$. Continue until you reach the target size. Each step makes a small reduction in model capacity, keeping the gap between teacher and student manageable.

The idea echoes how skilled trades have always transmitted knowledge: a master does not teach an apprentice everything at once, but works through graduated challenges, making sure each step builds on the last. The same principle applies here. The teacher at each stage has already compressed some of the master's knowledge into a slightly smaller form, making it a more approachable guide than the original large model would have been for the final small student.

The key insight is that distillation degrades gracefully when the capacity gap between teacher and student is small, but degrades sharply when the gap is large. A 1.5 billion parameter model trying to distill a 70 billion parameter model in one step faces an extreme compression challenge: the small student must somehow compress 46 times more knowledge than it has capacity for. Much of the teacher's knowledge is simply unrepresentable.

Progressive distillation sidesteps this by keeping each step's compression ratio modest, typically halving the number of parameters or layers at each stage. The mathematics here connects to rate-distortion theory: the information loss at each distillation step is bounded by the capacity difference between teacher and student. Smaller steps mean smaller per-step losses, and these compose more favorably than a single large compression.

Consider an analogy from information theory. Suppose you must encode a 100-bit message using only 20 bits. You will lose at least 80 bits of information in a single step, regardless of how cleverly you design the encoding. But if you first encode it as 80 bits, then as 60, then 40, then 20, you make four smaller losses. Each encoding can optimize for the most important information to retain given the current representation, and the chain of compressions can move toward a 20-bit representation that preserves more of the original 100-bit message's semantic content than a single 100-to-20-bit compression would have.

The same logic applies to neural network compression. A 3-billion-parameter model distilled from a 7-billion-parameter model can accurately reproduce most of the teacher's output distribution because the capacity gap is manageable. That 3-billion-parameter model, trained well, now carries a compressed but faithful version of the teacher's knowledge. Starting from this intermediate model rather than from scratch, a 1-billion-parameter student has a much better guide.

A concrete protocol for progressive distillation of a transformer language model:

1. Start with teacher $M_0$ (e.g., 7B parameters, 32 layers)
2. Initialize student $M_1$ from the first 24 layers of $M_0$ (reducing to approximately 5B parameters)
3. Train $M_1$ with KD loss using $M_0$ as teacher until convergence
4. Initialize student $M_2$ from alternate layers of $M_1$ (reducing to approximately 2.5B parameters)
5. Train $M_2$ with KD loss using $M_1$ as teacher
6. Repeat until target size is reached

Layer initialization from the teacher is important. Starting the student from teacher weights rather than random initialization gives it a large head start and reduces the number of training steps needed. The student is not starting from scratch; it is refining an already-good set of representations. When you remove layers from the teacher to initialize the student, you are essentially asking the student to reproduce the teacher's behavior using a subset of its layers. The remaining layers still carry the teacher's learned weights, so they already encode reasonable representations. The training then adjusts those representations to compensate for the missing layers.

This layer-dropping initialization is motivated by an empirical observation about transformers: neighboring layers in a deep network often learn redundant representations. A well-trained 32-layer model may have considerable redundancy between, say, layers 10 and 11, because both are exposed to similar input distributions during training. Dropping one of them and fine-tuning the other introduces only a modest quality loss, much less than random initialization would.

An effective training schedule changes what is distilled at different stages. Early in each stage, the feature-level loss dominates, aligning the student's internal representations to those of the current teacher. Later in training, the soft-label KD loss becomes more important, sharpening the student's final predictions. This mirrors a sensible learning curriculum: first, ensure the student builds the right kind of representations internally; then, refine how it uses those representations to make accurate predictions.

Progressive distillation also naturally applies to other model dimensions. You can progressively reduce the number of attention heads, the width of feedforward layers, or the sequence length capacity, making a different architectural tradeoff at each stage. In practice, layer count reduction and width reduction are the most common, since layers correspond to representational depth and width corresponds to representational capacity per layer, and both have well-understood relationships to model quality.

Progressive distillation requires training more models for more total epochs than single-step distillation. The question of whether that cost is justified depends on the compression ratio. For modest compression, say a 2x reduction in parameters, single-step distillation is usually sufficient and the added cost of a multi-stage pipeline is not warranted. For aggressive compression, reducing a model by 10x or more, progressive distillation frequently produces substantially better final models and is the standard industrial practice.

The threshold is not fixed but depends on architecture, task, and data. A model on a simple classification task with abundant training data may tolerate large single-step compression with minimal quality loss because the task itself is forgiving. A model on a complex generation task with limited fine-tuning data may require progressive staging even for moderate compression ratios because quality degrades steeply with compression.

All three methods above share a common assumption: the student is trained on the same data distribution as the teacher. For discriminative tasks like classification, this is fine. For generative tasks, it creates a fundamental problem.

Consider training a small student language model to distill a large teacher. The standard approach samples a batch of text from a dataset, runs both teacher and student on it, and minimizes the KL divergence between their output distributions at each position. But the text in the dataset was written by humans. The student, when deployed, will generate its own text. That text will follow its own distribution, which is different from the human-written distribution in ways that depend on the student's current capabilities and biases. The student is being trained to handle human-generated inputs, not the inputs it will encounter.

This is the distribution shift problem in sequence generation distillation, and on-policy distillation is the solution. The name comes from reinforcement learning terminology: "on-policy" means training on data generated by the current policy (the model being trained), as opposed to "off-policy" training on a static dataset generated by some other distribution. The key insight is that the model that will be evaluated at inference time is the current student, so it makes sense to train it on the kind of inputs it will produce.

The mismatch between training distribution and inference distribution is called exposure bias. During training on a static dataset, the student always sees gold-standard input prefixes: "The quick brown fox jumps". During inference, it generates its own prefixes: perhaps "The quick brown fox leaps" at step 4. The word "leaps" was never in the training data for position 4, so the student has never learned to continue from it. Small errors compound: each generation step produces a slightly wrong token, which pushes the model further off-distribution, where it has no reliable signal.

The compounding nature of this error is worth dwelling on. After one generation step, the input is slightly off the training distribution. After five steps, it is further off. After twenty steps, the accumulated deviation may be large enough that the student's outputs are qualitatively different from what it would produce on clean inputs. This is visible in practice as repetition loops (where the model gets stuck generating the same phrase repeatedly), incoherence (where later sentences contradict earlier ones), and degradation in quality over long sequences.

Exposure bias affects standard language model training too, not just distillation. Scheduled Sampling (Bengio et al., 2015) is a classical remedy: mix gold tokens with the model's own predictions at training time. The fraction of model-generated tokens starts low and increases over training, progressively exposing the model to its own distribution. On-policy distillation extends this idea specifically for the distillation setting, making it the default rather than a scheduled addition.

In on-policy distillation, the student generates sequences using its own current policy, then uses those generated sequences as training inputs. Let $x_{1:t}$ be the student's own generated prefix. At position $t+1$, the objective minimizes:

where:

• $\pi_S$: the student's own policy (probability distribution over sequences)
• $x_{1:t} \sim \pi_S$: sequences generated by sampling from the student
• $p_T(\cdot \mid x_{1:t})$: teacher's next-token distribution given the student-generated prefix
• $p_S(\cdot \mid x_{1:t})$: student's next-token distribution given the same prefix
• $D_{\text{KL}}$: KL divergence measuring how far the student's distribution is from the teacher's

The key difference from standard distillation is the expectation over $\pi_S$: training inputs come from the student itself, not from a static dataset. The student learns to mimic the teacher on the distribution it encounters during inference.

This objective has a natural interpretation: after each generation step by the student, ask the teacher "given what the student just produced, what would you predict next?" and train the student to match that prediction. The teacher provides a corrective signal at every position in the student's own generated sequences, guiding the student toward teacher-quality continuations even when it starts from its own text.

An important subtlety in on-policy distillation is the direction of the KL divergence. Standard forward KL, $D_{\text{KL}}(p_T \| p_S)$, is mode-covering: it penalizes the student heavily for assigning low probability to anything the teacher assigns high probability to. When the teacher is a large, capable model with a rich, multimodal distribution, this coverage requirement causes the student to spread probability mass across many low-probability teacher tokens. The student tries to cover all the teacher's modes, which leads to averaged-out, overly cautious generations that hedge across many possibilities rather than committing to any one good output.

MiniLLM (Gu et al., 2023) proposes using reverse KL, $D_{\text{KL}}(p_S \| p_T)$, in the on-policy setting:

Reverse KL is mode-seeking: it penalizes the student for placing probability mass where the teacher does not. This encourages the student to be precise about the modes it covers rather than spreading evenly. For language generation, this produces sharper, more confident output at the cost of some diversity. The student learns: "focus your probability mass where the teacher is confident." Rather than trying to cover every possible completion the teacher might give, it covers the most likely ones well.

The intuition behind forward versus reverse KL can be grounded in a simple example. Suppose the teacher assigns 40% probability to "cat" and 40% to "dog" as next tokens, with the remaining 20% spread across other words. Forward KL forces the student to assign significant probability to both "cat" and "dog" and all the minor modes. Reverse KL allows the student to pick one of the two dominant modes and concentrate its probability there, accepting that it will miss the other major mode (and all minor modes) in exchange for being precise about the mode it does cover.

Whether this tradeoff is desirable depends on the application. For open-ended generation, diversity matters and forward KL is often preferred. For instruction-following, where there is typically one best response, reverse KL's mode-seeking behavior produces crisper, more decisive outputs that users typically prefer.

The gradient of the reverse KL under the student's policy can be computed using the policy gradient theorem, connecting on-policy distillation to reinforcement learning:

where $\theta$ are the student's parameters. This is the REINFORCE estimator applied to the distillation objective. In practice, variance reduction techniques like baseline subtraction are needed to make this stable. The log-ratio $\log \frac{p_S(x)}{p_T(x)}$ is a reward signal: if the student's probability exceeds the teacher's, the gradient pushes the student to reduce that probability. If the teacher's probability exceeds the student's, the gradient pushes the student to increase it.

Training with on-policy distillation requires careful engineering. You must:

1. Generate from the student: sample a batch of sequences from the student model using autoregressive decoding
2. Query the teacher: run the teacher forward on the same sequences to get its probability distributions at each position
3. Compute the loss: compute the KL divergence between teacher and student at each generated token position
4. Update the student: backpropagate through the student only (teacher is frozen)
5. Repeat: the student's policy evolves, so the on-policy distribution changes throughout training

This loop is more expensive than offline distillation because every training step requires a generation pass. A standard offline distillation step processes a batch in one forward pass through both models. An on-policy step requires an autoregressive generation loop (one forward pass per token in the sequence) plus the forward pass for the loss computation. For a 256-token sequence, on-policy training is roughly 256 times more expensive per example than offline training. Techniques like caching teacher outputs for recently generated sequences can reduce overhead, but the fundamental cost difference remains significant.

Several engineering strategies have emerged to manage this cost. One is to use short generation prefixes: generate 20-50 tokens from the student, then compute the loss on those tokens rather than generating full sequences of hundreds of tokens. Another is mixed training: alternate between on-policy batches (expensive, for distribution coverage) and offline batches (cheap, for data efficiency). A third is to use a smaller, faster version of the teacher as the on-policy signal generator, with the full large teacher used only for periodic evaluation.

The implementation in production LLM distillation pipelines often uses a variant called SeqKD (sequence-level knowledge distillation, Kim and Rush, 2016), where the teacher generates full sequences, those sequences are added to the training corpus, and the student is trained on them with standard maximum likelihood. This is computationally equivalent to offline training after the one-time generation cost, but ensures the student's training data includes teacher-generated text rather than only human-generated text, partially addressing the distribution shift problem at lower computational overhead than true on-policy generation.

Let's implement each distillation variant with a small language model to see how the loss functions, training dynamics, and generation quality differ.

We build a compact transformer setup that makes the patterns clear without requiring large-scale compute.

We build a minimal transformer with configurable depth and width, giving us control over teacher vs. student capacity.

The feature alignment loss requires a projection from the student's lower-dimensional hidden states to the teacher's dimensionality. Without this projection, we cannot directly compare hidden states of size 64 (student) to hidden states of size 128 (teacher).

The student's CE loss after feature distillation reflects how well it has learned to predict next tokens while simultaneously aligning its hidden-state geometry to the teacher's. Notice that we pass gradients through both the student parameters and the projector parameters: the student learns representations that, when projected, match the teacher; the projector learns the mapping that best bridges the two geometric spaces.

Now we implement attention transfer. The loss averages MSE across all student layers, mapping each student layer to a proportional teacher layer index.

On-policy distillation requires generating from the student during training. This is the most structurally different variant: unlike the previous two, which add extra loss terms computed on static batches, on-policy distillation changes the data itself.

Note that the on-policy student's loss is a KL divergence rather than a cross-entropy, which puts it on a different scale than the other variants. The raw loss values therefore should not be ranked against one another. In this short toy run, the on-policy KL decreases more modestly than the two cross-entropy curves, reflecting the harder objective of matching a teacher across the student's changing generated-prefix distribution.

One way to assess feature distillation is to measure how well the student's projected hidden states match the teacher's after training. A high cosine similarity indicates the two models have developed directionally similar internal representations.

A high cosine similarity confirms that feature distillation successfully steered the student to learn directionally similar representations to the teacher. The projection layer absorbs the dimensionality difference, so the comparison is made in the teacher's geometric space. What this measures is not whether the student has the same information content as the teacher (it cannot; it has half the dimensions), but whether the information it does contain is organized in the same directional structure.

Comparing teacher and student attention heatmaps gives a visual sense of how well structural transfer succeeded. After attention transfer training, the student's attention maps should show similar concentration patterns to the teacher's, even though the student has fewer heads and layers.

The progressive distillation code shows the key structural difference from other methods: we train two models sequentially, where the second uses the first as its teacher.

Each distillation variant introduces computational and design overhead beyond standard KD, and understanding these tradeoffs is essential before committing to one approach. None of these methods is universally best; the right choice depends on your task, your architecture, and your compute budget.

Feature distillation requires designing the projection layer $W_r$ and choosing which teacher layers to align with which student layers. This layer correspondence problem has no universal solution. Incorrect alignments can harm the student, forcing it to mimic teacher representations that are architecturally incompatible with its capacity. For transformer models, a useful heuristic is to match layers by relative depth, aligning the student's $i$-th layer to the teacher's $\lfloor i \cdot N_T / N_S \rfloor$-th layer. But this is domain-dependent: in vision transformers, early layers specialize for local texture, and matching them strictly to the corresponding teacher layers may not benefit semantic tasks. A more principled approach is to search over layer correspondences using a small validation set, selecting the alignment that maximizes student validation performance, but this adds another hyperparameter search loop to an already expensive training process.

There is also the risk of over-regularization from feature distillation. The student has less representational capacity than the teacher. Forcing it to reproduce the teacher's high-dimensional geometry can cause the student to allocate its limited capacity to mimicking teacher structure rather than solving the task. In practice, researchers report that the feature loss coefficient $\gamma$ needs careful tuning: too high and the student becomes a teacher-geometry-replicator that fails at the task; too low and the feature alignment provides no benefit over standard KD.

Attention transfer faces a related correspondence problem when teacher and student have different numbers of heads. Averaging across heads before computing the loss loses head specialization information: some teacher heads specialize for syntactic dependencies while others handle positional patterns. Matching individual student heads to teacher heads preserves this, but requires a head alignment procedure that may be brittle. The benefit of attention transfer also diminishes for tasks where attention patterns are diffuse or task-agnostic, making it primarily valuable for structured prediction problems. For generative language modeling where attention patterns are inherently noisy and variable across sequences, the signal from attention transfer may add less than the computational overhead it introduces.

Progressive distillation multiplies the training cost by the number of stages. A three-stage pipeline requires training three models, each for multiple epochs, compared to a single training run for direct distillation. The intermediate models add memory overhead. In practice, this cost is offset by the improved final student quality, but for resource-constrained settings, it may be preferable to use a single-step approach with a more carefully tuned loss. Another subtlety is model selection at each stage: if an intermediate model is poorly trained, errors propagate forward through the chain, potentially making the final student worse than a direct one-step distillation would have been.

On-policy distillation is the most computationally expensive variant because it requires autoregressive generation at every training step. For long sequences, this is prohibitively slow without significant engineering effort. The standard remedy is to generate shorter prefixes and limit the on-policy component to a fraction of training batches. Mixed training, alternating between on-policy and offline steps, often achieves most of the benefit at lower cost. The choice between forward and reverse KL also has significant practical implications: forward KL produces more diverse generations but can result in overly hedged outputs, while reverse KL produces more focused, confident text but may be overconfident on inputs outside the student's learned distribution.

A recurring practical limitation across all variants is that they assume the teacher is a reliable guide. If the teacher is itself imperfect, biased, or overfit to the training distribution, the student faithfully replicates those flaws. This is particularly relevant for attention transfer: if the teacher's attention patterns encode spurious correlations, teaching the student to replicate them transfers the flaw. It is also relevant for on-policy distillation in a subtle way: the teacher is queried on student-generated sequences, which may be very different from the sequences the teacher was trained on. If the teacher's distribution is poorly calibrated on out-of-distribution inputs, the on-policy KL divergence signal becomes unreliable precisely in the regime where the student most needs guidance. Auditing teacher behavior before distillation, and monitoring the quality of teacher outputs on student-generated prefixes throughout on-policy training, is sound engineering practice.

Given the distinct costs and benefits, a practical decision guide:

• Use feature distillation when your student architecture mirrors the teacher's (e.g., both are transformers with comparable layer structure) and you have sufficient compute to handle the projection overhead. It provides the richest training signal for structural tasks.
• Use attention transfer when your task requires precise relational understanding and the teacher has developed well-defined structural attention patterns. It is most valuable for encoder models on parsing, coreference, and structured prediction tasks.
• Use progressive distillation when your target compression ratio exceeds roughly 8-10x. Below that threshold, single-step distillation usually suffices. Above it, the quality gains from staged compression typically justify the training cost.
• Use on-policy distillation for any generative task where inference quality on long sequences matters. The distribution shift problem is severe enough for autoregressive generation that on-policy training, even a mixed on-policy/offline schedule, almost always improves over purely offline distillation.

In production, these methods are routinely combined. A common pipeline for compressing a large instruction-tuned LLM into a deployable small model uses progressive staging to reach the target size, feature distillation and attention transfer within each stage for richer signal, and on-policy training in the final stage to align the student's generation distribution with the teacher's.

This chapter extended knowledge distillation from output-level soft labels into four richer transfer mechanisms:

• Feature distillation aligns intermediate hidden representations between teacher and student using a projection layer, giving a dense training signal that guides the student's internal geometry toward the teacher's representational structure.
• Attention transfer teaches the student to replicate the teacher's attention weight distributions, transferring structural knowledge about which tokens matter and how information flows through the network across all sequence positions.
• Progressive distillation chains multiple small compression steps to avoid the sharp quality degradation that occurs when teacher and student differ by a large capacity factor in a single step, motivated by the rate-distortion observation that smaller per-step capacity gaps compose more favorably than one large gap.
• On-policy distillation trains the student on sequences it generates itself, eliminating the exposure bias that causes standard distillation to fail at inference time when the student generates its own off-distribution text. The choice between forward KL (mode-covering, diverse) and reverse KL (mode-seeking, focused) shapes the generation characteristics of the trained student.

Each variant addresses a specific limitation of vanilla KD. Feature distillation and attention transfer enrich the training signal by exploiting internal structure. Progressive distillation makes compression less sensitive to large teacher-student gaps. On-policy distillation corrects distribution shift in generation tasks. In practice, these methods are often combined: production distillation pipelines for large language models typically use feature distillation together with on-policy training, and progressive staging is standard when the target model is more than 10 times smaller than the teacher.

The on-policy methods described here also connect directly to the broader alignment literature. RLHF (Reinforcement Learning from Human Feedback) can be understood as a generalization of on-policy distillation where the "teacher" signal comes from human preferences rather than a larger language model. The policy gradient machinery is the same; only the reward source changes. Understanding on-policy distillation therefore gives you a foundation for understanding a large class of post-training alignment techniques.

The next chapter examines pruning as a complementary compression technique. While distillation trains a new smaller model from scratch using teacher supervision, pruning removes unnecessary components from an existing model, giving a different route to the same goal of efficient inference.

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