Coverage for transformer_lens/tools/analysis/logit_lens.py: 95%
317 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Logit lens and vocabulary readout.
3The `logit lens <https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens>`_
4reads every intermediate residual-stream state through the model's own final
5norm and unembedding, so each layer can be asked "what would the model predict
6if it stopped here?". :func:`logit_lens` does that in one call over a cache's
7accumulated residual stream; :func:`logit_readout` is the kernel underneath it
8and reads *any* ``d_model`` vector (a head output, a steering direction, a
9sparse-coder decoder column) through the same path.
11Both go through the real ``ln_final`` module, ``W_U``, ``b_U`` and the
12adapter's post-unembedding transform (Gemma-style softcap, Cohere logit scale),
13so the final stack entry reproduces the model's own logits on a raw bridge and
14the model's log-probabilities in compatibility mode (``center_unembed`` shifts
15raw logits by a per-position constant). Logits are produced chunk by chunk, so
16``[layers, batch, pos, d_vocab]`` is never materialized unless the full
17vocabulary is explicitly requested.
19Example::
21 from transformer_lens.model_bridge import TransformerBridge
22 from transformer_lens.tools.analysis import logit_lens
24 model = TransformerBridge.boot_transformers("gpt2", device="cpu")
25 result = logit_lens(model, "The Eiffel Tower is in the city of", targets=" Paris")
26 for label, rank in zip(result.labels, result.rank_trajectory(" Paris").tolist()):
27 print(f"{label:>12}: rank {rank}")
29 # Rank of the actual next token at every position and layer: a per-row target.
30 tokens = model.to_tokens("The Eiffel Tower is in the city of Paris")
31 nxt = logit_lens(model, tokens[:, :-1], positions=None, targets=tokens[:, 1:])
32 nxt.rank_trajectory() # [layers, batch, pos]
33"""
35import itertools
36from dataclasses import dataclass
37from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
39import torch
40from jaxtyping import Float, Int
42from transformer_lens.ActivationCache import ActivationCache
43from transformer_lens.tools.analysis._model_state import require_eval_mode
44from transformer_lens.utilities import Slice, SliceInput
45from transformer_lens.utilities.logits_utils import logits_to_df
46from transformer_lens.utilities.quantization import require_readable_weight
48# Shared targets: one string / id or a sequence of them. A tensor means per-row
49# targets (see logit_readout).
50TokenInput = Union[str, int, Sequence[Union[str, int]], torch.Tensor]
52_RETURN_TYPES = ("logits", "log_probs", "probs")
53_INDEX_DTYPES = (torch.int64, torch.int32, torch.int16, torch.int8, torch.uint8)
54# Value-column names, matching utilities.logits_utils.logits_to_df.
55_VALUE_COLUMNS = {"logits": "logit", "log_probs": "log_prob", "probs": "probability"}
57# Normalization types for which the model has a final norm the lens must apply.
58_FINAL_NORM_TYPES = ("LN", "LNPre", "RMS", "RMSPre")
60# Block variants without a ``resid_mid`` hook; mirrors direct_logit_attribution.
61_HYBRID_VARIANT_NAMES = ("mamba", "ssm", "mixer", "linear_attn")
64@dataclass
65class LogitReadout:
66 """Result of :func:`logit_readout`.
68 Exactly one of the three selectors expands the trailing axis of ``values``:
69 ``targets`` gives ``[..., n_targets]`` (``[..., 1]`` for per-row targets),
70 ``top_k`` gives ``[..., k]`` with ``top_ids``, ``vocab`` gives
71 ``[..., len(vocab)]``; with none of them the trailing axis is the full
72 vocabulary. ``entropy`` and ``target_ranks`` are always computed over the
73 full vocabulary, whatever was selected.
75 Attributes:
76 values: The requested quantity (see ``return_type``) with the leading
77 dims of the input vectors and the selected trailing axis.
78 return_type: ``"logits"``, ``"log_probs"`` or ``"probs"``.
79 entropy: Full-vocabulary entropy of the softmax, in nats, ``[...]``.
80 top_ids: Token ids aligned with ``values`` when ``top_k`` was used.
81 vocab_ids: The requested subset ids when ``vocab`` was used.
82 target_ids: Shared targets as ``[n_targets]``; per-row targets as a
83 tensor with the leading dims of the vectors (one id per row).
84 target_ranks: Zero-based full-vocabulary rank of each target (0 is the
85 argmax; ties count only strictly larger logits), aligned with
86 ``values``.
87 per_row_targets: Whether ``targets`` was given per row (a tensor).
88 applied_ln: Whether the model's final norm was applied to the vectors.
89 applied_output_transform: Whether the adapter's post-unembedding
90 transform was invoked (it is the identity for most architectures).
91 """
93 values: Float[torch.Tensor, "... selected"]
94 return_type: str
95 entropy: Float[torch.Tensor, "..."]
96 top_ids: Optional[Int[torch.Tensor, "... k"]] = None
97 vocab_ids: Optional[Int[torch.Tensor, "v"]] = None
98 target_ids: Optional[torch.Tensor] = None
99 target_ranks: Optional[Int[torch.Tensor, "... t"]] = None
100 per_row_targets: bool = False
101 applied_ln: bool = True
102 applied_output_transform: bool = True
104 @property
105 def selected_ids(self) -> Optional[torch.Tensor]:
106 """Token ids along the trailing axis of ``values`` when it is not the full vocabulary."""
107 if self.target_ids is not None: 107 ↛ 108line 107 didn't jump to line 108 because the condition on line 107 was never true
108 return self.target_ids
109 if self.vocab_ids is not None:
110 return self.vocab_ids
111 return self.top_ids
113 def row_ids(self, index: Tuple[int, ...]) -> Optional[torch.Tensor]:
114 """Trailing-axis token ids for one row of ``values`` (``None`` = full vocabulary)."""
115 if self.top_ids is not None:
116 return self.top_ids[index]
117 if self.target_ids is not None and self.per_row_targets:
118 return self.target_ids[index].reshape(1)
119 return self.selected_ids
121 def decode(self, tokenizer: Any) -> Any:
122 """Decode the trailing-axis token ids to strings, keeping the leading structure.
124 ``top_k`` and per-row readouts decode per row; shared ``targets`` /
125 ``vocab`` readouts decode the id list once. Full-vocabulary readouts
126 have no id list to decode.
127 """
128 decode = getattr(tokenizer, "decode", None)
129 if not callable(decode): 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true
130 raise TypeError("tokenizer must provide a callable decode method")
131 if self.top_ids is not None:
132 return _decode_nested(self.top_ids, decode)
133 if self.target_ids is not None and self.per_row_targets: 133 ↛ 135line 133 didn't jump to line 135 because the condition on line 133 was always true
134 return _decode_nested(self.target_ids, decode)
135 ids = self.selected_ids
136 if ids is None:
137 raise ValueError(
138 "decode needs a targets, vocab or top_k readout; this one is full-vocabulary"
139 )
140 return [str(decode([int(i)])) for i in ids.tolist()]
143@dataclass
144class LogitLensResult:
145 """Result of :func:`logit_lens`.
147 Attributes:
148 readout: The :class:`LogitReadout` over the accumulated stack. Its
149 leading axis is the stack entry, aligned with ``labels``; the
150 remaining leading dims are ``[batch, pos]`` after the requested
151 ``positions`` / ``batch_slice`` (an integer selection drops its axis).
152 labels: Stack entry labels from ``accumulated_resid`` (``"0_pre"``,
153 ``"3_mid"``, ``"final_post"``).
154 has_batch_dim: Whether a batch axis follows the stack axis.
155 target_names: For shared targets, the string passed for each target
156 (``None`` where an id was passed), so :meth:`rank_trajectory`
157 accepts the strings again.
158 """
160 readout: LogitReadout
161 labels: List[str]
162 has_batch_dim: bool
163 target_names: Optional[List[Optional[str]]] = None
165 @property
166 def values(self) -> torch.Tensor:
167 """Shortcut for ``readout.values``."""
168 return self.readout.values
170 def rank_trajectory(self, target: Union[str, int, None] = None) -> torch.Tensor:
171 """Full-vocabulary rank of one target at every stack entry, ``[layer, ...]``.
173 ``target`` is the string passed to ``targets`` or a token id; ``None``
174 selects the only target when exactly one was requested, and is the only
175 form for per-row targets.
176 """
177 ranks = self.readout.target_ranks
178 ids = self.readout.target_ids
179 if ranks is None or ids is None: 179 ↛ 180line 179 didn't jump to line 180 because the condition on line 179 was never true
180 raise ValueError("rank_trajectory needs a lens run with targets=...")
181 if self.readout.per_row_targets:
182 if target is not None:
183 raise ValueError("per-row targets have one target per row; call rank_trajectory()")
184 return ranks[..., 0]
185 if target is None:
186 if ids.numel() != 1:
187 raise ValueError("pass target= when more than one target was requested")
188 return ranks[..., 0]
189 if isinstance(target, str):
190 names = self.target_names or []
191 if target not in names:
192 known = [n for n in names if n is not None]
193 raise ValueError(
194 f"unknown target {target!r}; string lookup needs the string as passed "
195 f"(known: {known}); pass the token id otherwise"
196 )
197 return ranks[..., names.index(target)]
198 matches = (ids == int(target)).nonzero(as_tuple=True)[0]
199 if matches.numel() == 0:
200 raise ValueError(f"token id {target} was not among the requested targets")
201 return ranks[..., int(matches[0])]
203 def top_tokens(self, tokenizer: Any) -> Dict[str, Any]:
204 """Decoded top-k tokens per stack entry, ``{label: nested lists}``."""
205 if self.readout.top_ids is None:
206 raise ValueError("top_tokens needs a lens run with top_k=...")
207 return {
208 label: _decode_nested(self.readout.top_ids[i], tokenizer.decode)
209 for i, label in enumerate(self.labels)
210 }
212 def to_dataframe(self, tokenizer: Optional[Any] = None, top_k: Optional[int] = None) -> Any:
213 """One row per (stack entry, batch, position, token) as a pandas DataFrame.
215 Rows within a (layer, batch, pos) group are sorted by descending value and
216 ``top_k`` keeps the leading ones. Full-vocabulary ``"logits"`` readouts
217 go through :func:`transformer_lens.utilities.logits_utils.logits_to_df`
218 (columns ``logit``, ``log_prob``, ``probability``); every other readout
219 carries ``token_index``, optional ``token_string`` and one value column
220 named like ``logits_to_df`` names it (``logit`` / ``log_prob`` /
221 ``probability``).
222 """
223 import pandas as pd
225 values = self.readout.values
226 lead = tuple(values.shape[1:-1])
227 axis_names = ("batch", "pos") if self.has_batch_dim else ("pos",)
228 column = _VALUE_COLUMNS[self.readout.return_type]
229 frames = []
230 for i, label in enumerate(self.labels):
231 for idx in itertools.product(*[range(n) for n in lead]):
232 row = values[(i, *idx)].detach().cpu()
233 row_ids = self.readout.row_ids((i, *idx))
234 if row_ids is None and self.readout.return_type == "logits":
235 frame = logits_to_df(row, tokenizer, top_k)
236 else:
237 ids = torch.arange(row.shape[-1]) if row_ids is None else row_ids.cpu()
238 order = torch.argsort(row, descending=True)
239 if top_k is not None:
240 order = order[:top_k]
241 frame = pd.DataFrame({"token_index": ids[order].tolist()})
242 if tokenizer is not None:
243 frame["token_string"] = [tokenizer.decode([t]) for t in ids[order].tolist()]
244 frame[column] = row[order].tolist()
245 frame.insert(0, "layer", label)
246 for axis, value in enumerate(idx):
247 name = axis_names[axis] if axis < len(axis_names) else f"dim{axis}"
248 frame.insert(1 + axis, name, value)
249 frames.append(frame)
250 return pd.concat(frames, ignore_index=True)
253def _decode_nested(ids: torch.Tensor, decode: Callable[[List[int]], str]) -> Any:
254 if ids.ndim == 0: 254 ↛ 255line 254 didn't jump to line 255 because the condition on line 254 was never true
255 return str(decode([int(ids)]))
256 if ids.ndim == 1:
257 return [str(decode([int(i)])) for i in ids.tolist()]
258 return [_decode_nested(row, decode) for row in ids]
261def _to_token_ids(model: Any, tokens: Union[str, int, Sequence[Union[str, int]]]) -> List[int]:
262 """Resolve strings / ids into a list of single-token ids."""
263 if isinstance(tokens, (str, int)):
264 tokens = [tokens]
265 ids: List[int] = []
266 for token in tokens:
267 if isinstance(token, str):
268 ids.append(int(model.to_single_token(token)))
269 elif isinstance(token, int) and not isinstance(token, bool): 269 ↛ 272line 269 didn't jump to line 272 because the condition on line 269 was always true
270 ids.append(token)
271 else:
272 raise TypeError(f"tokens must be strings or ints, got {type(token).__name__}")
273 return ids
276def _target_names(targets: Optional[TokenInput]) -> Optional[List[Optional[str]]]:
277 """Positional record of the strings in a shared-targets argument."""
278 if targets is None or isinstance(targets, torch.Tensor):
279 return None
280 if isinstance(targets, (str, int)):
281 return [targets if isinstance(targets, str) else None]
282 return [t if isinstance(t, str) else None for t in targets]
285def _final_norm(model: Any) -> Optional[torch.nn.Module]:
286 """The module to apply before unembedding, or ``None`` when the model has no final norm.
288 Raises rather than guessing: an unfolded affine norm applied through the
289 cached-scale path, or an unknown normalization type passed through
290 untouched, would both read the stream wrong silently.
291 """
292 norm_type = getattr(model.cfg, "normalization_type", None)
293 if norm_type is None: 293 ↛ 294line 293 didn't jump to line 294 because the condition on line 293 was never true
294 return None
295 if norm_type not in _FINAL_NORM_TYPES:
296 raise ValueError(
297 f"logit readout does not know how to apply normalization_type={norm_type!r}; "
298 "pass apply_ln=False to read the raw stream, which will not reproduce the model."
299 )
300 ln_final = getattr(model, "ln_final", None)
301 if not isinstance(ln_final, torch.nn.Module):
302 raise ValueError(
303 f"{getattr(model.cfg, 'model_name', 'this model')} declares "
304 f"normalization_type={norm_type!r} but exposes no `ln_final` component, so the "
305 "final norm cannot be applied exactly. Add the alias to its adapter; "
306 "apply_ln=False reads the raw stream instead, which only works when the "
307 "unembedding consumes d_model vectors directly and does not reproduce the model."
308 )
309 return ln_final
312def _require_decoder_unembed(model: Any) -> None:
313 modules = getattr(model, "_modules", {})
314 if "encoder_blocks" in modules or "decoder_blocks" in modules:
315 raise NotImplementedError(
316 "logit lens does not support encoder-decoder bridges yet; the decoder stack "
317 "with cross-attention needs its own accumulated residual stream."
318 )
319 if getattr(model, "unembed", None) is None:
320 raise NotImplementedError(
321 "logit readout needs an `unembed` component; encoder-only bridges expose a "
322 "task head instead of W_U."
323 )
326def _validate_selectors(
327 model: Any,
328 *,
329 return_type: str,
330 vocab: Optional[Union[Sequence[int], torch.Tensor]],
331 targets: Optional[TokenInput],
332 top_k: Optional[int],
333 chunk_size: int,
334) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], int]:
335 """Cheap argument checks shared by both entry points; run before any forward pass.
337 Returns ``(vocab_ids, target_ids, d_vocab)`` where ``target_ids`` is a flat
338 ``[n_targets]`` tensor for shared targets or the caller's tensor (long) for
339 per-row targets.
340 """
341 if return_type not in _RETURN_TYPES:
342 raise ValueError(f"return_type must be one of {_RETURN_TYPES}, got {return_type!r}")
343 if sum(x is not None for x in (vocab, targets, top_k)) > 1:
344 raise ValueError("pass at most one of vocab, targets and top_k")
345 if chunk_size < 1:
346 raise ValueError(f"chunk_size must be >= 1, got {chunk_size}")
347 if vocab is None and targets is None and top_k is None:
348 return None, None, -1
349 d_vocab = int(model.W_U.shape[1]) # read only when a selector needs the bound
350 if top_k is not None and not 1 <= top_k <= d_vocab:
351 raise ValueError(f"top_k must be in [1, {d_vocab}], got {top_k}")
352 for tensor, name in ((vocab, "vocab"), (targets, "targets")):
353 if isinstance(tensor, torch.Tensor) and tensor.dtype not in _INDEX_DTYPES:
354 raise TypeError(f"{name} tensors must hold integer token ids, got dtype {tensor.dtype}")
355 vocab_ids: Optional[torch.Tensor] = None
356 if vocab is not None:
357 vocab_ids = torch.as_tensor(vocab, dtype=torch.long).reshape(-1)
358 target_ids: Optional[torch.Tensor] = None
359 if isinstance(targets, torch.Tensor):
360 target_ids = targets.reshape(1) if targets.ndim == 0 else targets
361 target_ids = target_ids.to(torch.long)
362 elif targets is not None:
363 target_ids = torch.tensor(_to_token_ids(model, targets), dtype=torch.long)
364 for ids, name in ((vocab_ids, "vocab"), (target_ids, "targets")):
365 if ids is not None and (
366 ids.numel() == 0 or int(ids.min()) < 0 or int(ids.max()) >= d_vocab
367 ):
368 raise ValueError(f"{name} ids must be within [0, {d_vocab})")
369 return vocab_ids, target_ids, d_vocab
372def _apply_output_transform(model: Any, logits: torch.Tensor) -> Tuple[torch.Tensor, bool]:
373 adapter = getattr(model, "adapter", None)
374 transform = getattr(adapter, "apply_output_logits_transform", None)
375 if not callable(transform): 375 ↛ 376line 375 didn't jump to line 376 because the condition on line 375 was never true
376 return logits, False
377 return transform(logits), True
380def _run_ln_final(
381 ln_final: torch.nn.Module, flat: Float[torch.Tensor, "n d_model"]
382) -> torch.Tensor:
383 # Norm components contract on [batch, pos, d_model]; stats are per vector, so
384 # the row axis can be presented as the position axis.
385 return ln_final(flat.unsqueeze(0)).squeeze(0)
388def _broadcast_targets(targets: torch.Tensor, lead: Tuple[int, ...]) -> torch.Tensor:
389 """Right-align a per-row target tensor with the vectors' leading dims."""
390 if targets.ndim > len(lead):
391 raise ValueError(
392 f"per-row targets of shape {tuple(targets.shape)} have more dims than the "
393 f"vectors' leading dims {lead}"
394 )
395 shaped = targets.reshape((1,) * (len(lead) - targets.ndim) + tuple(targets.shape))
396 try:
397 return shaped.expand(*lead)
398 except RuntimeError as err:
399 raise ValueError(
400 f"per-row targets of shape {tuple(targets.shape)} do not broadcast to the "
401 f"vectors' leading dims {lead}; they are right-aligned, so a per-example "
402 f"target over several positions must be shaped [batch, 1]"
403 ) from err
406def logit_readout(
407 model: Any,
408 vectors: Float[torch.Tensor, "... d_model"],
409 *,
410 apply_ln: bool = True,
411 vocab: Optional[Union[Sequence[int], torch.Tensor]] = None,
412 targets: Optional[TokenInput] = None,
413 return_type: str = "logits",
414 top_k: Optional[int] = None,
415 chunk_size: int = 64,
416 enable_grad: bool = False,
417) -> LogitReadout:
418 """Read ``d_model`` vectors through the model's final norm and unembedding.
420 Computes ``transform(ln_final(x) @ W_U + b_U)`` for every vector, where
421 ``transform`` is the adapter's post-unembedding transform, in chunks of
422 ``chunk_size`` rows so at most ``chunk_size × d_vocab`` logits exist at once.
423 Everything derived from the full vocabulary (the softmax normalizer,
424 ``entropy``, ``target_ranks``) is computed per chunk and discarded; only the
425 selected columns are kept.
427 Args:
428 model: A ``TransformerBridge`` in eval mode.
429 vectors: Residual-stream-shaped vectors, any leading dims.
430 apply_ln: Apply the model's ``ln_final`` module first (recomputed
431 statistics per vector). ``False`` reads the raw vectors and does not
432 reproduce the model.
433 vocab: Token ids to keep along the trailing axis.
434 targets: Tokens whose values and full-vocabulary ranks to report.
435 Strings (single tokens) or ints, or a sequence of them, are
436 *shared* across every row. A tensor is *per row*: it is
437 right-aligned and broadcast to the vectors' leading dims (for
438 ``[batch, pos, d_model]`` vectors pass ``[batch, pos]`` next-token
439 ids, or ``[batch, 1]`` for one id per example) and the trailing
440 axis of ``values`` has length 1. Note the convention differs from
441 ``ActivationCache.logit_attrs``, which reads a 1-D tensor as
442 per-example first; here alignment is purely right-to-left.
443 return_type: ``"logits"``, ``"log_probs"`` or ``"probs"``.
444 top_k: Keep the ``k`` largest-logit tokens per vector, with their ids.
445 chunk_size: Rows of ``vectors`` processed per unembedding matmul.
446 enable_grad: Keep autograd on (for trainable lenses and attribution);
447 default runs under ``torch.no_grad()``.
449 Returns:
450 A :class:`LogitReadout`.
452 Raises:
453 ValueError: On a bad ``return_type``, more than one selector, a
454 ``chunk_size`` below 1, a wrong trailing dimension, empty vectors,
455 per-row targets that do not broadcast, or a model whose final norm
456 cannot be applied exactly.
457 NotImplementedError: For encoder-decoder and encoder-only bridges.
458 """
459 require_eval_mode(model, operation="logit_readout")
460 _require_decoder_unembed(model)
461 if not isinstance(vectors, torch.Tensor) or not vectors.is_floating_point(): 461 ↛ 462line 461 didn't jump to line 462 because the condition on line 461 was never true
462 raise TypeError("vectors must be a floating-point tensor")
463 ln_final = _final_norm(model) if apply_ln else None
464 vocab_ids, target_ids, _ = _validate_selectors(
465 model,
466 return_type=return_type,
467 vocab=vocab,
468 targets=targets,
469 top_k=top_k,
470 chunk_size=chunk_size,
471 )
472 d_model = int(model.cfg.d_model)
473 if vectors.shape[-1] != d_model:
474 raise ValueError(f"vectors must end in d_model={d_model}, got shape {tuple(vectors.shape)}")
475 lead = tuple(vectors.shape[:-1])
476 if vectors.numel() == 0:
477 raise ValueError(f"vectors has no rows to read (shape {tuple(vectors.shape)})")
479 W_U = require_readable_weight(model.W_U, operation="read logits through W_U", owner=model)
480 b_U = model.b_U
482 per_row = isinstance(targets, torch.Tensor) and targets.ndim >= 1
483 row_targets: Optional[torch.Tensor] = None
484 if per_row:
485 assert target_ids is not None
486 target_ids = _broadcast_targets(target_ids, lead)
487 row_targets = target_ids.reshape(-1).to(W_U.device)
488 select = vocab_ids if vocab_ids is not None else (None if per_row else target_ids)
489 select_dev = None if select is None else select.to(W_U.device)
491 flat = vectors.reshape(-1, d_model).to(device=W_U.device, dtype=W_U.dtype)
492 out_values: List[torch.Tensor] = []
493 out_top_ids: List[torch.Tensor] = []
494 out_ranks: List[torch.Tensor] = []
495 out_entropy: List[torch.Tensor] = []
496 applied_transform = False
498 grad_ctx = torch.enable_grad() if enable_grad else torch.no_grad()
499 with grad_ctx:
500 for start in range(0, flat.shape[0], chunk_size):
501 chunk = flat[start : start + chunk_size]
502 if ln_final is not None:
503 chunk = _run_ln_final(ln_final, chunk)
504 logits = chunk @ W_U + b_U
505 logits, applied_transform = _apply_output_transform(model, logits)
506 log_probs = torch.log_softmax(logits.float(), dim=-1)
507 out_entropy.append(-(log_probs.exp() * log_probs).sum(dim=-1))
508 if return_type == "logits":
509 quantity = logits
510 elif return_type == "log_probs":
511 quantity = log_probs
512 else:
513 quantity = log_probs.exp()
514 if row_targets is not None:
515 rows = row_targets[start : start + chunk_size].unsqueeze(-1) # [n, 1]
516 target_logits = torch.gather(logits, -1, rows)
517 out_ranks.append((logits > target_logits).sum(dim=-1, keepdim=True))
518 out_values.append(torch.gather(quantity, -1, rows))
519 elif target_ids is not None:
520 assert select_dev is not None
521 target_logits = logits[:, select_dev].unsqueeze(-1) # [n, t, 1]
522 out_ranks.append((logits.unsqueeze(1) > target_logits).sum(dim=-1))
523 out_values.append(quantity[:, select_dev])
524 elif top_k is not None:
525 top = torch.topk(logits, k=top_k, dim=-1, sorted=True)
526 out_top_ids.append(top.indices)
527 out_values.append(torch.gather(quantity, -1, top.indices))
528 elif select_dev is not None:
529 out_values.append(quantity[:, select_dev])
530 else:
531 out_values.append(quantity)
533 return LogitReadout(
534 values=torch.cat(out_values, dim=0).reshape(*lead, -1),
535 return_type=return_type,
536 entropy=torch.cat(out_entropy, dim=0).reshape(lead),
537 top_ids=torch.cat(out_top_ids, dim=0).reshape(*lead, -1) if out_top_ids else None,
538 vocab_ids=vocab_ids,
539 target_ids=target_ids,
540 target_ranks=torch.cat(out_ranks, dim=0).reshape(*lead, -1) if out_ranks else None,
541 per_row_targets=per_row,
542 applied_ln=ln_final is not None,
543 applied_output_transform=applied_transform,
544 )
547def _residual_hook_names(n_layers: int, incl_mid: bool) -> List[str]:
548 names = []
549 for layer in range(n_layers):
550 names.append(f"blocks.{layer}.hook_resid_pre")
551 if incl_mid:
552 names.append(f"blocks.{layer}.hook_resid_mid")
553 names.append(f"blocks.{n_layers - 1}.hook_resid_post")
554 return names
557def _resolve_layers(layers: Optional[Sequence[int]], labels: List[str]) -> List[int]:
558 if layers is None:
559 return list(range(len(labels)))
560 resolved = []
561 for layer in layers:
562 index = layer + len(labels) if layer < 0 else layer
563 if not 0 <= index < len(labels):
564 raise ValueError(
565 f"layer {layer} is out of range for the accumulated stack; valid entries are "
566 f"0..{len(labels) - 1} = {labels}"
567 )
568 resolved.append(index)
569 return resolved
572def logit_lens(
573 model: Union[Any, ActivationCache],
574 input: Union[str, List[str], torch.Tensor, None] = None,
575 *,
576 cache: Optional[ActivationCache] = None,
577 layers: Optional[Sequence[int]] = None,
578 incl_mid: bool = False,
579 positions: SliceInput = -1,
580 batch_slice: SliceInput = None,
581 lens: Optional[Callable[[torch.Tensor, int], torch.Tensor]] = None,
582 apply_ln: bool = True,
583 vocab: Optional[Union[Sequence[int], torch.Tensor]] = None,
584 targets: Optional[TokenInput] = None,
585 return_type: str = "logits",
586 top_k: Optional[int] = None,
587 chunk_size: int = 64,
588 enable_grad: bool = False,
589) -> LogitLensResult:
590 """Read the accumulated residual stream through the final norm and unembedding.
592 Builds the stack with
593 :meth:`~transformer_lens.ActivationCache.ActivationCache.accumulated_resid`,
594 optionally applies a per-entry ``lens`` translator, applies the final norm
595 through :meth:`~transformer_lens.ActivationCache.ActivationCache.apply_ln_to_stack`
596 with recomputed statistics, and hands the result to :func:`logit_readout`.
597 The final stack entry reproduces the model's logits (raw bridge) or
598 log-probabilities (compatibility mode).
600 Args:
601 model: A ``TransformerBridge``, or an :class:`ActivationCache` (then
602 ``input`` must be omitted).
603 input: Prompt(s) or token tensor to run when no ``cache`` is given. Only
604 the residual-stream hooks are cached. A batch of strings is padded
605 by the tokenizer, so with the default ``positions=-1`` the shorter
606 prompts are read at a pad position; pass a token tensor (or
607 per-prompt positions) when prompts differ in length.
608 cache: An existing cache to read instead of running the model.
609 layers: Indices into the accumulated stack (negative allowed); ``None``
610 keeps every entry.
611 incl_mid: Include ``resid_mid`` entries. Raises on hybrid stacks whose
612 SSM blocks have no ``resid_mid``.
613 positions: Position selection (``Slice`` semantics); default last token.
614 batch_slice: Batch selection (``Slice`` semantics).
615 lens: Optional ``f(entry, index) -> entry`` applied to each stack entry
616 before the norm, e.g. a tuned-lens translator. Combine with
617 ``enable_grad=True`` to fit or attribute through it.
618 targets: As in :func:`logit_readout`; a per-row tensor is aligned with
619 the stack's ``[batch, pos]`` dims *after* ``positions`` and
620 ``batch_slice`` are applied (``[batch, pos]`` next-token ids with
621 ``positions=None``, ``[batch]`` with the default last position).
622 apply_ln, vocab, return_type, top_k, chunk_size, enable_grad: As in
623 :func:`logit_readout`.
625 Returns:
626 A :class:`LogitLensResult`.
627 """
628 if isinstance(model, ActivationCache):
629 if input is not None or cache is not None:
630 raise ValueError("pass either a cache or a model with input, not both")
631 cache = model
632 model = cache.model
633 require_eval_mode(model, operation="logit_lens")
634 _require_decoder_unembed(model)
635 # Everything cheap is checked before the forward pass is paid for.
636 ln_final = _final_norm(model) if apply_ln else None
637 _validate_selectors(
638 model,
639 return_type=return_type,
640 vocab=vocab,
641 targets=targets,
642 top_k=top_k,
643 chunk_size=chunk_size,
644 )
645 if incl_mid:
646 layer_types = model.layer_types() if hasattr(model, "layer_types") else []
647 hybrid = [t for t in layer_types if any(p in _HYBRID_VARIANT_NAMES for p in t.split("+"))]
648 if hybrid:
649 raise ValueError(
650 f"incl_mid=True needs a resid_mid hook on every block, but this stack has "
651 f"block types {sorted(set(hybrid))} without one; use incl_mid=False."
652 )
653 if cache is None:
654 if input is None:
655 raise ValueError("logit_lens needs an input to run or a cache to read")
656 _, cache = model.run_with_cache(
657 input, names_filter=_residual_hook_names(int(model.cfg.n_layers), incl_mid)
658 )
659 assert cache is not None
661 pos_slice = positions if isinstance(positions, Slice) else Slice(positions)
662 batch = batch_slice if isinstance(batch_slice, Slice) else Slice(batch_slice)
663 stack, labels = cache.accumulated_resid(
664 layer=None, incl_mid=incl_mid, apply_ln=False, pos_slice=pos_slice, return_labels=True
665 )
666 keep = _resolve_layers(layers, labels)
667 stack = stack[keep]
668 labels = [labels[i] for i in keep]
669 if cache.has_batch_dim:
670 stack = batch.apply(stack, dim=1)
671 has_batch_dim = cache.has_batch_dim and batch.mode != "int"
672 grad_ctx = torch.enable_grad() if enable_grad else torch.no_grad()
673 with grad_ctx:
674 if lens is not None:
675 stack = torch.stack([lens(stack[i], keep[i]) for i in range(stack.shape[0])], dim=0)
676 if ln_final is not None:
677 # Recomputed-statistics path through the ln_final module (the #1076 fix).
678 stack = cache.apply_ln_to_stack(
679 stack,
680 layer=None,
681 pos_slice=pos_slice,
682 has_batch_dim=has_batch_dim,
683 recompute_ln=True,
684 )
685 readout = logit_readout(
686 model,
687 stack,
688 apply_ln=False,
689 vocab=vocab,
690 targets=targets,
691 return_type=return_type,
692 top_k=top_k,
693 chunk_size=chunk_size,
694 enable_grad=enable_grad,
695 )
696 readout.applied_ln = ln_final is not None
697 return LogitLensResult(
698 readout=readout,
699 labels=labels,
700 has_batch_dim=has_batch_dim,
701 target_names=_target_names(targets),
702 )