Part 2: Adaptive Rounding: Optimal Brain Quantizer (OBQ)

In Part 1, we looked at why rounding direction matters: a stochastic rounding experiment on a single ResNet-18 layer found assignments that beat round-to-nearest by over 10 points, just by choosing up vs. down more carefully. Then we looked into second order analysis of quantization, and reazlied that the damage is governed entirely by the Hessian, which is very expensive to compute for large models, and Adaround makes it tractable with three simplications: layer-wise, diagonal & constant curvature - this collapses the objective to a simple asymmetric M.S.E between the full-precision network and the quantized layer outputs. However, we also saw that Adaround requires 10k-15k gradient steps per layer to converge for 8 bits and 4 bits respectively.

In this post and the next, we will look at OBQ and GPTQ – the methods that made second-order adaptive rounding tractable at scale.

Problem Formulation¶

Recall the following optimization objective we derived using the Taylor series expansion in the AdaRound blogpost:

where:

• \(\Delta \mathbf{w}^{(\ell)} = \hat{\mathbf{w}}^{(\ell)} - \mathbf{w}^{(\ell)}\) is the perturbation introduced by quantizing the weights at layer \(\ell\) i.e, the vector of rounding errors
• \(\mathbf{H}(\mathbf{w}^{(\ell)})\) is the Hessian of the task loss with respect to the weights at layer \(\ell\), capturing how sensitive the loss is to perturbations in each weight and each pair of weights

Minimizing this objective means finding the quantized weights \(\hat{\mathbf{w}}^{(\ell)}\) whose rounding errors cause the least damage to the loss, weighted by the curvature of the loss landscape.

The challenge is that computing the full Hessian \(\mathbf{H}(\mathbf{w}^{(\ell)})\) is intractable for large layers. Let’s now understand how OBQ approximates it it efficiently enough to apply at scale.

From Layer Error to Row-wise Weight Errors¶

Consider a single output:

where:

• \(w_j\) is the \(j\)-th weight in a given row
• \(x_j\) is the corresponding input activation

Our goal is to preserve \(y\) as much as possible after quantization. Notice that a single output value depends only on one row of the weight matrix, which means we can treat each row independently.

Quantizing weight \(w_j\) introduces a rounding error:

The resulting output error is:

Goal¶

A natural proxy objective is to minimize the expected squared output error:

Substituting the expression for \(\Delta y\):

Defining:

gives:

Comparing this with the Taylor-series objective above, we can see they have identical structure — \(\Delta\mathbf{w}^\top \mathbf{H} \Delta\mathbf{w}\) — which means the input covariance matrix \(\mathbf{H} = \mathbb{E}[\mathbf{x}\mathbf{x}^\top]\) is the tractable proxy for the true Hessian. Rather than computing second derivatives of the task loss (expensive and requiring backprop through the entire network), we approximate the Hessian using only the layer’s input activations, which are cheap to collect with a small calibration dataset.

With that background, let’s unpack OBQ step by step.

OBQ: Motivation¶

At its core, OBQ asks the following question:

Quantizing weight \(w_j\) introduces a fixed rounding error \(\Delta w_j\). How should the remaining weights in the current row be adjusted to absorb this error?

Consider a simple example with three weights:

Suppose we quantize \(w_1\), introducing a fixed rounding error:

We would like to compensate by adjusting the remaining weights:

The challenge is to find \(\Delta w_2\) and \(\Delta w_3\) that minimize the output damage.

Recall the Proxy Loss¶

From the previous section, the proxy loss has the same structure as the Taylor-series objective, with the input covariance as the Hessian:

where \(\Delta\mathbf{w} = [\Delta w_1, \Delta w_2, \Delta w_3]^\top\).

Expanding the quadratic form:

Since \(\Delta w_1\) is fixed (determined when \(w_1\) was quantized), we minimize \(L\) with respect to \(\Delta w_2\) and \(\Delta w_3\).

Taking Partial Derivatives¶

Differentiating with respect to \(\Delta w_2\):

Setting to zero:

Differentiating with respect to \(\Delta w_3\):

Setting to zero:

Linear System¶

Collecting both equations:

Solving:

This tells us how the remaining weights should be adjusted to compensate for the error introduced when quantizing \(w_1\).

Notice that \(\Delta w_2\) and \(\Delta w_3\) cannot be computed independently. The optimal adjustment to \(w_2\) depends on the adjustment applied to \(w_3\), and vice versa. This happens because the corresponding input features are correlated, which is captured by the off-diagonal term \(H_{23}\).

