Distillation Variants › Feature Distillation

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:

$$\mathcal{L}_{\text{feat}} = \| \mathbf{h}^T_l - W_r \mathbf{h}^S_l \|_2^2$$

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:

$$\mathcal{L}_{\text{total}} = \alpha \cdot \mathcal{L}_{\text{CE}} + \beta \cdot \mathcal{L}_{\text{KD}} + \gamma \cdot \mathcal{L}_{\text{feat}}$$

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:

$$\mathcal{L}_{\text{feat}} = \left\| \frac{\mathbf{h}^T_l}{\|\mathbf{h}^T_l\|_2} - \frac{W_r \mathbf{h}^S_l}{\|W_r \mathbf{h}^S_l\|_2} \right\|_2^2$$

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.