Title: Unveil Benign Overfitting for Transformer in Vision: Training Dynamics, Convergence, and Generalization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Vision Transformers

Training Dynamics

Benign Overfitting

Self-Attention

Feature Learning
Abstract: Transformers have demonstrated great power in the recent development of large foundational models. In particular, the Vision Transformer (ViT) has brought revolutionary changes to the field of vision, achieving significant accomplishments on the experimental side. However, their theoretical capabilities, particularly in terms of generalization when trained to overfit training data, are still not fully understood. To address this gap, this work delves deeply into the \textit{benign overfitting} perspective of transformers in vision. To this end, we study the optimization of a Transformer composed of a self-attention layer with softmax followed by a fully connected layer under gradient descent on a certain data distribution model. By developing techniques that address the challenges posed by softmax and the interdependent nature of multiple weights in transformer optimization, we successfully characterized the training dynamics and achieved generalization in post-training. Our results establish a sharp condition that can distinguish between the small test error phase and the large test error regime, based on the signal-to-noise ratio in the data model. The theoretical results are further verified by experimental simulation. To the best of our knowledge, this is the first work to characterize benign overfitting for Transformers.

Unveil Benign Overfitting for Transformer in Vision: Training Dynamics, Convergence, and Generalization

Abstract

Transformers have demonstrated great power in the recent development of large foundational models. In particular, the Vision Transformer (ViT) has brought revolutionary changes to the field of vision, achieving significant accomplishments on the experimental side. However, their theoretical capabilities, particularly in terms of generalization when trained to overfit training data, are still not fully understood. To address this gap, this work delves deeply into the benign overfitting perspective of transformers in vision. To this end, we study the optimization of a Transformer composed of a self-attention layer with softmax followed by a fully connected layer under gradient descent on a certain data distribution model. By developing techniques that address the challenges posed by softmax and the interdependent nature of multiple weights in transformer optimization, we successfully characterized the training dynamics and achieved generalization in post-training. Our results establish a sharp condition that can distinguish between the small test error phase and the large test error regime, based on the signal-to-noise ratio in the data model. The theoretical results are further verified by experimental simulation. To the best of our knowledge, this is the first work to characterize benign overfitting for Transformers.

1 Introduction

Transformers (Vaswani et al., 2017) have revolutionized numerous fields in artificial intelligence, ranging from natural language processing (Devlin et al., 2018; OpenAI, 2023), computer vision Dosovitskiy et al. (2020); He et al. (2022), graph Yun et al. (2019); Hu et al. (2020) to reinforcement learning Chen et al. (2021a); Janner et al. (2021). In particular, Vision Transformer (ViT) Dosovitskiy et al. (2020) has been developed to advancing the computer vision tasks (Liu et al., 2021; Khan et al., 2022; Guo et al., 2022; Tolstikhin et al., 2021; Lin et al., 2014) compared to Convolutional neural networks. Since then, with the high performance, ViT emerged as a hot area of research and application. Numerous technologies and methods began to surface, including enhancing ViT's data and resource efficiency (Touvron et al., 2021; Hassani et al., 2021; Chen et al., 2021b; Touvron et al., 2022), reducing heavy computation cost Chen et al. (2022); Zhu et al. (2024); Wang et al. (2023).

Recent success brought by ViT has been inspiring an increasing number of works to understand the ViT through empirical and theoretical studies. The empirical investigation works study the robustness (Bhojanapalli et al., 2021; Naseer et al., 2021; Paul and Chen, 2022; Bai et al., 2021); and the role in

self-supervision (Caron et al., 2021; Chen et al., 2021c; Oquab et al., 2023) and others. On the other hand side, the theoretical works have been conducted to understand the ViT from the perspective of expressivity (Vuckovic et al., 2020; Edelman et al., 2022), optimization (Zhang et al., 2020; Tian et al., 2023a,b), generalization (Li et al., 2023a), and in-context learning (Huang et al., 2023c; Garg et al., 2022; Bai et al., 2024; Ahn et al., 2024).

Despite the insightful understanding provided by the above investigations, it remains unclear how ViTs generalize to unseen data when they are trained to overfit the training data. In particular, the conditions under which ViTs can exhibit benign overfitting (Bartlett et al., 2020), where the test error remains small despite overfitting to the training data, are not well understood. Moreover, prior work on the optimization and generalization of ViTs often focuses on simplified settings, such as linear transformers or unrealistic loss functions, due to technical limitations. In this work, we aim to fill this gap through a feature learning theory (Allen-Zhu and Li, 2020; Cao et al., 2022). Based on a data generative model with M tokens composed of signal tokens and noise tokens and a two-layer ViT with softmax attention, we characterize the training dynamics from a random initialization and the generalization ability of the ViT after convergence. We establish a sharp separation in condition to distinct the benign overfitting and harmful overfitting regime for ViTs. Our contributions are summarized as follows:

We successfully characterize the optimization of the ViT through three different phases in dynamics that exhibit rich and unique behaviors related to self-attention. Building on the convergence results, we further characterize the generalization bounds on unseen datasets.

Technically, we develop novel methods to handle non-linear attention and the inter-dependent nature of multiple weights in transformer optimization when training from scratch.

We establish a sharp separation condition in the signal-to-noise ratio in data to distinguish the benign overfitting and harmful overfitting regimes of ViTs. This separation is then verified by experimental simulation.

2 Related Work

Benign Overfiting in Deep Learning The most prominent behavior of deep learning is its breaking of bias-variance trade-off in statistical learning theory, i.e., good generalization on unseen data even with overfitting on training data. This line of research starts from Bartlett et al. (2020) who studied benign overfitting in linear regression, learning data generated by a linear model with additive Gaussian noises. It is shown that at large input dimensions (which leads to over-parameterization), the excessive risk of the interpolation can be asymptotically optimal. Hastie et al. (2022); Wu and Xu (2020) studied linear regression when both the dimension and the number of samples scale together with a fixed ratio. Timor et al. (2023) proposed a unified data model for linear predictors and studied conditions under which benign overfitting occurs in different problems. Besides the studies on linear models (Wang et al., 2021; Zou et al., 2021; Zhou and Ge, 2023), several recent works tried to study benign overfitting in neural networks (Frei et al., 2022; Kornowski et al., 2024; Chatterji et al., 2022; Xu et al., 2023). In particular, Li et al. (2021) investigated the benign overfitting in two-layer neural networks with the first layer parameters fixed at random initialization. Furthermore, benign overfitting was characterized in convolution neural networks (Cao et al., 2022; Kou et al., 2023), OOD (Chen et al., 2024), federated learning (Huang et al., 2023b), graph neural network (Huang et al., 2023a) by feature learning theory (Allen-Zhu and Li, 2020; Cao et al., 2022). Unlike existing research on benign overfitting, this work focuses on ViTs.

Optimization of Transformers Towards understanding the optimization of Transformers, Zhang et al. (2020) provided the analysis for adaptive gradient methods. Jelassi et al. (2022) proposed a spatially structured dataset and a simplified ViT model and showed that ViT implicitly learns the spatial structure of the dataset while generalizing. Besides, Li et al. (2023b) provided fine-grained mechanistic understanding of how transformers learn 'semantic structure' for BERT-like framework (Devlin et al., 2018). Furthermore, Tian et al. (2023a,b) characterizes the SGD training dynamics of a 1-layer Transformer and multi-layer Transformer respectively. They focused on the unique phenomena and the role of attention in the training dynamics. Huang et al. (2023c) studied the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent in order to in-context learn linear function classes. In addition, Li et al. (2024) provided a theoretical

analysis of the training dynamics of Transformers with nonlinear self-attention within In-Context Learning. In particular, Li et al. (2023a) is most relevant to us, as they also study the training dynamics of ViTs with a similar data model. However, they considered hinge loss and a specific initialization to simplify analysis and did not characterize the harmful overfitting regime. In summary, while the above work studies the optimization dynamics of attention-based models, they do not characterize benign overfitting as we investigate.

3 Problem Setup

In this section, we outline the data generation model, the Vision Transformer model, and the gradient descent training algorithm.

Notations. For two sequences { x n } and { y n } , we denote x n = O ( y n ) if there exist some absolute constant C > 0 and N > 0 such that | x n | ≤ C | y n | for all n ≥ N , denote x n = Ω( y n ) if y n = O ( x n ) , and denote x n = Θ( y n ) if x n = O ( y n ) and x n = Ω( y n ) . We also denote x n = o ( y n ) if lim | x n /y n | = 0 . Finally, we use ˜ O ( · ) , ˜ Ω( · ) and ˜ Θ( · ) to hide logarithmic factors in these notations respectively.

Definition 3.1 ( Data Generation Model ) . Let µ + , µ -∈ R d be fixed vectors representing the signals contained in data points, where ∥ µ + ∥ 2 = ∥ µ -∥ 2 = ∥ µ ∥ 2 and ⟨ µ + , µ -⟩ = 0 . Then each data point ( X , y ) with X = ( x 1 , x 2 , . . . , x M ) ⊤ ∈ R M × d and y ∈ {-1 , 1 } is generated from the following distribution D :

The label y is generated as a Rademacher random variable.

If y = 1 then x 1 is given as µ + , if y = -1 then x 1 is given as µ -, which represents signals.

A noise vector ξ 2 is generated from the Gaussian distribution N (0 , ˜ σ 2 p · ( I -µ + µ ⊤ + · ∥ µ ∥ -2 2 -µ -µ ⊤ -· ∥ µ ∥ -2 2 )) .

Noise vectors ξ 3 , . . . , ξ M is generated i.i.d from the Gaussian distribution N (0 , σ 2 p · ( I -µ + µ ⊤ + · ∥ µ ∥ -2 2 -µ -µ ⊤ -· ∥ µ ∥ -2 2 )) .

x 2 , . . . , x M are given by ξ 2 , . . . , ξ M , which represent noises.

Our data model is envisioned by considering M tokens of the data points: x 1 , x 2 , . . . , x M , which can be seen as patches derived from image data. x 1 embodies the signal that is inherently linked to the data's class label, while x 2 , x 3 , . . . , x M represents the noise, which is not associated with the label. For simplicity, we assume that the noise patches is independently drawn from Gaussian distribution N (0 , ˜ σ 2 p · ( I -µ + µ ⊤ + ·∥ µ ∥ -2 2 -µ -µ ⊤ -·∥ µ ∥ -2 2 )) and N (0 , σ 2 p · ( I -µ + µ ⊤ + ·∥ µ ∥ -2 2 -µ -µ ⊤ -·∥ µ ∥ -2 2 )) , ensuring the noise vectors remains orthogonal to the signal vectors µ + and µ -. Also, if ˜ σ p is sufficiently larger than σ p , the input tokens will become sparse, which is consistent with the empirical observation that the attention and activation values in Transformer-based models are usually sparse (Child et al., 2019; Robinson et al., 2023). We denote SNR = ∥ µ ∥ 2 / ( σ p √ d ) to represent the signal-to-noise ratio.

Two-layer Transformer. We consider a two layer Transformer network with a self-attention layer and a fixed linear layer, which is defined as:

Here, φ ( · ) : R M → R M denote the softmax function, W Q , W K ∈ R d × d h , W V ∈ R d × d v denote the query matrix, key matrix, and value matrix, respectively, and w O ∈ R d v denote the weight for the linear layer. We use θ to denote the collection of all the model weights. This model is a simplified version of the ViT model from Dosovitskiy et al. (2020), making our analysis focus on the self-attention mechanism which is the most critical component of the ViTs.

Given a training data set S = { ( X n , y n ) } N n =1 generated from the distribution D defined in Definition 3.1, where the subscript n represents the n -th sample, we train the two-layer Transformer by

minimizing the empirical cross-entropy loss function:

where ℓ ( z ) = log(1 + exp( -z )) and f ( X , θ ) is the two-layer Transformer. We further defined the population loss (test loss) L D ( θ ) := E ( X ,y ) ∼ D ℓ ( yf ( X , θ )) .

Training algorithm. We consider Gaussian initialization for the network weights W Q , W K and W V , where each entry of W Q and W K is sampled from a Gaussian distribution N (0 , σ 2 h ) , and W V is sampled from N (0 , σ 2 V ) at initialization. We use gradient descent to optimize training loss L S ( θ ) , and the update of W Q , W K and W V can be written as follows:

where W ( t ) can be designated as W Q , W K or W V .

Recalling the ViT defined in Eq. (1), we can intuitively analyze its training dynamics and generalization: if more attention is paid to signals, i.e., x ⊤ l W Q W ⊤ K µ ± ≥ x ⊤ l W Q W ⊤ K ξ i for i ∈ [ M ] / { 1 } , then the vector φ ( x ⊤ l W Q W ⊤ K X ⊤ ) X has a higher similarity with signals µ ± rather than with noises ξ i . In turn, W V together with w O have more chance to learn signals µ ± and utilize them to make prediction. At the same time, if vector W V w O learns signals µ ± more than noises ξ i , i.e., µ ± W V w O ≥ ξ i W V w O , then during the gradient descent training, the gradient send to x ⊤ l W Q W ⊤ K µ ± will be larger than that send to x ⊤ l W Q W ⊤ K ξ i . As a result, more and more attention is paid on signals and the similarity between signals µ ± and W V w O becomes increasingly high. Therefore, the Transformer can perform well on new data points. On the contrary, if much attention is paid to some noises in the training dataset and the model utilizes them for making prediction and optimizing the training loss, then the model can fit the training dataset well but might not perform well on the new test data.

4 Main Results

This section presents our main theoretical results which characterize the convergence and generalization of the ViT model under different sample size N and signal-to-noise ratio SNR = ∥ µ ∥ 2 / ( σ p √ d ) . The results are based on the following conditions.

Condition 4.1. Given a sufficiently small failure probability δ > 0 and a target training loss ϵ > 0 , suppose that:(1) Dimension d h = ˜ Ω 1 ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . (2) Dimension d = ˜ Ω ( ϵ -2 N 2 d h ) . (3) Training sample size N = Ω(polylog( d )) .(4) The number of input tokens M = Θ(1) . (5) The ℓ 2 -norm of linear layer weights ∥ w O ∥ 2 = Θ(1) . (6) The learning rate η ≤ ˜ O (min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h ) . (7) The standard deviation of Gaussian initialization σ V satisfies: σ V ≤ ˜ O ( ∥ w O ∥ -1 2 · min {∥ µ ∥ -1 2 , ( σ p √ d ) -1 } · d -1 4 h ) . (8) The variance of Gaussian initialization σ 2 h satisfies: min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -2 ≤ σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -3 2 . (9) Target training loss ϵ ≤ O (1 / polylog( d )) . (10) The relationship between σ p and ˜ σ p satisfies ˜ σ p = C p σ p and C p = 5 √ M .

Conditions (1) and (2) ensure an over-parametrized learning setting, and similar conditions have been made in the theoretical analysis of CNN models (Cao et al., 2022; Kou et al., 2023). Condition (3) ensures that there are enough samples in each class with high probability. Conditions (4)-(5) are intended to simplify the calculation, and can be easily generalized to M = Ω(1) , ∥ w O ∥ 2 = o (1) or ∥ w O ∥ 2 = ω (1) setting. Conditions (6)-(8) ensure that the Transformer can be effectively trained well. Condition (9) ensures that the Transformer is sufficiently overfitting the training data. Condition (10) ensures the sparsity of the input token features.

Theorem 4.1 ( Benign Overfiting ) . Under Condition 4.1, if N · SNR 2 = Ω(1) , then with probability at least 1 -d -1 , there exist T = Θ( η -1 ϵ -1 ∥ µ ∥ -2 2 ∥ w O ∥ -2 2 ) such that:

The training loss converges to ϵ : L S ( θ ( T )) ≤ ϵ .

The test loss is nearly zero: L D ( θ ( T )) ≤ o (1) .

Theorem 4.1 characterizes the case of benign overfitting . It shows that as long as N · SNR 2 = Ω(1) , the Transformer can generalize well, even though it overfit the training data. To complement the above result and highlight the sharpness of this condition, we present the following theorem for the regime of harmful overfitting .

Theorem 4.2 ( Harmful Overfiting ) . Under Condition 4.1, if N -1 · SNR -2 = Ω(1) , then with probability at least 1 -d -1 , there exist T = Θ( Nη -1 ϵ -1 σ -2 p d -1 ∥ w O ∥ -2 2 ) such that:

The training loss converges to ϵ : L S ( θ ( T )) ≤ ϵ .

The test loss is high: L D ( θ ( T )) = Θ(1) .

Theorem 4.2 shows that if N -1 · SNR -2 = Ω(1) , the trained Transformer has a high test loss as a result of overfitting the noises in the training data. Theorem 4.1 and Theorem 4.2 reveals a sharp phase transition between benign overfitting and harmful overfitting .

Comparison with other results. Firstly, compared with CNNs with ReLU q network result (Cao et al., 2022), when signal-to-noise ratio is small ( SNR ≤ 1 ), our results show that ViTs require less number of samples to gengeralize well, which reflects the advantage of Transformers. Secondly, previous theoretical benign overfitting analysis of CNNs with ReLU activations rely heavily on signal strength ∥ µ ∥ 2 and signal-to-noise ratio SNR (Kou et al., 2023; Meng et al., 2023), i.e., SNR 2 = ˜ O (1 /N ) and require ∥ µ ∥ 2 to be large enough, while our analysis does not need to impose restrictions on these conditions. Thirdly, previous works only show when Transformer can generalize well (Li et al., 2023a; Deora et al., 2023), while we demonstrate harmful overfitting as a complementary, which reflect that our benign overfitting condition is tighter and more precise.

5 Proof Sketch

This section discusses our main challenges in studying ViT's training dynamic, and presents our technical solutions to overcoming them. The complete proofs are given in the appendix.

5.1 Vectorized Q & K and scalarized V

Our first main challenge is to deal with the three matrices W Q , W K and W V . In contrast to CNN models, whose convolutional kernels can be treated as vectors and thus allow for a more straightforward analytical approach (Cao et al., 2022; Kou et al., 2023), the QKV matrices within the Transformer model are inherently complex to analyze. Moreover, their mutual interactions further complicate the analysis. To circumvent this complexity, Oymak et al. (2023) and Tian et al. (2023a) merge the key-query weights(e.g. W := W Q W ⊤ K ), Huang et al. (2023c); Zhang et al. (2024); Jelassi et al. (2022) employ a specific initialization (e.g. W (0) Q = W (0) K = I ). Although their approach simplifies the analysis, it is not conducive to showing the dynamics of QKV interactions, i.e., how W Q , W K , and W V affect each other.

In order to simplify the analysis of QKV dynamics without losing rigor, we propose key techniques called Vectorized Q & K and scalarized V . The basic idea comes from the property that the product of token feature vectors with the QKV matrices results in vectors, e.g., µ ⊤ + W Q , ξ ⊤ 2 W Q , etc., which are more amenable to analysis. Further, each entry of matrix XW Q W ⊤ K X ⊤ can be regarded as the inner product of two vectors, e.g., µ ⊤ + W Q W ⊤ K µ + = ⟨ µ ⊤ + W Q , µ ⊤ + W K ⟩ . Therefore, the dynamics of attention can be studied by analyzing the dynamics of the vectorized Q & K defined as follows:

