SQWA: Stochastic Quantized Weight Averaging for Improving
the Generalization Capability of Low-Precision Deep Neural Networks

Typically, the precision of parameters and data is reduced for efficient implementations in real-time signal processing system designs. While audio and video signal processing demands precision exceeding eight bits, many DNNs function well with lower precisions, such as one or two bits. Particularly, the performance of low-precision DNNs can be improved considerably by conducting retraining after quantization. Thus, quantization is a highly promising approach for the efficient implementation of DNNs. The quantization training algorithm, first proposed by [17] and [6], has been combined with various types of quantization methods such as symmetric uniform [8], asymmetric uniform [44], non-uniform [25], and differentiable [5, 39, 16, 15] quantizers. In recent years, a few elaborate techniques have been developed, such as employing knowledge distillation and carefully controlling the learning rate and bit-precision for improved generalization [24, 26, 31]. Weight normalization is adopted to avoid a long tail distribution of the model weights [4].

In this work, we focus on obtaining a good training scheme to optimize QDNNs. This approach is focused on developing well-generalized low-precision DNNs instead of developing elaborate quantization schemes. It is noteworthy that we only used the uniform quantization scheme for simplifying the hardware [34, 2] for inference. Non-uniform quantization can yield improved QDNN performance when the precision is the same; however, it demands additional operations, which can be time-consuming when hardware with conventional arithmetic blocks is involved.

Stochastic gradient descent (SGD) is the most widely used method for DNN training. However, the loss surface for SGD contains many sharp minima [14]; thus, SGD-based training exhibits overfitting frequently. Many regularization techniques can be applied for alleviating this problem, such as L2-loss, dropout, and cyclical learning rate scheduling [35, 33, 32]. The ensemble of models is known to increase the generalization capability. However, this method typically demands increased cost for training and inference. Fast geometric ensemble (FGE) is a recent technique for the ensemble of models [9]. Dropout [33] and dropconnect [36] can be interpreted as building an ensemble of models by weight averaging.

SWA is a recently developed regularization technique; that is based on the weight averaging of models captured during training with cyclical learning rate scheduling [18]. SWA demonstrates excellent performances in many CNN models on various datasets.

SWA can be explained using the loss visualization technique. The visualization method for representing three models in a single loss surface has been suggested in [9] and [18]. In those studies, training algorithms SWA and FGE were presented by discovering that the local minima trained with SGD were interconnected. Furthermore, because the loss surface for training and test were different, the minima found by SGD during training were not necessarily the best for the test data. Instead, the average of the models indicated a significantly improved generalization capability. Figure 2 (a) depicts three models captured during cyclical learning rate scheduling and shows that the average is located near the center in the loss surface [18].

SWA has been applied to low-precision training. Stochastic weight averaging in low-precision (SWALP) employs SWA for a cost-efficient training, where a low-precision (e.g., 8-bit) model is trained with cyclical learning rate scheduling and models are captured during training at the lowest learning rate in the cycle [38]. The captured models are then averaged to obtain the final full-precision model. Thus, SWALP is intended to design high-precision models and is vastly different from our work, which optimizes severely quantized models (e.g., 2-bit) for inference.

The weight vector, $\mathbf{w}$, of a deep neural network can be quantized in $b$-bit using a symmetric uniform quantizer as follows:

$\displaystyle\text{Q}^{b}(\mathbf{w})$

$\displaystyle=\text{sign}(\mathbf{w})\cdot\Delta\cdot\text{min}\Big{\{}\Big{\lfloor}\Big{(}\frac{|\mathbf{w}|}{\Delta}+0.5\Big{)}\Big{\rfloor},\frac{(M-1)}{2}\Big{\}},$

(1)

where $\text{sign}(\cdot)$ is the sign function, $\Delta$ is the quantization step size, $M$ is the number of quantization levels that can be computed with $2^{b}-1$. A trained full-precision model can be directly quantized with Equation (1), but the performance will be significantly degraded when severe quantization such as 1- or 2-bit is employed. To relieve this problem, retraining on quantization domain is adopted in previous studies [17, 6, 29, 43] as follows:

	$$\displaystyle l_{i}=\sum\limits_{j\in A_{i}}w_{ij}^{(q)}y_{j}^{(q)}$$		(2)
	$$\displaystyle y_{i}^{(q)}=\phi_{i}(l_{i})$$		(3)
	$$\displaystyle\delta_{j}=\phi_{j}^{{}^{\prime}}(l_{j})\sum\limits_{i\in P_{j}}\delta_{i}w_{ij}^{(q)}$$		(4)
	$$\displaystyle\dfrac{\partial E}{\partial w_{ij}}=-\delta_{i}y_{j}^{(q)}$$		(5)
	$$\displaystyle w_{ij,new}=w_{ij}-\eta\biggl{\langle}\dfrac{\partial E}{\partial w_{ij}}\biggr{\rangle}$$		(6)
	$$\displaystyle w_{ij,new}^{(q)}=Q_{ij}(w_{ij,new})$$		(7)

