On the Capacity of Self-Attention

ICLR 2026 Conference Submission20044 Authors

Keywords: Self-Attention, Transformers, Capacity, Scaling Laws, Multi-Head Attention, Information Theory, Mechanistic Interpretability, Deep Learning Theory, Superposition.

TL;DR: We derive and empirically validate a scaling law for self-attention's capacity to recognize relationships, revealing a new rationale for using multiple attention heads, even when each item attends to only a single target.

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 \emph{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 = \Theta(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.

Primary Area: interpretability and explainable AI

Submission Number: 20044