Coverage for transformer_lens/tools/analysis/backward_lens.py: 90%

1"""Backward Lens gradient-factor capture and vocabulary projection. 

2 

3The Backward Lens represents a linear weight gradient as a sum of token-position 

4outer products and projects residual-width factors into the model vocabulary. 

5The public API supports raw dense-MLP decoder-only ``TransformerBridge`` models 

6(for example GPT-2 and Pythia/GPT-NeoX), reading each MLP projection's weight 

7layout from the Bridge component rather than the model class. 

8""" 

9 

10from __future__ import annotations 

11 

12from collections.abc import Callable, Iterator, Sequence 

13from contextlib import contextmanager 

14from dataclasses import dataclass 

15from typing import Any, Literal, cast 

16 

17import torch 

18import torch.nn.functional as F 

19from jaxtyping import Bool, Float, Int 

20 

21WeightLayout = Literal["in_out", "out_in"] 

22ProjectedFactor = Literal["forward_inputs", "output_gradients"] 

23DEFAULT_TOP_K = 10 

24 

25 

26@dataclass(frozen=True) 

27class LinearGradientFactors: 

28 """Detached factors and reconstruction for one linear weight gradient. 

29 

30 ``forward_inputs`` and ``output_gradients`` have shapes ``[position, in]`` 

31 and ``[position, out]``. ``output_gradients`` and ``weight_gradient`` preserve 

32 the raw ``d(loss)/d(tensor)`` sign; they are not negated into update directions. 

33 Gradient tensors use the requested storage layout. All tensors are cloned to 

34 CPU in float32 so the result owns no autograd graph. 

35 """ 

36 

37 forward_inputs: Float[torch.Tensor, "position in_features"] 

38 output_gradients: Float[torch.Tensor, "position out_features"] 

39 weight_gradient: Float[torch.Tensor, "weight_dim_0 weight_dim_1"] 

40 reconstructed_gradient: Float[torch.Tensor, "weight_dim_0 weight_dim_1"] 

41 absolute_reconstruction_error: float 

42 relative_reconstruction_error: float 

43 weight_layout: WeightLayout 

44 

45 

46@dataclass(frozen=True) 

47class VocabularyRanking: 

48 """Owned CPU copies of signed vocabulary rankings with shape ``[..., k]``. 

49 

50 ``values`` preserves the floating dtype and sign of ``logits``; ``indices`` 

51 has dtype ``torch.int64``. Both tensors are detached. 

52 """ 

53 

54 values: Float[torch.Tensor, "*leading k"] 

55 indices: Int[torch.Tensor, "*leading k"] 

56 

57 

58@dataclass(frozen=True) 

59class BackwardLensMatrixResult: 

60 """Factors and vocabulary readouts for one dense MLP weight matrix. 

61 

62 ``factors`` contains the full linear factorization. ``projected_factor`` says 

63 whether its residual-width ``forward_inputs`` or raw-gradient 

64 ``output_gradients`` were decoded. ``factor_norms`` and ``zero_norm_mask`` 

65 have shape ``[position]`` with float32 and bool dtypes. Largest and smallest 

66 signed rankings are always retained. Full ``vocabulary_logits`` are present 

67 only when explicitly requested. Normalized rankings and optional full logits 

68 are present when the Normalized Logit Lens is requested. Every retained tensor 

69 is a detached CPU-owned value; gradient descent subtracts raw gradients. 

70 """ 

71 

72 factors: LinearGradientFactors 

73 projected_factor: ProjectedFactor 

74 factor_norms: Float[torch.Tensor, "position"] 

75 zero_norm_mask: Bool[torch.Tensor, "position"] 

76 vocabulary_size: int 

77 target_token_id: int 

78 top_ranking: VocabularyRanking 

79 bottom_ranking: VocabularyRanking 

80 target_largest_ranks: Int[torch.Tensor, "position"] 

81 target_smallest_ranks: Int[torch.Tensor, "position"] 

82 normalized_top_ranking: VocabularyRanking | None = None 

83 normalized_bottom_ranking: VocabularyRanking | None = None 

84 normalized_target_largest_ranks: Int[torch.Tensor, "position"] | None = None 

85 normalized_target_smallest_ranks: Int[torch.Tensor, "position"] | None = None 

86 vocabulary_logits: Float[torch.Tensor, "position d_vocab"] | None = None 

87 normalized_vocabulary_logits: Float[torch.Tensor, "position d_vocab"] | None = None 

88 

89 def logits(self, *, normalized: bool = False) -> Float[torch.Tensor, "position d_vocab"]: 

90 """Return opted-in raw or Normalized Logit Lens full logits.""" 

