Adaptive rounding architecture — layer reconstruction, learned masks, annealing, frozen integers

Why architecture matters here

The architectural reason this matters is that quantisation error is the one error in a serving stack you cannot fix later. Everything else in a deployment is recoverable: a bad scaling policy can be retuned, a slow kernel can be replaced, a wrong batch size can be changed at runtime. But once you have quantised a checkpoint and shipped it, the information destroyed by rounding is gone. Any accuracy you failed to preserve at build time becomes a permanent property of the artefact, and the only remedy is to re-run the whole pipeline. That asymmetry — cheap to do well once, impossible to fix in production — is what justifies spending GPU-minutes on a rounding search that changes nothing about how the model runs.

It matters more as bits go down, and it matters non-linearly. At 8 bits the gap between the two candidate integers is small relative to the weight distribution, so the damage round-to-nearest can do is bounded and adaptive rounding buys you a fraction of a point. At 4 bits the gap is large, the per-weight error is a serious fraction of the weight's own magnitude, and the accumulated output error is what separates a usable model from a broken one. This is why adaptive rounding was a curiosity in the 8-bit CNN era and is table stakes in the 4-bit LLM era: the same technique, applied where the errors are big enough for the search to have something to work with.

It also matters because of where it sits relative to everything else in the quantisation toolbox — a point that is routinely confused. Rounding is the last step, and it is downstream of every other decision. Techniques that reshape the weight or activation distributions before quantisation, or that choose scales more cleverly, or that change what the grid looks like, all operate on the inputs to rounding. They change the numbers being rounded. Adaptive rounding changes the rounding itself. That means it composes with those techniques rather than competing with them, and it means that improving them does not make it redundant: no matter how good your scales are, you still have a fractional number and a binary choice at the end.

Finally, it matters because of where it sits on the cost curve. Quantisation-aware training gets better results but needs a training loop, labelled data, hyperparameters, and hours to days of compute. Adaptive rounding needs a few hundred unlabelled forward passes and a local objective whose target is the FP model's own output, which you have for free — most of the gap between naive PTQ and QAT, at a small multiple of naive PTQ's cost, with none of QAT's machinery.

The architecture: every piece explained

Start with what is actually being decided. Given a weight w and a scale s, the fractional value is w/s. Round-to-nearest computes round(w/s). Adaptive rounding instead computes floor(w/s) + b, where b is a binary variable in {0, 1} chosen per weight. Note what this guarantees structurally: the result is always one of the two integers adjacent to the true value, so no weight can move more than one step. The search cannot invent a wild value. The worst case is bounded by the same grid spacing that bounds round-to-nearest; the only question is which side of the fence each weight lands on.

The objective is the piece that makes it work, and it is deliberately local. You do not optimise the network's end-to-end loss — that would need labels, would couple every layer to every other, and would be a training run. Instead you take one layer at a time and minimise the difference between the quantised layer's output and the original layer's output on real activations: ||Wx - W_q x||. This is layer-wise reconstruction, and it is the reason the method is post-training rather than training. It is unsupervised (the FP layer supplies the target), it is parallelisable across layers, and it is small enough to fit on one GPU regardless of model size.

The calibration data supplies the x in that objective, and this is where the technique gets its power and its main failure mode. A few hundred real input samples are enough to reveal which weights are multiplied by big numbers and which by small ones, and that ranking is what tells the optimiser where to spend its error budget. The data must be unlabelled but it must be representative: the search is explicitly fitting the rounding to the activation statistics you show it. Show it the wrong distribution and it will faithfully optimise for the wrong distribution.

The continuous relaxation is the trick that makes the search tractable. You cannot backpropagate through a binary variable, so b is replaced by h(V), where V is a real-valued learnable tensor of the same shape as the weights and h is a rectified sigmoid — a sigmoid stretched so that its output saturates hard at exactly 0 and exactly 1 over a range of inputs, rather than approaching them asymptotically. That detail matters: it means V can reach a genuinely binary state during optimisation rather than merely getting close to one, so the final hard-threshold step is not a lie about what was optimised.

The regulariser closes the gap. Left alone, gradient descent is happy to park h(V) at 0.5 for many weights — a half-rounded weight that minimises the objective but does not exist in any integer format. So an annealing term is added that starts near zero and grows over the optimisation, penalising values away from the extremes and progressively pushing every h(V) to 0 or 1. Early on the optimiser explores freely; late on it is forced to commit. When it finishes, h(V) is thresholded, the integers are frozen as W_int = floor(W/s) + round(h(V)), and V is discarded. Nothing about it survives into the artefact.

End-to-end flow

A pipeline run starts with an FP checkpoint and a calibration set of a few hundred unlabelled samples drawn from the real input distribution. The first pass is ordinary: run the FP model over the calibration batch and cache, for each layer you intend to quantise, the inputs that layer actually saw. These cached activations are the entire empirical content of the method. Alongside them you compute the scale and zero-point by whatever method the pipeline already uses; adaptive rounding takes these as given.

Then the pipeline walks the layers in order. For each one it constructs the objective: on one side, the FP layer's output on the cached inputs (the target, computed once); on the other, the quantised layer's output with rounding governed by the learnable V. V is initialised so that h(V) starts at the round-to-nearest solution — an important choice, because it means the optimisation begins from the current best-known answer and can only improve on it from there, rather than starting from noise and hoping to catch up.

