Paper: Read the Full Paper (PDF)

Authors: Chenghao Li, Chaoning Zhang

The environment configuration file has already been provided, and you can use the following script to create the required environment.

pip install -r requirements.txt

The success of Vision Transformer (ViT) in image recognition tasks has been well-documented. While ViT can capture global dependencies better than Convolutional Neural Networks (CNN), CNN's local characteristics are still valuable due to their resource efficiency. Recently, RetNet has demonstrated remarkable performance in language modeling, outperforming Transformers with explicit local modeling. This has led researchers to explore Transformers in the Computer Vision (CV) field.

This paper investigates the applicability of RetNet from a CNN perspective and introduces a customized RetNet variant for visual tasks. Similar to RetNet, we enhance ViT's local modeling by introducing a weight mask to the self-attention matrix. Our initial results with a learnable element-wise weight mask (ELM) show promise. However, ELM introduces additional parameters and optimization complexity.

To address this, our work proposes a novel Gaussian mixture mask (GMM) with only two learnable parameters, making it suitable for various ViT variants with adaptable attention mechanisms. Experiments on multiple small datasets illustrate the effectiveness of our Gaussian mask in enhancing ViTs at minimal additional cost in terms of parameters and computation.

Element-wise Conv Operation. If the mask is a circulant matrix with multi-diagonal properties, when such a matrix is added to the attention scores using the Hadamard product, it can be perceived as performing an operation akin to an element-wise convolution. The convolution kernel is the hidden feature vector, while the feature map is the weight mask. Examples of 1D and 2D cases are shown above.

Gaussian Mixture Mask (GMM) Attention From the perspective of CNN, by explicitly modeling and discovering two features of ViT using masks, we introduce the Gaussian Mixture Mask. We simply add the obtained GMM to the preceding softmax operation and validate it through experiments. Adding it before and after the softmax does not significantly affect the results. To maintain a structure similar to the original paper, we treat this operation as a straightforward masking operation to achieve the Gaussian mixture attention mechanism.

Building upon these insights, Gaussian Mixture Mask outperforms the Element-wise Learnable Mask significantly and minimizes the additional parameter and computation overhead to almost zero.

The following simple commands can be used to run experiments on your machine.

Quick Start

python main.py

Original model

python main.py --model [model_name]

GMM variants

python main.py --model [model_name] --is_GMM

We would like to express our great appreciation to the code authors of the pytorch-image-models and Vision Transformer for Small-Size Datasets repositories, for their great help to the machine learning community.