Coverage for transformer_lens/tools/analysis/projection_kernel.py: 88%

1"""Projection Kernel utilities for comparing linear subspaces. 

2 

3The Projection Kernel (PK) between subspaces with orthonormal bases ``U`` and 

4``V`` is ``||U.T @ V||_F^2``. It is invariant to basis choices within either 

5subspace and equals the sum of squared principal-angle cosines. 

6""" 

7 

8from __future__ import annotations 

9 

10import math 

11from dataclasses import dataclass 

12from numbers import Real 

13from typing import Any, List, Literal, Optional, Sequence, Tuple, cast 

14 

15import torch 

16from jaxtyping import Bool, Float, Int 

17 

18from transformer_lens.tools.analysis.jacobian_lens_decomposition import ( 

19 _linalg_on_cpu_if_mps, 

20) 

21 

22AttentionRole = Literal["Q", "K", "V", "O"] 

23LayerOrder = Literal["forward", "all"] 

24HeadKind = Literal["query", "kv"] 

25 

26_PAIRWISE_TEMP_BYTES = 64 * 1024 * 1024 

27 

28 

29@dataclass(frozen=True) 

30class SubspaceBasis: 

31 """An explicitly ranked orthonormal basis extracted from a matrix. 

32 

33 Attributes: 

34 basis: Orthonormal column-space basis, ``[ambient_dim, rank]``. 

35 singular_values: All reduced-SVD singular values, in descending order. 

36 rank: Number of retained basis directions. 

37 measured_rank: Numerical rank before optional caller truncation. 

38 rtol: Effective relative rank tolerance. 

39 threshold: Absolute singular-value threshold used for rank measurement. 

40 input_shape: Shape of the matrix from which the basis was extracted. 

41 """ 

42 

43 basis: Float[torch.Tensor, "ambient rank"] 

44 singular_values: Float[torch.Tensor, "spectrum"] 

45 rank: int 

46 measured_rank: int 

47 rtol: float 

48 threshold: float 

49 input_shape: Tuple[int, int] 

50 

51 @property 

52 def ambient_dim(self) -> int: 

53 """Dimension of the space containing the subspace.""" 

54 return self.basis.shape[0] 

55 

56 

57@dataclass(frozen=True) 

58class ProjectionKernelResult: 

59 """Projection Kernel score and its principal-angle decomposition.""" 

60 

61 score: Float[torch.Tensor, ""] 

62 normalized: Float[torch.Tensor, ""] 

63 cosines: Float[torch.Tensor, "principal_angle"] 

64 angles: Float[torch.Tensor, "principal_angle"] 

65 rank_a: int 

66 rank_b: int 

67 ambient_dim: int 

68 

69 

70@dataclass(frozen=True) 

71class RandomSubspaceReference: 

72 """Analytic PK moments for independent random equal-rank subspaces.""" 

73 

74 ambient_dim: int 

75 rank: int 

76 mean: float 

77 variance: float 

78 

79 

80@dataclass(frozen=True) 

81class AttentionHeadRef: 

82 """Structured identity for one attention-head weight subspace.""" 

83 

84 layer: int 

85 head: int 

86 role: AttentionRole 

87 kind: HeadKind 

88 

89 @property 

90 def label(self) -> str: 

91 """Return the conventional TransformerLens layer/head label.""" 

92 return f"L{self.layer}H{self.head}" 

93 

94 

95@dataclass(frozen=True) 

96class HeadAffinityPair: 

97 """One ranked source-target head pair.""" 

98 

99 source: AttentionHeadRef 

100 target: AttentionHeadRef 

101 score: float 

102 normalized: float 

103 

104 

105@dataclass(frozen=True) 

106class HeadAffinityResult: 

107 """Projection Kernel affinities between two attention-head roles. 

108 

109 Score tensors have shape 

110 ``[source_layer, source_head, target_layer, target_head]``. Layer index 

111 tuples map tensor positions to original model block numbers. 

112 

113 ``source_ranks`` and ``target_ranks`` are measured numerical ranks for each 

114 head before optional truncation. Scalar ``source_rank`` and ``target_rank`` 

115 are the retained basis widths used for their respective roles. 

116 """ 

117 

118 scores: Float[torch.Tensor, "source_layer source_head target_layer target_head"] 

