BRECQ: Pushing the Limit of Post-Training Quantization by Block Reconstruction

Introduction

The past decade has witnessed the rapid development of deep learning in many tasks, such as computer vision, autonomous driving, etc. However, the issue of huge computation cost and memory footprint requirements in deep learning has received considerable attention. Some works such as neural architecture search (zoph2016nas) try to design and search a tiny network, while others, like quantization (hubara2017qnn), and network pruning (han2015deepcompress) are designed to compress and accelerate off-the-shelf well-trained redundant networks.

Many popular quantization and network pruning methods follow a simple pipeline: training the original model and then finetune the quantized/pruned model. However, this pipeline requires a full training dataset and many computation resources to perform end-to-end backpropagation, which will greatly delay the production cycle of compressed models. Besides, not all training data are always ready-to-use considering the privacy problem. Therefore, there is more demand in industry for quantizing the neural networks without retraining, which is called Post-training Quantization. Although PTQ is fast and light, it suffers from severe accuracy degeneration when the quantization precision is low. For example, DFQ (nagel2019dfq) can quantize ResNet-18 to 8-bit without accuracy loss (69.7% top-1 accuracy) but in 4-bit quantization, it can only achieve 39% top-1 accuracy. The primary reason is the approximation in the parameter space is not equivalent to the approximation in model space thus we cannot assure the optimal minimization on the final task loss.

Recent works like (nagel2020adaround) recognized the problem and analyzed the loss degradation by Taylor series expansion. Analysis of the second-order error term indicates we can reconstruct each layer output to approximate the task loss degeneration. However, their work cannot further quantize the weights into INT2 because the cross-layer dependency in the Hessian matrix cannot be ignored when the perturbation on weight is not small enough. In this work, we analyze the second-order error based on the Gauss-Newton matrix. We show that the second-order error can be transformed into network final outputs but suffer from bad generalization. To achieve the best tradeoff, we adopt an intermediate choice, block reconstruction. In addition, our contributions are threefold:

Based on the second-order analysis, we define a set of reconstruction units and show that block reconstruction is the best choice with the support from theoretical and empirical evidence. We also use Fisher Information Matrix to assign each pre-activation with an importance measure during reconstruction.

We incorporate genetic algorithm and the well-defined intra-block sensitivity measure to generate latency and size guaranteed mixed precision quantized neural networks, which fulfills a general improvement on both specialized hardware (FPGA) and general hardware (ARM CPU).

We conduct extensive experiments to verify our proposed methods. We find that our method is applicable to a large variety of tasks and models. Moreover, we show that post-training quantization can quantize weights to INT2 without significant accuracy loss for the first time.

Preliminaries

Quantization Background Uniform symmetric quantization maps the floating-point numbers to several fixed-points. These points (or grids) have the same interval and are symmetrically distributed. We denote the set that contains these grids as $\mathcal{Q}_{b}^{\text{u,sym}}=s\times\{-2^{b-1},\dots,0,\dots,2^{b-1}-1\}$. Here, $s$ is the step size between two grids and $b$ is the bit-width. Quantization function, denoted by $q(\cdot):\mathcal{R}\rightarrow\mathcal{Q}_{b}^{\text{u,sym}}$, is generally designed to minimize the quantization error:

Solving this minimization problem, one can easily get the $q(\cdot)$ by leveraging the rounding-to-nearest operation $\lfloor{\cdot}\rceil$. Rounding-to-nearest is a prevalent method to perform quantization, e.g. PACT (choi2018pact). However, recently some empirical or theoretical evidence supports that simply minimizing the quantization error in parameter space does not bring optimal task performances. Specifically, Esser2020LEARNED propose to learn the step size $s$ by gradient descent in quantization-aware training (QAT). LAPQ (nahshan2019lapq) finds the optimal step size when the loss function is minimized without re-training the weights. Their motivations are all towards minimizing a final objective, which is the task loss, i.e.,

Taylor Expansion It turns out that the quantization imposed on weights could be viewed as a special case of weight perturbation. To quantitatively analyze the loss degradation caused by quantization, nagel2020adaround use Taylor series expansions and approximates the loss degradation by

At last, the objective is transformed into a practical proxy signal, the change in feature-maps ($\mathbf{z}=\mathbf{W}\mathbf{x}$), and the quantized model can be obtained by a layer-by-layer feature map reconstruction algorithm (with few calibration images). Recent works, like Bit-Split (wang2020bitsplit) and AdaQuant (hubara2020adaquant), also take this layer-wise objective to improve the post-training quantization. However, they failed to quantize weights to INT2. We think the inherent reason is that when $\Delta\mathbf{w}$ grows higher, the former assumptions do not hold and an accurate signal is required.

Proposed Method

