Source-linked AI summary
Block Pruning For Faster Transformers
François Lagunas, Ella Charlaix, Victor Sanh, Alexander M. Rush
TL;DR
Growing pre-trained transformers improve NLP performance but create larger, slower models, while pruning and distillation address size and speed imperfectly. The paper introduces Block Movement pruning, which combines block structure with movement pruning during fine-tuning. Across classification and generation tasks, it reports compact models with competitive speed and accuracy, including a 2.4x faster SQuAD model with a 1% F1 drop.
Problem
Pre-trained transformer models keep growing while pruning mainly reduces storage and distillation mainly improves inference speed.
Method
Block Movement pruning applies movement pruning to fixed-size parameter blocks during task-specific fine-tuning.
Results
Across classification and generation benchmarks, the approach achieves speedups with limited accuracy loss, including 2.4x faster SQuAD v1.1 with a 1% F1 drop.
Takeaways & Limitations
The method produces pruned models competitive with distilled models in speed and pruned models in size across diverse tasks and base models.
Takeaways & Limitations
Highly sparse movement-pruned models may not become substantially faster on standard hardware that cannot exploit sparse matrix-vector products.
Abstract
from arXiv · showhide
Pre-training has improved model accuracy for both classification and generation tasks at the cost of introducing much larger and slower models. Pruning methods have proven to be an effective way of reducing model size, whereas distillation methods are proven for speeding up inference. We introduce a block pruning approach targeting both small and fast models. Our approach extends structured methods by considering blocks of any size and integrates this structure into the movement pruning paradigm for fine-tuning. We find that this approach learns to prune out full components of the underlying model, such as attention heads. Experiments consider classification and generation tasks, yielding among other results a pruned model that is a 2.4x faster, 74% smaller BERT on SQuAD v1, with a 1% drop on F1, competitive both with distilled models in speed and pruned models in size.
1 Introduction
Pre-trained transformers improve benchmark performance but continue to grow in size and speed costs. Block pruning combines smaller models with efficient inference, achieving strong speedups with limited accuracy loss.
- Pre-trained transformer models are standard for classification and generation, while their growing size increases storage and efficiency demands.
- Pruning reduces task-specific model storage, but unstructured sparse models often require reconstructing dense shapes on standard hardware.
- Distillation produces dense models that run faster on standard hardware, but they are often larger than pruned models without careful engineering and size selection.
- Block pruning targets the gap by encouraging structures that can be optimized on dense hardware while integrating with movement pruning during fine-tuning.
- 2.4x speedup on SQuAD v1.1 accompanies a 1% drop in F1, while QQP reaches 2.3x speedup with a 1% F1 loss.
- 1.39x speedup on CNN/DailyMail accompanies an average 2-point drop across ROUGE metrics and a 3.5x reduction in decoder weights.
2 Related Work
Prior work compresses language models through distillation, unstructured pruning, and structured pruning. These approaches trade off model size, inference speed, hardware requirements, and preservation of useful transformer components.
- The related-work landscape considers distillation, pruning, and structured pruning as three main approaches to compressing pre-trained language models.
- Distillation creates smaller BERT models through teacher-student training, using either task-agnostic or task-specific layer-wise strategies.
- Unstructured pruning selects individual transformer weights by magnitude or importance scores, substantially reducing model size but requiring specialized hardware for sparse speedups.
- Structured pruning removes coherent weight groups, including attention heads that prior work finds can often be removed without significant performance degradation.
3 Background
The paper frames pruning as producing task-specific parameters that are both smaller and efficient on parallel hardware. It builds on score-based movement pruning while noting that high sparsity alone may not accelerate standard inference.
- The goal is to fine-tune parameters for an end task while making them smaller and efficiently computable on parallel hardware.
- Transformer parameter budgets are dominated by feed-forward and multi-head attention sub-layers, whose matrices account for their principal weights.
- Score-based pruning introduces per-parameter scores and masks weights through a function of those scores, with magnitude pruning zeroing low-absolute-value parameters.
- Movement pruning optimizes score parameters during fine-tuning using thresholded masks and a regularized objective that encourages sparsity.
- 94% pruning yields F1=87.5 on SQuAD v1.1 versus BERT-base F1=88.5, but remains insufficiently faster on standard hardware.
4 Model: Block Movement Pruning
Block Movement pruning extends movement pruning from individual parameters to fixed-size local blocks. Shared block scores provide a flexible structure intended to preserve pruning flexibility while improving data locality and inference efficiency.
- Each transformer matrix is partitioned into fixed-sized blocks, with each block acting as a regularized group rather than an independently pruned parameter.
- Block masks are generated by thresholding a smaller score matrix and expanding its values across corresponding weight blocks.
- The method trains with distillation to match a teacher model while requiring block sizes and shapes instead of a fully specified new architecture.
- Block size creates a trade-off: oversized blocks are difficult to prune, whereas undersized blocks do not support efficient inference.
- The experiments test shared attention and feed-forward block sizes, including (32, 32) square blocks and paired dimension-pruning patterns.
- Attention-head blocks and sufficiently large square blocks can remove complete components while leaving dense remaining matrices for computation.
5 Experimental Setup
The experiments evaluate task-specific pruning across classification, question answering, and summarization, comparing pruning methods with fast-inference baselines. The setup measures task performance, model size, and inference speed while testing several pruning structures and training procedures.
- Datasets and tasks: Experiments cover question answering, natural language inference, sentence similarity, sentiment classification, and abstractive summarization.The evaluated datasets include SQuAD v1.1, SQuAD v2, MNLI, QQP, SST-2, and CNN/DailyMail.
- Evaluation: Performance is reported with accuracy for MNLI and SST-2, F1 for QQP, EM and F1 for SQuAD, and ROUGE for CNN/DailyMail.
- Models: BERT is used for classification and question answering, while BART is used for summarization, with both base and large model configurations considered.BERT-base and BART-base contain 110M and 139M parameters, respectively; the large configurations contain 340M and 406M parameters.
- Baselines and measurement: The study compares movement pruning and fast-inference models including DistilBERT, TinyBERT, MobileBERT, and dBART using matched GPU and CPU timing setups.Inference speed is measured on an RTX 3090 GPU and Intel i7 CPU with batch size 128.
- Training procedure: The main hyperparameter is training duration, with SQuAD v1.1 using 20 epochs and approximately 12 hours of fine-tuning instead of 45 minutes for standard fine-tuning.Warmup and post-pruning cooldown phases help, but their exact lengths have limited impact on final performance.
- Pruning methods: Block pruning extends movement pruning with square blocks, while hybrid and structured variants target feed-forward dimensions or attention heads.Attention layers use square blocks or head-sized blocks; incomplete removal of components can recover sparsity without improving speed.
6 Experiments
Experiments show that hybrid block pruning can remove full feed-forward dimensions and attention heads, enabling favorable accuracy–speed–density trade-offs across encoder-only and encoder-decoder models.
- Main Results: Hybrid pruning learns to remove entire feed-forward dimensions and attention heads, enabling head removal during inference.The experiments therefore focus on the Hybrid approach, which targets both component types.
- SQuAD v1.1: 2.5x speedup is achieved at F1 87.5, while Hybrid Filled uses 25% of BERT’s encoder parameters at the same accuracy.TinyBERT and DistilBERT use 50% of BERT’s encoder parameters at that accuracy.
- Baselines: Movement pruning produces sparse models without significant speedups, square block pruning performs better, and hybrid blocks perform best among these intrinsic baselines.Movement pruning can even yield a small accuracy increase while providing little speed benefit.
- Compression and Speed: Hybrid Filled reaches a 2.25x speedup under minimal accuracy loss, while direct MHA block pruning is faster but degrades accuracy more.Measured speedups can underestimate those achievable with specialized implementations.
- Teacher Models: At 1.6x speedup, Hybrid pruning reaches F1 90.3, and at 2.2x speedup it matches BERT-base accuracy of F1 88.5.A larger teacher improves accuracy up to 80% sparsity, but becomes detrimental under heavier pruning.
- Encoder-Decoder: Hybrid pruning compresses BART decoders by ratios of 3.4 for BART-base and 3.5 for BART-large, with speedups reaching 1.4 times original speed.The authors report only a small ROUGE drop and expect further engineering could increase speedups.
7 Analysis
The analysis examines block pruning across model sizes, block sizes, quantization, and distillation. Results show strong compression and speed gains, with larger models and larger blocks offering favorable trade-offs.
- Large Model Pruning: 18% dense BERT-large reaches F1 90.2 with a 3.2x speedup versus BERT-large at F1 93.2.The pruned model is also faster than a BERT-base model.
- Block Size Influence: Block sizes from 4 to 32 preserve BERT-base original F1, while larger blocks increase speedup.A block size of 32 reaches a 1.3 speedup without F1 drop and 1.75 speedup for a 2% F1 drop.
- Quantization: The method is compatible with full 8-bit quantization without significant accuracy loss.Quantization was evaluated on pruned BERT-base models using SQuAD v1.1 dev Exact Match and F1.
- Impact of Distillation: Combining hybrid pruning with distillation consistently outperforms pruning alone, with larger benefits on smaller datasets such as SST-2.The authors attribute this pattern to regularization countering over-fitting from additional pruning steps.
8 Conclusion
The paper concludes that block pruning can extract compact models during fine-tuning, matching or exceeding distilled networks across diverse tasks and base models. It offers a simple alternative that avoids pre-training, data augmentation, and architecture search.
- Conclusion: Block pruning extracts small models that are equivalent to or better than distilled networks.The approach is evaluated across a diverse set of tasks and base models.
- Conclusion: The method operates during fine-tuning without data augmentation or architecture search.The authors describe it as simple and robust for accelerating models on specific tasks.
9 Impact
The paper links smaller pruned models to lower compute and energy requirements and to improved privacy when models run on user devices. These impacts are framed as especially relevant to anticipated NLP demand.
- Resource Use: Smaller pruned models are expected to reduce compute resources and energy needed for natural language tasks.The authors qualify this expectation by emphasizing preservation of original model performance.
- Privacy: Running smaller models on user devices instead of servers allows more information to remain private.The paper identifies this benefit as particularly relevant to future NLP applications.
A Reproducibility & Hyper-Parameters
The paper provides code, task-specific hyperparameters, checkpoints, and model artifacts through Hugging Face repositories, the model hub, and Amazon S3.
- Code: Complete experiment, analysis, figure, and table code is available in the Hugging Face nn_pruning repository.The repository is hosted at github.com/huggingface/nn_pruning.
- Hyper-Parameters: Task-specific JSON files contain fine-tuning and pruning hyperparameters, evaluation results, and final sparsity statistics.There is one hyperparameter file per task.
- Checkpoints: SQuAD V1 checkpoints and related hyperparameter information are listed in the repository.The entries include the referenced SQuAD V1 checkpoints.
- Model Artifacts: Some produced models are directly available through the Hugging Face model hub.Other models and intermediary training checkpoints are stored on Amazon S3.
B Additional Data
The experiments compare block movement pruning with distilled, hybrid, unstructured, and layer-dropping methods using speed and compression measures. They emphasize that block pruning can remove complete attention heads, with comparisons reported across benchmark tasks.
- Head pruning: Block pruning removes more complete attention heads than row/column pruning while improving accuracy on SST-2.This comparison is reported in the head-pruning evaluation.
- Experimental setup: Block Movement pruning is evaluated with different block patterns while pruning only the attention layers.Compression measures non-zero parameter reduction in attention linear layers; head compression measures reduction in complete non-zero heads.
- Reported evaluations: The reported tables cover head-pruning comparisons and broader distillation and pruning-method comparisons.Table 8 concerns SST-2 head pruning; Tables 9 and 10 concern distillation, structured pruning, and pruning methods.
- Experimental setup: Speed is selected as the main metric because it is the major practical measure of inference efficiency.The comparison includes MobileBERT and TinyBERT as strong speed/accuracy baselines.
- Comparisons: The study compares the method with distilled models, hybrid pruning variants, unstructured pruning, and full-layer dropping.The comparisons include TinyBERT, MobileBERT, Hybrid pruning, Movement Pruning, and methods from Gordon et al. and Sajjad et al.