Implementing Stand-Alone Self-Attention in Vision Models using Pytorch (13 Jun 2019)

• Stand-Alone Self-Attention in Vision Models paper
• Author:
  • Prajit Ramachandran (Google Research, Brain Team)
  • Niki Parmar (Google Research, Brain Team)
  • Ashish Vaswani (Google Research, Brain Team)
  • Irwan Bello (Google Research, Brain Team)
  • Anselm Levskaya (Google Research, Brain Team)
  • Jonathon Shlens (Google Research, Brain Team)
• Awesome :)

Method

• Attention Layer
  
  
  • Equation 1:
    
    
    
    
• Relative Position Embedding
  
  

  • The row and column offsets are associated with an embedding and respectively each with dimension . The row and column offset embeddings are concatenated to form . This spatial-relative attention is now defined as below equation.

  • Equation 2:
    
    

  • I refer to the following paper when implementing this part.

    • Attention Augemnted Convolutional Networks paper

1. Replacing Spatial Convolutions
   - A 2 × 2 average pooling with stride 2 operation follows the attention layer whenever spatial downsampling is required. - This work applies the transform on the ResNet family of architectures. The proposed transform swaps the 3 × 3 spatial convolution with a self-attention layer as defined in Equation 3.
2. Replacing the Convolutional Stem
   - The initial layers of a CNN, sometimes referred to as the stem, play a critical role in learning local features such as edges, which later layers use to identify global objects. - The stem performs self-attention within each 4 × 4 spatial block of the original image, followed by batch normalization and a 4 × 4 max pool operation.

Experiments

Setup

• Spatial extent: 7
• Attention heads: 8
• Layers:
  • ResNet 26: [1, 2, 4, 1]
  • ResNet 38: [2, 3, 5, 2]
  • ResNet 50: [3, 4, 6, 3]

Usage

Requirements

• torch==1.0.1

Todo

• Experiments
• IMAGENET
• Review relative position embedding, attention stem
• Code Refactoring

Reference

• ResNet Pytorch CIFAR
• ResNet Pytorch
• Thank you :)