Model Distillation: Training Compact Models from Large Ones

The Problem: A Brilliant Model That Doesn't Fit Anywhere

Imagine you have built an AI system that reads handwritten answer sheets and sorts them by which digit or letter was written — something an OMR-scanning company or an exam board might need. You train a huge neural network on millions of examples, and it becomes extremely good: 99.3% accurate. There's just one problem. The network has 60 million parameters, takes 300 megabytes of storage, and needs half a second and a GPU to process a single image. Now suppose this system has to run inside a scanning app on a low-cost tablet used in a district exam centre with no reliable internet connection, and it needs to process a scanned sheet in under 50 milliseconds on an ordinary ARM processor with no GPU at all. The huge model simply will not fit. It is too slow, too heavy, and too power-hungry.

The obvious fix — "just build a smaller network from scratch" — usually fails badly. If you train a small network directly on the same labelled data the big one used, it typically ends up far less accurate, because it has less capacity to discover the patterns on its own. So the real engineering question is not "how do I make a small model?" — that's easy. It's "how do I make a small model that is almost as good as the big one?" This is exactly the problem that model distillation was invented to solve: instead of training a small model in isolation, you train it under the guidance of the large model, transferring what the large model has already learned into a compact form. The big model is called the teacher; the small model being trained is called the student.

Start With an Analogy: Two Kinds of Answer Sheets

Think about how you'd actually learn a subject from two different kinds of teachers. The first teacher hands back your test with just a tick or a cross next to each answer — correct or incorrect, nothing else. If you wrote "5" for a digit that was actually "3," you just get a cross. You know you were wrong, but you learn nothing about how wrong you were, or which of the wrong answers were reasonable near-misses.

Now imagine a second teacher who, instead of a tick or cross, writes something like: "This looks 51% like a 3, 31% like a 5, and 19% like an 8 — the loop and the slant are close to a 5, but the closed top confirms it's a 3." That second response is far richer. It doesn't just tell you the right answer — it tells you which wrong answers were almost right and which were nowhere close. A student who sees answers like this repeatedly starts to understand the deeper structure of the problem: what actually makes a 3 look different from a 5, and what makes both of them different from an 8.

This is the core intuition behind distillation. A correct-label dataset (used for ordinary training) is like the first teacher: for every training image, it only tells you the one correct class. A large, already-trained teacher network, on the other hand, can produce something like the second teacher's response — a full probability distribution over all possible classes, revealing which wrong answers are close calls and which are absurd. We call the single correct-answer label a hard label, and the full probability distribution the teacher produces a soft label. Distillation trains the small student network to match both — but especially to pay attention to the richer soft labels, because that is where most of the extra information lives.

Turning the Analogy Into Numbers: Softmax and Temperature

To make "soft labels" precise, we need to see how a network turns its internal scores into probabilities. A neural network's last layer produces one raw number per class, called a logit — an unnormalised score where a bigger number means the network is more confident about that class. Logits can be any real number, including negative ones, and they don't add up to anything meaningful on their own. To turn logits into probabilities that sum to 1, we apply a function called softmax.

Softmax uses Euler's number, e ≈ 2.718, raised to the power of each logit. Raising e to a power always gives a positive result, and — importantly — a small increase in the input produces a much bigger jump in the output when the input is large. That's exactly the behaviour that makes the highest logit dominate the resulting probability. The plain softmax formula for a logit z_i among classes is:

softmax(z_i) = e^(z_i) / (e^(z_1) + e^(z_2) + ... + e^(z_n))

Let's compute this on real numbers. Suppose the teacher network looks at an image and produces these logits for three visually similar digit classes — "3", "5", and "8": [5.0, 3.0, 1.0]. Plugging into softmax:

e^5.0 = 148.41
e^3.0 = 20.09
e^1.0 = 2.72
sum   = 171.22

