Gradient-based Intra-attention Pruning on Pre-trained Language Models

TL;DR

This paper introduces GRAIN, a gradient-based intra-attention pruning method for pre-trained language models that enables highly effective, task-specific model compression by inspecting and pruning intra-attention structures, resulting in significant speedups and maintained performance.

Contribution

GRAIN is a novel structured pruning approach that prunes intra-attention components with a gradient separation strategy, improving flexibility and efficiency over traditional head pruning methods.

Findings

GRAIN achieves 6-7x speedups with 93-99% performance retention.

It outperforms existing pruning methods, especially at high sparsity levels.

Even with only 3% of weights remaining, the model remains competitive.

Abstract

Pre-trained language models achieve superior performance but are computationally expensive. Techniques such as pruning and knowledge distillation have been developed to reduce their sizes and latencies. In this work, we propose a structured pruning method GRAIN (Gradient-based Intra-attention pruning), which performs task-specific pruning with knowledge distillation and yields highly effective models. Different from common approaches that prune each attention head as a whole, GRAIN inspects and prunes intra-attention structures, which greatly expands the structure search space and enables more flexible models. We also propose a gradient separation strategy that reduces the interference of distillation on pruning for a better combination of the two approaches. Experiments on GLUE, SQuAD, and CoNLL 2003 show that GRAIN notably outperforms other methods, especially in the high sparsity regime, and achieves $6\sim7\times$ speedups while maintaining $93\%\sim99\%$ performance. Under extreme compression where only $3\%$ transformer weights remain, the pruned model is still competitive compared to larger models.