Coverage for transformer_lens/tools/analysis/svd_circuits.py: 98%

1"""Singular-vector decomposition of a single attention head's QK and OV maps. 

2 

3An attention head is characterized by two low-rank linear maps: the query-key 

4map ``W_Q W_K^T`` that scores source positions, and the output-value map 

5``W_V W_O`` that writes the attended value back into the residual stream. This 

6tool takes the singular value decomposition of each map for one head and exposes 

7the singular values (how much each direction matters) together with the left and 

8right singular vectors that span the map's input and output spaces. 

9 

10The decomposition is weight-space only: it reads ``W_Q``/``W_K``/``W_V``/``W_O`` 

11and needs no forward pass, no activation cache, and no compatibility mode. Each 

12map is kept factored through 

13:class:`~transformer_lens.FactoredMatrix.FactoredMatrix`, so the 

14``d_model x d_model`` product is never materialized and the returned rank is 

15bounded by ``d_head``. 

16 

17For a factored map ``A @ B`` (``A: [ldim, mdim]``, ``B: [mdim, rdim]``), the SVD's 

18``U`` columns live in ``A``'s input space (``ldim``) and ``V`` columns live in 

19``B``'s output space (``rdim``): feeding ``x = U[:, i]`` through the map gives 

20``x @ (A @ B) == S[i] * V[:, i]``, never the reverse. For OV (``A = W_V_h``, 

21``B = W_O_h``), ``U``'s columns are therefore the value-computation *input* 

22directions this head reads from the residual stream, and ``V``'s columns are the 

23*output* directions it writes back into the residual stream - the ones to 

24project through ``W_U`` for a vocab or logit readout. For QK (``A = W_Q_h``, 

25``B = W_K_h.transpose(-1, -2)``), both ``U`` (destination/query-read) and ``V`` 

26(source/key-read) are read directions; QK only ever produces a scalar attention 

27score, so neither is a write direction. The historical ``.Vh`` alias returns the 

28same tensor as ``.V`` and is never used here. 

29 

30Adjacent singular values closer than a relative gap ``eps`` leave their singular 

31directions defined only up to a rotation, so the result carries a per-direction 

32degeneracy report. Directions are grouped into contiguous blocks that end only 

33at a gap of at least ``eps``; callers attribute such a block as a subspace 

34instead of trusting a single, rotation-dependent direction inside it. 

35 

36A singular value near zero relative to the top of the spectrum is null rather 

37than near-equal: its singular vector is an arbitrary null-space direction, not a 

38rotation of a comparable neighbour. Null directions are flagged under their own 

39tolerance, ``null_rtol``, keyed to the spectrum's top value the way 

40:func:`torch.linalg.matrix_rank` keys its default tolerance. 

41 

42Example:: 

43 

44 from transformer_lens.model_bridge import TransformerBridge 

45 from transformer_lens.tools.analysis.svd_circuits import decompose_head 

46 

47 model = TransformerBridge.boot_transformers("gpt2", device="cpu") 

48 decomposition = decompose_head(model, layer=0, head=0) 

49 ov = decomposition.OV 

50 for row in ov.rank_report: 

51 print(f"direction {row.idx}: sigma={row.sigma:.3f} ratio={row.sigma_ratio:.3f}") 

52""" 

53 

54from dataclasses import dataclass 

55from typing import Callable, List, Literal, Optional, Sequence, Tuple, Union 

56 

57import torch 

58from jaxtyping import Float 

59 

60from transformer_lens.FactoredMatrix import FactoredMatrix 

61from transformer_lens.tools.analysis._model_state import require_eval_mode 

62 

63# Which of a head's two maps to decompose: the query-key map or the output-value map. 

64Which = Literal["QK", "OV"] 

65 

66_VALID_WHICH: Tuple[Which, ...] = ("QK", "OV") 

67 

68# Relative-gap threshold: neighbouring singular values closer than this (in 

69# relative terms) are treated as one rotation-ambiguous block. 

70_DEFAULT_EPS = 1e-2 

71 

72# Absolute floor so relative-gap and ratio computations never divide by ~0. 

73_SIGMA_FLOOR = 1e-12 

74 

75# Default seed for the causal gate's random baseline, so a bare patch_along_directions 

76# call is reproducible instead of drawing from the global RNG. Matches estimate_occupancy, 

77# which seeds its random controls by default. 

78_DEFAULT_BASELINE_SEED = 0 

79 

80# Number of random in-span control subspaces the gate averages its baseline delta over, 

81# so one lucky or unlucky draw does not decide the gate. The per-draw magnitudes are 

82# heavy-tailed, so the averaged threshold still varies with the seed and a direction whose 

83# delta sits near it can gate either way; pass an explicit rng and raise n_baseline when 

84# the verdict is close to the threshold. 

85_DEFAULT_N_BASELINE = 32 

86 

87 

88class DegenerateDirectionError(ValueError): 

89 """Raised when per-direction attribution is requested for a degenerate direction. 

90 

91 Subclasses ``ValueError`` so callers that already ``except ValueError`` keep 

92 working, mirroring how the other analysis tools raise ``ValueError`` for 

93 their input guards. 

94 """ 

95 

96 

97@dataclass 

98class RankReportRow: 

99 """One singular direction's summary, aligned with column ``idx`` of ``U``/``V``. 

100 

101 Attributes: 

102 idx: Position of the direction, matching the column index in ``U`` and ``V``. 

103 sigma: The singular value for this direction. 

104 sigma_ratio: ``sigma`` normalized by the largest singular value, in ``[0, 1]``. 

105 is_degenerate: True when this direction is not attributable on its own: it 

106 shares a block with a neighbour (defined only up to a rotation within 

107 that block) or it is numerically null. 

108 is_null: True when ``sigma_ratio`` falls below ``null_rtol``, so the singular 

109 vector is an arbitrary direction from the map's null space. 

110 block_id: Index of the contiguous block this direction belongs to. 

111 """ 

112 

