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

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 (the output and input directions they act on). 

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 

17Right singular vectors are read from :attr:`~transformer_lens.FactoredMatrix.FactoredMatrix.V` 

18(its columns are the right singular vectors). The historical ``.Vh`` alias is 

19deprecated and returns the same tensor, so it is never used here. 

20 

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

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

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

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

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

26 

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

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

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

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

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

32 

33Example:: 

34 

35 from transformer_lens.model_bridge import TransformerBridge 

36 from transformer_lens.tools.analysis.svd_circuits import decompose_head 

37 

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

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

40 ov = decomposition.OV 

41 for row in ov.rank_report: 

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

43""" 

44 

45from dataclasses import dataclass 

46from typing import List, Literal, Optional, Sequence, Tuple 

47 

48import torch 

49from jaxtyping import Float 

50 

51from transformer_lens.FactoredMatrix import FactoredMatrix 

52 

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

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

55 

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

57 

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

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

60_DEFAULT_EPS = 1e-2 

61 

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

63_SIGMA_FLOOR = 1e-12 

64 

65 

66class DegenerateDirectionError(ValueError): 

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

68 

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

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

71 their input guards. 

72 """ 

73 

74 

75@dataclass 

76class RankReportRow: 

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

78 

79 Attributes: 

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

81 sigma: The singular value for this direction. 

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

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

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

85 that block) or it is numerically null. 

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

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

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

89 """ 

90 

91 idx: int 

92 sigma: float 

93 sigma_ratio: float 

94 is_degenerate: bool 

95 is_null: bool 

96 block_id: int 

97 

98 

99@dataclass 

100class HeadSVD: 

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

102 

103 Attributes: 

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

105 layer: Layer of the decomposed head. 

106 head: Head index within the layer. 

107 U: Left singular vectors, ``[d_model, rank]`` (column i is output direction i). 

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

109 V: Right singular vectors, ``[d_model, rank]`` (column i is input direction i). 

110 The reconstruction is ``U @ S.diag() @ V.transpose(-2, -1)``. 

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

112 columns of ``U``/``V``. 

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

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

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

116 is numerically null. 

117 """ 

118 

119 which: Which 

120 layer: int 

121 head: int 

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

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

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

125 rank_report: List[RankReportRow] 

126 eps: float 

127 null_rtol: float 

128 

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

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

131 return self.rank_report[i].is_degenerate 

132 

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

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

135 

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

137 run for a degenerate one. 

138 """ 

139 block_id = self.rank_report[i].block_id 

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

141 

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

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

144 blocks: List[List[int]] = [] 

145 current: List[int] = [] 

146 current_block_id: Optional[int] = None 

147 degenerate = False 

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

149 for row in self.rank_report: 

150 if row.block_id != current_block_id: 

151 if degenerate: 

152 blocks.append(current) 

153 current = [] 

154 degenerate = False 

155 current_block_id = row.block_id 

156 current.append(row.idx) 

157 degenerate = degenerate or row.is_degenerate 

158 if degenerate: 

159 blocks.append(current) 

160 return blocks 

161 

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

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

164 

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

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

167 """ 

168 row = self.rank_report[i] 

169 if not row.is_degenerate: 

170 return 

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

172 block = self.block_of(i) 

173 if row.is_null: 

174 raise DegenerateDirectionError( 

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

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

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

178 f"or skip it." 

179 ) 

180 raise DegenerateDirectionError( 

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

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

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

184 ) 

185 

186 

187@dataclass 

188class HeadDecomposition: 

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

190 

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

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

193 """ 

194 

195 layer: int 

196 head: int 

197 QK: Optional[HeadSVD] = None 

198 OV: Optional[HeadSVD] = None 

199 

200 

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

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

203 

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

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

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

207 """ 

208 weight = weight.detach() 

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

210 

211 

212def _head_weights( 

213 model, layer: int, head: int 

214) -> Tuple[ 

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

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

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

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

219]: 

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

221 

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

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

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

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

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

227 head counts already match. 

228 """ 

229 attn = model.blocks[layer].attn 

230 n_kv_heads = attn.W_K.shape[0] 

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

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

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

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

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

236 return W_Q_h, W_K_h, W_V_h, W_O_h 

237 

238 

239def _degeneracy_blocks( 

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

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

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

243 

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

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

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

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

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

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

250 ratios finite when a divisor is ~0. 

251 """ 

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

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

254 return [] 

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

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

257 blocks: List[List[int]] = [] 

258 current = [0] 

259 for i in range(1, n): 

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

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

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

263 if near_equal or null_run: 

264 current.append(i) 

265 else: 

266 blocks.append(current) 

267 current = [i] 

268 blocks.append(current) 

269 return blocks 

270 

271 

272def _build_rank_report( 

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

274) -> List[RankReportRow]: 

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

276 

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

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

279 """ 

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

281 top = max(values) if values else 0.0 

282 denominator = top if top > _SIGMA_FLOOR else _SIGMA_FLOOR 

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

284 for block_id, block in enumerate(blocks): 

285 shared = len(block) > 1 

286 for idx in block: 

287 sigma_ratio = values[idx] / denominator 

288 is_null = sigma_ratio < null_rtol 

289 rows[idx] = RankReportRow( 

290 idx=idx, 

291 sigma=values[idx], 

292 sigma_ratio=sigma_ratio, 

293 is_degenerate=shared or is_null, 

294 is_null=is_null, 

295 block_id=block_id, 

296 ) 

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

298 

299 

300def _factored_head_svd( 

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

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

303 *, 

304 which: Which, 

305 layer: int, 

306 head: int, 

307 eps: float, 

308 null_rtol: Optional[float] = None, 

309) -> HeadSVD: 

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

311 

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

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

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

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

316 

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

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

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

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

321 """ 

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

323 d_model = U.shape[0] 

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

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

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

327 return HeadSVD( 

328 which=which, 

329 layer=layer, 

330 head=head, 

331 U=U, 

332 S=S, 

333 V=V, 

334 rank_report=rank_report, 

335 eps=eps, 

336 null_rtol=resolved_null_rtol, 

337 ) 

338 

339 

340def decompose_head( 

341 model, 

342 layer: int, 

343 head: int, 

344 *, 

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

346 eps: float = _DEFAULT_EPS, 

347 null_rtol: Optional[float] = None, 

348) -> HeadDecomposition: 

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

350 

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

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

353 compatibility mode. The returned factors are detached from the model. 

354 

355 Args: 

356 model: A ``TransformerBridge``. 

357 layer: Layer of the head to decompose. 

358 head: Head index within the layer. 

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

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

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

362 closer than this share a block. 

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

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

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

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

367 

368 Returns: 

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

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

371 

372 Raises: 

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

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

375 """ 

376 n_layers = model.cfg.n_layers 

377 n_heads = model.cfg.n_heads 

378 if not 0 <= layer < n_layers: 

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

380 if not 0 <= head < n_heads: 

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

382 if isinstance(which, str): 

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

384 requested = tuple(which) 

385 if not requested: 

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

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

388 if invalid: 

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

390 

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

392 qk = None 

393 ov = None 

394 if "QK" in requested: 

395 qk = _factored_head_svd( 

396 W_Q_h, 

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

398 which="QK", 

399 layer=layer, 

400 head=head, 

401 eps=eps, 

402 null_rtol=null_rtol, 

403 ) 

404 if "OV" in requested: 

405 ov = _factored_head_svd( 

406 W_V_h, W_O_h, which="OV", layer=layer, head=head, eps=eps, null_rtol=null_rtol 

407 ) 

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