[논문 리뷰] SimA: Simple Softmax-free Attention for Vision Transformers

Soroush Abbasi Koohpayegani, Hamed Pirsiavash|arXiv (Cornell University)|2022. 06. 17.

CCD and CMOS Imaging Sensors인용 수 8

한 줄 요약

이 논문은 시각 Transformer에서 계산 비용이 높은 Softmax를 제거하고 쿼리 및 키팽의 ℓ₁-정규화를 도입함으로써 효율성을 향상시키는 SimA라는 Softmax-무료 attention 메커니즘을 제안한다. 이는 입력 해상도에 따라 O(ND) 또는 O(D²N)의 복잡도로 복잡도를 감소시키며, 추론 시 동적 계산 순서를 결정함으로써 구현된다. SimA는 ImageNet, MS-COCO, 그리고 자기지도 학습 벤치마크에서 최신 기준 성능을 달성하면서도 exp(.) 연산을 완전히 제거하여 엣지 디바이스에서의 효율성을 크게 향상시킨다.

ABSTRACT

Recently, vision transformers have become very popular. However, deploying them in many applications is computationally expensive partly due to the Softmax layer in the attention block. We introduce a simple but effective, Softmax-free attention block, SimA, which normalizes query and key matrices with simple $\ell_1$-norm instead of using Softmax layer. Then, the attention block in SimA is a simple multiplication of three matrices, so SimA can dynamically change the ordering of the computation at the test time to achieve linear computation on the number of tokens or the number of channels. We empirically show that SimA applied to three SOTA variations of transformers, DeiT, XCiT, and CvT, results in on-par accuracy compared to the SOTA models, without any need for Softmax layer. Interestingly, changing SimA from multi-head to single-head has only a small effect on the accuracy, which simplifies the attention block further. The code is available here: https://github.com/UCDvision/sima

연구 동기 및 목표

• 엣지 디바이스에 배포하기 위해 시각 Transformer의 계산 비용, 특히 Softmax 내부의 비용이 높은 exp(.) 연산을 줄이기 위해.
• attention 메커니즘에서 Softmax 레이어를 제거함에도 불구하고 모델 정확도를 유지하거나 향상시키기 위해.
• 토큰 수(N)와 채널 수(D)에 기반하여 추론 시 동적 계산 재정렬을 가능하게 하여 효율성을 최적화하기 위해.
• 다중 헤드와 GELU를 대체할 수 있는 단일 헤드 attention과 ReLU 활성화 함수의 타당성을 탐색하기 위해.
• 수치적 안정성과 반정밀도 호환성을 유지하면서도 순열 불변성과 토큰 중요도 해석을 유지하는 attention 메커니즘을 제공하기 위해.

제안 방법

• self-attention의 Softmax 연산을 쿼리 및 키팽 행렬의 ℓ₁-정규화 이후의 내적 연산으로 대체한다.
• attention을 단순한 행렬 곱셈으로 계산: O = (Q_normalized) @ (K_normalized.T) @ V로, exp(.) 및 Softmax가 필요 없도록 한다.
• 추론 시 동적으로 계산 순서를 선택한다: N > D 이면 (Q @ K.T) @ V, D > N 이면 Q @ (K.T @ V)로 FLOPs를 최소화한다.
• ℓ₁-정규화를 통해 수치적 안정성을 확보하여 반정밀도 추론 시 오버플로우 없이 구현 가능하도록 한다.
• 토큰을 순서 없는 집합으로 간주함으로써 순열 불변성을 유지하여, 다중 소스 입력(예: 여러 카메라 또는 센서)과도 사용할 수 있도록 한다.
• 아키텍처 변경 없이 DeiT, XCiT, CvT 모델에 SimA 블록을 적용하며, 표준 학습 프로토콜을 유지한다.

실험 결과

연구 질문

• RQ1 Softmax-무료 attention 메커니즘은 시각 Transformer에서 최신 기준 성능을 유지할 수 있는가?
• RQ2 쿼리 및 키팽에 대한 ℓ₁-정규화로 Softmax를 대체할 경우 정확도는 유지되며 계산 비용은 감소하는가?
• RQ3 N과 D에 기반한 추론 시 동적 계산 순서 정렬이 효율성 향상에 상당한 기여를 할 수 있는가?
• RQ4 Softmax가 제거된 상황에서 단일 헤드 attention과 ReLU는 다중 헤드와 GELU에 비해 성능에 큰 영향을 주는가?
• RQ5 정규화된 쿼리 및 키팡 벡터의 크기 값이 모델의 해석을 위한 의미 있는 시각화 지도로 기능할 수 있는가?

주요 결과

• SimA는 ImageNet 분류, MS-COCO 객체 검출 및 세그멘테이션, 그리고 자기지도 학습 벤치마크에서 최신 기준 성능을 달성하며, 정확도는 유사 수준을 유지한다.
• 1536×1536 해상도 이미지에서 SimA는 DeiT 대비 58% 빠르게 동작한다. 이는 Softmax의 exp(.) 오버헤드 제거 덕분이다.
• 모델 차원이 8192일 경우, SimA는 XCiT 대비 추론 시간을 22% 감소시킨다. 이는 Softmax 계산을 회피했기 때문이다.
• 단일 헤드 SimA는 다중 헤드 최신 기준 모델과 유사한 정확도를 달성하여 헤드 수 감소에 따른 성능 손실가 최소화됨을 시사한다.
• GELU 대신 ReLU를 사용할 경우 정확도 저하가 미미하여, 활성화 함수 변경에 대해 높은 내구성을 보임을 시사한다.
• 정규화된 쿼리 및 키팡 벡터의 ℓ₂-노름은 표준 DeiT에서 관찰되는 평탄한 분포와는 달리 중요한 이미지 영역을 강조하는 의미 있는 시각화 지도를 생성한다.