Quantization-Aware Training Methods

• Quantization-aware training is a methodology that simulates low-precision computation by embedding fake quantization modules into neural network training for hardware-constrained deployment.
• It integrates quantization operators into both forward and backward passes, using surrogate gradients like STE and RDFS to mitigate instability and bias in ultra-low-bit regimes.
• StableQAT demonstrates improved accuracy and controlled gradient variance with minimal modifications to standard training pipelines, making it effective for LLM and vision models.

Quantization-aware training (QAT) is a class of neural network optimization methodologies that simulates low-precision computation during training to enable robust model deployment on memory- and latency-constrained hardware. QAT integrates quantization operators—most commonly low-bit uniform or non-uniform rounding—directly into the forward and backward passes, allowing parameters and internal statistics to adapt to discretization noise. The primary challenge, exacerbated in the ultra-low-bit (2–4 bit) regime, is achieving stable, unbiased gradient flow and maintaining generalization and accuracy. Recent advances have introduced theoretically grounded surrogates for gradient estimation, hyperparameter-efficient workflows, feature-based regularizations, and principled frameworks to address instability and variance explosion inherent in classical STE approaches (Chen et al., 27 Jan 2026).

1. Core Principles of Quantization-Aware Training

QAT modifies the training loop of a neural network model by embedding “fake quantization” modules at critical points—typically after weight and activation tensors—so that forward computation uses quantized values matching the target inference bit-width (e.g., INT4, INT8). The quantization operator is generally parameterized as

$$v_q = s \cdot \mathsf{clip}(\mathsf{round}(v/s) + z;\;q_\mathrm{min},\,q_\mathrm{max})$$

where $s$ and $z$ are trainable scale and zero-point parameters, and $q_\mathrm{min}$/$q_\mathrm{max}$ set the representable integer range (e.g., $[-7, 8]$ for 4-bit signed). In the backward pass, gradients must propagate through the non-differentiable $\mathsf{round}$ and $\mathsf{clip}$ operators. Naively, the Jacobian $\partial v_q/\partial v$ is zero almost everywhere, so QAT typically employs the Straight-Through Estimator (STE):

$$\frac{\partial v_q}{\partial v} \approx 1 \;\;\text{(inside quantization range)}$$

which treats quantization as an identity operation for gradient flow. While simple, STE introduces substantial forward-backward mismatch, especially at low bit-width, leading to bias, instability, and poorly conditioned optimization (Chen et al., 27 Jan 2026). More recent approaches propose surrogate gradients derived via discrete Fourier analysis, smooth continuous relaxations, or curvature-aware corrections to mitigate these pathologies.

2. Surrogate Gradient Formulations and Stability

Traditional QAT recipes (STE, soft quantizers) suffer pronounced instability in the 2–3 bit regime:

• The STE passes gradients as if $\mathsf{round}$ is the identity, ignoring threshold sensitivity and grid-induced bias.
• Gradient mismatch manifests as exploding or vanishing gradients near quantization thresholds.
• Variance in gradient estimators spikes, and training collapses under ill-conditioned hyperparameters.

StableQAT introduces the Rotated Damped Fourier Surrogate (RDFS), which interprets the rounding operator as a rotated triangle wave expanded in a Fourier series (Chen et al., 27 Jan 2026). The $M$-term truncation yields a smooth, bounded surrogate for the Jacobian:

$$g_M(x,x_q) = \frac{1 - A\sqrt{2}\pi \sum_{m=0}^M \frac{(-1)^m}{2m+1}\cos[(2m+1)\pi(x+x_q)]} {1 + A\sqrt{2}\pi \sum_{m=0}^M \frac{(-1)^m}{2m+1}\cos[(2m+1)\pi(x+x_q)]}$$

In practice, first-order ($M=0$) suffices:

$$g(x,x_q) = \frac{1 - A\sqrt{2}\pi \cos(\pi(x + x_q))}{1 + A\sqrt{2}\pi \cos(\pi(x + x_q))}$$

Properties:

• $A=0$ limit recovers the STE.
• Surrogate is $C^\infty$ in $x$ and strictly bounded for $A<1/(\sqrt{2}\pi)$.
• Gradient variance remains bounded even as sharpness increases ($\operatorname{Var} \approx 0.076$ for $A \to 1/(\sqrt{2}\pi)$).
• Computational cost is dominated by a single $\cos$ per activation, making RDFS as efficient as STE and 3–5$\times$ faster than exp-based soft surrogates.

StableQAT yields stable, robust convergence curves, controlled gradient norms, and error-bar dispersion minimization across hyperparameter choices, particularly in ultra-low-bit LLM and vision scenarios (Chen et al., 27 Jan 2026).

3. Integration and Hyperparameter Management

StableQAT requires only minimal modification to standard QAT pipelines: replacing STE with the RDFS gradient surrogate in custom PyTorch-style autograd functions. Usage example:

