DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

Victor Sanh; Lysandre Debut; Julien Chaumond; Thomas Wolf

Source-linked AI summary

DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

Victor Sanh, Lysandre Debut, Julien Chaumond, Thomas Wolf

arXiv:1910.01108v4Oct 2, 2019cs.CL

TL;DR

Large pre-trained NLP models improve downstream tasks but are costly to train and difficult to run on devices. The paper introduces DistilBERT, a distilled general-purpose BERT model, and reports 40% fewer parameters, 60% faster inference, and 97% retained language-understanding performance.

Problem
Large pre-trained NLP models improve performance but impose substantial computational, environmental, memory, and on-device deployment costs.
Source: §1 · p. 1
Method
DistilBERT uses pre-training knowledge distillation with a triple loss combining masked language modeling, distillation, and cosine embedding objectives.
Source: §1 · p. 2, §2 · p. 2
Results
40% fewer parameters and 60% faster inference are achieved while retaining 97% of BERT’s performance across varied downstream tasks.
Source: §1 · p. 2, §4 · p. 3, Size and inference speed · p. 4
Takeaways & Limitations
DistilBERT provides a smaller general-purpose model that preserves flexibility across tasks and is suitable for edge applications.
Source: §1 · p. 2, §6 · p. 5

Abstract

from arXiv · show

As Transfer Learning from large-scale pre-trained models becomes more prevalent in Natural Language Processing (NLP), operating these large models in on-the-edge and/or under constrained computational training or inference budgets remains challenging. In this work, we propose a method to pre-train a smaller general-purpose language representation model, called DistilBERT, which can then be fine-tuned with good performances on a wide range of tasks like its larger counterparts. While most prior work investigated the use of distillation for building task-specific models, we leverage knowledge distillation during the pre-training phase and show that it is possible to reduce the size of a BERT model by 40%, while retaining 97% of its language understanding capabilities and being 60% faster. To leverage the inductive biases learned by larger models during pre-training, we introduce a triple loss combining language modeling, distillation and cosine-distance losses. Our smaller, faster and lighter model is cheaper to pre-train and we demonstrate its capabilities for on-device computations in a proof-of-concept experiment and a comparative on-device study.

1 Introduction

Large-scale pre-trained NLP models improve downstream performance but create environmental, computational, memory, and on-device deployment challenges. The paper addresses these challenges with knowledge-distilled, general-purpose models that preserve downstream-task flexibility while reducing model size, inference time, and training cost.

Large-scale pre-trained language models improve NLP tasks but often contain several hundred million parameters, with larger models tending to perform better downstream.Source: §1 · p. 1
Growing computational and memory requirements increase environmental costs and may hinder real-time, on-device adoption.Source: §1 · p. 1
Knowledge distillation produces smaller general-purpose language models that retain good performance across downstream tasks, while reducing inference demands and computational training budget.Source: §1 · p. 2
40% smaller Transformer models achieve similar performance across downstream tasks while being 60% faster at inference time.The models are pre-trained through distillation under supervision from a larger Transformer language model.Source: §1 · p. 2
The triple loss combines language-modeling, distillation, and cosine-distance objectives, whose components are each important for best performance.Source: §1 · p. 2

2 Knowledge distillation

Knowledge distillation trains a compact student model to reproduce a larger teacher’s behavior using the teacher’s soft probability distribution. DistilBERT combines this distillation objective with masked language modeling and cosine alignment of hidden states.

Distillation principle

Knowledge distillation trains a compact student to reproduce the behavior of a larger teacher or model ensemble.The student is optimized to match the teacher rather than learning independently.Source: §2 · p. 2

Distillation loss

The distillation loss uses the teacher’s full soft target distribution, providing a richer training signal than one-hot labels.The student probabilities are trained against the teacher probabilities through cross-entropy.Source: §2 · p. 2

Temperature

A shared softmax temperature T smooths teacher and student output distributions during training, while inference sets T to 1.The temperature controls output-distribution smoothness.Source: §2 · p. 2

Combined objective

The final objective linearly combines distillation and masked language modeling losses with a cosine embedding loss aligning student and teacher hidden-state directions.The cosine term encourages the corresponding hidden-state vectors to point in similar directions.Source: §2 · p. 2

3 DistilBERT: a distilled version of BERT

DistilBERT keeps BERT’s general architecture while removing token-type embeddings and the pooler and halving the number of layers. It is initialized from alternating teacher layers and distilled with large-batch training on BERT’s corpus, retaining 97% of BERT’s GLUE performance.

Student architecture

