[논문 리뷰] Attention as an RNN

Leo Feng, Frederick Tung외 4명|arXiv (Cornell University)|2024. 05. 22.

Educational and Psychological AssessmentsPsychology인용 수 3

한 줄 요약

이 논문은 Aaren을 제안하며, 이는 자기주의 어텐션을 미분 가능한 RNN으로 재구성함으로써 기존 트랜스포머와 동일한 병렬 훈련을 가능하게 하면서도, 증분 토큰 업데이트를 통한 효율적이고 일정한 메모리 사용을 가능하게 하는 새로운 어텐션 메커니즘을 제공한다. Aaren은 38개의 순차적 모델링 작업 전반에서 트랜스포머 수준의 성능을 달성하면서도 계산 및 메모리 복잡도를 이차함수에서 선형으로 감소시킨다.

ABSTRACT

The advent of Transformers marked a significant breakthrough in sequence modelling, providing a highly performant architecture capable of leveraging GPU parallelism. However, Transformers are computationally expensive at inference time, limiting their applications, particularly in low-resource settings (e.g., mobile and embedded devices). Addressing this, we (1) begin by showing that attention can be viewed as a special Recurrent Neural Network (RNN) with the ability to compute its extit{many-to-one} RNN output efficiently. We then (2) show that popular attention-based models such as Transformers can be viewed as RNN variants. However, unlike traditional RNNs (e.g., LSTMs), these models cannot be updated efficiently with new tokens, an important property in sequence modelling. Tackling this, we (3) introduce a new efficient method of computing attention's extit{many-to-many} RNN output based on the parallel prefix scan algorithm. Building on the new attention formulation, we (4) introduce extbf{Aaren}, an attention-based module that can not only (i) be trained in parallel (like Transformers) but also (ii) be updated efficiently with new tokens, requiring only constant memory for inferences (like traditional RNNs). Empirically, we show Aarens achieve comparable performance to Transformers on $38$ datasets spread across four popular sequential problem settings: reinforcement learning, event forecasting, time series classification, and time series forecasting tasks while being more time and memory-efficient.

연구 동기 및 목표

• 모바일 및 임베디드 장치와 같은 자원이 제한된 환경에서 트랜스포머의 높은 추론 비용을 해결하기 위해.
• 트랜스포머의 병렬성과 RNN의 효율적이고 증분적인 추론 간 격차를 메우기 위해.
• 병렬 훈련과 일정한 메모리, 스트리밍 추론을 모두 지원하는 새로운 어텐션 공식화를 개발하기 위해.
• 병렬 프리픽스 스캔 알고리즘을 활용해 어텐션의 다대다 RNN 출력을 효율적으로 계산하는 방법을 도입하기 위해.
• 결과적으로 Aaren 모듈이 트랜스포머 성능를 유지하면서도 시간과 메모리 복잡도를 크게 감소시킬 수 있음을 입증하기 위해.

제안 방법

• 표준 자기주의 어텐션을 다대일 RNN의 특수한 케이스로 재구성함으로써 어텐션 계산의 순환적 해석을 가능하게 한다.
• 다수의 출력을 효율적으로 계산하기 위해 병렬 프리픽스 스캔 알고리즘을 활용한 새로운 RNN 기반 어텐션 공식화를 도입한다.
• 소프트맥스 기반 어텐션의 미분 가능하고 정확한 재구현을 RNN 형태로 유도함으로써 이전의 근사 방법에서 관찰된 지수적 타임스탬프 편향을 피한다.
• 일정한 은닉 상태 크기를 유지하고 새로운 토큰에 대한 증분 업데이트를 지원하는 학습 가능한 RNN 모듈로 Aaren을 설계한다.
• 과거 활성화를 저장하지 않고도 장거리 의존성을 포착할 수 있도록 학습 가능한 초기 은닉 상태 벡터 $ q $ 를 활용한다.
• 효율적인 병렬 프리픽스 합 알고리즘(예: 힐리스-스티리)을 활용하여 어텐션 가중치를 $ O(N \log N) $ 시간 내에 계산함으로써 확장 가능한 추론을 가능하게 한다.

실험 결과

연구 질문

• RQ1 표준 자기주의 어텐션을 공식적으로 순환 신경망으로 재해석할 수 있는가?
• RQ2 왜 어텐션 기반 모델인 트랜스포머는 병렬 훈련 능력은 갖추고 있음에도 불구하고 효율적인 증분 추론을 지원하지 못하는가?
• RQ3 트랜스포머의 병렬성과 RNN의 일정한 메모리 효율성을 모두 결합한 어텐션 메커니즘을 설계할 수 있는가?
• RQ4 병렬 프리픽스 스캔 기반의 새로운 어텐션 공식화가 트랜스포머 수준의 성능를 확보하면서도 계산 복잡도를 감소시킬 수 있는가?
• RQ5 제안된 Aaren 모듈은 다양한 순차적 작업에서 트랜스포머와 비교해 추론 효율성과 모델 성능 측면에서 어떻게 비교되는가?

주요 결과

• Aaren은 강화학습, 이벤트 예측, 시계열 분류, 시계열 예측 분야의 38개의 다양한 순차적 모델링 데이터셋에서 트랜스포머 수준의 성능를 달성한다.
• Aaren의 누적 추론 시간과 메모리 사용량은 시퀀스 길이에 따라 선형적으로 증가하지만, KV 캐싱을 사용하는 트랜스포머는 이차함수적으로 증가한다.
• 동등한 트랜스포머 대비 Aaren은 파라미터 수가 0.016%만 증가(3,152,896 대비 3,152,384)하여 효율성 향상의 비용이 매우 낮다.
• 추론 중에 일정한 메모리 사용량을 유지하므로, 모바일 및 임베디드 시스템과 같은 자원이 제한된 장치에 배포할 수 있다.
• 병렬 프리픽스 스캔 기반의 어텐션 계산은 어텐션 메커니즘의 표현력을 유지하면서도 효율적이고 미분 가능한 RNN 스타일의 추론을 가능하게 한다.
• 실험 결과 Aaren이 모든 평가된 작업에서 속도와 메모리 효율성 측면에서 트랜스포머를 능가하거나 동등하게 성능을 발휘함을 확인했다.