91 if not isinstance(normalized, bool): 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true

92 raise TypeError("normalized must be a bool") 

93 if normalized and self.normalized_top_ranking is None: 

94 raise ValueError("normalized logits were not requested during analysis") 

95 logits = self.normalized_vocabulary_logits if normalized else self.vocabulary_logits 

96 if logits is None: 

97 kind = "normalized " if normalized else "" 

98 raise ValueError( 

99 f"full {kind}logits were not retained; call " 

100 "BackwardLens.analyze(..., return_full_logits=True)" 

101 ) 

102 return logits 

103 

104 def _retained_ranking(self, *, normalized: bool, largest: bool) -> VocabularyRanking: 

105 if not isinstance(normalized, bool): 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true

106 raise TypeError("normalized must be a bool") 

107 if normalized: 

108 ranking = self.normalized_top_ranking if largest else self.normalized_bottom_ranking 

109 if ranking is None: 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true

110 raise ValueError("normalized logits were not requested during analysis") 

111 return ranking 

112 return self.top_ranking if largest else self.bottom_ranking 

113 

114 def top(self, *, k: int, normalized: bool = False) -> VocabularyRanking: 

115 """Return up to the retained largest signed logits and token ids.""" 

116 return _slice_vocabulary_ranking( 

117 self._retained_ranking(normalized=normalized, largest=True), k=k 

118 ) 

119 

120 def bottom(self, *, k: int, normalized: bool = False) -> VocabularyRanking: 

121 """Return up to the retained smallest signed logits and token ids.""" 

122 return _slice_vocabulary_ranking( 

123 self._retained_ranking(normalized=normalized, largest=False), k=k 

124 ) 

125 

126 def top_tokens(self, tokenizer: Any, *, k: int, normalized: bool = False) -> list[list[str]]: 

127 """Decode the largest-``k`` vocabulary ids for every position.""" 

128 return _decode_vocabulary_ranking(self.top(k=k, normalized=normalized), tokenizer) 

129 

130 def bottom_tokens(self, tokenizer: Any, *, k: int, normalized: bool = False) -> list[list[str]]: 

131 """Decode the smallest-``k`` vocabulary ids for every position.""" 

132 return _decode_vocabulary_ranking(self.bottom(k=k, normalized=normalized), tokenizer) 

133 

134 def target_ranks( 

135 self, 

136 target_token_id: int, 

137 *, 

138 largest: bool, 

139 normalized: bool = False, 

140 ) -> Int[torch.Tensor, "position"]: 

141 """Return zero-based target ranks per position in the requested ordering. 

142 

143 ``largest=True`` gives rank zero to the largest logit. ``largest=False`` 

144 gives rank zero to the smallest, which is the useful raw-gradient 

145 convention for the second MLP matrix because gradient descent subtracts it. 

146 Ties receive the same competition rank. The analyzed target's ranks are 

147 always retained; other token ids require opted-in full logits. 

148 """ 

149 if isinstance(target_token_id, bool) or not isinstance(target_token_id, int): 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true

150 raise TypeError("target_token_id must be an integer") 

151 if not 0 <= target_token_id < self.vocabulary_size: 

152 raise ValueError( 

153 f"target_token_id must be in [0, {self.vocabulary_size - 1}]; " 

154 f"got {target_token_id}" 

155 ) 

156 if not isinstance(largest, bool): 156 ↛ 157line 156 didn't jump to line 157 because the condition on line 156 was never true

157 raise TypeError("largest must be a bool") 

158 if not isinstance(normalized, bool): 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true

159 raise TypeError("normalized must be a bool") 

160 if normalized: 

161 ranks = ( 

162 self.normalized_target_largest_ranks 

163 if largest 

164 else self.normalized_target_smallest_ranks 

165 ) 

166 if ranks is None: 166 ↛ 167line 166 didn't jump to line 167 because the condition on line 166 was never true

167 raise ValueError("normalized logits were not requested during analysis") 

168 else: 

169 ranks = self.target_largest_ranks if largest else self.target_smallest_ranks 

170 if target_token_id == self.target_token_id: 

171 return ranks.clone() 

172 return _target_vocabulary_ranks( 

173 self.logits(normalized=normalized), 

174 target_token_id=target_token_id, 

175 largest=largest, 

176 ) 

177 

178 def gradient_descent_target_ranks( 

179 self, target_token_id: int, *, normalized: bool = False 

180 ) -> Int[torch.Tensor, "position"]: 

181 """Return ascending raw-gradient target ranks (rank zero is smallest).""" 

182 return self.target_ranks(target_token_id, largest=False, normalized=normalized) 

183 

184 

185@dataclass(frozen=True) 

