A Theoretical Analysis of Self-Supervised Learning for Vision Transformers

## Abstract

Abstract Masked reconstruction, which predicts randomly masked patches from unmasked ones, has emerged as an important approach in self-supervised pretraining. However, the theoretical understanding of masked pretraining is rather limited, especially for the foundational architecture of transformers. In this paper, to the best of
our knowledge, we provide the first end-to-end theoretical guarantee of learning one-layer transformers in masked reconstruction self-supervised pretraining. On the conceptual side, we posit a mechanism of how transformers trained with masked vision pretraining objectives produce empirically observed local and diverse attention patterns, on data distributions with spatial structures that highlight feature-position correlations . On the technical side, our end-to-end characterization of training dynamics in softmax-attention models simultaneously accounts for input and position embeddings, which is developed based on a careful analysis tracking the interplay between feature-wise and position-wise attention correlations.

## 1 Introduction

Self-supervised learning has been a leading approach to pretrain neural networks for downstream applications since the introduction of BERT [] and GPT [] in natural language processing (NLP). On the other hand, in vision, self-supervised learning was initially focused more on discriminative methods, which include contrastive learning [, ] and non-contrastive learning methods [, , , ]. Inspired by masked language models in NLP, together with the crucial progress made by [] in successfully implementing vision transformers (ViTs), generative approaches,
such as masked reconstruction-based methods, have become popular in self-supervised vision pretraining, especially since the emergence of masked autoencoders (MAE) [] and SimMIM [].

In masked vision pretraining, neural networks are instructed to reconstruct some or all patches of an image given a masked version, aiming to learn certain abstract semantics of visual contents when trained to fill in the missing patches.
In practice, this approach not only proves to be very successful with remarkable finetuning performance but also shows intriguing phenomena that differ significantly from the behaviors observed in other self-supervised learning methods. The seminal work of [] showed that MAE can conduct visual reasoning when filling in masked patches even with very high masking rates, suggesting that masked image modeling learns some complex relationships between visual objects. Some critical observations from recent research [, , ] have suggested that the ViTs trained via masked reconstruction objectives display diverse attention patterns: different query patches pay attention to distinct local areas. This is in sharp contrast to the discriminative self-supervised learning approaches, whose attention maps mostly capture the most significant global pattern regardless of where the query patches are, leading to a phenomenon known as “attention collapse”, as shown in Figure 1.
Given these empirical observations, it naturally prompts the question: from a theoretical
standpoint, how do ViTs learn these varied attention patterns during masked pretraining?

Despite the substantial empirical effort dedicated to investigating masked vision pretraining, our theoretical understanding of it is still nascent. Most existing theories for self-supervised pretraining focused on discriminative methods [, , , , , , , ], such as contrastive learning. Among very few attempts towards masked image modeling, [] studied the patch-based attention via an integral kernel perspective; [] analyzed MAE through an augmentation graph framework, which connects MAE to contrastive learning. [] characterized the optimization process of MAE with shallow convolutional neural networks (CNNs). However, transformers, as the prevailing architecture in current deep learning practice, were not addressed in the aforementioned theoretical studies of
vision pretraining and, more broadly, self-supervised pretraining, leaving a significant gap in the literature.

Figure: Figure 1: Visualization of attention maps in the last layer of ViT for query patches from two different spatial locations, trained by the generative self-supervised learning approach of masked reconstruction (MAE) and two discriminative self-supervised learning methods (DINO [] and MoCo []).
Refer to caption: x1.png

Building on the mind-blowing
empirical advances and recognizing the lack of theoretical understanding of masked reconstruction-based methods with transformers in self-supervised learning,
we are motivated to ask:
\mdfsetup
frametitle= Our Research Questions ,
shadow=true,shadowsize=1pt,
skipabove=2pt,
innertopmargin=-3pt,
innerbottommargin=7pt,
innerrightmargin=7pt,
innerleftmargin=7pt,
frametitleaboveskip=-frametitlealignment=,
linewidth=0.5pt
{mdframed}
Can we theoretically characterize
what solutions transformers converge to in masked vision pretraining? How do transformers learn diverse local attention patterns rather than object-focused global attention in such pretraining?

##### Contributions.

In this paper, we take a step towards answering the above questions by studying the training process of one-layer transformers trained by gradient descent (GD) for masked reconstruction objectives. We focus on spatially structured data distributions, where each image is sampled from different clusters characterized by different patch-wise associations.
Under such setting:

- 1.
We provide the global convergence guarantee of the masked reconstruction loss and
characterize how attention is distributed at convergence to demonstrate the non-collapsed diverse local attention patterns obtained by masked pretraining
, which to our knowledge is the first end-to-end theory of learning transformers for masked-reconstruction type self-supervised methods.
- 2.
We analyze the training dynamics of attention correlations (see Definition 3.1) and prove that ViTs manage to capture desirable diverse local patterns
in masked pretraining by learning the target feature-position attention correlations for all visual features irrespective of their significance i.e., whether they are global or local features. This marks the first result of learning the softmax attention model that jointly considers input and position encodings.
- 3.
We design a novel empirical metric, termed as the attention diversity metric, to probe vision transformers trained by different methods. Our new observations (see Figure 2) provide further evidence for the diverse local patterns revealed through masked image modeling.

Figure: Figure 2: Attention Diversity Metric: We examined the last layer of ViTs trained by masked reconstruction (MAE), discriminative self-suprvised learning (MoCo, 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 on different, localized areas. The results show that the MAE model excels in capturing diverse local patterns (see Section 6 for more details).
Refer to caption: extracted/5633500/output.png

##### Comparisons with prior works.

A few works [, ] have studied topics related to ours. Here we summarize the differences between our work and theirs regarding settings and analysis.

- •
[] pioneered the characterization of
the training dynamics for transformers in supervised learning. They provided the first convergence result for one-layer softmax-attention transformers trained on a simple visual data distribution, in which the patch association
is fixed (see Definition 2.1 in their paper). Their assumptions require learning only the position-position correlations, which is rather limited. In Section 3, after formally defining the feature-position correlations, we will highlight the inadequacy of such a setting to fully demonstrate masked pretraining’s capability to learn locality in detail. In this work, we draw inspiration from their data assumptions and generalize them to allow varying spatial structures of patches
(see Definition 2.1). Due to this generalization, we need to analyze the learning processes of different spatial correlations among visual features simultaneously, which poses key challenges in our overall analysis.
- •
[] analyzed the feature learning process of MAE with CNNs architectures, without any patch-level positional structure in the network.
The main implication of their theoretical result is that the trained CNN can provably identify all global patterns in masked pretraining, which does not capture the findings in recent self-attention investigations [, ] and Figure 1 that masked pretraining with transformers can learn diverse local patterns and avoid collapsing solutions. This suggests that the masked reconstruction-based method’s success under transformers may fall beyond what CNN-based analysis could reveal, emphasizing the importance of studying attention patterns from a theoretical point of view. This is the focus of this paper.

##### Notation.

We introduce notations to be used throughout the paper. For any two functions $h(x)$ and $g(x)$, we use $h(x)=\Omega(g(x))$ $\big{(}$resp. $h(x)=O(g(x))\big{)}$ to denote that there exist some universal constants $C_{1}>0$ and $a_{1}$, s.t. $|h(x)|\geq C_{1}|g(x)|$ $\big{(}$resp. $|h(x)|\leq C_{1}|g(x)|\big{)}$ for all $x\geq a_{1}$; Furthermore, $h(x)=\Theta(g(x))$ indicates $h(x)=\Omega(g(x))$ and $h(x)=O(g(x))$ hold simultaneously. We use $\operatorname{\mathds{1}}\{\cdot\}$ to denote the indicator function, and let $[N]\coloneqq\{1,2,\ldots,N\}$.
We use $\widetilde{O}$, $\widetilde{\Omega}$, and $\widetilde{\Theta}$ to further hide logarithmic factors in the respective notations.
We use $\operatorname{poly}(P)$ and $\operatorname{polylog}(P)$ to represent large constant-degree polynomials of $P$ and $\log(P)$, respectively.

## 2 Problem Setup

In this section, we present our problem formulations for studying the training process of transformers during masked vision pretraining. We begin with some background information, followed by a description of our dataset settings. We then detail the masked pretraining strategy and the specific transformer architecture considered in this paper.

### 2.1 Background

We follow the masked reconstruction frameworks in [, ]. Each data sample $X\in\mathbb{R}^{d\times P}$ has the form $X=(X_{\mathbf{p}})_{\mathbf{p}\in\mathcal{P}}$, which has $|\mathcal{P}|=P$ patches, and each patch $X_{\mathbf{p}}\in\mathbb{R}^{d}$. Given a collection of images $\{X_{i}\}_{i\in[n]}$, we select a masking set $\mathcal{M}_{i}\subset[P]$ for each image $X_{i}$, and mask these patches to a uniform value $\mathsf{M}\in\mathbb{R}^{d}$. The resulting masked images $\{\mathsf{M}(X_{i})\}_{i\in[n]}$ are given by

$$ $\displaystyle\mathsf{M}(X_{i})_{\mathbf{p}}=\left\{\begin{array}[]{cc}[X_{i}]_ {\mathbf{p}}&\mathbf{p}\in\mathcal{U}_{i}\\ \mathsf{M}&\mathbf{p}\in\mathcal{M}_{i}\end{array}\right.,\qquad i\in[n],$ (2.3) $$

where $\mathcal{U}_{i}=\mathcal{P}\setminus\mathcal{M}_{i}$ is the index set of unmasked patches.
Let $F:X\mapsto\widehat{X}$ be an architecture that outputs a reconstructed image $\widehat{X}\in\mathbb{R}^{d\times P}$ for any given input $X\in\mathbb{R}^{d\times P}$.
The pretraining objective is then defined as the mean-squared reconstruction loss over a series of subsets $\mathcal{P}^{\prime}_{i}\subset\mathcal{P}$ of the image as follows:

$$ $\widehat{\mathcal{L}}(F)=\frac{1}{n}\sum_{i=1}^{n}\sum_{\mathbf{p}\in\mathcal{ P}^{\prime}_{i}}\Big{\|}[X_{i}]_{\mathbf{p}}-[F(\mathsf{M}(X_{i}))]_{\mathbf{p }}\Big{\|}_{2}^{2}.$ (2.4) $$

[] chose the subset $\mathcal{P}_{i}^{\prime}$ as the set of masked patches $\mathcal{M}_{i}$, whereas [] aimed to reconstruct the full image $\mathcal{P}_{i}^{\prime}=\mathcal{P}$. We do not explore the trade-offs between these two approaches in our study. As shall be seen momentarily, our theoretical framework is based on $F$ being a simplified version of vision transformers [] which utilizes the attention mechanism [].

### 2.2 Data Distribution

We assume the data samples $X\in\mathbb{R}^{d\times P}$ are drawn independently based on some data distribution $\mathcal{D}$. To capture the *feature-position (FP) correlation* in the learning problem, we consider the following setup for vision data. We assume that the data distribution consists of many different clusters, where each cluster captures a distinct spatial pattern, and hence is defined by a different partition of patches with a different set of visual features. We define the data distribution $\mathcal{D}$ formally as follows. An intuitive illustration of data generation is given in Figure 3.

Figure: Figure 3: Illustration of our data distribution (see Definition 2.1). Each cluster $\mathcal{D}_{k}$ is segmented into distinct areas $\mathcal{P}_{k,j}$ , with squares in the same color representing the same area $\mathcal{P}_{k,j}$. The global area $\mathcal{P}_{k,1}$ (depicted in orange) contains a 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 is also applicable to non-contiguous cases.
Refer to caption: x2.png

###### Definition 2.1 (Data distribution $\mathcal{D}$ ) .

The data distribution $\mathcal{D}$ has $K=O(1)$ different clusters $\{\mathcal{D}_{k}\}_{k=1}^{K}$. For every cluster $\mathcal{D}_{k},k\in[K]$, there is a corresponding partition of $\mathcal{P}$ into $N_{k}$ disjoint subsets $\mathcal{P}=\bigcup^{N_{k}}_{j=1}\mathcal{P}_{k,j}$ which we call areas. For each sample $X=(X_{\mathbf{p}})_{\mathbf{p}\in\mathcal{P}}$, its sampling process is as follows:

- •
We draw $\mathcal{D}_{k}$ uniformly at random from all clusters and draw a sample $X$ from $\mathcal{D}_{k}$.
- •
Given $k\in[K]$, for any $j\in[N_{k}]$, all patches $X_{\mathbf{p}}$ in the area $\mathcal{P}_{k,j}$ are given the same content $X_{\mathbf{p}}=v_{k,j}z_{j}(X)$, where $v_{k,j}\in\mathbb{R}^{d}$ is the visual feature and $z_{j}(X)$ is the latent variable. We assume $\bigcup_{k=1}^{K}\bigcup^{N_{k}}_{j=1}\{v_{k,j}\}$ are orthogonal to each other with unit norm.
- •
Given $k\in[K]$, for any $j\in[N_{k}]$, $z_{j}(X)\in[L,U]$, where $0\leq L<U$ are on the order of $\Theta(1)$.(^3^33The distribution of $z_{j}(X)$ can be arbitrary within the above support set.)

###### Definition 2.1 (Data distribution 𝒟 𝒟 \mathcal{D} caligraphic_D ) .

##### Global and local features in an image.

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 as the texture of leaves in the background. Recent empirical studies on self-supervised pretraining with transformers [, ] and observations in Figure 1 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 a 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 $\mathcal{D}_{k}$ with $k\in[K]$ be a cluster from $\mathcal{D}$. We let $\mathcal{P}_{k,1}$ be the global area of cluster $\mathcal{D}_{k}$, and all the other areas $\mathcal{P}_{k,j},j\in[N_{k}]\setminus\{1\}$ be the local areas. Since each area corresponds to an assigned feature, we also call them the global and local features, respectively. Moreover, we assume:

- •
Global area: given $k\in[K]$,
we set $C_{k,1}=|\mathcal{P}_{k,1}|=\Theta(P^{\kappa_{c}})$ with $\kappa_{c}\in[0.5005,1]$, where $C_{k,1}$ is the number of patches in the global area $\mathcal{P}_{k,1}$.
- •
Local area: given $k\in[K]$, we choose $C_{k,j}=\Theta(P^{\kappa_{s}})$ with $\kappa_{s}\in[0.001,0.5]$ for $j>1$, where $C_{k,j}$ denotes the number of patches in the local area $\mathcal{P}_{k,j}$.

The rationale for defining the global feature in this manner stems from observing that patches representing global features ($C_{k,1}$) typically occur more frequently than those representing local features ($C_{k,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.

###### Assumption 2.2 (Global feature vs local feature) .

### 2.3 Masked Vision Pretraining with Transformers

##### Transformer architecture.

A transformer block [, ] consists of a 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^{Q}$, a key matrix $W^{K}$, and a value matrix $W^{V}$. Given an input $X$,
the output of one head in the self-attention layer can be described by the following mapping:

$$ $\displaystyle G(X;W^{Q},W^{K},W^{V})=\operatorname{softmax}\left({(W^{Q}X)^{ \top}W^{K}X}\right)\cdot(W^{V}X)^{\top},$ (2.5) $$

where the $\operatorname{softmax}(\cdot)$ function is applied row-wisely and for a vector input $z\in\mathbb{R}^{P}$, the $i$-th entry of $\operatorname{softmax}(z)$ is given by $\textstyle\frac{\exp(z_{i})}{\sum_{s=1}^{P}\exp(z_{s})}$.

To simplify the theoretical analysis, we consolidate the product of query and key matrices $(W^{Q})^{\top}W^{K}$ into one weight matrix denoted as $Q$. Furthermore, we set $W^{V}$ to be the identity matrix and fixed during the training. These simplifications are often taken in recent theoretical works [, , ] in order to allow tractable analysis. With these simplifications in place,
(2.5) can be rewritten as

$$ $\displaystyle G(X;Q)=\operatorname{softmax}\left({X^{\top}QX}\right)\cdot X^{ \top}.$ (2.6) $$

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 [, ]. Our assumptions regarding the positional encodings are as follows:

###### Assumption 2.3 (Positional encoding) .

We assume fixed positional encodings, which is consistent with the implementation in MAE []: $E=(e_{{\mathbf{p}}})_{{\mathbf{p}}\in\mathcal{P}}\in\mathbb{R}^{d\times P}$ where positional embedding vectors $e_{{\mathbf{p}}}$ are orthogonal to each other and to all the features $v_{k,j}$, and are of unit-norm.

We now include positional embeddings in (2.6) and introduce the network architecture for masked reconstruction used in this study.

###### Definition 2.4 (Transformer network for masked pretraining) .

We assume that our vision transformer $F(X;Q)$ consists of a single-head self-attention layer with an attention weight matrix $Q\in\mathbb{R}^{d\times d}$. For an input image $X\sim\mathcal{D}$, we add positional encoding by letting $\widetilde{X}=X+E$. The attention score from patch $X_{\mathbf{p}}$ to patch $X_{{\mathbf{q}}}$ is denoted by

$$ $\displaystyle\textstyle\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}(X;Q) \coloneqq\frac{e^{\widetilde{X}_{{\mathbf{p}}}^{\top}Q\widetilde{X}_{{\mathbf{ q}}}}}{\sum_{{\mathbf{r}}\in\mathcal{P}}e^{\widetilde{X}_{{\mathbf{p}}}^{\top} Q\widetilde{X}_{{\mathbf{r}}}}},\quad\textrm{for }\mathbf{p},\mathbf{q}\in \mathcal{P}.$ (2.7) $$

Then the output of the transformer is given by

$$ $\displaystyle[F(X;Q)]_{{\mathbf{p}}}=\sum_{{\mathbf{q}}\in\mathcal{P}}\mathbf{ attn}_{{\mathbf{p}}\to{\mathbf{q}}}(X;Q)\cdot X_{{\mathbf{q}}},\quad\textrm{ for }\mathbf{p}\in\mathcal{P}.$ (2.8) $$

Then we formally define the masking operation and the objective for our masked pretraining task.

###### Definition 2.5 (Random masking) .

Let $\mathsf{M}(X)\to\mathbb{R}^{d\times P}$ denote the random masking operation, which randomly selects (without replacement) a subset of patches $\mathcal{M}$ in $X$ with a masking ratio $\gamma=\Theta(1)\in(0,1)$ and masks them to be $\mathsf{M}:=\mathbf{0}\in\mathbb{R}^{d}$. The masked samples obey (2.3).

###### Assumption 2.3 (Positional encoding) .

###### Definition 2.4 (Transformer network for masked pretraining) .

###### Definition 2.5 (Random masking) .

##### Masked reconstruction objective.

To train the transformer model $F(\mathsf{M}(X);Q)$, we minimize the squared reconstruction error (2.4) only on masked patches, where $\mathsf{M}(X)$ follows Definition 2.5. The training objective thus can be written as

$$ $\displaystyle\mathcal{L}(Q)$ $\displaystyle\coloneqq\frac{1}{2}{\mathbb{E}}\left[\sum{}_{{\mathbf{p}}\in \mathcal{P}}\operatorname{\mathds{1}}\{{\mathbf{p}}\in\mathcal{M}\}\Big{\|}[F( \mathsf{M}(X);Q)]_{{\mathbf{p}}}-X_{{\mathbf{p}}}\Big{\|}^{2}\right],$ (2.9) $$

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 (2.9) is minimized via GD with learning rate $\eta>0$. At $t=0$, we initialize $Q^{(0)}:=\mathbf{0}_{d\times d}$ as the zero matrix. The parameter is updated as follows:

$$ $\displaystyle Q^{(t+1)}=Q^{(t)}-\eta\nabla_{Q}\mathcal{L}(Q^{(t)}).$ $$

In the following, any variable with a superscript ^(t) represents that variable at the $t$-th step of training. Note that the initialization of $Q^{(0)}$ results in any query patch uniformly attending to all patches.

## 3 Attention Patterns and Feature-Position Correlations

To show the significance of the data distribution design and understand the nature of the masked reconstruction task, in this section, we will provide some preliminary implications of the spatial structures in Definition 2.1. Intuitively, for a given cluster $\mathcal{D}_{k}$, to reconstruct a missing patch $\mathbf{p}\in\mathcal{P}_{k,j}\cap\mathcal{M}$, the attention head should exploit all unmasked patches in the target area $\mathcal{P}_{k,j}$ to find the same visual feature $v_{k,j}$ to fill in the blank, which emphasizes the locality for $\mathbf{p}$ in different areas. This approach remains effective even at high masking rates, provided that a small number of patches in $\mathcal{P}_{k,j}$ remain unmasked. We will elaborate on this point by describing the *area attentions* and illustrating the intuition about how they can be learned via attention correlations (Definition 3.1).

##### Area attention.

We first define a new notation for a cleaner presentation. For $X\sim\mathcal{D}$ and $\mathbf{p}\in\mathcal{P}$, we write the attention of patch $X_{\mathbf{p}}$ to a subset $\mathcal{A}\subset\mathcal{P}$ of patches by

$$ $\displaystyle\textstyle\widetilde{\mathbf{Attn}}_{{\mathbf{p}}\to\mathcal{A}}( X;Q)\coloneqq\sum_{\mathbf{q}\in\mathcal{A}}\mathbf{attn}_{\mathbf{p}\to \mathbf{q}}(X;Q).$ $$

Let us 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 $\mathcal{D}_{k}$, and the patch $X_{\mathbf{p}}$ with $\mathbf{p}\in\mathcal{P}_{k,j}$ is masked. Then the prediction of $X_{\mathbf{p}}$ given masked input $\mathsf{M}(X)$ can be written as

$$ $\displaystyle[F(\mathsf{M}(X);Q)]_{{\mathbf{p}}}$ $\displaystyle\textstyle=\sum_{{\mathbf{q}}\in\mathcal{P}}\mathsf{M}(X)_{{ \mathbf{q}}}\cdot\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}(\mathsf{M}(X);Q)$ $\displaystyle\textstyle=\sum_{i\in[N_{k}]}z_{i}(X)v_{k,i}\cdot\widetilde{ \mathbf{Attn}}_{{\mathbf{p}}\to\mathcal{U}\cap\mathcal{P}_{k,i}}(\mathsf{M}(X) ;Q)\text{ (since $\mathsf{M}(X)_{\mathbf{q}}=\mathbf{0}$ if $\mathbf{q}\in \mathcal{M}$)}.$ $$

To reconstruct the original patch $X_{\mathbf{p}}$, the transformer should not only focus on the correct area $\mathcal{P}_{k,j}$, but must also prioritize attention to the *unmasked* patches within this area. This specificity is denoted by the area attention $\widetilde{\mathbf{Attn}}_{{\mathbf{p}}\to\mathcal{U}\cap\mathcal{P}_{k,j}}$ over $\mathcal{U}\cap\mathcal{P}_{k,j}$, a requirement imposed by masking operations. To further explain how transformers perform such prioritization,
we introduce the following quantities, which capture the major insights of our analysis that differentiate from those in [].

###### Definition 3.1 .

(Attention correlations)
Let ${\mathbf{p}}\in\mathcal{P}$, and
we define attention correlations as:

- 1.
Feature-Position (FP) Correlation: $\Phi_{{\mathbf{p}}\to v_{k,m}}\coloneqq e^{\top}_{{\mathbf{p}}}Qv_{k,m}$, for $k\in[K]$ and $m\in[N_{k}]$;
- 2.
Position-Position (PP) Correlation: $\Upsilon_{{\mathbf{p}}\to{\mathbf{q}}}\coloneqq e_{{\mathbf{p}}}^{\top}Qe_{{
\mathbf{q}}},\ \forall\mathbf{q}\in\mathcal{P}.$

Due to our (zero) initialization of $Q^{(0)}$, we have $\Phi^{(0)}_{{\mathbf{p}}\to{v_{k,m}}}=\Upsilon^{(0)}_{{\mathbf{p}}\to{\mathbf{
q}}}=0$.

These two types of attention correlations, FP correlation $\Phi_{\mathbf{p}\to v_{k,m}}$ and PP correlation $\Phi_{\mathbf{p}\to\mathbf{q}}$, act as the exponent terms within the $\operatorname{softmax}$ calculations for attention scores. Given $\mathbf{p}\in\mathcal{P}_{k,j}$ is masked, the (unnormalized) attention $\mathbf{attn}_{\mathbf{p}\to\mathbf{q}}$ directed towards an unmasked patch $\mathbf{q}$ is influenced jointly by these correlations. Hence, the described attention pattern can emerge from either a substantial FP correlation $\Phi_{\mathbf{p}\to v_{k,j}}$ or a significant PP correlation $\Phi_{\mathbf{p}\to\mathbf{q}}$ for $\mathbf{q}$ in the same area as $\mathbf{p}$. However, in our setting, the latter mechanism—learning via PP correlation—fails to produce desired attention patterns, as explained below on a notable gap that prior work [] did not address.

###### Definition 3.1 .

##### Can pure positional attention explain the transformer’s ability to learn locality?

[] theoretically analyzed how ViTs can identify spatially localized patterns by minimizing the supervised cross-entropy loss with GD. Their analysis focused on a spatially structured dataset equivalent to our settings when $K=1$ without distinguishing the global and local features. Hence, the association between patch $\mathbf{p}$ and $\mathbf{q}$ is fixed since there is only a single cluster. Based on such invariant patch-wise relations, they showed that ViTs can achieve optimal attention patterns by solely learning the “patch association”, where $e_{\mathbf{p}}^{\top}Qe_{\mathbf{q}}$ is large for ${\mathbf{p}},{\mathbf{q}}$ coming from the same area, corresponding to large PP correlation $\Gamma_{\mathbf{p}\to\mathbf{q}}$ in our settings. However, such an assumption of invariant patch associations is often unrealistic for vision datasets in practice, for instance, a cube-shaped building typically requires a different attention pattern from a bird inside the woods. Therefore, when multiple patterns appear in the data distribution (e.g., in our settings with $K>1$ clusters), relying solely on PP correlations is insufficient, and is often undesirable due to variations of PP correlations across clusters. This highlights the necessity of examining FP correlations for a deeper understanding of the local representation power of transformers.

## 4 Statements of Main Results

In this section, we present the main theorems of this paper. 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.

##### Information gap and a technical condition.

Based on our data model in Section 2.2, we introduce a notion of information gap to quantify the difference of significance between global and local areas (cf. Assumption 2.2). Denoted as $\Delta$, the information gap is defined as follows:

$$ $\Delta\coloneqq(1-\kappa_{s})-2(1-\kappa_{c}).$ (4.1) $$

Our study focuses on the regime where $\Delta$ is not too close to zero, i.e., $|\Delta|=\Omega(1)$, which allows for cleaner induction arguments. This condition could be potentially relaxed via more involved analysis.

##### Notations for theorem and proof presentations.

Based on the crucial role of those unmasked patches for the reconstruction task, we further define the unmasked area attention as follows:

$$ $\displaystyle\textstyle\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}( \mathsf{M}(X);Q)\coloneqq\widetilde{\mathbf{Attn}}_{{\mathbf{p}}\to\mathcal{U} \cap\mathcal{P}_{k,m}}(\mathsf{M}(X);Q).$ $$

We abbreviate $\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}(\mathsf{M}(X);Q^{(t)})$ ($\mathbf{attn}_{\mathbf{p}\to\mathbf{q}}(\mathsf{M}(X);Q^{(t)})$) as $\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}^{(t)}$( $\mathbf{attn}_{\mathbf{p}\to\mathbf{q}}^{(t)}$). Moreover, to properly evaluate the reconstruction performance, we further introduce the following notion of the reconstruction loss with respect to a specific patch $\mathbf{p}\in\mathcal{P}$:

$$ $\displaystyle\textstyle\mathcal{L}_{\mathbf{p}}(Q)$ $\displaystyle=\frac{1}{2}\mathbb{E}\left[\operatorname{\mathds{1}}\{{\mathbf{p }}\in\mathcal{M}\}\Big{\|}[F(\mathsf{M}(X);Q)]_{{\mathbf{p}}}-X_{{\mathbf{p}}} \Big{\|}^{2}\right].$ (4.2) $$

Now we present our first main result regarding the convergence of masked pretraining.

###### Theorem 4.1 (Training convergence) .

Suppose the information gap $\Delta\in[-0.5,-\Omega(1)]\cup[\Omega(1),1]$. For any $0<\epsilon<1$, suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$. We train the transformer in Def. 2.4 by GD to minimize reconstruction loss in (2.9) with $\eta\ll\operatorname{poly}(P)$.
Then for each patch $\mathbf{p}\in\mathcal{P}$, we have

- 1.
Loss converges: $\mathcal{L}_{\mathbf{p}}(Q^{(T^{\star})})-\mathcal{L}_{\mathbf{p}}^{\star}\leq\epsilon$ in $T^{\star}=O\Big{(}\frac{1}{\eta}\log(P)P^{\max\{2(\frac{U}{L}-1),1\}(1-\kappa_
{s})}+\frac{1}{\eta\epsilon}\log\big{(}\frac{P}{\epsilon}\big{)}\Big{)}$ iterations, where ${\mathcal{L}}_{\mathbf{p}}^{\star}$ is the global minimum of the patch-level reconstruction loss in (4.2).
- 2.
Area-wide pattern of attention: given cluster $k\in[K]$, and $\mathbf{p}\in\mathcal{P}_{k,j}$ for some $j\in[N_{k}]$, if $X_{\mathbf{p}}$ is masked, then the one-layer transformer nearly “pays all attention” to all unmasked patches in the same area $\mathcal{P}_{k,j}$ as $\mathbf{p}$, i.e.,
$\textstyle\Big{(}1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,j}}^{(T^{\star
})}\Big{)}^{2}\leq O(\epsilon).$

Theorem 4.1 indicates that, at the time of convergence, for any masked query patch $X_{\mathbf{p}}$ in the $k$-th cluster, the transformer exhibits an *area-wide* pattern of attention, concentrating on those unmasked patches within the area that $\mathbf{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 $\mathbf{p}$ belongs to the global or local areas, which jointly highlight the diverse local patterns for masked vision pretraining.

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 area-wide attention patterns as suggested in Section 3. The intuitive explanations are provided in Section 5.

###### Theorem 4.2 (Learning Feature-Position correlations) .

Following the same assumptions in Theorem 4.1, for $\mathbf{p}\in\mathcal{P}$, given $k\in[K]$, if $\mathbf{p}\in\mathcal{P}_{k,j}$ for some $j\in[N_{k}]$, we have

For positive information gap $\Delta\in[\Omega(1),1]$:

- a.
Global areas ($j=1$) learn FP correlation in one-phase: $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ monotonically increases to $O(\log(P)/{\eta})$ throughout the training, with all other attention correlations remaining close to 0 0.
- b.
Local areas ($j>1$) learn FP correlation in two-phase:
In phase one, FP correlation $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ between the local area and the global area feature quickly decreases to $-\Theta(\log(P))$ whereas all other attention correlations stay close to zero;
In phase two, FP correlation $\Phi^{(t)}_{\mathbf{p}\to v_{k,j}}$ for the target local area starts to grow until convergence with all other attention correlations nearly unchanged.

For negative information gap $\Delta\in[-0.5,-\Omega(1)]$:

- c.
All areas
learn FP correlation through one-phase: $\Phi^{(t)}_{\mathbf{p}\to v_{k,j}}$ monotonically increases to $O(\log(P)/{\eta})$ throughout the training, with all other attention correlations remain close to 0 0.

The training dynamics are different depending on whether $\Delta$ is positive or negative, and further vary for positive $\Delta$ depending on whether $X_{\mathbf{p}}$ is situated in global or local areas. Typically, the target FP correlations are learned directly in a single phase. However, for a positive information gap $\Delta$, when patch $\mathbf{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. Such behavior of the learning dynamics will be further explained in Section 5 through our main proof ideas.

###### Theorem 4.1 (Training convergence) .

###### Theorem 4.2 (Learning Feature-Position correlations) .

## 5 Proof Sketch

The framework for constructing attention correlations, as outlined in Section 3, offers insights into how transformers may achieve specific localities during masked pretraining. Within this framework, we start this section with a warm-up discussion to further clarify why target FP correlations, rather than PP correlations, should be the focus of learning. We then offer an overview of our proof, which analyzes the dynamics of these attention correlations.

##### Warm-up intuition for attention correlations.

Let us start by exploring the role of attention correlations in determining the area attention scores.
Given a masked input $\mathsf{M}(X)$ with $\mathbf{p}\in\mathcal{P}_{k,j}\cap\mathcal{M}$, it holds that

$$ $\displaystyle\widetilde{\mathbf{Attn}}_{{\mathbf{p}}\to\mathcal{P}_{k,j}}( \mathsf{M}(X);Q)\propto\sum_{\mathbf{q}\in\mathcal{U}\cap\mathcal{P}_{k,j}}e^{ \Phi_{\mathbf{p}\to v_{k,j}}+\Upsilon_{\mathbf{p}\to\mathbf{q}}}+\sum_{\mathbf {q}\in\mathcal{M}\cap\mathcal{P}_{k,j}}e^{\Upsilon_{\mathbf{p}\to\mathbf{q}}}.$ $$

The first term on the RHS is proportional to our unmasked area attention $\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,j}}$, which is balanced by the relative magnitudes between Feature-Position and Position-Position correlations. Upon initial observation, large FP correlation $\Phi_{\mathbf{p}\to v_{k,j}}$ or PP correlation $\Upsilon_{{\mathbf{p}}\to{\mathbf{q}}}$ for patch $\mathbf{q}\in\mathcal{P}_{k,j}$ may both contribute to the attention towards the area $\mathcal{U}\cap\mathcal{P}_{k,j}$. However, two key issues can arise from the PP correlation, which jointly lead to the suppression of PP correlation for reconstruction:
i). such a mechanism inadvertently directs attention to the masked patches, which is not desirable; ii). such position association could be vulnerable to the variation across different clusters, i.e., $\mathbf{p},\mathbf{q}\in\mathcal{P}_{k,j}$ does not necessarily hold for all $k\in[K]$. Consequently, as illustrated in Figure 4, to obtain a perfect reconstruction of the masked patch, learning a relatively large FP correlation $\Phi_{\mathbf{p}\to v_{k,j}}$ from the patches to the correct features is required.

Figure: Figure 4: The mechanism of how the masked patch attends to other patches via attention correlations.
Refer to caption: extracted/5633500/attention.png

##### Main idea of the proof.

The main idea of our analysis is to track the dynamics of those attention correlations. We first provide the following lemma of GD updates for $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,m}}$ and $\Upsilon_{\mathbf{p}\to\mathbf{q}}^{(t)}$.

###### Lemma 5.1 (Gradient of attention correlations, informal) .

Given ${\mathbf{p}},{\mathbf{q}}\in\mathcal{P}$, let $\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,m}}}=\frac{1}{\eta}\big{(}\Phi^{(t+1)}_{{
\mathbf{p}}\to v_{k,m}}-\Phi^{(t)}_{{\mathbf{p}}\to v_{k,m}}\big{)}$ with $k\in[K],m\in[N_{k}]$, and $\beta^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}=\frac{1}{\eta}\big{(}\Upsilon^{(t+1)
}_{{\mathbf{p}}\to{\mathbf{q}}}-\Upsilon^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}
\big{)}$. We use $a_{k,\mathbf{p}}$ to indicate that the index of the area that patch $\mathbf{p}$ is located in the cluster $\mathcal{D}_{k}$.

- a.
For the same area, $\textstyle\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,a_{k,\mathbf{p}}}}}\approx\mathbf
{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}\left(1-\mathbf{
Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}\right)^{2};$
- b.
If $a_{k,\mathbf{p}}\not=1$, for the global area, $\textstyle\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,1}}}\approx-\big{(}\mathbf{Attn}^
{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}\big{)}^{2};$
- c.
For other area $m\notin\{a_{k,\mathbf{p}}\}\cup\{1\}$, $\textstyle|\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,m}}}|=O\Big{(}\frac{\alpha^{(t)}
_{{\mathbf{p}}\to{v_{k,a_{k,\mathbf{p}}}}}+\operatorname{\mathds{1}}\{n\not=1
\}|\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,1}}}|}{N_{k}}\Big{)};$
- d.
$\beta^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}=\sum_{k\in[K]}\beta^{(t)}_{k,{
\mathbf{p}}\to{\mathbf{q}}}$, where $|\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}|=O\Big{(}\frac{|\alpha^{(t)}_{k,{
\mathbf{p}}\to a_{k,\mathbf{q}}}|}{C_{k,a_{k,\mathbf{q}}}}\Big{)}$.

Here, $\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}$ can be interpreted as the projected gradient for PP correlation $\Upsilon^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}$ on the $k$-th cluster. Thus, Lemma 5.1.d. directly implies that the overall increment of PP correlation $\Upsilon^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}$ should be negligible compared to the corresponding FP correlations $\Phi^{(t)}_{{\mathbf{p}}\to{v_{k,a_{k,\mathbf{q}}}}}$ on certain cluster $\mathcal{D}_{k}$ since $C_{k,a_{k,\mathbf{q}}}\gg 1$ and $K=\Theta(1)$, which implies that all PP correlations are small.

Moreover, from Lemma 5.1, it is observed that the target FP correlation $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,a_{k,\mathbf{p}}}}$ exhibits a monotonically increasing trend, while FP-correlations for feature $v_{k,1}$ will monotonically decrease if $a_{k,\mathbf{p}}\not=1$. Moreover, FP gradients of the patch $\mathbf{p}$ to other unrelated local features are significantly smaller and become negligible since $N_{k}\gg 1$. Based on these observations, we now explain how the trend of FP-correlations determines the different learning behaviors varied across different locations $a_{k,\mathbf{p}}$ and information gap $\Delta$.

- •
For local areas with $a_{k,p}>1$: at the beginning of training, $\mathbf{attn}^{(0)}_{\mathbf{p}\to{\mathbf{q}}}=\frac{1}{P}$ for any ${\mathbf{q}}\in\mathcal{P}$ due to zero initialization. However, with high probability, the count of unmasked patches in the global area $\mathcal{P}_{k,1}$ significantly exceeds that in other areas.
1.
$\Delta=(1-\kappa_{s})-2(1-\kappa_{c})\geq\Omega(1)$: $|\alpha^{(0)}_{{\mathbf{p}}\to{v_{k,1}}}|\approx{\big{(}\mathbf{Attn}_{\mathbf
{p}\to\mathcal{P}_{k,1}}^{(0)}\big{)}}^{2}=\Omega(\frac{1}{P^{2(1-\kappa_{c})}
})\gg\Theta(\frac{1}{P^{1-\kappa_{s}}})=\mathbf{Attn}_{\mathbf{p}\to\mathcal{P
}_{k,a_{k,\mathbf{p}}}}^{(0)}\approx\alpha^{(0)}_{{\mathbf{p}}\to{v_{k,a_{k,
\mathbf{p}}}}}$. Thus, $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ enjoys a much larger decreasing rate initially, which defines the decoupling phase of non-target global FP correlations. The significant decrease in
$\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ leads to a reduction in $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(t)}$, which in turn makes a drop in $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|$. As the gradient decreases below $\alpha_{\mathbf{p}\to v_{k,a_{k,pb}}}^{(t)}$, $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ within the target local area begins to have a larger gradient, which triggers the growth of target FP correlation phase, leading to continuous growth until convergence.
2.
$\Delta\leq-\Omega(1)$: the above relationship reverses, and $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ starts with a fast initial increase, which moves directly into the growth phase, eliminating the need to decouple $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$.
- •
For global areas with $a_{k,\mathbf{p}}=1$: the global area is exactly the target area that FP correlations should be learned, and thus the second term in Lemma 5.1 will not exist in this case. As a result, the training process also skips the initial decoupling phase and starts with the growth phase for $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$. Such an effect is independent of the value of $\Delta$.

###### Lemma 5.1 (Gradient of attention correlations, informal) .

## 6 Experiments

Previous studies on the attention mechanisms of ViT-based pre-training approaches have mainly utilized a metric known as the attention distance []. Such a metric quantifies the average spatial distance between the query and key tokens, weighted by their self-attention coefficients. 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 a 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, we introduce a novel and revised version of average attention distance, called attention diversity metric, which is designed to assess whether various patches are concentrating on a similar region, thereby directly capturing global information.

##### Attention diversity metric, in distance.

This metric is computed for self-attention with a single head of the specific layer. For a given image divided into $P\times P$ patches, the process unfolds as follows: for each patch, it is employed as the query patch to calculate the attention weights towards all $P^{2}$ patches, and those with the top-$n$ attention weights are selected. Subsequently, the coordinates (e.g. $(i,j)$ with $i,j\in[P]$) of these top-$n$ patches are concatenated in sequence to form a $2\times n$-dimensional vector. The final step computes the average distance between all these $2n$-dimensional vectors, i.e., $P^{2}\times P^{2}$ vector pairs.

##### Setup.

In this work, we compare the performance of ViT-B/16 encoder pre-trained on ImageNet-1K [] among the following four models: masked reconstruction(MAE), contrastive learning (MoCo v3 []), other discriminative self-supervised (DINO []), and supervised learning (DeiT []). We focus on $12$ different attention heads in the last layer of ViT-B on different pre-trained models. The box plot visualizes the distribution of the top-10 averaged attention focus across 152 example images, as similarly done in [].

##### Implications.

The experiment results based on our new metric are provided in Figure 2. Lower values of the attention diversity metric signify a 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 more diverse attention patterns, setting it apart from other models that prioritize a uniform global information with less attention diversity. This aligns with and provides further evidence for the findings in [].

## 7 Additional 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 [, , ], robustness [, ], and role of positional embeddings [, ].
Recent studies [, , ] have delved into ViTs with self-supervision to uncover the mechanisms at play, particularly through visualization and analysis of metrics related to self-attention. [] 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 on tasks with weak semantics following fine-tuning. Building on MIM’s advantages,
[] further proposed a simple feature distillation method that incorporates locality into various self-supervised methods, leading to an overall improvement in the finetuning performance. [] 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 a 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-supervised learning approaches, contrastive learning [, , , ], and its variant non-contrastive self-supervised learning [, , ]. Some other works study the mask prediction approach [, , ], which is the focus of this paper. [] provided statistical downstream guarantees for reconstructing missing patches. [] studied the benefits of head and prompt tuning with masked pretraining under a Hidden Markov Model framework. [] 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 [, , , , ], internal mechanism [, ], limitations [, ], and PAC learning []. 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 [, , , , , , , ]. Regarding the training dynamics of attention-based models, [] studied the training process of shallow ViTs in a classification task. Subsequent research expanded on this by exploring the graph transformer with positional encoding [] and in-context learning performance of transformers with nonlinear self-attention and nonlinear MLP []. However, all of these analyses
rely crucially on stringent assumptions on the initialization of transformers and hardly generalize to our setting. [] 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 [] considered a generalized case with multiple layers. [, ] investigated the implicit bias for self-attention models trained with GD.
Furthermore, [] proved the in-context convergence of a one-layer softmax transformer trained via GD and illustrated the attention dynamics throughout the training process. More recently, [] studied GD dynamics on 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.

## 8 Conclusion

In this work, we study the training process of masked vision pretraining with a one-layer softmax-based transformer. Our key contribution lies in showing that transformers trained with masked reconstruction objectives exhibit local and diverse patterns by learning correlations between patch features and patch positions. To our knowledge, our work is the first in analyzing softmax self-attention on both patch and position embedding simultaneously. Our proof techniques leveraged ideas from phase decomposition based on the interplay between feature-position and position-wise correlations, which do not need to disentangle patches and positional encodings as in prior works. We anticipate that our theory can be useful for future studies of the spatial-or-temporal structures of state-of-the-art transformers and can promote theoretical studies relevant to deep learning practice.

## Acknowledgements

The work of Z. Wen and Y. Chi is supported in part by NSF under CCF-1901199, CCF-2007911, DMS-2134080, and by ONR under N00014-19-1-2404. The work of Y. Liang is supported in part by NSF under RINGS-2148253, CCF-1900145, and DMS-2134145.

## Appendix A Overview of the Proof Techniques

In this section, we explain our key proof techniques in analyzing the masked pretraining of transformers. We focus on the reconstruction of a specific patch $X_{\mathbf{p}}$ for $\mathbf{p}\in\mathcal{P}$. We aim to elucidate the training phases through which the model learns FP correlations related to the area associated with $\mathbf{p}$ across different clusters $k\in[K]$.

Our characterization of training phases differentiates between whether $X_{\mathbf{p}}$ is located in the global or local areas and further varies based on whether $\Delta$ is positive or negative. Specifically, for $\Delta\in[\Omega(1),1]$, we observe distinct learning dynamics for FP correlations between local and global areas:

- •
Local area attends to FP correlation in two-phase: given $k\in[K]$, if $a_{k,\mathbf{p}}\neq 1$, then
1.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ first quickly decreases whereas all other $\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}$ with $m\not=1$ and $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}$ do not change much;
2.
after some point, the increase of $\Phi^{(t)}_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}$ takes dominance. Such $\Phi^{(t)}_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}$ will keep growing until convergence with all other FP and PP attention correlations nearly unchanged.
- •
Global areas learn FP correlation in one-phase: given $k\in[K]$, if $a_{k,\mathbf{p}}=1$, the update of $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ will dominate throughout the training, whereas all other $\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}$ with $m\not=1$ and learned PP correlations remain close to 0 0.

For $\Delta\in[-0.5,-\Omega(1)]$, the behaviors of learning FP correlations are uniform for all areas. Namely, all areas learn FP correlation through one-phase: given $k\in[K]$, throughout the training, the increase of $\Phi^{(t)}_{\mathbf{p}\to v_{k,a_{k,\bm{p}}}}$ dominates, whereas all other $\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}$ with $m\not=a_{k,\mathbf{p}}$ and PP correlations $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}$ remain close to 0 0.

For clarity, this section will mainly focus on the learning of local feature correlations with a positive information gap $\Delta\geq\Omega(1)$ in Sections A.2 and A.3, which exhibits a two-phase process.
The other scenarios will be discussed briefly in Section A.4.

### A.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 to track the dynamics of those attention correlations. We first provide the following GD updates of $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,m}}$ and $\Upsilon_{\mathbf{p}\to\mathbf{q}}^{(t)}$ (see Section B.1.1 for formal statements).

###### Lemma A.1 (FP correlations, informal) .

Given $k\in[K]$, for ${\mathbf{p}}\in\mathcal{P}$, denote $n=a_{k,\mathbf{p}}$, let $\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,m}}}=\frac{1}{\eta}\big{(}\Phi^{(t+1)}_{{
\mathbf{p}}\to v_{k,m}}-\Phi^{(t)}_{{\mathbf{p}}\to v_{k,m}}\big{)}$ for $m\in[N_{k}]$, and suppose $X_{\mathbf{p}}$ is masked. Then

- 1.
for the same area, $\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,n}}}\approx\mathbf{Attn}^{(t)}_{{\mathbf{p}
}\to\mathcal{P}_{k,n}}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_
{k,n}}\right)^{2};$
- 2.
if $k\in\mathcal{B}_{\mathbf{p}}$, for the global area,
$\displaystyle\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,1}}}$
$\displaystyle\approx-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}
\cdot\Bigg{(}\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}\left(1-
\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}\right)+\mathbf{Attn}^{(
t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}
}\to\mathcal{P}_{k,n}}\right)\Bigg{)};$
- 3.
for other area $m\notin\{n\}\cup\{1\}$,
$\displaystyle\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,m}}}\approx\mathbf{Attn}^{(t)}
_{{\mathbf{p}}\to\mathcal{P}_{k,m}}\Bigg{(}\operatorname{\mathds{1}}{\{n\not=1
\}}\left(\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}\right)^{2}-
\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)\mathbf{
Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\Bigg{)}.$

From Lemma A.1, it is observed that for ${\mathbf{p}}\in\mathcal{P}_{k,n}$, the feature correlation $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,n}}$ exhibits a monotonically increasing trend over time because $\alpha_{{\mathbf{p}}\to v_{k,n}}^{(t)}\geq 0$.
Furthermore, if $n>1$, i.e., $\mathcal{P}_{k,n}$ is the local area, $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,1}}$ will monotonically decrease.

###### Lemma A.2 (PP attention correlations, informal) .

Given ${\mathbf{p}},{\mathbf{q}}\in\mathcal{P}$, let $\beta^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}=\frac{1}{\eta}\big{(}\Upsilon^{(t+1)
}_{{\mathbf{p}}\to{\mathbf{q}}}-\Upsilon^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}
\big{)}$, and suppose $X_{\mathbf{p}}$ is masked. Then
$\beta^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}=\sum_{k\in[N]}\beta^{(t)}_{k,{
\mathbf{p}}\to{\mathbf{q}}}$, where $\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}$ satisfies

