docs/methods/activation_patching.md

# Activation Patching

[Home](../../README.md) / [Methods](README.md) / 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](#complete-example-input-to-output) defines the model, inputs, and all other variables before running the method.

```text
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

```python
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](../models/README.md). 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](../../demos/activation_patching_demo.ipynb) · [Standalone script](../../examples/tutorials/causal_activation_patching.py) · [Reference](../REFERENCES.md#causal-methods) · [Implementation](../../src/vlm_probing/causal/activation_patching.py)

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

![Measured residual replacements and first-order token attribution](../assets/tutorials/causal-patching.png)