WinQ: Accelerating Quantization-Aware Training of Language Models Around Saddle Points

TL;DR

This paper analyzes the convergence issues in quantization-aware training of language models and proposes WinQ, an algorithm that accelerates training and improves quantization performance, especially at low bit-widths.

Contribution

The paper introduces WinQ, a novel method that accelerates quantization-aware training by weight resetting and gradient regularization, leading to significant speedups and accuracy improvements.

Findings

WinQ accelerates QAT by up to 4 times across various models and methods.

WinQ improves sub-4-bit quantization accuracy by up to 8.8%.

Hessian spectrum analysis reveals weights converge to saddle points with eigenvalues near zero.

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.