Quantization Aware Training, ERNIE and Kurtosis Regularizer: a short empirical study

TL;DR

This paper empirically investigates quantization aware training (QAT) for pre-trained models like Ernie, highlighting challenges and proposing a simple regularizer approach to improve INT8 accuracy in low-precision inference.

Contribution

It identifies why existing regularizers do not work out-of-the-box for Ernie and proposes a basic method to enhance quantization robustness for pre-trained models.

Findings

Initial results show increased INT8 accuracy with the proposed regularizer.

Analysis explains the incompatibility of existing regularizers with Ernie.

Provides practical insights for low-precision deployment of pre-trained models.

Abstract

Pre-trained language models like Ernie or Bert are currently used in many applications. These models come with a set of pre-trained weights typically obtained in unsupervised/self-supervised modality on a huge amount of data. After that, they are fine-tuned on a specific task. Applications then use these models for inference, and often some additional constraints apply, like low power-budget or low latency between input and output. The main avenue to meet these additional requirements for the inference settings, is to use low precision computation (e.g. INT8 rather than FP32), but this comes with a cost of deteriorating the functional performance (e.g. accuracy) of the model. Some approaches have been developed to tackle the problem and go beyond the limitations of the PTO (Post-Training Quantization), more specifically the QAT (Quantization Aware Training, see [4]) is a procedure that interferes with the training process in order to make it affected (or simply disturbed) by the quantization phase during the training itself. Besides QAT, recently Intel-Habana Labs have proposed an additional and more direct way to make the training results more robust to subsequent quantization which uses a regularizer, therefore changing the loss function that drives the training procedure. But their proposal does not work out-of-the-box for pre-trained models like Ernie, for example. In this short paper we show why this is not happening (for the Ernie case) and we propose a very basic way to deal with it, sharing as well some initial results (increase in final INT8 accuracy) that might be of interest to practitioners willing to use Ernie in their applications, in low precision regime.