saelens

Train sparse autoencoders to interpret model features.

  • Sparse Autoencoders
  • SAE
  • Mechanistic Interpretability
  • Feature Discovery
  • Superposition

Declared platforms: linux · macos · windows

Install
npx skills add 'https://github.com/NousResearch/hermes-agent/tree/main/optional-skills/mlops/saelens'
Download bundle ↓
main · 24fd22bScanned 2026-09-15

Contributors

GitHub-linked commit authors for this SKILL.md at the saved revision. Co-authors and history before file renames are not included.

File history ↗
View on GitHub
---name: saelensdescription: Train sparse autoencoders to interpret model features.version: 1.0.1author: Orchestra Researchlicense: MITdependencies: [sae-lens>=6.0.0, transformer-lens>=2.0.0, torch>=2.0.0]platforms: [linux, macos, windows]metadata:  hermes:    tags: [Sparse Autoencoders, SAE, Mechanistic Interpretability, Feature Discovery, Superposition] --- # SAELens: Sparse Autoencoders for Mechanistic Interpretability SAELens is the primary library for training and analyzing Sparse Autoencoders (SAEs) - a technique for decomposing polysemantic neural network activations into sparse, interpretable features. Based on Anthropic's groundbreaking research on monosemanticity. **GitHub**: [jbloomAus/SAELens](https://github.com/jbloomAus/SAELens) (1,100+ stars) ## The Problem: Polysemanticity & Superposition Individual neurons in neural networks are **polysemantic** - they activate in multiple, semantically distinct contexts. This happens because models use **superposition** to represent more features than they have neurons, making interpretability difficult. **SAEs solve this** by decomposing dense activations into sparse, monosemantic features - typically only a small number of features activate for any given input, and each feature corresponds to an interpretable concept. ## When to Use SAELens **Use SAELens when you need to:**- Discover interpretable features in model activations- Understand what concepts a model has learned- Study superposition and feature geometry- Perform feature-based steering or ablation- Analyze safety-relevant features (deception, bias, harmful content) **Consider alternatives when:**- You need basic activation analysis → Use **TransformerLens** directly- You want causal intervention experiments → Use **pyvene** or **TransformerLens**- You need production steering → Consider direct activation engineering ## Installation ```bashpip install sae-lens``` Requirements: Python 3.10+, transformer-lens>=2.0.0 ## Core Concepts ### What SAEs Learn SAEs are trained to reconstruct model activations through a sparse bottleneck: ```Input Activation → Encoder → Sparse Features → Decoder → Reconstructed Activation    (d_model)       ↓        (d_sae >> d_model)    ↓         (d_model)                 sparsity                      reconstruction                 penalty                          loss``` **Loss Function**: `MSE(original, reconstructed) + L1_coefficient × L1(features)` ### Key Validation (Anthropic Research) In "Towards Monosemanticity", human evaluators found **70% of SAE features genuinely interpretable**. Features discovered include:- DNA sequences, legal language, HTTP requests- Hebrew text, nutrition statements, code syntax- Sentiment, named entities, grammatical structures ## Workflow 1: Loading and Analyzing Pre-trained SAEs ### Step-by-Step ```pythonfrom transformer_lens import HookedTransformerfrom sae_lens import SAE # 1. Load model and pre-trained SAEmodel = HookedTransformer.from_pretrained("gpt2-small", device="cuda")# In sae-lens v6, SAE.from_pretrained() returns JUST the SAE (not a tuple).sae = SAE.from_pretrained(    release="gpt2-small-res-jb",    sae_id="blocks.8.hook_resid_pre",    device="cuda")# If you also need the cfg dict and feature sparsity, use:# sae, cfg_dict, sparsity = SAE.from_pretrained_with_cfg_and_sparsity(...) # 2. Get model activationstokens = model.to_tokens("The capital of France is Paris")_, cache = model.run_with_cache(tokens)activations = cache["resid_pre", 8]  # [batch, pos, d_model] # 3. Encode to SAE featuressae_features = sae.encode(activations)  # [batch, pos, d_sae]print(f"Active features: {(sae_features > 0).sum()}") # 4. Find top features for each positionfor pos in range(tokens.shape[1]):    top_features = sae_features[0, pos].topk(5)    token = model.to_str_tokens(tokens[0, pos:pos+1])[0]    print(f"Token '{token}': features {top_features.indices.tolist()}") # 5. Reconstruct activationsreconstructed = sae.decode(sae_features)reconstruction_error = (activations - reconstructed).norm()``` ### Available Pre-trained SAEs | Release | Model | Layers ||---------|-------|--------|| `gpt2-small-res-jb` | GPT-2 Small | Multiple residual streams || `gemma-2b-res` | Gemma 2B | Residual streams || Various on HuggingFace | Search tag `saelens` | Various | ### Checklist- [ ] Load model with TransformerLens- [ ] Load matching SAE for target layer- [ ] Encode activations to sparse features- [ ] Identify top-activating features per token- [ ] Validate reconstruction quality ## Workflow 2: Training a Custom SAE ### Step-by-Step ```pythonfrom sae_lens import (    LanguageModelSAETrainingRunner,    LanguageModelSAERunnerConfig,    StandardTrainingSAEConfig,    LoggingConfig,) # 1. Configure training (v6 uses a NESTED config: SAE-specific options live in a#    `sae=` sub-config, and logging options live in a `logger=` sub-config).#    Note: `architecture`, `d_sae`, `l1_coefficient` etc. are now on the SAE sub-config,#    and legacy flat options like `hook_layer`, `activation_fn`, `log_to_wandb` were removed.cfg = LanguageModelSAERunnerConfig(    # SAE architecture + sparsity (nested)    sae=StandardTrainingSAEConfig(        d_in=768,          # Model dimension        d_sae=768 * 8,     # Expansion factor of 8        l1_coefficient=8e-5,  # Sparsity penalty        apply_b_dec_to_input=True,        normalize_activations="expected_average_only_in",    ),     # Data-generating function (model + hook point)    model_name="gpt2-small",    hook_name="blocks.8.hook_resid_pre",  # layer is inferred from hook_name (no hook_layer)     # Training    lr=4e-4,    l1_warm_up_steps=1000,    train_batch_size_tokens=4096,    training_tokens=100_000_000,     # Data    dataset_path="monology/pile-uncopyrighted",    context_size=128,     # Logging (nested)    logger=LoggingConfig(        log_to_wandb=True,        wandb_project="sae-training",    ),     # Checkpointing    checkpoint_path="checkpoints",    n_checkpoints=5,) # 2. Traintrainer = LanguageModelSAETrainingRunner(cfg)  # SAETrainingRunner still works as an aliassae = trainer.run() # 3. Evaluateprint(f"L0 (avg active features): {trainer.metrics['l0']}")print(f"CE Loss Recovered: {trainer.metrics['ce_loss_score']}")``` > **v6 migration note:** For other SAE types swap the `sae=` sub-config —> `GatedTrainingSAEConfig`, `TopKTrainingSAEConfig` (set `k` directly), or> `JumpReLUTrainingSAEConfig` (uses `l0_coefficient`). Legacy flat options> (`architecture`, `expansion_factor`, `hook_layer`, `activation_fn`/`activation_fn_kwargs`,> `use_ghost_grads`, ghost grads, b_dec/decoder init options) were removed in v6. ### Key Hyperparameters | Parameter | Typical Value | Effect ||-----------|---------------|--------|| `d_sae` | 4-16× d_model | More features, higher capacity || `l1_coefficient` | 5e-5 to 1e-4 | Higher = sparser, less accurate || `lr` | 1e-4 to 1e-3 | Standard optimizer LR || `l1_warm_up_steps` | 500-2000 | Prevents early feature death | ### Evaluation Metrics | Metric | Target | Meaning ||--------|--------|---------|| **L0** | 50-200 | Average active features per token || **CE Loss Score** | 80-95% | Cross-entropy recovered vs original || **Dead Features** | <5% | Features that never activate || **Explained Variance** | >90% | Reconstruction quality | ### Checklist- [ ] Choose target layer and hook point- [ ] Set expansion factor (d_sae = 4-16× d_model)- [ ] Tune L1 coefficient for desired sparsity- [ ] Enable L1 warm-up to prevent dead features- [ ] Monitor metrics during training (W&B)- [ ] Validate L0 and CE loss recovery- [ ] Check dead feature ratio ## Workflow 3: Feature Analysis and Steering ### Analyzing Individual Features ```pythonfrom transformer_lens import HookedTransformerfrom sae_lens import SAEimport torch model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")sae = SAE.from_pretrained(  # v6 returns just the SAE    release="gpt2-small-res-jb",    sae_id="blocks.8.hook_resid_pre",    device="cuda") # Find what activates a specific featurefeature_idx = 1234test_texts = [    "The scientist conducted an experiment",    "I love chocolate cake",    "The code compiles successfully",    "Paris is beautiful in spring",] for text in test_texts:    tokens = model.to_tokens(text)    _, cache = model.run_with_cache(tokens)    features = sae.encode(cache["resid_pre", 8])    activation = features[0, :, feature_idx].max().item()    print(f"{activation:.3f}: {text}")``` ### Feature Steering ```pythondef steer_with_feature(model, sae, prompt, feature_idx, strength=5.0):    """Add SAE feature direction to residual stream."""    tokens = model.to_tokens(prompt)     # Get feature direction from decoder    feature_direction = sae.W_dec[feature_idx]  # [d_model]     def steering_hook(activation, hook):        # Add scaled feature direction at all positions        activation += strength * feature_direction        return activation     # Generate with steering    output = model.generate(        tokens,        max_new_tokens=50,        fwd_hooks=[("blocks.8.hook_resid_pre", steering_hook)]    )    return model.to_string(output[0])``` ### Feature Attribution ```python# Which features most affect a specific output?tokens = model.to_tokens("The capital of France is")_, cache = model.run_with_cache(tokens) # Get features at final positionfeatures = sae.encode(cache["resid_pre", 8])[0, -1]  # [d_sae] # Get logit attribution per feature# Feature contribution = feature_activation × decoder_weight × unembeddingW_dec = sae.W_dec  # [d_sae, d_model]W_U = model.W_U    # [d_model, vocab] # Contribution to "Paris" logitparis_token = model.to_single_token(" Paris")feature_contributions = features * (W_dec @ W_U[:, paris_token]) top_features = feature_contributions.topk(10)print("Top features for 'Paris' prediction:")for idx, val in zip(top_features.indices, top_features.values):    print(f"  Feature {idx.item()}: {val.item():.3f}")``` ## Common Issues & Solutions > All examples below use the v6 nested config: SAE-specific options go in the `sae=`> sub-config (`StandardTrainingSAEConfig` / `TopKTrainingSAEConfig` / etc.), training> knobs stay on the top-level `LanguageModelSAERunnerConfig`. ### Issue: High dead feature ratio```pythonfrom sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig # WRONG: no warm-up, features die earlycfg = LanguageModelSAERunnerConfig(    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),    l1_warm_up_steps=0,  # Bad!) # RIGHT: warm up the L1 penalty (v6 removed ghost grads; warm-up is the lever now)cfg = LanguageModelSAERunnerConfig(    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),    l1_warm_up_steps=1000,  # Gradually increase)``` ### Issue: Poor reconstruction (low CE recovery)```python# Reduce sparsity penalty and/or add capacity (both on the SAE sub-config)cfg = LanguageModelSAERunnerConfig(    sae=StandardTrainingSAEConfig(        d_in=768,        d_sae=768 * 16,       # More capacity        l1_coefficient=5e-5,  # Lower = better reconstruction    ),)``` ### Issue: Features not interpretable```pythonfrom sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig, TopKTrainingSAEConfig # Increase sparsity (higher L1)cfg = LanguageModelSAERunnerConfig(    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),)# Or use a TopK SAE (k is set directly in v6, not via activation_fn_kwargs)cfg = LanguageModelSAERunnerConfig(    sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50),  # Exactly 50 active features)``` ### Issue: Memory errors during training```pythoncfg = LanguageModelSAERunnerConfig(    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),    train_batch_size_tokens=2048,  # Reduce batch size    store_batch_size_prompts=4,    # Fewer prompts in buffer    n_batches_in_buffer=8,         # Smaller activation buffer)``` ## Integration with Neuronpedia Browse pre-trained SAE features at [neuronpedia.org](https://neuronpedia.org): ```python# Features are indexed by SAE ID# Example: gpt2-small layer 8 feature 1234# → neuronpedia.org/gpt2-small/8-res-jb/1234``` ## Key Classes Reference | Class | Purpose ||-------|---------|| `SAE` | Sparse Autoencoder model || `LanguageModelSAERunnerConfig` | Top-level training configuration (nests `sae=` and `logger=`) || `StandardTrainingSAEConfig` / `TopKTrainingSAEConfig` / `GatedTrainingSAEConfig` / `JumpReLUTrainingSAEConfig` | SAE-type-specific sub-configs (v6) || `LoggingConfig` | Logging/W&B sub-config (v6) || `LanguageModelSAETrainingRunner` | Training loop manager (alias: `SAETrainingRunner`) || `ActivationsStore` | Activation collection and batching || `HookedSAETransformer` | TransformerLens + SAE integration | ## Reference Documentation For detailed API documentation, tutorials, and advanced usage, see the `references/` folder: | File | Contents ||------|----------|| [references/README.md](references/README.md) | Overview and quick start guide || [references/api.md](references/api.md) | Complete API reference for SAE, TrainingSAE, configurations || [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for training, analysis, steering | ## External Resources ### Tutorials- [Basic Loading & Analysis](https://github.com/jbloomAus/SAELens/blob/main/tutorials/basic_loading_and_analysing.ipynb)- [Training a Sparse Autoencoder](https://github.com/jbloomAus/SAELens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb)- [ARENA SAE Curriculum](https://www.lesswrong.com/posts/LnHowHgmrMbWtpkxx/intro-to-superposition-and-sparse-autoencoders-colab) ### Papers- [Towards Monosemanticity](https://transformer-circuits.pub/2023/monosemantic-features) - Anthropic (2023)- [Scaling Monosemanticity](https://transformer-circuits.pub/2024/scaling-monosemanticity/) - Anthropic (2024)- [Sparse Autoencoders Find Highly Interpretable Features](https://arxiv.org/abs/2309.08600) - Cunningham et al. (ICLR 2024) ### Official Documentation- [SAELens Docs](https://jbloomaus.github.io/SAELens/)- [Neuronpedia](https://neuronpedia.org) - Feature browser ## SAE Architectures | Architecture | Description | Use Case ||--------------|-------------|----------|| **Standard** | ReLU + L1 penalty | General purpose || **Gated** | Learned gating mechanism | Better sparsity control || **TopK** | Exactly K active features | Consistent sparsity | ```pythonfrom sae_lens import LanguageModelSAERunnerConfig, TopKTrainingSAEConfig # TopK SAE (exactly 50 features active) — `k` is set on the SAE sub-config in v6cfg = LanguageModelSAERunnerConfig(    sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50),)``` 
Discovery context

Discovered by repository scan. No exact path reference found in the snapshot’s root AGENTS.md.