Add Jacobian lens (J-lens): published-artifact loading, readout, native fitting, and interventions

[danielxmed]

Description

Implements the Jacobian lens (J-lens) as a TransformerBridge-only analysis component, following the design proposal in #1505 and the estimator from Gurnee et al., Verbalizable Representations Form a Global Workspace in Language Models.

The J-lens characterizes an intermediate residual activation by its corpus-averaged first-order causal effect on the final residual stream. Each fitted source layer has one d_model × d_model transport matrix; readout then uses the model's own final norm, unembedding, and architecture-specific output-logit transform. The logit lens is available through the same path by disabling Jacobian transport.

Fixes #1505.

What is included

• JacobianLens.from_pretrained, load, and save, compatible with the published neuronpedia/jacobian-lens artifact format. TransformerLens artifacts may add recursively validated provenance metadata while remaining loadable with weights_only=True.
• readout, returning retained top-k values and token ids by default, with full vocabulary logits available as an explicit opt-in. The final layer reuses the model's actual logits rather than recomputing them.
• Native deterministic fitting with Bridge backward hooks. The estimator sums causal target effects t' >= t, averages valid source positions per prompt, then averages prompts equally. merge combines disjoint prompt shards using positive n_prompts weights and requires matching provenance.
• steering_hooks, ablation_hooks, and pseudoinverse-coordinate swap_hooks, with shared position, shape, dtype, and device validation.
• Per-layer/per-device Jacobian caching, Hub retry integration, optional Hub revision pinning, and complete fit provenance including model name, resolved model revision, TransformerLens version, raw-weight contract, hook convention, corpus, prompt count, estimator settings, and dtype.
• An executable Gemma-2 demo in demos/Jacobian_Lens_Demo.ipynb.

Supported model contract

JacobianLens requires a freshly booted, raw-weight, causal decoder-only TransformerBridge with a single rank-3 residual stream per block and a direct final-norm/unembedding path. HookedTransformer, compatibility mode, manually processed weights, encoder/seq2seq models, recurrent physical-block reuse, multi-stream block outputs, and unsupported final projections are rejected explicitly before fitting or readout.

Output-logit behavior is delegated to architecture adapters. This PR adds explicit contracts for Cohere/Cohere2 scaling, Granite division by logits_scaling, Falcon-H1 multiplication by lm_head_multiplier, Gemma-family softcaps, and fail-closed handling for ambiguous MPT logit_scale configurations.

Review-driven changes

• Corrected and pinned the exact causal estimator with a cross-position toy test.
• Removed HookedTransformer support and the configurable public target_layer path.
• Centralized model validation across fitting, readout, lens vectors, and all intervention builders.
• Hardened metadata safety, Hub downloads, shard merging, top-k memory behavior, device caching, interventions, and swap conditioning.
• Added fail-closed topology, processed-weight, output-transform, model-name, and resolved-revision validation.
• Added real TransformerBridge steering/fitting coverage, a small regular-CI Gemma2 test, and the requested slow google/gemma-2-2b-it published-artifact parity test covering top-k and softcapping.

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Validation

All results below were produced from the final reviewed tree on local WSL/Ubuntu with the local RTX GPU where applicable:

• make test-pr: 3,521 unit, 18 docstring, 150 acceptance, and 856 integration tests passed.
• Slow model tests: the published Gemma-2 parity/top-k/softcap test passed, and both GPT-2 fitting/merge/convergence tests passed.
• Notebook: 8/8 cells passed under nbval.
• make check-format: 640 files clean.
• uv run mypy .: 317 source files clean.
• uv build: sdist and wheel built successfully.
• git diff --check and a secret-pattern scan passed.

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 the feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces in a way that affects backward compatibility

[jlarson4]

Hi @danielxmed! Thanks for putting this together, finally getting around to a full review for you. This is a very comprehensive review with a lot of comments, apologies in advance. For something like this I want to be extra thorough. Please let me know if you run into any questions as you work through them

I have a couple general suggestions that don't tie to any specific line:

1. gemma parity test: #1505 committed to slow gemma-2-2b parity tests and itself noted the CI-cached -it variant has a published lens. A @pytest.mark.slow gemma-2-2b-it parity test is feasible and would give the softcap path and a second architecture CI coverage.
2. I noted in the $1505 that we should drop support for HookedTransformer from this tool. Please do that, and then begin making your way through my comments. Some of them may be obsoleted by dropping that support, ping me if you are uncertain at any point.

[jlarson4]

