Distillation Objectives

The core challenge in knowledge distillation (KD) is defining precisely what knowledge to transfer from the teacher to the student and how to measure the success of this transfer. The mechanism for this is the distillation objective or loss function. It quantifies the discrepancy between the teacher's behavior and the student's behavior, guiding the student's training process. While the original KD formulation focused on matching output distributions, several sophisticated objectives have been developed to capture richer information embedded within the teacher model.

Soft Targets: Matching Output Distributions

The most fundamental KD objective, proposed by Hinton et al. (2015), involves training the student to mimic the teacher's output probability distribution over classes or tokens. Simply matching the final predictions (hard labels) isn't sufficient, as the teacher's distribution often contains subtle information about class relationships or token likelihoods – knowledge that is lost when converting to a single, hard prediction.

To extract this richer information, a temperature scaling parameter, $T$, is introduced into the softmax function applied to the logits (the raw, unnormalized outputs before the final activation).

$$\sigma(z_i, T) = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$

Here, $z_i$ is the logit for class or token $i$. A higher temperature ($T > 1$) softens the probability distribution, pushing probabilities closer together and revealing more about the relative similarities the teacher assigns to different outputs. A temperature of $T=1$ recovers the standard softmax.

The distillation loss, $L_{KD}$, is typically the Kullback-Leibler (KL) divergence between the student's softened predictions ($p_S = \sigma(z_S, T)$) and the teacher's softened predictions ($p_T = \sigma(z_T, T)$):

$$L_{KD} = T^2 \cdot D_{KL}(p_S || p_T) = T^2 \sum_i p_T(i) \log \frac{p_T(i)}{p_S(i)}$$

The $T^2$ scaling factor is important. Because the gradients produced by the softened targets are scaled by $1/T^2$ relative to the gradients from hard targets, multiplying the KD loss by $T^2$ ensures that the relative contribution of the KD loss during training remains roughly constant even if the temperature $T$ is changed.

This objective encourages the student not just to predict the correct output but to understand why the teacher predicts it, learning the relationships between outputs. In practice, this $L_{KD}$ is often combined with a standard supervised loss (e.g., cross-entropy $L_{CE}$) on the true hard labels, using a weighting factor $\alpha$:

$$L_{Total} = \alpha L_{CE}(y_{true}, \sigma(z_S, T=1)) + (1 - \alpha) L_{KD}(\sigma(z_S, T), \sigma(z_T, T))$$

This ensures the student still learns to match the ground truth while benefiting from the teacher's soft targets. Choosing the optimal temperature $T$ and weight $\alpha$ often requires empirical tuning.

Intermediate Representation Matching

While matching output distributions is effective, knowledge isn't solely contained in the final layer. Intermediate layers of deep networks like LLMs learn hierarchical representations that capture syntactic, semantic, and contextual information. Distilling this intermediate knowledge can provide stronger guidance to the student.

The objective here is to minimize the difference between the hidden states or activations of selected intermediate layers in the teacher ($h_T^l$) and student ($h_S^l$) models. Common loss functions for this include:

• Mean Squared Error (MSE): Directly minimizes the Euclidean distance. $$L_{Intermediate}^{MSE} = \sum_{l \in L_{match}} || f_S(h_S^l) - f_T(h_T^l) ||_2^2$$
• Cosine Similarity Loss: Maximizes the cosine similarity, focusing on the angle between representation vectors rather than their magnitude. Useful when magnitudes might differ significantly. $$L_{Intermediate}^{Cosine} = \sum_{l \in L_{match}} (1 - \cos(f_S(h_S^l), f_T(h_T^l)))$$

Here, $L_{match}$ is the set of layer indices chosen for matching. The functions $f_S$ and $f_T$ represent optional transformation layers (e.g., linear projections) used to align the dimensions if the student and teacher layers have different hidden sizes.

Considerations for intermediate matching include:

