1️⃣ Setup¶

The task is indirect object identification (IOI). Each prompt names two people and then says that one of them gave something to, and the model has to complete it with the other name. Every prompt is paired with the one that swaps the two names: same tokens, same length, opposite answer. That pair is the clean/corrupted pair the rest of the notebook works on.

In [1]:

try:
    import google.colab
    is_colab = True
except ImportError:
    is_colab = False

if is_colab:
    !pip install -U nnsight

try: import google.colab is_colab = True except ImportError: is_colab = False if is_colab: !pip install -U nnsight

In [2]:

from IPython.display import clear_output
import einops
import torch
import plotly.express as px
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+notebook"


from nnsight import TransformersModel

from IPython.display import clear_output import einops import torch import plotly.express as px import plotly.io as pio pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+notebook" from nnsight import TransformersModel

/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

In [3]:

import nnsight
print(nnsight.__version__)

import nnsight print(nnsight.__version__)

0.8.0

In [4]:

model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
clear_output()
print(model)

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()
  )
)

In [5]:

prompts = [
    "When John and Mary went to the shops, John gave the bag to",
    "When John and Mary went to the shops, Mary gave the bag to",
    "When Tom and James went to the park, James gave the ball to",
    "When Tom and James went to the park, Tom gave the ball to",
    "When Dan and Sid went to the shops, Sid gave an apple to",
    "When Dan and Sid went to the shops, Dan gave an apple to",
    "After Martin and Amy went to the park, Amy gave a drink to",
    "After Martin and Amy went to the park, Martin gave a drink to",
]

# Answers are each formatted as (correct, incorrect):
answers = [
    (" Mary", " John"),
    (" John", " Mary"),
    (" Tom", " James"),
    (" James", " Tom"),
    (" Dan", " Sid"),
    (" Sid", " Dan"),
    (" Martin", " Amy"),
    (" Amy", " Martin"),
]

# Tokenize clean and corrupted inputs:
clean_tokens = model.tokenizer(prompts, return_tensors="pt")["input_ids"].to(model.device)
# The associated corrupted input is the prompt after the current clean prompt
# for even indices, or the prompt prior to the current clean prompt for odd indices
corrupted_tokens = clean_tokens[
    [(i + 1 if i % 2 == 0 else i - 1) for i in range(len(clean_tokens))]
]

# Tokenize answers for each prompt:
answer_token_indices = torch.tensor(
    [
        [model.tokenizer(answers[i][j])["input_ids"][0] for j in range(2)]
        for i in range(len(answers))
    ]
)

prompts = [ "When John and Mary went to the shops, John gave the bag to", "When John and Mary went to the shops, Mary gave the bag to", "When Tom and James went to the park, James gave the ball to", "When Tom and James went to the park, Tom gave the ball to", "When Dan and Sid went to the shops, Sid gave an apple to", "When Dan and Sid went to the shops, Dan gave an apple to", "After Martin and Amy went to the park, Amy gave a drink to", "After Martin and Amy went to the park, Martin gave a drink to", ] # Answers are each formatted as (correct, incorrect): answers = [ (" Mary", " John"), (" John", " Mary"), (" Tom", " James"), (" James", " Tom"), (" Dan", " Sid"), (" Sid", " Dan"), (" Martin", " Amy"), (" Amy", " Martin"), ] # Tokenize clean and corrupted inputs: clean_tokens = model.tokenizer(prompts, return_tensors="pt")["input_ids"].to(model.device) # The associated corrupted input is the prompt after the current clean prompt # for even indices, or the prompt prior to the current clean prompt for odd indices corrupted_tokens = clean_tokens[ [(i + 1 if i % 2 == 0 else i - 1) for i in range(len(clean_tokens))] ] # Tokenize answers for each prompt: answer_token_indices = torch.tensor( [ [model.tokenizer(answers[i][j])["input_ids"][0] for j in range(2)] for i in range(len(answers)) ] )

Next, we create a function to calculate the mean logit difference for the correct vs incorrect answer tokens.

In [6]:

def get_logit_diff(logits, answer_token_indices=answer_token_indices):
    logits = logits[:, -1, :]
    correct_logits = logits.gather(1, answer_token_indices[:, 0].unsqueeze(1))
    incorrect_logits = logits.gather(1, answer_token_indices[:, 1].unsqueeze(1))
    return (correct_logits - incorrect_logits).mean()

def get_logit_diff(logits, answer_token_indices=answer_token_indices): logits = logits[:, -1, :] correct_logits = logits.gather(1, answer_token_indices[:, 0].unsqueeze(1)) incorrect_logits = logits.gather(1, answer_token_indices[:, 1].unsqueeze(1)) return (correct_logits - incorrect_logits).mean()

We then calculate the logit difference for both the clean and the corrupted baselines.

In [7]:

clean_logits = model.trace(clean_tokens, trace=False).logits.cpu()
corrupted_logits = model.trace(corrupted_tokens, trace=False).logits.cpu()

CLEAN_BASELINE = get_logit_diff(clean_logits, answer_token_indices).item()
print(f"Clean logit diff: {CLEAN_BASELINE:.4f}")

CORRUPTED_BASELINE = get_logit_diff(corrupted_logits, answer_token_indices).item()
print(f"Corrupted logit diff: {CORRUPTED_BASELINE:.4f}")

clean_logits = model.trace(clean_tokens, trace=False).logits.cpu() corrupted_logits = model.trace(corrupted_tokens, trace=False).logits.cpu() CLEAN_BASELINE = get_logit_diff(clean_logits, answer_token_indices).item() print(f"Clean logit diff: {CLEAN_BASELINE:.4f}") CORRUPTED_BASELINE = get_logit_diff(corrupted_logits, answer_token_indices).item() print(f"Corrupted logit diff: {CORRUPTED_BASELINE:.4f}")

