PaperlayerView exact revision on arXiv

Source-linked AI summary

Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders

Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Conmy, Vikrant Varma, János Kramár, Neel Nanda

arXiv:2407.14435v3cs.LG

TL;DR

Sparse autoencoders must balance faithful reconstruction of language-model activations against sparse, interpretable decompositions. The paper introduces JumpReLU SAEs, trains their discontinuous thresholds and direct L0 objective with straight-through estimators, and finds strong fidelity without sacrificing average interpretability or training efficiency.

  • Problem

    SAEs need both sparse, interpretable decompositions and faithful reconstructions, but these objectives are in tension.

  • Method

    The paper replaces ReLU with thresholded JumpReLU and trains a weighted L2 reconstruction plus L0 sparsity loss using straight-through estimators.

  • Results

    JumpReLU SAEs consistently exceed Gated SAE fidelity and match or exceed TopK fidelity at fixed sparsity across Gemma 2 9B activation sites.

  • Takeaways & Limitations

    JumpReLU offers a mild improvement over prevailing SAE methods while retaining similar average feature interpretability and efficient elementwise computation.

  • Takeaways & Limitations

    Evaluations cover several sites and layers of only Gemma 2 9B, leaving transfer to other model families uncertain.

Abstract

from arXiv · show

Sparse autoencoders (SAEs) are a promising unsupervised approach for identifying causally relevant and interpretable linear features in a language model's (LM) activations. To be useful for downstream tasks, SAEs need to decompose LM activations faithfully; yet to be interpretable the decomposition must be sparse -- two objectives that are in tension. In this paper, we introduce JumpReLU SAEs, which achieve state-of-the-art reconstruction fidelity at a given sparsity level on Gemma 2 9B activations, compared to other recent advances such as Gated and TopK SAEs. We also show that this improvement does not come at the cost of interpretability through manual and automated interpretability studies. JumpReLU SAEs are a simple modification of vanilla (ReLU) SAEs -- where we replace the ReLU with a discontinuous JumpReLU activation function -- and are similarly efficient to train and run. By utilising straight-through-estimators (STEs) in a principled manner, we show how it is possible to train JumpReLU SAEs effectively despite the discontinuous JumpReLU function introduced in the SAE's forward pass. Similarly, we use STEs to directly train L0 to be sparse, instead of training on proxies such as L1, avoiding problems like shrinkage.

1. Introduction

SAEs seek sparse and faithful decompositions of language-model activations, but these objectives conflict. JumpReLU SAEs address this tension with thresholded activations, direct L0 training, and competitive fidelity, efficiency, and interpretability.

  • Motivation: SAEs decompose language-model activations into sparse combinations of learned feature directions for interpretability and downstream analysis.Applications include circuit analysis and model steering.
  • Motivation: Sparsity and faithfulness are naturally in tension because increasing sparsity typically increases reconstruction error for a fixed dictionary size and training method.
  • Method: JumpReLU replaces ReLU with a positive threshold that removes low pre-activations while preserving larger ones, separating feature selection from magnitude estimation.The modification is designed to reduce false positives without systematically shrinking active feature magnitudes.
  • Method: Straight-through estimators enable standard gradient-based training despite the thresholded loss providing no direct gradient for the threshold parameter.The expected-loss derivative is expressed using probability densities of feature activations.
  • Results: At a given sparsity level, JumpReLU reconstructions consistently outperform Gated SAEs and match or slightly exceed TopK SAEs across Gemma 2 9B activation sites.JumpReLU uses one forward and backward pass with an elementwise activation, unlike TopK’s partial sort.
  • Results: Manual and automated evaluations find average JumpReLU features similarly interpretable to Gated and TopK features, despite more very-high-frequency features than Gated SAEs.Fewer than 0.06% of features in a 131k-width SAE have very high frequencies.

2. Preliminaries