113 idx: int 

114 sigma: float 

115 sigma_ratio: float 

116 is_degenerate: bool 

117 is_null: bool 

118 block_id: int 

119 

120 

121@dataclass 

122class HeadSVD: 

123 """Factored SVD of one head map (``"QK"`` or ``"OV"``) with a degeneracy report. 

124 

125 Attributes: 

126 which: Which map this decomposes, ``"QK"`` or ``"OV"``. 

127 layer: Layer of the decomposed head. 

128 head: Head index within the layer. 

129 U: Left singular vectors, ``[d_model, rank]``: column i is the map's input 

130 direction i (for OV, the residual-stream direction this head's value 

131 computation reads from; for QK, the destination/query-read direction). 

132 S: Singular values, ``[rank]``, sorted descending. 

133 V: Right singular vectors, ``[d_model, rank]``: column i is the map's output 

134 direction i for OV (the residual-stream direction this head writes into, 

135 the one to project through ``W_U``), or the source/key-read direction for 

136 QK (QK produces no write direction). The reconstruction is 

137 ``U @ S.diag() @ V.transpose(-2, -1)``. 

138 rank_report: Per-direction :class:`RankReportRow` list, aligned with the 

139 columns of ``U``/``V``. 

140 eps: Relative gap below which adjacent directions share a block; every block 

141 boundary sits at a gap of at least ``eps``. 

142 null_rtol: Relative-to-top-singular-value tolerance below which a direction 

143 is numerically null. 

144 folded_ln: Whether the model's weights carried a folded final LayerNorm when 

145 this was decomposed. ``enable_compatibility_mode`` folds ``ln1`` into 

146 ``W_V`` and centres ``W_O``/``W_U``, so a decomposition describes the OV 

147 map only under the state it was built in; the readout and patch consumers 

148 refuse a decomposition whose state no longer matches the model. 

149 """ 

150 

151 which: Which 

152 layer: int 

153 head: int 

154 U: Float[torch.Tensor, "d_model rank"] 

155 S: Float[torch.Tensor, "rank"] 

156 V: Float[torch.Tensor, "d_model rank"] 

157 rank_report: List[RankReportRow] 

158 eps: float 

159 null_rtol: float 

160 folded_ln: bool = False 

161 

162 def is_degenerate(self, i: int) -> bool: 

163 """Whether direction ``i`` is refused: rotation-ambiguous inside a block, or null.""" 

164 return self.rank_report[i].is_degenerate 

165 

166 def block_of(self, i: int) -> List[int]: 

167 """Return every direction index sharing direction ``i``'s degeneracy block. 

168 

169 The result is a singleton ``[i]`` for an isolated direction and the full 

170 run for a degenerate one. 

171 """ 

172 block_id = self.rank_report[i].block_id 

173 return [row.idx for row in self.rank_report if row.block_id == block_id] 

174 

175 def degenerate_blocks(self) -> List[List[int]]: 

176 """Return every degenerate block's indices: the subspaces to attribute whole or skip.""" 

177 blocks: List[List[int]] = [] 

178 current: List[int] = [] 

179 current_block_id: Optional[int] = None 

180 degenerate = False 

181 # block_ids are contiguous and ascending, so one scan recovers the blocks. 

182 for row in self.rank_report: 

183 if row.block_id != current_block_id: 

184 if degenerate: 

185 blocks.append(current) 

186 current = [] 

187 degenerate = False 

188 current_block_id = row.block_id 

189 current.append(row.idx) 

190 degenerate = degenerate or row.is_degenerate 

191 if degenerate: 

192 blocks.append(current) 

193 return blocks 

194 

195 def require_isolated(self, i: int) -> None: 

196 """Raise :class:`DegenerateDirectionError` unless direction ``i`` is attributable alone. 

197 

198 The message names the cause, since a rotation-ambiguous block is still a 

199 subspace worth attributing while a null block carries no signal. 

200 """ 

201 row = self.rank_report[i] 

202 if not row.is_degenerate: 

203 return 

204 where = f"Direction {i} of the {self.which} SVD of head L{self.layer}H{self.head}" 

205 block = self.block_of(i) 

206 if row.is_null: 

207 raise DegenerateDirectionError( 

208 f"{where} is numerically null (sigma_ratio {row.sigma_ratio:.2e} < " 

209 f"null_rtol {self.null_rtol:.2e}), so its singular vector is an arbitrary " 

210 f"null-space direction. Attribute the null block {block} as a subspace, " 

211 f"or skip it." 

212 ) 

213 raise DegenerateDirectionError( 

214 f"{where} lies in block {block}, whose members are not separated by a " 

215 f"relative gap of eps={self.eps:g} and so are defined only up to a rotation. " 

216 f"Attribute the block as a subspace instead of the single direction." 

217 ) 

218 

219 

220@dataclass 

221class HeadDecomposition: 

222 """Container returned by :func:`decompose_head`. 

223 

224 ``QK`` and ``OV`` hold the :class:`HeadSVD` for each requested map, or 

225 ``None`` when that map was not requested. 

226 """ 

227 

228 layer: int 

229 head: int 

230 QK: Optional[HeadSVD] = None 

231 OV: Optional[HeadSVD] = None 

232 

233 

234def _read_weight(weight: torch.Tensor) -> torch.Tensor: 

235 """Detach one per-head weight and put it in SVD precision. 

236 

237 Detached so the returned factors carry no autograd graph into the model. fp16/bf16 

238 are promoted because reduced-precision SVD is unstable; float64 is kept so the 

239 default ``null_rtol`` matches the precision actually decomposed. 

240 """ 

241 weight = weight.detach() 

242 return weight if weight.dtype == torch.float64 else weight.float() 

243 

244 

245def _head_weights( 

246 model, layer: int, head: int 

247) -> Tuple[ 

248 Float[torch.Tensor, "d_model d_head"], 

249 Float[torch.Tensor, "d_model d_head"], 

250 Float[torch.Tensor, "d_model d_head"], 

251 Float[torch.Tensor, "d_head d_model"], 

252]: 

