Geometric Analysis of Token Selection in Multi-Head Attention

TL;DR

This work investigates token selection in multi-head attention through a geometric lens, modeling attention as a classifier in value–state space with a top-$N$ selection and defining Precision, Recall, and F-score via interpretable radii. It derives non-asymptotic, dimension- and margin-dependent bounds under empirically grounded assumptions and validates the theory on open LLMs, revealing a robust head taxonomy (Retriever, Mixer, Reset) with clear implications for sparsification and interpretability. Empirical results show the theory’s envelopes match observed behavior, with the strongest separability at small $N$ and a sink token shaping recall through geometric alignment. Overall, the framework provides a principled, geometry-aware approach to understanding and pruning attention in large language models.

Abstract

We present a geometric framework for analysing multi-head attention in large language models (LLMs). Without altering the mechanism, we view standard attention through a top-N selection lens and study its behaviour directly in value-state space. We define geometric metrics - Precision, Recall, and F-score - to quantify separability between selected and non-selected tokens, and derive non-asymptotic bounds with explicit dependence on dimension and margin under empirically motivated assumptions (stable value norms with a compressed sink token, exponential similarity decay, and piecewise attention weight profiles). The theory predicts a small-N operating regime of strongest non-trivial separability and clarifies how sequence length and sink similarity shape the metrics. Empirically, across LLaMA-2-7B, Gemma-7B, and Mistral-7B, measurements closely track the theoretical envelopes: top-N selection sharpens separability, sink similarity correlates with Recall. We also found that in LLaMA-2-7B heads specialize into three regimes - Retriever, Mixer, Reset - with distinct geometric signatures. Overall, attention behaves as a structured geometric classifier with measurable criteria for token selection, offering head level interpretability and informing geometry-aware sparsification and design of attention in LLMs.

Figures (15)

• Figure 1: Illustrative 2D example of geometric separability in value-state space.Left: Token embeddings lie on a circle. Middle: After scaling by their attention weights $\alpha_i$, both selected (magenta stars) and non-selected (black dots) points move toward the origin. Right: We see the points lying in the $B_{r_{\min}}(s)$ and $B_{r_{\max}}(s)$ as well as the $r_{\min}$ and $r_{\max}$.
• Figure 2: Except attention sink value–state norms remain stable.Left: Average norm of value states across all layer–head pairs. Right: Coefficient of variation of the non-sink norms.
• Figure 3: Exponential fitting provides a good approximation with low error. The distribution of mean absolute error across layer-heads shows the average 9% error.
• Figure 4: Attention mass follows a piecewise-structured distribution across sequence positions. This piecewise structure highlights that attention is highly non-uniform, reflecting positional and linguistic regularities.
• Figure 5: Attention achieves its strongest non-trivial separability at small selection sizes ($N\in[1,4]$), while separability is weakest at intermediate $N$. Global means of Precision (left) and Recall (right) versus $N$ (solid blue), with theoretical upper (dotted orange) and lower (dashed green, capped at $0$) envelopes.

Theorems & Definitions (6)

• Definition 1
• Definition 2
• Definition 3
• Theorem 1: Precision bound
• Theorem 2: Recall bound
• Corollary 1: F-score bounds