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