transformer_lens.tools.analysis.attribution_patching module¶

Attribution patching — linearized activation patching on TransformerBridge.

Attribution patching estimates the causal effect of every model component on a task metric with a gradient-based linearization of activation patching: rather than one forward pass per intervention, it reads a single gradient cache. For a clean/corrupt prompt pair it runs a clean forward (for a_clean) and a corrupt forward whose backward hooks capture g = d(metric)/d(a) (for a_corrupt and its gradient), and scores each node with the first-order Taylor estimate effect(node) = (a_clean - a_corrupt) . g. Scores over a batch of clean/corrupt pairs are averaged before ranking.

Only the TransformerBridge API is targeted; TransformerLens v4 deprecates HookedTransformer.

Sign/direction convention (denoising form): the gradient is taken on the corrupt run and the estimate points toward the clean activation, so a positive score means patching that node from corrupt toward clean moves the metric in the positive direction. An oracle-parity test maps this convention onto a pinned reference rather than assuming the two agree.

Memory note: gradients are retained only for hook points passing names_filter. Retaining gradients at every hook point roughly doubles cache memory, so callers should filter to the hook families their analysis actually reads.

Scope: this build ships node and edge granularity with plain attribution (ig_steps=1). The integrated-gradient path (EAP-IG, ig_steps>1) and ablate-outside faithfulness are not implemented yet; their API is declared here, and ig_steps>1 raises NotImplementedError, so downstream code can pin against a stable surface now.

class transformer_lens.tools.analysis.attribution_patching.AttributionResult(node_scores: dict[~transformer_lens.tools.analysis.attribution_patching.Node, float], edge_scores: dict[tuple[~transformer_lens.tools.analysis.attribution_patching.Node, ~transformer_lens.tools.analysis.attribution_patching.Node], float] = <factory>)¶

Bases: object

Scored output of an attribution-patching sweep.

node_scores¶

Signed first-order effect estimate per node, (a_clean - a_corrupt) . d(metric)/d(a). A positive score means patching that node from corrupt toward clean moves the metric in the positive direction (the denoising convention pinned in the module docstring). For an edge-granularity sweep this is instead each writer’s aggregate over its own outgoing edge scores. Because the graph has a terminal logits reader that every writer feeds, that aggregate equals the writer’s direct node score (see attribution_patch()).

Type:: dict[transformer_lens.tools.analysis.attribution_patching.Node, float]

edge_scores¶

Per-edge effect estimate keyed by (source, destination). Populated for an edge-granularity sweep; empty for a node sweep.

Type:: dict[tuple[transformer_lens.tools.analysis.attribution_patching.Node, transformer_lens.tools.analysis.attribution_patching.Node], float]

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

node_scores: dict[Node, float]¶

top_edges(k: int = 10) → list[tuple[Node, Node, float]]¶

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

Ranking is by absolute score, matching top_nodes: a large negative edge effect is as causally important as a large positive one. Ties keep enumeration order (stable sort). Requesting more than the available edges returns all of them.

top_nodes(k: int = 10) → list[tuple[Node, float]]¶

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

Ranking is by absolute score: a node with a large negative effect is as causally important as one with a large positive effect, so magnitude — not signed value — orders the circuit. Ties keep enumeration order (stable sort). Requesting more than the available nodes returns all of them.

class transformer_lens.tools.analysis.attribution_patching.EdgeAttributionConfig(granularity: Literal['node', 'edge'] = 'node', ig_steps: int = 1)¶

Bases: object

Configuration for an attribution-patching sweep.

The two axes are deliberately orthogonal:

granularity selects what is scored: "node" scores each residual-stream write, "edge" scores each (source, destination) write->read pair.
ig_steps selects gradient fidelity. ig_steps=1 is plain attribution patching / EAP: a single first-order Taylor gradient taken at the corrupt point. ig_steps>1 is EAP-IG: the integrated gradient averaged over that many points along the corrupt->clean path, which corrects the gradient saturation that makes plain attribution unfaithful.

There is intentionally no method field. An earlier design had both a method enum ("attribution"/"EAP"/"EAP-IG") and ig_steps, which overlap: the method is fully determined by granularity and whether ig_steps exceeds 1. Collapsing them removes the invalid states (e.g. method="attribution", ig_steps=5).

This build implements node and edge granularity with plain attribution. ig_steps>1 is accepted by the type but raises NotImplementedError at construction, so downstream code can import and reference this API now while the integrated-gradient path is not implemented yet. Once EAP-IG lands, the default flips to ig_steps=5 (EAP-IG is the faithful default); until then the default is the only executable value, ig_steps=1.

granularity¶

"node" or "edge". Defaults to "node".

Type:: Literal[‘node’, ‘edge’]

ig_steps¶

Integrated-gradient path steps (>=1). Defaults to 1.

Type:: int

granularity: Literal['node', 'edge'] = 'node'¶

ig_steps: int = 1¶

