A data efficient pyramid vision transformer for image classification

Abstract

Transformer-based models such as the vision transformer (ViT) have seen widespread use in computer vision tasks. However, while ViT demonstrates excellent performance, it requires a massive dataset and the process of computing self-attention between patches has quadratic complexity. To handle these challenges, the vision transformer needs to be data-efficient to train effectively on smaller datasets and its computational complexity should scale linearly with the number of image patches. In response, this paper introduces a novel data efficient pyramid vision transformer (DE-PVT), designed to train on limited datasets by utilizing a teacher-student approach and linear computational complexity relative to the number of patches, achieved through a linear spatial reduction mechanism. In this research, patches are extracted from images using soft-splitting mechanism to capture local continuity and fine-grained details of images. The teacher model is a customized lightweight convolutional neural network that imparts knowledge to the transformer-based student model for classifying images. Moreover, a weighted loss function is employed to compute the overall cross-entropy loss in the proposed model. The final classification of the test image is determined by considering both the teacher and student models, with the contribution of each model being mathematically defined. To validate the model, the DE-PVT framework was trained on well-known benchmark image datasets including ImageNet-1K (32 × 32), CIFAR-10 and CIFAR-100, achieving F1-scores of 90.92%, 97.76%, and 94.59% respectively. These results suggest that the proposed model could positively impact computer vision tasks, particularly where data availability and rapid computation are critical concerns.