Distillation

◆AI◆Advanced

Introduction

Knowledge distillation transfers the capabilities of a large, powerful teacher model to a smaller, more efficient student model. The student is trained to mimic the teacher's outputs, logits, or internal representations, achieving performance close to the teacher at a fraction of the computational cost.

Distillation is one of the most practical techniques in LLM deployment. It directly addresses the cost-quality trade-off: a distilled 7B model can match or exceed the performance of a 70B model on specific tasks while running 10-50x faster and using 90% less memory.

This guide covers the major distillation paradigms: logit-based (classic Hinton distillation), feature-based (distilling hidden representations), on-policy (student generates, teacher evaluates), dataset distillation (using teacher to generate training data), and quantization-aware distillation.

Teacher-Student Framework

The teacher-student framework is the foundation of knowledge distillation. A pre-trained teacher model (typically large: 70B+, 400B+) generates soft targets or feature representations that guide the training of a smaller student model. The student learns to match the teacher's output distribution, not just the hard labels from the training data.

Method	Knowledge Source	Student Objective	Cost
Logit-based	Teacher logits (softmax)	KL divergence	Low
Feature-based	Hidden states	MSE on features	Medium
On-policy	Student samples + teacher eval	Preference loss	High
Dataset	Teacher-generated data	Standard SFT	Low (one-time gen)

ℹ

info

Start with dataset distillation (simplest): generate training data from the teacher and train the student with standard SFT. If the student still underperforms, add logit-based distillation during training. Feature-based and on-policy methods add complexity for marginal gains in most scenarios.

Logit-Based Distillation

Logit-based distillation minimizes the KL divergence between the teacher's and student's output probability distributions. The key insight is that the teacher's soft labels contain richer information than hard labels — the relative probabilities of incorrect answers encode the teacher's understanding of task structure.

The temperature parameter controls the softness of the probability distribution. Higher temperatures produce softer distributions that reveal more fine-grained relationships between classes. During training, both teacher and student logits are divided by the same temperature before computing KL divergence.

logit-distill.py

Python

1	import torch
2	import torch.nn.functional as F
3	from transformers import AutoModelForCausalLM, AutoTokenizer
4
5	class DistillationTrainer:
6	def __init__(self, teacher_name: str, student_name: str, temperature: float = 2.0):
7	self.teacher = AutoModelForCausalLM.from_pretrained(teacher_name)
8	self.student = AutoModelForCausalLM.from_pretrained(student_name)
9	self.temperature = temperature
10	self.teacher.eval() # Freeze teacher
11
12	def distillation_loss(self, student_logits, teacher_logits):
13	# Apply temperature scaling
14	soft_teacher = F.log_softmax(
15	teacher_logits / self.temperature, dim=-1
16	)
17	soft_student = F.log_softmax(
18	student_logits / self.temperature, dim=-1
19	)
20
21	# KL divergence loss
22	kl_loss = F.kl_div(
23	soft_student, soft_teacher,
24	log_target=True,
25	reduction="batchmean"
26	)
27
28	# Scale by temperature squared
29	return (self.temperature ** 2) * kl_loss
30
31	def train_step(self, input_ids):
32	with torch.no_grad():
33	teacher_outputs = self.teacher(input_ids)
34	teacher_logits = teacher_outputs.logits
35
36	student_outputs = self.student(input_ids)
37	student_logits = student_outputs.logits
38
39	loss = self.distillation_loss(student_logits, teacher_logits)
40
41	loss.backward()
42	return loss.item()

Combined Distillation + SFT Loss

In practice, combine distillation loss with standard SFT cross-entropy loss on ground truth labels. The alpha parameter controls the trade-off between mimicking the teacher and learning from ground truth.

combined-distill-loss.py

Python

1	def combined_loss(
2	student_logits, teacher_logits,
3	labels, alpha=0.5, temperature=2.0
4	):
5	# Standard cross-entropy (SFT loss)
6	ce_loss = F.cross_entropy(
7	student_logits.view(-1, student_logits.size(-1)),
8	labels.view(-1),
9	ignore_index=-100
10	)
11
12	# Distillation loss (KL divergence)
13	soft_teacher = F.log_softmax(
14	teacher_logits / temperature, dim=-1
15	)
16	soft_student = F.log_softmax(
17	student_logits / temperature, dim=-1
18	)
19	kl_loss = F.kl_div(
20	soft_student, soft_teacher,
21	log_target=True,
22	reduction="batchmean"
23	)
24	kl_loss = (temperature ** 2) * kl_loss
25
26	# Weighted combination
27	return alpha * ce_loss + (1 - alpha) * kl_loss
28
29	# Training loop
30	optimizer = torch.optim.AdamW(student.parameters(), lr=1e-4)
31	for batch in dataloader:
32	logits = student(batch["input_ids"]).logits
33	with torch.no_grad():
34	t_logits = teacher(batch["input_ids"]).logits
35
36	loss = combined_loss(
37	logits, t_logits, batch["labels"],
38	alpha=0.3, temperature=3.0
39	)
40	loss.backward()
41	optimizer.step()

