2. Transformer Models for Image Classification

Transformer Models

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.

Data-Efficient Image Transformer (DeIT) [ICML 2021]

• Key Contributions: a teacher-student model for transformer which uses a distillation strategy to train the transformer on smaller datasets and achieve comparable performance to ViT on image classification tasks.
• Knowledge Distillation allows a smaller student network to learn from the soft labels (softmax labels) estimated by a larger teacher network

• Architecture

  • Teacher network: ResNet
  • Student network: Transformer (DeIT)
    • Similar to ViT, inputs a classification token and a sequence of image patch embeddings. Additionally, appends a distillation token to this sequence. All of these tokens are input to a multilayer transformer encoder with a stack of Multi-headed self-attention + Feed-Forward Network.
    • distillation token allows the student transformer network to use attention to learn from the teacher network
      • by using a convnet teacher network, inductive biases can be transferred to the student transformer network in a soft way. this results in a better performance than a transformer teacher.
      • soft distillation minimizes the KL-divergence between the softmax labels of the student and teacher networks
      • hard distillation minimizes the croos-entropy loss between the labels predicted by the student and teacher networks
  • class and distillation token vectors are learned from data through backprop during training and both interact with the image patch embeddings through self-attention
    • in the initial layers the class and distillation vectors are different but converge to being similar at the final layer
    • a linear output layer predicts the class and the distilled label
  • the model is trained on cross-entropy loss + distillation loss
    • distillation loss tries to ensure that the distillation token reproduces the hard labels output by the teacher network for the samples
    • cross-entropy loss tries to ensure that the class token matches the ground-truth labels for the samples
  • during inference, the predicted class can be a late fusion of the class and distilled label (for e.g. sum of the softmax labels)