SSA: Attenzione Sparsa Sparsa mediante Allineamento degli Output di Attenzione Completa e Sparsa nello Spazio delle Caratteristiche

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

November 25, 2025

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

cs.AI

Abstract

La complessità quadratica dell'attenzione completa limita l'elaborazione efficiente di contesti lunghi nei grandi modelli linguistici (LLM). L'attenzione sparsa mitiga questo costo limitando ogni query a considerare un sottoinsieme di token precedenti; tuttavia, approcci senza addestramento spesso portano a un grave degrado delle prestazioni. I metodi di attenzione sparsa nativa (ad es., NSA, MoBA) alleviano questo problema, ma presentano un paradosso critico: producono una minore sparsità dell'attenzione rispetto ai modelli a attenzione completa, nonostante mirino ad approssimarla, il che può limitarne l'efficacia. Attribuiamo questo paradosso a una carenza negli aggiornamenti del gradiente: le coppie chiave-valore a basso rango escluse durante l'addestramento sparso non ricevono né contributo in avanti né gradienti all'indietro, e quindi non apprendono mai una soppressione appropriata. Per superare questa limitazione, proponiamo SSA (Sparse Sparse Attention), un framework di addestramento unificato che considera sia l'attenzione sparsa che quella completa e applica un allineamento bidirezionale a ogni livello. Questo progetto preserva il flusso del gradiente verso tutti i token incoraggiando esplicitamente gli output dell'attenzione sparsa ad allinearsi con le loro controparti a attenzione completa, promuovendo così una sparsità più forte. Di conseguenza, SSA raggiunge prestazioni allo stato dell'arte nell'inferenza sia con attenzione sparsa che completa su molteplici benchmark di commonsense. Inoltre, SSA consente ai modelli di adattarsi agevolmente a budget di sparsità variabili; le prestazioni migliorano costantemente man mano che più token possono partecipare, supportando compromessi flessibili tra capacità di calcolo e prestazioni al momento dell'inferenza. Infine, mostriamo che l'addestramento con attenzione sparsa nativa migliora sorprendentemente l'estrapolazione di contesti lunghi mitigando l'over-allocation dei valori di attenzione nelle aree sink, con SSA che dimostra la capacità di estrapolazione più forte.

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.