Coverage for transformer_lens/tools/analysis/logit_lens.py: 95%

1"""Logit lens and vocabulary readout. 

2 

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. 

10 

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. 

18 

19Example:: 

20 

21 from transformer_lens.model_bridge import TransformerBridge 

22 from transformer_lens.tools.analysis import logit_lens 

23 

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}") 

28 

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""" 

34 

35import itertools 

36from dataclasses import dataclass 

37from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union 

38 

39import torch 

40from jaxtyping import Float, Int 

41 

42from transformer_lens.ActivationCache import ActivationCache 

43from transformer_lens.tools.analysis._model_state import ( 

44 hybrid_block_types, 

45 require_eval_mode, 

46 to_token_ids, 

47) 

48from transformer_lens.utilities import Slice, SliceInput 

49from transformer_lens.utilities.logits_utils import logits_to_df 

50from transformer_lens.utilities.quantization import require_readable_weight 

51 

52# Shared targets: one string / id or a sequence of them. A tensor means per-row 

53# targets (see logit_readout). 

54TokenInput = Union[str, int, Sequence[Union[str, int]], torch.Tensor] 

55 

56_RETURN_TYPES = ("logits", "log_probs", "probs") 

57_INDEX_DTYPES = (torch.int64, torch.int32, torch.int16, torch.int8, torch.uint8) 

58# Value-column names, matching utilities.logits_utils.logits_to_df. 

59_VALUE_COLUMNS = {"logits": "logit", "log_probs": "log_prob", "probs": "probability"} 

60 

61# Normalization types for which the model has a final norm the lens must apply. 

62_FINAL_NORM_TYPES = ("LN", "LNPre", "RMS", "RMSPre") 

63 

64# Relative to the reference's per-position spread, the identity error measures 0-2e-6 in 

65# fp32 and under 1e-2 in bf16; readouts that skip a model's transform measure 0.7-1.5. 

66_READOUT_IDENTITY_RTOL = 0.1 

67 

68# Readout-path modules ln_final -> unembed would skip: BERT/ModernBERT-decoder heads, OPT-350m's 

69# project_out, DeepSeek-V4's hc_head, BD3LM's adaLN sigma_map, AST's classifier LayerNorm. 

70_READOUT_TRANSFORM_COMPONENTS = ( 

71 "mlm_head", 

72 "prediction_head", 

73 "project_out", 

74 "hc_head", 

75 "sigma_map", 

76 "classifier_ln", 

77) 

78 

79 

80@dataclass 

81class LogitReadout: 

82 """Result of :func:`logit_readout`. 

83 

84 Exactly one of the three selectors expands the trailing axis of ``values``: 

85 ``targets`` gives ``[..., n_targets]`` (``[..., 1]`` for per-row targets), 

86 ``top_k`` gives ``[..., k]`` with ``top_ids``, ``vocab`` gives 

87 ``[..., len(vocab)]``; with none of them the trailing axis is the full 

88 vocabulary. ``entropy`` and ``target_ranks`` are always computed over the 

89 full vocabulary, whatever was selected. 

90 

91 Attributes: 

92 values: The requested quantity (see ``return_type``) with the leading 

93 dims of the input vectors and the selected trailing axis. 

94 return_type: ``"logits"``, ``"log_probs"`` or ``"probs"``. 

95 entropy: Full-vocabulary entropy of the softmax, in nats, ``[...]``. 

96 top_ids: Token ids aligned with ``values`` when ``top_k`` was used. 

97 vocab_ids: The requested subset ids when ``vocab`` was used. 

98 target_ids: Shared targets as ``[n_targets]``; per-row targets as a 

99 tensor with the leading dims of the vectors (one id per row). 

100 target_ranks: Zero-based full-vocabulary rank of each target (0 is the 

101 argmax; ties count only strictly larger logits), aligned with 

102 ``values``. 

103 per_row_targets: Whether ``targets`` was given per row (a tensor). 

104 applied_ln: Whether the model's final norm was applied to the vectors. 

105 applied_output_transform: Whether the adapter's post-unembedding 

106 transform was invoked (it is the identity for most architectures). 

107 """ 

108 

109 values: Float[torch.Tensor, "... selected"] 

110 return_type: str 

111 entropy: Float[torch.Tensor, "..."] 

112 top_ids: Optional[Int[torch.Tensor, "... k"]] = None 

113 vocab_ids: Optional[Int[torch.Tensor, "v"]] = None 

114 target_ids: Optional[torch.Tensor] = None 

