feat(analysis): causal-swap benchmark for coordinate_patch_hooks

[janmenjayap]

---

Summary

Adds the causal coordinate-swap benchmark for JacobianLens.coordinate_patch_hooks(...) (#1749):
corpus/metric primitives, a baseline-capability filter, a norm-matched control-token selector, a
per-trial runner that shares one decomposition_cache between the real and control conditions, a
bootstrap-CI / fingerprinted-artifact schema, a generation entry point, a checked-in frozen artifact
against GPT-2-small, one cached-model integration test, and a new demo notebook that only reads the
frozen artifact. The library layers are model-free or tested against the existing _ToyBridge
fixture; only the integration test and artifact generation touch a real model, and GPT-2-small needs
no HF_TOKEN.

Implements #1776. Jacobian_Lens_Demo.ipynb is untouched.

---

Motivation

#1749 shipped coordinate_patch_hooks but explicitly deferred measuring whether it causes a
directional change in model output, under controls the codebase doesn't otherwise have for this
mechanism: baseline capability filtering, a norm-matched random-atom control, and bootstrap
uncertainty on every reported rate. #1776 is the tracking issue for that deferred work (named "PR3"
in both #1749's own "Out of scope" section and this repo's roadmap notes); this PR ships it.

---

What ships

New module transformer_lens/tools/analysis/jacobian_lens_causal_swap_benchmark.py, six commits,
layered bottom-up:

1. feat(analysis): causal-swap benchmark corpus and answer metrics — FunctionSpec /
   BenchmarkCorpus / PromptTrialSpec / iter_prompt_trials, and AnswerMetrics /
   compute_answer_metrics (rank/margin/tie scoring against a target token, ported from
   Jacobian_Lens_Demo.ipynb's existing _target_metrics cell as tested library code).
2. feat(analysis): baseline-capability filter — BaselineRecord / filter_baseline_capable,
   splitting prompts into capable/excluded by whether the source's own answer is already the
   deterministic argmax, so a benchmark trial is never run against a prompt the model can't answer to
   begin with.
3. feat(analysis): norm-matched control-token selection — select_norm_matched_control_token,
   picking a token whose lens_vector_dictionary atom norm is within a relative tolerance of the real
   target's, deterministically given a seed, so a trial can distinguish "swapping toward this concept
   mattered" from "any similarly sized coordinate edit would have moved the logits."
4. feat(analysis): coordinate-patch causal-swap trial runner — TrialResult /
   run_causal_swap_trial / run_causal_swap_benchmark, which calls coordinate_patch_hooks for the
   real and control conditions against a shared decomposition_cache, and catches ValueError
   per condition (recording status="skipped_source_inactive") rather than aborting the sweep or
   touching coordinate_patch_hooks's own fail-fast contract.
5. feat(analysis): bootstrap CI and frozen artifact schema — BootstrapResult /
   bootstrap_success_rate_ci (percentile bootstrap, seeded), and build_protocol_manifest /
   fingerprint_manifest / serialize_artifact / load_artifact (the same manifest-plus-SHA256
   recipe Jacobian_Lens_Demo.ipynb already uses, with load_artifact re-verifying the fingerprint on
   read so a hand-edited artifact is rejected rather than silently trusted).
6. docs(notebook): causal-swap benchmark for coordinate_patch_hooks — the generation entry point
   (python -m transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark, no HF_TOKEN
   required), the checked-in frozen artifact
   (demos/data/jacobian_lens_causal_swap_benchmark_gpt2.json), the read-only demo notebook
   (demos/Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb), the docs subsection in
   jacobian_lens_fitting.md, notebook registration (docs/make_docs.py, docs/source/index.md,
   .github/workflows/checks.yml), and the public-symbol exports in
   transformer_lens/tools/analysis/__init__.py (BenchmarkCorpus, FunctionSpec,
   run_causal_swap_benchmark, bootstrap_success_rate_ci, load_artifact).

---

The frozen artifact

demos/data/jacobian_lens_causal_swap_benchmark_gpt2.json (schema version 1) is generated against the
published GPT-2-small Jacobian lens (neuronpedia/jacobian-lens,
gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt, revision a4114d77) over the
countries corpus and layers 0-10, with alpha=1.0, k=25, control_tolerance=0.1,
control_seed=0. It records 198 trials and 10 baseline-excluded prompts. On this model/lens/corpus
the current run reports a null effect: both the real-swap and the norm-matched control success
rates are 0.0, each with a 95% percentile-bootstrap CI of [0.0, 0.0] over 10,000 resamples. That is a
legitimate, honestly reported outcome, not a bug: the benchmark's contribution is the controlled
methodology and the uncertainty on every rate, not a headline success number. The notebook renders
the artifact and never calls the model.

---

Design decisions (from the tracking issue #1776)

• Single layer per trial, swept as a trial dimension, never a fixed multi-layer band —
  coordinate_patch_hooks re-decomposes the live activation at each hooked layer, and stacking layers
  is order-dependent and can raise (confirmed on GPT-2: [6,7]/[6,9]/[6,10] raise, [6,8]
  survives). Every hook installation uses exactly one layer and positions=[-1].
• Harness-level skip-and-record fail policy — the trial runner's own try/except ValueError,
  not a change to coordinate_patch_hooks's existing fail-fast contract.
• Shared decomposition_cache between the real and control conditions — both conditions patch the
  same pre-hook activation at the same (layer, batch, position); only the first call performs the
  vocabulary-scale scan, verified by an explicit test.
• GPT-2-small as the reference model for the frozen artifact and integration test, to avoid the
  HF_TOKEN gate the gemma family needs.

---

Non-goals

A successful swap in the artifact would show a directional causal effect under the stated controls on
GPT-2-small, not proof of unique causal mediation, exhaustive concept coverage, or a result that
transfers to closed-weight models. No new prompt corpora beyond the reused country/function set, no
multi-layer band protocol, and no change to Jacobian_Lens_Demo.ipynb.

---

Testing

• tests/unit/tools/test_jacobian_lens_causal_swap_benchmark.py — model-free: prompt-trial
  enumeration, answer-metric rank/margin/tie arithmetic and its validation errors, baseline-capability
  splitting, norm-matched control-token selection (determinism, tolerance, exclusion, self-exclusion),
  bootstrap CI bounds/determinism/empty-input rejection, manifest fingerprinting, and
  serialize/load round-trip including a tampered-fingerprint rejection.
• tests/unit/tools/test_jacobian_lens_causal_swap_benchmark_trials.py — _ToyBridge-backed: shared
  decomposition_cache reuse between real/control conditions, skip-and-record on an inactive source,
  baseline-incapable prompts excluded from the trial set, and seeded determinism of a full benchmark
  run.
• tests/integration/test_jacobian_lens_causal_swap_benchmark.py::test_causal_swap_benchmark_gpt2_smoke
  — cached GPT-2-small, structural assertions only (no specific success-rate claim), including that
  each surviving trial's control token is norm-matched to its target within tolerance.

Local gates verified:

• uv run pytest tests/unit/tools/test_jacobian_lens_causal_swap_benchmark.py tests/unit/tools/test_jacobian_lens_causal_swap_benchmark_trials.py -q — 27 passed.
• uv run mypy . — Success: no issues found in 399 source files.

---

Related work

• [Proposal] Causal coordinate-swap benchmark for coordinate_patch_hooks #1776 (open, tracking issue) — the proposal this PR implements.
• feat(jacobian_lens): add dynamic J-space coordinate-patch hooks #1749 (merged) — coordinate_patch_hooks, the mechanism under benchmark here; named this
  benchmark as its own deferred "PR3."
• feat(jacobian_lens): anchored J-space coordinate patching (offline primitive) #1741 (merged) — the offline coordinate_patch(...) primitive coordinate_patch_hooks wraps;
  no overlap.
• Add clean-coordinate clamping for JacobianLens swaps #1747 (merged) / JacobianLens.swap_hooks re-reads live coordinates, causing odd/even cancellation across layer bands #1746 (closed) — swap_hooks-specific clean-coordinate clamping; different
  mechanism, no functional conflict with this PR.

---

Checklist

• uv run mypy . clean.
• Unit tests (test_jacobian_lens_causal_swap_benchmark.py + _trials.py) pass.
• source .env && uv run pytest tests/integration/test_jacobian_lens_causal_swap_benchmark.py run.
• uv run build-docs and nbval on Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb run.
• make test-pr run (unit + docstring + acceptance + integration).
• Maintainer sign-off on single-layer-per-trial scope, skip-and-record fail policy, and
  GPT-2-small model/lens choice (carried over from [Proposal] Causal coordinate-swap benchmark for coordinate_patch_hooks #1776).

---