Vision transformers have recently made a breakthrough in computer vision showing excellent performance in terms of precision for numerous applications. However, their computational cost is very high compared to alternative approaches such as Convolutional Neural Networks. To address this problem, we propose a novel framework for image classification called RAViT based on a multi-branch network that operates on several copies of the same image with different resolutions to reduce the computational cost while preserving the overall accuracy. Furthermore, our framework includes an early exit mechanism that makes our model adaptive and allows to choose the appropriate trade-off between accuracy and computational cost at run-time. For example in a two-branch architecture, the original image is first resized to reduce its resolution, then a prediction is performed on it using a first transformer and the resulting prediction is reused together with the original-size image to perform a final prediction on a second transformer with less computation than a classical Vision transformer architecture. The early-exit process allows the model to make a final prediction at intermediate branches, saving even more computation. We evaluated our approach on CIFAR-10, Tiny ImageNet, and ImageNet. We obtained an equivalent accuracy to the classical Vision transformer model with only around 70% of FLOPs.
Figures & tables
Fig. 1 : Example of a RAViT architecture on a CIFAR-10 image.
Fig. 2 : Framework of our RAViT architecture with 3 branches. Note that the number of branches can be increased or decreased without changing radically the architecture
Fig. 3 : Performance analysis of different RAViT models with 2 branches used on CIFAR-10. (a) Accuracy for different layer numbers on CIFAR-10. (b) Number of FLOPs for different layer numbers on CIFAR-10.
Fig. 4 : Tiny ImageNet results. (a) Test accuracy in relation to FLOPs for different architectures. (b) Accuracy (written on top of each columns) and exit distributions for different values of the early-exit threshold applied to a 2-0-3 network.
Architecture
Accuracy
MFLOPs
% MFLOPs
3 layers ViT
82.3
83.16
75
4 layers ViT
84.3
110.88
100
5 layers ViT
85.3
138.61
125
our 1-3 (no EE)
84.9
89.99
81
our 1-3 (0.05 Eth )
84.6
80.81
72
our 1-3 (0.15 Eth )
82.6
67.72
61
TABLE I : Comparison of ViT with different number of layers and our models with several EE thresholds on Cifar-10.
Fig. 5 : Performance analysis of different RAViT trained on Tiny ImageNet with 3 branches and 3 layers on the third one. To be compared with a 4-layer classical ViT of 0.41 accuracy and 223.3 GFLOPs. (a) Accuracy for different models having 3 branch-3 layers. (b) GFLOPs for different models having 3 branch-3 layers.
Architecture
Accuracy
GFLOPs
% GFLOPs
3 layers ViT
38.8
11.52
75
4 layers ViT
41.0
15.36
100
6 layers ViT
41.4
23.04
150
our 2-0-3 (no EE)
40.7
12.00
78
our 2-0-3 (0.2 Eth )
40.4
10.98
71
our 2-0-3 (0.25 Eth )
39.1
9.72
63
TABLE II : Comparison of ViT with different number of layers and our models with several EE thresholds on Tiny ImageNet.
Fig. 6 : Accuracy (written on top of each columns) and exit distributions for different values of the early-exit threshold applied to different models on ImageNet. (a) RAViT 1-1-8 model. (b) RAViT 1-1-10 model.
Architecture
Accuracy
GFLOPs
% GFLOPs
8 layers ViT
71.08
23.26
67
10 layers ViT
72.80
29.07
83
12 layers ViT (ViT-B)
73.36
34.89
100
our 1-1-8 (no EE)
73.25
24.43
70
our 1-1-8 (1.0 Eth )
72.6
22.84
65
our 1-1-10 (no EE)
73.38
30.25
81
TABLE III : Comparison of ViT with different number of layers and our models on ImageNet.