[논문 리뷰] KVT: k-NN Attention for Boosting Vision Transformers

Pichao Wang, Xue Wang외 5명|arXiv (Cornell University)|2021. 05. 28.

Advanced Neural Network Applications참고 문헌 62인용 수 4

한 줄 요약

이 논문은 각 쿼리에 대해 가장 유사도가 높은 상위-k개의 키만 선택함으로써 밀도 높은 자기주의 어텐션을 대체하는 경량이며 미분 가능한 메커니즘인 k-NN 어텐션을 제안한다. 이는 국소적 인덕티브 바이어스를 강화하고 노이즈가 있는 토큰을 걸러내며, 아키텍처 수정 없이도 11종의 비전 트랜스포머 아키텍처에서 학습 속도를 향상시키고 성능을 향상시킨다.

ABSTRACT

Convolutional Neural Networks (CNNs) have dominated computer vision for years, due to its ability in capturing locality and translation invariance. Recently, many vision transformer architectures have been proposed and they show promising performance. A key component in vision transformers is the fully-connected self-attention which is more powerful than CNNs in modelling long range dependencies. However, since the current dense self-attention uses all image patches (tokens) to compute attention matrix, it may neglect locality of images patches and involve noisy tokens (e.g., clutter background and occlusion), leading to a slow training process and potential degradation of performance. To address these problems, we propose the $k$-NN attention for boosting vision transformers. Specifically, instead of involving all the tokens for attention matrix calculation, we only select the top-$k$ similar tokens from the keys for each query to compute the attention map. The proposed $k$-NN attention naturally inherits the local bias of CNNs without introducing convolutional operations, as nearby tokens tend to be more similar than others. In addition, the $k$-NN attention allows for the exploration of long range correlation and at the same time filters out irrelevant tokens by choosing the most similar tokens from the entire image. Despite its simplicity, we verify, both theoretically and empirically, that $k$-NN attention is powerful in speeding up training and distilling noise from input tokens. Extensive experiments are conducted by using 11 different vision transformer architectures to verify that the proposed $k$-NN attention can work with any existing transformer architectures to improve its prediction performance. The codes are available at \url{https://github.com/damo-cv/KVT}.

연구 동기 및 목표

• 모든 토큰을 집계하는 밀도 높은 자기주의 어텐션으로 인해 발생하는 느린 학습과 노이즈 민감도 문제를 해결한다.
• 전체 연결 어텐션의 국소적 이미지 구조를 포착하거나 혼잡한 배경이나 가림을 다루는 데서 한계를 극복한다.
• 전연성 장기 모델링을 유지하면서 인덕티브 바이어스를 향상시키는 간단하면서도 효과적인 컨볼루션 인덕티브 바이어스의 대안을 제안한다.
• k-NN 어텐션이 어떤 비전 트랜스포머 아키텍처에도 일반적으로 적용되어 정확도와 학습 속도를 향상시킬 수 있음을 입증한다.
• 이론적 및 실증적으로 k-NN 어텐션이 아키텍처 변경 없이 노이즈를 줄이고 수렴 속도를 높임을 검증한다.

제안 방법

• 표준 밀도 높은 자기주의 어텐션을 k-NN 어텐션으로 대체하며, 각 쿼리 토큰에 대해 가장 유사도가 높은 상위-k개의 키 토큰만 사용하여 어텐션 점수를 계산한다.
• 쿼리와 키 벡터 간의 유사도를 내적 어텐션을 통해 계산하고, 코사인 또는 내적 유사도 기반으로 k개의 가장 가까운 키를 선택한다.
• 어텐션 맵은 오직 선택된 k개의 키에 대해서만 계산되며, 이로 인해 계산 비용이 감소하고 관련 없거나 노이즈가 있는 토큰이 걸러진다.
• k-NN 메커니즘은 미분 가능하고 엔드 투 엔드로 학습 가능하므로, 아키텍처 변경 없이도 어떤 비전 트랜스포머에나 통합될 수 있다.
• 명시적인 컨볼루션 연산 없이도 근처 토큰이 더 자주 선택되므로, 자연스럽게 국소적 인덕티브 바이어스를 계승한다.
• k-NN 메커니즘은 특징 유사도에 기반해 집계 그래프를 동적으로 조정함으로써 장거리 의존성을 허용하면서도 먼 거리에 있는 유사도가 낮은 토큰을 억제한다.

실험 결과

연구 질문

• RQ1 밀도 높은 자기주의 어텐션을 k-NN 어텐션으로 대체하면 비전 트랜스포머의 학습 속도와 모델의 강건성이 향상되는가?
• RQ2 k-NN 어텐션은 학습 도중 배경 혼잡이나 가림 영역과 같은 노이즈가 있는 토큰을 효과적으로 걸러내는가?
• RQ3 표준 자기주의 어텐션과 비교해 k-NN 어텐션은 국소적 인덕티브 바이어스를 어느 정도 유지하거나 향상시키는가?
• RQ4 k-NN 어텐션은 아키텍처 변경 없이 다양한 비전 트랜스포머 아키텍처에 일반적으로 적용 가능한가?
• RQ5 성능 및 학습 역학 측면에서 k-NN 어텐션은 컨볼루션 인덕티브 바이어스와 어떻게 비교되는가?

주요 결과

• k-NN 어텐션은 11종의 비전 트랜스포머 아키텍처에서 학습 속도를 크게 향상시켜 수렴 시간을 단축시키지만 최종 정확도에 영향을 주지 않는다.
• 배경 혼잡이나 가려진 영역과 같은 노이즈가 있는 토큰을 걸러내어 일반화 성능을 향상시키며, 더 강건한 특징 학습을 이끈다.
• k-NN 어텐션은 ImageNet-1K에서 최고 성능 또는 경쟁력 있는 성능을 달성하며, 다양한 백본 아키텍처에서 일관된 향상이 관찰된다.
• 이론적 분석에 따르면, k-NN 어텐션은 외곽선과 노이즈 특징의 영향을 통제함으로써 강력한 일반화 보장을 유지한다.
• 실증 결과는 k-NN 어텐션이 어텐션 맵의 분산을 줄이고 어텐션 맵의 흩어짐을 향상시켜 더 집중적이고 의미 있는 어텐션 헤드를 만들어내는 것으로 확인되었다.
• 이 방법은 일반적으로 적용 가능하다: 아키텍처 변경 없이도 어떤 비전 트랜스포머에나 삽입할 수 있으며, 다양한 모델과 데이터셋에서 일관되게 성능 향상을 이룬다.