MaxViT: 다중 축 비전 트랜스포머

최근 Transformer는 컴퓨터 비전 분야에서 큰 주목을 받고 있다. 그러나 이미지 크기에 비례하여 확장성이 떨어지는 자기 주의(multi-head self-attention) 메커니즘의 한계로 인해 최신 비전 백본에서의 광범위한 채택이 제한되어 왔다. 본 논문에서는 선형 복잡도로 임의의 입력 해상도에서 전역-국소 공간 상호작용을 가능하게 하는, 효율적이고 확장 가능한 주의 모델인 ‘다축 주의(Multi-axis Attention)’를 제안한다. 이 모델은 두 가지 구성 요소로 이루어져 있다: 블록화된 국소 주의와 확장된 전역 주의. 이러한 설계 선택을 통해 전역과 국소 정보 간의 상호작용을 효율적으로 구현할 수 있다. 또한 제안된 주의 모델을 컨볼루션과 효과적으로 융합함으로써 새로운 아키텍처 요소를 도입하였으며, 이를 기반으로 여러 단계에 걸쳐 기본 블록을 단순 반복함으로써 단순한 계층 구조의 비전 백본인 MaxViT를 제안한다. 특히, MaxViT는 초기 고해상도 단계에서도 전체 네트워크 내에서 전역적인 시각을 유지할 수 있다. 다양한 비전 작업에서 본 모델의 효과성을 입증하였다. 이미지 분류 작업에서는 다양한 설정에서 최신 기준을 달성한다. 추가 데이터 없이도 MaxViT는 ImageNet-1K 데이터셋에서 86.5%의 top-1 정확도를 달성하였으며, ImageNet-21K에서 사전 훈련한 경우 88.7%의 top-1 정확도를 기록하였다. 하류 작업에서는 객체 탐지 및 시각적 미학 평가(task)에서도 우수한 성능을 보였다. 또한 본 모델이 ImageNet 데이터셋에서 강력한 생성 모델링 능력을 갖추고 있음을 보여주며, MaxViT 블록이 일반적인 비전 모듈로서의 높은 잠재력을 지닌다는 것을 입증하였다. 소스 코드 및 훈련된 모델은 https://github.com/google-research/maxvit 에 공개될 예정이다.