Atenção Cosseno: Transformadores Lineares com Atenção Cosseno

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

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

cs.AI

Resumo

Mecanismos de atenção, particularmente atenção softmax, têm sido fundamentais para o sucesso de modelos baseados em transformadores como o GPT. No entanto, a complexidade de memória quadrática da atenção softmax em relação ao comprimento da sequência apresenta desafios significativos para o processamento de sequências mais longas. Apresentamos o Cottention, um novo mecanismo de atenção que substitui a operação softmax pela similaridade de cosseno. Ao aproveitar as propriedades da similaridade de cosseno e reorganizar a equação de atenção, o Cottention alcança uma complexidade de memória linear nativa em relação ao comprimento da sequência, tornando-o inerentemente mais eficiente em termos de memória do que a atenção softmax. Demonstramos que o Cottention pode ser reformulado como uma rede neural recorrente (RNN) com um estado oculto finito, permitindo o uso de memória constante durante a inferência. Avaliamos o Cottention nas tarefas bidirecionais BERT e causal GPT, demonstrando desempenho comparável à atenção softmax, enquanto reduzimos significativamente os requisitos de memória. Para garantir uma computação eficiente, desenvolvemos um kernel CUDA personalizado para o Cottention. Nossos resultados mostram que o Cottention é uma alternativa promissora à atenção softmax, possibilitando o processamento de sequências mais longas sem sacrificar o desempenho, devido à sua complexidade de memória linear nativa e capacidade de manter uma pegada de memória constante durante a inferência.

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.