코텐션: 코사인 유사도로 구현한 선형 트랜스포머

초록

코텐션은 소프트맥스 대신 코사인 유사도를 이용해 어텐션 가중치를 계산함으로써 시퀀스 길이에 대해 선형 메모리 복잡도를 달성한다. 정규화와 학습 가능한 스케일링을 통해 훈련 안정성을 확보하고, 양방향 BERT와 인과형 GPT에 적용해 기존 소프트맥스 어텐션과 비슷한 성능을 보이며 메모리 사용량을 크게 절감한다. CUDA 커스텀 커널을 제공해 실시간 추론 시에도 고정 메모리 사용이 가능하도록 설계되었다.

상세 분석

본 논문은 어텐션 메커니즘의 근본적인 비용 구조를 재검토한다. 기존 소프트맥스 어텐션은 Q·Kᵀ를 계산한 뒤 행별로 소프트맥스를 적용해 가중치를 정규화하는데, 이 과정은 시퀀스 길이 s에 대해 O(s²) 메모리를 요구한다. 코텐션은 코사인 유사도 Sim(Q,K)=N(Q)·N(K)ᵀ 형태로 정의함으로써, Q와 K를 각각 L2 정규화(N)한 뒤 행렬 곱을 수행한다. 여기서 핵심은 정규화 단계가 행렬 곱과 완전히 분리될 수 있다는 점이다. 따라서 연산 순서를 바꾸어 (N(K)ᵀ·V) 를 먼저 계산하면 중간 결과의 차원은 d² (d는 헤드당 키 차원)만 필요하게 되어 메모리 사용량이 O(d²)로 감소한다. 이는 s≫d인 경우에 특히 유리하다.

하지만 코사인 어텐션은 행별 합이 s까지 커질 수 있어 훈련 초기에 그래디언트가 폭발하거나 소실될 위험이 있다. 이를 완화하기 위해 논문은 시퀀스 길이 s의 m제곱(σ(m)으로 스케일링)으로 나누는 정규화 항을 도입한다. m은 각 헤드마다 학습 가능한 스칼라이며, 시그모이드 함수를 통과해