where $l_{i}$ is the logit of the unit $i$, $\delta_{i}$ is the error signal of the unit $i$, $w_{ij}$ is the weight from the unit $j$ to the unit $i$, $y_{j}$ is the output activation of the unit $j$. $\eta$ is the learning rate, $A_{i}$ is the set of units anterior to the unit $i$, $P_{j}$ is the set of units posterior to the unit $j$, $Q(\cdot)$ is the weight quantizer, $\phi(\cdot)$ is the activation function. The superscript $(q)$ indicates the value is quantized, and $\langle\cdot\rangle$ is the average operation over the mini-batch. As described in Equation (2) to  (7), the forward, backward, and gradient calculation is conducted with the quantized weights, but weight update adopts the full-precision parameters. This is because that the quantization step size, $\Delta$, is usually much larger than the computed gradients, $\dfrac{\partial E}{\partial\mathbf{w}}$. The weights are not changed if the gradient is directly updated to the quantized weights.

The visualization method in [18, 9] shows the location of three weight vectors $\mathbf{w}_{1}$, $\mathbf{w}_{2}$, and $\mathbf{w}_{3}$. For locating these three models on the same loss surface, the projection vectors $\mathbf{u}$ and $\mathbf{v}$ are formed as follows:

	$$\displaystyle\mathbf{u}=(\mathbf{w}_{2}-\mathbf{w}_{1})$$		(8)
	$$\displaystyle\mathbf{v}=(\mathbf{w}_{3}-\mathbf{w}_{1})-\langle\mathbf{w}_{3}-\mathbf{w}_{1},\mathbf{w}_{2}-\mathbf{w}_{1}\rangle/\left\lVert\mathbf{w}_{2}-\mathbf{w}_{1}\right\rVert^{2}$$		(9)
	$$\displaystyle\mathbf{\hat{u}}=\frac{\mathbf{u}}{\left\lVert\mathbf{u}\right\rVert}$$		(10)
	$$\displaystyle\mathbf{\hat{v}}=\frac{\mathbf{v}}{\left\lVert\mathbf{v}\right\rVert}$$		(11)

The normalized vectors $\mathbf{\hat{u}}$ and $\mathbf{\hat{v}}$ form an orthonormal basis in the plane containing $\mathbf{w}_{1}$, $\mathbf{w}_{2}$, and $\mathbf{w}_{3}$. These three vectors can be visualized on a Cartesian grid in the basis $\mathbf{\hat{u}}$ and $\mathbf{\hat{v}}$ using a set of points $P$.

$\displaystyle P=\mathbf{w}_{1}+x\cdot\mathbf{\hat{u}}+y\cdot\mathbf{\hat{v}},$

(12)

where $x$ and $y$ are the coordinates.

We trained the 2-bit ternary quantized ResNet-20 [12] on the CIFAR-100 dataset [21] using the retraining algorithm [17]. After the performance has fully converged during the retraining process, we captured three quantized models $\mathbf{w}^{(q)}_{1}$, $\mathbf{w}^{(q)}_{2}$, and $\mathbf{w}^{(q)}_{3}$ with time intervals¹¹1$\mathbf{w}^{(q)}_{1}$, $\mathbf{w}^{(q)}_{2}$, and $\mathbf{w}^{(q)}_{3}$ that were captured at epochs 214, 232, and 250, respectively on the training epoch. Figure 1 (a) visualizes the three quantized networks using Equations (8) to (12).

