Vision Transformer (1)

Transformer 관련 포스트

Transformer

Transformer 구현

Transformer Positional Encoding

---

 

Transformer 모델은 발표 이후에 자연어 처리 (NLP) 분야에서의 새로운 기준이 되어 왔습니다. 자연스럽게 self-attention 개념을 컴퓨터 비젼 분야에 적용하려는 시도가 많았는데요, Vision Transformer (ViT) 등장 이전에는 convolution 구조와 함께 제한적으로 사용되었고 ResNet 계열 모델을 필두로 한 CNN 을 완벽하게 대체할 수는 없었습니다. 구글에서 발표한 Vision Transformer 모델은 convolution 구조를 완전히 배제하고 Transformer 모델과 비슷하게 순수한 self-attention 만을 기반으로 하여 image classification benchmark 데이터셋에 대해 제일 높은 성능을 달성했습니다. 

Vision Transformer

Patch embedding

모델의 기본적인 구조는 Figure 1과 같습니다. 하나의 이미지가 여러 개의 패치로 나누어져 왼쪽 상단 패치를 시퀀스의 첫 원소, 오른쪽 하단 패치를 시퀀스의 마지막 원소로 하여 인코더의 입력 시퀀스로 삽입됩니다. 패치 크기가 $P$ 일때, 하나의 이미지 $x\in R^{H\times W\times C}$는 각 패치가 1차원의 텐서로 펼쳐져서 (flatten) 구성된 $x_p\in R^{N\times (P^2\cdot C)}$ 시퀀스가 됩니다. ($N$은 패치의 개수입니다.) 예를 들어 이미지 크기가 $H=W=224$ 이고 $P=16$ 일때 $14\times 14=196$ 개의 $16\times 16\times 3=768$ 차원의 텐서로 이루어진 시퀀스가 되는 것이죠. 이후에는 시퀀스의 각 요소 별로 임베딩을 위한 선형변환을 수행하고 (patch embedding) Transformer 모델과 마찬가지로 모델 전체의 차원은 $D$로 통일합니다. 따라서 $H=W=224, P=16$인 경우에 patch embedding 결과 $196\times 768$ 크기의 패치 임베딩 행렬이 도출됩니다.

패치 임베딩을 그림과 표현하면 Figure 2와 같습니다. 1) 입력 이미지를 패치 크기로 잘라낸 후 2) 선형변환을 통해 패치 임베딩을 완성합니다. 하지만 이 과정을 2차원 convolution 을 이용하면 하나의 과정으로 구현할 수 있습니다. 다음 코드와 같이 패치의 채널은 RGB 3이고 임베딩 차원은 768 이니 Pytorch 에서 in_channel=3, out_channel=768 로 설정한 후 kernel_size, stride 를 모두 패치 크기만큼 지정하면 됩니다.

# Input image [B, C, H, W]
x = torch.randn(1, 3, 224, 224)

# 2D conv
conv = nn.Conv2d(3, 768, 16, 16)
x = conv(x) # [B, 768, 14, 14]
x = x.reshape(B, -1, 196).transpose(1, 2) # [B, 196, 768]

Figure 2

 

Class token & Positional embedding

BERT 의 'class token' 과 마찬가지로 입력 시퀀스 맨 앞에 학습가능한 임베딩 파라미터를 붙여줍니다. ($z_0^0=x_{class}$) 이 'class token' 은 이것의 인코더 출력이 ($z_L^0$) 전체 이미지의 representation 을 나타내는 특별한 token 으로 최종단의 classification 을 위한 MLP layer 가 $z_L^0$에 달리게 됩니다. 대용량 데이터셋을 이용한 사전훈련 시에는 $z_L^0$에 2-layer MLP 가 달리게 되고 fine-tuning 시에는 1-layer 가 달리게 됩니다. 결과적으로 class token 을 위한 임베딩은 $1\times 768$ 크기의 텐서가 되고 이것이 패치 임베딩 행렬에 맨 왼쪽에 붙게 되어 인코더에 들어가는 최종적인 텐서의 크기는 $(196+1)\times 768$ 이 됩니다.

또한, Transformer 모델처럼 시퀀스 각 요소의 위치 정보를 알려주기 위해 positional embedding 텐서를 더해줍니다. Positional embedding 은 학습가능한 $197\times 768$ 크기의 텐서로 1차원 position embedding 입니다. 논문에서는 ViT 인코더가 이미지의 패치 레벨에 대해 동작하기 때문에 픽셀 레벨 입력에 대해서 공간 정보가 덜 중요해져 (224x224 $\rightarrow$ 14x14) 2차원 position 임베딩과 결과가 크게 차이가 없었다고 합니다. 중요한 점은 'class token' 과 positional 임베딩은 학습가능한 임베딩 텐서로 Figure 3과 같이 Pytorch 에서는 단순히 nn.Parameter 함수로 선언한 학습가능한 텐서가 됩니다.

Figure 3

 

Transformer encoder

ViT 에서의 Transformer encoder 는 기존 Transformer 대비 순서가 약간 변형된 Equation 1과 같은 순서로 진행됩니다. (기존 Transformer 에서 LayerNormalization 은 residual addition 이후에 적용되었습니다.) $z_0$은 입력 이미지에 대해 구성한 전체 임베딩이고 MSA (Multihead Self-Attention), MLP, LN (Layer Normalization) 이 Figure 1의 오른쪽 부분과 같이 구성되어 있습니다. 위에서 언급했듯이 최종적인 분류를 위한 representation 은 Equation 1(4) 처럼 'class token' 의 최종 아웃풋 $z_L^0$이 사용됩니다.

