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.
翻译:量化感知训练(QAT)对于在严格的内存和延迟约束下部署大型模型至关重要,然而在超低位宽下实现稳定且鲁棒的优化仍然具有挑战性。基于直通估计器(STE)或软量化器的常见方法通常存在梯度失配、不稳定性或高计算开销等问题。为此,我们提出了StableQAT,一个统一且高效的QAT框架。该框架通过对舍入算子进行离散傅里叶分析,推导出一种新颖、轻量且理论依据充分的替代反向传播方法,从而稳定了超低位设置下的训练。StableQAT严格地将STE推广为一种特例,属于我们更具表达力的替代函数族,其产生的梯度平滑、有界且计算成本低,从而在各种超参数选择下提升了QAT的训练性能和稳定性。在实验中,StableQAT在2-4位量化区间内展现出稳定高效的QAT性能,相比标准QAT技术,在训练开销可忽略不计的情况下,表现出更高的训练稳定性、鲁棒性以及更优的性能。我们的代码可在 https://github.com/microsoft/StableQAT 获取。