[Proposal] Validate, generalize, and stress-test Projection-Kernel Head Affinity (follow-up to #1721)

[janmenjayap]

---

Proposal

#1721 shipped the correctness spine of Projection-Kernel (PK) head affinity — a data-free,
principal-angle subspace metric for comparing attention heads — as
attention_head_subspace_affinity(model, *, source_role="O", target_role, layer_order, rank, rtol)
(transformer_lens/tools/analysis/projection_kernel.py:591). Its own commit plan explicitly deferred
three pieces of follow-up work to "a separate issue, post-merge." This issue proposes all three as
independently-mergeable PRs, bundled under one tracking issue so a maintainer can sign off on the whole
slice in one pass rather than reviewing four small proposals separately:

1. Validate PK against the existing baseline. Add a checked-in, reproducible comparison of PK
   affinity (raw and min(p,q)-normalized) against composition_scores as classifiers of "known head
   cluster vs. not," reported as AUC across Q/K/V target roles on gpt2-small.
2. Generalize beyond source_role="O". Lift the current restriction so arbitrary
   (source_role, target_role) ∈ {Q,K,V,O}² pairs are supported, plus an optional OV-composite
   subspace extractor.
3. Stability report. Measure device/dtype/rtol sensitivity of the affinity matrices at
   integration scale (full gpt2-small, not synthetic unit-test matrices), and add regression tests
   pinned to the resulting tolerances.

Each is independently shippable and none blocks another, but all three gate on the same approval, so
they are proposed together.

---

Motivation

composition_scores (transformer_lens/utilities/matrix.py:16, surfaced as
TransformerBridge.all_composition_scores at transformer_lens/model_bridge/bridge.py:1955 and
HookedTransformer.all_composition_scores at transformer_lens/HookedTransformer.py:2694) is magnitude-
sensitive and TransformerLens's only existing tool for this kind of head-to-head comparison. PK is meant
to upgrade it, but #1721 shipped only the metric itself — no report demonstrating that PK actually
recovers a known head cluster better than composition_scores (e.g. the source paper's headline
gpt2-small finding that L4H7 acts as a hub/identity head, arXiv:2601.10266). Without that validation,
PK is an unvalidated primitive sitting next to the tool it's supposed to improve on.

Separately, source_role is currently hardcoded to "O" only
(projection_kernel.py:619), and the OV-composite subspace was named in the original proposal's design
section as a "TL extension beyond the paper" but explicitly scoped out of the merged PR. Generalizing
this is needed before PK can answer questions like "how does one head's output subspace couple with
another head's query subspace" for arbitrary role pairs.

Finally, the module's device/dtype/rtol handling (_compute_dtype at
projection_kernel.py:181, _rank_tolerance_dtype at :186, _validate_rtol at :191) is unit-tested
on small synthetic matrices only. #1721's own risk log flagged "SVD sign / device sensitivity" as a
real risk, but there is no integration-scale reproducibility report on real gpt2-small weights to back
it up.

