Up or Down? Adaptive Rounding for Post-Training Quantization

In plain words

To shrink a trained network, its weights are stored with fewer bits. The usual method rounds each weight to the closest value the smaller grid allows. This paper shows that the closest value is not always the best one. When a group of weights is rounded together, some should go up and some down, so that their combined effect on the network's output stays close to the original.

The method is called AdaRound. It decides the up-or-down choice for every weight in one layer at a time, looking at a small set of unlabelled images to see what the layer actually receives. It does not retrain the network. Instead it treats each choice as a dial that moves smoothly between "round down" and "round up", nudges the dials to reduce the layer's output error, and then snaps each one to a final yes or no. The payoff is that a network can keep most of its accuracy at 4-bit weights, where plain nearest rounding loses a great deal.

TL;DR

• AdaRound is a post-training weight-rounding method. For each weight it chooses between floor and ceil on the quantization grid, instead of always taking the nearest value §3.
• It minimises a per-layer reconstruction error, relaxed to a continuous problem with a rectified sigmoid and a regulariser that pushes each choice to 0 or 1 §3.3.
• It needs no fine-tuning and a small unlabelled calibration set (1024 ImageNet images in most ablations) §5.
• Author-reported: Resnet18 at 4-bit weights, FP32 activations, reaches 68.60% top-1 (FP32 69.68%) and nearest rounding reaches 23.99% tab. 4 tab. 6 §5.
• Relevant if you quantize convolutional or linear layers to 4 or 8 bits post-training and nearest rounding loses too much accuracy.

Problem

Post-training quantization is attractive because no retraining is needed, but the usual weight rounding (nearest value on a fixed grid) is taken as the default without being optimised for the task §1. The authors argue that this ignores how weight errors interact through the layer's inputs and the task loss §2. Their Taylor-expansion analysis of the task loss gives a quadratic binary problem over the rounding choices §3.1. Solving that exactly needs the full Hessian and is NP-hard, and the paper's experiments show a local per-layer MSE proxy gives nearly the same accuracy §3.2 tab. 2.

Assumptions stated in the paper: a trained network that has converged (so the gradient term is dropped) §3.1; the scale is fixed before rounding is optimised §3.1; batch-norm is absorbed into adjacent weights before quantization §5; activations stay in FP32 unless stated §5.

Method

Setup. A layer's weight W is quantized as floor or ceil of W / s on a fixed per-layer scale s, clipped to an integer range [§3.1, eqs. for the floor and ceil grid]. Rounding-to-nearest is the special case that picks whichever of the two is closer §1.

Step 1: local objective. A second-order Taylor expansion of the task loss, with the gradient term dropped at convergence, gives a quadratic cost in the weight perturbation §3.1. Two approximations reduce it to a per-layer problem. The Hessian of the loss with respect to the pre-activations is taken to be diagonal, and its diagonal is treated as a constant independent of the data. The result is a plain mean-squared error on the layer's pre-activations, which decomposes into independent problems for each output row (or, for convolutions, each output channel) §3.2 app. B.

Step 2: continuous relaxation. Each rounding choice becomes a soft variable V. A rectified sigmoid h(V) takes values in [0, 1] and replaces the binary choice inside the quantizer §3.3, rectified-sigmoid equation. The stretch parameters are zeta = 1.1 and gamma = -0.1 §3.3. The relaxed quantizer is floor(W / s) + h(V), clipped and rescaled.

Step 3: regulariser. A regulariser sums 1 - |2h - 1|^beta over all weights. It is zero only when every h is exactly 0 or 1 §3.3, regulariser equation. Beta is annealed from a large value, which lets h move freely at first, to a small value, which pushes h to 0 or 1 at the end §3.3 fig. 2.

Step 4: asymmetric reconstruction. The per-layer loss compares the output of the original float layer, after its activation function, with the output of the soft-quantized layer fed by the already-quantized earlier layers' outputs. This stops quantization error from building up across layers [§3.3, asymmetric equation]. The objective is optimised with stochastic gradient descent on calibration data §3.3.

Step 5: hard rounding. After optimisation each h is taken as 0 or 1 §3.3 fig. 3. The paper reports that the optimised values converge close to 0 or 1 fig. 3.

