NIPQ: Noise proxy-based Integrated Pseudo-Quantization

Neural network quantization is a representative optimization technique that reduces the memory footprint by storing the activation and weight in a low-precision domain. In addition, when hardware acceleration is available (e.g., low-precision arithmetics [57, 48, 37, 51] or bit-serial operations [50, 15, 19]), it also brings a substantial performance boost. These advantages make network inference affordable in large-scale servers as well as embedded devices [39, 2], which has helped popularize it in various applications. However, quantization has a critical disadvantage, quality degradation due to limited degrees of freedom. The way to train the networks accurately within limited precision is critical and receiving much attention these days.

To mitigate the accuracy degradation, quantization-aware training (QAT) has emerged that trains a neural network with quantization operators to adapt to low precision. While the quantization operator is not differentiable, the straight-through estimator (STE) [5] allows the backpropagation of the quantized data based on linear approximation [60, 20]. This approximation works well in redundant networks with moderate precision ($>$4-bit). Thus, not only early studies [9, 41, 60] but also advanced ones [8, 17, 54] have proposed diverse QAT schemes based on STE and shown that popular neural networks (i.e., ResNet-18 [21]) can be quantized into 4-bit without accuracy loss.

However, STE-based QAT bypasses the approximated gradient, not the true gradient, and many studies have pointed out that it can incur instability and bias during training [34, 20, 46, 38, 42, 36]. For instance, PROFIT[38] points out that the instability is a major source of accuracy drop for the optimized networks (e.g., MobileNet-v2/v3 [43, 24]). More recently, an alternative training scheme, pseudo-quantization training (PQT) based on pseudo-quantization noise (PQN), has been proposed [3, 10, 44] to address the instability of STE-based QAT. During PQT, the behavior of the quantization operator is simulated via PQN, and the learnable parameters are updated based on the proxy of quantization. While those studies are applied only to the weight, they can stabilize the training process significantly compared to QAT with STE and show the potential of PQT.

Nevertheless, the existing PQT algorithms have room for improvement. Various STE-based studies [60, 8, 17] have shown that truncation contributes significantly to reducing quantization errors, but even the advanced PQT studies [10, 44] use a naive min-max quantization. Integrating the truncation on the PQT framework can greatly reduce quantization error and enable us to exploit a PQT scheme for activation, which has an input-dependent distribution and requires a static quantization range. In addition, there is no theoretical support for whether PQT guarantees the optimal convergence of the quantization parameters. Intuitive interpretation exists, but proof of whether quantization parameters are optimized after PQT has yet to be provided.

In this paper, we propose an novel PQT method (Figure 1) called Noise proxy-based Integrated Pseudo-Quantization (NIPQ) that quantizes all activation and weight based on PQN. We present a novel idea, called Noise proxy, that shares the same quantization hyper-parameters (e.g., truncation boundary and bit-width) with the existing STE-based algorithm LSQ(+) [17, 6]. However, noise proxy allows to update the quantization parameters, as well as the network parameters, via gradient descent with PQN instead of STE. Subsequently, an arbitrary network can be optimized in a mixed-precision representation without STE instability. Our key contributions are summarized as follows:

• •

  NIPQ is the first PQT that integrates truncation in addition to discretization. This extension not only further reduces the weight quantization error but also enables PQT for activation quantization.

• •

  NIPQ optimizes an arbitrary network into the mixed-precision with awareness of the given resource constraint without human intervention.

• •

  We provide theoretical analysis showing that NIPQ updates the quantization hyperparameters toward minimizing the quantization error.

• •

  We provide extensive experimental results to validate the utility of NIPQ. It outperforms all existing mixed-precision quantization schemes by a large margin.

