[Paper Review] AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

Paper : https://arxiv.org/pdf/2010.11929

Github : https://github.com/google-research/vision_transformer

 

2021년에 발표된 논문이지만 ViT는 현재까지도 널리 활용되는 필수적인 논문이기에 다시 한번 읽고 정리해보고자 한다. 이 논문은 자연어 처리 분야에서 우수한 성능을 보인 Transformer 구조를 컴퓨터 비전 분야에 적용하는 방법을 제시했으며, 이미지를 패치 단위로 분할하여 각 패치를 하나의 단어처럼 취급하는 방식을 도입했다. 대규모 데이터셋으로 사전 학습한 뒤 중소형 데이터셋으로 Fine-tuning할 경우, Convolution Network보다 적은 컴퓨팅 자원으로도 최고 수준의 성능을 달성할 수 있었다.

ViT는 이후 Swin Transformer, DeiT, MAE 등 수많은 후속 연구의 기반이 되었고, 현재도 Computer Vision 분야에서 중요한 아키텍처로 사용되고 있다.

 

📚 사전 지식

이 논문을 효과적으로 이해하는 데 필요한 사전 지식을 정리했다.

 

Transformer

https://dream2reality.tistory.com/26

[Paper Review] Attention Is All You Need

 

Translation invariance & Translational equivariance

Computer Vision에서 Translation은 모든 점을 주어진 방향과 동일한 거리로 이동시키는 기하학적인 변환(Geometric transformation)을 의미한다.

Translation은 수학적으로 각 점 $x$에 대한 상수 벡터 $v$의 합으로 표현할 수 있다.

 

$$ T_v(x)=x+v $$

 

아래 예제에서는 이미지 (a)에서 각 픽셀을 왼쪽 150 픽셀 이동하여 이미지 (b)를 얻은 것이다.

 

Translation invariance

어떤 Translation에도 변경되지 않는 속성은 Translation invariance이다. 위의 이미지를 고려해 보겠습니다. 아래 예제처럼 픽셀 값이 이동했더라도 두 이미지 모두에서 여자임을 인식할 수 있다. Image classifier는 두 이미지 모두에 대해 "woman"이라는 레이블을 예측해야 한다. 실제로 classifier의 출력은 대상의 위치에 영향을 받아서는 안된다. 따라서 classifier의 출력은 Translation invariance입니다.

 

Translation Equivariance

함수 f와 함수 g가 있을 때 두 함수의 적용 순서가 결과를 바꾸지 않을 때 Translation Equivariance이다.

$$ f(g(x))=g(f(x)) $$

Convolution 연산은 Translation Equivariance이며 이미지 변환을 적용하여 Convolution을 수행한 결과와 Convolution 연산을 수행하고 변환을 적용한 결과는 동일하다.

Translation Invariance and Equivariance in Convolutional Neural Networks

Convolution Neural Network(CNN)은 자연적으로 회전(rotation), 크기 변화(scaling), 아핀 변환(affine transformation)에 대해서 Translation Invariance and Equivariance 을 가지지 않는다.

이러한 한계를 극복하기 위해 Data Augmentation을 사용해야 하며 이를 통해 CNN은 기하학적 변환에 더 강해질 수 있다.

Example

Translation invariance : 얼굴 인식 시스템에서 사람의 얼굴이 이미지의 어느 위치에 있든 동일한 사람으로 인식할 수 있다.

Translation equivariance : 객체 탐지에서 물체가 이미지 내에서 이동해도 해당 물체의 특징을 일관되게 감지할 수 있다.

 

 

Inductive bias

AI에서 Inductive bias(귀납적 편향)란 학습 모델이 학습 데이터에서 보지 못한 새로운 입력에 대해 예측할 때 사용하는 가정이나 선호도를 의미한다. 다시 말해, 모델이 일반화를 할 때 사용하는 '선입견' 또는 '가정'이라고 볼 수 있다.

 

예를 들어 **K-최근접 이웃(KNN)**의 Inductive bias는 가까운 데이터 포인트들은 비슷한 속성을 가진다고 가정하며, **Convolution Neural Network(CNN)**의 대표적인 Inductive bias은 Locality와 Translation Invarianced다. 픽셀들이 공간적으로 가까운 것들끼리 더 중요한 관계를 가진다고 가정하며 동일한 특징이 이미지의 어느 위치에 나타나도 동일하게 인식해야 한다고 가정한다.

 

적절한 Inductive bias은 모델이 새로운 데이터에 대해 더 잘 일반화하도록 도와주며, 모델이 학습 데이터에 과적합되는 것을 막아주는 역할을 한다. 또한, 모든 문제에 똑같이 잘 작동하는 모델은 없으며, 특정 문제에 적합한 Inductive bias을 가진 알고리즘을 선택하는 것이 중요하다.

 

