PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling and evaluating sparse autoencoders

Leo Gao, Tom Dupré la Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, Jeffrey Wu

arXiv:2406.04093v1cs.LGcs.AI

TL;DR

Sparse autoencoders are difficult to scale because they must balance reconstruction and extreme sparsity while avoiding dead latents. The paper uses k-sparse autoencoders and training modifications to control sparsity directly, studies scaling and feature quality, and finds generally improving metrics with larger autoencoders, including a 16 million latent GPT-4 model.

  • Problem

    Training sparse autoencoders at useful language-model scales is difficult because reconstruction, sparsity, and dead-latent behavior complicate scaling studies.

  • Method

    The paper uses k-sparse TopK autoencoders, dead-latent mitigation techniques, scaling-law analysis, and metrics for feature recovery, explanation, and downstream-effect sparsity.

  • Results

    Larger sparse autoencoders generally score better on the introduced feature-quality metrics, and the methodology scales to a 16 million latent GPT-4 autoencoder.

  • Takeaways & Limitations

    The reported methods make large sparse autoencoders practical to train and evaluate across reconstruction, sparsity, and feature-quality dimensions.

  • Takeaways & Limitations

    A large fraction of randomly activating features, especially in GPT-4, are not yet adequately monosemantic.

Abstract

from arXiv · show

Sparse autoencoders provide a promising unsupervised approach for extracting interpretable features from a language model by reconstructing activations from a sparse bottleneck layer. Since language models learn many concepts, autoencoders need to be very large to recover all relevant features. However, studying the properties of autoencoder scaling is difficult due to the need to balance reconstruction and sparsity objectives and the presence of dead latents. We propose using k-sparse autoencoders [Makhzani and Frey, 2013] to directly control sparsity, simplifying tuning and improving the reconstruction-sparsity frontier. Additionally, we find modifications that result in few dead latents, even at the largest scales we tried. Using these techniques, we find clean scaling laws with respect to autoencoder size and sparsity. We also introduce several new metrics for evaluating feature quality based on the recovery of hypothesized features, the explainability of activation patterns, and the sparsity of downstream effects. These metrics all generally improve with autoencoder size. To demonstrate the scalability of our approach, we train a 16 million latent autoencoder on GPT-4 activations for 40 billion tokens. We release training code and autoencoders for open-source models, as well as a visualizer.

1 Introduction

The paper addresses the difficulty of training sparse autoencoders at scale and develops methods for reliable scaling, evaluation, and open-source use.

  • Motivation: Sparse autoencoders are promising for finding language-model features and circuits, but extreme sparsity has limited prior work to small models.This motivates methods that make larger autoencoders practical.
  • Contributions: The authors develop a methodology for training extremely wide, sparse autoencoders with very few dead latents across language models.They demonstrate the approach with a 16 million latent autoencoder on GPT-4 residual-stream activations.
  • Evaluation: The paper introduces metrics for hypothesized-feature recovery, downstream-effect sparsity, and high-precision, high-recall explanation of activation patterns.These metrics target feature usefulness beyond reconstruction and sparsity alone.
  • Results: Larger sparse autoencoders are generally better according to the paper’s latent-quality metrics.The reported metrics assess feature recovery, explainability, and sparsity of downstream effects.
  • Resources: The authors release training code, GPT-2 small autoencoders, and visualizers for GPT-2 small and the 16 million latent GPT-4 autoencoder.The release supports inspection and reuse of the trained models.

2 Methods

