mirror of
https://github.com/Orchestra-Research/AI-Research-SKILLs.git
synced 2026-09-19 05:04:39 +08:00
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:
@@ -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)
|
||||
@@ -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>
|
||||
|
||||
|
||||
Reference in New Issue
Block a user