Module 4 — Quantization-aware training

Module 3 delivered a 2.35 MB int8 model with a 1.3-point accuracy drop, which was tolerable for our leaf classifier. On another model — a fine-grained skin-lesion classifier, say — the same procedure would routinely drop five or ten points, and no product would ship on that. Quantization-aware training exists for exactly those cases: it teaches the network to be a good int8 model while it is still learning, and it is the only technique that reliably brings the drop back to noise.

Why post-training quantization ever loses accuracy

To understand what QAT fixes, look at what post-training quantization does.

At every layer, the runtime maps a float value to the nearest int8 grid point. Between two grid points, all values collapse to the same integer — that is the quantization error. Averaged over a well-behaved distribution, this error is small: about 0.4 percent of the range. But three things concentrate the damage:

1. Layers whose activations have a very wide dynamic range. If a ReLU output goes from 0 to 60, one int8 step is roughly 0.24. Small activations that carried the signal get flattened to zero.
2. Layers followed by a hard non-linearity like sigmoid. A tiny shift in the pre-activation can flip the output side, and the model becomes less confident in a way that compounds through the network.
3. Rare but important patterns: a subtle texture that appears in one percent of images is exactly the kind of pattern a post-quantized model tends to lose, because it barely moved the activation ranges recorded by the representative dataset.

Post-training quantization treats the model as fixed and picks the least bad rounding. QAT lets the model learn around the rounding.

Fake quantization: the core idea

During QAT, we simulate quantization inside the training loop. Weights and activations are rounded to int8 grid points in the forward pass, but the backward pass treats the rounding as if it were the identity — a device called the straight-through estimator. The gradient flows back through as if quantization had not happened, so the weights can still learn; the forward pass sees the quantization error, so the loss reflects the model's real behaviour at deployment.

The result is a training that pushes weights toward values that quantize cleanly. A weight that would land halfway between two int8 grid points gets nudged toward one of them, because that is where the loss is now lower.

The API

TensorFlow Model Optimization Toolkit provides one function:

import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

qat_model = quantize_model(baseline_model)

qat_model.compile(
    optimizer=keras.optimizers.Adam(1e-5),   # very low: this is fine-tuning
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

qat_model.fit(train_ds, validation_data=val_ds, epochs=5)

Four things about this snippet.

First, quantize_model wraps every supported layer with a QuantizeWrapper that inserts fake quantization ops. The model is still a Keras model — fit, evaluate, predict all work — but its forward pass now simulates int8 arithmetic.

Second, the learning rate is very low. QAT is fine-tuning, not training from scratch: the weights are already good, and you only want them to shift enough to recover from quantization noise. 1e-5 is a good starting point; 1e-4 risks unlearning what the base model knew.

Third, five epochs is usually enough. QAT is a corrective pass, not a full re-training. If validation accuracy is still climbing at epoch 10, something is off — likely the learning rate is too low and the model has not yet reacted.

Fourth, you fine-tune on the same dataset you used for the base training. QAT is not the place to introduce new data; the goal is to preserve the behaviour learned earlier, minus the quantization damage.

Converting a QAT model

After training, the conversion is exactly the full-integer procedure from module 3:

converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

tflite_qat_model = converter.convert()

The representative dataset still matters — the fake-quant ops recorded ranges during training, but the converter refines them on real data. The resulting .tflite file is the same size as a post-training int8 model. Only the accuracy has changed.

The trade-off matrix, updated again

Variant	Size	Latency	Top-1 accuracy
Float32 baseline	9.20 MB	88 ms	0.941
Full integer, post-training	2.35 MB	22 ms	0.928
Full integer, QAT	2.35 MB	22 ms	0.939

QAT recovered 1.1 of the 1.3 points we lost. On a classifier where post-training had cost five points, QAT typically recovers four or five. The gap it does not close is genuine — some models are structurally hard to quantize, usually because they rely on activations with tails the int8 range cannot cover.

When QAT is worth the extra cost

QAT is not free. Compared to post-training:

• Training takes a few more hours (five epochs on the leaf dataset).
• The training pipeline has to work — a research checkpoint you were handed is not enough, you need reproducible data loading and the original loss.
• You have to schedule a second training pass whenever the model changes.

Use it when three conditions hold: the post-training drop is more than one point on the metric that matters, the drop is concentrated on important classes rather than uniform, and you have kept the training code in a runnable state. Skip it when post-training already lands under budget — pursuing an extra 0.2 percent is not worth reinstantiating a training run.

Try post-training first, always

Even if you know QAT will be needed, start with the post-training procedure from module 3. It sets a clean baseline, exercises the converter, produces a working .tflite file the mobile team can integrate against, and gives you a measured accuracy drop to justify the QAT effort. QAT is a targeted fix, not a default choice.

What can go wrong with QAT

The loss stays flat. Usually the learning rate is too low (or exceptionally, the model was already at its post-training optimum). Try 3e-5, then 1e-4, but watch validation accuracy: a jump above 1e-4 and the base weights start to unlearn.

One layer's accuracy collapses. The QAT wrapper does not cover every custom layer. Check the model with qat_model.summary() and confirm every trainable layer is wrapped. Custom layers may need tfmot.quantization.keras.QuantizeConfig to describe how their weights should be quantized.

Batch normalisation regresses. BN statistics interact with quantization ranges. On a small dataset, the safe recipe is to freeze BN (layer.trainable = False) before wrapping the model, which mirrors the technique from course 08, module 8.

Key takeaways

• QAT inserts fake quantization into the forward pass and uses the straight-through estimator on the backward pass, so the model learns weights that quantize cleanly.
• Use a very low learning rate (around 1e-5) and a few epochs: QAT is fine-tuning, not full training.
• Convert the QAT model with the same full-integer settings as post-training; the file size is identical, only the accuracy improves — typically recovering most of the post-training drop.
• Use QAT only when post-training loses more than one point, especially when the drop concentrates on important classes; a 0.2-point gain does not justify a new training run.

Next module: pruning and clustering, which attack the file size axis rather than the accuracy axis.