Softmax Linear Attention

• Softmax Linear Attention (SLA) is a hybrid attention mechanism that blends softmax and linear methods to optimize computational efficiency and maintain expressive power.
• It employs learned gating, injectivity-restoring kernel designs, and residual augmentations to overcome limitations of pure linear or softmax approaches.
• Empirical evaluations demonstrate SLA's capability to nearly match softmax performance with significantly lower compute, benefiting long-context language and vision tasks.

Softmax Linear Attention (SLA) encompasses a range of hybrid and refined attention mechanisms that unify the computational efficiency of linear attention with the expressivity and inductivity properties of softmax-based attention. SLA research converges on architectures, theoretical justifications, and practical implementations that address the fundamental limitations of both paradigms, typically by gating, hybridization, or crafting injective and locally-biased kernelizations.

1. Core Principles and Mathematical Formulation

SLA mechanisms are founded on the dichotomy between the standard dot-product softmax attention and linear kernel-based attention. Standard softmax attention, for an input sequence (length $L$, embedding dimension $d$) with queries $Q \in \mathbb{R}^{L \times d}$, keys $K$, and values $V$, forms an output as

$O^\mathrm{soft} = \mathrm{Softmax}(QK^\top) V,$

with quadratic complexity in sequence length. In canonical linear attention, the exponential kernel is replaced by an inner product of feature maps, yielding

$O^\mathrm{lin} = \phi(Q) [\phi(K)^\top V],$

for a suitable feature map $\phi$. This form allows associative computation for linear time and space complexity.

SLA frameworks blend these by employing gating mechanisms or hybridization strategies, so an output may be formed as a (possibly normalized) convex combination: $$O_i = g_i O^{\mathrm{soft}}_i + (1 - g_i) O^{\mathrm{lin}}_i,$$ where $g_i \in [0,1]$ is a learned gate at the token or chunk level (Deng et al., 3 Feb 2026), or, in vision models, injectivity and locality are restored by careful feature map design and residual augmentations (Han et al., 2024).

Table: Key Operations in SLA Layers

Variant	Output Computation	Complexity
Pure Softmax	$O^\mathrm{soft} = \mathrm{Softmax}(QK^\top)V$	$O(L^2 d)$
Pure Linear	$O^\mathrm{lin} = \phi(Q)[\phi(K)^\top V]$	$O(L dr)$
SLA (token-hybrid)	$$O_i = g_i O^{\mathrm{soft}}_i + (1 - g_i) O^{\mathrm{lin}}_i$$	$O(L_\mathrm{soft} L d + L_\mathrm{lin} dr)$

Hybridization can be intra-layer (Deng et al., 3 Feb 2026, Han et al., 2024, Meng et al., 2 Feb 2026), across layers (Li et al., 16 Jan 2026), or at the head level (implementing softmax among heads per token) (Xu et al., 2 Feb 2026).

2. Theoretical Properties: Injectivity, Locality, and Expressivity

A central theoretical insight is that softmax attention is injective in its query map for generic key matrices, while canonical linear attention is not: there exist $p \neq q$ with $L_K(p) = L_K(q)$ for any continuous feature map $\phi$, unless specific modifications are made (Han et al., 2024). Non-injectivity of linear attention leads to "semantic confusion," undermining fine-grained retrieval and token discrimination.

Softmax attention also exhibits strong local modeling bias, critical in vision transformers and sequence domains. SLA-type mechanisms, such as InLine, address these by introducing subtraction-normalized feature maps (restoring injectivity) and local window residuals (restoring locality) (Han et al., 2024).

Another theoretical axis is degree-of-freedom characterization and approximation regimes: SLA approaches can use statistical degrees of freedom (DoF) to set the dimensionality of the feature map per layer to minimize kernel approximation error subject to compute constraints (Nishikawa et al., 4 Jul 2025). Furthermore, in the large-prompt regime, softmax attention empirically and theoretically converges to its linearized counterpart, with quantifiable non-asymptotic concentration bounds (Boursier et al., 12 Dec 2025).

3. SLA Architecture Variants

Multiple instantiations of SLA exist:

• Gated Hybrid SLA ((Deng et al., 3 Feb 2026), NAtS-L): At each chunk or token, a gating network assigns a score that determines whether softmax or linear attention is applied. Outputs from both branches are normalized and merged via the gate.
• Injectivity & Locality-Augmented SLA ((Han et al., 2024), InLine): Subtraction-normalized linear kernels enforce injectivity; a local window residual restores the inductive local bias characteristic of softmax.
• Head-wise Softmax Linear Attention (Xu et al., 2 Feb 2026): "Global competition" is restored by applying the softmax operator over the head dimension rather than tokens, yielding a $O(L)$ scheme that mimics winner-take-all behavior without tokenwise softmax normalization.
• Agent Attention as Unified Framework (Han et al., 2023): By mediating attention via a set of agent tokens and two softmax operations (query$\to$agent, agent$\to$key), agent attention interpolates between softmax (quadratic) and linear (pure kernel) as special cases.
• Norm-Preserved and MLP-learned Kernelization (Zhang et al., 2024, Meng et al., 2 Feb 2026): Linear attention is enhanced via learned MLP feature maps that promote spikiness and dot-product monotonicity (Hedgehog), or by preserving pretrained norms to maintain distributional consistency in hybrid layers (STILL).

These architectural choices determine the complexity, expressivity, and robustness characteristics of SLA models.

4. Empirical Performance and Trade-offs