253 """Return ``(W_Q_h, W_K_h, W_V_h, W_O_h)`` for one head, detached, in SVD precision. 

254 

255 Reads the single block's per-head weights rather than the full-model 

256 ``W_Q``/``W_K``/``W_V``/``W_O`` stacks, so only one layer is materialized. On 

257 grouped-query attention ``W_K``/``W_V`` carry one row per key-value head, so the 

258 query head is mapped to its key-value head (query head ``h`` reads kv head 

259 ``h // (n_heads // n_kv_heads)``); a no-op for multi-head attention, where the 

260 head counts already match. 

261 """ 

262 attn = model.blocks[layer].attn 

263 n_kv_heads = attn.W_K.shape[0] 

264 kv_head = head // (model.cfg.n_heads // n_kv_heads) 

265 W_Q_h = _read_weight(attn.W_Q[head]) # [d_model, d_head] 

266 W_K_h = _read_weight(attn.W_K[kv_head]) # [d_model, d_head] 

267 W_V_h = _read_weight(attn.W_V[kv_head]) # [d_model, d_head] 

268 W_O_h = _read_weight(attn.W_O[head]) # [d_head, d_model] 

269 return W_Q_h, W_K_h, W_V_h, W_O_h 

270 

271 

272def _degeneracy_blocks( 

273 S: Float[torch.Tensor, "rank"], eps: float, null_rtol: float 

274) -> List[List[int]]: 

275 """Group singular directions into contiguous blocks separated by relative gaps of ``eps``. 

276 

277 ``S`` is sorted descending. Direction ``i`` joins the open block when 

278 ``1 - S[i]/S[i-1] < eps``, or when it and its predecessor are both null relative to 

279 the top value. Only the gap to the previous direction counts: a block may span far 

280 more than ``eps`` end to end, but every boundary is an ``eps`` gap, and that 

281 separation from the rest of the spectrum is what makes a block stable under 

282 perturbation; no smaller contiguous group inside it is. ``_SIGMA_FLOOR`` keeps the 

283 ratios finite when a divisor is ~0. 

284 """ 

285 n = int(S.shape[0]) 

286 if n == 0: 286 ↛ 287line 286 didn't jump to line 287 because the condition on line 286 was never true

287 return [] 

288 values = [float(x) for x in S.tolist()] 

289 top = max(values[0], _SIGMA_FLOOR) 

290 blocks: List[List[int]] = [] 

291 current = [0] 

292 for i in range(1, n): 

293 prev = max(values[i - 1], _SIGMA_FLOOR) 

294 near_equal = (1.0 - values[i] / prev) < eps 

295 null_run = (values[i] / top) < null_rtol and (values[i - 1] / top) < null_rtol 

296 if near_equal or null_run: 

297 current.append(i) 

298 else: 

299 blocks.append(current) 

300 current = [i] 

301 blocks.append(current) 

302 return blocks 

303 

304 

305def _build_rank_report( 

306 S: Float[torch.Tensor, "rank"], blocks: List[List[int]], null_rtol: float 

307) -> List[RankReportRow]: 

308 """Summarize each singular direction and tag its degeneracy block. 

309 

310 ``blocks`` must partition ``range(len(S))``. Nullness flags a direction on its own: 

311 a lone null singular vector is arbitrary even though nothing groups with it. 

312 """ 

313 values = [float(x) for x in S.tolist()] 

314 top = max(values) if values else 0.0 

315 denominator = top if top > _SIGMA_FLOOR else _SIGMA_FLOOR 

316 rows: List[Optional[RankReportRow]] = [None] * len(values) 

317 for block_id, block in enumerate(blocks): 

318 shared = len(block) > 1 

319 for idx in block: 

320 sigma_ratio = values[idx] / denominator 

321 is_null = sigma_ratio < null_rtol 

322 rows[idx] = RankReportRow( 

323 idx=idx, 

324 sigma=values[idx], 

325 sigma_ratio=sigma_ratio, 

326 is_degenerate=shared or is_null, 

327 is_null=is_null, 

328 block_id=block_id, 

329 ) 

330 return [row for row in rows if row is not None] 

331 

332 

333def _factored_head_svd( 

334 A: Float[torch.Tensor, "d_model d_head"], 

335 B: Float[torch.Tensor, "d_head d_model"], 

336 *, 

337 which: Which, 

338 layer: int, 

339 head: int, 

340 eps: float, 

341 null_rtol: Optional[float] = None, 

342 folded_ln: bool = False, 

343) -> HeadSVD: 

344 """Decompose the factored map ``A @ B`` for one head into a :class:`HeadSVD`. 

345 

346 For OV pass ``A = W_V_h`` and ``B = W_O_h``; for QK pass ``A = W_Q_h`` and 

347 ``B = W_K_h.transpose(-1, -2)``. The map stays factored through 

348 :class:`FactoredMatrix`, so the ``d_model x d_model`` product is never 

349 materialized and the rank is bounded by ``d_head``. 

350 

351 ``null_rtol`` of ``None`` resolves to ``d_model * torch.finfo(S.dtype).eps``, 

352 the same relative tolerance :func:`torch.linalg.matrix_rank` uses by default 

353 for a square ``d_model x d_model`` map, so the null cutoff tracks the 

354 decomposition's own numerical rank rather than a hand-tuned constant. 

355 

356 ``folded_ln`` is whether the model's weights carried a folded final LayerNorm, 

357 recorded on the result so the consumers can refuse a decomposition taken under a 

358 different state. 

359 """ 

360 U, S, V = FactoredMatrix(A, B).svd() 

361 d_model = U.shape[0] 

362 resolved_null_rtol = null_rtol if null_rtol is not None else d_model * torch.finfo(S.dtype).eps 

363 blocks = _degeneracy_blocks(S, eps, null_rtol=resolved_null_rtol) 

