Distill Useful Behaviour From a Teacher

Section 1: The Practical Outcome First

In Lesson 13, we sized our target architecture and selected a 1B parameter model (Llama 3.2 1B) to operate under a strict 2.0 GB container RAM limit on edge Kubernetes nodes. However, deploying a raw 1B base model directly to production fails: its zero-shot reasoning is brittle, its structured tool output violates schemas under edge cases, and its domain accuracy falls far short of an unconstrained 70B parameter frontier model.

The goal of Knowledge Distillation (KD) is to compress the reasoning, calibration, and domain judgment of a massive teacher model (e.g., Llama 3.1 70B) into our compact student, allowing the edge service to achieve frontier-grade reliability at a fraction of the serving cost and latency.

[Knowledge Distillation Architecture: Teacher -> Student]
┌────────────────────────────────────────────────────────┐
│ 70B Teacher Model (Hosted Centrally / Multi-GPU Pool)  │
│ - Massive parameter capacity & multi-step judgment    │
│ - Evaluates domain incident stack traces               │
└───────────────────────────┬────────────────────────────┘
                            │
            ┌───────────────┴───────────────┐
            ▼                               ▼
┌──────────────────────────────┐┌──────────────────────────────┐
│ White-Box Logit Distillation ││ Black-Box Sequence Distill.  │
│ (Hinton et al., 2015)        ││ (Kim & Rush, 2016)           │
│ Soft probability targets     ││ Greedy/Beam candidate paths  │
│ L_KD = T^2 * D_KL(P_T || P_S)││ Student trained on (x, y_hat)│
└──────────────┬───────────────┘└──────────────┬───────────────┘
               └───────────────┬───────────────┘
                               ▼
┌────────────────────────────────────────────────────────┐
│ 1B Compact Student Specialist (Llama 3.2 1B)           │
│ - Fits <= 2.0 GB RAM envelope                          │
│ - p99 TTFT < 50 ms on 4 CPU threads                    │
│ - Replicates teacher routing judgment                  │
└────────────────────────────────────────────────────────┘

Distillation is not merely “fine-tuning on model outputs.” Engineers must navigate three consequential decisions:

1. Access Constraints: Do you have access to the teacher’s internal vocabulary logits (white-box), or can you only query text completions via an API (black-box)?
2. Objective Formulation: Should the loss function enforce Forward KL (mode-covering) or Reverse KL (mode-seeking)? Choosing the wrong direction causes the student to hallucinate intermediate nonsense between valid reasoning modes.
3. Dark Knowledge Calibration: How much temperature scaling is needed to reveal the teacher’s secondary probabilities without amplifying background noise?

---

Section 2: White-Box Logit Distillation and the T-Squared Scale Factor

In white-box distillation (Hinton et al., 2015), the student learns not just the hard binary label (which class is correct), but the soft probability distribution across the entire vocabulary. This distribution contains what Geoffrey Hinton termed “dark knowledge”: the relative probabilities assigned to incorrect or alternative choices.

For example, when classifying an incident involving an unresponsive cache:

• Hard label: [1.0, 0.0, 0.0, 0.0] (Redis Failover).
• Standard teacher softmax (T = 1.0): [0.72, 0.18, 0.07, 0.03].

The hard label tells the model that expanding database connection pools is just as wrong as triggering an executive escalation. The teacher’s soft distribution reveals that expanding the database pool (0.18) is actually a viable secondary mitigation, while executive escalation (0.03) is an irrelevant emergency trigger.

1. Temperature-Scaled Softmax

To expose dark knowledge in the long tail of the distribution, logits are softened using a temperature parameter T > 1:

p_i = exp( v_i / T ) / sum_j exp( v_j / T )   [Teacher soft distribution]
q_i = exp( z_i / T ) / sum_j exp( z_j / T )   [Student soft distribution]

Where v represents teacher logits and z represents student logits.

