[Proposal] Jacobian Lens follow-up: validation deliverables + deferred PR2 scope

[jlarson4]

Proposal

The Jacobian lens vertical slice from #1505 has landed on dev (#1507, plus registry short-name resolution in #1537). This tracks the #1505 commitments that did not land with the first PR, in two tiers so the correctness items can ship independently:

• Tier 1 — the three unrealized validation-plan checks.
• Tier 2 — the second demo model + the explicitly-deferred PR2 scope.

Motivation

#1505 committed to a five-point validation plan ("all numbers to be posted in the PR") and named PR2+ follow-ups. #1507 shipped the code, tests, and a demo. This issue covers all follow-ups and any trailing items that didn't make the first PR

Pitch

Tier 1 – Validation Deliverables

• External-oracle parity – Slow test comparing readouts against the official anthropics/jacobian-lens package on fixed prompts, with per-layer top-k overlap + rank-correlation thresholds. Current parity tests are self-consistency + load/validate only. Lock in the [Proposal] Jacobian Lens (J-lens): load published lenses, fit natively, read out and intervene #1505 spike number: worst-case 7/8 top-8 overlap across 75 layer×prompt cells (top-1 identical in 70/75) on gemma-2-2b, using an overlap threshold rather than bitwise equality for bf16 near-ties. @mukund1985 working on this (test(integration): add oracle parity test for JacobianLens (#1539 Tier-1) #1543)
• Kurtosis profile – Check that excess kurtosis of readout logits is roughly 0 in the first third of depth and rises through the workspace band. @rbellec working on this
• Causal swap success-rate – swap_hooks works and the demo shows rank-of-answer qualitatively, but there's no quantified success rate across an intermediate-layer band. Report France→China swap success honestly on an open-weights model.

Tier 2 – Demo + Remaining PR2 scope

• Second demo model – The demo uses only gemma-2-2b; add the promised Bridge-only modern-arch demo (Qwen3.5-4B, the official-walkthrough model). @KYinXu Tackling this component
• Community-format converters – Import fit-checkpoint files / other lens zoos into the JacobianLens artifact format.
• Gradient-pursuit sparse decomposition – J-space coordinate / occupancy analysis over fitted lenses.
• Community fitting guide – Short doc on fitting + merge-parallelizing new-model lenses. Task already address by docs: add Jacobian Lens fitting and sharding guide #1544

Additional context

• Original proposal: [Proposal] Jacobian Lens (J-lens): load published lenses, fit natively, read out and intervene #1505 · Implementation: Add Jacobian lens (J-lens): published-artifact loading, readout, native fitting, and interventions #1507 · Registry follow-up: feat(jacobian_lens): add artifact registry + short-name resolution in from_pretrained #1537
• Reference implementation (Apache-2.0, frozen): https://github.com/anthropics/jacobian-lens — the Tier 1 parity oracle
• Paper: Gurnee et al., Verbalizable Representations Form a Global Workspace in Language Models, Transformer Circuits Thread, 2026 — https://transformer-circuits.pub/2026/workspace/index.html (§4 kurtosis; Methods for the swap protocol)
• Published artifacts: https://huggingface.co/neuronpedia/jacobian-lens

Checklist

• I have checked that there is no similar issue in the repo (required)

[rbellec]

Hi @jlarson4, I'd like to start with a single Tier-1 item: the kurtosis profile check, scoped to just that one so it stays independently shippable per the issue's framing.

My plan: assert that the excess kurtosis of the readout logits is ~0 through the first third of depth and rises across the workspace band (the §4 signature), reporting the measured per-layer profile in the PR rather than baking in a single magic threshold. I'm working through the #1507 slice to see how readouts are exposed and where the check should hook in.

I'd be glad to pick up the other Tier-1 items (causal-swap, parity) once this lands, if that's helpful. My time gets tighter in August, so I'd rather ship one solid check first. Does starting with just kurtosis work for you?

[jlarson4]

@rbellec sounds good! I will put you down as working on that item so that no one overlaps with you unintentionally.

[emerardd]

@jlarson4 I'd like to take the causal swap success-rate item as a focused Tier-1 contribution, independently of @rbellec's kurtosis work.

I've reviewed #1507, the current Jacobian Lens demo, and the paper's companion experiment configuration.

My proposed scope is:

• use the base google/gemma-2-2b model and its matching published lens artifact, intentionally following the existing demo rather than the separate gemma-2-2b-it integration test;
• run the complete country subset: four functions × twelve ordered source-target pairs, giving 48 trials;
• apply each swap at every token position across layers 10–24;
• run both alpha=1 and alpha=2, reporting their success rates separately;
• define success as the target-appropriate answer becoming top-1 in the next-token distribution.

I would report the full per-trial matrix, including baseline/swapped top-1 tokens and target-answer ranks, with France → China highlighted as the direct extension of the existing demo. The protocol and revisions would be fixed before running, and all failures would remain in the reported aggregate.

The demo already notes that alpha=2 can overshoot on this model, sometimes making China itself top-1. I would therefore treat alpha=1 as the primary result and alpha=2 as a separately labelled double-strength diagnostic. I expect this 2B base model to perform below the paper's Sonnet 4.5 country rate.

My default deliverable would be a reproducible 48-trial table and aggregate results in the existing demo, plus focused unit tests for model-independent metric and aggregation logic.

The repository already has a slow Jacobian Lens test using gemma-2-2b-it, but this experiment uses the base checkpoint. Would you prefer an additional slow base-model integration check, or should the pinned real-model results remain in the demo and PR evidence?

[mukund1985]

I'd like to take the external-oracle parity item (Tier 1, item 1).

I wrote #1507 and #1537, so I know the readout path and artifact format well. The parity check is the most critical correctness gate before any Tier 2 work builds on top.

My proposed scope:

• install the frozen anthropics/jacobian-lens reference (Apache-2.0) as a test-only dependency, pinned to a specific commit to ensure the oracle stays fixed
• run both implementations on the same fixed prompt set against google/gemma-2-2b, extracting per-layer top-k readout logits from both
• assert worst-case top-8 overlap >= 7/8 across all 75 layer×prompt cells (matching the [Proposal] Jacobian Lens (J-lens): load published lenses, fit natively, read out and intervene #1505 spike number), plus Spearman rank-correlation threshold >= 0.95 for the top-64 logits per cell
• use overlap threshold rather than bitwise equality to accommodate bf16 near-ties, consistent with the issue framing
• mark the test @pytest.mark.slow and gate it behind the existing real-model CI marker so it doesn't slow down fast CI

I'll file a PR once the numbers lock in. Happy to coordinate with @rbellec and @Dreamer431 so all three Tier 1 items land independently.

[KYinXu]

@jlarson4 I'd like to take the Tier-2 item for adding the second model Qwen3.5-4b to the demo walkthrough

My proposed scope:

• extend the existing Jacobian Lens notebook rather than create a second one;
• load Qwen/Qwen3.5-4B through TransformerBridge
• use the official walkthrough’s pinned n=1000 artifact
• compare J-lens and logit-lens readouts on the walkthrough’s two-hop prompt across representative layers
• report the observed tokens, runtime, and peak GPU memory
• keep the Gemma intervention and steering sections unchanged
• validate the notebook with nbval.

I would limit this PR to the notebook unless running Qwen exposes a source-level correctness issue.

[jlarson4]

@mukund1985 Sounds good, I see your PR is already up, I will take a look at it soon.

@KYinXu That is a great plan!! Thanks for tackling that, I look forward to seeing the PR, the item is all yours