364 rank_report = _build_rank_report(S, blocks, null_rtol=resolved_null_rtol) 

365 return HeadSVD( 

366 which=which, 

367 layer=layer, 

368 head=head, 

369 U=U, 

370 S=S, 

371 V=V, 

372 rank_report=rank_report, 

373 eps=eps, 

374 null_rtol=resolved_null_rtol, 

375 folded_ln=folded_ln, 

376 ) 

377 

378 

379def decompose_head( 

380 model, 

381 layer: int, 

382 head: int, 

383 *, 

384 which: Sequence[str] = ("QK", "OV"), 

385 eps: float = _DEFAULT_EPS, 

386 null_rtol: Optional[float] = None, 

387) -> HeadDecomposition: 

388 """Decompose a head's QK (``W_Q W_K^T``) and/or OV (``W_V W_O``) maps via SVD. 

389 

390 Weight-space only: this reads the head's per-block weights via the bridge's 

391 ``model.blocks[layer].attn`` accessors and needs no forward pass and no 

392 compatibility mode. The returned factors are detached from the model. Whether the 

393 model's weights carry a folded final LayerNorm is recorded on each returned 

394 :class:`HeadSVD` so the readout and patch consumers can refuse a decomposition 

395 taken under a different state. 

396 

397 Args: 

398 model: A ``TransformerBridge``. 

399 layer: Layer of the head to decompose. 

400 head: Head index within the layer. 

401 which: Which maps to decompose, a non-empty sequence drawn from 

402 ``("QK", "OV")``. A bare string is rejected rather than iterated. 

403 eps: Relative gap at which a block of adjacent directions ends; directions 

404 closer than this share a block. 

405 null_rtol: Relative-to-top-singular-value tolerance below which a direction 

406 counts as numerically null. Defaults to ``None``, which resolves to 

407 ``d_model * torch.finfo(S.dtype).eps`` per map, matching 

408 :func:`torch.linalg.matrix_rank`'s default tolerance. 

409 

410 Returns: 

411 A :class:`HeadDecomposition` whose ``QK``/``OV`` fields hold a 

412 :class:`HeadSVD` for each requested map. 

413 

414 Raises: 

415 ValueError: If ``layer`` or ``head`` is out of range, or ``which`` is 

416 empty, a bare string, or contains an unknown entry. 

417 """ 

418 n_layers = model.cfg.n_layers 

419 n_heads = model.cfg.n_heads 

420 if not 0 <= layer < n_layers: 

421 raise ValueError(f"layer must be in [0, {n_layers}), got {layer!r}") 

422 if not 0 <= head < n_heads: 

423 raise ValueError(f"head must be in [0, {n_heads}), got {head!r}") 

424 if isinstance(which, str): 

425 raise ValueError(f"which must be a sequence of map names such as ('QK',), got {which!r}") 

426 requested = tuple(which) 

427 if not requested: 

428 raise ValueError("which must request at least one of 'QK' or 'OV'") 

429 invalid = [entry for entry in requested if entry not in _VALID_WHICH] 

430 if invalid: 

431 raise ValueError(f"which entries must be in {_VALID_WHICH}, got {invalid!r}") 

432 

433 W_Q_h, W_K_h, W_V_h, W_O_h = _head_weights(model, layer, head) 

434 folded_ln = _folded_ln_state(model) 

435 qk = None 

436 ov = None 

437 if "QK" in requested: 

438 qk = _factored_head_svd( 

439 W_Q_h, 

440 W_K_h.transpose(-1, -2), 

441 which="QK", 

442 layer=layer, 

443 head=head, 

444 eps=eps, 

445 null_rtol=null_rtol, 

446 folded_ln=folded_ln, 

447 ) 

448 if "OV" in requested: 

449 ov = _factored_head_svd( 

450 W_V_h, 

451 W_O_h, 

452 which="OV", 

453 layer=layer, 

454 head=head, 

455 eps=eps, 

456 null_rtol=null_rtol, 

457 folded_ln=folded_ln, 

458 ) 

459 return HeadDecomposition(layer=layer, head=head, QK=qk, OV=ov) 

460 

461 

462def _folded_ln_state(model) -> bool: 

463 """Whether the model's weights currently carry a folded final LayerNorm. 

464 

465 ``enable_compatibility_mode`` sets ``compatibility_mode`` before it processes 

466 weights and skips processing entirely under ``no_processing``, so the flag alone 

467 does not imply a fold. The fold happens only when the bridge has processed its 

468 weights with ``fold_ln`` enabled *and* the adapter supports folding: adapters that 

469 declare ``supports_fold_ln = False`` have ``fold_ln`` downgraded to ``False`` 

470 inside ``ProcessWeights``, leaving ``_fold_ln_requested`` set while the norms stay 

471 unfolded. 

472 """ 

473 processed = getattr(model, "_weights_processed", False) 

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

475 fold_ln_requested = getattr(adapter, "_fold_ln_requested", True) 

476 supports_fold_ln = getattr(adapter, "supports_fold_ln", True) 

477 return bool(processed and fold_ln_requested and supports_fold_ln) 

478 

479 

480def _validate_bridge_compatibility(model) -> None: 

481 """Reject a ``TransformerBridge`` whose ``W_U`` would give a silently wrong projection. 

482 

483 Projecting an OV direction through ``W_U`` requires the final LayerNorm folded 

484 into ``W_U``. This mirrors :func:`direct_logit_attribution`, the other 

485 unembedding-touching analysis tool, but the check here is stricter: it reads the 

486 bridge's recorded weight-processing state rather than the ``compatibility_mode`` 

487 flag. ``enable_compatibility_mode`` sets that flag before it processes weights and 

488 skips processing entirely under ``no_processing``, so the flag can be ``True`` while 

489 ``W_U`` is still unfolded. The fold only happens once the bridge processes its 

490 weights with ``fold_ln`` enabled, which is what ``_weights_processed``, the 

491 adapter's ``_fold_ln_requested``, and the adapter's ``supports_fold_ln`` record: 

492 adapters that declare no fold support have ``fold_ln`` downgraded to ``False`` 

493 inside ``ProcessWeights``, so their norms stay unfolded even though the bridge 

494 reports processed weights. No hybrid-architecture restriction applies: 

495 projecting a rank-1 OV direction through ``W_U`` does not depend on the block-layout 

496 assumptions that check exists for elsewhere. 

497 """ 

