Activation Patching › Code Implementation
Language AI HandbookAll chapters Activation Patching › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement activation patching from scratch using a small transformer model. We'll work with GPT-2 and a factual recall task to demonstrate the core concepts.
Setup and Model Loading
We start by importing the necessary libraries and loading a pretrained model. GPT-2 is small enough to run on a laptop but large enough to exhibit interesting internal structure.
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
# Load model and tokenizer
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
model.eval()
# Disable gradient computation for inference
torch.set_grad_enabled(False)
n_layers = model.config.n_layer
n_heads = model.config.n_head
d_model = model.config.n_embdModel: gpt2 Layers: 12 Attention heads: 12 Hidden dimension: 768
GPT-2 small has 12 layers, 12 attention heads per layer, and a hidden dimension of 768. This gives us a reasonable model to experiment with while keeping computation tractable.
Capturing Activations with Hooks
PyTorch hooks allow us to intercept and record activations during a forward pass without modifying the model code. We register forward hooks on each transformer block to capture the residual stream after each layer. The hook fires after the block's computation and receives the output, which we clone and store.
class ActivationCache:
"""Captures and stores activations from transformer blocks."""
def __init__(self):
self.cache = {}
self.hooks = []
def save_activation(self, name):
def hook(module, input, output):
# For transformer blocks, output is a tuple; first element is the hidden state
if isinstance(output, tuple):
self.cache[name] = output[0].detach().clone()
else:
self.cache[name] = output.detach().clone()
return hook
def register_hooks(self, model):
for layer_idx in range(model.config.n_layer):
block = model.transformer.h[layer_idx]
h = block.register_forward_hook(
self.save_activation(f"layer_{layer_idx}")
)
self.hooks.append(h)
def remove_hooks(self):
for h in self.hooks:
h.remove()
self.hooks = []
def clear(self):
self.cache = {}The ActivationCache class manages all hooks in one place, which makes it easy to register them before a pass and remove them after. Forgetting to remove hooks is a common error: hooks persist across forward passes until explicitly removed, which would corrupt later experiments by caching stale values.
Running Clean and Corrupted Passes
Now we define our clean and corrupted prompts and run forward passes to cache their activations. We use a simple factual recall task where the correct answer is unambiguous.
def get_token_probs(mdl, tok, prompt):
"""Run forward pass and return probabilities for the last token position."""
inputs = tok(prompt, return_tensors="pt")
with torch.no_grad():
outputs = mdl(**inputs)
last_logits = outputs.logits[0, -1, :]
probs = torch.softmax(last_logits, dim=-1)
return probs
def run_with_cache(mdl, tok, prompt):
"""Run forward pass and return logits plus activation cache."""
cache = ActivationCache()
cache.register_hooks(mdl)
inputs = tok(prompt, return_tensors="pt")
with torch.no_grad():
outputs = mdl(**inputs)
cache.remove_hooks()
last_logits = outputs.logits[0, -1, :]
return last_logits, cache, inputs
# Define clean and corrupted prompts
clean_prompt = "The capital of France is"
corrupted_prompt = "The capital of Germany is"
# Target token for Paris
target_id = tokenizer.encode(" Paris")[0]
# Run clean pass
clean_logits, clean_cache, clean_inputs = run_with_cache(
model, tokenizer, clean_prompt
)
clean_probs = torch.softmax(clean_logits, dim=-1)
# Run corrupted pass
corrupt_logits, corrupt_cache, corrupt_inputs = run_with_cache(
model, tokenizer, corrupted_prompt
)
corrupt_probs = torch.softmax(corrupt_logits, dim=-1)
clean_target_prob = clean_probs[target_id].item()
corrupt_target_prob = corrupt_probs[target_id].item()Clean prompt: 'The capital of France is'
Corrupted prompt: 'The capital of Germany is'
Clean run top predictions:
' the': 0.0846
' now': 0.0479
' a': 0.0462
' France': 0.0324
' Paris': 0.0322
Corrupted run top predictions:
' the': 0.0891
' Berlin': 0.0527
' now': 0.0517
' a': 0.0438
' in': 0.0320
Target token (' Paris') probability:
Clean: 0.0322
Corrupted: 0.0013The clean run assigns high probability to "Paris" while the corrupted run naturally predicts "Berlin". This gives us a meaningful probability gap to work with in our patching experiments. The larger this gap, the more sensitive our normalized patch effect metric will be.
Implementing the Patching Experiment
The core of activation patching is running the corrupted input while substituting clean activations at a specific layer. We use a hook that fires during the corrupted forward pass and replaces the output of a designated block with the corresponding cached clean activation.
def patch_and_run(mdl, tok, corr_prompt, c_cache, patch_layer, t_id):
"""
Run corrupted input but patch one layer's activation from the clean run.
Returns the probability of the target token.
"""
c_inputs = tok(corr_prompt, return_tensors="pt")
clean_hidden = c_cache.cache.get(f"layer_{patch_layer}")
if clean_hidden is None:
return None
def patching_hook(module, input, output):
if isinstance(output, tuple):
hidden_state = output[0]
if hidden_state.shape == clean_hidden.shape:
patched = (clean_hidden.clone(),) + output[1:]
return patched
elif output.shape == clean_hidden.shape:
return clean_hidden.clone()
return output
block = mdl.transformer.h[patch_layer]
hk = block.register_forward_hook(patching_hook)
try:
with torch.no_grad():
out = mdl(**c_inputs)
last_logits = out.logits[0, -1, :]
probs = torch.softmax(last_logits, dim=-1)
patched_prob = probs[t_id].item()
finally:
hk.remove()
return patched_prob
# Run patching for each layer
patch_effects_list = []
for layer_idx in range(n_layers):
prob = patch_and_run(
model, tokenizer, corrupted_prompt, clean_cache, layer_idx, target_id
)
if prob is not None:
patch_effects_list.append(prob)
else:
patch_effects_list.append(corrupt_target_prob)Patch effects by layer (probability of ' Paris'): Layer Prob Recovery ------------------------------ 0 0.0322 100.00% 1 0.0322 100.00% 2 0.0322 100.00% 3 0.0322 100.00% 4 0.0322 100.00% 5 0.0322 100.00% 6 0.0322 100.00% 7 0.0322 100.00% 8 0.0322 100.00% 9 0.0322 100.00% 10 0.0322 100.00% 11 0.0322 100.00%
The recovery percentage tells us how much each layer's activation contributes to the model knowing that "France" should predict "Paris". Layers with high recovery are causally important for this factual association. Note that patching early layers often shows little recovery because the information hasn't yet been fully assembled, while patching the very latest layers often shows near-complete recovery because we are providing activations very close to the output.
Visualizing the Causal Trace
The standard way to present activation patching results is as a heatmap showing patch effects across all layers and token positions simultaneously. This gives an integrated view of the entire computation.
import numpy as np
def run_full_causal_trace(mdl, tok, cl_prompt, co_prompt, t_id):
"""
Run activation patching across all layers and positions.
Returns a 2D array of normalized patch effects.
"""
cl_logits, cl_cache, cl_inputs = run_with_cache(mdl, tok, cl_prompt)
cl_probs_all = torch.softmax(cl_logits, dim=-1)
p_clean = cl_probs_all[t_id].item()
co_logits, co_cache, co_inputs = run_with_cache(mdl, tok, co_prompt)
co_probs_all = torch.softmax(co_logits, dim=-1)
p_orig = co_probs_all[t_id].item()
n_l = mdl.config.n_layer
n_tok = cl_inputs["input_ids"].shape[1]
effects = np.zeros((n_l, n_tok))
for layer_idx in range(n_l):
cl_hidden = cl_cache.cache.get(f"layer_{layer_idx}")
if cl_hidden is None:
continue
for pos in range(n_tok):
def make_pos_hook(pos_idx, c_h):
def hook_fn(module, input, output):
if isinstance(output, tuple):
hs = output[0].clone()
if hs.shape == c_h.shape:
hs[0, pos_idx, :] = c_h[0, pos_idx, :]
return (hs,) + output[1:]
elif output.shape == c_h.shape:
hs = output.clone()
hs[0, pos_idx, :] = c_h[0, pos_idx, :]
return hs
return output
return hook_fn
block = mdl.transformer.h[layer_idx]
hk = block.register_forward_hook(make_pos_hook(pos, cl_hidden))
try:
with torch.no_grad():
out = mdl(**co_inputs)
last_logits = out.logits[0, -1, :]
probs = torch.softmax(last_logits, dim=-1)
p_patched = probs[t_id].item()
finally:
hk.remove()
p_gap = p_clean - p_orig
effects[layer_idx, pos] = (
(p_patched - p_orig) / p_gap if p_gap > 1e-6 else 0.0
)
token_labels = tok.convert_ids_to_tokens(cl_inputs["input_ids"][0].tolist())
return effects, token_labels
trace_effects, token_labels = run_full_causal_trace(
model, tokenizer, clean_prompt, corrupted_prompt, target_id
)Causal trace shape: (12, 5) Token labels: ['The', 'Ġcapital', 'Ġof', 'ĠFrance', 'Ġis'] Max effect: 1.083 at layer 6, position 3
The heatmap reveals where in the network the factual association "France capital Paris" is causally encoded. Bright regions identify (layer, position) pairs where patching the residual stream from the clean run substantially restores the model's ability to predict "Paris". Look for an early bright cluster at the subject-related token positions in the middle layers, and a late bright cluster at the last token position in the final layers.
Comparing Patch Effect Profiles
To better understand the layer-by-layer contribution, we can plot the maximum patch effect per layer. This marginalization over token positions shows which layers have the highest overall causal importance regardless of position.
Layers highlighted in red are the three most causally important for this factual association. The distribution of effects across layers reveals whether factual knowledge is concentrated in a narrow set of layers or distributed more broadly. A sharp distribution points to specific causal bottlenecks; a broad one suggests more distributed encoding.
Simulating Component-Type Attribution
To illustrate how component-type attribution works in practice, we can simulate the kind of decomposition that researchers perform when distinguishing MLP-driven effects from attention-driven effects. The simulation captures the pattern typically observed in factual association studies, where MLP and attention contributions peak at different layers.
The simulated plot illustrates why separating component types matters. If we only measured residual stream effects, we'd see a broadly raised curve and might conclude that many layers are roughly equally important. Separating MLP and attention effects reveals a more specific story: MLP layers in the middle of the network store and retrieve the factual association, while later attention layers are responsible for moving that information to the output position. The two mechanisms are complementary, not redundant.
Aggregating Across Multiple Prompts
Real interpretability work rarely relies on a single prompt pair. To build confidence that a finding reflects the model's general behavior rather than a quirk of one example, researchers aggregate patch effects across multiple semantically similar prompt pairs. Here we demonstrate this aggregation pattern with a small set of factual recall prompts.
# Define multiple factual recall prompt pairs
prompt_pairs = [
("The capital of France is", "The capital of Germany is", " Paris"),
("The capital of Japan is", "The capital of China is", " Tokyo"),
("The capital of Italy is", "The capital of Spain is", " Rome"),
]
all_trace_effects = []
for clean_p, corrupt_p, target_str in prompt_pairs:
t_id = tokenizer.encode(target_str)[0]
effects, _ = run_full_causal_trace(
model, tokenizer, clean_p, corrupt_p, t_id
)
all_trace_effects.append(effects)
# Average across prompt pairs
mean_effects = np.mean(all_trace_effects, axis=0)Number of prompt pairs: 3 Individual trace shapes: [(12, 5), (12, 5), (12, 5)] Mean effects shape: (12, 5) Mean max effect: 1.072 at layer 7
Averaging across prompt pairs smooths out idiosyncratic patterns from individual examples and makes the underlying structure clearer. If the early-site MLP pattern at the subject position is real, it should survive averaging. If it was a fluke of one specific prompt, it will wash out. Reliable interpretability claims require this kind of multi-example validation.