The methods replace difficult sparsity tuning with direct TopK control, address dead latents, and compare autoencoder behavior through reconstruction, sparsity, and downstream metrics.

  • Inputs: Inputs are residual-stream activations from GPT-2 small and GPT-4-series models, evaluated after mean subtraction and unit-norm normalization.Experiments use context length 64 and layers chosen near the end of each network.
  • Baseline: Baseline ReLU autoencoders optimize reconstruction MSE plus an L1 penalty on latent activations, requiring the sparsity coefficient λ to be tuned.MSE is normalized by the reconstruction error from always predicting mean activations.
  • TopK activation: TopK keeps only the k largest encoder activations, directly setting the number of active latents while leaving the decoder unchanged.This removes the L1 penalty and simplifies comparisons across sparsity levels.
  • Activation comparison: TopK significantly improves the reconstruction-sparsity Pareto frontier, with the best overall performance among compared activation functions.The comparison uses reconstruction MSE, active-latent count L0, and downstream cross-entropy loss.
  • Dead latents: Dead latents can substantially worsen MSE and waste computation; decoder-transpose initialization and an auxiliary top-kaux reconstruction loss are used to prevent them.The paper notes that unmitigated ablations reached up to 90% dead latents.

3 Scaling laws

The paper studies how sparse autoencoder reconstruction scales with compute, latent count, sparsity, and subject-model size. It finds power-law behavior and shows that larger subject models require larger autoencoders for matched reconstruction error, while noting that MSE alone is not a principled comparison across latent counts.

  • Compute scaling: MSE follows a power law of compute, although the smallest models deviate from the trend.
  • Evaluation boundary: Comparing MSE across different n is not a fair comparison because larger autoencoders impose looser information bottlenecks, making lower MSE easier to achieve.The paper therefore treats compute-optimized MSE training as arguably unprincipled when latents, rather than reconstruction predictions, are the main artifact.
  • Compute scaling: The largest learning rate that converges scales as 1/√n, while the optimal learning rate for L(N) is about four times smaller than for L(C).
  • Compute scaling: Tokens to convergence scale approximately as Θ(n0.6) for GPT-2 small and Θ(n0.65) for GPT-4.The paper notes this sublinear scaling must eventually break because each latent would otherwise receive vanishing gradient signal.
  • Sparsity scaling: MSE follows a joint scaling law in latent count n and sparsity k, but the fit applies only in the small-k regime.The fitted law includes interaction and irreducible-loss terms, with γ = −0.042 and η = −0.085; increasing k makes L(N) steeper and lowers irreducible loss.
  • Subject-model scaling: Larger subject models require larger autoencoders to achieve the same MSE at fixed k, with a worse scaling exponent.This result is reported for models in the GPT-4 family with k = 32.

4 Evaluation

The paper evaluates sparse autoencoders beyond reconstruction and sparsity using downstream behavior, feature recovery, explanation quality, and ablation-effect sparsity. These metrics generally favor larger, sparser autoencoders, although trends worsen when activations become dense.

  • Evaluation metrics: Autoencoder quality is evaluated with downstream loss, probe loss, explainability, and ablation sparsity rather than reconstruction and sparsity alone.The authors argue that reconstruction-sparsity improvements can become uninformative, so they measure whether learned features support language-model behavior and interpretation.
  • Scaling trends: Larger autoencoders generally improve the evaluation metrics, while increasing active latents has mixed effects and breaks down near dense regimes.Higher active-latent counts worsen token-pattern explanations but improve probe loss and ablation sparsity until L0 approaches dmodel.
  • Downstream behavior: k-sparse autoencoders improve downstream loss more than MSE over prior methods, with MSE scaling cleanly against KL divergence and delta cross-entropy at fixed sparsity.The scaling relationship can break when sparsity is modulated at test time.
  • Absolute evaluation: The 16 million latent autoencoder reaches a language-modeling loss corresponding to 10% of GPT-4's pretraining compute, offering an absolute behavioral interpretation.This compute-equivalent metric is proposed because zero-ablation baselines can produce high fidelity scores even for poor explanations; the zero-ablation fidelity is 98.2%.
  • Known-feature recovery: Probe scores increase and then decrease with k; TopK generally outperforms ReLU and residual-stream channels, while the metric improves during unsupervised GPT-4 autoencoder training.The probe metric tests whether hypothesized features are recovered using 61 binary classification tasks.
  • Explainability: N2G explanations are easiest for larger, sparser autoencoders, and TopK achieves over 1.5x recall with over 0.9x precision relative to same-size ReLU models.Dense token patterns are easy to explain but can reflect dense, less informative latents.
  • Explanation reconstruction: Using N2G explanations to simulate latents yields better downstream loss for larger and sparser autoencoders, combining explanation quality with behavioral relevance.The metric weights both precision and recall and gives more weight to recall for densely activating latents.
  • Ablation sparsity: Ablation effects are sparse, with latent-effect (L1/L2)^2 values of 10-14% versus 60% for residual-stream channels, but become less interpretable near k = 512.For GPT-2 small, the metric uses T = 16 future tokens and reverses its favorable trend when k approaches dmodel = 768.

