Profiles of attention heads

Interactive figures for On attention heads and bilinear forms.

Every attention head carries a bilinear form \(L = W_K^{\mathsf T} W_Q\) on the residual stream. Write it as \(L = S + T\) with \(S\) symmetric and \(T\) antisymmetric, and scale so that \(\lVert L\rVert = 1\). Let \(\lambda_+\) list the positive eigenvalues of \(S\) and \(\lambda_-\) the absolute values of its negative eigenvalues, both decreasing and padded with zeros to a common length. The profile \(\pi(L) = (a,b,c,d)\) is given by

\[ \begin{aligned} a &= \lVert T\rVert^2, \\ b &= 2\,\langle \lambda_+,\lambda_-\rangle, \\ c &= \lVert(\lambda_+ - \lambda_-)_+\rVert^2, \\ d &= \lVert(\lambda_+ - \lambda_-)_-\rVert^2, \end{aligned} \]

Here \(x_+ = \max(x,0)\) and \(x_- = \max(-x,0)\) for a real number \(x\); for vectors, these operations are applied entrywise. The four coordinates are nonnegative and sum to one, so the profile is a point of the \(3\)-simplex, drawn as a tetrahedron whose vertices are the antisymmetric forms (\(a = 1\)), the symmetric forms with spectrum symmetric about zero (\(b = 1\)), and the positive and negative semidefinite forms (\(c = 1\), \(d = 1\)). In the plots, \(H\) denotes a symmetric matrix with spectrum symmetric about zero.

model	layers	heads/layer	\(N\)	\(n\)	heads	positional
distilgpt2	6	12	768	64	72	learned absolute
gpt2	12	12	768	64	144	learned absolute
gpt2-medium	24	16	1024	64	384	learned absolute
gpt2-large	36	20	1280	64	720	learned absolute
opt-125m	12	12	768	64	144	learned absolute
opt-350m	24	16	1024	64	384	learned absolute
opt-1.3b	24	32	2048	64	768	learned absolute
opt-2.7b	32	32	2560	80	1,024	learned absolute
opt-6.7b	32	32	4096	128	1,024	learned absolute
opt-13b	40	40	5120	128	1,600	learned absolute
bloom-3b	30	32	2560	80	960	alibi
bloom-7b1	30	32	4096	128	960	alibi
pythia-70m-deduped	6	8	512	64	48	rope ^†
Mistral-Small-24B-Base	40	32	5120	128	1,280	rope ^†

^† With rotary embeddings the bilinear form varies with the distance between the two tokens; those pages carry a slider over the relative position \(d\), from 0 to 128. For the other models the form does not depend on position.

Code, data and the Lean formalizations are in the repository.