🤖 AI

Geometric Analysis of Token Selection in Multi-Head Attention

This paper introduces a geometric framework for analyzing multi-head attention in large language models by defining precision, recall, and F-score metrics to quantify token separability, deriving non-asymptotic bounds that predict optimal performance in small-N regimes, and empirically validating these findings across multiple models to reveal distinct head specialization patterns and inform geometry-aware sparsification.

Original authors: Timur Mudarisov, Mikhal Burtsev, Tatiana Petrova, Radu State

Published 2026-08-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Timur Mudarisov, Mikhal Burtsev, Tatiana Petrova, Radu State

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a giant, digital library where a robot librarian is trying to answer your questions. To do this, the librarian doesn't just read one book; it scans millions of pages at once. But reading everything at the same time is chaotic and slow. So, the librarian uses a special trick called "attention." Think of attention like a spotlight. When the librarian looks at a specific word in your question, the spotlight shines on the most important words in the history of the conversation to help form an answer.

In the world of artificial intelligence, this spotlight system is called "Multi-Head Attention." It's the engine that powers modern AI models, allowing them to understand context and nuance. However, for a long time, scientists treated this spotlight as a simple averaging machine—like a smoothie where every ingredient is blended together equally. But recent observations suggested the spotlight is actually more like a selective filter, picking out specific ingredients while ignoring others. The big question researchers have been asking is: Is the group of words the spotlight picks out actually a coherent, organized team, or is it just a random mix of words that happen to be nearby? Understanding this geometry—the shape and arrangement of the selected words—could help us make AI models smaller, faster, and more efficient without losing their smarts.

In this paper, the authors, Timur Mudarisov and his team, decided to stop looking at the spotlight as a blurry glow and start examining the individual beams of light it casts. They developed a new way to look at how AI models choose which words to focus on, treating the selection process like a game of "geometric sorting." Instead of asking if the selected words make sense linguistically, they asked a simpler, more visual question: Do the selected words cluster together in a neat, tight group, or are they scattered all over the place?

To answer this, they invented a set of "geometric rulers" called precision and recall. Imagine you have a basket of marbles (the words the AI picked) and a basket of pebbles (the words it ignored). If the marbles are all bunched up in one corner of the room and the pebbles are far away, that's a perfect score. If the marbles are mixed in with the pebbles, the score is low. The authors found that when AI models pick their top words, those words do indeed form a much tighter, more organized cluster than if you had just picked words at random. It's as if the AI has a secret instinct for grouping related ideas together, even if it's just looking at the math behind the scenes.

However, the story gets a bit weirder when they looked at the very first word of any sentence. In AI land, this is known as the "sink token." It's like a unique anchor in the middle of the room that receives a massive amount of attention, even though it doesn't really carry any useful information. The authors discovered that this "sink" is a geometric oddity. It has a distinct, smaller size (norm) and points in the opposite direction of all the other words. When the AI includes this sink in its top picks, it messes up the neat clustering of the other words. It's like trying to organize a team photo, but one person is standing in a different dimension and pulling everyone else off-balance. The paper shows that if you remove this "sink" from the selection, the remaining words snap into a much more perfect, organized shape.

Based on these geometric observations, the team created a new way to classify the different "spotlights" (or attention heads) inside the AI. They didn't use complex math formulas or guesswork; they just looked at which word each spotlight liked the most. This led to three distinct personalities for the AI's attention heads:

Retrievers: These are the focused detectives. They almost always pick the current word being discussed as their most important source. They are rare, making up only about 6% to 17% of the team.
Mixers: These are the social butterflies. They love the "sink" (the first word) and use it as a base to blend in information from other words. They are very common in some models, like Gemma and LLaMA-3.
Resets: These are the wanderers. They don't have a favorite word; they pick from all over the place. They are the most common type in models like Mistral and LLaMA-2.

The authors tested this new classification system by trying to "prune" or remove parts of the AI to see if it could still work. They found that knowing whether a head is a Retriever, Mixer, or Reset gives a helpful hint about which parts of the AI are important. It's like knowing which players on a sports team are the goal-scorers and which are the defenders. However, they were careful to note that this isn't a magic bullet. While it helps at moderate levels of cutting, it doesn't always beat other methods when you try to cut the AI down to the bare minimum. Also, they found one specific "super-head" in the Gemma model that was so critical that removing it broke the whole system, proving that even with these neat categories, some individual parts are irreplaceable.

Ultimately, this paper suggests that the way AI models select information is far more structured and geometric than we thought. It's not just a random blur; it's a carefully organized dance where some words lead, some follow, and one tiny "sink" word pulls the whole group in a strange direction. By understanding these shapes, we might be able to build smarter, leaner AI models in the future, though we still have to be careful not to cut the wrong pieces.

Technical Summary: Geometric Analysis of Token Selection in Multi-Head Attention

