Coverage for transformer_lens/patching.py: 96%
370 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Patching.
3A module for patching activations in a transformer model, and measuring the effect of the patch on
4the output. This implements the activation patching technique for a range of types of activation.
5The structure is to have a single :func:`generic_activation_patch` function that does everything,
6and to have a range of specialised functions for specific types of activation.
7:func:`get_act_patch_direct_path` patches one head's output into a single later head's Q/K/V input.
9Context:
11Activation Patching is technique introduced in the `ROME paper <http://rome.baulab.info/>`, which
12uses a causal intervention to identify which activations in a model matter for producing some
13output. It runs the model on input A, replaces (patches) an activation with that same activation on
14input B, and sees how much that shifts the answer from A to B.
16More details: The setup of activation patching is to take two runs of the model on two different
17inputs, the clean run and the corrupted run. The clean run outputs the correct answer and the
18corrupted run does not. The key idea is that we give the model the corrupted input, but then
19intervene on a specific activation and patch in the corresponding activation from the clean run (ie
20replace the corrupted activation with the clean activation), and then continue the run. And we then
21measure how much the output has updated towards the correct answer.
23- We can then iterate over many
24 possible activations and look at how much they affect the corrupted run. If patching in an
25 activation significantly increases the probability of the correct answer, this allows us to
26 localise which activations matter.
27- A key detail is that we move a single activation __from__ the clean run __to __the corrupted run.
28 So if this changes the answer from incorrect to correct, we can be confident that the activation
29 moved was important.
31Intuition:
33The ability to **localise** is a key move in mechanistic interpretability - if the computation is
34diffuse and spread across the entire model, it is likely much harder to form a clean mechanistic
35story for what's going on. But if we can identify precisely which parts of the model matter, we can
36then zoom in and determine what they represent and how they connect up with each other, and
37ultimately reverse engineer the underlying circuit that they represent. And, empirically, on at
38least some tasks activation patching tends to find that computation is extremely localised:
40- This technique helps us precisely identify which parts of the model matter for a certain
41 part of a task. Eg, answering “The Eiffel Tower is in” with “Paris” requires figuring out that
42 the Eiffel Tower is in Paris, and that it’s a factual recall task and that the output is a
43 location. Patching to “The Colosseum is in” controls for everything other than the “Eiffel Tower
44 is located in Paris” feature.
45- It helps a lot if the corrupted prompt has the same number of tokens
47This, unlike direct logit attribution, can identify meaningful parts of a circuit from anywhere
48within the model, rather than just the end.
49"""
51from __future__ import annotations
53import itertools
54from contextlib import contextmanager
55from dataclasses import dataclass
56from functools import partial
57from typing import (
58 Any,
59 Callable,
60 Dict,
61 Iterator,
62 List,
63 Mapping,
64 Optional,
65 Sequence,
66 Tuple,
67 Union,
68 cast,
69 overload,
70)
72import einops
73import pandas as pd
74import torch
75from jaxtyping import Float, Int
76from tqdm.auto import tqdm
77from typing_extensions import Literal
79import transformer_lens.utilities as utils
80from transformer_lens.ActivationCache import ActivationCache
81from transformer_lens.model_protocol import TransformerLensModel
82from transformer_lens.utilities.statistics import (
83 bootstrap_ci,
84 derive_generator,
85 sign_flip_permutation_pvalue,
86 standard_error,
87)
89# %%
90Logits = torch.Tensor
91AxisNames = Literal["layer", "pos", "head_index", "head", "src_pos", "dest_pos"]
94# %%
97def make_df_from_ranges(
98 column_max_ranges: Sequence[int], column_names: Sequence[str]
99) -> pd.DataFrame:
100 """
101 Takes in a list of column names and max ranges for each column, and returns a dataframe with the cartesian product of the range for each column (ie iterating through all combinations from zero to column_max_range - 1, in order, incrementing the final column first)
102 """
103 rows = list(itertools.product(*[range(axis_max_range) for axis_max_range in column_max_ranges]))
104 df = pd.DataFrame(rows, columns=column_names)
105 return df
108# %%
109CorruptedActivation = torch.Tensor
110PatchedActivation = torch.Tensor
111MetricFn = Callable[[Float[torch.Tensor, "batch pos d_vocab"]], torch.Tensor]
112Reduce = Literal["scalar", "mean", "none"]
113_REDUCE_MODES = ("scalar", "mean", "none")
114_RESULT_AXIS_RENAMES = {"head_index": "head", "dest_pos": "pos"}
117@dataclass
118class PatchingResult:
119 """Per-example (or per-metric) activation-patching results with statistics on top.
121 Returned by :func:`generic_activation_patch` and every ``get_act_patch_*``
122 function when ``reduce="none"``, ``baseline=True`` or several metrics are
123 requested. ``values`` is laid out ``[metric?, *axes | n_rows, batch?]``: the
124 metric axis exists only when a mapping of metrics was passed
125 (``has_metric_axis``), the middle axes are ``axis_names`` (or one flat
126 ``n_rows`` axis when an ``index_df`` was supplied), and the trailing batch
127 axis exists only when ``reduced`` is ``False``.
129 Attributes:
130 values: The patched metric values.
131 axis_names: Names of the swept axes, ``[]`` in flat ``index_df`` mode.
132 index_df: One row per swept cell, in the order the cells were run.
133 metric_names: Metric names, ``["metric"]`` for a single callable.
134 reduced: ``True`` when the batch axis has been averaged away.
135 baseline: Unpatched corrupted-run metric, ``[metric?, batch]`` (or
136 ``[metric?]`` when reduced), when the sweep was run with
137 ``baseline=True``.
138 has_metric_axis: Whether ``values`` carries a leading metric axis.
139 patch_type_names: Names of the ``patch_type`` axis entries when the
140 result came from a ``*_every`` helper (``"out"``, ``"q"``, ...).
141 GQA ``k``/``v`` sweeps are zero-padded to ``n_heads`` there, so the
142 padded cells appear as zero-valued rows in ``index_df`` and
143 :meth:`to_dataframe`.
144 """
146 values: torch.Tensor
147 axis_names: List[str]
148 index_df: pd.DataFrame
149 metric_names: List[str]
150 reduced: bool
151 baseline: Optional[torch.Tensor] = None
152 has_metric_axis: bool = False
153 patch_type_names: Optional[List[str]] = None
155 def __getitem__(self, metric_name: str) -> "PatchingResult":
156 """One metric's result, with the metric axis dropped."""
157 if metric_name not in self.metric_names:
158 raise KeyError(f"unknown metric {metric_name!r}; have {self.metric_names}")
159 if not self.has_metric_axis: 159 ↛ 160line 159 didn't jump to line 160 because the condition on line 159 was never true
160 return self
161 i = self.metric_names.index(metric_name)
162 return PatchingResult(
163 values=self.values[i],
164 axis_names=list(self.axis_names),
165 index_df=self.index_df,
166 metric_names=[metric_name],
167 reduced=self.reduced,
168 baseline=None if self.baseline is None else self.baseline[i],
169 has_metric_axis=False,
170 patch_type_names=self.patch_type_names,
171 )
173 def _require_per_example(self, what: str) -> None:
174 if self.reduced:
175 raise ValueError(f"{what} needs per-example values; run the sweep with reduce='none'")
177 def mean(self) -> torch.Tensor:
178 """Batch mean per cell (``values`` itself when already reduced)."""
179 return self.values if self.reduced else self.values.mean(dim=-1)
181 def stderr(self) -> torch.Tensor:
182 """Standard error of the batch mean per cell."""
183 self._require_per_example("stderr")
184 return standard_error(self.values, dim=-1)
186 def _broadcast_baseline(self) -> torch.Tensor:
187 if self.baseline is None:
188 raise ValueError("effect needs a baseline; run the sweep with baseline=True")
189 lead = 1 if self.has_metric_axis else 0
190 middle = self.values.ndim - lead - (0 if self.reduced else 1)
191 shape = (
192 tuple(self.baseline.shape[:lead]) + (1,) * middle + tuple(self.baseline.shape[lead:])
193 )
194 return self.baseline.reshape(shape)
196 def effect(self) -> torch.Tensor:
197 """``values - baseline``: the change each patch makes to the corrupted run, per example."""
198 return self.values - self._broadcast_baseline()
200 def normalized(self, clean_baseline: torch.Tensor) -> torch.Tensor:
201 """``(values - corrupted) / (clean - corrupted)`` per example.
203 ``clean_baseline`` is the same metric evaluated on the unpatched clean run,
204 shaped like ``baseline``. 0 means the patch did nothing, 1 means it
205 restored the clean metric. A prompt whose clean and corrupted metrics
206 coincide has no recovery scale and yields ``inf`` / ``nan`` for that
207 example; filter such prompts out before normalizing.
208 """
209 corrupted = self._broadcast_baseline()
210 baseline = self.baseline
211 assert baseline is not None # _broadcast_baseline raised otherwise
212 clean = torch.as_tensor(clean_baseline, dtype=self.values.dtype, device=self.values.device)
213 if clean.shape != baseline.shape:
214 raise ValueError(
215 f"clean_baseline must match baseline shape {tuple(baseline.shape)}, "
216 f"got {tuple(clean.shape)}"
217 )
218 clean = clean.reshape(corrupted.shape)
219 return (self.values - corrupted) / (clean - corrupted)
221 def bootstrap_ci(
222 self,
223 *,
224 confidence: float = 0.95,
225 n_resamples: int = 1000,
226 seed: int = 0,
227 statistic: str = "mean",
228 of_effect: bool = False,
229 ) -> Tuple[torch.Tensor, torch.Tensor]:
230 """Percentile bootstrap interval of the batch ``statistic`` per cell.
232 ``of_effect=True`` bootstraps ``effect()`` instead of ``values``. Seeded
233 through :func:`~transformer_lens.utilities.statistics.derive_generator`.
234 """
235 self._require_per_example("bootstrap_ci")
236 data = self.effect() if of_effect else self.values
237 return bootstrap_ci(
238 data,
239 dim=-1,
240 statistic=statistic,
241 confidence=confidence,
242 n_resamples=n_resamples,
243 generator=derive_generator(seed, "bootstrap", statistic),
244 )
246 def permutation_pvalue(
247 self,
248 *,
249 n_permutations: int = 1000,
250 seed: int = 0,
251 alternative: str = "two-sided",
252 ) -> torch.Tensor:
253 """Sign-flip permutation p-value that each cell's mean effect is zero.
255 Uses ``effect()`` (needs ``baseline=True``); exact for 12 or fewer
256 examples, Monte Carlo above that.
257 """
258 self._require_per_example("permutation_pvalue")
259 return sign_flip_permutation_pvalue(
260 self.effect(),
261 dim=-1,
262 n_permutations=n_permutations,
263 generator=derive_generator(seed, "permutation"),
264 alternative=alternative,
265 )
267 def to_dataframe(self) -> pd.DataFrame:
268 """Long form: one row per (metric, swept cell, example) with the value."""
269 rows = []
270 n_metrics = len(self.metric_names)
271 values = self.values if self.has_metric_axis else self.values.unsqueeze(0)
272 flat_cells = values.reshape(n_metrics, len(self.index_df), -1).detach().cpu()
273 for m, name in enumerate(self.metric_names):
274 for c, (_, index_row) in enumerate(self.index_df.iterrows()):
275 for b in range(flat_cells.shape[-1]):
276 row = {"metric": name, **index_row.to_dict()}
277 if self.patch_type_names is not None and "patch_type" in row:
278 row["patch_type_name"] = self.patch_type_names[int(row["patch_type"])]
279 if not self.reduced: 279 ↛ 281line 279 didn't jump to line 281 because the condition on line 279 was always true
280 row["example"] = b
281 row["value"] = flat_cells[m, c, b].item()
282 rows.append(row)
283 return pd.DataFrame(rows)
286@overload
287def generic_activation_patch(
288 model: TransformerLensModel,
289 corrupted_tokens: Int[torch.Tensor, "batch pos"],
290 clean_cache: ActivationCache,
291 patching_metric: MetricFn,
292 patch_setter: Callable[
293 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
294 ],
295 activation_name: str,
296 index_axis_names: Optional[Sequence[AxisNames]] = None,
297 index_df: Optional[pd.DataFrame] = None,
298 return_index_df: Literal[False] = False,
299 *,
300 reduce: Literal["scalar", "mean"] = "scalar",
301 baseline: Literal[False] = False,
302) -> torch.Tensor:
303 ...
306@overload
307def generic_activation_patch(
308 model: TransformerLensModel,
309 corrupted_tokens: Int[torch.Tensor, "batch pos"],
310 clean_cache: ActivationCache,
311 patching_metric: MetricFn,
312 patch_setter: Callable[
313 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
314 ],
315 activation_name: str,
316 index_axis_names: Optional[Sequence[AxisNames]],
317 index_df: Optional[pd.DataFrame],
318 return_index_df: Literal[True],
319 *,
320 reduce: Literal["scalar", "mean"] = "scalar",
321 baseline: Literal[False] = False,
322) -> Tuple[torch.Tensor, pd.DataFrame]:
323 ...
326@overload
327def generic_activation_patch(
328 model: TransformerLensModel,
329 corrupted_tokens: Int[torch.Tensor, "batch pos"],
330 clean_cache: ActivationCache,
331 patching_metric: Union[MetricFn, Mapping[str, MetricFn]],
332 patch_setter: Callable[
333 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
334 ],
335 activation_name: str,
336 index_axis_names: Optional[Sequence[AxisNames]] = None,
337 index_df: Optional[pd.DataFrame] = None,
338 return_index_df: bool = False,
339 *,
340 reduce: Literal["none"],
341 baseline: bool = False,
342) -> PatchingResult:
343 ...
346@overload
347def generic_activation_patch(
348 model: TransformerLensModel,
349 corrupted_tokens: Int[torch.Tensor, "batch pos"],
350 clean_cache: ActivationCache,
351 patching_metric: MetricFn,
352 patch_setter: Callable[
353 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
354 ],
355 activation_name: str,
356 index_axis_names: Optional[Sequence[AxisNames]] = None,
357 index_df: Optional[pd.DataFrame] = None,
358 return_index_df: bool = False,
359 *,
360 reduce: Reduce = "scalar",
361 baseline: Literal[True],
362) -> PatchingResult:
363 ...
366@overload
367def generic_activation_patch(
368 model: TransformerLensModel,
369 corrupted_tokens: Int[torch.Tensor, "batch pos"],
370 clean_cache: ActivationCache,
371 patching_metric: Mapping[str, MetricFn],
372 patch_setter: Callable[
373 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
374 ],
375 activation_name: str,
376 index_axis_names: Optional[Sequence[AxisNames]] = None,
377 index_df: Optional[pd.DataFrame] = None,
378 return_index_df: bool = False,
379 *,
380 reduce: Reduce = "scalar",
381 baseline: bool = False,
382) -> PatchingResult:
383 ...
386def generic_activation_patch(
387 model: TransformerLensModel,
388 corrupted_tokens: Int[torch.Tensor, "batch pos"],
389 clean_cache: ActivationCache,
390 patching_metric: Union[MetricFn, Mapping[str, MetricFn]],
391 patch_setter: Callable[
392 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
393 ],
394 activation_name: str,
395 index_axis_names: Optional[Sequence[AxisNames]] = None,
396 index_df: Optional[pd.DataFrame] = None,
397 return_index_df: bool = False,
398 *,
399 reduce: str = "scalar", # Reduce; widened so the ValueError below fires before beartype
400 baseline: bool = False,
401) -> Union[torch.Tensor, Tuple[torch.Tensor, pd.DataFrame], PatchingResult]:
402 """
403 A generic function to do activation patching, will be specialised to specific use cases.
405 Activation patching is about studying the counterfactual effect of a specific activation between a clean run and a corrupted run. The idea is have two inputs, clean and corrupted, which have two different outputs, and differ in some key detail. Eg "The Eiffel Tower is in" vs "The Colosseum is in". Then to take a cached set of activations from the "clean" run, and a set of corrupted.
407 Internally, the key function comes from three things: A list of tuples of indices (eg (layer, position, head_index)), a index_to_act_name function which identifies the right activation for each index, a patch_setter function which takes the corrupted activation, the index and the clean cache, and a metric for how well the patched model has recovered.
409 The indices can either be given explicitly as a pandas dataframe, or by listing the relevant axis names and having them inferred from the tokens and the model config. It is assumed that the first column is always layer.
411 This function then iterates over every tuple of indices, does the relevant patch, and stores it
413 Every patched forward runs the whole corrupted batch, so per-example values are
414 available for free: ``reduce="none"`` keeps them (the metric must then return a
415 ``[batch]`` tensor) and returns a :class:`PatchingResult` with mean, standard
416 error, bootstrap intervals and permutation p-values on top; ``reduce="mean"``
417 averages a ``[batch]`` metric back to today's tensor; a mapping of metrics
418 evaluates all of them on each patched forward. ``baseline=True`` adds one
419 unpatched corrupted forward so effects and normalized recovery can be reported.
421 Args:
422 model: The relevant model
423 corrupted_tokens: The input tokens for the corrupted run
424 clean_cache: The cached activations from the clean run
425 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc), or a mapping of names to such functions. Must return a scalar under ``reduce="scalar"`` and a ``[batch]`` tensor otherwise.
426 patch_setter: A function which acts on (corrupted_activation, index, clean_cache) to edit the activation and patch in the relevant chunk of the clean activation
427 activation_name: The name of the activation being patched
428 index_axis_names: The names of the axes to (fully) iterate over, implicitly fills in index_df
429 index_df: The dataframe of indices, columns are axis names and each row is a tuple of indices. Will be inferred from index_axis_names if not given. When this is input, the output will be a flattened tensor with an element per row of index_df
430 return_index_df: A Boolean flag for whether to return the dataframe of indices too (tensor return only; a ``PatchingResult`` always carries ``index_df``)
431 reduce: ``"scalar"`` (default, metric returns a scalar, tensor return), ``"mean"`` (metric returns ``[batch]``, averaged, tensor return) or ``"none"`` (metric returns ``[batch]``, kept, ``PatchingResult`` return)
432 baseline: Also evaluate the metric on the unpatched corrupted run and return a ``PatchingResult`` carrying it
434 Returns:
435 patched_output: The tensor of the patching metric for each patch. By default it has one dimension for each index dimension, via index_df set explicitly it is flattened with one element per row.
436 index_df *optional*: The dataframe of indices
437 Or a :class:`PatchingResult` when ``reduce="none"``, ``baseline=True`` or a metric mapping was given.
438 """
439 # Lazy import: the tools.analysis package imports this module, so top level would be circular.
440 from transformer_lens.tools.analysis._model_state import require_eval_mode
442 require_eval_mode(model, operation="activation patching")
443 if reduce not in _REDUCE_MODES:
444 raise ValueError(f"reduce must be one of {_REDUCE_MODES}, got {reduce!r}")
445 metrics: Dict[str, MetricFn]
446 if isinstance(patching_metric, Mapping):
447 has_metric_axis = True
448 metrics = dict(patching_metric)
449 if not metrics: 449 ↛ 450line 449 didn't jump to line 450 because the condition on line 449 was never true
450 raise ValueError("patching_metric mapping must contain at least one metric")
451 else:
452 has_metric_axis = False
453 metrics = {"metric": patching_metric}
454 want_result = has_metric_axis or reduce == "none" or baseline
455 if want_result and return_index_df:
456 raise ValueError(
457 "return_index_df only applies to the tensor return; a PatchingResult carries index_df"
458 )
460 if index_df is None:
461 assert index_axis_names is not None
462 number_of_heads = model.cfg.n_heads
463 # For some models, the number of key value heads is not the same as the number of attention heads
464 if activation_name in ["k", "v"] and model.cfg.n_key_value_heads is not None:
465 number_of_heads = model.cfg.n_key_value_heads
467 # Get the max range for all possible axes
468 max_axis_range = {
469 "layer": model.cfg.n_layers,
470 "pos": corrupted_tokens.shape[-1],
471 "head_index": number_of_heads,
472 }
473 max_axis_range["src_pos"] = max_axis_range["pos"]
474 max_axis_range["dest_pos"] = max_axis_range["pos"]
475 max_axis_range["head"] = max_axis_range["head_index"]
477 # Get the max range for each axis we iterate over
478 index_axis_max_range = [max_axis_range[axis_name] for axis_name in index_axis_names]
480 # Get the dataframe where each row is a tuple of indices
481 index_df = make_df_from_ranges(index_axis_max_range, index_axis_names)
483 flattened_output = False
484 else:
485 # A dataframe of indices was provided. Verify that we did not *also* receive index_axis_names
486 assert index_axis_names is None
487 index_axis_max_range = index_df.max().to_list()
488 flattened_output = True
490 batch_size = int(corrupted_tokens.shape[0])
491 per_example = reduce != "scalar"
492 cell_shape = (len(index_df),) if flattened_output else tuple(index_axis_max_range)
493 store_shape = (len(metrics),) + cell_shape + ((batch_size,) if per_example else ())
494 # Create an empty tensor to show the patched metric for each patch
495 store = torch.zeros(store_shape, device=model.cfg.device)
497 def evaluate(logits: torch.Tensor) -> torch.Tensor:
498 """Every metric on one set of logits, shaped [metric, batch?] after validation."""
499 rows = []
500 for name, fn in metrics.items():
501 out = torch.as_tensor(fn(logits))
502 if per_example:
503 if out.ndim != 1 or out.shape[0] != batch_size:
504 raise ValueError(
505 f"patching_metric {name!r} must return a [batch] tensor of length "
506 f"{batch_size} under reduce={reduce!r}, got shape {tuple(out.shape)}"
507 )
508 elif out.ndim != 0:
509 raise ValueError(
510 f"patching_metric {name!r} must return a scalar under reduce='scalar', got "
511 f"shape {tuple(out.shape)}; return a [batch] tensor and pass reduce='mean' "
512 "or reduce='none' to keep per-example values"
513 )
514 rows.append(out.detach())
515 return torch.stack(rows).to(store.dtype)
517 baseline_values: Optional[torch.Tensor] = None
518 if baseline: 518 ↛ 523line 518 didn't jump to line 523 because the condition on line 518 was always true
519 # Same entry point as the patched runs, with no hooks attached.
520 baseline_values = evaluate(model.run_with_hooks(corrupted_tokens, fwd_hooks=[]))
522 # A generic patching hook - for each index, it applies the patch_setter appropriately to patch the activation
523 def patching_hook(corrupted_activation, hook, index, clean_activation):
524 if corrupted_activation.requires_grad: 524 ↛ 526line 524 didn't jump to line 526 because the condition on line 524 was always true
525 corrupted_activation = corrupted_activation.clone()
526 return patch_setter(corrupted_activation, index, clean_activation)
528 for c, index_row in enumerate(tqdm((list(index_df.iterrows())))):
529 index = index_row[1].to_list()
531 # The current activation name is just the activation name plus the layer (assumed to be the first element of the input)
532 current_activation_name = utils.get_act_name(activation_name, layer=index[0])
534 # The hook function cannot receive additional inputs, so we use partial to include the specific index and the corresponding clean activation
535 current_hook = partial(
536 patching_hook,
537 index=index,
538 clean_activation=clean_cache[current_activation_name],
539 )
541 patched_logits = model.run_with_hooks(
542 corrupted_tokens, fwd_hooks=[(current_activation_name, current_hook)]
543 )
544 values = evaluate(patched_logits)
545 if flattened_output:
546 store[:, c] = values
547 else:
548 store[(slice(None), *index)] = values
550 if reduce == "mean":
551 store = store.mean(dim=-1)
552 if baseline_values is not None:
553 baseline_values = baseline_values.mean(dim=-1)
555 if not want_result:
556 patched_metric_output = store[0]
557 if return_index_df:
558 return patched_metric_output, index_df
559 return patched_metric_output
561 return PatchingResult(
562 values=store if has_metric_axis else store[0],
563 axis_names=[] if flattened_output else list(index_axis_names or []),
564 index_df=index_df,
565 metric_names=list(metrics),
566 reduced=reduce != "none",
567 baseline=None
568 if baseline_values is None
569 else (baseline_values if has_metric_axis else baseline_values[0]),
570 has_metric_axis=has_metric_axis,
571 )
574# %%
575# Defining patch setters for various shapes of activations
576def layer_pos_patch_setter(corrupted_activation, index, clean_activation):
577 """
578 Applies the activation patch where index = [layer, pos]
580 Implicitly assumes that the activation axis order is [batch, pos, ...], which is true of everything that is not an attention pattern shaped tensor.
581 """
582 assert len(index) == 2
583 layer, pos = index
584 corrupted_activation[:, pos, ...] = clean_activation[:, pos, ...]
585 return corrupted_activation
588def layer_pos_head_vector_patch_setter(
589 corrupted_activation,
590 index,
591 clean_activation,
592):
593 """
594 Applies the activation patch where index = [layer, pos, head_index]
596 Implicitly assumes that the activation axis order is [batch, pos, head_index, ...], which is true of all attention head vector activations (q, k, v, z, result) but *not* of attention patterns.
597 """
598 assert len(index) == 3
599 layer, pos, head_index = index
600 corrupted_activation[:, pos, head_index] = clean_activation[:, pos, head_index]
601 return corrupted_activation
604def layer_head_vector_patch_setter(
605 corrupted_activation,
606 index,
607 clean_activation,
608):
609 """
610 Applies the activation patch where index = [layer, head_index]
612 Implicitly assumes that the activation axis order is [batch, pos, head_index, ...], which is true of all attention head vector activations (q, k, v, z, result) but *not* of attention patterns.
613 """
614 assert len(index) == 2
615 layer, head_index = index
616 corrupted_activation[:, :, head_index] = clean_activation[:, :, head_index]
618 return corrupted_activation
621def layer_head_pattern_patch_setter(
622 corrupted_activation,
623 index,
624 clean_activation,
625):
626 """
627 Applies the activation patch where index = [layer, head_index]
629 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns.
630 """
631 assert len(index) == 2
632 layer, head_index = index
633 corrupted_activation[:, head_index, :, :] = clean_activation[:, head_index, :, :]
635 return corrupted_activation
638def layer_head_pos_pattern_patch_setter(
639 corrupted_activation,
640 index,
641 clean_activation,
642):
643 """
644 Applies the activation patch where index = [layer, head_index, dest_pos]
646 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns.
647 """
648 assert len(index) == 3
649 layer, head_index, dest_pos = index
650 corrupted_activation[:, head_index, dest_pos, :] = clean_activation[:, head_index, dest_pos, :]
652 return corrupted_activation
655def layer_head_dest_src_pos_pattern_patch_setter(
656 corrupted_activation,
657 index,
658 clean_activation,
659):
660 """
661 Applies the activation patch where index = [layer, head_index, dest_pos, src_pos]
663 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns.
664 """
665 assert len(index) == 4
666 layer, head_index, dest_pos, src_pos = index
667 corrupted_activation[:, head_index, dest_pos, src_pos] = clean_activation[
668 :, head_index, dest_pos, src_pos
669 ]
671 return corrupted_activation
674# %%
675# Defining activation patching functions for a range of common activation patches.
676get_act_patch_resid_pre = partial(
677 generic_activation_patch,
678 patch_setter=layer_pos_patch_setter,
679 activation_name="resid_pre",
680 index_axis_names=("layer", "pos"),
681)
682get_act_patch_resid_pre.__doc__ = """
683 Function to get activation patching results for the residual stream (at the start of each block) (by position). Returns a tensor of shape [n_layers, pos]
685 See generic_activation_patch for a more detailed explanation of activation patching
687 Args:
688 model: The relevant model
689 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
690 clean_cache (ActivationCache): The cached activations from the clean run
691 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
693 Returns:
694 patched_output (torch.Tensor): The tensor of the patching metric for each resid_pre patch. Has shape [n_layers, pos]
695 """
697get_act_patch_resid_mid = partial(
698 generic_activation_patch,
699 patch_setter=layer_pos_patch_setter,
700 activation_name="resid_mid",
701 index_axis_names=("layer", "pos"),
702)
703get_act_patch_resid_mid.__doc__ = """
704 Function to get activation patching results for the residual stream (between the attn and MLP layer of each block) (by position). Returns a tensor of shape [n_layers, pos]
706 See generic_activation_patch for a more detailed explanation of activation patching
708 Args:
709 model: The relevant model
710 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
711 clean_cache (ActivationCache): The cached activations from the clean run
712 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
714 Returns:
715 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos]
716 """
718get_act_patch_attn_out = partial(
719 generic_activation_patch,
720 patch_setter=layer_pos_patch_setter,
721 activation_name="attn_out",
722 index_axis_names=("layer", "pos"),
723)
724get_act_patch_attn_out.__doc__ = """
725 Function to get activation patching results for the output of each Attention layer (by position). Returns a tensor of shape [n_layers, pos]
727 See generic_activation_patch for a more detailed explanation of activation patching
729 Args:
730 model: The relevant model
731 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
732 clean_cache (ActivationCache): The cached activations from the clean run
733 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
735 Returns:
736 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos]
737 """
739get_act_patch_mlp_out = partial(
740 generic_activation_patch,
741 patch_setter=layer_pos_patch_setter,
742 activation_name="mlp_out",
743 index_axis_names=("layer", "pos"),
744)
745get_act_patch_mlp_out.__doc__ = """
746 Function to get activation patching results for the output of each MLP layer (by position). Returns a tensor of shape [n_layers, pos]
748 See generic_activation_patch for a more detailed explanation of activation patching
750 Args:
751 model: The relevant model
752 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
753 clean_cache (ActivationCache): The cached activations from the clean run
754 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
756 Returns:
757 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos]
758 """
759# %%
760get_act_patch_attn_head_out_by_pos = partial(
761 generic_activation_patch,
762 patch_setter=layer_pos_head_vector_patch_setter,
763 activation_name="z",
764 index_axis_names=("layer", "pos", "head"),
765)
766get_act_patch_attn_head_out_by_pos.__doc__ = """
767 Function to get activation patching results for the output of each Attention Head (by position). Returns a tensor of shape [n_layers, pos, n_heads]
769 See generic_activation_patch for a more detailed explanation of activation patching
771 Args:
772 model: The relevant model
773 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
774 clean_cache (ActivationCache): The cached activations from the clean run
775 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
777 Returns:
778 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads]
779 """
781get_act_patch_attn_head_q_by_pos = partial(
782 generic_activation_patch,
783 patch_setter=layer_pos_head_vector_patch_setter,
784 activation_name="q",
785 index_axis_names=("layer", "pos", "head"),
786)
787get_act_patch_attn_head_q_by_pos.__doc__ = """
788 Function to get activation patching results for the queries of each Attention Head (by position). Returns a tensor of shape [n_layers, pos, n_heads]
790 See generic_activation_patch for a more detailed explanation of activation patching
792 Args:
793 model: The relevant model
794 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
795 clean_cache (ActivationCache): The cached activations from the clean run
796 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
798 Returns:
799 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads]
800 """
802get_act_patch_attn_head_k_by_pos = partial(
803 generic_activation_patch,
804 patch_setter=layer_pos_head_vector_patch_setter,
805 activation_name="k",
806 index_axis_names=("layer", "pos", "head"),
807)
808get_act_patch_attn_head_k_by_pos.__doc__ = """
809 Function to get activation patching results for the keys of each Attention Head (by position). Returns a tensor of shape [n_layers, pos, n_heads] or [n_layers, pos, n_key_value_heads] if the model has a different number of key value heads than attention heads.
811 See generic_activation_patch for a more detailed explanation of activation patching
813 Args:
814 model: The relevant model
815 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
816 clean_cache (ActivationCache): The cached activations from the clean run
817 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
819 Returns:
820 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads] or [n_layers, pos, n_key_value_heads] if the model has a different number of key value heads than attention heads.
821 """
823get_act_patch_attn_head_v_by_pos = partial(
824 generic_activation_patch,
825 patch_setter=layer_pos_head_vector_patch_setter,
826 activation_name="v",
827 index_axis_names=("layer", "pos", "head"),
828)
829get_act_patch_attn_head_v_by_pos.__doc__ = """
830 Function to get activation patching results for the values of each Attention Head (by position). Returns a tensor of shape [n_layers, pos, n_heads] or [n_layers, pos, n_key_value_heads] if the model has a different number of key value heads than attention heads.
832 See generic_activation_patch for a more detailed explanation of activation patching
834 Args:
835 model: The relevant model
836 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
837 clean_cache (ActivationCache): The cached activations from the clean run
838 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
840 Returns:
841 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads] or [n_layers, pos, n_key_value_heads] if the model has a different number of key value heads than attention heads.
842 """
843# %%
844get_act_patch_attn_head_pattern_by_pos = partial(
845 generic_activation_patch,
846 patch_setter=layer_head_pos_pattern_patch_setter,
847 activation_name="pattern",
848 index_axis_names=("layer", "head_index", "dest_pos"),
849)
850get_act_patch_attn_head_pattern_by_pos.__doc__ = """
851 Function to get activation patching results for the attention pattern of each Attention Head (by destination position). Returns a tensor of shape [n_layers, n_heads, dest_pos]
853 See generic_activation_patch for a more detailed explanation of activation patching
855 Args:
856 model: The relevant model
857 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
858 clean_cache (ActivationCache): The cached activations from the clean run
859 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
861 Returns:
862 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads, dest_pos]
863 """
865get_act_patch_attn_head_pattern_dest_src_pos = partial(
866 generic_activation_patch,
867 patch_setter=layer_head_dest_src_pos_pattern_patch_setter,
868 activation_name="pattern",
869 index_axis_names=("layer", "head_index", "dest_pos", "src_pos"),
870)
871get_act_patch_attn_head_pattern_dest_src_pos.__doc__ = """
872 Function to get activation patching results for each destination, source entry of the attention pattern for each Attention Head. Returns a tensor of shape [n_layers, n_heads, dest_pos, src_pos]
874 See generic_activation_patch for a more detailed explanation of activation patching
876 Args:
877 model: The relevant model
878 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
879 clean_cache (ActivationCache): The cached activations from the clean run
880 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
882 Returns:
883 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads, dest_pos, src_pos]
884 """
886# %%
887get_act_patch_attn_head_out_all_pos = partial(
888 generic_activation_patch,
889 patch_setter=layer_head_vector_patch_setter,
890 activation_name="z",
891 index_axis_names=("layer", "head"),
892)
893get_act_patch_attn_head_out_all_pos.__doc__ = """
894 Function to get activation patching results for the outputs of each Attention Head (across all positions). Returns a tensor of shape [n_layers, n_heads]
896 See generic_activation_patch for a more detailed explanation of activation patching
898 Args:
899 model: The relevant model
900 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
901 clean_cache (ActivationCache): The cached activations from the clean run
902 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
904 Returns:
905 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads]
906 """
908get_act_patch_attn_head_q_all_pos = partial(
909 generic_activation_patch,
910 patch_setter=layer_head_vector_patch_setter,
911 activation_name="q",
912 index_axis_names=("layer", "head"),
913)
914get_act_patch_attn_head_q_all_pos.__doc__ = """
915 Function to get activation patching results for the queries of each Attention Head (across all positions). Returns a tensor of shape [n_layers, n_heads]
917 See generic_activation_patch for a more detailed explanation of activation patching
919 Args:
920 model: The relevant model
921 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
922 clean_cache (ActivationCache): The cached activations from the clean run
923 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
925 Returns:
926 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads]
927 """
929get_act_patch_attn_head_k_all_pos = partial(
930 generic_activation_patch,
931 patch_setter=layer_head_vector_patch_setter,
932 activation_name="k",
933 index_axis_names=("layer", "head"),
934)
935get_act_patch_attn_head_k_all_pos.__doc__ = """
936 Function to get activation patching results for the keys of each Attention Head (across all positions). Returns a tensor of shape [n_layers, n_heads] or [n_layers, n_key_value_heads] if the model has a different number of key value heads than attention heads.
938 See generic_activation_patch for a more detailed explanation of activation patching
940 Args:
941 model: The relevant model
942 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
943 clean_cache (ActivationCache): The cached activations from the clean run
944 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
946 Returns:
947 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads] or [n_layers, n_key_value_heads] if the model has a different number of key value heads than attention heads.
948 """
950get_act_patch_attn_head_v_all_pos = partial(
951 generic_activation_patch,
952 patch_setter=layer_head_vector_patch_setter,
953 activation_name="v",
954 index_axis_names=("layer", "head"),
955)
956get_act_patch_attn_head_v_all_pos.__doc__ = """
957 Function to get activation patching results for the values of each Attention Head (across all positions). Returns a tensor of shape [n_layers, n_heads] or [n_layers, n_key_value_heads] if the model has a different number of key value heads than attention heads.
959 See generic_activation_patch for a more detailed explanation of activation patching
961 Args:
962 model: The relevant model
963 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
964 clean_cache (ActivationCache): The cached activations from the clean run
965 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
967 Returns:
968 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads] or [n_layers, n_key_value_heads] if the model has a different number of key value heads than attention heads.
969 """
971get_act_patch_attn_head_pattern_all_pos = partial(
972 generic_activation_patch,
973 patch_setter=layer_head_pattern_patch_setter,
974 activation_name="pattern",
975 index_axis_names=("layer", "head_index"),
976)
977get_act_patch_attn_head_pattern_all_pos.__doc__ = """
978 Function to get activation patching results for the attention pattern of each Attention Head (across all positions). Returns a tensor of shape [n_layers, n_heads]
980 See generic_activation_patch for a more detailed explanation of activation patching
982 Args:
983 model: The relevant model
984 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
985 clean_cache (ActivationCache): The cached activations from the clean run
986 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
988 Returns:
989 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads]
990 """
992# %%
995_HEAD_TYPES = ("out", "q", "k", "v", "pattern")
996_BLOCK_TYPES = ("resid_pre", "attn_out", "mlp_out")
997_PER_EXAMPLE_NOTE = """
998 Keyword-only ``reduce`` ("scalar" | "mean" | "none") and ``baseline`` are forwarded to
999 generic_activation_patch; with reduce="none", baseline=True or a mapping of metrics the
1000 return is a PatchingResult carrying per-example values and statistics instead of a tensor.
1001"""
1002for _fn in (
1003 get_act_patch_resid_pre,
1004 get_act_patch_resid_mid,
1005 get_act_patch_attn_out,
1006 get_act_patch_mlp_out,
1007 get_act_patch_attn_head_out_by_pos,
1008 get_act_patch_attn_head_q_by_pos,
1009 get_act_patch_attn_head_k_by_pos,
1010 get_act_patch_attn_head_v_by_pos,
1011 get_act_patch_attn_head_pattern_by_pos,
1012 get_act_patch_attn_head_pattern_dest_src_pos,
1013 get_act_patch_attn_head_out_all_pos,
1014 get_act_patch_attn_head_q_all_pos,
1015 get_act_patch_attn_head_k_all_pos,
1016 get_act_patch_attn_head_v_all_pos,
1017 get_act_patch_attn_head_pattern_all_pos,
1018):
1019 _fn.__doc__ = (_fn.__doc__ or "") + _PER_EXAMPLE_NOTE
1022def _wrap_tensor_result(tensor: torch.Tensor, axis_names: Sequence[str]) -> PatchingResult:
1023 """Give a default tensor sweep the result shape so it can be stacked with typed ones."""
1024 return PatchingResult(
1025 values=tensor,
1026 axis_names=list(axis_names),
1027 index_df=make_df_from_ranges(list(tensor.shape), list(axis_names)),
1028 metric_names=["metric"],
1029 reduced=True,
1030 )
1033def _stack_patch_types(
1034 results: Sequence[Union[torch.Tensor, PatchingResult]],
1035 *,
1036 axis_order: Sequence[str],
1037 n_heads: int,
1038 type_names: Sequence[str],
1039) -> Union[torch.Tensor, PatchingResult]:
1040 """Stack per-activation sweeps on a leading ``patch_type`` axis.
1042 Tensors (the default return) are stacked as-is after the caller padded and
1043 rearranged them. ``PatchingResult`` values are first permuted to
1044 ``axis_order`` (``head_index`` / ``dest_pos`` renamed to ``head`` / ``pos``)
1045 and the head axis zero-padded to ``n_heads`` for GQA ``k``/``v`` sweeps,
1046 keeping the metric axis first and the batch axis last. ``index_df`` is
1047 rebuilt from the padded grid so it stays one row per cell of ``values``;
1048 padded GQA cells therefore appear as zero-valued rows.
1049 """
1050 if all(isinstance(r, torch.Tensor) for r in results): 1050 ↛ 1051line 1050 didn't jump to line 1051 because the condition on line 1050 was never true
1051 return torch.stack([r for r in results if isinstance(r, torch.Tensor)], dim=0)
1052 typed = [r for r in results if isinstance(r, PatchingResult)]
1053 if len(typed) != len(results): 1053 ↛ 1054line 1053 didn't jump to line 1054 because the condition on line 1053 was never true
1054 raise TypeError("cannot stack tensor and PatchingResult sweeps together")
1055 first = typed[0]
1056 for r in typed[1:]:
1057 if r.metric_names != first.metric_names or r.reduced != first.reduced: 1057 ↛ 1058line 1057 didn't jump to line 1058 because the condition on line 1057 was never true
1058 raise ValueError("every sweep must use the same metrics and reduce mode to be stacked")
1059 head_axis = list(axis_order).index("head") if "head" in axis_order else None
1060 values = []
1061 for r in typed:
1062 names = [_RESULT_AXIS_RENAMES.get(n, n) for n in r.axis_names]
1063 if sorted(names) != sorted(axis_order): 1063 ↛ 1064line 1063 didn't jump to line 1064 because the condition on line 1063 was never true
1064 raise ValueError(f"sweep axes {names} do not match {list(axis_order)}")
1065 lead = 1 if r.has_metric_axis else 0
1066 perm = list(range(lead)) + [lead + names.index(n) for n in axis_order]
1067 tail = list(range(lead + len(names), r.values.ndim))
1068 v = r.values.permute(*perm, *tail)
1069 pad = 0 if head_axis is None else n_heads - v.shape[lead + head_axis]
1070 if head_axis is not None and pad > 0:
1071 pad_shape = list(v.shape)
1072 pad_shape[lead + head_axis] = pad
1073 v = torch.cat([v, v.new_zeros(pad_shape)], dim=lead + head_axis)
1074 values.append(v)
1075 stacked = torch.stack(values, dim=1 if first.has_metric_axis else 0)
1076 lead = 1 if first.has_metric_axis else 0
1077 cell_shape = list(stacked.shape[lead : lead + 1 + len(axis_order)])
1078 return PatchingResult(
1079 values=stacked,
1080 axis_names=["patch_type", *axis_order],
1081 index_df=make_df_from_ranges(cell_shape, ["patch_type", *axis_order]),
1082 metric_names=list(first.metric_names),
1083 reduced=first.reduced,
1084 baseline=first.baseline,
1085 has_metric_axis=first.has_metric_axis,
1086 patch_type_names=list(type_names),
1087 )
1090def _run_sweeps(sweeps, axis_names_per_sweep, model, corrupted_tokens, clean_cache, metric, kwargs):
1091 """Run a helper's sweeps, computing a requested baseline once rather than once per sweep."""
1092 baseline = bool(kwargs.pop("baseline", False))
1093 results: List[Union[torch.Tensor, PatchingResult]] = []
1094 for i, (sweep, axis_names) in enumerate(zip(sweeps, axis_names_per_sweep)):
1095 out = sweep(
1096 model, corrupted_tokens, clean_cache, metric, baseline=baseline and i == 0, **kwargs
1097 )
1098 if baseline and isinstance(out, torch.Tensor):
1099 out = _wrap_tensor_result(out, axis_names)
1100 results.append(out)
1101 return results
1104def get_act_patch_attn_head_all_pos_every(
1105 model, corrupted_tokens, clean_cache, metric, **kwargs
1106) -> Union[Float[torch.Tensor, "patch_type layer head"], PatchingResult]:
1107 """Helper function to get activation patching results for every head (across all positions) for every act type (output, query, key, value, pattern). Wrapper around each's patching function, returns a stacked tensor of shape [5, n_layers, n_heads]
1109 Args:
1110 model: The relevant model
1111 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
1112 clean_cache (ActivationCache): The cached activations from the clean run
1113 metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
1114 **kwargs: ``reduce`` / ``baseline`` forwarded to each sweep; a ``PatchingResult`` is then returned with a leading ``patch_type`` axis (``patch_type_names`` = out, q, k, v, pattern)
1116 Returns:
1117 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [5, n_layers, n_heads]
1118 """
1119 results = _run_sweeps(
1120 (
1121 get_act_patch_attn_head_out_all_pos,
1122 get_act_patch_attn_head_q_all_pos,
1123 get_act_patch_attn_head_k_all_pos,
1124 get_act_patch_attn_head_v_all_pos,
1125 get_act_patch_attn_head_pattern_all_pos,
1126 ),
1127 (("layer", "head"),) * 4 + (("layer", "head_index"),),
1128 model,
1129 corrupted_tokens,
1130 clean_cache,
1131 metric,
1132 kwargs,
1133 )
1134 if all(isinstance(r, torch.Tensor) for r in results):
1135 tensors = [r for r in results if isinstance(r, torch.Tensor)]
1136 n_heads = tensors[0].size(-1)
1137 # Reshape k and v to be compatible with the rest of the results in case of n_key_value_heads != n_heads
1138 for i in (2, 3):
1139 tensors[i] = torch.nn.functional.pad(tensors[i], (0, n_heads - tensors[i].size(-1)))
1140 return torch.stack(tensors, dim=0)
1141 return _stack_patch_types(
1142 results, axis_order=("layer", "head"), n_heads=model.cfg.n_heads, type_names=_HEAD_TYPES
1143 )
1146def get_act_patch_attn_head_by_pos_every(
1147 model, corrupted_tokens, clean_cache, metric, **kwargs
1148) -> Union[Float[torch.Tensor, "patch_type layer pos head"], PatchingResult]:
1149 """Helper function to get activation patching results for every head (by position) for every act type (output, query, key, value, pattern). Wrapper around each's patching function, returns a stacked tensor of shape [5, n_layers, pos, n_heads]
1151 Args:
1152 model: The relevant model
1153 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
1154 clean_cache (ActivationCache): The cached activations from the clean run
1155 metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
1156 **kwargs: ``reduce`` / ``baseline`` forwarded to each sweep; a ``PatchingResult`` is then returned with a leading ``patch_type`` axis (``patch_type_names`` = out, q, k, v, pattern)
1158 Returns:
1159 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [5, n_layers, pos, n_heads]
1160 """
1161 results = _run_sweeps(
1162 (
1163 get_act_patch_attn_head_out_by_pos,
1164 get_act_patch_attn_head_q_by_pos,
1165 get_act_patch_attn_head_k_by_pos,
1166 get_act_patch_attn_head_v_by_pos,
1167 get_act_patch_attn_head_pattern_by_pos,
1168 ),
1169 (("layer", "pos", "head"),) * 4 + (("layer", "head_index", "dest_pos"),),
1170 model,
1171 corrupted_tokens,
1172 clean_cache,
1173 metric,
1174 kwargs,
1175 )
1176 if all(isinstance(r, torch.Tensor) for r in results):
1177 tensors = [r for r in results if isinstance(r, torch.Tensor)]
1178 n_heads = tensors[0].size(-1)
1179 # Reshape k and v to be compatible with the rest of the results in case of n_key_value_heads != n_heads
1180 for i in (2, 3):
1181 tensors[i] = torch.nn.functional.pad(tensors[i], (0, n_heads - tensors[i].size(-1)))
1182 # Reshape pattern to be compatible with the rest of the results
1183 tensors[4] = einops.rearrange(tensors[4], "batch head pos -> batch pos head")
1184 return torch.stack(tensors, dim=0)
1185 return _stack_patch_types(
1186 results,
1187 axis_order=("layer", "pos", "head"),
1188 n_heads=model.cfg.n_heads,
1189 type_names=_HEAD_TYPES,
1190 )
1193def get_act_patch_block_every(
1194 model, corrupted_tokens, clean_cache, metric, **kwargs
1195) -> Union[Float[torch.Tensor, "patch_type layer pos"], PatchingResult]:
1196 """Helper function to get activation patching results for the residual stream (at the start of each block), output of each Attention layer and output of each MLP layer. Wrapper around each's patching function, returns a stacked tensor of shape [3, n_layers, pos]
1198 Args:
1199 model: The relevant model
1200 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
1201 clean_cache (ActivationCache): The cached activations from the clean run
1202 metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
1203 **kwargs: ``reduce`` / ``baseline`` forwarded to each sweep; a ``PatchingResult`` is then returned with a leading ``patch_type`` axis (``patch_type_names`` = resid_pre, attn_out, mlp_out)
1205 Returns:
1206 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [3, n_layers, pos]
1207 """
1208 results = _run_sweeps(
1209 (get_act_patch_resid_pre, get_act_patch_attn_out, get_act_patch_mlp_out),
1210 (("layer", "pos"),) * 3,
1211 model,
1212 corrupted_tokens,
1213 clean_cache,
1214 metric,
1215 kwargs,
1216 )
1217 if all(isinstance(r, torch.Tensor) for r in results):
1218 return torch.stack([r for r in results if isinstance(r, torch.Tensor)], dim=0)
1219 return _stack_patch_types(
1220 results, axis_order=("layer", "pos"), n_heads=0, type_names=_BLOCK_TYPES
1221 )
1224# %%
1225def _check_direct_path_support(model: Any, src_layers: Sequence[int], component: str) -> None:
1226 """Refuse missing attention, source writes other than ``z @ W_O``, and K/V on GQA/MQA."""
1227 # Lazy import: the tools.analysis package imports this module, so top level would be circular.
1228 from transformer_lens.tools.analysis.attribution_patching import (
1229 _attention_write_obstruction,
1230 )
1232 cfg = model.cfg
1233 src_layer = min(src_layers)
1234 for layer in range(src_layer, cfg.n_layers):
1235 if getattr(model.blocks[layer], "attn", None) is None:
1236 raise ValueError(
1237 f"blocks.{layer} has no `attn` component (an SSM, mixer, linear-attention or "
1238 "AltUp block), and direct path patching needs one in every block from "
1239 f"src_layer={src_layer} on."
1240 )
1241 detail = _attention_write_obstruction(model, set(src_layers))
1242 if detail is not None:
1243 raise ValueError(
1244 "per-head attention outputs are normalized or scaled before the residual add on "
1245 f"this architecture ({detail}), so a source head's residual write is not "
1246 "hook_z @ W_O and direct path patching is unsupported here."
1247 )
1248 if component == "q":
1249 return
1250 n_heads = cfg.n_heads
1251 scalar_kv = getattr(cfg, "n_key_value_heads", None)
1252 per_layer_kv = getattr(cfg, "per_layer_num_key_value_heads", None) or []
1253 # Only destinations are K/V-patched; the source is read via hook_z and W_O, per query head.
1254 for layer in range(src_layer + 1, cfg.n_layers):
1255 kv = per_layer_kv[layer] if layer < len(per_layer_kv) else None
1256 if kv is None:
1257 kv = scalar_kv
1258 if kv is not None and kv != n_heads:
1259 raise ValueError(
1260 f"component={component!r} is unsupported here: blocks.{layer} has "
1261 f"n_key_value_heads={kv} for n_heads={n_heads} (GQA/MQA), so per-query-head "
1262 'K/V semantics are undefined. component="q" is available.'
1263 )
1266@contextmanager
1267def _split_qkv_input(model: Any) -> Iterator[None]:
1268 """Enable the per-head pre-norm Q/K/V forks, restoring the caller's flags on exit."""
1269 saved_split = bool(getattr(model.cfg, "use_split_qkv_input", False))
1270 saved_attn_in = bool(getattr(model.cfg, "use_attn_in", False))
1271 try:
1272 # The two forks are mutually exclusive, so use_attn_in must go off first.
1273 if saved_attn_in:
1274 model.set_use_attn_in(False)
1275 if not saved_split:
1276 model.set_use_split_qkv_input(True)
1277 yield
1278 finally:
1279 model.set_use_split_qkv_input(saved_split)
1280 if saved_attn_in:
1281 model.set_use_attn_in(True)
1284def _make_direct_path_hook(
1285 delta: Float[torch.Tensor, "batch pos d_model"], dst_head: int, fired: List[bool]
1286) -> Callable[..., torch.Tensor]:
1287 """Add the source head's clean-minus-corrupted write to one head's fork input."""
1289 def hook_fn(
1290 fork_input: Float[torch.Tensor, "batch pos n_heads d_model"], hook
1291 ) -> Float[torch.Tensor, "batch pos n_heads d_model"]:
1292 fired.append(True)
1293 patched = fork_input.clone()
1294 patched[:, :, dst_head] += delta
1295 return patched
1297 return hook_fn
1300def get_act_patch_direct_path(
1301 model: TransformerLensModel,
1302 corrupted_tokens: torch.Tensor,
1303 clean_cache: ActivationCache,
1304 corrupted_cache: ActivationCache,
1305 patching_metric: Callable[[torch.Tensor], torch.Tensor],
1306 src_layer: int,
1307 src_head: int,
1308 component: Literal["q", "k", "v"] = "q",
1309 verbose: bool = True,
1310) -> Float[torch.Tensor, "n_layers n_heads"]:
1311 """Patch one source head's output into each later head's Q, K or V input alone.
1313 Direct path patching (Wang et al. 2022) isolates the edge A -> B: A's clean-minus-corrupted
1314 write (``z @ W_O``) is added at every position to B's own pre-``ln1`` copy of the residual,
1315 while every other reader still sees A's corrupted output. The result is the exact edge
1316 intervention, folded LayerNorm or not.
1318 The sweep temporarily enables ``cfg.use_split_qkv_input`` (restoring the caller's flags on
1319 exit), which costs about ``3 * n_heads`` times the usual ``ln1`` activation memory per forward.
1321 Args:
1322 model: A ``TransformerBridge`` whose attention exposes the pre-Q/K/V fork.
1323 corrupted_tokens: Token ids for the corrupted run, ``[batch, pos]``.
1324 clean_cache: Cache of the clean run; must hold ``blocks.{src_layer}.attn.hook_z``.
1325 corrupted_cache: Cache of the corrupted run; must hold the same ``hook_z``.
1326 patching_metric: Maps the patched logits to a scalar.
1327 src_layer: Layer of the source head.
1328 src_head: Index of the source head.
1329 component: Destination input to patch: ``"q"`` (default), ``"k"`` or ``"v"``.
1331 Returns:
1332 ``[n_layers, n_heads]``: ``results[l, h]`` is the metric with the path into head
1333 ``(l, h)`` patched. Rows ``l <= src_layer`` are 0.0 placeholders, not measurements.
1335 Raises:
1336 ValueError: if a block from ``src_layer`` on has no attention, if the source block
1337 normalises, scales or joins its attention output before the residual add (post-norm,
1338 scaled-residual or parallel-mixer blocks), if ``component`` is ``"k"`` or ``"v"`` on
1339 a GQA/MQA model, or if any module is in training mode.
1340 NotImplementedError: if a destination layer's attention has no pre-Q/K/V fork.
1341 """
1342 # Lazy import: the tools.analysis package imports this module, so top level would be circular.
1343 from transformer_lens.tools.analysis._model_state import require_eval_mode
1345 _check_direct_path_support(model, [src_layer], component)
1346 require_eval_mode(model, operation="get_act_patch_direct_path()")
1348 n_layers = model.cfg.n_layers
1349 n_heads = model.cfg.n_heads
1350 results = torch.zeros(n_layers, n_heads, device=model.cfg.device)
1351 dst_pairs = [(dl, dh) for dl in range(src_layer + 1, n_layers) for dh in range(n_heads)]
1353 bridge = cast(Any, model)
1354 W_O = bridge.blocks[src_layer].attn.W_O[src_head]
1355 z_name = f"blocks.{src_layer}.attn.hook_z"
1356 clean_write = clean_cache[z_name][:, :, src_head] @ W_O
1357 delta = clean_write - corrupted_cache[z_name][:, :, src_head] @ W_O
1359 with _split_qkv_input(bridge):
1360 for dst_layer, dst_head in tqdm(
1361 dst_pairs,
1362 desc=f"Direct path patch ({src_layer},{src_head}) → * [{component}]",
1363 disable=not verbose,
1364 ):
1365 hook_name = f"blocks.{dst_layer}.attn.hook_{component}_input"
1366 fired: List[bool] = []
1367 logits = model.run_with_hooks(
1368 corrupted_tokens,
1369 fwd_hooks=[(hook_name, _make_direct_path_hook(delta, dst_head, fired))],
1370 )
1371 # Without a fork the hook never fires and the metric would silently be unpatched.
1372 if not fired:
1373 raise NotImplementedError(
1374 f"{hook_name} did not fire: blocks.{dst_layer}.attn has no pre-Q/K/V fork, so "
1375 "direct path patching cannot reach that head's input."
1376 )
1377 results[dst_layer, dst_head] = patching_metric(logits).item()
1378 return results
1381def get_act_patch_direct_path_all_sources(
1382 model: TransformerLensModel,
1383 corrupted_tokens: torch.Tensor,
1384 clean_cache: ActivationCache,
1385 corrupted_cache: ActivationCache,
1386 patching_metric: Callable[[torch.Tensor], torch.Tensor],
1387 component: Literal["q", "k", "v"] = "q",
1388 verbose: bool = True,
1389) -> Float[torch.Tensor, "n_layers n_heads n_layers n_heads"]:
1390 """Run :func:`get_act_patch_direct_path` from every source head.
1392 ``results[sl, sh, dl, dh]`` is the metric with the path from head ``(sl, sh)`` into head
1393 ``(dl, dh)`` patched, and 0.0 where ``dl <= sl``. One forward per (source, later destination)
1394 pair, so it suits small models or targeted sub-sweeps.
1395 """
1396 # Lazy import: the tools.analysis package imports this module, so top level would be circular.
1397 from transformer_lens.tools.analysis._model_state import require_eval_mode
1399 _check_direct_path_support(model, range(model.cfg.n_layers), component)
1400 require_eval_mode(model, operation="get_act_patch_direct_path_all_sources()")
1402 n_layers = model.cfg.n_layers
1403 n_heads = model.cfg.n_heads
1404 results = torch.zeros(n_layers, n_heads, n_layers, n_heads, device=model.cfg.device)
1405 src_pairs = [(sl, sh) for sl in range(n_layers) for sh in range(n_heads)]
1406 for src_layer, src_head in tqdm(
1407 src_pairs,
1408 desc=f"Direct path patch — all sources [{component}]",
1409 disable=not verbose,
1410 ):
1411 results[src_layer, src_head] = get_act_patch_direct_path(
1412 model=model,
1413 corrupted_tokens=corrupted_tokens,
1414 clean_cache=clean_cache,
1415 corrupted_cache=corrupted_cache,
1416 patching_metric=patching_metric,
1417 src_layer=src_layer,
1418 src_head=src_head,
1419 component=component,
1420 verbose=False,
1421 )
1422 return results