Structured Initialization for Vision Transformers

TL;DR

This paper introduces a novel initialization method for Vision Transformers that embeds CNN-like inductive biases through initialization alone, improving performance on small datasets while maintaining scalability.

Contribution

It proposes a new initialization strategy for ViTs that enhances small-data performance without architectural changes, outperforming existing heuristics.

Findings

Significantly outperforms standard ViT initialization on small/medium benchmarks.

Maintains competitive performance on large-scale datasets like ImageNet-1K.

Easily integrates into various transformer architectures with consistent gains.

Abstract

Convolutional Neural Networks (CNNs) inherently encode strong inductive biases, enabling effective generalization on small-scale datasets. In this paper, we propose integrating this inductive bias into ViTs, not through an architectural intervention but solely through initialization. The motivation here is to have a ViT that can enjoy strong CNN-like performance when data assets are small, but can still scale to ViT-like performance as the data expands. Our approach is motivated by our empirical results that random impulse filters can achieve commensurate performance to learned filters within a CNN. We improve upon current ViT initialization strategies, which typically rely on empirical heuristics such as using attention weights from pretrained models or focusing on the distribution of attention weights without enforcing structures. Empirical results demonstrate that our method significantly outperforms standard ViT initialization across numerous small and medium-scale benchmarks, including Food-101, CIFAR-10, CIFAR-100, STL-10, Flowers, and Pets, while maintaining comparative performance on large-scale datasets such as ImageNet-1K. Moreover, our initialization strategy can be easily integrated into various transformer-based architectures such as Swin Transformer and MLP-Mixer with consistent improvements in performance.