Swin Transformer (2021) : 새로운 CV 모델의 시작

안녕하세요!

 

이번에 리뷰할 논문은 Computer vision 모델에 어쩌면 새로운 패러다임을 제시할 수도 있는 

Transformer 기반 backbone 모델입니다. 

 

"Swin Transformer: Hierarchical Vision Transformer using Shifted Windows"

 

https://arxiv.org/abs/2103.14030v2

 

저의 관심분야와도 상당히 비슷하기 때문에 더 기대되는 논문 리뷰입니다.

흔히 vision transformer(Vit)라고도 불리는 데요. 자연어 처리와에 강건한 trasnformer들을 이용하여

다른 도메인인 vision간의 관계를 어떻게 해석했을 지 한번 알아보도록 하죠!

 

1. 서론

 Computer vision은 사실 CNN 기반의 모델들이 엄청난 성능을 발휘하는 영역(domain)입니다. CNN은 이미지가 주변 픽셀 영향을 받는다는 locality, 얻어낸 feature representation이 위치 정보에 반응하는 translation equivariance/invariance, 각 layer별로 얻을 수 있는 정보가 달라지는 high-low feature relationship 등 이미지 처리에서 훌륭한 성능을 보입니다.

여타 다른 trasnformer 기반의 vision 모델들이 그렇듯, 이 논문에서도 AlexNet(2012)에 대한 소개와 함께 CNN의 성능을 인정하며 시작합니다. 2012년 이후 현재의 CNN은 다양한 연구를 통해 좀 더 복잡하고 강인한 구조를 가지고 더 큰 scale을 처리할 수 있는 상태가 되었습니다. 

 

사실 이렇게나 훌륭하게 발전한 Computer vision과는 달리 NLP는 침체기가 꽤나 길었습니다. NLP에서 혁신을 가져온 모델은 바로 "transformer"(2017)인데요, 이 모델은 attention 만을 이용하여 sequence 데이터를 처리하고 시퀀스가 이어짐에 따라 정보가 손실되는 것을 줄이는 long-range dependency 문제를 해결하게 됩니다. 이 때부터 NLP에도 혁신이 찾아와 많은 발전이 이루어졌고, 자연스레 Computer vision에도 이와 같은 구조를 사용해보려는 시도가 늘게 됩니다. 

 

NLP에서 강인한 transformer를 computer vision에서도 generalized한 backbone로 사용한다는 것은 사실 쉽지 않을 거 같습니다. 두 개의 modality 사이의 차이점들도 존재하기 때문인데, 가장 먼저 언급되는 것은 "scale"의 문제입니다. word token은 scale이라는 개념 자체가 없지만 image의 구성 요소들은 이러한 scale에 민감하게 반응합니다. (그렇다면 scale을 해결하는 것이 swin transformer의 목표였을 수도 있겠네요) 그 다음 문제는 pixel level에서의 dense prediction 문제입니다. attention은 token(혹은 image 크기)에 제곱수에 해당하는 computational complexity를 갖기 때문에, 고해상도의 미지를 처리하는 것은 많은 computational cost를 요하게 됩니다.

 

Swin Transformer는 앞에서 언급한 두 문제점을 해결하기 위해 hierarchical representation을 제시합니다. small-size patch에서 시작하여 점점 그 patch를 넓혀가기 위해 sliding window(SWin)을 사용하는 모델입니다. 계층 구조를 갖는 feature map을 추출할 수 있을 것이며 이를 FPN이나 U-Net과 같은 기술들과 결합하여 feature들을 활용할 수 있습니다. 각 window 마다 해당하는 patch의 개수가 정해져 있기 때문에 complexity 역시 선형적으로 바뀌게 됩니다. 즉 앞서 언급한 문제점 두개를 모두 해결하기 위한 backbone 모델인 셈이죠.

 

위와 같이 window size를 바꿔가며 feature map을 만들어내기 때문에 dense prediction에 강할 것으로 예측되며 complexity도 줄일 수 있다

swin transformer의 핵심은 "shift window"입니다. window를 self-attention 각 layer마다 이동(shift)시키는 것이 그 방법론인데요, 이전 layer와의 다리(bridge) 역할을 한다고 보시면 될 거 같습니다. 즉, dense prediction을 위한 patch 별 정보 전달을 효율적으로 하는 수단인 셈입니다. 또한 real-world latency를 만족하기 위해 query와 key가 같은 set에 속해있는 apatch들로 구성하였습니다. (초기에는 key와 query set이 달라서 low latency 문제를 겪었다고 합니다)

sliding과는 달리 shifting은 mlp 등에서도 증명된 방법론이다. 

저는 무엇보다도 nlp와 cv 모델이 만약 같아질 수 있는 하나의 backbone이 존재한다면, multi-modal 문제나 여러 task 처리에 아주 효율적일 것이라는 생각이 듭니다. 

2. 관련 지식

CNN외에도 computer vision에는 다음과 같은 시도들이 있었습니다.

 

a. self-attention 기반의 backbone 모델 (sliding window 사용)

b. CNN 구조에 self-attention/transformer이 추가되는 방식