Definition 5.1 (Vectorized Q & K) . Let W ( t ) Q and W ( t ) K be the QK matrices of the ViT at the t-th iteration of gradient descent. Then we define the vectorized Q and vectorized K as follows

With Definition 5.1, denoting S + := { n ∈ [ N ] : y n = 1 } , S -:= { n ∈ [ N ] : y n = -1 } , we further analyze the dynamics of the vectorized Q & K . By carefully computing the dynamic of q ( t ) + , we have

where φ ( t ) n,i := φ ( x ⊤ n,i W ( t ) Q W ( t ) ⊤ K X ⊤ n ) is a shorthand notation, and w ⊤ O W ( t ) ⊤ V X ⊤ n and X n W ( t ) K can be viewed in the following forms

Thus q ( t +1) + -q ( t ) + can be decomposed into a linear combination of k ( t ) + and k ( t ) n,i . Therefore, Eq. (2) can be further expanded as q ( t +1) + -q ( t ) + = α ( t ) + , + k ( t ) + + ∑ n ∈ S + ∑ M i =2 α ( t ) n, + ,i k ( t ) n,i , where ( α ( t ) + , + , α ( t ) n, + , 2 , . . . , α ( t ) n, + ,M ) = η NM ∑ n ∈ S + -ℓ ′ ( t ) n ∥ µ ∥ 2 2 w ⊤ O W ( t ) ⊤ V X ⊤ n ( diag ( φ ( t ) n, 1 ) -φ ( t ) ⊤ n, 1 φ ( t ) n, 1 ) . In- tuitively, if α ( t ) + , + is larger enough than α ( t ) n, + ,i , then ⟨ q ( t ) + , k ( t ) + ⟩ will grow faster than ⟨ q ( t ) + , k ( t ) n,i ⟩ , which means token µ + will pay more attention to µ + rather than ξ n,i . The dynamics of q ( t ) -, q ( t ) n,i , k ( t ) + , k ( t ) -and k ( t ) n,i are similar to q ( t ) + , and we can study the dynamics of QK matries by analysing the linear combination coefficients such as α ( t ) + , + , α ( t ) n, + ,i .

As x ⊤ n,i W ( t ) V w O are scalars and their dynamics contain a factor ∥ w O ∥ 2 2 , we define the scalarized V as follows.

Definition 5.2 (Scalarized V) . Let W ( t ) V be the V matrix of the ViT at the t-th iteration of gradient descent. Then there exist coefficients γ ( t ) V, + , γ ( t ) V, -, ρ ( t ) V,n,i such that

for i ∈ [ M ] \{ 1 } , n ∈ [ N ] .

With the Vectorized Q & K and scalarized V , one can simplify the study of the Transformer learning process to a meticulous calculation of the coefficients such as α , γ , ρ throughout the training period. So how to calculate the dynamics of these coefficients is a key point in our analysis. In the next subsection we describe how to handle softmax function and further give bounds for these coefficients.

5.2 Dealing with the softmax function

Our second challenge is to deal with the softmax function, which is the critical component and introduces non-linear transformation in our Transformer model.

As W Q and W K are within the softmax function and W V is outside the softmax function, we divide the dynamic of training process into two key aspects: (1) How W Q and W K affect W V ; (2) How W V affects W Q and W K . Next, we present our approach to addressing these two critical issues.

1. How W Q and W K affect W V : Without loss of generality, we take a data point ( X n , y n ) with y n = 1 as an example and provide the dynamics for γ ( t ) V, + and ρ ( t ) V,n,i in W V as follows: Lemma 5.1 (Dynamics of γ and ρ ) .

for i ∈ [ M ] \{ 1 } , n ∈ S + , where ℓ ′ ( t ) n := ℓ ′ ( y n f ( X n , θ ( t ))) is a shorthand notation.

In benign overfitting regime where N · SNR 2 = Ω(1) , the factor ∑ n ∈ S + η ∥ µ ∥ 2 2 NM is larger than 2 ηC 2 p σ 2 p d NM .

Therefore, as long as x ⊤ n,s W ( t ) Q W ( t ) ⊤ K µ + are not less than x ⊤ n,s W ( t ) Q W ( t ) ⊤ K ξ n,i , γ ( t ) V, + will grow faster than ρ ( t ) V,n,i . In other words, if q ( t ) ± and q ( t ) n,i prefer to align with k ( t ) ± rather than k ( t ) n,i , then W V would prefer to learn signals rather than memorize noises.

How W V affects W Q and W K : Recalling that ( α ( t ) + , + , α ( t ) n, + , 2 , . . . , α ( t ) n, + ,M ) = η NM ∑ n ∈ S + -ℓ ′ ( t ) n ∥ µ ∥ 2 2 w ⊤ O W ( t ) ⊤ V X ⊤ n ( diag ( φ ( t ) n, 1 ) -φ ( t ) ⊤ n, 1 φ ( t ) n, 1 ) , where the most complicated part

is the matrix ( diag ( φ ( t ) n, 1 ) -φ ( t ) ⊤ n, 1 φ ( t ) n, 1 ) . We observe that the matrix diag ( φ ( t ) n, 1 ) -φ ( t ) ⊤ n, 1 φ ( t ) n, 1 has two important properties:

The diagonal elements are positive, while the elements on the off-diagonal are negative.

The sum of each row and column of this matrix is 0.

Based on properties 1 and 2, we can deduce the following conclusion: as long as µ ⊤ + W ( t ) V w O is larger enough than ξ ⊤ n,i W ( t ) V w O , we have α ( t ) + , + ≥ 0 and α ( t ) n, + ,i ≤ 0 , implying that q ( t ) + prefers to align k ( t ) + rather than k ( t ) n,i .

In summary, if more attention is paid to signals, W V learns the signals faster than it memorizes the noises; conversely, if W V learns the signals significantly more than it memorizes the noises, more attention will be paid to signals rather than noises. In other words, W Q , W K and W V promote each other.

5.3 Three-Stage Decoupling

Our third main challenge is to deal with the complicated relation among the coefficients, e.g., γ , ρ , α . Inspired by the two-stage analysis utilized by Cao et al. (2022) to decouple the coefficients in CNN models, we analyze the ViT training process in three stage. In the following, we explain the key steps for proving Theorem 4.1. The proof for Theorem 4.2 is similar and we detail it in the appendix.

Stage 1: The analysis in 5.2 shows that W Q , W K and W V can promote each other. But this process of mutual reinforcement requires some conditions, e.g. µ ⊤ + W ( t ) V w O is sufficiently larger than ξ ⊤ n,i W ( t ) V w O which does not necessarily hold for Gaussian initialization, complicating the proof. Stage 1 was introduced to solve this problem, and Lemma 5.2 formally describes this stage:

Lemma 5.2 (V's Beginning of Learning Signals) . Under the same conditions as Theorem 4.1, there exist T 1 = 10 M (3 M +1) N ηd 1 4 h ( N ∥ µ ∥ 2 2 -60 M 2 C 2 p σ 2 p d ) ∥ w O ∥ 2 2 such that the first element of vector X n W ( t ) V w O dominate its other elements, that is, µ ⊤ + W ( t ) V w O ≥ 3 M · | ξ ⊤ n,i W ( t ) V w O | for all n ∈ S + , i ∈ [ M ] \{ 1 } and µ ⊤ -W ( t ) V w O ≤ -3 M · | ξ ⊤ n,i W ( t ) V w O | for all n ∈ S -, i ∈ [ M ] \{ 1 } .

With Lemma 5.2, we can guarantee the monotonicity of attention on signals for t ≥ T 1 , which allows us to concentrate on analyzing the growth rate of attention in the following stages.

Stage 2: Note that the output of softmax function has an upper bound 1, thus the output of ViT can be bounded by max {| µ ⊤ ± W ( t ) V w O | , | ξ ⊤ n,i W ( t ) V w O |} . By Lemma 5.1, it can be proved that there exists T 2 = Θ ( η -1 ∥ µ ∥ -2 2 ∥ w O ∥ -2 2 log(6 N 2 M 2 /δ ) -1 ) such that: max {| µ ⊤ ± W ( t ) V w O | , | ξ ⊤ n,i W ( t ) V w O |} = o (1) , and further get 1 / 2 -o (1) ≤ -ℓ ′ ( t ) n ≤ 1 / 2 + o (1) . Therefore, one can simplify the dynamics of the coefficients (e.g., γ , ρ , α ) by plugging the tight bounds of -ℓ ′ ( t ) n . The following lemma provides some bounds for the dynamics of W Q , W K and W V in Stage 2.

Lemma 5.3 (Dynamics of QKV in Stage 2) . Under Condition 4.1, if N · SNR 2 = Ω(1) , then with probability at least 1 -d -1 , there exist constant C , T 1 = Θ ( η -1 d -1 4 h ∥ µ ∥ -2 2 ∥ w O ∥ -2 2 ) and

Lemma 5.3 presents the training dynamics of W Q , W K and W V under benign overfitting regime in two aspects: direction and speed , that is,

direction: W V w O prefer to learn the signals rather than memorize the noises; more and more attention is paid to signals, while the attention on noises is decreasing.

speed: The inner products of µ ± and W V w O grow at a linear rate; the inner products of label-related vectorized Q & K grow at a logarithmic rate.

Stage 3: As the training process going on, the loss begins to converge, and the loss derivatives -ℓ ′ ( t ) n no longer remain near 1 / 2 . Therefore, the increasing rate of the inner products of vectorize Q &K and µ ⊤ ± W V w O begins to diminish. Based on the analysis in Stage 2, the attention is sufficiently sparse at the beginning of Stage 3. In other words, the attention on the signals is nearly 1, while the attention on the noises is nearly 0. According to this sparsity, µ ⊤ ± W ( t ) V w O will grow much faster than ξ ⊤ n,i W ( t ) V w O by Lemma 5.1, and Lemma 5.4 provides the lower bound of µ ⊤ ± W ( t ) V w O .

Lemma 5.4. Under the same conditions as Theorem 4.1, there exists T 3 = Θ ( η -1 ϵ -1 ∥ µ ∥ -2 2 ∥ w O ∥ -2 2 log(6 N 2 M 2 /δ ) -1 ) such that:

µ ⊤ + W ( t ) V w O ≥ log ( exp( µ ⊤ + W ( T 2 ) V w O ) + ηC ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ( t -T 2 ) ) , µ ⊤ -W ( t ) V w O ≤ -log ( exp( -µ ⊤ -W ( T 2 ) V w O ) + ηC ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ( t -T 2 ) ) for t ∈ [ T 2 , T 3 ] , where C is a constant, T 2 is the last iteration of stage 2.

Convergence: Lemma 5.4 provides logarithmic lower bounds for µ ⊤ ± W ( t ) V w O . Plugging t = T 3 into these inequality, we have | µ ⊤ ± W ( t ) V w O | ≥ log ( Θ(1 /ϵ ) ) , then as long as | ξ ⊤ n,i W ( t ) V w O | is sufficiently small, it can be proved that y n f ( X n , θ ( t )) ≥ log ( 1 /ϵ ) for all n ∈ [ N ] , thus ℓ ( y n f ( X n , θ ( t ))) = log(1 + exp( -log(1 /ϵ ))) ≤ ϵ and L S ( θ ( t )) ≤ ϵ . The convergence of ViT is accordingly obtained.

Generalization: Consider a new data point ( X , y ) generated from the distribution defined in Definition 3.1. Without loss of generality, we suppose that the signal token is µ + and the label is 1, i.e. X = ( µ + , ξ 2 , . . . , ξ M ) , y = 1 . It is clear that ξ ⊤ i W Q W ⊤ K µ ± has a mean zero, which implies that the attention on the signals in the test data may not necessarily be as high as that in the training data. However, the sparsity of attention during training process creates conditions for W V w O to utilize the signals to make predictions, facilitating the generalization of the ViT. We provide the lower bound for the output of ViT on the unseen data with high probability as follows:

Lemma 5.5. Under the same conditions as Theorem 4.1, there exists T 3 = Θ ( η -1 ϵ -1 ∥ µ ∥ -2 2 ∥ w O ∥ -2 2 log(6 N 2 M 2 /δ ) -1 ) , with probability at least 1 -δ/N 2 M ,

where C and C ′ are constants. Lemma 5.5 shows that yf ( X , θ ( T 3 )) is large with high probability, and through careful calculation, we can provide a small bound for test loss. More details are in Appendix D.4.

6 Experimental Verification

We present simulation results on synthetic data and MNIST dataset to verify our theoretical results.

Synthetic data experiments setting: We follow Definition 3.1 to generate the training set and test set. Specifically, we set token size M = 16 and feature dimension d = 1024 . Without loss of generality, we set µ + = ∥ µ ∥ 2 · [1 , 0 , . . . , 0] ⊤ and µ -= ∥ µ ∥ 2 · [0 , 1 , 0 , . . . , 0] ⊤ . We generate noise vector ξ 2 from the Gaussian distribution N (0 , ˜ σ 2 p I ) , where ˜ σ p is fixed to 4. Similarly, we generate the other noise vectors ξ i for i ∈ [ M ] / { 1 , 2 } from the Gaussian distribution N (0 , σ 2 p I ) where σ p is fixed to 0.2.

We consider a two-layer Transformer defined in Section 3. The dimensions of matrix W Q , W K and W V are set as d h = d v = 512 . The ViT parameters are initialized using PyTorch's default initialization method, and then they are divided by 16 to ensure that the weights are initialized small enough. We train the ViT with full-batch gradient descent and learning rate η = 0 . 1 , target training loss ϵ = 0 . 01 . We consider different traning sample size N ranging from 2 to 20, and different signal-to noise-ratio SNR ranging from 0.16 to 15.6. We evaluate the test loss with 100 test data points after training loss converges to ϵ . All experiments are performed on an NVIDIA A100 GPU.

N·SNR²=1000

0.7

20

0.6

18

16

0.5

Sample size N

14

12

0.4

10

0.3

8

0.2

6

4

0.1

2

0.16

5.0

10.0

15.0

Signal-to-Noise Ratio SNR

1.0

N·SNR²=1000

20

18

0.8

16

Sample size N

14

0.6

12

10

0.4

8

6

0.2

4

2

0.0

0.16

5.0

10.0

15.0

Signal-to-Noise RatioSNR

Synthetic data experiments results: Figure 1a shows that as N and SNR increase, the test loss tends to decrease. As Figure 1b shows, the theoretically derived red curve ( N · SNR 2 = 1000 )

almost follows the experimental dividing line between the yellow area (test loss > 0.2) and the purple area (test loss < 0.2). These experimental results further validate our theoretical results, that is, the sharp condition separation between benign and harmful overfitting, where N · SNR 2 = Ω(1) is a precise condition for benign overfitting and N -1 · SNR -2 = Ω(1) is a precise condition for harmful overfitting. More experimental results on the dynamics of QKV can be found in Appendix A.

(a) SNR=0.2

(b) SNR=2.0

(c) Test Loss Heatmap

1.2

7800

- 1.0

10

10

Sample size N

6600

0.8

15

15

5400 -

0.6

20

20

4200

0.4

25

25

3000

5

10

15

20

25

5

10

15

20

25

0.2

Signal-to-Noise Ratio SNR

0.6

1.0

1.4

1.8

MNIST experiments setting: We add Gaussian random noises to the outer regions of the images with a width of 4. The original images and the noises are multiplied by diferent factors to generate diferent dataset with specifc SNR. For example, in Figure 2 (a), we mutiply the original image by 1/6 and the noises by 5/6, thus SNR=0.2; in Figure 2 (b),we mutiply the original image by 2/3 and the noises by 1/3, thus SNR=2.0. We consider a ViT model that consists of two attention layers, each equipped with four self-attention heads, followed by a MLP with ReLU activation. The training sample size N ranges from 3000 to 8400 and the SNR ranges from 0.2 to 2.0.

MNIST experiments results: Figure 2 (c) shows the heatmap of test loss, which shows a transition between benign and harmful overfitting regimes. The larger the sample size N and signal-to-noise ratio SNR, the better the generalization performance.

7 Conclusion

This paper studies the training dynamics, convergence, and generalization for a two-layer Transformer in vision. By analyzing the Vectorized Q & K and scalarized V in three-stage decomposition and carefully handling the softmax function, we give the precise increasing rate of the attention and output of the Transformer in the training process, and further analyze its generalization performance. Our theoretical results reveal a sharp condition separation between benign and harmful overfitting. One limitation of our methodology is our proof technique relies heavily on the sparsity of feature strength, e.g., we assume the standard deviation ˜ σ p is sufficiently larger than σ p , ensuring that more attention is paid to ξ 2 rather than other tokens. One future direction is to generalize our analysis to study other abilities of Transformers, such as in-context learning, prompt-tuning and time series forecasting.

References

Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. (2024). Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems , 36.

Allen-Zhu, Z. and Li, Y. (2020). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816 .

Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. (2024). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems , 36.

Bai, Y., Mei, J., Yuille, A. L., and Xie, C. (2021). Are transformers more robust than cnns? Advances in neural information processing systems , 34:26831-26843.

Bartlett, P. L., Long, P. M., Lugosi, G., and Tsigler, A. (2020). Benign overfitting in linear regression. Proceedings of the National Academy of Sciences , 117(48):30063-30070.

Bhojanapalli, S., Chakrabarti, A., Glasner, D., Li, D., Unterthiner, T., and Veit, A. (2021). Understanding robustness of transformers for image classification. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10231-10241.

Cao, Y., Chen, Z., Belkin, M., and Gu, Q. (2022). Benign overfitting in two-layer convolutional neural networks. Advances in neural information processing systems , 35:25237-25250.

Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., and Joulin, A. (2021). Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision , pages 9650-9660.

Chatterji, N. S., Long, P. M., and Bartlett, P. L. (2022). The interplay between implicit bias and benign overfitting in two-layer linear networks. Journal of machine learning research , 23(263):1-48.

Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. (2021a). Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097.

Chen, T., Cheng, Y., Gan, Z., Yuan, L., Zhang, L., and Wang, Z. (2021b). Chasing sparsity in vision transformers: An end-to-end exploration. Advances in Neural Information Processing Systems , 34:19974-19988.

Chen, W., Huang, W., Du, X., Song, X., Wang, Z., and Zhou, D. (2022). Auto-scaling vision transformers without training. arXiv preprint arXiv:2202.11921 .

Chen, X., Xie, S., and He, K. (2021c). An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision , pages 9640-9649.

Chen, Y., Huang, W., Zhou, K., Bian, Y., Han, B., and Cheng, J. (2024). Understanding and improving feature learning for out-of-distribution generalization. Advances in Neural Information Processing Systems , 36.

Child, R., Gray, S., Radford, A., and Sutskever, I. (2019). Generating long sequences with sparse

transformers. arXiv preprint arXiv:1904.10509 .

Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. (2023). On the optimization and generalization of multi-head attention. arXiv preprint arXiv:2310.12680 .

Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 .

Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 .

Edelman, B. L., Goel, S., Kakade, S., and Zhang, C. (2022). Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR.

Frei, S., Chatterji, N. S., and Bartlett, P. (2022). Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory , pages 2668-2703. PMLR.

Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. (2022). What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598.

Guo, M.-H., Xu, T.-X., Liu, J.-J., Liu, Z.-N., Jiang, P.-T., Mu, T.-J., Zhang, S.-H., Martin, R. R., Cheng, M.-M., and Hu, S.-M. (2022). Attention mechanisms in computer vision: A survey. Computational visual media , 8(3):331-368.

Hassani, A., Walton, S., Shah, N., Abuduweili, A., Li, J., and Shi, H. (2021). Escaping the big data paradigm with compact transformers. arXiv preprint arXiv:2104.05704 .

Hastie, T., Montanari, A., Rosset, S., and Tibshirani, R. J. (2022). Surprises in high-dimensional ridgeless least squares interpolation. Annals of statistics , 50(2):949.

He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R. (2022). Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 16000-16009.

Hu, Z., Dong, Y., Wang, K., and Sun, Y. (2020). Heterogeneous graph transformer. In Proceedings of the web conference 2020 , pages 2704-2710.

Huang, W., Cao, Y., Wang, H., Cao, X., and Suzuki, T. (2023a). Graph neural networks provably benefit from structural information: A feature learning perspective. arXiv preprint arXiv:2306.13926 .

Huang, W., Shi, Y., Cai, Z., and Suzuki, T. (2023b). Understanding convergence and generalization in federated learning through feature learning theory. In The Twelfth International Conference on Learning Representations .

Huang, Y., Cheng, Y., and Liang, Y. (2023c). In-context convergence of transformers. arXiv preprint arXiv:2310.05249 .

Janner, M., Li, Q., and Levine, S. (2021). Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:1273-1286.

Jelassi, S., Sander, M., and Li, Y. (2022). Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:37822-37836.

Khan, S., Naseer, M., Hayat, M., Zamir, S. W., Khan, F. S., and Shah, M. (2022). Transformers in vision: A survey. ACM computing surveys (CSUR) , 54(10s):1-41.

Kornowski, G., Yehudai, G., and Shamir, O. (2024). From tempered to benign overfitting in relu neural networks. Advances in Neural Information Processing Systems , 36.

Kou, Y., Chen, Z., Chen, Y., and Gu, Q. (2023). Benign overfitting in two-layer relu convolutional neural networks. In International Conference on Machine Learning , pages 17615-17659. PMLR.

Li, H., Wang, M., Liu, S., and Chen, P.-Y. (2023a). A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015 .

Li, H., Wang, M., Lu, S., Cui, X., and Chen, P.-Y. (2024). Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607 .

Li, Y., Li, Y., and Risteski, A. (2023b). How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning , pages 1968919729. PMLR.

Li, Z., Zhou, Z.-H., and Gretton, A. (2021). Towards an understanding of benign overfitting in neural networks. arXiv preprint arXiv:2106.03212 .

Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L. (2014). Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13 , pages 740-755. Springer.

Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. (2021). Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022.

Meng, X., Zou, D., and Cao, Y. (2023). Benign overfitting in two-layer relu convolutional neural networks for xor data. arXiv preprint arXiv:2310.01975 .

Naseer, M. M., Ranasinghe, K., Khan, S. H., Hayat, M., Shahbaz Khan, F., and Yang, M.-H. (2021). Intriguing properties of vision transformers. Advances in Neural Information Processing Systems , 34:23296-23308.

OpenAI (2023). Gpt-4 technical report.

Oquab, M., Darcet, T., Moutakanni, T., Vo, H., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F., El-Nouby, A., et al. (2023). Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193 .

Oymak, S., Rawat, A. S., Soltanolkotabi, M., and Thrampoulidis, C. (2023). On the role of attention in prompt-tuning. In International Conference on Machine Learning , pages 26724-26768. PMLR.

Paul, S. and Chen, P.-Y. (2022). Vision transformers are robust learners. In Proceedings of the AAAI conference on Artificial Intelligence , volume 36, pages 2071-2081.

Robinson, B. S., Drenkow, N., Conwell, C., and Bonner, M. (2023). A sparse null code emerges in deep neural networks. In UniReps: the First Workshop on Unifying Representations in Neural Models .

Tian, Y., Wang, Y., Chen, B., and Du, S. S. (2023a). Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in Neural Information Processing Systems , 36:71911-71947.

Tian, Y., Wang, Y., Zhang, Z., Chen, B., and Du, S. (2023b). Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535 .

Timor, N., Vardi, G., and Shamir, O. (2023). Implicit regularization towards rank minimization in relu networks. In International Conference on Algorithmic Learning Theory , pages 1429-1459. PMLR.

Tolstikhin, I. O., Houlsby, N., Kolesnikov, A., Beyer, L., Zhai, X., Unterthiner, T., Yung, J., Steiner, A., Keysers, D., Uszkoreit, J., et al. (2021). Mlp-mixer: An all-mlp architecture for vision. Advances in neural information processing systems , 34:24261-24272.

Touvron, H., Bojanowski, P., Caron, M., Cord, M., El-Nouby, A., Grave, E., Izacard, G., Joulin, A., Synnaeve, G., Verbeek, J., et al. (2022). Resmlp: Feedforward networks for image classification with data-efficient training. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45(4):5314-5321.

Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. (2021). Training data-efficient image transformers & distillation through attention. In International conference on machine learning , pages 10347-10357. PMLR.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems , 30.

Vuckovic, J., Baratin, A., and Combes, R. T. d. (2020). A mathematical theory of attention. arXiv preprint arXiv:2007.02876 .

Wang, K., Muthukumar, V., and Thrampoulidis, C. (2021). Benign overfitting in multiclass classification: All roads lead to interpolation. Advances in Neural Information Processing Systems , 34:24164-24179.

Wang, W., Ma, S., Xu, H., Usuyama, N., Ding, J., Poon, H., and Wei, F. (2023). When an image is worth 1,024 x 1,024 words: A case study in computational pathology. arXiv preprint arXiv:2312.03558 .

Wu, D. and Xu, J. (2020). On the optimal weighted ℓ 2 regularization in overparameterized linear regression. Advances in Neural Information Processing Systems , 33:10112-10123.

Xu, Z., Wang, Y., Frei, S., Vardi, G., and Hu, W. (2023). Benign overfitting and grokking in relu networks for xor cluster data. arXiv preprint arXiv:2310.02541 .

Yun, S., Jeong, M., Kim, R., Kang, J., and Kim, H. J. (2019). Graph transformer networks. Advances in neural information processing systems , 32.

Zhang, J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra, S. (2020). Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems , 33:15383-15393.

Zhang, R., Frei, S., and Bartlett, P. L. (2024). Trained transformers learn linear models in-context. Journal of Machine Learning Research , 25(49):1-55.

Zhou, M. and Ge, R. (2023). Implicit regularization leads to benign overfitting for sparse linear regression. In International Conference on Machine Learning , pages 42543-42573. PMLR.

Zhu, L., Liao, B., Zhang, Q., Wang, X., Liu, W., and Wang, X. (2024). Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.09417 .

Zou, D., Wu, J., Braverman, V., Gu, Q., and Kakade, S. (2021). Benign overfitting of constant-stepsize sgd for linear regression. In Conference on Learning Theory , pages 4633-4635. PMLR.

Contents

Appendix

A More Experimental Results on Training Dynamics

In this section, we follow a similar setting as the experiments in Section 6, and add more experiments. All the experiments in this paper can be performed within hours.

Experiments setting: We focus on tracking the dynamics of QKV under benign and harmful overfitting regime, and the key parameters are as follows:

M = 4

d = 2048

d h = 1024

d v = 256

σ p = 0 . 05

˜ σ p = 1

Specifically, we set n = 16 , ∥ µ ∥ 2 = 25 as benign overfitting setting, and set n = 10 , ∥ µ ∥ 2 = 15 as harmful overfitting setting.

0

100

200

300

400

500

Training Iterations t

0

2

4

6

8

µ

/latticetop

+

W

(

t

)

Q

W

(

t

)

/latticetop

K

µ

+

µ

/latticetop

+

W

(

t

)

Q

W

(

t

)

/latticetop

K

ξ

2

0

100

200

300

400

500

Training Iterations t

-

0

.

02

0

.

00

0

.

02

0

.

04

0

.

06

0

.

08

0

.

10

0

.

12

ξ

/latticetop

i

W

(

t

)

Q

W

(

t

)

/latticetop

K

µ

+

ξ

/latticetop

i

W

(

t

)

Q

W

(

t

)

/latticetop

K

ξ

2

0

100

200

300

400

500

Training Iterations t

0

1

2

3

4

5

6

7

8

µ

/latticetop

+

W

(

t

)

V

w

O

ξ

/latticetop

2

W

(

t

)

V

w

O

0

100

200

300

400

500

Training Iterations t

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

0

.

6

attention on signal

µ

+

attention on noise

ξ

2

0

100

200

300

400

500

Training Iterations t

0

1

2

3

4

5

6

7

µ

/latticetop

+

W

(

t

)

Q

W

(

t

)

/latticetop

K

µ

+

µ

/latticetop

+

W

(

t

)

Q

W

(

t

)

/latticetop

K

ξ

2

0

100

200

300

400

500

Training Iterations t

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

ξ

/latticetop

i

W

(

t

)

Q

W

(

t

)

/latticetop

K

µ

+

ξ

/latticetop

i

W

(

t

)

Q

W

(

t

)

/latticetop

K

ξ

2

0

100

200

300

400

500

Training Iterations t

0

2

4

6

8

µ

/latticetop

+

W

(

t

)

V

w

O

ξ

/latticetop

2

W

(

t

)

V

w

O

0

100

200

300

400

500

Training Iterations t

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

0

.

6

attention on signal

µ

+

attention on noise

ξ

2

Experiments results: Figure 3 shows the training dynamics under benign overfitting setting, and Figure 4 shows the training dynamics under harmful overfitting setting. In Figure 3, x ⊤ l W Q W ⊤ K µ + grow faster than x ⊤ l W Q W ⊤ K ξ 2 , and more and more attention are paid on signal µ + . Meanwhile, µ + W V w O grows faster than ξ 2 W V w O . In Figure 4, x ⊤ l W Q W ⊤ K ξ 2 grow faster than x ⊤ l W Q W ⊤ K µ + , and more and more attention are paid on noise ξ 2 . Meanwhile, ξ 2 W V w O grows faster than µ + W V w O . In other words, the ViT prefers to learn the signals rather than memorize the noises under benign overfitting setting, while it prefers to memorize the noises rather than learn the signals under harmful overfitting setting.

B Basic Calculation

In this section, we introduce some notations for this paper, and give expressions for the exact gradient of loss function L S ( θ ) with respect to W Q , W K and W V .

B.1 Notations

B.2 Gradient Calculation

Gradient of Softmax . Before computing the gradient of QKV, we solve for the gradient of Softmax function. Suppose there is an input vector a = ( a 1 , a 2 , . . . , a m ) , then the softmax function can be defined as follows:

Then we calculate the gradient of softmax ( a i ) with respect to a i and a j for i, j ∈ [ m ] , i = j .

̸

(3) and (4) show that the gradient of a Softmax function can be represented as a Jacobi matrix where the elements on the diagonal are softmax ( a i ) · ( 1 -softmax ( a i ) ) and the elements on the off-diagonal are softmax ( a i ) · softmax ( a j ) . With these properties, the gradient of our attention vector φ n,i ( i ∈ [ M ] , n ∈ [ N ] ) can be expressed as follows

Lemma B.1. the gradients of loss function L S ( θ ) with respect to W Q , W K and W V are given by

Proof of Lemma B.1. Considering that in the vector x ⊤ n,i W Q W ⊤ K X ⊤ n , the j -th element of x ⊤ n,i , i.e. ( x n,i ) [ j ] , is only used to multiply W Q [ j, :] , then we have

Considering all j ∈ [ d ] , we get

And then we get

where the third equality is by (6), the fourth equality is by chain rule, the fifth equality is by (5), For the last equality, we expand the equality by materializing all the x n,i (e.g., x n, 1 = µ + for n ∈ S + ). Using the similar method, we obtain the gradient of W K .

The gradient of W V can be obtained using the chain rule as follows

B.3 Update Rules

Definition B.1 (Scalarized V) . Let W ( t ) V be the V matrix of the ViT at the t-th iteration of gradient descent. Then there exist coefficients γ ( t ) V, + , γ ( t ) V, -, ρ ( t ) V,n,i such that

for i ∈ [ M ] \{ 1 } , n ∈ [ N ] . We further denote V ( t ) + := µ ⊤ + W ( t ) V w O , V ( t ) -:= µ ⊤ -W ( t ) V w O and V ( t ) n,i := ξ ⊤ n,i W ( t ) V w O . We refer to it as scalarized V .

Definition B.2 (Vectorized Q & K) . Let W ( t ) Q and W ( t ) K be the QK matrices of the ViT at the t-th iteration of gradient descent. Then we define the vectorized Q and vectorized K as follows

for i ∈ [ M ] \{ 1 } , n ∈ [ N ] .

Definition B.3 (Gradient Decomposition) . There exist coefficients α ( t ) + , + , α ( t ) n, + ,i , α ( t ) -, -, α ( t ) n, -,i , α ( t ) n,i, + , α ( t ) n,i, -, α ( t ) n,i,n ′ ,i ′ , β ( t ) + , + , β ( t ) n, + ,i , β ( t ) -, -, β ( t ) n, -,i , β ( t ) n,i, + , β ( t ) n,i, -, β ( t ) n,i,n ′ ,i ′ such that

Remark. With the scalarized V in B.1, we can portray the learning process of W V on signals and noises by analyzing the dynamics of γ ( t ) V, + , γ ( t ) V, -and ρ ( t ) V,n,i (or V ( t ) + , V ( t ) -and V ( t ) n,i ) . Also, as scalars, the effect of W V on the learning process of W Q and W K is more conveniently to describe. With the Vectorized Q & K in B.2, we can decompose x ⊤ W Q W K X ⊤ into the inner product of vectors q (a generalized reference to q + , q -and q n,i ) and k (a generalized reference to k + , k -and k n,i ), e.g., ⟨ q + , k + ⟩ = µ ⊤ + W Q W K µ + . Then the dynamics of x ⊤ W Q W K X ⊤ can be characterized by analyzing the dynamics of the inner products of q and k , and then further get the dynamics of attentions. The gradient decomposition in B.3 comes from (7) and (8). Note that (7) ends in X n W K and (8) ends in x ⊤ n,i W Q . This suggests that the dynamics of q ( t ) ± and q ( t ) n,i can be decomposed into linear combinations of k ( t ) ± and k ( t ) n,i , and similarly, the gradient of k ( t ) ± and k ( t ) n,i can be decomposed into linear combinations of q ( t ) ± and q ( t ) n,i . We calculate α (a generalized reference to α + , + , α n,i + , α -, -, etc.) and β (a generalized reference to β + , + , β n, + ,i , β -, -, etc.) in detail in Appendix F.

Lemma B.2 (Update Rule for V) . The coefficients γ ( t ) V, + , γ ( t ) V, -, ρ ( t ) V,n,i defined in Definition B.1 satisfy the following iterative equations:

∈

for

i

[

M

]

\{

1

}

, n

∈

[

N

]

.

Proof of Lemma B.2. Base on (9), we have

where the second equality we expand X n into vectors and make inner products with x , the third equality we materializing all the x n,l (e.g., x n, 1 = µ + for n ∈ S + ). Note the orthogonality between µ and ξ n,i , we can remove many of the terms in this equation. Take x = µ + as an example, we have

Then we have

Dividing by ∥ w O ∥ 2 2 we get

This proves the update rule for γ ( t ) V, + , The proof for γ ( t ) V, + and ρ ( t ) V,n,i is similar to it.

Lemma B.3 (Update Rule for QK) . The dynamics of x ⊤ W Q W K x can be characterized as follows

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] .

Proof of Lemma B.3. Based on Definition B.3, we have

Then we get

where ∆ W ( t ) Q := W ( t +1) Q -W ( t ) Q , ∆ W ( t ) K := W ( t +1) K -W ( t ) K , the second equality is by ∆ q ( t ) + = µ ⊤ + ∆ W ( t ) Q and ∆ k ( t ) + = µ ⊤ + ∆ W ( t ) K , the third equality is by plugging (14) and (15). This proves the update rule for ⟨ q ( t ) + , k ( t ) + ⟩ , the proof for ⟨ q ( t ) n,i , k ( t ) + ⟩ , ⟨ q ( t ) + , k ( t ) n,j ⟩ , ⟨ q ( t ) n,i , k ( t ) n,j ⟩ are exactly the same.

Remark. In (16), note that the magnitude of ⟨ k ( t ) + , k ( t ) n,i ⟩ , ⟨ q ( t ) + , q ( t ) n,i ⟩ is much smaller than that of ∥ k ( t ) + ∥ 2 2 , ∥ q ( t ) + ∥ 2 2 , and that the magnitude of α ( t ) n, + ,i , β ( t ) n, + ,i will not be greater than that of α ( t ) + , + , β ( t ) + , + (we will prove it later). Also, since the learning rate η is sufficiently small, the last term of (16) will be very small. Then (16) can be expressed in the following form

The dynamics of ⟨ q ( t ) n,i , k ( t ) + ⟩ , ⟨ q ( t ) + , k ( t ) n,j ⟩ , ⟨ q ( t ) n,i , k ( t ) n,j ⟩ can be expressed in a similar way. Later we will rigorously explain the so-call {lower order term}.

C Concentration Inequalities

In this section, we will give some concentration inequalities that show some important properties of the data and the ViT parameters at random initialization.

Lemma C.1 (Lemma B.1 in Cao et al. (2022)) . Suppose that δ > 0 and n ≥ 8 log(4 /δ ) . Then with probability at least 1 -δ ,

Lemma C.2 (Initialization of V) . Suppose that δ > 0 . Then with probability at least 1 -δ , for i ∈ [ M ] \{ 1 } , n ∈ [ N ] .

Proof of Lemma C.2. It is clear that µ ⊤ ± W (0) V w O is a random variable with mean zero and variance σ 2 V ∥ w O ∥ 2 2 ∥ µ ∥ 2 2 . Therefore, by Gaussian tail bound and the condition that σ V ≤ ˜ O ( ∥ w O ∥ -1 2 · min {∥ µ ∥ -1 2 , ( σ p √ d ) -1 } · d -1 4 h ) , with probability at least 1 -δ/NM ,

Moreover, each element of vector w ⊤ O W (0) V is a random variable with mean zero and variance σ 2 V ∥ w O ∥ 2 2 . Therefore, by Bernstein's inequality and the condition of σ V , with probability at least 1 -δ/NM

Applying a union bound completes the proof.

Lemma C.3 (Initialization of QK) . Suppose that δ > 0 . Then with probability at least 1 -δ , for

i, j

∈

[

M

]

\{

1

}

,

i

′

Proof of Lemma C.3. It is clear that each element of vector q (0) + is a random variable with mean zero and variance ∥ µ ∥ 2 2 σ 2 h . By Bernstein's inequality, with probability at least 1 -δ/ 3 N 2 M 2 , we have