Empirical studies show that SLA models, with properly tuned gates or hybridization schedules, can recover nearly all of the performance of full softmax attention in both language and vision tasks, while offering significant efficiency gains (Deng et al., 3 Feb 2026, Han et al., 2024, Zhang et al., 2024).

• Token-level hybrid SLA matches or exceeds softmax-only perplexity on long-context retrieval tasks with only $1.7\times$ the compute of pure linear variants, compared to $\sim 5\times$ for softmax-only layers, and achieves up to $2$–$5\times$ speedup in decoding throughput for $128$k-token contexts (Deng et al., 3 Feb 2026).
• SLA in vision (InLine) closes or overcomes the softmax-linear gap on ImageNet-1K, with InLine-Swin-T achieving 82.4% top-1 accuracy at the same FLOPs as softmax-Swin-T (81.3%) (Han et al., 2024), and SoLA-Vision yielding competitive accuracy with layerwise fraction of softmax layers (Li et al., 16 Jan 2026).
• Hedgehog achieves $99$% recovery of softmax performance after finetuned- or pretrained-conversion in both language and vision settings (Zhang et al., 2024).
• Chunk-wise routing (STILL) enables linearization of full LLMs while retaining reasoning and long-context retrieval (e.g., 86.2% RULER S-NIAH-1 recovery) and scaling to $64$k tokens with flat memory (Meng et al., 2 Feb 2026).

Empirically, the fraction of tokens requiring softmax attention in NAtS-L is minimized by the differentiable architecture search, focusing expensive computation only on tokens needed for long-range retrieval (Deng et al., 3 Feb 2026). In vision, the combination of injectivity and local bias is essential to go beyond standard linear kernelization (Han et al., 2024).

5. Limitations and Boundary of SLA Utility

Despite closing the gap, certain theoretical and empirical limitations persist:

• Expressivity: Softmax attention remains strictly more expressive than linear attention on tasks requiring precise, one-hot, or globally-competitive selection, as shown in statistical and separation-theorem analyses (Deng et al., 2023, Duranthon et al., 26 Sep 2025).
• Non-injectivity: Unless modified, linear attention remains non-injective, making it vulnerable to semantic collisions and reduced discriminative power (Han et al., 2024).
• Task specialization: In retrieval or single-location regression, softmax attention achieves Bayes-optimality while linear intrinsicly falls short, with the performance gap exponential in signal strength (Duranthon et al., 26 Sep 2025).
• SLA approximation in large-prompt regime: Measure-theoretic results justify using linear-analytic dynamics for softmax attention only when prompt lengths are sufficiently large, as quantified by explicit non-asymptotic error bounds (Boursier et al., 12 Dec 2025).

A plausible implication is that while SLA is suitable for massively long sequences and contexts where only a small fraction of tokens truly require global focus, modelers must quantify the performance–efficiency trade-off in task-specific settings.

6. Implementation Variants and Adaptation Strategies

Recent SLA methodologies offer various approaches for implementation:

• Layerwise assignments: Layer-by-layer scheduling of softmax and linear layers (SoLA-Vision) enables global context injection with sparse softmax layers (Li et al., 16 Jan 2026).
• Per-token routing: Per-token/ per-chunk gates (NAtS-L, STILL) allow fine-grained control, adjusted via differentiable neural architecture search (Deng et al., 3 Feb 2026, Meng et al., 2 Feb 2026).
• Learned MLP kernels/distillation: Hedgehog and DoF-based SLA distillation use supervised or unsupervised distillation from softmax attention to tailor strong feature maps for linearization, often layerwise and at fixed or allocated feature budget (Nishikawa et al., 4 Jul 2025, Zhang et al., 2024).
• Agent-based architectures: Agent Attention leverages an intermediate set of tokens to mediate global context and aggregation, reducing effective complexity as a function of agent-token count (Han et al., 2023).

Implementation details, such as chunk size, gate parameterization, and feature map initialization, influence performance and resource utilization.

7. Outlook and Design Guidelines

Research on SLA suggests several guiding principles for hybrid and linearized Transformer design:

• Allocate softmax attention only where necessary, using learning-based gates or contextually-aware scoring (Deng et al., 3 Feb 2026, Meng et al., 2 Feb 2026).
• Preserve or restore injectivity and locality by appropriate feature-map and residual design (subtraction normalization, local window augmentation) (Han et al., 2024).
• Distill softmax behavior into linear mechanisms using supervised kernel alignment or attention-weight mimicry, with dimensionality set to match per-layer effective complexity (Nishikawa et al., 4 Jul 2025, Zhang et al., 2024).
• For hierarchical or vision models, insert softmax layers sparsely after linear stacking; additive returns saturate rapidly, and early layers can remain fully linear (Li et al., 16 Jan 2026).
• Use agent or head-wise global competition to approximate winner-take-all properties at manageable (often $O(L)$) cost (Xu et al., 2 Feb 2026, Han et al., 2023).
• When context length is extreme, exploit measure-theoretic convergence of softmax to linear attention, but monitor the empirical regime for possible expressed limitations (Boursier et al., 12 Dec 2025).

The development of SLA continues to be driven by the dual imperatives of efficiency (scaling to very long contexts or high-resolution spatial representations) and expressivity (retaining global discrimination, local modeling, and injectivity). The field remains active, with leading directions including adaptive granularity of competition, learnable clustering for normalization, and further convergence of kernel-based and hybrid attention paradigms.