A Capacity-Based Rationale for Multi-Head Attention

We study the capacity of the self-attention key-query channel: for a fixed budget, how many distinct token-token relations can a single layer reliably encode? We introduce *Relational Graph Recognition*, where the key-query channel encodes a directed graph and, given a context (a subset of the vertices), must recover the neighbors of each vertex in the context. We measure resources by the total key dimension $D_K = hd_k$. In a tractable multi-head model, we prove matching information-theoretic lower bounds and upper bounds via explicit constructions showing that recovering a graph with $m'$ relations in $d_{\text{model}}$-dimensional embeddings requires $D_K$ to grow essentially as $m'/d_{\text{model}}$ up to logarithmic factors, and we obtain corresponding guarantees for scaled-softmax attention. This analysis yields a new, capacity-based rationale for multi-head attention: even in permutation graphs, where **all queries attend to a single target**, splitting a fixed $D_K$ budget into multiple heads increases capacity by reducing interference from embedding superposition. Controlled experiments mirror the theory, revealing sharp phase transitions at the predicted capacity, and the multi-head advantage persists when adding softmax normalization, value routing, and a full Transformer block trained with frozen GPT-2 embeddings.

Imagine a listener at a crowded party trying to follow who is talking to whom. Modern AI language systems face the same challenge through a mechanism called "attention," which links related items in a passage—say, connecting a pronoun to the name it refers to. We asked a basic question: for a fixed budget, how many such links can a single attention layer keep straight before the voices blur together and it starts making mistakes? To find out, we created a clean, controllable test in which the relationships to be learned form a graph the model must recover. We then proved mathematical limits showing that the required budget grows in proportion to the number of relationships divided by how much internal space the model uses to describe each item. This also explained a principle used in every Transformer, the design behind today's large language models: splitting attention into many parallel "heads." Because compact descriptions force unrelated items to share space and create background noise, dividing the work among many smaller listeners cuts the interference—like splitting a party into small conversations. Knowing this limit lets engineers size attention layers by principle rather than by trial and error.