115 target_ranks: Optional[Int[torch.Tensor, "... t"]] = None 

116 per_row_targets: bool = False 

117 applied_ln: bool = True 

118 applied_output_transform: bool = True 

119 

120 @property 

121 def selected_ids(self) -> Optional[torch.Tensor]: 

122 """Token ids along the trailing axis of ``values`` when it is not the full vocabulary.""" 

123 if self.target_ids is not None: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 return self.target_ids 

125 if self.vocab_ids is not None: 

126 return self.vocab_ids 

127 return self.top_ids 

128 

129 def row_ids(self, index: Tuple[int, ...]) -> Optional[torch.Tensor]: 

130 """Trailing-axis token ids for one row of ``values`` (``None`` = full vocabulary).""" 

131 if self.top_ids is not None: 

132 return self.top_ids[index] 

133 if self.target_ids is not None and self.per_row_targets: 

134 return self.target_ids[index].reshape(1) 

135 return self.selected_ids 

136 

137 def decode(self, tokenizer: Any) -> Any: 

138 """Decode the trailing-axis token ids to strings, keeping the leading structure. 

139 

140 ``top_k`` and per-row readouts decode per row; shared ``targets`` / 

141 ``vocab`` readouts decode the id list once. Full-vocabulary readouts 

142 have no id list to decode. 

143 """ 

144 decode = getattr(tokenizer, "decode", None) 

145 if not callable(decode): 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true

146 raise TypeError("tokenizer must provide a callable decode method") 

147 if self.top_ids is not None: 

148 return _decode_nested(self.top_ids, decode) 

149 if self.target_ids is not None and self.per_row_targets: 149 ↛ 151line 149 didn't jump to line 151 because the condition on line 149 was always true

150 return _decode_nested(self.target_ids, decode) 

151 ids = self.selected_ids 

152 if ids is None: 

153 raise ValueError( 

154 "decode needs a targets, vocab or top_k readout; this one is full-vocabulary" 

155 ) 

156 return [str(decode([int(i)])) for i in ids.tolist()] 

157 

158 

159@dataclass 

160class LogitLensResult: 

161 """Result of :func:`logit_lens`. 

162 

163 Attributes: 

164 readout: The :class:`LogitReadout` over the accumulated stack. Its 

165 leading axis is the stack entry, aligned with ``labels``; the 

166 remaining leading dims are ``[batch, pos]`` after the requested 

167 ``positions`` / ``batch_slice`` (an integer selection drops its axis). 

168 labels: Stack entry labels from ``accumulated_resid`` (``"0_pre"``, 

169 ``"3_mid"``, ``"final_post"``). 

170 has_batch_dim: Whether a batch axis follows the stack axis. 

171 target_names: For shared targets, the string passed for each target 

172 (``None`` where an id was passed), so :meth:`rank_trajectory` 

173 accepts the strings again. 

174 """ 

175 

176 readout: LogitReadout 

177 labels: List[str] 

178 has_batch_dim: bool 

179 target_names: Optional[List[Optional[str]]] = None 

180 

181 @property 

182 def values(self) -> torch.Tensor: 

183 """Shortcut for ``readout.values``.""" 

184 return self.readout.values 

185 

186 def rank_trajectory(self, target: Union[str, int, None] = None) -> torch.Tensor: 

187 """Full-vocabulary rank of one target at every stack entry, ``[layer, ...]``. 

188 

189 ``target`` is the string passed to ``targets`` or a token id; ``None`` 

190 selects the only target when exactly one was requested, and is the only 

191 form for per-row targets. 

192 """ 

193 ranks = self.readout.target_ranks 

194 ids = self.readout.target_ids 

195 if ranks is None or ids is None: 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true

196 raise ValueError("rank_trajectory needs a lens run with targets=...") 

197 if self.readout.per_row_targets: 

198 if target is not None: 

199 raise ValueError("per-row targets have one target per row; call rank_trajectory()") 

200 return ranks[..., 0] 

201 if target is None: 

202 if ids.numel() != 1: 

203 raise ValueError("pass target= when more than one target was requested") 

204 return ranks[..., 0] 

205 if isinstance(target, str): 

206 names = self.target_names or [] 

207 if target not in names: 

208 known = [n for n in names if n is not None] 

209 raise ValueError( 

210 f"unknown target {target!r}; string lookup needs the string as passed " 

211 f"(known: {known}); pass the token id otherwise" 

212 ) 

213 return ranks[..., names.index(target)] 

214 matches = (ids == int(target)).nonzero(as_tuple=True)[0] 

