docs/methods/activation_patching.md

Activation Patching

Home / Methods / Activation Patching

Replace selected residual activations with those from a paired input, then measure how the chosen output score changes. Model weights are unchanged.

API and parameters

probe is a Prober bound to a loaded model; inputs contains prepared model inputs. The blocks below are API signatures: * marks keyword-only parameters, not an argument to copy. The complete example defines the model, inputs, and all other variables before running the method.

method = probe.causal.patch(*, layers=None, tokens="visual")
result = method.run(inputs, *, source, metric, alignment="strict")

Parameter	Default	Meaning and constraints
layers	None	All available residual layers; pass one integer or a nonempty list of unique available layer IDs. Check probe.describe().
tokens	"visual"	Positions to replace/attribute. Use expanded token selectors described below.
inputs	required	Receiver/base input: native model keyword dictionary or ProbeInputs. Full sequence, no KV cache.
source	required	Aligned donor input providing replacement residuals. It need not be the clean condition; define the intervention direction explicitly.
metric	required	Output score to compare. See metric contract below.
alignment	"strict"	"strict" checks token IDs, masks, and declared alignment metadata. "position" permits different token IDs, retaining all structural checks. Neither establishes semantic correspondence automatically.

Selections use zero-based expanded model coordinates, including image tokens. Common token selectors are "all", "visual", "text", "last_prompt", "prediction", an integer/list of positions, or a boolean [B,T] mask. "last_prompt" selects one final context token; "prediction" additionally selects generated positions when TokenLayout.prompt preserves the original context. "visual"/"text" need a visual-token layout. For text-only examples, explicitly choose tokens="all" or "last_prompt" rather than the visual-token default.

A metric is a TokenMargin, SequenceLogProb, or a callable metric(logits) returning a scalar or [B]; it is evaluated on full-sequence logits. TokenMargin(positive, negative, position="last_prompt") compares two vocabulary IDs at one selected prediction position. For multi-token answers, explicitly append answer tokens, preserve the original prompt boundary, and use SequenceLogProb(answer_mask, reduction="sum") or "mean". First-subword margins are not complete-answer probabilities.

Computation and returned tensors

At each selected layer independently, compute the receiver and source residuals, copy source values only at selected positions, and execute a new receiver forward. This is a measured intervention, not a gradient estimate. A configured list of layers produces separate experiments; it does not edit all layers simultaneously.

For a per-example score, baseline_score, intervention_score, and effect are [L,B]; a scalar metric gives [L]. The sign is effect = intervention_score - baseline_score. If source is the corrupt reference and inputs is clean, negative values indicate reduced support for the selected clean answer. Outputs are detached on Prober(result_device="cpu") by default. result.save(directory) writes tensors and metadata, not model weights.

Install pip install "vlmprobe[transformers,visualization]". The following block is complete: it loads GPT-2, constructs an author-published IOI prompt pair, checks candidate tokenization, runs the methods, saves results, and draws the built-in plots. First use downloads GPT-2 (about 0.5 GB); CPU works and CUDA is faster. No method training is needed. VLM_PROBING_MODEL can point to an existing local checkpoint.

Complete example: input to output

import os
from pathlib import Path
import torch
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, GPT2LMHeadModel
from vlmprobe import Prober, TokenMargin
from vlmprobe.visualization import plot_causal_heatmap, plot_intervention_curves

# First run downloads GPT-2 (~0.5 GB). Use an existing local checkpoint if desired.
model_id = os.environ.get("VLM_PROBING_MODEL", "openai-community/gpt2")
device = os.environ.get("VLM_PROBING_DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu")
model = GPT2LMHeadModel.from_pretrained(model_id, attn_implementation="eager").to(device).eval()
tokenizer = AutoTokenizer.from_pretrained(model_id)
probe = Prober(model, processor=tokenizer)
print({name: info for name, info in probe.describe()["methods"].items()
       if name.startswith("causal.")})
