Source-linked AI summary
Compressing Large-Scale Transformer-Based Models: A Case Study on BERT
Prakhar Ganesh, Yao Chen, Xin Lou, Mohammad Ali Khan, Yin Yang, Hassan Sajjad, Preslav Nakov, Deming Chen, Marianne Winslett
TL;DR
Large Transformer models achieve strong NLP performance but are costly in memory, computation, energy, and deployment resources. This paper systematically surveys and compares compression methods for BERT, analyzes their mechanisms and best practices, and identifies open research directions. Across tasks, quantization can reduce BERT to 15% or 10.2% of its original size with accuracy drops of 0.6% or 0.9%, while practical speedups require specialized hardware or libraries.
Problem
Large Transformer models are too resource-hungry for low-capability devices and latency-constrained applications, while comprehensive comparisons of compression methods for Transformer-based NLP models were lacking.
Method
The paper conducts a thorough comparative study of BERT compression methods, surveys best practices, and analyzes how the methods work.
Results
Quantization reduces BERT to 15% and 10.2% of its original size with accuracy drops of 0.6% and 0.9%, respectively, across various tasks.
Takeaways & Limitations
Traditional methods such as quantization and pruning benefit BERT, while BERT-specific distillation and architectural methods provide competitive results and architectural insight.
Takeaways & Limitations
Quantization and unstructured pruning do not provide run-time speedups on standard devices without specialized hardware or libraries.
Abstract
from arXiv · showhide
Pre-trained Transformer-based models have achieved state-of-the-art performance for various Natural Language Processing (NLP) tasks. However, these models often have billions of parameters, and, thus, are too resource-hungry and computation-intensive to suit low-capability devices or applications with strict latency requirements. One potential remedy for this is model compression, which has attracted a lot of research attention. Here, we summarize the research in compressing Transformers, focusing on the especially popular BERT model. In particular, we survey the state of the art in compression for BERT, we clarify the current best practices for compressing large-scale Transformer models, and we provide insights into the workings of various methods. Our categorization and analysis also shed light on promising future research directions for achieving lightweight, accurate, and generic NLP models.
1 Introduction
Transformer-based pre-training delivers strong NLP performance but creates substantial resource demands, motivating model compression. The paper surveys and compares compression methods for BERT while identifying best practices and future research directions.
- Motivation: Pre-training generic Transformer models on large corpora followed by task-specific fine-tuning benefits diverse NLP tasks.The paper describes this workflow for sentiment analysis, paraphrase detection, reading comprehension, question answering, and summarization.
- Motivation: GPT-3 has over 175 billion parameters, causing high memory consumption, computational overhead, and energy use.These demands are especially problematic for lower-capacity devices and latency-sensitive applications.
- Motivation: A single GPT-3 training run costs well above 12 million USD after training on 300 billion tokens.The paper also notes that fine-tuning or inference cannot run on a 32GB Tesla V100 GPU.
- Model compression: Model compression addresses resource demands, and prior work suggests compressing over-parameterized models can reduce performance errors compared with directly training smaller models.The paper lists pruning, quantization, and knowledge distillation among established compression approaches.
- Scope and contribution: The paper provides a systematic comparative study of Transformer compression, focusing on BERT while noting that methods may extend to decoders and multilingual models.The discussion is restricted to BERT to provide detailed insight into the methods.
- Scope and contribution: The study summarizes compression techniques and best practices while highlighting future directions for lightweight, accurate, and generic NLP models.Its timeliness is linked to growing BERT-like model use and limited access to expensive GPUs or capable on-device hardware.
2 Breakdown & Analysis of BERT
BERT combines token, segment, and position embeddings with stacked Transformer encoder units containing self-attention and feed-forward sub-units. Breakdown experiments show that feed-forward networks dominate model size and computation, while practical costs depend strongly on hardware.
- BERT architecture: BERT is pretrained on Wikipedia and BookCorpus using masked language modeling and next sentence prediction.Masked language modeling learns sentence context, while next sentence prediction learns relationships between two sentences.
- BERT architecture: WordPiece tokenization segments complex words into subwords, reducing vocabulary size and improving robustness to out-of-vocabulary words.BERT adds a [CLS] token whose output supports classification tasks.
- BERT architecture: BERT sums token, segment, and position embeddings before passing representations through the Transformer backbone to an application-dependent output layer.The final layer can be a classifier for tasks such as sentiment analysis.
- BERT architecture: Each BERT encoder contains self-attention and feed-forward network sub-units with residual connections, and architecture is specified by L, H, and A.L is the number of encoder units, H is embedding-vector size, and A is the number of attention heads.
- Breakdown analysis: Figure 3 compares model size, theoretical FLOPs, run-time memory, and inference latency for BERTBASE on GPU and CPU hardware.The experiments use 256-token sentences, an Nvidia Titan X GPU, and an Intel Xeon CPU.
- Breakdown analysis: FFN sub-units consume the most memory and execute the highest number of FLOPs, while the embedding layer is substantial in size but requires zero inference FLOPs.The embedding layer is a lookup table, whereas attention incurs computation from dot products and softmax operations.
- Breakdown analysis: Attention-adjacent linear layers add relatively small overhead compared with FFN sub-units, with the pre-attention layer roughly three times larger than the post-attention layer.The difference reflects three attention inputs—key, value, and query—but one output.
- Hardware-dependent costs: GPU execution takes 57.1 ms per example versus 750.9 ms on CPU, while GPU run-time memory exceeds CPU-only memory because tensors are duplicated.In practice, multi-head self-attention is costlier than theoretical estimates because it combines several matrix transformations, multiplication, and softmax.
3 Compression Methods
BERT compression methods target different components because no single existing technique addresses the model’s self-attention, linear layers, embedding size, and depth together. The survey covers quantization, pruning, knowledge distillation, attention decomposition, and embedding compression, alongside methods that trade model size against inference speed.
- Combining Methods: Compression methods target different BERT components and can be combined when their implementations are orthogonal.The survey distinguishes model-size reduction and inference-speed methods from training-procedure compression.
- Quantization: Quantization reduces the number of unique weight values and required bits, lowering memory footprint and potentially improving low-precision inference speed.Naïve truncation can cause quantization noise and accuracy loss; quantization-aware training adjusts quantized weights to retain accuracy.
- Pruning: Pruning removes redundant weights or components through unstructured weight removal or structured reduction of encoder units, embedding size, attention heads, and other blocks.Attention-head pruning can retain high accuracy with only 1–2 heads per encoder unit, compared with the original model’s 16 heads.
- Knowledge Distillation: Knowledge distillation trains a smaller student from outputs of one or more larger teacher models, reducing embedding size H, encoder count L, or the Transformer backbone.Reducing H requires transformations to reconcile different intermediate-output dimensions, while reducing L makes student units learn from multiple teacher units.
- Attention Decomposition: Attention decomposition reduces computation by calculating attention in smaller groups, projecting key-query matrices to lower dimensions, or restricting softmax computation to top-k values.Because multi-head self-attention contains no weights, these methods improve runtime memory and execution speed rather than model size.
- Embedding Matrix Compression: Embedding matrix compression targets a lookup table comprising about 21% of BERT’s complete model size by reducing vocabulary size V from the original roughly 30k.A 5k vocabulary reproduces 94% of tokens created with a 30k vocabulary, supporting vocabulary reduction as a compression option.
4 Effectiveness of the Compression Methods
Compression methods trade off model size, inference speed, and accuracy differently. Quantization and distillation generally offer strong accuracy-size trade-offs, while hardware and embedding treatment strongly affect practical gains.
- Quantization and Pruning: Quantization reduces BERT to 15% or 10.2% of its original size with only 0.6% or 0.9% average accuracy loss across tasks.Its architecture-invariant design preserves the model’s components and connections while reducing weight precision.
- Quantization and Pruning: Unstructured pruning reaches 67.6% of the original size without accuracy loss, but frozen embeddings limit more extreme compression.Even 3% encoder weight density leaves the total model at 23.8% of the original size.
- Practical Advice: Quantization and unstructured pruning usually reduce model size without runtime speedups on standard devices, requiring specialized hardware or libraries for execution gains.A 1.3% model achieves only 1.83x GPU speedup and a 1.53% average accuracy drop, increasing to 2.1x with specialized accelerators.
- Structured Pruning: Structured pruning compresses BERT to 39.1% of its original size with a 1.86% average accuracy drop when depth and width are reduced together.Reducing hyperparameters in harmony performs better than focusing on only one.
- Distillation from Attention Maps: 60.7% model size with only 0.1% average accuracy loss is achieved by distilling attention maps alone, versus 1.73% loss from other distillation targets.The comparison indicates that intermediate attention maps are an important BERT distillation target.
- Combining Multiple Distillations: 13.3% model size and 9.4x GPU speedup are achieved with multiple distillation targets at a 1.0% accuracy drop.An even smaller 1.6% student is approximately 25x faster but loses 12.3% accuracy.
5 Open Issues & Research Directions
The analysis finds that established compression techniques and BERT-specific methods can produce competitive results, but BERT compression remains at an early stage. It identifies layer-wise, backbone, and compound compression as promising research directions.
- Traditional quantization and pruning methods provide benefits for compressing BERT.
- Layer-wise compression: Layer-wise compression is promising because BERT compression methods are often coupled across encoder units and inner architecture.Methods that compress each layer independently have shown promising results but remain under-explored.
- Alternative backbones: Replacing Transformer backbones with BiLSTMs or CNNs achieves high compression ratios but causes sizable accuracy drops.The paper suggests exploring more complex and hybrid BiLSTM/CNN/Transformer backbones to limit performance loss.
- Compound compression: Combining complementary compression methods can improve overall compression performance beyond individual methods.The analysis reports that compound compression methods perform better than their individual counterparts.