215 if matches.numel() == 0: 

216 raise ValueError(f"token id {target} was not among the requested targets") 

217 return ranks[..., int(matches[0])] 

218 

219 def top_tokens(self, tokenizer: Any) -> Dict[str, Any]: 

220 """Decoded top-k tokens per stack entry, ``{label: nested lists}``.""" 

221 if self.readout.top_ids is None: 

222 raise ValueError("top_tokens needs a lens run with top_k=...") 

223 return { 

224 label: _decode_nested(self.readout.top_ids[i], tokenizer.decode) 

225 for i, label in enumerate(self.labels) 

226 } 

227 

228 def to_dataframe(self, tokenizer: Optional[Any] = None, top_k: Optional[int] = None) -> Any: 

229 """One row per (stack entry, batch, position, token) as a pandas DataFrame. 

230 

231 Rows within a (layer, batch, pos) group are sorted by descending value and 

232 ``top_k`` keeps the leading ones. Full-vocabulary ``"logits"`` readouts 

233 go through :func:`transformer_lens.utilities.logits_utils.logits_to_df` 

234 (columns ``logit``, ``log_prob``, ``probability``); every other readout 

235 carries ``token_index``, optional ``token_string`` and one value column 

236 named like ``logits_to_df`` names it (``logit`` / ``log_prob`` / 

237 ``probability``). 

238 """ 

239 import pandas as pd 

240 

241 values = self.readout.values 

242 lead = tuple(values.shape[1:-1]) 

243 axis_names = ("batch", "pos") if self.has_batch_dim else ("pos",) 

244 column = _VALUE_COLUMNS[self.readout.return_type] 

245 frames = [] 

246 for i, label in enumerate(self.labels): 

247 for idx in itertools.product(*[range(n) for n in lead]): 

248 row = values[(i, *idx)].detach().cpu() 

249 row_ids = self.readout.row_ids((i, *idx)) 

250 if row_ids is None and self.readout.return_type == "logits": 

251 frame = logits_to_df(row, tokenizer, top_k) 

252 else: 

253 ids = torch.arange(row.shape[-1]) if row_ids is None else row_ids.cpu() 

254 order = torch.argsort(row, descending=True) 

255 if top_k is not None: 

256 order = order[:top_k] 

257 frame = pd.DataFrame({"token_index": ids[order].tolist()}) 

258 if tokenizer is not None: 

259 frame["token_string"] = [tokenizer.decode([t]) for t in ids[order].tolist()] 

260 frame[column] = row[order].tolist() 

261 frame.insert(0, "layer", label) 

262 for axis, value in enumerate(idx): 

263 name = axis_names[axis] if axis < len(axis_names) else f"dim{axis}" 

264 frame.insert(1 + axis, name, value) 

265 frames.append(frame) 

266 return pd.concat(frames, ignore_index=True) 

267 

268 

269def _decode_nested(ids: torch.Tensor, decode: Callable[[List[int]], str]) -> Any: 

270 if ids.ndim == 0: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true

271 return str(decode([int(ids)])) 

272 if ids.ndim == 1: 

273 return [str(decode([int(i)])) for i in ids.tolist()] 

274 return [_decode_nested(row, decode) for row in ids] 

275 

276 

277def _target_names(targets: Optional[TokenInput]) -> Optional[List[Optional[str]]]: 

278 """Positional record of the strings in a shared-targets argument.""" 

279 if targets is None or isinstance(targets, torch.Tensor): 

280 return None 

281 if isinstance(targets, (str, int)): 

282 return [targets if isinstance(targets, str) else None] 

283 return [t if isinstance(t, str) else None for t in targets] 

284 

285 

286def _final_norm(model: Any) -> Optional[torch.nn.Module]: 

287 """The module to apply before unembedding, or ``None`` when the model has no final norm. 

288 

289 Raises rather than guessing: an unfolded affine norm applied through the 

290 cached-scale path, or an unknown normalization type passed through 

291 untouched, would both read the stream wrong silently. 

292 """ 

293 norm_type = getattr(model.cfg, "normalization_type", None) 

294 if norm_type is None: 294 ↛ 295line 294 didn't jump to line 295 because the condition on line 294 was never true

295 return None 

296 if norm_type not in _FINAL_NORM_TYPES: 

297 raise ValueError( 

298 f"logit readout does not know how to apply normalization_type={norm_type!r}; " 

299 "pass apply_ln=False to read the raw stream, which will not reproduce the model." 

300 ) 

