ZeerAandacht: Het Leren van Intrinsieke Schaarse Aandacht in Jouw LLMs

SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

October 17, 2024

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

cs.AI

Samenvatting

Aandacht is de hoeksteen van moderne Large Language Models (LLMs). Toch beperkt de kwadratische complexiteit ervan de efficiëntie en schaalbaarheid van LLMs, vooral voor die met een lange contextvenster. Een veelbelovende aanpak om deze beperking aan te pakken, is het benutten van de spaarzaamheid in aandacht. Bestaande spaarzaamheidsgebaseerde oplossingen vertrouwen echter voornamelijk op vooraf gedefinieerde patronen of heuristieken om spaarzaamheid te benaderen. Deze praktijk schiet tekort om de dynamische aard van aandachtsspaarzaamheid volledig vast te leggen in taakgerichte taken op basis van taal. Dit artikel betoogt dat aandachtsspaarzaamheid geleerd moet worden in plaats van vooraf gedefinieerd. Hiertoe ontwerpen we SeerAttention, een nieuw Aandachtsmechanisme dat de conventionele aandacht aanvult met een leerbaar hek dat adaptief significante blokken selecteert in een aandachtskaart en de overige blokken als spaarzaam beschouwt. Een dergelijke blokniveau spaarzaamheid balanceert effectief nauwkeurigheid en versnelling. Om efficiënt leren van het hekwerknetwerk mogelijk te maken, ontwikkelen we een aangepaste FlashAttention-implementatie die de blokniveau grondwaarheid van de aandachtskaart met minimale overhead extraheren. SeerAttention is niet alleen van toepassing op post-training, maar blinkt ook uit in fijnafstemming met een lange context. Onze resultaten tonen aan dat SeerAttention in post-trainingfasen aanzienlijk beter presteert dan state-of-the-art statische of op heuristiek gebaseerde spaarzame aandachtsmethoden, terwijl het ook veelzijdiger en flexibeler is om zich aan te passen aan variërende contextlengtes en spaarzaamheidsverhoudingen. Wanneer toegepast op fijnafstemming met YaRN met een lange context, kan SeerAttention een opmerkelijke spaarzaamheidsverhouding van 90% bereiken bij een contextlengte van 32k met minimaal perplexiteitsverlies, wat een 5,67x versnelling biedt ten opzichte van 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.