Equation 1

Multihead self-attention 관련해서는 기존 Transformer 모델과 거의 비슷하게 동작합니다. 입력 시퀀스 $z\in R^{N\times D}$라 할 때 각 헤드별 self-attention 은 Equation 2와 같이 동작하고 $D_h$는 각 헤드별 feature 차원으로 전체 $D$에서 헤드 개수를 나누어줍니다. 이후에는 Equation 3같이 기존 Transformer 모델의 feed-forward networks 와 같이 각 헤드별 출력 $D_h$를 이어붙인 이후에 선형변환하여 최종적으로 $D$ 차원으로 만듭니다.

Equation 2

Equation 3

인코더 전체 구조를 그림으로 표현하면 Figure 4와 같습니다. 입력으로 들어가는 임베딩은 'class token' / positional embedding 을 포함하여 $197\times 768$ 크기가 되고 최종 출력도 같은 크기의 텐서가 도출됩니다. 먼저 768 차원의 query, key, value 텐서를 구성하기 위해 2304=768*3 크기의 선형변환을 적용하고 12 multi-head self-attention 을 사용하므로 각 head 당 query, key, value 텐서의 크기는 768/12=64가 됩니다. 이후에는 MLP 블락이 적용됩니다. MLP 블락은 2-layer 와 BERT 에서 사용된 GELU (Gaussian Error Linear Unit) 비선형 활성화 함수가 사용됩니다.

Figure 4

Transformer 모델과 마찬가지로 Figure 4처럼 구성된 Transformer 인코더를 여러개 적층합니다. 첫번째 층의 출력이 두번째 층의 입력으로 들어가고 이 과정이 반복되며 마지막 인코더 출력에 대해 최종 분류를 위한 MLP 를 적용합니다. 전체적인 그림은 Figure 5와 같습니다.

Figure 5

Hybrid architecture

이미지의 패치에 대해 바로 적용하기도 하고 CNN의 feature map 에 대해서 ViT를 적용하는 것을 hybrid architecture 라 칭합니다. Hybrid architecture 에서는 입력 이미지가 아닌 입력 이미지가 CNN 을 거친 feature map 에 대하여 패치를 구성하게 되는데요, 이때 feature map 한 셀별로 입력 이미지의 일정 부분에 대한 정보를 담고 있고 해상도가 낮으니 $P=1$의 패치 크기를 가질 수 있게 설계합니다. 이러한 경우는 feature map 의 공간 차원을 펼치고 바로 $D$ 차원으로 선형변환 하면 되겟죠. 'class token' / positional embedding 은 똑같이 적용합니다.

 

Inductive bias

Transformer 등장 이후 inductive bias 라는 용어가 레딧이나 논문에 자주 등장하고 있습니다. Inductive bias 란 우리가 머신러닝 / 딥러닝을 수행할 때 많은 데이터에 대해 귀납적으로 문제를 풀고 싶어하고 그 문제를 더 잘풀기 위해 설계한 모델 / 목적함수의 추가적인 가정을 말합니다. 예를 들어 우리가 이미지 관련 문제를 풀고 싶을때는 CNN 모델을 사용하는데, 이는 CNN이 이미지의 국소적인 부분을 동시에 바라보게 만들면 이미지 관련 태스크를 잘 풀것이다라고 하는 inductive bias 가 삽입된 것이라 볼 수 있습니다. 마찬가지로 시계열 데이터를 처리할 때 사용하는 RNN / LSTM / 1d convolution 등의 모델도 시간 의존성을 반영하고자 하는 inductive bias 가 삽입된 것이죠. 반면 Transformer 같은 경우는 순수히 self-attention 기반으로 동작하기에 모델적인 측면의 inductive bias 가 상대적으로 부족하고 이로 인해 높은 성능을 위해서는 대용량의 데이터가 필요하다고 알려져 있습니다.

• 그렇다고 Transformer 모델에 inductive bias 가 없다고는 할 수 없을 것 같습니다. 애초에 positional embedding (fixed sinusoidal) 삽입도 그러하고 multihead self-attention 이후에 MLP 를 적용하는 것이 각 헤드 별 attention 결과를 커플링하여 feature 차원으로 프로젝션하자는 의도가 담겨져 있다고 볼 수 있기 때문입니다.

위에서 살펴본 ViT는 CNN에 비해 이미지 특화된 1) 국소적인 2차원 부분 구조, 2) translation equivariance 의 inductive bias 가 부족합니다. 2차원 이웃을 고려하는 부분은 처음 입력 이미지를 패치로 구분할 때와 실험에서 살펴볼 fine-tuning 시에 사용되고 삽입한 position embedding 은 초기화시 패치별 위치에 대한 아무 정보를 제공하지 않고 처음부터 학습되어야만 합니다. 추후 실험 포스트에서 살펴보겠지만 ViT 또한 작은 데이터셋에 대해 스크래치로 훈련시키면 성능이 잘 나오지 않아 대용량의 훈련 데이터셋으로 설계상 작은 inductive bias 를 넘어서려는 시도를 하게 됩니다. 

 

---

Vision Transformer (2)

Vision Transformer (3) - Attention Map

Vision Transformer (4) - Pytorch 구현