498 # Lazy import - keeps the module importable without the bridge as a hard dependency. 

499 from transformer_lens.model_bridge import TransformerBridge 

500 

501 if not isinstance(model, TransformerBridge): 

502 return 

503 if not _folded_ln_state(model): 

504 raise ValueError( 

505 "Projecting an OV direction through W_U on a TransformerBridge requires the " 

506 "final LayerNorm folded into W_U, which only happens once the bridge has " 

507 "processed its weights with fold_ln enabled. Call " 

508 "`model.enable_compatibility_mode()` (its default folds LayerNorm) after " 

509 "loading the bridge, then retry." 

510 ) 

511 

512 

513def _validate_decomposition_matches_model(model, head_svd: HeadSVD) -> None: 

514 """Refuse a decomposition built under a different folded-LayerNorm state than the model. 

515 

516 ``enable_compatibility_mode`` folds ``ln1`` into ``W_V`` and centres ``W_O``/``W_U``, 

517 so a :class:`HeadSVD` decomposed before the call describes a different OV map than the 

518 model now computes: the cached ``U``/``S``/``V`` are stale, and the returned readout or 

519 patch would be silently wrong rather than raise a shape error. This guard is orthogonal 

520 to :func:`_validate_bridge_compatibility` (which requires the fold to be present): here 

521 the model may be in either state, only mismatched from the decomposition's. 

522 """ 

523 current = _folded_ln_state(model) 

524 if current != head_svd.folded_ln: 

525 raise ValueError( 

526 f"This HeadSVD was decomposed with folded_ln={head_svd.folded_ln} but the " 

527 f"model's weights now have folded_ln={current}; the cached singular vectors " 

528 f"describe a different OV map. Re-run decompose_head under the current state, " 

529 f"then retry." 

530 ) 

531 

532 

533def vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) -> Float[torch.Tensor, "d_vocab k"]: 

534 """Project the top-k OV output directions through the unembedding. 

535 

536 Requires ``head_svd.which == "OV"``: QK produces no write direction to project 

537 (see the module docstring). On a ``TransformerBridge``, the final LayerNorm must 

538 be folded into ``W_U``, which requires the bridge to have processed its weights 

539 with ``fold_ln`` enabled on an adapter that supports folding. 

540 

541 Does not call ``head_svd.require_isolated``: a degenerate direction's vocab 

542 readout is still a well-defined projection, unlike a per-direction causal claim, 

543 so it is not gated here. The contract that no direction is reported without a 

544 passing causal patch is enforced by :func:`patch_along_directions`. 

545 

546 Args: 

547 model: A ``TransformerBridge`` with the final LayerNorm folded into ``W_U``; 

548 only its ``W_U`` is read. 

549 head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. 

550 k: Number of top singular directions to project. 

551 

552 Returns: 

553 ``W_U.T @ head_svd.V[:, :k]``, shape ``[d_vocab, k]``: column i is 

554 direction i's projection through the unembedding. 

555 

556 Raises: 

557 ValueError: If ``head_svd.which != "OV"``, if ``k`` is not in 

558 ``(0, rank]``, if ``head_svd`` was decomposed under a different 

559 folded-LayerNorm state than ``model`` now has, or if ``model`` is a 

560 ``TransformerBridge`` whose weights do not carry a folded final 

561 LayerNorm. 

562 """ 

563 if head_svd.which != "OV": 

564 raise ValueError(f"vocab_readout requires an OV HeadSVD, got which={head_svd.which!r}") 

565 _validate_decomposition_matches_model(model, head_svd) 

566 rank = head_svd.V.shape[1] 

567 if not 0 < k <= rank: 

568 raise ValueError(f"k must be in (0, {rank}], got {k!r}") 

569 _validate_bridge_compatibility(model) 

570 # Promote V to W_U's dtype rather than forcing float32: matmul does not promote its 

571 # operands, so a bare .float() raises a dtype mismatch on any bf16, fp16, or fp64 model. 

572 W_U = model.W_U 

573 return W_U.T @ head_svd.V[:, :k].to(W_U.dtype) 

574 

575 

576@dataclass 

577class LogitSignature: 

578 """Rank-1-reconstruction logit effect for one OV direction, per requested token. 

579 

580 Attributes: 

581 direction: Which ``HeadSVD`` column this reconstructs. 

582 values: Signed logit contribution, aligned with the requested tokens. 

583 """ 

584 

585 direction: int 

586 values: Float[torch.Tensor, "token"] 

587 

588 

589def logit_signature( 

590 model, 

591 head_svd: HeadSVD, 

592 direction: int, 

593 tokens: Union[int, Sequence[int], torch.Tensor], 

594) -> LogitSignature: 

595 """Signed logit effect of one OV direction's rank-1 reconstruction on the given tokens. 

596 

597 Pure weight-space computation: reconstructs the head's OV output along a single 

598 singular direction (``S[direction] * V[:, direction]``, never ``U`` - see the 

599 module docstring) and projects it through ``W_U`` restricted to ``tokens``. Runs 

600 no forward pass and builds no cache. 

601 

602 Args: 

603 model: A ``TransformerBridge`` with the final LayerNorm folded into ``W_U``; 

604 only its ``W_U`` is read. 

605 head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. 

606 direction: Column index of the singular direction to reconstruct. 

607 tokens: Token id(s) to read the logit effect for. 

608 

609 Returns: 

610 A :class:`LogitSignature` with one value per requested token. 

611 

612 Raises: 

613 ValueError: If ``head_svd.which != "OV"``, if ``direction`` is not in 

614 ``[0, rank)``, if ``head_svd`` was decomposed under a different 

615 folded-LayerNorm state than ``model`` now has, or if ``model`` is a 

616 ``TransformerBridge`` whose weights do not carry a folded final 

617 LayerNorm. 

618 DegenerateDirectionError: If ``direction`` is not attributable alone (see 

619 :meth:`HeadSVD.require_isolated`). 

620 """ 