class QuantizeSTE(torch.autograd.Function): @staticmethod def forward(ctx, x, scale, zero_point, A): # Quantization: scale, round, clip x_scaled = x / scale x_q = torch.round(x_scaled) + zero_point x_q_clipped = x_q.clamp(qmin, qmax) ctx.save_for_backward(x_scaled, x_q_clipped, torch.tensor(A)) return (x_q_clipped - zero_point) * scale @staticmethod def backward(ctx, grad_output): x_scaled, x_q_clipped, A = ctx.saved_tensors c = A * math.sqrt(2) * math.pi arg = math.pi * (x_scaled + x_q_clipped) cosv = torch.cos(arg) g = (1 - c * cosv) / (1 + c * cosv) return grad_output * g / scale, None, None, None

Key hyperparameters:

• Bitwidth ($b$): Set quantization grid; stable optimization is most challenging at $b=2,3$.
• Amplitude ($A$): Controls surrogate sharpness ($A=0$ recovers STE). Default $A=0.21$ yields robustness across LLMs and vision models; avoid $A \gtrsim 0.25$–$0.3$ unless tuning for sharpness.
• Fourier order ($M$): $M=0$ is sufficient; higher orders give negligible accuracy gain and increased computation.
• Learning rate: Standard QAT schedules ($1\mathrm{e}{-5}$–$2\mathrm{e}{-4}$) are compatible due to bounded gradient variance.
• Warmup: 1–5% is typically sufficient.

No special scheduling, annealing, or surrogate parameter ramp-up is needed; StableQAT is plug-and-play in existing QAT workflows.

4. Comparative Empirical Performance

The efficacy of StableQAT is demonstrated across transformer and vision models at bit-widths of 2–4 bits (Chen et al., 27 Jan 2026):

• At 4 bits, StableQAT marginally surpasses full-precision (FP16) accuracy in certain configurations and outperforms ParetoQ and DSQ by $+0.3$–$1.4$ points on benchmarks including ARC-E, ARC-C, BoolQ, HellaSwag, OpenBookQA, PIQA, SciQ, Winogrande.
• Gains at 3 bits are pronounced: $+2.7$–$6.9$ points over baselines; at 2 bits, StableQAT achieves stability with $+1.7$ point improvement.
• Backward-pass latency and memory match that of STE, with substantial speedup against soft quantization relaxations (DSQ).
• Training-loss trajectories exhibit monotonic decrease and controlled gradient norm evolution, with minimal sensitivity to random seed or learning rate choice.

Robust deployment defaults for large models: bitwidth=4, $A=0.21$, $M=0$, standard QAT learning rates, 1–5% warmup, weight-only quantization for LLMs (forward path matches standard QAT).

5. Practical Implementation and Deployment

StableQAT is branch-free and fusion-friendly; the surrogate $g(x,x_q)$ depends only on $(x + x_q)$, allowing vectorized implementation suitable for CUDA kernel fusion and custom operators. Guidelines for latency-sensitive inference:

• Implement surrogate gradient as a fused operator for minimal backward overhead.
• Monitor gradient norm and training loss; reduce $A$ slightly if spikes occur.
• Weight-only quantization avoids activation quantization penalty at inference; StableQAT forward computation is identical to standard QAT.

Empirical benchmarks indicate efficient scaling for LLM and vision architectures at 2–4 bits, with negligible additional memory or compute load over STE.

6. Relationships to Prior QAT Methodologies

StableQAT strictly generalizes the STE: the latter emerges as the $A \to 0$, $M = -1$ limit of the surrogate family. Unlike exp-based soft relaxations, StableQAT achieves bounded variance and robust convergence without needing a schedule for temperature or surrogate annealing. RDFS surrogates are formally justified—the $M$th partial Fourier sum is the optimal $L^2$ approximation among degree-$M$ trigonometric polynomials to the rotated rounding operator. These theoretical properties endow StableQAT with stability, robustness, and efficiency at regimes where past STE and soft quantization surrogates are provably ill-conditioned (Chen et al., 27 Jan 2026).

7. Future Directions and Theoretical Insights

The combination of discrete Fourier analysis and surrogate gradient theory used in StableQAT offers a pathway for further extensions in ultra-low-bit quantization, architecture search, and hardware-efficient training. Its strictly bounded surrogate class may be generalized to alternative quantization topologies, mixed-precision networks, and non-uniform grids. The framework also clarifies the inherent limitations in classical STE and soft quantization relaxations, demonstrating that optimal stability and variance control necessitate analytic gradient surrogates constructed via functional approximation of rounding. Empirical validation suggests further exploration of high-order surrogate expansions or adaptive amplitude schedules in dynamic workload or transfer learning contexts.

---

StableQAT establishes a unified, theoretically grounded backbone for quantization-aware training in resource-constrained deployment scenarios, demonstrating unequivocal superiority in stability, robustness, and efficiency across diverse neural architectures and ultra-low-precision quantization regimes (Chen et al., 27 Jan 2026).