[Paper Review] An Image is Worth 16X16 Words : Transformers for Image Recognition at Scale

 

An Image is Worth 16X16 Words : Transformers for Image Recognition at Scale

 

위 논문은 자연어 처리(NLP) 분야에서 널리 사용되는 Transformer 개념을 컴퓨터 비전 분야에 적용해본 논문이다. 먼저 자연어 처리 분야에서 Transformer 개념은 어떤 것인지 살펴보는게 좋을 것 같다.

 

 

NLP 에서의 Transformer

자연어 처리 분야에서는 문장과 같은 연결성이 중요한 시퀀스 데이터에서 RNN 계열의 모델이 많이 사용되는데 이는 직전의 출력 결과를 입력으로 사용하여 재귀적으로 활용하게 된다. 하지만 이는 좀 더 멀리 떨어진 단어 보다는 무조건 가까운 단어가 연관성이 높게 나타난다는 단점이 있다.

 

이러한 문제를 Long-term Dependency 문제라고 하는데 이를 해결하기 위해 제안된 모델이 LSTM(Long Short Term Memory)라는 개념이다. 이는 새로운 입력값 뿐만 아니라 이전의 출력값을 잃어버릴지, 얼마나 유지할지를 고려하는 게이트(gate)들을 추가하여 상대적으로 멀리 있는 단어에 영향을 전달 할 수 있도록 한다. 하지만 이 또한 거리에 대한 한계가 존재했고 순차적으로 연산한다는 점에서 병렬처리에 어려움이 있어 연산량이 너무 많아 학습 속도가 느리다는 단점이 있었다.

 

이러한 문제를 해결하기 위해 등장한 것이 어텐션(Attention)이다. 자연어 처리 문제에서 기계 번역(Machine Translation)에 대한 예제로 이는 인코더(Encoder)-디코더(Decoder) 구조로 이루어져있으며, 이 구조에서 인코더는 입력 시퀀스 하나를 벡터 표현으로 압축하고, 디코더는 이 벡터 표현을 통해서 출력 시퀀스를 만들어냈다.

 

하지만 이러한 구조는 인코더가 입력 시퀀스를 하나의 벡터로 압축하는 과정에서 입력 시퀀스의 정보가 일부 손실 된다는 단점이 있었고, 이를 해결하기 위해 어텐션 개념이 사용되곤 했다. 이와 같은 어텐션으로 인코더와 디코더로 구현한 것이 바로 Transformer 이다.

 

 

Transformer 는 인코더에서 입력 시퀀스를 입력 받고, 디코더에서 출력 시퀀스를 출력하는 인코더-디코더 구조로 되어있는데 이전과 다른점은 인코더와 디코더라는 단위가 N개가 존재할 수 있다는 것이다. 논문에서는 6개를 쌓았다고 하고 이 갯수는 얼마든지 조절이 가능하다.

 

 

인코더 구조는 하나의 인코더가 Self-Attention 및 Feed Forward 인 두 개의 sub-layer 이루어져 있으며, 그들 간에 같은 가중치(weight)를 공유하진 않는다. 인코더에 들어온 입력은 일단 self-attention layer 를 지나게 되는데 이 인코더는 하나의 특정한 단어를 인코딩 하기 위해 입력 내의 모든 다른 단어들과의 관계를 살펴보게 된다. 이를 통과한 출력은 다시 feed-forward 신경망으로 들어가게 되고 똑같은 feed-forward 신경망이 각 위치의 단어마다 독립적으로 적영되어 출력을 만든다. 디코더 구조 또한 두 층 사이에 인코더-디코더 어텐션이 포함되어 있는데 이는 디코더가 입력 문장 중에서 각 타임 스텝에서 가장 관련 있는 부분에 집중 할 수 있도록 해준다. 

 

 

 

 

 

 

 

Abstract

