Swin Transformer 논문 리뷰 - Hierarchical Vision Transformer using Shifted Windows

Ze Liu† / Yutong Lin† / Yue Cao / Han Hu / Yixuan Wei† / Zheng Zhang / Stephen Lin / Baining Guo 

ICCV 2021 에 발표된 Microsoft Research Asia의 논문 Swin Transformer 리뷰입니다. 😊 

Motivation & Contributions

Language 분야와 달리 비전 분야에 고화질 픽셀의 규모에 대해 트랜스포머를 적용하기 어려움을 지적하고 Hierarchical transformer 구조를 활용하고 window 내 어텐션을 통해 효율적으로 성능을 개선합니다.

이 논문은 Hierarchical Transformer를 통해 다양한 스케일에 대해 flexible한 모델을 만들 수 있고, 이미지 크기가 늘어남에 따라 계산량이 선형적으로 늘어난다 (비교적 적게)는 장점을 가지고 있다고 주장하고 있습니다. 또한 Shifted window 를 따라 representation을 배우면서 겹치지 않는 local windows 의 셀프 어텐션에 효율성을 얻고, cross-window connection을 허용해 겹치지 않는 local window 사이 어텐션도 계산하고 있습니다.

Comparison with ViT Featuremaps

Method

먼저 작은 사이즈의 패치(회색 라인들)로 시작해 트랜스포머 layer를 지나며 패치를 합쳐 hierarchical representation을 학습합니다. 이를 통해 feature pyramid networks (FPN) 혹은 U-Net 구조처럼 활용될 수 있다고 하고, 이 non-overlapping windows 에서 셀프 어텐션을 local하게 계산함을 통해 이미지 크기에 따라 선형적인 계산량을 얻었다고 합니다. 이는 각 Window 내에 패치 수가 고정되어 있기 때문입니다. 따라서 기존의 트랜스포머 방식들이 이미지 크기의 제곱에 해당하는 계산량을 효과적으로 줄였다고 주장하고 있습니다.

An illustration of the shifted window approach for computing self-attention in the proposed Swin Transformer architecture.

또한 Swin Transformer의 핵심 요소는 연속 self-attention layers 간의 window를 shift하는 것입니다. Shifted window는 이전 layer의 window를 연결하므로 성능이 크게 향상됩니다.

이 전략은 실제 latency과 관련해서도 효율적인데, window 내의 모든 쿼리 패치는 하드웨어에서 메모리 액세스를 용이하게 하는 동일한 키를 공유합니다. 대조적으로, 이전의 슬라이딩 윈도우 self-attention 접근 방식은 다른 쿼리 픽셀에 대한 다른 키 세트 때문에 일반 하드웨어에서 낮은 latency으로 어려움을 겪습니다.

 

이 논문의 아키텍처는 아래 그림과 같은데 먼저 stage 1 에서 4by4 패치를 통해 48개 채널을 가진 4/h by 4/w 개의 패치를 만들고 linear embedding을 해줍니다. 이후 hierarchical representation을 만들기 위해 네트워크가 깊어지며 이웃한 2 × 2  패치를 merge하는 레이어를 거쳐 채널을 늘리고 다운샘플링을 통해 2C개의 채널로 만들어 주는 것을 stage 2 라고 하며 이후 반복해 stage 3 4를 통과해 더 작은 resolution의 representation을 만듭니다. 이는 다른 cnn 네트워크의 representation resolution 과 같으므로 백본으로 활용 가능하다고 합니다.

Self-attention in non-overlapped windows

겹치지 않는 윈도우의 self-attention은 아래와 같이 계산됩니다. M × M 패치에 대해, 

$$
\begin{aligned}
&\Omega(\mathrm{MSA})=4 h w C^{2}+2(h w)^{2} C \\
&\Omega(\mathrm{W}-\mathrm{MSA})=4 h w C^{2}+2 M^{2} h w C
\end{aligned}
$$

기존의 multi head self-attention (MSA) 는 hw 에 제곱에 비례해 계산량이 증가했지만 window-MSA는 선형적인 것을 확인할 수 있습니다.

Shifted window partitioning in successive blocks

