Knowledge Distillation

What is Knowledge Distillation?

Large deep neural networks (Teacher models like GPT-4 or 70B LLMs) achieve outstanding accuracy, but are too slow and expensive to deploy for high-throughput real-time APIs.

Knowledge Distillation (Geoffrey Hinton et al., 2015) compresses a large, high-capacity Teacher Model into a compact, ultra-fast Student Model (e.g. DistilBERT or 8B LLMs).

  TEACHER MODEL (Giant, 70B Params) ──► Soft Logits (T = 4.0) ──┐
                                                                 ├──► [ KL DIVERGENCE LOSS ] ──► Update STUDENT (8B)!
  STUDENT MODEL (Compact, 8B Params) ─► Soft Logits (T = 4.0) ──┘

The Student achieves $\sim 95%$ of Teacher accuracy while running $10\times$ faster on cheaper inference hardware!

The Concept of "Dark Knowledge"

Why train a Student model on Teacher probabilities instead of raw ground-truth hard labels ($[1, 0, 0]$)?

Consider classifying an image of a BMW Car:

• Hard One-Hot Label: [1.0 Car, 0.0 Truck, 0.0 Cat]
• Teacher Soft Probabilities: [0.85 Car, 0.14 Truck, 0.00001 Cat]

The Teacher's soft probabilities reveal Dark Knowledge: a BMW shares visual features with a Truck, but shares zero features with a Cat!

Hard one-hot labels discard this rich structural relationship. Soft targets convey rich geometric decision boundaries to the Student.

Temperature-Scaled Softmax

To extract dark knowledge from small non-target logits, apply Temperature Scaling ($T$) to Softmax outputs:

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

┌──────────────────────────┬──────────────────────────┐
│ TEMPERATURE T = 1.0      │ TEMPERATURE T = 4.0      │
├──────────────────────────┼──────────────────────────┤
│ Standard Softmax.        │ Softens probabilities!   │
│ Sharp probabilities.     │ Magnifies tiny non-target│
│ Non-target logits hidden.│ class logit relationships│
└──────────────────────────┴──────────────────────────┘

Setting $T > 1.0$ (typically $T = 3.0\text{--}5.0$) smooths the output distribution, making relative non-target logit probabilities visible to the Student.

Combined Loss Function

The Student is trained using a weighted combination of two losses:

$$\mathcal{L}{\text{Student}} = \alpha \cdot T^2 \cdot \mathcal{L}{\text{KL}}(q_s^T, q_t^T) + (1 - \alpha) \cdot \mathcal{L}{\text{CE}}(y_s, y{\text{true}})$$

1. Distillation Loss ($\mathcal{L}_{\text{KL}}$): Kullback-Leibler Divergence between Student soft logits $q_s^T$ and Teacher soft logits $q_t^T$ at temperature $T$. (Scaled by $T^2$ to balance gradient magnitudes).
2. Student Loss ($\mathcal{L}_{\text{CE}}$): Standard Cross-Entropy loss between Student outputs at $T=1.0$ and hard ground-truth labels $y_{\text{true}}$.

Distillation Types in Practice

• Response-Based Distillation: Student mimics final output probabilities of the Teacher (Hinton Distillation).
• Feature-Based Distillation (FitNets): Student hidden layers learn to match intermediate feature map activations of Teacher hidden layers.
• Sequence Level Distillation (MiniLM / DistilBERT): Student learns to match self-attention matrix distributions of Transformer blocks.

Say this out loud

Knowledge Distillation compresses a large Teacher model into a small Student model using soft probability targets. Temperature scaling softens logit distributions to reveal Dark Knowledge inter class similarity relationships. The Student is trained using KL divergence on soft Teacher logits combined with Cross Entropy on hard labels.

Followups to expect

1. What is Self-Distillation? Training a student model using a teacher model that shares the exact same network architecture, where the student iteratively matches past epoch predictions to improve generalization.
2. What is Synthetic Data Distillation for LLMs? Using a massive teacher LLM (GPT-4) to generate 100,000 synthetic reasoning outputs, training a smaller student LLM (LLaMA 8B) on the generated text using standard SFT.