119 normalized: Float[torch.Tensor, "source_layer source_head target_layer target_head"] 

120 valid_mask: Bool[torch.Tensor, "source_layer source_head target_layer target_head"] 

121 source_role: AttentionRole 

122 target_role: AttentionRole 

123 source_layer_indices: Tuple[int, ...] 

124 target_layer_indices: Tuple[int, ...] 

125 source_head_kind: HeadKind 

126 target_head_kind: HeadKind 

127 source_ranks: Int[torch.Tensor, "source_layer source_head"] 

128 target_ranks: Int[torch.Tensor, "target_layer target_head"] 

129 source_rank: int 

130 target_rank: int 

131 rank: Optional[int] 

132 rtol: float 

133 

134 def top_pairs(self, k: int = 20, *, normalized: bool = False) -> List[HeadAffinityPair]: 

135 """Return the highest-scoring valid pairs with deterministic tie order.""" 

136 if isinstance(k, bool) or not isinstance(k, int) or k < 1: 

137 raise ValueError(f"k must be a positive integer, got {k!r}") 

138 

139 selected_scores = self.normalized if normalized else self.scores 

140 entries: List[HeadAffinityPair] = [] 

141 for source_layer, source_head, target_layer, target_head in ( 141 ↛ 167line 141 didn't jump to line 167 because the loop on line 141 didn't complete

142 torch.nonzero(self.valid_mask, as_tuple=False).cpu().tolist() 

143 ): 

144 source = AttentionHeadRef( 

145 layer=self.source_layer_indices[source_layer], 

146 head=source_head, 

147 role=self.source_role, 

148 kind=self.source_head_kind, 

149 ) 

150 target = AttentionHeadRef( 

151 layer=self.target_layer_indices[target_layer], 

152 head=target_head, 

153 role=self.target_role, 

154 kind=self.target_head_kind, 

155 ) 

156 entries.append( 

157 HeadAffinityPair( 

158 source=source, 

159 target=target, 

160 score=float(self.scores[source_layer, source_head, target_layer, target_head]), 

161 normalized=float( 

162 self.normalized[source_layer, source_head, target_layer, target_head] 

163 ), 

164 ) 

165 ) 

166 

167 def sort_key(pair: HeadAffinityPair) -> Tuple[float, int, int, int, int]: 

168 value = pair.normalized if normalized else pair.score 

169 return ( 

170 -value, 

171 pair.source.layer, 

172 pair.source.head, 

173 pair.target.layer, 

174 pair.target.head, 

175 ) 

176 

177 entries.sort(key=sort_key) 

178 return entries[:k] 

179 

180 

181def _compute_dtype(dtype: torch.dtype) -> torch.dtype: 

182 """Return a dtype supported by stable SVD on common PyTorch backends.""" 

183 return torch.float64 if dtype == torch.float64 else torch.float32 

184 

185 

186def _validate_rtol(rtol: Optional[float], shape: Tuple[int, int], dtype: torch.dtype) -> float: 

187 if rtol is None: 

188 return max(shape) * torch.finfo(_compute_dtype(dtype)).eps 

189 if isinstance(rtol, bool) or not isinstance(rtol, Real): 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true

190 raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}") 

191 value = float(rtol) 

192 if not math.isfinite(value) or value < 0: 

193 raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}") 

194 return value 

195 

196 

197def _validate_rank(rank: Optional[int], max_rank: int) -> Optional[int]: 

198 if rank is None: 

199 return None 

200 if isinstance(rank, bool) or not isinstance(rank, int): 

201 raise ValueError(f"rank must be an integer or None, got {rank!r}") 

202 if not 1 <= rank <= max_rank: 

203 raise ValueError(f"rank must be between 1 and min(matrix.shape)={max_rank}, got {rank}") 

204 return rank 

205 

206 

207def orthonormal_subspace( 

208 matrix: Float[torch.Tensor, "ambient width"], 

209 *, 

210 rank: Optional[int] = None, 

211 rtol: Optional[float] = None, 

212) -> SubspaceBasis: 

