Swin Transformer : Hierarchical Vision Transformer using Shifted Windows

[논문리뷰]

ICCV 2021 

---

ABSTRACT

• 저자들은 기존 텍스트의 단어와 이미지의 픽셀 해상도와의 도메인 차이를 해결하고자 Shifted windows로 표현이 계산되는 계층적인 트랜스포머를 제안한다.
• Shifted windows 방식은 self-attention 계산을 중첩되지 않도록 제한함으로써 효율성을 높이는 동시에 cross-windows connetion을 허용함으로써 효율성을 높인다.
• 이러한 방식은 다양한 스케일로 모델링할 수 있는 유연성을 가지며 이미지 크기와 관련하여 선형 계산 복잡성을 가지고 있다.

Introduction

그림1

• (a)는 논문에서 제안한 Swin Transformer의 구조이다. 각 로컬 윈도우(빨간색 표시) 내에서만 자체 주의 계산으로 인해 이미지 크기를 입력하기 위한 선형 계산 복잡성을 갖는다. 
• 본 논문에서는 NLP 및 CNN이 시각에서 하는 것처럼 컴퓨터 비전을 위한 backbone 역할을 할 수 있도록 트랜스포머의 적용 가능성을 확장하고자 한다.
• 기존의 트랜스포머는 토큰을 모두 고정된 스케일로 이러한 vision task 적합하지 않은 특성이라고 저자는 말한다.
  
  또 다른 문제점은 텍스트 부분의 단어에 비해 이미지의 픽셀 해상도가 훨씬 높다는 것이다. 픽셀 수준에서 밀도 높은 예측이 필요한 분할과 같은 비전 작업이 많이 존재하며, self-attention의 계산 복잡성이 이미지 크기 때문에 고해상도 이미지의 트랜스포머에서는 다루기 어려울 것이다.
  
  이러한 문제점을 해결하고자 계층적 특징 맵을 구성하고 이미지 크기에 대한 선형 계산 복잡성을 갖는 범용 트랜스포머 백본인 Swin Transformer를 제안하였다.
• 그림 1-(a)에 나타낸 바와 같이, Swin Transformer는 작은 크기의 패치(회색으로 표시된)에서 시작하여 더 깊은 Transformer 계층에서 이웃 패치를 점진적으로 병합하여 계층적 표현을 구성한다.
  
  계산 복잡도는 이미지를 분할하는 중첩되지 않은 창(빨간색으로 표시된) 내에서 로컬로 자체 주의를 계산된다. 각 창에 있는 패치 수는 고정되어 있으므로 복잡성은 이미지 크기에 선형으로 된다.
  
  이러한 장점은 Swin Transformer를 단일 해상도의 특징 맵을 생성하고 2차 복잡성을 갖는 이전의 Transformer 기반 아키텍처와 달리 다양한 비전 작업에 범용 백본으로 적합하다고 저자는 말한다.
  

그림2

• Swin Transformer의 핵심 요소는 그림2와 같이, 연속적인 self-attention 계층들 사이의 윈도우 파티션의 이동이다.
  
  이동된 윈도우는 이전 계층의 윈도우들을 브리지하고, 윈도우 사이에 모델링 능력을 상당히 향상시키는 연결들을 제공한다고 한다.
  윈도우 내의 모든 쿼리 패치들은 동일한 Key 세트을 공유하고, 이는 하드웨어에서 메모리 액세스를 용이하게 하는 장점이 있다.
•  

---

Overall Architecture

그림3

• 위 구조처럼 Stage1을 거치고 계층적 표현을 생성하기 위해, 토큰의 수는 네트워크가 깊어질수록 Patch Merging 레이어에 의해 감소된다.
• W-MSA는 Window 안에서의 self-attention
• SW-MSA는 Window 간의 self-attention

3.1 Swin Transformer Block

• Swin Transformer는 Transformer 블록의 표준 MSA(multi-head self attention) 모듈을 Shifted window 기반의 모듈로 대체하여 구축되며, 다른 레이어는 동일하게 유지된다.

3.2 Shifted Window based Self-Attention

• 기존 Transformer는 모두 토큰과 다른 모든 토큰 사이의 관계가 계산되는 global self-attention을 수행한다. 글로벌 계산은 토큰 수에 대해 2차 복잡성을 초래하므로 밀도 높은 예측을 위해 또는 고해상도 이미지를 표현하기 위해 막대한 토큰 세트가 필요한 많은 비전 문제에 적합하지 않는다.

Self-attention in non-overlapped windows

• W-MAS의 경우 윈도우 안에서만 attention score를 계산하므로 연산량은 위 사진과 같다.
• MSA는 패치 수 hw에 2차이고 W-MSA는 M이 고정될 때 선형입니다(논문에서는 기본적으로 7로 설정됨)
• 상세한 연산 과정은 아래 블로그를 참고하였다.
  
  https://velog.io/@temoc/%EB%85%BC%EB%AC%B8%EB%A6%AC%EB%B7%B0-Swin-Transformer-Hierarchical-Vision-Transformer-using-Shifted-Windows#footnote_3

Shifted window partitioning in successive blocks 

• 윈도우 기반의 자체 attention 모듈은 윈도우 간 연결이 부족하여 모델링 능력에 한계가 있다. 따라서 overlapping이 아닌 윈도우의 효율적인 계산을 유지하기 위해 연속적인 Swin Transformer 블록에서 두 분할 구성을 교대로 사용하는 shifted window partitioning 접근법을 제안한다. 

Efficient batch computation for shifted configuration

그림4

• 그림 4와 같이 왼쪽 상단 방향으로 순환 shifting함으로써 보다 효율적인 배치 계산 접근법을 제안합니다. 이 시프트 후에 배치 윈도우는 특징 맵에 인접하지 않은 여러 하위 윈도우로 구성될 수 있으므로 마스킹 메커니즘을 사용하여 자체 attention 계산을 각 하위 윈도우 내로 제한한다.
• 그림4 처럼 Mask된 파티션은 attention을 다시 수행하지 않는다. 즉 원본 이미지에서 서로 인접하지 않은 부분은 Mask 처리를 진행 하는 것이다.

Relative position bias

• $M^2$ X $d$ 차원($M$은 윈도우 내 패치의 개수)을 가진 Q, K, V
• relative position bias $B$를 추가하여 연산한다.
• 입력에 position 임베딩을 추가하면 성능이 약간 떨어지므로 구현에 채택되지 않았다고 한다. 
  (SWIN은 ViT와 다르게 encoder 입력 부분에서 position embedding을 하지 않는다.)

---

Experiments

• ImageNet 1k, COCO, ADE20K에 대하여 Classification, object detection, sementic segmentation에 대하여 실험을 진행하였다.

표1 - Image Classfication

표2 - Object Detection

표3 - Sementic Segmentation

• Ablation study on the shifted windows approach and different position embedding methods