StableQAT: 극도로 낮은 비트폭에서 안정적인 양자화 인식 학습

StableQAT: Stable Quantization-Aware Training at Ultra-Low Bitwidths

양자화 인식 학습(QAT)은 엄격한 메모리 및 지연 제약 조건 하에서 대규모 모델을 배포하는 데 필수적이지만, 극도로 낮은 비트폭에서 안정적이고 강력한 최적화를 달성하는 것은 여전히 어려운 과제입니다. Straight-through estimator (STE) 또는 소프트 양자화기를 기반으로 하는 일반적인 방법은 종종 기울기 불일치, 불안정성 또는 높은 계산 오버헤드를 겪습니다. 따라서 우리는 StableQAT라는 통합적이고 효율적인 QAT 프레임워크를 제안합니다. StableQAT는 원본 연산의 이산 푸리에 분석에서 파생된 새로운, 경량적이며 이론적으로 뒷받침되는 대체 함수를 사용하여 극도로 낮은 비트 환경에서 학습을 안정화합니다. StableQAT는 STE를 엄격하게 일반화하며, STE는 우리 프레임워크의 더 표현력이 풍부한 대체 함수 패밀리의 특수한 경우로 나타납니다. 이를 통해 QAT 학습 성능과 안정성을 향상시키는 부드럽고 경계가 있으며 저렴한 기울기를 제공합니다. 실험 결과, StableQAT는 2~4비트 환경에서 안정적이고 효율적인 QAT를 보여주며, 표준 QAT 기술에 비해 향상된 학습 안정성, 견고성 및 우수한 성능을 제공하며, 동시에 무시할 만한 학습 오버헤드를 발생시킵니다. 저희 코드는 다음 링크에서 확인할 수 있습니다: https://github.com/microsoft/StableQAT.

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.