[논문리뷰 | CV] ViT : An Image Worth 16 x 16 Words : Transformers for Image Recognition At Scale

Transformer를 이미지에 적용하며 처음으로 괜찮은 성과를 낸 모델이자 그것의 시초가 된 모델이다. NLP를 공부하며 많이 봤었던 Transformer이기에 문제 없이 읽었으나 조금 어려웠던 부분들이 있어 해당 부분은 이후에 따로 정리해보고자 한다.

💡 1. Abstract

NLP분야에서 transformer 구조가 굉장히 지배적인 standard로써 자리매김해왔지만, vision 분야에서의 응용은 잘 되지 않았다.

vision 분야에서 attention을 사용하는 경우, 주로 convolutional networks와 혼합되어 사용되거나, convnet의 전체적인 구조를 기반으로 몇몇 요소들만 대체하는 식으로 사용되어왔다.

연구진은 해당 논문에서 이러한 CNN 의존 불필요성과, 오직 transformer가 image patch들의 시퀀스에 바로 적용해 이미지 분류 task에 매우 잘 동작할 수 있다는 것을 실험적으로 증명했다.

많은 양의 데이터에 pre-trained 한 뒤, 중 ~ 소 size의 image recognition 벤치마크 데이터셋에 전이되어 학습할 때, Vision Transformer(ViT)는 기존의 다른 성능 좋은 CNN들에 비해 요구되는 계산량은 상당히 적으면서도 성능은 좋았다

🚀 1. Introduction

self-attention기반 구조는 (특히 트랜스포머) NLP 분야에서 정말 많이 사용되고 있고 널리 사용되는 방식은 대규모 corpus에 대해 사전훈련 후 더 작은 task-specific한 데이터셋에 fine-tune을 진행하는 것이다.

transformer의 계산의 효율성과 확장성(scalability)을 통해 과거보다 더 큰 사이즈의 모델을 훈련하는 것이 가능해졌다.

그러는 와중에도 모델과 데이터의 사이즈 증가에도 성능이 saturating이 되지 않았다

• 깊게 쌓아도 쌓아도 괜찮음
• 사실 transformer를 깊게 쌓는 것에 대해서 resnet 같은 것을 듣지 못함
• 그 이유는 아무래도 hidden vector를 만들 때에 cnn과 달리 모든 input에 대한 고려가 간다
• 즉, 그 부분이 큰 문장에서 어떤 의미인지를 지속적으로 파악하는 것이기 때문
• 하지만 cnn은 도장을 찍으면서 종합해주는 방식 → 국소 위치에 대한 내용을 담게 되므로 소실이 발생
• 아마..?

하지만 여전히 computer vision 분야에서는 convolutional architecture가 대세이다. transformer의 NLP 영역에서의 성공을 보고 많은 연구들에서 CNN구조에 self-attention을 혼합하여 실험하였고, convolutional layer를 전부 self-attention으로 대체한 연구도 있었다.

후자의 모델들은 이론적으로는 효율적이지만, 어텐션의 독특한 계산방식 때문에 현대의 하드웨어 가속기에 효율적으로 확장되지는 못했다.

• 왜냐하면 이미지 전체를 그대로 넣어야 하기에 단어 벡터에 비해서 연산이 너무 클 것

그래서 large-scale 이미지 recognition에서는 여전히 클래식한 ResNet의 구조를 갖고 있는 모델들이 SotA로 남아있다.

ViT 연구진 역시 transformer의 성공에 영감을 받아 연구진은 표준 transformer의 구조를 조금 수정한 후 이미지에 직접 적용하는 실험을 진행

그것을 위해 저자들은 이미지를 작은 patch들로 분할하고 이 patch들의 linear embedding의 시퀀스를 트랜스포머의 인풋으로 전달

• 이를 통해 볼 때 이미지의 patch는 NLP에서의 token과 같

그런데 imageNet등의 중간 사이즈(데이터 양) 데이터셋에 대해 훈련을 진행했을 때, 비슷한 사이즈의 ResNet-like architecture 모델들에 비해 정확도가 떨어졌다. 이는 언뜻 보기에 아쉬운 결과였으에 이에 다음과 같은 생각을 했다.

→ generalize를 위한 모델 고유의 가정 즉, Inductive biases로써 CNN에는 내제되어있는 translation equivariance이나 locality와 같은 성질이 Transformer에는 부족하기 때문에 충분하지 못한 양의 데이터에 훈련될 경우 모델의 일반화 성능이 좋지 못할 것.

Inductive Bias란 무엇일까?

