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 granularity with plain attribution (ig_steps=1). Edge scoring (EAP), the integrated-gradient path (EAP-IG, ig_steps>1), and ablate-outside faithfulness are not implemented yet; their API is declared here — granularity="edge" and ig_steps>1 raise 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).

Type:

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

edge_scores

Per-edge effect estimate keyed by (source, destination). Declared here so the result API is stable across the PR series; it is populated only once edge scoring lands and is 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 highest-magnitude edges — populated once edge scoring lands.

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 granularity with plain attribution only. granularity="edge" and ig_steps>1 are accepted by the type but raise NotImplementedError at construction, so downstream code can import and reference this API now while edge scoring and the integrated-gradient path are 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'], position: int, layer: int | None = None, head: int | None = None)

Bases: object

A node in the residual-stream computational graph at node granularity.

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:

  • "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.

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); the per-head attn_head_out node slices head self.head out of the shared hook_z tensor.

kind: Literal['embed', 'attn_head_out', 'mlp_out']
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 node’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), then scores each node with the first-order Taylor estimate effect(node) = (a_clean - a_corrupt) . g.

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 scores are averaged across the batch before ranking.

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. This PR supports node granularity with plain attribution (ig_steps=1) only; other values raise at construction.

Returns:

An AttributionResult whose node_scores are averaged over the batch. edge_scores stays empty until edge scoring lands.

Raises:

ValueError – if clean/corrupt are not 2D, hold a different number of pairs, or a pair tokenizes to different lengths (activations must align position-by-position).

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