301 ln_final = getattr(model, "ln_final", None) 

302 if not isinstance(ln_final, torch.nn.Module): 

303 raise ValueError( 

304 f"{getattr(model.cfg, 'model_name', 'this model')} declares " 

305 f"normalization_type={norm_type!r} but exposes no `ln_final` component, so the " 

306 "final norm cannot be applied exactly. Add the alias to its adapter; " 

307 "apply_ln=False reads the raw stream instead, which only works when the " 

308 "unembedding consumes d_model vectors directly and does not reproduce the model." 

309 ) 

310 return ln_final 

311 

312 

313def _require_decoder_unembed(model: Any, *, operation: str) -> None: 

314 modules = getattr(model, "_modules", {}) 

315 if "encoder_blocks" in modules or "decoder_blocks" in modules: 

316 raise NotImplementedError( 

317 "logit lens does not support encoder-decoder bridges yet; the decoder stack " 

318 "with cross-attention needs its own accumulated residual stream." 

319 ) 

320 if getattr(model, "unembed", None) is None: 

321 raise NotImplementedError( 

322 "logit readout needs an `unembed` component; encoder-only bridges expose a " 

323 "task head instead of W_U." 

324 ) 

325 _require_direct_readout(model, operation=operation) 

326 

327 

328def _require_direct_readout(model: Any, *, operation: str) -> None: 

329 """Reject bridges whose logits are not ``unembed(ln_final(last block's hook_resid_post))``.""" 

330 modules = getattr(model, "_modules", None) or {} 

331 blocks = modules.get("blocks") 

332 for layer, block in enumerate(blocks if isinstance(blocks, torch.nn.ModuleList) else ()): 

333 if not getattr(block, "hook_out_is_single_residual_stream", False): 

334 raise ValueError( 

335 f"{operation} requires single-stream [batch, position, d_model] block outputs; " 

336 f"blocks.{layer} ({type(block).__name__}) carries a stack of residual streams " 

337 "that the model combines before ln_final, so its hook_resid_post is not the " 

338 "stream the model unembeds." 

339 ) 

340 mounted = [name for name in _READOUT_TRANSFORM_COMPONENTS if name in modules] 

341 if mounted: 

342 raise ValueError( 

343 f"{operation} reads the residual stream through ln_final -> unembed only, but this " 

344 f"bridge mounts {mounted} on its readout path, so the result would not be the " 

345 "model's prediction." 

346 ) 

347 

348 

349def _validate_selectors( 

350 model: Any, 

351 *, 

352 return_type: str, 

353 vocab: Optional[Union[Sequence[int], torch.Tensor]], 

354 targets: Optional[TokenInput], 

355 top_k: Optional[int], 

356 chunk_size: int, 

357) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], int]: 

358 """Cheap argument checks shared by both entry points; run before any forward pass. 

359 

360 Returns ``(vocab_ids, target_ids, d_vocab)`` where ``target_ids`` is a flat 

361 ``[n_targets]`` tensor for shared targets or the caller's tensor (long) for 

362 per-row targets. 

363 """ 

364 if return_type not in _RETURN_TYPES: 

365 raise ValueError(f"return_type must be one of {_RETURN_TYPES}, got {return_type!r}") 

366 if sum(x is not None for x in (vocab, targets, top_k)) > 1: 

367 raise ValueError("pass at most one of vocab, targets and top_k") 

368 if chunk_size < 1: 

369 raise ValueError(f"chunk_size must be >= 1, got {chunk_size}") 

370 if vocab is None and targets is None and top_k is None: 

371 return None, None, -1 

372 d_vocab = int(model.W_U.shape[1]) # read only when a selector needs the bound 

373 if top_k is not None and not 1 <= top_k <= d_vocab: 

374 raise ValueError(f"top_k must be in [1, {d_vocab}], got {top_k}") 

375 for tensor, name in ((vocab, "vocab"), (targets, "targets")): 

376 if isinstance(tensor, torch.Tensor) and tensor.dtype not in _INDEX_DTYPES: 

377 raise TypeError(f"{name} tensors must hold integer token ids, got dtype {tensor.dtype}") 

378 vocab_ids: Optional[torch.Tensor] = None 

379 if vocab is not None: 

380 vocab_ids = torch.as_tensor(vocab, dtype=torch.long).reshape(-1) 

381 target_ids: Optional[torch.Tensor] = None 

382 if isinstance(targets, torch.Tensor): 

383 target_ids = targets.reshape(1) if targets.ndim == 0 else targets 

