Knowledge Distillation: Theory & Methods

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Soft Targets and Temperature
5. The Distillation Loss
6. Feature-Based Distillation
7. Relation-Based Distillation
8. Why Distillation Works
9. Self-Distillation
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive the KD loss with temperature scaling.
2. Explain why soft targets carry more information than hard labels.
3. Describe feature-based and relation-based distillation variants.
4. Prove that distillation provides a better learning signal than standard training.
5. Explain self-distillation and why it improves even without a larger teacher.

Notation

• $T$ — temperature parameter
• $p^T$ — teacher's softmax output at temperature $T$
• $q^S$ — student's softmax output
• $\mathcal{L}_{\text{KD}}$ — distillation loss
• $\alpha$ — balancing coefficient

Core Intuition

A large "teacher" model knows more than its hard predictions reveal. Its soft probabilities encode dark knowledge: which wrong classes are similar to the correct one, confidence levels, inter-class relationships. By training a small "student" to match these soft distributions, the student learns richer information than from hard labels alone — achieving accuracy closer to the teacher with far fewer parameters.

Soft Targets and Temperature

Standard softmax: $p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$ — often near one-hot for trained models.

Temperature-scaled softmax:

• $T = 1$: standard (peaked).
• $T > 1$: softer (reveals inter-class structure).
• $T \to \infty$: uniform distribution.

Information content: At high temperature, the soft probabilities reveal which classes the teacher considers similar. E.g., for a "cat" image: P(dog)=0.1, P(car)=0.001 tells the student that cats are more like dogs than cars.

The Distillation Loss

Components:

• Distillation loss ($T^2$-scaled KL): Match teacher's soft distribution. The $T^2$ factor compensates for the reduced gradient magnitude at high temperature.
• Hard label loss (cross-entropy with true labels): Ensure student still predicts correctly.
• $\alpha$: Typically 0.5–0.9 (weight toward distillation).

Typical $T$: 2–20. Higher for more "knowledge transfer," lower for more "accuracy."

Feature-Based Distillation

Match intermediate representations, not just outputs:

where $\phi_l$ is a learned projection (since teacher and student may have different dimensions).

Variants:

• FitNets: Match hidden layer activations via a regressor.
• Attention transfer: Match attention maps across layers.
• CRD (Contrastive): Maximize mutual information between teacher and student features.

Relation-Based Distillation

Instead of matching individual features, match the relationships between samples:

where $r$ is a relation function (e.g., cosine similarity, distance).

Intuition: The student learns the teacher's geometry — which samples the teacher considers similar/dissimilar.

Why Distillation Works

Theorem (informal): Soft labels have higher effective "label smoothing" and provide more bits of information per training example than hard labels.

Arguments:

1. Richer gradient signal: Hard labels provide gradient only for the correct class. Soft labels provide gradients for ALL classes (weighted by teacher confidence).
2. Implicit regularization: Soft targets prevent overconfident predictions, acting as label smoothing.
3. Curriculum effect: The teacher's confidence indicates difficulty — easy examples get peaked distributions, hard examples get uncertain ones.

Self-Distillation

Surprising finding: Distilling a model into an identical architecture (same size) improves performance.

Mechanism: The teacher (same model from a previous training run or an ensemble of checkpoints) provides smoothed labels that regularize the student training.

Born-Again Networks: Train generation 1 normally, then distill into generation 2 (same size). Generation 2 often outperforms generation 1 by 0.5-1%.

Common Pitfalls

Pitfall 1. Setting temperature too low ($T=1$). The teacher's output is near one-hot — almost no dark knowledge is transferred. Use $T=4$–$20$.

Pitfall 2. Making the student too small. If the capacity gap is too large (e.g., 70B teacher → 100M student), distillation can't bridge it. Intermediate-sized students work better.

Pitfall 3. Forgetting the $T^2$ scaling factor. Without it, the distillation gradient is $T^2$ times too small relative to the hard label loss, making distillation ineffective.

Summary

• Knowledge distillation transfers soft probability distributions from teacher to student.
• Temperature $T$ controls softness: higher = more inter-class information.
• Loss: $\alpha T^2 \cdot \text{KL}(\text{teacher} \| \text{student}) + (1-\alpha) \cdot \text{CE}(\text{label}, \text{student})$.
• Feature distillation matches intermediate layers; relation distillation matches sample geometry.
• Distillation provides richer gradients and implicit regularization.
• Even self-distillation (same size) improves performance.

Exercises

Exercise 1. For a 3-class problem with teacher logits $[5, 3, 1]$: compute the soft probabilities at $T=1, 4, 10$.

Exercise 2. Derive why the $T^2$ factor is needed in the KD loss (hint: analyze the gradient magnitude of KL at temperature $T$).

Exercise 3. Design a feature distillation setup for distilling BERT-large into BERT-small (specify which layers to match and projection dimensions).

Exercise 4. Explain why distillation from an ensemble of teachers is more effective than from a single teacher.

Exercise 5. Prove that in the limit $T \to \infty$, the KL distillation loss reduces to MSE between logits.