# AdaRound for one conv or linear layer  [§3.3] [§5]
# Input:  weight (out, in); scale s; calibration inputs (N samples)
# Output: rounding decisions in {0, 1}; adjusted weight on the grid
def adaround_layer(weight, scale, inputs, iters, beta_range):
    # Lower grid level for each weight  [§3.3]
    base = floor(weight / scale)
    # Initialise V from the fractional part  [code: adaround_wrapper.py#L94]
    V = init_from_fraction(weight / scale - base)
    for step in range(iters):
        # Sample a mini-batch of 32 calibration inputs  [§5]
        batch = sample_batch(inputs, 32)
        # Relaxed weights with the rectified sigmoid  [§3.3]
        h = clip(sigmoid(V) * (1.1 + 0.1) - 0.1, 0, 1)
        soft_w = scale * clip(base + h, lo, hi)
        # Compare FP32 and soft outputs after the activation  [§3.3]
        recon = mse(act(layer(batch, soft_w)), act(fp_layer(batch)))
        # Push h towards 0 or 1, with beta annealed  [§3.3] [fig. 2]
        reg = sum(1 - abs(2 * h - 1) ** beta(step))
        loss = recon + lam * reg
        V = adam_step(V, loss)
    # Hard decision: round up where V >= 0  [§3.3]
    # [code: adaround_wrapper.py#L64]
    return where(V >= 0, base + 1, base)

lo, hi, lam and beta stand for the bounds, the regulariser weight and the annealing schedule. The values are in Hyperparameters and Paper vs code. The pseudocode is a summary, not the code.

Hyperparameters

Results

Author-reported. Setup: ImageNet validation top-1, models from torchvision, per-layer symmetric weight quantization, five seeds (mean ± std), 1024 calibration images, 10k iterations unless stated §5. Activations FP32 unless stated.

Ablation on Resnet18, first layer only at 4 bits tab. 2:

Stochastic rounding, first layer tab. 1: 100 samples give mean 52.06 ± 5.52 and best 63.06, against nearest 52.29. Floor and ceil for all weights give 0.10. The text says 48 of the 100 samples beat nearest §2.

Regulariser and reconstruction ablation, Resnet18, all layers tab. 3 tab. 4:

Straight-through estimator and grid choice, Resnet18, all layers tab. 5 tab. 6:

Comparison with other post-training methods, 4-bit weights, FP32 activations tab. 7. Full precision: Resnet18 69.68, Resnet50 76.07, InceptionV3 77.40, MobilenetV2 71.72.

Same table, 4-bit weights and 8-bit activations tab. 7: DFQ (own run) 38.98, 52.84, -, 46.57; bias correction (per-channel) 67.4, 74.8, 59.5, -; AdaRound with activation quantization 68.55 ± 0.01, 75.01 ± 0.05, 75.72 ± 0.09, 69.25 ± 0.06 (CLE). The CLE preprocessing row is marked by the paper's dagger.

Bias correction vs AdaRound, Resnet18, 4-bit weights tab. 8: nearest 23.99, empirical bias correction 38.87, AdaRound 68.60.

Semantic segmentation, DeeplabV3+ with MobileNetV2 backbone, Pascal VOC, mIoU tab. 9: FP32 72.94; nearest at 4/8 gives 6.09; DFQ own run at 4/8 gives 14.45; AdaRound at 4/32 gives 70.89 ± 0.33; AdaRound with 8-bit activations gives 70.86 ± 0.37. Uses 512 calibration images §5.

The abstract and conclusion put Resnet18 and Resnet50 at 4-bit weights within 1% of FP32 §6. Taking the table values: 69.68 - 68.60 = 1.08 for Resnet18 (or 0.97 using the 68.71 row), and 76.07 - 75.23 = 0.84 for Resnet50. So the "within 1%" wording is exact only for the 68.71 Resnet18 row, and the 68.60 row sits just above 1%.

Paper vs code

