Quantization-Aware Training (QAT): Simulated Quantization

Post-training quantization (PTQ) gets most INT8 CNNs within a point or two of FP32 accuracy, but for aggressively compressed models — sub-1MB keyword spotters, INT4 weights, or networks with sensitive layers like depthwise convolutions and squeeze-excite blocks — PTQ can lose 3–10 accuracy points. Quantization-aware training (QAT) recovers most of that gap by letting the network learn weights and activations that are robust to quantization noise, instead of quantizing a network that was never told it would be quantized.

Why PTQ Falls Short

PTQ calibrates scale/zero-point parameters from a small calibration set after training is complete. The weights themselves were optimized purely for FP32 accuracy — they may sit near quantization decision boundaries, have long tails that blow up the dynamic range, or rely on precision the integer grid can't represent. The rounding error introduced by quantization is then just noise injected once, with no chance for the network to compensate.

QAT instead simulates quantization during training, so gradient descent actively pushes weights and activations toward values that survive rounding with minimal loss. It's the difference between building a bridge and then discovering the truck weight limit, versus designing the bridge knowing the trucks in advance.

The Fake-Quantization Node

QAT doesn't run actual low-bit integer arithmetic during training — training still happens in FP32 on GPU. Instead, a fake-quantize (simulated quantization) node is inserted after each weight tensor and activation tensor in the forward pass:

x_fq = dequant(quant(x, scale, zero_point))
     = scale · ( clamp( round(x/scale) + zp, qmin, qmax ) − zp )

This rounds and clamps x to the target integer grid, then immediately converts back to floating point. The forward pass therefore "feels" the same rounding error the deployed INT8 kernel will produce, while the backward pass and weight updates still run in full floating-point precision.

Straight-Through Estimator (STE): round() has zero gradient everywhere except at integer boundaries (undefined derivative), which would kill backpropagation. QAT sidesteps this with the STE: on the backward pass, the gradient of round(x) is approximated as 1 inside the clamp range and 0 outside it:

∂x_fq/∂x ≈ 1   if qmin ≤ round(x/scale)+zp ≤ qmax
         ≈ 0   otherwise (saturated)

This is a biased approximation — it's not the true gradient of a step function — but empirically it works well because rounding error is small relative to the gradient signal, and it lets the optimizer treat quantization as a differentiable, if noisy, operation.

What Gets Learned

Two categories of parameters adapt during QAT:

• Weights shift slightly off their FP32-optimal values to land closer to grid points, reducing rounding error where it matters most for the loss.
• Quantization parameters (scale, zero-point) for activations are often made learnable themselves — this is the basis of methods like LSQ (Learned Step Size Quantization) — so the clipping range adapts to the true post-training activation distribution rather than a static calibration snapshot.
• Batch normalization statistics are typically frozen or folded early in QAT ("BN folding"), because BN's scale/shift would otherwise interact badly with the discrete activation grid and destabilize training.

Symmetric vs. Asymmetric, Per-Tensor vs. Per-Channel

QAT doesn't change these design choices from PTQ — it just makes the network adapt to whichever is chosen:

Per-channel weight quantization combined with QAT is the most common production recipe for CNN backbones (MobileNet, YOLO-Nano, person-detection nets) deployed on CMSIS-NN or ARM Ethos-U.

Worked Example: Simulating INT8 on a Single Weight

Take a weight w = 0.734, with a per-channel symmetric scale derived from that channel's max absolute weight |w|_max = 1.5 over an INT8 range [-127, 127]:

scale = |w|_max / 127 = 1.5 / 127 = 0.011811

Forward pass fake-quantization:

q = round(w / scale) = round(0.734 / 0.011811) = round(62.15) = 62
w_fq = q · scale = 62 · 0.011811 = 0.73228

Quantization error introduced this step: Δw = w − w_fq = 0.734 − 0.73228 = 0.00172.

During backprop, suppose the loss gradient with respect to this weight is ∂L/∂w_fq = -0.05. Under STE (q=62 is within [-127,127], not saturated), the gradient passes through unchanged:

∂L/∂w ≈ ∂L/∂w_fq · 1 = -0.05

The optimizer updates w using this gradient as normal (e.g., w_new = w − η·∂L/∂w), nudging w in FP32 space. Over many iterations, the network learns to place weights where Δw costs the least loss — not necessarily where FP32 training alone would have put them.

Check: at deployment, the real INT8 kernel computes with q = 62 directly (no floating point at all), producing exactly w_fq = 0.73228 in dequantized terms — matching what the network trained against. This is the core guarantee QAT provides: train-time and deploy-time numerics match, unlike PTQ where the network never saw the rounding during optimization.

QAT Workflow on a Cortex-M Target

1. Start from a pretrained FP32 model (QAT fine-tunes; it rarely trains from scratch).
2. Fold batch norm into preceding conv/FC weights, then insert fake-quant nodes on weights and activations (TensorFlow's tfmot.quantization.keras, PyTorch's torch.ao.quantization QAT APIs).
3. Fine-tune for a fraction of original training epochs (typically 5–15% of original schedule) at a reduced learning rate — full retraining is rarely necessary since the model only needs to adapt to rounding, not relearn features.
4. Convert to the true integer graph (TFLite Micro converter, ONNX QOperator/QDQ export) — fake-quant nodes are replaced with real INT8 ops; scales/zero-points become static kernel parameters.
5. Validate on-target: run the converted model through CMSIS-NN or the Ethos-U55 driver and confirm accuracy matches the simulated (fake-quantized) accuracy from step 3 — a large gap indicates an unsupported op, mismatched rounding mode, or BN-folding error.

When QAT Is Worth the Cost

QAT requires the original training pipeline, labeled data, and GPU time — much more expensive than PTQ's few hundred calibration images. It's justified when:

• PTQ accuracy loss exceeds your budget (commonly >1–2 points for classification, tighter for detection mAP).
• Target precision is below INT8 (INT4/INT2, binary/ternary networks) — PTQ essentially never works well below INT8; QAT is close to mandatory.
• Architecture has known quantization-fragile components — depthwise separable convolutions, squeeze-and-excite gates, and attention softmax layers all quantize poorly under PTQ.
• The model is deployed at scale, where the one-time training cost amortizes across many units (e.g., a wake-word model shipped in millions of devices).

Key Takeaways

• QAT inserts fake-quantization nodes (round → clamp → dequant) into the FP32 training graph so the network experiences rounding error during training, not just at deployment.
• The Straight-Through Estimator (STE) approximates the gradient of the non-differentiable round() as 1 inside the clip range, 0 outside — a biased but practically effective backprop rule.
• Weights, and often learnable activation scale/zero-point (as in LSQ), adapt jointly with BN folded/frozen early to avoid instability.
• QAT typically fine-tunes a pretrained FP32 model for a short schedule rather than training from scratch, converting to true integer ops only at export time.
• QAT is worth its extra training cost when PTQ accuracy loss is unacceptable, precision drops below INT8, the architecture has quantization-fragile layers, or the model ships at high volume.