Achtung: Lineare Transformer mit Cosinus-Aufmerksamkeit

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

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

cs.AI

Zusammenfassung

Aufmerksamkeitsmechanismen, insbesondere Softmax-Aufmerksamkeit, haben maßgeblich zum Erfolg von Transformer-basierten Modellen wie GPT beigetragen. Die quadratische Speicherkomplexität der Softmax-Aufmerksamkeit in Bezug auf die Sequenzlänge stellt jedoch erhebliche Herausforderungen für die Verarbeitung längerer Sequenzen dar. Wir stellen Cottention vor, einen neuartigen Aufmerksamkeitsmechanismus, der die Softmax-Operation durch die Kosinus-Ähnlichkeit ersetzt. Durch Nutzung der Eigenschaften der Kosinus-Ähnlichkeit und Umstrukturierung der Aufmerksamkeitsgleichung erreicht Cottention eine native lineare Speicherkomplexität in Bezug auf die Sequenzlänge, was es im Vergleich zur Softmax-Aufmerksamkeit inhärent speichereffizienter macht. Wir zeigen, dass Cottention als rekurrentes neuronales Netzwerk (RNN) mit einem endlichen Hidden State umformuliert werden kann, was einen konstanten Speicherverbrauch während der Inferenz ermöglicht. Wir evaluieren Cottention sowohl bei den bidirektionalen BERT- als auch bei den kausalen GPT-Aufgaben und zeigen eine vergleichbare Leistung zur Softmax-Aufmerksamkeit bei gleichzeitiger signifikanter Reduzierung des Speicherbedarfs. Zur Gewährleistung effizienter Berechnungen entwickeln wir einen benutzerdefinierten CUDA-Kern für Cottention. Unsere Ergebnisse zeigen, dass Cottention eine vielversprechende Alternative zur Softmax-Aufmerksamkeit darstellt, die die Verarbeitung längerer Sequenzen ohne Leistungseinbußen ermöglicht, aufgrund seiner nativen linearen Speicherkomplexität und der Fähigkeit, während der Inferenz eine konstante Speicherbelegung aufrechtzuerhalten.

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.