• When T = 1.0, probabilities are sharp and dominated by the top mode.
• When T = 2.5, probabilities soften, amplifying secondary semantic relationships.
• When T -> infinity, probabilities approach a uniform distribution where distillation mathematically simplifies to minimizing the mean squared error between zero-centered logits.

2. The T-Squared Gradient Compensation Factor

A critical mathematical property identified in Hinton et al. (2015) is that softening logits by temperature T attenuates the magnitude of the backpropagated gradients by a factor of 1 / T^2.

To understand why, examine the gradient of the distillation loss with respect to the student logit z_i:

d(L_KD) / d(z_i) = (1 / T) * (q_i - p_i)

At high temperatures, the Taylor expansion of the exponential function exp(x) ~= 1 + x yields:

p_i ~= (1 / V) * (1 + (v_i - v_mean) / T)
q_i ~= (1 / V) * (1 + (z_i - z_mean) / T)

Substituting these approximations into the gradient equation gives:

d(L_KD) / d(z_i) ~= (1 / (V * T^2)) * (z_i - v_i)

Notice that the gradient is scaled by 1 / T^2! If you combine soft distillation loss with standard hard cross-entropy without scaling, the distillation gradients will shrink dramatically as you increase T.

To keep the relative weighting between the distillation loss and the hard-label cross-entropy loss stable as you tune T, you must multiply the distillation loss by T^2:

L_total = alpha * (T^2 * L_KD) + (1 - alpha) * L_CE

Where:

• L_KD = D_KL( p || q ) = sum_i p_i * log( p_i / q_i )
• L_CE = -sum_i y_i * log( q_i ) (evaluated at standard T = 1.0)
• alpha balances distillation from the teacher against ground-truth supervision (typically alpha = 0.5 or alpha = 0.7).

The Missing T-Squared Multiplier

Omitting the T^2 scale factor is one of the most common implementation bugs in custom training loops. If you sweep temperature from T = 1.0 to T = 4.0 without multiplying by T^2, the distillation gradient norm shrinks by 16x. The student silently stops learning from the teacher’s soft logits, appearing as though high temperatures degrade model quality when in reality the optimizer received virtually zero gradient signal.

---

Section 3: Forward KL vs. Reverse KL: Mode-Covering vs. Mode-Seeking

When formulating the divergence objective between teacher P and student Q, the direction of the Kullback-Leibler divergence completely alters how the student compresses knowledge.

[Forward KL vs Reverse KL Divergence Geometry]

1. Forward KL: D_KL(P || Q) [Mode-Covering / Zero-Avoiding]
   Teacher P:   /\         /\        (Two distinct valid modes)
               /  \       /  \
   Student Q: /    \_____/    \      (Hedges across both modes)
   Penalty: Explodes if P > 0 and Q == 0.
   Failure: Averages conflicting reasoning paths -> generates hallucinated text.

2. Reverse KL: D_KL(Q || P) [Mode-Seeking / Zero-Forcing]
   Teacher P:   /\         /\        (Two distinct valid modes)
               /  \       /  \
   Student Q:   /\                   (Locks into ONE mode cleanly)
               /  \
   Penalty: Explodes if Q > 0 and P == 0.
   Benefit: Generates high-confidence, valid paths without hedging.

1. Forward KL: D_KL(P_teacher || P_student) (Mode-Covering)

Forward KL evaluates the expectation under the teacher’s distribution:

D_KL(P || Q) = E_{y ~ P}[ log( P(y) / Q(y) ) ]

• Behavior: Mode-covering (zero-avoiding). If the teacher assigns non-zero probability to a token or sequence (P(y) > 0), the student must assign non-zero probability (Q(y) > 0), otherwise the ratio P(y) / Q(y) explodes to infinity.
• Problem in Autoregressive Generative Models: A 70B teacher distribution has thousands of diverse, valid continuations (multiple modes). A 1B student lacks the parameter capacity to cover all modes simultaneously. Under Forward KL, the student attempts to stretch its distribution over all modes, assigning probability mass to “void regions” (unlikely sequences between valid modes). In production, this causes the student to hedge, stutter, and generate syntactically plausible hallucinations (MiniLLM, Gu et al., 2024). However, Reverse KL does not guarantee JSON/schema compliance; schema compliance requires grammar-constrained decoding or structural validation at inference time.

