[논문 리뷰] Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding

Shengjie Luo, Shanda Li외 7명|arXiv (Cornell University)|2021. 06. 23.

Advanced Neural Network Applications인용 수 10

한 줄 요약

이 논문은 빠른 푸리에 변환(Fast Fourier Transform, FFT)을 사용하여 O(n log n) 복잡도를 달성하는 새로운 커널화된 어텐션 메커니즘을 제안한다. 이는 상대적 위치 인코딩(Relative Positional Encoding, RPE)을 통합하여 트랜스포머에서 장수열 모델링을 안정적이고 빠르며 정확하게 가능하게 한다. 이 방법은 RPE 행렬의 토플리츠 구조를 활용하여 계산을 가속화하면서도 어텐션 점수의 분산을 줄여 학습 안정성을 향상시킨다.

ABSTRACT

The attention module, which is a crucial component in Transformer, cannot scale efficiently to long sequences due to its quadratic complexity. Many works focus on approximating the dot-then-exponentiate softmax function in the original attention, leading to sub-quadratic or even linear-complexity Transformer architectures. However, we show that these methods cannot be applied to more powerful attention modules that go beyond the dot-then-exponentiate style, e.g., Transformers with relative positional encoding (RPE). Since in many state-of-the-art models, relative positional encoding is used as default, designing efficient Transformers that can incorporate RPE is appealing. In this paper, we propose a novel way to accelerate attention calculation for Transformers with RPE on top of the kernelized attention. Based upon the observation that relative positional encoding forms a Toeplitz matrix, we mathematically show that kernelized attention with RPE can be calculated efficiently using Fast Fourier Transform (FFT). With FFT, our method achieves $\\mathcal{O}(n\\log n)$ time complexity. Interestingly, we further demonstrate that properly using relative positional encoding can mitigate the training instability problem of vanilla kernelized attention. On a wide range of tasks, we empirically show that our models can be trained from scratch without any optimization issues. The learned model performs better than many efficient Transformer variants and is faster than standard Transformer in the long-sequence regime.

연구 동기 및 목표

• 표준 자기어텐션의 이중 복잡도로 인해 장수열에서 비효율적인 문제를 해결하기 위해.
• 최신 모델에서 일반적으로 사용되는 상대적 위치 인코딩(Relative Positional Encoding, RPE)을 사용하는 트랜스포머에서 효율적인 어텐션 계산을 가능하게 하기 위해.
• 기존 근사 방법들이 도트-지수화 어텐션을 초월하는 더 넓은 표현력을 지닌 RPE를 처리할 수 없는 한계를 극복하기 위해.
• 실제로 자주 발생하는 최적화 불안정성 문제로 인해 학습이 어려운 커널화된 어텐션 모델의 학습 안정성을 향상시키기 위해.
• 세부 조정(fine-tuning)에 의존하지 않고도 장수열 환경에서 높은 계산 효율성과 성능을 동시에 달성하기 위해.

제안 방법

• 상대적 위치 인코딩 행렬의 토플리츠 구조를 활용하여 커널화된 어텐션과 RPE를 재구성한다.
• 빠른 푸리에 변환(Fast Fourier Transform, FFT)을 사용하여 RPE 행렬과 커널화된 특징 행렬 간의 행렬-벡터 곱셈을 가속화함으로써 시간 복잡도를 O(n log n)으로 감소시킨다.
• 표준 소프트맥스 어텐션을 특징 맵을 사용하여 어텐션 함수를 근사하는 커널화된 어텐션 메커니즘으로 대체한다.
• 질의와 키의 노름이 클 경우 특히 유의미한 분산 감소를 위해 정규화된 어텐션(Normalized Positional Receptive Function, NPRF)을 도입한다.
• 커널화된 어텐션 메커니즘을 인코더 및 디코더 블록에 모두 적용하여 트랜스포머 아키텍처에 완전히 통합할 수 있도록 한다.
• 다양한 특징 맵 함수(예: PRF, TRF, Sphere-PRF, ORF)를 지원하여 다양한 선택지에 대해 일반화 성능을 입증한다.

실험 결과

연구 질문

• RQ1 표준 도트-지수화 공식을 초월하는 커널화된 어텐션과 상대적 위치 인코딩(RPE)을 효율적으로 계산할 수 있는가?
• RQ2 커널화된 어텐션에 RPE를 도입할 경우, 기존 커널화된 어텐션 대비 학습 안정성이 향상되는가?
• RQ3 RPE 행렬의 토플리츠 구조를 활용하여 어텐션 계산에서 이중 이하의 시간 복잡도를 달성할 수 있는가?
• RQ4 제안된 방법은 장수열 작업에서 표준 및 효율적인 트랜스포머 변종과 비교해 성능가능한가?
• RQ5 제안된 방법은 커널화된 어텐션 메커니즘의 다양한 특징 맵 차원과 유형에 대해 강건한가?

주요 결과

• 제안된 방법은 RPE 행렬의 토플리츠 구조를 활용하여 FFT를 적용함으로써 O(n log n) 시간 복잡도를 달성하여 장수열 모델링을 효율적으로 가능하게 한다.
• 정규화된 커널화된 어텐션과 RPE를 결합한 모델은 최적화 문제 없이 초기 상태에서 안정적으로 학습이 가능하며, 기존의 커널화된 어텐션과는 달리 학습 불안정성이 발생하지 않는다.
• ImageNet 데이터셋에서 RPE를 적용한 NPRF-Transformer는 80.9%의 top-1 정확도를 달성하여 표준 DeiT-base(81.2%)와 PRF 변환 모델(79.5%)을 모두 초월한다.
• IWSLT14 독일어-영어 번역 작업에서 RPE와 정규화된 어텐션을 적용한 모델은 BLEU 점수 32.1을 기록하여 잘 트레이닝된 표준 트랜스포머(32.3)와 유사한 성능을 보였으며, 선형 복잡도를 유지했다.
• 특징 맵 차원이 16일 때조차도 성능 저하가 최소한이었으며, 다양한 특징 맵 함수(PRF, TRF, Sphere-PRF, ORF)에 대해 일반화 성능이 뛰어나다.
• 아블레이션 연구에서 RPE 사용 시 표준 어텐션에서 커널화된 어텐션으로 전환할 때의 근사 오차가 감소하며, 특히 정규화와 결합할 경우 성능 저하가 더 작아짐을 확인했다.