PTQ4ViT: Post-training Quantization for Vision Transformers with Twin Uniform Quantization

Abstract

Quantization is one of the most effective methods to compress neural networks, which has achieved great success on convolutional neural networks (CNNs). Recently, vision transformers have demonstrated great potential in computer vision. However, previous post-training quantization methods performed not well on vision transformer, resulting in more than 1% accuracy drop even in 8-bit quantization. Therefore, we analyze the problems of quantization on vision transformers. We observe the distributions of activation values after softmax and GELU functions are quite different from the Gaussian distribution. We also observe that common quantization metrics, such as MSE and cosine distance, are inaccurate to determine the optimal scaling factor. In this paper, we propose the twin uniform quantization method to reduce the quantization error on these activation values. And we propose to use a Hessian guided metric to evaluate different scaling factors, which improves the accuracy of calibration at a small cost. To enable the fast quantization of vision transformers, we develop an efficient framework, PTQ4ViT. Experiments show the quantized vision transformers achieve near-lossless prediction accuracy (less than 0.5% drop at 8-bit quantization) on the ImageNet classification task.

Z. Yuan and C. Xue—contribute equally to this paper.

Appendix

1.1 Derivation of Hessian guided metric

Hessian guided metric introduces as small an increment on task loss $L = CE(\hat{y}, y)$ as possible, in which $\hat{y}$ is the prediction of the quantized model and y is the ground truth. Here y is approximated by the prediction of the floating-point model $y_{FP}$, since no labels of input data are available in PTQ.

Quantization introduces a small perturbation $\epsilon$ on weight W, whose effect on task loss $\mathbb {E}[L(W)]$ could be analyzed with Taylor series expansion,

$$\begin{aligned} \mathbb {E}[L(\hat{W})]-\mathbb {E}[L(W)]\approx \epsilon ^T \bar{g}^{(W)}+\frac{1}{2}\epsilon ^T \bar{H}^{(W)}\epsilon . \end{aligned}$$

(8)

Since the pretrained model has converged to a local optimum, The gradients $\bar{g}^{(W)}$ is close to zero and could be ignored. The Hessian matrix $\bar{H}^{(W)}$ on weight could be computed by

$$\begin{aligned} \dfrac{\partial ^2\,L}{\partial w_i \partial w_j} = \dfrac{\partial }{\partial w_j} (\sum _{k=1}^{m} \dfrac{\partial L}{\partial O_k} \dfrac{\partial O_k}{\partial w_i}) = \sum _{k=1}^{m} \dfrac{\partial L}{\partial O_k} \dfrac{\partial ^2 O_k}{\partial w_i \partial w_j} + \sum _{k,l=1}^{m} \dfrac{\partial O_k}{\partial w_i} \dfrac{\partial ^2\,L}{\partial O_k \partial O_l} \dfrac{\partial O_l}{\partial w_j}. \end{aligned}$$

(9)

$O = W^TX \in R ^m$ is the output of the layer, and $\dfrac{\partial ^2 O_k}{\partial w_i \partial w_j} = 0$. So the first term of Eq. (9) is zero, and $\bar{H}^{(W)} = J_O(W)^T \bar{H}^{(O)} J_O(W)$. Therefore, Eq. 8 could be further written as,

$$\begin{aligned} \mathbb {E}[L(\hat{W})]-\mathbb {E}[L(W)]\approx \dfrac{1}{2}(J_{O}(W) \epsilon )^T \bar{H}^{(O)} J_{O}(W) \epsilon \approx \dfrac{1}{2}(\hat{O}-O)^T\bar{H}^{(O)} (\hat{O}-O) \end{aligned}$$

(10)

Following Liu et al. [14], we use the Diagonal Fisher Information Matrix to substitute $\bar{H}^{(O)}$. The optimization is formulated as:

$$\begin{aligned} \mathop {min}_{\varDelta _W} \mathbb {E}[ (\hat{O}-O)^T \text {diag}((\dfrac{\partial L}{\partial O_1})^2, \cdots , (\dfrac{\partial L}{\partial O_m})^2) (\hat{O}-O)]. \end{aligned}$$

(11)