Block Pruning For Faster Transformers
Introduces block movement pruning, combining it with movement pruning to accelerate Transformer models by 2.4x and reduce size by 74% while maintaining performance.
Key Findings
Methodology
This paper proposes Block Movement Pruning, which partitions Transformer weight matrices into fixed-size blocks with shared pruning scores. During fine-tuning, the method optimizes a combined loss to learn which blocks to prune, encouraging removal of entire attention heads and feed-forward components. It extends movement pruning by supporting multi-scale blocks, enabling hardware-friendly sparsity. Experiments on BERT and BART across tasks like SQuAD, MNLI, and CNN/DM show 2.4x speedup and 74% model size reduction with minimal performance loss. The approach leverages transfer learning, regularization, and thresholding to achieve effective structured sparsity.
Key Results
- On SQuAD v1.1, the pruned models achieve 2.4x inference speedup with only 1% F1 drop, and model size is reduced to 26% of the original.
- In QQP, similar speedup (2.3x) with 1% F1 loss; on summarization, 1.39x speedup with 2-point ROUGE drop, and decoder weights reduced 3.5x.
- Across multiple tasks, the method consistently outperforms unstructured movement pruning and approaches or surpasses distilled models in speed and compression, validating its efficiency and robustness.
Significance
This work addresses the critical challenge of deploying large pre-trained Transformers efficiently on hardware. By introducing hardware-friendly block sparsity, it bridges the gap between model compression and inference acceleration. The method enables practical deployment in resource-constrained environments, reduces energy consumption, and accelerates real-time NLP applications. It also opens avenues for integrating structured pruning with hardware-aware optimization, advancing the state-of-the-art in efficient deep learning models.
Technical Contribution
The core technical innovation is the extension of movement pruning to block structures with shared scores, supporting multi-scale block sizes. This facilitates hardware-compatible sparsity, enabling entire attention heads and feed-forward layers to be pruned effectively. The approach simplifies hyperparameter tuning, integrates seamlessly with transfer learning, and provides theoretical and empirical guarantees of sparsity and speedup. It surpasses prior methods by balancing model performance, sparsity, and hardware efficiency, offering a practical solution for large-scale Transformer acceleration.
Novelty
This is the first work to extend movement pruning to fixed-size blocks with shared scores, supporting multi-scale block configurations. Unlike previous unstructured or purely structural pruning, it combines the flexibility of movement pruning with hardware-aware block sparsity, enabling significant speedups without performance degradation. This dual focus on model accuracy and hardware efficiency marks a novel contribution in the field.
Limitations
- Block size selection critically impacts performance; overly large blocks reduce pruning granularity, while small blocks may hinder hardware acceleration.
- Extreme pruning ratios can cause performance drops, especially on complex tasks, requiring careful regularization and tuning.
- Current validation is primarily on GPU hardware; effectiveness on CPU or specialized accelerators remains to be explored.
Future Work
Future directions include adaptive block size strategies based on hardware profiles, combining pruning with quantization and distillation, and extending the approach to larger models and multi-modal architectures. Further research is needed to optimize block selection dynamically, improve generalization across hardware platforms, and integrate with end-to-end deployment pipelines for real-world NLP systems.
AI Executive Summary
The rapid growth of pre-trained Transformer models like BERT and BART has revolutionized natural language processing, but their enormous size poses significant challenges for deployment, especially in resource-constrained environments. While traditional compression techniques such as unstructured pruning and distillation have made strides, they often fall short in delivering practical inference speedups on standard hardware. This bottleneck stems from the difficulty of efficiently utilizing sparse matrices in real-world hardware accelerators.
To address this, the authors introduce a novel block movement pruning method that extends the movement pruning framework by partitioning weight matrices into fixed-size blocks. These blocks are assigned shared pruning scores, which are optimized during fine-tuning through a combined loss function that encourages the removal of entire blocks, including attention heads and feed-forward layers. This approach supports multi-scale block configurations, balancing pruning granularity with hardware efficiency.
Experimental results on multiple NLP tasks demonstrate the effectiveness of this method. On SQuAD v1.1, the pruned models achieve a 2.4x inference speedup with only a 1% drop in F1 score, and the model size is reduced by 74%. Similar gains are observed across other tasks like QQP and summarization, with models maintaining high accuracy while significantly reducing computational costs. These findings suggest that block-based structured pruning can unlock practical acceleration for large Transformer models, making them more suitable for deployment in real-world applications.
This work marks a significant step forward in model compression, bridging the gap between theoretical sparsity and hardware-friendly acceleration. By enabling entire attention heads and feed-forward components to be pruned efficiently, it paves the way for more energy-efficient, faster NLP systems. Future efforts will focus on adaptive block sizing, hardware-aware optimization, and extending the approach to larger models and multi-modal architectures, promising a new era of scalable, efficient deep learning models.
Deep Dive
Abstract
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.