• 결국 cnn은 특정 범위에 대한 지식이 훨씬 뛰어나기 때문에 주어지지 않은 samples들에 대해서도 학습을 잘 할 것이라는 의미
• 그에 반해 transformer는 이미지 patch에 대해서 attention을 구하는 것이기 때문에 이미지들이 적을 때에 그 효과가 강하지 못할 것

그러나 모델이 보다 큰 데이터 셋(14M ~ 300M장)에 대해 훈련되는 경우 결과가 변하는데, 연구진은 대규모의 훈련이 모델의 inductive bias를 찍어누를 수 있다는 것을 확인했다.

ViT가 충분히 큰 데이터에 사전훈련된 이후 더 작은 데이터 셋에 대해 전이되어 학습이 이루어 질 때 기존 SotA모델들의 성능에 필적하거나 혹은 능가하는 훌륭한 결과를 냈다고 한다.

• 결국에는 transformer가 모든 경우에 대해 고려할 수 있을 정도로 데이터를 먹을 경우 잘 할 것이라는 의미

📦 2. Related Works

transformer가 2017년 Vaswani에 의해 machine translation 분야에서 제안되었고, 많은 NLP task에서 sota가 되었다. 대규모의 transformer 기반 모델들은 일반적으로 대규모 말뭉치에 pre-train을 시킨 후에 당면한 과제에 맞게 fine-tuning을 시키는 방식으로 학습을 진행했다.

• BERT는 denoising self-supervised을, GPT는 language modeling을 pre-training task로써 활용

이미지에 대한 self-attention의 단순한 적용은 각 픽셀이 다른 모든 픽셀들에 attend 될 것을 요구한다.

• 하나의 픽셀을 임베딩 할 때에 다른 픽셀들도 임베딩에 참여할 것을 요구한다는 것

그리고 이는 픽셀 수에 대해 매우 높은 계산 복잡도를 야기하며, 이에 다양한 input size에 확장되는 것이 어렵다. 이에 transformer를 image processing에 적용시키기 위해서 몇가지 근사(approximation)방법들이 시도됨

• local self-attention, sparse attention, applying it in blocks of varying size 등

이러한 많은 specialized attention 구조들은 컴퓨터 비전 task에서 괜찮은 결과를 보여주지만 효율적으로 구현되기 위해서는 복잡한 엔지니어링이 필요하다.

self-attention을 이용하여 feature map들을 augmenting하거나, CNN의 output에 attention을 추가적으로 사용하는 방식 등으로 CNN과 self-attention을 결합하는데도 많은 연구가 이루어졌다.

하지만 연구진들이 보았을 때 전체 이미지에 대해 transformer가 global self-attention을 수행하는 사례는 없었고 ViT와 가장 가까운 모델은 iGPT 뿐이라고 한다.

iGPT는 이미지의 해상도와 color space를 줄인뒤 image pixel들에 transformer를 적용한 모델이며, 하나의 generative model로써 비지도 방식으로 훈련되었고 그로부터 얻어진 representation은 분류를 위해 fine-tuning 되거나 선형탐색될 수 있다. 이를 예시로 활용했을 때 imagenet에서 최대정확도 72%로 나타났다.

해당 논문은 "기존 ImageNet 데이터셋보다 더 많은 이미지를 보유한 recognition 데이터셋을 화룡한 연구"의 연장선으로 볼 수 있다.

추가적인 데이터소스의 사용은 표준 벤치마크들에서 sota의 결과를 달성하게한다는 연구들이 있으며, 더 나아가 CNN의 성능이 데이터셋 size에 따라 어떻게 달라질 수 있는지에 대한 연구와 대규모 데이터셋(ImageNet-21k, JFT-300M)에 대한 CNN의 전이학습을 경험적으로 탐구하는 연구도 있었다.

• 즉, 위 두가지가 사용할 더 큰 데이터셋

본 연구는 바로 위에 언급된 두 데이터셋에 초점을 맞추고있으나 이전의 연구들에서 사용한 ResNet 이 아닌 Transformer를 사용하여 학습을 진행하였다.

🛠 Method

모델 설계 시 연구진은 기존 transformer와 최대한 유사하게 구성하려 함

이렇게 하는 것의 이점은 쉽게 확장이 가능한 NLP transformer 구조와 그것의 효율적인 구현을 거의 바로 사용할 수 있다는 것이다.

3.1 Vision Transformer(ViT)

모델에 대한 전반적인 overview는 위 그림과 같다.

기존 transformer는 1차원 토큰 임베딩 시퀀스를 input으로써 받는다. 모델 내부에서 2차원 image를 다루기 위해, 본래에 3차원인 이미지 $x∈R^{H×W×C}$를 2차원의 패치들 $x_p∈R^{N×(P^2⋅C)}$로 flatten 실시했다.

