Softmax Attention

Follow this topic Ask about this topic

Momentum

7 papers in the last four weeks, up 133% on the four weeks before. 0.1% of all new papers.

Jul 13Week of Sep 28

Latest papers 47

May 10, 2026cs.LG

Nectar: Neural Estimation of Cached-Token Attention via Regression

Evaluating softmax attention over a fixed long context requires reading every cached key-value pair for each new query token. For a given context (a book, a manual, a legal corpus) the attention output is a deterministic function of the query. We propose Nectar, which fits a compact neural network to this function for queries drawn from a task-relevant distribution. Nectar fits two networks per layer and KV-head: a target network that predicts the attention output and a score network that predicts the log-normalizer. The pair plugs into the standard masked self-attention at inference time, replacing the

O(n)

attention over the cache with a forward pass whose cost does not depend on

n

. Each module carries on the order of

|θ|

parameters per layer and KV-head, typically much smaller than the

2nd

KV-cache footprint at the same granularity. We report experiments on models from 1.7B to 8B parameters across five long-context datasets. The approximation error tracks the next-token accuracy gap to full attention, and allocating capacity non-uniformly across layers reduces that gap in our ablation. Beyond this analysis of metrics, we check that the text generations (following a question prompt) of a model equipped with a Nectar module match in semantic content those obtained by giving the same model access to the full cache.

João Monteiro, Michal Klein, Pierre Ablin +1

Apple

Softmax Attention Self-Attention KV Caching Long-Context Language Model Inference Efficient Attention Multi-Head Attention

May 8, 2026cs.LG

Scaling Limits of Long-Context Transformers

We study the long-context limit of softmax self-attention with a fixed query and a random context of

n

i.i.d. keys on the sphere, viewing the inverse temperature

β_n

as the scaling parameter that decides whether attention degenerates into uniform averaging or collapses onto the single closest key. We show that the critical scale at which selectivity emerges is determined by the local exponent of the distance-to-query distribution near zero rather than by global features of the context, and scales like

β_n^\ast \asymp n^{2/(d-1)}

for uniform keys on

\mathbb{S}^{d-1}

. Furthermore, we characterize the limiting laws of the ordered attention weights and of the attention output across all regimes of

β_n

: a subcritical regime in which the output reduces to a local average around

q

with explicit deterministic bias and Gaussian fluctuations; a critical regime in which a finite collection of nearest keys retains macroscopic mass without single-key collapse; and a supercritical regime in which all mass concentrates on the closest key. Of notable interest is the subcritical case with identity value matrix where the attention map approximately implements a backward heat equation.

Giuseppe Bruno, Shi Chen, Zhengjiang Lin +2

Softmax Attention Self-Attention Long-Context Modeling Scaling Laws

May 8, 2026cs.LG

Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression

Mechanistic accounts of in-context learning (ICL) have identified iterative algorithms for linear regression and related linear prediction tasks, often using linear or ReLU attention variants. For nonlinear ICL, prior work has related softmax and kernelized attention to functional-gradient-type dynamics, but it remains unclear whether a standard transformer with softmax attention can implement a convergent solver with an end-to-end prediction-error guarantee. In this paper, we study in-context kernel ridge regression (KRR) with Gaussian kernels and show that a standard softmax-attention transformer can approximate the KRR predictor during its forward pass by implementing preconditioned Richardson iteration on the associated kernel linear system. Under bounded-data assumptions, we construct a single-head transformer with

O(\log(1/ε))

blocks and MLP width

O(\sqrt{N/ε})

that achieves

ε

-accurate prediction for prompts of length

N

