PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Dictionary Learning with Gated Sparse Autoencoders

Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Tom Lieberum, Vikrant Varma, János Kramár, Rohin Shah, Neel Nanda

arXiv:2404.16014v2cs.LGcs.AI

TL;DR

Sparse autoencoders offer sparse, linear decompositions for discovering interpretable language-model features, but their L1 sparsity penalty can shrink reconstructed activations. The paper introduces Gated SAEs, which separate feature detection from magnitude estimation and apply sparsity pressure only to detection. Across models and sites, Gated SAEs improve the sparsity–fidelity trade-off, solve shrinkage, and remain comparably interpretable.

  • Problem

    The L1 penalty used to encourage SAE sparsity introduces reconstruction biases, including shrinkage, that harm reconstruction accuracy.

  • Method

    Gated SAEs separate determining which dictionary directions are active from estimating their magnitudes, applying the sparsity penalty only to the detection function.

  • Results

    Gated SAEs Pareto improve sparsity and reconstruction fidelity across multiple models and activation sites, while remaining comparably interpretable to baseline SAE features.

  • Takeaways & Limitations

    At many sites, Gated SAEs require half the L0 to achieve the same loss recovered as baseline methods.

  • Takeaways & Limitations

    The approach assumes sparsity and linearity of language-model computation, and it complicates SAE training with a more complex encoder.

Abstract

from arXiv · show

Recent work has found that sparse autoencoders (SAEs) are an effective technique for unsupervised discovery of interpretable features in language models' (LMs) activations, by finding sparse, linear reconstructions of LM activations. We introduce the Gated Sparse Autoencoder (Gated SAE), which achieves a Pareto improvement over training with prevailing methods. In SAEs, the L1 penalty used to encourage sparsity introduces many undesirable biases, such as shrinkage -- systematic underestimation of feature activations. The key insight of Gated SAEs is to separate the functionality of (a) determining which directions to use and (b) estimating the magnitudes of those directions: this enables us to apply the L1 penalty only to the former, limiting the scope of undesirable side effects. Through training SAEs on LMs of up to 7B parameters we find that, in typical hyper-parameter ranges, Gated SAEs solve shrinkage, are similarly interpretable, and require half as many firing features to achieve comparable reconstruction fidelity.

1. Introduction

Sparse autoencoders help discover sparse decompositions of language-model activations, but prevailing L1 training introduces reconstruction bias. Gated SAEs separate feature detection from magnitude estimation and improve the sparsity–fidelity trade-off across models and activation sites.

  • Motivation: Sparse autoencoders seek interpretable, sparse decompositions of language-model activations into an overcomplete set of feature directions.This direction is motivated by superposition, in which many concepts are represented as linear directions but only a sparse subset is active for each input.
  • Motivation: The prevailing L1 penalty can trade reconstruction accuracy for lower sparsity loss, introducing biases that harm reconstruction accuracy.The paper identifies shrinkage as a clear consequence of this bias.
  • Approach: Gated SAEs use separate affine transformations to detect active dictionary elements and estimate their activation magnitudes, applying the sparsity penalty only to detection.Weight sharing limits increases in parameter count and inference-time compute relative to an equivalent-width baseline.
  • Results: Across GELU-1L, Pythia-2.8B, and Gemma-7B, Gated SAEs Pareto improve sparsity and reconstruction fidelity over baseline SAEs at fixed training compute.The evaluation covers MLP outputs, attention outputs, and residual-stream activations.
  • Results: Gated SAE features are comparably interpretable to baseline SAE features in a small double-blind study.A Mann-Whitney U test fails to reject equal overall interpretability distributions, with p = .84.
  • Caveat: Gated SAE training requires 50% more compute than a matched-width baseline for computing the training loss because of an auxiliary loss term.This caveat concerns training-loss computation rather than the equivalent-width inference-time comparison.

2. Sparse Autoencoder Background