Transformer 구조는 이미 자연어 처리 분야에서 널리 사용되고 있는 개념이다. 컴퓨터 비전 분야에는 아직 적용된 사례가 거의 없는데, 원래 컴퓨터 비전에서 어텐션 개념은 CNN 구조에 주로 적용되어 전체 구조를 유지하되 bottleneck 구조에서 어텐션을 하는 방식으로 진행되었었다. 본 논문에서는 이렇게 특정 구성 요소를 대체하는 것에 그치지 않고 이미지 패치의 시퀀스에 Transformer 가 적용 될 때 이미지 분류 작업에서 매우 우수한 성능을 보여준다는 것을 발견하였다. 많은 양의 데이터에 대해 사전 훈련(pre-trained) 되어 여러 벤치 마크(ImageNet, CIFAR-100, ...)에서 SOTA 보다 훨씬 적은 계산량으로 우수한 결과를 얻을 수 있다. 

 

 

 

1. Introduction 

NLP 분야에서 Transformer 는 large text corpus 에서 사전 훈련을 수행하고 task-specific dataset 에 대해 fine-tuning 을 수행하는 것이다. Transormer 의 계산 효율성 및 확장성으로 100B 이상의 파라미터(parameter)를 사용하여 큰 모델을 학습 할 수 있게 되었다. 그러나 컴퓨터 비전 분야에서는 CNN 구조가 여전히 많이 사용되고 있으며 여러 연구에서 CNN과 유사한 구조를 self-attention 과 결합하려고 시도해왔다.

 

본 논문에서는 NLP의 Transformer 에서 모델 수정을 최소화 하면서 표준 Transformer 를 이미지에 직접 적용해보는 실험을 시도하였다. 이를 위해 이미지를 패치로 분할하고 이러한 패치의 선형 임베딩 시퀀스(linear embedding sequence)를 Transformer 에 대한 입력으로 설정한다(feed). 이미지 패치는 NLP 에서 token(word)와 동일한 방식으로 처리된다. 이 때 지도 학습(uperviesed fashion)으로 이미지 분류에 대한 모델을 학습한다. 이러한 모델은 중간 규모(mid-sized)의 ImageNet 과 같은 데이터세트에서 학습할 때 적을한 결과를 산출하여 비슷한 크기의 ResNet 보다 몇 퍼센트 정도 낮은 정확도를 달성한다. 이는 equivariance 및 locality 즉 CNN 고유의 inductive bias 를 고려할 수 있는 기능이 transformer 에 없기 때문에 불충분한 양의 데이터에 대해 학습할 때 일반화가 잘 안된다는 문제가 있었다. 

 

그러나 large-scale 데이터 세트에서 모델을 학습하면 위와는 다른 결과를 내뿜는다. 엄청 대규모의 데이터 세트로 학습하게 된다면 inductive bias 를 능가 할 수 있다는 결론이 나온다. 즉, 충분한 규모로 사전 학습되고 더 적은 데이터로 fine tuning 할 때 좋은 결과를 얻을 수 있다는 것이다. 각 유명한 벤치마크에서 ImageNet 에서 88.36%, CIFAR-100 에서 94.55% 등의 성능으로 SOTA 에 접근하거나 이를 능가하는 결과를 확인하였다. 

 

(Related Work 는 생략한다...)

 

 

 

2. Method

2.1 Vision Transformer(ViT)

 

Vision Transformer

 

컴퓨터 비전 분야에서 Transformer 는 token embedding 의 1D 시퀀스를 입력으로 받으며 이미지 패치들을 처리하기 위해 아래와 같이 H x W x C 로 이루어진 이미지 $x\in \mathbb{R}^{H \times W\times C} $ 를 flatten 된 2D 이미지 패치인 시퀀스 $x_{p}\in \mathbb{R}^{N \times (P^{2}\cdot C)} $ 로 재구성 한다. 

 

