# DistilBERT: A Distilled Version of BERT

Source: https://zilliz.com/learn/distilbert-distilled-version-of-bert
Author: Haziqa Sajid
Published: 2024-09-22

Advancements in [transformer-based language models](https://zilliz.com/glossary/transformer-models) have significantly changed natural language processing ([NLP](https://zilliz.com/learn/A-Beginner-Guide-to-Natural-Language-Processing)). [BERT](https://zilliz.com/learn/what-is-bert), or Bidirectional Encoder Representations Transformers, is one of the most influential language models, introduced in 2018. Its bidirectional nature reads text in both directions, which makes it different from earlier models that processed text in only one direction. Because of this, BERT achieves state-of-the-art results in various NLP tasks.

However, one major disadvantage of [BERT](https://zilliz.com/learn/explore-bertopic-novel-neural-topic-modeling-technique) is its resource-intensive nature. Its base model has 110 million parameters, while the larger version has 340 million. Training and deploying BERT on [edge devices](https://zilliz.com/blog/unstructured-data-processing-from-cloud-to-edge) require a lot of memory and computational resources. 

To overcome these problems, [**DistilBERT**](https://arxiv.org/pdf/1910.01108) was introduced as a smaller, faster, and **distilled version of BERT**. It maintains 97% of BERT's language understanding capabilities while being 40% small and 60% faster. This article will discuss the methodologies used in DistilBERT development, including the training approach and its results. We will also see how to implement DistilBERT with Hugging Face.


## What is DistilBERT?

**DistilBERT** pre-trains a smaller and faster general-purpose language representation model, which can be fine-tuned with good performance on many downstream tasks.

The primary motivation behind DistilBERT is the need for compact and more efficient models for NLP applications that operate in real-time on edge devices. [Language models](https://zilliz.com/glossary/large-language-models-(llms)) like BERT often require substantial computational resources, limiting deployment in resource-constrained environments.

To address these challenges, DistilBERT focuses on the following key objectives:

- **Computational Efficiency:** While BERT requires more computational resources to operate due to its large number of parameters, ****DistilBERT reduces the size of a BERT model by [40%](https://arxiv.org/pdf/1910.01108). It requires less computation and time, which is especially useful when working with large datasets.

- **Faster Inference Speed:** BERT's complexity leads to slow inference times. DistilBERT addresses this problem by being smaller and optimized for speed and giving [60%](https://arxiv.org/pdf/1910.01108) faster inference times compared to BERT. On-device applications, such as mobile [question-answering apps](https://zilliz.com/blog/building-intelligent-chatbot-with-nlp-and-milvus), DistilBERT is [71%](https://arxiv.org/pdf/1910.01108) faster than BERT.

- **Comparable Performance:** Although DistilBERT is much smaller, it retains[ 97%](https://arxiv.org/pdf/1910.01108) of BERT’s accuracy on popular [NLP benchmarks](https://zilliz.com/learn/popular-datasets-for-natural-language-processing). This balance between size reduction and minimal performance degradation makes it a solid alternative to BERT.


## How Does DistilBERT Work?

Understanding the underlying methodologies that enable the DistilBERT model to deliver significant efficiency gains without sacrificing performance is important.

DistilBERT's key methodologies are the following:

- Knowledge Distillation 

- Triple Loss Function

- Architecture Modification


### Knowledge Distillation

Knowledge distillation, introduced by [Hinton et al. in 2015](https://arxiv.org/abs/1503.02531), is the compression and core technique behind DistilBERT. This technique trains a smaller, simpler model (the student) by replicating knowledge from a larger, pre-trained model (the teacher). Here, BERT serves as the teacher, and DistilBERT is the student. Generally, the distillation procedure is applied during the pertaining but can also be used during the fine-tuning.

Consider an example to understand knowledge distillation. Suppose a large, [state-of-the-art image classification](https://zilliz.com/learn/image-embeddings-for-enhanced-image-search) model, like [ResNet-50](https://zilliz.com/blog/combine-ai-models-for-image-search-using-onnx-and-milvus), which has been trained on a large dataset. We can use this model as the teacher network. Now, we want to deploy a smaller model on resource-constrained devices like smartphones without sacrificing accuracy. We can train a smaller student network like MobileNet via knowledge distillation. 

The student network learns from the teacher network soft probabilities and minimizes the cross-entropy between the student's predicted distribution. Soft probabilities refer to a model's output probabilities before converting them into hard, discrete class predictions. This improves performance and reduces memory and computation needs.

![Figure: Teacher-student model knowledge distillation](https://assets.zilliz.com/Figure_Teacher_student_model_knowledge_distillation_c309acdc2a.png)


In this illustration, the teacher network trains the student network by giving it soft probabilities. This helps the student network learn well and replicate the teacher network's behavior.


### Triple Loss Function in Knowledge Distillation

The knowledge distillation process in DistilBERT uses the triple loss function. This loss function helps the student model (DistilBERT) learn from the knowledge and patterns of the teacher model (BERT). But before discussing loss functions in detail, it's necessary to understand the concept of softmax temperature (temperature in softmax function). This hyperparameter controls the smoothness of the output probabilities of the teacher model. 

We usually use a [softmax function](https://zilliz.com/learn/decoding-softmax-understanding-functions-and-impact-in-ai) in standard [neural networks](https://zilliz.com/learn/Neural-Networks-and-Embeddings-for-Language-Models) to convert the model output scores into probabilities. The softmax function is defined as:

Softmax(z)=f(zi​)= exp(z) / sum(exp(z))

However, Softmax Temperature introduces a small variation to the standard softmax function, where the model outputs are divided by 𝑇 (temperature) before applying the softmax function. This is done to reduce the difficulty level for the student model, as the probability distribution learned by the teacher is often very close to 0 or 1, making it harder for the student to learn. The formula for Softmax temperature is:

Softmax(zi, T)=f(zi, T)= exp(z/T) / sum(exp(z/T))

The temperature 𝑇 controls the smoothness of the probability distribution:

- When 𝑇 = 1, this is equivalent to the standard softmax.

- When 𝑇 > 1, the output probabilities are smoother.

- When 𝑇 < 1, the output probability distribution becomes more rough.

In other words, the Higher the temperature, the smoother the probability distribution.

![Figure: Illustration of a neural network producing different probability distributions based on the temperature 𝑇](https://assets.zilliz.com/Figure_Illustration_of_a_neural_network_producing_different_probability_distributions_based_on_the_temperature_9fb5942c11.png)

[Source](https://towardsdatascience.com/distilbert-11c8810d29fc) 

Now, let's look at the loss functions used in DistilBERT one by one. These functions help DistilBERT learn from BERT and update its weights.

- Masked language modeling loss _Lmlm_

- Distillation loss Lce

- Cosine embedding loss Lcos


### Masked language modeling loss (MLM)

Like BERT, DistilBERT learns language patterns by using a technique called masked language modeling (MLM). In MLM, certain words in a sentence are randomly replaced with a special token, and the model learns to predict these masked words based on the surrounding context. 

The [cross-entropy loss](https://zilliz.com/learn/Cross-Entropy-Loss-Unraveling-its-Role-in-Machine-Learning) between the predicted and actual distributions is calculated, and the weights of the student’s model are updated through [backpropagation](https://zilliz.com/glossary/backpropagation).


![Figure: Illustration of Masked language modeling loss calculation](https://assets.zilliz.com/Figure_Illustration_of_Masked_language_modeling_loss_calculation_75c02e8ca4.png)

[Source](https://towardsdatascience.com/distilbert-11c8810d29fc)

### Distillation loss

Distillation loss is essential in training the student model during knowledge distillation. It's calculated by comparing the student's prediction probability to the soft probabilities provided by the teacher model.

Lce= iti*log(si)

Where ti the probability is estimated by the teacher and si the probability is estimated by the student.

While it is technically possible to use only the student loss to train the student model, relying only on student loss is not optimal because of the softmax transformation when the temperature 𝑇 is set to 1. When 𝑇 = 1, the resulting probability distribution is likely to be very skewed, with one label having a very high probability close to 1, while the probabilities for all other labels approach 0. This behavior becomes challenging when many classification labels may be valid for the same input.

To address this issue, DistilBERT uses a distillation loss, where the softmax probabilities are calculated with a temperature 𝑇 > 1. This change makes the output probabilities more evenly distributed, which helps the student model to consider multiple possible answers. 

![Figure: Illustration of Distillation loss calculation](https://assets.zilliz.com/Figure_Illustration_of_Distillation_loss_calculation_874c06bd64.png)

[Source](https://towardsdatascience.com/distilbert-11c8810d29fc)


### Cosine embedding loss

Cosine embedding loss is another important component used in the DistilBERT training process. Its primary purpose is to ensure that the student model learns to replicate the teacher model's internal knowledge representations ([embeddings](https://zilliz.com/learn/everything-you-should-know-about-vector-embeddings)) by minimizing the cosine distance between the embeddings generated by the student and teacher models. 

The [cosine similarity](https://zilliz.com/blog/similarity-metrics-for-vector-search) can be computed as the negative of the cosine similarity between the teacher’s and student’s generated embeddings(h).

Lcos= 1 - hteacher . hstudenthteacher hstudent

![Figure: Cosine embedding loss calculation](https://assets.zilliz.com/Figure_Cosine_embedding_loss_calculation_45465360c6.png)

 [Source](https://towardsdatascience.com/distilbert-11c8810d29fc)


### Triple loss

Finally, the loss function in DistilBERT is determined by adding the linear combinations of the three loss functions. Based on the loss value, the student model undergoes backpropagation to update its weights.

Loss = Lmlm+  Lce+ Lcos

![Figure: DistilBERT loss function](https://assets.zilliz.com/Figure_Distil_BERT_loss_function_61e774bfd1.png)

[Source](https://towardsdatascience.com/distilbert-11c8810d29fc)


### DistilBERT Architecture

The student model, DistilBERT, has the same general architecture as BERT but with some modifications. The token-type [embeddings](https://docs.zilliz.com.cn/reference/python/EmbeddingModels) and the pooler are removed, and the number of layers is reduced by a factor of 2. BERT Base has 12 layers, while DistilBERT has only 6. These architectural changes make DistilBERT more efficient, and further improvements were made during training, as explained next.

- **Student Initialization:** A critical part of training DistilBERT is determining how to initialize the student model so it learns well. Since both BERT (the teacher) and DistilBERT (the student) share the same dimensionality, the student is initialized by taking one layer out of two from the teacher model.

- **Distillation:** DistilBERT was distilled in very large batches (up to 4K examples per batch) using dynamic masking. The next sentence prediction (NSP) objective was removed so the model could focus more on understanding the masked words in each sentence. This concept was taken from [RoBERTa](https://zilliz.com/blog/roberta-optimized-method-for-pretraining-self-supervised-nlp-systems) to provide a more robust performance.

- **Data and Compute Power:** DistilBERT was trained on the same dataset as the original BERT model, a concatenation of [English Wikipedia](https://paperswithcode.com/dataset/wiki-en) (2500M words) and [Toronto Book Corpus](https://github.com/sgraaf/Replicate-Toronto-BookCorpus) (800M words), and on 8 16GB V100 GPUs for approximately 90 hours.


![Diagram of BERT BASE and DistilBERT model architecture](https://assets.zilliz.com/Diagram_of_BERT_BASE_and_Distil_BERT_model_architecture_facb5e7639.png)

[Source](https://www.researchgate.net/figure/Schematic-diagram-of-BERT-BASE-and-DistilBERT-model-architecture_fig1_382939584)


## Experimental Results

DistilBERT's language understanding and generalization capabilities were evaluated on the [General Language Understanding Evaluation](https://gluebenchmark.com/) (**GLUE**) benchmark, a collection of 9 datasets for evaluating natural language understanding systems. DistilBERT's performances were also evaluated on downstream tasks, such as sentiment classification ([**IMDb**](https://paperswithcode.com/sota/sentiment-analysis-on-imdb)) ****and question-answering tasks ([**SQuAD v1.1**](https://datarepository.wolframcloud.com/resources/SQuAD-v1.1) - Stanford Question Answering Dataset). 


### GLUE

![Figure: Comparison of the dev sets of the GLUE benchmark](https://assets.zilliz.com/Figure_Comparison_of_the_dev_sets_of_the_GLUE_benchmark_98148c1ef8.png)

[Source](https://arxiv.org/pdf/1910.01108)

DistilBERT achieves a score of 77.0 on the GLUE benchmark and is always on par with or improving over the ELMo baseline provided by the GLUE authors (up to 19 points of accuracy on the STS-B - Semantic Textual Similarity Benchmark). DistilBERT also compares to BERT, retaining 97% of the performance with 40% fewer parameters.


### IMDb & SQuAD

![Figure: DistilBERT yields comparable performance on downstream tasks](https://assets.zilliz.com/Figure_Distil_BERT_yields_comparable_performance_on_downstream_tasks_b96297dd06.png)

[Source](https://arxiv.org/pdf/1910.01108)

DistilBERT achieved results comparable to those of BERT on downstream tasks. It is only 0.6% behind BERT in test accuracy on the IMDb benchmark, while it is 40% smaller. On SQuAD, DistilBERT is within 3.9 points of the full BERT.


### Size and Inference Speed

![Figure: DistilBERT inference time and model parameters](https://assets.zilliz.com/Figure_Distil_BERT_inference_time_and_model_parameters_7c22057079.png)

[Source](https://arxiv.org/pdf/1910.01108)

DistilBERT has 40% fewer parameters than BERT and is 60% faster than BERT. The number of parameters and inference time for both models were compared on the STS-B development set, using a batch size of 1 on a CPU.


## DistillBERT Implementation with Hugging Face

DistilBERT can be implemented in Python using the Hugging Face Transformers library. Here are the steps for implementing DistilBERT. The code below is from [HuggingFace](https://huggingface.co/docs/transformers/en/model_doc/distilbert#transformers.DistilBertForSequenceClassification). 


### Step 1: Install the Transformers library

You can install the Transformers library using pip. Open your terminal and enter the following command.

```

    !pip install transformers

```


### Step 2: Load the Pre-trained DistilBERT Model

You can load the pre-trained DistilBERT model using the _DistilBertForSequenceClassification_ class.

```

    # Import libraries 
    import torch
    from transformers import AutoTokenizer,DistilBertForSequenceClassification
    #Load tokenizer and pre-trained model 
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
    model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")

```


### Step 3: Tokenize the Input Text

Tokenize the input text using _tokenizer_ function.

```

    Input_text = "Hello, my dog is cute"
    # Tokenize the input text
    inputs = tokenizer(Input_text, return_tensors="pt")

```


### Step 4: Inference using the pre-trained model on the Input Text

Finally, we get the predicted label by taking the _argmax_ of the model's output logits.

```

    with torch.no_grad():
        logits = model(**inputs).logits


    # Get predicted class ID 
    predicted_class_id = logits.argmax().item()

```

To train DistilBERT on a specific task, you need to set the number of output classes and provide appropriate training data. The following code demonstrates how to do this:

```

    # Set the number of output classes 
    num_labels = len(model.config.id2label)


    # Create a new model with the desired number of labels 
    model = DistilBertForSequenceClassification.from_pretrained(
        "distilbert-base-uncased", num_labels=num_labels
    )
    # Calculate loss 
    labels = torch.tensor([1])
    loss = model(**inputs, labels=labels).loss

```

The code demonstrates the steps for using DistilBERT for a SequenceClassification with Hugging Face, where the model predicts labels based on the input text. While the current example focuses on classification, DistilBERT, like BERT, is a general-purpose pre-trained language model that can be used to generate contextual embeddings for text. 

These embeddings can be stored in [vector databases](https://zilliz.com/learn/beginner-guide-to-implementing-vector-databases) such as [Milvus](https://milvus.io/) or [Zilliz Cloud](https://zilliz.com/cloud), designed for efficient vector retrieval and [semantic search](https://zilliz.com/glossary/semantic-search). Vector databases are optimized for handling high-dimensional vectors and are perfect for many modern AI applications, such as retrieval-augmented generation (RAG) systems, AI chatbots, recommendation engines, and natural language processing apps. 


## Conclusions 

DistilBERT made a huge step in advancing BERT by significantly compressing the model while maintaining comparable performance across various NLP tasks. Additionally, at just 207 MB, it becomes much easier to integrate on devices with limited resources. While knowledge distillation is the primary technique, DistilBERT can be further compressed through methods like [quantization](https://zilliz.com/learn/scalar-quantization-and-product-quantization) or pruning algorithms for even greater efficiency.


## Further resources 

- DistilBERT paper: [Distilled version of BERT](https://arxiv.org/pdf/1910.01108)

- [DistilBERT GitHub repository](https://github.com/huggingface/swift-coreml-transformers)

- [What is BERT (Bidirectional Encoder Representations from Transformers)?](https://zilliz.com/learn/what-is-bert)

- [RoBERTa: An Optimized Method for Pretraining Self-supervised NLP Systems](https://zilliz.com/blog/roberta-optimized-method-for-pretraining-self-supervised-nlp-systems)

- [What are Vector Databases and How Do They Work?](https://zilliz.com/learn/what-is-vector-database)

- [Build AI Apps with Retrieval Augmented Generation (RAG)](https://zilliz.com/learn/Retrieval-Augmented-Generation)

- [Generative AI Resource Hub | Zilliz](https://zilliz.com/learn/generative-ai)

- [What is a Transformer Model?](https://zilliz.com/glossary/transformer-models)

- [Semantic Search with Milvus and OpenAI](https://www.google.com/search?q=semantic+search+Zilliz&sca_esv=4d75a1f392f2e349&sca_upv=1&rlz=1C1JJTC_en-GBPK1104PK1104&sxsrf=ADLYWIJHm5QG1ryd070A9ih-2JAzKIps6w%3A1726862028229&ei=zNLtZtXQDar2i-gPg7mdkQw&ved=0ahUKEwjV5pKuptKIAxUq-wIHHYNcJ8IQ4dUDCA8&uact=5&oq=semantic+search+Zilliz&gs_lp=Egxnd3Mtd2l6LXNlcnAiFnNlbWFudGljIHNlYXJjaCBaaWxsaXoyBBAjGCcyCBAAGIAEGKIEMggQABiABBiiBDIIEAAYgAQYogQyCBAAGIAEGKIEMggQABiABBiiBEj3CVDtAVj2BXABeACQAQCYAb8CoAHyBqoBBTItMi4xuAEDyAEA-AEC-AEBmAIDoAKFBcICChAAGLADGNYEGEfCAgcQIxiwAhgnmAMAiAYBkAYFkgcHMS4wLjEuMaAHwA0&sclient=gws-wiz-serp)

- [What is Semantic Similarity?](https://zilliz.com/glossary/semantic-similarity)