High-Dimensional Analysis of Single-Layer Attention for Sparse-Token Classification
This work provides foundational theoretical insights into when and how attention mechanisms can learn to focus on sparse, weak signals, which is incremental but important for understanding adaptive token selection in machine learning.
The paper theoretically analyzes a single-layer attention classifier for sparse-token classification, showing it can achieve vanishing test error with logarithmic signal strength scaling, outperforming linear classifiers that require square-root scaling, and proves that just two gradient updates suffice for effective learning in high-dimensional regimes.
When and how can an attention mechanism learn to selectively attend to informative tokens, thereby enabling detection of weak, rare, and sparsely located features? We address these questions theoretically in a sparse-token classification model in which positive samples embed a weak signal vector in a randomly chosen subset of tokens, whereas negative samples are pure noise. In the long-sequence limit, we show that a simple single-layer attention classifier can in principle achieve vanishing test error when the signal strength grows only logarithmically in the sequence length $L$, whereas linear classifiers require $\sqrt{L}$ scaling. Moving from representational power to learnability, we study training at finite $L$ in a high-dimensional regime, where sample size and embedding dimension grow proportionally. We prove that just two gradient updates suffice for the query weight vector of the attention classifier to acquire a nontrivial alignment with the hidden signal, inducing an attention map that selectively amplifies informative tokens. We further derive an exact asymptotic expression for the test error and training loss of the trained attention-based classifier, and quantify its capacity -- the largest dataset size that is typically perfectly separable -- thereby explaining the advantage of adaptive token selection over nonadaptive linear baselines.