The header formula is a mean over (prompt, t, t' >= t) triples, but the implementation takes the sum over targets and averages only over sources and prompts (~56× difference at seq 128, skip 16). Anyone calibrating raw J magnitudes from this formula will be misled; the header is the only dissenter, so fix it here.

[jlarson4]

lens_vectors reads model.W_U with no validate_model(model). It's the gateway for steering_hooks/ablation_hooks/swap_hooks, so all three build and run on a fold_ln+center_unembed HT (or compat-mode Bridge; the hook names still resolve) while computing vectors in the wrong residual basis. That's the silent-wrongness L34–37 promises to refuse loudly, and #1505's plan validated on attach, the per-call model API change dropped that invariant. One validate_model(model) here restores it for all three builders.

[jlarson4]

save() stores arbitrary metadata, but load() (L213) uses weights_only=True. A numpy scalar in metadata saves fine and then fails this module's own load() with an opaque UnpicklingError, before the friendly no-'J'-key error is reached. Consider validating metadata values are weights-only-safe primitives at save time.

[jlarson4]

Two robustness gaps:

1. hf_hub_download here isn't wrapped by call_hf_with_retry, so lens downloads lack the 429 backoff every other Hub fetch in the repo gets.
2. revision=None tracks the mutable main of a third-party repo. Wrap the call, and consider documenting revision-pinning in the docstring.

[jlarson4]

load() defaults a missing n_prompts to 0 (L221) and this weighting multiplies that shard by zero. Raise or warn when any merged lens has n_prompts == 0.

[jlarson4]

With this threshold, random-init W_U trips the guard in 20/20 seeds at d_vocab=256k (0/20 at GPT-2's 50k). Random row means scale as ~1/sqrt(d_vocab) and cross 0.01 right around Gemma/Llama-3 vocabs, so fitting on an untrained control model (a standard interp baseline) errors out claiming center_unembed. Your own measurement on L836 (centered ≈ 1e-7 vs trained ≥ 3.4) supports a much tighter threshold with ~10^4 margin both ways:

if rows.mean(dim=-1).abs().max() < 1e-3 * rows.abs().mean():

[jlarson4]

zeros_like(target) makes the cotangent run in the model dtype. fp32 begins only at the post-hoc .float() on L971. For bf16 models that's ~8-bit-mantissa accumulation with no test pinning it and no recorded fit dtype. At minimum record the dtype in metadata and recommend fp32 fits in the docstring, upcasting the captured activations would fix it outright.

[jlarson4]

This line assumes every validated model owns an ln_final module, but the two ends of the chain disagree: _require_unprocessed_weights normalizes normalization_type=None to "" (see L812), so None-norm models pass the fold_ln check, While HookedTransformer.__init__ explicitly skips creating ln_final when normalization_type is None (see HookedTransformer.py:247-249). nn.Module.__getattr__ then raises on the missing submodule. The failure ordering is what we're aiming to fix here: fit() never unembeds anything, so it validates and completes, the crash waits for the first readout.

model = HookedTransformer(HookedTransformerConfig(..., normalization_type=None))  # attn-only toy
lens = JacobianLens.fit(model, prompts)  # validator passes; full fit succeeds
lens.readout(model, prompt)  # AttributeError: 'HookedTransformer' object has no attribute 'ln_final'

The user pays the whole fitting cost, then gets a bare AttributeError from deep inside _unembed that names neither normalization nor the lens contract.

Two possible solutions:

1. Reject normalization_type in (None, "") in _require_unprocessed_weights with an actionable message. This option is loudest, though it also bans fit, which works fine on these models
2. Mirror the model's own forward semantics: final = model.ln_final(batched) if getattr(model, "ln_final", None) is not None else batched. Two lines that keep attn-only toy models fully supported and unembed them exactly the way HookedTransformer itself does.

[jlarson4]

This fixture is only consumed by the readout-agreement test. fit() and all three intervention builders never execute on a real Bridge in CI, and those are exactly the code paths the module-side comments flag. One small Bridge fit-or-steering test here would cover them.

[jlarson4]

With per-position Linear blocks there's no cross-position interaction, so only the t' == t term is nonzero. A deliberately wrong fit planting cotangents at ALL positions passes the closed-form test bit-for-bit. The integration tests are cosine-based or fit-vs-fit, so nothing pins the target set. One toy variant with attention-style mixing closes the hole.

[danielxmed]

Thanks for the thorough review. I reworked the PR around a strict TransformerBridge-only contract and addressed all inline points.

• The estimator definition now matches the causal target sum, with a cross-position test pinning t' >= t.
• Validation now covers fitting, readout, lens vectors, and all interventions. HookedTransformer and configurable target_layer support were removed; incompatible processing, topology, attention, and output paths fail explicitly.
• Persistence, Hub retries and revision pinning, shard merging, device caching, top-k retention, final-logit reuse, intervention position/device handling, and swap conditioning are hardened.
• Fit provenance now records and validates model_name and the resolved model_revision, alongside the TransformerLens version, raw-weight contract, hook convention, corpus, prompt count, estimator settings, and dtype. Reduced-precision fits warn and recommend fp32.
• Coverage now includes real Bridge steering/fitting, causal position mixing, regular-CI Gemma2, and the requested slow google/gemma-2-2b-it published-artifact parity test covering top-k and softcapping.

Final validation is clean: make test-pr passed 3,521 unit, 18 docstring, 150 acceptance, and 856 integration tests; the Gemma slow test passed, both GPT-2 slow tests passed, and notebook validation passed 8/8. Mypy passed across 317 source files, formatting passed across 640 files, and both the sdist and wheel built successfully.