Distillation: small, fast students from large teachers

Why this module

A large model is accurate and slow; the car needs fast and good enough. Distillation trains a small or fast student to imitate a large or slow teacher, so most of the teacher's quality survives at a fraction of the cost. It is how foundation models get onto vehicles and robots (DeepMind's on-device robotics postings ask for distillation explicitly), how diffusion and flow models get down to a few sampling steps, and how rule-based planners you may have written become training signals for learned ones.

The module covers four kinds of distillation, each with what is copied from the teacher:

1. Knowledge distillation: learning from soft targets

A classifier's softmax says more than its top class. A teacher that says "70 percent car, 25 percent truck, 5 percent pedestrian" also says that this car looks a bit like a truck and nothing like a pedestrian. Hinton, Vinyals and Dean (1503.02531) call this dark knowledge and train the student to match it, softened by a temperature $T$:

where $z_t$ and $z_s$ are the teacher's and student's logits. A high $T$ flattens both distributions so the small probabilities matter; the $T^2$ keeps the gradient size steady as $T$ changes; the optional cross-entropy term keeps the student honest on real labels.

def distillation_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    labels: torch.Tensor | None,
    temperature: float = 4.0,
    alpha: float = 0.9,
) -> torch.Tensor:
    """Hinton-style knowledge distillation.

    The student matches the teacher's softened distribution softmax(z / T); a high
    temperature T exposes how the teacher ranks the wrong classes ("dark knowledge").
    The soft term is scaled by T^2 so its gradients keep the same size as T changes.
    If labels are given, a share 1 - alpha of the loss is ordinary cross-entropy.

    Raises:
        ValueError: If ``temperature`` is not positive or ``alpha`` is outside [0, 1].
    """
    if temperature <= 0 or not 0.0 <= alpha <= 1.0:
        raise ValueError(
            f"need temperature > 0 and alpha in [0, 1], got {temperature}, {alpha}"
        )
    soft_teacher = torch.softmax(teacher_logits / temperature, dim=-1)
    log_student = torch.log_softmax(student_logits / temperature, dim=-1)
    soft = nn.functional.kl_div(log_student, soft_teacher, reduction="batchmean")
    soft = soft * temperature**2
    if labels is None:
        return soft
    hard = nn.functional.cross_entropy(student_logits, labels)
    return alpha * soft + (1 - alpha) * hard

What the temperature does

The ranking of the wrong classes is the information a hard label throws away. Too low a temperature hides it; too high flattens everything towards uniform. Typical values are 2 to 10.

Why the $T^2$

Differentiate the soft term with respect to a student logit and the temperature appears twice: once from the softmax, giving a factor $1/T$, and once more inside the probabilities:

As $T$ grows the probabilities themselves flatten towards $1/K$, so $p_s - p_t$ shrinks like $1/T$, and the product stays about the same size: without the $T^2$ the soft term's gradient would fade as $T$ rises, and its balance with the hard-label term would depend on $T$. In the high-temperature limit, $T\,(p_s - p_t) \approx (z_s - z_t)/K$ once the mean logit is removed: matching soft targets becomes matching logits. A test checks the loss equals $T^2$ times the tempered KL.

A worked comparison

The toy: six overlapping classes on a ring. The teacher (256 hidden units) trains on 20,000 labels; the student (16 hidden units) gets only 60 labels.

With soft targets, 60 examples were enough for the small student to match the teacher; with hard labels it fell 9 points short. Each soft target carries the teacher's whole belief about that point, which here is close to the true class probabilities, so a few points pin down the class boundaries. The unlabelled transfer set added nothing in this two-dimensional toy; on high-dimensional inputs, where 60 points cover almost nothing, it is usually the larger win, and it needs no human labels because the teacher labels it.

In this toy the temperature hardly mattered: students distilled at $T$ = 1, 2, 4, 8 and 16 all reached 0.784 to 0.788 (3 seeds). The six classes overlap, so the teacher's outputs are already soft at $T = 1$ and there is little hidden ranking to reveal. Temperature matters when the teacher is confident, as a large network trained to near-zero loss usually is.

DistilBERT (1910.01108) is the reference result for language: distilling BERT during pre-training gave a model 40 percent smaller and 60 percent faster that keeps 97 percent of its language-understanding performance.

2. Feature distillation

The student can also copy what the teacher computes inside. FitNets (1412.6550) train a deeper, thinner student to match the teacher's intermediate representations as "hints", and report a student with about 10 times fewer parameters outperforming its larger teacher on CIFAR-10. For driving, the useful version is to match the teacher's BEV or token features, so the student inherits its scene representation even where the outputs alone would say little.

3. Policy distillation, and why on-policy matters

For a policy, the "labels" are actions. Policy Distillation (1511.06295) distilled Atari agents into much smaller networks at expert level, and showed that distilling several single-task teachers into one student can beat the teachers.

