Add Mechanistic Interpretability category with 4 skills

New category 04-mechanistic-interpretability with comprehensive skills for
reverse-engineering neural network internals:

- TransformerLens: HookPoints, activation caching, circuit analysis (346 lines + 3 refs)
- SAELens: Sparse Autoencoder training for feature discovery (386 lines + 3 refs)
- pyvene: Stanford's causal intervention library with DAS (473 lines + 3 refs)
- nnsight: Remote interpretability via NDIF for 70B+ models (436 lines + 3 refs)

~6,500 lines of documentation across 16 files. Updates README to 74 total skills.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
zechenzhangAGI
2025-12-17 11:48:12 -08:00
parent d9a48bb892
commit 5b0d8bfaaa
17 changed files with 4701 additions and 9 deletions
@@ -0,0 +1,436 @@
---
name: nnsight-remote-interpretability
description: Provides guidance for interpreting and manipulating neural network internals using nnsight with optional NDIF remote execution. Use when needing to run interpretability experiments on massive models (70B+) without local GPU resources, or when working with any PyTorch architecture.
version: 1.0.0
author: Orchestra Research
license: MIT
tags: [nnsight, NDIF, Remote Execution, Mechanistic Interpretability, Model Internals]
dependencies: [nnsight>=0.5.0, torch>=2.0.0]
---
# nnsight: Transparent Access to Neural Network Internals
nnsight (/ɛn.saɪt/) enables researchers to interpret and manipulate the internals of any PyTorch model, with the unique capability of running the same code locally on small models or remotely on massive models (70B+) via NDIF.
**GitHub**: [ndif-team/nnsight](https://github.com/ndif-team/nnsight) (730+ stars)
**Paper**: [NNsight and NDIF: Democratizing Access to Foundation Model Internals](https://arxiv.org/abs/2407.14561) (ICLR 2025)
## Key Value Proposition
**Write once, run anywhere**: The same interpretability code works on GPT-2 locally or Llama-3.1-405B remotely. Just toggle `remote=True`.
```python
# Local execution (small model)
with model.trace("Hello world"):
hidden = model.transformer.h[5].output[0].save()
# Remote execution (massive model) - same code!
with model.trace("Hello world", remote=True):
hidden = model.model.layers[40].output[0].save()
```
## When to Use nnsight
**Use nnsight when you need to:**
- Run interpretability experiments on models too large for local GPUs (70B, 405B)
- Work with any PyTorch architecture (transformers, Mamba, custom models)
- Perform multi-token generation interventions
- Share activations between different prompts
- Access full model internals without reimplementation
**Consider alternatives when:**
- You want consistent API across models → Use **TransformerLens**
- You need declarative, shareable interventions → Use **pyvene**
- You're training SAEs → Use **SAELens**
- You only work with small models locally → **TransformerLens** may be simpler
## Installation
```bash
# Basic installation
pip install nnsight
# For vLLM support
pip install "nnsight[vllm]"
```
For remote NDIF execution, sign up at [login.ndif.us](https://login.ndif.us) for an API key.
## Core Concepts
### LanguageModel Wrapper
```python
from nnsight import LanguageModel
# Load model (uses HuggingFace under the hood)
model = LanguageModel("openai-community/gpt2", device_map="auto")
# For larger models
model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto")
```
### Tracing Context
The `trace` context manager enables deferred execution - operations are collected into a computation graph:
```python
from nnsight import LanguageModel
model = LanguageModel("gpt2", device_map="auto")
with model.trace("The Eiffel Tower is in") as tracer:
# Access any module's output
hidden_states = model.transformer.h[5].output[0].save()
# Access attention patterns
attn = model.transformer.h[5].attn.attn_dropout.input[0][0].save()
# Modify activations
model.transformer.h[8].output[0][:] = 0 # Zero out layer 8
# Get final output
logits = model.output.save()
# After context exits, access saved values
print(hidden_states.shape) # [batch, seq, hidden]
```
### Proxy Objects
Inside `trace`, module accesses return Proxy objects that record operations:
```python
with model.trace("Hello"):
# These are all Proxy objects - operations are deferred
h5_out = model.transformer.h[5].output[0] # Proxy
h5_mean = h5_out.mean(dim=-1) # Proxy
h5_saved = h5_mean.save() # Save for later access
```
## Workflow 1: Activation Analysis
### Step-by-Step
```python
from nnsight import LanguageModel
import torch
model = LanguageModel("gpt2", device_map="auto")
prompt = "The capital of France is"
with model.trace(prompt) as tracer:
# 1. Collect activations from multiple layers
layer_outputs = []
for i in range(12): # GPT-2 has 12 layers
layer_out = model.transformer.h[i].output[0].save()
layer_outputs.append(layer_out)
# 2. Get attention patterns
attn_patterns = []
for i in range(12):
# Access attention weights (after softmax)
attn = model.transformer.h[i].attn.attn_dropout.input[0][0].save()
attn_patterns.append(attn)
# 3. Get final logits
logits = model.output.save()
# 4. Analyze outside context
for i, layer_out in enumerate(layer_outputs):
print(f"Layer {i} output shape: {layer_out.shape}")
print(f"Layer {i} norm: {layer_out.norm().item():.3f}")
# 5. Find top predictions
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
for token, prob in zip(top_tokens.indices, top_tokens.values):
print(f"{model.tokenizer.decode(token)}: {prob.item():.3f}")
```
### Checklist
- [ ] Load model with LanguageModel wrapper
- [ ] Use trace context for operations
- [ ] Call `.save()` on values you need after context
- [ ] Access saved values outside context
- [ ] Use `.shape`, `.norm()`, etc. for analysis
## Workflow 2: Activation Patching
### Step-by-Step
```python
from nnsight import LanguageModel
import torch
model = LanguageModel("gpt2", device_map="auto")
clean_prompt = "The Eiffel Tower is in"
corrupted_prompt = "The Colosseum is in"
# 1. Get clean activations
with model.trace(clean_prompt) as tracer:
clean_hidden = model.transformer.h[8].output[0].save()
# 2. Patch clean into corrupted run
with model.trace(corrupted_prompt) as tracer:
# Replace layer 8 output with clean activations
model.transformer.h[8].output[0][:] = clean_hidden
patched_logits = model.output.save()
# 3. Compare predictions
paris_token = model.tokenizer.encode(" Paris")[0]
rome_token = model.tokenizer.encode(" Rome")[0]
patched_probs = torch.softmax(patched_logits[0, -1], dim=-1)
print(f"Paris prob: {patched_probs[paris_token].item():.3f}")
print(f"Rome prob: {patched_probs[rome_token].item():.3f}")
```
### Systematic Patching Sweep
```python
def patch_layer_position(layer, position, clean_cache, corrupted_prompt):
"""Patch single layer/position from clean to corrupted."""
with model.trace(corrupted_prompt) as tracer:
# Get current activation
current = model.transformer.h[layer].output[0]
# Patch only specific position
current[:, position, :] = clean_cache[layer][:, position, :]
logits = model.output.save()
return logits
# Sweep over all layers and positions
results = torch.zeros(12, seq_len)
for layer in range(12):
for pos in range(seq_len):
logits = patch_layer_position(layer, pos, clean_hidden, corrupted)
results[layer, pos] = compute_metric(logits)
```
## Workflow 3: Remote Execution with NDIF
Run the same experiments on massive models without local GPUs.
### Step-by-Step
```python
from nnsight import LanguageModel
# 1. Load large model (will run remotely)
model = LanguageModel("meta-llama/Llama-3.1-70B")
# 2. Same code, just add remote=True
with model.trace("The meaning of life is", remote=True) as tracer:
# Access internals of 70B model!
layer_40_out = model.model.layers[40].output[0].save()
logits = model.output.save()
# 3. Results returned from NDIF
print(f"Layer 40 shape: {layer_40_out.shape}")
# 4. Generation with interventions
with model.trace(remote=True) as tracer:
with tracer.invoke("What is 2+2?"):
# Intervene during generation
model.model.layers[20].output[0][:, -1, :] *= 1.5
output = model.generate(max_new_tokens=50)
```
### NDIF Setup
1. Sign up at [login.ndif.us](https://login.ndif.us)
2. Get API key
3. Set environment variable or pass to nnsight:
```python
import os
os.environ["NDIF_API_KEY"] = "your_key"
# Or configure directly
from nnsight import CONFIG
CONFIG.API_KEY = "your_key"
```
### Available Models on NDIF
- Llama-3.1-8B, 70B, 405B
- DeepSeek-R1 models
- Various open-weight models (check [ndif.us](https://ndif.us) for current list)
## Workflow 4: Cross-Prompt Activation Sharing
Share activations between different inputs in a single trace.
```python
from nnsight import LanguageModel
model = LanguageModel("gpt2", device_map="auto")
with model.trace() as tracer:
# First prompt
with tracer.invoke("The cat sat on the"):
cat_hidden = model.transformer.h[6].output[0].save()
# Second prompt - inject cat's activations
with tracer.invoke("The dog ran through the"):
# Replace with cat's activations at layer 6
model.transformer.h[6].output[0][:] = cat_hidden
dog_with_cat = model.output.save()
# The dog prompt now has cat's internal representations
```
## Workflow 5: Gradient-Based Analysis
Access gradients during backward pass.
```python
from nnsight import LanguageModel
import torch
model = LanguageModel("gpt2", device_map="auto")
with model.trace("The quick brown fox") as tracer:
# Save activations and enable gradient
hidden = model.transformer.h[5].output[0].save()
hidden.retain_grad()
logits = model.output
# Compute loss on specific token
target_token = model.tokenizer.encode(" jumps")[0]
loss = -logits[0, -1, target_token]
# Backward pass
loss.backward()
# Access gradients
grad = hidden.grad
print(f"Gradient shape: {grad.shape}")
print(f"Gradient norm: {grad.norm().item():.3f}")
```
**Note**: Gradient access not supported for vLLM or remote execution.
## Common Issues & Solutions
### Issue: Module path differs between models
```python
# GPT-2 structure
model.transformer.h[5].output[0]
# LLaMA structure
model.model.layers[5].output[0]
# Solution: Check model structure
print(model._model) # See actual module names
```
### Issue: Forgetting to save
```python
# WRONG: Value not accessible outside trace
with model.trace("Hello"):
hidden = model.transformer.h[5].output[0] # Not saved!
print(hidden) # Error or wrong value
# RIGHT: Call .save()
with model.trace("Hello"):
hidden = model.transformer.h[5].output[0].save()
print(hidden) # Works!
```
### Issue: Remote timeout
```python
# For long operations, increase timeout
with model.trace("prompt", remote=True, timeout=300) as tracer:
# Long operation...
```
### Issue: Memory with many saved activations
```python
# Only save what you need
with model.trace("prompt"):
# Don't save everything
for i in range(100):
model.transformer.h[i].output[0].save() # Memory heavy!
# Better: save specific layers
key_layers = [0, 5, 11]
for i in key_layers:
model.transformer.h[i].output[0].save()
```
### Issue: vLLM gradient limitation
```python
# vLLM doesn't support gradients
# Use standard execution for gradient analysis
model = LanguageModel("gpt2", device_map="auto") # Not vLLM
```
## Key API Reference
| Method/Property | Purpose |
|-----------------|---------|
| `model.trace(prompt, remote=False)` | Start tracing context |
| `proxy.save()` | Save value for access after trace |
| `proxy[:]` | Slice/index proxy (assignment patches) |
| `tracer.invoke(prompt)` | Add prompt within trace |
| `model.generate(...)` | Generate with interventions |
| `model.output` | Final model output logits |
| `model._model` | Underlying HuggingFace model |
## Comparison with Other Tools
| Feature | nnsight | TransformerLens | pyvene |
|---------|---------|-----------------|--------|
| Any architecture | Yes | Transformers only | Yes |
| Remote execution | Yes (NDIF) | No | No |
| Consistent API | No | Yes | Yes |
| Deferred execution | Yes | No | No |
| HuggingFace native | Yes | Reimplemented | Yes |
| Shareable configs | No | No | Yes |
## 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 LanguageModel, tracing, proxy objects |
| [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for local and remote interpretability |
## External Resources
### Tutorials
- [Getting Started](https://nnsight.net/start/)
- [Features Overview](https://nnsight.net/features/)
- [Remote Execution](https://nnsight.net/notebooks/features/remote_execution/)
- [Applied Tutorials](https://nnsight.net/applied_tutorials/)
### Official Documentation
- [Official Docs](https://nnsight.net/documentation/)
- [NDIF Info](https://ndif.us/)
- [Community Forum](https://discuss.ndif.us/)
### Papers
- [NNsight and NDIF Paper](https://arxiv.org/abs/2407.14561) - Fiotto-Kaufman et al. (ICLR 2025)
## Architecture Support
nnsight works with any PyTorch model:
- **Transformers**: GPT-2, LLaMA, Mistral, etc.
- **State Space Models**: Mamba
- **Vision Models**: ViT, CLIP
- **Custom architectures**: Any nn.Module
The key is knowing the module structure to access the right components.
@@ -0,0 +1,78 @@
# nnsight Reference Documentation
This directory contains comprehensive reference materials for nnsight.
## Contents
- [api.md](api.md) - Complete API reference for LanguageModel, tracing, and proxy objects
- [tutorials.md](tutorials.md) - Step-by-step tutorials for local and remote interpretability
## Quick Links
- **Official Documentation**: https://nnsight.net/
- **GitHub Repository**: https://github.com/ndif-team/nnsight
- **NDIF (Remote Execution)**: https://ndif.us/
- **Community Forum**: https://discuss.ndif.us/
- **Paper**: https://arxiv.org/abs/2407.14561 (ICLR 2025)
## Installation
```bash
# Basic installation
pip install nnsight
# For vLLM support
pip install "nnsight[vllm]"
```
## Basic Usage
```python
from nnsight import LanguageModel
# Load model
model = LanguageModel("openai-community/gpt2", device_map="auto")
# Trace and access internals
with model.trace("The Eiffel Tower is in") as tracer:
# Access layer output
hidden = model.transformer.h[5].output[0].save()
# Modify activations
model.transformer.h[8].output[0][:] *= 0.5
# Get final output
logits = model.output.save()
# Access saved values outside context
print(hidden.shape)
```
## Key Concepts
### Tracing
The `trace()` context enables deferred execution - operations are recorded and executed together.
### Proxy Objects
Inside trace, module accesses return Proxies. Call `.save()` to retrieve values after execution.
### Remote Execution (NDIF)
Run the same code on massive models (70B+) without local GPUs:
```python
# Same code, just add remote=True
with model.trace("Hello", remote=True):
hidden = model.model.layers[40].output[0].save()
```
## NDIF Setup
1. Sign up at https://login.ndif.us/
2. Get API key
3. Set environment variable: `export NDIF_API_KEY=your_key`
## Available Remote Models
- Llama-3.1-8B, 70B, 405B
- DeepSeek-R1 models
- More at https://ndif.us/
@@ -0,0 +1,344 @@
# nnsight API Reference
## LanguageModel
Main class for wrapping language models with intervention capabilities.
### Loading Models
```python
from nnsight import LanguageModel
# Basic loading
model = LanguageModel("openai-community/gpt2", device_map="auto")
# Larger models
model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto")
# With custom tokenizer settings
model = LanguageModel(
"gpt2",
device_map="auto",
torch_dtype=torch.float16,
)
```
### Model Attributes
```python
# Access underlying HuggingFace model
model._model
# Access tokenizer
model.tokenizer
# Model config
model._model.config
```
---
## Tracing Context
The `trace()` method creates a context for deferred execution.
### Basic Tracing
```python
with model.trace("Hello world") as tracer:
# Operations are recorded, not executed immediately
hidden = model.transformer.h[5].output[0].save()
logits = model.output.save()
# After context, operations execute and saved values are available
print(hidden.shape)
```
### Tracing Parameters
```python
with model.trace(
prompt, # Input text or tokens
remote=False, # Use NDIF remote execution
validate=True, # Validate tensor shapes
scan=True, # Scan for shape info
) as tracer:
...
```
### Remote Execution
```python
# Same code works remotely
with model.trace("Hello", remote=True) as tracer:
hidden = model.transformer.h[5].output[0].save()
```
---
## Proxy Objects
Inside tracing context, accessing modules returns Proxy objects.
### Accessing Values
```python
with model.trace("Hello") as tracer:
# These are Proxy objects
layer_output = model.transformer.h[5].output[0]
attention = model.transformer.h[5].attn.output
# Operations create new Proxies
mean = layer_output.mean(dim=-1)
normed = layer_output / layer_output.norm()
```
### Saving Values
```python
with model.trace("Hello") as tracer:
# Must call .save() to access after context
hidden = model.transformer.h[5].output[0].save()
# Now hidden contains actual tensor
print(hidden.shape)
```
### Modifying Values
```python
with model.trace("Hello") as tracer:
# In-place modification
model.transformer.h[5].output[0][:] = 0
# Replace with computed value
model.transformer.h[5].output[0][:] = some_tensor
# Arithmetic modification
model.transformer.h[5].output[0][:] *= 0.5
model.transformer.h[5].output[0][:] += steering_vector
```
### Proxy Operations
```python
with model.trace("Hello") as tracer:
h = model.transformer.h[5].output[0]
# Indexing
first_token = h[:, 0, :]
last_token = h[:, -1, :]
# PyTorch operations
mean = h.mean(dim=-1)
norm = h.norm()
transposed = h.transpose(1, 2)
# Save results
mean.save()
```
---
## Module Access Patterns
### GPT-2 Structure
```python
with model.trace("Hello") as tracer:
# Embeddings
embed = model.transformer.wte.output.save()
pos_embed = model.transformer.wpe.output.save()
# Layer outputs
layer_out = model.transformer.h[5].output[0].save()
# Attention
attn_out = model.transformer.h[5].attn.output.save()
# MLP
mlp_out = model.transformer.h[5].mlp.output.save()
# Final output
logits = model.output.save()
```
### LLaMA Structure
```python
with model.trace("Hello") as tracer:
# Embeddings
embed = model.model.embed_tokens.output.save()
# Layer outputs
layer_out = model.model.layers[10].output[0].save()
# Attention
attn_out = model.model.layers[10].self_attn.output.save()
# MLP
mlp_out = model.model.layers[10].mlp.output.save()
# Final output
logits = model.output.save()
```
### Finding Module Names
```python
# Print model structure
print(model._model)
# Or iterate
for name, module in model._model.named_modules():
print(name)
```
---
## Multiple Prompts (invoke)
Process multiple prompts in a single trace.
### Basic Usage
```python
with model.trace() as tracer:
with tracer.invoke("First prompt"):
hidden1 = model.transformer.h[5].output[0].save()
with tracer.invoke("Second prompt"):
hidden2 = model.transformer.h[5].output[0].save()
```
### Cross-Prompt Intervention
```python
with model.trace() as tracer:
# Get activations from first prompt
with tracer.invoke("The cat sat on the"):
cat_hidden = model.transformer.h[6].output[0].save()
# Inject into second prompt
with tracer.invoke("The dog ran through the"):
model.transformer.h[6].output[0][:] = cat_hidden
output = model.output.save()
```
---
## Generation
Generate text with interventions.
### Basic Generation
```python
with model.trace() as tracer:
with tracer.invoke("Once upon a time"):
# Intervention during generation
model.transformer.h[5].output[0][:] *= 1.2
output = model.generate(max_new_tokens=50)
print(model.tokenizer.decode(output[0]))
```
---
## Gradients
Access gradients for analysis (not supported with remote/vLLM).
```python
with model.trace("The quick brown fox") as tracer:
hidden = model.transformer.h[5].output[0].save()
hidden.retain_grad()
logits = model.output
target_token = model.tokenizer.encode(" jumps")[0]
loss = -logits[0, -1, target_token]
loss.backward()
# Access gradient
grad = hidden.grad
```
---
## NDIF Remote Execution
### Setup
```python
import os
os.environ["NDIF_API_KEY"] = "your_key"
# Or configure directly
from nnsight import CONFIG
CONFIG.set_default_api_key("your_key")
```
### Using Remote
```python
model = LanguageModel("meta-llama/Llama-3.1-70B")
with model.trace("Hello", remote=True) as tracer:
hidden = model.model.layers[40].output[0].save()
logits = model.output.save()
# Results returned from NDIF
print(hidden.shape)
```
### Sessions (Batching Requests)
```python
with model.session(remote=True) as session:
with model.trace("First prompt"):
h1 = model.model.layers[20].output[0].save()
with model.trace("Second prompt"):
h2 = model.model.layers[20].output[0].save()
# Both run in single NDIF request
```
---
## Utility Methods
### Early Stopping
```python
with model.trace("Hello") as tracer:
hidden = model.transformer.h[5].output[0].save()
tracer.stop() # Don't run remaining layers
```
### Validation
```python
# Validate shapes before execution
with model.trace("Hello", validate=True) as tracer:
hidden = model.transformer.h[5].output[0].save()
```
### Module Access Result
```python
with model.trace("Hello") as tracer:
# Access result of a method call
result = tracer.result
```
---
## Common Module Paths
| Model | Embeddings | Layers | Attention | MLP |
|-------|------------|--------|-----------|-----|
| GPT-2 | `transformer.wte` | `transformer.h[i]` | `transformer.h[i].attn` | `transformer.h[i].mlp` |
| LLaMA | `model.embed_tokens` | `model.layers[i]` | `model.layers[i].self_attn` | `model.layers[i].mlp` |
| Mistral | `model.embed_tokens` | `model.layers[i]` | `model.layers[i].self_attn` | `model.layers[i].mlp` |
@@ -0,0 +1,300 @@
# nnsight Tutorials
## Tutorial 1: Basic Activation Analysis
### Goal
Load a model, access internal activations, and analyze them.
### Step-by-Step
```python
from nnsight import LanguageModel
import torch
# 1. Load model
model = LanguageModel("openai-community/gpt2", device_map="auto")
# 2. Trace and collect activations
prompt = "The capital of France is"
with model.trace(prompt) as tracer:
# Collect from multiple layers
activations = {}
for i in range(12): # GPT-2 has 12 layers
activations[i] = model.transformer.h[i].output[0].save()
# Get final logits
logits = model.output.save()
# 3. Analyze (outside context)
print("Layer-wise activation norms:")
for layer, act in activations.items():
print(f" Layer {layer}: {act.norm().item():.2f}")
# 4. Check predictions
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
print("\nTop predictions:")
for token_id, prob in zip(top_tokens.indices, top_tokens.values):
token_str = model.tokenizer.decode(token_id)
print(f" {token_str!r}: {prob.item():.3f}")
```
---
## Tutorial 2: Activation Patching
### Goal
Patch activations from one prompt into another to test causal relationships.
### Step-by-Step
```python
from nnsight import LanguageModel
import torch
model = LanguageModel("gpt2", device_map="auto")
clean_prompt = "The Eiffel Tower is in the city of"
corrupted_prompt = "The Colosseum is in the city of"
# 1. Get clean activations
with model.trace(clean_prompt) as tracer:
clean_hidden = model.transformer.h[8].output[0].save()
clean_logits = model.output.save()
# 2. Define metric
paris_token = model.tokenizer.encode(" Paris")[0]
rome_token = model.tokenizer.encode(" Rome")[0]
def logit_diff(logits):
return (logits[0, -1, paris_token] - logits[0, -1, rome_token]).item()
print(f"Clean logit diff: {logit_diff(clean_logits):.3f}")
# 3. Patch clean into corrupted
with model.trace(corrupted_prompt) as tracer:
# Replace layer 8 output with clean activations
model.transformer.h[8].output[0][:] = clean_hidden
patched_logits = model.output.save()
print(f"Patched logit diff: {logit_diff(patched_logits):.3f}")
# 4. Systematic patching sweep
results = torch.zeros(12) # 12 layers
for layer in range(12):
# Get clean activation for this layer
with model.trace(clean_prompt) as tracer:
clean_act = model.transformer.h[layer].output[0].save()
# Patch into corrupted
with model.trace(corrupted_prompt) as tracer:
model.transformer.h[layer].output[0][:] = clean_act
logits = model.output.save()
results[layer] = logit_diff(logits)
print(f"Layer {layer}: {results[layer]:.3f}")
print(f"\nMost important layer: {results.argmax().item()}")
```
---
## Tutorial 3: Cross-Prompt Activation Sharing
### Goal
Transfer activations between different prompts in a single trace.
### Step-by-Step
```python
from nnsight import LanguageModel
model = LanguageModel("gpt2", device_map="auto")
with model.trace() as tracer:
# First prompt - get "cat" representations
with tracer.invoke("The cat sat on the mat"):
cat_hidden = model.transformer.h[6].output[0].save()
# Second prompt - inject "cat" into "dog"
with tracer.invoke("The dog ran through the park"):
# Replace with cat's activations
model.transformer.h[6].output[0][:] = cat_hidden
modified_logits = model.output.save()
# The dog prompt now has cat's internal representations
print(f"Modified logits shape: {modified_logits.shape}")
```
---
## Tutorial 4: Remote Execution with NDIF
### Goal
Run the same interpretability code on massive models (70B+).
### Step-by-Step
```python
from nnsight import LanguageModel
import os
# 1. Setup API key
os.environ["NDIF_API_KEY"] = "your_key_here"
# 2. Load large model (runs remotely)
model = LanguageModel("meta-llama/Llama-3.1-70B")
# 3. Same code, just remote=True
prompt = "The meaning of life is"
with model.trace(prompt, remote=True) as tracer:
# Access layer 40 of 70B model!
hidden = model.model.layers[40].output[0].save()
logits = model.output.save()
# 4. Results returned from NDIF
print(f"Hidden shape: {hidden.shape}")
print(f"Logits shape: {logits.shape}")
# 5. Check predictions
import torch
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
print("\nTop predictions from Llama-70B:")
for token_id, prob in zip(top_tokens.indices, top_tokens.values):
print(f" {model.tokenizer.decode(token_id)!r}: {prob.item():.3f}")
```
### Batching with Sessions
```python
# Run multiple experiments in one NDIF request
with model.session(remote=True) as session:
with model.trace("What is 2+2?"):
math_hidden = model.model.layers[30].output[0].save()
with model.trace("The capital of France is"):
fact_hidden = model.model.layers[30].output[0].save()
# Compare representations
similarity = torch.cosine_similarity(
math_hidden.mean(dim=1),
fact_hidden.mean(dim=1),
dim=-1
)
print(f"Similarity: {similarity.item():.3f}")
```
---
## Tutorial 5: Steering with Activation Addition
### Goal
Add a steering vector to change model behavior.
### Step-by-Step
```python
from nnsight import LanguageModel
import torch
model = LanguageModel("gpt2", device_map="auto")
# 1. Get contrasting activations
with model.trace("I love this movie, it's wonderful") as tracer:
positive_hidden = model.transformer.h[6].output[0].save()
with model.trace("I hate this movie, it's terrible") as tracer:
negative_hidden = model.transformer.h[6].output[0].save()
# 2. Compute steering direction
steering_vector = positive_hidden.mean(dim=1) - negative_hidden.mean(dim=1)
# 3. Generate without steering
test_prompt = "This restaurant is"
with model.trace(test_prompt) as tracer:
normal_logits = model.output.save()
# 4. Generate with steering
with model.trace(test_prompt) as tracer:
# Add steering at layer 6
model.transformer.h[6].output[0][:] += 3.0 * steering_vector
steered_logits = model.output.save()
# 5. Compare predictions
def top_prediction(logits):
token = logits[0, -1].argmax()
return model.tokenizer.decode(token)
print(f"Normal: {top_prediction(normal_logits)}")
print(f"Steered (positive): {top_prediction(steered_logits)}")
```
---
## Tutorial 6: Logit Lens
### Goal
See what the model "believes" at each layer.
### Step-by-Step
```python
from nnsight import LanguageModel
import torch
model = LanguageModel("gpt2", device_map="auto")
prompt = "The quick brown fox jumps over the lazy"
with model.trace(prompt) as tracer:
# Collect residual stream at each layer
residuals = []
for i in range(12):
resid = model.transformer.h[i].output[0].save()
residuals.append(resid)
# Access model's unembedding and final layernorm
W_U = model._model.lm_head.weight.T # [d_model, vocab]
ln_f = model._model.transformer.ln_f
print("Layer-by-layer predictions for final token:")
for i, resid in enumerate(residuals):
# Apply final layernorm
normed = ln_f(resid)
# Project to vocabulary
layer_logits = normed @ W_U
# Get prediction
probs = torch.softmax(layer_logits[0, -1], dim=-1)
top_token = probs.argmax()
top_prob = probs[top_token].item()
print(f"Layer {i}: {model.tokenizer.decode(top_token)!r} ({top_prob:.3f})")
```
---
## External Resources
### Official Resources
- [Getting Started](https://nnsight.net/start/)
- [Features Overview](https://nnsight.net/features/)
- [Documentation](https://nnsight.net/documentation/)
- [Tutorials](https://nnsight.net/tutorials/)
### NDIF Resources
- [NDIF Homepage](https://ndif.us/)
- [Available Models](https://ndif.us/models)
- [API Key Signup](https://login.ndif.us/)
### Paper
- [NNsight and NDIF](https://arxiv.org/abs/2407.14561) - ICLR 2025
### Community
- [Discussion Forum](https://discuss.ndif.us/)
- [GitHub Issues](https://github.com/ndif-team/nnsight/issues)
@@ -0,0 +1,473 @@
---
name: pyvene-interventions
description: Provides guidance for performing causal interventions on PyTorch models using pyvene's declarative intervention framework. Use when conducting causal tracing, activation patching, interchange intervention training, or testing causal hypotheses about model behavior.
version: 1.0.0
author: Orchestra Research
license: MIT
tags: [Causal Intervention, pyvene, Activation Patching, Causal Tracing, Interpretability]
dependencies: [pyvene>=0.1.8, torch>=2.0.0, transformers>=4.30.0]
---
# pyvene: Causal Interventions for Neural Networks
pyvene is Stanford NLP's library for performing causal interventions on PyTorch models. It provides a declarative, dict-based framework for activation patching, causal tracing, and interchange intervention training - making intervention experiments reproducible and shareable.
**GitHub**: [stanfordnlp/pyvene](https://github.com/stanfordnlp/pyvene) (840+ stars)
**Paper**: [pyvene: A Library for Understanding and Improving PyTorch Models via Interventions](https://aclanthology.org/2024.naacl-demo.16) (NAACL 2024)
## When to Use pyvene
**Use pyvene when you need to:**
- Perform causal tracing (ROME-style localization)
- Run activation patching experiments
- Conduct interchange intervention training (IIT)
- Test causal hypotheses about model components
- Share/reproduce intervention experiments via HuggingFace
- Work with any PyTorch architecture (not just transformers)
**Consider alternatives when:**
- You need exploratory activation analysis → Use **TransformerLens**
- You want to train/analyze SAEs → Use **SAELens**
- You need remote execution on massive models → Use **nnsight**
- You want lower-level control → Use **nnsight**
## Installation
```bash
pip install pyvene
```
Standard import:
```python
import pyvene as pv
```
## Core Concepts
### IntervenableModel
The main class that wraps any PyTorch model with intervention capabilities:
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load base model
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# Define intervention configuration
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=8,
component="block_output",
intervention_type=pv.VanillaIntervention,
)
]
)
# Create intervenable model
intervenable = pv.IntervenableModel(config, model)
```
### Intervention Types
| Type | Description | Use Case |
|------|-------------|----------|
| `VanillaIntervention` | Swap activations between runs | Activation patching |
| `AdditionIntervention` | Add activations to base run | Steering, ablation |
| `SubtractionIntervention` | Subtract activations | Ablation |
| `ZeroIntervention` | Zero out activations | Component knockout |
| `RotatedSpaceIntervention` | DAS trainable intervention | Causal discovery |
| `CollectIntervention` | Collect activations | Probing, analysis |
### Component Targets
```python
# Available components to intervene on
components = [
"block_input", # Input to transformer block
"block_output", # Output of transformer block
"mlp_input", # Input to MLP
"mlp_output", # Output of MLP
"mlp_activation", # MLP hidden activations
"attention_input", # Input to attention
"attention_output", # Output of attention
"attention_value_output", # Attention value vectors
"query_output", # Query vectors
"key_output", # Key vectors
"value_output", # Value vectors
"head_attention_value_output", # Per-head values
]
```
## Workflow 1: Causal Tracing (ROME-style)
Locate where factual associations are stored by corrupting inputs and restoring activations.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("gpt2-xl")
tokenizer = AutoTokenizer.from_pretrained("gpt2-xl")
# 1. Define clean and corrupted inputs
clean_prompt = "The Space Needle is in downtown"
corrupted_prompt = "The ##### ###### ## ## ########" # Noise
clean_tokens = tokenizer(clean_prompt, return_tensors="pt")
corrupted_tokens = tokenizer(corrupted_prompt, return_tensors="pt")
# 2. Get clean activations (source)
with torch.no_grad():
clean_outputs = model(**clean_tokens, output_hidden_states=True)
clean_states = clean_outputs.hidden_states
# 3. Define restoration intervention
def run_causal_trace(layer, position):
"""Restore clean activation at specific layer and position."""
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=layer,
component="block_output",
intervention_type=pv.VanillaIntervention,
unit="pos",
max_number_of_units=1,
)
]
)
intervenable = pv.IntervenableModel(config, model)
# Run with intervention
_, patched_outputs = intervenable(
base=corrupted_tokens,
sources=[clean_tokens],
unit_locations={"sources->base": ([[[position]]], [[[position]]])},
output_original_output=True,
)
# Return probability of correct token
probs = torch.softmax(patched_outputs.logits[0, -1], dim=-1)
seattle_token = tokenizer.encode(" Seattle")[0]
return probs[seattle_token].item()
# 4. Sweep over layers and positions
n_layers = model.config.n_layer
seq_len = clean_tokens["input_ids"].shape[1]
results = torch.zeros(n_layers, seq_len)
for layer in range(n_layers):
for pos in range(seq_len):
results[layer, pos] = run_causal_trace(layer, pos)
# 5. Visualize (layer x position heatmap)
# High values indicate causal importance
```
### Checklist
- [ ] Prepare clean prompt with target factual association
- [ ] Create corrupted version (noise or counterfactual)
- [ ] Define intervention config for each (layer, position)
- [ ] Run patching sweep
- [ ] Identify causal hotspots in heatmap
## Workflow 2: Activation Patching for Circuit Analysis
Test which components are necessary for a specific behavior.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# IOI task setup
clean_prompt = "When John and Mary went to the store, Mary gave a bottle to"
corrupted_prompt = "When John and Mary went to the store, John gave a bottle to"
clean_tokens = tokenizer(clean_prompt, return_tensors="pt")
corrupted_tokens = tokenizer(corrupted_prompt, return_tensors="pt")
john_token = tokenizer.encode(" John")[0]
mary_token = tokenizer.encode(" Mary")[0]
def logit_diff(logits):
"""IO - S logit difference."""
return logits[0, -1, john_token] - logits[0, -1, mary_token]
# Patch attention output at each layer
def patch_attention(layer):
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=layer,
component="attention_output",
intervention_type=pv.VanillaIntervention,
)
]
)
intervenable = pv.IntervenableModel(config, model)
_, patched_outputs = intervenable(
base=corrupted_tokens,
sources=[clean_tokens],
)
return logit_diff(patched_outputs.logits).item()
# Find which layers matter
results = []
for layer in range(model.config.n_layer):
diff = patch_attention(layer)
results.append(diff)
print(f"Layer {layer}: logit diff = {diff:.3f}")
```
## Workflow 3: Interchange Intervention Training (IIT)
Train interventions to discover causal structure.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM
import torch
model = AutoModelForCausalLM.from_pretrained("gpt2")
# 1. Define trainable intervention
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=6,
component="block_output",
intervention_type=pv.RotatedSpaceIntervention, # Trainable
low_rank_dimension=64, # Learn 64-dim subspace
)
]
)
intervenable = pv.IntervenableModel(config, model)
# 2. Set up training
optimizer = torch.optim.Adam(
intervenable.get_trainable_parameters(),
lr=1e-4
)
# 3. Training loop (simplified)
for base_input, source_input, target_output in dataloader:
optimizer.zero_grad()
_, outputs = intervenable(
base=base_input,
sources=[source_input],
)
loss = criterion(outputs.logits, target_output)
loss.backward()
optimizer.step()
# 4. Analyze learned intervention
# The rotation matrix reveals causal subspace
rotation = intervenable.interventions["layer.6.block_output"][0].rotate_layer
```
### DAS (Distributed Alignment Search)
```python
# Low-rank rotation finds interpretable subspaces
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=8,
component="block_output",
intervention_type=pv.LowRankRotatedSpaceIntervention,
low_rank_dimension=1, # Find 1D causal direction
)
]
)
```
## Workflow 4: Model Steering (Honest LLaMA)
Steer model behavior during generation.
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
# Load pre-trained steering intervention
intervenable = pv.IntervenableModel.load(
"zhengxuanzenwu/intervenable_honest_llama2_chat_7B",
model=model,
)
# Generate with steering
prompt = "Is the earth flat?"
inputs = tokenizer(prompt, return_tensors="pt")
# Intervention applied during generation
outputs = intervenable.generate(
inputs,
max_new_tokens=100,
do_sample=False,
)
print(tokenizer.decode(outputs[0]))
```
## Saving and Sharing Interventions
```python
# Save locally
intervenable.save("./my_intervention")
# Load from local
intervenable = pv.IntervenableModel.load(
"./my_intervention",
model=model,
)
# Share on HuggingFace
intervenable.save_intervention("username/my-intervention")
# Load from HuggingFace
intervenable = pv.IntervenableModel.load(
"username/my-intervention",
model=model,
)
```
## Common Issues & Solutions
### Issue: Wrong intervention location
```python
# WRONG: Incorrect component name
config = pv.RepresentationConfig(
component="mlp", # Not valid!
)
# RIGHT: Use exact component name
config = pv.RepresentationConfig(
component="mlp_output", # Valid
)
```
### Issue: Dimension mismatch
```python
# Ensure source and base have compatible shapes
# For position-specific interventions:
config = pv.RepresentationConfig(
unit="pos",
max_number_of_units=1, # Intervene on single position
)
# Specify locations explicitly
intervenable(
base=base_tokens,
sources=[source_tokens],
unit_locations={"sources->base": ([[[5]]], [[[5]]])}, # Position 5
)
```
### Issue: Memory with large models
```python
# Use gradient checkpointing
model.gradient_checkpointing_enable()
# Or intervene on fewer components
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=8, # Single layer instead of all
component="block_output",
)
]
)
```
### Issue: LoRA integration
```python
# pyvene v0.1.8+ supports LoRAs as interventions
config = pv.RepresentationConfig(
intervention_type=pv.LoRAIntervention,
low_rank_dimension=16,
)
```
## Key Classes Reference
| Class | Purpose |
|-------|---------|
| `IntervenableModel` | Main wrapper for interventions |
| `IntervenableConfig` | Configuration container |
| `RepresentationConfig` | Single intervention specification |
| `VanillaIntervention` | Activation swapping |
| `RotatedSpaceIntervention` | Trainable DAS intervention |
| `CollectIntervention` | Activation collection |
## Supported Models
pyvene works with any PyTorch model. Tested on:
- GPT-2 (all sizes)
- LLaMA / LLaMA-2
- Pythia
- Mistral / Mixtral
- OPT
- BLIP (vision-language)
- ESM (protein models)
- Mamba (state space)
## 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 IntervenableModel, intervention types, configurations |
| [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for causal tracing, activation patching, DAS |
## External Resources
### Tutorials
- [pyvene 101](https://stanfordnlp.github.io/pyvene/tutorials/pyvene_101.html)
- [Causal Tracing Tutorial](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/Causal_Tracing.html)
- [IOI Circuit Replication](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/IOI_Replication.html)
- [DAS Introduction](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/DAS_Main_Introduction.html)
### Papers
- [Locating and Editing Factual Associations in GPT](https://arxiv.org/abs/2202.05262) - Meng et al. (2022)
- [Inference-Time Intervention](https://arxiv.org/abs/2306.03341) - Li et al. (2023)
- [Interpretability in the Wild](https://arxiv.org/abs/2211.00593) - Wang et al. (2022)
### Official Documentation
- [Official Docs](https://stanfordnlp.github.io/pyvene/)
- [API Reference](https://stanfordnlp.github.io/pyvene/api/)
## Comparison with Other Tools
| Feature | pyvene | TransformerLens | nnsight |
|---------|--------|-----------------|---------|
| Declarative config | Yes | No | No |
| HuggingFace sharing | Yes | No | No |
| Trainable interventions | Yes | Limited | Yes |
| Any PyTorch model | Yes | Transformers only | Yes |
| Remote execution | No | No | Yes (NDIF) |
@@ -0,0 +1,73 @@
# pyvene Reference Documentation
This directory contains comprehensive reference materials for pyvene.
## Contents
- [api.md](api.md) - Complete API reference for IntervenableModel, intervention types, and configurations
- [tutorials.md](tutorials.md) - Step-by-step tutorials for causal tracing, activation patching, and trainable interventions
## Quick Links
- **Official Documentation**: https://stanfordnlp.github.io/pyvene/
- **GitHub Repository**: https://github.com/stanfordnlp/pyvene
- **Paper**: https://arxiv.org/abs/2403.07809 (NAACL 2024)
## Installation
```bash
pip install pyvene
```
## Basic Usage
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load model
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# Define intervention
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.VanillaIntervention,
)
]
)
# Create intervenable model
intervenable = pv.IntervenableModel(config, model)
# Run intervention (swap activations from source to base)
base_inputs = tokenizer("The cat sat on the", return_tensors="pt")
source_inputs = tokenizer("The dog ran through the", return_tensors="pt")
_, outputs = intervenable(
base=base_inputs,
sources=[source_inputs],
)
```
## Key Concepts
### Intervention Types
- **VanillaIntervention**: Swap activations between runs
- **AdditionIntervention**: Add source to base activations
- **ZeroIntervention**: Zero out activations (ablation)
- **CollectIntervention**: Collect activations without modifying
- **RotatedSpaceIntervention**: Trainable intervention for causal discovery
### Components
Target specific parts of the model:
- `block_input`, `block_output`
- `mlp_input`, `mlp_output`, `mlp_activation`
- `attention_input`, `attention_output`
- `query_output`, `key_output`, `value_output`
### HuggingFace Integration
Save and load interventions via HuggingFace Hub for reproducibility.
@@ -0,0 +1,383 @@
# pyvene API Reference
## IntervenableModel
The core class that wraps PyTorch models for intervention.
### Basic Usage
```python
import pyvene as pv
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("gpt2")
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.VanillaIntervention,
)
]
)
intervenable = pv.IntervenableModel(config, model)
```
### Forward Pass
```python
# Basic intervention
original_output, intervened_output = intervenable(
base=base_inputs,
sources=[source_inputs],
)
# With unit locations (position-specific)
_, outputs = intervenable(
base=base_inputs,
sources=[source_inputs],
unit_locations={"sources->base": ([[[5]]], [[[5]]])}, # Position 5
)
# Return original output too
original, intervened = intervenable(
base=base_inputs,
sources=[source_inputs],
output_original_output=True,
)
```
### Generation
```python
# Generate with interventions
outputs = intervenable.generate(
base_inputs,
sources=[source_inputs],
max_new_tokens=50,
do_sample=False,
)
```
### Saving and Loading
```python
# Save locally
intervenable.save("./my_intervention")
# Load
intervenable = pv.IntervenableModel.load("./my_intervention", model=model)
# Save to HuggingFace
intervenable.save_intervention("username/my-intervention")
# Load from HuggingFace
intervenable = pv.IntervenableModel.load(
"username/my-intervention",
model=model
)
```
### Getting Trainable Parameters
```python
# For trainable interventions
params = intervenable.get_trainable_parameters()
optimizer = torch.optim.Adam(params, lr=1e-4)
```
---
## IntervenableConfig
Configuration container for interventions.
### Basic Config
```python
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(...)
]
)
```
### Multiple Interventions
```python
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(layer=3, component="block_output", ...),
pv.RepresentationConfig(layer=5, component="mlp_output", ...),
pv.RepresentationConfig(layer=7, component="attention_output", ...),
]
)
```
---
## RepresentationConfig
Specifies a single intervention target.
### Parameters
| Parameter | Type | Description |
|-----------|------|-------------|
| `layer` | int | Layer index |
| `component` | str | Component to intervene on |
| `intervention_type` | type | Intervention class |
| `unit` | str | Intervention unit ("pos", "h", etc.) |
| `max_number_of_units` | int | Max units to intervene |
| `low_rank_dimension` | int | For trainable interventions |
| `subspace_partition` | list | Dimension ranges |
### Components
| Component | Description |
|-----------|-------------|
| `block_input` | Input to transformer block |
| `block_output` | Output of transformer block |
| `mlp_input` | Input to MLP |
| `mlp_output` | Output of MLP |
| `mlp_activation` | MLP hidden activations |
| `attention_input` | Input to attention |
| `attention_output` | Output of attention |
| `attention_value_output` | Attention values |
| `query_output` | Query vectors |
| `key_output` | Key vectors |
| `value_output` | Value vectors |
| `head_attention_value_output` | Per-head values |
### Example Configs
```python
# Position-specific intervention
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.VanillaIntervention,
unit="pos",
max_number_of_units=1,
)
# Trainable low-rank intervention
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.LowRankRotatedSpaceIntervention,
low_rank_dimension=64,
)
# Subspace intervention
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.VanillaIntervention,
subspace_partition=[[0, 256], [256, 512]], # First 512 dims split
)
```
---
## Intervention Types
### Basic Interventions
#### VanillaIntervention
Replaces base activations with source activations.
```python
pv.RepresentationConfig(
intervention_type=pv.VanillaIntervention,
...
)
```
#### AdditionIntervention
Adds source activations to base.
```python
pv.RepresentationConfig(
intervention_type=pv.AdditionIntervention,
...
)
```
#### SubtractionIntervention
Subtracts source from base.
```python
pv.RepresentationConfig(
intervention_type=pv.SubtractionIntervention,
...
)
```
#### ZeroIntervention
Sets activations to zero (ablation).
```python
pv.RepresentationConfig(
intervention_type=pv.ZeroIntervention,
...
)
```
#### CollectIntervention
Collects activations without modification.
```python
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.CollectIntervention,
)
]
)
intervenable = pv.IntervenableModel(config, model)
_, collected = intervenable(base=inputs)
# collected contains the activations
```
### Trainable Interventions
#### RotatedSpaceIntervention
Full-rank trainable rotation.
```python
pv.RepresentationConfig(
intervention_type=pv.RotatedSpaceIntervention,
...
)
```
#### LowRankRotatedSpaceIntervention
Low-rank trainable intervention (DAS).
```python
pv.RepresentationConfig(
intervention_type=pv.LowRankRotatedSpaceIntervention,
low_rank_dimension=64,
...
)
```
#### BoundlessRotatedSpaceIntervention
Boundless DAS variant.
```python
pv.RepresentationConfig(
intervention_type=pv.BoundlessRotatedSpaceIntervention,
...
)
```
#### SigmoidMaskIntervention
Learnable binary mask.
```python
pv.RepresentationConfig(
intervention_type=pv.SigmoidMaskIntervention,
...
)
```
---
## Unit Locations
Specify exactly where to intervene.
### Format
```python
unit_locations = {
"sources->base": (source_locations, base_locations)
}
```
### Examples
```python
# Single position
unit_locations = {"sources->base": ([[[5]]], [[[5]]])}
# Multiple positions
unit_locations = {"sources->base": ([[[3, 5, 7]]], [[[3, 5, 7]]])}
# Different source and base positions
unit_locations = {"sources->base": ([[[5]]], [[[10]]])}
```
---
## Supported Models
pyvene works with any PyTorch model. Officially tested:
| Family | Models |
|--------|--------|
| GPT-2 | gpt2, gpt2-medium, gpt2-large, gpt2-xl |
| LLaMA | llama-7b, llama-2-7b, llama-2-13b |
| Pythia | pythia-70m to pythia-12b |
| Mistral | mistral-7b, mixtral-8x7b |
| Gemma | gemma-2b, gemma-7b |
| Vision | BLIP, LLaVA |
| Other | OPT, Phi, Qwen, ESM, Mamba |
---
## Quick Reference: Common Patterns
### Activation Patching
```python
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=layer,
component="block_output",
intervention_type=pv.VanillaIntervention,
)
]
)
```
### Causal Tracing (ROME-style)
```python
config = pv.IntervenableConfig(
representations=[
# First corrupt with noise
pv.RepresentationConfig(
layer=0,
component="block_input",
intervention_type=pv.NoiseIntervention,
),
# Then restore at target layer
pv.RepresentationConfig(
layer=target_layer,
component="block_output",
intervention_type=pv.VanillaIntervention,
),
]
)
```
### DAS (Distributed Alignment Search)
```python
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=layer,
component="block_output",
intervention_type=pv.LowRankRotatedSpaceIntervention,
low_rank_dimension=1, # Find 1D causal direction
)
]
)
```
@@ -0,0 +1,376 @@
# pyvene Tutorials
## Tutorial 1: Basic Activation Patching
### Goal
Swap activations between two prompts to test causal relationships.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# 1. Load model
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# 2. Prepare inputs
base_prompt = "The Colosseum is in the city of"
source_prompt = "The Eiffel Tower is in the city of"
base_inputs = tokenizer(base_prompt, return_tensors="pt")
source_inputs = tokenizer(source_prompt, return_tensors="pt")
# 3. Define intervention (patch layer 8)
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=8,
component="block_output",
intervention_type=pv.VanillaIntervention,
)
]
)
intervenable = pv.IntervenableModel(config, model)
# 4. Run intervention
_, patched_outputs = intervenable(
base=base_inputs,
sources=[source_inputs],
)
# 5. Check predictions
patched_logits = patched_outputs.logits
probs = torch.softmax(patched_logits[0, -1], dim=-1)
rome_token = tokenizer.encode(" Rome")[0]
paris_token = tokenizer.encode(" Paris")[0]
print(f"P(Rome): {probs[rome_token].item():.4f}")
print(f"P(Paris): {probs[paris_token].item():.4f}")
```
---
## Tutorial 2: Causal Tracing (ROME-style)
### Goal
Locate where factual associations are stored by corrupting inputs and restoring activations.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("gpt2-xl")
tokenizer = AutoTokenizer.from_pretrained("gpt2-xl")
# 1. Define prompts
clean_prompt = "The Space Needle is in downtown"
# We'll corrupt by adding noise to embeddings
clean_inputs = tokenizer(clean_prompt, return_tensors="pt")
seattle_token = tokenizer.encode(" Seattle")[0]
# 2. Get clean baseline
with torch.no_grad():
clean_outputs = model(**clean_inputs)
clean_prob = torch.softmax(clean_outputs.logits[0, -1], dim=-1)[seattle_token].item()
print(f"Clean P(Seattle): {clean_prob:.4f}")
# 3. Sweep over layers - corrupt input, restore at each layer
results = []
for restore_layer in range(model.config.n_layer):
# Config: add noise at input, restore at target layer
config = pv.IntervenableConfig(
representations=[
# Noise intervention at embedding
pv.RepresentationConfig(
layer=0,
component="block_input",
intervention_type=pv.NoiseIntervention,
),
# Restore clean at target layer
pv.RepresentationConfig(
layer=restore_layer,
component="block_output",
intervention_type=pv.VanillaIntervention,
),
]
)
intervenable = pv.IntervenableModel(config, model)
# Source is clean (for restoration), base gets noise
_, outputs = intervenable(
base=clean_inputs,
sources=[clean_inputs], # Restore from clean
)
prob = torch.softmax(outputs.logits[0, -1], dim=-1)[seattle_token].item()
results.append(prob)
print(f"Restore at layer {restore_layer}: P(Seattle) = {prob:.4f}")
# 4. Find critical layers (where restoration helps most)
import numpy as np
results = np.array(results)
critical_layers = np.argsort(results)[-5:]
print(f"\nMost critical layers: {critical_layers}")
```
---
## Tutorial 3: Trainable Interventions (DAS)
### Goal
Learn a low-rank intervention that achieves a target counterfactual behavior.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# 1. Define trainable intervention
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=6,
component="block_output",
intervention_type=pv.LowRankRotatedSpaceIntervention,
low_rank_dimension=64, # Learn 64-dim subspace
)
]
)
intervenable = pv.IntervenableModel(config, model)
# 2. Setup optimizer
optimizer = torch.optim.Adam(
intervenable.get_trainable_parameters(),
lr=1e-3
)
# 3. Training data (simplified example)
# Goal: Make model predict "Paris" instead of "Rome"
base_prompt = "The capital of Italy is"
target_token = tokenizer.encode(" Paris")[0]
base_inputs = tokenizer(base_prompt, return_tensors="pt")
# 4. Training loop
for step in range(100):
optimizer.zero_grad()
_, outputs = intervenable(
base=base_inputs,
sources=[base_inputs], # Self-intervention
)
# Loss: maximize probability of target token
logits = outputs.logits[0, -1]
loss = -torch.log_softmax(logits, dim=-1)[target_token]
loss.backward()
optimizer.step()
if step % 20 == 0:
prob = torch.softmax(logits.detach(), dim=-1)[target_token].item()
print(f"Step {step}: loss={loss.item():.4f}, P(Paris)={prob:.4f}")
# 5. Analyze learned rotation
rotation = intervenable.interventions["layer.6.comp.block_output.unit.pos.nunit.1#0"][0]
print(f"Learned rotation shape: {rotation.rotate_layer.weight.shape}")
```
---
## Tutorial 4: Position-Specific Intervention
### Goal
Intervene at specific token positions only.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# 1. Setup
base_prompt = "John and Mary went to the store"
source_prompt = "Alice and Bob went to the store"
base_inputs = tokenizer(base_prompt, return_tensors="pt")
source_inputs = tokenizer(source_prompt, return_tensors="pt")
# 2. Position-specific config
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.VanillaIntervention,
unit="pos",
max_number_of_units=1, # Single position
)
]
)
intervenable = pv.IntervenableModel(config, model)
# 3. Intervene at position 0 only (first name)
_, outputs = intervenable(
base=base_inputs,
sources=[source_inputs],
unit_locations={"sources->base": ([[[0]]], [[[0]]])},
)
# 4. Intervene at multiple positions
_, outputs = intervenable(
base=base_inputs,
sources=[source_inputs],
unit_locations={"sources->base": ([[[0, 2]]], [[[0, 2]]])},
)
```
---
## Tutorial 5: Collecting Activations
### Goal
Extract activations without modifying them.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# 1. Config with CollectIntervention
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=5,
component="block_output",
intervention_type=pv.CollectIntervention,
),
pv.RepresentationConfig(
layer=10,
component="attention_output",
intervention_type=pv.CollectIntervention,
),
]
)
intervenable = pv.IntervenableModel(config, model)
# 2. Run and collect
inputs = tokenizer("Hello world", return_tensors="pt")
_, collected = intervenable(base=inputs)
# 3. Access collected activations
layer5_output = collected[0]
layer10_attn = collected[1]
print(f"Layer 5 block output shape: {layer5_output.shape}")
print(f"Layer 10 attention output shape: {layer10_attn.shape}")
```
---
## Tutorial 6: Generation with Interventions
### Goal
Apply interventions during text generation.
### Step-by-Step
```python
import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
# 1. Get steering direction (happy vs sad)
happy_inputs = tokenizer("I am very happy and", return_tensors="pt")
sad_inputs = tokenizer("I am very sad and", return_tensors="pt")
# Collect activations
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=6,
component="mlp_output",
intervention_type=pv.CollectIntervention,
)
]
)
collector = pv.IntervenableModel(config, model)
_, happy_acts = collector(base=happy_inputs)
_, sad_acts = collector(base=sad_inputs)
steering_direction = happy_acts[0].mean(dim=1) - sad_acts[0].mean(dim=1)
# 2. Config for steering during generation
config = pv.IntervenableConfig(
representations=[
pv.RepresentationConfig(
layer=6,
component="mlp_output",
intervention_type=pv.AdditionIntervention,
)
]
)
intervenable = pv.IntervenableModel(config, model)
# 3. Generate with steering
prompt = "Today I feel"
inputs = tokenizer(prompt, return_tensors="pt")
# Create source with steering direction
# (This is simplified - actual implementation varies)
output = intervenable.generate(
inputs,
max_new_tokens=20,
do_sample=True,
temperature=0.7,
)
print(tokenizer.decode(output[0]))
```
---
## External Resources
### Official Tutorials
- [pyvene 101](https://stanfordnlp.github.io/pyvene/tutorials/pyvene_101.html)
- [Causal Tracing](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/Causal_Tracing.html)
- [DAS Introduction](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/DAS_Main_Introduction.html)
- [IOI Replication](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/IOI_Replication.html)
### Papers
- [pyvene Paper](https://arxiv.org/abs/2403.07809) - NAACL 2024
- [ROME](https://arxiv.org/abs/2202.05262) - Meng et al. (2022)
- [Inference-Time Intervention](https://arxiv.org/abs/2306.03341) - Li et al. (2023)
@@ -0,0 +1,386 @@
---
name: sparse-autoencoder-training
description: Provides guidance for training and analyzing Sparse Autoencoders (SAEs) using SAELens to decompose neural network activations into interpretable features. Use when discovering interpretable features, analyzing superposition, or studying monosemantic representations in language models.
version: 1.0.0
author: Orchestra Research
license: MIT
tags: [Sparse Autoencoders, SAE, Mechanistic Interpretability, Feature Discovery, Superposition]
dependencies: [sae-lens>=6.0.0, transformer-lens>=2.0.0, torch>=2.0.0]
---
# 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
```bash
pip 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
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
# 1. Load model and pre-trained SAE
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, cfg_dict, sparsity = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
# 2. Get model activations
tokens = 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 features
sae_features = sae.encode(activations) # [batch, pos, d_sae]
print(f"Active features: {(sae_features > 0).sum()}")
# 4. Find top features for each position
for 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 activations
reconstructed = 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
```python
from sae_lens import SAE, LanguageModelSAERunnerConfig, SAETrainingRunner
# 1. Configure training
cfg = LanguageModelSAERunnerConfig(
# Model
model_name="gpt2-small",
hook_name="blocks.8.hook_resid_pre",
hook_layer=8,
d_in=768, # Model dimension
# SAE architecture
architecture="standard", # or "gated", "topk"
d_sae=768 * 8, # Expansion factor of 8
activation_fn="relu",
# Training
lr=4e-4,
l1_coefficient=8e-5, # Sparsity penalty
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
log_to_wandb=True,
wandb_project="sae-training",
# Checkpointing
checkpoint_path="checkpoints",
n_checkpoints=5,
)
# 2. Train
trainer = SAETrainingRunner(cfg)
sae = trainer.run()
# 3. Evaluate
print(f"L0 (avg active features): {trainer.metrics['l0']}")
print(f"CE Loss Recovered: {trainer.metrics['ce_loss_score']}")
```
### 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
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
import torch
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, _, _ = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
# Find what activates a specific feature
feature_idx = 1234
test_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
```python
def 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 position
features = sae.encode(cache["resid_pre", 8])[0, -1] # [d_sae]
# Get logit attribution per feature
# Feature contribution = feature_activation × decoder_weight × unembedding
W_dec = sae.W_dec # [d_sae, d_model]
W_U = model.W_U # [d_model, vocab]
# Contribution to "Paris" logit
paris_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
### Issue: High dead feature ratio
```python
# WRONG: No warm-up, features die early
cfg = LanguageModelSAERunnerConfig(
l1_coefficient=1e-4,
l1_warm_up_steps=0, # Bad!
)
# RIGHT: Warm-up L1 penalty
cfg = LanguageModelSAERunnerConfig(
l1_coefficient=8e-5,
l1_warm_up_steps=1000, # Gradually increase
use_ghost_grads=True, # Revive dead features
)
```
### Issue: Poor reconstruction (low CE recovery)
```python
# Reduce sparsity penalty
cfg = LanguageModelSAERunnerConfig(
l1_coefficient=5e-5, # Lower = better reconstruction
d_sae=768 * 16, # More capacity
)
```
### Issue: Features not interpretable
```python
# Increase sparsity (higher L1)
cfg = LanguageModelSAERunnerConfig(
l1_coefficient=1e-4, # Higher = sparser, more interpretable
)
# Or use TopK architecture
cfg = LanguageModelSAERunnerConfig(
architecture="topk",
activation_fn_kwargs={"k": 50}, # Exactly 50 active features
)
```
### Issue: Memory errors during training
```python
cfg = LanguageModelSAERunnerConfig(
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` | Training configuration |
| `SAETrainingRunner` | Training loop manager |
| `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 |
```python
# TopK SAE (exactly 50 features active)
cfg = LanguageModelSAERunnerConfig(
architecture="topk",
activation_fn="topk",
activation_fn_kwargs={"k": 50},
)
```
@@ -0,0 +1,70 @@
# SAELens Reference Documentation
This directory contains comprehensive reference materials for SAELens.
## Contents
- [api.md](api.md) - Complete API reference for SAE, TrainingSAE, and configuration classes
- [tutorials.md](tutorials.md) - Step-by-step tutorials for training and analyzing SAEs
- [papers.md](papers.md) - Key research papers on sparse autoencoders
## Quick Links
- **GitHub Repository**: https://github.com/jbloomAus/SAELens
- **Neuronpedia**: https://neuronpedia.org (browse pre-trained SAE features)
- **HuggingFace SAEs**: Search for tag `saelens`
## Installation
```bash
pip install sae-lens
```
Requirements: Python 3.10+, transformer-lens>=2.0.0
## Basic Usage
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
# Load model and SAE
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, cfg_dict, sparsity = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
# Encode activations to sparse features
tokens = model.to_tokens("Hello world")
_, cache = model.run_with_cache(tokens)
activations = cache["resid_pre", 8]
features = sae.encode(activations) # Sparse feature activations
reconstructed = sae.decode(features) # Reconstructed activations
```
## Key Concepts
### Sparse Autoencoders
SAEs decompose dense neural activations into sparse, interpretable features:
- **Encoder**: Maps d_model → d_sae (typically 4-16x expansion)
- **ReLU/TopK**: Enforces sparsity
- **Decoder**: Reconstructs original activations
### Training Loss
`Loss = MSE(original, reconstructed) + L1_coefficient × L1(features)`
### Key Metrics
- **L0**: Average number of active features (target: 50-200)
- **CE Loss Score**: Cross-entropy recovered vs original model (target: 80-95%)
- **Dead Features**: Features that never activate (target: <5%)
## Available Pre-trained SAEs
| Release | Model | Description |
|---------|-------|-------------|
| `gpt2-small-res-jb` | GPT-2 Small | Residual stream SAEs |
| `gemma-2b-res` | Gemma 2B | Residual stream SAEs |
| Various | Search HuggingFace | Community-trained SAEs |
@@ -0,0 +1,333 @@
# SAELens API Reference
## SAE Class
The core class representing a Sparse Autoencoder.
### Loading Pre-trained SAEs
```python
from sae_lens import SAE
# From official releases
sae, cfg_dict, sparsity = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
# From HuggingFace
sae, cfg_dict, sparsity = SAE.from_pretrained(
release="username/repo-name",
sae_id="path/to/sae",
device="cuda"
)
# From local disk
sae = SAE.load_from_disk("/path/to/sae", device="cuda")
```
### SAE Attributes
| Attribute | Shape | Description |
|-----------|-------|-------------|
| `W_enc` | [d_in, d_sae] | Encoder weights |
| `W_dec` | [d_sae, d_in] | Decoder weights |
| `b_enc` | [d_sae] | Encoder bias |
| `b_dec` | [d_in] | Decoder bias |
| `cfg` | SAEConfig | Configuration object |
### Core Methods
#### encode()
```python
# Encode activations to sparse features
features = sae.encode(activations)
# Input: [batch, pos, d_in]
# Output: [batch, pos, d_sae]
```
#### decode()
```python
# Reconstruct activations from features
reconstructed = sae.decode(features)
# Input: [batch, pos, d_sae]
# Output: [batch, pos, d_in]
```
#### forward()
```python
# Full forward pass (encode + decode)
reconstructed = sae(activations)
# Returns reconstructed activations
```
#### save_model()
```python
sae.save_model("/path/to/save")
```
---
## SAEConfig
Configuration class for SAE architecture and training context.
### Key Parameters
| Parameter | Type | Description |
|-----------|------|-------------|
| `d_in` | int | Input dimension (model's d_model) |
| `d_sae` | int | SAE hidden dimension |
| `architecture` | str | "standard", "gated", "jumprelu", "topk" |
| `activation_fn_str` | str | Activation function name |
| `model_name` | str | Source model name |
| `hook_name` | str | Hook point in model |
| `normalize_activations` | str | Normalization method |
| `dtype` | str | Data type |
| `device` | str | Device |
### Accessing Config
```python
print(sae.cfg.d_in) # 768 for GPT-2 small
print(sae.cfg.d_sae) # e.g., 24576 (32x expansion)
print(sae.cfg.hook_name) # e.g., "blocks.8.hook_resid_pre"
```
---
## LanguageModelSAERunnerConfig
Comprehensive configuration for training SAEs.
### Example Configuration
```python
from sae_lens import LanguageModelSAERunnerConfig
cfg = LanguageModelSAERunnerConfig(
# Model and hook
model_name="gpt2-small",
hook_name="blocks.8.hook_resid_pre",
hook_layer=8,
d_in=768,
# SAE architecture
architecture="standard", # "standard", "gated", "jumprelu", "topk"
d_sae=768 * 8, # Expansion factor
activation_fn="relu",
# Training hyperparameters
lr=4e-4,
l1_coefficient=8e-5,
lp_norm=1.0,
lr_scheduler_name="constant",
lr_warm_up_steps=500,
# Sparsity control
l1_warm_up_steps=1000,
use_ghost_grads=True,
feature_sampling_window=1000,
dead_feature_window=5000,
dead_feature_threshold=1e-8,
# Data
dataset_path="monology/pile-uncopyrighted",
streaming=True,
context_size=128,
# Batch sizes
train_batch_size_tokens=4096,
store_batch_size_prompts=16,
n_batches_in_buffer=64,
# Training duration
training_tokens=100_000_000,
# Logging
log_to_wandb=True,
wandb_project="sae-training",
wandb_log_frequency=100,
# Checkpointing
checkpoint_path="checkpoints",
n_checkpoints=5,
# Hardware
device="cuda",
dtype="float32",
)
```
### Key Parameters Explained
#### Architecture Parameters
| Parameter | Description |
|-----------|-------------|
| `architecture` | SAE type: "standard", "gated", "jumprelu", "topk" |
| `d_sae` | Hidden dimension (or use `expansion_factor`) |
| `expansion_factor` | Alternative to d_sae: d_sae = d_in × expansion_factor |
| `activation_fn` | "relu", "topk", etc. |
| `activation_fn_kwargs` | Dict for activation params (e.g., {"k": 50} for topk) |
#### Sparsity Parameters
| Parameter | Description |
|-----------|-------------|
| `l1_coefficient` | L1 penalty weight (higher = sparser) |
| `l1_warm_up_steps` | Steps to ramp up L1 penalty |
| `use_ghost_grads` | Apply gradients to dead features |
| `dead_feature_threshold` | Activation threshold for "dead" |
| `dead_feature_window` | Steps to check for dead features |
#### Learning Rate Parameters
| Parameter | Description |
|-----------|-------------|
| `lr` | Base learning rate |
| `lr_scheduler_name` | "constant", "cosineannealing", etc. |
| `lr_warm_up_steps` | LR warmup steps |
| `lr_decay_steps` | Steps for LR decay |
---
## SAETrainingRunner
Main class for executing training.
### Basic Training
```python
from sae_lens import SAETrainingRunner, LanguageModelSAERunnerConfig
cfg = LanguageModelSAERunnerConfig(...)
runner = SAETrainingRunner(cfg)
sae = runner.run()
```
### Accessing Training Metrics
```python
# During training, metrics logged to W&B include:
# - l0: Average active features
# - ce_loss_score: Cross-entropy recovery
# - mse_loss: Reconstruction loss
# - l1_loss: Sparsity loss
# - dead_features: Count of dead features
```
---
## ActivationsStore
Manages activation collection and batching.
### Basic Usage
```python
from sae_lens import ActivationsStore
store = ActivationsStore.from_sae(
model=model,
sae=sae,
store_batch_size_prompts=8,
train_batch_size_tokens=4096,
n_batches_in_buffer=32,
device="cuda",
)
# Get batch of activations
activations = store.get_batch_tokens()
```
---
## HookedSAETransformer
Integration of SAEs with TransformerLens models.
### Basic Usage
```python
from sae_lens import HookedSAETransformer
# Load model with SAE
model = HookedSAETransformer.from_pretrained("gpt2-small")
model.add_sae(sae)
# Run with SAE in the loop
output = model.run_with_saes(tokens, saes=[sae])
# Cache with SAE activations
output, cache = model.run_with_cache_with_saes(tokens, saes=[sae])
```
---
## SAE Architectures
### Standard (ReLU + L1)
```python
cfg = LanguageModelSAERunnerConfig(
architecture="standard",
activation_fn="relu",
l1_coefficient=8e-5,
)
```
### Gated
```python
cfg = LanguageModelSAERunnerConfig(
architecture="gated",
)
```
### TopK
```python
cfg = LanguageModelSAERunnerConfig(
architecture="topk",
activation_fn="topk",
activation_fn_kwargs={"k": 50}, # Exactly 50 active features
)
```
### JumpReLU (State-of-the-art)
```python
cfg = LanguageModelSAERunnerConfig(
architecture="jumprelu",
)
```
---
## Utility Functions
### Upload to HuggingFace
```python
from sae_lens import upload_saes_to_huggingface
upload_saes_to_huggingface(
saes=[sae],
repo_id="username/my-saes",
token="hf_token",
)
```
### Neuronpedia Integration
```python
# Features can be viewed on Neuronpedia
# URL format: neuronpedia.org/{model}/{layer}-{sae_type}/{feature_id}
# Example: neuronpedia.org/gpt2-small/8-res-jb/1234
```
@@ -0,0 +1,318 @@
# SAELens Tutorials
## Tutorial 1: Loading and Analyzing Pre-trained SAEs
### Goal
Load a pre-trained SAE and analyze which features activate on specific inputs.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
import torch
# 1. Load model and SAE
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, cfg_dict, sparsity = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
print(f"SAE input dim: {sae.cfg.d_in}")
print(f"SAE hidden dim: {sae.cfg.d_sae}")
print(f"Expansion factor: {sae.cfg.d_sae / sae.cfg.d_in:.1f}x")
# 2. Get model activations
prompt = "The capital of France is Paris"
tokens = model.to_tokens(prompt)
_, cache = model.run_with_cache(tokens)
activations = cache["resid_pre", 8] # [1, seq_len, 768]
# 3. Encode to SAE features
features = sae.encode(activations) # [1, seq_len, d_sae]
# 4. Analyze sparsity
active_per_token = (features > 0).sum(dim=-1)
print(f"Average active features per token: {active_per_token.float().mean():.1f}")
# 5. Find top features for each token
str_tokens = model.to_str_tokens(prompt)
for pos in range(len(str_tokens)):
top_features = features[0, pos].topk(5)
print(f"\nToken '{str_tokens[pos]}':")
for feat_idx, feat_val in zip(top_features.indices, top_features.values):
print(f" Feature {feat_idx.item()}: {feat_val.item():.3f}")
# 6. Check reconstruction quality
reconstructed = sae.decode(features)
mse = ((activations - reconstructed) ** 2).mean()
print(f"\nReconstruction MSE: {mse.item():.6f}")
```
---
## Tutorial 2: Training a Custom SAE
### Goal
Train a Sparse Autoencoder on GPT-2 activations.
### Step-by-Step
```python
from sae_lens import LanguageModelSAERunnerConfig, SAETrainingRunner
# 1. Configure training
cfg = LanguageModelSAERunnerConfig(
# Model
model_name="gpt2-small",
hook_name="blocks.6.hook_resid_pre",
hook_layer=6,
d_in=768,
# SAE architecture
architecture="standard",
d_sae=768 * 8, # 8x expansion
activation_fn="relu",
# Training
lr=4e-4,
l1_coefficient=8e-5,
l1_warm_up_steps=1000,
train_batch_size_tokens=4096,
training_tokens=10_000_000, # Small run for demo
# Data
dataset_path="monology/pile-uncopyrighted",
streaming=True,
context_size=128,
# Dead feature prevention
use_ghost_grads=True,
dead_feature_window=5000,
# Logging
log_to_wandb=True,
wandb_project="sae-training-demo",
# Hardware
device="cuda",
dtype="float32",
)
# 2. Train
runner = SAETrainingRunner(cfg)
sae = runner.run()
# 3. Save
sae.save_model("./my_trained_sae")
```
### Hyperparameter Tuning Guide
| If you see... | Try... |
|---------------|--------|
| High L0 (>200) | Increase `l1_coefficient` |
| Low CE recovery (<80%) | Decrease `l1_coefficient`, increase `d_sae` |
| Many dead features (>5%) | Enable `use_ghost_grads`, increase `l1_warm_up_steps` |
| Training instability | Lower `lr`, increase `lr_warm_up_steps` |
---
## Tutorial 3: Feature Attribution and Steering
### Goal
Identify which SAE features contribute to specific predictions and use them for steering.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
import torch
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, _, _ = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
# 1. Feature attribution for a specific prediction
prompt = "The capital of France is"
tokens = model.to_tokens(prompt)
_, cache = model.run_with_cache(tokens)
activations = cache["resid_pre", 8]
features = sae.encode(activations)
# Target token
target_token = model.to_single_token(" Paris")
# Compute feature contributions to target logit
# contribution = feature_activation * decoder_weight * unembedding
W_dec = sae.W_dec # [d_sae, d_model]
W_U = model.W_U # [d_model, d_vocab]
# Feature direction projected to vocabulary
feature_to_logit = W_dec @ W_U # [d_sae, d_vocab]
# Contribution of each feature to "Paris" at final position
feature_acts = features[0, -1] # [d_sae]
contributions = feature_acts * feature_to_logit[:, target_token]
# Top contributing features
top_features = contributions.topk(10)
print("Top features contributing to 'Paris':")
for idx, val in zip(top_features.indices, top_features.values):
print(f" Feature {idx.item()}: {val.item():.3f}")
# 2. Feature steering
def steer_with_feature(feature_idx, strength=5.0):
"""Add a feature direction to the residual stream."""
feature_direction = sae.W_dec[feature_idx] # [d_model]
def hook(activation, hook_obj):
activation[:, -1, :] += strength * feature_direction
return activation
output = model.generate(
tokens,
max_new_tokens=10,
fwd_hooks=[("blocks.8.hook_resid_pre", hook)]
)
return model.to_string(output[0])
# Try steering with top feature
top_feature_idx = top_features.indices[0].item()
print(f"\nSteering with feature {top_feature_idx}:")
print(steer_with_feature(top_feature_idx, strength=10.0))
```
---
## Tutorial 4: Feature Ablation
### Goal
Test the causal importance of features by ablating them.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
import torch
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, _, _ = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
prompt = "The capital of France is"
tokens = model.to_tokens(prompt)
# Baseline prediction
baseline_logits = model(tokens)
target_token = model.to_single_token(" Paris")
baseline_prob = torch.softmax(baseline_logits[0, -1], dim=-1)[target_token].item()
print(f"Baseline P(Paris): {baseline_prob:.4f}")
# Get features to ablate
_, cache = model.run_with_cache(tokens)
activations = cache["resid_pre", 8]
features = sae.encode(activations)
top_features = features[0, -1].topk(10).indices
# Ablate top features one by one
for feat_idx in top_features:
def ablation_hook(activation, hook, feat_idx=feat_idx):
# Encode → zero feature → decode
feats = sae.encode(activation)
feats[:, :, feat_idx] = 0
return sae.decode(feats)
ablated_logits = model.run_with_hooks(
tokens,
fwd_hooks=[("blocks.8.hook_resid_pre", ablation_hook)]
)
ablated_prob = torch.softmax(ablated_logits[0, -1], dim=-1)[target_token].item()
change = (ablated_prob - baseline_prob) / baseline_prob * 100
print(f"Ablate feature {feat_idx.item()}: P(Paris)={ablated_prob:.4f} ({change:+.1f}%)")
```
---
## Tutorial 5: Comparing Features Across Prompts
### Goal
Find which features activate consistently for a concept.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
import torch
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, _, _ = SAE.from_pretrained(
release="gpt2-small-res-jb",
sae_id="blocks.8.hook_resid_pre",
device="cuda"
)
# Test prompts about the same concept
prompts = [
"The Eiffel Tower is located in",
"Paris is the capital of",
"France's largest city is",
"The Louvre museum is in",
]
# Collect feature activations
all_features = []
for prompt in prompts:
tokens = model.to_tokens(prompt)
_, cache = model.run_with_cache(tokens)
activations = cache["resid_pre", 8]
features = sae.encode(activations)
# Take max activation across positions
max_features = features[0].max(dim=0).values
all_features.append(max_features)
all_features = torch.stack(all_features) # [n_prompts, d_sae]
# Find features that activate consistently
mean_activation = all_features.mean(dim=0)
min_activation = all_features.min(dim=0).values
# Features active in ALL prompts
consistent_features = (min_activation > 0.5).nonzero().squeeze(-1)
print(f"Features active in all prompts: {len(consistent_features)}")
# Top consistent features
top_consistent = mean_activation[consistent_features].topk(min(10, len(consistent_features)))
print("\nTop consistent features (possibly 'France/Paris' related):")
for idx, val in zip(top_consistent.indices, top_consistent.values):
feat_idx = consistent_features[idx].item()
print(f" Feature {feat_idx}: mean activation {val.item():.3f}")
```
---
## External Resources
### Official Tutorials
- [Basic Loading & Analysis](https://github.com/jbloomAus/SAELens/blob/main/tutorials/basic_loading_and_analysing.ipynb)
- [Training SAEs](https://github.com/jbloomAus/SAELens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb)
- [Logits Lens with Features](https://github.com/jbloomAus/SAELens/blob/main/tutorials/logits_lens_with_features.ipynb)
### ARENA Curriculum
Comprehensive SAE course: https://www.lesswrong.com/posts/LnHowHgmrMbWtpkxx/intro-to-superposition-and-sparse-autoencoders-colab
### Key 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 Interpretable Features](https://arxiv.org/abs/2309.08600) - ICLR 2024
@@ -0,0 +1,346 @@
---
name: transformer-lens-interpretability
description: Provides guidance for mechanistic interpretability research using TransformerLens to inspect and manipulate transformer internals via HookPoints and activation caching. Use when reverse-engineering model algorithms, studying attention patterns, or performing activation patching experiments.
version: 1.0.0
author: Orchestra Research
license: MIT
tags: [Mechanistic Interpretability, TransformerLens, Activation Patching, Circuit Analysis]
dependencies: [transformer-lens>=2.0.0, torch>=2.0.0]
---
# TransformerLens: Mechanistic Interpretability for Transformers
TransformerLens is the de facto standard library for mechanistic interpretability research on GPT-style language models. Created by Neel Nanda and maintained by Bryce Meyer, it provides clean interfaces to inspect and manipulate model internals via HookPoints on every activation.
**GitHub**: [TransformerLensOrg/TransformerLens](https://github.com/TransformerLensOrg/TransformerLens) (2,900+ stars)
## When to Use TransformerLens
**Use TransformerLens when you need to:**
- Reverse-engineer algorithms learned during training
- Perform activation patching / causal tracing experiments
- Study attention patterns and information flow
- Analyze circuits (e.g., induction heads, IOI circuit)
- Cache and inspect intermediate activations
- Apply direct logit attribution
**Consider alternatives when:**
- You need to work with non-transformer architectures → Use **nnsight** or **pyvene**
- You want to train/analyze Sparse Autoencoders → Use **SAELens**
- You need remote execution on massive models → Use **nnsight** with NDIF
- You want higher-level causal intervention abstractions → Use **pyvene**
## Installation
```bash
pip install transformer-lens
```
For development version:
```bash
pip install git+https://github.com/TransformerLensOrg/TransformerLens
```
## Core Concepts
### HookedTransformer
The main class that wraps transformer models with HookPoints on every activation:
```python
from transformer_lens import HookedTransformer
# Load a model
model = HookedTransformer.from_pretrained("gpt2-small")
# For gated models (LLaMA, Mistral)
import os
os.environ["HF_TOKEN"] = "your_token"
model = HookedTransformer.from_pretrained("meta-llama/Llama-2-7b-hf")
```
### Supported Models (50+)
| Family | Models |
|--------|--------|
| GPT-2 | gpt2, gpt2-medium, gpt2-large, gpt2-xl |
| LLaMA | llama-7b, llama-13b, llama-2-7b, llama-2-13b |
| EleutherAI | pythia-70m to pythia-12b, gpt-neo, gpt-j-6b |
| Mistral | mistral-7b, mixtral-8x7b |
| Others | phi, qwen, opt, gemma |
### Activation Caching
Run the model and cache all intermediate activations:
```python
# Get all activations
tokens = model.to_tokens("The Eiffel Tower is in")
logits, cache = model.run_with_cache(tokens)
# Access specific activations
residual = cache["resid_post", 5] # Layer 5 residual stream
attn_pattern = cache["pattern", 3] # Layer 3 attention pattern
mlp_out = cache["mlp_out", 7] # Layer 7 MLP output
# Filter which activations to cache (saves memory)
logits, cache = model.run_with_cache(
tokens,
names_filter=lambda name: "resid_post" in name
)
```
### ActivationCache Keys
| Key Pattern | Shape | Description |
|-------------|-------|-------------|
| `resid_pre, layer` | [batch, pos, d_model] | Residual before attention |
| `resid_mid, layer` | [batch, pos, d_model] | Residual after attention |
| `resid_post, layer` | [batch, pos, d_model] | Residual after MLP |
| `attn_out, layer` | [batch, pos, d_model] | Attention output |
| `mlp_out, layer` | [batch, pos, d_model] | MLP output |
| `pattern, layer` | [batch, head, q_pos, k_pos] | Attention pattern (post-softmax) |
| `q, layer` | [batch, pos, head, d_head] | Query vectors |
| `k, layer` | [batch, pos, head, d_head] | Key vectors |
| `v, layer` | [batch, pos, head, d_head] | Value vectors |
## Workflow 1: Activation Patching (Causal Tracing)
Identify which activations causally affect model output by patching clean activations into corrupted runs.
### Step-by-Step
```python
from transformer_lens import HookedTransformer, patching
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
# 1. Define clean and corrupted prompts
clean_prompt = "The Eiffel Tower is in the city of"
corrupted_prompt = "The Colosseum is in the city of"
clean_tokens = model.to_tokens(clean_prompt)
corrupted_tokens = model.to_tokens(corrupted_prompt)
# 2. Get clean activations
_, clean_cache = model.run_with_cache(clean_tokens)
# 3. Define metric (e.g., logit difference)
paris_token = model.to_single_token(" Paris")
rome_token = model.to_single_token(" Rome")
def metric(logits):
return logits[0, -1, paris_token] - logits[0, -1, rome_token]
# 4. Patch each position and layer
results = torch.zeros(model.cfg.n_layers, clean_tokens.shape[1])
for layer in range(model.cfg.n_layers):
for pos in range(clean_tokens.shape[1]):
def patch_hook(activation, hook):
activation[0, pos] = clean_cache[hook.name][0, pos]
return activation
patched_logits = model.run_with_hooks(
corrupted_tokens,
fwd_hooks=[(f"blocks.{layer}.hook_resid_post", patch_hook)]
)
results[layer, pos] = metric(patched_logits)
# 5. Visualize results (layer x position heatmap)
```
### Checklist
- [ ] Define clean and corrupted inputs that differ minimally
- [ ] Choose metric that captures behavior difference
- [ ] Cache clean activations
- [ ] Systematically patch each (layer, position) combination
- [ ] Visualize results as heatmap
- [ ] Identify causal hotspots
## Workflow 2: Circuit Analysis (Indirect Object Identification)
Replicate the IOI circuit discovery from "Interpretability in the Wild".
### Step-by-Step
```python
from transformer_lens import HookedTransformer
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
# IOI task: "When John and Mary went to the store, Mary gave a bottle to"
# Model should predict "John" (indirect object)
prompt = "When John and Mary went to the store, Mary gave a bottle to"
tokens = model.to_tokens(prompt)
# 1. Get baseline logits
logits, cache = model.run_with_cache(tokens)
john_token = model.to_single_token(" John")
mary_token = model.to_single_token(" Mary")
# 2. Compute logit difference (IO - S)
logit_diff = logits[0, -1, john_token] - logits[0, -1, mary_token]
print(f"Logit difference: {logit_diff.item():.3f}")
# 3. Direct logit attribution by head
def get_head_contribution(layer, head):
# Project head output to logits
head_out = cache["z", layer][0, :, head, :] # [pos, d_head]
W_O = model.W_O[layer, head] # [d_head, d_model]
W_U = model.W_U # [d_model, vocab]
# Head contribution to logits at final position
contribution = head_out[-1] @ W_O @ W_U
return contribution[john_token] - contribution[mary_token]
# 4. Map all heads
head_contributions = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)
for layer in range(model.cfg.n_layers):
for head in range(model.cfg.n_heads):
head_contributions[layer, head] = get_head_contribution(layer, head)
# 5. Identify top contributing heads (name movers, backup name movers)
```
### Checklist
- [ ] Set up task with clear IO/S tokens
- [ ] Compute baseline logit difference
- [ ] Decompose by attention head contributions
- [ ] Identify key circuit components (name movers, S-inhibition, induction)
- [ ] Validate with ablation experiments
## Workflow 3: Induction Head Detection
Find induction heads that implement [A][B]...[A] → [B] pattern.
```python
from transformer_lens import HookedTransformer
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
# Create repeated sequence: [A][B][A] should predict [B]
repeated_tokens = torch.tensor([[1000, 2000, 1000]]) # Arbitrary tokens
_, cache = model.run_with_cache(repeated_tokens)
# Induction heads attend from final [A] back to first [B]
# Check attention from position 2 to position 1
induction_scores = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)
for layer in range(model.cfg.n_layers):
pattern = cache["pattern", layer][0] # [head, q_pos, k_pos]
# Attention from pos 2 to pos 1
induction_scores[layer] = pattern[:, 2, 1]
# Heads with high scores are induction heads
top_heads = torch.topk(induction_scores.flatten(), k=5)
```
## Common Issues & Solutions
### Issue: Hooks persist after debugging
```python
# WRONG: Old hooks remain active
model.run_with_hooks(tokens, fwd_hooks=[...]) # Debug, add new hooks
model.run_with_hooks(tokens, fwd_hooks=[...]) # Old hooks still there!
# RIGHT: Always reset hooks
model.reset_hooks()
model.run_with_hooks(tokens, fwd_hooks=[...])
```
### Issue: Tokenization gotchas
```python
# WRONG: Assuming consistent tokenization
model.to_tokens("Tim") # Single token
model.to_tokens("Neel") # Becomes "Ne" + "el" (two tokens!)
# RIGHT: Check tokenization explicitly
tokens = model.to_tokens("Neel", prepend_bos=False)
print(model.to_str_tokens(tokens)) # ['Ne', 'el']
```
### Issue: LayerNorm ignored in analysis
```python
# WRONG: Ignoring LayerNorm
pre_activation = residual @ model.W_in[layer]
# RIGHT: Include LayerNorm
ln_scale = model.blocks[layer].ln2.w
ln_out = model.blocks[layer].ln2(residual)
pre_activation = ln_out @ model.W_in[layer]
```
### Issue: Memory explosion with large models
```python
# Use selective caching
logits, cache = model.run_with_cache(
tokens,
names_filter=lambda n: "resid_post" in n or "pattern" in n,
device="cpu" # Cache on CPU
)
```
## Key Classes Reference
| Class | Purpose |
|-------|---------|
| `HookedTransformer` | Main model wrapper with hooks |
| `ActivationCache` | Dictionary-like cache of activations |
| `HookedTransformerConfig` | Model configuration |
| `FactoredMatrix` | Efficient factored matrix operations |
## Integration with SAELens
TransformerLens integrates with SAELens for Sparse Autoencoder analysis:
```python
from transformer_lens import HookedTransformer
from sae_lens import SAE
model = HookedTransformer.from_pretrained("gpt2-small")
sae = SAE.from_pretrained("gpt2-small-res-jb", "blocks.8.hook_resid_pre")
# Run with SAE
tokens = model.to_tokens("Hello world")
_, cache = model.run_with_cache(tokens)
sae_acts = sae.encode(cache["resid_pre", 8])
```
## 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 HookedTransformer, ActivationCache, HookPoints |
| [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for activation patching, circuit analysis, logit lens |
## External Resources
### Tutorials
- [Main Demo Notebook](https://transformerlensorg.github.io/TransformerLens/generated/demos/Main_Demo.html)
- [Activation Patching Demo](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Activation_Patching_in_TL_Demo.ipynb)
- [ARENA Mech Interp Course](https://arena-foundation.github.io/ARENA/) - 200+ hours of tutorials
### Papers
- [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html)
- [In-context Learning and Induction Heads](https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html)
- [Interpretability in the Wild (IOI)](https://arxiv.org/abs/2211.00593)
### Official Documentation
- [Official Docs](https://transformerlensorg.github.io/TransformerLens/)
- [Model Properties Table](https://transformerlensorg.github.io/TransformerLens/generated/model_properties_table.html)
- [Neel Nanda's Glossary](https://www.neelnanda.io/mechanistic-interpretability/glossary)
## Version Notes
- **v2.0**: Removed HookedSAE (moved to SAELens)
- **v3.0 (alpha)**: TransformerBridge for loading any nn.Module
@@ -0,0 +1,54 @@
# TransformerLens Reference Documentation
This directory contains comprehensive reference materials for TransformerLens.
## Contents
- [api.md](api.md) - Complete API reference for HookedTransformer, ActivationCache, and HookPoints
- [tutorials.md](tutorials.md) - Step-by-step tutorials for common interpretability workflows
- [papers.md](papers.md) - Key research papers and foundational concepts
## Quick Links
- **Official Documentation**: https://transformerlensorg.github.io/TransformerLens/
- **GitHub Repository**: https://github.com/TransformerLensOrg/TransformerLens
- **Model Properties Table**: https://transformerlensorg.github.io/TransformerLens/generated/model_properties_table.html
## Installation
```bash
pip install transformer-lens
```
## Basic Usage
```python
from transformer_lens import HookedTransformer
# Load model
model = HookedTransformer.from_pretrained("gpt2-small")
# Run with activation caching
tokens = model.to_tokens("Hello world")
logits, cache = model.run_with_cache(tokens)
# Access activations
residual = cache["resid_post", 5] # Layer 5 residual stream
attention = cache["pattern", 3] # Layer 3 attention patterns
```
## Key Concepts
### HookPoints
Every activation in the transformer has a HookPoint wrapper, enabling:
- Reading activations via `run_with_cache()`
- Modifying activations via `run_with_hooks()`
### Activation Cache
The `ActivationCache` stores all intermediate activations with helper methods for:
- Residual stream decomposition
- Logit attribution
- Layer-wise analysis
### Supported Models (50+)
GPT-2, LLaMA, Mistral, Pythia, GPT-Neo, OPT, Gemma, Phi, and more.
@@ -0,0 +1,362 @@
# TransformerLens API Reference
## HookedTransformer
The core class for mechanistic interpretability, wrapping transformer models with hooks on every activation.
### Loading Models
```python
from transformer_lens import HookedTransformer
# Basic loading
model = HookedTransformer.from_pretrained("gpt2-small")
# With specific device/dtype
model = HookedTransformer.from_pretrained(
"gpt2-medium",
device="cuda",
dtype=torch.float16
)
# Gated models (LLaMA, Mistral)
import os
os.environ["HF_TOKEN"] = "your_token"
model = HookedTransformer.from_pretrained("meta-llama/Llama-2-7b-hf")
```
### from_pretrained() Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `model_name` | str | required | Model name from OFFICIAL_MODEL_NAMES |
| `fold_ln` | bool | True | Fold LayerNorm weights into subsequent layers |
| `center_writing_weights` | bool | True | Center residual stream writer means |
| `center_unembed` | bool | True | Center unembedding weights |
| `dtype` | torch.dtype | None | Model precision |
| `device` | str | None | Target device |
| `n_devices` | int | 1 | Number of devices for model parallelism |
### Weight Matrices
| Property | Shape | Description |
|----------|-------|-------------|
| `W_E` | [d_vocab, d_model] | Token embedding matrix |
| `W_U` | [d_model, d_vocab] | Unembedding matrix |
| `W_pos` | [n_ctx, d_model] | Positional embedding |
| `W_Q` | [n_layers, n_heads, d_model, d_head] | Query weights |
| `W_K` | [n_layers, n_heads, d_model, d_head] | Key weights |
| `W_V` | [n_layers, n_heads, d_model, d_head] | Value weights |
| `W_O` | [n_layers, n_heads, d_head, d_model] | Output weights |
| `W_in` | [n_layers, d_model, d_mlp] | MLP input weights |
| `W_out` | [n_layers, d_mlp, d_model] | MLP output weights |
### Core Methods
#### forward()
```python
logits = model(tokens)
logits = model(tokens, return_type="logits")
loss = model(tokens, return_type="loss")
logits, loss = model(tokens, return_type="both")
```
Parameters:
- `input`: Token tensor or string
- `return_type`: "logits", "loss", "both", or None
- `prepend_bos`: Whether to prepend BOS token
- `start_at_layer`: Start execution from specific layer
- `stop_at_layer`: Stop execution at specific layer
#### run_with_cache()
```python
logits, cache = model.run_with_cache(tokens)
# Selective caching (saves memory)
logits, cache = model.run_with_cache(
tokens,
names_filter=lambda name: "resid_post" in name
)
# Cache on CPU
logits, cache = model.run_with_cache(tokens, device="cpu")
```
#### run_with_hooks()
```python
def my_hook(activation, hook):
# Modify activation
activation[:, :, 0] = 0
return activation
logits = model.run_with_hooks(
tokens,
fwd_hooks=[("blocks.5.hook_resid_post", my_hook)]
)
```
#### generate()
```python
output = model.generate(
tokens,
max_new_tokens=50,
temperature=0.7,
top_k=40,
top_p=0.9,
freq_penalty=1.0,
use_past_kv_cache=True
)
```
### Tokenization Methods
```python
# String to tokens
tokens = model.to_tokens("Hello world") # [1, seq_len]
tokens = model.to_tokens("Hello", prepend_bos=False)
# Tokens to string
text = model.to_string(tokens)
# Get string tokens (for debugging)
str_tokens = model.to_str_tokens("Hello world")
# ['<|endoftext|>', 'Hello', ' world']
# Single token validation
token_id = model.to_single_token(" Paris") # Returns int or raises error
```
### Hook Management
```python
# Clear all hooks
model.reset_hooks()
# Add permanent hook
model.add_hook("blocks.0.hook_resid_post", my_hook)
# Remove specific hook
model.remove_hook("blocks.0.hook_resid_post")
```
---
## ActivationCache
Stores and provides access to all activations from a forward pass.
### Accessing Activations
```python
logits, cache = model.run_with_cache(tokens)
# By name and layer
residual = cache["resid_post", 5]
attention = cache["pattern", 3]
mlp_out = cache["mlp_out", 7]
# Full name string
residual = cache["blocks.5.hook_resid_post"]
```
### Cache Keys
| Key Pattern | Shape | Description |
|-------------|-------|-------------|
| `hook_embed` | [batch, pos, d_model] | Token embeddings |
| `hook_pos_embed` | [batch, pos, d_model] | Positional embeddings |
| `resid_pre, layer` | [batch, pos, d_model] | Residual before attention |
| `resid_mid, layer` | [batch, pos, d_model] | Residual after attention |
| `resid_post, layer` | [batch, pos, d_model] | Residual after MLP |
| `attn_out, layer` | [batch, pos, d_model] | Attention output |
| `mlp_out, layer` | [batch, pos, d_model] | MLP output |
| `pattern, layer` | [batch, head, q_pos, k_pos] | Attention pattern (post-softmax) |
| `attn_scores, layer` | [batch, head, q_pos, k_pos] | Attention scores (pre-softmax) |
| `q, layer` | [batch, pos, head, d_head] | Query vectors |
| `k, layer` | [batch, pos, head, d_head] | Key vectors |
| `v, layer` | [batch, pos, head, d_head] | Value vectors |
| `z, layer` | [batch, pos, head, d_head] | Attention output per head |
### Analysis Methods
#### decompose_resid()
Decomposes residual stream into component contributions:
```python
components, labels = cache.decompose_resid(
layer=5,
return_labels=True,
mode="attn" # or "mlp" or "full"
)
```
#### accumulated_resid()
Get accumulated residual at each layer (for Logit Lens):
```python
accumulated = cache.accumulated_resid(
layer=None, # All layers
incl_mid=False,
apply_ln=True # Apply final LayerNorm
)
```
#### logit_attrs()
Calculate logit attribution for components:
```python
attrs = cache.logit_attrs(
residual_stack,
tokens=target_tokens,
incorrect_tokens=incorrect_tokens
)
```
#### stack_head_results()
Stack attention head outputs:
```python
head_results = cache.stack_head_results(
layer=-1, # All layers
pos_slice=None # All positions
)
# Shape: [n_layers, n_heads, batch, pos, d_model]
```
### Utility Methods
```python
# Move cache to device
cache = cache.to("cpu")
# Remove batch dimension (for batch_size=1)
cache = cache.remove_batch_dim()
# Get all keys
keys = cache.keys()
# Iterate
for name, activation in cache.items():
print(name, activation.shape)
```
---
## HookPoint
The fundamental hook mechanism wrapping every activation.
### Hook Function Signature
```python
def hook_fn(activation: torch.Tensor, hook: HookPoint) -> torch.Tensor:
"""
Args:
activation: Current activation value
hook: The HookPoint object (has .name attribute)
Returns:
Modified activation (or None to keep original)
"""
# Modify activation
return activation
```
### Common Hook Patterns
```python
# Zero ablation
def zero_hook(act, hook):
act[:, :, :] = 0
return act
# Mean ablation
def mean_hook(act, hook):
act[:, :, :] = act.mean(dim=0, keepdim=True)
return act
# Patch from cache
def patch_hook(act, hook):
act[:, 5, :] = clean_cache[hook.name][:, 5, :]
return act
# Add steering vector
def steer_hook(act, hook):
act += 0.5 * steering_vector
return act
```
---
## Utility Functions
### patching module
```python
from transformer_lens import patching
# Generic activation patching
results = patching.generic_activation_patch(
model=model,
corrupted_tokens=corrupted,
clean_cache=clean_cache,
patching_metric=metric_fn,
patch_setter=patch_fn,
activation_name="resid_post",
index_axis_names=("layer", "pos")
)
```
### FactoredMatrix
Efficient operations on factored weight matrices:
```python
from transformer_lens import FactoredMatrix
# QK circuit
QK = FactoredMatrix(model.W_Q[layer], model.W_K[layer].T)
# OV circuit
OV = FactoredMatrix(model.W_V[layer], model.W_O[layer])
# Get full matrix
full = QK.AB
# SVD decomposition
U, S, V = QK.svd()
```
---
## Configuration
### HookedTransformerConfig
Key configuration attributes:
| Attribute | Description |
|-----------|-------------|
| `n_layers` | Number of transformer layers |
| `n_heads` | Number of attention heads |
| `d_model` | Model dimension |
| `d_head` | Head dimension |
| `d_mlp` | MLP hidden dimension |
| `d_vocab` | Vocabulary size |
| `n_ctx` | Maximum context length |
| `act_fn` | Activation function name |
| `normalization_type` | "LN" or "LNPre" |
Access via:
```python
model.cfg.n_layers
model.cfg.d_model
```
@@ -0,0 +1,339 @@
# TransformerLens Tutorials
## Tutorial 1: Basic Activation Analysis
### Goal
Understand how to load models, cache activations, and inspect model internals.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
import torch
# 1. Load model
model = HookedTransformer.from_pretrained("gpt2-small")
print(f"Model has {model.cfg.n_layers} layers, {model.cfg.n_heads} heads")
# 2. Tokenize input
prompt = "The capital of France is"
tokens = model.to_tokens(prompt)
print(f"Tokens shape: {tokens.shape}")
print(f"String tokens: {model.to_str_tokens(prompt)}")
# 3. Run with cache
logits, cache = model.run_with_cache(tokens)
print(f"Logits shape: {logits.shape}")
print(f"Cache keys: {len(cache.keys())}")
# 4. Inspect activations
for layer in range(model.cfg.n_layers):
resid = cache["resid_post", layer]
print(f"Layer {layer} residual norm: {resid.norm().item():.2f}")
# 5. Look at attention patterns
attn = cache["pattern", 0] # Layer 0
print(f"Attention shape: {attn.shape}") # [batch, heads, q_pos, k_pos]
# 6. Get top predictions
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
for token_id, prob in zip(top_tokens.indices, top_tokens.values):
print(f"{model.to_string(token_id.unsqueeze(0))}: {prob.item():.3f}")
```
---
## Tutorial 2: Activation Patching
### Goal
Identify which activations causally affect model output.
### Concept
1. Run model on "clean" input, cache activations
2. Run model on "corrupted" input
3. Patch clean activations into corrupted run
4. Measure effect on output
### Step-by-Step
```python
from transformer_lens import HookedTransformer
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
# Define clean and corrupted prompts
clean_prompt = "The Eiffel Tower is in the city of"
corrupted_prompt = "The Colosseum is in the city of"
clean_tokens = model.to_tokens(clean_prompt)
corrupted_tokens = model.to_tokens(corrupted_prompt)
# Get clean activations
_, clean_cache = model.run_with_cache(clean_tokens)
# Define metric
paris_token = model.to_single_token(" Paris")
rome_token = model.to_single_token(" Rome")
def logit_diff(logits):
"""Positive = model prefers Paris over Rome"""
return (logits[0, -1, paris_token] - logits[0, -1, rome_token]).item()
# Baseline measurements
clean_logits = model(clean_tokens)
corrupted_logits = model(corrupted_tokens)
print(f"Clean logit diff: {logit_diff(clean_logits):.3f}")
print(f"Corrupted logit diff: {logit_diff(corrupted_logits):.3f}")
# Patch each layer
results = []
for layer in range(model.cfg.n_layers):
def patch_hook(activation, hook, layer=layer):
activation[:] = clean_cache["resid_post", layer]
return activation
patched_logits = model.run_with_hooks(
corrupted_tokens,
fwd_hooks=[(f"blocks.{layer}.hook_resid_post", patch_hook)]
)
results.append(logit_diff(patched_logits))
print(f"Layer {layer}: {results[-1]:.3f}")
# Find most important layer
best_layer = max(range(len(results)), key=lambda i: results[i])
print(f"\nMost important layer: {best_layer}")
```
### Position-Specific Patching
```python
import torch
seq_len = clean_tokens.shape[1]
results = torch.zeros(model.cfg.n_layers, seq_len)
for layer in range(model.cfg.n_layers):
for pos in range(seq_len):
def patch_hook(activation, hook, layer=layer, pos=pos):
activation[:, pos, :] = clean_cache["resid_post", layer][:, pos, :]
return activation
patched_logits = model.run_with_hooks(
corrupted_tokens,
fwd_hooks=[(f"blocks.{layer}.hook_resid_post", patch_hook)]
)
results[layer, pos] = logit_diff(patched_logits)
# Visualize as heatmap
import matplotlib.pyplot as plt
plt.figure(figsize=(12, 8))
plt.imshow(results.numpy(), aspect='auto', cmap='RdBu')
plt.xlabel('Position')
plt.ylabel('Layer')
plt.colorbar(label='Logit Difference')
plt.title('Activation Patching Results')
```
---
## Tutorial 3: Direct Logit Attribution
### Goal
Identify which components (heads, neurons) contribute to specific predictions.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
prompt = "The capital of France is"
tokens = model.to_tokens(prompt)
logits, cache = model.run_with_cache(tokens)
# Target token
target_token = model.to_single_token(" Paris")
# Get unembedding direction for target
target_direction = model.W_U[:, target_token] # [d_model]
# Attribution per attention head
head_contributions = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)
for layer in range(model.cfg.n_layers):
# Get per-head output at final position
z = cache["z", layer][0, -1] # [n_heads, d_head]
for head in range(model.cfg.n_heads):
# Project through W_O to get contribution to residual
head_out = z[head] @ model.W_O[layer, head] # [d_model]
# Dot with target direction
contribution = (head_out @ target_direction).item()
head_contributions[layer, head] = contribution
# Find top contributing heads
flat_idx = head_contributions.flatten().topk(10)
print("Top 10 heads for predicting 'Paris':")
for idx, val in zip(flat_idx.indices, flat_idx.values):
layer = idx.item() // model.cfg.n_heads
head = idx.item() % model.cfg.n_heads
print(f" L{layer}H{head}: {val.item():.3f}")
```
---
## Tutorial 4: Induction Head Detection
### Goal
Find attention heads that implement the [A][B]...[A] → [B] pattern.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
# Create repeated sequence pattern
# Pattern: [A][B][C][A] - model should attend from last A to B
seq = torch.randint(1000, 5000, (1, 20))
# Repeat first half
seq[0, 10:] = seq[0, :10]
_, cache = model.run_with_cache(seq)
# For induction heads: position i should attend to position (i - seq_len/2 + 1)
# At position 10 (second A), should attend to position 1 (first B)
induction_scores = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)
for layer in range(model.cfg.n_layers):
pattern = cache["pattern", layer][0] # [heads, q_pos, k_pos]
# Check attention from repeated positions to position after first occurrence
for offset in range(1, 10):
q_pos = 10 + offset # Position in second half
k_pos = offset # Should attend to corresponding position in first half
# Average attention to the "correct" position
induction_scores[layer] += pattern[:, q_pos, k_pos]
induction_scores[layer] /= 9 # Average over offsets
# Find top induction heads
print("Top induction heads:")
for layer in range(model.cfg.n_layers):
for head in range(model.cfg.n_heads):
score = induction_scores[layer, head].item()
if score > 0.3:
print(f" L{layer}H{head}: {score:.3f}")
```
---
## Tutorial 5: Logit Lens
### Goal
See what the model "believes" at each layer before final unembedding.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
prompt = "The quick brown fox jumps over the lazy"
tokens = model.to_tokens(prompt)
logits, cache = model.run_with_cache(tokens)
# Get accumulated residual at each layer
# Apply LayerNorm to match what unembedding sees
accumulated = cache.accumulated_resid(layer=None, incl_mid=False, apply_ln=True)
# Shape: [n_layers + 1, batch, pos, d_model]
# Project to vocabulary
layer_logits = accumulated @ model.W_U # [n_layers + 1, batch, pos, d_vocab]
# Look at predictions for final position
print("Layer-by-layer predictions for final token:")
for layer in range(model.cfg.n_layers + 1):
probs = torch.softmax(layer_logits[layer, 0, -1], dim=-1)
top_token = probs.argmax()
top_prob = probs[top_token].item()
print(f"Layer {layer}: {model.to_string(top_token.unsqueeze(0))!r} ({top_prob:.3f})")
```
---
## Tutorial 6: Steering with Activation Addition
### Goal
Add a steering vector to change model behavior.
### Step-by-Step
```python
from transformer_lens import HookedTransformer
import torch
model = HookedTransformer.from_pretrained("gpt2-small")
# Get activations for contrasting prompts
positive_prompt = "I love this! It's absolutely wonderful and"
negative_prompt = "I hate this! It's absolutely terrible and"
_, pos_cache = model.run_with_cache(model.to_tokens(positive_prompt))
_, neg_cache = model.run_with_cache(model.to_tokens(negative_prompt))
# Compute steering vector (positive - negative direction)
layer = 6
steering_vector = (
pos_cache["resid_post", layer].mean(dim=1) -
neg_cache["resid_post", layer].mean(dim=1)
)
# Generate with steering
test_prompt = "The movie was"
test_tokens = model.to_tokens(test_prompt)
def steer_hook(activation, hook):
activation += 2.0 * steering_vector
return activation
# Without steering
normal_output = model.generate(test_tokens, max_new_tokens=20)
print(f"Normal: {model.to_string(normal_output[0])}")
# With positive steering
steered_output = model.generate(
test_tokens,
max_new_tokens=20,
fwd_hooks=[(f"blocks.{layer}.hook_resid_post", steer_hook)]
)
print(f"Steered: {model.to_string(steered_output[0])}")
```
---
## External Resources
### Official Tutorials
- [Main Demo](https://transformerlensorg.github.io/TransformerLens/generated/demos/Main_Demo.html)
- [Exploratory Analysis](https://transformerlensorg.github.io/TransformerLens/generated/demos/Exploratory_Analysis_Demo.html)
- [Activation Patching Demo](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Activation_Patching_in_TL_Demo.ipynb)
### ARENA Course
Comprehensive 200+ hour curriculum: https://arena-foundation.github.io/ARENA/
### Neel Nanda's Resources
- [Getting Started in Mech Interp](https://www.neelnanda.io/mechanistic-interpretability/getting-started)
- [Mech Interp Glossary](https://www.neelnanda.io/mechanistic-interpretability/glossary)
- [YouTube Channel](https://www.youtube.com/@neelnanda)
+30 -9
View File
@@ -33,7 +33,7 @@ Modern AI research requires mastering dozens of specialized tools and frameworks
AI Researchers spend more time debugging infrastructure than testing hypotheses—slowing the pace of scientific discovery.
We provide a comprehensive library of expert-level research engineering skills that enable AI agents to autonomously implement and execute different stages of AI research experiments—from data preparation and model training to evaluation and deployment.
- Specialized Expertise - Each skill provides deep, production-ready knowledge of a specific framework (Megatron-LM, vLLM, TRL, etc.)
- End-to-End Coverage - 70 skills spanning model architecture, tokenization, fine-tuning, data processing, post-training, distributed training, optimization, inference, infrastructure, agents, RAG, multimodal, prompt engineering, MLOps, observability, and emerging techniques
- End-to-End Coverage - 74 skills spanning model architecture, tokenization, fine-tuning, mechanistic interpretability, data processing, post-training, distributed training, optimization, inference, infrastructure, agents, RAG, multimodal, prompt engineering, MLOps, observability, and emerging techniques
- Research-Grade Quality - Documentation sourced from official repos, real GitHub issues, and battle-tested production workflows
## Available AI Research Engineering Skills
@@ -70,6 +70,12 @@ Install individual skills directly from the marketplace using the Claude Code CL
- **[Unsloth](03-fine-tuning/unsloth/)** - 2x faster QLoRA fine-tuning (75 lines + 4 refs)
- **[PEFT](03-fine-tuning/peft/)** - Parameter-efficient fine-tuning with LoRA, QLoRA, DoRA, 25+ methods (431 lines + 2 refs)
### 🔬 Mechanistic Interpretability (4 skills)
- **[TransformerLens](04-mechanistic-interpretability/transformer-lens/)** - Neel Nanda's library for mech interp with HookPoints, activation caching (346 lines + 3 refs)
- **[SAELens](04-mechanistic-interpretability/saelens/)** - Sparse Autoencoder training and analysis for feature discovery (386 lines + 3 refs)
- **[pyvene](04-mechanistic-interpretability/pyvene/)** - Stanford's causal intervention library with declarative configs (473 lines + 3 refs)
- **[nnsight](04-mechanistic-interpretability/nnsight/)** - Remote interpretability via NDIF, run experiments on 70B+ models (436 lines + 3 refs)
### 📊 Data Processing (2 skills)
- **[Ray Data](05-data-processing/ray-data/)** - Distributed ML data processing, streaming execution, GPU support (318 lines + 2 refs)
- **[NeMo Curator](05-data-processing/nemo-curator/)** - GPU-accelerated data curation, 16× faster deduplication (375 lines + 2 refs)
@@ -161,12 +167,13 @@ Install individual skills directly from the marketplace using the Claude Code CL
- **[Model Pruning](19-emerging-techniques/model-pruning/)** - 50% sparsity with Wanda, SparseGPT, <1% accuracy loss (417 lines)
**Available skills in Claude marketplace** (70 total):
**Available skills in Claude marketplace** (74 total):
| Category | Skills |
|----------|--------|
| Model Architecture | `implementing-llms-litgpt`, `mamba-architecture`, `nanogpt`, `rwkv-architecture` |
| Tokenization | `huggingface-tokenizers`, `sentencepiece` |
| Fine-Tuning | `axolotl`, `llama-factory`, `peft-fine-tuning`, `unsloth` |
| Mechanistic Interpretability | `transformer-lens-interpretability`, `sparse-autoencoder-training`, `pyvene-interventions`, `nnsight-remote-interpretability` |
| Data Processing | `nemo-curator`, `ray-data` |
| Post-Training | `grpo-rl-training`, `openrlhf-training`, `simpo-training`, `fine-tuning-with-trl` |
| Safety | `constitutional-ai`, `llamaguard`, `nemo-guardrails` |
@@ -185,7 +192,7 @@ Install individual skills directly from the marketplace using the Claude Code CL
## Demo
All 70 skills in this repo are automatically synced to [Orchestra Research](https://www.orchestra-research.com/research-skills), where you can add them to your projects with one click and use them with AI research agents.
All 74 skills in this repo are automatically synced to [Orchestra Research](https://www.orchestra-research.com/research-skills), where you can add them to your projects with one click and use them with AI research agents.
**[Demo](https://www.orchestra-research.com/perspectives/LLM-with-Orchestra)**: With this `skills`, a physics PhD is able to [reproduce](https://www.orchestra-research.com/perspectives/LLM-with-Orchestra) Thinking Machines Lab's "LoRA Without Regret" findings.
The Orchestra agent autonomously wrote training code using TRL, provisioned H100 GPUs, ran GRPO experiments overnight, and generated publication-ready analysis, successfully validating that rank=16 LoRA achieves 99.4% of rank=256's SFT performance and that rank=1 LoRA outperforms full fine-tuning on RL tasks (52.1% vs 33.3% on GSM8k math reasoning). ([Video demo](https://www.youtube.com/watch?v=X0DoLYfXl5I))
@@ -264,7 +271,7 @@ skill-name/
## Roadmap
We're building towards 70 comprehensive skills across the full AI research lifecycle. See our [detailed roadmap](ROADMAP.md) for the complete development plan.
We're building towards 80 comprehensive skills across the full AI research lifecycle. See our [detailed roadmap](ROADMAP.md) for the complete development plan.
[View Full Roadmap →](ROADMAP.md)
@@ -273,14 +280,14 @@ We're building towards 70 comprehensive skills across the full AI research lifec
| Metric | Current | Target |
|--------|---------|--------|
| **Skills** | **70** (high-quality, standardized YAML) | 70 ✅ |
| **Skills** | **74** (high-quality, standardized YAML) | 80 |
| **Avg Lines/Skill** | **420 lines** (focused + progressive disclosure) | 200-600 lines |
| **Documentation** | **~115,000 lines** total (SKILL.md + references) | 100,000+ lines |
| **Gold Standard Skills** | **58** with comprehensive references | 50+ |
| **Documentation** | **~120,000 lines** total (SKILL.md + references) | 100,000+ lines |
| **Gold Standard Skills** | **62** with comprehensive references | 50+ |
| **Contributors** | 1 | 100+ |
| **Coverage** | Architecture, Tokenization, Fine-Tuning, Data Processing, Post-Training, Safety, Distributed, Optimization, Evaluation, Infrastructure, Inference, Agents, RAG, Multimodal, Prompt Engineering, MLOps, Observability | Full Lifecycle ✅ |
| **Coverage** | Architecture, Tokenization, Fine-Tuning, Mechanistic Interpretability, Data Processing, Post-Training, Safety, Distributed, Optimization, Evaluation, Infrastructure, Inference, Agents, RAG, Multimodal, Prompt Engineering, MLOps, Observability | Full Lifecycle ✅ |
**Recent Progress**: +4 skills (Lambda Labs, SAM, BLIP-2, AudioCraft) completing the 70-skill roadmap with GPU cloud and extended multimodal capabilities
**Recent Progress**: +4 skills (TransformerLens, SAELens, pyvene, nnsight) adding Mechanistic Interpretability category for reverse-engineering neural networks
**Philosophy**: Quality > Quantity. Following [Anthropic official best practices](anthropic_official_docs/best_practices.md) - each skill provides 200-500 lines of focused, actionable guidance with progressive disclosure.
@@ -300,6 +307,7 @@ claude-ai-research-skills/
├── 01-model-architecture/ (5 skills ✓ - Megatron, LitGPT, Mamba, RWKV, NanoGPT)
├── 02-tokenization/ (2 skills ✓ - HuggingFace Tokenizers, SentencePiece)
├── 03-fine-tuning/ (4 skills ✓ - Axolotl, LLaMA-Factory, Unsloth, PEFT)
├── 04-mechanistic-interpretability/ (4 skills ✓ - TransformerLens, SAELens, pyvene, nnsight)
├── 05-data-processing/ (2 skills ✓ - Ray Data, NeMo Curator)
├── 06-post-training/ (4 skills ✓ - TRL, GRPO, OpenRLHF, SimPO)
├── 07-safety-alignment/ (3 skills ✓ - Constitutional AI, LlamaGuard, NeMo Guardrails)
@@ -367,6 +375,19 @@ All contributors are featured in our [Contributors Hall of Fame](CONTRIBUTORS.md
## Recent Updates
<details>
<summary><b>December 2025 - v0.11.0 🔬 Mechanistic Interpretability</b></summary>
- 🔬 **NEW CATEGORY**: Mechanistic Interpretability (4 skills)
- 🔍 TransformerLens skill: Neel Nanda's library for mech interp with HookPoints, activation caching, circuit analysis
- 🧠 SAELens skill: Sparse Autoencoder training and analysis for feature discovery, monosemanticity research
- ⚡ pyvene skill: Stanford's causal intervention library with declarative configs, DAS, activation patching
- 🌐 nnsight skill: Remote interpretability via NDIF, run experiments on 70B+ models without local GPUs
- 📝 ~6,500 new lines of documentation across 16 files
- **74 total skills** (filling the missing 04 category slot)
</details>
<details>
<summary><b>November 25, 2025 - v0.10.0 🎉 70 Skills Complete!</b></summary>