Differentiable Subset Pruning of Transformer Heads (2021)

Abstract: Multi-head attention, a collection of several attention mechanisms that independently attend to different parts of the input, is the key ingredient in the Transformer. Recent work has shown, however, that a large proportion of the heads in a Transformer’s multi-head attention mechanism can be safely pruned away without significantly harming the performance of the model; such pruning leads to models that are noticeably smaller and faster in practice. Our work introduces a new head pruning technique that we term differentiable subset pruning. ntuitively, our method learns per- head importance variables and then enforces a user-specified hard constraint on the number of unpruned heads. he importance variables are learned via stochastic gradient descent. e conduct experiments on natural language inference and machine translation; we show that differentiable subset pruning performs comparably or better than previous works while offering precise control of the sparsity level.1

## introduction the transformer @xcite has become one of the most popular neural architectures used in nlp. adaptations of the transformer have been applied to nearly every popular nlp task, for example, parsing (zhou and zhao, 2019), machine translation @xcite , question answering @xcite inter alia. transformers also form the backbone of state-of-the-art pre-trained language models, for example, bert @xcite , gpt-2 @xcite , and gpt-3 @xcite , that have further boosted performance on various data-driven nlp problems. the key ingredient in the transformer architecture is the multi-head attention mechanism, which is an assembly of multiple attention functions @xcite applied in parallel. in outperforms two recently proposed transformer head pruners-michel et al. (2019) and voita et al. ( 2019 )-on both tasks in terms of sparsityperformance trade-off. our method recovers a pruned transformer that has ≈ 80% accuracy on mnli and ≈ 30 bleu score on iwslt when more than 90% of the heads are removed, which brings about ≈ 33% inference speedup and ≈ 24% model size shrinkage. 2 our experiments also suggest several broader conclusions about pruning transformers. in this paper, we taxonomize existing pruning methods into two pruning paradigms: pipelined pruning and joint pruning. pipelined pruning consists of two stages: (i) training or fine-tuning an overparameterized model on the target task and (ii) pruning the model after training. a number of techniques fall into this category @xcite @xcite . in contrast, joint pruning blends the pruning objective into the training objective by training or fine-tuning the over-parameterized model with a sparsity-enforcing regularizer, sometimes followed up by a trivial post-processing step to arrive at a final sparse model. @xcite and @xcite are examples of 2 see § 5.4. ## background: multi-head attention in this section, we provide a detailed overview of multi-head attention @xcite in order to develop the specific technical vocabulary to discuss our approaches for head pruning. we omit details about other parts of the transformer and refer the reader back to the original work of @xcite . first, let @xmath0, . . . , z t be a sequence of t real vectors where z @xmath1, and let @xmath2. an attention mechanism is defined as the projection matrices w o , w v , w q , w @xmath3. in self-attention, query q comes from the same sequence z. a transformer is composed of l identical layers. in layer 1 ≤ @xmath4, h l different attention mechanisms are applied in parallel; importantly, it is this parallelism that has lead to the rise of the transformer-it is a more efficient architecture in practice so it can be trained on more data. each individual attention mechanism is referred to as a head; thus, multi-head attention is the simultaneous application of multiple attention heads in a single architecture. in @xcite , the multiple heads are combined through summation: we also introduce a gate variable g lh that takes values in the interval [0, 1]: inserting g lh into the multi-head attention enables our pruning approach: setting the gate variable to g @xmath5. in the following sections, for the sake of notational simplicity, we ignore the layer structure of heads and label heads with a single index @xmath6, . . . , h}, where @xmath7. ## differentiable subset pruning in this section, we propose a new head pruning technique that we term differentiable subset pruning. the key insight behind our approach is that head pruning can be viewed as subset selection. concretely, our goal is to find a subset of k heads (where k is a user-specified positive integer) that still allows the model to achieve high performance. many neural network pruners, for example, @xcite proposed head pruning technique, make it notably difficult to pre-specify the number of pruned heads k 3 . to make our subset pruner differentiable, we apply the gumbel-softmax trick @xcite 3 later discussed in § 5.2. and its extension to subset selection @xcite . this gives us a pruning scheme that always returns the specified number of heads and can be applied in a pipelined or a joint setting. in both cases, the differentiability is necessary to learn the head weights.