5 Understanding the TopK activation function

The TopK activation function avoids L1-induced activation shrinkage and improves the reconstruction-sparsity frontier, but standard TopK overfits its training sparsity when used at test time. Multi-TopK largely restores generalization across fixed and dynamic sparsity levels.

  • Activation shrinkage: TopK prevents activation shrinkage by selecting the k largest activations without an L1 penalty.Refinement increases ReLU activations on average but not TopK activations, indicating less shrinkage in TopK models.
  • Activation refinement: Activation refinement improves reconstruction MSE and downstream next-token cross-entropy for both ReLU and TopK, but only partially closes their performance gap.The refinement is smaller and unbiased for TopK, unlike the positive-biased refinement observed for ReLU.
  • Frontier performance: TopK improves the reconstruction-sparsity Pareto frontier over ReLU, with the best overall performance among the compared approaches.The comparison includes reconstruction MSE, active-latent count L0, and downstream cross-entropy loss.
  • Test-time sparsity: Standard TopK provides a progressive code only up to its training k, after which reconstruction worsens as test-time k′ increases.MSE improves slightly above the training k before becoming substantially worse, indicating overfitting to the training sparsity.
  • Multi-TopK: Multi-TopK produces a progressive code across test-time k′ values, though it performs slightly worse than TopK at the training k.The example loss L(k) + L(4k)/8 is sufficient to obtain this behavior.
  • Fixed versus dynamic sparsity: With Multi-TopK, TopK and JumpReLU test-time curves nearly coincide, supporting either fixed or token-dependent latent counts without reconstruction loss.JumpReLU selects latents above a threshold θ, allowing the number of active latents to vary across tokens.

6 Limitations and Future Directions

The paper identifies optimization, metric, interpretability, and scaling limitations that motivate several future research directions.

  • TopK forces every token to use exactly k latents, whereas constraining expected sparsity E[L0] would likely be preferable.
  • The optimization could improve through learning-rate scheduling, better optimizers, and better auxiliary losses for preventing dead latents.
  • Future work should better establish which metrics track downstream relevance and study applications such as steering, anomaly detection, and circuit identification.
  • Combining mixture-of-experts methods with autoencoders could reduce asymptotic training cost and enable much larger autoencoders.
  • A large fraction of discovered feature activations, especially in GPT-4, are not yet adequately monosemantic.The authors believe improved techniques and greater scale in latent count and training tokens may overcome this issue.
  • The probe metric is noisy, n2g captures only simple patterns, and 64-token contexts may miss interesting GPT-4 behaviors.

7 Related work

This section describes training choices and implementation findings for sparse autoencoders, including auxiliary losses, optimization, batching, normalization, and input preprocessing.

  • Auxiliary loss: AuxK models reconstruction error with selected dead latents, and adding its auxiliary loss increases computational cost by about 10%.Latents are flagged dead after typically 10 million inactive tokens; the usual coefficient is α = 1/32.
  • Optimization and batching: GPT-4 scaling requires gradient clipping to prevent instability and divergence at higher learning rates, affecting L(C) substantially but not L(N).
  • Optimization and batching: A batch size of 131,072 tokens is used for most experiments to gain parallelism, while appropriately tuned hyperparameters keep L(N) largely batch-size independent.
  • Input preprocessing: Inputs are mean-centered across dmodel and normalized to unit norm before autoencoder processing or reconstruction-error computation.

