Knowledge Distillation Temperature Sweep on MNIST

# Reproducing Section 3 of *Distilling the Knowledge in a Neural Network* Hinton, Vinyals & Dean (2015), §3 "Preliminary experiments on MNIST" reports three numbers, as **absolute test errors out of 10,000**: | model | reported test errors | |---|---| | teacher: 784–1200–1200–10, dropout + max-norm + ±2px jitter, all 60k cases | **67** | | student: 784–800–800–10, **no** regularization, hard labels | **146** | | same student, regularized *solely* by matching the teacher's soft targets at **T=20** | **74** | > "This net achieved 67 test errors whereas a smaller net with two hidden layers of 800 rectified linear hidden units and no regularization achieved 146 errors. But if the smaller net was regularized solely by adding the additional task of matching the soft targets produced by the large net at a temperature of 20, it achieved 74 test errors." ## Loss $$L = \alpha \cdot \mathrm{CE}(y,\, z_s) \;+\; (1-\alpha)\cdot T^2 \cdot \mathrm{KL}\!\left(\mathrm{softmax}(z_t/T)\;\|\;\mathrm{softmax}(z_s/T)\right)$$ The $T^2$ factor keeps soft-target gradient magnitudes comparable across temperatures. The paper's §3 student is the $\alpha=0$ case ("regularized **solely** by … soft targets"), and it sees the **un-jittered** transfer set — the point being that knowledge about translation invariance transfers through soft targets even though the transfer set contains no translations. ## Measurement discipline - Error **counts**, not accuracy %, so results are directly comparable to…

import os, math, json, time, random
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("cuda available:", torch.cuda.is_available())
if dev.type == "cuda":
    print("device:", torch.cuda.get_device_name(0))
else:
    print("WARNING: no GPU — training will be slow")
print("torch:", torch.__version__)

def set_seed(s):
    random.seed(s); np.random.seed(s)
    torch.manual_seed(s); torch.cuda.manual_seed_all(s)
…

cuda available: True
device: NVIDIA A100-SXM4-40GB
torch: 2.7.1+cu128
ready

try:
    from torchvision import datasets
    tr = datasets.MNIST("/home/user/data", train=True, download=True)
    te = datasets.MNIST("/home/user/data", train=False, download=True)
    Xtr = tr.data.clone(); ytr = tr.targets.clone()
    Xte = te.data.clone(); yte = te.targets.clone()
    src = "torchvision.datasets.MNIST"
except Exception as e:
    print("torchvision MNIST failed:", repr(e)[:300], "-> falling back to HF datasets")
    from datasets import load_dataset
    ds = load_dataset("mnist")
    Xtr = torch.tensor(np.stack([np.array(im) for im in ds["train"]["image"]]))
    ytr = torch.tensor(ds["train"]["label"])
    Xte = torch.tensor(np.stack([np.array(im) for im in ds["test"]["image"]]))
    yte = torch.tensor(ds["test"]["label"])
    src = "huggingface datasets: mnist"

# to GPU, scaled to [0,1], kept as 28x28 so we can jitter cheaply
…

source: torchvision.datasets.MNIST
train: (60000, 28, 28) test: (10000, 28, 28)
value range: 0.0 - 1.0
train label counts: [5923, 6742, 5958, 6131, 5842, 5421, 5918, 6265, 5851, 5949]
test  label counts: [980, 1135, 1032, 1010, 982, 892, 958, 1028, 974, 1009]
test-set size = 10000 -> error counts are directly comparable to 67/146/74

