Quantisation-Aware Training and the Straight-Through Estimator

Post-training quantisation takes a finished model and finds the best possible integer approximation of it. Quantisation-aware training does something different and stronger: it trains a model that is good after being quantised, letting the weights move to places where the rounding error does less damage. At 4 bits and below the gap between the two approaches is often the difference between a usable model and an unusable one.

The gradient problem

QAT inserts fake quantisation into the forward pass: quantise, immediately dequantise, and continue in floating point. The network computes with values restricted to the integer grid, so it experiences the quantisation error during training and can adapt to it.

The backward pass has a problem. The round function is piecewise constant, so its derivative is zero almost everywhere and undefined at the steps. Backpropagating honestly through it gives zero gradient to every weight and nothing trains.

The straight-through estimator resolves this by lying. In the backward pass, treat the quantiser as the identity:

Gradients pass through unchanged for values inside the representable range and are zeroed outside it, so clipped values stop receiving updates that would push them further out. It is not the true gradient of the forward computation. It is the gradient of a surrogate, and it works because the quantiser is close to the identity on average and the error it introduces is roughly unbiased noise.

What QAT actually changes

Two things, and they are worth separating.

The weights move. Given the chance, training pulls weights toward positions on the grid where the loss is flat, so the same nominal rounding error costs less. This is why QAT can beat post-training quantisation even when both use identical grids.

The scales are learned. Treating \(s\) and \(z\) as trainable parameters, as in learned step size quantisation, lets the network choose its own clipping thresholds per layer by gradient descent rather than by a calibration heuristic. This usually accounts for a large share of QAT's advantage and it is the part most easily bolted onto an existing recipe.

Why it is used less than it should be

QAT requires the training pipeline, the training data, and a meaningful fraction of the original training compute. For an open-weights model that a downstream team wants to compress, none of those are available. That is the entire reason post-training quantisation dominates practice: it needs a few hundred calibration samples and a few GPU-minutes, and for weight-only 4-bit quantisation of a large model it is usually good enough.

QAT keeps its place where the bit width is aggressive, at 2 or 3 bits, where post-training methods degrade sharply; where the target is an edge accelerator with fixed integer arithmetic and no fallback path; and where the model is small enough that the accuracy loss from post-training methods is proportionally larger, which it is, since smaller models have less redundancy to spare.

When it breaks

Batch normalisation folding must happen before, not after. If BN is folded into the preceding convolution after QAT, the folded weights have a different distribution from the ones that were trained, and the calibrated scales no longer fit. Folding first, then training, is the correct order, and getting it backwards produces a model that is fine in the training harness and broken on export.

The STE's bias is real at low bit width. The surrogate gradient is a good approximation when the quantisation error is small relative to the weight magnitude. At 2 bits it is not, and training becomes unstable in a way that manifests as loss spikes and oscillating weights near decision boundaries. Softer relaxations that anneal toward a hard quantiser exist for this regime.

Weights oscillate across grid boundaries. A weight sitting near the midpoint between two levels can flip between them on successive steps, which injects gradient noise into every downstream computation and slows convergence. Oscillation dampening or freezing weights that have flipped repeatedly is a standard mitigation and is rarely mentioned in method descriptions.

QAT on a fine-tuned model can undo the fine-tuning. Running QAT with the original pretraining objective on a model that was aligned or instruction-tuned pulls it back toward the base distribution. The QAT phase has to use the same data and objective as the final training stage, which teams reliably discover after the first attempt.