Fully Quantized Transformer for Machine Translation (2020)

Abstract: State-of-the-art neural machine translation methods employ massive amounts of parameters. Drastically reducing computational costs of such methods without affecting performance has been up to this point unsuccessful. To this end, we propose FullyQT: an all-inclusive quantization strategy for the Transformer. To the best of our knowledge, we are the first to show that it is possible to avoid any loss in translation quality with a fully quantized Transformer. Indeed, compared to full-precision, our 8-bit models score greater or equal BLEU on most tasks. Comparing ourselves to all previously proposed methods, we achieve state-of-the-art quantization results.

## introduction the idea of using neural networks for machine translation was only recently proposed @xcite @xcite . nonetheless, the approach became the state-of-the-art in the field @xcite @xcite . a key element of this success was to allow the decoder to attend to all hidden states of the encoder @xcite . a few variations to this additive attention mechanism have been proposed, such as multiplicative and self-attention @xcite @xcite . the latter formed the basis of the transformer network @xcite , which achieved state-of-the-art results in machine translation. inspiring a new wave of work, numerous natural language processing tasks reached new heights @xcite . unfortunately, these models use an enormous amount of parameters. inference on resource-limited hardware such as edge-devices is thus impractical. a solution to reduce the computational burden of these networks is to lower numerical precision. consequently, numerical values can be represented using fewer bits @xcite . this method called quantization has the advantage of providing good compression rates with minimal loss in accuracy. it is also conveniently supported by most hardware. properly quantizing the transformer would allow computational speed gains at inference, as well as deployment on more constrained devices. in this work, we propose a quantization-aware training strategy for the entire transformer architecture. our method is easy to implement and results are consistent with the full-precision transformer. we test our approach on multiple translation tasks such as wmt14 en-fr and wmt14 en-de and obtain state-of-the-art quantization results. in comparison with full-precision, our quantized models score equal or higher bleu on most tasks. we are, to the best of our knowledge, the first to show that the transformer architecture can be fully quantized without impairing translation quality. we also perform an ablation study and show that quantizing specific components of the transformer improves bleu score. ## background in this section, we review a broad spectrum of quantization and pruning methods for neural network compression. ## related work recently, simple quantization solutions have been applied to the transformer. @xcite apply k-means quantization and binarization with two centroids over the weights of the network. for both methods, a look up table associated with each quantized layer is used to map indices to their corresponding centroids. similarly, @xcite compares binary, 4 and 8-bit uniform quantization of the transformer weights. a big disadvantage with quantizing only the weights of a network is that operations must still be performed in full-precision. even though the parameters' memory usage is reduced, these constantly have to be converted back to full-precision. achieving quantization of both weights and activations is much more beneficial. the first attempt at doing so for the transformer applies 8-bit quantization on weights and inputs of feed forward layers and binarizes the (q, k) input of the multi-head attention @xcite . the scaling factor √ d k is approximated by a constant which can be computed as a right bitshift. the method resulted in a huge drop in translation accuracy. achieving better performance, bhandare et al. (2019) quantize certain matmul operations and use the kl divergence to estimate the most suited parameters for each quantization range. they restrain from quantizing all matmuls, reporting poorer results in accuracy. aside from translation, the concurrent work by @xcite quantizes the embedding and fully connected layers of bert @xcite . the softmax and layernorm operations are kept in full-precision. on the glue benchmark, their loss in accuracy is minimal compared to the original model. all of these methods omit quantizing the whole transformer architecture, resulting in suboptimal computational efficiency. furthermore, these solutions all fail to avoid impairing translation quality. our method achieves both. ## experiments in this section, we present the results of our full quantization scheme on various tasks. we first compare our method on a machine translation setup. then we present the results of numerous ablation studies. we also compare the impact of delaying quantization on translation quality. finally, we evaluate our method on two language model tasks and experiment with node pruning. ## pruning useless nodes we experiment with node pruning our transformer models. once the model is fully trained and quantized, we can further compress it by removing useless nodes. by useless, we mean nodes which do not cause any loss in translation quality when removed. we choose to prune nodes instead of independently pruning weights. the latter method usually requires special hardware or software to the only nodes of the transformer which can be removed without causing alterations to other components of the network are the nodes in between the two layers of each feed-forward network. fortunately, these consist of a substantial portion of the model's weights. in the case of the base transformer, for a respective vocabulary of size 37000 and 32000, 39.96% and 41.65% of the total weights are owned by the feed-foward networks. this number grows to 47.03% and 48.18% in the big transformer. to evaluate which nodes can be safely pruned without affecting translation quality, we estimate x max for each node of the relu output over a few hundred steps. this is done on the training set, using the fully trained model and keeping all other weights frozen. these x max are computed before quantizing the relu output and do not replace the ones used by the quantization process. figure 3 in the appendix shows the histogram of these running estimates for one relu layer in the encoder and one in the decoder. all other relu layers share the same pattern, where in the encoder there are always multiple x max close to 0. this does not happen in the decoder. once the running estimates are computed, we prune its corresponding node if x @xmath0. we empirically found @xmath1.025 to work well, with higher thresholds causing bleu to quickly decay. no retraining of the model is performed after pruning nodes. using this method, we can further compress the transformer without affecting bleu scores. our approach has the advantage of being adaptive, ## conclusion we proposed a full quantization strategy for the transformer architecture. our objective was to ex-ploit hardware resources as efficiently as possible, quantizing all operations which could provide a computational speed gain. with fullyqt, we achieve higher bleu scores than all other quantization methods for the transformer on multiple translation tasks and avoid any loss in bleu compared to full-precision. specifically, out of 35 experiments, 8-bit quantization performed better than full-precision in 21 cases. if instead of minimizing inference time, one wants to maximize translation accuracy, then applying quantization to only certain components of the transformer seems to be the best option. indeed, our ablation study showed than bleu score could increase even more when only specific elements of the transformer were quantized. further gains might be possible, but supplementary experiments would be necessary to determine the best combination. we plan on extending our work to variations of the transformer, as well as further exploring the compression of these networks. multi-head attention add & norm feed forward input embedding add & norm positional encoding n× inputs masked multi-head attention add & norm feed forward output embedding add & norm positional encoding n× outputs (shifted right) multi-head attention add & norm linear softmax output probabilities quantize relu x quantize linear quantize linear figure 1: (left) fully quantized transformer, (right) feed-forward. matmul scale mask (opt.) softmax matmul q k v quantize quantize quantize quantize linear linear scaled dot-product attention linear linear scaled dot-product attention concat h k q v linear linear quantize linear quantize quantize quantize figure 2: (left) scaled dot-product attention, (right) multi-head attention.