Source-linked AI summary
An Empirical Study of Training Self-Supervised Vision Transformers
Xinlei Chen, Saining Xie, Kaiming He
TL;DR
Self-supervised Vision Transformers lack established training recipes, particularly compared with mature convolutional-network practice. This paper empirically studies foundational training components across self-supervised frameworks and finds that instability can hide 1∼3% accuracy degradation, while stabilization improves accuracy and supports strong scaling evidence.
Problem
Training recipes for self-supervised ViT are not yet established as maturely as those for standard convolutional networks.
Method
The paper benchmarks self-supervised ViT across MoCo v3 and other frameworks, probing learning rate, optimizer, architecture, and model-scale choices.
Results
The study finds that instability can cause hidden 1∼3% accuracy degradation, while freezing the patch projection layer alleviates instability and consistently increases accuracy.
Takeaways & Limitations
Self-supervised ViT can achieve strong results, including competitive performance against convolutional ResNets and improved accuracy with larger models.
Takeaways & Limitations
Freezing the first layer alleviates but does not solve instability, which can persist when the learning rate is too large and may concern all layers.
Abstract
from arXiv · showhide
This paper does not describe a novel method. Instead, it studies a straightforward, incremental, yet must-know baseline given the recent progress in computer vision: self-supervised learning for Vision Transformers (ViT). While the training recipes for standard convolutional networks have been highly mature and robust, the recipes for ViT are yet to be built, especially in the self-supervised scenarios where training becomes more challenging. In this work, we go back to basics and investigate the effects of several fundamental components for training self-supervised ViT. We observe that instability is a major issue that degrades accuracy, and it can be hidden by apparently good results. We reveal that these results are indeed partial failure, and they can be improved when training is made more stable. We benchmark ViT results in MoCo v3 and several other self-supervised frameworks, with ablations in various aspects. We discuss the currently positive evidence as well as challenges and open questions. We hope that this work will provide useful data points and experience for future research.
1. Introduction
The paper studies self-supervised Vision Transformers as a baseline for narrowing the pre-training gap between vision and language. It finds that instability can quietly reduce accuracy, while stabilization improves results and supports strong scaling and framework comparisons.
- Self-supervised ViT training recipes remain less established than those for convolutional networks, motivating investigation of batch size, learning rate, and optimizer.
- Instability can cause a mild 1∼3% accuracy degradation without catastrophic divergence, making partial failure difficult to notice.
- Freezing the patch projection layer alleviates instability and consistently increases accuracy across several scenarios.
- The study benchmarks multiple frameworks, architectural choices, and model scales while discussing positive evidence, challenges, and open questions.
- Self-supervised Transformers achieve strong contrastive-learning results, and larger models can improve accuracy unlike supervised ViT results reported previously.
- The models obtain competitive results against large convolutional ResNets, while removing positional embeddings causes only a small accuracy decrease.
2. Related Work
Related work connects self-supervised vision to contrastive Siamese learning and positions ViT as a Transformer-based vision baseline. The paper focuses on contrastive training rather than masked auto-encoding.
- Contrastive learning trains representations by attracting positive samples and separating negative samples, with supervised-level comparisons reported for some tasks.
- Recent Siamese methods remove negative samples while retaining the architecture, emphasizing invariant features learned by matching positive samples.
- ViT is purely Transformer-based and largely closes the architectural gap between NLP and vision, making it a baseline for self-supervised computer vision.
- Earlier self-supervised vision Transformers generally used masked auto-encoding, whereas this work studies the contrastive/Siamese paradigm.
3. MoCo v3
MoCo v3 is an incremental, scalable contrastive framework used to study self-supervised Transformers. It uses two augmented views, momentum keys, batch negatives, and a symmetrized InfoNCE-style loss.
- MoCo v3 is an incremental improvement of MoCo v1/2 designed to balance simplicity, accuracy, and scalability.
- Two augmented crops are encoded into query and key vectors, with corresponding views forming positive pairs and other batch keys serving as negatives.
- MoCo v3 abandons the memory queue for sufficiently large batches and uses the symmetrized loss ctr(q1, k2)+ctr(q2, k1).
- The query encoder includes backbone, projection, and prediction heads, while the key encoder omits the prediction head and updates by momentum.
- 73.8 linear accuracy is reported for R50 MoCo v3, versus 71.1 for MoCo v2 and 72.2 for MoCo v2+.
- The reported MoCo v3 improvement is mainly attributed to the extra prediction head and large-batch training with batch size 4096.
4. Stability of Self-Supervised ViT Training
Replacing ResNets with ViTs in self-supervised frameworks exposes training instability that can remain hidden behind apparently decent accuracy. The study examines basic factors and shows that stabilizing training, especially with random patch projection, improves results across several frameworks.
- Empirical observations: Self-supervised ViT training can be unstable even when replacing a ResNet backbone appears straightforward.The instability is identified as a main practical challenge in contrastive and Siamese self-supervised frameworks.
- Batch size: 71.5% and 72.6% linear probing accuracy are obtained with batch sizes of 1k and 2k, whereas 4k becomes noticeably unstable at 72.2%.The 6k batch shows larger dips and reaches 69.7%, despite not diverging.
- Learning rate: 70.4 vs. 72.2 linear probing accuracy shows that lr=0.5e-4 under-fits relative to lr=1.0e-4; lr=1.5e-4 is less stable and less accurate.Smaller learning rates improve stability but can under-fit, while larger rates increase instability.
- Optimizer: 72.5% accuracy is achieved by LAMB at lr=5e-4, but accuracy drops by 1.6% and 6.0% at lr=6e-4 and 8e-4.The curves remain smooth while degrading gradually, suggesting that smoothness alone does not guarantee reliable optimization.
- Gradient dynamics: Gradient spikes appear earlier in the patch projection layer and are delayed by tens of iterations in later layers, motivating its freezing.A sudden gradient change is associated with a dip in the training curve.
- Stability intervention: Random patch projection raises MoCo v3 accuracy by 1.7% to 73.4% at lr=1.5e-4 and improves accuracy in SimCLR, BYOL, and SwAV.The reported gains are 0.8% and 1.3% for SimCLR and BYOL, and SwAV improves from 65.8% to 66.4% at its largest stable learning rate.
- Alternative interventions: BatchNorm and WeightNorm do not improve instability on the learnable patch projection, while sufficiently strong gradient clipping can become equivalent to freezing it.The authors interpret freezing as alleviating, rather than solving, an optimization problem.
5. Implementation Details
The implementation follows standard ViT designs and specifies the training, architecture, evaluation, and model-configuration details used throughout the experiments.
- Optimization: The study uses AdamW with a default batch size of 4096, searches learning rate and weight decay, and applies a 40-epoch warmup followed by cosine decay.The authors note that the long warmup helps alleviate instability, although unstable results remain possible.
- MLP heads: The MoCo v3 projection head is a 3-layer MLP, while its prediction head is a 2-layer MLP with 4096-dimensional hidden layers and 256-dimensional outputs.Both heads use ReLU in hidden layers; MoCo v3 applies batch normalization to all MLP layers.
- ViT architecture: The ViT uses 16×16 or 14×14 input patches, two-dimensional sine-cosine position embeddings, a learnable class token, Transformer blocks, and the final normalized class token as backbone output.For 224×224 inputs, the patch sequences have length 196 or 256.
- Linear probing: Representation quality is evaluated by removing the self-supervised MLP heads, freezing the features, and training a supervised linear classifier on ImageNet for 90 epochs.The classifier uses SGD with batch size 4096 and evaluates single-crop top-1 validation accuracy.
- Model configurations: The experiments study ViT-S, ViT-B, ViT-L, and ViT-H configurations, with ViT-B used by default for ablations.The configurations are characterized by Transformer block count, channel dimension, attention-head count, and an MLP hidden dimension of 4×dim.
- Compute cost: Training cost is reported using model FLOPs and TPU hours per 100 ImageNet epochs, with FLOPs measured per 224×224 crop and compared against ResNet-50.The table reports relative FLOPs versus the 4.1G ResNet-50 baseline.
6. Experimental Results
Experiments benchmark and ablate self-supervised ViT across architectures, normalization, heads, training length, scale, and comparisons with prior art. Results show competitive representations, sensitivity to design choices, and accuracy saturation at larger scales.
- Self-supervised ViT was evaluated on ImageNet-1k using linear probing across several configurations, with ViT-B as the default ablation model.
- Architecture and normalization: Removing BatchNorm required a batch size of 2048 for convergence and reduced linear accuracy from 76.5% to 74.4%.The authors conclude that BatchNorm is not necessary for contrastive learning, but appropriate use can improve accuracy.
- Architecture and normalization: Removing the prediction MLP reduced linear accuracy from 76.5% to 75.5%, although MoCo remained functional without the predictor.
- Training length: Training ViT-S for 600 epochs improved accuracy by 0.9%, while longer training produced diminishing gains for ViT-B.
- Comparisons with prior art: Smaller patches improved accuracy by approximately 2–3% at roughly 6× the FLOPs, with ViT-BN-L/7 reaching 81.0% under MoCo v3 linear probing.
- Comparisons with prior art: Self-supervised ViT models achieved competitive results against large convolutional ResNets and appeared to rely less on positional information than convolutional networks.Removing position embeddings caused only a small accuracy decrease.
- Scaling and transfer: Accuracy improved as self-supervised ViT models grew from ViT-B to ViT-L but saturated at ViT-H, while transfer results exceeded supervised counterparts on three of four datasets.
- Scaling and transfer: Self-supervised pre-training reduced overfitting and narrowed the disadvantage of ViT on small datasets, although larger-data supervised pre-training could still perform better.
7. Conclusion
The paper consolidates empirical evidence on self-supervised ViT training while identifying instability, scaling saturation, and unresolved comparisons as continuing challenges. It aims to provide data and experience for narrowing the pre-training gap between vision and language.
- The study compares ViT with convolutional networks, supervised and self-supervised training, and contrastive learning versus masked auto-encoding.
- The authors report both positive evidence and challenges, open questions, and opportunities for future research.
- The postscript notes that similar instability behavior was observed for supervised BiT-ResNet, suggesting the instability issue may extend beyond this paper’s scope.
A. Additional Implementation Details
The appendix specifies augmentation, BatchNorm, optimizer, MLP-head, and monitoring choices used to implement and assess the self-supervised ViT experiments. It also describes an approximate dense kNN monitor for exposing training stability.
- Data and normalization: Each training iteration uses two 224×224 crops with random resized cropping, flipping, color jittering, grayscale conversion, blurring, and solarization.
- Data and normalization: SyncBN is the default BatchNorm implementation, with batching choices differing in whether paired crops share a batch.
- Optimization and heads: The AdamW implementation uses decoupled weight decay consistent with the cited formulation, including framework-specific scaling considerations.
- Optimization and heads: BYOL MLP heads use BatchNorm in hidden layers but not output layers, whereas the best SwAV configuration uses no BatchNorm in its MLP heads.
- Stability monitoring: Dense kNN monitoring uses a memory bank containing 10% of ImageNet and is sufficient to reflect training stability.
- Reference implementation: The ResNet-50 MoCo v3 implementation uses LARS with a 4096 batch, lr=0.3, wd=1.5e-6, temperature τ=1.0, and a scheduled momentum coefficient.