Know What You Don't Need: Single-Shot Meta-Pruning for Attention Heads

TL;DR

This paper introduces Single-Shot Meta-Pruning, a method to efficiently prune attention heads in pre-trained Transformer models, reducing computational costs while maintaining or improving task performance.

Contribution

It presents a novel meta-learning based approach for adaptive, single-shot pruning of attention heads in Transformers before fine-tuning, enhancing efficiency and effectiveness.

Findings

Prunes 50% of attention heads with minimal performance loss

Reduces both fine-tuning and inference overheads

Improves quality of text representations

Abstract

Deep pre-trained Transformer models have achieved state-of-the-art results over a variety of natural language processing (NLP) tasks. By learning rich language knowledge with millions of parameters, these models are usually overparameterized and significantly increase the computational overhead in applications. It is intuitive to address this issue by model compression. In this work, we propose a method, called Single-Shot Meta-Pruning, to compress deep pre-trained Transformers before fine-tuning. Specifically, we focus on pruning unnecessary attention heads adaptively for different downstream tasks. To measure the informativeness of attention heads, we train our Single-Shot Meta-Pruner (SMP) with a meta-learning paradigm aiming to maintain the distribution of text representations after pruning. Compared with existing compression methods for pre-trained models, our method can reduce the overhead of both fine-tuning and inference. Experimental results show that our pruner can selectively prune 50% of attention heads with little impact on the performance on downstream tasks and even provide better text representations. The source code will be released in the future.