[논문 리뷰] ViT: AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE(2021)

이번에 리뷰할 논문은 AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE 입니다.

 

https://paperswithcode.com/paper/an-image-is-worth-16x16-words-transformers-1

Papers with Code - An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

 

 

---

ViT를 이해하기 위해서는 기본적으로 transformer에 대한 이해가 필요합니다.

혹여나 접하지 않으셨다면 우선 transformer를 이해하고 보시는 것을 추천드립니다.

---

 

 

 

 

NLP분야에서는 2017년도에 나온 Self-attention을 base로 하는 Transformer 모델의 computational efficiency와 scalability 덕에 좋은 성능을 보여주었습니다. 기존의 RNN기반 모델들의 경우 일정량의 데이터 이상을 학습하게 되면 오히려 성능이 떨어지는 모습을 보여주었지만, transformer의 경우 scalability가 좋기 때문에 대용량 데이터에 대해서도 학습이 가능했죠. 이런 특징 덕에 transformer를 기반으로 한 BERT, GPT와 같은 LLM모델들이 나올 수 있었습니다.

 

물론 computer vision 분야에도 transformer의 핵심인 self-attention을 적용하려는 시도가 있었습니다. 하지만 현재까지의 시도로는 scalability의 한계가 있었고, 대규모 데이터셋에 대해서는 ResNet이 여전히 지배적이었습니다.

 

저자는 가능한 transformer를 최소한으로 수정하면서 이미지를 standard transformer에 적용하려 하였으며, 이를 위해 image를 여러 patch로 split한 후 linear embedding 변환기의 input으로 사용하였습니다. 이는 마치 각 patch들을 token으로 취급하는 셈이죠.

 

위와 같이 만든 모델에 대해 ImageNet과 같은 적당한 크기의 데이터셋으로 학습을 하였더니 ResNet은 넘지 못하지만 적당한 accuracy를 보여주었다고 합니다. 이는 CNN에는 있는 translation equivariance(이동 불변성)과 locality(지역성)의 inductive bias가 부족했기 때문입니다. 여기에서 inductive bias란 기계 학습에서 모델이 학습 데이터 외에 데이터를 일반화하는데 사용하는 가정 이라고 생각하시면 됩니다. 예를 들어, 선형 모델의 inductive bias로는 데이터가 선형적으로 분포해야한다는 가정을 들 수 있겠죠. 

 

transformer 역시 inductive bias가 존재하는데 바로 translation equivariance와 locality 입니다. translation equivariance의 경우 객체의 위치에 상관없이 동일한 객체를 식별하고 분류할 수 있는 능력을 의미하며, 이는 CNN의 경우 자연스럽게 내제되어있음을 알 수 있습니다. 또한 locality의 경우 이미지의 한 픽셀은 주변 픽셀과의 관계가 출력에 더 큰 영향을 미칠 것이라는 가정이죠. 이 역시 CNN에는 자연스럽게 반영되어있습니다. 하지만 transformer는 기본적으로 sequence 데이터를 처리하기 위해 만들어진 모델이기에 입력 데이터의 순서에 민감하게 작용하며, 각 입력 요소들은 모두 독립적으로 처리됩니다. 따라서 해당 inductive bias가 약한 편입니다. 

 

하지만 아주 큰 대규모 데이터셋(1,400만 ~ 3억 개의 이미지)에서 훈련을 하게 되면 inductive bias를 능가하여 오히려 더 좋은 성능을 보여준다고 합니다.

 

 

 

 

 

 

 

 

 

 

 

우선 transformer는 자연어를 위한 모델이라고 하였습니다. 즉, input data를 sequence data로 받는다는 것을 의미합니다. ViT에서는 이미지를 sequence data로 변환하기 위해 이미지를 patch들로 나누어 flatten 을 시킵니다. 

 

original image의 경우 $ x \in \mathbb{R} ^{H \times W \times C} $ 의 형태를 가집니다. 이를 P 사이즈를 갖는 patch로 split하여 sequence 데이터로 만들게 되면  $ x \in \mathbb{R}^{N\times(P^{2} \cdot C)} $ 와 같은 형태를 가지게 됩니다.

 

• N : patch의 개수 ($ HW/P^{2} $)
• P : 각 patch의 size (P^2 는 한 patch의 픽셀 개수)
• C : channel의 개수

 

위 형태를 보면 마치 자연어 데이터를 embedding 한 형태와 비슷해보입니다. 각 patch들을 token으로 취급하는 것이죠. 하지만 해당 sequence data의 차원 수인 $ P^{2} \cdot C $의 경우 transformer가 요구하는 embedding 차원수인 D와 일치하지 않을 수 있습니다. 즉, 요구하는 차원 수로 고정시키기 위해 linear projection을 사용해 D차원으로 mapping 을 거칩니다. 이로써 transformer를 학습할 수 있는 상태가 되었고 이러한 과정을 patch embedding 이라고 합니다.

 

 

 

 

