A Theoretical Analysis of Self-Supervised Learning for Vision Transformers
Authors: Yu Huang
Zixin Wen
Yuejie Chi
Yingbin Liang
Venue: ICLR 2025 Poster (2025)
Keywords: Theory of transformers
Convergence analysis
Nonconvex optimization
Theory of self-supervised learning
Primary area: learning theory

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.

Published as a conference paper at ICLR 2025
A THEORETICAL ANALYSIS OF SELF-SUPERVISED LEARNING FOR VISION TRANSFORMERS
Yu Huang': Zixin Wen'; Yuejie Chi Yingbin Liang
University of Pennsylvania *Carnegie Mellon University The Ohio State University
ABSTRACT
Self-supervised learning has become cornerstone in computer vision, primarily divided into reconstruction-base methods like masked autoencoders (MAF) and discriminative methods such as contrastive learning (CL). Recent empirical ob- servations 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 inves- tigations 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 minus- cule local features, and study the impact of imbalance among these features. We analyze the training dynamics of one-layer softmax-based ViTs 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.
INTRODUCTION
Self-supervised learning (SSL) has been a leading approach to pretrain neural networks for down- stream applications since the introduction of BERT (Devlin et al., 2018) and GPT (Radford et al., 2018) in natural language processing (NLP). On the other hand, in vision, self-supervised learning focused more on discriminative methods, which include contrastive learning (CL) (He et al., 2020; Chen et al., 2020) and non-contrastive learning methods (Grill al., 2020; Chen et al., 2020; Caron et al.. 2021; Zbontar et al., 2021). Inspired by masked language models in NLP and the seminal work of vision transformers (ViTs) (Dosovitskiy al., 2020), generative approaches, such as masked reconstruction-bassee methods, have gained prominence in self-supervised vision pretraining. The masked autoencoders (MAE) (He et al., 2022) and SimMIM (Xie et al., 2022) have demonstrated the effectiveness of visual representation learning via reconstruction-bassee objectives.
Contrastive learning-like objectives promote instance discrimination among samples in the same batch of training. With suitable data augmentation, CL returns well-trained vision encoders like CLIP (Radford et al., 2021) and DINO (Caron et al., 2021) that can serve as backbones for state- of-the-art multimodal large language models (MLLMs)) (Tong et al., 2024). Masked reconstruction objectives (e.g.. MAE), on the other hand, enforce neural networks to reconstruct some or all patches of an image given masked inputs. In practice, the MAE-like approach proves to have intriguing generalization properties that differ significantly from the behaviors in CL. The seminal work (He al., 2022) showed that MAE can visibly conduct visual reasoning fill missing patches even under very high masking rates. Some critical observations from recent research (Wei et al., 2022b; Park al., 2023; Xie et al., 2023) provide comparative studies of these SSL approaches. They concluded that the ViTs trained via generative objectives display diverse attention patterns: different query patches pay attention to distinct local areas. This is in sharp contrast the discriminative approaches,
'The first authors contributed equally.

---

Published as a conference paper at ICLR 2025
DINO
MoCo
MAE
Query 2
Generative: diverse
local patterns associated with query patches
Query I
Discriminative: collapse into analogous atlention patterns
Figure I: Visualization of attention maps in the last layer of VIT for query patches from two different spatial locations, similar to those presented in Park et al. (2023). The ViTs were trained by the generative self-supervised learning approach of masked reconstruction (MAE) and discriminative methods: DINO (Caron et al., 2021) and MoCo (Chen et al., 2021b).
whose attention heads focus primarily the most significant global pattern regardless of where the query patches are, as shown in Figure 1. These empirical observations motivate the question: from theoretical standpoint, how do ViTs pick up these observed attention patterns during the training process, respectively for different SSL methods?
Despite extensive empirical efforts of studying SSL in vision pretraining, its theoretical understanding still nascent. Most existing theories of SSL. focused on the discriminative approach (Arora et al., 2019; Chen al., 2021a; Robinson al., 2021; HaoChen al., 2021; Tian et al., 2021: Wang et al., 2021; Wen & Li, 2021; 2022), especially (non-)contrastive learning. There are also a few attempts towards understanding methods using the generative approach like masked reconstructions (Cao et al., 2022; Zhang et al., 2022; HaoChen et al., 2021; Pan et al., 2022), which mainly adapt the theories developed for CL to their context. fact, there are two major limitations these prior works: Transformers, as the dominant architecture in practice, were not studied in the aforementioned works self-supervised learning and ii) there still lacks a suitable thcoretical framework that can provide convincing explanations for the empirical findings in Park et al. (2023); Xie et al. (2023), especially on the difference the attention patterns learned by the different approaches of SSL. The above limitations highlight significant gap in the literature SSL for vision pretraining
Motivated by the limited theoretical characterization of SSL for vision with transformers, especially in comparing CL and masked reconstruction objectives, aim to address the following research questions:
Our Research Questions
Can we theoretically characterize the solutions that ViTs converge to in these two mainstream self-supervised learning approaches? How do differences in attention patterns emerge during their respective training processes?
Contributions. In this paper, we take step toward answering the above questions. We study the gradient descent (GD) training process of one-layer softmax-based ViTs for both masked reconstruc- tion and contrastive learning, focusing spatially structured data distributions generalized from supervised learning settings (Jelassi et al., 2022). In our setting, each image is sampled from distinct clusters characterized by unique patch-wise feature associations. Each cluster contains two types of features: a large portion of patches reside in a global area and share global features, while the remaining local areas contain relatively few patches with their own local features. We measure the imbalance of feature distribution by a condition called the information gap A, which is defined in eq. (4. Under such setting:
More detailed discussions for related work found in Appendix
2

---

Published as a conference paper at ICLR 2025
1. We provide global convergence guarantees for training ViTs both the MAE and the CL loss fucntions. To the best of our knowledge, this is the first end-to-end guarantee for learning VITs with self-supervised learning objectives;
2. We provide comprehensive characterization of the training dynamics of attention corre- lations (see Definition 3.1) to illustrate the attention patterns to which ViTs converge: i). MAE provably learns diverse attention patterns, with each patch concentrating its attention on its designated area based on its position, even under a substantial information gap A; ii). CL primarily learns global attention pattern, causing all patches to focus on the global area regardless of their locations, even with minor A. These qualitative differences in the solutions learned by these two SSL methods provide strong theoretical support for the empirical behavior gaps observed Park et al. (2023); Xie et al. (2023), and highlight the theoretical advantage of MAE in handling highly imbalanced data structures.
Notation. We introduce notations to be used throughout the paper. For any two functions h(x) and g(z), we use h((r)) = (x))) (resp. kx(x:) = O(g(z)) to denote that there exist some universal
constants C1 0 and 01. s.t. [h(x)) 2 C1lg(x)) (resp. |h(z)) Ci Cz))(=))) for all 2 01; Furthermore, h(x)) (g(x)) indicates ((x)) = ((((=) and h(x) Ozg(x)) hold simultaneously. } to denote the indicator function, and let [MV] 21, We use Õ, Z2 and O to further hide logarithmic factors in the respective notations. We use poly (P) and polylog(P) to represent large constant-degree polynomials of P and log(P), respectively.
We use }
PROBLEM SETUP
In this section, we present our problem formulations for studying the training process of ViTs in self-supervised pretraining. We begin with some background information, followed by a description of our data distribution. We then detail the pretraining strategies using MAE and CL. respectively with the specific transformer architecture considered in this paper.
2.1 BACKGROUND ON SELF-SUPERVISED LEARNING
Masked reconstruction- based learning. We follow the masked reconstruction frameworks He et al. (2022); Xie et al. (2022). Each data sample X E dR P has the form X = (Xp)peP, which has IP| = P patches, and each patch Xp E IRd, Given collection of images [ei)e we select masking set M, C P for each image X,, and mask these patches uniform value M Rd The resulting masked images (M(X,) are given by
[Xi]p M
DEU PEM,
M(X,), =
i [n]]
(2.1)
where U4 P M; is the index set of unmasked patches. Let F X + X be architecture that outputs reconstructed image X d for any given input X 12dxP, The pretraining objective is then defined as the mean equn-equared reconstruction loss over series of subsets Pi C P of the image as follows:
Caasked ; E
-
|
(2.2)
MAE (He et al., 2022) chose the subset P/ as the set of masked patches M, whereas SimMIM (Xie et al., 2022) aimed to reconstruct the full image P/ = P. We do not explore the trade-offs between these two approaches in our study.
Contrastive learning. Contrastive learning (Chen et al., 2020) aims to learn meaningful representa- tions F by distinguishing between similar and dissimilar data points. For given batch (X.) generate a positive pair (xx"),xx xx for each by applying random augmentations to X,. Negative pairs (x(", (2) for j * are formed from different data points. The model F is trained to minimize the following contrastive loss:
Smy (x x(P))
]
Lcontrastive (F) -
(2.3)
Smp (x(1)xx11))
ESejn]
3

---