Due to its practical usefulness, various studies have been proposed for multi-bit linear quantization [60, 8, 26, 17]. Most of these studies are based on STE, and the proposed schemes have evolved by designing the quantization function to enable the optimization of the quantization parameters via gradient descent. However, in optimized networks such as MobileNet-v2, accuracy loss induced by STE instability has been reported for both activation [20] and weight [38]. They addressed the instability via a newly designed pipeline and non-linear approximation but suffered from the increased complexity and cost of QAT. In this work, NIPQ updates the quantization parameters without STE approximation, enabling stable convergence without additional cost or complexity and, most importantly, outperforming the existing methods by a large margin.

Mixed-precision studies focus on allocating layer-wise or group-wise bit-width in consideration of the precision sensitivity of each layer to minimize accuracy drop within the given resource constraints. Various methods have been proposed, e.g., RL-based [53, 16, 33], Hessian-based [14, 13, 58], and differentiable [52, 54] algorithms, but RL-based and Hessian-based methods are relatively complex, requiring a lot of parameter tuning, and differentiable algorithms still suffer from STE approximation error. In the present work, we reinterpret the differentiable bit-width tuning in terms of PQT framework, enabling the layer-wise bit-width tuning via gradient descent without STE instability.

Robust quantization [47, 1, 7] aims to guide the convergence of the network toward a smooth and flat loss surface based on additional regularization. The robustness of neural networks is highly beneficial for deploying noisy devices or low-precision ALUs. In the case of noise proxy, it inherently improves the robustness during PQT with PQN. As far as we know, we observe for the first time that the robustness of activation quantization can be enhanced (Section 7.2).

The benefit of the PQT has been demonstrated in diverse studies [10, 44] in various perspectives. However, previous studies utilize PQT with PQN in a limited domain only for weight and have yet to provide theoretical support for the convergence of PQT corresponding to the quantization. We extend the idea of PQT to both activation and weight by integrating the idea of truncation in addition to discretization and prove that the optimization of quantization parameters through noise proxy minimizes the actual quantization error. In addition, our work is the first to demonstrate that the PQT-based pipeline outperforms STE-based mixed-precision quantization algorithms by a large margin.

Before introducing the details, we first propose a simple problem for straightforward explanation. The objective is to minimize the difference from the target data $t\in\mathbb{R}^{N}$ and the quantized value of the learnable parameters $x\in\mathbb{R}^{N}$. The example loss function is defined as follows:

where subscript $i$ means the i-th element of the data, $t$ is arbitrary distribution, and $Q(\cdot)$ is a given quantization function, parameterized $\alpha$ (truncation boundary) and $b$ (bit-width). To minimize the loss, we need to update $x,\alpha$, and $b$ judiciously. We utilize this example for the rest of the paper.

In this work, we focus on linear quantization, especially for a more specific case that has an edge on hardware acceleration (e.g., affine layer-wise quantization for activation and symmetry channel-wise quantization for weight) [51]. To get optimal quality within restricted resources, we need to tune the quantization parameters, shown in Figure 3. In this section, we analyze the limitations of existing STE-based and PQT-based algorithms and provide insight for next-step innovation.

In this work, we propose a novel PQT pipeline to compensate for the quality degradation after quantization without STE instability. Moreover, the proposed scheme includes the support of updating truncation, resulting in accurate quantization not only for weight but also for activation. Besides, we adopt the gradient-based bit-width optimization via continuous relaxation [54, 10]. This extension allows us to update all of the hyper-parameters for linear quantization via gradient descent, which greatly simplifies the layer-wise optimization process.

In this section, we define the proposed idea and present the analytical support for the optimal convergence of quantization parameters based on PQT. Our method is composed of a pair of algorithms: STE-based quantization $Q(x|\alpha,b)$ [17, 6] and the corresponding noise proxy $\tilde{Q}(x|\alpha,b)$. The algorithms are expressed as follows:

where $\Delta=\frac{\alpha}{2^{b}-1}$ represents the quantization step size or the difference between quantization levels, $\epsilon$ represents the PQN, and $\lfloor\cdot\rceil$ means the rounding operation¹¹1Please note that we present the quantization algorithm for non-negative distribution data, e.g., activation after ReLU, for simplicity. However, this approach holds the same for weight, except that quantization index is used in $[-2^{bit-1},2^{bit-1}-1]$..

