Swin Transformer V2 : Scaling Up Capacity and Resolution

기존의 Shifted Window Transformer 를 더욱 큰 모델, 고해상도의 이미지에서 더욱 안정적으로 학습하는 방식과 모델의 성능을 높이는 기술들을 이야기하고있습니다.

Abstract

기존의 Swin Transformer를 30억 개의 parameter를 사용할 때에도, 1536×1536 크기의 이미지에도 적용할 수 있도록 개선한 논문입니다. 이 논문에서 제시하는 vision model의 크기를 키우는 방법은 다른 모델에서 충분히 적용 가능하다고 이야기하고 있습니다.

기존의 Swin Transformer에서 해결해야 할 부분은

• Vision model은 크기를 키울 때 학습이 불안정한 문제가 있고
• 고해상도의 이미지 혹은 window를 필요로 하는 많은 downstream vision task 의 경우 어떻게 낮은 해상도를 처리하는 모델에서 높은 해상도를 처리하는 모델로 전이학습(transfer learninng) 이 효과적으로 잘 될 수 있는지 불분명하다는 것입니다.

특히 모델을 돌릴 때 큰 모델, 고해상도의 이미지를 처리할 때는 GPU 메모리 사용량 도 중요한 부분인데요, 이를 해결하기 위해 여러 기법을 적용해보았습니다.

1. 큰 모델의 학습 안정성을 높이기 위해 normalization을 attention 이전이 아닌 다음에 적용하고( post normalization ) scaled cosine attention 접근법 을 적용합니다.
2. 저해상도에서 고해상도 모델로 전이학습할 시 positional information를 log-scale로 continuous하게 표현합니다.
3. 이외에 GPU 메모리 사용량을 줄이는 여러 기법을 이야기합니다.

Introduction

Introduction 에서는 세가지의 문제점을 설명하며 이에 대한 해결방안을 말해줍니다.

1. 작은 크기의 vision model에서 큰 vision model로 전이학습할 때 학습이 불안정한 문제가 존재합니다.

논문에서 저자들은 activation의 진폭(amplitude)이 불일치한다는 것을 발견했다고 합니다.

이는 기존의 Swin Transformer 에서는 residual unit 이 main branch에 그대로 더해져서 layer를 지나갈수록 값이 누적되어 점차 크기가 커지기 때문이라고 합니다. 그래서 이를 해결하기 위한 방법으로 post normalization 기법을 적용하였습니다.

위의 그림에서와 같이 layer norm을 각 residual unit의 처음 부분에서 끝 부분으로 옮기는데요, 이렇게 하면 activation 값이 한층 낮아지게 됩니다.

또한 기존 dot-product attention 대신 scaled cosine attention 을 사용하였는데, 이는 block 입력의 진폭에 별 관계없이 연산이 수행되고 따라서 attention value는 안정된 값을 유지한다는 이점을 가집니다.

이러한 방법들을 통해 큰 모델로 전이학습 시 학습의 안정성을 높이고 최종 성능 또한 향상된다고 이야기하고있습니다.

2. 작은 크기의 이미지를 다룰 때의 window size와 큰 크기의 이미지에서 fine-tuning할 때의 window size의 차이가 큰 문제입니다.

기존의 방식은 bi-cubit interpolation 을 진행하였으나 문제 해결은 되지 않았기에, 여기서는 log-scale continuous position bias(Log-CPB) 를 제안합니다.

이는 작은 meta network를 하나 만들어 적용하는 방법인데요, meta net은 임의의 좌표를 받아 log-space로 변환을 진행하기 때문에 extrapolation 비율이 작으며, 모델 간 공유가 가능하기 때문에 이미지의 해상도가 변해도 문제 없이 사용할 수 있습니다.

3. 또한 고해상도 이미지를 처리하려면 GPU 사용량이 크게 증가하는 문제점이 있습니다.

이를 줄이기 위해 다음과 같은 방법을 사용합니다.

• Zero-Optimizer
• Activation Check Pointing
• a Novel implementation of sequential self-attention computation

위의 방법들을 사용하여, 더 큰 모델을 안정적으로 학습, 더 좋은 성능을 얻을 수 있습니다.

기존의 Swin Transformer와 마찬가지로 Image Classification, Object Detection, Semantic Segmentation task에서 실험하였고, 기존보다 더 좋은 결과를 얻었습니다.

Related Works

Language Networks and scaling up

NLP 분야에서는 이미 큰 모델일수록 좋다는 것이 입증되어 있다는 부분을 언급합니다.
BERT-340M, Megatron-Turing-530B, Switch-Transformer-1.6T 등이 제안되었습니다.

Vision networks and scaling up

하지만 이미지 분야에서는 모델 크기를 키우는 시도가 별로 없었다고 언급하고 있습니다. JFT-3B 정도가 존재합니다.

Transferring across window / kernel resolution

CNN 논문들은 고정된 크기의 kernel(ex. 1, 3, 5) 만을 사용하였습니다.

ViT와 같은 global vision transformer에서는 이미지 전체에 attention을 적용하며 입력 이미지의 해상도에 따라 window size가 정비례하여 증가합니다.
기존 Swin Transformer와 같은 local vision transformer는 window size가 고정되어 있거나 fine-tuning 중에 변화할 수 있습니다.

