SSA: 特徴量空間における完全注意とスパース注意の出力を整列させるスパーススパース注意機構

要旨

大規模言語モデル（LLM）における完全注意機構の二次計算量は、長文脈処理の効率性を制限する。疎注意機構は各クエリが関与する過去トークンのサブセットを制限することでこのコストを軽減するが、訓練不要な手法では性能の大幅な低下が生じやすい。ネイティブ疎注意手法（NSA、MoBAなど）はこの問題を緩和するものの、重要なパラドックスが存在する：完全注意を近似することを目的としながら、かえって完全注意モデルよりも注意の疎性が低下し、有効性が制限される可能性がある。我々はこのパラドックスの原因を勾配更新の欠陥に帰する：疎訓練中に除外される低ランクのキー・バリューペアは順方向の寄与も逆方向の勾配も受けず、適切な抑制を学習できない。この制限を克服するため、本論文ではSSA（Sparse Sparse Attention）を提案する。これは疎注意と完全注意の両方を考慮し、層ごとに双方向のアライメントを強化する統一的訓練フレームワークである。この設計により、全てのトークンへの勾配流れを維持しつつ、疎注意出力が完全注意出力に整合することを明示的に促進し、より強力な疎性を実現する。その結果、SSAは複数の常識ベンチマークにおいて、疎注意と完全注意の両推論で最先端の性能を達成した。さらにSSAはモデルが様々な疎性バジェットに適応することを可能にし、関与トークン数が増えるにつれて性能が一貫して向上し、推論時の計算量と性能の柔軟なトレードオフを実現する。最後に、ネイティブ疎注意訓練がシンク領域への注意値の過剰割り当てを緩和することで、長文脈外挿能力を向上させることを示し、SSAが最も強力な外挿能力を発揮することを確認する。

English

The quadratic complexity of full attention limits efficient long-context processing in large language models (LLMs). Sparse attention mitigates this cost by restricting each query to attend to a subset of previous tokens; however, training-free approaches often lead to severe performance degradation. Native sparse-attention methods (e.g., NSA, MoBA) alleviate this issue, yet exhibit a critical paradox: they produce lower attention sparsity than full-attention models, despite aiming to approximate full attention, which may constrain their effectiveness. We attribute this paradox to gradient update deficiency: low-ranked key-value pairs excluded during sparse training receive neither forward contribution nor backward gradients, and thus never learn proper suppression. To overcome this limitation, we propose SSA (Sparse Sparse Attention), a unified training framework that considers both sparse and full attention and enforces bidirectional alignment at every layer. This design preserves gradient flow to all tokens while explicitly encouraging sparse-attention outputs to align with their full-attention counterparts, thereby promoting stronger sparsity. As a result, SSA achieves state-of-the-art performance under both sparse and full attention inference across multiple commonsense benchmarks. Furthermore, SSA enables models to adapt smoothly to varying sparsity budgets; performance improves consistently as more tokens are allowed to attend, supporting flexible compute-performance trade-offs at inference time. Finally, we show that native sparse-attention training surprisingly improves long-context extrapolation by mitigating the over-allocation of attention values in sink areas, with SSA demonstrating the strongest extrapolation capability.

SSA: 特徴量空間における完全注意とスパース注意の出力を整列させるスパーススパース注意機構

SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space

要旨

Support