[논문 리뷰] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows(2021)

이번에 리뷰할 논문은 Swin Transformer: Hierarchical Vision Transformer using Shifted Windows 입니다.

 

https://paperswithcode.com/paper/swin-transformer-hierarchical-vision

Papers with Code - Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

 

 

 

 

 

 

 

2017년도에 transformer 모델이 등장하면서 NLP 분야에서 큰 혁신이 일어났다고 해도 과언이 아닐만큼 엄청난 performance를 보여주었습니다. 이런 transformer를 NLP뿐만 아닌 CV에도 활용하려 했으며, 2021년도에 tranformer를 vision에 적합시킨 ViT 모델이 나오게 됩니다.

• [논문 리뷰] AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE(2021)

 

하지만 NLP를 위해 탄생한 transformer를 CV로 변환하는데 2가지 challenge가 존재했습니다.

 

1. Difference involves scale

NLP의 경우 데이터를 각 단어 token으로 처리하는 반면 시각적 요소의 경우 크기가 다양합니다(multi-scale). 즉, token의 규모가 모두 고정된 transformer가 vision과 응용하는데 적합하지 않은 문제가 발생합니다.

 

2. much higher resolution of pixels in images compared to words in passages of text

transformer의 핵심인 self-attention의 경우 모든 token에 대해 서로 상호작용합니다. ViT의 경우 이미지를 각 patch로 나누어 하나의 patch를 token으로 취급하기 때문에 image size가 증가함에 따라 계산 복잡도가 quadratic하게 증가한다는 문제가 발생합니다.

 

 

 

위 두가지 문제를 해결하기 위해 본 논문에서는 representation이 shifted window로 계산되는 계층적 transformer를 제안합니다. 점차 계층적으로 view를 바꿔주기에 multi-scale을 반영할 수 있으며, 각 window끼리 독립적으로 attention을 수행하기에 계산복잡도를 선형적으로 낮춘다는 것이죠.

 

 

위 figure의 빨간색 경계는 window간의 경계를 의미하며, 회색 경계는 patch간의 경계를 의미합니다. 기존 ViT의 경우 patch의 size가 일정하게 유지되어 전역적인 정보만을 학습할 수 있던 반면, Swin transformer는 다양한 scale의 정보를 학습할 수 있어 segmentation이나 object detection 분야에서 활용될 수 있습니다. 

 

 

 

 

 

 

 

 

Overall Architecture

전체적인 Swin transformer의 구조는 다음과 같습니다.

 

 

우선 ViT와 같이 image를 겹치지 않게 patch로 split 합니다 . 이러한 patch들은 하나의 token으로서 역할을 수행하죠. 위 그림과 같이 input image의 shape이 HxWx3 이고, patch size가 4x4 라면, 해상도는 4배만큼 감소할 것이며, channel 수는 4x4x3 인 48로 증가하게 됩니다.

 

 

다음으로 linear embedding을 적용하여 임의의 C 차원으로 projection 합니다. 즉, feature map의 shape은 $ \frac {H} {4} \times \frac {W} {4} \times C $ 가 되겠죠. 

 

 

 

그 후에 본 논문에서 제안하는 Swin Transformer Block을 수행합니다. 

 

 

위와 같이 총 두번의 서로 다른 W-MSA와 SW-MSA로 구성된 self-attention을 수행합니다. 이 뒤에는 사이에 GELU 비선형성을 준 2-layer MLP를 적용합니다. 이에 대한 설명은 뒤에서 자세히 하도록 하겠습니다. 물론 해당 block은 self-attention만으로 이루어져있기에 shape의 변화는 없습니다.

 

 

 

다음으로 계층적 구조를 위해 깊어질수록 토큰 수를 감소시킵니다. 이를 위해 이웃하는 2x2 patch를 하나의 patch로 merge하는 Patch Merging layer를 통과합니다. 예를 들어, 8x8의 patch들로 이루어진 feature map이 존재한다고 했을 때 서로 이웃하는 2x2 patch를 merge했으므로 해상도는 2배만큼 감소하고, 그만큼 channel 수는 4배만큼 증가하게 됩니다. 원래 channel수가 C였으니 4C가 되는 셈이죠. 하지만linear layer를 통해 2C로 reduction 시킨다고 합니다.

 

 

