Quantization-Aware Training (QAT)

QAT inserts FakeQuant operations into the computational graph during training. These nodes simulate the clamping, scaling, and rounding behavior of target integer formats (e.g., INT8) on the forward pass, while maintaining full-precision weights for the backward pass update.

• Mechanism: Uses the Straight-Through Estimator (STE) to bypass the non-differentiable rounding operation during backpropagation.
• Result: The optimizer finds wide, flat minima in the loss landscape that are robust to the noise floor introduced by quantization error.

To map floating-point tensors to integer grids, QAT dynamically tracks the min/max ranges of activations using moving average observers. This ensures the quantization grid covers the dynamic range of the signal without excessive clipping.

• Clipping Mitigation: Learns optimal thresholds to balance saturation error (clipping) against rounding error (granularity).
• Per-Channel Quantization: Often applies distinct scales per convolutional filter channel, preserving variance across depth dimensions that would otherwise be homogenized by a single tensor-wide scale.

Before quantization, Batch Normalization (BN) layers are mathematically absorbed into the preceding convolutional weights and biases. This is a critical pre-processing step for QAT because BN parameters operate at a different numerical scale.

• Graph Optimization: Converts y = γ * (x - μ) / σ + β into a single affine operation on the weights.
• Inference Efficiency: Eliminates redundant runtime operations and prevents the quantization of the small, high-precision BN parameters from destabilizing the integer arithmetic pipeline.

The Straight-Through Estimator (STE) is the mathematical trick that makes QAT possible. It treats the hard rounding function as an identity map for gradient calculation, allowing error signals to pass through the discrete quantization barrier.

• Backward Pass: ∂L/∂x ≈ ∂L/∂y (ignoring the zero-gradient of the step function).
• Constraint: Effective only when the quantization step size is small relative to the weight variance; otherwise, the gradient mismatch introduces significant bias, leading to unstable convergence.

While PTQ applies quantization to a static, pre-trained model using a small calibration dataset, QAT retrains the model to anticipate the error. This distinction is critical for low-bit precision (≤ 4-bit).

• Accuracy Retention: QAT typically recovers >99% of FP32 accuracy at INT8, whereas PTQ often suffers a 1-5% drop on complex tasks.
• Use Case: PTQ is sufficient for benign models; QAT is mandatory for aggressive compression or models with highly non-linear activation distributions, such as dense RF signal processing networks.

Advanced QAT pipelines incorporate the target hardware's operator support matrix directly into the training loop. If a specific layer (e.g., a custom activation) is not supported by the NPU's integer instruction set, it is left in floating-point or replaced with a hardware-friendly approximation.

• Mixed-Precision QAT: Assigns different bit-widths (e.g., INT4 for weights, INT8 for activations) to different layers based on sensitivity analysis.
• Energy Modeling: Integrates a differentiable power model to optimize not just for accuracy, but for the TOPS/Watt efficiency of the final compiled graph.