Sparse autoencoders encode model activations into sparse feature vectors and decode them into reconstructions. Their training balances squared reconstruction error against an L1 sparsity penalty, while L0 and loss recovered evaluate sparsity and fidelity.

  • Representation: A sparse autoencoder represents an activation as a sparse linear combination of M≫n learned, unit-norm feature directions.The encoder produces sparse nonnegative feature activations, which the decoder linearly combines to reconstruct the input.
  • Architecture: The baseline SAE is a single-layer autoencoder trained to reconstruct model activations while constraining its hidden representation to be sparse.The decoder weights correspond to dictionary directions, and the latent representation supplies feature activations.
  • Training objective: Training jointly minimizes squared reconstruction distance and an L1 norm penalty on active features, balanced by coefficient λ.The reconstruction term encourages fidelity, while the L1 term encourages sparsity.
  • Training bias: L1 regularization can bias reconstructions toward lower feature magnitudes, producing shrinkage even when perfect reconstruction is possible.Figure 2 gives a single-feature example that reconstructs 1/2 rather than 1 when λ = 1.
  • Metrics: SAE sparsity is evaluated with L0, the average number of active features per input.L0 is used for evaluation because it is not differentiable and cannot be directly optimized.
  • Metrics: Loss recovered measures language-model reconstruction fidelity after splicing SAE reconstructions into the model, ranging from 0% for zero output to 100% for perfect reconstruction.It is calculated from average cross-entropy loss on an evaluation dataset.

3. Gated SAEs

Gated SAEs separate feature detection from magnitude estimation, applying sparsity pressure to detection while reducing shrinkage in activation magnitudes. Weight sharing controls parameter growth, and the tied architecture admits a Jump ReLU interpretation.

  • Motivation: L1 regularization can trade reconstruction fidelity for sparsity, systematically underestimating feature activation magnitudes through shrinkage.The sparsity penalty pushes activations toward zero while reconstruction pushes them high enough for accuracy.
  • Motivation: The baseline encoder both detects active features and estimates their magnitudes, although L1 is useful for the former and biases the latter.Separating these functions narrows the parameters directly affected by the sparsity penalty.
  • Architecture: Gated SAEs use separate affine transformations for feature detection and magnitude estimation, with sparsity applied only to the detection pathway.The gating pathway determines active features, while the magnitude pathway estimates coefficients for features that pass the gate.
  • Architecture: Weight sharing prevents the two encoder pathways from doubling encoder parameters by sharing projection directions while allowing distinct norms and biases.The paper reports that tied weights add only 2 × M parameters relative to a baseline SAE and slightly improve performance in an ablation.
  • Limitations: A simple rescaling of shrunk activations may not remove L1-induced bias because training can learn sub-optimal encoder and decoder directions.The paper states that Section 5.2 and Figure 11 provide empirical evidence for this limitation.
  • Architecture: With tied weights, the gated encoder is equivalent to a single-layer linear encoder with a discontinuous Jump ReLU activation.The paper presents this reinterpretation as an intuition for how Gated SAEs reconstruct activations.

4. Evaluation

Across models and activation sites, Gated SAEs improve the sparsity–reconstruction trade-off, remove shrinkage, and remain comparably interpretable to baseline SAEs.

  • 4.1. Comprehensive Benchmarking: Gated SAEs are evaluated on GELU-1L, Pythia-2.8B, and Gemma-7B activations across MLP outputs, attention outputs, and residual streams.
  • 4.1. Comprehensive Benchmarking: Gated SAEs Pareto improve baseline SAEs, yielding sparser decompositions at any desired level of reconstruction fidelity under matched training compute.Baseline comparisons use 50% more learned features to account for Gated SAE training costs of at most 1.5×.
  • 4.1. Comprehensive Benchmarking: Figure 5 compares loss recovered against L0 and reports better reconstruction fidelity for Gated SAEs at every displayed sparsity level.
  • 4.2. Shrinkage: Gated SAE reconstructions are unbiased with γ≈1, whereas baseline SAEs show shrinkage with γ<1 that worsens as λ increases and L0 decreases.The relative reconstruction bias uses γ=1 for unbiased reconstructions and γ<1 for shrinkage.
  • 4.3. Statistical Analysis: The interpretability comparison did not reject equal distributions across models or equal overall interpretability, with p=.95 and p=.84 respectively.
  • 4.3. Statistical Analysis: A blinded paired-rater study across Pythia-2.8B and Gemma-7B found Gated features at least comparable in interpretability, without conclusively showing superiority.The study rated features as Yes, No, or Maybe and paired observations by model, layer, site, and rater.

5. Why do Gated SAEs improve SAE training?

