ShuyaVision
[논문 리뷰] Swin Transformer : Hierarchical Vision Transformer using Shifted Windows 본문
[논문 리뷰] Swin Transformer : Hierarchical Vision Transformer using Shifted Windows
Sangbro 2023. 11. 18. 19:07Abstract & Introduction
기존의 Transformer를 Vision task에 적용하기에는 크게 2가지 문제점이 있다고 말한다. 첫번째로는 visual entities의 scale이 text에 비해 다양하다는 점이다. NLP에서는 word를 embedding한 token 형태로 사용하기 때문에 크기를 고정적으로 활용하고, 기존의 Transformer 들은 대부분 이를 기반으로 구성되어 있다. 하지만 visual entities는 환경, 이미지의 종류에 따라 scale이 다양하게 잡히는 경우가 많다. 특히 Object Detection과 같은 경우, DETR에서는 작은 물체를 잡지 못하는 문제점이 존재했듯, visual elements는 scale에 따라 크기가 따라 변하기 때문에 NLP와는 다른 환경이라 할 수 있다.
두번째로는 images의 pixel 수가 embedded text에 비해 high resolution이라는 점이다. Semantic segmentation과 같은 task는, pixel level에서늬 dense prediction을 필요로 하며 이는 Transformer의 self-attention 연산 방식에서는 image size에 따라 연산량이 제곱으로 커지게 한다. 이러한 domain간의 차이를 극복하여 vision task에도 transformer를 잘 활용하기 위해서, 본 논문에서는 Shifted windows를 활용하는 hierarchical Transformer인 Swin Transformer를 제시한다.
Swin Transformer는 hierarchival feature maps를 구성해 사용하며, image size에 따라 연산량이 선형적으로 상승하도록 설계되었다. 위 Figure 1에 나와있듯, small-sized patches들로부터 출발해 이를 점진적으로 이웃 patches들과 합쳐 활용한다. 이러한 hierarchical feature maps들을 통해 Swin Transformer는 FPN이나 U-Net등이 목표로 하는 dense prediction을 잘 수행할 수 있다. 또한 빨간색 윤곽선에 해당하는 patches 끼리 self-attention을 계산하도록 설계한다. 이를 통해 겹쳐서 self-attention을 수행하는 부분이 없어지기 때문에, linear complexity를 달성할 수 있다. 이러한 장점을 기반으로 하여, 기존의 ViT를 대체할 수 있는 장점이 된다.
Swin Transformer의 key design은 위 Figure 2에서 표시되는 것 처럼, 연속적인 self-attention 레이어 사이의 shift of the window 이다. 본 논문에서는 shifted windows가 이전 레이어의 window를 연결하여 modeling 성능을 크게 향상시키는 방안이라고 주장하며, Table 4에서 비교군을 제시한다. 위 전략은 latency 측면에서도 효율적이라고 한다. 각 window에 들어있는 모든 query, 즉 각 patches들은 같은 key set을 가지게 되기 때문에, HW의 memery access를 용이하게 까지 한다. 이와 대조적으로, sliding windows의 방법을 채택한 연구들은 서로 다른 query에 대해 다른 key set을 가지게 되어 HW에서 latency 측면에서 비효율적인 모습을 보인다.
*** (개인적인 생각으로 이 정도까지 주장하는 건 논리의 비약이라고 생각함). 이전 Layer의 output에 대해, transformer이기 때문에 squential한 데이터일테지만, positional encoding을 고려해 feature map이라고 가정하겠다. 이전 Layer에서 선택한 왼쪽 상단의 4x4 window는 다음 레이어에서 right-down 방향으로 (2, 2) 이동한다. 이를 통해 이전 Layer에서 고려하지 못했던, pixel 간의 연결성을 self-attention으로 고려해주겠다는 뜻이다. 이를 통해 모든 pixel 간의 데이터를 연결하는 방식이 된다. 그런데 이전 Layer에서 나온 feature map의 경우 그 정보를 이미 포함하고 있을 것이기 때문에, window간에 상호 소통하는 것은 아니다. 이번 layer의 shifted window가 이전 레이어의 window를 연결하는 것은 아니라는 뜻이다. 다만, window를 통해 모든 feature map을 고려하는 transformer에 비해, receptive field가 작아지는 부분을 shifted window로 "보완" 한다는 표현이 적절해 보인다. ***
Methods
Overall architecture
위 그림은 Swin Transformer architecture 중 tiny version을 그린 그림이다. 먼저 ViT와 같이, RGB image를 non-overlapping patches로 나눈다. 각 patch는 token으로 취급되며, 각 feature는 각 RGB values 값을 사용하여 구성된다. 논문에서는 patch size를 4x4로 사용하여, 각 패치의 feature dimension은 4x4x3=48이 된다. 이 patches에 Linear Embedding을 적용하여 차원 C로 project(투영)한다. 그 다음 Swin Transformer Block이 이에 적용된다. 각 blocks은 H/4 x W/4 x C의 tokens를 가지게 되며 linear embedding과 함께 Stage 1이라 명명한다.
Hierarchical 한 구조를 만들기 위해, layer가 깊어질 수록 patch들을 merge해서 tokens의 수는 줄어들게 된다. 첫번째 패치를 merge하는 layer에서는 2x2 이웃 패치의 features들을 merge하는데 이 때, 좌우로 합치는 것이 아니라 3차원 방향, 즉 Channel 방향으로 합친다. 결과물은 4C-dimensional이 되고, 이를 linear layer를 적용해 출력 차원은 2C가 되도록 설정한다. 이는 tokens의 수를 2 x downsampling of resolution (여기서는 2x2) 만큼 줄이고 이는 Stage 2가 되며 이 프로세스를 Stage 3, Stage 4까지 반복한다.
위 과정은 VGG나 ResNet과 같은 CNN과 거의 흡사하게 feature map resolutions을 줄인다. 이를 통해 vision tasks의 backbone을 대체하는 방안이 될 것이라고 주장한다.
Swin Transformer block
Swin Transformer는 multi-haed self attention(그림에서는 MSA)를 shitfed windows를 base로 한 Transformer block으로 대체하며, 다른 layer들은 그대로 둔다. 위 그림의 (b)에 나와있듯, Swin Transformer block은 shifted window를 베이스로한 MSA 모듈과 2-layer MLP를 두며, 사이의 비선형함수로는 GELU를 사용한다. 각 MSA 모듈과 MLP 이전에는 LayerNorm layer가 존재하며, residual connection이 각 모듈 이후에 적용되도록 구성되어 있다.
참고로, 그림의 W-MSA는 Window Multi-head Self-Attention이고 SW-MSA는 Shifed Window Multi-haed Self-Attention이다.
Self-attention in non-overlapped windows
효율적인 모델링을 위해, 논문에서는 local-window 내에서 self-attention을 연산하도록 제안한다. Windows들은 non-overlapping manner로 arranged 된다. 각 윈도우가 MxM의 patches들을 가지고 있다고 가정했을 때, h x w의 patches들에 대해 global MSA 모듈의 계산 복잡도와 W-MSA의 계산 복잡도는 아래와 같다.
global MSA는 hw에 대해 제곱으로 증가하는 반면, W-MSA는 linear하게 증가한다.
Shifted window partitioning in successive blocks
Window 기반 self-attention 모듈은 window 간 연결이 부족해 모델링 power가 부족할 수 있다. Non-overlapping window간의 효율적인 계산을 유지하면서, window 간의 연결성을 도입하기 위해, 연속되는 Swin Transformer blocks에서 두개의 partition 구성을 사용하는 shifted window partitioning을 제안한다.
Figure 2에 나와있듯, 첫번째 모듈은 top-left에 위치하며, 8x8 feature map이 4x4(M=4)인 2x2 구성으로 고르게 분할 된다. 그리고 다음 모듈에서는 이전 layer에서의 window 구성에서 ([M/2], [M/2]) 만큼 shift한 구성을 사용한다. 이는 아래와 같이 계산될 수 있다.
여기서 ^l 과
Efficient batch computation for shifted configuration
그런데, Shifted window partitioning은 더 많은 window를 만들어낸다는 단점이 존재하며, 일부 창은 MxM 보다 작아야하는 문제가 존재한다.
Naive하게는 작아진 window들에 padding을 설정해 MxM 크기로 맞추어 주고, attention을 계산 시, padding 된 부분에 대해서는 masking을 한다. 하지만 이는 결국 window 수를 늘리게 되고 이로 인한 computation은 꽤나 큰 편이다.
그래서 논문에서는 왼쪽-위 방향으로 cyclic-shifting 하여 보다 효율적인 계산 방식을 채택한다. Shift 이후에는, window는 인접하지 않았던 여러 부분들(several sub-windows)와 함께 구성되게 된다. (실제로 계산되어야하는 window는 위 그림에서 보면 작은 window로 분할, 물리적으로 떨어져있음) 이 때, self-attention 계산을 각 하위창 내로 제한하기 위해 masking 메커니즘이 사용된다. cyclic-shifting을 활용하면, 이런 방식을 통해 window의 수를 똑같이 맞추어 computation의 크기를 동일하게 유지한다.
Relative position bias
self-attention을 계산할 때, relative position bias(B)를 더해준다. 이로 인해, position embedding을 사용하지 않고도, position 정보를 유지할 수 있고, 오히려 더 좋은 성능을 낸다고 한다. 이 relative position bias는 bi-cubic interpolation을 통해 다른 크기의 window를 가진 모델을 초기화하는 데도 사용할 수 있다고 한다.
