Efficient Training of Visual Transformers with Small Datasets

TL;DR

This paper introduces a self-supervised training task for Visual Transformers that enhances their performance on small datasets by encouraging learning of spatial relations, making training more robust and improving accuracy.

Contribution

The paper proposes a novel self-supervised task that can be integrated with existing Visual Transformers to improve their data efficiency and robustness in small dataset regimes.

Findings

Self-supervised task improves VT accuracy on small datasets

Method is architecture-agnostic and easy to implement

Significant accuracy gains demonstrated across multiple datasets

Abstract

Visual Transformers (VTs) are emerging as an architectural paradigm alternative to Convolutional networks (CNNs). Differently from CNNs, VTs can capture global relations between image elements and they potentially have a larger representation capacity. However, the lack of the typical convolutional inductive bias makes these models more data-hungry than common CNNs. In fact, some local properties of the visual domain which are embedded in the CNN architectural design, in VTs should be learned from samples. In this paper, we empirically analyse different VTs, comparing their robustness in a small training-set regime, and we show that, despite having a comparable accuracy when trained on ImageNet, their performance on smaller datasets can be largely different. Moreover, we propose a self-supervised task which can extract additional information from images with only a negligible computational overhead. This task encourages the VTs to learn spatial relations within an image and makes the VT training much more robust when training data are scarce. Our task is used jointly with the standard (supervised) training and it does not depend on specific architectural choices, thus it can be easily plugged in the existing VTs. Using an extensive evaluation with different VTs and datasets, we show that our method can improve (sometimes dramatically) the final accuracy of the VTs. Our code is available at: https://github.com/yhlleo/VTs-Drloc.