. Our construction reveals a functional decomposition within the transformer architecture: softmax attention produces a row-normalized Gaussian-kernel operator needed for cross-token interactions, while ReLU MLP layers act locally to approximate the intra-token scalar arithmetic required by the update. Empirically, we train GPT-2-style transformers on Gaussian-process regression tasks to further test the preconditioned Richardson interpretation. Through linear probing, we compare the transformer's layer-wise predictions with the step-wise outputs of classical KRR solvers and find that its error profiles align most consistently with preconditioned Richardson iteration. Ablation studies further support this interpretation. Together, our theory and experiments identify preconditioned Richardson iteration as a concrete mechanism that softmax-attention transformers can realize for nonlinear in-context Gaussian-kernel regression.

Mingsong Yan, Dongyang Li, Charles Kulick +1

Department of Mathematics, University of California, Santa Barbara, CA

Softmax Attention Transformer Interpretability Kernel Ridge Regression Kernel Regression Self-Attention In-Context Learning+2

May 8, 2026cs.LG

Vertex-Softmax: Tight Transformer Verification via Exact Softmax Optimization

Certified verification of transformer attention requires bounding the softmax function over interval constraints on the pre-softmax scores. Existing verifiers relax softmax ndependently of the downstream objective, leaving avoidable slack. We prove that the exact optimum of this score-box problem is attained at a vertex of the constraint box, and establish a threshold structure theorem showing that, after sorting the objective coefficients, the optimum lies among only linearly many candidates, yielding the Vertex-Softmax primitive with log-linear complexity in the sequence length. We further prove a formal optimality result showing that Vertex-Softmax is the tightest sound bound obtainable from score intervals alone, characterizing precisely what additional structure (score correlations, score-value coupling) is needed for further improvement. Integrated into a CROWN Convex Relaxation based Optimization for Worst-case Neurons)-style verifier with a formal soundness guarantee, Vertex-Softmax significantly improves certified rates and substantially tightens lower bounds across MNIST, Fashion-MNIST, and CIFAR-10 attention models, while consistently matching or outperforming alpha-CROWN and branch-and-bound baselines at a fraction of their cost.

Navid Rezazadeh, Arash Gholami Davoodi

University of California, Irvine · Carnegie Mellon University

Softmax Attention Neural Network Verification Certified Robustness

May 7, 2026cs.DS

Nearly Optimal Attention Coresets

We consider the problem of estimating the Attention mechanism in small space, and prove the existence of coresets for it of nearly optimal size. Specifically, we show that for any set of unit-norm keys and values

(K,V)

in

\mathbb{R}^d

, there exists a subset

(K',V')

of size at most

O({\sqrt{d} e^{ρ+o(ρ)}/\varepsilon})

such that

\left\| \operatorname{Attn}(q,K,V)- \operatorname{Attn}(q,K',V') \right\| \le \varepsilon

simultaneously for all queries whose norm is bounded by

ρ

. This outperforms the best known results for this problem. We also offer an improved lower bound showing that

\varepsilon

-coresets must have size

Ω({\sqrt{d} e^ρ/ε})

.

Edo Liberty, Alexandr Andoni, Eldar Kleiner

Softmax Attention Efficient Attention Coreset Selection

May 5, 2026cs.LG

Gradient Flow Structure and Quantitative Dynamics of Multi-Head Self-Attention

Transformer self-attention can be interpreted as a gradient flow on the unit sphere, in which tokens evolve under softmax interaction potentials and tend to form clusters. While prior work has established clustering behavior for single-head attention, the multi-head setting remains less understood due to geometric interference between heads, which invalidates standard monotonicity arguments. In this work, we develop a theoretical framework for multi-head self-attention dynamics and resolve several open questions. We show that, under suitable conditions on the score matrices, a natural multi-head energy functional is non-decreasing along both flat and spherical dynamics. We identify the key obstruction to per-head monotonicity as radial shadow terms, which are projections of each head's output onto token directions, persisting even under orthogonality assumptions. We introduce a sufficient condition ensuring monotonicity and establish robustness to approximate orthogonality. In a simplified scalar-head regime with equiangular token configurations, we derive a closed-form expression for the critical inverse temperature governing clustering behavior, and show that heterogeneous heads exhibit super-additive clustering rates. In this regime, we also prove a separation in clustering time between ReLU and softmax attention in the linearized dynamics. Finally, we establish an entropy production identity and show that attention entropy increases monotonically toward equilibrium as clustering progresses. Our results provide a unified perspective on the dynamics of multi-head attention and clarify the mechanisms underlying clustering and stability in transformer models.

