Atención Coseno: Transformadores Lineales con Atención Coseno

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

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

cs.AI

Resumen

Los mecanismos de atención, en particular la atención softmax, han sido fundamentales en el éxito de los modelos basados en transformadores como GPT. Sin embargo, la complejidad cuadrática de la memoria de la atención softmax con respecto a la longitud de la secuencia plantea desafíos significativos para el procesamiento de secuencias más largas. Presentamos Cottention, un nuevo mecanismo de atención que reemplaza la operación softmax con similitud coseno. Al aprovechar las propiedades de la similitud coseno y reorganizar la ecuación de atención, Cottention logra una complejidad de memoria lineal nativa con respecto a la longitud de la secuencia, lo que lo hace inherentemente más eficiente en memoria que la atención softmax. Demostramos que Cottention puede reformularse como una red neuronal recurrente (RNN) con un estado oculto finito, lo que permite un uso de memoria constante durante la inferencia. Evaluamos Cottention en las tareas bidireccionales de BERT y GPT causal, demostrando un rendimiento comparable a la atención softmax mientras se reducen significativamente los requisitos de memoria. Para garantizar una computación eficiente, desarrollamos un kernel CUDA personalizado para Cottention. Nuestros resultados muestran que Cottention es una alternativa prometedora a la atención softmax, que permite el procesamiento de secuencias más largas sin sacrificar el rendimiento, debido a su complejidad de memoria lineal nativa y su capacidad para mantener una huella de memoria constante durante la inferencia.

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.