[논문] An Image Is Worth 16x16 Words: Transformers For Image Recognition At Scale 설명, 정리

• Introduction
  
  Self attention 기반의 구조들, 대표적으로 transformer은 NLP에서 좋은 성능을 보인다. Transformer는 연산 효율성이 좋아 매우 큰 크기의 모델들을 학습시킬 수 있었다. 하지만 컴퓨터 비젼 분야에서는, convolutional 구조들이 아직 우세하다. NLP 분야에서의 성적을 보고, CNN 구조를 self attention과 결합하거나 CNN을 완전히 대체하려는 연구들이 진행되고 있다.
  본 논문에서는 전형적인 transformer 모델에 거의 수정을 하지 않고 이미지를 직접 적용하는 방법을 실험한다. 이를 위해 이미지를 패치로 나누어 이 패치들의 linear embedding의 sequence를 transformer의 입력으로 사용한다. 이미지 패치들은 NLP에서 단어 토큰들과 동일히 여겨지고, 모델을 supervised image classification으로 훈련시켰다.
  중간 정도 크기의 데이터셋(ImageNet 등)으로 강한 regularizatoin 없이 훈련시켰을 때, ResNet을 사용한 결과보다 약간 낮은 결과가 나온다. 이는 트랜스포머가 CNN에는 있는 inductive bias가 없기 때문이다. 하지만 큰 데이터셋(14M~300M 이미지)에서 학습할 때는 다르다. 큰 규모의 학습이 inductive bias를 뛰어넘는다. VIT는 충분한 크기로 pretrain하고 더 적은 데이터포인트를 가진 task로 tranfer 했을 때 최고의 결과를 보였다.
  
  
  
• Related Work
  
  큰 규모의 트랜스포머 기반 모델들은 큰 corpora(말뭉치)에서 pretrain하고 주어진 task로 finetune한다. naive하게 self attention을 생각해 보면 각 픽셀을 서로 다른 픽셀들에게 적용할 수 있다. 하지만 이는 quadratic cost가 들고, 현실적인 input size에 적용할 수 없다. 기존에 제시되었던 많은 모델들은 attention 구조에 특화되어 computer vision task에 대한 결과를 보였었다. 하지만 기존의 방법들은 hardware accelerator을 위해 복잡한 기법이 필요하다. 또한 self attention과 CNN을 결합하는 방법들도 제안되었다.
  
  
  
• Method
  
  모델 디자인에는 기존 트랜스포머를 최대한 따른다. 이 때문에 NLP 외의 분야에도 자유롭게 사용할 수 있다.
  
  

  

  Fig 1

  
  트랜스포머는 토큰 embedding의 1D sequence를 입력받는다. 2D 이미지를 다루기 위해, 이미지 \(x\in R^{H\times W\times C}\)를 flatten된 2D 패치의 sequence인 \(x_{p} \in R^{N\times (P^{2}\cdot C)}\)로 reshape한다. \((H,W)\)는 원본 이미지의 해상도이고, \(C\)는 채널의 수, \((P,P)\)는 각 이미지 패치의 해상도이다. \(N=HW/P^{2}\)은 패치들의 수다. 트랜스포머는 모든 layer들에 대해 일정한 latent vector size D를 사용한다. 따라서 trainable linear projection (1)을 사용해 패치들을 flatten해서 D차원으로 map한다. (1)에서 \(x_{p}^{1}\)은 patch sequence인 \(x_{p}\)의 첫 패치를 나타낸다(\(x_{p}^{1} \in R^{1\times (P^{2}\times C)}\)). 그리고 이 output을 patch embedding이라 부른다.$$ z_{0} = [x_{class};x_{p}^{1}E;x_{p}^{2}E;...;x_{p}^{N}E]\ +\ E_{pos},\ \ \ \ E\in R^{(P^{2}\cdot C)\times D},E_{pos}\in R^{(N+1)\times D}\ \ \ \ (1) $$
  BERT의 [class] 토큰과 유사하게, embed된 패치들의 sequence에 learnable embedding을 덧붙인다(\(z_{0}^{0} = x_{class}\)). 트랜스포머 인코더 \(z_{L}^{0}\)의 output에서의 state는 image representation y (4)의 역할을 한다.$$ y = LN(z_{L}^{0})\ \ \ \ (4) $$
  Pretraining과 fine tuning 모두에서, classification head는 \(z_{L}^{0}\)에 부착되어 있다. Classification head는 pretraining 단계에서는 MLP로 하나의 hidden layer로 적용되고, fine tuning 시에는 하나의 linear layer로 적용된다.
  Position embedding은 위치정보를 알려주기 위해 patch embedding에 더해진다. 본 논문에서는 standard learnable 1D position embedding을 사용한다. 이로써 나오는 embedding vector의 sequence는 인코더에 입력된다.
  
  트랜스포머 인코더는 multiheaded self attention(MSA)와 MLP 블록((2),(3))이 번갈아 나타난다.$$ z_{l}' = MSA(LN(z_{l-1}))+z_{l-1},\ \ \ \ l = 1...L\ \ \ \ (2) $$$$ z_{l} = MLP(LN(z_{l}'))+z_{l}',\ \ \ \ l = 1...L\ \ \ \ (3) $$
  모든 블록 앞에는 LayerNorm(LN)이 가해지고, 매 블록 뒤에 residual connection이 위치한다(Fig 1). MLP는 GELU non-linearity를 두 layer 포함한다.
  
  Inductive Bias : VIT는 CNN보다 image specific inductive bias를 훨씬 덜 가지고 있다. CNN에서는 locality, 2dim neighborhood 구조, translation equivariance는 전체 모델의 각 layer에 적용된다. VIT에서는 MLP층만 local하고 translationally equivariant하다(반면 self attention layer는 global하다). 2 dimension neighborhood 구조는 매우 드물게 사용된다. 모델의 초기에 이미지를 patch로 자를 때와, 다른 해상도를 가진 이미지의 position embedding을 fine tune하는데 사용된다. 그것 외에는 초기화 할 때의 position embedding은 2d 위치 정보를 담고있지 못하고, 패치들 간의 공간적 관계는 scratch부터 학습되어야 한다.
  
  Hybrid Architecture : raw 이미지 패치를 사용하는 것 대신, input sequence는 CNN의 feature map에서부터 만들어질 수 있다. Patch embedding projection E는 CNN feature map에서 얻은 패치들에 적용된다. 특별한 경우에는 spatial size 1x1을 가질 수 있는데, 이러면 input sequence는 단순히 feature map의 spatial dimension을 flatten하고 트랜스포머의 차원에 project함으로써 얻어진다.
  
  VIT는 큰 데이터셋으로 pretrain하고, 작은 downstream task로 fine tune한다. 이를 위해, pre-trained prediction head를 제거하고 0으로 초기화된 D x K feedforward layer을 붙인다. 이때 K는 downstream class의 수를 의미한다. Pretrain하는 것 보다 더 높은 해상도에서 finetune하는 것이 좋다. 높은 해상도의 이미지를 feed할 때, 패치 크기를 동일하게 유지함으로써 더 효과적인 sequence length를 만든다. VIT는 메모리가 뒷받침해주는 한 임의의 sequence 길이를 다룰 수 있지만, pretrain된 position embedding은 더이상 의미가 없을 수 있다. 따라서 pre-train된 position embedding의 2D interpolation을 진행한다. 이런 해상도 조절과 패치 추출은 VIT에 이미지의 2D 구조가 수동적으로 입력되었을 때와 관련된 inductive bias다.
  
  
  
