Quantify Jacobian Lens causal-swap success rate

[emerardd]

Description

Quantifies Jacobian Lens causal-swap success on an open-weights model in the
existing executable demo.

• Pins the base google/gemma-2-2b model, matching published lens artifact, and
  companion country configuration to exact revisions.
• Runs four functions over all twelve ordered country source-target pairs
  (48 trials per alpha), applying each swap at every token position across
  layers 10-24.
• Reports alpha=1 and alpha=2 separately, including baseline/swapped top-1
  tokens, target-answer ranks, and the France to China slice.
• Reports the unperturbed source-answer top-1 rate beside unconditional swap
  success, plus swap success conditioned on a correct baseline.
• Splits the evaluation into focused protocol, metric, sweep, aggregation, and
  visualization cells. A grouped bar chart compares baseline and swap rates,
  while two heatmaps show every target-answer rank.
• Defines success as the target-answer token being the deterministic argmax of
  the swapped next-token logits, retains all failures, and reports exact
  maximum ties separately.
• Keeps metric and aggregation logic local to the notebook, with executable
  assertions rejecting duplicate, missing, or alpha-mismatched trial matrices.

The measured base-model results are:

• Baseline source-answer top-1: 33/48 trial rows (68.75%), equivalent to
  11/16 unique function-source prompts.
• alpha=1: 16/48 successes (33.33%).
• alpha=1, conditioned on a correct baseline: 13/33 successes (39.39%).
• alpha=2: 0/48 successes; the country token itself becomes top-1 in every
  trial, consistent with double-strength overshoot.
• France to China baseline: 3/4 source answers top-1.
• France to China at alpha=1: 1/4 successes unconditionally and 1/3 among
  baseline-correct trials (language only).
• France to China at alpha=2: 0/4 successes.

This uses TransformerBridge with raw Hugging Face weights. It does not use the
deprecated HookedTransformer path, adds no dependency, and keeps the pinned
real-model evaluation in the demo rather than adding a separate slow test.

Addresses the causal swap success-rate item in #1539.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Screenshots

Not attached separately. The executed notebook contains the aggregate chart,
per-trial rank heatmaps, a compact summary table, and a collapsible full result
matrix.

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

The unit-test checklist is not applicable to this notebook-only change. The
notebook-specific CI-equivalent check passes locally.

Local verification

• uv run pytest --nbval-sanitize-with demos/doc_sanitize.cfg demos/Jacobian_Lens_Demo.ipynb
  • 17 passed in 274.47s after the final chart-coordinate adjustment.
• git diff --check
  • passed.

[jlarson4]

Great work on this, it looks awesome @emerardd! Approved, holding until dev releases so that all items for #1539 land on the same release.