Published as a conference paper at ICLR 2025
Sampling
P1,
Pk,
PKA
Global
P1,
Local
Px
D1
P**
P1.3
Pr.
(De3k=1
D1
Dx
Figure 2: Illustration of our data distribution (sce Definition 2.1). Each cluster Dy is segmented into distinct areas Pk.j with squares in the same color representing the same area Px.j- The global area Pk51 (depicted in orange) contains larger count of patches compared to any other local areas. It is important to note that while we use spatially contiguous partitions for clarity in this illustration, our data model also applicable non-contiguous cases.
where F measures the similarity between two representations, and is temperature parameter controlling the sharpness of the distribution.
2.2 DATA DISTRIBUTION
We assume the data samples X E IR are drawn independently based on some data distribution D. To capture the feature-position (FP) correlation in the learning problem, consider the following setup for vision data. assume that the data distribution consists many different clusters, where each cluster captures distinct spatial pattern, and hence is defined by different partition of patches with a different set of visual features. We define the data distribution DD formally as follows. An intuitive illustration of data generation is given in Figure 2.
Definition 2.1 (Data distribution D). The data distribution D has K Oppolylog (P) different clusters {Dx} 1 k= For every cluster DR, k E [K], there is a corresponding partition of P into Nu disjoint subsets P U Pk.j which we call areas. For each sample X (Xp) its sampling process is as follows:
We draw D& uniformly at random from all clusters and draw a sample X from Dj.
Given [K], for any j E [Nu]. all patches Xp in area Pk.j are given the same content Xp (x)zz((); where va, E IRd is the visual feature and (X) the latent variable. We assume U 1==1 1 {uk.j} are orthogonal to each other with unit norm.
Given & E [K]. for any j [Na]. (X) E U.. where 0 < L < U are on the order of @(1).
Global and local features, and empirical observations in prior works. Image data naturally contains two types of features: the global features and the local features. For instance, in an image of an object, global features can capture the shape and texture of the object, such as the fur color of an animal, whereas local features describe specific details of local areas, such the texture of leaves in the background. Recent empirical studies on self-supervised pretraining with ViTs (Park et al., 2023; Wei al., 2022b) and observations in Figure I collectively show that masked pretraining exhibits the capacity to avoid attention collapse concentrating towards those global shapes by identifying diverse local attention patterns. Consequently, unraveling their mechanisms necessitates thorough examination of data characteristics that embody both global and local features. In this paper, we characterize these two types of features by the following assumption on the data.
Assumption 2.2 (Global feature vs local feature). Let Dj with is E [K] be cluster from D. We let Pk,1 be the global area of cluster Di, and all the other areas Pk.j.. E [Nu] \ [1} be the local areas. Since each area corresponds respectively. Moreover, we assume:
assigned feature, also call them the global and local features,
Global area: given & [K], we Ck IPx.1l O(pr)) with N. e [0.5005, 1], where Ck.1 is the number of patches in the global area PA.L.
The distribution of 2, (X) can be arbitrary within the ahove support set.
4

---

Published as a conference paper at ICLR 2025
Attention Diveersity Metric, distance
Ang attn distence Cuntt: pat-hh
Pread Ltdlei
Nodel
(supervised)
We design novel empirical metric, the attention diversity metric, to probe the last layer of ViTs trained by masked reconstructions (MAE), CL (MoCo), another discriminative SSL approach (DINO), and supervised learning (DeiT). Lower values of this metric signify focused attention on a similar area across different patches, reflecting a global pattern of focus. Conversely, higher values suggest that attention is dispersed, focusing different, localized areas. The results show that the MAE model excels in capturing diverse local patterns compared to discriminative methods like CL. (see Appendix B for details).
Figure 3: Attention Diversity Metric:
Local area: given [K]. we choose Ck.j ((P*.) with Ki [0.001, 0.5] for 1, where Ck.j denotes the number of patches in the local area Pk.j
The rationale for defining the global feature this manner stems from observing that patches representing global features (C*,1) typically occur more frequently than those representing local features (Ck.j. for j > 1), since global features capture the primary visual information of an image, offering a dominant view, while local features focus on subtler details within the image. Our empirical observations (see Figure 3) further substantiate the significance of distinguishing between global and local patterns data distributions, which is essential for elucidating the distinct behaviors exhibited by MAE and CL.
2.3 MASKED RECONSTRUCTION WITH TRANSFORMERS
Transformer architecture. A transformer block ( (Vaswani et al., 2017; Dosovitskiy et al., 2020) consists of self-attention layer followed by an MLP layer. The self-attention layer has multiple heads, each of which consists of the following components: a query matrix W a key matrix W and value matrix WV. Given an input X, the output of one head in the self-attention layer can be described by the following mapping:
G(X; W W wV) softmax ((WX)) w*x), (ovXx)Tx<x) (WVX)
(2.4)
where the softmax(-) function is applied row-wisely and for a vector input E R",, the i-th entry of softmax(2) is given by
exp(2;)
Eln exp(x.)"
To simplify the theoretical analysis, we consolidate the product of query and key matrices (WQ) W into one weight matrix denoted as Q. Furthermore, we set W to be the identity matrix and fixed during the training. These simplifications are often taken in recent theoretical works (Jelassi et al.. 2022; Huang et al., 2023; Zhang et al., 2023a) in order to allow tractable analysis. With these simplifications place, eq. (2.4) can be rewritten as
G(X; Q) = softmax (xT XT.
(2.5)
Input tokens in transformers are indistinguishable without explicit spatial information. Therefore, positional encodings should be added to the input embeddings to retain this crucial positional context as in practices (Dosovitskiy et al., 2020; He et al., 2022). Our assumptions regarding the positional follows:
encodings
Assumption 2.3 (Positional encoding). We assume fixed positional encodings, which is consistent with the implementation in MAE (He et al., 2022): embedding vectors ep orthogonal each other and to all the features Ukjj and are of unit-norm.
(ep)pEp ddxP where positional
5

---

Published as a conference paper at ICLR 2025
We now include positional embeddings in eq. (2.5) and introduce the network architecture for masked reconstruction used in this study.
Definition 2.4 (ViTs for MAE). We assume that our vision transformer Fmae (x consists of a single-head self-attention layer with an attention weight matrix Q E Rdxd For an input image
X ~ D, we add positional encoding by letting X = x E. The attention score from patch Xp to patch Xq is denoted by
attn th p- (X;Q) == x 2
for P.
(2.6)
YeEp x Qx,
Then the output of the transformer is given by
(x:
Σ qep
attn (X: Xq. for forpe
(2.7)
Then formally define the masking operation and the objective for masked pretraining task. Definition 2.5 (Random masking). Let M(X) + Rdx denote the random masking operation, which randomly selects (without replacement) subset of patches M in X with masking ratio (1)) E (0, 1) and masks them to be M == 0 E IR The masked samples obey eq. (2.1).
MAE objective. To train the model Em (M(X; ; Q), following the methodology described in MAE practice (He et al., 2022), we minimize the squared reconstruction error in eq. (2.2) only on masked patches, where M(X) follows Def. 2.5. The training objective thus can be written as
(Q) =
1
E1(e E M) (M))) pEP
(M(X):
xpl
(2.8)
where the expectation is with respect to both the data distribution and the masking. Note that our objective remains highly nonconvex with the model defined in Definition 2.4.
Training algorithm. The learning objective in eq. (2.8) is minimized via GD with learning rate initialize (()) Odxd dl as the zero matrix. The parameter is updated follows:
Q ((++1) Q() Vaao (Q
Note that the initialization of Q (0) results in any query patch uniformly attending all patches.
2.4 CONTRASTIVE LEARNING WITH TRANSFORMERS
The transformer architecture used for CL is similar to that of MAE, but with a minor modification to accommodate contrastive loss, as outlined below.
Definition 2.6 (ViTs for CL). We consider vision transformer Fel (X; Q) consisting of single- head self-attention layer with an attention weight matrix Q E Rdxd For an input image X, the attention score from patch Xp patch Xg is denoted by
IT 4x4
attn (atnn-q(X;; (X;Q) Dxx oxx
for p. EP.
(2.9)
The output of the transformer is then computed as
Fe1 (X:Q) P
Σ attn (X:Q) -Xg
ERd.
(2.10)
D.qEP
which represents the average pooling of all the patches.
The key distinction is that we separate the positional and patch embeddings within the attention mechanism for technical simplicity. However, it is important to emphasize that these two types of embeddings remain coupled for attention calculations.
Definition 2.7 (Data augmentation). For sample X E Rd generate new samples X and X++ by independently applying random masking as in Def. 2.5 with a ratio "o = e(1), similar to the crop-resize operations used in practice. The unmasked sets for them are denoted as 2/+ and U++

---

Published as a conference paper at ICLR 2025
CL objective. Given sample X, we first generate pair of positive samples {x+ x++} via Def. 2.7. Then we generate batch of i.dd. negative samples 91 = {X <E[.. Denoting
2 = NUU [ we minimize the expected contrastive loss in eq. (2.3) with fz-regularization:
eim (x+ (x+x++))/+
]]
1Q1
Le1 = Ex+,x++m - log
Σ Exeg Sim (x+,)
(2.11)
where || lIe denotes the Frobenius norm, À > is the regularization parameter, and the similarity of the representations of X and X' obtained by F () is defined as
pi X') (Fo1 (X; StopGrad (F") (X':
The StopGrad operator ensures that gradient is computed for this term. Additionally, augmentation is applied to the negative samples. Both practices are standard in the literature on the theory of contrastive learning (Wen & Li, 2021; 2022). Similar MAE, update Q by GD with zero-initialization:
((+11) = Q -/:0~1/Veee V/Voaa (Q
(2.12)
In the following, any variable with superscript (t) represents that variable at the 1-th step of training.
3 ATTENTION PATTERNS AND FEATURE-POSITION CORRELATIONS
To show the significance of the data distribution design and understand the nature of our self- supervised learning tasks, in this section, we will provide some preliminary implications of the spatial structures in Def. 2.1. Intuitively, for MAE, for a given cluster Dx, to reconstruct missing patch p Pk. M, the attention head should exploit all unmasked patches in the target area Pk.j j to find the same visual feature 0k, to fill in the blank, which emphasizes the locality for in different areas. However, CL focuses on any discriminative patterns regardless of the location of p. which can align positive pairs but may lead to collapsed attention patterns. We will elaborate on these points by describing the area attentions and illustrating the intuition about how they can be learned attention correlations (Def. 3.1).
Area attention. We first define new notation for a cleaner presentation. For X ~ D and p E P, we write the attention of patch Xp to subset A C P of patches by
Attn - A (X; = =yEEA
for E [m,c].
MAE's ability to learn locality with ViTs. Let us first explain why the above notion of area attention matters in understanding how attention works in masked reconstruction. Suppose we have a sample X picked from Dr. and the patch Xp with Pa. is masked. Then the prediction of Xp given masked input M(X) can be written as
pmse (M(X); Q))p EgEP M(X) attn (M(()))
= =E ((()ekk Attn (M(X): Q) (since M(X) if E M).
To reconstruct the original patch Xp, the transformer should not only focus on the correct area Pk.j but must also prioritize attention the unmasked patches within this area. This specificity denoted by the area attention Attnp-unPL, over 2 PR.j. requirement imposed by masking operations. We refer these location-dependent attention patterns locality.
To further explain how ViTs perform such prioritization, we introduce the following quantities, which capture the major insights of our analysis distinguish between MAE and contrastive learning.
Definition 3.1. (Attention correlations) Let P, and we define attention correlations as:
Feature-Position (FP) Correlation: pp = ep Qua,m. for k E [K] and m E [Ns];
PositionPosiition (PP) Correlation: T Tp-qq Qeqe Vg P.
(0) Due (zero) initialization of (()), have (1)) ppruk,m p-q 0.
7

---

Published as a conference paper at ICLR 2025
These two types of attention correlations, FP correlation Pp-rer. and PP correlation Pp-+q. act as the exponent terms within the softmax calculations for attention scores. Given p E Pk.jj is masked, the (unnormalized) attention at directed towards unmasked patch q is influenced jointly by these correlations. Hence, the described attention pattern for MAE can emerge from either a substantial FP correlation $ p--.o., significant PP correlation O pp-- for q in the same area as p. However, in our setting, the latter aethaannsm-maarniii via PP correlation-
to produce desired attention patterns: i). such mechanism inadvertently directs attention the masked patches, which is not desirable; it). such position association could be vulnerable to the variation across different clusters, i.e., p. q E Pk., does not necessarily hold for all h E [K]. This also highlights that prior work (Jelassi et al., 2022) that relied solely on pure positional attention cannot fully explain the ViTs' ability to learn locality when the patch-wise associations are not fixed.
Why CL may fail to explore the locality. Now turning to CL., for X' E DE, we have the following form of similarity between the positive pair:
(Fo) (X++ Q), Fel
Q))
Pr Ep.'EP E Attnp--u-nan (X+; Attn (X++:Q).
Thus, to align the positive representations effectively, the optimal strategy is also to direct atten- tion toward specific area for each patch p, i.e., greedily ensuring that only one area attention
Attnp-uutnPe, is activated for some i E [N&x]- However, the above expression suggests that the selected area by the optimal strategy may not necessarily depend on the location p. which could lead to a collapsed attention scenario where all patches focus on the same area. Regarding attention correlations, the attention mechanism defined in eq. (2.9) requires to handle only the FP corre- lations among different features for CL. Theorem 4.4 in the next section confirms that collapsed solution indeed occurs: VIT' trained with CL concentrate attention the global area across all patches by exclusively capturing global FP correlations across all patches, i.c., Pp-rur1 becomes large for Vp E P.
STATEMENTS OF MAIN RESULTS
In this section, we present our main theorems on the learning processes of ViTs in MAE and CL. We begin by introducing notations that will be used in theorem presentations.
Information gap and technical condition. Based on our data model in Section 2.2, we introduce a notion of information gap to quantify the degree of imbalance between global and local areas (ef. Assumption 2.2). Denoted as A, the information gap is defined as follows:
(4.1) Broadly speaking, a larger A means that the number of global features is much greater than local ones, indicating a significant imbalance. In contrast, smaller value reflects only a slight imbalance.
A= (1 - k2) 2(1 ke).
Unmasked area attention. Based on the crucial role of those unmasked patches for both reconstruc- tion task and positive contrastive pairs, we further define the unmasked area attention follows:
Attn p P (X; Q): ;= Attnp-urrpp p-rPPr.. (X;Q), for i {m,cc.
4.1 MAE LEARNS DIVERSE ATTENTION PATTERNS
Our results are structured into two parts: i). analysis of convergence (Theorem 4.1), which includes the global convergence guarantee of the masked reconstruction loss and characterization of the attention pattern at the end of training to demonstrate the diverse locality; ii). learning dynamics of attention correlations (Theorem 4.2), which shows how transformers capture target FP correlations while downplaying PP correlations as discussed in Section 3.
To properly evaluate the reconstruction performance, we further introduce the following notion of the reconstruction loss with respect to a specific patch p E P:
1 Cuae (Q) E 1pp E (MMX)
Xpp
(4.2)
*Our study focuses on the regime where A is not too close to zero, i.e., | 2(11, which allows for cleaner induction arguments. This condition could he potentially relaxed via more involved analysis.
8

---

Published as a conference paper at ICLR 2025
Nearly no attention areas
Pk.
k,
Nearly atcention the masked patches the same area: small Yp-q for Pk.
k,1
Mask
PPP Quiny
syPks
Focus the unmaskeed patches the same area: large P-vy.
D
Figure 4: The mechanism of how the masked patch attends other patches via attention correlations.
Now we present our first main result regarding the convergence of MAE.
Theorem 4.1 (Training convergence). Suppose the information gap A (0()) [2(1), 1]. For any 0 < < < 1, suppose polylog(P) > log(2). We train the ViTs in Def. 2.4 by GD to minimize reconstruction loss in eq. (2.8) with vj < pollyyy (P). Then for each patch p P, we have
1. Loss converges: Lm (Q( Ceee.p S Tn = - =
(og(P)
Pm 1),1)(1xx))
q log (?)) iterations, where miaa. is the global minimum of the patch-level reconstruction loss in equation 4.2.
2. Area-wide pattern of atlention: given cluster k E [K], and p E Pk. for some j E [N&]. if Xp is masked. then the one-layer transformer nearly "pays all attention" all unmasked patches the same area Pk. as i.e.,
O(e).
Theorem 4.1 indicates that, at the time of convergence, for any masked query patch Xp in the k-th cluster, the transformer exhibits an area-wide pattern attention, concentrating on those unmasked patches within the area that p lies in, as demonstrated in Section 3. The location of the patch determines such area-wide attention and can be achieved no matter if p belongs the global or local areas, which jointly highlight the diverse local patterns for masked vision pretraining matter degree of the imbalance.
Next, we detail the training phases of attention correlations in the following theorem, which explicitly confirms that the model learns target FP correlations while ignoring PP correlations to achieve the desirable arca-wide attention patterns suggested in Section 3 (illustrated in Figure 4).
Theorem 4.2 (Learning Feature-Posinee
correlations). Following the same assumptions in Theo- rem 4.1, for p P. given k E [K], if E Pk. for some j E [Ne],, we have
For positive information gap A [](1).1]:
Global areas (j learn FP correlation in one-phase: & p-rru. monotonically increases to (P/e) throughout the training, with all other attention correlations remain close to O. b. Local areas (j > 1) learn FP correlation in two-phase: In phase one, FP correlation p between local area and the global area feature quickly decreases (log (P)) whereas all
(1)
other attention correlation stay close to zero; In phase two, FP correlation (t) + for the target local area starts grow until convergence with all other attention correlations nearly unchanged.
For negative information gap A E 1-0.5, -((()]
All areas learn FP correlation through one-phase: T @p-v..s monotonically increases to O(log(P/e)) throughout the training. with all other attention correlations remain close to O.
The training dynamics are different depending whether A is positive negative, and further vary for positive A depending on whether Xp is situated in global or local areas. Typically, the target FP correlations learned directly in single phase. However, for positive information gap A, when patch p is located in a local area, the learning process contains an additional decoupling phase, to reduce the FP correlation with the non-target global features.

---

Published as a conference paper at ICLR 2025
4.2 CONTRASTIVE LEARNING COLLAPSES TO GLOBAL ATTENTION PATTERNS
In contrast to MAE'S ability to learn diverse local features regardless of the information gap, our results in this section demonstrate that CL inevitably collapses to global attention patterns by solely learning global FP correlations, even under slight structural imbalance.
To prevent trivial solutions in CL, we adopt a noisy variant of the data distribution.
Assumption 4.3 (Noisy data). We assume that the data used for contrastive learning is sampled from D= Specifically, to generate a sample X ~ De1, we first draw D, then add independent and identically distributed (i.id.)) noise Gp ~ (N(0, of to each patch Zp. The resulting sample is defined as Xp = Zp + Sp- We denote X E D if Dj--
Theorem 4.4 (Learning with contrastive objective). Suppose the information gap E [-0.5, Q11) U [S](L,,]] We train the ViTs in Def. 2.6 by GD to minimize eq. (2.11) with " < poly(P), ob = T= O (o Then after T 0 ( l iterations, we have
"
1. Loss converges: Lei (o)(0( s C + poly(p poly(P). where Ca is the global minimum of the contrastive loss in eq. (2.11).
2. Attention concentration on global area : given X E D with E [K], for any p P, with high probability, we have Attn P QX))(())) = poly(r) for X' E {x+ ,X +.+
3. All patches learn global FP correlation: given E { for any p E P, E [", 4 P
& pp with m >> 1, and at the convergence, 4 (TT O(log P), 2 o = o(1).
Intuition behind learning global correlations. As discussed in Section 3, the optimal alignment of x involves directing attention towards the same feature for each patch p. possibly irrespective of its location. As long the imbalanced structure, where global features dominate the data distribution, exists-eve even to a small degree- it leads to an order-wise stronger concentration of attention on global areas initialization. Consequently, global FP correlations receive larger gradients compared to local ones. Therefore, global FP correlations are learned first, and focusing on these global correlations is sufficient for the CL objective converge.
two positive representations 1²s (Q; x and F=
(Q; F=t
Significance of the results. Theorem 4.1 and Theorem 4.4 address critical gap in understanding self- supervised pretraining by offering the first theoretical framework for learning with ViTs, one of the most advanced architectures in vision practice, whereas prior studies have primarily focused on linear models, CNNs, or MLPs (Wen & Li, 2021; Ji et al., 2023; Pan et al., 2022). Moreover, by identifying the collapsed solution in CL and emphasizing the effectiveness of MAE in capturing diverse attention patterns, we provide qualitative comparison between MAE and contrastive learning, validating a non-trivial empirical observation (Park et al., 2023). This offers comprehensive theoretical analysis of self self-supervised learning with ViTs.
CONCLUSION
In this work, we study the training process of MAE and CL with one-layer softmax-based ViTs. Our key contribution is providing the first end-to-end convergence guarantees for these two prominent self-supervised approaches with transformer architectures. We characterize the attention patterns at convergence and show that MAE exhibits diverse attention patterns by learning feature-position correlations across all features, even with highly skewed feature distributions. In contrast, CL collapses to global attention patterns by focusing solely global feature-position correlations, despite minimal distributional deviations between features. This provides theoretical justification for the behavior gap of MAE and CL observed in practice. Our proof techniques use phase decomposition based on the interplay between feature- posittin and position wise correlations, avoiding the need to disentangle patches and positional encodings in prior work. We anticipate that theory will be valuable for future studies of spatial or temporal structures in state-of-the-art transformers and will advance theoretical research in deep learning.
"This also holds when no data augmentation is applied to X.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
The work of Y. Huang is partially supported in part by the ONR grant N00014-22-1-2355 and the NSF grants CCF-2418 8156 and DMS-2143215. The work of Wen and Y. Chi supported in part by NSF under CCF- 1901 199, CCF-2007911, DMS-2134080, DMS-213415 N00014-19-1-2404. The work of Liang supported in part by NSF under RINGS-2148253, CCF-1900145, and DMS-2134145.
and by ONR under
REPRODUCIBILITY STATEMENT:
The main body of the paper presents only theoretical results, with all proofs provided in the appendices. Additionally, the appendices include proof sketches that offer intuitive explanations of the proof steps. The appendix also contains experimental results, with detailed descriptions of the experimental settings to facilitate result reproduction.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. arXiv preprint WiXx:23300000299 2023.
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saun- shi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229. 2019.
Srinadh Bhojanapalli, Ayan Chakrabarti, Daniel Glasner, Daliang Li, Thomas Unterthiner, and Andreas Veit. Understanding robustness of transformers for image classification. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10231-10241, 2021.
Shuhao Cao, Peng Xu, and David A Clifton. How to understand masked autoencoders. arXiv preprint arXiv:2202.036700 2022.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties self-supervised vision transformers. Proceedings of the IEEF/CVF international conference computer vision, pp. 9650-9660, 2021.
Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer. arXiv preprint arXiv:2402.04084, 2024.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Ting Chen, Calvin Luo, and Lala Li. Intriguing properties of contrastive losses. Advances in Neural Information Processing Systems, 34:11884-11844 2021a.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. 2021 IEEE/CVF International Conference Computer Vision (ICCV), pp. 9620- 9629, 2021b. URL htttps:// //pp..pp...//// wew.eteeaantiii www.mmattiicccola..oroooo.. wowp:/ppss///ooooroooo.o : 233024 948.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04886 2018.
Ming-Wei
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.1 1929, 2020.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation self-attention mechanisms. International Conference on Machine Learning, pp. 5793-5831. PMLR, 2022.
11

---

Published as a conference paper at ICLR 2025
Amin Ghiasi, Hamid Kazemi, Eitan Borgnia, Steven Reich, Manli Shu, Micah Goldblum, An- drew Gordon Wilson, and Tom Goldstein. What do vision transformers learn? a visual exploration. arXiv preprint arXiv:2212.06672/ 2022.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Pa- pailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13199 2023.
Evan Greene and Jon Wellner. Exponential bounds for the hypergeometric distribution. Bernoulli: official journal of the Bernoulli Society for Mathematical Statistics and Probubility, 23(3):1911, 2017.
Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-super learning. Advances in neural information processing systems, 33:21271 2020.
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 171, 2020.
Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. In Advances in Neural Information Processing Systems, volume 34, pp. 5000-5011, 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9738, 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Yu Huang. Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems, volume 35, pp. 37822-37836, 2022.
Wenlong Ji, Zhun Deng, Ryumei Nakada, James Zou, and Linjun Zhang. The power of contrast for feature learning: theoretical analysis. Journal of Machine Learning Research, 24(330):1-78, 2023.
Jason D Lee. Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo Predicting what you already know helps: Provable self-supervised learning. In Advances in Neural Information Processing Systems, volume 34, pp. 309-323, 2021.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015. 2023a.
Hongkang Li, Meng Wang, Tengfei Ma. Sijia Liu, Zaixi Zhang, and Pin-Yt Yu Chen. What improves the generalization of graph transformer? a theoretical dive into self-attention and positional encoding. In NeurlPS 2023 Workshop: New Frontiers Graph Learning. 2023b.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis, 2024.
Bingbin Liu, Daniel J Hsu, Pradeep Ravikumar, and Andrej Risteski. Masked prediction: A parameter identifiability view. In Advances in Neural Information Processing Systems, pp. 21241-21254, 2022.
Luke Melas Melas-Kyriazi. Do you even need attention? a stack of feed- ferd-forward layers does surprisingly well imagenet. arXiv preprint arXiv:2105.00722. 2021.
12

---

Published as a conference paper at ICLR 2025
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 2024.
Jiachun Pan, Pan Zhou, and YAN Shuicheng. Towards understanding why mask reconstruction
pretraining helps in downstream tasks. In The Eleventh International Conference on Learning Representations, 2022.
Namuk Park and Songkuk Kim. How do vision transformers work? arXiv preprint arXiv:2202.06709, 2022.
Namuk Park, Wonjae Kim, Byeongho Heo, Taekyung Kim, and Sangdoo Yun. What do self- supervised vision transformers learn? In The Eleventh International Conference Learning Representations, 2023.
Sayak Paul and Pin-Yu Chen. Vision transformers are robust learners, In Proceedings of the AAAI conference on Artificial Intelligence, volume 36, pp. 2071-2081, 2022.
Ashwini Pokle, Jinjin Tian, Yuchen Li, and Andrej Risteski. Contrasting the landscape of contrastive and non-contrastive learning. arXiv preprint 2203:22203.1500,, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, et al. Improving language understanding by generative pre-training 2018
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- 8748-8766 PMLR, 2021.
Maithra Raghu, Thomas Unterthiner, Simon Korblithh Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks? Advances in Neural Information Processing Systems, 14:12116-1212 2021.
Joshua Robinson, Li Sun, Ke Yu, Kayhan Batmanghelich, Stefanie Jegelka, and Suvrit Sra. Can contrastive learning avoid shortcut solutions? Advances in neural information processing systems, 34:4974-4986, 2021.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-252. 2015.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarith- mic depth. arXiv preprint 240.:22400200268, 2024a.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36, 2024b.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308, 16898, 2023a.
Davoud Atace Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max Max-margin token selection in attention mechanism. In Thiirty-eventth Conference Neural Information Processing Systems, 2023b. URL https::////......... /ttp///peeeeeeeeeeee/ee/eeeee.em net. / forum? i ddWxxc808ghH. httpp:://pperreiiieee
Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning, pp. 10268-10278 PMLR, 2021.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv preprint arXiv:2305. 16380, 2023.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. JoMA: Demystify- ing multilayer transformers via joint dynamics of MLP and attention. In The Twelfth International Conference on Learning Representations, 2024. URL https: /ttp//ppeeeviieeee net forum? i id=iibbggRNYCF.
13

---

Published as a conference paper at ICLR 2025
Shengbang Tong, Ellis Brown, Penghao Wu, Sanghyun Woo, Manoj Middepogu, Sai Charitha Akula, Jihan Yang, Shusheng Yang, Adithya Iyer, Xichen Pan, et al. Cambrian-Ie 1: A fully open, vision-centric exploration of multimodal Ilms. arXiv preprint arXiv:2406.16880 16860, 2024.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Asher Trockman and J Zico Kolter. Patches are all you need? arXiv preprint arXiv:2201.09792 2022.
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis, Implicit bias and fast convergence rates for self-attention. arXiv preprint M..Kvv22402.057::5 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.
James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. A mathematical theory of attention.
arXiv preprint arXiv:2007.02876 2020.
Xiang Wang, Xinlei Chen, Simon S Du, and Yuandong Tian. Towards demystifying representation learning with non-contrastive self-supervision. arXiv preprint arXiv:21 10.04947, 2021.
Colin Wei, Sang Michael Xie, and Tengyu Ma. Why do pretrained language models help in downstream tasks? analysis of head and prompt tuning. Advances in Neural Information Processing Systems, 34:16158-16170, 2021.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35: 12711-2083,, 2022a.
Yixuan Wei, Han Hu, Zhenda Xie, Zheng Zhang, Yue Cao, Jianmin Bao, Dong Chen, and Baining Guo. Contrastive learning rivals masked image modeling in fine-tuning via feature distillation. arXiv preprint arXiv:2205.1 141441, 2022b.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. International Conference on Machine Learning, pp. 11080-11090. PMLR, 2021.
Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning, pp. 11112-11122. PMLR, 2021.
Zixin Wen and Yuanzhi Li. The mechanism of prediction head in non-contrastive self-supervised learning. Advances in Neural Information Processing Systems, pp. 14794-24880 2022.
Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 9653 0233,966,,22 2022.
Zhenda Xie, Zigang Geng, Jingeheng Hu, Zheng Zhang, Han Hu, and Yue Cao. Revealing the dark secrets of masked image modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14475-14485, 2023.
Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. In The Thirtt-eighth Annual Conference on Neural Information Processing Systems, 2024.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of eqquence-to-sqqenne functions? arXiv preprint arXiv: 191199921.11710 2019.
14

---

Published as a conference paper at ICLR 2025
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins:
Self-supervised
learning via redundancy reduction. International Conference on Machine Learning, pp. 12310- 12320. PMLR, 2021.
Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. Advances in Neural Information Processing Systems, 35:27127-27139)
2022.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context.
arXiv preprint urXiv:2306.09927, 2023a.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang. and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization, arXiv preprint arXiv:2305. 19420, 2023b.
15

---

Published as a conference paper at ICLR 2025
APPENDIX
Related Work
17
B Experiments
17
Overview of the Proof Techniques
18
C.I GD Dynamics of Attention Correlations
19
C.2 Phase I: Decoupling the Global FP Correlations
20
C.3 Phase II: Growth of Target Local FP Correlation
21
C.4 Learning Processes in Other Scenarios
Preliminaries
22
D.I Gradient Computations
22
D.2 High-probability Event
29
D.3 Properties of Loss Function
29
Overall Induction Hypotheses and Proof Plan for MAE
E.1 Positive Information Gap
32
E.2 Negative Information Gap
................................
32
E.3 Proof Outline
32
Analysis for the Local Area with Positive Information Gap
32
F.1
Phase I, Stage
33
Phase I, Stage
................................
38
F3 Phase II, Stage
42
F4 Phase II. Stage
45
G Analysis for Local Areas with Negative Information Gap
48
H Analysis for the Global area
50
Proof of Main Theorems for MAE
1.1 Proof of Induction Hypotheses
51
1.2 Proof of Theorem 4.1 and Theorem 4.2 with Positive Information Gap
52
1.3 Proof of Theorem 4.1 and Theorem 4.2 with Negative Information Gap
53
J Proof of Main Theorems in Contrastive Learning
53
J.1 Preliminaries
54
J.2 Initial Stage: Global Correlations Emerge
55
J.3 Convergence
61
16

---

Published as a conference paper at ICLR 2025
A RELATED WORK
Empirical studies of transformers in vision. A number of works have aimed to understand the transformers in vision from different perspectives: comparison with CNNs (Raghu et al., 2021; Ghiasi et al., 2022; Park & Kim, 2022), robustness (Bhojanapalli et al., 2021; Paul & Chen, 2022), and role of positional embeddings -Melas-Kyriazi, 2021; Trockman & Kolter, 2022). Recent studies (Xie et al., 2023; Wei et al., 2022b; Park et al., 2023) have delved into ViTs with self suservuperiiion to uncover the mechanisms at play, particularly through visualization and analysis of metrics related to self-attention. Xie et al. (2023) compared the masked image modeling (MIM) method with supervised models, revealing MIM's capacity to enhance diversity and locality across all ViT layers, w which significantly boosts performance
tasks with weak semanties following fine-tuning. Building on MIM's advantages, Wei et al. (2022b) further proposed a simple feature distillation method that incorporates locality into various self-supervised methods, leading to an overall improvement in the finetuning performance. Park et al. (2023) conducted a detailed comparison between masked image modeling (MIM) and contrastive learning. They demonstrated that contrastive learning will make the self-attentions collapse into homogeneity for all query patches due to the nature of discriminative learning, while MIM leads to diverse self-attention map since it focuses on local patterns.
Theory of self-supervised learning. A major line of theoretical studies falls into one of the most successful self- saf-ssuervvied learning approaches, contrastive learning (Wen & Li, 2021; Robinson et al., 2021; Chen et al., 2021a; Arora et al., 2019), and its variant non non-contrastive self-superviser learning (Wen & Li, 2022: Pokle et al., 2022; Wang et al., 2021). Some other works study the mask prediction approach (I (Lee et al., 2021; Wei et al., 2021; Liu et al., 2022), which is the focus of this paper. Lee et al. (2021) provided statistical downstream guarantees for reconstructing missing patches. Wei et al. (2021) studied the benefits of head and prompt tuning with masked pretraining under Hidden Markov Model framework. Liu et al. (2022) provided a parameter identifiability view to understand the benefit of masked prediction tasks, which linked the masked reconstruction tasks to the informativeness of the representation via identifiability techniques from tensor decomposition.
Theory of transformers and attention models. Prior work has studied the theoretical properties of transformers from various aspects: representational power (Yun et al., 2019; Edelman et al., 2022; Vuckovic et al., 2020; Wei et al., 2022a; Sanford et al., 2024a), internal mechanism (Tarzanagh et al., 2023a; Weiss et al., 2021), limitations (Hahn, 2020; Sanford et al., 2024b), and PAC learning (Chen & Li, 2024). Recently, there has been a growing body of research studying in-context learning with transformers due to the remarkable emergent in-context ability of large language models (Zhang et al., 2023b; Von Oswald et al., 2023; Giannou et al., 2023; Ahn et al., 2023; Zhang et al., 2023a; Huang et al., 2023; Nichani et al., 2024; Li et al., 2024). Regarding the training dynamics of attention- based models, Li et al. (2023a) studied the training process of shallow ViTs in a classification task. Subsequent research expanded on this by exploring the graph transformer with positional encoding (Li et al., 2023b) and in-context learning performance of transformers with nonlinear self-attention and nonlinear MLP (Li al., 2024). However, all of these analyses rely crucially on stringent assumptions on the initialization of transformers and hardly generalize to our setting. Tian al. (2023) mathematically described how the attention map evolves trained by SGD for one-layer transformer but did not provide any convergence guarantee, and the follow-up work Tian et al. (2024) considered a generalized case with multiple layers. Tarzanagh et al. (2023b); Vasudeva et al. (2024) investigated the implicit bias for self-attention models trained with GD. Furthermore, Huang et al. (2023) proved the in-context convergence of one-layer softmax transformer trained via GD and illustrated the attention dynamics throughout the training process. Yang al. (2024) generalized such an in-context learning problem to a mult-head setting with non-linear task functions. Nichani et al. (2024) studied GD dynamics a simplified two-layer attention-only transformer and proved that it can encode the causal structure in the first attention layer. However, none of the previous studies analyzed the training of transformers under self-supervised learning, which is the focus of this paper.
B EXPERIMENTS
Previous studies the attention mechanisms of biT-based pre-training approaches have mainly utilized a metric known as the attention distance (Dosovitskiy et al., 2020). Such a metric quantifies the average spatial distance between the query and key tokens, weighted by their self-attention coefficients.
17

---

Published as a conference paper at ICLR 2025
The general interpretation is that larger attention distances indicate global understanding, and smaller values suggest a focus on local features. However, such a metric does not adequately determine if the self-attention mechanism is identifying unique global pattern. A high attention distance could result from different patches focusing on varied distant areas, which does not necessarily imply that global information is being effectively synthesized. To address this limitation, revised version of average attention distance, called the attention diversity metric, which is designed assess whether various patches are concentrating on similar region, thereby directly capturing global information.
introduce a novel and
Attention diversity metric, in distance. This metric computed for self-attention with a single head of the specific layer. For given image divided into N x N patches, the process unfolds as follows: for each patch, employed as the query patch to calculate the attention weights towards all N2 patches, and those with the top-n attention weights are selected. Subsequently, the coordinates (e.g. (i.j) with j E [M] these top-n patches are concatenated sequence to form 2 x n- dimensional vector. The final step computes the average distance between all these 2n-dimensional vectors, Le., N2 x N2 vector pairs.
Setup. In this work, compare the performance of ViT-B/16 encoder pre-trained ImageNet- IK (Russakovsky et al., 2015) among the following four models: masked reconstruction model (MAE), contrastive learning model (MoCo v3 (Chen al., 2021b)), other self-supervised model (DINO Caron et al. (2021)), and supervised model (DeiT Touvron et al. (2021)). We focus on 12 different attention heads in the last layer of VYT-B on different pre-trained models. The box plot visualizes the distribution of the top-10 10 averaged attention focus across 152 example images, as similarly done in Dosovitskiy al. (2020).
Implications. The experiment results based on new metric are provided in Figure 3. Lower values the attention diversity metric signify focused attention on a coherent area across different patches, reflecting a global pattern of focus. On the other hand, higher values suggest that attention is dispersed, focusing on different, localized areas. It can be seen that the masked pretraining model is particularly effective in learning that prioritize a uniform global information with less attention diversity. This aligns with and provides further evidence for the findings in Park et al. (2023).
diverse attention patterns, setting it apart from other models
C OVERVIEW OF THE PROOF TECHNIQUES
In this section, explain our key proof techniques in analyzing the self-supervised pretraining of transformers, using MAE as an example. We focus on the reconstruction p P. We aim to elucidate the training phases through which the model learns FP correlations related to the area associated with p across different clusters & E [K].
specific patch Xp for
Our characterization of training phases differentiates between whether Xp is located in the global or local areas and further varies based on whether is positive or negative. Specifically, for A E [9(1), 1}, observe distinct learning dynamics for FP correlations between local and global areas:
Local area attends FP correlation in two-phase: given k [K], if at.p * 1, then
1. & ( first quickly decreases whereas all other T Pp-rk... with + 1 and T(" do not change much;
after some point, the increase of P takes dominance. Such & pp- will keep growing until convergence with all other FP and PP attention correlations nearly unchanged.
(1)
Global areas learn FP correlation in one one-phase: given & E [K]. if aokp = 1, the update of (t) (() p will dominate throughout the training, whereas all other pp with m * and learned PP correlations remain close to 0.
For A E 0.5, (1)), the behaviors of learning FP correlations are uniform for all areas. Namely, all areas leam FP correlation through one-phase: given & [K], throughout the training, the increase
18

---

Published as a conference paper at ICLR 2025
of $ Pp-u... dominates, whereas all other o (1) pp- with m o aR and PP correlations T T p->q
(1)
remain close to O.
For clarity, this section will mainly focus on the learning of local feature correlations with a positive information gap A ((1) in Appendices C.2 and C.3, which exhibits two-phase process. The other scenarios will be discussed briefly in Appendix C.4.
C.1 GD DYNAMICS OF ATTENTION CORRELATIONS
Based on the crucial roles that attention correlations play in determining the reconstruction loss, the main idea of our analysis is track the dynamics of those attention correlations. We first provide the following GD updates of ( pp-re.m and T (see Appendix D.I.1 for formal statements).
Lemma C.1 (FP correlations, informal). Given E for, p E P, denote n = ak.p, let ap-vv..
(t++) 4 +v..... for m E [Nx).. and suppose Xp masked. Then p
(2) 1. for the same area, a tok
Attn (1)
(1- 1-Attn
if k E Bp. for the global area,
Attn (1) P-PPn1 Atti (1 11-Att P + Attn
p-Pan (1 (1- Attm )
0 ~ Attn p-P
3. for other area m & (nf
(Attmn -11)))))--111011111 Attm -1 Attn p P
(0) p
Attn
From Lemma C.I. it is observed that for p Pk.ns the feature correlation (1) p exhibits monotonically increasing trend over time because Op-rve,n 2 O. Furthermore, if n > 1, i.e. Pk, is the local area, ( p-ee, will monotonically decrease.
Lemma C.2 (PP attention correlations, informal). Given p. P. let 8p-ng = [(T p
and suppose Xp is masked. Then pp = EKE(N) k p.pp- where a(f) kk.p-g satisfies
1. d' dk.q n1, p attnping (1) (1 Attn (1)
pttmp--P...) )²
if k E Bp Cq. where and = > and axiq = 1:
(1) Pk.p-q = attn Attn
Att (tt
3. if ffakq = =m (n) U (1}, where 0k.p
(1) o xs attn
n fm (Attm -
(0) p-PPa. Attn
(e)
Based on the above gradient update for r pp we further introduce the following auxiliary quantity r(()
k.p-aq which can be interpreted as the PP attention correlation "projected" on the k-th cluster Dr. and will be useful in the later proof.
(1+1) : (1)) k.p-qq k. + 1
with (0)
k = 0.
(C.1)
We directly verify that r ExETK] r(() k.p-
19

---

(t) Published as a conference paper at ICLR 2025
The key observation by comparing Lemma C.1 and C.2 is that the gradient of projected PP attention
a(t)
P.p->q is smaller than the corresponding FP gradient 0p 0p- in magnitude since attnp Ro
Aten PPEaLa We will show that the interplay between the increase of dp-vve. and the decrease
(1-7)Ck:ag.
(t) of p determines the learning behaviors for the local patch p E Px., with > 1, and which effect will happen first depends on the initial attention, which also determined by the value of information gap A.
C.2 PHASE I: DECOUPLING THE GLOBAL FP CORRELATIONS
We now explain how the attention correlations evolve the initial phase of the training decouple the correlations of the non-target global features when p is located in the local area for the k-th cluster. This phase can be further divided into the following two stages.
Stage 1. At the beginning of training, $ Pp-uk.. T k.p-i 0, and hence attmp-1 for any q E P. which implies that the transformer equally attends to each patch. However, with high probability, the number of unmasked global features in the global area Pk,1 is much larger than
(0)
(0) others. Hence, Attn PPP
unnaaaa 2 (p-r)))) >> O (p
Attn (0)
PP.m for > 1.
Therefore, by Lemma C.I and C.2, we immediately obtain
1 = 0 p whereas app
dp--.... PLOOE = + (Pi--.):
all other FP correlation gradients ap-ree, with # ax.p small;
(0)
all projected PP correlation gradients kk.p are small.
Since = (1 - 2 (1)1), it can be seen that 4 p enjoys much larger decreasing rate initially. This captures the decoupling process of the feature correlations with the global feature 2k,1 the global area for p. It can be shown that such an effect will dominate over a certain period
that defines stage of phase I. At the end of this stage, we will have p S - (log( PP), whereas all FP attention correlation & Pp+ with m > and all projected PP correlations r(t)
k.p->q
stay close to 0 (see Appendix F.1).
During stage 1, the significant decrease of the global FP correlation T pp-r,, leads to reduction
in the attention score Attn P- Meanwhile, attention scores Attn p-PP.. (where m > 1) for other patches remain consistent, reflecting uniform distribution over unmasked patches within each area. By the end of stage 1, Attn P-P,, drops certain level, resulting in decrease in |of2
as it approaches ap - which indicates that stage begins.
Stage 2. Soon as stage begins, the dominant effect switches as (1 ja reaches the same order (t) of magnitude as ap- The following result shows that (1()) pp- must update during stage
Lemma C.3 (Switching of dominant effects (See Appendix F.2)). Under the same conditions as Theorem for p P, there exists Tı. such that iteration t Ti + 1, we have
(1-1) PP-.ve.kkpp
2 (log())), and to (T+1) -O((lg(P)):
T
b. all other FP correlations $12 - with m + 1, ak are small:
c. all projected PP correlations (t) k.p are small.
(t) b further is more challenging A log( P), we observe that (ap-m.. is Intuition of the transition. Once ( decreases 21
(t) approximately equal to ap-vvL.k.p"
After this point, reducing o
compared to the increase in T (t)
pp- To illustrate, minimal decrease of क ( by amount
20

---

Published as a conference paper at ICLR 2025
(
of 0.001 (P)(P) will yield pap- S ( dominant effect.
Such a discrepancy triggers the switch of the
020002
C.3 PHASE II: GROWTH OF TARGET LOCAL FP CORRELATION
Moving beyond phase I, FP correlation p Pp-eo within the target local area already enjoys (4)
larger gradient Op Op- than other pp pp- with m + ak.p and all projected PP correlations
r(t) k,p-'q' We can show that the growth of + (1) ppttkpp will continue to dominate until the end of training by recognizing the following two stages.
Rapid growth stage.
(1) At the beginning of phase II, ap-+veekk. is mainly driven by Attn (t)
Pr. since 1 Attn P-Proee remains at the constant order. Therefore, the growth of P p naturally
(e)
results in a boost in Attup P.P..kkp thereby promoting an increase in its own gradient ap Opp.o....p.
which defines the rapid growth stage. On the other hand, we can prove that the following gap holds for FP and projected PP correlation gradients (see Appendix F.3):
all other FP correlation gradients ap-vve.m with m * at.p are small:
all projected PP correlation gradients ((() Pkp- are small.
Convergence stage. After the rapid growth stage. the desired local pattern with a high target
feature-position correlation & opp-vx..x is learned. In this last stage, it is demonstrated that the above conditions for non-target FP and projected PP correlations remain valid, while the growth of (1
(t) Pp PUk.o.p starts decelerate pp-vx.eu, reaches (llog())) resulting in Attn P-P... ((1).
(*)
which leads convergence (see Appendix F.4).
C.4 LEARNING PROCESSES OTHER SCENARIOS
In this section, talk about the learning process in other settings, including learning FP correlations for the local area when the information gap is negative, learning FP correlations for the global area, and failure learn PP correlations.
What is the role of positive information gap? As described in stage 1 of phase in Appendix C.2, (0)
the decoupling effect happens at the beginning of the training because ap-ru.1 ap-ror.eg.r attributed to A 2 (11). However, in cases where A S (11,, this relationship reverses, with (0) Op-rus,1 becoming significantly smaller than (tp->"k.ee.p" Similarly, other FP gradients ap-rm..
(0)
(0)
with * 1, ak.p and all the projected gradients of PP correlation 3pth are small in magnitude. Consequently, D pp-rer.p,p starts with a larger gradient, eliminating the need to decouple FP correla- tions for the global area. As result, training skips the initial phase, and moves directly into Phase II, during which (1) p continues increase until converges (see Appendix G).
Learning FP correlations for the global area. When the patch Xp is located in the global area of cluster k, i.e., ap.p 1, the attention score Attn directed towards the target area Pk,
(0)
initially higher compared other attention scores due to the presence of a significant number of unmasked patches in the global area. This leads to an initially larger gradient ap-ive.x.p" Such an effect is independent of the value of A. As result, the training process skips the initial phase, which is typically necessary for the cases where ak.p > with a positive information gap, and moves directly into Phase II (see Appendix H).
(0)
All PP correlations are small. Integrating the analysis from all previous discussions, we establish that for every cluster l: E [K], regardless of its association with Cp (global area) B, (local area),
21

---

Published as a conference paper at ICLR 2025
and for any patch Xq with q P, the projected PP correlation X k.p-~q remains nearly zero in comparison the significant changes observed in the FP correlation, because the gradient k/.p pk.p-q is relatively negligible. Therefore, the overall PP correlation Tp-q- EA=1 k.p-g
T(
also stays close
to zero, given that the number of clusters K = (())
D PRELIMINARIES
In this section, we will introduce warm-up gradient computations and probabilistic lemmas that establish essential properties of the data and the loss function, which are pivotal for the technical proofs in the upcoming sections for masked pretraining. Throughout the appendix, we assume N₂ = N and Ck.n = Cn for all k E [K] for simplicity. We will also omit the explicit dependence on X for (x)).. We use kx [K] to denote the cluster index that a given image is drawn from. Furthermore,
will abbreviate Lmae .p) C(Cp), and pmae as F for simplicity, when the context makes it clear. We abbreviate Attmp-PP... (X; Q(9) (attn (X; Q(1)) as
(1) Attn P Pk.m. (attn)"ag). when the context makes it clear.
D.I GRADIENT COMPUTATIONS
We first calculate the gradient with respect to Q. We omit the superscript '(t)" and write L(Q) as here for simplicity.
Lemma D.1. The gradient of the loss function with respect to Q given by
ac
-E
Σ Σ M(X)) - (F(M(X); PEM
aQ
M(X) - attn Eattnp-.M())) 7
M(X)p
Σ
Proof. We begin with the chain rule and obtain
ac =
([F(((X):
Q)lp (IF(M(X):Q)) -xp)]
E Σ
8Q
eQ
PEM
@attn Dattnp-a
E
0)1 Xx)]..
(D.I)
aQ
PEM
We focus on the gradient for each attention score:
MM(1aa))))))) M(X)aa)
dattnp-g
Σ cxp
(MX)) -M(X).T
iQ
Σ. (XXX((()]
-tttn--
attmp
M(X) Σ M(ttn
attmm M(X)p
Substituting the above equation into equation D.I, we complete the proof.
Recall that the quantities (1) &f and r are defined in Definition 3.1. These quantities are associated with the attention weights for each token, and they play a crucial role in our analysis of learning dynamics. We will restate their definitions here for clarity.
Definition D.2. (Attention correlations) Given p. q E P. for 2 0, we define two types of attention correlations follows:
22

---

Published as a conference paper at ICLR 2025
1. Feature Attention Correlation: $ (1)
pu = cp (00000000 Ukm for k: [K] and E [N]:
2. Positional Attention Correlation: r ppq - Ep 1(0000 Eq.
T (0) pp-q O. By our initialization, we have & uk
Next, we will apply the expression in Lemma D.1 compute the gradient dynamics of these attention
correlations.
D.1.1 FORMAL STATEMENTS AND PROOF OF LEMMA C.I AND C.2
We first introduce some notations. Given r E U. for p E P, k [K] and n E [N] define the following quantities:
JP MMX) (F(M()())))
T
Ip.k.. M(X), Σ
attnp- M(X)
Uk,n
WEP
T
KP. = M(X),
Σ
attn M(X)w
eq
WEP
Lemma D.3 (Formal statement of Lemma C.I). Given & E [K], for p P, denote n ak.p let (t)
(0
pm 7 for E [Ne], then
(t+1) (t)
Op-7Uk.
a. for m n,
(1)
= E
Attt p-
]
Fn
1 - Attn (5)
Attm
p-PPunn )²
+ Σ
(Attn (4)
p-P
afn
b. for m # n,
E 1fpe M, =, =kak=]]Atti Attn (1) p-Pm
Σ zm (Attn (1)
(t )
ampmm
Zm
2 (1- Attm (t) pt- Attn + am (1 Attn (t)
Attmp--Pum
Attn
Proof. From Lemma D.I, we have
0L
P-tuk.
(0)
= = aQ
jk
=E(1(p M)
Σ
attn Ip.k.m)
reWl
= =E[1(p M, k) Σ
atttn-rdd? Ip.k.m)
rell
where the last equality holds since when lx + k, _..k.. 0 due orthogonality. Thus, in the following, we only need to consider the case kx = k.
Case 1: m n.
23

---

Published as a conference paper at ICLR 2025
For r U Pk.ns since Uk,n' 1 Uk,n for n' # n, and uk,n 1 feqlqER we have
Jp 2n Vk.n Znlkkn
qElPPk.. km
zm (1 Attmp-PPp PP,..)
p. UKnn
at uknn JF/2a /
qetinPN..
For r E U Pk with n' Fn
JPP = Ukm
En Un
attup- 2n Uk,n
qelmP
-zz Attmp Pinn
Azd Ukin!
att pp Ukin
Uk,n
qEunPeea
- 2n, AttupttPPP...
Putting together, then we obtain:
aL
(- 8Q
Uk,n
E [[(p x 2 (9)
=
3
1 Attm )² + Σ Lozn (Atto (1)
(2)
) 2
Case 2: m n. Similarly
For r E 2 2 Pk..
JP = Un
UnU
attm Uk
qeunP
= 2.. - Attn PPP,,,
pp....m
Zx Uk,n
attnp 2m Dk.m
qElinPe.
2m Attmttt--i PP.m
For r U UnPk
JP = 2m Ukmm
Uk,
(tt 2 2m Uk
qElnnkk.
- 2 Attmp PP..
T
p..k.n
Zm Ukm
Σ
attn 22 p =m Uk.m
Uk.
qEMrPAm
21
AttDp-Pe..
24

---

Published as a conference paper at ICLR 2025
For U Pk.as # n.
JP
Za zu
Zn Uk,
attn (1)
2 Uk,a
qElnP
Ip.k,"
attn 2 Uk.m
Uk.m
qEUrNPA..
-m Attup-
Putting them together, then we complete the proof.
Lemma D.4 (Formal statement of Lemma C.2), Given E P, let pp-aa ((r (t) then
p>>q) p p
8 p
Σ
alt) P.p-
where ((0 Bkp-- satisfies
kE(K
if ak.p
a4.q n,
E
1(pe M,Kx kJattn
Σ z Attmp-Prnn) )²+
(0)
aftnn
2 Attn (HQEE4)AAtt!" /
b. for as. = n m ak.q
0 = 1/p E kx = =x)aam
Σ apnn
(Att
(t)
2
z
pop )
(1) ( p Attn (0)
[1q Attn (1)
P-PP.e
Proof.
0L
((t)
(- leg = - M} Σ
attn JP
p-q
"Q
rell
Then we let
310) k-p-q
= E11pppppp M, =k} Σ
attn 12
rell
In the following, we denote ak.p and og m for simplicity.
Case 1: m =m If E unnu SK.P P
F
JP = 2n 2n
T
Un
p wn Uk.n
WEunP
= km (1 Atttlp--Paa
KP.q = (eq - (attnp-qe +
Σ
attn eq
w
- attnn--qg-
25

---

Published as a conference paper at ICLR 2025
For r u Pk,ns and q
J.P 2nkki
attnp-w 2n Vk,
wEunPk,
KP (er (attup--qeq
Σ
attnp--eww) Tea eq
w/q
=-attnp-i
Thus
attmp-rrJP KP
rEunPA
attnp-w
n
weunPk,
attnp-, attnp-+
PEMMPE.. k.
Attn
attu
For E u Pk.ss a fn
Jp = Zallka
Uk,
Σ
at
WEUNPE.a
2a
attnp-ww
wEUnP kM
kP (er - (attnp--qea + Σ attn eq
w/qq
-ttttpp-
Thus
Σ ptttnp-r.PPp attnp-qq' FEll
2 [11 Attn :P.n)) 3² Σ za (Attn
a#n
If E Mn Pk.
. For E unnk
T JP 2n, 2ntlkin
attmp-w22p Uk,n
WEunPE...
2. (1 Attmp--PP....
KP4 (er (attup-gee
Σ attup- ea
w
attnnp
For E 1 n Pk.as a + n
JP = Zooka
Untkin
attnp-w22 Vk
wEunP
26

---

Published as a conference paper at ICLR 2025
- WEUnPEE
attlp-ww
KP = (er - (attMp-pqea +
Σ
attn Tee
w$q
-atttpp-
Thus
Σ
rell
attnp--q''
FR (1 Attn Attnp-Pm..) -2 Attnp--Pa.
Σ
= (Attup-PPull
apn
Putting it together,
((0)
Ppp-+q E = =kjattup-qa
2 (1 Attnp--Paal 1/9E + + Attn .Pinn)+
Σ
zm (Atto - )ram))
mff
Case 2: m n. Similarly, if u U
For E uunuPA Pk...
JP 2n0kkn Un
attnp-we Vk.
=22 Attnp
kP (er atttpp-qee
Σ
Tea eq
w
--tttp--
For r = q
JP Zm kkm
ZnUkin
attnp -w 2m Ukmm
WEZNNPA,.
2m Attnpttppp
kP.q (eq attttp-po cpp
Σ attn eg
www
1 1--ath -tttnp-q
For r E u n a # n, and r + q
JD = zolkka
Un
wEunP
= AttD
KPa = (er - attnpp-gge
Σ
attn Tea eq
w/q
--attnp-g
27

---

Published as a conference paper at ICLR 2025
Thus
Σ Cattnp-rJfRK Kpp4
rell
(-o7 Attnp-sPen) Att 2 Attnp-pp2 Pp
Σ 2 (Attnn--Pe. )
attnp->q
q E P
. For r E CUMPA P
T JF kn.kk
Znikkn
attmp Zn Uk,
wEunPk.r
=2 21 Attnpp-Pr..
kP. (e, attnpp-gq
Σ attnp-weww) cc
w/q
= =attttp--q
For r E U n Pk.a. a +
J = kol T
Untkkn
attnp--w3a'l Zn
WEUnPE,
Attttt
KP = (er - attn
Σ
attn, -we))-wwew) Tea eq
w/q
-tttnp-p
Thus
Σ Cattnp-r-Pkpp
rell
= attnp--q
- (1-(111 (1-AAtp--Paa) Attup +
Σ z (Attmp-PPe..
afnn
Therefore
((0) k
= =E E M, = =kjattnp-aa
(- (1 Attup-Pand Attnp-PPa. l E u)== Attnp- P
Σ 26 (Attnp-77 -P...)
as
Based the above gradient update for r p we further introduce the following auxiliary quantity, which will be useful the later proof.
(t+1) kpp-q p m
((()
with (0)
k.p-i =
(D.2)
It is easy to verify that T ExEIK) T(() -h.p-va
28

---

Published as a conference paper at ICLR 2025
D.2 HIGH-PROBABILITY EVENT
We first introduce the following exponential bounds for the hypergeometric distribution Hyper (m, D,M). Hyper (m, D, .) describes the probability of certain successes (random draws for which the object drawn has a specified feature) in m draws, without replacement, from a finite population of size M that contains exactly D objects with that feature, wherein each draw is either success or a failure.
Proposition D.5 (Greene & Wellner (2017)). Suppose S Hyper (m, (,, D,M) with m, D M. Define PM = D/M. Then all >>> 0
t
P ( - mpm t) S exx
4muM (- + 2t
We then utilize this property to prove the high-probability set introduced in Appendix C.I.
Lemma D.6. For l E [K] E [N], define
Ex P) = {M: (Pk. nu = ((C.),
(D.3)
have
P(M E Ek,n) 2 - 2 (x.,
(D.4)
where Cn,0 > is some constant.
Proof. Under the random masking strategy, given ht E [K] and E [N]. Yk,, = -u Pk.n follows the hypergeometric distribution, i.e. Yu,, ~ Hyper - P, C., P). Then by tail bounds, for > 0, have:
Fl - (1- 1) S 2exp(
4(1 - C) + 21
Letting O(C,), we have
P|YK,- = O 2 1 - -2e
Ca,1 1 Co.
We further have the following fact, which will be useful for proving the property of loss objective in the next subsection. Lemma D.7. For ls E [K] and n E [N], we have
P(U Pk.nl 0) s exp((cn,oCn).
(D.5)
where Co.0 > 0 is some constant.
Proof. By the form of probability density for Hyper((1 P) Cn, P), we have
(c:) ((P-Co)
P(U = =0
(1-))))
$ c exp(- App(-nnoCa)).
D.3 PROPERTIES OF LOSS FUNCTION
Recall the training and regional reconstruction loss we consider are given by:
((Q)
Σ 1(p e M] Q,, Xpl12
(D.6)
pEp
Q)) = E [ E MIPEMMHN [[F((((X, E) - xp
(D.7)
In this part, we will present several important lemmas for such training objective. We first single the following lemma, which connects the loss form with the attention score.
29

---

Published as a conference paper at ICLR 2025
Lemma D.8 (Loss Calculation). The population loss L(Q) can be decomposed into the following form:
((Q) Σ
(Q(
where
PEP
K
(p) 2
Σ
=
k=1
]
(1- Attn Attup--Praapp
2 +
E z (Attn
P
aFa
Proof.
Lp(Q)
K
1
Σ
E
= (IF(((X), E)) -xp1²)
K
Σ
E
Σ
Attnp-PPam 2om Uk.m
Zan. ak
2
MEIN
1
E
M,px =k] kb}
2 )²+ (1- (1-Athn -utnp--Ph.eppp
Σ
zm (Attup pappppp-Pee. P
m
where (i) follows since the features are orthogonal.
We then introduce some additional crucial notations for the loss objectives.
c
min QERdxa
Cp(Q),
(D.8a)
L
Clow (02 N-1
Σ P [un ( PkIzOp. =00)
(D.8b)
P
2
RE(K]
(,(()= = Σ (Q),
where
k=1
E (x) 2
1 = k,M E Ex }
2 Zarn
(t) [ttt
p- ptt
(Attn )
(D.8c)
astar.
Here o = (X))( f denotes the minimum value of the population loss in equation D.7, and flow represents the unavoidable errors for p E P. given that all the patches in Ps.ae.p are masked. We will show that clow serves as a lower bound for Li. and demonstrate that the network trained with GD will attain nearly zero error compared to </ow. Our convergence will be established by the sub-optimality gap with respect to Clow, which necessarily implies the convergence Li- (It also implies Cp- -Clow is small.)
Lemma D.9. For C and Chtw defined in equation D.8a and equation D.8b, respectively, we have Clow < C. and they are both at the order of exp (-(c + 1l {1 & bke (o2pp)))
where Cı, C2 > 0 are some constants.
30

---

Published as a conference paper at ICLR 2025
Proof. We first prove L
Cp:
cp
min QeRaxa
Σ
=k}
(1- Attn (t)
)²
2 +
2 ZOk.
(Attm (0
pisppp )
2 min
M, k Pk.ap.
= =0}-
QeRdxd
(1- Attn (2)
p Pa.an
2 +
Σ
2/ =ou
Attn
aptax.
Notice that when all patches in Pk.az. masked, Attn P.Pok.pp
(1)
0. Moreover,
E mGLar
Am Attn (1)
P-PM.m
N-1
by Cauchy-Schwand inequality. Thus
1 cp 2 Σ
L
)P (u Pkoa.pl =0)
Clow
Cp
O exp
Po + 1f1 & Useln tok, )oop%.)) immediately comes from
Lemma D.7. Furthermore, we only need to show Lp O exp ( (e + 1[ &
(ax.k) (o2k)))))PP)))) This can be directly obtained by choosing Q = ols for some suffi- ciently large a and hence omitted here.
Lemma D.10. Given p E P.for any Q, we have
(+( S LQ) c S E, (Q) + o exp (-
(-(cpPre pre +1 {1 & (ax.p) pr.)))
where C3, > 0 are some constants.
Proof. The lower bound is directly obtained by the definition and thus we only prove the upper bound.
(Q)
(Q)
K
1/p M, = K, Eh } (zeen (1 p
(2)
Σ
2 (Attn
(1)
E
axaax
K UPP(M E
k=1
$ 0
cxp (cyPre + +1 (1 & (csp.)))
where the last inequality follows from Lemma D.6.
E OVERALL INDUCTION HYPOTHESES AND PROOF PLAN FOR MAE
Our main proof utilizes the induction hypotheses. In this section, we introduce the main induction hypotheses for the positive and negative information gaps, which will later be proven to be valid throughout the entire learning process.
31

---

Published as a conference paper at ICLR 2025
E. POSITIVE INFORMATION GAP
We first state our induction hypothesis for the that the information gap A is positive. Induction Hypothesis E.1. For S T, given p.q P. for ls E [K]] the following holds
(t) monotonically increasing, and p Pk
E [0, Ō(1):
a. ORo
b. if Ok.p # 1, then pp is monotonically decreasing and Dp E 11-)
c. 1 pp- = =0(pp- for m & {1} U fax.p):
d. for q 97 p. pp-q ((p,))
p-p = -
E.2 NEGATIVE INFORMATION GAP
Now we turn to the case that A < -((1).
Induction Hypothesis E.2. For < T, given q. P, for [K], the following holds
monotonically increasing, and pp E [0, O(1)::
a. S.
b. if a + 1, then $ D is monotonically decreasing and 1 D E {1-Ō(pLas) 0]:
c. ( 1p = Ō (pTt)))) for m $ {1} U {ax.p):
d. for q p. T p~g- = Ō (pr)):
e. r r pp - =
E.3 PROOF OUTLINE
In both settings, we can classify the process through which transformers learn the feature attention correlation & 0p-vv..u. into two distinct scenarios. These scenarios hinge on the spatial relation of the area within the context of the k-th partition DR. specifically, whether p is located in the global area of the k-tth cluster, i.e. whether akpp 1. The learning dynamics exhibit different behaviors of learning the local FP correlation in the local area with different A, while the behaviors for features located in the global area are very similar, unaffected by the value of A. Therefore, through Appendices F to H, we delve into the learning phases and provide technical proofs for the local area with A 2 ((1), local area with A (1( and the global area respectively. Finally, we will put this analysis together to prove that the Induction Hypothesis E. (resp. Induction Hypothesis E.2) holds during the entire training process, thereby validating the main theorems in Appendix
F ANALYSIS FOR THE LOCAL AREA WITH POSITIVE INFORMATION GAP
In this section, we focus specific patch E P with the k-th cluster for k E [K]. and present the analysis for the case that Xp is located in the local area for the k-th cluster, i.e. ag,p > 1. We will analyze the case that A > 2(1). Throughout this section, we denote Gs.p will analyze the convergence of the training process via two phases of dynamics. At the beginning of each phase, will establish induction hypothesis, which expect to remain valid throughout that phase. Subsequently, we will analyze the dynamics under such a hypothesis within the phase, aiming provide proof of the hypothesis by the end of the phase.
for simplicity. We
32

---

Published as a conference paper at ICLR 2025
F.1 PHASE I, STAGE 1
In this section, we shall discuss the initial stage of phase I. Firstly, we present the induction hypothesis in this stage.
We define the stage 1 of phase I as all iterations t < Ti, where
1 A U
( - 0.011 log(P))
T - max t:
(1) Ull.
We state the following induction hypotheses, which will hold throughout this period: Induction Hypothesis F.1. For each S s1STi.qE {p). the following holds:
a. $ (8) p uk.n
is monotonically increasing, and 4 (t) puknnn E $0,0
( -0.01 log(P)
po.02
b. $ p p
(1)
is monotonically decreasing and p 10
E El
+ ( 0.01) log( P),
4(5)
c.
1 1 pp = =0
(1)
pl-ns
06,1 for m + 1,m;
d. ((())
= =0
41t)
Povkn for ak.q = n. [ k.p-pp
/. =0
(()) km
-4(t)
p p
k.p->q
On
P
e. (r(t) [ = =0(
(1/000
o
o siime for ak.q = 1;
Ci
P
f. k.p-aql
(t)
for au.a F 1, n2
F.1.1 PROPERTY OF ATTENTION SCORES
We first introduce several properties of the attention score if Induction Hypothesis E.1 and Induction Hypothesis F.1 hold.
Lemma F.1. For 1, if Induction Hypothesis E.I and Induction Hypothesis F. hold iteration t < Tı, then the following holds
1. Attn (1) P-P... Attn P 2 2(1):
2. IfM Ek,,, Attn =0
(1)
(p-t.)) =((ptt))
3. Moreover, if M E ER, we have Attn() (t)
Ω
Ph.1
p
0.01
4. For q E MN (PE, U UK.1), attn = (1 p-A -Attn (t)
Pk.
Lemma F22 For 1, if Induction Hypothesis E. and Induction Hypothesis F.1 hold iteration t < Tı, then for m. # n, 1, the following holds:
Atttn P4.1
Attn POPKIN
(1) 1. For any q Pk.m attn S S0
-Att (t) P 5 1-Attn Pe
2. Moreover, Attn (4)
PL,
PL.
N
The above properties can be easily verified through direct calculations by using the definition in equation 2.6 and conditions in Induction Hypothesis F.1,, which are omitted here for brevity.
33

---

Published as a conference paper at ICLR 2025
F.1.2 BOUNDING THE GRADIENT UPDATES FOR FP CORRELATIONS
Lemma F.3. For 1, if Induction Hypothesis E.1 and Induction Hypothesis F.I hold iteration
0 S t Tí, then uu 2 and satisfies:
Cn 0
(A) pokm
0
P pI-Na
Proof. By Lemma C.2, we have
(1) Pu
E
1 = k, M)tttn P-Pann
2x
Attn
2 p-sPe,n )
+
+
2m 2n (Attn
)
myn
1(kx &x, upe M) E
(1) xx (1- Attn
Pkm
2
+
am on
(1)
2
1[kx = &E, Ope M)Attn (1) E
24 Attn
(t)
2
(Atnn
(1)
P-PPa.
P-PPra
p-
mige
<P(M
E Ek,n)
E
1fkx =,pE M)
2/ (2(11- Attop-
(1)
)²+
+ Σ
2m on
AAtn (1)
) &enm
myfn
+
+0(1) PPM E))
Cn
$ 0
+ Cex)
P
$ 20
where the second inequality invokes Lemma F.I and Lemma D.6, and the last inequality is due to erp((Ca,11.) Cn) << Co. Similarly, we can show that Op-rrr...
(t)
C. 2 +
Lemma E4. For n 1, if Induction Hypothesis E.1 and Induction Hypothesis F.1 hold at iteration 0 s s Tú, then u < O and satisfies
. p2 1---a---n 0.01) -000)) = 0
pu.98-
Proof. We first single out the following fact:
2 11-tttn p-Pp..))
Attn PeP -27 -zp(1 Attn Attpptpe.
(1-AAthn p+PPe.
Attn
Σ 2/
Attn (1)
(tttn )²
P
a#1.
2 max =e Pttttp--p af11,
22 Attnn (1)
2n Attmp--Pen
2 (1)
p-P. ("
(t)
Attn (t)
Attn (t)
P
(1) =121 Attn
Attnp-pul Atttt (1)
+ 2 Attn (0)
Pttnnp-pe.
max 2 Attn
(t)
0z1.
(F.1)
Therefore, by Lemma C.I, we have
(0) P
E E 1fkx = &, C M)Attn
34

---

Published as a conference paper at ICLR 2025
P-Paa )
-21 (1 Attn (1) P-
Attn (1)
p-Pa)) ( [2AAtt P-Pa + Afttt (0) p-PP,
max 20 Atttt
a1
1fkx k & a MP (1)
Pet
Σ
2/20
(ttnn (8)
p-ska
2
a
x P(M E Ex,1) (- (2((21) 22(
1)) + 0(1) P(M E7
p2x -
<-0( p2x( - p0.98
where the second inequality invokes Lemma F.1 and the last inequality comes from Lemma D.6.
Lemma F5. At each iteration S TT, if Induction Hypothesis E.1 and Induction Hypothesis F. hold, then for any m > 1 with m. + n,, the following holds
(1)
(2)
l0 (8)
0
Op th
app 201- ) =0
ap -v +p..
P
N
Proof. By Lemma C.I, for j we have
pitap-pped ))
a p
(t)
SE E
1fkx k. MJAttn (1) Pam
Z (Attn (1)
2
(F22
a.ptee, 12
- m
(1)
< E
(1) 1fkx - p P.PPk.m
(zm 2
(2)
Attn (1)
p-tk.
+ (1 (1-Attn (1)
p- Atttpp-p...
Attn (1)
(F33
For equation F.2, we have
Ma.m
< E
1 1skx = k Ex,1 E Npe M) (
P-.m. Pk.m
Σ
am
(Attn
(1)
atmmm
+ E
1skx k. ( npe A)tttn P-Pa.m
Σ afmm..
2 2m
(Attn
Attn (t)
Attn (t)
< E
1[kx- E, E ADE MJO
P-PP
p Phhn
N
(Attm )² + +0 + NV
)
+ +0(1) P(M E (Ex,1 n )9
am
0
+ 0(11 P(M E (Ex,1 n Ex)
N
< 0
app- pI-K,
where the second inequality is due to Lemma F.2, the last incquality follows from Lemma F.4 and Lemma D.6.
On the other hand, for equation F.3, we can use the similar argument by invoking Lemma F.2 and Lemma F.3, and thus obtain
<0(
ap ap-rvun
0 7mmm
pl-ke
35

---

Published as a conference paper at ICLR 2025
Putting them together, have
(1) (0 t02
(1)
la mp S 0
pl-ka
F.1.3 BOUNDING THE GRADIENT UPDATES FOR POSITIONAL CORRELATIONS
Lemma F6. For n 1, if Induction Hypothesis E.1 and Induction Hypothesis F.I hold at iteration
0 S t < Tú, then for P {p} and ak.q n, we have a!t)
Pk,p-q 2 0 and satisfie::
(t)
(11)
P = θ
Cn
Furthermore, we have 1310 k 0
(4)
-ap ppok
(2)
Proof. By Lemma C.2, for q E Pk.n with # p. we have
g(t)
1fkx k.p M.q UJattn 2n (1- Attn (0)
)'+ Σ 2
(Attn (0)
E
p
myn
H1
1fkx = k.p M.qe M}attn M (- Attn (1) E
PeePe,n (1 (1- 1ttttn
+ 1skx= = k E Ma p
Σ zm (Attn
2
f3
Firstly, for H1, notice that
2 + + am (Attn Pa.m 2
(C, - 10 = 1[kx= = k.p M)
ma
=(aa
For H2, since p. E M, by Lemma F.1, we can upper bound attn by 0 ( thus
(2)
(1- (1- (= (t) H2 S. = =.px=== (S,AAt))
10) pt ()) < <0(
Further notice that H3 can be upper bounded by O(Hi), putting it together, we have
(
Pk.p
A
ap-tve.n
C₂
Turn to g(t)
k.p--p' when q =
8 = E 1[kx k, M)attm ()) (- Attn (1- Att
36

---

Published as a conference paper at ICLR 2025
+ E lfix ki MJattn
Σ z
(0) (Attn
2
P
We can bound J2 in similar way as Hz. Thus, we only focus on further bounding J3:
1[kx kE M)0( Attn - Attn (5) p-
p-Pam )
J3 S
Σ
Fm
(Attn (t)
|ap Jap-omal
where the first inequality holds by invoking Lemma F.1 and the last inequality follows similar arguments as analysis for equation F.2.
Lemma F.7. For 1, if Induction Hypothesis E. and Induction Hypothesis F.I hold iteration
S t Tú, then for P {p) and az.q 1, we have Sk.p- satisfies:
lap app- ap-ezal
|ap ap-
(1) PRpp = 0
0
P
C1
Proof. By Lemma C.2, for q E Pk,1, we have
g(1)
k. k.p p
- M.x=
27 Attn Attm P-P (1
Attn P-PPe
) 2 (1 Attn ) Attn
Σ
(Attn p
)²)
af1
(F.4)
E 1fxx = k M,qe MJattn
(i) (1
Attn Pan
Attn
62
+E 1lkx k, M,q M)attn
Σ z
(0) Attn
For equation F.4 denoted as G1 following the direct calculations, we have
-(1 - 1)G1 =((a p-
We can further bound G2 and Gg in a similar way as H2 and Hg in Lemma F.6 and thus obtain
(0)
G2 s 0
Pp
(1)
apptraa
G3 S 0
which completes the proof.
Lemma F8. For n > 1, if Induction Hypothesis E.1 and Induction Hypothesis F.I hold at iteration
0 S l < Tú, then for q P (p) and + ak.q. Pk.p satisfies:
(1)
(1)
|stpp_ =
ap tuen -pp-ta.11
P
37

---

Published as a conference paper at ICLR 2025
Proof.
By Lemma C.2, for q E Pk.m we have
((1)
K.p p
- M,
Attn P- P-P (1- Attn (0)
p- ) + z (1-Alt P-Pmn Attn
Σ
z (Attop-72
mm
ayn,m
(F.S)
1fkx k, pe M,qe M3attn p (2) +1 Attn
Attup-
1fkx = k.p M.qe Maattn
Σ 2 Attn azo
(c)
pooko Pkio
F.5 can be upper bounded by o (t) ppp-..mm
(
(1) p--.%
(8)
equation
Cm
NC..
-op ) where the first equality holds by invoking Lemma F.5. 12 and la can P be bounded similarly as G2 and G3, which is omitted here.
o
F.1.4
AT THE END OF PHASE I, STAGE
Lemma F.9. For > 1, if Induction Hypothesis E.1 and Induction Hypothesis F. hold for all o 0 (og(P)
) At iteration t Ti + 1, we have
"
(T++1) S - a. 4
1 ( -0.01) log(P);
(1++1 ( p(i-xs)+ f (4 -0.01)
b. Attn
P-P..
Proof. By comparing Lemma E.3 and Lemma F.4, we have la (ap2tu,,1
>>
(8)
Then the
log (o..0)) directly follows from Lemma F.4.
existence of Ti.k O
F.2 PHASE I, STAGE 2
During stage 1, of2...., significantly decreases to decouple the FP correlations with the global feature, resulting in a decrease in Attn while other Attn with m 1 remain approximately at the order of O ( (0 (p.L)))) By the end of phase I. (Attn )² decreases O 222.002-2. ). leading to a decrease lap- as approaches towards (1) ap-+ve... At this point, stage 2 begins. Shortly after entering this phase, the prior dominant role of the decrease (8)
of 4 4p 1 in learning dynamics diminishes as lap (ap- reaches the same order of magnitude as (0) ap-
(0)
define stage of phase I as iterations Ti < 1 Ti, where
(0) 7 max t > Ti p
(1)) p
<
A
+
0.01
+
ci - Ka)
log(P) }
2L
L
U
for some small constant of > 0.
For computational convenience, we make the following assumptions for ki and Kg. which can be easily relaxed with the cost of additional calculations.
1 1 0.01
0.01
c3(1 Kx))
-- 1
<
(F.6a)
2
U
38

---

Published as a conference paper at ICLR 2025
+ (-)))==))) s - + 0.01)
(F.6b)
Here ep is some small. We state the following induction hypotheses, which will hold throughout this period:
Induction Hypothesis F.2. For each Zi < t - 7ú, E {p}. the following holds:
a. L Phen is monotonically increasing, and 4 pm
[0, Si
log(P)
b.
is ( + 0.01) log( (P). - ( - 0.01) log(P));
monotonically
decreasing
and
p
p-
(e)
110) 1v.n
(0 1
c. 1 (1(
0
for m + n,
pt-
d. r(()
0
(1(9]
p-ueee for akgq
ni 1Tppp k =0
1
k.p->q
On
P
[0((1)
100 0
2011 panken
c. -.ppp-og
(1)
0
p-
for ak.a 1.:
C₁
-0())
f. K.ppp-
(t)
=0(
for Ok.q F 1,
F.2.1 PROPERTY OF ATTENTION SCORES
We first single out several properties of attention scores that will be used for the proof of Induction Hypothesis F.2.
Lemma F.1. if Induction Hypothesis E.1 and Induction Hypothesis F.2 hold at iteration Ti + 1 S
t < Tú. then the following holds
1. 1 Attn P-Pee Attn (1)
P-Pa, 2 0(1);
2. if M E an Attn (1)
P-P E
[((xtta)), O ),0)) (
3. Moreover Attn (1) P-PPl
=0 =0
- -003-0000))) .011
ifM Ex, we have Attn (t)
(pll-nelt /
l pa nell 4 10.00)
4. for q E M kn UPa.), p 0 Attn
Attn A
Pe.
Lemma F11. if Induction Hypothesis E. and Induction Hypothesis F.2 hold at iteration Ti + 1 t < Tú, then for m # n, the following holds:
1. for any q E Pk.mo attnp-g <0(- 1 1 - AAttn
(
Pk.n
Attup--F..
2. Moreover, Attn (t) Pk.m <0(1 Atttn P-P -Aten
N
F.2.2 BOUNDING THE GRADIENT UPDATES OF FP CORRELATIONS
Lemma F.12. For n > 1, if Induction Hypothesis E., and Induction Hypothesis F.2 hold at iteration (8)
T 515 Tí, then o Op-rrunn and satisfies:
= (1(p1-m))
Op
p1-ma)
39

---

Published as a conference paper at ICLR 2025
Proof.
By Lemma C.2, we have
P Uk.
(Attm p- )
E
= p. p-
P-Pa
23 2n
Attn (0)
p-PPe )
p
+
+- =m
(Attn
1 = Ek, Op Mtttt P-P E
2n (1 11 Attn P-PPe,n )
2
+
Σ
2
En
AAtn
2
myln
+ E
1fxx = k En E M)Attn (1)
2n ( (- Attn (t)
p P-PPin
2 +
Σ
zm
(Attn (1)
P-ok.m
2
PMM
E Ek,n)
E
1(kx =k (0
P-PPen
2n Attmp--aen
3 (1-
(1)
p-sPain
2 )+
Σ
= =n
Attt EER.
(1)
On P
2 1
where the last inequality invokes Lemma F.10.
Lemma F.13. For I, if Induction Hypothesis E.1 and Induction Hypothesis F:2 hold iteration (t)
T + --1st < Tú, then ap < and satisfies
|a u
p2(1- (((At000))))
Proof. Following equation F.1, we have
-212 (1- Attn (1-Attn P- Attn P-P
(1)
2/ (1- Attn (1)
(1-Att D-PPu.1
Attn
Σ 2
(Attn
2
-
P-PPA.
af1,1
--1111 Attn (t) P-Pkn
Attn p.pp..) Atttt (t)
P-P + +2 (1)
max = Attn (t) af1,n
P
P-Pri
Therefore, by Lemma C.I, obtain
P - E
1[kx k, & DPt Mttt P-P
(
(-2 (1- Attn (1)
P Attn (t)
p.P..)
2 Attn (1)
P-P. + #
(t)
max 2 Attn (1) af1,n
+ E
4 1lkx K6 P P-P
Σ
zîzo (Attm (1)
P-Pk.a))
af
(- ( 2(1 Ω ke(1-xc)+ # # )
< P(M E Ex.1)
+0(1) PMM E E
< <-2 p2(1-x.)+((A40..00
where the second inequality invokes Lemma F.10 and the last inequality comes from Lemma D.6. The upper bound can be obtained by using similar arguments and invoking the upper bound for Attn (1)
PP.. in Lemma F.10.
40

---

Published as a conference paper at ICLR 2025
Lemma F.14. For 1, if Induction Hypothesis E.1 and Induction Hypothesis F.2 hold iteration T + Sis Tı, then for any with # n, the following holds
la (0 p S 0 - - -0p--up +Uk1
P1-ka
The proof is similar to Lemma F.5, and thus omitted here.
F.2.3 BOUNDING THE GRADIENT UPDATES POSITIONAL CORRELATIONS
We then summarize the properties for gradient updates of positional correlations, which utilize the identical calculations in Section F.1.3.
Lemma F.1.. For >> 1, if Induction Hypothesis E.1 and Induction Hypothesis F:2 hold at iteration T1 + Sts Tú, then
a. if 0n,9 = and q +
2 0,
((t)
Pk.p-a
+
P
On
and |33.01
0 p.r. -
(t)
b. if ax,q
-1, k.ppp-- = =0("0 p- =0("
-a -p-
1 k1.1
+ 0
11 (6)
c. if "x m and m # 1,m, (3(() 10 0
pp.
F.2.4 END OF PHASE I, STAGE 2
Lemma F.16. Induction Hypothesis F:2 holds for all iteration +1 S 1 s T,
0
l) ne) and at iteration = Ti 1. have
9
(1+
(1(,-x,) log(P))
a. T P-
2
U
I +1)
b.
p 26
+
L
log(P).
Proof. The existence of Ti = Ti + 0 +0 (P)((P)
directly follows from Lemma F.12 and
"
(t)
Lemma F.13. Moreover, since ap 1 < 0, then
(++ p P-kukn
<
0.01
c1(1
Kig)
log(P) P)
A T 0.01) <
(cn + ci) Ka)
log(P) ())
2L
L
U
U
where the last inequality invokes equation F.6a. Now suppose Pp-ven _[(I=n))oR(( log(P) then
ob (1+1)
o P-
(7++1)
- 2L + 0.01 log( P). Denote the first time that of reaches
-22 + 0.001 log( P)
(t) T. Note that T Ti since ap- the change of & p satisfies (ap- log( (P). Then
(t)
for 2 T, the following holds:
(8) 1. Attno-op 2 ll
(p:));;
2. Attn (1)
P-PPA 50
I-xa 10.001
Therefore,
la (0
$ 1fkx= = PpE M P-P
(0)
1ap-
2 Attn P (1- (1- -ttt Attn (1)
P-PP, (1 Attn (9
P (
21
41

---

Published as a conference paper at ICLR 2025
+
1 1 = K, Op
M)Attn (0)
P-,
Σ ztt0 (Attn
(t)
2
a#1,1
0
Op-ven.1
(4)
+ P(M E &
(0(1)-
+0(1) P(M Ex21
p p +0.001
p 1K
(t)
0
0
p p +0.001
p(1-N,)+0.002
Lemma F.12 still holds, and thus
(1)
la 1ap --E | S 150(
p0.002
+1) Since [ pm D 20,,
Ω (log(P)) we have
$ (1++1) $
+1)
T p 1. 52 (2pp.002)
+ ((T)
> ?(pb.0u2 log(P)).
p-
p
p-
which contradicts the assumption that T +1+ pete
< of (111-x))og((P
F.3 PHASE II, STAGE
For n > 1, we define stage of phase II as all iterations
Ti + +1sst S T2, where
Ty - max t: (1) p-ue - (1-kss log (og(P))
L
We state the following induction hypotheses, which will hold throughout this stage: Induction Hypothesis F.3. For each Ti +1 T2, q
{p), the following holds:
a. $ Pkk
is monotonically increasing, and 4 (t) p
c (1-x)) log(P), (1-xa) log (P)
U
L.
b. + (8) 1 is monotonically decreasing and
E
-
[[-
1
A
0.01
log(P)-
( A -0.01 log(p)}:
2
U
2
((t))
(0)
c.
pp-vve,
(t)
vk.n
pi-se
for m # 1,m; pi
d. (())
k.p->q
0
(11)
p-opp for
dk. forak.
11, (1) kpp- 0
p
(1)
P
p
((((()
+ 0 2(0 pa spmrm for ak.q = 1.;
e. [
0
Ci
P
((1)
- -00
f.
(t)
for Ob.q * 1, n.
F.3. PROPERTY OF ATTENTION SCORES
We first single out several properties of attention scores that will be used for the proof of Induction Hypothesis F.3.
Lemma F.17. if Induction Hypothesis E.I and Induction Hypothesis F.3 hold at iteration T, +1
L < T2, then the following holds
1. if M E EK,,. Attn 11) Pu. 2 2
p0-- y Moreover, if Attn (1) P does not
reach the constant level, - Att 2(1); otherwise, 1 Attn (1)
(1)
P
Ω (1
42

---

Published as a conference paper at ICLR 2025
1-Attn C
2. Attn (9)
p-tk = o
p- h.n
); if M E Ek,1, we have Attn
pil-neht 617 0.01)
of (in--e) 2 +0.01)
3. for q M (Pk,n U Pk..), attn = =0( |AAtt (8) ( p- po-Tk.n
(2)
Lemma F.18. if Induction Hypothesis E. and Induction Hypothesis F33 hold at iteration Ti t < T2, then for m + n, the following holds:
1. for any q E Pk. attnp-n S 50( 1 --tt (c)
p-
Pk.n
00 AAttn (2)
-Attn
2. Moreover, Attn (8)
P-P
F.3.2 BOUNDING THE GRADIENT UPDATES FP CORRELATIONS
Lemma F.19 if Induction Hypothesis E.1 and Induction Hypothesis F33 hold at iteration Ti + 1 S
t < T2. then av- 2 0 and satisfies:
a p 2 min
) ( 52( p - )(1 )1- -2(4 (1(1-1)(1xx)) -
1
Proof. By Lemma C.2, we have
(t) pp-io Ukkm
E
1[kx = k.pE 10)
3 2x
Attn (0)
Σ
Am
2
[Attmp-pae Pk.m
(1)
mfn
E
1[kx = Ex, Ope M)Attn
P-Pnn
3 { (2((1 (1- Attn p-PPen
+ Σ
2 (Attn
p-sklem
+ E
1fxx = k, Ei Op M)Attn p-
22 (1 Attn (t)
1-Attt
puctppp-i P-Pan
+
Σ
2 2n 2n
(Attn p- P
PM E E.
E
1fkx k,P M]Attn P-Pa-
(1 Attn (1)
p-PPA.
)²
2 + Σ
2 2n
(Attn
p-sp,c)
)² Ek.
mfnn
+0(1) PPMMPPMM E E
{
.0 pa14 - -)()) (pa))-ha-m))
min
p(1 of X1-xx)
where the last inequality invokes Lemma F.17 by observing that for E Expes
Attn Attnp-pan Pn (1 Attn p- )² 2 min 2 p(1 (0)0))))
2(1), pex(t-1X(11
Lemma F.20. For 1, if Induction Hypothesis E.I and Induction Hypothesis F.3 hold teraaion (t)
T + 1 565 T2, then ap- and satisfies
la (x)
mu 1 min (
-101 (
}
1
p(1 46) (0 (1-)))))
-1
p(--x)+ K 0.01)
43

---

Published as a conference paper at ICLR 2025
(4) Opre
(1) ap-ue.
1a pu S max {o(
p(1-8.)+ (AA22-001) +(a--00)+ 11
- (5)((-n))) }
Proof. Following equation F.1, we have
-21 (1 Attn Attnp--pppp P Attmp-PP.. P - =z =1 (1-Attn
Σ = (Att )
2
Attn p-PP.l
a
$ --1 Attn P-Pknn Attn Pp..)) AtAttn p-Pk.n + + Atttt
(1)
max 2 Attn (+)
p-P
af1,
Therefore, by Lemma C.I, we obtain
( a
E 1fxx- k, Ep, Pp E MMAttm (1)
P- )
(4) 21 (1 P
Attn (8) P-PPA.
2 Atttupp- +
(1)
A2ttt (1)
max egAtt 1lnn
E
1kxx = K Ope (pE
(0)
20 Attup--Pael
P-P
af11,1
< - min {
p(1 (1)))x))
-IC--)(--N))
where the second inequality invokes Lemma F.17 and equation F.bb. Moreover,
Ja laf2
S
1[kx = k & EE Op ENEMM A]tttn
P
2 - Attn Attn P-.Pi..
Attn P-PP.. + (1-Attt P-PPt.1)
Attn (1)
Attn P-Pra.)
E
[kkx = k Ex, E
OpEMME 22 ALt (1)
p-Pa, (1 Attn Phin
Attn P-P.n.
1skx = k. 1 Πpe M32) (Attn, (8)
p-P.. }2 ( Attn
+ E
<max {o p(l-ke)+ ((
ap-erm
ap-vken
p2(1-k.)+
-o.0))
0.01)
- (1 (1-xx,)
where the second inequality invokes Lemma F.17.
Lemma F.21. For n > 1, if Induction Hypothesis E.1 and Induction Hypothesis F.3 hold at iteration T + t < T2 for any m > 1 with m + the following holds
(1)
(1)
m S 0 0p-+v U2 Cp
la
pl-ka
The proof is similar to Lemma F.5, and thus omitted here.
F.3.3 BOUNDING THE GRADIENT UPDATES OF POSITIONAL CORRELATIONS
We then summarize the properties for gradient updates of positional correlations, which utilizes the identical calculations as in Section F.I.3.
Lemma F.22. For >> 1, if Induction Hypothesis E. and Induction Hypothesis F3 hold at iteration T + 1 stt T2, then
if 0k.a and q # p.
(0)
0:
((0)
O
P vkk.n
and 1,300
0
a un
(t) (e)
10211
p
(1) -o
(1
[aft)
b. foo 1, 1,
0
p-tuk
0
C2
c. if ookq m and m * 1, 1/3(t)
0
(2)
(1)
44

---

Published as a conference paper at ICLR 2025
F.3.4 END OF PHASE II, STAGE
Lemma F223, Induction Hypothesis F.3 holds for all T, + 1 S t < T2, and at iteration = T2 + 1,
have
(8) 4 Mhin
(1-R,) log(P):
b. Attn p- [1) ifM Ek
Proof. By comparing Lemma F.19 and Lemma F.20-F233, we have (0 ap-rre, > 10p kppp
Then the existence of T2 Ti + ( log(P))r
directly follows from Lemma F.19, where
"
A = max {1-
-1--11
G 2 -
Kg). (1-Kg).
(1
The second statement can directly verified by noticing that + (1-ma) log(P) while all other attention correlations are sufficiently small.
F.4 PHASE II, STAGE 2
In this final stage, we establish that these structures indeed represent the solutions toward which the algorithm converges.
Given any 0 < c << 1, for > 1, define
3 1/4 - N }
T5 max t T2 (10)
I log C5
(F.7)
where cs is some largely enough constant.
state the following induction hypotheses, which will hold throughout this stage: Induction Hypothesis F.4. For n > suppose polylog(P) P) log(z)) for each T2 + S Tý
E P (p), the following holds:
a. $ m0am
is monotonically increasing, and (1) p [ log( (P), (P/)))
b.
is monotonically decreasing and 4 p E
( + 0.011 log( (P)
1
(11)
1 0
(
0.01) log(P))
500 for m + 1, n; 1921
@ (1)
0
c.
pi-aa
d. ((())
=0
4(1))
iivnn for akgq n, (T(()) k =0
(5)
bek
Cn
(0(1)
((0)
503 Di for ak.q 1.;
k.p
0
0
C1
O(0)
((0)
(!))
k.p
0
for Og.q * 1, n.
p
F.4.1 PROPERTY OF ATTENTION SCORES
We first single out several properties of attention scores that will be used for the proof of Induction Hypothesis F.4.
Lemma F24. if Induction Hypothesis E.1 and Induction Hypothesis E.4 hold at iteration Tn, < S
The then the following holds
45

---

Published as a conference paper at ICLR 2025
1. i/M an Attm 2(1 and Attm 2² O(e).
(0)
(2)
( nell b1 0.01) ): ifM E Ek,1, we have Attn (1)
=0 1-Attn Pain
2. Moreover, Attn P-PP. 1-Attn (t) P-P p 10.001 (0.01)
p- =
1---- Attn (c) 3. MEQEMME (PR,n U PK,)). a -0
(1)
p- Pk.m
Lemma F25. if Induction Hypothesis I.1 and Induction Hypothesis F.4 hold at iteration Tu,2 < <t< TV 2. then for m # n, the following holds:
1 Attnn p-
1. for any q E Pk.m, attnp-1 50((
Attn
2. Attn p-Pan
and if M E Ek Attn Pua N
+ 1- A-Att P-
F.4.2 BOUNDING THE GRADIENT UPDATES OF FP CORRELATIONS
Lemma F.26. For > 1, if Induction Hypothesis E.1 and Induction Hypothesis F.4 hold at iteration T3 + St < T2, then di tpkn 2 0 and satisfies:
apre s(()).
Proof. By Lemma C.2, have (t) Pur.
E
lkx k.p M))
2x Attn (0)
)²
Σ
zm 2n
(Attn
(
P-PP.r"
p-PPe,a
mnnnn
1(kx Ek, Op Mttttn p-Pen E
3 in (1 (1- Attn (0)
Attmp-Pee.
)²
Σ
2m
2
(1) (tttn
mofn
( 2/ (1-Attt
+ E
1flx = K Op M) p- Pk.m
Attn (8)
P-
+
22
(t) pAtto--p
2P P(M E E.
1(kx = ki -.. Mtttm (9) pttp 3 (1- Attn (1) P-Pee
)²
Σ
(Attn p-Pmm
(Attn
)² Exen
zn
mm
+0(1) (MM)) ) 2 (())
where the last inequality invokes Lemma F.24, Lemma D.6 and the fact that
2 (op(-polylog(X)) >> exp C-Cn,
Lemma F.27. For >> 1, if Induction Hypothesis E.1 and Induction Hypothesis F.4 hold iteration (t) Tn,3 < 1 S x then ap-ovt,1 and satisfies
(1) 1a max {o
ap Uk.
)
Op Uk Uk.n
p(1-kc)+ (1/2 -0011)
(2((1-x)) (A 0.02) (1 3 (1-k.))
46

---

Published as a conference paper at ICLR 2025
The proof follows the similar arguments Lemma F.20 by noticing that P(M (k,mm for any
m # n.
Lemma F.28. For n > 1, if Induction Hypothesis E. I and Induction Hypothesis F.4 hold at iteration T3 < s T2, thenfor any m >> 1 with m # n, the following holds
(t) ap Puk. S 00 p-p.k.m 0
0
pl-xa
Proof. We first note that
-2 2 1 - Attn p- Pk.
(1)
Attn Attu
- Zym A1-Attu
Attn (1)
Σ
(tttn
(1)
2
p-PPea
p-
+
am
af1,
2m
2a a.
max
2 Attn (4)
P P-PPa.n
2 Attn
(1)
- Attn (1)
p-Pk,m (1-Attn
P-PPE,n
(1)
Attn
A-tnp-PPe p-
< <-0(1 - Attn P.....
since when M Ek have Attn (t)
P-PPu. = -2(1 Attn
(1)
Thus, we have
0 2 0 (1)
P- 2 E | = k, Ek Attn (0)
p -211 Attn (
P- !
P-P..m
2 -0 Optu
P1-ka
F.4.3 BOUNDING THE GRADIENT UPDATES OF POSITIONAL CORRELATIONS
We then summarize the properties for gradient updates of positional correlations, which utilizes the identical calculations as Section F.1.3.
Lemma F.29. For >> 1, if Induction Hypothesis E.1 and Induction Hypothesis F.4 hold iteration Th + St S T5, then
a. if 0k.q = n and q # p. Ph.p-n 2 0;
((t)) Bhpp-g = O
aaf-eee and 13!50) = Ca
(1) (c)
p-rkkn p-
b. fkkkq 1, 1113%%-p-o =0 =0
(ab (upn bukin
-a(t)
20n,1 + 0
C
(6) -a
ifok m and m + 1, n, (astt) pppp--al 0
p uo
p 16,1
F.4.4 END OF PHASE II, STAGE
Lemma F33. For > 1, and 0 < < 1, suppose polylog(P) >> log(?)- Then Induction
Hypothesis F.4 holds for all T2 < t < T5 T2 O log(P( and at iteration t T + have
1/c
1. Cn.p(Q"i+)) 2k
2. If M E Ex,,, we have (1 - Attn )² < O
(T5+1)
The existence of Tú T2,x + O log((P derive
Proof.
directly follows from Lemma F.26. We further
1/c
Ex. (Q
47

---

Published as a conference paper at ICLR 2025
E
1f4x k. K.
o (1- Attmp-
-Pam)) 2 Σ
Am (m )
$ 1
.y.U2
.. o1)) O())
2K
2K
where the first inequality is due direct calculations by the definition of T2. and the second inequality can be obtained by setting Cn in equation F.7 sufficiently large.
G ANALYSIS FOR LOCAL AREAS WITH NEGATIVE INFORMATION GAP
In this section, we focus on a specific patch p P with the k-th cluster for & E [K]. and present the analysis for the case that Xp is located in the local area for the k-th cluster, i.c. ak.p > 1. Throughout this section, denote ak.p n for simplicity. When A < -1((1), we can show that the gap of attention correlation changing rate for the positive case does not exist anymore, and conversely
(t) ap Op- ap- 0 from the beginning. We can reuse most of the gradient calculations in the
(t)
previous section and only sketch them in this section.
Stage 1: we define stage as all iterations 0 I L S Tng,,1, where
(1 (1-Ra)
K.) log (P)
T Fneg,1 max (()
p-
L
We state the following induction hypothesis, which will hold throughout this stage:
Induction Hypothesis G.I. For each 0 < <LS 1 1, E {p), the following holds:
(t) a. p p monotonically increasing, and $ P (, (1=Na) log( P)
b. $ ( p- is monotonically decreasing and $ P
P
b (t) ppmu
p 1i.
for o 1, 1, p- m.l
(1))
(2)
(8) p 1g.
d. (1))
h.p->q
= =0
Cn
for
ak.q forak:
=n, (r) k = =
p
(n(t) m
= 0
(0(0)
2 p-
(00 pposee for ak.q = 1;
e.
Ci
((()) 0
for ax.q f n.
k.p
Through similar calculations for phase II, stage I in Appendix F.3, we obtain the following lemmas control the gradient updates for attention correlations.
Lemma G.1. If Induction Hypothesis E.2 and Induction Hypothesis G.I hold for S t S T Tneg. then we have
D
1 (
min
(G.la)
P
p214
-)())))))))))
(1)
-0
(G.1b)
P P-A
(1)
(1)
la
50
aptrun
Op-rua.1
for all m + n,
(G.lc)
p1-ka
48

---

Published as a conference paper at ICLR 2025
(t)
P.p-aq
=
O
for DR.q = =n, -1,
(G.1d)
1
0
ap-tvkin
1
for ak.9 =1, 1
(G.le)
C₁
(1)
1 k.ppppp
1kppp~g =0
apttun
for all ak.p F n, 1.
(G.1f)
P
Here A < 0 implies vu,,
< Up-tve,n' Induction Hypothesis G.1 can be directly proved by
(t)
Lemma G.1 and we have
p (1nnxx)11111)11(11--1111 1))-111-8.) log(P)) (
Tneg.1 0
(G.2)
Stage 2: Given any 0 E < 1, define
(())
ThE I max
1 Ti
log
CB
(G.3)
where ce is some largely enough constant. We then state the following induction hypotheses, which will hold throughout this stage:
Induction Hypothesis G.2. For n > suppose polylog( P) >> log(?), for q E P!! (p), and each
T 1l < t S 1 the following holds:
a. T (1)
pu monotonically increasing, and + (t)
(1-x.)
log( P), ((log(/())))
b. T (8)
pp is monotonically decreasing and $ p
D 4
0
$60 = =
c. 1 pm
vAll pi-ny
for m + 1,
4(1)
d. ((()
Pove for ax.q
n2 ( k.p->pl - 0
k.p-iq
Cn
(t)
0
(11000
(0(00)
O
1l0) lukn
P 1ukk for ak.q 1;
C1
(1)
(c) for ak.q # 1, n.
f. (t)
k.p-tq
Lemma G.2. Induction Hypothesis E.2 and Induction Hypothesis G.2 hold for T Tueg,1 <t S T I then we have
2 >((E),
(G.4a)
Pu Uk.n
(1)
0 2 a (1)
2 -0
ap-+vk.
ap
(G.4b)
P-A
(t)
la 1
100
ap
for all m + n, 1
(G.4c)
p1-ka
(t)
0
for forak.
(G.4d)
On
[0(0 1p---1 = 0
Op-sven
(t)
+ 00
/ap-ve.al
for as.. = 1,
(G.4e)
P
C2
(1)
(1)
kpp-pl
11111p = =0
ap-rven ap-rmm))
for all ak,p + n, 1.
(G.4f)
P
49

---

Published as a conference paper at ICLR 2025
Induction Hypothesis G.2 can be directly proved by Lemma G.2. Furthermore, the end of this stage, we will have:
Lemma G.3. Suppose polylog(P) P) >> log(?), then Induction Hypothesis G.2 holds for all Tueg.1 <
1. < Teg..1 = Tneg,1 + ) ) and at iteration L = T neg.1 + 1, we have
1. Cn.p(QTiexiite 2K
(T (Tmee (1-Attup pomp S O(e).
2. If M Ex.n. we have
H ANALYSIS FOR THE GLOBAL AREA
When #p,6 1, i.e. the patch lies in the global area, the analysis is much simpler and does not depend on the value of A. We can reuse most of the gradient calculations in Appendix and only sketch them in this section.
For Xp in the global region PA,L, since the overall attention Attn p PA.1 to the target feature already
reaches 2 C 12 pi -p(ptte) due to the large number of unmasked patches featuring 2k21 when
(0) M E E which is significantly larger than Attn P-Pr..m θ
for all other m > 1. This
(t) results in large ap- initially, and thus the training directly enters phase II.
Stage 1: we define stage 1 as all iterations 0 1 Te, where
Te,1 - max
(t) P
S
(1 Kie)
1og(P1}
We state the following induction hypotheses, which will hold throughout this stage: Induction Hypothesis H.1. For each S t Te,1, q {p), the following holds:
a. + P is monotonically increasing, and p E To. (1-me) log(P) (P)):
I.
b. $ p-sm..m is monotonically decreasing for m > and p - (ogg ).0
c. ((()
h.p->q =0
(11
for Ok.a
O(0 ( k.ppppp O
Ci
(t) d.
k. =0(
for Ok.q + 1.
Through similar calculations for phase stage I in Appendix F.3, we obtain the following lemmas to control the gradient updates for attention correlations.
Lemma H.1. If Induction Hypothesis E.1 (or Induction Hypothesis E.2) and Induction Hypothesis H.l hold for < t S Te,1, then we have
2 min ( (ppp-m))..0 (poLt S0C4--D1--0 -ko) ((pr))--))))
( u
(H.
(1)
1a
(1)
<
ap-va
for all m + 1,
(H.
+u
pl-ke
((t)
Pkppiq
= =
(p--ve,1
for akk. 1,9
(H.1c)
C2
(t)
[ Appppl
= 0
(ap-ou.1
for all at,q > 1.
(H.1d)
P
50

---

Published as a conference paper at ICLR 2025
Induction Hypothesis H.1 can be directly proved by Lemma H.1 and we have pmax{1,2( - 11)--1)----1- xxe) (P)
Te,1 =0
(H.2)
7
Stage 2: Given any < 1, define
(()
T A max
t To,
20 pu S log
pt pl-re)
(H.3)
where e7 is some largely enough constant. We then state the following induction hypotheses, which will hold throughout this stage:
Induction Hypothesis H.2. For > 1, suppose polylog(P) >> log(?), q E P (p). for each
Te,1 + -ists T2 the following holds:
a. pp t is monotonically increasing, and $ pp-res E [1-ka] log( P), O(log
b. + pu is monotonically decreasing for > and d pm
[lo((P) ),0
419) = =0
((9) 1p-mn.11 for Ok.q = ( kpp-pl = o
c. (())
k.p->q
C1
( for 0k.q + 1.
d.
( (t)
k.p- 0
p-vya.1
We also have the following lemmas to control the gradient updates for attention correlations. Lemma H.2. If Induction Hypothesis E. (or Induction Hypothesis E.2) and Induction Hypothesis H.1
hold for To,1 + S T then we have
(t) 0
(ap+Uk.1 ) pl
2 (a(c))
for all m +
(H.4a)
1 1
Pk. Mk.m
(2)
3(1)
=
for dz.q 1
(H.4b)
kpp-
1.1. 1 = 0 (pp
for all azaq 1.
(H.4c)
Induction Hypothesis H.2 can be directly proved by Lemma H.2. Furthermore, at the end this stage, will have:
Lemma H.3. Suppose polylog(P) >> log(z)) then Induction Hypothesis H.2 holds for all Te.
t t T5, = TT,1 + +0
log(Pr 1
and at iteration t = To,1 + 1, we have
1. (x.p(QTait)) 2K
(1)+
2. I/M E Ex,1, we have
Attn < ((e)..
PROOF OF MAIN THEOREMS FOR MAE
I.1 PROOF OF INDUCTION HYPOTHESES
We are now ready to show Induction Hypothesis E.1 (resp. Induction Hypothesis E.2) holds through the learning process. Theorem I.1 (Positive Information Gap). For sufficiently large P > 0, 17 & log( P), 2(1) A
1,
Induction Hypothesis E.1 holds for all iterations 0, .
O (olylox(P)
Theorem I.2 (Negative Information Gap). For sufficiently large P 0. <
log( (P),
-
0.5 < A <
-O(1, Induction Hypothesis E.2 holds for all iterations L = 0, 1, =
n. palylog(
51

---

Published as a conference paper at ICLR 2025
Proof of Theorem L1. is easy verify Induction Hypothesis E.1 holds at iteration f 0 due to the initialization Q Odxd- At iteration / >>> 0:
Induction Hypothesis E. can be proven by Induction Hypothesis F.1-F.4 a and Induction Hypothesis H.1-H.2 a, combining with the fact that log(1/c) < polylog(P).
Induction Hypothesis E. lb. can be obtained by invoking Induction Hypothesis F.1-F.4
Induction Hypothesis E. cc be obtained by invoking Induction Hypothesis F.1-F44 and Induction Hypothesis H.1-H.2 b.
To prove Induction Hypothesis E.1d. for q o p. p--q = - T k.p-'q By item d-f in Induction Hypothesis F.1-F.4 and item c-d in Induction Hypothesis H.1-H.2, conclude that no matter the relative areas q and p belong to for specific cluster, for all & [K] throughout the entire learning process, the following upper bound always holds:
r(() k.p-9
max 1
( poee, max
(()
S {
KE[T] pu
Moreover, since K O(polylog(P)) (P)), we then have T Ō(tt), which completes the proof.
The proof for Induction Hypothesis E. is similar as before, by noticing that r(t)
k.p->p
((p) for each & E [K]. which is due to Induction Hypothesis F.1-F.4 d and Induction Hypothesis H.1.
The proof of Theorem 1.2 mirrors that of Theorem I.1,, with the only difference being the substitution of relevant sections with Induction Hypothesis E.2. For the sake of brevity, this part of the proof is not reiterated here.
1.2 PROOF OF THEOREM 4.1 AND THEOREM 4.2 WITH POSITIVE INFORMATION GAP
Theorem 1.3. Suppose D(1) < A S 1. For any e 1, suppose polylog(P). > log( t We apply GD to train the loss function given in equation 2.8 with n < poly(P). Then for each E P, we have
1. The loss converges: after T 0( (log (2 D0.10 -m)
log(Pc-))
iterations, (p(((T)) - C S where Cj is the global minimum of patch-level construction loss in equation 4.2.
"
tpe
2. Attention score concentrates: given cluster k E [K], if Xp is masked, then the one-layer transformer nearly "pays all attention" " all unmasked patches the same area Pk.au.p'
(T") i.e.. 11 11-Attin --ttt Attnp--Peane p-PPkoN.)) ) <O(c).
3. Local area learning feature attention correlation through two-phase: given E [K], if
ak.p 1, then we have
Tp not changing much; (b) after some point, the increase of (t) p takes dominance. Such & (t) Pp-rerop, will
(a) p first quickly decrease with all other D
keep growing until convergence with all other feature and positional attention correla- tions nearly unchanged.
4. Global area learning feature attention correlation through one-phase: given k E [K], if ak.p 1, throughout the training. the increase of " dominates, whereas all ( with m * 1 and position attention correlations remain close to O.
Proof. The first statement is obtained by letting T max ( T + 1 in Lemma F.30 and Lemma H.3, combining wth Lemma D.9 and Lemma 10 which lead to
(T)) c. < (T()(T))
Clow
52

---

Published as a conference paper at ICLR 2025
$ =p( 0 exp Pop + {1 Ukelk fax.p) f.Pp)))
+ 0 exp (eg + {1 (ke(k ccsp..
- 2K
<.
The second statement follows from Lemma F.30 and Lemma H.3. The third and fourth statements directly follow from the learning process described in Appendix and Appendix H when Induction Hypothesis E. holds.
1.3 PROOF OF THEOREM 4.1 AND THEOREM 4.2 WITH NEGATIVE INFORMATION GAP
Theorem 1.4. Suppose -0.5 < A < Q(1). For any 0 < < 1, suppose polylog( (P) >> log( We apply GD to train the loss function given equation 2.8 with 1) & poly(P). Then for each p E P. we have
logl 41 = (1 11,300--a))
1. The loss converges: after T
iterations,
4
( f S 6, where L is the global minimum of patch-level construction loss in equation 4.2.
2. Attention score concentrates: given cluster l [K], if Xp is masked, then the one-layer transformer nearty "pays all attention" to all unmasked patches in the same area Pk.ax.p'
(T") i.e. (1- Attn p S O(c).
3. All areas learning feature atlention correlation through one-phase: given k [K], through- (t) out the training, the increase of & p Po-ce dominates, whereas all $ pp with m + 1 and position attention correlations T p-q remuin close to O.
Proof. The first statement obtained by letting T* = max (Te 1e T + in Lemma G.3 and Lemma H.3, combining wth Lemma D.9 and Lemma 1.10, which lead to
(p()(T))) Cp < (p(Q(T)) c/o
(²)) ++0 exp (-(cyph (x + {1 & Use[K (orp))/ppp)pppp
o exp (- Peg 11 {1 UkeiK fak (ca
K 2K
<c
The second statement follows from Lemma G.3 and Lemma H.3. The third and fourth statements directly follow from the learning process described in Appendix and Appendix H when Induction Hypothesis E.2 holds.
J PROOF OF MAIN THEOREMS IN CONTRASTIVE LEARNING
Notations. Throughout this section, we abbreviate attnp-q(X;Q()) as attnp"ng(X) We Fel as F and Lei L for simplicity. We abbreviate At (X:Q())
also write
(attn (X; Q(0)) as Attn PP.... (at) when the context makes it clear. Furthermore, denote
Sime (im
Sim (ximp(*+) (x+ X x)/T
(x,
ExeØ (X+, lx 2) = ExeØ Sime (X (X/)
Theorem J.1 (Learning with contrastive objective). Suppose the information gap A [00.5 -Q(1)) ([()))(((())LL
[[1,, 11. We train the ViTs in Def. 2.6 by GD to minimize (2.11) with 1j < pP)y
of = 1. O ( after T O poly log iterations, have.
53

---

Published as a conference paper at ICLR 2025
1. Objective converges: Lai(Q(T)) I Cai poly(P) where Coa is the global minimum of the contrastive objective in (2.11).
2. Attention concentration on global area : given X E D with [K], for any p P, with high probability, we have Attnp--Pa, (X'QQTT) (1) X' E (x+ x++}.
3. All patches learn global FP correlation: given k E for any E P, P. T", pp- with m > 1, and at the convergence, p (llog P), p).2p-11 (TT)) - o().
(t)
In the following, will sketch the proof of the above theorem. Indeed, the roadmap of the analysis similar to the masked reconstruction loss by using the induction argument, where the key difference is the properties for the gradient of the contrastive objective.
J.1 PRELIMINARIES
In the following, we denote the contrastive loss without regularization as Sim (x+ (x+,x+*)/-
((Q)= A Ex+,x++00 2t T log
Σx eim Fo (X+,x'))
Lemma J.2 (feature gradient of contrastive loss). Given E for p E P, let ap ap-ve.m
~it)
(
(1++)
1 )p for m E [N&,]. then
p-
0c (Q(1)) Uk.m =a (1)(t)
0 tp =ep (-
aQ
p-
where
0L
(0)
pp (- ()(1)))km. kk
8Q'
Simp X'1//r (x1)/
E
attn (1)
p (X (x ) (F(X++:Q(0)
Σ
EXee (XI Q(X):
p
qEP
X1e
x Σ attn (0)
(X+)
Ukmm
Proof. Notice that
OC
OC
AQ.
iQ
eQ
Then for
2C σQ
we begin with the chain rule and obtain
OC
aQ
(i (x)x)) )]
E
(Simx(Xt,x't llg
Σ
EQQ
X
TF (F((x+:Q)
(F(x++:Q)
Σ
Sim ,x)//
eimm (X+, F(X'; 2))
E
aQ
Σ
Exeem
X
(attnp-sq(X*)
x+ (f(xtt:0) Σ
Exee Simp x')/ (')/= F(X'; Q1)
(X ,')/T
E
P
aQ
P.qEP
X1e8
(J.1)
We focus on the gradient for each attention score:
Pattnp-q(X*
exp (eT + +X+))
(x -
aQ
Er (xx))
54

---

Published as a conference paper at ICLR 2025
attilp-i Σ
Mettup--rep(Xp - Xit)
att (x (Xt)ep X
attn attup-or( (
Substituting the above equation into equation J.1, we have
0C
1 F F
Σ
attup-pa
(r(xt+:0)- (xxxx)((((xx))))))))))) Σ (x+xxx
Σx Exeg ( X0)// F(X:;
aQ
P.qEP
X
attnp (ettnp-r(Xt))X) ]
cp Xd
Σ
Therefore,
0C eQ
Op Pu-
(1)
cp
kk
e S (X/ X/T/T Exee Exeege (Smm F(X::Q)
1
E
Σ
attn p p- papppp-gxxxxta (x+)xx (F(x++:Q)
Σ
P
qEP
X'e
X Σ
attmp- attnp-r(X*)
Vk.m
We then present high-probability event ensuring that the number common unmasked patches in each area between positive augmented data pairs is proportional to the total number of patches in that area.
Lemma J.3 (masking overlap). Given a sample X Doi, with propbability 1- - over the randomness of masking augmentation to obtain X+, X++ supposing X belongs to the k-th cluster, holds that
Σ 10 + 0) @(Ck.m) Vm E [N&] PEP
We denote the event that the above inequalities hold A1.com- Similarly, we have the following event for X+ and X +++ hols with high probability:
AI,
Σ p
[
+ p
+ 0 (Ck,m)) Vm [Nul]
AI,+
Σ PEPA.
1xx 40} = ((CK,,))Vmm E [Nd).
Proof. The proof is similar to the analysis Lemma D.6 by using the concentration property of hypergeometric distribution.
J.2 INITIAL STAGE: GLOBAL CORRELATIONS EMERGE
For the training process at the initial stage, we define the stage transition time T₁ to be the iteration when & ( P (1 Ke) log(P) (P) for all p P and [K].
We state the following induction hypothesis, which will hold throughout this stage:
Induction Hypothesis J.1. For each < < Ti O log( P) &e [K]. letting
"
pa-ane log(P) the following holds:
2
a. (10 + is monotonically increasing, and & p E
(8)
[o, (1 Kc) log(P))
55

---

Published as a conference paper at ICLR 2025
b. $ (1)
mm S 0 (max pa(mu-ne) (0) (p)ee.)) for 1.
Lemma J.4 (bounding the noise correlation). Let us define noiseless version of the attention score and the network output as
attn xx)) (X) =
E Erep ((( for P.
(J.2)
F(x:Q)= = p Σ
attn (X) Xa
E ER.
(J.3)
P.qEP
e Simp (x+ xx++//+
SSi (x+,x-1)//-
(x, 23)
(x,))))= Exee e (X+ X)/T"
ExEE g (x+ x
Then supposing Induction Hypothesis J.1 holds for S Tr, with high probability over the randomness of x * M. then for X E B, any p, q E P, s E [Ne]. it holds that
1
attnp-og (X) attn (2ttn/2gg (X)
S
poly(d)
| (x: 1r0 (x: 01
1
poly(d)
750 1 (X, -X) (X,2))
poly(d)
We denote the event that the above inequalities hold as Az-
Proof. The result follows directly from the concentration of Gaussian random variables, the bound- edness of the feature vectors and the boundedness of [|epQ112 S k due to the Induction Hypothesis J.1
Lemma J.5 (attention score). Suppose the Induction Hypothesis J.1 holds for 1 S Tr given (X X++, assuming X E De with k E [K], then for m E [Nk]. E P, we have
1. for E (+, ++, if Xa E A2,0. then (t)
(a) 1- Attn P (Xa)= 2 211 and Attn (t)
P-PPa, p (Xa) 2 0( )))))
(b) for > 1, Attn (1)
(X²) O
1- Atttn p p (x").
):
P..
pi-ne
2. for X' E R. we have
(1)
(a) 1- - -Attn PL. (X') 2 ((1) and Attnp-PPe (X') 2 s( -((p-tm)):
1- -tAtn sP,, (X')
(b) for > 1, Attn (X') = O
pi-ss
The intuition behind this lemma is that, due to the zero initialization of Q, the attention scores are nearly uniform. As a result, the area attention score Attn (X+) is proportional to the number of unmasked patches in this area. If Induction Hypothesis J.1 holds, we can easily conclude that only the area attention score for the global area will increase, while the relative relationships among the local area attention scores will be preserved.
Lemma J.6 (logit score). Suppose the Induction Hypothesis J.1 holds for t S Tı, given {X xx++, N.. suppose E DEi we have
1) 14 (X, B) 2 22( No E B D
1 - e0 (X,
1 B) 2 2(1),
) else.
(x, < 0 No
2)
56

---

Published as a conference paper at ICLR 2025
Lemma J.7 (feature gradient near initialization). Suppose the Induction Hypothesis J.1 holds for t I To. then for t I Tı, given & E [K], m E for P E P.
For the global feature m = L.
10 = θ
1 P
21 -xt (tt)) -Pp (x++))
For the local feature mn > 1
0p mp
(1)
=
1
E
3m (1 - Attnp-pppPp
Attmp PP P (xttt)
p
1
2m (1 Att)ttppp-ra Pa.m AttM Paa Att)) Patta Paa (Pa (xctt))
+
0
P
Proof.
No attup -a (xt)xx (1 -
E
l, F
qEP
x
Σ
attnp attap-r(Xt)xx Uhmm(1. + +as)
N.
(a) 1
attn p (x xx+ T (11 Σ (
E
P
qEp
x+
attnp-r(Xt)x
(1. + + + +pp.mmm.1
N2
No
(Xt (2sen.s) (a- 2
E
p
ZiUkit
attn attnp- (X (x+) xt)xzu0n)))
Uk.m
(J1)
Σ
attnp (xx)xxxxxx) (ziV).)) ((1 (1 Σ T.F(( ; =Q)
Eq
Σ rePrit
attm (x (xt)er
Uk
(J2)
No attn p b (Xt)eT ((x -
+
R
P
qePnil
N
attn attn (x)))
vx.m
(Jz)
j=1 NM!
No
1
E
Σ
attn q (X+)ET (1 -))
p
t
Eq
Σ
attn (X+)sr
Ukmm
(J.)
rePrut
57

---

Published as a conference paper at ICLR 2025
where (a) is bounded by Lemma J.4 with error up Ep.k.m.1 S poly(d) U is the set of masked patches for X We first look at the term J1, notice that &q is the random Gaussian noise with zero mean, and is independent of attn and L, we then have
E J4 p2
attnp o (X )eg (11-
Σ
attn -r 2mmkkmm Uk.
mett
D'EPIEPE.
Eq
attnp attnp--((X*)) 0xr
Ukm
rEPrUt
E
attup- 1x+)5 (11 (1-
Σ
E
( (c.)
p2
qEPx.m n
p'EP rEM
Eq
Σ
attmp--r (x+))
Uk,m
rePrU
2m Σ
attnp p 1ampp-alXtt (1 2p)
E Σ
attmp' Coump--((xt)) (x attnp--g(X
p2
qEPa,mnl m
D'EP rE PAiM!!!!!!
Attnp (x+)
attn ,q (X (x-)
((1 ([
E (pmpp-))__((tt)))) + Epp
Zm
p2
qEPk. -
P'EP
1 2m(11 Attmpp Pe.
Pe.m (x AttuppppPe c,m (xtt))
=
where (a) is bounded by invoking Lemma J.4 with error up to Ep. I poly(d)* and the last equality is due to Lemma J.5.
N
J/ 1
E
attmp-a(A)) (x+ (2(Un.)))))))
PE
i+++
[gq
Σ
attn - (X |5] ck m[]
rEPnLt
N
Σ
attn attnp-gg (X (x- (2(UR.)) (1 (1-
Σ
attnp'-r(> (X++)
1=1 qEPt. nut
P'EPrrEPrutt
Σ
attn |Er] vk
rEPnIst
E
attn putnn--gg s x (2m T(1 (np
Σ
p2
QEPR.m N4
P'EP-rEPnIAt
Sq
attn p dr (X (x+)5t] [kmlk]
rEPrIl
1
E
2m
attn attnp-g(X* (X+ ](1 Ep)
p2
qE Pt
attnip-q(X (X+)1qqu11
Σ
attn (
P'EP.rEPnuttrut
1
E
7m
attn (xtnp-g(x+)(11-)))
p2
qEPt
(attnp--q (Xt 1quu++
Σ
attn attnp--r(X (X (xt))attt jattup-r(XX) Epp.km.
D'ERrEPnUttru#
58

---

Published as a conference paper at ICLR 2025
Thus, by invoking Lemma J.5, we have
|J2| S
1 E
attn (x+ (p) (rEPIMA
max attup-r a(tmp-n(xt))) + P.
mm
qe DEPK.m nut
< o
1
2m Attnp Phim (X7) Attmp (ttmp-)Paa((tt))
E
P. CK,1
J/ =
Σ
attnp
(X+)) ((x
qEPnU+
T attmp-r )2y(trs)) ukm
Σ
E
E
attnp-ag(X*) (1)
p
attup-o [[s]
Zm
rEPu, 1
Σ
attnp->g(X) (1 - Epp
2m
p2
qe Pk nd nextt
attm a (X++)
Σ
attnp-
PLEP
rEP,, nut met
Σ
attup-so soppp--gX" (X)((- Ep)
Zm
p2
qEpr nu++
attn (x(xtt)(( Attn [x+)] (X+]] E...im.
P'EP
Attup Phim (x (x+)(1- Attnp (xtm)))))((())))) (1-))
Σ 0 ((
Attup Attnpp-PPan (x++))
<0
E
p2
Ck
PEEP
Ja Ck.m
<
<0(
where the last inequality is due to Lemma J.5.
No
J1 = p2
Σ
attnp p - (X
(2.U).)) T
1=1
N₂ P'EP j=1 rut
Σ
attn,' -g (xxtit)
Σ 2.
Σ
attup' - (X - -2) Uk.))
"EMNDA
P'EP j=l
Σ
p 2 2k
J=1 ,E
No
E
Σ
attn (X* (2(e).))) T
p2
1=1 nut
59

---

Published as a conference paper at ICLR 2025
N2
(11 (1- Σ Σ
E
attnp' - (X (x+++))2+ 2g
P'EP j=l j=1 q' SPEPPtt
E 4
Σ
attup .g (X -0) )a Ukes)
P'EP j
No
2.U Σ
attnp-r2y1 ukmm + Ep.kmm.
J=1 NEE
E Attnp +P
(1- 2m
Attup- Attmp-Pe.((Xt) (Xx)) uk
Σ
2m Attt (X (X+))os))
Nx
(11 (11 Σ Σ
Σ
attnp' (x+it) (xttt)zzkkks 2zjkkks
D'EP j=1 q'
E
Σ
Σ
attnp attnp'+g (X ())s. Mkess
Ep.km..
X E
P'EP j'Eq'E
T
E
Attnp Pe (x+)
22 (1
P Ptt (xt)) uk.m
Σ
Zm Antnntttpppa +Pa, (X )okes
p2
iymm
N2
N₂
(1-
Σ Σ
Attnp
-Po (X
Σ
Σ
Attmp'- (X (x=+) ()
j=1
X x
P'EP j=1
mmk,mm,
E Attup 4 (X+)
(1-
Attup- (x+))
am
Σ 2mm) (xx+))
Σ
Zs Attnnn'p-P (X (X-+)))
P'EP
x
(J1,1)
(x+)
Σ
3m Attn-:tttppppppp (X+)
p2
Σ (1- (4=6))attapiippp. (X (X=)-
Σ
Za Lattl.Attpp Pr. (X -1))
P'EP
"ERNND
(J1.2
m.k.m.5
Notice that J1,1 (J4.). Furthermore, when m 1, J1.2 is negligible compared to J1, else
|J12| S
0
E
Attn PPe (x+)
2m Attnp-pp (x+)
p2
Σ
((1 Attnp'--p (X++)
Σ
2x Attmp'-PP. (X
p'ep
so 1
2m (pp)) Attnp (X+ Attn (x+ Attn (X++t)
Putting all the terms together, and noticed that
E 0 P
(1 -t))Attt Pr.m AttD (X+) Attnp- . -x., (X*+t)
60

---

Published as a conference paper at ICLR 2025
0
2m 2m - Amt Attn (x.
Ck.1
then we complete the proof.
Proof of Induction Hypothesis J.1. By Lemma J.7 and Lemma J.6, at the initial stage of the learning process, we have
(0)
1 Attnp PA. (x. Att
E
PPa (xt+)
U1
P
(0)
S
1
max B Attmp (x ( Attnp Paim (Pmm)
P
Attup-PP.... (X (Xt)Attp--PP.... (X AAtn) PA. (
E
Then by the relations of attention score in Lemma J.5, focusing on the high-propbability event AI, and A A we have
(4) -oe ( max{ pk.-1 papp.p na) $10) p
fa for m 1 (J.4)
By Lemma J.7 and Lemma J.5, we have P 2 22 ((ppt)))))
(t)
Ag (t)
P-.Pa. which implies the
regularization in this stage is not violated for the dominated FP correlation ob (1)
P Hence,
(t) could focus on the relation between (t) P and a P- for m 1.
Therefore, the existence Ti can be directly obtained by the gradient estimation Lemma J.7 and
the lower bound for the area attention of the global area in Lemma J.5. The induction argument follows directly from J.4.
The key takeaway from the first stage is that the growth of feature posiition attention correlation for the global area is dominant, specifically, Op- > 1 After this initial stage, T P- reaches ((log(() (P)), Attnp-pP., has reached (2(1) and 1 - l still keeps at constant level. The dominance of global FP correlation will be preserved in the following and the learning process will enter the convergence stage.
J.3 CONVERGENCE
At this stage, we are going to prove that as long as the VITs have already learned the global correlations, they will indeed converge to these global solutions, which leads to the collapsed global representation. We present the statement of our convergence theorem below.
Theorem J.8 (Convergence guarantees). Letting T2
2 pa log P for any T
7
[T2, O(( (oly(P) log log letting x O plog)) we have
"
Σ ((Q()) C. poly P
T=
where Ct is the global minimum of the regularized contrastive objective.
We have the following hypothesis for the end of the learning process.
Induction Hypothesis J.2. For & E [ pa log P. O(( PPly(P) log P. })}, we have the following resutls:
For any & E [K], p E P, and m E [N&]
E [C C2] log P
(t)
m
p5.
where CH, C 0 some constants and 6, (0, 1) is small constant.
61

---

Published as a conference paper at ICLR 2025
Attention score from the global area: given X Dx, 1 Attn (X9) S poly(P] for
a E ++, and 1 Attm (X 2 poly(P) for & E [No] with high probability.
Bounded gradient for the loss:
(oL(Q S
poly p).
We can reuse most of the calculations in the proof of Induction Hypothesis J.1 to prove the hypothesis and here we only discuss how to bound the gradient of the objective. If the regularization is not
4.60
violated, i.e., Op
(1)
pp-ck.m
we have T (L)
p S O(log(P)). For T2, denote the first
((t)
(1) time when ap-ve,
1 ((() as Ti, by Lemma J.7, we have ap ap-vvk. 2 for
I E [Ti, Ji and 1 p- = C Colog// P for some constant C > 0. Then we have Ti S 0 p log Thus, for 2 Ti.
(T1)
(V So
)2 S S poly(P)
p P- 18.1
DEP
Proof of convergence. We first define a learning network that we deem as the "optimal" network with the global feature-position attention pattern. Specifically, define Q* matrix satisfied ep 0' = o. with o² PEE No) and QK Uk.m = for p E P and & E [K], m E [N&]-
[
Furthermore, w @² = O, where w1, W2 E Span (fep)pEP n (uk. + Here we suppose C achieved at the matrix Q = Q.
Moreover, We consider the following pseudo losses and objective: define the linearized learner F(0 x) ((Q(1) X) + VQF(Q(), XX( Q(9),
Ž((Q) E
-T log
Exees Ex (
Obj,(Q) = ž((2) + 11
and
(
((Q)) = -T llg
Then we discuss the values of different losses at Q = Q"
We have the following properties:
1
f(Q') 5 Ca 0
(J.5)
poly
(Z(0*) - I.(Q7) < 0
(J.6)
poly(d)
[Z((Q)) -2(0)) poly d'
1
(J.7)
For the first property, we only need to consider the contrastive loss at the global minimum. Notice that for our data distribution, the global minimum of the contrastive loss is achieved when the network can perfectly distinguish the samples from different clusters. Thus, we have L O Ne ). Notice that on the event Al.oom supposing X DE. which happens with prob > 1 -- pss we have
F(O;)(((xxxt))) = (Uk.1. (b.1 +
Σ
1 + o(1)
[((C,.))
DEPtinutnutt
62

---

Published as a conference paper at ICLR 2025
(F(Q", .x"), (Uk,1 vx,1) + o(1) for X' E Di
Furthermore, by Bernstein's inequality, we have with probability 2 1 poly(d) we have 11cp112
ofd + Ō polyta)
Ō( poly(aj ). we denote such an event as Ag. Suppose we consider the temperature 7 = O log ). then conditioned on A1.com A,3s, we have (F(Q", x+), ;(Q(") X');
k ) up to wlog d) + (11) for X" 2 n Di', which could minimize the loss the level of (log Ng
the error of ( poly(d) Then we have
N K
((Q') s (1 poly(d)
O
poly(d) ) < ((Q) + 0
poly(d)
The second property follows from the observation that
[2,00)) L(Q") s = 0(IV2 - (t)
<0((vaZ(\")lz (1 Attt P-rF, (x++)) s =0( Poly(P)) <O(poly(P))
Similarly, the third property follows from the fact that
(0)) 00
F(Q", s Ō poly(P)
Now we will use the tools from online learning to obtain a loss guarantee:
(VVL(Q())) Q(1) Q")
" (VL(
)| 1 -0 11Qut Q
n
1 110 Qr + - 12.
Poly
2
Notice that Obj.(Q) (Q) is a convex function over Q and (Q = L(Q(1),, thus
(Vo[(Q)), Q Q*) (V))) Q Q")
Obj, (Q - (Q*)
(by convexity)
(b) - ((Q") 5 pPly(P))
(by J.6 and J.7)
(Q) - f. -Ö( poly(P))
(by J.5)
L(Q()) - C. - -ō( 'oly(P)) _((oo)(P))
(by definition of Obj)
Thus by telescoping summation, we have
((Q(t)) cn1
T T2 T=
( -()))))())))))))))))))))))))))))
1 < 0
-0) 0 poly(P))
t=
(Q())) - Q
1
00
Tn
poly(P))
which completes the proof.
63


Reviews:
[{"review_id": "JiHsPwUeet", "reviewer": "Reviewer_i2aQ", "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.", "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. \n\n- Orthonormality assumption: line 193. Why do you make that assumption? To maximize span?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "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?\n- 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?\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3wvHdagUKQ", "reviewer": "Reviewer_3zvq", "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.", "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. \n- Is the proposed theory would still be valid for datasets where the global area is small? \n- 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.  \nMinors: \n- what is $k_s$? It seems the sum of weights for global area.\n- What $\\Delta$ means? Can you explain in details, such as when $\\Delta=-1,0,1$?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n2. 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.\n3. 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. \n2. There are a lot of assumptions, such as the input data distribution and positional encoding. It is not clear the differences in real problems.\n3. What is the insight from this work to develop new methods? How to prevent the attention collapse in CL?\n4. The discussion is based on SGD. However, Adam is the popular optimizer for ViTs. \n5. 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. \n\n[1] Atito etal., \"SiT: Self-supervised vIsion Transformer\", arxiv 04, 2021. \n[2] Zhou etal., \"iBOT: Image BERT Pre-Training with Online Tokenizer\", arxiv 11, 2021. \n[3] Oquab etal., \"DINOv2: Learning Robust Visual Features without Supervision\", arxiv 04, 2023. \n[4] He etal., \"Masked Autoencoders Are Scalable Vision Learners\", arxiv 11, 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vp8uTi3u65", "reviewer": "Reviewer_WFmo", "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.", "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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n2. 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. \n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aRMBZ6zDbY", "reviewer": "Reviewer_ChhH", "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.", "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]$?\n2. 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.\n3. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)