[논문 리뷰] KDEformer: Accelerating Transformers via Kernel Density Estimation

Amir Zandieh, In‐Su Han외 2명|arXiv (Cornell University)|2023. 02. 05.

Advanced Neural Network Applications인용 수 5

한 줄 요약

KDEformer는 도트 곱 어텐션을 커널 밀도 추정(KDE) 문제로 재구성함으로써, 증명된 스펙트럴 노름 오차 경계를 갖는 서브-이차 계산을 가능하게 하여 트랜스포머의 계산 속도를 높인다. ImageNet에서 최대 18.3배의 속도 향상을 이끌었으며 정확도 저하가 0.5% 미만이었고, BigGAN에서 FID 점수 측면에서 정확한 어텐션보다 뛰어난 생성 품질을 달성했다. FLOPs는 4.14배 적게 사용하였다.

ABSTRACT

Dot-product attention mechanism plays a crucial role in modern deep architectures (e.g., Transformer) for sequence modeling, however, naïve exact computation of this model incurs quadratic time and memory complexities in sequence length, hindering the training of long-sequence models. Critical bottlenecks are due to the computation of partition functions in the denominator of softmax function as well as the multiplication of the softmax matrix with the matrix of values. Our key observation is that the former can be reduced to a variant of the kernel density estimation (KDE) problem, and an efficient KDE solver can be further utilized to accelerate the latter via subsampling-based fast matrix products. Our proposed KDEformer can approximate the attention in sub-quadratic time with provable spectral norm bounds, while all prior results merely provide entry-wise error bounds. Empirically, we verify that KDEformer outperforms other attention approximations in terms of accuracy, memory, and runtime on various pre-trained models. On BigGAN image generation, we achieve better generative scores than the exact computation with over $4 imes$ speedup. For ImageNet classification with T2T-ViT, KDEformer shows over $18 imes$ speedup while the accuracy drop is less than $0.5\%$.

연구 동기 및 목표

• 트랜스포머의 자기어텐션 메커니즘에서 발생하는 이차 계산 및 메모리 병목 현상을 해결하기 위해.
• 단순히 항목별 오차 경계가 아니라 증명된 스펙트럴 노름 오차 보장을 갖는 어텐션 근사 방법을 개발하기 위해.
• 계산 복잡도를 감소시켜 장거리 시퀀스 모델의 효율적인 추론 및 훈련을 가능하게 하기 위해.
• 기존의 근사 어텐션 방법들보다 비전 및 언어 모델에서 생성 성능 및 효율성을 향상시키기 위해.

제안 방법

• 다이아몬드 스케일링 행렬 D와 어텐션 행렬 A에 대해 모두 일반화된 커널 밀도 추정(KDE) 문제로 어텐션 계산을 축소한다.
• KDE 기반 확률을 사용한 중요도 샘플링을 통해 근사 어텐션 계산을 위한 작은 열 부분집합을 선택한다.
• 지속적인 최신 기술의 KDE 솔버를 활용하기 위해 지수 커널 밀도 문제를 가우시안 KDE 문제로 변환하는 변환을 적용한다.
• 전체 어텐션 행렬 A의 명시적 계산을 피하기 위해 KDE를 통해 다이아몬드 스케일링 행렬 D를 추정한다.
• 추정된 샘플링 행렬 Π와 대각행렬 eD를 사용한 부분 샘플링 행렬 곱셈을 적용하여 출력을 효율적으로 계산한다.
• 어 attention 출력의 안정성과 정밀도를 보장하는 스펙트럴 노름 오차 경계(식 1)를 제공한다.

실험 결과

연구 질문

• RQ1 어떻게 어텐션 메커니즘을 커널 밀도 추정 문제로 재구성하여 서브-이차 계산을 가능하게 할 수 있는가?
• RQ2 KDE 기반 샘플링과 스케일링은 어텐션의 계산 복잡도를 감소시키면서도 모델 성능을 유지할 수 있는가?
• RQ3 어떻게 어텐션 근사에서 스펙트럴 노름 오차 경계를 달성할 수 있으며, 이는 항목별 경계보다 더 강력한 이론적 보장을 제공하는가?
• RQ4 제안된 방법은 장거리 시퀀스 작업에서 정확도, 속도, 메모리 효율성 측면에서 기존의 근사 어텐션 메커니즘을 초월하는가?
• RQ5 KDEformer는 FLOPs를 줄이며 BigGAN과 같은 모델의 생성 품질을 향상시킬 수 있는가?

주요 결과

• BigGAN 이미지 생성 작업에서, KDEformer는 정확한 어텐션보다 FID 점수를 낮게 유지하면서도 FLOPs를 4.14배 줄였다.
• T2T-ViT를 사용한 ImageNet 분류에서, KDEformer는 정확도가 82.08%로 정확한 계산 대비 0.47% 떨어졌지만 FLOPs는 18.3배 감소시켰다.
• Long Range Arena에서의 엔드 투 엔드 훈련에서, KDEformer는 텍스트 데이터셋에서 정확한 어텐션 대비 8배 빠르게 실행되었으며 정확도 저하가 최소였다.
• LRA 데이터셋 전반에서 KDEformer는 평균 2.28 GB의 최대 메모리를 사용했으며, 정확한 계산(9.21 GB)과 경쟁자들보다 크게 낮았다.
• LRA에서 평균 테스트 정확도 57.15%를 기록하며 두 번째로 높은 성능을 달성했고, 모든 근사 방법 중에서 가장 빠르고 메모리 효율성이 뛰어났다.
• 실험 결과는 KDEformer가 Reformer, Performer, ScatterBrain보다 여러 벤치마크에서 정확도와 효율성 측면에서 뛰어나다는 것을 확인했다.