[Proposal] Attribution Patching + EAP/EAP-IG: linearized activation patching and faithful edge-level circuit discovery

[janmenjayap]

---

Proposal

Add a self-contained tools/analysis/ module that estimates the causal effect of every node and edge in a model's computational graph in a single 2-forward-plus-1-backward pass, instead of one forward pass per intervention. The module implements three layered methods on the same substrate:

1. Attribution patching — a first-order Taylor estimate of activation patching: effect ≈ (clean_act − corrupt_act) · grad_of_metric, computed for all components at once from one gradient cache.
2. EAP (Edge Attribution Patching) — score every edge in the residual-stream computational graph (node = head / MLP / input at a position; edge = residual write→read), not just nodes, yielding a ranked edge list.
3. EAP-IG — replace the single saturating gradient with an integrated-gradients path between corrupt and clean activations, giving markedly more faithful circuits (the ablate-everything-outside-the-circuit metric).
4. Faithfulness evaluation — given a candidate edge set, mean-ablate every edge outside it and report how much of the task metric (e.g. IOI logit-diff) is recovered.
5. A demo on gpt2-small IOI reproducing the known circuit from EAP/EAP-IG scores alone.

This is the linearized-patching substrate that ACDC (Automatic Circuit Discovery — a separate, not-yet-filed sibling proposal for a patch-and-prune circuit-discovery tool that would consume this module's ranked edge list as its pre-filter) and attribution-graph work build on — the single highest-leverage addition to TL's circuit-discovery toolkit.

Status: exploratory proposal, no code changes.
Suggested labels: enhancement, tooling, TransformerBridge, complexity: moderate-to-high

---

Motivation

What TL has today. TL already exposes every primitive this method needs but assembles none of them into an all-at-once attributor. run_with_cache caches activations and can retain their gradients for a manual backward; get_act_patch_direct_path (in tools/analysis/direct_path_patching.py) does a forward-mode linear perturbation of a single A→B edge (no gradients); patching.py:generic_activation_patch does exact (expensive) node patching; direct_logit_attribution maps a component to the logit; the hook graph enumerates the nodes. What is missing is the tool that reads one gradient cache and produces every node/edge score in a single sweep.

Why it extends the existing baseline. Exact activation patching (generic_activation_patch) needs one forward pass per patched node/edge — quadratic in graph size and infeasible for full-graph edge discovery. Attribution patching is its gradient-based linearization: two forward passes (clean, corrupt) plus one backward pass estimate all effects simultaneously (Nanda 2023). get_act_patch_direct_path already scores a single edge, but by a forward-mode linear perturbation (no gradients/backward) that coincides with the first-order gradient estimate only when the intervening LayerNorm is treated as linear; EAP instead reads the full edge-scoring sweep from one gradient cache. EAP-IG then fixes attribution patching's known failure mode — gradient saturation near the clean point makes the first-order estimate unfaithful — by integrating the gradient along a corrupt→clean path.

Why no maintained TL implementation exists. Three reference implementations exist; none is TransformerLens. Live-checked 2026-09-03:

Repo	License	Stars	Last push	Note
hannamw/eap-ig	MIT	86	2026-05-23	Most active; strongest pinned oracle
hannamw/MIB-circuit-track	Apache-2.0	24	2025-06-30	Built on nnsight, not TL; MIB benchmark harness
Aaquib111/edge-attribution-patching	none (unlicensed)	49	2024-05-31	NeurIPS 2024 ATTRIB paper code; stale; unlicensed so cannot be vendored

All three are single-paper research scripts, not packaged as a reusable library API, with no hook-native TransformerLens integration — confirming the audit's original claim. hannamw/eap-ig's recent activity and permissive MIT license make it the strongest pinned oracle of the three (upgrade from an EAP-only oracle). None of this supersedes filing: the method exists elsewhere, but the TL-native tool does not.

SAELens boundary. This tool operates entirely on the model's own activations and gradients through the residual-stream hook graph. It trains no dictionary and loads no external artifact, so there is no SAELens dependency. (Edge attribution into SAE features is a separate, artifact-side extension and is explicitly out of scope.)

---

Pitch

Proposed API

Names adjustable to maintainer preference.

from transformer_lens import HookedTransformer  # or TransformerBridge
from transformer_lens.tools.analysis import (
    attribution_patch, EdgeAttributionConfig, faithfulness,
)

model = HookedTransformer.from_pretrained("gpt2")  # Bridge works unchanged

clean = model.to_tokens("When John and Mary went to the store, John gave a drink to")
corrupt = model.to_tokens("When John and Mary went to the store, Mary gave a drink to")

def metric(logits):                       # IOI logit-diff (Mary - John) on final pos
    return logits[0, -1, mary_id] - logits[0, -1, john_id]

# EAP: score every edge in the computational graph, all at once.
result = attribution_patch(
    model, clean, corrupt, metric,
    config=EdgeAttributionConfig(
        method="EAP-IG",      # {"attribution", "EAP", "EAP-IG"}
        ig_steps=5,           # integrated-gradients path steps (EAP-IG only)
        granularity="edge",   # {"node", "edge"}
    ),
)

result.edge_scores          # sparse [n_edges] signed effect estimates
result.top_edges(k=50)      # ranked (src_node, dst_node, score)
result.node_scores          # [n_nodes] aggregated

# Build a circuit by thresholding, then measure faithfulness.
circuit = result.top_edges(k=50)
faith = faithfulness(
    model, clean, corrupt, metric, circuit,
    ablation="mean",         # mean-ablate everything OUTSIDE the circuit
)
faith.recovered             # fraction of clean-vs-corrupt metric recovered by circuit
faith.full_metric, faith.corrupt_metric   # baselines reported alongside

Design (algorithm)

Let m be the task metric, a_clean/a_corrupt the activation at a node under the clean/corrupt run.

1. Two forward + one backward. Run clean (cache activations), run corrupt (cache activations, retaining their gradients), then compute the scalar m = metric(logits) and call m.backward() manually to populate ∂m/∂a at every node. (run_with_cache(..., incl_bwd=True) only backpropagates the model's own scalar output and does not support a custom non-scalar metric such as logit-diff, so the metric backward is done by hand.)
2. Node attribution (Taylor). effect(node) ≈ (a_clean − a_corrupt) · ∂m/∂a — first-order estimate of the exact patch m(patch node to clean) − m(corrupt).
3. Edge attribution (EAP). In the residual-stream graph, an edge u→v carries u's write into v's read. Its score is (a_clean[u] − a_corrupt[u]) · ∂m/∂(input of v), evaluated with the gradient at v's read point. Sweeping all (u, v) from the two caches gives every edge score at once.
4. EAP-IG (faithful gradient). Replace the single gradient at the clean point with the integrated gradient along the straight path a(α) = a_corrupt + α·(a_clean − a_corrupt), α ∈ (0,1]: IG = (a_clean − a_corrupt) · (1/N) Σ_k ∂m/∂a |_{a(α_k)}. This averages out gradient saturation and is what buys faithfulness.
5. Faithfulness metric. Given a candidate edge set C, mean-ablate every edge outside C (replace with dataset-mean activation) and measure the recovered fraction of the clean→corrupt metric gap. This is the paper's headline evaluation and the honest check that a high-scoring edge set is actually a circuit.
6. Dataset averaging. Steps 1–4 as described operate on a single clean/corrupt pair. A single pair is too noisy to recover the canonical IOI circuit (check 4) — the original IOI and EAP work average scores/gradients over many templates and name pairs (15+ templates, multiple names each). attribution_patch must accept a batch of clean/corrupt pairs and average per-edge scores across the batch before ranking; this is part of the vertical slice, not a follow-up.

Reuse-map table

Sub-step	TL primitive	file:symbol
Clean/corrupt activation caches	run_with_cache	transformer_lens/HookedTransformer.py:run_with_cache
Gradient cache (∂m/∂a)	grad-retaining forward + manual metric.backward() (not incl_bwd=True, which only backprops the scalar model output)	transformer_lens/HookedRootModule.py:run_with_cache (incl_bwd mechanics), exposed via transformer_lens/HookedTransformer.py:run_with_cache
Per-head Q/K/V and output hooks (required for edge granularity)	use_split_qkv_input=True (separate hook_q_input/hook_k_input/hook_v_input per head) and use_attn_result=True (per-head hook_result before the sum into resid) — both must be set on the model config or these hook points don't exist	HookedTransformerConfig
Single-edge forward-mode linearization (context; EAP is gradient-based)	get_act_patch_direct_path	transformer_lens/tools/analysis/direct_path_patching.py
Enumerate nodes/edges of the graph	hook graph / HookPoint, get_full_resid_decomposition	transformer_lens/hook_points.py, transformer_lens/ActivationCache.py
Exact-patch oracle (validation)	generic_activation_patch + setters	transformer_lens/patching.py:generic_activation_patch
Metric readout (logit-diff)	logit_attrs, W_U + final LN	transformer_lens/ActivationCache.py:logit_attrs
IG path interpolation on activations	setter hooks via add_hook	transformer_lens/hook_points.py:HookPoint.add_hook
Which heads/edges are canonical (context)	detect_head	transformer_lens/head_detector.py

Main correctness risk (named)

Risk 1 — Edge-graph bookkeeping: attributing an edge's effect to the wrong (src, dst) pair, or double-counting a residual write. Because every node writes into the residual stream and every downstream node reads the sum, the edge decomposition must map each writer's contribution to each reader's gradient exactly once, with consistent position/head indexing between the activation cache and the gradient cache. A silent off-by-one in layer/position/head indexing, or summing a write into both a direct edge and a through-MLP edge, produces plausible-looking but wrong scores that no smoke test catches.

Mitigations (raise/validate, do not merely warn):

• Sum-to-total assertion. For a small graph, assert that the summed node attributions reconstruct the first-order metric change m(clean) − m(corrupt) to a tolerance; raise on mismatch — this catches double-counting and index misalignment directly.
• Exact-patch parity on a tiny model. On a 2-layer toy model, assert every edge's attribution matches the sign (and, within tolerance, the magnitude) of the exact generic_activation_patch effect for that edge.
• Explicit graph object. Represent nodes/edges as a typed graph with asserted (layer, position, head) keys shared by both caches; raise if a cache is missing a required hook (e.g. hook_pattern/hook_result) rather than silently skipping edges.
• Honesty in the return type. Attribution scores are estimates; faithfulness (ablate-outside-circuit) is the causal check and is reported next to a random-edge-set baseline.

Risk 2 — Edge-level ablation has no ready-made hook. Standard TL hook points (e.g. hook_resid_post) expose the summed contribution of every upstream writer; there is no existing hook that lets you zero out one writer→reader edge while leaving every other edge into that reader untouched. faithfulness() therefore cannot be built by toggling existing node hooks — it needs the same per-edge patching machinery get_act_patch_direct_path uses for a single edge (reconstruct a component's input as the sum of only the included edges' contributions, substituting the corrupt-run value for excluded edges), generalized to sweep every edge in the circuit at once. This is harder than the indexing risk above and should be scoped and estimated explicitly during implementation, not treated as a detail of the bookkeeping fix.

Risk 3 — Attention-pattern nonlinearity degrades edges into Q/K. Edges that feed a head's query or key input pass through the softmax before affecting the metric; the first-order (or IG) linear estimate is measurably less faithful there than for edges into V or MLP inputs, a caveat the EAP paper itself notes. Faithfulness scores should be expected to be worse for pattern-affecting edges, and this should be reported as a known limitation rather than debugged as a bug if the demo shows it.

Validation plan (falsifiable)

1. Reconstruction (unit). On a tiny synthetic model, summed node attributions ≈ m(clean) − m(corrupt) to atol (linearization identity).
2. Exact-patch parity (unit). On a 2-layer toy model with a small edge set, each edge's attribution-patch score matches the exact generic_activation_patch effect in sign and within tolerance in magnitude.
3. IG improves faithfulness (integration). On gpt2-small IOI, a fixed-size circuit built from EAP-IG scores achieves higher ablate-outside faithfulness than one built from plain EAP scores at the same edge budget.
4. IOI circuit recovery (oracle, @pytest.mark.slow). Top-scoring edges, averaged over a batch of IOI templates/name pairs (not a single clean/corrupt pair), on gpt2-small IOI overlap the established IOI head set (name-movers 9.9/9.6/10.0, S-inhibition 7.3/7.9/8.6/8.10, etc.) above a threshold, compared against the pinned reference repo's scores.
5. Baseline honesty (integration). A random equal-size edge set recovers markedly less metric than the discovered circuit; both numbers reported.

Scope of the first PR (vertical slice)

Attribution patching + EAP edge scores + EAP-IG + a faithfulness check on gpt2-small IOI, plus one demo. Exact paths:

• transformer_lens/tools/analysis/attribution_patching.py — new module (attribution_patch, EdgeAttributionConfig, faithfulness, result dataclass with .top_edges()).
• transformer_lens/tools/analysis/__init__.py — export.
• tests/unit/tools/test_attribution_patching.py — synthetic 2-layer fixtures (checks 1–2), no HF downloads.
• tests/integration/test_attribution_patching.py — gpt2-small IOI EAP vs EAP-IG faithfulness + random baseline (checks 3, 5).
• tests/integration/test_attribution_patching_oracle_parity.py — slow IOI recovery vs pinned reference (check 4), @pytest.mark.slow.
• demos/Attribution_Patching_Demo.ipynb — IOI edge scoring → circuit → faithfulness (nbval).

Follow-up work

Full auto-pruning into a minimal subgraph is deliberately deferred to the separate ACDC proposal noted above (this tool is its fast pre-filter). Attribution graphs, SAE-feature edges, multi-task datasets, and speed optimizations are tiered as follow-up work in my own commit plan (not part of this issue; I'll open a tracking issue for these once the vertical slice lands).

Model coverage

• CI reference: gpt2-small — cacheable, canonical IOI circuit for the oracle.
• Demo/extension: Pythia-160m noted as an extension (follow-up).
• Honesty caveat: small open models will not reproduce large-model paper figures verbatim; the demo reports the discovered circuit's faithfulness next to the random-edge-set baseline and the full/corrupt metric, and states that circuit quality is measured by faithfulness, not by score magnitude alone.

---

Alternatives

• Exact activation patching for every edge (generic_activation_patch). Rejected: one forward pass per edge is quadratic in graph size and infeasible for full-graph discovery; it is kept only as the correctness oracle.
• Node-only attribution patching (Nanda 2023). Rejected as the sole deliverable: it misses edge structure that ACDC/EAP need; we include it as the base case but ship edge scoring.
• Plain EAP without integrated gradients. Rejected as the headline method: gradient saturation makes it less faithful (the EAP-IG paper's central finding); we ship EAP as a baseline and EAP-IG as the default.
• SAE-feature edge attribution. Rejected for the slice: requires trained dictionaries (SAELens artifact side); orthogonal to the residual-stream substrate this PR establishes.

---

Correctness oracle

Two-tier:

• Brute-force / analytic reference (primary, fast unit): on a tiny synthetic model, (i) summed node attributions must reconstruct the first-order metric change, and (ii) each edge's attribution-patch score must match the exact generic_activation_patch effect in sign and within tolerance. Deterministic, no downloads. This is the mandatory oracle where no cheap external reference applies.
• External frozen impl (slow, @pytest.mark.slow): hannamw/eap-ig (MIT, most actively maintained — preferred), with Aaquib111/edge-attribution-patching and hannamw/MIB-circuit-track as secondary references, pinned and installed test-only, used for a threshold-based comparison of edge scores / recovered IOI circuit on gpt2-small. Tolerant comparison (rank overlap above threshold), not bit-exact. MIB-circuit-track runs on nnsight; treat only as a cross-check of expected head sets, not an installable oracle dependency.

---

Additional context

• Papers: Neel Nanda, "Attribution Patching: Activation Patching At Industrial Scale" (2023 blog). Aaquib Syed, Can Rager, Arthur Conmy, "Attribution Patching Outperforms Automated Circuit Discovery," arXiv:2310.10348. Michael Hanna, Sandro Pezzelle, Yonatan Belinkov, "Have Faith in Faithfulness: Going Beyond Circuit Overlap When Finding Model Mechanisms" (EAP-IG), arXiv:2403.17806.
• External implementation status: three research repos, none TransformerLens (live-checked 2026-09-03): hannamw/eap-ig (MIT, 86 stars, pushed 2026-05-23 — strongest oracle), hannamw/MIB-circuit-track (Apache-2.0, 24 stars, pushed 2025-06-30, nnsight-based), Aaquib111/edge-attribution-patching (unlicensed, 49 stars, pushed 2024-05-31, stale). All single-paper scripts, no packaged API, no TL hook integration; suitable only as pinned, test-only oracles, not as a substitute for this proposal.
• Repo scope fit confirmed: matches AGENTS.md (repo root) §1's definition of TL — "exposes internal activations through a hook system for caching, editing, and ablating intermediate state" — via existing run_with_cache/ActivationCache/generic_activation_patch primitives, no new runtime dependency.
• Artifact sources: none — method uses only activations + gradients from the model's own hook graph; no dictionaries, no registry for the vertical slice.
• Algorithm note for reviewers. Attribution patching is a first-order Taylor estimate of exact patching; it is exact only in the linear limit and saturates near the clean point. EAP-IG's integrated-gradients path is what restores faithfulness and must be the default. The edge decomposition requires each residual write to be attributed to each downstream read exactly once — the sum-to-total assertion is the guard.
• Provisional-citation flags: arXiv IDs 2310.10348, 2403.17806, and 2211.00593 are pre-2026 and established; no [verify] needed. IOI head set is Wang, Variengien, Conmy, Shlegeris, Steinhardt 2022, arXiv:2211.00593.

---

Checklist

• Checked no similar issue / tool exists (grep of tools/analysis/; only patching-family + direct_path_patching + J-lens present, none doing full-graph edge scoring).
• Vertical slice only: attribution + EAP + EAP-IG + faithfulness on gpt2-small IOI. Auto-pruning deferred to the separate ACDC proposal (pre-filter consumer, not yet filed); attribution graphs / SAE edges / multi-task deferred to a tiered follow-up issue I'll open post-merge.
• Claim-in-comments workflow; each follow-up tier item independently shippable.
• Correctness oracle: brute-force/exact-patch (unit) + pinned hannamw/eap-ig (primary), Aaquib111/edge-attribution-patching / hannamw/MIB-circuit-track (secondary) (slow, threshold).

---

Reviewer note

cc @jlarson4 — could you take a look and let me know if this scope and API shape look right?
If this gets a green light, I'd like to pick it up myself. Could you assign this issue to me once it's approved?
Thank you

---

[jlarson4]

Thanks for putting this together @janmenjayap! You have the green light on scope and API shape, with the amendments below

1. Most Important Build against TransformerBridge. I am approaching completion on the next major release, TransformerLens v4, which includes officially deprecating HookedTransformer. Your Pitch and Reuse-map both reference only Hooked* classes, this will need to be amended to function as part of the Bridge system.

2. Demo naming issue – Attribution_Patching_Demo.ipynb is already an existing demo that demonstrates part of what you're discussing here, just not as its own API (node-level attribution patching on IOI; no EAP, no IG, no faithfulness). We will need to give the new demo a unique name and leave the existing demo in place.

3. Risk 1 – edge bookkeeping

   • The sum-to-total identity is exact only for a linear network. Make the exact-equality assertion on a linear synthetic model
   • Add token-length parity validation: raise if a clean/corrupt pair tokenizes to different lengths
   • Pin the sign/direction convention (your step 1–2 is the denoising form: gradients from the corrupt run, patch-toward-clean) and map it explicitly to the pinned oracle's convention rather than assuming they match

4. Risk 2 – edge ablation
   No single hook does it, but no new model machinery is needed either: capture writer outputs in the same forward pass (hook_result per head, hook_mlp_out, hook_embed), then rewrite each reader's input at the pre-LN fork hooks (hook_q/k/v_input, MLP in-hook) by subtracting excluded writers' live contributions and adding their replacement values. The residual sum is linear, so the correction is exact, and this is how hannamw/EAP-IG itself implements faithfulness. What needs to be scoped explicitly are the following:

   • Memory – use_attn_result + use_split_qkv_input materialize per-head [batch, seq, n_heads, d_model] tensors. Fine on gpt2-small; document that it does not scale to large models rather than discovering it later.
   • Ablation default: faithfulness(ablation="mean") as proposed defaults away from the evaluation your checks 3–4 compare against. EAP-IG's headline faithfulness patches corrupt activations into out-of-circuit edges. Default must match whatever the pinned oracle computes (corrupt-patching), with mean as an option. Otherwise the oracle-parity check compares unlike quantities.
   • The correction-term indexing inherits Risk 1's failure modes and the same mutation-checked reconstruction test should cover it.

5. Risk 3 – Q/K nonlinearity: I agree with your intention on this. One cheap addition: have the demo/faithfulness report break scores out by edge class (into-Q/K vs into-V vs into-MLP) so the limitation is measured and visible.

6. An additional risk I want to note: Gradient-cache memory – Retaining gradients at every hook point roughly doubles cache memory on top of the per-head tensors above. The tool should cache only the hook families it needs (writer outputs + reader inputs) via a names filter, not the full cache.

7. An API note – method={"attribution","EAP","EAP-IG"} overlaps granularity={"node","edge"}
attribution is node granularity, EAP is edge granularity with a 1-point gradient, EAP-IG is edge with an integrated path. Two axes encoding one choice invites invalid combos (method="EAP", granularity="node"). I'd suggest collapsing to granularity + ig_steps (ig_steps=1 ≡ plain EAP/attribution).

Let me know if you have questions about any of the above, happy to answer them. Thanks again for putting together such a thorough proposal!

[janmenjayap]

Hi @jlarson4 — #1742 was closed on merge of #1750, but that PR only shipped PR1 of the five-PR vertical slice we agreed on in review (node-level attribution patching; granularity="edge" and ig_steps>1 both still raise NotImplementedError in the merged code). Per the slice plan, the proposal's "vertical slice only" checklist item isn't satisfied until PR5 (edge scoring + faithfulness + EAP-IG + oracle parity + demo) lands.

Could you reopen #1742 so it keeps tracking the remaining PR2–PR5 work? I'd rather keep this under the existing issue than file a new one — it already holds your amendment comment (edge bookkeeping, sign convention, memory caveats) that PR2 builds on directly, and splitting it would fragment that context. I'll open a fresh tracking issue only for the tiered follow-up work (ACDC hand-off, attribution graphs, SAE-feature edges, etc.) once PR5 merges, per the plan already on the issue.

[jlarson4]

@janmenjayap Reopened! My apologies for closing this early.