A Theoretical Analysis of Self-Supervised Learning for Vision Transformers

We thank the reviewer for providing the valuable feedback! In our revised version, we have highlighted all changes using red-colored text for clarity. Below we address your concerns.

Q1. The finding within the Sec. B should appear in the main paper instead of the Appendix.

Thank you for your helpful suggestions! In the revision, we have included our empirical findings (see Figure 3) and discussed insights for local and global differences (see line 237-240) in the main body of the paper as suggested by the reviewer.

Q2. Should it be $\mathcal{M}_i\in \mathcal{P}$ instead of $\mathcal{M}_i\in [P]$

Thank you for pointing out this issue. We have updated this notation in the revision to maintain consistency and clarity throughout the paper.

Q3. Does the same cluster exist in different $\mathcal{D}_k$ defined in Sec. 2.2 and depicted in Fig.2?

We would like to clarify that the area partitions of different $\mathcal{D}_k$ need not be the same. The only assumption we impose on the area partitions is about the number of patches of the areas within one cluster, without specifying any fixed spatial arrangements across different clusters.

Q4. From the proof, can we determine the required ratio between positive and negative samples to ensure the global convergence properties?

Thank you for your question. Our current analysis only estimates the number of negative samples needed to ensure certain concentration properties essential for our analysis. A more fine-grained characterization of the positive-to-negative sample ratio, particularly in relation to global convergence rates, would require a detailed examination of data properties and complex calculations of the dynamics. This is beyond our current goal but is an interesting question on its own. In fact, we would like to see some extensions or investigations into how transformers detect the most significant features, not just in the sense of number of patches, but maybe also how much magnitude of the activations they produce as compared to other features. We believe there are still many unrevealed mysteries on the inductive bias of vision transformers and self-supervised learning.

We thank the reviewer again for your highly inspiring comments. We hope that our responses have resolved your concerns. In case you have further questions, please feel free to let us know, and we will be more than happy to answer them.

We thank the reviewer for the helpful comments! In our revised version, we have highlighted all changes using red-colored text for clarity. Below we address your concerns.

Q1. Regarding deeper or more complex transformer architectures?

a. We want to point out that a rigorous analysis of the training dynamics of deep transformers remains an extremely challenging open problem, currently beyond known technical tools. The difficulty can arise from different directions. The interactions between the weights and gradients in multiple layers are unclear and unlikely to be fully understood by the current tools. The architectural inductive bias may be more compounded with more layers involved, blurring the analysis of individual components. And as we known currently, there is hardly any general technical tools for non-convex optimizations, and the existing general tools rarely produces results that connects to the practice of deep learning. This is evidenced by the fact that most of the recent theoretical studies [1-7] of characterizing the training dynamics of transformers and ViTs focused on the single attention layer, which already required highly non-trivial techniques.

b. A rigorous theoretical analysis for training dynamics of multi-layer transformers is only possible when stricter assumptions are added (for example, when assuming layer-wise training or using even more specialized training recipes), and is in general beyond what the current theoretical tools can handle.

c. Lastly, we would like to emphasize that the goal of this paper is to offer a theoretical explanation for an interesting behavior gap of ViTs pretrained with MAE and CL. Our theoretical findings show that even with one-layer ViTs, we can provably characterize a similar behavior separation which aligns well with the empirical observations of full ViTs models. This suggests that our theoretical insights are not limited to a single self-attention layer but may apply more broadly to full ViTs.

We thank the reviewer again for your highly inspiring comments. We hope that our responses have resolved your concerns. In case you have further questions, please feel free to let us know, and we will be more than happy to answer them.

[1] Jelassi et al. Vision transformers provably learn spatial structure

[2] Tian et al. Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

[3] Zhang et al. Trained transformers learn linear models in-context

[4] Huang et al. In-Context Convergence of Transformers

[5] Tarzanagh et al. Max-Margin Token Selection in Attention Mechanism

[6] Li et al., A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity.

[7] Li et al., How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?

Thank you for your time in reviewing our paper. In our revised version, we have highlighted all changes using red-colored text for clarity. Below we address your concerns

Q1. Not sure where they conducted their experiments and where the results are.

a. First, we would like to clarify that the observation mentioned on line 377 is meant to explain the theoretical expressions given in lines 370-373 (lines 398-401 in the revision), suggesting that the area selected by the theoretically optimal solution may not depend on the location of $p$. We reworded the sentence to make it easier to understand from the context.

b. As we mention in the main body of the paper, we have included the experimental results in Appendix Section B. Please note that the primary focus of our paper is on the theoretical analysis of phenomena, which has been well supported by the extensive experiments provided in [1, 2]. Our own experiments in Appendix B mainly serve to further validate our theoretical findings as complement to the existing experiments. Further, our experiments also validate our proposed new empirical metric – "attention diversity metric"-- designed to assess the attention patterns observed after MAE and CL pretraining.