The inner loop is a small optimisation, typically a few thousand steps of Adam over V only. Weights are frozen. Scales are frozen. The activations are fixed tensors in memory. Nothing is being trained in any meaningful sense — the only parameters are the rounding decisions, and the only gradient signal is 'does flipping this weight up or down make this layer's output closer to the original?'. The regulariser weight is annealed on a schedule across those steps, so the loss surface itself changes shape as the run proceeds: permissive at the start, increasingly insistent on binary values toward the end.

The subtlety in the layer walk is which inputs you feed. If each layer is reconstructed using the FP model's activations, every layer is optimised in isolation against a clean input it will never actually see in production, because in the real quantised network its input comes from the previous quantised layer and carries that layer's error. The alternative is to feed each layer the activations produced by the already-quantised prefix, which makes every layer's rounding search aware of the error it is inheriting and gives it the chance to compensate. This costs a re-run of the prefix per layer and is strictly the better choice at low bit-widths, where accumulated error is the dominant term.

When the layer's optimisation converges, h(V) is hard-thresholded at 0.5, the integer tensor is written, and V is thrown away. The layer is now an ordinary quantised layer, indistinguishable in format from one produced by round-to-nearest. The pipeline moves to the next layer, and — if it is using quantised-prefix inputs — recomputes the cached activations through the layer it just froze.

What ships at the end is an ordinary checkpoint whose only trace of the whole exercise is that some fraction of its integers — a minority, weighted toward the channels the calibration data said mattered — sit on the far side of the fence from where nearest would have put them.

Failure modes and mitigations

• Unrepresentative calibration data. The single largest risk, and it fails silently. The search fits rounding to the activation statistics you provide; feed it clean web text and serve code, and you have optimised for the wrong distribution while every build-time metric looks fine. Draw calibration samples from production traffic, and re-draw them when the traffic changes.
• Overfitting to the calibration set. A few hundred samples and a few thousand optimisation steps over a tensor with as many parameters as the layer is a textbook overfitting setup. The layer's reconstruction loss keeps improving while generalisation quietly degrades. Hold out a slice of calibration data that the optimiser never sees and watch its reconstruction error, not just the training slice's.
• Half-rounded weights at freeze time. If the annealing schedule is too short or too weak, a population of h(V) values is still sitting near 0.5 when you threshold. Those weights get an essentially arbitrary assignment that the objective never endorsed. Instrument the fraction of h(V) outside a band around 0.5 at the end of every layer and fail the run if it is not overwhelming.
• FP-input reconstruction at low bit-width. Optimising every layer against clean FP activations is cheaper and is the default in several implementations, but it means no layer ever compensates for inherited error. At 4 bits this is the difference between a working model and a broken one. Use quantised-prefix inputs and pay the re-run.
• Layer-local success, network-level failure. Reconstruction error is a proxy. It is possible for every layer to report an excellent local fit while end-to-end quality drops, because the objective never saw the network's actual task. Always close the loop with a real end-to-end evaluation; never ship on reconstruction loss alone.
• Interaction with distribution-reshaping techniques. Methods that move outliers between weights and activations change the numbers rounding sees. Applied in the wrong order — rounding search first, reshaping second — the search has optimised against a distribution that no longer exists. Rounding is last, always.
• Cost surprise on very large models. The per-layer optimisation is small, but there is one per layer and the activation cache is large enough to exhaust host memory. Stream the cache and treat the run as a scheduled build job, not something you do interactively.

Operational playbook

• Make it a build artefact, not a notebook. Adaptive rounding is deterministic given data, seed, and schedule. Pin all three, run it in CI, and version the output checkpoint with a manifest recording the calibration set hash. A rounding search you cannot reproduce is a checkpoint you cannot debug.
• Version the calibration set explicitly. It is an input to the artefact as surely as the weights are. Store it, hash it, and record which set produced which checkpoint. When quality regresses after a re-quantisation, the calibration set is the first thing you will want to diff.
• Track flip rate per layer. The fraction of weights the search moved off nearest is the technique's own vital sign. Near zero means the search found nothing and you are paying for nothing. Implausibly high means something is wrong with the scales or the data. Both are actionable; neither is visible in the loss.
• Evaluate end to end on a task you care about. Reconstruction loss goes in the logs, task metrics go in the gate. The whole point of the method is a proxy objective, and the only way to know the proxy held is to check the thing it was standing in for.
• Compare against round-to-nearest every time. Keep the naive checkpoint as the control. It is nearly free to produce and it is the only way to attribute a quality change to the rounding search rather than to a scale change, a data change, or a library upgrade.
• Budget the run honestly. Wall-clock scales with layer count and with whether you recompute quantised-prefix activations. Measure it once on the real model and put the number in the pipeline's documentation, so nobody plans a release around a quantisation step they think takes minutes.
• Re-run when the input distribution moves. The rounding is fitted to a distribution. When production traffic shifts materially, the fit is stale in a way that no serving metric will announce. Tie re-quantisation to the same cadence as any other distribution-sensitive component.