Knowledge Distillation: Training Small Models

Open Google Pay or PhonePe and send someone ₹500. Somewhere between your tap and the "Payment Successful" screen, a machine learning model has to decide whether this transaction looks like fraud, and it has to decide well under a second — UPI transactions in India run into the billions every month, and the switching systems behind them are built around tight response-time budgets for every single one. A slow fraud check is a failed payment, and a failed payment at a roadside stall or a kirana store is a lost customer.

Here is the catch. The most accurate fraud-detection model a bank or payments company can build is usually large: a deep network with tens of millions of parameters, trained offline on years of transaction history, where it has all the time and compute it wants. That big model is excellent at telling real transactions from fraudulent ones. But it is too slow and too heavy to sit in the live path of every UPI transaction happening across the country. What actually screens your ₹500 payment in real time has to be small, fast, and cheap to run, and yet it still has to be almost as good as the giant model that isn't in the room.

This is the exact problem knowledge distillation was invented to solve: how do you take what a large, accurate model has learned and pour it into a small, fast model, without losing much of what made the large model good in the first place?

Teacher, Student, and the Compression Problem

This chapter sits inside a subject called model compression: the set of techniques for shrinking a trained neural network's size, memory footprint, and inference cost while keeping as much of its accuracy as possible. Model compression matters because accuracy and efficiency usually pull in opposite directions. Larger networks, with more layers, more parameters, wider hidden dimensions, tend to fit complex patterns better and generalize better, up to a point. But every extra parameter costs memory to store, energy to compute, and milliseconds to run. Model compression has three main families of technique: pruning (deleting weights or whole neurons that contribute little to the output), quantization (storing and computing with lower-precision numbers, such as 8-bit integers instead of 32-bit floats), and knowledge distillation, which is the subject of this chapter.

Knowledge distillation takes a different approach from the other two. Instead of shrinking an existing network by cutting pieces out of it, you train an entirely new, smaller network from scratch, using the original labelled data plus something extra: the output of a larger, already-trained network. The large, accurate network is called the teacher. The small network being trained is called the student. The technique was formalized in a short but hugely influential 2015 paper, "Distilling the Knowledge in a Neural Network," by Geoffrey Hinton, Oriol Vinyals, and Jeff Dean at Google. Their central claim was that a trained network's output carries far more useful information than just its final answer, and a student that learns to reproduce that fuller output can end up more capable than a student trained on the raw labels alone.

What the Teacher Knows That the Labels Don't Say

Picture a large, well-trained classifier looking at a photograph of a handwritten digit that is actually a "3". A typical training label for this image is a hard label: a one-hot vector like [0, 1, 0, 0] over the classes {0, 3, 5, 8}, which says "it is class 3, full stop, and every other class is impossible." But that is not what a good trained network actually believes when it looks at the image. It might assign 84% probability to "3", but also 11% to "5" and 4% to "8", because a slightly untidy handwritten 3 shares curves with both, and it correctly assigns almost nothing to "0", because a 3 and a 0 do not look alike at all.

That extra information, the relative probabilities the network assigns to the wrong classes, is not noise. It reflects something real the network has learned about how classes relate to each other. Hinton called this dark knowledge: knowledge that is present inside a trained model's output distribution but invisible if you only ever look at its single most likely class. A hard label can never contain this information, because by construction it assigns zero probability to everything except the correct answer. A student trained purely on hard labels has to rediscover every one of these class relationships on its own, using only a correct-or-incorrect signal on each example. A student that is shown the teacher's full probability distribution learns the relationships directly, which packs a far richer training signal into every single example.

Softmax Temperature: Turning Up the Contrast

There is a practical problem with using a trained network's raw output directly as a training target: a confident, well-trained teacher usually produces probabilities that are already extremely close to a hard label. If the teacher is 99.97% sure an image is a "3" and assigns 0.0002% to "8", that tiny probability carries almost no usable gradient signal — it is nearly as uninformative as a one-hot vector. To make the dark knowledge easier for a student to learn from, distillation "softens" the teacher's output using a technique called temperature scaling.

Recall that a classifier's last layer produces raw scores called logits, one real number per class before any normalization, and the standard softmax function turns logits z into probabilities:

q_i = exp(z_i) / Σ_j exp(z_j)

Distillation generalizes this with a temperature parameter T:

q_i = exp(z_i / T) / Σ_j exp(z_j / T)

When T = 1, this is exactly the ordinary softmax. As T grows past 1, dividing every logit by a larger number squeezes the differences between logits closer together before exponentiating, which spreads probability mass more evenly across classes: the distribution softens. As T → ∞, the output approaches a uniform distribution over all classes. As T shrinks below 1, the opposite happens: the distribution sharpens, and as T → 0 it approaches the one-hot hard label. Temperature is purely a training-time knob for shaping the loss signal. Once the student is trained and ready to deploy, it runs with the ordinary T = 1 softmax like any other classifier; nobody ships a model that only works at T = 4.

Building the Distillation Loss

