🏎 Smaller, faster, cheaper, lighter: Introducing DistilBERT, a distilled version of BERT
By: Victor Sanh
Tags: Machine Learning, Distillation, Transformers, NLP, Bert

At Hugging Face, we experienced first-hand the growing popularity of these models as our NLP library — which encapsulates most of them — got installed more than 400,000 times in just a few months.

However, as these models were reaching a larger NLP community, an important and challenging question started to emerge. How should we put these monsters in production? How can we use such large models under low latency constraints? Do we need (costly) GPU servers to serve at scale?

For many researchers and developers, these can be deal-breaking issues 💸

To build more privacy-respecting systems, we noticed an increasing need to have machine learning systems operate on the edge rather than calling a cloud API and sending possibly private data to servers. Running models on devices like your smartphone 📲 also requires light-weight, responsive and energy-efficient models!

Last but not least, we are more and more concerned about the environmental cost of scaling exponentially computing requirements of these models.

So, how can we reduce the size of these monster models⁉️

There are many techniques available to tackle the previous questions. The most common tools include quantization (approximating the weights of a network with a smaller precision) and weights pruning (removing some connections in the network). For these technics, you can have a look at the excellent blog post of Rasa on quantizing BERT.

We decided to focus on distillation: a technique you can use to compress a large model, called the teacher, into a smaller model, called the student.

⚗️ Knowledge Distillation — Transferring generalization capabilities

Knowledge distillation (sometimes also referred to as teacher-student learning) is a compression technique in which a small model is trained to reproduce the behavior of a larger model (or an ensemble of models). It was introduced by Bucila et al. and generalized by Hinton et al. a few years later. We will follow the latter method.

In supervised learning, a classification model is generally trained to predict a gold class by maximizing its probability (softmax of logits) using the log-likelihood signal. In many cases, a good performance model will predict an output distribution with the correct class having a high probability, leaving other classes with probabilities near zero.

But, some of these “almost-zero” probabilities are larger than the others, and this reflects, in part, the generalization capabilities of the model.

For instance, a desk chair might be mistaken with an armchair but should usually not be mistaken with a mushroom. This uncertainty is sometimes referred to as the “dark knowledge” 🌚

Another way to understand distillation is that it prevents the model to be too sure about its prediction (similarly to label smoothing).

Here is an example to see this idea in practice. In language modeling, we can easily observe this uncertainty by looking at the distribution over the vocabulary. Here are the top 20 guesses by BERT for completing this famous quote from the Casablanca movie:

The top 20 guesses from BERT (base) for the masked token. The Language model identified two highly probable tokens (day & life) followed by a long tail of valid tokens.

👯‍♂️ How can we copy this dark knowledge?

In the teacher-student training, we train a student network to mimic the full output distribution of the teacher network (its knowledge).

We are training the student to generalize the same way as the teacher by matching the output distribution.

Rather than training with a cross-entropy over the hard targets (one-hot encoding of the gold class), we transfer the knowledge from the teacher to the student with a cross-entropy over the soft targets (probabilities of the teacher). Our training loss thus becomes:

With t the logits from the teacher and s the logits of the student

This loss is a richer training signal since a single example enforces much more constraint than a single hard target.

To further expose the mass of the distribution over the classes, Hinton et al. introduce a softmax-temperature:

T is the temperature parameter.

When T → 0, the distribution becomes a Kronecker (and is equivalent to the one-hot target vector), when T →+∞, it becomes a uniform distribution. The same temperature parameter is applied both to the student and the teacher at training time, further revealing more signals for each training example. At inference, T is set to 1 and recover the standard Softmax.

🗜Hands-on coding in PyTorch — Compressing BERT

We want to compress a large language model using distilling. For distilling, we’ll use the Kullback-Leibler loss since the optimizations are equivalent:

When computing the gradients with respect to q (the student distribution) we obtain the same gradients. It allows us to leverage PyTorch implementation for faster computation:

A Knowledge distillation training step in PyTorch. Copy the gist from here.

Using the teacher signal, we are able to train a smaller language model, we call DistilBERT, from the supervision of BERT 👨‍👦 (we used the English bert-base-uncased version of BERT).

Following Hinton et al., the training loss is a linear combination of the distillation loss and the masked language modeling loss. Our student is a small version of BERT in which we removed the token-type embeddings and the pooler (used for the next sentence classification task) and kept the rest of the architecture identical while reducing the numbers of layers by a factor of two.

Overall, our distilled model, DistilBERT, has about half the total number of parameters of BERT base and retains 95% of BERT’s performances on the language understanding benchmark GLUE.