Coverage for transformer_lens/tools/analysis/attribution_patching.py: 92%

1"""Attribution patching — linearized activation patching on ``TransformerBridge``. 

2 

3Attribution patching estimates the causal effect of every model component on a 

4task metric with a *gradient-based linearization* of activation patching: rather 

5than one forward pass per intervention, it reads a single gradient cache. For a 

6clean/corrupt prompt pair it runs a clean forward (for ``a_clean``) and a corrupt 

7forward whose backward hooks capture ``g = d(metric)/d(a)`` (for ``a_corrupt`` and 

8its gradient), and scores each node with the 

9first-order Taylor estimate ``effect(node) = (a_clean - a_corrupt) . g``. Scores 

10over a batch of clean/corrupt pairs are averaged before ranking. 

11 

12Only the ``TransformerBridge`` API is targeted; TransformerLens v4 deprecates 

13``HookedTransformer``. 

14 

15Sign/direction convention (denoising form): the gradient is taken on the 

16*corrupt* run and the estimate points *toward* the clean activation, so a 

17positive score means patching that node from corrupt toward clean moves the 

18metric in the positive direction. An oracle-parity test maps this convention 

19onto a pinned reference rather than assuming the two agree. 

20 

21Memory note: gradients are retained only for hook points passing ``names_filter``. 

22Retaining gradients at every hook point roughly doubles cache memory, so callers 

23should filter to the hook families their analysis actually reads. 

24 

25Scope: this build ships node and edge granularity with plain attribution 

26(``ig_steps=1``). The integrated-gradient path (EAP-IG, ``ig_steps>1``) and 

27ablate-outside faithfulness are not implemented yet; their API is declared here, 

28and ``ig_steps>1`` raises :class:`NotImplementedError`, so downstream code can pin 

29against a stable surface now. 

30""" 

31 

32from __future__ import annotations 

33 

34from contextlib import contextmanager 

35from dataclasses import dataclass, field 

36from typing import Any, Callable, Iterator, Literal, Optional, Sequence, Union 

37 

38import torch 

39 

40from transformer_lens.tools.analysis._model_state import require_eval_mode 

41 

42MetricFn = Callable[[torch.Tensor], torch.Tensor] 

43NamesFilter = Union[str, Sequence[str], Callable[[str], bool], None] 

44 

45NodeKind = Literal[ 

46 "embed", "attn_head_out", "mlp_out", "q_input", "k_input", "v_input", "mlp_in", "logits" 

47] 

48Granularity = Literal["node", "edge"] 

49 

50 

51@dataclass 

52class GradientCache: 

53 """Activations and their metric-gradients from one forward + backward pass. 

54 

55 Attributes: 

56 activations: Detached activation tensor per cached hook name. 

57 gradients: ``d(metric)/d(activation)`` per cached hook name. 

58 metric: The scalar metric value at this run (detached). 

59 """ 

60 

61 activations: dict[str, torch.Tensor] 

62 gradients: dict[str, Optional[torch.Tensor]] 

63 metric: torch.Tensor 

64 

65 

66@dataclass(frozen=True) 

67class Node: 

68 """A node in the residual-stream computational graph. 

69 

70 Nodes are the typed, hashable keys the attribution sweep scores. Each node 

71 is identified by ``(kind, layer, position, head)``; ``kind`` selects the node 

72 family and constrains which of ``layer``/``head`` apply. Three kinds are 

73 *writers* -- they contribute a value into the residual stream: 

74 

75 - ``"embed"``: the token embedding write. ``layer`` and ``head`` are ``None``. 

76 - ``"attn_head_out"``: one attention head's output. ``layer`` and ``head`` set. 

77 - ``"mlp_out"``: one layer's MLP output. ``layer`` set, ``head`` is ``None``. 

78 

79 Five kinds are *readers* -- they consume the residual stream as an edge's 

80 destination (see :func:`enumerate_edges`): 

81 

82 - ``"q_input"`` / ``"k_input"`` / ``"v_input"``: one attention head's split 

83 Q/K/V input. ``layer`` and ``head`` set. 

84 - ``"mlp_in"``: one layer's MLP entry. ``layer`` set, ``head`` is ``None``. 

85 - ``"logits"``: the terminal readout of the final residual, read at 

86 ``blocks.{n_layers-1}.hook_resid_post``. ``layer`` is that final layer and 

87 ``head`` is ``None``. Every writer feeds this reader, so a writer's 

88 aggregate over its outgoing edges equals its direct node score. 

89 

90 ``position`` is the sequence index the node is read at. The invariants above 

91 are enforced in ``__post_init__`` so a malformed key raises rather than 

92 silently producing a wrong graph. 

93 """ 

94 

95 kind: NodeKind 

96 position: int 

97 layer: Optional[int] = None 

98 head: Optional[int] = None 

99 

100 def __post_init__(self) -> None: 

101 if self.kind == "embed": 

102 if self.layer is not None or self.head is not None: 

103 raise ValueError("embed nodes take neither layer nor head") 

104 elif self.kind == "attn_head_out": 

105 if self.layer is None or self.head is None: 

106 raise ValueError("attn_head_out nodes need both layer and head") 

107 elif self.kind == "mlp_out": 

108 if self.layer is None: 108 ↛ 109line 108 didn't jump to line 109 because the condition on line 108 was never true

109 raise ValueError("mlp_out nodes need a layer") 

110 if self.head is not None: 

111 raise ValueError("mlp_out nodes take no head") 

112 elif self.kind in ("q_input", "k_input", "v_input"): 

113 if self.layer is None or self.head is None: 

114 raise ValueError(f"{self.kind} nodes need both layer and head") 

115 elif self.kind == "mlp_in": 

116 if self.layer is None: 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true

117 raise ValueError("mlp_in nodes need a layer") 

118 if self.head is not None: 

119 raise ValueError("mlp_in nodes take no head") 

120 elif self.kind == "logits": 120 ↛ 126line 120 didn't jump to line 126 because the condition on line 120 was always true

121 if self.layer is None: 121 ↛ 122line 121 didn't jump to line 122 because the condition on line 121 was never true

122 raise ValueError("logits nodes need a layer") 

123 if self.head is not None: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 raise ValueError("logits nodes take no head") 

125 else: 

126 raise ValueError(f"unknown node kind {self.kind!r}") 

127 

128 @property 

129 def hook_name(self) -> str: 

130 """The cache hook point this node reads from. 

131 

132 Uses the standard ``TransformerBridge`` alias names (``hook_embed``, 

133 ``blocks.{l}.attn.hook_z``, ``blocks.{l}.hook_mlp_out``, 

134 ``blocks.{l}.attn.hook_q_input``/``hook_k_input``/``hook_v_input``, 

135 ``blocks.{l}.hook_mlp_in``, ``blocks.{l}.hook_resid_post``); the per-head 

136 nodes (``attn_head_out``, ``q_input``, ``k_input``, ``v_input``) slice 

137 head ``self.head`` out of the shared per-head tensor. 

138 """ 

139 if self.kind == "embed": 

140 return "hook_embed" 

141 if self.kind == "attn_head_out": 

142 return f"blocks.{self.layer}.attn.hook_z" 

143 if self.kind == "mlp_out": 

144 return f"blocks.{self.layer}.hook_mlp_out" 

145 if self.kind in ("q_input", "k_input", "v_input"): 

146 return f"blocks.{self.layer}.attn.hook_{self.kind}" 

147 if self.kind == "logits": 

148 return f"blocks.{self.layer}.hook_resid_post" 

149 return f"blocks.{self.layer}.hook_mlp_in" 

150 

151 

152@dataclass(frozen=True) 

153class EdgeAttributionConfig: 

154 """Configuration for an attribution-patching sweep. 

155 

156 The two axes are deliberately orthogonal: 

157 

158 - ``granularity`` selects what is scored: ``"node"`` scores each residual-stream 

159 write, ``"edge"`` scores each ``(source, destination)`` write->read pair. 

160 - ``ig_steps`` selects gradient fidelity. ``ig_steps=1`` is plain attribution 

161 patching / EAP: a single first-order Taylor gradient taken at the corrupt 

162 point. ``ig_steps>1`` is EAP-IG: the integrated gradient averaged over that 

163 many points along the corrupt->clean path, which corrects the gradient 

164 saturation that makes plain attribution unfaithful. 

165 

166 There is intentionally **no** ``method`` field. An earlier design had both a 

167 ``method`` enum (``"attribution"``/``"EAP"``/``"EAP-IG"``) and ``ig_steps``, 

168 which overlap: the method is fully determined by ``granularity`` and whether 

169 ``ig_steps`` exceeds 1. Collapsing them removes the invalid states (e.g. 

170 ``method="attribution", ig_steps=5``). 

171 

172 This build implements node and edge granularity with plain attribution. 

173 ``ig_steps>1`` is accepted by the type but raises :class:`NotImplementedError` 

174 at construction, so downstream code can import and reference this API now 

175 while the integrated-gradient path is not implemented yet. Once EAP-IG lands, 

176 the default flips to ``ig_steps=5`` (EAP-IG is the faithful default); until 

177 then the default is the only executable value, ``ig_steps=1``. 

178 

179 Attributes: 

180 granularity: ``"node"`` or ``"edge"``. Defaults to ``"node"``. 

181 ig_steps: Integrated-gradient path steps (``>=1``). Defaults to ``1``. 

182 """ 

183 

184 granularity: Granularity = "node" 

185 ig_steps: int = 1 

186 

187 def __post_init__(self) -> None: 

188 if self.ig_steps < 1: 

189 raise ValueError(f"ig_steps must be >= 1, got {self.ig_steps}") 

190 if self.ig_steps > 1: 

191 raise NotImplementedError( 

192 "ig_steps>1 (EAP-IG integrated gradients) is not implemented yet; this " 

193 "build supports ig_steps=1 (plain attribution) only." 

194 ) 

195 

196 

197@dataclass 

198class AttributionResult: 

199 """Scored output of an attribution-patching sweep. 

200 

201 Attributes: 

202 node_scores: Signed first-order effect estimate per node, 

203 ``(a_clean - a_corrupt) . d(metric)/d(a)``. A positive score means 

204 patching that node from corrupt toward clean moves the metric in the 

205 positive direction (the denoising convention pinned in the module 

206 docstring). For an edge-granularity sweep this is instead each 

207 writer's aggregate over its own outgoing edge scores. Because the 

208 graph has a terminal logits reader that every writer feeds, that 

209 aggregate equals the writer's direct node score (see 

210 :func:`attribution_patch`). 

211 edge_scores: Per-edge effect estimate keyed by ``(source, destination)``. 

212 Populated for an edge-granularity sweep; empty for a node sweep. 

213 """ 

214 

215 node_scores: dict[Node, float] 

216 edge_scores: dict[tuple[Node, Node], float] = field(default_factory=dict) 

217 

218 def top_nodes(self, k: int = 10) -> list[tuple[Node, float]]: 

219 """The ``k`` nodes with the largest effect magnitude, strongest first. 

220 

221 Ranking is by absolute score: a node with a large negative effect is as 

222 causally important as one with a large positive effect, so magnitude — not 

223 signed value — orders the circuit. Ties keep enumeration order (stable 

224 sort). Requesting more than the available nodes returns all of them. 

225 """ 

226 ranked = sorted(self.node_scores.items(), key=lambda item: abs(item[1]), reverse=True) 

227 return ranked[:k] 

228 

229 def top_edges(self, k: int = 10) -> list[tuple[Node, Node, float]]: 

230 """The ``k`` edges with the largest effect magnitude, strongest first. 

231 

232 Ranking is by absolute score, matching ``top_nodes``: a large negative 

233 edge effect is as causally important as a large positive one. Ties keep 

234 enumeration order (stable sort). Requesting more than the available 

235 edges returns all of them. 

236 """ 

237 ranked = sorted(self.edge_scores.items(), key=lambda item: abs(item[1]), reverse=True) 

238 return [(writer, reader, score) for (writer, reader), score in ranked[:k]] 

239 

240 

241def _required_hook_names(n_layers: int, granularity: Granularity = "node") -> list[str]: 

242 """Hook points a sweep at ``granularity`` reads. 

243 

244 Node granularity needs the embed write plus each layer's attn-z and 

245 mlp-out. Edge granularity additionally needs the per-head hook points on 

246 both sides of an edge into or out of an attention head: ``attn.hook_result`` 

247 (writer -- a head's own contribution before the sum into the residual 

248 stream) and the split ``attn.hook_q_input``/``hook_k_input``/``hook_v_input`` 

249 (reader -- the residual each head's Q/K/V projection reads separately). 

250 """ 

251 names = ["hook_embed"] 

252 for layer in range(n_layers): 

253 names.append(f"blocks.{layer}.attn.hook_z") 

254 names.append(f"blocks.{layer}.hook_mlp_out") 

255 if granularity == "edge": 

256 names.append(f"blocks.{layer}.attn.hook_result") 

257 names.append(f"blocks.{layer}.attn.hook_q_input") 

258 names.append(f"blocks.{layer}.attn.hook_k_input") 

259 names.append(f"blocks.{layer}.attn.hook_v_input") 

260 return names 

261 

262 

263def _ensure_edge_hook_flags(model: Any) -> None: 

264 """Turn on the Bridge flags edge granularity's hook points require. 

265 

266 ``attn.hook_result``, the split ``attn.hook_q_input``/``hook_k_input``/ 

267 ``hook_v_input``, and ``hook_mlp_in`` all exist on the Bridge 

268 unconditionally but only fire when their owning flag 

269 (``cfg.use_attn_result`` / ``cfg.use_split_qkv_input`` / 

270 ``cfg.use_hook_mlp_in``) is on, so an edge sweep must enable all three 

271 before caching or the writer- and reader-side hook points it needs never 

272 populate. 

273 

274 Memory caveat: enabling ``use_attn_result``/``use_split_qkv_input`` makes 

275 every cached per-head tensor ``[batch, seq, n_heads, d_model]`` instead of 

276 the summed ``[batch, seq, d_model]`` residual. That is fine on a model the 

277 size of gpt2-small; it does not scale to models with many heads or layers. 

278 

279 Args: 

280 model: A ``TransformerBridge`` (or compatible) exposing ``cfg`` and 

281 ``set_use_attn_result``/``set_use_split_qkv_input``/ 

282 ``set_use_hook_mlp_in``. 

283 """ 

284 if not model.cfg.use_attn_result: 284 ↛ 286line 284 didn't jump to line 286 because the condition on line 284 was always true

285 model.set_use_attn_result(True) 

286 if not model.cfg.use_split_qkv_input: 286 ↛ 288line 286 didn't jump to line 288 because the condition on line 286 was always true

287 model.set_use_split_qkv_input(True) 

288 if not model.cfg.use_hook_mlp_in: 288 ↛ exitline 288 didn't return from function '_ensure_edge_hook_flags' because the condition on line 288 was always true

289 model.set_use_hook_mlp_in(True) 

290 

291 

292@contextmanager 

293def _edge_hook_flags(model: Any) -> Iterator[None]: 

294 """Enable the Bridge flags an edge sweep needs, then restore the caller's state. 

295 

296 ``attn.hook_result``, the split ``attn.hook_q_input``/``hook_k_input``/ 

297 ``hook_v_input``, and ``hook_mlp_in`` all exist on the Bridge unconditionally 

298 but only fire when their owning flag (``cfg.use_attn_result`` / 

299 ``cfg.use_split_qkv_input`` / ``cfg.use_hook_mlp_in``) is on, so an edge sweep 

300 must enable all three before caching or the writer- and reader-side hook 

301 points it needs never populate. 

302 

303 ``use_split_qkv_input`` is mutually exclusive with ``use_attn_in``, so a 

304 caller who arrives with ``use_attn_in`` on would otherwise trip the 

305 exclusivity error. This turns ``use_attn_in`` off before enabling the split 

306 input, and restores it after ``use_split_qkv_input`` has been turned back off. 

307 

308 Invariant: the caller's flag state (``use_attn_result``, 

309 ``use_split_qkv_input``, ``use_hook_mlp_in``, ``use_attn_in``) is unchanged on 

310 return, including when the body raises. Restoration runs in a ``finally`` 

311 block so a raise mid-sweep -- or in the caller's own later code -- cannot 

312 leave the per-head tensors materialized on the model. 

313 

314 Memory caveat: enabling ``use_attn_result``/``use_split_qkv_input`` makes 

315 every cached per-head tensor ``[batch, seq, n_heads, d_model]`` instead of 

316 the summed ``[batch, seq, d_model]`` residual. That is fine on a model the 

317 size of gpt2-small; it does not scale to models with many heads or layers. 

318 

319 Args: 

320 model: A ``TransformerBridge`` (or compatible) exposing ``cfg`` and 

321 ``set_use_attn_result``/``set_use_split_qkv_input``/ 

322 ``set_use_hook_mlp_in``/``set_use_attn_in``. 

323 """ 

324 cfg = model.cfg 

325 saved_attn_result = cfg.use_attn_result 

326 saved_split_qkv_input = cfg.use_split_qkv_input 

327 saved_hook_mlp_in = cfg.use_hook_mlp_in 

328 saved_attn_in = bool(getattr(cfg, "use_attn_in", False)) 

329 try: 

330 # use_attn_in and use_split_qkv_input are mutually exclusive; clear 

331 # use_attn_in first so enabling the split input cannot raise. 

332 if saved_attn_in: 

333 model.set_use_attn_in(False) 

334 if not cfg.use_attn_result: 

335 model.set_use_attn_result(True) 

336 if not cfg.use_split_qkv_input: 336 ↛ 338line 336 didn't jump to line 338 because the condition on line 336 was always true

337 model.set_use_split_qkv_input(True) 

338 if not cfg.use_hook_mlp_in: 

339 model.set_use_hook_mlp_in(True) 

340 yield 

341 finally: 

342 model.set_use_attn_result(saved_attn_result) 

343 model.set_use_split_qkv_input(saved_split_qkv_input) 

344 model.set_use_hook_mlp_in(saved_hook_mlp_in) 

345 # Re-enabling use_attn_in requires use_split_qkv_input already off; the 

346 # line above restored it, and a caller with use_attn_in on cannot also 

347 # have had use_split_qkv_input on, so this cannot trip the exclusivity. 

348 if saved_attn_in: 

349 model.set_use_attn_in(True) 

350 

351 

352def _check_required_hooks( 

353 cache: GradientCache, required_names: Sequence[str], graph: str, hint: str 

354) -> None: 

355 """Raise if any of ``required_names`` is absent from ``cache.activations``. 

356 

357 Shared by every granularity's graph-construction step, so a cache built 

358 with too narrow a ``names_filter`` -- or one produced while a required 

359 Bridge flag was off -- fails loudly instead of silently producing a 

360 truncated graph. 

361 """ 

362 missing = [name for name in required_names if name not in cache.activations] 

363 if missing: 

364 raise ValueError( 

365 f"{graph} requires hook points missing from the cache: " 

366 + ", ".join(missing) 

367 + f". {hint}" 

368 ) 

369 

370 

371def enumerate_nodes(model: Any, cache: GradientCache) -> list[Node]: 

372 """Enumerate the full node-granularity graph from the Bridge hook graph. 

373 

374 The graph is *explicit*: for ``n_layers`` layers it always contains the embed 

375 write, every attention head's output, and every layer's MLP output, at every 

376 sequence position. Sequence length and head count are read from the cached 

377 tensor shapes; ``n_layers`` from ``model.cfg``. 

378 

379 Args: 

380 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``. 

381 cache: A :class:`GradientCache` holding at least the required hook points. 

382 

383 Returns: 

384 The node list, ordered embed-then-layerwise for deterministic ranking. 

385 

386 Raises: 

387 ValueError: if any required hook point is absent from ``cache`` — the 

388 graph is never silently truncated. 

389 """ 

390 n_layers = int(model.cfg.n_layers) 

391 _check_required_hooks( 

392 cache, 

393 _required_hook_names(n_layers), 

394 "node graph", 

395 "Cache with a names_filter that keeps hook_embed, blocks.*.attn.hook_z, " 

396 "and blocks.*.hook_mlp_out.", 

397 ) 

398 

399 seq_len = cache.activations["hook_embed"].shape[1] 

400 

401 nodes: list[Node] = [Node(kind="embed", position=position) for position in range(seq_len)] 

402 for layer in range(n_layers): 

403 n_heads = cache.activations[f"blocks.{layer}.attn.hook_z"].shape[2] 

404 for position in range(seq_len): 

405 for head in range(n_heads): 

406 nodes.append(Node(kind="attn_head_out", layer=layer, head=head, position=position)) 

407 nodes.append(Node(kind="mlp_out", layer=layer, position=position)) 

408 return nodes 

409 

410 

411def _assert_edges_unique(edges: Sequence[tuple[Node, Node]]) -> None: 

412 """Raise if any writer -> reader pair appears more than once in ``edges``. 

413 

414 A writer feeding two distinct readers (a head's output feeding both the 

415 next layer's attention input and this layer's MLP input, say) is two 

416 edges; this guards the enumeration against a construction bug that 

417 collapses or duplicates a single ``(writer, reader)`` pair instead. 

418 """ 

419 seen: set[tuple[Node, Node]] = set() 

420 for edge in edges: 

421 if edge in seen: 

422 raise ValueError(f"edge {edge} enumerated more than once") 

423 seen.add(edge) 

424 

425 

426def _required_edge_reader_hook_names(n_layers: int) -> list[str]: 

427 """The reader hook points edge enumeration additionally requires. 

428 

429 ``_required_hook_names(..., granularity="edge")`` covers the per-head 

430 attention hooks a writer/reader pair into or out of a head needs. Edge 

431 enumeration reads two reader points those miss: 

432 

433 - each layer's MLP entry, ``attn.hook_mlp_in``'s layer-level sibling 

434 ``hook_mlp_in``, gated on ``cfg.use_hook_mlp_in`` the same way the per-head 

435 hooks are gated on their own flags, and 

436 - the terminal ``blocks.{n_layers-1}.hook_resid_post``, where the logits 

437 reader takes its gradient. This final residual hook fires unconditionally, 

438 so no Bridge flag gates it. 

439 """ 

440 names = [f"blocks.{layer}.hook_mlp_in" for layer in range(n_layers)] 

441 names.append(f"blocks.{n_layers - 1}.hook_resid_post") 

442 return names 

443 

444 

445def _edge_hook_names(n_layers: int) -> list[str]: 

446 """Every hook point an edge-granularity sweep must cache. 

447 

448 The per-head attention hooks from ``_required_hook_names(..., granularity="edge")`` 

449 plus each layer's MLP-entry reader hook and the terminal logits reader hook. 

450 """ 

451 return _required_hook_names(n_layers, granularity="edge") + _required_edge_reader_hook_names( 

452 n_layers 

453 ) 

454 

455 

456def enumerate_edges(model: Any, cache: GradientCache) -> list[tuple[Node, Node]]: 

457 """Enumerate every writer -> reader edge in the residual-stream graph. 

458 

459 At a fixed sequence position, the residual stream is a running sum: a 

460 reader (a head's split Q/K/V input, a layer's MLP entry, or the terminal 

461 logits readout) is fed by every writer (the embed write, every attention 

462 head's output, every layer's MLP output) that precedes it. Building the 

463 graph position-by-position tracks which writers are "available" so far and 

464 connects each new reader to all of them, then adds that layer's writers to 

465 the available set before moving on -- so a writer never edges to a reader 

466 upstream of it, and a writer feeding both a direct edge and a through-MLP 

467 edge produces two distinct ``(u, v)`` pairs rather than one summed together. 

468 

469 A terminal ``logits`` reader (read at the final ``hook_resid_post``) closes 

470 the graph: after the per-layer loop every remaining writer -- including the 

471 final layer's MLP output, which no per-layer reader sees -- edges to it. That 

472 edge carries the writer's direct skip-connection contribution to the metric, 

473 so a writer's aggregate over its outgoing edges equals its direct node score. 

474 

475 On ``cfg.parallel_attn_mlp`` models (Pythia, GPT-J, Falcon, Phi) the MLP 

476 reads the layer input, not the post-attention residual, so a layer's own 

477 heads are not writers into that layer's MLP; those same-layer head->mlp_in 

478 edges are dropped while the heads still feed later readers and the logits 

479 reader. 

480 

481 Args: 

482 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers`` 

483 and, optionally, ``cfg.parallel_attn_mlp``. 

484 cache: A :class:`GradientCache` holding at least the required hook points 

485 for edge granularity. 

486 

487 Returns: 

488 The edge list as ``(writer, reader)`` node pairs; no pair repeats. 

489 

490 Raises: 

491 ValueError: if any required hook point is absent from ``cache`` -- the 

492 graph is never silently truncated. 

493 """ 

494 n_layers = int(model.cfg.n_layers) 

495 _check_required_hooks( 

496 cache, 

497 _edge_hook_names(n_layers), 

498 "edge graph", 

499 "Cache with a names_filter that keeps the edge-granularity hook set: " 

500 "hook_embed, blocks.*.attn.hook_z, blocks.*.hook_mlp_out, " 

501 "blocks.*.attn.hook_result, blocks.*.attn.hook_q_input, " 

502 "blocks.*.attn.hook_k_input, blocks.*.attn.hook_v_input, " 

503 "blocks.*.hook_mlp_in, and blocks.{n_layers-1}.hook_resid_post.", 

504 ) 

505 

506 parallel_attn_mlp = bool(getattr(model.cfg, "parallel_attn_mlp", False)) 

507 

508 seq_len = cache.activations["hook_embed"].shape[1] 

509 edges: list[tuple[Node, Node]] = [] 

510 

511 for position in range(seq_len): 

512 available: list[Node] = [Node(kind="embed", position=position)] 

513 for layer in range(n_layers): 

514 n_heads = cache.activations[f"blocks.{layer}.attn.hook_z"].shape[2] 

515 

516 attn_reader_kinds: tuple[NodeKind, NodeKind, NodeKind] = ( 

517 "q_input", 

518 "k_input", 

519 "v_input", 

520 ) 

521 attn_readers = [ 

522 Node(kind=kind, layer=layer, head=head, position=position) 

523 for kind in attn_reader_kinds 

524 for head in range(n_heads) 

525 ] 

526 for reader in attn_readers: 

527 edges.extend((writer, reader) for writer in available) 

528 

529 layer_heads = [ 

530 Node(kind="attn_head_out", layer=layer, head=head, position=position) 

531 for head in range(n_heads) 

532 ] 

533 

534 mlp_reader = Node(kind="mlp_in", layer=layer, position=position) 

535 if parallel_attn_mlp: 

536 # The MLP reads the layer input, not the post-attention residual, 

537 # so this layer's heads are not writers into its MLP. They still 

538 # become available to later readers and the logits reader. 

539 edges.extend((writer, mlp_reader) for writer in available) 

540 available = available + layer_heads 

541 else: 

542 available = available + layer_heads 

543 edges.extend((writer, mlp_reader) for writer in available) 

544 

545 available = available + [Node(kind="mlp_out", layer=layer, position=position)] 

546 

547 logits_reader = Node(kind="logits", layer=n_layers - 1, position=position) 

548 edges.extend((writer, logits_reader) for writer in available) 

549 

550 _assert_edges_unique(edges) 

551 return edges 

552 

553 

554def _as_predicate(names_filter: NamesFilter) -> Callable[[str], bool]: 

555 if names_filter is None: 555 ↛ 556line 555 didn't jump to line 556 because the condition on line 555 was never true

556 return lambda name: True 

557 if isinstance(names_filter, str): 557 ↛ 558line 557 didn't jump to line 558 because the condition on line 557 was never true

558 target = names_filter 

559 return lambda name: name == target 

560 if callable(names_filter): 560 ↛ 561line 560 didn't jump to line 561 because the condition on line 560 was never true

561 return names_filter 

562 wanted = set(names_filter) 

563 return lambda name: name in wanted 

564 

565 

566def cache_activation_and_gradient( 

567 model: Any, 

568 tokens: torch.Tensor, 

569 metric_fn: MetricFn, 

570 names_filter: NamesFilter = None, 

571 compute_gradient: bool = True, 

572) -> GradientCache: 

573 """Run one forward and capture activations plus (optionally) their metric-gradients. 

574 

575 Registers a forward hook and a backward hook at each cached point, runs one 

576 grad-enabled forward, then drives a single backward with 

577 :func:`torch.autograd.grad` to fire the backward hooks. 

578 ``run_with_cache(..., incl_bwd=True)`` only backpropagates the model's own 

579 scalar output, so a custom (non-scalar-output) metric such as a logit-diff 

580 needs the backward driven here. 

581 

582 Gradients come from the backward hooks, not from ``.grad`` on the cached 

583 tensors. ``TransformerBridge`` reshapes the tensor handed to a forward hook at 

584 a converted point (``attn.hook_z``, ``hook_q/k/v``, ``hook_attn_out``), so that 

585 tensor is a view the model's forward never consumes: ``retain_grad()`` on it is 

586 inert and its ``.grad`` stays ``None``. A backward hook goes through the same 

587 conversion and delivers the real gradient in canonical shape. Driving the 

588 backward with :func:`torch.autograd.grad` instead of ``metric.backward()`` 

589 keeps it off every parameter's ``.grad`` buffer — no caller grads clobbered, no 

590 model-sized buffer allocated. 

591 

592 With ``compute_gradient=False`` the backward is skipped entirely: only forward 

593 hooks are registered, no backward is driven, and every cached point's gradient 

594 is ``None``. Callers that read activations only (the clean pass of 

595 :func:`attribution_patch`, which pairs these with the *corrupt* run's gradients) 

596 use this to run a plain forward instead of a needless forward + backward. 

597 

598 Args: 

599 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``, 

600 ``hook_dict``, and the ``hooks()`` context manager. 

601 tokens: Input token ids for a single forward pass. 

602 metric_fn: Maps the model logits to a scalar to differentiate. 

603 names_filter: Restricts which hook points are cached (and have gradients 

604 retained). ``None`` (the default) caches the node-granularity hook set — 

605 ``hook_embed`` plus each layer's ``attn.hook_z`` and ``hook_mlp_out``; 

606 pass an explicit filter to cache any hook point outside that set. On a 

607 real Bridge, ``None`` cannot mean "every hook point": the gated points 

608 (``hook_mlp_in``, ``attn.hook_result``, the split-QKV inputs) that 

609 ``hook_dict`` exposes raise in ``add_hook`` unless their ``set_use_*`` 

610 flag is on. 

611 compute_gradient: When ``True`` (default) capture gradients via backward 

612 hooks. When ``False`` run an activation-only forward and leave every 

613 gradient ``None``. 

614 

615 Returns: 

616 A :class:`GradientCache` with per-hook activations, and gradients when 

617 ``compute_gradient`` is ``True`` (all ``None`` otherwise). 

618 """ 

619 if compute_gradient and not torch.is_grad_enabled(): 619 ↛ 620line 619 didn't jump to line 620 because the condition on line 619 was never true

620 raise ValueError( 

621 "cache_activation_and_gradient needs autograd, but gradient tracking " 

622 "is off (torch.no_grad(), set_grad_enabled(False), or inference mode)." 

623 ) 

624 

625 if names_filter is None: 

626 # Default to the node-granularity hook set rather than every hook point: a 

627 # real Bridge's hook_dict holds gated points (hook_mlp_in, attn.hook_result, 

628 # the split-QKV inputs) that add_hook rejects unless their set_use_* flag is 

629 # on, so caching "everything" raises. Pass an explicit filter to reach them. 

630 names_filter = _required_hook_names(int(model.cfg.n_layers)) 

631 predicate = _as_predicate(names_filter) 

632 names = [name for name in model.hook_dict if predicate(name)] 

633 if not names: 633 ↛ 634line 633 didn't jump to line 634 because the condition on line 633 was never true

634 raise ValueError("names_filter matched no hook points") 

635 

636 live: dict[str, torch.Tensor] = {} 

637 activations: dict[str, torch.Tensor] = {} 

638 gradients: dict[str, Optional[torch.Tensor]] = {} 

639 

640 def make_fwd_hook(name: str) -> Callable[..., None]: 

641 def hook(tensor: torch.Tensor, *, hook: Any) -> None: 

642 del hook 

643 if isinstance(tensor, torch.Tensor): 643 ↛ 646line 643 didn't jump to line 646 because the condition on line 643 was always true

644 live[name] = tensor 

645 activations[name] = tensor.detach().clone() 

646 return None 

647 

648 return hook 

649 

650 def make_bwd_hook(name: str) -> Callable[..., None]: 

651 def hook(grad: torch.Tensor, *, hook: Any) -> None: 

652 del hook 

653 if isinstance(grad, torch.Tensor): 653 ↛ 655line 653 didn't jump to line 655 because the condition on line 653 was always true

654 gradients[name] = grad.detach().clone() 

655 return None 

656 

657 return hook 

658 

659 fwd_hooks = [(name, make_fwd_hook(name)) for name in names] 

660 

661 if not compute_gradient: 

662 with model.hooks(fwd_hooks=fwd_hooks): 

663 logits = model(tokens) 

664 metric = metric_fn(logits) 

665 if metric.dim() != 0: 665 ↛ 666line 665 didn't jump to line 666 because the condition on line 665 was never true

666 raise ValueError( 

667 f"metric_fn must return a scalar tensor, got shape {tuple(metric.shape)}" 

668 ) 

669 return GradientCache( 

670 activations=activations, 

671 gradients={name: None for name in activations}, 

672 metric=metric.detach(), 

673 ) 

674 

675 bwd_hooks = [(name, make_bwd_hook(name)) for name in names] 

676 

677 with model.hooks(fwd_hooks=fwd_hooks, bwd_hooks=bwd_hooks): 

678 logits = model(tokens) 

679 metric = metric_fn(logits) 

680 if metric.dim() != 0: 680 ↛ 681line 680 didn't jump to line 681 because the condition on line 680 was never true

681 raise ValueError( 

682 f"metric_fn must return a scalar tensor, got shape {tuple(metric.shape)}" 

683 ) 

684 captured_names = [name for name in names if name in live] 

685 # torch.autograd.grad only *drives* the backward; the gradients we keep are 

686 # the converted-shape ones the backward hooks write into `gradients`. Unlike 

687 # metric.backward() it touches no parameter `.grad` buffer. 

688 returned = ( 

689 torch.autograd.grad( 

690 metric, 

691 inputs=[live[name] for name in captured_names], 

692 allow_unused=True, 

693 retain_graph=False, 

694 ) 

695 if captured_names 

696 else () 

697 ) 

698 

699 # The most-upstream cached point's own backward hook never fires — nothing 

700 # requested is upstream of it, so the backward stops at it — but its cached 

701 # tensor is on-path and unconverted, so torch.autograd.grad returns that 

702 # gradient directly. Backfill any point the hooks missed from that return. 

703 for name, grad in zip(captured_names, returned): 

704 if gradients.get(name) is None: 

705 gradients[name] = None if grad is None else grad.detach().clone() 

706 

707 return GradientCache(activations=activations, gradients=gradients, metric=metric.detach()) 

708 

709 

710def _node_effects( 

711 clean_cache: GradientCache, 

712 corrupt_cache: GradientCache, 

713 nodes: Sequence[Node], 

714) -> dict[Node, float]: 

715 """Score every node with the first-order attribution ``(a_clean - a_corrupt) . g``. 

716 

717 The gradient ``g`` is taken from the corrupt cache (denoising convention). For 

718 each node the feature dimension (``d_model``, or ``d_head`` for an attention 

719 head) is contracted at the node's position/head, giving one signed scalar. 

720 """ 

721 scores: dict[Node, float] = {} 

722 for node in nodes: 

723 name = node.hook_name 

724 grad = corrupt_cache.gradients.get(name) 

725 if grad is None: 725 ↛ 726line 725 didn't jump to line 726 because the condition on line 725 was never true

726 raise ValueError( 

727 f"node {node} reads {name!r}, but the corrupt cache holds no gradient " 

728 "there; cache with a names_filter that retains this hook point." 

729 ) 

730 delta = clean_cache.activations[name] - corrupt_cache.activations[name] 

731 contribution = delta * grad 

732 if node.kind == "attn_head_out": 

733 value = contribution[0, node.position, node.head].sum() 

734 else: 

735 value = contribution[0, node.position].sum() 

736 scores[node] = float(value) 

737 return scores 

738 

739 

740def _writer_hook_name(node: Node) -> str: 

741 """The residual-stream hook point holding a writer node's own contribution. 

742 

743 Distinct from ``Node.hook_name``: an ``attn_head_out`` node's ``hook_name`` 

744 resolves to ``attn.hook_z``, the pre-``hook_result`` value node granularity 

745 scores. An edge's writer contribution must instead be measured in the same 

746 ``d_model`` space a reader's gradient lives in, which is ``attn.hook_result`` 

747 -- the per-head decomposition of the head's contribution after it is 

748 projected into the residual stream. 

749 """ 

750 if node.kind == "embed": 

751 return "hook_embed" 

752 if node.kind == "attn_head_out": 

753 return f"blocks.{node.layer}.attn.hook_result" 

754 if node.kind == "mlp_out": 754 ↛ 756line 754 didn't jump to line 756 because the condition on line 754 was always true

755 return f"blocks.{node.layer}.hook_mlp_out" 

756 raise ValueError(f"{node.kind} is a reader kind and has no writer contribution") 

757 

758 

759def _edge_effects( 

760 clean_cache: GradientCache, 

761 corrupt_cache: GradientCache, 

762 edges: Sequence[tuple[Node, Node]], 

763) -> dict[tuple[Node, Node], float]: 

764 """Score every edge with ``(a_clean[u] - a_corrupt[u]) . d(metric)/d(input of v)``. 

765 

766 Mirrors ``_node_effects``: the delta is the writer's own residual 

767 contribution (clean minus corrupt cache), dotted with the reader's 

768 corrupt-run gradient -- the same denoising convention ``_node_effects`` 

769 uses. Unlike a node score, the delta and the gradient are read from two 

770 different hook points (the writer's and the reader's), since an edge 

771 measures how much of one component's output reaches another component's 

772 input. A ``logits`` reader takes its gradient at the final 

773 ``hook_resid_post`` and contracts over ``d_model`` at its position, the same 

774 shape path as an ``mlp_in`` reader. 

775 

776 Two structural facts keep this off a per-edge recompute: a writer hook's 

777 clean-minus-corrupt delta is the same tensor for every edge that writer 

778 feeds, and every edge into one reader shares that reader's gradient vector. 

779 Each writer delta is therefore computed once, and each reader's incoming 

780 edges are scored with a single batched matrix-vector product over the 

781 stacked writer deltas -- reducing to the same per-edge scalars a 

782 ``float((delta_vec * grad_vec).sum())`` per edge would, reader gradient 

783 reused across all of that reader's edges. 

784 """ 

785 scores: dict[tuple[Node, Node], float] = {} 

786 writer_deltas: dict[str, torch.Tensor] = {} 

787 

788 def writer_delta_vec(writer: Node) -> torch.Tensor: 

789 name = _writer_hook_name(writer) 

790 delta = writer_deltas.get(name) 

791 if delta is None: 

792 delta = clean_cache.activations[name] - corrupt_cache.activations[name] 

793 writer_deltas[name] = delta 

794 if writer.kind == "attn_head_out": 

795 return delta[0, writer.position, writer.head] 

796 return delta[0, writer.position] 

797 

798 writers_by_reader: dict[Node, list[Node]] = {} 

799 reader_order: list[Node] = [] 

800 for writer, reader in edges: 

801 bucket = writers_by_reader.get(reader) 

802 if bucket is None: 

803 writers_by_reader[reader] = bucket = [] 

804 reader_order.append(reader) 

805 bucket.append(writer) 

806 

807 for reader in reader_order: 

808 reader_name = reader.hook_name 

809 grad = corrupt_cache.gradients.get(reader_name) 

810 writers = writers_by_reader[reader] 

811 if grad is None: 811 ↛ 812line 811 didn't jump to line 812 because the condition on line 811 was never true

812 raise ValueError( 

813 f"edge {(writers[0], reader)} reads its gradient at {reader_name!r}, but " 

814 "the corrupt cache holds none there; cache with a names_filter that " 

815 "retains this hook point." 

816 ) 

817 if reader.kind in ("q_input", "k_input", "v_input"): 

818 grad_vec = grad[0, reader.position, reader.head] 

819 else: 

820 grad_vec = grad[0, reader.position] 

821 delta_matrix = torch.stack([writer_delta_vec(writer) for writer in writers]) 

822 edge_values = delta_matrix @ grad_vec 

823 for writer, value in zip(writers, edge_values.tolist()): 

824 scores[(writer, reader)] = value 

825 return scores 

826 

827 

828def _aggregate_edge_scores_to_writer_nodes( 

829 edge_scores: dict[tuple[Node, Node], float], 

830) -> dict[Node, float]: 

831 """Sum each writer's outgoing edge scores into that writer's aggregate node score. 

832 

833 A writer's aggregate is the sum of its effects along every edge it feeds. 

834 enumerate_edges gives every writer an edge to the terminal logits reader, so 

835 the aggregate includes the writer's direct skip-connection contribution to 

836 the metric -- the part of its residual-stream write that no intermediate 

837 component reads, only the final readout does -- and therefore equals the 

838 quantity a node-granularity sweep measures directly at the writer's own hook 

839 point. Every writer has at least that one outgoing edge, including the final 

840 layer's MLP output, whose only reader is the logits terminal, so no writer is 

841 missing from the aggregate. 

842 """ 

843 totals: dict[Node, float] = {} 

844 for (writer, _reader), score in edge_scores.items(): 

845 totals[writer] = totals.get(writer, 0.0) + score 

846 return totals 

847 

848 

849def attribution_patch( 

850 model: Any, 

851 clean: torch.Tensor, 

852 corrupt: torch.Tensor, 

853 metric_fn: MetricFn, 

854 config: EdgeAttributionConfig = EdgeAttributionConfig(), 

855) -> AttributionResult: 

856 """Estimate every component's causal effect on ``metric_fn`` in two forwards + one backward. 

857 

858 For each clean/corrupt pair this runs a clean forward (for ``a_clean``) and a 

859 corrupt forward whose backward hooks capture ``g = d(metric)/d(a)`` (for 

860 ``a_corrupt`` and its gradient). At node granularity (``config.granularity == 

861 "node"``) each node is scored with the first-order Taylor estimate 

862 ``effect(node) = (a_clean - a_corrupt) . g``. At edge granularity 

863 (``config.granularity == "edge"``) each writer -> reader edge is scored with 

864 ``effect(edge) = (a_clean[writer] - a_corrupt[writer]) . d(metric)/d(input of 

865 reader)``, and ``node_scores`` holds each writer's aggregate effect (the sum 

866 of its outgoing edge scores). 

867 

868 Sign/direction convention (denoising form): gradients are taken on the *corrupt* 

869 run and the estimate points *toward* the clean activation, so a positive score 

870 means patching that node from corrupt toward clean moves the metric in the 

871 positive direction. An oracle-parity test maps this convention onto a pinned 

872 reference rather than assuming the two agree. 

873 

874 Dataset averaging: ``clean``/``corrupt`` may hold a batch of prompt pairs. Each 

875 pair is scored independently (per-example forward/backward, so its own 

876 reconstruction identity holds) and per-node (or per-edge) scores are averaged 

877 across the batch before ranking. 

878 

879 The model and every submodule must be in evaluation mode. Separate clean and 

880 corrupt forwards cannot produce meaningful activation differences if stochastic 

881 training layers such as dropout remain active. 

882 

883 Args: 

884 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``, 

885 ``hook_dict``, and ``hooks()``. 

886 clean: Clean token ids, shape ``[batch, seq]``. 

887 corrupt: Corrupt token ids, shape ``[batch, seq]``, paired row-by-row with 

888 ``clean``. 

889 metric_fn: Maps single-example logits to a scalar to differentiate. 

890 config: Sweep configuration. Node and edge granularity are both 

891 supported with plain attribution (``ig_steps=1``); ``ig_steps>1`` 

892 raises at construction. 

893 

894 Returns: 

895 An :class:`AttributionResult` whose ``node_scores`` are averaged over the 

896 batch. For an edge-granularity sweep, ``edge_scores`` is populated too and 

897 ``node_scores`` is the per-writer aggregate of those edge scores. The 

898 graph's terminal logits reader gives every writer an edge carrying its 

899 direct skip-connection contribution to the metric, so this aggregate 

900 equals the quantity a node-granularity sweep on the same model returns 

901 (including for the final layer's MLP output, whose only outgoing edge is 

902 the one to the logits reader). 

903 

904 Raises: 

905 ValueError: if ``clean``/``corrupt`` are not 2D, hold a different number of 

906 pairs, a pair tokenizes to different lengths (activations must align 

907 position-by-position), or the model or one of its submodules is in 

908 training mode. 

909 """ 

910 if clean.ndim != 2 or corrupt.ndim != 2: 910 ↛ 911line 910 didn't jump to line 911 because the condition on line 910 was never true