Training and Using Sparse Autoencoders with SAELens in Hermes Agent

{ "title": "Sparse Autoencoder Training and Analysis with SAE-Lens", "body": "## Sparse Autoencoder Training and Analysis with SAE-Lens\n\nThis document describes how to train and analyze Sparse Autoencoders (SAEs) to decompose polysemantic neural network activations into sparse, interpretable features. SAEs help discover interpretable features in model activations, understand what concepts a model has learned, study superposition and feature geometry, perform feature-based steering or ablation, and analyze safety-relevant features such as deception, bias, or harmful content.\n\n## Prerequisites\n\n- Python 3.10 or later\n- transformer-lens>=2.0.0\n- sae-lens installed via pip:\n \n pip install sae-lens\n \n- Access to a transformer model via TransformerLens (for example, gpt2-small)\n- For training: a dataset such as monology/pile-uncopyrighted and sufficient GPU memory\n- For logging: a Weights & Biases account (optional)\n\n## Architecture Overview\n\nAn SAE is a two-layer neural network that maps input activations to sparse features and then reconstructs the original activation. The flow is:\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\nThe loss function combines a reconstruction term and a sparsity penalty:\nMSE(original, reconstructed) + L1_coefficient × L1(features)\n\n## Loading and Analyzing Pre-trained SAEs\n\nPre-trained SAEs are available for GPT-2 Small (multiple residual streams) and Gemma 2B (residual streams). Search the HuggingFace tag saelens for more models. Neuronpedia feature URLs follow the pattern neuronpedia.org/{model}/{layer}-res-jb/{feature_id}.\n\n### Procedure\n\n1. Load the model with HookedTransformer.from_pretrained().\n2. Load the matching SAE with SAE.from_pretrained(release, sae_id, device). In v6, this returns just the SAE object, not a tuple. If you need the config dictionary and feature sparsity, use SAE.from_pretrained_with_cfg_and_sparsity().\n3. Get model activations: tokens = model.to_tokens(text); _, cache = model.run_with_cache(tokens); activations = cache['resid_pre', layer].\n4. Encode activations to SAE features: sae_features = sae.encode(activations).\n5. Find top features per position: sae_features[0, pos].topk(5).\n6. Reconstruct activations: reconstructed = sae.decode(sae_features); reconstruction_error = (activations - reconstructed).norm().\n\n### Example\n\npython\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\")\n# In sae-lens v6, SAE.from_pretrained() returns JUST the SAE (not a tuple).\nsae = SAE.from_pretrained(\n release=\"gpt2-small-res-jb\",\n sae_id=\"blocks.8.hook_resid_pre\",\n device=\"cuda\"\n)\n# If you also need the cfg dict and feature sparsity, use:\n# sae, cfg_dict, sparsity = SAE.from_pretrained_with_cfg_and_sparsity(...)\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## Training a Custom SAE\n\nTraining a custom SAE requires configuring a LanguageModelSAERunnerConfig with nested sub-configs: SAE-specific options go in sae= (for example, StandardTrainingSAEConfig), and logging options go in logger= (a LoggingConfig). Legacy flat options such as architecture, expansion_factor, hook_layer, activation_fn, activation_fn_kwargs, use_ghost_grads, and log_to_wandb have been removed in v6. The layer is inferred from hook_name; the hook_layer parameter is removed. Ghost grads are removed; use l1_warm_up_steps to prevent dead features.\n\n### Parameters\n\n| Parameter | Meaning | Required |\n|---|---|---|\n| d_sae | Number of features in the SAE (expansion factor × d_model) | Yes |\n| l1_coefficient | Sparsity penalty weight in loss function | Yes |\n| lr | Learning rate for optimizer | Yes |\n| l1_warm_up_steps | Steps to gradually increase L1 penalty from 0 to full value | No |\n| train_batch_size_tokens | Number of tokens per training batch | Yes |\n| training_tokens | Total number of tokens to train on | Yes |\n| context_size | Sequence length for training data | Yes |\n| k | Number of active features for TopK SAE (set on TopKTrainingSAEConfig) | For TopK architecture |\n| apply_b_dec_to_input | Whether to apply decoder bias to input before encoding | No |\n| normalize_activations | Strategy for normalizing activations (e.g., 'expected_average_only_in') | No |\n| log_to_wandb | Whether to log metrics to Weights & Biases | No |\n| wandb_project | W&B project name for logging | If log_to_wandb=True |\n| checkpoint_path | Directory to save model checkpoints | No |\n| n_checkpoints | Number of checkpoints to save during training | No |\n| store_batch_size_prompts | Number of prompts to store in activation buffer | No |\n| n_batches_in_buffer | Number of batches to keep in activation buffer | No |\n\n### Procedure\n\n1. Configure training with LanguageModelSAERunnerConfig, nesting SAE-specific options in sae= and logging in logger=.\n2. Set model_name, hook_name (layer inferred from hook_name), lr, l1_warm_up_steps, train_batch_size_tokens, training_tokens, dataset_path, context_size, checkpoint_path, n_checkpoints.\n3. Create trainer: trainer = LanguageModelSAETrainingRunner(cfg).\n4. Run training: sae = trainer.run().\n5. Evaluate: print(trainer.metrics['l0'], trainer.metrics['ce_loss_score']).\n\n### Example\n\npython\nfrom sae_lens import (\n LanguageModelSAETrainingRunner,\n LanguageModelSAERunnerConfig,\n StandardTrainingSAEConfig,\n LoggingConfig,\n)\n\n# 1. Configure training (v6 uses a NESTED config: SAE-specific options live in a\n# `sae=` sub-config, and logging options live in a `logger=` sub-config).\n# Note: `architecture`, `d_sae`, `l1_coefficient` etc. are now on the SAE sub-config,\n# and legacy flat options like `hook_layer`, `activation_fn`, `log_to_wandb` were removed.\ncfg = LanguageModelSAERunnerConfig(\n # SAE architecture + sparsity (nested)\n sae=StandardTrainingSAEConfig(\n d_in=768, # Model dimension\n d_sae=768 * 8, # Expansion factor of 8\n l1_coefficient=8e-5, # Sparsity penalty\n apply_b_dec_to_input=True,\n normalize_activations=\"expected_average_only_in\",\n ),\n\n # Data-generating function (model + hook point)\n model_name=\"gpt2-small\",\n hook_name=\"blocks.8.hook_resid_pre\", # layer is inferred from hook_name (no hook_layer)\n\n # Training\n lr=4e-4,\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 (nested)\n logger=LoggingConfig(\n log_to_wandb=True,\n wandb_project=\"sae-training\",\n ),\n\n # Checkpointing\n checkpoint_path=\"checkpoints\",\n n_checkpoints=5,\n)\n\n# 2. Train\ntrainer = LanguageModelSAETrainingRunner(cfg) # SAETrainingRunner still works as an alias\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\nFor other SAE types in v6, swap the sae= sub-config: use GatedTrainingSAEConfig, TopKTrainingSAEConfig (set k directly), or JumpReLUTrainingSAEConfig (uses l0_coefficient).\n\n## Feature Analysis and Steering\n\n### Finding What Activates a Specific Feature\n\nLoad the model and SAE, then encode activations for test texts and extract activation values for a specific feature index.\n\npython\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( # v6 returns just the SAE\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### Steering Model Outputs with Features\n\nDefine a steering hook that adds strength * sae.W_dec[feature_idx] to the activation, then call model.generate() with fwd_hooks.\n\npython\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### Attributing Output Logits to Features\n\nCompute features at the final position, then compute feature contributions as features * (sae.W_dec @ model.W_U[:, target_token]). Find the top contributing features.\n\npython\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## Failure Modes and Solutions\n\n### High Dead Feature Ratio\n\nCaused by no L1 warm-up (l1_warm_up_steps=0). Fix by setting l1_warm_up_steps=1000. In v6, ghost grads are removed; warm-up is the primary lever. A dead feature ratio above 70% indicates a problem.\n\npython\nfrom sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig\n\n# WRONG: no warm-up, features die early\ncfg = LanguageModelSAERunnerConfig(\n sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),\n l1_warm_up_steps=0, # Bad!\n)\n\n# RIGHT: warm up the L1 penalty (v6 removed ghost grads; warm-up is the lever now)\ncfg = LanguageModelSAERunnerConfig(\n sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),\n l1_warm_up_steps=1000, # Gradually increase\n)\n\n\n### Poor Reconstruction (Low CE Recovery)\n\nCaused by too high sparsity penalty or insufficient capacity. Reduce l1_coefficient and/or increase d_sae. For example, increase d_sae to 16× d_model and reduce l1_coefficient to 5e-5. A CE loss score below 95% indicates poor reconstruction.\n\npython\n# Reduce sparsity penalty and/or add capacity (both on the SAE sub-config)\ncfg = LanguageModelSAERunnerConfig(\n sae=StandardTrainingSAEConfig(\n d_in=768,\n d_sae=768 * 16, # More capacity\n l1_coefficient=5e-5, # Lower = better reconstruction\n ),\n)\n\n\n### Features Not Interpretable\n\nCaused by insufficient sparsity. Increase l1_coefficient or switch to TopK SAE with lower k. For example, increase l1_coefficient to 1e-4 or use TopK SAE with k=50. An L0 value above 90 indicates too many active features.\n\npython\nfrom sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig, TopKTrainingSAEConfig\n\n# Increase sparsity (higher L1)\ncfg = LanguageModelSAERunnerConfig(\n sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),\n)\n# Or use a TopK SAE (k is set directly in v6, not via activation_fn_kwargs)\ncfg = LanguageModelSAERunnerConfig(\n sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50), # Exactly 50 active features\n)\n\n\n### Memory Errors During Training\n\nReduce train_batch_size_tokens, store_batch_size_prompts, and/or n_batches_in_buffer. For example, set train_batch_size_tokens to 2048, store_batch_size_prompts to 4, and n_batches_in_buffer to 8. A typical GPU can handle up to 1100 tokens per batch before memory issues arise.\n\npython\ncfg = LanguageModelSAERunnerConfig(\n sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),\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## Pre-trained SAE Releases\n\nTwo notable releases are gpt2-small-res-jb and gemma-2b-res. Features are indexed by SAE ID. For example, GPT-2 Small layer 8 feature 1234 corresponds to neuronpedia.org/gpt2-small/8-res-jb/1234. The gpt2-small-res-jb release contains features from 2023, while gemma-2b-res was released in 2024.\n\npython\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## TopK SAE Example\n\nFor TopK SAE, set k on the SAE sub-config directly:\n\npython\nfrom sae_lens import LanguageModelSAERunnerConfig, TopKTrainingSAEConfig\n\n# TopK SAE (exactly 50 features active) — `k` is set on the SAE sub-config in v6\ncfg = LanguageModelSAERunnerConfig(\n sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50),\n)\n\n\n## Additional Notes\n\n- In v6, SAE.from_pretrained() returns just the SAE object, not a tuple. Use SAE.from_pretrained_with_cfg_and_sparsity() if the cfg dict and feature sparsity are needed.\n- Training config uses nested sub-configs: SAE-specific options go in sae= (e.g., StandardTrainingSAEConfig), logging in logger= (LoggingConfig).\n- Legacy flat options removed in v6: architecture, expansion_factor, hook_layer, activation_fn, activation_fn_kwargs, use_ghost_grads, b_dec/decoder init options, log_to_wandb (now in logger=).\n- Layer is inferred from hook_name; hook_layer parameter removed.\n- Ghost grads removed in v6; use l1_warm_up_steps to prevent dead features.\n- For other SAE types in v6, swap sae= sub-config: GatedTrainingSAEConfig, TopKTrainingSAEConfig (set k directly), or JumpReLUTrainingSAEConfig (uses l0_coefficient).\n- Pre-trained SAEs available for GPT-2 Small (multiple residual streams) and Gemma 2B (residual streams); search HuggingFace tag saelens for more.\n- Neuronpedia feature URLs follow pattern: neuronpedia.org/{model}/{layer}-res-jb/{feature_id}.\n- The SAETrainingRunner alias still works as an alternative to LanguageModelSAETrainingRunner.\n- The ActivationsStore and HookedSAETransformer classes are available in the references/ module for advanced use." }