class MLP(nn.Module):
    """784 -> h -> h -> 10 ReLU net. p_in/p_h > 0 enables the teacher's dropout recipe."""
    def __init__(self, h=800, p_in=0.0, p_h=0.0):
        super().__init__()
        self.p_in, self.p_h = p_in, p_h
        self.fc1 = nn.Linear(784, h)
        self.fc2 = nn.Linear(h, h)
        self.fc3 = nn.Linear(h, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        if self.p_in: x = F.dropout(x, self.p_in, self.training)
        x = F.relu(self.fc1(x))
        if self.p_h: x = F.dropout(x, self.p_h, self.training)
        x = F.relu(self.fc2(x))
        if self.p_h: x = F.dropout(x, self.p_h, self.training)
        return self.fc3(x)

…

harness ready — MLP / max_norm_ / jitter / distill_loss / test_errors / train_net
jitter check: (4, 28, 28) mean before/after: 0.119 0.119

t0 = time.time()
# Teacher: 784-1200-1200-10 ReLU, dropout (0.2 input / 0.5 hidden), max-norm 3.5,
# +-2px jitter, all 60,000 training cases, long SGD-momentum schedule.
# A 250-epoch cosine run reached 85 errors and was still improving at anneal, so the
# schedule is extended to 800 epochs (the paper's teacher is also trained "for a long time").
final_t, best_t, hist_t, teacher = train_net(
    h=1200, seed=0, epochs=800, bs=128, lr=0.1, momentum=0.9,
    p_in=0.2, p_h=0.5, max_norm=3.5, use_jitter=True,
    log_every=50, tag="teacher-1200")

print(f"\nTEACHER 1200x1200  final={final_t}  best={best_t}  (paper: 67)")
print(f"trained in {(time.time()-t0)/60:.1f} min")
RESULTS["teacher_1200_dropout_jitter"] = {"errors": [final_t], "best": best_t, "paper": 67}

[teacher-1200 seed=0] epoch   1/800  test errors=439  best=439
  [teacher-1200 seed=0] epoch  50/800  test errors=170  best=163
  [teacher-1200 seed=0] epoch 100/800  test errors=159  best=141
  [teacher-1200 seed=0] epoch 150/800  test errors=145  best=137
  [teacher-1200 seed=0] epoch 200/800  test errors=150  best=118
  [teacher-1200 seed=0] epoch 250/800  test errors=127  best=108
  [teacher-1200 seed=0] epoch 300/800  test errors=112  best=102
  [teacher-1200 seed=0] epoch 350/800  test err…

teacher.eval()
with torch.no_grad():
    TEACHER_LOGITS = torch.cat([teacher(Xtr[i:i+2000]) for i in range(0, Xtr.size(0), 2000)]).detach()

torch.save({"state_dict": teacher.state_dict(),
            "final_errors": final_t, "best_errors": best_t,
            "logits": TEACHER_LOGITS.cpu()},
           "/home/user/checkpoints/teacher_1200.pt")

print("cached teacher logits on the UN-JITTERED 60k transfer set:", tuple(TEACHER_LOGITS.shape))
train_err = (TEACHER_LOGITS.argmax(1) != ytr).sum().item()
print(f"teacher train errors: {train_err}/60000   test errors: {final_t}")

# how soft are the targets at each temperature? (mean max-probability)
for T in (1, 4, 8, 20, 40):
    p = F.softmax(TEACHER_LOGITS / T, dim=-1)
    ent = -(p * p.clamp_min(1e-12).log()).sum(1).mean().item()
    print(f"  T={T:>2}: mean max-prob={p.max(1).values.mean():.4f}  mean entropy={ent:.4f} nats")
…

cached teacher logits on the UN-JITTERED 60k transfer set: (60000, 10)
teacher train errors: 104/60000   test errors: 75
  T= 1: mean max-prob=0.9941  mean entropy=0.0215 nats
  T= 4: mean max-prob=0.8412  mean entropy=0.6353 nats
  T= 8: mean max-prob=0.5429  mean entropy=1.5575 nats
  T=20: mean max-prob=0.2405  mean entropy=2.1824 nats
  T=40: mean max-prob=0.1586  mean entropy=2.2759 nats

teacher frozen — every student below distills from these identical logits

# Student 800x800, HARD LABELS ONLY: no dropout, no max-norm, no jitter, plain CE.
hard_errs, hard_best = [], []
for s in (0, 1, 2):
    f, b, h, _ = train_net(h=800, seed=s, epochs=150, bs=128, lr=0.1, momentum=0.9,
                           alpha=1.0, log_every=75, tag="student-hard")
    hard_errs.append(f); hard_best.append(b)
    print(f"  seed {s}: final={f}  best={b}")

print(f"\nSTUDENT 800x800 hard labels: mean={np.mean(hard_errs):.1f} +- {np.std(hard_errs):.1f} "
      f"(runs {hard_errs}, best-epoch {hard_best})   paper: 146")
RESULTS["student_800_hard"] = {"errors": hard_errs, "best": min(hard_best), "paper": 146}