Error-Compensated Rounding at Two and Three Bits

Nearest rounding is optimal per element and not optimal per layer. Each weight is moved to the closest grid point, which minimizes that weight’s own error, but the layer’s output is a sum of products, and the errors of individual weights combine through the input distribution. Two weights with errors of opposite sign partially cancel in the output; two with the same sign reinforce. A rounding procedure that considers the layer’s output rather than each weight’s distance can do better than nearest rounding at the same bit width, and at two or three bits the difference is the difference between a working model and a broken one.

The objective, stated precisely

For a linear layer with weight matrix $W$ and a calibration set of inputs $X$, the layer’s output error from quantization is

$$| (W - \hat{W}) X |_F^2 = \mathrm{tr}\left( (W - \hat{W})^\top (W - \hat{W}) , X X^\top \right)$$

The matrix $H = X X^\top$ is the second moment of the inputs, the same quantity that appears in the sensitivity analysis used for bit allocation. It is positive semi-definite and, for a calibration set with more samples than the input dimension, positive definite. The error is not the sum of squared weight errors; it is a weighted sum, where the weights are the correlations between input features. A weight error along a direction in which the inputs have little variance costs almost nothing, and a weight error along a high-variance direction costs in proportion.

This changes the problem completely. Nearest rounding solves a diagonal version of the objective, where $H$ is a multiple of the identity. The real $H$ is far from diagonal: input features are correlated, and some directions in input space carry most of the variance.

The optimal unconstrained solution to this objective, for a fixed number of nonzero weights, is the subject of the classic pruning analysis that produced Optimal Brain Damage and its successors. For quantization the constraint is not sparsity but a grid, and the grid makes the problem combinatorial. The practical algorithms are greedy approximations to it, and they work by choosing each weight’s rounding in an order that lets earlier decisions be corrected by later ones.

Rounding one column at a time, with a residual

The sequential procedure processes the weight matrix one column at a time, in an order chosen in advance. For column $j$, every weight in that column is rounded to its nearest grid point, and the error introduced is recorded. The recorded error is then propagated to the not-yet-processed columns: for each remaining column $k$, the weights are adjusted by an amount that would compensate for column $j$‘s error, given the correlation between input features $j$ and $k$.

In matrix terms, if $e_j$ is the error vector for column $j$, the update to the remaining columns is $e_j$ times a row of the inverse of $H$ restricted to the remaining indices. The inverse appears because the compensation must be expressed in weight space while the objective is in output space, and $H^{-1}$ is the map between them.

The procedure is equivalent to a Cholesky-style elimination of the objective. It processes the quadratic form one variable at a time, and at each step it has the choice of which variable to eliminate next. The choice matters because the compensation is only applied to unprocessed columns: a column processed last has no later columns to absorb its error, so its error is final.

The result is that each weight’s rounding decision is made in the context of the compensation that later weights can provide. A weight that would round down and leave a positive error might round up instead if the compensation available from the remaining columns is small, and the decision is made by comparing the objective value under each choice, not the distance to each grid point.

import numpy as np

def quantize_columnwise(W, X, levels, blocksize=128):
    """Greedy error-compensating quantization of W [out, in] against X [in, n]."""
    H = X @ X.T                      # input second moment, [in, in]
    Hinv = np.linalg.inv(H)
    Wq = W.copy()
    n_in = W.shape[1]
    for start in range(0, n_in, blocksize):
        stop = min(start + blocksize, n_in)
        for j in range(start, stop):
            col = Wq[:, j]
            amax = np.abs(col).max()
            if amax == 0:
                continue
            scale = amax / levels
            q = np.clip(np.round(col / scale), -levels, levels)
            err = col - q * scale          # per-output-channel error
            Wq[:, j] = q * scale
            # Compensate the remaining columns using the correlation row.
            if j + 1 < n_in:
                factor = Hinv[j, j + 1:] / Hinv[j, j]
                Wq[:, j + 1:] -= np.outer(err, factor)
    return Wq

Two things in that code are worth naming. The scale is recomputed per column from the column’s own maximum, which is a per-input-channel scale rather than the per-output-channel scale used elsewhere in this cluster; the choice follows the compensation, which operates along the input axis. And the compensation is applied to all remaining columns at once, which is a dense rank-one update and the reason the algorithm is expensive.

Cost, and why blocking is not optional

A rank-one update of the remaining matrix costs $C_{\text{out}} \times (C_{\text{in}} - j)$ multiply-adds at step $j$, and the sum over $j$ is $C_{\text{out}} C_{\text{in}}^2 / 2$. For a 4096-by-4096 matrix that is about 34 billion operations per layer, which is orders of magnitude more than a forward pass. The whole procedure is offline and one-shot, so it is affordable, but only because it is not repeated.

