Coverage for transformer_lens/tools/analysis/jacobian_lens.py: 95%
692 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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 warnings
64from dataclasses import dataclass
65from importlib.metadata import version
66from typing import (
67 Any,
68 Callable,
69 Dict,
70 List,
71 Mapping,
72 MutableMapping,
73 Optional,
74 Sequence,
75 Tuple,
76 Union,
77)
79import torch
80from jaxtyping import Float, Int
81from tqdm.auto import tqdm
83from transformer_lens.ActivationCache import ActivationCache
84from transformer_lens.tools.analysis._model_state import require_eval_mode
85from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import (
86 CoordinatePatch,
87 solve_coordinate_patch,
88 solve_coordinate_patch_positions,
89)
90from transformer_lens.tools.analysis.jacobian_lens_decomposition import (
91 DEFAULT_K,
92 JSpaceDecomposition,
93 JSpaceOccupancy,
94 JSpaceVarianceProfile,
95 _diagnose_intervention_pair,
96 estimate_occupancy,
97 get_sparse_decomposition,
98)
99from transformer_lens.utilities.hf_utils import call_hf_with_retry
101TokenInput = Union[str, int]
103# Backward-provider seam for the fit drive loop. A provider takes the target
104# residual, the source residuals it differentiates against, one batched one-hot
105# cotangent, and the ``retain_graph`` flag, and returns one gradient per source
106# (the same tuple ``torch.autograd.grad`` returns). Keeping the backward step
107# behind this narrow callable lets the estimator-independent driver run an
108# alternate estimator without touching its capture / cotangent-batching /
109# averaging machinery; the ordinary J-lens path passes :func:`_ordinary_vjp`.
110BackwardProvider = Callable[
111 [torch.Tensor, List[torch.Tensor], torch.Tensor, bool],
112 Tuple[torch.Tensor, ...],
113]
115# ---------------------------------------------------------------------------
116# Registry helpers
117# ---------------------------------------------------------------------------
119_REGISTRY_CACHE: Optional[Dict[str, Any]] = None
122def _load_registry() -> Dict[str, Any]:
123 """Return the bundled artifact registry, loading it once on first call."""
124 global _REGISTRY_CACHE
125 if _REGISTRY_CACHE is None:
126 import json
127 import pathlib
129 registry_path = pathlib.Path(__file__).with_name("jacobian_lens_registry.json")
130 with registry_path.open() as fh:
131 _REGISTRY_CACHE = json.load(fh)
132 return _REGISTRY_CACHE
135def _resolve_registry_entry(name_or_path: str) -> Optional[Tuple[str, str]]:
136 """Return ``(repo_id, filename)`` if *name_or_path* matches the registry.
138 Matching is tried in two passes:
139 1. Direct key match against short model names (e.g. ``"gemma-2-2b"``).
140 2. Alias match against Hugging Face model IDs (e.g. ``"google/gemma-2-2b"``).
142 Returns ``None`` when there is no match so callers can fall through to the
143 generic Hub download path.
144 """
145 registry = _load_registry()
146 if name_or_path in registry:
147 entry = registry[name_or_path]
148 return entry["repo_id"], entry["filename"]
149 for entry in registry.values():
150 if isinstance(entry, dict) and name_or_path in entry.get("aliases", []):
151 return entry["repo_id"], entry["filename"]
152 return None
155# Fitting excludes early positions (attention sinks with atypical residual statistics)
156# and the final position (no next-token target), matching the reference implementation.
157DEFAULT_SKIP_FIRST_POSITIONS = 16
158DEFAULT_TOP_K = 10
160# Keys written by fit() that must not appear in converted-lens metadata so that
161# merge() can refuse to mix TL-fitted lenses with externally converted ones.
162# Note: "target_layer" is intentionally NOT listed here — it must survive
163# conversion so that validate_model() can detect and refuse checkpoints that
164# were fitted against a non-final target layer.
165_FIT_RESERVED_KEYS: frozenset = frozenset(
166 {
167 "transformer_lens_fit",
168 "transformer_lens_version",
169 "model_system",
170 "processing",
171 "hook_convention",
172 "fit_dtype",
173 "dim_batch",
174 "max_seq_len",
175 "skip_first_positions",
176 }
177)
179# Top-level payload keys that some checkpoint writers store as flat provenance
180# (rather than nested under a "metadata" key).
181_CHECKPOINT_FLAT_PROVENANCE: frozenset = frozenset({"model_name", "model_revision", "corpus"})
184@dataclass
185class JacobianLensReadout:
186 """Result of a :meth:`JacobianLens.readout` call.
188 Attributes:
189 lens_topk_values:
190 Per-layer retained top-k pre-softmax values, on CPU.
191 lens_topk_indices:
192 Per-layer retained top-k vocabulary ids, on CPU.
193 model_topk_values:
194 The model output's retained top-k pre-softmax values, on CPU.
195 model_topk_indices:
196 The model output's retained top-k vocabulary ids, on CPU.
197 lens_logits:
198 Optional full per-layer logits, on CPU. Present only when
199 ``readout(return_full_logits=True)`` was requested.
200 model_logits:
201 Optional full model logits, on CPU. Present only when
202 ``readout(return_full_logits=True)`` was requested.
203 tokens:
204 The token ids of the run prompt, ``[seq]``.
205 positions:
206 The (normalized, non-negative) positions the readout covers, aligned
207 with the ``pos`` axis of retained top-k and optional full logits.
208 use_jacobian:
209 Whether the Jacobian transport was applied (``False`` = logit lens).
210 """
212 lens_topk_values: Dict[int, Float[torch.Tensor, "pos k"]]
213 lens_topk_indices: Dict[int, Int[torch.Tensor, "pos k"]]
214 model_topk_values: Float[torch.Tensor, "pos k"]
215 model_topk_indices: Int[torch.Tensor, "pos k"]
216 tokens: Int[torch.Tensor, "seq"]
217 positions: List[int]
218 use_jacobian: bool = True
219 lens_logits: Optional[Dict[int, Float[torch.Tensor, "pos d_vocab"]]] = None
220 model_logits: Optional[Float[torch.Tensor, "pos d_vocab"]] = None
222 def top_tokens(self, tokenizer: Any, k: int = 5) -> Dict[int, List[List[str]]]:
223 """Decode the top-``k`` tokens per layer and position.
225 Args:
226 tokenizer: The model's tokenizer (``model.tokenizer``).
227 k: Number of top tokens to decode per (layer, position).
229 Returns:
230 ``{layer: [ [top-k strings] per position ]}``, positions aligned with
231 :attr:`positions`.
232 """
233 out: Dict[int, List[List[str]]] = {}
234 retained = next(iter(self.lens_topk_indices.values())).shape[-1]
235 if not 1 <= k <= retained:
236 raise ValueError(f"k must be between 1 and the retained top_k={retained}, got {k}")
237 for layer, ids in self.lens_topk_indices.items():
238 out[layer] = [[tokenizer.decode([t]) for t in row[:k].tolist()] for row in ids]
239 return out
242class JacobianLens:
243 """A fitted Jacobian lens: one transport matrix per source layer.
245 Layer convention (matching the reference implementation and the published
246 artifacts): index ``l`` refers to the **output of block** ``l`` at the
247 Bridge-native hook ``blocks.{l}.hook_out``. ``J[l]`` maps that activation
248 to the final block's output, pre final norm. The final layer itself is never
249 fitted (its transport is the identity), so
250 ``source_layers == [0, ..., n_layers - 2]`` for a full fit.
252 Attributes:
253 jacobians: ``{layer: [d_model, d_model]}`` transport matrices, fp32, CPU.
254 n_prompts: Number of prompts averaged into the fit.
255 d_model: Residual stream width the lens was fitted for.
256 metadata: Optional provenance (model name, TransformerLens version, fit
257 hyperparameters). Preserved by :meth:`save`/:meth:`load`; artifacts
258 from the reference implementation load with empty metadata.
259 """
261 def __init__(
262 self,
263 jacobians: Dict[int, Float[torch.Tensor, "d_model d_model"]],
264 *,
265 n_prompts: int,
266 d_model: int,
267 metadata: Optional[Dict[str, Any]] = None,
268 ) -> None:
269 if not jacobians: 269 ↛ 270line 269 didn't jump to line 270 because the condition on line 269 was never true
270 raise ValueError("jacobians must contain at least one layer")
271 for layer, matrix in jacobians.items():
272 if matrix.shape != (d_model, d_model): 272 ↛ 273line 272 didn't jump to line 273 because the condition on line 272 was never true
273 raise ValueError(
274 f"jacobians[{layer}] has shape {tuple(matrix.shape)}, "
275 f"expected ({d_model}, {d_model})"
276 )
277 self.jacobians: Dict[int, torch.Tensor] = {
278 int(layer): matrix.detach().float().cpu() for layer, matrix in jacobians.items()
279 }
280 self.n_prompts = int(n_prompts)
281 self.d_model = int(d_model)
282 self.metadata: Dict[str, Any] = dict(metadata or {})
283 self._device_jacobians: Dict[Tuple[int, torch.device], torch.Tensor] = {}
284 self._dictionary_cache: Dict[Tuple[int, torch.device], torch.Tensor] = {}
285 self._unembedding_snapshots: Dict[torch.device, torch.Tensor] = {}
287 @property
288 def source_layers(self) -> List[int]:
289 """Sorted list of layers this lens has transport matrices for."""
290 return sorted(self.jacobians)
292 def __repr__(self) -> str:
293 layers = self.source_layers
294 return (
295 f"JacobianLens(layers={layers[0]}..{layers[-1]} ({len(layers)}), "
296 f"d_model={self.d_model}, n_prompts={self.n_prompts})"
297 )
299 # ------------------------------------------------------------------ #
300 # persistence #
301 # ------------------------------------------------------------------ #
303 def save(self, path: str, *, dtype: torch.dtype = torch.float16) -> None:
304 """Save the lens in the reference implementation's artifact format.
306 The four official keys (``J``, ``n_prompts``, ``source_layers``,
307 ``d_model``) are written unchanged so the file stays loadable by the
308 reference package; TransformerLens provenance is stored under an
309 additive ``metadata`` key.
311 Args:
312 path: Destination ``.pt`` path.
313 dtype: Storage dtype. Defaults to fp16 like the reference
314 implementation — Jacobian entries are order-one, so the smaller
315 dtype costs little precision and halves the artifact on disk.
316 """
317 _validate_metadata(self.metadata)
318 payload: Dict[str, Any] = {
319 "J": {layer: matrix.to(dtype) for layer, matrix in self.jacobians.items()},
320 "n_prompts": self.n_prompts,
321 "source_layers": self.source_layers,
322 "d_model": self.d_model,
323 }
324 if self.metadata:
325 payload["metadata"] = self.metadata
326 torch.save(payload, path)
328 @classmethod
329 def load(cls, path: str) -> "JacobianLens":
330 """Load a lens artifact or fit checkpoint saved in a supported schema.
332 Two file schemas are accepted:
334 **Artifact** (the reference format, written by :meth:`save` or the
335 Anthropic reference package): must contain a ``J`` key mapping layer
336 indices to transport matrices, plus ``n_prompts``, ``d_model``, and an
337 optional ``metadata`` dict.
339 **Fit checkpoint** (running-sum format, written by the reference
340 implementation's ``write_checkpoint()`` during fitting): must contain
341 a ``jacobian_sum`` key mapping layer indices to *running-sum* matrices
342 (i.e. the sum over prompts, not yet divided by the prompt count), plus
343 ``n_done``. ``d_model`` is inferred from the first matrix's shape;
344 no explicit ``d_model`` key is required or expected. The per-layer
345 means are reconstructed on load. A
346 ``converted_from: "jacobian_lens_checkpoint"`` key is added to
347 metadata so :meth:`merge` refuses to silently combine checkpoints with
348 natively TL-fitted lenses. Fit-reserved provenance keys
349 (``transformer_lens_fit``, etc.) are stripped; scalar fields that can
350 be serialised under ``weights_only=True`` are preserved.
351 Tensor-valued metadata fields that would fail :func:`_validate_metadata`
352 are recorded by name in a ``dropped_fields`` list.
354 Fit checkpoint schema (reference ``write_checkpoint()`` format)
355 ---------------------------------------------------------------
356 The reference implementation writes exactly six top-level keys; all
357 other keys in the payload are ignored::
359 {
360 "jacobian_sum": {<layer int>: <float32 tensor [d, d]>, ...},
361 "n_done": <int>, # prompts accumulated into jacobian_sum
362 "next_idx": <int>, # next prompt index (informational)
363 "source_layers": [<int>, ...], # documented layer indices (informational)
364 "target_layer": <int>, # target layer — harvested into metadata
365 "skip_first": <int>, # leading positions skipped (informational)
366 # optional flat provenance accepted from alternative checkpoint writers:
367 "model_name": <str>,
368 "model_revision": <str>,
369 "corpus": <str>,
370 # optional nested provenance accepted from alternative writers:
371 "metadata": {<str>: <scalar/list/dict>, ...},
372 }
374 Args:
375 path: Path to the ``.pt`` file.
377 Raises:
378 ValueError: If the file lacks both a ``J`` key (artifact) and a
379 ``jacobian_sum`` key (checkpoint), or if a checkpoint records a
380 non-positive ``n_prompts``.
381 """
382 payload = torch.load(path, map_location="cpu", weights_only=True)
383 if "J" in payload:
384 return cls(
385 {int(layer): matrix for layer, matrix in payload["J"].items()},
386 n_prompts=int(payload.get("n_prompts", 0)),
387 d_model=int(payload["d_model"]),
388 metadata=payload.get("metadata"),
389 )
390 if "jacobian_sum" in payload:
391 return cls._from_checkpoint_payload(path, payload)
392 raise ValueError(
393 f"{path} does not look like a Jacobian lens artifact or fit checkpoint. "
394 "Expected a 'J' key (artifact) or 'jacobian_sum' key (fit checkpoint). "
395 "See JacobianLens.load() for the supported file schemas."
396 )
398 @classmethod
399 def _from_checkpoint_payload(cls, path: str, payload: Dict[str, Any]) -> "JacobianLens":
400 """Reconstruct a JacobianLens from a fit-checkpoint payload.
402 Divides the running Jacobian sums by ``n_prompts``, strips fit-reserved
403 provenance keys, harvests safe scalar metadata, records dropped tensor
404 fields, and marks the result as converted so :meth:`merge` refuses to
405 mix it with natively TL-fitted lenses.
406 """
407 n_prompts = int(payload.get("n_done", payload.get("n_prompts", 0)))
408 if n_prompts <= 0:
409 raise ValueError(
410 f"{path} is a fit checkpoint with n_prompts={n_prompts}; "
411 "a positive prompt count is required to reconstruct the Jacobian mean."
412 )
413 if not payload["jacobian_sum"]:
414 raise ValueError(
415 f"{path} is a fit checkpoint with an empty jacobian_sum; "
416 "at least one layer matrix is required to reconstruct d_model."
417 )
418 first_matrix = next(iter(payload["jacobian_sum"].values()))
419 d_model = first_matrix.shape[0]
420 jacobians = {
421 int(layer): matrix.float() / n_prompts
422 for layer, matrix in payload["jacobian_sum"].items()
423 }
425 # Collect raw provenance: first from nested "metadata", then supplement
426 # with flat top-level keys that some checkpoint writers place directly
427 # in the payload (model_name, model_revision, corpus).
428 raw_meta: Dict[str, Any] = dict(payload.get("metadata") or {})
429 for key in _CHECKPOINT_FLAT_PROVENANCE:
430 if key in payload and key not in raw_meta:
431 raw_meta[key] = payload[key]
432 # target_layer lives at the top level in the reference checkpoint format;
433 # harvest it into metadata so validate_model() can check the fitting target.
434 if "target_layer" in payload and "target_layer" not in raw_meta:
435 raw_meta["target_layer"] = payload["target_layer"]
437 # Build clean metadata: drop fit-reserved keys, record tensor-valued
438 # fields that _validate_metadata would reject (they cannot survive a
439 # weights_only=True reload), and keep everything else that validates.
440 dropped_fields: List[str] = []
441 clean_meta: Dict[str, Any] = {}
442 for key, value in raw_meta.items():
443 if key in _FIT_RESERVED_KEYS:
444 continue
445 if isinstance(value, torch.Tensor):
446 dropped_fields.append(f"{key}: shape={tuple(value.shape)} dtype={value.dtype}")
447 continue
448 try:
449 _validate_metadata({key: value})
450 clean_meta[key] = value
451 except ValueError:
452 dropped_fields.append(key)
454 clean_meta["converted_from"] = "jacobian_lens_checkpoint"
455 if dropped_fields:
456 clean_meta["dropped_fields"] = dropped_fields
458 return cls(jacobians, n_prompts=n_prompts, d_model=d_model, metadata=clean_meta)
460 @classmethod
461 def from_pretrained(
462 cls,
463 name_or_path: str,
464 *,
465 filename: str = "lens.pt",
466 revision: Optional[str] = None,
467 model: Any = None,
468 ) -> "JacobianLens":
469 """Load a lens from a local path, a short model name, or a Hub repo.
471 Resolution order
472 ----------------
473 1. **Local file** — if *name_or_path* is an existing ``.pt`` file,
474 load it directly.
475 2. **Local directory** — if *name_or_path* is a directory, load
476 ``<name_or_path>/<filename>``.
477 3. **Registry short name or HF model ID** — if *name_or_path* matches
478 a key or alias in the bundled ``jacobian_lens_registry.json`` (e.g.
479 ``"gemma-2-2b"`` or ``"google/gemma-2-2b"``), the corresponding
480 artifact in ``neuronpedia/jacobian-lens`` is fetched automatically.
481 The *filename* argument is ignored in this case because the registry
482 already encodes the correct subpath.
483 4. **Explicit Hub repo** — otherwise *name_or_path* is treated as a Hub
484 repo id and *filename* is used as-is, preserving full backward
485 compatibility (e.g. ``from_pretrained("neuronpedia/jacobian-lens",
486 filename="gpt2-small/jlens/...")``).
488 Args:
489 name_or_path: A local ``.pt`` file, a local directory, a short
490 model name such as ``"gemma-2-2b"`` or ``"llama3.1-8b"``, a
491 Hugging Face model ID such as ``"google/gemma-2-2b"``, or an
492 explicit Hub repo id paired with *filename*.
493 filename: File (or subpath) inside a local directory or an explicit
494 Hub repo. Ignored when *name_or_path* resolves via the
495 registry.
496 revision: Optional Hub revision (branch, tag, or commit) to pin.
497 When omitted, the Hub repository's mutable default branch is
498 followed; pin a commit hash for reproducible analyses.
499 model: If given, :meth:`validate_model` is called so dimension or
500 weight-processing mismatches fail here rather than at first use.
502 Returns:
503 The loaded (and, if ``model`` was given, validated) lens.
505 Examples::
507 # Short model name — no need to remember the HF subpath
508 lens = JacobianLens.from_pretrained("gemma-2-2b", model=model)
510 # HF model ID also works
511 lens = JacobianLens.from_pretrained("google/gemma-2-2b", model=model)
513 # Explicit Hub repo + subpath (backward-compatible)
514 lens = JacobianLens.from_pretrained(
515 "neuronpedia/jacobian-lens",
516 filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
517 model=model,
518 )
519 """
520 import os
522 if os.path.isfile(name_or_path): 522 ↛ 523line 522 didn't jump to line 523 because the condition on line 522 was never true
523 lens = cls.load(name_or_path)
524 elif os.path.isdir(name_or_path): 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true
525 lens = cls.load(os.path.join(name_or_path, filename))
526 else:
527 from huggingface_hub import hf_hub_download
529 resolved = _resolve_registry_entry(name_or_path)
530 if resolved is not None:
531 repo_id, resolved_filename = resolved
532 else:
533 repo_id, resolved_filename = name_or_path, filename
535 local_path = call_hf_with_retry(
536 hf_hub_download,
537 repo_id=repo_id,
538 filename=resolved_filename,
539 revision=revision,
540 )
541 lens = cls.load(local_path)
542 if model is not None:
543 lens.validate_model(model)
544 return lens
546 @classmethod
547 def merge(cls, lenses: Sequence["JacobianLens"]) -> "JacobianLens":
548 """Combine lenses fitted on disjoint prompt slices.
550 The per-layer matrices are averaged weighted by each lens's
551 ``n_prompts``, matching the reference implementation, so fitting can be
552 parallelized across processes or machines and merged afterwards.
553 Provenance must match across shards (apart from ``n_prompts``), so a
554 merge cannot silently relabel matrices fitted with different models,
555 corpora, dtypes, or estimator settings. The merged count replaces the
556 per-shard count.
558 Args:
559 lenses: Lenses that agree exactly on ``source_layers`` and
560 ``d_model``.
562 Raises:
563 ValueError: On an empty sequence or mismatched lenses.
564 """
565 if not lenses:
566 raise ValueError("cannot merge an empty sequence of lenses")
567 invalid_counts = [
568 (index, lens.n_prompts) for index, lens in enumerate(lenses) if lens.n_prompts <= 0
569 ]
570 if invalid_counts:
571 raise ValueError(
572 "every lens passed to merge() must have positive n_prompts; "
573 f"invalid shards: {invalid_counts}"
574 )
575 first = lenses[0]
576 for lens in lenses:
577 _validate_metadata(lens.metadata)
578 first_provenance = {
579 key: value for key, value in first.metadata.items() if key != "n_prompts"
580 }
581 for other in lenses[1:]:
582 if other.source_layers != first.source_layers or other.d_model != first.d_model:
583 raise ValueError(
584 "all lenses being merged must share the same source_layers and d_model"
585 )
586 other_provenance = {
587 key: value for key, value in other.metadata.items() if key != "n_prompts"
588 }
589 if other_provenance != first_provenance:
590 raise ValueError(
591 "all lenses being merged must share the same provenance metadata "
592 "apart from n_prompts"
593 )
594 total = sum(lens.n_prompts for lens in lenses)
595 merged = {
596 layer: torch.stack([lens.jacobians[layer] * lens.n_prompts for lens in lenses]).sum(
597 dim=0
598 )
599 / total
600 for layer in first.source_layers
601 }
602 metadata = dict(first.metadata)
603 if metadata:
604 metadata["n_prompts"] = total
605 return cls(merged, n_prompts=total, d_model=first.d_model, metadata=metadata)
607 # ------------------------------------------------------------------ #
608 # model validation #
609 # ------------------------------------------------------------------ #
611 def validate_model(self, model: Any) -> "JacobianLens":
612 """Check that ``model`` matches this lens; raise loudly if not.
614 Requires a raw causal ``TransformerBridge`` with the standard direct
615 final-norm/unembed path, verifies recorded model provenance, residual
616 width and layer range, and enforces the published final-block target
617 convention.
619 Args:
620 model: A raw ``TransformerBridge``.
622 Returns:
623 ``self``, for chaining.
625 Raises:
626 TypeError: If model is not a ``TransformerBridge``.
627 ValueError: On model provenance or ``d_model`` mismatch,
628 out-of-range source layers, compatibility mode, unsupported
629 attention/output paths, or a non-final target convention.
630 """
631 _require_raw_bridge(model)
632 artifact_model_name = self.metadata.get("model_name")
633 current_model_name = getattr(model.cfg, "model_name", None)
634 if artifact_model_name is not None and artifact_model_name != current_model_name:
635 raise ValueError(
636 f"lens was fitted for model {artifact_model_name!r}, but the supplied "
637 f"model is {current_model_name!r}."
638 )
639 artifact_revision = self.metadata.get("model_revision")
640 current_revision = _get_model_revision(model)
641 if artifact_revision is not None and artifact_revision != current_revision:
642 raise ValueError(
643 f"lens was fitted for model revision {artifact_revision!r}, but the "
644 f"supplied model revision is {current_revision!r}."
645 )
646 d_model = model.cfg.d_model
647 if d_model != self.d_model:
648 raise ValueError(
649 f"lens was fitted for d_model={self.d_model}, but the model has "
650 f"d_model={d_model} — this lens belongs to a different model."
651 )
652 n_layers = model.cfg.n_layers
653 final_layer = n_layers - 1
654 out_of_range = [layer for layer in self.source_layers if not 0 <= layer < final_layer]
655 if out_of_range:
656 raise ValueError(
657 f"lens has source layers {out_of_range} outside the model's "
658 f"0..{final_layer - 1} source range — this lens belongs to a different model."
659 )
660 target_layer = int(self.metadata.get("target_layer", final_layer))
661 if target_layer != final_layer:
662 raise ValueError(
663 f"lens targets layer {target_layer}, but readout supports only the "
664 f"published final-layer convention ({final_layer}); refit without "
665 "a custom target layer."
666 )
667 return self
669 # ------------------------------------------------------------------ #
670 # reading #
671 # ------------------------------------------------------------------ #
673 def clear_device_cache(self) -> None:
674 """Release cached Jacobians, dictionaries, and unembedding snapshots on devices."""
675 self._device_jacobians.clear()
676 self._dictionary_cache.clear()
677 self._unembedding_snapshots.clear()
679 def _matrix_on(self, layer: int, device: Union[str, torch.device]) -> torch.Tensor:
680 """Return one cached fp32 Jacobian copy for a layer/device pair."""
681 if layer not in self.jacobians:
682 raise ValueError(
683 f"layer {layer} is not in this lens's source layers "
684 f"({self.source_layers[0]}..{self.source_layers[-1]})"
685 )
686 resolved_device = torch.device(device)
687 key = (layer, resolved_device)
688 matrix = self._device_jacobians.get(key)
689 if matrix is None:
690 matrix = self.jacobians[layer].to(device=resolved_device, dtype=torch.float32)
691 self._device_jacobians[key] = matrix
692 return matrix
694 def transport(
695 self,
696 residual: Float[torch.Tensor, "... d_model"],
697 layer: int,
698 ) -> Float[torch.Tensor, "... d_model"]:
699 """Map layer-``layer`` activations into the final block's output basis.
701 Computes ``J[layer] @ h`` per activation vector, in fp32.
703 Args:
704 residual: Activations from the output of block ``layer``.
705 layer: Source layer index.
706 """
707 matrix = self._matrix_on(layer, residual.device)
708 return residual.float() @ matrix.T
710 @torch.no_grad()
711 def readout(
712 self,
713 model: Any,
714 input: Union[str, Int[torch.Tensor, "batch seq"]],
715 *,
716 layers: Optional[Sequence[int]] = None,
717 positions: Optional[Sequence[int]] = None,
718 use_jacobian: bool = True,
719 top_k: int = DEFAULT_TOP_K,
720 return_full_logits: bool = False,
721 ) -> JacobianLensReadout:
722 """Read per-layer vocabulary logits for a prompt.
724 Runs the model once with caching, transports the residual stream at each
725 requested layer through ``J[layer]`` (or the identity when
726 ``use_jacobian=False`` — the logit lens), and applies the model's own
727 final norm, unembedding, architecture logit scaling, and logit soft cap.
729 Args:
730 model: A raw ``TransformerBridge``.
731 input: A prompt string, or a ``[1, seq]`` token tensor.
732 layers: Layers to read. Defaults to every fitted layer plus the
733 final layer. The final layer (``n_layers - 1``) is always read
734 with the identity transport — by construction its lens equals
735 the model's own output distribution.
736 positions: Token positions to read (negative indices allowed).
737 Defaults to all positions.
738 use_jacobian: Apply the Jacobian transport. ``False`` gives the
739 logit-lens baseline through the identical code path.
740 top_k: Number of values and vocabulary ids retained per layer and
741 position. Defaults to 10.
742 return_full_logits: Also retain full vocabulary tensors on CPU.
743 This is opt-in because a 64-token Gemma readout across all
744 layers is roughly 1.7 GB.
746 Returns:
747 A :class:`JacobianLensReadout`.
749 Raises:
750 ValueError: If the model fails :meth:`validate_model`, ``input`` is
751 batched, ``top_k`` is invalid, or a requested layer has no
752 transport matrix.
753 """
754 self.validate_model(model)
755 tokens = model.to_tokens(input) if isinstance(input, str) else input
756 if tokens.ndim != 2 or tokens.shape[0] != 1: 756 ↛ 757line 756 didn't jump to line 757 because the condition on line 756 was never true
757 raise ValueError(f"readout expects a single prompt; got shape {tuple(tokens.shape)}")
758 n_layers = model.cfg.n_layers
759 final_layer = n_layers - 1
760 if layers is None:
761 layers = self.source_layers + [final_layer]
762 layers = [_normalize_layer(layer, n_layers) for layer in layers]
763 for layer in layers:
764 if use_jacobian and layer != final_layer and layer not in self.jacobians: 764 ↛ 765line 764 didn't jump to line 765 because the condition on line 764 was never true
765 raise ValueError(
766 f"layer {layer} is not in this lens's source layers; "
767 f"available: {self.source_layers} (+{final_layer} as identity)"
768 )
770 if top_k < 1: 770 ↛ 771line 770 didn't jump to line 771 because the condition on line 770 was never true
771 raise ValueError(f"top_k must be at least 1, got {top_k}")
772 seq_len = tokens.shape[1]
773 norm_positions = _normalize_positions(positions, seq_len)
775 hook_names = {
776 layer: _resid_post_hook_name(layer) for layer in layers if layer != final_layer
777 }
778 wanted = set(hook_names.values())
779 logits, cache = model.run_with_cache(tokens, names_filter=lambda name: name in wanted)
780 selected_model_logits = logits[0, norm_positions, :].float()
781 if top_k > selected_model_logits.shape[-1]: 781 ↛ 782line 781 didn't jump to line 782 because the condition on line 781 was never true
782 raise ValueError(
783 f"top_k={top_k} exceeds the model vocabulary size "
784 f"{selected_model_logits.shape[-1]}"
785 )
786 model_topk = selected_model_logits.topk(top_k, dim=-1)
787 full_model_logits = selected_model_logits.cpu() if return_full_logits else None
788 lens_topk_values: Dict[int, torch.Tensor] = {}
789 lens_topk_indices: Dict[int, torch.Tensor] = {}
790 full_lens_logits: Optional[Dict[int, torch.Tensor]] = {} if return_full_logits else None
791 for layer in layers:
792 if layer == final_layer:
793 layer_logits = selected_model_logits
794 layer_topk = model_topk
795 else:
796 activation = cache[hook_names[layer]]
797 _validate_residual_activation(
798 activation,
799 d_model=model.cfg.d_model,
800 hook_name=hook_names[layer],
801 )
802 residual = activation[0, norm_positions, :]
803 transported = self.transport(residual, layer) if use_jacobian else residual.float()
804 layer_logits = _unembed(model, transported)
805 layer_topk = layer_logits.topk(top_k, dim=-1)
806 lens_topk_values[layer] = layer_topk.values.cpu()
807 lens_topk_indices[layer] = layer_topk.indices.cpu()
808 if full_lens_logits is not None:
809 if layer == final_layer:
810 assert full_model_logits is not None
811 full_lens_logits[layer] = full_model_logits
812 else:
813 full_lens_logits[layer] = layer_logits.cpu()
814 return JacobianLensReadout(
815 lens_topk_values=lens_topk_values,
816 lens_topk_indices=lens_topk_indices,
817 model_topk_values=model_topk.values.cpu(),
818 model_topk_indices=model_topk.indices.cpu(),
819 tokens=tokens[0].cpu(),
820 positions=norm_positions,
821 use_jacobian=use_jacobian,
822 lens_logits=full_lens_logits,
823 model_logits=full_model_logits,
824 )
826 @torch.no_grad()
827 def lens_vectors(
828 self,
829 model: Any,
830 tokens: Union[TokenInput, Sequence[TokenInput]],
831 layer: int,
832 ) -> Float[torch.Tensor, "n d_model"]:
833 """Residual-stream directions for vocabulary tokens at a layer.
835 The J-lens vector for token ``t`` is row ``t`` of ``W_U J[layer]``
836 expressed in layer-``layer`` residual coordinates:
837 ``v_t = J[layer]^T W_U[:, t]``.
839 Args:
840 model: The model supplying ``W_U``.
841 tokens: A token string / id, or a sequence of them. Strings must
842 encode to a single token.
843 layer: Source layer for the vectors.
845 Returns:
846 One vector per token, fp32, on the model's device.
847 """
848 self.validate_model(model)
849 layer = _normalize_layer(layer, model.cfg.n_layers)
850 token_ids = _to_token_ids(model, tokens)
851 unembed_columns = model.W_U[:, token_ids].float() # [d_model, n]
852 matrix = self._matrix_on(layer, unembed_columns.device)
853 return (matrix.T @ unembed_columns).T
855 @torch.no_grad()
856 def lens_vector_dictionary(
857 self, model: Any, layer: int
858 ) -> Float[torch.Tensor, "d_vocab d_model"]:
859 """Full-vocabulary J-lens dictionary at ``layer``: ``[d_vocab, d_model]``.
861 Row ``t`` is the J-lens vector ``v_t = J[layer]^T W_U[:, t]`` -- this is
862 :meth:`lens_vectors` over the entire vocabulary. The result is cached while the
863 model's unembedding is unchanged so a sparse decomposition can reuse it;
864 :meth:`clear_device_cache` releases it.
866 The dictionary is vocabulary-sized and cached on the model's device
867 (``d_vocab * d_model`` fp32 values, on the order of gigabytes for a large
868 vocabulary), one entry per requested layer. One detached copy of ``W_U`` is
869 retained per device to detect changes without transferring weights to the host.
871 Args:
872 model: The model supplying ``W_U``.
873 layer: Source layer for the dictionary (must be a fitted source layer).
875 Returns:
876 The dictionary, fp32, on the model's device.
877 """
878 self.validate_model(model)
879 layer = _normalize_layer(layer, model.cfg.n_layers)
880 unembed = model.W_U
881 device = torch.device(unembed.device)
882 snapshot = self._unembedding_snapshots.get(device)
883 if snapshot is None or not torch.equal(snapshot, unembed):
884 self._unembedding_snapshots[device] = unembed.detach().clone()
885 stale_keys = [key for key in self._dictionary_cache if key[1] == device]
886 for key in stale_keys:
887 del self._dictionary_cache[key]
889 key = (layer, device)
890 dictionary = self._dictionary_cache.get(key)
891 if dictionary is None:
892 matrix = self._matrix_on(layer, device) # [d_model, d_model]
893 dictionary = (matrix.T @ unembed.float()).T # [d_vocab, d_model]
894 self._dictionary_cache[key] = dictionary
895 return dictionary
897 @torch.no_grad()
898 def decompose(
899 self,
900 model: Any,
901 activation_or_prompt: Union[torch.Tensor, str],
902 layer: int,
903 *,
904 position: Optional[int] = None,
905 k: int = DEFAULT_K,
906 algorithm: str = "nonnegative_orthogonal_matching_pursuit",
907 ) -> JSpaceDecomposition:
908 """Decompose an activation into its J-space content at ``layer``.
910 ``activation_or_prompt`` is either:
912 - a raw activation vector of shape ``[d_model]`` (leave ``position`` as ``None``), or
913 - a prompt -- a string or a ``[1, seq]`` token tensor -- in which case ``position``
914 selects the token whose ``blocks.{layer}.hook_out`` activation is decomposed.
916 The full-vocabulary dictionary at ``layer`` is built (and cached) via
917 :meth:`lens_vector_dictionary`, then :func:`get_sparse_decomposition` solves for a
918 ``k``-sparse nonnegative combination of J-lens vectors.
920 Args:
921 model: A raw ``TransformerBridge``.
922 activation_or_prompt: An activation vector, or a prompt (string / token tensor).
923 layer: Source layer (must be a fitted source layer).
924 position: Token position when a prompt is given; must be ``None`` for a raw
925 activation vector.
926 k: Upper bound on the number of J-lens vectors to select; selection stops early
927 once no unselected vector is materially positively correlated, so fewer may be
928 returned.
929 algorithm: Coefficient-update rule; see :func:`get_sparse_decomposition`.
931 Returns:
932 A :class:`JSpaceDecomposition`. Its ``support`` (token ids here) holds only the
933 numerically active J-lens vectors and ``selected_support`` every selected vector,
934 with ``support.numel() <= selected_support.numel() <= k``. ``support`` and its
935 token-decoding tensors are on CPU; the vector-valued outputs stay on the model's
936 device.
938 Raises:
939 ValueError: On an invalid model, a mismatched activation shape, a batched prompt,
940 an unfitted layer, or an invalid ``k`` / ``algorithm``.
941 RuntimeError: If the default nonnegative least-squares solver cannot validate its
942 result against the KKT conditions.
943 """
944 activation, resolved_layer = self._resolve_activation(
945 model, activation_or_prompt, layer, position
946 )
947 dictionary = self.lens_vector_dictionary(model, resolved_layer)
948 return get_sparse_decomposition(
949 activation.float().to(dictionary.device), dictionary, k, algorithm=algorithm
950 )
952 @torch.no_grad()
953 def coordinate_patch(
954 self,
955 model: Any,
956 activation_or_prompt: Union[torch.Tensor, str],
957 layer: int,
958 source_token: TokenInput,
959 target_token: TokenInput,
960 *,
961 position: Optional[int] = None,
962 decomposition: Optional[JSpaceDecomposition] = None,
963 k: int = DEFAULT_K,
964 mode: str = "substitute",
965 alpha: float = 1.0,
966 algorithm: str = "nonnegative_orthogonal_matching_pursuit",
967 ) -> CoordinatePatch:
968 """Patch sparse J-space coordinates for an activation at ``layer``.
970 The activation may be a raw ``[d_model]`` vector or a prompt paired with ``position``,
971 exactly as in :meth:`decompose`. ``source_token`` must occur in the active sparse support.
972 ``substitute`` replaces the target coordinate with the source coordinate and zeros the
973 source; ``swap`` exchanges them. Other coordinates and ``x - reconstruction`` are fixed.
975 A supplied ``decomposition`` avoids repeating the vocabulary-scale sparse solve after its
976 activation and dictionary compatibility have been validated.
978 Args:
979 model: A raw ``TransformerBridge``.
980 activation_or_prompt: An activation vector, or a prompt (string / token tensor).
981 layer: Source layer (must be a fitted source layer).
982 source_token: Active source concept, as a single-token string or token id.
983 target_token: Distinct target concept, as a single-token string or token id.
984 position: Token position when a prompt is given; ``None`` for a raw activation.
985 decomposition: Optional compatible decomposition of this activation and dictionary.
986 k: Sparse-solver upper bound when ``decomposition`` is not supplied.
987 mode: ``"substitute"`` or ``"swap"``.
988 alpha: Finite interpolation strength; zero is an exact no-op.
989 algorithm: Sparse coefficient-update rule when solving a fresh decomposition.
991 Returns:
992 A :class:`CoordinatePatch` containing the edited frame, diagnostics, and activation.
993 """
994 activation, resolved_layer = self._resolve_activation(
995 model, activation_or_prompt, layer, position
996 )
997 dictionary = self.lens_vector_dictionary(model, resolved_layer)
998 source_id, target_id = _to_token_ids(model, [source_token, target_token])
999 return solve_coordinate_patch(
1000 activation.float().to(dictionary.device),
1001 dictionary,
1002 source_id,
1003 target_id,
1004 decomposition=decomposition,
1005 k=k,
1006 mode=mode,
1007 alpha=alpha,
1008 algorithm=algorithm,
1009 )
1011 def _resolve_activation(
1012 self,
1013 model: Any,
1014 activation_or_prompt: Union[torch.Tensor, str],
1015 layer: int,
1016 position: Optional[int],
1017 ) -> Tuple[torch.Tensor, int]:
1018 """Validate the model and layer, then resolve either a raw ``[d_model]`` activation or a
1019 prompt plus ``position`` to the activation vector to analyse.
1021 Returns ``(activation, resolved_layer)``. Shared by :meth:`decompose`,
1022 :meth:`coordinate_patch`, and :meth:`occupancy` so all accept the same input forms with
1023 identical validation.
1024 """
1025 self.validate_model(model)
1026 resolved_layer = _normalize_layer(layer, model.cfg.n_layers)
1027 if resolved_layer not in self.jacobians:
1028 raise ValueError(
1029 f"layer {layer} is not in this lens's source layers; "
1030 f"available: {self.source_layers}"
1031 )
1033 if position is None:
1034 if not isinstance(activation_or_prompt, torch.Tensor):
1035 raise ValueError(
1036 "analysis expects a raw activation tensor when position is None; pass a "
1037 "prompt together with a position to analyze a model activation"
1038 )
1039 activation = activation_or_prompt
1040 if activation.ndim != 1 or activation.shape[0] != self.d_model:
1041 raise ValueError(
1042 f"activation must be 1-D of length d_model={self.d_model}, "
1043 f"got shape {tuple(activation.shape)}"
1044 )
1045 else:
1046 if isinstance(activation_or_prompt, torch.Tensor) and activation_or_prompt.ndim == 1:
1047 raise ValueError(
1048 "position is only valid with a prompt (string or [1, seq] tokens); "
1049 "pass a raw 1-D activation with position=None instead"
1050 )
1051 tokens = (
1052 model.to_tokens(activation_or_prompt)
1053 if isinstance(activation_or_prompt, str)
1054 else activation_or_prompt
1055 )
1056 if tokens.ndim != 2 or tokens.shape[0] != 1:
1057 raise ValueError(
1058 f"analysis expects a single prompt; got shape {tuple(tokens.shape)}"
1059 )
1060 hook_name = _resid_post_hook_name(resolved_layer)
1061 _, cache = model.run_with_cache(tokens, names_filter=lambda name: name == hook_name)
1062 norm_position = _normalize_positions([position], tokens.shape[1])[0]
1063 activation = cache[hook_name][0, norm_position, :]
1064 return activation, resolved_layer
1066 @torch.no_grad()
1067 def occupancy(
1068 self,
1069 model: Any,
1070 activation_or_prompt: Union[torch.Tensor, str],
1071 layer: int,
1072 *,
1073 position: Optional[int] = None,
1074 max_atoms: int = DEFAULT_K,
1075 num_control_dictionaries: int = 32,
1076 seed: int = 0,
1077 ) -> JSpaceOccupancy:
1078 """Estimate how many J-lens vectors are meaningfully active in an activation at ``layer``.
1080 Resolves ``activation_or_prompt`` (a raw ``[d_model]`` vector, or a prompt plus
1081 ``position``) exactly as :meth:`decompose`, builds the cached full-vocabulary dictionary
1082 via :meth:`lens_vector_dictionary`, and calls :func:`estimate_occupancy`.
1084 Args:
1085 model: A raw ``TransformerBridge``.
1086 activation_or_prompt: An activation vector, or a prompt (string / token tensor).
1087 layer: Source layer (must be a fitted source layer).
1088 position: Token position when a prompt is given; ``None`` for a raw activation.
1089 max_atoms: Maximum number of J-lens vectors to consider.
1090 num_control_dictionaries: Number of random control dictionaries to average over.
1091 seed: Seed for the random control dictionaries (reproducibility).
1093 Returns:
1094 A :class:`JSpaceOccupancy`.
1095 """
1096 activation, resolved_layer = self._resolve_activation(
1097 model, activation_or_prompt, layer, position
1098 )
1099 dictionary = self.lens_vector_dictionary(model, resolved_layer)
1100 return estimate_occupancy(
1101 activation.float().to(dictionary.device),
1102 dictionary,
1103 max_atoms=max_atoms,
1104 num_control_dictionaries=num_control_dictionaries,
1105 seed=seed,
1106 )
1108 @torch.no_grad()
1109 def fraction_of_variance(
1110 self,
1111 model: Any,
1112 prompts: Union[str, torch.Tensor, Sequence[Union[str, torch.Tensor]]],
1113 layers: Optional[Sequence[int]] = None,
1114 *,
1115 k: int = DEFAULT_K,
1116 skip_first: int = 16,
1117 positions: Optional[Sequence[int]] = None,
1118 show_progress: bool = False,
1119 ) -> JSpaceVarianceProfile:
1120 """Profile the J-space share of activation variance over a prompt corpus.
1122 Each prompt is run once (caching ``blocks.{layer}.hook_out`` for every requested layer).
1123 At each sampled position the activation is decomposed and its J-space variance fraction
1124 ``||j_space_component||^2 / ||activation||^2`` is recorded. The numerator is the
1125 ``j_space_component`` -- the orthogonal projection of the activation onto the span of the
1126 selected support (the paper's appendix "J-space component"), *not* the nonnegative
1127 ``reconstruction``; the two coincide only when every selected atom stays active. Per layer
1128 the profile reports the median of those fractions and the pooled ratio
1129 ``sum(||j_space_component||^2) / sum(||activation||^2)`` (the paper's "fraction of total
1130 variance").
1132 A layer that samples no positions -- every prompt shorter than ``skip_first``, or only
1133 zero-norm activations -- contributes no fractions: its ``median`` and ``pooled`` are
1134 ``float("nan")`` and its ``per_position`` tensor is empty.
1136 Args:
1137 model: A raw ``TransformerBridge``.
1138 prompts: A prompt, or a sequence of prompts. Each token tensor must represent exactly
1139 one prompt and have shape ``[1, seq]``.
1140 layers: Source layers to profile; defaults to all fitted ``source_layers``.
1141 k: Number of J-lens vectors per decomposition.
1142 skip_first: Non-negative index before which positions are skipped (mirrors the fit's
1143 early-position skip); not used for sampling when ``positions`` is given.
1144 positions: Explicit positions to sample instead of ``skip_first`` onward.
1145 show_progress: Show a tqdm progress bar over prompts.
1147 Returns:
1148 A :class:`JSpaceVarianceProfile`.
1150 Raises:
1151 ValueError: On an invalid model, an unfitted layer, an empty corpus, a negative
1152 ``skip_first``, or a token tensor that does not have shape ``[1, seq]``.
1153 """
1154 self.validate_model(model)
1155 if skip_first < 0:
1156 raise ValueError(f"skip_first must be non-negative, got {skip_first}")
1157 if layers is None:
1158 resolved_layers = list(self.source_layers)
1159 else:
1160 resolved_layers = [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]
1161 for layer in resolved_layers:
1162 if layer not in self.jacobians:
1163 raise ValueError(
1164 f"layer {layer} is not in this lens's source layers; "
1165 f"available: {self.source_layers}"
1166 )
1167 prompt_list: List[Union[str, torch.Tensor]] = (
1168 [prompts] if isinstance(prompts, (str, torch.Tensor)) else list(prompts)
1169 )
1170 if not prompt_list:
1171 raise ValueError("prompts must be a non-empty prompt or sequence of prompts")
1173 hook_names = {layer: _resid_post_hook_name(layer) for layer in resolved_layers}
1174 wanted_hooks = set(hook_names.values())
1175 dictionaries = {
1176 layer: self.lens_vector_dictionary(model, layer) for layer in resolved_layers
1177 }
1178 fractions: Dict[int, List[float]] = {layer: [] for layer in resolved_layers}
1179 pooled_j_space: Dict[int, float] = {layer: 0.0 for layer in resolved_layers}
1180 pooled_total: Dict[int, float] = {layer: 0.0 for layer in resolved_layers}
1182 for prompt in tqdm(prompt_list, desc="J-space variance", disable=not show_progress):
1183 tokens = model.to_tokens(prompt) if isinstance(prompt, str) else prompt
1184 if tokens.ndim != 2 or tokens.shape[0] != 1:
1185 raise ValueError(
1186 "fraction_of_variance expects each tokenized prompt to have shape "
1187 f"[1, seq], got {tuple(tokens.shape)}"
1188 )
1189 _, cache = model.run_with_cache(tokens, names_filter=lambda name: name in wanted_hooks)
1190 seq_len = tokens.shape[1]
1191 sampled = (
1192 list(range(skip_first, seq_len))
1193 if positions is None
1194 else _normalize_positions(positions, seq_len)
1195 )
1196 for layer in resolved_layers:
1197 dictionary = dictionaries[layer]
1198 activations = cache[hook_names[layer]][0] # [seq, d_model]
1199 for position in sampled:
1200 activation = activations[position].float().to(dictionary.device)
1201 total = float(activation @ activation)
1202 if total <= 0.0: 1202 ↛ 1203line 1202 didn't jump to line 1203 because the condition on line 1202 was never true
1203 continue
1204 decomposition = get_sparse_decomposition(activation, dictionary, k)
1205 j_space = float(
1206 decomposition.j_space_component @ decomposition.j_space_component
1207 )
1208 fractions[layer].append(j_space / total)
1209 pooled_j_space[layer] += j_space
1210 pooled_total[layer] += total
1212 median = {
1213 layer: float(torch.tensor(fractions[layer]).median())
1214 if fractions[layer]
1215 else float("nan")
1216 for layer in resolved_layers
1217 }
1218 pooled = {
1219 layer: pooled_j_space[layer] / pooled_total[layer]
1220 if pooled_total[layer] > 0
1221 else float("nan")
1222 for layer in resolved_layers
1223 }
1224 per_position = {layer: torch.tensor(fractions[layer]) for layer in resolved_layers}
1225 return JSpaceVarianceProfile(
1226 layers=resolved_layers, median=median, pooled=pooled, per_position=per_position
1227 )
1229 # ------------------------------------------------------------------ #
1230 # interventions #
1231 # ------------------------------------------------------------------ #
1233 def steering_hooks(
1234 self,
1235 model: Any,
1236 token: TokenInput,
1237 layers: Sequence[int],
1238 *,
1239 alpha: float = 4.0,
1240 positions: Optional[Sequence[int]] = None,
1241 ) -> List[Tuple[str, Any]]:
1242 """Hooks that steer the residual stream along a token's J-lens vector.
1244 At each layer the unit-normalized lens vector is added, scaled by
1245 ``alpha`` times the activation's **median** per-position residual norm:
1246 ``h <- h + alpha * median||h|| * v̂``. This norm-matched
1247 parameterization follows the steering description in the reference
1248 implementation's experiment protocols; the paper's minimal form is
1249 the unscaled ``h <- h + alpha * v_t``, recoverable by passing the raw
1250 :meth:`lens_vectors` output to your own hook. The median (not mean) is
1251 used so attention-sink positions — whose residual norms run orders of
1252 magnitude above typical positions — do not inflate the scale.
1254 Args:
1255 model: The model the hooks will run on.
1256 token: The concept token to steer toward.
1257 layers: Layers to intervene at.
1258 alpha: Steering strength scalar; ``0`` disables. Because of the
1259 norm-matched scale, values of order 1 already perturb the
1260 stream by roughly its own magnitude.
1261 positions: Chunk-local positions to steer (negative indices allowed
1262 and normalized on every hook invocation). Defaults to all.
1264 Returns:
1265 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)`` or
1266 ``model.run_with_hooks(fwd_hooks=...)``.
1267 """
1268 self.validate_model(model)
1269 hooks = []
1270 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
1271 direction = self.lens_vectors(model, token, layer)[0]
1272 unit = _unit_rows(direction.unsqueeze(0), layer=layer)[0]
1273 device_units: Dict[torch.device, torch.Tensor] = {}
1275 def transform(
1276 selected: Float[torch.Tensor, "batch pos d_model"],
1277 unit: torch.Tensor = unit,
1278 device_units: Dict[torch.device, torch.Tensor] = device_units,
1279 ) -> Float[torch.Tensor, "batch pos d_model"]:
1280 local_unit = _cached_on_device(unit, device_units, selected.device)
1281 scale = alpha * selected.float().norm(dim=-1).median()
1282 return selected.float() + scale * local_unit
1284 hooks.append(
1285 (
1286 _resid_post_hook_name(layer),
1287 _make_intervention_hook(transform, positions, model.cfg.d_model),
1288 )
1289 )
1290 return hooks
1292 def ablation_hooks(
1293 self,
1294 model: Any,
1295 tokens: Union[TokenInput, Sequence[TokenInput]],
1296 layers: Sequence[int],
1297 *,
1298 positions: Optional[Sequence[int]] = None,
1299 ) -> List[Tuple[str, Any]]:
1300 """Hooks that project token directions out of the residual stream.
1302 For each token's unit lens vector ``v̂``: ``h <- h - (h·v̂) v̂``,
1303 applied sequentially when several tokens are given.
1305 Args:
1306 model: The model the hooks will run on.
1307 tokens: Concept token(s) to suppress.
1308 layers: Layers to intervene at.
1309 positions: Chunk-local positions to ablate (negative indices allowed
1310 and normalized on every hook invocation). Defaults to all.
1312 Returns:
1313 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
1314 """
1315 self.validate_model(model)
1316 hooks = []
1317 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
1318 vectors = self.lens_vectors(model, tokens, layer)
1319 units = _unit_rows(vectors, layer=layer)
1320 device_units: Dict[torch.device, torch.Tensor] = {}
1322 def transform(
1323 selected: Float[torch.Tensor, "batch pos d_model"],
1324 units: torch.Tensor = units,
1325 device_units: Dict[torch.device, torch.Tensor] = device_units,
1326 ) -> Float[torch.Tensor, "batch pos d_model"]:
1327 local_units = _cached_on_device(units, device_units, selected.device)
1328 result = selected.float()
1329 for unit in local_units:
1330 coeff = result @ unit
1331 result = result - coeff.unsqueeze(-1) * unit
1332 return result
1334 hooks.append(
1335 (
1336 _resid_post_hook_name(layer),
1337 _make_intervention_hook(transform, positions, model.cfg.d_model),
1338 )
1339 )
1340 return hooks
1342 def swap_hooks(
1343 self,
1344 model: Any,
1345 source_token: TokenInput,
1346 target_token: TokenInput,
1347 layers: Sequence[int],
1348 *,
1349 alpha: float = 1.0,
1350 positions: Optional[Sequence[int]] = None,
1351 ) -> List[Tuple[str, Any]]:
1352 """Hooks that swap two concepts' *live* coordinates in lens space.
1354 The paper's patching-in-lens-coordinates intervention: with
1355 ``V = [v_s, v_t]`` and lens coordinates ``c = V⁺ h`` (pseudoinverse),
1356 the update is ``h <- h + alpha * V (sigma(c) - c)`` where ``sigma``
1357 exchanges the two coordinates. The component of ``h`` orthogonal to
1358 ``span{v_s, v_t}`` is untouched. ``alpha=2`` is the paper's
1359 "double-strength" swap.
1361 This transform re-reads ``c`` from the activation seen by every hook.
1362 It is therefore an involution when applied repeatedly in a subspace
1363 whose coordinates are preserved between layers: a second application
1364 can undo the first. For the paper's multi-layer clamp protocol, use
1365 :meth:`swap_clamp_hooks` with activations cached from the clean run.
1367 Args:
1368 model: The model the hooks will run on.
1369 source_token: The concept to remove (e.g. ``" France"``).
1370 target_token: The concept to install (e.g. ``" China"``).
1371 layers: Layers to intervene at.
1372 alpha: Swap strength.
1373 positions: Chunk-local positions to swap (negative indices allowed
1374 and normalized on every hook invocation). Defaults to all.
1376 Returns:
1377 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
1378 """
1379 self.validate_model(model)
1380 source_id, target_id = _to_token_ids(model, [source_token, target_token])
1381 if source_id == target_id:
1382 raise ValueError(
1383 "source_token and target_token resolve to the same token id; "
1384 "a coordinate swap would be a silent no-op"
1385 )
1387 hooks = []
1388 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
1389 vectors = self.lens_vectors(model, [source_id, target_id], layer)
1390 units = _unit_rows(vectors, layer=layer)
1391 _diagnose_intervention_pair(
1392 units, description=f"swap vectors at layer {layer}", stacklevel=3
1393 )
1394 basis = vectors.T # [d, 2]
1395 pinv = torch.linalg.pinv(basis) # [2, d]
1396 device_basis: Dict[torch.device, torch.Tensor] = {}
1397 device_pinv: Dict[torch.device, torch.Tensor] = {}
1399 def transform(
1400 selected: Float[torch.Tensor, "batch pos d_model"],
1401 basis: torch.Tensor = basis,
1402 pinv: torch.Tensor = pinv,
1403 device_basis: Dict[torch.device, torch.Tensor] = device_basis,
1404 device_pinv: Dict[torch.device, torch.Tensor] = device_pinv,
1405 ) -> Float[torch.Tensor, "batch pos d_model"]:
1406 local_basis = _cached_on_device(basis, device_basis, selected.device)
1407 local_pinv = _cached_on_device(pinv, device_pinv, selected.device)
1408 coords = selected.float() @ local_pinv.T # [..., 2]
1409 delta = alpha * ((coords[..., [1, 0]] - coords) @ local_basis.T)
1410 return selected.float() + delta
1412 hooks.append(
1413 (
1414 _resid_post_hook_name(layer),
1415 _make_intervention_hook(transform, positions, model.cfg.d_model),
1416 )
1417 )
1418 return hooks
1420 def swap_clamp_hooks(
1421 self,
1422 model: Any,
1423 source_token: TokenInput,
1424 target_token: TokenInput,
1425 layers: Sequence[int],
1426 clean_cache: Union[ActivationCache, Mapping[str, torch.Tensor]],
1427 *,
1428 positions: Optional[Sequence[int]] = None,
1429 ) -> List[Tuple[str, Any]]:
1430 """Hooks that clamp lens coordinates to their clean-run exchange.
1432 For each layer, this projects the corresponding activation from
1433 ``clean_cache`` into that layer's lens basis, exchanges its source and
1434 target coordinates once, and holds the live activation at that fixed
1435 target. Unlike :meth:`swap_hooks`, the update is idempotent at each
1436 layer: ``h <- h + V (c_target - V⁺h)``.
1438 Args:
1439 model: The model the hooks will run on.
1440 source_token: The concept to remove (e.g. ``" France"``).
1441 target_token: The concept to install (e.g. ``" China"``).
1442 layers: Layers to intervene at.
1443 clean_cache: Activations from an unmodified ``run_with_cache`` at
1444 each requested layer's ``blocks.{layer}.hook_out`` name. Either
1445 the ``ActivationCache`` it returns by default or the plain dict
1446 from ``return_cache_object=False`` is accepted.
1447 positions: Chunk-local positions to clamp (negative indices allowed
1448 and normalized against the clean activations). Defaults to all.
1450 Returns:
1451 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
1452 """
1453 self.validate_model(model)
1454 source_id, target_id = _to_token_ids(model, [source_token, target_token])
1455 if source_id == target_id:
1456 raise ValueError(
1457 "source_token and target_token resolve to the same token id; "
1458 "a coordinate clamp would be a silent no-op"
1459 )
1461 hooks = []
1462 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
1463 hook_name = _resid_post_hook_name(layer)
1464 if hook_name not in clean_cache:
1465 raise ValueError(f"clean_cache is missing activation {hook_name!r}")
1466 clean = clean_cache[hook_name]
1467 _validate_residual_activation(
1468 clean, d_model=model.cfg.d_model, hook_name=f"clean_cache[{hook_name!r}]"
1469 )
1470 normalized = _normalize_positions(positions, clean.shape[1])
1471 clean_selected = clean if positions is None else clean[:, normalized, :]
1473 vectors = self.lens_vectors(model, [source_id, target_id], layer)
1474 units = _unit_rows(vectors, layer=layer)
1475 _diagnose_intervention_pair(
1476 units, description=f"swap vectors at layer {layer}", stacklevel=3
1477 )
1478 basis = vectors.T # [d, 2]
1479 pinv = torch.linalg.pinv(basis) # [2, d]
1480 target_coords = (clean_selected.float() @ pinv.T)[..., [1, 0]]
1481 device_basis: Dict[torch.device, torch.Tensor] = {}
1482 device_pinv: Dict[torch.device, torch.Tensor] = {}
1483 device_targets: Dict[torch.device, torch.Tensor] = {}
1485 def transform(
1486 selected: Float[torch.Tensor, "batch pos d_model"],
1487 basis: torch.Tensor = basis,
1488 pinv: torch.Tensor = pinv,
1489 target_coords: torch.Tensor = target_coords,
1490 device_basis: Dict[torch.device, torch.Tensor] = device_basis,
1491 device_pinv: Dict[torch.device, torch.Tensor] = device_pinv,
1492 device_targets: Dict[torch.device, torch.Tensor] = device_targets,
1493 ) -> Float[torch.Tensor, "batch pos d_model"]:
1494 local_basis = _cached_on_device(basis, device_basis, selected.device)
1495 local_pinv = _cached_on_device(pinv, device_pinv, selected.device)
1496 local_targets = _cached_on_device(target_coords, device_targets, selected.device)
1497 if local_targets.shape[1] != selected.shape[1] or local_targets.shape[0] not in (
1498 1,
1499 selected.shape[0],
1500 ):
1501 raise ValueError(
1502 "clean_cache activation shape is incompatible with the live activation: "
1503 f"target coordinates have shape {tuple(local_targets.shape)}, "
1504 f"live activation has shape {tuple(selected.shape)}"
1505 )
1506 coords = selected.float() @ local_pinv.T
1507 delta = (local_targets - coords) @ local_basis.T
1508 return selected.float() + delta
1510 hooks.append(
1511 (
1512 hook_name,
1513 _make_intervention_hook(transform, positions, model.cfg.d_model),
1514 )
1515 )
1516 return hooks
1518 def coordinate_patch_hooks(
1519 self,
1520 model: Any,
1521 source_token: TokenInput,
1522 target_token: TokenInput,
1523 layers: Sequence[int],
1524 *,
1525 positions: Sequence[int],
1526 decomposition_cache: Optional[
1527 MutableMapping[Tuple[int, int, int], JSpaceDecomposition]
1528 ] = None,
1529 k: int = DEFAULT_K,
1530 mode: str = "substitute",
1531 alpha: float = 1.0,
1532 algorithm: str = "nonnegative_orthogonal_matching_pursuit",
1533 ) -> List[Tuple[str, Any]]:
1534 """Hooks that anchor-patch one J-space coordinate live, per forward-pass position.
1536 Unlike :meth:`coordinate_patch`, which edits one already-captured activation offline,
1537 this installs a forward hook that solves :func:`solve_coordinate_patch` independently for
1538 every ``(batch_idx, position)`` pair at each requested layer -- a vocabulary-scale sparse
1539 decomposition per pair, per hook firing, unless ``decomposition_cache`` supplies one
1540 already validated for that ``(layer, batch_idx, position)`` key.
1542 Args:
1543 model: The model the hooks will run on.
1544 source_token: Active source concept, as a single-token string or token id.
1545 target_token: Distinct target concept, as a single-token string or token id.
1546 layers: Layers to intervene at.
1547 positions: Chunk-local positions to patch (negative indices allowed and normalized on
1548 every hook invocation). Required -- there is no full-sequence default, because a
1549 silent default would trigger a vocabulary-scale solve at every position.
1550 decomposition_cache: Optional caller-owned mapping from ``(layer, batch_idx,
1551 position)`` to a previously validated
1552 :class:`~transformer_lens.tools.analysis.jacobian_lens_decomposition.JSpaceDecomposition`.
1553 A hit skips the vocabulary-scale scan; a miss solves and populates the cache.
1554 Purely a performance path -- correctness does not depend on it.
1555 k: Sparse-solver upper bound on a cache miss.
1556 mode: ``"substitute"`` or ``"swap"``.
1557 alpha: Finite interpolation strength; zero is an exact no-op.
1558 algorithm: Sparse coefficient-update rule on a cache miss.
1560 Returns:
1561 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
1563 Raises:
1564 ValueError: If ``positions`` is empty, if ``source_token`` and ``target_token``
1565 resolve to the same id, or if ``source_token`` is not in the top-``k`` active
1566 support of every patched ``(batch_idx, position)`` pair *at the moment its hook
1567 fires* -- the whole forward pass fails rather than silently patching a subset.
1568 This precondition is stronger and more order-dependent than "active on a clean
1569 forward pass": in a band of layers an earlier hook's patch edits the residual
1570 that a later layer re-decomposes, and ``substitute``/``swap`` zero or move the
1571 source coordinate, so the source can be removed from a later layer's active
1572 support even though it was active on an unhooked pass. Stacking layers or
1573 positions therefore makes this progressively harder to satisfy.
1575 Warns:
1576 UserWarning: Once per call, naming the number of layers and positions that will
1577 perform a live vocabulary-scale solve on every cache miss.
1578 """
1579 self.validate_model(model)
1580 if not positions:
1581 raise ValueError("positions must contain at least one index")
1582 resolved_layers = [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]
1583 source_id, target_id = _to_token_ids(model, [source_token, target_token])
1584 if source_id == target_id:
1585 raise ValueError(
1586 "source_token and target_token resolve to the same token id; "
1587 "a coordinate patch would be a silent no-op"
1588 )
1589 warnings.warn(
1590 f"coordinate_patch_hooks installs {len(resolved_layers)} layer(s) x "
1591 f"{len(positions)} position(s) of coordinate-patch hooks; every (batch, position) "
1592 "pair not already present in decomposition_cache performs a vocabulary-scale sparse "
1593 "decomposition on every forward pass",
1594 UserWarning,
1595 stacklevel=2,
1596 )
1598 requested = tuple(positions)
1599 hooks = []
1600 for layer in resolved_layers:
1601 dictionary = self.lens_vector_dictionary(model, layer)
1603 def hook_fn(
1604 activation: Float[torch.Tensor, "batch pos d_model"],
1605 hook: Any,
1606 layer: int = layer,
1607 dictionary: torch.Tensor = dictionary,
1608 ) -> Float[torch.Tensor, "batch pos d_model"]:
1609 hook_name = getattr(hook, "name", "intervention hook")
1610 _validate_residual_activation(
1611 activation, d_model=model.cfg.d_model, hook_name=hook_name
1612 )
1613 normalized = _normalize_positions(requested, activation.shape[1])
1614 selected = activation[:, normalized, :].float().to(dictionary.device)
1615 patched, _ = solve_coordinate_patch_positions(
1616 selected,
1617 dictionary,
1618 normalized,
1619 source_id,
1620 target_id,
1621 layer=layer,
1622 decomposition_cache=decomposition_cache,
1623 k=k,
1624 mode=mode,
1625 alpha=alpha,
1626 algorithm=algorithm,
1627 )
1628 output = activation.clone()
1629 output[:, normalized, :] = patched.to(
1630 device=activation.device, dtype=activation.dtype
1631 )
1632 return output
1634 hooks.append((_resid_post_hook_name(layer), hook_fn))
1635 return hooks
1637 # ------------------------------------------------------------------ #
1638 # fitting #
1639 # ------------------------------------------------------------------ #
1641 @classmethod
1642 def fit(
1643 cls,
1644 model: Any,
1645 prompts: Sequence[str],
1646 *,
1647 corpus: str,
1648 source_layers: Optional[Sequence[int]] = None,
1649 dim_batch: int = 8,
1650 max_seq_len: int = 128,
1651 skip_first_positions: int = DEFAULT_SKIP_FIRST_POSITIONS,
1652 show_progress: bool = True,
1653 metadata: Optional[Dict[str, Any]] = None,
1654 ) -> "JacobianLens":
1655 """Fit a Jacobian lens on a hooked model.
1657 Implements the reference estimator exactly. For each prompt: one forward
1658 pass (the prompt replicated ``dim_batch`` times along the batch axis),
1659 then ``ceil(d_model / dim_batch)`` backward passes. Each backward plants
1660 a one-hot cotangent for one output dimension at *every* valid target
1661 position simultaneously — causal attention guarantees the gradient at
1662 source position ``t`` is then the sum over target positions
1663 ``t' >= t`` with no explicit masking. Rows are averaged over valid
1664 source positions (the first ``skip_first_positions`` and the final
1665 position are excluded), and prompts contribute equally to the final
1666 mean. There is no randomness: the computation is deterministic given
1667 the prompts.
1669 The reference implementation reports that fit quality saturates
1670 quickly — on the order of 100 prompts of 128 tokens is usable; the
1671 published lenses use up to 1000. Use :meth:`merge` to parallelize
1672 across prompt slices.
1674 Args:
1675 model: A raw ``TransformerBridge``. Model parameters are temporarily
1676 frozen (``requires_grad=False``) during fitting and restored
1677 after. Cotangents and activation gradients use the model dtype;
1678 fit with a float32 model for the highest-fidelity estimator. The
1679 model and all of its submodules must be in evaluation mode.
1680 prompts: Prompt strings. Prompts too short to contain a valid
1681 position (``seq_len <= skip_first_positions + 1``) are skipped
1682 with a warning and do not count toward ``n_prompts``.
1683 corpus: Stable identifier for the prompt corpus or slice, recorded
1684 in artifact provenance.
1685 source_layers: Layers to fit. Defaults to every layer below
1686 the final layer. Negative indices count from ``n_layers``.
1687 dim_batch: Output dimensions per backward pass. Higher is faster
1688 but replicates the prompt ``dim_batch`` times in memory; total
1689 backward FLOPs are unchanged.
1690 max_seq_len: Prompts are truncated to this many tokens.
1691 skip_first_positions: Leading positions excluded from the source
1692 average.
1693 show_progress: Show a tqdm progress bar over prompts.
1694 metadata: Extra provenance merged into :attr:`metadata`.
1696 Returns:
1697 The fitted :class:`JacobianLens`.
1699 Raises:
1700 TypeError: If model is not a ``TransformerBridge``.
1701 ValueError: On compatibility mode, training mode, invalid provenance
1702 or layer indices, or if no prompt was long enough to fit on.
1703 """
1704 _require_raw_bridge(model)
1705 require_eval_mode(model, operation="JacobianLens.fit()")
1706 if not isinstance(corpus, str) or not corpus.strip(): 1706 ↛ 1707line 1706 didn't jump to line 1707 because the condition on line 1706 was never true
1707 raise ValueError("corpus must be a non-empty provenance identifier")
1708 n_layers = model.cfg.n_layers
1709 d_model = model.cfg.d_model
1710 resolved_target = n_layers - 1
1711 if source_layers is None:
1712 resolved_sources = list(range(resolved_target))
1713 else:
1714 resolved_sources = sorted(
1715 {_normalize_layer(layer, n_layers) for layer in source_layers}
1716 )
1717 if not resolved_sources: 1717 ↛ 1718line 1717 didn't jump to line 1718 because the condition on line 1717 was never true
1718 raise ValueError("source_layers is empty")
1719 if resolved_sources[-1] >= resolved_target: 1719 ↛ 1720line 1719 didn't jump to line 1720 because the condition on line 1719 was never true
1720 raise ValueError(
1721 f"every source layer must be below target_layer={resolved_target}; "
1722 f"got {resolved_sources}"
1723 )
1724 if dim_batch < 1: 1724 ↛ 1725line 1724 didn't jump to line 1725 because the condition on line 1724 was never true
1725 raise ValueError(f"dim_batch must be >= 1, got {dim_batch}")
1726 if skip_first_positions < 0:
1727 raise ValueError(f"skip_first_positions must be >= 0, got {skip_first_positions}")
1728 fit_dtype = model.W_U.dtype
1729 if fit_dtype in (torch.float16, torch.bfloat16):
1730 warnings.warn(
1731 f"fitting in {fit_dtype} accumulates Jacobian gradients at reduced "
1732 "precision; use a float32 TransformerBridge for the highest-fidelity fit",
1733 UserWarning,
1734 stacklevel=2,
1735 )
1737 transport_matrices, n_done = _fit_transport_matrices(
1738 model,
1739 prompts,
1740 source_layers=resolved_sources,
1741 dim_batch=dim_batch,
1742 max_seq_len=max_seq_len,
1743 skip_first_positions=skip_first_positions,
1744 show_progress=show_progress,
1745 backward_provider=_ordinary_vjp,
1746 )
1748 fit_metadata: Dict[str, Any] = {
1749 "model_name": getattr(model.cfg, "model_name", None),
1750 "model_revision": _get_model_revision(model),
1751 "transformer_lens_version": version("transformer-lens"),
1752 "model_system": "TransformerBridge",
1753 "processing": {
1754 "compatibility_mode": False,
1755 "weight_basis": "raw_huggingface",
1756 },
1757 "hook_convention": "blocks.{layer}.hook_out",
1758 "corpus": corpus,
1759 "n_prompts": n_done,
1760 "fit_dtype": str(fit_dtype).removeprefix("torch."),
1761 "target_layer": resolved_target,
1762 "dim_batch": dim_batch,
1763 "max_seq_len": max_seq_len,
1764 "skip_first_positions": skip_first_positions,
1765 "transformer_lens_fit": True,
1766 }
1767 reserved = sorted(set(fit_metadata).intersection(metadata or {}))
1768 if reserved: 1768 ↛ 1769line 1768 didn't jump to line 1769 because the condition on line 1768 was never true
1769 raise ValueError(f"metadata cannot override fit provenance keys: {reserved}")
1770 full_metadata = dict(metadata or {})
1771 full_metadata.update(fit_metadata)
1772 _validate_metadata(full_metadata)
1773 return cls(
1774 transport_matrices,
1775 n_prompts=n_done,
1776 d_model=d_model,
1777 metadata=full_metadata,
1778 )
1781# ---------------------------------------------------------------------- #
1782# helpers #
1783# ---------------------------------------------------------------------- #
1786def _resid_post_hook_name(layer: int) -> str:
1787 """Bridge-native hook for the output of block ``layer``."""
1788 return f"blocks.{layer}.hook_out"
1791def _get_model_revision(model: Any) -> Optional[str]:
1792 """Return the resolved Hugging Face commit recorded on a booted model."""
1793 original_model = getattr(model, "original_model", None)
1794 hf_config = getattr(original_model, "config", None)
1795 revision = getattr(hf_config, "_commit_hash", None)
1796 return revision if isinstance(revision, str) and revision else None
1799def _require_raw_bridge(model: Any) -> None:
1800 """Require the causal raw-Bridge contract used by fit and readout."""
1801 from transformer_lens.model_bridge import TransformerBridge
1803 if not isinstance(model, TransformerBridge):
1804 raise TypeError(
1805 "JacobianLens supports TransformerBridge only; load a fresh model with "
1806 "TransformerBridge.boot_transformers(...)."
1807 )
1808 if getattr(model, "compatibility_mode", False):
1809 raise ValueError(
1810 "compatibility mode is enabled on this TransformerBridge and changes "
1811 "the residual basis. Use a freshly booted "
1812 "TransformerBridge.boot_transformers(...) model with raw weights."
1813 )
1814 if getattr(model, "_weights_processed", False):
1815 raise ValueError(
1816 "process_weights was called on this TransformerBridge and changed the "
1817 "raw HuggingFace weight basis. Use a freshly booted "
1818 "TransformerBridge.boot_transformers(...) model."
1819 )
1820 adapter = model.adapter
1821 if not adapter.supports_generation:
1822 raise ValueError(
1823 "JacobianLens requires a causal decoder-only Bridge whose adapter "
1824 f"supports text generation; {type(adapter).__name__} declares "
1825 "supports_generation=False."
1826 )
1827 adapter.validate_output_logits_transform()
1828 attention_dir = getattr(model.cfg, "attention_dir", "causal")
1829 if attention_dir != "causal":
1830 raise ValueError(
1831 "JacobianLens requires causal attention because its estimator relies on "
1832 "causality to exclude target positions before each source position; "
1833 f"got attention_dir={attention_dir!r}."
1834 )
1835 total_ut_steps = int(getattr(model.cfg, "total_ut_steps", 1) or 1)
1836 if total_ut_steps != 1:
1837 raise ValueError(
1838 "JacobianLens requires each physical block hook to fire once per forward; "
1839 f"this looped-depth Bridge runs total_ut_steps={total_ut_steps}."
1840 )
1841 component_mapping = adapter.get_component_mapping()
1842 required_components = ("blocks", "ln_final", "unembed")
1843 missing_components = [
1844 component
1845 for component in required_components
1846 if component not in component_mapping or not hasattr(model, component)
1847 ]
1848 if missing_components:
1849 raise ValueError(
1850 "JacobianLens requires the standard direct ln_final -> unembed output path; "
1851 f"this Bridge is missing {missing_components}."
1852 )
1853 blocks_component = component_mapping["blocks"]
1854 if not getattr(blocks_component, "hook_out_is_single_residual_stream", False):
1855 raise ValueError(
1856 "JacobianLens requires single-stream [batch, position, d_model] block "
1857 f"outputs; {type(blocks_component).__name__} does not provide that contract."
1858 )
1859 if "project_out" in component_mapping:
1860 raise ValueError(
1861 "JacobianLens does not yet support a final output projection between "
1862 "the residual stream and unembedding."
1863 )
1864 unembed_width = model.W_U.shape[0]
1865 if unembed_width != model.cfg.d_model:
1866 raise ValueError(
1867 "JacobianLens requires a direct d_model-width unembedding after ln_final; "
1868 f"got W_U input width {unembed_width} for d_model={model.cfg.d_model}. "
1869 "Architectures with a final output projection are not yet supported."
1870 )
1873def _validate_metadata(metadata: Dict[str, Any]) -> None:
1874 """Reject values that ``torch.load(weights_only=True)`` cannot reload."""
1876 def validate(value: Any, path: str) -> None:
1877 if value is None or type(value) in (bool, int, float, str):
1878 return
1879 if type(value) in (list, tuple):
1880 for index, item in enumerate(value):
1881 validate(item, f"{path}[{index}]")
1882 return
1883 if type(value) is dict:
1884 for key, item in value.items():
1885 if type(key) is not str:
1886 raise ValueError(
1887 f"{path} has non-string key {key!r}; metadata keys must be strings"
1888 )
1889 validate(item, f"{path}.{key}")
1890 return
1891 raise ValueError(
1892 f"{path} has unsupported type {type(value).__name__}; use only "
1893 "None, bool, int, float, str, lists, tuples, and string-keyed dicts"
1894 )
1896 validate(metadata, "metadata")
1899def _normalize_positions(positions: Optional[Sequence[int]], seq_len: int) -> List[int]:
1900 """Normalize negative chunk-local positions and raise before indexing."""
1901 if positions is None:
1902 return list(range(seq_len))
1903 normalized = [position + seq_len if position < 0 else position for position in positions]
1904 out_of_range = [position for position in normalized if not 0 <= position < seq_len]
1905 if out_of_range:
1906 raise ValueError(
1907 f"positions {out_of_range} out of range for an activation chunk of length {seq_len}"
1908 )
1909 return normalized
1912def _cached_on_device(
1913 tensor: torch.Tensor,
1914 cache: Dict[torch.device, torch.Tensor],
1915 device: Union[str, torch.device],
1916) -> torch.Tensor:
1917 """Cache a small fp32 intervention tensor on each activation device."""
1918 resolved_device = torch.device(device)
1919 local = cache.get(resolved_device)
1920 if local is None: 1920 ↛ 1923line 1920 didn't jump to line 1923 because the condition on line 1920 was always true
1921 local = tensor.to(device=resolved_device, dtype=torch.float32)
1922 cache[resolved_device] = local
1923 return local
1926def _unit_rows(vectors: torch.Tensor, *, layer: int) -> torch.Tensor:
1927 """Normalize intervention vectors and reject zero/non-finite directions."""
1928 vectors = vectors.float()
1929 norms = vectors.norm(dim=-1, keepdim=True)
1930 if (~torch.isfinite(norms) | (norms <= torch.finfo(torch.float32).eps)).any(): 1930 ↛ 1931line 1930 didn't jump to line 1931 because the condition on line 1930 was never true
1931 raise ValueError(f"lens vectors at layer {layer} contain a zero or non-finite direction")
1932 return vectors / norms
1935def _make_intervention_hook(
1936 transform: Callable[[torch.Tensor], torch.Tensor],
1937 positions: Optional[Sequence[int]],
1938 d_model: int,
1939) -> Callable[..., torch.Tensor]:
1940 """Apply a transform with shared position, dtype, and device hardening."""
1941 requested = None if positions is None else tuple(positions)
1942 if requested == (): 1942 ↛ anywhereline 1942 didn't jump anywhere: it always raised an exception.
1943 raise ValueError("positions must contain at least one index")
1945 def hook_fn(
1946 activation: Float[torch.Tensor, "batch pos d_model"], hook: Any
1947 ) -> Float[torch.Tensor, "batch pos d_model"]:
1948 hook_name = getattr(hook, "name", "intervention hook")
1949 _validate_residual_activation(activation, d_model=d_model, hook_name=hook_name)
1950 normalized = _normalize_positions(requested, activation.shape[1])
1951 selected = activation if requested is None else activation[:, normalized, :]
1952 transformed = transform(selected)
1953 if transformed.shape != selected.shape: 1953 ↛ 1954line 1953 didn't jump to line 1954 because the condition on line 1953 was never true
1954 raise ValueError(
1955 f"intervention returned shape {tuple(transformed.shape)}, "
1956 f"expected {tuple(selected.shape)}"
1957 )
1958 transformed = transformed.to(device=activation.device, dtype=activation.dtype)
1959 if requested is None:
1960 return transformed
1961 output = activation.clone()
1962 output[:, normalized, :] = transformed
1963 return output
1965 return hook_fn
1968def _unembed(
1969 model: Any, residual: Float[torch.Tensor, "pos d_model"]
1970) -> Float[torch.Tensor, "pos d_vocab"]:
1971 """Apply the model's own final norm, unembedding, logit scale, and soft cap.
1973 The norm/unembed components contract on ``[batch, pos, d_model]``, so the
1974 position rows are passed through with a singleton batch axis. The residual
1975 is moved to the unembedding's device for sharded/multi-GPU models.
1976 """
1977 unembed_weight = model.W_U
1978 compute_dtype = unembed_weight.dtype
1979 batched = residual.to(device=unembed_weight.device, dtype=compute_dtype).unsqueeze(0)
1980 logits = model.unembed(model.ln_final(batched)).squeeze(0)
1981 return model.adapter.apply_output_logits_transform(logits).float()
1984def _to_token_ids(model: Any, tokens: Union[TokenInput, Sequence[TokenInput]]) -> List[int]:
1985 """Convert token strings / ids into a list of single-token ids."""
1986 if isinstance(tokens, (str, int)):
1987 tokens = [tokens]
1988 ids: List[int] = []
1989 for token in tokens:
1990 if isinstance(token, str):
1991 ids.append(model.to_single_token(token))
1992 elif isinstance(token, bool):
1993 raise ValueError(f"token {token!r} must be a string or integer token id, not bool")
1994 else:
1995 ids.append(int(token))
1996 if not ids:
1997 raise ValueError("tokens must contain at least one token")
1998 d_vocab = model.W_U.shape[1]
1999 invalid = [token_id for token_id in ids if not 0 <= token_id < d_vocab]
2000 if invalid:
2001 raise ValueError(f"token ids {invalid} out of range for vocabulary size {d_vocab}")
2002 return ids
2005def _validate_residual_activation(
2006 activation: torch.Tensor,
2007 *,
2008 d_model: int,
2009 hook_name: str,
2010) -> None:
2011 """Fail before interpreting a non-standard block output as a residual stream."""
2012 if activation.ndim != 3 or activation.shape[-1] != d_model:
2013 raise ValueError(
2014 f"{hook_name} must have shape [batch, position, {d_model}], "
2015 f"got {tuple(activation.shape)}"
2016 )
2019def _normalize_layer(layer: int, n_layers: int) -> int:
2020 """Resolve negative layer indices and bounds-check."""
2021 resolved = layer + n_layers if layer < 0 else layer
2022 if not 0 <= resolved < n_layers:
2023 raise ValueError(f"layer {layer} out of range for a {n_layers}-layer model")
2024 return resolved
2027class _frozen_parameters:
2028 """Context manager: freeze all parameters, restore their flags on exit.
2030 Freezing keeps the autograd graph rooted at the residual stream rather than
2031 at the weights, so fitting retains only the blocks between the earliest
2032 source layer and the target layer.
2033 """
2035 def __init__(self, model: Any) -> None:
2036 self.model = model
2037 self.saved: List[Tuple[torch.nn.Parameter, bool]] = []
2039 def __enter__(self) -> None:
2040 self.saved = [(param, param.requires_grad) for param in self.model.parameters()]
2041 for param, _ in self.saved:
2042 param.requires_grad_(False)
2044 def __exit__(self, *exc: Any) -> None:
2045 for param, flag in self.saved:
2046 param.requires_grad_(flag)
2049def _ordinary_vjp(
2050 target: torch.Tensor,
2051 sources: List[torch.Tensor],
2052 cotangent: torch.Tensor,
2053 retain_graph: bool,
2054) -> Tuple[torch.Tensor, ...]:
2055 """Ordinary vector-Jacobian product backing the J-lens estimator.
2057 Thin wrapper over ``torch.autograd.grad`` so the drive loop takes its
2058 backward step through the :data:`BackwardProvider` seam. This is the
2059 identity-preserving provider: routing the ordinary fit through it changes no
2060 numerics.
2061 """
2062 return torch.autograd.grad(
2063 outputs=target,
2064 inputs=sources,
2065 grad_outputs=cotangent,
2066 retain_graph=retain_graph,
2067 )
2070def _jacobian_for_prompt(
2071 model: Any,
2072 tokens: Int[torch.Tensor, "one seq"],
2073 *,
2074 source_layers: List[int],
2075 dim_batch: int,
2076 skip_first_positions: int,
2077 backward_provider: BackwardProvider,
2078) -> Dict[int, Float[torch.Tensor, "d_model d_model"]]:
2079 """Exact per-prompt Jacobian rows via batched one-hot cotangents.
2081 Assumes parameters are already frozen (see :class:`_frozen_parameters`) so
2082 that marking the earliest source activation ``requires_grad`` roots the
2083 graph there.
2085 The backward pass is taken through ``backward_provider`` rather than calling
2086 ``torch.autograd.grad`` directly, so the capture / cotangent-batching /
2087 averaging mechanics are shared by any estimator. The ordinary J-lens path
2088 passes :func:`_ordinary_vjp`, which reproduces the original numerics exactly.
2089 """
2090 d_model = model.cfg.d_model
2091 target_layer = model.cfg.n_layers - 1
2092 seq_len = tokens.shape[1]
2093 valid_positions = list(range(skip_first_positions, seq_len - 1))
2094 replicated = tokens.expand(dim_batch, -1)
2096 captured: Dict[str, torch.Tensor] = {}
2097 root_name = _resid_post_hook_name(min(source_layers))
2098 hook_layers = sorted(set(source_layers) | {target_layer})
2100 def capture_fn(
2101 activation: Float[torch.Tensor, "batch pos d_model"], hook: Any
2102 ) -> Float[torch.Tensor, "batch pos d_model"]:
2103 _validate_residual_activation(activation, d_model=d_model, hook_name=hook.name)
2104 if hook.name == root_name and not activation.requires_grad:
2105 activation.requires_grad_(True)
2106 captured[hook.name] = activation
2107 return activation
2109 fwd_hooks = [(_resid_post_hook_name(layer), capture_fn) for layer in hook_layers]
2110 with torch.enable_grad(), model.hooks(fwd_hooks=fwd_hooks):
2111 model(replicated, return_type=None)
2113 target = captured[_resid_post_hook_name(target_layer)]
2114 sources = [captured[_resid_post_hook_name(layer)] for layer in source_layers]
2115 device = target.device
2116 positions_index = torch.tensor(valid_positions, device=device)
2117 batch_index = torch.arange(dim_batch, device=device)
2119 jacobians = {
2120 layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in source_layers
2121 }
2122 cotangent = torch.zeros_like(target)
2123 n_passes = -(-d_model // dim_batch) # ceil division
2124 for pass_index in range(n_passes):
2125 dim_start = pass_index * dim_batch
2126 n_dims = min(dim_batch, d_model - dim_start)
2127 cotangent.zero_()
2128 cotangent[
2129 batch_index[:n_dims, None],
2130 positions_index[None, :],
2131 dim_start + batch_index[:n_dims, None],
2132 ] = 1.0
2133 grads = backward_provider(
2134 target,
2135 sources,
2136 cotangent,
2137 pass_index < n_passes - 1,
2138 )
2139 for layer, grad in zip(source_layers, grads):
2140 # each gradient lives on its layer's device under sharded/device_map setups
2141 rows = grad[:n_dims, positions_index.to(grad.device), :].float().mean(dim=1)
2142 jacobians[layer][dim_start : dim_start + n_dims, :] = rows.cpu()
2143 del grads
2144 return jacobians
2147def _fit_transport_matrices(
2148 model: Any,
2149 prompts: Sequence[str],
2150 *,
2151 source_layers: List[int],
2152 dim_batch: int,
2153 max_seq_len: int,
2154 skip_first_positions: int,
2155 show_progress: bool,
2156 backward_provider: BackwardProvider,
2157) -> Tuple[Dict[int, Float[torch.Tensor, "d_model d_model"]], int]:
2158 """Estimator-independent J-lens drive loop.
2160 Owns the mechanics shared by every estimator: the frozen-parameter
2161 lifecycle, the per-prompt forward/backward accumulation (source/target
2162 residual capture, one-hot cotangent batching, ``dim_batch`` chunking,
2163 valid-position selection, and source-position averaging all live in
2164 :func:`_jacobian_for_prompt`), the prompt-accumulation running sum, and the
2165 final division into per-prompt means. Only the backward step varies: it is
2166 taken through ``backward_provider``, so an alternate estimator reuses this
2167 loop unchanged. Callers own input validation, provenance, and lens
2168 construction.
2170 Args:
2171 model: A raw ``TransformerBridge`` whose parameters are frozen for the
2172 duration of the loop.
2173 prompts: Prompt strings; prompts too short to contain a valid position
2174 (``seq_len <= skip_first_positions + 1``) are skipped with a warning
2175 and do not count toward the returned prompt total.
2176 source_layers: Resolved, in-range source layers to fit.
2177 dim_batch: Output dimensions per backward pass.
2178 max_seq_len: Prompts are truncated to this many tokens.
2179 skip_first_positions: Leading positions excluded from the source average.
2180 show_progress: Show a tqdm progress bar over prompts.
2181 backward_provider: Backward-step seam; pass :func:`_ordinary_vjp` for the
2182 ordinary J-lens numerics.
2184 Returns:
2185 ``(transport_matrices, n_prompts)`` where ``transport_matrices`` maps each
2186 source layer to its prompt-averaged ``[d_model, d_model]`` matrix and
2187 ``n_prompts`` is the number of prompts that contributed.
2189 Raises:
2190 ValueError: If no prompt was long enough to contribute valid positions.
2191 """
2192 d_model = model.cfg.d_model
2193 jacobian_sum = {
2194 layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in source_layers
2195 }
2196 n_done = 0
2197 iterator = tqdm(prompts, desc="fitting J-lens", disable=not show_progress)
2198 with _frozen_parameters(model):
2199 for prompt in iterator:
2200 tokens = model.to_tokens(prompt)[:, :max_seq_len]
2201 seq_len = tokens.shape[1]
2202 if seq_len <= skip_first_positions + 1:
2203 warnings.warn(
2204 f"skipping prompt with only {seq_len} tokens "
2205 f"(need > {skip_first_positions + 1})",
2206 stacklevel=3,
2207 )
2208 continue
2209 per_prompt = _jacobian_for_prompt(
2210 model,
2211 tokens,
2212 source_layers=source_layers,
2213 dim_batch=dim_batch,
2214 skip_first_positions=skip_first_positions,
2215 backward_provider=backward_provider,
2216 )
2217 for layer in source_layers:
2218 jacobian_sum[layer] += per_prompt[layer]
2219 n_done += 1
2220 if n_done == 0:
2221 raise ValueError(
2222 "every prompt was too short to contribute valid positions; nothing was fitted"
2223 )
2224 means = {layer: jacobian_sum[layer] / n_done for layer in source_layers}
2225 return means, n_done