High-Dimensional Analysis of Single-Layer Attention for Sparse-Token Classification

2025-09-29

AI Summary

This work investigates how attention mechanisms efficiently detect weak, sparse, and rare features in long sequences. Method: Leveraging high-dimensional random matrix theory and mean-field analysis, we characterize the training dynamics and generalization performance of single-layer attention classifiers in the asymptotic regime where sample size and embedding dimension grow proportionally. Contribution/Results: We derive the first exact asymptotic expressions for both test error and training loss, and quantify the data separability capacity of attention. Theoretically, we prove that only two gradient steps suffice to effectively align latent weak signals; moreover, when signal strength scales logarithmically with sequence length, the test error vanishes—substantially outperforming linear classifiers, which require signal strength scaling as the square root of sequence length.

Technology Category

• Machine Learning: Learning Theory
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs
• Security and Privacy: Large-scale security measurements
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Analyzing attention mechanisms for sparse-token classification tasks
• Studying learnability of attention in high-dimensional regimes
• Quantifying attention's advantage over linear classifiers in feature selection

Innovation

• Single-layer attention enables logarithmic signal strength scaling
• Two gradient updates align query weights with hidden signal
• Attention map selectively amplifies informative sparse tokens