DistilBERT: A Distilled Version of BERT for Efficient NLP

The BERT model, released by Google AI researchers in 2018, marked a major milestone in natural language processing (NLP). By pre-training on a massive corpus of unlabeled text, BERT learned rich linguistic representations that could then be fine-tuned for a wide variety of downstream NLP tasks, achieving state-of-the-art results.

However, the power of BERT came with a cost: the model was very large, with hundreds of millions of parameters. This made BERT computationally expensive to train and run inference on, limiting its practical applicability, especially in resource-constrained environments like mobile devices or real-time services.

To address this challenge, researchers at Hugging Face proposed DistilBERT in 2019 as a compressed version of BERT that retains most of its knowledge and capabilities while being significantly faster and lighter. In this article, we‘ll take a deep dive into how DistilBERT works and how it compares to the original BERT.

Knowledge Distillation: The Key to DistilBERT

At the heart of DistilBERT is the technique of knowledge distillation. The basic idea is to "distill" or transfer knowledge from a large, complex model (the teacher) to a smaller, simpler model (the student). By learning to mimic the outputs of the teacher model, the student model can achieve comparable performance while being more efficient.

In general, knowledge distillation works by adding a loss term to the student model‘s training objective that minimizes the difference between the student‘s predictions and the teacher‘s predictions. This forces the student to not just learn to predict the correct labels, but to emulate the teacher‘s full output distribution, capturing more of its knowledge.

There are a few key benefits to knowledge distillation:

1. The student model is typically much smaller than the teacher, in terms of number of parameters and architectural complexity. This leads to lower memory usage, faster inference, and cheaper training.

2. The student model learns both from the training data and from the outputs of the teacher model. This allows the student to pick up additional knowledge and generalization capabilities that may not be apparent from the raw data alone.

3. The teacher model can be an ensemble of multiple models, enabling the student to learn from a diverse set of "experts". This has been shown to boost the student‘s performance.

Distilling BERT into DistilBERT

The DistilBERT model follows the same general Transformer architecture as BERT, but with a few key differences:

• DistilBERT has 40% fewer parameters than BERT-base. This is achieved by reducing the number of layers in the Transformer encoder from 12 to 6, while keeping the hidden size and feed-forward dimension the same.

• The token-type embeddings and pooler layers are removed. The researchers found these to be less important for many downstream tasks.

• The number of attention heads is kept at 12, same as BERT-base. This was found to be important for maintaining comparable performance.

To distill knowledge from BERT into this more compact architecture, a few techniques are used:

1. Larger batches: DistilBERT is trained with a batch size of 4K, enabled by gradient accumulation, compared to a batch size of 256 for BERT. This provides more negative examples within each batch, leading to better distillation.

2. Dynamic masking: Instead of performing masking once during data preprocessing as in BERT, DistilBERT generates different masks for each input sequence on the fly. This results in more diverse examples over the course of training.

3. Removing the next sentence prediction objective: The researchers found the NSP loss proposed in the original BERT model was not beneficial for DistilBERT, so they removed this term from the training loss.

4. Cosine embedding loss: In addition to the standard distillation loss that matches the student‘s output distribution to the teacher‘s, DistilBERT adds a cosine embedding loss that aligns the directions of the student‘s and teacher‘s hidden activations. This was found to boost performance on several downstream tasks.

5. Initialization from BERT layers: Each layer of DistilBERT is initialized from one of the layers of the BERT-base checkpoint, taking every other layer. This provides a good initialization point that "warm-starts" the distillation process.

DistilBERT‘s Performance

So how well does DistilBERT stack up against its teacher model? The researchers evaluated DistilBERT on the GLUE benchmark, a collection of diverse NLU tasks, as well as question answering with SQuAD.

On GLUE, DistilBERT achieves 97% of BERT-base‘s performance on average while having 40% fewer parameters. It consistently outperforms the previous state-of-the-art models like OpenAI GPT and ELMo.

On SQuAD 1.1, DistilBERT reaches 86.9 F1, which is 3.9 points behind BERT-base but still significantly better than non-BERT models.

Inference speed is where DistilBERT really shines. On a CPU, DistilBERT is around 60% faster than BERT-base. On a mobile device (a 2018 iPad Pro), DistilBERT is 71% faster while having a much smaller memory footprint, making it feasible to run on-device.

Using DistilBERT in Practice

Thanks to the Transformers library from Hugging Face, using a pre-trained DistilBERT model in Python is quite straightforward. First, install the library:

pip install transformers

Then, you can load a DistilBERT model and its associated tokenizer like this:

from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

tokenizer = DistilBertTokenizer.from_pretrained(‘distilbert-base-uncased‘)
model = DistilBertForSequenceClassification.from_pretrained(‘distilbert-base-uncased‘)

Here we‘re loading the pre-trained DistilBERT model from the Hugging Face Model Hub. The tokenizer is responsible for preprocessing raw text into a sequence of token IDs that can be fed into the model. The model itself is an instance of DistilBertForSequenceClassification, which adds a classification head on top of the base DistilBERT model for tasks like sentiment analysis.

To run the model on some text:

text = "I love this movie! The acting was great and the plot kept me engaged from start to finish."
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)

The model outputs raw logits, which you can softmax to get class probabilities:

import torch.nn.functional as F

probs = F.softmax(outputs.logits, dim=-1)

And to get the predicted class:

predicted_class_id = probs.argmax().item()
print("Predicted class:", model.config.id2label[predicted_class_id])

This will print out the predicted sentiment for the input text (e.g. "POSITIVE").

The Future of Efficient NLP Models

DistilBERT showcases the potential of knowledge distillation and model compression techniques to make large NLP models more practical and accessible. However, it‘s just one step in an ongoing journey towards more efficient NLP architectures.

Since the release of DistilBERT, there have been many other efforts to build compact, fast models that retain the power of BERT and its successors. For example:

• MobileBERT from Microsoft uses a more advanced knowledge distillation setup to further reduce model size while maintaining performance.
• TinyBERT from Huawei employs a two-stage distillation process that transfers knowledge at both the pre-training and task-specific fine-tuning stages.
• ALBERT from Google uses cross-layer parameter sharing and factorized embeddings to achieve a significant reduction in parameters without losing much performance.

In the future, we can expect to see even more innovation in this space as NLP models continue to grow in size and capability. Some promising directions include:

• Improved distillation techniques that better capture the knowledge in large models
• Novel architectures designed from the ground up for efficiency, such as sparse Transformers or models with adaptive computation
• Combining model compression with other approaches like quantization and pruning for further speedups and memory savings
• Automated methods for architecture search and compression, like neural architecture search and AutoML

As these techniques advance, the gap between large, powerful NLP models and efficient, deployable ones will continue to narrow, unlocking new applications for natural language technology.

Ultimately, the goal is to democratize NLP by making state-of-the-art models accessible to everyone, regardless of their computational resources or technical expertise. DistilBERT is an important milestone on this path, demonstrating that it‘s possible to have our cake and eat it too: to enjoy the benefits of large-scale language models without the prohibitive costs. As the field progresses, we can look forward to a future where natural language understanding is not just powerful, but truly universal.