[논문 리뷰] Swin Transformer 간단하게 리뷰해보기

 

구현은 조금 빡빡해서 리뷰를 먼저하고 구현은 천천히 업로드 할 예정이다.

 

Swin Transformer는 Shifted Window의 줄임말에 Transformer를 붙인 용어이다.

 

용어 그대로 Shifted window를 활용해 성능을 개선시킨 모델이라 논문을 읽지 않아도 추측해 볼 수 있다.

 

나오게 된 것은 다양한 scale을 갖는 visual entities를 커버하고, Large resolution의 이미지 인풋에 대해서 patch words를 잘 구성해야한다는 점을 motive로 하여 개발이 된 모델이다.

 

논문 제목 : Swin Transformer : Hierachical Vision Transformer using Shifted Windows

 

Abstract

 

NLP와 비전에서의 domain 차이로 인하여 NLP에 좀 더 specific한 transformer를 도입하는데 어려움이 있다. 

 

어려움은 Text 단어와 비교하면 image pixel들의 high resolution이나 visual entities의 large variation 등을 말한다.

 

이를 해결하기위해 나온 모델이 Swin Transformer이다. Hierarchical 구조와 Shifted Window를 통한 representation(attention) 계산을 통해 성능을 향상시켰다.

 

구체적으로는 self attention 계산은 ViT와 다르게 local window로 쪼개서 계산하여 (cross window connection 도 추가) 성능을 향상시켰으며, 다양한 scale에서도 좋은 성능을 유지하고, image size를 증가시킬 떄 기존 ViT가 quadratic하게 복잡도가 증가했으나 local 로 쪼갠 방법을 통해 linear한 증가로 복잡도의 증가폭이 줄어드는 효과를 보여줬다.

 

또한 분류, 객체탐지, semantic segmentation task 에서도 좋은 성능을 나타낸 모델이다.

 

 

 

 

Intro

 

비전분야는 CNN이 대부분의 지배하고 있는 상황이다. 그러나 NLP는 다양한 architecture를 활용하면서 성능향상을 고민했고, 현재는 Transformer가 대세가 된 상황이다. 그러면서 CV에서도 Transformer에 관심을 주기 시작했고, 연구도 활발하게 이뤄지고 있다.

 

하지만 두 domain의 차이로 인하여 transferring이 어려운데 크게 2가지로 볼 수 있다.

 

- Scale 차이

Transformer는 고정된 scale의 token이기에 vision에서는 한계가 있다. (vision에서는 entities마다 scale이 다름)

 

- Text word보다 image는 higher resolution of pixels

또한 픽셀단위에서도 높은 prediction을 요구하는 task가 많은데, size가 커질수록 quadratic하게 복잡도가 증가하는 ViT의 경우처럼 Transformer를 도입할 때 문제점들이 있다.

 

그래서 이를 해결하고자 Swin Transformer를 개발한 것이다. 

 

주요 특징은 2가지이다.

 

- Hierarchical feature maps를 만드는 것으로 작은 patch부터 시작해서 점점 patch size를 키워가는 방식으로 attention을 계산하게 된다.

 

- 이미지 size를 키울 때 linear computational complexity 증가를 보여주는데 이는 local self attention 계산을 활용한 것이다.

각 window안 patch개수는 고정이므로 복잡도가 linear하게 증가하는 것이다. 이에 따라 기존 ViT 보다 task 적용도 용이하다.

 

이외에도 논문에서는 Hierarchical 구조를 통해 FPN, U-Net등에서 적용가능하다고 언급하고 있다.

 

Self-attention Layer 사이에 window partition을 shift하는 block을 넣어 성능을 향상시켰는데 이는 real-world latency에서 기존 transformer 계열 모델들보다 더 낮은 값을 보여준다. 같은 window안의 patch들은 같은 key set을 공유해 메모리 적으로도 좋다. 그래서 latency는 낮추면서도 성능은 비슷한 결과를 확인했다.

 

 

 

 

Architecture

 

이 논문에서는 Tiny 버전을 가정하고 설명하고 있다.

 

전체 Overview Architecture는 다음과 같다.

 

 

Input

 

Input은 RGB image로 해서 ViT 처럼 patch split을 진행한 후 각 patch는 token처럼 사용되고, RGB channel dimension은 feature에 해당된다. patch size는 4x4로 고정시키기 때문에 각 token마다 16개의 patch가 있으므로 4x4x3 = 48개의 feature dim을 갖게 된다.

 

Stage 1

 