• 즉, 패치로 만든 후에 일렬로 squeeze시키는 방식

이때 H, W는 원본이미지의 해상도이며, C는 이미지의 채널 수

그리고 P, P는 패치들의 해상도이다. 이때 $**N=HW/P^2$은 패치의 수를 의미하는데, Transformer Encoder에 input되는 시퀀스 길이-개수(하이퍼파리미터)**로써 바라볼 수도 있다.

transformer는 내부의 모든 layer들의 latent vector의 size가 D로 통일시켰으며 2차원 patch들을 다시 1차원으로 flatten하고(size : $NP^2C$) 이를 trainable한 linear projection을 거쳐 D차원 벡터로 매핑시켰다.

• 이러한 projection의 결과벡터가 앞서 말한 patch embeddings들

BERT의 [CLS]토큰과 유사하게 임베딩된 패치들의 맨 앞에 하나의 학습가능한 class 토큰 임베딩 벡터를 추가했다.

이 [CLS]의 임베딩벡터($z^0_0$=$x_{class}$)는 Transformer의 여러 encoder층을 거쳐 최종 output($z^L_0$)으로 나왔을 때, 이미지에 대한 1차원 representation vector로써의 역할을 수행

• L은 마지막 차원임을 의미

pre-training시와 fine-tuning시 둘다, 이 image representation vector 위에 classification head가 부착된다.

classification head는 사전학습 시에는 하나의 은닉층을 가진 MLP로 구현되고, 파인튜닝 시에는 단일 linear layer로 구현된다.

이미지의 위치정보를 유지하기 위해 patch embeddings에 trainable한 position embeddings이 더해진다. 

이미지를 위한 advanced한 2D-aware position embedding을 사용했으나 전통적인 1D position embedding과 비교했을 때, 유의미한 성능향상을 가져다주지 않아 1D 위치임베딩을 사용했다고 한다.

위의 과정을 통해 최종 임베딩 벡터들이 나오고 이 시퀀스가 encoder에 input으로 제공된다.

transformer 인코더는 multi-headed self-attention(MSA) layer들과 MLP 블록들이 교차되어 구성되었다.

• 하나의 이미지도 차원을 쪼개서 multi-head로 관찰하는 모습
• encoder 직후에 mlp를 다는 식으로 쭉 진행

Layer-normalization이 모든 block의 전에 적용되며, residual connection이 모든 블록 이후에 붙는다.

• 여기서도 residual connection이 사용되는 모습

MLP는 GELU(Gaussian Error Linear Unit)를 activation으로 사용하는 2개의 layer를 포함하여 구성된다.

• GELU가 무엇인가?

Inductive bias.

Vision Transformer가 CNN보다 inductive bias 가 크다는 것을 알 수 있었다.

In CNNs, locality, two-dimensional neighborhood structure, and translation equivariance are baked into each layer throughout the whole model.

모든 모델의 CNN의 경우 각 레이어에서 주변부와 함께 들어가는 2차원 구조에서 동일한 변형을 적용시켜줌

• 하나의 layer에 커널을 하나만 찍으므로
• 그렇기에 그 부분 부분에 대한 정보를 종합해서 나타내주는 방식

In ViT, only MLP layers are local and translationally equivariant, while the self-attention layers are global.

하지만 ViT의 경우 patch들로 짤라서 그 부분에 대해 다른 patch를 보았을 때 어떻게 임베딩을 해야 하는지를 그것만 해당해서 배우게 된다.
그것 뿐만이 아니라 scratch 구조에서 모델을 학습시켜야 함

그에 반해 CNN은 이미지의 커널을 찍어가며 각 이미지를 어떻게 임베딩시키면 좋을지를 학습
attention이라는 구조 자체가 데이터가 많을 때에 학습이 잘 되는 구조
왜냐하면 그러한 벡터를 어떤 attention으로 바라보아야 하는가를 여러 번에 걸쳐서 학습하는 것이 좋기 때문이다.
그러한 구조에서 나온 것이 아닌가 싶다

Hybrid Architecture

raw image patches를 대신해 CNN의 feature map으로부터 input sequence가 형성될 수 있다.

• 그니까 cnn의 output에 transformer를 붙이는 방식

이러한 하이브리드 모델에서 식(1)의 patch embedding projection E는 CNN의 feature map으로부터 추출된 패치들에 적용된다.