621 if head_svd.which != "OV": 621 ↛ 622line 621 didn't jump to line 622 because the condition on line 621 was never true

622 raise ValueError(f"logit_signature requires an OV HeadSVD, got which={head_svd.which!r}") 

623 _validate_decomposition_matches_model(model, head_svd) 

624 # Bounds-check before indexing: a negative direction would wrap into V/rank_report and a 

625 # too-large one would raise a bare IndexError, both hiding a caller mistake as a wrong or 

626 # cryptic result rather than a clear refusal. 

627 rank = head_svd.V.shape[1] 

628 direction = int(direction) 

629 if not 0 <= direction < rank: 

630 raise ValueError(f"direction index {direction} out of range [0, {rank})") 

631 head_svd.require_isolated(direction) 

632 _validate_bridge_compatibility(model) 

633 token_ids = torch.as_tensor(tokens, dtype=torch.long).reshape(-1) 

634 # Match the rank-1 reconstruction to W_U's dtype for the projection: matmul does not 

635 # promote its operands, so a float32 reconstruction breaks on bf16, fp16, or fp64 models. 

636 W_U = model.W_U 

637 reconstruction = (head_svd.S[direction] * head_svd.V[:, direction]).to(W_U.dtype) 

638 values = reconstruction @ W_U[:, token_ids] 

639 return LogitSignature(direction=direction, values=values) 

640 

641 

642@dataclass 

643class ActivationProjection: 

644 """Per-position coefficients of a head's actual output in its OV output basis. 

645 

646 Attributes: 

647 head_svd: The OV decomposition this was projected against. 

648 coefficients: ``[pos, rank]``; ``coefficients[:, i]`` is the signed amount of 

649 singular direction ``i`` (``head_svd.V[:, i]``) present in the head's actual 

650 output at each position. Summing ``coefficients[:, i] * head_svd.V[:, i]`` 

651 over ``i`` reconstructs the head's real per-position output to numerical 

652 precision, since ``V``'s columns are orthonormal and this projects onto the 

653 exact basis the head writes in. 

654 str_tokens: Tokenized prompt, aligned with the position axis, for display. 

655 """ 

656 

657 head_svd: HeadSVD 

658 coefficients: Float[torch.Tensor, "pos rank"] 

659 str_tokens: List[str] 

660 

661 

662def project_activations( 

663 model, head_svd: HeadSVD, prompt: Union[str, torch.Tensor] 

664) -> ActivationProjection: 

665 """Project a head's actual per-position output onto its OV singular directions. 

666 

667 Requires ``head_svd.which == "OV"``: this projects onto the write/output basis 

668 ``V``, and QK has no such vector (see the module docstring). Runs a real forward 

669 pass with ``use_attn_result`` enabled to read the per-head output 

670 (``hook_result``), then projects it onto ``head_svd.V``. Restores the model's 

671 prior ``use_attn_result`` setting afterward, since flipping that config flag as a 

672 side effect of a read-only analysis call would surprise a caller who already had 

673 hooks or a cache built around its prior state. 

674 

675 Args: 

676 model: A ``TransformerBridge``. 

677 head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. 

678 prompt: A single prompt (not a batch): a string or a ``[pos]`` or ``[1, pos]`` 

679 token tensor. 

680 

681 Returns: 

682 An :class:`ActivationProjection` with the per-position coefficients. 

683 

684 Raises: 

685 ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under 

686 a different folded-LayerNorm state than ``model`` now has, if ``prompt`` 

687 is not a single prompt (a tensor that is not 1-D, or is 2-D with a leading 

688 dimension greater than one), or if the model is in training mode. 

689 NotImplementedError: If the model's attention adapter exposes no per-head result, 

690 so ``set_use_attn_result(True)`` cannot fork the attention output. 

691 """ 

692 if head_svd.which != "OV": 

693 raise ValueError( 

694 f"project_activations requires an OV HeadSVD, got which={head_svd.which!r}" 

695 ) 

696 _validate_decomposition_matches_model(model, head_svd) 

697 # A single prompt is a string or a 1-D [pos] / 2-D [1, pos] token tensor. Reject any 

698 # other rank here, before the forward: a 0-D tensor would raise a bare IndexError from 

699 # the shape probe below, and a [1, 1, pos] tensor would slip past a leading-dimension 

700 # test and only fail after the forward. A list of prompt strings reveals its batch size 

701 # only once tokenized, so the post-cache check below still guards that path. 

702 if isinstance(prompt, torch.Tensor) and ( 

703 prompt.ndim not in (1, 2) or (prompt.ndim == 2 and prompt.shape[0] > 1) 

704 ): 

705 raise ValueError( 

706 "project_activations requires a single prompt, got a token tensor of shape " 

707 f"{tuple(prompt.shape)}; pass a [pos] or [1, pos] tensor, or a single string." 

708 ) 

709 require_eval_mode(model, operation="project_activations()") 

710 # Cache only the one hook this reads. run_with_cache otherwise retains every hook point of 

711 # the forward pass (about 1 GB against 19 MB on a 512-token gpt2-small prompt) to read a 

712 # single head's output. 

713 hook_name = f"blocks.{head_svd.layer}.attn.hook_result" 

714 previous = getattr(model.cfg, "use_attn_result", False) 

715 model.set_use_attn_result(True) 

716 try: 

717 _, cache = model.run_with_cache(prompt, names_filter=lambda name: name == hook_name) 

718 finally: 

719 model.set_use_attn_result(previous) 

720 result = cache[hook_name][..., head_svd.head, :] 

