Network Quantization with Element-wise Gradient Scaling

We design a uniform quantizer $Q$ that converts a full-precision input $x$ to a quantized output $Q({x})$, where we denote by $x$ a scalar element of either a weight or an input activation tensor $\mathbf{x}$ in a layer. We learn a quantization interval [21, 14] using lower and upper bounds, denoted by $l$ and $u$, respectively. Specifically, the quantizer first generates a full-precision latent value $x_{n}$ by normalizing and clipping the input value $x$ as follows:

where $\textrm{clip}(\cdot,0,1)$ is a clipping function with lower and upper bounds of $0$ and $1$, respectively. Note that weight and/or activation quantizers in every quantized layer use separate parameters for the quantization intervals (\ie, $l$ and $u$). For $b$-bit quantization, the latent value $x_{n}$ is converted to a discrete value $x_{q}$ using a round function with pre-/post-scaling as follows:

Finally, the quantizer outputs a quantized weight $Q_{W}(x)$ or activation $Q_{A}(x)$ as follows:

where we restrict the quantized activation $Q_{A}(x)$ to be non-negative [5] considering the pre-activation by a ReLU. To adjust an output scale of the layer, we train an additional parameter $\alpha$ for each quantized layer, which is multiplied by the output activations of convolutional or fully-connected layers.

The main difficulty of training a quantized network arises from the round function in Eq. (2), since its derivative is zero at almost everywhere. Most quantization methods avoid zero gradients using STE [3]. It approximates the derivative of the round function by an identity function, that is, $\mathcal{G}_{\mathbf{x}_{n}}=\mathcal{G}_{\mathbf{x}_{q}}$, where we denote by $\mathbf{x}_{n}$ and $\mathbf{x}_{q}$ tensors containing latent and discrete values, respectively, and by $\mathcal{G}_{\mathbf{x}_{n}}$ and $\mathcal{G}_{\mathbf{x}_{q}}$ corresponding gradients. Propagating the same gradient from discrete to latent values is, however, sub-optimal for the following reasons: (1) Multiple latent values can produce the same discrete value; (2) The same gradient provided by the discrete value affects differently to each of the latent values. To overcome this problem, we introduce an EWGS method, an effective alternative to STE, defined as follows:

where $g_{x_{n}}$ and $g_{x_{q}}$ are the elements of the gradients $\mathcal{G}_{\mathbf{x}_{n}}$ and $\mathcal{G}_{\mathbf{x}_{q}}$, corresponding to the partial derivatives of a task loss w.r.t $x_{n}$ and $x_{q}$, respectively. $\mathrm{sign}(\cdot)$ is a signum function and $\delta\geq 0$ is a scaling factor. EWGS adjusts the gradient element of discrete values $g_{x_{q}}$ adaptively using the sign of the element, $\mathrm{sign}(g_{x_{q}})$, and a discretization error, $x_{n}-x_{q}$. Note that STE is a special case of EWGS, that is, Eq. (4) corresponds to STE, when a scaling factor $\delta$ is zero.

We visualize in Fig. 2 1-D examples illustrating an effect of EWGS. We can see that EWGS encourages a gradient element for the discrete value $g_{x_{q}}$ to decrease with the scale of $(1+\delta\mathrm{sign}(g_{x_{q}})(x_{n}-x_{q}))$ to update the latent value $x_{n}$, when the latent value $x_{n}$ to update is already located farther than the discrete value $x_{q}$ in the direction of change (\ie, $-\mathrm{sign}(g_{x_{q}})$) as shown in Fig. 2(a) (right) and Fig. 2(b) (middle), and to increase in the opposite case as shown in Fig. 2(a) (middle) and Fig. 2(b) (right). Note that we use non-negative values for the scaling factor $\delta$, since negative ones lead to opposite effects. To sum up, EWGS resolves discrepancies between latent and discrete values during backpropagation by considering discretization errors between these values and their direction of change. As will be shown in Sec. 4.3, this not only stabilizes the training of a quantized network but also encourages better convergence compared to STE.

It is crucial to determine a scaling factor $\delta$, since an improper value would hinder the training process, and weight/activation quantizers in different layers may require different degrees of scaling. Let us consider the following equation:

where $\epsilon=x_{n}-x_{q}$ is a discretization error from the round function in Eq. (2). Since an absolute value of the error is bounded by a small number \ie, $|\epsilon|\leq\frac{0.5}{2^{b}-1}$, we assume that the error is small enough to approximate Eq. (5) as follows:

where $g_{x_{q}}^{\prime}=\frac{\partial g_{x_{q}}}{\partial x_{q}}$ is a second-order derivative of a task loss w.r.t the discrete value $x_{q}$. This can be represented as follows:

