PCAInit-ViT: Leveraging Data Geometry for Robust Vision Transformer Initialization

Abstract

Vision Transformers (ViTs) have emerged as a powerful alternative to convolutional networks for image classification, but their performance remains highly sensitive to weight initialization, especially in data-scarce or low-regularization settings. We introduce PCAInit-ViT, a novel, data-aware initialization method that uses principal component analysis (PCA) to initialize all trainable layers of a ViT, including patch embeddings, attention projections (QKV and output), MLP blocks, and the classification head. Unlike standard schemes such as Glorot or He, PCAInit-ViT aligns weights with the dominant directions of variation in the data by computing PCA from layer-wise activations after a forward pass over a small training subset. This strategy introduces an inductive bias aligned with the data distribution, partially compensating for ViTs’ lack of architectural priors like locality and translation equivariance. Experiments on CIFAR-10 with a ViT-Tiny architecture across three regularization regimes demonstrate consistent improvements in convergence speed and test accuracy. Our results suggest that PCAInit-ViT acts as an implicit regularizer and offers a lightweight, principled alternative to large-scale pretraining.