Research Reports

Analyzing SAELens

6 min read

Analyzing SAELens

Source: decoderesearch/SAELens @ 5c68c34

How It Works

SAELens is a library for training and analyzing Sparse Autoencoders (SAEs) on language model activations — a core technique in mechanistic interpretability. The goal is to decompose the dense, opaque activations of a transformer into sparse, interpretable features. Given an LLM and a dataset, SAELens trains an autoencoder whose latent space is forced to be sparse (most neurons fire near-zero most of the time), then provides tools to analyze what those sparse features represent.

The system has three major phases: training, loading, and analysis. During training, activations are extracted from a target LLM using TransformerLens, stored in a mixing buffer, and fed to an SAE that learns to reconstruct them under sparsity constraints. After training, the SAE can be saved (via safetensors) and loaded for inference or wrapped around a model to intercept activations at specific layers. The analysis layer provides HookedSAETransformer, which attaches SAEs as hooks directly into TransformerLens’s forward pass, enabling seamless feature inspection.

Architecture

The codebase follows a clean three-layer design:

sae_lens/
├── saes/           # 9 SAE architecture implementations
│   ├── sae.py      # Base classes: SAE, TrainingSAE, SAEConfig
│   ├── standard_sae.py, gated_sae.py, topk_sae.py, ...
│   └── registry.py # Architecture name → class mapping
├── training/       # Training pipeline (8 modules)
│   ├── activations_store.py   # Buffer for model activations
│   ├── sae_trainer.py         # Core training loop
│   └── train_sae_on_language_model.py  # Runner orchestrator
├── loading/        # Model & SAE loading (2,148 LOC)
│   └── pretrained_sae_loaders.py  # Converters for GemmaScope, OpenAI, etc.
├── analysis/       # HookedSAETransformer, Neuronpedia integration
├── synthetic/      # Synthetic benchmarking framework (14 files)
├── config.py       # Dataclass configs (749 LOC)
└── evals/          # Downstream eval metrics

Total: ~19,200 LOC source, ~37,400 LOC tests across 80 test files.

The Spine

The primary user flow traces through three stages:

1. Configuration. Users create a LanguageModelSAERunnerConfig — a 60+ field dataclass that is generic over the SAE architecture type. The config specifies the model, dataset, layer to hook, SAE architecture, hyperparameters, and logging. CLI users pass these via simple_parsing, which maps nested dataclass fields to --flag.subflag arguments.

2. Training. LanguageModelSAETrainingRunner.__init__() loads the LLM via TransformerLens, creates an ActivationsStore (which runs the model on a dataset and buffers activations in a MixingBuffer), and instantiates the SAE. The SAETrainer.fit() loop runs: activations are batched, the SAE’s training_forward_pass() encodes and decodes them, and losses (reconstruction + sparsity auxiliary) are backpropagated. Mixed precision (bfloat16), torch.compile, and GPU prefetching are available.

3. Inference/Analysis. After training, save_final_sae() converts the TrainingSAE into a lightweight inference SAE and saves via safetensors. Users can load pretrained SAEs with SAE.from_pretrained(release, sae_id) or wrap a model with HookedSAETransformer.from_pretrained() to intercept activations at specific layers during any forward pass.

Key Patterns

Inference/Training Split

Every SAE architecture has two classes: an inference variant (extends SAE) and a training variant (extends TrainingSAE). The inference class provides encode()/decode()/forward(). The training class adds training_forward_pass(), calculate_aux_loss(), and get_coefficients(). This cleanly separates deployment concerns from training internals — a pattern few ML libraries bother with.

Architecture Registry

registry.py maintains two global dicts (SAE_CLASS_REGISTRY, SAE_TRAINING_CLASS_REGISTRY) mapping string names to class tuples. Nine architectures are supported: standard, gated, TopK, batch TopK, Matryoshka batch TopK, JumpReLU, matching pursuit, temporal, and transcoder variants. Adding a new architecture requires implementing the two classes and registering them — no modification to core code.

