ViT의 등장 이후 이를 베이스로 다양한 형태의 variants가 쏟아져 나오기 시작했다. 이번에는 그중 가장 대표적으로 알려진 backbone 중 하나인 Swin Transformer에 대해 리뷰해 보겠다.
Background
이전까지는 Abstract, Introduction 등 논문의 진행을 그대로 따라갔는데 ViT base backbone 연구모델은 대게 흐름이 비슷하다고 느껴 이번 글에서는 주요한 부분만 빠르게 짚어보겠다.
ViT는 JFT-300M을 pretrain해 IM1K에 대해 SOTA를 달성했다. 하지만 이는 pretrain에 한정된 이야기이고 ViT의 고질적 단점 이자 장점 low bias to image = concerate to general information으로 인해 IM1K pretrain->IM1K predict에 대한 성능은 기존 CNN보다 낮았다.
ViT는 CNN보다 gpu memory 사용량이 많고 고정된 image size [224x224]만을 사용한다. 이는 적은 데이터 혹은 컴퓨팅 용량을 사용해야 하는 시나리오에는 적합하지 못했으며 이후에 ViT가 CNN처럼 image에 대해 높은 편향을 가지도록 만드는 연구가 여럿 진행되었다.
Swin Transformer는 이문제에 대하여 Architecture적으로 접근하였으며 이에 사용한 대표적인 방법 shifted window는 이미지에 대해 계층적인 정보를 파악하며 효율성과 성능 모두 끌어올릴 수 있었으며 ViT base model 역시 backbone으로 활용할 수 있음을 시사했다.
이러한(Make ViT general) 연구를 진행한 사람들은 대부분 NLP와 CV의 도메인적 차이를 언급하고 이를 바탕으로 연구의 배경을 설명한다.
이 논문에서 저자는 NLP와 CV의 차이를 scale에 두었다. 즉 고정된 사이즈의 자연어(word embedding)에 적용되는 Transformer는 이미지(can have many resolution)에 곧바로 적용되기에 적합하지 않다고 언급했다.
이건 단순히 입력이 아니다 예측에도 똑같이 적용된다. 한 단어를 맞추는 것과 한 물체를 맞추는 것을 생각해 보면 쉽게 연상할 수 있다. 논문에서는 이러한 task를 dense prediction으로 명명하고 CV에 backbone으로 사용되기 위해 필요한 특성이라 언급하였다.
Method
Overall Architecture
Swin의 전체적인 아키텍처는 아래와 같다.
논문에서 왼쪽의 Image input부터 오른쪽으로 하나씩 설명하므로 이 글에서도 똑같은 순서로 리뷰해 보겠다.
Patch embedding & Linear embedding(Stage 1)
이 부분은 ViT와 큰 차이가 없다. 4*4 patch size를 사용하여 이미지를 쪼개는 데 주의할 점은 16*16을 사용하는 ViT보다 매우 작은 사이즈라는 것이다. 이후에는 embedding layer를 거쳐 C차원으로 만든다.
Patch merging
Stage1 이후부터 각 Stage를 넘어갈때마다 patch merging을 거치는데 그때마다 patch의 개수와 채널이 변하는 걸 확인할 수 있다.
계층적인 정보를 넣어주기 위해 주위의 2*2 patch를 4개로 겹쳐 쌓아 4C로 만든 다음 다시 2C로 보내주는데 이 과정을 patch merging이라고 정의하였다.
Swin Transformer blocks
그림(b)와 같이 Swin Transformer block은 기존 multi head attention(MHA)를 대체하는 블록이다. 구조는 MHA와 동일하고 MSA가 W-MSA, SW-MSA가 사용되는 것으로 그리고 activation으로 GELU를 사용한다
Shifted Window based Self Attention
Swin Transformer block에 기재된 W-MSA, SW-MSA를 의미하고 기존에 global attention의 큰 연산량을 줄이기 위해 고안한 방법이다.
기존 global self attention이 모든 patch 간의 attention을 진행했던 것과 달리 window size [M x M]을 정하고 이 사이즈만큼의 패치들 끼리만 attention을 진행한다.
공식은 위와 같은데 오른쪽의 h, w(전체를 보는)에서 M, M(일정부분만 보는)으로 변한다. CNN에서 수용장을 h,w로 잡은 것 에서 3,3로 잡은 느낌으로 이해해도 될 것 같다.
이렇게 시야를 제한하면 트랜스포머는 pooling같은게 없으니 CNN보다 전역적인 정보를 학습하기는 어렵다. 저자 역시 이를 이해했으며 이를 방지하기 위해 Shifted window partitioning을 제안한다.
방법은 간단하다 위처럼 기존 윈도우의 절반만큼 이동해 window간의 정보를 공유해 준다.그림(b)에 기재된 것처럼 W-MSA와 SW-MSA가 번갈아 적용된다.
Efficient batch computation for shifted configuration
SW-MSA로 patch를 여러 시점에서 보는건 좋지만 기존 W-MSA보다 많은 패치가 attention에 참여하게 된다. 그림에서는 [2x2]에서 [3x3]으로 증가했다.
위의 그림처럼 left-top의 자투리 부분을 잘라서 right bottom에 붙인다. 그림에 표기된 A, B, C는 원래 있던 이미지의 회색 부분과 연관이 있는 부분은 아니므로 마스킹하고 진행한다.(left top이 연산에서 무시되는 게 아닌지 검색했는데 이러한 부분은 크게 신경 쓰지 않는 것 같다)
MHA는 연산시 (batch_size, n_head, seq_len, D)로 기억한다. 아마 [3x3]으로 쪼개면 결국 적은 사이즈의 patch에도 mask를 더해 정사각형으로 만들어야만 병렬적인 연산이 가능한데 이 기법으로 이러한 문제를 방지한 것으로 보인다.
Experiments
Swin은 기존 ViT와 달리 IM1K pretrain만으로도 CNN 아키텍처 모델보다 좋은 성능을 달성했으며 ImageNet22K에 대해서도 ViT보다 좋은 성능과 자원 효율을 보여준다.
저자는 CNN 모델이 neural architecture search(NAS)와 같은 방법으로 IM1K에 최적화되어 있음을 지적했고 Swin 역시 추가적인 성능향상이 가능함을 암시했다.
분류보다 눈에 띄었던 부분은 detection과 segmentation인데 Swin은 백본으로 사용되면 ViT base 아키텍처보다 월등히 좋은 성능을 보여준다.
ViT가 백본으로 쓰일 때 어떤 식으로 적용되는지 상세히는 모르지만 저자의 주장으로는 기존 ViT 아키텍처는 하나의 채널을 그대로 백본으로 사용하지만 Swin은 여러 채널 출력을 사용가능한 점이 유리하게 작용했을 것이라 추측했다.
ViT를 읽고 이를 base로 하는 여러 논문을 읽어봤는데 전반적으로 ViT의 image에 대한 낮은 편향을 극복하려는 시도가 많았다. Swin은 AlexNet->VGGNet으로의 발전처럼 feature map을 잘게 쪼개는 접근법을 취해 이러한 편향을 더했다고 생각된다.
주관적으로도 이전 경진대회에서 detection, segmentation에서 swin의 성능, 효율성이 다른 ViT base 모델보다 좋은 기억이 있어서 아키텍처 선정을 할 때 빠지지 않고 사용하곤 한다.
