[Proposal] Causal coordinate-swap benchmark for coordinate_patch_hooks

[janmenjayap]

---

Proposal

Add a benchmark that measures whether JacobianLens.coordinate_patch_hooks(...) (#1749) causes a directional change in model output, under three controls the codebase does not yet have for this mechanism: baseline capability filtering, norm-matched random-atom controls, and bootstrap uncertainty on every reported rate. Ship the harness as tested library code, a frozen per-trial artifact generated once against a published lens, and a notebook that consumes that artifact rather than recomputing it.

Suggested labels: enhancement, complexity-moderate, TransformerBridge

---

Motivation

coordinate_patch(...) (#1741) and coordinate_patch_hooks(...) (#1749) establish that an anchored J-space coordinate edit can be computed and, now, installed live inside a forward pass. Neither establishes whether performing that edit changes the model's output in the direction the edited concept predicts. proposal.md's own Alternative 5 deferred this deliberately when it was first raised against the offline primitive: "model capability and prompt design would obscure review of the algebraic primitive." That reasoning was carried forward unchanged into #1749's own "Out of scope" section, which named this exact benchmark as PR3.

---

Relationship to the existing swap_hooks benchmark (Jacobian_Lens_Demo.ipynb section 4)

demos/Jacobian_Lens_Demo.ipynb already contains a 48-trial country-swap benchmark (source in the notebook's section "4. Quantifying causal-swap success across the country subset"), shipped with #1507/#1545 well before coordinate_patch_hooks existed. It is the nearest precedent for this issue and this issue reuses several of its conventions directly, but it is a different benchmark, not an earlier version of this one:

• It exercises swap_hooks (a closed-form pseudoinverse exchange over a fixed 2-vector basis, re-applied at every position across a 15-layer band, layers=range(10, 25)), not coordinate_patch_hooks (a full sparse-solve anchored edit, one (batch, position) pair at a time). The two mechanisms have different failure modes under multi-layer composition -- see "Layer scope" below -- and are not interchangeable in a shared harness without conflating what each measures.
• It reports baseline correctness per function (e.g. "currency 0/12") but does not use it as a filter: every trial's swap "success" is counted even on prompts the unperturbed model already answers wrong on 100% of the time (all 12 currency prompts). This issue's baseline-capability filter would drop that entire function's trials rather than count them.
• It has no norm-matched random-atom control: there is no way, from that notebook alone, to tell "swapping France mattered" apart from "any coordinate edit of similar magnitude would have moved the logits this much."
• It reports point-estimate success rates ("31.25%") with no uncertainty interval.
• It recomputes all 96 conditions live inside the notebook, gated only by # NBVAL_IGNORE_OUTPUT, so nbval still executes the model on every CI run of that notebook. This issue's frozen-artifact-plus-notebook split is the opposite of that by design (see "Frozen artifact" below), matching proposal.md's original "no notebook in PR1" reasoning: recomputation lets a displayed number silently drift from the number a reader is told to trust.

This issue reuses that section's country/function/template/answer content (the anthropics/jacobian-lens flexible-generalization.json configuration, pinned by revision, per future-prs.md scope item 1: "reuse an existing TransformerLens prompt corpus where practical"), its _target_metrics-style rank/margin/tie helper, and its protocol-manifest-plus-SHA256-fingerprint convention. It does not modify Jacobian_Lens_Demo.ipynb itself -- that notebook's outputs are already reviewed and pinned to a specific model/lens/config revision triple by #1747's most recent re-execution, and re-running it is out of scope here.

---

Proposed scope

1. Fixed prompt set. Reuse the flexible-generalization.json country subset (4 countries, 4 function templates, 12 ordered source/target pairs per function -- the same content already embedded as COUNTRY_CONFIG in Jacobian_Lens_Demo.ipynb). A genuinely new corpus is out of scope unless this one proves unusable for coordinate_patch_hooks specifically.
2. Baseline capability filtering. For each (function, source) prompt, run one baseline forward pass. Only prompts whose source-appropriate answer is the deterministic argmax proceed to the intervention trials; every other prompt is recorded in a separate exclusion ledger with its baseline rank, not silently dropped and not counted as a trial. On the existing gemma-2-2b numbers this would remove the entire currency function (0/12 baseline-correct) from the trial set -- the concrete case that motivates having this filter, rather than reporting it as descriptive context the way the current notebook does.
3. Norm-matched random-atom controls, operationalized against the actual token-level API. coordinate_patch_hooks takes target_token: TokenInput, not a raw vector, and lens_vector_dictionary(model, layer) is a [d_vocab, d_model] matrix indexed by token id (one atom per vocabulary token) -- so "a random dictionary atom of matched norm" is realizable as a random control token id, drawn so that its atom's norm at the trial's layer is within a configurable relative tolerance of the real target token's atom norm, excluding the source token, the real target token, and both prompts' answer tokens. The control trial then calls coordinate_patch_hooks identically except for target_token=control_token, and is scored against the same real-target-answer metric as the real trial. This isolates "swapping toward this specific concept mattered" from "any coordinate edit of this magnitude would have mattered," without extending the public coordinate_patch_hooks signature to accept raw vectors.
4. Per-trial outputs plus bootstrap uncertainty. Every surviving trial's full record (baseline metrics, real-condition metrics, control-condition metrics, the control token chosen) is retained, not just an aggregate rate. Real and control success rates are each reported with a seeded percentile-bootstrap confidence interval, not a bare fraction.
5. Open-weight-models-only claims. The frozen artifact and its interpretation prose describe only the open-weight model actually run (see "Model and lens choice" below); no claim about closed-weight models' magnitudes, matching status.md correction 5 and feat(jacobian_lens): add dynamic J-space coordinate-patch hooks #1749's own established discipline.
6. Frozen artifact, notebook consumes it only. A generation step (documented, run once, out-of-band) produces a versioned JSON artifact with a protocol manifest and SHA256 fingerprint (same recipe Jacobian_Lens_Demo.ipynb already uses: hashlib.sha256(json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode()).hexdigest()). The new notebook loads that artifact file and renders tables/plots from it; it does not call coordinate_patch_hooks over the full trial matrix itself, so nbval never re-executes 96+ model forward passes for this notebook.

Layer scope: single layer per trial, not a multi-layer band (needs sign-off)

The existing swap_hooks benchmark installs across a 15-layer band (range(10, 25)) because swap_hooks re-reads live coordinates at each layer and (per #1746/#1747) is explicitly designed to be re-applied that way. coordinate_patch_hooks was not: per the review fix already merged into #1749 (see jacobian_lens.py's Raises: note on coordinate_patch_hooks), stacking hooked layers is order-dependent and can raise -- an earlier layer's substitute/swap edit changes the residual a later layer re-decomposes, which can remove the source coordinate from that later layer's active support even though it was active pre-hook. This was observed concretely on GPT-2: [6,7], [6,9], and [6,10] raise; [6,8] survives. A benchmark protocol that installs a fixed multi-layer band across 48+ source/target pairs would therefore raise unpredictably depending on which pairs happen to survive at which layer combinations, which is not a fixed, reproducible protocol.

Recommendation: each trial installs coordinate_patch_hooks at exactly one layer and one position (positions=[-1], the final prompt token, matching the next-token success metric). Layer becomes a swept trial dimension ((function, source, target, layer)) instead of a fixed band, so every trial is an independent single-layer intervention and the compounding failure mode above cannot occur inside a trial. This is a deliberate scope narrowing relative to the existing swap_hooks benchmark's band protocol, not an oversight, and is the first thing that should get explicit sign-off, since it changes what the benchmark can claim (single-layer causal effect, not a multi-layer workspace-broadcast claim the way section 3/4 of Jacobian_Lens_Demo.ipynb do for swap_clamp_hooks/swap_hooks).

---

Harness-level fail policy: skip-and-record, not abort-the-sweep (needs sign-off)

coordinate_patch_hooks itself is fail-fast by design (#1749 design decision 2): if the source is inactive at any (batch, position) pair in one forward pass, that forward pass raises. A benchmark sweeping many (source, target, layer) combinations will hit this for some combinations -- the [6,7]/[6,9]/[6,10] example above shows it is not rare. If the trial runner let that exception propagate, a single inactive-source combination would abort the entire sweep, which is incompatible with "record every trial" (scope item 4). Recommendation: the trial runner catches the ValueError per trial (not inside coordinate_patch_hooks, which keeps its existing contract unchanged) and records that trial as skipped: source inactive with the layer/pair that raised, rather than either aborting the sweep or silently omitting the trial from the output. This is a new, harness-only policy -- it does not relax or contradict coordinate_patch_hooks's own fail-fast contract, which callers outside this benchmark still get unchanged.

---

Model and lens choice (needs sign-off)

Recommendation: GPT-2-small (gpt2-small in jacobian_lens_registry.json), the same model and published lens test_coordinate_patch_hooks_gpt2_no_op_and_leaves_other_positions_unchanged already uses. This avoids requiring the HF_TOKEN gate the gemma family needs (AGENTS.md section 3), and lets the frozen artifact be regenerated by anyone with the existing integration-test fixtures already set up. The alternative -- gemma-2-2b, matching the existing swap_hooks benchmark's model exactly for side-by-side comparison -- is heavier to regenerate and gated; flagging both options rather than assuming the answer. Reusing the country-config content does not require reusing gemma-2-2b as the model: the baseline-capability filter (scope item 2) exists precisely to make the benchmark honest about which of those prompts a given model actually gets right before intervening, so a different base model changes which trials survive filtering, not whether the protocol is valid.

The exact layer(s) swept are an empirical choice made when the artifact is generated (which of published_gpt2_lens's fitted source_layers produce a non-empty, non-degenerate trial set after baseline filtering) and are recorded in the artifact's protocol manifest, the same way the existing swap_hooks benchmark's 10-24 band is a recorded, not hardcoded-in-this-issue, choice.

---

Non-goals

• New prompt corpora beyond the fixed country/function set reused here.
• Any claim of exhaustive concept coverage.
• Closed-weight model results or claims.
• Claims that a successful swap proves unique causal mediation -- report a directional causal effect under the stated controls, nothing stronger (reuses status.md correction 6's language, the same citation feat(jacobian_lens): add dynamic J-space coordinate-patch hooks #1749's own "Design decisions" table points to).
• A multi-layer band protocol for coordinate_patch_hooks (see "Layer scope" above); a future issue could propose one once single-layer results establish whether it is worth the added compounding risk.
• Modifying demos/Jacobian_Lens_Demo.ipynb's existing swap_hooks benchmark cells.
• HookedTransformer support (separate follow-up, PR4, orthogonal to this work).

---

Proposed repository changes

• New module transformer_lens/tools/analysis/jacobian_lens_causal_swap_benchmark.py: corpus loading, baseline-capability filter, norm-matched control-token selection, the per-trial runner built on coordinate_patch_hooks, bootstrap confidence intervals, and artifact serialization/deserialization.
• Export the module's public entry points from transformer_lens/tools/analysis/__init__.py.
• New model-free unit tests in tests/unit/tools/, reusing the existing _ToyBridge fixture from tests/unit/tools/conftest.py (no new fixture needed -- it already supports real forward passes and .hooks(fwd_hooks=...)).
• One cached-model integration test in tests/integration/, GPT-2-small, a deliberately tiny trial set (structural assertions only, no specific success-rate claims).
• A documented, reproducible generation command that produces the frozen JSON artifact, checked into the repository.
• New demo notebook consuming the frozen artifact, registered with the notebook-checks CI job.
• Short docs: the benchmark's protocol, its controls, and the interpretation caveats above.

---

Validation

Unit tests use small synthetic dictionaries and the existing _ToyBridge fixture with hand-computed expected filters, control selections, and bootstrap bounds. The one integration test asserts structural/algebraic properties only (artifact schema valid, control atom norm within the configured tolerance, at least one non-excluded trial produced) -- not a specific success rate, per #1749's own established policy against behavior-dependent test assertions. The frozen artifact's actual numbers and their interpretation are documented prose and notebook content, not test assertions.

---

Checklist

• Confirm no competing open issue/PR for a coordinate_patch_hooks benchmark (repeat duplicate search immediately before filing; JacobianLens.swap_hooks re-reads live coordinates, causing odd/even cancellation across layer bands #1746/Add clean-coordinate clamping for JacobianLens swaps #1747 are swap_hooks-specific and non-overlapping).
• Maintainer sign-off on: single-layer-per-trial scope (not a multi-layer band), the skip-and-record harness fail policy, the norm-matched control-token construction, and the GPT-2-small model/lens choice.
• Tracking issue number added to the PR description once filed.