SKILL.md
SKILL.mdBrowse 4 files
3,940 tokens
15,003 bytes
Token encoding: o200k_base
Snapshot 24fd22b
1---2name: saelens3description: Train sparse autoencoders to interpret model features.4version: 1.0.15author: Orchestra Research6license: MIT7dependencies: [sae-lens>=6.0.0, transformer-lens>=2.0.0, torch>=2.0.0]8platforms: [linux, macos, windows]9metadata:10 hermes:11 tags: [Sparse Autoencoders, SAE, Mechanistic Interpretability, Feature Discovery, Superposition]12 13---14 15# SAELens: Sparse Autoencoders for Mechanistic Interpretability16 17SAELens 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.18 19**GitHub**: [jbloomAus/SAELens](https://github.com/jbloomAus/SAELens) (1,100+ stars)20 21## The Problem: Polysemanticity & Superposition22 23Individual 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.24 25**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.26 27## When to Use SAELens28 29**Use SAELens when you need to:**30- Discover interpretable features in model activations31- Understand what concepts a model has learned32- Study superposition and feature geometry33- Perform feature-based steering or ablation34- Analyze safety-relevant features (deception, bias, harmful content)35 36**Consider alternatives when:**37- You need basic activation analysis → Use **TransformerLens** directly38- You want causal intervention experiments → Use **pyvene** or **TransformerLens**39- You need production steering → Consider direct activation engineering40 41## Installation42 43```bash44pip install sae-lens45```46 47Requirements: Python 3.10+, transformer-lens>=2.0.048 49## Core Concepts50 51### What SAEs Learn52 53SAEs are trained to reconstruct model activations through a sparse bottleneck:54 55```56Input Activation → Encoder → Sparse Features → Decoder → Reconstructed Activation57 (d_model) ↓ (d_sae >> d_model) ↓ (d_model)58 sparsity reconstruction59 penalty loss60```61 62**Loss Function**: `MSE(original, reconstructed) + L1_coefficient × L1(features)`63 64### Key Validation (Anthropic Research)65 66In "Towards Monosemanticity", human evaluators found **70% of SAE features genuinely interpretable**. Features discovered include:67- DNA sequences, legal language, HTTP requests68- Hebrew text, nutrition statements, code syntax69- Sentiment, named entities, grammatical structures70 71## Workflow 1: Loading and Analyzing Pre-trained SAEs72 73### Step-by-Step74 75```python76from transformer_lens import HookedTransformer77from sae_lens import SAE78 79# 1. Load model and pre-trained SAE80model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")81# In sae-lens v6, SAE.from_pretrained() returns JUST the SAE (not a tuple).82sae = SAE.from_pretrained(83 release="gpt2-small-res-jb",84 sae_id="blocks.8.hook_resid_pre",85 device="cuda"86)87# If you also need the cfg dict and feature sparsity, use:88# sae, cfg_dict, sparsity = SAE.from_pretrained_with_cfg_and_sparsity(...)89 90# 2. Get model activations91tokens = model.to_tokens("The capital of France is Paris")92_, cache = model.run_with_cache(tokens)93activations = cache["resid_pre", 8] # [batch, pos, d_model]94 95# 3. Encode to SAE features96sae_features = sae.encode(activations) # [batch, pos, d_sae]97print(f"Active features: {(sae_features > 0).sum()}")98 99# 4. Find top features for each position100for pos in range(tokens.shape[1]):101 top_features = sae_features[0, pos].topk(5)102 token = model.to_str_tokens(tokens[0, pos:pos+1])[0]103 print(f"Token '{token}': features {top_features.indices.tolist()}")104 105# 5. Reconstruct activations106reconstructed = sae.decode(sae_features)107reconstruction_error = (activations - reconstructed).norm()108```109 110### Available Pre-trained SAEs111 112| Release | Model | Layers |113|---------|-------|--------|114| `gpt2-small-res-jb` | GPT-2 Small | Multiple residual streams |115| `gemma-2b-res` | Gemma 2B | Residual streams |116| Various on HuggingFace | Search tag `saelens` | Various |117 118### Checklist119- [ ] Load model with TransformerLens120- [ ] Load matching SAE for target layer121- [ ] Encode activations to sparse features122- [ ] Identify top-activating features per token123- [ ] Validate reconstruction quality124 125## Workflow 2: Training a Custom SAE126 127### Step-by-Step128 129```python130from sae_lens import (131 LanguageModelSAETrainingRunner,132 LanguageModelSAERunnerConfig,133 StandardTrainingSAEConfig,134 LoggingConfig,135)136 137# 1. Configure training (v6 uses a NESTED config: SAE-specific options live in a138# `sae=` sub-config, and logging options live in a `logger=` sub-config).139# Note: `architecture`, `d_sae`, `l1_coefficient` etc. are now on the SAE sub-config,140# and legacy flat options like `hook_layer`, `activation_fn`, `log_to_wandb` were removed.141cfg = LanguageModelSAERunnerConfig(142 # SAE architecture + sparsity (nested)143 sae=StandardTrainingSAEConfig(144 d_in=768, # Model dimension145 d_sae=768 * 8, # Expansion factor of 8146 l1_coefficient=8e-5, # Sparsity penalty147 apply_b_dec_to_input=True,148 normalize_activations="expected_average_only_in",149 ),150 151 # Data-generating function (model + hook point)152 model_name="gpt2-small",153 hook_name="blocks.8.hook_resid_pre", # layer is inferred from hook_name (no hook_layer)154 155 # Training156 lr=4e-4,157 l1_warm_up_steps=1000,158 train_batch_size_tokens=4096,159 training_tokens=100_000_000,160 161 # Data162 dataset_path="monology/pile-uncopyrighted",163 context_size=128,164 165 # Logging (nested)166 logger=LoggingConfig(167 log_to_wandb=True,168 wandb_project="sae-training",169 ),170 171 # Checkpointing172 checkpoint_path="checkpoints",173 n_checkpoints=5,174)175 176# 2. Train177trainer = LanguageModelSAETrainingRunner(cfg) # SAETrainingRunner still works as an alias178sae = trainer.run()179 180# 3. Evaluate181print(f"L0 (avg active features): {trainer.metrics['l0']}")182print(f"CE Loss Recovered: {trainer.metrics['ce_loss_score']}")183```184 185> **v6 migration note:** For other SAE types swap the `sae=` sub-config —186> `GatedTrainingSAEConfig`, `TopKTrainingSAEConfig` (set `k` directly), or187> `JumpReLUTrainingSAEConfig` (uses `l0_coefficient`). Legacy flat options188> (`architecture`, `expansion_factor`, `hook_layer`, `activation_fn`/`activation_fn_kwargs`,189> `use_ghost_grads`, ghost grads, b_dec/decoder init options) were removed in v6.190 191### Key Hyperparameters192 193| Parameter | Typical Value | Effect |194|-----------|---------------|--------|195| `d_sae` | 4-16× d_model | More features, higher capacity |196| `l1_coefficient` | 5e-5 to 1e-4 | Higher = sparser, less accurate |197| `lr` | 1e-4 to 1e-3 | Standard optimizer LR |198| `l1_warm_up_steps` | 500-2000 | Prevents early feature death |199 200### Evaluation Metrics201 202| Metric | Target | Meaning |203|--------|--------|---------|204| **L0** | 50-200 | Average active features per token |205| **CE Loss Score** | 80-95% | Cross-entropy recovered vs original |206| **Dead Features** | <5% | Features that never activate |207| **Explained Variance** | >90% | Reconstruction quality |208 209### Checklist210- [ ] Choose target layer and hook point211- [ ] Set expansion factor (d_sae = 4-16× d_model)212- [ ] Tune L1 coefficient for desired sparsity213- [ ] Enable L1 warm-up to prevent dead features214- [ ] Monitor metrics during training (W&B)215- [ ] Validate L0 and CE loss recovery216- [ ] Check dead feature ratio217 218## Workflow 3: Feature Analysis and Steering219 220### Analyzing Individual Features221 222```python223from transformer_lens import HookedTransformer224from sae_lens import SAE225import torch226 227model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")228sae = SAE.from_pretrained( # v6 returns just the SAE229 release="gpt2-small-res-jb",230 sae_id="blocks.8.hook_resid_pre",231 device="cuda"232)233 234# Find what activates a specific feature235feature_idx = 1234236test_texts = [237 "The scientist conducted an experiment",238 "I love chocolate cake",239 "The code compiles successfully",240 "Paris is beautiful in spring",241]242 243for text in test_texts:244 tokens = model.to_tokens(text)245 _, cache = model.run_with_cache(tokens)246 features = sae.encode(cache["resid_pre", 8])247 activation = features[0, :, feature_idx].max().item()248 print(f"{activation:.3f}: {text}")249```250 251### Feature Steering252 253```python254def steer_with_feature(model, sae, prompt, feature_idx, strength=5.0):255 """Add SAE feature direction to residual stream."""256 tokens = model.to_tokens(prompt)257 258 # Get feature direction from decoder259 feature_direction = sae.W_dec[feature_idx] # [d_model]260 261 def steering_hook(activation, hook):262 # Add scaled feature direction at all positions263 activation += strength * feature_direction264 return activation265 266 # Generate with steering267 output = model.generate(268 tokens,269 max_new_tokens=50,270 fwd_hooks=[("blocks.8.hook_resid_pre", steering_hook)]271 )272 return model.to_string(output[0])273```274 275### Feature Attribution276 277```python278# Which features most affect a specific output?279tokens = model.to_tokens("The capital of France is")280_, cache = model.run_with_cache(tokens)281 282# Get features at final position283features = sae.encode(cache["resid_pre", 8])[0, -1] # [d_sae]284 285# Get logit attribution per feature286# Feature contribution = feature_activation × decoder_weight × unembedding287W_dec = sae.W_dec # [d_sae, d_model]288W_U = model.W_U # [d_model, vocab]289 290# Contribution to "Paris" logit291paris_token = model.to_single_token(" Paris")292feature_contributions = features * (W_dec @ W_U[:, paris_token])293 294top_features = feature_contributions.topk(10)295print("Top features for 'Paris' prediction:")296for idx, val in zip(top_features.indices, top_features.values):297 print(f" Feature {idx.item()}: {val.item():.3f}")298```299 300## Common Issues & Solutions301 302> All examples below use the v6 nested config: SAE-specific options go in the `sae=`303> sub-config (`StandardTrainingSAEConfig` / `TopKTrainingSAEConfig` / etc.), training304> knobs stay on the top-level `LanguageModelSAERunnerConfig`.305 306### Issue: High dead feature ratio307```python308from sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig309 310# WRONG: no warm-up, features die early311cfg = LanguageModelSAERunnerConfig(312 sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),313 l1_warm_up_steps=0, # Bad!314)315 316# RIGHT: warm up the L1 penalty (v6 removed ghost grads; warm-up is the lever now)317cfg = LanguageModelSAERunnerConfig(318 sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),319 l1_warm_up_steps=1000, # Gradually increase320)321```322 323### Issue: Poor reconstruction (low CE recovery)324```python325# Reduce sparsity penalty and/or add capacity (both on the SAE sub-config)326cfg = LanguageModelSAERunnerConfig(327 sae=StandardTrainingSAEConfig(328 d_in=768,329 d_sae=768 * 16, # More capacity330 l1_coefficient=5e-5, # Lower = better reconstruction331 ),332)333```334 335### Issue: Features not interpretable336```python337from sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig, TopKTrainingSAEConfig338 339# Increase sparsity (higher L1)340cfg = LanguageModelSAERunnerConfig(341 sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),342)343# Or use a TopK SAE (k is set directly in v6, not via activation_fn_kwargs)344cfg = LanguageModelSAERunnerConfig(345 sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50), # Exactly 50 active features346)347```348 349### Issue: Memory errors during training350```python351cfg = LanguageModelSAERunnerConfig(352 sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),353 train_batch_size_tokens=2048, # Reduce batch size354 store_batch_size_prompts=4, # Fewer prompts in buffer355 n_batches_in_buffer=8, # Smaller activation buffer356)357```358 359## Integration with Neuronpedia360 361Browse pre-trained SAE features at [neuronpedia.org](https://neuronpedia.org):362 363```python364# Features are indexed by SAE ID365# Example: gpt2-small layer 8 feature 1234366# → neuronpedia.org/gpt2-small/8-res-jb/1234367```368 369## Key Classes Reference370 371| Class | Purpose |372|-------|---------|373| `SAE` | Sparse Autoencoder model |374| `LanguageModelSAERunnerConfig` | Top-level training configuration (nests `sae=` and `logger=`) |375| `StandardTrainingSAEConfig` / `TopKTrainingSAEConfig` / `GatedTrainingSAEConfig` / `JumpReLUTrainingSAEConfig` | SAE-type-specific sub-configs (v6) |376| `LoggingConfig` | Logging/W&B sub-config (v6) |377| `LanguageModelSAETrainingRunner` | Training loop manager (alias: `SAETrainingRunner`) |378| `ActivationsStore` | Activation collection and batching |379| `HookedSAETransformer` | TransformerLens + SAE integration |380 381## Reference Documentation382 383For detailed API documentation, tutorials, and advanced usage, see the `references/` folder:384 385| File | Contents |386|------|----------|387| [references/README.md](references/README.md) | Overview and quick start guide |388| [references/api.md](references/api.md) | Complete API reference for SAE, TrainingSAE, configurations |389| [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for training, analysis, steering |390 391## External Resources392 393### Tutorials394- [Basic Loading & Analysis](https://github.com/jbloomAus/SAELens/blob/main/tutorials/basic_loading_and_analysing.ipynb)395- [Training a Sparse Autoencoder](https://github.com/jbloomAus/SAELens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb)396- [ARENA SAE Curriculum](https://www.lesswrong.com/posts/LnHowHgmrMbWtpkxx/intro-to-superposition-and-sparse-autoencoders-colab)397 398### Papers399- [Towards Monosemanticity](https://transformer-circuits.pub/2023/monosemantic-features) - Anthropic (2023)400- [Scaling Monosemanticity](https://transformer-circuits.pub/2024/scaling-monosemanticity/) - Anthropic (2024)401- [Sparse Autoencoders Find Highly Interpretable Features](https://arxiv.org/abs/2309.08600) - Cunningham et al. (ICLR 2024)402 403### Official Documentation404- [SAELens Docs](https://jbloomaus.github.io/SAELens/)405- [Neuronpedia](https://neuronpedia.org) - Feature browser406 407## SAE Architectures408 409| Architecture | Description | Use Case |410|--------------|-------------|----------|411| **Standard** | ReLU + L1 penalty | General purpose |412| **Gated** | Learned gating mechanism | Better sparsity control |413| **TopK** | Exactly K active features | Consistent sparsity |414 415```python416from sae_lens import LanguageModelSAERunnerConfig, TopKTrainingSAEConfig417 418# TopK SAE (exactly 50 features active) — `k` is set on the SAE sub-config in v6419cfg = LanguageModelSAERunnerConfig(420 sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50),421)422```423 Discovery context
Discovered by repository scan. No exact path reference found in the snapshot’s root AGENTS.md.