Measuring Affinity between Attention-Head Weight Subspaces via the Projection Kernel

TL;DR

Measuring subspace overlap between attention-head weight matrices using the Projection Kernel (PK), defined as $\text{PK}(\mathcal{S},\mathcal{S}')=\sum_{i=1}^m \cos^2\theta_i = \|\mathbf{U}^{\top}\mathbf{U}'\|_{\mathrm{F}}^2$, PK provides a rotation-invariant, principal-angle-based view of head relationships. The study shows PK better reproduces IOI-related head interactions than the Composition Score (CS) on GPT2-small, and introduces a framework to compare PK distributions against random baselines via KL divergence. It identifies L4H7 as an Identity Head acting as a hub and demonstrates how projecting the unembedding matrix onto head subspaces yields interpretable, token-level insights. Overall, PK offers a robust, input-agnostic lens on global head structure that complements input-dependent mechanistic interpretability analyses.

Abstract

Understanding relationships between attention heads is essential for interpreting the internal structure of Transformers, yet existing metrics do not capture this structure well. We focus on the subspaces spanned by attention-head weight matrices and quantify head-to-head relationships using the Projection Kernel (PK), a principal-angle-based measure of subspace similarity. Experiments show that PK reproduces known head-to-head interactions on the IOI task more clearly than prior metrics such as the Composition Score. We further introduce a framework to quantify the informativeness of PK distributions by comparing them with a reference distribution derived from random orthogonal subspaces. As an application, we analyze a directed graph constructed from PK and show that, in GPT2-small, L4H7 acts as a hub by functioning as an identity head.

Figures (21)

• Figure 1: For the attention-head weights of GPT2-small, we compute similarities between heads using (top) the Projection Kernel (PK) and (bottom) the Composition Score (CS), and visualize the top 20 scores as edges for each weight pairing type (OQ, OK, and OV). Edge opacity corresponds to the score normalized by the maximum score. The seven head classes excluding Identity Heads follow wang2023interpretability. With PK, many heads from IOI-relevant head classes appear, and edges concentrate among these heads. In contrast, this tendency is weaker with CS; instead, many edges connect heads that do not belong to any head class.
• Figure 2: Illustration of measuring similarity using the Projection Kernel (PK) in (\ref{['eq:pk']}) between the $d_\text{head}$-dimensional subspace $\mathcal{S}^{(h, \text{O})}\subset\mathbb{R}^d$ spanned by the output weights $\bm{W}_\text{O}^h\in\mathbb{R}^{d\times d_\text{head}}$ of head $h$ and the $d_\text{head}$-dimensional subspace $\mathcal{S}^{(h', \text{Q})}\subset\mathbb{R}^d$ spanned by the query weights ${\bm{W}_\text{Q}^{h'}}^\top \in\mathbb{R}^{d\times d_\text{head}}$ of head $h'$. The layer indices of heads $h$ and $h'$ are denoted by $\ell_h$ and $\ell_h'$ with $\ell_h < \ell_h'$. A larger $\text{PK}(\mathcal{S}^{(h,\text{O})}, \mathcal{S}^{(h',\text{Q})})$ indicates that the subspaces share more directions with small principal angles in (\ref{['eq:pa-def']}), and thus that the output of head $h$ is more strongly incorporated into the query of head $h'$.
• Figure 3: Principal angles $\theta_1=0$ and $\theta_2$ between two planes $\mathcal{S},\mathcal{S}'\subset \mathbb{R}^3$ with $\dim \mathcal{S}=\dim \mathcal{S}' = 2$, and the corresponding principal vectors $\bm{s}_1,\bm{s}_2\in\mathcal{S}$ and $\bm{s}_1',\bm{s}_2'\in\mathcal{S}'$. (Inspired by mandolesi2021blade)
• Figure 4: For a two-layer attention-only Transformer, we visualize head-to-head similarity strengths using (a) PK and (b) CS. In the first layer, we show the top two Previous Token Heads, and in the second layer, the top two Induction Heads, and color them by their head scores (see Appendix \ref{['app:attn-only-2l']} for details).
• Figure 5: We compute PK and CS scores for OQ, OK, and OV and plot them as scatter plots. The Spearman's $\rho$ values over all pairs are $0.885$, $0.844$, and $0.866$, indicating strong correlations. In contrast, when we restrict to pairs that belong to the top $3\%$ in either metric (the $\blacklozenge$, $\blacksquare$, and $\blacktriangle$ symbols), the coefficients drop to $-0.251$, $-0.109$, and $0.186$, and the correlations become weak. In particular, for pairs that belong only to the top $3\%$ in CS (the $\blacktriangle$ symbol), many have small PK scores. For visibility, we plot values up to the $99.99$th percentile.