Mais Esparsa é Mais Rápida e Menos é Mais: Atenção Esparsa Eficiente para Transformers de Longo Alcance

Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers

June 24, 2024

Autores: Chao Lou, Zixia Jia, Zilong Zheng, Kewei Tu

cs.AI

Resumo

Acomodar sequências longas de forma eficiente em Transformadores autoregressivos, especialmente dentro de uma janela de contexto estendida, apresenta desafios significativos devido à complexidade computacional quadrática e aos substanciais requisitos de memória KV inerentes aos mecanismos de autoatenção. Neste trabalho, introduzimos a Atenção SPARSEK, um novo mecanismo de atenção esparsa projetado para superar esses obstáculos computacionais e de memória, mantendo o desempenho. Nossa abordagem integra uma rede de pontuação e um operador de máscara top-k diferenciável, SPARSEK, para selecionar um número constante de pares KV para cada consulta, permitindo assim a otimização baseada em gradientes. Como resultado, a Atenção SPARSEK oferece complexidade de tempo linear e pegada de memória constante durante a geração. Resultados experimentais revelam que a Atenção SPARSEK supera os métodos de atenção esparsa anteriores e proporciona melhorias significativas de velocidade tanto durante o treinamento quanto na inferência, especialmente em modelagem de linguagem e tarefas subsequentes. Além disso, nosso método pode ser integrado de forma transparente em Modelos de Linguagem Grandes (LLMs) pré-treinados com ajustes mínimos, oferecendo uma solução prática para gerenciar efetivamente dependências de longo alcance em diversas aplicações.

English

Accommodating long sequences efficiently in autoregressive Transformers, especially within an extended context window, poses significant challenges due to the quadratic computational complexity and substantial KV memory requirements inherent in self-attention mechanisms. In this work, we introduce SPARSEK Attention, a novel sparse attention mechanism designed to overcome these computational and memory obstacles while maintaining performance. Our approach integrates a scoring network and a differentiable top-k mask operator, SPARSEK, to select a constant number of KV pairs for each query, thereby enabling gradient-based optimization. As a result, SPARSEK Attention offers linear time complexity and constant memory footprint during generation. Experimental results reveal that SPARSEK Attention outperforms previous sparse attention methods and provides significant speed improvements during both training and inference, particularly in language modeling and downstream tasks. Furthermore, our method can be seamlessly integrated into pre-trained Large Language Models (LLMs) with minimal fine-tuning, offering a practical solution for effectively managing long-range dependencies in diverse applications.