배경과 문제의식
NLP 분야에서는 pre-train에는 대규모 데이터를 이용하고, smaller task-specific dataset으로 fine-tuning을 하는 Transformer를 사용하는 접근법이 주류가 되었다. 이와 달리, Computer Vision에서는 CNN 구조가 여전히 주류였다. 물론 CNN의 구조와 Self-Attention과의 결합을 시도해보려는 연구들은 존재했지만, 이마저도 CNN 구조에 의존하는 것이었다. CNN이 오랜 시간 CV를 지배해온 데에는 CNN 구조의 Inductive Bias 덕분이었다.
Inductive Bias란 모델이 학습하기 전에 이미지들의 특성을 고려하여 모델을 설계한 것이다.
CNN의 설계 구조에는 크게 Locality, Translation Equivariance, Hirarchical Composition을 볼 수 있다.
Locality
NxN 필터를 통해 주변 N^2개의 픽셀만 보는 것은, “중요한 시각적 패턴은 근처 픽셀들 사이에 존재한다”는 가정을 기반으로 설계된 것이다.
Translation Equivariance
같은 사이즈의 필터를 이미지 전체 슬라이딩하는데, 이는 Weight Sharing이다. 해당 구조는 “같은 종류의 패턴은 이미지 어디에 나타나도 같은 방식으로 감지되어야 한다”는 가정을 기반으로 설계되었다.
Hierarchical Compsition
얕은 레이어는 에지와 색상, 중간 레이어는 텍스처와 모양, 깊은 레이더는 물체의 부분과 전체를 학습하는데, 해당 구조는 “복잡한 시각 패턴은 단순한 패턴들의 조합이다”라는 가정을 기반으로 설계되었다.
하지만, CNN에는 문제가 있었는데, 이는 이미지 내부에 멀리 떨어진 요소들 사이의 관계가 중요할 때 생기게 된다. 예를 들어, 배경과 전경의 관계를 파악하거나, ‘이 사람이 무엇을 들고 있는가’를 판단하려면 이미지 전체를 동시에 봐야 한다. CNN은 이를 해결하기 위해서는 많은 레이어들을 쌓아야 해결할 수 있으며 이는 Computation Cost가 상당히 많이 든다.
따라서 본 논문에서는 CNN의 Inductive Bias를 버리고, 순수한 Transformer로 이미지를 처리해보려는 시도를 해보게된 것이다.
핵심 아이디어
CNN의 Inductive Bias는 데이터가 부족할 때의 선택이었을 뿐, 이미지를 패치 토큰으로 변환한 순수 Transformer는 충분한 데이터 앞에서 그 가정 없이도 CNN을 넘어설 수 있다.
아키텍처 구조
ViT는 224x224 이미지를 받으면 16x16 크기의 패치 196개로 자르게 된다.
각 패치를 1D로 펼치면 16x16x3=768차원 벡터가 된다. 이걸 선형 변환으로 d차원 임베딩으로 만드는게 Patch Embedding이다.
여기에서 눈 여겨봐야 할 것은, 시퀀스 맨 앞에 [CLS] 토큰을 붙이는데 이는 BERT에서 가져온 개념으로 이미지 정보가 없는 학습 가능한 벡터이다.
다른 하나는 Position Enbedding이다. 이는 Transformer가 입력 순서를 원래 모르기에, 각 패치가 이미지에서 어디 있었는지 알려주는 학습 가능한 위치 벡터이다.
[CLS] 토큰의 추가로 총 197개 토큰을 Transformer Encoder에 통과시키는데, 이때 핵심은 Self-Attention에서 197개 토큰이 전부 서로를 바라보게 된다. 해당 과정을 L번 반복하고 최종 레이어에서 [CLS] 토큰만 뽑으면, 해당 토큰에는 전체 이미지 정보를 응축한 벡터가 된다.
이를 MLP 헤드에 통과시켜 최종 클래스 확률로 변환하는 과정을 거친다.
직관적 이해
196명의 분석가 회의가 있다고 해보자.
이미지를 받으면 196명의 분석가가 각자 하나씩 패치를 맡는다. 이 분석가들이 처음엔 자기 담당 구역만 알고 있고, 그 다음 전원 회의를 열게 된다.
해당 회의에서는 모든 분석가가 서로에게 “당신 구역엔 뭐가 있나요?”를 동시에 물어보고 답한다. 이게 Self-Attention이다.
해당 회의를 L번 반복하면서 각자의 이해가 점점 깊어지게 된다.
해당 회의에 처음부터 아예 빈 노트를 갖고 들어온 회의 서기 ([CLS] 토큰)가 있다. 본인의 담당 구역은 없지만, 모든 분석가의 발언을 듣고 기록하고 L번의 회의가 끝나면 서기의 노트가 이미지 전체를 요약한 가장 압축된 정보가 된다.
이때 이를 보고 최종적으로 본 이미지가 뭔지 판단을 내리게 된다.
CNN이라면 분석가들이 옆 구역 담당자하고만 얘기하고 멀리 있는 사람들과는 여러 단계를 거쳐야만 소통할 수 있게 된다.
한계
- 데이터 의존성이 가장 큰 한계이다. 약 130만 장이 되는 ImageNet 수준에서는 CNN보다 성능이 낮고, JFT-300M이라는 3억 장짜리 데이터셋이 있어야 CNN의 성능을 넘는다. DeiT가 이후 연구에서 Knowledge Distillation으로 본 문제를 일부 해결하려 했지만, 데이터 의존성이라는 구조적 한계는 여전하다.
- 해상도가 올라가면 연산량이 폭발한다. Self-Attention은 토큰 수의 제곱으로 연산량이 늘어나는데, 224x224에서는 196개 토큰이라 괜찮지만, 448x448이 되면 토큰이 1,024개로 늘고 연산량은 약 27배 뛴다. 이러한 이유로 Swin Transformer가 등장하게 된다.
- 위치 인코딩이 1D이다. 이미지는 2D 구조인데 패치를 1D 시퀀스로 늘어놓고 위치를 1D로 인코딩하는 방식인데, 본 논문에서는 2D 위치 인코딩과 성능 차이가 없다고 했지만 직관적으로는 2D 구조 정보의 손실이 존재할 것 같다는 의문이 든다.