Therefore, as long as d h = Ω(log(6 N 2 M 2 /δ )) , we have

Similarly, we have

Moreover, ⟨ q (0) + , q (0) -⟩ has mean zero. By Bernstein's inequality, with probability at least 1 -δ/ 3 N 2 M 2 , we have

We can prove the rest of this lemma in a similar way.

Lemma C.4 (Lemma B.2 in Cao et al. (2022)) . Suppose that δ > 0 and d = Ω(log(4 NM/δ )) . Then with probability at least 1 -δ

for i, i ′ ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , i = i ′ or n = n ′ .

D Benign Overfitting

In this section, we consider the benign overfitting regime under the condition that N · SNR 2 = Ω(1) . We analyze the dynamics of V ± , V n,i , the inner product of q ± , q n,i and k ± , k n,i during gradient descent training, and further give the upper bound for population loss. The proofs in this section are based on the results in Section C, which hold with high probability.

D.1 Stage I

In Stage I, V ( t ) ± , V ( t ) n,i begin to pull apart until | V ( t ) ± | is sufficiently larger than | V ( t ) n,i | . At the same time, the inner products of q and k maintain their magnitude.

Lemma D.1 (Gradient of Loss) . As long as max {| V ( t ) + | , | V ( t ) -| , | V ( t ) n,i |} = o (1) , we have -ℓ ′ ( y n f ( X n , θ ( t ))) remains 1 / 2 ± o (1) .

Proof of Lemma D.1. Note that ℓ ( z ) = log(1 + exp( -z )) and -ℓ ′ = exp( -z ) / (1 + exp( -z )) , without loss of generality, we assume y n = 1 , we have

Note that

Then we have

̸

̸

LemmaD.2 (Bound of Attention) . As long as |⟨ q ( t ) ± , k ( t ) ± ⟩| , |⟨ q ( t ) n,i , k ( t ) ± ⟩| , |⟨ q ( t ) ± , k ( t ) n,j ⟩| , |⟨ q ( t ) n,i , k ( t ) n,j ⟩| = o (1) , we have

Proof of Lemma D.2. It is clear that exp( o (1)) = 1+ o (1) . Therefore, as long as |⟨ q ( t ) ± , k ( t ) ± ⟩| = o (1) , we have

Similarly, we have

Theorem 4.1, we have for t ∈ [0 , T 0 ] .

Proof of Lemma D.3. By Lemma B.2, we have

where the second inequality is by Lemma C.1 and -ℓ ′ ( t ) n ≤ 1 . Similarly, we have

Lemma D.3 (Upper bound of V) . Let T 0 = O ( 1 ηd 1 4 h ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ) . Then under the same conditions as

By Definition B.1, we have

where the first inequality is by triangle inequality, the second inequality is by Lemma C.2. Similarly, we have | V ( t ) -| = O ( d -1 4 h ) .

where the second inequality is by Lemma C.4 and -ℓ ′ ( t ) n ≤ 1 , the third inequality is by d = ˜ Ω ( ϵ -2 N 2 d h ) ≥ 4 N √ log(4 N 2 M 2 /δ ) , the last inequality is by N · SNR 2 = Ω(1) . Then by Definition B.1, we have

where the first inequality is by triangle inequality, the second inequality is by Lemma C.2, which completes the proof.

Lemma D.4 (Inner Products Hold Magnitude) . Let T 0 = O ( 1 ηd 1 4 h ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ) . Then under the same conditions as Theorem 4.1, we have

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] and t ∈ [0 , T 0 ]

The proof for Lemma D.4 is in Section F.3. Note that σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -3 2 , thus O ( max {∥ µ ∥ 2 2 , σ 2 p d } · σ 2 h · √ d h log(6 N 2 M 2 /δ ) ) = o (1) .

Lemma D.5 (V's Beginning of Learning Signals) . Under the same conditions as Theorem 4.1, there exist T 1 = 10 M (3 M +1) N ηd 1 4 h ( N ∥ µ ∥ 2 2 -60 M 2 C 2 p σ 2 p d ) ∥ w O ∥ 2 2 such that the first element of vector X n W ( t ) V w O dominate its other elements, that is, V ( t ) + ≥ 3 M · | V ( t ) n,i | for all n ∈ S + , i ∈ [ M ] \{ 1 } and V ( t ) -≤ -3 M · | V ( t ) n,i | for all n ∈ S -, i ∈ [ M ] \{ 1 } .

Proof of Lemma D.5. Let C be a constant larger than 10 M (3 M +1) , then as long as N · SNR 2 = N ∥ µ ∥ 2 2 σ 2 p d ≥ 60 CM 2 C 2 p C -10 M (3 M +1) , we have N ∥ µ ∥ 2 2 -60 M 2 C 2 p σ 2 p d ≥ 10 M (3 M +1) N ∥ µ ∥ 2 2 C , and further get

which satisfies the time condition in Lemma D.3 and Lemma D.4. Then by Lemma D.1 and Lemma D.2 we have

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] and t ∈ [0 , T 1 ] . Plugging them in the update rule for γ ( t ) V, + showed in Lemma B.2 and we have

Then by Definition B.1 and taking a summation, we have

Similarly, we have

Similarly, by | ρ ( t +1) V,n,i -ρ ( t ) V,n,i | ≤ 2 η ˜ σ 2 p d N in (17), we have

According to (19), (20) and (21), it is easy to verify that V ( T 1 ) + -3 M · | V ( T 1 ) n,i | ≥ 0 and V ( T 1 ) -+3 M · | V ( T 1 ) n,i | ≤ 0 , which completes the proof.

D.2 Stage II

In stage II, ⟨ q + , k + ⟩ , ⟨ q n,i , k + ⟩ grows while ⟨ q + , k n,j ⟩ , ⟨ q n,i , k n,j ⟩ decreases, resulting in attention focusing more and more on the signals and less on the noises. By the results of stage I, we have the following conditions at the beginning of stage II

Some of the proofs at this stage are based on the above conditions.

Notations. To better characterize the gap between different inner products, we define the following notations:

Lemma D.6 (Upper bound of V) . Let T 0 = O ( 1 η ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 log(6 N 2 M 2 /δ ) ) . Then under the same conditions as Theorem 4.1, we have

for t ∈ [0 , T 0 ] .

The proof of Lemma D.6 is similar to that of Lemma D.3, except that the time T 0 is changed.

Let T 2 = Θ ( 1 η ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 log(6 N 2 M 2 /δ ) ) , then by Lemma D.6 and Lemma D.1 we have 1 2 -o (1) ≤ -ℓ ′ ( t ) n ≤ 1 2 + o (1) for n ∈ [ N ] , t ∈ [ T 1 , T 2 ] , which can simplify the calculations of α and β defined in Definition B.3 by replacing -ℓ ′ ( t ) n by their bounds. Next we prove the following four propositions B ( t ) , C ( t ) , D ( t ) , E ( t ) by induction on t for t ∈ [ T 1 , T 2 ] :

C ( t ) :

̸

By the results of Stage I, we know that B ( T 1 ) , C ( T 1 ) , E ( T 1 ) are true. To prove that B ( t ) , C ( t ) , D ( t ) and E ( t ) are true in stage 2, we will prove the following claims holds for t ∈ [ T 1 , T 2 ] :

D.2.1 Proof of Claim 1

By the results of Stage I, we have

Assume that D ( T 1 ) , . . . , D ( t -1) ( t ∈ [ T 1 , T 2 ] ) are true, then ⟨ q ( s ) ± , k ( s ) ± ⟩ , ⟨ q ( s ) n,i , k ( s ) ± ⟩ are monotonically non-decreasing and ⟨ q ( s ) ± , k ( s ) n,j ⟩ , ⟨ q ( s ) n,i , k ( s ) n,j ⟩ are monotonically non-increasing for s ∈ [ T 1 , t -1] , so we have

̸

̸

for s ∈ [ T 1 , t ] . Further we have the lower bounds for the attention on signal µ ± as follows for s ∈ [ T 1 , t ] :

where the second equality is by exp( o (1)) = 1 + o (1) . Similarly, we have softmax ( ⟨ q ( s ) n,i , k ( s ) ± ⟩ ) ≥ 1 M -o (1) . Plugging them in the update rule for γ ( s ) V, + showed in Lemma B.2 and we have

for s ∈ [ T 1 , t ] . Then by Definition B.1 and taking a summation, we have

where the last inequality is by V ( T 1 ) + ≥ 0 and M = Θ(1) . Similarly, we have

By | ρ ( t +1) V,n,i -ρ ( t ) V,n,i | ≤ 2 η ˜ σ 2 p d N in (17) and taking a summation, we have

Combining (24) and (25) we have

where the last inequality is by V ( T 1 ) + ≥ 3 M · | V ( T 1 ) n,i | and requires N · SNR 2 ≥ 60 M 2 C 2 p . The proof for V ( t +1) -≤ -3 M · | V ( t +1) n,i | is the same.

Next, we prove the upper bound for V ± and V n,i . Based on the upper bound of attention(<1) and -ℓ ′ n ≤ 1 , we have

Then we can get that

where the first inequality is by the monotonicity of γ V, + and the definition of V + , the last inequality is by the result of stage 1 where V ( T 1 ) + = O ( d -1 ) . Similarly, we have

which completes the proof for the upper bound of V ± . Expanding (25) yields

where the last inequality is by the result of phase 1 where | V ( T 1 ) n,i | = O ( d -1 4 h ) and the condition that N · SNR 2 > Ω(1) .

D.2.2 Proof of Claim 2

By the results of F.5, we have the dynamic of ⟨ q , k ⟩ as follows

for s ∈ [ T 1 , t ] . The seven equations above show that ⟨ q ( s ) ± , k ( s ) ± ⟩ , ⟨ q ( s ) n,i , k ( s ) ± ⟩ are monotonically increasing and ⟨ q ( s ) ± , k ( s ) n,j ⟩ , ⟨ q ( s ) n,i , k ( s ) n,j ⟩ are monotonically decreasing. Next, we provide the logarithmic increasing lower bounds of Λ ( s +1) n, ± ,j and Λ ( s +1) n,i, ± ,j . By (31) and (33), we have

where the last inequality is by σ 2 h ≥ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } d -1 2 h (log(6 N 2 M 2 /δ )) -2 . Multiply both sides simultaneously by exp(Λ ( s ) n, + ,j ) and get

Taking a summation from T 1 to t and get

By the property that Λ ( t ) n, + ,j is monotonically increasing, we have

Similarly, we have

By (35) and (37), we have

Then using the similar method as for Λ ( t ) n, ± ,j , we get

which complete the proof. The proof for Claim 3 is in Section F.8

D.2.3 Proof of Claim 4

By the results of F.6, we have

for s ∈ [ T 1 , t ] . Further we have

For the last inequality, by η ≤ ˜ O (min { µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h ) , σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } d -1 2 h (log(6 N 2 M 2 /δ )) -3 2 , ⟨ q ( T 1 ) + , k ( T 1 ) + ⟩ = o (1) and the monotonicity of ⟨ q ( s ) + , k ( s ) + ⟩ for s ∈ [ T 1 , t ] , we have exp ( ηC 10 ∥ µ ∥ 4 2 σ 2 h d h exp( ⟨ q ( t ) + , k ( t ) + ⟩ ) ) ≤ exp( o (1)) ≤ C 11 . Multiplying both

where the last inequality is by plugging (46). Taking a summation we have

where the first inequality is due to ⟨ q ( s ) + , k ( s ) + ⟩ is monotone increasing, the last inequality is by T 2 = Θ( η -1 ∥ µ ∥ -2 2 ∥ w O ∥ -2 2 log(6 N 2 M 2 /δ ) -1 ) and σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } d -1 2 h (log(6 N 2 M 2 /δ )) -3 2 . By ∫ ⟨ q ( t +1) + , k ( t +1) + ⟩ ⟨ q ( T 1 ) + , k ( T 1 ) + ⟩ exp( x ) dx = exp( ⟨ q ( t +1) + , k ( t +1) + ⟩ ) -exp( ⟨ q ( T 1 ) + , k ( T 1 ) + ⟩ ) , we have

By the results of F.6, we also have

Then using the similar method as for ⟨ q ( t +1) + , k ( t +1) + ⟩ , we get

Next we provide the upper bound for |⟨ q ( t +1) ± , k ( t +1) ∓ ⟩| , |⟨ q ( t +1) n,i , k ( t +1) n ′ ,j ⟩| . By the results of F.7, we have

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i, j ∈ [ M ] \{ 1 } , n ∈ S ± .

̸

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , n = n ′ . Plugging these into the update rule of ⟨ q ( t ) ± , k ( t ) ∓ ⟩ , ⟨ q ( t ) n,i , k ( t ) n,j ⟩ and assume that propositions C ( T 1 ) , . . . , C ( t ) hold, we have

where the first inequality is by triangle inequality, the second inequality is by (213), the last equality is by |⟨ q ( T 1 ) + , k ( T 1 ) -⟩| = o (1) and d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . Similarly we have

̸

̸

̸

̸

where the first inequality is by triangle inequality, the second inequality is by (196), the second equality is by σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -3 2 , the last equality is by |⟨ q ( T 1 ) n,i , k ( T 1 ) n,j ⟩| = o (1) , d = ˜ Ω ( ϵ -2 N 2 d h ) and d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) .

̸

̸

D.3 Stage III

In Stage III, the outputs of ViT grow up and the loss derivatives are no longer at o (1) . We will carefully compute the growth rate of V ± and V n,i while keeping monitoring the monotonicity of ⟨ q , k ⟩ . By substituting t = T 2 = Θ ( 1 η ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ) into propositions B ( t ) , C ( t ) , D ( t ) , E ( t ) in Stage II, we have the following conditions at the beginning of stage III

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , i = j or n = n ′ .

̸

̸

Let T 3 = Θ ( 1 ηϵ ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ) , Next we prove the following four propositions F ( t ) , G ( t ) , H ( t ) , I ( t ) by induction on t for t ∈ [ T 2 , T 3 ] :

for i, j ∈ [ M ] \{ 1 } , n, n ∈ [ N ] , n = n .

· F ( t ) :

̸

∈ \{ } ∈

̸

̸

· H ( t ) :

̸

By the results of Stage II, we know that F ( T 1 ) , G ( T 2 ) , I ( T 2 ) are true. To prove that F ( t ) , G ( t ) , H ( t ) and I ( t ) are true in stage 3, we will prove the following claims holds for t ∈ [ T 2 , T 3 ] :

D.3.1 Proof of Claim 5

The proofs for V ( t ) + ≥ 3 M · | V ( t ) n,i | and V ( t ) -≤ -3 M · | V ( t ) n,i | are the same as for D.2.1. Based on H ( T 2 ) , . . . , H ( t ) where ⟨ q ( s ) ± , k ( s ) ± ⟩ and ⟨ q ( s ) n,i , k ( s ) ± ⟩ are monotonically non-decreasing and max j ⟨ q ( s ) ± , k ( s ) n,j ⟩ , max j ⟨ q ( s ) n,i , k ( s ) n,j ⟩ are monotonically non-increasing for s ∈ [ T 2 , t -1] , we have

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] , s ∈ [ T 2 , t ] . We further get

For the first inequality, by (62) and the monotonicity of ⟨ q ( s ) , k ( s ) ⟩ ( ⟨ q ( s ) ± , k ( s ) ± ⟩ is increasing and ⟨ q ( s ) ± , k ( s ) n,j ⟩ is decreasing), there exist a constant C such that C exp( ⟨ q ( s ) ± , k ( s ) ± ⟩ ) ≥ exp( ⟨ q ( s ) ± , k ( s ) ± ⟩ ) + M ∑ j ′ =2 exp( ⟨ q ( s ) ± , k ( s ) n,j ′ ⟩ ) . The second inequality is by plugging (62). Similarly, we have

Plugging (64) and (65) into the update rule of ρ V,n,i in Lemma B.2 and get

̸

̸

where the second inequality is by Lemma C.4 and | ℓ ′ ( t ) n | ≤ 1 . For the last inequality, since d = ˜ Ω ( ϵ -2 N 2 d h ) , we have N · M · 2˜ σ 2 p √ d log(4 N 2 M 2 /δ ) ≤ 1 2 ˜ σ 2 p d . By Definition B.1 and taking a summation we have

(67)

where the first equality is by N · SNR 2 ≥ Ω(1) , the second equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . Then we have a constant upper bound for the sum of V n,i as follows:

for n ∈ [ N ] , s ∈ [ T 2 , t ] .

Expanding (64) and (65) we have

where the equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) .

where the equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . Then we have

for i ∈ [ M ] \{ 1 } , n ∈ [ N ] , s ∈ [ T 2 , t ] . Next we provide the bounds for -ℓ ′ ( s ) n . Note that ℓ ( z ) = log(1+exp( -z )) and -ℓ ′ = exp( -z ) / (1+exp( -z )) , without loss of generality, we assume y n = 1 , we have

for s ∈ [ T 2 , t ] . The second equality is by plugging (68), (69), (70) and (71), the inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | for s ∈ [ T 2 , t ] . Similarly, we have

Plugging (73) and (74) into (72) we have

where the first inequality is by plugging (74). For the last inequality, note that V ( T 2 ) + ≥ 0 and V ( s ) + is monotonically increasing, so there exist a constant C 16 such that 1 1+exp( V ( s ) + ) ≥ C 16 exp( V ( s ) + ) . We also have the upper bound

Plugging (70), (71) and (75) into the update rule of γ ( t ) V, + and in Lemma B.2 and get

where the second inequality is by (75). Then by definition B.1, we get

Multiply both sides simultaneously by exp( V ( s ) + ) and get

Taking a summation from T 2 to t and get

By the property that V ( s ) + is monotonically increasing, we have

Similarly, we have

Next we provide upper bounds for V ( t +1) + and V ( t +1) -. By the update rule of γ ( t ) V, + and in Lemma B.2 we have

where the second inequality is by (76). Then by definition B.1, we get

Further we have

