Source-linked AI summary
DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
Victor Sanh, Lysandre Debut, Julien Chaumond, Thomas Wolf
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.
Method
DistilBERT uses pre-training knowledge distillation with a triple loss combining masked language modeling, distillation, and cosine embedding objectives.
Results
40% fewer parameters and 60% faster inference are achieved while retaining 97% of BERT’s performance across varied downstream tasks.
Takeaways & Limitations
DistilBERT provides a smaller general-purpose model that preserves flexibility across tasks and is suitable for edge applications.
Abstract
from arXiv · showhide
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.
- Growing computational and memory requirements increase environmental costs and may hinder real-time, on-device adoption.
- Knowledge distillation produces smaller general-purpose language models that retain good performance across downstream tasks, while reducing inference demands and computational training budget.
- 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.
- The triple loss combines language-modeling, distillation, and cosine-distance objectives, whose components are each important for best performance.
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.
- 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.
- 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.
- 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.
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.
- 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.
- 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.
- 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.
- 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.
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.
- 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.
- Downstream tasks: 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.
- 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.
- 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.
- 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.
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.
- 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.
- 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.
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.
- Conclusion: 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.
- Future work: DistilBERT is presented as a compelling option for edge applications.This extends the model’s relevance beyond conventional fine-tuning settings.