이는 하나의 특수한 case로, 그 패치들이 1*1의 공간적 크기 - 즉 픽셀 단위를 가질 수 있는데, 이는 input 시퀀스가 단지 feature map의 공간차원들을 flattening 한 뒤, Trasnformer의 차원으로 projecting함으로써 얻어진다는 것을 의미한다.

• 그니까 픽셀 단위로 squeeze해서 Transformer로 넣는 것과 동일
• 하지만 CNN을 활용하기에 spatial 정보도 담길 것

위에서와 동일하게 classification input 임베딩과 position embedding들이 추가된다.

3.2 Fine - Tuning and Higher Resolution

연구진은 ViT를 large dataset들에 pre-train하고 더 작은규모의 downstream task 데이터셋에 fine-tune하였다. 이 과정에서 pre-trained prediction head는 제거되고 0으로 초기화된 D×K의 feed forward layer가 부착되었다. 즉, classification layer는 다시 학습을 시키는 전통적인 fine-tuning 방식을 따랐다.

• 이때 K는 downstream classification task의 class 개수

최근 연구들에서 pre-training시 보다 더 높은 해상도로 fine-tune 하는것이 종종 benefical 하다는 사실이 밝혀졌다.

이에 fine-tuning시 사전훈련 데이터보다 더 높은 해상도에서 진행했으며, 이러한 고해상도 image를 전달할 때 patch size (P×P)는 유지하고, 트랜스포머 유효 input 시퀀스 길이인 N을 늘리는 방향으로 조절했다. ViT는 어떠한 길이의 시퀀스도 처리할 수 있으나 이 경우 pre-trained position embeddings는 의미를 잃게된다.

이를 극복하고자 위치 임베딩이 사전학습에서 학습되고 그것을 fine-tuning에 적용될 때에 그것의 원본 이미지에서의 위치를 기준으로 하는 2D interpolation을 적용하고자 했다

• 즉 위치 임베딩이 fine-tuning시에 변화되기는 하지만 종합적으로 어느 정도의 위치인지 변환시켜줌

위와 같은 해상도 조정 및 patch extraction이 image의 2차원 구조에 대한 inductive bias가 ViT에 수동적으로 주입되는 유일한 과정이다.

• 즉, 2D 위치에 대한 inductive bias가 없는 transformer에게 위치에 대한 개념을 주입시켜주는 것이다.

🔬 4. Experiments

연구진들은 ResNet, ViT, 그리고 하이브리드 모델의 representation learning capabilities를 평가했다. 각 모델이 요구하는 데이터의 양상을 파악하기 위해 다양한 사이즈의 데이터셋으로 사전훈련을 진행했고 많은 벤치마크 테스크에 대해 평가 실시했다.

ViT가 모델의 사전훈련 연산 대비 성능 부분에서 다른 모델들보다 더 낮은 pre-training 비용으로 가장 높은 성능을 달성했다. 마지막으로 self-supervised를 이용한 작은 실험을 수행해 ViT가 self-supervised에서도 가능성이 있음을 보였다.

4.1 Setup

Datasets

뭐 데이터셋 쓴 얘기

Model Variants

BERT에 사용된 뼈대를 base로 삼았으며 'Base'와 'Large'모델은 BERT와 완전히 동일하다. 여기에 더 큰 규모의 'Huge'모델을 추가하였다.

ViT-L/16 → 16*16 patch size를 사용하는 Large variant 모델.

패치 사이즈와 transformer input sequence의 길이는 반비례하므로 패치 크기가 작을수록 계산 코스트가 증가

비교 대상을 위한 baseline CNN으로 ResNet을 사용하였는데, 오리지날 ResNet의 Batch Normalization을 Group Normalization으로 대체하였고 standardized convolution을 사용하였다.

• Group Normalization이란?

group normalization - Google Search

이러한 작은 변형이 전이학습의 성능을 더 향상시키며, 이 modified model을 ResNet(BiT)로 표기한다.

하이브리드 모델에서는 CNN의 중간 feature map들을 ViT에 1*1 패치로 나누어서 넣었다.

• 그렇게 되면 들어가는 sequence가 매우 길 것

Training & Fine-tuning

pre-training에는 모든 모델을 Adam optimizer를 활용했으며 하이퍼파라미터 세팅은 $β_1$= 0.9, $β_2$=0.999, batch_size = 4096로 두었고, weight decay를 적용하였는데 이것이 모든 모델에 대해 transfer시 유용한 도움을 주는것을 발견하였다.

학습률 스케쥴로는 linear learning rate warmup and decay를 사용하였다.

한편 fine-tuning에는 SGD with momentum 옵티마이저를 사용했고, batch_size = 512으로 두고 훈련하였다.

Metrics