🔥

pro tip

The alpha parameter controls the distillation-accuracy trade-off. Start with alpha=0.5 for balanced learning. If the student overfits to the teacher's mistakes, increase alpha. If the student ignores the teacher's soft labels, decrease alpha. The temperature should be tuned between 2 and 8 — higher temperatures reveal more teacher knowledge but also amplify noise.

On-Policy Distillation

On-policy distillation (also called online distillation or self-distillation) trains the student on its own generated samples, with the teacher providing feedback. Unlike logit-based distillation (which is off-policy — the teacher guides on fixed data), on-policy distillation allows the student to explore regions of the output space it actually visits, leading to better generalization.

on-policy-distill.py

Python

1	import torch
2
3	def on_policy_distill_step(
4	student, teacher, tokenizer,
5	prompts, max_new_tokens=128, beta=0.5
6	):
7	# Student generates responses
8	student_outputs = student.generate(
9	prompts, max_new_tokens=max_new_tokens,
10	do_sample=True, temperature=0.7,
11	pad_token_id=tokenizer.eos_token_id
12	)
13
14	# Teacher scores the student's generations
15	with torch.no_grad():
16	teacher_scores = teacher(
17	student_outputs,
18	labels=student_outputs
19	).loss # Lower perplexity = better
20
21	# Student loss = negative log-likelihood of teacher-preferred tokens
22	student_logits = student(student_outputs).logits
23	student_loss = F.cross_entropy(
24	student_logits.view(-1, student_logits.size(-1)),
25	student_outputs.view(-1),
26	ignore_index=tokenizer.pad_token_id
27	)
28
29	# Weight by teacher preference score
30	weight = torch.exp(-beta * teacher_scores)
31	weighted_loss = weight * student_loss
32
33	return weighted_loss.mean()

Quantization as Distillation

Quantization-aware training (QAT) can be viewed as a form of self-distillation where the full-precision model is the teacher and the quantized model is the student. The student learns to compensate for information loss from quantization by minimizing the divergence between its outputs and those of the full-precision teacher.

qat-distill.py

Python

1	import torch.ao.quantization as quant
2
3	class QATDistillation:
4	def __init__(self, full_precision_model, quantized_model):
5	self.teacher = full_precision_model.eval()
6	self.student = quantized_model.train()
7	self.student.qconfig = quant.get_default_qat_qconfig("x86")
8
9	def qat_step(self, input_ids):
10	with torch.no_grad():
11	teacher_logits = self.teacher(input_ids).logits
12
13	student_logits = self.student(input_ids).logits
14	loss = F.kl_div(
15	F.log_softmax(student_logits, dim=-1),
16	F.log_softmax(teacher_logits, dim=-1),
17	log_target=True,
18	reduction="batchmean"
19	)
20
21	loss.backward()
22	return loss.item()
23
24	def convert(self):
25	# Convert QAT model to quantized
26	self.student.eval()
27	quantized_model = quant.convert(self.student)
28	return quantized_model

✓

best practice

Quantization-aware distillation typically recovers 1-3% of the accuracy lost during post-training quantization. It is most effective for lower bit-widths (4-bit, 3-bit) where the information loss is significant. For 8-bit quantization, simple post-training quantization usually suffices.

Practical Benefits

Distillation delivers concrete, measurable benefits in production deployments. The following table compares a distilled 7B model against its 70B teacher on common deployment metrics.

Metric	Teacher (70B)	Student (7B)	Improvement
Inference Latency	2.1s per token	0.12s per token	17x faster
GPU Memory	140 GB (4x A100)	14 GB (1x A100)	90% reduction
Cost per Token	$0.0035	$0.0002	17x cheaper
Serving Capacity	8 concurrent users	256 concurrent users	32x more

⚠

warning

Distillation narrows the student's capabilities to match the teacher on the distillation data distribution. The student may not generalize to out-of-distribution inputs as well as the teacher. Always evaluate the distilled student on your full evaluation suite, not just the distillation benchmark, before deploying.

$Blueprint — Engineering Documentation·Section ID: AI-DIST-01·Revision: 1.0

X GitHub YouTube Instagram Facebook Buy me a coffee Sponsor