Efficient Machine Translation with Model Pruning and Quantization (2021)

Abstract: We participated in all tracks of the WMT 2021 efficient machine translation task: single-core CPU, multi-core CPU, and GPU hardware with throughput and latency conditions. Our submissions combine several efficiency strategies: knowledge distillation, a simpler simple recurrent unit (SSRU) decoder with one or two layers, lexical shortlists, smaller numerical formats, and pruning. For the CPU track, we used quantized 8-bit models. For the GPU track, we experimented with FP16 and 8-bit integers in tensorcores. Some of our submissions optimize for size via 4-bit log quantization and omitting a lexical shortlist. We have extended pruning to more parts of the network, emphasizing component- and block-level pruning that actually improves speed unlike coefficient-wise pruning.

## introduction this paper describes the university of edinburgh's submission to sixth conference on machine translation (wmt2021) 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 edinburgh's en-de system submitted to the wmt2021 news translation task as the teacher model. 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 blocklevel parameters to improve speed. we further experimented with quantizing into smaller numerical formats, including fixed point 8-bit quantization on the cpu, and both 8-bit and log based 4-bit quantization on the gpu, as well as post-quantization fine-tuning of 4-bit quantized models. 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 exploring tensorcores on the gpu. ## training teacher models we used edinburgh's en↔de systems submitted to the wmt 2021 news translation task as teacher models @xcite . we trained transformerbig models @xcite , using a shared 32k sentencepiece @xcite vocabulary, built in three stages: corpus filtering, back-translation and fine-tuning. the models achieved 29.90 and 51.78 bleu on @xmath0, with multiple references). we used sequence-level knowledge distillation @xcite to synthesize forward, backward, and backward-forward translations using the teachers. we filtered the synthesized parallel data using handcrafted rules foot_1 , followed by removing bottom 5% according to cross-entropy per word on the generated side using kenlm @xcite . ## knowledge distillation we ran experiments using different combinations of teacher-synthesized corpora. one variant included all of the synthesized data: parallel, monolingual backward and forward as well as backwardforward (aji and heafield, 2020b). another variant excludes only the fully-synthetic monolingual backward-forward data, while the final variant used parallel data only. all student models were trained using a validation set consisting of the subset of sentences in the english-german wmt test sets from 2015-2019 that were originally in english. training concluded after reaching 20 consecutive validations without an improvement in bleu score. the student models used the same shared vocabulary as the teacher ensemble. during decoding, we used a lexical shortlist @xcite @xcite of the top 50 most probable alignments, combined through a union with the top 50 most frequent vocabulary items. other than this, we used the default training hyperparameters from marian for the transformer-base model. each of the student models used transformer encoders @xcite and rnn-based decoders with simpler simple recurrent unit (ssru) @xcite . several different architectures were explored; these differ in the number of encoder and decoder blocks as well as in the sizes of the embedding and ffn layers. further to this, some of our transformer architectures use a modified attention matrix of shape (d emb , n head ×d head ) rather than the typical (d emb , d emb ). in all cases we use 8 transformer heads per layer, and set d @xmath1. the student architectures are summarized in ## conclusion we participated in all tracks of the wmt 2021 efficiency tracks and we submitted multiple systems that have different trade-offs between speed and translation quality. we performed ample hyperparameter tuning and exploration in order to take advantage of gpu tensorcores for decoding, but unfortunately we couldn't beat our optimised fp16 baseline. for the cpu submission we used 8bit 6 https://github.com/marian-nmt/ marian-dev/pull/743 7 https://developer.nvidia.com/blog/ programming-tensor-cores-cuda-9/ 8 https://github.com/nvidia/cutlass 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 largest gemm in decoding -the output layer.