No Equivariant Architecture Covers All Equivariant Attention

2026-08-31Machine Learning

Machine Learning
AI summary

The authors studied multi-head self-attention (MHSA) layers when they are designed to be symmetric under certain group actions, meaning they behave consistently when inputs are transformed by those groups. They found that for such symmetry to hold, the group can only rearrange groups of attention heads, with specific constraints on internal matrices. They also showed that if you fix an MHSA architecture to exactly preserve this symmetry, it limits the kinds of functions the layer can represent. In particular, the space of symmetric MHSA parameters splits into many separate parts, and a single architecture can only cover one part, leading to reduced expressiveness. The authors quantified this complexity for a specific group and setup, showing an enormous number of these parts exist even for modest model sizes.

EquivarianceMulti-head self-attentionSymmetry groupsRepresentation theoryEquivariance constraintParameter spaceZariski-irreducible componentsD4 groupRegular representationExpressivity
Authors
Tīkun Ông
Abstract
We give a complete characterization of equivariant multi-head self-attention (MHSA): if an MHSA layer is equivariant to a symmetry group $G$, then $G$ can only act by permuting head-clusters, with QK and OV matrices satisfying an equivariance constraint tied to the group action. As a consequence, we prove that any fixed MHSA architecture that achieves exact equivariance by polynomially parameterizing unconstrained MHSA parameters inevitably leads to expressivity loss within the class of equivariant maps: the equivariance locus of unconstrained MHSA forms a union of extremely many Zariski-irreducible components in a reduced parameter space, and any single architecture covers at most one. For $G=D_4$ acting on $C$ copies of the regular representation as the token feature space, we show that there are $Ω(C^{64})$ components for eight attention heads.