DistilBERT removes token-type embeddings and the pooler while reducing the number of Transformer layers by a factor of 2.The student retains BERT’s general architecture, and the authors report that hidden-size variations are less beneficial than reducing layers.Source: §3 · p. 2

Student initialization

The student is initialized from the teacher by taking one layer out of two, leveraging their common dimensionality.This initialization is intended to help the student sub-network converge.Source: §3 · p. 2

Evaluation

97% of BERT performance is retained on GLUE development sets, with results reported as medians over 5 runs with different seeds.The comparison includes ELMo, BERT, and DistilBERT results.Source: Table 1 · p. 3

Distillation

DistilBERT is distilled with gradient accumulation, dynamic masking, no next-sentence prediction objective, and batches of up to 4K examples.The training procedure follows recently proposed best practices for BERT training.Source: §3 · p. 3

Data and compute power

DistilBERT uses English Wikipedia and the Toronto Book Corpus and trains for approximately 90 hours on 8 16GB V100 GPUs.RoBERTa is provided as a compute comparison, requiring 1 day on 1024 32GB V100 GPUs.Source: §3 · p. 3

4 Experiments

Experiments show that DistilBERT retains most of BERT’s performance across GLUE and downstream tasks while substantially reducing model size and inference time. Additional distillation improves SQuAD results, and on-device tests demonstrate practical speed and memory characteristics.

General Language Understanding

97% performance retention with 40% fewer parameters is achieved by DistilBERT relative to BERT on GLUE.Across all 9 GLUE tasks, DistilBERT is on par with or improves over the ELMo baseline.Source: §4 · p. 3

Downstream tasks

0.6% point separates DistilBERT from BERT in IMDb test accuracy, while DistilBERT is 40% smaller.On SQuAD, DistilBERT is within 3.9 points of full BERT.Source: §4.1 · p. 3
79.8 F1 and 70.4 EM are reached on SQuAD after adding adaptation-phase distillation from a BERT model fine-tuned on SQuAD.These results are within 3 points of the full model, following distillation during both pre-training and adaptation.Source: §4.1 · pp. 3–4

Speed and size trade-off

40% fewer parameters and 60% faster inference distinguish DistilBERT from BERT on a CPU STS-B development-set pass.The comparison uses an Intel Xeon E5-2690 v3 CPU at 2.9GHz with batch size 1.Source: Size and inference speed · p. 4

On device computation

71% faster inference than BERT is measured for DistilBERT on an iPhone 7 Plus, excluding tokenization, with a whole-model size of 207 MB.The experiment evaluates a mobile question-answering application.Source: Size and inference speed · p. 4

Ablation study

The two distillation losses account for a large portion of GLUE performance, whereas removing the Masked Language Modeling loss has little impact.The ablation study reports macro-score deltas relative to the full triple loss.Source: §4.2 · p. 4

5 Related work

Prior work mainly explores task-specific or multi-teacher distillation, while this work favors general-purpose pre-training distillation. Other compression approaches include pruning and quantization, which are orthogonal to the proposed method.

Task-specific distillation

Prior distillation studies mainly build task-specific setups, whereas this work finds general-purpose pre-training distillation beneficial.Examples include transferring fine-tuned BERT classification models to LSTM classifiers and distilling fine-tuned SQuAD models into smaller Transformers.Source: §5 · p. 4

Multi-distillation

Multi-distillation combines knowledge from teacher ensembles through multi-task learning to regularize distillation and has been applied to compact question-answering models.Tsai et al. also pre-train a multilingual model from scratch solely through distillation.Source: §5 · p. 4

Other compression techniques

Other compression techniques include pruning self-attention heads and quantization to derive smaller models; both are orthogonal to this work.Pruning can remove some self-attention heads at test time without significantly degrading performance, and some layers can be reduced to one head.Source: §5 · p. 4

6 Conclusion and future work

The paper introduces DistilBERT as a smaller, faster general-purpose BERT model that retains most of BERT’s language understanding capabilities. It also shows that distillation can train general-purpose language models and supports DistilBERT for edge applications.

Conclusion

DistilBERT is 40% smaller and 60% faster than BERT while retaining 97% of its language understanding capabilities.These results establish its efficiency and retained performance relative to BERT.Source: §6 · p. 5
The study demonstrates that distillation can successfully train a general-purpose language model.The authors also analyzed the method’s components through an ablation study.Source: §6 · p. 5

Future work

DistilBERT is presented as a compelling option for edge applications.This extends the model’s relevance beyond conventional fine-tuning settings.Source: §6 · p. 5