2. Reverse KL: D_KL(P_student || P_teacher) (Mode-Seeking)

Reverse KL evaluates the expectation under the student’s own distribution:

D_KL(Q || P) = E_{y ~ Q}[ log( Q(y) / P(y) ) ]

• Behavior: Mode-seeking (zero-forcing). If the teacher assigns zero probability to an action (P(y) = 0), the student is heavily penalized for generating it (Q(y) > 0). However, if the teacher has two valid modes, the student is free to concentrate 100% of its capacity on one of those modes while completely ignoring the other.
• Application: For generative tasks (incident summaries, code synthesis, structured JSON generation), Reverse KL (implemented via policy gradient objectives in frameworks like MiniLLM) significantly reduces hallucination and exposure bias. The student learns to generate outputs that are always valid according to the teacher, even if it does not know every alternative phrasing the teacher knew.

Practical Engineering Objective Selection Matrix

Deployment Task	Recommended Objective	Primary Rationale
Fixed-Category Triage / Classification	Forward KL (Standard KD)	Discrete, small label space; student needs full coverage of all mitigation classes.
Open-Ended Text Generation	Reverse KL (MiniLLM)	Mode-seeking prevents averaging divergent reasoning paths; substantially reduces hallucinations.
Structured Output / Schema Extraction	Reverse KL + Format MFT	Zero-forcing guarantees strict compliance with JSON keys and enum boundaries.
API Black-Box (No Logit Access)	Sequence Distillation (Seq-KD)	Train on teacher beam-search completions using standard autoregressive cross-entropy.

---

Section 4: Black-Box Sequence-Level Distillation (Seq-KD)

When distilling from proprietary APIs or third-party models where logit vectors are inaccessible, engineers use Sequence-Level Distillation (Seq-KD) (Kim & Rush, 2016).

Rather than aligning per-token probability vectors, the pipeline treats the teacher as an expert data synthesizer:

Step 1: Dataset Generation
Query Prompts X ───► [70B Teacher Model (Beam Search / Nucleus Sampling)]
                               │
                               ▼
               Synthetic Expert Completions Y_hat

Step 2: Student Training
Training Pairs (X, Y_hat) ───► [1B Student Model (Standard Cross-Entropy Loss)]

The Perplexity-BLEU Paradox

Kim & Rush (2016) discovered a fundamental paradox in sequence distillation:

• A student trained on teacher beam-search completions (Seq-KD) exhibits higher (worse) perplexity on ground-truth validation data than a student trained on raw human text.
• Yet, when evaluated on task metrics (BLEU, routing accuracy, human preference), the Seq-KD model significantly outperforms the human-trained baseline.

Why does this happen? Human training text contains idiosyncratic noise, typos, and divergent writing styles. Ground-truth cross-entropy forces the student to distribute probability mass across all human idiosyncrasies. When trained on teacher beam-search sequences, the training distribution is unimodal and deterministic. The student’s probability distribution “collapses” onto the teacher’s most probable mode. For operational specialists (e.g. SRE triage), mode collapse onto high-accuracy predictions is not a defect—it is the desired engineering outcome.

---

Section 5: Worked Example: Incident Triage Distillation Calculation

Let us walk through the exact calculations for our SRE incident triage classification task.

Operational Setup

• Classes (4): 0: Redis Failover, 1: DB Pool Expand, 2: Kafka Rebalance, 3: Safety Escalation.
• Raw Teacher Logits: [4.0, 3.8, -1.0, -2.0].
• Initial Student Logits: [2.0, 1.0, 0.5, -0.5].

Run the automated verification suite to verify the temperature softening and Hinton loss calculations:

# Save as scripts/verify-distillation-math.py and run via python3
import math

def softmax(logits, temp=1.0):
    scaled = [l / temp for l in logits]
    max_s = max(scaled)
    exps = [math.exp(s - max_s) for s in scaled]
    sum_e = sum(exps)
    return [e / sum_e for e in exps]

