[2023-2] 백승우 - AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

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

0. Abstract

트랜스포머 아키텍처는 자연어 처리 작업의 사실상의 표준이 되었지만, 컴퓨터 비전에 대한 적용은 아직 제한적이다. 컴퓨터 비전 분야에서 attention은 CNN와 함께 적용되거나 전체 구조를 그대로 유지하면서 CNN의 특정 구성 요소를 대체하는 데 사용되는데 그친다.

 

CNN에 대한 이러한 의존도가 필요하지 않으며 이미지 패치 시퀀스에 직접 적용되는 순수한 변환기가, 이미지 분류 작업에서 매우 우수한 성능을 발휘할 수 있음을 보여준다. 대량의 데이터에 대해 사전 학습하고 여러 중소형 이미지 인식 벤치마크(ImageNet, CIFAR-100, VTAB 등)로 전송하면 비전 트랜스포머(ViT)는 훈련에 훨씬 적은 컴퓨팅 리소스를 필요로 하면서도 최첨단 컨볼루션 네트워크에 비해 우수한 결과를 얻을 수 있다.

 

1. Information

Self attention-based architecture, 특히 transformer는 NLP에서 주로 선택되는 모델이 되었습니다. Transformer는 대규모 텍스트 코퍼스를 사전 학습한 다음, 소규모 작업별 데이터 세트에서 미세 조정한다. 100억 개가 넘는 파라미터로 전례 없는 크기의 모델을 훈련할 수 있게 되었다. 그리고 모델과 데이터 세트가 증가함에 따라 성능은 여전히 포화 상태에 도달할 기미가 보이지 않는다.

 

그러나 컴퓨터 비전에서는 컨볼루션 아키텍처가 여전히 지배적이다. 자연어 처리의 성공에 영감을 받아 여러 연구에서 CNN과 유사한 아키텍처를 self-attention과 결합하려고 시도하고 있으며, 일부는 컨볼루션을 완전히 대체한다. 컨볼루션을 완전히 대체하는 방안은 이론적으로는 효율적이지만, 특수한 attention 패턴을 사용하기 때문에 최신 하드웨어 가속기에서는 효과적이지 않다.

---

Stand-alone self-attention layer (CNN을 attention으로 완전히 대체한 방안)

Stand-Alone Self-Attention in Vision Models

Attention layer로 모든 convolutional layer를 대체한 'stand-alone attention network'를 만들고, 이러한 모델이 baseline CNN 모델의 성능을 넘어서면서 파라미터와 컴퓨팅 효율까지 좋다는 결과를 보여준다.

 

Local Attention Layer

 

먼저 convolution과 비슷하게 self-attention을 할 windown size k를 정해준다. 논문에서는 이렇게 정한 block을 'memory block'이라고 부른다. 또한, 예시에서는 보기 쉽도록 2차원상에 그려놓았지만 실제로는 3차원 tensor형태를 가진다. 

Single head attention의 output

 

결국 self-attention이 하고 싶은 것은 query 픽셀과 key 픽셀들 간의 내적과 softmax를 취해 similarity matrix를 만들고 이를 values픽셀들과 내적하여 output을 냄으로써, 해당 픽셀과 주변 픽셀들 간의 relational information을 배우도록 유도하는 것이다.

또한 실제로는 single-head가 아닌 multe-head attention을 사용하고, 이는 head당 query, key, values의 depth를 나누어 각각 다른 transforms을 적용함으로써 다양한 representation을 배우도록 한다.

 

Positional Encoding

 

self-attention mechanism은 positional에 대한 정보를 알려주지 못하기 때문에 별도로 이러한 정보를 넣어주어야 한다. 대표적인 방법으로 'Attention is All You Need'에서 나온 sinusodial 함수를 사용한 embedding이 있는데, 여러 연구를 통해 이러한 absolute positional encoding보다는 relative positional embedding이 더 좋은 성능을 낸다는 것이 밝혀졌다.