B.1 TopK training details

TopK training combines direct sparsity control with auxiliary dead-latent mitigation and systems optimizations, enabling large-scale experiments while revealing several scaling patterns.

  • Dead-latent prevention: AuxK empirically eliminates almost all dead latents by training completion, while baseline models generally have few dead latents as well.
  • Dead-latent prevention: The reduction in dead latents is mostly attributable to combining AuxK loss with tied initialization.
  • Initialization: Tied initialization substantially improves MSE, whereas the encoder initialization scheme has no effect with tied initialization and slightly hurts independently.
  • Ablations: Decoder normalization slightly improves loss, and benc does not affect MSE at convergence although removing it slows convergence slightly.
  • Systems optimization: For the largest 16-million-latent autoencoder, the implementation uses 512-way sharding and exploits sparse operations to reduce compute and memory.
  • Systems optimization: Sparse kernels theoretically improve compute efficiency by up to 6x in the sparsity limit because only the encoder forward pass remains dense.
  • Scaling patterns: Log latent density is approximately Gaussian, while modal density and modal feature importance decrease as total latent count increases at fixed k.

E.6 Recurring dense features in GPT-2 small

Recurring dense latents include context-position, first-token, and repetition features, with some opposing pairs suggesting optimization failures. Latents also separate into two fuzzily distinct sub-spaces with different vocabulary-token diversity.

  • Recurring dense features: The two densest latents encode increasing activation toward later or earlier context positions, reaching approximately 40% and 35% density.TopK may prevent these features from activating at lower values.
  • Recurring dense features: First-token-position features recur at approximately 30% density, with large activations concentrated at the first token and smaller later activations appearing uninterpretable.The authors conjecture that the smaller values reflect interference with the first-position direction.
  • Recurring dense features: A recurring repetition latent is approximately 20% dense and activates strongly on repetitive sequences including dates, numbers, punctuation, and repeated phrases.Random activations are typically unrelated and uninterpretable.
  • Recurring dense features: Opposing latent pairs often have decoder cosine similarity close to −1, and the less dense member is typically uninterpretable.The authors conjecture that these pairs reflect optimization failure through cancellation of spurious activations.
  • Latent sub-spaces: The latent projections reveal two fuzzily separated sub-spaces using approximately 25% and 75% of the vector-space dimensions.The sub-space distinction is fairly consistent with ordering by singular values and may be guided by different encoder norms.
  • Latent sub-spaces: Latents in the smaller cluster activate on fewer than 100 vocabulary tokens, whereas some latents in the larger cluster activate on up to 1000.Effective token number quantifies the vocabulary-token diversity of each latent.

F.1 Impact of different locations

Autoencoder behavior varies substantially by layer and context position: later layers and tokens can be harder to reconstruct or more consequential for downstream predictions. Different evaluation metrics favor different locations.

  • Layer and location: Learning-rate optima vary by layer and location type, but remain within a factor of two across the GPT-2 small sweep.The location types include MLP delta, attention delta, MLP post, and attention post.
  • Layer and location: Earlier layers require noticeably more tokens for convergence, while residual-stream MSE generally increases toward the final layer before dropping there.Residual-stream deltas peak around layer 6, and attention-delta MSE falls sharply at the final layer.
  • Layer and location: Downstream loss and KL worsen with layer during reconstruction ablation despite lower normalized MSE at late layers, with an exception around the final layers.Layers 11/12 and especially 12/12 can have better normalized MSE but more severe downstream effects.
  • Layer and location: Earlier layers yield better Neuron2graph explanations, whereas later layers perform better on probe loss and sparsity.This demonstrates that layer choice affects evaluation metrics differently.
  • Context position: Later context positions are harder to reconstruct, while first positions are especially easy by normalized MSE but have residual-stream norms more than an order of magnitude larger.The first-position norm exception occurs at early and final GPT-2 small layers.
  • Context position: First-position reconstruction ablations produce significantly worse downstream loss and KL at layers with large first-position norms despite better normalized MSE.The authors relate this to large-norm directions important for loss on other tokens rather than the current token.

