Core ML

Quantization-Aware Training

Round weights to low precision during training itself, not after, so gradients push the network toward values that survive rounding instead of being surprised.

The forward pass rounds weights to a low-precision grid so the network trains against the error it will actually ship with, while the backward pass pretends that rounding step was the identity function to keep gradients flowing

Why Does This Exist?

A retailer's shelf-scanning app needs to run a product classifier fully offline, on the same mid-range phones its warehouse staff already carry. The trained model runs comfortably in 32-bit floating point on a training server. Shrinking every weight down to 8-bit integers, a standard move called post-training quantization, would cut the model's memory footprint by 4x and speed up inference on phone hardware built for integer math — and it's usually applied after training finishes, with no retraining involved.

Applied here, it doesn't go well. Accuracy drops from 94% to 81%. Nothing about the quantization step was done wrong; the network simply was never trained with any awareness that its weights would eventually be rounded onto a coarse int8 grid, so a handful of small floating-point weight differences that mattered a great deal to the trained model land on the same rounded value, and the distinctions they carried disappear.

Quantization-aware training fixes this by simulating the rounding during training, not after. The network then has the chance to learn weights that are already robust to the exact rounding it will face at deployment.

Think of It Like This

Rehearsing on the actual stage, not a full-size mockup

A theater troupe rehearses a play in a full rehearsal hall with generous space to move — and the actual venue is a much smaller stage, with tighter blocking and less room between actors. Rehearsing entirely in the big hall and only encountering the small stage on opening night means every actor discovers the constraints for the first time in front of an audience, and the choreography that worked in the hall doesn't fit.

Rehearsing on a stage taped out to the real venue's dimensions, from early rehearsals onward, means the blocking that gets locked in already respects the real constraint. Nothing about the choreography needs to change on opening night, because it was never built assuming space that wasn't going to be there.

How It Actually Works

Fake quantization in the forward pass

During each forward pass in training, every weight (and often each activation) is rounded to the nearest value on a low-precision grid — say, one of 256 evenly spaced values for int8 — then immediately used in that rounded form for the rest of the forward computation. This simulated rounding is called fake quantization: the value is stored and gradients are still computed in full floating-point precision, but the network experiences the same rounding error at every step of training that it will experience at deployment.

q(x) = \text{clip}\!\left(\text{round}\!\left(\frac{x - x_{\min}}{s}\right) \cdot s + x_{\min},\; x_{\min},\; x_{\max}\right)

$x$ is the original full-precision value, $s$ is the spacing between grid points (the scale), and $x_{\min}$ , $x_{\max}$ are the range being quantized into. The forward pass uses $q(x)$ everywhere $x$ would normally appear — the loss the network optimizes is the loss it would actually get at low precision, not an optimistic full-precision estimate of it.

The straight-through estimator

The round() function inside $q(x)$ has a gradient of exactly zero almost everywhere — it's a staircase, flat between steps — so backpropagating through it literally would stop nearly all gradient flow, and training would grind to a halt. Quantization-aware training sidesteps this with the straight-through estimator (STE): during the backward pass, pretend $q(x)$ was the identity function, and let the gradient pass through unchanged, as if no rounding had happened at all.

This is mathematically inconsistent — the forward and backward passes are computing derivatives of two different functions — and it works anyway. The full-precision weight underneath still receives a real gradient signal and still updates normally; only the forward computation sees the rounding. Over many steps, the optimizer settles on full-precision weights whose rounded versions perform well, which is the only thing that actually matters at deployment.

Why this beats quantizing after the fact

Post-training quantization asks a question the network was never trained to answer: "how much does rounding hurt a value you didn't know would be rounded?" Quantization-aware training changes the question to "find a value that is good even after being rounded" — and because the rounding happens on every single forward pass throughout training, the optimizer has thousands of chances to route around wherever quantization error would otherwise land hardest.

Show Me the Code

Fake-quantizing a batch of weights at three bit-widths, and measuring how much error each introduces.

1import numpy as np2
3
4def fake_quantize(x: np.ndarray, num_levels: int, x_min: float, x_max: float) -> np.ndarray:5    """Round to the nearest of num_levels evenly spaced values, then clamp."""6    scale = (x_max - x_min) / (num_levels - 1)7    q = np.round((x - x_min) / scale) * scale + x_min8    return np.clip(q, x_min, x_max)9
10
11rng = np.random.default_rng(2)12w = np.clip(rng.normal(0, 1, size=2000), -3, 3)13for bits in (8, 4, 2):14    q = fake_quantize(w, 2 ** bits, -3, 3)15    mse = float(np.mean((w - q) ** 2))16    print(f"{bits}-bit ({2 ** bits} levels): quantization MSE = {mse:.6f}")17# -> 8-bit (256 levels): quantization MSE = 0.00004518# -> 4-bit (16 levels): quantization MSE = 0.01329219# -> 2-bit (4 levels): quantization MSE = 0.341346

Going from 8 bits to 4 bits multiplies the quantization error nearly 300-fold, and 2 bits multiplies it further still — this is the error a network trained without quantization awareness simply absorbs as a surprise at deployment, and the error quantization-aware training gives the optimizer thousands of chances to train around.

Watch Out For

Forgetting to quantize activations, not just weights

Simulating rounding on the weights while leaving activations in full precision trains a network that's only robust to half of the error it will actually face at deployment, since inference hardware typically quantizes both. Confirm the fake-quantization step covers activation tensors at every layer boundary the real deployment target will quantize, not just the stored weights.

Assuming the straight-through estimator's inconsistency causes instability

It's reasonable to expect that backpropagating through a function the forward pass didn't actually use would make training diverge, but in practice the STE is remarkably stable — the mismatch between forward and backward functions rarely derails convergence, because the underlying full-precision weight is still receiving a directionally correct gradient. Treat unexpected instability during quantization-aware training as a learning-rate or initialization issue first; it's rarely the STE itself.

The Quick Version

Post-training quantization rounds a trained network's weights afterward and can cost significant accuracy, since the network never learned to be robust to that rounding.
Quantization-aware training simulates rounding — fake quantization — during every forward pass in training, so the optimizer finds weights that are good even after rounding.
The straight-through estimator lets gradients pass through the rounding step as if it were the identity function, since round()'s real gradient is zero almost everywhere.
The technique is mathematically inconsistent between forward and backward passes, and it works well in practice regardless.