P(3) = 148.41 / 171.22 = 0.867
P(5) =  20.09 / 171.22 = 0.117
P(8) =   2.72 / 171.22 = 0.016

At these settings, the teacher is almost completely certain it's a "3" (86.7%) and barely acknowledges "5" (11.7%) or "8" (1.6%). This is close to a hard label already — not much richer than a tick or cross. This is where temperature comes in. Distillation modifies softmax by dividing every logit by a temperature value T before exponentiating:

softmax(z_i, T) = e^(z_i / T) / sum_j e^(z_j / T)

When T = 1, this is identical to plain softmax. As T grows larger, the logits are squashed closer together before the exponential is applied, which makes the resulting probabilities more even — "softer." Let's recompute the same logits with T = 4:

Divide logits by T=4: [1.25, 0.75, 0.25]

e^1.25 = 3.49
e^0.75 = 2.12
e^0.25 = 1.28
sum    = 6.89

P(3) = 3.49 / 6.89 = 0.507
P(5) = 2.12 / 6.89 = 0.307
P(8) = 1.28 / 6.89 = 0.186

Now the picture is completely different. The teacher still favours "3" (50.7%), but it clearly signals that "5" is a serious contender (30.7%) while "8" is a distant third (18.6%). This softened distribution is the "second teacher" from our analogy — it reveals the relationships between classes that a hard label of just "3" could never show. Choosing T is a design decision: values roughly between 2 and 10 are common in practice; too low and you're back to near-hard labels, too high and even the wrong classes look equally plausible and the useful ranking gets washed out.

Reading the Code

Here is the softmax-with-temperature function written in Python using NumPy, along with the two computations above, so you can check it against the hand-worked numbers:

import numpy as np

def softmax(logits, T=1.0):
    scaled_logits = np.array(logits) / T
    # subtract the max value first only to avoid overflow;
    # it does not change the final probabilities
    exp_vals = np.exp(scaled_logits - np.max(scaled_logits))
    return exp_vals / np.sum(exp_vals)

teacher_logits = [5.0, 3.0, 1.0]   # scores for classes "3", "5", "8"

print(softmax(teacher_logits, T=1))   # [0.867 0.117 0.016]
print(softmax(teacher_logits, T=4))   # [0.507 0.307 0.186]