where $\mathbf{z}^{(n)}\in\mathcal{R}^{m}$. Since the pretrained full precision model is converged to a local minimum, we can assume the Hessian is positive-semidefinite (PSD). Specifically, the converged model has $\nabla_{\mathbf{z}^{(n)}}L$ close to $\mathbf{0}$ so the first term in Eq. (missing) 5 is neglected and Hessian becomes the Gauss-Newton (GN) matrix $\mathbf{G}^{(\theta)}$. GN matrix can be written in matrix form (botev2017) as

where $\mathbf{J}_{\mathbf{z}^{(n)}}({\theta})$ is the Jacobian matrix of the network output with respect to the network parameters. However, in practice, we cannot explicitly compute and store the Jacobian for each input data point in such a raw form. To reduce the computation and memory budget, we will transform the second-order error into the network output, as shown in the following theorem.

Consider an $n$-layer feedforward neural network with ReLU activation function. Assuming all weights are quantized, the second-order error optimization can be transformed by:

We prove the theorem using the quadratic form, details can be found in Appendix LABEL:append_proof. Here we provide a sketch of the proof by matrix form. The product of perturbation and Jacobian can be thought as the first-order Taylor approximation of the change in network output $\Delta\mathbf{z}^{(n)}$:

Therefore, combining Eq. (missing) 8 and Eq. (missing) 6 we can transform the large-scale second-order error into the change in network outputs characterized by the output Hessian $\mathbf{H}^{(\mathbf{z}^{(n)})}$. The theorem indicates a simple observation, suppose a well-trained teacher model and an initialized student model, we can minimize their discrepancy by reconstructing the network’s final output $\mathbf{z}^{(n)}$, which coincides with and generalizes the distillation (hinton2015distilling; polino2018distillation). LAPQ (nahshan2019lapq) also considers the dependency but their optimization does not rely on second-order information. However, we should emphasize that distillation requires the same computation and data resources as in normal training procedure, which is impractical for PTQ with limited data.

2 Block Reconstruction

Although the network output reconstruction has an accurate estimation of the second-order error, we find in practice it is worse than the layer-by-layer reconstruction in PTQ. The primary reason for this is optimizing the whole networks over 1024 calibration data samples leads to over-fitting easily. As jakubovitz2019generalizationerror explained, the networks can have perfect expressivity when the number of parameters exceeds the number of data samples during training, but lower training error does not ensure lower test error. We find layer-wise reconstruction acts like a regularizer which reduces the generalization error by matching each layer’s output distribution. In other words, both layer-wise and network-wise output reconstruction has their own drawbacks. And there should be a better bias-variance trade-off choice to conduct reconstruction at an intermediate granularity.

To this end, we define 2 extra kinds of intermediate reconstruction granularity: Stage-wise reconstruction and Block-wise reconstruction. These 4 reconstruction granularities are described below:

Layer-wise Reconstruction: Assume the Hessian matrix is layer-diagonal and optimize the layer output one-by-one. It does not consider cross-layer dependency and resemble existing methods (nagel2020adaround; hubara2020adaquant; wang2020bitsplit).

Block-wise Reconstruction: A block is the core component in modern CNN, such as the Residual Bottleneck Block as shown in 1(a). This method assumes the Hessian matrix is block-diagonal and block-by-block perform reconstruction, which ignores inter-block dependencies.

Stage-wise Reconstruction: A stage is where the featuremaps will be downsampled and generate more channels, which is believed to produce higher-level features. Typical CNN in ImageNet dataset contains 4 or 5 different stages. This method simultaneously optimizes all layers within a stage and thus considers more dependencies than the block-wise method.

Network-wise Reconstruction: Optimize the whole quantized network by reconstructing the output of the final layers. This method resembles distillation but does not result in good performances with few images because of high generalization error.

The relationship between network, stage, block, and layer is illustrated in 1(a). We test these 4 kinds of reconstruction granularity and find that block-wise optimization outperforms others. We think this is because the main off-diagonal loss in the Hessian is concentrated in each block, as 1(b) orange part illustrated, while the inter-block loss is small and can be ignored in the optimization. The shortcut connections, which is proposed in (he2016resnet), may also increase the dependencies within a block. Also, the stage-wise or net-wise optimization suffer from the bad generalization on the validation set and degenerate the final performances. We report the quantitative comparison in Sec. 4.1. We name our algorithm Brecq, because we choose block as our base reconstruction unit. It is necessary to point out that our analysis does not give the optimal configuration of the reconstruction granularity. The choice of block-wise optimization comes from our experiments and we find this choice has two merits. (1) No hyper-parameters included and (2) applicable for all models and all tasks we tested.

3 Approximating Pre-activation Hessian

We use the diagonal Fisher Information Matrix (FIM) to replace the pre-activation Hessian. Formally, given a probabilistic model $p(x|\theta)$, the FIM is defined as:

The FIM is equal to the negative expected Hessian of the log-likelihood function, therefore, a simple corollary is that the Hessian of task loss will become FIM if the model distribution matches the true data distribution (lecun2012efficientbackprop). Although matching true data distribution seems impossible, this is the best we can do since the pretrained model is converged.