721 if result.shape[0] != 1: 721 ↛ 722line 721 didn't jump to line 722 because the condition on line 721 was never true

722 raise ValueError( 

723 f"project_activations requires a single prompt, got batch={result.shape[0]}" 

724 ) 

725 result = result.squeeze(0).to(head_svd.V.dtype) 

726 coefficients = result @ head_svd.V 

727 str_tokens = model.to_str_tokens(prompt) 

728 return ActivationProjection(head_svd=head_svd, coefficients=coefficients, str_tokens=str_tokens) 

729 

730 

731def _validate_retained_blocks(head_svd: HeadSVD, retained: Sequence[int]) -> None: 

732 """Raise if ``retained`` splits a degenerate block instead of keeping it whole or empty. 

733 

734 A degenerate block's members are defined only up to a rotation (or, for a null 

735 block, arbitrary null-space vectors), so attributing a causal effect to part of the 

736 block while dropping the rest would let a caller route around the guard 

737 :meth:`HeadSVD.require_isolated` already enforces per direction. 

738 """ 

739 retained_set = set(retained) 

740 for block in head_svd.degenerate_blocks(): 

741 block_set = set(block) 

742 overlap = retained_set & block_set 

743 if overlap and overlap != block_set: 

744 raise DegenerateDirectionError( 

745 f"retained directions {sorted(overlap)} split block {sorted(block_set)} of " 

746 f"the {head_svd.which} SVD of head L{head_svd.layer}H{head_svd.head}, whose " 

747 f"members are not separated by a relative gap of eps={head_svd.eps:g} or are " 

748 f"jointly null. Keep or ablate the whole block, not part of it." 

749 ) 

750 

751 

752def _resolve_retained( 

753 head_svd: HeadSVD, keep: Optional[Sequence[int]], ablate: Optional[Sequence[int]] 

754) -> List[int]: 

755 """Resolve ``keep``/``ablate`` to the sorted list of retained direction indices. 

756 

757 Exactly one of ``keep``/``ablate`` must be given; ``ablate``'s complement over the 

758 map's full rank becomes the retained set. Every supplied index is coerced to ``int`` 

759 and bounds-checked against ``[0, rank)`` before use, so a negative index raises rather 

760 than wrapping into ``V[:, ...]`` and an out-of-range ``ablate`` raises rather than 

761 silently subtracting nothing from the complement. Raises 

762 :class:`DegenerateDirectionError` if the result would split a degenerate block (see 

763 :func:`_validate_retained_blocks`). 

764 """ 

765 if (keep is None) == (ablate is None): 

766 raise ValueError("patch_along_directions requires exactly one of keep or ablate") 

767 rank = head_svd.V.shape[1] 

768 

769 def _checked(indices: Sequence[int], name: str) -> List[int]: 

770 resolved: List[int] = [] 

771 for raw in indices: 

772 index = int(raw) 

773 if not 0 <= index < rank: 

774 raise ValueError(f"{name} index {index} out of range [0, {rank})") 

775 resolved.append(index) 

776 return resolved 

777 

778 if keep is not None: 

779 retained = sorted(set(_checked(keep, "keep"))) 

780 else: 

781 assert ablate is not None 

782 ablate_set = set(_checked(ablate, "ablate")) 

783 retained = [i for i in range(rank) if i not in ablate_set] 

784 _validate_retained_blocks(head_svd, retained) 

785 return retained 

786 

787 

788def _make_subspace_hook(head: int, projector: Float[torch.Tensor, "d_model d_model"]): 

789 """Build a ``hook_result`` hook that reconstructs one head's output onto ``span(projector)``. 

790 

791 Leaves every other head's slice of the ``[batch, pos, head_index, d_model]`` tensor 

792 untouched. Clones before mutating so the hook never writes into the activation the 

793 forward pass itself is still using. 

794 """ 

795 

796 def hook_fn(activation: torch.Tensor, hook) -> torch.Tensor: 

797 activation = activation.clone() 

798 # The projector is drawn on CPU (QR is unimplemented on MPS and a CUDA generator 

799 # cannot feed a CPU randn), so move it to the activation's device and dtype here. 

800 activation[:, :, head, :] = activation[:, :, head, :] @ projector.to( 

801 device=activation.device, dtype=activation.dtype 

802 ) 

803 return activation 

804 

805 return hook_fn 

806 

807 

808@dataclass 

809class PatchResult: 

810 """Result of causally patching a head's output onto a chosen OV singular subspace. 

811 

812 Attributes: 

813 head_svd: The OV decomposition patched against. 

814 retained: Direction indices whose span the head's output was reconstructed 

815 onto; the complement was zeroed. 

816 original_metric: Metric value on the unmodified prompt. 

817 patched_metric: Metric value after the subspace reconstruction. 

818 delta_metric: ``patched_metric - original_metric``. 

819 baseline_delta_metric: the mean of the per-draw ``delta_metric`` magnitudes over 

820 several random control subspaces of the same width as ``retained``, each drawn 

821 inside the head's own OV span ``span(V)`` rather than from the full residual 

822 stream, so the control is the effect of an arbitrary same-size subspace of this 

823 head's output rather than of an unrelated residual-stream direction. Averaging 

824 the magnitudes rather than the signed deltas keeps this a typical control 

825 effect that does not shrink when the controls mix sign, so it is non-negative. 

826 gated: whether the retained subspace passed the causal test for the mode it was 

827 expressed in, comparing ``abs(delta_metric)`` against ``baseline_delta_metric`` 

828 (or an explicit threshold, if one was passed). For ``ablate`` (retain the complement), 

829 removing a load-bearing subspace should move the metric more than removing an 

830 arbitrary same-size one, so ``gated`` is ``abs(delta_metric) > threshold``. 

831 For ``keep`` (retain only the given subspace), a subspace that reconstructs 

832 the head's behavior should move the metric less than keeping an arbitrary 

833 same-size one, so ``gated`` is ``abs(delta_metric) < threshold``. A single 

834 "moved more than baseline" test cannot answer both, since ``keep=S`` and 

835 ``ablate=complement(S)`` resolve to the same retained set. 

836 """ 