Knowledge distillation trains the student against two targets at once, combined into a single loss function:

• A soft loss: how well the student's softened output (at some temperature T > 1) matches the teacher's softened output at the same temperature. This is usually measured with Kullback-Leibler (KL) divergence, written KL(p‖q), which measures how different a distribution q is from a reference distribution p.
• A hard loss: the ordinary cross-entropy between the student's normal (T = 1) prediction and the true ground-truth label, exactly as if the student were being trained the standard way with no teacher involved.

The two are combined as:

L_total = α · T² · L_soft(teacher, student; T) + (1 − α) · L_hard(y_true, student)

where α is a weight between 0 and 1 controlling how much the student should trust the teacher's soft guidance versus the raw ground truth. The T² term is not optional decoration. As T grows, the gradients produced by the soft loss shrink, roughly in proportion to 1/T². Left uncorrected, raising the temperature to expose more dark knowledge would simultaneously weaken the soft loss's pull on training, tangling two effects that should be independent. Multiplying the soft loss by T² keeps its contribution to the total gradient roughly stable as T is tuned, so temperature and loss weighting stay separate, controllable knobs. This detail comes directly from Hinton, Vinyals, and Dean's original analysis of how the soft-target gradient behaves at large T.

One more useful fact ties the classic formulation to what most libraries implement today: cross-entropy and KL divergence between the same two distributions differ only by the teacher's own entropy, a quantity that does not depend on the student at all. Minimizing KL(teacher‖student) and minimizing the cross-entropy between them therefore produce identical gradients with respect to the student's weights. The two formulations you will meet in different papers and frameworks are not competing ideas, just the same objective written two ways.

Worked Example: Distilling a Digit Classifier

Let's trace the arithmetic through one training example, using the four-class digit setup from earlier: classes {"0", "3", "5", "8"}, true label "3". These are constructed numbers, chosen to make the mechanics easy to trace by hand, not measurements from a real dataset. Suppose that for one training image, the teacher and student networks produce these logits:

• Teacher logits: 0 → 1.0, 3 → 6.0, 5 → 4.0, 8 → 3.0
• Student logits: 0 → 0.5, 3 → 3.0, 5 → 2.5, 8 → 2.0

First, look at the teacher's ordinary softmax at T = 1, which is what you would use if you only cared about its final answer: 0 → 0.57%, 3 → 83.90%, 5 → 11.35%, 8 → 4.18%. The teacher is confidently right, but almost all of the interesting structure, the fact that 5 is a far more plausible second guess than 0, has been squashed down toward single-digit percentages. Now soften both networks' logits at temperature T = 4:

• Teacher softmax at T = 4: 0 → 12.11%, 3 → 42.28%, 5 → 25.64%, 8 → 19.97%
• Student softmax at T = 4: 0 → 16.74%, 3 → 31.28%, 5 → 27.61%, 8 → 24.36%

Notice two things. First, both distributions now agree on the full ranking, 3 > 5 > 8 > 0. The student agrees with the teacher on which class is correct, and it has also absorbed the same relative ordering among the wrong classes — dark knowledge, successfully transferred. Second, the student's distribution is flatter than the teacher's; it is a smaller, less confident network this early in training, and its job right now is to move toward the teacher's shape, not to match it exactly on one update.

Now compute the two loss terms. The soft loss is the KL divergence between the teacher's and student's T = 4 distributions:

KL(teacher ‖ student) = Σ p_i · log(p_i / q_i) ≈ 0.0294

Scaled by T² = 16, this becomes the soft-loss contribution: 16 × 0.0294 ≈ 0.4707. The hard loss is the student's ordinary cross-entropy against the true label "3", read off its T = 1 softmax, where 3 gets 48.63% probability, so the hard loss is −log(0.4863) ≈ 0.7210. Weighting both terms equally with α = 0.5:

L_total = 0.5 × 0.4707 + 0.5 × 0.7210 ≈ 0.2354 + 0.3605 ≈ 0.5959

That single number, 0.5959, is what backpropagation actually uses to update the student's weights for this one example, pulling it simultaneously toward the teacher's soft beliefs and toward the ground truth. Here is the full calculation as runnable Python, reproducing every figure above:

import math

def softmax_with_temperature(logits, T=1.0):
    scaled = [z / T for z in logits]
    m = max(scaled)                      # subtract max for numerical stability
    exps = [math.exp(s - m) for s in scaled]
    total = sum(exps)
    return [e / total for e in exps]

def kl_divergence(p, q):
    return sum(pi * math.log(pi / qi) for pi, qi in zip(p, q))

# Logits for one handwritten-digit image, over classes ["0", "3", "5", "8"]
teacher_logits = [1.0, 6.0, 4.0, 3.0]      # large, accurate model
student_logits = [0.5, 3.0, 2.5, 2.0]      # small model being trained
true_class_index = 1                       # ground-truth label is "3"

T = 4.0
teacher_soft = softmax_with_temperature(teacher_logits, T)
student_soft = softmax_with_temperature(student_logits, T)
student_hard = softmax_with_temperature(student_logits, 1.0)