213 """Extract an explicitly ranked orthonormal column-space basis. 

214 

215 Low-precision inputs are promoted to float32 before the reduced SVD. With no 

216 explicit ``rtol``, numerical rank uses the compute-SVD error scale relative 

217 to the largest singular value, so the default does not depend on the storage 

218 dtype — e.g. a 4096-row half-precision input gets ``4096 * float32_eps`` 

219 (about ``4.9e-4``), tight enough to flag an exactly repeated column as 

220 rank-deficient. An explicit ``rank`` truncates the measured subspace but may 

221 not exceed its measured rank. 

222 

223 Args: 

224 matrix: Finite floating-point matrix with shape ``[ambient_dim, width]``. 

225 rank: Optional number of leading singular directions to retain. 

226 rtol: Optional non-negative relative singular-value threshold. 

227 

228 Returns: 

229 Basis, complete singular spectrum, and rank metadata. 

230 

231 Raises: 

232 ValueError: If the matrix or rank policy is invalid. 

233 """ 

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

235 raise ValueError(f"matrix must be a torch.Tensor, got {type(matrix).__name__}") 

236 if matrix.ndim != 2: 236 ↛ 237line 236 didn't jump to line 237 because the condition on line 236 was never true

237 raise ValueError(f"matrix must be two-dimensional, got shape {tuple(matrix.shape)}") 

238 if matrix.shape[0] == 0 or matrix.shape[1] == 0: 

239 raise ValueError(f"matrix dimensions must be non-empty, got shape {tuple(matrix.shape)}") 

240 if not torch.is_floating_point(matrix): 240 ↛ 241line 240 didn't jump to line 241 because the condition on line 240 was never true

241 raise ValueError(f"matrix must have a real floating-point dtype, got {matrix.dtype}") 

242 if not bool(torch.isfinite(matrix).all()): 

243 raise ValueError("matrix must contain only finite values") 

244 

245 input_shape = (matrix.shape[0], matrix.shape[1]) 

246 requested_rank = _validate_rank(rank, min(input_shape)) 

247 compute_dtype = _compute_dtype(matrix.dtype) 

248 effective_rtol = _validate_rtol(rtol, input_shape, matrix.dtype) 

249 work = matrix.to(dtype=compute_dtype) 

250 left, singular_values, _ = torch.linalg.svd(work, full_matrices=False) 

251 threshold = float(singular_values[0].item()) * effective_rtol 

252 measured_rank = int((singular_values > threshold).sum().item()) 

253 if measured_rank == 0: 

254 raise ValueError( 

255 "matrix numerical rank is zero " 

256 f"for shape {input_shape}, dtype {matrix.dtype}, and threshold {threshold:.6g}" 

257 ) 

258 if requested_rank is not None and requested_rank > measured_rank: 

259 raise ValueError( 

260 f"requested rank {requested_rank} exceeds measured rank {measured_rank} " 

261 f"at threshold {threshold:.6g}" 

262 ) 

263 

264 selected_rank = measured_rank if requested_rank is None else requested_rank 

265 return SubspaceBasis( 

266 basis=left[:, :selected_rank], 

267 singular_values=singular_values, 

268 rank=selected_rank, 

269 measured_rank=measured_rank, 

270 rtol=effective_rtol, 

271 threshold=threshold, 

272 input_shape=input_shape, 

273 ) 

274 

275 

276def _validate_subspace(value: SubspaceBasis, name: str) -> None: 

277 if not isinstance(value, SubspaceBasis): 277 ↛ 278line 277 didn't jump to line 278 because the condition on line 277 was never true

278 raise ValueError(f"{name} must be a SubspaceBasis, got {type(value).__name__}") 

279 basis = value.basis 

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

281 raise ValueError(f"{name}.basis must be a two-dimensional tensor") 

282 if not torch.is_floating_point(basis) or not bool(torch.isfinite(basis).all()): 

283 raise ValueError(f"{name}.basis must be finite and real floating-point") 

284 singular_values = value.singular_values 

285 if ( 

286 not isinstance(singular_values, torch.Tensor) 

287 or singular_values.ndim != 1 

288 or not torch.is_floating_point(singular_values) 

289 or not bool(torch.isfinite(singular_values).all()) 

290 ): 

291 raise ValueError(f"{name}.singular_values must be a finite floating-point vector") 

292 if singular_values.device != basis.device: 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true

