One-Step Forward and Backtrack: Overcoming Zig-Zagging in Loss-Aware Quantization Training

TL;DR

This work identifies a zig-zagging-like instability arising from gradient quantization errors in loss-aware quantization (LAQ) and shows it can severely slow convergence, especially at extremely low bitwidth. It then proposes BLAQ, a two-stage, one-step forward and backtrack framework that uses a forward trial gradient to guide the current step and a backtracking update that blends trial and current gradients to stabilize updates. The authors provide convergence analysis under standard convex assumptions and demonstrate through extensive experiments on CIFAR10, MNIST, SVHN, and ImageNet that BLAQ achieves faster convergence and higher accuracy than LAQ and many baselines, particularly in 1-bit and 2-bit settings. The approach offers a robust, pruning-free path to effective ultra-low-bit quantization with practical implications for edge deployments of DNNs.

Abstract

Weight quantization is an effective technique to compress deep neural networks for their deployment on edge devices with limited resources. Traditional loss-aware quantization methods commonly use the quantized gradient to replace the full-precision gradient. However, we discover that the gradient error will lead to an unexpected zig-zagging-like issue in the gradient descent learning procedures, where the gradient directions rapidly oscillate or zig-zag, and such issue seriously slows down the model convergence. Accordingly, this paper proposes a one-step forward and backtrack way for loss-aware quantization to get more accurate and stable gradient direction to defy this issue. During the gradient descent learning, a one-step forward search is designed to find the trial gradient of the next-step, which is adopted to adjust the gradient of current step towards the direction of fast convergence. After that, we backtrack the current step to update the full-precision and quantized weights through the current-step gradient and the trial gradient. A series of theoretical analysis and experiments on benchmark deep models have demonstrated the effectiveness and competitiveness of the proposed method, and our method especially outperforms others on the convergence performance.

Figures (6)

• Figure 1: Trajectories of weight updating obtained by LAQ and BLAQ in 1-bit quantization on the loss function $\ell(\omega)=5\left(\omega_{1}-0.054\right)^{2}+\left(\omega_{2}+0.055\right)^{2}$. The black line shows the evolution of full-precision weight values via iterations during training, while the red and yellow lines are the evolution of 1-bit quantization case.
• Figure 2: Regression of randomly selected weights during the last 100 epochs of training ResNet18 on ILSVRC12 (ImageNet) during LAQ 2-bit quantization. The (a) show the full-precision weights, and the (b) plot the corresponding quantized weights.
• Figure 3: Regression with the same weights as Fig. \ref{['figVGG1']} selected during the last 100 epochs of training ResNet18 on ILSVRC12 (ImageNet) during BLAQ 2-bit quantization. The (a) show the full-precision weights, and the (b) plot the corresponding quantized weights.
• Figure 4: Two-stage updating process: the blue rectangles denote the real weights, and the red ones are the trial weights.
• Figure 5: Training loss trained with LAQ/baseline and BLAQ. The red curve shows the evolution of mean loss values obtained by BLAQ, while the blue one by LAQ. The shaded areas indicate the fluctuation (error) range of the mean loss curve.