feat(projection_kernel): add attention-head subspace affinity

docs/source/content/projection_kernel.md

@@ -0,0 +1,139 @@
+# 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 and a TransformerBridge wrapper for comparing attention-head weight spaces.
+
+## 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.
+
+## 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$,
+`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).

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

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
+        )

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
+    )