또한 ViT는 classification을 수행하기 위해 BERT에서 사용되었던 [class] 토큰(CLS token)을 사용합니다. BERT의 사용법과 마찬가지로 첫 번째 layer에 들어가는 sequence data의 첫 부분에 들어가게 됩니다. 

 

 

위 그림을 보면 가장 상단에 1 row가 추가된 것을 확인할 수 있습니다. 기본적으로 squence는 위에서 아래로 내려오기 때문에 가장 상단에 추가된 것이며, 이를 $ z^{0}_{0} $ 으로 표현합니다. $ z $는 embedding된 patch들을 가르키는데 사용되며 $ z^{a}_{b} $가 있을 때 a는 squence 내의 위치를 indexing하며, b는 layer 를 indexing합니다. 즉, $ z^{0}_{0} $ 는 0번째 layer의 0번째 sequence에 위치한 embeded patch를 의미하죠.

 

또한 position embedding 역시 더해주어야 합니다. image를 patch로 split하여 flatten을 시켰기 때문에 각 patch에 대한 위치정보가 모두 사라진 상태입니다. 다시 말해 2d 정보가 사라진 상태이기 때문에 position embedding을 통해 해당 patch가 어떤 위치에 존재했는 지에 대한 정보를 전달해줍니다. 물론 position embedding을 더해준다고 정확한 2d 정보를 얻을 수 있는 것은 아니며, 이는 학습되는 파라미터입니다. 

 

 

여기까지 ViT의 첫 번째 layer로 들어가는 input을 완성했습니다. 이 과정을 식으로 표현하면 다음과 같습니다.

 

 

$ x_{class} $는 CLS 토큰을 의미하며, $ x^{1}_{p}E $는 D차원으로 mapping 하기 위해 $ E $를 내적하는 것으로 확인할 수 있습니다.

 

 

다음으로 각 patch에 대해 position정보를 주기 위해 $ E_{pos} $를 더해 position embedding을 적용하는 것을 확인할 수 있습니다.

 

 

 

 

이후에는 embedded된 patch에 대해 transformer의 encoder연산을 진행합니다.

 

 

이를 식으로 표현하면 다음과 같습니다.

 

• MSA : Multihead Self-Attention
• MLP : Multi Layer Perceptron
• LN : Layer Normalization

 

추가로 $ z_{l-1}, z'_{l} $을 더해주는 것은 residual connection을 수행하는 것입니다. 이 역시 transformer encoder의 한 과정이죠. 

 

 

 

 

다음으로 CLS 토큰을 포함한 input sequence data가 transformer를 거쳐 나온 output의 $ z^{0}_{L} $ 인 CLS 토큰 출력값을 LN(Layer Normalization)을 적용하여 image representation인 y를 반환할 수 있습니다. 이를 식으로 표현하면 다음과 같죠.

 

 

저는 CLS 토큰에 대해 한가지 의문점이 들었습니다. "어떻게 CLS 토큰은 이미지의 전체 특성을 나타내어 분류의 측도가 되는 것일까?" 사실 이 부분은 self-attention 매커니즘을 이해하고 있다면 당연한 결과입니다. 

 

 

self-attention의 경우 query에 대해 해당 sequence의 모든 token들과의 attention score를 구하여 적용합니다. 다시 ViT로 넘어와 문장 앞에 CLS 토큰을 추가해준다면 해당 토큰은 뒤에 오는 모든 patch에 대해 attention 을 적용하게 되기 때문에 해당 이미지의 모든 patch에 대한 관계를 학습하는 것이 가능해집니다.

 

 

 

 

 

 

이렇게 구한 CLS 토큰은 1xD vector 로 반환될 것입니다. (transformer의 encoder는 shape에 대해 멱등(Idempotent)하기 때문) 후에는 1xD vector를 다시 MLP 연산을 하여 class 분류를 진행합니다. 

• 일반적으로 MLP는 1개 이상의 fully-connected layer로 구성된 신경망을 의미

 

 

 

 

 

 

 

 

 

일반적으로 ViT는 fine tuning하여 사용됩니다.(inductive bias를 충족하기 위해 대규모 데이터셋으로 훈련 필요) 만약 fine tuning으로 사용하게 된다면 pretrain된 예측 layer를 제거한 후 0으로 초기화한 DxK feedforward layer를 추가합니다.

 

• K : class 개수
• CLS토큰인 1xD 벡터를 1xK 벡터로 변환

 

 

또한 pretrain할 때 더 높은 resolution으로 fine tuning하는 것이 더 좋은 성능을 낸다고 합니다. 이를 수행하기 위해서 patch size를 고정하면 늘어난 해상도만큼 sequence 길이는 증가할 것이고, pretrain된 position embedding은 더 이상 의미가 사라질 것입니다.

 

이때는 단순히 pretrain된 position embedding을 2D interpolation을 수행하여 원래 이미지에 따라 조정하여 사용하게 됩니다. 이 과정은 유일하게 inductive bias가 수동으로 transformer에 전달되는 요소라고 합니다.