PaperlayerView exact revision on arXiv

Source-linked AI summary

When Pruning Meets Interpretability: Preserving Sparse Autoencoder Robustness in LLMs

Suchit Gupte, Xueru Zhang, Mohammad Mahdi Khalili

arXiv:2608.25941v1cs.LG

TL;DR

The paper examines when sparse autoencoders trained on dense LLMs remain reliable after post-hoc pruning without retraining. It develops a perturbation-energy framework and evaluates pruning methods across SAE metrics, finding that activation-aware pruning better preserves behavior while middle layers remain especially vulnerable.

  • Problem

    The paper addresses the limited understanding of whether SAEs trained on dense LLMs remain valid after pruning without retraining.

  • Method

    The paper combines perturbation-theoretic analysis with four-category SAEBench evaluation across pruning methods, model architectures, and layer locations.

  • Results

    Activation-aware Wanda and SparseGPT preserve SAE behavior substantially better than magnitude pruning, while middle layers show consistently greater sensitivity across methods and sparsity levels.

  • Takeaways & Limitations

    Pruning choices should account for activation geometry, and layer-wise sparsity allocation can preserve SAE quality while reducing perplexity at the same average sparsity.

  • Takeaways & Limitations

    The conclusions rely on four relatively homogeneous model–SAE pairs and unstructured pruning, leaving generalization to other architectures, tuning settings, outputs, and pruning structures unclear.

Abstract

from arXiv · show

Sparse autoencoders (SAEs) are widely used to interpret the internal representations of large language models (LLMs), yet their reliability under post-hoc model compression remains poorly understood. We present a systematic study of how pruning affects SAE behavior and theoretically show that, for a fixed SAE, its impact is governed by perturbation energy, a covariance-weighted norm. This perspective exposes a key limitation of magnitude pruning: by ignoring activation geometry, it distorts the learned representation space and degrades SAE functionality. Activation-aware methods such as Wanda and SparseGPT, in contrast, implicitly control perturbation energy and are therefore substantially more robust at preserving SAE behavior. We further reveal a consistent structural vulnerability across all pruning methods: middle layers are significantly more sensitive to pruning than early or late layers. Guided by this insight, we propose a layer-wise sparsity allocation strategy, achieving lower perplexity under the same average pruning sparsity. Experiments across four model architectures validate our theoretical findings. Code is publicly available at https://github.com/osu-srml/sae-robustness-under-pruning/tree/main.

1 Introduction

The paper asks when an SAE trained on a dense model remains valid after pruning without retraining. It develops a perturbation-energy account, compares pruning methods, and identifies middle-layer sensitivity as a basis for sparsity allocation.

  • Existing pruning evaluations emphasize predictive performance, while giving little attention to how pruning criteria perturb internal activation distributions.
  • Pruning can shift downstream activations away from an SAE’s training distribution, causing silent latent failures that standard language-modeling benchmarks may miss.
  • The paper studies whether an existing dense-model SAE transfers to a pruned model, avoiding the cost of retraining interpretability tooling for each model variant.
  • Magnitude, SparseGPT, and Wanda pruning are theoretically compared through perturbation energy, with magnitude pruning inducing the largest energy and degradation.
  • Across four SAEBench categories, middle layers are more pruning-sensitive than early or late layers, motivating layer-wise sparsity allocation with lower perplexity at equal average sparsity.

2 Preliminaries

