PGB: One-Shot Pruning for BERT via Weight Grouping and Permutation

TL;DR

PGB addresses the high inference cost of large BERT models by introducing a one-shot, semi-structured pruning method that groups weights via permutation and prunes non-group weights, with dynamic per-layer grouping and optional layer dropping. It combines permutation-based grouping, weight pruning, re-permutation, and weight compensation followed by re-finetuning to preserve task accuracy while achieving substantial compression. Empirical results on BERT_BASE show that PGB outperforms state-of-the-art structured pruning methods in both efficiency and accuracy on GLUE and SQuAD benchmarks, and extends well to DistilBERT and RoBERTa. The approach offers a practical, hardware-friendly path to fast, accurate compression of transformer architectures without lengthy distillation or iterative pruning pipelines.

Abstract

Large pretrained language models such as BERT suffer from slow inference and high memory usage, due to their huge size. Recent approaches to compressing BERT rely on iterative pruning and knowledge distillation, which, however, are often too complicated and computationally intensive. This paper proposes a novel semi-structured one-shot pruning method for BERT, called $\textit{Permutation and Grouping for BERT}$ (PGB), which achieves high compression efficiency and sparsity while preserving accuracy. To this end, PGB identifies important groups of individual weights by permutation and prunes all other weights as a structure in both multi-head attention and feed-forward layers. Furthermore, if no important group is formed in a particular layer, PGB drops the entire layer to produce an even more compact model. Our experimental results on BERT$_{\text{BASE}}$ demonstrate that PGB outperforms the state-of-the-art structured pruning methods in terms of computational cost and accuracy preservation.

Figures (7)

• Figure 1: Permutation and Grouping for BERT (PGB), where weight matrices are grouped and all individual weights not belonging to groups are pruned.
• Figure 2: The grouping procedure of PGB for weight matrices in BERT, which applies to both MHA and FFN layers
• Figure 3: Performance comparison with structured pruning methods varying the reduced FLOPs ratio on $\text{BERT}_{\text{BASE}}$.
• Figure 4: Performance comparison with structured pruning methods varying the Reduced FLOPs ratio on $\text{SQuAD}$ benchmarks.
• Figure 5: PGB with $\text{BERT}_{\text{BASE}}$ and $\text{RoBERTa}_{\text{BASE}}$