[논문 리뷰] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

Angelos Katharopoulos, Apoorv Vyas외 2명|arXiv (Cornell University)|2020. 06. 29.

Neural Networks and Applications참고 문헌 35인용 수 338

한 줄 요약

논문은 소프트맥스 주의(attention)를 커널 기반의 선형 주의로 대체하는 선형 트랜스포머를 소개하여 O(N) 시간/메모리 및 자기회귀 추론을 수천 배 빠르게 하면서도 바닐라 트랜스포머와 유사한 성능을 유지합니다.

ABSTRACT

Transformers achieve remarkable performance in several tasks but due to their quadratic complexity, with respect to the input's length, they are prohibitively slow for very long sequences. To address this limitation, we express the self-attention as a linear dot-product of kernel feature maps and make use of the associativity property of matrix products to reduce the complexity from $\\mathcal{O}\\left(N^2\ ight)$ to $\\mathcal{O}\\left(N\ ight)$, where $N$ is the sequence length. We show that this formulation permits an iterative implementation that dramatically accelerates autoregressive transformers and reveals their relationship to recurrent neural networks. Our linear transformers achieve similar performance to vanilla transformers and they are up to 4000x faster on autoregressive prediction of very long sequences.

연구 동기 및 목표

• 긴 시퀀스에 대한 트랜스포머의 자기 주의 비용이 제곱으로 증가하는 문제를 동기부여하고 해결합니다.
• 커널 특징 맵을 사용한 선형 주의 제형을 제안하여 복잡도를 줄입니다.
• 인과 마스킹과 함께 선형 주의를 보여주고 효율적인 자기회귀 추론을 위한 RNN 유사 재귀를 도출합니다.
• 이미지 생성과 음성 인식에서 벤치마크와 비교하여 경험적으로 검증합니다.

제안 방법

• phi라는 특징 맵을 이용해 self-attention을 커널화된 내적(dot-product)으로 재정식화하여 V' = (phi(Q) (phi(K)^T V)) / (phi(Q)^T sum_j phi(K_j)).
• 모든 쿼리에 대해 합 S = sum_j phi(K_j) V_j^T 와 Z = sum_j phi(K_j) 를 사용하여 합을 계산하고, 앞으로의 패스에서 O(N) 시간으로 처리합니다.
• 선형 시간과 자동 회귀 생성을 위한 S_i 및 Z_i의 점진적 업데이트를 포함한 인과 마스킹을 도입합니다.
• 인과 마스킹을 가진 트랜스포머 계층의 RNN 유사 시각화, 주의 메모리 s와 정규화 메모리 z, 그리고 두 단계 업데이트를 특징으로 합니다.

실험 결과

연구 질문

• RQ1 트랜스포머의 자기 주의가 성능을 희생하지 않으면서 선형(O(N)) 시간과 메모리로 재구성될 수 있는가?
• RQ2 양의 유사도 함수를 갖는 커널 기반 선형 주의가 소프트맥스 주의와 비교해 효과적인 자기회귀 추론을 가능하게 하는가?
• RQ3 인과 마스킹이 있는 트랜스포머 계층을 RNN으로 해석할 수 있어 순차적 생성을 효율화할 수 있는가?
• RQ4 선형 주의 대 표준 트랜스포머 및 Reformer를 이용한 긴 시퀀스 작업(이미지 생성, 음성 인식)에서 정확도와 속도의 실험적 트레이드오프는 어떠한가?
• RQ5 자기회귀 과제의 학습 상태와 추론 상태에서 선형 주의의 성능은 어떠한 차이가 있는가?

주요 결과

Method	Bits/dim	Images/sec
Softmax	0.621	0.45
LSH-1	0.745	0.68
LSH-4	0.676	0.27
Linear (ours)	0.644	142.8

• 선형 트랜스포머는 바닐라 트랜스포머와 유사한 성능을 달성하면서도 긴 시퀀스의 자기회귀 예측에서 최대 4000배까지 더 빠르다.
• 선형 주의는 시퀀스 길이에 대해 메모리와 시간이 선형적으로 스케일하며, 소프트맥스 주의의 제곱 증가와 달리 증가한다.
• MNIST, CIFAR-10, WSJ 실험에서 선형 트랜스포머는 경쟁력 있는 비트/디멘션 또는 PER 성능을 달성하면서 훨씬 높은 처리량과 더 낮은 메모리를 보인다.
• 선형 주의로의 자기회귀 추론은 내부 상태(S_i, Z_i)를 유지하고 업데이트함으로써 매 단계의 비용을 상수로 유지할 수 있다.
• 이 연구는 트랜스포머와 인과 마스킹 간의 이론적 연결을 RNN 유사 재귀를 통해 제시한다.