The LLM Surgeon

How do we obtain a good surrogate of the loss $q$? One of the easiest approaches is to locally expand the log loss through a second-order Taylor expansion around the pretrained weights ${\bm{\theta}}^{*}$, yielding:

where $[\nabla\mathcal{L}({\bm{\theta}}^{*})]_{i}=\frac{\partial}{\partial{\bm{\theta}}_{i}}\mathcal{L}({\bm{\theta}}_{i}^{*})$ denotes the Jacobian and $[{\bm{H}}_{{\bm{\theta}}}]_{ij}=\frac{\partial^{2}}{\partial{\bm{\theta}}_{i}{\bm{\theta}}_{j}}\mathcal{L}({\bm{\theta}}_{ij})$ denotes the Hessian. The first-order term vanishes $[\nabla\mathcal{L}({\bm{\theta}}^{*})]_{i}={\bm{0}}$ at the optimum. Note that in practice the first order term may not vanish. While we follow this assumption initially, we consider interleaved first-order corrections to mitigate the issue in section 3.6. The quadratic expansion of eq. 3 forms the basis of the optimal brain damage (LeCun et al., 1989) and optimal brain surgeon (Hassibi & Stork, 1992) pruning methods. Note that from a probabilistic perspective, a quadratic approximation of the log likelihood implies a Gaussian approximation of the likelihood, as also observed by (Wang et al., 2019) and illustrated in fig. 2. This is well-known (Bishop & Nasrabadi, 2006), (MacKay, 2003) as the Laplace approximation $q({\bm{\theta}})=\mathcal{N}({\bm{\theta}}\mid{\bm{\theta}}^{*}+\nabla\mathcal{L}({\bm{\theta}}^{*}),{\bm{H}}_{{\bm{\theta}}^{*}}^{-1}$), with pretrained weights are the mean and the local inverse Hessian is the covariance matrix capturing correlations between weights.

For a network trained with negative log-likehood loss, the Hessian is identical to the Fisher matrix:

which has the benefit of always being positive semi-definite, with the inverse thus forming a proper covariance matrix for $q$, and can be approximated with Monte Carlo samples of $p_{{\bm{\theta}}}(y|x_{n})$. For most LLMs, this would be treating the softmax output of the network as categorical distribution $p_{{\bm{\theta}}}(y|x_{n})$, and sampling from that. In practice, we use the ‘empirical Fisher’ replacing the expectation over $y$ with target data $y_{n}$ (Kunstner et al., 2019). The full (empirical) Fisher ${\bm{F}}_{{\bm{\theta}}}\in\mathbb{R}^{P\times P}$ scales quadratically in the number of parameters $P$. To overcome this, the Fisher is often written in terms of layer-wise blocks ${\bm{F}}_{lk}=\sum_{n=1}^{N}\mathbb{E}\left[\text{vec}(\nabla_{{\bm{W}}_{l}}\log p_{{\bm{\theta}}}(y|x_{n}))\text{vec}(\nabla_{{\bm{W}}_{k}}\log p_{{\bm{\theta}}}(y|x_{n}))^{T}\right]$, and approximated by only treating layers independently (Martens & Grosse, 2015; Botev et al., 2017):

where $\otimes$ denotes the Kronecker product and $\text{vec}(\cdot)$ the matrix vectorisation operation. Because we disregard cross-layer interactions we write ${\bm{F}}_{l}$ instead of ${\bm{F}}_{ll}$ for Fisher blocks associated with the weight matrix ${\bm{W}}_{l}{\in}\mathbb{R}^{R\times C}$ producing outputs ${\bm{y}}_{l,n}={\bm{W}}_{l}{\bm{a}}_{l,n}{\in}\mathbb{R}^{R}$ from inputs ${\bm{a}}_{l,n}{\in}\mathbb{R}^{C}$, for each layer $l$ and datapoint $n$. Consequently, we can compute Fisher blocks from input activations ${\bm{a}}_{l,n}{\in}\mathbb{R}^{C}$ of forward-passed data $x_{n}$ and output gradients ${\bm{g}}_{l,n}{=}\nabla_{{\bm{y}}_{l,n}}\mathcal{L}{\in}\mathbb{R}^{R}$ from backpropagation.

