[Vision Transformer] AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

안녕하세요, 데이터 과학 전공하고 있는 황경하입니다.

오늘은 저번 시간에 Transformer에 대해 공부했기 때문에 이를 Vision 분야에 접목한 Vision Transformer에 대해 공부해보려 합니다. 데이터를 구하기 비교적 쉬워진 지금 시점에서는 데이터만 더 많다면 기존 모델보다 더 좋은 성능을 보일 수 있다고 한다면, 그 모델은 각광을 받을 것입니다. 오늘 읽을 논문인 ViT(Vision Transformer)는 데이터 수가 많을 때 SOTA 모델보다 더 뛰어난 성능을 보인 논문입니다.

 

* 리뷰할 논문: https://arxiv.org/pdf/2010.11929/1000

* 해당 글에 첨부된 사진들은 모두 위 논문에서 발췌한 것임을 미리 밝힙니다.

* 해당 글에 설명은 논문을 기반으로 한 주관적 해석이므로 틀린 부분이 있다면, 댓글로 지적 부탁드립니다!

* 참고 자료: https://www.youtube.com/watch?v=KXBohWurGOU

https://youtu.be/91Qipj5NMnk?si=klhwL6eVXXyl51zG

 

---

# Abstract

Transformer는 자연어 처리 분야에서는 사실상 표준이 되었지만, 비전 분야에서는 사용된 적이 없습니다.

비전 분야에서는 Attention + CNN 방식은 사용된 적이 있죠. ViT 모델은 CNN에 대한 의존도를 없애고, 단순히 트랜스포머만으로 이미지 분류 작업에서 우수한 성능을 보임을 증명했습니다. Large-Scale의 데이터로 사전 학습되고, 이후에 Mid-Scale의 데이터로 전이 학습되었을 때, ViT는 Cost는 훨씬 적으면서 SOTA CNN 모델보다 더 뛰어난 성능을 보입니다.

---

# Introduction

자연어 처리 분야에서는 Transformer를 통해 대규모 데이터셋으로 Pre-Training 한 후, 더 적은 데이터셋으로 Fine-Tuning 하여 뛰어난 성능을 보였습니다. Transformer는 계산 효율성과 확장성이 좋아 1000억 개 이상의 파라미터를 가진 모델을 생성하고 훈련하는 것도 가능해졌습니다.

 

그러나, 컴퓨터 비전 분야에서는 아직 CNN 아키텍처가 더 우세합니다. Transformer의 성능이 인정을 받음에 따라서 CNN에 Attention을 도입하거나 혹은 Convolution을 완전히 대체하는 연구가 이루어졌으나 이론적인 내용뿐 실제로 적용되긴 어렵습니다. 따라서, 아직은 ResNet을 활용한 아키텍처 구조가 SOTA를 기록했습니다.

 

본 논문에서는 Transformer의 구조를 최대한 유지하면서 이미지에 적용해보려 합니다. Super-vised Learning 방식을 채택하며, NLP의 Transformer에서 입력 토큰을 이미지의 패치로 대체하여 생성합니다. 이렇게 생성된 ViT 모델은 Mid-Scale Image Dataset에서 비슷한 크기의 ResNet보다 더 낮은 성능을 보입니다. 그러나, 데이터셋이 더 커지면 성능은 SOTA 모델에 근접하건 혹은 능가합니다.

 

# Related Works

트랜스포머를 컴퓨터 비전에 적용할 때에 단순히 이미지 자체에 Self-Attention을 진행하면, 각 픽셀이 다른 모든 픽셀에 대해 계산을 진행해야 합니다. 이는 픽셀 수에 대한 2차 비용 때문에 계산 비용이 어마어마하게 커집니다. 따라서, 이를 해결하기 위한 여러 논문들이 나왔습니다. 그러나, 위 논문들은 이론적으로는 매우 효율적이나 GPU에서 구현하기 위해서는 복잡한 엔지니어링이 필요합니다.

 

