Extremely Low Bit Transformer Quantization for On-Device Neural Machine Translation (2020)

Abstract: The deployment of widely used Transformer architecture is challenging because of heavy computation load and memory overhead during inference, especially when the target device is limited in computational resources such as mobile or edge devices. Quantization is an effective technique to address such challenges. Our analysis shows that for a given number of quantization bits, each block of Transformer contributes to translation quality and inference computations in different manners. Moreover, even inside an embedding block, each word presents vastly different contributions. Correspondingly, we propose a mixed precision quantization strategy to represent Transformer weights by an extremely low number of bits (e.g., under 3 bits). For example, for each word in an embedding block, we assign different quantization bits based on statistical property. Our quantized Transformer model achieves 11.8× smaller model size than the baseline model, with less than -0.5 BLEU. We achieve 8.3× reduction in run-time memory footprints and 3.5× speed up (Galaxy N10+) such that our proposed compression strategy enables efficient implementation for on-device NMT.

## introduction transformer @xcite is one of the state-of-the-art approaches for neural machine translation (nmt), and hence, being widely accepted. for example, in wmt19 machine translation tasks, it is reported that 80% of submitted systems have adopted the transformer architecture @xcite . note that high translation quality of transformer models entails a large number of parameters. moreover, the transformer model is inherently much slower than conventional * equal contribution. machine translation approaches (e.g., statistical approaches) mainly due to the auto-regressive inference scheme @xcite incrementally generating each token. as a result, deploying the transformer model to mobile devices with limited resources involves numerous practical implementation issues. to address such implementation challenges with little degradation in translation quality, we study a low-bit quantization strategy for transformer to accomplish high-performance on-device nmt. we note that most previous studies to compress transformer models utilize uniform quantization (e.g. int8 or int4). while uniform quantization may be effective for memory footprint savings, it would face various issues to improve inference time and to maintain reasonable bleu score. for example, even integer arithmetic units for inference operations present limited speed up @xcite and resulting bleu score of quantized transformer can be substantially degraded with low-bit quantization such as int4 @xcite . while determining the number of quantization bits for transformer, it is crucial to consider that each component of transformer may exhibit varied sensitivity of quantization error toward degradation in translation quality @xcite . accordingly, a mixed precision quantization can be suggested as an effort to assign different numbers of quantization bits depending on how each component after quantization is sensitive to the loss function. in addition, as we illustrate later, even assigning different quantization bits for each row of an embedding block can further reduce the overall number of quantization bits of the entire transformer model. our proposed quantization strategy, thus, provides a finer-grained mixed precision approach compared to previous methods, such as @xcite @xcite ) that consider layerwise or matrix-wise mixed precision. accommodating distinguished implementation properties (e.g., latency and translation quality drop) of each component in transformer, we propose the following methodologies to decide precision of a block: 1) in the case of embedding block, statistical importance of each word is taken into account and 2) for encoder and decoder blocks, sensitivity of each quantized sub-layer is considered. the main contributions of this paper are as follows: ## quantization strategy for transformer for transformer, we suggest the following two techniques to decide the number of quantization bits for each block: 1) in the case of embedding block, frequency of each word is taken into account and 2) for encoder and decoder blocks, we find the minimum number of quantization bits for each type of sub-layers that allows reasonable degradation in bleu score after quantization. ## related work previous researches proposed various model compression techniques to reduce the size of transformer models. @xcite apply pruning @xcite to eliminate redundant weights of transformer and report that higher pruning rates lead to greater bleu score degradation. as for pruning, achieving inference speed up is more challenging because unstructured pruning method is associated with irregular data formats, and hence, low parallelism @xcite . uniform quantization for transformer is explored within reasonable degradation in bleu score at int8, while bleu score can be severely damaged at low bit-precision such as int4 @xcite . in order to exploit efficient integer arithmetic units with uniformly quantized models, activations need to be quantized as well @xcite . furthermore, probability mapping operations in transformer, such as layer norm. and softmax, could exhibit significant amount of error in computational results with low precision data type @xcite . ## conclusion in this work, we analyze each block and sub-layer of the transformer and propose an extremely lowbit quantization strategy for transformer architecture. our 2.6-bit quantized transformer model achieves 11.8× model compression ratio with reasonable -0.5 bleu. we also achieve the compression ratio of 8.3× in memory footprints and 3.5× speed up on a mobile device (galaxy n10+).