You are designing a G-equivariant attention layer that can express all equivariant functions expressible by an unconstrained attention layer, with as few parameters as possible.
We proved that in general this is not possible without increasing the head count or head dimension (you can try this and see why). That is, if you only constrain the vanilla attention's parameters, it is not possible to cover all possible ways the vanilla attention can be equivariant, if you want your layer to be equivariant. (There are some exceptions in degenerate cases.)
The “Equivariant attention theorem” says roughly: an attention layer is G-equivariant if and only if its heads organize themselves into G-orbits, with the QK matrices Mh and OV matrices Rh satisfying Mg.h = gMhg−1 and Rg.h = gRhg−1. Several heads could have the same QK matrix, in which case we say those heads form a cluster, and only the cluster-sum of the Rh enters in the equivariance condition.
On the left hand side, you choose:
The parameter count of the vanilla attention with token dimension dim V, head count H, and head dimension d is shown.
Then, click on the play button .
A list of clusterings appears below. Each box shows (1) how the heads are clustered, and (2) how G acts on the clusters.
On the right panel, a list of G-orbits shows up.
Click on one to add a new G-orbit G/K. Each dot corresponds to one head in your architecture, and the orbit dictates how G acts on them.
For each orbit, you also specify a K-representation, where K is the stabilizer. This is done again by specifying the multiplicities of irreps of K.
The parameter count of the equivariant architecture is updated live.
Once you're done, click on the play button on the right panel. You will then see which G-clusterings on the left are covered. The goal is to cover everything with minimal parameter count.
For each G, we fixed a list of representatives of conjugacy classes of subgroups of G. These index all possible G-orbits. For each such subgroup K, we also fixed an isomorphism from K to some standard group (e.g. Cn, Dn), and the irrep labels are with respect to this identification.
A strategy that always works: take H copies of the regular orbit (size |G|) with width d (i.e., d copies of the trivial irrep). You can check this by clicking on the regular orbit, putting d in the one multiplicity slot, and cloning it until you have H of them.
This always results in the same parameter count as the vanilla architecture, and you can convince yourself that this is the same (up to a constant) as averaging the vanilla one over G.
However, this is not always optimal.