T ernary BERT : Distillation-aware Ultra-low Bit BERT (2020)

Abstract: Transformer-based pre-training models like BERT have achieved remarkable performance in many natural language processing tasks. However, these models are both computation and memory expensive, hindering their deployment to resource-constrained devices. In this work, we propose TernaryBERT, which ternarizes the weights in a fine-tuned BERT model. Specifically, we use both approximation-based and loss-aware ternarization methods and empirically investigate the ternarization granularity of different parts of BERT. Moreover, to reduce the accuracy degradation caused by lower capacity of low bits, we leverage the knowledge distillation technique in the training process. Experiments on the GLUE benchmark and SQuAD show that our proposed TernaryBERT outperforms the other BERT quantization methods, and even achieves comparable performance as the full-precision model while being 14.9x smaller.

## introduction transformer-based models have shown great power in various natural language processing (nlp) tasks. pre-trained with gigabytes of unsupervised data, these models usually have hundreds of millions of parameters. for instance, the bert-base model has 109m parameters, with the model size of 400+mb if represented in 32-bit floating-point format, which is both computation and memory expensive during inference. this poses great challenges for these models to run on resource-constrained devices like cellphones. to alleviate this problem, various methods are proposed to compress these models, like using low-rank approximation @xcite , weight-sharing @xcite , knowledge distillation @xcite @xcite ), ⇤ authors contribute equally. pruning @xcite @xcite , adaptive depth and/or width @xcite , and quantization @xcite @xcite . compared with other compression methods, quantization compresses a neural network by using lower bits for weight values without changing the model architecture, and is particularly useful for carefully-designed network architectures like transformers. in addition to weight quantization, further quantizing activations can speed up inference with target hardware by turning floating-point operations into integer or bit operations. in @xcite , 8-bit quantization is successfully applied to transformer-based models with comparable performance as the full-precision baseline. however, quantizing these models to ultra low bits (e.g., 1 or 2 bits) can be much more challenging due to significant reduction in model capacity. to avoid severe accuracy drop, more complex quantization methods, like mixed-precision quantization @xcite and product quantization (pq) @xcite , are used. however, mixed-precision quantization is unfriendly to some hardwares, and pq requires extra clustering operations. besides quantization, knowledge distillation @xcite which transfers knowledge learned in the prediction layer of a cumbersome teacher model to a smaller student model, is also widely used to compress bert @xcite @xcite . instead of directly being used to compress bert, the distillation loss can also be used in combination with other compression methods @xcite @xcite , to fully leverage the knowledge of teacher model. in this work, we propose ternarybert, whose weights are restricted to { 1, 0, +1}. instead of directly using knowledge distillation to compress a model, we use it to improve the performance of ternarized student model with the same size as the teacher model. in this way, we wish to transfer the knowledge from the highly-accurate teacher model to the ternarized student model with smaller capacity, and to fully explore the compactness by combining quantization and distillation. we investigate the ternarization granularity of different parts of the bert model, and apply various distillation losses to improve the performance of ternarybert. figure 1 summarizes the accuracy versus model size on mnli, where our proposed method outperforms other bert compression methods. more empirical results on the glue benchmark and squad show that our proposed ternarybert outperforms other quantization methods, and even achieves comparable performance as the full-precision baseline, while being much smaller. ## approach in this section, we elaborate on the method of using knowledge distillation to train ternarybert, the weights of which take values in { 1, 0, +1}. let the full-precision weight in the bert model be w, where @xmath0. the corresponding ternarized weight is denoted as ŵ = q w (w) where q w is the weight ternarization function. the whole framework, which we call distillation-aware ternarization, is shown in figure 2 . specifically, at the t-th training iteration, we first ternarize the weights w t in the student bert model to ŵt . then we do the forward pass with the ternarized model. after that, the gradient of the distillation loss w.r.t. the quantized weights @l @ ŵt is computed. as is shown in @xcite , it is important to keep the full-precision weight during training. hence, we use the full-precision weight for parameter update: w t+1 = updateparameter(w t , @l @ ŵt , ⌘ t ), where ⌘ t is the learning rate at the t-th iteration. in the following, we will first introduce what and how to quantize in section 3.1. then in section 3.2, we introduce the distillation loss used to improve the performance of the ternarized model. ## experiments in this section, we evaluate the efficacy of the proposed ternarybert on both the glue benchmark @xcite and squad @xcite . the experimental code is modified from the huggingface transformer library. 2 we use both twn and lat to ternarize the weights. we use layer-wise ternarization for weights in transformer layers while row-wise ternarization we compare our proposed method with q-bert @xcite and q8bert @xcite using their reported results. we also compare with a weight-ternarized bert baseline q2bert by modifying the min-max 8-bit quantization to min-max ternarization using the released code of q8bert. 3 for more direct comparison, we also evaluate the proposed method under the same 8-bit quantization settings as q-bert and q8bert. when the weights are quantized to 8bit, we use layer-wise scaling for both the weights in transformer layers and the word embedding as 8-bit quantization already has high resolution. ## conclusion in this paper, we proposed to use approximationbased and loss-aware ternarization to ternarize the weights in the bert model, with different granularities for word embedding and weights in the transformer layers. distillation is also used to reduce the accuracy drop caused by lower capacity due to quantization. empirical experiments show that the proposed ternarybert outperforms stateof-the-art bert quantization methods and even performs comparably as the full-precision bert. (a) full-precision bert. (b) ternarybert. (a) full-precision bert. (b) ternarybert. (a) full-precision bert. (b) ternarybert. (a) full-precision bert. (b) ternarybert. figure 9 : attention patterns of full-precision and ternary bert trained on sst-2. the input sentence is "old-form moviemaking at its best."