G Irreducible loss term

The authors find an irreducible loss term in autoencoder scaling fits, despite the possibility of perfect reconstruction with an arbitrarily large autoencoder. They hypothesize that less structured activation components scale more slowly.

  • Irreducible loss: Fit quality is substantially worse when the scaling law omits an irreducible loss term, despite arbitrarily large autoencoders being able to perfectly reconstruct inputs.The authors contrast this with language-model loss, which has irreducible uncertainty from unpredictable text.
  • Irreducible loss: The authors hypothesize that activations contain components with different amounts of structure, with less structured data having a worse scaling exponent.This remains a hypothesis rather than a fully understood explanation.
  • Irreducible loss: An L(N) exponent of −0.04 appears on 768-dimensional Gaussian noise, compared with approximately −0.26 on similarly dimensional GPT-2 small activations.The Gaussian-data scaling law is much shallower.

H Further probe based evaluations results

The probe-based evaluation follows autoencoder quality during training and examines the 16M model across tasks. The supplied passages identify the baseline and task breakdown but do not state numerical probe outcomes.

  • Probe evaluation: Probe evaluation scores are tracked through training for 128k, 1M, and 16M autoencoders against a residual-stream-channel baseline score of 0.600.The figure compares score trajectories across autoencoder sizes.
  • Probe evaluation: The 16M autoencoder’s probe scores are examined from approximately 10B tokens elapsed, when probe features begin developing.This focuses the analysis on the later training period.
  • Probe evaluation: The 16M probe evaluation is broken down by task, with Europarl, bigrams, occupations, and AG News aggregating multiple tasks.The supplied passage does not report the resulting scores.
  • Probe evaluation: The evaluation suite consists of tasks used for probe-based evaluation.The supplied table caption identifies the table but does not enumerate the tasks.

I Contributions

The project combined systems and scaling work, metric and visualization infrastructure, feature analysis, and research leadership across contributors. Several researchers also advised on optimization, sparsity, and probe-based evaluation.

  • Systems and scaling: Leo Gao implemented the training codebase and infrastructure for GPT-4 experiments and conducted most scaling and architecture experiments.His systems work covered kernels, parallelism, numerics, and data processing, while experiments included latent counts, model size, batch size, token budget, and learning rate.
  • Methods and analysis: Tom Dupré studied activation shrinkage, progressive recovery, and Multi-TopK, while implementing Gated and ProLU baselines and analyzing GPT-2 small layers and locations.He also helped refine scaling laws for L(N, K) and Ls(N) and discovered latent sub-spaces.
  • Explanations and features: Henk Tillman worked on N2G explanations, language-model-based explainer scoring, activation scraping infrastructure, and qualitatively interesting and safety-related features.His work included finding features of qualitative interest, including safety-related features.
  • Evaluation and visualization: Jeff Wu studied ablation effects, sparsity, and ablation reconstruction; managed metrics infrastructure; and built the visualizer’s data collation and website code.He also analyzed the cost of a fully sparse bottleneck, recurring dense features, and safety-relevant features using attribution.
  • Research advice: Gabriel Goh suggested TopK and contributed intuitions about TopK, AuxK, and optimization, while Rajan Troll advised on optimization, scaling, and systems.Alec Radford contributed the irreducible loss term and advised on probe-based metrics, optimization, and scaling.
  • Leadership: Jan Leike and Ilya Sutskever managed and led the Superalignment team.
Loading 2406.04093v1…