📺 Vision Transformer

Vision Transformer(ViT)의 전체 구조를 보면 아래 그림과 같다.

ViT가 내부에서 어떻게 동작하는지 하나씩 살펴본다.

 

Model Architecture

아래 그림은 ($H, W, C$) 크기의 이미지를 ($P^2\cdot C$) 크기의 패치 N 개로 나눈 후 Flatten 시키는 과정이다.

($224, 224, 3$)의 이미지를 ($16^2\times 3$) 패치로 나누는 것을 예로 들면 가로 세로를 각각 16으로 나누면 ($14\times16,14\times16, 3$)이 ($196,16, 16, 3)$으로 되는 것을 볼 수 있다. 즉, ($16, 16, 3$) 크기의 패치가 196개. 이것을 Flatten하면 ($196, 16^2\times3)$이므로 ($196, 768$)의 패치($x_p$)가 나온다.

 

 

이렇게 Flatten한 패치를 Linear Projection을 거쳐 D차원의 Feature Vector로 매핑한다. 이렇게 나온 결과를 Patch Embedding이라 부른다.

($196, 768$)의 패치를 ($768,512$)의 Linear Layer를 거쳐 ($196,512$)의 Feature Vector($x_pE$)로 만든다.

 

이후에 기존 Patch Embedding($x_pE)$에 학습 가능한 Embedding Vector $x_{class}$를 추가한다.

($196,512)$ → ($197, 512$)

 

그다음에 학습 가능한 Position Embedding을 각 Patch Embedding에 더해준다.

($197,512$) + ($197,512$) = ($197,512$)

 

수식을 살펴보면 아래와 같이 정리할 수 있다.

 

여기까지가 아래 그림에 대한 내용이다.

1. 이미지를 패치 단위로 나눔
2. 패치를 Flatten 한 후 Linear Projection 진행
3. 학습 가능한 [Class] Embdding을 추가
4. Positional Embedding을 각 Patch에 더함

 

Patch+Position Embedding을 이제 Transformer Encoder에 입력하게 된다. Transformer 내부 구조에 대해 알아본다.

 

Embedded Patches가 들어오면 Layer Normalization, Multi-Head Attention을 수행한 후 Residual Connection을 진행한다. 이후에 Layer Normalization, MLP 수행 후 Residual Connection을 진행하며 이 과정을 Transformer Layer 수(L)만큼 진행한다.

 

수식은 아래와 같다.

 

마지막 Transformer 블록에서 나온 [class] 토큰의 Representation인 $z^0_L$에 Classification Head를 부착하여 Classification을 진행한다. ($z^0_L$ 토큰은 이미지의 특징을 종합적으로 담고 있다.)

 

아래 그림에서 왼쪽은 ViT에서 Transformer Encoder고 오른쪽은 기존 Transformer에서 Encoder 구조이다.

ViT에서는 Normalization을 먼저 해줬는데 이는 기존에 Attention 이후에 Normalization을 했던 것과는 다르다. 대부분의 최신 모델들을 Pre-LN을 채택하고 있으며 이는 학습 안정성 향상, 더 깊은 네트워크 학습 가능, Warm-up 필요성 감소 등의 이유다.

최근에는 Layer Normalization이 아닌 RMSNorm과 같은 경량화된 정규화 기법 도입 및 Normalization을 최소화하는 시도가 있다.

 

Inductive bias

ViT는 CNN보다 이미지 데이터에 대해 Inductive bias가 훨씬 적다. CNN은 이미지의 공간적 특성이 구조적으로 내재되어 있는 반면, ViT는 이러한 특성들을 데이터를 통해 학습해 나가기 때문이다. 이는 ViT가 더 유연한 학습이 가능하지만, 학습시키려면 대규모 데이터셋이 필요한 이유다.

(Attention 메커니즘 특성상 Patches가 한 번에 입력되기 모든 Patch 간의 관계를 직접 학습하게 되며 각 순서 정보를 따로 학습해야 한다.)

 

Hybrid Architecture

ViT에서는 원본 이미지를 직접 패치로 나누지만 Hybrid 구조에서는 CNN으로 Feature Map을 생성하여 패치를 추출한다. 해당 방식을 사용하면 CNN의 Local Feature 추출과 Transformer의 Global Attention 메커니즘을 결합할 수 있다.

($224, 224, 3$) 크기의 이미지를 CNN 네트워크에 통과시킨 후, 마지막 Feature map이 ($14,14,512$)일 경우 이를 14x14 개의 패치를 사용하는 방식

 

Finetuning and Higher Resolution