- 1.
if $a_{k,\mathbf{p}}=a_{k,\mathbf{q}}=n$, $\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}\approx\mathbf{attn}^{(t)}_{{
\mathbf{p}}\to{\mathbf{q}}}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to
\mathcal{P}_{k,n}}\right)^{2};$
- 2.
if $k\in\mathcal{B}_{\mathbf{p}}\cap\mathcal{C}_{\mathbf{q}}$, where $a_{k,\mathbf{p}}=n>1$ and $a_{k,\mathbf{q}}=1$:
$\displaystyle\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}$
$\displaystyle\approx-\mathbf{attn}^{(t)}_{{\mathbf{p}}\to\mathbf{q}}\cdot\Bigg
{(}\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}\left(1-\mathbf{Attn}
^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}\right)+\mathbf{Attn}^{(t)}_{{\mathbf
{p}}\to\mathcal{P}_{k,n}}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{
P}_{k,n}}\right)\Bigg{)};$
- 3.
if $a_{k,\mathbf{q}}=m\notin\{n\}\cup\{1\}$, where $a_{k,\mathbf{p}}=n$,
$\displaystyle\beta^{(t)}_{k,{\mathbf{p}}\to\mathbf{q}}\approx\mathbf{attn}^{(t
)}_{{\mathbf{p}}\to\mathbf{q}}\cdot\Bigg{(}\operatorname{\mathds{1}}{\{n\not=1
\}}\left(\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,1}}\right)^{2}-
\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)\mathbf{
Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\Bigg{)}.$

Based on the above gradient update for $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}$, we further introduce the following auxiliary quantity $\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}$, which can be interpreted as the PP attention correlation “projected” on the $k$-th cluster $\mathcal{D}_{k}$, and will be useful in the later proof.

$$ $\displaystyle\Upsilon^{(t+1)}_{k,\mathbf{p}\to\mathbf{q}}\coloneqq\Upsilon^{(t )}_{k,\mathbf{p}\to\mathbf{q}}+\eta\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}, \quad\text{ with }\Upsilon^{(0)}_{k,\mathbf{p}\to\mathbf{q}}=0.$ (A.1) $$

We can directly verify that $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}=\sum_{k\in[K]}\Upsilon^{(t)}_{k,
\mathbf{p}\to\mathbf{q}}$.

The key observation by comparing Lemma A.1 and A.2 is that the gradient of projected PP attention $\beta^{(t)}_{k,{\mathbf{p}}\to\mathbf{q}}$ is smaller than the corresponding FP gradient $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{q}}}}^{(t)}$ in magnitude since $\textstyle\mathbf{attn}^{(t)}_{{\mathbf{p}}\to\mathbf{q}}\approx\frac{\mathbf{
Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{q}}}}}{(1-\gamma)C_{k,
a_{k,\mathbf{q}}}}$. We will show that the interplay between the increase of $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,n}}$ and the decrease of $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,1}}$ determines the learning behaviors for the local patch $\mathbf{p}\in\mathcal{P}_{k,n}$ with $n>1$, and which effect will happen first depends on the initial attention, which is also determined by the value of information gap $\Delta$.

###### Lemma A.1 (FP correlations, informal) .

###### Lemma A.2 (PP attention correlations, informal) .

### A.2 Phase I: Decoupling the Global FP Correlations

We now explain how the attention correlations evolve at the initial phase of the training to decouple the correlations of the non-target global features when $\mathbf{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, $\Phi^{(0)}_{\mathbf{p}\to v_{k,m}}=\Upsilon^{(0)}_{k,\mathbf{p}\to\mathbf{q}}=0$, and hence $\mathbf{attn}^{(0)}_{\mathbf{p}\to{\mathbf{q}}}=\frac{1}{P}$ for any ${\mathbf{q}}\in\mathcal{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 $\mathcal{P}_{k,1}$ is much larger than others. Hence, $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(0)}=\frac{|\mathcal{U}\cap
\mathcal{P}_{k,1}|}{P}\geq\Omega(\frac{1}{P^{1-\kappa_{c}}})\gg\Theta(\frac{1}
{P^{1-\kappa_{s}}})=\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,m}}^{(0)}$ for $m>1$. Therefore, by Lemma A.1 and A.2, we immediately obtain

- •
$\alpha_{\mathbf{p}\to v_{k,1}}^{(0)}=-\Theta\Big{(}\frac{1}{P^{2(1-\kappa_{c})
}}\Big{)}$, whereas $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(0)}=\Theta\Big{(}\frac{1}{P^{(
1-\kappa_{s})}}\Big{)}$;
- •
all other FP correlation gradients $\alpha_{\mathbf{p}\to v_{k,m}}^{(0)}$ with $m\not=1,a_{k,\mathbf{p}}$ are small;
- •
all projected PP correlation gradients $\beta^{(0)}_{k,\mathbf{p}\to\mathbf{q}}$ are small.

Since $\Delta=(1-\kappa_{s})-2(1-\kappa_{c})\geq\Omega(1)$, it can be seen that $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ enjoys a much larger decreasing rate initially. This captures the decoupling process of the feature correlations with the global feature $v_{k,1}$ in the global area for $\mathbf{p}$. It can be shown that such an effect will dominate over a certain period that defines stage 1 of phase I. At the end of this stage, we will have $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}\leq-\Omega\left(\log(P)\right)$, whereas all FP attention correlation $\Phi_{\mathbf{p}\to v_{k,m}}^{(t)}$ with $m>1$ and all projected PP correlations $\Upsilon_{k,\mathbf{p}\to\mathbf{q}}^{(t)}$ stay close to 0 0 (see Section D.1).

During stage 1, the significant decrease of the global FP correlation $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ leads to a reduction in the attention score $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(t)}$. Meanwhile, attention scores $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}$ (where $m>1$) for other patches remain consistent, reflecting a uniform distribution over unmasked patches within each area. By the end of stage 1, $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(t)}$ drops to a certain level, resulting in a decrease in $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|$ as it approaches $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}$, which indicates that stage 2 begins.

##### Stage 2.

Soon as stage 2 begins, the dominant effect switches as $|\alpha^{(t)}_{\mathbf{p}\to v_{k,1}}|$ reaches the same order of magnitude as $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$. The following result shows that $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ must update during stage 2.

###### Lemma A.3 (Switching of dominant effects (See Section D.2 )) .

Under the same conditions as Theorem 4.1, for $\mathbf{p}\in\mathcal{P}$, there exists $\widetilde{T}_{1}$, such that at iteration $t=\widetilde{T}_{1}+1$, we have

- a.
$\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(\widetilde{T}_{1}+1)}\geq\Omega
\left(\log(P)\right)$, and $\Phi_{\mathbf{p}\to v_{k,1}}^{(\widetilde{T}_{1}+1)}=-\Theta(\log(P))$;
- b.
all other FP correlations $\Phi_{\mathbf{p}\to v_{k,m}}^{(t)}$ with $m\not=1,a_{k,\mathbf{p}}$ are small;
- c.
all projected PP correlations $\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}$ are small.

###### Lemma A.3 (Switching of dominant effects (See Section D.2 )) .

##### Intuition of the transition.

Once $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ decreases to $-\frac{\Delta}{2L}\log(P)$, we observe that $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|$ is approximately equal to $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$. After this point, reducing $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ further is more challenging compared to the increase in $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$.
To illustrate, a minimal decrease of $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ by an amount of $\frac{0.001}{L}\log(P)$ will yield $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\leq O(\frac{\alpha_{\mathbf{p}\to v_{k,
n}}^{(t)}}{P^{0.002}})$. Such a discrepancy triggers the switch of the dominant effect.

### A.3 Phase II: Growth of Target Local FP Correlation

Moving beyond phase I, FP correlation $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ within the target local area $\mathbf{p}$ already enjoys a larger gradient $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ than other $\Phi_{\mathbf{p}\to v_{k,m}}^{(t)}$ with $m\not=a_{k,\mathbf{p}}$ and all projected PP correlations $\Upsilon_{k,\mathbf{p}\to\mathbf{q}}^{(t)}$. We can show that the growth of $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ will continue to dominate until the end of training by recognizing the following two stages.

##### Rapid growth stage.

At the beginning of phase II, $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ is mainly driven by $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}^{(t)}$ since $1-\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}^{(t)}$ remains at the constant order. Therefore, the growth of $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ naturally results in a boost in $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}^{(t)}$, thereby promoting an increase in its own gradient $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$, 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 Section D.3):

- •
all other FP correlation gradients $\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}$ with $m\not=a_{k,\mathbf{p}}$ are small;
- •
all projected PP correlation gradients $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}$ are small.

##### Convergence stage.

After the rapid growth stage, the desired local pattern with a high target feature-position correlation $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ 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 $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ starts to decelerate as $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ reaches $\Theta(\log(P))$, resulting in $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}\approx\Omega(1)$, which leads to convergence (see Section D.4).

### A.4 Learning Processes in Other Scenarios

In this section, we 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 to learn PP correlations.

##### What is the role of positive information gap?

As described in stage 1 of phase 1 in Section A.2, the decoupling effect happens at the beginning of the training because $\alpha_{\mathbf{p}\to v_{k,1}}^{(0)}\gg\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf
{p}}}}^{(0)}$ attributed to $\Delta\geq\Omega(1)$. However, in cases where $\Delta\leq-\Omega(1)$, this relationship reverses, with $\alpha_{\mathbf{p}\to v_{k,1}}^{(0)}$ becoming significantly smaller than $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(0)}$. Similarly, other FP gradients $\alpha_{\mathbf{p}\to v_{k,m}}^{(0)}$ with $m\not=1,a_{k,\mathbf{p}}$ and all the projected gradients of PP correlation $\beta_{\mathbf{p}\to\mathbf{q}}^{(0)}$ are small in magnitude. Consequently, $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ starts with a larger gradient, eliminating the need to decouple FP correlations for the global area. As a result, training skips the initial phase, and moves directly into Phase II, during which $\Phi_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(t)}$ continues to increase until it converges (see Appendix E).

##### Learning FP correlations for the global area.

When the patch $X_{\mathbf{p}}$ is located in the global area of cluster $k$, i.e., $a_{k,\mathbf{p}}=1$, the attention score $\mathbf{Attn}^{(0)}_{\mathbf{p}\to\mathcal{P}_{k,1}}$ directed towards the target area $\mathcal{P}_{k,1}$ is initially higher compared to 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 $\alpha_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}^{(0)}$. Such an effect is independent of the value of $\Delta$. As a result, the training process skips the initial phase, which is typically necessary for the cases where $a_{k,\mathbf{p}}>1$ with a positive information gap, and moves directly into Phase II (see Appendix F).

##### All PP correlations are small.

Integrating the analysis from all previous discussions, we establish that for every cluster $k\in[K]$, regardless of its association with $\mathcal{C}_{\mathbf{p}}$ (global area) or $\mathcal{B}_{\mathbf{p}}$ (local area), and for any patch $X_{\mathbf{q}}$ with $\mathbf{q}\in\mathcal{P}$, the projected PP correlation $\Upsilon_{k,\mathbf{p}\to\mathbf{q}}^{(t)}$ remains nearly zero in comparison to the significant changes observed in the FP correlation, because the gradient $\beta_{k,\mathbf{p}\to\mathbf{q}}^{(t)}$ is relatively negligible. Therefore, the overall PP correlation $\Upsilon_{\mathbf{p}\to\mathbf{q}}^{(t)}=\sum_{k=1}^{K}\Upsilon_{k,\mathbf{p}
\to\mathbf{q}}^{(t)}$ also stays close to zero, given that the number of clusters $K=\Theta(1)$.

## Appendix B 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. Throughout the appendix, we assume $N_{k}=N$ and $C_{k,n}=C_{n}$ for all $k\in[K]$ for simplicity. We will also omit the explicit dependence on $X$ for $z_{n}(X)$. We use $k_{X}\in[K]$ to denote the cluster index that a given image $X$ is drawn from.

### B.1 Gradient Computations

We first calculate the gradient with respect to $Q$. We omit the superscript ‘$(t)$’ and write $\mathcal{L}(Q)$ as $\mathcal{L}$ here for simplicity.

###### Lemma B.1 .

The gradient of the loss function with respect to $Q$ is given by

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial Q}=-\mathbb{E}$ $\displaystyle\left[\sum_{{\mathbf{p}}\in\mathcal{M}}\right.\sum_{{\mathbf{q}}} \mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\mathsf{M}(X)^{\top}_{{\mathbf{q}}} (X_{{\mathbf{p}}}-[F(\mathsf{M}(X);Q)]_{{\mathbf{p}}})\cdot$ $\displaystyle\qquad\left.\widetilde{\mathsf{M}}(X)_{{\mathbf{p}}}\left( \widetilde{\mathsf{M}}(X)_{{\mathbf{q}}}-\sum_{{\mathbf{r}}}\mathbf{attn}_{{ \mathbf{p}}\to{\mathbf{r}}}\widetilde{\mathsf{M}}(X)_{{\mathbf{r}}}\right)^{ \top}\right].$ $$

###### Proof.

We begin with the chain rule and obtain

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial Q}$ $\displaystyle=\mathbb{E}[\sum_{{\mathbf{p}}\in\mathcal{M}}\frac{\partial[F( \mathsf{M}(X);Q)]_{{\mathbf{p}}}}{\partial Q}([F(\mathsf{M}(X);Q)]_{{\mathbf{p }}}-X_{{\mathbf{p}}})]$ $\displaystyle=\mathbb{E}[\sum_{{\mathbf{p}}\in\mathcal{M}}\sum_{{\mathbf{q}}} \frac{\partial\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}}{\partial Q}\mathsf{ M}(X)^{\top}_{{\mathbf{q}}}([F(\mathsf{M}(X);Q)]_{{\mathbf{p}}}-X_{{\mathbf{p} }})].$ (B.1) $$

We focus on the gradient for each attention score:

$$ $\displaystyle\frac{\partial\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}}{ \partial Q}$ $\displaystyle=\sum_{{\mathbf{r}}}\frac{\exp\left(\widetilde{\mathsf{M}}(X)^{ \top}_{{\mathbf{p}}}Q(\widetilde{\mathsf{M}}(X)_{{\mathbf{r}}}+\widetilde{ \mathsf{M}}(X)_{{\mathbf{q}}})\right)}{\left(\sum_{{\mathbf{r}}}\exp( \widetilde{\mathsf{M}}(X)^{\top}_{{\mathbf{p}}}Q\widetilde{\mathsf{M}}(X)_{{ \mathbf{r}}})\right)^{2}}\widetilde{\mathsf{M}}(X)_{{\mathbf{p}}}(\widetilde{ \mathsf{M}}(X)_{{\mathbf{q}}}-\widetilde{\mathsf{M}}(X)_{{\mathbf{r}}})^{\top}$ $\displaystyle=\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\sum_{{\mathbf{r}}} \mathbf{attn}_{{\mathbf{p}}\to{\mathbf{r}}}\widetilde{\mathsf{M}}(X)_{{\mathbf {p}}}(\widetilde{\mathsf{M}}(X)_{{\mathbf{q}}}-\widetilde{\mathsf{M}}(X)_{{ \mathbf{r}}})^{\top}$ $\displaystyle=\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\widetilde{\mathsf{M} }(X)_{{\mathbf{p}}}\cdot\left[\widetilde{\mathsf{M}}(X)_{{\mathbf{q}}}-\sum_{{ \mathbf{r}}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{r}}}\widetilde{\mathsf{M}}( X)_{{\mathbf{r}}}\right]^{\top}.$ $$

Substituting the above equation into (B.1), we complete the proof.
∎

Recall that the quantities $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,m}}$ and $\Upsilon^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}$ 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 B.2 .

(Attention correlations)
Given ${\mathbf{p}},{\mathbf{q}}\in\mathcal{P}$
, for $t\geq 0$, we define two types of attention correlations as follows:

- 1.
Feature Attention Correlation: $\Phi^{(t)}_{{\mathbf{p}}\to v_{k,m}}\coloneqq e^{\top}_{{\mathbf{p}}}Q^{(t)}v_
{k,m}$ for $k\in[K]$ and $m\in[N]$;
- 2.
Positional Attention Correlation: $\Upsilon^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}\coloneqq e_{{\mathbf{p}}}^{\top}Q
^{(t)}e_{{\mathbf{q}}}$.

By our initialization, we have $\Phi^{(0)}_{{\mathbf{p}}\to v_{k,m}}=\Upsilon^{(0)}_{{\mathbf{p}}\to{\mathbf{q
}}}=0$.

Next, we will apply the expression in Lemma B.1 to compute the gradient dynamics of these attention correlations.

###### Lemma B.1 .

###### Proof.

###### Definition B.2 .

#### B.1.1 Formal Statements and Proof of Lemma A.1 and A.2

We first introduce some notations. Given ${\mathbf{r}}\in\mathcal{U}$, for ${\mathbf{p}}\in\mathcal{P}$, $k\in[K]$ and $n\in[N]$ define the following quantities:

$$ $\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$ $\displaystyle\coloneqq\mathsf{M}(X)_{{\mathbf{r}}}^{\top}(X_{{\mathbf{p}}}-[F( \mathsf{M}(X);Q)]_{{\mathbf{p}}})$ $\displaystyle I^{{\mathbf{p}},k,n}_{{\mathbf{r}}}$ $\displaystyle\coloneqq\left(\widetilde{\mathsf{M}}(X)_{{\mathbf{r}}}-\sum_{ \mathbf{w}\in\mathcal{P}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}\widetilde{ \mathsf{M}}(X)_{\mathbf{w}}\right)^{\top}v_{k,n}$ $\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$ $\displaystyle\coloneqq\left(\widetilde{\mathsf{M}}(X)_{{\mathbf{r}}}-\sum_{ \mathbf{w}\in\mathcal{P}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}\widetilde{ \mathsf{M}}(X)_{\mathbf{w}}\right)^{\top}e_{{\mathbf{q}}}$ $$

###### Lemma B.3 (Formal statement of Lemma A.1 ) .

Given $k\in[K]$, for ${\mathbf{p}}\in\mathcal{P}$, denote $n=a_{k,\mathbf{p}}$, let $\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,m}}}=\frac{1}{\eta}\big{(}\Phi^{(t+1)}_{{
\mathbf{p}}\to v_{k,m}}-\Phi^{(t)}_{{\mathbf{p}}\to v_{k,m}}\big{)}$ for $m\in[N_{k}]$, then

- a.
for $m=n$,
$\displaystyle\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,n}}}=\mathbb{E}$
$\displaystyle\Bigg{[}\mathds{1}\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}\mathbf{
Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\cdot$
$\displaystyle\quad\left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to
\mathcal{P}_{k,n}}\right)^{2}+\sum_{a\not=n}z_{a}^{2}z_{n}\left(\mathbf{Attn}^
{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}\right)\Bigg{]};$
- b.
for $m\not=n$,
$\displaystyle\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,m}}}=\mathbb{E}\Bigg{[}\mathds
{1}{\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}}\mathbf{Attn}^{(t)}_{{\mathbf{p}}
\to\mathcal{P}_{k,m}}\cdot\Bigg{(}\sum_{a\not=m,n}z_{a}^{2}z_{m}\left(\mathbf{
Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}-$
$\displaystyle\quad\left(z_{m}z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}
}\to\mathcal{P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{
k,n}}+z_{m}^{3}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}
\right)\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}\right)\Bigg{)}
\Bigg{]}.$

###### Proof.

From Lemma B.1, we have

$$ $\displaystyle\alpha^{(t)}_{{\mathbf{p}}\to{v_{k,m}}}$ $\displaystyle=e_{{\mathbf{p}}}^{\top}(-\frac{\partial\mathcal{L}}{\partial Q}) {v_{k,m}}$ $\displaystyle=\mathbb{E}[\mathds{1}{\{{\mathbf{p}}\in\mathcal{M}\}}\sum_{{ \mathbf{r}}\in\mathcal{U}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{r}}}J^{{ \mathbf{p}}}_{{\mathbf{r}}}\cdot I^{{\mathbf{p}},k,m}_{{\mathbf{r}}}]$ $\displaystyle=\mathbb{E}[\mathds{1}{\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}} \sum_{{\mathbf{r}}\in\mathcal{U}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{r}}}J^ {{\mathbf{p}}}_{{\mathbf{r}}}\cdot I^{{\mathbf{p}},k,m}_{{\mathbf{r}}}]$ $$

where the last equality holds since when $k_{X}\not=k$, $I^{{\mathbf{p}},k,m}_{{\mathbf{r}}}=0$ due to orthogonality. Thus, in the following, we only need to consider the case $k_{X}=k$.

###### Lemma B.3 (Formal statement of Lemma A.1 ) .

###### Proof.

##### Case 1: m = n 𝑚 𝑛 m=n italic_m = italic_n .

- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}$, since $v_{k,n^{\prime}}\perp v_{k,n}$ for $n^{\prime}\not=n$, and $v_{k,n}\perp\{e_{{\mathbf{q}}}\}_{{\mathbf{q}}\in\mathcal{P}}$ we have
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}v_{k,n}^{\top}\left(z_{n}v_{k,n}-\sum_{{\mathbf{q}}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}z_
{n}v_{k,n}\right)$
$\displaystyle=z_{n}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}
}\right)$
$\displaystyle I^{{\mathbf{p}},k,n}_{{\mathbf{r}}}$
$\displaystyle=(z_{n}v_{k,n}-\sum_{{\mathbf{q}}\in\mathcal{U}\cap\mathcal{P}_{k
,n}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}z_{n}v_{k,n})^{\top}v_{k,n}=J^{
{\mathbf{p}}}_{{\mathbf{r}}}/z_{n}$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n^{\prime}}$ with $n^{\prime}\not=n$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n^{\prime}}v_{k,n^{\prime}}^{\top}\left(z_{n}v_{k,n}-\sum_{{
\mathbf{q}}\in\mathcal{U}\cap\mathcal{P}_{k,n^{\prime}}}\mathbf{attn}_{{
\mathbf{p}}\to{\mathbf{q}}}z_{n^{\prime}}v_{k,n^{\prime}}\right)$
$\displaystyle=-z_{n^{\prime}}^{2}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,
n^{\prime}}}$
$\displaystyle I^{{\mathbf{p}},k,n}_{{\mathbf{r}}}$
$\displaystyle=\left(z_{n^{\prime}}v_{k,n^{\prime}}-\sum_{{\mathbf{q}}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}z_
{n}v_{k,n}\right)^{\top}v_{k,n}$
$\displaystyle=-z_{n}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}$

Putting it together, then we obtain:

$$ $\displaystyle e_{{\mathbf{p}}}^{\top}(-\frac{\partial L}{\partial Q})v_{k,n}= \mathbb{E}$ $\displaystyle\left[\mathds{1}\{\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}\}\mathbf {Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right.\cdot$ $\displaystyle\left.\quad\left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p }}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{a\not=n}z_{a}^{2}z_{n}\left(\mathbf{ Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]$ $$

##### Case 2 2 2 2 : m ≠ n 𝑚 𝑛 m\not=n italic_m ≠ italic_n .

Similarly

- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}v_{k,n}^{\top}\left(z_{n}v_{k,n}-\sum_{{\mathbf{q}}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}z_
{n}v_{k,n}\right)$
$\displaystyle=z_{n}^{2}(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}})$
$\displaystyle I^{{\mathbf{p}},k,m}_{{\mathbf{r}}}$
$\displaystyle=\left(z_{n}v_{k,n}-\sum_{{\mathbf{q}}\in\mathcal{U}\cap\mathcal{
P}_{k,m}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}z_{m}v_{k,m}\right)^{\top}
v_{k,m}$
$\displaystyle=-z_{m}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,m}$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{m}v_{k,m}^{\top}\left(z_{n}v_{k,n}-\sum_{{\mathbf{q}}\in
\mathcal{U}\cap\mathcal{P}_{k,m}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{
q}}}z_{m}v_{k,m}\right)$
$\displaystyle=-z_{m}^{2}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}$
$\displaystyle I^{{\mathbf{p}},k,n}_{{\mathbf{r}}}$
$\displaystyle=\left(z_{m}v_{k,m}-\sum_{{\mathbf{q}}\in\mathcal{U}\cap\mathcal{
P}_{k,m}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}z_{m}v_{k,m}\right)^
{\top}v_{k,m}$
$\displaystyle=z_{n}(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}})$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,a}$, $a\not=n,m$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{a}v_{k,a}^{\top}\left(z_{n}v_{k,n}-\sum_{{\mathbf{q}}\in
\mathcal{U}\cap\mathcal{P}_{k,a}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{
q}}}z_{a}v_{k,a}\right)$
$\displaystyle=-z_{a}^{2}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}$
$\displaystyle I^{{\mathbf{p}},k,n}_{{\mathbf{r}}}$
$\displaystyle=\left(z_{a}v_{k,a}-\sum_{{\mathbf{q}}\in\mathcal{U}\cap\mathcal{
P}_{k,m}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}z_{m}v_{k,m}\right)^
{\top}v_{k,m}$
$\displaystyle=-z_{m}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}$

Putting them together, then we complete the proof. ∎

###### Lemma B.4 (Formal statement of Lemma A.2 ) .

Given ${\mathbf{p}},{\mathbf{q}}\in\mathcal{P}$, let $\beta^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}=\frac{1}{\eta}\big{(}\Upsilon^{(t+1)
}_{{\mathbf{p}}\to{\mathbf{q}})}-\Upsilon^{(t)}_{{\mathbf{p}}\to{\mathbf{q}})}
\big{)}$, then

