Attenzione Cosinusoidale: Trasformatori Lineari con Attenzione Cosinusoidale

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

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

cs.AI

Abstract

I meccanismi di attenzione, in particolare l'attenzione softmax, sono stati fondamentali per il successo dei modelli basati su trasformatori come GPT. Tuttavia, la complessità di memoria quadratica dell'attenzione softmax rispetto alla lunghezza della sequenza pone significativi problemi per il trattamento di sequenze più lunghe. Introduciamo Cottention, un nuovo meccanismo di attenzione che sostituisce l'operazione softmax con la similarità cosinica. Sfruttando le proprietà della similarità cosinica e riorganizzando l'equazione di attenzione, Cottention raggiunge una complessità di memoria lineare nativa rispetto alla lunghezza della sequenza, rendendolo intrinsecamente più efficiente dal punto di vista della memoria rispetto all'attenzione softmax. Dimostriamo che Cottention può essere riformulato come una rete neurale ricorrente (RNN) con uno stato nascosto finito, consentendo un utilizzo costante della memoria durante l'inferenza. Valutiamo Cottention sia nei compiti bidirezionali di BERT che nei compiti causali di GPT, dimostrando prestazioni comparabili all'attenzione softmax riducendo significativamente i requisiti di memoria. Per garantire un calcolo efficiente, sviluppiamo un kernel CUDA personalizzato per Cottention. I nostri risultati mostrano che Cottention è una promettente alternativa all'attenzione softmax, consentendo il trattamento di sequenze più lunghe senza sacrificare le prestazioni, grazie alla sua complessità di memoria lineare nativa e alla capacità di mantenere una footprint di memoria costante durante l'inferenza.

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.