논문에서도 'relative positional embedding'을 사용하고, 이는 위 이미지와 같이 query기준으로 주변 픽셀들의 상대적인 위치 index들을 사용한다. 실제로는 위와 같이 (0, 0), (-1, 0) 등의 index value를 가지는 matrix가 아니고, learnable look-up table을 만들고서 해당 index의 값들을 가져오게 된다. 좀 더 자세히 설명하면, row offset과 column offset에 대해서 각 따로 learnable matrix가 look-up table로 존재하고 해당하는 offset의 벡터를 가져와서 이를 concat해준다. Row offset embedding과 column offset embedding 각 1/2 d_out 차원을 가지므로, 합치면 d_out 차원이 된다.

---

 

NLP에서 transformer의 성공적인 확장에 영감을 받아, 최소한의 수정만으로 transformer를 이미지에 직접 적용했다. 이를 위해 이미지를 패치(NLP에서의 토큰(단어)과 같은 방식으로 취급)로 분할하고 이러한 패치의 선형 임베딩 시퀀스를 Transformer에 입력으로 제공한다. 지도 학습 방식으로 이미지 분류에 대해 모델을 훈련시킨다.

 

강력한 정규화 없이 ImageNet과 같은 중간 규모의 데이터 세트에 대해 훈련할 경우, 이 모델은 비슷한 크기의 ResNet보다 몇 퍼센트 포인트 낮은 수준의 정확도를 보이지만, 대규모 데이터 세트(1400만~3억 개의 이미지)로 모델을 학습시키면 ViT의 성능이 압도적으로 오른다(대규모 훈련이 귀납적 편향(정규화)보다 더 효과적). ViT(비전 트랜스포머)는 충분한 규모로 사전 훈련한 후 데이터 포인트가 적은 작업에 적용하면 뛰어난 결과를 얻을 수 있다.

 

[공개 이미지넷-21k 데이터 세트 or 사내 JFT-300M 데이터 세트로 사전 훈련된 ViT의 정확도]

ImageNet: 88.55%

ImageNet-ReaL: 90.72%

CIFAR-100: 94.55%

19개 작업으로 구성된 VTAB 제품군: 77.63%

 

2. Method

Model Overview

(1) 이미지를 고정된 크기의 패치로 분할

(2) 각 패치를 선형적으로 임베딩하고, position embedding을 추가

(3) Vector seqence를  Transformer encoder에 입력

(4) Classification를 수행하기 위해 학습 가능한 'classification token'을 시퀀스에 추가

 

Vision Transformer(ViT)

Before Input

기존 트랜스포머에서는 input을 1차원의 token embedding sequence로 받았습니다. 하지만 2차원인 이미지를 input으로 받기 위해서는 이미지를 reshape하는 과정이 필요하였습니다. 이미지를 고정된 크기의 패치들로 쪼갠 후에 각각의 패치를 1차원으로 flatten시켜 ViT에 input으로 넣어주게 됩니다. 즉, 이미지를 쪼개서 만든 패치들이 NLP에서의 token들과 같은 의미를 띄게 됩니다.

Embedding

 

(BERT에서의 class token과 같이) 학습 가능한 임베딩을 첫 번째 패치에 추가하였다. 해당 임베딩은 이미지 전체에 대한 표현을 한다. N+1개의 학습 가능한 1D 포지션 임베딩 벡터 (이미지 패치의 위치를 특정하기 위함)를 만들어준 후, 이를 각 이미지 패치 벡터와 합한다(두 matrix의 concat).

 

Transformer block의 수행 과정

Multihead Self-Attention

Transformer의 QKV self-attention과 거의 동일

 

 

Input seqence인 z에 대하여 query, key, value를 나눕니다. 이후 attention wights인 A_ij​는 query q_i와 key k_j의 pairwise similarity를 기초하여 만들어집니다. q와 k의 행렬곱에 스케일을 거친 후, softmax를 거쳐서 attention matrix인 A를 얻습니다. 마지막으로 A와 v의 weighted sum과의 행렬 곱으로 나온 결괏값이 Self-Attention matrix SA가 됩니다.

