SSA : Attention Sparse Sparse par Alignement des Sorties de l'Attention Complète et de l'Attention Sparse dans l'Espace des Caractéristiques

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

November 25, 2025

papers.authors: Zhenyi Shen, Junru Lu, Lin Gui, Jiazheng Li, Yulan He, Di Yin, Xing Sun

cs.AI

papers.abstract

La complexité quadratique de l'attention complète limite le traitement efficace des contextes longs dans les grands modèles de langage (LLM). L'attention creuse atténue ce coût en restreignant chaque requête à un sous-ensemble de tokens précédents ; cependant, les approches non entraînées conduisent souvent à une dégradation sévère des performances. Les méthodes d'attention creuse native (par exemple, NSA, MoBA) pallient ce problème, mais présentent un paradoxe critique : elles produisent une sparsité d'attention inférieure à celle des modèles à attention complète, bien qu'elles visent à l'approximer, ce qui peut limiter leur efficacité. Nous attribuons ce paradoxe à une déficience des mises à jour de gradient : les paires clé-valeur de faible rang exclues pendant l'entraînement creux ne reçoivent ni contribution avant ni gradients arrière, et n'apprennent donc jamais une suppression appropriée. Pour surmonter cette limitation, nous proposons SSA (Sparse Sparse Attention), un cadre d'entraînement unifié qui considère à la fois l'attention creuse et complète et impose un alignement bidirectionnel à chaque couche. Cette conception préserve le flux de gradient vers tous les tokens tout en encourageant explicitement les sorties de l'attention creuse à s'aligner sur leurs équivalents en attention complète, favorisant ainsi une sparsité plus forte. En conséquence, SSA atteint des performances de pointe en inférence creuse et complète sur plusieurs benchmarks de bon sens. De plus, SSA permet aux modèles de s'adapter fluidement à différents budgets de sparsité ; les performances s'améliorent constamment à mesure que plus de tokens peuvent participer, permettant des compromis flexibles entre calcul et performance lors de l'inférence. Enfin, nous montrons que l'entraînement par attention creuse native améliore de manière surprenante l'extrapolation en contexte long en atténuant la sur-allocation des valeurs d'attention dans les zones "sink", SSA démontrant la plus forte capacité d'extrapolation.

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.