293 raise ValueError(f"{name}.singular_values must be on the basis device") 

294 if ( 294 ↛ 302line 294 didn't jump to line 302 because the condition on line 294 was never true

295 not isinstance(value.input_shape, tuple) 

296 or len(value.input_shape) != 2 

297 or any( 

298 isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 1 

299 for dimension in value.input_shape 

300 ) 

301 ): 

302 raise ValueError(f"{name}.input_shape must contain two positive dimensions") 

303 if singular_values.shape[0] != min(value.input_shape): 303 ↛ 304line 303 didn't jump to line 304 because the condition on line 303 was never true

304 raise ValueError(f"{name}.singular_values length must equal min(input_shape)") 

305 if ( 305 ↛ 311line 305 didn't jump to line 311 because the condition on line 305 was never true

306 isinstance(value.rank, bool) 

307 or not isinstance(value.rank, int) 

308 or value.rank < 1 

309 or basis.shape[1] != value.rank 

310 ): 

311 raise ValueError(f"{name}.rank must equal its positive basis width") 

312 if ( 312 ↛ 318line 312 didn't jump to line 318 because the condition on line 312 was never true

313 isinstance(value.measured_rank, bool) 

314 or not isinstance(value.measured_rank, int) 

315 or value.measured_rank < value.rank 

316 or value.measured_rank > min(value.input_shape) 

317 ): 

318 raise ValueError(f"{name}.measured_rank must be at least its selected rank") 

319 if value.input_shape[0] != basis.shape[0]: 319 ↛ 320line 319 didn't jump to line 320 because the condition on line 319 was never true

320 raise ValueError(f"{name}.input_shape must share its basis ambient dimension") 

321 if ( 321 ↛ 327line 321 didn't jump to line 327 because the condition on line 321 was never true

322 isinstance(value.rtol, bool) 

323 or not isinstance(value.rtol, Real) 

324 or not math.isfinite(value.rtol) 

325 or value.rtol < 0 

326 ): 

327 raise ValueError(f"{name}.rtol must be finite and non-negative") 

328 if ( 328 ↛ 334line 328 didn't jump to line 334 because the condition on line 328 was never true

329 isinstance(value.threshold, bool) 

330 or not isinstance(value.threshold, Real) 

331 or not math.isfinite(value.threshold) 

332 or value.threshold < 0 

333 ): 

334 raise ValueError(f"{name}.threshold must be finite and non-negative") 

335 

336 

337def _clamp_projection_scores(scores: torch.Tensor, upper_bound: int) -> torch.Tensor: 

338 """Clamp roundoff-scale PK bound violations and reject larger violations.""" 

339 tolerance = 100 * torch.finfo(scores.dtype).eps * max(1, upper_bound) 

340 minimum = float(scores.detach().min().item()) 

341 maximum = float(scores.detach().max().item()) 

342 if minimum < -tolerance or maximum > upper_bound + tolerance: 

343 raise ValueError( 

344 "Projection Kernel score lies outside its theoretical bounds: " 

345 f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, {upper_bound}] " 

346 f"within tolerance {tolerance:.6g}" 

347 ) 

348 return scores.clamp(min=0.0, max=float(upper_bound)) 

349 

350 

351def _clamp_principal_cosines(cosines: torch.Tensor) -> torch.Tensor: 

352 """Clamp roundoff-scale cosine violations and reject larger violations.""" 

353 tolerance = 100 * torch.finfo(cosines.dtype).eps * max(1, cosines.numel()) 

354 minimum = float(cosines.detach().min().item()) 

355 maximum = float(cosines.detach().max().item()) 

356 if minimum < -tolerance or maximum > 1 + tolerance: 

357 raise ValueError( 

358 "Principal-angle cosine lies outside its theoretical bounds: " 

359 f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, 1] " 

360 f"within tolerance {tolerance:.6g}" 

361 ) 

362 return cosines.clamp(min=0.0, max=1.0) 

363 

364 

365def _singular_values(matrix: torch.Tensor) -> torch.Tensor: 

366 """Compute singular values with an explicit MPS CPU fallback.""" 

367 return _linalg_on_cpu_if_mps(torch.linalg.svdvals, matrix) 

368 

369 