(H, W)는 원본 이미지의 해상도(resolution)였다면 (P, P)는 각 이미지 패치의 해상도이다. $ N = HW / P^{2} $ 는 Transformer 의 시퀀스 길이를 뜻한다. Transformer 는 모든 레이어를 일정한 width 의 길이로 사용하기 때문에 학습 가능한 linear projection 은 각 vectorized patch 를 dimension D 에 매핑하며 식 1과 같이 그 결과를 "patch embedding" 이라고 한다. 

 

BERT의 [class] token 과 유사하게 Transformer 인코더의 출력 상태(output state)인 $(Z_{0}^{L})$가 image representation y 로 사용되는 embedded patch 의 시퀀스 $ (z_{0}^{0}= x_{class}) $ 앞에 학습 가능한 embedding 을 추가한다. 이는 식 4에 나타나있으며, pre-train 및 fine-tuning 중에 classification head 는 $(Z_{L}^{0})$ 에 추가된다. 

 

 

position embedding 은 위치 정보를 유지하기 위해 patch embedding 에 더해진다. Transformer 인코더는 Multi-headed self-attention 및 MLP block 으로 구성된다. Layernorm(LN)은 모든 block 이전에 적용되고 residual connection 은 모든 block 이후에 적용된다. 

 

Hybrid Architecture

이미지를 패치로 분할하는 대신 ResNet 의 intermediate feature map 에서 입력 시퀀스를 형성할 수 있다. 이 하이브리드 모델에서 patch embedding projection E (Eq. 1)는 ResNet 의 early stage 로 대체된다. ResNet 의 intermediate 2D feature map 중 하나는 시퀀스로 flat 되고 Transformer 차원에 projection 된 다음 Transformer 에 입력 시퀀스로 공급된다. classification input embedding 및 position embedding 은 위에서 설명한대로 Transformer 에 대한 입력으로 추가된다. 

 

 

 

2.2 Fine-tuning and Higher Resolution

일반적으로 대규모 데이터 세트에 대해 ViT 를 pre-train 하고 downstream task 를 fine-tuning 한다. 이를 위해 pre-train 된 prediction head 를 제거하고 0으로 초기화 된 D x K feedforward layer 를 연결한다. 여기서 K 는 downstream class 의 수를 뜻한다. pre-train 보다 높은 해상도로 fine-tuning 을 하는 것이 성능에 도움이 된다. 

 

더 높은 해상도의 이미지가 주어질 때 패치 크기를 동일하게 유지하게 되므로 유효한 시퀀스 길이가 더 커지게 된다. Vision Transformer 는 임의의 시퀀스 길이(up to memory constraint)를 처리 할 수 있지만 pre-train 된 position embedding 은 더 이상 의미가 없기 때문에 원본 이미지에서의 position 에 따라 pre-train 된 position embedding의 2D interpolation 을 수행한다. 이러한 해상도 조정 및 패치 추출은 이미지의 2D 구조에 대한 inductive bias 가 Vision Transformer 에 수동으로 적용되는 유일한 지점이라고 할 수 있다. 

 

 

 

3. Experiments

 

ResNet, Vision Transformer(ViT) 및 하이브리드의 represiontation learning capability 를 평가하였다. 각 모델의 데이터 요구사항을 이해하기위해 다양한 크기의 데이터 세트를 사전 학습하고 많은 벤치 마크에서 작업을 수행하였다. 모델을 pre-train 하는 계산 비용을 고려해볼 때 ViT 는 더 낮은 pre-train 비용으로 대부분의 recognition 벤치마크에서 SOTA 를 달성하였다.

 

 

3.1 Setup 

Datasets

pre-train 을 위한 데이터 세트는 ILSVRC-2012 ImageNet 의 -k 클래스 및 1.3M 이미지를 사용하였고 ImageNet-21k 및 14M 이미지를 사용하였으며 JFT 의 18k 클래스 및 303M 이미지를 사용하였다. 또한 Transfer Learning 을 위한 데이터세트로는 ImageNet 및 ReaL labels, CIFAR 10/10, Oxford-IIIT Pets, Oxford Flowers-102, 19-task VTAB classification suite 를 사용하였다. 

 