$$ $\beta^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}=\sum_{k\in[K]}\beta^{(t)}_{k,{ \mathbf{p}}\to{\mathbf{q}}},\qquad\text{ where $\beta^{(t)}_{k,{\mathbf{p}}\to {\mathbf{q}}}$ satisfies}$ $$

- a.
if $a_{k,\mathbf{p}}=a_{k,\mathbf{q}}=n$,
$\displaystyle\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}=\mathbb{E}$
$\displaystyle\Bigg{[}\mathds{1}{\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}}\mathbf
{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}\cdot\Bigg{(}\sum_{a\not=n}z^{2}_{a}
\left(\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}+$
$\displaystyle z^{2}_{n}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}
_{k,n}}\right)\left(\mathds{1}{\{{\mathbf{q}}\in\mathcal{U}\}}-\mathbf{Attn}^{
(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)\Bigg{)}\Bigg{]};$
- b.
for $a_{k,\mathbf{p}}=n\not=m=a_{k,\mathbf{q}}$,
$\displaystyle\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}$
$\displaystyle=\mathbb{E}\Bigg{[}\mathds{1}{\{{\mathbf{p}}\in\mathcal{M},,k_{X}
=k\}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}\cdot\Bigg{(}\sum_{a\not
=n}z_{a}^{2}\left(\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right
)^{2}-$
$\displaystyle\left(z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to
\mathcal{P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}
}+\mathds{1}{\{{\mathbf{q}}\in\mathcal{U}\}}z_{m}^{2}\mathbf{Attn}^{(t)}_{{
\mathbf{p}}\to\mathcal{P}_{k,m}}\right)\Bigg{)}\Bigg{]}.$

###### Proof.

$$ $\displaystyle\beta^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}$ $\displaystyle=e_{{\mathbf{p}}}^{\top}(-\frac{\partial\mathcal{L}}{\partial Q}) e_{{\mathbf{q}}}=\mathbb{E}[\mathds{1}{\{{\mathbf{p}}\in\mathcal{M}\}}\sum_{{ \mathbf{r}}\in\mathcal{U}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{r}}}J^{ {\mathbf{p}}}_{{\mathbf{r}}}K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}]$ $$

Then we let

$$ $\displaystyle\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}\coloneqq\mathbb{E}[ \mathds{1}{\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}}\sum_{{\mathbf{r}}\in \mathcal{U}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{r}}}J^{{\mathbf{p}}}_ {{\mathbf{r}}}K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}].$ $$

In the following, we denote $a_{k,\mathbf{p}}=n$ and $a_{k,\mathbf{q}}=m$ for simplicity.

###### Lemma B.4 (Formal statement of Lemma A.2 ) .

###### Proof.

##### Case 1: m = n 𝑚 𝑛 m=n italic_m = italic_n .

If ${\mathbf{q}}\in\mathcal{U}\cap\mathcal{P}_{k,n}$:

- •
For ${\mathbf{r}}={\mathbf{q}}$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}v_{k,n}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{n
}v_{k,n}\right)$
$\displaystyle=z_{n}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}
}\right)$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{q}}}-(\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_
{{\mathbf{q}}}+\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}
\to\mathbf{w}}e_{\mathbf{w}}))^{\top}e_{{\mathbf{q}}}$
$\displaystyle=1-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}.$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}$, and ${\mathbf{r}}\not={\mathbf{q}}$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}v_{k,n}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{n
}v_{k,n}\right)$
$\displaystyle=z_{n}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}
}\right)$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-(\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_
{{\mathbf{q}}}+\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}
\to\mathbf{w}}e_{\mathbf{w}}))^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$
Thus
$\displaystyle\sum_{{\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{
attn}_{{\mathbf{p}}\to{\mathbf{r}}}J^{{\mathbf{p}}}_{{\mathbf{r}}}\cdot K^{{
\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}^{2}\left(1-\sum_{\mathbf{w}\in\mathcal{U}\cap\mathcal{P}_{
k,n}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}\right)$
$\displaystyle\cdot\left(-\sum_{{\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}
}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{r}}}\mathbf{attn}_{{\mathbf{p}}\to{
\mathbf{q}}}+\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\right)$
$\displaystyle=z_{n}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}
}\right)^{2}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,a}$, $a\not=n$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{a}v_{k,a}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,a}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{a
}v_{k,a}\right)$
$\displaystyle=-z_{a}^{2}\sum_{\mathbf{w}\in\mathcal{U}\cap\mathcal{P}_{k,a}}
\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-(\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_
{{\mathbf{q}}}+\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}
\to\mathbf{w}}e_{\mathbf{w}}))^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$

Thus

$$ $\displaystyle\sum_{{\mathbf{r}}\in\mathcal{U}}\mathbf{attn}_{{\mathbf{p}}\to{ \mathbf{r}}}J^{{\mathbf{p}}}_{{\mathbf{r}}}K^{{\mathbf{p}},{\mathbf{q}}}_{{ \mathbf{r}}}$ $\displaystyle=\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\cdot\left(z_{n}^{2} \left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{a \not=n}z_{a}^{2}\left(\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^ {2}\right)$ $$

If ${\mathbf{q}}\in\mathcal{M}\cap\mathcal{P}_{k,n}$:

- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}$,
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}v_{k,n}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{n
}v_{k,n}\right)$
$\displaystyle=z_{n}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}
}\right)$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-(\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_
{{\mathbf{q}}}+\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}
\to\mathbf{w}}e_{\mathbf{w}}))^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,a}$, $a\not=n$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{a}v_{k,a}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,a}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{a
}v_{k,a}\right)$
$\displaystyle=-z_{a}^{2}\sum_{\mathbf{w}\in\mathcal{U}\cap\mathcal{P}_{k,a}}
\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-(\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_
{{\mathbf{q}}}+\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}
\to\mathbf{w}}e_{\mathbf{w}}))^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$

Thus

$$ $\displaystyle\sum_{{\mathbf{r}}\in\mathcal{U}}\mathbf{attn}_{{\mathbf{p}}\to{ \mathbf{r}}}J^{{\mathbf{p}}}_{{\mathbf{r}}}K^{{\mathbf{p}},{\mathbf{q}}}_{{ \mathbf{r}}}$ $\displaystyle=\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\cdot\left(z_{n}^{2} \left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)^{2}-z_{n}^{2} \left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)+\sum_{a\not=n} z_{a}^{2}\left(\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}\right)$ $$

Putting it together,

$$ $\displaystyle\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}=\mathbb{E}\left[ \mathds{1}{\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}}\mathbf{attn}_{{\mathbf{p}} \to{\mathbf{q}}}\cdot\right.$ $\displaystyle\left.\left(-z_{n}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to \mathcal{P}_{k,n}}\right)\mathds{1}{\{{\mathbf{q}}\in\mathcal{M}\}}+z_{n}^{2} \left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m \not=n}z_{m}^{2}\left(\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}\right)^ {2}\right)\right]$ $$

##### Case 2: m ≠ n 𝑚 𝑛 m\not=n italic_m ≠ italic_n .

Similarly, if ${\mathbf{q}}\in\mathcal{U}\cap\mathcal{P}_{k,m}$:

- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}$,
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}v_{k,n}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{n
}v_{k,n}\right)$
$\displaystyle=z_{n}^{2}(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}})$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_{
{\mathbf{q}}}-\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}\to
\mathbf{w}}e_{\mathbf{w}})^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$
- •
For ${\mathbf{r}}={\mathbf{q}}$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{m}v_{k,m}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,m}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{m
}v_{k,m}\right)$
$\displaystyle=-z_{m}^{2}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{q}}}-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_{
{\mathbf{q}}}-\sum_{\mathbf{w}\not=\mathbf{w}}\mathbf{attn}_{{\mathbf{p}}\to
\mathbf{w}}e_{\mathbf{w}})^{\top}e_{{\mathbf{q}}}$
$\displaystyle=1-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,a}$, $a\not=n$, and ${\mathbf{r}}\not={\mathbf{q}}$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{a}v_{k,a}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,a}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{a
}v_{k,a}\right)$
$\displaystyle=-z_{a}^{2}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_{
{\mathbf{q}}}-\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}\to
\mathbf{w}}e_{\mathbf{w}})^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$

Thus

$$ $\displaystyle\sum_{{\mathbf{r}}\in\mathcal{U}}\mathbf{attn}_{{\mathbf{p}}\to{ \mathbf{r}}}J^{{\mathbf{p}}}_{{\mathbf{r}}}K^{{\mathbf{p}},{\mathbf{q}}}_{{ \mathbf{r}}}$ $\displaystyle=\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\cdot\left(-z_{n}^{2} \left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)\mathbf{Attn}_{ {\mathbf{p}}\to\mathcal{P}_{k,n}}-z_{m}^{2}\mathbf{Attn}_{{\mathbf{p}}\to \mathcal{P}_{k,m}}+\sum_{a\not=n}z_{a}^{2}\left(\mathbf{Attn}_{{\mathbf{p}}\to \mathcal{P}_{k,a}}\right)^{2}\right)$ $$

If ${\mathbf{q}}\in\mathcal{M}\cap\mathcal{P}_{k,m}$:

- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,n}$,
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{n}v_{k,n}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,n}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{n
}v_{k,n}\right)$
$\displaystyle=z_{n}^{2}(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}})$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_{
{\mathbf{q}}}-\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}\to
\mathbf{w}}e_{\mathbf{w}})^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$
- •
For ${\mathbf{r}}\in\mathcal{U}\cap\mathcal{P}_{k,a}$, $a\not=n$
$\displaystyle J^{{\mathbf{p}}}_{{\mathbf{r}}}$
$\displaystyle=z_{a}v_{k,a}^{\top}\left(z_{n}v_{k,n}-\sum_{\mathbf{w}\in
\mathcal{U}\cap\mathcal{P}_{k,a}}\mathbf{attn}_{{\mathbf{p}}\to\mathbf{w}}z_{a
}v_{k,a}\right)$
$\displaystyle=-z_{a}^{2}\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}$
$\displaystyle K^{{\mathbf{p}},{\mathbf{q}}}_{{\mathbf{r}}}$
$\displaystyle=(e_{{\mathbf{r}}}-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}e_{
{\mathbf{q}}}-\sum_{\mathbf{w}\not={\mathbf{q}}}\mathbf{attn}_{{\mathbf{p}}\to
\mathbf{w}}e_{\mathbf{w}})^{\top}e_{{\mathbf{q}}}$
$\displaystyle=-\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}$

Thus

$$ $\displaystyle\sum_{{\mathbf{r}}\in\mathcal{U}}\mathbf{attn}_{{\mathbf{p}}\to{ \mathbf{r}}}J^{{\mathbf{p}}}_{{\mathbf{r}}}K^{{\mathbf{p}},{\mathbf{q}}}_{{ \mathbf{r}}}$ $\displaystyle=\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\cdot\left(-z_{n}^{2} \left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)\mathbf{Attn}_{ {\mathbf{p}}\to\mathcal{P}_{k,n}}+\sum_{a\not=n}z_{a}^{2}\left(\mathbf{Attn}_{ {\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}\right).$ $$

Therefore

$$ $\displaystyle\beta^{(t)}_{k,{\mathbf{p}}\to{\mathbf{q}}}$ $\displaystyle=\mathbb{E}\left[\mathds{1}{\{{\mathbf{p}}\in\mathcal{M},k_{X}=k \}}\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}\cdot\right.$ $\displaystyle\left.\left(-z_{n}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to \mathcal{P}_{k,n}}\right)\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}- \mathds{1}{\{{\mathbf{q}}\in\mathcal{U}\}}z_{m}^{2}\mathbf{Attn}_{{\mathbf{p}} \to\mathcal{P}_{k,m}}\right.\right.$ $\displaystyle\left.\left.+\sum_{a\not=n}z_{a}^{2}\left(\mathbf{Attn}_{{\mathbf {p}}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right].$ $$

∎

Based on the above gradient update for $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}$, we further introduce the following auxiliary quantity, which will be useful in the later proof.

$$ $\displaystyle\Upsilon^{(t+1)}_{k,\mathbf{p}\to\mathbf{q}}\coloneqq\Upsilon^{(t )}_{k,\mathbf{p}\to\mathbf{q}}+\eta\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}, \quad\text{ with }\Upsilon^{(0)}_{k,\mathbf{p}\to\mathbf{q}}=0$ (B.2) $$

It is easy to verify that $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}=\sum_{k\in[K]}\Upsilon^{(t)}_{k,
\mathbf{p}\to\mathbf{q}}$.

### B.2 High-probability Event

We first introduce the following exponential bounds for the hypergeometric
distribution Hyper $(m,D,M)$. Hyper $(m,D,M)$ 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 a success or a failure.

###### Proposition B.5 ( [ 16 ] ) .

Suppose $S\sim$ Hyper $(m,D,M)$ with $1\leq m,D\leq M$. Define $\mu_{M}\coloneqq D/M$. Then for all $t>0$

$$ $P\left(|S-m\mu_{M}|>t\right)\leq 2\exp\left(-\frac{t^{2}}{4m\mu_{M}+2t}\right).$ $$

We then utilize this property to prove the high-probability set introduced in Section A.1.

###### Lemma B.6 .

For $k\in[K]$ $n\in[N]$, define

$$ $\displaystyle\mathcal{E}_{k,n}(\gamma,P)\coloneqq\{\mathsf{M}:|\mathcal{P}_{k, n}\cap\mathcal{U}|=\Theta(C_{n})\},$ (B.3) $$

we have

$$ $\displaystyle\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n})\geq 1-2\exp(-c_{n,1}C_ {n})$ (B.4) $$

where $c_{n,0}>0$ is some constant.

###### Proof.

Under the random masking strategy, given $k\in[K]$ and $n\in[N]$, $Y_{k,n}=|\mathcal{U}\cap\mathcal{P}_{k,n}|$ follows the hypergeometric distribution, i.e. $Y_{k,n}\sim\text{Hyper}((1-\gamma)P,C_{n},P)$. Then by tail bounds, for $t>0$, we have:

$$ $\displaystyle\mathbb{P}[|Y_{k,n}-(1-\gamma)C_{n}|>t]$ $\displaystyle\leq 2\exp(-\frac{t^{2}}{4(1-\gamma)C_{n}+2t})$ $$

Letting $t=\Theta(C_{n})$, we have

$$ $\displaystyle\mathbb{P}[Y_{k,n}=\Theta(C_{n})]\geq 1-2e^{-c_{n,1}C_{n}}.$ $$

∎

We further have the following fact, which will be useful for proving the property of loss objective in the next subsection.

###### Lemma B.7 .

For $k\in[K]$ and $n\in[N]$, we have

$$ $\displaystyle\mathbb{P}(|\mathcal{U}\cap\mathcal{P}_{k,n}|=0)\leq\exp(-c_{n,0} C_{n}).$ (B.5) $$

where $c_{n,0}>0$ is some constant.

###### Proof.

By the form of probability density for $\text{Hyper}((1-\gamma)P,C_{n},P)$, we have

$$ $\displaystyle\mathbb{P}(|\mathcal{U}\cap\mathcal{P}_{k,n}|=0)$ $\displaystyle=\frac{{C_{n}\choose 0}{(P-C_{n})\choose(1-\gamma)P}}{{P\choose(1 -\gamma)P}}$ $\displaystyle\leq\gamma^{C_{n}}=\exp(-c_{n,0}C_{n})).$ $$

∎

###### Proposition B.5 ( [ 16 ] ) .

###### Lemma B.6 .

###### Proof.

###### Lemma B.7 .

###### Proof.

### B.3 Properties of Loss Function

Recall the training and regional reconstruction loss we consider are given by:

$$ $\displaystyle\mathcal{L}(Q)$ $\displaystyle\coloneqq\frac{1}{2}\mathbb{E}\left[\sum_{{\mathbf{p}}\in\mathcal {P}}\operatorname{\mathds{1}}\{{\mathbf{p}}\in\mathcal{M}\}\left\|[F(\mathsf{M }(X);Q,E)]_{{\mathbf{p}}}-X_{{\mathbf{p}}}\right\|^{2}\right]$ (B.6) $\displaystyle\mathcal{L}_{\mathbf{p}}(Q)$ $\displaystyle=\frac{1}{2}\mathbb{E}\left[\operatorname{\mathds{1}}\{{\mathbf{p }}\in\mathcal{M}\}\left\|[F(\mathsf{M}(X),E)]_{{\mathbf{p}}}-X_{{\mathbf{p}}} \right\|^{2}\right]$ (B.7) $$

In this part, we will present several important lemmas for such a training objective. We first single out the following lemma, which connects the loss form with the attention score.

###### Lemma B.8 (Loss Calculation) .

The population loss $L(Q)$ can be decomposed into the following form:

$$ $\displaystyle\mathcal{L}(Q)$ $\displaystyle=\sum_{\mathbf{p}\in\mathcal{P}}{\mathcal{L}}_{\mathbf{p}}(Q), \text{ where }$ $\displaystyle\mathcal{L}_{\mathbf{p}}(Q)$ $\displaystyle=\frac{1}{2}\sum_{k=1}^{K}\mathbb{E}\left[\mathds{1}\{{\mathbf{p} }\in\mathcal{M},k_{X}=k\}\right.\cdot$ $\displaystyle\left.\quad\left(z_{a_{k,\mathbf{p}}}^{2}\left(1-\mathbf{Attn}^{( t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}\right)^{2}+\sum_{a\not=a _{k,\mathbf{p}}}z_{a}^{2}\left(\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P} _{k,a}}\right)^{2}\right)\right]$ $$

###### Proof.

$$ $\displaystyle\mathcal{L}_{\mathbf{p}}(Q)$ $\displaystyle=\frac{1}{2}\sum_{k=1}^{K}\mathbb{E}\left[\operatorname{\mathds{1 }}\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}\left\|[F(\mathsf{M}(X),E)]_{{\mathbf{ p}}}-X_{{\mathbf{p}}}\right\|^{2}\right]$ $\displaystyle=\frac{1}{2}\sum_{k=1}^{K}\mathbb{E}\left[\operatorname{\mathds{1 }}\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}\left\|\sum_{m\in[N]}\mathbf{Attn}_{{ \mathbf{p}}\to\mathcal{P}_{k,m}}z_{m}v_{k,m}-z_{a_{k,\mathbf{p}}}v_{k,a_{k, \mathbf{p}}}\right\|^{2}\right]$ $\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}\frac{1}{2}\sum_{k=1}^{K}\mathbb {E}\left[\operatorname{\mathds{1}}\{{\mathbf{p}}\in\mathcal{M},k_{X}=k\}\left( z_{a_{k,\mathbf{p}}}^{2}\left(1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a _{k,\mathbf{p}}}}\right)^{2}+\sum_{m\not=a_{k,\mathbf{p}}}z_{m}^{2}\left( \mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $$

where $(i)$ follows since the features are orthogonal.
∎

We then introduce some additional crucial notations for the loss objectives.

$$ $\displaystyle\mathcal{L}^{*}_{\mathbf{p}}=\min_{Q\in\mathbb{R}^{d\times d}} \mathcal{L}_{\mathbf{p}}(Q),$ (B.8a) $\displaystyle\mathcal{L}^{\text{low}}_{\mathbf{p}}=\frac{1}{2}(\sigma_{z}^{2}+ \frac{L^{2}}{N-1})\sum_{k\in[K]}\mathbb{P}\left(|\mathcal{U}\cap\mathcal{P}_{k ,z_{a_{k,\mathbf{p}}}}|=0\right)$ (B.8b) $\displaystyle\widetilde{\mathcal{L}}_{\mathbf{p}}(Q)=\sum_{k=1}^{K}\widetilde{ \mathcal{L}}_{k,\mathbf{p}}(Q),\quad\text{ where }$ $\displaystyle\widetilde{\mathcal{L}}_{k,\mathbf{p}}(Q)=\frac{1}{2}\mathbb{E} \left[\mathds{1}\{{\mathbf{p}}\in\mathcal{M},k_{X}=k,\mathsf{M}\in\mathcal{E}_ {k,z_{a_{k,\mathbf{p}}}}\}\right.\cdot$ $\displaystyle\left.\quad\left(z_{a_{k,\mathbf{p}}}^{2}\left(1-\mathbf{Attn}^{( t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}\right)^{2}+\sum_{a\not=a _{k,\mathbf{p}}}z_{a}^{2}\left(\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P} _{k,a}}\right)^{2}\right)\right]$ (B.8c) $$

Here $\sigma_{z}^{2}=\mathbb{E}[Z_{n}(X)^{2}]$. $\mathcal{L}_{\mathbf{p}}^{\star}$ denotes the minimum value of the population loss in (B.7), and $\mathcal{L}^{\text{low}}_{\mathbf{p}}$ represents the unavoidable errors for $\mathbf{p}\in\mathcal{P}$, given that all the patches in $\mathcal{P}_{k,a_{k,\mathbf{p}}}$ are masked. We will show that $\mathcal{L}^{\text{low}}_{\mathbf{p}}$ serves as a lower bound for $\mathcal{L}_{\mathbf{p}}^{\star}$, and demonstrate that the network trained with GD will attain nearly zero error compared to $\mathcal{L}^{\text{low}}_{\mathbf{p}}$. Our convergence will be established by the sub-optimality gap with respect to $\mathcal{L}^{\text{low}}_{\mathbf{p}}$, which necessarily implies the convergence to $\mathcal{L}_{\mathbf{p}}^{\star}$. (It also implies $\mathcal{L}_{\mathbf{p}}^{\star}-\mathcal{L}^{\text{low}}_{\mathbf{p}}$ is small.)

###### Lemma B.9 .

For ${\mathcal{L}}_{\mathbf{p}}^{\star}$ and $\mathcal{L}^{\text{low}}_{\mathbf{p}}$ defined in (B.8a) and (B.8b), respectively, we have $\mathcal{L}^{\text{low}}_{\mathbf{p}}\leq{\mathcal{L}}_{\mathbf{p}}^{\star}$ and they are both at the order of $\Theta\Big{(}\exp\Big{(}-\big{(}c_{1}P^{\kappa_{c}}+\operatorname{\mathds{1}}
\big{\{}1\not\in\cup_{k\in[K]}\{a_{k,\mathbf{p}}\}\big{\}}c_{2}P^{\kappa_{s}}
\big{)}\Big{)}\Big{)}$ where $c_{1},c_{2}>0$ are some constants.

###### Proof.

We first prove $\mathcal{L}^{\text{low}}_{\mathbf{p}}\leq{\mathcal{L}}_{\mathbf{p}}^{\star}$:

$$ $\displaystyle{\mathcal{L}}_{\mathbf{p}}^{\star}=\min_{Q\in\mathbb{R}^{d\times d}}$ $\displaystyle\frac{1}{2}\sum_{k=1}^{K}\mathbb{E}\left[\mathds{1}\{{\mathbf{p}} \in\mathcal{M},k_{X}=k\}\right.\cdot$ $\displaystyle\left.\quad\left(z_{a_{k,\mathbf{p}}}^{3}\left(1-\mathbf{Attn}^{( t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}\right)^{2}+\sum_{a\not=a _{k,\mathbf{p}}}z_{a}^{2}z_{a_{k,\mathbf{p}}}\left(\mathbf{Attn}^{(t)}_{{ \mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]$ $\displaystyle\geq\min_{Q\in\mathbb{R}^{d\times d}}$ $\displaystyle\frac{1}{2}\sum_{k=1}^{K}\mathbb{E}\left[\mathds{1}\{{\mathbf{p}} \in\mathcal{M},k_{X}=k\}\operatorname{\mathds{1}}\{|\mathcal{U}\cap\mathcal{P} _{k,a_{\mathbf{p},k}}|=0\}\right.\cdot$ $\displaystyle\left.\quad\left(z_{a_{k,\mathbf{p}}}^{3}\left(1-\mathbf{Attn}^{( t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}\right)^{2}+\sum_{a\not=a _{k,\mathbf{p}}}z_{a}^{2}z_{a_{k,\mathbf{p}}}\left(\mathbf{Attn}^{(t)}_{{ \mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]$ $$

Notice that when all patches in $\mathcal{P}_{k,a_{k,\mathbf{p}}}$ are masked, $\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}=0$. Moreover,

$$ $\sum\limits_{m\not=a_{k,\mathbf{p}}}z_{m}^{2}\mathbf{Attn}^{(t)}_{{\mathbf{p}} \to\mathcal{P}_{k,m}}\geq\frac{L^{2}}{N-1}$ $$

by Cauchy–Schwarz inequality. Thus

$$ $\displaystyle{\mathcal{L}}_{\mathbf{p}}^{\star}$ $\displaystyle\geq\frac{1}{2}\sum_{k=1}^{K}(\sigma_{z}^{2}+\frac{L^{2}}{N-1}) \mathbb{P}\left(|\mathcal{U}\cap\mathcal{P}_{k,a_{k,\mathbf{p}}}|=0\right)= \mathcal{L}^{\text{low}}_{\mathbf{p}}.$ $$

$\mathcal{L}^{\text{low}}_{\mathbf{p}}=\Theta\Big{(}\exp\Big{(}-\big{(}c_{1}P^{
\kappa_{c}}+\operatorname{\mathds{1}}\big{\{}1\not\in\cup_{k\in[K]}\{a_{k,
\mathbf{p}}\}\big{\}}c_{2}P^{\kappa_{s}}\big{)}\Big{)}\Big{)}$ immediately comes from Lemma B.7.
Furthermore, we only need to show ${\mathcal{L}}_{\mathbf{p}}^{\star}=O\Big{(}\exp\Big{(}-\big{(}c_{1}P^{\kappa_{
c}}+\operatorname{\mathds{1}}\big{\{}1\not\in\cup_{k\in[K]}\{a_{k,\mathbf{p}}
\}\big{\}}c_{2}P^{\kappa_{s}}\big{)}\Big{)}\Big{)}$. This can be directly obtained by choosing $Q=\sigma I_{d}$ for some sufficiently large $\sigma$ and hence omitted here.
∎

###### Lemma B.10 .

Given $\mathbf{p}\in\mathcal{P}$, for any $Q$, we have

$$ $\displaystyle\widetilde{\mathcal{L}}_{\mathbf{p}}(Q)\leq{L}_{\mathbf{p}}(Q)- \mathcal{L}^{\text{low}}_{\mathbf{p}}\leq\widetilde{\mathcal{L}}_{\mathbf{p}}( Q)+O\Big{(}\exp\Big{(}-\big{(}c_{3}P^{\kappa_{c}}+\operatorname{\mathds{1}} \big{\{}1\not\in\cup_{k\in[K]}\{a_{k,\mathbf{p}}\}\big{\}}c_{4}P^{\kappa_{s}} \big{)}\Big{)}\Big{)}.$ $$

where $c_{3},c_{4}>0$ are some constants.

###### Proof.

The lower bound is directly obtained by the definition and thus we only prove the upper bound.

$$ $\displaystyle{L}_{\mathbf{p}}(Q)-\widetilde{\mathcal{L}}_{\mathbf{p}}(Q)$ $\displaystyle=\frac{1}{2}\sum_{k=1}^{K}\mathbb{E}\left[\mathds{1}\{{\mathbf{p} }\in\mathcal{M},k_{X}=k,\mathsf{M}\in\mathcal{E}^{c}_{k,z_{a_{k,\mathbf{p}}}} \}\cdot\Big{(}z_{a_{k,\mathbf{p}}}^{2}\left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p} }\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}\right)^{2}+\sum_{a\not=a_{k,\mathbf{p}}} z_{a}^{2}\left(\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,a}}\right)^{ 2}\Big{)}\right]$ $\displaystyle\leq\sum_{k=1}^{K}U^{2}\mathbb{P}(\mathsf{M}\in\mathcal{E}^{c}_{k ,z_{a_{k,\mathbf{p}}}})$ $\displaystyle\leq O\Big{(}\exp\Big{(}-\big{(}c_{3}P^{\kappa_{c}}+\operatorname {\mathds{1}}\big{\{}1\not\in\cup_{k\in[K]}\{a_{k,\mathbf{p}}\}\big{\}}c_{4}P^{ \kappa_{s}}\big{)}\Big{)}\Big{)}.$ $$

where the last inequality follows from Lemma B.6.

∎

###### Lemma B.8 (Loss Calculation) .

###### Proof.

###### Lemma B.9 .

###### Proof.

###### Lemma B.10 .

###### Proof.

## Appendix C Overall Induction Hypotheses and Proof Plan

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.

### C.1 Positive Information Gap

We first state our induction hypothesis for the case that the information gap $\Delta$ is positive.

###### Induction Hypothesis C.1 .

For $t\leq T$, given $\mathbf{p},\mathbf{q}\in\mathcal{P}$, for $k\in[K]$, the following holds

- a.
$\Phi^{(t)}_{\bm{p}\to v_{k,a_{k,\mathbf{p}}}}$ is monotonically increasing, and $\Phi^{(t)}_{\bm{p}\to v_{k,a_{k,\mathbf{p}}}}\in[0,\widetilde{O}(1)]$;
- b.
if $a_{k,\mathbf{p}}\not=1$, then $\Phi^{(t)}_{\bm{p}\to v_{k,1}}$ is monotonically decreasing and $\Phi^{(t)}_{\bm{p}\to v_{k,1}}\in[-\widetilde{O}(1),0]$;
- c.
$|\Phi^{(t)}_{\bm{p}\to v_{k,m}}|=\widetilde{O}(\frac{1}{P^{1-\kappa_{s}}})$ for $m\notin\{1\}\cup\{a_{k,\mathbf{p}}\}$;
- d.
for $\mathbf{q}\not=\mathbf{p}$, $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}=\widetilde{O}(\frac{1}{P^{\kappa_{s}}})$;
- e.
$\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{p}}=\widetilde{O}(\frac{1}{P})$.

###### Induction Hypothesis C.1 .

### C.2 Negative Information Gap

Now we turn to the case that $\Delta\leq-\Omega(1)$.

###### Induction Hypothesis C.2 .

For $t\leq T$, given $\mathbf{p},\mathbf{q}\in\mathcal{P}$, for $k\in[K]$, the following holds

- a.
$\Phi^{(t)}_{\bm{p}\to v_{k,a_{k,\mathbf{p}}}}$ is monotonically increasing, and $\Phi^{(t)}_{\bm{p}\to v_{k,a_{k,\mathbf{p}}}}\in[0,\widetilde{O}(1)]$;
- b.
if $a_{k,\mathbf{p}}\not=1$, then $\Phi^{(t)}_{\bm{p}\to v_{k,1}}$ is monotonically decreasing and $\Phi^{(t)}_{\bm{p}\to v_{k,1}}\in[-\widetilde{O}(\frac{1}{P^{-\Delta}}),0]$;
- c.
$|\Phi^{(t)}_{\bm{p}\to v_{k,m}}|=\widetilde{O}(\frac{1}{P^{1-\kappa_{s}}})$ for $m\notin\{1\}\cup\{a_{k,\mathbf{p}}\}$;
- d.
for $\mathbf{q}\not=\mathbf{p}$, $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}=\widetilde{O}(\frac{1}{P^{\kappa_{s}}})$;
- e.
$\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{p}}=\widetilde{O}(\frac{1}{P})$.

###### Induction Hypothesis C.2 .

### C.3 Proof Outline

In both settings, we can classify the process through which transformers learn the feature attention correlation $\Phi^{(t)}_{\bm{p}\to v_{k,a_{k,\mathbf{p}}}}$ into two distinct scenarios. These scenarios hinge on the spatial relation of the area $\bm{p}$ within the context of the $k$-th partition ${\mathcal{D}_{k}}$, specifically, whether $\bm{p}$ is located in the global area of the $k$-th cluster, i.e. whether $a_{k,\bm{p}}=1$. The learning dynamics exhibit different behaviors of learning the local FP correlation in the local area with different $\Delta$, while the behaviors for features located in the global area are very similar, unaffected by the value of $\Delta$. Therefore, through Appendices D, E and F, we delve into the learning phases and provide technical proofs for the local area with $\Delta\geq\Omega(1)$, local area with $\Delta\leq-\Omega(1)$ and the global area respectively. Finally, we will put this analysis together to prove that the C.1 (resp. C.2) holds during the entire training process, thereby validating the main theorems in Appendix G.

## Appendix D Analysis for the Local Area with Positive Information Gap

In this section, we focus on a specific patch $\mathbf{p}\in\mathcal{P}$ with the $k$-th cluster for $k\in[K]$, and present the analysis for the case that $X_{\mathbf{p}}$ is located in the local area for the $k$-th cluster, i.e. $a_{k,\bm{p}}>1$. We will analyze the case that $\Delta\geq\Omega(1)$. Throughout this section, we denote $a_{k,\bm{p}}=n$ for simplicity. We will analyze the convergence of the training process via two phases of dynamics.
At the beginning of each phase, we will establish an induction hypothesis, which we expect to remain valid throughout that phase. Subsequently, we will analyze the dynamics under such a hypothesis within the phase, aiming to provide proof of the hypothesis by the end of the phase.

### D.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\leq T_{1}$, where

$$ $T_{1}\triangleq\max\left\{t:\Phi_{\mathbf{p}\to v_{k,n}}^{(t)}\geq-\frac{1}{U} \left(\frac{\Delta}{2}-0.01\right)\log(P)\right\}.$ $$

We state the following induction hypotheses, which will hold throughout this period:

###### Induction Hypothesis D.1 .

For each $0\leq t\leq T_{1}$, $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, the following holds:

- a.
$\Phi^{(t)}_{\bm{p}\to v_{k,n}}$ is monotonically increasing, and $\Phi^{(t)}_{\bm{p}\to v_{k,n}}\in[0,O\left(\frac{\left(\frac{\Delta}{2}-0.01
\right)\log(P)}{P^{0.02}}\right)]$;
- b.
$\Phi^{(t)}_{\bm{p}\to v_{k,1}}$ is monotonically decreasing and $\Phi^{(t)}_{\bm{p}\to v_{k,1}}\in[-\frac{1}{U}\left(\frac{\Delta}{2}-0.01
\right)\log(P),0]$;
- c.
$|\Phi^{(t)}_{\bm{p}\to v_{k,m}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to v_{k,n}}-
\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}$ for $m\not=1,n$;
- d.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to
v
_{k,n}}}{C_{n}}\Big{)}$ for $a_{k,\mathbf{q}}=n$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$;
- e.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{|\Phi^{(t)}_{\bm{p}
\to v_{k,1}}|}{C_{1}}\Big{)}+O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to v_{k,n}}-\Phi
^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}=1$;
- f.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1,n$.

###### Induction Hypothesis D.1 .

#### D.1.1 Property of Attention Scores

We first introduce several properties of the attention score if C.1 and D.1 hold.

###### Lemma D.1 .

For $n>1$, if C.1 and D.1 hold at iteration $t\leq T_{1}$, then the following holds

- 1.
$1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\mathbf{Attn}^{(t)}_{
\mathbf{p}\to\mathcal{P}_{k,1}}\geq\Omega(1)$;
- 2.
If $\mathsf{M}\in\mathcal{E}_{k,n}$,
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}=\Theta\Big{(}\frac{1}{P^{
1-\kappa_{s}}}\Big{)}$
;
- 3.
Moreover, if $\mathsf{M}\in\mathcal{E}_{k,1}$, we have $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=\Omega\Big{(}\frac{1}{P^{
\frac{1-\kappa_{s}}{2}-0.01}}\Big{)}$;
- 4.
For ${\mathbf{q}}\in\mathcal{M}\cap(\mathcal{P}_{k,n}\cup\mathcal{P}_{k,1})$, $\mathbf{attn}_{\mathbf{p}\to{\mathbf{q}}}^{(t)}=O\Big{(}\frac{1-\mathbf{Attn}^
{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\mathbf{Attn}^{(t)}_{\mathbf{p}\to
\mathcal{P}_{k,n}}}{P}\Big{)}$.

###### Lemma D.2 .

For $n>1$, if C.1 and D.1 hold at iteration $t\leq T_{1}$, then for $m\not=n,1$, the following holds:

- 1.
For any ${\mathbf{q}}\in\mathcal{P}_{k,m}$, $\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{q}}}\leq O\Big{(}\frac{1-\mathbf{
Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\mathbf{Attn}_{\mathbf{p}\to
\mathcal{P}_{k,n}}^{(t)}\Big{)}}{P})$.
- 2.
Moreover, $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\leq O\Big{(}\frac{1-
\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\mathbf{Attn}^{(t)}_{
\mathbf{p}\to\mathcal{P}_{k,n}}}{N}\Big{)}.$

The above properties can be easily verified through direct calculations by using the definition in (2.7) and conditions in D.1, which are omitted here for brevity.

###### Lemma D.1 .

###### Lemma D.2 .

#### D.1.2 Bounding the Gradient Updates for FP Correlations

###### Lemma D.3 .

For $n>1$, if C.1 and D.1 hold at iteration $0\leq t\leq T_{1}$, then $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}\geq 0$ and satisfies:

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}=\Theta\Big{(}\frac{C_{n}}{P} \Big{)}=\Theta\Big{(}\frac{1}{P^{1-\kappa_{s}}}\Big{)}.$ $$

###### Proof.

By Lemma A.2, we have

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in\mathcal{M}\} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_{n}^{3}\left( 1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not= n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}} \right)^{2}\right)\right]$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}\cap\mathbf {p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot \left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} \right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}^{c} \cap\mathbf{p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k ,n}}\cdot\left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_ {k,n}}\right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle{\leq}\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n})$ $\displaystyle\quad\cdot\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in \mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_ {n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2} +\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{ P}_{k,m}}\right)^{2}\right)\Big{|}\mathcal{E}_{k,n}\right]$ $\displaystyle\quad+O(1)\cdot\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n}^{c})$ $\displaystyle\leq O\Big{(}\frac{C_{n}}{P})+O(\exp(-c_{n,1}C_{n})\Big{)}$ $\displaystyle\leq O\Big{(}\frac{C_{n}}{P}\Big{)},$ $$

where the second inequality invokes Lemma D.1 and Lemma B.6, and the last inequality is due to $exp(-c_{n,1}C_{n})\ll\frac{C_{n}}{P}$.
Similarly, we can show that $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}\geq\Omega(\frac{C_{n}}{P})$.

∎

###### Lemma D.4 .

For $n>1$, if C.1 and D.1 hold at iteration $0\leq t\leq T_{1}$, then $\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}<0$ and satisfies

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\geq\Omega\Big{(}\frac{1}{P ^{2(\frac{1-\kappa_{s}}{2}-0.01)}}\Big{)}=\Omega\Big{(}\frac{1}{P^{0.98-\kappa _{s}}}\Big{)}.$ $$

###### Proof.

We first single out the following fact:

$$ $\displaystyle-z_{1}z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal {P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-z_{1}^{3 }\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\right)\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}+\sum_{a\not=1,n}z_{a}^{2}z_{1} \left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}$ $\displaystyle\leq z_{1}\left(\max_{a\not=1,n}{z_{a}^{2}\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,a}}}-z_{n}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}-z_{1}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1 }}\right)(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\mathbf{Attn} ^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})$ $\displaystyle=-z_{1}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}- \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})\left(z_{n}^{2}\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}+z_{1}^{2}\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,1}}-\max_{a\not=1,n}{z_{a}^{2}\mathbf{Attn}^{(t)}_ {\mathbf{p}\to\mathcal{P}_{k,a}}}\right).$ (D.1) $$

Therefore, by Lemma A.1, we have

$$ $\displaystyle\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}$ $\displaystyle\leq\mathbb{E}\Bigg{[}\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\cdot$ $\displaystyle\quad\left(-z_{1}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P} _{k,n}}-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})\left(z_{n}^{2} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}+z_{1}^{2}\mathbf{Attn}^{( t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\max_{a\not=1,n}{z_{a}^{2}\mathbf{Attn}^{ (t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}}\right)\right)\Bigg{]}$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}^{c} \cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{ k,1}}\cdot\sum_{a\not=1,n}z_{1}^{2}z_{a}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,a}}\right)^{2}\right]$ $\displaystyle\leq\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,1})\cdot\Big{(}-\big{( }\Omega(1)\cdot\Omega(\frac{1}{P^{2\times(\frac{1-\kappa_{s}}{2}-0.01)}})\big{ )}\Big{)}+O(1)\cdot\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,1}^{c})$ $\displaystyle\leq-\Omega\Big{(}\frac{1}{P^{2\times(\frac{1-\kappa_{s}}{2}-0.01 )}}\Big{)}=-\Omega\Big{(}\frac{1}{P^{0.98-\kappa_{s}}}\Big{)}$ $$

where the second inequality invokes Lemma D.1 and the last inequality comes from Lemma B.6.
∎

###### Lemma D.5 .

At each iteration $t\leq T_{1}$, if C.1 and D.1 hold, then for any $m>1$ with $m\not=n$, the following holds

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}|\leq O\Big{(}\frac{\alpha_{ \mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{N}\Big{)}=O \Big{(}\frac{\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}-\alpha_{\mathbf{p}\to v_{k,1 }}^{(t)}}{P^{1-\kappa_{s}}}\Big{)}.$ $$

###### Proof.

By Lemma A.1, for $m\not=n$, we have

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}$ $\displaystyle\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in\mathcal{M} \}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\cdot\left(\sum_{a\not= m,n}z_{a}^{2}z_{m}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}} \right)^{2}\right)\right]$ (D.2) $\displaystyle-\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}$ $\displaystyle\leq\mathbb{E}\Bigg{[}\mathds{1}{\{k_{X}=k,\mathbf{p}\in\mathcal{ M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\cdot\left(z_{m}z_{n} ^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)\mathbf {Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right.$ $\displaystyle\qquad\left.+z_{m}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,m}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}} \right)\Bigg{]}$ (D.3) $$

For (D.2), we have

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}$ $\displaystyle\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathcal{E}_{k,n}\cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p }\to\mathcal{P}_{k,m}}\cdot\left(\sum_{a\not=m,n}z_{a}^{2}z_{m}\left(\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,(\mathcal{E}_{k,1}\cap \mathcal{E}_{k,n})^{c}\cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,m}}\cdot\left(\sum_{a\not=m,n}z_{a}^{2}z_{m}\left( \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]$ $\displaystyle\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathcal{E}_{k,n}\cap\mathbf{p}\in\mathcal{M}\}}O\Bigg{(}\frac{1-\mathbf{Attn} ^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}}{N}\Bigg{)}\right.$ $\displaystyle\qquad\left.\cdot\left(z_{1}^{2}z_{m}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,1}}\right)^{2}+O\Big{(}\frac{1}{N}\Big{)}\right) \right]+O(1)\cdot\mathbb{P}(\mathsf{M}\in(\mathcal{E}_{k,1}\cap\mathcal{E}_{k, n})^{c})$ $\displaystyle\leq O\Big{(}\frac{|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{N}\Big {)}+O(1)\cdot\mathbb{P}(\mathsf{M}\in(\mathcal{E}_{k,1}\cap\mathcal{E}_{k,n})^ {c})$ $\displaystyle\leq O\Big{(}\frac{|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{P^{1- \kappa_{s}}}\Big{)}$ $$

where the second inequality is due to Lemma D.2, the last inequality follows from Lemma D.4 and Lemma B.6.

On the other hand, for (D.3), we can use the similar argument by invoking Lemma D.2 and Lemma D.3, and thus obtain

$$ $\displaystyle-\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}$ $\displaystyle\leq O\Big{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P^{1- \kappa_{s}}}\Big{)}.$ $$

Putting them together, we have

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}|\leq O\Big{(}\frac{\alpha^{ (t)}_{\mathbf{p}\to v_{k,n}}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P^{1-\kappa _{s}}}\Big{)}.$ $$

∎

###### Lemma D.3 .

###### Proof.

###### Lemma D.4 .

###### Proof.

###### Lemma D.5 .

###### Proof.

#### D.1.3 Bounding the Gradient Updates for Positional Correlations

###### Lemma D.6 .

For $n>1$, if C.1 and D.1 hold at iteration $0\leq t\leq T_{1}$, then for $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$ and $a_{k,\mathbf{q}}=n$, we have $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}\geq 0$ and satisfies:

$$ $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta\Big{(}\frac{\alpha _{\mathbf{p}\to v_{k,n}}^{(t)}}{C_{n}}\Big{)}.$ $$

Furthermore, we have $|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v
_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}$.

###### Proof.

By Lemma A.2, for ${\mathbf{q}}\in\mathcal{P}_{k,n}$ with ${\mathbf{q}}\not=\mathbf{p}$,
we have

$$ $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}$ $\displaystyle=$ $\displaystyle\underbrace{\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M},{\mathbf{q}}\in\mathcal{U}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{ \mathbf{q}}}\cdot\left(z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not=n}z_{m}^{2}\left(\mathbf{Attn}^{(t)} _{\mathbf{p}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]}_{H_{1}}$ $\displaystyle+\underbrace{\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M},{\mathbf{q}}\in\mathcal{M}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{ \mathbf{q}}}\cdot\left(-z_{n}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P} _{k,n}}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right) \right)\right]}_{H_{2}}$ $\displaystyle+\underbrace{\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M},{\mathbf{q}}\in\mathcal{M}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{ \mathbf{q}}}\cdot\left(\sum_{m\not=n}z_{m}^{2}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]}_{H_{3}}.$ $$

Firstly, for $H_{1}$, notice that

$$ $\displaystyle(C_{n}-1)H_{1}$ $\displaystyle=\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in\mathcal{M}\}} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_{n}^{2}\left( 1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not= n}z_{m}^{2}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\right)^{ 2}\right)\right]$ $\displaystyle=\Theta(\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}).$ $$

For $H_{2}$, since $\mathbf{p},{\mathbf{q}}\in\mathcal{M}$, by Lemma D.1, we can upper bound $\mathbf{attn}^{(t)}_{\mathbf{p}\to{\mathbf{q}}}$ by $O\Big{(}\frac{1}{P}\Big{)}$, thus

$$ $\displaystyle-H_{2}\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M}\}}O\Big{(}\frac{1}{P}\Big{)}\cdot\left(z_{n}^{2}\mathbf{Attn}^{(t) }_{\mathbf{p}\to\mathcal{P}_{k,n}}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}\right)\right)\right]\leq O\Big{(}\frac{\alpha_{\mathbf{p} \to v_{k,n}}^{(t)}}{P}\Big{)}.$ $$

Further notice that $H_{3}$ can be upper bounded by $O(H_{1})$, putting it together, we have

$$ $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta\Big{(}\frac{\alpha _{\mathbf{p}\to v_{k,n}}^{(t)}}{C_{n}}\Big{)}.$ $$

Turn to $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{p}}$, when ${\mathbf{q}}=\mathbf{p}$,

$$ $\displaystyle\beta^{(t)}_{n}=\underbrace{\mathbb{E}\left[\mathds{1}{\{k_{X}=k, \mathbf{p}\in\mathcal{M}\}}\mathbf{attn}^{(t)}_{{\mathbf{p}}\to{\mathbf{p}}} \cdot\left(-z_{n}^{2}\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}} \left(1-\mathbf{Attn}^{(t)}_{{\mathbf{p}}\to\mathcal{P}_{k,n}}\right)\right) \right]}_{J_{2}}$ $\displaystyle+\underbrace{\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to\mathbf{p}}\cdot\left(\sum_{m \not=n}z_{m}^{2}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}} \right)^{2}\right)\right]}_{J_{3}}.$ $$

We can bound $J_{2}$ in a similar way as $H_{2}$. Thus, we only focus on further bounding $J_{3}$:

$$ $\displaystyle J_{3}$ $\displaystyle\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in\mathcal{M} \}}O(\frac{1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\mathbf{Attn }^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}}{P})\cdot\left(\sum_{m\not=n}z_{m}^{2 }\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle\leq O\Bigg{(}\frac{|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{P} \Bigg{)}.$ $$

where the first inequality holds by invoking Lemma D.1 and the last inequality follows similar arguments as analysis for (D.2).
∎

###### Lemma D.7 .

For $n>1$, if C.1 and D.1 hold at iteration $0\leq t\leq T_{1}$, then for $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$ and $a_{k,\mathbf{q}}=1$, we have $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}$ satisfies:

$$ $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Bigg{(}\frac{|\alpha_ {\mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{P}\Bigg{) }+O\Bigg{(}\frac{|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{C_{1}}\Bigg{)}.$ $$

###### Proof.

By Lemma A.2, for ${\mathbf{q}}\in\mathcal{P}_{k,1}$, we have

$$ $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=$ $\displaystyle-\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in\mathcal{M},{ \mathbf{q}}\in\mathcal{U}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{\mathbf{q}}} \cdot\right.$ $\displaystyle\left.\left(z_{1}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P }_{k,1}}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})+z_{n}^{2} \left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\sum_{a\not=1,n}z_{a}^{2}\left( \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]$ (D.4) $\displaystyle\underbrace{-\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M},{\mathbf{q}}\in\mathcal{M}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{ \mathbf{q}}}\cdot\left(z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} \right)\right]}_{G_{2}}$ $\displaystyle\underbrace{+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M},{\mathbf{q}}\in\mathcal{M}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{ \mathbf{q}}}\cdot\left(\sum_{a\not=n}z_{a}^{2}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]}_{G_{3}}.$ $$

For (D.4) denoted as $G_{1}$, following the direct calculations, we have

$$ $\displaystyle-(C_{1}-1)G_{1}$ $\displaystyle=\Theta(\alpha_{\mathbf{p}\to v_{k,1}}^{(t)})$ $$

