Histopathology Classification and Knowledge Distillation

Two small, self-contained PyTorch experiments from 2024, each one script that downloads its own data and runs end to end. The first asks whether architecture capacity helps on a small medical-imaging benchmark: a two-block CNN against an ImageNet-pretrained ResNet18 on PathMNIST, colorectal-cancer histology tiles in nine tissue classes. The second asks whether knowledge distillation helps a small network: a teacher CNN on MNIST, then a half-width student trained twice, once against the teacher and once without. Both are in the repository under MIT.

Left: PathMNIST test accuracy, SimpleCNN 77.23 percent against a pre-trained ResNet18 at 81.91 percent. Right: MNIST test accuracy, teacher 98.81, student trained with distillation 98.72, student trained without it 98.64
Two separate experiments on different datasets. Left: PathMNIST, a two-block CNN against a pre-trained ResNet18 after an identical budget. Right: MNIST, a teacher and two students of the same architecture, one trained with a distillation loss against the teacher and one without.

Two questions, one design

Both experiments ask the same kind of question—where the accuracy of a small network actually comes from—and both answer it the same way: change exactly one thing, hold everything else fixed, and read the gap. On PathMNIST the thing that changes is the architecture and what it was initialised from. On MNIST the architecture is held constant and what changes is the signal the loss is computed against. In each pair the data, the preprocessing, the optimiser, the batch size and the number of epochs are identical, so a difference in the final number has one candidate explanation rather than several.

That constraint is what shaped the code. Each experiment is a single script that downloads its own data, builds both models, trains them under one budget defined in one place, and prints the two numbers side by side. There is no configuration surface for the two arms of a comparison to drift apart through, and no state carried between runs: the result on the page is what the script produces from an empty directory.

The two experiments

Two panels. (a) PathMNIST tiles are converted to tensors and normalised, then split into two branches: a two-block SimpleCNN of 0.42 million parameters with a fully connected head, and an ImageNet-pretrained ResNet18 of 11.2 million parameters, a 7 by 7 convolution, four residual stages, global average pooling and a fully connected head. Both are evaluated on the same data for the same number of steps, giving 77.23 against 81.91 percent. (b) MNIST goes to a teacher CNN of 1.20 million parameters whose softmax at temperature 5 produces soft targets, and to a student CNN of 0.30 million parameters; the loss is alpha times the KL divergence to the teacher plus one minus alpha times the cross entropy against the hard labels, and the same student is also trained with the first term removed.
The two runs end to end. They share the preprocessing and differ in what is being compared: two architectures on PathMNIST, and one architecture trained twice on MNIST, with and without the distillation term in the loss.
Experiment Models Training Test accuracy
PathMNIST, 9 classes, 28 × 28 RGB SimpleCNN, two conv blocks, ~0.42 M parameters, from scratch
ResNet18, ~11.2 M, ImageNet weights, fully fine-tuned		 Adam 1e-3, batch 64, 5 epochs, no augmentation, no scheduler, no seed 77.23%
81.91%
MNIST distillation Teacher CNN, ~1.2 M
Student at half width, ~0.30 M, with and without a distillation loss (T = 5, α = 0.7)		 Adam 1e-3, batch 128, 5 epochs each, no seed 98.81%
98.72% / 98.64%

The four networks

SimpleCNN is two convolution blocks and a fully connected head, about 0.42 M parameters, initialised from nothing. It sets the floor: what a model small enough to train from scratch in minutes can reach on 28 × 28 tissue tiles.

ResNet18 sets the other end of the same budget: a 7 × 7 stem, four residual stages, global average pooling, and a new nine-way head on about 11.2 M parameters carrying ImageNet weights. Every layer is fine-tuned rather than frozen, so the comparison is between two models that both spent the whole budget learning this data—one starting from random weights, one starting from features learned on photographs. Whether those features are worth anything on stained tissue, which looks nothing like ImageNet, is the question the run answers.

Teacher and student on MNIST are the same design at two widths: a convolutional feature stack and a classifier head at about 1.20 M parameters for the teacher and 0.30 M for the student, a quarter of the size. Keeping the shape fixed and moving only the width is what makes the student's two runs comparable to each other and to the teacher.

What distillation moves across

A one-hot label says which class an image belongs to and nothing else. A trained network's output says more than that: a four that leans towards nine scores nine higher than it scores seven, and that ordering is a statement about the shape of the input, learned from the whole training set. Under an ordinary softmax the winning class takes almost all the probability mass and the rest of that structure is numerically invisible.

Temperature is what makes it visible. Dividing the logits by T = 5 before the softmax flattens the distribution until the relative sizes of the losing classes matter, and the student is asked to match that flattened distribution rather than the label. The loss puts both signals side by side: α = 0.7 on the divergence from the teacher's softened output, and the remaining 0.3 on the ordinary cross entropy against the hard label, so the student is pulled towards the teacher's judgement without being allowed to drift away from the ground truth.

The control is the same student, the same data and the same budget with the first term removed. Two runs that differ in one term of one loss are what turn “distillation helps” from a claim into a measurement.

What the numbers show

PathMNIST. A 4.7-point gap after an identical, short budget on a test split of about 7,180 images is well outside noise: the pre-trained ResNet18 separates cleanly from the two-block CNN trained from scratch, and the comparison is set up so that both models see exactly the same data for exactly the same number of steps.

MNIST. The same student architecture is trained twice, once against the teacher's temperature-softened outputs and once on the labels alone, so the only thing that differs between the two runs is the distillation term. With it, the 0.30 M student closes part of the distance to a teacher four times its size.

Resources