연구진은 성능을 평가하기 위해 downstream 데이터셋에 대한 결과를 few-shot accuracy 혹은 fine-tuning accuracy로 성능 비교를 진행했다. 이때 Fine-tuning acc는 각 모델의 각 데이터셋에 대한 성능을 반영한다고 볼 수 있다.

Few-shot accuracy는 훈련이미지들의 subset의 representation을 $\{−1,1\}^K$의 타겟벡터로 매핑하는 규제된 linear regression 문제의 해를 구하는 과정에서 얻어진다.

• 흠 왜 {-1, 1}일까
• 0과 1이 아닌 것은 조금 의외이다.

주로 fine-tuning 성능으로 비교에 집중했으나 fine-tuning이 너무 고비용인 경우 보다 가볍게 평가할 수 있도록 few-shot accuracy를 사용했다.

4.2 Comparison to SoTA

저자들은 먼저 ViT variants 모델들 중 가장 큰 size의 ViT-H/14와 ViT-L/16을 SOTA 성능의 CNN들과 비교했다.

첫 비교대상은 대규모의 ResNet으로 지도방식의 전이학습을 수행시킨 Big Transfer(BiT)모델이고, 두 번째 대상은 큰 EfficientNet으로 self-supervised learning을 이용해 학습시킨 Noisy Student이다.

논문에 따르면 ImageNet에 대해서는 Noisy Student가 SOTA고, 이외 데이터셋들에 대해서는 BiT-L이 SOTA이었다고 한다.

위 그림이 비교 결과이며 동일한 데이터셋에 대해 pre-train을 했을 때, 모든 downstream task에서 ViT-L/16의 성능이 BiT-L을 능가하면서 동시에 사전훈련에 드는 계산상의 cost는 확연하게 낮음을 확인할 수 있다.

더 큰 모델인 ViT-H/14은 모든 실험 데이터셋에 대해 더욱 향상된 성능을 보여준다. 특히 더욱 어려운 데이터셋인 ImageNet, CIFAR-100, VTAB 세트 등에서 두드러진다.

흥미롭게도 이 Huge한 모델도 이전의 SOTA들과 비교했을 때 더 낮은 pre-train 비용이 요구된다. 그러나 연구진들은 이러한 비용 차이에 사전학습의 효율성이 데이터 구조의 선택 뿐 아니라 training schedule이나 optimizer, weight decay등의 다른 파라미터들에도 영향을 받는다는 것을 고려했기에 이 부분에 대한 것은 따로 4.4 부분에서 다루었다.

• 실제로 compute cost가 낮은 것 맞는가에 대한 연구

마지막으로 ViT-L/16을 JFT보다 작은 규모의 데이터셋인 ImageNet-21k에 pre-train 하였을 때에도 대부분에 데이터 셋에 대해 상당히 좋은 성능을 보였다.

4.3 Pre-Training Data Requirements

ViT는 대규모의 JFT-300M 데이터셋에 사전학습될 때 기존의 SotA ResNet들보다 더 적은 inductive biases들을 가지고도 더 잘 작동한다. 그렇다면 데이터셋의 크기가 얼마나 큰 영향이 있는지를 연구했다.

먼처 첫 번째로 ImageNet → ImageNet-21k → JFT-300M의 흐름으로 데이터셋의 크기를 증가시키며 사전학습을 진행했다. 작은 데이터셋에도 가능한 높은 성능을 얻게하기 위해 weight decay, dropout & label smoothing와 같은 규제를 적용했다.

위 그림은 각 모델들을 각 데이터셋에 pre-train시킨 뒤 ImageNet에 fine-tuning한 성능측정 결과이다. 가장 작은 사이즈의 표준 ImageNet에 사전학습 시켰을 때에 강한 규제를 가했음에도 오히려 ViT-Large가 ViT-Base보다 성능이 낮은 모습을 보였다.

조금 더 큰 규모의 ImageNet-21k에 사전훈련 시킨 경우에도 크게 다르지 못함

오직 JFT-300M(가장 대규모)으로 pre-train 했을 때 큰 모델의 완전한 이점을 확인할 수 있었다.

위 그림에서 색칠된 회색 영역은 BiT-50과 152의 성능 구간이다.

이때에 작은 데이터셋으로 사전훈련했을 때에는 BiT가 앞서지만 사전훈련 데이터셋 규모가 커질수록 ViT가 BiT를 이긴다

두 번째 실험으로써, 저자들은 full JFT-300M 데이터셋 뿐만아니라 그것을 랜덤하게 선택해 9M, 30M, 90M 크기의 데이터로 훈련시켰다.

