Pruning Neural Machine Translation for Speed Using Group Lasso

Proceedings of the Sixth Conference on Machine Translation (WMT), pages 1074–1086 November 10–11, 2021. ©2021 Association for Computational Linguistics 1074 Pruning Neural Machine Translation for Speed Using Group Lasso Maximiliana Behnke University of Edinburgh maximiliana.behnke@ed.ac.uk Kenneth Heafield University of Edinburgh kheafiel@inf.ed.ac.uk Abstract Unlike most work on pruning neural networks, we make inference faster. Group lasso regular- isation 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. Exper- iments on top of English→German models, which already have state-of-the-art speed and size, show that two-thirds of feedforward con- nections 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 addi- tional 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.7× faster at the cost 0.9–1.7 BLEU in comparison to the unoptimised baselines. Across language pairs, we see similar sparsity patterns: an ascend- ing or U-shaped distribution in encoder feed- forward and attention layers and an ascending distribution in the decoder. 1 Introduction Making transformer-based machine translation models (Vaswani et al., 2017) faster and smaller is a common requirement for server and mobile deploy- ment. We focus on pruning methods that actually improve speed upon strong baselines. There is a va- riety of work on pruning individual parameters (See et al., 2016; Brix et al., 2020), structures like atten- tion heads (Voita et al., 2019; Behnke and Heafield, 2020), and even whole layers (Sajjad et al., 2020). Unfortunately, much of the prior work on pruning does not report speed or makes inference slower: Brix et al. (2020) achieved no speed-up while Yao et al. (2019) report a 87.5% sparse model took 1.6× as long using cuSPARSE. However, Gale et al. (2020) point out that coefficient-sparse ker- nels like cuSPARSE are highly unoptimised. Even block-sparse kernels are 1.8× slower at 70% spar- sity (Gray et al., 2017) 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 prob- lem remains dense, we sidestep the need for sparse kernels to improve speed. We use group lasso (Yuan and Lin, 2006) regu- larisation, which encourages groups to diminish to- gether, during the usual training procedure. Murray et al. (2019) used group lasso to prune feedforward layers in their submission to the Efficient Transla- tion Task at WNGT 2019 (Kim et al., 2019a). Their submissions, which “eliminate more than 25% of the …

Content truncated to 500 words