Token Sample Complexity of Attention

TL;DR

This paper introduces the concept of token-sample complexity to analyze how attention mechanisms in large language models converge at extreme sequence lengths, providing theoretical bounds and empirical validation.

Contribution

It establishes new convergence bounds for attention maps and token distributions, addressing limitations of previous analyses at large context sizes.

Findings

Attention map convergence rate is $C(R)/\sqrt{n}$ for compactly supported distributions.

Convergence rate for moments is $C'(R)/n^{eta}$ with $eta<rac{1}{2}$, depending on distribution support.

Experimental results on Gaussian data and BERT models validate theoretical predictions.

Abstract

As context windows in large language models continue to expand, it is essential to characterize how attention behaves at extreme sequence lengths. We introduce token-sample complexity: the rate at which attention computed on $n$ tokens converges to its infinite-token limit. We estimate finite-$n$ convergence bounds at two levels: pointwise uniform convergence of the attention map, and convergence of moments for the transformed token distribution. For compactly supported (and more generally sub-Gaussian) distributions, our first result shows that the attention map converges uniformly on a ball of radius $R$ at rate $C(R)/\sqrt{n}$, where $C(R)$ grows exponentially with $R$. For large $R$, this estimate loses practical value, and our second result addresses this issue by establishing convergence rates for the moments of the transformed distribution (the token output of the attention layer). In this case, the rate is $C'(R)/n^{\beta}$ with $\beta<\tfrac{1}{2}$, and $C'(R)$ depends polynomially on the size of the support of the distribution. The exponent $\beta$ depends on the attention geometry and the spectral properties of the tokens distribution. We also examine the regime in which the attention parameter tends to infinity and the softmax approaches a hardmax, and in this setting, we establish a logarithmic rate of convergence. Experiments on synthetic Gaussian data and real BERT models on Wikipedia text confirm our predictions.