On the Capacity of Self-Attention
Micah Adler
TL;DR
The paper formalizes the capacity of self-attention via Relational Graph Recognition (RGR), modeling edge recovery among $m$ items with $m'$ directed edges under a total key dimension budget $D_K=h d_k$. It derives a tight capacity law, $D_K = \Theta\left(\dfrac{m' \log m'}{d_{\text{model}}}\right)$, showing this is both necessary and sufficient for broad graph families and contexts, and provides explicit constructions and empirical validation. A central insight is a capacity-based rationale for multi-head attention: distributing budget across many small heads mitigates interference when embeddings are compressed, increasing the number of recoverable relations even when each item attends to a single target. The work combines information‑theoretic lower bounds with constructive algorithms under Gaussian and restricted‑incoherence embedding models, and confirms predictions with controlled experiments, yielding a principled design rule for allocating key–query budget across heads in self‑attention systems. These results illuminate when and how self‑attention capacity is achieved and offer a concrete, falsifiable framework to guide architecture choices in capacity‑constrained settings.
Abstract
While self-attention is known to learn relations among tokens, we lack a formal understanding of its capacity: how many distinct relations can a single layer reliably recover for a given budget? To formalize this, we introduce Relational Graph Recognition (RGR), where the key-query channel represents a graph on $m$ items with $m'$ directed edges, and, given a context of items, must recover the neighbors of each item. We measure resources by the total key dimension $D_K = h\,d_k$. Within this framework, we analytically derive a capacity scaling law and validate it empirically. We show that $D_K = Θ(m' \log m' / d_{\text{model}})$ is both necessary (information-theoretic lower bound) and sufficient (explicit construction) in a broad class of graphs to recover $m'$ relations. This scaling law directly leads to a new, capacity-based rationale for multi-head attention that applies even when each item only attends to a single target. When embeddings are uncompressed ($m = d_{\text{model}}$) and the graph is a permutation, a single head suffices. However, compression ($m > d_{\text{model}}$) forces relations into overlapping subspaces, creating interference that a single large head cannot disentangle. Our analysis shows that allocating a fixed $D_K$ across many small heads mitigates this interference, increasing the number of recoverable relations. Controlled single-layer experiments mirror the theory, revealing a sharp performance threshold that matches the predicted capacity scaling and confirms the benefit of distributing $D_K$ across multiple heads. Altogether, these results provide a concrete scaling law for self-attention capacity and a principled design rule for allocating key-query budget across heads.