186class BackwardLensLayerResult: 

187 """Vocabulary-facing input/output MLP matrix results for one indexed layer.""" 

188 

189 layer: int 

190 input_projection: BackwardLensMatrixResult 

191 output_projection: BackwardLensMatrixResult 

192 

193 

194@dataclass(frozen=True) 

195class BackwardLensResult: 

196 """Detached result of one :meth:`BackwardLens.analyze` call. 

197 

198 ``prompt`` and ``target_token`` echo the analyzed inputs; ``target_token_id`` 

199 is the single vocabulary id the target text encodes to. ``loss`` is the raw 

200 scalar cross-entropy of the final-position next-token prediction against the 

201 target; it preserves the ``d(loss)/d(...)`` sign convention and is not negated. 

202 ``prompt_token_ids`` is an owned CPU int64 tensor with shape ``[position]``; 

203 every residual-width factor in ``layers`` is aligned to these same positions. 

204 Position zero is a prepended BOS only when the model and tokenizer configuration 

205 requests one. ``layers`` preserves requested order. Maximum errors summarize 

206 both matrices over every requested layer. 

207 ``includes_normalized_logits`` records whether the Normalized Logit Lens was 

208 computed. ``includes_full_logits`` records whether full vocabulary tensors 

209 were retained in addition to bounded rankings. No model or tokenizer reference 

210 is retained. 

211 """ 

212 

213 prompt: str 

214 prompt_token_ids: Int[torch.Tensor, "position"] 

215 target_token: str 

216 target_token_id: int 

217 loss: float 

218 layers: tuple[BackwardLensLayerResult, ...] 

219 max_absolute_reconstruction_error: float 

220 max_relative_reconstruction_error: float 

221 includes_normalized_logits: bool 

222 includes_full_logits: bool 

223 

224 def layer(self, layer: int) -> BackwardLensLayerResult: 

225 """Return one requested layer result or raise ``KeyError``.""" 

226 for result in self.layers: 

227 if result.layer == layer: 

228 return result 

229 raise KeyError(f"layer {layer} was not analyzed") 

230 

231 

232@dataclass(frozen=True) 

233class _MLPLayerGradientFactors: 

234 """Detached gradient factors for both MLP projections in one layer.""" 

235 

236 layer: int 

237 input_projection: LinearGradientFactors 

238 output_projection: LinearGradientFactors 

239 

240 

241@dataclass(frozen=True) 

242class _DenseMLPGradientCapture: 

243 """Private capture result for one dense-MLP next-token loss. 

244 

245 Tensor fields are detached, owned CPU copies. 

246 """ 

247 

248 prompt_token_ids: Int[torch.Tensor, "1 position"] 

249 target_token_id: int 

250 loss: float 

251 layers: tuple[_MLPLayerGradientFactors, ...] 

252 

253 

254def _validate_floating_matrix(name: str, tensor: Any) -> Float[torch.Tensor, "rows columns"]: 

255 if not isinstance(tensor, torch.Tensor): 

256 raise TypeError(f"{name} must be a torch.Tensor") 

257 if tensor.ndim != 2: 

258 raise ValueError(f"{name} must be rank 2; got shape {tuple(tensor.shape)}") 

259 if 0 in tensor.shape: 

260 raise ValueError(f"{name} must have no empty dimensions; got shape {tuple(tensor.shape)}") 

261 if not tensor.is_floating_point(): 

262 raise TypeError(f"{name} must have a floating dtype; got {tensor.dtype}") 

263 if not bool(torch.isfinite(tensor).all()): 

264 raise ValueError(f"{name} must contain only finite values") 

265 return cast(Float[torch.Tensor, "rows columns"], tensor) 

266 

267 

268def _reconstruction_errors( 

269 reference: Float[torch.Tensor, "rows columns"], 

270 reconstruction: Float[torch.Tensor, "rows columns"], 

271) -> tuple[float, float]: 

272 """Return max absolute and symmetric scale-aware relative errors. 

273 

274 The relative error is ``||reference - reconstruction||_F`` divided by the 

275 maximum of the two input Frobenius norms and float32 epsilon. This remains 

276 finite when one or both gradients are zero. 

277 """ 

278 difference = reference - reconstruction 

279 absolute = float(difference.abs().max()) 

280 scale = torch.maximum(reference.norm(), reconstruction.norm()).clamp_min( 

281 torch.finfo(reference.dtype).eps 

282 ) 

283 relative = float(difference.norm() / scale) 

284 return absolute, relative 

285 

286 

287def _to_detached_float32( 

288 name: str, tensor: Float[torch.Tensor, "rows columns"] 

289) -> Float[torch.Tensor, "rows columns"]: 