작은 부분집합에 훈련시킬 때에도 기본 JFT에 대해 훈련시 사용한 하이퍼파라미터들을 그대로 사용했다. 이를 통해 규제의 효과가 아닌 모델의 본질적인 성질들을 평가하고자 했다.

해당 실험에서는 효율적으로 평가하기 위해 그냥 accuracy를 보는 대신 few-shot linear accuracy로 모델 간의 비교를 진행했다.

결과는 위 그래프와 같으며 상대적으로 작은 데이터셋에 대해서 사전훈련시에는 동일하게 BiT가 ViT를 앞었다.

ViT는 이런 소규모 데이터셋에 대해 BiT보다 더 많이 overfit되었다.

예를들면 ResNet 50보다 ViT-B/32가 살짝 더 빨랐지만, 9M보다 작은 데이터셋에 대해서 성능이 더 낮았으며, 90M보다 큰 데이터셋에 대해서는 성능이 더 나았다.

이 양상은 ResNet 152 vs ViT-L/16의 비교에서도 동일하게 나타났다. 이 결과는 'convolutional inductive bias'(지역을 바탕으로 예측하는 방향)이 상대적으로 작은 데이터셋에 대해서 useful하지만, 큰 데이터셋에 대해서는 단지 적절한 패턴을 학습하는것 만으로 충분하거나 심지어 더 이롭다는 직관을 공고하게 만들어주었다.

위의 성능표를 보았을 때 VTAB을 사용했을 때 보이듯이 더 많은 데이터셋에 대해 매우 효과적임을 알 수 있다.

ViT의 few-shot 특성의 추가적인 분석은 하나의 흥미로운 향후 연구방향이 될 것 같다.

4.4 Scaling Study

JFT-300M데이터셋으로부터의 전이학습 성능평가를 수행함에 있어 각 모델의 performance 대비 pre-training cost를 평가했다.

ResNet 7개와 ViT 6개, 하이브리드모델 5개의 모델들을 고려하였다.

위 그래프로 나타난 패턴은 첫 번째로 모든 ViT가 ResNet 모델들을 성능/계산코스트 trade-off에서 압도한다는 것이다. 동일한 성능을 달성하기위해 드는 계산비용이 ViT가 2 ~ 4배는 더 적다.

두 번째로 하이브리드 모델이 비교적 작을 때에는 ViT의 성능을 능가한다. 이 차이는 모델 사이즈가 커짐에 따라 점차 소실된다. 이는 어떤 사이즈의 ViT에도 convoluational local feature processing이 성능 향상에 도움을 줄 수 있다는 것을 암시하기에 의미가 있다.

마지막 세 번째로 ViT는 실험을 시도해본 range 내에서 saturated 되지 않았기에, 미래의 feature scaling efforts를 고무시킨다.

4.5 Inspecting Vision Transformer

ViT가 이미지를 어떻게 처리하는지 이해하기 위해 내부의 representation들을 분석했다.

ViT의 첫 번째 층은 flatten된 patch들을 저차원 공간으로 linearly project한다. 위 그림의 왼쪽 사진이 해당 층에서 학습된 임베딩 필터들의 상위 주 성분들을 보여주고있다.

• 상위 주 성분..?
• 이미지 patch를 flatten시키고 2차원으로 projection을 한 후에 그 행렬에 대한 주 성분을 분석한 것

이러한 성분들은 각 패치 내의 미세 구조를 저 차원으로 표현하기 위한 그럴듯한 기저함수를 닮았다고 볼 수 있다.

• 그러니까 이러한 요소들이 닮았다고 하는 것

projection이후에는 학습된 position embedding이 patch representation에 추가된다.

위 그림의 가운데 그림은 모델이 pos-embedding의 유사도를 통해 이미지 내부의 거리 개념을 인코딩하는 방법을 학습하고 있다는 것을 보여준다. 즉, 같은 행/열에 있는 patch는 유사한 임베딩을 갖는다.

또한 pos-encoding을 더하고 안 더하느냐에 따라 모델의 성능차이가 큰데 1d pos embedding과 2d pos embedding의 차이는 거의 없으며 오히려 1d가 더 낫다.

self-attention은 ViT가 이미지 전체의 정보를 낮은 layer에서도 통합하여 사용할 수 있게 한다. 연구진은 네트워크가 이 광활한 수용력을 얼마나 이용하는지 그 정도를 조사해보았다고 한다.

구체적으로는 정보가 attention weights에 의해 통합되는 image space내의 평균거리를 계산했다. 그것의 결과가 위 그림에서 오른쪽이다.

