PaperlayerView exact revision on arXiv

Source-linked AI summary

Output-aware Residual Stream Pruning for Large Language Models

Chayne Thrash, Kevin Chen, Soheil Kolouri

arXiv:2609.35579v1cs.LG

TL;DR

Activation-only residual-stream pruning overlooks that equally sized perturbations can affect model outputs differently. This paper combines activation covariance and output sensitivity in an efficient spectral approach, improving output-distribution preservation and downstream performance over activation-only pruning across evaluated models and compression rates.

  • Problem

    Activation reconstruction error alone treats perturbation directions alike, although downstream computations can respond differently to perturbations of similar magnitude.

    Source: §1 · p. 2

  • Method

    A second-order output-KL approximation combines activation covariance with local output sensitivity, and a tractable spectral upper bound generates candidate pruning subspaces.

    Source: §1 · p. 2, §4 · p. 4, §4.1 · p. 5

  • Results

    The method achieves lower KL divergence and perplexity than SliceGPT across every evaluated model and compression rate, with higher average commonsense reasoning accuracy in almost all settings.

    Source: §5 · p. 8

  • Takeaways & Limitations

    Across evaluated instruction-tuned models and settings, sensitivity-aware bases preserve model quality more effectively than covariance-only residual-stream pruning, particularly as sparsity increases.

    Source: §6 · p. 10

    The layerwise approximation neglects both cross-token and cross-layer curvature interactions.

    Source: §4.1 · p. 5

Abstract

from arXiv · show

Residual stream pruning methods reduce inference cost by shrinking the model's hidden dimension, but existing approaches typically choose these dimensions by minimizing activation reconstruction error. This criterion implicitly treats all perturbation directions as equally important, ignoring the sensitivity of downstream layers. We introduce a sensitivity-aware approach to residual-stream pruning that directly accounts for this direction-dependent sensitivity. Using a second-order approximation to the output KL divergence, we characterize the effect of a residual-stream perturbation through both its activation covariance and the local sensitivity of the model output. The resulting subspace selection objective couples these two quantities, but is difficult to optimize directly. We derive a tractable spectral upper bound that reduces subspace selection to an eigendecomposition of a sensitivity-weighted covariance matrix, retaining the efficiency and structural simplicity of rotation-based pruning methods. Across several instruction-tuned language model families, our method consistently reduces calibration KL divergence relative to activation-only pruning and improves perplexity and downstream task performance over a range of compression levels. Our results show that preserving activation energy alone is insufficient for residual-stream pruning, and that explicitly accounting for how perturbations propagate to the model output provides a more effective criterion for selecting dimensions to remove.

1 INTRODUCTION

Growing LLMs make deployment increasingly expensive, motivating compression that retains strong performance. Residual-stream pruning reduces model dimensions, and this work selects removable directions using both activation error and output sensitivity.

1 INTRODUCTION

  • LLM growth has made deployment increasingly expensive, motivating smaller, more efficient models that retain strong performance on complex tasks.Structured pruning can reduce model size and inference cost without customized kernels or costly retraining.Source: §1 · p. 1
  • Unlike pruning neurons or attention heads, residual-stream pruning reduces the input and output dimensions of transformer blocks.It exploits orthogonal-basis invariance to remove residual directions, but activation reconstruction alone overlooks that equally sized perturbations can affect output distributions differently.Source: §1 · pp. 1–2
  • Sensitivity-aware pruning balances induced activation error with output sensitivity and consistently improves output-distribution preservation and downstream performance over activation-only residual-stream pruning.The method derives from a second-order approximation to output KL and uses a tractable upper bound solved with a small number of eigendecompositions.Source: §1 · p. 2

2 RELATED WORK

Prior work compresses LLMs through unstructured or structured pruning, including methods that remove residual-stream dimensions. Output-aware compression also considers model behavior, and this approach combines output sensitivity with activation statistics for direction selection.

2 RELATED WORK

  • Unstructured methods such as SparseGPT and Wanda generally need specialized kernels or hardware for inference speedups, while structured pruning removes components to reduce computation directly.Structured examples include LLM-Pruner's first-order importance estimates and LLM-Surgeon's Kronecker-factored empirical Fisher.Source: §2 · p. 2
  • SliceGPT (Ashkboos et al., 2024) uses rotational invariance and PCA to remove low-variance residual-stream directions, yielding smaller dense weights compatible with standard hardware.Figure 1 depicts rotating unimportant directions into removable coordinates and selecting with activation variance and output sensitivity, balanced through τ.Source: §2 · p. 2, Figure 1 · p. 3
  • Output-aware compression includes first-order Taylor and empirical-Fisher approaches, while YAQA, EvoPress, and T´yr-the-pruner optimize objectives tied to predictive-distribution changes.The paper identifies YAQA as its closest relation, combining second-order output-sensitivity information with activation statistics to select residual-stream directions.Source: §2 · p. 2

3 METHOD

The method builds on rotation-based residual-stream pruning and selects removable directions using an efficient approximation to an objective that accounts for activation reconstruction and output sensitivity.

