Knowledge distillation

WHAT is knowledge distillation?

WHY do soft targets carry more information? Consider an image classifier over {cat, dog, car}. The hard label for a photo of a cat is $[1,0,0]$ . But the teacher might output $[0.9, 0.099, 0.001]$ . That tells the student something the hard label cannot: "a cat looks much more like a dog than like a car." This inter-class similarity structure is called dark knowledge.

The mechanism: softmax with temperature

Standard softmax over logits $z_i$ : $p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$

Softmax with temperature $T$ : $p_i(T) = \frac{e^{z_i/T}}{\sum_j e^{z_j/T}}$

$T = 1$ : normal softmax.
$T \to \infty$ : all logits divided by huge number $\Rightarrow$ probabilities $\to$ uniform $1/K$ .
$T \to 0$ : sharpens toward a one-hot (argmax).

The loss function — derived from first principles

We want the student to do two things at once:

(1) Match the teacher's soft targets. Use cross-entropy (equivalently KL divergence) between teacher soft probs $p^T(T)$ and student soft probs $q^T(T)$ , both at temperature $T$ : $\mathcal{L}_{\text{soft}} = -\sum_i p_i^{T}(T)\,\log q_i^{T}(T)$

(2) Still get the true answer right. Standard cross-entropy with hard labels $y$ at $T=1$ : $\mathcal{L}_{\text{hard}} = -\sum_i y_i \,\log q_i(1)$

Combine them with a mixing weight $\alpha$ :

WHY the $T^2$ factor? (This is the step everyone forgets.) When you backprop through softmax-with-temperature, the gradient of $\mathcal{L}_{\text{soft}}$ w.r.t. student logits scales like $1/T^2$ (one $1/T$ from the student softmax, one from the chain rule through the logit division). To keep the magnitude of the soft-target gradient comparable to the hard-target gradient, we multiply by $T^2$ . Without it, raising $T$ silently shrinks the soft loss into irrelevance.

Gradient sketch (WHY $1/T^2$ ): For softmax cross-entropy, $\partial \mathcal{L}/\partial z_i = q_i - p_i$ . With temperature, the student prob is $q_i(T)$ and each logit is scaled by $1/T$ , so $\frac{\partial \mathcal{L}_{\text{soft}}}{\partial z_i} = \frac{1}{T}\big(q_i(T) - p_i(T)\big).$ For small logits (Taylor expand $e^{z/T}\approx 1+z/T$ ), $q_i(T)-p_i(T) \approx \frac{1}{KT}(z_i - v_i)$ where $v_i$ are teacher logits, giving an overall $\sim 1/T^2$ . Multiplying the loss by $T^2$ cancels it. ✅

Worked example 1 — softening a distribution

Teacher logits $z = [4, 1, -2]$ over {cat, dog, car}.

At $T=1$ : $e^4=54.6,\ e^1=2.72,\ e^{-2}=0.135$ , sum $=57.5$ . $p = [0.950,\ 0.047,\ 0.0023]$ . Why this step? Confidence is crushed onto "cat"; dog/car info is nearly invisible.

At $T=2$ : logits become $[2, 0.5, -1]$ . $e^2=7.39,\ e^{0.5}=1.65,\ e^{-1}=0.368$ , sum $=9.40$ . $p = [0.786,\ 0.175,\ 0.039]$ . Why this step? Now "dog" is a clearly non-trivial 17.5% — the student can learn that cats resemble dogs. That's the dark knowledge surfacing.

Worked example 2 — computing the combined loss

Say $\alpha = 0.5$ , $T = 2$ . Student soft probs $q^T(2) = [0.70, 0.25, 0.05]$ , teacher $p^T(2)=[0.786,0.175,0.039]$ , true label = cat, student normal probs $q(1)=[0.80,0.15,0.05]$ .

Soft term: $\mathcal{L}_{\text{soft}} = -(0.786\ln 0.70 + 0.175\ln 0.25 + 0.039\ln 0.05)$ $= -(0.786(-0.357) + 0.175(-1.386) + 0.039(-3.00)) = 0.640$ .

Hard term: $\mathcal{L}_{\text{hard}} = -\ln(0.80) = 0.223$ .

Total: $\mathcal{L} = 0.5\cdot 2^2 \cdot 0.640 + 0.5\cdot 0.223 = 1.280 + 0.112 = 1.392$ . Why this step? The $T^2=4$ scaling makes the soft term dominate here — exactly its purpose: force the student to absorb the teacher's structure while the hard label keeps it honest.

KD for LLMs specifically

Common LLM KD variants:

Response-based (logit) KD: match next-token logit distributions. ← classic Hinton KD.
Feature-based KD: match intermediate hidden states/attention maps (student learns how the teacher thinks, not just what it outputs).
Sequence-level KD: teacher generates text, student trained on that generated text as data (used heavily for instruction-following distillation, e.g. Alpaca-style).

Common mistake Steel-manning the common errors

Mistake 1: "Just use $T=1$ , why complicate it?" Why it feels right: $T=1$ is the "real" probability the teacher believes. The fix: At $T=1$ the teacher is near one-hot, so the soft targets ≈ hard targets and you gain almost nothing. You need $T>1$ (often 2–10) to expose dark knowledge.