384 target_ids = target_ids.to(torch.long) 

385 for ids, name in ((vocab_ids, "vocab"), (target_ids, "targets")): 

386 if ids is not None and ( 

387 ids.numel() == 0 or int(ids.min()) < 0 or int(ids.max()) >= d_vocab 

388 ): 

389 raise ValueError(f"{name} ids must be within [0, {d_vocab})") 

390 if targets is not None and not isinstance(targets, torch.Tensor): 

391 target_ids = torch.tensor(to_token_ids(model, targets), dtype=torch.long) 

392 return vocab_ids, target_ids, d_vocab 

393 

394 

395def _apply_output_transform(model: Any, logits: torch.Tensor) -> Tuple[torch.Tensor, bool]: 

396 adapter = getattr(model, "adapter", None) 

397 transform = getattr(adapter, "apply_output_logits_transform", None) 

398 if not callable(transform): 398 ↛ 399line 398 didn't jump to line 399 because the condition on line 398 was never true

399 return logits, False 

400 return transform(logits), True 

401 

402 

403def _run_ln_final( 

404 ln_final: torch.nn.Module, flat: Float[torch.Tensor, "n d_model"] 

405) -> torch.Tensor: 

406 # Norm components contract on [batch, pos, d_model]; stats are per vector, so 

407 # the row axis can be presented as the position axis. 

408 return ln_final(flat.unsqueeze(0)).squeeze(0) 

409 

410 

411def _require_readout_identity( 

412 model: Any, 

413 cache: ActivationCache, 

414 logits: Any, 

415 ln_final: Optional[torch.nn.Module], 

416 positions: Slice, 

417 *, 

418 operation: str, 

419 chunk_size: int = 64, 

420) -> None: 

421 """Refuse when ``ln_final -> unembed`` on the last block's output does not give ``logits``.""" 

422 if not isinstance(logits, torch.Tensor): 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true

423 return 

424 resid = cache[f"blocks.{int(model.cfg.n_layers) - 1}.hook_resid_post"] 

425 resid = positions.apply(resid, dim=1) 

426 reference = positions.apply(logits, dim=1) 

427 W_U = require_readable_weight(model.W_U, operation=f"{operation} readout check", owner=model) 

428 if reference.shape != (*resid.shape[:-1], W_U.shape[1]): 

429 raise ValueError( 

430 f"{operation} reads ln_final -> unembed into {W_U.shape[1]} logits per position, " 

431 f"but the model returned logits of shape {tuple(logits.shape)} for a residual " 

432 f"stream of shape {tuple(resid.shape)}, so its readout is not that path." 

433 ) 

434 # Compatibility-mode logits are only promised up to center_unembed's per-position shift. 

435 log_probs = bool(getattr(model, "compatibility_mode", False)) 

436 flat = resid.reshape(-1, resid.shape[-1]) 

437 flat_reference = reference.reshape(-1, reference.shape[-1]) 

438 worst = spread = 0.0 

439 with torch.no_grad(): 

440 for start in range(0, flat.shape[0], chunk_size): 

441 chunk = flat[start : start + chunk_size].to(device=W_U.device, dtype=W_U.dtype) 

442 if ln_final is not None: 442 ↛ 444line 442 didn't jump to line 444 because the condition on line 442 was always true

443 chunk = _run_ln_final(ln_final, chunk) 

444 ours, _ = _apply_output_transform(model, chunk @ W_U + model.b_U) 

445 ours = ours.float() 

446 theirs = flat_reference[start : start + chunk_size].to(ours.device).float() 

447 if log_probs: 

448 ours, theirs = torch.log_softmax(ours, -1), torch.log_softmax(theirs, -1) 

449 worst = max(worst, (ours - theirs).abs().max().item()) 

450 spread = max(spread, (theirs - theirs.mean(-1, keepdim=True)).abs().max().item()) 

451 if worst > _READOUT_IDENTITY_RTOL * spread: 

452 quantity = "log-probabilities" if log_probs else "logits" 

453 raise ValueError( 

454 f"{operation} reads the last block's hook_resid_post through ln_final -> unembed, " 

455 f"but that differs from the model's own {quantity} by up to {worst:.3g} against a " 

456 f"spread of {spread:.3g} (tolerance {_READOUT_IDENTITY_RTOL:g} of the spread): the " 

457 "model transforms the residual stream before its logits in a way this tool does not " 

458 "reproduce." 

459 ) 

460 

461 

462def _broadcast_targets(targets: torch.Tensor, lead: Tuple[int, ...]) -> torch.Tensor: 