290 """Detach and convert a validated tensor, rejecting float32 overflow.""" 

291 converted = tensor.detach().float() 

292 if not bool(torch.isfinite(converted).all()): 

293 raise ValueError(f"{name} must remain finite when converted to float32") 

294 return converted 

295 

296 

297def _build_linear_gradient_factors( 

298 forward_inputs: Any, 

299 output_gradients: Any, 

300 weight_gradient: Any, 

301 *, 

302 weight_layout: Any, 

303) -> LinearGradientFactors: 

304 """Reconstruct a weight gradient from aligned token-position factors. 

305 

306 Args: 

307 forward_inputs: Linear inputs with shape ``[position, in_features]``. 

308 output_gradients: Loss gradients with respect to linear outputs, shape 

309 ``[position, out_features]``. 

310 weight_gradient: Independently computed gradient in ``weight_layout``. 

311 weight_layout: ``"in_out"`` for GPT-2 ``Conv1D`` storage or 

312 ``"out_in"`` for ``torch.nn.Linear`` storage. 

313 

314 Returns: 

315 Detached factors, the independent gradient, its reconstruction, and 

316 reconstruction errors, all on CPU with float32 tensor values. 

317 """ 

318 validated_inputs = _validate_floating_matrix("forward_inputs", forward_inputs) 

319 validated_gradients = _validate_floating_matrix("output_gradients", output_gradients) 

320 validated_weight = _validate_floating_matrix("weight_gradient", weight_gradient) 

321 if weight_layout not in ("in_out", "out_in"): 

322 raise ValueError("weight_layout must be 'in_out' or 'out_in'") 

323 validated_layout = cast(WeightLayout, weight_layout) 

324 if validated_inputs.shape[0] != validated_gradients.shape[0]: 

325 raise ValueError( 

326 "forward_inputs and output_gradients must have the same number of positions; " 

327 f"got {validated_inputs.shape[0]} and {validated_gradients.shape[0]}" 

328 ) 

329 devices = {validated_inputs.device, validated_gradients.device, validated_weight.device} 

330 if len(devices) != 1: 330 ↛ 331line 330 didn't jump to line 331 because the condition on line 330 was never true

331 raise ValueError( 

332 "forward_inputs, output_gradients, and weight_gradient must share a device" 

333 ) 

334 

335 inputs = _to_detached_float32("forward_inputs", validated_inputs) 

336 gradients = _to_detached_float32("output_gradients", validated_gradients) 

337 canonical = inputs.T @ gradients 

338 reconstruction = canonical if validated_layout == "in_out" else canonical.T 

339 reference = _to_detached_float32("weight_gradient", validated_weight) 

340 if not bool(torch.isfinite(reconstruction).all()): 

341 raise ValueError("the float32 outer-product reconstruction must contain only finite values") 

342 if reference.shape != reconstruction.shape: 

343 raise ValueError( 

344 f"weight_gradient shape {tuple(reference.shape)} does not match the " 

345 f"{validated_layout} reconstruction shape {tuple(reconstruction.shape)}" 

346 ) 

347 absolute, relative = _reconstruction_errors(reference, reconstruction) 

348 return LinearGradientFactors( 

349 forward_inputs=inputs.cpu().clone(), 

350 output_gradients=gradients.cpu().clone(), 

351 weight_gradient=reference.cpu().clone(), 

352 reconstructed_gradient=reconstruction.cpu().clone(), 

353 absolute_reconstruction_error=absolute, 

354 relative_reconstruction_error=relative, 

355 weight_layout=validated_layout, 

356 ) 

357 

358 

359def _rank_vocabulary_logits(logits: Any, *, k: Any, largest: Any) -> VocabularyRanking: 

360 """Return largest or smallest vocabulary logits and token ids per row. 

361 

362 Ordering among exactly tied logits is intentionally unspecified and follows 

363 :func:`torch.topk`. 

364 """ 

365 if not isinstance(logits, torch.Tensor): 

366 raise TypeError("logits must be a torch.Tensor") 

367 if logits.ndim not in (1, 2) or logits.shape[-1] == 0: 

368 raise ValueError( 

369 f"logits must have shape [vocab] or [position, vocab]; got {tuple(logits.shape)}" 

370 ) 

371 if not logits.is_floating_point(): 

372 raise TypeError(f"logits must have a floating dtype; got {logits.dtype}") 

373 if not bool(torch.isfinite(logits).all()): 

374 raise ValueError("logits must contain only finite values") 

375 if not isinstance(largest, bool): 

376 raise TypeError(f"largest must be a bool; got {type(largest).__name__}") 

