Gradient Flow Polarizes Softmax: Why Transformers Are Obsessed with Sparse Attention

TL;DR

Why do Transformers focus so much on the "first token" or create "massive activations" in a few dimensions? This paper reveals that it’s not just the data—it’s the math. By analyzing the value-softmax model, the authors prove that gradient flow acts like a "survival of the fittest" algorithm, naturally driving attention towards low-entropy (one-hot) solutions even when many dense solutions are available.

Background: The Mystery of the Attention Sink

In modern LLMs, we often observe "attention sinks"—where a head exerts nearly all its attention on a single, often semantically irrelevant token (like a BOS token or a period). While we've known what these are, we haven't fully understood why the optimization process prefers them. Is it a functional necessity, or a side effect of our architecture?

The "Value-Softmax" Lens

To answer this, the authors isolate the core operation of an attention head: $\beta =V\sigma (a)$ Where $V$ is a learnable value matrix and $a$ is an attention score vector. They look at how gradient flow (the continuous version of Gradient Descent) evolves these parameters.

The Replicator Dynamics Intuition

The breakthrough insight is mirroring Replicator Dynamics from evolutionary biology. In these systems, the rate of change of a population trait is proportional to its "fitness" relative to the average.

The authors show that in the value-softmax model:

• Order Preservation: If token A starts with a higher score than token B, it stays higher throughout training.
• Repulsion: The gap between the "best" token and the rest grows over time.

Figure 1: As training progresses, the winner-take-all dynamic forces one attention score toward 1.0 while others decay to zero.

Why Softmax is the Culprit

The paper proves that for Logistic Loss (standard in classification/language modeling), the attention vector $\sigma (a)$ must converge to a one-hot vector.

Interestingly, this doesn't happen with:

1. Sigmoid Attention: Lacks the mean-centering term that creates "competition" between tokens.
2. Linear Networks: While they favor rank-1 solutions, they don't force the scores to align with specific coordinate axes (tokens).

Impact: Sinks, Outliers, and Adversarial Vulnerability

The theoretical polarization carries heavy practical implications:

1. Massive Activations: This is the "flip side" of attention sinks. As attention becomes one-hot, the values in $V$ associated with the "sink" token must grow to satisfy the loss, leading to activation outliers.
2. Token Influence Imbalance: Because the model focuses on so few tokens, it becomes fragile.

Figure 2: In single-head models, swapping the "max-logit" token can flip the model's entire prediction with high success rates.

Conclusion: A Built-in Bias

The study highlights that the Softmax parameterization provides an inductive bias toward sparsity. While sparsity is often useful for efficiency (like KV cache compression), it can result in "pathological" behaviors where the model ignores the majority of its context. For researchers building the next generation of LLMs, this suggests that if we want "denser" attention, we need to look beyond Softmax and consider alternative normalization functions that don't satisfy the polarizing conditions of replicator dynamics.