463 """Right-align a per-row target tensor with the vectors' leading dims.""" 

464 if targets.ndim > len(lead): 

465 raise ValueError( 

466 f"per-row targets of shape {tuple(targets.shape)} have more dims than the " 

467 f"vectors' leading dims {lead}" 

468 ) 

469 shaped = targets.reshape((1,) * (len(lead) - targets.ndim) + tuple(targets.shape)) 

470 try: 

471 return shaped.expand(*lead) 

472 except RuntimeError as err: 

473 raise ValueError( 

474 f"per-row targets of shape {tuple(targets.shape)} do not broadcast to the " 

475 f"vectors' leading dims {lead}; they are right-aligned, so a per-example " 

476 f"target over several positions must be shaped [batch, 1]" 

477 ) from err 

478 

479 

480def logit_readout( 

481 model: Any, 

482 vectors: Float[torch.Tensor, "... d_model"], 

483 *, 

484 apply_ln: bool = True, 

485 vocab: Optional[Union[Sequence[int], torch.Tensor]] = None, 

486 targets: Optional[TokenInput] = None, 

487 return_type: str = "logits", 

488 top_k: Optional[int] = None, 

489 chunk_size: int = 64, 

490 enable_grad: bool = False, 

491) -> LogitReadout: 

492 """Read ``d_model`` vectors through the model's final norm and unembedding. 

493 

494 Computes ``transform(ln_final(x) @ W_U + b_U)`` for every vector, where 

495 ``transform`` is the adapter's post-unembedding transform, in chunks of 

496 ``chunk_size`` rows so at most ``chunk_size × d_vocab`` logits exist at once. 

497 Everything derived from the full vocabulary (the softmax normalizer, 

498 ``entropy``, ``target_ranks``) is computed per chunk and discarded; only the 

499 selected columns are kept. 

500 

501 Args: 

502 model: A ``TransformerBridge`` in eval mode. 

503 vectors: Residual-stream-shaped vectors, any leading dims. 

504 apply_ln: Apply the model's ``ln_final`` module first (recomputed 

505 statistics per vector). ``False`` reads the raw vectors and does not 

506 reproduce the model. 

507 vocab: Token ids to keep along the trailing axis. 

508 targets: Tokens whose values and full-vocabulary ranks to report. 

509 Strings (single tokens) or ints, or a sequence of them, are 

510 *shared* across every row. A tensor is *per row*: it is 

511 right-aligned and broadcast to the vectors' leading dims (for 

512 ``[batch, pos, d_model]`` vectors pass ``[batch, pos]`` next-token 

513 ids, or ``[batch, 1]`` for one id per example) and the trailing 

514 axis of ``values`` has length 1. Note the convention differs from 

515 ``ActivationCache.logit_attrs``, which reads a 1-D tensor as 

516 per-example first; here alignment is purely right-to-left. 

517 return_type: ``"logits"``, ``"log_probs"`` or ``"probs"``. 

518 top_k: Keep the ``k`` largest-logit tokens per vector, with their ids. 

519 chunk_size: Rows of ``vectors`` processed per unembedding matmul. 

520 enable_grad: Keep autograd on (for trainable lenses and attribution); 

521 default runs under ``torch.no_grad()``. 

522 

523 Returns: 

524 A :class:`LogitReadout`. 

525 

526 Raises: 

527 ValueError: On a bad ``return_type``, more than one selector, a 

528 ``chunk_size`` below 1, a wrong trailing dimension, empty vectors, 

529 per-row targets that do not broadcast, a model whose final norm 

530 cannot be applied exactly, or a bridge whose readout is not 

531 ``ln_final -> unembed`` on the last block's output. 

532 NotImplementedError: For encoder-decoder and encoder-only bridges. 

533 """ 

534 require_eval_mode(model, operation="logit_readout") 

535 _require_decoder_unembed(model, operation="logit_readout") 

536 if not isinstance(vectors, torch.Tensor) or not vectors.is_floating_point(): 536 ↛ 537line 536 didn't jump to line 537 because the condition on line 536 was never true

537 raise TypeError("vectors must be a floating-point tensor") 

538 ln_final = _final_norm(model) if apply_ln else None 

539 vocab_ids, target_ids, _ = _validate_selectors( 

540 model, 

541 return_type=return_type, 

542 vocab=vocab, 

543 targets=targets, 

544 top_k=top_k, 

545 chunk_size=chunk_size, 

546 ) 

547 d_model = int(model.cfg.d_model) 

