상세 컨텐츠

본문 제목

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

Paper Review

by SONG, Ph.D 2021. 9. 24. 17:08

본문

 

 

https://arxiv.org/abs/2010.11929

 

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

While the Transformer architecture has become the de-facto standard for natural language processing tasks, its applications to computer vision remain limited. In vision, attention is either applied in conjunction with convolutional networks, or used to rep

arxiv.org

 

첫번째로 리뷰 할 논문은, 2020년에 Google Research, Brain Team에서 공개한 논문이다.

 

자연어 처리(Natural Language Processing, NLP) 분야에서 널리 사용되고 있는 "Transformer" 개념을 Computer Vision 분야에 적용한 논문이다. 소위 Vision Transformer로 통용되고 있다.


Transformer

먼저 자연어 처리 분야에서 사용되는 Transformer에 대해 알아보자!

 

Transformer는 2017년 구글이 발표한 "Attention is all you need"에서 나온 모델이다.

 

인코더(Encoder)-디코더(Decoder)로 이루어진 seq2seq 구조를 Attention만으로 구현한 모델로, recurrent neural networks (RNN) 기반 모델들을 가볍게 뛰어넘었다.

https://papers.nips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html

 

Attention is All you Need

Attention is All you Need Part of Advances in Neural Information Processing Systems 30 (NIPS 2017) Authors Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin Abstract The dominant seque

papers.nips.cc

자연어 처리 분야에서는 sequence 데이터를 처리하기 위해 RNN 모델이 많이 사용 되는데,

이는 t-1 시점의 output이 t 시점의 입력으로 사용되는 재귀적인 구조로 이루어져 있다.

하지만, 이는 가까운 단어가 연관성이 높게 나타날 수 밖에 없는 구조로, long-term dependency(hidden layer의 과거 정보가 마지막까지 전달되지 못하는 현상) 라는 문제를 갖는다.


이를 해결하기 위해 제안된 모델이 long short term memory(LSTM)이다.

LSTM은 여러 게이트(gate)를 추가하여, 상대적으로 과거에 있는 단어의 영향까지 고려될 수 있도록 한다.

하지만, LSTM 또한 거리에 대한 한계가 존재하며, 순차적으로 연산한다는 점에서 parameter 수가 많아 연산량 크고, 학습 속도가 느리다는 단점이 있다.


이러한 문제를 해결하기 위해 등장한 것이 Attention이다.

인코더-디코더로 이루어진 구조에서 인코더는 input sequence 하나를 벡터 표현으로 압축하고,

디코더는 이 벡터 표현을 통해서 ouput sequence를 만들어낸다.

 

하지만, 이러한 구조는 인코더가 input sequence를 하나의 벡터로 압축하는 과정에서 input sequence 정보가 일부 손실 된다는 단점이 존재한다.

이를 해결하기 위해 Attention 개념이 사용된 것이다.

이와 같이 Attention으로 인코더와 디코더를 구현한 것이 바로 Transformer이다!

그림 출처: NLP in Korean

Transformer는 앞서 설명한 것과 같이,

Sequence가 인코더의 input으로 들어가고, 디코더를 통해 sequence가 output 되는 인코더-디코더 구조로 되어 있다.

 

이전과 다른점은 인코더-디코더 블럭이 N개 존재할 수 있다는 것이다(N은 hyper-parameter).

그림 출처: NLP in Korean

인코더들은 모두 동일한 구조를 가지고 있으나(Self-Attention 및 Feed Forward 인 두 개의 sub-layer), 같은 가중치(weight)를 공유하진 않는다.

인코더에 들어온 input은 일단 Self-Attention layer 를 지나게 되는데,

하나의 특정한 단어를 인코딩 하기 위해 input 내의 모든 다른 단어들과의 관계를 살펴본다.

이를 통과한 output(Self-Attention의 output)은 다시 Feed-Forward 신경망으로 들어가고, 동일한 Feed-Forward 신경망이 각 위치의 단어마다 독립적으로 적용되어 output을 만든다.

디코더 구조 또한, 인코더에 있는 두 layer를 모두 가지고 있다.

