Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Vision Transformer for Small-Size Datasets

Seung Hoon Lee, Seunghyun Lee, Byung Cheol Song

TL;DR

This work tackles the data inefficiency of Vision Transformers on small datasets by introducing two add-on modules that increase locality inductive bias. Shifted Patch Tokenization expands the receptive field of tokens by spatially shifting input patches, while Locality Self-Attention sharpens inter-token focus through diagonal masking and a learnable softmax temperature. Empirical results show consistent gains across ViT variants on Tiny-ImageNet and CIFAR datasets, with notable improvements on Swin, and modest improvements on mid-size ImageNet, reducing the performance gap with CNNs. The methods are lightweight, easily pluggable into existing ViTs, and enable learning from scratch on small datasets, accelerating practical deployment of ViTs in data-constrained settings.

Abstract

Recently, the Vision Transformer (ViT), which applied the transformer structure to the image classification task, has outperformed convolutional neural networks. However, the high performance of the ViT results from pre-training using a large-size dataset such as JFT-300M, and its dependence on a large dataset is interpreted as due to low locality inductive bias. This paper proposes Shifted Patch Tokenization (SPT) and Locality Self-Attention (LSA), which effectively solve the lack of locality inductive bias and enable it to learn from scratch even on small-size datasets. Moreover, SPT and LSA are generic and effective add-on modules that are easily applicable to various ViTs. Experimental results show that when both SPT and LSA were applied to the ViTs, the performance improved by an average of 2.96% in Tiny-ImageNet, which is a representative small-size dataset. Especially, Swin Transformer achieved an overwhelming performance improvement of 4.08% thanks to the proposed SPT and LSA.

Vision Transformer for Small-Size Datasets

TL;DR

This work tackles the data inefficiency of Vision Transformers on small datasets by introducing two add-on modules that increase locality inductive bias. Shifted Patch Tokenization expands the receptive field of tokens by spatially shifting input patches, while Locality Self-Attention sharpens inter-token focus through diagonal masking and a learnable softmax temperature. Empirical results show consistent gains across ViT variants on Tiny-ImageNet and CIFAR datasets, with notable improvements on Swin, and modest improvements on mid-size ImageNet, reducing the performance gap with CNNs. The methods are lightweight, easily pluggable into existing ViTs, and enable learning from scratch on small datasets, accelerating practical deployment of ViTs in data-constrained settings.

Abstract

Recently, the Vision Transformer (ViT), which applied the transformer structure to the image classification task, has outperformed convolutional neural networks. However, the high performance of the ViT results from pre-training using a large-size dataset such as JFT-300M, and its dependence on a large dataset is interpreted as due to low locality inductive bias. This paper proposes Shifted Patch Tokenization (SPT) and Locality Self-Attention (LSA), which effectively solve the lack of locality inductive bias and enable it to learn from scratch even on small-size datasets. Moreover, SPT and LSA are generic and effective add-on modules that are easily applicable to various ViTs. Experimental results show that when both SPT and LSA were applied to the ViTs, the performance improved by an average of 2.96% in Tiny-ImageNet, which is a representative small-size dataset. Especially, Swin Transformer achieved an overwhelming performance improvement of 4.08% thanks to the proposed SPT and LSA.
Paper Structure (23 sections, 10 equations, 6 figures, 7 tables)

This paper contains 23 sections, 10 equations, 6 figures, 7 tables.

Table of Contents

  1. INTRODUCTION
  2. RELATED WORK
  3. PROPOSED METHOD
  4. Preliminary
  5. Shifted Patch Tokenization
  6. Formulation
  7. Patch Embedding Layer
  8. Pooling Layer
  9. Locality Self-Attention Mechanism
  10. Diagonal Masking
  11. Learnable Temperature Scaling
  12. EXPERIMENT
  13. SETTING
  14. Environment and Dataset
  15. Model Configurations
  16. ...and 8 more sections

Figures (6)

  • Figure 1: Effect of the proposed method on the overall performance when learning Tiny-ImageNet from scratch. Throughput refers to how many images can be processed per unit of time. The stars and dots indicate after and before the proposed method are applied, respectively.
  • Figure 1: Various Shift Directions.
  • Figure 2: Architectures of the proposed SPT and LSA.
  • Figure 3: The learned temperature according to depth. Here, the red dashed line indicates the temperature of standard ViT.
  • Figure 4: Kullback–Leibler Divergence (KLD) of attention score distributions. The average KLDs were measured on Tiny-ImageNet.
  • ...and 1 more figures