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 by it on the device.

Fake quantization snaps each weight to its nearest int8 grid point on every forward pass, so the network trains against the rounding error it will ship with while gradients pass straight through

Why Does This Exist?

Picture a shelf-scanning app that has to run its product classifier fully offline, on the mid-range phones warehouse staff already carry. The trained model lives in 32-bit floating point, four bytes per weight. Rounding every weight to an 8-bit integer, one byte each, cuts the model to a quarter of its size and runs on phone hardware built for integer math. That rounding step, applied after training with no retraining, is called post-training quantization, and it is the obvious move.

Applied here, it stings. The network never knew its weights would land on a coarse grid, so small floating-point differences it depended on collapse onto the same rounded value, and the distinctions they carried disappear. Accuracy drops, and nothing was done wrong. The rounding was simply a surprise.

Quantization-aware training removes the surprise by simulating the rounding during training itself. Every forward pass feels the error the deployed model will feel, so the optimizer settles on weights that are good even after rounding.

The one idea you need first. Quantization means replacing a continuous value with the nearest point on a fixed grid, like rounding dollars to whole five-dollar bills. An int8 grid holds 256 evenly spaced values, so anything learned must survive that coarseness. The precision vocabulary behind fp32, fp16 and int8 is covered in mixed-precision training. This page is about training through the rounding; knowledge distillation teaches a small model to imitate a big one instead, which is the other way to shrink a model.

Think of It Like This

Rehearsing on a stage taped to the real size

A theater troupe rehearses in a generous hall, then opens in a much smaller venue. Blocking that worked in the hall does not fit the real stage, and the cast discovers that in front of an audience. The fix is to tape the real stage dimensions onto the rehearsal floor from the first week, so every move locked in already respects the true constraint. Opening night changes nothing, because the show never assumed space that was never there.

The full-precision weight is the rehearsal, the int8 grid is the taped outline, and opening night is deployment. The mapping holds point for point: rehearsing inside the tape is the forward pass with fake quantization, and locked-in blocking is the converged weight.

It stops here: actors adjust on purpose, while weights adjust because a gradient pushes them. And nothing on a stage corresponds to the straight-through estimator, the backward-pass trick below, which is pure math rather than rehearsal.

How It Actually Works

Fake quantization in the forward pass

On every forward pass, each weight (and usually each activation) is rounded to the nearest point of the low-precision grid, then used in that rounded form for the rest of the computation. The storage stays full precision, but the network computes with rounded values. This is fake quantization: simulate the deployment rounding now, so the loss being optimized is the loss the shipped model will actually get.

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

In plain words: $x$ is the full-precision value, $x_{\min}$ and $x_{\max}$ are the range being squeezed into the grid, and $s$ is the scale, the spacing between neighbouring grid points. The fraction inside measures how many steps $x$ sits above the range floor, round snaps it to a whole step, and multiplying back by $s$ lands on the grid point. clip holds outliers at the range edges.

The straight-through estimator

There is a catch. The round() in that formula is a staircase, flat between steps, so its true gradient is zero almost everywhere. Backpropagating honestly would kill nearly all gradient flow and training would stall. The fix is the straight-through estimator (STE): on the backward pass, pretend the rounding step was the identity function and let the gradient pass through unchanged.

That is openly inconsistent. Forward and backward compute with two different functions. It works because the full-precision weight underneath still receives a directionally correct signal and keeps updating, while only the forward computation sees the rounding. Over thousands of steps the optimizer drifts toward weights whose rounded versions perform well, which is the only thing deployment cares about.

Worked example

Take one weight $x = 0.37$ , a range of $[-1, 1]$ , and a tiny 2-bit grid with 4 levels. The scale is $s = (1 - (-1)) / (4 - 1) = 2/3 \approx 0.6667$ . Steps above the floor: $(0.37 + 1) / 0.6667 = 2.055$ , which rounds to $2$ . Back to a value: $2 \times 0.6667 - 1 = 0.3333$ . The four grid points are $-1$ , $-0.3333$ , $0.3333$ and $1$ , so $0.37$ snaps to $0.3333$ with an error of $0.0367$ . That error is what the forward pass feels on every step, and what the optimizer learns to route around.

Now compare with real int8, 256 levels over the same range. The step is $2/255 \approx 0.0078$ , so the worst rounding error on any weight is half a step, about $0.0039$ . The 2-bit grid in this example inflicts roughly ten times the error of the coarsest step int8 ever takes, which is why the toy grid makes the mechanism visible while int8 quietly preserves accuracy.

Watch Out For

Simulating rounding on weights but not activations

The fake-quantization step is easy to add to stored weights and easy to forget on activations, and a network trained that way has only practiced against half the error it will face, since deployment hardware quantizes both. The symptom is specific: validation accuracy during training looks fine, then accuracy drops on the device with no other change. Confirm the simulation covers activation tensors at every layer boundary the deployment target quantizes, not just the weights.

Batch-norm statistics that no longer match deployment

Batch normalization layers normalize with batch statistics during training but with frozen statistics once deployed, and fake quantization interacts with both. If the normalization statistics shift underneath the simulated rounding, on-device accuracy lands below what training evaluation promised. The symptom is a consistent gap between the last training evaluation and the first device run. Fold normalization into the neighbouring weights before the quantization-aware phase, or freeze its statistics, so training and deployment normalize identically.

The Quick Version

Post-training quantization rounds a trained network afterward, and the network never practiced against that rounding, so accuracy can drop hard.
Quantization-aware training fake-quantizes on every forward pass, so the optimizer finds weights that are good even after rounding.
The straight-through estimator passes gradients through the rounding step as if it were the identity, because the true gradient of round() is zero almost everywhere.
A 2-bit grid snaps 0.37 to 0.3333 with error 0.0367, while int8 over the same range never errs by more than about 0.0039 per weight.
Simulate rounding on activations too, and keep normalization statistics identical between training and deployment.