Source-linked AI summary
FastBERT: a Self-distilling BERT with Adaptive Inference Time
Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Haotang Deng, Qi Ju
TL;DR
BERT-like models are costly and slow to run under constrained or fluctuating deployment demands. FastBERT combines sample-wise adaptive inference with self-distillation to adjust computation, achieving strong efficiency-accuracy trade-offs across twelve NLP datasets. Its speedup is tunable across a broad range, with accuracy loss controlled by the tolerated threshold.
Problem
BERT-like models have high computation and slow inference, making deployment difficult under limited resources and changing request volumes.
Method
FastBERT dynamically adjusts executed layers by sample and uses self-distillation during training for adaptive, efficient inference.
Results
Across twelve NLP datasets, FastBERT is 2 to 3 times faster than BERT without performance degradation, with speedup tunable between 1 and 12 times when more accuracy loss is tolerated.
Takeaways & Limitations
FastBERT offers a speed-tunable BERT-like model for balancing inference efficiency and accuracy under varying deployment demands.
Takeaways & Limitations
Future work includes extending FastBERT to other pre-training architectures and a wider range of NLP tasks.
Abstract
from arXiv · showhide
Pre-trained language models like BERT have proven to be highly performant. However, they are often computationally expensive in many practical scenarios, for such heavy models can hardly be readily implemented with limited resources. To improve their efficiency with an assured model performance, we propose a novel speed-tunable FastBERT with adaptive inference time. The speed at inference can be flexibly adjusted under varying demands, while redundant calculation of samples is avoided. Moreover, this model adopts a unique self-distillation mechanism at fine-tuning, further enabling a greater computational efficacy with minimal loss in performance. Our model achieves promising results in twelve English and Chinese datasets. It is able to speed up by a wide range from 1 to 12 times than BERT if given different speedup thresholds to make a speed-performance tradeoff.
1 Introduction
FastBERT addresses the practical cost and fixed-computation limitations of BERT-like models with sample-wise adaptive inference and self-distillation. Across twelve English and Chinese NLP tasks, it targets adjustable speed and competitive accuracy.
- Motivation: BERT-like models improve NLP performance but incur high computation and slow inference, limiting deployment under constrained resources.The introduction highlights fluctuating request volumes and the need to reserve substantial server capacity for peak periods.
- Motivation: Knowledge distillation reduces computation but typically requires a separate student model and trades task accuracy for efficiency.The paper also identifies fixed model computation as inadequate for inputs with varying difficulty.
- Proposal: FastBERT dynamically adjusts the number of executed layers for each sample and combines this adaptive mechanism with self-distillation.The design aims to avoid overcomputing simple inputs while retaining capacity for complex samples.
- Evidence: Twelve English and Chinese NLP tasks show substantial computation reduction with very little accuracy loss.The experiments cover six Chinese and six English tasks.
- Proposal: FastBERT provides a practical speed-tunable BERT model that balances inference speed and accuracy as request amounts vary.Its speed can be adjusted to match changing operational demands.
2 Related work
Related work covers larger pre-trained language models, knowledge distillation, and adaptive computation. FastBERT is positioned against approaches that increase model capacity, use separate students, or adapt computation in other modalities.
- Pre-trained language models: BERT and later models such as RoBERTa, T5, UER, and K-BERT improve accuracy through larger corpora, longer training, scaling, or added knowledge.These approaches generally use heavier settings or more data.
- Pre-trained language models: BERT-base contains 110 million parameters across twelve Transformer blocks, while BERT-large uses 24 layers, increasing inference cost.ALBERT reduces model size by sharing parameters, but the paper still identifies redundant computation as a concern.
- Knowledge distillation: Knowledge-distillation methods transfer information from heavy teacher models to lighter students through approaches such as intermediate-layer extraction, staged learning, or multiple losses.Examples include PKD-BERT, TinyBERT, and DistilBERT.
- Knowledge distillation: Classic distillation commonly uses a separate student structure whose effectiveness depends on the teacher and may reduce performance while improving speed.The paper contrasts this setup with its integrated approach.
- Adaptive inference: Prior adaptive-computation methods operate token-wise or image-patch-wise, while adaptive mechanisms had not been applied to NLP pre-training models for efficiency.The cited approaches add recurrent token steps or adjust layers within image regions.
3 Methodology
FastBERT fuses adaptive inference and self-distillation into a single speed-up approach. The paper presents this combination as targeting competitive accuracy and efficiency together.
- Core design: FastBERT fuses adaptive inference and self-distillation into a novel approach for improving both accuracy and efficiency.The combined design is illustrated in Figure 2.
3.1 Model architecture
FastBERT combines a 12-layer Transformer backbone with teacher and student classifiers to support sample-wise early exits. Student branches estimate uncertainty at intermediate layers, allowing simple samples to stop early while complex samples continue.
- Backbone: FastBERT uses a 12-layer Transformer encoder with a teacher classifier in its backbone.The backbone includes an embedding layer, Transformer blocks, and a teacher classifier for downstream inference.
- Adaptive inference: Samples with low uncertainty exit the batch, whereas higher-uncertainty samples proceed to deeper Transformer layers.This sample-wise mechanism assigns different numbers of executed layers according to input complexity.
- Backbone: The embedding layer converts each input sentence into vector representations by combining word, position, and segment embeddings.These representations initialize the layer-by-layer Transformer feature extraction.
- Branches: Student classifiers are attached to Transformer outputs to provide early predictions, particularly for simple samples.Their architecture is designed to balance prediction accuracy against inference speed.
3.2 Model training
FastBERT trains its backbone and student classifiers in separate stages, freezing one module while training the other. Self-distillation then trains student classifiers to match the backbone teacher using unlabeled data and KL-divergence losses.
- Training procedure: FastBERT training proceeds through backbone pre-training, backbone fine-tuning, and self-distillation for student classifiers.The backbone and student modules are trained separately, with one module frozen while the other is updated.
- Backbone training: FastBERT can directly use BERT-like pre-training methods or load an existing high-quality pre-trained model.The teacher classifier remains unaffected during backbone pre-training.
- Backbone training: For each downstream task, the backbone and teacher classifier are fine-tuned on task-specific data while student classifiers remain disabled.Student branches are enabled later for self-distillation.
- Self-distillation: Student classifiers are trained by comparing their predictions with the teacher’s soft-label using KL-divergence.The total self-distillation loss sums the divergences across all L − 1 student classifiers.
- Self-distillation: Self-distillation can use unlimited unlabeled data because it requires only the teacher’s output, and teacher and student outputs come from the same model.This avoids additional pre-training structures used by conventional teacher-student distillation.
3.3 Adaptive inference
FastBERT performs adaptive inference by routing samples through different numbers of Transformer layers according to uncertainty and a tunable Speed threshold. Its acceleration relies on lightweight classifiers enabling early exits while preserving a speed–accuracy trade-off.
- Adaptive inference: FastBERT measures whether each sample’s current inference is credible enough to terminate at every Transformer layer.This sample-wise mechanism adjusts the number of executed encoding layers according to sample complexity.
- Adaptive inference: Speed is the threshold distinguishing high and low uncertainty and serves as a halt value for balancing inference accuracy and efficiency.Higher Speed sends fewer samples to higher layers and increases overall inference speed.
- Adaptive inference: Samples with uncertainty below the Speed threshold exit early, while higher-uncertainty samples continue to the next layer.Uncertainty is computed from a student classifier’s output probability distribution using normalized entropy.
- Computational basis: FastBERT reports adjustable accuracy and FLOPs outcomes under multiple Speed settings, including speedups ranging from 2.02x to 11.74x in the listed results.The table includes Speed values of 0.1, 0.5, and 0.8.
4 Experimental results
The experiments evaluate FastBERT on twelve NLP datasets spanning Chinese and English classification and sentence-matching tasks. They measure accuracy and computational complexity while comparing FastBERT with BERT and DistilBERT baselines.
- Datasets: FastBERT is evaluated on twelve NLP datasets, comprising six Chinese and six English tasks.The Chinese collection includes classification and sentence-matching datasets, while the English collection contains six sentence-classification datasets.
- Metrics: The evaluation reports accuracy and sample-averaged FLOPs under different Speed values.FLOPs measure the number of floating-point operations performed for a single process and indicate computational complexity.
- FLOPs analysis: FastBERT’s acceleration is grounded in the classifier requiring much less computational load than the Transformer.This supports reducing Transformer computation even after adding classifiers.
- Baselines: The comparisons use BERT-base and DistilBERT baselines, including DistilBERT variants with three and one layers.BERT-base has twelve layers, while the cited DistilBERT baseline has six layers.
4.3 Performance comparison
FastBERT evaluates accuracy and computational efficiency across adjustable inference speeds, using uncertainty to determine layer execution and self-distillation to reduce computation. Experiments analyze classifier uncertainty, exit-layer distributions, and convergence on the Book Review dataset.
- Performance comparison: FastBERT reports accuracy and sample-averaged FLOPs for BERT, DistilBERT, and FastBERT under different Speed values.Speedup is calculated with BERT as the benchmark.
- Performance comparison: At Speed = 0.1, FastBERT achieves 2 to 5 times speedup without accuracy loss for most datasets.
- Performance comparison: With a small tolerated accuracy loss, FastBERT reaches 7 to 11 times speedup, with its speedup ratio adjustable between 1 and 12.
- Uncertainty analysis: Classifier accuracy follows the LUHA pattern: lower uncertainty corresponds to higher accuracy for classifiers at the bottom, middle, and top of FastBERT.The analysis uses Student-Classifier0, Student-Classifier5, and Teacher-Classifier on the Book Review dataset.
- Exit-layer distribution: At Speed = 0.8, 61% of samples exit after the first Transformer layer, eliminating unnecessary calculations in the next eleven layers.
- Self-distillation convergence: Self-distillation decreases FLOPs while accuracy increases during fine-tuning and then remains nearly unchanged as computation is reduced.On the Book Review dataset, FastBERT is fine-tuned for three epochs and self-distilled for five more epochs.
4.6 Ablation study
Ablation studies on Book Review and Yelp.P show that adaptive inference and self-distillation are both important for maintaining accuracy while accelerating FastBERT.
- Ablation study: FastBERT without self-distillation or adaptive inference performs worse at almost the same speedup.
- Ablation study: When acceleration exceeds five times, downstream accuracy degrades dramatically without adaptive inference.
- Ablation study: The ablations support adaptation and self-distillation as key mechanisms for achieving speedups with low accuracy loss.
5 Conclusion
FastBERT combines adaptive inference during prediction with self-distillation during training to improve efficiency while preserving accuracy. Across twelve NLP datasets, it supports tunable speedups and remains compatible with other BERT-like pretrained models.
- Conclusion: FastBERT combines self-distillation during training with adaptive inference during prediction to improve efficiency with less accuracy loss.
- Conclusion: Across twelve NLP datasets, FastBERT is 2 to 3 times faster than BERT without performance degradation.
- Conclusion: When more accuracy loss is tolerated, FastBERT can tune its speedup between 1 and 12 times.
- Conclusion: FastBERT can load publicly available BERT-like models, including BERT-WWM, ERNIE, and RoBERTa, for initialization.
6 Future work
Future work targets improving the speed-speedup relationship, extending FastBERT to other pretrained architectures, and evaluating it on additional NLP tasks.
- Future work: Future work includes linearizing the Speed-Speedup curve and extending FastBERT to XLNet and ELMo.
- Future work: The approach is also proposed for wider NLP coverage, including named entity recognition and machine translation.