Research & AlgorithmsDANIEL SCHMITTER · RESEARCH & ENGINEERING

Research & Algorithms / ML Engineering

Representation learning · E31 · Implementation

Distill a representation, not just a class prediction

A compact CNN learns against teacher features with an optional classification loss. The feature interface is the real contract between teacher and student.

Daniel Schmitter · September 19, 2026 · 6 min read

PyTorch CNNCached teacher arraysMSE / cross-entropy
A compact student matches teacher features, not just class probabilities. The two feature-map stacks meet at the training loss.
Figure 1. Transfer the representation. A compact student matches teacher features, not just class probabilities. The two feature-map stacks meet at the training loss. Distillation schematic. Original vector illustration.

Follow the information

From input to outcome

The teacher supplies cached feature targets for the same image identities. MSE sends gradients into the student feature path, not into the frozen teacher. Optional label cross-entropy is an additional branch documented in the training diagram.

The teacher supplies cached feature targets for the same image identities. MSE sends gradients into the student feature path, not into the frozen teacher. Optional label cross-entropy is an additional branch documented in the training diagram.
Figure 2. Information flow. Solid arrows carry observations, tensors or artifacts; other routes are explicitly labelled. Signal shapes, matrices and network icons are schematic, not measured samples or literal neuron counts. Open full-size SVG ↗ On narrow screens, scroll the diagram horizontally.

Read this alongside Figure 1: A compact student matches teacher features, not just class probabilities. The two feature-map stacks meet at the training loss. The module map and layer-level figures below expand the operations in this route.

Distill a representation, not just a class prediction: architectureTraining image: CIFAR-100 input → Student CNN: 32 → 64 → 128 channels → Feature projection: 512 dimensions → Teacher feature target: Cached representation → MSE / cross-entropy: Configurable learning objective → Student evaluation: Classifier or ridge probe. A high-level module map; comparison branches and training details are explained in the article.REPRESENTATION LEARNING / E31 / MODULE MAP01 INPUTTraining imageCIFAR-100 input02 MODULEStudent CNN32 → 64 → 128 channels03 MODULEFeature projection512 dimensions04 MODULETeacher feature targetCached representation05 MODULEMSE / cross-entropyConfigurable learning objective06 OUTPUTStudent evaluationClassifier or ridge probe
Source-grounded module map. Boxes summarize operations, not individual neurons; comparison arms and training paths are detailed below. On a small screen, scroll the diagram horizontally.
Training image — CIFAR-100 input

The architecture in context

The system we are building

Feature distillation asks the student to reconstruct an internal representation rather than merely imitate the teacher’s winning class. The compact network uses three convolutional stages and a projected feature vector, with a classifier available as a separate head. The training function supports feature MSE, classification cross-entropy and a combination.

Who does what in the stack

PyTorch CNN
Defines the compact student and feature projection.
Cached teacher arrays
Supply fixed targets with explicit sample identities.
MSE / cross-entropy
Separate representation imitation from label learning.

The custom loop aligns teacher features with student images and chooses which loss reaches which parameters. PyTorch supplies the network and optimizer; cached teacher arrays remove the need to execute the teacher at every student update. They do not remove the cost of producing or storing the teacher representation from the whole system.

Framework responsibility map. Each row maps a library or custom component to its job; rows are not a sequential inference graph.
Framework responsibility map. Each row maps a library or custom component to its job; rows are not a sequential inference graph. Open full-size SVG ↗

From module map to executable structure

Inside Compact distillation student

CIFAR-100 images 3×32×32; teacher feature width 512; student classifier 100 classes.

Layer-level implementation. B denotes batch size; parameter and shape conventions are expanded in the table.
Layer-level implementation. B denotes batch size; parameter and shape conventions are expanded in the table. Open full-size SVG ↗
Layer / tensor / operation ledger
Layer or branchOutput shapeImplementation detail
Conv 3→32 + ReLU + pool 2B × 32 × 16 × 163×3, stride 1, padding 1.
Conv 32→64 + ReLU + pool 2B × 64 × 8 × 83×3, stride 1, padding 1.
Conv 64→128 + ReLUB × 128 × 8 × 8Adaptive average pool 1×1 → flatten 128.
Feature projectionB × 512Linear 128→512; raw output is the feature-MSE target space.
Optional classifierB × 100ReLU(features) → Linear 512→100. Unused by feature-only distillation.

The teacher features are cached targets, so the teacher is not in the student’s backward graph. The feature projection is linear; applying its later classifier ReLU before computing MSE would change the experiment. Downstream Gram evaluation uses frozen features and must be reported separately from the student’s direct classifier accuracy.

