Source-linked AI summary
Structured Pruning Learns Compact and Accurate Models
Mengzhou Xia, Zexuan Zhong, Danqi Chen
TL;DR
Large language models create storage, memory, and computation costs, while pruning has limited speedups and distillation requires expensive training on unlabeled data. CoFi jointly prunes coarse- and fine-grained Transformer units with layerwise distillation, and reports competitive accuracy with over 10× speedup on GLUE and SQuAD. The results position task-specific structured pruning as an alternative to distillation for extreme compression without expensive pre-training or data augmentation.
Problem
Pruning often provides limited latency gains, whereas distillation requires large unlabeled corpora and expensive training for compact language models.
Method
CoFi jointly prunes coarse-grained layers and fine-grained heads or hidden dimensions using multiple masks, with layerwise distillation between unpruned and pruned models.
Results
CoFi delivers more accurate models across speedup and model-size levels on GLUE and SQuAD, achieving over 10× speedups while preserving more than 90% of accuracy across datasets.
Takeaways & Limitations
Task-specific structured pruning can produce highly compact, fast models without additional unlabeled data for general distillation or expensive pre-training and data augmentation.
Takeaways & Limitations
Experiments with data augmentation cover only four datasets because augmented-data training is very expensive, with MNLI requiring more than 200 GPU hours.
Abstract
from arXiv · showhide
The growing size of neural language models has led to increased attention in model compression. The two predominant approaches are pruning, which gradually removes weights from a pre-trained model, and distillation, which trains a smaller compact model to match a larger one. Pruning methods can significantly reduce the model size but hardly achieve large speedups as distillation. However, distillation methods require large amounts of unlabeled data and are expensive to train. In this work, we propose a task-specific structured pruning method CoFi (Coarse- and Fine-grained Pruning), which delivers highly parallelizable subnetworks and matches the distillation methods in both accuracy and latency, without resorting to any unlabeled data. Our key insight is to jointly prune coarse-grained (e.g., layers) and fine-grained (e.g., heads and hidden units) modules, which controls the pruning decision of each parameter with masks of different granularity. We also devise a layerwise distillation strategy to transfer knowledge from unpruned to pruned models during optimization. Our experiments on GLUE and SQuAD datasets show that CoFi yields models with over 10x speedups with a small accuracy drop, showing its effectiveness and efficiency compared to previous pruning and distillation approaches.
1 Introduction
CoFi addresses the compression trade-off between pruning's limited speedups and distillation's costly unlabeled-data training by jointly pruning Transformer structures at multiple granularities. It reports compact, fast models with competitive accuracy on GLUE and SQuAD, while using task-specific pruning and layerwise distillation.
- Pre-trained language models have high storage, memory, and computation costs, motivating model-compression research for real-world use.
- Distillation can provide compelling speedup-performance tradeoffs, but TinyBERT requires 2,500M tokens and 3.5 days of training on 4 GPUs.
- CoFi jointly prunes coarse-grained layers and fine-grained heads or hidden dimensions, controlling each parameter with multiple masks.
- Layerwise distillation dynamically learns mappings between unpruned and pruned models, improving beyond simple prediction-layer distillation.
- Over 10× speedups and more accurate models across speedup and model-size levels are reported on GLUE and SQuAD compared with strong pruning and distillation baselines.
- CoFi preserves more than 90% of accuracy across datasets while producing smaller and faster models without additional unlabeled data for general distillation.
2 Background
Transformer compression targets models whose storage, memory, and computation costs hinder practical use. Background work compares distillation and structured pruning, including their trade-offs in data, training cost, flexibility, and inference speed.
- 2.1 Transformers: A Transformer has L blocks, each containing multi-head self-attention and feed-forward layers, followed by residual connections and layer normalization.The attention layer processes input X across Nh heads, while the feed-forward layer uses an up-projection, GELU activation, and down-projection.
- 2.1 Transformers: MHA and FFN layers account for 1/3 and 2/3 of Transformer parameters, respectively, while both take similar GPU time and FFNs bottleneck CPUs.
- 2.2 Distillation: Distillation transfers knowledge from a larger teacher to a smaller student, commonly combining general distillation on unlabeled data with task-specific distillation.General distillation or student pre-training helps retain performance but is computationally expensive.
- 2.2 Distillation: TinyBERT uses a randomly initialized student and two-step distillation, including general distillation that takes 3.5 days on 4 GPUs.CoFi instead prunes a fine-tuned BERT model and takes at most 20 hours on 1 GPU across GLUE datasets.
- 2.3 Pruning: Distillation methods generally assume a fixed student structure, whereas pruning searches for task-specific subnetworks by removing Transformer layers, heads, dimensions, or blocks.Existing structured pruning trends toward fine-grained units, but pruned models rarely achieve large speedups; unstructured sparsity can reach 97% while remaining difficult to accelerate on current hardware.
- 2.3 Pruning: Removing heads alone provides limited latency gains, with one remaining head per layer yielding a reported 1.4× speedup.Pruning work also considers entire layers, FFN layers, and intermediate dimensions as alternative structured units.
3 Method
CoFi jointly prunes coarse- and fine-grained Transformer units while using layerwise distillation to transfer knowledge from an unpruned teacher to a changing pruned model. Multiple masks control parameter pruning, producing compressed structures with large inference speedups.
- Coarse- and Fine-Grained Pruning: CoFi jointly prunes coarse-grained layers and fine-grained units such as heads, intermediate dimensions, and hidden dimensions.It explicitly prunes MHA and FFN layers alongside finer units.
- Coarse- and Fine-Grained Pruning: Explicitly pruning MHA and FFN layers addresses the difficulty of obtaining such coarse pruning through fine-grained pruning alone, especially at high sparsity.MHA and FFN layers are dropped separately rather than as one combined layer.
- Coarse- and Fine-Grained Pruning: Hidden-dimension masks are shared across layers because residual connections connect corresponding representation dimensions between successive layers.Although few dimensions may be pruned, the paper reports that this improves performance significantly.
- Coarse- and Fine-Grained Pruning: Multiple masks jointly control each parameter’s pruning decision, allowing a weight to be removed through its layer, intermediate dimension, or hidden dimension.This differs from approaches applying one pruning strategy separately to MHAs and FFNs.
- Distillation to Pruned Models: Layerwise distillation dynamically maps teacher layers to student layers because the pruned model structure changes during training.The mapping is learned by matching teacher layers to student layers, with constraints used to address layer mismatch.
- Distillation to Pruned Models: The final distillation objective combines prediction-layer and layerwise losses, with λ controlling their relative contributions.Dynamic matching is reported to avoid frequent layer mismatch and improve performance beyond prediction-layer distillation alone.
4 Experiments
Experiments evaluate CoFi across GLUE and SQuAD against pruning and distillation baselines, showing strong accuracy-speedup tradeoffs and the effects of pruning units, distillation objectives, and resulting structures.
- Setup: CoFi is evaluated on eight GLUE tasks and SQuAD v1.1 against pruning and distillation baselines.Experiments start from BERTbase, use development sets, and report target sparsities from 60% to 95%.
- Overall performance: CoFi delivers more accurate models than distillation and pruning baselines at every speedup level and model size.The comparison uses accuracy against inference speedup and model size, excluding embedding size from model-size calculations.
- Comparison with TinyBERT4: Over 10× inference speedup is achieved with performance comparable to or better than TinyBERT4, while CoFi trains for at most 20 GPU hours on task-specific data.TinyBERT4 relies on general distillation from a large unlabeled corpus, which can take up to hundreds of GPU hours.
- Comparison with TinyBERT4: With task-specific data augmentation, CoFi still outperforms TinyBERT4 on most datasets under similar 11-12× speedups.Only four datasets use augmentation because training on augmented data is expensive; MNLI requires more than 200 GPU hours in total.
- Pruning units: Removing hidden-unit masks causes accuracy drops, while removing layer masks greatly reduces speedups at 95% sparsity.At 60% and 95% sparsity, the corresponding model sizes are 34M and 5M; equal parameter counts can yield different speedups.
- Distillation objectives: Removing distillation causes performance drops of 1.9-6.8 points, and dynamic layer matching outperforms fixed hidden distillation.Dynamic matching aligns teacher and student layers as pruning changes the student structure.
5 Related Work
Related work spans structured and unstructured pruning, upstream pruning for task-agnostic models, and other Transformer acceleration techniques such as distillation and quantization.
- Structured pruning: Structured pruning in Transformers removes layers, heads, intermediate dimensions, or blocks in weight matrices.The literature includes both coarse and fine-grained structural choices.
- Upstream pruning: Upstream pruning applies pruning during masked language modeling to produce task-agnostic models rather than task-specific subnetworks.The paper identifies applying CoFi upstream as a promising future direction.
- Other acceleration methods: Other Transformer acceleration methods include distillation, quantization, dynamic inference acceleration, and matrix decomposition.The paper does not directly compare against some methods because their evaluation environments differ.
6 Conclusion
CoFi combines coarse- and fine-grained structured pruning with layerwise distillation to produce compact Transformer models with large speedups and competitive results.
- CoFi jointly prunes MHA/FFN layers, individual heads, and hidden dimensions in Transformer-based models.
- More than 10× speedup is achieved while maintaining competitive results with a structure different from standard distillation models.
- Task-specific structured pruning is presented as an alternative to distillation that avoids expensive pre-training or data augmentation.
- The work scopes CoFi to task-specific pruning because upstream, task-agnostic pruning involves complex design choices.
A Reproducibility & Hyperparameters
The experiments specify training schedules, hard-concrete mask optimization, and a sparsity constraint to control the final pruned model.
- Hyperparameters for the experiments are reported in Table 7.
- Training schedules differ by dataset size: larger GLUE datasets and SQuAD use 20 training epochs plus 20 finetuning epochs, while small GLUE datasets use 100 epochs.
- L0-compression masks are sampled using hard-concrete distributions, with sigmoid stretching controlled by l and r and steepness controlled by β.
- The learnable mask-distribution parameters are updated during the forward pass to learn pruning masks.
- A Lagrangian violation penalty imposes equality between expected model sparsity ŝ and target sparsity t.
C Details of Baseline Methods
The evaluation compares CoFi with established pruning and distillation systems, using Hybrid Filled as the pruning baseline.
- The comparison includes DistillBERT, TinyBERT4, TinyBERT6, DynaBERT, and Block Pruning.
- TinyBERT4 and TinyBERT6 use both general distillation for pretraining and task-specific distillation.
- DynaBERT provides dynamic-sized models by specifying width and depth.
- Block Pruning couples pruning with prediction-layer distillation.
- The strongest Hybrid Filled approach is selected as the baseline.
D Data Statistics
Dataset statistics for the GLUE and SQuAD datasets used in the experiments are reported in Table 8.
- Train sizes and evaluation metrics for each dataset are reported in Table 8.
- Table 8 presents data statistics for the GLUE and SQuAD datasets.
E TinyBERT4 w/ Data Augmentation
The experiments compare CoFi with pruning and distillation baselines under matched or documented settings. CoFi generally preserves strong accuracy while providing substantial inference speedups, and its dynamic layer distillation improves performance across sparsities.
- Comparison to Movement Pruning: CoFi achieves 2−11× inference speedups, whereas Movement Pruning achieves no speedup gains.CoFi largely outperforms Movement Pruning on MNLI and is comparable on SQuAD for models larger than 10M parameters.
- Comparison to Block Pruning: CoFi outperforms or matches Block Pruning under the same prediction-layer distillation objective.Block Pruning never achieves a speedup of 10 for a similarly sized SST-2 model.
- Ablation on Distillation Objectives: Dynamic layer matching learns task-specific correspondences between pruned-student and unpruned-teacher layers during training.For SST-2, student layers 7, 9, 10, and 11 align with teacher layers 3, 6, 9, and 12; QQP learns a different mapping.
- Ablation on Distillation Objectives: Adding dynamic layer distillation improves performance across all sparsity rates and tasks.The ablation adds layer distillation to prediction distillation and reports clear gains across sparsities.
- Model Structures: Increasing sparsity produces shallower models, which usually contain more MHA layers than FFN layers.The observed layer counts may indicate that MHA layers are more important for downstream tasks than FFN layers.
- Results Across Backbones: At 60% sparsity, RoBERTa models maintain full-model performance, while BERT surpasses RoBERTa at sparsities above 90%.RoBERTa pruning outperforms BERT below 90% sparsity.
J Training Time Measurement
The training-time comparison measures TinyBERT using RTX 2080Ti GPUs and scales timing from a small corpus to the original corpus. The section also presents comparison figures and an ablation table for pruning and distillation methods.
- Training Time Measurement: TinyBERT training time is measured on NVIDIA RTX 2080Ti GPUs using a small corpus and scaled to the original corpus size.The small corpus contains 10.6M tokens, while the original corpus contains 2500M tokens.
- Training Time Measurement: 430s are required for one epoch on the 10.6M-token corpus used for TinyBERT timing.The reported timing is used to estimate training time on the larger original corpus.
- Comparison Materials: Figure 6 compares CoFi with Block Pruning using prediction-layer distillation, while Table 11 reports layer-distillation ablations across sparsities.These comparisons accompany the training-time measurement discussion.