Dataclass-First Configs

All configuration uses Python dataclasses with __post_init__ validation, not Pydantic. simple_parsing bridges these to CLI arguments. The main config is generic: LanguageModelSAERunnerConfig[T_TRAINING_SAE_CONFIG], which gives type-safe access to architecture-specific parameters after parsing.

Activation Normalization Modes

Four normalization strategies handle the fact that different transformer layers have different activation scales: none, constant norm rescale, layer norm, and Anthropic’s “expected average” approach. The latter estimates a scaling factor during training then permanently folds it into weights post-training via fold_activation_norm_scaling_factor().

Non-Obvious Details

Meta-device loading trick. SAEs are instantiated on the meta device (zero memory), then load_state_dict(assign=True) loads weights directly into the target device. This avoids the common pattern of allocating then copying, cutting peak memory in half for large SAEs.

Sparse COO tensors in TopK. The TopK implementation uses PyTorch sparse COO tensors with a custom SparseHookPoint that only converts to dense when analysis hooks are active. A custom _sparse_matmul_nd() reshapes N-dimensional sparse tensors to 2D for sparse.mm. This is a significant memory optimization — TopK SAEs can have millions of features but only k active per token.

Custom autograd for JumpReLU. JumpReLU uses a hand-written torch.autograd.Function called Step with a rectangle-based gradient approximation. The soft threshold’s bandwidth is a learnable hyperparameter, and the gradient is designed to be smooth in a band around the threshold while being zero elsewhere — a clever balance between straight-through estimation and full softplus.

_SAEWrapper.__dict__ trick. The analysis layer stores SAEs in __dict__ rather than as nn.Module submodules. This preserves TransformerLens’s hook cache path compatibility — the SAE doesn’t appear in the module tree, so cached activations aren’t duplicated.

Synthetic benchmarking framework. The synthetic/ package provides a complete ground-truth testing environment: FeatureDictionary defines orthogonal feature directions, ActivationGenerator fires them with configurable probabilities (Zipfian), and a Hierarchy system creates tree-structured feature dependencies. This allows measuring SAE recovery quality against known features — far more rigorous than evaluating only on real model activations.

V5 equivalence tests. tests/_comparison/ retains old v5 implementations and verifies that v6 produces identical outputs for the same inputs. This regression safety net is unusual for a research library and speaks to the seriousness of the v6 rewrite.

Assessment

Code quality: Excellent. Strict Pyright type checking (with only 9 targeted relaxations), comprehensive ruff linting, and a CLAUDE.md that enforces simplicity and correctness. The code is consistently readable despite implementing complex math.

Architecture fitness: Very high. The inference/training split, generic config system, and registry pattern create an extensible framework that doesn’t sacrifice type safety. Deep TransformerLens integration is both a strength (seamless hook-based activation extraction) and a coupling risk — migrating to other frameworks would require significant work.

Testing: Exceptional. At 2× source LOC, the test suite is rigorous. Per-architecture config builders in tests/helpers.py, statistical tests with tight bounds, v5 equivalence tests, and CI disk usage tracking demonstrate mature testing practices.

Concerns:

  • pretrained_sae_loaders.py (2,035 LOC) is a monolith handling format conversion for many external SAE formats. It would benefit from splitting into per-provider modules.
  • The 60+ field LanguageModelSAERunnerConfig has a steep learning curve. The __post_init__ validation block is substantial and some defaults depend on each other in non-obvious ways.
  • The coupling to TransformerLens means the library can’t easily work with raw HuggingFace models or other transformer implementations.

Recommendation: SAELens is the most mature open-source SAE training framework available. The v6 rewrite is a significant improvement in type safety and extensibility. It’s the right choice for researchers who need to train or analyze SAEs on TransformerLens-compatible models. Those working outside that ecosystem should evaluate the coupling cost before committing.

Graph View