Ayan Pendharkar

Softmax Attention Attention Head Analysis Self-Attention Clustering Attention Mechanisms Multi-Head Attention

Apr 29, 2026cs.LG

Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models

Training stable biological foundation models requires rethinking attention mechanisms: we find that using sigmoid attention as a drop in replacement for softmax attention a) produces better learned representations: on six diverse single-cell datasets, sigmoid achieves 25% higher cell-type separation, better cell-type cohesion metrics, and lower validation loss, b) faster training, models with sigmoid attention train up to 10% faster than their softmax counterparts, and c) more stable training by eliminating inherent sources of instability in softmax attention. We establish that sigmoid attention has globally bounded derivatives (

\leq 0.25

) as opposed to softmax, and a diagonal Jacobian structure in contrast with softmax's dense coupling, which together help alleviate training instabilities. In stress tests on 160M-parameter bidirectional attention models trained without gradient clipping on 8K-token sequences, softmax diverges catastrophically, with gradients exploding by four orders of magnitude, while sigmoid remains stable. Finally, we implement and open-source TritonSigmoid, an efficient GPU kernel that achieves 515 TFLOPS on H100 GPUs, outperforming both FlashAttention-2 and FlashSigmoid, with native padding support, which is essential for biological sequences. Our results establish sigmoid attention as both theoretically grounded and empirically superior for biological foundation models. Code is available at https://github.com/MSDLLCpapers/triton-sigmoid

Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller +1

Merck & Co., Inc., Cambridge, MA, USA

Softmax Attention Self-Attention Efficient Attention Attention Mechanisms Neural Network Training Dynamics

Apr 28, 2026physics.data-an

Emergent Self-Attention from Astrocyte-Gated Associative Memory Dynamics

We introduce a Hopfield-type associative memory in which effective connectivity is multiplicatively modulated by astrocytic gains evolving under an entropy-regularized replicator equation. The coupled neuron-astrocyte dynamics admit a Lyapunov function, ensuring global convergence. At fixed points, astrocytic gains implement a softmax-normalized allocation over pattern similarity scores, yielding a mechanistic realization of self-attention as emergent routing on the gain simplex. In regimes of high memory load and interference, the model significantly improves retrieval accuracy relative to classical Hopfield dynamics and recent neuron-astrocyte baselines. These results establish a dynamical systems framework linking glial modulation, competitive resource allocation, and attention-like computation.

Arnau Vivet, Alex Arenas

1Departament d’Enginyeria Inform`atica i Matem`atiques, Universitat Rovira i Virgili, 43007 Tarragona, Spain · 2Complexity Science Hub Vienna, Metternichgasse 8, 1030 Vienna, Austria

Softmax Attention Dynamical Systems Self-Attention Replicator Dynamics Neural Processes Hopfield Networks+2

Apr 28, 2026cs.LG

QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention

FlashAttention improves efficiency through tiling, but its online softmax still relies on floating-point arithmetic for numerical stability, making full quantization difficult. We identify three main obstacles to integer-only FlashAttention: (1) scale explosion during tile-wise accumulation, (2) inefficient shift-based exponential operations on GPUs, and (3) quantization granularity constraints requiring uniform scales for integer comparison. To address these challenges, we propose \textit{QFlash}, an end-to-end integer FlashAttention design that performs softmax entirely in the integer domain and runs as a single Triton kernel. On seven attention workloads from ViT, DeiT, and Swin models, QFlash achieves up to 6.73

\times

