[논문] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

본 논문에서는 vision Transformer의 일종인 Swin Transformer를 제시한다. 이 Transformer는 Shifted window를 통해 계층적인 representation을 계산한다. Shifted windowing 기법은 중복되지 않는 local window에서만 self-attention을 계산해 효율성을 높이면서 cross-window 연결을 활용한다. 이 계층적인 구조는 다양한 scale의 이미지를 유연하게 나타내고 이미지 크기에 따라 linear 계산 복잡도를 가지게 한다.

 

 

 

1. Introduction

Computer vision의 모델링은 오랫동안 convolutional neural networks (CNN)으로 모델링 되어왔다.

하지만 natural language processing (NLP)분야는 다른 network 구조를 통해 진화해왔고 현재 가장 많이 쓰이는 구조는 Transformer이다.

Transformer는 데이터의 long-range dependency들을 잘 나타낼 수 있다. 이런 language 분야에서의 성공은 computer vision에서의 Transformer 활용 연구로 이어지게 된다.

 

언어 분야에서의 높은 성과를 이미지로 옮기는데는 다음과 같은 문제점이 있다.

• Scale 문제
  • Language Transformer에서 word 토큰을 고정된 기본 단위가 되는 것에 비해
  • Visual 단위는 다양한 scale에 따라 크게 차이가 난다
• High resolution 문제
  • Passage에 있는 단어에 비해 이미지의 pixel을 더 높은 resolution을 가진다
  • 따라서 self-attention 계산 복잡도는 이미지 크기에 제곱으로 커진다

위 문제들을 해결하기 위해 Swin Transformer를 제시한다.

• 계층적인 feature map을 만들어 이미지 크기에 대해 선형적인 계산 복잡도를 가진다
  • 먼저 작은 크기를 가진 patch부터 시작해 이웃 patch들을 합치게된다.
  • Self-attention을 이미지를 나눈 local partition (빨간색 칸)에서만 계산해 선형 계산 복잡도를 가지게 한다.

 

Swin Transformer의 주요 디자인 요소는 연속적인 self-attention layer 간 window shift이다.

• Shifted된 window는 이전 layer와 연결되어 모델링 파워를 크게 증진시킨다.
• 또한 모든 query patch들은 같은 key set을 공유하기 때문에 메모리 접근을 용이하게 한다.
  • 따라서 shifted window 기법이 sliding window 기법에 비해 더 낮은 latency를 가지게 한다

제시된 Swin Transformer 모델은 다양한 vision 작업에서 기존 SOTA를 뛰어넘는 좋은 성능을 보인다.

 

 

 

2. Related Work

CNN and variants

AlexNet 이후 CNN은 하나의 독립적인 분야로 연구된다.

더 깊고 효율적인 CNN구조인 VGG, GoogleNet, ResNet, DenseNet, HRNet, EfficientNet 등이 제시되었다.

또 각 convolution layer를 개선한 depthwise convolution과 deormable convolution도 제안되었다.

 

이런 CNN이 computer vision 응용에 기본적인 backbone 구조로 사용되었지만 Transformer-like 구조의 높은 가능성 또한 존재한다.

 

 

Self-attention based backbone architectrures

Transformer의 self-attention 구조에 영감받아 ResNet의 몇몇 또는 모든 layer를 self-attention layer로 바꾸는 연구들도 존재한다.

• 이런 연구들에서는 local window 안 pixel들에 대해서 self-attention을 계산해 최적화하고 기존 ResNet구조에 비해 조금 더 높은 acurracy/FLOPs trade-off를 보여준다
• 실제로는 비싼 memory access 때문에 기존 convolution network보다 더 높은 latency를 보인다

 

Self-attention/Transformers to complement CNN

CNN구조를 self-attention이나 Tranformer를 이용해 보강하는 연구들이다.

• 기존 backbone 구조나 head network들을 self-attention layer로 보완한다
  • Tranformer가 멀리 떨어진 의존도나 상호작용을 잘 encode하는 능력을 활용한다

 

Transformer based vision backbones

본 논문과 가장 관련성이 높은 연구는 Vision Transformer (ViT)와 그 후속 연구들이다.

• ViT는 곂치지 않는 중간 크기 image patch들을 사용해 image classification을 진행한다
• DeiT는 많은 training dataset을 필요로 하는 ViT와 다르게 작은 dataset을 이용한 효율적인 학습 방법을 제시한다

하지만 ViT는 general-purpose backbone으로 적합하지 않다.

