Quantization-Aware Training (QAT)

Quantization-aware training (QAT) is a training technique that prepares a model for low-precision inference by simulating quantization during training. Models normally train and store weights as 32-bit floating-point numbers (FP32), and model quantization converts them to formats such as 8-bit integers (INT8) for deployment, which makes weights 4 times smaller and lets CPUs, GPUs, and neural processing units (NPUs) use fast integer arithmetic. The conversion rounds every value and clips outliers, which costs accuracy. QAT inserts that rounding into the forward pass while the model is still learning, so gradient descent finds weights that keep working after quantization. It is usually run as a short fine-tuning stage on an already trained model. Google's integer-arithmetic-only inference scheme for mobile CPUs paired INT8 inference with this kind of training procedure, and QAT has since become a standard tool for edge AI deployment.

How Quantization-Aware Training Works#

QAT adds fake quantization operations, also called quantize-dequantize (Q/DQ) pairs, to a model's weights and activations. Each one maps a float to the nearest INT8 level using a scale and zero point, then immediately maps it back to float, so the network keeps computing in floating point but sees the same rounded values an INT8 engine would use. The scales come from calibration, which records the typical range of each tensor on a few batches of data, and some methods also learn them during training.

Rounding has a gradient of zero almost everywhere, which would stop learning. QAT uses the straight-through estimator (STE), which treats the rounding step as the identity function during backpropagation, so gradients pass through to the underlying FP32 weights:

import torch

# FP32 weights that training keeps and updates
w = torch.tensor([0.3137, -0.7214, 0.0521, 1.1800], requires_grad=True)

# Symmetric INT8 fake quantization: snap each weight to one of 255 levels, then map back to float
scale = w.detach().abs().max().item() / 127
w_q = torch.fake_quantize_per_tensor_affine(w, scale, 0, -127, 127)
print(f"Fake-quantized: {[round(v, 4) for v in w_q.tolist()]}")

# The forward pass sees the rounded weights; the straight-through estimator still delivers gradients to w
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
loss = (w_q * x).sum()
loss.backward()
print(f"Gradient on FP32 weights: {w.grad.tolist()}")

The forward pass uses the rounded weights 0.3159, −0.7247, 0.0557, and 1.18, each within half a quantization step (0.0046) of its FP32 value, yet the gradient reaching the FP32 weights is exactly the input, 1.0, 2.0, 3.0, and 4.0, as if no rounding had happened. Over many steps the optimizer moves weights to positions where the rounding error does the least harm, and the trained model exports to INT8 with its learned scales.

QAT vs. Post-Training Quantization#

Post-training quantization (PTQ) converts a finished FP32 model directly, using a few hundred calibration images to set the ranges. It takes minutes and needs no training pipeline, but the weights never get a chance to adapt to rounding, so sensitive models can lose noticeable accuracy. QAT costs a fine-tuning run with the full training data, and in return recovers most of that loss. Teams usually try PTQ first and switch to QAT only when the INT8 accuracy drop is too large:

Ultralytics YOLO runs QAT when training receives quantize=8: it calibrates INT8 fake quantization through NVIDIA Model Optimizer, fine-tunes with it in the loop, and saves ranges that ONNX and TensorRT exports write as Q/DQ nodes without separate calibration data. In TensorRT engines on COCO val2017, PTQ dropped the medium YOLO26 model from 0.527 to 0.470 mAP50-95 while QAT held it at 0.514, and the nano model lost only 0.010 with PTQ, so QAT added nothing there.

from ultralytics import YOLO

# Start from a trained FP32 YOLO26 checkpoint
model = YOLO("yolo26n.pt")

# Fine-tune with INT8 fake quantization in the loop, using a small learning rate
model.train(data="coco8.yaml", epochs=5, quantize=8, optimizer="AdamW", lr0=0.00001, cos_lr=True)

# Export to ONNX; the learned INT8 ranges become Q/DQ nodes, so no calibration data is needed
model.export(format="onnx", quantize=8)

Where QAT Is Used#

• Edge vision: INT8 detectors on mobile phones, smart cameras, drones, and NPUs, where integer-only accelerators cannot run floating-point models at all.
• Inference servers: GPU engines such as TensorRT run INT8 Tensor Core kernels for higher throughput, and QAT keeps accuracy close to FP32.
• Language models: QAT now reaches below 8 bits. Google released Gemma 3 QAT checkpoints in 4-bit Q4_0 format that preserve quality close to bfloat16 at a fraction of the memory.

Major frameworks ship QAT tooling, including PyTorch's torchao, TensorFlow's Model Optimization Toolkit, and Intel's NNCF for OpenVINO.

QAT vs. Related Terms#

• Model quantization: The umbrella term for lowering numeric precision. PTQ and QAT are its two main workflows.
• Mixed precision training: Trains in FP16 or BF16 to save memory and time, without preparing the model for INT8 inference.
• Half precision: FP16 inference, which most models tolerate without retraining, so QAT is rarely needed for it.
• Model pruning: Removes weights or channels instead of lowering their precision, and the two compressions stack.
• Knowledge distillation: Trains a smaller architecture to imitate a larger one. It is often combined with QAT, with the full-precision model acting as teacher for its own quantized copy.