speedup over I-ViT and up to 8.69

\times

speedup on Swin, while reducing energy consumption by 18.8% compared to FP16 FlashAttention, without sacrificing Top-1 accuracy on ViT/DeiT and remaining competitive on Swin under per-tensor quantization. Our code is publicly available at https://github.com/EfficientCompLab/qflash.

Sehyeon Oh, Yongin Kwon, Jemin Lee

University of Science and Technology, Daejeon, Republic of Korea · Electronics and Telecommunications Research Institute, Daejeon, Republic of Korea · Pusan National University, Busan, Republic of Korea +1

Softmax Attention GPU Acceleration Vision Transformer Efficient ViTs Attention Quantization Attention Mechanisms+2

Apr 27, 2026cs.LG

Transformer Approximations from ReLUs

We provide a systematic recipe for translating ReLU approximation results to softmax attention mechanism. This recipe covers many common approximation targets. Importantly, it yields target-specific, economic resource bounds beyond universal approximation statements. We showcase the recipe on multiplication, reciprocal computation, and min/max primitives. These results provide new analytical tools for analyzing softmax transformer models.

Jerry Yao-Chieh Hu, Mingcheng Lu, Yi-Chen Lee +1

Center for Foundation Models and Generative AI, Northwestern University, Evanston, IL 60208, USA · Department of Computer Science, Northwestern University, Evanston, IL 60208, USA · Department of Physics, National Taiwan University, Taipei 10617, Taiwan +1

Softmax Attention Neural Network Approximation Theory Attention Mechanisms

Apr 26, 2026cs.LG

ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers

Existing attention accelerators often trade exact softmax semantics, depend on fused Tensor Core kernels, or incur sequential depth that limits FP32 throughput on long sequences. We present \textbf{ELSA}, an algorithmic reformulation of online softmax attention that (i)~preserves exact softmax semantics in real arithmetic with a \emph{provable}

\mathcal{O}(u\log n)

FP32 relative error bound; (ii)~casts the online softmax update as a prefix scan over an associative monoid

(m,S,W)

, yielding

O(n)

extra memory and

O(\log n)

parallel depth; and (iii)~is Tensor-Core independent, implemented in Triton and CUDA C++, and deployable as a \emph{drop-in replacement} requiring no retraining or weight modification. Unlike FlashAttention-2/3, which rely on HMMA/GMMA Tensor Core instructions and provide no compatible FP32 path, ELSA operates identically on A100s and resource-constrained edge devices such as Jetson TX2 -- making it the only hardware-agnostic exact-attention kernel that reduces parallel depth to

O(\log n)

at full precision. On A100 FP32 benchmarks (1K--16K tokens), ELSA delivers

1.3

--

3.5\times

speedup over memory-efficient SDPA and

1.97

--

2.27\times

on BERT; on Jetson TX2, ELSA achieves

1.5

--

1.6\times

over Math (64--900 tokens), with

17.8

--

20.2\%

throughput gains under LLaMA-13B offloading at

\ge

32K. In FP16, ELSA approaches hardware-fused baselines at long sequences while retaining full FP32 capability, offering a unified kernel for high-precision inference across platforms. Our code and implementation are available at https://github.com/ming053l/ELSA.

Chih-Chung Hsu, Xin-Di Ma, Wo-Ting Liao +1

Advanced Computer Vision Laboratory, National Yang Ming Chiao Tung University

Efficient Transformer Inference Softmax Attention GPU Acceleration Self-Attention High-Performance Computing Vision Transformer+3

Apr 9, 2026cs.LG

Sinkhorn doubly stochastic attention rank decay analysis