Ablations show that Gated SAE gains depend on limiting the auxiliary L1 penalty, retaining the rmag parameter, and tying encoder directions. Comparisons with shrinkage-only correction indicate that improved learned directions, not shrinkage correction alone, account for much of the performance advantage.

  • Ablation study: Freezing the decoder during the auxiliary loss improves performance by limiting the L1 penalty to feature-detection parameters.
  • Ablation study: Removing rmag slightly reduces performance, indicating that its contribution is useful but not critical.
  • Ablation study: Untied encoders do not improve performance and slightly worsen it, while tied encoder directions avoid a 50% increase in parameter count and inference-time compute.
  • Shrinkage comparison: Resolving shrinkage alone slightly improves baseline SAEs, but a significant performance gap remains relative to Gated SAEs.
  • Shrinkage comparison: The remaining gap suggests that Gated SAEs benefit from learning better encoder and decoder directions, beyond correcting feature-magnitude estimation.

6. Related Work

Related work connects Gated SAEs to mechanistic interpretability, classical sparse coding, dictionary learning in language models, and disentanglement. It also situates the method as an alternative response to shrinkage-related limitations in prior SAE training.

  • Mechanistic Interpretability: Mechanistic interpretability seeks to explain neural-network outputs through learned algorithms, motivating methods for finding interpretable directions in activations.
  • Classical Dictionary Learning: Classical sparse coding and neuroscience study sparse representations with more representations than basis vectors, providing precedents for dictionary learning.
  • Dictionary Learning in Language Models: Prior dictionary-learning work applied sparse or discrete feature methods to language-model representations across GPT-2-like, BERT, Pythia, and pretrained-model settings.
  • Dictionary Learning’s Limitations and Improvements: Wright and Sharkey identified shrinkage and proposed decoder finetuning, but finetuning all SAE parameters can sacrifice sparsity or interpretability.
  • Disentanglement: Unlike disentanglement methods that often align features with a latent-space basis, this work decomposes pretrained-language-model activations into sparse linear combinations of overcomplete dictionary elements.

7. Conclusion

The paper concludes that Gated SAEs improve the reconstruction–sparsity trade-off while retaining comparable interpretability, with substantial reductions in the number of firing features. The authors also identify assumptions, evaluation gaps, and possible inference-time alternatives that bound these conclusions.

  • Conclusion: Gated SAEs Pareto-improve reconstruction quality and sparsity while remaining comparably interpretable to baseline SAEs.
  • Conclusion: Gated SAEs require half the L0 to achieve the same loss recovered at many sites.
  • Limitations: The approach assumes that computation in large language models is sufficiently sparse and linear for SAE-based conclusions to be reliable.
  • Limitations: Interpretability evaluations use held-out data and manual checks but do not comprehensively test whether learned dictionaries contain causally meaningful intermediate variables.
  • Limitations: Inference-time pruning of low-activating baseline features could potentially narrow the performance gap without changing the baseline SAE.
  • Future work: Future work should test whether Gated SAEs continue to improve dictionary learning beyond 7B base language models, including larger chat, multimodal, and mixture-of-experts models.

9. Author contributions

The paper’s architecture, experiments, interpretability study, figures, and writing were distributed across the listed contributors.

  • Author contributions: Senthooran Rajamanoharan developed the Gated SAE architecture and training methodology and co-led the main experiments and writing.
  • Author contributions: Arthur Conmy and Senthooran Rajamanoharan performed the mainline experiments and led the writing of all sections.
  • Author contributions: Tom Lieberum implemented the manual interpretability study and created Figure 3, while János Kramár designed and analyzed that study.

A. Inference-time optimization

Inference-time optimization replaces a trained SAE encoder with a sparse approximation algorithm, allowing dictionaries to be compared independently of encoder quality across target sparsities. Gated dictionaries generally recover slightly more loss and resolve shrinkage, though standard dictionaries can sometimes perform better at higher test-time sparsity.

  • Method: Inference-time optimization uses a trained SAE decoder as a dictionary while ignoring the encoder.This separates dictionary quality from encoder quality during evaluation.
  • Method: A single dictionary can produce a full loss-recovered-versus-sparsity Pareto frontier by sweeping target sparsity without retraining.The test-time algorithm enables this sweep directly.
  • Results: Gated dictionaries have a small but real advantage over standard dictionaries in loss recovered at most target sparsity levels.The comparison suggests Gated SAEs learn better dictionaries in addition to addressing shrinkage.
  • Results: Baseline dictionaries can outperform dictionaries trained at the same low test-time sparsity, with higher-training-sparsity dictionaries often performing better under ITO.For example, Figure 10 reports that training L0 around 100 can outperform training L0 around 10 across many test-time sparsities.
  • Results: Gated SAEs outperform baseline SAEs at all but one tested MLP-output or residual-stream site, with an anomaly at one Gemma attention-output site.At that site, dataset mismatch and mean ablation explain why SAE reconstructions had lower loss than the original model.
  • Results: Gated SAEs resolve shrinkage in Pythia-2.8B as measured by relative reconstruction bias.Baseline SAEs exhibit shrinkage, whereas Gated SAE reconstructions are unbiased with γ≈1.

