[논문 리뷰] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows (2021)
나의 정리
논문이 지적한 문제점: NLP에서 사용되는 Transformer가 vision task에 적용되는데 큰 문제점이 두 가지 있다.
- visual entity의 scale이 큰 variation을 가진다.
- image resolution이 커지면 computation cost가 매우 커진다.
해결 방안: hierarchical feature를 생성하는 transformer를 사용한다.
계층 구조를 사용하기 때문에 다양한 image size에 대해서 사용을 할 수 있고 shift window를 사용해서 고정된 patch에서만 self-attention이 적용되는 것을 방지한다.
- shift window도 padding 방식을 효율적으로 적용하였다.
- positional embedding도 절대 좌표가 아닌 상대 좌표를 사용해 더 좋은 결과를 냈다고 한다.
- 기존 CNN기반 ResNet을 backbone으로 주로 사용했는데 backbone으로 transformer를 사용할 수 있는 좋은 논문인 것 같다.
Abstract
NLP에서 주로 사용되는 Transformer가 image task에서 적용이 어려운 이유가 2가지 있다.
- visual entity의 크기(scale)가 매우 큰 variation을 가진다는 점
- text에 비해 image의 pixel은 매우 높은 해상도를 가진다는 점
따라서 이러한 문제점들을 shifted window를 사용하는 계층 transformer 구조로 해결한다.
Hierarchical 구조의 경우 다양한 크기에 대해서 유연하기 때문에 이미지 사이즈에 대한 계산 복잡도가 기존엔 제곱에 비례했지만 linear 하게 변한다.
Introduction
기존의 image task에선 CNN 기반의 연구가 활발하고 지배적이었지만, Transformer를 적용시켜 backbone으로 사용하려는 노력을 한 연구이다.
NLP와 Computer Vision은 위에서 언급한 scale과 resolution에서 차이점이 존재한다.
scale이 다양한 문제는 NLP에서 사용되던 transformer의 경우 일정한 크기의 token을 사용하다 보니 적절하게 적용이 되지 않았다.
resolution이 더 큰 문제는 image의 size가 커질수록 self-attention의 계산 복잡도가 급격하게 커지게 된다. (quadratic)
기존의 ViT와는 다르게 계층 구조를 가진다.
빨간색 선으로 나뉜 것을 window라고 칭하고 window 내부에서만 self-attention을 진행하게 되어 연산량을 줄인다.
작은 크기의 patch로 시작하여서 layer가 깊어짐에 따라 주변 patch들과 병합하는 방식으로 진행해서 FPN, U-Net과 같은 dense prediction이 가능하다. (segmentation task)
Swin Transformer에서 가장 key design은 위 그림의 shifted window다.
왼쪽 그림과 같이 patch를 나누어 self-attention을 진행한 뒤 다음 layer에서 window를 patch 1/2 크기만큼 오른쪽 밑으로 shift 시킨 뒤 self-attention을 진행한다.
더 자세한 내용은 뒤에서 다시 설명한다.
이전 sliding window 방식은 다른 query pixel에 대해 다른 key set을 가지게 되어 낮은 latency를 가지게 되는 반면, shifted window 방식은 모든 query patch는 같은 key set을 공유해서 memory access면에서 latency가 더 적다. (real-time latency)
Related Work
- CNN and Variants
그동안 CNN이 mainstream이었고 VGG, GoogleNet, ResNet, DenseNet, HRNet, EfficientNet 등 많은 논문들이 나오면서 발전했다. - Self-attention based backbone architectures
ResNet 기반의 spatial convolution layer를 self-attention layer로 대체하려는 시도가 있었으나 앞서 언급한 sliding window 방식의 memory access 문제가 있어서 어려움을 겪었다. - Self-attention/Transformers to complement CNNs
CNN을 보완하기 위해서 self-attention을 사용하려는 방법을 시도해보았다. - Transformer based vision backbones
Transformer를 vision task의 backbone으로 사용하기 위해서 연구되는 방향이다. 대표적으로 ViT, DeiT들이 있다.
Method
Overall Architecture
위 그림이 전체 architecture인데 크게 4가지 stage로 구성되어 있다.
먼저 Patch Partition 과정을 거치게 된다.
여기선 먼저 ViT와 같은 patch splitting module을 이용해 non-overlapping patch로 RGB 이미지를 나누어준다.
여기서 나눠진 patch는 token으로 사용되고 이 feature는 raw pixel의 RGB 값을 이어 붙인 것이 된다.
즉 patch가 4*4로 나누어진다면 RGB value를 곱해주어 최종적으로 4*4*3 = 48 dimension을 가지게 된다.
- stage 1 - Linear Embedding
이후 linear layer를 통하여 H/4 * W/4 * 48 텐서를 H/4 * W/4 * C 텐서로 변환한다.
여기서 C는 model의 크기에 따라서 달라지게 된다. (96-T, 96-S, 128-B, 192-L 왼쪽부터 tiny 모델)
위에서 언급했듯 layer가 깊어짐에 따라 주변 patch를 병합하면서 token의 개수를 줄여나가게 된다.
→Patch Merging layer를 통해서 병합을 해준다.
- stage 2~4 - Patch merging
처음 4*4 크기의 작은 patch들을 점점 합쳐가며 계층 구조를 만드는 과정이다.
인접한 2*2개의 patch를 concat 하여 채널 수가 4배가 된다.
이를 linear layer를 사용해 최종적으로 2C를 만들어 준다. - Swin Transformer block
주목할 점은 처음 block에서는 Window Multi-head Self Attention(W-MSA)를 통과하고 그다음 block에서는 Shifted Window based Multi-haed Self Attention(SW-MSA)를 통과하게 된다.
각각 2-layer MLP, Layer Norm(LN), GELU를 통과하게 된다.
Shifted Window based Self-Attention
계속 언급했듯 기존 self-attention은 모든 query와 key set에 대해 연산을 해야 돼서 해상도가 높아짐에 따라 계산 복잡도가 제곱에 비례해 매우 커진다.
M(window size)는 h*w(image size)에 비해 훨씬 작기 때문에 연산량이 적고 image size가 커지더라도 ViT에 비해 연산량을 줄일 수 있다.
하지만 window가 fix 되어 있기 때문에 self-attention시에 고정된 부분만을 수행한다는 단점이 있어 window를 shift 하여 한번 더 self-attention을 수행해서 문제를 해결하였다.
Shifted window partitioning in successive blocks
앞에서 언급한 shift 시키는 연산에 대한 설명이다.
위의 그림으로 shift가 이루어진다.
window를 shift 시키는 것을 cyclic shift라고 칭한다.
window size의 1/2 만큼 우측 하단으로 shift 하고 A, B, C 구역을 padding 시킨다.
이때 padding 시키는 부분은 반대편인 좌 상단에서 온 것이므로 A, B, C를 포함해서 self-attention을 진행하는 것은 의미가 없다. (실제로는 떨어져 있는 pixel들이므로)
따라서 A, B, C에 mask를 씌운 뒤 self-attention을 수행한다.
그 뒤에 reverse cyclic shift를 진행해 원래 값으로 되돌린다.
그냥 padding을 사용해서 이 방식을 대신할 수 있지만, computational cost가 증가될 수 있어 이 방법을 사용했다고 한다.
Relative position bias
Swin Transformer는 ViT와 다르게 position embedding을 처음부터 더하지 않았다.
self-attention 과정에서 relative position bias를 추가함으로써 그 역할을 대체하였다.
SoftMax를 취하기 전에 B를 더해주는데 이 B가 Relative position bias이다.
기존의 absolute position embedding보다 relative position bias를 주는 것이 더 효과적이라고 주장한다.
Conclusion
- hierarchical feature를 생성하는 새로운 vision Transformer를 제시한다.