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 (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 pip install transformer-lens

For development version:

pip install git+https://github.com/TransformerLensOrg/TransformerLens

Core Concepts HookedTransformer

The main class that wraps transformer models with HookPoints on every activation:

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:

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 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 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.

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

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

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

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

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:

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 Overview and quick start guide references/api.md Complete API reference for HookedTransformer, ActivationCache, HookPoints references/tutorials.md Step-by-step tutorials for activation patching, circuit analysis, logit lens External Resources Tutorials Main Demo Notebook Activation Patching Demo ARENA Mech Interp Course - 200+ hours of tutorials Papers A Mathematical Framework for Transformer Circuits In-context Learning and Induction Heads Interpretability in the Wild (IOI) Official Documentation Official Docs Model Properties Table Neel Nanda's Glossary Version Notes v2.0: Removed HookedSAE (moved to SAELens) v3.0 (alpha): TransformerBridge for loading any nn.Module