The Approximation Rank of Softmax Attention: Sharp Geometric Laws and Robust Interaction Dimension
Machine LearningArtificial Intelligence
Summary
The authors investigate what shapes affect the complexity when approximating softmax attention, a key part of transformer models. They find that the geometry of the input data's support—whether points lie on a sphere or fill a ball—predicts how the approximation rank scales with parameters like dimension and error tolerance. They also study how individual attention heads behave, showing a precise link between the effective dimension of query-key interactions and approximation difficulty. Experiments on BERT confirm moderate reductions in complexity consistent with their theoretical predictions. Overall, their work clarifies how geometric properties influence attention approximation limits.
Authors
Yuhe Sui, Jianing Zhang
Abstract
Which geometry controls the rank complexity of normalized softmax attention? We study maximum-row-$\ell_1$ approximation rank, exactly the least unrestricted rank preserving every bounded vector-valued output. Two sharp worst-case laws isolate support geometry: for fixed $d$ and error $\varepsilon$, spherical self-attention has rank $Θ_{d,\varepsilon}(\min\{n,(1+β)^{(d-1)/2}\})$, while full-ball geometry adds one radial degree and, for $β\geβ_0(d,\varepsilon)$ and $n\ge C_d e^{β/8}$, gives $Θ_{d,\varepsilon}(β^{d/2})$. For a fixed head, row-softmax quotients out row-scalar logit directions: the remaining visible query--key interaction dimension $r$ yields an $r/2$ per-instance upper law, and bounded constructions show this exponent is minimax sharp. Approximate interaction subspaces incur an explicit residual output error and yield a tolerance-indexed SVD dimension. On an 84-head BERT-base calibration set, we observe modest effective-dimension reductions across many head--temperature settings, together with positive associations with finite constructive rank upper certificates. Together, these results separate support geometry, which sets worst-case temperature scaling, from softmax-visible interaction geometry, which controls per-head approximation complexity.