提出新方法让低比特模型训练更稳定高效
StableQAT: Stable Quantization-Aware Training at Ultra-Low Bitwidths
- 基于傅里叶分析设计新型梯度代理，替代传统STE
- 2-4比特下训练稳定，性能优于现有方法
- 无需额外计算开销，适合部署在资源受限场景
量化感知训练（QAT）对在严格内存与延迟约束下部署大模型至关重要，但在超低比特设置下实现稳定可靠的优化仍具挑战。基于直通估计器（STE）或软量化的方法常因梯度不匹配、不稳定或高计算开销而受限。为此，我们提出StableQAT，一种统一且高效的QAT框架，通过离散傅里叶分析导出的轻量级、理论严谨的反向传播代理，实现超低比特下的训练稳定。该代理严格涵盖STE作为特例，生成平滑、有界且低成本的梯度，在多种超参数下均提升训练性能与稳定性。实验表明，StableQAT在2-4比特区间表现出色，训练稳定、鲁棒性强，且相比标准QAT方法性能更优，训练开销可忽略。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。