Coverage for transformer_lens/tools/analysis/jacobian_lens.py: 94%
611 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Jacobian Lens (J-lens).
3The Jacobian lens characterizes an intermediate residual-stream activation by its
4first-order causal effect on the model's output, averaged over a corpus of contexts.
5For each layer :math:`\\ell` it fits a single :math:`d_{model} \\times d_{model}` matrix
7.. math::
9 J_\\ell = \\mathbb{E}_{\\text{prompt}}\\left[
10 \\frac{1}{|V|}\\sum_{t \\in V}\\sum_{t' \\in V,\\,t' \\geq t}
11 \\frac{\\partial h_{\\text{final},t'}}{\\partial h_{\\ell,t}}
12 \\right]
14where ``V`` is the set of valid positions after the configured leading skip
15and final-position exclusion. Target-position effects are summed for each
16source; only source positions and prompts are averaged.
18mapping the output of block :math:`\\ell` to the final block's output (pre final
19norm). Reading the lens applies the model's own final norm and unembedding:
20:math:`\\text{lens}(h_\\ell) = W_U\\,\\mathrm{norm}(J_\\ell h_\\ell)`.
21The logit lens is the special case :math:`J_\\ell = I`. The rows of
22:math:`W_U J_\\ell` ("J-lens vectors") are residual-stream directions associated
23with single vocabulary tokens, and support causal interventions: steering,
24ablation, and exchanging one concept for another via a pseudoinverse coordinate
25swap.
27Introduced in `Verbalizable Representations Form a Global Workspace in Language
28Models <https://transformer-circuits.pub/2026/workspace/index.html>`_ (Gurnee et
29al., Transformer Circuits Thread, 2026). The fitting estimator and the artifact
30format follow Anthropic's Apache-2.0 reference implementation
31(`anthropics/jacobian-lens <https://github.com/anthropics/jacobian-lens>`_), so
32lenses fitted here interoperate with artifacts published on the Hugging Face Hub
33(e.g. `neuronpedia/jacobian-lens
34<https://huggingface.co/neuronpedia/jacobian-lens>`_); the interventions are
35implemented from the paper's Methods section.
37Warning:
38 Published lens artifacts are fitted on **raw** HuggingFace activations.
39 Jacobian lens supports only a freshly booted
40 ``TransformerBridge.boot_transformers`` model, whose weights are raw by
41 default. Compatibility mode and direct ``process_weights`` calls change the
42 residual basis and are refused rather than returning silently wrong
43 readouts. The model must also be a causal decoder whose adapter supports
44 text generation, use single-stream block outputs, and expose the standard
45 direct ``ln_final -> d_model-width unembed`` output path.
47Example::
49 import torch
50 from transformer_lens.model_bridge import TransformerBridge
51 from transformer_lens.tools.analysis import JacobianLens
53 model = TransformerBridge.boot_transformers("gpt2", device="cpu")
54 lens = JacobianLens.from_pretrained(
55 "neuronpedia/jacobian-lens",
56 filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
57 model=model,
58 )
59 result = lens.readout(model, "The Eiffel Tower is in the city of")
60 print(result.top_tokens(model.tokenizer, k=5)[8][-1]) # layer 8, final position
61"""
63import math
64import warnings
65from dataclasses import dataclass
66from importlib.metadata import version
67from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
69import torch
70from jaxtyping import Float, Int
71from tqdm.auto import tqdm
73from transformer_lens.tools.analysis.jacobian_lens_decomposition import (
74 DEFAULT_K,
75 JSpaceDecomposition,
76 JSpaceOccupancy,
77 JSpaceVarianceProfile,
78 estimate_occupancy,
79 get_sparse_decomposition,
80)
81from transformer_lens.utilities.hf_utils import call_hf_with_retry
83TokenInput = Union[str, int]
85# ---------------------------------------------------------------------------
86# Registry helpers
87# ---------------------------------------------------------------------------
89_REGISTRY_CACHE: Optional[Dict[str, Any]] = None
92def _load_registry() -> Dict[str, Any]:
93 """Return the bundled artifact registry, loading it once on first call."""
94 global _REGISTRY_CACHE
95 if _REGISTRY_CACHE is None:
96 import json
97 import pathlib
99 registry_path = pathlib.Path(__file__).with_name("jacobian_lens_registry.json")
100 with registry_path.open() as fh:
101 _REGISTRY_CACHE = json.load(fh)
102 return _REGISTRY_CACHE
105def _resolve_registry_entry(name_or_path: str) -> Optional[Tuple[str, str]]:
106 """Return ``(repo_id, filename)`` if *name_or_path* matches the registry.
108 Matching is tried in two passes:
109 1. Direct key match against short model names (e.g. ``"gemma-2-2b"``).
110 2. Alias match against Hugging Face model IDs (e.g. ``"google/gemma-2-2b"``).
112 Returns ``None`` when there is no match so callers can fall through to the
113 generic Hub download path.
114 """
115 registry = _load_registry()
116 if name_or_path in registry:
117 entry = registry[name_or_path]
118 return entry["repo_id"], entry["filename"]
119 for entry in registry.values():
120 if isinstance(entry, dict) and name_or_path in entry.get("aliases", []):
121 return entry["repo_id"], entry["filename"]
122 return None
125# Fitting excludes early positions (attention sinks with atypical residual statistics)
126# and the final position (no next-token target), matching the reference implementation.
127DEFAULT_SKIP_FIRST_POSITIONS = 16
128DEFAULT_TOP_K = 10
129_SWAP_WARN_COSINE = 0.99
130_SWAP_ERROR_COSINE = 0.999
132# Keys written by fit() that must not appear in converted-lens metadata so that
133# merge() can refuse to mix TL-fitted lenses with externally converted ones.
134# Note: "target_layer" is intentionally NOT listed here — it must survive
135# conversion so that validate_model() can detect and refuse checkpoints that
136# were fitted against a non-final target layer.
137_FIT_RESERVED_KEYS: frozenset = frozenset(
138 {
139 "transformer_lens_fit",
140 "transformer_lens_version",
141 "model_system",
142 "processing",
143 "hook_convention",
144 "fit_dtype",
145 "dim_batch",
146 "max_seq_len",
147 "skip_first_positions",
148 }
149)
151# Top-level payload keys that some checkpoint writers store as flat provenance
152# (rather than nested under a "metadata" key).
153_CHECKPOINT_FLAT_PROVENANCE: frozenset = frozenset({"model_name", "model_revision", "corpus"})
156@dataclass
157class JacobianLensReadout:
158 """Result of a :meth:`JacobianLens.readout` call.
160 Attributes:
161 lens_topk_values:
162 Per-layer retained top-k pre-softmax values, on CPU.
163 lens_topk_indices:
164 Per-layer retained top-k vocabulary ids, on CPU.
165 model_topk_values:
166 The model output's retained top-k pre-softmax values, on CPU.
167 model_topk_indices:
168 The model output's retained top-k vocabulary ids, on CPU.
169 lens_logits:
170 Optional full per-layer logits, on CPU. Present only when
171 ``readout(return_full_logits=True)`` was requested.
172 model_logits:
173 Optional full model logits, on CPU. Present only when
174 ``readout(return_full_logits=True)`` was requested.
175 tokens:
176 The token ids of the run prompt, ``[seq]``.
177 positions:
178 The (normalized, non-negative) positions the readout covers, aligned
179 with the ``pos`` axis of retained top-k and optional full logits.
180 use_jacobian:
181 Whether the Jacobian transport was applied (``False`` = logit lens).
182 """
184 lens_topk_values: Dict[int, Float[torch.Tensor, "pos k"]]
185 lens_topk_indices: Dict[int, Int[torch.Tensor, "pos k"]]
186 model_topk_values: Float[torch.Tensor, "pos k"]
187 model_topk_indices: Int[torch.Tensor, "pos k"]
188 tokens: Int[torch.Tensor, "seq"]
189 positions: List[int]
190 use_jacobian: bool = True
191 lens_logits: Optional[Dict[int, Float[torch.Tensor, "pos d_vocab"]]] = None
192 model_logits: Optional[Float[torch.Tensor, "pos d_vocab"]] = None
194 def top_tokens(self, tokenizer: Any, k: int = 5) -> Dict[int, List[List[str]]]:
195 """Decode the top-``k`` tokens per layer and position.
197 Args:
198 tokenizer: The model's tokenizer (``model.tokenizer``).
199 k: Number of top tokens to decode per (layer, position).
201 Returns:
202 ``{layer: [ [top-k strings] per position ]}``, positions aligned with
203 :attr:`positions`.
204 """
205 out: Dict[int, List[List[str]]] = {}
206 retained = next(iter(self.lens_topk_indices.values())).shape[-1]
207 if not 1 <= k <= retained:
208 raise ValueError(f"k must be between 1 and the retained top_k={retained}, got {k}")
209 for layer, ids in self.lens_topk_indices.items():
210 out[layer] = [[tokenizer.decode([t]) for t in row[:k].tolist()] for row in ids]
211 return out
214class JacobianLens:
215 """A fitted Jacobian lens: one transport matrix per source layer.
217 Layer convention (matching the reference implementation and the published
218 artifacts): index ``l`` refers to the **output of block** ``l`` at the
219 Bridge-native hook ``blocks.{l}.hook_out``. ``J[l]`` maps that activation
220 to the final block's output, pre final norm. The final layer itself is never
221 fitted (its transport is the identity), so
222 ``source_layers == [0, ..., n_layers - 2]`` for a full fit.
224 Attributes:
225 jacobians: ``{layer: [d_model, d_model]}`` transport matrices, fp32, CPU.
226 n_prompts: Number of prompts averaged into the fit.
227 d_model: Residual stream width the lens was fitted for.
228 metadata: Optional provenance (model name, TransformerLens version, fit
229 hyperparameters). Preserved by :meth:`save`/:meth:`load`; artifacts
230 from the reference implementation load with empty metadata.
231 """
233 def __init__(
234 self,
235 jacobians: Dict[int, Float[torch.Tensor, "d_model d_model"]],
236 *,
237 n_prompts: int,
238 d_model: int,
239 metadata: Optional[Dict[str, Any]] = None,
240 ) -> None:
241 if not jacobians: 241 ↛ 242line 241 didn't jump to line 242 because the condition on line 241 was never true
242 raise ValueError("jacobians must contain at least one layer")
243 for layer, matrix in jacobians.items():
244 if matrix.shape != (d_model, d_model): 244 ↛ 245line 244 didn't jump to line 245 because the condition on line 244 was never true
245 raise ValueError(
246 f"jacobians[{layer}] has shape {tuple(matrix.shape)}, "
247 f"expected ({d_model}, {d_model})"
248 )
249 self.jacobians: Dict[int, torch.Tensor] = {
250 int(layer): matrix.detach().float().cpu() for layer, matrix in jacobians.items()
251 }
252 self.n_prompts = int(n_prompts)
253 self.d_model = int(d_model)
254 self.metadata: Dict[str, Any] = dict(metadata or {})
255 self._device_jacobians: Dict[Tuple[int, torch.device], torch.Tensor] = {}
256 self._dictionary_cache: Dict[Tuple[int, torch.device], torch.Tensor] = {}
258 @property
259 def source_layers(self) -> List[int]:
260 """Sorted list of layers this lens has transport matrices for."""
261 return sorted(self.jacobians)
263 def __repr__(self) -> str:
264 layers = self.source_layers
265 return (
266 f"JacobianLens(layers={layers[0]}..{layers[-1]} ({len(layers)}), "
267 f"d_model={self.d_model}, n_prompts={self.n_prompts})"
268 )
270 # ------------------------------------------------------------------ #
271 # persistence #
272 # ------------------------------------------------------------------ #
274 def save(self, path: str, *, dtype: torch.dtype = torch.float16) -> None:
275 """Save the lens in the reference implementation's artifact format.
277 The four official keys (``J``, ``n_prompts``, ``source_layers``,
278 ``d_model``) are written unchanged so the file stays loadable by the
279 reference package; TransformerLens provenance is stored under an
280 additive ``metadata`` key.
282 Args:
283 path: Destination ``.pt`` path.
284 dtype: Storage dtype. Defaults to fp16 like the reference
285 implementation — Jacobian entries are order-one, so the smaller
286 dtype costs little precision and halves the artifact on disk.
287 """
288 _validate_metadata(self.metadata)
289 payload: Dict[str, Any] = {
290 "J": {layer: matrix.to(dtype) for layer, matrix in self.jacobians.items()},
291 "n_prompts": self.n_prompts,
292 "source_layers": self.source_layers,
293 "d_model": self.d_model,
294 }
295 if self.metadata:
296 payload["metadata"] = self.metadata
297 torch.save(payload, path)
299 @classmethod
300 def load(cls, path: str) -> "JacobianLens":
301 """Load a lens artifact or fit checkpoint saved in a supported schema.
303 Two file schemas are accepted:
305 **Artifact** (the reference format, written by :meth:`save` or the
306 Anthropic reference package): must contain a ``J`` key mapping layer
307 indices to transport matrices, plus ``n_prompts``, ``d_model``, and an
308 optional ``metadata`` dict.
310 **Fit checkpoint** (running-sum format, written by the reference
311 implementation's ``write_checkpoint()`` during fitting): must contain
312 a ``jacobian_sum`` key mapping layer indices to *running-sum* matrices
313 (i.e. the sum over prompts, not yet divided by the prompt count), plus
314 ``n_done``. ``d_model`` is inferred from the first matrix's shape;
315 no explicit ``d_model`` key is required or expected. The per-layer
316 means are reconstructed on load. A
317 ``converted_from: "jacobian_lens_checkpoint"`` key is added to
318 metadata so :meth:`merge` refuses to silently combine checkpoints with
319 natively TL-fitted lenses. Fit-reserved provenance keys
320 (``transformer_lens_fit``, etc.) are stripped; scalar fields that can
321 be serialised under ``weights_only=True`` are preserved.
322 Tensor-valued metadata fields that would fail :func:`_validate_metadata`
323 are recorded by name in a ``dropped_fields`` list.
325 Fit checkpoint schema (reference ``write_checkpoint()`` format)
326 ---------------------------------------------------------------
327 The reference implementation writes exactly six top-level keys; all
328 other keys in the payload are ignored::
330 {
331 "jacobian_sum": {<layer int>: <float32 tensor [d, d]>, ...},
332 "n_done": <int>, # prompts accumulated into jacobian_sum
333 "next_idx": <int>, # next prompt index (informational)
334 "source_layers": [<int>, ...], # documented layer indices (informational)
335 "target_layer": <int>, # target layer — harvested into metadata
336 "skip_first": <int>, # leading positions skipped (informational)
337 # optional flat provenance accepted from alternative checkpoint writers:
338 "model_name": <str>,
339 "model_revision": <str>,
340 "corpus": <str>,
341 # optional nested provenance accepted from alternative writers:
342 "metadata": {<str>: <scalar/list/dict>, ...},
343 }
345 Args:
346 path: Path to the ``.pt`` file.
348 Raises:
349 ValueError: If the file lacks both a ``J`` key (artifact) and a
350 ``jacobian_sum`` key (checkpoint), or if a checkpoint records a
351 non-positive ``n_prompts``.
352 """
353 payload = torch.load(path, map_location="cpu", weights_only=True)
354 if "J" in payload:
355 return cls(
356 {int(layer): matrix for layer, matrix in payload["J"].items()},
357 n_prompts=int(payload.get("n_prompts", 0)),
358 d_model=int(payload["d_model"]),
359 metadata=payload.get("metadata"),
360 )
361 if "jacobian_sum" in payload:
362 return cls._from_checkpoint_payload(path, payload)
363 raise ValueError(
364 f"{path} does not look like a Jacobian lens artifact or fit checkpoint. "
365 "Expected a 'J' key (artifact) or 'jacobian_sum' key (fit checkpoint). "
366 "See JacobianLens.load() for the supported file schemas."
367 )
369 @classmethod
370 def _from_checkpoint_payload(cls, path: str, payload: Dict[str, Any]) -> "JacobianLens":
371 """Reconstruct a JacobianLens from a fit-checkpoint payload.
373 Divides the running Jacobian sums by ``n_prompts``, strips fit-reserved
374 provenance keys, harvests safe scalar metadata, records dropped tensor
375 fields, and marks the result as converted so :meth:`merge` refuses to
376 mix it with natively TL-fitted lenses.
377 """
378 n_prompts = int(payload.get("n_done", payload.get("n_prompts", 0)))
379 if n_prompts <= 0:
380 raise ValueError(
381 f"{path} is a fit checkpoint with n_prompts={n_prompts}; "
382 "a positive prompt count is required to reconstruct the Jacobian mean."
383 )
384 if not payload["jacobian_sum"]:
385 raise ValueError(
386 f"{path} is a fit checkpoint with an empty jacobian_sum; "
387 "at least one layer matrix is required to reconstruct d_model."
388 )
389 first_matrix = next(iter(payload["jacobian_sum"].values()))
390 d_model = first_matrix.shape[0]
391 jacobians = {
392 int(layer): matrix.float() / n_prompts
393 for layer, matrix in payload["jacobian_sum"].items()
394 }
396 # Collect raw provenance: first from nested "metadata", then supplement
397 # with flat top-level keys that some checkpoint writers place directly
398 # in the payload (model_name, model_revision, corpus).
399 raw_meta: Dict[str, Any] = dict(payload.get("metadata") or {})
400 for key in _CHECKPOINT_FLAT_PROVENANCE:
401 if key in payload and key not in raw_meta:
402 raw_meta[key] = payload[key]
403 # target_layer lives at the top level in the reference checkpoint format;
404 # harvest it into metadata so validate_model() can check the fitting target.
405 if "target_layer" in payload and "target_layer" not in raw_meta:
406 raw_meta["target_layer"] = payload["target_layer"]
408 # Build clean metadata: drop fit-reserved keys, record tensor-valued
409 # fields that _validate_metadata would reject (they cannot survive a
410 # weights_only=True reload), and keep everything else that validates.
411 dropped_fields: List[str] = []
412 clean_meta: Dict[str, Any] = {}
413 for key, value in raw_meta.items():
414 if key in _FIT_RESERVED_KEYS:
415 continue
416 if isinstance(value, torch.Tensor):
417 dropped_fields.append(f"{key}: shape={tuple(value.shape)} dtype={value.dtype}")
418 continue
419 try:
420 _validate_metadata({key: value})
421 clean_meta[key] = value
422 except ValueError:
423 dropped_fields.append(key)
425 clean_meta["converted_from"] = "jacobian_lens_checkpoint"
426 if dropped_fields:
427 clean_meta["dropped_fields"] = dropped_fields
429 return cls(jacobians, n_prompts=n_prompts, d_model=d_model, metadata=clean_meta)
431 @classmethod
432 def from_pretrained(
433 cls,
434 name_or_path: str,
435 *,
436 filename: str = "lens.pt",
437 revision: Optional[str] = None,
438 model: Any = None,
439 ) -> "JacobianLens":
440 """Load a lens from a local path, a short model name, or a Hub repo.
442 Resolution order
443 ----------------
444 1. **Local file** — if *name_or_path* is an existing ``.pt`` file,
445 load it directly.
446 2. **Local directory** — if *name_or_path* is a directory, load
447 ``<name_or_path>/<filename>``.
448 3. **Registry short name or HF model ID** — if *name_or_path* matches
449 a key or alias in the bundled ``jacobian_lens_registry.json`` (e.g.
450 ``"gemma-2-2b"`` or ``"google/gemma-2-2b"``), the corresponding
451 artifact in ``neuronpedia/jacobian-lens`` is fetched automatically.
452 The *filename* argument is ignored in this case because the registry
453 already encodes the correct subpath.
454 4. **Explicit Hub repo** — otherwise *name_or_path* is treated as a Hub
455 repo id and *filename* is used as-is, preserving full backward
456 compatibility (e.g. ``from_pretrained("neuronpedia/jacobian-lens",
457 filename="gpt2-small/jlens/...")``).
459 Args:
460 name_or_path: A local ``.pt`` file, a local directory, a short
461 model name such as ``"gemma-2-2b"`` or ``"llama3.1-8b"``, a
462 Hugging Face model ID such as ``"google/gemma-2-2b"``, or an
463 explicit Hub repo id paired with *filename*.
464 filename: File (or subpath) inside a local directory or an explicit
465 Hub repo. Ignored when *name_or_path* resolves via the
466 registry.
467 revision: Optional Hub revision (branch, tag, or commit) to pin.
468 When omitted, the Hub repository's mutable default branch is
469 followed; pin a commit hash for reproducible analyses.
470 model: If given, :meth:`validate_model` is called so dimension or
471 weight-processing mismatches fail here rather than at first use.
473 Returns:
474 The loaded (and, if ``model`` was given, validated) lens.
476 Examples::
478 # Short model name — no need to remember the HF subpath
479 lens = JacobianLens.from_pretrained("gemma-2-2b", model=model)
481 # HF model ID also works
482 lens = JacobianLens.from_pretrained("google/gemma-2-2b", model=model)
484 # Explicit Hub repo + subpath (backward-compatible)
485 lens = JacobianLens.from_pretrained(
486 "neuronpedia/jacobian-lens",
487 filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
488 model=model,
489 )
490 """
491 import os
493 if os.path.isfile(name_or_path): 493 ↛ 494line 493 didn't jump to line 494 because the condition on line 493 was never true
494 lens = cls.load(name_or_path)
495 elif os.path.isdir(name_or_path): 495 ↛ 496line 495 didn't jump to line 496 because the condition on line 495 was never true
496 lens = cls.load(os.path.join(name_or_path, filename))
497 else:
498 from huggingface_hub import hf_hub_download
500 resolved = _resolve_registry_entry(name_or_path)
501 if resolved is not None:
502 repo_id, resolved_filename = resolved
503 else:
504 repo_id, resolved_filename = name_or_path, filename
506 local_path = call_hf_with_retry(
507 hf_hub_download,
508 repo_id=repo_id,
509 filename=resolved_filename,
510 revision=revision,
511 )
512 lens = cls.load(local_path)
513 if model is not None:
514 lens.validate_model(model)
515 return lens
517 @classmethod
518 def merge(cls, lenses: Sequence["JacobianLens"]) -> "JacobianLens":
519 """Combine lenses fitted on disjoint prompt slices.
521 The per-layer matrices are averaged weighted by each lens's
522 ``n_prompts``, matching the reference implementation, so fitting can be
523 parallelized across processes or machines and merged afterwards.
524 Provenance must match across shards (apart from ``n_prompts``), so a
525 merge cannot silently relabel matrices fitted with different models,
526 corpora, dtypes, or estimator settings. The merged count replaces the
527 per-shard count.
529 Args:
530 lenses: Lenses that agree exactly on ``source_layers`` and
531 ``d_model``.
533 Raises:
534 ValueError: On an empty sequence or mismatched lenses.
535 """
536 if not lenses:
537 raise ValueError("cannot merge an empty sequence of lenses")
538 invalid_counts = [
539 (index, lens.n_prompts) for index, lens in enumerate(lenses) if lens.n_prompts <= 0
540 ]
541 if invalid_counts:
542 raise ValueError(
543 "every lens passed to merge() must have positive n_prompts; "
544 f"invalid shards: {invalid_counts}"
545 )
546 first = lenses[0]
547 for lens in lenses:
548 _validate_metadata(lens.metadata)
549 first_provenance = {
550 key: value for key, value in first.metadata.items() if key != "n_prompts"
551 }
552 for other in lenses[1:]:
553 if other.source_layers != first.source_layers or other.d_model != first.d_model:
554 raise ValueError(
555 "all lenses being merged must share the same source_layers and d_model"
556 )
557 other_provenance = {
558 key: value for key, value in other.metadata.items() if key != "n_prompts"
559 }
560 if other_provenance != first_provenance:
561 raise ValueError(
562 "all lenses being merged must share the same provenance metadata "
563 "apart from n_prompts"
564 )
565 total = sum(lens.n_prompts for lens in lenses)
566 merged = {
567 layer: torch.stack([lens.jacobians[layer] * lens.n_prompts for lens in lenses]).sum(
568 dim=0
569 )
570 / total
571 for layer in first.source_layers
572 }
573 metadata = dict(first.metadata)
574 if metadata:
575 metadata["n_prompts"] = total
576 return cls(merged, n_prompts=total, d_model=first.d_model, metadata=metadata)
578 # ------------------------------------------------------------------ #
579 # model validation #
580 # ------------------------------------------------------------------ #
582 def validate_model(self, model: Any) -> "JacobianLens":
583 """Check that ``model`` matches this lens; raise loudly if not.
585 Requires a raw causal ``TransformerBridge`` with the standard direct
586 final-norm/unembed path, verifies recorded model provenance, residual
587 width and layer range, and enforces the published final-block target
588 convention.
590 Args:
591 model: A raw ``TransformerBridge``.
593 Returns:
594 ``self``, for chaining.
596 Raises:
597 TypeError: If model is not a ``TransformerBridge``.
598 ValueError: On model provenance or ``d_model`` mismatch,
599 out-of-range source layers, compatibility mode, unsupported
600 attention/output paths, or a non-final target convention.
601 """
602 _require_raw_bridge(model)
603 artifact_model_name = self.metadata.get("model_name")
604 current_model_name = getattr(model.cfg, "model_name", None)
605 if artifact_model_name is not None and artifact_model_name != current_model_name:
606 raise ValueError(
607 f"lens was fitted for model {artifact_model_name!r}, but the supplied "
608 f"model is {current_model_name!r}."
609 )
610 artifact_revision = self.metadata.get("model_revision")
611 current_revision = _get_model_revision(model)
612 if artifact_revision is not None and artifact_revision != current_revision:
613 raise ValueError(
614 f"lens was fitted for model revision {artifact_revision!r}, but the "
615 f"supplied model revision is {current_revision!r}."
616 )
617 d_model = model.cfg.d_model
618 if d_model != self.d_model:
619 raise ValueError(
620 f"lens was fitted for d_model={self.d_model}, but the model has "
621 f"d_model={d_model} — this lens belongs to a different model."
622 )
623 n_layers = model.cfg.n_layers
624 final_layer = n_layers - 1
625 out_of_range = [layer for layer in self.source_layers if not 0 <= layer < final_layer]
626 if out_of_range:
627 raise ValueError(
628 f"lens has source layers {out_of_range} outside the model's "
629 f"0..{final_layer - 1} source range — this lens belongs to a different model."
630 )
631 target_layer = int(self.metadata.get("target_layer", final_layer))
632 if target_layer != final_layer:
633 raise ValueError(
634 f"lens targets layer {target_layer}, but readout supports only the "
635 f"published final-layer convention ({final_layer}); refit without "
636 "a custom target layer."
637 )
638 return self
640 # ------------------------------------------------------------------ #
641 # reading #
642 # ------------------------------------------------------------------ #
644 def clear_device_cache(self) -> None:
645 """Release lazily cached Jacobian copies and full-vocabulary dictionaries on
646 accelerator devices."""
647 self._device_jacobians.clear()
648 self._dictionary_cache.clear()
650 def _matrix_on(self, layer: int, device: Union[str, torch.device]) -> torch.Tensor:
651 """Return one cached fp32 Jacobian copy for a layer/device pair."""
652 if layer not in self.jacobians:
653 raise ValueError(
654 f"layer {layer} is not in this lens's source layers "
655 f"({self.source_layers[0]}..{self.source_layers[-1]})"
656 )
657 resolved_device = torch.device(device)
658 key = (layer, resolved_device)
659 matrix = self._device_jacobians.get(key)
660 if matrix is None:
661 matrix = self.jacobians[layer].to(device=resolved_device, dtype=torch.float32)
662 self._device_jacobians[key] = matrix
663 return matrix
665 def transport(
666 self,
667 residual: Float[torch.Tensor, "... d_model"],
668 layer: int,
669 ) -> Float[torch.Tensor, "... d_model"]:
670 """Map layer-``layer`` activations into the final block's output basis.
672 Computes ``J[layer] @ h`` per activation vector, in fp32.
674 Args:
675 residual: Activations from the output of block ``layer``.
676 layer: Source layer index.
677 """
678 matrix = self._matrix_on(layer, residual.device)
679 return residual.float() @ matrix.T
681 @torch.no_grad()
682 def readout(
683 self,
684 model: Any,
685 input: Union[str, Int[torch.Tensor, "batch seq"]],
686 *,
687 layers: Optional[Sequence[int]] = None,
688 positions: Optional[Sequence[int]] = None,
689 use_jacobian: bool = True,
690 top_k: int = DEFAULT_TOP_K,
691 return_full_logits: bool = False,
692 ) -> JacobianLensReadout:
693 """Read per-layer vocabulary logits for a prompt.
695 Runs the model once with caching, transports the residual stream at each
696 requested layer through ``J[layer]`` (or the identity when
697 ``use_jacobian=False`` — the logit lens), and applies the model's own
698 final norm, unembedding, architecture logit scaling, and logit soft cap.
700 Args:
701 model: A raw ``TransformerBridge``.
702 input: A prompt string, or a ``[1, seq]`` token tensor.
703 layers: Layers to read. Defaults to every fitted layer plus the
704 final layer. The final layer (``n_layers - 1``) is always read
705 with the identity transport — by construction its lens equals
706 the model's own output distribution.
707 positions: Token positions to read (negative indices allowed).
708 Defaults to all positions.
709 use_jacobian: Apply the Jacobian transport. ``False`` gives the
710 logit-lens baseline through the identical code path.
711 top_k: Number of values and vocabulary ids retained per layer and
712 position. Defaults to 10.
713 return_full_logits: Also retain full vocabulary tensors on CPU.
714 This is opt-in because a 64-token Gemma readout across all
715 layers is roughly 1.7 GB.
717 Returns:
718 A :class:`JacobianLensReadout`.
720 Raises:
721 ValueError: If the model fails :meth:`validate_model`, ``input`` is
722 batched, ``top_k`` is invalid, or a requested layer has no
723 transport matrix.
724 """
725 self.validate_model(model)
726 tokens = model.to_tokens(input) if isinstance(input, str) else input
727 if tokens.ndim != 2 or tokens.shape[0] != 1: 727 ↛ 728line 727 didn't jump to line 728 because the condition on line 727 was never true
728 raise ValueError(f"readout expects a single prompt; got shape {tuple(tokens.shape)}")
729 n_layers = model.cfg.n_layers
730 final_layer = n_layers - 1
731 if layers is None:
732 layers = self.source_layers + [final_layer]
733 layers = [_normalize_layer(layer, n_layers) for layer in layers]
734 for layer in layers:
735 if use_jacobian and layer != final_layer and layer not in self.jacobians: 735 ↛ 736line 735 didn't jump to line 736 because the condition on line 735 was never true
736 raise ValueError(
737 f"layer {layer} is not in this lens's source layers; "
738 f"available: {self.source_layers} (+{final_layer} as identity)"
739 )
741 if top_k < 1: 741 ↛ 742line 741 didn't jump to line 742 because the condition on line 741 was never true
742 raise ValueError(f"top_k must be at least 1, got {top_k}")
743 seq_len = tokens.shape[1]
744 norm_positions = _normalize_positions(positions, seq_len)
746 hook_names = {
747 layer: _resid_post_hook_name(layer) for layer in layers if layer != final_layer
748 }
749 wanted = set(hook_names.values())
750 logits, cache = model.run_with_cache(tokens, names_filter=lambda name: name in wanted)
751 selected_model_logits = logits[0, norm_positions, :].float()
752 if top_k > selected_model_logits.shape[-1]: 752 ↛ 753line 752 didn't jump to line 753 because the condition on line 752 was never true
753 raise ValueError(
754 f"top_k={top_k} exceeds the model vocabulary size "
755 f"{selected_model_logits.shape[-1]}"
756 )
757 model_topk = selected_model_logits.topk(top_k, dim=-1)
758 full_model_logits = selected_model_logits.cpu() if return_full_logits else None
759 lens_topk_values: Dict[int, torch.Tensor] = {}
760 lens_topk_indices: Dict[int, torch.Tensor] = {}
761 full_lens_logits: Optional[Dict[int, torch.Tensor]] = {} if return_full_logits else None
762 for layer in layers:
763 if layer == final_layer:
764 layer_logits = selected_model_logits
765 layer_topk = model_topk
766 else:
767 activation = cache[hook_names[layer]]
768 _validate_residual_activation(
769 activation,
770 d_model=model.cfg.d_model,
771 hook_name=hook_names[layer],
772 )
773 residual = activation[0, norm_positions, :]
774 transported = self.transport(residual, layer) if use_jacobian else residual.float()
775 layer_logits = _unembed(model, transported)
776 layer_topk = layer_logits.topk(top_k, dim=-1)
777 lens_topk_values[layer] = layer_topk.values.cpu()
778 lens_topk_indices[layer] = layer_topk.indices.cpu()
779 if full_lens_logits is not None:
780 if layer == final_layer:
781 assert full_model_logits is not None
782 full_lens_logits[layer] = full_model_logits
783 else:
784 full_lens_logits[layer] = layer_logits.cpu()
785 return JacobianLensReadout(
786 lens_topk_values=lens_topk_values,
787 lens_topk_indices=lens_topk_indices,
788 model_topk_values=model_topk.values.cpu(),
789 model_topk_indices=model_topk.indices.cpu(),
790 tokens=tokens[0].cpu(),
791 positions=norm_positions,
792 use_jacobian=use_jacobian,
793 lens_logits=full_lens_logits,
794 model_logits=full_model_logits,
795 )
797 @torch.no_grad()
798 def lens_vectors(
799 self,
800 model: Any,
801 tokens: Union[TokenInput, Sequence[TokenInput]],
802 layer: int,
803 ) -> Float[torch.Tensor, "n d_model"]:
804 """Residual-stream directions for vocabulary tokens at a layer.
806 The J-lens vector for token ``t`` is row ``t`` of ``W_U J[layer]``
807 expressed in layer-``layer`` residual coordinates:
808 ``v_t = J[layer]^T W_U[:, t]``.
810 Args:
811 model: The model supplying ``W_U``.
812 tokens: A token string / id, or a sequence of them. Strings must
813 encode to a single token.
814 layer: Source layer for the vectors.
816 Returns:
817 One vector per token, fp32, on the model's device.
818 """
819 self.validate_model(model)
820 layer = _normalize_layer(layer, model.cfg.n_layers)
821 token_ids = _to_token_ids(model, tokens)
822 unembed_columns = model.W_U[:, token_ids].float() # [d_model, n]
823 matrix = self._matrix_on(layer, unembed_columns.device)
824 return (matrix.T @ unembed_columns).T
826 @torch.no_grad()
827 def lens_vector_dictionary(
828 self, model: Any, layer: int
829 ) -> Float[torch.Tensor, "d_vocab d_model"]:
830 """Full-vocabulary J-lens dictionary at ``layer``: ``[d_vocab, d_model]``.
832 Row ``t`` is the J-lens vector ``v_t = J[layer]^T W_U[:, t]`` -- this is
833 :meth:`lens_vectors` over the entire vocabulary. The result is cached per
834 (layer, device) so a sparse decomposition can reuse it; :meth:`clear_device_cache`
835 releases it.
837 The dictionary is vocabulary-sized and cached on the model's device
838 (``d_vocab * d_model`` fp32 values, on the order of gigabytes for a large
839 vocabulary), one entry per requested layer.
841 Args:
842 model: The model supplying ``W_U``.
843 layer: Source layer for the dictionary (must be a fitted source layer).
845 Returns:
846 The dictionary, fp32, on the model's device.
847 """
848 self.validate_model(model)
849 layer = _normalize_layer(layer, model.cfg.n_layers)
850 device = torch.device(model.W_U.device)
851 cached = self._dictionary_cache.get((layer, device))
852 if cached is None:
853 matrix = self._matrix_on(layer, device) # [d_model, d_model]
854 cached = (matrix.T @ model.W_U.float()).T # [d_vocab, d_model]
855 self._dictionary_cache[(layer, device)] = cached
856 return cached
858 @torch.no_grad()
859 def decompose(
860 self,
861 model: Any,
862 activation_or_prompt: Union[torch.Tensor, str],
863 layer: int,
864 *,
865 position: Optional[int] = None,
866 k: int = DEFAULT_K,
867 algorithm: str = "nonnegative_orthogonal_matching_pursuit",
868 ) -> JSpaceDecomposition:
869 """Decompose an activation into its J-space content at ``layer``.
871 ``activation_or_prompt`` is either:
873 - a raw activation vector of shape ``[d_model]`` (leave ``position`` as ``None``), or
874 - a prompt -- a string or a ``[1, seq]`` token tensor -- in which case ``position``
875 selects the token whose ``blocks.{layer}.hook_out`` activation is decomposed.
877 The full-vocabulary dictionary at ``layer`` is built (and cached) via
878 :meth:`lens_vector_dictionary`, then :func:`get_sparse_decomposition` solves for a
879 ``k``-sparse nonnegative combination of J-lens vectors.
881 Args:
882 model: A raw ``TransformerBridge``.
883 activation_or_prompt: An activation vector, or a prompt (string / token tensor).
884 layer: Source layer (must be a fitted source layer).
885 position: Token position when a prompt is given; must be ``None`` for a raw
886 activation vector.
887 k: Upper bound on the number of J-lens vectors to select; selection stops early
888 once no unselected vector is materially positively correlated, so fewer may be
889 returned.
890 algorithm: Coefficient-update rule; see :func:`get_sparse_decomposition`.
892 Returns:
893 A :class:`JSpaceDecomposition`. Its ``support`` (token ids here) holds only the
894 numerically active J-lens vectors and ``selected_support`` every selected vector,
895 with ``support.numel() <= selected_support.numel() <= k``. ``support`` and its
896 token-decoding tensors are on CPU; the vector-valued outputs stay on the model's
897 device.
899 Raises:
900 ValueError: On an invalid model, a mismatched activation shape, a batched prompt,
901 an unfitted layer, or an invalid ``k`` / ``algorithm``.
902 RuntimeError: If the default nonnegative least-squares solver cannot validate its
903 result against the KKT conditions.
904 """
905 activation, resolved_layer = self._resolve_activation(
906 model, activation_or_prompt, layer, position
907 )
908 dictionary = self.lens_vector_dictionary(model, resolved_layer)
909 return get_sparse_decomposition(
910 activation.float().to(dictionary.device), dictionary, k, algorithm=algorithm
911 )
913 def _resolve_activation(
914 self,
915 model: Any,
916 activation_or_prompt: Union[torch.Tensor, str],
917 layer: int,
918 position: Optional[int],
919 ) -> Tuple[torch.Tensor, int]:
920 """Validate the model and layer, then resolve either a raw ``[d_model]`` activation or a
921 prompt plus ``position`` to the activation vector to analyse.
923 Returns ``(activation, resolved_layer)``. Shared by :meth:`decompose` and
924 :meth:`occupancy` so both accept the same input forms with identical validation.
925 """
926 self.validate_model(model)
927 resolved_layer = _normalize_layer(layer, model.cfg.n_layers)
928 if resolved_layer not in self.jacobians:
929 raise ValueError(
930 f"layer {layer} is not in this lens's source layers; "
931 f"available: {self.source_layers}"
932 )
934 if position is None:
935 if not isinstance(activation_or_prompt, torch.Tensor):
936 raise ValueError(
937 "decompose expects a raw activation tensor when position is None; pass a "
938 "prompt together with a position to decompose a model activation"
939 )
940 activation = activation_or_prompt
941 if activation.ndim != 1 or activation.shape[0] != self.d_model:
942 raise ValueError(
943 f"activation must be 1-D of length d_model={self.d_model}, "
944 f"got shape {tuple(activation.shape)}"
945 )
946 else:
947 if isinstance(activation_or_prompt, torch.Tensor) and activation_or_prompt.ndim == 1:
948 raise ValueError(
949 "position is only valid with a prompt (string or [1, seq] tokens); "
950 "pass a raw 1-D activation with position=None instead"
951 )
952 tokens = (
953 model.to_tokens(activation_or_prompt)
954 if isinstance(activation_or_prompt, str)
955 else activation_or_prompt
956 )
957 if tokens.ndim != 2 or tokens.shape[0] != 1:
958 raise ValueError(
959 f"decompose expects a single prompt; got shape {tuple(tokens.shape)}"
960 )
961 hook_name = _resid_post_hook_name(resolved_layer)
962 _, cache = model.run_with_cache(tokens, names_filter=lambda name: name == hook_name)
963 norm_position = _normalize_positions([position], tokens.shape[1])[0]
964 activation = cache[hook_name][0, norm_position, :]
965 return activation, resolved_layer
967 @torch.no_grad()
968 def occupancy(
969 self,
970 model: Any,
971 activation_or_prompt: Union[torch.Tensor, str],
972 layer: int,
973 *,
974 position: Optional[int] = None,
975 max_atoms: int = DEFAULT_K,
976 num_control_dictionaries: int = 32,
977 seed: int = 0,
978 ) -> JSpaceOccupancy:
979 """Estimate how many J-lens vectors are meaningfully active in an activation at ``layer``.
981 Resolves ``activation_or_prompt`` (a raw ``[d_model]`` vector, or a prompt plus
982 ``position``) exactly as :meth:`decompose`, builds the cached full-vocabulary dictionary
983 via :meth:`lens_vector_dictionary`, and calls :func:`estimate_occupancy`.
985 Args:
986 model: A raw ``TransformerBridge``.
987 activation_or_prompt: An activation vector, or a prompt (string / token tensor).
988 layer: Source layer (must be a fitted source layer).
989 position: Token position when a prompt is given; ``None`` for a raw activation.
990 max_atoms: Maximum number of J-lens vectors to consider.
991 num_control_dictionaries: Number of random control dictionaries to average over.
992 seed: Seed for the random control dictionaries (reproducibility).
994 Returns:
995 A :class:`JSpaceOccupancy`.
996 """
997 activation, resolved_layer = self._resolve_activation(
998 model, activation_or_prompt, layer, position
999 )
1000 dictionary = self.lens_vector_dictionary(model, resolved_layer)
1001 return estimate_occupancy(
1002 activation.float().to(dictionary.device),
1003 dictionary,
1004 max_atoms=max_atoms,
1005 num_control_dictionaries=num_control_dictionaries,
1006 seed=seed,
1007 )
1009 @torch.no_grad()
1010 def fraction_of_variance(
1011 self,
1012 model: Any,
1013 prompts: Union[str, torch.Tensor, Sequence[Union[str, torch.Tensor]]],
1014 layers: Optional[Sequence[int]] = None,
1015 *,
1016 k: int = DEFAULT_K,
1017 skip_first: int = 16,
1018 positions: Optional[Sequence[int]] = None,
1019 show_progress: bool = False,
1020 ) -> JSpaceVarianceProfile:
1021 """Profile the J-space share of activation variance over a prompt corpus.
1023 Each prompt is run once (caching ``blocks.{layer}.hook_out`` for every requested layer).
1024 At each sampled position the activation is decomposed and its J-space variance fraction
1025 ``||j_space_component||^2 / ||activation||^2`` is recorded. The numerator is the
1026 ``j_space_component`` -- the orthogonal projection of the activation onto the span of the
1027 selected support (the paper's appendix "J-space component"), *not* the nonnegative
1028 ``reconstruction``; the two coincide only when every selected atom stays active. Per layer
1029 the profile reports the median of those fractions and the pooled ratio
1030 ``sum(||j_space_component||^2) / sum(||activation||^2)`` (the paper's "fraction of total
1031 variance").
1033 A layer that samples no positions -- every prompt shorter than ``skip_first``, or only
1034 zero-norm activations -- contributes no fractions: its ``median`` and ``pooled`` are
1035 ``float("nan")`` and its ``per_position`` tensor is empty.
1037 Args:
1038 model: A raw ``TransformerBridge``.
1039 prompts: A prompt, or a sequence of prompts. Each token tensor must represent exactly
1040 one prompt and have shape ``[1, seq]``.
1041 layers: Source layers to profile; defaults to all fitted ``source_layers``.
1042 k: Number of J-lens vectors per decomposition.
1043 skip_first: Non-negative index before which positions are skipped (mirrors the fit's
1044 early-position skip); not used for sampling when ``positions`` is given.
1045 positions: Explicit positions to sample instead of ``skip_first`` onward.
1046 show_progress: Show a tqdm progress bar over prompts.
1048 Returns:
1049 A :class:`JSpaceVarianceProfile`.
1051 Raises:
1052 ValueError: On an invalid model, an unfitted layer, an empty corpus, a negative
1053 ``skip_first``, or a token tensor that does not have shape ``[1, seq]``.
1054 """
1055 self.validate_model(model)
1056 if skip_first < 0:
1057 raise ValueError(f"skip_first must be non-negative, got {skip_first}")
1058 if layers is None:
1059 resolved_layers = list(self.source_layers)
1060 else:
1061 resolved_layers = [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]
1062 for layer in resolved_layers:
1063 if layer not in self.jacobians:
1064 raise ValueError(
1065 f"layer {layer} is not in this lens's source layers; "
1066 f"available: {self.source_layers}"
1067 )
1068 prompt_list: List[Union[str, torch.Tensor]] = (
1069 [prompts] if isinstance(prompts, (str, torch.Tensor)) else list(prompts)
1070 )
1071 if not prompt_list:
1072 raise ValueError("prompts must be a non-empty prompt or sequence of prompts")
1074 hook_names = {layer: _resid_post_hook_name(layer) for layer in resolved_layers}
1075 wanted_hooks = set(hook_names.values())
1076 dictionaries = {
1077 layer: self.lens_vector_dictionary(model, layer) for layer in resolved_layers
1078 }
1079 fractions: Dict[int, List[float]] = {layer: [] for layer in resolved_layers}
1080 pooled_j_space: Dict[int, float] = {layer: 0.0 for layer in resolved_layers}
1081 pooled_total: Dict[int, float] = {layer: 0.0 for layer in resolved_layers}
1083 for prompt in tqdm(prompt_list, desc="J-space variance", disable=not show_progress):
1084 tokens = model.to_tokens(prompt) if isinstance(prompt, str) else prompt
1085 if tokens.ndim != 2 or tokens.shape[0] != 1:
1086 raise ValueError(
1087 "fraction_of_variance expects each tokenized prompt to have shape "
1088 f"[1, seq], got {tuple(tokens.shape)}"
1089 )
1090 _, cache = model.run_with_cache(tokens, names_filter=lambda name: name in wanted_hooks)
1091 seq_len = tokens.shape[1]
1092 sampled = (
1093 list(range(skip_first, seq_len))
1094 if positions is None
1095 else _normalize_positions(positions, seq_len)
1096 )
1097 for layer in resolved_layers:
1098 dictionary = dictionaries[layer]
1099 activations = cache[hook_names[layer]][0] # [seq, d_model]
1100 for position in sampled:
1101 activation = activations[position].float().to(dictionary.device)
1102 total = float(activation @ activation)
1103 if total <= 0.0: 1103 ↛ 1104line 1103 didn't jump to line 1104 because the condition on line 1103 was never true
1104 continue
1105 decomposition = get_sparse_decomposition(activation, dictionary, k)
1106 j_space = float(
1107 decomposition.j_space_component @ decomposition.j_space_component
1108 )
1109 fractions[layer].append(j_space / total)
1110 pooled_j_space[layer] += j_space
1111 pooled_total[layer] += total
1113 median = {
1114 layer: float(torch.tensor(fractions[layer]).median())
1115 if fractions[layer]
1116 else float("nan")
1117 for layer in resolved_layers
1118 }
1119 pooled = {
1120 layer: pooled_j_space[layer] / pooled_total[layer]
1121 if pooled_total[layer] > 0
1122 else float("nan")
1123 for layer in resolved_layers
1124 }
1125 per_position = {layer: torch.tensor(fractions[layer]) for layer in resolved_layers}
1126 return JSpaceVarianceProfile(
1127 layers=resolved_layers, median=median, pooled=pooled, per_position=per_position
1128 )
1130 # ------------------------------------------------------------------ #
1131 # interventions #
1132 # ------------------------------------------------------------------ #
1134 def steering_hooks(
1135 self,
1136 model: Any,
1137 token: TokenInput,
1138 layers: Sequence[int],
1139 *,
1140 alpha: float = 4.0,
1141 positions: Optional[Sequence[int]] = None,
1142 ) -> List[Tuple[str, Any]]:
1143 """Hooks that steer the residual stream along a token's J-lens vector.
1145 At each layer the unit-normalized lens vector is added, scaled by
1146 ``alpha`` times the activation's **median** per-position residual norm:
1147 ``h <- h + alpha * median||h|| * v̂``. This norm-matched
1148 parameterization follows the steering description in the reference
1149 implementation's experiment protocols; the paper's minimal form is
1150 the unscaled ``h <- h + alpha * v_t``, recoverable by passing the raw
1151 :meth:`lens_vectors` output to your own hook. The median (not mean) is
1152 used so attention-sink positions — whose residual norms run orders of
1153 magnitude above typical positions — do not inflate the scale.
1155 Args:
1156 model: The model the hooks will run on.
1157 token: The concept token to steer toward.
1158 layers: Layers to intervene at.
1159 alpha: Steering strength scalar; ``0`` disables. Because of the
1160 norm-matched scale, values of order 1 already perturb the
1161 stream by roughly its own magnitude.
1162 positions: Chunk-local positions to steer (negative indices allowed
1163 and normalized on every hook invocation). Defaults to all.
1165 Returns:
1166 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)`` or
1167 ``model.run_with_hooks(fwd_hooks=...)``.
1168 """
1169 self.validate_model(model)
1170 hooks = []
1171 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
1172 direction = self.lens_vectors(model, token, layer)[0]
1173 unit = _unit_rows(direction.unsqueeze(0), layer=layer)[0]
1174 device_units: Dict[torch.device, torch.Tensor] = {}
1176 def transform(
1177 selected: Float[torch.Tensor, "batch pos d_model"],
1178 unit: torch.Tensor = unit,
1179 device_units: Dict[torch.device, torch.Tensor] = device_units,
1180 ) -> Float[torch.Tensor, "batch pos d_model"]:
1181 local_unit = _cached_on_device(unit, device_units, selected.device)
1182 scale = alpha * selected.float().norm(dim=-1).median()
1183 return selected.float() + scale * local_unit
1185 hooks.append(
1186 (
1187 _resid_post_hook_name(layer),
1188 _make_intervention_hook(transform, positions, model.cfg.d_model),
1189 )
1190 )
1191 return hooks
1193 def ablation_hooks(
1194 self,
1195 model: Any,
1196 tokens: Union[TokenInput, Sequence[TokenInput]],
1197 layers: Sequence[int],
1198 *,
1199 positions: Optional[Sequence[int]] = None,
1200 ) -> List[Tuple[str, Any]]:
1201 """Hooks that project token directions out of the residual stream.
1203 For each token's unit lens vector ``v̂``: ``h <- h - (h·v̂) v̂``,
1204 applied sequentially when several tokens are given.
1206 Args:
1207 model: The model the hooks will run on.
1208 tokens: Concept token(s) to suppress.
1209 layers: Layers to intervene at.
1210 positions: Chunk-local positions to ablate (negative indices allowed
1211 and normalized on every hook invocation). Defaults to all.
1213 Returns:
1214 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
1215 """
1216 self.validate_model(model)
1217 hooks = []
1218 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
1219 vectors = self.lens_vectors(model, tokens, layer)
1220 units = _unit_rows(vectors, layer=layer)
1221 device_units: Dict[torch.device, torch.Tensor] = {}
1223 def transform(
1224 selected: Float[torch.Tensor, "batch pos d_model"],
1225 units: torch.Tensor = units,
1226 device_units: Dict[torch.device, torch.Tensor] = device_units,
1227 ) -> Float[torch.Tensor, "batch pos d_model"]:
1228 local_units = _cached_on_device(units, device_units, selected.device)
1229 result = selected.float()
1230 for unit in local_units:
1231 coeff = result @ unit
1232 result = result - coeff.unsqueeze(-1) * unit
1233 return result
1235 hooks.append(
1236 (
1237 _resid_post_hook_name(layer),
1238 _make_intervention_hook(transform, positions, model.cfg.d_model),
1239 )
1240 )
1241 return hooks
1243 def swap_hooks(
1244 self,
1245 model: Any,
1246 source_token: TokenInput,
1247 target_token: TokenInput,
1248 layers: Sequence[int],
1249 *,
1250 alpha: float = 1.0,
1251 positions: Optional[Sequence[int]] = None,
1252 ) -> List[Tuple[str, Any]]:
1253 """Hooks that swap two concepts' coordinates in lens space.
1255 The paper's patching-in-lens-coordinates intervention: with
1256 ``V = [v_s, v_t]`` and lens coordinates ``c = V⁺ h`` (pseudoinverse),
1257 the update is ``h <- h + alpha * V (sigma(c) - c)`` where ``sigma``
1258 exchanges the two coordinates. The component of ``h`` orthogonal to
1259 ``span{v_s, v_t}`` is untouched. ``alpha=2`` is the paper's
1260 "double-strength" swap.
1262 Args:
1263 model: The model the hooks will run on.
1264 source_token: The concept to remove (e.g. ``" France"``).
1265 target_token: The concept to install (e.g. ``" China"``).
1266 layers: Layers to intervene at (the paper clamps the swap across an
1267 intermediate-layer band).
1268 alpha: Swap strength.
1269 positions: Chunk-local positions to swap (negative indices allowed
1270 and normalized on every hook invocation). Defaults to all.
1272 Returns:
1273 ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.
1274 """
1275 self.validate_model(model)
1276 source_id, target_id = _to_token_ids(model, [source_token, target_token])
1277 if source_id == target_id:
1278 raise ValueError(
1279 "source_token and target_token resolve to the same token id; "
1280 "a coordinate swap would be a silent no-op"
1281 )
1283 hooks = []
1284 for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
1285 vectors = self.lens_vectors(model, [source_id, target_id], layer)
1286 units = _unit_rows(vectors, layer=layer)
1287 cosine = abs(float((units[0] @ units[1]).item()))
1288 if not math.isfinite(cosine) or cosine >= _SWAP_ERROR_COSINE:
1289 raise ValueError(
1290 f"swap vectors at layer {layer} are numerically near-parallel "
1291 f"(abs cosine={cosine:.6f}); choose better-separated concepts"
1292 )
1293 if cosine >= _SWAP_WARN_COSINE:
1294 warnings.warn(
1295 f"swap vectors at layer {layer} are poorly conditioned "
1296 f"(abs cosine={cosine:.6f}); the intervention may be amplified",
1297 UserWarning,
1298 stacklevel=2,
1299 )
1300 basis = vectors.T # [d, 2]
1301 pinv = torch.linalg.pinv(basis) # [2, d]
1302 device_basis: Dict[torch.device, torch.Tensor] = {}
1303 device_pinv: Dict[torch.device, torch.Tensor] = {}
1305 def transform(
1306 selected: Float[torch.Tensor, "batch pos d_model"],
1307 basis: torch.Tensor = basis,
1308 pinv: torch.Tensor = pinv,
1309 device_basis: Dict[torch.device, torch.Tensor] = device_basis,
1310 device_pinv: Dict[torch.device, torch.Tensor] = device_pinv,
1311 ) -> Float[torch.Tensor, "batch pos d_model"]:
1312 local_basis = _cached_on_device(basis, device_basis, selected.device)
1313 local_pinv = _cached_on_device(pinv, device_pinv, selected.device)
1314 coords = selected.float() @ local_pinv.T # [..., 2]
1315 delta = alpha * ((coords[..., [1, 0]] - coords) @ local_basis.T)
1316 return selected.float() + delta
1318 hooks.append(
1319 (
1320 _resid_post_hook_name(layer),
1321 _make_intervention_hook(transform, positions, model.cfg.d_model),
1322 )
1323 )
1324 return hooks
1326 # ------------------------------------------------------------------ #
1327 # fitting #
1328 # ------------------------------------------------------------------ #
1330 @classmethod
1331 def fit(
1332 cls,
1333 model: Any,
1334 prompts: Sequence[str],
1335 *,
1336 corpus: str,
1337 source_layers: Optional[Sequence[int]] = None,
1338 dim_batch: int = 8,
1339 max_seq_len: int = 128,
1340 skip_first_positions: int = DEFAULT_SKIP_FIRST_POSITIONS,
1341 show_progress: bool = True,
1342 metadata: Optional[Dict[str, Any]] = None,
1343 ) -> "JacobianLens":
1344 """Fit a Jacobian lens on a hooked model.
1346 Implements the reference estimator exactly. For each prompt: one forward
1347 pass (the prompt replicated ``dim_batch`` times along the batch axis),
1348 then ``ceil(d_model / dim_batch)`` backward passes. Each backward plants
1349 a one-hot cotangent for one output dimension at *every* valid target
1350 position simultaneously — causal attention guarantees the gradient at
1351 source position ``t`` is then the sum over target positions
1352 ``t' >= t`` with no explicit masking. Rows are averaged over valid
1353 source positions (the first ``skip_first_positions`` and the final
1354 position are excluded), and prompts contribute equally to the final
1355 mean. There is no randomness: the computation is deterministic given
1356 the prompts.
1358 The reference implementation reports that fit quality saturates
1359 quickly — on the order of 100 prompts of 128 tokens is usable; the
1360 published lenses use up to 1000. Use :meth:`merge` to parallelize
1361 across prompt slices.
1363 Args:
1364 model: A raw ``TransformerBridge``. Model parameters are temporarily
1365 frozen (``requires_grad=False``) during fitting and restored
1366 after. Cotangents and activation gradients use the model dtype;
1367 fit with a float32 model for the highest-fidelity estimator.
1368 prompts: Prompt strings. Prompts too short to contain a valid
1369 position (``seq_len <= skip_first_positions + 1``) are skipped
1370 with a warning and do not count toward ``n_prompts``.
1371 corpus: Stable identifier for the prompt corpus or slice, recorded
1372 in artifact provenance.
1373 source_layers: Layers to fit. Defaults to every layer below
1374 the final layer. Negative indices count from ``n_layers``.
1375 dim_batch: Output dimensions per backward pass. Higher is faster
1376 but replicates the prompt ``dim_batch`` times in memory; total
1377 backward FLOPs are unchanged.
1378 max_seq_len: Prompts are truncated to this many tokens.
1379 skip_first_positions: Leading positions excluded from the source
1380 average.
1381 show_progress: Show a tqdm progress bar over prompts.
1382 metadata: Extra provenance merged into :attr:`metadata`.
1384 Returns:
1385 The fitted :class:`JacobianLens`.
1387 Raises:
1388 TypeError: If model is not a ``TransformerBridge``.
1389 ValueError: On compatibility mode, invalid provenance or layer
1390 indices, or if no prompt was long enough to fit on.
1391 """
1392 _require_raw_bridge(model)
1393 if not isinstance(corpus, str) or not corpus.strip(): 1393 ↛ 1394line 1393 didn't jump to line 1394 because the condition on line 1393 was never true
1394 raise ValueError("corpus must be a non-empty provenance identifier")
1395 n_layers = model.cfg.n_layers
1396 d_model = model.cfg.d_model
1397 resolved_target = n_layers - 1
1398 if source_layers is None:
1399 resolved_sources = list(range(resolved_target))
1400 else:
1401 resolved_sources = sorted(
1402 {_normalize_layer(layer, n_layers) for layer in source_layers}
1403 )
1404 if not resolved_sources: 1404 ↛ 1405line 1404 didn't jump to line 1405 because the condition on line 1404 was never true
1405 raise ValueError("source_layers is empty")
1406 if resolved_sources[-1] >= resolved_target: 1406 ↛ 1407line 1406 didn't jump to line 1407 because the condition on line 1406 was never true
1407 raise ValueError(
1408 f"every source layer must be below target_layer={resolved_target}; "
1409 f"got {resolved_sources}"
1410 )
1411 if dim_batch < 1: 1411 ↛ 1412line 1411 didn't jump to line 1412 because the condition on line 1411 was never true
1412 raise ValueError(f"dim_batch must be >= 1, got {dim_batch}")
1413 if skip_first_positions < 0:
1414 raise ValueError(f"skip_first_positions must be >= 0, got {skip_first_positions}")
1415 fit_dtype = model.W_U.dtype
1416 if fit_dtype in (torch.float16, torch.bfloat16):
1417 warnings.warn(
1418 f"fitting in {fit_dtype} accumulates Jacobian gradients at reduced "
1419 "precision; use a float32 TransformerBridge for the highest-fidelity fit",
1420 UserWarning,
1421 stacklevel=2,
1422 )
1424 jacobian_sum = {
1425 layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in resolved_sources
1426 }
1427 n_done = 0
1428 iterator = tqdm(prompts, desc="fitting J-lens", disable=not show_progress)
1429 with _frozen_parameters(model):
1430 for prompt in iterator:
1431 tokens = model.to_tokens(prompt)[:, :max_seq_len]
1432 seq_len = tokens.shape[1]
1433 if seq_len <= skip_first_positions + 1:
1434 warnings.warn(
1435 f"skipping prompt with only {seq_len} tokens "
1436 f"(need > {skip_first_positions + 1})",
1437 stacklevel=2,
1438 )
1439 continue
1440 per_prompt = _jacobian_for_prompt(
1441 model,
1442 tokens,
1443 source_layers=resolved_sources,
1444 dim_batch=dim_batch,
1445 skip_first_positions=skip_first_positions,
1446 )
1447 for layer in resolved_sources:
1448 jacobian_sum[layer] += per_prompt[layer]
1449 n_done += 1
1450 if n_done == 0:
1451 raise ValueError(
1452 "every prompt was too short to contribute valid positions; nothing was fitted"
1453 )
1455 fit_metadata: Dict[str, Any] = {
1456 "model_name": getattr(model.cfg, "model_name", None),
1457 "model_revision": _get_model_revision(model),
1458 "transformer_lens_version": version("transformer-lens"),
1459 "model_system": "TransformerBridge",
1460 "processing": {
1461 "compatibility_mode": False,
1462 "weight_basis": "raw_huggingface",
1463 },
1464 "hook_convention": "blocks.{layer}.hook_out",
1465 "corpus": corpus,
1466 "n_prompts": n_done,
1467 "fit_dtype": str(fit_dtype).removeprefix("torch."),
1468 "target_layer": resolved_target,
1469 "dim_batch": dim_batch,
1470 "max_seq_len": max_seq_len,
1471 "skip_first_positions": skip_first_positions,
1472 "transformer_lens_fit": True,
1473 }
1474 reserved = sorted(set(fit_metadata).intersection(metadata or {}))
1475 if reserved: 1475 ↛ 1476line 1475 didn't jump to line 1476 because the condition on line 1475 was never true
1476 raise ValueError(f"metadata cannot override fit provenance keys: {reserved}")
1477 full_metadata = dict(metadata or {})
1478 full_metadata.update(fit_metadata)
1479 _validate_metadata(full_metadata)
1480 return cls(
1481 {layer: jacobian_sum[layer] / n_done for layer in resolved_sources},
1482 n_prompts=n_done,
1483 d_model=d_model,
1484 metadata=full_metadata,
1485 )
1488# ---------------------------------------------------------------------- #
1489# helpers #
1490# ---------------------------------------------------------------------- #
1493def _resid_post_hook_name(layer: int) -> str:
1494 """Bridge-native hook for the output of block ``layer``."""
1495 return f"blocks.{layer}.hook_out"
1498def _get_model_revision(model: Any) -> Optional[str]:
1499 """Return the resolved Hugging Face commit recorded on a booted model."""
1500 original_model = getattr(model, "original_model", None)
1501 hf_config = getattr(original_model, "config", None)
1502 revision = getattr(hf_config, "_commit_hash", None)
1503 return revision if isinstance(revision, str) and revision else None
1506def _require_raw_bridge(model: Any) -> None:
1507 """Require the causal raw-Bridge contract used by fit and readout."""
1508 from transformer_lens.model_bridge import TransformerBridge
1510 if not isinstance(model, TransformerBridge):
1511 raise TypeError(
1512 "JacobianLens supports TransformerBridge only; load a fresh model with "
1513 "TransformerBridge.boot_transformers(...)."
1514 )
1515 if getattr(model, "compatibility_mode", False):
1516 raise ValueError(
1517 "compatibility mode is enabled on this TransformerBridge and changes "
1518 "the residual basis. Use a freshly booted "
1519 "TransformerBridge.boot_transformers(...) model with raw weights."
1520 )
1521 if getattr(model, "_weights_processed", False):
1522 raise ValueError(
1523 "process_weights was called on this TransformerBridge and changed the "
1524 "raw HuggingFace weight basis. Use a freshly booted "
1525 "TransformerBridge.boot_transformers(...) model."
1526 )
1527 adapter = model.adapter
1528 if not adapter.supports_generation:
1529 raise ValueError(
1530 "JacobianLens requires a causal decoder-only Bridge whose adapter "
1531 f"supports text generation; {type(adapter).__name__} declares "
1532 "supports_generation=False."
1533 )
1534 adapter.validate_output_logits_transform()
1535 attention_dir = getattr(model.cfg, "attention_dir", "causal")
1536 if attention_dir != "causal":
1537 raise ValueError(
1538 "JacobianLens requires causal attention because its estimator relies on "
1539 "causality to exclude target positions before each source position; "
1540 f"got attention_dir={attention_dir!r}."
1541 )
1542 total_ut_steps = int(getattr(model.cfg, "total_ut_steps", 1) or 1)
1543 if total_ut_steps != 1:
1544 raise ValueError(
1545 "JacobianLens requires each physical block hook to fire once per forward; "
1546 f"this looped-depth Bridge runs total_ut_steps={total_ut_steps}."
1547 )
1548 component_mapping = adapter.get_component_mapping()
1549 required_components = ("blocks", "ln_final", "unembed")
1550 missing_components = [
1551 component
1552 for component in required_components
1553 if component not in component_mapping or not hasattr(model, component)
1554 ]
1555 if missing_components:
1556 raise ValueError(
1557 "JacobianLens requires the standard direct ln_final -> unembed output path; "
1558 f"this Bridge is missing {missing_components}."
1559 )
1560 blocks_component = component_mapping["blocks"]
1561 if not getattr(blocks_component, "hook_out_is_single_residual_stream", False):
1562 raise ValueError(
1563 "JacobianLens requires single-stream [batch, position, d_model] block "
1564 f"outputs; {type(blocks_component).__name__} does not provide that contract."
1565 )
1566 if "project_out" in component_mapping:
1567 raise ValueError(
1568 "JacobianLens does not yet support a final output projection between "
1569 "the residual stream and unembedding."
1570 )
1571 unembed_width = model.W_U.shape[0]
1572 if unembed_width != model.cfg.d_model:
1573 raise ValueError(
1574 "JacobianLens requires a direct d_model-width unembedding after ln_final; "
1575 f"got W_U input width {unembed_width} for d_model={model.cfg.d_model}. "
1576 "Architectures with a final output projection are not yet supported."
1577 )
1580def _validate_metadata(metadata: Dict[str, Any]) -> None:
1581 """Reject values that ``torch.load(weights_only=True)`` cannot reload."""
1583 def validate(value: Any, path: str) -> None:
1584 if value is None or type(value) in (bool, int, float, str):
1585 return
1586 if type(value) in (list, tuple):
1587 for index, item in enumerate(value):
1588 validate(item, f"{path}[{index}]")
1589 return
1590 if type(value) is dict:
1591 for key, item in value.items():
1592 if type(key) is not str:
1593 raise ValueError(
1594 f"{path} has non-string key {key!r}; metadata keys must be strings"
1595 )
1596 validate(item, f"{path}.{key}")
1597 return
1598 raise ValueError(
1599 f"{path} has unsupported type {type(value).__name__}; use only "
1600 "None, bool, int, float, str, lists, tuples, and string-keyed dicts"
1601 )
1603 validate(metadata, "metadata")
1606def _normalize_positions(positions: Optional[Sequence[int]], seq_len: int) -> List[int]:
1607 """Normalize negative chunk-local positions and raise before indexing."""
1608 if positions is None:
1609 return list(range(seq_len))
1610 normalized = [position + seq_len if position < 0 else position for position in positions]
1611 out_of_range = [position for position in normalized if not 0 <= position < seq_len]
1612 if out_of_range:
1613 raise ValueError(
1614 f"positions {out_of_range} out of range for an activation chunk of length {seq_len}"
1615 )
1616 return normalized
1619def _cached_on_device(
1620 tensor: torch.Tensor,
1621 cache: Dict[torch.device, torch.Tensor],
1622 device: Union[str, torch.device],
1623) -> torch.Tensor:
1624 """Cache a small fp32 intervention tensor on each activation device."""
1625 resolved_device = torch.device(device)
1626 local = cache.get(resolved_device)
1627 if local is None: 1627 ↛ 1630line 1627 didn't jump to line 1630 because the condition on line 1627 was always true
1628 local = tensor.to(device=resolved_device, dtype=torch.float32)
1629 cache[resolved_device] = local
1630 return local
1633def _unit_rows(vectors: torch.Tensor, *, layer: int) -> torch.Tensor:
1634 """Normalize intervention vectors and reject zero/non-finite directions."""
1635 vectors = vectors.float()
1636 norms = vectors.norm(dim=-1, keepdim=True)
1637 if (~torch.isfinite(norms) | (norms <= torch.finfo(torch.float32).eps)).any(): 1637 ↛ 1638line 1637 didn't jump to line 1638 because the condition on line 1637 was never true
1638 raise ValueError(f"lens vectors at layer {layer} contain a zero or non-finite direction")
1639 return vectors / norms
1642def _make_intervention_hook(
1643 transform: Callable[[torch.Tensor], torch.Tensor],
1644 positions: Optional[Sequence[int]],
1645 d_model: int,
1646) -> Callable[..., torch.Tensor]:
1647 """Apply a transform with shared position, dtype, and device hardening."""
1648 requested = None if positions is None else tuple(positions)
1649 if requested == (): 1649 ↛ 1650line 1649 didn't jump to line 1650 because the condition on line 1649 was never true
1650 raise ValueError("positions must contain at least one index")
1652 def hook_fn(
1653 activation: Float[torch.Tensor, "batch pos d_model"], hook: Any
1654 ) -> Float[torch.Tensor, "batch pos d_model"]:
1655 hook_name = getattr(hook, "name", "intervention hook")
1656 _validate_residual_activation(activation, d_model=d_model, hook_name=hook_name)
1657 normalized = _normalize_positions(requested, activation.shape[1])
1658 selected = activation if requested is None else activation[:, normalized, :]
1659 transformed = transform(selected)
1660 if transformed.shape != selected.shape: 1660 ↛ 1661line 1660 didn't jump to line 1661 because the condition on line 1660 was never true
1661 raise ValueError(
1662 f"intervention returned shape {tuple(transformed.shape)}, "
1663 f"expected {tuple(selected.shape)}"
1664 )
1665 transformed = transformed.to(device=activation.device, dtype=activation.dtype)
1666 if requested is None:
1667 return transformed
1668 output = activation.clone()
1669 output[:, normalized, :] = transformed
1670 return output
1672 return hook_fn
1675def _unembed(
1676 model: Any, residual: Float[torch.Tensor, "pos d_model"]
1677) -> Float[torch.Tensor, "pos d_vocab"]:
1678 """Apply the model's own final norm, unembedding, logit scale, and soft cap.
1680 The norm/unembed components contract on ``[batch, pos, d_model]``, so the
1681 position rows are passed through with a singleton batch axis. The residual
1682 is moved to the unembedding's device for sharded/multi-GPU models.
1683 """
1684 unembed_weight = model.W_U
1685 compute_dtype = unembed_weight.dtype
1686 batched = residual.to(device=unembed_weight.device, dtype=compute_dtype).unsqueeze(0)
1687 logits = model.unembed(model.ln_final(batched)).squeeze(0)
1688 return model.adapter.apply_output_logits_transform(logits).float()
1691def _to_token_ids(model: Any, tokens: Union[TokenInput, Sequence[TokenInput]]) -> List[int]:
1692 """Convert token strings / ids into a list of single-token ids."""
1693 if isinstance(tokens, (str, int)):
1694 tokens = [tokens]
1695 ids: List[int] = []
1696 for token in tokens:
1697 if isinstance(token, str):
1698 ids.append(model.to_single_token(token))
1699 else:
1700 ids.append(int(token))
1701 if not ids:
1702 raise ValueError("tokens must contain at least one token")
1703 d_vocab = model.W_U.shape[1]
1704 invalid = [token_id for token_id in ids if not 0 <= token_id < d_vocab]
1705 if invalid:
1706 raise ValueError(f"token ids {invalid} out of range for vocabulary size {d_vocab}")
1707 return ids
1710def _validate_residual_activation(
1711 activation: torch.Tensor,
1712 *,
1713 d_model: int,
1714 hook_name: str,
1715) -> None:
1716 """Fail before interpreting a non-standard block output as a residual stream."""
1717 if activation.ndim != 3 or activation.shape[-1] != d_model:
1718 raise ValueError(
1719 f"{hook_name} must have shape [batch, position, {d_model}], "
1720 f"got {tuple(activation.shape)}"
1721 )
1724def _normalize_layer(layer: int, n_layers: int) -> int:
1725 """Resolve negative layer indices and bounds-check."""
1726 resolved = layer + n_layers if layer < 0 else layer
1727 if not 0 <= resolved < n_layers:
1728 raise ValueError(f"layer {layer} out of range for a {n_layers}-layer model")
1729 return resolved
1732class _frozen_parameters:
1733 """Context manager: freeze all parameters, restore their flags on exit.
1735 Freezing keeps the autograd graph rooted at the residual stream rather than
1736 at the weights, so fitting retains only the blocks between the earliest
1737 source layer and the target layer.
1738 """
1740 def __init__(self, model: Any) -> None:
1741 self.model = model
1742 self.saved: List[Tuple[torch.nn.Parameter, bool]] = []
1744 def __enter__(self) -> None:
1745 self.saved = [(param, param.requires_grad) for param in self.model.parameters()]
1746 for param, _ in self.saved:
1747 param.requires_grad_(False)
1749 def __exit__(self, *exc: Any) -> None:
1750 for param, flag in self.saved:
1751 param.requires_grad_(flag)
1754def _jacobian_for_prompt(
1755 model: Any,
1756 tokens: Int[torch.Tensor, "one seq"],
1757 *,
1758 source_layers: List[int],
1759 dim_batch: int,
1760 skip_first_positions: int,
1761) -> Dict[int, Float[torch.Tensor, "d_model d_model"]]:
1762 """Exact per-prompt Jacobian rows via batched one-hot cotangents.
1764 Assumes parameters are already frozen (see :class:`_frozen_parameters`) so
1765 that marking the earliest source activation ``requires_grad`` roots the
1766 graph there.
1767 """
1768 d_model = model.cfg.d_model
1769 target_layer = model.cfg.n_layers - 1
1770 seq_len = tokens.shape[1]
1771 valid_positions = list(range(skip_first_positions, seq_len - 1))
1772 replicated = tokens.expand(dim_batch, -1)
1774 captured: Dict[str, torch.Tensor] = {}
1775 root_name = _resid_post_hook_name(min(source_layers))
1776 hook_layers = sorted(set(source_layers) | {target_layer})
1778 def capture_fn(
1779 activation: Float[torch.Tensor, "batch pos d_model"], hook: Any
1780 ) -> Float[torch.Tensor, "batch pos d_model"]:
1781 _validate_residual_activation(activation, d_model=d_model, hook_name=hook.name)
1782 if hook.name == root_name and not activation.requires_grad:
1783 activation.requires_grad_(True)
1784 captured[hook.name] = activation
1785 return activation
1787 fwd_hooks = [(_resid_post_hook_name(layer), capture_fn) for layer in hook_layers]
1788 with torch.enable_grad(), model.hooks(fwd_hooks=fwd_hooks):
1789 model(replicated, return_type=None)
1791 target = captured[_resid_post_hook_name(target_layer)]
1792 sources = [captured[_resid_post_hook_name(layer)] for layer in source_layers]
1793 device = target.device
1794 positions_index = torch.tensor(valid_positions, device=device)
1795 batch_index = torch.arange(dim_batch, device=device)
1797 jacobians = {
1798 layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in source_layers
1799 }
1800 cotangent = torch.zeros_like(target)
1801 n_passes = -(-d_model // dim_batch) # ceil division
1802 for pass_index in range(n_passes):
1803 dim_start = pass_index * dim_batch
1804 n_dims = min(dim_batch, d_model - dim_start)
1805 cotangent.zero_()
1806 cotangent[
1807 batch_index[:n_dims, None],
1808 positions_index[None, :],
1809 dim_start + batch_index[:n_dims, None],
1810 ] = 1.0
1811 grads = torch.autograd.grad(
1812 outputs=target,
1813 inputs=sources,
1814 grad_outputs=cotangent,
1815 retain_graph=pass_index < n_passes - 1,
1816 )
1817 for layer, grad in zip(source_layers, grads):
1818 # each gradient lives on its layer's device under sharded/device_map setups
1819 rows = grad[:n_dims, positions_index.to(grad.device), :].float().mean(dim=1)
1820 jacobians[layer][dim_start : dim_start + n_dims, :] = rows.cpu()
1821 del grads
1822 return jacobians