[Proposal] Projection-kernel affinity for attention-head weight subspaces

[janmenjayap]

Proposal

Add a TransformerBridge-first tool for measuring overlap between the residual-stream
subspaces read or written by attention-head weights. For orthonormal bases U and V,

PK(S, T) = ||U.T @ V||_F^2 = sum_i cos(theta_i)^2,

where theta_i are principal angles. The Projection Kernel (PK) ignores singular-value
magnitude and is invariant to basis changes inside either subspace. It complements
Composition Score (CS): PK measures overlapping support, while CS measures scale-sensitive
composition of full linear maps under an isotropic-input assumption.

The work should separate:

1. a small, maintained numerical and TransformerBridge API; and
2. a versioned research artifact for GPT-2 labels, baselines, plots, reproduction, and
   causal validation.

Motivation

TransformerLens exposes Q/K/V/O weights and implements Q-, K-, and V-composition, but it
has no principal-angle or subspace-overlap analysis. Several results motivate this gap:

• Elhage et al. introduced weight composition for head-to-head paths.
• Merullo et al. found causally important one- to three-dimensional communication
  channels by decomposing QK and OV matrices.
• Yamagiwa, Takase, and Shimodaira report that PK over individual Q/K/V/O subspaces more
  clearly recovers several known GPT-2-small relationships.
• MAPS shows that static weight scores can predict head outputs and causal ablations, but
  also that heads are multifunctional and context dependent.

The PK paper studies one model, does not directly test whether PK predicts path-specific
causal effects, and explicitly says PK does not identify head function. Work on
representation similarity likewise warns that geometric and functional similarity can
dissociate. TransformerLens can provide the primitive and test its limits with existing
hooks and direct path patching.

Pitch

Stable core API

Names are open to maintainer preference. A minimal model API could be:

from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis import attention_head_subspace_affinity

bridge = TransformerBridge.boot_transformers("gpt2", device="cpu")
result = attention_head_subspace_affinity(
    bridge, source_role="O", target_role="Q", layer_order="forward"
)

result.scores       # [source_layer, source_head, target_layer, target_head]
result.normalized   # PK / d_head for equal-rank heads
result.valid_mask
result.top_pairs(20)

The numerical layer should remain model independent:

• orthonormal_subspace(matrix, *, rank=None, rtol=None) extracts a documented
  numerical column-space basis and measured rank;
• projection_kernel(basis_a, basis_b) returns raw PK, projector-normalized PK, and
  principal-angle cosines.

Explicit basis extraction avoids hiding rank decisions inside each score. For ranks r
and s, raw PK lies in [0, min(r, s)]. Projector normalization is
PK / sqrt(r*s), the cosine between projection matrices; at equal rank this is the
paper's PK / d_head.

The first model wrapper should expose OQ, OK, and OV without assuming equal source and
target head counts.

Weight-layout and architecture contract

The paper uses left-multiplication matrices. TransformerLens basis-generating matrices are:

Role	Matrix	Per-head shape
Q	W_Q[layer, head]	[d_model, d_head]
K	W_K[layer, kv_head]	[d_model, d_head]
V	W_V[layer, kv_head]	[d_model, d_head]
O	W_O[layer, head].T	[d_model, d_head]

The O transpose is essential. Results carry original layer indices and separate
source/target head metadata.

For grouped-query attention, K/V retain native KV-head identity rather than being silently
repeated to query-head labels. OQ has query-head × query-head axes; OK and OV have
query-head × KV-head axes. An expanded view can be added later for comparison with CS,
but duplicated K/V views must not be treated as independent heads.

Hybrid models include only attention layers and preserve original block indices.
Architectures without standard readable Q/K/V/O projections, including MLA and opaque
native-forward adapters, raise a precise unsupported-architecture error.

Numerical contract

• Validate floating, finite, two-dimensional matrices with a shared ambient dimension.
• Use reduced SVD for rank measurement and basis extraction; never materialize a
  [d_model, d_model] projector.
• Use a documented dtype- and shape-aware rank tolerance; permit caller-pinned rank.
• Clamp overlap singular values only for roundoff before acos.
• Preserve device and document accumulation/return dtype for low-precision inputs.
• Precompute each head basis once before batched pair scoring.