c. Transformer 기반의 backbone 모델 

 

각 분야에 해당하는 연구들은 논문 각주에 보다 자세히 나와있으니 관심 있으신 분들은 확인해보면 좋을 거 같습니다.

3. 모델

Swin Transformer는 입력 image를 겹치지 않는(non-overlapping) patch들로 쪼개는 것으로 시작합니다. 각각의 patch는 nlp에서 처럼 token(단어 혹은 sub-word 하나)으로 취급됩니다. patch size의 경우 pixel RGB를 고려하여 결정되는 데요, 예를 들어 4x4의 패치 크기를 갖는다면 이에 해당하는 feature dimension은 RGB를 고려하여 4x4x3이 되게 됩니다. 이러한 raw 값을 linear한 embedding layer에 통과시켜 C의 차원을 갖는 벡터를 project합니다.

 

임베딩된 feature vector는 transformer block에 들어가게 됩니다. hierarchical representation을 위하여 token은 네트워크가 깊어짐에 따라 merging을 통해 줄어들 게 됩니다. (예를 들어 첫 matching layer에서는 2x2의 주변 패치들을 통합시켜 만든 4C차원의 feature를 layer에 통과시킵니다). 앞 예시대로 진행될 때 해상도가 자연스레 downsample되는 데, 2C의 차원을 갖는 결과물이 나오게 됩니다. 

 

transformer block은 MSA(multi-head self attention)으로 이루어져 있습니다. shifted window 방식을 사용한 것이 특징이라 할 수 있고, 나머지는 transformer와 동일합니다. MSA외에도 2개의 MLP와 GELU를 사용하여 비선형성을 확보합니다. 아래 그림과 같이 Layer Normalization과 residual network도 사용되었고요.

transformer block의 형태

그렇다면 shifted window를 기반으로 한 self-attention이 대체 뭐길래 이렇게나 강조하는 것일까요? 일단 단순한 개념은 global attention에서 벗어난 sparse attention이라 보면 될 거 같습니다. 이야기인 즉슨, compuation을 linear하게 줄여보기 위함과 동시에 dense prediction을 위한 방법이었다고 합니다. 각 window는 겹치지 않는 형태로 M x M의 패치들로 구성되게 됩니다. 기존의 경우 hxw 패치들에 대하여 hw * (2C)^2 + (hw)^2 * 2C의 complexity, 즉 hw에 대하여 quadraitc한 형태입니다. MxM의 패치를 쓸 경우, hw * (2C)^2는 그대로이고 뒤의 항이 (M)^2 * 2C * hw가 되어 hw에 대하여 linear한 형태로 변합니다.

swin의 transformer block은 두 개를 연속적으로 사용하는 데 하나는 위에서 언급했듯이 Shifted한 형태가 됩니다. 위에서 살펴본 아래 그림을 다시 보시면,

window가 W의 반 정도만큼 이전 layer보다 움직인 것을 알 수 있습니다. (height과 width 모두의 방향으로 움직입니다) 정확히는 기존 layer의 window가 top-left에서 시작하여 bottom-right방향으로 움직였음을 확인할 수 있습니다. 위와 같은 방식으로 주변에 있는 window와의 연결성을 확보하였다(cross-window connection) 주장합니다.

 

이 때, window size가 부족해서 M x M인 경우가 생길 수 있는데, 이 경우에는 padding을 하고 masked out을 하는 것이 naive한 해결책이라 합니다. 좀 더 확실한 방법은 아래와 같은데요, 

cyclic shift를 이용하여 만들어진 adjacent feature map들을 masking함으로 성능을 높일 수 있다 합니다.

 

이후, relative position bias라는 말이 등장합니다. 윈도우에 존재하는 patch들을 고려하여 M^2 x M^2 차원을 갖는 bias term이 attention에 추가되게 됩니다. 기존 transformer 모델 등에서 사용하는 absolute position embedding은 오히려 성능을 약화시키기 때문에 학습하는(learnt) relative position bias만을 사용하였다 합니다. 

4. 결과

 

위와 같이 모델을 구성하였다 합니다. Swin-T는 ViT-B/DeiT-B와 비슷한 복잡도를 갖는다고 합니다. 

ImageNet-1K와 COCO, ADE20K등을 사용하여 결과를 냈는 데 아래와 같습니다. (throughput도 구하는 등 상당히 computational 문제에 신경을 쓴 모습이 보입니다)

 

ablation study를 위해 window의 효과와 postional embedding과 관련하여 실험을 진행하였는데요, 결과는 아래와 같습니다. 

5. 결론

 

접근 방법과 문제 해결 방식이 나쁘지 않았던 논문입니다. 접근 방법의 경우 transformer의 한계점을 명확히 분석하면서 CNN의 장점을 살리고 CNN이 갖지 못하는 장점마저도 얻을 수 있다 생각합니다. 실제 결과도 좋았다 생각하고요! 가장 주목할 점은 아무래도 shifted window 알고리즘이었던 거 같습니다. 나머지는 사실 엄청나다 보기 힘들지만 shifted window를 통해 많은 문제를 해결할 수 있었던 거 같습니다.