Skip to content
ssunil.dev ssunil.dev

// writing

Teaching Small Models to Punch Above Their Weight

In August 2024, Philipp Schmid of Hugging Face posted on X about NVIDIA’s Mistral NeMo Minitron 8B. The headline: NVIDIA had taken a 12-billion parameter model, pruned it, distilled it, and produced an 8-billion parameter model that outperformed both Mistral-7B and LLaMA-3.1-8B across a suite of benchmarks. Using 380 billion tokens for distillation instead of training from scratch. The AI community treated this as breaking news.

Jeff Dean’s reply was essentially: we published this in 2015.

He was not wrong. In 2015, Geoffrey Hinton, Oriol Vinyals, and Dean himself published a paper called Distilling the Knowledge in a Neural Network. It is one of those papers that changed everything and got cited by everyone but somehow never made it into the popular conversation. Which is a shame, because the core idea is one of the most elegant tricks in all of machine learning, and you do not need a PhD to understand it.

Here is the setup. You have a very large, very accurate neural network. It is brilliant. It is also enormous. It lives on a server rack and drinks electricity like it has somewhere to be. You cannot put it on a phone. You cannot put it on an edge device. You cannot put it anywhere that does not have its own power substation and a cooling budget.

You also have a small neural network. It is fast, cheap, and fits on a microcontroller. But it is not very smart. Train it on the same data and it produces worse results. This is not a surprise. Smaller models have less capacity. They cannot memorize as many patterns. The conventional wisdom in 2015 was: if you want better accuracy, you need a bigger model. End of story.

Hinton’s insight was that this is not actually the end of the story. It is just the end of the obvious part.

What the teacher knows that the labels don’t

Here is what makes the paper clever.

When you train a neural network to classify images, you give it labels. This is a cat. This is a dog. This is a truck. The labels are binary. A picture is either a cat or it is not. There is no room for nuance.

But the large model, after training, does not think in binaries. Ask it about a picture of a cat and it will say: 92% cat, 4% tiger, 2% dog, 1% lynx, and a scattering of negligible probabilities across everything else. That 4% for tiger is not a mistake. It is information. The model has learned that this particular cat looks a bit like a tiger. Maybe the stripes, maybe the posture, maybe something no human would notice but the model picked up from seeing a million examples.

Hinton called this “dark knowledge.” Not the answer the model gives, but the uncertainty around the answer. The almost-right guesses. The relationships between wrong answers. A cat that looks like a tiger is fundamentally different from a cat that looks like a toaster, and the large model knows this even though the label just says “cat” in both cases.

The question is: can you transfer that dark knowledge to the small model?

The temperature trick

The answer is yes, and the mechanism is almost disappointingly simple.

Neural networks produce their predictions through a function called softmax, which takes a bunch of raw numbers and converts them into probabilities that add up to one. The problem is that softmax is aggressive. If one number is even slightly larger than the rest, softmax pushes its probability close to 1 and crushes everything else toward zero. All that nuanced dark knowledge gets flattened into near-certainty.

Softmax with Temperature

For the mathematically inclined: softmax converts raw model outputs (logits) into probabilities. The temperature parameter T divides each logit before exponentiation. At T=1, the winner takes almost all the probability mass. Raise T and the distribution softens, exposing the relative confidence across every class. The formula is simple. The consequences are not.

The fix is a single parameter called temperature. Raise the temperature and the softmax relaxes. Instead of 92% cat, 4% tiger, and zeros everywhere else, you get something like 40% cat, 20% tiger, 15% dog, 10% lynx. The relative relationships are preserved, but now the small model can actually see them.

The training recipe is straightforward. Run the training data through the large model at high temperature. Capture those soft, nuanced probability distributions. Then train the small model to match them, also at high temperature. The small model is not learning “this is a cat.” It is learning “this is a cat, but it looks a bit like a tiger, and nothing at all like a truck.” That is a much richer learning signal than a binary label could ever provide.