SAEs encode activations as sparse nonnegative feature coefficients over an overcomplete learned dictionary and reconstruct them through encoder and decoder functions. Their architectures, losses, activation functions, and sparsity penalties determine the sparsity–fidelity trade-off.

  • SAE architectures: SAEs represent an activation x in R^n using a sparse feature vector f in R^M over M much greater than n learned dictionary directions.The decoder columns represent the feature directions used to reconstruct x.
  • Activation functions: JumpReLU zeroes pre-activations below a positive threshold and acts as the identity above it, allowing activity selection to be separated from magnitude estimation.The threshold is implemented through a Heaviside step function.
  • Loss functions: SAE losses combine reconstruction error with a sparsity penalty, while the sparsity coefficient controls the trade-off between sparsity and reconstruction fidelity.Auxiliary loss terms can help train otherwise poorly differentiated parameters or revive dead features.
  • Activation functions: TopK activations enforce sparsity by retaining only the top K pre-activations, so TopK SAEs do not require a separate sparsity penalty.
  • Sparsity penalties: L1 sparsity penalties provide gradients but also penalize feature magnitudes, can harm reconstruction fidelity, and are not invariant to encoder–decoder rescaling.Reparameterization-invariant L1 removes the need to constrain decoder-column norms.
  • Kernel density estimation: Kernel density estimation empirically estimates probability densities from finite samples and is used here to estimate quantities involving feature-activation distributions.

3. JumpReLU SAEs

JumpReLU SAEs modify standard ReLU SAEs with per-feature positive thresholds and train them using an L2 reconstruction loss plus a direct L0 sparsity penalty. Straight-through pseudo-derivatives provide threshold gradients despite the piecewise-constant forward functions.

  • JumpReLU replaces the encoder’s ReLU with an activation that zeroes pre-activations below each feature’s positive threshold.The threshold vector specifies when each feature becomes active.
  • The threshold separates deciding which features are active from estimating the magnitudes of active features.This parallels the gating mechanism in Gated SAEs and the TopK activation function.
  • The training loss combines L2 reconstruction error with an L0 sparsity penalty instead of using an L1 proxy or auxiliary threshold-training task.The L0 penalty is expressed using Heaviside step functions on encoder pre-activations.
  • Pseudo-derivatives approximate the weak derivatives of JumpReLU and Heaviside functions, providing threshold gradients when pre-activations lie within an ε-wide window.The pseudo-derivatives use a kernel K, with ε functioning as a KDE bandwidth; too little or too much smoothing increases noise or bias.
  • Because the threshold appears inside Heaviside functions, the loss supplies no ordinary gradient for training it.The remaining SAE parameters are encouraged purely to reconstruct faithfully, while sparsity is handled by the threshold parameter.
  • JumpReLU SAEs are trained with ordinary batch-gradient updates while applying the pseudo-derivatives during the backward pass.The paper provides pseudocode for the forward pass, loss, and straight-through estimators in automatic-differentiation frameworks.

4. How STEs enable training through the jump

The JumpReLU loss is piecewise constant for individual inputs but has a generally non-zero expected-loss gradient. Kernel-density estimates implemented through straight-through estimators recover this gradient for threshold updates.

  • Training seeks the gradient of the expected loss, which is not generally zero even though each individual loss is piecewise constant in the threshold.This distinction makes threshold learning possible despite zero pointwise gradients.
  • The expected loss can be differentiated analytically with respect to each threshold using the feature pre-activation density.The density is evaluated at the corresponding threshold, with decoder columns entering the expression.
  • A generalized kernel-density estimator estimates the expected-loss gradient from batches of input activations.The estimator uses sampled activations x_1, x_2, …, x_N.
  • Using the defined pseudo-derivatives in the backward pass produces the batch-wise mean gradient used by the optimizer to update the thresholds.The pseudo-derivatives therefore connect the discontinuous forward pass to standard gradient-based optimization.
  • Straight-through training is equivalent to estimating the true expected-loss gradient with the paper’s kernel-density estimator.This establishes the training procedure as a gradient-estimation method rather than an arbitrary surrogate.

5. Evaluation