The self-attention mechanism is central to the success of Transformer architectures. However, standard row-stochastic attention has been shown to suffer from significant signal degradation across layers. In particular, it can induce rank collapse, resulting in increasingly uniform token representations, as well as entropy collapse, characterized by highly concentrated attention distributions. Recent work has highlighted the benefits of doubly stochastic attention as a form of entropy regularization, promoting a more balanced attention distribution and leading to improved empirical performance. In this paper, we study rank collapse across network depth and show that doubly stochastic attention matrices normalized with Sinkhorn algorithm preserve rank more effectively than standard softmax row-stochastic ones. As previously shown for softmax, skip connections are crucial to mitigate rank collapse. We empirically validate this phenomenon on both sentiment analysis and image classification tasks. Moreover, we derive a theoretical bound for the pure self-attention rank decay when using Sinkhorn normalization and find that rank decays to one doubly exponentially with depth, a phenomenon that has already been shown for softmax.

Michela Lapenna, Rita Fioresi, Bahman Gharesifard

Department of Physics and Astronomy, University of Bologna, Bologna, Italy · Department of Pharmacy and Biotechnologies, University of Bologna Bologna, Italy · Department of Mathematics and Statistics, Queen’s University, Kingston, ON, Canada

Softmax Attention Self-Attention Transformer Attention Attention Mechanisms

Feb 9, 2026cs.LG

Entropy-Generated Attention Beyond Softmax and Entmax: Kaniadakis and Reciprocal-Symmetric Abe Operators

We derive two attention operators from generalized statistical entropies. Kaniadakis entropy yields an exact full-support normalization whose weights and low-score sensitivities decay algebraically, rather than exponentially as in Softmax or by exact truncation as in entmax. Classical Abe entropy yields an implicit reciprocal-symmetric operator. With

q=e^ε

, the involution

q\leftrightarrow q^{-1}

removes every odd correction about Softmax; we obtain the normalized second- and fourth-order terms, including the deformation of the normalization multiplier. These stationary laws follow from a Fisher-metric Lagrangian on the probability simplex, whose Shannon sector recovers scaled dot-product Softmax. We also give a tangent-gradient test for deciding whether changing the entropy changes the attention profile or only its scale. Rényi and two-parameter Sharma--Mittal entropies retain the Tsallis--entmax inverse-gradient shape, but their global moments make the effective temperature input dependent when the external temperature is fixed. Distinguishing profile-shape equivalence from fixed-parameter operator equivalence separates new normalization shapes from adaptive rescalings and organizes the operators by support, tail behavior, and realization complexity.

Gunn Kim

Department of Physics, Sejong University, Seoul 05006, Republic of Korea

Softmax Attention

Nov 21, 2025cs.CL

Selective Rotary Position Embedding

Position information is essential for language modeling. In softmax transformers, Rotary Position Embeddings (\textit{RoPE}) encode positions through \textit{fixed-angle} rotations, while in linear transformers, order is handled via input-dependent (selective) gating that decays past key-value associations. Selectivity has generally been shown to improve language-related tasks. Inspired by this, we introduce \textit{Selective RoPE}, an \textit{input-dependent} rotary embedding mechanism, that generalizes \textit{RoPE}, and enables rotation in \textit{arbitrary angles} for both linear and softmax transformers. We show that softmax attention already performs a hidden form of these rotations on query-key pairs, uncovering an implicit positional structure. We further show that in state-space models and gated linear transformers, the real part manages forgetting while the imaginary part encodes positions through rotations. We validate our method by equipping gated transformers with \textit{Selective RoPE}, demonstrating that its input-dependent rotations improve performance in language modeling and on difficult sequence tasks like copying, state tracking, and retrieval.

Sajad Movahedi, Timur Carstensen, Arshia Afzal +3

ELLIS Institute Tübingen1, LIONS, EPFL2, University of Freiburg3, Max-Planck-Institute for Intelligent Systems4, Prior Labs5

Softmax Attention Rotary Positional Embeddings Autoregressive Language Modeling Linear Attention Attention Mechanisms

Oct 4, 2025cs.LG

Performance-Efficiency Tradeoffs in Transformers: An Approximation Theory Perspective

