Softmax as Linear Attention in the Large-Prompt Regime: a Measure-based Perspective

2025/12/12 by Boursier, Etienne, Boyer, Claire
Computer Science · Engineering · Neuroscience · #Advanced Memory and Neural Computing #Adversarial Robustness in Machine Learning #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Neural dynamics and brain function

paper · doi:10.48550/arxiv.2512.11784

openalex publication_date 2025/12/12 · openalex created_date 2025/12/16 · openalex updated_date 2026/07/29

Abstract

Softmax attention is a central component of transformer architectures, yet its nonlinear structure poses significant challenges for theoretical analysis. We develop a unified, measure-based framework for studying single-layer softmax attention under both finite and infinite prompts. For i.i.d. Gaussian inputs, we lean on the fact that the softmax operator converges in the infinite-prompt limit to a linear operator acting on the underlying input-token measure. Building on this insight, we establish non-asymptotic concentration bounds for the output and gradient of softmax attention, quantifying how rapidly the finite-prompt model approaches its infinite-prompt counterpart, and prove that this concentration remains stable along the entire training trajectory in general in-context learning settings with sub-Gaussian tokens. In the case of in-context linear regression, we use the tractable infinite-prompt dynamics to analyze training at finite prompt length. Our results allow optimization analyses developed for linear attention to transfer directly to softmax attention when prompts are sufficiently long, showing that large-prompt softmax attention inherits the analytical structure of its linear counterpart. This, in turn, provides a principled and broadly applicable toolkit for studying the training dynamics and statistical behavior of softmax attention layers in large prompt regimes.

Citations

• Token Sample Complexity of Attention
• Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descent
• Statistical Advantage of Softmax Attention: Insights from Single-Location Regression
• Learning In-context \pmbn-grams with Transformers: Sub-\pmbn-grams Are Near-stationary Points
• In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention
• A Unified Perspective on the Dynamics of Deep Transformers
• Training Dynamics of In-Context Learning in Linear Attention
• Analysis of mean-field models arising from self-attention dynamics in transformer architectures with layer normalization
• Emergence of meta-stable clustering in mean-field transformer models
• Asymptotic theory of in-context learning by linear attention
• The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains
• A mathematical perspective on Transformers
• In-Context Convergence of Transformers
• Linear attention is (maybe) all you need (to understand transformer optimization)
• One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention
• Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression
• Trained Transformers Learn Linear Models In-Context
• Transformers learn in-context by gradient descent
• What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
• Data Distributional Properties Drive Emergent In-Context Learning in Transformers
• Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?
• Sinkformers: Transformers with Doubly Stochastic Attention
• Gradient Descent on Infinitely Wide Neural Networks: Global Convergence and Generalization
• A Functional Perspective on Learning Symmetric Functions with Neural Networks
• A Mathematical Theory of Attention
• Language Models are Few-Shot Learners
• Exact lower and upper bounds on the incomplete gamma function
• Stochastic Deep Networks
• Approximation capability of neural networks on spaces of probability measures and tree-structured domains
• Stochastic Particle Gradient Descent for Infinite Ensembles
• Attention Is All You Need
• Bandits with heavy tail

Cited by

• Token Sample Complexity of Attention

Related

• Statistical Advantage of Softmax Attention: Insights from Single-Location Regression
• Why Softmax Attention Outperforms Linear Attention
• Self-Adjust Softmax
• Scalable-Softmax Is Superior for Attention
• On the Expressiveness of Softmax Attention: A Recurrent Neural Network Perspective