def kl_divergence(p, q):
    """D_KL(P || Q) = sum P(i) * log(P(i) / Q(i))"""
    eps = 1e-12
    return sum(p_i * math.log((p_i + eps) / (q_i + eps)) for p_i, q_i in zip(p, q) if p_i > 0)

# 1. Evaluate probabilities at T=1.0 vs T=2.5
teacher_logits = [4.0, 3.8, -1.0, -2.0]
student_logits = [2.0, 1.0, 0.5, -0.5]

p_t10 = softmax(teacher_logits, temp=1.0)
p_t25 = softmax(teacher_logits, temp=2.5)

print("Teacher Probabilities (T=1.0):", [round(x, 3) for x in p_t10])
print("Teacher Soft Probabilities (T=2.5):", [round(x, 3) for x in p_t25])

# 2. Compute Hinton Soft Loss at T=2.0
T = 2.0
p_soft = softmax(teacher_logits, temp=T)
q_soft = softmax(student_logits, temp=T)
raw_kl = kl_divergence(p_soft, q_soft)
hinton_loss = (T ** 2) * raw_kl

# 3. Check Gradient Norms: unscaled vs T^2 scaled
grad_unscaled = [(1.0 / T) * (q_i - p_i) for p_i, q_i in zip(p_soft, q_soft)]
grad_scaled = [T * (q_i - p_i) for p_i, q_i in zip(p_soft, q_soft)]
norm_unscaled = math.sqrt(sum(g ** 2 for g in grad_unscaled))
norm_scaled = math.sqrt(sum(g ** 2 for g in grad_scaled))

print(f"Hinton Loss at T={T}: {hinton_loss:.4f} (Raw KL: {raw_kl:.4f})")
print(f"Gradient Norms: Unscaled = {norm_unscaled:.4f}, Scaled (xT^2) = {norm_scaled:.4f}")

Execution Output and Mathematical Audit

Teacher Probabilities (T=1.0): [0.547, 0.448, 0.004, 0.001]
Teacher Soft Probabilities (T=2.5): [0.465, 0.430, 0.063, 0.042]
Hinton Loss at T=2.0: 0.8560 (Raw KL: 0.2140)
Gradient Norms: Unscaled = 0.1367, Scaled (xT^2) = 0.5466

Engineering Interpretation

1. Revealing Dark Knowledge: At T = 1.0, Kafka Rebalance receives only 0.4% probability. At T = 2.5, it expands to 6.3%—a 15.7x amplification. The student observes that Kafka rebalancing has non-trivial correlation with connection pool exhaustion, allowing it to generalize when network partitions occur.
2. Gradient Stability: At T = 2.0, raw KL divergence between softened teacher and student distributions is 0.2140. Multiplying by T^2 = 4.0 yields a normalized Hinton distillation loss of 0.8560, ensuring that the backpropagated gradients remain in the same dynamic range as standard cross-entropy loss.

---

Section 6: Detecting and Diagnosing Transfer Failures

Distillation does not always succeed on the first training run. Staff engineers monitor five diagnostic metrics to isolate four distinct transfer failure modes:

[Diagnostic Decision Tree for Distillation Failure]
                         Validation Loss High / Accuracy Regressed
                                             │
               ┌─────────────────────────────┴─────────────────────────────┐
               ▼                                                           ▼
       Spearman rho < 0.70                                         Spearman rho >= 0.85
     (Probability Rank Broken)                                   (Probability Rank Preserved)
               │                                                           │
       ┌───────┴───────┐                                           ┌───────┴───────┐
       ▼               ▼                                           ▼               ▼
[Capacity Gap]  [Miscalibration]                            [Format Drift]  [Grounding Loss]
Student 1B too  Temperature T                              Schema parser   Hallucinates on
small for task  too high/low;                               fails; regex   unanswerables;
-> Expand to 3B sweep T in [1, 3]                           breaks -> Add  calibration ECE
                                                            format penalty  high -> Add Negatives