For the last inequality, by η ≤ ˜ O (min { µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h ) , V ( T 2 ) + = Θ(1) and the monotonicity of V ( s ) + , we have exp( 3 η ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 8 exp( V ( s ) + / 2) ) ≤ exp( o (1)) ≤ C 18 . Multiplying both sides by ( V ( s +1) + / 2 -V ( s ) + / 2 ) simultaneously gives

where the last inequality is by plugging (82). Taking a summation we have

Similarly, we have

D.3.2 Proof of Claim 6

By H ( T 2 ) , . . . , H ( t -1) , we have softmax ( ⟨ q ( t ) ± , k ( t ) ± ⟩ ) , softmax ( ⟨ q ( t ) n,i , k ( t ) ± ⟩ ) = 1 -o (1) and softmax ( ⟨ q ( t ) ± , k ( t ) n,j ⟩ ) , softmax ( ⟨ q ( t ) n,i , k ( t ) n,j ⟩ ) = o (1) , which have been proved in D.3.1. By the results of F.4, we have the signs of α and β as follows:

Then combined with G ( T ) and we have the dynamics of ⟨ q , k ⟩ as follows:

Similarly, we have

Similarly, we have

Similarly, we have

which completes the proof. The proof for Claim 7 is in Section F.11

D.3.3 Proof of Claim 8

By the results of F.9, we have

Further we have

For the last inequality, by η ≤ ˜ O (min { µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h ) , σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -3 2 , ⟨ q ( T 1 ) + , k ( T 1 ) + ⟩ = o (1) and the monotonicity of ⟨ q ( s ) + , k ( s ) + ⟩ for s ∈ [ T 1 , t ] , we have exp ( ηC 10 ∥ µ ∥ 4 2 σ 2 h d h log ( O ( 1 ϵ ) ) exp( ⟨ q ( t ) + , k ( t ) + ⟩ ) ) ≤ exp( o (1)) ≤ C 11 . Multiplying both sides by ( ⟨ q ( t +1) + , k ( t +1) + ⟩ - ⟨ q ( t ) + , k ( t ) + ⟩ ) simultaneously gives

where the last inequality is by plugging (93). Taking a summation we have

where the first inequality is due to ⟨ q ( s ) + , k ( s ) + ⟩ is monotone increasing, the last equality is by T 3 = Θ( η -1 ϵ -1 ∥ µ ∥ -2 2 ∥ w O ∥ -2 2 ) , ∥ w O ∥ 2 2 = Θ(1) and σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -3 2 . By ∫ ⟨ q ( t +1) + , k ( t +1) + ⟩ ⟨ q ( T 2 ) + , k ( T 2 ) + ⟩ exp( x ) dx = exp( ⟨ q ( t +1) + , k ( t +1) + ⟩ ) -exp( ⟨ q ( T 2 ) + , k ( T 2 ) + ⟩ ) , we have

where the last inequality is by ⟨ q ( T 2 ) + , k ( T 2 ) + ⟩ ≤ log ( d 1 2 h ) . By the results of F.9, we also have

Then using the similar method as for ⟨ q ( t +1) + , k ( t +1) + ⟩ , we get

Next we provide the upper bound for |⟨ q ( t +1) ± , k ( t +1) ∓ ⟩| , |⟨ q ( t +1) n,i , k ( t +1) n ′ ,j ⟩| . By the results of F.10, we have

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i, j ∈ [ M ] \{ 1 } , n ∈ S ± .

̸

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , n = n ′ . Plugging these and proposition G ( t ) into the update rule of |⟨ q ( t ) ± , k ( t ) ∓ ⟩| , |⟨ q ( t ) n,i , k ( t ) n,j ⟩| and get

where the first inequality is by triangle inequality, the second inequality is by (213), the last equality is by |⟨ q ( T 2 ) + , k ( T 2 ) -⟩| = o (1) and d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . Similarly we have

̸

̸

̸

̸

̸

̸

where the first inequality is by triangle inequality, the second inequality is by (196), the second equality is by σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -3 2 , the last equality is by |⟨ q ( T 2 ) n,i , k ( T 2 ) n,j ⟩| = o (1) , d = ˜ Ω ( ϵ -2 N 2 d h ) and d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) .

Lemma D.7 (Convergence of Training Loss) . There exist T = C 19 ηϵ ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 such that

Proof of Lemma D.7. Recall (68), (69), (70) and (68), we have

Substituting t = T = C 19 ηϵ ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 into propositions F ( t ) and get

For n ∈ S + , we bound f ( X n , θ ( t )) as follows

where the first inequality is by plugging (108), (109), (110) and (111). For the second inequality, by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) , we have O ( N · SNR 2 · ( log(6 N 2 M 2 /δ ) ) 3 log ( C 20 ϵ ) d 1 2 h ) = o (1) = log(1 + o (1)) ≤ log( C 20 ) as long as C 20 is sufficiently large. Then we have

Similarly, we have ℓ ( t ) n ≤ ϵ for n ∈ S -. Therefore, we have L S ( θ ( T )) = 1 N N ∑ n =1 ℓ ( t ) n ≤ ϵ .

D.4 Population Loss

Consider a new data point ( X , y ) generated from the distribution defined in Definition 3.1. Without loss of generality, we suppose that the signal token is µ + and the label is 1, i.e. X = ( µ + , ξ 2 , . . . , ξ M ) , y = 1 . We will first give the bounds of attention, then the bounds of the output of ViT, and finally the bound of the expectation of loss.

Lemma D.8 (Bound of Attention) . Under the same conditions as Theorem 4.1, with probability at least 1 -δ/ 2 N 2 M , we have that

Proof of Lemma D.8. By the result of D.2 and D.3, µ ⊤ + W ( T 1 ) Q W ( T 1 ) ⊤ K µ + = o (1) , and µ ⊤ + W ( t ) Q W ( t ) ⊤ K µ + is monotonically non-decreasing for t ∈ [ T 1 , T 3 ] , we have µ ⊤ + W ( T 3 ) Q W ( T 3 ) ⊤ K µ + ≥ -o (1) .

By proposition G ( t ) in D.3, we have ∥ µ ⊤ + W ( T 3 ) Q ∥ 2 2 = Θ ( ∥ µ ∥ 2 2 σ 2 h d h ) . In order to give bounds for µ ⊤ + W ( T 3 ) Q W ( T 3 ) ⊤ K ξ i , we provide a bound for the F-norm of W ( T 3 ) K -W (0) K . According to the properties of F-norm, we have

∥ W ( T 3 ) K -W (0) K ∥ F ≤ ∥ W ( T 1 ) K -W (0) K ∥ F + ∥ W ( T 2 ) K -W ( T 1 ) K ∥ F + ∥ W ( T 3 ) K -W ( T 2 ) K ∥ F (116) Next, we will provide the bounds of ∥ W ( T 1 ) K -W (0) K ∥ F , ∥ W ( T 2 ) K -W ( T 1 ) K ∥ F , and ∥ W ( T 3 ) K -W ( T 2 ) K ∥ F respectively.

Bounding ∥ W ( T 1 ) K -W (0) K ∥ F : Recall that so we have

We bound | ℓ ′ ( t ) n | , ∥ X n ∥ F , ∥ ( diag ( φ ( t ) n,i ) -φ ( t ) ⊤ n,i φ ( t ) n,i ) ∥ F , ∥ X n W ( t ) V w O ∥ F , ∥ µ ⊤ ± W ( t ) Q ∥ F and ∥ ξ ⊤ n,i W ( t ) Q ∥ F respectively.

| ℓ ′ ( t ) n | ≤ 1

Note that each element of vector φ n,i has an upper bound 1, and the number of tokens has a constant upper bound ( M = Θ(1) ), then we have ∥ ( diag ( φ ( t ) n,i ) -φ ( t ) ⊤ n,i φ ( t ) n,i ) ∥ F = O (1) .

By Lemma D.3, we have | V ( t ) + | , | V ( t ) -| , | V ( t ) n,i | = O ( d -1 4 h ) , so we have

By Lemma D.4, we have ∥ µ ⊤ ± W ( t ) Q ∥ F = Θ ( ∥ µ ∥ 2 σ h d 1 2 h ) , ∥ ξ ⊤ n,i W ( t ) Q ∥ F = Θ ( σ p σ h d 1 2 d 1 2 h )

Plugging them into (117) and get

Taking a summation we have

where the last inequality is by N · SNR 2 = Ω(1) .

Bounding ∥ W ( T 2 ) K -W ( T 1 ) K ∥ F : By F.7, we know that each element of matrix ( diag ( φ ( t ) n, 1 ) -φ ( t ) ⊤ n, 1 φ ( t ) n, 1 ) have upper bound O ( 1 1+ η 2 ∥ µ ∥ 4 2 ∥ w O ∥ 2 2 d 1 2 h N ( log(6 N 2 M 2 /δ ) ) 2 · ( t -T 1 )( t -T 1 -1) ) , and each element of matrix

that the number of tokens has a constant upper bound ( M = Θ(1) ), we have

By B ( t ) in D.2, we have | V ( t ) ± | , | V ( t ) n,i | ≤ O ( d -1 4 h ) + ηC 4 ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ( t -T 1 ) for t ∈ [ T 1 , T 2 ] . We further have

Then we have

The reason for the last equation is similar to (256). We also have

Plugging them into (117) and get

Taking a summation we have

By F ( t ) in D.3, we have we further have

Plugging them into (117) and get

Taking a summation we have

Combining (116), (119), (125) and (129) and get

where the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . Bounding µ ⊤ + W ( T 3 ) Q W ( T 3 ) ⊤ K ξ i : We decompose W ( T 3 ) K into W (0) K + ( W ( T 3 ) K -W (0) K ) , then we have

Now we bound µ ⊤ + W ( T 3 ) Q W (0) ⊤ K ξ i and µ ⊤ + W ( T 3 ) Q ( W ( T 3 ) ⊤ K -W (0) ⊤ K ) ξ i respectively. Bounding µ ⊤ + W ( T 3 ) Q W (0) ⊤ K ξ i : Note that ∥ µ ⊤ + W ( T 3 ) Q ∥ 2 2 = Θ ( ∥ µ ∥ 2 2 σ 2 h d h ) and each element of W (0) ⊤ K is sampled from a Gaussian distribution N (0 , σ 2 h ) , we have that each element of

µ ⊤ + W ( T 3 ) Q W (0) ⊤ K is a random variable with mean zero and standard deviation smaller than C 21 ∥ µ ∥ 2 σ 2 h d 1 2 h . Therefore, for any i ∈ [ M ] \{ 1 } by Bernstein's inequality, with probability at least 1 -δ/ 3 N 2 M 2 ,

Therefore µ ⊤ + W ( T 3 ) Q ( W ( T 3 ) ⊤ K -W (0) ⊤ K ) ξ i is a Gaussian distribution with mean zero and standard deviation smaller than C 22 ∥ µ ∥ 2 σ p σ 2 h d h . By Gaussian tail bound, for any i ∈ [ M ] \{ 1 } , with probability at least 1 -δ/ 6 N 2 M 2 ,

where the last equality is by d = ˜ Ω ( ϵ -2 N 2 d h ) and σ 2 h ≤ min {∥ µ ∥ -2 2 , ( σ 2 p d ) -1 } · d -1 2 h · ( log(6 N 2 M 2 /δ ) ) -3 2 . Applying a union bound, with probability at least 1 -δ/ 2 N 2 M ,

Therefore, we have

which complete the proof.

Lemma D.9 (Bound of V) . Under the same conditions as Theorem 4.1, with probability at least 1 -δ/ 2 N 2 M , we have that

Proof of Lemma D.9. By (9), we have

Then we have

where the second inequality is by ∣ ∣ y n ℓ ′ ( t ) n ∣ ∣ ≤ 1 , ∥ X n ∥ F = O ( max {∥ µ ∥ 2 , σ p √ d } ) and ∥ ∥ φ ( x n,l W ( t ) Q W ( t ) ⊤ K ( X n ) ⊤ ) ∥ ∥ 2 = O (1) . Taking a summation and get

Then ξ ⊤ i W ( T 3 ) V w O is a Gaussian random variable with mean zero and standard deviation smaller than O ( σ V ∥ w O ∥ 2 σ p √ d + max {∥ µ ∥ 2 ,σ p √ d }· σ p ϵ ∥ µ ∥ 2 2 ) . By Gaussian tail bound, for any i ∈ [ M ] \{ 1 } , with probability at least 1 -δ/ 2 N 2 M 2 , where the last inequality is by σ V ≤ ˜ O ( ∥ w O ∥ -1 2 · min {∥ µ ∥ -1 2 , ( σ p √ d ) -1 }· d -1 4 h ) , N · SNR 2 = Ω(1) and d = ˜ Ω ( ϵ -2 N 2 d h ) and d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . Applying a union bound completes the proof.

Lemma D.10 (Population Loss) . Under the same conditions as Theorem 4.1, we have

Proof of Lemma D.10. Let event A to be the event that Lemma D.8 holds, and event B to be the event that Lemma D.9 holds. By a union bound, P ( A ∧ B ) ≥ 1 -δ/N 2 M . We decompose L D ( θ ( T 3 ) ) into the following form:

L D ( θ ( T 3 ) ) = E [ ✶ ( A ∧ B ) ℓ ( yf ( X , θ ( T 3 )))] + E [ ✶ (( A ∧ B ) c ) ℓ ( yf ( X , θ ( T 3 )))] (135) Then we bound E [ ✶ ( A ∧ B ) ℓ ( yf ( X , θ ( T 3 )))] and E [ ✶ (( A ∧ B ) c ) ℓ ( yf ( X , θ ( T 3 )))] respectively. Bounding E [ ( A ∧ B ) ℓ ( yf ( X , θ ( T 3 )))] : When event A and event B hold, we have

where the first inequality is by Lemma D.8, the second inequality is by (112) and Lemma D.9, the last inequality is by M = Θ(1) . Therefore,

Bounding E [ (( A ∧ B ) c ) ℓ ( yf ( X , θ ( T 3 )))] :

where the first inequality is by Cauchy-Schwartz inequality, the second inequality is by log(1 + exp( -z )) ≤ exp( -z ) .

Next we provide a bound for E [exp( -2 f ( X , θ ( T 3 )))] . Note that

and µ + W T 3 V w O ≥ 0 , we further have

Therefore,

where the third inequality is by | x | ≤ max { x, -x } , The fourth inequality is by the independence of ξ i W ( T 3 ) V w O .

We denote σ = ∥ W ( T 3 ) V w O ∥ 2 , then ξ i W ( T 3 ) V w O is a Gaussian distribution with mean zero and variance smaller than C 2 p σ 2 σ 2 p . Therefore,

where the first equality is by (134), the third equality is by σ V ≤ ˜ O ( ∥ w O ∥ -1 2 · min {∥ µ ∥ -1 2 , ( σ p √ d ) -1 } · d -1 4 h ) , N · SNR 2 ≥ Ω(1) and d = ˜ Ω ( ϵ -2 N 2 d h ) . Then we have

We further get

Plugging (137) and (142) into (135), we have

E Harmful Overfitting

In this section, we consider harmful overfitting case under the condition that N -1 · SNR -2 ≥ Ω(1) . The proofs in this section are based on the results in Section C, which hold with high probability.

E.1 Stage I

In Stage I, V ( t ) ± , V ( t ) n,i begin to pull apart until | V ( t ) ± | is sufficiently larger than | V ( t ) n,i | . At the same time, the inner products of q and k maintain their magnitude.

Lemma E.1 (Upper bound of V) . Let T 0 = O ( N ηd 1 4 h σ 2 p d ∥ w O ∥ 2 2 ) . Then under the same conditions as

Theorem 4.2, we have for t ∈ [0 , T 0 ] .

Proof of Lemma E.1. By Lemma B.2, we have

where the second inequality is by Lemma C.1 and -ℓ ′ ( t ) n ≤ 1 . Similarly, we have

By Definition B.1, we have

where the first inequality is by triangle inequality and t ≤ T 0 , the second inequality is by Lemma C.2, the last equality is by N -1 · SNR -2 = Ω(1) . Similarly, we have | V ( t ) -| = O ( d -1 4 h ) .

where the second inequality is by Lemma C.4 and -ℓ ′ ( t ) n ≤ 1 , the third inequality is by d = ˜ Ω ( ϵ -2 N 2 d h ) , the last inequality is by N · SNR 2 = Ω(1) . Then by Definition B.1, we have

where the first inequality is by triangle inequality, the second inequality is by Lemma C.2, which completes the proof.

Lemma E.2 (Inner Products Hold Magnitude) . Let T 0 = O ( N ηd 1 4 h σ 2 p d ∥ w O ∥ 2 2 ) . Then under the same conditions as Theorem 4.2, we have

The proof for Lemma E.2 is similar to D.4.

LemmaE.3 (V's Beginning of Learning Signals) . There exist T 1 = 4 M (3 M +1) N ηd 1 4 h ( C 2 p -24 M ) σ 2 p d ∥ w O ∥ 2 2 such that the second element of vector X n W ( T 1 ) V w O dominate its other elements, which means V ( T 1 ) n, 2 ≥ 3 M · max {| V ( T 1 ) + | , | V ( T 1 ) n,i |} for all n ∈ S + , i ∈ [ M ] \{ 1 , 2 } and V ( T 1 ) n, 2 ≤ -3 M · max {| V ( T 1 ) -| , | V ( T 1 ) n,i |} for all n ∈ S -, i ∈ [ M ] \{ 1 , 2 } .

Proof of Lemma E.3. According to Condition 4.1 where C 2 p = 25 M , we have Then we have

which satisfies the time condition in Lemma E.1 and Lemma E.2. Then by Lemma D.1 and Lemma D.2 we have

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] and t ∈ [0 , T 1 ] . Plugging them in the update rule for γ ( t ) V, + showed in Lemma B.2 and we have

where the last inequality is by N -1 · SNR -2 = Ω(1) .Then by Definition B.1 and taking a summation, we have

Similarly, we have

ρ

For n ∈ S + , by Lemma B.2, we have

(

t

+1)

(

t

)

where the second inequality is by Lemma C.4 and -ℓ ′ ( t ) n ≤ 1 , the third inequality is by d = ˜ Ω ( ϵ -2 N 2 d h ) . Then by Definition B.1, we have

for n ∈ S + . Similarly, for n ∈ S -, we have

for i ∈ [ M ] \{ 1 , 2 } , where the second inequality is by Lemma C.4 and -ℓ ′ ( t ) n ≤ 1 , the last inequality is by d = ˜ Ω ( ϵ -2 N 2 d h ) . Then by Definition B.1, we have

for i ∈ [ M ] \{ 1 , 2 } , where the first inequality is by triangle inequality, the second inequality is by Lemma C.2.

According to (145), (146), (148), (149) and (151), it is easy to verify that V ( T 1 ) n, 2 ≥ 3 M · max {| V ( T 1 ) + | , | V ( T 1 ) n,i |} for all n ∈ S + , i ∈ [ M ] \{ 1 , 2 } and V ( T 1 ) n, 2 ≤ -3 M · max {| V ( T 1 ) -| , | V ( T 1 ) n,i |} for all n ∈ S -, i ∈ [ M ] \{ 1 , 2 } , which completes the proof.

E.2 Stage II

In stage II, ⟨ q + , k + ⟩ , ⟨ q n,i , k + ⟩ grows while ⟨ q + , k n,j ⟩ , ⟨ q n,i , k n,j ⟩ decreases, resulting in attention focusing more and more on the signal and less on the noise. By the results of stage I, we have the following condition at the beginning of stage II

for all n S , i [ M ] 1 , 2 .