c. As also suggested by Reviewer ChhH, in the revision, we have added the experimental findings including Figure 3 and its explanations (line 237-240) in the main body of the paper.

Q2. 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?

a. As we mention in the main body of the paper, we have included the experimental results in Appendix Section B. Please note that the primary focus of our paper is on the theoretical analysis of the observed phenomenon, which is supported by the extensive experiments in prior works [1, 2]. Our own experiments in Appendix B mainly serve to further validate our theoretical findings as a complement to the existing experiments. Further, our experiments also validate our proposed new empirical metric – "attention diversity metric"-- designed to assess the attention patterns observed after MAE and CL pretraining.

b. Regarding the notion of information gap, our definition here is not related to the notion of information bottleneck. It refers to the (theoretical) difference between the sizes of the global and local feature areas. If such a difference is above a certain threshold (i.e., the information gap is larger than a threshold), then there is an imbalance of local and global features. Our experiments have been conducted with respect to such a metric. Information bottleneck-type experiments do not fit the concept that we aim to demonstrate.

Q3. 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.

We respectfully disagree with the reviewer to count this point as a weakness of the paper. As our title suggests, we focus on theoretical analysis rather than empirical analysis. Our abstract also discusses our contribution in the context of “theoretical understanding remains limited”. Further, as the reviewer commented in the strength part”, the paper contains “Extensive, well-formulated, and thorough proofs.” Overall, it should be clear that the main contributions here should be on the theory side. Nevertheless, we added theoretical analysis in several places in our abstract to avoid future confusion.

Q9. what is $\kappa_s$? It seems the sum of weights for global area

As outlined in Assumption 2.2, $P^{\kappa_s}$ represents the number of patches in each local area. Therefore, $\kappa_s$ specifically measures the number of patches in each local area. Instead, $P^{\kappa_c}$ refers to the number of patches in the global area.

Q10. what $\Delta$ means?

a. As outlined in Assumption 2.2, $\kappa_c$ and $\kappa_s$ determine the number of patches belonging to the global ($P^{\kappa_c}$) and each local ($P^{\kappa_s}$) area respectively. Consequently, the difference $\Delta = (1-\kappa_s) - 2(1-\kappa_c) \approx \kappa_c - \kappa_s$, as specified in Equation (4.1), quantifies the disparity in the number of patches between these areas. This metric allows us to gauge to what extent the global area contains more information than each local area, due to each patch carrying feature information specific to its designated area.

b. Specifically, a larger $\Delta$, such as 1, indicates that the number of global features substantially exceeds the number of local features, highlighting a significant imbalance. Conversely, a smaller $\Delta$, such as smaller than 0, suggests a more balanced distribution between global and local features.

We thank the reviewer again for your highly inspiring comments. We hope that our responses have resolved your concerns. If so, we wonder if the reviewer could kindly consider increasing your score. Certainly, we are more than happy to answer your further questions.

[1] Jelassi et al. Vision transformers provably learn spatial structure

[2] Tian et al. Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

[3] Zhang et al. Trained transformers learn linear models in-context

[4] Huang et al. In-Context Convergence of Transformers

[5] Tarzanagh et al. Max-Margin Token Selection in Attention Mechanism

[6] Park et al., What Do Self-Supervised Vision Transformers Learn?

[7] Wei et al., Contrastive learning rivals masked image modeling in fine-tuning via feature distillation

[8] Xie et al., Revealing the Dark Secrets of Masked Image Modeling

[9]Wang et al., Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot

We thank the reviewer for the time and thoughtful comments on our work. In our revised version, we have highlighted all changes using red-colored text for clarity. Below we address your concerns.

Q1. based on only one-layer of self-attention, cannot fully explain the complex behavior of ViTs

a. Along the research line of characterizing the training dynamics of transformers and ViTs, even the theoretical analysis of a single attention layer with one weight matrix simplification requires highly non-trivial techniques, evidenced by the most recent such studies [1-5] on the training dynamics of one-layer transformers under various settings and data models.

b. Even with the one-layer attention-only setting, our objective is highly non-convex, leading to complicated multi-phase decompositions in analyzing the training dynamics. A rigorous theoretical analysis for multi-layer ViTs is only possible when stricter assumptions are added (for example, when assuming layer-wise training or using even more specialized training recipes), and is in general beyond what the current theoretical tools can handle.

c. The goal of this paper is to offer a theoretical explanation for an interesting behavior gap of ViTs pretrained with MAE and CL. Our theoretical findings show that even with one-layer ViTs, we can provably characterize a similar behavior separation which aligns well with the empirical observations of full ViTs models. This suggests that our theoretical insights are not limited to a single self-attention layer but may apply more broadly to full ViTs. We hope this work could inspire interests in the theory community to extend our results.

Q2. a lot of assumptions:

a. We would like to emphasize two key points about these assumptions: i). Most assumptions, such as the input data distribution, are abstracted from real-world datasets that exhibit both global and local features, as evidenced by various empirical studies [6-8]; ii). Our assumptions align with those commonly made in existing theoretical work. For instance, positional encoding is a standard assumption in many studies, as referenced in [1,9].

b. Our assumptions represent necessary compromises to make complex real-world phenomena amenable to mathematical analysis. For instance, precisely characterizing real data distributions [1,4] or optimizing multi-layer transformers [1-5] mathematically is challenging. Our focus is on distilling the core architecture—specifically the attention mechanism— to better understand and simulate these phenomena. Our approach simplifies the complex designs of self-supervised learning and ViTs to uncover fundamental dynamics using mathematical tools, serving as a prototype for practical observations.

Q3. insight from this work to develop new methods? How to prevent the attention collapse in CL?

a. Our theoretical results implied that CL by nature tends to capture the most prominent features in an image, which are often global, while MIM can better capture minor, local features. This complementary capability motivates the potential advantage for hybrid objectives that merge the strengths of both approaches. By doing so, we can harness CL's efficiency at highlighting critical global information alongside MIM's precision in local feature detection.

b. It is important to note that avoiding attention collapse in CL is not universally advantageous. There are scenarios where such collapse might be beneficial, particularly when the goal is to highlight dominant image features explicitly. However, for applications requiring a broader exploration of feature space, integrating the CL framework with objectives that prioritize local detail—similar to those used in MAE methods—can be beneficial. Such an approach will ensure a balanced attention mechanism that leverages the strengths of both global and local feature extraction.

Q4. The discussion is based on SGD. However, Adam is the popular optimizer for ViTs.

Although Adam is commonly used for training ViTs, Theoretical works typically focus on from SGD/GD to provide a cleaner analysis and clearer presentations of the training behaviors. Adam, due to its momentum and normalization, results in significantly more complex training dynamics. To the best of our knowledge, nearly all of the existing theories for analyzing training dynamics for transformers and ViTs e.g., [1-5,9], including ours, focus on gradient descent. However, we believe our analysis techniques can be further extended to study Adam by incorporating the previous optimization tools for handling Adam.

We sincerely hope the reviewer could evaluate the contributions of the paper by positioning the paper in the state-of-the-art of the theoretical study of transformers. Please note that the theoretical progress of transformers is still highly restricted by the existing mathematical tools, which can handle only limited architectures and under certain assumptions. Expecting the theoretical setting to fully capture the complicated experimental settings is somewhat unrealistic given the current techniques that have been developed. This said, this paper makes a significant step forward in developing rigorous techniques and frameworks for analyzing transformers in SSL. The technical problem that we have handled is already one of the most challenging mathematical problems compared to the existing theoretical studies on the training dynamics of transformers in other tasks.

Q1. This theoretical framework introduces several assumptions.

We would like to highlight that our work aims to provide a theoretical framework explaining the behavior gap between two SSL approaches. The results serve as proof techniques to offer insights, not as definitive representations of practical phenomena. Expecting absolute scientific accuracy in such complex systems exceeds the scope of theoretical studies. Moreover:

1. Generality vs Specificity: Theoretical works that aim to provide very general theorems often fail to provide fine grained characterization of an example with special structures. We carefully chose our assumptions that align with standard theoretical studies of deep learning to enable analytical tractability. Simplified settings, such as ours, are widely used to provide foundational insights [1–9]. And even though they have added certain ad hoc assumptions to enable mathematical analysis, the shared mechanisms are still manifested in the proofs and experiments. It is our contribution, not drawbacks, to find a set of assumptions to enable these analyses and make them simple, to reach sufficiently interesting results in the theoretical sense. We have explained the role of specific assumptions in the paper and our initial response. Finally, due to the approach we took, our theory is highly interpretable, and we are open to addressing further specific questions if needed.

2. Comparison Study via Theory is difficult: Creating a unified setting to compare both SSL methods is significantly more challenging than developing separate theoretical results for each with fewer assumptions. This effort necessitates compromises, which are reflected in our design.

3. Empirical Validation: The claim regarding “global FP correlations” has been supported by existing empirical explorations [11-13], including our own attention distance experiments. Unfortunately, we cannot offer new experiments given the short fuse, and we kindly remind the reviewer that per ICLR policy, requests for additional experiments at this stage are not permissible.

Q2. This work is limited in theoretical analysis, lacking new messages for SSL design

1. As noted in our initial response, our work includes empirical insights that could inform SSL algorithm design, but its primary goal is to provide a theoretical explanation. Expecting every theoretical contribution to directly inspire state-of-the-art designs is beyond the foundational scope of theoretical research.

2. We respectfully disagree with the critique regarding the omission of the paper mentioned by the reviewer, as it does not reflect a lack of awareness of key developments. Our focus is on the theoretical comparison of the two methods, rather than on the empirical design of combining them. Therefore, while the referenced paper may be tangentially related, it is not core to our contribution and does not directly align with the objectives of our work.