An Image is Worth 16x16 Words:Transformers for Image Recognition at Scale(ViT)

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (ViT) 를 읽고 정리한 글입니다.


An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Overview

  • image classification task에 기존의 transformer 모델을 이용한다.
    • transformer의 장점들을 사용할 수 있음.
    • simple, computational efficiency, scalability
  • Vision Transformer(ViT)
    • 원본 이미지를 patches들로 split한다.
      • 이 때 이미지 패치들을 NLP에서의 token(word)와 같은 역할임.
    • 이 패치들의 sequence of linear embedding을 Transformer의 input으로 feed한다.
  • 크지않은 데이터 셋에서는 ResNet보다 약간 낮은 정확도
    • Transformer는 CNN과 다르게 translation equivariance, locality 같은 inductive biases(=weight sharing)이 없기 때문에 데이터셋이 적으면 generalize되기 어려움.
  • 하지만 large scale 데이터에서는 CNN의 inductive bias를 능가함.
    • image recognition benchmarks에서 여러 SOTA 달성함.

제안한 모델이 최근 연구 중에서는 iGPT 논문과 유사하다고 함.

  • reducing image resolution and color space한 후 transformer를 적용

Methods

(Explain one of the methods that the thesis used.)

  • Model overview 스크린샷 2020-11-01 오전 11 51 34

  • Vision Transformer(ViT)

    • image를 patch 단위로 잘라서 flatten시킨 후, 그것을 linear projection하여 encoder에 feed한다. image

    • Equation 스크린샷 2020-11-01 오후 12 32 37

    • (Eq.1) trainable linear projection은 flatten된 patch들을 D dimension에 mapping시킨다.
    • (Eq.4) BERT의 [class] 토큰처럼, embedded patch의 sequence(image representation y) 전에 learnable embedding을 추가한다. (prepend)
      • both during pre-training and fine-tuning, a classification head is attached to $\left({z}_{L}^{0} \right)$
    • Position embeddings
      • positional information을 유지하기 위해 patch embedding에 붙여짐. (자세한건 Appendix.D.3)
      • standard learnable 1D position embedding 사용
      • 이렇게 position embedding까지 더해진 sequence of embedding vector들은 encoder의 input이 됨.
    • Encoder
      • alternating layers of multiheaded self-attention(MSA)
      • MLP blocks
        • contains 2 layer with a GELU non-linearity
        • MLP 부분 code
        1
        2
        3
        4
        5
        6
        7
        8
        9
        10
        11
        12
        13
        14
        15
        16
        17
        18
        19
        20
        21
        22
        23
        24
        25
        26
        27
        28
        29
        30
              class MlpBlock(nn.Module):
              """Transformer MLP / feed-forward block."""
        
          def apply(self,
                      inputs,
                      mlp_dim,
                      dtype=jnp.float32,
                      out_dim=None,
                      dropout_rate=0.1,
                      deterministic=True,
                      kernel_init=nn.initializers.xavier_uniform(),
                      bias_init=nn.initializers.normal(stddev=1e-6)):
              """Applies Transformer MlpBlock module."""
              actual_out_dim = inputs.shape[-1] if out_dim is None else out_dim
              x = nn.Dense(
                  inputs,
                  mlp_dim,
                  dtype=dtype,
                  kernel_init=kernel_init,
                  bias_init=bias_init)
              x = nn.gelu(x)
              x = nn.dropout(x, rate=dropout_rate, deterministic=deterministic)
              output = nn.Dense(
                  x,
                  actual_out_dim,
                  dtype=dtype,
                  kernel_init=kernel_init,
                  bias_init=bias_init)
              output = nn.dropout(output, rate=dropout_rate, deterministic=deterministic)
              return output
        
      • Layernorm(LN) is applied before every block
      • residual connections after every block
    • Hybrid Architecture
      • patch embedding projection E (Eq.1)이 CNN feature map으로 대체될 수 있다.
        • 즉, ResNet과 같은 CNN구조의 모델을 가지고, 2D feature map 중 하나를 1D로 flatten시킨 후 transformer dimension에 projecting 시킴.
        • 위에서 만들어진 sequence에 classification input embedding, position embedding를 추가시켜 encoder에 input으로써 feed 시킬 수 있음.

Fine-Tuning and Higher Resolution

  • large dataset으로 pre-training하고, smaller downstream task에 대해 fine-tune 하려고 함.
  • 이를 위해서 pre-trained prediction head를 지우고, 0으로 initializedgks D x K feedforward layer를 추가함.
    • K : # of downstream classes
  • pre-training 보다 높은 resolution으로 fine-tuning하는 것은 beneficial할 때도 있음.
  • higher resolution을 feed하게 되면, patch size는 동일하므로 sequence length가 길어짐
  • ViT는 임의적인 sequence length를 다룰 수 있음(메모리 제약에 따라서)
  • 하지만 pre-trained position embedding이 의미 없어질 수 있음
    • 원본 이미지의 location에 따라 pre-trained position embedding의 2D interpolation을 수행함.
  • 위와 같은 resolution adjustment와 patch extraction은 이미지의 2D 구조에 대해 inductive bias를 manually ViT에 주입시키는 유일한 포인트임.

Code

References

(References for your additional studies)

https://jeonsworld.github.io/vision/vit/