주의: 코사인 어텐션을 사용한 선형 트랜스포머

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.