이러한 연구 중 하나의 연구는 이미지는 2*2 크기의 패치로 분할하고 이를 Self-Attention을 적용한 모델이 있습니다. 이는 현재 저희가 보고 있는 ViT 모델과 매우 유사합니다. 다만, 본 논문에서는 Large-Scale 데이터셋을 통한 사전 학습이 Transformer를 최신 CNN 모델과 유사하게 혹은 더 능가하게 만든다는 것을 보여주며, 기존에는 2*2 크기의 패치로 저해상도 이미지만 적용이 가능했지만, 본 논문은 16*16 크기의 패치를 사용함으로써 중간 해상도 이미지도 처리할 수 있다는 것이 차이점입니다.

---

# Method

## Vision Transformer (ViT)

모델의 구조는 Transformer에서 크게 벗어나지 않습니다. 먼저, 입력 데이터는 H*W*C shape을 가진다고 가정해 봅시다. 우리는 이 이미지를 패치단위로 쪼개어 Linear Projection을 수행하고 이를 transformer Encoder에 넣을 것입니다.

 

입력 이미지가 H*W*C이고, N개의 패치를 만들 것이며, 그 패치의 크기는 P라고 해봅시다. 그러면, 각 패치의 크기는 P*P 일 것이고,  채널은 C 일 것입니다. 이를 Flatten 한다고 했을 때, 각 패치는 P*P*C의 크기를 가지게 되겠죠. 이런 패치가 N개 존재하니 N*(P^2*C)의 2D Patch가 존재할 것입니다. 그런데, 생각해 보면 우리가 정할 수 있는 하이퍼 파라미터는 P밖에 없습니다. 이미지의 크기가 고정되어 있기 때문에 N은 P에 의해 결정되죠. 따라서, 논문에서는 아래와 같이 표현합니다.

그리고 이를 Linear Projection을 통해 D차원으로 매핑합니다. 이 과정을 통틀어서 패치 임베팅이라고 표현합니다. 이렇게 통과된 값은 Transformer에 입력되기 전에 Positional Embedding을 거치게 됩니다. 이는 Transformer에서 사용한 개념과 동일하게 Transformer에는 Recurrence 혹은 Convolution이 존재하지 않기 때문에 위치 정보를 따로 포함해주어야 하기 때문입니다. 이 Positional Embedding 과정에 사용되는 값은 2D 값이 아닌 1D 값을 사용합니다.

 

마지막으로, BERT에서 사용한 것과 같이 클래스 토큰을 하나 추가합니다. 이는 랜덤 한 값으로 구성되어 있습니다. 즉, 아무런 정보도 가지고 있지 않은 데이터입니다. 이를 추가하는 과정이 잘 이해가 되지 않을 수 있습니다. 이를 이해하기 위해선, Attention이 병렬적으로 진행됨을 이해해야 합니다. 즉, 패치별로 Attention이 진행되고, 이 Attention을 거친 값은 자기 자신에 대한 정보가 가장 많이 포함되어 있을 것입니다. 그렇기 때문에 이 모든 정보를 MLP Layer에 넣어 예측을 진행하면 오히려 모델이 무엇이 중요한 정보인지를 파악하기 어렵게 만듭니다. 따라서, 본 논문에서는 클래스 토큰을 하나 추가하여 Attention 과정에서 다른 패치 정보들과 정보를 혼합하고 이렇게 나온 값만을 MLP에 입력하여 예측을 진행합니다. 이렇게 되면, 클래스 토큰은 모든 패치에 대한 정보를 혼합한 정보를 가진 데이터가 되고 올바른 예측을 수행할 수 있게 됩니다.

 

이렇게 하면, 입력 데이터와 전체적인 아키텍처를 본 것입니다. 좀 더 자세히 들어가서, 내부 구조를 살펴보면 ViT의 Transformer Encoder의 구조는 Vanila Transformer의 인코더 구조와 다릅니다.

 

왼쪽: ViT 인코더, 오른쪽: Vanila Transforemr 인코더

 

