[Proposal] Pretrained J++ lenses through JacobianLens: artifact provenance and filtered readouts

Start with JacobianLens.load() and readout() in transformer_lens.tools.analysis, then inspect existing JacobianLens tests and the pretrained artifact schemas described in the issue. Add fixtures for provenance handling and readout filtering, including padded vocabulary rows and unchanged raw logits. Done means both published schemas can round-trip provenance and optional exclusions affect lens top-k rankings as specified; the issue also calls for a parity case and executable comparison example.

Proposal

Extend JacobianLens to support published J++ lenses with preserved artifact provenance and optional pre-top-k vocabulary filtering. Reuse its existing transport, final-norm/unembedding and intervention paths. The first implementation would cover pretrained inference, without adding another lens class or a native J++ fitting backend.

J++ Lens combines clustered Jacobian estimates, relevance-propagation rules during fitting and readout filtering. Once fitted, it remains one fixed transport matrix per layer, making the existing JacobianLens a natural consumer. Author implementation, original artifacts, Neuronpedia-converted artifacts, including Qwen3.5-0.8B.

Motivation

I tested the published Qwen3.5-0.8B J++ artifact against the authors' code. Its matrices already work through TransformerLens, but two pieces of the workflow are missing:

• JacobianLens.load() reads the converted artifact's J matrices but drops its top-level provenance. On the actual released file, the resulting metadata is {}, losing the recorded method, model configuration and activation-capture convention. The authors' original parameters.jacobians schema is not currently accepted.
• readout() ranks the full vocabulary before retaining top-k. Removing punctuation from those retained results cannot recover the top-k over a filtered vocabulary.

Supporting these contracts would let researchers use the released lenses through the same readout and intervention APIs they already use for J-lenses. It also fits the consolidation in #1907: this extends an existing tool without restoring the removed relevance-rule hooks or adding model-forward overhead.

Pitch

1. Artifact interoperability and provenance. Recognize the original parameters.jacobians/config schema and the converted J/provenance schema. Preserve source attribution and normalize recorded model identity/revision, method, source/target layers, capture point and stream-reduction convention for validation and save/load round trips. Check recorded values against the supplied model, without inventing missing fitting revisions. Keep imported, expert-combined J++ artifacts out of the prompt-shard merge() path unless they provide appropriate additive statistics.

2. Optional filtered lens rankings. Add an exclude_token_ids option to readout(), applied before lens top-k selection, with the exclusions recorded in the result. Model prediction rankings and optional full logits remain unfiltered. The default empty exclusion set preserves current behavior; logit-lens comparisons can use the same policy. Invalid IDs or too few eligible tokens for the requested k should raise clearly.

A small tokenizer utility would build exclusions using the authors' Unicode letter/digit rule, with explicit keep_token_ids for evaluated concepts. It must scan the LM-head vocabulary width, including padding rows, rather than only len(tokenizer). Special-token exclusion is a separate policy: the authors' base rule does not remove every special token.

API sketch; the exclusion utility and argument are proposed:

from transformer_lens.tools.analysis import JacobianLens, readout_token_exclusions

lens = JacobianLens.from_pretrained(
    "neuronpedia/jacobian-lens",
    filename="qwen3.5-0.8b/jpp/Salesforce-wikitext/Qwen3.5-0.8B_jpp_lens.pt",
    revision="e166deb5ad540932100ef803c8c9bcc2a5453694",
    model=model,
)
excluded = readout_token_exclusions(
    model.tokenizer,
    vocab_size=model.cfg.d_vocab,
    keep_token_ids=scored_concept_ids,
)
result = lens.readout(
    model, prompt, layers=[4, 12, 20], positions=[-1],
    exclude_token_ids=excluded,
)

3. Tests and a worked example. Add schema/provenance and merge-contract fixtures, ranking tests for default behavior, Unicode, padding, candidate exceptions and unchanged raw logits, plus a fixed real-model parity case. A small executable example would compare J++/J/logit-lens readouts with explicit filtering settings. The initial model scope would follow the existing single-stream JacobianLens contract.

Alternatives

Users can already load the Neuronpedia-converted matrices with an explicit Hub filename and filter full logits themselves. That works for local experiments, but loses the artifact's provenance and repeats vocabulary/ranking logic. The standalone author package is also available; this proposal makes its pretrained maps usable within existing Bridge workflows. A separate lens class or a training backend is unnecessary for this scope.

Additional context

Completed a local CUDA/fp32 parity gate on Qwen3.5-0.8B: four prompts, including Chinese, at blocks 4/12/20, for 12 prompt/layer pairs. Compared raw HF before Bridge wrapping, the authors' lens implementation and existing TransformerLens readout:

Comparison	Maximum absolute difference
Captured residuals	0
Transported residuals	0
Full lens logits	0
Actual raw-HF vs Bridge model logits	0

All raw and base-filtered top-10 ID lists matched. The authors' separate final-residual unembedding helper differed from the actual model path by at most 2.10e-5, within the preset tolerance. This validates numerical compatibility, not a J++-versus-J accuracy claim.

The vocabulary checks also reproduced a meaningful policy distinction: the base rule excluded 6,133 IDs, while the authors' evaluation rule excluded 6,116 because it preserves 17 scored symbols such as +, × and ÷. All 243 padded LM-head rows were excluded.

Reproducibility pins: author code d7efc0de, lens revision e166deb5ad540932100ef803c8c9bcc2a5453694, model revision 2fc06364715b967f1860aea9cf38778875588b17. The approximately 46 MiB converted lens matched Hub LFS SHA-256 39bf7690b7e1d917ccff9c0fd7418abeaa02fe334691edfb7e814ee890b99c06. I can include the focused reproduction with the implementation.

Does this pretrained-artifact and filtered-readout extension fit the intended JacobianLens scope? I would start with the loader/provenance fixtures, then the optional ranking API and example.

Checklist

• I have checked that there is no similar issue in the repo.