837 

838 head_svd: HeadSVD 

839 retained: List[int] 

840 original_metric: float 

841 patched_metric: float 

842 delta_metric: float 

843 baseline_delta_metric: float 

844 gated: bool 

845 

846 

847@torch.no_grad() 

848def patch_along_directions( 

849 model, 

850 head_svd: HeadSVD, 

851 prompt: Union[str, torch.Tensor], 

852 metric: Callable[[torch.Tensor], float], 

853 *, 

854 keep: Optional[Sequence[int]] = None, 

855 ablate: Optional[Sequence[int]] = None, 

856 threshold: Optional[float] = None, 

857 rng: Optional[torch.Generator] = None, 

858 n_baseline: int = _DEFAULT_N_BASELINE, 

859) -> PatchResult: 

860 """Causally validate a claimed OV subfunction by reconstructing the head's output onto it. 

861 

862 Requires ``head_svd.which == "OV"``: this reconstructs the write/output basis 

863 ``V``, and QK has no such vector (see the module docstring). Runs the prompt with 

864 ``use_attn_result`` enabled once unmodified, once with the head's ``hook_result`` 

865 slice reconstructed onto ``span(head_svd.V[:, retained])``, and once per random 

866 control subspace. Each control subspace is drawn *inside the head's own OV span* 

867 ``span(V)`` (not from the full residual stream, where a width-``w`` random subspace 

868 would keep only ``w/d_model`` of a head output that itself occupies only 

869 ``rank/d_model`` of the stream; an in-span control of width ``w`` keeps ``w/rank`` 

870 of it), so a moved metric is compared against the effect of an arbitrary 

871 subspace of this head's output of the same width. The per-draw control delta magnitudes 

872 are averaged over ``n_baseline`` draws so one lucky or unlucky draw does not decide the 

873 gate, and so controls that mix sign do not cancel into a smaller threshold. Averaging 

874 narrows but does not remove the seed dependence of the threshold; see the ``n_baseline`` 

875 entry below. Restores the model's prior ``use_attn_result`` setting afterward. 

876 

877 The gate's success condition depends on the mode the caller expressed, because 

878 ``keep=S`` and ``ablate=complement(S)`` resolve to the same retained set and a single 

879 "moved more than the control" test would answer only the ``ablate`` question. See 

880 :attr:`PatchResult.gated`. 

881 

882 Args: 

883 model: A ``TransformerBridge``. 

884 head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. 

885 prompt: A single prompt: a string or a ``[1, pos]`` token tensor. 

886 metric: A function from the model's logits to a scalar. 

887 keep: Direction indices to retain; the rest are zeroed. Exactly one of 

888 ``keep``/``ablate`` must be given. 

889 ablate: Direction indices to zero; the rest are retained. 

890 threshold: Explicit gate threshold. Defaults to ``None``, which uses 

891 ``baseline_delta_metric`` (already a magnitude) instead. 

892 rng: Optional generator for the random control subspaces, for reproducibility. 

893 Defaults to a generator seeded with ``_DEFAULT_BASELINE_SEED`` so a bare 

894 call is reproducible rather than drawing from the global RNG. 

895 n_baseline: Number of random in-span control subspaces to average the baseline 

896 delta over. Must be at least 1. The per-draw control magnitudes are 

897 heavy-tailed, so the averaged threshold still varies with the seed and a 

898 direction whose delta sits near it can gate either way; pass an explicit 

899 ``rng`` for a reproducible verdict and raise ``n_baseline`` when the verdict 

900 is close to the threshold. Each draw costs one forward pass. 

901 

902 Returns: 

903 A :class:`PatchResult` describing the patched, baseline, and original metrics. 

904 

905 Raises: 

906 ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under a 

907 different folded-LayerNorm state than ``model`` now has, if ``keep``/``ablate`` 

908 are both given or both omitted, if any index is out of ``[0, rank)``, if the 

909 retained set is empty (``keep=[]`` or ``ablate`` over the full rank) or spans the 

910 full rank (``keep`` over every direction) and no explicit ``threshold`` is 

911 supplied, if ``n_baseline < 1``, or if the model is in training mode. 

912 NotImplementedError: If the model's attention adapter exposes no per-head result, 

913 so ``set_use_attn_result(True)`` cannot fork the attention output. 

914 DegenerateDirectionError: If the retained directions split a degenerate block 

915 (see :func:`_validate_retained_blocks`). 

916 """ 

917 if head_svd.which != "OV": 

918 raise ValueError( 

919 f"patch_along_directions requires an OV HeadSVD, got which={head_svd.which!r}" 

920 ) 

921 if n_baseline < 1: 

922 raise ValueError(f"n_baseline must be at least 1, got {n_baseline}") 

923 _validate_decomposition_matches_model(model, head_svd) 

924 retained = _resolve_retained(head_svd, keep, ablate) 

925 # _resolve_retained has confirmed exactly one of keep/ablate is set, so the mode the 

926 # caller expressed is unambiguous and selects the gate's success condition (see below). 

927 mode = "keep" if keep is not None else "ablate" 

928 V = head_svd.V 

929 rank = V.shape[1] 

930 if (not retained or len(retained) == rank) and threshold is None: 

931 # An empty retained set reconstructs the head onto the zero subspace; a full-rank 

932 # retained set reconstructs it onto span(V) itself (V V^T). In both cases the kept 

933 # projector and the equal-width random control projector coincide: for the empty set 

934 # both are the zero projector, and for the full set every width-rank in-span control 

935 # is V Q Q^T V^T = V V^T. The two deltas then tie by construction and the gate 

936 # compares a quantity against itself. Require an explicit threshold to gate either 

937 # degenerate set on purpose. 

938 raise ValueError( 

939 "keep/ablate retains either no directions or every direction, so the "