Optimal brain surgery relies on removing and adapting weights such that the loss is least negatively affected, thus it behooves us to write the problem as a constrained optimization problem. From the Gaussian approximation discussed in section 2.1 obtained by quadratically expanding the log likelihood loss ${-}\log p{\approx}\frac{1}{2}{\bm{\theta}}^{T}{\bm{F}}{\bm{\theta}}$, the optimal update $\Delta{\bm{\theta}}{=}\hat{{\bm{\theta}}}{-}{\bm{\theta}}$ (and thus also $\hat{{\bm{\theta}}}{=}{\bm{\theta}}{+}\Delta{\bm{\theta}}$) becomes the following equality constrained quadratic optimization problem (Hassibi & Stork, 1992):

where ${\bm{F}}$ is positive semi-definite and $\mathcal{K}$ is the set of $K$ indices that are pruned (i.e., set to zero).

Accurate pruning relies on approximating the local curvature accurately while overcoming the memory cost associated with storing the true curvature. Specifically, even with the block-wise approximation of eq. 5, ${\bm{F}}\in\mathbb{R}^{RC\times RC}$ requires summing $N$ large $RC\times RC$ matrices, too large to practically fit in memory. Instead, we adapt the KFAC approximation (Martens & Grosse, 2015) that assumes independence of activations and derivatives, approximating an expectation of Kronecker products as a Kronecker product of two expectations $\mathbb{E}[{\bm{g}}_{l,n}{\bm{g}}_{l,n}^{T}\otimes{\bm{a}}_{l,n}{\bm{a}}_{l,n}^{T}]\approx\mathbb{E}[{\bm{g}}_{l,n}{\bm{g}}_{l,n}^{T}]\otimes\mathbb{E}[{\bm{a}}_{l,n}{\bm{a}}_{l,n}^{T}]$, allowing layer-wise Fisher blocks to be approximated as ${\bm{F}}_{l}\approx\widetilde{{\bm{F}}}_{l}$, where

constructed from activations ${\bm{a}}_{l,n}\in\mathbb{R}^{C}$ from forward passes and gradients ${\bm{g}}_{l,n}\in\mathbb{R}^{R}$ from backward passes (Eschenhagen et al., 2024). The approximation originates from optimization literature, but has recently gained popularity for other problems that require curvature approximations (Immer et al., 2022; van der Ouderaa et al., 2023), including structured pruning in Wang et al. (2019).

An additional advantage of approximating Fisher blocks as Kronecker products is that the inverse becomes particularly easy to compute $\vspace{-0.2em}\widetilde{{\bm{F}}}^{-1}={\bm{G}}^{-1}\otimes{\bm{A}}^{-1}$, thus only requires inverting the factors. This fact allows us to never explicitly construct large $RC{\times}RC$ matrices in memory that make up $\widetilde{{\bm{F}}}$ and $\widetilde{{\bm{F}}}^{-1}$, but rather directly work with the much smaller matrices ${\bm{G}}$ and ${\bm{A}}$.

The number of possible combinations in which weights can be removed grows (supra-)exponentially in parameter count, making it infeasible to estimate a separate cost $\mathcal{L}$ for each such removal. A common strategy, therefore, is to treat weights independently when computing removal costs $\mathcal{L}$. We also follow this strategy, but note that this does not necessarily imply that we have to make such same strong independence assumption for the weight updates $\Delta{\bm{\theta}}$ after selecting weights to be removed. Unlike most prior work, we present correlated weight updates by taking into account off-diagonal elements of the Fisher approximation in section 3.4.