• 여기서 이 "어텐션 거리"는 CNN에서의 receptive field size와 유사한 개념이라고 보면된다.

• 즉 내가 해당 부분을 고려할 때에 다른 것들을 얼마나 고려하는지

그림을 보면 우선 층이 깊어질수록 어텐션거리가 증가한다. 그러면서도 가장 아래 layer에서도 몇몇 attention head가 이미지 대부분에 attend하고있는 것을 확인할 수 있는데, 이는 정보를 global하게 integrate할 수 있는 능력을 모델이 실제로 사용하고 있음을 보여준다.

• 즉 특정 부분은 전체적으로 보고 일정 부분은 특정 요소로만 파악하고

또한, 하위층에서 다른 몇몇 어텐션 헤드들은 작은 어텐션거리를 갖는다.

하지만 이 고도로 local한 어텐션 능력은 hybrid 모델에서는 덜 나타난다.

이는 이 local한 어텐션이 CNN의 초기 convolutional layer들과 유사한 기능을 수행하기 때문

• 그니까, CNN = low layer Transformer

그리고 attention 길이는 depth가 깊어질 수록 증가하게 된다.

• 층이 높아지면서 더 넓은 데이터를 이해하는 것

또한 저자들은 ViT가 실제로 분류에 의미적으로 관련있는 영역에 attend한다는 것을 발견하였다.(Figure6)

또한 모델이 image 내부의 요소에 따라 그 부분을 분리해서 파악한다고 한다.

• 내용은 Appendix D.7 참고

4.6 Self-supervision

트랜스포머기반 모델들은 NLP task들에서 인상깊은 퍼포먼스를 보여주었다.

그러나 그것들의 많은 성공은 뛰어난 확장성 뿐만 아니라 대규모의 자기지도 사전훈련(self-supervised pre-training)에서 기인한다.

• 그니까 MLM이나 NSP 같은 것들

연구진들은 BERT에서 사용된 masked language modeling task를 모방하여 masked patch prediction for self-supervision를 실험해보았고, 그 결과는 scratch로부터 학습시키는 것보다 유의미한 성능향상을 가져다주었다.

그러나 여전히 지도방식의 사전훈련(supervised pre-training)에는 많이 못미치는 성능이었다.

이러한 self-supervised pre-training에 대한 탐험은 미래의 연구로 남겨두고있다.

🧪 Appendix

Appendix는 모든 부분은 정리하지 않았고 공부를 하며 정리하면 좋을 것이라 판단되었던 부분들만 선별해 정리했다.

B.1.2 Self-Supervision

MLM처럼 이미지에도 적용하고자 함

50% 확률로 patch 임베딩이 변화하는데 80%로는 [Mask] 토큰, 10%로는 다른 임베딩, 10%로는 자신의 것을 유지하도록 바꿈

• MLM과 완전히 동일

그리고 그것을 예측하는데 있어 3bit 512 color에 대한 평균 색깔을 예측하도록 했다.

JFT에 대해서 14epochs로 1M steps 정도 학습을 실시

예측은 아래 세가지 종류를 통해서 이루어졌다.

1) 3bit 512 color에 대한 평균 색깔에 대한 예측만을 실시하도록 했다.

2) 4 * 4 크기에 대해 예측하도록 했고 16개의 부분에 대해 모두 예측을 하는 것

• 이때에도 512 color에 대한 예측 실시

3) 3가지 색깔에 대한 0 ~ 255의 regression을 실시

위의 3가지 모두에 잘 작동했으나 L2가 가장 성능이 낮았고 1번이 가장 성능이 높았다.

최근에 나온 논문 중에 MLM을 15%로 적용하는 방법도 시도해보았으나 성능이 좋게 나오지는 못했다.

마지막으로 연구진은 위와 같이 많은 데이터와 오랜 시간의 pre-train이 필수적이지 않다고 말했다.

100K로 학습을 줄이고 데이터도 줄였음에도 거의 비슷한 성능을 보였다.

D.2 Transformer Shape

여러가지 Transformer Dimension에 대한 실험을 실시

위의 그림을 통해 결과를 표현

모든 모델은 8개의 layer를 가진 ViT 모델이며 D = 1024와 D = 2024에 대한 그림이고 patch size는 둘이 동일하게 32를 실시했다.

결과를 보았을 때 Depth가 가장 큰 효과가 있었다.

하지만 16개의 layer부터는 성능의 향상이 감소하기 시작했다.

MLP의 width를 늘려주는 것의 효과가 가장 적어 보였고 patch size를 줄여주는 것은 파라미터수를 증가시키지 않고 input으로 들어가는 것들에 대한 길이만 증가했을 때 robust한 효과 상승을 보였다.