The preliminaries define SAEs as overcomplete encoder–decoder maps trained on model activations, then formulate pruning as a perturbation to linear-layer outputs. They also state the Lipschitz assumption used for the analysis.

  • 2.1 Sparse autoencoders for mechanistic interpretability: A sparse autoencoder maps R^d to R^d through an encoder into R^m and decoder from R^m, where m ≫ d.
  • 2.1 Sparse autoencoders for mechanistic interpretability: SAEs are trained on activations sampled from a model’s forward pass to minimize a reconstruction objective.
  • 2.1 Sparse autoencoders for mechanistic interpretability: The analysis assumes the trained SAE is L-Lipschitz, an assumption satisfied by standard architectures using Lipschitz activations such as ReLU and GELU.
  • 2.2 Weight pruning as activation perturbation: Pruning zeros selected entries of a linear-layer weight matrix, producing W′ = W + ∆W and an activation shift δ = ∆W x_in.
  • 2.2 Weight pruning as activation perturbation: Magnitude, Wanda, and SparseGPT use distinct pruning objectives based respectively on weight magnitude, activation statistics, and second-order layer-wise regression.

3 Perturbation Theory: SAE Degradation under Pruning

The theory bounds fixed-SAE degradation using perturbation energy, a covariance-weighted measure of pruning-induced activation changes. It explains why activation-aware methods better preserve SAE behavior than magnitude pruning.

  • Perturbation-theoretic framework: The unified analysis bounds SAE degradation using perturbation energy together with the SAE’s Lipschitz constant.
  • Perturbation-theoretic framework: Theorem 3.2 decomposes expected reconstruction degradation into intrinsic SAE error, direct perturbation energy, and their interaction.
  • Perturbation-theoretic framework: For a fixed SAE, perturbation energy is the operative variable when comparing pruning methods because the Lipschitz constant remains constant.
  • Perturbation-theoretic framework: Perturbation energy is a covariance-weighted squared norm: changes along high-variance input directions contribute more than changes along low-variance directions.
  • Separation among pruning methods: Magnitude pruning ignores covariance structure, Wanda uses a diagonal approximation, and SparseGPT approximately minimizes the full perturbation energy through second-order information.

4 From Theory to Evaluation: SAEBench as a Sufficient Validation Suite

The paper uses four SAEBench metric categories as a joint validation suite for distinct SAE failure modes under activation-distribution shifts. Together, they test reconstruction, feature integrity, causal interventions, and task-level performance.

  • Joint Sufficiency: Core, Feature Absorption, SCR, and TPP jointly cover orthogonal failure modes, and omitting any category would leave a validation blind spot.
  • Failure Mode 1: Reconstruction Degrades: Core metrics test whether SAE reconstruction preserves model behavior when pruning shifts the activation distribution.
  • Failure Mode 2: Features Break: Feature Absorption detects latent false negatives and false positives that can leave aggregate reconstruction intact while breaking interpretability.
  • Failure Mode 3: Interventions Stop Working: Spurious Correlation Removal tests whether latent-level interventions retain their intended causal influence without collateral effects on unrelated concepts.

5 Experimental Setup

The study fixes pretrained SAEs and applies three post-hoc pruning methods across four model–SAE pairs, evaluating layer-level SAE degradation at 50% sparsity.

  • Models and SAEs: Four model–SAE pairs spanning roughly two orders of magnitude in parameter count are evaluated using fixed pretrained SAEs.Fixing SAE weights isolates degradation attributable to pruning rather than SAE training variability.
  • Pruning methods: MAGNITUDE, WANDA, and SPARSEGPT are applied post-hoc without retraining or weight updates after pruning.WANDA and SPARSEGPT use 128 OpenWebText calibration samples.
  • Pruning methods: Experiments use 50% sparsity, a practical operating point where all three pruning methods remain functional and SAE degradation differences are visible.This setting supports direct comparison among the pruning methods.
  • Evaluation: SAE behavior is measured with Core, Feature Absorption, SCR, and TPP metrics on every layer’s residual output.For SCR and TPP, k = 10 is used as the canonical ablation threshold.

6 Results