377 if isinstance(k, bool) or not isinstance(k, int) or not 1 <= k <= logits.shape[-1]: 

378 raise ValueError(f"k must be in [1, {logits.shape[-1]}]; got {k!r}") 

379 validated_logits = cast(Float[torch.Tensor, "*leading d_vocab"], logits) 

380 ranked = torch.topk(validated_logits.detach(), k=k, dim=-1, largest=largest, sorted=True) 

381 return VocabularyRanking( 

382 values=ranked.values.cpu().clone(), indices=ranked.indices.cpu().clone() 

383 ) 

384 

385 

386def _slice_vocabulary_ranking(ranking: VocabularyRanking, *, k: int) -> VocabularyRanking: 

387 """Return an owned prefix of an already sorted vocabulary ranking.""" 

388 retained = ranking.indices.shape[-1] 

389 if isinstance(k, bool) or not isinstance(k, int) or not 1 <= k <= retained: 

390 raise ValueError(f"k must be in [1, retained top_k={retained}]; got {k!r}") 

391 return VocabularyRanking( 

392 values=ranking.values[..., :k].clone(), 

393 indices=ranking.indices[..., :k].clone(), 

394 ) 

395 

396 

397def _target_vocabulary_ranks( 

398 logits: Float[torch.Tensor, "position d_vocab"], *, target_token_id: int, largest: bool 

399) -> Int[torch.Tensor, "position"]: 

400 """Return zero-based competition ranks for one vocabulary id per row.""" 

401 target = logits[:, target_token_id].unsqueeze(-1) 

402 comparisons = logits > target if largest else logits < target 

403 return comparisons.sum(dim=-1, dtype=torch.int64).cpu().clone() 

404 

405 

406def _decode_vocabulary_ranking(ranking: VocabularyRanking, tokenizer: Any) -> list[list[str]]: 

407 """Decode a two-dimensional vocabulary ranking without retaining a tokenizer.""" 

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

409 if not callable(decode): 

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

411 if ranking.indices.ndim != 2: 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true

412 raise ValueError("decoded matrix rankings must have shape [position, k]") 

413 return [[str(decode([token_id])) for token_id in row.tolist()] for row in ranking.indices] 

414 

415 

416@torch.no_grad() 

417def _project_residual_factors( 

418 model: Any, factors: Float[torch.Tensor, "position d_model"] 

419) -> Float[torch.Tensor, "position d_vocab"]: 

420 """Apply fresh final normalization and unembedding to residual-width rows.""" 

421 _validate_floating_matrix("factors", factors) 

422 if factors.shape[-1] != int(model.cfg.d_model): 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true

423 raise ValueError( 

424 f"factors must have width d_model={model.cfg.d_model}; got {factors.shape[-1]}" 

425 ) 

426 unembed_weight = model.W_U 

427 if not isinstance(unembed_weight, torch.Tensor) or unembed_weight.ndim != 2: 427 ↛ 428line 427 didn't jump to line 428 because the condition on line 427 was never true

428 raise ValueError("the GPT-2 Bridge must expose a rank-2 unembedding weight") 

429 batched = ( 

430 factors.detach().to(device=unembed_weight.device, dtype=unembed_weight.dtype).unsqueeze(0) 

431 ) 

432 logits = model.unembed(model.ln_final(batched)).squeeze(0) 

433 if logits.ndim != 2 or logits.shape != (factors.shape[0], unembed_weight.shape[1]): 433 ↛ 434line 433 didn't jump to line 434 because the condition on line 433 was never true

434 raise RuntimeError( 

435 "final normalization and unembedding must return [position, d_vocab]; " 

436 f"got {tuple(logits.shape)}" 

437 ) 

438 projected = logits.detach().float() 

439 if not bool(torch.isfinite(projected).all()): 439 ↛ 440line 439 didn't jump to line 440 because the condition on line 439 was never true

440 raise ValueError("vocabulary projection must remain finite in float32") 

441 return projected.clone() 

442 

443 

444def _factor_norms_and_normalized_rows( 

445 factors: Float[torch.Tensor, "position width"], 

446) -> tuple[ 

447 Float[torch.Tensor, "position"], 

448 Bool[torch.Tensor, "position"], 

449 Float[torch.Tensor, "position width"], 

450]: 

451 """Return original L2 norms, exact-zero mask, and safely unit-normalized rows.""" 

452 _validate_floating_matrix("factors", factors) 

453 rows = _to_detached_float32("factors", factors).cpu().clone() 

454 norms = rows.norm(dim=-1) 

455 zero_mask = norms == 0 

456 denominators = torch.where(zero_mask, torch.ones_like(norms), norms) 

457 normalized = rows / denominators.unsqueeze(-1) 

