StableQAT: Stable Quantization-Aware Training at Ultra-Low Bitwidths
This addresses the problem of deploying large models under strict memory and latency constraints for AI practitioners, though it appears incremental as it builds on existing QAT methods.
The paper tackles the challenge of achieving stable and robust optimization in quantization-aware training at ultra-low bitwidths, proposing StableQAT, which demonstrates improved training stability and superior performance with negligible overhead in 2-4 bit regimes.
Quantization-aware training (QAT) is essential for deploying large models under strict memory and latency constraints, yet achieving stable and robust optimization at ultra-low bitwidths remains challenging. Common approaches based on the straight-through estimator (STE) or soft quantizers often suffer from gradient mismatch, instability, or high computational overhead. As such, we propose StableQAT, a unified and efficient QAT framework that stabilizes training in ultra low-bit settings via a novel, lightweight, and theoretically grounded surrogate for backpropagation derived from a discrete Fourier analysis of the rounding operator. StableQAT strictly generalizes STE as the latter arises as a special case of our more expressive surrogate family, yielding smooth, bounded, and inexpensive gradients that improve QAT training performance and stability across various hyperparameter choices. In experiments, StableQAT exhibits stable and efficient QAT at 2-4 bit regimes, demonstrating improved training stability, robustness, and superior performance with negligible training overhead against standard QAT techniques. Our code is available at https://github.com/microsoft/StableQAT.