It is noted that $\mathbf{w}^{(q)}_{1}$, $\mathbf{w}^{(q)}_{2}$, and $\mathbf{w}^{(q)}_{3}$ are located at ‘$\mathbf{w1}$’, ‘$\mathbf{w2}$’, and ‘$\mathbf{w3}$’, respectively. The limitation of this visualization method when applied to QDNNs is obvious. Even though $\mathbf{w}^{(q)}_{1}$, $\mathbf{w}^{(q)}_{2}$, and $\mathbf{w}^{(q)}_{3}$ are quantized weights, the other points between them are represented in full-precision. Thus, the exact shape of the loss surface cannot be determined when the weights are quantized. Hence, we plot the loss surface after quantizing the high-precision location vector, $P$. It is noteworthy that we can employ the full-precision weight vectors from Equation (6) to compute Equation (10) and (11). We denote these two normalized vectors as $\mathbf{\hat{u}}^{f}$ and $\mathbf{\hat{v}}^{f}$ to avoid confusion; subsequently, the three quantized vectors can be visualized on a Cartesian grid using a set of quantized points, $P^{q}$, for QDNNs as follows:

	$$\displaystyle\mathbf{w}^{f}=\mathbf{w}^{f}_{1}+x\cdot\mathbf{\hat{u}}^{f}+y\cdot\mathbf{\hat{v}}^{f}$$		(13)
	$$\displaystyle P^{q}=\text{sign}(\mathbf{w}^{f})\cdot\Delta\cdot\text{min}\Big{\{}\Big{\lfloor}\Big{(}\frac{|\mathbf{w}^{f}|}{\Delta}+0.5\Big{)}\Big{\rfloor},\frac{(M-1)}{2}\Big{\}}$$		(14)

where $\mathbf{w}^{f}_{1}$ is the full-precision weight vector that can be employed during retraining. We report the relationship of the three vectors $\mathbf{w}_{1}^{(q)}$, $\mathbf{w}_{2}^{(q)}$, and $\mathbf{w}_{3}^{(q)}$ obtained using the modified visualization method in Figure 1 (b). The relationship of the three quantized weight vectors, which cannot be observed in Figure 1 (a), is well represented. As they were captured in the epoch order (‘$\mathbf{w1}$’ $\rightarrow$ ‘$\mathbf{w2}$’ $\rightarrow$ ‘$\mathbf{w3}$’) during the retraining, a path along $\mathbf{w}_{1}^{(q)}$, $\mathbf{w}_{2}^{(q)}$ and $\mathbf{w}_{3}^{(q)}$ appeared. Because all of the points, $P^{(q)}$, were expressed in 2-bit quantized weights, the surface fluctuated strongly owing to quantization noise.

Network and hyperparameter configuration: We trained ResNet-20 [12] and MobileNetV2 [28] for the CIFAR-100 dataset. The training hyperparameters are as follows. For full-precision training, the batch size was 128 and the number of epochs trained was 175. An SGD optimizer with a momentum of 0.9 was used. The learning rate began at 0.1 and decreased by 0.1 times at the 75th and 125th epochs. Additionally, L2-loss was added with a scale of 5e-4. For QDNN retraining, the batch size and optimizer were the same as those of the full-precision one. The cyclical learning rate scheduling for retraining is described in Figure 3 (Top).

We captured the quantized models at the minimum points of the cyclical learning rate scheduling and obtained their average. To fine-tune the averaged model, the initial learning rate was set as 0.001 and decreased by 0.1 times at every epoch. We only performed three to four epochs for the fine-tuning. We did not employ L2-loss for the QDNN training as it conflicted with the clipping of the quantization.

In the next step of the SQWA training process, we performed QDNN training with cyclical learning rate scheduling, as depicted in Figure 3 (Top). We captured the models when the learning rates were the lowest in the cycles (i.e., blue diamonds in Figure 3 (Top and Middle)). To select the models for averaging, we considered two groups of the networks, as depicted in Figure 3 (Bottom). The first group (dashed red box) was selected at the beginning of the training and the other group (solid blue box) was captured after a sufficient number of training epochs has elapsed. We employed seven models for both groups, and their performances are compared in Table 2.

More specifically, the models in the first group were captured between the 40th and 76th²²2The training performance was too low to capture for models earlier than 40th epochs. epochs. Their test accuracies were approximately 64.7% and the averaged model demonstrated an accuracy of 67.94%. It is noteworthy that we took the average of seven 2-bit ternary QDNNs ($-\Delta$, 0, and $\Delta$), of which the resultant model was a 4-bit ($-7\hat{\Delta}$, $-6\hat{\Delta}$, …, 0, …, $6\hat{\Delta}$, and $7\hat{\Delta}$) QDNN. We conducted re-quantization and fine-tuned the averaged model to obtain a final 2-bit QDNN, which yielded a test accuracy of 66.75%. The result of the second group was significantly better than that of the first group. The averaged 4-bit model yielded a test accuracy of 68.81%. After the fine-tuning, the final performance of the 2-bit QDNN was 67.75%. From these results, we can deduce the following:

