Structured Pruning Learns Compact and Accurate Models (2022)

Abstract: 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.

## introduction pre-trained language models @xcite @xcite , inter alia) have become the mainstay in natural language processing. these models have high costs in terms of storage, memory, and computation time and it has motivated a large body of work on model compression to make them smaller and faster to use in real-world applications @xcite . the the inference speedups (↗) are reported against a bert base model and we evaluate all the models on an nvidia v100 gpu ( §4.1). the models labeled as ‡ use a different teacher model and are not a direct comparison. models are one order of magnitude faster. 3 two predominant approaches to model compression are pruning and distillation ( pruning methods search for an accurate subnetwork in a larger pre-trained model. recent work has investigated how to structurally prune transformer networks @xcite , from removing entire layers @xcite , to pruning heads @xcite , intermediate dimensions @xcite and blocks in weight matrices @xcite . the trend of structured pruning leans towards removing finegrained units to allow for flexible final structures. however, thus far, pruned models rarely achieve large speedups (2-3× improvement at most). by contrast, distillation methods usually first specify a fixed model architecture and perform a general distillation step on an unlabeled corpus, before further fine-tuning or distillation on task-specific data @xcite @xcite . welldesigned student architectures achieve compelling speedup-performance tradeoffs, yet distillation to these randomly-initialized student networks on large unlabeled data is prohibitively slow. 4 for instance, tinybert @xcite ) is first trained on 2,500m tokens for 3 epochs, which requires training 3.5 days on 4 gpus (figure 1 ). 5 in this work, we propose a task-specific, structured pruning approach called cofi (coarse and fine-grained pruning) and show that structured pruning can achieve highly compact subnetworks and obtain large speedups and competitive accuracy as distillation approaches, while requiring much less computation. our key insight is to jointly prune coarse-grained units (e.g., self-attention or feed-forward layers) and fine-grained units (e.g., heads, hidden dimensions) simultaneously. different from existing works, our approach controls the pruning decision of every single parameter by multiple masks of different granularity. this is the key to large compression, as it allows the greatest flexibility of pruned structures and eases the optimization compared to only pruning small units. it is known that pruning with a distillation objective can substantially improve performance @xcite . unlike a fixed student architecture, pruned structures are unkown prior to training and it is challenging to distill between intermediate layers of the unpruned and pruned models @xcite . hence, we propose a layerwise distillation method, which dynamically learns the layer mapping between the two structures. we show that this strategy can better lead to performance gains beyond simple prediction-layer distillation. our experiments show that cofi delivers more accurate models at all levels of speedups and model sizes on the glue @xcite and squad v1.1 @xcite datasets, compared to strong pruning and distillation baselines. concretely, it achieves over 10× speedups and a 95% sparsity across all the datasets while preserving more than 90% of accuracy. our results suggest that task-specific structured pruning is an appealing solution in practice, yielding smaller and faster models without requiring additional unlabeled data for general distillation. ## method we propose a structured pruning approach cofi, which jointly prunes coarse-grained and finegrained units ( §3.1) with a layerwise distillation objective transferring knowledge from unpruned to pruned models ( §3.2). a combination of the two leads to highly compressed models with large inference speedups. ## conclusion we propose cofi, a structured pruning approach that incorporates all levels of pruning, including mha/ffn layers, individual heads, and hidden dimensions for transformer-based models. coupled with a distillation objective tailored to structured pruning, we show that cofi compresses models into a rather different structure from standard distillation models but still achieves competitive results with more than 10× speedup. we conclude that task-specific structured pruning from large-sized models could be an appealing replacement for distillation to achieve extreme model compression, without resorting to expensive pre-training or data augmentation. though cofi can be directly applied to structured pruning for task-agnostic models, we frame the scope of this work to task-specific pruning due to the complexity of the design choices for upstream pruning. we hope that future research continues this line of work, given that pruning from a large pre-trained model could possibly incur less computation compared to general distillation and leads to more flexible model structures.