Tracing this line by line: scaled_logits divides each logit by T. When T=1, this is [5.0, 3.0, 1.0] unchanged; when T=4, this becomes [1.25, 0.75, 0.25]. Subtracting the maximum (a standard numerical-stability trick, since large exponents can overflow a computer's floating-point range) shifts every value down without changing the final ratio — dividing every term of a fraction by the same constant leaves the fraction unchanged. np.exp then applies e to each entry, and dividing by the sum normalises everything so the three numbers add up to 1.0. Running this produces exactly the two probability triples computed by hand above.

The Teacher-Student Training Pipeline

With softened probabilities defined, we can describe the full distillation procedure as a pipeline with four stages: (1) train a large teacher model the normal way, using the full labelled dataset, until it reaches high accuracy; (2) freeze the teacher's weights completely — it will not be updated again; (3) for every training example, feed it through the frozen teacher to record its softened output at a chosen temperature, alongside the original correct hard label; (4) train a much smaller, separately-initialised student network on the same examples, using a loss function that pulls its predictions toward both the true hard label and the teacher's soft label. The diagram below shows exactly how a single training image flows through this pipeline and how both signals combine to update the student.

Notice what the diagram makes explicit: the teacher's blue box has no arrow flowing back into it — it is frozen and never changes during this process. Only the student (orange) receives an update, shown by the dashed red feedback arrow travelling from the loss box back into the student. The student's own bar chart starts out nearly flat (0.35 / 0.34 / 0.31) because an untrained small network has no real opinion yet; as training proceeds across many images, repeated feedback from the loss pushes the student's bars to gradually resemble the teacher's — sharper on "3," moderate on "5," low on "8."

Combining Two Kinds of Feedback: The Distillation Loss

We now need a single number that measures "how wrong" the student currently is, combining both the hard label and the teacher's soft label, so that standard training (which adjusts weights to reduce a loss number) can be used. This combined score is called the distillation loss, and it has two parts, each computed using cross-entropy — a standard way of scoring how far a predicted probability distribution is from a target one.

The soft loss compares the student's softened output (at the same temperature T used for the teacher) against the teacher's softened output. The hard loss compares the student's ordinary output (at T=1) against the true one-hot label. A weighting factor alpha (a number between 0 and 1) controls how much the total loss leans on each source — this is a choice the person training the model makes, often through experimentation:

def distillation_loss(student_logits, teacher_logits, true_label_index,
                       T=4.0, alpha=0.3):
    teacher_soft = softmax(teacher_logits, T)
    student_soft = softmax(student_logits, T)
    soft_loss = -np.sum(teacher_soft * np.log(student_soft))

    student_hard = softmax(student_logits, T=1.0)
    hard_loss = -np.log(student_hard[true_label_index])

    return alpha * hard_loss + (1 - alpha) * soft_loss

Let's trace this with real numbers, continuing our digit example. Teacher logits stay [5.0, 3.0, 1.0] (classes "3","5","8" in that order, so true_label_index = 0). Suppose the student, still early in training, produces logits [1.0, 0.8, 0.5] — much less confident and less separated than the teacher's.

teacher_soft (T=4) = [0.507, 0.307, 0.186]   (computed earlier)

student_soft (T=4): divide by 4 -> [0.25, 0.20, 0.125]
  e^0.25=1.284, e^0.20=1.221, e^0.125=1.133, sum=3.638
  student_soft = [0.353, 0.336, 0.311]

soft_loss = -(0.507*ln(0.353) + 0.307*ln(0.336) + 0.186*ln(0.311))
          = -(0.507*(-1.042) + 0.307*(-1.091) + 0.186*(-1.167))
          = -(-0.528 - 0.335 - 0.217)
          = 1.080

student_hard (T=1): [1.0,0.8,0.5]
  e^1.0=2.718, e^0.8=2.226, e^0.5=1.649, sum=6.593
  student_hard = [0.412, 0.338, 0.250]

hard_loss = -ln(0.412) = 0.886

total = alpha*hard_loss + (1-alpha)*soft_loss
      = 0.3*0.886 + 0.7*1.080
      = 0.266 + 0.756
      = 1.02

That single number, 1.02, is what a training loop tries to shrink toward zero over thousands of images, by adjusting the student's weights via backpropagation (the same weight-updating mechanism used to train any neural network — distillation only changes what the loss is computed against, not how weights get updated). As training continues, the student's own logits gradually spread out and align with the teacher's, and both the soft loss and hard loss shrink together. One refinement from the original research (mentioned here, not required for the arithmetic above): the soft loss is often additionally multiplied by T² before being combined, because softening with a larger temperature also shrinks the size of the gradient signal coming from that term, and the T² factor corrects for that so the two loss terms stay comparable in scale.

Why Softened Probabilities Teach More Than Correct Answers Alone

Here is the deeper reason this works, sometimes called dark knowledge — a term used in the original distillation research. A hard label only ever says "the answer is 3." It treats "guessing 5" and "guessing an upside-down teacup" as equally wrong, even though one is a reasonable mistake and the other is absurd. A well-trained teacher's soft label, by contrast, has implicitly learned — from seeing millions of examples — exactly how the classes relate to each other: which digits share visual features, which words share meanings, which images share textures. This relational information is not present anywhere in the original hard-labelled dataset; it only exists inside the trained teacher's weights. Distillation is precisely the mechanism for extracting that relational knowledge and re-encoding it into a smaller network, which is why a distilled student can often reach much higher accuracy than an identically-sized network trained from scratch on the same hard labels alone.

Beyond Matching Outputs: A Quick Look at Feature-Based Distillation

Everything above is called response-based distillation — the student only ever tries to match the teacher's final output layer. Researchers have also explored feature-based distillation, where the student is additionally trained to match the teacher's intermediate hidden-layer activations, not just its final answer — an approach introduced under the name "FitNets" (Romero and colleagues, 2014–15), sometimes described as giving the student "hints" partway through the network, not just at the end. This tends to help when the student's architecture is very different in shape from the teacher's, since matching only the final output gives the student little guidance about how to organise its internal layers. For a Grade 9 level, the key fact to hold onto is simpler: distillation is not limited to one exact recipe — matching final soft probabilities is the classic and most common version, but the same underlying idea (train a small network using signals from a large trained one, rather than raw labels alone) can be applied at different points inside the network.

Common Misconception: Distillation Is Not Pruning, and It Is Not Quantization

A very common mix-up is to treat "model distillation," "pruning," and "quantization" as three names for the same thing — "making a model smaller." They are not the same technique, and confusing them will cost marks in any exam that tests conceptual understanding. Pruning starts with one already-trained network and deletes individual weights or whole neurons that contribute little to its output, keeping the same basic architecture minus some connections. Quantization keeps every weight and the entire architecture exactly as it is, but stores each number using fewer bits — for example, representing weights with 8-bit integers instead of 32-bit floating-point numbers, trading a little precision for a large reduction in memory and faster arithmetic. Distillation is different in kind from both: it does not start by editing an existing trained network at all. It trains an entirely separate, independently-initialised network — often with a completely different architecture and a different number of layers — from the beginning, using a second network's soft outputs as an additional training signal. In real deployment pipelines, these three techniques are often combined: you might first distil a large teacher into a smaller student, then prune a few more redundant connections from that student, then quantize the final result — three independent compression strategies stacked together, not three words for one idea.

Where This Matters: Real Systems Built for Indian Conditions

India's scale makes distillation especially relevant, because so much AI here has to run where connectivity and hardware are constrained rather than assumed. Consider a real-time payment system processing UPI transactions: fraud-detection scoring has to complete in a few tens of milliseconds per transaction, at a volume of many transactions every second, which makes a slow, giant model impractical no matter how accurate it is in isolation — a distilled, fast model that keeps most of the teacher's accuracy is often the only version that can actually be deployed at that speed and scale. Or consider an offline exam-preparation app meant to work for students in areas with patchy mobile data — any on-device AI feature (say, checking handwritten practice answers, or giving instant feedback on a spoken English sentence) has to fit and run entirely on a budget smartphone's processor, with no cloud round-trip possible. Satellite and space applications face an even stricter version of the same constraint: onboard processing on a satellite has a fixed, small compute budget and cannot simply request more GPU time from the ground, so any AI model analysing imagery onboard has to already be compact before launch, not compressed later. In every one of these situations, the underlying pattern is identical to the one this chapter walked through: start with the most accurate model you can train where compute is abundant, then distil what it has learned into a model shaped to fit where compute is scarce. A well-known example from natural language processing illustrates just how effective this can be: DistilBERT, a distilled version of the BERT language model published by Hugging Face researchers in 2019, is roughly 40% smaller and runs about 60% faster than the original BERT, while retaining approximately 97% of its performance on standard language-understanding benchmarks — a striking demonstration that most of a large model's capability really can survive compression into a much smaller one, if the compression is done through distillation rather than naive shrinking.

Test Your Understanding

1. A teacher model produces logits [4.0, 2.0, 0.0] for three classes. Compute the softmax probabilities at T=1 and again at T=5. Which of the two distributions carries more "dark knowledge" about which wrong class is closest to being right, and why?
2. In the distillation loss formula alpha * hard_loss + (1-alpha) * soft_loss, what does training reduce to if you set alpha = 1? What does it reduce to if you set alpha = 0?
3. Explain, in your own words, why a student model trained only on one-hot hard labels can never learn that "3" and "5" are more visually similar than "3" and "8," even after seeing a million training images.
4. True or False, with a correction if false: "Model distillation is just another word for pruning a neural network."
5. A hospital currently runs a large diagnostic model on cloud servers and wants a compact version running offline on a tablet in a rural clinic. List the steps you would follow to prepare that compact model using distillation.
6. Why must the teacher model be fully trained and frozen before the student begins training, rather than training teacher and student together as equal partners from the start?

Answers

1. At T=1: probabilities are approximately [0.867, 0.117, 0.016]. At T=5: dividing logits by 5 gives [0.8, 0.4, 0.0], giving probabilities of approximately [0.472, 0.316, 0.212]. The T=5 distribution carries more dark knowledge: at T=1 the second class is barely visible (0.117) and the third looks almost impossible (0.016), while at T=5 the gap between the classes is far smaller, clearly showing that the second class is a serious near-miss rather than an irrelevant option.
2. With alpha=1, the soft loss term is multiplied by zero and disappears entirely — the student trains as if it were an ordinary network learning only from hard labels, with no distillation happening at all. With alpha=0, the hard loss disappears completely — the student learns purely by imitating the teacher's soft output and never directly sees the ground-truth labels during that part of training.
3. A one-hot hard label only ever contains a single 1 and the rest 0s — for a "3," the label vector looks the same ([1,0,0] over classes "3,5,8") every single time, regardless of whether that particular "3" was neatly written or scrawled to look almost like a "5." There is no channel in a hard label through which "closeness to 5" could ever be represented, no matter how many examples the network sees — the information simply isn't present in the target it's being trained to match.
4. False. Pruning starts from an already-trained network and deletes some of its existing weights or neurons, keeping the same underlying architecture. Distillation trains a brand-new, independently-initialised network (often with a different architecture altogether) from scratch, using a separate trained network's soft outputs as an extra training signal. They achieve a similar goal — smaller, faster models — through completely different mechanisms.
5. Train (or use an existing) large, accurate teacher model on the full hospital dataset using cloud-scale compute. Freeze the teacher and run it over the training images to record its softened probability outputs at a chosen temperature. Design a small student architecture sized to fit the tablet's memory and processing limits. Train the student using the combined distillation loss (hard label plus teacher soft label) until its validation accuracy is within an acceptable margin of the teacher's — checked especially carefully here, since a medical application demands very high reliability. Deploy the trained student onto the tablet, applying quantization afterward if additional size reduction is still needed.
6. The entire method depends on the teacher's soft outputs already containing genuine, learned relationships between classes. An untrained (or still-training) teacher's soft outputs would just be close to random noise, giving the student nothing meaningful to imitate. The soft-label signal only becomes useful once the teacher itself has already learned real structure in the data — which is why it must be trained and frozen first.

Summary: What Distillation Actually Does

Model distillation compresses the knowledge of a large, accurate "teacher" network into a small, fast "student" network — not by editing the teacher's weights, but by training a separate, smaller network to imitate the teacher's behaviour. The key technical trick is softening the teacher's output probabilities with a temperature parameter before using them as a training target, because a softened distribution reveals which wrong answers are near-misses and which are nonsense — information a plain correct/incorrect label can never carry. The student is trained on a combined loss that weighs this soft-label imitation against ordinary training on the true hard labels, using the same backpropagation mechanism as any other neural network. Distillation is a distinct compression strategy from pruning (deleting weights from an existing network) and quantization (storing existing weights with less precision), and in real systems the three are often combined rather than used as substitutes for one another. Wherever accurate AI needs to run somewhere compute is scarce — a budget smartphone, a real-time payments pipeline, an offline rural clinic, an onboard satellite processor — distillation is the standard technique for keeping most of a large model's intelligence while discarding most of its size.