🤖 AI Summary
This work addresses the performance plateau and slow convergence commonly observed in sub-4-bit quantization-aware training (QAT), which arises as weights converge to flat saddle regions of the loss landscape. The study reveals, for the first time, that this bottleneck is closely tied to the spectral properties of the Hessian: weights tend to settle in regions where Hessian eigenvalues approach zero. To overcome this, the authors propose WinQ, an algorithm combining weight interpolation-based resetting with noise-injected gradient regularization to effectively escape saddle points and accelerate optimization. Evaluated across 16 experimental settings, WinQ consistently improves performance—yielding up to an 8.8% gain for sub-4-bit models under identical training costs—and accelerates QAT convergence by as much as fourfold.
📝 Abstract
Quantization-aware training (QAT) is widely adopted to quantize language models by training full-precision weights using gradients from the quantized model. The main bottleneck is its slow convergence and early performance plateau, particularly below 4-bit-widths. While this problem has been observed in prior work, its precise cause remains unclear. In this paper, we analyze the convergence of QAT by estimating the spectrum of the loss-surface Hessians. We find that the weights converge to flat regions around saddle points, where a large fraction of the Hessian eigenvalues are both positive and negative. During training, an increasing fraction of Hessian eigenvalues concentrates around zero, whose magnitude decreases. At lower bit-widths, the magnitude of eigenvalues in the Hessian spectrum is significantly smaller. To mitigate these issues, we propose an algorithm called WinQ to accelerate QAT, which involves: (1) periodically resetting weights to the linear interpolation of full-precision and quantized weights, reducing the distance to the quantization grid and increasing eigenvalue magnitude, and (2) computing gradients of noise-injected weights to regularize the Hessian. Extensive experiments show that WinQ accelerates QAT by up to 4 times across various quantization methods and models. Under the same training cost, WinQ improves state-of-the-art sub-4-bit quantization by up to 8.8%. These results are consistent across 16 settings with different language models, quantization methods, and bit widths.