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

Abstract: The attention mechanism is the crucial component of the transformer architecture. Recent research shows that most attention heads are not confident in their decisions and can be pruned. However, removing them before training a model results in lower quality. In this paper, we apply the lottery ticket hypothesis to prune heads in the early stages of training. Our experiments on machine translation show that it is possible to remove up to three-quarters of attention heads from transformer-big during early training 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. Our method is complementary to other approaches, such as teacher-student, with English→German student model gaining an additional 10% speed-up with 75% encoder attention removed and 0.2 BLEU loss.

## introduction the transformer model @xcite performs well for a variety of tasks, including neural machine translation @xcite @xcite . however, like many neural networks, it is overparametrised, and inference is costly. attention heads are the headline feature of the transformer model, essential to learning relationships between words as well as complex structural representations. @xcite showed that many of these heads could be pruned in a fully trained model, but removing the same heads before training yielded lower quality. we investigate a third way: pruning heads in early training. empirically our method enables even more pruning, which is useful for faster machine translation. reinitialising a model with the same pruned structure underperformed in @xcite , which is consistent with the lottery ticket hypothesis @xcite . according to the lottery ticket hypothesis, randomly initialising 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 transformer heads. we follow lottery ticket training strategies @xcite to prune in early training, achieving a better trade-off between pruning and quality than pruning after training @xcite . our main goal is faster inference speed for machine translation deployment with minimal impact on quality. pruning heads means they can be removed from the model entirely (with other heads shifted down), resulting in a layer configured to have fewer heads. unlike most work on pruning @xcite , there is no need for sparse matrices, block-sparse matrix operators, or additional masking. in particular, we go further than @xcite by removing rather than masking. in this paper we combine findings of both @xcite ("what") and @xcite ("how") to prune attention heads. first, we define a training scheme based on an iterative approach that does not require full convergence of a model each time partial pruning takes place. to analyse the impact of pruning in a variety of settings, we experiment with a stock and highly optimised system across two language pairs: @xmath0. we present and analyse our results in sections 7 and 8. our key findings show that: 1. the lottery ticket hypothesis can be applied to prune whole blocks of parameters, instead of removing them separately. ## related work magnitude pruning is one of the simplest algorithms, in which the smallest weights are removed. successfully applied to nmt @xcite , this method works on a coefficient level and often requires retraining to recover the damage done by pruning. further research shows that training a model from scratch with the same structure as the pruned one yields subpar results. finishing training is a necessary step to reduce the size of a model @xcite without too much damage. however, the sparsity of singular weights is generally too low to be efficiently exploited by a cpu or gpu. block sparsity @xcite is more hardware friendly because masked blocks can be skipped entirely. in this paper, we concentrate on a specific case of block sparsity that removes entire attention heads from a model without masking. @xcite applied the lottery ticket hypothesis and other techniques to prune individual coefficients from a transformer for machine translation. in their experiments, a stabilised version of lottery ticket pruning damages translation quality by 2 bleu points while removing 80% of all parameters. they improve upon that further by proposing a mix of lottery ticket and magnitude pruning. in their work, all layers are pruned the same amount, whereas our work prunes globally to reveal which layers can be pruned more aggressively. they aimed to compress the model and did not report any speed results, subsequently clarifying after their presentation that they did not achieve a speed improvement. here, we aim for speed and only marginal improvements to size. rather than prune individual coefficients, we pruned entire heads which can then be removed from the model entirely without even calling a sparse matrix routine. pruning is usually done at the end of training and then requires either retraining or tuning. there is an ongoing research field on integrating pruning into training for example, @xcite pruned weights that have accumulated the lowest total gradients and reduces the memory footprint to allow training much larger models than possible on available hardware. our lottery ticket method does not require to modify a training algorithm and can be easily scripted to work "out of the box" with existing toolkits. @xcite observed that numerous computations in the attention mechanism are redundant with many layers sharing similar distributions. they proposed reusing attention output within adjacent layers in a model, which requires a model to learn which layers should be allowed to share outputs. this reuse of parameters could be understood as a pruning method that concentrates on removing vertical redundancy, in contrast to our research, which is more horizontal. since the attention mechanism is expensive to use in a decoder -with o(n 2 ) complexity looping when generating translations -the better option would be to replace it with less expensive equivalent. in our teacher-student experiments, simpler simple recurrent unit (ssru) @xcite replaces the decoder self-attention mechanism. still, this approach leaves an encoder and context between them unchanged. the lottery ticket pruning is complementary and can remove encoder and context heads on the top of it. looking into an impact of attention on output, @xcite analysed a text classification task whether "high attention weights correlate with greater impact on model predictions". they argued that, in contrast to a simple classification, "for tasks with a much larger output space (such as language modelling or machine translation) . . . almost anything may flip the decision". however, according to our experiments, careful head removal based on their importance does not damage quality. ## background the usual approach to pruning assumes that a model is converged first and pruned second, optionally with continued training. @xcite have shown that iteratively pruning a model uncovers smaller and better quality subnetworks in comparison to pruning just once at the end. still, training a model until convergence at every pruning iteration is too expensive to utilise for most architectures. for this reason, @xcite introduced late resetting and early turnaround. both of these methods combined shorten training time of each step in the iterative lottery scheme. late resetting reverts parameters after pruning back to a checkpoint from early stages of training, not to the starting initialisation. early turnaround means a model does not need to be fully trained to make a pruning decision but can approximate that by doing short training loops. lottery ticket pruning has been applied to natural language processing (nlp) tasks, including nmt @xcite . the winning ticket for that task was "remarkably robust to pruning" of singular weights if embeddings were spared from pruning. however, @xcite noted a linear drop in bleu with sparsit. @xcite analysed the attention mechanism and noticed that the majority of heads are useless: they either do not have linguistically interpretable roles or cannot make reliable choices when making alignments. those heads were pruned by tuning a model with a l 0 regulariser that progressively switched off less essential heads. the l 0 regulariser needs a model to be fully trained first and then pruned while tuned. in contrast, our paper focuses on pruning heads as early as possible in training so that a model can converge with them removed. using their selection heuristic, empirically we can safely prune more heads overall. ## methodology in this section, we describe the lottery ticket approach as well as the decision heuristic based on attention importance @xcite to remove heads in our models. ## setup in order to investigate how effectively pruning works, we concentrate on two language pairs: @xmath1. the first one is considered a low-resource, even with additional back-translated data. in contrast, @xmath2. we trained and decoded our models using the marian machine translation toolkit @xcite . @xmath3. the corpus consists of ~200 000 parallel sentences plus an additional 800 000 sampled from news crawl and backtranslated using a shallow nmt model trained on the existing small bilingual corpora @xcite . we use the development and test sets provided in 2016. we also evaluate on the 2017 and 2018 testsets. the preprocessing follows the steps of normalisation, tokenisation, truecasing using moses scripts, and bpe segmentation @xcite . the vocabulary is shared and contains 36000 words. the architecture is transformer-big @xcite , trained using default recommended settings for such a model in marian toolkit. 4 the models trained until cross-entropy has stopped improving for 10 consecutive validations, and select model checkpoints with highest bleu scores. @xmath4, we follow the workshop on neural generation and translation 2020 efficiency shared task. 5 the shared task specified @xmath5. as is standard for efficient translation, we applied teacher-student training @xcite using the sentence-level system submitted by microsoft to the wmt 2019 news translation task (junczys-dowmunt, 2019). the student models have a standard 6-layers transformer encoder @xcite but the decoder is a faster two-layer simpler simple recurrent unit (ssru) @xcite . the embedding dimension is 256, feed-forward network size is 1536. the models use shared vocabulary of 32,000 subword units created with sentencepiece @xcite . all student models were trained on 13m sentences of available parallel data, using the concatenated english-german wmt testsets from 2016-2018 as a validation set. 6 the models were trained until bleu stopped improving for 20 consecutive validations to overfit the teacher, and the checkpoint with highest bleu scores was selected. since a student model should mimic the teacher as closely as possible, we did not use regularization like dropout and label smoothing. other training hyperparameters were marian defaults for training a transformer base model. 7 student models have sharp probability distributions so we translate using beam size 1. thanks to those settings, the baseline translates about 2335 words per second on a single cpu core. ## experiments the goal is to prune as many heads as possible without damaging translation quality. the pruning procedure has some hyperparameters: the late resetting point, how long to train before making a pruning decision and how many heads to prune each iteration. exploring this space is expensive; we arbitrarily set these to 5-6 saving checkpoints (25k batches for en-de, 12k for tr-en) each pruning iteration have run for 3-4 checkpoints (15k batches for en-de, 8k for tr-en) after which selected attention heads are removed. the number of heads removed is roughly a total number of layers containing attention divided by 2. removing less than that makes pruning slow and removing more in one go results in a unified distribution of attention heads (it usually picks one head per layer) and may be too aggressive in some cases. in each iteration, we change a seed value to make a model see data in different order. we focus on results roughly within 50% to 85% heads removed. this range covers the interesting part from minor to noticeable degradation in translation quality. to evaluate an iteration, heads are pruned as usual then we reset the model back to the late resetting checkpoint and continue training to completion. ## analysis in this section, we further analyse our pruning results in terms of pruning progress and head distribution. we reinitialise our pruned @xmath6, not the architecture itself. head distribution in both architecture or initialisation? to check if the lottery ticket hypothesis is right in the context of our paper, we reinitialise our pruned models while keeping their structure. we compare average bleu difference between pruned ( there is a consistent quality gap between pruned and reinitialised models that widens with sparsity. it confirms the assumptions made by the lottery ticket hypothesis: starting with a larger model and then deliberately selecting attention heads reveals which are the "winning tickets" in the initialisation lottery. ## speed the main objective of our research is to remove heads from a transformer to make inference faster. for this reason, we make a trade-off between a total training time and inference speed, which is particularly useful in an industry production environment. in to compare translation speed, we select the despite attention heads being just a small fraction of all parameters (~5% fewer parameters with about 10% size reduction), pruning them lessens the burden on inference significantly. since all three attention types were pruned in transformerbig experiments, the speed-up is considerablethe model is 1.5 times faster with 0.3 bleu loss. in their paper among many reported models, junczys-dowmunt et al. (2018b) achieved 8.57× speed-up with -0.8 bleu loss on gpu when scaling down from transformer-big teacher to transformer-base student. in another experiment, they gained 1.31× speed-up with -0.6 bleu when using int8 quantisation on cpu. our method is complementary to those as lottery ticket pruning can always remove heads on the top of existing solutions. continuing that line of thought, our small student model translates about 10% faster when pruned. however, it is important to remember that decoder is the key reason why the transformer is slow and it has already been optimized with an ssru. this means there is a smaller margin of improvement in this type of a model. again, attention pruning in this case is complementary and pushes the state-of-the-art even further. just for comparison, we also include the baseline models trained with half (4) and one (1) attention head in each layer (including context attention). the model with just one head everywhere is slightly to compare our work with the state-of-theart in machine translation speed, we submitted @xmath9. these submissions were converged on a larger amount of data to maximize quality. since our method usually selects one head to remove per layer, we experimented with more aggressive and lenient pruning by removing 3 and 6 heads per iteration respectively. these submissions were on the pareto frontier for speed and quality, meaning that no other submission was simultaneously faster and higher quality. the speed-up is about 10% on cpu with 75% encoder heads removed (tab. 8). in terms of on gpu, our best pruned model gains 15% speed-up w.r.t. words per second (wps) losing 0.1 bleu in comparison to an unpruned model (tab. 9). these results show that even when tested on a ## future work in this paper, we applied block-wise pruning to the transformer and its attention mechanism in particular. the natural progress of this research would be to prune other parts of the network -with the lottery ticket approach or not -to see how far block pruning can go without too much impact on quality. furthermore, the heuristic algorithm we chose that decides which heads are not to be removed can definitely be improved on and extended to other types of block-sparsity cases. ## conclusions this paper investigated block-wise pruning of attention heads in the transformer by applying the lottery ticket hypothesis to the problem. we used an iterative approach with pruning done in early stages training. our experiments on nmt have proved that it is possible to remove a significant percentage of all heads (50-72%) in a large transformer with no significant damage to translation quality. since attention mechanism is expensive, especially during inference, reducing the number of heads in a model led to 1.5× speed-up and more if one is willing to sacrifice quality for speed. in the teacher-student regime, the student model with a reduced decoder can be pruned of 75% encoder heads with 0.1-0.2 bleu loss and 10-15% faster translation speed. this shows that lottery ticket pruning is complementary to other methods that reduce inference workload. no matter how a model is trained like, attention heads can be easily removed from it. we hope our paper will inspire further work on attention-sparse architectures. in our paper, we have only shown one example of a heuristic approach -there may be yet to be identified more efficient algorithms better suited to specific tasks, which will result in no need to train overly parametrised models.