The diagonal of the pre-activation FIM is equal to the squared gradients of each elements, which is successfully used in Adam (kingma2014adam) for the second momentum. The optimization objective becomes

Compared with the MSE minimization, the above minimization incorporates the squared gradient information. If the output has higher absolute gradients, it will receive more attention when being reconstructed. A similar method for pruning the pre-activation has been proposed in theis2018fisherprune.

Note that Brecq is compatible with any optimization method, like STE (hubara2017qnn). Here we adopt adaptive rounding (nagel2020adaround) for weights and learned step size (Esser2020LEARNED) for activation step size because we observe they generally perform better in PTQ, see details in Appendix LABEL:sec_learning. We formulate the overall calibration algorithm for a unified precision model in algorithm 1. We should emphasize that we only need a small subset (1024 in our experiments) of the whole training dataset to calibrate the quantized model. And we can obtain a quantized ResNet-18 within 20 minutes on a single GTX 1080TI GPU.

4 Mixed Precision

To further push the limit of post-training quantization, we employ mixed precision techniques, which can be formulated by

Here, $\mathbf{c}$ is the bit-width vector with the shape of number of layers. $H(\cdot)$ is a hardware performance measurement function, which is used to ensure the mixed precision model has the same or lower hardware performance (e.g., memory and speed) than a predefined threshold $\delta$. We choose 2, 4, 8-bit for mixed precision because they are most common in practical deployment.

Regarding the training loss $L$, we find that nearly all existing literature (cai2020zeroq; hubara2020adaquant; dong2019hawq) uses layer-wise measurement. They all assume the sensitivity within a layer is independent and can be summed together. Therefore, the mixed precision problem becomes an integer programming problem. However, we argue that the loss measurement should contain two parts: diagonal loss and off-diagonal loss, the first is the same with previous works and measure the sensitivity of each layer independently, while the off-diagonal loss is used to measure the cross-layer sensitivity. Theoretically, we should examine all permutations, which results in $3^{n}$ possibilities and prohibits the search algorithm. Our first attempt is to reduce the off-diagonal loss into the block-level as we mentioned that the Hessian can be approximated to a block-diagonal matrix. Granted, we still find the search space is large, for example, if a block has four layers, then we have to consider the $3^{4}=81$ permutations for a single block. Based on our preliminary experiments, we find that 4-bit and 8-bit quantization nearly do not drop the final accuracy. Hence we only take 2-bit permutations into consideration and drastically reduce the search space. We use genetic algorithm (guo2020single) to search the optimal bitwidth configuration with hardware performance threshold, the algorithm is located in LABEL:alg:2. Due to space limits, we put related works in Appendix LABEL:sec_related_work. Readers can refer to related works for a brief discussion on quantization and second-order analysis.

Experiments

In this section, we report experimental results for the ImageNet classification task and MS COCO object detection task. The detailed implementation of the experiments can be found in the Appendix LABEL:append_implement. The rest of this section will contain ablation study on reconstruction granularity, classification and detection results, mixed precision results and comparison with quantization-aware training. In Appendix LABEL:sec_append_exp, we conduct more experiments, including the impact of the first and the last layer, the impact of calibration dataset size and data source.

We test four kinds of reconstruction granularity: Net-wise, Stage-wise, Block-wise, and Layer-wise

reconstruction. We conduct ImageNet experiments using MobileNetV2 and ResNet-18 with 2-bit weight quantization for all layers except for the first and the last layer. It can be seen from Table 1 that Block-wise optimization outperforms other methods. This result implies that the generalization error in net-wise and stage-wise optimization outweighs their off-diagonal loss. In ResNet-18, we find the difference is not significant, this can be potentially attributed to that ResNet-18 only has 19 layers in the body and the block size, as well as the stage size, is small, therefore leading to indistinct results.

2 ImageNet

We conduct experiments on a variety of modern deep learning architectures, including ResNet (he2016resnet) with normal convolution, MobileNetV2 (sandler2018mobilenetv2) with depthwise separable convolution and RegNet (radosavovic2020regnet) with group convolution. Last but not least important, we also investigate the neural architecture searched (NAS) models, MNasNet (tan2019mnasnet). In Table 2, we only quantize weights into low-bit integers and keep activations full precision. We compare with strong baselines including Bias Correction, optimal MSE, AdaRound, AdaQuant, and Bit-split. Note that the first and the last layer are kept with 8-bit. While most of the existing methods have good performances in 4-bit quantization, they cannot successfully quantize the model into 2-bit. Our method consistently achieves the lowest accuracy degradation for ResNets (within 5%) and other compact models. We further quantize activations into 4-bit to make the quantized model run on integer-arithmetic hardware platforms. We find that 4-bit activation quantization can have a huge impact on RegNet and MobileNet. Nonetheless, our methods produce higher performance than other state-of-the-arts. To be noted, Brecq is the first to promote the 2W4A accuracy of PTQ to a usable level while all other existing methods crashed.

3 Comparison with Quantization-aware Training