Multihead self-attention은 위에서의 self-attention의 연장 버전입니다. k개의 헤드에서 self attention을 병렬적으로 수행한 후에 concatenated outputs들을 project합니다. (B x heads x N x Dim_heads → B x N x D)

 

3. Experiment

Image classification accuracy

 

당시의 SOTA 모델들과 비교하였을 때 더 좋은 성능을 보여주고 있다. 이때 JFT는 구글에서 내부적으로 사용하는 JFT-300M dataset을 의미하며, 이름에서도 알 수 있듯이 약 3억 개의 이미지로 이루어진 거대한 데이터 세트이다.

ImageNet. 작은 데이터 세트에 대해 사전 학습된 경우 대형 ViT 모델은 BiT ResNet(음영 처리된 영역)보다 성능이 떨어지지만, 큰 데이터 세트에 대해 사전 학습된 경우 성능이 뛰어나다. 마찬가지로, 데이터 세트가 커질수록 더 큰 ViT 변형이 더 작은 변형을 추월한다.

 

ImageNet에 대한 선형적인 몇 개의 샷 평가와 사전 훈련 크기 비교. ResNets는 더 작은 사전 학습 데이터 세트에서 더 나은 성능을 보이지만, 더 큰 사전 학습에서 더 나은 성능을 보이는 ViT보다 더 빨리 정체된다. ViT-b는 모든 숨겨진 차원이 절반으로 줄어든 ViT-B이다.

 

다양한 아키텍처에 대한 성능과 사전 훈련 컴퓨팅 비교: ViT, ResNet, 하이브리드. 비전 트랜스포머는 일반적으로 동일한 컴퓨팅 예산으로 ResNet보다 성능이 뛰어나다. 하이브리드는 모델 크기가 작은 경우 순수 트랜스포머보다 성능이 향상되지만, 모델 크기가 큰 경우 그 격차가 사라진다.

 

ViT는 많은 양의 데이터셋에서 좋은 성능을 보인다. 반면, 적은 양의 데이터셋에서 pre-training할 경우, BiT ResNets보다 좋은 성능을 내지 못하게 됩니다. 또한 attention 기법을 사용했기 때문에 연산 효율성이 ResNet 모델보다 월등히 뛰어나다.

 

Inspecting ViT

왼쪽: ViT-L/32의 RGB 값의 초기 선형 임베딩 필터. 가운데: ViT-L/32의 위치 임베딩 유사도. 타일은 표시된 행과 열이 있는 패치의 위치 임베딩과 다른 모든 패치의 위치 임베딩 간의 코사인 유사성을 보여줍니다. 오른쪽: 헤드 및 네트워크 깊이에 따른 관심 영역의 크기. 각 점은 한 레이어에 있는 16개의 헤드 중 하나에 대한 이미지 전반의 평균 attention 거리를 나타냅니다.

 

(1) RGB embedding fileters

ViT의 첫 번째 layer인 linear layer를 분석하였다. 해당 레이어는 flattened patch들을 낮은 차원으로 projection할 때 사용된다. 해당 필터들을 시각화하여 보여지는 구성 요소들이 마치 CNN에서 보여지는 lower dim에서의 필터 모습과 흡사한 점을 확인할 수 있다.

 

(2) Position embedding simliarity

Projection 이후 position embedding을 patch들에 더하는 과정에서, 이미지 내의 거리를 기반으로 position embedding의 유사도를 시각화하였다. 거리가 가까운 패치들일 수록 유사한 position embedding을 나타내는 것을 확인할 수 있다. 또한 같은 행/열에 해당하는 패칭들도 유사한 embedding을 가지고 있다.

 

(3) Size of attened area and Network depth