which corresponds to EWGS in Eq. (4). This suggests that we can set the scaling factor $\delta$ as $\frac{g_{x_{q}}^{\prime}}{|g_{x_{q}}|}$, but calculating an exact Hessian matrix $H$ to obtain the second-order derivative $g_{x_{q}}^{\prime}$ is computationally demanding. We instead approximate the second-order derivative by an average of diagonal elements in the Hessian matrix $H$, with an assumption that the main diagonal dominates the matrix $H$, and the discrete values $\mathbf{x}_{q}$ obtained from the same weight or activation quantizer in a layer influence similarly to the loss function [11, 10]. To this end, we compute a Hessian trace with an efficient algorithm [10, 40] using the Hutchinson’s method [1]:

where $I$ is an identity matrix, $\mathbb{E}$ is an expectation operator, and $\mathbf{v}$ is a random vector drawn from the Rademacher distribution, satisfying $\mathbb{E}[\mathbf{vv}^{T}]=I$. This implies that we can estimate the trace of a Hessian matrix $\mathrm{Tr}(H)$ with $\mathbb{E}[\mathbf{v}^{T}H\mathbf{v}]$, where we can obtain $H\mathbf{v}$ efficiently without forming an exact Hessian matrix as follows:

We then define the scale parameter $\delta$ as follows:

where $N$ is the number of diagonal elements in the Hessian matrix and $G$ is a gradient representative determined from a distribution of the gradients $\mathcal{G}_{\mathbf{x}_{q}}$. Based on Eq. (7), we could take an average over the absolute values of gradient elements \ie, $\mathbb{E}[|g_{x_{q}}|]$ for setting $G$, but we empirically found that most gradients are concentrated near zero, such that the average value tends to be biased to small gradient elements. We instead set $G$ to a sufficiently large value. A plausible reason is that considering large gradient elements is more important, since they dominate the training. Specifically, we use $3\sigma(\mathcal{G}_{\mathbf{x}_{q}})$ as the gradient representative $G$, where $\sigma(\cdot)$ computes a standard deviation. It enables finding a sufficiently large gradient element (\eg, $3\sigma(\cdot)$ accounts for roughly 99 percent of data in case of the Gaussian distribution). We could take the maximum over the absolute values of gradient elements, but it often corresponds to an outlier of a distribution.

Assuming that the loss function is locally convex, we take a non-negative value for the scaling factor (\ie, $\textrm{max}(0,\delta)$), which coincides with the condition in Eq. (4). We use individual scaling factors for all weight and activation quantizers in a network, and update them periodically during training for efficiency. We summarize in Algorithm 1 an overall quantization procedure of our approach.

We compare in Table 1 the validation accuracy on ImageNet [8] using the ResNet-18 [16] architecture under various bit-width settings. All numbers for other methods, except for LSQ [12], are taken from corresponding papers including the performance of full-precision models. Note that LSQ reports the results with a pre-activation structure [17] of ResNet, which is different from ours. We thus take the results from the work of [4] in which LSQ is reproduced using the same network structure as ours. We summarize the findings from Table 1 as follows: (1) Our quantization method with EWGS achieves the state of the art. For 4-bit weights and 4-bit activations, our method shows the classification accuracy slightly lower than LSQ+ [4], but the performance gain w.r.t the full-precision model is on a par with LSQ+. In particular, our model achieves the performance comparable to the full-precision one with only 3-bit representations. (2) Our method performs better than QuantNet [39] and DSQ [14], which attempt to address the problem of STE using soft quantizers, indicating that EWGS is a better alternative to STE than soft quantizers. (3) Our method exploiting EWGS brings significant performance improvement in a binary setting, surpassing other methods including the ones specially designed for binary quantization [34, 33]. This suggests that EWGS works favorably even with large discretization errors. (4) We verify the effectiveness of our method over a wide range of quantization bit-widths, outperforming the state of the art consistently, whereas other methods report the results selectively in few settings.

We show in Tables 2 and 3 quantization results for ResNet-34 [16] and MobileNet-V2 [35], respectively, on ImageNet [8]. As mentioned earlier, LSQ [12] uses a different network structure¹¹1The full-precision baseline of ResNet-34 used in LSQ shows the top-1 validation accuracy of 74.1, which is higher than that of our full-precision baseline (73.3)., but we include its performance in Table 2 in order to compare relative performance drops or gains for network quantization. We can observe from Table 2 similar findings in Table 1. Our method outperforms the state of the art over all bit-width settings. With 3-bit weights and 3-bit activations, the quantized network trained with our method does not degrade the performance at all, compared with the full-precision model. Ours also gives better results than LSQ in terms of performance drops or gains after quantization. We can see from Table 3 that our model performs better than PACT [7, 37] and DSQ [14], but it is slightly outperformed by PROFIT [31] for 4-bit quantization of both weights and activations. Note that PROFIT exploits many training heuristics, such as knowledge distillation [18], progressive quantization [47], an exponential moving average of weights, batch normalization post-training, and iterative training with incremental weight freezing. We achieve a comparable result using a simple gradient scaling without bells and whistles, which confirms that EWGS is simple yet effective method for network quantization. Moreover, PROFIT is effective to quantize light-weight networks only, while ours can be applied to various network architectures under a wide range of bit-widths.