[트랜스포머] 비전 트랜스포머 (Vision Transformer)
📄 관련 논문: [ICLR 2021] An Image is Worth 16X16 Words: Transformers for Image Recognition at Scale
Overview
비전 트랜스포머 (ViT)는 이미지 분류 문제를 해결하기 위해 고안된 모델이다.
분류는 클래스를 예측하는 비생성적 문제이기 때문에 생성 과정이 필요하지 않다.
따라서 ViT는 트랜스포머의 인코더만 사용한다.
전체 구조는 언어 모델의 인코더와 거의 동일하며, Embedding, Attention, MLP로 구성되어 있다.
단지 이미지를 처리하기 때문에, 단어 대신 이미지의 패치가 토큰 역할을 한다.
패치 임베딩 (Patch Embedding)
ViT에서는 이미지를 픽셀 단위로 직접 처리하지 않고, 사전에 고정된 크기의 작은 패치(Patch) 단위로 분할하여 1차원 토큰 시퀀스로 변환한다.
각 패치는 독립적인 토큰으로 처리된다.
따라서 학습 때 사용한 이미지 해상도와 추론 때 사용할 이미지 해상도가 다르더라도, 패치의 개수만 달라지기 때문에 모델 자체는 정상적으로 동작한다.
패치 임베딩의 전체 과정은 다음과 같다.
1. Patchify
패치 크기 $P$는 사전에 고정되며, 일반적으로 $16\times16$ 크기의 패치를 많이 사용한다.
이미지에서 분할된 $i$번째 2D 패치는 다음과 같은 텐서 형태를 띤다.
\[\mathbf{x}_i= \begin{bmatrix} x_{1,1} &x_{1,2} &\cdots&x_{1,P}\\ x_{2,1} &x_{2,2} &\cdots&x_{2,P}\\ \vdots &\vdots&\ddots&\vdots\\ x_{P,1} &x_{P,2} &\cdots&x_{P,P}\\ \end{bmatrix}\in\mathbb{R}^{C\times P\times P} ,\quad\text{where }x_{row,col}\in\mathbb{R}^{C\times1\times 1} \tag{1}\]2. Flatten
트랜스포머 모델에 입력하기 위해, 2D 패치 $\mathbf{x}_i$를 1차원 벡터 $\mathbf{x}_p^i$로 Flatten한다.
\[\mathbf{x}_p^i=\text{Flatten}(\mathbf{x}_i)\in\mathbb{R}^{P^2C} \tag{2}\]입력 이미지에서 생성된 전체 패치의 개수가 $N$개일 때, 최종적으로 이미지는 다음과 같은 1차원 패치들의 Sequence로 재구성된다.
\[\mathbf{x}_{1:N}= \begin{bmatrix} \mathbf{x}_p^1,&\mathbf{x}_p^2,&\cdots&,\mathbf{x}_p^N \end{bmatrix} ,\quad\text{where }N=\frac{HW}{P^2} \tag{3}\]이때 픽셀값은 $[0,255]$ 범위를 가지기 때문에, 학습의 안정성을 위해 $[0,1]$ 범위로 정규화를 한다.
\[\mathbf{x}_p^i=\text{Normalizaiton}(\mathbf{x}_p^i) \tag{4}\]3. Linear Projection
Flatten된 벡터를 바로 트랜스포머의 입력으로 사용하지 않고, 신경망에 통과시켜 $D$ 차원의 Patch Embedding $\mathbf{e}_i$로 변환한다.
\[\vphantom{\Big(} \mathbf{e}_{1:N}=\mathbf{x}_{1:N}\cdot W_E\in\mathbb{R}^{N\times D} \tag{5}\]이때, 학습 가능한 클래스 토큰 (cls token) $\mathbf{e}_{cls}$를 시퀀스의 맨 앞에 추가한다.
클래스 토큰은 학습 과정에서 이미지의 전역적인 표현을 학습하며, 이후 분류 단계에서 최종적으로 사용된다.
4. 위치 임베딩 (Positional Embedding)
패치를 1차원으로 펼치면서 잃어버린 2차원 공간 정보를 보완하기 위해, $N+1$개의 패치 임베딩에 각각 위치 임베딩 (positional embedding) $\mathbf{p}_i$를 더해준다.
이를 통해 모델이 각 패치의 공간적 위치 정보를 인식할 수 있게 된다.
아래 그림은 각 입력 패치의 위치 임베딩이 다른 모든 패치들의 위치 임베딩과 얼마나 유사한지를 시각화한 맵이다. 이는 트랜스포머 모델이 패치들 간의 공간적 관계를 어떻게 인코딩하고 있는지를 보여준다.
예를 들어, 아래 그림에서 중앙 4행 4열에 위치한 작은 히트맵은 $(4,4)$ 위치의 패치를 기준점으로 삼아, 다른 모든 패치들과의 유사도를 나타낸다.
이 히트맵의 정중앙은 가장 밝은 노란색으로, 자기 자신과의 위치 임베딩 유사도가 가장 높음을 의미한다. 기준 패치에서 가까운 위치에 있는 패치들일수록 밝은 색상을 띠며 (유사도가 높음), 모서리나 가장자리로 갈수록 거리가 멀어지기 때문에 어두운 색상을 띠는 것 (유사도가 낮음)을 확인할 수 있다.
Multi-Head Self-Attention
언어 모델에서 사용되는 Multi-Head Self-Attention 매커니즘을 ViT에도 동일하게 사용한다.
각각의 Head는 서로 다른 가중치 행렬을 가지며, 이를 통해 자신만의 고유한 $Q,K,V$를 독립적으로 만든다.
\[Q,K,V\in\mathbb{R}^{h\times(N+1)\times d_h} ,\quad\text{where }d_h=\frac{D}{h} \tag{8}\]이로 인해, 1번 Head는 물체의 윤곽에만 집중하고, 2번 Head는 색감에 집중하는 등 다양한 feature를 포착할 수 있다.
Attention 연산은 동일하게 아래와 같이 정의된다.
\[\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^\top}{\sqrt d_k}\right)V\in\mathbb{R}^{(N+1)\times(N+1)} \tag{9}\]언어 모델에서는 Attention map이 단어 간의 관계를 파악하듯, ViT에서는 패치 간의 관계를 학습한다.
최종적으로 ViT Encoder에서는 총 $L$번의 $[\text{Norm}\to\text{Multi-Head Attn}\to\text{Norm}\to\text{MLP}]$ 과정을 거쳐, 패치들 간의 관계를 반영한 contextual feature를 출력한다.
클래스 토큰 (Class Token)
클래스 토큰은 입력 이미지 전체의 정보를 종합하여, 최종적으로 이 이미지는 오토바이다.와 같은 결정을 내리는 역할을 한다.
트랜스포머 인코더를 통과한 임베딩 벡터 중, 클래스 토큰 벡터만 MLP Head를 통과하여 최종 클래스를 예측한다.
아래의 그림은 클래스 토큰의 Attention map을 시각화한 것이다.
직관적으로, 클래스 토큰은 이 이미지가 어떤 클래스인지 결론을 내리기 위해 어떤 패치의 정보가 가장 중요해?라는 질문을 던지는 것으로 볼 수 있다.
- 밝은 노란색 영역 (High Score): 클래스 토큰이
너의 정보가 결정을 내리는 데 매우 중요해.라고 강하게 주목하는 패치들이다. - 어두운 보라색 영역 (Low Score): 클래스 토큰이
너의 정보는 별로 중요하지 않으니 무시해도 되겠어.라고 판단한 패치들이다.
위와 같은 Spatial Attention Map은 연산 결과에서 클래스 토큰 자기 자신과의 Attention 가중치를 제외하고, 나머지 실제 이미지 패치들에 대한 가중치 배열을 원래의 2차원 이미지 해상도에 맞게 Reshape하여 시각화할 수 있다.
\[A_{\text{cls}}=\text{softmax}\left(\frac{\mathbf{q}_{\text{cls}}\cdot K^\top}{\sqrt{d_k}}\right) =\begin{bmatrix} a_{\text{cls}},&a_1,&a_2,&\cdots,&a_N \end{bmatrix}\in\mathbb{R}^{1\times(N+1)} \tag{10}\]예를 들어 $N=9$ 일 경우, 첫 번째 원소인 $a_{\text{cls}}$를 버리고 나머지 9개의 가중치를 $3 \times 3$ 크기의 2차원 공간 형태로 변환하여 시각화할 수 있다.
\[\text{Spatial Attention Map} =\begin{bmatrix}\vphantom{\Big(} a_1&a_2&a_3\\a_4&a_5&a_6\\a_7&a_8&a_9 \end{bmatrix}\in\mathbb{R}^{3\times3}\]예측 (Prediction)
최종적으로, Transformer Block을 통과한 출력 클래스 토큰 $\mathbf{z}_{\text{cls}}$는 MLP Head를 통과하여 클래스 개수 $K$와 동일한 차원의 벡터로 변환된다.
\[\text{Logits}=\text{MLP}(\mathbf{z}_{\text{cls}})\in\mathbb{R}^K \tag{11}\]이후 Softmax 함수를 적용하여, 클래스별 확률 분포로 변환한다.
\[P=\text{softmax}(\text{Logits})\in\mathbb{R}^K \tag{13}\]결과적으로 이 확률 분포에서 가장 높은 확률 값을 가진 라벨이 선택된다.
\[\hat{y}=\underset{i}{\arg\max}(p_i) ,\quad\text{where }p_i\in P \tag{13}\]
Inductive Bias
Inductive Bias는 모델이 처음 보는 데이터에 대해 예측을 수행할 때 사용하는 내장된 기본 가정이다.
모델이 태어날 때부터 가지고 있는 일종의 ‘편견’이나 ‘세상을 보는 방식’이라고 생각하면 된다.
CNN
CNN은 이미지 처리에 특화된 두 가지 강력한 가정을 가지고 있다. (Inductive Bias가 큼)
이 가정들은 이미지 데이터의 구조적 특성에 매우 잘 맞기 때문에, 비교적 적은 데이터로도 빠르고 효율적으로 학습할 수 있다.
1. 지역성 (Locality)
이는 ‘중요한 정보는 보통 서로 가까이 붙어있다.’는 가정이다.
CNN의 필터는 이미지의 작은 영역만을 관찰하면서, 그 영역 안에서 특징을 추출한다.
예를 들어 고양이의 눈, 코, 입이 서로 가까이 붙어있는 것처럼, 이미지의 주요 시각적 패턴은 local pattern을 가진다고 보는 것이다.
2. 이동 불변성 (Translation Invariance)
이는 ‘특징의 위치가 달라져도 동일하게 인식할 수 있다.’라는 가정이다.
하나의 필터가 이미지 전체를 슬라이딩하며 동일한 연산을 적용하기 때문에, 모델은 특정 패턴이 이미지의 어디에 있든 동일한 특징으로 인식할 수 있다.
예를 들어 고양이가 이미지의 왼쪽에 있든 오른쪽에 있든, CNN은 그것을 같은 ‘고양이’로 인식한다.
Transformer
반면 트랜스포머는 이러한 사전 가정이 없고, 모델이 이미지 패치들 간의 관계를 Self-Attention을 통해 데이터로부터 직접 학습해야만 한다. (Inductive Bias가 작음)
트랜스포머는 ‘모든 요소는 다른 모든 요소와 어떤 식으로든 관련이 있을 수 있다.’라는 기본 전제를 가지고 있다.
즉, 모델은 입력 이미지의 한 패치가 다른 모든 패치와 어떤 관계를 맺고 있는지를 전역적으로 계산하며, 이는 Locality 같은 제한적인 가정을 하지 않기 때문에 이미지의 왼쪽 끝 픽셀과 오른쪽 끝 픽셀 사이에도 직접적인 상호작용이 가능하다.
트랜스포머가 대규모 데이터셋을 필요로 하는 이유
이처럼 트랜스포머는 약한 Inductive Bias를 가지기 때문에, 데이터의 구조나 패턴에 대한 사전 지식이 거의 없다.
따라서 모델이 스스로 모든 관계와 구조를 처음부터 학습해야 하므로, 효과적으로 학습하기 위해서는 방대한 양의 데이터가 필요한 것이다.
아래 그림에서 회색 영역은 CNN 기반의 모델의 성능 범위를 나타낸다. 적은 데이터셋인 ImageNet에 대해 학습했을 때는 ViT의 성능이 CNN보다 낮지만, 대규모 데이터셋인 JFT-300M에 대해 학습했을 때는 ViT의 성능이 CNN을 능가하는 것을 확인할 수 있다.
출처: An Image is Worth 16X16 Words: Transformers for Image Recognition at Scale