48의 feature를 임의의 값 C로 Linear projection을 한 이후 뒤에서 언급할 수정된 Self-attention computation block을 통과시킨다. 이때 토큰 수는 $\frac{H}{4} x \frac{W}{4} $ 를 유지한다.

 

Stage 2

 

patch merging layers를 활용해 Hierarchical representation 구조를 만들기위해 token 수가 감소시킨다. 

 

다음 그림의 왼쪽(a) 처럼 이웃한 2x2크기의 patch들을 겹치지 않게 merge한다.

 

 

merge는 concat 방식을 사용하는데 그렇게 되면 4C의 feature를 갖게 된다. 그 후 2C로 linear projection하고 난 후 Swin Transformation Block을 통과시키는데 이때 token 수는 concat하면서 줄어들었기 때문에 $ \frac{H}{8} x \frac{W}{8} $ 만큼이 남게 된다.

 

 

Stage3, 4 는 2의 과정을 반복한다.

 

그렇게되면 각각 $\frac{H}{16} x \frac{W}{16}$, $\frac{H}{32}x\frac{W}{32}$ 의 token 수가 된다.

 

이러한 구조는 VGG, ResNet 처럼 feature map의 Hierarchical representation 을 갖기때문에 FPN, U-Net 구조 등에 활용될 수 있다.

 

Swin Transformer Block은 Multi-head Self-attention module을 대체하는 block으로 나머지는 ViT의 Block과 동일하게 진행된다.

 

 

 

 

Shifted Window Based Self-attention

 

ViT는 Global Self-attention 계산으로 inductive bias를 줄이고 Large-scale의 dataset에서 뛰어난 성능을 보여줬지만 Global 계산은 size 증가에 따라 복잡도가 token수의 관점에서 quadratic하게 증가한다는 문제가 있다. 

 

그래서 Swin Transformer에서는 Global을 Local window로 계산하고, Window들을 partition image에 각각 할당한다.

 

논문에 따르면 Global 과 Local Window 활용했을 때 complexity를 수식으로 계산하면 다음과 같이 나온다고 언급한다.

 

$$ \Omega (MSA) = 4hwC^2 + 2(hw)^2C $$

$$ \Omega (W-MSA) = 4hwC^2 + 2M^2hwC $$

여기서 M은 patch size로 4로 고정시키기 때문에 Linear하게 복잡도 증가폭을 줄인다는 것이다.

 

그러나 Window-based self-attention 계산은 window 간 connection이 없다. 따라서 cross window connection을 하면서도 non-overlapping을 유지할 수 있는(기존 Window-based의 특징을 살리면서) Shifted window partitioning을 적용했다.

이것을 Window based self-attention 계산 다음 바로 계산하도록 붙여서 block단위로 설계를 했고, 성능향상을 확인할 수 있었다.

 

 

Efficient Batch Computation for Shifted Configuration

 

이 부분이 살짝 어려웠는데, Shifted Window는 patch size보다 작은 경우가 있기에, 더 좋은 partition은 어떻게 할지에 대해서 언급하고 있다.

 

먼저 simple한 방법으로 Shift로 인하여 생긴 작은 Window는 padding된 값으로 masking을 하여 attention을 계산하는 것이다. 하지만 2x2처럼 작은 Window partition에서는 computation이 기존 window based 보다 증가한다는 점이 있다. (다음 그림처럼 sub window가 늘어났기 때문이다.)

 

 

그래서 그림과 같이 표시된 부분을 cyclic shift를 통해 옮긴 후 각각의 Window에 해당되지 않는 부분은 masking 하여 attention을 계산하고 나온 값을 다시 reverse cyclic shift로 원래대로 옮겨놓는 과정을 통해 복잡도를 늘리지 않게 한다.

 

 

 

Relative Position bias

 

ViT와 다르게 Attention 계산시 다음 수식의 B값 처럼 bias term을 추가하는데 이때 B값은 relative position 에 해당된다. 

상대적 위치는 각 축마다 [-M+1, M-1] 의 범위안에 값이므로 $B \in R^{M^2xM^2} $ 의 값은 $ B` \in R^{2M-1 x 2M-1} $ 에서 구할 수 있으며 이를 학습 parameter로 넣어준 후 MLP를 통과시키게 된다. Attention Term은 다음과 같다.

 

$$ Attention(Q,K,V) = Softmax\left(\frac{QK^T}{\sqrt{d}} + B\right)V $$

 

다음 글에서는 구현 코드를 통해 다음 내용들을 직접 짜보고 이해한 내용을 작성하도록 하겠다.

 

 

참고 논문 - https://arxiv.org/abs/2103.14030

Swin Transformer: Hierarchical Vision Transformer using Shifted Windows