From 69653575f94fd8a9a7b98ac479867a9408d97103 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 24 Aug 2026 18:00:56 +0530 Subject: [PATCH 1/2] Add projection-kernel subspace geometry --- docs/source/content/projection_kernel.md | 83 +++++ docs/source/index.md | 1 + tests/unit/tools/test_projection_kernel.py | 282 ++++++++++++++++ transformer_lens/tools/analysis/__init__.py | 15 + .../tools/analysis/projection_kernel.py | 306 ++++++++++++++++++ 5 files changed, 687 insertions(+) create mode 100644 docs/source/content/projection_kernel.md create mode 100644 tests/unit/tools/test_projection_kernel.py create mode 100644 transformer_lens/tools/analysis/projection_kernel.py diff --git a/docs/source/content/projection_kernel.md b/docs/source/content/projection_kernel.md new file mode 100644 index 000000000..a09ffb7f1 --- /dev/null +++ b/docs/source/content/projection_kernel.md @@ -0,0 +1,83 @@ +# Projection Kernel + +Projection Kernel (PK) measures overlap between two linear subspaces without depending on +the choice of basis within either subspace. TransformerLens provides a model-independent +numerical API for comparing linear subspaces. + +## Definition + +For subspaces $S,T \subseteq \mathbb{R}^d$ with orthonormal bases +$U \in \mathbb{R}^{d \times r}$ and $V \in \mathbb{R}^{d \times s}$, + +$$ +\operatorname{PK}(S,T) = \lVert U^\top V \rVert_F^2 + = \sum_i \cos^2(\theta_i), +$$ + +where $\theta_i$ are the principal angles. Raw PK lies in $[0,\min(r,s)]$. +TransformerLens also returns + +$$ +\frac{\operatorname{PK}(S,T)}{\sqrt{rs}}, +$$ + +the cosine between the two projection matrices. For equal rank $m$, this is +PK divided by $m$. + +PK measures shared geometric support. It does not measure weight magnitude, prove that one +head composes with another, identify head function, or establish a causal pathway. + +## Model-independent API + +Extract rank explicitly before scoring: + +```python +import torch + +from transformer_lens.tools.analysis import orthonormal_subspace, projection_kernel + +matrix_a = torch.randn(32, 4) +matrix_b = torch.randn(32, 6) +basis_a = orthonormal_subspace(matrix_a) +basis_b = orthonormal_subspace(matrix_b) +result = projection_kernel(basis_a, basis_b) + +print(result.score) +print(result.normalized) +print(result.cosines) +print(result.angles) +``` + +`orthonormal_subspace` uses a reduced SVD. Its default relative rank tolerance is +`max(matrix.shape) * eps` in the computation dtype. Supplying `rank` selects the leading +singular subspace, but the requested rank cannot exceed the measured numerical rank. + +Float64 inputs remain float64. Float32 inputs remain float32. Float16 and bfloat16 inputs are +promoted to float32 before SVD, and outputs remain float32. Inputs must be finite, +two-dimensional, real floating-point tensors. + +## Random-subspace reference + +For independent Haar-distributed rank-$m$ planes in $\mathbb{R}^d$, +`random_projection_kernel_moments(d, m)` returns + +$$ +\mathbb{E}[\operatorname{PK}] = \frac{m^2}{d}, \qquad +\operatorname{Var}(\operatorname{PK}) = +\frac{2m^2(d-m)^2}{d^2(d-1)(d+2)}. +$$ + +These moments are descriptive. Trained heads are dependent and anisotropic, so the helper +does not return a p-value or claim a calibrated significance test. + +## Relationship to Composition Score + +PK discards singular-value magnitude and asks whether two read/write spaces overlap. +Composition Score retains the scale of the full linear maps and asks how strongly they +compose under its assumptions. They are complementary metrics and can rank pairs +differently. High values from either metric should be treated as candidate relationships +for activation-level or causal follow-up. + +The method follows Hiroaki Yamagiwa, Yusuke Takase, and Hidetoshi Shimodaira, +“Measuring Affinity between Attention-Head Weight Subspaces via the Projection Kernel,” +[arXiv:2601.10266](https://arxiv.org/abs/2601.10266). \ No newline at end of file diff --git a/docs/source/index.md b/docs/source/index.md index cb1f19692..6f85aea9d 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -57,6 +57,7 @@ content/contributing content/hook_system content/compatibility_mode content/ssm_interpretability +content/projection_kernel content/jacobian_lens_fitting content/debugging_numerical_divergence generated/demos/Main_Demo diff --git a/tests/unit/tools/test_projection_kernel.py b/tests/unit/tools/test_projection_kernel.py new file mode 100644 index 000000000..9054e92c8 --- /dev/null +++ b/tests/unit/tools/test_projection_kernel.py @@ -0,0 +1,282 @@ +"""Unit tests for Projection Kernel subspace geometry.""" + +import math + +import pytest +import torch + +from transformer_lens.tools.analysis.projection_kernel import ( + SubspaceBasis, + orthonormal_subspace, + projection_kernel, + random_projection_kernel_moments, +) + + +class TestOrthonormalSubspace: + def test_extracts_full_rank_tall_basis(self): + matrix = torch.tensor([[1.0, 0.0], [0.0, 2.0], [1.0, 1.0], [0.0, 1.0]], dtype=torch.float64) + + result = orthonormal_subspace(matrix) + + assert result.rank == 2 + assert result.measured_rank == 2 + assert result.ambient_dim == 4 + assert result.input_shape == (4, 2) + assert result.basis.dtype == torch.float64 + assert torch.allclose(result.basis.T @ result.basis, torch.eye(2, dtype=torch.float64)) + + def test_explicit_rank_truncates_but_records_measured_rank(self): + matrix = torch.diag(torch.tensor([3.0, 2.0, 1.0])) + + result = orthonormal_subspace(matrix, rank=2) + + assert result.rank == 2 + assert result.measured_rank == 3 + assert result.basis.shape == (3, 2) + assert result.singular_values.tolist() == pytest.approx([3.0, 2.0, 1.0]) + + def test_threshold_boundary_is_excluded(self): + matrix = torch.diag(torch.tensor([1.0, 0.25], dtype=torch.float64)) + + result = orthonormal_subspace(matrix, rtol=0.25) + + assert result.measured_rank == 1 + assert result.threshold == pytest.approx(0.25) + + def test_invariant_to_scale_and_invertible_right_transform(self): + matrix = torch.tensor( + [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, -1.0]], dtype=torch.float64 + ) + transform = torch.tensor([[2.0, 1.0], [-1.0, 3.0]], dtype=torch.float64) + original = orthonormal_subspace(matrix) + scaled = orthonormal_subspace(7.0 * matrix) + transformed = orthonormal_subspace(matrix @ transform) + + assert projection_kernel(original, scaled).normalized.item() == pytest.approx(1.0) + assert projection_kernel(original, transformed).normalized.item() == pytest.approx(1.0) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_low_precision_promotes_to_float32(self, dtype): + result = orthonormal_subspace(torch.eye(3, dtype=dtype)) + + assert result.basis.dtype == torch.float32 + assert result.singular_values.dtype == torch.float32 + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) + def test_preserves_supported_dtype_and_cpu_device(self, dtype): + result = orthonormal_subspace(torch.eye(3, dtype=dtype, device="cpu")) + + assert result.basis.dtype == dtype + assert result.singular_values.dtype == dtype + assert result.basis.device == torch.device("cpu") + + @pytest.mark.parametrize( + ("matrix", "message"), + [ + (torch.ones(3), "two-dimensional"), + (torch.ones(2, 2, dtype=torch.int64), "floating-point"), + (torch.ones(2, 2, dtype=torch.complex64), "floating-point"), + (torch.empty(0, 2), "non-empty"), + (torch.tensor([[1.0, float("nan")]]), "finite"), + (torch.zeros(2, 2), "numerical rank is zero"), + ], + ) + def test_rejects_invalid_matrices(self, matrix, message): + with pytest.raises(ValueError, match=message): + orthonormal_subspace(matrix) + + @pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"rank": True}, "rank must be an integer"), + ({"rank": 0}, "rank must be between"), + ({"rank": 3}, "rank must be between"), + ({"rtol": -1.0}, "rtol must be"), + ({"rtol": float("inf")}, "rtol must be"), + ], + ) + def test_rejects_invalid_rank_options(self, kwargs, message): + with pytest.raises(ValueError, match=message): + orthonormal_subspace(torch.eye(2), **kwargs) + + def test_rejects_rank_above_measured_rank(self): + matrix = torch.tensor([[1.0, 0.0], [0.0, 0.0], [0.0, 0.0]]) + + with pytest.raises(ValueError, match="exceeds measured rank 1"): + orthonormal_subspace(matrix, rank=2) + + +class TestProjectionKernel: + def test_identical_and_orthogonal_spaces(self): + first = orthonormal_subspace(torch.eye(4)[:, :2]) + same = projection_kernel(first, first) + orthogonal = projection_kernel(first, orthonormal_subspace(torch.eye(4)[:, 2:])) + + assert same.score.item() == pytest.approx(2.0) + assert same.normalized.item() == pytest.approx(1.0) + assert same.cosines.tolist() == pytest.approx([1.0, 1.0]) + assert same.angles.tolist() == pytest.approx([0.0, 0.0]) + assert orthogonal.score.item() == pytest.approx(0.0) + assert orthogonal.angles.tolist() == pytest.approx([math.pi / 2, math.pi / 2]) + + def test_recovers_known_principal_angles(self): + theta = torch.tensor(math.pi / 3, dtype=torch.float64) + basis_a = torch.eye(4, dtype=torch.float64)[:, :2] + basis_b = torch.stack( + [ + torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=torch.float64), + torch.stack([torch.tensor(0.0), theta.cos(), theta.sin(), torch.tensor(0.0)]), + ], + dim=1, + ) + + result = projection_kernel(orthonormal_subspace(basis_a), orthonormal_subspace(basis_b)) + + assert result.cosines.tolist() == pytest.approx([1.0, 0.5]) + assert result.angles.tolist() == pytest.approx([0.0, math.pi / 3]) + assert result.score.item() == pytest.approx(1.25) + + def test_nested_unequal_rank_spaces(self): + small = orthonormal_subspace(torch.eye(4)[:, :2]) + large = orthonormal_subspace(torch.eye(4)[:, :3]) + + result = projection_kernel(small, large) + + assert result.score.item() == pytest.approx(2.0) + assert result.normalized.item() == pytest.approx(math.sqrt(2 / 3)) + assert result.rank_a == 2 + assert result.rank_b == 3 + + def test_equivalent_definitions_and_symmetry(self): + generator = torch.Generator().manual_seed(3) + first = orthonormal_subspace(torch.randn(6, 2, generator=generator, dtype=torch.float64)) + second = orthonormal_subspace(torch.randn(6, 3, generator=generator, dtype=torch.float64)) + + result = projection_kernel(first, second) + reverse = projection_kernel(second, first) + projector_trace = torch.trace( + (first.basis @ first.basis.T) @ (second.basis @ second.basis.T) + ) + + assert result.score.item() == pytest.approx(result.cosines.square().sum().item()) + assert result.score.item() == pytest.approx(projector_trace.item()) + assert result.score.item() == pytest.approx(reverse.score.item()) + + def test_checks_orthonormality(self): + invalid = SubspaceBasis( + basis=torch.ones(3, 1) * 2, + singular_values=torch.ones(1), + rank=1, + measured_rank=1, + rtol=1e-6, + threshold=1e-6, + input_shape=(3, 1), + ) + + with pytest.raises(ValueError, match="orthonormal"): + projection_kernel(invalid, invalid) + + def test_rejects_invalid_singular_value_metadata(self): + valid = orthonormal_subspace(torch.eye(3)) + invalid = SubspaceBasis( + basis=valid.basis, + singular_values=torch.tensor([1.0, float("nan"), 1.0]), + rank=valid.rank, + measured_rank=valid.measured_rank, + rtol=valid.rtol, + threshold=valid.threshold, + input_shape=valid.input_shape, + ) + + with pytest.raises(ValueError, match="singular_values must be a finite"): + projection_kernel(invalid, valid) + + def test_rejects_nonfinite_basis_metadata(self): + valid = orthonormal_subspace(torch.eye(2)) + nonfinite = SubspaceBasis( + basis=torch.tensor([[float("nan"), 0.0], [0.0, 1.0]]), + singular_values=valid.singular_values, + rank=2, + measured_rank=2, + rtol=valid.rtol, + threshold=valid.threshold, + input_shape=(2, 2), + ) + + with pytest.raises(ValueError, match="basis must be finite"): + projection_kernel(nonfinite, valid) + + def test_clamps_tiny_score_overshoot(self): + valid = orthonormal_subspace(torch.eye(2)) + almost_orthonormal = SubspaceBasis( + basis=valid.basis * (1 + 5e-7), + singular_values=valid.singular_values, + rank=valid.rank, + measured_rank=valid.measured_rank, + rtol=valid.rtol, + threshold=valid.threshold, + input_shape=valid.input_shape, + ) + + result = projection_kernel(almost_orthonormal, almost_orthonormal) + + assert result.score.item() == 2.0 + assert result.normalized.item() == 1.0 + + def test_rejects_large_score_bound_violation(self): + valid = orthonormal_subspace(torch.eye(2)) + invalid = SubspaceBasis( + basis=valid.basis * 1.01, + singular_values=valid.singular_values, + rank=valid.rank, + measured_rank=valid.measured_rank, + rtol=valid.rtol, + threshold=valid.threshold, + input_shape=valid.input_shape, + ) + + with pytest.raises(ValueError, match="outside its theoretical bounds"): + projection_kernel(invalid, invalid, check_orthonormal=False) + + def test_rejects_ambient_dimension_mismatch(self): + first = orthonormal_subspace(torch.eye(3)) + second = orthonormal_subspace(torch.eye(4)) + + with pytest.raises(ValueError, match="ambient dimensions"): + projection_kernel(first, second) + + +class TestRandomProjectionKernelMoments: + def test_formula(self): + result = random_projection_kernel_moments(8, 2) + + assert result.mean == pytest.approx(0.5) + assert result.variance == pytest.approx(2 * 4 * 36 / (64 * 7 * 10)) + + def test_full_space_is_constant(self): + result = random_projection_kernel_moments(5, 5) + + assert result.mean == pytest.approx(5.0) + assert result.variance == pytest.approx(0.0) + + @pytest.mark.parametrize(("ambient_dim", "rank"), [(1, 1), (4, 0), (4, 5), (True, 1)]) + def test_rejects_invalid_dimensions(self, ambient_dim, rank): + with pytest.raises(ValueError): + random_projection_kernel_moments(ambient_dim, rank) + + def test_seeded_monte_carlo_matches_mean(self): + ambient_dim, rank, samples = 8, 2, 1000 + generator = torch.Generator().manual_seed(17) + first, _ = torch.linalg.qr( + torch.randn(samples, ambient_dim, rank, generator=generator, dtype=torch.float64) + ) + second, _ = torch.linalg.qr( + torch.randn(samples, ambient_dim, rank, generator=generator, dtype=torch.float64) + ) + overlap = torch.einsum("sdr,sdk->srk", first, second) + empirical_mean = overlap.square().sum(dim=(-2, -1)).mean().item() + reference = random_projection_kernel_moments(ambient_dim, rank) + standard_error = math.sqrt(reference.variance / samples) + + assert empirical_mean == pytest.approx(reference.mean, abs=6 * standard_error) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index a02d225a6..781913d02 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -12,6 +12,7 @@ - jacobian_lens: The Jacobian lens (J-lens) — per-layer causal transport to the output vocabulary basis, with loading of published lens artifacts, native fitting, readouts, interventions, and J-space sparse decomposition. + - projection_kernel: Basis-invariant subspace overlap and principal angles. """ from transformer_lens.tools.analysis.direct_logit_attribution import ( @@ -33,6 +34,14 @@ estimate_occupancy, get_sparse_decomposition, ) +from transformer_lens.tools.analysis.projection_kernel import ( + ProjectionKernelResult, + RandomSubspaceReference, + SubspaceBasis, + orthonormal_subspace, + projection_kernel, + random_projection_kernel_moments, +) __all__ = [ "DirectLogitAttribution", @@ -41,9 +50,15 @@ "JSpaceVarianceProfile", "JacobianLens", "JacobianLensReadout", + "ProjectionKernelResult", + "RandomSubspaceReference", + "SubspaceBasis", "direct_logit_attribution", "estimate_occupancy", "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", "get_sparse_decomposition", + "orthonormal_subspace", + "projection_kernel", + "random_projection_kernel_moments", ] diff --git a/transformer_lens/tools/analysis/projection_kernel.py b/transformer_lens/tools/analysis/projection_kernel.py new file mode 100644 index 000000000..5d55bc58f --- /dev/null +++ b/transformer_lens/tools/analysis/projection_kernel.py @@ -0,0 +1,306 @@ +"""Projection Kernel utilities for comparing linear subspaces. + +The Projection Kernel (PK) between subspaces with orthonormal bases ``U`` and +``V`` is ``||U.T @ V||_F^2``. It is invariant to basis choices within either +subspace and equals the sum of squared principal-angle cosines. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from numbers import Real +from typing import Optional, Tuple + +import torch + + +@dataclass(frozen=True) +class SubspaceBasis: + """An explicitly ranked orthonormal basis extracted from a matrix.""" + + basis: torch.Tensor + singular_values: torch.Tensor + rank: int + measured_rank: int + rtol: float + threshold: float + input_shape: Tuple[int, int] + + @property + def ambient_dim(self) -> int: + """Dimension of the space containing the subspace.""" + return self.basis.shape[0] + + +@dataclass(frozen=True) +class ProjectionKernelResult: + """Projection Kernel score and its principal-angle decomposition.""" + + score: torch.Tensor + normalized: torch.Tensor + cosines: torch.Tensor + angles: torch.Tensor + rank_a: int + rank_b: int + ambient_dim: int + + +@dataclass(frozen=True) +class RandomSubspaceReference: + """Analytic PK moments for independent random equal-rank subspaces.""" + + ambient_dim: int + rank: int + mean: float + variance: float + + +def _compute_dtype(dtype: torch.dtype) -> torch.dtype: + """Return a dtype supported by stable SVD on common PyTorch backends.""" + return torch.float64 if dtype == torch.float64 else torch.float32 + + +def _validate_rtol(rtol: Optional[float], shape: Tuple[int, int], dtype: torch.dtype) -> float: + if rtol is None: + return max(shape) * torch.finfo(dtype).eps + if isinstance(rtol, bool) or not isinstance(rtol, Real): + raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}") + value = float(rtol) + if not math.isfinite(value) or value < 0: + raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}") + return value + + +def _validate_rank(rank: Optional[int], max_rank: int) -> Optional[int]: + if rank is None: + return None + if isinstance(rank, bool) or not isinstance(rank, int): + raise ValueError(f"rank must be an integer or None, got {rank!r}") + if not 1 <= rank <= max_rank: + raise ValueError(f"rank must be between 1 and min(matrix.shape)={max_rank}, got {rank}") + return rank + + +def orthonormal_subspace( + matrix: torch.Tensor, + *, + rank: Optional[int] = None, + rtol: Optional[float] = None, +) -> SubspaceBasis: + """Extract an explicitly ranked orthonormal column-space basis. + + Low-precision inputs are promoted to float32 before the reduced SVD. With no + explicit ``rtol``, numerical rank uses ``max(matrix.shape) * eps`` relative + to the largest singular value. An explicit ``rank`` truncates the measured + subspace but may not exceed its measured rank. + + Args: + matrix: Finite floating-point matrix with shape ``[ambient_dim, width]``. + rank: Optional number of leading singular directions to retain. + rtol: Optional non-negative relative singular-value threshold. + + Returns: + Basis, complete singular spectrum, and rank metadata. + + Raises: + ValueError: If the matrix or rank policy is invalid. + """ + if not isinstance(matrix, torch.Tensor): + raise ValueError(f"matrix must be a torch.Tensor, got {type(matrix).__name__}") + if matrix.ndim != 2: + raise ValueError(f"matrix must be two-dimensional, got shape {tuple(matrix.shape)}") + if matrix.shape[0] == 0 or matrix.shape[1] == 0: + raise ValueError(f"matrix dimensions must be non-empty, got shape {tuple(matrix.shape)}") + if not torch.is_floating_point(matrix): + raise ValueError(f"matrix must have a real floating-point dtype, got {matrix.dtype}") + if not bool(torch.isfinite(matrix).all()): + raise ValueError("matrix must contain only finite values") + + input_shape = (matrix.shape[0], matrix.shape[1]) + requested_rank = _validate_rank(rank, min(input_shape)) + compute_dtype = _compute_dtype(matrix.dtype) + effective_rtol = _validate_rtol(rtol, input_shape, compute_dtype) + work = matrix.to(dtype=compute_dtype) + left, singular_values, _ = torch.linalg.svd(work, full_matrices=False) + threshold = float(singular_values[0].item()) * effective_rtol + measured_rank = int((singular_values > threshold).sum().item()) + if measured_rank == 0: + raise ValueError( + "matrix numerical rank is zero " + f"for shape {input_shape}, dtype {matrix.dtype}, and threshold {threshold:.6g}" + ) + if requested_rank is not None and requested_rank > measured_rank: + raise ValueError( + f"requested rank {requested_rank} exceeds measured rank {measured_rank} " + f"at threshold {threshold:.6g}" + ) + + selected_rank = measured_rank if requested_rank is None else requested_rank + return SubspaceBasis( + basis=left[:, :selected_rank], + singular_values=singular_values, + rank=selected_rank, + measured_rank=measured_rank, + rtol=effective_rtol, + threshold=threshold, + input_shape=input_shape, + ) + + +def _validate_subspace(value: SubspaceBasis, name: str) -> None: + if not isinstance(value, SubspaceBasis): + raise ValueError(f"{name} must be a SubspaceBasis, got {type(value).__name__}") + basis = value.basis + if not isinstance(basis, torch.Tensor) or basis.ndim != 2: + raise ValueError(f"{name}.basis must be a two-dimensional tensor") + if not torch.is_floating_point(basis) or not bool(torch.isfinite(basis).all()): + raise ValueError(f"{name}.basis must be finite and real floating-point") + singular_values = value.singular_values + if ( + not isinstance(singular_values, torch.Tensor) + or singular_values.ndim != 1 + or not torch.is_floating_point(singular_values) + or not bool(torch.isfinite(singular_values).all()) + ): + raise ValueError(f"{name}.singular_values must be a finite floating-point vector") + if singular_values.device != basis.device: + raise ValueError(f"{name}.singular_values must be on the basis device") + if ( + not isinstance(value.input_shape, tuple) + or len(value.input_shape) != 2 + or any( + isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 1 + for dimension in value.input_shape + ) + ): + raise ValueError(f"{name}.input_shape must contain two positive dimensions") + if singular_values.shape[0] != min(value.input_shape): + raise ValueError(f"{name}.singular_values length must equal min(input_shape)") + if ( + isinstance(value.rank, bool) + or not isinstance(value.rank, int) + or value.rank < 1 + or basis.shape[1] != value.rank + ): + raise ValueError(f"{name}.rank must equal its positive basis width") + if ( + isinstance(value.measured_rank, bool) + or not isinstance(value.measured_rank, int) + or value.measured_rank < value.rank + or value.measured_rank > min(value.input_shape) + ): + raise ValueError(f"{name}.measured_rank must be at least its selected rank") + if value.input_shape[0] != basis.shape[0]: + raise ValueError(f"{name}.input_shape must share its basis ambient dimension") + if ( + isinstance(value.rtol, bool) + or not isinstance(value.rtol, Real) + or not math.isfinite(value.rtol) + or value.rtol < 0 + ): + raise ValueError(f"{name}.rtol must be finite and non-negative") + if ( + isinstance(value.threshold, bool) + or not isinstance(value.threshold, Real) + or not math.isfinite(value.threshold) + or value.threshold < 0 + ): + raise ValueError(f"{name}.threshold must be finite and non-negative") + + +def _clamp_projection_scores(scores: torch.Tensor, upper_bound: int) -> torch.Tensor: + """Clamp roundoff-scale PK bound violations and reject larger violations.""" + tolerance = 100 * torch.finfo(scores.dtype).eps * max(1, upper_bound) + minimum = float(scores.detach().min().item()) + maximum = float(scores.detach().max().item()) + if minimum < -tolerance or maximum > upper_bound + tolerance: + raise ValueError( + "Projection Kernel score lies outside its theoretical bounds: " + f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, {upper_bound}] " + f"within tolerance {tolerance:.6g}" + ) + return scores.clamp(min=0.0, max=float(upper_bound)) + + +def projection_kernel( + subspace_a: SubspaceBasis, + subspace_b: SubspaceBasis, + *, + check_orthonormal: bool = True, +) -> ProjectionKernelResult: + """Measure overlap between two explicitly extracted subspaces. + + Raw PK lies in ``[0, min(rank_a, rank_b)]``. The normalized value is + ``PK / sqrt(rank_a * rank_b)``, the cosine between the two projection + matrices. Principal angles are returned in radians. + """ + if not isinstance(check_orthonormal, bool): + raise ValueError(f"check_orthonormal must be a Boolean, got {check_orthonormal!r}") + _validate_subspace(subspace_a, "subspace_a") + _validate_subspace(subspace_b, "subspace_b") + if subspace_a.ambient_dim != subspace_b.ambient_dim: + raise ValueError( + "subspaces must have equal ambient dimensions, got " + f"{subspace_a.ambient_dim} and {subspace_b.ambient_dim}" + ) + if subspace_a.basis.device != subspace_b.basis.device: + raise ValueError( + "subspace bases must be on the same device, got " + f"{subspace_a.basis.device} and {subspace_b.basis.device}" + ) + + dtype = torch.promote_types(subspace_a.basis.dtype, subspace_b.basis.dtype) + dtype = _compute_dtype(dtype) + first = subspace_a.basis.to(dtype=dtype) + second = subspace_b.basis.to(dtype=dtype) + if check_orthonormal: + eps = torch.finfo(dtype).eps + tolerance = 10 * max(first.shape[0], first.shape[1], second.shape[1]) * eps + first_identity = torch.eye(first.shape[1], dtype=dtype, device=first.device) + second_identity = torch.eye(second.shape[1], dtype=dtype, device=second.device) + if not torch.allclose(first.T @ first, first_identity, rtol=tolerance, atol=tolerance): + raise ValueError("subspace_a basis columns must be orthonormal") + if not torch.allclose(second.T @ second, second_identity, rtol=tolerance, atol=tolerance): + raise ValueError("subspace_b basis columns must be orthonormal") + + overlap = first.T @ second + score = _clamp_projection_scores(overlap.square().sum(), min(subspace_a.rank, subspace_b.rank)) + cosines = torch.linalg.svdvals(overlap) + angles = torch.acos(cosines.clamp(min=0.0, max=1.0)) + denominator = math.sqrt(subspace_a.rank * subspace_b.rank) + return ProjectionKernelResult( + score=score, + normalized=score / denominator, + cosines=cosines, + angles=angles, + rank_a=subspace_a.rank, + rank_b=subspace_b.rank, + ambient_dim=subspace_a.ambient_dim, + ) + + +def random_projection_kernel_moments(ambient_dim: int, rank: int) -> RandomSubspaceReference: + """Return PK moments for independent Haar-distributed rank-``rank`` planes. + + These idealized descriptive moments are not calibrated p-values for trained + model weights, whose head subspaces are dependent and anisotropic. + """ + if isinstance(ambient_dim, bool) or not isinstance(ambient_dim, int) or ambient_dim < 2: + raise ValueError(f"ambient_dim must be an integer at least 2, got {ambient_dim!r}") + if isinstance(rank, bool) or not isinstance(rank, int) or not 1 <= rank <= ambient_dim: + raise ValueError(f"rank must be an integer between 1 and {ambient_dim}, got {rank!r}") + + mean = rank**2 / ambient_dim + variance = ( + 2 + * rank**2 + * (ambient_dim - rank) ** 2 + / (ambient_dim**2 * (ambient_dim - 1) * (ambient_dim + 2)) + ) + return RandomSubspaceReference( + ambient_dim=ambient_dim, + rank=rank, + mean=float(mean), + variance=float(variance), + ) From ec0575023ca75605702c6fc169ad15c93f187d2e Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 24 Aug 2026 18:04:07 +0530 Subject: [PATCH 2/2] Add Bridge attention-head subspace affinity --- docs/source/content/projection_kernel.md | 58 ++- .../test_attention_weight_accessors.py | 27 ++ .../model_bridge/test_projection_kernel.py | 35 ++ tests/unit/tools/test_projection_kernel.py | 217 +++++++++- transformer_lens/tools/analysis/__init__.py | 11 +- .../tools/analysis/projection_kernel.py | 383 +++++++++++++++++- 6 files changed, 726 insertions(+), 5 deletions(-) create mode 100644 tests/integration/model_bridge/test_projection_kernel.py diff --git a/docs/source/content/projection_kernel.md b/docs/source/content/projection_kernel.md index a09ffb7f1..edfef6c4e 100644 --- a/docs/source/content/projection_kernel.md +++ b/docs/source/content/projection_kernel.md @@ -2,7 +2,7 @@ Projection Kernel (PK) measures overlap between two linear subspaces without depending on the choice of basis within either subspace. TransformerLens provides a model-independent -numerical API for comparing linear subspaces. +numerical API and a TransformerBridge wrapper for comparing attention-head weight spaces. ## Definition @@ -56,6 +56,62 @@ Float64 inputs remain float64. Float32 inputs remain float32. Float16 and bfloat promoted to float32 before SVD, and outputs remain float32. Inputs must be finite, two-dimensional, real floating-point tensors. +## Attention-head affinity + +The TransformerBridge wrapper computes OQ, OK, or OV affinity for every selected head pair: + +```python +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.tools.analysis import attention_head_subspace_affinity + +model = TransformerBridge.boot_transformers("gpt2", device="cpu") +result = attention_head_subspace_affinity( + model, + source_role="O", + target_role="Q", + layer_order="forward", +) + +print(result.scores.shape) +print(result.normalized.shape) +print(result.top_pairs(20, normalized=True)) +``` + +The axes are +`[source_attention_layer, source_head, target_attention_layer, target_head]`. +`source_layer_indices` and `target_layer_indices` map tensor positions to original block +numbers. `valid_mask` has the same shape as the score tensors. Invalid entries are zero. + +With `layer_order="forward"`, only strict earlier-to-later pairs are valid. Use +`layer_order="all"` to include every source-target layer pair. + +## Weight orientation + +TransformerLens exposes each basis-generating matrix in residual-stream coordinates: + +| Role | Matrix used as a basis | Per-head shape | +|---|---|---| +| Q | `W_Q` | `[d_model, d_head]` | +| K | `W_K` | `[d_model, d_head]` | +| V | `W_V` | `[d_model, d_head]` | +| O | `W_O.T` | `[d_model, d_head]` | + +The O transpose is required: `W_O` itself has shape `[d_head, d_model]`. + +## MHA, GQA, and hybrid models + +- Multi-head attention produces query-head axes for O and Q and K/V role axes of the same + size. +- Grouped-query attention preserves native K/V heads. OK and OV are therefore rectangular: + query heads by KV heads. K/V weights are not repeated to query-head count. +- Hybrid models include only blocks exposing bridged attention and retain original block + numbers in their layer metadata. +- Architectures without readable standard Q/K/V/O projections, such as MLA or opaque + native-forward attention, fail with a role- and layer-specific error. + +By default, every head must be full column rank. An explicit `rank` applies the same +truncation to both roles and must not exceed any participating head's measured rank. + ## Random-subspace reference For independent Haar-distributed rank-$m$ planes in $\mathbb{R}^d$, diff --git a/tests/integration/model_bridge/test_attention_weight_accessors.py b/tests/integration/model_bridge/test_attention_weight_accessors.py index 8afc91dbd..c81592255 100644 --- a/tests/integration/model_bridge/test_attention_weight_accessors.py +++ b/tests/integration/model_bridge/test_attention_weight_accessors.py @@ -182,3 +182,30 @@ def test_mha_circuits_untouched(self, tiny_gpt2_bridge): assert torch.equal(QK.B, bridge.W_K.transpose(-2, -1)) assert torch.equal(OV.A, bridge.W_V) assert torch.equal(OV.B, bridge.W_O) + + +class TestProjectionKernelGQA: + """Head affinity keeps grouped K/V heads native rather than expanding them.""" + + @pytest.mark.parametrize("role", ["K", "V"]) + def test_native_kv_axes_and_sample_parity(self, llama_bridge, role): + from transformer_lens.tools.analysis.projection_kernel import ( + attention_head_subspace_affinity, + orthonormal_subspace, + projection_kernel, + ) + + bridge, _ = llama_bridge + result = attention_head_subspace_affinity(bridge, target_role=role) + target_weight = getattr(bridge.blocks[1].attn, f"W_{role}")[1] + expected = projection_kernel( + orthonormal_subspace(bridge.blocks[0].attn.W_O[0].T), + orthonormal_subspace(target_weight), + ) + + assert result.scores.shape == (2, 4, 2, 2) + assert int(result.valid_mask.sum()) == 8 + assert result.target_head_kind == "kv" + assert result.scores[0, 0, 1, 1].item() == pytest.approx( + expected.score.item(), rel=1e-5, abs=1e-5 + ) diff --git a/tests/integration/model_bridge/test_projection_kernel.py b/tests/integration/model_bridge/test_projection_kernel.py new file mode 100644 index 000000000..a3194ff80 --- /dev/null +++ b/tests/integration/model_bridge/test_projection_kernel.py @@ -0,0 +1,35 @@ +"""Real-Bridge integration checks for Projection Kernel head affinity.""" + +import pytest + +from transformer_lens.tools.analysis.projection_kernel import ( + attention_head_subspace_affinity, + orthonormal_subspace, + projection_kernel, +) + + +@pytest.mark.parametrize(("role", "attribute"), [("Q", "W_Q"), ("K", "W_K"), ("V", "W_V")]) +def test_gpt2_head_affinity_contract_and_sample_parity(gpt2_bridge, role, attribute): + result = attention_head_subspace_affinity(gpt2_bridge, target_role=role) + + assert result.scores.shape == (12, 12, 12, 12) + assert result.source_layer_indices == tuple(range(12)) + assert result.target_layer_indices == tuple(range(12)) + assert int(result.valid_mask.sum()) == 9504 + assert result.source_head_kind == "query" + assert result.target_head_kind == ("query" if role == "Q" else "kv") + assert bool((result.scores[result.valid_mask] >= -1e-5).all()) + assert bool((result.scores[result.valid_mask] <= 64 + 1e-4).all()) + assert bool((result.normalized[result.valid_mask] >= -1e-6).all()) + assert bool((result.normalized[result.valid_mask] <= 1 + 1e-5).all()) + + source = orthonormal_subspace(gpt2_bridge.blocks[0].attn.W_O[0].T) + target_weight = getattr(gpt2_bridge.blocks[1].attn, attribute)[1] + expected = projection_kernel(source, orthonormal_subspace(target_weight)) + assert result.scores[0, 0, 1, 1].item() == pytest.approx( + expected.score.item(), rel=1e-5, abs=1e-5 + ) + assert result.normalized[0, 0, 1, 1].item() == pytest.approx( + expected.normalized.item(), rel=1e-5, abs=1e-6 + ) diff --git a/tests/unit/tools/test_projection_kernel.py b/tests/unit/tools/test_projection_kernel.py index 9054e92c8..9ebef5c7c 100644 --- a/tests/unit/tools/test_projection_kernel.py +++ b/tests/unit/tools/test_projection_kernel.py @@ -1,18 +1,46 @@ -"""Unit tests for Projection Kernel subspace geometry.""" +"""Unit tests for Projection Kernel subspace geometry and head affinity.""" import math +from types import SimpleNamespace import pytest import torch from transformer_lens.tools.analysis.projection_kernel import ( SubspaceBasis, + _pairwise_projection_kernel, + attention_head_subspace_affinity, orthonormal_subspace, projection_kernel, random_projection_kernel_moments, ) +class SyntheticAttention: + def __init__(self, seed: int, *, d_model: int = 5, d_head: int = 2): + generator = torch.Generator().manual_seed(seed) + self.W_Q = torch.randn(2, d_model, d_head, generator=generator, dtype=torch.float64) + self.W_K = torch.randn(1, d_model, d_head, generator=generator, dtype=torch.float64) + self.W_V = torch.randn(1, d_model, d_head, generator=generator, dtype=torch.float64) + self.W_O = torch.randn(2, d_head, d_model, generator=generator, dtype=torch.float64) + + +class SyntheticBlock: + def __init__(self, attention): + self.attn = attention + + +class SyntheticBridge: + def __init__(self, layer_indices=(0, 2, 5)): + self.cfg = SimpleNamespace(device="cpu", original_architecture="SyntheticArchitecture") + self.attention_blocks = [ + (layer, SyntheticBlock(SyntheticAttention(10 + layer))) for layer in layer_indices + ] + + def blocks_with(self, submodule): + return self.attention_blocks if submodule == "attn" else [] + + class TestOrthonormalSubspace: def test_extracts_full_rank_tall_basis(self): matrix = torch.tensor([[1.0, 0.0], [0.0, 2.0], [1.0, 1.0], [0.0, 1.0]], dtype=torch.float64) @@ -280,3 +308,190 @@ def test_seeded_monte_carlo_matches_mean(self): standard_error = math.sqrt(reference.variance / samples) assert empirical_mean == pytest.approx(reference.mean, abs=6 * standard_error) + + +class TestPairwiseProjectionKernel: + @pytest.mark.parametrize("max_temp_bytes", [1, 48, 96, 10_000]) + def test_tiled_scores_match_nested_loop_for_unequal_ranks(self, max_temp_bytes): + generator = torch.Generator().manual_seed(23) + source, _ = torch.linalg.qr(torch.randn(5, 7, 2, generator=generator, dtype=torch.float64)) + target, _ = torch.linalg.qr(torch.randn(4, 7, 3, generator=generator, dtype=torch.float64)) + + actual = _pairwise_projection_kernel(source, target, max_temp_bytes=max_temp_bytes) + expected = torch.empty(5, 4, dtype=torch.float64) + for source_index in range(5): + for target_index in range(4): + overlap = source[source_index].T @ target[target_index] + expected[source_index, target_index] = overlap.square().sum() + + assert torch.allclose(actual, expected) + + +class TestAttentionHeadSubspaceAffinity: + def test_gqa_roles_have_native_rectangular_axes_and_hybrid_indices(self): + model = SyntheticBridge() + + oq = attention_head_subspace_affinity(model, target_role="Q") + ok = attention_head_subspace_affinity(model, target_role="K") + ov = attention_head_subspace_affinity(model, target_role="V") + + assert oq.scores.shape == (3, 2, 3, 2) + assert ok.scores.shape == (3, 2, 3, 1) + assert ov.scores.shape == (3, 2, 3, 1) + assert oq.source_layer_indices == (0, 2, 5) + assert oq.target_layer_indices == (0, 2, 5) + assert oq.source_head_kind == "query" + assert oq.target_head_kind == "query" + assert ok.target_head_kind == "kv" + assert ov.target_head_kind == "kv" + assert int(oq.valid_mask.sum()) == 12 + assert int(ok.valid_mask.sum()) == 6 + assert int(ov.valid_mask.sum()) == 6 + + @pytest.mark.parametrize(("role", "attribute"), [("Q", "W_Q"), ("K", "W_K"), ("V", "W_V")]) + def test_sample_matches_independent_o_to_target_calculation(self, role, attribute): + model = SyntheticBridge() + + result = attention_head_subspace_affinity(model, target_role=role) + source_matrix = model.attention_blocks[0][1].attn.W_O[0].T + target_matrix = getattr(model.attention_blocks[1][1].attn, attribute)[-1] + expected = projection_kernel( + orthonormal_subspace(source_matrix), orthonormal_subspace(target_matrix) + ) + + assert result.scores[0, 0, 1, -1].item() == pytest.approx(expected.score.item()) + assert result.normalized[0, 0, 1, -1].item() == pytest.approx(expected.normalized.item()) + + def test_all_layer_order_includes_every_pair(self): + result = attention_head_subspace_affinity( + SyntheticBridge((1, 4)), target_role="K", layer_order="all" + ) + + assert bool(result.valid_mask.all()) + assert int(result.valid_mask.sum()) == 8 + + def test_invalid_entries_are_zero(self): + result = attention_head_subspace_affinity(SyntheticBridge(), target_role="Q") + + assert torch.equal(result.scores[~result.valid_mask], torch.zeros(24, dtype=torch.float64)) + assert torch.equal( + result.normalized[~result.valid_mask], torch.zeros(24, dtype=torch.float64) + ) + + def test_wrapper_does_not_retain_weight_autograd_graph(self): + model = SyntheticBridge((0, 1)) + for _, block in model.attention_blocks: + block.attn.W_Q.requires_grad_() + block.attn.W_O.requires_grad_() + + result = attention_head_subspace_affinity(model, target_role="Q") + + assert not result.scores.requires_grad + assert not result.normalized.requires_grad + + def test_result_dtype_device_and_rank_metadata(self): + model = SyntheticBridge((0, 1)) + for _, block in model.attention_blocks: + block.attn.W_Q = block.attn.W_Q.float() + block.attn.W_O = block.attn.W_O.float() + + result = attention_head_subspace_affinity(model, target_role="Q") + + assert result.scores.dtype == torch.float32 + assert result.scores.device == torch.device("cpu") + assert result.source_ranks.dtype == torch.long + assert result.source_ranks.shape == (2, 2) + assert result.target_ranks.shape == (2, 2) + assert result.source_rank == 2 + assert result.target_rank == 2 + + def test_rank_deficiency_names_role_layer_and_head(self): + model = SyntheticBridge() + deficient = model.attention_blocks[1][1].attn.W_Q[1] + deficient[:, 1] = deficient[:, 0] + + with pytest.raises(ValueError, match="role Q at layer 2, head 1"): + attention_head_subspace_affinity(model, target_role="Q") + + def test_explicit_common_truncation_accepts_rank_one_head(self): + model = SyntheticBridge() + deficient = model.attention_blocks[1][1].attn.W_Q[1] + deficient[:, 1] = deficient[:, 0] + + result = attention_head_subspace_affinity(model, target_role="Q", rank=1) + + assert result.source_rank == 1 + assert result.target_rank == 1 + assert result.target_ranks[1, 1].item() == 1 + assert bool((result.normalized[result.valid_mask] <= 1 + 1e-12).all()) + + def test_top_pairs_excludes_masked_entries_and_breaks_ties_lexicographically(self): + model = SyntheticBridge((0, 3)) + common_q = torch.eye(4, dtype=torch.float64)[:, :2].expand(2, -1, -1).clone() + common_o = common_q.transpose(-2, -1).clone() + for _, block in model.attention_blocks: + block.attn.W_Q = common_q + block.attn.W_O = common_o + + result = attention_head_subspace_affinity(model, target_role="Q") + pairs = result.top_pairs(10) + + assert len(pairs) == 4 + assert [ + (pair.source.layer, pair.source.head, pair.target.layer, pair.target.head) + for pair in pairs + ] == [(0, 0, 3, 0), (0, 0, 3, 1), (0, 1, 3, 0), (0, 1, 3, 1)] + assert all(pair.score == pytest.approx(2.0) for pair in pairs) + + @pytest.mark.parametrize("k", [0, -1, True]) + def test_top_pairs_rejects_invalid_k(self, k): + result = attention_head_subspace_affinity(SyntheticBridge(), target_role="Q") + + with pytest.raises(ValueError, match="k must be"): + result.top_pairs(k) + + def test_rejects_unsupported_pairing_and_layer_order(self): + model = SyntheticBridge() + + with pytest.raises(ValueError, match="source_role must be 'O'"): + attention_head_subspace_affinity(model, source_role="Q", target_role="K") + with pytest.raises(ValueError, match="target_role must be one of"): + attention_head_subspace_affinity(model, target_role="O") + with pytest.raises(ValueError, match="layer_order must be one of"): + attention_head_subspace_affinity(model, target_role="Q", layer_order="backward") + + def test_rejects_no_attention_and_invalid_weight_shape(self): + empty = SyntheticBridge(()) + invalid = SyntheticBridge() + invalid.attention_blocks[0][1].attn.W_K = torch.ones(5, 2) + + with pytest.raises(ValueError, match="No attention layers"): + attention_head_subspace_affinity(empty, target_role="Q") + with pytest.raises(ValueError, match="role K at layer 0.*expected a three-dimensional"): + attention_head_subspace_affinity(invalid, target_role="K") + + def test_rejects_missing_and_nonfinite_weights(self): + missing = SyntheticBridge() + nonfinite = SyntheticBridge() + del missing.attention_blocks[0][1].attn.W_K + nonfinite.attention_blocks[1][1].attn.W_K[0, 0, 0] = float("nan") + + with pytest.raises(ValueError, match="cannot expose role K at layer 0"): + attention_head_subspace_affinity(missing, target_role="K") + with pytest.raises(ValueError, match="role K at layer 2 must contain only finite"): + attention_head_subspace_affinity(nonfinite, target_role="K") + + @pytest.mark.parametrize("shape", [(2, 5, 2), (1, 6, 2), (1, 5, 3)]) + def test_rejects_inconsistent_role_shapes(self, shape): + model = SyntheticBridge() + model.attention_blocks[1][1].attn.W_K = torch.ones(shape, dtype=torch.float64) + + with pytest.raises(ValueError, match="consistent.*shape"): + attention_head_subspace_affinity(model, target_role="K") + + def test_rejects_empty_head_axis(self): + model = SyntheticBridge((0,)) + model.attention_blocks[0][1].attn.W_K = torch.empty(0, 5, 2) + + with pytest.raises(ValueError, match="must have non-empty head"): + attention_head_subspace_affinity(model, target_role="K", layer_order="all") diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index 781913d02..e6c3a02b2 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -12,7 +12,8 @@ - jacobian_lens: The Jacobian lens (J-lens) — per-layer causal transport to the output vocabulary basis, with loading of published lens artifacts, native fitting, readouts, interventions, and J-space sparse decomposition. - - projection_kernel: Basis-invariant subspace overlap and principal angles. + - projection_kernel: Basis-invariant subspace overlap and TransformerBridge + attention-head OQ/OK/OV affinity. """ from transformer_lens.tools.analysis.direct_logit_attribution import ( @@ -35,9 +36,13 @@ get_sparse_decomposition, ) from transformer_lens.tools.analysis.projection_kernel import ( + AttentionHeadRef, + HeadAffinityPair, + HeadAffinityResult, ProjectionKernelResult, RandomSubspaceReference, SubspaceBasis, + attention_head_subspace_affinity, orthonormal_subspace, projection_kernel, random_projection_kernel_moments, @@ -45,6 +50,9 @@ __all__ = [ "DirectLogitAttribution", + "AttentionHeadRef", + "HeadAffinityPair", + "HeadAffinityResult", "JSpaceDecomposition", "JSpaceOccupancy", "JSpaceVarianceProfile", @@ -53,6 +61,7 @@ "ProjectionKernelResult", "RandomSubspaceReference", "SubspaceBasis", + "attention_head_subspace_affinity", "direct_logit_attribution", "estimate_occupancy", "get_act_patch_direct_path", diff --git a/transformer_lens/tools/analysis/projection_kernel.py b/transformer_lens/tools/analysis/projection_kernel.py index 5d55bc58f..8dca25f4b 100644 --- a/transformer_lens/tools/analysis/projection_kernel.py +++ b/transformer_lens/tools/analysis/projection_kernel.py @@ -10,14 +10,30 @@ import math from dataclasses import dataclass from numbers import Real -from typing import Optional, Tuple +from typing import Any, List, Literal, Optional, Sequence, Tuple, cast import torch +AttentionRole = Literal["Q", "K", "V", "O"] +LayerOrder = Literal["forward", "all"] +HeadKind = Literal["query", "kv"] + +_PAIRWISE_TEMP_BYTES = 64 * 1024 * 1024 + @dataclass(frozen=True) class SubspaceBasis: - """An explicitly ranked orthonormal basis extracted from a matrix.""" + """An explicitly ranked orthonormal basis extracted from a matrix. + + Attributes: + basis: Orthonormal column-space basis, ``[ambient_dim, rank]``. + singular_values: All reduced-SVD singular values, in descending order. + rank: Number of retained basis directions. + measured_rank: Numerical rank before optional caller truncation. + rtol: Effective relative rank tolerance. + threshold: Absolute singular-value threshold used for rank measurement. + input_shape: Shape of the matrix from which the basis was extracted. + """ basis: torch.Tensor singular_values: torch.Tensor @@ -56,6 +72,103 @@ class RandomSubspaceReference: variance: float +@dataclass(frozen=True) +class AttentionHeadRef: + """Structured identity for one attention-head weight subspace.""" + + layer: int + head: int + role: AttentionRole + kind: HeadKind + + @property + def label(self) -> str: + """Return the conventional TransformerLens layer/head label.""" + return f"L{self.layer}H{self.head}" + + +@dataclass(frozen=True) +class HeadAffinityPair: + """One ranked source-target head pair.""" + + source: AttentionHeadRef + target: AttentionHeadRef + score: float + normalized: float + + +@dataclass(frozen=True) +class HeadAffinityResult: + """Projection Kernel affinities between two attention-head roles. + + Score tensors have shape + ``[source_layer, source_head, target_layer, target_head]``. Layer index + tuples map tensor positions to original model block numbers. + """ + + scores: torch.Tensor + normalized: torch.Tensor + valid_mask: torch.Tensor + source_role: AttentionRole + target_role: AttentionRole + source_layer_indices: Tuple[int, ...] + target_layer_indices: Tuple[int, ...] + source_head_kind: HeadKind + target_head_kind: HeadKind + source_ranks: torch.Tensor + target_ranks: torch.Tensor + source_rank: int + target_rank: int + rank: Optional[int] + rtol: float + + def top_pairs(self, k: int = 20, *, normalized: bool = False) -> List[HeadAffinityPair]: + """Return the highest-scoring valid pairs with deterministic tie order.""" + if isinstance(k, bool) or not isinstance(k, int) or k < 1: + raise ValueError(f"k must be a positive integer, got {k!r}") + + selected_scores = self.normalized if normalized else self.scores + entries: List[HeadAffinityPair] = [] + for source_layer, source_head, target_layer, target_head in ( + torch.nonzero(self.valid_mask, as_tuple=False).cpu().tolist() + ): + source = AttentionHeadRef( + layer=self.source_layer_indices[source_layer], + head=source_head, + role=self.source_role, + kind=self.source_head_kind, + ) + target = AttentionHeadRef( + layer=self.target_layer_indices[target_layer], + head=target_head, + role=self.target_role, + kind=self.target_head_kind, + ) + entries.append( + HeadAffinityPair( + source=source, + target=target, + score=float(self.scores[source_layer, source_head, target_layer, target_head]), + normalized=float( + self.normalized[source_layer, source_head, target_layer, target_head] + ), + ) + ) + + def sort_key(pair: HeadAffinityPair) -> Tuple[float, int, int, int, int]: + value = pair.normalized if normalized else pair.score + return ( + -value, + pair.source.layer, + pair.source.head, + pair.target.layer, + pair.target.head, + ) + + entries.sort(key=sort_key) + return entries[:k] + + def _compute_dtype(dtype: torch.dtype) -> torch.dtype: """Return a dtype supported by stable SVD on common PyTorch backends.""" return torch.float64 if dtype == torch.float64 else torch.float32 @@ -304,3 +417,269 @@ def random_projection_kernel_moments(ambient_dim: int, rank: int) -> RandomSubsp mean=float(mean), variance=float(variance), ) + + +def _pairwise_projection_kernel( + source_bases: torch.Tensor, + target_bases: torch.Tensor, + *, + max_temp_bytes: int = _PAIRWISE_TEMP_BYTES, +) -> torch.Tensor: + """Compute pairwise PK scores while bounding the overlap-tensor allocation.""" + if source_bases.ndim != 3 or target_bases.ndim != 3: + raise ValueError("source_bases and target_bases must be three-dimensional") + if source_bases.shape[1] != target_bases.shape[1]: + raise ValueError("source_bases and target_bases must share an ambient dimension") + if source_bases.device != target_bases.device or source_bases.dtype != target_bases.dtype: + raise ValueError("source_bases and target_bases must share a device and dtype") + if ( + isinstance(max_temp_bytes, bool) + or not isinstance(max_temp_bytes, int) + or max_temp_bytes < 1 + ): + raise ValueError("max_temp_bytes must be a positive integer") + + source_count, _, source_rank = source_bases.shape + target_count, _, target_rank = target_bases.shape + scores = torch.empty( + source_count, target_count, dtype=source_bases.dtype, device=source_bases.device + ) + bytes_per_overlap = source_rank * target_rank * source_bases.element_size() + pair_capacity = max(1, max_temp_bytes // bytes_per_overlap) + source_tile = max(1, min(source_count, pair_capacity // max(1, target_count))) + target_tile = max(1, min(target_count, pair_capacity // source_tile)) + + for source_start in range(0, source_count, source_tile): + source_stop = min(source_start + source_tile, source_count) + for target_start in range(0, target_count, target_tile): + target_stop = min(target_start + target_tile, target_count) + overlap = torch.einsum( + "adr,bds->abrs", + source_bases[source_start:source_stop], + target_bases[target_start:target_stop], + ) + scores[source_start:source_stop, target_start:target_stop] = overlap.square().sum( + dim=(-2, -1) + ) + return scores + + +def _architecture_name(model: Any) -> str: + cfg = getattr(model, "cfg", None) + architecture = getattr(cfg, "original_architecture", None) + return str(architecture) if architecture is not None else type(model).__name__ + + +def _read_role_matrix(model: Any, block: Any, layer: int, role: AttentionRole) -> torch.Tensor: + attn = getattr(block, "attn", None) + if attn is None: + raise ValueError(f"attention block {layer} does not expose an attn component") + attribute = f"W_{role}" + try: + matrix = getattr(attn, attribute) + except (AttributeError, RuntimeError, ValueError) as error: + raise ValueError( + f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}" + ) from error + if not isinstance(matrix, torch.Tensor): + raise ValueError( + f"role {role} at layer {layer} must be a tensor, got {type(matrix).__name__}" + ) + if matrix.ndim != 3: + raise ValueError( + f"role {role} at layer {layer} expected a three-dimensional per-head weight, " + f"got shape {tuple(matrix.shape)}" + ) + if role == "O": + matrix = matrix.transpose(-2, -1) + if not torch.is_floating_point(matrix): + raise ValueError(f"role {role} at layer {layer} must have a floating-point dtype") + if not bool(torch.isfinite(matrix).all()): + raise ValueError(f"role {role} at layer {layer} must contain only finite values") + return matrix.detach() + + +def _validate_role_shapes( + matrices: Sequence[torch.Tensor], layers: Sequence[int], role: AttentionRole +) -> Tuple[int, int, int]: + expected = (matrices[0].shape[0], matrices[0].shape[1], matrices[0].shape[2]) + if min(expected) < 1: + raise ValueError( + f"role {role} at layer {layers[0]} must have non-empty head, d_model, and width " + f"dimensions, got shape {expected}" + ) + for matrix, layer in zip(matrices[1:], layers[1:]): + if tuple(matrix.shape) != expected: + raise ValueError( + f"role {role} at layer {layer} has shape {tuple(matrix.shape)}, expected " + f"the consistent [heads, d_model, width] shape {expected}" + ) + return expected + + +def _extract_bases( + matrices: Sequence[torch.Tensor], + layers: Sequence[int], + role: AttentionRole, + *, + selected_rank: int, + rtol: float, + dtype: torch.dtype, + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor]: + layer_bases: List[torch.Tensor] = [] + layer_ranks: List[List[int]] = [] + for matrix, layer in zip(matrices, layers): + head_bases: List[torch.Tensor] = [] + head_ranks: List[int] = [] + for head in range(matrix.shape[0]): + try: + subspace = orthonormal_subspace( + matrix[head].to(dtype=dtype), rank=selected_rank, rtol=rtol + ) + except ValueError as error: + raise ValueError( + f"Could not extract role {role} at layer {layer}, head {head}: {error}" + ) from error + head_bases.append(subspace.basis.to(device=device)) + head_ranks.append(subspace.measured_rank) + layer_bases.append(torch.stack(head_bases)) + layer_ranks.append(head_ranks) + return torch.stack(layer_bases), torch.tensor(layer_ranks, dtype=torch.long, device=device) + + +def attention_head_subspace_affinity( + model: Any, + *, + source_role: str = "O", + target_role: str, + layer_order: str = "forward", + rank: Optional[int] = None, + rtol: Optional[float] = None, +) -> HeadAffinityResult: + """Compute OQ, OK, or OV Projection Kernel affinities for a TransformerBridge. + + K/V axes preserve native key-value heads on grouped-query attention models; + they are never expanded to query-head count. Hybrid models include only + attention blocks and report their original block indices. + + Args: + model: A TransformerBridge exposing readable per-head attention weights. + source_role: Source role; v1 supports only ``"O"``. + target_role: One of ``"Q"``, ``"K"``, or ``"V"``. + layer_order: ``"forward"`` keeps strict earlier-to-later pairs; ``"all"`` + keeps every pair. + rank: Optional common truncation rank. By default every head must be full + column rank. + rtol: Optional relative numerical-rank tolerance. + + Returns: + Affinity tensors, validity mask, original layer indices, and rank metadata. + """ + if source_role != "O": + raise ValueError(f"source_role must be 'O' in v1, got {source_role!r}") + if target_role not in ("Q", "K", "V"): + raise ValueError(f"target_role must be one of ['Q', 'K', 'V'], got {target_role!r}") + if layer_order not in ("forward", "all"): + raise ValueError(f"layer_order must be one of ['forward', 'all'], got {layer_order!r}") + validated_source_role = cast(AttentionRole, source_role) + validated_target_role = cast(AttentionRole, target_role) + validated_layer_order = cast(LayerOrder, layer_order) + blocks_with = getattr(model, "blocks_with", None) + if not callable(blocks_with): + raise ValueError("model must be a TransformerBridge exposing blocks_with('attn')") + attention_blocks = list(blocks_with("attn")) + if not attention_blocks: + raise ValueError("No attention layers found — cannot compute head subspace affinity.") + + layer_indices = [int(layer) for layer, _ in attention_blocks] + source_matrices = [ + _read_role_matrix(model, block, layer, validated_source_role) + for layer, block in attention_blocks + ] + target_matrices = [ + _read_role_matrix(model, block, layer, validated_target_role) + for layer, block in attention_blocks + ] + source_heads, source_ambient, source_width = _validate_role_shapes( + source_matrices, layer_indices, validated_source_role + ) + target_heads, target_ambient, target_width = _validate_role_shapes( + target_matrices, layer_indices, validated_target_role + ) + if source_ambient != target_ambient: + raise ValueError( + f"roles {validated_source_role} and {validated_target_role} must share d_model, got " + f"{source_ambient} and {target_ambient}" + ) + + dtype = source_matrices[0].dtype + for matrix in source_matrices[1:] + target_matrices: + dtype = torch.promote_types(dtype, matrix.dtype) + dtype = _compute_dtype(dtype) + effective_rtol = _validate_rtol(rtol, (source_ambient, max(source_width, target_width)), dtype) + requested_rank = _validate_rank(rank, min(source_ambient, source_width, target_width)) + source_rank = source_width if requested_rank is None else requested_rank + target_rank = target_width if requested_rank is None else requested_rank + + cfg = getattr(model, "cfg", None) + configured_device = getattr(cfg, "device", None) + result_device = ( + torch.device(configured_device) + if configured_device is not None + else source_matrices[0].device + ) + source_bases, source_ranks = _extract_bases( + source_matrices, + layer_indices, + validated_source_role, + selected_rank=source_rank, + rtol=effective_rtol, + dtype=dtype, + device=result_device, + ) + target_bases, target_ranks = _extract_bases( + target_matrices, + layer_indices, + validated_target_role, + selected_rank=target_rank, + rtol=effective_rtol, + dtype=dtype, + device=result_device, + ) + + layer_count = len(layer_indices) + flat_source = source_bases.reshape(layer_count * source_heads, source_ambient, source_rank) + flat_target = target_bases.reshape(layer_count * target_heads, target_ambient, target_rank) + scores = _pairwise_projection_kernel(flat_source, flat_target).reshape( + layer_count, source_heads, layer_count, target_heads + ) + scores = _clamp_projection_scores(scores, min(source_rank, target_rank)) + normalized = scores / math.sqrt(source_rank * target_rank) + if validated_layer_order == "forward": + layer_tensor = torch.tensor(layer_indices, device=result_device) + layer_mask = layer_tensor[:, None] < layer_tensor[None, :] + valid_mask = layer_mask[:, None, :, None].expand_as(scores) + else: + valid_mask = torch.ones_like(scores, dtype=torch.bool) + scores = torch.where(valid_mask, scores, torch.zeros_like(scores)) + normalized = torch.where(valid_mask, normalized, torch.zeros_like(normalized)) + + target_kind: HeadKind = "query" if validated_target_role == "Q" else "kv" + return HeadAffinityResult( + scores=scores, + normalized=normalized, + valid_mask=valid_mask, + source_role=validated_source_role, + target_role=validated_target_role, + source_layer_indices=tuple(layer_indices), + target_layer_indices=tuple(layer_indices), + source_head_kind="query", + target_head_kind=target_kind, + source_ranks=source_ranks, + target_ranks=target_ranks, + source_rank=source_rank, + target_rank=target_rank, + rank=rank, + rtol=effective_rtol, + )