Limitations of Normalization in Attention Mechanism

TL;DR

This work analyzes limitations of normalization in attention mechanisms, with a focus on softmax. It introduces a general normalization framework and proves a distance bound (Theorem 1), a geometric separability bound under uniform spherical embeddings (Theorem 2), and a gradient bound (Lemma) that yields a $\| abla_l \boldsymbol{\alpha}\|_2 \le \min\{\tfrac{1}{4T}, \sqrt{2}\}$ for softmax, clarifying the trade-off between sharpness and gradient stability. Empirical validation on GPT-2 demonstrates distance collapse as the top-$N$ grows with sequence length, saturation of geometric separability around 70–85%, and $1/T$-scaled gradient sensitivity at low temperatures. The results motivate robust normalization or sparsity-guided alternatives (e.g., Sparsemax, Scalable-Softmax, Self-Adjusted Softmax) and provide practical guidelines for designing long-context Transformers with reliable selective capabilities.

Abstract

This paper investigates the limitations of the normalization in attention mechanisms. We begin with a theoretical framework that enables the identification of the model's selective ability and the geometric separation involved in token selection. Our analysis includes explicit bounds on distances and separation criteria for token vectors under softmax scaling. Through experiments with pre-trained GPT-2 model, we empirically validate our theoretical results and analyze key behaviors of the attention mechanism. Notably, we demonstrate that as the number of selected tokens increases, the model's ability to distinguish informative tokens declines, often converging toward a uniform selection pattern. We also show that gradient sensitivity under softmax normalization presents challenges during training, especially at low temperature settings. These findings advance current understanding of softmax-based attention mechanism and motivate the need for more robust normalization and selection strategies in future attention architectures.

Figures (4)

• Figure 1: Illustrative example of the geometric separation. Left: Token embeddings lie on a circle. Middle: After scaling by their attention weights $\alpha_i$, both attended (magenta stars) and non-attended (black dots) points move toward the origin. Right: Only the selected tokens that remain inside the ball $B_{r}(s)$ (shaded) are deemed distinguishable.
• Figure 2: Distance statistics validate Theorem \ref{['th1']}. (a) With $N=5$, both the true distance (green) and its expectation (blue) grow roughly linearly in $L$; the red upper bound is safe but conservative. (b) With $L=1024$, increasing $N$ beyond 20 yields diminishing returns: the distance plateaus while the bound tightens. (c) Critical top-$N$ obtained by a KS test ($\alpha=0.01$); fewer than 6 % of the tokens need to be selected before the empirical and expected distances become statistically indistinguishable.
• Figure 3: Geometric separability saturates at 70–85%. For increasing top-$N$, the empirical fraction of distinguishable embeddings $N_s/N$ (green boxes) quickly plateaus; roughly one-fifth of selected tokens remain outside $B_r(s)$.
• Figure 4: Gradient sensitivity decays as $\mathbf{1/T}$. Maximum finite-difference Jacobian norm $g(T,\varepsilon)$ for three perturbation magnitudes (coloured curves, log–log scale). The dashed black curve is the theoretical bound $\min\{1/(4T),\,\sqrt2\}$ from Corollary \ref{['cor:softmax_grad']}.

Theorems & Definitions (10)

• Lemma 1
• Corollary 1
• Theorem 1
• Corollary 2
• Theorem 2
• Lemma 2
• Corollary 3
• proof : Lemma 1
• proof : Theorem 1
• proof : Theorem 2