Distillation Variants › Code Implementation
Michael BrenndoerferFebruary 26, 202615 min read
Language AI HandbookAll chapters Part 41 of 67Model Compression
Distillation Variants › Code ImplementationOne section · Sections in this chapter6 / 9 ↓
Sections in this chapter9
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement each distillation variant with a small language model to see how the loss functions, training dynamics, and generation quality differ.
We build a compact transformer setup that makes the patterns clear without requiring large-scale compute.
In[3]:
Codeimport numpy as np
import torchToy Language Model Setup
We build a minimal transformer with configurable depth and width, giving us control over teacher vs. student capacity.
In[4]:
Codeimport torch
import torch.nn as nn
class MultiHeadSelfAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
assert d_model % num_heads == 0
self.num_heads = num_heads
self.d_k = d_model // num_heads
self.qkv = nn.Linear(d_model, 3 * d_model)
self.proj = nn.Linear(d_model, d_model)
def forward(self, x):
B, T, C = x.shape
qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, self.d_k)
q, k, v = qkv.unbind(dim=2)
q, k, v = [t.transpose(1, 2) for t in (q, k, v)]
scale = self.d_k**-0.5
attn = torch.softmax(q @ k.transpose(-2, -1) * scale, dim=-1)
out = (attn @ v).transpose(1, 2).reshape(B, T, C)
return self.proj(out), attn
class TransformerBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff):
super().__init__()
self.attn = MultiHeadSelfAttention(d_model, num_heads)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model),
)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x):
attn_out, attn_weights = self.attn(self.norm1(x))
x = x + attn_out
x = x + self.ff(self.norm2(x))
return x, attn_weights
class TinyTransformer(nn.Module):
def __init__(
self, vocab_size, d_model, num_heads, num_layers, d_ff, seq_len
):
super().__init__()
self.embed = nn.Embedding(vocab_size, d_model)
self.pos_embed = nn.Embedding(seq_len, d_model)
self.blocks = nn.ModuleList(
[
TransformerBlock(d_model, num_heads, d_ff)
for _ in range(num_layers)
]
)
self.norm = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, vocab_size, bias=False)
self.seq_len = seq_len
def forward(self, idx):
B, T = idx.shape
pos = torch.arange(T, device=idx.device).unsqueeze(0)
x = self.embed(idx) + self.pos_embed(pos)
attn_maps = []
hidden_states = [x]
for block in self.blocks:
x, attn = block(x)
attn_maps.append(attn)
hidden_states.append(x)
x = self.norm(x)
logits = self.head(x)
return logits, hidden_states, attn_mapsIn[5]:
Code# Build teacher and student models
VOCAB_SIZE = 64
SEQ_LEN = 16
teacher = TinyTransformer(
vocab_size=VOCAB_SIZE,
d_model=128,
num_heads=4,
num_layers=4,
d_ff=512,
seq_len=SEQ_LEN,
)
student_feat = TinyTransformer(
vocab_size=VOCAB_SIZE,
d_model=64,
num_heads=2,
num_layers=2,
d_ff=256,
seq_len=SEQ_LEN,
)
student_attn = TinyTransformer(
vocab_size=VOCAB_SIZE,
d_model=64,
num_heads=2,
num_layers=2,
d_ff=256,
seq_len=SEQ_LEN,
)
def count_params(model):
return sum(p.numel() for p in model.parameters())
teacher_params = count_params(teacher)
student_params = count_params(student_feat)
compression_ratio = teacher_params / student_paramsOut[6]:
ConsoleTeacher parameters: 811,776 Student parameters: 109,312 Compression ratio: 7.4x
The teacher is roughly 4 times larger than each student, creating a meaningful compression challenge. We then pretrain the teacher on a synthetic next-token prediction task so it develops real, learnable representations for the students to distill.
In[7]:
Codeimport torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
# Generate synthetic token sequences for training
num_samples = 1000
data = torch.randint(0, VOCAB_SIZE, (num_samples, SEQ_LEN))
targets = torch.roll(data, -1, dims=1) # Next-token targets
targets[:, -1] = data[:, 0] # Wrap-around for completeness
dataset = TensorDataset(data, targets)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
# Pretrain the teacher to give it meaningful representations
teacher_optimizer = torch.optim.Adam(teacher.parameters(), lr=1e-3)
teacher.train()
for epoch in range(3):
total_loss = 0.0
for x, y in loader:
logits, _, _ = teacher(x)
loss = F.cross_entropy(logits.reshape(-1, VOCAB_SIZE), y.reshape(-1))
teacher_optimizer.zero_grad()
loss.backward()
teacher_optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(loader)
teacher_final_loss = avg_lossOut[8]:
ConsoleTeacher pretraining complete. Teacher final training loss: 0.0605
Feature Distillation Loss
The feature alignment loss requires a projection from the student's lower-dimensional hidden states to the teacher's dimensionality. Without this projection, we cannot directly compare hidden states of size 64 (student) to hidden states of size 128 (teacher).
In[9]:
Codeclass FeatureProjection(nn.Module):
"""Projects student hidden states to teacher dimensionality."""
def __init__(self, d_student, d_teacher):
super().__init__()
self.proj = nn.Linear(d_student, d_teacher, bias=False)
def forward(self, h_student):
return self.proj(h_student)
def feature_distillation_loss(h_teacher, h_student, projector, normalize=True):
"""
Compute L2 feature alignment loss between teacher and student hidden states.
h_teacher: (B, T, d_teacher) - teacher activations (detached)
h_student: (B, T, d_student) - student activations
"""
h_student_proj = projector(h_student)
if normalize:
h_t = F.normalize(h_teacher.detach(), dim=-1)
h_s = F.normalize(h_student_proj, dim=-1)
else:
h_t = h_teacher.detach()
h_s = h_student_proj
return F.mse_loss(h_s, h_t)
# Feature projectors: student (64-dim) -> teacher (128-dim)
projector_feat = FeatureProjection(d_student=64, d_teacher=128)
# Training loop with feature distillation
student_feat_optimizer = torch.optim.Adam(
list(student_feat.parameters()) + list(projector_feat.parameters()), lr=1e-3
)
teacher.eval()
feat_losses = []
FEAT_WEIGHT = 0.5 # Weight for feature loss
KD_TEMP = 4.0
for epoch in range(5):
total_loss = 0.0
student_feat.train()
for x, y in loader:
with torch.no_grad():
t_logits, t_hidden, _ = teacher(x)
s_logits, s_hidden, _ = student_feat(x)
# Cross-entropy loss on ground truth
ce_loss = F.cross_entropy(
s_logits.reshape(-1, VOCAB_SIZE), y.reshape(-1)
)
# Soft-label KD loss
t_soft = F.log_softmax(t_logits / KD_TEMP, dim=-1)
s_soft = F.log_softmax(s_logits / KD_TEMP, dim=-1)
kd_loss = F.kl_div(s_soft, t_soft.exp(), reduction="batchmean") * (
KD_TEMP**2
)
# Feature distillation: align last transformer block hidden states
feat_loss = feature_distillation_loss(
t_hidden[-1], s_hidden[-1], projector_feat
)
loss = 0.5 * ce_loss + 0.3 * kd_loss + FEAT_WEIGHT * feat_loss
student_feat_optimizer.zero_grad()
loss.backward()
student_feat_optimizer.step()
total_loss += ce_loss.item()
feat_losses.append(total_loss / len(loader))
feat_final_loss = feat_losses[-1]Out[10]:
ConsoleFeature distillation training complete. Student (feature distill) final CE loss: 0.4165
The student's CE loss after feature distillation reflects how well it has learned to predict next tokens while simultaneously aligning its hidden-state geometry to the teacher's. Notice that we pass gradients through both the student parameters and the projector parameters: the student learns representations that, when projected, match the teacher; the projector learns the mapping that best bridges the two geometric spaces.
Attention Transfer Loss
Now we implement attention transfer. The loss averages MSE across all student layers, mapping each student layer to a proportional teacher layer index.
In[11]:
Codedef attention_transfer_loss(attn_teacher, attn_student):
"""
Compute MSE loss between teacher and student attention distributions.
attn_teacher: list of (B, H, T, T) tensors, one per teacher layer
attn_student: list of (B, H, T, T) tensors, one per student layer
We match student layer i to the corresponding teacher layer
by striding through teacher layers.
"""
n_teacher = len(attn_teacher)
n_student = len(attn_student)
total_loss = 0.0
for s_idx, a_s in enumerate(attn_student):
# Map student layer index to teacher layer index proportionally
t_idx = int(round(s_idx * (n_teacher - 1) / max(n_student - 1, 1)))
a_t = attn_teacher[t_idx]
# Teacher may have more heads; average both to a single head
a_t_avg = a_t.mean(dim=1) # (B, T, T)
a_s_avg = a_s.mean(dim=1) # (B, T, T)
total_loss = total_loss + F.mse_loss(a_s_avg, a_t_avg.detach())
return total_loss / n_student
# Training loop with attention transfer
student_attn_optimizer = torch.optim.Adam(student_attn.parameters(), lr=1e-3)
attn_losses = []
ATTN_WEIGHT = 0.5
for epoch in range(5):
total_loss = 0.0
student_attn.train()
for x, y in loader:
with torch.no_grad():
t_logits, _, t_attn = teacher(x)
s_logits, _, s_attn = student_attn(x)
ce_loss = F.cross_entropy(
s_logits.reshape(-1, VOCAB_SIZE), y.reshape(-1)
)
t_soft = F.log_softmax(t_logits / KD_TEMP, dim=-1)
s_soft = F.log_softmax(s_logits / KD_TEMP, dim=-1)
kd_loss = F.kl_div(s_soft, t_soft.exp(), reduction="batchmean") * (
KD_TEMP**2
)
attn_loss = attention_transfer_loss(t_attn, s_attn)
loss = 0.5 * ce_loss + 0.3 * kd_loss + ATTN_WEIGHT * attn_loss
student_attn_optimizer.zero_grad()
loss.backward()
student_attn_optimizer.step()
total_loss += ce_loss.item()
attn_losses.append(total_loss / len(loader))
attn_final_loss = attn_losses[-1]Out[12]:
ConsoleAttention transfer training complete. Student (attention transfer) final CE loss: 0.3820
On-Policy Distillation
On-policy distillation requires generating from the student during training. This is the most structurally different variant: unlike the previous two, which add extra loss terms computed on static batches, on-policy distillation changes the data itself.
In[13]:
Codestudent_onpolicy = TinyTransformer(
vocab_size=VOCAB_SIZE,
d_model=64,
num_heads=2,
num_layers=2,
d_ff=256,
seq_len=SEQ_LEN,
)
student_onpolicy_optimizer = torch.optim.Adam(
student_onpolicy.parameters(), lr=5e-4
)
def generate_student_sequences(model, batch_size, seq_len, vocab_size):
"""Sample sequences from the student's autoregressive policy."""
model.eval()
with torch.no_grad():
# Start with random first tokens
tokens = torch.randint(0, vocab_size, (batch_size, 1))
for _ in range(seq_len - 1):
logits, _, _ = model(tokens)
next_token_logits = logits[:, -1, :]
probs = torch.softmax(next_token_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
tokens = torch.cat([tokens, next_token], dim=1)
return tokens
onpolicy_losses = []
ONPOLICY_GEN_BATCH = 32
for epoch in range(5):
total_loss = 0.0
count = 0
student_onpolicy.train()
# On-policy: generate from student, then distill teacher onto those inputs
for _ in range(len(loader)):
# Step 1: Generate sequences from current student policy
gen_seqs = generate_student_sequences(
student_onpolicy, ONPOLICY_GEN_BATCH, SEQ_LEN, VOCAB_SIZE
)
# Step 2: Get teacher distributions on student-generated text
student_onpolicy.train()
with torch.no_grad():
t_logits, _, _ = teacher(gen_seqs)
# Step 3: Get student distributions on the same text
s_logits, _, _ = student_onpolicy(gen_seqs)
# Step 4: KL divergence from student to teacher (forward KL)
t_probs = F.softmax(t_logits / KD_TEMP, dim=-1)
s_log_probs = F.log_softmax(s_logits / KD_TEMP, dim=-1)
kl_loss = F.kl_div(s_log_probs, t_probs, reduction="batchmean") * (
KD_TEMP**2
)
student_onpolicy_optimizer.zero_grad()
kl_loss.backward()
student_onpolicy_optimizer.step()
total_loss += kl_loss.item()
count += 1
onpolicy_losses.append(total_loss / count)
onpolicy_final_loss = onpolicy_losses[-1]Out[14]:
ConsoleOn-policy distillation training complete. Student (on-policy) final KL loss: 35.0270
Note that the on-policy student's loss is a KL divergence rather than a cross-entropy, which puts it on a different scale than the other variants. The raw loss values therefore should not be ranked against one another. In this short toy run, the on-policy KL decreases more modestly than the two cross-entropy curves, reflecting the harder objective of matching a teacher across the student's changing generated-prefix distribution.
Comparing Training Dynamics
Out[15]:
Visualization
Visualizing Feature Alignment
One way to assess feature distillation is to measure how well the student's projected hidden states match the teacher's after training. A high cosine similarity indicates the two models have developed directionally similar internal representations.
In[16]:
Codedef compute_feature_cosine_similarities(
teacher_model, student_model, projector, data_loader
):
"""
Compute mean cosine similarity between teacher and projected student
hidden states across all batch elements.
"""
sims = []
teacher_model.eval()
student_model.eval()
projector.eval()
with torch.no_grad():
for x, _ in data_loader:
t_logits, t_hidden, _ = teacher_model(x)
s_logits, s_hidden, _ = student_model(x)
h_t = F.normalize(t_hidden[-1], dim=-1) # (B, T, d_teacher)
h_s = F.normalize(
projector(s_hidden[-1]), dim=-1
) # (B, T, d_teacher)
# Cosine similarity per position then average over batch and seq
sim = (h_t * h_s).sum(dim=-1) # (B, T)
sims.append(sim.mean().item())
return np.mean(sims)
cos_sim = compute_feature_cosine_similarities(
teacher, student_feat, projector_feat, loader
)Out[17]:
ConsoleMean cosine similarity (teacher vs projected student): 0.6611 A value near 1.0 indicates near-perfect directional alignment. A value near 0.0 indicates uncorrelated representations.
A high cosine similarity confirms that feature distillation successfully steered the student to learn directionally similar representations to the teacher. The projection layer absorbs the dimensionality difference, so the comparison is made in the teacher's geometric space. What this measures is not whether the student has the same information content as the teacher (it cannot; it has half the dimensions), but whether the information it does contain is organized in the same directional structure.
Attention Pattern Comparison
Comparing teacher and student attention heatmaps gives a visual sense of how well structural transfer succeeded. After attention transfer training, the student's attention maps should show similar concentration patterns to the teacher's, even though the student has fewer heads and layers.
Out[18]:
Visualization
Progressive Distillation: Multi-Stage Compression
The progressive distillation code shows the key structural difference from other methods: we train two models sequentially, where the second uses the first as its teacher.
In[19]:
Codedef progressive_distill_step(
teacher_model, student_model, loader, epochs=3, lr=1e-3
):
"""
Single stage of progressive distillation: student learns from current teacher.
Returns per-epoch CE losses.
"""
optimizer = torch.optim.Adam(student_model.parameters(), lr=lr)
teacher_model.eval()
stage_losses = []
for epoch in range(epochs):
total_loss = 0.0
student_model.train()
for x, y in loader:
with torch.no_grad():
t_logits, _, _ = teacher_model(x)
s_logits, _, _ = student_model(x)
ce_loss = F.cross_entropy(
s_logits.reshape(-1, VOCAB_SIZE), y.reshape(-1)
)
t_soft = F.log_softmax(t_logits / KD_TEMP, dim=-1)
s_soft = F.log_softmax(s_logits / KD_TEMP, dim=-1)
kd_loss = F.kl_div(s_soft, t_soft.exp(), reduction="batchmean") * (
KD_TEMP**2
)
loss = 0.5 * ce_loss + 0.5 * kd_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += ce_loss.item()
stage_losses.append(total_loss / len(loader))
return stage_losses
# Stage 1: Teacher (4 layers, 128-dim) -> Intermediate (3 layers, 96-dim)
intermediate = TinyTransformer(
vocab_size=VOCAB_SIZE,
d_model=96,
num_heads=4,
num_layers=3,
d_ff=384,
seq_len=SEQ_LEN,
)
stage1_losses = progressive_distill_step(
teacher, intermediate, loader, epochs=3
)
# Stage 2: Intermediate -> Final Small (2 layers, 64-dim)
final_progressive = TinyTransformer(
vocab_size=VOCAB_SIZE,
d_model=64,
num_heads=2,
num_layers=2,
d_ff=256,
seq_len=SEQ_LEN,
)
stage2_losses = progressive_distill_step(
intermediate, final_progressive, loader, epochs=3
)
# Single-step baseline: direct from teacher to same-size student
single_step = TinyTransformer(
vocab_size=VOCAB_SIZE,
d_model=64,
num_heads=2,
num_layers=2,
d_ff=256,
seq_len=SEQ_LEN,
)
single_step_losses = progressive_distill_step(
teacher, single_step, loader, epochs=6
)
all_prog_losses = stage1_losses + stage2_lossesOut[20]:
ConsoleProgressive distillation stages: Stage 1 final CE loss (teacher -> intermediate): 1.4941 Stage 2 final CE loss (intermediate -> student): 4.0537 Single-step baseline final CE loss: 0.1754
Out[21]:
Visualization
Comparing All Variants
Out[22]:
Visualization