distillation_loss = kl_divergence(teacher_soft, student_soft)
hard_label_loss = -math.log(student_hard[true_class_index])

alpha = 0.5
total_loss = alpha * (T ** 2) * distillation_loss + (1 - alpha) * hard_label_loss

print(f"Teacher (T=4): {[round(p, 3) for p in teacher_soft]}")
print(f"Student (T=4): {[round(p, 3) for p in student_soft]}")
print(f"Distillation loss (KL):  {distillation_loss:.4f}")
print(f"Hard-label loss (CE):    {hard_label_loss:.4f}")
print(f"Total loss:              {total_loss:.4f}")

# Output:
# Teacher (T=4): [0.121, 0.423, 0.256, 0.2]
# Student (T=4): [0.167, 0.313, 0.276, 0.244]
# Distillation loss (KL):  0.0294
# Hard-label loss (CE):    0.7210
# Total loss:              0.5959

As training continues over thousands of examples like this one, gradient descent nudges the student's logits closer to the teacher's, including on the wrong classes, until the small network's whole belief structure starts to resemble the large one's — the rankings among wrong answers along with the right one.

Choosing a Student, and What It Costs to Train One

Knowledge distillation needs a student architecture before it needs a loss function. Two approaches are common. One is to shrink the same architecture family the teacher uses: fewer layers, narrower hidden dimensions, fewer attention heads, so the student is structurally a smaller cousin of the teacher. The other is to switch to a fundamentally more efficient architecture family and use distillation to help it catch up to the teacher's accuracy. Both are used heavily in production.

The clearest real-world case study is DistilBERT, released by Hugging Face in 2019. The teacher was BERT-base, a widely used language model with 12 transformer layers and about 110 million parameters. The student, DistilBERT, keeps the same hidden size (768) and the same general transformer design, but halves the depth to 6 layers, drops components BERT used only for its original pretraining setup, and is initialized by copying every other layer's weights directly from the teacher rather than starting from random values. The result has about 66 million parameters, 40% fewer than the teacher, runs roughly 60% faster, and retains about 97% of BERT-base's language understanding performance on the standard GLUE benchmark suite. DistilBERT's training loss blends three terms rather than two: the soft-target distillation loss (using a comparatively low temperature of T = 2), the standard masked-language-modelling loss on hard labels, and a cosine-similarity loss that pushes the student's internal hidden-state vectors, layer by layer, to point in the same direction as the teacher's.

That third term hints at a broader family of techniques known as feature-based distillation: instead of only matching the teacher's final softmax output, you also match its intermediate representations, such as hidden states or attention patterns. TinyBERT, from Huawei's Noah's Ark Lab, pushes this further with a multi-stage distillation process that matches embeddings, hidden states, and attention matrices layer by layer, in addition to the final logits. The lesson holds across all of these systems: transferring more of the teacher's internal reasoning brings the student closer to the teacher's accuracy at a fraction of the size.

Distillation is not the same thing as an efficient architecture. A network like MobileNet is small mainly because of its design: it replaces standard convolutions with a cheaper operation called a depthwise separable convolution, and it can be trained the ordinary way, with no teacher at all. Distillation and efficient architecture design are complementary rather than competing — you can design a small, efficient student and still train it by distillation from a large teacher. In practice, the best compressed models usually combine several compression techniques, a leaner architecture, distillation, pruning, and quantization together, rather than relying on any single one.

Back to the Payment Screen

Return to that UPI transaction. In a real fraud-detection pipeline, the setup looks a lot like the digit example above, just with far more classes, features, and training examples. A large teacher model, perhaps an ensemble of gradient-boosted trees and deep networks, is trained offline on years of transaction history with no real-time constraints on it at all; its only job is to be as accurate as possible. A much smaller student network is then trained to match the teacher's soft risk scores across millions of past transactions, plus the actual fraud and not-fraud hard labels, using exactly the combined loss built in this chapter. The student that ends up running in the live payment path never consults the teacher at inference time — it does not need to. Everything useful the teacher worked out during its slow, expensive offline training has already been poured into the student's weights during distillation, and the student now carries that knowledge alone, fast enough to return a verdict before you have finished looking at your phone.

This is also why knowledge distillation matters more in India than the raw technique might suggest at first glance. Indian users run their apps across an unusually wide spread of hardware, from flagship phones to budget Android devices with a few gigabytes of RAM, over mobile networks whose quality varies sharply between a metro apartment and a village with patchy connectivity. A keyboard app predicting the next word in Hindi or Hinglish, a UPI app scoring a payment, or a voice assistant handling a regional language all face the same choice this chapter has been about: ship the big, accurate model and make some users wait, or ship a distilled model that fits the weakest device in the user base and stays fast for everyone. Distillation is what lets a team choose accuracy and speed together, instead of accuracy or speed.

Think About It

Think about this: How would you explain knowledge distillation: training small models to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind knowledge distillation: training small models, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.