We can further bound $G_{2}$ and $G_{3}$ in a similar way as $H_{2}$ and $H_{3}$ in Lemma D.6 and thus obtain

$$ $\displaystyle-G_{2}$ $\displaystyle\leq O\Big{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P}\Big{)},$ $\displaystyle G_{3}$ $\displaystyle\leq O\Bigg{(}\frac{|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{P} \Bigg{)}.$ $$

which completes the proof.
∎

###### Lemma D.8 .

For $n>1$, if C.1 and D.1 hold at iteration $0\leq t\leq T_{1}$, then for $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$ and $n\not=a_{k,\mathbf{q}}$, $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}$ satisfies:

$$ $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{ \mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}.$ $$

###### Proof.

By Lemma A.2, for ${\mathbf{q}}\in\mathcal{P}_{k,m}$, we have

$$ $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=$ $\displaystyle-\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in\mathcal{M},{ \mathbf{q}}\in\mathcal{U}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{\mathbf{q}}} \cdot\right.$ $\displaystyle\left.\left(z_{m}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P }_{k,m}}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}})+z_{n}^{2} \left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\sum_{a\not=n,m}z_{a}^{2}\left( \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]$ (D.5) $\displaystyle\underbrace{-\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M},{\mathbf{q}}\in\mathcal{M}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{ \mathbf{q}}}\cdot\left(z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} \right)\right]}_{I_{2}}$ $\displaystyle\underbrace{+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathbf{p}\in \mathcal{M},{\mathbf{q}}\in\mathcal{M}\}}\mathbf{attn}^{(t)}_{\mathbf{p}\to{ \mathbf{q}}}\cdot\left(\sum_{a\not=n}z_{a}^{2}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}\right)\right]}_{I_{3}}.$ $$

(D.5) can be upper bounded by $O\Big{(}\frac{|\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}|}{C_{m}}\Big{)}=O\Big{(}
\frac{|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t
)}|}{NC_{m}}\Big{)}=O\Big{(}\frac{|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}-\alpha
_{\mathbf{p}\to v_{k,1}}^{(t)}|}{P}\Big{)}$, where the first equality holds by invoking Lemma D.5. $I_{2}$ and $I_{3}$ can be bounded similarly as $G_{2}$ and $G_{3}$, which is omitted here.
∎

###### Lemma D.6 .

###### Proof.

###### Lemma D.7 .

###### Proof.

###### Lemma D.8 .

###### Proof.

#### D.1.4 At the end of Phase I, Stage 1

###### Lemma D.9 .

For $n>1$, if C.1 and D.1 hold for all $0\leq t\leq T_{1}=O\Big{(}\frac{\log(P)P^{0.98-\kappa_{s}}}{\eta}\Big{)}$, At iteration $t=T_{1}+1$, we have

- a.
$\Phi^{(T_{1}+1)}_{\mathbf{p}\to v_{k,1}}\leq-\frac{1}{U}\left(\frac{\Delta}{2}
-0.01\right)\log(P)$;
- b.
$\mathbf{Attn}^{(T_{1}+1)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=O\Big{(}\frac{1}{P^
{(1-\kappa_{c})+\frac{L}{U}(\frac{\Delta}{2}-0.01)}}\Big{)}$.

###### Proof.

By comparing Lemma D.3 and Lemma D.4, we have $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\gg\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}$. Then the existence of $T_{1,k}=O\Big{(}\frac{\log(P)P^{0.98-\kappa_{s}}}{\eta}\Big{)}$ directly follows from Lemma D.4.
∎

###### Lemma D.9 .

###### Proof.

### D.2 Phase I, Stage 2

During stage 1, $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ significantly decreases to decouple the FP correlations with the global feature, resulting in a decrease in $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(t)}$, while other $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}$ with $m>1$ remain approximately at the order of $O\left(\frac{1}{P^{1-\kappa_{s}}}\right)$ ($\Theta\left(\frac{1}{P^{1-\kappa_{s}}}\right)$). By the end of phase I, $(\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(t)})^{2}$ decreases to $O(\frac{1}{P^{1.96-2\kappa_{s}}})$, leading to a decrease in $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|$ as it approaches towards $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}$. At this point, stage 2 begins. Shortly after entering this phase, the prior dominant role of the decrease of $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ in learning dynamics diminishes as $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|$ reaches the same order of magnitude as $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}$.

We define stage 2 of phase I as all iterations $T_{1}<t\leq\widetilde{T}_{1}$, where

$$ $\widetilde{T}_{1}\triangleq\max\left\{t>T_{1}:\Phi_{\mathbf{p}\to v_{k,n}}^{(t )}-\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}\leq\left(\frac{\Delta}{2L}+\frac{0.01}{L }+\frac{c_{1}^{*}(1-\kappa_{s})}{U}\right)\log(P)\right\}.$ $$

for some small constant $c^{*}_{1}>0$.

For computational convenience, we make the following assumptions for $\kappa_{c}$ and $\kappa_{s}$, which can be easily relaxed with the cost of additional calculations.

$$ $\displaystyle\frac{\Delta}{2}\Big{(}\frac{1}{L}-\frac{1}{U}\Big{)}+\frac{0.01} {L}+\frac{0.01}{U}$ $\displaystyle\leq\frac{c_{0}^{*}(1-\kappa_{s})}{U}$ (D.6a) $\displaystyle(1-\frac{c_{1}^{*}L}{U})(1-\kappa_{s})$ $\displaystyle\leq(1-\kappa_{c})+\frac{U}{L}(\frac{\Delta}{2}+0.01)$ (D.6b) $$

Here $c_{0}^{*}$ is some small.
We state the following induction hypotheses, which will hold throughout this period:

###### Induction Hypothesis D.2 .

For each $T_{1}<t\leq\widetilde{T}_{1}$, $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, the following holds:

- a.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}$ is monotonically increasing, and $\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}\in[0,\frac{c^{*}_{0}+c^{*}_{1}}{U}\log(P)]$;
- b.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ is monotonically decreasing and $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}\in[-\frac{1}{L}\left(\frac{\Delta}{2}+0.01
\right)\log(P),-\frac{1}{U}\left(\frac{\Delta}{2}-0.01\right)\log(P)]$;
- c.
$|\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p}\to v
_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}$ for $m\not=1,n$;
- d.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to
v
_{k,n}}}{C_{n}}\Big{)}$ for $a_{k,\mathbf{q}}=n$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$;
- e.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{|\Phi^{(t)}_{\bm{p}
\to v_{k,1}}|}{C_{1}}\Big{)}+O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to v_{k,n}}-\Phi
^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}=1$.;
- f.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1,n$.

###### Induction Hypothesis D.2 .

#### D.2.1 Property of Attention Scores

We first single out several properties of attention scores
that will be used for the proof of D.2.

###### Lemma D.10 .

if C.1 and D.2 hold at iteration $T_{1}+1\leq t\leq\widetilde{T}_{1}$, then the following holds

- 1.
$1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\mathbf{Attn}^{(t)}_{
\mathbf{p}\to\mathcal{P}_{k,1}}\geq\Omega(1)$;
- 2.
if $\mathsf{M}\in\mathcal{E}_{k,n}$,
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\in\Big{[}\Omega\big{(}
\frac{1}{P^{1-\kappa_{s}}}\big{)},O\big{(}\frac{1}{P^{(1-c_{1}^{*}-c_{0}^{*})(
1-\kappa_{s})}}\big{)}\Big{]}$
;
- 3.
Moreover, $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=O\Big{(}\frac{1}{P^{(1-
\kappa_{c})+\frac{L}{U}(\frac{\Delta}{2}-0.01)}}\Big{)}$; if $\mathsf{M}\in\mathcal{E}_{k,1}$, we have $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=\Omega\Big{(}\frac{1}{P^{
(1-\kappa_{c})+\frac{U}{L}(\frac{\Delta}{2}+0.01)}}\Big{)}$;
- 4.
for ${\mathbf{q}}\in\mathcal{M}\cap(\mathcal{P}_{k,n}\cup\mathcal{P}_{k,1})$, $\mathbf{attn}_{{\mathbf{p}}\to{\mathbf{q}}}^{(t)}=O\Big{(}\frac{1-\mathbf{Attn
}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\mathbf{Attn}^{(t)}_{\mathbf{p}\to
\mathcal{P}_{k,1}}}{P}\Big{)}$.

###### Lemma D.11 .

if C.1 and D.2 hold at iteration $T_{1}+1\leq t\leq\widetilde{T}_{1}$, then for $m\not=n$, the following holds:

- 1.
for any ${\mathbf{q}}\in\mathcal{P}_{k,m}$, $\mathbf{attn}^{(t)}_{\mathbf{p}\to{\mathbf{q}}}\leq O\Big{(}\frac{1-\mathbf{
Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\mathbf{Attn}^{(t)}_{\mathbf{p}
\to\mathcal{P}_{k,1}}}{P}\Big{)}$;
- 2.
Moreover, $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\leq O\Big{(}\frac{1-
\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\mathbf{Attn}^{(t)}_{
\mathbf{p}\to\mathcal{P}_{k,n}}}{N}\Big{)}.$

###### Lemma D.10 .

###### Lemma D.11 .

#### D.2.2 Bounding the Gradient Updates of FP Correlations

###### Lemma D.12 .

For $n>1$, if C.1 and D.2 hold at iteration $T_{1}+1\leq t\leq\widetilde{T}_{1}$, then $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}\geq 0$ and satisfies:

$$ $\displaystyle\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}=\Omega\Big{(}\frac{1}{P^{1- \kappa_{s}}}\Big{)}.$ $$

###### Proof.

By Lemma A.2, we have

$$ $\displaystyle\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in\mathcal{M}\} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_{n}^{3}\left( 1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not= n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}} \right)^{2}\right)\right]$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}\cap\mathbf {p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot \left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} \right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}^{c} \cap\mathbf{p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k ,n}}\cdot\left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_ {k,n}}\right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle{\gtrsim}\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n})$ $\displaystyle\quad\cdot\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in \mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_ {n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2} +\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{ P}_{k,m}}\right)^{2}\right)\Big{|}\mathcal{E}_{k,n}\right]$ $\displaystyle\geq\Omega\Big{(}\frac{C_{n}}{P}\Big{)}$ $$

where the last inequality invokes Lemma D.10.

∎

###### Lemma D.13 .

For $n>1$, if C.1 and D.2 hold at iteration $T_{1}+1\leq t\leq\widetilde{T}_{1}$, then $\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}<0$ and satisfies

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\geq\Omega\Big{(}\frac{1}{P ^{2(1-\kappa_{c})+\frac{U}{L}(\Delta+0.02)}}\Big{)}.$ $$

###### Proof.

Following (D.1), we have

$$ $\displaystyle-z_{1}z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal {P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-z_{1}^{3 }\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\right)\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}+\sum_{a\not=1,n}z_{a}^{2}z_{1} \left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}$ $\displaystyle\leq-z_{1}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} -\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})\left(z_{n}^{2}\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}+z_{1}^{2}\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,1}}-\max_{a\not=1,n}{z_{a}^{2}\mathbf{Attn}^{(t)}_ {\mathbf{p}\to\mathcal{P}_{k,a}}}\right)$ $$

Therefore, by Lemma A.1, we obtain

$$ $\displaystyle\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}$ $\displaystyle\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1} }\cdot\right.$ $\displaystyle\quad\left(-z_{1}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P} _{k,n}}-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})\right.$ $\displaystyle\left.\cdot\left(z_{n}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}+z_{1}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1 }}-\max_{a\not=1,n}{z_{a}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k, a}}}\right)\right)\left.\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}^{c} \cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{ k,1}}\cdot\sum_{a\not=1,n}z_{1}^{2}z_{a}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,a}}\right)^{2}\right]$ $\displaystyle\leq\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,1})\cdot\Bigg{(}- \Omega(1)\cdot\Omega\Big{(}\frac{1}{P^{2(1-\kappa_{c})+\frac{2U}{L}(\frac{ \Delta}{2}+0.01)}}\Big{)}\Bigg{)}+O(1)\cdot\mathbb{P}(\mathsf{M}\in\mathcal{E} _{k,1}^{c})$ $\displaystyle\leq-\Omega\Big{(}\frac{1}{P^{2(1-\kappa_{c})+\frac{U}{L}(\Delta+ 0.02)}}\Big{)}$ $$

where the second inequality invokes Lemma D.10 and the last inequality comes from Lemma B.6. The upper bound can be obtained by using similar arguments and invoking the upper bound for $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(t)}$ in Lemma D.10.
∎

###### Lemma D.14 .

For $n>1$, if C.1 and D.2 hold at iteration $T_{1}+1\leq t\leq\widetilde{T}_{1}$, then for any $m>1$ with $m\not=n$, the following holds

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}|\leq O\Big{(}\frac{\alpha^{ (t)}_{\mathbf{p}\to v_{k,n}}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P^{1-\kappa _{s}}}\Big{)}.$ $$

The proof is similar to Lemma D.5, and thus omitted here.

###### Lemma D.12 .

###### Proof.

###### Lemma D.13 .

###### Proof.

###### Lemma D.14 .

#### D.2.3 Bounding the Gradient Updates of Positional Correlations

We then summarize the properties for gradient updates of positional correlations, which utilize the identical calculations as in Section  D.1.3.

###### Lemma D.15 .

For $n>1$, if C.1 and D.2 hold at iteration $T_{1}+1\leq t\leq\widetilde{T}_{1}$, then

- a.
if $a_{k,\mathbf{q}}=n$ and $\mathbf{q}\not=\mathbf{p}$, $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}\geq 0$; $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta\Big{(}\frac{\alpha_{\mathbf{p}
\to v_{k,n}}^{(t)}}{C_{n}}\Big{)}\text{ and }|\beta^{(t)}_{n}|=O(\frac{\alpha_
{\mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P})$.
- b.
if $a_{k,\mathbf{q}}=1$, $|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v
_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}+O\Big{(}\frac{|
\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{C_{1}}\Big{)}$.
- c.
if $a_{k,\mathbf{q}}=m$ and $m\not=1,n$, $|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v
_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}$.

###### Lemma D.15 .

#### D.2.4 End of Phase I, Stage 2

###### Lemma D.16 .

D.2 holds for all iteration $T_{1}+1\leq t\leq\widetilde{T}_{1}=T_{1}+O\Big{(}\frac{\log(P)P^{1-\kappa_{s}}
}{\eta}\Big{)}$, and at iteration $t=\widetilde{T}_{1}+1$, we have

- a.
$\Phi_{\mathbf{p}\to v_{k,n}}^{(\widetilde{T}_{1}+1)}\geq\frac{c_{1}^{*}(1-
\kappa_{s})\log(P)}{U}$;
- b.
$\Phi_{\mathbf{p}\to v_{k,1}}^{(\widetilde{T}_{1}+1)}\geq-(\frac{\Delta}{2L}+
\frac{0.01}{L})\log(P)$.

###### Proof.

The existence of $\widetilde{T}_{1}=T_{1}+O\Big{(}\frac{\log(P)P^{1-\kappa_{s}}}{\eta}\Big{)}$ directly follows from Lemma D.12 and Lemma D.13. Moreover, since $\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}<0$, then

$$ $\Phi_{\mathbf{p}\to v_{k,n}}^{(\widetilde{T}_{1}+1)}\leq\left(\frac{\Delta}{2L }+\frac{0.01}{L}+\frac{c_{1}^{*}(1-\kappa_{s})}{U}\right)\log(P)-\frac{1}{U}( \frac{\Delta}{2}-0.01)\leq\frac{(c_{0}^{*}+c_{1}^{*})(1-\kappa_{s})}{U}\log(P)$ $$

where the last inequality invokes (D.6a).
Now suppose $\Phi_{\mathbf{p}\to v_{k,n}}^{(\widetilde{T}_{1}+1)}<\frac{c_{1}^{*}(1-\kappa_
{s})\log(P)}{U}$, then $\Phi_{\mathbf{p}\to v_{k,1}}^{(\widetilde{T}_{1}+1)}<-(\frac{\Delta}{2L}+\frac
{0.01}{L})\log(P)$. Denote the first time that $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$ reaches $-(\frac{\Delta}{2L}+\frac{0.001}{L})\log(P)$ as $\widetilde{T}$. Note that $\widetilde{T}<\widetilde{T}_{1}$ since $\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}$, the change of $\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}$, satisfies $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\ll\log(P)$. Then for $t\geq\widetilde{T}$, the following holds:

- 1.
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\geq\Omega\left(\frac{1}{P
^{1-\kappa_{s}}}\right)$;
- 2.
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\leq O\Big{(}\frac{1}{P^{
\frac{1-\kappa_{s}}{2}+0.001}}\Big{)}$.

Therefore,

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|$ $\displaystyle\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1} }\cdot\right.$ $\displaystyle\quad z_{1}\left(z_{n}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}})+z_{ 1}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}(1-\mathbf{Attn}^{(t )}_{\mathbf{p}\to\mathcal{P}_{k,1}})\right)\left.\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}^{c} \cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{ k,1}}\cdot\sum_{a\not=1,n}z_{1}^{2}z_{a}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,a}}\right)^{2}\right]$ $\displaystyle\leq O(\frac{\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P^{\frac{1- \kappa_{s}}{2}+0.001}})+\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,1})\cdot\bigg{( }O(1)\cdot O\Big{(}\frac{1}{P^{\frac{1-\kappa_{s}}{2}+0.001}}\Big{)}\bigg{)}+O (1)\cdot\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,1}^{c})$ $\displaystyle\leq O\Big{(}\frac{\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P^{\frac {1-\kappa_{s}}{2}+0.001}}\Big{)}+O\Big{(}\frac{1}{P^{(1-\kappa_{s})+0.002}} \Big{)}.$ $$

Lemma D.12 still holds, and thus

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\leq O\Big{(}\frac{\alpha_{ \mathbf{p}\to v_{k,n}}^{(t)}}{P^{0.002}}\Big{)}.$ $$

Since $|\Phi_{\mathbf{p}\to v_{k,1}}^{(\widetilde{T}_{1}+1)}-\Phi_{\mathbf{p}\to v_{k
,1}}^{(\widetilde{T})}|\geq\Omega\left(\log(P)\right)$, we have

$$ $\Phi_{\mathbf{p}\to v_{k,n}}^{(\widetilde{T}_{1}+1)}\geq|\Phi_{\mathbf{p}\to v _{k,1}}^{(\widetilde{T}_{1}+1)}-\Phi_{\mathbf{p}\to v_{k,1}}^{(\widetilde{T})} |\cdot\Omega(P^{0.002})+\Phi_{\mathbf{p}\to v_{k,n}}^{(\widetilde{T})}\gg \Omega(P^{0.002}\log(P)),$ $$

which contradicts the assumption that $\Phi_{\mathbf{p}\to v_{k,n}}^{(\widetilde{T}_{1}+1)}<\frac{c_{1}^{*}(1-\kappa_
{s})\log(P)}{U}$.
∎

###### Lemma D.16 .

###### Proof.

### D.3 Phase II, Stage 1

For $n>1$, we define stage 1 of phase II as all iterations $\widetilde{T}_{1}+1\leq t\leq T_{2}$, where

$$ $T_{2}\triangleq\max\left\{t:\Phi_{\mathbf{p}\to v_{k,n}}^{(t)}\leq\frac{(1- \kappa_{s})}{L}\log(P)\right\}.$ $$

We state the following induction hypotheses, which will hold throughout this stage:

###### Induction Hypothesis D.3 .

For each $\widetilde{T}_{1}+1\leq t\leq T_{2}$, $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, the following holds:

- a.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}$ is monotonically increasing, and $\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}\in\Big{[}\frac{c^{*}_{1}(1-\kappa_{s})}{U}
\log(P),\frac{(1-\kappa_{s})}{L}\log(P)\Big{]}$;
- b.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ is monotonically decreasing and
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}\in\Big{[}-\frac{1}{L}\left(\frac{\Delta}{2}
+0.01\right)\log(P)-o(1),-\frac{1}{U}\left(\frac{\Delta}{2}-0.01\right)\log(P)
\Big{]};$
- c.
$|\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p}\to v
_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}$ for $m\not=1,n$;
- d.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to
v
_{k,n}}}{C_{n}}\Big{)}$ for $a_{k,\mathbf{q}}=n$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$;
- e.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{|\Phi^{(t)}_{\bm{p}
\to v_{k,1}}|}{C_{1}}\Big{)}+O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to v_{k,n}}-\Phi
^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}=1$.;
- f.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1,n$.

###### Induction Hypothesis D.3 .

#### D.3.1 Property of Attention Scores

We first single out several properties of attention scores
that will be used for the proof of D.3.

###### Lemma D.17 .

if C.1 and D.3 hold at iteration $\widetilde{T}_{1}+1\leq t\leq T_{2}$, then the following holds

- 1.
if $\mathsf{M}\in\mathcal{E}_{k,n}$,
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\geq\Omega\Big{(}\frac{1}{
P^{(1-\frac{c_{1}^{*}L}{U})(1-\kappa_{s})}}\Big{)}$
. Moreover, if $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}$ does not reach the constant level, $1-\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}=\Omega(1)$; otherwise, $1-\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}=\Omega\left(\frac{1}{P^
{(\frac{U}{L}-1)(1-\kappa_{s})}}\right)$.
- 2.
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=O\Big{(}\frac{1-\mathbf{
Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}}{P^{(1-\kappa_{c})+\frac{L}{U}(
\frac{\Delta}{2}-0.01)}}\Big{)}$; if $\mathsf{M}\in\mathcal{E}_{k,1}$, we have $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=\Omega\Big{(}\frac{1}{P^{
(1-\kappa_{c})+\frac{U}{L}(\frac{\Delta}{2}+0.01)}}\Big{)}$;
- 3.
for ${\mathbf{q}}\in\mathcal{M}\cap(\mathcal{P}_{k,n}\cup\mathcal{P}_{k,1})$, $\mathbf{attn}_{\mathbf{p}\to{\mathbf{q}}}^{(t)}=O\Big{(}\frac{1-\mathbf{Attn}_
{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}}{P}\Big{)}$

###### Lemma D.18 .

if C.1 and D.3 hold at iteration $\widetilde{T}_{1}+1\leq t\leq T_{2}$, then for $m\not=n$, the following holds:

- 1.
for any ${\mathbf{q}}\in\mathcal{P}_{k,m}$, $\mathbf{attn}^{(t)}_{\mathbf{p}\to{\mathbf{q}}}\leq O\Big{(}\frac{1-\mathbf{
Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}}{P}\Big{)}$.
- 2.
Moreover, $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\leq O\Big{(}\frac{1-
\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\mathbf{Attn}^{(t)}_{
\mathbf{p}\to\mathcal{P}_{k,n}}}{N}\Big{)}.$

###### Lemma D.17 .

###### Lemma D.18 .

#### D.3.2 Bounding the Gradient Updates of FP Correlations

###### Lemma D.19 .

if C.1 and D.3 hold at iteration $\widetilde{T}_{1}+1\leq t\leq T_{2}$, then $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}\geq 0$ and satisfies:

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}\geq\min\bigg{\{}\Omega(\frac {1}{P^{(1-\frac{c_{1}^{*}L}{U})(1-\kappa_{s})}}),\Omega\left(\frac{1}{P^{2( \frac{U}{L}-1)(1-\kappa_{s})}}\right)\bigg{\}}.$ $$

###### Proof.

By Lemma A.2, we have

$$ $\displaystyle\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in\mathcal{M}\} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_{n}^{3}\left( 1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not= n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}} \right)^{2}\right)\right]$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}\cap\mathbf {p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot \left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} \right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}^{c} \cap\mathbf{p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k ,n}}\cdot\left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_ {k,n}}\right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle{\gtrsim}\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n})\cdot$ $\displaystyle~{}~{}\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in\mathcal{M }\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_{n}^{3} \left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{ m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m }}\right)^{2}\right)\Big{|}\mathcal{E}_{k,n}\right]$ $\displaystyle\quad+O(1)\cdot\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n}^{c})$ $\displaystyle\gtrsim\min\bigg{\{}\Omega\bigg{(}\frac{1}{P^{(1-\frac{c_{1}^{*}L }{U})(1-\kappa_{s})}}\bigg{)},\Omega\left(\frac{1}{P^{2(\frac{U}{L}-1)(1- \kappa_{s})}}\right)\bigg{\}}$ $$

where the last inequality invokes Lemma D.17 by observing that for $\mathsf{M}\in\mathcal{E}_{k,n}$,

$$ $\displaystyle\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}(1-\mathbf{ Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)})^{2}$ $\displaystyle\geq\min\bigg{\{}\Omega\left(\frac{1}{P^{(1-\frac{c_{1}^{*}L}{U}) (1-\kappa_{s})}}\right)\cdot\Omega(1),\Omega(1)\cdot\Omega\left(\frac{1}{P^{2 \times(\frac{U}{L}-1)(1-\kappa_{s})}}\right)\bigg{\}}.$ $$

