Activation Patching¶
Activation patching copies an activation from one forward pass into another and measures what the substitution does to the output. It answers a question that reading activations cannot: whether a component causes a behavior, or merely correlates with it.
This tutorial works through the indirect object identification (IOI) task on GPT-2 small, building a (layer x token) map of where the answer is carried, and then runs the same experiment on Llama-3.1-8B through NDIF.
Setup¶
try:
import google.colab
is_colab = True
except ImportError:
is_colab = False
if is_colab:
!pip install -U nnsight
from IPython.display import clear_output
import torch
import nnsight
from nnsight import TransformersModel
import plotly.express as px
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
# Load gpt2
model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
clear_output()
print(model)
GPT2LMHeadModel(
(transformer): GPT2Model(
(wte): Embedding(50257, 768)
(wpe): Embedding(1024, 768)
(drop): Dropout(p=0.1, inplace=False)
(h): ModuleList(
(0-11): 12 x GPT2Block(
(ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
(attn): GPT2Attention(
(c_attn): Conv1D()
(c_proj): Conv1D()
(attn_dropout): Dropout(p=0.1, inplace=False)
(resid_dropout): Dropout(p=0.1, inplace=False)
)
(ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
(mlp): GPT2MLP(
(c_fc): Conv1D()
(c_proj): Conv1D()
(act): NewGELUActivation()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
(ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
)
(lm_head): Linear(in_features=768, out_features=50257, bias=False)
(generator): Generator(
(streamer): Streamer()
)
)
Next up, we define our clean prompt and our corrupted prompt. As prompts may be associated with many different feature circuits (i.e., circuits responsible for IOI, deciding if the language is English, or prompt refusal), choosing a counterfactual prompt with only changes directly related your feature of interest is essential.
Here, we switch the name of the repeated subject, thus swapping out the indirect object for our IOI task:
clean_prompt = "After John and Mary went to the store, Mary gave a bottle of milk to"
corrupted_prompt = "After John and Mary went to the store, John gave a bottle of milk to"
We then use the tokenizer on the two words of interest (“John” and “Mary”) to find the token that represents them. That way we can grab the prediction for these two tokens and compare. Because our prompts don't end in a space, make sure to add a space before each word (i.e., the combined space + word token is what we're looking for).
correct_index = model.tokenizer(" John")["input_ids"][0] # includes a space
incorrect_index = model.tokenizer(" Mary")["input_ids"][0] # includes a space
print(f"' John': {correct_index}")
print(f"' Mary': {incorrect_index}")
' John': 1757 ' Mary': 5335
Patching Experiment¶
Now we can run the actual patching intervention! What does this even mean?
We now have two prompts, a "clean" one and a "corrupted" one. Intuitively, the model output for each of these prompts should be different: we'd expect the model to answer "John" for the clean prompt and "Mary" for the corrupted prompt.
In this experiment, we run the model with the clean prompt as an input and then get each layer's output value (i.e., residual stream) and calculate the logit difference between the correct and incorrect answers for this run. Next, we calculate the logit difference between the correct and incorrect answers for the corrupted prompt. Finally, we patch in the residual stream from the clean prompt into the corrupted prompt and collect the logit difference.
Typically, you would do each of these steps in different forward passes, but with NNsight we can batch our operations and do these three runs in one forward pass! This saves us a lot of time.
Step 1: Clean Run
First, we'll run the model with the clean prompt:
"After John and Mary went to the store, Mary gave a bottle of milk to"
During this clean run, we collect the final output of each layer. We also record the logit difference in the final model output between the correct answer token " John" and the incorrect token " Mary".
Step 1 code:
We will be using NNsight's batching functionality to run this experiment, but if you were to do it in separate forward passes (less optimal code!), here is what the clean forward pass would look like.N_LAYERS = len(model.transformer.h)
clean_hs = []
# Clean run
with model.trace() as tracer:
with tracer.invoke(clean_prompt) as invoker:
clean_tokens = invoker.inputs[1]['input_ids'][0].save()
# Get hidden states of all layers in the network.
# We index the output at 0 because it's a tuple where the first index is the hidden state.
for layer_idx in range(N_LAYERS):
clean_hs.append(model.transformer.h[layer_idx].output[0].save())
# Get logits from the lm_head.
clean_logits = model.lm_head.output
# Calculate the difference between the correct answer and incorrect answer for the clean run and save it.
clean_logit_diff = (
clean_logits[0, -1, correct_index] - clean_logits[0, -1, incorrect_index]
).save()
Step 2: Corrupted Run
Next, we run the model using the corrupted input prompt:
"After John and Mary went to the store, John gave a bottle of milk to"
During this corrupted run, we collect the logit difference in the final model output between the correct and incorrect answer tokens
Note: because we are testing changes induced by the corrupted prompt, the target answers remain the same as in the clean run. That is, the correct token is still " John" and the incorrect token is still " Mary".
Step 2 code:
We will be using NNsight's batching functionality to run this experiment, but if you were to do activation patching in three separate forward passes, here is what the corrupted forward pass would look like.# Corrupted run
with model.trace(corrupted_prompt) as tracer:
corrupted_logits = model.lm_head.output
# Calculate the difference between the correct answer and incorrect answer for the corrupted run and save it.
corrupted_logit_diff = (
corrupted_logits[0, -1, correct_index]
- corrupted_logits[0, -1, incorrect_index]
).save()
Step 3: Activation Patching Intervention
Finally, we perform our activation patching procedure. For each token position in the clean prompt, we loop through all layers of the model. Within each layer, we run a forward pass using the corrupted prompt, and patch in the corresponding activation from our clean run at the given token position. We then collect the final output difference between the correct and incorrect answer tokens for each patched activation.
Step 3 code:
We will be using NNsight's batching functionality to run this experiment, but if you were to do it in separate forward passes (less optimal code!), here is what the patching portion would look like. Note that this code will run many forward passes, one for each layer and token position combination within the nested for loops. This is why we recommend optimizing your code using invoke batching.
# Activation Patching Intervention
ioi_patching_results = []
# Iterate through all the layers
for layer_idx in range(len(model.transformer.h)):
_ioi_patching_results = []
# Iterate through all tokens
for token_idx in range(len(clean_tokens)):
# Patching corrupted run at given layer and token
with model.trace(corrupted_prompt) as tracer:
# Apply the patch from the clean hidden states to the corrupted hidden states.
model.transformer.h[layer_idx].output[0][:, token_idx] = clean_hs[layer_idx][:,token_idx,:]
patched_logits = model.lm_head.output
patched_logit_diff = (
patched_logits[0, -1, correct_index]
- patched_logits[0, -1, incorrect_index]
)
# Calculate the improvement in the correct token after patching.
patched_result = (patched_logit_diff - corrupted_logit_diff) / (
clean_logit_diff - corrupted_logit_diff
)
_ioi_patching_results.append(patched_result.item())
_ioi_patching_results.save()
ioi_patching_results.append(_ioi_patching_results)
Now that we understand each of the steps in the activation patching workflow, let's try implementing the whole experiment in one forward pass by breaking up each of our inputs into multiple invocation calls and batching them.
tokenized_prompt = model.tokenizer(clean_prompt)["input_ids"]
print(model.tokenizer.decode(tokenized_prompt))
After John and Mary went to the store, Mary gave a bottle of milk to
To batch our operations we use an empty model.trace() together with several tracer.invoke(prompt) calls. Each invoke adds a separate row to our batch, so we can run all necessary computations in one forward pass. Because the patching invokes need a clean-run activation that is only produced partway through the forward pass, we synchronize the invokes with tracer.barrier(n): a barrier makes n invokes wait at the same point so the clean activation is ready before the patched runs read it. For a gentler introduction to running several prompts in one trace, see the Batching feature guide.
So let's define a function that does that for us.
def gpt2_activation_patching(clean_prompt, corrupted_prompt, N_LAYERS):
ioi_patching_results = []
with torch.no_grad():
with model.trace() as tracer:
#STEP 0: Define barriers for each layer and token
barriers_per_layer = [tracer.barrier(len(tokenized_prompt)+1) for _ in range(N_LAYERS)]
# STEP 1: Clean run
with tracer.invoke(clean_prompt) as invoker:
# At each layer save the output of the current layer.
for layer_idx in range(N_LAYERS):
hidden_state = model.transformer.h[layer_idx].output
# Call barrier so nnsight can prepare the hidden states for the patching invoke.
# Now, all other invoker will complete their code until the same barrier is called.
barriers_per_layer[layer_idx]()
# Get logits from the lm_head and calculate the difference between the correct answer and incorrect answer for the clean run and save it.
clean_logits = model.lm_head.output
clean_logit_diff = (
clean_logits[0, -1, correct_index] - clean_logits[0, -1, incorrect_index]
).save()
# STEP 2: Corrupt run
with tracer.invoke(corrupted_prompt) as invoker:
corrupted_logits = model.lm_head.output
# Calculate the difference between the correct answer and incorrect answer for the corrupted run and save it.
corrupted_logit_diff = (
corrupted_logits[0, -1, correct_index]
- corrupted_logits[0, -1, incorrect_index]
).save()
# STEP 3: Activation Patching Intervention – across all layers & token positions
for layer_idx in range(len(model.transformer.h)):
local_ioi_patching_results = []
for token_idx in range(len(tokenized_prompt)):
# Call a seperate invoker for each layer and token postion
with tracer.invoke(corrupted_prompt) as invoker:
# Call the barrier so nnsight can grab the hidden states for this invoke
barriers_per_layer[layer_idx]()
# Patch (replace) the clean hidden states over the corrupted hidden states.
model.transformer.h[layer_idx].output[:, token_idx, :] = hidden_state[:, token_idx, :]
# Get logits from the lm_head and calculate the difference between the correct answer and incorrect answer for this invoker.
patched_logits = model.lm_head.output
patched_logit_diff = (
patched_logits[0, -1, correct_index]
- patched_logits[0, -1, incorrect_index]
)
# Calculate the improvement in the correct token after patching and append it to the local result list.
patched_result = (patched_logit_diff - corrupted_logit_diff) / (
clean_logit_diff - corrupted_logit_diff
)
local_ioi_patching_results.append(patched_result.item())
# Append this layer's per-token results to the global list.
ioi_patching_results.append(local_ioi_patching_results)
# Save the whole results container so it survives the trace. For plain
# Python objects (lists, dicts, ints) prefer `nnsight.save(...)` over
# `.save()`; save the container, not each element inside it.
ioi_patching_results = nnsight.save(ioi_patching_results)
return clean_logit_diff, corrupted_logit_diff, ioi_patching_results
Now lets run our function to obtain the activation patching results
N_LAYERS = len(model.transformer.h)
clean_logit_diff, corrupted_logit_diff, ioi_patching_results = gpt2_activation_patching(clean_prompt, corrupted_prompt, N_LAYERS)
Visualize Results¶
Let's define a function to plot our activation patching results.
def plot_ioi_patching_results(ioi_patching_results,
x_labels,
plot_title="Normalized Logit Difference After Patching Residual Stream on the IOI Task"):
fig = px.imshow(
ioi_patching_results,
color_continuous_midpoint=0.0,
color_continuous_scale="RdBu",
labels={"x": "Position", "y": "Layer","color":"Norm. Logit Diff"},
x=x_labels,
title=plot_title,
)
return fig
Let's see how the patching intervention changes the logit difference! Let's use a heatmap to examine how the logit difference changes after patching each layer's output across token positions.
print(f"Clean logit difference: {clean_logit_diff:.3f}")
print(f"Corrupted logit difference: {corrupted_logit_diff:.3f}")
clean_decoded_tokens = [model.tokenizer.decode(token) for token in tokenized_prompt]
token_labels = [f"{token}_{index}" for index, token in enumerate(clean_decoded_tokens)]
fig = plot_ioi_patching_results(ioi_patching_results,token_labels,"Patching GPT-2-small Residual Stream on IOI task")
fig.show()
Clean logit difference: 4.124 Corrupted logit difference: -2.272
In the above plot, we see that patching the clean residual stream into the corrupted model does not change much in the final token difference for input tokens 0-8. This is expected, as there is no difference in the clean vs. corrupted prompt for these tokens, so patching in the clean activations at this point shouldn't change the model prediction.
However, when we get to token #9, " Mary", where the subject is introduced for the second time, there is a sharp increase in output logit difference, indicating that the patch changes how the model predicts the outcome downstream, particularly for the earlier layers. Towards the middle layers, the logit differences are decreasing. We are thus seeing how the network is tracking information about the indirect object as the layers progress.
A similar but opposite effect is observed when the activations for the final prompt token are patched: the normalized logit difference increases after a transition period in the middle layers.
Limitations¶
Although activation patching is an effective technique for circuit localization, it requires running a forward pass through the model for every patch, making it computationally expensive.
Attribution patching is an approximation of activation patching that helps scale the technique to larger experiments and models. See our Attribution Patching tutorial to try it out!
Trying on a bigger model¶
Although the original IOI experiment was performed on GPT-2 small, NDIF lets researchers explore the same problems on large-scale models that they can't host themselves!
Let's see how the activations of Llama-3.1-8B contribute to the IOI task using activation patching over NDIF's remote infrastructure. The only change from the code above is that the model runs on NDIF's servers instead of your machine - the tracing API is identical.
To make it a bit more interesting, we will implement attention-head activation patching and make a few changes to our previous code to run it more efficiently.
Note: The cells in this section run against NDIF's public deployment of
meta-llama/Llama-3.1-8B. They require an NDIF API key and access to that model, and will only run where such a deployment is reachable - the shown outputs were produced against NDIF. See the Remote Execution feature guide for setup.
NNsight Remote Setup¶
Make sure you have obtained your NDIF API key and configured your workspace for remote execution. You should also create a Hugging Face account, create an access token, and request access to the Llama 3.1 8B model.
from nnsight import CONFIG
if is_colab:
# include your HuggingFace Token and NNsight API key on Colab secrets
from google.colab import userdata
NDIF_API = userdata.get('NDIF_API')
HF_TOKEN = userdata.get('HF_TOKEN')
CONFIG.set_default_api_key(NDIF_API)
!huggingface-cli login -token HF_TOKEN
clear_output()
Let's load the llama 3.1 8B model.
Let's define some IOI prompts. Each of these prompts can be used as a 'clean' and as a 'corrupted' prompt, as each prompt has a related corrupted version with the IO switched out. I included prompts of different length so we have to deal with padding.
prompts = [
"After John and Mary went to the store, John gave a bottle of milk to",
"After Mary and John went to the store, Mary gave a bottle of milk to",
"When Lisa and Sarah went to the cinema, Lisa gave the ticket to",
"When Lisa and Sarah went to the cinema, Sarah gave the ticket to"
]
# Load model. For remote execution the weights live on NDIF, so no device_map is needed.
model = TransformersModel("meta-llama/Llama-3.1-8B")
print(model)
LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(128256, 4096)
(layers): ModuleList(
(0-31): 32 x LlamaDecoderLayer(
(self_attn): LlamaAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=1024, bias=False)
(v_proj): Linear(in_features=4096, out_features=1024, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
)
(mlp): LlamaMLP(
(gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
(up_proj): Linear(in_features=4096, out_features=14336, bias=False)
(down_proj): Linear(in_features=14336, out_features=4096, bias=False)
(act_fn): SiLUActivation()
)
(input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
(post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
)
)
(norm): LlamaRMSNorm((4096,), eps=1e-05)
(rotary_emb): LlamaRotaryEmbedding()
)
(lm_head): Linear(in_features=4096, out_features=128256, bias=False)
(generator): Generator(
(streamer): Streamer()
)
)
Let's now have a look at the tokenization of the prompts. As you can see LLama 3.1 includes the <|begin_of_text|> (bos) special token at the start of each prompt. Also, the two shorter prompts include padding tokens. For our activation patching experiment we want to exclude these position.
tokenized_prompts = model.tokenizer(prompts, padding=True)["input_ids"]
for i in range(len(tokenized_prompts)):
print(model.tokenizer.convert_ids_to_tokens(tokenized_prompts[i]))
['<|begin_of_text|>', 'After', 'ĠJohn', 'Ġand', 'ĠMary', 'Ġwent', 'Ġto', 'Ġthe', 'Ġstore', ',', 'ĠJohn', 'Ġgave', 'Ġa', 'Ġbottle', 'Ġof', 'Ġmilk', 'Ġto'] ['<|begin_of_text|>', 'After', 'ĠMary', 'Ġand', 'ĠJohn', 'Ġwent', 'Ġto', 'Ġthe', 'Ġstore', ',', 'ĠMary', 'Ġgave', 'Ġa', 'Ġbottle', 'Ġof', 'Ġmilk', 'Ġto'] ['<|end_of_text|>', '<|end_of_text|>', '<|begin_of_text|>', 'When', 'ĠLisa', 'Ġand', 'ĠSarah', 'Ġwent', 'Ġto', 'Ġthe', 'Ġcinema', ',', 'ĠLisa', 'Ġgave', 'Ġthe', 'Ġticket', 'Ġto'] ['<|end_of_text|>', '<|end_of_text|>', '<|begin_of_text|>', 'When', 'ĠLisa', 'Ġand', 'ĠSarah', 'Ġwent', 'Ġto', 'Ġthe', 'Ġcinema', ',', 'ĠSarah', 'Ġgave', 'Ġthe', 'Ġticket', 'Ġto']
Define the answers to these prompts, formatted as (correct, incorrect)
answers = [
(" Mary", " John"),
(" John", " Mary"),
(" Sarah", " Lisa"),
(" Lisa", " Sarah")
]
From each example in our prompts we will now creat a clean and a corrupted version. We further use the special_token_mask to create a mask to later exclude padding and bos tokens from the experiment
To streamline the process, we create a function.
def prepare_patching_examples(model, prompts, answers):
# Tokenize prompts and crete the patching mask
clean_inputs = model.tokenizer(prompts, return_tensors="pt", padding=True, return_special_tokens_mask=True)
clean_tokens = clean_inputs['input_ids']
attention_maks = clean_inputs['attention_mask']
patching_mask = ~clean_inputs['special_tokens_mask'].bool()
# To get the corrupted inputs we just switch the postions of the prompt of all pairs.
corrupted_tokens = clean_tokens[
[(i + 1 if i % 2 == 0 else i - 1) for i in range(len(prompts))]
]
# Tokenize answers for each prompt:
answer_token_indices = torch.tensor([
[model.tokenizer(answers[i][j])["input_ids"][1] for j in range(2)]
for i in range(len(answers))
])
return clean_tokens, corrupted_tokens, attention_maks, patching_mask, answer_token_indices
clean_tokens, corrupted_tokens, attention_maks,\
patching_mask, answer_token_indices = prepare_patching_examples(model, prompts, answers)
print("clean_tokens.shape = ", clean_tokens.shape)
print("corrupted_tokens.shape = ", corrupted_tokens.shape)
print("patching_mask.shape = ", patching_mask.shape)
print("answer_tokens = " , answer_token_indices)
clean_tokens.shape = torch.Size([4, 17])
corrupted_tokens.shape = torch.Size([4, 17])
patching_mask.shape = torch.Size([4, 17])
answer_tokens = tensor([[10455, 3842],
[ 3842, 10455],
[21077, 29656],
[29656, 21077]])
Patching Attention Heads¶
The residual stream isn't the only model component you can apply activation patching on: let's try patching Llama's attention heads to see how they influence the IOI task! Here, we apply our patching intervention on the attention output, o_proj.output in Llama models.
Because the multihead attention of Llama models are stored in a projection matrix containing all attention heads, we will need to resize the tensor to reveal individual attention head contributions. The einops library is a handy way to resize tensors.
import einops
Okay, now let's apply our three activation patching steps to our attention heads during an IOI task. Because Llama-8B is a bigger model, we're going to do this step inside a session in multiple different traces so that we don't run into memory issues. In total we have to run N_HEADS x N_LAYERS patching interventions. For LLama 3.1 8B that are 32 x 32 = 1024 interventions. We can further improve our result by skipping layers previous to our interventions by caching the original corrupted hidden states from previous runs.
def activation_patching(model, prompts, answers):
N_LAYERS = model.config.num_hidden_layers
N_HEADS = model.config.num_attention_heads
batch_size = len(prompts)
# Prepare patching inputs with our previously defined function
clean_tokens, corrupted_tokens, attention_maks,\
patching_mask, answer_token_indices = prepare_patching_examples(model, prompts, answers)
with torch.no_grad():
# With setting remote=True, we send all code in the session to ndif's servers.
with model.session(remote=True) as session:
ioi_patching_results_all = [].save()
# STEP 1: Clean run, grab clean activations for each attention head
z_hs = {}
with model.trace(**{"input_ids": clean_tokens, "attention_mask": attention_maks}) as tracer:
for layer_idx, layer in enumerate(model.model.layers):
# attention output for llama models needs to be reshaped to look at individual heads
z = layer.self_attn.o_proj.input # dimensions [batch x seq x D_MODEL]
z_reshaped = einops.rearrange(z, 'b s (nh dh) -> b s nh dh',nh=N_HEADS) # dimensions [batch x seq x N_HEADS x D_HEADS]
for head_idx in range(N_HEADS):
z_hs[layer_idx,head_idx] = z_reshaped[:,:,head_idx,:]
# Get logits from the lm_head and calcute the logit diff.
clean_logits = model.lm_head.output
clean_logit_diff = (
clean_logits[range(batch_size), -1, answer_token_indices[:, 0]] - clean_logits[range(batch_size), -1, answer_token_indices[:, 1]]
).save()
# STEP 2: Corrupted run, grab corrupted logits for later comparison.
corrupted_hs = {}
with model.trace(**{"input_ids": corrupted_tokens, "attention_mask": attention_maks}) as tracer:
# cache layer outputs so we can skip pre-intervention layers in step 3.
for layer_idx, layer in enumerate(model.model.layers):
corrupted_hs[layer_idx] = layer.output
# Get logits from the lm_head and calcute the logit diff.
corrupted_logits = model.lm_head.output
corrupted_logit_diff = (
corrupted_logits[range(batch_size), -1, answer_token_indices[:, 0]] - corrupted_logits[range(batch_size), -1, answer_token_indices[:, 1]]
).save()
# STEP 3: Patching runs, apply 'clean' model state at each layer and head,
for layer_idx in range(N_LAYERS):
local_ioi_patching_results = []
# Create a new tracer for each layer.
with model.trace() as tracer:
for head_idx in range(N_HEADS):
# Add a seperate batch dimension for each attention head.
with tracer.invoke(**{"input_ids": corrupted_tokens, "attention_mask": attention_maks}) as invoker:
# Layers preceeding our intervention can be skipped by incerting the cached hidden states from step 2.
for skip_layer_idx, layer in enumerate(model.model.layers[:layer_idx]):
layer.skip(corrupted_hs[skip_layer_idx])
# Apply the patch from the clean hidden states to the corrupted hidden state for given layer and head.
z = model.model.layers[layer_idx].self_attn.o_proj.input # dimensions [batch x seq x D_MODEL]
z_patched = einops.rearrange(z, 'b s (nh dh) -> b s nh dh',nh=N_HEADS)
# Patch in clean state. We apply the patching_mask to excude special tokens form our experiment.
z_patched[:,:,head_idx,:][patching_mask] = z_hs[layer_idx,head_idx][patching_mask]
z_patched = einops.rearrange(z_patched, 'b s nh dh -> b s (nh dh)') # reshape
model.model.layers[layer_idx].self_attn.o_proj.input = z_patched # apply to model
# Get logits from the lm_head and calcute the logit diff.
patched_logits = model.lm_head.output
patched_logit_diff = (
patched_logits[range(batch_size), -1, answer_token_indices[:, 0]] - patched_logits[range(batch_size), -1, answer_token_indices[:, 1]]
)
# Calculate the improvement in the correct token after patching.
patched_result = (patched_logit_diff - corrupted_logit_diff) / (
clean_logit_diff - corrupted_logit_diff
)
local_ioi_patching_results.append(patched_result.mean().item())
ioi_patching_results_all.append(local_ioi_patching_results)
return clean_logit_diff.mean().item(), corrupted_logit_diff.mean().item(), ioi_patching_results_all
Now lets run our function.
clean_logit_diff, corrupted_logit_diff, ioi_patching_results_all = activation_patching(model, prompts, answers)
✓ [1f2dfec1624f4a27a409f1bd9cbdbb47] COMPLETED (7.2s) Your job has been completed.
⬇ Downloading result: 0%| | 0.00/2.57k [00:00<?]
⬇ Downloading result: 100%|██████████| 2.57k/2.57k [00:00<00:00]
Visualize Results¶
Let's use the same plotting function from earlier to visualize how patching the Llama-3.1-8B attention heads influenced model output during the IOI task.
print(f"Clean logit difference: {clean_logit_diff:.3f}")
print(f"Corrupted logit difference: {corrupted_logit_diff:.3f}")
print(ioi_patching_results_all)
N_HEADS = 32
x_labels = [f"Head {i}" for i in range(N_HEADS)]
fig2 = plot_ioi_patching_results(ioi_patching_results_all, x_labels, "Patching Llama Attention Heads on IOI task")
fig2.show()
Clean logit difference: 4.312 Corrupted logit difference: -4.312 [[-0.005645751953125, -0.0003509521484375, 0.003997802734375, 0.0003509521484375, -0.005645751953125, 0.0019989013671875, -0.00762939453125, 0.0, -0.00164794921875, -0.0032958984375, -0.00762939453125, -0.00494384765625, -0.005645751953125, -0.007293701171875, -0.00531005859375, -0.0019989013671875, -0.0089111328125, -0.005645751953125, -0.00494384765625, -0.002349853515625, -0.0032958984375, 0.0019989013671875, -0.0036468505859375, -0.0032958984375, -0.0012969970703125, 0.0019989013671875, -0.00531005859375, 0.0, -0.00164794921875, -0.003997802734375, -0.005645751953125, -0.005645751953125], [-0.00164794921875, -0.003997802734375, -0.00531005859375, -0.0032958984375, -0.00457763671875, -0.0019989013671875, -0.00164794921875, 0.0, -0.0012969970703125, -0.01092529296875, -0.0036468505859375, -0.003997802734375, -0.007293701171875, -0.00762939453125, -0.0036468505859375, -0.0003509521484375, -0.00164794921875, 0.0, 0.0, -0.0036468505859375, -0.00531005859375, -0.00531005859375, -0.0036468505859375, -0.01092529296875, -0.00762939453125, -0.00799560546875, -0.0036468505859375, -0.0036468505859375, -0.006591796875, -0.002349853515625, -0.0032958984375, -0.0069580078125], [0.0, -0.00164794921875, -0.0019989013671875, -0.0036468505859375, -0.0032958984375, -0.0003509521484375, -0.00164794921875, 0.0019989013671875, -0.007293701171875, -0.0036468505859375, -0.0089111328125, 0.0, -0.007293701171875, -0.00762939453125, -0.007293701171875, -0.00762939453125, -0.003997802734375, -0.0012969970703125, 0.0019989013671875, 0.00799560546875, 0.0003509521484375, -0.00164794921875, -0.00164794921875, -0.0036468505859375, -0.00762939453125, -0.002349853515625, -0.00531005859375, -0.01092529296875, -0.0032958984375, 0.011962890625, -0.00494384765625, -0.007293701171875], [-0.00164794921875, -0.005645751953125, -0.0036468505859375, -0.007293701171875, -0.0036468505859375, 0.0003509521484375, -0.0019989013671875, -0.01202392578125, -0.0003509521484375, -0.0036468505859375, -0.003997802734375, -0.007293701171875, -0.00531005859375, 0.0019989013671875, -0.0059814453125, 0.0, -0.0036468505859375, -0.00164794921875, -0.0032958984375, 0.0, -0.005645751953125, -0.0036468505859375, 0.0, -0.0069580078125, -0.0012969970703125, -0.006011962890625, -0.0069580078125, -0.00762939453125, -0.002349853515625, -0.005645751953125, -0.003997802734375, -0.00164794921875], [-0.007293701171875, -0.0106201171875, -0.003997802734375, -0.00164794921875, -0.0036468505859375, -0.002349853515625, 0.0003509521484375, -0.0036468505859375, 0.002349853515625, -0.00164794921875, -0.0032958984375, -0.007659912109375, -0.003997802734375, -0.002349853515625, -0.0019989013671875, -0.00494384765625, 0.0341796875, -0.00164794921875, -0.007293701171875, -0.00164794921875, -0.005645751953125, 0.0003509521484375, -0.00762939453125, 0.00164794921875, -0.00927734375, -0.00531005859375, -0.002349853515625, -0.0019989013671875, -0.0019989013671875, -0.003997802734375, -0.0019989013671875, -0.0019989013671875], [-0.0003509521484375, -0.003997802734375, -0.007293701171875, -0.005645751953125, -0.007293701171875, -0.007293701171875, -0.0032958984375, -0.0019989013671875, -0.0032958984375, -0.0032958984375, -0.00164794921875, -0.00164794921875, -0.0036468505859375, -0.0032958984375, -0.0019989013671875, -0.00164794921875, 0.0, 0.0, -0.002349853515625, -0.01165771484375, 0.000701904296875, -0.00164794921875, -0.003997802734375, -0.003997802734375, -0.00164794921875, -0.00164794921875, -0.0059814453125, -0.0012969970703125, -0.005645751953125, -0.00927734375, -0.0036468505859375, 0.0003509521484375], [-0.0003509521484375, -0.005645751953125, -0.007293701171875, 0.0003509521484375, -0.00164794921875, -0.0036468505859375, -0.0036468505859375, 0.0, -0.0036468505859375, -0.00164794921875, -0.005645751953125, 0.0003509521484375, -0.005645751953125, -0.0089111328125, 0.0, -0.00494384765625, -0.0036468505859375, -0.0019989013671875, -0.005645751953125, -0.005645751953125, 0.003997802734375, 0.00164794921875, 0.0003509521484375, -0.00927734375, -0.00531005859375, -0.003997802734375, -0.005645751953125, -0.0125732421875, -0.005645751953125, -0.003997802734375, -0.0036468505859375, -0.0089111328125], [0.0003509521484375, -0.0059814453125, 0.0003509521484375, 0.0019989013671875, -0.0089111328125, -0.0059814453125, -0.0059814453125, -0.0032958984375, 0.00164794921875, -0.005645751953125, -0.00164794921875, 0.0036468505859375, -0.0036468505859375, -0.0003509521484375, -0.0019989013671875, -0.007293701171875, -0.00762939453125, -0.00494384765625, 0.0003509521484375, -0.0019989013671875, -0.003997802734375, -0.005645751953125, 0.0032958984375, 0.0, -0.0019989013671875, 0.0003509521484375, 0.005645751953125, -0.007659912109375, -0.00164794921875, -0.003997802734375, 0.0, -0.002349853515625], [0.0032958984375, 0.0003509521484375, 0.0, -0.00164794921875, -0.0032958984375, -0.00531005859375, -0.00164794921875, 0.003997802734375, -0.0036468505859375, -0.00164794921875, -0.00164794921875, -0.00494384765625, -0.005645751953125, -0.00531005859375, -0.0019989013671875, -0.0019989013671875, -0.0036468505859375, -0.00164794921875, -0.005645751953125, 0.000701904296875, 0.0, -0.005645751953125, -0.007293701171875, -0.0019989013671875, 0.0019989013671875, 0.000701904296875, 0.003997802734375, -0.0036468505859375, -0.0032958984375, 0.002349853515625, 0.0, 0.0036468505859375], [-0.0036468505859375, 0.00164794921875, -0.00762939453125, -0.00164794921875, 0.0, -0.00164794921875, -0.00531005859375, 0.0003509521484375, -0.0036468505859375, -0.005645751953125, -0.0036468505859375, 0.0, -0.00531005859375, 0.00201416015625, -0.0012969970703125, -0.0003509521484375, -0.005645751953125, -0.0032958984375, -0.0012969970703125, -0.0019989013671875, -0.00164794921875, -0.0036468505859375, 0.0019989013671875, 0.00164794921875, -0.0036468505859375, 0.0036468505859375, -0.0036468505859375, -0.00164794921875, -0.0019989013671875, -0.00531005859375, -0.005645751953125, -0.0032958984375], [0.002593994140625, -0.025390625, -0.0201416015625, 0.0, -0.0069580078125, -0.0036468505859375, -0.0019989013671875, 0.0, -0.00531005859375, -0.005645751953125, 0.00164794921875, -0.0012969970703125, -0.00164794921875, -0.0019989013671875, -0.00164794921875, -0.0059814453125, -0.00164794921875, -0.00164794921875, 0.0, -0.0036468505859375, -0.0036468505859375, -0.005645751953125, -0.0036468505859375, 0.0019989013671875, -0.007293701171875, -0.00164794921875, -0.0032958984375, 0.0, -0.0003509521484375, 0.00164794921875, 0.0, -0.0036468505859375], [-0.00531005859375, -0.00927734375, 0.002349853515625, 0.0, -0.0032958984375, -0.00457763671875, -0.00531005859375, -0.0019989013671875, -0.00531005859375, -0.005645751953125, -0.0059814453125, 0.0, -0.0125732421875, -0.0003509521484375, -0.0069580078125, -0.0032958984375, 0.0003509521484375, -0.0019989013671875, -0.00164794921875, -0.0032958984375, -0.005645751953125, -0.003997802734375, 0.017578125, 0.002349853515625, 0.0, -0.0019989013671875, -0.0059814453125, -0.0089111328125, -0.0036468505859375, -0.005645751953125, -0.00164794921875, -0.0036468505859375], [-0.007293701171875, -0.0036468505859375, -0.00164794921875, -0.0069580078125, -0.002349853515625, -0.0032958984375, -0.00164794921875, 0.006591796875, 0.0, -0.0003509521484375, -0.0032958984375, -0.0032958984375, 0.00860595703125, -0.0089111328125, -0.00164794921875, -0.0012969970703125, -0.0036468505859375, -0.005645751953125, -0.00927734375, -0.0036468505859375, -0.00164794921875, -0.00164794921875, 0.00164794921875, 0.0003509521484375, -0.0032958984375, -0.00823974609375, -0.00164794921875, -0.0012969970703125, -0.0125732421875, -0.00164794921875, -0.00164794921875, 0.0019989013671875], [-0.0019989013671875, -0.005645751953125, -0.00164794921875, -0.00531005859375, -0.0003662109375, 0.0, -0.002349853515625, -0.0019989013671875, -0.0036468505859375, 0.00164794921875, -0.0019989013671875, 0.0036468505859375, 0.0036468505859375, 0.0005950927734375, 0.0012969970703125, -0.00164794921875, 0.00164794921875, 0.0, -0.00927734375, -0.0019989013671875, 0.002349853515625, -0.0003509521484375, -0.005645751953125, 0.0, -0.0019989013671875, -0.0036468505859375, -0.0089111328125, -0.00164794921875, 0.0012969970703125, -0.0032958984375, -0.00164794921875, 0.0059814453125], [0.0, 0.0, 0.0, 0.0, -0.0003509521484375, -0.00531005859375, 0.0238037109375, -0.0106201171875, 0.0, -0.0003509521484375, -0.00164794921875, 0.0, 0.0036468505859375, 0.007293701171875, -0.0003509521484375, -0.0019989013671875, 0.033935546875, -0.0036468505859375, 0.0, 0.044921875, -0.00164794921875, 0.017578125, -0.007293701171875, -0.0032958984375, -0.0059814453125, -0.005645751953125, -0.00164794921875, 0.0, -0.005645751953125, -0.00164794921875, -0.00531005859375, -0.007293701171875], [0.0, 0.002349853515625, -0.0003509521484375, -0.003997802734375, 0.0, -0.002349853515625, 0.00164794921875, -0.0003509521484375, 0.00164794921875, 0.0, -0.01220703125, 0.0032958984375, 0.0019989013671875, 0.0019989013671875, 0.00164794921875, -0.003997802734375, 0.05908203125, 0.0036468505859375, 0.0, 0.0, -0.00164794921875, 0.00164794921875, 0.0019989013671875, -0.003997802734375, 0.0019989013671875, -0.00164794921875, -0.003997802734375, 0.0, -0.002349853515625, -0.0019989013671875, -0.0003509521484375, 0.00164794921875], [-0.016845703125, 0.09130859375, 0.032470703125, -0.0162353515625, 0.002349853515625, -0.0019989013671875, -0.00164794921875, -0.0019989013671875, -0.005645751953125, 0.0019989013671875, -0.0019989013671875, -0.003997802734375, -0.005645751953125, -0.003997802734375, 0.0, -0.0036468505859375, -0.029052734375, 0.0089111328125, 0.0036468505859375, 0.0, 0.0019989013671875, 0.00164794921875, 0.0233154296875, 0.0019989013671875, 0.00164794921875, -0.003997802734375, 0.0036468505859375, -0.00164794921875, 0.0128173828125, 0.00164794921875, -0.01483154296875, -0.00164794921875], [-0.0036468505859375, 0.0019989013671875, 0.0, 0.0019989013671875, 0.0019989013671875, 0.0019989013671875, -0.00164794921875, -0.005645751953125, 0.0019989013671875, 0.0019989013671875, 0.0, -0.00164794921875, 0.0019989013671875, 0.0019989013671875, 0.0036468505859375, 0.0, -0.00164794921875, 0.0, 0.0, -0.0019989013671875, 0.0, 0.0019989013671875, 0.0, 0.0003509521484375, 0.0223388671875, -0.01220703125, -0.007293701171875, 0.126953125, 0.0019989013671875, 0.0019989013671875, 0.0, -0.00164794921875], [0.0019989013671875, 0.0003509521484375, 0.0059814453125, 0.0019989013671875, 0.0003509521484375, -0.0036468505859375, 0.0019989013671875, 0.0003509521484375, -0.003997802734375, 0.0036468505859375, -0.00469970703125, -0.00164794921875, -0.0019989013671875, 0.00164794921875, 0.00164794921875, 0.0003509521484375, 0.003997802734375, -0.0059814453125, -0.0019989013671875, -0.0036468505859375, -0.00531005859375, 0.0019989013671875, 0.00433349609375, -0.012939453125, -0.0059814453125, -0.0019989013671875, 0.00164794921875, 0.005645751953125, -0.0012969970703125, 0.064453125, 0.0019989013671875, 0.002349853515625], [-0.00164794921875, -0.005645751953125, 0.0019989013671875, -0.021484375, 0.0036468505859375, 0.0003509521484375, -0.0003509521484375, 0.0019989013671875, 0.005645751953125, -0.00164794921875, 0.0, -0.005645751953125, 0.00164794921875, -0.0019989013671875, 0.00164794921875, -0.00860595703125, -0.0003509521484375, 0.0, 0.0003509521484375, 0.0, -0.05615234375, 0.00164794921875, 0.0, 0.08935546875, 0.0, -0.0003509521484375, -0.0003509521484375, 0.00164794921875, 0.0, 0.00164794921875, -0.0019989013671875, 0.0036468505859375], [-0.0059814453125, -0.002349853515625, -0.00164794921875, 0.041259765625, 0.0, 0.0019989013671875, 0.0, 0.00164794921875, -0.003997802734375, 0.0, -0.0019989013671875, 0.0, -0.0771484375, 0.05224609375, -0.166015625, -0.0036468505859375, -0.00164794921875, -0.005645751953125, -0.007293701171875, 0.0019989013671875, -0.00531005859375, 0.0, 0.0, -0.0019989013671875, -0.0019989013671875, -0.00799560546875, -0.0059814453125, 0.0003509521484375, -0.00164794921875, -0.0019989013671875, 0.0, -0.00164794921875], [0.0, -0.0019989013671875, 0.01458740234375, -0.0228271484375, 0.0, 0.0, 0.0003509521484375, 0.0, -0.0135498046875, 0.00164794921875, 0.02490234375, 0.0, 0.0, 0.0036468505859375, 0.0019989013671875, -0.0019989013671875, 0.0, 0.0, 0.00164794921875, 0.0, 0.0, 0.00164794921875, -0.0036468505859375, 0.0, -0.0019989013671875, 0.0019989013671875, -0.00164794921875, 0.0, -0.00164794921875, 0.0, 0.0, 0.0], [0.0019989013671875, -0.0019989013671875, 0.0019989013671875, 0.0019989013671875, -0.0003509521484375, 0.0, 0.0, 0.0, 0.0, 0.0, -0.0019989013671875, -0.0003509521484375, 0.021484375, 0.0, -0.0184326171875, -0.0012969970703125, 0.0036468505859375, -0.0019989013671875, -0.030517578125, 0.1455078125, 0.0, 0.0, -0.0032958984375, 0.0, 0.0, 0.0, 0.0, -0.0019989013671875, -0.0019989013671875, 0.0019989013671875, 0.0206298828125, 0.0], [0.0, -0.0019989013671875, -0.0019989013671875, 0.0, 0.0019989013671875, 0.00164794921875, 0.0, -0.00531005859375, 0.0019989013671875, 0.0019989013671875, 0.0, 0.00164794921875, 0.0019989013671875, -0.0019989013671875, 0.01153564453125, 0.0, -0.0003509521484375, 0.0, -0.0019989013671875, -0.0012969970703125, -0.0003509521484375, 0.0059814453125, 0.0, 0.0019989013671875, 0.0, 0.00164794921875, 0.00164794921875, 0.0019989013671875, 0.00164794921875, 0.0019989013671875, -0.005645751953125, -0.00164794921875], [0.0, -0.0019989013671875, -0.0019989013671875, -0.0003509521484375, 0.0, 0.003997802734375, 0.0, 0.0, 0.0019989013671875, 0.0, 0.0, 0.0, -0.0019989013671875, -0.0019989013671875, -0.0019989013671875, 0.0, 0.0019989013671875, -0.0019989013671875, -0.0019989013671875, -0.0019989013671875, -0.0019989013671875, -0.0032958984375, 0.02197265625, -0.0036468505859375, -0.0106201171875, -0.0019989013671875, 0.00927734375, 0.0859375, 0.0, 0.003997802734375, -0.0019989013671875, 0.003997802734375], [-0.0036468505859375, 0.005645751953125, -0.0019989013671875, -0.0019989013671875, 0.053955078125, 0.0012969970703125, 0.0, 0.0, -0.0019989013671875, 0.0019989013671875, -0.0019989013671875, 0.0, -0.0019989013671875, 0.029296875, 0.00164794921875, -0.0019989013671875, -0.0036468505859375, -0.0019989013671875, -0.00164794921875, 0.0, -0.0032958984375, -0.0019989013671875, 0.0019989013671875, 0.0, -0.0019989013671875, 0.0019989013671875, -0.0019989013671875, 0.0019989013671875, 0.0, 0.0, 0.0, 0.0], [-0.00762939453125, -0.0003509521484375, 0.0, 0.0169677734375, -0.00164794921875, -0.0019989013671875, 0.0019989013671875, 0.0, 0.0019989013671875, 0.0, -0.0019989013671875, 0.0003509521484375, 0.045654296875, -0.00164794921875, 0.0019989013671875, -0.07666015625, 0.0, 0.0, 0.0, -0.0019989013671875, -0.00164794921875, 0.0019989013671875, 0.0, 0.0019989013671875, 0.0, 0.0003509521484375, 0.0059814453125, 0.0, -0.005645751953125, -0.0019989013671875, 0.0019989013671875, -0.0019989013671875], [0.0, 0.0, 0.0019989013671875, 0.0019989013671875, 0.0, 0.0, -0.0036468505859375, 0.002349853515625, 0.0, 0.003997802734375, 0.0, 0.0, -0.0019989013671875, 0.0, 0.0, 0.003997802734375, -0.003997802734375, -0.0019989013671875, 0.003997802734375, 0.0, 0.126953125, 0.01190185546875, 0.019775390625, -0.0019989013671875, 0.0, 0.0, 0.0, 0.0019989013671875, 0.003997802734375, 0.003997802734375, -0.0019989013671875, 0.0], [0.0, -0.003997802734375, 0.0, 0.0019989013671875, -0.0019989013671875, -0.003997802734375, 0.0, -0.0019989013671875, 0.0003509521484375, -0.0019989013671875, 0.0, -0.0019989013671875, 0.0223388671875, -0.00164794921875, 0.0029449462890625, -0.046630859375, 0.0019989013671875, -0.0036468505859375, 0.00164794921875, -0.0019989013671875, 0.00164794921875, 0.0, 0.0, -0.0019989013671875, -0.0019989013671875, -0.0019989013671875, 0.0019989013671875, 0.0019989013671875, 0.0, -0.0019989013671875, -0.0019989013671875, -0.0019989013671875], [0.0019989013671875, -0.0019989013671875, -0.00164794921875, -0.0019989013671875, 0.0019989013671875, 0.0019989013671875, 0.0003509521484375, -0.00164794921875, 0.022216796875, 0.0, -0.12060546875, 0.03369140625, -0.0036468505859375, -0.0019989013671875, 0.0, -0.0036468505859375, 0.0, -0.0036468505859375, -0.0019989013671875, -0.0019989013671875, -0.0019989013671875, 0.0, 0.0003509521484375, 0.0019989013671875, 0.0003509521484375, 0.0019989013671875, 0.0, 0.002349853515625, 0.0019989013671875, -0.0019989013671875, -0.0036468505859375, 0.0019989013671875], [-0.0036468505859375, 0.0, 0.0, -0.0019989013671875, 0.0019989013671875, 0.0, 0.0, 0.0, 0.0019989013671875, 0.0, 0.0019989013671875, 0.0, 0.006591796875, 0.05615234375, 0.01129150390625, -0.0003509521484375, 0.0019989013671875, 0.0, -0.0019989013671875, 0.00799560546875, -0.0019989013671875, 0.0019989013671875, 0.0, 0.0, 0.0059814453125, -0.07373046875, 0.012939453125, 0.018310546875, -0.003997802734375, -0.0036468505859375, 0.0390625, 0.0289306640625], [-0.00531005859375, 0.0185546875, -0.0198974609375, 0.015869140625, 0.0019989013671875, -0.00164794921875, 0.0019989013671875, 0.005645751953125, 0.0019989013671875, -0.0036468505859375, 0.0019989013671875, -0.00927734375, 0.0019989013671875, -0.00531005859375, 0.00531005859375, 0.0019989013671875, 0.0019989013671875, 0.0, 0.0, 0.0019989013671875, -0.00457763671875, 0.0242919921875, 0.0478515625, -0.00762939453125, 0.0, -0.0036468505859375, 0.0019989013671875, 0.0, 0.00531005859375, -0.00164794921875, 0.0019989013671875, -0.003997802734375]]