• Dense vision 작업 (object detection, semantic segmentation)에 적합하지 않다
• High resolution 이미지가 주어질 때 복잡도가 이미지 크기의 제곱으로 커지게 된다

Swin Transformer 구조는 classification 작업에 특화되어있지 않음에도 가장 좋은 speed-accuracy trade-off 결과를 보여준다.

또한 COCO object detection과 ADE20K semantic segmentation에도 SOTA 성능을 보여준다.

 

 

 

3. Method

3.1 Overall Architecture

전체적인 모델 구조는 위 그림과 같다 (Swin-T)

• 먼저 input RGB 이미지를 곂치지 않는 patch들로 나눈다
  • 각 patch는 하나의 token으로 취급되고 feature는 pixel RGB 값을 연결해 나타내진다
  • 구현에서는 4x4 patch 크기를 사용한다 (각 feature: 4x4x3=48)
• Stage 1
  • Linear embedding layer를 사용해 각 feature가 특정 차원으로 embed되게 한다
  • 각 patch token에 대해서 변형된 self-attention 계산을 수행한다 (Swin Transformer block)
• Stage 2
  • 계층적인 representation을 생성하기 위해서 patch merging을 통해 layer가 깊어질수록 token 수는 줄어들게 된다
  • 첫번째 patch merging layer는 2x2 이웃하는 patch들을 하나의 그룹으로 묶는다
    • 그리고 연결된 4C 차원 features에 linear layer를 적용한다
    • 이는 token 수를 4배로 줄이고 output 차원을 2C로 만든다
  • Feature transformation 이후에 Swin Transformer block을 적용한다
    • 이때 출력 resolution은 H/8 x W/8로 고정된다
• Stage 2를 반복해 Stage 3(output: H/16 x H/16)과 Stage 4(output: H/32 x H/32)를 만든다

연결된 stage들은 계층적인 구조를 만들고 각 feature map 크기는 일반적인 convolutional network와 동일하다.

 

 

Swin Transformer block

Swin Transformer는 기존 multi-head self attention (MSA) 모듈을 shifted windows에 기반한 모듈로 대체한다.

• 위 Figure 3(b)에서 묘사되어 있듯이 Swin Tansformer block은
  • MSA에 기반한 shifted window
  • 2-layer MLP와 GELU non-linearity
  • 각 MSA와 MLP 이전에 적용된 LayerNorm (LN)
  • 그리고 residual connection으로 구성되어 있다

 

 

3.2 Shifted Window based Self-Attention

 기존 Transformer 구조와 image classification을 위한 변형들은 모두 global self-attention을 계산한다.

• 따라서 이미지 크기의 제곱으로 복잡도가 커지게 된다.

 

Self-attention in non-overlapped windows

더 효율적인 모델링을 위해 local window 안에서만 self-attention을 계산하는 기법을 제시한다.

• 각 window는 동일하게 나눠져 곂치지 않게 이미지에 배치된다
• 각 window가 M x M patch를 가지고 있다 가정하면 global 계산 복잡도와 window-based 계산 복잡도는 다음과 같다:
  • $ \omega(MSA)=4hwC^2+2(hw)^2C$
  • $\omega(W-MSA)=4hwC^2+2M^2hwC$
  • 전자는 hw에 대해 quadratic한 복잡도를 가지고 후자는 linear한 복잡도를 가진다

 

Shifted window partitioning in successive blocks

제시된 window-based self-attention은 window 간 연결점이 부족하다.

따라서 효율적인 계산을 그대로 사용하는 corss-window connection 기법을 소개한다.

• 연속적인 Swin Transformer block들에 대해 서로 다른 두 개의 partitioning 구성을 사용한다

위 Figure 2에 묘사되어 있듯이 

• 첫번째 모듈은 top-left pixel부터 8x8 feature map을 4x4 크기의 window로 2x2 개 생성한다. (M=4)
• 다음 모듈은 window를 이전 layer에 비해 ($ \left( \left\lfloor \frac{M}{2} \right\rfloor , \left\lfloor \frac{M}{2} \right\rfloor \right) $) 만큼 shift한다

 

Shifted window 접근 방법을 사용해 연속적인 Swin Transformer block은 다음과 같이 계산된다:

$$ \hat{\mathbf{z}}^l = \text{W-MSA} \left( \text{LN} \left( \mathbf{z}^{l-1} \right) \right) + \mathbf{z}^{l-1}, $$

