transformer_lens.tools.analysis.jacobian_lens module

Jacobian Lens (J-lens).

The Jacobian lens characterizes an intermediate residual-stream activation by its first-order causal effect on the model’s output, averaged over a corpus of contexts. For each layer \(\ell\) it fits a single \(d_{model} \times d_{model}\) matrix

\[J_\ell = \mathbb{E}_{\text{prompt}}\left[ \frac{1}{|V|}\sum_{t \in V}\sum_{t' \in V,\,t' \geq t} \frac{\partial h_{\text{final},t'}}{\partial h_{\ell,t}} \right]\]

where V is the set of valid positions after the configured leading skip and final-position exclusion. Target-position effects are summed for each source; only source positions and prompts are averaged.

mapping the output of block \(\ell\) to the final block’s output (pre final norm). Reading the lens applies the model’s own final norm and unembedding: \(\text{lens}(h_\ell) = W_U\,\mathrm{norm}(J_\ell h_\ell)\). The logit lens is the special case \(J_\ell = I\). The rows of \(W_U J_\ell\) (“J-lens vectors”) are residual-stream directions associated with single vocabulary tokens, and support causal interventions: steering, ablation, and exchanging one concept for another via a pseudoinverse coordinate swap.

Introduced in Verbalizable Representations Form a Global Workspace in Language Models (Gurnee et al., Transformer Circuits Thread, 2026). The fitting estimator and the artifact format follow Anthropic’s Apache-2.0 reference implementation (anthropics/jacobian-lens), so lenses fitted here interoperate with artifacts published on the Hugging Face Hub (e.g. neuronpedia/jacobian-lens); the interventions are implemented from the paper’s Methods section.

Warning

Published lens artifacts are fitted on raw HuggingFace activations. Jacobian lens supports only a freshly booted TransformerBridge.boot_transformers model, whose weights are raw by default. Compatibility mode and direct process_weights calls change the residual basis and are refused rather than returning silently wrong readouts. The model must also be a causal decoder whose adapter supports text generation, use single-stream block outputs, and expose the standard direct ln_final -> d_model-width unembed output path.

Example:

import torch
from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis import JacobianLens

model = TransformerBridge.boot_transformers("gpt2", device="cpu")
lens = JacobianLens.from_pretrained(
    "neuronpedia/jacobian-lens",
    filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
    model=model,
)
result = lens.readout(model, "The Eiffel Tower is in the city of")
print(result.top_tokens(model.tokenizer, k=5)[8][-1])  # layer 8, final position
class transformer_lens.tools.analysis.jacobian_lens.JacobianLens(jacobians: Dict[int, Float[Tensor, 'd_model d_model']], *, n_prompts: int, d_model: int, metadata: Dict[str, Any] | None = None)

Bases: object

A fitted Jacobian lens: one transport matrix per source layer.

Layer convention (matching the reference implementation and the published artifacts): index l refers to the output of block l at the Bridge-native hook blocks.{l}.hook_out. J[l] maps that activation to the final block’s output, pre final norm. The final layer itself is never fitted (its transport is the identity), so source_layers == [0, ..., n_layers - 2] for a full fit.

jacobians

{layer: [d_model, d_model]} transport matrices, fp32, CPU.

n_prompts

Number of prompts averaged into the fit.

d_model

Residual stream width the lens was fitted for.

metadata

Optional provenance (model name, TransformerLens version, fit hyperparameters). Preserved by save()/load(); artifacts from the reference implementation load with empty metadata.

ablation_hooks(model: Any, tokens: str | int | Sequence[str | int], layers: Sequence[int], *, positions: Sequence[int] | None = None) → List[Tuple[str, Any]]

Hooks that project token directions out of the residual stream.

For each token’s unit lens vector v̂: h <- h - (h·v̂) v̂, applied sequentially when several tokens are given.

Parameters:

  • model – The model the hooks will run on.

  • tokens – Concept token(s) to suppress.

  • layers – Layers to intervene at.

  • positions – Chunk-local positions to ablate (negative indices allowed and normalized on every hook invocation). Defaults to all.

Returns:

[(hook_name, fn), ...] for model.hooks(fwd_hooks=...).

clear_device_cache() → None

Release lazily cached Jacobian copies on accelerator devices.

classmethod fit(model: Any, prompts: Sequence[str], *, corpus: str, source_layers: Sequence[int] | None = None, dim_batch: int = 8, max_seq_len: int = 128, skip_first_positions: int = 16, show_progress: bool = True, metadata: Dict[str, Any] | None = None) → JacobianLens

Fit a Jacobian lens on a hooked model.

Implements the reference estimator exactly. For each prompt: one forward pass (the prompt replicated dim_batch times along the batch axis), then ceil(d_model / dim_batch) backward passes. Each backward plants a one-hot cotangent for one output dimension at every valid target position simultaneously — causal attention guarantees the gradient at source position t is then the sum over target positions t' >= t with no explicit masking. Rows are averaged over valid source positions (the first skip_first_positions and the final position are excluded), and prompts contribute equally to the final mean. There is no randomness: the computation is deterministic given the prompts.

The reference implementation reports that fit quality saturates quickly — on the order of 100 prompts of 128 tokens is usable; the published lenses use up to 1000. Use merge() to parallelize across prompt slices.

Parameters:

  • model – A raw TransformerBridge. Model parameters are temporarily frozen (requires_grad=False) during fitting and restored after. Cotangents and activation gradients use the model dtype; fit with a float32 model for the highest-fidelity estimator.

  • prompts – Prompt strings. Prompts too short to contain a valid position (seq_len <= skip_first_positions + 1) are skipped with a warning and do not count toward n_prompts.

  • corpus – Stable identifier for the prompt corpus or slice, recorded in artifact provenance.

  • source_layers – Layers to fit. Defaults to every layer below the final layer. Negative indices count from n_layers.

  • dim_batch – Output dimensions per backward pass. Higher is faster but replicates the prompt dim_batch times in memory; total backward FLOPs are unchanged.

  • max_seq_len – Prompts are truncated to this many tokens.

  • skip_first_positions – Leading positions excluded from the source average.

  • show_progress – Show a tqdm progress bar over prompts.

  • metadata – Extra provenance merged into metadata.

Returns:

The fitted JacobianLens.

Raises:

  • TypeError – If model is not a TransformerBridge.

  • ValueError – On compatibility mode, invalid provenance or layer indices, or if no prompt was long enough to fit on.

classmethod from_pretrained(name_or_path: str, *, filename: str = 'lens.pt', revision: str | None = None, model: Any = None) → JacobianLens

Load a lens from a local path, a short model name, or a Hub repo.

Resolution order

  1. Local file — if name_or_path is an existing .pt file, load it directly.

  2. Local directory — if name_or_path is a directory, load <name_or_path>/<filename>.

  3. Registry short name or HF model ID — if name_or_path matches a key or alias in the bundled jacobian_lens_registry.json (e.g. "gemma-2-2b" or "google/gemma-2-2b"), the corresponding artifact in neuronpedia/jacobian-lens is fetched automatically. The filename argument is ignored in this case because the registry already encodes the correct subpath.

  4. Explicit Hub repo — otherwise name_or_path is treated as a Hub repo id and filename is used as-is, preserving full backward compatibility (e.g. from_pretrained("neuronpedia/jacobian-lens", filename="gpt2-small/jlens/...")).

param name_or_path:

A local .pt file, a local directory, a short model name such as "gemma-2-2b" or "llama3.1-8b", a Hugging Face model ID such as "google/gemma-2-2b", or an explicit Hub repo id paired with filename.

param filename:

File (or subpath) inside a local directory or an explicit Hub repo. Ignored when name_or_path resolves via the registry.

param revision:

Optional Hub revision (branch, tag, or commit) to pin. When omitted, the Hub repository’s mutable default branch is followed; pin a commit hash for reproducible analyses.

param model:

If given, validate_model() is called so dimension or weight-processing mismatches fail here rather than at first use.

returns:

The loaded (and, if model was given, validated) lens.

Examples:

# Short model name — no need to remember the HF subpath
lens = JacobianLens.from_pretrained("gemma-2-2b", model=model)

# HF model ID also works
lens = JacobianLens.from_pretrained("google/gemma-2-2b", model=model)

# Explicit Hub repo + subpath (backward-compatible)
lens = JacobianLens.from_pretrained(
    "neuronpedia/jacobian-lens",
    filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
    model=model,
)
lens_vectors(model: Any, tokens: str | int | Sequence[str | int], layer: int) → Float[Tensor, 'n d_model']

Residual-stream directions for vocabulary tokens at a layer.

The J-lens vector for token t is row t of W_U J[layer] expressed in layer-layer residual coordinates: v_t = J[layer]^T W_U[:, t].

Parameters:

  • model – The model supplying W_U.

  • tokens – A token string / id, or a sequence of them. Strings must encode to a single token.

  • layer – Source layer for the vectors.

Returns:

One vector per token, fp32, on the model’s device.

classmethod load(path: str) → JacobianLens

Load a lens artifact saved by this class or the reference package.

Parameters:

path – Path to the .pt artifact.

Raises:

ValueError – If the file lacks the J key (e.g. a fit checkpoint rather than a saved lens).

classmethod merge(lenses: Sequence[JacobianLens]) → JacobianLens

Combine lenses fitted on disjoint prompt slices.

The per-layer matrices are averaged weighted by each lens’s n_prompts, matching the reference implementation, so fitting can be parallelized across processes or machines and merged afterwards. Provenance must match across shards (apart from n_prompts), so a merge cannot silently relabel matrices fitted with different models, corpora, dtypes, or estimator settings. The merged count replaces the per-shard count.

Parameters:

lenses – Lenses that agree exactly on source_layers and d_model.

Raises:

ValueError – On an empty sequence or mismatched lenses.

readout(model: Any, input: str | Int[Tensor, 'batch seq'], *, layers: Sequence[int] | None = None, positions: Sequence[int] | None = None, use_jacobian: bool = True, top_k: int = 10, return_full_logits: bool = False) → JacobianLensReadout

Read per-layer vocabulary logits for a prompt.

Runs the model once with caching, transports the residual stream at each requested layer through J[layer] (or the identity when use_jacobian=False — the logit lens), and applies the model’s own final norm, unembedding, architecture logit scaling, and logit soft cap.

Parameters:

  • model – A raw TransformerBridge.

  • input – A prompt string, or a [1, seq] token tensor.

  • layers – Layers to read. Defaults to every fitted layer plus the final layer. The final layer (n_layers - 1) is always read with the identity transport — by construction its lens equals the model’s own output distribution.

  • positions – Token positions to read (negative indices allowed). Defaults to all positions.

  • use_jacobian – Apply the Jacobian transport. False gives the logit-lens baseline through the identical code path.

  • top_k – Number of values and vocabulary ids retained per layer and position. Defaults to 10.

  • return_full_logits – Also retain full vocabulary tensors on CPU. This is opt-in because a 64-token Gemma readout across all layers is roughly 1.7 GB.

Returns:

A JacobianLensReadout.

Raises:

ValueError – If the model fails validate_model(), input is batched, top_k is invalid, or a requested layer has no transport matrix.

save(path: str, *, dtype: dtype = torch.float16) → None

Save the lens in the reference implementation’s artifact format.

The four official keys (J, n_prompts, source_layers, d_model) are written unchanged so the file stays loadable by the reference package; TransformerLens provenance is stored under an additive metadata key.

Parameters:

  • path – Destination .pt path.

  • dtype – Storage dtype. Defaults to fp16 like the reference implementation — Jacobian entries are order-one, so the smaller dtype costs little precision and halves the artifact on disk.

property source_layers: List[int]

Sorted list of layers this lens has transport matrices for.

steering_hooks(model: Any, token: str | int, layers: Sequence[int], *, alpha: float = 4.0, positions: Sequence[int] | None = None) → List[Tuple[str, Any]]

Hooks that steer the residual stream along a token’s J-lens vector.

At each layer the unit-normalized lens vector is added, scaled by alpha times the activation’s median per-position residual norm: h <- h + alpha * median||h|| * v̂. This norm-matched parameterization follows the steering description in the reference implementation’s experiment protocols; the paper’s minimal form is the unscaled h <- h + alpha * v_t, recoverable by passing the raw lens_vectors() output to your own hook. The median (not mean) is used so attention-sink positions — whose residual norms run orders of magnitude above typical positions — do not inflate the scale.

Parameters:

  • model – The model the hooks will run on.

  • token – The concept token to steer toward.

  • layers – Layers to intervene at.

  • alpha – Steering strength scalar; 0 disables. Because of the norm-matched scale, values of order 1 already perturb the stream by roughly its own magnitude.

  • positions – Chunk-local positions to steer (negative indices allowed and normalized on every hook invocation). Defaults to all.

Returns:

[(hook_name, fn), ...] for model.hooks(fwd_hooks=...) or model.run_with_hooks(fwd_hooks=...).

swap_hooks(model: Any, source_token: str | int, target_token: str | int, layers: Sequence[int], *, alpha: float = 1.0, positions: Sequence[int] | None = None) → List[Tuple[str, Any]]

Hooks that swap two concepts’ coordinates in lens space.

The paper’s patching-in-lens-coordinates intervention: with V = [v_s, v_t] and lens coordinates c = V⁺ h (pseudoinverse), the update is h <- h + alpha * V (sigma(c) - c) where sigma exchanges the two coordinates. The component of h orthogonal to span{v_s, v_t} is untouched. alpha=2 is the paper’s “double-strength” swap.

Parameters:

  • model – The model the hooks will run on.

  • source_token – The concept to remove (e.g. " France").

  • target_token – The concept to install (e.g. " China").

  • layers – Layers to intervene at (the paper clamps the swap across an intermediate-layer band).

  • alpha – Swap strength.

  • positions – Chunk-local positions to swap (negative indices allowed and normalized on every hook invocation). Defaults to all.

Returns:

[(hook_name, fn), ...] for model.hooks(fwd_hooks=...).

transport(residual: Float[Tensor, '... d_model'], layer: int) → Float[Tensor, '... d_model']

Map layer-layer activations into the final block’s output basis.

Computes J[layer] @ h per activation vector, in fp32.

Parameters:

  • residual – Activations from the output of block layer.

  • layer – Source layer index.

validate_model(model: Any) → JacobianLens

Check that model matches this lens; raise loudly if not.

Requires a raw causal TransformerBridge with the standard direct final-norm/unembed path, verifies recorded model provenance, residual width and layer range, and enforces the published final-block target convention.

Parameters:

model – A raw TransformerBridge.

Returns:

self, for chaining.

Raises:

  • TypeError – If model is not a TransformerBridge.

  • ValueError – On model provenance or d_model mismatch, out-of-range source layers, compatibility mode, unsupported attention/output paths, or a non-final target convention.

class transformer_lens.tools.analysis.jacobian_lens.JacobianLensReadout(lens_topk_values: Dict[int, Float[Tensor, 'pos k']], lens_topk_indices: Dict[int, Int[Tensor, 'pos k']], model_topk_values: Float[Tensor, 'pos k'], model_topk_indices: Int[Tensor, 'pos k'], tokens: Int[Tensor, 'seq'], positions: List[int], use_jacobian: bool = True, lens_logits: Dict[int, Float[Tensor, 'pos d_vocab']] | None = None, model_logits: Float[Tensor, 'pos d_vocab'] | None = None)

Bases: object

Result of a JacobianLens.readout() call.

lens_topk_values

Per-layer retained top-k pre-softmax values, on CPU.

Type:

Dict[int, jaxtyping.Float[Tensor, ‘pos k’]]

lens_topk_indices

Per-layer retained top-k vocabulary ids, on CPU.

Type:

Dict[int, jaxtyping.Int[Tensor, ‘pos k’]]

model_topk_values

The model output’s retained top-k pre-softmax values, on CPU.

Type:

jaxtyping.Float[Tensor, ‘pos k’]

model_topk_indices

The model output’s retained top-k vocabulary ids, on CPU.

Type:

jaxtyping.Int[Tensor, ‘pos k’]

lens_logits

Optional full per-layer logits, on CPU. Present only when readout(return_full_logits=True) was requested.

Type:

Dict[int, jaxtyping.Float[Tensor, ‘pos d_vocab’]] | None

model_logits

Optional full model logits, on CPU. Present only when readout(return_full_logits=True) was requested.

Type:

jaxtyping.Float[Tensor, ‘pos d_vocab’] | None

tokens

The token ids of the run prompt, [seq].

Type:

jaxtyping.Int[Tensor, ‘seq’]

positions

The (normalized, non-negative) positions the readout covers, aligned with the pos axis of retained top-k and optional full logits.

Type:

List[int]

use_jacobian

Whether the Jacobian transport was applied (False = logit lens).

Type:

bool

lens_logits: Dict[int, Float[Tensor, 'pos d_vocab']] | None = None
lens_topk_indices: Dict[int, Int[Tensor, 'pos k']]
lens_topk_values: Dict[int, Float[Tensor, 'pos k']]
model_logits: Float[Tensor, 'pos d_vocab'] | None = None
model_topk_indices: Int[Tensor, 'pos k']
model_topk_values: Float[Tensor, 'pos k']
positions: List[int]
tokens: Int[Tensor, 'seq']
top_tokens(tokenizer: Any, k: int = 5) → Dict[int, List[List[str]]]

Decode the top-k tokens per layer and position.

Parameters:

  • tokenizer – The model’s tokenizer (model.tokenizer).

  • k – Number of top tokens to decode per (layer, position).

Returns:

{layer: [ [top-k strings] per position ]}, positions aligned with positions.

use_jacobian: bool = True