Code: the AIMET library, official repository qualcomm/aimet, at commit 1b4844d1a1b8e4b2e2c9d86822d54f1045ee7dd8, BSD-3-Clause. Its AdaRound regression runner cites this paper and its docs link it [code: AIMETRegression/features/onnx/adaround_runner.py#L38-L39] [code: Docs/ptq_techniques/adaround.rst#L9]. This is the authors' employer's repository, not a third-party copy, but the paper itself links no code, so this match rests on the citations inside the repository. The clone is at raw/code/2004.10568.

Differences found by reading the code:

• Iterations. Paper: 10k iterations §5. Code: 10k only for weights of 8 bits or more and 15k for lower bit-widths, unless overridden adaround_weight.py adaround_weight.py. A 4-bit run in the code uses 15k.
• Regulariser and beta. Paper gives the regulariser form and says beta is annealed, but no numbers §3.3. Code defaults: reg_param 0.01, beta from 20 to 2, warm start 20% adaround_weight.py. Beta decays with a cosine schedule adaround_loss.py.
• Warm start. During the first 20% of iterations the rounding loss is zero adaround_loss.py. The paper has no warm-up step in the text read §3.3.
• Reconstruction norm. The paper writes a squared Frobenius norm of the layer output difference §3.3. The code computes torch.norm with p="fro" and dim=1, squares it and takes the mean adaround_loss.py. With an integer dim, PyTorch treats the norm as a 2-norm along that dimension. For convolution outputs that is the channel dimension, so the value is a mean over batch and spatial positions rather than a sum over all elements. This changes the scale of the reconstruction term relative to lambda. Not verified by a numerical test.
• Grid default. The paper's experiments set the scale by minimising weight MSE §5. The code's default quantization scheme is post_training_tf_enhanced adaround_weight.py. The paper's grid is not the default unless the user chooses it. I did not check how the TF-enhanced grid is computed.
• Calibration size. The paper uses 1024 images (2048 for the literature table) §5 §5.2. The code's docstring recommends the smaller of the loader length and ceil(2000 / batch size) batches, counted in batches, not images adaround_weight.py.
• Batch size and optimiser. Batch size 32 matches §5 adaround_optimizer.py. The code creates Adam without a learning rate, so PyTorch's default applies adaround_optimizer.py. The paper says "default hyper-parameters" §5, which is consistent but leaves the rate unstated.
• Activation function. The code applies the layer's following activation to both the target and the soft output before the loss, as the asymmetric objective does adaround_optimizer.py.
• Alpha initialisation. The code sets the soft variable from the fractional part of W / s, inverting the rectified sigmoid adaround_wrapper.py. The paper does not give an initialisation.
• Final weights. After optimisation the code sets the wrapper to hard rounding adaround_optimizer.py, then folds the weights by applying the soft rounding again adaround_weight.py. Folded weights therefore keep the soft values unless h has reached exactly 0 or 1. I did not check whether the regulariser makes that happen in practice.
• Asymmetric input on the default path. The code uses the quantized-side cached inputs only in the checkpoint path adaround_weight.py. On the default path it calls one cached dataset adaround_weight.py. I did not verify in the sampler which input the quantized side receives there.
• Supported layers. Conv2d, ConvTranspose2d and Linear only adaround_optimizer.py. Other modules raise an error.
• Stretch constants. Zeta and gamma come from AdaroundConstants. I did not open that file, so I did not check the values against 1.1 and -0.1 in the paper adaround_wrapper.py.

No discrepancy was found for the rectified-sigmoid form [code: TrainingExtensions/torch/src/python/aimet_torch/_base/adaround/adaround_wrapper.py#L53-L62] or for the 4-bit default [code: TrainingExtensions/torch/src/python/aimet_torch/_base/adaround/adaround_weight.py#L98].

Pitfalls

• Iteration count. The paper says 10k, the code uses 15k at 4 bits by default. Results depend on which you pick, so state it when comparing.
• Two Resnet18 numbers. 68.60 ± 0.09 is the main setting tab. 4 tab. 6. 68.71 ± 0.06 appears in the literature table with 2048 images and 20k iterations tab. 7. Do not mix them.
• Unstated regulariser and beta values. The regulariser weight, beta range and warm-up are not in the paper. Copy the code defaults or tune them, and say which you used.
• Reconstruction scale. The code's mean-of-channel-norms reduction and the paper's squared Frobenius norm can differ by a spatial factor. This changes the balance against the regulariser. Check it numerically before porting the loss.
• Soft weights after folding. The folded weights may not sit exactly on the integer grid if h did not reach 0 or 1. Check that the exported weights are on the grid before deploying.
• Grid choice. Using the code's default grid gives a different grid from the paper's weight-MSE grid. Table 6 shows AdaRound's gain depends on the grid, with min-max much worse before and after AdaRound.
• Memory. The code caches activations in CPU RAM if they fit within 3/4 of free memory, otherwise it reads inputs from disk adaround_optimizer.py adaround_optimizer.py.
• Hidden cost. Each layer is optimised in turn, so runtime grows with depth. The paper reports about 10 minutes for Resnet18 §5. It reports no timing for larger models.
• Sensitivity to data. Fewer calibration images work, down to 256 within 2% of FP32 §5.1 fig. 4. Images from other datasets (Pascal VOC, COCO) cost less than 0.2% on average §5.1 fig. 4.
• Activations. Most results keep FP32 activations. The W4A8 rows use min/max ranges for activations, and the paper says this rarely hurts §5.2 tab. 7.
• Scale is fixed. The scale is not optimised jointly with the rounding, because changing it changes the problem §5.1.

Lineage

• Builds on: the per-layer MSE and reconstruction ideas from earlier compression work, cited as §4 (names only, no card for them).
• Builds on: the rectified sigmoid and regulariser from Louizos et al. (2018), cited in §3.3.
• Compared against: bias correction (DFQ and the banner post-training work), OMSE + optimisation, and OCS, all named in tab. 7 (no cards).
• QUBO background: the cross-entropy solver used in the ablations §5.1 and app. A.
• Superseded by: not known from content in this wiki.