From 69653575f94fd8a9a7b98ac479867a9408d97103 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 24 Aug 2026 18:00:56 +0530 Subject: [PATCH 1/4] 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/4] 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, + ) From f4801cbaee659af985f216eefb561a33cbc56470 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Thu, 27 Aug 2026 14:28:28 +0530 Subject: [PATCH 3/4] fix(projection_kernel): harden numerical behavior and MPS support - Add an MPS CPU fallback for principal-angle singular values while preserving result devices. - Make rank tolerance storage-aware without collapsing realistic FP16 and BF16 matrices, and reject material cosine-bound violations. - Preserve unsupported-weight errors, return independent contiguous masks, and expand numerical regression coverage. --- tests/mps/test_mps_basic.py | 30 +++++ tests/unit/tools/test_projection_kernel.py | 104 +++++++++++++++++- .../tools/analysis/projection_kernel.py | 59 ++++++++-- 3 files changed, 179 insertions(+), 14 deletions(-) diff --git a/tests/mps/test_mps_basic.py b/tests/mps/test_mps_basic.py index 63f6e6b5d..b25720bbf 100644 --- a/tests/mps/test_mps_basic.py +++ b/tests/mps/test_mps_basic.py @@ -18,6 +18,11 @@ import pytest import torch +from transformer_lens.tools.analysis.projection_kernel import ( + SubspaceBasis, + projection_kernel, +) + # Skip the entire module on non-MPS runners (Linux CI, CPU-only Macs) pytestmark = pytest.mark.skipif( not torch.backends.mps.is_available(), @@ -138,6 +143,31 @@ def test_mps_tensor_basic_operations(): _cleanup() +def test_mps_projection_kernel_principal_angles(): + """Projection Kernel computes principal angles and preserves the MPS device.""" + try: + basis = torch.eye(4, device="mps", dtype=torch.float32)[:, :2] + subspace = SubspaceBasis( + basis=basis, + singular_values=torch.ones(2, device="mps"), + rank=2, + measured_rank=2, + rtol=4 * torch.finfo(torch.float32).eps, + threshold=4 * torch.finfo(torch.float32).eps, + input_shape=(4, 2), + ) + + result = projection_kernel(subspace, subspace) + + assert result.score.device.type == "mps" + assert result.normalized.device.type == "mps" + assert result.cosines.device.type == "mps" + assert result.angles.device.type == "mps" + assert result.cosines.cpu().tolist() == pytest.approx([1.0, 1.0]) + finally: + _cleanup() + + def test_mps_softmax_and_layernorm(): """Softmax and LayerNorm — core transformer ops — work on MPS.""" x = torch.randn(4, 16, 64, device="mps", dtype=torch.float32) diff --git a/tests/unit/tools/test_projection_kernel.py b/tests/unit/tools/test_projection_kernel.py index 9ebef5c7c..5794dfec5 100644 --- a/tests/unit/tools/test_projection_kernel.py +++ b/tests/unit/tools/test_projection_kernel.py @@ -90,6 +90,30 @@ def test_low_precision_promotes_to_float32(self, dtype): assert result.basis.dtype == torch.float32 assert result.singular_values.dtype == torch.float32 + assert result.rtol == pytest.approx(torch.finfo(dtype).eps) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_low_precision_default_detects_rank_deficiency(self, dtype): + matrix = torch.tensor([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [4.0, 4.0]], dtype=dtype) + + result = orthonormal_subspace(matrix) + + assert result.measured_rank == 1 + assert result.rtol == pytest.approx(torch.finfo(dtype).eps) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_low_precision_default_handles_realistic_ambient_dimension(self, dtype): + generator = torch.Generator().manual_seed(7) + full_rank = torch.randn(4096, 8, generator=generator, dtype=torch.float64) + rank_deficient = full_rank.clone() + rank_deficient[:, -1] = 0.3 * full_rank[:, 0] + 0.7 * full_rank[:, 1] + + full_result = orthonormal_subspace(full_rank.to(dtype)) + deficient_result = orthonormal_subspace(rank_deficient.to(dtype)) + + assert full_result.measured_rank == 8 + assert deficient_result.measured_rank == 7 + assert full_result.rtol == pytest.approx(torch.finfo(dtype).eps) @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) def test_preserves_supported_dtype_and_cpu_device(self, dtype): @@ -165,6 +189,19 @@ def test_recovers_known_principal_angles(self): assert result.angles.tolist() == pytest.approx([0.0, math.pi / 3]) assert result.score.item() == pytest.approx(1.25) + def test_rank_deficient_input_has_exact_score_and_angle(self): + rank_deficient = torch.tensor([[1.0, 2.0], [0.0, 0.0], [0.0, 0.0]], dtype=torch.float64) + tilted = torch.tensor([[3.0], [4.0], [0.0]], dtype=torch.float64) + + result = projection_kernel( + orthonormal_subspace(rank_deficient), orthonormal_subspace(tilted) + ) + + assert result.rank_a == 1 + assert result.score.item() == pytest.approx(0.36) + assert result.cosines.tolist() == pytest.approx([0.6]) + assert result.angles.tolist() == pytest.approx([math.acos(0.6)]) + def test_nested_unequal_rank_spaces(self): small = orthonormal_subspace(torch.eye(4)[:, :2]) large = orthonormal_subspace(torch.eye(4)[:, :3]) @@ -251,6 +288,8 @@ def test_clamps_tiny_score_overshoot(self): assert result.score.item() == 2.0 assert result.normalized.item() == 1.0 + assert result.cosines.tolist() == [1.0, 1.0] + assert result.angles.tolist() == [0.0, 0.0] def test_rejects_large_score_bound_violation(self): valid = orthonormal_subspace(torch.eye(2)) @@ -267,6 +306,21 @@ def test_rejects_large_score_bound_violation(self): with pytest.raises(ValueError, match="outside its theoretical bounds"): projection_kernel(invalid, invalid, check_orthonormal=False) + def test_rejects_large_principal_cosine_bound_violation(self): + valid = orthonormal_subspace(torch.eye(2)) + invalid = SubspaceBasis( + basis=torch.diag(torch.tensor([1.01, 0.1])), + 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="Principal-angle cosine"): + projection_kernel(valid, invalid, check_orthonormal=False) + def test_rejects_ambient_dimension_mismatch(self): first = orthonormal_subspace(torch.eye(3)) second = orthonormal_subspace(torch.eye(4)) @@ -293,8 +347,8 @@ 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 + def test_seeded_monte_carlo_matches_moments(self): + ambient_dim, rank, samples = 8, 2, 5000 generator = torch.Generator().manual_seed(17) first, _ = torch.linalg.qr( torch.randn(samples, ambient_dim, rank, generator=generator, dtype=torch.float64) @@ -303,11 +357,14 @@ def test_seeded_monte_carlo_matches_mean(self): 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() + scores = overlap.square().sum(dim=(-2, -1)) + empirical_mean = scores.mean().item() + empirical_variance = scores.var(unbiased=False).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) + assert empirical_variance == pytest.approx(reference.variance, rel=0.1) class TestPairwiseProjectionKernel: @@ -370,6 +427,17 @@ def test_all_layer_order_includes_every_pair(self): assert bool(result.valid_mask.all()) assert int(result.valid_mask.sum()) == 8 + def test_forward_valid_mask_is_contiguous_and_independently_writable(self): + result = attention_head_subspace_affinity(SyntheticBridge((0, 3)), target_role="Q") + + assert result.valid_mask.is_contiguous() + assert result.valid_mask[0, 0, 1, 0] + assert result.valid_mask[0, 1, 1, 0] + + result.valid_mask[0, 0, 1, 0] = False + + assert result.valid_mask[0, 1, 1, 0] + def test_invalid_entries_are_zero(self): result = attention_head_subspace_affinity(SyntheticBridge(), target_role="Q") @@ -405,6 +473,16 @@ def test_result_dtype_device_and_rank_metadata(self): assert result.source_rank == 2 assert result.target_rank == 2 + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_wrapper_uses_least_precise_storage_dtype_for_default_rtol(self, dtype): + model = SyntheticBridge((0, 1)) + for _, block in model.attention_blocks: + block.attn.W_Q = block.attn.W_Q.to(dtype=dtype) + + result = attention_head_subspace_affinity(model, target_role="Q") + + assert result.rtol == pytest.approx(torch.finfo(dtype).eps) + def test_rank_deficiency_names_role_layer_and_head(self): model = SyntheticBridge() deficient = model.attention_blocks[1][1].attn.W_Q[1] @@ -481,6 +559,26 @@ def test_rejects_missing_and_nonfinite_weights(self): with pytest.raises(ValueError, match="role K at layer 2 must contain only finite"): attention_head_subspace_affinity(nonfinite, target_role="K") + def test_preserves_not_implemented_error_from_weight_property(self): + model = SyntheticBridge() + original = model.attention_blocks[0][1].attn + + class UnsupportedAttention: + W_O = original.W_O + + @property + def W_K(self): + raise NotImplementedError("K weights are not supported") + + model.attention_blocks[0] = (0, SyntheticBlock(UnsupportedAttention())) + + with pytest.raises( + NotImplementedError, match="cannot expose role K at layer 0" + ) as exc_info: + attention_head_subspace_affinity(model, target_role="K") + + assert isinstance(exc_info.value.__cause__, NotImplementedError) + @pytest.mark.parametrize("shape", [(2, 5, 2), (1, 6, 2), (1, 5, 3)]) def test_rejects_inconsistent_role_shapes(self, shape): model = SyntheticBridge() diff --git a/transformer_lens/tools/analysis/projection_kernel.py b/transformer_lens/tools/analysis/projection_kernel.py index 8dca25f4b..e13e7f65c 100644 --- a/transformer_lens/tools/analysis/projection_kernel.py +++ b/transformer_lens/tools/analysis/projection_kernel.py @@ -174,9 +174,16 @@ def _compute_dtype(dtype: torch.dtype) -> torch.dtype: return torch.float64 if dtype == torch.float64 else torch.float32 +def _rank_tolerance_dtype(dtypes: Sequence[torch.dtype]) -> torch.dtype: + """Return the least precise storage dtype for a shared rank tolerance.""" + return max(dtypes, key=lambda dtype: torch.finfo(dtype).eps) + + 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 + compute_epsilon = torch.finfo(_compute_dtype(dtype)).eps + storage_epsilon = torch.finfo(dtype).eps + return max(max(shape) * compute_epsilon, storage_epsilon) 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) @@ -204,9 +211,10 @@ def orthonormal_subspace( """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. + explicit ``rtol``, numerical rank uses the larger of the compute-SVD error + scale and one input-storage epsilon, 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]``. @@ -233,7 +241,7 @@ def orthonormal_subspace( 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) + effective_rtol = _validate_rtol(rtol, input_shape, matrix.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 @@ -336,6 +344,27 @@ def _clamp_projection_scores(scores: torch.Tensor, upper_bound: int) -> torch.Te return scores.clamp(min=0.0, max=float(upper_bound)) +def _clamp_principal_cosines(cosines: torch.Tensor) -> torch.Tensor: + """Clamp roundoff-scale cosine violations and reject larger violations.""" + tolerance = 100 * torch.finfo(cosines.dtype).eps * max(1, cosines.numel()) + minimum = float(cosines.detach().min().item()) + maximum = float(cosines.detach().max().item()) + if minimum < -tolerance or maximum > 1 + tolerance: + raise ValueError( + "Principal-angle cosine lies outside its theoretical bounds: " + f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, 1] " + f"within tolerance {tolerance:.6g}" + ) + return cosines.clamp(min=0.0, max=1.0) + + +def _singular_values(matrix: torch.Tensor) -> torch.Tensor: + """Compute singular values with an explicit MPS CPU fallback.""" + if matrix.device.type == "mps": + return torch.linalg.svdvals(matrix.cpu()).to(device=matrix.device) + return torch.linalg.svdvals(matrix) + + def projection_kernel( subspace_a: SubspaceBasis, subspace_b: SubspaceBasis, @@ -379,8 +408,8 @@ def projection_kernel( 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)) + cosines = _clamp_principal_cosines(_singular_values(overlap)) + angles = torch.acos(cosines) denominator = math.sqrt(subspace_a.rank * subspace_b.rank) return ProjectionKernelResult( score=score, @@ -477,6 +506,10 @@ def _read_role_matrix(model: Any, block: Any, layer: int, role: AttentionRole) - attribute = f"W_{role}" try: matrix = getattr(attn, attribute) + except NotImplementedError as error: + raise NotImplementedError( + f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}" + ) from error except (AttributeError, RuntimeError, ValueError) as error: raise ValueError( f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}" @@ -613,11 +646,15 @@ def attention_head_subspace_affinity( f"{source_ambient} and {target_ambient}" ) - dtype = source_matrices[0].dtype - for matrix in source_matrices[1:] + target_matrices: + matrices = source_matrices + target_matrices + tolerance_dtype = _rank_tolerance_dtype([matrix.dtype for matrix in matrices]) + dtype = matrices[0].dtype + for matrix in matrices[1:]: dtype = torch.promote_types(dtype, matrix.dtype) dtype = _compute_dtype(dtype) - effective_rtol = _validate_rtol(rtol, (source_ambient, max(source_width, target_width)), dtype) + effective_rtol = _validate_rtol( + rtol, (source_ambient, max(source_width, target_width)), tolerance_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 @@ -659,7 +696,7 @@ def attention_head_subspace_affinity( 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) + valid_mask = layer_mask[:, None, :, None].expand_as(scores).contiguous() else: valid_mask = torch.ones_like(scores, dtype=torch.bool) scores = torch.where(valid_mask, scores, torch.zeros_like(scores)) From bf62a8e1c4352242cdb77e92e2784c457d4800b8 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Thu, 27 Aug 2026 22:39:25 +0530 Subject: [PATCH 4/4] fix(projection_kernel): clarify API contracts and scaling limits - Add jaxtyping shape annotations and runtime-validation coverage for public tensor interfaces. - Clarify measured ranks, retained basis widths, and storage-aware tolerance behavior. - Restore alphabetical exports and document memory usage and quadratic pair-ranking costs. --- docs/source/content/projection_kernel.md | 29 +++++++++++++++++-- tests/unit/tools/test_projection_kernel.py | 23 +++++++++++++-- transformer_lens/tools/analysis/__init__.py | 2 +- .../tools/analysis/projection_kernel.py | 29 +++++++++++-------- 4 files changed, 65 insertions(+), 18 deletions(-) diff --git a/docs/source/content/projection_kernel.md b/docs/source/content/projection_kernel.md index edfef6c4e..89b619f9b 100644 --- a/docs/source/content/projection_kernel.md +++ b/docs/source/content/projection_kernel.md @@ -49,8 +49,10 @@ 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. +`max(max(matrix.shape) * compute_eps, storage_eps)`. This accounts for both SVD roundoff +and input quantization without letting large low-precision matrices produce tolerances above +one. 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, @@ -111,6 +113,29 @@ The O transpose is required: `W_O` itself has shape `[d_head, d_model]`. 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. +`source_ranks` and `target_ranks` report each head's measured numerical rank before +truncation. Scalar `source_rank` and `target_rank` are the retained basis widths used for +their respective roles. +When `rtol` is omitted, the wrapper uses the least-precise participating storage dtype to +derive one shared tolerance for both roles. + +## Memory and scaling + +The Bridge wrapper materializes dense orthonormal basis stacks before pairwise scoring. +Approximate basis memory per role is + +`layers * heads * d_model * rank * element_size`. + +For 32 layers, 32 heads, `d_model=4096`, rank 128, and float32, this is about **2.15 GB +per role** before allocating the other role, score tensors, masks, temporary pairwise +products, or ranking objects. Use an explicit lower rank or CPU execution when appropriate, +and start with smaller models. + +`HeadAffinityResult.top_pairs()` currently creates and sorts one Python +`HeadAffinityPair` for every valid pair. On large head grids this can add substantial +memory and runtime even when only a small `k` is requested: pair enumeration scales as +`O((layers * heads)^2)`. Streaming or tiled basis extraction and tensor-only top-k ranking +are possible future optimizations; they are not implemented by the current API. ## Random-subspace reference diff --git a/tests/unit/tools/test_projection_kernel.py b/tests/unit/tools/test_projection_kernel.py index 5794dfec5..eaa35e5f4 100644 --- a/tests/unit/tools/test_projection_kernel.py +++ b/tests/unit/tools/test_projection_kernel.py @@ -5,6 +5,7 @@ import pytest import torch +from beartype.roar import BeartypeCallHintParamViolation from transformer_lens.tools.analysis.projection_kernel import ( SubspaceBasis, @@ -16,6 +17,12 @@ ) +def test_analysis_exports_are_alphabetized(): + from transformer_lens.tools import analysis + + assert analysis.__all__ == sorted(analysis.__all__) + + class SyntheticAttention: def __init__(self, seed: int, *, d_model: int = 5, d_head: int = 2): generator = torch.Generator().manual_seed(seed) @@ -123,12 +130,22 @@ def test_preserves_supported_dtype_and_cpu_device(self, dtype): assert result.singular_values.dtype == dtype assert result.basis.device == torch.device("cpu") + @pytest.mark.parametrize( + "matrix", + [ + [[1.0]], + torch.ones(3), + torch.ones(2, 2, dtype=torch.int64), + torch.ones(2, 2, dtype=torch.complex64), + ], + ) + def test_runtime_typecheck_rejects_invalid_matrices(self, matrix): + with pytest.raises(BeartypeCallHintParamViolation): + orthonormal_subspace(matrix) + @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"), diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index e6c3a02b2..16e6c752c 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -49,8 +49,8 @@ ) __all__ = [ - "DirectLogitAttribution", "AttentionHeadRef", + "DirectLogitAttribution", "HeadAffinityPair", "HeadAffinityResult", "JSpaceDecomposition", diff --git a/transformer_lens/tools/analysis/projection_kernel.py b/transformer_lens/tools/analysis/projection_kernel.py index e13e7f65c..cb92fb9e0 100644 --- a/transformer_lens/tools/analysis/projection_kernel.py +++ b/transformer_lens/tools/analysis/projection_kernel.py @@ -13,6 +13,7 @@ from typing import Any, List, Literal, Optional, Sequence, Tuple, cast import torch +from jaxtyping import Bool, Float, Int AttentionRole = Literal["Q", "K", "V", "O"] LayerOrder = Literal["forward", "all"] @@ -35,8 +36,8 @@ class SubspaceBasis: input_shape: Shape of the matrix from which the basis was extracted. """ - basis: torch.Tensor - singular_values: torch.Tensor + basis: Float[torch.Tensor, "ambient rank"] + singular_values: Float[torch.Tensor, "spectrum"] rank: int measured_rank: int rtol: float @@ -53,10 +54,10 @@ def ambient_dim(self) -> int: class ProjectionKernelResult: """Projection Kernel score and its principal-angle decomposition.""" - score: torch.Tensor - normalized: torch.Tensor - cosines: torch.Tensor - angles: torch.Tensor + score: Float[torch.Tensor, ""] + normalized: Float[torch.Tensor, ""] + cosines: Float[torch.Tensor, "principal_angle"] + angles: Float[torch.Tensor, "principal_angle"] rank_a: int rank_b: int ambient_dim: int @@ -104,19 +105,23 @@ class HeadAffinityResult: Score tensors have shape ``[source_layer, source_head, target_layer, target_head]``. Layer index tuples map tensor positions to original model block numbers. + + ``source_ranks`` and ``target_ranks`` are measured numerical ranks for each + head before optional truncation. Scalar ``source_rank`` and ``target_rank`` + are the retained basis widths used for their respective roles. """ - scores: torch.Tensor - normalized: torch.Tensor - valid_mask: torch.Tensor + scores: Float[torch.Tensor, "source_layer source_head target_layer target_head"] + normalized: Float[torch.Tensor, "source_layer source_head target_layer target_head"] + valid_mask: Bool[torch.Tensor, "source_layer source_head target_layer target_head"] 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_ranks: Int[torch.Tensor, "source_layer source_head"] + target_ranks: Int[torch.Tensor, "target_layer target_head"] source_rank: int target_rank: int rank: Optional[int] @@ -203,7 +208,7 @@ def _validate_rank(rank: Optional[int], max_rank: int) -> Optional[int]: def orthonormal_subspace( - matrix: torch.Tensor, + matrix: Float[torch.Tensor, "ambient width"], *, rank: Optional[int] = None, rtol: Optional[float] = None,