Efficient Two-Stage Progressive Quantization of BERT (2022)

Abstract: The success of large BERT models has raised the demand for model compression methods to reduce model size and computational cost. Quantization can reduce the model size and inference latency, making inference more efficient, without changing its stucture, but it comes at the cost of performance degradation. Due to the complex loss landscape of ternarized/binarized BERT, we present an efficient two-stage progressive quantization method in which we fine tune the model with quantized weights and progressively lower its bits, and then we fine tune the model with quantized weights and activations. At the same time, we strategically choose which bitwidth to fine-tune on and to initialize from, and which bitwidth to fine-tune under augmented data to outperform the existing BERT binarization methods without adding an extra module, compressing the binary model 18% more than previous binarization methods or compressing BERT by 31x w.r.t. to the full-precision model. Our method without data augmentation can outperform existing BERT ternarization methods.

## introduction bert @xcite models have demonstrated remarkable performance on nlp tasks. however, their memory and high computational cost make it difficult to fit them onto edge devices with limited resources for inference. recently, a number of methods have been developed to reduce the number of trainable parameters of bert such as @xcite . however, some edge devices require low-precision models for deployment due to the structure of their arithmetic units (cortex-m, 2020). therefore, quantization of bert-based models is needed to avoid/minimize costly floating-point operations. previous works have tried to quantize bert models. ternarybert @xcite used knowledge distillation to transfer the knowledge of the full-precision bert to a bert model with weights that are ternarized into {-1, 0, 1} using 2 bits. binarybert @xcite binarized the weights into {-1, 1} using 1 bit by exploiting the adaptive width property of dynabert @xcite using a method called ternary weight splitting. @xcite managed to outperform @xcite when binarizing acitvations and binarizing weights. however, there is always a performance loss associated with all the aforementioned methods when using ternary weights and to keep competitive performance, data augmentation from @xcite is needed, requiring up to 60x more data as shown in a binarized/ternarized bert has a lot more complex landscape than a full-precision bert @xcite as shown in figure 1 , so training ternary/binary model from scratch could cause the model to get trapped in a poor local minima. therefore, we use a two-stage progressive quantization method inspired by @xcite . in the first stage, we progressively lower the bitwidth of weights, and subsequently fine-tune the model. more specifically, the model with a larger bitwidth is used to initialize the model for fine-tuning on a smaller bitwidth of weights. after reaching the best performance with quantized weights (e.g., binary weights), the bert model is then fine-tuned with quantized activations as the second stage of the quantization method. activations are also progressively quantized until the desired quantization bitwidth in a similar fashion. however, for example, applying the method from @xcite to ternarize bert, naively, costs too much time (or up to 1.67x longer than necessary) with data augmentation used for every bitwidth of weights, with no performance gain as shown in we are the first to ternarize the weights of bert, compressing it by 14.9x while outperforming the full-precision model performance across nearly all glue datasets and without data augmentation it can outperform existing bert ternarization methods. also, we outperform the state-of-the-art bert binarization in performance without adding an extra module, compressing the model 18% more than @xcite . ## preliminaries a quantized model has a set of full precision weights, w. for the forward pass, a quantization function from @xcite is used to ternarize each element w i in the weight matrix w into ŵt i by for binarization, w i is binarized into ŵb i during inference using the binarization method from @xcite by for training, we calculate the gradient of the distillation loss from eq. 16 w.r.t the quantized weights and update the full-precision weight of the quantized model. since the quantization function in eq 1, 2, and 3 are not differentiable, a straightthrough estimator adopted by @xcite is used to back propagate through the quantization function 1, 2, and/or 3. with this, the gradient ∂ x ∂x is approximated as an identity, so ∂l ∂x can be calculated as in our work, we use the knowledge distillation method from @xcite during fine-tuning as shown in eq. 16. which parts of the bert are quantized can be seen on appendix a. ## efficient two stage progressive quantization method with the complex loss landscape of binary/ternary models, training with binary/ternary weights and with quantized activations directly could lead the bert model to converge to a poor local minimum. therefore, we use an efficient two-stage progressive quantization (etspq) method. the details of the etspq method are provided in alg 1. in the first stage, we progressively lower the bitwidth of weights down to binary (or to our desired bitwidth) while fine-tuning it on the target downstream task. the first stage allows us to have a good starting point as we are going to use our binarized/ternarized model weights to then fine-tune the model with quantized activations. we then progressively lower the activation bits. however, to save training time at the same time getting the best possible performance, we choose which bitwidth to finetune on and to initialize from and which bitwidth to fine-tune under augmented data based on the performance of bert under that bitwidth, unlike @xcite . algorithm 1: two-stage progressive quantization input: train/dev data or aug train data; 32-bit teacher model; student model initialized from the teacher model output: student model with k-bit weights and p-bit activations / * progressively reduce bitwidth of weights * / 1 @xmath0, 8, 2, 1]; @xmath1, 8, 4] 2 for k in range(1,4) do 3 -initialize the student model weights from a model with wbits[k-1] bit weights 4 -use augmented data if wbits[k] ≤ 2, else use non-augmented data 5 for @xmath2,...,max epoch do 6 for @xmath3,...,data size do 7 train student and quantize student with wbits[k] bit weights and save the best student model / * progressively reduce bitwidth of activations * / 8 for p in range(1,3) do 9 -initialize the student model weights from a model with wbits[k] bit weights and abits[p-1] bit activations 10 for @xmath4,...,max epoch do 11 for @xmath5,...,data size do 12 train the quantized student model with abits[p] bit activations and wbits[k] bit weights and save the best student model ## experiments we measure the performance of etspq on glue @xcite and compare etspq with the latest ternarization/binarization methods. using the v100 gpu, we fine tune with a batch size of 16 for cola, and 32 for other datasets and we use adamw @xcite with weight decay of 0.01 and learning rate of 5e-5 with a linear learning rate scheduler for five epochs. ## conclusion in this work, we introduced an efficient two-stage progressive quantization method to solve the problem of irregular and complex landscape of bert incurred by the binarization/ternarization of its weights and to reduce the training burden of progressive quantization. we are the first to outperform the state-of-the-art binarization method of bert model without adding an extra module. moreover, our ternary bert can outperform the performance of its full-precision counterpart. also, we can outperform the state-of-the-art bert model with ternary weights on almost all glue datasets without the use of data augmentation.