Внимание: Линейные трансформеры с косинусным вниманием

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

Авторы: Gabriel Mongaras, Trevor Dohm, Eric C. Larson

cs.AI

Аннотация

Механизмы внимания, особенно механизм внимания softmax, сыграли ключевую роль в успехе моделей на основе трансформера, таких как GPT. Однако квадратичная сложность памяти механизма внимания softmax относительно длины последовательности представляет существенные вызовы для обработки более длинных последовательностей. Мы представляем Cottention, новый механизм внимания, который заменяет операцию softmax на косинусное сходство. Используя свойства косинусного сходства и переставляя уравнение внимания, Cottention достигает линейной сложности памяти относительно длины последовательности, что делает его более эффективным с точки зрения памяти, чем механизм внимания softmax. Мы показываем, что Cottention может быть переформулирован как рекуррентная нейронная сеть (RNN) с конечным скрытым состоянием, позволяющим использовать постоянный объем памяти во время вывода. Мы оцениваем Cottention на задачах как двунаправленного BERT, так и причинного GPT, демонстрируя сравнимую производительность с механизмом внимания softmax при значительном снижении требований к памяти. Для обеспечения эффективных вычислений мы разрабатываем специальное ядро CUDA для Cottention. Наши результаты показывают, что Cottention является многообещающей альтернативой механизму внимания softmax, позволяя обрабатывать более длинные последовательности без ущерба производительности благодаря линейной сложности памяти и способности поддерживать постоянный объем памяти во время вывода.

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.