The blocksize parameter in the code above is the standard mitigation. The inverse of $H$ is computed once for a block of columns rather than for the whole matrix, and the compensation is confined within the block. This makes the update cost proportional to $C_{\text{out}} \times \text{blocksize}^2$ per block, and the total linear in $C_{\text{in}}$ rather than quadratic. The price is that compensation cannot cross block boundaries, so a block’s last columns have less help than they would in a global pass.

The choice of blocksize is a trade between error and time, and it is not a free parameter: a block that is too large makes the inverse expensive and numerically fragile, and a block that is too small removes most of the compensation. Values in the low hundreds are typical, and the accuracy difference between a global pass and a blocked pass is small for well-conditioned $H$.

Numerical fragility

The inverse of $H$ is the fragile part. $X X^\top$ is positive semi-definite but can be ill-conditioned: if two input features are nearly collinear, the matrix has a small eigenvalue and its inverse has a large entry. A large entry in the inverse produces a large compensation, which pushes a weight far from its rounded value, which produces a large error in that column, which produces a larger compensation downstream. The procedure diverges, and it diverges on exactly the layers where the inputs are redundant.

The standard fix is to add a small multiple of the identity to $H$ before inverting, which is ridge regularization. It bounds the inverse’s largest entry at the cost of biasing the compensation. The damping constant is chosen relative to the mean diagonal of $H$, and its value determines how much of the compensation survives.

A second fix is to reorder the columns so that the ill-conditioned directions are processed last, where they have the least compensation to trigger. The ordering is derived from the diagonal of the inverse: columns with large inverse diagonal are processed first, because their rounding decision is the most consequential and the compensation available from the rest is the largest. This is the same pivot ordering that appears in Cholesky factorization, and for the same reason.

The third fix is to detect divergence and stop. If the compensation for a column exceeds some multiple of the column’s own scale, the compensation is capped or skipped. Skipping it converts the procedure back into nearest rounding for that column, which is worse locally but stable.

What the method cannot fix

The compensation moves error between columns. It does not remove it: the total error after processing is determined by the last columns, which have no successors to absorb their error, and by the block boundaries. For a matrix quantized to 2 bits, the grid has four levels, and a column whose values are spread across a wide range cannot be represented at all; the compensation shifts the residual into other columns, which then have their own four-level constraint. The procedure reaches a fixed point where every column’s error is as small as the grid allows given the others, and that fixed point can still be far from the original.

The second limitation is that the objective is the layer’s output error under the calibration distribution. A quantization that minimizes that error on the calibration set can be worse on inputs outside it, because the compensation is tuned to the correlations present in $X$. The failure is a form of overfitting, and it appears as a model that validates well on in-distribution data and degrades on inputs whose feature correlations differ.

The third is that the procedure quantizes weights against an exact input. In an integer inference path the input is itself quantized, so the error the compensation was designed to cancel is not the error that occurs. The compensation assumes $X$ is exact and adjusts $W$; if $X$ carries its own quantization error, the two errors combine without the designed cancellation. Weight-only quantization avoids this by keeping $X$ exact, which is another reason the technique is usually paired with the weight-only path rather than with full integer quantization.

The other way to change the problem

Compensation is one response to a coarse grid. A second is to change the basis so that the grid suits the data better, and the two are complementary.

A rotation applied to both the weights and the inputs leaves the product unchanged and changes the distribution of the weight columns. A Hadamard transform, which is an orthogonal matrix of $\pm 1$ entries and can be applied in $O(n \log n)$ operations by a fast transform of the same shape as the fast Fourier transform, spreads a heavy-tailed distribution across coordinates and tends to reduce the maximum magnitude. A column whose largest entry dominated its scale, and therefore starved the rest of the column of resolution, becomes a column with a more even spread after rotation, and the same grid represents it better. The transform is exact and reversible, so the inference path applies the inverse to the output, which costs one extra transform per layer.

The interaction with compensation is that rotation changes $H$. After a Hadamard transform the input second moment is closer to a multiple of the identity, because the transform mixes the correlated features, so the compensation has less to do and the objective is closer to the diagonal case. Running the transform first and the compensation second is therefore more effective than either alone, and it is the combination that makes 2-bit weight quantization viable for large matrices.

What remains at two bits is the irreducible fact that four levels cannot represent a distribution with a wide spread, no matter how the error is distributed. The techniques here rearrange which values lose information, and they can be remarkably effective at preserving the directions that matter for the layer’s output. They cannot create resolution that the format does not have. Below that, the options are to store a small number of weights at higher precision, which is the mixed-precision allocation applied within a tensor rather than across layers, or to change the model’s structure so that fewer parameters carry the load.