SSA: Atenção Esparsa Esparsa por meio do Alinhamento de Saídas de Atenção Plena e Esparsa no Espaço 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

Resumo

A complexidade quadrática da atenção completa limita o processamento eficiente de contextos longos em grandes modelos de linguagem (LLMs). A atenção esparsa mitiga esse custo ao restringir cada consulta a atender a um subconjunto de tokens anteriores; no entanto, abordagens *training-free* frequentemente resultam em severa degradação de desempenho. Métodos nativos de atenção esparsa (por exemplo, NSA, MoBA) aliviam esse problema, mas exibem um paradoxo crítico: eles produzem menor esparsidade de atenção do que os modelos de atenção completa, apesar de visarem aproximar a atenção completa, o que pode restringir sua eficácia. Atribuímos este paradoxo a uma deficiência na atualização do gradiente: pares chave-valor de baixo posto excluídos durante o treinamento esparso não recebem contribuição direta nem gradientes retroativos e, portanto, nunca aprendem a supressão adequada. Para superar essa limitação, propomos a SSA (*Sparse Sparse Attention*), uma estrutura de treinamento unificada que considera tanto a atenção esparsa quanto a completa e impõe um alinhamento bidirecional em cada camada. Este projeto preserva o fluxo do gradiente para todos os tokens, enquanto incentiva explicitamente que as saídas da atenção esparsa se alinhem com suas contrapartes de atenção completa, promovendo assim uma esparsidade mais forte. Como resultado, a SSA alcança desempenho de última geração sob inferência com atenção esparsa e completa em múltiplos benchmarks de *commonsense*. Adicionalmente, a SSA permite que os modelos se adaptem suavemente a diferentes orçamentos de esparsidade; o desempenho melhora consistentemente à medida que mais tokens têm permissão para atender, suportando compensações flexíveis entre computação e desempenho no momento da inferência. Finalmente, mostramos que o treinamento nativo com atenção esparsa surpreendentemente melhora a extrapolação de contexto longo ao mitigar a superalocação de valores de atenção em áreas *sink*, com a SSA demonstrando a mais forte capacidade de extrapolação.

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.