Three things everyone should know about Vision Transformers

TL;DR

This paper examines practical design choices to improve Vision Transformers (ViTs) in terms of efficiency and transferability. It demonstrates three key ideas: (i) parallelizing residual blocks to create wider/shallower ViTs without increasing parameter count or FLOPs, (ii) fine-tuning only the multi-head self-attention (MHSA) layers to adapt to higher resolutions and new tasks, and (iii) introducing a hierarchical MLP (hMLP) patch stem that enables Bert-like masked self-supervised pretraining with patch masking. Experiments on ImageNet-1k, ImageNet-V2, and six transfer datasets show competitive accuracy with reduced memory and compute, and effective self-supervised pretraining with patch masking. The findings offer practical guidance for deploying ViTs efficiently and for integrating patch-based self-supervised learning in vision models.

Abstract

After their initial success in natural language processing, transformer architectures have rapidly gained traction in computer vision, providing state-of-the-art results for tasks such as image classification, detection, segmentation, and video analysis. We offer three insights based on simple and easy to implement variants of vision transformers. (1) The residual layers of vision transformers, which are usually processed sequentially, can to some extent be processed efficiently in parallel without noticeably affecting the accuracy. (2) Fine-tuning the weights of the attention layers is sufficient to adapt vision transformers to a higher resolution and to other classification tasks. This saves compute, reduces the peak memory consumption at fine-tuning time, and allows sharing the majority of weights across tasks. (3) Adding MLP-based patch pre-processing layers improves Bert-like self-supervised training based on patch masking. We evaluate the impact of these design choices using the ImageNet-1k dataset, and confirm our findings on the ImageNet-v2 test set. Transfer performance is measured across six smaller datasets.

Figures (8)

• Figure 1: Impact of the parallelism on performance for a given model size (ViT-S36, -S60 and -B36) and 1--4 parallel branches.
• Figure 2: Impact of the training on parallel and sequential models.
• Figure 3: Comparison of parallel models with more blocks with models with a higher working dimensionality. L24$\times$1, B36$\times$1 and B18$\times$2 trained with LS.
• Figure 4: Fine-tuning the weights of the self-attention layer only (middle panel) leads to savings during fine-tuning in peak memory usage and computational cost. It also leads to important savings in the number of parameters when a model is fine-tuned for multiple resolutions or multiple downstream classification tasks.
• Figure 5: Design of our hMLP-stem: we start from subpatches and progressively merge them with linear layers interleaved by GELU non-linearities. The design of our stem is such that the patches are processed independently. Hence it commutes with masking.