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

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. centered 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.

Transformers rely on softmax attention, a powerful but mathematically difficult mechanism at the core of modern AI systems. In this work, we show that when prompts become very long, softmax attention behaves increasingly like a simpler and more tractable form, given by its measure based-limit. We develop a mathematical framework that quantifies this convergence, both for model predictions and for training dynamics. As an application, we prove new guarantees for in-context linear regression with softmax attention, including settings with anisotropic Gaussian data. Our results provide a principled explanation for why simplified linear-attention models can successfully capture the behavior of practical transformer architectures in large-prompt regimes with Gaussian data.