3 METHOD

  • The method derives an objective incorporating activation reconstruction and output sensitivity, then approximates it with candidate bases generated using a small number of eigendecompositions.Figure 2 describes validation through sensitivity–KL correlation and comparison of τ-selected bases with SliceGPT across layers at 25% pruning.Source: §3 · p. 3, Figure 2 · p. 4

3.1 RESIDUAL-STREAM PRUNING

  • Orthogonal transformations can be absorbed into linear layers along the residual stream, allowing rotations to identify directions for removal while retaining others.The retained basis V and removed basis U define the directions kept and pruned.Source: §3.1 · p. 3
  • Removing U-coordinates leaves the reduced representation V ⊤x and, after projections are absorbed into adjacent layers, produces smaller dense weight matrices.Source: §3.1 · p. 3
  • Layer-dependent orthogonal bases require transition matrices along the residual connection, adding storage and computation whose cost decreases quadratically with retained width.Source: §3.1 · p. 3

4 DIRECTION SELECTION IN SLICEGPT

SliceGPT chooses removed residual-stream directions by activation energy, whereas this section develops an output-sensitivity-aware criterion and an eigendecomposition-based approximation. The sensitivity signal comes from local output KL curvature, and the higher-dimensional construction minimizes a tractable upper bound.

4 DIRECTION SELECTION IN SLICEGPT

  • SliceGPT selects removed directions by minimizing discarded activation energy, equivalent to uncentered PCA retaining the leading principal directions; this ignores direction-dependent downstream sensitivity.The method uses residual-stream activations from calibration sequences to form the activation second-moment matrix.Source: §4 · p. 4

4.1 SENSITIVITY-AWARE DIRECTION SELECTION

  • The proposed criterion scores removed directions by their effect on model output, using a local quadratic approximation to KL divergence based on activation perturbations and their curvature.The expansion assumes twice differentiability at the unperturbed activation; under usual regularity conditions, the curvature is the Fisher information with respect to that activation.Source: §4.1 · pp. 4–5
  • A shared layerwise curvature H avoids token-specific curvature storage but discards correlations between activation outer products and their corresponding curvature matrices.The layerwise treatment also neglects cross-token and cross-layer curvature interactions.Source: §4.1 · p. 5
  • The objective stores layerwise C and H, and when H = I it exactly recovers SliceGPT’s objective; with activation reconstruction error fixed, larger curvature-assigned error generally accompanies larger KL divergence.The passage reports this directional comparison as evidence that H captures output sensitivity not distinguished by reconstruction error alone.Source: §4.1 · p. 5

4.2 EFFICIENT APPROXIMATION OF THE OBJECTIVE AND ITS THEORETICAL ANALYSIS

  • The rank-one objective has an exact spectral characterization, while the higher-dimensional construction selects eigenvectors minimizing an upper bound and thus provides an efficient approximation.The higher-dimensional bound is interpreted as a regularized form of the desired loss; directly optimizing the objective does not reduce to a standard eigenspace problem.Source: §4.1 · p. 5, §4.2 · p. 6
  • For different values of τ, candidate subspaces attain different original-objective values, and an appropriate τ yields lower L(U) than the SliceGPT solution based only on C.The passage reports this validation in Figure 2b.Source: §4.2 · p. 7

5 EXPERIMENTS

Across instruction-tuned models and task types, experiments compare sensitivity-aware pruning with SliceGPT and report stronger performance preservation across compression settings, alongside calibration-size, optimization, and runtime analyses.

5 EXPERIMENTS

  • The evaluation covers five instruction-tuned models from three families, language-modeling, common-sense and generative tasks, and calibration using 1,024 sequences or examples.Sensitivity for generative tasks is computed over assistant-response tokens.Source: §5 · pp. 7–8

Language modeling and common-sense reasoning results

  • At 30% sparsity, Llama-3.2-3B PPL falls from 57.17 to 30.14 and average accuracy rises from 47.40 to 49.43 versus SliceGPT; gains extend to math and instruction following.The reported improvements span language modeling, common-sense reasoning, mathematical reasoning, and instruction following, and generally grow at higher sparsity.Source: §5 · p. 8

Effect of calibration set size

  • Across calibration-set sizes, the proposed method outperforms SliceGPT; gains from more data are largest between 128 and 512–1,024 sequences, with small, occasionally non-monotonic gains beyond 1,024.Figure 3 compares average common-sense reasoning performance on Llama-3.1 8B Instruct using varying calibration-sample counts.Source: §5 · pp. 8–9, Figure 3 · p. 8

Effect of additional optimization

  • With 10,000 optimization steps per pruning site, additional optimization yields only small calibration-KL reductions and modest accuracy improvements.This supports the spectral candidate selection as already strong, with limited benefit from further minimizing the calibration objective.Source: §5 · p. 9, Table 3 · p. 9

