{"id":"saelens","name":"sparse-autoencoder-training","summary":"SAELensを用いたスパースオートエンコーダー(SAE)の訓練および解析に関する指針を提供し、ニューラルネットワークの活性化を解釈可能な特徴に分解する。","body":"# SAELens: Sparse Autoencoders for Mechanistic Interpretability\n\nSAELens 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.\n\n**GitHub**: [jbloomAus/SAELens](https://github.com/jbloomAus/SAELens) (1,100+ stars)\n\n## The Problem: Polysemanticity & Superposition\n\nIndividual 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.\n\n**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.\n\n## When to Use SAELens\n\n**Use SAELens when you need to:**\n- Discover interpretable features in model activations\n- Understand what concepts a model has learned\n- Study superposition and feature geometry\n- Perform feature-based steering or ablation\n- Analyze safety-relevant features (deception, bias, harmful content)\n\n**Consider alternatives when:**\n- You need basic activation analysis → Use **TransformerLens** directly\n- You want causal intervention experiments → Use **pyvene** or **TransformerLens**\n- You need production steering → Consider direct activation engineering\n\n## Installation\n\n```bash\npip install sae-lens\n```\n\nRequirements: Python 3.10+, transformer-lens>=2.0.0\n\n## Core Concepts\n\n### What SAEs Learn\n\nSAEs are trained to reconstruct model activations through a sparse bottleneck:\n\n```\nInput Activation → Encoder → Sparse Features → Decoder → Reconstructed Activation\n    (d_model)       ↓        (d_sae >> d_model)    ↓         (d_model)\n                 sparsity                      reconstruction\n                 penalty                          loss\n```\n\n**Loss Function**: `MSE(original, reconstructed) + L1_coefficient × L1(features)`\n\n### Key Validation (Anthropic Research)\n\nIn \"Towards Monosemanticity\", human evaluators found **70% of SAE features genuinely interpretable**. Features discovered include:\n- DNA sequences, legal language, HTTP requests\n- Hebrew text, nutrition statements, code syntax\n- Sentiment, named entities, grammatical structures\n\n## Workflow 1: Loading and Analyzing Pre-trained SAEs\n\n### Step-by-Step\n\n```python\nfrom transformer_lens import HookedTransformer\nfrom sae_lens import SAE\n\n# 1. Load model and pre-trained SAE\nmodel = HookedTransformer.from_pretrained(\"gpt2-small\", device=\"cuda\")\nsae, cfg_dict, sparsity = SAE.from_pretrained(\n    release=\"gpt2-small-res-jb\",\n    sae_id=\"blocks.8.hook_resid_pre\",\n    device=\"cuda\"\n)\n\n# 2. Get model activations\ntokens = model.to_tokens(\"The capital of France is Paris\")\n_, cache = model.run_with_cache(tokens)\nactivations = cache[\"resid_pre\", 8]  # [batch, pos, d_model]\n\n# 3. Encode to SAE features\nsae_features = sae.encode(activations)  # [batch, pos, d_sae]\nprint(f\"Active features: {(sae_features > 0).sum()}\")\n\n# 4. Find top features for each position\nfor pos in range(tokens.shape[1]):\n    top_features = sae_features[0, pos].topk(5)\n    token = model.to_str_tokens(tokens[0, pos:pos+1])[0]\n    print(f\"Token '{token}': features {top_features.indices.tolist()}\")\n\n# 5. Reconstruct activations\nreconstructed = sae.decode(sae_features)\nreconstruction_error = (activations - reconstructed).norm()\n```\n\n### Available Pre-trained SAEs\n\n| Release | Model | Layers |\n|---------|-------|--------|\n| `gpt2-small-res-jb` | GPT-2 Small | Multiple residual streams |\n| `gemma-2b-res` | Gemma 2B | Residual streams |\n| Various on HuggingFace | Search tag `saelens` | Various |\n\n### Checklist\n- [ ] Load model with TransformerLens\n- [ ] Load matching SAE for target layer\n- [ ] Encode activations to sparse features\n- [ ] Identify top-activating features per token\n- [ ] Validate reconstruction quality\n\n## Workflow 2: Training a Custom SAE\n\n### Step-by-Step\n\n```python\nfrom sae_lens import SAE, LanguageModelSAERunnerConfig, SAETrainingRunner\n\n# 1. Configure training\ncfg = LanguageModelSAERunnerConfig(\n    # Model\n    model_name=\"gpt2-small\",\n    hook_name=\"blocks.8.hook_resid_pre\",\n    hook_layer=8,\n    d_in=768,  # Model dimension\n\n    # SAE architecture\n    architecture=\"standard\",  # or \"gated\", \"topk\"\n    d_sae=768 * 8,  # Expansion factor of 8\n    activation_fn=\"relu\",\n\n    # Training\n    lr=4e-4,\n    l1_coefficient=8e-5,  # Sparsity penalty\n    l1_warm_up_steps=1000,\n    train_batch_size_tokens=4096,\n    training_tokens=100_000_000,\n\n    # Data\n    dataset_path=\"monology/pile-uncopyrighted\",\n    context_size=128,\n\n    # Logging\n    log_to_wandb=True,\n    wandb_project=\"sae-training\",\n\n    # Checkpointing\n    checkpoint_path=\"checkpoints\",\n    n_checkpoints=5,\n)\n\n# 2. Train\ntrainer = SAETrainingRunner(cfg)\nsae = trainer.run()\n\n# 3. Evaluate\nprint(f\"L0 (avg active features): {trainer.metrics['l0']}\")\nprint(f\"CE Loss Recovered: {trainer.metrics['ce_loss_score']}\")\n```\n\n### Key Hyperparameters\n\n| Parameter | Typical Value | Effect |\n|-----------|---------------|--------|\n| `d_sae` | 4-16× d_model | More features, higher capacity |\n| `l1_coefficient` | 5e-5 to 1e-4 | Higher = sparser, less accurate |\n| `lr` | 1e-4 to 1e-3 | Standard optimizer LR |\n| `l1_warm_up_steps` | 500-2000 | Prevents early feature death |\n\n### Evaluation Metrics\n\n| Metric | Target | Meaning |\n|--------|--------|---------|\n| **L0** | 50-200 | Average active features per token |\n| **CE Loss Score** | 80-95% | Cross-entropy recovered vs original |\n| **Dead Features** | <5% | Features that never activate |\n| **Explained Variance** | >90% | Reconstruction quality |\n\n### Checklist\n- [ ] Choose target layer and hook point\n- [ ] Set expansion factor (d_sae = 4-16× d_model)\n- [ ] Tune L1 coefficient for desired sparsity\n- [ ] Enable L1 warm-up to prevent dead features\n- [ ] Monitor metrics during training (W&B)\n- [ ] Validate L0 and CE loss recovery\n- [ ] Check dead feature ratio\n\n## Workflow 3: Feature Analysis and Steering\n\n### Analyzing Individual Features\n\n```python\nfrom transformer_lens import HookedTransformer\nfrom sae_lens import SAE\nimport torch\n\nmodel = HookedTransformer.from_pretrained(\"gpt2-small\", device=\"cuda\")\nsae, _, _ = SAE.from_pretrained(\n    release=\"gpt2-small-res-jb\",\n    sae_id=\"blocks.8.hook_resid_pre\",\n    device=\"cuda\"\n)\n\n# Find what activates a specific feature\nfeature_idx = 1234\ntest_texts = [\n    \"The scientist conducted an experiment\",\n    \"I love chocolate cake\",\n    \"The code compiles successfully\",\n    \"Paris is beautiful in spring\",\n]\n\nfor text in test_texts:\n    tokens = model.to_tokens(text)\n    _, cache = model.run_with_cache(tokens)\n    features = sae.encode(cache[\"resid_pre\", 8])\n    activation = features[0, :, feature_idx].max().item()\n    print(f\"{activation:.3f}: {text}\")\n```\n\n### Feature Steering\n\n```python\ndef steer_with_feature(model, sae, prompt, feature_idx, strength=5.0):\n    \"\"\"Add SAE feature direction to residual stream.\"\"\"\n    tokens = model.to_tokens(prompt)\n\n    # Get feature direction from decoder\n    feature_direction = sae.W_dec[feature_idx]  # [d_model]\n\n    def steering_hook(activation, hook):\n        # Add scaled feature direction at all positions\n        activation += strength * feature_direction\n        return activation\n\n    # Generate with steering\n    output = model.generate(\n        tokens,\n        max_new_tokens=50,\n        fwd_hooks=[(\"blocks.8.hook_resid_pre\", steering_hook)]\n    )\n    return model.to_string(output[0])\n```\n\n### Feature Attribution\n\n```python\n# Which features most affect a specific output?\ntokens = model.to_tokens(\"The capital of France is\")\n_, cache = model.run_with_cache(tokens)\n\n# Get features at final position\nfeatures = sae.encode(cache[\"resid_pre\", 8])[0, -1]  # [d_sae]\n\n# Get logit attribution per feature\n# Feature contribution = feature_activation × decoder_weight × unembedding\nW_dec = sae.W_dec  # [d_sae, d_model]\nW_U = model.W_U    # [d_model, vocab]\n\n# Contribution to \"Paris\" logit\nparis_token = model.to_single_token(\" Paris\")\nfeature_contributions = features * (W_dec @ W_U[:, paris_token])\n\ntop_features = feature_contributions.topk(10)\nprint(\"Top features for 'Paris' prediction:\")\nfor idx, val in zip(top_features.indices, top_features.values):\n    print(f\"  Feature {idx.item()}: {val.item():.3f}\")\n```\n\n## Common Issues & Solutions\n\n### Issue: High dead feature ratio\n```python\n# WRONG: No warm-up, features die early\ncfg = LanguageModelSAERunnerConfig(\n    l1_coefficient=1e-4,\n    l1_warm_up_steps=0,  # Bad!\n)\n\n# RIGHT: Warm-up L1 penalty\ncfg = LanguageModelSAERunnerConfig(\n    l1_coefficient=8e-5,\n    l1_warm_up_steps=1000,  # Gradually increase\n    use_ghost_grads=True,   # Revive dead features\n)\n```\n\n### Issue: Poor reconstruction (low CE recovery)\n```python\n# Reduce sparsity penalty\ncfg = LanguageModelSAERunnerConfig(\n    l1_coefficient=5e-5,  # Lower = better reconstruction\n    d_sae=768 * 16,       # More capacity\n)\n```\n\n### Issue: Features not interpretable\n```python\n# Increase sparsity (higher L1)\ncfg = LanguageModelSAERunnerConfig(\n    l1_coefficient=1e-4,  # Higher = sparser, more interpretable\n)\n# Or use TopK architecture\ncfg = LanguageModelSAERunnerConfig(\n    architecture=\"topk\",\n    activation_fn_kwargs={\"k\": 50},  # Exactly 50 active features\n)\n```\n\n### Issue: Memory errors during training\n```python\ncfg = LanguageModelSAERunnerConfig(\n    train_batch_size_tokens=2048,  # Reduce batch size\n    store_batch_size_prompts=4,    # Fewer prompts in buffer\n    n_batches_in_buffer=8,         # Smaller activation buffer\n)\n```\n\n## Integration with Neuronpedia\n\nBrowse pre-trained SAE features at [neuronpedia.org](https://neuronpedia.org):\n\n```python\n# Features are indexed by SAE ID\n# Example: gpt2-small layer 8 feature 1234\n# → neuronpedia.org/gpt2-small/8-res-jb/1234\n```\n\n## Key Classes Reference\n\n| Class | Purpose |\n|-------|---------|\n| `SAE` | Sparse Autoencoder model |\n| `LanguageModelSAERunnerConfig` | Training configuration |\n| `SAETrainingRunner` | Training loop manager |\n| `ActivationsStore` | Activation collection and batching |\n| `HookedSAETransformer` | TransformerLens + SAE integration |\n\n## Reference Documentation\n\nFor detailed API documentation, tutorials, and advanced usage, see the `references/` folder:\n\n| File | Contents |\n|------|----------|\n| [references/README.md](references/README.md) | Overview and quick start guide |\n| [references/api.md](references/api.md) | Complete API reference for SAE, TrainingSAE, configurations |\n| [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for training, analysis, steering |\n\n## External Resources\n\n### Tutorials\n- [Basic Loading & Analysis](https://github.com/jbloomAus/SAELens/blob/main/tutorials/basic_loading_and_analysing.ipynb)\n- [Training a Sparse Autoencoder](https://github.com/jbloomAus/SAELens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb)\n- [ARENA SAE Curriculum](https://www.lesswrong.com/posts/LnHowHgmrMbWtpkxx/intro-to-superposition-and-sparse-autoencoders-colab)\n\n### Papers\n- [Towards Monosemanticity](https://transformer-circuits.pub/2023/monosemantic-features) - Anthropic (2023)\n- [Scaling Monosemanticity](https://transformer-circuits.pub/2024/scaling-monosemanticity/) - Anthropic (2024)\n- [Sparse Autoencoders Find Highly Interpretable Features](https://arxiv.org/abs/2309.08600) - Cunningham et al. (ICLR 2024)\n\n### Official Documentation\n- [SAELens Docs](https://jbloomaus.github.io/SAELens/)\n- [Neuronpedia](https://neuronpedia.org) - Feature browser\n\n## SAE Architectures\n\n| Architecture | Description | Use Case |\n|--------------|-------------|----------|\n| **Standard** | ReLU + L1 penalty | General purpose |\n| **Gated** | Learned gating mechanism | Better sparsity control |\n| **TopK** | Exactly K active features | Consistent sparsity |\n\n```python\n# TopK SAE (exactly 50 features active)\ncfg = LanguageModelSAERunnerConfig(\n    architecture=\"topk\",\n    activation_fn=\"topk\",\n    activation_fn_kwargs={\"k\": 50},\n)\n```","author":"@Orchestra-Research","ownerProfile":null,"authorContacts":null,"sourceUrl":"https://github.com/Orchestra-Research/AI-Research-SKILLs/tree/main/04-mechanistic-interpretability/saelens","license":"MIT","category":"research","lang":"en","tokens":3137,"stars":0,"calls30d":2,"claimed":false,"visibility":"public","origin":"crawler","version":"0.1.0","createdAt":"2026-08-22","updatedAt":"2026-08-22","files":[{"path":"references/api.md","size":6968,"sha256":"528edc8369d0c8b419b037eae65a4363d6f41f9671c7f7f739327700393a5a35"},{"path":"references/README.md","size":2155,"sha256":"1ae9c3f499b5db57c344f05a654a8ab012d10dd59c059f55615e0ada28f349e0"},{"path":"references/tutorials.md","size":9383,"sha256":"7b4c7a2ba9a11116886af33ea11cf3e279e3994a171e4632363114eaa46000fb"}],"requires":{"mcp":[],"tools":[]},"safety":{"flags":[],"scannedAt":"2026-08-22","hasScripts":false,"networkEndpoints":["arxiv.org","jbloomaus.github.io","neuronpedia.org","transformer-circuits.pub","www.lesswrong.com"]}}