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:

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) * compute_eps, where compute_eps is the epsilon of the dtype the SVD actually runs in. Because float16 and bfloat16 inputs are promoted to float32 before the SVD, the default is the same for every storage dtype that promotes to float32. Deriving the tolerance from the compute dtype rather than the storage dtype keeps the rank decision tied to the precision of the singular values actually computed; a storage-dtype floor would instead discard that precision and collapse the measured rank of a full-rank half-precision head. Concretely, a half-precision matrix with 4096 rows gets rtol = 4096 * float32_eps ≈ 4.9e-4 once promoted — tight enough that a column duplicated from another is still measured as structurally rank-deficient while an independent full-rank matrix keeps its full rank. 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:

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. 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 derives one shared tolerance for both roles from the promoted compute dtype, so half-precision checkpoints do not need a manual rtol override. An explicit rtol still takes precedence.

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¶

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.