The LLM Surgeon

TL;DR

The paper tackles the challenge of deploying large language models under resource constraints by introducing LLM Surgeon, a data-driven pruning framework that uses Kronecker-factored curvature approximations of the empirical Fisher to prune unstructured, semi-structured, and structured patterns. It extends optimal pruning theory by deriving closed-form pruning costs and correlated weight updates, enabling global-threshold, multi-shot pruning with optional low-rank corrections. Empirical results on OPT and Llama-v2 show state-of-the-art performance for unstructured and semi-structured pruning, and demonstrate practical structured pruning of up to 30% with minimal loss, along with learned sparsity structures that concentrate pruning on early layers and FC blocks. The approach provides a principled, scalable path to compact, deployable LLMs without retraining from scratch, balancing accuracy and compute through configurable correlations and shot schedules.

Abstract

State-of-the-art language models are becoming increasingly large in an effort to achieve the highest performance on large corpora of available textual data. However, the sheer size of the Transformer architectures makes it difficult to deploy models within computational, environmental or device-specific constraints. We explore data-driven compression of existing pretrained models as an alternative to training smaller models from scratch. To do so, we scale Kronecker-factored curvature approximations of the target loss landscape to large language models. In doing so, we can compute both the dynamic allocation of structures that can be removed as well as updates of remaining weights that account for the removal. We provide a general framework for unstructured, semi-structured and structured pruning and improve upon weight updates to capture more correlations between weights, while remaining computationally efficient. Experimentally, our method can prune rows and columns from a range of OPT models and Llamav2-7B by 20%-30%, with a negligible loss in performance, and achieve state-of-the-art results in unstructured and semi-structured pruning of large language models.

Figures (5)

• Figure 1: LLM Surgeon allows interpolation of model size between existing pretrained models.
• Figure 2: Pruning as equality constrained optimization of quadratic approximation of the loss landscape (left), or equivalently, maximising the likelihood under a Laplace approximation (right).
• Figure 3: General framework for structured, semi-structured and unstructured compression.
• Figure 4: Sparsity levels obtained with structured pruning on OPT-125m by layer depth and type.
• Figure 5: Example illustration of nearest Kronecker factor approximations $\widetilde{{\bm{F}}} {\approx} \sum_{r=1}^{R_K} {\bm{G}}_i \otimes {\bm{A}}_i$, compared to classical KFAC with the IAD assumption. Larger $R_K$ yields better approximations to the true Fisher ${\bm{F}}$ for larger $R_K$, as measured by the root mean squared error (rmse).