Coverage for transformer_lens/patching.py: 95%

1"""Patching. 

2 

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 

8Context: 

9 

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. 

14 

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. 

21 

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. 

29 

30Intuition: 

31 

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: 

38 

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 

45 

46This, unlike direct logit attribution, can identify meaningful parts of a circuit from anywhere 

47within the model, rather than just the end. 

48""" 

49 

50from __future__ import annotations 

51 

52import itertools 

53from dataclasses import dataclass 

54from functools import partial 

55from typing import ( 

56 Callable, 

57 Dict, 

58 List, 

59 Mapping, 

60 Optional, 

61 Sequence, 

62 Tuple, 

63 Union, 

64 overload, 

65) 

66 

67import einops 

68import pandas as pd 

69import torch 

70from jaxtyping import Float, Int 

71from tqdm.auto import tqdm 

72from typing_extensions import Literal 

73 

74import transformer_lens.utilities as utils 

75from transformer_lens.ActivationCache import ActivationCache 

76from transformer_lens.model_protocol import TransformerLensModel 

77from transformer_lens.utilities.statistics import ( 

78 bootstrap_ci, 

79 derive_generator, 

80 sign_flip_permutation_pvalue, 

81 standard_error, 

82) 

83 

84# %% 

85Logits = torch.Tensor 

86AxisNames = Literal["layer", "pos", "head_index", "head", "src_pos", "dest_pos"] 

87 

88 

89# %% 

90 

91 

92def make_df_from_ranges( 

93 column_max_ranges: Sequence[int], column_names: Sequence[str] 

94) -> pd.DataFrame: 

95 """ 

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

97 """ 

98 rows = list(itertools.product(*[range(axis_max_range) for axis_max_range in column_max_ranges])) 

99 df = pd.DataFrame(rows, columns=column_names) 

100 return df 

101 

102 

103# %% 

104CorruptedActivation = torch.Tensor 

105PatchedActivation = torch.Tensor 

106MetricFn = Callable[[Float[torch.Tensor, "batch pos d_vocab"]], torch.Tensor] 

107Reduce = Literal["scalar", "mean", "none"] 

108_REDUCE_MODES = ("scalar", "mean", "none") 

109_RESULT_AXIS_RENAMES = {"head_index": "head", "dest_pos": "pos"} 

110 

111 

112@dataclass 

113class PatchingResult: 

114 """Per-example (or per-metric) activation-patching results with statistics on top. 

115 

116 Returned by :func:`generic_activation_patch` and every ``get_act_patch_*`` 

117 function when ``reduce="none"``, ``baseline=True`` or several metrics are 

118 requested. ``values`` is laid out ``[metric?, *axes | n_rows, batch?]``: the 

119 metric axis exists only when a mapping of metrics was passed 

120 (``has_metric_axis``), the middle axes are ``axis_names`` (or one flat 

121 ``n_rows`` axis when an ``index_df`` was supplied), and the trailing batch 

122 axis exists only when ``reduced`` is ``False``. 

123 

124 Attributes: 

125 values: The patched metric values. 

126 axis_names: Names of the swept axes, ``[]`` in flat ``index_df`` mode. 

127 index_df: One row per swept cell, in the order the cells were run. 

128 metric_names: Metric names, ``["metric"]`` for a single callable. 

129 reduced: ``True`` when the batch axis has been averaged away. 

130 baseline: Unpatched corrupted-run metric, ``[metric?, batch]`` (or 

131 ``[metric?]`` when reduced), when the sweep was run with 

132 ``baseline=True``. 

133 has_metric_axis: Whether ``values`` carries a leading metric axis. 

134 patch_type_names: Names of the ``patch_type`` axis entries when the 

135 result came from a ``*_every`` helper (``"out"``, ``"q"``, ...). 

136 GQA ``k``/``v`` sweeps are zero-padded to ``n_heads`` there, so the 

137 padded cells appear as zero-valued rows in ``index_df`` and 

138 :meth:`to_dataframe`. 

139 """ 

140 

141 values: torch.Tensor 

142 axis_names: List[str] 

143 index_df: pd.DataFrame 

144 metric_names: List[str] 

145 reduced: bool 

146 baseline: Optional[torch.Tensor] = None 

147 has_metric_axis: bool = False 

148 patch_type_names: Optional[List[str]] = None 

149 

150 def __getitem__(self, metric_name: str) -> "PatchingResult": 

151 """One metric's result, with the metric axis dropped.""" 

152 if metric_name not in self.metric_names: 

153 raise KeyError(f"unknown metric {metric_name!r}; have {self.metric_names}") 

154 if not self.has_metric_axis: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true

155 return self 

156 i = self.metric_names.index(metric_name) 

157 return PatchingResult( 

158 values=self.values[i], 

159 axis_names=list(self.axis_names), 

160 index_df=self.index_df, 

161 metric_names=[metric_name], 

162 reduced=self.reduced, 

163 baseline=None if self.baseline is None else self.baseline[i], 

164 has_metric_axis=False, 

165 patch_type_names=self.patch_type_names, 

166 ) 

167 

168 def _require_per_example(self, what: str) -> None: 

169 if self.reduced: 

170 raise ValueError(f"{what} needs per-example values; run the sweep with reduce='none'") 

171 

172 def mean(self) -> torch.Tensor: 

173 """Batch mean per cell (``values`` itself when already reduced).""" 

174 return self.values if self.reduced else self.values.mean(dim=-1) 

175 

176 def stderr(self) -> torch.Tensor: 

177 """Standard error of the batch mean per cell.""" 

178 self._require_per_example("stderr") 

179 return standard_error(self.values, dim=-1) 

180 

181 def _broadcast_baseline(self) -> torch.Tensor: 

182 if self.baseline is None: 

183 raise ValueError("effect needs a baseline; run the sweep with baseline=True") 

184 lead = 1 if self.has_metric_axis else 0 

185 middle = self.values.ndim - lead - (0 if self.reduced else 1) 

186 shape = ( 

187 tuple(self.baseline.shape[:lead]) + (1,) * middle + tuple(self.baseline.shape[lead:]) 

188 ) 

189 return self.baseline.reshape(shape) 

190 

191 def effect(self) -> torch.Tensor: 

192 """``values - baseline``: the change each patch makes to the corrupted run, per example.""" 

193 return self.values - self._broadcast_baseline() 

194 

195 def normalized(self, clean_baseline: torch.Tensor) -> torch.Tensor: 

196 """``(values - corrupted) / (clean - corrupted)`` per example. 

197 

198 ``clean_baseline`` is the same metric evaluated on the unpatched clean run, 

199 shaped like ``baseline``. 0 means the patch did nothing, 1 means it 

200 restored the clean metric. A prompt whose clean and corrupted metrics 

201 coincide has no recovery scale and yields ``inf`` / ``nan`` for that 

202 example; filter such prompts out before normalizing. 

203 """ 

204 corrupted = self._broadcast_baseline() 

205 baseline = self.baseline 

206 assert baseline is not None # _broadcast_baseline raised otherwise 

207 clean = torch.as_tensor(clean_baseline, dtype=self.values.dtype, device=self.values.device) 

208 if clean.shape != baseline.shape: 

209 raise ValueError( 

210 f"clean_baseline must match baseline shape {tuple(baseline.shape)}, " 

211 f"got {tuple(clean.shape)}" 

212 ) 

213 clean = clean.reshape(corrupted.shape) 

214 return (self.values - corrupted) / (clean - corrupted) 

215 

216 def bootstrap_ci( 

217 self, 

218 *, 

219 confidence: float = 0.95, 

220 n_resamples: int = 1000, 

221 seed: int = 0, 

222 statistic: str = "mean", 

223 of_effect: bool = False, 

224 ) -> Tuple[torch.Tensor, torch.Tensor]: 

225 """Percentile bootstrap interval of the batch ``statistic`` per cell. 

226 

227 ``of_effect=True`` bootstraps ``effect()`` instead of ``values``. Seeded 

228 through :func:`~transformer_lens.utilities.statistics.derive_generator`. 

229 """ 

230 self._require_per_example("bootstrap_ci") 

231 data = self.effect() if of_effect else self.values 

232 return bootstrap_ci( 

233 data, 

234 dim=-1, 

235 statistic=statistic, 

236 confidence=confidence, 

237 n_resamples=n_resamples, 

238 generator=derive_generator(seed, "bootstrap", statistic), 

239 ) 

240 

241 def permutation_pvalue( 

242 self, 

243 *, 

244 n_permutations: int = 1000, 

245 seed: int = 0, 

246 alternative: str = "two-sided", 

247 ) -> torch.Tensor: 

248 """Sign-flip permutation p-value that each cell's mean effect is zero. 

249 

250 Uses ``effect()`` (needs ``baseline=True``); exact for 12 or fewer 

251 examples, Monte Carlo above that. 

252 """ 

253 self._require_per_example("permutation_pvalue") 

254 return sign_flip_permutation_pvalue( 

255 self.effect(), 

256 dim=-1, 

257 n_permutations=n_permutations, 

258 generator=derive_generator(seed, "permutation"), 

259 alternative=alternative, 

260 ) 

261 

262 def to_dataframe(self) -> pd.DataFrame: 

263 """Long form: one row per (metric, swept cell, example) with the value.""" 

264 rows = [] 

265 n_metrics = len(self.metric_names) 

266 values = self.values if self.has_metric_axis else self.values.unsqueeze(0) 

267 flat_cells = values.reshape(n_metrics, len(self.index_df), -1).detach().cpu() 

268 for m, name in enumerate(self.metric_names): 

269 for c, (_, index_row) in enumerate(self.index_df.iterrows()): 

270 for b in range(flat_cells.shape[-1]): 

271 row = {"metric": name, **index_row.to_dict()} 

272 if self.patch_type_names is not None and "patch_type" in row: 

273 row["patch_type_name"] = self.patch_type_names[int(row["patch_type"])] 

274 if not self.reduced: 274 ↛ 276line 274 didn't jump to line 276 because the condition on line 274 was always true

275 row["example"] = b 

276 row["value"] = flat_cells[m, c, b].item() 

277 rows.append(row) 

278 return pd.DataFrame(rows) 

279 

280 

281@overload 

282def generic_activation_patch( 

283 model: TransformerLensModel, 

284 corrupted_tokens: Int[torch.Tensor, "batch pos"], 

285 clean_cache: ActivationCache, 

286 patching_metric: MetricFn, 

287 patch_setter: Callable[ 

288 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation 

289 ], 

290 activation_name: str, 

291 index_axis_names: Optional[Sequence[AxisNames]] = None, 

292 index_df: Optional[pd.DataFrame] = None, 

293 return_index_df: Literal[False] = False, 

294 *, 

295 reduce: Literal["scalar", "mean"] = "scalar", 

296 baseline: Literal[False] = False, 

297) -> torch.Tensor: 

298 ... 

299 

300 

301@overload 

302def generic_activation_patch( 

303 model: TransformerLensModel, 

304 corrupted_tokens: Int[torch.Tensor, "batch pos"], 

305 clean_cache: ActivationCache, 

306 patching_metric: MetricFn, 

307 patch_setter: Callable[ 

308 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation 

309 ], 

310 activation_name: str, 

311 index_axis_names: Optional[Sequence[AxisNames]], 

312 index_df: Optional[pd.DataFrame], 

313 return_index_df: Literal[True], 

314 *, 

315 reduce: Literal["scalar", "mean"] = "scalar", 

316 baseline: Literal[False] = False, 

317) -> Tuple[torch.Tensor, pd.DataFrame]: 

318 ... 

319 

320 

321@overload 

322def generic_activation_patch( 

323 model: TransformerLensModel, 

324 corrupted_tokens: Int[torch.Tensor, "batch pos"], 

325 clean_cache: ActivationCache, 

326 patching_metric: Union[MetricFn, Mapping[str, MetricFn]], 

327 patch_setter: Callable[ 

328 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation 

329 ], 

330 activation_name: str, 

331 index_axis_names: Optional[Sequence[AxisNames]] = None, 

332 index_df: Optional[pd.DataFrame] = None, 

333 return_index_df: bool = False, 

334 *, 

335 reduce: Literal["none"], 

336 baseline: bool = False, 

337) -> PatchingResult: 

338 ... 

339 

340 

341@overload 

342def generic_activation_patch( 

343 model: TransformerLensModel, 

344 corrupted_tokens: Int[torch.Tensor, "batch pos"], 

345 clean_cache: ActivationCache, 

346 patching_metric: MetricFn, 

347 patch_setter: Callable[ 

348 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation 

349 ], 

350 activation_name: str, 

351 index_axis_names: Optional[Sequence[AxisNames]] = None, 

352 index_df: Optional[pd.DataFrame] = None, 

353 return_index_df: bool = False, 

354 *, 

355 reduce: Reduce = "scalar", 

356 baseline: Literal[True], 

357) -> PatchingResult: 

358 ... 

359 

360 

361@overload 

362def generic_activation_patch( 

363 model: TransformerLensModel, 

364 corrupted_tokens: Int[torch.Tensor, "batch pos"], 

365 clean_cache: ActivationCache, 

366 patching_metric: Mapping[str, MetricFn], 

367 patch_setter: Callable[ 

368 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation 

369 ], 

370 activation_name: str, 

371 index_axis_names: Optional[Sequence[AxisNames]] = None, 

372 index_df: Optional[pd.DataFrame] = None, 

373 return_index_df: bool = False, 

374 *, 

375 reduce: Reduce = "scalar", 

376 baseline: bool = False, 

377) -> PatchingResult: 

378 ... 

379 

380 

381def generic_activation_patch( 

382 model: TransformerLensModel, 

383 corrupted_tokens: Int[torch.Tensor, "batch pos"], 

384 clean_cache: ActivationCache, 

385 patching_metric: Union[MetricFn, Mapping[str, MetricFn]], 

386 patch_setter: Callable[ 

387 [CorruptedActivation, Sequence[int], ActivationCache], PatchedActivation 

388 ], 

389 activation_name: str, 

390 index_axis_names: Optional[Sequence[AxisNames]] = None, 

391 index_df: Optional[pd.DataFrame] = None, 

392 return_index_df: bool = False, 

393 *, 

394 reduce: str = "scalar", # Reduce; widened so the ValueError below fires before beartype 

395 baseline: bool = False, 

396) -> Union[torch.Tensor, Tuple[torch.Tensor, pd.DataFrame], PatchingResult]: 

397 """ 

398 A generic function to do activation patching, will be specialised to specific use cases. 

399 

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

401 

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

403 

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

405 

406 This function then iterates over every tuple of indices, does the relevant patch, and stores it 

407 

408 Every patched forward runs the whole corrupted batch, so per-example values are 

409 available for free: ``reduce="none"`` keeps them (the metric must then return a 

410 ``[batch]`` tensor) and returns a :class:`PatchingResult` with mean, standard 

411 error, bootstrap intervals and permutation p-values on top; ``reduce="mean"`` 

412 averages a ``[batch]`` metric back to today's tensor; a mapping of metrics 

413 evaluates all of them on each patched forward. ``baseline=True`` adds one 

414 unpatched corrupted forward so effects and normalized recovery can be reported. 

415 

416 Args: 

417 model: The relevant model 

418 corrupted_tokens: The input tokens for the corrupted run 

419 clean_cache: The cached activations from the clean run 

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

421 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 

422 activation_name: The name of the activation being patched 

423 index_axis_names: The names of the axes to (fully) iterate over, implicitly fills in index_df 

424 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 

425 return_index_df: A Boolean flag for whether to return the dataframe of indices too (tensor return only; a ``PatchingResult`` always carries ``index_df``) 

426 reduce: ``"scalar"`` (default, metric returns a scalar, tensor return), ``"mean"`` (metric returns ``[batch]``, averaged, tensor return) or ``"none"`` (metric returns ``[batch]``, kept, ``PatchingResult`` return) 

427 baseline: Also evaluate the metric on the unpatched corrupted run and return a ``PatchingResult`` carrying it 

428 

429 Returns: 

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

431 index_df *optional*: The dataframe of indices 

432 Or a :class:`PatchingResult` when ``reduce="none"``, ``baseline=True`` or a metric mapping was given. 

433 """ 

434 # Lazy import: tools.analysis imports this module's neighbours at package import. 

435 from transformer_lens.tools.analysis._model_state import require_eval_mode 

436 

437 require_eval_mode(model, operation="activation patching") 

438 if reduce not in _REDUCE_MODES: 

439 raise ValueError(f"reduce must be one of {_REDUCE_MODES}, got {reduce!r}") 

440 metrics: Dict[str, MetricFn] 

441 if isinstance(patching_metric, Mapping): 

442 has_metric_axis = True 

443 metrics = dict(patching_metric) 

444 if not metrics: 444 ↛ 445line 444 didn't jump to line 445 because the condition on line 444 was never true

445 raise ValueError("patching_metric mapping must contain at least one metric") 

446 else: 

447 has_metric_axis = False 

448 metrics = {"metric": patching_metric} 

449 want_result = has_metric_axis or reduce == "none" or baseline 

450 if want_result and return_index_df: 

451 raise ValueError( 

452 "return_index_df only applies to the tensor return; a PatchingResult carries index_df" 

453 ) 

454 

455 if index_df is None: 

456 assert index_axis_names is not None 

457 number_of_heads = model.cfg.n_heads 

458 # For some models, the number of key value heads is not the same as the number of attention heads 

459 if activation_name in ["k", "v"] and model.cfg.n_key_value_heads is not None: 

460 number_of_heads = model.cfg.n_key_value_heads 

461 

462 # Get the max range for all possible axes 

463 max_axis_range = { 

464 "layer": model.cfg.n_layers, 

465 "pos": corrupted_tokens.shape[-1], 

466 "head_index": number_of_heads, 

467 } 

468 max_axis_range["src_pos"] = max_axis_range["pos"] 

469 max_axis_range["dest_pos"] = max_axis_range["pos"] 

470 max_axis_range["head"] = max_axis_range["head_index"] 

471 

472 # Get the max range for each axis we iterate over 

473 index_axis_max_range = [max_axis_range[axis_name] for axis_name in index_axis_names] 

474 

475 # Get the dataframe where each row is a tuple of indices 

476 index_df = make_df_from_ranges(index_axis_max_range, index_axis_names) 

477 

478 flattened_output = False 

479 else: 

480 # A dataframe of indices was provided. Verify that we did not *also* receive index_axis_names 

481 assert index_axis_names is None 

482 index_axis_max_range = index_df.max().to_list() 

483 flattened_output = True 

484 

485 batch_size = int(corrupted_tokens.shape[0]) 

486 per_example = reduce != "scalar" 

487 cell_shape = (len(index_df),) if flattened_output else tuple(index_axis_max_range) 

488 store_shape = (len(metrics),) + cell_shape + ((batch_size,) if per_example else ()) 

489 # Create an empty tensor to show the patched metric for each patch 

490 store = torch.zeros(store_shape, device=model.cfg.device) 

491 

492 def evaluate(logits: torch.Tensor) -> torch.Tensor: 

493 """Every metric on one set of logits, shaped [metric, batch?] after validation.""" 

494 rows = [] 

495 for name, fn in metrics.items(): 

496 out = torch.as_tensor(fn(logits)) 

497 if per_example: 

498 if out.ndim != 1 or out.shape[0] != batch_size: 

499 raise ValueError( 

500 f"patching_metric {name!r} must return a [batch] tensor of length " 

501 f"{batch_size} under reduce={reduce!r}, got shape {tuple(out.shape)}" 

502 ) 

503 elif out.ndim != 0: 

504 raise ValueError( 

505 f"patching_metric {name!r} must return a scalar under reduce='scalar', got " 

506 f"shape {tuple(out.shape)}; return a [batch] tensor and pass reduce='mean' " 

507 "or reduce='none' to keep per-example values" 

508 ) 

509 rows.append(out.detach()) 

510 return torch.stack(rows).to(store.dtype) 

511 

512 baseline_values: Optional[torch.Tensor] = None 

513 if baseline: 513 ↛ 518line 513 didn't jump to line 518 because the condition on line 513 was always true

514 # Same entry point as the patched runs, with no hooks attached. 

515 baseline_values = evaluate(model.run_with_hooks(corrupted_tokens, fwd_hooks=[])) 

516 

517 # A generic patching hook - for each index, it applies the patch_setter appropriately to patch the activation 

518 def patching_hook(corrupted_activation, hook, index, clean_activation): 

519 if corrupted_activation.requires_grad: 519 ↛ 521line 519 didn't jump to line 521 because the condition on line 519 was always true

520 corrupted_activation = corrupted_activation.clone() 

521 return patch_setter(corrupted_activation, index, clean_activation) 

522 

523 for c, index_row in enumerate(tqdm((list(index_df.iterrows())))): 

524 index = index_row[1].to_list() 

525 

526 # The current activation name is just the activation name plus the layer (assumed to be the first element of the input) 

527 current_activation_name = utils.get_act_name(activation_name, layer=index[0]) 

528 

529 # The hook function cannot receive additional inputs, so we use partial to include the specific index and the corresponding clean activation 

530 current_hook = partial( 

531 patching_hook, 

532 index=index, 

533 clean_activation=clean_cache[current_activation_name], 

534 ) 

535 

536 patched_logits = model.run_with_hooks( 

537 corrupted_tokens, fwd_hooks=[(current_activation_name, current_hook)] 

538 ) 

539 values = evaluate(patched_logits) 

540 if flattened_output: 

541 store[:, c] = values 

542 else: 

543 store[(slice(None), *index)] = values 

544 

545 if reduce == "mean": 

546 store = store.mean(dim=-1) 

547 if baseline_values is not None: 

548 baseline_values = baseline_values.mean(dim=-1) 

549 

550 if not want_result: 

551 patched_metric_output = store[0] 

552 if return_index_df: 

553 return patched_metric_output, index_df 

554 return patched_metric_output 

555 

556 return PatchingResult( 

557 values=store if has_metric_axis else store[0], 

558 axis_names=[] if flattened_output else list(index_axis_names or []), 

559 index_df=index_df, 

560 metric_names=list(metrics), 

561 reduced=reduce != "none", 

562 baseline=None 

563 if baseline_values is None 

564 else (baseline_values if has_metric_axis else baseline_values[0]), 

565 has_metric_axis=has_metric_axis, 

566 ) 

567 

568 

569# %% 

570# Defining patch setters for various shapes of activations 

571def layer_pos_patch_setter(corrupted_activation, index, clean_activation): 

572 """ 

573 Applies the activation patch where index = [layer, pos] 

574 

575 Implicitly assumes that the activation axis order is [batch, pos, ...], which is true of everything that is not an attention pattern shaped tensor. 

576 """ 

577 assert len(index) == 2 

578 layer, pos = index 

579 corrupted_activation[:, pos, ...] = clean_activation[:, pos, ...] 

580 return corrupted_activation 

581 

582 

583def layer_pos_head_vector_patch_setter( 

584 corrupted_activation, 

585 index, 

586 clean_activation, 

587): 

588 """ 

589 Applies the activation patch where index = [layer, pos, head_index] 

590 

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

592 """ 

593 assert len(index) == 3 

594 layer, pos, head_index = index 

595 corrupted_activation[:, pos, head_index] = clean_activation[:, pos, head_index] 

596 return corrupted_activation 

597 

598 

599def layer_head_vector_patch_setter( 

600 corrupted_activation, 

601 index, 

602 clean_activation, 

603): 

604 """ 

605 Applies the activation patch where index = [layer, head_index] 

606 

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

608 """ 

609 assert len(index) == 2 

610 layer, head_index = index 

611 corrupted_activation[:, :, head_index] = clean_activation[:, :, head_index] 

612 

613 return corrupted_activation 

614 

615 

616def layer_head_pattern_patch_setter( 

617 corrupted_activation, 

618 index, 

619 clean_activation, 

620): 

621 """ 

622 Applies the activation patch where index = [layer, head_index] 

623 

624 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns. 

625 """ 

626 assert len(index) == 2 

627 layer, head_index = index 

628 corrupted_activation[:, head_index, :, :] = clean_activation[:, head_index, :, :] 

629 

630 return corrupted_activation 

631 

632 

633def layer_head_pos_pattern_patch_setter( 

634 corrupted_activation, 

635 index, 

636 clean_activation, 

637): 

638 """ 

639 Applies the activation patch where index = [layer, head_index, dest_pos] 

640 

641 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns. 

642 """ 

643 assert len(index) == 3 

644 layer, head_index, dest_pos = index 

645 corrupted_activation[:, head_index, dest_pos, :] = clean_activation[:, head_index, dest_pos, :] 

646 

647 return corrupted_activation 

648 

649 

650def layer_head_dest_src_pos_pattern_patch_setter( 

651 corrupted_activation, 

652 index, 

653 clean_activation, 

654): 

655 """ 

656 Applies the activation patch where index = [layer, head_index, dest_pos, src_pos] 

657 

658 Implicitly assumes that the activation axis order is [batch, head_index, dest_pos, src_pos], which is true of attention scores and patterns. 

659 """ 

660 assert len(index) == 4 

661 layer, head_index, dest_pos, src_pos = index 

662 corrupted_activation[:, head_index, dest_pos, src_pos] = clean_activation[ 

663 :, head_index, dest_pos, src_pos 

664 ] 

665 

666 return corrupted_activation 

667 

668 

669# %% 

670# Defining activation patching functions for a range of common activation patches. 

671get_act_patch_resid_pre = partial( 

672 generic_activation_patch, 

673 patch_setter=layer_pos_patch_setter, 

674 activation_name="resid_pre", 

675 index_axis_names=("layer", "pos"), 

676) 

677get_act_patch_resid_pre.__doc__ = """ 

678 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] 

679 

680 See generic_activation_patch for a more detailed explanation of activation patching  

681 

682 Args: 

683 model: The relevant model 

684 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

685 clean_cache (ActivationCache): The cached activations from the clean run 

686 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

687 

688 Returns: 

689 patched_output (torch.Tensor): The tensor of the patching metric for each resid_pre patch. Has shape [n_layers, pos] 

690 """ 

691 

692get_act_patch_resid_mid = partial( 

693 generic_activation_patch, 

694 patch_setter=layer_pos_patch_setter, 

695 activation_name="resid_mid", 

696 index_axis_names=("layer", "pos"), 

697) 

698get_act_patch_resid_mid.__doc__ = """ 

699 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] 

700 

701 See generic_activation_patch for a more detailed explanation of activation patching  

702 

703 Args: 

704 model: The relevant model 

705 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

706 clean_cache (ActivationCache): The cached activations from the clean run 

707 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

708 

709 Returns: 

710 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos] 

711 """ 

712 

713get_act_patch_attn_out = partial( 

714 generic_activation_patch, 

715 patch_setter=layer_pos_patch_setter, 

716 activation_name="attn_out", 

717 index_axis_names=("layer", "pos"), 

718) 

719get_act_patch_attn_out.__doc__ = """ 

720 Function to get activation patching results for the output of each Attention layer (by position). Returns a tensor of shape [n_layers, pos] 

721 

722 See generic_activation_patch for a more detailed explanation of activation patching  

723 

724 Args: 

725 model: The relevant model 

726 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

727 clean_cache (ActivationCache): The cached activations from the clean run 

728 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

729 

730 Returns: 

731 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos] 

732 """ 

733 

734get_act_patch_mlp_out = partial( 

735 generic_activation_patch, 

736 patch_setter=layer_pos_patch_setter, 

737 activation_name="mlp_out", 

738 index_axis_names=("layer", "pos"), 

739) 

740get_act_patch_mlp_out.__doc__ = """ 

741 Function to get activation patching results for the output of each MLP layer (by position). Returns a tensor of shape [n_layers, pos] 

742 

743 See generic_activation_patch for a more detailed explanation of activation patching  

744 

745 Args: 

746 model: The relevant model 

747 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

748 clean_cache (ActivationCache): The cached activations from the clean run 

749 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

750 

751 Returns: 

752 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos] 

753 """ 

754# %% 

755get_act_patch_attn_head_out_by_pos = partial( 

756 generic_activation_patch, 

757 patch_setter=layer_pos_head_vector_patch_setter, 

758 activation_name="z", 

759 index_axis_names=("layer", "pos", "head"), 

760) 

761get_act_patch_attn_head_out_by_pos.__doc__ = """ 

762 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] 

763 

764 See generic_activation_patch for a more detailed explanation of activation patching  

765 

766 Args: 

767 model: The relevant model 

768 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

769 clean_cache (ActivationCache): The cached activations from the clean run 

770 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

771 

772 Returns: 

773 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads] 

774 """ 

775 

776get_act_patch_attn_head_q_by_pos = partial( 

777 generic_activation_patch, 

778 patch_setter=layer_pos_head_vector_patch_setter, 

779 activation_name="q", 

780 index_axis_names=("layer", "pos", "head"), 

781) 

782get_act_patch_attn_head_q_by_pos.__doc__ = """ 

783 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] 

784 

785 See generic_activation_patch for a more detailed explanation of activation patching  

786 

787 Args: 

788 model: The relevant model 

789 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

790 clean_cache (ActivationCache): The cached activations from the clean run 

791 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

792 

793 Returns: 

794 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, pos, n_heads] 

795 """ 

796 

797get_act_patch_attn_head_k_by_pos = partial( 

798 generic_activation_patch, 

799 patch_setter=layer_pos_head_vector_patch_setter, 

800 activation_name="k", 

801 index_axis_names=("layer", "pos", "head"), 

802) 

803get_act_patch_attn_head_k_by_pos.__doc__ = """ 

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

805 

806 See generic_activation_patch for a more detailed explanation of activation patching  

807 

808 Args: 

809 model: The relevant model 

810 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

811 clean_cache (ActivationCache): The cached activations from the clean run 

812 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

813 

814 Returns: 

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

816 """ 

817 

818get_act_patch_attn_head_v_by_pos = partial( 

819 generic_activation_patch, 

820 patch_setter=layer_pos_head_vector_patch_setter, 

821 activation_name="v", 

822 index_axis_names=("layer", "pos", "head"), 

823) 

824get_act_patch_attn_head_v_by_pos.__doc__ = """ 

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

826 

827 See generic_activation_patch for a more detailed explanation of activation patching  

828 

829 Args: 

830 model: The relevant model 

831 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

832 clean_cache (ActivationCache): The cached activations from the clean run 

833 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

834 

835 Returns: 

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

837 """ 

838# %% 

839get_act_patch_attn_head_pattern_by_pos = partial( 

840 generic_activation_patch, 

841 patch_setter=layer_head_pos_pattern_patch_setter, 

842 activation_name="pattern", 

843 index_axis_names=("layer", "head_index", "dest_pos"), 

844) 

845get_act_patch_attn_head_pattern_by_pos.__doc__ = """ 

846 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] 

847 

848 See generic_activation_patch for a more detailed explanation of activation patching  

849 

850 Args: 

851 model: The relevant model 

852 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

853 clean_cache (ActivationCache): The cached activations from the clean run 

854 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

855 

856 Returns: 

857 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads, dest_pos] 

858 """ 

859 

860get_act_patch_attn_head_pattern_dest_src_pos = partial( 

861 generic_activation_patch, 

862 patch_setter=layer_head_dest_src_pos_pattern_patch_setter, 

863 activation_name="pattern", 

864 index_axis_names=("layer", "head_index", "dest_pos", "src_pos"), 

865) 

866get_act_patch_attn_head_pattern_dest_src_pos.__doc__ = """ 

867 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] 

868 

869 See generic_activation_patch for a more detailed explanation of activation patching  

870 

871 Args: 

872 model: The relevant model 

873 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

874 clean_cache (ActivationCache): The cached activations from the clean run 

875 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

876 

877 Returns: 

878 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads, dest_pos, src_pos] 

879 """ 

880 

881# %% 

882get_act_patch_attn_head_out_all_pos = partial( 

883 generic_activation_patch, 

884 patch_setter=layer_head_vector_patch_setter, 

885 activation_name="z", 

886 index_axis_names=("layer", "head"), 

887) 

888get_act_patch_attn_head_out_all_pos.__doc__ = """ 

889 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] 

890 

891 See generic_activation_patch for a more detailed explanation of activation patching  

892 

893 Args: 

894 model: The relevant model 

895 corrupted_tokens (torch.Tensor): The input tokens for the corrupted run. Has shape [batch, pos] 

896 clean_cache (ActivationCache): The cached activations from the clean run 

897 patching_metric: A function from the model's output logits to some metric (eg loss, logit diff, etc) 

898 

899 Returns: 

900 patched_output (torch.Tensor): The tensor of the patching metric for each patch. Has shape [n_layers, n_heads] 

901 """ 

902 

903get_act_patch_attn_head_q_all_pos = partial( 

904 generic_activation_patch, 

905 patch_setter=layer_head_vector_patch_setter, 

906 activation_name="q", 

907 index_axis_names=("layer", "head"), 

908) 

909get_act_patch_attn_head_q_all_pos.__doc__ = """ 

910 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] 

911 

912 See generic_activation_patch for a more detailed explanation of activation patching  

913 

914 Args: 

915 model: The relevant model