• Layer Selection: Which layers contain the most valuable information to transfer? Early layers might capture basic features, while deeper layers capture more abstract concepts. Matching multiple layers is common.
• Architecture Mismatch: If the student has fewer layers than the teacher, strategies are needed to map teacher layers to student layers (e.g., matching the last layer of the student to the last layer of the teacher, or using learned projections).
• Computational Cost: Calculating these losses adds overhead to the training process, especially if matching large activation tensors.

Attention Transfer

For Transformer-based LLMs, the self-attention mechanism is a defining component. Attention maps, which represent the weighted relationships between tokens at different positions, encode significant structural and contextual information. Transferring this attention knowledge can help the student learn similar relational patterns.

The Attention Transfer (AT) objective minimizes the difference between attention maps ($A_T^l, A_S^l$) from corresponding layers:

$$L_{Attention} = \sum_{l \in L_{match}} \frac{1}{N_h} \sum_{h=1}^{N_h} || A_{S,h}^l - A_{T,h}^l ||_F^2$$

Where $N_h$ is the number of attention heads, and $|| \cdot ||_F^2$ denotes the squared Frobenius norm (sum of squared elements) of the difference between the attention matrices for head $h$ in layer $l$.

Challenges include:

• Head Mapping: If the student and teacher have different numbers of attention heads, a mapping strategy is needed (e.g., averaging teacher heads, using learned projections).
• Computational Overhead: Storing and comparing attention maps adds memory and computation requirements.
• Interpretation: While attention patterns are informative, directly forcing the student to match them might be overly restrictive compared to matching hidden states.

Contrastive Learning Objectives

Rather than directly matching representations element-wise (like MSE) or angle-wise (like Cosine), contrastive objectives focus on learning similarities and dissimilarities. Contrastive Representation Distillation (CRD) aims to teach the student to produce representations that are close to the teacher's representation for the same input (positive pair) but far from the teacher's representations for different inputs (negative pairs).

Speaking, the loss encourages $sim(z_S, z_T)$ to be high for the same input, while $sim(z_S, z_{T,neg})$ should be low, where $z_{T,neg}$ are teacher representations for other inputs in the batch or a memory bank. A typical loss function like InfoNCE can be adapted:

$$L_{Contrastive} \propto -\log \frac{\exp(sim(f_S(h_S), f_T(h_T))/\tau)}{\sum_{h_{T,neg}} \exp(sim(f_S(h_S), f_T(h_{T,neg}))/\tau)}$$

Here, $sim$ is a similarity function (e.g., dot product or cosine similarity), $\tau$ is a temperature parameter controlling the sharpness of the distribution over negative samples, and the sum is over negative teacher representations. $f_S$ and $f_T$ are again potential projection heads.

Contrastive objectives can be powerful for learning the underlying structure of the teacher's representation space without requiring strict element-wise alignment, potentially offering more flexibility to the student.

Combining Distillation Objectives

Often, the most effective distillation strategy involves combining multiple objectives. Different objectives capture complementary aspects of the teacher's knowledge. For example, one might combine soft target matching with intermediate feature matching and attention transfer:

$$L_{Total} = \lambda_{KD} L_{KD} + \lambda_{Inter} L_{Intermediate} + \lambda_{Attn} L_{Attention} + \dots$$

The hyperparameters $\lambda_i$ control the relative importance of each objective. Selecting the right combination and tuning these weights is a critical part of designing a successful distillation pipeline and typically involves significant experimentation based on the specific task, teacher-student architecture pairing, and performance metrics.

Knowledge transfer points between teacher and student models, illustrating common distillation objectives: matching output logits ($L_{KD}$), intermediate hidden states ($L_{Intermediate}$), and attention maps ($L_{Attention}$).

Choosing the appropriate objective, or combination of objectives, depends heavily on the specific characteristics of the teacher and student models, the available data, computational budget, and the desired trade-offs between student model size, speed, and fidelity.