feat(attribution_patching): edge attribution (EAP) scoring on TransformerBridge

[janmenjayap]

---

Summary

Extends attribution_patching with edge-level attribution (EAP): every residual-stream writer -> reader edge is scored from the same single gradient cache that node-level attribution already produces, with no extra forward/backward passes. granularity="edge" no longer raises NotImplementedError; ig_steps>1 (EAP-IG) still does and lands in a later
PR.

An edge u -> v carries writer u's contribution into reader v's input. Its score is (a_clean[u] - a_corrupt[u]) . d(metric)/d(input of v), evaluated at v's read point. Sweeping every (u, v) pair from the clean and corrupt caches yields all edge scores in one pass.

---

What's included

• Per-head hook requirements. Edges into/out of attention heads need hook_result (writer) and split hook_q_input/hook_k_input/hook_v_input (reader), which only fire once cfg.use_attn_result and cfg.use_split_qkv_input are on. These are now turned on before caching, and the missing-hook check is generalized so any granularity's cache is validated the same way rather than duplicating it per graph.
• Edge enumeration in the typed graph. enumerate_edges(model, cache) builds every writer -> reader pair position-by-position over the residual stream's running sum, sharing (layer, position, head) keys with the activation and gradient caches. NodeKind gains reader kinds (q_input/k_input/v_input/mlp_in); a writer feeding both a direct edge and a through-MLP edge yields two distinct pairs, enforced by a uniqueness guard rather than silently summed.
• Edge scoring wired into attribution_patch. granularity="edge" populates AttributionResult.edge_scores; node_scores becomes each writer's aggregate over its own outgoing edges. top_edges() ranks by absolute score instead of raising.
• Exact-patch parity, mutation-checked reconstruction, and sum-to-total (edge form) tests on a fully linear 2-layer toy Bridge: a genuine single-edge activation patch matches _edge_effects' sign and (on this linear toy) magnitude; perturbing one writer's captured contribution changes only the edges it feeds and nothing else; a reader's incoming edge scores sum to the same value as scoring its cached input-delta directly. generic_activation_patch itself isn't used for the parity check — its model: HookedTransformer parameter is enforced at runtime by this repo's jaxtyping/beartype test configuration, which rejects a TransformerBridge instance outright, so the check patches directly through the Bridge's own hooks(), the same mechanism generic_activation_patch uses internally.
• Memory caveat documented: per-head [batch, seq, n_heads, d_model] tensors are fine on gpt2-small and do not scale to large models.

---

Files

• transformer_lens/tools/analysis/attribution_patching.py (extended)
• tests/unit/tools/test_attribution_patching.py (extended)

---

Scope / follow-ups

Edge-level EAP only, ig_steps stays 1. The following land in follow-on PRs: ablate-outside faithfulness, integrated gradients (EAP-IG, ig_steps>1), oracle parity, and the demo notebook.

Addresses part of #1742 — this PR covers edge-granularity attribution (EAP); faithfulness, EAP-IG, oracle parity, and the demo are tracked as follow-on work above.

---

Checklist

• make unit-test green
• uv run mypy . clean
• make format applied
• granularity="edge" produces edge_scores/top_edges; ig_steps>1 still raises NotImplementedError
• Exact-patch parity (sign, and magnitude within tolerance on the 2-layer toy) and mutation-checked reconstruction both hold
• Per-head hooks required and raised on when missing; memory caveat documented
• No HookedTransformer reference added; all model interaction via the Bridge

---

[koriyoshi2041]

Reviewed exact head 7a436b73586455681c46705c4a7b056453d45ecc. I traced the graph in residual-stream order and checked the three core invariants: writers are only connected to downstream readers, attention writer deltas are taken in hook_result d_model space while Q/K/V reader gradients use the matching head/position index, and incoming edge scores reconstruct each reader input-delta dot gradient. The final-readout omission means edge-mode node_scores is intentionally an outgoing-edge aggregate rather than node-granularity attribution; the code and docs state that boundary explicitly.

Local focused result: 34/34 attribution-patching tests pass at this exact head. The hosted compatibility, type, format, full-coverage, benchmark, and notebook matrix is green. I did not find a correctness blocker for the scoped EAP implementation.

[jlarson4]

Thanks for picking up the edge portion of #1742, the typed edge graph and the reader-sum test are a solid base. A couple comments below, let me know if you have any questions

[koriyoshi2041]

Exact-head review receipt for 9fb101b42b2acf9223ac2c7425dd098e320971e5 (the approval immediately above was submitted without its intended body):

• terminal hook_resid_post reader closes the graph, including the final MLP writer, and outgoing-edge aggregates reconstruct direct node scores;
• parallel_attn_mlp excludes same-layer head → MLP edges while preserving later-reader and terminal edges;
• reader-grouped batched matrix-vector scoring matches the per-edge reference;
• edge-only hook flags restore across normal and exceptional exits, including use_attn_in.

Local exact-head result: uv run pytest tests/unit/tools/test_attribution_patching.py -q → 50 passed. Hosted compatibility, coverage, type, format, benchmark, and notebook checks are green. I found no remaining correctness blocker for the scoped EAP implementation.

[jlarson4]

@janmenjayap All requests resolved expertly, thank you for taking care of those. Approved and merging now