Aandacht: Lineaire Transformatoren met Cosinus Aandacht

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

Auteurs: Gabriel Mongaras, Trevor Dohm, Eric C. Larson

cs.AI

Samenvatting

Aandachtsmechanismen, met name softmax-aandacht, zijn essentieel geweest voor het succes van op transformer gebaseerde modellen zoals GPT. De kwadratische geheugencomplexiteit van softmax-aandacht ten opzichte van de sequentielengte vormt echter aanzienlijke uitdagingen voor het verwerken van langere sequenties. We introduceren Cottention, een nieuw aandachtsmechanisme dat de softmax-operatie vervangt door cosinusgelijkenis. Door gebruik te maken van de eigenschappen van cosinusgelijkenis en de aandachtsvergelijking te herschikken, bereikt Cottention een lineaire geheugencomplexiteit ten opzichte van de sequentielengte, waardoor het inherent geheugenefficiënter is dan softmax-aandacht. We tonen aan dat Cottention kan worden herschreven als een recurrent neuraal netwerk (RNN) met een eindige verborgen toestand, waardoor constante geheugenbenutting tijdens inferentie mogelijk is. We evalueren Cottention voor zowel de bidirectionele BERT- als causale GPT-taken, waarbij vergelijkbare prestaties als softmax-aandacht worden aangetoond, terwijl de geheugenvereisten aanzienlijk worden verminderd. Om efficiënte berekening te garanderen, ontwikkelen we een aangepaste CUDA-kernel voor Cottention. Onze resultaten tonen aan dat Cottention een veelbelovend alternatief is voor softmax-aandacht, waardoor de verwerking van langere sequenties mogelijk is zonder prestatieverlies, dankzij de lineaire geheugencomplexiteit en het vermogen om een constante geheugenafdruk tijdens inferentie te behouden.

English

Attention mechanisms, particularly softmax attention, have been instrumental in the success of transformer-based models such as GPT. However, the quadratic memory complexity of softmax attention with respect to sequence length poses significant challenges for processing longer sequences. We introduce Cottention, a novel attention mechanism that replaces the softmax operation with cosine similarity. By leveraging the properties of cosine similarity and rearranging the attention equation, Cottention achieves native linear memory complexity with respect to sequence length, making it inherently more memory-efficient than softmax attention. We demonstrate that Cottention can be reformulated as a recurrent neural network (RNN) with a finite hidden state, allowing for constant memory usage during inference. We evaluate Cottention on both the bidirectional BERT and causal GPT tasks, demonstrating comparable performance to softmax attention while significantly reducing memory requirements. To ensure efficient computation, we develop a custom CUDA kernel for Cottention. Our results show that Cottention is a promising alternative to softmax attention, enabling the processing of longer sequences without sacrificing performance, due to its native linear memory complexity and ability to maintain a constant memory footprint during inference.