SSA: Atención Dispersa Dispersa mediante la Alineación de las Salidas de Atención Completa y Dispersa en el Espacio de Características

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

November 25, 2025

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

cs.AI

Resumen

La complejidad cuadrática de la atención completa limita el procesamiento eficiente de contextos largos en los modelos de lenguaje grandes (LLM). La atención dispersa mitiga este costo al restringir cada consulta a atender a un subconjunto de tokens anteriores; sin embargo, los enfoques que no requieren entrenamiento a menudo conducen a una degradación severa del rendimiento. Los métodos nativos de atención dispersa (por ejemplo, NSA, MoBA) alivian este problema, pero exhiben una paradoja crítica: producen una menor dispersión de atención que los modelos de atención completa, a pesar de tener como objetivo aproximarse a la atención completa, lo que puede limitar su efectividad. Atribuimos esta paradoja a una deficiencia en la actualización del gradiente: los pares clave-valor de bajo rango excluidos durante el entrenamiento disperso no reciben contribución directa ni gradientes hacia atrás y, por lo tanto, nunca aprenden la supresión adecuada. Para superar esta limitación, proponemos SSA (Atención Dispersa Dispersa), un marco de entrenamiento unificado que considera tanto la atención dispersa como la completa y aplica una alineación bidireccional en cada capa. Este diseño preserva el flujo de gradientes hacia todos los tokens mientras fomenta explícitamente que las salidas de atención dispersa se alineen con sus contrapartes de atención completa, promoviendo así una mayor dispersión. Como resultado, SSA logra un rendimiento de vanguardia bajo inferencia tanto de atención dispersa como completa en múltiples benchmarks de sentido común. Además, SSA permite a los modelos adaptarse suavemente a diferentes presupuestos de dispersión; el rendimiento mejora consistentemente a medida que se permite atender a más tokens, respaldando compensaciones flexibles entre computación y rendimiento durante la inferencia. Finalmente, demostramos que el entrenamiento nativo de atención dispersa mejora sorprendentemente la extrapolación de contextos largos al mitigar la sobreasignación de valores de atención en las áreas "sumidero", siendo SSA el que demuestra la capacidad de extrapolación más sólida.

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.