A common misconception: distillation does not use different data. It uses the same data, but with a different target. Instead of the original labels, the student trains on the teacher’s soft predictions. The result is that the student needs far fewer training examples to reach the same accuracy, because each example carries more information. This is exactly what NVIDIA exploited with Minitron. They width-pruned the MLP and embedding layers of Mistral NeMo 12B, left the attention heads untouched, and distilled the result using 380 billion tokens. Training a comparable model from scratch would have taken trillions. Same quality. A fraction of the compute. The soft targets are simply a more efficient way to learn.

When you are done training, you set the temperature back to 1 and deploy the small model normally. It is still small. It is still fast. But it is meaningfully smarter than it would have been if you had just trained it on the original labels. The large model’s years of education, compressed into a form factor that fits in your pocket.

Why this is not just an academic trick

Knowledge distillation is the reason your phone can run AI features without melting.

Every time you use on-device speech recognition, face detection, real-time translation, or any of the local AI features that Apple, Google, and Samsung ship on their phones, you are almost certainly using a distilled model. The cloud has the enormous teacher. Your phone has the student. The student was trained on the teacher’s soft knowledge, not just the raw data. That is why it works as well as it does despite running on a chip that draws less power than a reading lamp.

This matters even more at the edge. Not the phone edge. The real edge. Industrial sensors running inference on microcontrollers with kilobytes of memory. Autonomous drones making decisions at thirty frames per second on a power budget measured in milliwatts. Medical devices that need to classify signals locally because there is no network connection in an operating room.

In all of these cases, the constraint is the same: you cannot ship the big model. You have to ship something smaller. And distillation is how you make something smaller without making it stupid.

The part that connects to everything else

Knowledge distillation is not an isolated technique. It sits at the center of a web of ideas that are reshaping how AI gets deployed.

Quantization takes a model and reduces the precision of its numbers. Instead of 32-bit floating point, you use 8-bit integers, or 4-bit, or even binary. You lose some accuracy. But if you distill first and then quantize, the accuracy loss is significantly smaller. The distilled model has learned better internal representations, which makes it more robust to the brutal rounding that quantization demands.

Pruning takes a model and removes the connections that contribute least to the output. A neural network might have millions of parameters, but most of them are barely doing anything. Remove them and the model gets smaller and faster. Again, distillation first, prune second. The teacher’s knowledge acts as a guide, ensuring the pruned model preserves what matters.

And then there is hardware. Custom AI accelerators, the kind I work on, are designed for specific model architectures running at specific precisions. A distilled, quantized, pruned model is not just smaller. It maps better onto the hardware. Fewer parameters means less memory access. Lower precision means cheaper arithmetic. Sparser connections mean less wasted computation. The efficiency gains are multiplicative, and they start with distillation.

This is not a coincidence. It is a stack.

What Hinton saw in 2015

The paper is over a decade old. It was presented at the NIPS 2014 Deep Learning Workshop, not a main conference. The math is approachable. The experiments are on MNIST, which is the “hello world” of machine learning datasets. Nothing about it screams “this will change the industry.”

But Hinton saw something that most people missed. He saw that the knowledge in a neural network is not in its parameters. It is in the relationships between its predictions. And those relationships can be separated from the architecture that produced them. A large model’s understanding of the world can be poured into a smaller vessel, and most of it survives the transfer.

That insight turned out to be one of the most consequential ideas in applied machine learning. Not because it made models more accurate. Because it made accuracy portable.

Every edge device running a distilled model is proof that intelligence does not have to be expensive. That the cost of training and the cost of inference are fundamentally different problems. And that a good teacher, even an enormous and impractical one, is worth building if the students it produces can go places it never could.

The large model lives in the data center. The small model lives in the world. And the bridge between them is a temperature parameter and a paper from 2015 that most people have never read.


Share this post on:

Previous Post
The Hidden Cost of Zero-Cost Abstractions