During PQT, $\tilde{Q}(x|\alpha,b)$ is applied on top of the quantization targets, activation and weight, and the learnable parameters are updated via gradient descent. Because $\alpha$ and $b$ are included in the gradient path, the training with noise proxy updates all of the quantization parameters as well as the network’s learnable parameters jointly. If we set the target loss function, $\mathcal{L}_{target}$ as the mixture of the task loss $\mathcal{L}_{task}$, and the cost loss $\mathcal{L}_{cost}$, for memory footprint or computation as follows:

where $\lambda$ is a hyper-parameter of balancing two loss functions. The layer-wise bit-width $b$ is assigned within the available resource budget, and the corresponding $\alpha$ is tuned appropriately to achieve the best quality possible.

After PQT, $Q(x|\alpha,b)$ is used for the true quantization during inference. To achieve the high-quality output with the true quantization operator, the tuned $\alpha$ and $b$ through PQT with noise proxy should also have the lowest loss with $Q(x|\alpha,b)$. In the following subsections, we provide theoretical support that the optimization of $\alpha$ and $b$ through PQT guarantees the optimal convergence $Q(x|\alpha,b)$ when we sample PQN judiciously.

As explained above, NIPQ enables the compensation of quality degradation of the quantized network without STE approximation. However, in practice, we need to consider several details to maximize obtainable accuracy.

In order to validate the performance of NIPQ, we conduct comprehensive studies to measure the output quality and observe diverse properties in various applications. The results in this paper are obtained based on PQN sampled from the quantization noise, unless otherwise specified explicitly. Likewise, QAT finetune is used at the last few epochs, otherwise specified explicitly. We implement our algorithm on PyTorch [40] library, and the code is available at our repository ³³3https://github.com/ECoLab-POSTECH/NIPQ. The details of target loss and training pipeline are available in the Supplementary Materials.

To achieve reliable output in low-precision, we propose NIPQ, which enhances the benefit of PQT pipeline by integrating the idea of truncation and providing theoretical support for it. NIPQ automates the layer-wise low-precision optimization for an arbitrary network by updating the network parameters and quantization parameters jointly via gradient descent toward minimizing the target loss. Our extensive results show that the proposed scheme outperforms all of the existing studies by a large margin for various vision and language applications.

In this supplementary material, we present the details of our implementation and additional experimental results for various tasks and different datasets. We provide the following items:

• •

  The detailed implementation of the cost loss function in Section S-II.

• •

  Detailed Configurations of our experiments in Section S-III

• •

  The results of quantization for super resolution task in Section S-IV.

• •

  Additional experimental results of object detection on MS-COCO dataset in Section S-V.

• •

  An ablation study on the effect of stochastic rounding in Section S-VI.

• •

  An ablation study on the effect of late training stage in Section S-VII.

• •

  Experimental results on quantization parameter robustness in Section S-VIII.

• •

  The visualization of the quantization noise distribution in Section S-IX

In order to restrict the utilization of memory and computation resources, we introduce an additional cost loss function in addition to the target loss, as explained in Equation (4) of the main paper. The cost functions for the memory consumption $\mathcal{L}_{cost-MP}$ and the computation cost $\mathcal{L}_{cost-BOP}$ are defined as follows:

where $h(\cdot)$ denotes Huber loss, $b^{w}_{i}$/$b^{a}_{i}$ denote the bit-width of i-th layer’s weight/activation, $e^{w}_{i}$/$e^{a}_{i}$ are the number of elements in the i-th layer’s weight/activation, $OPS_{i}$ is FLOPS of the i-th layer and $b_{t}$ denotes the target bit-width. $\mathcal{L}_{cost-MP}$ regularizes the average bit-width of activation/weight to the target bit, and $\mathcal{L}_{cost-BOP}$ regularizes the sum of overall bit-operation (BOPs) to the target BOPs. Note that we utilize the bit-operations (BOPs) as a representative metric to measure the computation cost of a neural network, which is commonly used in many previous studies [14, 13, 58]. However, any arbitrary differentiable function can be used as a drop-in replacement for the cost function, and NIPQ automatically optimizes the layer-wise bit-width to the sweet spot.