output_dir = Path(os.environ.get("VLM_PROBING_OUTPUT", "outputs/tutorials/causal"))
output_dir.mkdir(parents=True, exist_ok=True)
# An explicit author-published IOI pair (row 826); no result snapshots are used.
# Source: hannamw/eap-ig-faithfulness, data/ioi/gpt2.csv; see docs/REFERENCES.md.
clean_prompt = "When Amy and Laura got a snack at the house, Laura decided to give it to"
reference_prompt = "When Amy and Laura got a snack at the house, Nicholas decided to give it to"

def prepare_text(text):
    # The authors prepend GPT-2's BOS token. Keep the convention explicit.
    ids = [tokenizer.bos_token_id, *tokenizer.encode(text, add_special_tokens=False)]
    input_ids = torch.tensor([ids], device=device)
    return {"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids)}

inputs = prepare_text(clean_prompt)
reference = prepare_text(reference_prompt)
assert inputs["input_ids"].shape == reference["input_ids"].shape, "Pair token lengths must match"
positive_ids = tokenizer.encode(" Amy", add_special_tokens=False)
negative_ids = tokenizer.encode(" Laura", add_special_tokens=False)
assert len(positive_ids) == len(negative_ids) == 1, "TokenMargin requires one-token candidates"
metric = TokenMargin(positive=positive_ids[0], negative=negative_ids[0])
print("Clean:", clean_prompt)
print("Reference:", reference_prompt)
print("Candidates:", tokenizer.decode(positive_ids), tokenizer.decode(negative_ids))
# Check the role alignment, not just the equal length.
for position, (a, b) in enumerate(zip(inputs["input_ids"][0], reference["input_ids"][0])):
    print(position, repr(tokenizer.decode([a.item()])), repr(tokenizer.decode([b.item()])))
layers = sorted(probe.spec.residuals)
# Both methods replace/estimate the same reference-minus-clean change.
# This is corruption of the clean input, not restoration of the reference.
exact = probe.causal.patch(layers=layers, tokens="last_prompt").run(
    inputs, source=reference, metric=metric, alignment="position",
)
estimated = probe.causal.attribute(layers=layers, tokens="last_prompt").run(
    inputs, source=reference, metric=metric, alignment="position",
)
print("Exact effect [layer,batch]:", exact.tensors["effect"])
print("Estimated effect [layer]:", estimated.tensors["estimated_effect"])
# A separate all-position attribution scan supplies the token heatmap.
token_map = probe.causal.attribute(layers=layers, tokens="all").run(
    inputs, source=reference, metric=metric, alignment="position",
)
print("Token scores [layer,batch,token]:", token_map.tensors["token_scores"].shape)
exact.save(output_dir / "activation_patching")
estimated.save(output_dir / "attribution_patching")
token_map.save(output_dir / "attribution_patching_all_positions")
fig, axes = plt.subplots(1, 2, figsize=(14, 4.8), constrained_layout=True)
plot_intervention_curves(
    {"Measured patch": exact.tensors["effect"][:, 0],
     "First-order estimate": estimated.tensors["estimated_effect"]},
    x=layers, ax=axes[0], xlabel="Decoder layer", ylabel="Change in Amy-minus-Laura logit margin",
    title="The approximation can differ from the measured effect",
)
labels = ["<BOS>" if i == 0 else tokenizer.decode([token])
          for i, token in enumerate(inputs["input_ids"][0].tolist())]
plot_causal_heatmap(
    token_map, tensor="token_scores", column_labels=labels, ax=axes[1],
    title="Separate attribution scan across all positions",
    value_label="Estimated margin change",
)
axes[1].tick_params(axis="x", labelsize=7)
fig.savefig(output_dir / "causal-patching.png", dpi=150, bbox_inches="tight")
plt.show()

Model-specific use

The factory/run interface is identical across supported LM and VLM residual adapters. Replace only model loading and native input preparation using the model guides. For a gray-image reference, preserve prompt IDs, image resolution, masks and grid metadata, changing only pixels. Patching requires matching expanded coordinates; arbitrary token/region remapping is not implemented. Equal grid shapes do not prove equal semantic regions.

Live notebook · Standalone script · Reference · Implementation

Example visualization

This preview was generated by running the code on pretrained GPT-2 with the explicit IOI pair below. It is an example-level result, not a full-dataset reproduction. The notebook reruns inference and does not load the preview.