Network depth에 따라 attention distance를 측정한 그래프이다. attention distance는 attention weights에 기반하여 측정되었고 CNN에서의 receptive field size와 유사한 개념이다. 초기 레이어에서 Average attention distance가 매우 다양하게 나타나는데, 특정한 부분에서는 초기 레이어임에도 불구하고 이미지의 대부분을 반영하고 있는 반면 특정 head에서는 일관성 있게 작은 attention distance를 나타내기도 합니다. 레이어가 깊어질수록 attention distance는 모든 head에서 증가한다.

 

Self Supervision

Transformer는 NLP 작업에서 인상적인 성능을 보인다. 그러나 이러한 성공의 대부분은 뛰어난 확장성뿐만 아니라 대규모의 self-supervised pre-training에서 비롯됐다. 또한 BERT에서 사용되는 masked language modeling 작업을 모방하여 self-supervised를 위한 masked patch prediction에 대한 예비 탐색을 수행한다.

 

Self-supervised pre-training을 통해 더 작은 규모의 ViT-B/16 모델은 이미지넷에서 79.9%의 정확도를 달성했으며, 이는 처음부터 훈련할 때보다 2% 향상된 것이지만 여전히 supervised pre-training에 비해 4% 뒤처지는 수준이다.

 

4. Appendix

Inductive bias

학습으로 얻는 bias와 달리 모델이 정확한 예측을 하게 하기 위한 추가적인 가정을 의미한다. CNN은 지역 특성(locality), translation eqivariance 등을 가정하여 합성 곱 연산이 이루어지는 반면, Transformer는 attention 기법으로 데이터 요소 간의 관계를 파악하는 연산이 이루어진다. 따라서 Transformer로 이미지 데이터를 다루게 될 경우, CNN에 비해 Inductive bias가 부족하다고 볼 수 있다. Introduction에서 소개했던 것처럼, Inductive bias가 부족하다는 단점은 충분한 scale로 ViT를 pre-training한 후에 적은 데이터들로 fine-tunning하는 방법을 통해서 해결하였다.

Position Embedding

Position Embedding은 이미지의 위치 정보를 인코딩하기 위해서 사용된 기법이다. Transformer는 1차원의 임베딩 벡터를 input으로 받기 때문에, 이미지를 쪼개 만든 patch들을 flatten하여 1차원의 벡터로 만드는 과정이 요구된다. 벡터를 1차원으로 펼쳐버리니, 2차원인 이미지의 공간 정보를 잃어버리게 된다. 따라서 Position Embedding 벡터들을 1차원의 패치들에 더하여 이미지의 위치 정보를 유지할 수 있다.

 

(1) 1-dimensional positional embedding (default)

Input을 raster order sequence로 고려하여 1차원으로 positional embedding을 구현하였다.

 

(2) 2-dimensional positional embedding

Input을 이미지의 그리드로 고려하여 X-embedding과 Y-embedding을 나누어 positional embedding을 구현하였다.

 

(3) Relative positional embeddings

패치간의 상대적인 거리를 이용하여 positonal embedding을 구현하였다. attention mechanism 시에, 추가적인 original query와 positional embeddings을 key로 활용하여 추가적인 attention 연산을 수행하고 bias로 취급하여 main attention logis에 추가하는 방식이다.

Self supervision

Self-supervised learning은 라벨링되지 않은 데이터셋을 이용해 해당 데이터셋에 대한 전반적인 pre-training을 하는 학습 방법이다. Pre-training 후에는, downstream task에 라벨링된 작은 데이터셋으로 fine-tunning이 이루어진다. 많은 양의 데이터들을 라벨링하지 않고도 pre-training이 가능하다는 장점을 가지고 있다. ViT에서는 BERT에서 사용하는 masked patch prediction을 따서 self-supervisied learning을 진행하였다. 50%의 패치 임베딩들 중 80%는 [mask] embedding, 10%는 random other patch embedding, 10%는 가만히 놔두는 방식으로 masked prediction을 진행하였다.