The evaluation compares JumpReLU, Gated, and TopK SAEs across sparsity, reconstruction fidelity, feature frequencies, dead features, and interpretability. JumpReLU consistently matches or exceeds competing methods in fidelity while retaining comparable interpretability and few dead features.

  • 5.1. Evaluating the sparsity-fidelity trade-off: The evaluation measures sparsity with mean L0-norm and reconstruction fidelity primarily with delta LM loss, with FVU as a secondary metric.Delta LM loss is the increase in language-model cross-entropy when the SAE is spliced into the forward pass.
  • 5.1. Evaluating the sparsity-fidelity trade-off: JumpReLU SAEs consistently offer similar or better reconstruction fidelity at a given sparsity level than Gated or TopK SAEs.This pattern holds for Gemma 2 9B residual-stream activations and similarly for MLP- and attention-output activations.
  • 5.2. Feature activation frequencies: TopK and JumpReLU SAEs have more features active on over 10% of tokens than Gated SAEs, while over-1%-active feature counts are more mixed.Gated SAEs can have considerably more over-1%-active features than JumpReLU SAEs in the low-delta-LM-loss, higher-L0 regime.
  • 5.2. Feature activation frequencies: Fewer than 0.06% of features activate on over 10% of tokens even for the highest-L0 SAEs, and JumpReLU and AuxK-trained TopK SAEs consistently have few dead features without resampling.Dead features are defined as activating fewer than one in 10^7 tokens.
  • 5.3.1. Manual Interpretability: All three SAE architectures exhibit similar manual interpretability rating distributions across the evaluated layers, sublayers, and sparsity levels.The study selected 81 SAEs and assessed features at average active-feature counts near 20, 75, and 150.

6. Related work

SAE research seeks sparse representations that expose useful LM concepts and computational structure, while recent methods improve the sparsity–fidelity trade-off through thresholding and gating.

  • SAEs are studied as a way to uncover linearly represented concepts and computational subgraphs in language-model activations.
  • Recent SAE architectures and training methods aim to improve the sparsity–fidelity trade-off while preserving interpretability and functional relevance.
  • Thresholding mechanisms determine which features contribute to reconstructions in TopK, Gated, and related SAE approaches.
  • Straight-through estimators have been applied to train through discontinuities in quantized networks and other models.

7. Discussion

JumpReLU SAEs match or surpass competing methods in reconstruction faithfulness while retaining comparable interpretability and efficient training, but their scope and evaluation remain limited.

  • JumpReLU reconstructions consistently match or exceed TopK fidelity and exceed Gated fidelity at a given sparsity level.
  • JumpReLU features are similarly interpretable on average to Gated and TopK features in manual and automated evaluations.
  • JumpReLU training requires no auxiliary terms and avoids relatively expensive TopK operations.
  • The evaluations cover several sites and layers of only Gemma 2 9B, leaving transfer to other model families uncertain.
  • JumpReLU has relatively more very high-frequency features than Gated, while resembling TopK in this respect.
  • The method introduces initial-threshold and bandwidth hyperparameters, although default settings transfer reliably across tested models, sites, and layers.

9. Author contributions

The supplied material covers the paper’s evaluation and training methodology, including comparisons among Gated, TopK, and JumpReLU SAEs and methods for differentiating discontinuous objectives.

  • Straight-through estimators are used to estimate gradients for JumpReLU’s discontinuous threshold and sparsity objectives through expected-loss derivatives and kernel density estimation.
  • Gated SAEs use separate gating and magnitude computations, whose outputs are combined elementwise before standard decoding.
  • The evaluations compare Gated and TopK SAE variants, including approximate versus exact TopK and the use of the AuxK loss.
  • Approximate TopK performs similarly to exact TopK, while AuxK slightly improves fidelity at a given sparsity level.
  • The sparsity penalty trains JumpReLU SAEs toward target L0 values while minimizing reconstruction loss.

F. Training JumpReLU SAEs to match a desired level of sparsity

JumpReLU SAEs can be trained to target a desired sparsity level rather than leaving sparsity implicit in the sparsity coefficient and reconstruction loss. The resulting models achieve target-adjacent L0 values while maintaining a similar fidelity–sparsity relationship.

  • Targeting a specific sparsity level can replace the implicit L0 determined by the sparsity coefficient and reconstruction loss.This mirrors the controllability available when setting K for TopK SAEs.
  • A proposed loss trains JumpReLU SAEs toward a chosen sparsity target.The passage introduces this loss as a simple way to achieve target-controlled sparsity.
  • On Gemma 2 9B residual-stream activations after layer 20, the target-controlled loss produced a similar fidelity-to-sparsity relationship to the original JumpReLU loss.The comparison is reported in Fig. 9.
  • The converged L0 values were close to their specified targets.This is indicated by red points lying near their corresponding vertical target lines.