다양한 scale의 이미지를 처리하기 위해서는 아무래도 window size가 가변적인 것이 편하기에, 본 논문에서는 log-CPB 를 통해 전이학습이 좀 더 부드럽게 진행되도록 합니다.

Study on bias terms, Continuous convolution and variants

기존 Transformer는 절대적 위치 embedding 을 사용하였으나 요즘의 이미지 분야 연구에서는 상대 위치를 사용합니다. 이 논문에서 계속 강조하는 log-CPB는 임의의 window size에서 더 잘 작동한다고 이야기하고 있습니다.

Swin Transformer V2

Introduction에서 설명하였던 문제들을 제기하며, 모델에 대한 구체적인 내용을 서술하고 있습니다.

Post Normalization

Post Noramlization 방식은 layer norm을 block의 처음에서 끝으로 옮기는 방식입니다. 그러면 layer를 계속 통과해도 통과 후에 normalization을 거치니 activation 진폭이 크게 커지지 않게 됩니다.

위의 그림에서 pre- 방식과 post- 방식의 activation amplitude 차이를 볼 수 있습니다.

Scaled cosine attention

기존 Transformer는 query와 key 간 유사도를 계산할 때 dot-product로 계산합니다. 하지만, 큰 규모의 vision model에서, 일부 pixel에 의해 그 전체 값이 dominated 되는 문제가 발생합니다.

따라서 dot-product 대신 scaled cosine 연산 을 적용합니다. 물론 기존의 Swin Transformer처럼 position bias는 이 과정에서 더해집니다.

수식은 다음과 같습니다.

이 때, $\tau$ 는 학습가능한 scalar 값이며 layer 혹은 head 간 공유 되지 않습니다.

Scaling UpWindow Resolution

여러 window resolution에서 부드럽게 잘 넘어갈 수 있는 relative bias 방식인 log-CPB 를 살펴봅시다.

Continuous relative position bias

parameterized bias를 직접 최적화하는 대신 continuous position bias 를 상대 좌표 하에서 작은 meta network에 맞춥니다.

여기서 $\mathcal{G}$ 는 2-layer MLP(activation func : RELU) 와 같은 작은 network 입니다.

이러한 $\mathcal{G}$ 가 임의의 상대 좌표에 대해 bias value 를 생성하고, 어떤 크기의 window 에 적당한 값을 생성할 수 있습니다.

inference 때 각 상대 위치를 사전에 계산할 수 있고 모델 parameter로 저장할 수 있으며, 이는 원래 parameterized bias 접근법과 같은 수준으로 용이하다고합니다.

Log-spaced coordinates

해상도가 커질수록 상대 좌표의 차이도 커지는데, 이를 정비례하게 잡으면 그 차이가 너무 커지게 됩니다. 따라서 log-scale 로 변경하여 조절합니다.

예를 들어 8×8 크기의 window를 16×16 으로 키우면, 선형 비례하게 차이를 잡을 경우 [-7,7] 에서 [-15,15] 가 되므로 8/7 = 1.14배만큼 extrapolate 해야합니다.
하지만 log-scale의 경우 0.33배로 줄어들게됩니다.

아래 표에서 position bias 정책 별 성능 차이를 볼 수 있습니다. 물론 log-space CPB가 가장 좋다고 합니다.

Other Implementation

여기에서는 GPU 메모리 사용량을 줄이기 위한 방법들을 설명합니다.

Zero-Redundancy Optimizer (ZeRO)

기본 optimizer는 data-parallel mode에서 모델 parameter와 optimization state를 모든 GPU나 master node에 broadcast합니다.

이는 큰 모델에서 매우 비효율적인데요, 예를 들어 3B개의 parameter + Adam optimizer, fp32 정밀도를 사용할 경우 48G의 메모리를 잡아먹게됩니다.

ZeRO optimizer 는 모델 parameter와 그에 연관되는 optimizer state가 여러 개의 GPU에 분산 저장되어 사용량이 매우 줄어듭니다.

본 논문에서는 DeepSpeed framework와 ZeRO stage-1 option을 사용하였습니다.

Activation check-pointing

Transformer의 feature maps 역시 많은 메모리를 잡아먹는데, 이는 이미지와 window가 고해상도일 경우 병목효과를 발생시킬 수 있습니다. 이 최적화는 학습속도를 30% 가량 향상시킬 수 있다고합니다. (??)

Sequential self-attention computation

큰 모델을 초고해상도의 이미지, 예를 들어 1536×1536 해상도의 이미지와 32×32의 window size를 사용할 때는, 위의 두 최적화 전략을 사용해도 여전히 부담스럽습니다(40G).

여기서 병목은 self-attention module인데, 이를 해결하기 위해 이전의 batch 연산을 (전부) 사용하는 대신 self-attention 연산을 순차적으로 수행하도록 구현하였습니다.

이는 첫 2개의 stage에서만 적용하였고, 이는 전체 학습 속도에는 큰 영향을 주지는 않습니다.

Experiments

학습 시간을 줄이기 위해 192×192 해상도의 이미지로 사전학습을 진행하였습니다.

먼저 ImageNet-22K로 20 epoch 동안 self-supervised 학습을 시키고 같은 데이터셋에 classification task로 30 epoch 동안 학습시킵니다.

SWIN TRANSFORMER V2 - SCALING UP CAPACITY AND RESOULTION - 논문 보기