∎

###### Lemma D.20 .

For $n>1$, if C.1 and D.3 hold at iteration $\widetilde{T}_{1}+1\leq t\leq T_{2}$, then $\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}<0$ and satisfies

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}|$ $\displaystyle\geq\min\bigg{\{}\Omega\left(\frac{1}{P^{(1-\frac{c_{1}^{*}L}{U}) (1-\kappa_{s})}}\right),\Omega\left(\frac{1}{P^{(\frac{U}{L}-1)(1-\kappa_{s})} }\right)\bigg{\}}\cdot\Omega\Big{(}\frac{1}{P^{(1-\kappa_{c})+\frac{L}{U}( \frac{\Delta}{2}-0.01)}}\Big{)},$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}|$ $\displaystyle\leq\max\Big{\{}O\Big{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t) }}{P^{(1-\kappa_{c})+\frac{L}{U}(\Delta/2-0.01)}}\Big{)},O\Big{(}\frac{\alpha_ {\mathbf{p}\to v_{k,n}}^{(t)}}{P^{2(1-\kappa_{c})+\frac{L}{U}(\Delta-0.02)-(1- \frac{c_{1}^{*}L}{U})(1-\kappa_{s})}}\Big{)}\Big{\}}.$ $$

###### Proof.

Following (D.1), we have

$$ $\displaystyle-z_{1}z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal {P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-z_{1}^{3 }\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\right)\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}+\sum_{a\not=1,n}z_{a}^{2}z_{1} \left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}$ $\displaystyle\leq-z_{1}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} -\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})\left(z_{n}^{2}\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}+z_{1}^{2}\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,1}}-\max_{a\not=1,n}{z_{a}^{2}\mathbf{Attn}^{(t)}_ {\mathbf{p}\to\mathcal{P}_{k,a}}}\right).$ $$

Therefore, by Lemma A.1, we obtain

$$ $\displaystyle\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}$ $\displaystyle\leq\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1} }\cdot\right.$ $\displaystyle\quad\left(-z_{1}(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P} _{k,n}}-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}})\left(z_{n}^{2} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}+z_{1}^{2}\mathbf{Attn}^{( t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}-\max_{a\not=1,n}{z_{a}^{2}\mathbf{Attn}^{ (t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}}\right)\right)\left.\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}^{c} \cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{ k,1}}\cdot\sum_{a\not=1,n}z_{1}^{2}z_{a}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,a}}\right)^{2}\right]$ $\displaystyle\leq-\min\bigg{\{}\Omega\left(\frac{1}{P^{(1-\frac{c_{1}^{*}L}{U} )(1-\kappa_{s})}}\right),\Omega\left(\frac{1}{P^{(\frac{U}{L}-1)(1-\kappa_{s}) }}\right)\bigg{\}}\cdot\Omega\Big{(}\frac{1}{P^{(1-\kappa_{c})+\frac{L}{U}( \frac{\Delta}{2}-0.01)}}\Big{)}$ $$

where the second inequality invokes Lemma D.17 and (D.6b).
Moreover,

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\lesssim$ $\displaystyle\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathcal{E}_{k,n}\cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p }\to\mathcal{P}_{k,1}}\cdot\right.$ $\displaystyle\quad\left(z_{1}z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n }}+z_{1}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\right )\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\right)\left.\right]$ $\displaystyle=\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathcal{E}_{k,n}\cap\mathbf{p}\in\mathcal{M}\}}z_{1}z_{n}^{2}\mathbf{Attn}^{( t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}\cdot\left(1-\mathbf{Attn}^{(t)}_{\mathbf{ p}\to\mathcal{P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k ,n}}\right]$ $\displaystyle+\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,1}\cap \mathcal{E}_{k,n}\cap\mathbf{p}\in\mathcal{M}\}}z_{1}^{3}(\mathbf{Attn}^{(t)}_ {\mathbf{p}\to\mathcal{P}_{k,1}})^{2}\cdot\left(1-\mathbf{Attn}^{(t)}_{\mathbf {p}\to\mathcal{P}_{k,1}}\right)\right]$ $\displaystyle\leq\max\Big{\{}O\Bigg{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t )}}{P^{(1-\kappa_{c})+\frac{L}{U}(\frac{\Delta}{2}-0.01)}}\Big{)},O\Bigg{(} \frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P^{2(1-\kappa_{c})+\frac{2L}{U}( \frac{\Delta}{2}-0.01)-(1-\frac{c_{1}^{*}L}{U})(1-\kappa_{s})}}\Bigg{)}\Bigg{\}}$ $$

where the second inequality invokes Lemma D.17.
∎

###### Lemma D.21 .

For $n>1$, if C.1 and D.3 hold at iteration $\widetilde{T}_{1}+1\leq t\leq T_{2}$ for any $m>1$ with $m\not=n$, the following holds

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}|\leq O\Big{(}\frac{\alpha^{ (t)}_{\mathbf{p}\to v_{k,n}}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P^{1-\kappa _{s}}}\Big{)}.$ $$

The proof is similar to Lemma D.5, and thus omitted here.

###### Lemma D.19 .

###### Proof.

###### Lemma D.20 .

###### Proof.

###### Lemma D.21 .

#### D.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  D.1.3.

###### Lemma D.22 .

For $n>1$, if C.1 and D.3 hold at iteration $\widetilde{T}_{1}+1\leq t\leq{T}_{2}$, then

- a.
if $a_{k,\mathbf{q}}=n$ and $\mathbf{q}\not=\mathbf{p}$, $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}\geq 0$; $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta(\frac{\alpha_{\mathbf{p}\to v_{
k,n}}^{(t)}}{C_{n}})\text{ and }|\beta^{(t)}_{n}|=O\Big{(}\frac{\alpha_{
\mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}$.
- b.
if $a_{k,\mathbf{q}}=1$, $|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v
_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}+O\Big{(}\frac{|
\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{C_{1}}\Big{)}$.
- c.
if $a_{k,\mathbf{q}}=m$ and $m\not=1,n$, $|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v
_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}$.

###### Lemma D.22 .

#### D.3.4 End of Phase II, Stage 1

###### Lemma D.23 .

D.3 holds for all $\widetilde{T}_{1}+1\leq t\leq T_{2}$, and at iteration $t=T_{2}+1$, we have

- a.
$\Phi_{\mathbf{p}\to v_{k,n}}^{(t)}>\frac{(1-\kappa_{s})}{L}\log(P)$;
- b.
$\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}=\Omega(1)$ if $\mathsf{M}\in\mathcal{E}_{k,n}$.

###### Proof.

By comparing Lemma D.19 and Lemma D.20-D.23, we have $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}\gg|\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}|,
|\beta_{k,\mathbf{p}\to\mathbf{q}}^{(t)}|$. Then the existence of $T_{2}=\widetilde{T}_{1}+O\Big{(}\frac{\log(P)P^{\Lambda}}{\eta}\Big{)}$ directly follows from Lemma D.19, where

$$ $\Lambda=\max\Big{\{}(1-\frac{c_{1}^{*}L}{U}),2(\frac{U}{L}-1)\Big{\}}\cdot(1- \kappa_{s}).$ $$

The second statement can be directly verified by noticing that $\Phi_{\mathbf{p}\to v_{k,n}}^{(t)}>\frac{(1-\kappa_{s})}{L}\log(P)$ while all other attention correlations are sufficiently small.
∎

###### Lemma D.23 .

###### Proof.

### D.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<\epsilon<1$, for $n>1$, define

$$ $\displaystyle T^{\epsilon}_{2}\triangleq\max\left\{t>T_{2}:\Phi_{\mathbf{p}\to v _{k,n}}^{(t)}\leq\log\left(c_{5}\left(\left(\frac{3}{\epsilon}\right)^{\frac{1 }{2}}-1\right)N\right)\right\}.$ (D.7) $$

where $c_{5}$ is some largely enough constant.

We state the following induction hypotheses, which will hold throughout this stage:

###### Induction Hypothesis D.4 .

For $n>1$, suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$, for each $T_{2}+1\leq t\leq T_{2}^{\epsilon}$, $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, the following holds:

- a.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}$ is monotonically increasing, and $\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}\in[\frac{(1-\kappa_{s})}{L}\log(P),O(\log(P
/\epsilon))]$;
- b.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ is monotonically decreasing and $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}\in\Big{[}-\frac{1}{L}\left(\frac{\Delta}{2}
+0.01\right)\log(P)-o(1),-\frac{1}{U}\left(\frac{\Delta}{2}-0.01\right)\log(P)
\Big{]}$;
- c.
$|\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p}\to v
_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}$ for $m\not=1,n$;
- d.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to
v
_{k,n}}}{C_{n}}\Big{)}$ for $a_{k,\mathbf{q}}=n$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$;
- e.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{|\Phi^{(t)}_{\bm{p}
\to v_{k,1}}|}{C_{1}}\Big{)}+O\Big{(}\frac{\Phi^{(t)}_{\bm{p}\to v_{k,n}}-\Phi
^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}=1$.;
- f.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\bm{p}
\to v_{k,n}}-\Phi^{(t)}_{\bm{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1,n$.

###### Induction Hypothesis D.4 .

#### D.4.1 Property of Attention Scores

We first single out several properties of attention scores that will be used for the proof of D.4.

###### Lemma D.24 .

if C.1 and D.4 hold at iteration $T_{n,2}<t\leq T_{n,2}^{\epsilon}$, then the following holds

- 1.
if $\mathsf{M}\in\mathcal{E}_{k,n}$,
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}=\Omega(1)$ and $(1-\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)})^{2}\geq O(\epsilon)$.
- 2.
Moreover, $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=O\Big{(}\frac{1-\mathbf{
Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}}{P^{(1-\kappa_{c})+\frac{L}{U}(
\frac{\Delta}{2}-0.01)}}\Big{)}$; if $\mathsf{M}\in\mathcal{E}_{k,1}$, we have $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}=\Omega\Big{(}\frac{1-
\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}}{P^{(1-\kappa_{c})+\frac{
U}{L}(\frac{\Delta}{2}+0.01)}}\Big{)}$;
- 3.
for ${\mathbf{q}}\in\mathcal{M}\cap(\mathcal{P}_{k,n}\cup\mathcal{P}_{k,1})$, $\mathbf{attn}_{\mathbf{p}\to{\mathbf{q}}}^{(t)}=O\Big{(}\frac{1-\mathbf{Attn}_
{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}}{P}\Big{)}$.

###### Lemma D.25 .

if C.1 and D.4 hold at iteration $T_{n,2}<t\leq T^{\epsilon}_{n,2}$,then for $m\not=n$, the following holds:

- 1.
for any ${\mathbf{q}}\in\mathcal{P}_{k,m}$, $\mathbf{attn}^{(t)}_{\mathbf{p}\to{\mathbf{q}}}\leq O\Big{(}\frac{1-\mathbf{
Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(t)}}{P}\Big{)}$.
- 2.
$\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\leq O\Big{(}\frac{1-
\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}}{N}\Big{)}$, and if $\mathsf{M}\in\mathcal{E}_{k,m}$, $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}=\Theta\Big{(}\frac{1-
\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}}{N}\Big{)}$.

###### Lemma D.24 .

###### Lemma D.25 .

#### D.4.2 Bounding the Gradient Updates of FP Correlations

###### Lemma D.26 .

For $n>1$, if C.1 and D.4 hold at iteration $T_{2}+1\leq t\leq T^{\epsilon}_{2}$, then $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}\geq 0$ and satisfies:

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}\geq\Omega(\epsilon).$ $$

###### Proof.

By Lemma A.2, we have

$$ $\displaystyle\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in\mathcal{M}\} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_{n}^{3}\left( 1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not= n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}} \right)^{2}\right)\right]$ $\displaystyle=\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}\cap\mathbf {p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot \left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}} \right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p} \to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle\quad+\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathcal{E}_{k,n}^{c} \cap\mathbf{p}\in\mathcal{M}\}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k ,n}}\cdot\left(z_{n}^{3}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_ {k,n}}\right)^{2}+\sum_{m\not=n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,m}}\right)^{2}\right)\right]$ $\displaystyle{\gtrsim}\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n})\cdot$ $\displaystyle\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in\mathcal{M}\} \mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\cdot\left(z_{n}^{3}\left( 1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not= n}z_{m}^{2}z_{n}\left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}} \right)^{2}\right)\Big{|}\mathcal{E}_{k,n}\right]$ $\displaystyle\quad+O(1)\cdot\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,n}^{c})$ $\displaystyle\gtrsim\Omega(\epsilon)$ $$

where the last inequality invokes Lemma D.24, Lemma B.6 and the fact that

$$ $\displaystyle{\epsilon}\geq{\exp(-\operatorname{polylog}(K))}\gg\exp\left(-c_{ n,1}C_{n}\right).$ $$

∎

###### Lemma D.27 .

For $n>1$, if C.1 and D.4 hold at iteration $T_{n,3}<t\leq T^{\epsilon}_{n,4}$, then $\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}<0$ and satisfies

$$ $\displaystyle|\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}|\leq\max\bigg{\{}O\Big{(} \frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P^{(1-\kappa_{c})+\frac{L}{U}( \Delta/2-0.01)}}\Big{)},O\bigg{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P ^{2(1-\kappa_{c})+\frac{L}{U}(\Delta-0.02)-(1-\frac{c_{1}^{*}L}{U})(1-\kappa_{ s})}}\bigg{)}\bigg{\}}$ $$

The proof follows the similar arguments Lemma D.20 by noticing that $\epsilon\gg\mathbb{P}(\mathsf{M}\in\mathcal{E}_{k,m}^{c})$ for any $m\not=n$.

###### Lemma D.28 .

For $n>1$, if C.1 and D.4 hold at iteration $T_{2}<t\leq T^{\epsilon}_{2}$, then for any $m>1$ with $m\not=n$, the following holds

$$ $\displaystyle-O(\frac{\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}}{P^{1-\kappa_{s}}}) \leq\alpha_{\mathbf{p}\to v_{k,m}}^{(t)}\leq 0$ $$

###### Proof.

We first note that

$$ $\displaystyle-z_{1}z_{n}^{2}\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal {P}_{k,n}}\right)\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-z_{m}^{3 }\left(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\right)\mathbf{ Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,1}}+\sum_{a\not=1,n}z_{a}^{2}z_{m} \left(\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}\right)^{2}$ $\displaystyle\leq z_{m}\left(\max_{a\not=m,n}{z_{a}^{2}\mathbf{Attn}^{(t)}_{ \mathbf{p}\to\mathcal{P}_{k,a}}}-z_{n}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to \mathcal{P}_{k,n}}-z_{m}^{2}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m }}\right)\Big{(}1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}-\mathbf {Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,m}}\Big{)}$ $\displaystyle\lesssim-\Omega(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{ k,n}})$ $$

since when $\mathsf{M}\in\mathcal{E}_{k,n}$, we have $\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}}=\Omega(1)\gg\mathbf{Attn}
^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,a}}$. Thus, we have

$$ $\displaystyle 0\geq\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}$ $\displaystyle\gtrsim-\mathbb{E}\left[\mathds{1}{\{k_{X}=k,\mathcal{E}_{k,n} \cap\mathbf{p}\in\mathcal{M}\}}\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{ k,m}}\cdot\Omega(1-\mathbf{Attn}^{(t)}_{\mathbf{p}\to\mathcal{P}_{k,n}})\right]$ $\displaystyle\geq-O\Big{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P^{1- \kappa_{s}}}\Big{)}.$ $$

∎

###### Lemma D.26 .

###### Proof.

###### Lemma D.27 .

###### Lemma D.28 .

###### Proof.

#### D.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 in Section  D.1.3.

###### Lemma D.29 .

For $n>1$, if C.1 and D.4 hold at iteration ${T}_{2}+1\leq t\leq{T}^{\epsilon}_{2}$, then

- a.
if $a_{k,\mathbf{q}}=n$ and $\mathbf{q}\not=\mathbf{p}$, $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}\geq 0$; $\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta\Big{(}\frac{\alpha_{\mathbf{p}
\to v_{k,n}}^{(t)}}{C_{n}}\Big{)}\text{ and }|\beta^{(t)}_{n}|=O\Big{(}\frac{
\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}
\Big{)}$.
- b.
if $a_{k,\mathbf{q}}=1$, $|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v
_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}+O\Big{(}\frac{|
\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|}{C_{1}}\Big{)}$.
- c.
if $a_{k,\mathbf{q}}=m$ and $m\not=1,n$, $|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v
_{k,n}}^{(t)}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}$.

###### Lemma D.29 .

#### D.4.4 End of Phase II, Stage 2

###### Lemma D.30 .

For $n>1$, and $0<\epsilon<1$, suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$. Then D.4 holds for all $T_{2}<t\leq T^{\epsilon}_{2}=T_{2}+O\Big{(}\frac{\log(P\epsilon^{-1})}{\eta
\epsilon}\Big{)}$, and at iteration $t=T^{\epsilon}_{2}+1$, we have

- 1.
$\widetilde{\mathcal{L}}_{k,\mathbf{p}}(Q^{T^{\epsilon}_{2}+1})<\frac{\epsilon}
{2K}$;
- 2.
If $\mathsf{M}\in\mathcal{E}_{k,n}$ , we have $(1-\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(T^{\epsilon}_{2}+1)})^{2}
\leq O(\epsilon)$.

###### Proof.

The existence of $T_{2,k}^{\epsilon}=T_{2,k}+O(\frac{\log(P\epsilon^{-1})}{\eta\epsilon})$ directly follows from Lemma D.26. We further derive

$$ $\displaystyle\widetilde{\mathcal{L}}_{k,\mathbf{p}}(Q^{T^{\epsilon}_{2}+1})=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[\mathds{1}\{k_{X}=k,\mathbf{p}\in \mathcal{M}\cap\mathsf{M}\in\mathcal{E}_{k,n}\}\left(z_{n}^{2}\left(1-\mathbf{ Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}\right)^{2}+\sum_{m\not=n}z_{m}^{2}\left (\mathbf{Attn}_{n,m}\right)^{2}\right)\right]$ $\displaystyle\leq\frac{1}{2K}\cdot\gamma\cdot U^{2}\cdot(1+o(1))\cdot O(\epsilon)$ $\displaystyle\leq\frac{\epsilon}{2K}$ $$

where the first inequality is due to direct calculations by the definition of $T_{2}^{\epsilon}$, and the second inequality can be obtained by setting $c_{n,2}$ in (D.7) sufficiently large.
∎

###### Lemma D.30 .

###### Proof.

## Appendix E Analysis for Local Areas with Negative Information Gap

In this section, we focus on a specific patch $\mathbf{p}\in\mathcal{P}$ with the $k$-th cluster for $k\in[K]$, and present the analysis for the case that $X_{\mathbf{p}}$ is located in the local area for the $k$-th cluster, i.e. $a_{k,\bm{p}}>1$. Throughout this section, we denote $a_{k,\bm{p}}=n$ for simplicity. When $\Delta\leq-\Omega(1)$, we can show that the gap of attention correlation changing rate for the positive case does not exist anymore, and conversely $\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}\gg\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}$ from the beginning. We can reuse most of the gradient calculations in the previous section and only sketch them in this section.

##### Stage 1:

we define stage 1 as all iterations $0\leq t\leq T_{\mathrm{neg},1}$, where

$$ $T_{\mathrm{neg},1}\triangleq\max\left\{t:\Phi_{\mathbf{p}\to v_{k,n}}^{(t)} \leq\frac{(1-\kappa_{s})}{L}\log(P)\right\}.$ $$

We state the following induction hypothesis, which will hold throughout this stage:

###### Induction Hypothesis E.1 .

For each $0\leq t\leq T_{\mathrm{neg},1}$, $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, the following holds:

- a.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}$ is monotonically increasing, and $\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}\in\Big{[}0,\frac{(1-\kappa_{s})}{L}\log(P)
\Big{]}$;
- b.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ is monotonically decreasing and $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}\in\bigg{[}-O\Big{(}\frac{\Phi^{(t)}_{
\mathbf{p}\to v_{k,n}}}{P^{-\Delta}}\Big{)},0\bigg{]};$
- c.
$|\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p}\to v
_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}$ for $m\not=1,n$;
- d.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p
}\to v_{k,n}}}{C_{n}}\Big{)}$ for $a_{k,\mathbf{q}}=n$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{)}$;
- e.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{|\Phi_{\mathbf{p}
\to v_{k,1}}^{(t)}|}{C_{1}}\Big{)}+O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p}\to v_{
k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}=1$;
- f.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1,n$.

Through similar calculations for phase II, stage 1 in Section D.3, we obtain the following lemmas to control the gradient updates for attention correlations.

###### Lemma E.1 .

If C.2 and E.1 hold for $0\leq t\leq T_{\mathrm{neg},1}$, then we have

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}\geq\min\Bigg{\{}\Omega\Big{( }\frac{1}{P^{(1-\kappa_{s})}}\Big{)},\Omega\left(\frac{1}{P^{2(\frac{U}{L}-1)( 1-\kappa_{s})}}\right)\Bigg{\}},$ (E.1a) $\displaystyle 0\geq\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}\geq-O\Big{(}\frac{ \alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P^{-\Delta}}\Big{)},$ (E.1b) $\displaystyle|\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}|\leq O\Big{(}\frac{\alpha^{ (t)}_{\mathbf{p}\to v_{k,n}}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P^{1-\kappa _{s}}}\Big{)}\text{ for all }m\not=n,1$ (E.1c) $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta\Big{(}\frac{\alpha _{\mathbf{p}\to v_{k,n}}^{(t)}}{C_{n}}\Big{)}\text{ for }a_{k,\mathbf{q}}=n, \mathbf{q}\not=\mathbf{p}$ (E.1d) $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{ \mathbf{p}\to v_{k,n}}^{(t)}}{P}\Big{)}+O\Big{(}\frac{|\alpha_{\mathbf{p}\to v _{k,1}}^{(t)}|}{C_{1}}\Big{)}\text{ for }a_{k,\mathbf{q}}=1,$ (E.1e) $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|,|\beta^{(t)}_{k,\mathbf {p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{ \mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}\quad\text{ for all }a_{k,\mathbf{p}} \not=n,1.$ (E.1f) $$

Here $\Delta<0$ implies $|\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}|\ll\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}$. E.1 can be directly proved by Lemma E.1 and we have

$$ $\displaystyle T_{\mathrm{neg},1}=O\Big{(}\frac{P^{\max\{1,2(\frac{U}{L}-1)\} \cdot(1-\kappa_{s})}\log(P)}{\eta}\Big{)}.$ (E.2) $$

###### Induction Hypothesis E.1 .

###### Lemma E.1 .

##### Stage 2:

Given any $0<\epsilon<1$, define

$$ $\displaystyle T_{\mathrm{neg},1}^{\epsilon}\triangleq\max\left\{t>T_{1}:\Phi_{ \mathbf{p}\to v_{k,n}}^{(t)}\leq\log\left(c_{6}\left(\left(\frac{3}{\epsilon} \right)^{\frac{1}{2}}-1\right)P^{1-\kappa_{s}}\right)\right\}.$ (E.3) $$

where $c_{6}$ is some largely enough constant. We then state the following induction hypotheses, which will hold throughout this stage:

###### Induction Hypothesis E.2 .

For $n>1$, suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$, for $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, and each $T_{\mathrm{neg},1}<t\leq T_{\mathrm{neg},1}^{\epsilon}$, the following holds:

- a.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}$ is monotonically increasing, and $\Phi^{(t)}_{\mathbf{p}\to v_{k,n}}\in\Big{[}\frac{(1-\kappa_{s})}{L}\log(P),O(
\log(P/\epsilon))\Big{]}$;
- b.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ is monotonically decreasing and $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}\in\Bigg{[}-O\Big{(}\frac{\Phi^{(t)}_{
\mathbf{p}\to v_{k,n}}}{P^{-\Delta}}\Big{)},0\Bigg{]};$
- c.
$|\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p}\to v
_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}$ for $m\not=1,n$;
- d.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p
}\to v_{k,n}}}{C_{n}}\Big{)}$ for $a_{k,\mathbf{q}}=n$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{)}$;
- e.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{|\Phi_{\mathbf{p}
\to v_{k,1}}^{(t)}|}{C_{1}}\Big{)}+O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p}\to v_{
k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}=1$;
- f.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,n}}-\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1,n$.

###### Lemma E.2 .

