II · THE IDEA · ARTIFICIAL INTELLIGENCE
Quantization with Learned Rounding
▶ Listen · narrated
When weights live in fewer bits, most damage is done when real values are forced onto a grid. If that moment is fixed outside training, the model never braces for it.
At a glance
- Core idea
- Make rounding decisions influenceable during training
- Obstacle
- Hard rounding blocks gradient flow
- Aim
- Adapt weights to low-precision arithmetic
- Payoff
- Less accuracy loss than naïve post-hoc rounding
Think of each weight as a precise dial reading that must be written on a coarse ruler with only a few marks. Ordinary rounding just picks the nearest mark at the end. If you only do that after training, the dial was tuned as if every tiny mark existed, then suddenly forced onto a blunt ruler.
Learned rounding lets training notice the ruler. While the network is still learning, borderline readings can be nudged toward the side of a mark that hurts the answers less. The final run still uses the blunt ruler; the difference is that the dials were practised with that bluntness in mind, so fewer important settings sit awkwardly between marks.
Hard quantisation q(w) = Δ · round(w/Δ) (or an asymmetric affine variant) is piecewise constant, so ∂q/∂w is zero almost everywhere. Post-training quantisation therefore applies a non-differentiable map to a minimiser that never saw it. Straight-through estimators restore a backward signal by substituting ∂q/∂w ≈ 1 (or a clipped variant) while keeping hard q in the forward pass; rounding thresholds themselves remain non-learnable.
Learned rounding introduces a trainable influence on the discrete decision: for example a per-weight offset or rounding probability, a stochastic quantiser whose parameters enter the loss, or a continuation that softens the step and anneals toward hard rounding. Gradients then flow not only into w but into whatever parameterises bin assignment, so mass can move off unstable bin edges. At convergence the forward pass should implement the same grid used at inference; differentiability is a training device.
Limitations remain. If the training relaxation never agrees with the deployment quantiser, a train–serve gap persists. Very low bit-widths can still dominate regardless of rounding policy. Layer sensitivity is uneven, so uniform grids waste error budget. None of this specifies an optimiser, a bit-width, or a guaranteed recovery margin; those are empirical and architecture-dependent.
Look closer
The cut that training cannot see
Standard rounding maps a continuous value onto the nearest representable level and stops. That map is piecewise constant: almost everywhere its derivative is zero, and at the thresholds it is undefined. Gradient-based training therefore receives no useful signal about how a small change in a weight would have changed which bin it fell into. The quantiser sits outside the learning loop.
Softening the decision
Learned rounding keeps a differentiable path—through a relaxation, a noise model, or a parameter that steers which way borderline values tip—so that the loss can pull on rounding behaviour as well as on the underlying weights. The model is no longer only fitting a full-precision surrogate; it is rehearsing the discrete arithmetic it will face at inference.
Adaptation, not just compression
Once rounding participates in training, weights can drift away from values that sit awkwardly on bin edges and settle where quantisation error hurts the loss less. The claim is modest: degradation can be kept small relative to freezing the network and rounding afterwards. How small depends on bit-width, architecture, and the particular relaxation—none of which is fixed by the idea alone.
The story
Quantisation for neural networks is easy to describe and awkward to train through. A full-precision weight is replaced by one of a finite set of levels so that multiplication and storage cost less. The replacement step is rounding, or a close cousin of rounding. At inference that step is harmless: the network is frozen, the grid is fixed, and one simply evaluates. During training it is a different matter. Learning needs gradients. Rounding, taken literally, erases them.
The usual practical responses sit at two extremes. One is to ignore the problem until the end: train in full precision, then round. The network has never experienced the grid it must live on, so accuracy often drops, sometimes sharply at aggressive bit-widths. The other is to insert a straight-through estimator or similar surrogate so that the forward pass rounds while the backward pass pretends the quantiser was the identity. That restores a gradient path, but the path is a fiction: the true rounding decisions still do not depend on learnable parameters in a way the optimiser can reshape.
Learned rounding occupies a third position. Instead of treating the assignment of each value to a grid point as an immutable rule applied after the fact, it treats aspects of that assignment as things training may influence. The continuous weights remain the primary parameters, but an additional differentiable mechanism—however implemented—lets the loss favour configurations that round cleanly. Borderline weights can be nudged toward the side of a threshold that harms the objective less. The network is, in effect, practising under the arithmetic it will use later.
Nothing in this requires claiming a single canonical algorithm. The shared editorial move is to make the rounding decision participate in optimisation rather than merely conclude it. Some approaches soften the step with a temperature or noise schedule that anneals toward hard rounding. Others introduce explicit parameters that bias each weight’s rounding direction and train those parameters jointly with the model. In all such cases the forward computation at the end of training should match the discrete grid used at deployment; the differentiability is a training scaffold, not a permanent change to inference arithmetic.
The accuracy argument is comparative, not absolute. Full precision remains the gentler regime. Post-training quantisation without adaptation remains the cheaper engineering path when it happens to be good enough. Learned rounding is motivated where those two facts collide: the budget demands low precision, and naïve rounding costs more accuracy than the application can spare. By letting the model feel the grid while it can still move, the method aims to spend capacity on representations that survive the grid, rather than on ones that only looked good before the cut.
Why it mattered then
As models grew and deployment moved toward tighter memory and compute envelopes, reducing numeric precision stopped being a minor implementation detail and became a central efficiency lever. The snag was familiar from any discrete decision inside a continuous optimiser: the operation you want at run time is exactly the operation that blocks learning. Methods that left rounding entirely outside training repeatedly showed the same pattern—acceptable at mild compression, brittle when the grid grew coarse. Making the rounding step itself something training could adapt to was a direct response to that brittleness. It reframed quantisation from a post-hoc injury the model had to tolerate into a constraint the model could rehearse against while weights were still plastic.
Why it matters now
Open-weight models are routinely run on hardware where memory bandwidth and resident footprint dominate wall-clock cost. Four-bit and similar regimes are no longer exotic; they are how many people actually serve large networks locally. In that setting, the difference between a model that was merely rounded and a model that was trained to live on the grid is practical rather than academic: it shows up as usable quality at a bit-width that otherwise forces a larger machine or a smaller network. The same logic applies whenever a fixed quantiser is applied across layers and workloads that were never seen during a one-shot calibration pass. Learned rounding is one family of answers to a standing question: if inference will be discrete and low-precision, when should the model find that out?
The surprising detail
The intellectually odd part is not that rounding hurts—that is obvious—but that the hurt is often concentrated on values near decision boundaries. A weight sitting comfortably inside a bin can move a little and still quantise to the same level; a weight balanced on a threshold can flip bins under a tiny perturbation and change the forward pass discontinuously. Learned rounding spends effort on those unstable edges. In that sense the method is less about inventing new numeric formats and more about teaching the optimiser to stop leaving mass where the grid is least forgiving.
What is disputed
Public materials supplied for this lesson do not fix a single algorithm, bit-width table, or measured accuracy delta. Effects vary with architecture, grid coarseness, and how closely the training relaxation matches inference rounding; treat comparative claims as directional, not universal.
Remember this
If inference will round, training should feel the grid—otherwise the model adapts to a world it will never run in.
Test yourself
A team trains fully in floating point, then rounds weights to a coarse grid at deploy time and sees a sharp accuracy drop. They propose adding a differentiable rounding path during a fine-tuning stage so the loss can influence borderline decisions. What failure mode remains if the fine-tuning still uses a soft relaxation that never matches the hard grid used at inference?
The model may adapt to the soft surrogate rather than to true discrete rounding. At deployment the hard grid reappears, and any mismatch between the relaxed forward pass and the real quantiser becomes a new source of error—the same train–serve gap, only subtler. The relaxation has to anneal to, or otherwise agree with, the inference quantiser, or the learned adaptation is partly fictional.
Go deeper
- [1906.01021] Coresets for Estimating Means and Mean Square Error with Limited Greedy Samples · arxiv.org
- [2004.10536] Learning Sampling and Model-Based Signal Recovery for Compressed Sensing MRI · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.