논문: https://arxiv.org/pdf/2010.11929
저자: Alexey Dosovitskiy∗$^†$, Lucas Beyer∗, Alexander Kolesnikov∗, Dirk Weissenborn∗, Xiaohua Zhai∗, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby∗$^†$
∗ equal technical contribution, $^†$ equal advising, Google Research, Brain Team
인용: Dosovitskiy, Alexey, et al. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).
코드: https://github.com/google-research/vision_transformer
튜토리얼: https://johyeongseob.tistory.com/71
0. 초록
Transformer 구조가 NLP의 대표적인 기반이 되었지만, 컴퓨터 비전 영역에는 아직 적용되지 않았다. 컴퓨터 비전에서 attention의 개념은 CNN과 함께 사용되거나, 혹은 CNN의 특정 요소를 대체하는데 사용된다. 이에 저자들은 이미지 분류 작업에 CNN없이 순수한 transformer 구조로만 이루어진 모델을 제안한다. 그리고 대규모 이미지 데이터로 사전학습한 Vision Transformer (ViT)가 다양한 벤치마크 데이터셋 (ImageNet, CIFAR-100, VTAB)에 대해 기존 CNN 모델 대비 더 적은 훈련 계산비용으로 SOTA를 달성함을 보였다.
1. 서론
Self-attention 구조인 Transformer가 NLP에서 주류를 이룬다. Transformer의 계산 효율성과 확장성으로, 100B 파라미터 이상 크기를 가지는 모델 훈련이 가능하다. 하지만 컴퓨터 분야에서는 CNN이 우세하다. 왜냐하면 기존 CNN에 attention을 추가하려는 시도는 해당 attention이 병렬화가 힘들기 때문이다. 저자들은 Transformer에 영감을 얻어, Transformer의 구조와 유사하면서 image에 적용할 수 있도록 연구하였다. 그 결과 저자들은 이미지를 patch로 나누고 Transformer의 입력 token으로 사용하였다. 대규모가 아닌 데이터셋 ImageNet으로 훈련한 결과는 기존 SOTA보다 성능이 안 좋았으나, 대규모 데이터셋 (14M-300M 이미지)으로 훈련한 결과 SOTA를 달성하였다. 저자들은 대규모 데이터셋 훈련이 inductive bias (내제된 규칙, CNN 구조 등) 보다 우세함을 발견하였다.
2. 관련 연구
Transformer과 NLP 분야에서 SOTA를 달성하면서, 대규모 Transformer 기반 모델이 large corpora를 통해 사전학습 한 뒤, 테스크에 맞는 fine-tune으로 사용된다. 예를 들어, BERT (Bidirectional Encoder Representations from Transformers)는 사전 훈련에 denosing self-supervise 학습 방식 (e.g. masked language modeling)을 사용하거나, GPT (Generative Pre-trained Transformer)는 사전 훈련에 언어 모델링 (자기 회귀)을 사용한다.
Self-attention이 이미지에 적용된 초기 방법들은 각 픽셀을 모든 다른 픽셀에 가중치를 적용했다. 하지만, 해당 방법은 픽셀 수의 제곱 만큼 비용이 발생하므로 현실적으로 적용하기 어렵다. 대안으로 local neighborhoods에만 attention을 적용하거나, 다양한 크기의 block내에서만 attnetion을 수행하는 방법 등이 있으나, 구조가 복잡하여 GPU에 효과적으로 적용하기 어려웠다.
저자들이 제안하는 방법 (ViT)과 가장 유사한 기존 연구는 입력 이미지를 2x2로 자른 패치를 attention으로 이용하는 방법이다. 이 방법은 ViT와 매우 유사하나, 저자들의 목표는 대규모 이미지를 Transformer에 사전훈련하여, CNN SOTA와 경쟁력있는 모델을 만드는 것이다. (ViT는 제목처럼 이미지를 16x16으로 패치화함)
3. 제안 방법
저자는 제안 모델이 vanilla transformer와 가장 유사하도록 설계하였다.
3.1 Vision Transformer (ViT)
기존 Transformer는 1차원 벡터 시퀸스를 입력으로 받는다. 이미지는 2차원 매트릭스이므로, 저자들은 이미지 $x \in \mathbb{R}^{H \times W \times C}$를 Transformer의 입력에 적절하게 벡터 (token) 시퀸스 $x_p \in \mathbb{R}^{N \times (P^2 \cdot C)}$ 형태로 변환한다. $N=HW / P^2$은 패치 이미지 개수이다. Transformer는 입력으로 고정된 벡터 크기 $D$를 사용하므로 저자들은 patch 벡터를 D차원으로 projection한다 (Eq. 1). 위치 정보를 추가하기 위해 position embeddings은 패치 embeddings에 element-wise sum 한다.
BERT 의 [class] 토큰과 비슷하게, 저자들은 입력 토큰 시퀸스에 학습가능한 embedding (=벡터, 토큰) 을 추가한다 ($z^0_0$). 이후 transformer의 encoder를 지난 해당 벡터 ($z^0_L$) 는 이미지 분류 ($y$)에 사용된다 (Eq. 4). $z^0_L$은 이후 classification head를 통과하여 예측값 logit을 생성한다. classification head는 사전 훈련시 MLP, 전이 학습에는 단일 선형 레이어를 사용한다.
ViT의 transformer encoder는 multiheaded self-attention (MSA, see Appendix A)과 MLP block (Eq. 2, 3)이 반복되는 구조이다. Layernorm (LN)은 각 block 전에, residual connection은 각 block 후에 적용된다. MLP에서 활성화함수는 GELU를 사용한다 (Eq. 3).
\begin{align}
\mathbf{z}_0 &= [\mathbf{x}_{\text{class}}; \mathbf{x}_p^1 \mathbf{E}; \mathbf{x}_p^2 \mathbf{E}; \cdots; \mathbf{x}_p^N \mathbf{E}] + \mathbf{E}_{\text{pos}}, \quad \mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D},\ \mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times D} \tag{1} \\
\mathbf{z}'_\ell &= \text{MSA}(\text{LN}(\mathbf{z}_{\ell-1})) + \mathbf{z}_{\ell-1}, \quad \ell = 1 \dots L \tag{2} \\
\mathbf{z}_\ell &= \text{MLP}(\text{LN}(\mathbf{z}'_\ell)) + \mathbf{z}'_\ell, \quad \ell = 1 \dots L \tag{3} \\
\mathbf{y} &= \text{LN}(\mathbf{z}_L^0) \tag{4}
\end{align}
Appendix A: Multihead Self-Attention
일반적인 qkv self-attention 에서 저자는 입력 시퀸스 $z \in \mathbb{R}^{N \times D}$에서 가중치가 적용된 v (value)를 계산한다. 가중치 매트릭스 $A_{ij}$는 query $q^i$와 key $k^j$의 유사도를 계산한다.
\begin{align}
[\mathbf{q}, \mathbf{k}, \mathbf{v}] &= \mathbf{z} \mathbf{U}_{qkv}, \quad \mathbf{U}_{qkv} \in \mathbb{R}^{D \times 3D_h} \tag{5} \\
\mathbf{A} &= \text{softmax} \left( \frac{\mathbf{q} \mathbf{k}^\top}{\sqrt{D_h}} \right), \quad \mathbf{A} \in \mathbb{R}^{N \times N} \tag{6} \\
\text{SA}(\mathbf{z}) &= \mathbf{A} \mathbf{v} \tag{7}
\end{align}
Multihead self-attention (MSA)는 self-attention (SA)의 확장이다. 저자들은 head라고 부르는 self-attention을 k개 병렬계산한 다음, 결과들을 다시 채널별 연결 (concatenate)한다.
\begin{equation}
\text{MSA}(\mathbf{z}) = [\text{SA}_1(\mathbf{z}); \text{SA}_2(\mathbf{z}); \cdots; \text{SA}_k(\mathbf{z})] \, \mathbf{U}_{\text{msa}}, \quad \mathbf{U}_{\text{msa}} \in \mathbb{R}^{k \cdot D_h \times D} \tag{8}
\end{equation}
3.2 Fine-tuning and Highter Resolution
ViT는 기본적으로 대규모 데이터셋에 대해 사전 훈련한 뒤, 원하는 task에 맞는 소규모 전이 학습을 수행한다. 저자들은 전이 학습 시, prediction head를 $K$개 클래스를 가진 데이터셋에 맞게 $D \times K$ 순전파 layer로 교체한다. 또한 ViT는 고해상도 이미지가 입력인 경우에도 같은 크기의 패치로 나누므로, 패치 수가 늘어난다. 이때는 기존 positional embedding을 보간하여 패치 수에 맞게 조절한다.
4. 실험
저자들은 대규모의 데이터셋을 사전 학습을 진행한 후, 다양한 벤치마크 데이터셋으로 전이 학습 능력을 평가하였다.
4.1 Setup
Datasets. 사전 학습 데이터셋으로 ImageNet-1K (1k classes and 1.3M images), ImageNet-21K (21k classes and 14M images), JET (18k classes and 303M high-resolution images) 을 사용한다. 이 중 데이터셋 간 중복데이터는 제거한다. 전이학습 데이터셋으로 ImageNet validation, ImageNet ReaL, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-103을 사용한다. 전처리 방법은 Kolesnikov et al. (2020)*을 따른다. 저자들은 전이 학습 능력을 평가할 때, VTAB 방식도 사용한다. 이는 기존 19개 데이터셋을 3개 카테고리로 분류한 다음, 각 데이터셋의 훈련 데이터를 오직 1000장씩 제공하여, 개별적으로 원본 테스트 데이터셋에 대해 전이 학습을 평가하는 방법이다.
*Big transfer (BiT): General visual representation learning
Model Variants. 저자들은 BERT를 참고하여 ViT의 조건을 구성하였다. "Base"와 "Large"는 BERT 구조와 동일하며, 이보다 더 큰 모델은 "Huge"이다. 표기 방법은 다음과 같다. ViT-L/16은 "Large" 크기에 16*16 크기 이미지를 패치로 사용하는 모델이다.