SeerAttention: Aprendizado de Atenção Esparsa Intrínseca em Seus LLMs

SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

October 17, 2024

Autores: Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang, Mao Yang

cs.AI

Resumo

A atenção é a pedra angular dos Modelos de Linguagem Grandes (LLMs) modernos. No entanto, sua complexidade quadrática limita a eficiência e escalabilidade dos LLMs, especialmente para aqueles com uma janela de contexto longa. Uma abordagem promissora para lidar com essa limitação é aproveitar a esparsidade na atenção. No entanto, as soluções baseadas em esparsidade existentes predominantemente dependem de padrões predefinidos ou heurísticas para aproximar a esparsidade. Essa prática não consegue capturar totalmente a natureza dinâmica da esparsidade de atenção em tarefas baseadas em linguagem. Este artigo argumenta que a esparsidade de atenção deve ser aprendida em vez de predefinida. Para isso, projetamos o SeerAttention, um novo mecanismo de Atenção que complementa a atenção convencional com um portão aprendível que seleciona adaptativamente blocos significativos em um mapa de atenção e considera os demais blocos esparsos. Essa esparsidade em nível de bloco equilibra efetivamente a precisão e a aceleração. Para permitir a aprendizagem eficiente da rede de portões, desenvolvemos uma implementação FlashAttention personalizada que extrai a verdade fundamental em nível de bloco do mapa de atenção com um mínimo de sobrecarga. O SeerAttention não se aplica apenas ao pós-treinamento, mas também se destaca no ajuste fino de longo contexto. Nossos resultados mostram que nas fases pós-treinamento, o SeerAttention supera significativamente os métodos de atenção esparsa estáticos ou baseados em heurísticas de última geração, sendo também mais versátil e flexível para se adaptar a comprimentos de contexto variáveis e taxas de esparsidade. Quando aplicado ao ajuste fino de longo contexto com YaRN, o SeerAttention pode atingir uma notável taxa de esparsidade de 90% em um comprimento de contexto de 32k com uma perda mínima de perplexidade, oferecendo uma aceleração de 5,67 vezes em relação ao FlashAttention-2.

English

Attention is the cornerstone of modern Large Language Models (LLMs). Yet its quadratic complexity limits the efficiency and scalability of LLMs, especially for those with a long-context window. A promising approach addressing this limitation is to leverage the sparsity in attention. However, existing sparsity-based solutions predominantly rely on predefined patterns or heuristics to approximate sparsity. This practice falls short to fully capture the dynamic nature of attention sparsity in language-based tasks. This paper argues that attention sparsity should be learned rather than predefined. To this end, we design SeerAttention, a new Attention mechanism that augments the conventional attention with a learnable gate that adaptively selects significant blocks in an attention map and deems the rest blocks sparse. Such block-level sparsity effectively balances accuracy and speedup. To enable efficient learning of the gating network, we develop a customized FlashAttention implementation that extracts the block-level ground truth of attention map with minimum overhead. SeerAttention not only applies to post-training, but also excels in long-context fine-tuning. Our results show that at post-training stages, SeerAttention significantly outperforms state-of-the-art static or heuristic-based sparse attention methods, while also being more versatile and flexible to adapt to varying context lengths and sparsity ratios. When applied to long-context fine-tuning with YaRN, SeerAttention can achieve a remarkable 90% sparsity ratio at a 32k context length with minimal perplexity loss, offering a 5.67x speedup over FlashAttention-2.