Atenção Esparsa Nativa: Alinhada ao Hardware e Intrinsecamente Treinável com EsparsidadeNative Sparse Attention: Hardware-Aligned and Natively Trainable Sparse
Attention
A modelagem de contexto longo é crucial para os modelos de linguagem de próxima geração, mas o alto custo computacional dos mecanismos de atenção padrão apresenta desafios computacionais significativos. A atenção esparsa oferece uma direção promissora para melhorar a eficiência enquanto mantém as capacidades do modelo. Apresentamos o NSA, um mecanismo de Atenção Esparsa Nativamente Treinável que integra inovações algorítmicas com otimizações alinhadas ao hardware para alcançar uma modelagem eficiente de contexto longo. O NSA emprega uma estratégia esparsa hierárquica dinâmica, combinando compressão de tokens em nível grosseiro com seleção de tokens em nível fino para preservar tanto a consciência do contexto global quanto a precisão local. Nossa abordagem avança o design da atenção esparsa com duas inovações principais: (1) Alcançamos acelerações substanciais por meio de um design algorítmico balanceado em intensidade aritmética, com otimizações de implementação para hardware moderno. (2) Habilitamos o treinamento de ponta a ponta, reduzindo a computação de pré-treinamento sem sacrificar o desempenho do modelo. Como mostrado na Figura 1, os experimentos demonstram que o modelo pré-treinado com NSA mantém ou supera os modelos de Atenção Completa em benchmarks gerais, tarefas de contexto longo e raciocínio baseado em instruções. Enquanto isso, o NSA alcança acelerações substanciais em relação à Atenção Completa em sequências de 64k de comprimento ao longo da decodificação, propagação direta e propagação reversa, validando sua eficiência em todo o ciclo de vida do modelo.