for all n ∈ S -, i ∈ [ M ] \{ 1 , 2 } . | V ( T 1 ) + | , | V ( T 1 ) -| , | V ( T 1 ) n,i | = O ( d -1 4 h ) , |⟨ q ( T 1 ) ± , k ( T 1 ) ± ⟩| , |⟨ q ( T 1 ) n,i , k ( T 1 ) ± ⟩| , |⟨ q ( T 1 ) ± , k ( T 1 ) n,j ⟩| , |⟨ q ( T 1 ) n,i , k ( T 1 ) n ′ ,j ⟩| = O ( max {∥ µ ∥ 2 2 , σ 2 p d } · σ 2 h · √ d h log(6 N 2 M 2 /δ ) ) , |⟨ q ( T 1 ) ± , q ( T 1 ) ∓ ⟩| , |⟨ q ( T 1 ) n,i , q ( T 1 ) ± ⟩| , |⟨ q ( T 1 ) n,i , q ( T 1 ) n ′ ,j ⟩| = O ( max {∥ µ ∥ 2 2 , σ 2 p d } · σ 2 h · √ d h log(6 N 2 M 2 /δ ) ) , |⟨ k ( T 1 ) ± , k ( T 1 ) ∓ ⟩| , |⟨ k ( T 1 ) n,i , k ( T 1 ) ± ⟩| , |⟨ k ( T 1 ) n,i , k ( T 1 ) n ′ ,j ⟩| = O ( max {∥ µ ∥ 2 2 , σ 2 p d } · σ 2 h · √ d h log(6 N 2 M 2 /δ ) ) , ∥ q ( T 1 ) ± ∥ 2 2 , ∥ k ( T 1 ) ± ∥ 2 2 = Θ( ∥ µ ∥ 2 2 σ 2 h d h ) , ∥ q ( T 1 ) n,i ∥ 2 2 , ∥ k ( T 1 ) n,i ∥ 2 2 = Θ( σ 2 p σ 2 h dd h ) i, j [ M ] 1 , n, n ′ [ N ] .

Some of the proofs at this stage are based on the above conditions.

Notations. To better characterize the gap between different inner products, we define the following notations:

Lemma E.4 (Upper bound of V) . Let T 0 = O ( N ησ 2 p d ∥ w O ∥ 2 2 log(6 N 2 M 2 /δ ) ) . Then under the same conditions as Theorem 4.2, we have

for t ∈ [0 , T 0 ] .

The proof of Lemma E.4 is similar to that of Lemma E.1, except that the time T 0 is changed.

Let T 2 = Θ ( N ησ 2 p d ∥ w O ∥ 2 2 log(6 N 2 M 2 /δ ) ) , then by Lemma D.6 and Lemma D.1 we have 1 2 -o (1) ≤ -ℓ ′ ( t ) n ≤ 1 2 + o (1) for n ∈ [ N ] , t ∈ [ T 1 , T 2 ] , which can simplify the calculations of α and β defined in Definition B.3 by replacing -ℓ ′ ( t ) n by their bounds. Next we prove the following four propositions J ( t ) , K ( t ) , L ( t ) , M ( t ) by induction on t for t ∈ [ T 1 , T 2 ] :

̸

̸

̸

V ( t ) n, 2 ≥ 3 M · max {| V ( T 1 ) + | , | V ( T 1 ) n,i |} for all n ∈ S + , i ∈ [ M ] \{ 1 , 2 } . V ( t ) n, 2 ≤ -3 M · max {| V ( T 1 ) -| , | V ( T 1 ) n,i |} for all n ∈ S -, i ∈ [ M ] \{ 1 , 2 } . | V ( t ) ± | ≤ O ( d -1 4 h ) + ηC 4 σ 2 p d ∥ w O ∥ 2 2 ( t -T 1 ) N | V ( t ) n,i | ≤ O ( d -1 4 h ) + ηC 4 σ 2 p d ∥ w O ∥ 2 2 ( t -T 1 ) N for i ∈ [ M ] \{ 1 } , n ∈ [ N ] . · K ( t ) : ∥ q ( t ) ± ∥ 2 2 , ∥ k ( t ) ± ∥ 2 2 = Θ ( ∥ µ ∥ 2 2 σ 2 h d h ) , ∥ q ( t ) n,i ∥ 2 2 , ∥ k ( t ) n,i ∥ 2 2 = Θ ( σ 2 p σ 2 h dd h ) , |⟨ q ( t ) + , q ( t ) -⟩| , |⟨ q ( t ) ± , q ( t ) n,i ⟩| , |⟨ q ( t ) n,i , q ( t ) n ′ ,j ⟩| = o (1) , |⟨ k ( t ) + , k ( t ) -⟩| , |⟨ k ( t ) ± , k ( t ) n,i ⟩| , |⟨ k ( t ) n,i , k ( t ) n ′ ,j ⟩| = o (1) , for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , i = j or n = n ′ . · L ( t ) : ⟨ q ( t +1) ± , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) ± , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) ± , k ( t +1) ± ⟩ - ⟨ q ( t ) ± , k ( t ) ± ⟩ ) ⟨ q ( t +1) ± , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) ± , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) ± , k ( t +1) n,j ⟩ - ⟨ q ( t ) ± , k ( t ) n,j ⟩ ) ⟨ q ( t +1) n,i , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) n,i , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) n,i , k ( t +1) ± ⟩ - ⟨ q ( t ) n,i , k ( t ) ± ⟩ ) ⟨ q ( t +1) n,i , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) n,i , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) n,i , k ( t +1) n,j ⟩ - ⟨ q ( t ) n,i , k ( t ) n,j ⟩ ) Ψ ( t +1) ± ≥ log ( exp(Ψ ( T 1 ) ± ) + η 2 C 7 ∥ µ ∥ 2 2 σ 2 p d ∥ w O ∥ 2 2 d 1 2 h N 2 ( log(6 N 2 M 2 /δ ) ) 2 · ( t -T 1 )( t -T 1 +1) ) Ψ ( t +1) n, ± ,j ≥ log ( exp(Ψ ( T 1 ) n, ± ,j ) + η 2 C 7 ∥ µ ∥ 2 2 σ 2 p d ∥ w O ∥ 2 2 d 1 2 h N 2 ( log(6 N 2 M 2 /δ ) ) 2 · ( t -T 1 )( t -T 1 +1) Ψ ( t +1) n,i, ± ≥ log ( exp(Ψ ( T 1 ) n,i, ± ) + η 2 C 7 σ 4 p d 2 ∥ w O ∥ 2 2 d 1 2 h N 2 ( log(6 N 2 M 2 /δ ) ) 2 · ( t -T 1 )( t -T 1 +1) ) Ψ ( t +1) n,i,j ≥ log ( exp(Ψ ( T 1 ) n,i,j ) + η 2 C 7 σ 4 p d 2 ∥ w O ∥ 2 2 d 1 2 h N 2 ( log(6 N 2 M 2 /δ ) ) 2 · ( t -T 1 )( t -T 1 +1) ) for i ∈ [ M ] \{ 1 } , j ∈ [ M ] \{ 1 , 2 } , n ∈ S ± . · M ( t ) : | Ψ ( t ) ± | , | Ψ ( t ) n, ± ,j ′ | , | Ψ ( t ) n,i, ± | , | Ψ ( t ) n,i,j ′ | ≤ log( d 1 2 h ) |⟨ q ( t ) ± , k ( t ) ± ⟩| , |⟨ q ( t ) ± , k ( t ) n,j ⟩| , |⟨ q ( t ) n,i , k ( t ) ± ⟩| , |⟨ q ( t ) n,i , k ( t ) n,j ⟩| ≤ 2 log( d 1 2 h ) |⟨ q ( t ) ± , k ( t ) ∓ ⟩| , |⟨ q ( t ) n,i , k ( t ) n,j ⟩| = o (1) for i, j ∈ [ M ] \{ 1 } , j ′ ∈ [ M ] \{ 1 , 2 } , n, n ∈ [ N ] , n = n .

By the results of Stage I, we know that J ( T 1 ) , K ( T 1 ) , M ( T 1 ) are true. To prove that J ( t ) , K ( t ) , L ( t ) and M ( t ) are true in stage 2, we give the following claims holds for t ∈ [ T 1 , T 2 ] :

The proofs of Claim 9-12 are similar to the proofs of Claim 5-8. Next, we show some important procedures for proving Claim 10.

Proof of Claim 10:

Similar to D.2.2, we have the dynamic of Ψ as follows

for n ∈ [ N ] , i ∈ [ M ] \{ 1 } , j ∈ [ M ] \{ 1 , 2 } , s ∈ [ T 1 , t ] . Next, we provide the logarithmic increasing lower bounds of Ψ . Recall that

Multiply both sides simultaneously by exp(Ψ ( s ) + ) and get

Taking a summation from T 1 to t and get

By the property that Ψ ( s ) + is monotonically increasing, we have

Similarly, we have

for n ∈ [ N ] , i ∈ [ M ] \{ 1 } , j ∈ [ M ] \{ 1 , 2 } .

E.3 Stage III

In Stage III, the outputs of ViT grow up and the loss derivatives are no longer at o (1) . We will carefully compute the growth rate of V ± and V n,i while keeping monitoring the monotonicity of ⟨ q , k ⟩ . By substituting t = T 2 = Θ ( N ησ 2 p d ∥ w O ∥ 2 2 ) into propositions J ( t ) , K ( t ) , L ( t ) , M ( t ) in Stage II, we have the following conditions at the beginning of stage III

̸

̸

Let T 3 = Θ ( N ηϵσ 2 p d ∥ w O ∥ 2 2 ) , Next we prove the following four propositions N ( t ) , O ( t ) , P ( t ) , Q ( t ) by induction on t for t ∈ [ T 2 , T 3 ] :

for i, j ∈ [ M ] \{ 1 } , j ′ ∈ [ M ] \{ 1 , 2 } , n, n ∈ [ N ] , n = n .

N ( t ) :

̸

V ( t ) n, 2 ≥ 3 M · max {| V ( T 1 ) + | , | V ( T 1 ) n,i |} for all n ∈ S + , i ∈ [ M ] \{ 1 , 2 } . V ( t ) n, 2 ≤ -3 M · max {| V ( T 1 ) -| , | V ( T 1 ) n,i |} for all n ∈ S -, i ∈ [ M ] \{ 1 , 2 } . | V ( t ) ± | , | V ( t ) n,i | = o (1) , for i ∈ [ M ] \{ 1 , 2 } , n ∈ [ N ] . log ( exp( V ( T 2 ) n, 2 ) + ηC 17 σ 2 p d ∥ w O ∥ 2 2 ( t -T 2 ) N ) ≤ V ( t ) n, 2 ≤ 2 log ( O ( 1 ϵ ) ) , for all n ∈ S + . -2 log ( O ( 1 ϵ ) ) ≤ V ( t ) n, 2 ≤ -log ( exp( -V ( T 2 ) n, 2 ) + ηC 17 σ 2 p d ∥ w O ∥ 2 2 ( t -T 2 ) N ) for all n ∈ S -. · O ( t ) : ∥ q ( t ) ± ∥ 2 2 , ∥ k ( t ) ± ∥ 2 2 = Θ( ∥ µ ∥ 2 2 σ 2 h d h ) , ∥ q ( t ) n,i ∥ 2 2 , ∥ k ( t ) n,i ∥ 2 2 = Θ ( σ 2 p σ 2 h dd h ) , |⟨ q ( t ) + , q ( t ) -⟩| , |⟨ q ( t ) ± , q ( t ) n,i ⟩| , |⟨ q ( t ) n,i , q ( t ) n ′ ,j ⟩| = o (1) , |⟨ k ( t ) + , k ( t ) -⟩| , |⟨ k ( t ) ± , k ( t ) n,i ⟩| , |⟨ k ( t ) n,i , k ( t ) n ′ ,j ⟩| = o (1) for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , i = j or n = n ′ . · P ( t ) : ⟨ q ( t +1) ± , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) ± , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) ± , k ( t +1) ± ⟩ - ⟨ q ( t ) ± , k ( t ) ± ⟩ ) ⟨ q ( t +1) ± , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) ± , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) ± , k ( t +1) n,j ⟩ - ⟨ q ( t ) ± , k ( t ) n,j ⟩ ) ⟨ q ( t +1) n,i , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) n,i , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) n,i , k ( t +1) ± ⟩ - ⟨ q ( t ) n,i , k ( t ) ± ⟩ ) ⟨ q ( t +1) n,i , k ( t +1) n, 2 ⟩ - ⟨ q ( t ) n,i , k ( t ) n, 2 ⟩ ≥ 2 ( ⟨ q ( t +1) n,i , k ( t +1) n,j ⟩ - ⟨ q ( t ) n,i , k ( t ) n,j ⟩ ) Ψ ( t +1) ± ≥ Ψ ( t ) ± Ψ ( t +1) n, ± ,j ≥ Ψ ( t ) n, ± ,j Ψ ( t +1) n,i, ± ≥ Ψ ( t ) n,i, ± Ψ ( t +1) n,i,j ≥ Ψ ( t ) n,i,j for i ∈ [ M ] \{ 1 } , j ∈ [ M ] \{ 1 , 2 } , n ∈ S ± .

̸

̸

̸

By the results of Stage II, we know that N ( T 1 ) , O ( T 2 ) , Q ( T 2 ) are true. To prove that N ( t ) , O ( t ) , P ( t ) and Q ( t ) are true in Stage III, we will prove the following claims holds for t ∈ [ T 2 , T 3 ] :

Lemma E.5 (Convergence of Training Loss) . There exist T = C 19 N ηϵσ 2 p d ∥ w O ∥ 2 2 such that

Proof of Lemma E.5. Similar to D.7, we have

for i ∈ [ M ] \{ 1 } , j ∈ [ M ] \{ 1 , 2 } , n ∈ [ N ] .

For n ∈ S + , substituting t = T = C 19 N ηϵσ 2 p d ∥ w O ∥ 2 2 into propositions N ( t ) and get

we bound f ( X n , θ ( t )) as follows

For the second inequality, by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) , we have O ( σ 2 p d ( log(6 N 2 M 2 /δ ) ) 3 ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 d 1 2 h ) = o (1) = log(1 + o (1)) ≤ log( C 20 ) as long as C 20 is sufficiently large.

Then we have

Similarly, we have ℓ ( t ) n ≤ ϵ for n ∈ S -. Therefore, we have L S ( θ ( T )) = 1 N N ∑ n =1 ℓ ( t ) n ≤ ϵ .

E.4 Population Loss

Lemma E.6 (Population Loss) . Under the same conditions as Theorem 4.2, let T 3 = Θ ( N ηϵσ 2 p d ∥ w O ∥ 2 2 ) , we have

Proof of Lemma E.6. Similar to (133)

Taking a summation, we have

Then ξ ⊤ i W ( T 3 ) V w O is a Gaussian random variable with mean zero and standard deviation smaller than O ( σ V ∥ w O ∥ 2 σ p √ d + N ϵ √ d ) . By Gaussian tail bound, for any i ∈ [ M ] \{ 1 } , with probability at least 1 -1 / 2 M ,

where the last inequality is by σ V ≤ ˜ O ( ∥ w O ∥ -1 2 · min {∥ µ ∥ -1 2 , ( σ p √ d ) -1 } · d -1 4 h ) and d = ˜ Ω ( ϵ -2 N 2 d h ) . Applying a union bound, with probability at least 1 -1 / 2 , | ξ ⊤ i W ( T 3 ) V w O | ≤ 1 / 2 . Recall that V ( T 3 ) ± = o (1) , with probability at least 1 -1 / 2 , we have

D ≥ · ≥

F Complete Calculation Process For Benign Overfitting

In this section, we show more calculation process under benign overfitting regime. The calculactions for harmful overfitting is similar.

F.1 Calculactions for α and β

In this subsection, we give the calculactions for α and β defined in Definition B.3.

Restatement of Lemma B.1. the gradients of loss function L S ( θ ) with respect to W Q , W K and W V are given by

We will give expressions for α and β defined in Definition B.3 based on these two equations above. By (178) and the orthogonal relation between µ and ξ , we have

where w ⊤ O W ( t ) ⊤ V X ⊤ n and X n W ( t ) K can be viewed in the following forms

Then we can express α ( t ) + , + and α ( t ) n, + ,i as follows

̸

Using the similar method as for ∆ q ( t ) + , we get the other α and β as follows

̸

for n ∈ S + ,

̸

for n ∈ S -,

̸

for n ∈ S + , for n ∈ S -.

F.2 Update Rules for Inner Products

In this subsection, we give the update rules for the inner products of q and k .

Restatement of Lemma B.3. The dynamics of xW Q W K X ⊤ can be characterized as follows

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] .

In addition to these equations above, we give the complete inner product update rule as follows

for n = n ,

̸

for n = n ,

̸

for n = n ,

̸

F.3 Proof of Lemma D.4

Let T 0 = O ( 1 ηd 1 4 h ∥ µ ∥ 2 2 ∥ w O ∥ 2 2 ) . By Lemma D.3, we have | V ( t ) + | , | V ( t ) -| , | V ( t ) n,i | = O ( d -1 4 h ) for t ∈ [0 , T 0 ] by Lemma D.3. Plugging this into the expression for α and β gives

where the inequality is by -ℓ ′ ( t ) n ≤ 1 and the property that attention is smaller than 1 ( e.g. exp( ⟨ q ( t ) + , k ( t ) + ⟩ ) exp( ⟨ q ( t ) + , k ( t ) + ⟩ )+ M ∑ j =2 exp( ⟨ q ( t ) + , k ( t ) n,j ⟩ ) ≤ 1 ). We also have

̸

where the inequality is by -ℓ ′ ( t ) n ≤ 1 and the property that attention is smaller than 1. We also have

where the inequality is by -ℓ ′ ( t ) n ≤ 1 , Lemma C.4 and the property that attention is smaller than 1. Similarly, we have

for t ∈ [0 , T 0 ] .

Next we use induction to show that the following proposition A ( t ) holds for t ∈ [0 , T 0 ] A ( t ) :

By Lemma C.3 we know that A (0) is true. Now we assume A (0) , . . . , A ( T ) is true, then we need to proof that A ( T + 1) is true. We first proof |⟨ q ( T +1) + , k ( T +1) + ⟩| = O ( max {∥ µ ∥ 2 2 , σ 2 p d } · σ 2 h · √ d h log(6 N 2 M 2 /δ ) ) , as an example.

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] .

Taking a summation, we obtain that

Similarly to ⟨ q ( t ) + , k ( t ) + ⟩ , it is easy to know that the inner product does not change by a magnitude more than the product of max { α, β } and max {⟨ q , q ⟩ , ⟨ k , k ⟩} in a single iteration, which can be expressed as follows

where the second equality is by the condition that N · SNR 2 = Ω(1) . Taking a summation, we obtain that

It is clear that the magnitude of ⟨ q ( T +1) , k ( T +1) ⟩ - ⟨ q (0) , k (0) ⟩ is smaller than max {∥ µ ∥ 2 2 , σ 2 p d } · σ 2 h · √ d h log(6 N 2 M 2 /δ ) , Thus the magnitude of the bound for ⟨ q ( T +1) , k ( T +1) ⟩ is the same as that of ⟨ q ( T ) , k ( T ) ⟩ . The proof for ⟨ q ( T +1) , q ( T +1) ⟩ and ⟨ k ( T +1) , k ( T +1) ⟩ is exactly the same, and we can conclude the proof by an induction.

F.4 Lower Bounds of α and β

In this subsection, we present some bounds for α and β which can be used in D.2 and D.3. All the calculations in this subsection are based on the precise expression for α and β in F.1 and assume that B ( T 1 ) , . . . , B ( s ) , D ( T 1 ) , . . . , D ( s -1) hold ( s ∈ [ T 1 , t ] ). Then the following propositions hold:

Now we give the bounds respectively for α ( s ) + , + , α ( s ) n, + ,i , α ( s ) -, -, α ( s ) n, -,i , α ( s ) n,i, + , α ( s ) n,i, -, α ( s ) n,i,n ′ ,i ′ , β ( s ) + , + , β ( s ) n, + ,i , β ( s ) -, -, β ( s ) n, -,i , β ( s ) n,i, + , β ( s ) n,i, -, β ( s ) n,i,n ′ ,i ′ .

where the first inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | , the second inequality is by the fact that the sum of attention equal to 1. Similarly, we have

̸

̸

where the third inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | and softmax ( ⟨ q ( s ) + , k ( s ) + ⟩ ) ≥ 1 M -o (1) . Similarly, we have

where the first inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | , the second inequality is by M ∑ k =2 softmax ( ⟨ q ( s ) n,i , k ( s ) n,k ⟩ ) = ( 1 -softmax ( ⟨ q ( s ) n,i , k ( s ) + ⟩ ) ) . Similarly, we have

where the { lower order term } is by the property that ∥ ξ n ′ ,i ′ ∥ 2 2 is much larger than ⟨ ξ n ′ ,i ′ , ξ n,i ⟩ in Lemma C.4 and the condition d = ˜ Ω ( ϵ -2 N 2 d h ) , the first inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | , the last inequality is by ∥ ξ n ′ ,i ′ ∥ 2 2 ≥ σ 2 p d 2 and we absorb the { lower order term } . Similarly, we have

where the { lower order term } is by the property that ∥ ξ n ′ ,i ′ ∥ 2 2 is much larger than ⟨ ξ n ′ ,i ′ , ξ n,i ⟩ in Lemma C.4 and the condition d = ˜ Ω ( ϵ -2 N 2 d h ) , the first inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | , the last inequality is by ∥ ξ n ′ ,i ′ ∥ 2 2 ≥ σ 2 p d 2 and we absorb the { lower order term } . Similarly, we have

̸

(

s

)

̸

̸

where the { lower order term } is by the property that ∥ ξ n,i ∥ 2 2 is much larger than ⟨ ξ n ′ ,i ′ , ξ n,i ⟩ in Lemma C.4 and the condition d = ˜ Ω ( ϵ -2 N 2 d h ) , the third inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | and softmax ( ⟨ q ( s ) n,i , k ( s ) + ⟩ ) ≥ 1 M -o (1) , the last inequality is by ∥ ξ n,i ∥ 2 2 ≥ σ 2 p d 2 and we absorb the { lower order term } . Similarly, we have

where the { lower order term } is by the property that ∥ ξ n,i ∥ 2 2 is much larger than ⟨ ξ n ′ ,i ′ , ξ n,i ⟩ in Lemma C.4 and the condition d = ˜ Ω ( ϵ -2 N 2 d h ) , the second inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | , the last inequality is by ∥ ξ n,j ∥ 2 2 ≥ σ 2 p d 2 and we absorb the { lower order term } . Similarly, we have

Summarizing the above equations, we have the signs of α and β as follows:

F.5 Lower Bounds of ⟨ q, k ⟩

In order to give the lower bounds for ⟨ q , k ⟩ , we need to rewrite the bounds of α and β in a more concise form. We first expand the equations in F.4 under the assumption that B ( s ) holds for s ∈ [ T 1 , t ] .

where the ⟨ q ( s ) + , k ( s ) n,j ⟩ in the second inequality is a particular choice (we will characterize the dynamic of ⟨ q ( s ) + , k ( s ) + ⟩ - ⟨ q ( s ) + , k ( s ) n,j ⟩ ), the third inequality is by softmax ( ⟨ q ( s ) + , k ( s ) + ⟩ ) ≥ ( 1 M -o (1) ) . In the fourth inequality, by ⟨ q ( T 1 ) , k ( T 1 ) ⟩ = o (1) and the monotonicity of ⟨ q ( s ) , k ( s ) ⟩ ( ⟨ q ( s ) + , k ( s ) + ⟩ is increasing and ⟨ q ( s ) + , k ( s ) n,j ⟩ is decreasing), there exist a constant C such that C exp( ⟨ q ( s ) + , k ( s ) + ⟩ ) ≥ exp( ⟨ q ( s ) + , k ( s ) + ⟩ ) + M ∑ j ′ =2 exp( ⟨ q ( s ) + , k ( s ) n,j ′ ⟩ ) . In the last inequality, we plugging the lower bounds of

V ( s ) + and -ℓ ′ ( s ) n and then absorb all the constant factors. Similarly, we have

where the inequalities is similar to (223). Similarly, we have

where the inequalities is similar to (223), and ⟨ q ( s ) n,i , k ( s ) n,j ⟩ is a particular choice (we will characterize the dynamic of ⟨ q ( s ) n,i , k ( s ) + ⟩ - ⟨ q ( s ) n,i , k ( s ) n,j ⟩ ). Similarly, we have

With the concise lower bounds for α and β above and proposition C ( s ) , we will give the lower bounds for the dynamics of ⟨ q , k ⟩ .

where the first inequality is by (223), (224), the bound of ∥ q ( s ) + ∥ 2 2 , ∥ k ( s ) + ∥ 2 2 in stage II, the second inequality is by absorbing the { lower order term } and the constant factors. Similarly, we have

where the first inequality is by (227), (232), the bound of ∥ k ( s ) n,j ∥ 2 2 , ∥ q ( s ) + ∥ 2 2 in stage II, the second inequality is by absorbing the { lower order term } and the constant factors. Similarly, we have

where the first inequality is by (230), (228), the bound of ∥ k ( s ) + ∥ 2 2 , ∥ q ( s ) n,i ∥ 2 2 in stage II, the second inequality is by absorbing the { lower order term } and the constant factors. Similarly, we have

where the first inequality is by (234), (235), the bound of ∥ k ( s ) n,j ∥ 2 2 , ∥ q ( s ) n,i ∥ 2 2 in stage II, the second inequality is by absorbing the { lower order term } and the constant factors.

F.6 Upper Bounds of ⟨ q, k ⟩

In order to give the upper bounds of ⟨ q , k ⟩ in stage II, we need to give the upper bounds of α and β based on the equations in F.1 under the assumption that D ( T 1 ) , . . . , D ( s -1) hold for s ∈ [ T 1 , t ] .

where the first inequality is by -ℓ ′ ( s ) n ≤ 1 and softmax ( ⟨ q ( s ) + , k ( s ) + ⟩ ) ≤ 1 . For the second inequality, we first consider M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) + , k ( s ) + ⟩ )+ M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) ≤ M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) + , k ( s ) + ⟩ ) , then by the monotonicity of ⟨ q ( s ) + , k ( s ) n,j ⟩ and ⟨ q ( T 1 ) + , k ( T 1 ) n,j ⟩ = o (1) we have M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) ≤ C for s ∈ [ T 1 , t ] . The last inequality is by V ( s ) + , V ( s ) n,i = o (1) for s ∈ [ T 1 , t ] and absorbing the constant factors. Similarly, we have

̸

̸

where the second inequality is by V ( s ) + ≥ 3 M · | V ( s ) n,i | , -ℓ ′ ( s ) n ≤ 1 and the property that attention < 1 . For the third inequality, we consider exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) + , k ( s ) + ⟩ )+ M ∑ j ′ =2 exp( ⟨ q ( s ) + , k ( s ) n,j ′ ⟩ ) ≤ exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) + , k ( s ) + ⟩ ) first, then by the monotonicity of ⟨ q ( s ) + , k ( s ) + ⟩ and ⟨ q ( T 1 ) + , k ( T 1 ) + ⟩ we have exp( ⟨ q ( s ) + , k ( s ) + ⟩ ) ≥ C for s ∈ [ T 1 , t ] . The last inequality is by V ( s ) + = o (1) for s ∈ [ T 1 , t ] and absorbing the constant factors. Similarly, we have

where the second inequality is by softmax ( ⟨ q ( s ) n,i , k ( s ) + ⟩ ) ≤ 1 . For the second inequality, we first consider M ∑ j =2 exp( ⟨ q ( s ) n,i , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) n,i , k ( s ) + ⟩ )+ M ∑ j =2 exp( ⟨ q ( s ) n,i , k ( s ) n,j ⟩ ) ≤ M ∑ j =2 exp( ⟨ q ( s ) n,i , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) n,i , k ( s ) + ⟩ ) , then by the monotonicity of ⟨ q ( s ) n,i , k ( s ) n,j ⟩ and ⟨ q ( T 1 ) n,i , k ( T 1 ) n,j ⟩ = o (1) we have M ∑ j =2 exp( ⟨ q ( s ) n,i , k ( s ) n,j ⟩ ) ≤ C for s ∈ [ T 1 , t ] . The last inequality is by V ( s ) + , V ( s ) n,i = o (1) for s ∈ [ T 1 , t ] and absorbing the constant factors. Similarly, we have

where most of these processes are similar to the other equations above, and the last inequality we absorb the constant factors and the { lower order term } . Similarly, we have

Similar to F.5, we apply the bounds of α and β above to give the upper bounds for the dynamics ⟨ q , k ⟩ .

similarly, we have

similarly, we have

similarly, we have

F.7 Bounds for the Sum of α and β

The gradients of the inner products of q and k contain a lot of coefficients α and β , and in order to conveniently give the upper bounds of some lower order inner products, we will give upper bounds t ).

for the summation of α and β (e.g. ∑ s = T 1 | α ( s ) + , + |

Note that in the Jacobi matrix of the Softmax function, the elements on the diagonal are softmax ( a i ) · ( 1 -softmax ( a i ) ) and the elements on the off-diagonal are softmax ( a i ) · softmax ( a j ) . In Stage II, the attentions on signals µ ± increase and the attentions on noises ξ decrease, then we can consider the following cases

if a i = ⟨ q + , k + ⟩ or a i = ⟨ q i , k + ⟩ , softmax ( a i ) has a constant upper bound 1, ( 1 -softmax ( a i ) ) decreases as softmax ( a i ) increases. So the upper bound of softmax ( a i ) · ( 1 -softmax ( a i ) ) decreases as softmax ( a i ) increases.

if a i = ⟨ q + , k j ⟩ or a i = ⟨ q i , k j ⟩ , ( 1 -softmax ( a i ) ) has a constant upper bound 1. So the upper bound of softmax ( a i ) · ( 1 -softmax ( a i ) ) decreases as softmax ( a i ) decreases.

if a j = ⟨ q + , k j ⟩ or a j = ⟨ q i , k j ⟩ , softmax ( a i ) has a constant upper bound 1. So the upper bound of softmax ( a i ) · softmax ( a j ) decreases as softmax ( a j ) decreases.

Based on the above cases, we first study the bounds of the following terms

1 -softmax ( ⟨ q ( s ) + , k ( s ) + ⟩ )

1 -softmax ( ⟨ q ( s ) n,i , k ( s ) + ⟩ )

softmax ( ⟨ q ( s ) + , k ( s ) n,j ⟩ )

softmax ( ⟨ q ( s ) n,i , k ( s ) n,j ⟩ )

Note that 1 -softmax ( ⟨ q ( s ) + , k ( s ) + ⟩ ) = ∑ j softmax ( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) and 1 -softmax ( ⟨ q ( s ) n,i , k ( s ) + ⟩ ) = ∑ j softmax ( ⟨ q ( s ) n,i , k ( s ) n,j ⟩ ) , we only need to give the upper bounds for softmax ( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) and softmax ( ⟨ q ( s ) n,i , k ( s ) n,j ⟩ ) .

Assume that the propositions B ( T 1 ) , . . . , B ( s ) , D ( T 1 ) , . . . , D ( s -1) hold ( s ∈ [ T 1 , t ] ), we have

Then we have

For the first inequality, by ⟨ q ( T 1 ) , k ( T 1 ) ⟩ = o (1) and the monotonicity of ⟨ q ( s ) , k ( s ) ⟩ ( ⟨ q ( s ) ± , k ( s ) ± ⟩ is increasing and ⟨ q ( s ) ± , k ( s ) n,j ⟩ is decreasing), there exist a constant C such that C exp( ⟨ q ( s ) ± , k ( s ) ± ⟩ ) ≥ exp( ⟨ q ( s ) ± , k ( s ) ± ⟩ ) + M ∑ j ′ =2 exp( ⟨ q ( s ) ± , k ( s ) n,j ′ ⟩ ) . The second inequality is by plugging (251). For the last inequality, by Λ ( T 1 ) n, ± ,j = o (1) , there exist a constant C 13 such that C 13 ≤ C exp(Λ ( T 1 ) n, ± ,j ) and C 13 ≤ C 8 C . Similarly, we have

Plugging (250),(253) and (254) into the expressions of α , β we have

where the inequality is by ax + b x ≥ 2 √ ab for x > 0 , the third equality is by absorbing the lower order term η 2 C 13 ∥ µ ∥ 4 2 ∥ w O ∥ 2 2 d 1 2 h N ( log(6 N 2 M 2 /δ ) ) 2 , the last equality is by ∥ w O ∥ 2 = Θ(1) . Plugging this into (255) and get

Similarly, we have

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i ∈ [ M ] \{ 1 } , n ∈ S ± , the last equality is by N · SNR 2 ≥ Ω(1) .

for i ∈ [ M ] \{ 1 } , n ∈ S ± , the last equality is by N · SNR 2 ≥ Ω(1) .

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] , the last equality is by N · SNR 2 ≥ Ω(1) .

̸

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , n = n ′ , the last equality is by N · SNR 2 ≥ Ω(1) . Taking a summation we obtain that

where the last equality is by ∥ w O ∥ = Θ(1) . Similarly, we have

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

t

∑

|

α

(

s

)

n,i,

+

,

s

=

T

1

|

for i, j ∈ [ M ] \{ 1 } , n ∈ S ± .

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , n = n ′ .

With these sums of α and β above, we can easily prove Claim 3 and Claim 4.

F.8 Proof of Claim 3

In this subsection, we assume that E ( T 1 ) , . . . , E ( t ) hold, and then proof that C ( t +1) is true with the result of F.7.

̸

t

∑

s

=

T

1

|

α

(

s

)

n,i,

-

,

|

t

∑

s

=

T

1

|

α

(

s

)

n,i,n,j

|

,

t

∑

s

=

T

1

|

β

(

s

)

n,j,n,i

|

=

O

(

-

d

h

1

4

)

(269)

where the first inequality is by triangle inequality, the second inequality is by 199. Since σ 2 h ≥ ( max { σ 2 p d, ∥ µ ∥ 2 2 } ) -1 · d -1 2 h (log(6 N 2 M 2 /δ )) -2 and d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) , we have N 1 2 d -1 4 h log( d 1 2 h ) = o ( ∥ µ ∥ 2 2 σ 2 h d h ) , so ∥ q ( t +1) + ∥ 2 2 = ∥ q ( T 1 ) + ∥ 2 2 + o ( ∥ µ ∥ 2 2 σ 2 h d h ) = Θ( ∥ µ ∥ 2 2 σ 2 h d h ) . Similarly, we have

for i ∈ [ M ] \{ 1 } , n ∈ [ N ] .

̸

where the first inequality is triangle inequality, the second inequality is by (198), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) .

̸

where the first inequality is triangle inequality, the second inequality is by (202), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) . Similarly, we have |⟨ q ( t +1) -, q ( t +1) n,i ⟩| = o (1) .

̸

̸

̸

for i, j ∈ [ M ] \{ 1 } , i = j, n ∈ [ N ] . The first inequality is triangle inequality, the second inequality is by (204), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) .

̸

̸

̸

for i, j ∈ [ M ] \{ 1 } , n, n ∈ [ N ] , n = n . The first inequality is triangle inequality, the second inequality is by (205), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) .

̸

̸

where the first inequality is triangle inequality, the second inequality is by (214), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) .

∣

∣

∣

̸

̸

where the first inequality is triangle inequality, the second inequality is by (209), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) . Similarly, we have |⟨ k ( t +1) , k ( t +1) = o (1) .

̸

̸

̸

for i, j ∈ [ M ] \{ 1 } , i = j, n ∈ [ N ] . The first inequality is triangle inequality, the second inequality is by (211), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) .

̸

̸

for i, j ∈ [ M ] \{ 1 } , n, n ∈ [ N ] , n = n . The first inequality is triangle inequality, the second inequality is by (212), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) .

F.9 Upper Bounds of ⟨ q, k ⟩

In order to give the upper bounds for ⟨ q , k ⟩ in stage III, we need to give the upper bounds of α and β based on the equations in F.1. The main difference between this subsection and F.6 is that the bounds of | V ± | , | V n,i | is log ( O ( 1 ϵ ) ) in this subsection, while the bounds of | V ± | , | V n,i | is log ( O ( 1 ϵ ) ) in F.6, resulting in different bounds for α and β . Now we take α ( s ) + , + as an example

̸

where the first inequality is by -ℓ ′ ( s ) n ≤ 1 and softmax ( ⟨ q ( s ) + , k ( s ) + ⟩ ) ≤ 1 . For the second inequality, we first consider M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) + , k ( s ) + ⟩ )+ M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) ≤ M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) exp( ⟨ q ( s ) + , k ( s ) + ⟩ ) , then by the monotonicity of ⟨ q ( s ) + , k ( s ) n,j ⟩ and ⟨ q ( T 1 ) + , k ( T 1 ) n,j ⟩ = o (1) we have M ∑ j =2 exp( ⟨ q ( s ) + , k ( s ) n,j ⟩ ) ≤ C for t ∈ [ T 1 , T 3 ] . The

last inequality is by V ( s ) + , | V ( s ) n,i | ≤ 2 log ( O ( 1 ϵ ) ) for t ∈ [ T 2 , T 3 ] and absorbing the constant factors. Similar to F.6, we can give the bounds for the other α and β as follows:

Similar to F.5, we apply the bounds of α and β above to give the upper bounds for the dynamics ⟨ q , k ⟩ .

similarly, we have

similarly, we have

similarly, we have

F.10 Bounds for the Sum of α and β

Assume that the propositions F ( T 2 ) , . . . , F ( s ) , H ( T 2 ) , . . . , H ( s -1) hold ( s ∈ [ T 1 , t ] ), we have

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] , s ∈ [ T 2 , t ] . Similar to (64) and (65), we have

Plugging (282),(283) into the expressions of α , β and letting O ( log ( O ( 1 ϵ ) ) ) be the upper bound for | V ( s ) ± | , | V ( s ) n,i | we have

Similarly, we have

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i ∈ [ M ] \{ 1 } , n ∈ S ± , the last equality is by N · SNR 2 ≥ Ω(1) .

for i, j ∈ [ M ] \{ 1 } , n ∈ [ N ] .

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , n = n ′ . Taking a summation we obtain that

̸

where the last equality is by ∥ w O ∥ = Θ(1) . Similarly, we have

for i ∈ [ M ] \{ 1 } , n ∈ S ± .

for i, j ∈ [ M ] \{ 1 } , n ∈ S ± .

