Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Tianyi Chen, Sihan Chen, Xiaoyi Qu, Dan Zhao, Ruomei Yan, Jongwoo Ko, Luming Liang, Pashmina Cameron

TL;DR

This work tackles forward–backward mismatch in quantization-aware training (QAT) at ultra-low bitwidths by introducing StableQAT, which uses a Rotated Damped Fourier Surrogate (RDFS) to model the rounding operator. By rotating the coordinate system and applying a truncated Fourier series with a tunable amplitude, RDFS yields smooth, bounded gradients that generalize STE and improve training stability without extra computational overhead. The authors provide theoretical results showing that the surrogate achieves smaller $L^2$ approximation error than STE and bounded gradient variance compared to DSQ, with an asymptotic variance limit of about $0.077$ for maximal sharpness. Empirically, StableQAT delivers stable and improved performance on large language models (LLaMA-3 series) at 2–4 bits, often surpassing FP16 baselines, while maintaining efficiency and fusion-friendliness. The approach offers a practical, scalable solution for deploying ultra-low-bit quantized models with strong stability and accuracy gains.

Abstract

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.

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

TL;DR

This work tackles forward–backward mismatch in quantization-aware training (QAT) at ultra-low bitwidths by introducing StableQAT, which uses a Rotated Damped Fourier Surrogate (RDFS) to model the rounding operator. By rotating the coordinate system and applying a truncated Fourier series with a tunable amplitude, RDFS yields smooth, bounded gradients that generalize STE and improve training stability without extra computational overhead. The authors provide theoretical results showing that the surrogate achieves smaller approximation error than STE and bounded gradient variance compared to DSQ, with an asymptotic variance limit of about for maximal sharpness. Empirically, StableQAT delivers stable and improved performance on large language models (LLaMA-3 series) at 2–4 bits, often surpassing FP16 baselines, while maintaining efficiency and fusion-friendliness. The approach offers a practical, scalable solution for deploying ultra-low-bit quantized models with strong stability and accuracy gains.

Abstract

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.
Paper Structure (51 sections, 6 theorems, 67 equations, 8 figures, 3 tables)

This paper contains 51 sections, 6 theorems, 67 equations, 8 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Forward-Backward Mismatch in QAT.
  3. Contributions
  4. Preliminary and Related Work
  5. Quantization
  6. Straight-Through Estimator (STE)
  7. Soft Surrogate Quantization
  8. StableQAT
  9. Backward Pass via RDFS:
  10. Rotated Damped Fourier Surrogate (RDFS)
  11. Coordinate Rotation.
  12. Fourier Approximation.
  13. Inverse Rotation.
  14. Damped Amplitude and Gradient Stabilization
  15. Role of Amplitude.
  16. ...and 36 more sections

Key Result

Theorem 4.1

Let $f \in L^2([0,T])$ and let $f_n$ be the $n$th partial Fourier sum of $f$. Then, for all $n \in \mathbb{N}$, (i) $f_n$ is the unique minimizer of the $L^2$ approximation error among trigonometric polynomials of degree at most $n$; (ii) For $n \geq 1$, $\|f - f_n\|_2 < \|f - f_0\|_2$ if and only i

Figures (8)

  • Figure 1: StableQAT procedure.
  • Figure 2: Rotated Damped Fourier Surrogate (RDFS) under different orders and amplitudes.
  • Figure 3: Gradient spread of DSQ and StableQAT. As the surrogate sharpens toward rounding operator, DSQ exhibits exploding variance, while StableQAT shows bounded variance.
  • Figure 4: Time and space cost comparison (lower is better). They are measured on an input tensor with batch size 4 and sequence length 128. Each method is repeated 20 times.
  • Figure 5: Training loss (left) and gradient norm (right) comparison for Llama-3-1B under different learning rates.
  • ...and 3 more figures

Theorems & Definitions (13)

  • Theorem 4.1
  • proof
  • Theorem 4.2
  • proof
  • Theorem 3.1
  • proof
  • Lemma 4.1: DSQ Gradient
  • proof
  • Lemma 4.2
  • proof
  • ...and 3 more