Coverage for transformer_lens/tools/analysis/projection_kernel.py: 89%
324 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Projection Kernel utilities for comparing linear subspaces.
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"""
8from __future__ import annotations
10import math
11from dataclasses import dataclass
12from numbers import Real
13from typing import Any, List, Literal, Optional, Sequence, Tuple, cast
15import torch
16from jaxtyping import Bool, Float, Int
18from transformer_lens.tools.analysis._model_state import compute_dtype, require_int
19from transformer_lens.tools.analysis.jacobian_lens_decomposition import (
20 _linalg_on_cpu_if_mps,
21)
23AttentionRole = Literal["Q", "K", "V", "O"]
24LayerOrder = Literal["forward", "all"]
25HeadKind = Literal["query", "kv"]
27_PAIRWISE_TEMP_BYTES = 64 * 1024 * 1024
30@dataclass(frozen=True)
31class SubspaceBasis:
32 """An explicitly ranked orthonormal basis extracted from a matrix.
34 Attributes:
35 basis: Orthonormal column-space basis, ``[ambient_dim, rank]``.
36 singular_values: All reduced-SVD singular values, in descending order.
37 rank: Number of retained basis directions.
38 measured_rank: Numerical rank before optional caller truncation.
39 rtol: Effective relative rank tolerance.
40 threshold: Absolute singular-value threshold used for rank measurement.
41 input_shape: Shape of the matrix from which the basis was extracted.
42 """
44 basis: Float[torch.Tensor, "ambient rank"]
45 singular_values: Float[torch.Tensor, "spectrum"]
46 rank: int
47 measured_rank: int
48 rtol: float
49 threshold: float
50 input_shape: Tuple[int, int]
52 @property
53 def ambient_dim(self) -> int:
54 """Dimension of the space containing the subspace."""
55 return self.basis.shape[0]
58@dataclass(frozen=True)
59class ProjectionKernelResult:
60 """Projection Kernel score and its principal-angle decomposition."""
62 score: Float[torch.Tensor, ""]
63 normalized: Float[torch.Tensor, ""]
64 cosines: Float[torch.Tensor, "principal_angle"]
65 angles: Float[torch.Tensor, "principal_angle"]
66 rank_a: int
67 rank_b: int
68 ambient_dim: int
71@dataclass(frozen=True)
72class RandomSubspaceReference:
73 """Analytic PK moments for independent random equal-rank subspaces."""
75 ambient_dim: int
76 rank: int
77 mean: float
78 variance: float
81@dataclass(frozen=True)
82class AttentionHeadRef:
83 """Structured identity for one attention-head weight subspace."""
85 layer: int
86 head: int
87 role: AttentionRole
88 kind: HeadKind
90 @property
91 def label(self) -> str:
92 """Return the conventional TransformerLens layer/head label."""
93 return f"L{self.layer}H{self.head}"
96@dataclass(frozen=True)
97class HeadAffinityPair:
98 """One ranked source-target head pair."""
100 source: AttentionHeadRef
101 target: AttentionHeadRef
102 score: float
103 normalized: float
106@dataclass(frozen=True)
107class HeadAffinityResult:
108 """Projection Kernel affinities between two attention-head roles.
110 Score tensors have shape
111 ``[source_layer, source_head, target_layer, target_head]``. Layer index
112 tuples map tensor positions to original model block numbers.
114 ``source_ranks`` and ``target_ranks`` are measured numerical ranks for each
115 head before optional truncation. Scalar ``source_rank`` and ``target_rank``
116 are the retained basis widths used for their respective roles.
117 """
119 scores: Float[torch.Tensor, "source_layer source_head target_layer target_head"]
120 normalized: Float[torch.Tensor, "source_layer source_head target_layer target_head"]
121 valid_mask: Bool[torch.Tensor, "source_layer source_head target_layer target_head"]
122 source_role: AttentionRole
123 target_role: AttentionRole
124 source_layer_indices: Tuple[int, ...]
125 target_layer_indices: Tuple[int, ...]
126 source_head_kind: HeadKind
127 target_head_kind: HeadKind
128 source_ranks: Int[torch.Tensor, "source_layer source_head"]
129 target_ranks: Int[torch.Tensor, "target_layer target_head"]
130 source_rank: int
131 target_rank: int
132 rank: Optional[int]
133 rtol: float
135 def top_pairs(self, k: int = 20, *, normalized: bool = False) -> List[HeadAffinityPair]:
136 """Return the highest-scoring valid pairs with deterministic tie order."""
137 k = require_int(k, "k")
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 )
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 )
177 entries.sort(key=sort_key)
178 return entries[:k]
181def _validate_rtol(rtol: Optional[float], shape: Tuple[int, int], dtype: torch.dtype) -> float:
182 if rtol is None:
183 return max(shape) * torch.finfo(compute_dtype(dtype)).eps
184 if isinstance(rtol, bool) or not isinstance(rtol, Real): 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true
185 raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}")
186 value = float(rtol)
187 if not math.isfinite(value) or value < 0:
188 raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}")
189 return value
192def _validate_rank(rank: Optional[int], max_rank: int) -> Optional[int]:
193 if rank is None:
194 return None
195 return require_int(rank, "rank", maximum=max_rank)
198def orthonormal_subspace(
199 matrix: Float[torch.Tensor, "ambient width"],
200 *,
201 rank: Optional[int] = None,
202 rtol: Optional[float] = None,
203) -> SubspaceBasis:
204 """Extract an explicitly ranked orthonormal column-space basis.
206 Low-precision inputs are promoted to float32 before the reduced SVD. With no
207 explicit ``rtol``, numerical rank uses the compute-SVD error scale relative
208 to the largest singular value, so the default does not depend on the storage
209 dtype — e.g. a 4096-row half-precision input gets ``4096 * float32_eps``
210 (about ``4.9e-4``), tight enough to flag an exactly repeated column as
211 rank-deficient. An explicit ``rank`` truncates the measured subspace but may
212 not exceed its measured rank.
214 Args:
215 matrix: Finite floating-point matrix with shape ``[ambient_dim, width]``.
216 rank: Optional number of leading singular directions to retain.
217 rtol: Optional non-negative relative singular-value threshold.
219 Returns:
220 Basis, complete singular spectrum, and rank metadata.
222 Raises:
223 ValueError: If the matrix or rank policy is invalid.
224 """
225 if not isinstance(matrix, torch.Tensor): 225 ↛ 226line 225 didn't jump to line 226 because the condition on line 225 was never true
226 raise ValueError(f"matrix must be a torch.Tensor, got {type(matrix).__name__}")
227 if matrix.ndim != 2: 227 ↛ 228line 227 didn't jump to line 228 because the condition on line 227 was never true
228 raise ValueError(f"matrix must be two-dimensional, got shape {tuple(matrix.shape)}")
229 if matrix.shape[0] == 0 or matrix.shape[1] == 0:
230 raise ValueError(f"matrix dimensions must be non-empty, got shape {tuple(matrix.shape)}")
231 if not torch.is_floating_point(matrix): 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true
232 raise ValueError(f"matrix must have a real floating-point dtype, got {matrix.dtype}")
233 if not bool(torch.isfinite(matrix).all()):
234 raise ValueError("matrix must contain only finite values")
236 input_shape = (matrix.shape[0], matrix.shape[1])
237 requested_rank = _validate_rank(rank, min(input_shape))
238 effective_rtol = _validate_rtol(rtol, input_shape, matrix.dtype)
239 work = matrix.to(dtype=compute_dtype(matrix.dtype))
240 left, singular_values, _ = torch.linalg.svd(work, full_matrices=False)
241 threshold = float(singular_values[0].item()) * effective_rtol
242 measured_rank = int((singular_values > threshold).sum().item())
243 if measured_rank == 0:
244 raise ValueError(
245 "matrix numerical rank is zero "
246 f"for shape {input_shape}, dtype {matrix.dtype}, and threshold {threshold:.6g}"
247 )
248 if requested_rank is not None and requested_rank > measured_rank:
249 raise ValueError(
250 f"requested rank {requested_rank} exceeds measured rank {measured_rank} "
251 f"at threshold {threshold:.6g}"
252 )
254 selected_rank = measured_rank if requested_rank is None else requested_rank
255 return SubspaceBasis(
256 basis=left[:, :selected_rank],
257 singular_values=singular_values,
258 rank=selected_rank,
259 measured_rank=measured_rank,
260 rtol=effective_rtol,
261 threshold=threshold,
262 input_shape=input_shape,
263 )
266def _validate_subspace(value: SubspaceBasis, name: str) -> None:
267 if not isinstance(value, SubspaceBasis): 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true
268 raise ValueError(f"{name} must be a SubspaceBasis, got {type(value).__name__}")
269 basis = value.basis
270 if not isinstance(basis, torch.Tensor) or basis.ndim != 2: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 raise ValueError(f"{name}.basis must be a two-dimensional tensor")
272 if not torch.is_floating_point(basis) or not bool(torch.isfinite(basis).all()):
273 raise ValueError(f"{name}.basis must be finite and real floating-point")
274 singular_values = value.singular_values
275 if (
276 not isinstance(singular_values, torch.Tensor)
277 or singular_values.ndim != 1
278 or not torch.is_floating_point(singular_values)
279 or not bool(torch.isfinite(singular_values).all())
280 ):
281 raise ValueError(f"{name}.singular_values must be a finite floating-point vector")
282 if singular_values.device != basis.device: 282 ↛ 283line 282 didn't jump to line 283 because the condition on line 282 was never true
283 raise ValueError(f"{name}.singular_values must be on the basis device")
284 if not isinstance(value.input_shape, tuple) or len(value.input_shape) != 2: 284 ↛ 285line 284 didn't jump to line 285 because the condition on line 284 was never true
285 raise ValueError(f"{name}.input_shape must contain two positive dimensions")
286 for dimension in value.input_shape:
287 require_int(dimension, f"{name}.input_shape dimension")
288 if singular_values.shape[0] != min(value.input_shape): 288 ↛ 289line 288 didn't jump to line 289 because the condition on line 288 was never true
289 raise ValueError(f"{name}.singular_values length must equal min(input_shape)")
290 if require_int(value.rank, f"{name}.rank") != basis.shape[1]: 290 ↛ 291line 290 didn't jump to line 291 because the condition on line 290 was never true
291 raise ValueError(f"{name}.rank must equal its positive basis width")
292 require_int(
293 value.measured_rank,
294 f"{name}.measured_rank",
295 minimum=value.rank,
296 maximum=min(value.input_shape),
297 )
298 if value.input_shape[0] != basis.shape[0]: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 raise ValueError(f"{name}.input_shape must share its basis ambient dimension")
300 if ( 300 ↛ 306line 300 didn't jump to line 306 because the condition on line 300 was never true
301 isinstance(value.rtol, bool)
302 or not isinstance(value.rtol, Real)
303 or not math.isfinite(value.rtol)
304 or value.rtol < 0
305 ):
306 raise ValueError(f"{name}.rtol must be finite and non-negative")
307 if ( 307 ↛ 313line 307 didn't jump to line 313 because the condition on line 307 was never true
308 isinstance(value.threshold, bool)
309 or not isinstance(value.threshold, Real)
310 or not math.isfinite(value.threshold)
311 or value.threshold < 0
312 ):
313 raise ValueError(f"{name}.threshold must be finite and non-negative")
316def _clamp_projection_scores(scores: torch.Tensor, upper_bound: int) -> torch.Tensor:
317 """Clamp roundoff-scale PK bound violations and reject larger violations."""
318 tolerance = 100 * torch.finfo(scores.dtype).eps * max(1, upper_bound)
319 minimum = float(scores.detach().min().item())
320 maximum = float(scores.detach().max().item())
321 if minimum < -tolerance or maximum > upper_bound + tolerance:
322 raise ValueError(
323 "Projection Kernel score lies outside its theoretical bounds: "
324 f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, {upper_bound}] "
325 f"within tolerance {tolerance:.6g}"
326 )
327 return scores.clamp(min=0.0, max=float(upper_bound))
330def _clamp_principal_cosines(cosines: torch.Tensor) -> torch.Tensor:
331 """Clamp roundoff-scale cosine violations and reject larger violations."""
332 tolerance = 100 * torch.finfo(cosines.dtype).eps * max(1, cosines.numel())
333 minimum = float(cosines.detach().min().item())
334 maximum = float(cosines.detach().max().item())
335 if minimum < -tolerance or maximum > 1 + tolerance:
336 raise ValueError(
337 "Principal-angle cosine lies outside its theoretical bounds: "
338 f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, 1] "
339 f"within tolerance {tolerance:.6g}"
340 )
341 return cosines.clamp(min=0.0, max=1.0)
344def _singular_values(matrix: torch.Tensor) -> torch.Tensor:
345 """Compute singular values with an explicit MPS CPU fallback."""
346 return _linalg_on_cpu_if_mps(torch.linalg.svdvals, matrix)
349def projection_kernel(
350 subspace_a: SubspaceBasis,
351 subspace_b: SubspaceBasis,
352 *,
353 check_orthonormal: bool = True,
354) -> ProjectionKernelResult:
355 """Measure overlap between two explicitly extracted subspaces.
357 Raw PK lies in ``[0, min(rank_a, rank_b)]``. The normalized value is
358 ``PK / sqrt(rank_a * rank_b)``, the cosine between the two projection
359 matrices. Principal angles are returned in radians.
360 """
361 if not isinstance(check_orthonormal, bool): 361 ↛ 362line 361 didn't jump to line 362 because the condition on line 361 was never true
362 raise ValueError(f"check_orthonormal must be a Boolean, got {check_orthonormal!r}")
363 _validate_subspace(subspace_a, "subspace_a")
364 _validate_subspace(subspace_b, "subspace_b")
365 if subspace_a.ambient_dim != subspace_b.ambient_dim:
366 raise ValueError(
367 "subspaces must have equal ambient dimensions, got "
368 f"{subspace_a.ambient_dim} and {subspace_b.ambient_dim}"
369 )
370 if subspace_a.basis.device != subspace_b.basis.device: 370 ↛ 371line 370 didn't jump to line 371 because the condition on line 370 was never true
371 raise ValueError(
372 "subspace bases must be on the same device, got "
373 f"{subspace_a.basis.device} and {subspace_b.basis.device}"
374 )
376 dtype = torch.promote_types(subspace_a.basis.dtype, subspace_b.basis.dtype)
377 dtype = compute_dtype(dtype)
378 first = subspace_a.basis.to(dtype=dtype)
379 second = subspace_b.basis.to(dtype=dtype)
380 if check_orthonormal:
381 eps = torch.finfo(dtype).eps
382 tolerance = 10 * max(first.shape[0], first.shape[1], second.shape[1]) * eps
383 first_identity = torch.eye(first.shape[1], dtype=dtype, device=first.device)
384 second_identity = torch.eye(second.shape[1], dtype=dtype, device=second.device)
385 if not torch.allclose(first.T @ first, first_identity, rtol=tolerance, atol=tolerance):
386 raise ValueError("subspace_a basis columns must be orthonormal")
387 if not torch.allclose(second.T @ second, second_identity, rtol=tolerance, atol=tolerance): 387 ↛ 388line 387 didn't jump to line 388 because the condition on line 387 was never true
388 raise ValueError("subspace_b basis columns must be orthonormal")
390 overlap = first.T @ second
391 score = _clamp_projection_scores(overlap.square().sum(), min(subspace_a.rank, subspace_b.rank))
392 cosines = _clamp_principal_cosines(_singular_values(overlap))
393 angles = torch.acos(cosines)
394 denominator = math.sqrt(subspace_a.rank * subspace_b.rank)
395 return ProjectionKernelResult(
396 score=score,
397 normalized=score / denominator,
398 cosines=cosines,
399 angles=angles,
400 rank_a=subspace_a.rank,
401 rank_b=subspace_b.rank,
402 ambient_dim=subspace_a.ambient_dim,
403 )
406def random_projection_kernel_moments(ambient_dim: int, rank: int) -> RandomSubspaceReference:
407 """Return PK moments for independent Haar-distributed rank-``rank`` planes.
409 These idealized descriptive moments are not calibrated p-values for trained
410 model weights, whose head subspaces are dependent and anisotropic.
411 """
412 ambient_dim = require_int(ambient_dim, "ambient_dim", minimum=2)
413 rank = require_int(rank, "rank", maximum=ambient_dim)
415 mean = rank**2 / ambient_dim
416 variance = (
417 2
418 * rank**2
419 * (ambient_dim - rank) ** 2
420 / (ambient_dim**2 * (ambient_dim - 1) * (ambient_dim + 2))
421 )
422 return RandomSubspaceReference(
423 ambient_dim=ambient_dim,
424 rank=rank,
425 mean=float(mean),
426 variance=float(variance),
427 )
430def _pairwise_projection_kernel(
431 source_bases: torch.Tensor,
432 target_bases: torch.Tensor,
433 *,
434 max_temp_bytes: int = _PAIRWISE_TEMP_BYTES,
435) -> torch.Tensor:
436 """Compute pairwise PK scores while bounding the overlap-tensor allocation."""
437 if source_bases.ndim != 3 or target_bases.ndim != 3: 437 ↛ 438line 437 didn't jump to line 438 because the condition on line 437 was never true
438 raise ValueError("source_bases and target_bases must be three-dimensional")
439 if source_bases.shape[1] != target_bases.shape[1]: 439 ↛ 440line 439 didn't jump to line 440 because the condition on line 439 was never true
440 raise ValueError("source_bases and target_bases must share an ambient dimension")
441 if source_bases.device != target_bases.device or source_bases.dtype != target_bases.dtype: 441 ↛ 442line 441 didn't jump to line 442 because the condition on line 441 was never true
442 raise ValueError("source_bases and target_bases must share a device and dtype")
443 max_temp_bytes = require_int(max_temp_bytes, "max_temp_bytes")
445 source_count, _, source_rank = source_bases.shape
446 target_count, _, target_rank = target_bases.shape
447 scores = torch.empty(
448 source_count, target_count, dtype=source_bases.dtype, device=source_bases.device
449 )
450 bytes_per_overlap = source_rank * target_rank * source_bases.element_size()
451 pair_capacity = max(1, max_temp_bytes // bytes_per_overlap)
452 source_tile = max(1, min(source_count, pair_capacity // max(1, target_count)))
453 target_tile = max(1, min(target_count, pair_capacity // source_tile))
455 for source_start in range(0, source_count, source_tile):
456 source_stop = min(source_start + source_tile, source_count)
457 for target_start in range(0, target_count, target_tile):
458 target_stop = min(target_start + target_tile, target_count)
459 overlap = torch.einsum(
460 "adr,bds->abrs",
461 source_bases[source_start:source_stop],
462 target_bases[target_start:target_stop],
463 )
464 scores[source_start:source_stop, target_start:target_stop] = overlap.square().sum(
465 dim=(-2, -1)
466 )
467 return scores
470def _architecture_name(model: Any) -> str:
471 cfg = getattr(model, "cfg", None)
472 architecture = getattr(cfg, "original_architecture", None)
473 return str(architecture) if architecture is not None else type(model).__name__
476def _read_role_matrix(model: Any, block: Any, layer: int, role: AttentionRole) -> torch.Tensor:
477 attn = getattr(block, "attn", None)
478 if attn is None: 478 ↛ 479line 478 didn't jump to line 479 because the condition on line 478 was never true
479 raise ValueError(f"attention block {layer} does not expose an attn component")
480 attribute = f"W_{role}"
481 try:
482 matrix = getattr(attn, attribute)
483 except NotImplementedError as error:
484 raise NotImplementedError(
485 f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}"
486 ) from error
487 except (AttributeError, RuntimeError, ValueError) as error:
488 raise ValueError(
489 f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}"
490 ) from error
491 if not isinstance(matrix, torch.Tensor): 491 ↛ 492line 491 didn't jump to line 492 because the condition on line 491 was never true
492 raise ValueError(
493 f"role {role} at layer {layer} must be a tensor, got {type(matrix).__name__}"
494 )
495 if matrix.ndim != 3:
496 raise ValueError(
497 f"role {role} at layer {layer} expected a three-dimensional per-head weight, "
498 f"got shape {tuple(matrix.shape)}"
499 )
500 if role == "O":
501 matrix = matrix.transpose(-2, -1)
502 if not torch.is_floating_point(matrix): 502 ↛ 503line 502 didn't jump to line 503 because the condition on line 502 was never true
503 raise ValueError(f"role {role} at layer {layer} must have a floating-point dtype")
504 if not bool(torch.isfinite(matrix).all()):
505 raise ValueError(f"role {role} at layer {layer} must contain only finite values")
506 return matrix.detach()
509def _validate_role_shapes(
510 matrices: Sequence[torch.Tensor], layers: Sequence[int], role: AttentionRole
511) -> Tuple[int, int, int]:
512 expected = (matrices[0].shape[0], matrices[0].shape[1], matrices[0].shape[2])
513 if min(expected) < 1:
514 raise ValueError(
515 f"role {role} at layer {layers[0]} must have non-empty head, d_model, and width "
516 f"dimensions, got shape {expected}"
517 )
518 for matrix, layer in zip(matrices[1:], layers[1:]):
519 if tuple(matrix.shape) != expected:
520 raise ValueError(
521 f"role {role} at layer {layer} has shape {tuple(matrix.shape)}, expected "
522 f"the consistent [heads, d_model, width] shape {expected}"
523 )
524 return expected
527def _extract_bases(
528 matrices: Sequence[torch.Tensor],
529 layers: Sequence[int],
530 role: AttentionRole,
531 *,
532 selected_rank: int,
533 rtol: float,
534 dtype: torch.dtype,
535 device: torch.device,
536) -> Tuple[torch.Tensor, torch.Tensor]:
537 layer_bases: List[torch.Tensor] = []
538 layer_ranks: List[List[int]] = []
539 for matrix, layer in zip(matrices, layers):
540 head_bases: List[torch.Tensor] = []
541 head_ranks: List[int] = []
542 for head in range(matrix.shape[0]):
543 try:
544 subspace = orthonormal_subspace(
545 matrix[head].to(dtype=dtype), rank=selected_rank, rtol=rtol
546 )
547 except ValueError as error:
548 raise ValueError(
549 f"Could not extract role {role} at layer {layer}, head {head}: {error}"
550 ) from error
551 head_bases.append(subspace.basis.to(device=device))
552 head_ranks.append(subspace.measured_rank)
553 layer_bases.append(torch.stack(head_bases))
554 layer_ranks.append(head_ranks)
555 return torch.stack(layer_bases), torch.tensor(layer_ranks, dtype=torch.long, device=device)
558def attention_head_subspace_affinity(
559 model: Any,
560 *,
561 source_role: str = "O",
562 target_role: str,
563 layer_order: str = "forward",
564 rank: Optional[int] = None,
565 rtol: Optional[float] = None,
566) -> HeadAffinityResult:
567 """Compute OQ, OK, or OV Projection Kernel affinities for a TransformerBridge.
569 K/V axes preserve native key-value heads on grouped-query attention models;
570 they are never expanded to query-head count. Hybrid models include only
571 attention blocks and report their original block indices.
573 Args:
574 model: A TransformerBridge exposing readable per-head attention weights.
575 source_role: Source role; v1 supports only ``"O"``.
576 target_role: One of ``"Q"``, ``"K"``, or ``"V"``.
577 layer_order: ``"forward"`` keeps strict earlier-to-later pairs; ``"all"``
578 keeps every pair.
579 rank: Optional common truncation rank. By default every head must be full
580 column rank.
581 rtol: Optional relative numerical-rank tolerance.
583 Returns:
584 Affinity tensors, validity mask, original layer indices, and rank metadata.
585 """
586 if source_role != "O":
587 raise ValueError(f"source_role must be 'O' in v1, got {source_role!r}")
588 if target_role not in ("Q", "K", "V"):
589 raise ValueError(f"target_role must be one of ['Q', 'K', 'V'], got {target_role!r}")
590 if layer_order not in ("forward", "all"):
591 raise ValueError(f"layer_order must be one of ['forward', 'all'], got {layer_order!r}")
592 validated_source_role = cast(AttentionRole, source_role)
593 validated_target_role = cast(AttentionRole, target_role)
594 validated_layer_order = cast(LayerOrder, layer_order)
595 blocks_with = getattr(model, "blocks_with", None)
596 if not callable(blocks_with): 596 ↛ 597line 596 didn't jump to line 597 because the condition on line 596 was never true
597 raise ValueError("model must be a TransformerBridge exposing blocks_with('attn')")
598 attention_blocks = list(blocks_with("attn"))
599 if not attention_blocks:
600 raise ValueError("No attention layers found — cannot compute head subspace affinity.")
602 layer_indices = [int(layer) for layer, _ in attention_blocks]
603 source_matrices = [
604 _read_role_matrix(model, block, layer, validated_source_role)
605 for layer, block in attention_blocks
606 ]
607 target_matrices = [
608 _read_role_matrix(model, block, layer, validated_target_role)
609 for layer, block in attention_blocks
610 ]
611 source_heads, source_ambient, source_width = _validate_role_shapes(
612 source_matrices, layer_indices, validated_source_role
613 )
614 target_heads, target_ambient, target_width = _validate_role_shapes(
615 target_matrices, layer_indices, validated_target_role
616 )
617 if source_ambient != target_ambient: 617 ↛ 618line 617 didn't jump to line 618 because the condition on line 617 was never true
618 raise ValueError(
619 f"roles {validated_source_role} and {validated_target_role} must share d_model, got "
620 f"{source_ambient} and {target_ambient}"
621 )
623 matrices = source_matrices + target_matrices
624 dtype = matrices[0].dtype
625 for matrix in matrices[1:]:
626 dtype = torch.promote_types(dtype, matrix.dtype)
627 dtype = compute_dtype(dtype)
628 effective_rtol = _validate_rtol(rtol, (source_ambient, max(source_width, target_width)), dtype)
629 requested_rank = _validate_rank(rank, min(source_ambient, source_width, target_width))
630 source_rank = source_width if requested_rank is None else requested_rank
631 target_rank = target_width if requested_rank is None else requested_rank
633 cfg = getattr(model, "cfg", None)
634 configured_device = getattr(cfg, "device", None)
635 result_device = (
636 torch.device(configured_device)
637 if configured_device is not None
638 else source_matrices[0].device
639 )
640 source_bases, source_ranks = _extract_bases(
641 source_matrices,
642 layer_indices,
643 validated_source_role,
644 selected_rank=source_rank,
645 rtol=effective_rtol,
646 dtype=dtype,
647 device=result_device,
648 )
649 target_bases, target_ranks = _extract_bases(
650 target_matrices,
651 layer_indices,
652 validated_target_role,
653 selected_rank=target_rank,
654 rtol=effective_rtol,
655 dtype=dtype,
656 device=result_device,
657 )
659 layer_count = len(layer_indices)
660 flat_source = source_bases.reshape(layer_count * source_heads, source_ambient, source_rank)
661 flat_target = target_bases.reshape(layer_count * target_heads, target_ambient, target_rank)
662 scores = _pairwise_projection_kernel(flat_source, flat_target).reshape(
663 layer_count, source_heads, layer_count, target_heads
664 )
665 scores = _clamp_projection_scores(scores, min(source_rank, target_rank))
666 normalized = scores / math.sqrt(source_rank * target_rank)
667 if validated_layer_order == "forward":
668 layer_tensor = torch.tensor(layer_indices, device=result_device)
669 layer_mask = layer_tensor[:, None] < layer_tensor[None, :]
670 valid_mask = layer_mask[:, None, :, None].expand_as(scores).contiguous()
671 else:
672 valid_mask = torch.ones_like(scores, dtype=torch.bool)
673 scores = torch.where(valid_mask, scores, torch.zeros_like(scores))
674 normalized = torch.where(valid_mask, normalized, torch.zeros_like(normalized))
676 target_kind: HeadKind = "query" if validated_target_role == "Q" else "kv"
677 return HeadAffinityResult(
678 scores=scores,
679 normalized=normalized,
680 valid_mask=valid_mask,
681 source_role=validated_source_role,
682 target_role=validated_target_role,
683 source_layer_indices=tuple(layer_indices),
684 target_layer_indices=tuple(layer_indices),
685 source_head_kind="query",
686 target_head_kind=target_kind,
687 source_ranks=source_ranks,
688 target_ranks=target_ranks,
689 source_rank=source_rank,
690 target_rank=target_rank,
691 rank=rank,
692 rtol=effective_rtol,
693 )