G. Additional benchmarking results

Additional benchmarks examine reconstruction fidelity, feature activity, interpretability-related measures, and implementation choices across JumpReLU, Gated, and TopK SAEs. Results generally favor JumpReLU or show close behavior with Gated, while some metrics vary by layer, site, and sparsity.

  • Interpretability: Automated interpretability correlations vary by layer and site, with layer 20 and pre-linear attention output performing worst on this metric.The analysis reports distributions and means of correlations between LM-simulated and ground-truth activations.
  • Reconstruction fidelity: JumpReLU SAEs consistently provide more faithful reconstructions than Gated and at least as good, often slightly better, reconstructions than TopK at matched sparsity.This pattern is reported for MLP outputs, attention activations, and other Gemma 2 9B activation sites.
  • Feature activity: JumpReLU and TopK SAEs have few dead features without resampling, whereas the original Gated loss has around 40% dead features at layer 20.Dead features activate on fewer than one in 10^7 tokens.
  • Feature activity: JumpReLU and TopK feature-frequency histograms include a small high-frequency mode, while Gated has a shoulder at frequencies between 10^-2 and 10^-1.The histograms compare SAEs with L0 approximately 70 and use a logarithmic y-axis.
  • Training robustness: Changing the kernel used for pseudo-derivatives has little effect on fidelity-versus-sparsity curves.The comparison uses Gemma 2 9B post-layer 20 residual-stream SAEs trained on 2B tokens.
  • Attribution-weighted sparsity: Increasing L0 makes active-feature effects more diffuse for all SAE types, with the strongest effect appearing for TopK.JumpReLU and Gated behave mostly identically except for layer 31 residual-stream SAEs, and the authors caution against premature conclusions about feature quality.

I. Further details on our training methodology

The training methodology specifies activation normalization, optimizer settings, initialization choices, resampling policy, and JumpReLU threshold parameters. An explored constraint was not necessary for fidelity curves, but its broader consequences were not fully investigated.

  • Training setup: LM activations are normalized to have mean squared L2 norm of one during training.
  • Training setup: The optimizer is Adam with β1 = 0, β2 = 0.999, and ϵ = 10^-8.The learning rate starts at 7 × 10^-6, which is 10% of its final value.
  • Training setup: A pre-encoder bias is used during training, and ablations found no impact or a small performance improvement depending on model, site, and layer.
  • Initialization and scheduling: Encoder weights are initialized as the transpose of decoder weights, and λ is warmed up over the first 10,000 steps for non-TopK SAEs.The matrices are untied after initialization.
  • Dead-feature handling: Resampling is used for original-loss Gated SAEs but not for Gated RI-L1, TopK, or JumpReLU SAEs.
  • Caveat: Removing the explored constraint does not change fidelity-versus-sparsity curves, but its broader consequences were not fully explored.

J. Pseudo-code for implementing and training JumpReLU SAEs

The pseudo-code implements JumpReLU with custom backward passes, constructs the SAE reconstruction, and combines reconstruction error with an L0 sparsity penalty. Positive log-threshold parameterization and pre-activation handling support stable training behavior.

  • Components: The implementation includes custom backward passes for the Heaviside step function and JumpReLU, plus the SAE forward pass and loss.
  • Portability: The pseudo-code is written most closely for JAX but is portable to frameworks such as PyTorch with minimal changes.
  • Parameterization: The trainable parameter is log(threshold), ensuring a positive threshold during training.ReLU pre-activations prevent negative values from influencing gradient computation when thresholds become smaller than half the bandwidth.
  • JumpReLU activation: JumpReLU outputs x when x exceeds the threshold and zero otherwise.
  • JumpReLU activation: The JumpReLU backward pass differentiates with respect to x above threshold and uses a rectangle-kernel pseudo-derivative for the threshold.
  • SAE forward pass: The SAE applies an encoder, positive thresholding, and a decoder to produce reconstructed activations and feature magnitudes.A pre-encoder bias can optionally be applied, and the encoder uses ReLU pre-activations before JumpReLU.
  • Loss: The loss averages per-example squared reconstruction error plus a sparsity coefficient multiplied by the number of active features.
Loading 2407.14435v3…