가장 큰 차이점은 Normalization과 Residual Connection의 위치입니다. Vanila Transformer는 Attention 후에 Residual Connection과 Normalization을 거친 반면에, ViT는 Attention 전에 Normalization을 취한 후 Attetnion을 거치고 Residual Connection을 거칩니다. 이렇게 하는 이유는, 이미지의 경우 자연어보다 더 복잡한 특성을 가지고 있기 때문에 데이터가 불안정합니다. 따라서, Normalization을 통해 안정성을 확보한 후 Attention을 수행하고 이후에 다시 불안정한 데이터를 합쳐줌으로써 성능 향상이 있었다고 합니다.

 

이 중에서 Normalization은 Layer Norm을 사용했다고 합니다. 이는 Batch Norm과 달리 패치별로 Normalization을 해주는 것으로, 이미지 하나하나에 더 집중하겠다는 의미를 포함합니다.

 

이후 진행되는 MLP는 단순히 은닉층을 한 개 포함한 레이어로, Linear 층 두 개, 활성화 함수 한 개로 구성되어 있습니다. 여기서 활성화 함수는 GELU를 사용합니다. (최근 트렌드가 ReLU보다 GELU를 사용하는 쪽으로 변하고 있다고 합니다.)

 

이를 수식으로 표현하면 아래와 같습니다.

 

 

- Inductive bias: 이는 귀납적 편향이라고 부릅니다. 즉, 어떤 가정을 함으로써 생길 수 있는 편향입니다. 일반적으로 Convolution 레이어는 널리 알려진 것처럼 지역성을 가지고 있으며, translation equivariance이 모든 레이어에 내장되어 있습니다. 따라서, CNN은 이러한 Inductive bias를 가지고 있죠. 그러나, ViT는 MLP 레이어에 Convolution만 존재하고, 나머지 레이어에서는 Attention을 사용하여 Inductive bias를 가지고 있지 않습니다. 수동적으로 넣는 부분이라고 한다면, 이미지를 패치로 자르고, 미세 조정 시 위치 임베딩 값을 다시 설정해야 할 때에 2-dimensional neghborhood bias는 발생할 수 있습니다. 이 외에는 Bias를 가지고 있지 않아서 적은 데이터셋에서는 CNN보다 더 낮은 성능을 보입니다.

 

- Hybrid Architecture: 입력 이미지의 패치를 생성하는 과정을 CNN으로 대체할 수 있습니다. 이는 CNN과 ViT를 결합할 수 있다는 것으로, 패치 임베딩 과정 대신에 CNN을 통한 Feature Map으로 대체하는 방법을 뜻합니다.

---

## FINE-TUNING AND HIGHER RESOLUTION

일반적으로 ViT는 대규모 데이터셋으로 사전 학습을 진행한 후 Downstream task를 위해 미세조정하여 사용합니다. 이때에는 ViT의 MLP를 제거하고, 0으로 초기화된 가중치를 가진 FFN을 추가합니다. 이 FFN은 입력은 D차원을 가지며, 출력은 K (예측 클래스 수)를 가집니다.

 

미세조정 시에는 사전학습 데이터보다 더 높은 해상도를 가진 이미지로 하는 것이 더 유용할 수 있다고 합니다. 이는 패치 크기는 동일하므로 시퀀스 길이 (패치 개수)가 더 많아지기 때문이라고 합니다. 그러나, 입력 크기가 달라짐에 따라 Positional Embedding 값을 다시 조정해줘야 할 필요가 있습니다. 

---

이 이후에는 Experiments, Research Result가 나오는데, 사실 그냥 결과물들의 나열이라 딱히 배울만한 내용은 없습니다. 따라서, 본 블로깅에 포함하지 않겠습니다.

---

## 코드 실행 프로세스

다양한 논문리뷰 글이 있어 저도 살펴봤지만, 자세히 어떻게 구동되는지를 설명한 글이 없더라구요. 그래서, 한 번 작성해보려 합니다.

 

1. Multi-Head Attention에 넣기 위한 데이터 전처리

 