윈도우 기반 셀프 어텐션 모듈은 윈도우 간 연결이 부족하여 모델링 성능이 제한됩니다. non-overlapping windows의 효율적인 계산을 유지하면서 cross-window connections 을 도입하기 위해 연속 Swin Transformer 블록에서 두 개의 분할 구성을 번갈아 사용하는 shifted window partitioning 접근 방식을 제안합니다.

첫 번째 모듈은 왼쪽 상단 픽셀에서 시작하는 일반 창 분할 전략을 사용하며 8 × 8 피처맵은 크기가 4 × 4(M = 4)인 2 × 2 윈도우로 균등하게 분할됩니다. 그 후, 다음 모듈은 정규 분할된 창에서 창을  (⌊ M/2 ⌋, ⌊ M/2 ⌋) 픽셀만큼 이동시킴으로써 이전 레이어의 윈도우 구성으로부터 이동되는 윈도우 구성을 채택합니다.

$$
\begin{aligned}
&\hat{\mathbf{z}}^{l}=\mathrm{W}-\mathrm{MSA}\left(\mathrm{LN}\left(\mathbf{z}^{l-1}\right)\right)+\mathbf{z}^{l-1}, \\
&\mathbf{z}^{l}=\operatorname{MLP}\left(\mathbf{L N}\left(\hat{\mathbf{z}}^{l}\right)\right)+\hat{\mathbf{z}}^{l} \\
&\hat{\mathbf{z}}^{l+1}=\operatorname{SW}-\operatorname{MSA}\left(\operatorname{LN}\left(\mathbf{z}^{l}\right)\right)+\mathbf{z}^{l} \\
&\mathbf{z}^{l+1}=\operatorname{MLP}\left(\operatorname{LN}\left(\hat{\mathbf{z}}^{l+1}\right)\right)+\hat{\mathbf{z}}^{l+1}
\end{aligned}
$$

각 모듈에서 self attention 전에 LayerNorm을 적용하고 이때 multi head self attention은 window 내에서 계산되며, shifted window를 통해 두개 연속 트랜스포머에 의해 계산되는데,

 

An illustration of an Efficient Batch Computation for Self-attention in Shifted Window Partitioning

shifted window partitioning 의 문제는 ⌈M/h ⌉×⌈M/w ⌉부터 (⌈M/h ⌉+1)× (⌈ M/w ⌉+1) 까지 더 많은 창이 생기는데 이때 windows will be smaller than M × M 보다 작은 윈도우가 생기기도 합니다. 이때 Shifted window를 더 효율적으로 계산하기 위해 naïve 하게 패딩을 한것이 아니라, cyclic shift와 그 reverse 연산을 통해 배치 연산을 수행합니다. 일반 파티션된 윈도우 수가 작을때 패딩을 사용하면 계산량 증가가 (2.25배) 커서 위 그림과 같이 왼쪽 상단 방향으로 순환 시프트를 통해 보다 효율적인 배치 계산 접근법을 제안하고 있습니다. 이 shift 후, 배치 윈도우은 특징 맵에 인접하지 않은 여러 하위 윈도우으로 구성될 수 있으므로, self-attn 계산을 각 하위 윈도우 내에서 제한하기 위해 마스킹 메커니즘이 사용되고 있습니다. 순환 시프트를 사용하면 일괄 처리된 윈도우의 수가 일반 윈도우 파티션 수와 동일하게 유지되므로 효율적입니다.

$$
\operatorname{Attention}(Q, K, V)=\operatorname{SoftMax}\left(Q K^{T} / \sqrt{d}+B\right) V
$$

마지막으로 self attention 계산 시 relative position embedding 에 bias B를 추가해 성능을 더 높였다고 합니다.

Experiements

• 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}

이 논문에서는 윈도우 사이즈를 기본 7로 사용하고 있고, 위 리스트 처럼 4가지 variants를 가지고 실험한 결과를 보여주고 있습니다. ImageNet-22K로 학습한 뒤 여러 태스크 (ImageNet-1K image classification, COCO object detection, ADE20K semantic segmentation)에 대한 성능을 보고하고 있는데, 성능은 아래 참조 링크 확인하시면 좋을 것 같습니다.

 

---

논문 - Hierarchical Vision Transformer using Shifted Windows