A Theoretical Analysis of Self-Supervised Learning for Vision Transformers

Authors: Yu Huang, Zixin Wen, Yuejie Chi, Yingbin Liang
Primary area: learning theory
Keywords: Theory of transformers, Convergence analysis, Nonconvex optimization, Theory of self-supervised learning

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.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1729663455494

thread_type:
official_review

<thread_object>
creator: reviewer_chhh

<content>
summary:
This paper discusses the reasons for the different attention mechanisms between discriminative (contrastive-learning) and generative (MAE) approaches for self-supervised learning with solid proof. This mainly contributes to the community of how these two learning strategies can be chosen and also further pushes the detailed understanding of the mechanism behind these two methods.

strengths:
1. The quality of the paper is great; details for the problem setup, statement, and proof are given in the easy-to-understand setup.
2. The theoretical analysis provides a solid step to understanding the detailed differences between the mechanisms and can help future researchers develop more approaches based on this analysis.

weaknesses:
1. The finding within the Sec. B should appear in the main paper instead of the Appendix. Some of the sections, such as Sec. 2.3 and 2.4 can be shortened since there are some overlapped concepts. Afterward, the Fig. 4 and the explanation can be moved up to provide more insight into local and global differences.

questions:
1. In the first paragraph of Sec. 2.1 about **Masked reconstruction-based learning**, should it be $\mathcal{M}_i \subset \mathcal{P}$ instead of $\mathcal{M}_i \subset [P]$?
2. Does the same cluster exist in different $\mathcal{D}_k$ defined in Sec. 2.2 and depicted in Fig.2? If this is true, should this assumption be a little bit overwhelming? For CL, only entities obtained from the same images are considered positive. As for MAE, only local features are important. I am sure changing this assumption won't affect the proof, but I want to clarify this.
3. From the proof, can we determine the required ratio between positive and negative samples to ensure the global convergence properties? Since approaches such as MoCo require more negative samples to prevent collapsing.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730106345882

thread_type:
official_review

<thread_object>
creator: reviewer_wfmo

<content>
summary:
This paper provides a theoretical framework for understanding self-supervised learning (SSL) in vision transformers, comparing masked autoencoders (MAE) and contrastive learning (CL). It examines how these methods handle imbalances between global and local visual features, revealing that MAE learns both global and local features, resulting in diverse attention patterns, whereas CL tends to focus on global features, often collapsing into uniform attention patterns. This distinction supports empirical observations, showcasing MAE’s effectiveness in capturing spatially varied data, while CL is better suited for global pattern recognition.

strengths:
1. This paper introduces a novel theoretical framework for analyzing self-supervised learning (SSL) in vision transformers, filling a gap in the field that has been primarily empirical. 
2. The paper effectively illustrates the differences in feature capture between MAE and CL, showing how MAE learns both global and local features, while CL focuses more on global features. 
3. Through a detailed gradient descent analysis of a single-layer transformer, the study offers concrete findings on how MAE and CL converge to distinct attention patterns, ensuring result reliability and reproducibility. The study focuses on a single-layer transformer model, which may overlook the dynamics of self-supervised learning in deeper or more complex networks.

weaknesses:
The study focuses on a single-layer transformer model, which may overlook the dynamics of self-supervised learning in deeper or more complex networks.

questions:
Given that this study primarily focuses on single-layer transformer models, how might the dynamics of self-supervised learning differ in deeper or more complex transformer architectures? Specifically, Could the authors discuss potential approaches or challenges in extending their analysis to multi-layer transformers? Are there specific aspects of the current proof techniques that may or may not generalize to deeper architectures?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730683510809

thread_type:
official_review

<thread_object>
creator: reviewer_3zvq

<content>
summary:
This paper claim to present a theoretical analysis of self-supervised learning for Vision Transformers (ViTs), focusing on the differences between contrastive learning (CL) and masked autoencoders (MAE). The study introduces a novel framework to model visual data distribution, considering both dominant global features and minuscule local features. By analyzing the training dynamics of one-layer softmax-based self-attention using gradient descent on spatial structured data, the authors demonstrate that MAE effectively learns both global and local features, while CL tends to focus predominantly on global features. The research provides a rigorous theoretical explanation for empirically observed behaviors of MAE and CL in ViTs. It highlights how the degree of feature imbalance affects the learning process, offering insights into why MAE can capture both global and subtle local information simultaneously, while CL tends to emphasize global patterns.

strengths:
1. This paper proposes a comprehensive framework to theoretically analyse attention mechanism for CL and MIM. The framework simplify the analysis, avoiding disentangling patches and positional encoding.
2. This work attempt to answers the empirical observations for CL and MIM with ViTs: the MIM captures diverse local patterns, while the CL focuses on the dominant patterns.
3. Base on the proposed theory, this paper finds that there are two phases to learn the feature-position correlations when information gap is positive.

