Distillation Loss

The fundamental role of distillation loss is to serve as the training objective that guides the student model to mimic the teacher. Unlike standard supervised loss (e.g., cross-entropy with hard labels), distillation loss uses the teacher's softened outputs or intermediate features as the learning target. This transfers not just the final decision but the teacher's generalization capabilities and understanding of inter-class relationships.

The most common function, using Kullback-Leibler (KL) Divergence, measures how one probability distribution diverges from another. In logit distillation, the loss minimizes the KL divergence between the student's and teacher's softened output distributions. This teaches the student the teacher's confidence across all classes, including non-target classes, which is the essence of dark knowledge transfer.

• Key Formula: (L_{KD} = T^2 \cdot KL(\sigma(z_s / T) , || , \sigma(z_t / T))), where (z) are logits, (\sigma) is softmax, and (T) is the temperature.

Beyond final outputs, distillation loss can align intermediate feature representations. This function involves designing a loss (e.g., Mean Squared Error, Cosine Similarity) to minimize the distance between the teacher's and student's activations from selected hidden layers. This transfers the teacher's internal data transformations and abstract representations, often leading to more robust student models. Techniques like Attention Transfer and Hint Learning fall under this category.

Distillation loss incorporates temperature scaling (parameter T > 1) to soften the teacher's output probability distribution. This function creates a richer, smoother training signal compared to one-hot encoded hard labels. Higher temperatures produce a more uniform distribution, emphasizing the relative similarities between classes. The student learns a broader, more nuanced decision boundary, which is critical for combating overfitting and improving generalization on limited data.

In practice, distillation loss is rarely used alone. Its core function is to be combined with the standard supervised task loss (e.g., cross-entropy with ground truth labels). The total loss is a weighted sum: (L_{total} = \alpha L_{task} + \beta L_{distill}). This dual-objective function ensures the student both learns the correct task from true labels and absorbs the teacher's refined knowledge from soft targets, balancing accuracy and generalization.

A critical function of distillation loss is to act as a regularizer and guide during aggressive model compression. When a student model is significantly smaller or has undergone quantization or pruning, its capacity to learn from hard labels alone is limited. Distillation loss provides a high-quality, information-rich signal that helps the compact student recover accuracy that would otherwise be lost, making extreme compression for on-device deployment viable.

Kullback-Leibler Divergence is the primary statistical measure used as the distillation loss for logit-based knowledge transfer. It quantifies how one probability distribution (the student's softened outputs) diverges from a second, reference distribution (the teacher's softened outputs).

• Mechanism: Minimizing the KL divergence encourages the student's output probabilities to become statistically indistinguishable from the teacher's.
• Application: It is the standard loss function in the original distillation formulation, enabling the student to learn the teacher's 'dark knowledge'—the nuanced inter-class relationships embedded in the softened probability distribution.

Soft targets are the probability distributions generated by applying a softened softmax function to the teacher model's output logits. They are the supervisory signal for the student in logit distillation.

• Key Feature: Compared to hard, one-hot labels, soft targets provide a richer training signal by preserving the relative probabilities the teacher assigns to incorrect classes.
• Example: For an image of a cat, a teacher might output [cat: 0.9, dog: 0.09, lion: 0.01]. This vector, not just [1, 0, 0], is the soft target that teaches the student about visual similarities between classes.

Temperature scaling is a hyperparameter technique used to control the 'softness' of the probability distributions produced by the teacher and student models during logit distillation.

• Function: A temperature parameter T > 1 is introduced into the softmax function: softmax(z_i / T). Higher values of T produce smoother, more uniform probability distributions.
• Purpose: This smoothing amplifies the dark knowledge in the teacher's outputs, making it easier for the student model to learn the relative similarities between classes. The final inference is performed with T = 1.

Feature distillation is a knowledge transfer method where the student model is trained to replicate the intermediate feature representations or activations from specific layers of the teacher model, rather than just its final outputs.

• Objective: The distillation loss is applied to the outputs of intermediate layers (e.g., using Mean Squared Error or Cosine Similarity).
• Benefit: It provides a more direct and granular learning signal, guiding the student's internal representations to align with the teacher's, which can be especially effective for compressing deep or wide architectures.

Attention transfer is a specific form of feature distillation designed for architectures that use attention mechanisms, such as Transformers or certain CNNs. The student learns to mimic the teacher's spatial or self-attention maps.

• Mechanism: The distillation loss minimizes the difference between the attention matrices generated by teacher and student layers. For vision transformers, this can involve distilling the attention from the [CLS] token's interactions with image patches.
• Outcome: This forces the student to learn which parts of the input the teacher deems most important for the task, leading to more sample-efficient learning.

Mean Squared Error Loss is a common alternative distillation objective, particularly in feature-based and regression-style distillation tasks.

• Usage: It directly minimizes the squared L2 distance between the teacher and student outputs, whether those outputs are logits, feature maps, or attention scores.
• Comparison to KL Divergence: While KL divergence is probabilistic, MSE is a distance metric in the output space. MSE is often used when the goal is direct regression of the teacher's representations, or in data-free distillation scenarios where generating valid probability distributions is challenging.