E dinburgh’s Submission to the WMT 2022 Efficiency Task (2022)

Abstract: We participated in all tracks of the WMT 2022 efficient machine translation task: single-core CPU, multi-core CPU, and GPU hardware with throughput and latency conditions. Our submissions explores a number of several efficiency strategies: knowledge distillation, a simpler simple recurrent unit (SSRU) decoder with one or two layers, shortlisting, deep encoder, shallow decoder, pruning and bidirectional decoder. For the CPU track, we used quantized 8-bit models. For the GPU track, we used FP16 quantisation. We explored various pruning strategies and combination of one or more of the above methods.

## introduction this paper describes the university of edinburgh's submission to seventh conference on machine translation (wmt2022) efficiency task foot_0 , which measures performance on latency and throughput on both cpu and gpu, in addition to translation quality. our submission focused on the trade-off between these metrics and quality. our submission builds upon the work of last year's submission @xcite . we trained our models in a teacher-student setting @xcite , using the data provided by the organisers. for the students, we used a simpler simple recurrent unit (ssru) @xcite decoder, used a target vocabulary shortlist, and experimented with pruning the student models by removing component and block-level parameters to improve speed. we used 8-bit quantisation for the cpu submission and fp16 quantisation for the gpu submission. we further experimented with ibdecoder @xcite . for running our experiments, we improved upon the marian @xcite machine translation framework by incorporating speed ups for 8-bit matrix multiplication operations, optimizations for pruning neural network parameters on intel cpus, and profiler aided optimisation of various components. ## knowledge distillation we used the provided distilled data to build different student systems. the provided data was distilled through two different processes; for the monolingual input, distilled data was generated using a beam-size of 6. for the parallel data, the teacher ensemble was used to produce the 6 best candidate translations for each input sentence, the candidate most similar to the parallel reference in bleu score was kept as the distilled sentence. we trained student models using just the provided parallel data and identical systems using parallel+monolingual data. our early comparisons showed that the full corpora produced higher quality student systems according to automatic metrics; submitted systems therefore use both parallel and monolingual data. the student models were trained using a validation set consisting of the subset of sentences in the english-german wmt test sets from 2014-2019 that were originally in english. training concluded after reaching 20 consecutive validations without an improvement in bleu score. the student models all used the provided shared sentencepiece vocabulary. we used the default training hyperparameters from marian for the transformer-base model with the learning rate reduced to 0.0002. we explored a number of different configuration in order to find the optimal system on the pareto frontier for speed-quality. we experimented with the following configurations: deep encoders/shallow decoders the majority of the computational cost of the machine translation system falls to the decoder. we can therefore increase drastically the number of encoder layers and decrease the decoder layers without noticeable drop in quality @xcite . tied decoder layers since matrix multiplication is a memory bound problem, we can increase the number of decoder layers, as long as we don't add extra parameters. tied decoder layers allow us to maintain the same memory footprint and keep all the traversed matrices in cache. ssru we replace the self-attention in the decoder with an rnn using the less computation and memory intensive cell ssru. ## pruning attention is a crucial part of the transformer architecture, but it is also computationally expensive. research has shown that many heads can be pruned after training; with further work suggesting that pruning during training can be less damaging to quality. feedforward layers are also expensive and could be reduced. we expand upon our work from the previous year on the group lasso regularisation. we build upon the standard group lasso with a novel approach of aided regularisation. the idea behind it is to use supplementary information to scale the penalties per layer to steer them towards a specific behaviour. in practice, it means adding a new scalar γ alongside an already existing λ: as shown in the equation above, each layer has its individual γ, which gets updated after every backpropagation pass. in order to avoid sudden shits in γ between individual batches, which could make a ratio between perplexities and penalties even more unstable, γ are exponentially smoothed as training progresses: after every batch i, we calculate a local scalar γ j for each layer j based on information gathered during this specific update, which then updates a smoothed global scalar. α is a constant used in exponential average that controls the contribution of a new element in a sequence towards the overall average. we use α = 1e-4 in my experiments. we explore gradient-aided regularisation which scales penalties based on layer gradients. γ scalars should increase as gradients stop flowing through a layer since it indicates that this layer does not contribute to training as much, possibly stopping learning altogether. a layer with small gradients is a good candidate to be regularised more aggressively and vice versa. with w i being a regularised layer and ∇w as accumulated gradients in a model, the gradientaided γ function is defined as: we follow the training regime outlined by @xcite al. (2021 the results include quantised inference with models finetuned for the best quality performance. the 12-1.base model was regularised for 300k batches with λ = 0.05. the 12-1.tiny model was regularised for 200k batches with λ = 0.5. both aforementioned models were pruned in the encoder only. the 8-4.tiny.tied model was regularised for 200k batches with λ = 0.3 with both encoder and decoder layers being penalised. the quality gap becomes larger the harsher pruning is. the base transformer model with 12 pruned encoder layers gets 1.3× faster at the cost of 0.0017 comet point. applying quantisation on the top of it makes translation 3.1× faster in exchange of 0.13 comet points. we applied regularisation onto both encoder and decoder with the "8-4" tied tiny transformer architecture. this pruned and quantised model speeds up by a factor of 2.3× at a 0.8 bleu drop. the most aggressive pruning among the presented results is a tiny transformer with 12 encoder layers with more than 70% parameters removed. this model is 4× faster in comparison to its baseline with 3.3 bleu and 0.09 comet points drop. we note that quantisation struggles with quality on smaller models, both when trained from scratch or pruned. fine-tuning rectifies the problem to some degree, but quality is sacrificed for faster translation in the end. ## fixed point 8-bit quantisation quantising fp32 models into 8-bit integers is a known strategy to reduce decoding time, specifically on cpu, with a minimal impact on quality @xcite @xcite . this year's submission closely follows the quantisation scheme of last year's work @xcite . quantisation entails computing a scaling factor to collapse the range of values to @xcite . for parameters, this scaling factor is computed offline using the maximum absolute value but activation tensors change at runtime. to compute a scaling factor for them, we decoded the wmt16-20 datasets and recorded the scaling factor α(a i ) = 127/max(|a i |) for each instance a i of an activation tensor a. then, for production, we fixed the scaling factor for activation tensor a to the mean scaling factor plus 1.1 standard deviation: these scaling factors were baked into the model file so that statistics were not computed at runtime. we used predominantly intgemm foot_1 for our 8-bit gemm operations, including for the shortlisted output layer. all parameter matrices are quantised to 8-bit offline and the activations get quantised dynamically before a gemm operation. we only perform the gemm operation and the following activation in 8-bit integer mode. after a gemm operation, the output is de-quantized back to fp32. more formally we perform dequantize(σ(a * b+bias)), where the addition of the bias, the activation function σ, and the de-quantisation are applied in a streaming fashion to prevent a round trip to memory. furthermore we make use of intel's dnnl foot_2 for our pruned models, as it performs better than intgemm for irregular sized matrices. unfortunately, dnnl doesn't support streaming de-quantisation, bias addition or activation function application. quantisation does not extend to the attention layer, which is still computed in fp32. the reason being is that in the attention layer, both the a and b matrices of the gemm operation would need to be quantised at runtime, which makes the quantisation too expensive. we note that we only perform the gemm operations in 8-bit integers. similar to previous' year's submission, we performed quantisation fine-tuning for some 8-bit models, where we perform a small amount of training with low learning rate and a damaged gemm implementation that simulates the quantised output. we found that this helps regain some quality, especially in smaller models. ## shortlisting the single most expensive computation in machine translation is the cost of the output layer. we can reduce the computation if we only take into account likely output tokens, reducing the output layer size from 32000 to something much more manageable like 500-2000. we used ibm model based shortlisting @xcite . this lexical shortlist is straightforward to work with, but it is limiting in the sense that it doesn't capture well idioms and favours more literal translations. the hyper-parameters that control the size of this shortlist are: the number of most frequently targeted words included, and the number of probable translations for each token in the input. this year we increased the number of most-frequent and aligned tokens to 100,100 (from 50,50 in the previous year) in order to improve quality. the shortlist is built using alignment models trained on a specific corpora. the total number of tokens in a shortlist considered is influenced by the size of the current batch: the shortlist produced is the union of probable translations for each input token and overall most-likely candidates. in latency scenarios, where batches are a single sentence, a small shortlist is more detrimental to the quality than for larger batches, such as in throughput scenarios, that benefit from inclusion of more candidate tokens. similarly, this approach benefits when inputs are batched. ## ibdecoder the sequential nature of autoregressive decoding forms an inference bottleneck, hurting decoding parallelisation and latency. a popular method to break this bottleneck is to allow the parallel prediction of multiple target tokens per step through semi-or non-autoregressive modelling @xcite with a quality tradeoff. we experimented with interleaved bidirectional decoder @xcite , a variant of semiautoregressive decoder that predicts target tokens from the left-to-right and the right-to-left directions ## quality issues quantisation applied to small models, especially those that were pruned, struggles with maintaining the quality. for example, as can be seen in tab. 2, quantisation on top of pruned models damages the quality from 0.1 to 0.3 comet points. this gap is more evident in smaller architectures such as tiny or tied tiny. we hypothesise that the fewer parameters there are in a model, the more difficult it is to optimise through pruning and/or quantisation, or using a bidirectional generation. ibdecoder suffers from the pruning and quantisation particularly on the comet scores as shown in tab. 3. we compared several sentences that showed little difference in bleu but significant difference in the comet scores and tried to see what went wrong ( it seems quantisation doesn't always result in a a worse transaltion. in the second example the quantised ibdecoder produces a more complicated, but overall much better translation than the ibdecoder, which also suffers from a repetition error. this suggests that the model is quite brittle and very susceptible to small changes. ## software improvements we built our work using the marian machine translation framework, making some improvements on top of the submission from last year: avx512 inrinsics we implemented hand crafted intrinsics for various arithmetic operations, resulting in .5% improvement in performance. ## conclusion we participated in all tracks of the wmt 2022 efficiency task and we submitted multiple systems that have different trade-offs between speed and translation quality. for the cpu submission we used 8-bit integer decoding and a combination of pruned and non-pruned system, together with a lexical shortlist in order to reduce the computational cost of the output layer. we also experimented with ibdecoder in both cpu and gpu setting.