For independent random m-planes in R^d, the paper gives:

• mean m^2 / d;
• variance 2*m^2*(d-m)^2 / (d^2*(d-1)*(d+2)).

Expose these, if useful, only as descriptive moments. The normal approximation is not an
automatic p-value; trained head pairs are dependent and anisotropic.

Research reproduction and extension

The research artifact should pin model revision, processing mode, dtype, seed, pair
construction, and annotation data. It should not be imported by the library.

Study A — exact reproduction

Using raw GPT-2-small weights and the paper's 36 annotated heads across seven classes,
reproduce:

• 9,504 earlier-to-later pairs for each of OQ, OK, and OV;
• top-20 wiring diagrams and named edges;
• Table 1 PR-AUC against CS, Simple-CS, Linear CKA, and Procrustes;
• 10,296 non-identical layer-ordered QQ/KK/VV/OO pairs and Table 2;
• inlet/outlet rankings, especially L4H7;
• the random-subspace moment simulation; and
• Appendix K's raw-versus-processed comparison.

Table 1 targets are:

Method	OQ	OK	OV	Mean
PK	0.446	0.451	0.289	0.395
CS	0.227	0.352	0.215	0.265
Simple-CS	0.245	0.359	0.305	0.303
Linear CKA	0.337	0.323	0.409	0.356
Procrustes	0.328	0.269	0.428	0.342

This is not a blanket “PK wins” target: PK leads OQ and OK, while CKA and Procrustes beat
it on OV. Table 2 reports mean PR-AUC 0.047 and ROC-AUC 0.809 for PK, versus 0.038 and
0.787 for the strongest reported PR-AUC baseline, Procrustes.

Study B — causal-validity extension

Test:

Does static O→Q/K/V subspace overlap predict path-specific causal communication on
held-out prompts better than scale-sensitive or generic weight-similarity baselines?

Use clean/corrupted paired prompts and direct path patching from source-head output to a
target head's Q, K, or V input in at least two domains:

1. IOI for inhibition/name-mover and duplicate/inhibition relationships;
2. repeated-token induction for previous-token/induction relationships.

For every valid pair, record normalized absolute logit-difference recovery and uncertainty
across prompts. Evaluate Spearman correlation, top-k enrichment, PR-AUC at a preregistered
effect threshold, bootstrap confidence intervals, and paired bootstrap metric differences.

Baselines are PK, CS, Simple-CS, Linear CKA, Procrustes, source-head effect, target-head
effect, layer distance, weight norm, and seeded random ranking. Endpoint effects are
essential controls: otherwise a metric may merely rank globally important heads.

Primary hypotheses:

• H1: raw-weight GPT-2 results reproduce within declared numerical/ranking tolerances.
• H2: PK adds held-out predictive value beyond endpoint effects and layer distance for
  OQ and OK.
• H3: validity differs by role pairing; PK is not expected to dominate OV.
• H4: PK rankings are more stable than CS under behavior-preserving TransformerLens
  preprocessing.

Falsifiers include failure to reproduce Table 1, no held-out association after controls,
or equal performance from random norm-matched subspaces. Report negative results rather
than redefining the target.

Staged delivery

PR 1 — numerical primitive

• basis extraction, PK, principal angles, normalization, and null moments;
• model-free unit tests and API docs;
• no model wrapper, plotting, labels, or baseline bundle.

PR 2 — TransformerBridge wrapper

• OQ/OK/OV extraction, native MHA/GQA metadata, hybrid handling, masks, and ranking;
• synthetic tests plus cached-model Bridge integration tests;
• explicit unsupported-layout errors.

Research artifact/report

• exact GPT-2 reproduction and causal-validity study;
• baselines local to the artifact unless maintainers request promotion;
• machine-readable results and limitations.

Correctness and test plan

1. Known-angle constructions recover exact principal-angle cosines.
2. Cosine, overlap-Frobenius, and projector definitions agree on small fixtures.
3. Scores are invariant to orthogonal rebasing and nonsingular right transforms before
   basis extraction.
