Attention Cosinus: Transformateurs Linéaires Avec Attention Cosinus

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

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

cs.AI

Résumé

Les mécanismes d'attention, en particulier l'attention softmax, ont été essentiels dans le succès des modèles basés sur les transformers tels que GPT. Cependant, la complexité quadratique de la mémoire de l'attention softmax par rapport à la longueur de la séquence pose des défis significatifs pour le traitement de séquences plus longues. Nous introduisons Cottention, un nouveau mécanisme d'attention qui remplace l'opération softmax par la similarité cosinus. En exploitant les propriétés de la similarité cosinus et en réarrangeant l'équation d'attention, Cottention atteint une complexité de mémoire linéaire native par rapport à la longueur de la séquence, le rendant intrinsèquement plus efficace en mémoire que l'attention softmax. Nous démontrons que Cottention peut être reformulé en un réseau de neurones récurrent (RNN) avec un état caché fini, permettant une utilisation de mémoire constante lors de l'inférence. Nous évaluons Cottention sur les tâches bidirectionnelles BERT et causales GPT, démontrant des performances comparables à l'attention softmax tout en réduisant significativement les besoins en mémoire. Pour assurer un calcul efficace, nous développons un noyau CUDA personnalisé pour Cottention. Nos résultats montrent que Cottention est une alternative prometteuse à l'attention softmax, permettant le traitement de séquences plus longues sans sacrifier les performances, grâce à sa complexité de mémoire linéaire native et à sa capacité à maintenir une empreinte mémoire constante lors de l'inférence.

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.