458 return norms.clone(), zero_mask.clone(), normalized 

459 

460 

461def _build_matrix_result( 

462 model: Any, 

463 factors: LinearGradientFactors, 

464 *, 

465 projected_factor: ProjectedFactor, 

466 include_normalized_logits: bool, 

467 target_token_id: int, 

468 top_k: int, 

469 return_full_logits: bool, 

470) -> BackwardLensMatrixResult: 

471 if projected_factor not in ("forward_inputs", "output_gradients"): 471 ↛ 472line 471 didn't jump to line 472 because the condition on line 471 was never true

472 raise ValueError("projected_factor must be 'forward_inputs' or 'output_gradients'") 

473 if not isinstance(include_normalized_logits, bool): 473 ↛ 474line 473 didn't jump to line 474 because the condition on line 473 was never true

474 raise TypeError("include_normalized_logits must be a bool") 

475 if not isinstance(return_full_logits, bool): 475 ↛ 476line 475 didn't jump to line 476 because the condition on line 475 was never true

476 raise TypeError("return_full_logits must be a bool") 

477 rows = ( 

478 factors.forward_inputs if projected_factor == "forward_inputs" else factors.output_gradients 

479 ) 

480 norms, zero_mask, normalized_rows = _factor_norms_and_normalized_rows(rows) 

481 raw_logits = _project_residual_factors(model, rows) 

482 normalized_logits = ( 

483 _project_residual_factors(model, normalized_rows) if include_normalized_logits else None 

484 ) 

485 top_ranking = _rank_vocabulary_logits(raw_logits, k=top_k, largest=True) 

486 bottom_ranking = _rank_vocabulary_logits(raw_logits, k=top_k, largest=False) 

487 target_largest_ranks = _target_vocabulary_ranks( 

488 raw_logits, target_token_id=target_token_id, largest=True 

489 ) 

490 target_smallest_ranks = _target_vocabulary_ranks( 

491 raw_logits, target_token_id=target_token_id, largest=False 

492 ) 

493 normalized_top_ranking = None 

494 normalized_bottom_ranking = None 

495 normalized_target_largest_ranks = None 

496 normalized_target_smallest_ranks = None 

497 if normalized_logits is not None: 

498 normalized_top_ranking = _rank_vocabulary_logits(normalized_logits, k=top_k, largest=True) 

499 normalized_bottom_ranking = _rank_vocabulary_logits( 

500 normalized_logits, k=top_k, largest=False 

501 ) 

502 normalized_target_largest_ranks = _target_vocabulary_ranks( 

503 normalized_logits, target_token_id=target_token_id, largest=True 

504 ) 

505 normalized_target_smallest_ranks = _target_vocabulary_ranks( 

506 normalized_logits, target_token_id=target_token_id, largest=False 

507 ) 

508 return BackwardLensMatrixResult( 

509 factors=factors, 

510 projected_factor=projected_factor, 

511 factor_norms=norms, 

512 zero_norm_mask=zero_mask, 

513 vocabulary_size=raw_logits.shape[-1], 

514 target_token_id=target_token_id, 

515 top_ranking=top_ranking, 

516 bottom_ranking=bottom_ranking, 

517 target_largest_ranks=target_largest_ranks, 

518 target_smallest_ranks=target_smallest_ranks, 

519 normalized_top_ranking=normalized_top_ranking, 

520 normalized_bottom_ranking=normalized_bottom_ranking, 

521 normalized_target_largest_ranks=normalized_target_largest_ranks, 

522 normalized_target_smallest_ranks=normalized_target_smallest_ranks, 

523 vocabulary_logits=raw_logits.cpu().clone() if return_full_logits else None, 

524 normalized_vocabulary_logits=( 

525 normalized_logits.cpu().clone() 

526 if return_full_logits and normalized_logits is not None 

527 else None 

528 ), 

529 ) 

530 

531 

532def _validate_requested_layers(model: Any, layers: Sequence[int]) -> tuple[int, ...]: 

533 if isinstance(layers, (str, bytes)) or not isinstance(layers, Sequence): 533 ↛ 534line 533 didn't jump to line 534 because the condition on line 533 was never true

534 raise TypeError("layers must be a sequence of integer layer indices") 

535 requested = tuple(layers) 

536 if not requested: 

537 raise ValueError("layers must contain at least one layer index") 

538 for layer in requested: 

539 if isinstance(layer, bool) or not isinstance(layer, int): 

540 raise TypeError(f"each layer must be an integer; got {layer!r}") 

541 if len(set(requested)) != len(requested): 

542 raise ValueError("layers must not contain duplicate indices") 

543 n_layers = int(model.cfg.n_layers) 

