Attention heads forming orbits

One self-attention layer, penalized toward equivariance, learning modular addition

Each dot is one of the 32 heads of a single self-attention layer, drawn in a plane of its weight space while the layer learns modular addition. The experiment’s code is due to Georg Bökman.

1. Setup

model: single 32-head self-attention layer sandwiched between an embedding $W_E$ and an unembedding $W_U$
task: modular addition with $p = 24$ , $(x_1, x_2, x_3) \mapsto(x_2+x_3, x_3+x_1, x_1+x_2)$ . - symmetry: shifting all three inputs by $s\in C_{24}$ shifts every output by $2s$
loss: cross-entropy plus an equivariance loss
result: 100 % train accuracy, 76 % on the held-out inputs, equivariance error $10^{-4}$

2. Prediction

The structure theorem for attention (preprint) says that the layer is exactly equivariant if and only if its heads can be organized into orbits of the group $G = C_{24}$ of shifts. That is, the shift by $s$ sends head $h$ to a head $s \cdot h$ with $M_{s \cdot h} = T_s M_h T_s^{-1}, \qquad R_{s \cdot h} = T_{2s} R_h T_s^{-1},$ where

$M_h$ is the QK matrix of head $h$
$R_h$ is the OV matrix
$T_s$ is the cyclic shift of $\mathbb{R}^{24}$ by $s$ (heads with the same $M_h$ are merged, and only the sum of their $R_h$ enters).

The movie shows that a layer that is pressured towards equivariance arranges its heads in such orbits.

3. Movies

$\mathbb{R}^{24}$ decomposes into irreps of $C_{24}$ labeled by frequencies. For QK matrices living in $\mathbb{R}^{24\times 24}$ , we pick a pair of (query frequency $\omega_q$ , key frequency $\omega_k$ ) and look at a 2D subspace of $\omega_q\otimes \omega_k$ isomorphic to $\omega_q+\omega_k$ or $\omega_q-\omega_k$ .

top left: query frequency 0, key frequency 10. shift by $s$ rotates the dot by $10 s \cdot 2\pi/24$ .
bottom left: $R_h$ in output frequency 10 and input frequency 10. Shift by $s$ rotates a dot by $(2 \cdot 10 - 10)\, s \cdot 2\pi/24 = 10 s \cdot 2\pi/24$ (the factor 2 is the $2s$ by which the output shifts).
right column: QK and OV at frequency 6, where the rotation is by $6 s \cdot 2\pi/24$ .

What to look for: a $G$ -orbit appears in every panel as a regular polygon centered at the origin, with one vertex per value of the rotation. $10 s \bmod 24$ takes 12 values, so an orbit on frequency 10 is a 12-gon, and $6 s \bmod 24$ takes 4, so an orbit on frequency 6 is a square. The 12 blue heads and the 4 orange heads are such orbits. This was checked on the full matrices (a polygon in one plane says nothing about the other blocks of $M_h$ and $R_h$ ): after any of the 24 shifts, each of these heads lands on another head of its set, to within 2 % in relative Frobenius norm. The gray dots are the other 16 heads.

Technical notes

The axes of each panel rescale with the largest radius in view, so that every dot stays on screen.
The thin cross marks the origin, and each dot trails its last 250 epochs.
One frame per 25 epochs up to epoch 4000, then one per 100. The last frame is held before the loop restarts.