Optimal Brain Compression

Frantar, Singh & Alistarh, NeurIPS 2022 — a unified framework for accurate post-training pruning and quantization.

Setup

Post-training compression is done layer-wise: for each layer $\ell$ with weights $W_\ell$ and calibration inputs $X_\ell$, find compressed weights $\widehat{W}_\ell$ (sparse, quantized, or both) minimizing the layer-wise reconstruction error

\begin{equation} \min_{\widehat{W}_\ell} \; \big\lVert W_\ell X_\ell - \widehat{W}_\ell X_\ell \big\rVert_2^2 . \end{equation}

Write the layer size as $d = d_{\text{out}} \times d_{\text{in}}$. The punchline of the paper is an exact greedy OBS solver whose cost drops from the naive $O(d^4)$ to $O(d \cdot d_{\text{in}}^2)$.

Background: Optimal Brain Surgeon

Choosing which weights to prune is NP-hard, so OBS is a greedy algorithm that prunes weights one-by-one: at each step, remove the weight with the smallest loss increase and update all remaining weights to compensate,

\begin{equation} w_p = \operatorname*{argmin}_{w_p} \; \frac{w_p^2}{[H^{-1}]_{pp}}, \qquad \delta_p = - \frac{w_p}{[H^{-1}]_{pp}} \, H^{-1}_{:,p}, \end{equation}

where $H^{-1}_{:,p}$ denotes column $p$ of the inverse Hessian. Applied directly to a full layer this is hopeless: the Hessian is in $\mathbb{R}^{d \times d}$, which we can't even materialize, and inverting it once would be $O(d^3)$ — so algorithms that re-invert after pruning each weight would be $O(d^4)$.

Naive OBS cost time $O(d^4)$,   space $O(d^2)$.

Efficient OBS solver (ExactOBS)

Idea 1: work with per-row Hessians

For the squared reconstruction error, the problem decomposes over the $d_{\text{out}}$ rows of $W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$: each row has its own Hessian $H = 2 X X^{T} \in \mathbb{R}^{d_{\text{in}} \times d_{\text{in}}}$ — and all rows share the same Hessian initially.

W din columns dout rows H din × din
Each row of $W$ gets its own $d_{\text{in}} \times d_{\text{in}}$ Hessian, small enough to actually store and invert.
Insight

The layer-wise problem decomposes over rows: run OBS on each row of $W$ with its own $d_{\text{in}} \times d_{\text{in}}$ Hessian, which is small enough to actually store and invert.

Note that at each global step we don't want to commit to some pre-determined row and its current inverse Hessian — we want to look at all rows and prune the best weight overall. Hence all $d_{\text{out}}$ per-row inverse Hessians are kept in memory.

Cost total time $O(d_{\text{out}} \cdot d_{\text{in}}^4)$,   memory $O(d_{\text{out}} \cdot d_{\text{in}}^2)$.

Idea 2: rank-1 updates of the inverse Hessian

Look at two consecutive steps of the iterative process. At first we have some $H$ and its inverse $H^{-1}$; pruning weight $p$ removes row and column $p$ from $H$, giving $H_{-p}$.

Insight

$(H_{-p})^{-1}$ can be computed from $H^{-1}$ directly — no need to re-invert from scratch — with a simple rank-1 correction, one step of Gaussian elimination of row and column $p$:

\begin{equation} (H_{-p})^{-1} = \Big( H^{-1} - \frac{1}{[H^{-1}]_{pp}} \, {\color{#1a56b0}{H^{-1}_{:,p}}} \, {\color{#c0392b}{H^{-1}_{p,:}}} \Big)_{-p} . \end{equation}
col p row p
Gaussian elimination of row and column $p$ of $H^{-1}$.

In practice we don't even need to resize the matrix: the Gaussian elimination formula leaves row and column $p$ full of zeros, so they don't influence any future calculations.

Cost total time $O(d_{\text{out}} \cdot d_{\text{in}}^3)$,   memory $O(d_{\text{out}} \cdot d_{\text{in}}^2)$.

Idea 3: jointly consider all rows

Observation: within one row, the order in which its weights get pruned does not depend on the other rows — the issue is only that row operations can get interleaved in the global greedy order. Even more, the loss increase $\frac{w_p^2}{[H^{-1}]_{pp}}$ of each pruning step can also be computed one row at a time.

Insight

Completely prune each row on its own (keeping its full per-row inverse Hessian only while processing that row) and record the loss increase of every step — $O(d)$ numbers total. Then just sort these globally to create the pruning mask $M$.

Note that we still need to look at the whole layer because the target sparsity is per layer, not per row — the global sort of loss increases is what decides how many weights each row loses.

Once the per-row masks $M_i$ are known, the remaining weights of each row can be updated in a single step with the group OBS formula:

\begin{equation} \delta_{M_i} = - H^{-1}_{:,M_i} \big( (H^{-1})_{M_i} \big)^{-1} w_{M_i} . \end{equation}
Cost total time $O(d_{\text{out}} \cdot d_{\text{in}}^3)$,   memory $O(d_{\text{out}} \cdot d_{\text{in}} + d_{\text{in}}^2)$.

Structured sparsity extensions

Extension to quantization