Transformer가 computer vision에서 어떻게 쓰였는지, 관련 모델이 어떻게 발전하고 있는지 여러 포스팅에 걸쳐서 알아보고자 한다.
Computer vision domain에서 Transformer를 사용하는 대표적인 아키텍쳐는 Vision Transformer(ViT, An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, Alexy et al., 2021)이고, 이를 보완한 것이 Swin Transformer(Swin Transformer: Hierarchical Vision Transformer Using Shifted Windows, Liu et al., 2021)이다.
Computer vision에서의 두 Transformer-based 모델을 알아보자.
기초 내용인 Attention과 Transformer는 아래 글을 참조하자.
Transformers in Vision - (1) Attention & Transformer
Transformer가 computer vision에서 어떻게 쓰였는지, 관련 모델이 어떻게 발전하고 있는지 여러 포스팅에 걸쳐서 알아보고자 한다. 이번 포스팅에서는 가장 중요한 기초 내용인 attention과 transformer에 대
jjuke-brain.tistory.com
목차
ViT (Vision Transformer)
2021년 ICLR 학회에서 발표된 'An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale' 논문에서 제안한 Vision Transformer는 이미 NLP에서는 보편화된 transforemr 개념을 computer vision 분야에 적용하는 방법이다.
Introduction
당시 computer vision 분야에서는 여전히 CNN 기반 아키텍쳐가 여전히 주류를 이루었다.
그런데 이 논문에서는 CNN없이 오직 transformer만 사용하는데, image를 여러 patch로 나누어 각 patch를 NLP의 token(sequence)을 다루듯이 transformer를 적용한다. Image classification task에서 기존 방법들보다 훨씬 적은 computational cost로 (당시에) 최고 성능을 달성했다.
원리를 간단히 살펴보자.
- Image를 정해진 크기의 patch로 나눈다.
- Linear하게 embedding한 후,
- Position embedding을 추가한다. (Image classification task의 경우, 학습 가능한 classification token을 sequence에 추가해준다.)
- 위 결과 sequence를 Transformer encoder에 입력해준다.
- (Image classification task를 수행하는 경우) MLP head를 거쳐 class를 예측하도록 한다.
Method
저자들은 최대한 원래의 Transformer와 가까운 방법으로 모델을 설계했다. Transformer의 '간단하면서 scalable한 setup, 효율적인 구현'이라는 장점을 최대한 활용하기 위해서이다.
Patch Embedding and Position Embedding
먼저 Patch embedding에 대해 알아보자.
NLP에서의 Transformer는 token embedding인 1D sequence를 입력으로 받는데, 이미지는 2D이므로 이를 reshape하여 2D patch의 sequence로 만든다. 또한 모든 layer에서 latent vector size가 일정하므로, patch를 flatten한 후 일정 차원(/(D\))으로 맵핑시켰다. (이때의 linear projection은 학습 가능한 parameter이다.) 그 결과가 바로 patch embedding이다.
BERT의 token에서처럼 patch embedding에 학습 가능한 embedding을 추가한다. 이에 따라 Transformer encoder의 결과 벡터는 image representation의 역할을 할 수 있게 된다.
Position embedding은 위치 정보를 보존하기 위해 patch embedding에 추가되는 임베딩이다. 굳이 2D position embedding을 사용해도 큰 성능 개선이 없어서 1D (learnable) position embedding을 사용한다.
이렇게 Patch embedding과 Position embedding이 결합한 형태의 embedding이 Transformer encoder의 입력으로 주어지게 된다.
Pre-training과 fine-tuning 과정에서 classification head가 Transformer encoder의 결과 벡터에 attach되는데, 이떄 classification head는 pre-training의 경우에는 hidden layer가 하나인 MLP, fine-tuning의 경우에는 linear layer 하나인 MLP이다.
전형적으로 ViT는 큰 dataset에서 pre-train하여 작은 downstream task에 fine-tuning을 거쳐 적용한다. Fine-tuning을 통해 pre-training에서보다 큰 해상도의 이미지를 다룰 수 있다.
(같은 patch size에서) 이미지의 해상도가 클수록 sequence length가 커질텐데, ViT는 memory만 충분하다면 sequence length에 제한이 없다. 하지만, pre-training에서의 position embedding이 의미가 없어진다. 따라서 원래 이미지에서의 위치에 따라 pre-training했던 position embedding에 2D interpolation을 적용한다.
Resolution adjustment와 patch extraction은 이미지의 2차원 구조와 관련된 유일한 inductive bias이므로 중요하다!
Transformer Encoder
Transformer encoder는 Multi-head Self-Attention(MSA)과 MLP block으로 구성된다. 또한 모든 block 앞에는 LayerNorm(LN)이 적용되고, 뒤에는 residual connection이 적용된다. (MLP는 두 layer와 GELU non-linearity로 이루어진다.)
Inductive Bias
Vision Transformer는 image task에서 CNN보다 inductive bias가 훨씬 적다.
CNN에서는 locality와 2차원적인 neighborhood structure, 그리고 translation equivariance가 모델 전체에 걸쳐서 각 layer에 부여된다. 하지만 ViT의 경우 MLP layer에서만 locality와 translation equivariance가 있고, self-attention layer에서는 globality를 갖는다. (2차원 neighborhood structure는 image를 patch로 splitting할 때와 fine-tuning에서 매우 조금만 존재한다.)
게다가 position embedding 과정에서 patch의 2D 위치 정보를 전혀 반영하지 못하고, patch간의 모든 spatial relation은 처음부터(from scratch) 학습된다.
Hybrid Architecture
Image patch를 단순히 나누어 사용하는 대신에 CNN의 feature map으로 input sequence를 만들어낼 수 있다. 이를 hybrid model이라 하는데, 여기서는 patch embedding 과정에서 patch를 CNN feature map에서 추출한다.
이외 과정은 위에서 설명한 것과 같다.
Results
Table 1과 같이 ViT도 BERT(NLP에서 좋은 성능을 나타내는 transformer 기반 모델)와 마찬가지로 configuration을 갖고 있다. 다양한 ViT를 나타내기 위해 다음과 같은 표기법을 사용한다.
- ViT-L/16 : "Large" variant with \(16 \times 16\) input patch size
Transformer의 sequence length는 patch size의 제곱에 반비례한다는 사실에 주목하자. 즉, patch size가 작을수록 computational cost가 높아진다.
Table 2에서는 ResNet, EfficientNet과 ViT를 비교한 표이다. 값은 3번의 fine-tuning을 실행하여 정확도의 평균과 표준편차의 평균을 나타낸 값이다. 각 열에서의 JFT, I21k, BiT-L, Noisy Student는 각 모델에 대해 pre-training을 진행할 때 사용한 데이터셋이고, 이러한 모델들로 ImageNet, CIFAR, Oxford, VTAB 등의 데이터셋에 대해 classification을 진행했다.
모든 모델은 TPUv3 하드웨어로 학습했으며, TPUv3-core-days는 '하루동안 학습에 사용된 core 수 * 학습 시간'을 나타낸다. 예를 들어, ImageNet-21k 데이터셋으로 pre-train한 ViT-L/16 모델은 TPUv3 8코어로 약 30일정도(곱하면 약 240) 학습했다고 한다.
그리고, ViT는 inductive bias가 CNN에 비해 훨씬 약하다고 하였으므로, pre-training 시에 아주 큰 데이터셋을 사용해야 한다. 얼마나 큰 데이터셋을 사용해야 Vision Transformer가 ResNet보다 좋아지는지에 대한 실험 결과를 살펴보자.
Figure 4는 JFT-300M 데이터셋에서 일정 수만큼의 데이터를 샘플링한 subset으로 pre-training을 진행한 후, ImageNet 데이터에 대해 fine-tuning하여 few-shot accuracy를 구한 결과이다.
서로 학습 시간이 비슷했던 ViT-B/32와 ResNet50, 그리고 ViT-L/16과 ResNet152x2를 비교해보면 pre-training 데이터 수가 많아질수록 ViT의 성능이 더 좋아짐을 볼 수 있다. 즉, 작은 데이터셋에서는 convolutional inductive bias가 더 유용하지만, 데이터수가 많아지면 데이터에서 직접적으로 관련 pattern을 학습하는 것이 더 유용하다고 볼 수 있다.
Vision Transformer가 학습하는 원리
Vision Transformer가 이미지를 처리하는 방법을 이해하기 위해, 모델이 학습하는 과정에서의 representation을 분석해보자.
Figure 5는 patch들을 Linear projection하여 저차원 공간으로 맵핑할 때, 학습된 embedding filter의 top principal component(주성분)이다. Component들은 patch들이 갖는 미세한 구조의 저차원 representation에 대한 기저 함수(basis function)와 닮아 있다.
Projection 이후에는 patch representation에 position embedding이 추가된다고 하였다.
Figure 6은 모델이 position embedding의 유사성에서 (이미지 내에서의) 거리를 인코딩하는 방법을 학습한다는 것을 보여준다. 즉, 모델이 Position embedding을 학습하게 되면 위치가 가까운 patch일수록 비슷한 position embedding을 갖게 된다.
1D position embedding을 사용하는데도 2D image topology에 대한 representation을 학습한다. (그림에서 행, 열의 구조가 나타난다. 즉, 같은 행이나 열의 patch는 비슷한 position embedding을 갖는다.)
Self-attention 덕분에 ViT는 가장 낮은 layer에서도 이미지 전체에 대한 정보를 통합한다. Figure 7의 Attention distance는 CNN에서의 receptive field size와 같은 개념이다.
어떤 head들은 0번째 layer에서 이미 이미지 대부분에 attend한다. (예를 들어 Head 1은 0번째 layer에서 평균 attention distance가 아주 높다.) 반면 다른 head들(빨강, 진갈색 등)은 낮은 layer에서 작은 attention distance를 갖는데, 이러한 head들은 CNN의 convolutional layer와 같은 역할을 할 것으로 볼 수 있다.
Network가 깊어지면서 attention distance가 전체적으로 오르는데, 이는 곧 Figure 8과 같이 모델이 이미지를 분류하는 데 관련된 부분에 더 집중(attends to)한다는 것을 의미한다.
Swin Transformer
2021년 ICCV 학회에서 발표된 'Swin Transformer: Hierarchical Vision Transformer Using Shifted Windows'논문에서 제안한 Swin Transformer는 'shifted window'라는 개념을 사용하여 hierarchical representation을 계산하는 Vision Transformer이다. 이를 통해 Transformer를 image에서의 pixel에 적용할 때와 text에서의 word에 적용할 때의 차이를 극복하고자 했다.
Introduction
Vision vs Language for Transformer
컴퓨터비전 분야와 자연어 처리 분야에서 Transformer를 적용할 때, 기본 요소가 서로 다르다. 자연어 처리에서는 기본 요소가 word token이며, scale이 고정되어 있다. 하지만 컴퓨터 비전에서는 resolution에 따라 이미지의 size(scale)가 바뀌므로 attention을 적용할 때 문제가 된다.
그리고 두 분야에서는 resolution이 다르다. Text의 단어들에 비해 image에서의 pixel의 해상도가 훨씬 높아서, self-attention의 계산 복잡도가 image size의 제곱에 비례하므로 pixel-level의 dense prediction에 적용하기가 매우 힘들다.
Swin Transformer의 특징
Swin Transformer는 컴퓨터 비전 도메인의 다양한 task에 적용할 수 있는 Transformer 기반의 general backbone 아키텍쳐이다.
위에서 언급한 두 가지 문제점을 다음 방법으로 해결하였다.
- Scale 문제 : Hierarchical feature map으로 구성하여 해결
- Transformer layer가 깊어짐에 따라 작은 patch부터 점진적으로 patch size를 키워 간다. (Figure 9에서의 회색 line)
- FPN(Feature Pyramid Networks)이나 U-Net을 사용하여 dense prediction을 할 수 있다.
- Resolution 문제 : Shifted window 방법으로 해결
- Image를 windows(빨간색 line)로 나누고, self-attention을 각 window에 대해 계산한다.
- 각 window의 patch 개수는 고정되므로, 계산 복잡도는 image size에 선형이다.
Figure 9의 (b)에서 볼 수 있듯, ViT는 유일한 resolution의 feature map만 제공하고, 계산 복잡도가 image size의 제곱에 비례한다.
Shifted Window
Shifted window approach는 두 연속된 self-attention layer 간에 window를 이동시키는 방법으로, Swin Transformer의 주요 design 요소이다.
Figure 10에서 Layer \(l\)에서는 ViT에서와 같이 일반적으로 window partitioning을 진행한다. Layer \(l+1\)에서는 이 window를 이동시킨 것으로, 이 결과를 shifted window partitioning이라 한다.
이전 window에서 나누어졌던 부분이 shifted window에서 계산되므로 layer간의 connection이 생기고, 이에 따라 modeling power가 대폭 강화된다.
또한, window 내의 모든 query patch들이 같은 key set을 공유하는데, 이에 따라 memory access가 가능해진다. (기존 sliding window 기반의 self-attention 방법은 query pixel마다 서로 다른 key set을 가져 hardware에서 계산 시 low latency 문제가 있었다.)
이는 모든 MLP 아키텍쳐에서 유효한 방법이다.
이처럼 Swin Transformer는 컴퓨터 비전과 자연어 처리를 통합한 아키텍쳐로, visual & textual signal의 joint modeling이나 두 도메인의 knowledge를 공유하는 modeling을 가능하게 한다.
Method
Overall Architecture
Swin Transformer의 전체적인 구조는 다음과 같다.
Architecture의 동작 과정을 자세히 알아보자.
1. input RGB image를 받아 겹치지 않는 patch로 나눈다.
ViT와 동일한 방법으로 patch를 나눈다. 각 patch는 token으로 간주하고, feature는 pixel의 RGB 값들을 연결(concatenate)한 형태이다.
Patch size를 \(4 \times 4\)라 하면, 각 patch의 feature dimension은 \( 4 \times 4 \times 3 = 48\)일 것이다.
그리고 feature에 linear embedding layer를 적용하여 임의의 dimension \(C\) (channel)로 projection한다.
그 후, Swin Transformer block은 token의 resolution을 \(\left( \cfrac{H}{4} \times \cfrac{W}{4} \right) \)로 유지시키면서 self-attention을 계산한다.
2. Hierarchical representation을 적용한다.
Patch Merging Layer는 network가 깊어짐에 따라 token 개수를 줄이는 역할을 한다.
\(2 \times 2\)개의 이웃 patch의 feature를 연결(concatenate)한 \(4C\)차원의 feature에 linear layer를 적용한다. 따라서 patch(token, Figure 15에서 회색 outline)의 개수가 \(4 \times 4\)에서 \(2 \times 2\)가 되며, patch의 차원은 \(2C\)가 된다.
참고로 Figure 15에서 빨간 outline은 window이고, 같은 개수(그림에서는 \(4 \times 4 = 16\))의 patch를 갖는다.
이후에 Swin Transformer Block을 적용하여 feature transformation을 수행한다. 마찬가지로 resolution은 \(\left( \cfrac{H}{8} \times \cfrac{W}{8} \right)\)을 유지한다.
이러한 Stage(Patch Merging Layer + Swin Transformer Block)를 총 3번 반복하면 resolution은 \(\left( \cfrac{H}{4} \times \cfrac{W}{4} \right) \rightarrow \left( \cfrac{H}{8} \times \cfrac{W}{8} \right) \rightarrow \left( \cfrac{H}{16} \times \cfrac{W}{16} \right) \rightarrow \left( \cfrac{H}{32} \times \cfrac{W}{32} \right) \)가 된다. (전체 pixel 수는 일정하므로 channel이 늘어나게 된다.)
이러한 hierarchical representation은 CNN 기반인 VGG, ResNet 등의 feature map resolution과 같다. 즉, CNN처럼 Swin Transformer를 backbone으로 활용하여 다양한 컴퓨터 비전 task에 적용이 가능하다.
Swin Transformer Block을 간단히 살펴보자.
Figure 16에서 볼 수 있듯, Swin Transformer Block의 내부 구조는 기존 ViT의 MSA(Multi-Head self-attention) 모듈을 shifted window 기반의 W-MSA와 SW-MSA로 대체하였다. W-MSA는 shift되기 이전의 regular partitioned window에서의 self-attention 모듈이고, SW-MSA는 shifted partitioned window에서의 self-attention 모듈이다. (다른 layer는 동일하다.)
Shifted Window based Self-attention
컴퓨터 비전에서의 기존 transformer 기반 아키텍쳐는 global self-attention을 사용하였다. 하지만 위에서 언급했듯이, 계산 복잡도가 image size의 제곱에 비례하여 고해상도 이미지 처리나 dense prediction이 불가능한 문제점이 있었다.
이를 해결하기 위해 도입한 shifted window partitioning에 대해 자세히 알아보자.
1. Self-attention in non-overlapped windows
Self-attention의 계산 복잡도를 알아보자.
겹치지 않게 나눈 windows에 대해 self-attention을 계산하는데, 만약 각 window가 \(M \times M\)개의 patch를 갖는다면, MSA 모듈과 W-MSA 모듈의 계산 복잡도는 다음과 같다.
\( \Omega (\text{MSA}) = 4hwC^2 + 2(hw)^2C \)
\( \Omega (\text{W-MSA}) = 4hwC^2 + 2M^2hwC \)
- \( \Omega (\text{MSA}) \) : image의 size를 나타내는 \(hw\)에 대한 2차식 (quadratic)
- \( \Omega (\text{W-MSA}) \) : \(M\)이 고정되었을 때, \(hw\)에 대한 1차식 (linear)
2. Shifted Window partitioning in successive blocks
Window들 간의 connection이 잘 되어 있을수록 modeling power가 강해지는데, shifted window partitioning은 연속된 Swin Transformer Block의 partitioning 설정을 바꿔줌으로써 connection을 강화한다.
다시한 번 Figure 10을 보면, Layer \(l\)에서는 regular window partitioning을 진행하며, Layer \(l+1\)에서는 window들을 \((\lfloor\frac{M}{2}\rfloor, \lfloor\frac{M}{2}\rfloor)\)만큼 옮긴(shift) window들을 사용한다.
연속된 Swin Transformer block에서는 shifted windows를 사용하여 다음과 같이 계산을 진행한다. Swin Transformer Block 그림과 함께 이해해보자.
\( \hat{\mathbf{z}}^l = \text{W-MSA} (\text{LN}(\mathbf{z}^{l-1} )) + \mathbf{z}^{l-1} \)
\( \mathbf{z}^l = \text{MLP} (\text{LN} (\hat{\mathbf{z}}^l)) + \hat{\mathbf{z}}^l \)
\( \hat{\mathbf{z}}^{l+1} = \text{SW-MSA} (\text{LN} (\mathbf{z}^l)) + \mathbf{z}^l \)
\( \mathbf{z}^{l+1} = \text{MLP} ( \text{LN}( \hat{\mathbf{z}}^{l+1} )) + \hat{\mathbf{z}}^{l+1} \)
- \(\hat{\mathbf{z}}^l\) : Block \(l\)의 W-MSA, SW-MSA 모듈의 output feature
- \(\mathbf{z}^l\) : Block \(l\)의 MLP 모듈의 output feature
3. Efficient batch computation for shifted configuration
Shifted window partitioning의 문제점은 Figure 10에서 볼 수 있듯 window를 옮기는 과정에서 \(h, w\)를 따라 가각가 하나씩 window의 개수가 늘어나고(\((\lceil \frac{h}{M} \rceil \times \lceil \frac{c}{M} \rceil) \rightarrow (\lceil \frac{h}{M} \rceil + 1) \times (\lceil \frac{c}{M} \rceil + 1)\)), 가장자리 부분의 window의 크기는 \(M \times M\)보다 작아진다는 것이다.
이를 해결하기위한 단순한 방법으로는 작은 window에 padding을 적용하여 \(M \times M\)으로 사이즈를 맞춰주고, attention 계산 시에는 padding된 값을 masking(계산하지 않음)하는 것이다.
하지만, 이 방법은 window 개수가 적은 경우에는 계산량이 많이 늘어난다. Figure 10을 예로 들면, \(2 \times 2\)개의 window에서 \(3 \times 3\)개의 window로, 계산량이 2.25배가 늘어난다.
따라서 efficient batch computation 방법을 사용한다.
먼저, 왼쪽 위 방향으로 cyclic-shift(\(M \times M\) 사이즈의 window가 가장자리에 오도록 옮기는 것)를 진행한다. 그러면 Figure 18에서의 A, B, C와 같이 feature map과 겹치지 않는 sub-windows가 생기는데, masking을 통해 sub-window들은 계산하지 않는다.
그 후 반대 방향으로 reverse cyclic-shift를 하여 같은 방법으로 계산하면 계산된 batched window는 regular window의 크기와 같아진다.
이를 통해 효율적으로, low latency(실험으로 증명)로 계산을 진행할 수 있다.
4. Relative position bias
Self-attention을 계산할 때, 각 head에서 similarity를 계산하는 과정에 relative position bias \(B \in \mathbb{R}^{M^2 \times M^2} \)를 적용한다.
\( \text{Attention}(Q, K, V) = \text{SoftMax}(QK^\top / \sqrt{d} + B)V \)
- \(Q, K, V \in \mathbb{R}^{M^2 \times d} \) : query, key, value matrices
- \(d\) : query, key의 dimension
- \(M^2\) : window 하나에 있는 patch의 수
Bias term이 없거나 position embedding만 사용하는 경우보다 성능이 훨씬 향상되었고, 이는 실험으로 증명되었다.
또한, Pre-train된 relative position bias를 bi-cubic interpolation을 사용하여 window size에 따라 다르게 조정하여 사용이 가능하다.
Architecture Variants
BERT, ViT 등과 마찬가지로 Swin Transformer 또한 크기에 따라 다양한 모델이 존재한다. (window size \(M = 7\), query dimension \(d = 32\), 각 MLP의 expansion layer \(\alpha = 4\)로 고정)
- Swin-T (Tiny version) : \(C=96, \text{layer numbers} = \{2, 2, 6, 2\}\)
- 모델 size : ResNet-50, DeiT-S와 비슷 (base의 0.25배)
- Swin-S (Small version) : \(C=96, \text{layer numbers} = \{2, 2, 18, 2\}\)
- 모델 size : ResNet-101과 비슷 (base의 0.5배)
- Swin-B (Base) : \(C=96, \text{layer numbers} = \{2, 2, 18, 2\}\)
- 모델 size : ViT-B/DeiT-B와 비슷
- Swin-L (Large version) : \(C=192, \text{layer numbers} = \{2, 2, 18, 2\}\)
- 모델 size : base의 2배
\(C\)는 first stage의 hidden layer의 channel 수이고, 이론적인 계산 복잡도(FLOPs)를 기준으로 하였다.
최근댓글