Understanding and Overcoming the Challenges of Efficient Transformer Quantization (2021)

Abstract: Transformer-based architectures have become the de-facto standard models for a wide range of Natural Language Processing tasks. However, their memory footprint and high latency are prohibitive for efficient deployment and inference on resource-limited devices. In this work, we explore quantization for transformers. We show that transformers have unique quantization challenges – namely, high dynamic activation ranges that are difficult to represent with a low bit fixed-point format. We establish that these activations contain structured outliers in the residual connections that encourage specific attention patterns, such as attending to the special separator token. To combat these challenges, we present three solutions based on post-training quantization and quantization-aware training, each with a different set of compromises for accuracy, model size, and ease of use. In particular, we introduce a novel quantization scheme – per-embedding-group quantization. We demonstrate the effectiveness of our methods on the GLUE benchmark using BERT, establishing state-of-the-art results for post-training quantization. Finally, we show that transformer weights and embeddings can be quantized to ultra-low bit-widths, leading to significant memory savings with a minimum accuracy loss. Our source code is available at https://github.com/qualcomm-ai-research/transformer-quantization .

## introduction recently, transformer architectures have shown remarkable improvement in many natural language processing (nlp) tasks and beyond. based on the original transformer @xcite , language models pre-trained from large corpora of unlabeled text, such as bert @xcite , roberta @xcite , xlnet (yang * qualcomm ai research is an initiative of qualcomm technologies, inc. et al., 2019), transformer-xl @xcite , gpt family @xcite @xcite , have become an indispensable building block in modern nlp pipelines. they are also increasingly adopted in other areas, including vision @xcite @xcite and audio @xcite . despite cutting edge results in many applications, pre-trained transformer-based models are extremely large, sometimes exceeding billions of parameters. hence, efficient deployment of these models on resource-constrained embedded systems, and even sometimes in data centers, has become an important problem due to high latency and prohibitively large memory footprint and energy consumption. one effective method to tackle this problem is neural network quantization. quantization reduces memory consumption by using low-bit precision for weight and activation tensors. is also reduces inference time, and improves energy efficiency by employing low-bit fixed-point arithmetic instead of floating-point arithmetic @xcite . quantization, however, is not free. it introduces additional noise in the network that can lead to a drop in the model's performance. while prior work has demonstrated the feasibility of integeronly inference for computer vision models @xcite @xcite @xcite @xcite @xcite , there is relatively little work done on quantizing nlp models @xcite , and specifically on transformer models. understanding the challenges of transformer quantization and designing a robust and easy-touse quantization pipeline for them constitute the primary goal of this paper. the contributions of our work include: we evaluate our proposed solutions on eight different nlp tasks from the well-known glue benchmark. our techniques set a new state-ofthe-art of post-training quantization and per-tensor quantization-aware training for the bert model. to the best of our knowledge, this is the first work for the bert-like transformer quantization with a strong focus on post-training quantization. the presented method is not exclusive to bert and is easily applicable to other pre-trained transformer models. quantization one of the most powerful ways to decrease the computational time and memory consumption of neural networks is quantization, which uses low-bit representations for weight and/or activation tensors. when moving from 32 to 8 bits, the memory overhead of storing tensors decreases by a factor of 4, while the computational cost for matrix multiplication reduces quadratically by a factor of 16. low-bit fixed-point representations, such as int8, further reduce the energy consumption since the fixed-point operations are more efficient than their floating-point counterparts @xcite . however, exact latency improvements and energy savings are highly dependent on the target hardware. therefore, we focus in this work on achieving high memory and compute reduction while maintaining acceptable model accuracy and do not measure actual on-device performance gains. we will cover relevant basics of quantization here, for a more comprehensive overview of neural network quantization please refer to @xcite . a commonly used scheme for quantization is uniform affine or asymmetric quantization @xcite @xcite because it allows for efficient implementation of fixed-point arithmetic. it is defined by bitwidth @xmath0, scale factor @xmath1, and zero-point @xmath2. we simulate the quantization process in floating-point according to @xcite . quantizing a real-valued tensor x is performed by first mapping it to an unsigned integer grid: it is possible to approximately recover the realvalued input x through an operation that is often referred to as de-quantization: in the case of symmetric quantization, we restrict the quantization grid to be symmetric around z. it is common to have a single set of quantization parameters per tensor, known as per-tensor quantization. one could also increase the quantization granularity by defining separate quantizers for individual segments of a tensor. this will improve the accuracy of a network, but at the cost of an additional compute and memory overhead. an important class of quantization methods is post-training quantization (ptq) algorithms, which take a pre-trained fp32 network and convert it directly into a fixed-point network without the need for the original training pipeline (krishnamoorthi, 2018). a vital step in the ptq process is finding good quantization ranges for each quantizer. one way of doing this is static range estimation, which determines quantization parameters for the network by passing a few batches of calibration data through the model before inference. it yields more efficient inference since all the quantization parameters are known in advance and fixed. several of the most common range estimators include: current min-max or simply min-max, uses the full dynamic range of the tensor @xcite @xcite ; running min-max uses exponential moving average of the min and max over multiple batches @xcite ; mse finds quantization parameters that minimize mean squared error between quantized and floating-point tensors @xcite ). an alternative to ptq is to train a neural network with the simulated quantization operations in the network, known as quantization-aware training (qat, @xcite @xcite . it allows the model to better adapt to the introduced quantization noise compared to ptq, at the cost of longer training times, the need for labeled data and doing a hyper-parameter search. gradients through the nondifferentiable quantization step are usually approximated using the straight-through estimator @xcite . ranges for both weights and activations can be set using ptq range estimators or learned jointly with the weights during training, as in esser et al. (2019); jain et al. (2019). finally, it is possible to assign different bitwidths for different layers or parts of the network, a technique known as mixed precision @xcite @xcite @xcite . transformer quantization junczys-dowmunt et al. (2018) applied knowledge distillation and 8-bit post-training quantization to speed up transformer models for neural machine translation. @xcite also applied 8-bit posttraining quantization to the transformer model for machine translation and demonstrated how to utilize specialized hardware to accelerate the inference process. @xcite proposed an 8-bit quantization scheme for bert-like models and achieves compression of up to 25% of the original model size. @xcite applies mixed-precision quantization on bert, where they assign a different precision to different layers according to their sensitivity defined by hessian information. @xcite proposed a fully integer-only arithmetic inference scheme based on second-order polynomial approximations for gelu, softmax, and layernorm non-linearities. some examples of transformer-based model quantization with alternative quantization schemes include @xcite ; darvish @xcite . note that all the mentioned approaches for bert-like transformer quantization employ some form of qat and either do not discuss ptq alternatives or only use them as weak baselines. ## problem investigation first, we investigate what happens when we apply standard 8-bit post-training quantization to the bert model and evaluate it on eight downstream tasks from the glue benchmark @xcite . to quantize fine-tuned models, we use uniform affine quantization with static range in to find which part of the network is the most problematic, we perform an ablation study in which we do not quantize specific activations. the results are summarized in to understand why quantizing the residual ffn sum is so detrimental, we look at activation tensors in the problematic 11th layer. first, from figure 2a , we note that ffn's input and output have radically different dynamic ranges (note the scale for y-axes) due to strong outliers in the output tensor. applying per-tensor quantization for the ffn's residual sum is likely to cause a notable error because of the following trade-off between the range and the precision. on the one hand, using the full dynamic range for the ffn's output will lead to a very coarse quantization of its input. on the other hand, using higher precision for the input will cause informa-tion loss in the output due to aggressive clipping of its range. we also notice a correlation of outliers with special [sep] tokens. in addition to that, from figure 2b , we observe that only a few embedding dimensions are consistently responsible for these outliers across many data points. in appendix d, we show that this is the case for all layers of bert-base and all glue tasks. furthermore, we show that a similar issue is also present in multiple architectures and training objectives, including pre-trained bert-large, roberta, and distilroberta @xcite , and mo-bilebert @xcite . further analysis suggests that structured outliers in the ffn's residual connections lead to structured outliers in query-key multiplications in specific attention heads in the next attention layer, causing most of the tokens to attend to the special [sep] token. see appendix a for more details on this. ## methodology in this section, we introduce our proposed techniques for bert-like model quantization. motivated by our findings from section 3, we consider three ways of efficient bert quantization -posttraining mixed precision, a new per-embeddinggroup activation quantization, and quantizationaware training. each of these three methods comes with its own set of trade-offs, which is why we present all three. the reader can pick an appropriate solution for their practice. as before, we employ uniform affine quantization and static activation ranges, which are either estimated in ptq or learned during qat, as described in section 2. mixed precision ptq as seen in section 3, not all parts of bert are equally sensitive to the quantization noise. thus, selecting a higher bit-width for sensitive tensors can lead to better accuracy while efficiently keeping all the other tensors in 8-bit or per-embedding-group ptq as discussed in section 2, another way of improving the performance of the quantized model is to increase the quantization granularity. based on our observation from figure 2b , that the most problematic outliers in activation tensors are in few designated embedding dimensions, we consider having distinct quantization parameters for individual embedding dimensions or groups of embedding dimensions, as shown in figure 3 . we start by describing per-embedding activation quantization. in bert-like models, an intermediate hidden activation tensor x has a shape (b, t, d), where b is the batch size, t is the sequence length, and d is the number of embedding dimensions (@xmath3, @xcite . inspired by per-channel weight quantization @xcite , we can have distinct scaling factors and zero-points per embedding dimension instead of having two scalars for the whole tensor. in this case, we can collectively denote the quantization parameters by vectors s, @xmath4. the rest of the quantization machinery works as before, including range estimation, with the only difference that equations ( 1 ) and ( 2 ) are now with broadcasting along the last dimension. the proposed scheme should alleviate the activation quantization issue since the outlier embedding dimensions will no longer dominate the ranges of other embedding dimensions. note, however, that full per-embedding activation quantization will lead to a more expensive computational graph. to illustrate why, consider a matrix-vector multiplication wx, which in case of per-tensor quantization (and assuming @xmath5: a crucial detail here is that we can factor a common factor s w s x out of the summation. the sum is then efficiently calculated using integer-only arithmetic. in case of per-embedding activation quantization ( @xmath6, the matrix-vector multiplication becomes instead: here it is no longer possible to take the scaling factor out of the summation and perform a single re-scaling of the result. instead, one has to perform repeated intermediate re-scalings on the accumulator. to alleviate the overhead of constant re-scaling, we introduce per-embedding-group (peg) quantization, where we split the activation tensor into k evenly sized groups along the embedding dimension and share quantization parameters among elements in the same group: the proposed scheme might not be natively supported on all target device, but there is an efficient way to implement it using only per-tensor quantization. before quantizing the output of the first lay-ernorm (figure 1 ), we split the output tensor based on embedding groups into k individual tensors. we also accordingly split columns of the first linear layer and rows of the second layer and decompose them into k smaller linear layers each. the outputs of the first set of layers are elementwisesummed, and the outputs of the second set of layers are concatenated before the residual sum. with this functional equivalent rewriting, all operations can be performed using standard per-tensor quantization. to ensure all outliers end up in the same group, we employ a deterministic range-based permutation of the embedding dimensions. similar to range estimation for the activation quantization, we pass some calibration data through the unquantized network and record the dynamic range r j := max(x :,:,j ) -min(x :,:,j ) for each embedding dimension j. next, we define k evenly sized groups based on indices in arg sort(r). during the range estimation phase, we determine a separate quantization range for each group. the sorting and grouping need to happen only once before the range estimation phase and deployment to the target. peg quantization with permutation can still be simulated on hardware that only supports pertensor operations. first, we can share the same permutation for ffn's input, output and sum since we expect the outliers in the output dominate the ones from the input. second, we use the permutationequivariant properties of layernorm and linear layers (weights are permuted accordingly before inference). we first permute the output of the first layernorm, proceed as described above, and then apply inverse permutation before the next layer-norm. note that the peg quantization has a negligible memory overhead, introducing only d + 2 • 3 • k extra parameters per attention layer (permutation indices and scale & zero points per group for ffn's input, output, and sum), which is less than 0.04% of the total size of bert-base model. quantization-aware training finally, we consider a variant of qat with learnable ranges for both weights and activations by adapting the procedure from esser et al. (2019 comparison of methods we summarize different trade-offs for the proposed techniques in for qat experiments, we initialize all quantization parameters from the ptq setup described above. similarly to full-precision fine-tuning, we use adam (kingma and ba, 2014) and a maximum sequence length of 128, with padding using a special [pad] token for shorter sequences. we use a typical learning rate schedule from the transformer literature @xcite @xcite -a linear warmup for the first 10% of training steps followed by a linear decay to zero. we perform a hyper-parameter search over the maximum learning rate, batch size, number of epochs, and the self-attention dropout rate for every task and report the best median score over three runs with different random seeds. for reproducibility, we included more details on the search space and selected hyper-parameters in appendix b.3. ## conclusions in this paper, we explored quantization for bertlike transformers. we showed that these models have unique quantization challenges -namely, high dynamic activations ranges that are difficult to represent with a low bit fixed-point format. these activations contain structured outliers in the residual connections that encourage specific model behavior, such as attending to the special [sep] token. motivated by our findings, we proposed three solutions, one based on mixed precision quantization, a novel per-embedding-group quantization, and quantization-aware training. each of these methods has its own set of trade-offs in terms of accuracy, ease of use, and model size. our techniques overcome the dynamic range issues and set a new state-of-the-art for ptq and per-tensor qat on glue downstream tasks. finally, we achieved 4-bit weight and 2-bit token embedding quantization with less than 0.8% drop in terms of glue score, leading to significant memory and compute savings.