Logit Distillation

The central goal of logit distillation is to transfer the dark knowledge embedded in a teacher model's output logits. Unlike hard labels that only indicate the correct class, the softened probability distribution reveals the teacher's relative confidence across all classes, including non-target similarities (e.g., that a 'cat' is more similar to a 'dog' than to an 'airplane'). This provides a richer, more informative training signal for the student model, leading to better generalization and calibration than training on labels alone.

Logit distillation operates by applying temperature scaling to the teacher's logits. A temperature parameter T is introduced into the softmax function: softmax(z_i / T).

• High T (T > 1): 'Softens' the probability distribution, making it more uniform and amplifying the dark knowledge in non-target class probabilities.
• T = 1: Standard softmax, producing the model's normal prediction distribution. The student is trained to match these soft targets from the teacher, typically using a Kullback-Leibler (KL) Divergence loss, while also being trained on the ground truth labels with a standard cross-entropy loss.

The standard objective for logit distillation is the Kullback-Leibler (KL) Divergence between the softened output distributions of the teacher and student models. The total loss is often a weighted sum: L_total = α * L_CE(y, σ(z_s)) + β * T^2 * L_KL(σ(z_t / T) || σ(z_s / T)) Where:

• L_CE is cross-entropy with true labels y.
• L_KL is the KL divergence.
• z_t and z_s are teacher and student logits.
• α and β are balancing weights. The T^2 term is used to scale the gradients, ensuring the magnitude of the distillation loss is invariant to the temperature setting.

In continual learning, logit distillation is a critical tool for catastrophic forgetting mitigation. Algorithms like Learning without Forgetting (LwF) use self-distillation: when learning a new task, the model's own predictions (logits) on data from previous tasks are used as soft targets. By adding a distillation loss that penalizes deviation from these old outputs, the model is regularized to retain its original functionality while integrating new knowledge. This makes logit distillation a primary method for enforcing the stability side of the plasticity-stability trade-off.

Logit distillation offers distinct practical advantages compared to feature distillation (matching intermediate layer activations):

• Simplicity: It only requires access to the model's final output layer, making it architecture-agnostic and easy to implement.
• Efficiency: It avoids the computational and memory overhead of aligning high-dimensional intermediate feature maps.
• Effectiveness: For many model compression tasks (e.g., creating DistilBERT, TinyBERT), matching softened logits alone is sufficient to achieve significant compression with minimal performance drop. It is often the first and most widely used baseline in distillation research.

While powerful, logit distillation has limitations that lead to its use alongside other methods:

• Representational Bottleneck: Distilling only the final output may miss crucial representational knowledge encoded in the teacher's intermediate layers.
• Task Specificity: The logits are tied to the teacher's specific output head; for transferring general feature extraction abilities, feature distillation or attention distillation can be more effective.
• Data Dependency: It typically requires a transfer dataset, which may not be available in all scenarios. Consequently, in advanced compression pipelines, logit distillation is frequently combined with intermediate layer matching (e.g., as in TinyBERT's two-stage process) for optimal knowledge transfer.

Logit distillation facilitates knowledge transfer from a model trained on a large, general-source dataset (e.g., ImageNet) to a student specialized for a target domain with limited labeled data.

• Process: The teacher's generalized logits provide a robust starting point, and the student is fine-tuned on the target data using a combined loss of target labels and teacher logits.
• Outcome: The student benefits from the teacher's broad feature representations while specializing, often outperforming training from scratch or standard fine-tuning on small target datasets.
• Related Technique: Can be combined with feature distillation (matching intermediate activations) for an even stronger transfer of representational knowledge.

Neural networks are often poorly calibrated, meaning their predicted confidence scores do not reflect true likelihood. Distillation with temperature scaling produces softer probability distributions, which the student learns to mimic.

• Result: Student models trained via logit distillation frequently exhibit better calibration than those trained on hard labels alone. Their confidence scores are more reliable indicators of prediction correctness.
• Why It Works: The teacher's softened outputs (e.g., 'cat: 0.8, dog: 0.15, lion: 0.05') convey uncertainty and class relationships, teaching the student to output less overconfident, more nuanced probabilities.
• Importance: Critical for high-stakes applications like medical diagnosis or autonomous systems where understanding model certainty is as important as the prediction itself.

Distillation is a powerful method for condensing the knowledge of a computationally expensive model ensemble into a single, efficient student model.

• Procedure: The logits from multiple teacher models are averaged or combined to create a single, higher-quality set of soft targets. The student is then trained to match this consensus.
• Advantage: The student captures the diversity and robustness of the ensemble but requires only the computational cost of a single model at inference time.
• Statistical Benefit: The ensemble's aggregated logits often have lower variance and better generalization than any single teacher, transferring this improved inductive bias to the student.

Knowledge distillation is the overarching model compression and transfer learning paradigm where a compact student model is trained to mimic the behavior of a larger, more capable teacher model. Logit distillation is its most common form.

• Core Mechanism: The teacher provides a rich supervisory signal beyond simple labels.
• Primary Goals: Model compression (smaller, faster models) and performance improvement (better training signals).
• Variants: Includes logit, feature, and attention distillation.

Dark knowledge refers to the rich, comparative information embedded within a trained neural network's logit vectors (pre-softmax activations). This is the specific 'knowledge' transferred during logit distillation.

• Contrast with Hard Labels: A one-hot label only indicates the correct class. Dark knowledge encodes the teacher's relative confidence across all classes (e.g., 'this is very likely a cat, somewhat similar to a lynx, and not a truck').
• Accessing Dark Knowledge: The temperature scaling parameter (T > 1) in the softmax function is used to soften the output probabilities, making this comparative information more pronounced and accessible for the student to learn.

Temperature scaling is a critical hyperparameter technique used to modulate the output probability distribution of the teacher model in logit distillation.

• Formula: Softmax is modified as (\text{softmax}(z_i) = \frac{e^{z_i/T}}{\sum_j e^{z_j/T}}), where (z) are logits and (T) is the temperature.
• Effect of High T (T > 1): 'Softens' the distribution, increasing the entropy and amplifying the dark knowledge in non-true classes.
• Effect of Low T (T -> 0): Hardens the distribution towards a one-hot encoding.
• Standard Practice: A temperature like T=3 or T=4 is often used during distillation, while T=1 is used for final student inference.

Kullback-Leibler (KL) Divergence is the primary loss function used to train the student model in logit distillation. It measures the statistical distance between two probability distributions.

• Objective: Minimize the KL divergence between the soft targets from the teacher (with temperature) and the soft predictions from the student (with the same temperature).
• Standard Formulation: (\mathcal{L}{KD} = T^2 \cdot D{KL}(\mathbf{p}\tau^\text{teacher} | \mathbf{p}\tau^\text{student})), where (\mathbf{p}_\tau) are temperature-scaled softmax outputs. The (T^2) scaling factor ensures gradient magnitudes are consistent when temperature varies.
• Combined Loss: Often used alongside a standard cross-entropy loss with true labels: (\mathcal{L} = \alpha \cdot \mathcal{L}{CE} + \beta \cdot \mathcal{L}{KD}).

Learning without Forgetting is a foundational continual learning algorithm that directly applies logit distillation to mitigate catastrophic forgetting.

• Core Idea: When learning a new task, the model's own predictions (logits) on new task data, recorded before new training begins, are used as soft targets for the old tasks. This self-distillation signal regularizes the model to retain its previous functionality.
• Key Advantage: Does not require storing raw training data from previous tasks (exemplar-free), enhancing privacy and reducing storage.
• Process: 1) Record network outputs for new data using old parameters. 2) Update network on new task with labels. 3) Use a distillation loss (e.g., KL) to match new outputs to the recorded old outputs, preserving past knowledge.

Soft targets are the probability distributions generated by the teacher model after temperature scaling. They are the direct learning objective for the student in logit distillation.

• Nature: A vector of probabilities (e.g., [0.7, 0.25, 0.05]) instead of a hard target (e.g., [1, 0, 0]).
• Pedagogical Role: Provide a richer, more informative training signal by indicating inter-class relationships. The student learns not just 'what is correct' but 'what is similar and what is dissimilar'.
• Contrast with Hard Labels: Training with soft targets often leads to better generalization, faster convergence, and a more calibrated student model compared to training solely on hard labels.