예를 들어, 224 * 224 이미지가 입력되었다고 해봅시다. 우리가 나눌 패치 크기는 16이에요. 그러면, 224/16 = 14로 가로, 세로 14개의 패치가 생성되었을 것입니다. 그러면, 이를 linear projection을 거쳐 1D shape으로 변환하고 positional embedding 값을 추가해줄 것입니다. (Linear Projection은 코드로 구현하면, 한 줄입니다. 파이토치의 경우 nn.Conv2D(3, 256, kernel_size = 16, stride = 16)로 3채널을 256차원으로 변경해주는 것이죠.) 그러면, Flatten된 효과처럼 1*1 shape이며 크기가 16*16 = 256인 텐서가 14*14 = 196개 만들어질 것입니다. 최종적으로 196 * 256의 크기를 가진 데이터가 되겠네요. (이 196은 positional embedding 값이 더해진 상태입니다.) 그리고, 이에 Class Tocken 하나를 추가해주면, 준비는 끝입니다.

 

최종적으로, Multi-Head Attention에 넣는 데이터는 197 * 256 shape을 가집니다.

 

2. Multi-Head Attention

 

Attention은 기본은 Query, Key, Value를 만드는 것입니다. 또한, Attetion의 가장 큰 목적은 패치들끼리의 관계를 학습함에 있죠. 따라서, 모든 패치를 Attention 모듈에 넣어줄 것입니다.

 

먼저, 197 * 256 shape의 input data를 X 행렬이라고 하겠습니다. 그리고, W_q, W_k, W_v가 합쳐져 있는 행렬을 W라고 하죠. 그러면, 두 행렬을 곱하여 WQ 행렬을 만들겠습니다. (참고로, Attention의 경우 입력과 출력의 크기가 동일해야 합니다. 따라서, W_q, W_k, W_v는 모두 256 * 256의 shape을 가집니다. 그러면, W 행렬은 256 * 768의 shape을 가지게 되겠네요.)

 

WQ를 수행하게 되면, 결과는 197 * 768의 shape을 가지게 될 것입니다. 이를 3등분 하겠습니다. (이 때는 순서가 매우 중요합니다. W_q, W_k, W_v 행렬은 가중치를 공유하지 않기 때문에 앞 1/3은 XW_q = Query의 결과물, 중간 1/3은 XW_k = Key의 결과물, 마지막 1/3은 XW_v = Value의 결과물입니다.) 

 

그러면, Query, Key, Value가 정의되었기에 기존에 설명드린 대로 Query와 Key를 내적하고 Softmax를 취해 Attention Score를 만들어주고 이후에 Value를 곱해 Attention Output을 만들어냅니다.

 

단, 여기서는 Multi-Head Attention을 사용하기에 n_head = 12인 경우, W 행렬의 크기가 1/12가 됩니다. 즉, Query, Key, Value가 현재는 256 * 256이지만, Multi-Head Attention의 경우에는 256/12 = 21으로, 256 * 21의 shape을 가지게 될 것입니다. 그 후에, 각 결과물을 Concatenate 시켜 197 * 256 Shape을 만들겠지요.

 

3. MLP

위 과정을 거치게 되면, 각 패치들은 자신의 값이 가장 가중치가 높기에 자신을 중심으로 하여 다른 패치들간의 관계를 학습한 결과물들이 나오게 될 것입니다. 입력 패치가 197개였기에 결과도 197개의 토큰이 나오겠지요. 다만, Class Tocken을 제외한 196개의 결과는 자기 중심적으로 계산되었기에 예측에 편향이 있을 수 있습니다. 따라면, 아무 정보도 없던 Class Tocken의 정보를 MLP의 입력으로 사용합니다.

 

MLP는 단순히 은닉층 1개의 Layer로 구성되어 있습니다. 즉, Linear 층 2개와 GELU 활성화 함수 하나가 연결된 구조이죠.

마무리...

오늘은 ViT에 대해 공부해 봤습니다. 사실, 내용이 그리 쉽지만은 않습니다. 전체적인 아키텍처를 좀 더 자세히 보기 위해 노력했는데 여러분들의 이해에 도움이 될지 모르겠습니다. 논문 내에 코드가 포함된 Github 주소가 있기 때문에 궁금하신 분들은 직접 코드를 보면서 이해한 것이 맞는지 확인해 보시는 것이 좋을 것 같습니다.

 

제 자그마한 공부 내용의 공유가 여러분의 이해에 도움이 되었으면 합니다. 감사합니다.