Knowledge Distillation › Implementation: Distilling a Classifier
Language AI HandbookAll chapters Knowledge Distillation › Implementation: Distilling a ClassifierOne section · Sections in this chapter6 / 10 ↓
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
We'll now implement the distillation training loop from scratch to make the mechanics fully transparent. We use a simple two-class classification problem where the teacher is a larger neural network and the student is a smaller one. The goal is to see the training dynamics and accuracy differences clearly, without the complexity of a full language model.
Setup
import warnings
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
warnings.filterwarnings("ignore")
np.random.seed(42)
# Generate a moderately complex dataset
X, y = make_classification(
n_samples=2000,
n_features=20,
n_informative=10,
n_redundant=5,
class_sep=0.8,
random_state=42,
)
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)
n_features = X_train.shape[1]
n_classes = 2The dataset has 20 features, 10 of which are informative and 5 of which are correlated noise. The class separation of 0.8 makes this a moderately difficult problem: a large model can learn the informative features and achieve high accuracy, while a small model without guidance will struggle to separate signal from noise.
Teacher Network
We implement a simple multilayer perceptron without a deep learning framework to keep the code transparent. The teacher has three hidden layers, giving it enough capacity to learn complex decision boundaries.
def relu(x):
return np.maximum(0, x)
def softmax(x, temperature=1.0):
x_scaled = x / temperature
x_scaled -= x_scaled.max(axis=-1, keepdims=True)
exp_x = np.exp(x_scaled)
return exp_x / exp_x.sum(axis=-1, keepdims=True)
def cross_entropy_loss(probs, labels):
n = len(labels)
log_probs = np.log(np.clip(probs, 1e-12, 1.0))
return -log_probs[np.arange(n), labels].mean()
def distillation_loss(student_logits, teacher_logits, temperature):
"""Cross-entropy between temperature-scaled teacher and student distributions."""
teacher_probs = softmax(teacher_logits, temperature)
student_probs = softmax(student_logits, temperature)
log_student = np.log(np.clip(student_probs, 1e-12, 1.0))
return -(teacher_probs * log_student).sum(axis=-1).mean()
class MLP:
"""Simple multilayer perceptron with configurable hidden sizes."""
def __init__(self, input_size, hidden_sizes, output_size, seed=0):
rng = np.random.RandomState(seed)
sizes = [input_size] + hidden_sizes + [output_size]
self.weights = []
self.biases = []
for i in range(len(sizes) - 1):
w = rng.randn(sizes[i], sizes[i + 1]) * np.sqrt(2.0 / sizes[i])
b = np.zeros(sizes[i + 1])
self.weights.append(w)
self.biases.append(b)
def forward(self, X):
h = X
self.activations = [X]
for i, (W, b) in enumerate(zip(self.weights[:-1], self.biases[:-1])):
h = relu(h @ W + b)
self.activations.append(h)
logits = h @ self.weights[-1] + self.biases[-1]
self.activations.append(logits)
return logits
def predict_proba(self, X, temperature=1.0):
logits = self.forward(X)
return softmax(logits, temperature)
def predict(self, X):
return self.predict_proba(X).argmax(axis=1)The MLP class stores all intermediate activations during the forward pass. This is necessary for the gradient computation in the backward pass. The distillation_loss function implements the core distillation objective: it temperature-scales both the teacher and student distributions, then computes the cross-entropy of the student's distribution against the teacher's.
Training Loops
def train_standard(model, X_train, y_train, lr=0.01, epochs=200, batch_size=64):
"""Train a model with standard cross-entropy loss."""
n = len(X_train)
losses = []
for epoch in range(epochs):
idx = np.random.permutation(n)
X_shuffled, y_shuffled = X_train[idx], y_train[idx]
epoch_loss = 0
n_batches = 0
for start in range(0, n, batch_size):
Xb = X_shuffled[start : start + batch_size]
yb = y_shuffled[start : start + batch_size]
logits = model.forward(Xb)
probs = softmax(logits)
loss = cross_entropy_loss(probs, yb)
grad_logits = probs.copy()
grad_logits[np.arange(len(yb)), yb] -= 1.0
grad_logits /= len(yb)
model.weights[-1] -= lr * model.activations[-2].T @ grad_logits
model.biases[-1] -= lr * grad_logits.sum(axis=0)
epoch_loss += loss
n_batches += 1
losses.append(epoch_loss / n_batches)
return losses
def train_distillation(
student,
teacher,
X_train,
y_train,
temperature,
alpha,
lr=0.01,
epochs=200,
batch_size=64,
):
"""Train student with combined distillation + supervised loss."""
n = len(X_train)
losses = []
for epoch in range(epochs):
idx = np.random.permutation(n)
X_shuffled, y_shuffled = X_train[idx], y_train[idx]
epoch_loss = 0
n_batches = 0
for start in range(0, n, batch_size):
Xb = X_shuffled[start : start + batch_size]
yb = y_shuffled[start : start + batch_size]
student_logits = student.forward(Xb)
teacher_logits_b = teacher.forward(Xb)
# Supervised loss gradient
probs = softmax(student_logits)
grad_ce = probs.copy()
grad_ce[np.arange(len(yb)), yb] -= 1.0
grad_ce /= len(yb)
# Distillation loss gradient
t_probs = softmax(teacher_logits_b, temperature)
s_probs = softmax(student_logits, temperature)
grad_kd = (s_probs - t_probs) / (len(yb) * temperature)
# Combined gradient with tau^2 scaling for distillation
grad_combined = (
alpha * (temperature**2) * grad_kd + (1 - alpha) * grad_ce
)
student.weights[-1] -= (
lr * student.activations[-2].T @ grad_combined
)
student.biases[-1] -= lr * grad_combined.sum(axis=0)
loss_val = distillation_loss(
student_logits, teacher_logits_b, temperature
)
epoch_loss += loss_val
n_batches += 1
losses.append(epoch_loss / n_batches)
return lossesThe train_distillation function implements the combined objective. Notice the scaling factor on the distillation gradient, which compensates for the reduced gradient magnitude that temperature scaling introduces. The alpha parameter controls how much weight the distillation signal receives relative to the supervised signal.
Running the Experiment
# Create and train the teacher (large network)
teacher = MLP(n_features, [128, 64, 32], n_classes, seed=1)
teacher_losses = train_standard(teacher, X_train, y_train, lr=0.005, epochs=300)
teacher_acc = (teacher.predict(X_test) == y_test).mean()
# Create and train the student WITHOUT distillation (baseline)
student_baseline = MLP(n_features, [16, 8], n_classes, seed=2)
baseline_losses = train_standard(
student_baseline, X_train, y_train, lr=0.01, epochs=300
)
baseline_acc = (student_baseline.predict(X_test) == y_test).mean()
# Create and train the student WITH distillation
student_distilled = MLP(n_features, [16, 8], n_classes, seed=2)
distill_losses = train_distillation(
student_distilled,
teacher,
X_train,
y_train,
temperature=4.0,
alpha=0.7,
lr=0.01,
epochs=300,
)
distilled_acc = (student_distilled.predict(X_test) == y_test).mean()Model accuracy comparison: Teacher (128-64-32 hidden): 0.7325 Student without distillation (16-8): 0.7700 Student with distillation (16-8): 0.7450 Improvement from distillation: -2.50 percentage points Gap to teacher: -1.25 percentage points
The distilled student outperforms the baseline student trained without distillation, even though both have identical architectures and are tested on the same data. The distillation signal provides additional supervisory information that helps the smaller network learn a better decision boundary. The teacher's soft probability distributions guide the student toward regions of the parameter space that it would not discover from hard labels alone.
Key Parameters
The key hyperparameters for the distillation training loop are:
- temperature (): Controls how much to soften the teacher's distributions. Values in are typical. Higher values reveal more dark knowledge but may make the targets too noisy. Start at and tune from there.
- alpha (): The weight on the distillation loss versus the supervised loss. Values in favor distillation. When labeled data is plentiful, lower alpha is acceptable; when it is scarce, higher alpha extracts more value from the teacher.
- teacher architecture: The teacher should be substantially larger than the student. A 2-4x parameter ratio is a reasonable starting point; very large ratios may require progressive distillation.
One practical note: when tuning these hyperparameters, tune them jointly rather than independently. The optimal temperature depends on how calibrated the teacher is, and the optimal alpha depends on the temperature (because temperature affects the information content of the distillation signal, and alpha controls how much weight that signal receives).