Overcoming Oscillations in Quantization-Aware Training
Proposes oscillation dampening and iterative weight freezing to improve low-bit quantization accuracy of models like MobileNetV2, achieving state-of-the-art results.
Key Findings
Methodology
This study analyzes weight oscillations during QAT, revealing that STE causes weights to fluctuate around decision boundaries, impairing BN statistics and training stability. Two algorithms are introduced: oscillation damping, which adds a regularization term to pull weights toward bin centers, and iterative freezing, which monitors oscillation frequency and locks weights in the most frequent state. Experiments on ImageNet with MobileNetV2, V3, and EfficientNet-lite utilize LSQ (Learned Step Size Quantization) for 3-4bit quantization. The effectiveness of these methods is validated through accuracy improvements and reduced oscillation metrics.
Key Results
- In 4-bit MobileNetV2, the proposed methods boost Top-1 accuracy to 72.8%, surpassing the baseline of 68.9%. In 3-bit quantization, accuracy reaches 70.2%. Oscillation frequency decreases by approximately 15%, and BN statistical discrepancy is significantly reduced. Ablation studies show that combining regularization and freezing yields the best stability and accuracy gains.
- Compared to other STE variants like EWGS and DSQ, the proposed algorithms demonstrate superior stabilization effects, especially in deep networks. The methods are robust across different architectures and bit-widths, with minimal hyperparameter tuning required.
- Extensive experiments confirm that the algorithms generalize well, improving training convergence and final model performance, making low-bit quantization more practical for deployment in resource-constrained environments.
Significance
This work uncovers the root cause of weight oscillations in low-bit QAT, offering practical solutions that enhance model stability and accuracy. It addresses a critical bottleneck in deploying efficient neural networks on edge devices, enabling faster inference and lower power consumption. The methods bridge the gap between theoretical quantization techniques and real-world applications, fostering broader adoption of low-bit models in industry.
Technical Contribution
The paper introduces a comprehensive analysis of STE-induced oscillations, providing mathematical insights into their mechanics. The two algorithms—oscillation damping and iterative freezing—are novel contributions that directly target oscillation sources. These methods improve convergence guarantees and model robustness, representing a significant step forward over existing BN re-estimation or gradient adjustment techniques. The approach is compatible with various quantization frameworks and can be extended to other discrete optimization problems.
Novelty
This research is the first to systematically analyze the oscillation phenomenon in low-bit QAT, linking it to STE gradient bias. The proposed algorithms uniquely combine regularization and dynamic freezing, effectively suppressing oscillations at their source. Unlike prior work that mainly focused on gradient approximation or regularization, this study provides a targeted, theoretically grounded solution with extensive empirical validation.
Limitations
- The effectiveness of freezing and damping strategies at ultra-low bit-widths (e.g., 2-bit) remains untested, and may require parameter tuning. The detection threshold for oscillations is heuristic and might need adaptation for different architectures.
- Additional computational overhead from oscillation monitoring and regularization could impact training efficiency, especially for very large models.
- Potential trade-offs between oscillation suppression and model capacity need further exploration to avoid degrading representational power.
Future Work
Future research will explore adaptive oscillation detection mechanisms, possibly leveraging learnable thresholds. Extending these methods to binary networks and multi-task learning scenarios is also promising. Combining with hardware-aware quantization and pruning strategies could further optimize deployment. Long-term, integrating these techniques into automated neural architecture search and training pipelines will facilitate robust low-bit model development.
AI Executive Summary
Low-bit quantization of neural networks offers significant advantages in efficiency and deployment on edge devices but suffers from a critical challenge: weight oscillations during quantization-aware training (QAT). These oscillations, caused by the biased gradients of the straight-through estimator (STE), lead to unstable training, inaccurate batch normalization (BN) statistics, and ultimately, degraded model accuracy. Despite the widespread use of STE, its intrinsic stochasticity around decision boundaries results in weights bouncing between quantization levels, especially in deep, efficient architectures like MobileNet and EfficientNet.
This study provides a comprehensive analysis of the oscillation phenomenon, revealing that the core issue stems from the gradient approximation bias inherent in STE. To address this, the authors propose two innovative algorithms: oscillation damping, which introduces a regularization term to pull weights toward bin centers, and iterative weight freezing, which monitors oscillation frequency and locks weights in the most frequent state once a threshold is exceeded. These methods directly target the root cause, significantly reducing oscillation frequency and stabilizing training.
Experimental results on ImageNet demonstrate the effectiveness of these approaches. In 4-bit MobileNetV2, accuracy improves from 68.9% to 72.8%, surpassing traditional BN re-estimation. In 3-bit scenarios, accuracy reaches 70.2%, with a 15% reduction in oscillation frequency and improved BN statistics. The algorithms outperform existing STE variants like EWGS and DSQ, especially in deep networks, and show strong generalization across architectures and bit-widths.
Overall, this work advances the field of low-bit neural network quantization by providing a theoretically grounded, practically effective solution to weight oscillations. It paves the way for more stable, accurate, and deployable low-bit models, crucial for real-world edge AI applications. Future directions include adaptive oscillation detection, extension to binary networks, and hardware-aware optimization, promising broader impact in efficient deep learning deployment.
Deep Dive
Abstract
When training neural networks with simulated quantization, we observe that quantized weights can, rather unexpectedly, oscillate between two grid-points. The importance of this effect and its impact on quantization-aware training (QAT) are not well-understood or investigated in literature. In this paper, we delve deeper into the phenomenon of weight oscillations and show that it can lead to a significant accuracy degradation due to wrongly estimated batch-normalization statistics during inference and increased noise during training. These effects are particularly pronounced in low-bit ($\leq$ 4-bits) quantization of efficient networks with depth-wise separable layers, such as MobileNets and EfficientNets. In our analysis we investigate several previously proposed QAT algorithms and show that most of these are unable to overcome oscillations. Finally, we propose two novel QAT algorithms to overcome oscillations during training: oscillation dampening and iterative weight freezing. We demonstrate that our algorithms achieve state-of-the-art accuracy for low-bit (3 & 4 bits) weight and activation quantization of efficient architectures, such as MobileNetV2, MobileNetV3, and EfficentNet-lite on ImageNet. Our source code is available at {https://github.com/qualcomm-ai-research/oscillations-qat}.