Model Variants

ViT 구성은 BERT 에 사용되는 구성을 기반으로 한다. ViT-L / 16 모델은 16 x 16 입력 패치 크기의 Large Transformer 라는 것을 의미한다. 참고로 Transformer 의 시퀀스 길이는 $ image\: resolution^{2} / patch\: size^{2} $ 이며, B 는 Base, L은 Large, H는 Huge 를 뜻한다. 

baseline 인 CNN 의 경우 ResNet 을 사용하지만 Batch Norm 레이어를 Group Norm 으로 바꾸고 표준화 된 컨볼루션을 사용하였다. 이러한 수정은 Transfer 를 개선할 수 있으며 이러한 수정된 모델은 ResNet(BiT)라고 한다. 하이브리드의 경우 패치 크기가 픽셀인 ViT 에 intermediate feature map 을 제공하게 된다. 다른 시퀀스 길이를 비교하기 위해 (i) 에서는 R50 의 output of state 4 를 가져오고, (ii) 에서는 output of stage 4 를 제거하고 stage 3 에서 동일한 수의 레이어를 배치한다. (50개) 이는 시퀀스의 길이가 길고 ViT 헤드가 더 expensive 하다. 

 

 

Training & Fine-tuning

 

Pre-triain

• Adam optimizer : $ \beta _{1}=0.9 , \:  \beta _{21}=0.999, \: batchsize = 4096 $
• weight decay : 0.1

Fine-turning

• SGD with momentum, $ batchsize = 512 $
• linear learning rate warmup, decay
• Higher resolution :512 for ViT-L/16 , 518 for ViT-H/14

 

Metrics

few-shot 및 fine-tuning 정확도를 통해 downstream 데이터세트에 대한 결과를 산출하였다. fine-tuning 정확도는 각 데이터 세트에서 fine-tuning 을 수행한 후 각 모델의 성능을 측정하였다. 학습 이미지 하위 집합의 representation 을 $ \left \{ -1, 1 \right \}^{2}  $ target vector 에 매핑하는 정규화된 linear regression 문제를 해결하여 few-shot 정확도를 얻는다. 성능을 fine-tuning 하는데 중점을 두지만 이에 너무 많은 비용이 드는 빠른 실시간 평가를 위해 linear few-shot 정확도를 사용하기도 한다.

 

 

 

3.2 Comparison to State of the Arts

모든 모델은 TPUv3 에서 학습되었으며 pre-train 하는데 소요된 일수를 확인 할 수 있다. ViT-L/16 모델은 모든 데이터세트에 대해 BiT-L 과 일치하거나 더 좋은 성능을 보여줌과 동시에 훨씬 적은 computational resource 를 필요로 한다. 더 큰 모델은 ViT-H/14 는 ImageNet 및 CIFAR-100 과 VTAB 에서 성능이 더욱 향상되었다. 

 

 

3.3 Pre-training Data Requirements

ViT 는 대규모 사이즈의 데이터스테인 JFT-300M 에 대해 pre-train 하였을 때 더 좋은 성능을 보여준다. ResNet 보다 vision 에 대한 inductive bias 가 적을 때 데이터 세트의 크기가 얼마나 중요한지에 대한 실험을 수행하였다. 가장 작은 데이터 세트에 대해 pre-train 하였을 때 ViT-Large 모델은 정규화에도 불구하고 ViT-Base 보다 성능이 떨어진다. 그러나 ImageNet-21k 를 사용하면 성능이 비슷한 것을 볼 수 있다. JFT-300M 데이터세트에는 ViT 가 BiT 보다 더 좋은 성능을 보여준다는 것을 확인 할 수 있다. 결론은 더 작은 규모의 데이터세트의 경우 convolutional inductive bias 가 유리하지만 더 큰 대규모의 데이터세트의 경우 CNN 과 같은 패턴을 학습하는데 충분하다는 것을 알 수 있다. 

 

 

