Coverage for transformer_lens/tools/analysis/jacobian_lens.py: 93%
522 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""Jacobian Lens (J-lens).
3The Jacobian lens characterizes an intermediate residual-stream activation by its
4first-order causal effect on the model's output, averaged over a corpus of contexts.
5For each layer :math:`\\ell` it fits a single :math:`d_{model} \\times d_{model}` matrix
7.. math::
9 J_\\ell = \\mathbb{E}_{\\text{prompt}}\\left[
10 \\frac{1}{|V|}\\sum_{t \\in V}\\sum_{t' \\in V,\\,t' \\geq t}
11 \\frac{\\partial h_{\\text{final},t'}}{\\partial h_{\\ell,t}}
12 \\right]
14where ``V`` is the set of valid positions after the configured leading skip
15and final-position exclusion. Target-position effects are summed for each
16source; only source positions and prompts are averaged.
18mapping the output of block :math:`\\ell` to the final block's output (pre final
19norm). Reading the lens applies the model's own final norm and unembedding:
20:math:`\\text{lens}(h_\\ell) = W_U\\,\\mathrm{norm}(J_\\ell h_\\ell)`.
21The logit lens is the special case :math:`J_\\ell = I`. The rows of
22:math:`W_U J_\\ell` ("J-lens vectors") are residual-stream directions associated
23with single vocabulary tokens, and support causal interventions: steering,
24ablation, and exchanging one concept for another via a pseudoinverse coordinate
25swap.
27Introduced in `Verbalizable Representations Form a Global Workspace in Language
28Models <https://transformer-circuits.pub/2026/workspace/index.html>`_ (Gurnee et
29al., Transformer Circuits Thread, 2026). The fitting estimator and the artifact
30format follow Anthropic's Apache-2.0 reference implementation
31(`anthropics/jacobian-lens <https://github.com/anthropics/jacobian-lens>`_), so
32lenses fitted here interoperate with artifacts published on the Hugging Face Hub
33(e.g. `neuronpedia/jacobian-lens
34<https://huggingface.co/neuronpedia/jacobian-lens>`_); the interventions are
35implemented from the paper's Methods section.
37Warning:
38 Published lens artifacts are fitted on **raw** HuggingFace activations.
39 Jacobian lens supports only a freshly booted
40 ``TransformerBridge.boot_transformers`` model, whose weights are raw by
41 default. Compatibility mode and direct ``process_weights`` calls change the
42 residual basis and are refused rather than returning silently wrong
43 readouts. The model must also be a causal decoder whose adapter supports
44 text generation, use single-stream block outputs, and expose the standard
45 direct ``ln_final -> d_model-width unembed`` output path.
47Example::
49 import torch
50 from transformer_lens.model_bridge import TransformerBridge
51 from transformer_lens.tools.analysis import JacobianLens
53 model = TransformerBridge.boot_transformers("gpt2", device="cpu")
54 lens = JacobianLens.from_pretrained(
55 "neuronpedia/jacobian-lens",
56 filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
57 model=model,
58 )
59 result = lens.readout(model, "The Eiffel Tower is in the city of")
60 print(result.top_tokens(model.tokenizer, k=5)[8][-1]) # layer 8, final position
61"""
63import math
64import warnings
65from dataclasses import dataclass
66from importlib.metadata import version
67from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
69import torch
70from jaxtyping import Float, Int
71from tqdm.auto import tqdm
73from transformer_lens.utilities.hf_utils import call_hf_with_retry
75TokenInput = Union[str, int]
77# ---------------------------------------------------------------------------
78# Registry helpers
79# ---------------------------------------------------------------------------
81_REGISTRY_CACHE: Optional[Dict[str, Any]] = None
84def _load_registry() -> Dict[str, Any]:
85 """Return the bundled artifact registry, loading it once on first call."""
86 global _REGISTRY_CACHE
87 if _REGISTRY_CACHE is None:
88 import json
89 import pathlib
91 registry_path = pathlib.Path(__file__).with_name("jacobian_lens_registry.json")
92 with registry_path.open() as fh:
93 _REGISTRY_CACHE = json.load(fh)
94 return _REGISTRY_CACHE
97def _resolve_registry_entry(name_or_path: str) -> Optional[Tuple[str, str]]:
98 """Return ``(repo_id, filename)`` if *name_or_path* matches the registry.
100 Matching is tried in two passes:
101 1. Direct key match against short model names (e.g. ``"gemma-2-2b"``).
102 2. Alias match against Hugging Face model IDs (e.g. ``"google/gemma-2-2b"``).
104 Returns ``None`` when there is no match so callers can fall through to the
105 generic Hub download path.
106 """
107 registry = _load_registry()
108 if name_or_path in registry:
109 entry = registry[name_or_path]
110 return entry["repo_id"], entry["filename"]
111 for entry in registry.values():
112 if isinstance(entry, dict) and name_or_path in entry.get("aliases", []):
113 return entry["repo_id"], entry["filename"]
114 return None
117# Fitting excludes early positions (attention sinks with atypical residual statistics)
118# and the final position (no next-token target), matching the reference implementation.
119DEFAULT_SKIP_FIRST_POSITIONS = 16
120DEFAULT_TOP_K = 10
121_SWAP_WARN_COSINE = 0.99
122_SWAP_ERROR_COSINE = 0.999
124# Keys written by fit() that must not appear in converted-lens metadata so that
125# merge() can refuse to mix TL-fitted lenses with externally converted ones.
126# Note: "target_layer" is intentionally NOT listed here — it must survive
127# conversion so that validate_model() can detect and refuse checkpoints that
128# were fitted against a non-final target layer.
129_FIT_RESERVED_KEYS: frozenset = frozenset(
130 {
131 "transformer_lens_fit",
132 "transformer_lens_version",
133 "model_system",
134 "processing",
135 "hook_convention",
136 "fit_dtype",
137 "dim_batch",
138 "max_seq_len",
139 "skip_first_positions",
140 }
141)
143# Top-level payload keys that some checkpoint writers store as flat provenance
144# (rather than nested under a "metadata" key).
145_CHECKPOINT_FLAT_PROVENANCE: frozenset = frozenset({"model_name", "model_revision", "corpus"})
148@dataclass
149class JacobianLensReadout:
150 """Result of a :meth:`JacobianLens.readout` call.
152 Attributes:
153 lens_topk_values:
154 Per-layer retained top-k pre-softmax values, on CPU.
155 lens_topk_indices:
156 Per-layer retained top-k vocabulary ids, on CPU.
157 model_topk_values:
158 The model output's retained top-k pre-softmax values, on CPU.
159 model_topk_indices:
160 The model output's retained top-k vocabulary ids, on CPU.
161 lens_logits:
162 Optional full per-layer logits, on CPU. Present only when
163 ``readout(return_full_logits=True)`` was requested.
164 model_logits:
165 Optional full model logits, on CPU. Present only when
166 ``readout(return_full_logits=True)`` was requested.
167 tokens:
168 The token ids of the run prompt, ``[seq]``.
169 positions:
170 The (normalized, non-negative) positions the readout covers, aligned
171 with the ``pos`` axis of retained top-k and optional full logits.
172 use_jacobian:
173 Whether the Jacobian transport was applied (``False`` = logit lens).
174 """
176 lens_topk_values: Dict[int, Float[torch.Tensor, "pos k"]]
177 lens_topk_indices: Dict[int, Int[torch.Tensor, "pos k"]]
178 model_topk_values: Float[torch.Tensor, "pos k"]
179 model_topk_indices: Int[torch.Tensor, "pos k"]
180 tokens: Int[torch.Tensor, "seq"]
181 positions: List[int]
182 use_jacobian: bool = True
183 lens_logits: Optional[Dict[int, Float[torch.Tensor, "pos d_vocab"]]] = None
184 model_logits: Optional[Float[torch.Tensor, "pos d_vocab"]] = None
186 def top_tokens(self, tokenizer: Any, k: int = 5) -> Dict[int, List[List[str]]]:
187 """Decode the top-``k`` tokens per layer and position.
189 Args:
190 tokenizer: The model's tokenizer (``model.tokenizer``).
191 k: Number of top tokens to decode per (layer, position).
193 Returns:
194 ``{layer: [ [top-k strings] per position ]}``, positions aligned with
195 :attr:`positions`.
196 """
197 out: Dict[int, List[List[str]]] = {}
198 retained = next(iter(self.lens_topk_indices.values())).shape[-1]
199 if not 1 <= k <= retained:
200 raise ValueError(f"k must be between 1 and the retained top_k={retained}, got {k}")
201 for layer, ids in self.lens_topk_indices.items():
202 out[layer] = [[tokenizer.decode([t]) for t in row[:k].tolist()] for row in ids]
203 return out
206class JacobianLens:
207 """A fitted Jacobian lens: one transport matrix per source layer.
209 Layer convention (matching the reference implementation and the published
210 artifacts): index ``l`` refers to the **output of block** ``l`` at the
211 Bridge-native hook ``blocks.{l}.hook_out``. ``J[l]`` maps that activation
212 to the final block's output, pre final norm. The final layer itself is never
213 fitted (its transport is the identity), so
214 ``source_layers == [0, ..., n_layers - 2]`` for a full fit.
216 Attributes:
217 jacobians: ``{layer: [d_model, d_model]}`` transport matrices, fp32, CPU.
218 n_prompts: Number of prompts averaged into the fit.
219 d_model: Residual stream width the lens was fitted for.
220 metadata: Optional provenance (model name, TransformerLens version, fit
221 hyperparameters). Preserved by :meth:`save`/:meth:`load`; artifacts
222 from the reference implementation load with empty metadata.
223 """
225 def __init__(
226 self,
227 jacobians: Dict[int, Float[torch.Tensor, "d_model d_model"]],
228 *,
229 n_prompts: int,
230 d_model: int,
231 metadata: Optional[Dict[str, Any]] = None,
232 ) -> None:
233 if not jacobians: 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true
234 raise ValueError("jacobians must contain at least one layer")
235 for layer, matrix in jacobians.items():
236 if matrix.shape != (d_model, d_model): 236 ↛ 237line 236 didn't jump to line 237 because the condition on line 236 was never true
237 raise ValueError(
238 f"jacobians[{layer}] has shape {tuple(matrix.shape)}, "
239 f"expected ({d_model}, {d_model})"
240 )
241 self.jacobians: Dict[int, torch.Tensor] = {
242 int(layer): matrix.detach().float().cpu() for layer, matrix in jacobians.items()
243 }
244 self.n_prompts = int(n_prompts)
245 self.d_model = int(d_model)
246 self.metadata: Dict[str, Any] = dict(metadata or {})
247 self._device_jacobians: Dict[Tuple[int, torch.device], torch.Tensor] = {}
249 @property
250 def source_layers(self) -> List[int]:
251 """Sorted list of layers this lens has transport matrices for."""
252 return sorted(self.jacobians)
254 def __repr__(self) -> str:
255 layers = self.source_layers
256 return (
257 f"JacobianLens(layers={layers[0]}..{layers[-1]} ({len(layers)}), "
258 f"d_model={self.d_model}, n_prompts={self.n_prompts})"
259 )
261 # ------------------------------------------------------------------ #
262 # persistence #
263 # ------------------------------------------------------------------ #
265 def save(self, path: str, *, dtype: torch.dtype = torch.float16) -> None:
266 """Save the lens in the reference implementation's artifact format.
268 The four official keys (``J``, ``n_prompts``, ``source_layers``,
269 ``d_model``) are written unchanged so the file stays loadable by the
270 reference package; TransformerLens provenance is stored under an
271 additive ``metadata`` key.
273 Args:
274 path: Destination ``.pt`` path.
275 dtype: Storage dtype. Defaults to fp16 like the reference
276 implementation — Jacobian entries are order-one, so the smaller
277 dtype costs little precision and halves the artifact on disk.
278 """
279 _validate_metadata(self.metadata)
280 payload: Dict[str, Any] = {
281 "J": {layer: matrix.to(dtype) for layer, matrix in self.jacobians.items()},
282 "n_prompts": self.n_prompts,
283 "source_layers": self.source_layers,
284 "d_model": self.d_model,
285 }
286 if self.metadata:
287 payload["metadata"] = self.metadata
288 torch.save(payload, path)
290 @classmethod
291 def load(cls, path: str) -> "JacobianLens":
292 """Load a lens artifact or fit checkpoint saved in a supported schema.
294 Two file schemas are accepted:
296 **Artifact** (the reference format, written by :meth:`save` or the
297 Anthropic reference package): must contain a ``J`` key mapping layer
298 indices to transport matrices, plus ``n_prompts``, ``d_model``, and an
299 optional ``metadata`` dict.
301 **Fit checkpoint** (running-sum format, written by the reference
302 implementation's ``write_checkpoint()`` during fitting): must contain
303 a ``jacobian_sum`` key mapping layer indices to *running-sum* matrices
304 (i.e. the sum over prompts, not yet divided by the prompt count), plus
305 ``n_done``. ``d_model`` is inferred from the first matrix's shape;
306 no explicit ``d_model`` key is required or expected. The per-layer
307 means are reconstructed on load. A
308 ``converted_from: "jacobian_lens_checkpoint"`` key is added to
309 metadata so :meth:`merge` refuses to silently combine checkpoints with
310 natively TL-fitted lenses. Fit-reserved provenance keys
311 (``transformer_lens_fit``, etc.) are stripped; scalar fields that can
312 be serialised under ``weights_only=True`` are preserved.
313 Tensor-valued metadata fields that would fail :func:`_validate_metadata`
314 are recorded by name in a ``dropped_fields`` list.
316 Fit checkpoint schema (reference ``write_checkpoint()`` format)
317 ---------------------------------------------------------------
318 The reference implementation writes exactly six top-level keys; all
319 other keys in the payload are ignored::
321 {
322 "jacobian_sum": {<layer int>: <float32 tensor [d, d]>, ...},
323 "n_done": <int>, # prompts accumulated into jacobian_sum
324 "next_idx": <int>, # next prompt index (informational)
325 "source_layers": [<int>, ...], # documented layer indices (informational)
326 "target_layer": <int>, # target layer — harvested into metadata
327 "skip_first": <int>, # leading positions skipped (informational)
328 # optional flat provenance accepted from alternative checkpoint writers:
329 "model_name": <str>,
330 "model_revision": <str>,
331 "corpus": <str>,
332 # optional nested provenance accepted from alternative writers:
333 "metadata": {<str>: <scalar/list/dict>, ...},
334 }
336 Args:
337 path: Path to the ``.pt`` file.
339 Raises:
340 ValueError: If the file lacks both a ``J`` key (artifact) and a
341 ``jacobian_sum`` key (checkpoint), or if a checkpoint records a
342 non-positive ``n_prompts``.
343 """
344 payload = torch.load(path, map_location="cpu", weights_only=True)
345 if "J" in payload:
346 return cls(
347 {int(layer): matrix for layer, matrix in payload["J"].items()},
348 n_prompts=int(payload.get("n_prompts", 0)),
349 d_model=int(payload["d_model"]),
350 metadata=payload.get("metadata"),
351 )
352 if "jacobian_sum" in payload:
353 return cls._from_checkpoint_payload(path, payload)
354 raise ValueError(
355 f"{path} does not look like a Jacobian lens artifact or fit checkpoint. "
356 "Expected a 'J' key (artifact) or 'jacobian_sum' key (fit checkpoint). "
357 "See JacobianLens.load() for the supported file schemas."
358 )
360 @classmethod
361 def _from_checkpoint_payload(cls, path: str, payload: Dict[str, Any]) -> "JacobianLens":
362 """Reconstruct a JacobianLens from a fit-checkpoint payload.
364 Divides the running Jacobian sums by ``n_prompts``, strips fit-reserved
365 provenance keys, harvests safe scalar metadata, records dropped tensor
366 fields, and marks the result as converted so :meth:`merge` refuses to
367 mix it with natively TL-fitted lenses.
368 """
369 n_prompts = int(payload.get("n_done", payload.get("n_prompts", 0)))
370 if n_prompts <= 0:
371 raise ValueError(
372 f"{path} is a fit checkpoint with n_prompts={n_prompts}; "
373 "a positive prompt count is required to reconstruct the Jacobian mean."
374 )
375 if not payload["jacobian_sum"]:
376 raise ValueError(
377 f"{path} is a fit checkpoint with an empty jacobian_sum; "
378 "at least one layer matrix is required to reconstruct d_model."
379 )
380 first_matrix = next(iter(payload["jacobian_sum"].values()))
381 d_model = first_matrix.shape[0]
382 jacobians = {
383 int(layer): matrix.float() / n_prompts
384 for layer, matrix in payload["jacobian_sum"].items()
385 }
387 # Collect raw provenance: first from nested "metadata", then supplement
388 # with flat top-level keys that some checkpoint writers place directly
389 # in the payload (model_name, model_revision, corpus).
390 raw_meta: Dict[str, Any] = dict(payload.get("metadata") or {})
391 for key in _CHECKPOINT_FLAT_PROVENANCE:
392 if key in payload and key not in raw_meta:
393 raw_meta[key] = payload[key]
394 # target_layer lives at the top level in the reference checkpoint format;
395 # harvest it into metadata so validate_model() can check the fitting target.
396 if "target_layer" in payload and "target_layer" not in raw_meta:
397 raw_meta["target_layer"] = payload["target_layer"]
399 # Build clean metadata: drop fit-reserved keys, record tensor-valued
400 # fields that _validate_metadata would reject (they cannot survive a
401 # weights_only=True reload), and keep everything else that validates.
402 dropped_fields: List[str] = []
403 clean_meta: Dict[str, Any] = {}
404 for key, value in raw_meta.items():
405 if key in _FIT_RESERVED_KEYS:
406 continue
407 if isinstance(value, torch.Tensor):
408 dropped_fields.append(f"{key}: shape={tuple(value.shape)} dtype={value.dtype}")
409 continue
410 try:
411 _validate_metadata({key: value})
412 clean_meta[key] = value
413 except ValueError:
414 dropped_fields.append(key)
416 clean_meta["converted_from"] = "jacobian_lens_checkpoint"
417 if dropped_fields:
418 clean_meta["dropped_fields"] = dropped_fields
420 return cls(jacobians, n_prompts=n_prompts, d_model=d_model, metadata=clean_meta)
422 @classmethod
423 def from_pretrained(
424 cls,
425 name_or_path: str,
426 *,
427 filename: str = "lens.pt",
428 revision: Optional[str] = None,
429 model: Any = None,
430 ) -> "JacobianLens":
431 """Load a lens from a local path, a short model name, or a Hub repo.
433 Resolution order
434 ----------------
435 1. **Local file** — if *name_or_path* is an existing ``.pt`` file,
436 load it directly.
437 2. **Local directory** — if *name_or_path* is a directory, load
438 ``<name_or_path>/<filename>``.
439 3. **Registry short name or HF model ID** — if *name_or_path* matches
440 a key or alias in the bundled ``jacobian_lens_registry.json`` (e.g.
441 ``"gemma-2-2b"`` or ``"google/gemma-2-2b"``), the corresponding
442 artifact in ``neuronpedia/jacobian-lens`` is fetched automatically.
443 The *filename* argument is ignored in this case because the registry
444 already encodes the correct subpath.
445 4. **Explicit Hub repo** — otherwise *name_or_path* is treated as a Hub
446 repo id and *filename* is used as-is, preserving full backward
447 compatibility (e.g. ``from_pretrained("neuronpedia/jacobian-lens",
448 filename="gpt2-small/jlens/...")``).
450 Args:
451 name_or_path: A local ``.pt`` file, a local directory, a short
452 model name such as ``"gemma-2-2b"`` or ``"llama3.1-8b"``, a
453 Hugging Face model ID such as ``"google/gemma-2-2b"``, or an
454 explicit Hub repo id paired with *filename*.
455 filename: File (or subpath) inside a local directory or an explicit
456 Hub repo. Ignored when *name_or_path* resolves via the
457 registry.
458 revision: Optional Hub revision (branch, tag, or commit) to pin.
459 When omitted, the Hub repository's mutable default branch is
460 followed; pin a commit hash for reproducible analyses.
461 model: If given, :meth:`validate_model` is called so dimension or
462 weight-processing mismatches fail here rather than at first use.
464 Returns:
465 The loaded (and, if ``model`` was given, validated) lens.
467 Examples::
469 # Short model name — no need to remember the HF subpath
470 lens = JacobianLens.from_pretrained("gemma-2-2b", model=model)
472 # HF model ID also works
473 lens = JacobianLens.from_pretrained("google/gemma-2-2b", model=model)
475 # Explicit Hub repo + subpath (backward-compatible)
476 lens = JacobianLens.from_pretrained(
477 "neuronpedia/jacobian-lens",
478 filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
479 model=model,
480 )
481 """
482 import os
484 if os.path.isfile(name_or_path): 484 ↛ 485line 484 didn't jump to line 485 because the condition on line 484 was never true
485 lens = cls.load(name_or_path)
486 elif os.path.isdir(name_or_path): 486 ↛ 487line 486 didn't jump to line 487 because the condition on line 486 was never true
487 lens = cls.load(os.path.join(name_or_path, filename))
488 else:
489 from huggingface_hub import hf_hub_download
491 resolved = _resolve_registry_entry(name_or_path)
492 if resolved is not None:
493 repo_id, resolved_filename = resolved
494 else:
495 repo_id, resolved_filename = name_or_path, filename
497 local_path = call_hf_with_retry(
498 hf_hub_download,
499 repo_id=repo_id,
500 filename=resolved_filename,
501 revision=revision,
502 )
503 lens = cls.load(local_path)
504 if model is not None:
505 lens.validate_model(model)
506 return lens
508 @classmethod
509 def merge(cls, lenses: Sequence["JacobianLens"]) -> "JacobianLens":
510 """Combine lenses fitted on disjoint prompt slices.
512 The per-layer matrices are averaged weighted by each lens's
513 ``n_prompts``, matching the reference implementation, so fitting can be
514 parallelized across processes or machines and merged afterwards.
515 Provenance must match across shards (apart from ``n_prompts``), so a
516 merge cannot silently relabel matrices fitted with different models,
517 corpora, dtypes, or estimator settings. The merged count replaces the
518 per-shard count.
520 Args:
521 lenses: Lenses that agree exactly on ``source_layers`` and
522 ``d_model``.
524 Raises:
525 ValueError: On an empty sequence or mismatched lenses.
526 """
527 if not lenses:
528 raise ValueError("cannot merge an empty sequence of lenses")
529 invalid_counts = [
530 (index, lens.n_prompts) for index, lens in enumerate(lenses) if lens.n_prompts <= 0
531 ]
532 if invalid_counts:
533 raise ValueError(
534 "every lens passed to merge() must have positive n_prompts; "
535 f"invalid shards: {invalid_counts}"
536 )
537 first = lenses[0]
538 for lens in lenses:
539 _validate_metadata(lens.metadata)
540 first_provenance = {
541 key: value for key, value in first.metadata.items() if key != "n_prompts"
542 }
543 for other in lenses[1:]:
544 if other.source_layers != first.source_layers or other.d_model != first.d_model:
545 raise ValueError(
546 "all lenses being merged must share the same source_layers and d_model"
547 )
548 other_provenance = {
549 key: value for key, value in other.metadata.items() if key != "n_prompts"
550 }
551 if other_provenance != first_provenance:
552 raise ValueError(
553 "all lenses being merged must share the same provenance metadata "
554 "apart from n_prompts"
555 )
556 total = sum(lens.n_prompts for lens in lenses)
557 merged = {
558 layer: torch.stack([lens.jacobians[layer] * lens.n_prompts for lens in lenses]).sum(
559 dim=0
560 )
561 / total
562 for layer in first.source_layers
563 }
564 metadata = dict(first.metadata)
565 if metadata:
566 metadata["n_prompts"] = total
567 return cls(merged, n_prompts=total, d_model=first.d_model, metadata=metadata)
569 # ------------------------------------------------------------------ #
570 # model validation #
571 # ------------------------------------------------------------------ #
573 def validate_model(self, model: Any) -> "JacobianLens":
574 """Check that ``model`` matches this lens; raise loudly if not.
576 Requires a raw causal ``TransformerBridge`` with the standard direct
577 final-norm/unembed path, verifies recorded model provenance, residual
578 width and layer range, and enforces the published final-block target
579 convention.
581 Args:
582 model: A raw ``TransformerBridge``.
584 Returns:
585 ``self``, for chaining.
587 Raises:
588 TypeError: If model is not a ``TransformerBridge``.
589 ValueError: On model provenance or ``d_model`` mismatch,
590 out-of-range source layers, compatibility mode, unsupported
591 attention/output paths, or a non-final target convention.
592 """
593 _require_raw_bridge(model)
594 artifact_model_name = self.metadata.get("model_name")
595 current_model_name = getattr(model.cfg, "model_name", None)
596 if artifact_model_name is not None and artifact_model_name != current_model_name:
597 raise ValueError(
598 f"lens was fitted for model {artifact_model_name!r}, but the supplied "
599 f"model is {current_model_name!r}."
600 )
601 artifact_revision = self.metadata.get("model_revision")
602 current_revision = _get_model_revision(model)
603 if artifact_revision is not None and artifact_revision != current_revision:
604 raise ValueError(
605 f"lens was fitted for model revision {artifact_revision!r}, but the "
606 f"supplied model revision is {current_revision!r}."
607 )
608 d_model = model.cfg.d_model
609 if d_model != self.d_model:
610 raise ValueError(
611 f"lens was fitted for d_model={self.d_model}, but the model has "
612 f"d_model={d_model} — this lens belongs to a different model."
613 )
614 n_layers = model.cfg.n_layers
615 final_layer = n_layers - 1
616 out_of_range = [layer for layer in self.source_layers if not 0 <= layer < final_layer]
617 if out_of_range:
618 raise ValueError(
619 f"lens has source layers {out_of_range} outside the model's "
620 f"0..{final_layer - 1} source range — this lens belongs to a different model."
621 )
622 target_layer = int(self.metadata.get("target_layer", final_layer))
623 if target_layer != final_layer:
624 raise ValueError(
625 f"lens targets layer {target_layer}, but readout supports only the "
626 f"published final-layer convention ({final_layer}); refit without "
627 "a custom target layer."
628 )
629 return self
631 # ------------------------------------------------------------------ #
632 # reading #
633 # ------------------------------------------------------------------ #
635 def clear_device_cache(self) -> None:
636 """Release lazily cached Jacobian copies on accelerator devices."""
637 self._device_jacobians.clear()
639 def _matrix_on(self, layer: int, device: Union[str, torch.device]) -> torch.Tensor:
640 """Return one cached fp32 Jacobian copy for a layer/device pair."""
641 if layer not in self.jacobians: 641 ↛ 642line 641 didn't jump to line 642 because the condition on line 641 was never true
642 raise ValueError(
643 f"layer {layer} is not in this lens's source layers "
644 f"({self.source_layers[0]}..{self.source_layers[-1]})"
645 )
646 resolved_device = torch.device(device)
647 key = (layer, resolved_device)
648 matrix = self._device_jacobians.get(key)
649 if matrix is None:
650 matrix = self.jacobians[layer].to(device=resolved_device, dtype=torch.float32)
651 self._device_jacobians[key] = matrix
652 return matrix
654 def transport(
655 self,
656 residual: Float[torch.Tensor, "... d_model"],
657 layer: int,
658 ) -> Float[torch.Tensor, "... d_model"]:
659 """Map layer-``layer`` activations into the final block's output basis.
661 Computes ``J[layer] @ h`` per activation vector, in fp32.
663 Args:
664 residual: Activations from the output of block ``layer``.
665 layer: Source layer index.
666 """
667 matrix = self._matrix_on(layer, residual.device)
668 return residual.float() @ matrix.T
670 @torch.no_grad()
671 def readout(
672 self,
673 model: Any,
674 input: Union[str, Int[torch.Tensor, "batch seq"]],
675 *,
676 layers: Optional[Sequence[int]] = None,
677 positions: Optional[Sequence[int]] = None,
678 use_jacobian: bool = True,
679 top_k: int = DEFAULT_TOP_K,
680 return_full_logits: bool = False,
681 ) -> JacobianLensReadout:
682 """Read per-layer vocabulary logits for a prompt.
684 Runs the model once with caching, transports the residual stream at each
685 requested layer through ``J[layer]`` (or the identity when
686 ``use_jacobian=False`` — the logit lens), and applies the model's own
687 final norm, unembedding, architecture logit scaling, and logit soft cap.
689 Args:
690 model: A raw ``TransformerBridge``.
691 input: A prompt string, or a ``[1, seq]`` token tensor.
692 layers: Layers to read. Defaults to every fitted layer plus the
693 final layer. The final layer (``n_layers - 1``) is always read
694 with the identity transport — by construction its lens equals
695 the model's own output distribution.
696 positions: Token positions to read (negative indices allowed).
697 Defaults to all positions.
698 use_jacobian: Apply the Jacobian transport. ``False`` gives the
699 logit-lens baseline through the identical code path.
700 top_k: Number of values and vocabulary ids retained per layer and
701 position. Defaults to 10.
702 return_full_logits: Also retain full vocabulary tensors on CPU.
703 This is opt-in because a 64-token Gemma readout across all
704 layers is roughly 1.7 GB.
706 Returns:
707 A :class:`JacobianLensReadout`.
709 Raises:
710 ValueError: If the model fails :meth:`validate_model`, ``input`` is
711 batched, ``top_k`` is invalid, or a requested layer has no
712 transport matrix.
713 """
714 self.validate_model(model)
715 tokens = model.to_tokens(input) if isinstance(input, str) else input
716 if tokens.ndim != 2 or tokens.shape[0] != 1: 716 ↛ 717line 716 didn't jump to line 717 because the condition on line 716 was never true
717 raise ValueError(f"readout expects a single prompt; got shape {tuple(tokens.shape)}")
718 n_layers = model.cfg.n_layers
719 final_layer = n_layers - 1
720 if layers is None:
721 layers = self.source_layers + [final_layer]
722 layers = [_normalize_layer(layer, n_layers) for layer in layers]
723 for layer in layers:
724 if use_jacobian and layer != final_layer and layer not in self.jacobians: 724 ↛ 725line 724 didn't jump to line 725 because the condition on line 724 was never true
725 raise ValueError(
726 f"layer {layer} is not in this lens's source layers; "
727 f"available: {self.source_layers} (+{final_layer} as identity)"
728 )
730 if top_k < 1: 730 ↛ 731line 730 didn't jump to line 731 because the condition on line 730 was never true
731 raise ValueError(f"top_k must be at least 1, got {top_k}")
732 seq_len = tokens.shape[1]
733 norm_positions = _normalize_positions(positions, seq_len)
735 hook_names = {
736 layer: _resid_post_hook_name(layer) for layer in layers if layer != final_layer
737 }
738 wanted = set(hook_names.values())
739 logits, cache = model.run_with_cache(tokens, names_filter=lambda name: name in wanted)
740 selected_model_logits = logits[0, norm_positions, :].float()
741 if top_k > selected_model_logits.shape[-1]: 741 ↛ 742line 741 didn't jump to line 742 because the condition on line 741 was never true
742 raise ValueError(
743 f"top_k={top_k} exceeds the model vocabulary size "
744 f"{selected_model_logits.shape[-1]}"
745 )
746 model_topk = selected_model_logits.topk(top_k, dim=-1)
747 full_model_logits = selected_model_logits.cpu() if return_full_logits else None
748 lens_topk_values: Dict[int, torch.Tensor] = {}
749 lens_topk_indices: Dict[int, torch.Tensor] = {}
750 full_lens_logits: Optional[Dict[int, torch.Tensor]] = {} if return_full_logits else None
751 for layer in layers:
752 if layer == final_layer:
753 layer_logits = selected_model_logits
754 layer_topk = model_topk
755 else:
756 activation = cache[hook_names[layer]]
757 _validate_residual_activation(
758 activation,
759 d_model=model.cfg.d_model,
760 hook_name=hook_names[layer],
761 )
762 residual = activation[0, norm_positions, :]
763 transported = self.transport(residual, layer) if use_jacobian else residual.float()
764 layer_logits = _unembed(model, transported)
765 layer_topk = layer_logits.topk(top_k, dim=-1)
766 lens_topk_values[layer] = layer_topk.values.cpu()
767 lens_topk_indices[layer] = layer_topk.indices.cpu()
768 if full_lens_logits is not None:
769 if layer == final_layer:
770 assert full_model_logits is not None
771 full_lens_logits[layer] = full_model_logits
772 else:
773 full_lens_logits[layer] = layer_logits.cpu()
774 return JacobianLensReadout(
775 lens_topk_values=lens_topk_values,
776 lens_topk_indices=lens_topk_indices,
777 model_topk_values=model_topk.values.cpu(),
778 model_topk_indices=model_topk.indices.cpu(),
779 tokens=tokens[0].cpu(),
780 positions=norm_positions,
781 use_jacobian=use_jacobian,
782 lens_logits=full_lens_logits,
783 model_logits=full_model_logits,
784 )
786 @torch.no_grad()
787 def lens_vectors(
788 self,
789 model: Any,
790 tokens: Union[TokenInput, Sequence[TokenInput]],
791 layer: int,
792 ) -> Float[torch.Tensor, "n d_model"]:
793 """Residual-stream directions for vocabulary tokens at a layer.
795 The J-lens vector for token ``t`` is row ``t`` of ``W_U J[layer]``
796 expressed in layer-``layer`` residual coordinates:
797 ``v_t = J[layer]^T W_U[:, t]``.
799 Args:
800 model: The model supplying ``W_U``.
801 tokens: A token string / id, or a sequence of them. Strings must
802 encode to a single token.
803 layer: Source layer for the vectors.
805 Returns:
806 One vector per token, fp32, on the model's device.
807 """
808 self.validate_model(model)
809 layer = _normalize_layer(layer, model.cfg.n_layers)
810 token_ids = _to_token_ids(model, tokens)
811 unembed_columns = model.W_U[:, token_ids].float() # [d_model, n]
812 matrix = self._matrix_on(layer, unembed_columns.device)
813 return (matrix.T @ unembed_columns).T
815 # ------------------------------------------------------------------ #
816 # interventions #
817 # ------------------------------------------------------------------ #
819 def steering_hooks(
820 self,
821 model: Any,
822 token: TokenInput,
823 layers: Sequence[int],
824 *,
825 alpha: float = 4.0,
826 positions: Optional[Sequence[int]] = None,
827 ) -> List[Tuple[str, Any]]:
828 """Hooks that steer the residual stream along a token's J-lens vector.
830 At each layer the unit-normalized lens vector is added, scaled by
831 ``alpha`` times the activation's **median** per-position residual norm:
832 ``h <- h + alpha * median||h|| * v̂``. This norm-matched
833 parameterization follows the steering description in the reference
834 implementation's experiment protocols; the paper's minimal form is
835 the unscaled ``h <- h + alpha * v_t``, recoverable by passing the raw
836 :meth:`lens_vectors` output to your own hook. The median (not mean) is
837 used so attention-sink positions — whose residual norms run orders of
838 magnitude above typical positions — do not inflate the scale.
840 Args:
841 model: The model the hooks will run on.
842 token: The concept token to steer toward.
843 layers: Layers to intervene at.
844 alpha: Steering strength scalar; ``0`` disables. Because of the
845 norm-matched scale, values of order 1 already perturb the
846 stream by roughly its own magnitude.
847 positions: Chunk-local positions to steer (negative indices allowed
848 and normalized on every hook invocation). Defaults to all.
850 Returns:
851 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)`` or
852 ``model.run_with_hooks(fwd_hooks=...)``.
853 """
854 self.validate_model(model)
855 hooks = []
856 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
857 direction = self.lens_vectors(model, token, layer)[0]
858 unit = _unit_rows(direction.unsqueeze(0), layer=layer)[0]
859 device_units: Dict[torch.device, torch.Tensor] = {}
861 def transform(
862 selected: Float[torch.Tensor, "batch pos d_model"],
863 unit: torch.Tensor = unit,
864 device_units: Dict[torch.device, torch.Tensor] = device_units,
865 ) -> Float[torch.Tensor, "batch pos d_model"]:
866 local_unit = _cached_on_device(unit, device_units, selected.device)
867 scale = alpha * selected.float().norm(dim=-1).median()
868 return selected.float() + scale * local_unit
870 hooks.append(
871 (
872 _resid_post_hook_name(layer),
873 _make_intervention_hook(transform, positions, model.cfg.d_model),
874 )
875 )
876 return hooks
878 def ablation_hooks(
879 self,
880 model: Any,
881 tokens: Union[TokenInput, Sequence[TokenInput]],
882 layers: Sequence[int],
883 *,
884 positions: Optional[Sequence[int]] = None,
885 ) -> List[Tuple[str, Any]]:
886 """Hooks that project token directions out of the residual stream.
888 For each token's unit lens vector ``v̂``: ``h <- h - (h·v̂) v̂``,
889 applied sequentially when several tokens are given.
891 Args:
892 model: The model the hooks will run on.
893 tokens: Concept token(s) to suppress.
894 layers: Layers to intervene at.
895 positions: Chunk-local positions to ablate (negative indices allowed
896 and normalized on every hook invocation). Defaults to all.
898 Returns:
899 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
900 """
901 self.validate_model(model)
902 hooks = []
903 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
904 vectors = self.lens_vectors(model, tokens, layer)
905 units = _unit_rows(vectors, layer=layer)
906 device_units: Dict[torch.device, torch.Tensor] = {}
908 def transform(
909 selected: Float[torch.Tensor, "batch pos d_model"],
910 units: torch.Tensor = units,
911 device_units: Dict[torch.device, torch.Tensor] = device_units,
912 ) -> Float[torch.Tensor, "batch pos d_model"]:
913 local_units = _cached_on_device(units, device_units, selected.device)
914 result = selected.float()
915 for unit in local_units:
916 coeff = result @ unit
917 result = result - coeff.unsqueeze(-1) * unit
918 return result
920 hooks.append(
921 (
922 _resid_post_hook_name(layer),
923 _make_intervention_hook(transform, positions, model.cfg.d_model),
924 )
925 )
926 return hooks
928 def swap_hooks(
929 self,
930 model: Any,
931 source_token: TokenInput,
932 target_token: TokenInput,
933 layers: Sequence[int],
934 *,
935 alpha: float = 1.0,
936 positions: Optional[Sequence[int]] = None,
937 ) -> List[Tuple[str, Any]]:
938 """Hooks that swap two concepts' coordinates in lens space.
940 The paper's patching-in-lens-coordinates intervention: with
941 ``V = [v_s, v_t]`` and lens coordinates ``c = V⁺ h`` (pseudoinverse),
942 the update is ``h <- h + alpha * V (sigma(c) - c)`` where ``sigma``
943 exchanges the two coordinates. The component of ``h`` orthogonal to
944 ``span{v_s, v_t}`` is untouched. ``alpha=2`` is the paper's
945 "double-strength" swap.
947 Args:
948 model: The model the hooks will run on.
949 source_token: The concept to remove (e.g. ``" France"``).
950 target_token: The concept to install (e.g. ``" China"``).
951 layers: Layers to intervene at (the paper clamps the swap across an
952 intermediate-layer band).
953 alpha: Swap strength.
954 positions: Chunk-local positions to swap (negative indices allowed
955 and normalized on every hook invocation). Defaults to all.
957 Returns:
958 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
959 """
960 self.validate_model(model)
961 source_id, target_id = _to_token_ids(model, [source_token, target_token])
962 if source_id == target_id:
963 raise ValueError(
964 "source_token and target_token resolve to the same token id; "
965 "a coordinate swap would be a silent no-op"
966 )
968 hooks = []
969 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
970 vectors = self.lens_vectors(model, [source_id, target_id], layer)
971 units = _unit_rows(vectors, layer=layer)
972 cosine = abs(float((units[0] @ units[1]).item()))
973 if not math.isfinite(cosine) or cosine >= _SWAP_ERROR_COSINE:
974 raise ValueError(
975 f"swap vectors at layer {layer} are numerically near-parallel "
976 f"(abs cosine={cosine:.6f}); choose better-separated concepts"
977 )
978 if cosine >= _SWAP_WARN_COSINE:
979 warnings.warn(
980 f"swap vectors at layer {layer} are poorly conditioned "
981 f"(abs cosine={cosine:.6f}); the intervention may be amplified",
982 UserWarning,
983 stacklevel=2,
984 )
985 basis = vectors.T # [d, 2]
986 pinv = torch.linalg.pinv(basis) # [2, d]
987 device_basis: Dict[torch.device, torch.Tensor] = {}
988 device_pinv: Dict[torch.device, torch.Tensor] = {}
990 def transform(
991 selected: Float[torch.Tensor, "batch pos d_model"],
992 basis: torch.Tensor = basis,
993 pinv: torch.Tensor = pinv,
994 device_basis: Dict[torch.device, torch.Tensor] = device_basis,
995 device_pinv: Dict[torch.device, torch.Tensor] = device_pinv,
996 ) -> Float[torch.Tensor, "batch pos d_model"]:
997 local_basis = _cached_on_device(basis, device_basis, selected.device)
998 local_pinv = _cached_on_device(pinv, device_pinv, selected.device)
999 coords = selected.float() @ local_pinv.T # [..., 2]
1000 delta = alpha * ((coords[..., [1, 0]] - coords) @ local_basis.T)
1001 return selected.float() + delta
1003 hooks.append(
1004 (
1005 _resid_post_hook_name(layer),
1006 _make_intervention_hook(transform, positions, model.cfg.d_model),
1007 )
1008 )
1009 return hooks
1011 # ------------------------------------------------------------------ #
1012 # fitting #
1013 # ------------------------------------------------------------------ #
1015 @classmethod
1016 def fit(
1017 cls,
1018 model: Any,
1019 prompts: Sequence[str],
1020 *,
1021 corpus: str,
1022 source_layers: Optional[Sequence[int]] = None,
1023 dim_batch: int = 8,
1024 max_seq_len: int = 128,
1025 skip_first_positions: int = DEFAULT_SKIP_FIRST_POSITIONS,
1026 show_progress: bool = True,
1027 metadata: Optional[Dict[str, Any]] = None,
1028 ) -> "JacobianLens":
1029 """Fit a Jacobian lens on a hooked model.
1031 Implements the reference estimator exactly. For each prompt: one forward
1032 pass (the prompt replicated ``dim_batch`` times along the batch axis),
1033 then ``ceil(d_model / dim_batch)`` backward passes. Each backward plants
1034 a one-hot cotangent for one output dimension at *every* valid target
1035 position simultaneously — causal attention guarantees the gradient at
1036 source position ``t`` is then the sum over target positions
1037 ``t' >= t`` with no explicit masking. Rows are averaged over valid
1038 source positions (the first ``skip_first_positions`` and the final
1039 position are excluded), and prompts contribute equally to the final
1040 mean. There is no randomness: the computation is deterministic given
1041 the prompts.
1043 The reference implementation reports that fit quality saturates
1044 quickly — on the order of 100 prompts of 128 tokens is usable; the
1045 published lenses use up to 1000. Use :meth:`merge` to parallelize
1046 across prompt slices.
1048 Args:
1049 model: A raw ``TransformerBridge``. Model parameters are temporarily
1050 frozen (``requires_grad=False``) during fitting and restored
1051 after. Cotangents and activation gradients use the model dtype;
1052 fit with a float32 model for the highest-fidelity estimator.
1053 prompts: Prompt strings. Prompts too short to contain a valid
1054 position (``seq_len <= skip_first_positions + 1``) are skipped
1055 with a warning and do not count toward ``n_prompts``.
1056 corpus: Stable identifier for the prompt corpus or slice, recorded
1057 in artifact provenance.
1058 source_layers: Layers to fit. Defaults to every layer below
1059 the final layer. Negative indices count from ``n_layers``.
1060 dim_batch: Output dimensions per backward pass. Higher is faster
1061 but replicates the prompt ``dim_batch`` times in memory; total
1062 backward FLOPs are unchanged.
1063 max_seq_len: Prompts are truncated to this many tokens.
1064 skip_first_positions: Leading positions excluded from the source
1065 average.
1066 show_progress: Show a tqdm progress bar over prompts.
1067 metadata: Extra provenance merged into :attr:`metadata`.
1069 Returns:
1070 The fitted :class:`JacobianLens`.
1072 Raises:
1073 TypeError: If model is not a ``TransformerBridge``.
1074 ValueError: On compatibility mode, invalid provenance or layer
1075 indices, or if no prompt was long enough to fit on.
1076 """
1077 _require_raw_bridge(model)
1078 if not isinstance(corpus, str) or not corpus.strip(): 1078 ↛ 1079line 1078 didn't jump to line 1079 because the condition on line 1078 was never true
1079 raise ValueError("corpus must be a non-empty provenance identifier")
1080 n_layers = model.cfg.n_layers
1081 d_model = model.cfg.d_model
1082 resolved_target = n_layers - 1
1083 if source_layers is None:
1084 resolved_sources = list(range(resolved_target))
1085 else:
1086 resolved_sources = sorted(
1087 {_normalize_layer(layer, n_layers) for layer in source_layers}
1088 )
1089 if not resolved_sources: 1089 ↛ 1090line 1089 didn't jump to line 1090 because the condition on line 1089 was never true
1090 raise ValueError("source_layers is empty")
1091 if resolved_sources[-1] >= resolved_target: 1091 ↛ 1092line 1091 didn't jump to line 1092 because the condition on line 1091 was never true
1092 raise ValueError(
1093 f"every source layer must be below target_layer={resolved_target}; "
1094 f"got {resolved_sources}"
1095 )
1096 if dim_batch < 1: 1096 ↛ 1097line 1096 didn't jump to line 1097 because the condition on line 1096 was never true
1097 raise ValueError(f"dim_batch must be >= 1, got {dim_batch}")
1098 if skip_first_positions < 0:
1099 raise ValueError(f"skip_first_positions must be >= 0, got {skip_first_positions}")
1100 fit_dtype = model.W_U.dtype
1101 if fit_dtype in (torch.float16, torch.bfloat16):
1102 warnings.warn(
1103 f"fitting in {fit_dtype} accumulates Jacobian gradients at reduced "
1104 "precision; use a float32 TransformerBridge for the highest-fidelity fit",
1105 UserWarning,
1106 stacklevel=2,
1107 )
1109 jacobian_sum = {
1110 layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in resolved_sources
1111 }
1112 n_done = 0
1113 iterator = tqdm(prompts, desc="fitting J-lens", disable=not show_progress)
1114 with _frozen_parameters(model):
1115 for prompt in iterator:
1116 tokens = model.to_tokens(prompt)[:, :max_seq_len]
1117 seq_len = tokens.shape[1]
1118 if seq_len <= skip_first_positions + 1:
1119 warnings.warn(
1120 f"skipping prompt with only {seq_len} tokens "
1121 f"(need > {skip_first_positions + 1})",
1122 stacklevel=2,
1123 )
1124 continue
1125 per_prompt = _jacobian_for_prompt(
1126 model,
1127 tokens,
1128 source_layers=resolved_sources,
1129 dim_batch=dim_batch,
1130 skip_first_positions=skip_first_positions,
1131 )
1132 for layer in resolved_sources:
1133 jacobian_sum[layer] += per_prompt[layer]
1134 n_done += 1
1135 if n_done == 0:
1136 raise ValueError(
1137 "every prompt was too short to contribute valid positions; nothing was fitted"
1138 )
1140 fit_metadata: Dict[str, Any] = {
1141 "model_name": getattr(model.cfg, "model_name", None),
1142 "model_revision": _get_model_revision(model),
1143 "transformer_lens_version": version("transformer-lens"),
1144 "model_system": "TransformerBridge",
1145 "processing": {
1146 "compatibility_mode": False,
1147 "weight_basis": "raw_huggingface",
1148 },
1149 "hook_convention": "blocks.{layer}.hook_out",
1150 "corpus": corpus,
1151 "n_prompts": n_done,
1152 "fit_dtype": str(fit_dtype).removeprefix("torch."),
1153 "target_layer": resolved_target,
1154 "dim_batch": dim_batch,
1155 "max_seq_len": max_seq_len,
1156 "skip_first_positions": skip_first_positions,
1157 "transformer_lens_fit": True,
1158 }
1159 reserved = sorted(set(fit_metadata).intersection(metadata or {}))
1160 if reserved: 1160 ↛ 1161line 1160 didn't jump to line 1161 because the condition on line 1160 was never true
1161 raise ValueError(f"metadata cannot override fit provenance keys: {reserved}")
1162 full_metadata = dict(metadata or {})
1163 full_metadata.update(fit_metadata)
1164 _validate_metadata(full_metadata)
1165 return cls(
1166 {layer: jacobian_sum[layer] / n_done for layer in resolved_sources},
1167 n_prompts=n_done,
1168 d_model=d_model,
1169 metadata=full_metadata,
1170 )
1173# ---------------------------------------------------------------------- #
1174# helpers #
1175# ---------------------------------------------------------------------- #
1178def _resid_post_hook_name(layer: int) -> str:
1179 """Bridge-native hook for the output of block ``layer``."""
1180 return f"blocks.{layer}.hook_out"
1183def _get_model_revision(model: Any) -> Optional[str]:
1184 """Return the resolved Hugging Face commit recorded on a booted model."""
1185 original_model = getattr(model, "original_model", None)
1186 hf_config = getattr(original_model, "config", None)
1187 revision = getattr(hf_config, "_commit_hash", None)
1188 return revision if isinstance(revision, str) and revision else None
1191def _require_raw_bridge(model: Any) -> None:
1192 """Require the causal raw-Bridge contract used by fit and readout."""
1193 from transformer_lens.model_bridge import TransformerBridge
1195 if not isinstance(model, TransformerBridge):
1196 raise TypeError(
1197 "JacobianLens supports TransformerBridge only; load a fresh model with "
1198 "TransformerBridge.boot_transformers(...)."
1199 )
1200 if getattr(model, "compatibility_mode", False):
1201 raise ValueError(
1202 "compatibility mode is enabled on this TransformerBridge and changes "
1203 "the residual basis. Use a freshly booted "
1204 "TransformerBridge.boot_transformers(...) model with raw weights."
1205 )
1206 if getattr(model, "_weights_processed", False):
1207 raise ValueError(
1208 "process_weights was called on this TransformerBridge and changed the "
1209 "raw HuggingFace weight basis. Use a freshly booted "
1210 "TransformerBridge.boot_transformers(...) model."
1211 )
1212 adapter = model.adapter
1213 if not adapter.supports_generation:
1214 raise ValueError(
1215 "JacobianLens requires a causal decoder-only Bridge whose adapter "
1216 f"supports text generation; {type(adapter).__name__} declares "
1217 "supports_generation=False."
1218 )
1219 adapter.validate_output_logits_transform()
1220 attention_dir = getattr(model.cfg, "attention_dir", "causal")
1221 if attention_dir != "causal":
1222 raise ValueError(
1223 "JacobianLens requires causal attention because its estimator relies on "
1224 "causality to exclude target positions before each source position; "
1225 f"got attention_dir={attention_dir!r}."
1226 )
1227 total_ut_steps = int(getattr(model.cfg, "total_ut_steps", 1) or 1)
1228 if total_ut_steps != 1:
1229 raise ValueError(
1230 "JacobianLens requires each physical block hook to fire once per forward; "
1231 f"this looped-depth Bridge runs total_ut_steps={total_ut_steps}."
1232 )
1233 component_mapping = adapter.get_component_mapping()
1234 required_components = ("blocks", "ln_final", "unembed")
1235 missing_components = [
1236 component
1237 for component in required_components
1238 if component not in component_mapping or not hasattr(model, component)
1239 ]
1240 if missing_components:
1241 raise ValueError(
1242 "JacobianLens requires the standard direct ln_final -> unembed output path; "
1243 f"this Bridge is missing {missing_components}."
1244 )
1245 blocks_component = component_mapping["blocks"]
1246 if not getattr(blocks_component, "hook_out_is_single_residual_stream", False):
1247 raise ValueError(
1248 "JacobianLens requires single-stream [batch, position, d_model] block "
1249 f"outputs; {type(blocks_component).__name__} does not provide that contract."
1250 )
1251 if "project_out" in component_mapping:
1252 raise ValueError(
1253 "JacobianLens does not yet support a final output projection between "
1254 "the residual stream and unembedding."
1255 )
1256 unembed_width = model.W_U.shape[0]
1257 if unembed_width != model.cfg.d_model:
1258 raise ValueError(
1259 "JacobianLens requires a direct d_model-width unembedding after ln_final; "
1260 f"got W_U input width {unembed_width} for d_model={model.cfg.d_model}. "
1261 "Architectures with a final output projection are not yet supported."
1262 )
1265def _validate_metadata(metadata: Dict[str, Any]) -> None:
1266 """Reject values that ``torch.load(weights_only=True)`` cannot reload."""
1268 def validate(value: Any, path: str) -> None:
1269 if value is None or type(value) in (bool, int, float, str):
1270 return
1271 if type(value) in (list, tuple):
1272 for index, item in enumerate(value):
1273 validate(item, f"{path}[{index}]")
1274 return
1275 if type(value) is dict:
1276 for key, item in value.items():
1277 if type(key) is not str:
1278 raise ValueError(
1279 f"{path} has non-string key {key!r}; metadata keys must be strings"
1280 )
1281 validate(item, f"{path}.{key}")
1282 return
1283 raise ValueError(
1284 f"{path} has unsupported type {type(value).__name__}; use only "
1285 "None, bool, int, float, str, lists, tuples, and string-keyed dicts"
1286 )
1288 validate(metadata, "metadata")
1291def _normalize_positions(positions: Optional[Sequence[int]], seq_len: int) -> List[int]:
1292 """Normalize negative chunk-local positions and raise before indexing."""
1293 if positions is None:
1294 return list(range(seq_len))
1295 normalized = [position + seq_len if position < 0 else position for position in positions]
1296 out_of_range = [position for position in normalized if not 0 <= position < seq_len]
1297 if out_of_range:
1298 raise ValueError(
1299 f"positions {out_of_range} out of range for an activation chunk of length {seq_len}"
1300 )
1301 return normalized
1304def _cached_on_device(
1305 tensor: torch.Tensor,
1306 cache: Dict[torch.device, torch.Tensor],
1307 device: Union[str, torch.device],
1308) -> torch.Tensor:
1309 """Cache a small fp32 intervention tensor on each activation device."""
1310 resolved_device = torch.device(device)
1311 local = cache.get(resolved_device)
1312 if local is None: 1312 ↛ 1315line 1312 didn't jump to line 1315 because the condition on line 1312 was always true
1313 local = tensor.to(device=resolved_device, dtype=torch.float32)
1314 cache[resolved_device] = local
1315 return local
1318def _unit_rows(vectors: torch.Tensor, *, layer: int) -> torch.Tensor:
1319 """Normalize intervention vectors and reject zero/non-finite directions."""
1320 vectors = vectors.float()
1321 norms = vectors.norm(dim=-1, keepdim=True)
1322 if (~torch.isfinite(norms) | (norms <= torch.finfo(torch.float32).eps)).any(): 1322 ↛ 1323line 1322 didn't jump to line 1323 because the condition on line 1322 was never true
1323 raise ValueError(f"lens vectors at layer {layer} contain a zero or non-finite direction")
1324 return vectors / norms
1327def _make_intervention_hook(
1328 transform: Callable[[torch.Tensor], torch.Tensor],
1329 positions: Optional[Sequence[int]],
1330 d_model: int,
1331) -> Callable[..., torch.Tensor]:
1332 """Apply a transform with shared position, dtype, and device hardening."""
1333 requested = None if positions is None else tuple(positions)
1334 if requested == (): 1334 ↛ 1335line 1334 didn't jump to line 1335 because the condition on line 1334 was never true
1335 raise ValueError("positions must contain at least one index")
1337 def hook_fn(
1338 activation: Float[torch.Tensor, "batch pos d_model"], hook: Any
1339 ) -> Float[torch.Tensor, "batch pos d_model"]:
1340 hook_name = getattr(hook, "name", "intervention hook")
1341 _validate_residual_activation(activation, d_model=d_model, hook_name=hook_name)
1342 normalized = _normalize_positions(requested, activation.shape[1])
1343 selected = activation if requested is None else activation[:, normalized, :]
1344 transformed = transform(selected)
1345 if transformed.shape != selected.shape: 1345 ↛ 1346line 1345 didn't jump to line 1346 because the condition on line 1345 was never true
1346 raise ValueError(
1347 f"intervention returned shape {tuple(transformed.shape)}, "
1348 f"expected {tuple(selected.shape)}"
1349 )
1350 transformed = transformed.to(device=activation.device, dtype=activation.dtype)
1351 if requested is None:
1352 return transformed
1353 output = activation.clone()
1354 output[:, normalized, :] = transformed
1355 return output
1357 return hook_fn
1360def _unembed(
1361 model: Any, residual: Float[torch.Tensor, "pos d_model"]
1362) -> Float[torch.Tensor, "pos d_vocab"]:
1363 """Apply the model's own final norm, unembedding, logit scale, and soft cap.
1365 The norm/unembed components contract on ``[batch, pos, d_model]``, so the
1366 position rows are passed through with a singleton batch axis. The residual
1367 is moved to the unembedding's device for sharded/multi-GPU models.
1368 """
1369 unembed_weight = model.W_U
1370 compute_dtype = unembed_weight.dtype
1371 batched = residual.to(device=unembed_weight.device, dtype=compute_dtype).unsqueeze(0)
1372 logits = model.unembed(model.ln_final(batched)).squeeze(0)
1373 return model.adapter.apply_output_logits_transform(logits).float()
1376def _to_token_ids(model: Any, tokens: Union[TokenInput, Sequence[TokenInput]]) -> List[int]:
1377 """Convert token strings / ids into a list of single-token ids."""
1378 if isinstance(tokens, (str, int)):
1379 tokens = [tokens]
1380 ids: List[int] = []
1381 for token in tokens:
1382 if isinstance(token, str):
1383 ids.append(model.to_single_token(token))
1384 else:
1385 ids.append(int(token))
1386 if not ids:
1387 raise ValueError("tokens must contain at least one token")
1388 d_vocab = model.W_U.shape[1]
1389 invalid = [token_id for token_id in ids if not 0 <= token_id < d_vocab]
1390 if invalid:
1391 raise ValueError(f"token ids {invalid} out of range for vocabulary size {d_vocab}")
1392 return ids
1395def _validate_residual_activation(
1396 activation: torch.Tensor,
1397 *,
1398 d_model: int,
1399 hook_name: str,
1400) -> None:
1401 """Fail before interpreting a non-standard block output as a residual stream."""
1402 if activation.ndim != 3 or activation.shape[-1] != d_model:
1403 raise ValueError(
1404 f"{hook_name} must have shape [batch, position, {d_model}], "
1405 f"got {tuple(activation.shape)}"
1406 )
1409def _normalize_layer(layer: int, n_layers: int) -> int:
1410 """Resolve negative layer indices and bounds-check."""
1411 resolved = layer + n_layers if layer < 0 else layer
1412 if not 0 <= resolved < n_layers:
1413 raise ValueError(f"layer {layer} out of range for a {n_layers}-layer model")
1414 return resolved
1417class _frozen_parameters:
1418 """Context manager: freeze all parameters, restore their flags on exit.
1420 Freezing keeps the autograd graph rooted at the residual stream rather than
1421 at the weights, so fitting retains only the blocks between the earliest
1422 source layer and the target layer.
1423 """
1425 def __init__(self, model: Any) -> None:
1426 self.model = model
1427 self.saved: List[Tuple[torch.nn.Parameter, bool]] = []
1429 def __enter__(self) -> None:
1430 self.saved = [(param, param.requires_grad) for param in self.model.parameters()]
1431 for param, _ in self.saved:
1432 param.requires_grad_(False)
1434 def __exit__(self, *exc: Any) -> None:
1435 for param, flag in self.saved:
1436 param.requires_grad_(flag)
1439def _jacobian_for_prompt(
1440 model: Any,
1441 tokens: Int[torch.Tensor, "one seq"],
1442 *,
1443 source_layers: List[int],
1444 dim_batch: int,
1445 skip_first_positions: int,
1446) -> Dict[int, Float[torch.Tensor, "d_model d_model"]]:
1447 """Exact per-prompt Jacobian rows via batched one-hot cotangents.
1449 Assumes parameters are already frozen (see :class:`_frozen_parameters`) so
1450 that marking the earliest source activation ``requires_grad`` roots the
1451 graph there.
1452 """
1453 d_model = model.cfg.d_model
1454 target_layer = model.cfg.n_layers - 1
1455 seq_len = tokens.shape[1]
1456 valid_positions = list(range(skip_first_positions, seq_len - 1))
1457 replicated = tokens.expand(dim_batch, -1)
1459 captured: Dict[str, torch.Tensor] = {}
1460 root_name = _resid_post_hook_name(min(source_layers))
1461 hook_layers = sorted(set(source_layers) | {target_layer})
1463 def capture_fn(
1464 activation: Float[torch.Tensor, "batch pos d_model"], hook: Any
1465 ) -> Float[torch.Tensor, "batch pos d_model"]:
1466 _validate_residual_activation(activation, d_model=d_model, hook_name=hook.name)
1467 if hook.name == root_name and not activation.requires_grad:
1468 activation.requires_grad_(True)
1469 captured[hook.name] = activation
1470 return activation
1472 fwd_hooks = [(_resid_post_hook_name(layer), capture_fn) for layer in hook_layers]
1473 with torch.enable_grad(), model.hooks(fwd_hooks=fwd_hooks):
1474 model(replicated, return_type=None)
1476 target = captured[_resid_post_hook_name(target_layer)]
1477 sources = [captured[_resid_post_hook_name(layer)] for layer in source_layers]
1478 device = target.device
1479 positions_index = torch.tensor(valid_positions, device=device)
1480 batch_index = torch.arange(dim_batch, device=device)
1482 jacobians = {
1483 layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in source_layers
1484 }
1485 cotangent = torch.zeros_like(target)
1486 n_passes = -(-d_model // dim_batch) # ceil division
1487 for pass_index in range(n_passes):
1488 dim_start = pass_index * dim_batch
1489 n_dims = min(dim_batch, d_model - dim_start)
1490 cotangent.zero_()
1491 cotangent[
1492 batch_index[:n_dims, None],
1493 positions_index[None, :],
1494 dim_start + batch_index[:n_dims, None],
1495 ] = 1.0
1496 grads = torch.autograd.grad(
1497 outputs=target,
1498 inputs=sources,
1499 grad_outputs=cotangent,
1500 retain_graph=pass_index < n_passes - 1,
1501 )
1502 for layer, grad in zip(source_layers, grads):
1503 # each gradient lives on its layer's device under sharded/device_map setups
1504 rows = grad[:n_dims, positions_index.to(grad.device), :].float().mean(dim=1)
1505 jacobians[layer][dim_start : dim_start + n_dims, :] = rows.cpu()
1506 del grads
1507 return jacobians