The trap is the same as in imitation learning: if the student only learns on states the teacher visited, its small errors take it to states it never saw. On-policy distillation fixes this by letting the student act, then asking the teacher what it would have done there. Agarwal et al. (2306.13649) do this for language models (training on the student's own generated sequences with teacher feedback) to reduce the mismatch between training and inference. It is DAgger with a model as the expert, and the same reasoning applies to a small driving policy distilled from a large one: roll out the student in simulation, label its states with the teacher.

4. Step distillation: many sampling steps in one

Diffusion and flow models generate by integrating many small steps (Flow matching). Each step is a network call, which the Deploying module counts against the latency budget. Fewer, larger steps cut corners: with one Euler step, the flow below lands every sample between the two clusters, the average of where it could go.

But once trained, the teacher's map from a noise vector $z$ to a finished sample $x$ is deterministic. A student can learn that map by plain regression on pairs $(z, x)$ and then generate in one call:

def distill_sampler(
    teacher: VelocityMLP, teacher_steps: int = 16, steps: int = 2000, seed: int = 0
) -> nn.Module:
    """Trains a one-step generator to reproduce the teacher's multi-step samples.

    The teacher integrates its flow from noise z to a sample x with ``teacher_steps``
    Euler steps. That map z -> x is deterministic, so the student can learn it by
    plain regression on (z, x) pairs and then generate in one network call.
    """
    torch.manual_seed(seed)
    student = nn.Sequential(
        nn.Linear(2, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2)
    )
    optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
    for _ in range(steps):
        noise = torch.randn(256, 2)
        target = euler_sample(teacher, noise, teacher_steps)
        loss = ((student(noise) - target) ** 2).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return student

The student gets within 0.02 of the 16-step teacher at a sixteenth of the cost. What it learns is bounded by what it copies. Training the same one-step student on targets made with fewer teacher steps:

Up to 4 steps the student reproduces the teacher exactly, mistakes included. Beyond 8 steps the teacher keeps improving but the student stays near 0.28: from there on the limit is the student's own capacity to represent a sharp map with one small network, not the quality of its targets. The few points it leaves between the clusters are where the teacher's map changes fastest, which a smooth network finds hardest to copy. Progressive distillation (2202.00512) does this in stages, halving the step count each time, and consistency models (2303.01469) learn to map noise directly to data while still allowing more steps for quality (both in Sampling and guidance).

5. Distillation in driving

Patterns you will meet, and why each exists:

• Rule-based teachers. Hydra-MDP (2406.06978) distils from human and rule-based teachers, with separate heads for different evaluation metrics, so the student learns how the environment shapes planning end to end. For an engineer from a modular stack, this is the most direct use of existing planners and cost functions (From modular stacks to end-to-end).
• Large to small. Train the largest model the data supports offline, distil it into a model that fits the onboard budget. The teacher can also label unlabelled fleet data, which is the transfer set of section 1 at fleet scale.
• Slow decoder to fast decoder. Step distillation of a diffusion or flow trajectory decoder, as in section 4.
• Reasoning teacher to fast student. A model that writes a reasoning trace before acting (the Reasoning module) is too slow for every frame; distilling its actions into a trajectory-only student keeps much of the behaviour at a fraction of the latency. This is a design pattern, not a result we have measured.

What to watch for:

• The student inherits the teacher's mistakes, confidently. Evaluate the student on the scenes that matter, not on agreement with the teacher.
• Off-policy gaps. A student that only saw teacher states drifts; prefer on-policy data.
• Capacity limits. A student far smaller than the teacher loses the rare cases first, which are the safety-relevant ones. Track long-tail metrics separately (Scaling laws and data recipes).

Reference implementation (PyTorch)

libs/tutorial/distillation.py; tests are in libs/tutorial/tests/test_distillation.py: the loss is zero when student and teacher agree; two students with the same top class but a different ranking of the rest get different losses; hard labels enter with weight $1-\alpha$; a distilled student beats one trained on 60 labels by more than 5 points and comes within 2 points of the teacher; and a one-step student lands within 0.1 of the 16-step teacher while one teacher step lands between the clusters. Run uv run --group tutorial python -m libs.tutorial.distillation.

Check yourself

1. What does a teacher's softmax carry that a hard label does not, and what does the temperature do?
2. Why is the soft term multiplied by $T^2$?
3. Why did 60 soft-labelled points suffice in the toy, and when would a transfer set matter more?
4. What problem does on-policy distillation solve, and which imitation-learning method does it resemble?
5. Why does one Euler step of the teacher land between the clusters, and why can a one-step student avoid that?

Quiz

Answer each question; feedback appears as soon as you choose.

Sources

• Hinton, Vinyals and Dean, Distilling the Knowledge in a Neural Network, 2015.
• Romero et al., FitNets: Hints for Thin Deep Nets, 2014.
• Rusu et al., Policy Distillation, 2015.
• Sanh et al., DistilBERT, 2019.
• Agarwal et al., On-Policy Distillation of Language Models, 2023.
• Salimans and Ho, Progressive Distillation for Fast Sampling of Diffusion Models, 2022.
• Song et al., Consistency Models, 2023.
• Li et al., Hydra-MDP, 2024.