for i, j ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , n = n ′ .

With these sums of α and β above, we can easily prove Claim 7 and Claim 8.

F.11 Proof of Claim 7

In this subsection, we assume that I ( T 2 ) , . . . , I ( t ) hold, and then proof that G ( t +1) is true with the result of F.10.

where the first inequality is by triangle inequality, the second inequality is by 199, the third inequality is by t ≤ T 3 . Since σ 2 h ≥ ( max { σ 2 p d, ∥ µ ∥ 2 2 } ) -1 · d -1 2 h (log(6 N 2 M 2 /δ )) -2 and d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) , we have N ( log(6 N 2 M 2 /δ ) ) 3 log ( O ( 1 ϵ ) ) log( ϵ -1 d 1 2 h ) ϵd 1 2 h =

̸

o ( ∥ µ ∥ 2 2 σ 2 h d h ) , so ∥ q ( t +1) + ∥ 2 2 = ∥ q ( T 2 ) + ∥ 2 2 + o ( ∥ µ ∥ 2 2 σ 2 h d h ) = Θ( ∥ µ ∥ 2 2 σ 2 h d h ) . Similarly, we have ∣ ∣ ∣ ∥ q ( t +1) -∥ 2 2 -∥ q ( t +1) -∥ 2 2 ∣ ∣ ∣ = O ( N ( log(6 N 2 M 2 /δ ) ) 3 log ( O ( 1 ϵ ) ) log( ϵ -1 d 1 2 h ) ϵd 1 2 h ) = o ( ∥ µ ∥ 2 2 σ 2 h d h ) , ∣ ∣ ∣ ∥ k ( t +1) ± ∥ 2 2 -∥ k ( t +1) ± ∥ 2 2 ∣ ∣ ∣ = O ( ( 1 + SNR 2 ) N ( log(6 N 2 M 2 /δ ) ) 3 log ( O ( 1 ϵ ) ) log( ϵ -1 d 1 2 h ) ϵd 1 2 h ) = o ( ∥ µ ∥ 2 2 σ 2 h d h ) , ∣ ∣ ∣ ∥ q ( t +1) n,i ∥ 2 2 -∥ q ( t +1) n,i ∥ 2 2 ∣ ∣ ∣ = O ( ( log(6 N 2 M 2 /δ ) ) 3 log ( O ( 1 ϵ ) ) log( ϵ -1 d 1 2 h ) ϵd 1 2 h ) = o ( σ 2 p σ 2 h dd h ) , ∣ ∣ ∣ ∥ k ( t +1) n,i ∥ 2 2 -∥ k ( t +1) n,i ∥ 2 2 ∣ ∣ ∣ = O ( N ( log(6 N 2 M 2 /δ ) ) 3 log ( O ( 1 ϵ ) ) log( ϵ -1 d 1 2 h ) ϵd 1 2 h ) = o ( σ 2 p σ 2 h dd h ) , so we have ∥ q ( t +1) ± ∥ 2 2 , ∥ k ( t +1) ± ∥ 2 2 = Θ( ∥ µ ∥ 2 2 σ 2 h d h ) , ∥ q ( t +1) n,i ∥ 2 2 , ∥ k ( t +1) n,i ∥ 2 2 = Θ( σ 2 p σ 2 h dd h )

for i ∈ [ M ] \{ 1 } , n ∈ [ N ] .

̸

where the first inequality is triangle inequality, the second inequality is by (198), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) .

̸

= o (1) , where the first inequality is triangle inequality, the second inequality is by (202), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) . Similarly, we have |⟨ q ( t +1) -, q ( t +1) n,i ⟩| = o (1) .

̸

̸

̸

̸

̸

̸

for i, j ∈ [ M ] \{ 1 } , n, n ∈ [ N ] , n = n . The first inequality is triangle inequality, the second inequality is by (205), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d =

̸

̸

where the first inequality is triangle inequality, the second inequality is by (214), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) .

̸

where the first inequality is triangle inequality, the second inequality is by (209), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) . Similarly, we have |⟨ k ( t +1) -, k ( t +1) n,i ⟩| = o (1) .

̸

̸

̸

̸

for i, j ∈ [ M ] \{ 1 } , i = j, n ∈ [ N ] . The first inequality is triangle inequality, the second inequality is by (211), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) . |⟨ k ( t +1) n,i , k ( t +1) n,j ⟩| ≤ |⟨ k ( T 2 ) n,i , k ( T 2 ) n,j ⟩| + t ∑ s = T 2 ∣ ∣ ∣ ⟨ k ( s +1) n,i , k ( s +1) n,j ⟩ - ⟨ k ( s ) n,i , k ( s ) n,j ⟩ ∣ ∣ ∣ ≤ |⟨ k ( T 2 ) n,i , k ( T 2 ) n,j ⟩| + t ∑ s = T 2 ∣ ∣ ∣ β ( s ) n,i, + ⟨ q ( s ) + , k ( s ) n,j ⟩ + β ( s ) n,i, -⟨ q ( s ) -, k ( s ) n,j ⟩ + N ∑ n ′ =1 M ∑ l =2 β ( s ) n,i,n ′ ,l ⟨ q ( s ) n ′ ,l , k ( s ) n,j ⟩ + β ( s ) n,j, + ⟨ q ( s ) + , k ( s ) n,i ⟩ + β ( s ) n,j, -⟨ q ( s ) -, k ( s ) n,i ⟩ + N ∑ n ′ =1 M ∑ l =2 β ( s ) n,j,n ′ ,l ⟨ q ( s ) n ′ ,l , k ( s ) n,i ⟩

̸

̸

for i, j ∈ [ M ] \{ 1 } , n, n ∈ [ N ] , n = n . The first inequality is triangle inequality, the second inequality is by (212), the last equality is by d h = ˜ Ω ( max { SNR 4 , SNR -4 } N 2 ϵ -2 ) and d = ˜ Ω ( ϵ -2 N 2 d h ) .

F.12 Explanations of Lower Order Terms

In this section, we provide some explanations of lower order terms to demonstrate the rigor of our proof.

To bound the so-call { lower order term } , we condition that dimensions d , d h are sufficiently large and learning rate η is sufficiently small. Next, we show how we utilize these three parameters.

Sufficiently mall learning rate η : Recall the dynamics of QK

Note that terms ⟨ ∆ q ( t ) , k ( t ) ⟩ , ⟨ q ( t ) , ∆ k ( t ) ⟩ contain factor η , and term ⟨ ∆ q ( t ) , ∆ k ( t ) ⟩ contains factor η 2 . Therefore, as long as η is sufficiently small, ⟨ ∆ q ( t ) , ∆ k ( t ) ⟩ is sufficiently small than ⟨ ∆ q ( t ) , k ( t ) ⟩ and ⟨ q ( t ) , ∆ k ( t ) ⟩ . Now we take the dynamic of ⟨ q ( t ) + , k ( t ) + ⟩ as an example.

̸

∈ +

Under benign overfitting regime, we have ∥ k ( t ) + ∥ 2 2 = Θ( ∥ µ ∥ 2 2 σ 2 h d h ) , ⟨ q ( t ) + , k ( t ) + ⟩ ≤ log( ϵ -1 d 1 2 h ) and β ( t ) + , + = O ( η ∥ µ ∥ 2 2 ) . Therefore, as long as η = o ( σ 2 h d h (log( ϵ -1 d 1 2 h )) -1 ) , we have

Similar method can be used for other items in ( α ( t ) + , + k ( t ) + + ∑ n ∈ S + M ∑ i =2 α ( t ) n, + ,i k ( t ) n,i ) · ( β ( t ) + , + q ( t ) ⊤ + +

∑ n ∈ S + M ∑ i =2 β ( t ) n, + ,i q ( t ) ⊤ n,i ) . At last, we have

Sufficiently large dimension d h : Take (297) as an example, Noting that ∥ k ( t ) + ∥ 2 2 = O ( η ∥ µ ∥ 2 2 ) and ⟨ k ( t ) + , k ( t ) n,i ⟩ = o (1) . Therefore, as long as d h is sufficiently large, ∥ k ( t ) + ∥ 2 2 is much larger than ⟨ k ( t ) + , k ( t ) n,i ⟩ . Besides, by the property that the sum of each row and column of matrix ( diag ( φ ( t ) n, 1 ) -φ ( t ) ⊤ n, 1 φ ( t ) n, 1 ) is 0, we have α ( t ) + , + + ∑ n ∈ S + M ∑ i =2 α ( t ) n, + ,i = 0 . We also prove α ( t ) + , + ≥ 0 and α ( t ) n, + ,i ≤ 0 under benign overfitting regime, thus the magnitude of α ( t ) n, + ,i is smaller than α ( t ) + , + . All in all, it can be proved that ∑ n ∈ S + M ∑ i =2 α ( t ) n, + ,i ⟨ k ( t ) + , k ( t ) n,i ⟩ = o ( α ( t ) + , + ∥ k ( t ) + ∥ 2 2 ) . This method can be further applied to β ( t ) + , + ∥ q ( t ) + ∥ 2 2 + ∑ n S M ∑ i =2 β ( t ) n, + ,i ⟨ q ( t ) + , q ( t ) n,i ⟩ . At last, we can simplify (297) as follows:

Sufficiently large dimension d : Take α ( t ) n ′ ,i ′ , + as an example:

̸

̸

Note that ⟨ ξ n ′ ,i ′ , ξ n,i ⟩ can be divided into two types: ∥ ξ n ′ ,i ′ ∥ 2 2 and ⟨ ξ n ′ ,i ′ , ξ n,i ⟩ for i = i ′ or n = n ′ By Lemma C.4, we have for i, i ′ ∈ [ M ] \{ 1 } , n, n ′ ∈ [ N ] , i = i ′ or n = n ′ .

As long as d is sufficiently large, ∥ ξ n ′ ,i ′ ∥ 2 2 is much larger than ⟨ ξ n ′ ,i ′ , ξ n,i ⟩ for i = i ′ or n = n ′ . Therefore, (299) can be further simplified as follows:

̸

̸

̸

̸

G Takeaways for Practitioners

Our theoretical results mainly focuse on the impact of different N and SNR on the generalization performance. So we can provide guidance from the perspective of increasing N , SNR and N · SNR 2 . The following are some practical scenarios

.

Data Augmentation: Researchers sometimes employ the technique of data augmentation by introducing controlled noise into their datasets. From the perspective of our paper's results, this method reduce SNR but improve N because we generate "new" data point by adding noises. As reducing SNR may be harmful to generalization performance, we must make sure that we use enough data points to train the model( enough sample size N ).

Semi-Supervised Learning: Semi-supervised learning is useful when you have a small amount of labeled data and a large amount of unlabeled data. Labeled data can be seen as data with high SNR, while unlabeled data with low SNR because for some unlabeled samples, we may mistake their labels, making them equivalent to noises. In this scenario, we need to ensure that we have sufficient unlabeled data (enough sample size N ) and make full use of labeled data (high SNR data points).

Overall, we need to consider both the sample size N and the signal-to-noise ratio SNR to train the model.

H Broader Impacts

This work focus on theoretically studying the training dynamics and generalization of Transformer in Vision. The techniques used in this paper may be generalized to study other abilities of Transformer or other network models. Besides, the theoretical results in this paper may inspire more attempts at training large foundational models with high-quality data. We do not foresee any form of negative social impact induced by our work.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract reflects the paper's scope, and the introduction reflects the paper's contributions.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss our limitation in the conclusion section.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide the assumptions in Condition 4.1. The detailed proofs are in the appendix. We also provide a proof sketch in the main paper.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide the experiments setting in this paper, and the codes are submitted as supplementary material.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The codes are submitted as supplementary material, which include a readme file.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide detailed data generation methods and hyperparameters in the experiments setting.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: we plot the 2-sigma error bar for the training dynamics in Appendix A Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide the information on the computer resources in experiments setting. All the experiments in this paper can be performed within hours.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This work focuses on theoretical study of ViT, and all the data is synthesized. We see no ethical or immediate negative societal consequence of our work, and it will not violate the Code Of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the societal impacts in Appendix H. There are no negative social impacts in this paper.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: All the data in this paper is synthesized. This paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: All the data in this paper is synthesized. This paper does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: All the data in this paper is synthesized. This paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
The paper investigates the benign overfitting phenomenon in Vision Transformers. By examining the training dynamics and generalization of a two-layer Transformer model, it establishes a condition to differentiate between benign and harmful overfitting based on the signal-to-noise ratio in the data model. Theoretical results are supported by experimental simulations.

STRENGTHS:
1. The paper provides a deep theoretical understanding of how Vision Transformers can achieve benign overfitting, filling a gap in the literature.
2. The derivation of the conditions for benign and harmful overfitting is mathematically rigorous and well-founded. Strong conclusions drawn based on appropriate assumptions.
3. The theoretical findings are supported by experimental simulations, which confirm the sharp condition separating benign and harmful overfitting.

WEAKNESSES:
1. The sparsity assumption is too strong, as it relies on the signal being contained within one patch and the noise being contained within another, which may not align with real-world data distributions.
2. Experiments are limited to settings that perfectly align with theoretical assumptions. The paper should also explore scenarios that deviate from these assumptions to understand the limitations of the theoretical results.

QUESTIONS:
The authors are suggested to explore scenarios that deviate from these assumptions to understand the limitations of the theoretical results. For example, more signal patches and more noise patches are contained in all patches. Tests on real-world tasks are better.

LIMITATIONS:
See Questions and Weaknesses.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-08

comments: SUMMARY:
This paper provides a sharp theoretical characterization of the transition between benign and harmful overfitting regimes for Vision Transformers trained on linearly separable data. The authors carefully analyze the optimization dynamics and provide generalization bounds that depend on the signal-to-noise ratio of the data. Extensive experiments validate the theory.

STRENGTHS:
* Provides a precise characterization of benign vs harmful overfitting regimes for Vision Transformers
* Writing is clear and easy to follow
* Novel results on the harmful overfitting regime
* Extensive experiments align well with and validate the theory

WEAKNESSES:
* Only considers the linearly separable setting, which is already solvable by existing vision and language models, so the conclusions are not very surprising even if initialization and model details differ
* Lacks clear takeaways for practitioners - how can these theoretical insights be used to improve real-world vision models?
* Could the authors provide a more rigorous perspective on transformer bias (e.g. low rank structure) from an optimization perspective under their model?

QUESTIONS:
* If an MLP with ReLU activation was used instead of the simple linearly separable task, would the optimization dynamics converge faster or slower? Would the SNR requirements be stricter or more relaxed?
* What guidance can the authors provide to practitioners based on these theoretical results?

LIMITATIONS:
The authors acknowledge focusing only on the simplified linearly separable setting. More discussion on how the insights could extend to real-world nonlinear problems would be valuable.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates **the benign overfitting phenomenon in Vision Transformers**. The authors adopt a theoretical framework similar to that proposed by Cao et al. (2022), which use **a data model consisting of label-dependent signal and label-independent noise**, but employ **a two-layer Transformer architecture** instead of a two-layer convolutional neural network. They provide conditions under which benign and harmful overfitting occur.\n\nSTRENGTHS:\n* The paper provides **conditions under which benign and harmful overfitting occur** in a two-layer transformer, with results that are **tight up to a constant**.\n* The authors overcome challenges in analyzing the highly complex training dynamics of transformers by introducing **a novel technique called vectorized Q&K and scalarized V**, successfully addressing the learning dynamics. To the best of my knowledge, this is **the first work to successfully address the learning dynamics of even a simple transformer architecture without unrealistic assumptions** (e.g., merging the key-query weights). If all proofs are correct, this represents a significant technical contribution.\n\nWEAKNESSES:\nWhile I believe that the vectorized Q&K and scalarized V techniques are significant contributions, **it is difficult to verify the correctness of the proofs** due to readability issues in both the main text and the appendix. I suggest that the authors **improve the clarity and writing of their technical terms and proofs**.\n\nMinors and Typos\n* The notion of $\\mu$ in Definition 3.1 seems unnecessary\n* Line 9: modal→model\n* Line 541: Eexperimantal Rresults → Experimental Results\n\nQUESTIONS:\n* In the data distribution (Definition 3.1), what is **the role of the larger noise $\\xi_2$ in the analysis**? What would happen if the data distribution were the same as that considered in Cao et al. (2022), which consists of a single signal patch and a single noise patch?\n* In the numerical results section, it would be beneficial to compare the results with those for a two-layer convolutional neural network, as considered in Cao et al. (2022), to emphasize the advantages of the transformer architecture.\n\nReference\n\n[1] Cao, Y., Chen,Z., Belkin,M., and Gu,Q. Benign over fitting in two-layer convolutional neural networks. In NeurIPS 2022\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

comments: SUMMARY:
This study investigates the theoretical aspects of Vision Transformers (ViT) with a focus on their generalization capabilities, particularly under conditions of benign overfitting. Through a detailed analysis of the optimization process involving a self-attention layer and a fully connected layer, optimized using gradient descent on a specific data distribution modal, this work addresses the complexities introduced by softmax functions and the interdependencies of multiple weight configurations in transformer models. By developing novel techniques, the researchers delineate the training dynamics that lead to effective generalization in post-training scenarios. A key contribution is the establishment of a sharp condition based on the signal-to-noise ratio within the data, which predicts whether a small or large test error will occur. These theoretical findings are supported by experimental simulations, enhancing our understanding of transformers' performance in vision tasks.

STRENGTHS:
This paper rigorously analyzes the training dynamics of a simplified ViT model. Specifically, its technical contributions related to "Vectorized Q & K and Scalarized V" and "Dealing with the Softmax Function" may have a broader impact on subsequent theoretical research.

WEAKNESSES:
The current experiments only validate the theoretical results on synthetic datasets. It is recommended that the authors consider adding some experiments on real datasets to test the effects, such as experiments on benign overfitting of the ViT model on MNIST and CIFAR10.

QUESTIONS:
See in Weaknesses.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

comments: SUMMARY:
The paper studies the benign overfitting phenomenon for a two-layer Transformer in vision. The paper characterizes the optimization of the ViT through three different phases in training dynamics and finds a sharp separation condition of the signal-to-noise ratio to distinguish the benign and harmful overfitting of the ViT.

STRENGTHS:
1. The paper gives a sharp transition between benign and harmful overfitting for ViT, which can be verified by a simulation experiment.
2. The paper proposes a novel method of vectorized $QK$ and scalarized $V$ to simplify the study of Transformer.
3. The paper successfully deals with the challenges caused by softmax and multiple weights.

WEAKNESSES:
1. Lack the introduction of the benign overfitting phenomenon in the first section.
2. It is better to clarify some notations like $\Omega(\cdot)$, $\Theta(\cdot)$, $\omega(\cdot)$...
3. Data generation is specified, thus this model may be a little bit limited.
4. A small typo: (Line 149) $\mu||_2^{-2}$ --> $||\mu||_2^{-2}$

QUESTIONS:
In Theorem 4.1 and Theorem 4.2, the requirement is related to SNR. Why does $\tilde{\sigma}_p$ not occur in the requirement?

LIMITATIONS:
The authors have stated the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12