Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Free Lunch in LLM Compression: Revisiting Retraining after Pruning

Moritz Wagner, Christophe Roux, Max Zimmer, Sebastian Pokutta

TL;DR

This paper reevaluates post-pruning recovery for large language models by emphasizing local reconstruction over full retraining. It systematically analyzes propagation strategies, loss functions, and reconstruction granularity, finding that reconstructing attention and MLP components within each transformer block (block size 1/2) offers a Pareto-optimal balance of accuracy and memory usage. Surprisingly, per-matrix reconstruction underperforms, while full retraining often yields worse perplexity than locality-aware reconstruction; simple pruning methods like Wanda can even surpass complex approaches like SparseGPT when reconstruction is applied properly. Together, these results propose a practical, memory-efficient recipe for post-pruning recovery that challenges the prevailing wisdom of avoiding retraining in LLM compression.

Abstract

While Neural Network pruning typically requires retraining the model to recover pruning-induced performance degradation, state-of-the-art Large Language Models (LLMs) pruning methods instead solve a layer-wise mask selection and reconstruction problem on a small set of calibration data to avoid full retraining, as it is considered computationally infeasible for LLMs. Reconstructing single matrices in isolation has favorable properties, such as convexity of the objective and significantly reduced memory requirements compared to full retraining. In practice, however, reconstruction is often implemented at coarser granularities, e.g., reconstructing a whole transformer block against its dense activations instead of a single matrix. In this work, we study the key design choices when reconstructing or retraining the remaining weights after pruning. We conduct an extensive computational study on state-of-the-art GPT architectures, and report several surprising findings that challenge common intuitions about retraining after pruning. In particular, we observe a free lunch scenario: reconstructing attention and MLP components separately within each transformer block is nearly the most resource-efficient yet achieves the best perplexity. Most importantly, this Pareto-optimal setup achieves better performance than full retraining, despite requiring only a fraction of the memory. Furthermore, we demonstrate that simple and efficient pruning criteria such as Wanda can outperform much more complex approaches when the reconstruction step is properly executed, highlighting its importance. Our findings challenge the narrative that retraining should be avoided at all costs and provide important insights into post-pruning performance recovery for LLMs.

A Free Lunch in LLM Compression: Revisiting Retraining after Pruning

TL;DR

This paper reevaluates post-pruning recovery for large language models by emphasizing local reconstruction over full retraining. It systematically analyzes propagation strategies, loss functions, and reconstruction granularity, finding that reconstructing attention and MLP components within each transformer block (block size 1/2) offers a Pareto-optimal balance of accuracy and memory usage. Surprisingly, per-matrix reconstruction underperforms, while full retraining often yields worse perplexity than locality-aware reconstruction; simple pruning methods like Wanda can even surpass complex approaches like SparseGPT when reconstruction is applied properly. Together, these results propose a practical, memory-efficient recipe for post-pruning recovery that challenges the prevailing wisdom of avoiding retraining in LLM compression.

Abstract

While Neural Network pruning typically requires retraining the model to recover pruning-induced performance degradation, state-of-the-art Large Language Models (LLMs) pruning methods instead solve a layer-wise mask selection and reconstruction problem on a small set of calibration data to avoid full retraining, as it is considered computationally infeasible for LLMs. Reconstructing single matrices in isolation has favorable properties, such as convexity of the objective and significantly reduced memory requirements compared to full retraining. In practice, however, reconstruction is often implemented at coarser granularities, e.g., reconstructing a whole transformer block against its dense activations instead of a single matrix. In this work, we study the key design choices when reconstructing or retraining the remaining weights after pruning. We conduct an extensive computational study on state-of-the-art GPT architectures, and report several surprising findings that challenge common intuitions about retraining after pruning. In particular, we observe a free lunch scenario: reconstructing attention and MLP components separately within each transformer block is nearly the most resource-efficient yet achieves the best perplexity. Most importantly, this Pareto-optimal setup achieves better performance than full retraining, despite requiring only a fraction of the memory. Furthermore, we demonstrate that simple and efficient pruning criteria such as Wanda can outperform much more complex approaches when the reconstruction step is properly executed, highlighting its importance. Our findings challenge the narrative that retraining should be avoided at all costs and provide important insights into post-pruning performance recovery for LLMs.
Paper Structure (11 sections, 3 equations, 4 figures, 9 tables)

This paper contains 11 sections, 3 equations, 4 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Background and Methodology
  4. Background and related work
  5. Key design choices when reconstructing after pruning
  6. Experiments
  7. Propagation strategies and loss functions have no systematic impact.
  8. Granularity matters: challenging the avoid-retraining-at-all-costs narrative
  9. Reconstruction enables simple pruning methods to outperform complex ones
  10. Conclusion
  11. Additional Experiments

Figures (4)

  • Figure 1: Perplexity and memory usage of four different models pruned to 2:4 sparsity with Wanda and reconstructed with different granularities. The number of calibration samples is 1024. $\frac{1}{4}$ of the model is 6, 8, 8, and 10 transformer blocks for OPT-1.3B, OPT-6.7B, LLaMA-3-8B, and LLaMA-2-13B, respectively. Displayed are the best results propagation strategies and loss functions.
  • Figure 2: Histograms of recovery differences (\ref{['eq:recovery']}) between methods. Orange bars compare methods under identical hyperparameter configurations. Green bars compare each method's optimal setting for a given model and sparsity. The data used to generate these plots is shown in \ref{['tab:ppl_opt_1_3b']}, \ref{['tab:ppl_opt_6_7b']}, \ref{['tab:LLaMA-2_13b']}, and \ref{['tab:LLaMA-3_8b']} in \ref{['appendix:additional_experiments']}.
  • Figure 3: OPT-1.3B pruned to 50% unstructured sparsity with Wanda, showing mean perplexity vs. number of calibration samples (left) and number of epochs (right). The shaded areas indicate the min-max range over random seeds. Reconstruction uses MSE loss and DP. Full fine-tuning uses cross-entropy loss and true labels.
  • Figure 4: Learning rate ablation for reconstruction after pruning with Wanda to 50% unstructured sparsity (left) and 2:4 semi-structured sparsity (right). The $y$-axis is the recovery defined in \ref{['eq:recovery']}. Reconstruction was performed with block size $\frac{1}{2}$, MP, and MSE loss.