Calibration time comparison

  • On Llama-3.1-8B, sensitivity estimation takes 9 minutes, or 12% of total calibration time; overall calibration time rises only modestly.Direction selection dominates runtime, largely because of double-precision activation-covariance accumulation; additional eigendecompositions add comparatively little overhead.Source: §5 · p. 9, Table 4 · p. 9

Comparison with additional structured pruning methods

  • Residual-stream pruning remains competitive with structured methods across evaluated sparsity rates, but the comparison is not parameter-equivalent because reported rates exclude rotation parameters.The results are presented as context for comparing pruning paradigms rather than as a direct comparison.Source: Table 5 · p. 9, §5 · p. 10

6 CONCLUSION

The paper presents sensitivity-aware residual-stream pruning that accounts for how removed activation directions affect the output distribution. Across multiple instruction-tuned language models and evaluation settings, its bases preserve model quality more effectively than covariance-only pruning, particularly as sparsity increases.

  • Sensitivity-aware pruning combines activation covariance with output sensitivity to account for how removed directions affect the model output distribution.Source: §6 · p. 10
  • Across multiple instruction-tuned language models and evaluation settings, the resulting bases preserve model quality more effectively than covariance-only pruning, particularly as sparsity increases.Source: §6 · p. 10
  • The method uses an efficient approximation requiring only a few eigendecompositions to construct pruning bases.Source: §6 · p. 10
  • Future work includes exploring shared or more parameter-efficient transformations.Source: §6 · p. 10

A DERIVATIONS

The appendix connects activation-only pruning to PCA, derives spectral solutions and upper bounds for sensitivity-aware subspace selection, and reviews the orthogonal reparameterization that makes residual-stream rotations possible.

A.1 SLICEGPT EQUIVALENCE TO EQUATION 1

  • SliceGPT’s activation reconstruction objective selects the removed subspace through the smallest eigenvalues of the uncentered activation covariance, equivalently retaining its leading principal directions.Source: §A.1 · p. 14

A.2 DERIVATION OF PROPOSITION 1

  • For rank-one pruning, a smallest-eigenvalue eigenvector of τC + τ^-1H globally minimizes the objective when C and H are positive definite and the optimal scale is attained.Source: §A.2 · p. 15

A.3 DERIVATION OF PROPOSITION 2

  • For fixed τ, the upper bound is minimized by choosing the removed directions as eigenvectors corresponding to the k smallest eigenvalues of Mτ.Source: §A.3 · p. 15

A.4 DERIVATION OF PROPOSITION 3

  • Joint optimization of the spectral upper bound over the subspace and scale yields an equivalent regularized problem with a bottom eigenspace as a minimizer.Source: §A.4 · p. 16

B ORTHOGONAL REPARAMETERIZATION INVARIANCE OF LLMS

  • Orthogonal reparameterization preserves transformer-block computations by transforming hidden representations and adjacent linear weights, while LayerNorm requires conversion to an equivalent RMSNorm-connected form.Source: B ORTHOGONAL REPARAMETERIZATION… · p. 17

C IMPLEMENTATION DETAILS

Implementation uses sequential pruning with calibration activations and estimates output sensitivity efficiently; evaluation protocols are specified for reasoning and instruction-following tasks.

  • Pruning bases are constructed sequentially, with activations propagated through the partially pruned model before computing each next basis.Source: C IMPLEMENTATION DETAILS · p. 17
  • Sensitivity matrices use one sampled output token per token and a single backward pass, while covariance matrices and eigendecompositions use double precision and C and H are Frobenius-normalized.Source: C IMPLEMENTATION DETAILS · p. 17
  • Common-sense tasks use zero-shot normalized accuracy when available, GSM8K uses 8-shot chain-of-thought prompting, and IFEval is zero-shot; each uses the model’s chat template.Source: C IMPLEMENTATION DETAILS · p. 17

D PER-TASK COMMON-SENSE REASONING RESULTS

The section reports per-task common-sense reasoning results for five instruction-tuned models, which are used to produce the average accuracies in Table 1.

E RESULTS WITH REDPAJAMA CALIBRATION

The method is evaluated with RedPajama calibration on instruction-tuned models and against SliceGPT on non-instruction-tuned models. On the non-instruction-tuned models, it achieves lower KL divergence and perplexity and higher average common-sense reasoning accuracy across every model and sparsity level.

E RESULTS WITH REDPAJAMA CALIBRATION

F RESULTS ON NON-INSTRUCTION TUNED MODELS

  • Across every tested model and sparsity level, the method achieves lower KL divergence and perplexity than SliceGPT and higher average common-sense reasoning accuracy.The comparison covers base Llama-2 7B, Llama-3.2 3B, and Llama-3.1 8B.Source: F RESULTS ON NON-INSTRUCTION TUNED… · p. 19

G ANALYSIS OF DATASET-AVERAGE CURVATURE

The dataset-averaged curvature matrix largely preserves relative sensitivity across pruning-relevant directions, although it can miss token-specific variation when sensitive directions vary substantially across tokens.

Related papers

Loading 2609.35579v1…