$$ \mathbf{z}^l = \text{MLP} \left( \text{LN} \left( \hat{\mathbf{z}}^l \right) \right) + \hat{\mathbf{z}}^l, $$

$$ \hat{\mathbf{z}}^{l+1} = \text{SW-MSA} \left( \text{LN} \left( \mathbf{z}^l \right) \right) + \mathbf{z}^l, $$

$$ \mathbf{z}^{l+1} = \text{MLP} \left( \text{LN} \left( \hat{\mathbf{z}}^{l+1} \right) \right) + \hat{\mathbf{z}}^{l+1}, $$

• $\hat{z^l}$: W-MSA 모듈의 output feature
• $z^l$: MLP 모듈의 output feature
• W-MSA: Window based MSA
• SW-MSA: Shifted window based MSA

 

Efficient batch computation for shifted configuration

위 제시된 shifted window 기법의 문제는 더 많은 window와 M x M 보다 더 작은 크기의 window를 가지게 된다는 것이다.

• 간단한 해결 방법으로 작은 window를 M x M 크기로 pad하고 attention을 계산할 때 pad된 부분은 masking 하는 방법이다
  • 이 방법은 기존 window partitioning이 작을 때 고려할만하다 (ex. 2x2 -> 3x3)

본 논문에서는 cyclic-shifting을 사용한 더 효율적인 batch computation 접근법을 제시한다

 

• 먼저 위 figure 4처럼 좌상단 방향으로 shift를 진행한다.
• 이후 하나의 batched window는 여러개의 인접하는 feature map이 아닌 sub-window를 가지게 된다
  • 따라서 masking 기법을 사용해 self-attention 계산을 각 sub-window로 제한시킨다
• 이 방법을 사용하면 각 batched window 크기는 기존 window paritioning과 같게 되고 더 효율적인 계산이 가능하게 한다.

 

Relative position bias

Self-attention을 계산할 때 position bias를 각 head에 더한다:

• $ \text{Attention}(Q, K, V) = \text{SoftMax} \left( \frac{QK^T}{\sqrt{d}} + B \right) V, $
  • $ B \in \mathbb{R}^{M^2 \times M^2} $
  • $ M^2$: 각 window 안에 patch 수

상대적인 위치는 [-M + 1, M - 1] 사이에 있기 때문에

• 더 작은 bias matrix $ \hat{B} \in \mathbb{R}^{(2M - 1) \times (2M - 1)} $를 사용해
• B의 값은 $\hat{B}$에서 가져오게 한다

실험을 통해 이런 bias term이 없는 경우와 절대적인 위치 embedding을 사용한 bias term에 비해 위 제시된 방법이 큰 성능 향상을 보이는 것을 확인했다.

또한 pre-training 과정과 다른 window size를 가진 모델을 fine-tuning할 때에 bi-cubic interpolation을 사용해 bias를 초기화 할 수 있다는 사실도 발견했다.

 

 

3.3 Architecture Variants

실험에서는 가장 기본 모델인 Swin-B를 ViT-B/DeiT-B와 같은 계산 복잡도를 가지도록 만들었다.

또 Swin-T, Swin-S, Swin-L은 각각 0.25x, 0.5x, 2x배 더 큰 계산 복잡도를 가지는 모델들이다.

기본적으로

• Window 크기 M = 7
• 각 head의 query dimension d = 32
• 각 MLP의 expansion layer $\alpha=4$로 설정한다

각 변형된 모델의 hyper-parameter 구조는

• Swin-T: C = 96, layer numbers = {2, 2, 6, 2}
• Swin-S: C = 96, layer numbers = {2, 2, 18, 2}
• Swin-B: C = 128, layer numbers = {2, 2, 18, 2}
• Swin-L: C = 192, layer numbers = {2, 2, 18, 2}
  • C: 첫번째 stage의 hidden layer의 channel 수

 

 

 

4. Experiments

ImageNet-1K 이미지 분류, COCO object detection, ADE20K semantic segmentation으로 실험을 진행한다.

먼저 Swin Transformer 구조와 기존 SOTA 기법들을 비교한다.

그 이후 Swin Transformation의 주요한 디자인 요소들을 제거하는 실험을 진행한다.

 

 

4.1 Image Classification on ImageNet-1K

Settings

이미지 분류 작업은 ImageNet-1K를 사용한다.

• 1.28M training, 50K validation, 1000 classes
• 성능은 top-1 accuracy single crop로 보고한다

다음 두 가지 학습 환경을 사용한다:

• Regular ImageNet-1K training
  • AdamW optimizer
    • 300 epoch cosine decay learning rate
    • 20 epoch linear warm-up
  • Batch size 1024
  • 초기 learning rate: 0.001
  • Weight decay 0.05
  • Augmetation
    • Training data-efficient image transformers & distillation through attention
      • EMA와 repeated augmentation을 제외한 나머지 증강과 정규화 기법
• Pre-training on ImageNet-22K and fine-tuning on ImageNet-1K
  • 먼저 더 큰 데이터셋인 ImageNet-22K로 pre-train을 진행한다
    • 14.2M image, 22K classes
  • Pre-training
    • AdamW optimizer
      • 90 epoch linear decay learning rate
      • 5 epoch linear warm-up
    • Batch size 4096
    • 초기 learning rate 0.001
    • Weight decay 0.01
  • Fine-tuning
    • 30 epoch
    • Batch size 1024
    • Constant learning rate 10e-5
    • Weight decay 10e-8

 

Results with regular ImageNet-1K training

위 Table 1(a)는 다른 Transformer-based와 ConvNet-based 모델들과 비교 결과를 보여준다.

• 기존 SOTA Transformer-based 구조인 DeiT보다 확연히 높은 성능을 보여준다
• 기존 SOTA ConvNet들과 비교해 약간 더 좋은 speed-accuracy trade-off를 보여준다.

 

Results with ImageNet-22K pre-training

위 Table(b)는 pre-train 이후 fine-tunning 한 결과를 보여준다.

• Swin-B 모델의 경우 pre-training이 1.8~1.9% 정확도 향상을 보여준다
• 다른 Transformer-based 구조에 비해 약간 더 높은 speed-accuracy trade-off를 보여준다

 

 

4.2 Object Detection on COCO

Settings

COCO 2017 데이터를 사용해 학습을 진행한다.

• 118K training, 5K validation, 20K test-dev
  • Ablation 연구는 validation을 이용해 진행하고 system-level 비교는 test-dev 데이터를 이용한다

Ablation 연구에는 4개의 기본적인 object detection framework를 사용한다

• Cascade Mask R-CNN
• ATSS
• RepPoints v2
• Sparse   RCNN in mmdetection

모든 실험은 같은 setting을 사용한다

• multi-scale training (짧은 변은 480에서 800사이, 긴 변은 1333이하)
• AdamW optimizer
  • 초기 0.0001, weight decay 0.05, batch size 16
• 3x schedule (36 epoch)

System-level 비교에는

• Improved HTC with instaboost
• Stronger multi-scale training
• 6x scheduler (72 epoch)
• soft-NMS
• 그리고 ImageNet-22K pre-trained 모델을 초기값으로 사용한다

---

 

 

4.3 Semantic Segmentation on ADE20K

---

 

 

4.4 Ablation Study

이 섹션에서는 Swin Transformer의 주요 디자인 요소들을 제거하는 실험을 진행한다.

 

Shifted windows

Table 4에 따르면 모든 vision task에서 shifted window를 사용했을 때 더 높은 정확도를 보여준다.

Table 5에 따르면 shifted window의 latency도 크게 차이나지 않는것을 보여준다.

 

 

Relatvie position bias

Table 4에 따르면 relative position bias가 모든 vision task에서 역시 더 높은 정확도를 보여주는 것을 확인할 수 있다.

Absolute positional embedding을 사용했을 때 classification 작업에서는 정확도가 소폭 상승하지만 나머지 task에서는 오히려 정확도가 낮아지는 것을 볼 수 있다.

 

 

Different self-attention methods

Table 5는 다른 self-attention 기법들의 실제 속도를 비교한다.

본 논문에서 제시된 cyclic implementation은 naive padding 기법에 비해 13%, 18%, 18% (Swin-T, Swin-S, Swin-B) 더 빠른 속도를 보여준다.

또한 sliding window 기반 Swin Transformer보다 shifted window 기반 Swin Transformer도 훨씬 더 빠른 속도를 보여준다.

또한 Table 6에 따르면 두 기법의 정확도는 거의 비슷한 것을 확인할 수 있다.

 

 

 

5. Conclusion

본 논문에서는 새로운 Vision Transformer인 Swin Transformer를 소개한다.

특징으로 계층적인 representation을 사용하고 이미지 크기에 대해 linear한 복잡도를 가진다.

 

Shifted-window 기반 self-attention은 visual task들에 효용성과 효율성을 보여주고 이를 natural language processing 분야에 적용하는 기법을 탐구하는 것도 기대한다.