Generic Fibers and Functional Dimension of Multi-Head Attention

Weight-space symmetries are pervasive across neural architectures, where symmetries often generate positive dimensional families of parameter tuples that realize the same function. Multi-head self-attention has two known symmetries, including linear changes of basis within Q-K and O-V products and permutations of heads. We prove that for generic parameters, multi-head self-attention has no additional hidden symmetries, characterizing the generic fibers of the parametrization. As a consequence, we derive a formula for the functional dimension of multi-head attention. While an architecture's number of parameters is often used as a measure of expressivity as models are scaled, our functional dimension formula demonstrates when this heuristic is valid and when it is less reliable. We study a setting where the number of attention heads is varied while keeping the number of parameters fixed, and experimentally demonstrate that the model's memorization capacity in bits per parameter varies with the number of heads, but that most of this variation is explained away when the bits memorized is normalized by the functional dimension instead of the number of parameters.