Across models, metrics, and layers, activation-aware pruning preserves SAE behavior better than MAGNITUDE pruning. Middle layers are consistently most sensitive, motivating non-uniform sparsity that lowers perplexity at equivalent average sparsity.

  • Pruning-method comparison: MAGNITUDE causes the most severe SAE degradation, while WANDA and SPARSEGPT remain substantially more robust across SAEBench metrics.SPARSEGPT closely tracks the dense baseline, and the ordering holds across all eight metrics and 26 gemma-2-2b layers.
  • Pruning-method comparison: MAGNITUDE degrades SCR by up to 60% and TPP by up to 80%, whereas WANDA and SPARSEGPT remain within 15–25% of baseline on the same metrics.The result aligns with perturbation-energy theory: activation-aware methods better approximate the covariance-weighted objective.
  • Pruning-method comparison: Aggregate Core metrics can understate MAGNITUDE’s damage: middle-layer KL divergence remains ≥0.96 while SCR and TPP collapse.Interventional metrics therefore reveal failures that aggregate reconstruction measures may miss.
  • Layer sensitivity: Middle-layer vulnerability reflects accumulated upstream perturbation through residual connections despite middle layers having the lowest local ε2.Protecting early layers is therefore expected to reduce perturbations cascading into vulnerable middle layers.
  • Layer sensitivity: Middle layers 9–17 show higher sensitivity than early layers 0–8 or late layers 18–25 at 25%, 40%, and 50% sparsity.At 50% sparsity, MAGNITUDE middle-layer degradation reaches 10–14% in explained variance, 5–7% in cosine similarity, and 40–60% in SCR and TPP.
  • Layer-wise sparsity allocation: The layer-wise schedule assigns lower sparsity to early layers and higher sparsity to late layers to preserve activations used by middle layers.This allocation is motivated by the depth-dependent sensitivity profile rather than uniform compression.
  • Layer-wise sparsity allocation: Lower perplexity is obtained with the layer-wise schedule than with uniform sparsity at equivalent average sparsity for both WANDA and SPARSEGPT.The authors describe this as a preliminary evaluation rather than a fully developed method.
  • Cross-architecture generalisation: Across pythia-70m, gemma-2-2b, gemma-2-9b, and mistral-7b, cross-architecture results preserve the same method ordering.For pythia-70m, MAGNITUDE lowers KL divergence from 0.926 to 0.642, versus 0.880 for WANDA and 0.882 for SPARSEGPT.

7 Conclusion

The paper explains fixed-SAE degradation after pruning through perturbation energy and identifies middle layers as structurally vulnerable. It also reports limitations concerning SAE retraining, model-family coverage, pruning scope, and preliminary sparsity allocation.

  • 7 Conclusion: For a fixed SAE, pruning-induced degradation is governed by perturbation energy, a covariance-weighted measure of activation-distribution distortion.The framework also accounts for the SAE’s Lipschitz constant.
  • 7 Conclusion: Activation-aware WANDA and SPARSEGPT better preserve SAE behavior than MAGNITUDE pruning by controlling perturbation energy.MAGNITUDE pruning ignores activation geometry and induces larger perturbation energy.
  • 7 Conclusion: Middle layers are structurally vulnerable, while layer-wise sparsity allocation preserves representations under equivalent compression.The allocation strategy protects sensitive middle layers more effectively.
  • 7 Conclusion: The study fixes SAE weights, evaluates four relatively homogeneous model–SAE pairs, and focuses on unstructured pruning.Generalization to retrained SAEs, other architectures, and structured compression remains open.
  • 7 Conclusion: The proposed layer-wise allocation is only a preliminary scheduling illustration, not a fully benchmarked allocation method.It is not compared against dedicated or more non-intuitive allocation methods.

A.1 Detailed Proofs for the Perturbation-Theoretic Framework

The appendix supplies fully detailed, self-contained proofs for the perturbation theory developed in Section 3.

  • A.1 Detailed Proofs for the Perturbation-Theoretic Framework: The appendix contains detailed proofs of the perturbation theory.These proofs correspond to the framework developed in Section 3.
  • A.1 Detailed Proofs for the Perturbation-Theoretic Framework: The proofs are presented as self-contained arguments.The appendix is intended to stand independently for the theory’s derivations.
  • A.1 Detailed Proofs for the Perturbation-Theoretic Framework: The appendix’s proof material concerns the perturbation theory presented in Section 3.It does not introduce a separate theoretical framework in the supplied passage.