4. Identical, nested, orthogonal, unequal-rank, and deficient spaces have explicit values.
5. Seeded Monte Carlo moments match analytic formulas within sampling error.
6. Synthetic attention weights prove Q/K/V and W_O.T orientation.
7. MHA, native GQA, hybrid indices, masks, and tie ordering are pinned.
8. Selected Bridge scores match independent Torch calculations.
9. Unit tests load no models; integration tests reuse read-only cached Bridge fixtures.
10. Paper headline values live in the versioned reproduction, not brittle library CI.

Risks and interpretation

• Similarity is not causality: high PK only proposes a candidate relationship.
• Full head spaces may be too coarse: narrow causal channels can be diluted by
  rank-d_head PK; truncated variants are follow-ups, not post-hoc fixes.
• Subspace interventions can mislead: causal work needs controls for dormant-path
  activation described by Makelov et al.
• Preprocessing changes the estimand: raw Bridge weights match the primary paper.
• Rank is a choice: tolerance, spectrum, and dtype must be recorded.
• GQA creates pseudoreplication: repeated K/V views are not independent evidence.
• The random null is idealized: it is not calibrated significance for trained heads.

Alternatives

• Keep CS: retain it as a baseline; it measures scale-sensitive composition.
• Use decomposed CS: valuable for narrow high-gain channels, but not basis-invariant
  whole-subspace overlap.
• Use CKA/Procrustes: important alternatives that outperform PK in some paper settings.
• Flatten and cosine: depends on arbitrary head-basis choices.
• Use activation similarity: complementary and data conditional, not the same estimand.
• Keep this notebook-only: simpler, but provides no tested cross-architecture primitive.

Additional context

• Primary paper: Hiroaki Yamagiwa, Yusuke Takase, and Hidetoshi Shimodaira,
  “Measuring Affinity between Attention-Head Weight Subspaces via the Projection Kernel,”
  arXiv:2601.10266v1 (2026): https://arxiv.org/abs/2601.10266.
• https://github.com/ymgw55/proj-kernel-attn-subspace contained only a README and figure
  images, with no source or license, when checked on 2026-08-24.
• Related but non-duplicate TransformerLens work includes merged Composition Score PR Composition Scores #39,
  SVDInterpreter issue [Proposal] Implement SVD Decomposition of OV Circuits, Win, Wout #256 / merged PR Add SVD interpretability feature to TransformerLens #311, SVD-semantics issue [Proposal] Update and discuss behavior of SVD #341, Bridge analysis
  PR Added additional bridge analysis tools #1237, GQA circuit fix PR fix(bridge): expand grouped K/V heads in QK/OV and composition circuits on GQA models #1593, and gated-Qwen analysis fix PR Fix gated Qwen W_Q analysis weights #1653. None implements
  PK, principal angles, or projector-overlap head affinity.
• Bridge GQA weight-circuit fix PR fix(bridge): expand grouped K/V heads in QK/OV and composition circuits on GQA models #1593 is merged. This proposal still preserves native
  KV-head identity rather than counting expanded views independently.
• No new runtime dependency is required.

Checklist

• I have checked that there is no similar issue in the repo (required).
• No matching implementation or pull request was found on 2026-08-24.
• Orientation, pair counts, benchmarks, annotations, and null moments were checked.
• Stable API and paper reproduction are separated.
• The proposal includes falsifiable causal validation.
• No dependency or external code artifact is required.

Status: implemented in #1721

[janmenjayap]

Thank you, @jlarson4, for taking a look.

Implementation status

PR #1721 implements both proposed stages in two reviewable commits:

1. 69653575 — Add projection-kernel subspace geometry
2. ec057502 — Add Bridge attention-head subspace affinity

The PR is open against dev, is no longer a draft, and GitHub reports it as clean and mergeable.
It has seven changed files, 1,408 additions, and no deletions. CI currently reports 22 successful
checks, three conditionally skipped checks, and no failures or pending checks. The detailed local
validation results and limitations are recorded in the PR description.

#1720 originally proposed delivering the numerical layer and Bridge integration as separate PRs.
They are combined in #1721 as separate commits. When convenient, could you please let me know
whether reviewing the combined PR works for you? I would also be happy to restore the original
two-PR split if that would make the review easier.