class transformer_lens.tools.analysis.attribution_patching.GradientCache(activations: dict[str, Tensor], gradients: dict[str, Tensor | None], metric: Tensor)¶

Bases: object

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

activations¶

Detached activation tensor per cached hook name.

Type:: dict[str, torch.Tensor]

gradients¶

d(metric)/d(activation) per cached hook name.

Type:: dict[str, torch.Tensor | None]

metric¶

The scalar metric value at this run (detached).

Type:: torch.Tensor

activations: dict[str, Tensor]¶

gradients: dict[str, Tensor | None]¶

metric: Tensor¶

class transformer_lens.tools.analysis.attribution_patching.Node(kind: Literal['embed', 'attn_head_out', 'mlp_out', 'q_input', 'k_input', 'v_input', 'mlp_in', 'logits'], position: int, layer: int | None = None, head: int | None = None)¶

Bases: object

A node in the residual-stream computational graph.

Nodes are the typed, hashable keys the attribution sweep scores. Each node is identified by (kind, layer, position, head); kind selects the node family and constrains which of layer/head apply. Three kinds are writers – they contribute a value into the residual stream:

"embed": the token embedding write. layer and head are None.
"attn_head_out": one attention head’s output. layer and head set.
"mlp_out": one layer’s MLP output. layer set, head is None.

Five kinds are readers – they consume the residual stream as an edge’s destination (see enumerate_edges()):

"q_input" / "k_input" / "v_input": one attention head’s split Q/K/V input. layer and head set.
"mlp_in": one layer’s MLP entry. layer set, head is None.
"logits": the terminal readout of the final residual, read at blocks.{n_layers-1}.hook_resid_post. layer is that final layer and head is None. Every writer feeds this reader, so a writer’s aggregate over its outgoing edges equals its direct node score.

position is the sequence index the node is read at. The invariants above are enforced in __post_init__ so a malformed key raises rather than silently producing a wrong graph.

head: int | None = None¶

property hook_name: str¶

The cache hook point this node reads from.

Uses the standard TransformerBridge alias names (hook_embed, blocks.{l}.attn.hook_z, blocks.{l}.hook_mlp_out, blocks.{l}.attn.hook_q_input/hook_k_input/hook_v_input, blocks.{l}.hook_mlp_in, blocks.{l}.hook_resid_post); the per-head nodes (attn_head_out, q_input, k_input, v_input) slice head self.head out of the shared per-head tensor.

kind: Literal['embed', 'attn_head_out', 'mlp_out', 'q_input', 'k_input', 'v_input', 'mlp_in', 'logits']¶

layer: int | None = None¶

position: int¶

transformer_lens.tools.analysis.attribution_patching.attribution_patch(model: Any, clean: Tensor, corrupt: Tensor, metric_fn: Callable[[Tensor], Tensor], config: EdgeAttributionConfig = EdgeAttributionConfig(granularity='node', ig_steps=1)) → AttributionResult¶

Estimate every component’s causal effect on metric_fn in two forwards + one backward.

For each clean/corrupt pair this runs a clean forward (for a_clean) and a corrupt forward whose backward hooks capture g = d(metric)/d(a) (for a_corrupt and its gradient). At node granularity (config.granularity == "node") each node is scored with the first-order Taylor estimate effect(node) = (a_clean - a_corrupt) . g. At edge granularity (config.granularity == "edge") each writer -> reader edge is scored with effect(edge) = (a_clean[writer] - a_corrupt[writer]) . d(metric)/d(input of reader), and node_scores holds each writer’s aggregate effect (the sum of its outgoing edge scores).

Sign/direction convention (denoising form): gradients are taken on the corrupt run and the estimate points toward the clean activation, so a positive score means patching that node from corrupt toward clean moves the metric in the positive direction. An oracle-parity test maps this convention onto a pinned reference rather than assuming the two agree.

Dataset averaging: clean/corrupt may hold a batch of prompt pairs. Each pair is scored independently (per-example forward/backward, so its own reconstruction identity holds) and per-node (or per-edge) scores are averaged across the batch before ranking.

The model and every submodule must be in evaluation mode. Separate clean and corrupt forwards cannot produce meaningful activation differences if stochastic training layers such as dropout remain active.

Parameters:

model – A TransformerBridge (or compatible) exposing cfg.n_layers, hook_dict, and hooks().
clean – Clean token ids, shape [batch, seq].
corrupt – Corrupt token ids, shape [batch, seq], paired row-by-row with clean.
metric_fn – Maps single-example logits to a scalar to differentiate.
config – Sweep configuration. Node and edge granularity are both supported with plain attribution (ig_steps=1); ig_steps>1 raises at construction.

Returns:

An AttributionResult whose node_scores are averaged over the batch. For an edge-granularity sweep, edge_scores is populated too and node_scores is the per-writer aggregate of those edge scores. The graph’s terminal logits reader gives every writer an edge carrying its direct skip-connection contribution to the metric, so this aggregate equals the quantity a node-granularity sweep on the same model returns (including for the final layer’s MLP output, whose only outgoing edge is the one to the logits reader).