For semi-structured and unstructured we use independent costs for individual weight elements $k{\in}[1,RC]$, and for structured use independent costs for all rows $r{\in}[1,R]$ and columns $c{\in}[1,C]$. We find that we can derive the appropriate costs from the general cost formula eq. 7 by letting ${\bm{E}}{=}{\bm{e}}_{k}\in\mathbb{R}^{RC}$ where the single one-hot element at index $k$ of canonical basis vector ${\bm{e}}_{k}$ selects the weight to remove. For structured pruning, we similarly select rows $r$ and columns $c$, by setting ${\bm{E}}{=}{\bm{e}}_{r}^{T}{\otimes}{\bm{I}}{\in}\mathbb{R}^{C\times RC}$ or ${\bm{E}}{=}{\bm{I}}{\otimes}{\bm{e}}_{c}{\in}\mathbb{R}^{R\times RC}$ with ${\bm{e}}_{r}{\in}\mathbb{R}^{R}$, ${\bm{e}}_{c}{\in}\mathbb{R}^{C}$. Plugging into eq. 7, we find:

Full derivations can be found in sections A.2 and A.3. The costs for single elements $\mathcal{L}_{k}$ are equivalent to those found in optimal brain surgeon (Hassibi & Stork, 1992) and $\mathcal{L}_{r}$ and $\mathcal{L}_{c}$ closely resemble structured brain surgeon of (Wang et al., 2019), but in our case derived for matrix rows and columns (see section A.3). Given curvature estimates, costs for either removing all weights or all rows and columns can be computed in parallel. In addition, we derive costs for the more general sum of Kronecker factor approximation $\widetilde{{\bm{F}}}\approx{\bm{G}}_{1}\otimes{\bm{A}}_{1}+{\bm{G}}_{2}\otimes{\bm{A}}_{2}$ in appendix I through an eigendecomposition.

Unlike prior works that compress layer-by-layer (Frantar & Alistarh, 2023), we use a global threshold $\tau$ enabling a dynamic allocation of sparsity levels across layers, pruning most where it hurts the least. Our method can compress a model to a specifically chosen target size $\alpha$, defined as the fraction of weights that should remain, i.e. stay non-zero after compression. In all structured, semi-structured, and unstructured pruning (fig. 3), we select as many weights for removal so that the target size $\alpha$ is reached that inflict the least possible costs $\mathcal{L}$, as computed according to section 3.2. For unstructured pruning, this is as simple as sorting the costs for all weights $\mathcal{L}_{k}$ in the network and setting a global threshold $\tau$ such that $\alpha$ fraction of weights fall within the threshold $\mathcal{L}_{k}\leq\tau$. For M:N semi-structured pruning, we sort the M costs of each N consecutive weights and select the M weights with lowest cost. In case of a multi shot schedule (see section 3.5) we also sum the M lowest costs in each block to find a cost per block, sort costs per block across the entire network, and similar to the unstructured case set a global threshold $\tau$ such that an $\alpha$ fraction of weights fall within threshold. Lastly for structured pruning, we perform a sorting appropriately weighted by the number of elements that make up a row or column and set the global threshold $\tau$ such that $\alpha$ fraction of all weights fall within the threshold. Then we remove all rows and columns that fall within the threshold $\mathcal{L}_{r},\mathcal{L}_{c}\leq\tau$.

Like most other pruning methods, we prune multiple weights at once (Frantar & Alistarh, 2023; Wang et al., 2019). To arrive at pruning costs and weight updates for pruning multiple weights, it is common to compute costs and updates for individual weights (or sets of weights) independently and add them together to arrive at a joint pruning cost. In LLM Surgeon, we argue that it’s better to consider weight updates jointly instead of independently. After selecting the set of weights for pruning, we can often afford to compute a single correlated weight update associated to the joint removal of multiple weights, instead of naively summing weight updates associated to individual removals. We derive such correlated weight updates below. Note that, for the expected cost computation, we do assume that the row, column or weight costs are independent, as the number of possible combinations of weights to prune grows too large to compute within reasonable time.