370def projection_kernel( 

371 subspace_a: SubspaceBasis, 

372 subspace_b: SubspaceBasis, 

373 *, 

374 check_orthonormal: bool = True, 

375) -> ProjectionKernelResult: 

376 """Measure overlap between two explicitly extracted subspaces. 

377 

378 Raw PK lies in ``[0, min(rank_a, rank_b)]``. The normalized value is 

379 ``PK / sqrt(rank_a * rank_b)``, the cosine between the two projection 

380 matrices. Principal angles are returned in radians. 

381 """ 

382 if not isinstance(check_orthonormal, bool): 382 ↛ 383line 382 didn't jump to line 383 because the condition on line 382 was never true

383 raise ValueError(f"check_orthonormal must be a Boolean, got {check_orthonormal!r}") 

384 _validate_subspace(subspace_a, "subspace_a") 

385 _validate_subspace(subspace_b, "subspace_b") 

386 if subspace_a.ambient_dim != subspace_b.ambient_dim: 

387 raise ValueError( 

388 "subspaces must have equal ambient dimensions, got " 

389 f"{subspace_a.ambient_dim} and {subspace_b.ambient_dim}" 

390 ) 

391 if subspace_a.basis.device != subspace_b.basis.device: 391 ↛ 392line 391 didn't jump to line 392 because the condition on line 391 was never true

392 raise ValueError( 

393 "subspace bases must be on the same device, got " 

394 f"{subspace_a.basis.device} and {subspace_b.basis.device}" 

395 ) 

396 

397 dtype = torch.promote_types(subspace_a.basis.dtype, subspace_b.basis.dtype) 

398 dtype = _compute_dtype(dtype) 

399 first = subspace_a.basis.to(dtype=dtype) 

400 second = subspace_b.basis.to(dtype=dtype) 

401 if check_orthonormal: 

402 eps = torch.finfo(dtype).eps 

403 tolerance = 10 * max(first.shape[0], first.shape[1], second.shape[1]) * eps 

404 first_identity = torch.eye(first.shape[1], dtype=dtype, device=first.device) 

405 second_identity = torch.eye(second.shape[1], dtype=dtype, device=second.device) 

406 if not torch.allclose(first.T @ first, first_identity, rtol=tolerance, atol=tolerance): 

407 raise ValueError("subspace_a basis columns must be orthonormal") 

408 if not torch.allclose(second.T @ second, second_identity, rtol=tolerance, atol=tolerance): 408 ↛ 409line 408 didn't jump to line 409 because the condition on line 408 was never true

409 raise ValueError("subspace_b basis columns must be orthonormal") 

410 

411 overlap = first.T @ second 

412 score = _clamp_projection_scores(overlap.square().sum(), min(subspace_a.rank, subspace_b.rank)) 

413 cosines = _clamp_principal_cosines(_singular_values(overlap)) 

414 angles = torch.acos(cosines) 

415 denominator = math.sqrt(subspace_a.rank * subspace_b.rank) 

416 return ProjectionKernelResult( 

417 score=score, 

418 normalized=score / denominator, 

419 cosines=cosines, 

420 angles=angles, 

421 rank_a=subspace_a.rank, 

422 rank_b=subspace_b.rank, 

423 ambient_dim=subspace_a.ambient_dim, 

424 ) 

425 

426 

427def random_projection_kernel_moments(ambient_dim: int, rank: int) -> RandomSubspaceReference: 

428 """Return PK moments for independent Haar-distributed rank-``rank`` planes. 

429 

430 These idealized descriptive moments are not calibrated p-values for trained 

431 model weights, whose head subspaces are dependent and anisotropic. 

432 """ 

433 if isinstance(ambient_dim, bool) or not isinstance(ambient_dim, int) or ambient_dim < 2: 

434 raise ValueError(f"ambient_dim must be an integer at least 2, got {ambient_dim!r}") 

435 if isinstance(rank, bool) or not isinstance(rank, int) or not 1 <= rank <= ambient_dim: 

436 raise ValueError(f"rank must be an integer between 1 and {ambient_dim}, got {rank!r}") 

437 

438 mean = rank**2 / ambient_dim 

439 variance = ( 

440 2 

441 * rank**2 

442 * (ambient_dim - rank) ** 2 

443 / (ambient_dim**2 * (ambient_dim - 1) * (ambient_dim + 2)) 

444 ) 