• Experiments
  
  본 논문에서는 ResNet, Vision Transformer(ViT), hybrid의 학습 능력을 실험해 보았다. 다양한 크기의 데이터셋으로 pretrain하고 다양한 task를 실험해 보았다. Pretrain하는 cost는 ViT가 효율적이였다. 또한 self supervision을 사용한 실험도 진행해 보았다.
  Pre-train을 위해서 ILSVRC-2012(1K class,1.3M images), ImageNet-21K(21K class,14M images), JFT(18K class,303M high resolution images)를 사용했다. ViT는 BERT에 사용된 설정을 기반으로 사용했다(Table 1).
  
  

  

  
  'Base'와 'Large' 모델은 BERT에서 그대로 사용했고, 추가적으로 더 큰 'Huge' 모델을 추가했다. ViT-L/16의 경우는 'Large' 변형을 사용했고, 16x16 입력 패치 크기를 사용했다는 뜻이다. 트랜스포머의 sequence 길이는 패치 크기의 제곱에 반비례하기 때문에, 작은 패치 크기를 가진 모델들은 연산 비용이 많이 든다.
  CNN의 baseline으로는 ResNet을 사용했고, Batch Normalization layer을 Group Normalization으로 바꾸고, standardized convolution을 사용했다. 이러한 수정들은 transfer을 개선하는 효과가 있고, 이 수정된 모델을 ResNet(BiT)라 부른다.
  Hybrid를 위해, ViT에 중간 feature map을 feed하고, 한 픽셀짜리 패치 크기를 사용한다. 다른 sequence 길이들로 실험하기 위해, 일반적인 ResNet 50의 stage 4의 출력을 사용하기도 하고(i), stage 4를 제거하고 stage 3과 동일한 수의 layer을 배치시켜(layer의 수 유지) 확장된 stage 3의 출력을 사용하기도 한다(ii). (ii)는 4배 더 긴 sequence 길이와 더 비싼 ViT 모델을 만든다.
  Downstream dataset의 결과는 few-shot이나 fine-tuning accuracy로 보고한다. Fine-tuning accuracy는 각 데이터셋으로 각 모델을 fine-tune한 뒤 성능으로 보고한다. Few shot accuracy는 training image의 subset representation을 \(\left \{-1,1\right \}^{K}\) target vector로 map하는 least square regression problem을 풀어나감으로써 얻어진다.
  
  우선 본 논문에서 제시한 가장 큰 모델인 ViT-H/14와 ViT-L/16을 SOTA CNN과 비교한다. 첫 비교는 large ResNet으로 supervised transfer learning을 수행하는 Big Transfer(BiT)다. 둘째는 Noisy Student로, Efficient Net으로 ImageNet과 JFT-300M으로 semi-supervised learning을 통해 얻은 큰 EfficientNet이다. 논문을 쓴 시점에서는 Noisy Student가 SOTA였다고 한다. 비교결과가 Table 2에 제시되어있다. 비교적 작은 모델인  JFT-300M으로 pre-train된 ViT-L/16은 BiT-L의 성능을 능가했고, 더 큰 모델인 ViT-H/14는 성능을 더 끌어올리고, 특히 더 까다로운 데이터셋에 대해 좋은 성능을 보였다. 이 모델들은 기존의 SOTA보다 적은 연산을 필요로 했다.
  
  

  

  Table 2

  
  
  ResNet보다 vision에 대해 적은 inductive bias를 가지는 ViT에서, 데이터셋의 크기는 얼마나 중요한지를 알아보기 위해 두가지 실험을 진행해보았다.
  우선, ViT모델을 증가하는 크기의 데이터셋으로 pre-train시킨다. 작은 데이터셋에서 성능을 올리기 위해, 세가지 regularization 파라미터를 optimize한다(weight decay, dropout, label smoothing). Fig 3에서 보면, 가장 작은 데이터셋인 ImageNet에서 finetune한 결과를 볼 수 있다. Fig 3에서 알 수 있듯, ViT-Large 모델들은 regularization에도 불구하고 ViT-Base 모델들보다 성능이 나오지 않는다. JFT-300M을 사용할 때만, 큰 모델의 이점을 잘 살릴 수 있었다. 작은 데이터셋(ImageNet)에서는 BiT CNN이 ViT보다 성능이 좋았지만, 큰 데이터셋에는 ViT가 성능이 더 좋아졌다.
  
  

  

  
  둘째로, Full JFT-300M 데이터셋 뿐 아니라 9M, 30M, 90M 랜덤 subset으로도 모델을 train시켜 보았다. 작은 subset이라고 추가적인 regularization을 수행하지 않았고, 모든 세팅에서 동일한 하이퍼파라미터를 사용했다. 이를 통해 regularization의 효과가 아닌 내재된 모델 property를 평가하려 했다. Fig 4를 보면, ViT는 ResNet보다 적은 데이터셋에 대한 연산비용에 overfit했다. 예를 들어, ViT-B/32는 9M subset에서는 ResNet50보다 성능이 안좋았지만, 90M+ subset에서는 성능이 더 좋아졌다. 이를 통해 작은 데이터셋에는 convolutional inductive bias가 유용하지만, 큰 데이터셋에는 데이터로부터 직접 relevant pattern을 학습하는 것이 유용하다는 것을 확인할 수 있다.
  
  ViT의 작동 원리를 확인하기 위해 내부에서의 표현을 분석해 보았다. ViT의 첫 layer은 flatten된 패치를 lower dimension 공간으로 linearly project한다(1). Fig 7-Left는 학습된 embedding 필터의 주요 compoenet를 보여준다. Component들은 각 패치들 안에서 fine structure의 저차원 representation을 나타내는 그럴듯한 basis function을 닮았다. Fig 7-Center은 모델이 이미지 내의 거리를 encode하는 법을 position embedding과 유사하게 학습한다는 것을 보여준다. 예를 들어, 가까운 패치들은 유사한 position embedding을 가지는 경향을 보인다. 또한, 동일한 row나 column에 있는 패치들은 유사한 embedding을 보인다. 마지막으로 sinusoidal 구조는 더 큰 grid에서 명백하다. Position embedding은 2D 이미지 배치를 표현하도록 학습한다는 것은 수작업된 2D aware embedding이 왜 발전을 보이지 않는지 설명한다.
  
  

  

  
  Self Attention은 ViT가 가장 낮은 layer에서도 이미지 전체에서 정보를 수합하도록 한다. 그리고 어느 정도로 네트워크가 이 능력을 활용하는지 조사한다. 이는 attention weight를 기반으로 해서 이미지 공간에서 어느 정도로 정보가 통합되는지 평균 거리를 구함으로써 알 수 있다(Fig 7-Right). 이를 보면 어떤 head들은 lowest layer에도 등장하며 전체적으로 정보를 통합하려는 능력이 실제로 모델에 의해 쓰이고 있음을 확인할 수 있다. 다른 attention head들은 low layer에서 지속적으로 작은 attention distance를 가지고 있다. 이렇게 매우 localize된 attention은 트랜스포머 전에 ResNet을 적용하는 hybrid model에서 더욱 확연하고, 이는 이것들이 CNN의 convolutional layer와 유사한 함수로 작용한다는 것을 제시한다.
  
  트랜스포머의 높은 성능은 완벽한 확장성 뿐 아니라 큰 규모의 self supervised pre-training에서 온다. 본 논문에서는 self supervision을 위해 masked patch prediction에 대한 사전 탐구를 진행했다. 이는 BERT에서 제시된 masked language modeling task와 유사하게 진행된다.