Note: the GQA bridge-composition bug (#1553) and its fix (#1593) are both merged/closed, so none of
this work is blocked on grouped-query-attention support.

---

Pitch

1. IOI-cluster AUC report. Add a comparator that scores PK against composition_scores as binary
classifiers of "same known cluster vs. not." HeadAffinityResult.top_pairs(...)
(projection_kernel.py:143) already ranks pairs, but nothing compares that ranking against
all_composition_scores. Ground truth comes from two sources already available in this repository or
its cited literature rather than newly-derived labels: the source paper's L4H7 hub/identity-head finding
(arXiv:2601.10266), and the IOI circuit's name-mover heads (9,9), (9,6), (10,0) and S-inhibition
heads (7,3), (7,9), (8,6), (8,10) (Wang et al. 2022, arXiv:2211.00593), already cited in-repo in
demos/direct_path_patching_ioi.ipynb. One caveat worth flagging up front: head_detector.detect_head(...)
(transformer_lens/head_detector.py:40) only recognizes previous_token_head / duplicate_token_head /
induction_head patterns, none of which name "name-mover" or "S-inhibition," so that particular
cross-check only applies partially.

2. Cross-subspace coupling. Remove the source_role != "O" rejection
(projection_kernel.py:619) so any (source_role, target_role) ∈ {Q,K,V,O}² pair is supported, and add
an optional OV-composite output subspace (FactoredMatrix(W_V[h], W_O[h]), reusing the factored-SVD
machinery already in transformer_lens/FactoredMatrix.py:214 so nothing d_model × d_model is
materialized), gated behind an explicit "OV" role name rather than silently changing what "O" means.
One thing worth flagging to a maintainer before work starts: target_head_kind
(projection_kernel.py:712) currently hardcodes "kv" for anything that isn't "Q". That's correct
today because the only non-Q target roles are K/V (GQA-grouped), but once O and OV become valid
targets this needs to become "kv" if role in ("K", "V") else "query" — O and OV are per-full-
attention-head (query-head cardinality), not GQA-grouped.

3. Stability report. A checked-in report plus regression tests covering how much
attention_head_subspace_affinity(...) output varies across devices (CPU/CUDA/MPS), dtypes
(float32/float64/bfloat16/float16), and rtol choices, at integration scale. One concrete thing
to check, surfaced by reading the shipped module rather than a generic "test more devices" ask:
_singular_values (projection_kernel.py:370) routes through _linalg_on_cpu_if_mps (imported from
jacobian_lens_decomposition.py) for the principal-angle SVD, but orthonormal_subspace's own
torch.linalg.svd call (projection_kernel.py:255) — the per-head basis-extraction step that runs on
real per-head weight matrices — does not go through that wrapper. There is no existing MPS-specific test
for this module (grep -i mps tests/unit/tools/test_projection_kernel.py returns nothing). Whether this
is an oversight or intentional (basis extraction may be numerically safer on MPS than the cosine-overlap
SVD) is exactly the kind of question this report should answer with data.

---

Alternatives

File three separate issues instead of one. Rejected: all three gate on the same approval and are
each small/moderate in scope, so three separate issues would just mean the maintainer reviewing near-
identical context three times. #1539 (J-lens) used the single-tracking-issue-for-several-PRs pattern
before its own follow-up PRs were drafted, so this follows existing precedent in this repo.

Skip validation and generalization, ship only the stability report. Rejected: the AUC report and the
source_role generalization are both small, and leaving PK unvalidated against composition_scores
indefinitely undercuts the reason PK was proposed in the first place.

---

Additional context

Non-goals (this issue and all three PRs). Everything scoped out of #1721 for later, separate
issues: automatic clustering (connected_groups, spectral/agglomerative grouping, auto-labeling against
head_detector); cross-model/cross-checkpoint affinity; additional/larger models and a scale study
(gpt2-medium/gpt2-large, GQA models); an activation-based behavioral head-similarity cross-check; an
affinity-matrix registry/cache. Also out of scope for the cross-subspace-coupling piece: any new "QK"
composite type — this issue reads "a.OV against b.QK" from the original proposal as pairing the new
OV composite against existing single Q/K roles, not as a request for a second new composite type
(see open question below if that reading is wrong).

Open questions for the maintainer:

1. Approve filing this as scoped above, or reprioritize/cut one of the three pieces, before work starts
   on any of them.
2. Is citing the existing IOI head labels already present in demos/direct_path_patching_ioi.ipynb
   (rather than deriving a fresh positive set) acceptable for the AUC report, given that the
   head_detector cross-check only partially applies?
3. Confirm the reading that "a.OV against b.QK" means pairing the new OV composite against
   existing single Q/K roles, not a request for a literal new QK-composite role — if a literal QK
   composite is wanted, that changes the cross-subspace-coupling scope materially.
4. Should (source_role, target_role) be validated generically the same way for all 16 role-pair
   combinations, or should some combination remain restricted (e.g. is O↔O meaningful, or should
   same-role-as-O be rejected as degenerate)?
5. What tolerance bounds should the stability report's regression test pin for cross-device max/mean
   absolute difference in scores/normalized? This needs a first empirical run to propose sensible
   numbers, not a number picked in advance — will report back with data before pinning thresholds.

---

Checklist

• I have checked that there is no similar issue in the repo (required)

---

Reviewer note

Status: exploratory proposal, no code changes.
Suggested labels: enhancement, complexity-moderate, TransformerBridge

cc @jlarson4 — could you take a look and let me know if this scope and API shape look right?
If this gets a green light, I'd like to pick it up myself. Could you assign this issue to me once it's approved?
Thank you

---

[jlarson4]

Thanks for the detailed writeup as always @janmenjayap, and for catching the target_head_kind issue. I'd take a smaller slice of this.

Validation: yes, as a reproduction. The paper already compares PK and Composition Score on GPT-2 (Table 1), so this checks that our implementation matches it. Structure it the way the Jacobian Lens validation was built in #1539: a demo notebook, plus optionally a slow test.

• The test should assert relative structure (PK above CS on OQ/OK), not exact numbers. The paper has no reference code and leaves some protocol details unspecified.
• Use the paper's Table 8 labels (36 heads).
• Causal validation ([Proposal] Projection-kernel affinity for attention-head weight subspaces #1720's Study B) could land the same way.

Role generalization: same-role pairings only. Most of the proposed pairings add nothing new:

• PK is symmetric, so reversed pairings under layer_order="all" are just transposes.
• Under "forward", reversed pairings don't describe composition, because an earlier head can't read a later head's output.
• OV/QK composites span the same subspaces as the single roles whenever the other factor is full rank, and the wrapper already requires that.

Same-role similarity (QQ/KK/VV/OO, the paper's Table 2) is genuinely new. It needs:

• an unordered-distinct pair mode, since neither "forward" nor "all" produces Table 2's 10,296 pairs;
• head-kind labels on both axes: :721 hardcodes the source kind, alongside the :712 case you found.

Stability: one bug, no report. The default rank rtol is floored at the storage dtype's epsilon.

• Qwen2.5-0.5B loaded in bf16 fails the default O→Q call (requested rank 64 exceeds measured rank 20). Passing rtol=1e-4 works. That needs a fix and a regression test.
• Everything else I checked was stable. PK is invariant to basis sign and rotation, and GPT-2's top-20 pairs don't change across fp32, fp16, bf16 or MPS.

Your questions

1. See my requested scope below
2. No. Table 1 pools all 36 annotated heads as positives, so a 7-head subset can't be compared with it. Those heads do suit Study B.
3. Neither composite (see above).
4. O->{Q,K,V} plus same-role pairs. O<->O is meaningful. I'd leave the cross read-read pairs out for now.
5. No bounds needed for scores; the rtol bug wants a pass/fail test instead.

Suggested scope

• PR A: same-role pairings, the distinct-pair mode, head-kind labels on both axes.
• PR B: tensor top-k in top_pairs (the follow-up from feat(projection_kernel): add attention-head subspace affinity #1721 review), keeping the return type and tie order.
• PR C: a default rtol fix for half-precision checkpoints, with a regression test.
• Reproduction notebook, plus an optional slow test. Study B the same way if you pursue it.

Let me know if you have any questions!

[jlarson4]

I also do want to call out: you currently have 4 active proposals at various stages of development. I am going to leave this open to the rest of our contributors for the moment, and once we close out some of your other proposals we can come back around to this.

[Canonik]

I can take this one on!

[jlarson4]

@Canonik I am going to give @janmenjayap some time to wrap up his open PRs this week before I hand this off to someone else. I'll ping you if anything changes, or they can sign off on you taking it over if they're interested in working on this in a more advisory capacity.