On the other hand, while the per-layer (or per-tensor) bit-width also requires rounding operation during forward operation, NIPQ is not applicable for the bit-width because it relies on the statistics of quantization error, but it is improvable to achieve the statistics for the scalar value. To overcome this limitation, we propose to update the bit-width via stochastic rounding with STE (Section S-VI).

In this paper, all experiments are conducted using GPU servers having 8 x NVIDIA GTX3090 with 24 GB VRAM with 2 x AMD 7313 (16 Core 32 T). The number of GPUs is selected to satisfy the minimum requirement of GPU memory for the target task. All of the experiments are implemented based on the PyTorch [40] framework (v1.12.1) [40]. Our source code is also provided. The additional details of training configuration, e.g., optimizer type, initial learning rate, decay policy, etc., are determined depending on the characteristics of applications and provided in the following paragraphs.

Table S4 shows the configurations of ImageNet training for NIPQ results. In this experiment, we apply quantization to every convolution and linear layer, including the first and last layers. One exception is that the input of the first convolution layer is fixed as 8-bit. We use SGD with momentum optimizer and cosine annealing with warmup scheduling for learning rate adjustment [35]. $\eta_{min}$ is the final LR multiplier of cosine annealing, and $\lambda_{w}$, $\lambda_{a}$, and $\lambda_{b}$ are the hyper-parameter of resource constraints for the bit-width of weight, bit-width of activation, and BOPs, respectively. When knowledge distillation is triggered, we use EfficientNet-B0 [49] as a teacher network. We use the conventional dark-knowledge-based distillation [22].

Tables S4 and S4 show the detailed configurations of super-resolution task and object detection task, respectively. In both experiments, we keep the precision of the first and last layers as full-precision and apply low-precision quantization to the rest of the layers. In the super-resolution task, we use Adam optimizer [27] and cosine annealing scheduling for learning rate adjustment. In the object detection task, we use SGD with momentum optimizer and cosine annealing with warmup scheduling for learning rate adjustment. Like the image classification task, $\eta_{min}$ is the final LR multiplier of cosine annealing, and $\lambda_{w}$, $\lambda_{a}$, and $\lambda_{b}$ represent the hyper-parameters of resource constraints for the bit-width of weight, bit-width of activation, and BOPs, respectively.

Table S4 shows the detailed configurations of BERT-base [12] on the GLUE Task dataset. In this experiment, we modified the code from the huggingface-transformer [56] library. We apply weight quantization to every linear layer except the last classification head. Note that we do not quantize activation or word embedding. We use the AdamW optimizer and CosineLR scheduler for fine-tuning BERT except for the bit parameters because we find that AdamW can induce instability during training when the magnitude of the cost loss is too large. We use SGD with momentum optimizer for the bit parameters. Besides, we also find that $\alpha$ and $b$ parameters are not well trained when a single global learning rate is utilized ($1e-5$). For fast and reliable convergence, we use the learning rate of $1e-2$ for bit parameters. In addition, the gradient of $\alpha$ is multiplied $2^{b}-1$ times over the global learning rate.

Table S5 shows the quantitative analysis of NIPQ on super resolution task, and Figure S1 visualizes the quality of the generated figures. We report PSNR as a quantitative measure, one of the well-known metrics in the area of super resolution. NIPQ outmatches the specialized quantization algorithm for super resolution, DDTB[59], which applies dynamic quantization that adjusts the quantization step size depending on the input data. These experimental results indicate that NIPQ works well in the regression task as well.

We conduct an additional experiment on object detection task with the COCO dataset and report mAP on Table S6. NIPQ obtains the best results compared to existing quantization studies in the same average bit-width.