In other words, the remaining weights must coordinate their adjustments to minimize the overall output damage.

The OBS / GPTQ Trick¶

A naive implementation of the above procedure would look something like this:

The problem with this implementation is that matrix inversion is expensive. Inverting an \(n \times n\) matrix typically costs \(O(n^3)\), which quickly becomes prohibitive for modern neural network layers containing thousands of weights.

The key insight behind OBQ, later leveraged by GPTQ, is that these repeated inversions are unnecessary. Instead of recomputing the inverse Hessian after every quantization step, OBQ proves that the inverse corresponding to the remaining weights can be obtained directly from the original inverse Hessian by subtracting an outer-product from it to perfectly simulate having removed one weight.

Suppose weight \(p\) has just been quantized. Let \(H^{-1}\) denote the inverse Hessian before removing that weight. The inverse Hessian for the remaining (unquantized) weights is

where

• \(H^{-1}_{:,p}\) is the \(p\)-th column of the inverse Hessian,
• \(H^{-1}_{p,:}\) is the \(p\)-th row of the inverse Hessian, and
• \(\left(H^{-1}\right)_{pp}\) is the diagonal entry corresponding to the quantized weight.

The outer product \(H^{-1}_{:,p} H^{-1}_{p,:}\) costs only \(O(n^2)\), a significant improvement over the \(O(n^3)\) cost of performing a fresh matrix inversion.

This update can be interpreted as a form of Gaussian elimination. Rather than explicitly removing the quantized weight and reinverting the reduced Hessian, we perform elementary row and column operations on the inverse Hessian that eliminate the influence of weight \(p\). The resulting matrix is exactly the inverse Hessian we would have obtained had we removed that weight and recomputed the inverse from scratch.

Generalized Weight Update¶

Equation (1) showed how to optimally redistribute the quantization error from a single weight to the remaining weights in our three-weight example. The authors generalized this result to an arbitrary number of weights.

Suppose weight row \(w_p\) has just been quantized, introducing a rounding error:

The optimal adjustment to all remaining (unquantized) weights is:

The remaining weights are then updated as:

Greedy Search¶

So far, we’ve formulated quantization as a layer-wise output reconstruction problem, derived the optimal weight compensation rule, and shown how to efficiently update the inverse Hessian after each quantization step.

The remaining question is:

How do we choose which weight to quantize at each step?

OBQ answers this with a greedy search. At every step, it selects the weight whose quantization incurs the smallest increase in the proxy loss.

From the OBQ derivation, quantizing weight \(w_p\) increases the loss by:

where \(\Delta w_p = \hat{w}_p - w_p\) is the rounding error introduced by quantizing \(w_p\). At every iteration, OBQ selects the unquantized weight with the smallest \(\Delta L_p\).

Limitations of OBQ¶

OBQ produces excellent quantization quality but does not scale to large language models. Let’s understand why:

1. Greedy Search Cost¶

Consider a layer with weight matrix of size \(4096 \times 4096\). For a single row, OBQ evaluates all 4096 unquantized weights via \(\Delta L_p\), selects the best, quantizes it, compensates the remaining 4095 weights, and updates \(\mathbf{H}^{-1}\). This is repeated 4096 times for all the columns. Across 4096 rows, and multiple layers, the total work becomes enormous.

2. Parallelism Across Rows Is Broken by Greedy Search¶

In principle, each row of the weight matrix can be quantized independently but the greedy search breaks this.

Suppose we quantize two rows simultaneously. Row 1’s greedy search picks column 1 first; row 2’s picks column 200. After these decisions, the two rows need different inverse Hessians i.e, one with row and column 1 eliminated, the other with row and column 200 eliminated. They can no longer share a single \(\mathbf{H}^{-1}\).

For a \(4096 \times 4096\) layer, one copy of \(\mathbf{H}^{-1}\) holds \(4096^2 \approx 16.8\text{M}\) entries, roughly 64 MB in float32. Quantizing \(k\) rows in parallel requires \(k\) separate copies, so memory scales linearly with the degree of parallelism. At any meaningful scale, this is prohibitive.

Conclusion¶

Thank you for reading! I hope the post helped demystify OBQ, the foundation on which the very popular technique GPTQ is built. We’ll cover GPTQ in the next post. See you there!

Here’s a quick summary of OBQ’s design choices:

References and Further Reading¶

1. Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning
2. Optimal Brain Surgeon and General Network Pruning
3. Up or Down? Adaptive Rounding for Post-Training Quantization
4. GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers