Coverage for transformer_lens/tools/analysis/attribution_patching.py: 95%
585 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Attribution patching — linearized activation patching on ``TransformerBridge``.
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.
12Only the ``TransformerBridge`` API is targeted.
14Sign/direction convention (denoising form): the gradient is taken on the
15*corrupt* run and the estimate points *toward* the clean activation, so a
16positive score means patching that node from corrupt toward clean moves the
17metric in the positive direction. An oracle-parity test maps this convention
18onto a pinned reference rather than assuming the two agree.
20Memory note: gradients are retained only for hook points passing ``names_filter``.
21Retaining gradients at every hook point roughly doubles cache memory, so callers
22should filter to the hook families their analysis actually reads.
24Scope: this build ships node and edge granularity with plain attribution
25(``ig_steps=1``), plus ablate-outside faithfulness. The integrated-gradient path
26(EAP-IG, ``ig_steps>1``) is not implemented yet; its API is declared here, and
27``ig_steps>1`` raises :class:`NotImplementedError`, so downstream code can pin
28against a stable surface now.
30Faithfulness intervenes rather than estimating: it ablates every edge outside a
31candidate circuit and reports how much of the clean-to-corrupt metric gap the
32circuit recovers. Because the residual stream is a running sum and the reader
33hooks fire before the layer norm, each reader's input is rebuilt exactly.
34"""
36from __future__ import annotations
38from contextlib import contextmanager
39from dataclasses import dataclass, field
40from typing import (
41 Any,
42 Callable,
43 Collection,
44 Iterator,
45 Literal,
46 Optional,
47 Sequence,
48 Union,
49 cast,
50 get_args,
51)
53import torch
55from transformer_lens.model_bridge.generalized_components.block import (
56 VARIANT_SUBMODULE_NAMES,
57)
58from transformer_lens.tools.analysis._model_state import require_eval_mode
60MetricFn = Callable[[torch.Tensor], torch.Tensor]
61NamesFilter = Union[str, Sequence[str], Callable[[str], bool], None]
63NodeKind = Literal[
64 "embed", "attn_head_out", "mlp_out", "q_input", "k_input", "v_input", "mlp_in", "logits"
65]
66Granularity = Literal["node", "edge"]
67AblationMode = Literal["corrupt", "mean"]
68EdgeClass = Literal["into_qk", "into_v", "into_mlp", "into_logits"]
71@dataclass
72class GradientCache:
73 """Activations and their metric-gradients from one forward + backward pass.
75 Attributes:
76 activations: Detached activation tensor per cached hook name.
77 gradients: ``d(metric)/d(activation)`` per cached hook name.
78 metric: The scalar metric value at this run (detached).
79 """
81 activations: dict[str, torch.Tensor]
82 gradients: dict[str, Optional[torch.Tensor]]
83 metric: torch.Tensor
86@dataclass(frozen=True)
87class Node:
88 """A node in the residual-stream computational graph.
90 Nodes are the typed, hashable keys the attribution sweep scores. Each node
91 is identified by ``(kind, layer, position, head)``; ``kind`` selects the node
92 family and constrains which of ``layer``/``head`` apply. Three kinds are
93 *writers* -- they contribute a value into the residual stream:
95 - ``"embed"``: the token embedding write. ``layer`` and ``head`` are ``None``.
96 - ``"attn_head_out"``: one attention head's output. ``layer`` and ``head`` set.
97 - ``"mlp_out"``: one layer's MLP output. ``layer`` set, ``head`` is ``None``.
99 Five kinds are *readers* -- they consume the residual stream as an edge's
100 destination (see :func:`enumerate_edges`):
102 - ``"q_input"`` / ``"k_input"`` / ``"v_input"``: one attention head's split
103 Q/K/V input. ``layer`` and ``head`` set.
104 - ``"mlp_in"``: one layer's MLP entry. ``layer`` set, ``head`` is ``None``.
105 - ``"logits"``: the terminal readout of the final residual, read at
106 ``blocks.{n_layers-1}.hook_resid_post``. ``layer`` is that final layer and
107 ``head`` is ``None``. Every writer feeds this reader, so a writer's
108 aggregate over its outgoing edges equals its direct node score.
110 ``position`` is the sequence index the node is read at. The invariants above
111 are enforced in ``__post_init__`` so a malformed key raises rather than
112 silently producing a wrong graph.
113 """
115 kind: NodeKind
116 position: int
117 layer: Optional[int] = None
118 head: Optional[int] = None
120 def __post_init__(self) -> None:
121 if self.kind == "embed":
122 if self.layer is not None or self.head is not None:
123 raise ValueError("embed nodes take neither layer nor head")
124 elif self.kind == "attn_head_out":
125 if self.layer is None or self.head is None:
126 raise ValueError("attn_head_out nodes need both layer and head")
127 elif self.kind == "mlp_out":
128 if self.layer is None: 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true
129 raise ValueError("mlp_out nodes need a layer")
130 if self.head is not None:
131 raise ValueError("mlp_out nodes take no head")
132 elif self.kind in ("q_input", "k_input", "v_input"):
133 if self.layer is None or self.head is None:
134 raise ValueError(f"{self.kind} nodes need both layer and head")
135 elif self.kind == "mlp_in":
136 if self.layer is None: 136 ↛ 137line 136 didn't jump to line 137 because the condition on line 136 was never true
137 raise ValueError("mlp_in nodes need a layer")
138 if self.head is not None:
139 raise ValueError("mlp_in nodes take no head")
140 elif self.kind == "logits": 140 ↛ 146line 140 didn't jump to line 146 because the condition on line 140 was always true
141 if self.layer is None: 141 ↛ 142line 141 didn't jump to line 142 because the condition on line 141 was never true
142 raise ValueError("logits nodes need a layer")
143 if self.head is not None: 143 ↛ 144line 143 didn't jump to line 144 because the condition on line 143 was never true
144 raise ValueError("logits nodes take no head")
145 else:
146 raise ValueError(f"unknown node kind {self.kind!r}")
148 @property
149 def hook_name(self) -> str:
150 """The cache hook point this node reads from.
152 Uses the standard ``TransformerBridge`` alias names (``hook_embed``,
153 ``blocks.{l}.attn.hook_z``, ``blocks.{l}.hook_mlp_out``,
154 ``blocks.{l}.attn.hook_q_input``/``hook_k_input``/``hook_v_input``,
155 ``blocks.{l}.hook_mlp_in``, ``blocks.{l}.hook_resid_post``); the per-head
156 nodes (``attn_head_out``, ``q_input``, ``k_input``, ``v_input``) slice
157 head ``self.head`` out of the shared per-head tensor.
158 """
159 if self.kind == "embed":
160 return "hook_embed"
161 if self.kind == "attn_head_out":
162 return f"blocks.{self.layer}.attn.hook_z"
163 if self.kind == "mlp_out":
164 return f"blocks.{self.layer}.hook_mlp_out"
165 if self.kind in ("q_input", "k_input", "v_input"):
166 return f"blocks.{self.layer}.attn.hook_{self.kind}"
167 if self.kind == "logits":
168 return f"blocks.{self.layer}.hook_resid_post"
169 return f"blocks.{self.layer}.hook_mlp_in"
172@dataclass(frozen=True)
173class EdgeAttributionConfig:
174 """Configuration for an attribution-patching sweep.
176 The two axes are deliberately orthogonal:
178 - ``granularity`` selects what is scored: ``"node"`` scores each residual-stream
179 write, ``"edge"`` scores each ``(source, destination)`` write->read pair.
180 - ``ig_steps`` selects gradient fidelity. ``ig_steps=1`` is plain attribution
181 patching / EAP: a single first-order Taylor gradient taken at the corrupt
182 point. ``ig_steps>1`` is EAP-IG: the integrated gradient averaged over that
183 many points along the corrupt->clean path, which corrects the gradient
184 saturation that makes plain attribution unfaithful.
186 There is intentionally **no** ``method`` field. An earlier design had both a
187 ``method`` enum (``"attribution"``/``"EAP"``/``"EAP-IG"``) and ``ig_steps``,
188 which overlap: the method is fully determined by ``granularity`` and whether
189 ``ig_steps`` exceeds 1. Collapsing them removes the invalid states (e.g.
190 ``method="attribution", ig_steps=5``).
192 This build implements node and edge granularity with plain attribution.
193 ``ig_steps>1`` is accepted by the type but raises :class:`NotImplementedError`
194 at construction, so downstream code can import and reference this API now
195 while the integrated-gradient path is not implemented. Once EAP-IG lands,
196 the default flips to ``ig_steps=5`` (EAP-IG is the faithful default); until
197 then the default is the only executable value, ``ig_steps=1``.
199 Attributes:
200 granularity: ``"node"`` or ``"edge"``. Defaults to ``"node"``.
201 ig_steps: Integrated-gradient path steps (``>=1``). Defaults to ``1``.
202 """
204 granularity: Granularity = "node"
205 ig_steps: int = 1
207 def __post_init__(self) -> None:
208 if self.granularity not in get_args(Granularity):
209 raise ValueError(
210 f"granularity must be one of {get_args(Granularity)}, got {self.granularity!r}"
211 )
212 if self.ig_steps < 1:
213 raise ValueError(f"ig_steps must be >= 1, got {self.ig_steps}")
214 if self.ig_steps > 1:
215 raise NotImplementedError(
216 "ig_steps>1 (EAP-IG integrated gradients) is not implemented yet; this "
217 "build supports ig_steps=1 (plain attribution) only."
218 )
221@dataclass
222class AttributionResult:
223 """Scored output of an attribution-patching sweep.
225 Attributes:
226 node_scores: Signed first-order effect estimate per node,
227 ``(a_clean - a_corrupt) . d(metric)/d(a)``. A positive score means
228 patching that node from corrupt toward clean moves the metric in the
229 positive direction (the denoising convention pinned in the module
230 docstring). For an edge-granularity sweep this is instead each
231 writer's aggregate over its own outgoing edge scores. Because the
232 graph has a terminal logits reader that every writer feeds, that
233 aggregate equals the writer's direct node score (see
234 :func:`attribution_patch`).
235 edge_scores: Per-edge effect estimate keyed by ``(source, destination)``.
236 Populated for an edge-granularity sweep; empty for a node sweep.
237 """
239 node_scores: dict[Node, float]
240 edge_scores: dict[tuple[Node, Node], float] = field(default_factory=dict)
242 def top_nodes(self, k: int = 10) -> list[tuple[Node, float]]:
243 """The ``k`` nodes with the largest effect magnitude, strongest first.
245 Ranking is by absolute score: a node with a large negative effect is as
246 causally important as one with a large positive effect, so magnitude — not
247 signed value — orders the circuit. Ties keep enumeration order (stable
248 sort). Requesting more than the available nodes returns all of them.
249 """
250 ranked = sorted(self.node_scores.items(), key=lambda item: abs(item[1]), reverse=True)
251 return ranked[:k]
253 def top_edges(self, k: int = 10) -> list[tuple[Node, Node, float]]:
254 """The ``k`` edges with the largest effect magnitude, strongest first.
256 Ranking is by absolute score, matching ``top_nodes``: a large negative
257 edge effect is as causally important as a large positive one. Ties keep
258 enumeration order (stable sort). Requesting more than the available
259 edges returns all of them.
260 """
261 ranked = sorted(self.edge_scores.items(), key=lambda item: abs(item[1]), reverse=True)
262 return [(writer, reader, score) for (writer, reader), score in ranked[:k]]
265@dataclass(frozen=True)
266class FaithfulnessConfig:
267 """Configuration for an ablate-outside faithfulness measurement.
269 Attributes:
270 ablation: How an out-of-circuit edge's writer is replaced.
271 ``"corrupt"`` (default) substitutes the corrupt run's own contribution;
272 ``"mean"`` substitutes that writer's mean over the held-out
273 ``mean_tokens`` batch passed to :func:`faithfulness`.
274 """
276 ablation: AblationMode = "corrupt"
278 def __post_init__(self) -> None:
279 if self.ablation not in get_args(AblationMode):
280 raise ValueError(
281 f"ablation must be one of {get_args(AblationMode)}, got {self.ablation!r}"
282 )
285@dataclass
286class FaithfulnessResult:
287 """How much of the clean-to-corrupt metric gap a candidate circuit recovers.
289 Attributes:
290 recovered: ``(m_ablated - m_corrupt) / (m_clean - m_corrupt)``, a fraction
291 (``1.0`` = clean metric, ``0.0`` = corrupt metric); values outside
292 ``[0, 1]`` are possible, since a circuit can overshoot the clean run.
293 full_metric: The clean run's metric, the top of the gap.
294 corrupt_metric: The corrupt run's metric, the bottom of the gap.
295 circuit_size: Number of edges kept.
296 total_edges: Number of edges in the graph.
297 edge_class_recovered: Recovery after also removing this class's retained
298 edges, with the same replacements and gap; a class absent from the
299 circuit has the aggregate recovery. Conditional dependence, not
300 additive and not bounded to ``[0, 1]``.
301 edge_class_counts: Full-graph edge counts per class, not retained counts.
302 """
304 recovered: float
305 full_metric: float
306 corrupt_metric: float
307 circuit_size: int
308 total_edges: int
309 edge_class_recovered: dict[EdgeClass, float] = field(default_factory=dict)
310 edge_class_counts: dict[EdgeClass, int] = field(default_factory=dict)
313def _required_hook_names(n_layers: int, granularity: Granularity = "node") -> list[str]:
314 """Hook points a sweep at ``granularity`` reads.
316 Node granularity needs the embed write plus each layer's attn-z and
317 mlp-out. Edge granularity additionally needs the per-head hook points on
318 both sides of an edge into or out of an attention head: ``attn.hook_result``
319 (writer -- a head's own contribution before the sum into the residual
320 stream) and the split ``attn.hook_q_input``/``hook_k_input``/``hook_v_input``
321 (reader -- the residual each head's Q/K/V projection reads separately).
322 """
323 names = ["hook_embed"]
324 for layer in range(n_layers):
325 names.append(f"blocks.{layer}.attn.hook_z")
326 names.append(f"blocks.{layer}.hook_mlp_out")
327 if granularity == "edge":
328 names.append(f"blocks.{layer}.attn.hook_result")
329 names.append(f"blocks.{layer}.attn.hook_q_input")
330 names.append(f"blocks.{layer}.attn.hook_k_input")
331 names.append(f"blocks.{layer}.attn.hook_v_input")
332 return names
335def _ensure_edge_hook_flags(model: Any) -> None:
336 """Turn on the Bridge flags edge granularity's hook points require.
338 ``attn.hook_result``, the split ``attn.hook_q_input``/``hook_k_input``/
339 ``hook_v_input``, and ``hook_mlp_in`` all exist on the Bridge
340 unconditionally but only fire when their owning flag
341 (``cfg.use_attn_result`` / ``cfg.use_split_qkv_input`` /
342 ``cfg.use_hook_mlp_in``) is on, so an edge sweep must enable all three
343 before caching or the writer- and reader-side hook points it needs never
344 populate.
346 Memory caveat: enabling ``use_attn_result``/``use_split_qkv_input`` makes
347 every cached per-head tensor ``[batch, seq, n_heads, d_model]`` instead of
348 the summed ``[batch, seq, d_model]`` residual. That is fine on a model the
349 size of gpt2-small; it does not scale to models with many heads or layers.
351 Args:
352 model: A ``TransformerBridge`` (or compatible) exposing ``cfg`` and
353 ``set_use_attn_result``/``set_use_split_qkv_input``/
354 ``set_use_hook_mlp_in``.
355 """
356 if not model.cfg.use_attn_result: 356 ↛ 358line 356 didn't jump to line 358 because the condition on line 356 was always true
357 model.set_use_attn_result(True)
358 if not model.cfg.use_split_qkv_input: 358 ↛ 360line 358 didn't jump to line 360 because the condition on line 358 was always true
359 model.set_use_split_qkv_input(True)
360 if not model.cfg.use_hook_mlp_in: 360 ↛ exitline 360 didn't return from function '_ensure_edge_hook_flags' because the condition on line 360 was always true
361 model.set_use_hook_mlp_in(True)
364@contextmanager
365def _edge_hook_flags(model: Any) -> Iterator[None]:
366 """Enable the Bridge flags an edge sweep needs, then restore the caller's state.
368 ``attn.hook_result``, the split ``attn.hook_q_input``/``hook_k_input``/
369 ``hook_v_input``, and ``hook_mlp_in`` all exist on the Bridge unconditionally
370 but only fire when their owning flag (``cfg.use_attn_result`` /
371 ``cfg.use_split_qkv_input`` / ``cfg.use_hook_mlp_in``) is on, so an edge sweep
372 must enable all three before caching or the writer- and reader-side hook
373 points it needs never populate.
375 ``use_split_qkv_input`` is mutually exclusive with ``use_attn_in``, so a
376 caller who arrives with ``use_attn_in`` on would otherwise trip the
377 exclusivity error. This turns ``use_attn_in`` off before enabling the split
378 input, and restores it after ``use_split_qkv_input`` has been turned back off.
380 Invariant: the caller's flag state (``use_attn_result``,
381 ``use_split_qkv_input``, ``use_hook_mlp_in``, ``use_attn_in``) is unchanged on
382 return, including when the body raises. Restoration runs in a ``finally``
383 block so a raise mid-sweep -- or in the caller's own later code -- cannot
384 leave the per-head tensors materialized on the model.
386 Memory caveat: enabling ``use_attn_result``/``use_split_qkv_input`` makes
387 every cached per-head tensor ``[batch, seq, n_heads, d_model]`` instead of
388 the summed ``[batch, seq, d_model]`` residual. That is fine on a model the
389 size of gpt2-small; it does not scale to models with many heads or layers.
391 Args:
392 model: A ``TransformerBridge`` (or compatible) exposing ``cfg`` and
393 ``set_use_attn_result``/``set_use_split_qkv_input``/
394 ``set_use_hook_mlp_in``/``set_use_attn_in``.
395 """
396 cfg = model.cfg
397 saved_attn_result = cfg.use_attn_result
398 saved_split_qkv_input = cfg.use_split_qkv_input
399 saved_hook_mlp_in = cfg.use_hook_mlp_in
400 saved_attn_in = bool(getattr(cfg, "use_attn_in", False))
401 try:
402 # use_attn_in and use_split_qkv_input are mutually exclusive; clear
403 # use_attn_in first so enabling the split input cannot raise.
404 if saved_attn_in:
405 model.set_use_attn_in(False)
406 if not cfg.use_attn_result:
407 model.set_use_attn_result(True)
408 if not cfg.use_split_qkv_input:
409 model.set_use_split_qkv_input(True)
410 if not cfg.use_hook_mlp_in:
411 model.set_use_hook_mlp_in(True)
412 yield
413 finally:
414 model.set_use_attn_result(saved_attn_result)
415 model.set_use_split_qkv_input(saved_split_qkv_input)
416 model.set_use_hook_mlp_in(saved_hook_mlp_in)
417 # Re-enabling use_attn_in requires use_split_qkv_input already off; the
418 # line above restored it, and a caller with use_attn_in on cannot also
419 # have had use_split_qkv_input on, so this cannot trip the exclusivity.
420 if saved_attn_in:
421 model.set_use_attn_in(True)
424def _check_required_hooks(
425 cache: GradientCache, required_names: Sequence[str], graph: str, hint: str
426) -> None:
427 """Raise if any of ``required_names`` is absent from ``cache.activations``.
429 Shared by every granularity's graph-construction step, so a cache built
430 with too narrow a ``names_filter`` -- or one produced while a required
431 Bridge flag was off -- fails loudly instead of silently producing a
432 truncated graph.
433 """
434 missing = [name for name in required_names if name not in cache.activations]
435 if missing:
436 raise ValueError(
437 f"{graph} requires hook points missing from the cache: "
438 + ", ".join(missing)
439 + f". {hint}"
440 )
443def enumerate_nodes(model: Any, cache: GradientCache) -> list[Node]:
444 """Enumerate the full node-granularity graph from the Bridge hook graph.
446 The graph is *explicit*: for ``n_layers`` layers it always contains the embed
447 write, every attention head's output, and every layer's MLP output, at every
448 sequence position. Sequence length and head count are read from the cached
449 tensor shapes; ``n_layers`` from ``model.cfg``.
451 Args:
452 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``.
453 cache: A :class:`GradientCache` holding at least the required hook points.
455 Returns:
456 The node list, ordered embed-then-layerwise for deterministic ranking.
458 Raises:
459 ValueError: if any required hook point is absent from ``cache`` — the
460 graph is never silently truncated.
461 """
462 n_layers = int(model.cfg.n_layers)
463 _check_required_hooks(
464 cache,
465 _required_hook_names(n_layers),
466 "node graph",
467 "Cache with a names_filter that keeps hook_embed, blocks.*.attn.hook_z, "
468 "and blocks.*.hook_mlp_out.",
469 )
471 seq_len = cache.activations["hook_embed"].shape[1]
473 nodes: list[Node] = [Node(kind="embed", position=position) for position in range(seq_len)]
474 for layer in range(n_layers):
475 n_heads = cache.activations[f"blocks.{layer}.attn.hook_z"].shape[2]
476 for position in range(seq_len):
477 for head in range(n_heads):
478 nodes.append(Node(kind="attn_head_out", layer=layer, head=head, position=position))
479 nodes.append(Node(kind="mlp_out", layer=layer, position=position))
480 return nodes
483def _assert_edges_unique(edges: Sequence[tuple[Node, Node]]) -> None:
484 """Raise if any writer -> reader pair appears more than once in ``edges``.
486 A writer feeding two distinct readers (a head's output feeding both the
487 next layer's attention input and this layer's MLP input, say) is two
488 edges; this guards the enumeration against a construction bug that
489 collapses or duplicates a single ``(writer, reader)`` pair instead.
490 """
491 seen: set[tuple[Node, Node]] = set()
492 for edge in edges:
493 if edge in seen:
494 raise ValueError(f"edge {edge} enumerated more than once")
495 seen.add(edge)
498def _required_edge_reader_hook_names(n_layers: int) -> list[str]:
499 """The reader hook points edge enumeration additionally requires.
501 ``_required_hook_names(..., granularity="edge")`` covers the per-head
502 attention hooks a writer/reader pair into or out of a head needs. Edge
503 enumeration reads two reader points those miss:
505 - each layer's MLP entry, ``attn.hook_mlp_in``'s layer-level sibling
506 ``hook_mlp_in``, gated on ``cfg.use_hook_mlp_in`` the same way the per-head
507 hooks are gated on their own flags, and
508 - the terminal ``blocks.{n_layers-1}.hook_resid_post``, where the logits
509 reader takes its gradient. This final residual hook fires unconditionally,
510 so no Bridge flag gates it.
511 """
512 names = [f"blocks.{layer}.hook_mlp_in" for layer in range(n_layers)]
513 names.append(f"blocks.{n_layers - 1}.hook_resid_post")
514 return names
517def _edge_hook_names(n_layers: int) -> list[str]:
518 """Every hook point an edge-granularity sweep must cache.
520 The per-head attention hooks from ``_required_hook_names(..., granularity="edge")``
521 plus each layer's MLP-entry reader hook and the terminal logits reader hook.
522 """
523 return _required_hook_names(n_layers, granularity="edge") + _required_edge_reader_hook_names(
524 n_layers
525 )
528# Pre-norm blocks measured at most 2.7e-3 (gpt2, pythia-70m and Qwen2.5-0.5B in bfloat16, CPU
529# and MPS; float32 stays below 1e-7); post-norm blocks measured 0.05 to 19.
530_ATTN_WRITE_RTOL = 2.5e-2
533def _attn_write_hook_names(n_layers: int) -> list[str]:
534 """Each layer's attention residual write, cached only to check it against the head sum."""
535 return [f"blocks.{layer}.hook_attn_out" for layer in range(n_layers)]
538_EDGE_UNSUPPORTED = (
539 "edge granularity and faithfulness are unsupported here. Node granularity "
540 '(EdgeAttributionConfig(granularity="node")) is available.'
541)
544def _non_additive_attention_error(detail: str) -> ValueError:
545 return ValueError(
546 "per-head attention outputs are normalized or scaled before the residual add on this "
547 f"architecture ({detail}), so attn.hook_result is not a head's residual contribution: "
548 f"{_EDGE_UNSUPPORTED}"
549 )
552def _attention_write_obstruction(
553 model: Any, layers: Optional[Collection[int]] = None
554) -> Optional[str]:
555 """Describe the first block declaring a norm, scale or parallel mixer after its head sum."""
556 for layer, block in enumerate(getattr(model, "blocks", ())):
557 if layers is not None and layer not in layers:
558 continue
559 submodules = getattr(block, "submodules", None)
560 if not isinstance(submodules, dict) or "attn" not in submodules:
561 continue
562 scale = getattr(block, "residual_contribution_scale", 1.0)
563 if scale != 1.0:
564 return f"blocks.{layer} scales its attention output by {scale}"
565 # Falcon-H1 scales attention and adds a parallel Mamba output inline, after
566 # attn.hook_out, so the cached head-sum check cannot see either.
567 mixers = [name for name in VARIANT_SUBMODULE_NAMES if name != "attn" and name in submodules]
568 if mixers:
569 return f"blocks.{layer} shares that add with a parallel {mixers[0]} mixer"
570 aliases = getattr(block, "hook_aliases", None)
571 target = aliases.get("hook_attn_out") if isinstance(aliases, dict) else None
572 candidates = [target] if isinstance(target, str) else list(target or [])
573 # A fallback chain resolves to its first target whose owner the block declares.
574 for candidate in candidates:
575 owner = candidate.split(".", 1)[0]
576 if owner not in submodules: 576 ↛ 577line 576 didn't jump to line 577 because the condition on line 576 was never true
577 continue
578 if owner != "attn":
579 return f"blocks.{layer}.hook_attn_out reads {candidate}, after the head sum"
580 break
581 return None
584def _check_attention_write_structure(model: Any) -> None:
585 """Refuse blocks that declare a norm, scale or parallel mixer between head sum and residual."""
586 detail = _attention_write_obstruction(model)
587 if detail is not None:
588 raise _non_additive_attention_error(detail)
591def _check_kv_heads_match_query_heads(model: Any) -> None:
592 """Refuse GQA/MQA attention layers, whose K/V reader inputs are shared across query heads."""
593 cfg = model.cfg
594 blocks = list(getattr(model, "blocks", ()))
595 scalar_kv = getattr(cfg, "n_key_value_heads", None)
596 per_layer_kv = getattr(cfg, "per_layer_num_key_value_heads", None) or []
597 for layer in range(cfg.n_layers):
598 submodules = getattr(blocks[layer], "submodules", None) if layer < len(blocks) else None
599 # A hybrid's SSM or linear-attention block inherits the scalar count but has no K/V heads.
600 if isinstance(submodules, dict) and "attn" not in submodules:
601 continue
602 kv = per_layer_kv[layer] if layer < len(per_layer_kv) else None
603 if kv is None:
604 kv = scalar_kv
605 if kv is not None and kv != cfg.n_heads:
606 raise ValueError(
607 f"blocks.{layer} has n_key_value_heads={kv} for n_heads={cfg.n_heads} (GQA/MQA), "
608 f"so its K/V inputs are not per query head: {_EDGE_UNSUPPORTED}"
609 )
612def _check_attention_write_is_head_sum(cache: GradientCache, n_layers: int) -> None:
613 """Refuse a layer whose ``hook_attn_out`` is not its head-result sum plus a constant."""
614 for layer in range(n_layers):
615 per_head = cache.activations.get(f"blocks.{layer}.attn.hook_result")
616 written = cache.activations.get(f"blocks.{layer}.hook_attn_out")
617 if per_head is None or written is None:
618 continue
619 work_dtype = torch.promote_types(written.dtype, torch.float32)
620 written = written.to(work_dtype)
621 head_sum = per_head.sum(dim=2, dtype=work_dtype)
622 if head_sum.shape != written.shape: 622 ↛ 623line 622 didn't jump to line 623 because the condition on line 622 was never true
623 raise _non_additive_attention_error(
624 f"blocks.{layer}.hook_attn_out has shape {tuple(written.shape)}, the summed "
625 f"attn.hook_result {tuple(head_sum.shape)}"
626 )
627 # The output bias belongs to no head and cancels in every clean-minus-corrupt delta,
628 # so only the position-varying part of the remainder breaks additivity.
629 remainder = written - head_sum
630 remainder = remainder - remainder.mean(dim=(0, 1), keepdim=True)
631 scale = written.norm().clamp_min(torch.finfo(work_dtype).tiny)
632 relative = float(remainder.norm() / scale)
633 if relative > _ATTN_WRITE_RTOL:
634 raise _non_additive_attention_error(
635 f"blocks.{layer}.hook_attn_out differs from the summed attn.hook_result by "
636 f"{relative:.2g} of its norm"
637 )
640def _writer_hook_names(n_layers: int) -> list[str]:
641 """Hook points holding each writer's residual contribution: embed, head result, MLP out."""
642 names = ["hook_embed"]
643 for layer in range(n_layers):
644 names.append(f"blocks.{layer}.attn.hook_result")
645 names.append(f"blocks.{layer}.hook_mlp_out")
646 return names
649def _reader_hook_names(n_layers: int) -> list[str]:
650 """Hook points holding each reader's input; the logits reader's is listed separately."""
651 names: list[str] = []
652 for layer in range(n_layers):
653 names.append(f"blocks.{layer}.attn.hook_q_input")
654 names.append(f"blocks.{layer}.attn.hook_k_input")
655 names.append(f"blocks.{layer}.attn.hook_v_input")
656 names.append(f"blocks.{layer}.hook_mlp_in")
657 return names
660def enumerate_edges(model: Any, cache: GradientCache) -> list[tuple[Node, Node]]:
661 """Enumerate every writer -> reader edge in the residual-stream graph.
663 At a fixed sequence position, the residual stream is a running sum: a
664 reader (a head's split Q/K/V input, a layer's MLP entry, or the terminal
665 logits readout) is fed by every writer (the embed write, every attention
666 head's output, every layer's MLP output) that precedes it. Building the
667 graph position-by-position tracks which writers are "available" so far and
668 connects each new reader to all of them, then adds that layer's writers to
669 the available set before moving on -- so a writer never edges to a reader
670 upstream of it, and a writer feeding both a direct edge and a through-MLP
671 edge produces two distinct ``(u, v)`` pairs rather than one summed together.
673 A terminal ``logits`` reader (read at the final ``hook_resid_post``) closes
674 the graph: after the per-layer loop every remaining writer -- including the
675 final layer's MLP output, which no per-layer reader sees -- edges to it. That
676 edge carries the writer's direct skip-connection contribution to the metric,
677 so a writer's aggregate over its outgoing edges equals its direct node score.
679 On ``cfg.parallel_attn_mlp`` models (Pythia, GPT-J, Falcon, Phi) the MLP
680 reads the layer input, not the post-attention residual, so a layer's own
681 heads are not writers into that layer's MLP; those same-layer head->mlp_in
682 edges are dropped while the heads still feed later readers and the logits
683 reader.
685 Args:
686 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``
687 and, optionally, ``cfg.parallel_attn_mlp``.
688 cache: A :class:`GradientCache` holding at least the required hook points
689 for edge granularity.
691 Returns:
692 The edge list as ``(writer, reader)`` node pairs; no pair repeats.
694 Raises:
695 ValueError: if any required hook point is absent from ``cache`` — the
696 graph is never silently truncated.
697 """
698 n_layers = int(model.cfg.n_layers)
699 _check_required_hooks(
700 cache,
701 _edge_hook_names(n_layers),
702 "edge graph",
703 "Cache with a names_filter that keeps the edge-granularity hook set: "
704 "hook_embed, blocks.*.attn.hook_z, blocks.*.hook_mlp_out, "
705 "blocks.*.attn.hook_result, blocks.*.attn.hook_q_input, "
706 "blocks.*.attn.hook_k_input, blocks.*.attn.hook_v_input, "
707 "blocks.*.hook_mlp_in, and blocks.{n_layers-1}.hook_resid_post.",
708 )
710 parallel_attn_mlp = bool(getattr(model.cfg, "parallel_attn_mlp", False))
712 seq_len = cache.activations["hook_embed"].shape[1]
713 edges: list[tuple[Node, Node]] = []
715 for position in range(seq_len):
716 available: list[Node] = [Node(kind="embed", position=position)]
717 for layer in range(n_layers):
718 n_heads = cache.activations[f"blocks.{layer}.attn.hook_z"].shape[2]
720 attn_reader_kinds: tuple[NodeKind, NodeKind, NodeKind] = (
721 "q_input",
722 "k_input",
723 "v_input",
724 )
725 attn_readers = [
726 Node(kind=kind, layer=layer, head=head, position=position)
727 for kind in attn_reader_kinds
728 for head in range(n_heads)
729 ]
730 for reader in attn_readers:
731 edges.extend((writer, reader) for writer in available)
733 layer_heads = [
734 Node(kind="attn_head_out", layer=layer, head=head, position=position)
735 for head in range(n_heads)
736 ]
738 mlp_reader = Node(kind="mlp_in", layer=layer, position=position)
739 if parallel_attn_mlp:
740 # The MLP reads the layer input, not the post-attention residual,
741 # so this layer's heads are not writers into its MLP. They still
742 # become available to later readers and the logits reader.
743 edges.extend((writer, mlp_reader) for writer in available)
744 available = available + layer_heads
745 else:
746 available = available + layer_heads
747 edges.extend((writer, mlp_reader) for writer in available)
749 available = available + [Node(kind="mlp_out", layer=layer, position=position)]
751 logits_reader = Node(kind="logits", layer=n_layers - 1, position=position)
752 edges.extend((writer, logits_reader) for writer in available)
754 _assert_edges_unique(edges)
755 return edges
758def _as_predicate(names_filter: NamesFilter) -> Callable[[str], bool]:
759 if names_filter is None: 759 ↛ 760line 759 didn't jump to line 760 because the condition on line 759 was never true
760 return lambda name: True
761 if isinstance(names_filter, str): 761 ↛ 762line 761 didn't jump to line 762 because the condition on line 761 was never true
762 target = names_filter
763 return lambda name: name == target
764 if callable(names_filter): 764 ↛ 765line 764 didn't jump to line 765 because the condition on line 764 was never true
765 return names_filter
766 wanted = set(names_filter)
767 return lambda name: name in wanted
770def cache_activation_and_gradient(
771 model: Any,
772 tokens: torch.Tensor,
773 metric_fn: MetricFn,
774 names_filter: NamesFilter = None,
775 compute_gradient: bool = True,
776) -> GradientCache:
777 """Run one forward and capture activations plus (optionally) their metric-gradients.
779 Registers a forward hook and a backward hook at each cached point, runs one
780 grad-enabled forward, then drives a single backward with
781 :func:`torch.autograd.grad` to fire the backward hooks.
782 ``run_with_cache(..., incl_bwd=True)`` only backpropagates the model's own
783 scalar output, so a custom (non-scalar-output) metric such as a logit-diff
784 needs the backward driven here.
786 Gradients come from the backward hooks, not from ``.grad`` on the cached
787 tensors. ``TransformerBridge`` reshapes the tensor handed to a forward hook at
788 a converted point (``attn.hook_z``, ``hook_q/k/v``, ``hook_attn_out``), so that
789 tensor is a view the model's forward never consumes: ``retain_grad()`` on it is
790 inert and its ``.grad`` stays ``None``. A backward hook goes through the same
791 conversion and delivers the real gradient in canonical shape. Driving the
792 backward with :func:`torch.autograd.grad` instead of ``metric.backward()``
793 keeps it off every parameter's ``.grad`` buffer — no caller grads clobbered, no
794 model-sized buffer allocated.
796 With ``compute_gradient=False`` the backward is skipped entirely: only forward
797 hooks are registered, no backward is driven, and every cached point's gradient
798 is ``None``. Callers that read activations only (the clean pass of
799 :func:`attribution_patch`, which pairs these with the *corrupt* run's gradients)
800 use this to run a plain forward instead of a needless forward + backward.
802 Args:
803 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``,
804 ``hook_dict``, and the ``hooks()`` context manager.
805 tokens: Input token ids for a single forward pass.
806 metric_fn: Maps the model logits to a scalar to differentiate.
807 names_filter: Restricts which hook points are cached (and have gradients
808 retained). ``None`` (the default) caches the node-granularity hook set —
809 ``hook_embed`` plus each layer's ``attn.hook_z`` and ``hook_mlp_out``;
810 pass an explicit filter to cache any hook point outside that set. On a
811 real Bridge, ``None`` cannot mean "every hook point": the gated points
812 (``hook_mlp_in``, ``attn.hook_result``, the split-QKV inputs) that
813 ``hook_dict`` exposes raise in ``add_hook`` unless their ``set_use_*``
814 flag is on.
815 compute_gradient: When ``True`` (default) capture gradients via backward
816 hooks. When ``False`` run an activation-only forward and leave every
817 gradient ``None``.
819 Returns:
820 A :class:`GradientCache` with per-hook activations, and gradients when
821 ``compute_gradient`` is ``True`` (all ``None`` otherwise).
822 """
823 if compute_gradient and not torch.is_grad_enabled():
824 raise ValueError(
825 "cache_activation_and_gradient needs autograd, but gradient tracking "
826 "is off (torch.no_grad(), set_grad_enabled(False), or inference mode)."
827 )
829 if names_filter is None:
830 # Default to the node-granularity hook set rather than every hook point: a
831 # real Bridge's hook_dict holds gated points (hook_mlp_in, attn.hook_result,
832 # the split-QKV inputs) that add_hook rejects unless their set_use_* flag is
833 # on, so caching "everything" raises. Pass an explicit filter to reach them.
834 names_filter = _required_hook_names(int(model.cfg.n_layers))
835 predicate = _as_predicate(names_filter)
836 names = [name for name in model.hook_dict if predicate(name)]
837 if not names: 837 ↛ 838line 837 didn't jump to line 838 because the condition on line 837 was never true
838 raise ValueError("names_filter matched no hook points")
840 live: dict[str, torch.Tensor] = {}
841 activations: dict[str, torch.Tensor] = {}
842 gradients: dict[str, Optional[torch.Tensor]] = {}
844 def make_fwd_hook(name: str) -> Callable[..., None]:
845 def hook(tensor: torch.Tensor, *, hook: Any) -> None:
846 del hook
847 if isinstance(tensor, torch.Tensor): 847 ↛ 850line 847 didn't jump to line 850 because the condition on line 847 was always true
848 live[name] = tensor
849 activations[name] = tensor.detach().clone()
850 return None
852 return hook
854 def make_bwd_hook(name: str) -> Callable[..., None]:
855 def hook(grad: torch.Tensor, *, hook: Any) -> None:
856 del hook
857 if isinstance(grad, torch.Tensor): 857 ↛ 859line 857 didn't jump to line 859 because the condition on line 857 was always true
858 gradients[name] = grad.detach().clone()
859 return None
861 return hook
863 fwd_hooks = [(name, make_fwd_hook(name)) for name in names]
865 if not compute_gradient:
866 with model.hooks(fwd_hooks=fwd_hooks):
867 logits = model(tokens)
868 metric = metric_fn(logits)
869 if metric.dim() != 0: 869 ↛ 870line 869 didn't jump to line 870 because the condition on line 869 was never true
870 raise ValueError(
871 f"metric_fn must return a scalar tensor, got shape {tuple(metric.shape)}"
872 )
873 return GradientCache(
874 activations=activations,
875 gradients={name: None for name in activations},
876 metric=metric.detach(),
877 )
879 bwd_hooks = [(name, make_bwd_hook(name)) for name in names]
881 with model.hooks(fwd_hooks=fwd_hooks, bwd_hooks=bwd_hooks):
882 logits = model(tokens)
883 metric = metric_fn(logits)
884 if metric.dim() != 0: 884 ↛ 885line 884 didn't jump to line 885 because the condition on line 884 was never true
885 raise ValueError(
886 f"metric_fn must return a scalar tensor, got shape {tuple(metric.shape)}"
887 )
888 captured_names = [name for name in names if name in live]
889 # torch.autograd.grad only *drives* the backward; the gradients we keep are
890 # the converted-shape ones the backward hooks write into `gradients`. Unlike
891 # metric.backward() it touches no parameter `.grad` buffer.
892 returned = (
893 torch.autograd.grad(
894 metric,
895 inputs=[live[name] for name in captured_names],
896 allow_unused=True,
897 retain_graph=False,
898 )
899 if captured_names
900 else ()
901 )
903 # The most-upstream cached point's own backward hook never fires — nothing
904 # requested is upstream of it, so the backward stops at it — but its cached
905 # tensor is on-path and unconverted, so torch.autograd.grad returns that
906 # gradient directly. Backfill any point the hooks missed from that return.
907 for name, grad in zip(captured_names, returned):
908 if gradients.get(name) is None:
909 gradients[name] = None if grad is None else grad.detach().clone()
911 return GradientCache(activations=activations, gradients=gradients, metric=metric.detach())
914def capture_writer_outputs(
915 model: Any,
916 tokens: torch.Tensor,
917 names_filter: NamesFilter = None,
918) -> dict[str, torch.Tensor]:
919 """Capture writer contributions in one forward-only pass, with no gradient retained.
921 Memory caveat: ``attn.hook_result`` is a per-head ``[batch, seq, n_heads, d_model]``
922 tensor, so this does not scale to models with many heads or layers.
924 Args:
925 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``,
926 ``hook_dict``, and the ``hooks()`` context manager.
927 tokens: Input token ids for a single forward pass.
928 names_filter: Restricts which hook points are captured. ``None`` (the
929 default) captures ``hook_embed``, ``attn.hook_result`` and
930 ``hook_mlp_out``. Gated points (``attn.hook_result``, the split-QKV
931 inputs, ``hook_mlp_in``) raise in ``add_hook`` unless their
932 ``set_use_*`` flag is on; see :func:`_edge_hook_flags`.
934 Returns:
935 Detached activation tensors keyed by hook name, for every hook point
936 matching ``names_filter``.
938 Raises:
939 ValueError: if ``names_filter`` matches no hook point.
940 """
941 if names_filter is None:
942 names_filter = _writer_hook_names(int(model.cfg.n_layers))
943 predicate = _as_predicate(names_filter)
944 names = [name for name in model.hook_dict if predicate(name)]
945 if not names:
946 raise ValueError("names_filter matched no hook points")
948 captured: dict[str, torch.Tensor] = {}
950 def make_fwd_hook(name: str) -> Callable[..., None]:
951 def hook(tensor: torch.Tensor, *, hook: Any) -> None:
952 del hook
953 if isinstance(tensor, torch.Tensor): 953 ↛ 955line 953 didn't jump to line 955 because the condition on line 953 was always true
954 captured[name] = tensor.detach().clone()
955 return None
957 return hook
959 fwd_hooks = [(name, make_fwd_hook(name)) for name in names]
960 with torch.no_grad(), model.hooks(fwd_hooks=fwd_hooks):
961 model(tokens)
963 return captured
966def _node_effects(
967 clean_cache: GradientCache,
968 corrupt_cache: GradientCache,
969 nodes: Sequence[Node],
970) -> dict[Node, float]:
971 """Score every node with the first-order attribution ``(a_clean - a_corrupt) . g``.
973 The gradient ``g`` is taken from the corrupt cache (denoising convention). For
974 each node the feature dimension (``d_model``, or ``d_head`` for an attention
975 head) is contracted at the node's position/head, giving one signed scalar.
976 """
977 scores: dict[Node, float] = {}
978 for node in nodes:
979 name = node.hook_name
980 grad = corrupt_cache.gradients.get(name)
981 if grad is None: 981 ↛ 982line 981 didn't jump to line 982 because the condition on line 981 was never true
982 raise ValueError(
983 f"node {node} reads {name!r}, but the corrupt cache holds no gradient "
984 "there; cache with a names_filter that retains this hook point."
985 )
986 delta = clean_cache.activations[name] - corrupt_cache.activations[name]
987 contribution = delta * grad
988 if node.kind == "attn_head_out":
989 value = contribution[0, node.position, node.head].sum()
990 else:
991 value = contribution[0, node.position].sum()
992 scores[node] = float(value)
993 return scores
996def _writer_hook_name(node: Node) -> str:
997 """The residual-stream hook point holding a writer node's own contribution.
999 Distinct from ``Node.hook_name``: an ``attn_head_out`` node's ``hook_name``
1000 resolves to ``attn.hook_z``, the pre-``hook_result`` value node granularity
1001 scores. An edge's writer contribution must instead be measured in the same
1002 ``d_model`` space a reader's gradient lives in, which is ``attn.hook_result``
1003 -- the per-head decomposition of the head's contribution after it is
1004 projected into the residual stream.
1005 """
1006 if node.kind == "embed":
1007 return "hook_embed"
1008 if node.kind == "attn_head_out":
1009 return f"blocks.{node.layer}.attn.hook_result"
1010 if node.kind == "mlp_out": 1010 ↛ 1012line 1010 didn't jump to line 1012 because the condition on line 1010 was always true
1011 return f"blocks.{node.layer}.hook_mlp_out"
1012 raise ValueError(f"{node.kind} is a reader kind and has no writer contribution")
1015def _edge_effects(
1016 clean_cache: GradientCache,
1017 corrupt_cache: GradientCache,
1018 edges: Sequence[tuple[Node, Node]],
1019) -> dict[tuple[Node, Node], float]:
1020 """Score every edge with ``(a_clean[u] - a_corrupt[u]) . d(metric)/d(input of v)``.
1022 Mirrors ``_node_effects``: the delta is the writer's own residual
1023 contribution (clean minus corrupt cache), dotted with the reader's
1024 corrupt-run gradient -- the same denoising convention ``_node_effects``
1025 uses. Unlike a node score, the delta and the gradient are read from two
1026 different hook points (the writer's and the reader's), since an edge
1027 measures how much of one component's output reaches another component's
1028 input. A ``logits`` reader takes its gradient at the final
1029 ``hook_resid_post`` and contracts over ``d_model`` at its position, the same
1030 shape path as an ``mlp_in`` reader.
1032 Two structural facts keep this off a per-edge recompute: a writer hook's
1033 clean-minus-corrupt delta is the same tensor for every edge that writer
1034 feeds, and every edge into one reader shares that reader's gradient vector.
1035 Each writer delta is therefore computed once, and each reader's incoming
1036 edges are scored with a single batched matrix-vector product over the
1037 stacked writer deltas -- reducing to the same per-edge scalars a
1038 ``float((delta_vec * grad_vec).sum())`` per edge would, reader gradient
1039 reused across all of that reader's edges.
1040 """
1041 scores: dict[tuple[Node, Node], float] = {}
1042 writer_deltas: dict[str, torch.Tensor] = {}
1044 def writer_delta_vec(writer: Node) -> torch.Tensor:
1045 name = _writer_hook_name(writer)
1046 delta = writer_deltas.get(name)
1047 if delta is None:
1048 delta = clean_cache.activations[name] - corrupt_cache.activations[name]
1049 writer_deltas[name] = delta
1050 if writer.kind == "attn_head_out":
1051 return delta[0, writer.position, writer.head]
1052 return delta[0, writer.position]
1054 writers_by_reader: dict[Node, list[Node]] = {}
1055 reader_order: list[Node] = []
1056 for writer, reader in edges:
1057 bucket = writers_by_reader.get(reader)
1058 if bucket is None:
1059 writers_by_reader[reader] = bucket = []
1060 reader_order.append(reader)
1061 bucket.append(writer)
1063 for reader in reader_order:
1064 reader_name = reader.hook_name
1065 grad = corrupt_cache.gradients.get(reader_name)
1066 writers = writers_by_reader[reader]
1067 if grad is None: 1067 ↛ 1068line 1067 didn't jump to line 1068 because the condition on line 1067 was never true
1068 raise ValueError(
1069 f"edge {(writers[0], reader)} reads its gradient at {reader_name!r}, but "
1070 "the corrupt cache holds none there; cache with a names_filter that "
1071 "retains this hook point."
1072 )
1073 if reader.kind in ("q_input", "k_input", "v_input"):
1074 grad_vec = grad[0, reader.position, reader.head]
1075 else:
1076 grad_vec = grad[0, reader.position]
1077 delta_matrix = torch.stack([writer_delta_vec(writer) for writer in writers])
1078 edge_values = delta_matrix @ grad_vec
1079 for writer, value in zip(writers, edge_values.tolist()):
1080 scores[(writer, reader)] = value
1081 return scores
1084def _aggregate_edge_scores_to_writer_nodes(
1085 edge_scores: dict[tuple[Node, Node], float],
1086) -> dict[Node, float]:
1087 """Sum each writer's outgoing edge scores into that writer's aggregate node score.
1089 A writer's aggregate is the sum of its effects along every edge it feeds.
1090 enumerate_edges gives every writer an edge to the terminal logits reader, so
1091 the aggregate includes the writer's direct skip-connection contribution to
1092 the metric -- the part of its residual-stream write that no intermediate
1093 component reads, only the final readout does -- and therefore equals the
1094 quantity a node-granularity sweep measures directly at the writer's own hook
1095 point. Every writer has at least that one outgoing edge, including the final
1096 layer's MLP output, whose only reader is the logits terminal, so no writer is
1097 missing from the aggregate.
1098 """
1099 totals: dict[Node, float] = {}
1100 for (writer, _reader), score in edge_scores.items():
1101 totals[writer] = totals.get(writer, 0.0) + score
1102 return totals
1105def _writer_contribution_slice(writer: Node, tensor: torch.Tensor) -> torch.Tensor:
1106 """The writer's ``d_model`` contribution at its position, and head if it has one."""
1107 if writer.kind == "attn_head_out":
1108 return tensor[0, writer.position, writer.head]
1109 return tensor[0, writer.position]
1112def _excluded_writers_by_reader(
1113 edges: Sequence[tuple[Node, Node]],
1114 circuit: Collection[tuple[Node, Node]],
1115) -> dict[Node, list[Node]]:
1116 """Group the out-of-circuit writers feeding each reader node, so each gets one hook."""
1117 circuit_set = set(circuit)
1118 excluded: dict[Node, list[Node]] = {}
1119 for writer, reader in edges:
1120 if (writer, reader) not in circuit_set:
1121 excluded.setdefault(reader, []).append(writer)
1122 return excluded
1125def _make_writer_capture_hook(name: str, sink: dict[str, torch.Tensor]) -> Callable[..., None]:
1126 """Record a writer's live contribution for the reader hooks that follow it.
1128 Detached, not cloned: it is only read, and the reader hook clones its own input.
1129 """
1131 def hook(tensor: torch.Tensor, *, hook: Any) -> None:
1132 del hook
1133 if isinstance(tensor, torch.Tensor): 1133 ↛ 1135line 1133 didn't jump to line 1135 because the condition on line 1133 was always true
1134 sink[name] = tensor.detach()
1135 return None
1137 return hook
1140def _make_edge_ablation_hook(
1141 reader: Node,
1142 excluded_writers: Sequence[Node],
1143 live_writers: dict[str, torch.Tensor],
1144 replacement_writers: dict[str, torch.Tensor],
1145) -> Callable[..., torch.Tensor]:
1146 """Rebuild one reader's input with its out-of-circuit writers replaced.
1148 Exact, not approximate: the pre-norm residual is a plain sum of writer contributions.
1149 """
1150 if reader.kind in ("q_input", "k_input", "v_input"):
1151 # Node.__post_init__ guarantees a head for the per-head reader kinds.
1152 index: tuple[int, ...] = (0, reader.position, cast(int, reader.head))
1153 else:
1154 index = (0, reader.position)
1156 def hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
1157 del hook
1158 rebuilt = tensor.clone()
1159 for writer in excluded_writers:
1160 name = _writer_hook_name(writer)
1161 live = live_writers.get(name)
1162 replacement = replacement_writers.get(name)
1163 if live is None or replacement is None:
1164 raise ValueError(
1165 f"cannot ablate the edge from {writer} into {reader}: no "
1166 f"contribution captured at {name!r}. Enable the writer's hook "
1167 "family before running the ablation."
1168 )
1169 rebuilt[index] = (
1170 rebuilt[index]
1171 - _writer_contribution_slice(writer, live)
1172 + _writer_contribution_slice(writer, replacement)
1173 )
1174 return rebuilt
1176 return hook
1179def _ablate_edges(
1180 model: Any,
1181 tokens: torch.Tensor,
1182 edges: Sequence[tuple[Node, Node]],
1183 circuit: Collection[tuple[Node, Node]],
1184 replacement_writers: dict[str, torch.Tensor],
1185) -> torch.Tensor:
1186 """Return logits from one no-grad forward with every edge outside ``circuit`` ablated.
1188 The caller must have the writer and reader hook families enabled (see :func:`_edge_hook_flags`).
1189 """
1190 excluded_by_reader = _excluded_writers_by_reader(edges, circuit)
1192 live_writers: dict[str, torch.Tensor] = {}
1193 fwd_hooks: list[tuple[str, Callable[..., Any]]] = [
1194 (name, _make_writer_capture_hook(name, live_writers))
1195 for name in _writer_hook_names(int(model.cfg.n_layers))
1196 ]
1197 fwd_hooks.extend(
1198 (
1199 reader.hook_name,
1200 _make_edge_ablation_hook(reader, writers, live_writers, replacement_writers),
1201 )
1202 for reader, writers in excluded_by_reader.items()
1203 )
1205 with torch.no_grad(), model.hooks(fwd_hooks=fwd_hooks):
1206 logits = model(tokens)
1207 return logits
1210def _validate_pair(
1211 clean: torch.Tensor, corrupt: torch.Tensor, *, operation: str, method: str
1212) -> None:
1213 """Require 2D clean/corrupt token tensors holding equal numbers of equal-length pairs."""
1214 if clean.ndim != 2 or corrupt.ndim != 2:
1215 raise ValueError(
1216 f"{operation} expects 2D [batch, seq] token tensors, got clean "
1217 f"{tuple(clean.shape)} and corrupt {tuple(corrupt.shape)}"
1218 )
1219 if clean.shape[0] != corrupt.shape[0]:
1220 raise ValueError(
1221 "clean and corrupt must hold the same number of prompt pairs, got "
1222 f"{clean.shape[0]} and {corrupt.shape[0]}"
1223 )
1224 if clean.shape[1] != corrupt.shape[1]:
1225 raise ValueError(
1226 "each clean/corrupt pair must tokenize to the same length; got clean "
1227 f"length {clean.shape[1]} and corrupt length {corrupt.shape[1]}. "
1228 f"{method} aligns activations position-by-position."
1229 )
1232def attribution_patch(
1233 model: Any,
1234 clean: torch.Tensor,
1235 corrupt: torch.Tensor,
1236 metric_fn: MetricFn,
1237 config: EdgeAttributionConfig = EdgeAttributionConfig(),
1238) -> AttributionResult:
1239 """Estimate every component's causal effect on ``metric_fn`` in two forwards + one backward.
1241 For each clean/corrupt pair this runs a clean forward (for ``a_clean``) and a
1242 corrupt forward whose backward hooks capture ``g = d(metric)/d(a)`` (for
1243 ``a_corrupt`` and its gradient). At node granularity (``config.granularity ==
1244 "node"``) each node is scored with the first-order Taylor estimate
1245 ``effect(node) = (a_clean - a_corrupt) . g``. At edge granularity
1246 (``config.granularity == "edge"``) each writer -> reader edge is scored with
1247 ``effect(edge) = (a_clean[writer] - a_corrupt[writer]) . d(metric)/d(input of
1248 reader)``, and ``node_scores`` holds each writer's aggregate effect (the sum
1249 of its outgoing edge scores).
1251 Sign/direction convention (denoising form): gradients are taken on the *corrupt*
1252 run and the estimate points *toward* the clean activation, so a positive score
1253 means patching that node from corrupt toward clean moves the metric in the
1254 positive direction. An oracle-parity test maps this convention onto a pinned
1255 reference rather than assuming the two agree.
1257 Dataset averaging: ``clean``/``corrupt`` may hold a batch of prompt pairs. Each
1258 pair is scored independently (per-example forward/backward, so its own
1259 reconstruction identity holds) and per-node (or per-edge) scores are averaged
1260 across the batch before ranking.
1262 The model and every submodule must be in evaluation mode. Separate clean and
1263 corrupt forwards cannot produce meaningful activation differences if stochastic
1264 training layers such as dropout remain active.
1266 Args:
1267 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``,
1268 ``hook_dict``, and ``hooks()``.
1269 clean: Clean token ids, shape ``[batch, seq]``.
1270 corrupt: Corrupt token ids, shape ``[batch, seq]``, paired row-by-row with
1271 ``clean``.
1272 metric_fn: Maps single-example logits to a scalar to differentiate.
1273 config: Sweep configuration. Node and edge granularity are both
1274 supported with plain attribution (``ig_steps=1``); ``ig_steps>1``
1275 raises at construction.
1277 Returns:
1278 An :class:`AttributionResult` whose ``node_scores`` are averaged over the
1279 batch. For an edge-granularity sweep, ``edge_scores`` is populated too and
1280 ``node_scores`` is the per-writer aggregate of those edge scores. The
1281 graph's terminal logits reader gives every writer an edge carrying its
1282 direct skip-connection contribution to the metric, so this aggregate
1283 equals the quantity a node-granularity sweep on the same model returns
1284 (including for the final layer's MLP output, whose only outgoing edge is
1285 the one to the logits reader).
1287 Raises:
1288 ValueError: if ``clean``/``corrupt`` are not 2D, hold a different number of
1289 pairs, a pair tokenizes to different lengths (activations must align
1290 position-by-position), or the model or one of its submodules is in
1291 training mode. At edge granularity, also if per-head attention outputs
1292 are normalized or scaled before the residual add (post-norm,
1293 scaled-residual or parallel-mixer blocks), where ``attn.hook_result`` is not
1294 a head's residual contribution, or if an attention layer's K/V head count
1295 differs from ``n_heads`` (GQA/MQA).
1296 """
1297 _validate_pair(clean, corrupt, operation="attribution_patch", method="Attribution patching")
1298 require_eval_mode(model, operation="attribution_patch()")
1300 batch = int(clean.shape[0])
1301 n_layers = int(model.cfg.n_layers)
1303 if config.granularity == "edge":
1304 _check_attention_write_structure(model)
1305 _check_kv_heads_match_query_heads(model)
1306 hook_names = _edge_hook_names(n_layers)
1307 edge_totals: dict[tuple[Node, Node], float] = {}
1309 # Scope the per-head hook-flag mutation to the caching loop so the
1310 # caller's flag state is restored on both the normal and the error path.
1311 with _edge_hook_flags(model):
1312 for index in range(batch):
1313 clean_cache = cache_activation_and_gradient(
1314 model,
1315 clean[index : index + 1],
1316 metric_fn,
1317 names_filter=hook_names + _attn_write_hook_names(n_layers),
1318 compute_gradient=False,
1319 )
1320 _check_attention_write_is_head_sum(clean_cache, n_layers)
1321 corrupt_cache = cache_activation_and_gradient(
1322 model, corrupt[index : index + 1], metric_fn, names_filter=hook_names
1323 )
1324 edges = enumerate_edges(model, corrupt_cache)
1325 for edge, score in _edge_effects(clean_cache, corrupt_cache, edges).items():
1326 edge_totals[edge] = edge_totals.get(edge, 0.0) + score
1328 edge_scores = {edge: total / batch for edge, total in edge_totals.items()}
1329 node_scores = _aggregate_edge_scores_to_writer_nodes(edge_scores)
1330 return AttributionResult(node_scores=node_scores, edge_scores=edge_scores)
1332 node_hook_names = _required_hook_names(n_layers)
1333 totals: dict[Node, float] = {}
1335 for index in range(batch):
1336 clean_cache = cache_activation_and_gradient(
1337 model,
1338 clean[index : index + 1],
1339 metric_fn,
1340 names_filter=node_hook_names,
1341 compute_gradient=False,
1342 )
1343 corrupt_cache = cache_activation_and_gradient(
1344 model, corrupt[index : index + 1], metric_fn, names_filter=node_hook_names
1345 )
1346 nodes = enumerate_nodes(model, corrupt_cache)
1347 for node, score in _node_effects(clean_cache, corrupt_cache, nodes).items():
1348 totals[node] = totals.get(node, 0.0) + score
1350 node_scores = {node: total / batch for node, total in totals.items()}
1351 return AttributionResult(node_scores=node_scores)
1354def _normalize_circuit(
1355 circuit: Sequence[tuple[Node, Node]] | Sequence[tuple[Node, Node, float]],
1356) -> list[tuple[Node, Node]]:
1357 """Deduplicate circuit pairs in first-occurrence order, dropping ``top_edges`` scores."""
1358 normalized: list[tuple[Node, Node]] = []
1359 seen: set[tuple[Node, Node]] = set()
1360 for entry in circuit:
1361 if len(entry) == 2:
1362 writer, reader = entry
1363 elif len(entry) == 3: 1363 ↛ 1366line 1363 didn't jump to line 1366 because the condition on line 1363 was always true
1364 writer, reader, _score = entry
1365 else:
1366 raise ValueError(
1367 f"circuit entries must be (writer, reader) or (writer, reader, score), "
1368 f"got a {len(entry)}-tuple"
1369 )
1370 edge = (writer, reader)
1371 if edge not in seen:
1372 normalized.append(edge)
1373 seen.add(edge)
1374 return normalized
1377def _edge_class(edge: tuple[Node, Node]) -> EdgeClass:
1378 """The edge's class, keyed by what its reader consumes; Q and K share one (both feed scores)."""
1379 reader = edge[1]
1380 if reader.kind in ("q_input", "k_input"):
1381 return "into_qk"
1382 if reader.kind == "v_input":
1383 return "into_v"
1384 if reader.kind == "mlp_in":
1385 return "into_mlp"
1386 if reader.kind == "logits":
1387 return "into_logits"
1388 raise ValueError(f"{reader.kind} is a writer kind and has no edge class")
1391def _mean_writer_contributions(
1392 model: Any,
1393 mean_tokens: torch.Tensor,
1394) -> dict[str, torch.Tensor]:
1395 """Average calibration examples, preserving position, head, and singleton batch axes."""
1396 count = int(mean_tokens.shape[0])
1397 if count == 0: 1397 ↛ 1398line 1397 didn't jump to line 1398 because the condition on line 1397 was never true
1398 raise ValueError("mean_tokens must contain at least one calibration example")
1399 totals: dict[str, torch.Tensor] = {}
1400 for index in range(count):
1401 captured = capture_writer_outputs(model, mean_tokens[index : index + 1])
1402 for name, tensor in captured.items():
1403 running = totals.get(name)
1404 totals[name] = tensor if running is None else running + tensor
1405 return {name: total / count for name, total in totals.items()}
1408def faithfulness(
1409 model: Any,
1410 clean: torch.Tensor,
1411 corrupt: torch.Tensor,
1412 metric_fn: MetricFn,
1413 circuit: Sequence[tuple[Node, Node]] | Sequence[tuple[Node, Node, float]],
1414 config: FaithfulnessConfig = FaithfulnessConfig(),
1415 *,
1416 mean_tokens: Optional[torch.Tensor] = None,
1417) -> FaithfulnessResult:
1418 """Measure how much of the clean-to-corrupt metric gap a circuit recovers.
1420 Ablates every edge *outside* ``circuit`` (forward-only, exact rather than linearized) and
1421 reports the metric the model then produces. Feed a ranked edge set straight in from
1422 ``attribution_patch(..., config=EdgeAttributionConfig(granularity="edge")).top_edges(k=...)``.
1424 Under mean ablation recovery is still normalized to the clean/corrupt gap, so an empty
1425 circuit need not have zero recovery.
1427 Args:
1428 model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``,
1429 ``hook_dict``, and ``hooks()``.
1430 clean: Clean token ids for one evaluated example, shape ``[1, seq]``.
1431 corrupt: Paired corrupt token ids, shape ``[1, seq]``.
1432 metric_fn: Maps single-example logits to a scalar.
1433 circuit: The edges to keep, as ``(writer, reader)`` pairs or the
1434 ``(writer, reader, score)`` triples ``top_edges`` returns. Every
1435 other edge in the graph is ablated. Duplicate entries count once.
1436 config: Ablation configuration. Defaults to replacing an ablated writer
1437 with the corrupt run's own contribution.
1438 mean_tokens: Required only for ``ablation="mean"``. Held-out ``int32``/``int64``
1439 calibration token ids, ``[n_calibration, seq]`` at the evaluated length and
1440 on the clean/corrupt device; writer activations are averaged per position
1441 and head. Rows equal to either endpoint are rejected; avoid semantic leakage too.
1443 Returns:
1444 A :class:`FaithfulnessResult`. The per-edge-class breakout costs one extra
1445 forward per retained class.
1447 Raises:
1448 ValueError: if ``clean``/``corrupt`` are not 2D, hold a different number
1449 of pairs, hold more than one pair, a pair tokenizes to different
1450 lengths, the model or a submodule is in training mode, a circuit edge
1451 is not in the graph, the clean and corrupt metrics are equal (so
1452 the recovered fraction is undefined), mean calibration is missing,
1453 empty, misaligned, on a different input device, has an unsupported
1454 dtype, contains an evaluated endpoint, or is supplied in corrupt mode;
1455 or per-head attention outputs are normalized or scaled before the
1456 residual add (post-norm, scaled-residual or parallel-mixer blocks), or an
1457 attention layer's K/V head count differs from ``n_heads`` (GQA/MQA).
1458 """
1459 _validate_pair(clean, corrupt, operation="faithfulness", method="Faithfulness")
1460 if clean.shape[0] != 1:
1461 raise ValueError(
1462 "faithfulness measures one clean/corrupt pair at a time, since the "
1463 f"metric reads a single example; got {clean.shape[0]} pairs. Loop over "
1464 "pairs and aggregate the recovered fractions yourself."
1465 )
1467 if config.ablation == "mean":
1468 if mean_tokens is None:
1469 raise ValueError("mean ablation requires independent calibration via mean_tokens")
1470 if mean_tokens.ndim != 2:
1471 raise ValueError("mean_tokens must be a 2D [n_calibration, seq] token tensor")
1472 if mean_tokens.shape[0] == 0:
1473 raise ValueError("mean_tokens must contain at least one calibration example")
1474 if mean_tokens.shape[1] != clean.shape[1]:
1475 raise ValueError("mean_tokens must have the same sequence length as the evaluated pair")
1476 if mean_tokens.dtype not in (torch.int32, torch.int64):
1477 raise ValueError(
1478 "mean_tokens must contain token ids with dtype torch.int32 or torch.int64"
1479 )
1480 if mean_tokens.device != clean.device or mean_tokens.device != corrupt.device:
1481 raise ValueError("mean_tokens must be on the same device as the clean/corrupt tokens")
1482 if (mean_tokens == clean).all(dim=1).any() or (mean_tokens == corrupt).all(dim=1).any():
1483 raise ValueError("mean_tokens must not contain either evaluated clean/corrupt endpoint")
1484 elif mean_tokens is not None:
1485 raise ValueError("mean_tokens is only accepted with mean ablation")
1487 require_eval_mode(model, operation="faithfulness()")
1488 _check_attention_write_structure(model)
1489 _check_kv_heads_match_query_heads(model)
1491 kept = _normalize_circuit(circuit)
1492 n_layers = int(model.cfg.n_layers)
1493 hook_names = _edge_hook_names(n_layers) + _attn_write_hook_names(n_layers)
1495 with torch.no_grad(), _edge_hook_flags(model):
1496 corrupt_cache = cache_activation_and_gradient(
1497 model, corrupt, metric_fn, names_filter=hook_names, compute_gradient=False
1498 )
1499 _check_attention_write_is_head_sum(corrupt_cache, n_layers)
1500 edges = enumerate_edges(model, corrupt_cache)
1501 graph = set(edges)
1502 unknown = [edge for edge in kept if edge not in graph]
1503 if unknown:
1504 raise ValueError(
1505 f"circuit holds {len(unknown)} edge(s) that are not in the graph, "
1506 f"starting with {unknown[0]}; a stale or mistyped circuit would "
1507 "otherwise ablate nothing."
1508 )
1510 if mean_tokens is not None:
1511 replacements = _mean_writer_contributions(model, mean_tokens)
1512 else:
1513 replacements = capture_writer_outputs(model, corrupt)
1515 with torch.no_grad():
1516 clean_metric = float(metric_fn(model(clean)))
1517 corrupt_metric = float(corrupt_cache.metric)
1518 gap = clean_metric - corrupt_metric
1519 if gap == 0.0:
1520 raise ValueError(
1521 "the clean and corrupt runs produce the same metric, so the recovered "
1522 "fraction is undefined; choose a pair the metric separates."
1523 )
1524 ablated_metric = float(metric_fn(_ablate_edges(model, clean, edges, kept, replacements)))
1526 recovered = (ablated_metric - corrupt_metric) / gap
1527 class_counts: dict[EdgeClass, int] = {}
1528 for edge in edges:
1529 edge_class = _edge_class(edge)
1530 class_counts[edge_class] = class_counts.get(edge_class, 0) + 1
1532 kept_classes = {_edge_class(edge) for edge in kept}
1533 class_recovered: dict[EdgeClass, float] = {}
1534 for edge_class in class_counts:
1535 if edge_class not in kept_classes:
1536 class_recovered[edge_class] = recovered
1537 continue
1538 # Removing a class must not restore edges excluded by the candidate.
1539 reduced = [edge for edge in kept if _edge_class(edge) != edge_class]
1540 class_metric = float(
1541 metric_fn(_ablate_edges(model, clean, edges, reduced, replacements))
1542 )
1543 class_recovered[edge_class] = (class_metric - corrupt_metric) / gap
1545 return FaithfulnessResult(
1546 recovered=recovered,
1547 full_metric=clean_metric,
1548 corrupt_metric=corrupt_metric,
1549 circuit_size=len(kept),
1550 total_edges=len(edges),
1551 edge_class_recovered=class_recovered,
1552 edge_class_counts=class_counts,
1553 )