In addition, in Figure S2, we visualize the qualitative results of NIPQ on the VOC dataset. Bounding box regression and classification results of the quantized network are presented. As shown in the figure, NIPQ works surprisingly well in the 3-bit domain on the challenging object detection problem. YoloV5-S has a complicated structure, and the sensitivity of each layer is highly different. Because NIPQ has the ability to allocate the bit-width aware of the sensitivity automatically and enable stable convergence without STE instability, the quality of the quantized network outperforms all of the previous methods by a large margin.

While we propose an alternative training scheme for quantization instead of using STE, updating the bit-width is a remaining problem that is not addressed in the NIPQ pipeline. The proposed noise proxy is designed to update the learnable parameters by emulating the quantization operator based on PQN. However, the bit-width is assigned as a scalar value per the target tensor, and thereby it is impossible to aggregate the coarse-grained effect of the quantization operator. When we use rounding-based QAT with STE approximation, the bit-width also suffers from the instability of STE, resulting in highly unreliable result, as shown in Figure S3. Due to this limitation, many previous studies rely on the continuous approximation of bit-width during training [54, 10] to avoid the instability problem. However, the representation mismatches to the domain of bit-width, resulting in suboptimal convergence in practice, especially when the target bit-width is in a sub-4-bit domain. In this paper, we propose an alternative idea to utilize the stochastic rounding of bit-width during training. Stochastic rounding is an unbiased estimator, so the learnable bit-width converges to the optimal point as the learning progresses. In addition, the bit-width is evaluated in the discrete domain during training, which mitigates the domain gap between training and inference. As shown in Figure S3, the stochastic rounding consistently draws the pareto-front line with small variance, which enables us to search for the best quantization configurations within the given resource budget.

In Table S7, we show the results of NIPQ with different late training stage policies. As shown in the table, BN update offers a large performance benefit compared to the accuracy of the NIPQ training without the late stage tuning. Because PQN of NIPQ disturbs the statistics of normalization layers, the correction of the statistics is essential to maximize the accuracy in the inference phase. In addition, QAT finetune offers an additional performance improvement by giving an additional chance to adjust the learnable parameters of the entire network without the effect of PQN with a small learning rate, which enables the stabilization of network parameters near the optimal point with the tiniest effect of STE instability.

NIPQ also enhances the robustness of the quantization parameters as well as the network parameter. Figure S4 visualizes the results of measuring the accuracy while changing the quantization step size or the truncation interval. The more robust the network, the more it can endure the change of the quantization configuration. As shown in the figure, NIPQ shows comparable or superior results to the previous best algorithm for robustness, KURE [47]. It is especially worthy that existing studies have focused on improving the robustness of weight only [1, 47], but NIPQ also improves the robustness of activation as well by a large margin. To the best of our knowledge, this is for the first time that activation robustness can be improved, which is a crucial benefit of deploying networks in a noisy environment.

In Figure S5, we visualize the quantization noise distribution of ResNet-18’s 12-th convolution weight in different bit-widths. When applying quantization, not only rounding but also truncation is applied. The previous study argues that the quantization noise distribution follows a uniform distribution regardless of input distribution when the number of bits is sufficiently large [55]. According to our observation, the statement is held in practice when the bit-width is larger than 4-bit. However, in sub-4-bit precision, the distribution of noise seems to follow a bell-shaped curve instead of a uniform distribution. Due to these characteristics, conventionally uniform or gaussian distributions are often used to approximate PQN [54, 10]. However, as presented in this paper, the precise sampling of PQN following the quantization error distribution is essential to guarantee the convergence on the optimal point, while the uniform distribution shows comparable results in practice empirically. In this work, we realize the sampling process of quantization error distribution on GPU with practical performance as follows: first, the probability density function (PDF) of the quantization error distribution is estimated based on the histogram with 256 bins. Then, the distribution is sampled from the estimated PDF of the histogram. As shown in Figure S5, the sampled distribution precisely follows the quantization error distribution.