548 if vectors.shape[-1] != d_model: 

549 raise ValueError(f"vectors must end in d_model={d_model}, got shape {tuple(vectors.shape)}") 

550 lead = tuple(vectors.shape[:-1]) 

551 if vectors.numel() == 0: 

552 raise ValueError(f"vectors has no rows to read (shape {tuple(vectors.shape)})") 

553 

554 W_U = require_readable_weight(model.W_U, operation="read logits through W_U", owner=model) 

555 b_U = model.b_U 

556 

557 per_row = isinstance(targets, torch.Tensor) and targets.ndim >= 1 

558 row_targets: Optional[torch.Tensor] = None 

559 if per_row: 

560 assert target_ids is not None 

561 target_ids = _broadcast_targets(target_ids, lead) 

562 row_targets = target_ids.reshape(-1).to(W_U.device) 

563 select = vocab_ids if vocab_ids is not None else (None if per_row else target_ids) 

564 select_dev = None if select is None else select.to(W_U.device) 

565 

566 flat = vectors.reshape(-1, d_model).to(device=W_U.device, dtype=W_U.dtype) 

567 out_values: List[torch.Tensor] = [] 

568 out_top_ids: List[torch.Tensor] = [] 

569 out_ranks: List[torch.Tensor] = [] 

570 out_entropy: List[torch.Tensor] = [] 

571 applied_transform = False 

572 

573 grad_ctx = torch.enable_grad() if enable_grad else torch.no_grad() 

574 with grad_ctx: 

575 for start in range(0, flat.shape[0], chunk_size): 

576 chunk = flat[start : start + chunk_size] 

577 if ln_final is not None: 

578 chunk = _run_ln_final(ln_final, chunk) 

579 logits = chunk @ W_U + b_U 

580 logits, applied_transform = _apply_output_transform(model, logits) 

581 log_probs = torch.log_softmax(logits.float(), dim=-1) 

582 out_entropy.append(-(log_probs.exp() * log_probs).sum(dim=-1)) 

583 if return_type == "logits": 

584 quantity = logits 

585 elif return_type == "log_probs": 

586 quantity = log_probs 

587 else: 

588 quantity = log_probs.exp() 

589 if row_targets is not None: 

590 rows = row_targets[start : start + chunk_size].unsqueeze(-1) # [n, 1] 

591 target_logits = torch.gather(logits, -1, rows) 

592 out_ranks.append((logits > target_logits).sum(dim=-1, keepdim=True)) 

593 out_values.append(torch.gather(quantity, -1, rows)) 

594 elif target_ids is not None: 

595 assert select_dev is not None 

596 target_logits = logits[:, select_dev].unsqueeze(-1) # [n, t, 1] 

597 out_ranks.append((logits.unsqueeze(1) > target_logits).sum(dim=-1)) 

598 out_values.append(quantity[:, select_dev]) 

599 elif top_k is not None: 

600 top = torch.topk(logits, k=top_k, dim=-1, sorted=True) 

601 out_top_ids.append(top.indices) 

602 out_values.append(torch.gather(quantity, -1, top.indices)) 

603 elif select_dev is not None: 

604 out_values.append(quantity[:, select_dev]) 

605 else: 

606 out_values.append(quantity) 

607 

608 return LogitReadout( 

609 values=torch.cat(out_values, dim=0).reshape(*lead, -1), 

610 return_type=return_type, 

611 entropy=torch.cat(out_entropy, dim=0).reshape(lead), 

612 top_ids=torch.cat(out_top_ids, dim=0).reshape(*lead, -1) if out_top_ids else None, 

613 vocab_ids=vocab_ids, 

614 target_ids=target_ids, 

615 target_ranks=torch.cat(out_ranks, dim=0).reshape(*lead, -1) if out_ranks else None, 

616 per_row_targets=per_row, 

617 applied_ln=ln_final is not None, 

618 applied_output_transform=applied_transform, 

619 ) 

620 

621 

622def _residual_hook_names(n_layers: int, incl_mid: bool) -> List[str]: 

623 names = [] 

624 for layer in range(n_layers): 

625 names.append(f"blocks.{layer}.hook_resid_pre") 

626 if incl_mid: 

627 names.append(f"blocks.{layer}.hook_resid_mid") 

628 names.append(f"blocks.{n_layers - 1}.hook_resid_post") 

629 return names 

630 

631 

632def _resolve_layers(layers: Optional[Sequence[int]], labels: List[str]) -> List[int]: 

633 if layers is None: 

634 return list(range(len(labels))) 

