2. Transformer Models for Image Classification

Transformer Models for Image Classification

Vision Transformer (ViT) [ICLR 2021]

• Key Contributions: a pure transformer model without any conv layers that extracts features from a sequence of image patches and encodes them for image classification tasks.

• vision transformers perform well on image classification tasks

• ViT has lower degree of inductive bias, such as locality and translational invariance, compared to convnets, which results in lower performance than conv-nets for smaller datasets. But training on large datasets compensates for the lack of inductive bias and helps transformers perform well on classification tasks.

• Use a multi-layer transformer encoder network consisting of multiple MHSA and FFN layers

• ViT treats an image as a sequence of tokens and inputs these tokens to multiple transformer layers for classification.

  • a H X W image is divided into a sequence of N=H*W/p**2 patches of size p. N is the sequence length of the transformer input.
  • ViT divides an image $x \in R^{HxWxC}$ into 16x16 patches.
  • each image patch of size P is embedded into a D-dim vector through linear projection, creating N x P X P X D patch embeddings for the entire image, where N is the number of patches
  • a cls token is prepended to the patch sequence to generate a seq of length N+1.
  • learnable 1-D positional encoding is added to each of the N+1 tokens. Using 2D positional embedding does not produce significant performance gains.

• ViT Encoder

  • the N+1 token vectors are input to the transformer encoder, which generates N+1 encoded vectors as output.
  • encoder consists of L alternating layers of MHSA and MLP blocks
    • $z_0 = [x_{class}; x^0_pE,\ldots,x^N_pE] + E_{pos}, E_{pos} \in R^{(N+1)\times D}, E \in R^{(P^2*C) \times D} $.
    • $z_{l}^{sa} = MHSA(LN(z_{l-1})) + z_{l-1}, l= 1,\ldots,L$ (MHSA + residual connection)
    • $z_{l} = MLP(LN(z^{sa}_{l})) + z^{sa}_l, l= 1,\ldots,L$ (MLP + residual connection)
    • $y = LN(z_{L}^0)$ (final image rep that goes into the classification head)
    • layernorm is applied before each block and residual connections after each block
      • normalization is performed before MHSA

• hybrid architecture:

  • transformer can be combined with CNN by using the feature map outout by the CNN as the input image for patching instead of the raw image.

• ViT vs. CNN

  • in CNN, inductive bias and 2-dim neighborhood is inherent in each layer. In ViT, only the MLP layer retains locality and translational equivariance while the MSA layer operate on global features (all-to-all).
  • in ViT, the inductive bias about 2-dim image neighborhood is used
    • initially during splitting the input image into patches
    • then when adapting the position embeddings during fine-tuning/inference for different image resolutions. In this case, given the same patch size, the sequence length would differ from that used during pretraining, so the number of position embeddings will be different. So 2D interpolation of the pre-trained position embeddings are performed according to their location in the original image.
  • The patch embeddings do not store information about the 2D patch positions during initialization and all spatial positions between the patches are learned during training.