[論文レビュー] 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）は線形注意にヘッドレベルのソフトマックス競合を再導入し、セマンティックヘッド間の winner-take-all 型の選択性を実現しつつ、線形時間・線形メモリを維持します。複数の線形ベースラインに対して長い文脈タスクでの検索信頼性と頑健性を向上させます。

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.

研究の動機と目的

• 線形注意からソフトマックスを除去することによって生じる表現力のギャップ（Magnitude Neglect、Context Collapse）を識別する。
• 線形計算量を維持しつつヘッド間競合を導入するSLAを提案する。
• Magnitude感度の回復と漸近的な winner-take-all 动態を理論的に分析する。
• SLAの有効性を最先端の線形ベースライン（RetNet、GLA、GDN）に適用し、言語モデリングと長文脈タスクで示す。

提案手法

• QとKに対してヘッドレベルのソフトマックスゲートを追加し、ヘッド間競合を作成してマルチヘッド集約を再定義する。
• 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 を用いて計算する。
• 再帰的およびチャンク単位の訓練実装を提供して線形計算量を維持する。
• 軽量なパラメータ追加（レイヤーごとに2つの射影行列）を提供し、オーバーヘッドは極小。
• Magnitude感度の回復と漸近的な winner-take-all 挙動を示す理論的結果を提示する。

実験結果

リサーチクエスチョン

• RQ1 ヘッドレベルのソフトマックス競合は線形注意で失われたグローバルな選択性を回復できるのか？
• RQ2 SLAは線形時間・空間計算量を維持しつつ鋭い注意割り当てを提供できるのか？
• RQ3 SLA搭載の線形ベースライン（RetNet、GLA、GDN）はより良い検索と長文脈性能を達成するのか？
• RQ4 SLAにおける magnitude感度と winner-take-all動作に理論的裏付けがあるのか？
• RQ5 モデルサイズを跨いだ訓練/推論の効率とスケーラビリティにSLAはどう影響するのか？

主な発見

• SLAはクエリ/キーの大きさに反応するヘッドゲートを作動させることで magnitude感度を回復し、自信を持った鋭いフォーカシングを可能にする。
• モデルの信頼度が高まるにつれてSLAのヘッドゲートは単一のヘッドに集中し、one-hot競合に近い挙動をとる。
• 実証的な結果は、実世界タスクでの検索精度が基準線の線形モデル（例：Softmax-GLA、Softmax-RetNet、Softmax-GDN）より向上することを示す。
• 複数の長文脈ベンチマークで、SLAは線形ベースラインの性能を一貫して向上させ、完全なソフトマックス変換器とのギャップを縮小する。
• アブレーションによりヘッド数を増やすほどSLAの利得が大きくなることを示し、意味領域競合仮説を裏付ける。
• 訓練と推論には控えめなオーバーヘッドで、スループットを維持し、メモリ使用を拡張可能に保つ。