635 resolved = [] 

636 for layer in layers: 

637 index = layer + len(labels) if layer < 0 else layer 

638 if not 0 <= index < len(labels): 

639 raise ValueError( 

640 f"layer {layer} is out of range for the accumulated stack; valid entries are " 

641 f"0..{len(labels) - 1} = {labels}" 

642 ) 

643 resolved.append(index) 

644 return resolved 

645 

646 

647def logit_lens( 

648 model: Union[Any, ActivationCache], 

649 input: Union[str, List[str], torch.Tensor, None] = None, 

650 *, 

651 cache: Optional[ActivationCache] = None, 

652 layers: Optional[Sequence[int]] = None, 

653 incl_mid: bool = False, 

654 positions: SliceInput = -1, 

655 batch_slice: SliceInput = None, 

656 lens: Optional[Callable[[torch.Tensor, int], torch.Tensor]] = None, 

657 apply_ln: bool = True, 

658 vocab: Optional[Union[Sequence[int], torch.Tensor]] = None, 

659 targets: Optional[TokenInput] = None, 

660 return_type: str = "logits", 

661 top_k: Optional[int] = None, 

662 chunk_size: int = 64, 

663 enable_grad: bool = False, 

664) -> LogitLensResult: 

665 """Read the accumulated residual stream through the final norm and unembedding. 

666 

667 Builds the stack with 

668 :meth:`~transformer_lens.ActivationCache.ActivationCache.accumulated_resid`, 

669 optionally applies a per-entry ``lens`` translator, applies the final norm 

670 through :meth:`~transformer_lens.ActivationCache.ActivationCache.apply_ln_to_stack` 

671 with recomputed statistics, and hands the result to :func:`logit_readout`. 

672 The final stack entry reproduces the model's logits (raw bridge) or 

673 log-probabilities (compatibility mode). Bridges whose readout is not 

674 ``ln_final -> unembed`` on the last block's output are refused; a prompt run here 

675 with ``apply_ln=True`` is also checked against the model's own output, a supplied 

676 ``cache`` only structurally. 

677 

678 Args: 

679 model: A ``TransformerBridge``, or an :class:`ActivationCache` (then 

680 ``input`` must be omitted). 

681 input: Prompt(s) or token tensor to run when no ``cache`` is given. Only 

682 the residual-stream hooks are cached. A batch of strings is padded 

683 by the tokenizer, so with the default ``positions=-1`` the shorter 

684 prompts are read at a pad position; pass a token tensor (or 

685 per-prompt positions) when prompts differ in length. 

686 cache: An existing cache to read instead of running the model. 

687 layers: Indices into the accumulated stack (negative allowed); ``None`` 

688 keeps every entry. 

689 incl_mid: Include ``resid_mid`` entries. Raises on hybrid stacks whose 

690 SSM blocks have no ``resid_mid``. 

691 positions: Position selection (``Slice`` semantics); default last token. 

692 batch_slice: Batch selection (``Slice`` semantics). 

693 lens: Optional ``f(entry, index) -> entry`` applied to each stack entry 

694 before the norm, e.g. a tuned-lens translator. Combine with 

695 ``enable_grad=True`` to fit or attribute through it. 

696 targets: As in :func:`logit_readout`; a per-row tensor is aligned with 

697 the stack's ``[batch, pos]`` dims *after* ``positions`` and 

698 ``batch_slice`` are applied (``[batch, pos]`` next-token ids with 

699 ``positions=None``, ``[batch]`` with the default last position). 

700 apply_ln, vocab, return_type, top_k, chunk_size, enable_grad: As in 

701 :func:`logit_readout`. 

702 

703 Returns: 

704 A :class:`LogitLensResult`. 

705 """ 

706 if isinstance(model, ActivationCache): 

707 if input is not None or cache is not None: 

708 raise ValueError("pass either a cache or a model with input, not both") 

709 cache = model 

710 model = cache.model 

711 require_eval_mode(model, operation="logit_lens") 

712 _require_decoder_unembed(model, operation="logit_lens") 

713 # Everything cheap is checked before the forward pass is paid for. 

714 ln_final = _final_norm(model) if apply_ln else None 

715 _validate_selectors( 

716 model, 

717 return_type=return_type, 

718 vocab=vocab, 

719 targets=targets, 

720 top_k=top_k, 

721 chunk_size=chunk_size, 

722 ) 

723 if incl_mid: 

724 hybrid = hybrid_block_types(model) 

725 if hybrid: 

726 raise ValueError(