Decorrelation Speeds Up Vision Transformers

TL;DR

This paper tackles the high computational cost of Masked Autoencoder pre-training for Vision Transformers in low-label settings by applying Decorrelated Backpropagation (DBP) to the encoder. DBP learns per-layer decorrelation matrices to reduce input correlations, accelerating convergence while maintaining stability, and is fused into the model so downstream fine-tuning sees decorrelated weights. On 10% ImageNet-1K pre-training and 10% ADE20K fine-tuning, DBP-MAE achieves a 21.1% reduction in wall-clock time to reach peak performance and a 21.4% drop in carbon emissions, with an additional 1.1 point gain in ADE20K mIoU; industrial data also show consistent gains. The approach offers a practical path to faster, greener large-scale ViT pre-training reusable across industrial vision tasks, with potential extensions to medical imaging and automated DBP component selection.

Abstract

Masked Autoencoder (MAE) pre-training of vision transformers (ViTs) yields strong performance in low-label data regimes but comes with substantial computational costs, making it impractical in time- and resource-constrained industrial settings. We address this by nitegrating Decorrelated Backpropagation (DBP) into MAE pre-training, an optimization method that iteratively reduces input correlations at each layer to accelerate convergence. Applied selectively to the encoder, DBP achieves faster pre-training without loss of stability. To mimic constrained-data scenarios, we evaluate our approach on ImageNet-1K pre-training and ADE20K fine-tuning using randomly sampled subsets of each dataset. Under this setting, DBP-MAE reduces wall-clock time to baseline performance by 21.1%, lowers carbon emissions by 21.4%, and improves segmentation mIoU by 1.1 points. We observe similar gains when pre-training and fine-tuning on proprietary industrial data, confirming the method's applicability in real-world scenarios. These results demonstrate that DBP can reduce training time and energy use while improving downstream performance for large-scale ViT pre-training. Keywords: Deep learning, Vision transformers, Efficient AI, Decorrelation

Figures (12)

• Figure 1: Model overview of DBP application to the vision transformer Encoder. Input images are split into patches, linearly embedded, and the corresponding patch embeddings are decorrelated before entering the encoder. In each of the $L$ encoder blocks, the decorrelation matrix $\textbf{R}$ is applied to the inputs of the multi-head self-attention layers and to the inputs of the feed-forward MLP layers.
• Figure 2: Pre-training performance of BP-MAE and DBP-MAE on ImageNet-1K. Reported results are the average of five random initialized networks where the variability is indicated by the error shading. Panels show training loss (a) and validation loss (b) as a function of the number of epochs. The vertical dotted line indicates the point where DBP is turned off for DBP-MAE.
• Figure 3: Pre-training performance of BP-MAE and DBP-MAE on ImageNet-1K. Reported results are the average of five random initialized networks where the variability is indicated by the error shading. Panels (a) and (b) show training loss and validation loss respectively as a function of wall-clock time for the full training duration. Panels (c) and (d) show the training loss and validation loss respectively as a function of wall-clock time for the first 10 hours of training.
• Figure 4: Performance comparison of semantic segmentation fine-tuning on ADE20K, using pre-trained DBP-MAE and BP-MAE. Reported results are the average of fine-tuning on five random initialized pre-trained network checkpoints where the variability is indicated by the error shading. Panels show performance for the two regimes in training loss (a), validation loss (b), and validation mean IoU (c) as a function of the number of epochs.
• Figure 5: Pre-training performance of BP-MAE and DBP-MAE on semiconductor wire bonding data. Reported results are the average of five random initialized networks where the variability is indicated by the error shading. Panels show training loss (a) and validation loss (b) as a function of the number of epochs.