• •

  SQWA can be fully utilized when the models are captured after a sufficient convergence.

• •

  Based on the observation of the direct quantization results in Table 2, the second group forms a wider minima in the loss surface than the first group. Because direct quantization can be interpreted as a noise injection operation, less performance degradation suggests that the model is laying in a wider minimum or at the center of the loss surface.

We compare our SQWA results with those of previous studies in Table 3. Our proposed SQWA method outperforms the methods of previous studies. In particular, SQWA indicated 0.8%, 1.78%, 1.22%, and 0.33% higher test accuracies than DoReFa-Net [43], Residual [11], LQ-Net [41], and WNQ [4], respectively. This result is encouraging as the previous studies employed 2-bit 4-level quantizers, whereas we adopted 2-bit ternary and 1-bit binary quantizer. Furthermore, we compare our result to those involving 2-bit ternary and 1-bit binary quantizer. Our method with 2-bit ternary quantizer outperformed HLHLp [31] and KDQ [30] in terms of test accuracy by 1.31% and 0.75%, respectively. For the binary weights, SQWA achieves 2.18% higher accuracy than KDQ. It should be noted that KDQ improves the performance of the QDNN using the KD. Additionally, we exploit the KD technique to obtain a high-performance full-precision model. Because SQWA outperforms KDQ, it suggests that SQWA training methods can combine with KD. Furthermore, we evaluated the proposed SQWA method using MobileNetV2, which has a larger number of parameters than ResNet20. We trained a full-precision MobileNetV2 with KD and SWA and achieved a test accuracy of 77.64%. We exploited SQWA with the same cyclical learning rate scheduling used in the ResNet20 experiment. After a sufficient number of epochs, we captured seven models to establish a 4-bit averaged model and fine-tuned it. Our final 2-bit MobileNetV2 yielded the test accuracy that was 1.76% and 1.22% higher than those of L2Quant [3] and HLHLp [31], respectively.

Network and hyperparameter configuration: We trained ResNet-18 [12] for the ILSVRC 2012 classification dataset [22]. The training hyperparameters are as follows. We trained the full-precision model with a batch size of 1024 on 90 epochs, with the initial learning rate of 0.4 and decreased it by 0.1 times at the 30th, 60th, and 80th epochs. It is noteworthy that the initial learning rate of 0.4 was determined using the linear scaling rule, as suggested in [10]. We used the SGD optimizer with a momentum of 0.9. Additionally, L2-loss was added with a scale of 1e-4. For the QDNN retraining, the batch size and optimizer were the same as those of the full-precision training. Because we employed SQWA training, cyclical learning rate scheduling was employed, as shown in Figure 5 (Top). The maximum and minimum values of the learning rates were determined by considering the learning rate of the full-precision training, as suggested in Section 4.

We captured the models at the minimum points of the cyclical learning rate scheduling and obtained their average. To fine-tune the averaged model, the initial learning rate was set to 0.004 and decreased by 0.1 times at every epoch. We executed five epochs for the fine-tuning and did not employ L2-loss for the QDNN training.

We conducted additional experiments to investigate the effect of number of models on averaging. As discussed in Section 4, the number of captured models is related to the precision of the averaged model. More specifically, the averaged model using three 2-bit ternary models becomes a 3-bit QDNN. Thus, we employed 3, 7, 15, and 31 models such that the precision of the averaged model was 3, 4, 5, and 6 bits, respectively, and fine-tuned each model. The results are reported in Table 5. Adopting three models afforded a top-1 accuracy of 69.18%, which is 0.22% worse than the seven models. When 15 models were employed, the top-1 accuracy was similar to that of the 7 models but the top-5 accuracy was 0.2% higher. Using 31 models did not improve the performance.

We compare our results with those of previous studies in Table 6. Our result outperformed those of previous studies including the 2-bit 4-level (LQ-NET [41] and WNQ [4]) and ternary (TWN [8], TTQ [45], INQ [42], ADMM [23], QNet [39], QIL [20], and Apprentice [24]). More specifically, we achieved a top-1 accuracy of 69.4%. Only the QNet result is comparable with our result, although it is 0.3% lower. This result is significant because QNet employs non-linear quantizer while we adopt uniform quantization.