뒤의 Stage 3와 Stage 4는 Stage 2를 반복하여 진행됩니다. 이러한 patch merging 방식은 층이 깊어질수록 연산량을 감소시킬 뿐만 아니라 CNN의 pooling layer와 같이 해상도를 줄여 좀 더 global한 information을 추출할 수 있도록 하는 역할을 수행합니다. 

 

 

 

 

 

 

 

 

Self-attention in non-overlapped windows

기존의 transformer의 경우 token과 다른 모든 token과의 전역적인 계산을 수행하기 때문에 token 수에 대한 quadratic 계산 복잡도를 초래한다고 위에서 설명했었습니다. 이는 vision 분야에 적합하지 않은 이유 중 하나이기도 하죠. 해당 문제를 해결하기 위해 window가 서로 겹치기 않도록 균등하게 patch를 분배한 후에 window 내에서만 self-attention을 수행할 수 있도록 하였습니다. 

window 내에서만 self-attention을 수행하는 것은 계산 복잡도를 image size에 대해 선형적으로 증가할 수 있도록 변환해줍니다. image size가 증가한다고 해서 self-attention에 수행되는 patch의 수가 늘어나지 않으며, 단순히 window의 개수만 증가하기 때문이죠. (window 내의 patch 수는 항상 일정) 

 

이러한 self-attention을 본 논문에서는 W-MSA 라고 notation 합니다.

 

 

 

 

 

 

 

 

Shifted window partitioning in successive blocks

계산의 효율을 위해 window 내에서만 self-attention을 진행하게 되면 window간의 상호작용에 대해서는 학습을 할 수 없겠죠. 이를 해결하고자 window를 다음과 같이 새로 구성하여 다시 적용합니다.

 

 

위와 같이 가로 세로에 window를 하나씩 더 추가함으로써 shifted window partitioning을 수행합니다. 하지만 이 경우 window의 수가 증가하여 그만큼 계산량이 증가한다는 단점이 존재합니다. 이를 해결하고자 본 논문에서는 cyclic shift 기법을 소개합니다.

 

 

원래 SW-MSA를 위해 나누었던 window 수는 총 9개 였습니다. 하지만 위 figure와 같이 A, B, C 부분을 아래로 이동시켰으며 이를 cyclic shift라고 합니다. 이렇게 되면 다시 window 수는 총 4개가 되겠죠. 

 

 

하지만 위 경우를 생각해보면 회색 부분과 C 부분은 따로 self-attention이 수행되어야합니다. 이를 위해 masked MSA를 더해주어 수행하게 됩니다.

 

 

위의 예시를 보면 $ M^{2} $을 20이라 가정했을 때 1 ~ 10 과 11 ~ 20 은 각각 self-attention이 적용되어야 하므로 해당 부분에 대해서만 attention이 수행될 수 있도록 masked MSA를 적용하는 모습을 확인할 수 있습니다.

 

 

 

 

 

 

 

 

Efficient batch computation for shifted configuration

기본적으로 Swin Transformer는 window 끼리 독립적으로 attention을 수행합니다. 그렇기에 만약 window의 수가 4개가 있다면 같은 attention 연산을 4번 반복해서 하는 것이죠. 하지만 이는 계산 복잡도 적으로 보았을 때 비효율적입니다. 본 논문에서는 이 window들을 하나의 batch로 묶어 한번에 병렬적으로 연산을 진행합니다. 

 

 

 

 

 

 

 

 

Relative position bias

기존의 transformer 기반 ViT 모델의 경우 position embedding을 더해주어 위치 정보를 반영하였습니다. 하지만 swin transformer의 경우 position embedding 대신 relative position bias를 더해주어 이를 대체합니다.

 

 

기존의 position embedding이 위치 정보에 대해 절대좌표를 더해주었다고 한다면, 본 논문에서 제안하는 relative position bias의 경우 상대적인 좌표를 제안한다고 볼 수 있습니다. 상대적이기에 relative position bias의 범위는 window size를 M이라고 했을 때 [ -M+1, M-1] 으로 설정됩니다. 만약 window size(M)이 3이라면, Relative position bias의 범위는 [-2, 2]가 되는 것이죠. 

 

 

 

 

 

 

 

 

Architecture Variants

본 논문에서는 window size(M)을 7로 설정하였으며, query dimension(d)를 32로 설정했습니다. 또한 모델의 크기는 stage 3의 반복 횟수를 조절하여 다음과 같이 구성했습니다.