A Theoretical Analysis of Self-Supervised Learning for Vision Transformers (2403.02233v3)

Abstract: Self-supervised learning has become a cornerstone in computer vision, primarily divided into reconstruction-based methods like masked autoencoders (MAE) and discriminative methods such as contrastive learning (CL). Recent empirical observations reveal that MAE and CL capture different types of representations: CL tends to focus on global patterns, while MAE adeptly captures both global and subtle local information simultaneously. Despite a flurry of recent empirical investigations to shed light on this difference, theoretical understanding remains limited, especially on the dominant architecture vision transformers (ViTs). In this paper, to provide rigorous insights, we model the visual data distribution by considering two types of spatial features: dominant global features and comparatively minuscule local features, and study the impact of imbalance among these features. We analyze the training dynamics of one-layer softmax-based ViTs on both MAE and CL objectives using gradient descent. Our analysis shows that as the degree of feature imbalance varies, ViTs trained with the MAE objective effectively learn both global and local features to achieve near-optimal reconstruction, while the CL-trained ViTs favor predominantly global features, even under mild imbalance. These results provide a theoretical explanation for distinct behaviors of MAE and CL observed in empirical studies.

• The paper introduces an end-to-end theoretical model for one-layer transformers in MIM, ensuring global loss convergence with gradient descent.
• The paper reveals that transformers learn diverse, local feature-position correlations by leveraging softmax attention and inherent data structures.
• The paper validates its claims with numerical experiments, showing that MIM-trained transformers outperform alternatives by focusing on varied local attention patterns.

Transformers Learn Through Feature-Position Correlations in Masked Image Modeling

Recent advancements in deep learning have spotlighted the capabilities of transformers, particularly within the field of NLP. Their utility has since expanded into the visual domain, with masked image modeling (MIM) emerging as a notable application. This blog post explores the theoretical exploration of how transformers, when applied to MIM tasks, learn by developing correlations between features and their positions, as elucidated in recent research.

Theoretical Insights into Masked Image Modeling (MIM)

The paper under discussion provides a substantial theoretical framework for understanding the learning dynamics of transformers trained on masked image modeling tasks. It is highlighted that transformers not only capture global features but are also quite effective at discerning local patterns. The fundamental contribution of this research lies in its theoretical exposition of how transformers are predisposed to learning the correlation between visual features and their spatial positions during the MIM pretraining phase.

Key Findings from the Analysis

1. End-to-End Theoretical Model: For the first time, an end-to-end theoretical model has been proposed for learning one-layer transformers with softmax attention in an MIM setting. This provides a global convergence guarantee of the loss function trained by gradient descent (GD), examining both input and position embeddings simultaneously.
2. Feature-Position (FP) Correlations: The paper intricately models the learning process of FP correlations, revealing that the transformers, during the MIM pretraining, show a marked preference for local and diverse attention patterns. This inclination towards locality is attributed to the underlying data distribution and the inherently structured training dynamics of the softmax-based transformers.
3. Diverse Local Patterns: Through theoretical proofs and empirical validation, the paper illustrates how pretrained transformers via MIM exhibit local attention patterns that significantly diverge from global uniform patterns emphasized in discriminative self-supervised learning approaches. This is further substantiated by introducing an attention diversity metric to assess the manifold attention focus across different patches of an image.
4. Numerical Results and Claims: Numerical experiments affirm the proposed theoretical findings, showcasing that MIM-trained transformers outperform other models by focusing on diverse local patterns—a trait not uniformly observed in models trained through contrasts or supervised learning methods.

Implications and Future Directions

The theoretical insights into transformer learning dynamics in MIM tasks have pivotal implications. Firstly, they open pathways to designing more efficient transformer models tailored for visual tasks by leveraging feature-position correlations. Additionally, the introduced attention diversity metric offers a novel avenue for evaluating and comparing the inductive biases of various self-supervised learning frameworks.

Future research could explore extending the theoretical framework to multi-layer transformers and encompassing broader classes of data distributions. Furthermore, investigating the role of different masking strategies and their effects on the learning of FP correlations could yield deeper understandings of optimal pretraining strategies for visual transformers.

Conclusion

This paper underscores the significance of understanding the theoretical aspects underlying the empirical successes of MIM with transformers. By proving that transformers learn to attend to local features through FP correlations, it sheds light on why certain pretraining methods excel in downstream tasks. These findings not only contribute to the theoretical literature but also guide the design and evaluation of more sophisticated models for image processing and beyond.