If C.2 and E.2 hold for $T_{\mathrm{neg},1}<t\leq T_{\mathrm{neg},1}^{\epsilon}$, then we have

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,n}}\geq\Omega({\epsilon}),$ (E.4a) $\displaystyle 0\geq\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}\geq-O\Big{(}\frac{ \alpha_{\mathbf{p}\to v_{k,n}}^{(t)}}{P^{-\Delta}}\Big{)},$ (E.4b) $\displaystyle|\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}|\leq O\Big{(}\frac{\alpha^{ (t)}_{\mathbf{p}\to v_{k,n}}-\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}}{P^{1-\kappa _{s}}}\Big{)}\text{ for all }m\not=n,1$ (E.4c) $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta(\frac{\alpha_{ \mathbf{p}\to v_{k,n}}^{(t)}}{C_{n}})\text{ for }a_{k,\mathbf{q}}=n,\mathbf{q} \not=\mathbf{p}$ (E.4d) $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{ \mathbf{p}\to v_{k,n}}^{(t)}}{P}\Big{)}+O\Big{(}\frac{|\alpha_{\mathbf{p}\to v _{k,1}}^{(t)}|}{C_{1}}\Big{)}\text{ for }a_{k,\mathbf{q}}=1,$ (E.4e) $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|,|\beta^{(t)}_{k,\mathbf {p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha_{\mathbf{p}\to v_{k,n}}^{(t)}-\alpha_{ \mathbf{p}\to v_{k,1}}^{(t)}}{P}\Big{)}\quad\text{ for all }a_{k,\mathbf{p}} \not=n,1.$ (E.4f) $$

E.2 can be directly proved by Lemma E.2. Furthermore, at the end of this stage, we will have:

###### Lemma E.3 .

Suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$, then E.2 holds for all $T_{\mathrm{neg},1}<t\leq T_{\mathrm{neg},1}^{\epsilon}=T_{\mathrm{neg},1}+O
\Big{(}\frac{\log(P\epsilon^{-1})}{\eta\epsilon}\Big{)}$, and at iteration $t=T_{\mathrm{neg},1}^{\epsilon}+1$, we have

- 1.
$\widetilde{\mathcal{L}}_{k,\mathbf{p}}(Q^{T_{\mathrm{neg},1}^{\epsilon}+1})<
\frac{\epsilon}{2K}$;
- 2.
If $\mathsf{M}\in\mathcal{E}_{k,n}$ , we have $\Big{(}1-\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,n}}^{(T_{\mathrm{neg},1}^{
\epsilon}+1)}\Big{)}^{2}\leq O(\epsilon)$.

###### Induction Hypothesis E.2 .

###### Lemma E.2 .

###### Lemma E.3 .

## Appendix F Analysis for the Global area

When $a_{\mathbf{p},k}=1$, i.e. the patch lies in the global area, the analysis is much simpler and does not depend on the value of $\Delta$. We can reuse most of the gradient calculations in Appendix D and only sketch them in this section.

For $X_{\mathbf{p}}$ in the global region $\mathcal{P}_{k,1}$, since the overall attention $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(0)}$ to the target feature already reaches $\Omega\Big{(}\frac{C_{1}}{P}\Big{)}=\Omega\Big{(}\frac{1}{P^{1-\kappa_{c}}}
\Big{)}$ due to the large number of unmasked patches featuring $v_{k,1}$ when $\mathsf{M}\in\mathcal{E}_{k,1}$, which is significantly larger than $\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,m}}^{(0)}=\Theta\Big{(}\frac{1}{P^{
1-\kappa_{s}}}\Big{)}$ for all other $m>1$. This results in large $\alpha_{\mathbf{p}\to v_{k,1}}^{(t)}$ initially, and thus the training directly enters phase II.

##### Stage 1:

we define stage 1 as all iterations $0\leq t\leq{T_{c,1}}$, where

$$ ${T_{c,1}}\triangleq\max\left\{t:\Phi_{\mathbf{p}\to v_{k,1}}^{(t)}\leq\frac{(1 -\kappa_{c})}{L}\log(P)\right\}.$ $$

We state the following induction hypotheses, which will hold throughout this stage:

###### Induction Hypothesis F.1 .

For each $0\leq t\leq{T_{c,1}}$, $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, the following holds:

- a.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ is monotonically increasing, and $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}\in\Big{[}0,\frac{(1-\kappa_{c})}{L}\log(P)
\Big{]}$;
- b.
$\Phi_{\mathbf{p}\to v_{k,m}}$ is monotonically decreasing for $m>1$ and $\Phi_{\mathbf{p}\to v_{k,m}}\in\Big{[}-O\big{(}\frac{\log(P)}{N}\big{)},0\Big{]}$;
- c.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p
}\to v_{k,1}}}{C_{1}}\Big{)}$ for $a_{k,\mathbf{q}}=1$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,1}}}{P}\Big{)}$;
- d.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1$.

Through similar calculations for phase II, stage 1 in Section D.3, we obtain the following lemmas to control the gradient updates for attention correlations.

###### Lemma F.1 .

If C.1 (or C.2) and F.1 hold for $0\leq t\leq{T_{c,1}}$, then we have

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,1}}\geq\min\Bigg{\{}\Omega\Big{( }\frac{1}{P^{(1-\kappa_{c})}}\Big{)},\Omega\left(\frac{1}{P^{2(\frac{U}{L}-1)( 1-\kappa_{c})}}\right)\Bigg{\}},$ (F.1a) $\displaystyle|\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}|\leq O\Big{(}\frac{\alpha^{ (t)}_{\mathbf{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}\quad\text{ for all }m \not=1,$ (F.1b) $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta\Big{(}\frac{\alpha ^{(t)}_{\mathbf{p}\to v_{k,1}}}{C_{1}}\Big{)},\text{ for }a_{k,\mathbf{q}}=1, \mathbf{q}\not=\mathbf{p},$ (F.1c) $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|,|\beta^{(t)}_{k,\mathbf {p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{ )}\quad\text{ for all }a_{k,\mathbf{q}}>1.$ (F.1d) $$

F.1 can be directly proved by Lemma F.1 and we have

$$ $\displaystyle{T_{c,1}}=O\Bigg{(}\frac{P^{\max\{1,2(\frac{U}{L}-1)\}\cdot(1- \kappa_{c})}\log(P)}{\eta}\Bigg{)}.$ (F.2) $$

###### Induction Hypothesis F.1 .

###### Lemma F.1 .

##### Stage 2:

Given any $0<\epsilon<1$, define

$$ $\displaystyle T^{\epsilon}_{c,1}\triangleq\max\left\{t>{T_{c,1}}:\Phi_{\mathbf {p}\to v_{k,1}}^{(t)}\leq\log\left(c_{7}\left(\left(\frac{3}{\epsilon}\right)^ {\frac{1}{2}}-1\right)P^{1-\kappa_{c}}\right)\right\}.$ (F.3) $$

where $c_{7}$ is some largely enough constant. We then state the following induction hypotheses, which will hold throughout this stage:

###### Induction Hypothesis F.2 .

For $n>1$, suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$, $\mathbf{q}\in\mathcal{P}\setminus\{\mathbf{p}\}$, for each ${T_{c,1}}+1\leq t\leq T^{\epsilon}_{c,1}$, the following holds:

- a.
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ is monotonically increasing, and $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}\in\Big{[}\frac{(1-\kappa_{c})}{L}\log(P),O(
\log(P/\epsilon))\Big{]}$;
- b.
$\Phi_{\mathbf{p}\to v_{k,m}}$ is monotonically decreasing for $n>1$ and $\Phi_{\mathbf{p}\to v_{k,m}}\in\Big{[}-O\big{(}\frac{\log(P)}{N}\big{)},0\Big{]}$;
- c.
$\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=O\Big{(}\frac{\Phi^{(t)}_{\mathbf{p
}\to v_{k,1}}}{C_{1}}\Big{)}$ for $a_{k,\mathbf{q}}=1$, $|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,1}}}{P}\Big{)}$;
- d.
$|\Upsilon^{(t)}_{k,\mathbf{p}\to\mathbf{q}}|=O\Big{(}\frac{\Phi^{(t)}_{\mathbf
{p}\to v_{k,1}}}{P}\Big{)}$ for $a_{k,\mathbf{q}}\not=1$.

We also have the following lemmas to control the gradient updates for attention correlations.

###### Lemma F.2 .

If C.1 (or C.2) and F.1 hold for ${T_{c,1}}+1\leq t\leq{T^{\epsilon}_{c,1}}$, then we have

$$ $\displaystyle\alpha^{(t)}_{\mathbf{p}\to v_{k,1}}\geq\Omega\left({\epsilon} \right),|\alpha^{(t)}_{\mathbf{p}\to v_{k,m}}|\leq O\Big{(}\frac{\alpha^{(t)}_ {\mathbf{p}\to v_{k,1}}}{P^{1-\kappa_{s}}}\Big{)}\quad\text{ for all }m\not=1$ (F.4a) $\displaystyle\beta^{(t)}_{k,\mathbf{p}\to\mathbf{q}}=\Theta\Big{(}\frac{\alpha ^{(t)}_{\mathbf{p}\to v_{k,1}}}{C_{1}}\Big{)},\text{ for }a_{k,\mathbf{q}}=1, \mathbf{q}\not=\mathbf{p}$ (F.4b) $\displaystyle|\beta^{(t)}_{k,\mathbf{p}\to\mathbf{p}}|,|\beta^{(t)}_{k,\mathbf {p}\to\mathbf{q}}|=O\Big{(}\frac{\alpha^{(t)}_{\mathbf{p}\to v_{k,1}}}{P}\Big{ )}\quad\text{ for all }a_{k,\mathbf{q}}>1.$ (F.4c) $$

F.2 can be directly proved by Lemma F.2. Furthermore, at the end of this stage, we will have:

###### Lemma F.3 .

Suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$, then F.2 holds for all ${T_{c,1}}<t\leq T^{\epsilon}_{c,1}={T_{c,1}}+O\Big{(}\frac{\log(P\epsilon^{-1}
)}{\eta\epsilon}\Big{)}$, and at iteration $t=T^{\epsilon}_{c,1}+1$, we have

- 1.
$\widetilde{\mathcal{L}}_{k,\mathbf{p}}(Q^{T^{\epsilon}_{c,1}+1})<\frac{
\epsilon}{2K}$;
- 2.
If $\mathsf{M}\in\mathcal{E}_{k,1}$ , we have $\Big{(}1-\mathbf{Attn}_{\mathbf{p}\to\mathcal{P}_{k,1}}^{(T^{\epsilon}_{c,1}+1
)}\Big{)}^{2}\leq O(\epsilon)$.

###### Induction Hypothesis F.2 .

###### Lemma F.2 .

###### Lemma F.3 .

## Appendix G Proof of Main Theorems

### G.1 Proof of Induction Hypotheses

We are now ready to show C.1 (resp. C.2) holds through the learning process.

###### Theorem G.1 (Positive Information Gap) .

For sufficiently large $P>0$, $\eta\ll\log(P)$, $\Omega(1)\leq\Delta<1$, C.1 holds for all iterations $t=0,1,\cdots,T=O\Big{(}\frac{e^{\operatorname{polylog}(P)}}{\eta}\Big{)}$.

###### Theorem G.2 (Negative Information Gap) .

For sufficiently large $P>0$, $\eta\ll\log(P)$, $-0.5<\Delta\leq-\Omega(1)$, C.2 holds for all iterations $t=0,1,\cdots,T=O\Big{(}\frac{e^{\operatorname{polylog}(P)}}{\eta}\Big{)}$.

###### Theorem G.1 (Positive Information Gap) .

###### Theorem G.2 (Negative Information Gap) .

##### Proof of Theorem G.1 .

It is easy to verify C.1 holds at iteration $t=0$ due to the initialization $Q^{(0)}=\mathbf{0}_{d\times d}$. At iteration $t>0$:

- •
C.1a. can be proven by D.1-D.4 a and F.1-F.2 a, combining with the fact that $\log(1/\epsilon)\ll\operatorname{polylog}(P)$.
- •
C.1b. can be obtained by invoking D.1-D.4 b.
- •
C.1c. can be obtained by invoking D.1-D.4 c and F.1-F.2 b.
- •
To prove C.1d., for $\mathbf{q}\not=\mathbf{p}$, $\Upsilon_{\mathbf{p}\to\mathbf{q}}^{(t)}=\sum_{k=1}^{K}\Upsilon_{k,\mathbf{p}
\to\mathbf{q}}^{(t)}$. By item d-f in D.1-D.4 and item c-d in F.1-F.2, we can conclude that no matter the relative areas $\mathbf{q}$ and $\mathbf{p}$ belong to for a specific cluster, for all $k\in[K]$, throughout the entire learning process, the following upper bound always holds:
$\Upsilon_{k,\mathbf{p}\to\mathbf{q}}^{(t)}\leq\max_{t\in[T]}(|\Phi^{(t)}_{\bm{
p}\to v_{k,n}}|+|\Phi^{(t)}_{\bm{p}\to v_{k,1}}|)\max\Bigg{\{}O\Big{(}\frac{1}
{C_{1}}\Big{)},O\Big{(}\frac{1}{C_{n}}\Big{)},O\Big{(}\frac{1}{P}\Big{)}\Bigg{
\}}\leq\widetilde{O}\Big{(}\frac{1}{C_{n}}\Big{)}.$
Moreover, since $K=\Theta(1)$, we then have $\Upsilon_{\mathbf{p}\to\mathbf{q}}^{(t)}=\widetilde{O}(\frac{1}{C_{n}})$, which completes the proof.
- •
The proof for C.1d. is similar as before, by noticing that $\Upsilon_{k,\mathbf{p}\to\mathbf{p}}^{(t)}=\widetilde{O}(\frac{1}{P})$ for each $k\in[K]$, which is due to D.1-D.4 d and F.1-F.2 c.

The proof of Theorem G.2 mirrors that of Theorem G.1, with the only difference being the substitution of relevant sections with C.2. For the sake of brevity, this part of the proof is not reiterated here.

### G.2 Proof of Theorem 4.1 and Theorem 4.2 with Positive Information Gap

###### Theorem G.3 .

Suppose $\Omega(1)\leq\Delta\leq 1$. For any $0<\epsilon<1$, suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$. We apply GD to train the loss function given in (2.9) with $\eta\ll\operatorname{poly}(P)$.
Then for each $\mathbf{p}\in\mathcal{P}$, we have

- 1.
The loss converges: after $T^{\star}=O\Big{(}\frac{\log(P)P^{\max\{2(\frac{U}{L}-1),1\}(1-\kappa_{s})}}{
\eta}+\frac{\log(P\epsilon^{-1})}{\eta\epsilon}\Big{)}$ iterations, ${\mathcal{L}}_{\mathbf{p}}(Q^{(T^{\star})})-{\mathcal{L}}_{\mathbf{p}}^{*}\leq\epsilon$, where ${\mathcal{L}}_{\mathbf{p}}^{\star}$ is the global minimum of patch-level construction loss in (4.2).
- 2.
Attention score concentrates: given cluster $k\in[K]$, if $X_{\mathbf{p}}$ is masked, then
the one-layer transformer nearly “pays all attention” to all unmasked patches in the same area $\mathcal{P}_{k,a_{k,\mathbf{p}}}$, i.e., $\Big{(}1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}^{(T^{
\star})}\Big{)}^{2}\leq O(\epsilon)$.
- 3.
Local area learning feature attention correlation through two-phase: given $k\in[K]$, if $a_{k,\mathbf{p}}>1$, then we have
(a)
$\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ first quickly decrease with all other $\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}$ , $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}$ not changing much;
(b)
after some point, the increase of $\Phi^{(t)}_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}$ takes dominance. Such $\Phi^{(t)}_{\mathbf{p}\to v_{k,a_{k,\mathbf{p}}}}$ will keep growing until convergence with all other feature and positional attention correlations nearly unchanged.
- 4.
Global area learning feature attention correlation through one-phase: given $k\in[K]$, if $a_{k,\mathbf{p}}=1$, throughout the training, the increase of $\Phi^{(t)}_{\mathbf{p}\to v_{k,1}}$ dominates, whereas all $A^{(t)}_{1,m}$ with $m\not=1$ and position attention correlations remain close to 0 0.

###### Proof.

The first statement is obtained by letting $T^{\star}=\max\{T_{2}^{\epsilon},T_{c,1}^{\epsilon}\}+1$ in Lemma D.30 and Lemma F.3, combining wth Lemma B.9 and Lemma B.10, which lead to

$$ $\displaystyle{\mathcal{L}}_{\mathbf{p}}(Q^{(T^{\star})})-{\mathcal{L}}_{ \mathbf{p}}^{*}$ $\displaystyle\leq{\mathcal{L}}_{\mathbf{p}}(Q^{(T^{\star})})-\mathcal{L}^{ \text{low}}_{\mathbf{p}}$ $\displaystyle\leq\widetilde{\mathcal{L}}_{\mathbf{p}}(Q^{T^{\star}})+O\Big{(} \exp\Big{(}-\big{(}c_{3}P^{\kappa_{c}}+\operatorname{\mathds{1}}\big{\{}1\not \in\cup_{k\in[K]}\{a_{k,\mathbf{p}}\}\big{\}}c_{4}P^{\kappa_{s}}\big{)}\Big{)} \Big{)}$ $\displaystyle\leq K\cdot\frac{\epsilon}{2K}+O\Big{(}\exp\Big{(}-\big{(}c_{3}P^ {\kappa_{c}}+\operatorname{\mathds{1}}\big{\{}1\not\in\cup_{k\in[K]}\{a_{k, \mathbf{p}}\}\big{\}}c_{4}P^{\kappa_{s}}\big{)}\Big{)}\Big{)}$ $\displaystyle<\epsilon.$ $$

The second statement follows from Lemma D.30 and Lemma F.3. The third and fourth statements directly follow from the learning process described in Appendix D and Appendix F when C.1 holds.
∎

###### Theorem G.3 .

###### Proof.

### G.3 Proof of Theorem 4.1 and Theorem 4.2 with Negative Information Gap

###### Theorem G.4 .

Suppose $-0.5\leq\Delta\leq\Omega(1)$. For any $0<\epsilon<1$, suppose $\operatorname{polylog}(P)\gg\log(\frac{1}{\epsilon})$. We apply GD to train the loss function given in (2.9) with $\eta\ll\operatorname{poly}(P)$.
Then for each $\mathbf{p}\in\mathcal{P}$, we have

- 1.
The loss converges: after $T^{\star}=O\Big{(}\frac{\log(P)P^{\max\{2(\frac{U}{L}-1),1\}(1-\kappa_{s})}}{
\eta}+\frac{\log(P\epsilon^{-1})}{\eta\epsilon}\Big{)}$ iterations, ${\mathcal{L}}_{\mathbf{p}}(Q^{(T^{\star})})-{\mathcal{L}}_{\mathbf{p}}^{*}\leq\epsilon$, where ${\mathcal{L}}_{\mathbf{p}}^{\star}$ is the global minimum of patch-level construction loss in (4.2).
- 2.
Attention score concentrates: given cluster $k\in[K]$, if $X_{\mathbf{p}}$ is masked, then
the one-layer transformer nearly “pays all attention” to all unmasked patches in the same area $\mathcal{P}_{k,a_{k,\mathbf{p}}}$, i.e., $\Big{(}1-\mathbf{Attn}_{{\mathbf{p}}\to\mathcal{P}_{k,a_{k,\mathbf{p}}}}^{(T^{
\star})}\Big{)}^{2}\leq O(\epsilon)$.
- 3.
All areas learning feature attention correlation through one-phase: given $k\in[K]$, throughout the training, the increase of $\Phi^{(t)}_{\mathbf{p}\to v_{k,a_{k,\bm{p}}}}$ dominates, whereas all $\Phi^{(t)}_{\mathbf{p}\to v_{k,m}}$ with $m\not=1$ and position attention correlations $\Upsilon^{(t)}_{\mathbf{p}\to\mathbf{q}}$ remain close to 0 0.

###### Proof.

The first statement is obtained by letting $T^{\star}=\max\{T_{\mathrm{neg},1}^{\epsilon},T_{c,1}^{\epsilon}\}+1$ in Lemma E.3 and Lemma F.3, combining wth Lemma B.9 and Lemma B.10, which lead to

$$ $\displaystyle{\mathcal{L}}_{\mathbf{p}}(Q^{(T^{\star})})-{\mathcal{L}}_{ \mathbf{p}}^{*}$ $\displaystyle\leq{\mathcal{L}}_{\mathbf{p}}(Q^{(T^{\star})})-\mathcal{L}^{ \text{low}}_{\mathbf{p}}$ $\displaystyle\leq\widetilde{\mathcal{L}}_{\mathbf{p}}(Q^{T^{\star}})+O\Big{(} \exp\Big{(}-\big{(}c_{3}P^{\kappa_{c}}+\operatorname{\mathds{1}}\big{\{}1\not \in\cup_{k\in[K]}\{a_{k,\mathbf{p}}\}\big{\}}c_{4}P^{\kappa_{s}}\big{)}\Big{)} \Big{)}$ $\displaystyle\leq K\cdot\frac{\epsilon}{2K}+O\Big{(}\exp\Big{(}-\big{(}c_{3}P^ {\kappa_{c}}+\operatorname{\mathds{1}}\big{\{}1\not\in\cup_{k\in[K]}\{a_{k, \mathbf{p}}\}\big{\}}c_{4}P^{\kappa_{s}}\big{)}\Big{)}\Big{)}$ $\displaystyle<\epsilon.$ $$

The second statement follows from Lemma E.3 and Lemma F.3. The third and fourth statements directly follow from the learning process described in Appendix E and Appendix F when C.2 holds.
∎

###### Theorem G.4 .

###### Proof.

## Peer Reviews

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

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

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - The theory particularly for CL seems to be build on the notion of dominant concept which takes most of the area of the input. This assumption is not valid for almost all practical datasets for SSL. For instance, the average object to total image area ratio is 0.35 in ImageNet. Which means the so called background is almost always going to be dominant concept according this paper's definition. In that case the attention from any part of the image should focus on the background most of the time. If this is the case the kNN evaluation of DINO and MoCo should not be so high for imagenet and more importantly for other datasets. 
- Is the proposed theory would still be valid for datasets where the global area is small? 
- What will happen if the local areas are larger than the global area? This situation is going to be the dominant case in most practical datasets.  
Minors: 
- what is $k_s$? It seems the sum of weights for global area.
- What $\Delta$ means? Can you explain in details, such as when $\Delta=-1,0,1$?
soundness: 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.
2. This work attempt to answers the empirical observations for CL and MIM with ViTs: the MIM captures diverse local patterns, while the CL focuses on the dominant patterns.
3. Base on the proposed theory, this paper finds that there are two phases to learn the feature-position correlations when information gap is positive.
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.
weaknesses: 1. Vision transformers and transformers in general consists of two main blocks multi-head self-attention block as well as so called FFN consisting of two point-wise convolutions. The discussion is based on only one-layer of self-attention (with over simplification to only one weight matrix for self-attention) within vision transformer. Thus it cannot fully explain the complex behaviour of ViTs. 
2. There are a lot of assumptions, such as the input data distribution and positional encoding. It is not clear the differences in real problems.
3. What is the insight from this work to develop new methods? How to prevent the attention collapse in CL?
4. The discussion is based on SGD. However, Adam is the popular optimizer for ViTs. 
5. The original MIM (first MIM work with ViTs which outperformed supervised pretraining) work SiT [1] from which idea of MAE [4] is derived/copied from (only MIM part of SiT is taken in MAE) combines both MIM and CL, a strategy proven to be better and more principled time and again. Furthermore, state-of-the-art methods like iBoT [2] and its extension DINOv2 [3] also follow combined MIM and CL. It would be more fruitful to build theoretical framework for this combined MIM and CL setting (which seems to be the way forward in SSL) rather than building framework for MIM and CL separately. 

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

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 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. 
2. The paper effectively illustrates the differences in feature capture between MAE and CL, showing how MAE learns both global and local features, while CL focuses more on global features. 
3. Through a detailed gradient descent analysis of a single-layer transformer, the study offers concrete findings on how MAE and CL converge to distinct attention patterns, ensuring result reliability and reproducibility. The study focuses on a single-layer transformer model, which may overlook the dynamics of self-supervised learning in deeper or more complex networks.
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.
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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. In the first paragraph of Sec. 2.1 about **Masked reconstruction-based learning**, should it be $\mathcal{M}_i \subset \mathcal{P}$ instead of $\mathcal{M}_i \subset [P]$?
2. Does the same cluster exist in different $\mathcal{D}_k$ defined in Sec. 2.2 and depicted in Fig.2? If this is true, should this assumption be a little bit overwhelming? For CL, only entities obtained from the same images are considered positive. As for MAE, only local features are important. I am sure changing this assumption won't affect the proof, but I want to clarify this.
3. From the proof, can we determine the required ratio between positive and negative samples to ensure the global convergence properties? Since approaches such as MoCo require more negative samples to prevent collapsing.
soundness: 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.
2. The theoretical analysis provides a solid step to understanding the detailed differences between the mechanisms and can help future researchers develop more approaches based on this analysis.
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.
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.

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

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

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

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

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

Decision: Accept (Poster)