Knowledge Distillation › Knowledge Distillation

Modern language models keep growing. GPT-3 has 175 billion parameters. GPT-4 is estimated to be considerably larger. These models perform remarkably well, but deploying them at scale is expensive: each inference call touches hundreds of gigabytes of weights, requires high-end GPUs, and adds latency that users notice. You cannot run a 175B parameter model on a smartphone. You cannot serve it cheaply for millions of requests per day without a datacenter budget that most teams simply do not have.

The instinctive response to this problem is compression: make the model smaller. But naive size reduction, like simply removing layers or pruning weights, destroys the model's learned representations. The accuracy collapse is severe. Something important is being lost, and it is not immediately obvious what. The problem is not just that the model has fewer weights. The problem is that those weights were arranged in a specific way that encoded a rich, interlocking structure of knowledge. Removing pieces of that structure without care collapses the whole.

Knowledge distillation answers a different question: rather than compressing a model, can we train a small model to behave like a large one? Instead of shrinking GPT-3, can we train a much smaller model that has absorbed what GPT-3 knows? The large model becomes a teacher, and the small model becomes a student. The student learns from labeled data and from the teacher's outputs, specifically from the soft probability distributions the teacher assigns to each possible output.

The key insight, due to Geoffrey Hinton, Oriol Vinyals, and Jeff Dean in their 2015 paper "Distilling the Knowledge in a Neural Network," is that a model's output distribution carries far more information than a hard label. When a classifier says a handwritten "7" is 73% likely to be a 7, 18% likely to be a 1, and 5% likely to be a 4, those secondary probabilities encode structural knowledge about which concepts are similar. A student trained on these soft targets learns richer representations than one trained only on the correct label "7." The teacher's uncertainty is itself instructive, because it reveals the geometry of the problem space.

In the context of language models, the teacher's output distribution over a vocabulary of 50,000 tokens at each position communicates which words are semantically interchangeable, which completions are plausible, and what syntactic patterns are expected. A student trained to match this distribution inherits a compressed form of the teacher's understanding of language, without needing to replicate all 175 billion of the teacher's parameters.

This chapter develops the distillation framework from first principles. We start with the mathematical objective, work through temperature scaling in detail, examine how to choose and evaluate teacher models, and build a working implementation. We then scale these ideas to the LLM setting, examining DistilBERT, TinyBERT, and instruction-following distillation as concrete case studies. The chapter closes with an honest assessment of where distillation succeeds and where it falls short.