D.1.1. General training details

The training setup evaluates Gated and baseline SAEs across multiple models, sites, and hyperparameters while accounting for Gated SAEs’ higher training cost. The implementation includes large activation datasets, resampling, normalized training choices, and an equivalent parameterized activation-function view.

  • Data and evaluation: Experiments use activations from hundreds of millions to billions of language-model forward-pass activations.The activations are stored in shuffled buffers for optimization.
  • Training procedure: Training uses periodic resampling of rarely active features, with learning-rate warm-up reapplied after each resampling event.After resampling, the learning rate is reduced to 0.1× and cosine-warmed over the next 1000 steps.
  • Experimental scope: The study evaluates MLP outputs, attention outputs, and residual-stream activations across Pythia-2.8B and Gemma-7B sites.The broader experiments include five Pythia layers and four Gemma layers.
  • Architecture: Gated encoder equations are equivalent to a linear layer with a parameterized non-standard activation function under the stated weight-sharing scheme.The resulting activation is identified as JumpReLU in another context.

F. A toy setting where Jump ReLU SAEs outperform baseline SAEs

The toy setting illustrates why separating firing thresholds from activation magnitudes can make Jump ReLU more suitable than standard ReLU for sparse reconstruction. In this constructed example, Jump ReLU nearly perfectly reconstructs the signal, but the model is pedagogical and untested.

  • Motivation: Gated SAEs are equivalent at inference time to SAEs using a potentially discontinuous Jump ReLU instead of ReLU.The paper presents this greater expressivity as an additional reason they may outperform baseline SAEs.
  • Toy problem: In the toy model, ReLU’s single threshold must both decide whether to fire and determine the reconstruction origin.Choosing a low threshold increases false firing, while a high threshold distorts active-feature magnitudes.
  • Jump ReLU mechanism: Jump ReLU separates the firing threshold t from the origin d, allowing sparse detection without forcing the magnitude estimate to start at t.Its form is 𝟙(x>t)m(x−d), with d≤t.
  • Caveats: The toy model has not been empirically tested, and its pedagogical parameter values are not claimed to be mathematically optimal.The authors state that the model could be totally wrong while remaining useful for intuition.
  • Caveats: The toy analysis assumes a sparsely activating continuous feature, while the paper notes that the prevalence of such features in models remains open.The authors’ intuition is that many features are essentially binary but have confidence-valued coefficients.

H. Further analysis of the human interpretability study

Further analysis finds interpretability differences vary across sites and layers, while rater judgments retain a subjective component. The strongest gated-versus-baseline difference occurs for MLP-output SAEs.

  • Sites: A Friedman test finds significant differences between sites in Gated-vs-Baseline labels at p=.047, but not in raw labels at p=.92.The site effect appears in the paired difference rather than the unadjusted interpretability labels.
  • Sites: MLP-output SAEs show the strongest site-specific result, with mean .40, p=.003, and CI [.18, .63].Attention-output and final-residual SAEs have weaker, nonsignificant results with means .05 and −.07.
  • Layers: Layer analyses use Page’s trend test and Friedman tests to assess monotone trends and any layer differences separately for each model.The two tests address distinct forms of layer variation.
  • Layers: The layer results suggest some significant nonmonotone differences in both raw interpretability labels and Gated-vs-Baseline differences.The paper visualizes 90% BCa bootstrap confidence intervals per layer.
  • Raters: Rater significance tests weakly suggest that raters differed in their judgments, underscoring subjectivity in interpretability labeling.Different raters also saw different proportions of Pythia and Gemma features, partially confounding aggregation across models.
Loading 2404.16014v2…