MSCViT: A Small-size ViT architecture with Multi-Scale Self-Attention Mechanism for Tiny Datasets

TL;DR

MSCViT is a compact Vision Transformer architecture with multi-scale self-attention and convolutional features, designed specifically for small datasets, achieving high accuracy without large-scale pre-training.

Contribution

Introduces MSCViT, a parameter-efficient ViT variant with multi-scale attention and wavelet convolution, optimized for tiny datasets and reducing computational costs.

Findings

Achieves 84.68% accuracy on CIFAR-100 without pre-training.

Uses wavelet convolution for local feature extraction.

Reduces model parameters and FLOPs compared to original ViT.

Abstract

Vision Transformer (ViT) has demonstrated significant potential in various vision tasks due to its strong ability in modelling long-range dependencies. However, such success is largely fueled by training on massive samples. In real applications, the large-scale datasets are not always available, and ViT performs worse than Convolutional Neural Networks (CNNs) if it is only trained on small scale dataset (called tiny dataset), since it requires large amount of training data to ensure its representational capacity. In this paper, a small-size ViT architecture with multi-scale self-attention mechanism and convolution blocks is presented (dubbed MSCViT) to model different scales of attention at each layer. Firstly, we introduced wavelet convolution, which selectively combines the high-frequency components obtained by frequency division with our convolution channel to extract local features. Then, a lightweight multi-head attention module is developed to reduce the number of tokens and computational costs. Finally, the positional encoding (PE) in the backbone is replaced by a local feature extraction module. Compared with the original ViT, it is parameter-efficient and is particularly suitable for tiny datasets. Extensive experiments have been conducted on tiny datasets, in which our model achieves an accuracy of 84.68% on CIFAR-100 with 14.0M parameters and 2.5 GFLOPs, without pre-training on large datasets.