Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Softmax Linear Attention: Reclaiming Global Competition

Mingwei Xu, Xuan Lin, Xinnan Guo, Wanqing Xu, Wanyun Cui

TL;DR

Softmax Linear Attention (SLA) addresses the expressivity gap between linear and full attention by introducing inter-head competition. It preserves linear time complexity $O(L)$ while reintroducing global selectivity through head-wise softmax gates, leveraging the multi-head structure as coarse semantic slots. The authors prove that SLA restores magnitude sensitivity and asymptotic winner-take-all dynamics, and demonstrate consistent improvements over state-of-the-art linear baselines (RetNet, GLA, GDN) across retrieval, language modeling, and zero-shot reasoning in long-context settings. Empirically, SLA yields sharper focus and robustness with only modest parameter and compute overhead, offering a practical path to scalable, precise long-context modeling.

Abstract

While linear attention reduces the quadratic complexity of standard Transformers to linear time, it often lags behind in expressivity due to the removal of softmax normalization. This omission eliminates \emph{global competition}, a critical mechanism that enables models to sharply focus on relevant information amidst long-context noise. In this work, we propose \textbf{Softmax Linear Attention (SLA)}, a framework designed to restore this competitive selection without sacrificing efficiency. By lifting the softmax operation from the token level to the head level, SLA leverages attention heads as coarse semantic slots, applying a competitive gating mechanism to dynamically select the most relevant subspaces. This reintroduces the ``winner-take-all'' dynamics essential for precise retrieval and robust long-context understanding. Distinct from prior methods that focus on refining local kernel functions, SLA adopts a broader perspective by exploiting the higher-level multi-head aggregation structure. Extensive experiments demonstrate that SLA consistently enhances state-of-the-art linear baselines (RetNet, GLA, GDN) across language modeling and long-context benchmarks, particularly in challenging retrieval scenarios where it significantly boosts robustness against noise, validating its capability to restore precise focus while maintaining linear complexity.

Softmax Linear Attention: Reclaiming Global Competition

TL;DR

Softmax Linear Attention (SLA) addresses the expressivity gap between linear and full attention by introducing inter-head competition. It preserves linear time complexity while reintroducing global selectivity through head-wise softmax gates, leveraging the multi-head structure as coarse semantic slots. The authors prove that SLA restores magnitude sensitivity and asymptotic winner-take-all dynamics, and demonstrate consistent improvements over state-of-the-art linear baselines (RetNet, GLA, GDN) across retrieval, language modeling, and zero-shot reasoning in long-context settings. Empirically, SLA yields sharper focus and robustness with only modest parameter and compute overhead, offering a practical path to scalable, precise long-context modeling.

Abstract

While linear attention reduces the quadratic complexity of standard Transformers to linear time, it often lags behind in expressivity due to the removal of softmax normalization. This omission eliminates \emph{global competition}, a critical mechanism that enables models to sharply focus on relevant information amidst long-context noise. In this work, we propose \textbf{Softmax Linear Attention (SLA)}, a framework designed to restore this competitive selection without sacrificing efficiency. By lifting the softmax operation from the token level to the head level, SLA leverages attention heads as coarse semantic slots, applying a competitive gating mechanism to dynamically select the most relevant subspaces. This reintroduces the ``winner-take-all'' dynamics essential for precise retrieval and robust long-context understanding. Distinct from prior methods that focus on refining local kernel functions, SLA adopts a broader perspective by exploiting the higher-level multi-head aggregation structure. Extensive experiments demonstrate that SLA consistently enhances state-of-the-art linear baselines (RetNet, GLA, GDN) across language modeling and long-context benchmarks, particularly in challenging retrieval scenarios where it significantly boosts robustness against noise, validating its capability to restore precise focus while maintaining linear complexity.
Paper Structure (37 sections, 3 theorems, 14 equations, 4 figures, 5 tables)

This paper contains 37 sections, 3 theorems, 14 equations, 4 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Efficient Transformers and Linear Attention
  4. The Cost of Removing Softmax: Expressivity Bottlenecks
  5. Multi-Head Mechanism and Feature Diversity
  6. Method
  7. The Competition Gap in Linear Attention
  8. Softmax Linear Attention
  9. Recurrent Implementation.
  10. Chunkwise Parallel Training.
  11. Parameter Efficiency.
  12. Theoretical Analysis
  13. Restoring Magnitude Sensitivity via Head Competition
  14. Asymptotic Winner-Take-All Dynamics
  15. Experiments
  16. ...and 22 more sections

Key Result

Proposition 4.1

Let $w_{\text{lin}}(k) = \frac{\phi(q)^\top \phi(k)}{\sum_j \phi(q)^\top \phi(k_j)}$ be the normalized attention weight for key $k$ in linear attention. If $\phi(\cdot)$ is homogeneous (e.g., ReLU), then $w_{\text{lin}}(k)$ is invariant to scalar scaling of $q$. Consequently, the entropy of the atte

Figures (4)

  • Figure 1: Scaling curves on WikiText perplexity for GLA and Softmax GLA.
  • Figure 2: WikiText perplexity under different numbers of attention heads.
  • Figure 3: Overall comparison on training throughput, memory footprint and inference latency.
  • Figure 4: Softmax RetNet training loss curves under the same setup in Section \ref{['subsec:setup']}.

Theorems & Definitions (4)

  • Proposition 4.1
  • Theorem 4.2
  • Theorem 4.3
  • proof