Transformers have achieved remarkable successes across a wide range of applications, yet the theoretical foundation of their model efficiency remains underexplored. In this work, we investigate how the model parameters -- mainly attention heads and head dimensions -- should be allocated across layers to balance expressivity and efficiency. We first provide mathematical analysis on the role of early layers in information extraction from an approximation perspective, with a theoretical characterization on the trade-off between the number of heads and head dimension under a fixed parameter budget. In addition, we uncover and prove the \emph{saturation} behavior of softmax activations: Continuously increasing head dimensions can lead to diminishing returns in learning errors, particularly for long sequences. Supported by both theory and experiments, this saturation pattern suggests that later layers can operate more efficiently with reduced parameters. Combining these insights, we propose principled strategies for allocating attention heads and dimensions across Transformers' layers, shedding light on theoretically-grounded model efficiency of Transformer-based architectures.

Ruoxi Yu, Haotian Jiang, Jingpu Cheng +3

Center for Data Science, Peking University · Institute for Functional Intelligent Materials, National University of Singapore · Department of Mathematics, National University of Singapore +1

Softmax Attention Attention Head Analysis Transformer Attention Mechanisms

Sep 12, 2025cs.LG

Multipole Semantic Attention: A Fast Approximation of Softmax Attention for Pretraining

Pretraining transformers on long sequences (entire code repositories, collections of related documents) is bottlenecked by quadratic attention costs. We present Multipole Semantic Attention (MuSe), which accelerates 64k-context pretraining by 36% while matching baseline loss, requiring no architectural changes. MuSe clusters queries and keys separately in representation space. This yields query-specific summaries that substantially outperform spatial blocking at matched sparsity, while also enabling drop-in compatibility with existing pretrained models; we validate on Llama 3.1-8B and 3.2-1B without retraining. We pretrain language models up to 1B parameters at 64k context on code and scientific documents, confirming that MuSe preserves quality and long-context utilization during training.

Rupert Mitchell, Kristian Kersting

Department of Computer Science, TU Darmstadt, Darmstadt, Germany · Hessian Center for AI (hessian.AI), Darmstadt, Germany · German Research Center for Artificial Intelligence (DFKI), Darmstadt, Germany +1

Language Model Pretraining Softmax Attention Self-Attention Long-Context Language Modeling Transformer Attention Sparse Attention+1

Date pendingcs.LG

Attention by Synchronization in Coupled Oscillator Networks

We address transformer attention on energy-constrained physical substrates. Softmax attention requires exponentiation and global reduction, operations with high energy cost on von Neumann hardware and no natural physical analog. We show that Kuramoto synchronization dynamics (which arise in electrical, mechanical, superconducting, and charge-density-wave oscillator arrays, among other physical systems) implement a well-defined attention operation. The resulting mechanism, \emph{fixed-query oscillator attention}, replaces softmax's arithmetic with the equilibration of a gradient flow on the sphere: queries are learned anchors fixed on the sphere, and free oscillators evolve under Kuramoto--Lohe dynamics until they settle at positions encoding attention weights via cosine similarity. Because the computation is equilibration, no global exponential normalization is needed. The fixed point is provably unique and globally attractive from almost every initial condition, a guarantee that holds across every physical realization. Empirically, at the minimal hardware configuration (oscillator dimension

d_{\mathrm{osc}} = 2

), oscillator attention matches softmax on keyword spotting, and on subject-verb agreement it trains more reliably while reaching softmax's accuracy. Softmax retains an advantage on causal language modeling, but the gap decays as a power law in

d_{\mathrm{osc}}

. The main objective of this work is not to replace softmax in software but to provide a mathematically grounded blueprint for accurate attention on physical substrates.

Fabio Pasqualetti, Taosha Guo

Electrical Engineering and Computer Science University of California, Irvine

Softmax Attention Dynamical Systems Transformer Attention Attention Mechanisms