Mistake 2: Forgetting the $T^2$ factor. Why it feels right: The loss "looks fine" and still decreases. The fix: Without $T^2$ , the soft-loss gradient is $\sim 1/T^2$ smaller, so higher temperatures accidentally turn off distillation. Always scale.

Mistake 3: Distilling with a teacher weaker than needed, or on wrong data. Why it feels right: Any teacher gives soft labels. The fix: A poor teacher transfers poor dark knowledge; and if you only distill on data the teacher never saw well, its soft labels are noise. Match distillation data to teacher competence.

Mistake 4: Confusing soft targets with label smoothing. Why it feels right: Both produce non-one-hot targets. The fix: Label smoothing spreads probability uniformly (no info about which classes are similar). KD's soft targets are structured — that structure is the whole point.

Recall Feynman: explain to a 12-year-old

Imagine a genius chef (teacher) and a student cook. If the chef only says "this dish is 'pizza'," the student learns a little. But if the chef says "it's mostly pizza, kind of like a calzone, and nothing like soup," the student learns way more — including how foods relate to each other. Knowledge distillation is teaching the small model using the big model's full opinion, not just its final answer. The "temperature" is like asking the chef to slow down and share the subtle in-between thoughts instead of blurting the answer.

Recall checkpoints

Flashcards

What is knowledge distillation?

Training a small student model to imitate a large teacher model's softened output distribution (soft targets) rather than only hard labels.

What are "soft targets"?

The teacher's full probability distribution over classes (softened with temperature), carrying inter-class similarity info.

What is "dark knowledge"?

The relational information hidden in a teacher's non-argmax probabilities (e.g. cat resembles dog more than car).

Write softmax with temperature.

p_i(T)=\dfrac{e^{z_i/T}}{\sum_j e^{z_j/T}}

What happens as

T\to\infty

The distribution becomes uniform (

1/K

each).

What happens as

T\to 0

The distribution sharpens toward one-hot (argmax).

Write the distillation loss.

\mathcal{L}=\alpha T^2\mathcal{L}_{\text{soft}}+(1-\alpha)\mathcal{L}_{\text{hard}}

Why the

T^2

factor?

The soft-loss gradient scales like

1/T^2

; multiplying by

T^2

keeps its magnitude comparable to the hard-loss gradient.

Why not use

T=1

A trained teacher is near one-hot at

T=1

, so soft targets collapse to hard targets and dark knowledge is hidden.

KD vs label smoothing?

Label smoothing adds uniform noise (no structure); KD soft targets encode which classes are actually similar.

Name three LLM distillation variants.

Response/logit-based, feature-based (hidden states/attention), sequence-level (train on teacher-generated text).

Which term keeps the student factually correct?

The hard-label cross-entropy term

\mathcal{L}_{\text{hard}}

weighted by

(1-\alpha)

Connections

Softmax and Temperature
Model Compression
Pruning and Quantization
Cross-Entropy and KL Divergence
DistilBERT
Instruction Tuning
Transfer Learning

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, knowledge distillation ka funda simple hai: ek bada powerful "teacher" model hai jo already achha train ho chuka hai, aur ek chhota "student" model banana hai jo fast aur lightweight ho. Ab agar hum student ko sirf hard label (jaise "yeh cat hai", ek-hot vector) se sikhayein, toh woh utni hi info paata hai. Lekin teacher jab predict karta hai toh woh poori probability distribution deta hai — jaise "90% cat, 9% dog, 0.1% car". Yeh distribution batata hai ki cat dog jaisi lagti hai par car jaisi bilkul nahi. Isko hum "dark knowledge" kehte hain, aur yehi cheez student ko extra smart banati hai.

Problem yeh hai ki achha teacher over-confident hota hai — woh 0.9999 type numbers deta hai, toh useful similarity info chhup jaati hai. Isliye hum temperature $T$ use karte hain. Logits ko $T$ se divide karke softmax lete hain, jisse distribution flat ho jaati hai aur woh chhoti-chhoti values (dog wali) clearly visible ho jaati hain. $T$ bada karo toh distribution uniform ki taraf jaati hai, chhota karo toh sharp (one-hot) ban jaati hai.

Loss mein do cheezein combine hoti hain: ek soft loss (teacher ki soft probs match karo) aur ek hard loss (sahi answer bhi sahi rakho), aur inko $\alpha$ se mix karte hain. Ek important detail — soft loss ko $T^2$ se multiply karna padta hai, kyunki temperature ki wajah se gradient $1/T^2$ chhota ho jaata hai, aur $T^2$ usko wapas balance kar deta hai. Yeh factor log bhool jaate hain aur phir distillation kaam hi nahi karta.

LLMs ke liye yeh technique bahut powerful hai kyunki har token pe poore vocabulary (50k words) ki distribution milti hai — matlab bahut rich signal. Isi se DistilBERT, TinyLlama jaise chhote models bante hain jo bade models ki competence sasti mein inherit kar lete hain. Yaad rakho: Soft, Warm, Squared.

WHAT is knowledge distillation?

The mechanism: softmax with temperature

The loss function — derived from first principles

Worked example 1 — softening a distribution

Worked example 2 — computing the combined loss

KD for LLMs specifically

Recall checkpoints

Flashcards

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Pretraining & Fine-Tuning LLMs

Connections