weaknesses:
1. Vision transformers and transformers in general consists of two main blocks multi-head self-attention block as well as so called FFN consisting of two point-wise convolutions. The discussion is based on only one-layer of self-attention (with over simplification to only one weight matrix for self-attention) within vision transformer. Thus it cannot fully explain the complex behaviour of ViTs. 
2. There are a lot of assumptions, such as the input data distribution and positional encoding. It is not clear the differences in real problems.
3. What is the insight from this work to develop new methods? How to prevent the attention collapse in CL?
4. The discussion is based on SGD. However, Adam is the popular optimizer for ViTs. 
5. The original MIM (first MIM work with ViTs which outperformed supervised pretraining) work SiT [1] from which idea of MAE [4] is derived/copied from (only MIM part of SiT is taken in MAE) combines both MIM and CL, a strategy proven to be better and more principled time and again. Furthermore, state-of-the-art methods like iBoT [2] and its extension DINOv2 [3] also follow combined MIM and CL. It would be more fruitful to build theoretical framework for this combined MIM and CL setting (which seems to be the way forward in SSL) rather than building framework for MIM and CL separately. 

[1] Atito etal., "SiT: Self-supervised vIsion Transformer", arxiv 04, 2021. 
[2] Zhou etal., "iBOT: Image BERT Pre-Training with Online Tokenizer", arxiv 11, 2021. 
[3] Oquab etal., "DINOv2: Learning Robust Visual Features without Supervision", arxiv 04, 2023. 
[4] He etal., "Masked Autoencoders Are Scalable Vision Learners", arxiv 11, 2021.

questions:
- The theory particularly for CL seems to be build on the notion of dominant concept which takes most of the area of the input. This assumption is not valid for almost all practical datasets for SSL. For instance, the average object to total image area ratio is 0.35 in ImageNet. Which means the so called background is almost always going to be dominant concept according this paper's definition. In that case the attention from any part of the image should focus on the background most of the time. If this is the case the kNN evaluation of DINO and MoCo should not be so high for imagenet and more importantly for other datasets. 
- Is the proposed theory would still be valid for datasets where the global area is small? 
- What will happen if the local areas are larger than the global area? This situation is going to be the dominant case in most practical datasets.  
Minors: 
- what is $k_s$? It seems the sum of weights for global area.
- What $\Delta$ means? Can you explain in details, such as when $\Delta=-1,0,1$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730733361756

thread_type:
official_review

<thread_object>
creator: reviewer_i2aq

<content>
summary:
This paper is about understanding and contrasting the behavior of two pretraining techniques used for ViT architectures: Masked Auto-encoding and Contrastive Learning. They aim at providing convergence guarantees for this behavior as well. The overall conclusion is that the MAE learning objective learn both global and local behaviors, while contrastive learning has a bias towards only learning global behaviors.

strengths:
Extensive, well-formulated, and thorough proofs.

weaknesses:
- Not sure where they conducted their experiments and where the results are. They alluded to some experiments here: “line 377: However, we observe that the selected area does not necessarily depend on the location of p”. Where are these experiments located?
- This work lacks serious experimental validation. They need some kind of information bottleneck need experiments which play with the information gap and the size of the ViT tokens/patches. How do you measure this imbalance between local and global features?
- The abstract mentions “analysis” and never mentions “convergence guarantees”. Initially, I thought this paper would have extensive experimentation documenting the global and local behavior of ViT pre-trained with different objectives. But the paper isn’t about that.

questions:
- Information gap: What are the bounds/range on this? Can we see a visualization or even a histogram that shows how this value changes for different sizes of local image patches/tokens? Maybe apply it to different samples from various datasets and comment on the behaviour observed? The reason I am asking this is because a lot of theorems you mentioned make assumptions about the ranges of information gap and I think some approximations on the information gap’s behavior on some standard or simulated datasets might help. 

- Orthonormality assumption: line 193. Why do you make that assumption? To maximize span?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731978883009

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision Uploaded

comment:
We have uploaded a revised version of the paper, which incorporates the reviewers' suggestions. Below is a summary of the changes made:

1. The experimental findings, including Figure 3 and its explanations, have been moved to the main body of the paper.  
2. Notations and wording have been revised in several places as per the reviewers' suggestions, with all changes highlighted in red-colored texts in the revision.  

We sincerely thank all the reviewers for their valuable comments and suggestions. We are happy to address any further questions or concerns during the discussion period.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary:  
This paper presents a theoretical analysis of self-supervised learning for Vision Transformers (ViTs), focusing on the differences between contrastive learning (CL) and masked autoencoders (MAE). By analyzing the training dynamics of one-layer softmax-based self-attention using gradient descent on spatially structured data, the authors demonstrate that MAE effectively learns both global and local features, while CL tends to focus predominantly on global features, offering insights into why MAE can capture both global and subtle local information simultaneously, while CL tends to emphasize global patterns.

The paper's strengths are providing a comprehensive theoretical explanation for a behavior gap of ViTs pretrained with MAE and CL.  
The main weaknesses raised by the reviewers are 1) too many assumptions in the explanation (using a single-layer transformer model) and 2) lacking serious experimental validation.

All reviewers recognized the contribution of this paper, introducing a novel theoretical framework and providing a theoretical explanation for analyzing the behavior gap of ViTs pretrained with MAE and CL.

Although concerns were raised about the assumptions (specifically, the one-layer self-attention model) used in the theoretical explanation, as noted by Reviewer 3zvq and WFmo, this paper makes significant strides in uncovering insights into the behavior of the dominant architecture, vision transformers (ViTs).

Recognizing the importance of encouraging theoretical analysis of ViT models, the Area Chair recommends accepting this paper. Additionally, the authors are encouraged to incorporate the necessary clarifications and discussions into the final version.