Coverage for transformer_lens/patching.py: 73%
142 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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.
8Context:
10Activation Patching is technique introduced in the `ROME paper <http://rome.baulab.info/>`, which
11uses a causal intervention to identify which activations in a model matter for producing some
12output. It runs the model on input A, replaces (patches) an activation with that same activation on
13input B, and sees how much that shifts the answer from A to B.
15More details: The setup of activation patching is to take two runs of the model on two different
16inputs, the clean run and the corrupted run. The clean run outputs the correct answer and the
17corrupted run does not. The key idea is that we give the model the corrupted input, but then
18intervene on a specific activation and patch in the corresponding activation from the clean run (ie
19replace the corrupted activation with the clean activation), and then continue the run. And we then
20measure how much the output has updated towards the correct answer.
22- We can then iterate over many
23 possible activations and look at how much they affect the corrupted run. If patching in an
24 activation significantly increases the probability of the correct answer, this allows us to
25 localise which activations matter.
26- A key detail is that we move a single activation __from__ the clean run __to __the corrupted run.
27 So if this changes the answer from incorrect to correct, we can be confident that the activation
28 moved was important.
30Intuition:
32The ability to **localise** is a key move in mechanistic interpretability - if the computation is
33diffuse and spread across the entire model, it is likely much harder to form a clean mechanistic
34story for what's going on. But if we can identify precisely which parts of the model matter, we can
35then zoom in and determine what they represent and how they connect up with each other, and
36ultimately reverse engineer the underlying circuit that they represent. And, empirically, on at
37least some tasks activation patching tends to find that computation is extremely localised:
39- This technique helps us precisely identify which parts of the model matter for a certain
40 part of a task. Eg, answering “The Eiffel Tower is in” with “Paris” requires figuring out that
41 the Eiffel Tower is in Paris, and that it’s a factual recall task and that the output is a
42 location. Patching to “The Colosseum is in” controls for everything other than the “Eiffel Tower
43 is located in Paris” feature.
44- It helps a lot if the corrupted prompt has the same number of tokens
46This, unlike direct logit attribution, can identify meaningful parts of a circuit from anywhere
47within the model, rather than just the end.
48"""
50from __future__ import annotations
52import itertools
53from functools import partial
54from typing import Callable, Optional, Sequence, Tuple, Union, overload
56import einops
57import pandas as pd
58import torch
59from jaxtyping import Float, Int
60from tqdm.auto import tqdm
61from typing_extensions import Literal
63import transformer_lens.utilities as utils
64from transformer_lens.ActivationCache import ActivationCache
65from transformer_lens.HookedTransformer import HookedTransformer
67# %%
68Logits = torch.Tensor
69AxisNames = Literal["layer", "pos", "head_index", "head", "src_pos", "dest_pos"]
72# %%
75def make_df_from_ranges(
76 column_max_ranges: Sequence[int], column_names: Sequence[str]
77) -> pd.DataFrame:
78 """
79 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)
80 """
81 rows = list(itertools.product(*[range(axis_max_range) for axis_max_range in column_max_ranges]))
82 df = pd.DataFrame(rows, columns=column_names)
83 return df
86# %%
87CorruptedActivation = torch.Tensor
88PatchedActivation = torch.Tensor
91@overload
92def generic_activation_patch(
93 model: HookedTransformer,
94 corrupted_tokens: Int[torch.Tensor, "batch pos"],
95 clean_cache: ActivationCache,
96 patching_metric: Callable[[Float[torch.Tensor, "batch pos d_vocab"]], Float[torch.Tensor, ""]],
97 patch_setter: Callable[
98 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
99 ],
100 activation_name: str,
101 index_axis_names: Optional[Sequence[AxisNames]] = None,
102 index_df: Optional[pd.DataFrame] = None,
103 return_index_df: Literal[False] = False,
104) -> torch.Tensor:
105 ...
108@overload
109def generic_activation_patch(
110 model: HookedTransformer,
111 corrupted_tokens: Int[torch.Tensor, "batch pos"],
112 clean_cache: ActivationCache,
113 patching_metric: Callable[[Float[torch.Tensor, "batch pos d_vocab"]], Float[torch.Tensor, ""]],
114 patch_setter: Callable[
115 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
116 ],
117 activation_name: str,
118 index_axis_names: Optional[Sequence[AxisNames]],
119 index_df: Optional[pd.DataFrame],
120 return_index_df: Literal[True],
121) -> Tuple[torch.Tensor, pd.DataFrame]:
122 ...
125def generic_activation_patch(
126 model: HookedTransformer,
127 corrupted_tokens: Int[torch.Tensor, "batch pos"],
128 clean_cache: ActivationCache,
129 patching_metric: Callable[[Float[torch.Tensor, "batch pos d_vocab"]], Float[torch.Tensor, ""]],
130 patch_setter: Callable[
131 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation
132 ],
133 activation_name: str,
134 index_axis_names: Optional[Sequence[AxisNames]] = None,
135 index_df: Optional[pd.DataFrame] = None,
136 return_index_df: bool = False,
137) -> Union[torch.Tensor, Tuple[torch.Tensor, pd.DataFrame]]:
138 """
139 A generic function to do activation patching, will be specialised to specific use cases.
141 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.
143 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.
145 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.
147 This function then iterates over every tuple of indices, does the relevant patch, and stores it
149 Args:
150 model: The relevant model
151 corrupted_tokens: The input tokens for the corrupted run
152 clean_cache: The cached activations from the clean run
153 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
154 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
155 activation_name: The name of the activation being patched
156 index_axis_names: The names of the axes to (fully) iterate over, implicitly fills in index_df
157 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
158 return_index_df: A Boolean flag for whether to return the dataframe of indices too
160 Returns:
161 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.
162 index_df *optional*: The dataframe of indices
163 """
165 if index_df is None: 165 ↛ 192line 165 didn't jump to line 192 because the condition on line 165 was always true
166 assert index_axis_names is not None
168 number_of_heads = model.cfg.n_heads
169 # For some models, the number of key value heads is not the same as the number of attention heads
170 if activation_name in ["k", "v"] and model.cfg.n_key_value_heads is not None: 170 ↛ 171line 170 didn't jump to line 171 because the condition on line 170 was never true
171 number_of_heads = model.cfg.n_key_value_heads
173 # Get the max range for all possible axes
174 max_axis_range = {
175 "layer": model.cfg.n_layers,
176 "pos": corrupted_tokens.shape[-1],
177 "head_index": number_of_heads,
178 }
179 max_axis_range["src_pos"] = max_axis_range["pos"]
180 max_axis_range["dest_pos"] = max_axis_range["pos"]
181 max_axis_range["head"] = max_axis_range["head_index"]
183 # Get the max range for each axis we iterate over
184 index_axis_max_range = [max_axis_range[axis_name] for axis_name in index_axis_names]
186 # Get the dataframe where each row is a tuple of indices
187 index_df = make_df_from_ranges(index_axis_max_range, index_axis_names)
189 flattened_output = False
190 else:
191 # A dataframe of indices was provided. Verify that we did not *also* receive index_axis_names
192 assert index_axis_names is None
193 index_axis_max_range = index_df.max().to_list()
195 flattened_output = True
197 # Create an empty tensor to show the patched metric for each patch
198 if flattened_output: 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true
199 patched_metric_output = torch.zeros(len(index_df), device=model.cfg.device)
200 else:
201 patched_metric_output = torch.zeros(index_axis_max_range, device=model.cfg.device)
203 # A generic patching hook - for each index, it applies the patch_setter appropriately to patch the activation
204 def patching_hook(corrupted_activation, hook, index, clean_activation):
205 if corrupted_activation.requires_grad: 205 ↛ 207line 205 didn't jump to line 207 because the condition on line 205 was always true
206 corrupted_activation = corrupted_activation.clone()
207 return patch_setter(corrupted_activation, index, clean_activation)
209 for c, index_row in enumerate(tqdm((list(index_df.iterrows())))):
210 index = index_row[1].to_list()
212 # The current activation name is just the activation name plus the layer (assumed to be the first element of the input)
213 current_activation_name = utils.get_act_name(activation_name, layer=index[0])
215 # The hook function cannot receive additional inputs, so we use partial to include the specific index and the corresponding clean activation
216 current_hook = partial(
217 patching_hook,
218 index=index,
219 clean_activation=clean_cache[current_activation_name],
220 )
222 patched_logits = model.run_with_hooks(
223 corrupted_tokens, fwd_hooks=[(current_activation_name, current_hook)]
224 )
226 if flattened_output: 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true
227 patched_metric_output[c] = patching_metric(patched_logits).item()
228 else:
229 patched_metric_output[tuple(index)] = patching_metric(patched_logits).item()
231 if return_index_df: 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true
232 return patched_metric_output, index_df
233 else:
234 return patched_metric_output
237# %%
238# Defining patch setters for various shapes of activations
239def layer_pos_patch_setter(corrupted_activation, index, clean_activation):
240 """
241 Applies the activation patch where index = [layer, pos]
243 Implicitly assumes that the activation axis order is [batch, pos, ...], which is true of everything that is not an attention pattern shaped tensor.
244 """
245 assert len(index) == 2
246 layer, pos = index
247 corrupted_activation[:, pos, ...] = clean_activation[:, pos, ...]
248 return corrupted_activation
251def layer_pos_head_vector_patch_setter(
252 corrupted_activation,
253 index,
254 clean_activation,
255):
256 """
257 Applies the activation patch where index = [layer, pos, head_index]
259 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.
260 """
261 assert len(index) == 3
262 layer, pos, head_index = index
263 corrupted_activation[:, pos, head_index] = clean_activation[:, pos, head_index]
264 return corrupted_activation
267def layer_head_vector_patch_setter(
268 corrupted_activation,
269 index,
270 clean_activation,
271):
272 """
273 Applies the activation patch where index = [layer, head_index]
275 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.
276 """
277 assert len(index) == 2
278 layer, head_index = index
279 corrupted_activation[:, :, head_index] = clean_activation[:, :, head_index]
281 return corrupted_activation
284def layer_head_pattern_patch_setter(
285 corrupted_activation,
286 index,
287 clean_activation,
288):
289 """
290 Applies the activation patch where index = [layer, head_index]
292 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns.
293 """
294 assert len(index) == 2
295 layer, head_index = index
296 corrupted_activation[:, head_index, :, :] = clean_activation[:, head_index, :, :]
298 return corrupted_activation
301def layer_head_pos_pattern_patch_setter(
302 corrupted_activation,
303 index,
304 clean_activation,
305):
306 """
307 Applies the activation patch where index = [layer, head_index, dest_pos]
309 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns.
310 """
311 assert len(index) == 3
312 layer, head_index, dest_pos = index
313 corrupted_activation[:, head_index, dest_pos, :] = clean_activation[:, head_index, dest_pos, :]
315 return corrupted_activation
318def layer_head_dest_src_pos_pattern_patch_setter(
319 corrupted_activation,
320 index,
321 clean_activation,
322):
323 """
324 Applies the activation patch where index = [layer, head_index, dest_pos, src_pos]
326 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns.
327 """
328 assert len(index) == 4
329 layer, head_index, dest_pos, src_pos = index
330 corrupted_activation[:, head_index, dest_pos, src_pos] = clean_activation[
331 :, head_index, dest_pos, src_pos
332 ]
334 return corrupted_activation
337# %%
338# Defining activation patching functions for a range of common activation patches.
339get_act_patch_resid_pre = partial(
340 generic_activation_patch,
341 patch_setter=layer_pos_patch_setter,
342 activation_name="resid_pre",
343 index_axis_names=("layer", "pos"),
344)
345get_act_patch_resid_pre.__doc__ = """
346 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]
348 See generic_activation_patch for a more detailed explanation of activation patching
350 Args:
351 model: The relevant model
352 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
353 clean_cache (ActivationCache): The cached activations from the clean run
354 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
356 Returns:
357 patched_output (torch.Tensor): The tensor of the patching metric for each resid_pre patch. Has shape [n_layers, pos]
358 """
360get_act_patch_resid_mid = partial(
361 generic_activation_patch,
362 patch_setter=layer_pos_patch_setter,
363 activation_name="resid_mid",
364 index_axis_names=("layer", "pos"),
365)
366get_act_patch_resid_mid.__doc__ = """
367 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]
369 See generic_activation_patch for a more detailed explanation of activation patching
371 Args:
372 model: The relevant model
373 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
374 clean_cache (ActivationCache): The cached activations from the clean run
375 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
377 Returns:
378 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos]
379 """
381get_act_patch_attn_out = partial(
382 generic_activation_patch,
383 patch_setter=layer_pos_patch_setter,
384 activation_name="attn_out",
385 index_axis_names=("layer", "pos"),
386)
387get_act_patch_attn_out.__doc__ = """
388 Function to get activation patching results for the output of each Attention layer (by position). Returns a tensor of shape [n_layers, pos]
390 See generic_activation_patch for a more detailed explanation of activation patching
392 Args:
393 model: The relevant model
394 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
395 clean_cache (ActivationCache): The cached activations from the clean run
396 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
398 Returns:
399 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos]
400 """
402get_act_patch_mlp_out = partial(
403 generic_activation_patch,
404 patch_setter=layer_pos_patch_setter,
405 activation_name="mlp_out",
406 index_axis_names=("layer", "pos"),
407)
408get_act_patch_mlp_out.__doc__ = """
409 Function to get activation patching results for the output of each MLP layer (by position). Returns a tensor of shape [n_layers, pos]
411 See generic_activation_patch for a more detailed explanation of activation patching
413 Args:
414 model: The relevant model
415 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
416 clean_cache (ActivationCache): The cached activations from the clean run
417 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
419 Returns:
420 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos]
421 """
422# %%
423get_act_patch_attn_head_out_by_pos = partial(
424 generic_activation_patch,
425 patch_setter=layer_pos_head_vector_patch_setter,
426 activation_name="z",
427 index_axis_names=("layer", "pos", "head"),
428)
429get_act_patch_attn_head_out_by_pos.__doc__ = """
430 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]
432 See generic_activation_patch for a more detailed explanation of activation patching
434 Args:
435 model: The relevant model
436 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
437 clean_cache (ActivationCache): The cached activations from the clean run
438 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
440 Returns:
441 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads]
442 """
444get_act_patch_attn_head_q_by_pos = partial(
445 generic_activation_patch,
446 patch_setter=layer_pos_head_vector_patch_setter,
447 activation_name="q",
448 index_axis_names=("layer", "pos", "head"),
449)
450get_act_patch_attn_head_q_by_pos.__doc__ = """
451 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]
453 See generic_activation_patch for a more detailed explanation of activation patching
455 Args:
456 model: The relevant model
457 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
458 clean_cache (ActivationCache): The cached activations from the clean run
459 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
461 Returns:
462 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads]
463 """
465get_act_patch_attn_head_k_by_pos = partial(
466 generic_activation_patch,
467 patch_setter=layer_pos_head_vector_patch_setter,
468 activation_name="k",
469 index_axis_names=("layer", "pos", "head"),
470)
471get_act_patch_attn_head_k_by_pos.__doc__ = """
472 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.
474 See generic_activation_patch for a more detailed explanation of activation patching
476 Args:
477 model: The relevant model
478 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
479 clean_cache (ActivationCache): The cached activations from the clean run
480 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
482 Returns:
483 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.
484 """
486get_act_patch_attn_head_v_by_pos = partial(
487 generic_activation_patch,
488 patch_setter=layer_pos_head_vector_patch_setter,
489 activation_name="v",
490 index_axis_names=("layer", "pos", "head"),
491)
492get_act_patch_attn_head_v_by_pos.__doc__ = """
493 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.
495 See generic_activation_patch for a more detailed explanation of activation patching
497 Args:
498 model: The relevant model
499 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
500 clean_cache (ActivationCache): The cached activations from the clean run
501 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
503 Returns:
504 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.
505 """
506# %%
507get_act_patch_attn_head_pattern_by_pos = partial(
508 generic_activation_patch,
509 patch_setter=layer_head_pos_pattern_patch_setter,
510 activation_name="pattern",
511 index_axis_names=("layer", "head_index", "dest_pos"),
512)
513get_act_patch_attn_head_pattern_by_pos.__doc__ = """
514 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]
516 See generic_activation_patch for a more detailed explanation of activation patching
518 Args:
519 model: The relevant model
520 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
521 clean_cache (ActivationCache): The cached activations from the clean run
522 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
524 Returns:
525 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads, dest_pos]
526 """
528get_act_patch_attn_head_pattern_dest_src_pos = partial(
529 generic_activation_patch,
530 patch_setter=layer_head_dest_src_pos_pattern_patch_setter,
531 activation_name="pattern",
532 index_axis_names=("layer", "head_index", "dest_pos", "src_pos"),
533)
534get_act_patch_attn_head_pattern_dest_src_pos.__doc__ = """
535 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]
537 See generic_activation_patch for a more detailed explanation of activation patching
539 Args:
540 model: The relevant model
541 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
542 clean_cache (ActivationCache): The cached activations from the clean run
543 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
545 Returns:
546 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads, dest_pos, src_pos]
547 """
549# %%
550get_act_patch_attn_head_out_all_pos = partial(
551 generic_activation_patch,
552 patch_setter=layer_head_vector_patch_setter,
553 activation_name="z",
554 index_axis_names=("layer", "head"),
555)
556get_act_patch_attn_head_out_all_pos.__doc__ = """
557 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]
559 See generic_activation_patch for a more detailed explanation of activation patching
561 Args:
562 model: The relevant model
563 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
564 clean_cache (ActivationCache): The cached activations from the clean run
565 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
567 Returns:
568 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads]
569 """
571get_act_patch_attn_head_q_all_pos = partial(
572 generic_activation_patch,
573 patch_setter=layer_head_vector_patch_setter,
574 activation_name="q",
575 index_axis_names=("layer", "head"),
576)
577get_act_patch_attn_head_q_all_pos.__doc__ = """
578 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]
580 See generic_activation_patch for a more detailed explanation of activation patching
582 Args:
583 model: The relevant model
584 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
585 clean_cache (ActivationCache): The cached activations from the clean run
586 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
588 Returns:
589 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads]
590 """
592get_act_patch_attn_head_k_all_pos = partial(
593 generic_activation_patch,
594 patch_setter=layer_head_vector_patch_setter,
595 activation_name="k",
596 index_axis_names=("layer", "head"),
597)
598get_act_patch_attn_head_k_all_pos.__doc__ = """
599 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.
601 See generic_activation_patch for a more detailed explanation of activation patching
603 Args:
604 model: The relevant model
605 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
606 clean_cache (ActivationCache): The cached activations from the clean run
607 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
609 Returns:
610 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.
611 """
613get_act_patch_attn_head_v_all_pos = partial(
614 generic_activation_patch,
615 patch_setter=layer_head_vector_patch_setter,
616 activation_name="v",
617 index_axis_names=("layer", "head"),
618)
619get_act_patch_attn_head_v_all_pos.__doc__ = """
620 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.
622 See generic_activation_patch for a more detailed explanation of activation patching
624 Args:
625 model: The relevant model
626 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
627 clean_cache (ActivationCache): The cached activations from the clean run
628 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
630 Returns:
631 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.
632 """
634get_act_patch_attn_head_pattern_all_pos = partial(
635 generic_activation_patch,
636 patch_setter=layer_head_pattern_patch_setter,
637 activation_name="pattern",
638 index_axis_names=("layer", "head_index"),
639)
640get_act_patch_attn_head_pattern_all_pos.__doc__ = """
641 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]
643 See generic_activation_patch for a more detailed explanation of activation patching
645 Args:
646 model: The relevant model
647 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
648 clean_cache (ActivationCache): The cached activations from the clean run
649 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
651 Returns:
652 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads]
653 """
655# %%
658def get_act_patch_attn_head_all_pos_every(
659 model, corrupted_tokens, clean_cache, metric
660) -> Float[torch.Tensor, "patch_type layer head"]:
661 """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]
663 Args:
664 model: The relevant model
665 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
666 clean_cache (ActivationCache): The cached activations from the clean run
667 metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
669 Returns:
670 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [5, n_layers, n_heads]
671 """
672 act_patch_results: list[torch.Tensor] = []
673 act_patch_results.append(
674 get_act_patch_attn_head_out_all_pos(model, corrupted_tokens, clean_cache, metric)
675 )
676 act_patch_results.append(
677 get_act_patch_attn_head_q_all_pos(model, corrupted_tokens, clean_cache, metric)
678 )
680 # Reshape k and v to be compatible with the rest of the results in case of n_key_value_heads != n_heads
681 k_results = get_act_patch_attn_head_k_all_pos(model, corrupted_tokens, clean_cache, metric)
682 act_patch_results.append(
683 torch.nn.functional.pad(k_results, (0, act_patch_results[-1].size(-1) - k_results.size(-1)))
684 )
685 v_results = get_act_patch_attn_head_v_all_pos(model, corrupted_tokens, clean_cache, metric)
686 act_patch_results.append(
687 torch.nn.functional.pad(v_results, (0, act_patch_results[-1].size(-1) - v_results.size(-1)))
688 )
690 act_patch_results.append(
691 get_act_patch_attn_head_pattern_all_pos(model, corrupted_tokens, clean_cache, metric)
692 )
693 return torch.stack(act_patch_results, dim=0)
696def get_act_patch_attn_head_by_pos_every(
697 model, corrupted_tokens, clean_cache, metric
698) -> Float[torch.Tensor, "patch_type layer pos head"]:
699 """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]
701 Args:
702 model: The relevant model
703 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
704 clean_cache (ActivationCache): The cached activations from the clean run
705 metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
707 Returns:
708 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [5, n_layers, pos, n_heads]
709 """
710 act_patch_results = []
711 act_patch_results.append(
712 get_act_patch_attn_head_out_by_pos(model, corrupted_tokens, clean_cache, metric)
713 )
714 act_patch_results.append(
715 get_act_patch_attn_head_q_by_pos(model, corrupted_tokens, clean_cache, metric)
716 )
718 # Reshape k and v to be compatible with the rest of the results in case of n_key_value_heads != n_heads
719 k_results = get_act_patch_attn_head_k_by_pos(model, corrupted_tokens, clean_cache, metric)
720 act_patch_results.append(
721 torch.nn.functional.pad(k_results, (0, act_patch_results[-1].size(-1) - k_results.size(-1)))
722 )
723 v_results = get_act_patch_attn_head_v_by_pos(model, corrupted_tokens, clean_cache, metric)
724 act_patch_results.append(
725 torch.nn.functional.pad(v_results, (0, act_patch_results[-1].size(-1) - v_results.size(-1)))
726 )
728 # Reshape pattern to be compatible with the rest of the results
729 pattern_results = get_act_patch_attn_head_pattern_by_pos(
730 model, corrupted_tokens, clean_cache, metric
731 )
732 act_patch_results.append(einops.rearrange(pattern_results, "batch head pos -> batch pos head"))
733 return torch.stack(act_patch_results, dim=0)
736def get_act_patch_block_every(
737 model, corrupted_tokens, clean_cache, metric
738) -> Float[torch.Tensor, "patch_type layer pos"]:
739 """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]
741 Args:
742 model: The relevant model
743 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos]
744 clean_cache (ActivationCache): The cached activations from the clean run
745 metric: A function from the model's output logits to some metric (eg loss, logit diff, etc)
747 Returns:
748 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [3, n_layers, pos]
749 """
750 act_patch_results = []
751 act_patch_results.append(get_act_patch_resid_pre(model, corrupted_tokens, clean_cache, metric))
752 act_patch_results.append(get_act_patch_attn_out(model, corrupted_tokens, clean_cache, metric))
753 act_patch_results.append(get_act_patch_mlp_out(model, corrupted_tokens, clean_cache, metric))
754 return torch.stack(act_patch_results, dim=0)