Differentiable Subset Pruning of Transformer Heads

Differentiable Subset Pruning of Transformer Heads Jiaoda Li♣ Ryan Cotterell♣♠ Mrinmaya Sachan♣ ♣ETH Z¨urich, Switzerland ♠University of Cambridge, UK {jiaoda.li,ryan.cotterell,mrinmaya.sachan}@inf.ethz.ch Abstract Multi-head attention, a collection of several attention mechanisms that independently at- tend 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 at- tention 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 prac- tice. Our work introduces a new head pruning technique that we term differentiable subset pruning. Intuitively, our method learns per- head importance variables and then enforces a user-specified hard constraint on the number of unpruned heads. The importance variables are learned via stochastic gradient descent. We conduct experiments on natural language in- ference and machine translation; we show that differentiable subset pruning performs com- parably or better than previous works while offering precise control of the sparsity level.1 1 Introduction The Transformer (Vaswani et al., 2017) has be- come 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 (Ng et al., 2019), question answering (Yang et al., 2019) inter alia. Transformers also form the backbone of state-of-the-art pre-trained language models, for example, BERT (Devlin et al., 2019), GPT-2 (Radford et al., 2019), and GPT-3 (Brown et al., 2020), that have further boosted per- formance on various data-driven NLP problems. The key ingredient in the Transformer architec- ture is the multi-head attention mechanism, which is an assembly of multiple attention functions (Bahdanau et al., 2015) applied in parallel. In 1Our code is available here: https://github.com /rycolab/differentiable-subset-pruning. practice, each attention head works indepen- dently, which allows the heads to capture different kinds of linguistic phenomena (Clark et al., 2019; Goldberg, 2019; Ettinger, 2020; Jawahar et al., 2019). A natural question arises in this context: How many heads does a transformer need? Michel et al. (2019) offer the insight that a large portion of the Transformer’s heads can be pruned without significantly degrading the test accuracy on the desired task. The experimental evidence behind their claim is a simple greedy pro- cedure that sequentially removes heads. This sug- gests that a better pruner could reveal that a much larger portion of the heads can be safely removed. To provide a more robust answer to Michel et al.’s question, we build a high-performance pruner and show that their approach itself significantly under- estimates the number of Transformer heads that can be pruned away. From a bird’s eye view, our paper contributes the proposal that Transformer head pruning is best viewed as a subset selection problem. Subset se- lection is common across many areas of NLP, from extractive summarization (Gillenwater et al., 2012) to vowel typology (Cotterell and Eisner, 2017). In the case of head pruning, the concrete idea is that …

Content truncated to 500 words