Geometric Analysis of Token Selection in Multi-Head Attention

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.

💡 Research Summary

The paper proposes a geometric framework for analyzing token selection in multi‑head attention without modifying the underlying mechanism. By viewing standard attention as a top‑N selector operating in the value‑state space, the authors define three classification‑style metrics—Precision, Recall, and F‑score—to quantify how well the selected tokens are separated from the non‑selected ones. Two radii, r_min (the minimal distance from any non‑selected token to the representative vector) and r_max (the maximal distance from any selected token), provide concrete bounds for these metrics.

Empirically, the authors observe three consistent patterns across several open‑source LLMs (LLaMA‑2‑7B, Gemma‑7B, Mistral‑7B): (1) value‑state norms are highly stable across positions except for the first token, which acts as an “attention sink” with a compressed norm; (2) cosine similarity between token embeddings decays approximately exponentially with positional distance; (3) attention weights follow a piecewise profile consisting of a sink, a flat plateau, an oscillatory segment, and an exponential recency bias. These observations motivate three assumptions: (i) stable norms with a compressed sink, (ii) exponential decay of cross‑token similarity, and (iii) a four‑phase attention weight function parameterised by sink probability, sensitivity, frequency, and phase transition points.

Under these assumptions the authors introduce a margin Δ and a scale B, then derive non‑asymptotic bounds for the expected Precision and Recall that depend explicitly on the embedding dimension d, the margin, and the scale. Theorem 1 shows that when Δ>0 and d is large, the expected Precision is at least 1 − exp(−κd·Δ/B²), with a complementary upper bound involving the worst‑case pairwise probability p_{ij}. Theorem 2 provides an analogous bound for Recall. Corollary 1 translates these into bounds for the F‑score. The analysis predicts a “small‑N” regime (selecting only a few top‑weighted tokens) where separability is strongest, a degradation at intermediate N where the margin shrinks, and a trivial return to perfect scores when N approaches the full sequence length.

Experiments on OpenWebText with the three models confirm the theoretical envelopes: Precision and Recall curves lie within the predicted bounds, and the strongest non‑trivial separability indeed occurs for N≈1–4. The authors further cluster heads based on their Precision/Recall profiles into three functional types—Retriever (high Recall, low Precision), Mixer (balanced), and Reset (high Precision, low Recall). This taxonomy is robust across layers and models.

To demonstrate practical relevance, the paper applies the taxonomy to head sparsification. By ranking heads according to their type‑specific importance and removing a fraction of the least important heads, the authors achieve lower increases in negative log‑likelihood (ΔNLL) compared with random or score‑based pruning, especially at long context lengths (L=1024).

Overall, the work offers a unified geometric classifier view of attention, provides analytically tractable metrics for token selection quality, validates the theory with extensive empirical evidence, and shows how the insights can guide model compression. Limitations include the reliance on globally fixed parameters (λ, β, η, etc.) that may vary across models or training stages, and the lack of a concrete link between the exponential similarity decay and linguistic structure. Future work could focus on adaptive estimation of these parameters, deeper connections to syntax/semantics, and extending the framework to alternative attention variants (e.g., sparsemax, entmax).