544 invalid = [layer for layer in requested if not 0 <= layer < n_layers] 

545 if invalid: 

546 raise ValueError(f"layers must be in [0, {n_layers - 1}]; got {invalid}") 

547 return requested 

548 

549 

550def _require_raw_dense_mlp_bridge(model: Any) -> None: 

551 """Require the raw dense-MLP Bridge capabilities used by gradient capture. 

552 

553 The architecture is not constrained by class; support is decided per 

554 projection from the Bridge weight-layout oracle in projection discovery. 

555 """ 

556 from transformer_lens.model_bridge import TransformerBridge 

557 

558 if not isinstance(model, TransformerBridge): 

559 raise TypeError( 

560 "Backward Lens supports TransformerBridge only; load the model with " 

561 "TransformerBridge.boot_transformers(...)." 

562 ) 

563 if getattr(model, "compatibility_mode", False): 

564 raise ValueError( 

565 "Backward Lens requires a raw TransformerBridge; compatibility mode is enabled" 

566 ) 

567 if getattr(model, "_weights_processed", False): 

568 raise ValueError( 

569 "Backward Lens requires original model weights; this Bridge processed its weights" 

570 ) 

571 if int(model.cfg.n_devices) > 1: 

572 raise ValueError( 

573 "Backward Lens requires a single-device TransformerBridge because the MLP " 

574 "projections, final normalization, and unembed must be co-located; " 

575 f"device-map dispatch with cfg.n_devices={model.cfg.n_devices} is not supported" 

576 ) 

577 if bool(getattr(model.cfg, "gated_mlp", False)): 

578 raise NotImplementedError("Backward Lens currently requires dense, non-gated MLPs") 

579 if model.tokenizer is None: 

580 raise ValueError("Backward Lens requires a TransformerBridge with a tokenizer") 

581 for component in ("blocks", "ln_final", "unembed"): 

582 if not hasattr(model, component): 582 ↛ 583line 582 didn't jump to line 583 because the condition on line 582 was never true

583 raise ValueError(f"Backward Lens requires the standard {component} component") 

584 

585 

586@dataclass(frozen=True) 

587class _MLPLinear: 

588 """One validated dense-MLP linear projection with its resolved weight layout.""" 

589 

590 projection: Any 

591 weight_layout: WeightLayout 

592 

593 

594def _get_dense_mlp_projections( 

595 model: Any, layers: tuple[int, ...] 

596) -> dict[int, tuple[_MLPLinear, _MLPLinear]]: 

597 """Return validated live dense-MLP input/output projection bridges. 

598 

599 Each projection's storage orientation is resolved from the Bridge weight-layout 

600 oracle, so Conv1D ``[in, out]`` and ``torch.nn.Linear`` ``[out, in]`` weights are 

601 both accepted without inspecting the model class. Projections whose wrapped 

602 module the oracle cannot orient are rejected. 

603 """ 

604 from transformer_lens.hook_points import HookPoint 

605 from transformer_lens.model_bridge.generalized_components import ( 

606 LinearBridge, 

607 MLPBridge, 

608 ) 

609 from transformer_lens.model_bridge.generalized_components.mlp import ( 

610 weight_layout_in_out, 

611 ) 

612 

613 d_model = int(model.cfg.d_model) 

614 d_mlp = int(model.cfg.d_mlp) 

615 # Feature counts are fixed by the MLP role; storage order follows the layout. 

616 # Input maps d_model -> d_mlp, output maps d_mlp -> d_model. 

617 feature_pairs = ((d_model, d_mlp), (d_mlp, d_model)) 

618 projections: dict[int, tuple[_MLPLinear, _MLPLinear]] = {} 

619 for layer in layers: 

620 mlp = model.blocks[layer].mlp 

621 if not isinstance(mlp, MLPBridge) or getattr(mlp, "gate", None) is not None: 

622 raise ValueError(f"layer {layer} must have a dense, non-gated MLPBridge") 

623 pair = (getattr(mlp, "in", None), getattr(mlp, "out", None)) 

624 records: list[_MLPLinear] = [] 

625 for name, projection, feature_pair in zip( 

626 ("input", "output"), pair, feature_pairs, strict=True 

627 ): 

628 if not isinstance(projection, LinearBridge): 628 ↛ 629line 628 didn't jump to line 629 because the condition on line 628 was never true

629 raise ValueError(f"layer {layer} {name} projection must be a LinearBridge") 

630 layout_flag = weight_layout_in_out(projection) 

631 if layout_flag is None: 

632 raise ValueError( 

633 f"layer {layer} {name} projection has an unknown weight layout; " 

634 "Backward Lens supports Conv1D or torch.nn.Linear MLP projections" 

635 ) 