이러한 결과는 파라미터를 증가시키는 것보다는 patch size를 일정 수준 작게 해주는 것이 성능을 늘리는데 효과가 있으며 당연하게도 depth는 가장 큰 효과가 있다.

전반적으로 일단 이러한 것들이 모두 성능 향상에 도움이 도기는 했다.

D.3 Head Type and Class Token

NLP에서 사용하는 것처럼 average-pooling(mean pooling)을 사용하고자 함. GAP이라고 불림

Resnet에 대해서는 이러한 적용이 성능을 매우 떨어뜨렸다.

하지만 이는 CLS토큰에서도 마찬가지로 나타나기 때문에 token을 추가하는 것에 대해 성능 영향이 있다고 판단

다른 것에 대한 그래프는 아래와 같다.

• 최종적으로 epoch이 증가함에 따라 CLS토큰이 가장 좋은 성능을 보였다.

D.4 Positional Embedding

여러 가지 종류의 pos-embedding에 대해서도 실험을 진행

1) pos embedding을 제공하지 않기

2) 2D에 대한 embedding 제공

• X와 Y에 대한 임베딩이 다르게 학습된다. 그리고 input의 pathdp akwcndjtj X와 Y의 임베딩을 concat을 실시해준다.

3) 상대적 pos embedding 방식

• patch에 대해서 각 거리마다 그것의 embedding을 다르게 설정한다. 이를 위해 1-D 상대적 attention을 실시한다.
• 따라서 모든 pair에 대해 query + key/value 형식으로 주어지고 $p_q - p_k$를 계산할 수 있게 된다.
• 이를 통해 key에 대한 query patch의 extra attention을 실시하고 bias로 사용할 logit 값을 구하고 softmax에 적용하기 전에 main attention에 더해준다.
• 그니까 멀면 멀수록 attention을 낮춰주는 방식 등으로 사용

위의 변형과 동시에 3가지의 방식에 대해서도 연구를 더 진행했다.

1) input이 들어올 때에 encoder에 들어오기 이전에 바로 더해주는 방식

2) 각 레이어의 첫번째로 pos-embedding에 대해 학습하는 방식

3) 각 레이어에 pos-embedding이 있지만 모든 layer에 공유되는 방식으로 학습되는 방식

pos-embedding이 없는 것에 비해서는 효과 차이가 크지만 다른 것들에 대해서는 차이가 미미하다.

연구진의 transformer가 patch를 base로 동작하기 때문에 이러한 것이 크게 중요하지 않다는 것

• 14 * 14에 대한 pos-embedding을 학습하는 것이므로 어려운 과제가 아니기 때문

따라서 hyperparameter에 따라 약간씩의 차이가 있을 뿐 큰 차이가 존재하지는 않았다.

D.6 Axial Attention

Axial Attention은 multi dimensional tensor를 다루기에 단순하지만 매우 유용한 방법

이는 flatten을 시켜서 1-D attention을 수행하는 방식이 아니라 single axis마다 attention을 실시해 multiple attention을 수행하는 방식이다.

다른 axis는 독립적으로 유지하면서 axial attention을 실시해줄 때마다 특정 두 개의 attention을 결합시키는 방식이다.

연구진은 이를 적용해 ViT를 구성해보고자 했다. 이를 위해 ViT가 2차원의 데이터를 input으로 받도록 구성해주었고 axial transformer block을 구성해주었으며 이는 row-self-attention을 base로 column-self-attention이 동작하도록 구성했다.

위의 그래프로 결과가 나타나는데 axial-ViT가 성능이 더 높지만 compute cost가 더 높다.

이는 axial-ViT의 경우 MLP가 적용이 될 때에 2-D shape에 대한 attention이 진행되어야 하기 때문이다.

axial-ResNet50이 왼쪽 그래프에서는 매우 reasonable 해보이지만 오른쪽 그래프에 따르면 axial-resnet은 speed만 빠르고 성능이 낮기에 trade-off가 너무 강하다

✏️ 후기

생각보다 내용이 없다고 느낄 수도 있지만 axial attention, 2D-pos embedding 등 정말 많은 내용을 담고 있어 논문 자체를 완벽하게 이해하지는 못한 것 같다. 따로 이 부분에 대한 공부도 진행해야 할 것 같다. NLP에서는 단일 sequence만 다루다보니 2차원적인 구조에서의 attention에 대해서는 고려해보지 못했는데 이를 NLP에서도 활용할 방안이 있을까라는 생각도 하게 되었다. 여러 모로 생각할 지점들을 주었던 논문이었다.