II · THE IDEA · ARTIFICIAL INTELLIGENCE
Blockwise Quantization-Aware Training
▶ Listen · narrated
Full-precision training followed by a hard cast to eight-bit integers can erase hard-won accuracy. Simulating the cast during training, and introducing it gradually by block, softens that loss.
At a glance
- What it is
- Fine-tuning with fake quantisation ops, applied progressively by network block
- Forward pass
- Weights and activations are rounded to low-bit integers, then dequantised
- Backward pass
- Gradients cross the round step via a straight-through estimator
- Usual target
- Eight-bit integer arithmetic for inference hardware
- Contrast
- Post-training quantisation needs no retraining but often loses more accuracy
Think of a musician who must perform on a piano that only has every fourth key. If you hand them that broken piano on concert night, the music falls apart. If you make them rehearse on it for weeks, they learn fingerings that still work. Quantisation-aware training is that rehearsal: during training the network’s numbers are deliberately rounded to the coarse levels they will use on the device, so the model adjusts.
Blockwise means you do not wreck the whole piano at once. You disable keys in one section, let the player adapt, then move to the next section. Full-precision copies of the weights are still updated behind the scenes; only the forward pass pretends the values have already been rounded. Gradients ignore the mathematical fact that rounding is flat, and pass through roughly as if nothing happened — a useful fiction called the straight-through estimator.
The usual goal is eight-bit integers: far less memory and faster arithmetic on many chips, with accuracy close to the original floating-point network when the rehearsal has been long enough and staged carefully.
Blockwise quantisation-aware training inserts fake quantisation nodes on weights and activations inside successive network blocks during fine-tuning. Each fake quant applies an affine map x → round(clamp(x/s + z)) with scale s and zero-point z (z = 0 for symmetric schemes), then dequantises back to floating point so the rest of the graph is unchanged in type. Master weights remain FP32 (or FP16) and receive SGD or Adam updates.
Because round is non-differentiable, backprop uses a straight-through estimator: ∂L/∂x_pre ≈ ∂L/∂x_post through the round, sometimes with a straight-through variant that zeros the gradient outside the clamp range. Scales may be fixed from a calibration pass or learned jointly.
A blockwise schedule enables these nodes only for the active block (or stage), fine-tunes, then advances, keeping earlier blocks quantised and later blocks in full precision until their turn. This reduces the simultaneous distribution shift relative to full-model QAT. At export, fake quant is replaced by true integer ops targeting int8 kernels. Compared with post-training quantisation, QAT adapts weight values to the grid at the cost of training time; blockwise scheduling is a stability heuristic when global QAT diverges or underfits.
Look closer
Fake quantisation is deliberate theatre
During blockwise quantisation-aware training the network still stores and updates full-precision master weights. In the forward pass those weights — and the activations that flow from them — are scaled, rounded to the integer grid, and scaled back. The arithmetic the rest of the block sees is therefore already corrupted by quantisation noise. The master copy is updated from the resulting gradients, so the network learns parameters that remain useful after rounding, not parameters that only work in floating point.
The round step has no true gradient
Rounding is piecewise constant: its derivative is zero almost everywhere. Training would stall if that fact were taken literally. The usual fix is the straight-through estimator, which pretends the round operation is the identity in the backward pass and passes the upstream gradient through unchanged. The approximation is crude, yet it is stable enough that networks can adapt their weight distributions to the discrete grid they will face at inference.
Blocks absorb the shock in stages
Quantising every layer at once can produce a large, sudden drop in accuracy that is hard to recover from in a short fine-tune. A blockwise schedule freezes or leaves in full precision the parts not yet under training, inserts fake quantisation only inside the active block, and fine-tunes until that block settles. The frontier then moves. Earlier blocks have already adapted; later blocks still have full-precision headroom. The same idea appears in layer-wise and stage-wise variants; the shared aim is to keep the loss landscape traversable.
The story
Neural networks are almost always trained in floating-point arithmetic. At deployment, especially on mobile CPUs, DSPs and dedicated accelerators, those same networks are far more efficient when weights and activations live as low-bit integers. The conversion is not free. Mapping a continuous value to a small set of integer levels discards information, and if the network has never experienced that loss, accuracy can fall sharply.
Quantisation-aware training (QAT) attacks the problem by making the loss visible during training. Fake quantisation operations are inserted on weights and activations in the forward pass. A scale (and, for asymmetric schemes, a zero-point) maps a real-valued tensor onto an integer grid; values are rounded; the integers are mapped back to the original scale so that subsequent layers still see floating-point tensors. Mathematically the forward computation already matches what integer inference will do. The stored parameters remain full precision so that small gradient updates can still accumulate.
The backward pass must cross the round operation. Because rounding is flat almost everywhere, its true gradient is zero and would stop learning. Practitioners therefore use a straight-through estimator: in the backward direction the round is treated as if it were the identity, and the gradient of the loss with respect to the quantised value is copied onto the full-precision parameter. The network slowly reshapes its weight distributions so that mass sits where the discrete levels will land cleanly.
Blockwise QAT is a scheduling discipline on top of that mechanism. Instead of enabling fake quantisation across the whole model in one step, the trainer activates it inside one block — a residual stage, a transformer layer group, or another natural partition — fine-tunes until the loss recovers, then advances. Blocks that have already been quantised keep their fake ops; blocks still ahead remain in full precision. The hope is that each local shock is small enough for ordinary optimisation to absorb, and that earlier blocks, once adapted, provide a stable base for later ones.
This sits in contrast to post-training quantisation, which calibrates scales on a small data sample and converts a finished floating-point model without further gradient steps. Post-training methods are cheap and often sufficient at eight bits for many vision models, but they leave the weights themselves unadapted to the grid. When the gap is too large — aggressive bit-widths, sensitive architectures, or domains with little calibration data — QAT, and especially a staged blockwise schedule, is the route that recovers accuracy at the cost of extra training time.
The white-paper literature on neural network quantisation and the empirical studies of integer inference both treat eight-bit weights and activations as the practical sweet spot for general-purpose integer kernels: wide hardware support, substantial memory and bandwidth savings relative to float32, and, with care, accuracy close to the full-precision baseline. Blockwise QAT is one of the training-time tools used to reach that regime when a single global quantisation pass is not enough.
Why it mattered then
As mobile and edge accelerators standardised on integer matrix units, the industry needed networks that were not merely compressed after the fact but trained to tolerate that compression. Post-training conversion was fast, yet for many production models the accuracy drop was unacceptable. Quantisation-aware training closed the gap by letting the optimiser see the same rounding the device would apply. Doing the conversion block by block made the procedure workable on deep residual and multi-stage networks, where a single global switch often left the loss in a poor basin. The method mattered because it turned an inference-only constraint into something the training loop could negotiate.
Why it matters now
Integer eight-bit inference remains the default efficient path on phones, NPUs and many server accelerators. Open-weight models are routinely quantised for local use, and the same tension persists: aggressive post-training schemes are convenient, yet some architectures and tasks still need the network to practise under quantisation noise. Blockwise and layer-wise QAT schedules remain a practical pattern whenever a full-model fake-quant pass is unstable or when memory forces the trainer to touch only part of the graph at a time. Anyone shipping a model that must run in low precision still faces the choice between a cheap cast and a slower, staged adaptation.
The surprising detail
The straight-through estimator is almost embarrassingly wrong as a gradient: it claims that a function whose slope is zero has slope one. Yet without that lie, QAT would not train at all. The literature accepts the mismatch as an engineering necessity rather than a theoretically justified derivative, and the blockwise schedule can be read as a second pragmatic fix — reducing how much incorrect gradient signal the optimiser must swallow at once.
What is disputed
How finely to slice blocks, how long to fine-tune each stage, and whether to quantise weights before activations are practical choices rather than settled theory. Results vary by architecture and bit-width; the sources describe principles and empirical patterns, not a single mandatory schedule.
Remember this
Blockwise QAT teaches the network the integer grid in stages, so rounding becomes a condition of learning rather than a shock after the fact.
Test yourself
A team fine-tunes with fake quantisation on every layer from the first step and sees the loss spike and never fully recover. They switch to enabling fake quantisation one residual block at a time. Mechanically, what two things have changed about the optimisation problem?
First, only a fraction of the forward path injects quantisation noise at any moment, so the immediate distortion of activations and of the loss surface is smaller. Second, blocks that are not yet quantised still provide full-precision capacity, giving the optimiser room to compensate while each active block adapts; previously adapted blocks keep their fake ops and act as a stable, already-discretised base. The straight-through estimator is unchanged — what changed is how much of the network is asked to absorb incorrect gradients and rounding error in one go.
Go deeper
- [2106.08295] A White Paper on Neural Network Quantization · arxiv.org
- [2004.09602] Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.