Source-linked AI summary
Compressing BERT: Studying the Effects of Weight Pruning on Transfer Learning
Mitchell A. Gordon, Kevin Duh, Nicholas Andrews
TL;DR
BERT is effective but costly, motivating the question of whether compression preserves transfer learning and whether fine-tuning changes compressibility. The paper applies magnitude weight pruning during pre-training and evaluates transfer across downstream tasks, finding that 30-40% pruning has no observed cost while stronger pruning creates distinct degradation regimes. It concludes that BERT can be pruned once during pre-training rather than separately for each task.
Problem
BERT’s size creates costly memory and deployment requirements, raising whether compression preserves transfer learning and whether fine-tuning changes its compressibility.
Method
The paper applies magnitude weight pruning to English BERT and evaluates pruned models on downstream GLUE tasks, including tests after fine-tuning.
Results
30-40% pruning does not increase pre-training loss or affect downstream transfer; medium pruning disrupts information transfer, while high pruning can also prevent downstream fitting.
Takeaways & Limitations
BERT can be pruned before distribution without affecting its universality, rather than being compressed separately for each task.
Takeaways & Limitations
The paper notes that pruning-related sparsity may depend on dropout or the SGD training regime, and leaves this explanation for future work.
Abstract
from arXiv · showhide
Pre-trained universal feature extractors, such as BERT for natural language processing and VGG for computer vision, have become effective methods for improving deep learning models without requiring more labeled data. While effective, feature extractors like BERT may be prohibitively large for some deployment scenarios. We explore weight pruning for BERT and ask: how does compression during pre-training affect transfer learning? We find that pruning affects transfer learning in three broad regimes. Low levels of pruning (30-40%) do not affect pre-training loss or transfer to downstream tasks at all. Medium levels of pruning increase the pre-training loss and prevent useful pre-training information from being transferred to downstream tasks. High levels of pruning additionally prevent models from fitting downstream datasets, leading to further degradation. Finally, we observe that fine-tuning BERT on a specific task does not improve its prunability. We conclude that BERT can be pruned once during pre-training rather than separately for each task without affecting performance.
1 Introduction
BERT-style pre-training improves downstream performance but creates substantial computational and deployment costs. This paper studies whether magnitude pruning preserves transfer learning and whether fine-tuning changes BERT’s compressibility.
- BERT-style pre-trained feature extractors improve deep learning performance across many NLP tasks without requiring additional labeled data.
- BERT’s scale makes training and deployment costly, with BERT-Large requiring a Google TPU and BERT-Base needing hardware optimizations.Training BERT-Base from scratch costs approximately $7k and emits approximately 1438 pounds of CO2.
- The paper asks whether compressing BERT impedes transfer to new tasks and whether fine-tuning makes it more or less compressible.
- The authors compress English BERT with magnitude weight pruning and evaluate transfer on the diverse GLUE benchmark.Magnitude pruning removes weights close to zero.
- 30-40% pruning leaves pre-training loss and downstream transfer unchanged, whereas medium pruning disrupts transferred information and high pruning can also prevent downstream fitting.Tasks degrade at different rates as pre-training loss increases.
- Fine-tuning on a specific task does not meaningfully improve BERT’s prunability or change the pruning order.
- BERT can therefore be pruned before distribution rather than separately for each downstream task without affecting universality.The paper also considers over-pruning for accuracy-memory trade-offs on certain tasks.
2 Pruning: Compression, Regularization, Architecture Search
The paper frames pruning as compression, regularization, and sparse architecture search, then focuses on magnitude-based removal of near-zero weights. Its procedure progressively removes weights and continues training to recover accuracy.
- Compression: Pruning can reduce model parameters, storage, and GPU usage, enabling deployment on edge devices and potentially increasing inference speed when structures are removed.
- Regularization: Pruning also acts as permanent dropout or an L0-style regularizer, narrowing the range of functions the network can express.
- Sparse Architecture Search: As sparse architecture search, pruning makes SGD-induced sparsity explicit and relates to lottery-ticket subnetworks.
- Magnitude Weight Pruning: Magnitude pruning removes weights close to zero because their inputs are effectively ignored.
- Magnitude Weight Pruning: The procedure selects a target sparsity, thresholds weight magnitudes, removes weights, continues training, and optionally repeats at higher sparsity.
- Pruning Scope: Global pruning thresholds all parameters together, whereas matrix-local pruning thresholds each weight matrix separately; this work uses matrix-local pruning.
- Sparse Architecture Search: The authors note that sparse networks are difficult to train from scratch, while related methods allow SGD to search over subnetworks.
3 Experimental Setup
The experiments prune BERT-Base’s attention, feed-forward, and embedding parameters, continue pre-training at controlled sparsities, and evaluate transfer across selected GLUE tasks. Additional interventions separate information deletion from regularization and test pruning after fine-tuning.
- BERT Pruning: BERT-Base has 12 encoder layers, each with six prunable matrices covering self-attention and the feed-forward network.
- BERT Pruning: Each self-attention layer is pruned separately for key, value, and query projection matrices, along with the output projection.
- BERT Pruning: Word embeddings are pruned because near-zero values can be sparsified, and embeddings account for approximately 21% of BERT-Base’s memory.
- Pruning During Pre-Training: The study evaluates sparsities from 0% to 90%, gradually pruning over 10k steps and continuing pre-training for another 90k steps on English Wikipedia and BookCorpus.
- Downstream Evaluation: Pruned models are fine-tuned for three epochs on five GLUE tasks with four learning rates, using the best evaluation accuracies averaged for Figure 1.
- Downstream Evaluation: All downstream experiments fine-tune BERT weights in every layer rather than using BERT only as a static feature extractor.
- Control Experiments: A zeroing control allows deleted weights to vary during downstream training, separating pre-training information deletion from the regularization effect of pruning.
- Pruning After Fine-tuning: To test post-fine-tuning pruning, models are fine-tuned for three epochs, pruned at various sparsities, and trained for five additional epochs or seven at 80/90% sparsity.
4 Pruning Regimes
Pruning reveals three regimes: up to 30–40% sparsity has no observed effect, medium pruning removes transferable pre-training information, and high pruning can additionally restrict downstream fitting. The degradation reflects both information deletion and, at high sparsity, complexity restriction, with task sensitivity varying by dataset and task.
- 4.1 30-40% of Weights Are Discardable: 30-40% of weights can be pruned without affecting pre-training loss or downstream inference.The authors interpret this as pruning weights that SGD has effectively pushed toward zero.
- 4.2 Medium Pruning Levels Prevent Information Transfer: Past 40% pruning, pre-training loss rises, hidden-layer features diverge, and downstream accuracy begins to degrade.The feature divergence is measured relative to models with low pruning levels.
- 4.2 Medium Pruning Levels Prevent Information Transfer: Between 40-60% sparsity, information deletion is the main cause of performance degradation because pruning and information deletion reduce performance by the same amount.The pre-training dataset contains more information than labeled downstream datasets, supporting more robust representations.
- 4.2 Medium Pruning Levels Prevent Information Transfer: The main compression obstacle is preserving the pre-training inductive bias, which requires more weights than fitting downstream datasets.The paper attributes recovery limits to an information gap between pre-training and downstream datasets.
- 4.3 High Pruning Levels Also Prevent Fitting Downstream Datasets: At 70% sparsity and above, complexity restriction becomes a secondary degradation cause, while information-deleted models recover some but not all evaluation accuracy.Recovery depends on downstream data size: information deletion performs better on MNLI and QQP than on SST-2 and CoLA.
- 4.2 Medium Pruning Levels Prevent Information Transfer: Pre-training loss linearly predicts downstream accuracy after information deletion, but each task loses accuracy at a different rate.The paper reports that an extra 30% pruning is worth one QQP accuracy point but 10 CoLA points.
5 Downstream Fine-tuning Does Not Improve Prunability
Downstream fine-tuning does not substantially change which weights magnitude pruning selects, so task-specific fine-tuning provides little improvement in BERT’s prunability.
- Fine-tuned models do not surpass the development accuracies of models pruned during pre-training despite achieving similar training losses.
- Weights move only 0-4% on average from their original magnitude-sorting positions across downstream tasks.
- Higher-magnitude weights are more stable during downstream fine-tuning than lower-magnitude weights.
- After an epoch, weights quickly re-stabilize in a nearby sorting order, and longer pre-pruning training has only a marginal effect on selected weights and accuracy.
6 Related Work
Related work examines alternative compression strategies and different forms of redundancy in BERT, while this paper emphasizes pruning’s interaction with transfer across many tasks.
- Alternative methods seek to discard non-task-specific information through information bottlenecks or knowledge distillation.
- Pruning may carry a hidden cost of decreased model robustness.
- Prior attention-head pruning found that up to 40% of BERT’s attention heads could be removed after MNLI fine-tuning without affecting test accuracy.
- Weight pruning does not disproportionately remove any individual attention head relative to the whole-model pruning rate.
- The paper suggests that attention-head pruning and weight pruning may remove different redundancies, but whether recovery training changes head prunability remains unclear.
7 Conclusion And Future Work
The conclusion identifies BERT’s inductive bias as the central compression concern and argues that moderate magnitude pruning can preserve universality without task-specific pruning.
- Encoding BERT’s inductive bias requires more weights than fitting downstream data.
- Future work should preserve inductive bias and quantify its relevance across tasks during accuracy-memory trade-offs.
- Thirty to forty percent of BERT’s weights can be discarded without affecting its universality, while the remaining weights matter differently across tasks.
- Fine-tuning downstream tasks does not change the nature of the pruning trade-off by substantially changing which weights are selected.
- Generalization to other pre-trained language models remains to be shown in future work.
A Appendix
The appendix provides detailed views of pruning levels, sorting-order stability, attention-head allocation, and where pruned weights occur within BERT’s parameters.
- Table 1 reports pre-training development losses and GLUE development accuracies, pairing each accuracy with achieved training loss across pruning levels.
- Weight magnitude-sorting positions move 0-4% on average after downstream fine-tuning, with larger datasets and learning rates producing more movement per epoch.
- Higher-magnitude weights move less in the sorting order during fine-tuning than lower-magnitude weights.
- Magnitude heatmaps show banding in attention projections and earlier pruning of shorter-subword embeddings because their magnitudes are smaller.