그러나 두 layer 사이에 인코더-디코더 attention이 포함되어 있는데,

이는 디코더가 input sequence 중 각 타임 스텝에서 가장 관련 있는 부분에 집중 할 수 있도록 해준다.

그림 출처: NLP in Korean

자 그럼 이번 논문으로 다시 돌아가보자.


* Summary

Vision AI도 convolutional neural networks (CNN)없이 해결 가능하다!

 

기존 Computer Vision에서 Attention 개념은 CNN 구조에 주로 적용되어,

전체 구조를 유지하되 bottleneck 구조에서 Attention을 수행하는 방식으로 진행되었다.

본 논문에서는 이렇게 특정 구성 요소를 대체하는 것에 그치지 않고, 이미지 패치의 sequence에 Transformer 가 적용 될 때 이미지 분류 작업에서 매우 우수한 성능을 보여준다는 것을 발견하였다.

- 이미지는 이미지 조각의 sequence로 처리

- 대량의 데이터에 대해 사전 학습(pre-training)한 후, 작은 이미지 인식 벤치마크(이미지넷, CIFAR-100, VTAB)에 적용

Vision Transformer(ViT)는 state-of-the-arts (SOTA) CNN 기반의 모델과 비교했을 때 훌륭한 성능을 보이며, 계산량은 최소화


* Introduction

[Transformer의 계산 효율성과 scalability를 vision에 활용]

NLP 분야에서 Transformer는 large text corpus에서 pre-training을 수행하고, task-specific dataset에 대해 fine-tuning을 수행하는 것

Transormer 의 계산 효율성 및 확장성으로 인해, 100B 이상의 파라미터(parameter)를 사용하여 큰 모델을 학습 할 수 있게 되었다.

그러나 Computer Vision 분야에서는 CNN 구조가 여전히 많이 사용되고 있으며,

여러 연구에서 CNN과 유사한 구조를 Self-Attention 과 결합하려고 시도해왔다.

 

본 논문에서는 표준 Transformer을 최대한 변형 없이 직접적으로 이미지에 적용하였다.

이를 위해 이미지를 패치 조각으로 쪼개고, 각각의 패치에 대해 선형 임베딩 시퀀스(linear embedding sequence)를 Transformer 에 대한 input으로 설정한다.

이때 이미지 패치는 마치 NLP 에서의 토큰(token)과 동일한 방식으로 처리된다(모델은 지도 학습 방식으로 학습).

중간 규모의 ImageNet 과 같은 dataset에 모델을 적용해본 결과,

동일 크기의 ResNet 모델 보다 몇 퍼센트 정도 낮은 정확도를 달성한다.

어라? 정확도가 낮다고? 이는 모델이 기대에 미치지 못한 것 처럼 보인다.

==> 이는 Transformer에는 입-출력 위치의 동일성 (translation equivariance)locality 등의 CNN에서 이미지를 이해하는 inductive bias를 학습하지 못하였기 때문에, 데이터가 불충분할 때 일반화 성능이 떨어지는 것처럼 보인다.

이러한 문제는 large-scale dataset(1,400만~3억개)에 대해 모델을 학습했을 때 진면목을 볼 수 있다.

결과적으로, 대규모의 dataset은 CNN의 inductive bias를 능가 할 수 있다!

즉, 대규모의 dataset에서 ViT를 pre-training 결과, 적은 dataset 가진 task에 transfer learning하여 좋은 성능을 얻을 수 있었다.

ImageNet-21k 혹은 JFT-300M data 에 대해 사전 학습했을 때,

다양한 이미지 인식 벤치마크에 대해 ViT는 SOTA이거나, 혹은 근접한 성능을 보였다.

  ImageNet ImageNet-ReaL CIFAR-100 VTAB
ViT 88.55% 90.72% 94.55% 77.63%
Prev-SOTA 88.61%
(EfficientNet-L2-475)
90.9%
(FixEfficientNet-L2)
96.08%
(EfficientNet-L2)
76.3%
(BiT-L, Big Transfer)

* Method

 

 

