Softmax Linear Attention: Reclaiming Global Competition
This work addresses the problem of reduced expressivity in efficient linear attention models for machine learning practitioners, offering an incremental improvement by enhancing existing methods without sacrificing linear complexity.
The paper tackles the expressivity gap in linear attention Transformers by proposing Softmax Linear Attention (SLA), which restores global competition through a head-level softmax mechanism, resulting in consistent improvements over state-of-the-art linear baselines in language modeling and long-context benchmarks, with significant boosts in retrieval robustness against noise.
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.