445 return RandomSubspaceReference( 

446 ambient_dim=ambient_dim, 

447 rank=rank, 

448 mean=float(mean), 

449 variance=float(variance), 

450 ) 

451 

452 

453def _pairwise_projection_kernel( 

454 source_bases: torch.Tensor, 

455 target_bases: torch.Tensor, 

456 *, 

457 max_temp_bytes: int = _PAIRWISE_TEMP_BYTES, 

458) -> torch.Tensor: 

459 """Compute pairwise PK scores while bounding the overlap-tensor allocation.""" 

460 if source_bases.ndim != 3 or target_bases.ndim != 3: 460 ↛ 461line 460 didn't jump to line 461 because the condition on line 460 was never true

461 raise ValueError("source_bases and target_bases must be three-dimensional") 

462 if source_bases.shape[1] != target_bases.shape[1]: 462 ↛ 463line 462 didn't jump to line 463 because the condition on line 462 was never true

463 raise ValueError("source_bases and target_bases must share an ambient dimension") 

464 if source_bases.device != target_bases.device or source_bases.dtype != target_bases.dtype: 464 ↛ 465line 464 didn't jump to line 465 because the condition on line 464 was never true

465 raise ValueError("source_bases and target_bases must share a device and dtype") 

466 if ( 466 ↛ 471line 466 didn't jump to line 471 because the condition on line 466 was never true

467 isinstance(max_temp_bytes, bool) 

468 or not isinstance(max_temp_bytes, int) 

469 or max_temp_bytes < 1 

470 ): 

471 raise ValueError("max_temp_bytes must be a positive integer") 

472 

473 source_count, _, source_rank = source_bases.shape 

474 target_count, _, target_rank = target_bases.shape 

475 scores = torch.empty( 

476 source_count, target_count, dtype=source_bases.dtype, device=source_bases.device 

477 ) 

478 bytes_per_overlap = source_rank * target_rank * source_bases.element_size() 