The 4 Transfer Failure Modes

1. The Capacity Gap

• Symptom: The student’s task accuracy plateaus far below the teacher, and the Spearman rank-order correlation of vocabulary probabilities between teacher and student remains below 0.70.
• Root Cause: The student architecture simply lacks the parameter capacity to represent the teacher’s latent decision boundary.
• Remediation: Do not continue wasting GPU hours on hyperparameter sweeps. Upscale the student architecture (e.g., from 1B to 3B parameters as evaluated in Lesson 13) or prune the label space.

2. Temperature Miscalibration

• Symptom: At T = 5.0, the student converges slowly and displays high entropy across predictions; at T = 1.0, it memorizes only the argmax class and fails on near-tie edge cases.
• Root Cause: When T is too low, dark knowledge is hidden; when T is too high, the distribution becomes uniform noise, drowning out signal.
• Remediation: Perform a controlled grid sweep across T in {1.0, 1.5, 2.0, 2.5, 3.0} while holding learning rate constant. Select the temperature that maximizes validation task F1.

3. Format Drift

• Symptom: The student correctly classifies incidents with 95% accuracy in text, but automated API callers crash because the output JSON omits required keys or adds markdown code fences.
• Root Cause: The teacher possessed thousands of hours of general instruction tuning that maintained rigid syntax; the student’s small parameter count prioritized semantic tokens over syntax delimiters.
• Remediation: Incorporate a strict structural loss penalty during distillation, or append 500 hand-crafted formatting invariance examples (as demonstrated in Lesson 15).

4. Grounding Loss on Unanswerable Queries

• Symptom: When presented with corrupt or ambiguous log fragments where the teacher responds "Insufficient telemetry to determine root cause", the distilled student hallucinates a confident diagnosis.
• Root Cause: Standard logit distillation forces the student to mimic teacher probabilities on valid samples, but the student loses the teacher’s epistemic uncertainty on Out-of-Distribution (OOD) data.
• Remediation: Explicitly inject negative and unanswerable samples into the distillation set with an elevated cross-entropy weight (alpha = 0.8 on refusal targets).

---

Key takeaways

• Dark Knowledge in Soft Logits: Soft logit distillation transfers rich secondary probability relationships (dark knowledge) that hard binary labels discard.
• The T-Squared Rule: Gradients of soft logit loss scale as 1/T^2; always multiply L_KD by T^2 when combining with standard cross-entropy loss to prevent gradient collapse.
• Forward KL Causes Generative Hedging: Forward KL (mode-covering) penalizes zero student probabilities, forcing a capacity-limited student to assign probability mass to invalid void regions between teacher modes.
• Reverse KL Guarantees Mode Purity: Reverse KL (mode-seeking) enables the student to commit cleanly to the teacher’s highest-confidence modes, eliminating hallucinations during open-ended generation.
• Seq-KD Mode Collapse is a Feature: Training on teacher beam-search completions produces higher perplexity but superior task BLEU and routing accuracy by stripping human idiosyncrasies.
• Four-Stage Diagnostic Triage: When distillation fails, evaluate rank correlation (Spearman rho) and schema compliance to isolate Capacity Gaps, Miscalibration, Format Drift, and Grounding Loss.

Try it yourself

1. 01An engineer trains a student using Hinton soft logit distillation at temperature T = 3.0 combined with hard-label cross entropy: L = 0.5 * L_KD + 0.5 * L_CE. They forget to multiply L_KD by T^2. By what factor are the distillation gradients reduced relative to the hard-label cross-entropy gradients?

   Show a hint

   Softmax gradients with respect to logits scale inversely with T^2. At T = 3.0, T^2 = 9. Omitting the T^2 multiplier reduces the distillation gradient magnitude by 9x (an 88.9% attenuation).

