[论文解读] Softmax Linear Attention: Reclaiming Global Competition

Xu, Mingwei, Xuan Lin+3 more|arXiv (Cornell University)|Feb 2, 2026

Topic Modeling被引用 0

一句话总结

Softmax Linear Attention（SLA）通过头级 softmax 门重新引入线性注意力中的全局竞争，在保留线性时间与内存的同时实现语义头之间的赢家通吃式选择性。它在跨多种线性基线的长上下文任务中提升了检索可靠性与鲁棒性。

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 extbf{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 引发的表达能力缺口（数量感知 Magnitude Neglect、上下文坍缩 Context Collapse）来识别问题点。
• 提出 SLA，在保持线性复杂度的同时引入跨头竞争。
• 理论分析幅度敏感性恢复与渐近式赢家通吃动态。
• 通过将 SLA 应用于最先进的线性基线（RetNet、GLA、GDN），在语言建模与长上下文任务上展示其有效性。

提出的方法

• 通过在 Q、K 上添加头级 softmax 门来重新定义多头聚合，以创建跨头竞争。
• 将 SLA 输出表示为 O_SLA = Concat_h ((G^Q_h ⊙ φ(Q_h)) (G^K_h ⊙ φ(K_h))^T V_h) W^O。
• 计算 G^Q_h = softmax(Q W_GQ)_h 和 G^K_h = softmax(K W_GK)_h，使用低秩头投影 W_GQ、W_GK。
• 提供递归和分块训练实现，以维持线性复杂度。
• 提供轻量级参数增加（每层两个投影矩阵），开销可忽略。
• 给出理论结果，展示幅度敏感性恢复与渐近式赢家通吃行为。

实验结果

研究问题

• RQ1 头级 softmax 竞争是否能够恢复线性注意力中丢失的全局选择性？
• RQ2 SLA 在提供明显注意力分配的同时，是否保持线性时间/空间复杂度？
• RQ3 配备 SLA 的线性基线（RetNet、GLA、GDN）在检索和长上下文性能上是否更优？
• RQ4 是否存在对 SLA 的幅度敏感性与赢家通吃动态的理论支撑？
• RQ5 SLA 对不同模型规模的训练/推理效率与可扩展性有何影响？

主要发现

• SLA 通过使头门响应查询/键的幅度来恢复幅度敏感性，从而实现自信且尖锐的聚焦。
• 随着模型置信度的提升，SLA 的头门聚焦于单一头，近似实现单热（one-hot）竞争。
• 实证结果显示，与基线线性模型（如 Softmax-GLA、Softmax-RetNet、Softmax-GDN）相比，SLA 提高了实际任务的检索准确性。
• 在多个长上下文基准上，SLA 持续提升线性基线的性能，缩小了与全 Softmax 变换器之间的差距。
• 消融实验表明头数越多（H），SLA 的收益越大，验证了语义槽位竞争假说。
• 训练和推理开销适度，吞吐量维持，内存使用可扩展。