479 pair_capacity = max(1, max_temp_bytes // bytes_per_overlap) 

480 source_tile = max(1, min(source_count, pair_capacity // max(1, target_count))) 

481 target_tile = max(1, min(target_count, pair_capacity // source_tile)) 

482 

483 for source_start in range(0, source_count, source_tile): 

484 source_stop = min(source_start + source_tile, source_count) 

485 for target_start in range(0, target_count, target_tile): 

486 target_stop = min(target_start + target_tile, target_count) 

487 overlap = torch.einsum( 

488 "adr,bds->abrs", 

489 source_bases[source_start:source_stop], 

490 target_bases[target_start:target_stop], 

491 ) 

492 scores[source_start:source_stop, target_start:target_stop] = overlap.square().sum( 

493 dim=(-2, -1) 

494 ) 

495 return scores 

496 

497 

498def _architecture_name(model: Any) -> str: 

499 cfg = getattr(model, "cfg", None) 

500 architecture = getattr(cfg, "original_architecture", None) 

501 return str(architecture) if architecture is not None else type(model).__name__ 

502 

503 

504def _read_role_matrix(model: Any, block: Any, layer: int, role: AttentionRole) -> torch.Tensor: 

505 attn = getattr(block, "attn", None) 

506 if attn is None: 506 ↛ 507line 506 didn't jump to line 507 because the condition on line 506 was never true

507 raise ValueError(f"attention block {layer} does not expose an attn component") 

508 attribute = f"W_{role}" 

509 try: 

510 matrix = getattr(attn, attribute) 

511 except NotImplementedError as error: 

512 raise NotImplementedError( 

513 f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}" 

514 ) from error 

515 except (AttributeError, RuntimeError, ValueError) as error: 

516 raise ValueError( 

517 f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}" 

518 ) from error 

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

520 raise ValueError( 

521 f"role {role} at layer {layer} must be a tensor, got {type(matrix).__name__}" 

522 ) 

523 if matrix.ndim != 3: 

524 raise ValueError( 

525 f"role {role} at layer {layer} expected a three-dimensional per-head weight, " 

526 f"got shape {tuple(matrix.shape)}" 

527 ) 

528 if role == "O": 

529 matrix = matrix.transpose(-2, -1) 

530 if not torch.is_floating_point(matrix): 530 ↛ 531line 530 didn't jump to line 531 because the condition on line 530 was never true

531 raise ValueError(f"role {role} at layer {layer} must have a floating-point dtype") 

532 if not bool(torch.isfinite(matrix).all()): 

533 raise ValueError(f"role {role} at layer {layer} must contain only finite values") 

534 return matrix.detach() 

535 

536 

537def _validate_role_shapes( 

538 matrices: Sequence[torch.Tensor], layers: Sequence[int], role: AttentionRole 

539) -> Tuple[int, int, int]: 

540 expected = (matrices[0].shape[0], matrices[0].shape[1], matrices[0].shape[2]) 

541 if min(expected) < 1: 

542 raise ValueError( 

543 f"role {role} at layer {layers[0]} must have non-empty head, d_model, and width " 

544 f"dimensions, got shape {expected}" 

545 ) 

546 for matrix, layer in zip(matrices[1:], layers[1:]): 

547 if tuple(matrix.shape) != expected: 

548 raise ValueError( 

549 f"role {role} at layer {layer} has shape {tuple(matrix.shape)}, expected " 

550 f"the consistent [heads, d_model, width] shape {expected}" 

551 ) 

552 return expected 

553 

554 

555def _extract_bases( 

556 matrices: Sequence[torch.Tensor], 

557 layers: Sequence[int], 

558 role: AttentionRole, 

559 *, 

560 selected_rank: int, 

561 rtol: float, 

562 dtype: torch.dtype, 

563 device: torch.device, 

564) -> Tuple[torch.Tensor, torch.Tensor]: 

565 layer_bases: List[torch.Tensor] = [] 

566 layer_ranks: List[List[int]] = [] 

567 for matrix, layer in zip(matrices, layers): 

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

569 head_ranks: List[int] = [] 

570 for head in range(matrix.shape[0]): 

571 try: 

572 subspace = orthonormal_subspace( 

573 matrix[head].to(dtype=dtype), rank=selected_rank, rtol=rtol 

574 ) 

575 except ValueError as error: 

576 raise ValueError( 

577 f"Could not extract role {role} at layer {layer}, head {head}: {error}" 

578 ) from error 

579 head_bases.append(subspace.basis.to(device=device)) 

580 head_ranks.append(subspace.measured_rank) 

581 layer_bases.append(torch.stack(head_bases)) 

582 layer_ranks.append(head_ranks) 

583 return torch.stack(layer_bases), torch.tensor(layer_ranks, dtype=torch.long, device=device) 

584 

585 

586def attention_head_subspace_affinity( 

587 model: Any, 

588 *, 

589 source_role: str = "O", 

590 target_role: str, 

591 layer_order: str = "forward", 

592 rank: Optional[int] = None, 

593 rtol: Optional[float] = None, 

594) -> HeadAffinityResult: 

595 """Compute OQ, OK, or OV Projection Kernel affinities for a TransformerBridge. 

596 

597 K/V axes preserve native key-value heads on grouped-query attention models; 

598 they are never expanded to query-head count. Hybrid models include only 

599 attention blocks and report their original block indices. 

600 

601 Args: 

602 model: A TransformerBridge exposing readable per-head attention weights. 

603 source_role: Source role; v1 supports only ``"O"``. 

604 target_role: One of ``"Q"``, ``"K"``, or ``"V"``. 

605 layer_order: ``"forward"`` keeps strict earlier-to-later pairs; ``"all"`` 

606 keeps every pair. 

607 rank: Optional common truncation rank. By default every head must be full 

608 column rank. 

609 rtol: Optional relative numerical-rank tolerance. 

610 

611 Returns: 

612 Affinity tensors, validity mask, original layer indices, and rank metadata. 

613 """ 

614 if source_role != "O": 

615 raise ValueError(f"source_role must be 'O' in v1, got {source_role!r}") 

616 if target_role not in ("Q", "K", "V"): 

617 raise ValueError(f"target_role must be one of ['Q', 'K', 'V'], got {target_role!r}") 

618 if layer_order not in ("forward", "all"): 

619 raise ValueError(f"layer_order must be one of ['forward', 'all'], got {layer_order!r}") 

620 validated_source_role = cast(AttentionRole, source_role) 

621 validated_target_role = cast(AttentionRole, target_role) 

622 validated_layer_order = cast(LayerOrder, layer_order) 

623 blocks_with = getattr(model, "blocks_with", None) 

624 if not callable(blocks_with): 624 ↛ 625line 624 didn't jump to line 625 because the condition on line 624 was never true

625 raise ValueError("model must be a TransformerBridge exposing blocks_with('attn')") 

626 attention_blocks = list(blocks_with("attn")) 

627 if not attention_blocks: 

628 raise ValueError("No attention layers found — cannot compute head subspace affinity.") 

629 

630 layer_indices = [int(layer) for layer, _ in attention_blocks] 

631 source_matrices = [ 

632 _read_role_matrix(model, block, layer, validated_source_role) 

633 for layer, block in attention_blocks 

634 ] 

635 target_matrices = [ 

636 _read_role_matrix(model, block, layer, validated_target_role) 

637 for layer, block in attention_blocks 

638 ] 

639 source_heads, source_ambient, source_width = _validate_role_shapes( 

640 source_matrices, layer_indices, validated_source_role 

641 ) 

642 target_heads, target_ambient, target_width = _validate_role_shapes( 

643 target_matrices, layer_indices, validated_target_role 

644 ) 

645 if source_ambient != target_ambient: 645 ↛ 646line 645 didn't jump to line 646 because the condition on line 645 was never true

646 raise ValueError( 

647 f"roles {validated_source_role} and {validated_target_role} must share d_model, got " 

648 f"{source_ambient} and {target_ambient}" 

649 ) 

650 

651 matrices = source_matrices + target_matrices 

652 dtype = matrices[0].dtype 

653 for matrix in matrices[1:]: 

654 dtype = torch.promote_types(dtype, matrix.dtype) 

655 dtype = _compute_dtype(dtype) 

656 effective_rtol = _validate_rtol(rtol, (source_ambient, max(source_width, target_width)), dtype) 

657 requested_rank = _validate_rank(rank, min(source_ambient, source_width, target_width)) 

658 source_rank = source_width if requested_rank is None else requested_rank 

659 target_rank = target_width if requested_rank is None else requested_rank 

660 

661 cfg = getattr(model, "cfg", None) 

662 configured_device = getattr(cfg, "device", None) 

663 result_device = ( 

664 torch.device(configured_device) 

665 if configured_device is not None 

666 else source_matrices[0].device 

667 ) 

668 source_bases, source_ranks = _extract_bases( 

669 source_matrices, 

670 layer_indices, 

671 validated_source_role, 

672 selected_rank=source_rank, 

673 rtol=effective_rtol, 

674 dtype=dtype, 

675 device=result_device, 

676 ) 

677 target_bases, target_ranks = _extract_bases( 

678 target_matrices, 

679 layer_indices, 

680 validated_target_role, 

681 selected_rank=target_rank, 

682 rtol=effective_rtol, 

683 dtype=dtype, 

684 device=result_device, 

685 ) 

686 

687 layer_count = len(layer_indices) 

688 flat_source = source_bases.reshape(layer_count * source_heads, source_ambient, source_rank) 

689 flat_target = target_bases.reshape(layer_count * target_heads, target_ambient, target_rank) 

690 scores = _pairwise_projection_kernel(flat_source, flat_target).reshape( 

691 layer_count, source_heads, layer_count, target_heads 

692 ) 

693 scores = _clamp_projection_scores(scores, min(source_rank, target_rank)) 

694 normalized = scores / math.sqrt(source_rank * target_rank) 

695 if validated_layer_order == "forward": 

696 layer_tensor = torch.tensor(layer_indices, device=result_device) 

697 layer_mask = layer_tensor[:, None] < layer_tensor[None, :] 

698 valid_mask = layer_mask[:, None, :, None].expand_as(scores).contiguous() 

699 else: 

700 valid_mask = torch.ones_like(scores, dtype=torch.bool) 

701 scores = torch.where(valid_mask, scores, torch.zeros_like(scores)) 

702 normalized = torch.where(valid_mask, normalized, torch.zeros_like(normalized)) 

703 

704 target_kind: HeadKind = "query" if validated_target_role == "Q" else "kv"