Knowledge Distillation and Pruning: Getting a Small Model to Do a Large Model’s Job

Models Larger Than Their Task Requires

The most accurate neural networks are usually the largest ones, and large models are expensive to run. They need memory that a phone or a microcontroller does not have, they answer slowly, and at the scale of millions of requests per day, every extra millisecond and watt adds to the bill.

Researchers noticed early that trained networks are heavily redundant. Much of what makes a network large is needed to train it well, not to represent what it learned. That observation led to a family of techniques, grouped under the name model compression, that start from a large trained model and produce a smaller one with nearly the same accuracy. The three main ones are knowledge distillation, pruning and quantization. This post focuses on the first two.

Teachers and Students

The core idea of distillation appeared in 2006, when Cristian Bucilă, Rich Caruana and Alexandru Niculescu-Mizil published a paper called «Model Compression.» They had large ensembles of models that were accurate but slow, and they showed that a single small neural network could be trained to imitate the ensemble’s predictions and keep most of its accuracy.

Geoffrey Hinton, Oriol Vinyals and Jeff Dean gave the method its current form and name in «Distilling the Knowledge in a Neural Network» (2015). A large, already trained model is the teacher, and a smaller model is the student. The student is trained to reproduce the teacher’s outputs, not only the correct labels from the dataset.

Soft Targets and Temperature

Imitating a teacher works better than learning from labels alone because the teacher’s output contains more information than the label does. A label says an image is a dog and nothing else. A trained classifier’s output might say 90% dog, 9% wolf, 0.9% cat, and almost zero for truck. Those small probabilities encode which classes look alike, something the label alone never tells the student. Hinton called this «dark knowledge.»

In a well-trained network, though, the wrong-class probabilities are often so tiny that they barely affect the loss. The fix is a temperature parameter T in the softmax. Dividing the logits by a temperature above 1 flattens the distribution, so the small probabilities become large enough to matter. The student is trained with the same temperature to match these softened outputs, usually alongside the ordinary loss on the true labels, with a weight balancing the two. At inference the temperature goes back to 1.

Distilling Internal Representations

Matching only the final outputs gives the student a limited signal. FitNets, from Adriana Romero and colleagues (ICLR 2015), added «hints»: the student is also trained to reproduce the teacher’s activations at an intermediate layer, through a small adapter that accounts for the difference in size. This made it possible to train students that were thinner and deeper than the teacher.

Later methods transfer other internal quantities, such as attention maps or the relationships between examples in a batch. In Transformers, it is common to match hidden states and attention matrices layer by layer, as TinyBERT (Jiao and colleagues, 2020) does.

DistilBERT

The best-known example in language processing is DistilBERT, published in 2019 by Victor Sanh, Lysandre Debut, Julien Chaumond and Thomas Wolf at Hugging Face. The student has the same architecture as BERT-base with half the layers, 6 in place of 12, and was initialized by taking one layer out of every two from the teacher. It was trained with a combination of three losses: the usual masked language modeling loss, a distillation loss on the teacher’s softened output distribution, and a cosine loss that aligns the direction of the student’s hidden states with the teacher’s.

The paper reports a model 40% smaller and 60% faster than BERT that retains 97% of its language understanding performance as measured on the GLUE benchmark. DistilBERT became one of the most downloaded models on the Hugging Face Hub, largely because it could run on CPUs at reasonable speed.

Pruning

Pruning removes parts of a trained network. The idea dates back to 1989, when Yann LeCun, John Denker and Sara Solla published «Optimal Brain Damage,» which used second-derivative information to estimate which weights could be deleted with the least effect on the error.

The modern version is simpler. In 2015, Song Han and colleagues showed that removing the weights with the smallest absolute values, then retraining the remaining ones to recover accuracy, and repeating the cycle, could remove most of the connections in networks like AlexNet and VGG-16 without loss of accuracy. Their follow-up, «Deep Compression,» which won a best paper award at ICLR 2016, chained pruning with weight quantization and Huffman coding. It reduced the storage of AlexNet by 35 times, from 240 MB to 6.9 MB, and of VGG-16 by 49 times, from 552 MB to 11.3 MB, with no loss of accuracy.

Unstructured and Structured Pruning

There is an important practical distinction. Unstructured pruning deletes individual weights anywhere in the network. It reaches very high sparsity, but the result is a matrix of the same shape with many zeros in it, and ordinary hardware does not run a sparse matrix any faster unless the software and chips are built for it. The file gets smaller, and the speed often does not change.

Structured pruning removes whole units: neurons, convolutional filters, attention heads or entire layers. The resulting network is smaller in the ordinary sense and runs faster on any hardware. The cost is that accuracy drops sooner, because removing a whole filter is a much coarser operation than removing scattered weights. Paul Michel and colleagues showed in 2019 that many attention heads in Transformers can be removed at test time with little effect, which made head pruning a common technique for language models.

The Lottery Ticket Hypothesis

If a pruned network works well, a natural question is whether the small network could have been trained directly from the start. For a long time the answer seemed to be no: sparse networks trained from scratch did worse than pruned ones.

Jonathan Frankle and Michael Carbin proposed an explanation in «The Lottery Ticket Hypothesis,» which received a best paper award at ICLR 2019. They trained a network, pruned the smallest weights, and then reset the surviving weights to the exact values they had at initialization. Trained again from there, these sparse subnetworks, often 10 to 20% of the original size on MNIST and CIFAR-10, matched the accuracy of the full network. With a fresh random initialization, the same sparse structure did not.

Their hypothesis is that a large random network contains small subnetworks whose initial weights happen to be well suited for training, the winning tickets, and that part of the benefit of a large network is that it holds many such candidates. On larger networks, later work found that the weights need to be rewound to a point early in training instead of to the very beginning. Finding the tickets still requires training the big network first, so the result is more important for understanding why over-parameterized networks train well than as a practical shortcut.

Combining the Techniques

Distillation, pruning and quantization act on different things, so they are usually combined. Distillation produces a smaller architecture. Pruning removes what is left of the redundancy. Quantization then stores each remaining weight with fewer bits, for example 8-bit integers in place of 32-bit floating-point numbers. A typical pipeline for deployment on a device applies all three, checking accuracy after each step.

Distillation has also taken on a second role with large language models. A large model can generate answers, explanations or labels for a big set of inputs, and a much smaller model is then trained on that synthetic data. The student never sees the teacher’s internal probabilities, only its text, but the principle is the one from 2006: use an expensive model to teach a cheap one.

Limits

Compression always costs some accuracy. Average accuracy on a benchmark can stay almost unchanged while performance on rare cases gets worse. Sara Hooker and colleagues showed in 2019 that pruning affects a small subset of examples disproportionately, often atypical or underrepresented ones, even when the top-line accuracy barely moves. A student can also inherit its teacher’s errors and biases, since it is trained to copy them. And when the student is far smaller than the teacher, the gap can be too wide to bridge in one step, which is why some methods distill through an intermediate-size model first.

Reflection

Compression works because training a model and running it need different things. Training benefits from a large network, which has many possible paths to a good solution. Once the solution is found, most of that capacity is no longer needed. This is why the common practice is to train a large model first and then shrink it, and not to design a small model and train it directly. The lottery ticket experiments are the clearest evidence for this: the small network that works was inside the large one from the start, but it could only be identified after training the large one.

By: Max Johnson B.

Comentarios

Deja una respuesta

Tu dirección de correo electrónico no será publicada. Los campos obligatorios están marcados con *