Vision Transformer

- 이미지는 고정된 크기의 패치로 쪼갠다.

- 패치 각각을 선형적으로 임베딩한 후, 위치 임베딩을 더하여 ouput vector를 일반적인 Transformer 인코더의 input으로 사용한다.

- Classification tasks를 수행하기 위해 추가적으로 학습되는 "classification token"을 만들어 시퀀스에 더한다.

 

Vision Transformer (ViT)

[ Input ]

- 일반적으로 Transformer는 토큰 임베딩에 대한 1차원 sequence를 입력으로 받음

- ViT는 2차원의 이미지를 다루기 위해, 이미지를 flatten된 2차원의 이미지 패치인 sequence로 변환한다.

- 즉, 다음과 같이 표현된다.

 

$$ x∈{R}^{H\times W\times C}\ \to \ x_p∈{R}^{N\times (P^2∙C)}\\  $$

 

- (H, W)는 원본 이미지의 해상도(resolution)였다면 (P, P)는 각 이미지 패치의 해상도를 의미하며, C는 채널 개수를 의미한다.

- 여기서 $N=HW/P^2$는 Transformer의 sequence length를 의미한다. 

- Transformer는 모든 layer에서 고정 벡터 크기 D를 사용하므로, 이미지 패치는 D차원 벡터로 linear projection(patch embedding), BERT의 [CLS]토큰과 비슷하게 임베딩 된 패치의 시퀀스에 $z_0^0=x_{class}$ 임베딩을 추가로 붙여 넣는다.

$z_0^0=x_{class}$

- 이후 이 패치에 대해 나온 인코더 아웃풋은 이미지 representation으로 해석하여 분류에 사용된다.

 

[ Position Embedding ]

- 각각의 patch embedding에 위치 정보를 활용할 수 있도록 position embedding이 더해진다.

- 학습 가능한 1차원의 임베딩을 사용한다(2차원 정보를 유지하는 위치 임베딩도 활용해 보았으나, 유의미한 성능 향상은 없었기 때문).

- Transformer 인코더는 Multi-headed self-attention 및 MLP block 으로 구성된다.

- 추가적으로, Layernorm(LN)은 모든 block 이전에 적용되고 residual connection 은 모든 block 이후에 적용된다.

 

 

<전체 구조에 대한 formulation>

[ Hybrid Architecture ]

- 이미지를 패치 분할하는 대신, CNN기반 모델(ResNet)의 intermediate feature map의 결과를 input sequence로 사용할 수 있다.

- 하이브리드 모델에서는patch embedding projection은 ResNet 의 early stage 로 대체된다.

 

- 특수 케이스로 패치는 1x1 크기를 가질 수 있는데, 이는 input sequence를 단순히 feature map에 대한 차원으로 flatten 한 후 Transformer의 차원으로 projection 한 결과로 볼 수 있다.

- [CLS]에 해당하는 인풋 임베딩과 위치 임베딩은 기존 모델과 동일하게 적용한다.

 

Fine-tuning and Higher Resolution

일반적으로 ViT는 대규모 dataset에 대해 pre-train 하고 downstream task를 fine-tuning 하는 방법을 사용한다.

 

Fine-tuning시에는 pre-train 된 prediction head를 제거하고,