636 weight_layout: WeightLayout = "in_out" if layout_flag else "out_in" 

637 in_features, out_features = feature_pair 

638 expected_shape = ( 

639 (in_features, out_features) 

640 if weight_layout == "in_out" 

641 else (out_features, in_features) 

642 ) 

643 weight = getattr(projection.original_component, "weight", None) 

644 if not isinstance(weight, torch.nn.Parameter): 644 ↛ 645line 644 didn't jump to line 645 because the condition on line 644 was never true

645 raise ValueError( 

646 f"layer {layer} {name} original weight must be a trainable Parameter" 

647 ) 

648 if not weight.is_floating_point() or tuple(weight.shape) != expected_shape: 

649 raise ValueError( 

650 f"layer {layer} {name} weight must have shape {expected_shape} " 

651 f"and floating dtype; got {tuple(weight.shape)} and {weight.dtype}" 

652 ) 

653 if not weight.requires_grad: 

654 raise ValueError( 

655 f"layer {layer} {name} original weight must be a trainable Parameter" 

656 ) 

657 if not isinstance(projection.hook_in, HookPoint) or not isinstance( 657 ↛ 660line 657 didn't jump to line 660 because the condition on line 657 was never true

658 projection.hook_out, HookPoint 

659 ): 

660 raise ValueError(f"layer {layer} {name} projection is missing Bridge hook points") 

661 records.append(_MLPLinear(projection=projection, weight_layout=weight_layout)) 

662 projections[layer] = (records[0], records[1]) 

663 return projections 

664 

665 

666def _capture_once( 

667 captured: dict[tuple[int, str, str], torch.Tensor], key: tuple[int, str, str] 

668) -> Callable[[torch.nn.Module, tuple[Any, ...], Any], None]: 

669 """Build a non-modifying PyTorch hook that records one tensor.""" 

670 

671 def capture(_module: torch.nn.Module, _inputs: tuple[Any, ...], output: Any) -> None: 

672 if key in captured: 672 ↛ 673line 672 didn't jump to line 673 because the condition on line 672 was never true

673 raise RuntimeError(f"Backward Lens hook {key} fired more than once") 

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

675 raise RuntimeError(f"Backward Lens hook {key} returned a non-tensor output") 

676 captured[key] = output 

677 

678 return capture 

679 

680 

681@contextmanager 

682def _capture_projection_tensors( 

683 projections: dict[int, tuple[_MLPLinear, _MLPLinear]], 

684) -> Iterator[dict[tuple[int, str, str], torch.Tensor]]: 

685 """Capture exact linear boundaries while preserving every pre-existing hook.""" 

686 captured: dict[tuple[int, str, str], torch.Tensor] = {} 

687 handles: list[Any] = [] 

688 try: 

689 for layer, pair in projections.items(): 

690 for name, record in zip(("input", "output"), pair, strict=True): 

691 projection = record.projection 

692 input_key = (layer, name, "forward_input") 

693 output_key = (layer, name, "output") 

694 # Existing hook_in edits must run first so this is the actual linear input. 

695 handles.append( 

696 projection.hook_in.register_forward_hook(_capture_once(captured, input_key)) 

697 ) 

698 # Capture the raw linear output before any existing hook_out edits. 

699 handles.append( 

700 projection.hook_out.register_forward_hook( 

701 _capture_once(captured, output_key), prepend=True 

702 ) 

703 ) 

704 yield captured 

705 finally: 

706 for handle in reversed(handles): 

707 handle.remove() 

708 

709 

710def _single_batch_matrix(name: str, tensor: Any) -> Float[torch.Tensor, "position width"]: 

711 if not isinstance(tensor, torch.Tensor): 

712 raise RuntimeError(f"{name} must be a torch.Tensor") 

713 if tensor.ndim != 3 or tensor.shape[0] != 1: 

714 raise RuntimeError( 

715 f"{name} must have shape [1, position, width]; got {tuple(tensor.shape)}" 

716 ) 

717 matrix = tensor[0] 

718 if not matrix.is_floating_point(): 

719 raise RuntimeError(f"{name} must have a floating dtype; got {matrix.dtype}") 

720 return cast(Float[torch.Tensor, "position width"], matrix) 

721 

722 

723@contextmanager 

724def _preserve_model_rng(model: Any) -> Iterator[None]: 

725 """Preserve CPU and every CUDA/MPS RNG used by the wrapped model.""" 

726 parameter_devices = {parameter.device for parameter in model.original_model.parameters()} 

727 cuda_devices = sorted( 

728 { 

729 device.index 

730 for device in parameter_devices