Raises:

ValueError – if clean/corrupt are not 2D, hold a different number of pairs, a pair tokenizes to different lengths (activations must align position-by-position), or the model or one of its submodules is in training mode.

transformer_lens.tools.analysis.attribution_patching.cache_activation_and_gradient(model: Any, tokens: Tensor, metric_fn: Callable[[Tensor], Tensor], names_filter: str | Sequence[str] | Callable[[str], bool] | None = None, compute_gradient: bool = True) → GradientCache¶

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

Registers a forward hook and a backward hook at each cached point, runs one grad-enabled forward, then drives a single backward with torch.autograd.grad() to fire the backward hooks. run_with_cache(..., incl_bwd=True) only backpropagates the model’s own scalar output, so a custom (non-scalar-output) metric such as a logit-diff needs the backward driven here.

Gradients come from the backward hooks, not from .grad on the cached tensors. TransformerBridge reshapes the tensor handed to a forward hook at a converted point (attn.hook_z, hook_q/k/v, hook_attn_out), so that tensor is a view the model’s forward never consumes: retain_grad() on it is inert and its .grad stays None. A backward hook goes through the same conversion and delivers the real gradient in canonical shape. Driving the backward with torch.autograd.grad() instead of metric.backward() keeps it off every parameter’s .grad buffer — no caller grads clobbered, no model-sized buffer allocated.

With compute_gradient=False the backward is skipped entirely: only forward hooks are registered, no backward is driven, and every cached point’s gradient is None. Callers that read activations only (the clean pass of attribution_patch(), which pairs these with the corrupt run’s gradients) use this to run a plain forward instead of a needless forward + backward.

Parameters:

model – A TransformerBridge (or compatible) exposing cfg.n_layers, hook_dict, and the hooks() context manager.
tokens – Input token ids for a single forward pass.
metric_fn – Maps the model logits to a scalar to differentiate.
names_filter – Restricts which hook points are cached (and have gradients retained). None (the default) caches the node-granularity hook set — hook_embed plus each layer’s attn.hook_z and hook_mlp_out; pass an explicit filter to cache any hook point outside that set. On a real Bridge, None cannot mean “every hook point”: the gated points (hook_mlp_in, attn.hook_result, the split-QKV inputs) that hook_dict exposes raise in add_hook unless their set_use_* flag is on.
compute_gradient – When True (default) capture gradients via backward hooks. When False run an activation-only forward and leave every gradient None.

Returns:

A GradientCache with per-hook activations, and gradients when compute_gradient is True (all None otherwise).

transformer_lens.tools.analysis.attribution_patching.enumerate_edges(model: Any, cache: GradientCache) → list[tuple[Node, Node]]¶

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

At a fixed sequence position, the residual stream is a running sum: a reader (a head’s split Q/K/V input, a layer’s MLP entry, or the terminal logits readout) is fed by every writer (the embed write, every attention head’s output, every layer’s MLP output) that precedes it. Building the graph position-by-position tracks which writers are “available” so far and connects each new reader to all of them, then adds that layer’s writers to the available set before moving on – so a writer never edges to a reader upstream of it, and a writer feeding both a direct edge and a through-MLP edge produces two distinct (u, v) pairs rather than one summed together.

A terminal logits reader (read at the final hook_resid_post) closes the graph: after the per-layer loop every remaining writer – including the final layer’s MLP output, which no per-layer reader sees – edges to it. That edge carries the writer’s direct skip-connection contribution to the metric, so a writer’s aggregate over its outgoing edges equals its direct node score.

On cfg.parallel_attn_mlp models (Pythia, GPT-J, Falcon, Phi) the MLP reads the layer input, not the post-attention residual, so a layer’s own heads are not writers into that layer’s MLP; those same-layer head->mlp_in edges are dropped while the heads still feed later readers and the logits reader.

Parameters:

model – A TransformerBridge (or compatible) exposing cfg.n_layers and, optionally, cfg.parallel_attn_mlp.
cache – A GradientCache holding at least the required hook points for edge granularity.

Returns:

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

Raises:

ValueError – if any required hook point is absent from cache – the graph is never silently truncated.

transformer_lens.tools.analysis.attribution_patching.enumerate_nodes(model: Any, cache: GradientCache) → list[Node]¶

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

The graph is explicit: for n_layers layers it always contains the embed write, every attention head’s output, and every layer’s MLP output, at every sequence position. Sequence length and head count are read from the cached tensor shapes; n_layers from model.cfg.

Parameters:

model – A TransformerBridge (or compatible) exposing cfg.n_layers.
cache – A GradientCache holding at least the required hook points.

Returns:

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

Raises:

ValueError – if any required hook point is absent from cache — the graph is never silently truncated.