3.4 Scaling Study

JFT-300M 데이터세트에서 transfer 성능에 대해 다양한 모델로 확장된 연구를 수행한 결과는 다음 그림 5와 같다. 데이터세트의 크기는 모델 성능에 병목 현상이 없으며 각 모델의 pre-train 비용 대비 성능을 평가하였다. 일단 ViT 성능 및 계산량의 trade-off 에서 ResNet 보다 효율적이며 ViT 는 동일한 성능을 달성하기 위해 약 2배 정도 적은 컴퓨팅을 사용한다. 또한 하이브리드는 적은 컴퓨팅 비용에서 ViT 를 능가하지만 더 큰 컴퓨팅 비용에서는 차이가 사라지게 된다. ViT 는 시도된 범위 내에서 포화(saturate)되지 않는 것으로 보이며 향후 확장 가능성이 있다. 

 

 

3.5 Inspecting Vision Transformer

ViT 가 이미지를 처리하는 방법을 이해하기 위해 분석을 수행하였다. ViT 의 첫번째 레이어는 flatten patch 를 더 낮은 차원 공간에 projection 한다. 위 그림 7에서 왼쪽은 학습된 embedding filter 의 구성 요소들을 보여준다. 구성요소는 각 patch 내 미세 구조의 low-dimensional representation 에서 basic function 과 유사하다. projection 된 이후 학습된 position embedding 이 patch representation 에 추가된다.

 

가운데 그림은 모델이 position embedding 의 유사성에서 이미지 내 거리를 인코딩 하는 방법을 학습한다는 것을 보여준다. 더 가까운 patch 는 더 유사한 position embedding 을 갖는 경향이 있으며 행-열 구조가 나타난다. Self-attention 을 통해 ViT 는 가장 낮은 레이어에서도 전체 이미지에 대한 정보를 통합할 수 있다.

 

Self-attention 의 weight 를 기반으로 정보가 통합되는 이미지 공간의 평균 거리를 오른쪽 그림과 같이 계산 할 수 있다. attention distance 는 CNN의 receptive field size 와 유사하다. 일부 head 에서는 최하위 레이어에 있는 대부분의 이미지에 attention 을 적용하여 global 하게 모델을 사용할 수 있다는 것을 보여준다. 또한 신경망 깊이에 따라 attention distance 가 증가한다. 

 

아래 그림을 통해 output token 에서 입력 공간에 대한 attention 의 예시를 확인 할 수 있다. 

 

 

4. Conclusion 

이미지 인식 분야에서 Transformer 를 직접 적용하는 방법을 제안하였다. 컴퓨터 비전에서 Self-attention 을 사용하는 이전 연구들과는 달리 본 논문에서는 구조에 image-specific inductive bias 를 사용하지 않았다. 대신 이미지를 patch 로 해석하고 NLP에서 사용되는 standard transformer encoder 로 처리한다. 간단하면서도 확장성 있는 이러한 방법은 대규모의 데이터 세트에 대해 pre-train 될 때 잘 작동한다. 따라서 ViT 는 많은 이미지 분류 데이터세트에서 SOTA 를 능가하며, pre-train 비용이 굉장히 저렴하다는 장점을 가진다. 하지만 많은 challenge 가 다음과 같이 남아있다.

 

• detection 및 segmentation 분야에도 ViT 를 적용하는 것
• pre-training method 에 대해 연구하는 것
• 모델의 크기가 증가해도 포화(saturate) 상태가 아닌 것으로 보이기 때문에 ViT 를 확장할 필요성이 있다는 것

 

Attention Map 예시 

 

 

 

 

참고자료 1 : arxiv.org/abs/2010.11929

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

 

참고자료 2 : nlpinkorean.github.io/illustrated-transformer/

The Illustrated Transformer

 

참고자료 3 : jeonsworld.github.io/vision/vit/

[논문리뷰] An Image is Worth 16X16 Words: Transformers for Image Recognition at Scale