Add Jacobian lens (J-lens): published-artifact loading, readout, native fitting, and interventions#1507
Conversation
…nterventions Implements the Jacobian lens (Gurnee et al., Transformer Circuits 2026) as a tools/analysis component working with both TransformerBridge (raw weights) and HookedTransformer (from_pretrained_no_processing): - load/save in the official anthropics/jacobian-lens artifact format (fp16 storage, fp32 in memory), including artifacts published on the HF Hub - per-layer readout through the model\x27s own final norm, unembed, and logit soft cap, with the logit lens as the J=I baseline in the same code path - native fitting reproducing the reference estimator (one-hot cotangents at all valid target positions, sum over targets / mean over sources / mean over prompts), plus n_prompts-weighted merge() for parallel fitting - interventions from the paper: steering, ablation, pseudoinverse coordinate swap in lens coordinates - loud guards refusing fold_ln/compatibility-mode models (the published lenses are fitted on raw HF activations) Unit tests pin the estimator to a closed-form linear toy model. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Covers Hub artifact loading/validation, the final-row identity, the rank-collapse late-layer invariant, logit-lens divergence mid-stack, HookedTransformer/TransformerBridge raw-basis agreement, a directional steering smoke test, processed-model refusal, the exact fit==merge identity, and small-fit convergence toward the published lens (cosine 0.91 at L10 from 2 prompts). Fixes _unembed to respect the [batch, pos, d_model] component contract under runtime type checking. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
gemma-2-2b via raw TransformerBridge + the published Hub lens: J-lens vs logit-lens readout on a two-hop prompt with an unspoken intermediate (boot -> Italy -> euro), rank-trajectory plot for pinned tokens, the France->China pseudoinverse coordinate swap (top-1 flip on the language template at alpha=1; Beijing rank 968 -> 8 on the capital template), and J-lens steering. Outputs are baked; verified with nbval locally. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…nance wording - Detect center_unembed processing via the exactly-mean-centered W_U signature: from_pretrained(fold_ln=False) keeps normalization_type LN while still centering weights, which the fold_ln marker alone missed (lens vectors on such a model are silently wrong; readout ranks survive only because LN absorbs the centered component). - Bounds-check readout layers on the logit-lens path and positions after normalization (out-of-range values previously crashed with a raw KeyError or wrapped silently). - Steering: scale by the median per-position residual norm instead of the mean (attention-sink positions run orders of magnitude hot and made the effective strength prompt-length dependent), and attribute the norm-matched parameterization to the reference implementation experiments, not the paper (whose minimal form is h += alpha*v_t). - Move readout residuals to the unembedding device for sharded models; document merge() metadata handling; reword docstrings/error strings that tracked the reference implementation too closely. - Demo notebook: kernelspec metadata, honest alpha=2 overshoot phrasing (the paper reports alpha=2 recovering failures; gemma-2-2b overshoots), measured-GPU wording; re-executed, nbval clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
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:
- gemma parity test: #1505 committed to slow
gemma-2-2bparity tests and itself noted the CI-cached-itvariant has a published lens. A@pytest.mark.slow gemma-2-2b-itparity test is feasible and would give thesoftcappath and a second architecture CI coverage. - I noted in the $1505 that we should drop support for
HookedTransformerfrom 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.
|
|
||
| .. math:: | ||
|
|
||
| J_\\ell = \\mathbb{E}_{\\text{prompt},\\,t,\\,t' \\geq t} |
There was a problem hiding this comment.
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.
| f"layer {layer} is not in this lens's source layers " | ||
| f"({self.source_layers[0]}..{self.source_layers[-1]})" | ||
| ) | ||
| unembed_columns = model.W_U[:, token_ids].float() # [d_model, n] |
There was a problem hiding this comment.
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.
| } | ||
| if self.metadata: | ||
| payload["metadata"] = self.metadata | ||
| torch.save(payload, path) |
There was a problem hiding this comment.
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.
| else: | ||
| from huggingface_hub import hf_hub_download | ||
|
|
||
| local_path = hf_hub_download(name_or_path, filename, revision=revision) |
There was a problem hiding this comment.
Two robustness gaps:
hf_hub_downloadhere isn't wrapped bycall_hf_with_retry, so lens downloads lack the 429 backoff every other Hub fetch in the repo gets.revision=Nonetracks the mutable main of a third-party repo. Wrap the call, and consider documenting revision-pinning in the docstring.
| raise ValueError( | ||
| "all lenses being merged must share the same source_layers and d_model" | ||
| ) | ||
| total = sum(lens.n_prompts for lens in lenses) |
There was a problem hiding this comment.
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.
| # raw checkpoints sit orders of magnitude above the 0.01 threshold | ||
| # (measured: gpt2 3.4, gemma-2-2b 17.2; centered gpt2 ~1e-7). | ||
| rows = unembed_weight[:64].float() | ||
| if rows.mean(dim=-1).abs().max() < 0.01 * rows.abs().mean(): |
There was a problem hiding this comment.
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():
| jacobians = { | ||
| layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in source_layers | ||
| } | ||
| cotangent = torch.zeros_like(target) |
There was a problem hiding this comment.
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.
| """ | ||
| compute_dtype = model.W_U.dtype | ||
| batched = residual.to(device=model.W_U.device, dtype=compute_dtype).unsqueeze(0) | ||
| logits = model.unembed(model.ln_final(batched)).float().squeeze(0) |
There was a problem hiding this comment.
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:
- Reject
normalization_type in (None, "")in_require_unprocessed_weightswith an actionable message. This option is loudest, though it also bansfit, which works fine on these models - 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 wayHookedTransformeritself does.
|
|
||
|
|
||
| @pytest.fixture(scope="module") | ||
| def published_lens(ht_gpt2): |
There was a problem hiding this comment.
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.
| nn.init.normal_(self.linear.weight, std=0.2) | ||
| self.hook_resid_post = HookPoint() | ||
|
|
||
| def forward(self, resid): |
There was a problem hiding this comment.
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.
2 participants
Description
Implements the Jacobian lens (J-lens) as a
tools/analysiscomponent, per thedesign proposal in #1505.
The J-lens (Gurnee et al., Verbalizable Representations Form a Global
Workspace in Language Models,
Transformer Circuits 2026) characterizes intermediate residual-stream
activations by their corpus-averaged, first-order causal effect on the model's
output: one
d_model x d_modeltransport matrix per layer, read through themodel's own final norm and unembedding. The logit lens is the
J = Ispecialcase and is available through the identical code path.
What's included (
transformer_lens/tools/analysis/jacobian_lens.py):JacobianLens.from_pretrained/load/save: the officialanthropics/jacobian-lensartifact format (fp16 storage / fp32 in memory), including the 38-model zoo
published at
neuronpedia/jacobian-lens.Artifacts saved here add a provenance
metadatakey while staying loadableby the official package.
readout(): per-layer, per-position vocabulary logits via the model's ownln_final+unembedplus the config'soutput_logits_soft_capwhereconfigured (e.g. Gemma-2), returning a
JacobianLensReadoutdataclass.fit(): native fitting with TL backward hooks, reproducing the officialestimator (one-hot cotangents planted at every valid target position;
causal masking turns the per-source gradient into the sum over targets
t' >= t; mean over valid source positions; unweighted mean over prompts).merge()combines shard fits exactly (n_prompts-weighted).none):
steering_hooks(norm-matched variant, documented as such),ablation_hooks, and the pseudoinverse coordinateswap_hooks, allreturning
(hook_name, fn)pairs formodel.hooks(...).validate_model/fit/readoutrefuse (a) fold_ln'd HookedTransformers(
normalization_typeending inPre), (b) any compatibility-modeTransformerBridge —
enable_compatibility_mode()processes weights bydefault and a
no_processing=Truecall can't be reliably distinguishedafterwards (the mirror image of DLA's requirement), and (c) models whose
W_Uis exactly mean-centered, the signature ofcenter_unembed— whichfrom_pretrainedapplies even withfold_ln=False. Centering of writingweights alone leaves no post-hoc signature, so the docs state plainly that
from_pretrained_no_processing/ rawboot_transformersare the supportedpaths. Open to a different contract here if preferred.
Works with both
TransformerBridge(rawboot_transformersdefault — therecommended path) and
HookedTransformer(from_pretrained_no_processing).Fixes #1505.
Type of change
Validation
All numbers from real runs (local RTX GPU + a rented 24GB GPU; harness and raw
results retained by the author). Produced at the branch head; the only commit
after the main validation pass is a device-indexing robustness fix re-verified
by the full unit/integration suite.
fit()vs officialjlens.fit()on gpt2 (fp32 CPU, same 3 wikitext prompts, dim_batch 16,seq 64): worst per-layer relative Frobenius error 7.7e-07 across all
11 layers — numerically identical estimator.
prompts, final position, every fitted layer):
top-1 identical 70/75 (Bridge) / 71/75 (HT no_processing); model rows
top-1 identical. Residue is bf16 near-ties.
overlap 8/8 in 85 of 93 cells and at least 7/8 in all; top-1 identical
91/93; tokenization and model rows identical. No architecture
special-casing needed in the tool.
J -> I): gpt2 rank of the model's top-1 token underthe J-lens collapses 8577 (L0) -> 12 (L10) -> 0 (L11); gemma-2-2b
J-lens/logit-lens top-10 overlap goes 0/10 (L13) -> 7/10 (L24) -> 10/10 (L25).
n=277: cosine 0.51/0.66/0.91 at L0/L5/L10. gemma-2-2b (seq 128, official
wikitext corpus) vs published n=454: n=25 cosine from 0.868 (L0, minimum
across all layers) to 0.998 (L24); n=100 from 0.941 to 1.000 — rising with
depth apart from a shallow dip around L2-L5. n=100 readouts reach 8/8 top-8
agreement with the published lens at late band layers. Qwen3.5-4B: even a
5-prompt fit reaches cosine 0.645 (L0) to 0.998 (L30) vs the published
n=1000 lens, crossing 0.90 by L17.
through the mid-to-late band reproduces on both models (gemma-2-2b:
~0.17 -> 0.73 over L20 -> L24; Qwen3.5-4B: ~1.0 -> ~2.2 over L12 -> L27),
but the paper's near-zero early-layer plateau does not — early layers show
elevated kurtosis on these small open models, consistent with their known
noisy early readouts (the same heavy-tailed junk-token attractors are
visible in Neuronpedia's readouts; the paper's numbers are from
Claude-scale models). Not encoded as a test threshold.
swap " France" -> " China" at every position, layers 10-24): top-1 flips
' French' -> ' Chinese' at alpha=1 on "Most people in France speak";
' Beijing' moves rank 968 -> 8 on "The capital of France is" (whose
unperturbed top-1 on this base model is the filler ' a'); ' Asia' 8 -> 1
on the continent template. The paper reports alpha=2 recovering some
alpha=1 failures on Sonnet 4.5; on this small model alpha=2 instead
overshoots into verbalizing ' China' directly, so the demo uses alpha=1.
Tests: 20 unit (a closed-form linear toy model pins the estimator and
transport orientation exactly; guards incl. the centered-unembed signature)
steering, exact fit==merge identity, small-fit convergence; 2 marked
slow).make format,mypy, andmake docstring-testpass locally. Fullmake unit-test: 3249 passed with 4 failures confined totests/unit/test_generate_no_tokenizer.py, which fail identically on a cleanorigin/devcheckout in the same environment (py3.12, torch 2.10,transformers 5.8.1):
generate()concatenates a cuda:0 tensor with cputensors at
HookedTransformer.py:2187for tokenizer-free models onCUDA-visible machines, which CPU-only CI never hits. Unrelated to this
additive change; I can open a separate bug report/fix for it.
The issue proposed
@pytest.mark.slowgemma-2-2b parity tests; since basegemma-2-2b is HF-gated, that coverage shipped instead as the offline
validation above, with gpt2 (CI-cached, published lens available) carrying
the in-repo parity tests.
Demo:
demos/Jacobian_Lens_Demo.ipynb(gemma-2-2b, ~6GB GPU): the two-hopboot -> Italy -> euro readout vs the logit lens, rank trajectories, the
France -> China swap, and steering. Executed end-to-end with baked outputs;
nbval passes locally. Not added to the CI notebook matrix (gated model + GPU);
can add it if preferred.
Limitations / follow-ups (PR2 roadmap in #1505)
no artifact registry/aliases, no fit checkpoint/resume (
merge()coversparallelism; official checkpoint files are rejected with a clear error).
steering_hooksuses a norm-matched scale (unit lens vector times theactivation's median per-position residual norm), following the reference
implementation's experiment protocols rather than the paper's minimal
h += alpha * v_t; the docstring spells this out and raw vectors areavailable via
lens_vectorsfor the minimal form.position-targeted protocols are exposed via
positions=.paper's Claude-scale results; the numbers above are what gemma-2-2b and
Qwen3.5-4B actually produce.
Checklist:
docs are generated)
backward compatibility