Problem Statement
While the attention mechanism is the core of modern Large Language Models (LLMs), existing theoretical frameworks often focus on universal approximation capabilities or view attention as a normalized kernel smoother. These perspectives clarify expressivity but fail to describe the specific geometry of tokens selected by trained attention heads. Empirical observations, such as the "attention sink" phenomenon where early tokens receive disproportionate mass, suggest that attention involves structured selection rather than diffuse averaging. However, a simple framework to determine whether the tokens selected by attention form a geometrically separable class in value-state space has been lacking. This paper addresses the question: Do the top- $N$ attention-selected contributions form a coherent, separable set in the value-state space, and how does the attention sink influence this geometry?

Methodology
The authors propose a geometric framework that treats a single attention head as a classifier over attention-scaled value contributions ( $y_i = \alpha_i v_i$ ) rather than analyzing raw attention weights.

Geometric Metrics: For a fixed query position, the top- $N$ tokens by attention weight are selected. The authors define the aggregate vector $s_N = \sum_{i \in I_N} \alpha_i v_i$ . They introduce Geometric Precision and Geometric Recall to measure how well the selected contribution vectors separate from the unselected ones around $s_N$ . These are summarized into an Extremal Geometric F-score ( $F_N$ ).
Theoretical Bounds: The authors derive deterministic lower bounds for precision and recall based on "selected-unselected inversions" (cases where an unselected token is closer to the aggregate than a selected one). Crucially, they introduce a sink-aware constant-norm regime. This theory assumes non-sink tokens share a constant norm and bounded directional coherence, while the sink token (first token) has a distinct, smaller norm and a signed directional relationship (often negative alignment) with non-sink tokens.
Head Taxonomy: A parameter-free taxonomy is introduced based on the "source-winner" rule. For each head, the source position maximizing the contribution ( $\alpha_j \|v_j\|$ $α_{j} ∥ v_{j} ∥$ ) is identified. Heads are classified as:
- Retriever: The current token is the dominant source.
- Mixer: The sink token is the dominant source.
- Reset: No single privileged source dominates (winning sources are distributed).
Empirical Validation: The framework is evaluated on LLaMA, Gemma, and Mistral models using WikiText-103 and OpenWebText corpora. Experiments include comparing top- $N$ selection against random subsets, analyzing sink geometry (norms and cosines), and testing the taxonomy's utility in structured head sparsification.

Key Contributions

Geometric Classifier View: The paper recasts top- $N$ attention as a geometric classification problem in value-state space, introducing precision, recall, and an extremal separability score ( $F_N$ ).
Sink-Aware Bounds: It derives deterministic lower bounds on geometric separability that explicitly account for the anomalous geometry of the sink token, providing conservative certificates of separability.
Parameter-Free Head Taxonomy: It establishes a taxonomy of Retriever, Mixer, and Reset regimes based solely on contribution-maximizing sources, requiring no fitted thresholds or margins.
Empirical Characterization: The study provides empirical evidence that top-attention subsets are more geometrically separable than random subsets and that the sink token acts as a distinct geometric component that disrupts separability when included in top- $N$ selections.

Results

Separability: Across all models (LLaMA, Gemma, Mistral), top-attention subsets are consistently more geometrically separable than random subsets of the same size, particularly for sparse and intermediate $N$ .
Sink Geometry: The sink token exhibits a distinct geometry: its value norm is 3–10 $\times$ smaller than non-sink tokens, and its direction is systematically negatively aligned with non-sink directions. Including the sink in top- $N$ selection significantly reduces the geometric separability score ( $F_N$ ), often by a factor of two.
Taxonomy Distribution: Retriever heads are consistently a minority (5.8–17.2% across models). The dominant regime (Mixer vs. Reset) varies by architecture but remains stable across corpora (OpenWebText vs. WikiText). Mixer heads show strong sink alignment, while Retriever heads align with the current token.
Sparsification Utility: The taxonomy provides a useful coarse signal for structured head removal, particularly at moderate retention levels. However, it does not universally outperform single-metric baselines (like weight magnitude or entropy) under aggressive pruning. A critical anomaly was identified in Gemma-7B (Layer 0, Head 12), where a single head's removal caused a distinct failure mode, highlighting that regime-level averages can mask essential individual heads.

Significance and Claims
The paper claims to provide a "simple framework" that moves beyond viewing attention as mere averaging, revealing structured selection and head-level specialization. The significance lies in:

Diagnostic Capability: The geometric metrics and sink-aware bounds serve as conservative certificates for separability, validating that attention-selected tokens are not just semantically ranked but geometrically coherent.
Mechanistic Insight: The taxonomy refines the binary "active-dormant" view of attention heads into three distinct mechanistic regimes (Retriever, Mixer, Reset), showing that the balance between sink-mediated and diffuse computation is an architectural property rather than a universal constant.
Practical Limitations: The authors modestly claim that while the taxonomy offers a useful signal for head sparsification, it is not a universally optimal pruning rule. It should be used as a diagnostic signal combined with head-level sensitivity analysis, especially given the existence of critical, individually essential heads that regime-level averages may obscure.

Technical Summary: Geometric Analysis of Token Selection in Multi-Head Attention

More like this