[논문 리뷰] Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

Meng-Hao Guo, Zheng-Ning Liu외 2명|ArXiv.org|2021. 05. 05.

Advanced Neural Network Applications참고 문헌 98인용 수 79

한 줄 요약

외부 어텐션을 두 개의 외부 메모리로 구현한 작은 학습 가능한 선형 계층으로 제안하고, 선형 복잡도와 다양한 비전 태스크에서 경쟁력 있는 결과를 달성하며, EAMLP로 불리는 all-MLP 변형까지 포함한다.

ABSTRACT

Attention mechanisms, especially self-attention, have played an increasingly important role in deep feature representation for visual tasks. Self-attention updates the feature at each position by computing a weighted sum of features using pair-wise affinities across all positions to capture the long-range dependency within a single sample. However, self-attention has quadratic complexity and ignores potential correlation between different samples. This paper proposes a novel attention mechanism which we call external attention, based on two external, small, learnable, shared memories, which can be implemented easily by simply using two cascaded linear layers and two normalization layers; it conveniently replaces self-attention in existing popular architectures. External attention has linear complexity and implicitly considers the correlations between all data samples. We further incorporate the multi-head mechanism into external attention to provide an all-MLP architecture, external attention MLP (EAMLP), for image classification. Extensive experiments on image classification, object detection, semantic segmentation, instance segmentation, image generation, and point cloud analysis reveal that our method provides results comparable or superior to the self-attention mechanism and some of its variants, with much lower computational and memory costs.

연구 동기 및 목표

• 시각 태스크에서 self-attention의 2차 복잡도와 샘플에 의존하지 않는 특성을 동기부여하고 해결한다.
• 데이터셋 수준의 외부 메모리 유닛을 갖춘 external attention을 도입하여 데이터셋 수준의 상관관계를 포착한다.
• 인기 있는 아키텍처에서 self-attention을 더 낮은 계산 및 메모리 비용으로 대체할 수 있음을 보여준다.
• 이미지 분류, 탐지, 분할, 생성, 3D 포인트 클라우드 태스크 전반에 걸친 external attention의 다재다능함을 입증한다.
• 경쟁력 있는 성능을 가진 all-MLP 아키텍처를 만들기 위해 multi-head external attention(EAMLP)을 제안한다.

제안 방법

• M_k와 M_v를 키 및 값 메모리로 사용하는 external attention을 정의한다.
• A = Norm(F M_k^T)이고 F_out = A M_v로 주의를 계산하며, M_k와 M_v를 구현하는 선형 계층을 사용한다.
• 행과 열 전반에서 어텐션 스코어를 안정시키기 위해 이중 정규화를 사용한다.
• 더 풍부한 표현을 위해 multi-head external attention으로 확장한다.
• 기존 아키텍처에 external attention을 도입하고 all-MLP 모델(EAMLP)을 구축한다.

실험 결과

연구 질문

• RQ1 external attention이 시각 아키텍처에서 self-attention을 선형 계산 비용으로 대체할 수 있는가?
• RQ2 데이터셋 수준의 external memories를 도입하면 다양한 비전 태스크에서 일반화 및 성능이 향상되는가?
• RQ3 멀티-헤드 external attention(MEA)가 정확도와 효율성 면에서 self-attention 및 다른 어텐션 변형과 어떻게 비교되는가?
• RQ4 external attention이 ImageNet에서 CNN/Transformer 성능에 맞는 모든-MLP 비전 모델을 가능하게 할 수 있는가?
• RQ5 정규화 전략이 external attention의 안정성 및 성능에 미치는 영향은 무엇인가?

주요 결과

• External attention은 계산 및 메모리 사용이 더 적은 비용으로, 태스크 전반에서 self-attention과 동등하거나 우수한 성능을 달성한다.
• 작은 공유 메모리(S ~ 64 등)를 사용하면 입력 크기에 대해 선형 복잡도(O(dSN))를 얻는다.
• Multi-head external attention은 competitive ImageNet 정확도(보고된 설정에서 상위 1 등 79.4%까지)를 가진 all-MLP 아키텍처(EAMLP)를 가능하게 한다.
• self-attention을 external attention으로 교체하면 backbone 네트워크에 통합될 때 여러 벤치마크(VOC, COCO 등)에서 분할 및 탐지 지표가 향상된다.
• External attention은 의미 있는 객체와 영역에 초점을 보이는 해석 가능한 어텐션 맵을 제공하며, 헤드가 서로 다른 영역에 주의를 기울인다.