Pruning Neural Machine Translation for Speed Using Group Lasso (2021)

Abstract: Unlike most work on pruning neural networks, we make inference faster. Group lasso regularisation enables pruning entire rows, columns or blocks of parameters that result in a smaller dense network. Because the network is still dense, efficient matrix multiply routines are still used and only minimal software changes are required to support variable layer sizes. Moreover, pruning is applied during training so there is no separate pruning step. Experiments on top of English->German models, which already have state-of-the-art speed and size, show that two-thirds of feedforward connections can be removed with 0.2 BLEU loss. With 6 decoder layers, the pruned model is 34% faster; with 2 tied decoder layers, the pruned model is 14% faster. Pruning entire heads and feedforward connections in a 12–1 encoder-decoder architecture gains an additional 51% speed-up. These push the Pareto frontier with respect to the trade-off between time and quality compared to strong baselines. In the WMT 2021 Efficiency Task, our pruned and quantised models are 1.9–2.7x faster at the cost 0.9–1.7 BLEU in comparison to the unoptimised baselines. Across language pairs, we see similar sparsity patterns: an ascending or U-shaped distribution in encoder feedforward and attention layers and an ascending distribution in the decoder.

## introduction making transformer-based machine translation models @xcite faster and smaller is a common requirement for server and mobile deployment. we focus on pruning methods that actually improve speed upon strong baselines. there is a variety of work on pruning individual parameters @xcite , structures like attention heads @xcite , and even whole layers @xcite . unfortunately, much of the prior work on pruning does not report speed or makes inference slower: @xcite achieved no speed-up while @xcite report a 87.5% sparse model took 1.6× as long using cusparse. however, @xcite point out that coefficient-sparse kernels like cusparse are highly unoptimised. even block-sparse kernels are 1.8× slower at 70% sparsity @xcite though they did eventually achieve a 1.4× speed-up with "balanced pruning". we propose pruning entire rows or columns and even whole submatrices of a tensor, resulting in a smaller dense matrix. because the inference problem remains dense, we sidestep the need for sparse kernels to improve speed. we use group lasso @xcite regularisation, which encourages groups to diminish together, during the usual training procedure. @xcite used group lasso to prune feedforward layers in their submission to the efficient translation task at wngt 2019 @xcite . their submissions, which "eliminate more than 25% of the model's parameters while suffering a decrease of only 1.1 bleu" were at best 6% faster than their baseline. when tuned for the same quality loss, our method reduces size by 66% and translates 52% faster. moreover, their submissions were slower than higher-quality competitors by an order of magnitude, whereas our baselines are state-of-the-art. too much work @xcite @xcite on efficiency compares a baseline unoptimized system with their optimized system, which is smaller or faster in exchange for some reduction in bleu. what these papers fail to prove is that their method works better than existing methods that also make models smaller or faster in exchange for some reduction in bleu like knowledge distillation @xcite , quantisation, reducing the number of layers, prior work on pruning, or simply training a smaller model. in other words, is the trade-off offered by their method any better than the trade-offs already available, regardless of the type of method? stacking the exist-ing methods produces a variety of data points with different speed and quality. the pareto frontier is the set of data points that a practitioner would choose from: no other data point is simultaneously faster (or smaller) and of higher quality. we argue that a new method's empirical justification should advance the pareto frontier. in this work, we build upon and compare to strong baselines to show the frontier advances. to compare with the state-of-the-art in terms of speed and to investigate the impact this pruning makes on different languages, we build upon @xmath0, @xmath1. we experiment with four architecture variations: a typical decoder with 6 layers and faster variations with shallow decoder of only 1-2 layers. we also include our experiments from the wmt 2021 efficiency shared task. our key findings show that: 1. it is possible to prune entire nodes from feedforward layers early during training, resulting in pareto optimal architectures (quality vs speed). similarly, pruning entire heads on the top of it results in even faster models. 2. different language pairs exhibit similar structural sparsity patterns. 3. pruning during training matches, and sometimes outperforms, retraining the pruned model from scratch. 4. among the @xmath2, the notable examples include a model with a 6-layered decoder being 34% faster at the cost of 0.2 bleu and a model with 12-1 encoder-decoder ratio gaining additional 51% speed-up costing 0.3 bleu. 5. this type of pruning combined with quantisation gives a significant speed boost. our models are 1.9-2.7× faster at the cost of 0.9-1.7 bleu. ## related work extensive research to reduce workload, compress and speed-up neural machine translation models includes methods such as knowledge distillation @xcite , quantisation @xcite , layer approximation @xcite and pruning. for the best results, they can be stacked together to train the efficient state-of-the-art model. in their analysis, @xcite claim that 85% of transformer neurons are redundant across the network. using transfer learning, they find the minimal set of neurons that achieve optimum performance given the task. however, that method requires a fully pretrained model to perform a bruteforce search on it, making overall training time too long. pruning techniques are usually split into two groups: unstructured and structured. unstructured removes individual coefficients. it is straightforward to apply and yields good quality results simultaneously, which makes it popular. unstructured magnitude pruning, while successfully applied to nmt @xcite , often needs retraining to recover from quality damage. moreover, it also requires an efficient matrix multiplication routine to get any speed-up besides size compression. the latest research on combining lottery ticket hypothesis with other methods @xcite sparsified nmt models by 70 to 90% while losing between 0.6 to 3.3 bleu points in quality. they used a sparse matrix representation for compression but did not report any speed gains. on the other hand, structured pruning removes whole layers or groups of parameters, such as blocks @xcite , which makes it easier to optimise on hardware via a special block-sparse matrix multiply @xcite . we apply block sparsity, but the blocks are entire rows or columns so that the usual dense matrix multiply can be used with less overhead. another line of work prunes entire attention heads from a model @xcite ), which we also explore in our approach. @xcite combine unstructured sparsity with a light structure that aims to balance parallel workloads. they introduce a specialised kernel for their structured sparsity. our workloads are easier to balance because they retain density. the idea that different levels of coarseness can be combined may also extend to prune coefficients, rows, columns, and layers simultaneously in future work. @xcite parametrise weight matrices with low-rank factorisations and remove rank-1 components during training, which is said to better preserve linear transformation of uncompressed matrices. they report compressing a transformer-xl language model by 90% with 1.6× speed-up during inference. low-rank approximations preserve density, albeit at the cost of doing serial ma-trix multiplications. @xcite explored a structured dropout that allows users to prune models for inference. unfortunately they fail short on nmt experiments. they call wmt14 en-de a 'competitive benchmark' which it has not been for many years. most problematically, they use tokenised bleu, which has been noted to be harmful and gives false 'boost' of multiple points on tokenised data. they do not specify the tokeniser or script they use either. again, there do not include any report on speed or model sizes despite claiming to have much smaller models as a result. @xcite used group lasso to sparsify a variant of rnn for text classification, which is an easier task to learn than nmt. they have to train until convergence twice, which we avoid. they provide no speed or model size analysis, suggesting that there is no improvement or proper implementation. group lasso has also been previously used by @xcite to compress the delta between a base model and a domain adapted version of the model. they still have to run a full-size model in inference, so they have no overall speed gain. they also have to store the full base model; compression only refers to the delta. in contrast, our work makes the base model faster and smaller. the different goals also mean different groups: they focused on embeddings that update in domain adaptation while we focus on costly parts of the architecture. though we use the same algorithm of group lasso, our method differs in several ways from @xcite . we prune submatrices in addition to rows and columns, though experiments on just rows and columns show better performance than theirs. they pruned only feedforward layers; we see more speed-up from feedforward layers and additionally prune attention. finally, we use the normal adam optimiser (kingma and ba, 2014) instead of proximal gradient descent @xcite . empirically, we find turning regularisation off after some training is important to quality. overall, we achieve a much better tradeoff between quality and speed/compression. most of the methods above need either tuning or retraining, often multiple times. they are usually treated as techniques to compress already existing models. still, there are ongoing research efforts on training a reduced model from start to finish in one go. for example, @xcite pruned weights with the lowest total accumulated gradients and reduced the memory footprint to allow training much larger models than possible on available hardware. some methods prune immediately after initialisation, in either unstructured @xcite or structured @xcite way. our method is orthogonal and is integrated into a training scheme instead. using regularisation to sparsify groups of parameters was introduced by @xcite and has been since then built upon in the machine learning field @xcite . in this paper, we use group lasso in its simplest form to achieve structural sparsity in transformer layers, focusing on inference speed of machine translation. ## methodology to allow regularisation to remove parameters structurally, we need to define how we group parameters. depending on which matrix it is, we treat parameters in its rows, columns or heads as one entity to be pruned together. thus, we apply a group lasso over them. a bias term, if necessary, is treated as a part of regularised groups. we want such a sparsity pattern to emerge early into training so that there is no need to retrain or tune it later. ## pareto trade-off (@xmath3 in this section, we look at the pareto trade-off between the translation quality and speed for all our @xmath4. 3 ). to be fair in our comparison, we trained several simpler baselines with uniformly smaller feedforward dimensions set to {768, 384, 192, 96}. for"12-1" we additionally set heads per layer to 4 to roughly reflect sparsity percentages of pruned models. we pit against each other the said baselines, the pruned models and their reinitialised counterparts. naturally, the models with 6 decoder layers are slower but of a higher quality. however, it is better to switch to fewer decoder layers than to prune too far. our experiments on "12-1" architecture show that its pruned models outperform all others (including all simpler baselines), being a leader in the pareto frontier. ## wmt2021 efficiency shared task to put our method to the final test, we participated in wmt2021 efficiency task foot_3 @xcite . under the task constraints, we trained, pruned and quantised 12-1.tiny and 12-1.micro architectures. we tried two pruning settings, following the directions set in sect. 5.3 and 5.4: rowcol- ## analysis to analyse sparse architecture patterns, we experiment with @xmath5. we prune individual connections in all attention and feedforward layers. in fig. 4 , we present the distribution of remaining parameters for the baseline and the models regularised with λ = {0.1, 0.2, 0.3, 0.4}. since the decoder self-attention is replaced with ssru @xcite , we only show two "pairs" of parameters: encoder self-attention and decoder context attention, with their feedforward counterparts. both encoder layer types follow a similar sparsity pattern making a "u-shape", with the second and third ones being the most aggressively pruned. on the other hand, the decoder parameters are pruned less and less with each subsequent layer. this arrangement of parameters is identical to that exhibited by pruned attention heads in @xcite . in that paper, the attention in the encoder also prunes middle layers, and the context attention retains more heads in further layers. it strongly indicates that the decoder prefers to attend to itself first and confront context later. the estonian architectures, in which we pruned entire attention heads, exhibit a roughly similar structure. for us, this is a strong signal that structural pruning with its architecture search may have a broader generalisation. ## conclusions this paper investigated the structural pruning of a transformer incorporated into a typical training routine. we focused on shredding nodes in feedforward layers and whole attention heads as training progresses. our experiments on knowledgedistilled models with deep and shallow decoders have shown that this type of pruning leads to pareto optimal architectures in quality and speed. moreover, it converges in just one "pass" like a baseline since there is no need to repeat an entire or a part of the training. the resulting sparsity patterns are similar across different languages, with the first and middle layers being the most prioritised during pruning. on the other hand, our experiments on pruning both feedforward and attention layers reveal that some of them, such as the last context attention layer, distinctively avoid being pruned.