The equation and the update

Ldistill=1512B∑i∥fθ(xi)−ti∥22\mathcal L_{\rm distill}=\frac{1}{512B}\sum_i\|f_\theta(x_i)-t_i\|_2^2

Adam 1e-3, batch 256; default 8 epochs, two seeds. Feature-only loss is MSE; hybrid adds supervised cross-entropy; scratch uses cross-entropy alone. These objectives update different parts of the same allocated model.

Learning or solution path. A parameter-update path is different from the forward inference path; see text for target-network, frozen-feature and local-loss boundaries.
Learning or solution path. A parameter-update path is different from the forward inference path; see text for target-network, frozen-feature and local-loss boundaries. Open full-size SVG ↗

Implementation card / no invented benchmarks

Capacity, budget and execution evidence

Parameters / retained state
210,596 allocated scalars; 51,300 classifier scalars receive no gradient in feature-only mode. The feature path has 159,296.
Duration and hardware evidence
The training function returns elapsed time, but E32’s JSON does not retain it; no duration is supplied as if measured.
Source coordinates
E31 lines 34–42, 71–91; E32 saved metrics
Current reproduction context
Current workstation, supplied by the author: Apple M4, 128 GB unified RAM, 40 GPU cores and 16 CPU cores. This is context for prospective reproduction, not attribution of every archived run. Python and framework versions are not fully locked for these historical sources; declarations, when available, are identified separately.

Counts above are calculated from the stated layer shapes unless identified as saved measurements. They exclude optimizer state and nontrainable buffers. No archived training was rerun for this revision.

What these design choices change

Global pooling makes the feature projection much smaller than flattening the last spatial map. A512-wide representation is chosen to match the teacher target, not because every student needs 512 features. Reducing it requires a new teacher projection or a different distillation objective.

Reproduction and measurement protocol

Assert teacher/student row identity before comparing features; cached targets in a different shuffle order create a valid-shaped but meaningless loss. Check classifier gradients are absent in feature-only mode. The saved teacher-to-student accuracy gap shows that parameter reduction did not preserve the teacher’s usefulness.

For a new run, save the resolved Python/framework versions, backend, dtype, seed, input shapes, batch size and exact source revision. Start with one batch and one update. Log training steps separately from epochs or environment steps. Do not equate the configured maximum with a completed budget or convergence.

Measure initialization/compilation, data preparation, warmed forward pass, training updates and evaluation separately. Synchronize accelerator work around timed regions using the chosen framework’s supported mechanism. Report peak process memory and framework allocation separately; parameter bytes exclude activations, gradients, optimizer state and input buffers. On a shared machine, begin with a single CPU worker and a small batch rather than claiming all available resources.

A closer look at the implementation

The code that carries the idea

The excerpt distinguishes feature regression from label supervision. Matching feature dimensions is only a shape check: normalization, scale and sample ordering also determine what MSE means. A classifier trained on those features and a separately fitted Gram readout are different deployment heads.

Python · file · lines 71–85
def train_student(Xtr, ytr, Ftr, mode, epochs, seed):
    torch.manual_seed(seed); net = StudentCNN(); opt = torch.optim.Adam(net.parameters(), 1e-3)
    n = len(Xtr); t0 = time.time()
    for ep in range(epochs):
        perm = torch.randperm(n)
        for i in range(0, n, 256):
            bi = perm[i:i + 256]; opt.zero_grad()
            if mode == "distill":
                loss = ((net.features(Xtr[bi]) - Ftr[bi]) ** 2).mean()  # feature distillation only
            elif mode == "distillsup":                                 # hybrid: feature-distill + supervised (best shot for compact)
                loss = ((net.features(Xtr[bi]) - Ftr[bi]) ** 2).mean() + nn.functional.cross_entropy(net(Xtr[bi]), ytr[bi])
            else:
                loss = nn.functional.cross_entropy(net(Xtr[bi]), ytr[bi])  # from scratch supervised
            loss.backward(); opt.step()
    return net, sum(p.numel() for p in net.parameters()), time.time() - t0

Verbatim archive excerpt from closed_form_distill_v3_compact.py. Context-dependent historical code, not a standalone runnable program. Comments retain their original wording; the article distinguishes implemented behavior from stale or overbroad comments.

The boundary that matters

A small parameter count is not proof of successful knowledge transfer. The saved student falls substantially short of its teacher-feature comparator; E32 reports that outcome. The archive does not establish a teacher-equivalent edge model.

Keep building

Other posts of interest