0으로 초기화된 D x K FC-Layer를 연결한다. (K=# of downstream class)

 

이때, fine-tuning 을 진행할 때, pre-train보다 높은 해상도로 진행하는 것이 성능 향상에 도움이 된다.

 

더 높은 해상도의 이미지가 주어질 때, 이미지 패치 크기를 동일하게 유지하므로 유효 sequence length가 길어지게 된다.

 

ViT는 더 높은 하드웨어의 메모리가 허용하는 한, 임의의 길이의 시퀀스를 처리할 수 있다는 장점이 존재한다.

(단, 이 경우 pre-train 된 position embedding이 더이상 의미 없게 되므로, 원본 이미지 position에 따라 2D interpolation을 수행)

 

이러한 해상도 조절과 패치 추출은 ViT에서 이미지의 2D 구조에 대한 inductive bias를 수동적으로 다루는 unique points이다.

 


 

* Experiments

ResNet, ViT 및 Hybride에 대한 represiontation learning capability 를 평가하였다.

데이터 규모에 대한 내용을 실험적으로 평가하기 위해 , 다양한 크기의 dataset를 pre-train 하고 많은 banchmark에서 작업을 수행하였다.

 

Pre-train을 위한 datasets는 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 를 사용하였다.

 

Pre-train 계산 비용을 고려해볼 때, ViT 는 더 낮은 pre-train 비용으로 대부분의 recognition banchmark에서 SOTA 를 달성하였다.

 Model variants

 

ViT 구성은 BERT 에 사용되는 구성을 기반으로 한다(iT-L / 16 모델은 16 x 16 입력 패치 크기의 Large Transformer 라는 것을 의미).

 

Baseline 인 CNN 의 경우 ResNet 을 사용하지만 Batch Norm을 Group Norm로 변경한 표준화된 convolution을 사용하였다.

이는 Transfer를 개선할 수 있으며, 이러한 수정된 모델은 ResNet(BiT)라고 한다.

 

-> Batch Norm을 Group Norm으로 변경한 부분이 Transfer를 개선하는 것에 대해서는 생각해봐야 할 것 같다.

 

논문에서는 Few-shot 및 fine-tuning 정확도를 통해 downstream dataset에 대한 결과를 산출하였다.

Fine-tuning 정확도는 각 dataset에서 Fine-tuning 을 수행한 후, 각 모델의 성능을 측정하였다.

 

Benchmark dataset performance

 

TPUv3 에서 학습되며, Pre-train 시 소요된 일수를 확인 가능하다.

 

ViT-L/16 모델은 모든 dataset에 대해 BiT-L 과 일치하거나, 더 좋은 성능을 보여주는 것을 확인할 수 있고,

동시에 훨씬 적은 computational resource 를 필요로 한다.

 

ViT-H/14 는 ImageNet 및 CIFAR-100 과 VTAB 에서 성능이 더욱 향상되었다.

 

 

[ Pre-training Data Requirements ]

ViT 는 대규모 사이즈의 dataset인 JFT-300M 에 대해 pre-train 하였을 때 더 좋은 성능을 보여주는 것을 확인하였다.

 

Inductive bias에서 데이터 세트의 크기가 얼마나 중요한지에 대한 실험을 수행하였다.

 

결과적으로, 작은 규모의 dataset의 경우 convolutional inductive bias 가 더 큰 비중을 차지하지만,

대규모의 dataset의 경우 CNN 과 같은 패턴을 학습하는데 충분하다는 것을 알 수 있다.

 

[ Scaling Study ]

아래 그림은, JFT-300M dataset에서 transfer 성능에 대해 확장된 연구를 수행한 결과를 나타낸다.

 

Dataset의 크기는 모델 성능에 병목 현상이 없으며, 각 모델의 pre-train 비용 대비 성능을 평가하였다.

 

ViT가 trade-off( accuracy vs. compute )에서 ResNet 보다 효율적이며,

동일한 성능을 달성하기 위해 약 2배 낮은 computational budget을 사용하는 것을 알 수 있다.

 

Hybrid는 적은 computational budget에서 ViT 를 능가하지만, 큰 computational budget에서는 차이가 사라지게 된다.

ViT 는 시도된 범위 내에서 포화(saturate)되지 않는 것으로 보이며, 향후 확장 가능성이 있음을 확인할 수 있다.

[ Inspecting Vision Transformer ]

ViT가 이미지 데이터를 어떻게 처리하는지 이해하기 위해 모델의 representation을 분석하였다.

 

[ Embedding Projection ]

- ViT 의 first layer는 flatten patch 를 더 낮은 차원 공간에 projection 한다.

 

- 아래 그림은 학습된 임베딩 필터 중 중요한 몇 가지 구성 요소들을 시각화한 것이다.

- 이러한 구성요소는 각 patch 내 미세 구조의 low-dimensional representation 에서 기본 함수들을 나타내는 것으로 볼 수 있다.

- Projection 된 이후 학습된 position embedding 이 patch representation 에 추가된다

 

[ Position Embedding ]

- Linear projection 후, 각각의 패치 representation에는 position embedding이 더해지게 된다(이미지 내 거리를 인코딩).

 

- 가까운 patch는 유사한 position embedding 을 갖는 경향이 있으며, 이는 행-열 구조로 보인다(같은 행, 열에 동일한 embedding).

- Self-attention 을 통해 ViT 는 가장 낮은 레이어에서도 전체 이미지에 대한 정보를 통합할 수 있다는 장점이 존재한다.

[ Self-attention ]

- Self-attentio의 weight를 기반으로, 이미지 공간의 평균 거리를 아래 그림과 같이 계산할 수 있다(ViT가 전체 이미지 정보를 통합).

 

- Attention distance 는 CNN의 receptive field size 와 유사하게 볼 수 있다.

- Head 중 일부에서 최하위 layer에 존재하는 대부분의 이미지에 attention 을 적용하여 global 하게 정보를 통합 할 수 있다.

 

- 또 다른 attention head는 밑단 레이어에서 일관적으로 작은 거리의 패치에 집중하는 모습을 보였는데, 이렇게 지역적인 attention은 하이브리드 모델에서는 좀처럼 나타나지 않았다. 즉, 이 attention head는 CNN의 밑단에서 일어나는 것과 비슷한 작용을 하는 것이라고 유추할 수 있다. (생각해보기)

 

- 또한, Attention이 일어나는 거리는 네트워크의 깊이가 깊어질수록 늘어난다는 것도 찾을 수 있었다.

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

 

- 또한 전체적으로 모델은 의미적으로 classification task에 필요한 부분에 attention 되는 것을 찾을 수 있었다.

 

 

[ Self-supervision ]

- Labelling 된 dataset에 대한 pre-training 대신 self-supervision을 pre-training task로 실험을 진행하였다.

 

Masked patch prediction task를 이미지에 적용하여 ViT의 작은 모델 (ViT-B/16)을 학습하여 ImageNet에서 79.9%의 성능을 얻음 (이는 from-scratch로 학습했을 때보다 2% 높은 정확도지만, supervised pre-training에 비해 4% 떨어지는 정확도)

 


* Conclusion

본 논문은 이미지 classification task에 Transformer를 직접 적용하는 방법을 제안하였다.

 

Computer Vision에서 Self-attention 을 사용하는 이전 연구들과는 달리, 구조에 image-specific inductive bias를 사용하지 않음.

이미지를 patch 쪼개어 NLP에서 사용되는 Transformer 인코더로 사용한다.

 

간단하며, 확장성이 풍부한 이 방법은 대규모의 dataset에 대해 pre-train 될 때 잘 작동한다.

따라서, ViT 는 많은 이미지 분류 dataset에서 SOTA 를 능가하거나 비슷한 성능을 보이고, pre-train 비용이 굉장히 저렴하다는 장점을 가진다.

[ Challenge ]

- Object detection이나 semantic segmentation 분야에도 ViT 를 적용

- 기존 pre-training은 단순한 방법을 사용하는 것인데, 이러한 method 에 대해 연구가 더 필요할 것으로 생각됨

 

- 모델의 크기가 증가해도 포화(saturate) 상태가 아닌 것으로 보이기 때문에 ViT 를 확장할 필요성이 있다는 것 (How?)

- Patch 단위로 쪼개는 것을 Wafer Map에 적용한다면 어떠한 것을 고려해야 할지


첫번째 논문으로 Vision Transformer를 리뷰하였는데, 굉장히 어려운 논문이며, 추가적으로 공부할 부분이 필요하다고 생각된다.

 

단순히, pre-training이라고 한다면, 반도체, 철강, 섬유 등의 제조 산업에서는 어느정도 제한성이 존재할 것으로 보인다.

 

이 부분이 보완된다면, 제조 공정 내 ADC 분야에서 엄청난 성능을 보일 것으로 기대한다.

 

 

 

 

 

댓글 영역