A.1.1 Setup and Notation

The setup models pruning as a weight perturbation in a linear layer and studies how a fixed Lipschitz SAE reconstructs the resulting activations through perturbation energy.

  • A.1.1 Setup and Notation: The layer input x_in is a random vector drawn from the data distribution D.It enters a linear layer with weight matrix W.
  • A.1.1 Setup and Notation: Pruning changes W by zeroing selected entries, producing a weight change ΔW whose nonzero entries equal the negatives of the removed weights.For every pruned entry, (ΔW)_ij = −W_ij.
  • A.1.1 Setup and Notation: The pruned activation is x′ = (W + ΔW)x_in = x + δ, where δ = ΔW x_in is the perturbation vector.Thus pruning is represented as an activation shift from x to x′.
  • A.1.1 Setup and Notation: The fixed SAE f_θ maps R^d to R^d and is assumed to be L-Lipschitz.The analysis bounds reconstruction error on pruned activations using perturbation energy.

A.1.2 Proof of Lemma 3.1

Lemma 3.1 relates reconstruction error on pruned activations to error on dense activations and the activation perturbation. The bound uses the triangle inequality and the SAE’s Lipschitz property.

  • A.1.2 Proof of Lemma 3.1: Lemma 3.1 considers a fixed L-Lipschitz SAE and pruned activation x′ = x + δ.The proof targets reconstruction error at the perturbed activation.
  • A.1.2 Proof of Lemma 3.1: The proof applies the triangle inequality to insert the dense-model reconstruction f_θ(x) between x′ and f_θ(x′).This separates baseline reconstruction error from perturbation-dependent terms.
  • A.1.2 Proof of Lemma 3.1: The activation difference contributes a term bounded by the perturbation norm, while the SAE-output difference is bounded using Lipschitz continuity.The latter is at most L∥δ∥ when the input changes by δ.
  • A.1.2 Proof of Lemma 3.1: The factor 1 + L combines the direct activation shift with the SAE’s response to that shift.Both contributions add to the reconstruction-error bound.

Proof of Theorem 3.2 (Expected Bound and Covariance Decomposition)

The theorem bounds SAE degradation after pruning using intrinsic reconstruction error, perturbation energy, and their interaction, then decomposes perturbation energy according to activation covariance. This framework distinguishes pruning methods and motivates evaluating their effects across SAE failure modes and layer sensitivities.

  • Covariance decomposition: The perturbation energy is ε2 = tr(∆W Σxin ∆W⊤), weighting perturbations according to the covariance of layer inputs.In the covariance eigenbasis, each contribution is λi∥∆W vi∥2, so high-variance directions receive greater penalty.
  • Pruning-method comparison: MAGNITUDE, WANDA, and SPARSEGPT form a hierarchy of approximations: none, diagonal, and second-order/full covariance information, respectively.MAGNITUDE matches perturbation energy only under identity-proportional input covariance, whereas WANDA and SPARSEGPT account for activation structure more closely.
  • Evaluation: The evaluation combines Core, Feature Absorption, SCR, and TPP metrics to cover distinct SAE failure modes under shifted activation distributions.SCR and TPP are evaluated across seven ablation thresholds, with k treated as an evaluation hyperparameter rather than a model or pruning-method parameter.
  • Empirical implications: Across methods, MAGNITUDE produces the largest SAE degradation, while WANDA and SPARSEGPT remain closer to the dense baseline; middle layers are most vulnerable despite lower local ε2.The reported interpretation attributes middle-layer vulnerability to upstream perturbations propagating through residual connections, motivating protection of early layers.
Loading 2608.25941v1…