ViT를 Pre-train시킬 때, 큰 데이터셋으로 학습시킨 후 이보다 더 작은 데이터로 Fine-tuning을 한다. Fine-tuning을 할 때 Patch 크기는 유지하면서 Pre-train에 사용한 데이터보다 해상도가 큰 데이터로 학습하면 성능이 더 좋다는 연구가 있다. 이 경우 기존에 학습된 Position Embedding의 의미가 없을 수 있다. 이때 Position Embedding에 대해서 2D Interpolation을 진행하게 되는데, 이때 유일하게 ViT에서 수동으로 inductive bias가 주입되게 된다.

 

예를 들어 이미지 A(224, 224, 3)와 이미지 B(384, 384, 3)에서 Patch 사이즈가 $16\times16$일 경우에 이미지 A에서는 패치 개수가 196개가 나오지만, 이미지 B에서는 575개가 나오기 때문에 기존 Position Embedding이 의미가 없어진다.

 

 

🧪 실험 및 결과

• ResNet, Vision Transformer, Hybrid에 대해서 평가를 진행
• 각 모델의 데이터 요구 사항을 이해하기 위해 다양한 크기의 데이터셋에 대해 Pre-train하고 여러 benchmark를 평가
• ViT는 모델 3가지가 있으며 Base, Large, Huge로 구성된다.
• ViT-L/16은 입력 패치가 16 x 16인 Large 모델을 뜻한다.

 

• JFT-300M 데이터셋에서 Pre-train된 모델을 3번의 fine-tuning을 거쳐 성능 비교
• ViT 모델이 모든 데이터셋에서 다른 모델보다 성능이 우수한 것을 볼 수 있으며 적은 컴퓨팅 리소스가 소요된 것을 볼 수 있다.

 

• Visual Task Adaptation Benchmark(VTAB)는 이미지 모델의 일반화 능력과 적응력을 평가하기 위한 Benchmark
• 아래 이미지는 VTAB를 Natural, Specialized, Structed 그룹으로 나누고, 각 Task에서 성능을 비교한 것
  • Natural Tasks : 자연 이미지 태스크
  • Specialized Tasks : 특수 도메인 태스크
  • Structured Tasks : 구조적 태스크

 

• ViT Huge 모델은 작은 데이터셋에서 사전 학습된 경우 BiT(ResNet)보다 성능이 떨어지지만, 대규모 데이터 세트에서 사전 학습된 경우 성능이 뛰어나며 데이터셋이 커질수록 ViT Huge 모델이 ViT Large 보다 성능이 좋아진다.
• ImageNet의 Linear Few-shot과 Pre-training 크기 비교에서도 ResNet은 더 작은 Pre-training 데이터셋에서 더 나은 성능을 보이지만, 더 큰 Pre-training 데이터셋에서는 성능이 ViT보다 좋지 않고 더 빨리 정체된다.

 

• 다양한 모델 아키텍처(ViT, ResNet, Hybrid)의 성능을 사전 학습에 사용된 컴퓨팅 자원과 비교한 결과
• 동일한 컴퓨팅 자원을 사용했을 때 ViT가 ResNet보다 더 좋은 성능을 보인다.
• 작은 모델 크기에서는 Hybrid 모델이 순수 ViT보다 낫다.

 

• Low Layer에서도 다양한 거리의 특징을 볼 수 있다.
• CNN은 작은 영역부터 시작해서 넓은 영역을 보게 되는 것과 대조된다.

 

• Position Embedding은 초기에 값이 그냥 랜덤 값인데 학습된 Position Embedding의 Similarity Matrix를 구하면 아래와 같다.
• 어느 정도 위치 정보에 대한 값이 학습된 것을 볼 수 있다.

 

하지만 아래와 같이 학습 파라미터에 따라 민감하게 변하는 것을 볼 수 있다.

 

• Position Embedding에 대한 실험 결과이며 1D와 2D의 성능은 비슷했다.
• Position Embedding을 주지 않았을 때는 성능이 많이 떨어지는 것을 볼 수 있다.

 

🏁 결론

• ViT는 이미지를 패치로 분할해서 기존 NLP용 Transformer를 적용
• 대규모 데이터로 사전학습 시 매우 효과적이며 기존 SOTA 대비 좋은 성능을 보이면서도 사전학습 비용이 저렴
• 객체 탐지, 세그멘테이션 등 다른 Vision Task로 확장에 대한 연구 필요
• Self-supervised Learning 방법 개선 필요
• 모델 Scale 확장을 통한 성능 향상 가능성

 

📄 참고자료

https://www.baeldung.com/cs/translation-invariance-equivariance

https://www.appliedaicourse.com/blog/inductive-bias-in-machine-learning/

https://google-research.github.io/task_adaptation/