注意: コサインアテンションを用いた線形トランスフォーマー

Cottention: Linear Transformers With Cosine Attention

September 27, 2024

著者: Gabriel Mongaras, Trevor Dohm, Eric C. Larson

cs.AI

要旨

注意メカニズム、特にソフトマックスアテンションは、GPTなどのトランスフォーマーベースのモデルの成功に重要な役割を果たしてきました。ただし、シーケンスの長さに関連するソフトマックスアテンションの二次メモリ複雑度は、より長いシーケンスの処理において重要な課題を提起しています。本研究では、コサイン類似性によってソフトマックス演算を置き換えた新しいアテンションメカニズムであるCottentionを紹介します。コサイン類似性の特性を活用し、アテンション方程式を再配置することで、Cottentionはシーケンスの長さに関連するネイティブな線形メモリ複雑度を実現し、ソフトマックスアテンションよりもメモリ効率が高くなります。Cottentionを有限の隠れ状態を持つ再帰ニューラルネットワーク（RNN）として再定式化できることを示し、推論時の一定のメモリ使用を可能にします。Cottentionを双方向BERTおよび因果GPTのタスクで評価し、ソフトマックスアテンションと比較可能な性能を示しながら、メモリ要件を大幅に削減します。効率的な計算を確保するために、Cottention用のカスタムCUDAカーネルを開発します。結果は、Cottentionがソフトマックスアテンションに対する有望な代替手段であり、ネイティブな線形メモリ複雑度と推論中の一定のメモリフットプリントを維持する能力により、性能を犠牲にすることなく、より長いシーケンスの処理を可能にすることを示しています。

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.