Clean logit diff: 2.8138
Corrupted logit diff: -2.8138

Now let's define an ioi_metric function to evaluate patched IOI changes normalized to our clean and corruped baselines.

In [8]:

def ioi_metric(
    logits,
    answer_token_indices=answer_token_indices,
):
    return (get_logit_diff(logits, answer_token_indices) - CORRUPTED_BASELINE) / (
        CLEAN_BASELINE - CORRUPTED_BASELINE
    )

print(f"Clean Baseline is 1: {ioi_metric(clean_logits).item():.4f}")
print(f"Corrupted Baseline is 0: {ioi_metric(corrupted_logits).item():.4f}")

def ioi_metric( logits, answer_token_indices=answer_token_indices, ): return (get_logit_diff(logits, answer_token_indices) - CORRUPTED_BASELINE) / ( CLEAN_BASELINE - CORRUPTED_BASELINE ) print(f"Clean Baseline is 1: {ioi_metric(clean_logits).item():.4f}") print(f"Corrupted Baseline is 0: {ioi_metric(corrupted_logits).item():.4f}")

Clean Baseline is 1: 1.0000
Corrupted Baseline is 0: 0.0000

2️⃣ Attribution Patching Over Components¶

Attribution patching is a technique that uses gradients to take a linear approximation to activation patching. The key assumption is that the corrupted run is a locally linear function of its activations.

We take the gradient of the patch metric (ioi_metric) with respect to its activations, where a patch of activations means moving corrupted_x to corrupted_x + (clean_x - corrupted_x). The metric's change is then (corrupted_grad_x * (clean_x - corrupted_x)).sum(). All we need is one backward pass on the corrupted prompt with respect to the patch metric, caching every gradient with respect to the activations.

The activations are already non-leaf tensors inside the trace, so none of them needs requires_grad_(True) — reading .grad inside with value.backward(): is enough.

A note on c_proj: Most HuggingFace models don’t have nice individual attention head representations to hook. Instead, we have the linear layer c_proj which implicitly combines the “projection per attention head” and the “sum over attention head” operations. See this snippet from ARENA for more information.

TL;DR: We will use the input to c_proj for causal interventions on a particular attention head.

In [9]:

clean_out = []
corrupted_out = []
corrupted_grads = []

# Pass 1 ” clean run.
with model.trace(clean_tokens):
    for layer in model.transformer.h:
        # Clean attention output for this layer, across all heads
        attn_out = layer.attn.c_proj.input
        clean_out.append(attn_out.save())

# Pass 2” corrupted run + backward.
with model.trace(corrupted_tokens):
    corrupted_refs = []
    for layer in model.transformer.h:
        attn_out = layer.attn.c_proj.input
        corrupted_refs.append(attn_out)
        corrupted_out.append(attn_out.save())

    logits = model.lm_head.output.save()

    value = ioi_metric(logits.cpu())

    with value.backward():
        for attn_out in reversed(corrupted_refs):
            corrupted_grads.insert(0, attn_out.grad.save())

clean_out = [] corrupted_out = [] corrupted_grads = [] # Pass 1 ” clean run. with model.trace(clean_tokens): for layer in model.transformer.h: # Clean attention output for this layer, across all heads attn_out = layer.attn.c_proj.input clean_out.append(attn_out.save()) # Pass 2” corrupted run + backward. with model.trace(corrupted_tokens): corrupted_refs = [] for layer in model.transformer.h: attn_out = layer.attn.c_proj.input corrupted_refs.append(attn_out) corrupted_out.append(attn_out.save()) logits = model.lm_head.output.save() value = ioi_metric(logits.cpu()) with value.backward(): for attn_out in reversed(corrupted_refs): corrupted_grads.insert(0, attn_out.grad.save())

Next, for a given activation we compute (corrupted_grad_act * (clean_act - corrupted_act)).sum(). We use einops.reduce to rearrange and sum activations over the correct dimension. In this case, we want to estimate the effect of specific attention heads, so we sum over heads rather than token position.

In [10]:

patching_results = []

for corrupted_grad, corrupted, clean, layer in zip(
    corrupted_grads, corrupted_out, clean_out, range(len(clean_out))
):

    residual_attr = einops.reduce(
        corrupted_grad[:,-1,:] * (clean[:,-1,:] - corrupted[:,-1,:]),
        "batch (head dim) -> head",
        "sum",
        head = 12,
        dim = 64,
    )

    patching_results.append(
        residual_attr.detach().cpu().numpy()
    )

patching_results = [] for corrupted_grad, corrupted, clean, layer in zip( corrupted_grads, corrupted_out, clean_out, range(len(clean_out)) ): residual_attr = einops.reduce( corrupted_grad[:,-1,:] * (clean[:,-1,:] - corrupted[:,-1,:]), "batch (head dim) -> head", "sum", head = 12, dim = 64, ) patching_results.append( residual_attr.detach().cpu().numpy() )

In [11]:

fig = px.imshow(
    patching_results,
    color_continuous_scale="RdBu",
    color_continuous_midpoint=0.0,
    title="Attribution Patching Over Attention Heads",
    labels={"x": "Head", "y": "Layer","color":"Norm. Logit Diff"},

)

fig.show()

fig = px.imshow( patching_results, color_continuous_scale="RdBu", color_continuous_midpoint=0.0, title="Attribution Patching Over Attention Heads", labels={"x": "Head", "y": "Layer","color":"Norm. Logit Diff"}, ) fig.show()