SeerAttention: Apprendimento dell'Attenzione Sparso Intrinseca nei Tuoi LLMs

SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

October 17, 2024

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

cs.AI

Abstract

L'attenzione è la pietra angolare dei moderni Grandi Modelli Linguistici (LLM). Tuttavia, la sua complessità quadratica limita l'efficienza e la scalabilità dei LLM, specialmente per quelli con una finestra di contesto lunga. Un approccio promettente che affronta questa limitazione è sfruttare la sparso nell'attenzione. Tuttavia, le soluzioni basate sulla sparso esistenti si basano principalmente su pattern predefiniti o euristiche per approssimare la sparso. Questa pratica non riesce a catturare appieno la natura dinamica della sparso dell'attenzione nei compiti basati sul linguaggio. Questo articolo sostiene che la sparso dell'attenzione dovrebbe essere appresa piuttosto che predefinita. A tal fine, progettiamo SeerAttention, un nuovo meccanismo di Attenzione che potenzia l'attenzione convenzionale con un gate apprendibile che seleziona in modo adattivo blocchi significativi in una mappa di attenzione e considera gli altri blocchi sparso. Tale sparso a livello di blocco bilancia efficacemente precisione e accelerazione. Per consentire un apprendimento efficiente della rete di gating, sviluppiamo un'implementazione FlashAttention personalizzata che estrae la verità fondamentale a livello di blocco della mappa di attenzione con un overhead minimo. SeerAttention non si applica solo al post-training, ma eccelle anche nel fine-tuning a lungo contesto. I nostri risultati mostrano che nelle fasi post-training, SeerAttention supera significativamente i metodi di attenzione sparso statici o basati su euristiche all'avanguardia, essendo anche più versatile e flessibile per adattarsi a lunghezze di contesto e rapporti di sparso variabili. Applicato al fine-tuning a lungo contesto con YaRN, SeerAttention può raggiungere un notevole rapporto di sparso del 90% a una lunghezza di contesto di 32k con una perdita minima di perplessità, offrendo un'accelerazione di 5,67x rispetto a 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.