To improve the performance-to-sparsity ratio, we propose pruning in multiple shots. We theoretically justify this multi-shot approach by noting that the surrogate loss landscape $q$ relies on a Taylor expansion (eq. 3) that only holds locally and thus becomes unreliable for larger jumps $\Delta{\bm{\theta}}$ in parameter space. We mitigate this by pruning in multiple $T{>}1$ shots, $t\in[1,2,\ldots,T]$, each resulting in a smaller weight update $\Delta{\bm{\theta}}$ after which the curvature of the loss surface can be re-estimated. When pruning to target size $\alpha$, ie. removing $1{-}\alpha$ of total weights, we choose a schedule $\alpha_{t}$ starting at $\alpha_{0}=1$ and ends with $\alpha_{T}{=}\alpha$, such that after $T$ shots, exactly $\alpha$ fraction of the total weight remain. Empirically, we find that a linear schedule for $\alpha_{t}$, as formulated in section 4, monotonically improves pruning performance with more shots, and that higher sparsity levels typically require more shots (see section F.1). Multi-shot pruning allows one to spend (linearly in $T$) more computation to improve the final compression performance.

We propose optional interleaved low-rank first-order corrections to further improve compression performance. So far, we assumed parameters are in a local optimum when finding a closed-form solution to the quadratic constraint problem. In practice, however, this assumption likely does not hold since (i) the neural network may not be optimised to the minimum, (ii) a different loss may be used for compression than used for training, or (iii) we prune in multiple shots (section 3.5) inevitably causing weights to diverge from the optimum. To mitigate this, we consider first-order corrections by interleaving pruning shots with low-rank adaptations of weights ${\bm{W}}_{l}{+}{\bm{U}}{\bm{V}}$ (LoRA, by (Hu et al., 2021)), commonly used in LLM finetuning. We always absorb updates after each shot, so that the next loss estimate $q$ is closer to the optimum and underlying assumptions are likely to hold more closely. By absorbing LoRA updates between shots, the sum of low-rank updates can have a higher rank than individual updates. That is, we have $\text{rank}({\bm{U}}^{1}{\bm{V}}^{1}{+}{\bm{U}}^{2}{\bm{V}}^{2}{+}\ldots{+}{\bm{U}}^{T}{\bm{V}}^{T})\geq\text{rank}({\bm{U}}^{t}{\bm{V}}^{t})$ for the updates ${\bm{U}}^{t}{\bm{V}}^{t}$ at any shot $t$, with equality only arising if updates lie exactly in the same subspace which is unlikely to ever occur in practice. This insight could also be used during regular LoRA finetuning and may therefore be useful outside the context of model compression to allow more expressive low-rank model adaptation, at negligible cost.

Structured compression of rows and columns enables direct savings in memory and compute through a straight reduction of matrix dimensions in the model. For LLM surgeon, we consider in section 3.4 weight updates with different levels of correlations: limited to correlations within rows and columns, and correlations both within and between rows and columns. We further compare against magnitude pruning, which only uses weight magnitudes, L-OBD, which only uses activations, and K-OBD, which also uses Kronecker-factored curvature but assumes full independence and thus only prunes without updating remaining weights. We report results in table 1, and observe that more correlations results in better performance, with the largest improvements for the Llama-v2 model family.

While a 50% structured compression is not better than a smaller model of similar size, LLM Surgeon allows us to reduce model size by up to 30% with minimal loss, without training a smaller model from scratch fig. 1. In our structured compression experiments our proposed LLM Surgeon method outperforms all baselines and achieves the best performance for each compression target size.

Additionally, we assess compression performance in conjunction with the proposed first-order corrections using the interleaved low-rank adaptation described in section 3.6. We find that LoRA improves compression performance in the smallest 125m model, but not in larger models. We hypothesise that larger models are more prone to overfitting on the relatively few batches of wikitext-2 data used to compress the model. Nevertheless, we conclude that interleaved LoRA can be useful in cases, and recommend first using the proposed method without interleaved updates and, if enough data is available for compression, optionally using it if it improves performance.

For 2:4 semi-structured pruning, we compare LLM Surgeon with magnitude pruning, which only uses weight magnitudes, single-shot L-OBD, which only uses activations, and single-shot K-OBD, which also uses Kronecker-factored curvature but assumes full independence and thus only prunes without updating remaining weights as well as the recent state-of-the-art SparseGPT (Frantar & Alistarh, 2023). We report test performance after 50 % (2:4) semi-structured compression on wikitext-2 data in table 3. We empirically find that considering more weight correlations results in improved final performance after compression. Our proposed LLM Surgeon is competitive with prior work outperforming all baselines in terms of test set perplexity (PPL).