2. 02A team uses Forward KL divergence to distill a 70B teacher into a 1B student for an automated incident postmortem generator. The student frequently generates mixed, nonsensical sentences that blend two different outage causes. Why does Forward KL cause this behavior, and what objective resolves it?

   Show a hint

   Forward KL is mode-covering; it heavily penalizes the student for assigning zero probability where the teacher has probability mass. Because the 1B student cannot represent both modes, it places probability in the intermediate void, creating hallucinations. Switching to Reverse KL (mode-seeking) allows the student to lock cleanly into a single valid explanation.

3. 03During an SRE triage distillation run, teacher logits are [3.0, 1.0, -1.0]. Compute the softmax probability for class 0 at T = 1.0 versus T = 2.0. Does temperature scaling increase or decrease the relative probability of the top class?

   Show a hint

   At T=1.0: exp(3)=20.09, exp(1)=2.72, exp(-1)=0.37 -> sum=23.18 -> P(0)=20.09/23.18=0.867. At T=2.0: exp(1.5)=4.48, exp(0.5)=1.65, exp(-0.5)=0.61 -> sum=6.74 -> P(0)=4.48/6.74=0.665. Temperature scaling decreases top-class probability, transferring mass to secondary classes.

Interview questions

Why did Geoffrey Hinton introduce temperature scaling and the T^2 multiplier in knowledge distillation?

In well-trained neural networks, the argmax class often receives >99% probability, hiding the relative relationships between incorrect classes (dark knowledge). Applying temperature T > 1 softens the distribution, revealing secondary semantic correlations. However, softening divides the logit gradients by T^2. Multiplying the distillation loss by T^2 restores gradient magnitude, ensuring stable optimization when combining soft distillation loss with standard hard cross-entropy.

What is the fundamental mathematical difference between Forward KL and Reverse KL for distilling autoregressive LLMs?

Forward KL (D_KL(P_teacher || P_student)) is mode-covering: it penalizes the student infinitely if P > 0 but Q = 0, forcing a small student to spread probability mass across all teacher modes and causing hallucinations in void regions. Reverse KL (D_KL(P_student || P_teacher)) is mode-seeking: it penalizes the student if Q > 0 but P = 0, allowing the student to concentrate its limited capacity strictly on the teacher’s highest-probability modes without hedging.

How does Sequence-Level Distillation (Seq-KD) differ from standard token-level logit distillation, and when is it preferred?

Token-level distillation minimizes KL divergence between vocabulary logit vectors at every token step, requiring white-box access to identical teacher and student vocabularies. Sequence-level distillation (Kim & Rush, 2016) uses the teacher to generate full text completions via beam search or greedy decoding, then trains the student on these sequences using standard cross-entropy. Seq-KD is preferred when distilling across different tokenizers, when accessing closed-source APIs without logit endpoints, or when seeking mode collapse for deterministic generation.

How can an MLOps team detect whether a distilled model failure is due to a Capacity Gap versus Temperature Miscalibration?

Evaluate the Spearman rank-order correlation of token probabilities between teacher and student. If the rank correlation is low (rho < 0.70) across all temperatures, the student architecture lacks the parameter capacity to represent the teacher’s latent representations (Capacity Gap). If rank correlation is moderate (0.70 ≤ rho < 0.85), investigate temperature miscalibration (too low or too high). If rank correlation is high (rho ≥ 0.85) but validation task accuracy is poor, the temperature is likely too low (ignoring dark knowledge) or too high (fitting uniform noise), which can be resolved via a temperature sweep between 1.0 and 3.0. **Note:** These correlation thresholds are illustrative heuristics for debugging; teams should validate thresholds for their specific domain and task.

Related lessons

• Fine-Tuning & Domain-Specific Models · Prove a Specialist Is Worth ShippingChoose a Compact Student Model13 min
• Fine-Tuning & Domain-Specific Models · Make Training Data TrustworthyFormat Conversations and Mask the Loss13 min
• Fine-Tuning & Domain-Specific Models · Adapt Within a Compute BudgetDebug and Reproduce a Training Run14 min
• Fine-Tuning & Domain-Specific Models · Adapt Within a Compute BudgetFull Fine-Tuning Versus Adapters14 min