Losing Heads in the Lottery: Pruning Transformer Attention in Neural Machine Translation

Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, pages 2664–2674, November 16–20, 2020. c⃝2020 Association for Computational Linguistics 2664 Losing Heads in the Lottery: Pruning Transformer Attention in Neural Machine Translation Maximiliana Behnke The University of Edinburgh maximiliana.behnke@ed.ac.uk Kenneth Heafield The University of Edinburgh kheafiel@inf.ed.ac.uk Abstract The attention mechanism is the crucial com- ponent of the transformer architecture. Re- cent research shows that most attention heads are not confident in their decisions and can be pruned after training. However, remov- ing them before training a model results in lower quality. In this paper, we apply the lot- tery ticket hypothesis to prune heads in the early stages of training, instead of doing so on a fully converged model. Our experiments on machine translation show that it is possi- ble to remove up to three-quarters of all at- tention heads from a transformer-big model with an average −0.1 change in BLEU for Turkish→English. The pruned model is 1.5 times as fast at inference, albeit at the cost of longer training. The method is complementary to other approaches, such as teacher-student, with our English→German student losing 0.2 BLEU at 75% encoder attention sparsity. 1 Introduction The transformer model (Vaswani et al., 2017) per- forms well for a variety of tasks, including neu- ral machine translation (Dong et al., 2018; Grund- kiewicz and Junczys-Dowmunt, 2019; Junczys- Dowmunt, 2018). However, like many neural net- works, it is overparametrised, and inference is costly. Attention heads are the headline feature of the transformer model, essential to learning rela- tionships between words as well as complex struc- tural representations. Voita et al. (2019) showed that many of these heads could be pruned in a fully trained model, but removing the same heads be- fore training yielded lower quality. We investigate a third way: pruning heads in early training. Em- pirically our method enables even more pruning, which is useful for faster machine translation. Reinitialising a model with the same pruned structure underperformed in Voita et al. (2019), which is consistent with the lottery ticket hypoth- esis (Frankle and Carbin, 2019). According to the lottery ticket hypothesis, randomly initialis- ing a model is akin to buying lottery tickets and a smaller network, such as a pruned model, buys fewer tickets. Prior lottery ticket research prunes individual parameters to form a sparse network; we show that this logic extends to entire trans- former heads. We follow lottery ticket train- ing strategies (Frankle et al., 2019) to prune in early training, achieving a better trade-off between pruning and quality than pruning after training (Voita et al., 2019). Our main goal is faster inference speed for ma- chine translation deployment with minimal impact on quality. Pruning heads means they can be re- moved from the model entirely (with other heads shifted down), resulting in a layer configured to have fewer heads. Unlike most work on pruning (Zhu and Gupta, 2017; Gale et al., 2019), there is no need for sparse matrices, block-sparse ma- trix operators, or additional …

Content truncated to 500 words