Original: A T HEORETICAL U NDERSTANDING OF SHALLOW V ISION T RANSFORMERS : L EARNING , G ENERALIZATION , AND S AMPLE C OMPLEXITY
Revised: A T HEORETICAL U NDERSTANDING OF SHALLOW V ISION T RANSFORMERS : L EARNING , G ENERALIZATION , AND S AMPLE C OMPLEXITY

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Vision Transformers (ViTs) with self-attention modules have recently achieved great empirical success in many vision tasks.
Revised: Vision Transformers (ViTs) with self-attention modules have recently achieved great empirical success in many vision tasks.

Original: Due to non-convex interactions across layers, however, the theoretical learning and generalization analysis is mostly elusive.
Revised: Due to non-convex interactions across layers, however, the theoretical learning and generalization analysis is mostly elusive.

Original: Based on a data model characterizing both label-relevant and label-irrelevant tokens, this paper provides the first theoretical analysis of training a shallow ViT, i.e., one self-attention layer followed by a two-layer perceptron, for a classification task.
Revised: Based on a data model characterizing both label-relevant and label-irrelevant tokens, this paper provides the first theoretical analysis of training a shallow ViT, i.e., one self-attention layer followed by a two-layer perceptron, for a classification task.

Original: We characterize the sample complexity to achieve a zero generalization error.
Revised: We characterize the sample complexity to achieve a zero generalization error.

Original: Our sample complexity bound is positively correlated with the inverse of the fraction of label-relevant tokens, the token noise level, and the initial model error.
Revised: Our sample complexity bound is positively correlated with the inverse of the fraction of label-relevant tokens, the token noise level, and the initial model error.

Original: We also prove that a training process using stochastic gradient descent (SGD) leads to a sparse attention map, which is a formal verification of the general intuition about the success of attention.
Revised: We also prove that a training process using stochastic gradient descent (SGD) leads to a sparse attention map, which is a formal verification of the general intuition about the success of attention.

Original: Moreover, this paper indicates that a proper token sparsification can improve the test performance by removing label-irrelevant and/or noisy tokens, including spurious correlations.
Revised: Moreover, this paper indicates that a proper token sparsification can improve the test performance by removing label-irrelevant and/or noisy tokens, including spurious correlations.

Original: Empirical experiments on synthetic data and CIFAR-10 dataset justify our theoretical results and generalize to deeper ViTs.
Revised: Empirical experiments on synthetic data and CIFAR-10 dataset justify our theoretical results and generalize to deeper ViTs.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: As the backbone of Transformers (Vaswani et al., 2017), the self-attention mechanism (Bahdanau et al., 2014) computes the feature representation by globally modeling long-range interactions within the input.
Revised: As the backbone of Transformers (Vaswani et al., 2017) the self-attention mechanism (Bahdanau et al., 2014) computes the feature representation by globally modeling long-range interactions within the input.

Original: Transformers have demonstrated tremendous empirical success in numerous areas, including nature language processing (Kenton & Toutanova, 2019;
Revised: Transformers have demonstrated tremendous empirical success in numerous areas, including nature language processing (Kenton & Toutanova, 2019;

Original: Radford et al., 2019; 2018; Brown et al., 2020), recommendation system (Zhou et al., 2018; Chen et al., 2019; Sun et al., 2019), and reinforcement learning (Chen et al., 2021; Janner et al., 2021; Zheng et al., 2022).
Revised: Radford et al., 2019; 2018; Brown et al., 2020), recommendation system (Zhou et al., 2018; Chen et al., 2019; Sun et al., 2019) and reinforcement learning (Chen et al., 2021; Janner et al., 2021; Zheng et al., 2022).

Original: Starting from the advent of Vision Transformer (ViT) (Dosovitskiy et al., 2020), Transformer-based models (Touvron et al., 2021; Jiang et al., 2021; Wang et al., 2021; Liu et al., 2021a) gradually replace convolutional neural network (CNN) architectures and become prevalent in vision tasks.
Revised: Starting from the advent of Vision Transformer (ViT) (Dosovitskiy et al., 2020), Transformer-based models (Touvron et al., 2021; Jiang et al., 2021; Wang et al., 2021; Liu et al., 2021a) gradually replace convolutional neural network (CNN) architectures and become prevalent in vision tasks.

Original: Various techniques have been developed to train ViT efficiently.
Revised: Various techniques have been developed to train ViT efficiently.

Original: Among them, token sparsification (Pan et al., 2021; Rao et al., 

Revised: Among them, token sparsification (Pan et al., 2021; Rao et al.,

Original: ; Liang et al., 2022; Tang et al., 2022; Yin et al., 2022) removes redundant tokens (image patches) of data to improve the computational complexity while maintaining a comparable learning performance.
Revised: ; Liang et al., 2022; Tang et al., 2022; Yin et al., 2022) removes redundant tokens (image patches) of data to improve the computational complexity while maintaining a comparable learning performance.

Original: For example, Liang et al.
Revised: For example, Liang et al.

Original: ; Tang et al.
Revised: ; Tang et al.

Original: (2022) prune tokens following criteria designed based on the magnitude of the attention map.
Revised: (2022) prune tokens following criteria designed based on the magnitude of the attention map.

Original: Despite the remarkable empirical success, one fundamental question about training Transformers is still vastly open, which is
Revised: Despite the remarkable empirical success, one fundamental question about training Transformers is still vastly open, which is

Original: Under what conditions does a Transformer achieve satisfactory generalization?
Revised: Under what conditions does a Transformer achieve satisfactory generalization?

Original: Some recent works analyze Transformers theoretically from the perspective of proved Lipschitz constant of self-attention (James Vuckovic, 2020;
Revised: Some recent works analyze Transformers theoretically from the perspective of proved Lipschitz constant of self-attention (James Vuckovic, 2020;

Original: Kim et al., 2021), properties of the neural tangent kernel (Hron et al., 2020; Yang, 2020) and expressive power and Turing-completeness (Dehghani et al., 2018;
Revised: Kim et al., 2021), properties of the neural tangent kernel (Hron et al., 2020; Yang, 2020) and expressive power and Turing-completeness (Dehghani et al., 2018;

Original: Yun et al., 2019; Bhattamishra et al., 2020a;b; Edelman et al., 2022; Dong et al., 2021;
Revised: Yun et al., 2019; Bhattamishra et al., 2020a;b; Edelman et al., 2022; Dong et al., 2021;

Original: Likhosherstov et al., 2021;
Revised: Likhosherstov et al., 2021;

Original: Cordonnier et al., 2019; Levine et al., 2020) with statistical guarantees (Snell et al., 2021;
Revised: Cordonnier et al., 2019; Levine et al., 2020) with statistical guarantees (Snell et al., 2021;

Original: Wei et al., 2021).
Revised: Wei et al., 2021).

Original: Likhosherstov et al.
Revised: Likhosherstov et al.

Original: (2021) showed a model complexity for the function approximation of the self-attention module.
Revised: (2021) showed a model complexity for the function approximation of the self-attention module.

Original: Cordonnier et al.
Revised: Cordonnier et al.

Original: (2019) provided sufficient and necessary conditions for multi-head self-attention structures to simulate convolution layers.
Revised: (2019) provided sufficient and necessary conditions for multi-head self-attention structures to simulate convolution layers.

Original: None of these works, however, characterize the generalization performance of the learned model theoretically.
Revised: None of these works, however, characterize the generalization performance of the learned model theoretically.

Original: Only Edelman et al.
Revised: Only Edelman et al.

Original: (2022) theoretically proved that a single self-attention head can represent a sparse function of the input with a sample complexity for a generalization gap between the training loss and the test loss, but no discussion is provided regarding what algorithm to train the Transformer to achieve a desirable loss.
Revised: (2022) theoretically proved that a single self-attention head can represent a sparse function of the input with a sample complexity for a generalization gap between the training loss and the test loss, but no discussion is provided regarding what algorithm to train the Transformer to achieve a desirable loss.

Original: Contributions : To the best of our knowledge, this paper provides the first learning and generalization analysis of training a basic shallow Vision Transformer using stochastic gradient descent (SGD).
Revised: Contributions : To the best of our knowledge, this paper provides the first learning and generalization analysis of training a basic shallow Vision Transformer using stochastic gradient descent (SGD).

Original: This paper focuses on a binary classification problem on structured data, where tokens with discriminative patterns determine the label from a majority vote, while tokens with non-discriminative patterns do not affect the labels.
Revised: This paper focuses on a binary classification problem on structured data, where tokens with discriminative patterns determine the label from a majority vote, while tokens with non-discriminative patterns do not affect the labels.

Original: We train a ViT containing a self-attention layer followed by a two-layer perceptron using SGD from a proper initial model.
Revised: We train a ViT containing a self-attention layer followed by a two-layer perceptron using SGD from a proper initial model.

Original: This paper explicitly characterizes the required number of training samples to achieve a desirable generalization performance, referred to as the sample complexity.
Revised: This paper explicitly characterizes the required number of training samples to achieve a desirable generalization performance, referred to as the sample complexity.

Original: Our sample complexity bound is positively correlated with the inverse of the fraction of label-relevant tokens, the token noise level, and the error from the initial model, indicating a better generalization performance on data with fewer label-irrelevant patterns and less noise from a better initial model.
Revised: Our sample complexity bound is positively correlated with the inverse of the fraction of label-relevant tokens, the token noise level, and the error from the initial model, indicating a better generalization performance on data with fewer label-irrelevant patterns and less noise from a better initial model.

Original: The highlights of our technical contributions include:
Revised: The highlights of our technical contributions include:

Original: First, this paper proposes a new analytical framework to tackle the non-convex optimization and generalization for shallow ViTs.
Revised: First, this paper proposes a new analytical framework to tackle the non-convex optimization and generalization for shallow ViTs.

Original: Due to the more involved non-convex interactions of learning parameters and diverse activation functions across layers, the ViT model, i.e., a three-layer neural network with one self-attention layer, considered in this paper is more complicated to analyze than three-layer CNNs considered in Allen-Zhu et al. (2019a); Allen-Zhu & Li (2019), the most complicated neural network model that has been analyzed so far for across-layer nonconvex interactions.
Revised: Due to the more involved non-convex interactions of learning parameters and diverse activation functions across layers, the ViT model, i.e., a three-layer neural network with one self-attention layer, considered in this paper is more complicated to analyze than three-layer CNNs considered in Allen-Zhu et al. (2019a); Allen-Zhu & Li (2019), the most complicated neural network model that has been analyzed so far for across-layer nonconvex interactions.

Original: We consider a structured data model with relaxed assumptions from existing models and establish a new analytical framework to overcome the new technical challenges to handle ViTs.
Revised: We consider a structured data model with relaxed assumptions from existing models and establish a new analytical framework to overcome the new technical challenges to handle ViTs.

Original: Second, this paper theoretically depicts the evolution of the attention map during the training and characterizes how “attention” is paid to different tokens during the training.
Revised: Second, this paper theoretically depicts the evolution of the attention map during the training and characterizes how “attention” is paid to different tokens during the training.

Original: Specifically, we show that under the structured data model, the learning parameters of the self-attention module grow in the direction that projects the data to the label-relevant patterns, resulting in an increasingly sparse attention map.
Revised: Specifically, we show that under the structured data model, the learning parameters of the self-attention module grow in the direction that projects the data to the label-relevant patterns, resulting in an increasingly sparse attention map.

Original: This insight provides a theoretical justification of the magnitude-based token pruning methods such as (Liang et al., 2022; Tang et al., 2022) for efficient learning.
Revised: This insight provides a theoretical justification of the magnitude-based token pruning methods such as (Liang et al., 2022; Tang et al., 2022) for efficient learning.

Original: Third, we provide a theoretical explanation for the improved generalization using token sparsification.
Revised: Third, we provide a theoretical explanation for the improved generalization using token sparsification.

Original: We quantitatively show that if a token sparsification method can remove class-irrelevant and/or highly noisy tokens, then the sample complexity is reduced while achieving the same testing accuracy.
Revised: We quantitatively show that if a token sparsification method can remove class-irrelevant and/or highly noisy tokens, then the sample complexity is reduced while achieving the same testing accuracy.

Original: Moreover, token sparsification can also remove spurious correlations to improve the testing accuracy (Likhomanenko et al., 2021; Zhu et al., 2021a).
Revised: Moreover, token sparsification can also remove spurious correlations to improve the testing accuracy (Likhomanenko et al., 2021; Zhu et al., 2021a).

Original: This insight provides a guideline in designing token sparsification and few-shot learning methods for Transformer (He et al., 2022;
Revised: This insight provides a guideline in designing token sparsification and few-shot learning methods for Transformer (He et al., 2022;

Original: Guibas et al., 2022).
Revised: Guibas et al., 2022).

Original: 1. B ACKGROUND AND R ELATED W ORK
Revised: 1. B ACKGROUND AND R ELATED W ORK

Original: Efficient ViT learning.
Revised: Efficient ViT learning.

Original: To alleviate the memory and computation burden in training (Dosovitskiy et al., 2020; Touvron et al., 2021; Wang et al., 2022), various acceleration techniques have been developed other than token sparsification.
Revised: To alleviate the memory and computation burden in training (Dosovitskiy et al., 2020; Touvron et al., 2021; Wang et al., 2022), various acceleration techniques have been developed other than token sparsification.

Original: Zhu et al.
Revised: Zhu et al.

Original: (2021b) identifies the importance of different dimensions in each layer of ViTs and then executes model pruning.
Revised: (2021b) identifies the importance of different dimensions in each layer of ViTs and then executes model pruning.

Original: Liu et al.
Revised: Liu et al.

Original: (2021b)
Revised: (2021b)

Original: ; Lin et al.
Revised: ; Lin et al.

Original: ; Li et al.
Revised: ; Li et al.

Original: (2022d) quantize weights and inputs to compress the learning model.
Revised: (2022d) quantize weights and inputs to compress the learning model.

Original: Li et al.
Revised: Li et al.

Original: (2022a) studies automated progressive learning that automatically increases the model capacity onthe-fly.
Revised: (2022a) studies automated progressive learning that automatically increases the model capacity onthe-fly.

Original: Moreover, modifications of attention modules, such as the network architecture based on local attention (Wang et al., 2021; Liu et al., 2021a; Chu et al., 2021), can simplify the computation of global attention for acceleration.
Revised: Moreover, modifications of attention modules, such as the network architecture based on local attention (Wang et al., 2021; Liu et al., 2021a; Chu et al., 2021), can simplify the computation of global attention for acceleration.

Original: Theoretical analysis of learning and generalization of neural networks.
Revised: Theoretical analysis of learning and generalization of neural networks.

Original: One line of research (Zhong et al., 2017b; Fu et al., 2020; Zhong et al., 2017a;
Revised: One line of research (Zhong et al., 2017b; Fu et al., 2020; Zhong et al., 2017a;

Original: Zhang et al., 2020a;b; Li et al., 2022c) analyzes the generalization performance when the number of neurons is smaller than the number of training samples.
Revised: Zhang et al., 2020a;b; Li et al., 2022c) analyzes the generalization performance when the number of neurons is smaller than the number of training samples.

Original: The neural-tangent-kernel (NTK) analysis (Jacot et al., 2018; Allen-Zhu et al., 2019a;b; Arora et al., 2019; Cao & Gu, 2019; Zou & Gu, 2019; Du et al., 2019; Chen et al., 2020;
Revised: The neural-tangent-kernel (NTK) analysis (Jacot et al., 2018; Allen-Zhu et al., 2019a;b; Arora et al., 2019; Cao & Gu, 2019; Zou & Gu, 2019; Du et al., 2019; Chen et al., 2020;

Original: Li et al., 2022b) considers strongly overparameterized networks and eliminates the nonconvex inter- actions across layers by linearizing the neural network around the initialization.
Revised: Li et al., 2022b) considers strongly overparameterized networks and eliminates the nonconvex interactions across layers by linearizing the neural network around the initialization.

Original: The generalization performance is independent of the feature distribution and cannot explain the advantages of selfattention modules.
Revised: The generalization performance is independent of the feature distribution and cannot explain the advantages of selfattention modules.

Original: Neural network learning on structured data.
Revised: Neural network learning on structured data.

Original: Li & Liang (2018) provide the generalization analysis of a fully-connected neural network when the data comes from separated distributions.
Revised: Li & Liang (2018) provide the generalization analysis of a fully-connected neural network when the data comes from separated distributions.

Original: Daniely & Malach (2020); Shi et al.
Revised: Daniely & Malach (2020); Shi et al.

Original: ; Karp et al.
Revised: ; Karp et al.

Original: (2021); Brutzkus & Globerson (2021)  study fully connected networks and convolutional neural networks assuming that data contains discriminative patterns and background patterns.
Revised: (2021); Brutzkus & Globerson (2021) ; Zhang et al. (2023) study fully connected networks and convolutional neural networks assuming that data contains discriminative patterns and background patterns.

Original: Allen-Zhu & Li (2022) illustrates the robustness of adversarial training by introducing the feature purification mechanism, in which neural networks with non-linear activation functions can memorize the data-dependent features.
Revised: Allen-Zhu & Li (2022) illustrates the robustness of adversarial training by introducing the feature purification mechanism, in which neural networks with non-linear activation functions can memorize the data-dependent features.

Original: Wen & Li (2021) extends this framework to the area of self-supervised contrastive learning.
Revised: Wen & Li (2021) extends this framework to the area of self-supervised contrastive learning.

Original: All these works consider one-hidden-layer neural networks without self-attention.
Revised: All these works consider one-hidden-layer neural networks without self-attention.

Original: Notations : Vectors are in bold lowercase, and matrices and tensors are in bold uppercase.
Revised: Notations : Vectors are in bold lowercase, and matrices and tensors are in bold uppercase.

Original: Scalars are in normal fonts.
Revised: Scalars are in normal fonts.

Original: Sets are in calligraphy font.
Revised: Sets are in calligraphy font.

Original: For instance, Z is a matrix, and z is a vector.
Revised: For instance, Z is a matrix, and z is a vector.

Original: z i denotes the i -th entry of z , and Z i,j denotes the ( i, j ) -th entry of Z .
Revised: z i denotes the i -th entry of z , and Z i,j denotes the ( i, j ) -th entry of Z .

Original: [ K ] ( K > 0 ) denotes the set including integers from 1 to K .
Revised: [ K ] ( K > 0 ) denotes the set including integers from 1 to K .

Original: I d ∈ R d × d
Revised: I d ∈ R d × d

Original: and e i represent the identity matrix in R d × d
Revised: and e i represent the identity matrix in R d × d

Original: and the i -th standard basis vector, respectively.
Revised: and the i -th standard basis vector, respectively.

Original: We follow the convention that f ( x ) =
Revised: We follow the convention that f ( x ) =

Original: O ( g ( x ))
Revised: O ( g ( x ))

Original: (or Ω( g ( x )) ,
Revised: (or Ω( g ( x )) ,

Original: Θ( g ( x ))) means that f ( x ) increases at most, at least, or in the order of g ( x ) , respectively.
Revised: Θ( g ( x ))) means that f ( x ) increases at most, at least, or in the order of g ( x ) , respectively.

Original:  P ROBLEM F ORMULATION AND L EARNING A LGORITHM
Revised:  P ROBLEM F ORMULATION AND L EARNING A LGORITHM

Original: We study a binary classification problem 1 following the common setup in (Dosovitskiy et al., 2020;
Revised: We study a binary classification problem 1 following the common setup in (Dosovitskiy et al., 2020;

Original: Touvron et al., 2021; Jiang et al., 2021).
Revised: Touvron et al., 2021; Jiang et al., 2021).

Original: Given N training samples { ( X n , y n ) } Nn =1 generated from an unknown distribution D and a fair initial model, the goal is to find an improved model that maps X to y for any ( X , y ) ∼ D .
Revised: Given N training samples { ( X n , y n ) } Nn =1 generated from an unknown distribution D and a fair initial model, the goal is to find an improved model that maps X to y for any ( X , y ) ∼ D .

Original: Here each data point contains L tokens x n 1 , x n 2 , · · · , x nL , i.e.,
Revised: Here each data point contains L tokens x n 1 , x n 2 , · · · , x nL , i.e.,

Original: X n =
Revised: X n =

Original: [ x n 1 , · · · , x nL ] ∈ R d × L , where each token is d -dimensional and unit-norm.
Revised: [ x n 1 , · · · , x nL ] ∈ R d × L , where each token is d -dimensional and unit-norm.

Original: y n ∈ { +1 , − 1 } is a scalar.
Revised: y n ∈ { +1 , − 1 } is a scalar.

Original: A token can be an image patch (Dosovitskiy et al., 2020).
Revised: A token can be an image patch (Dosovitskiy et al., 2020).

Original: We consider a general setup that also applies to token sparsification, where some tokens are set to zero to reduce the computational time.
Revised: We consider a general setup that also applies to token sparsification, where some tokens are set to zero to reduce the computational time.

Original: Let S n ⊆
Revised: Let S n ⊆

Original: [ L ] denote the set of indices of remaining tokens in X n after sparsification.
Revised: [ L ] denote the set of indices of remaining tokens in X n after sparsification.

Original: Then |S n | ≤ L , and S n =
Revised: Then |S n | ≤ L , and S n =

Original: [ L ] without token sparsification.
Revised: [ L ] without token sparsification.

Original: Learning is performed over a basic three-layer Vision Transformer, a neural network with a singlehead self-attention layer and a two-layer fully connected network, as shown in (1).
Revised: Learning is performed over a basic three-layer Vision Transformer, a neural network with a singlehead self-attention layer and a two-layer fully connected network, as shown in (1).

Original: This is a simplified model of practical Vision Transformers (Dosovitskiy et al., 2020) to avoid unnecessary complications in analyzing the most critical component of ViTs, the self-attention.
Revised: This is a simplified model of practical Vision Transformers (Dosovitskiy et al., 2020) to avoid unnecessary complications in analyzing the most critical component of ViTs, the self-attention.

Original: [Equation]
Revised: [Equation]

Original: where the queue weights W Q in R m b × d , the key weights W K in R m b × d , and the value weights
Revised: where the queue weights W Q in R m b × d , the key weights W K in R m b × d , and the value weights

Original: W V in R m a × d in the attention unit are multiplied with X n to obtain the queue vector W Q X n , the key vector W K X n , and the value vector W V X n , respectively (Vaswani et al., 2017).
Revised: W V in R m a × d in the attention unit are multiplied with X n to obtain the queue vector W Q X n , the key vector W K X n , and the value vector W V X n , respectively (Vaswani et al., 2017).

Original: W O is in R m × m a and A =
Revised: W O is in R m × m a and A =

Original: ( a (1) , a (2) , · · · , a L ) where a ( l ) ∈ R m ,
Revised: ( a (1) , a (2) , · · · , a L ) where a ( l ) ∈ R m ,

Original: l ∈
Revised: l ∈

Original: [ L ] are the hidden-layer and output-layer weights of the two-layer perceptron, respectively.
Revised: [ L ] are the hidden-layer and output-layer weights of the two-layer perceptron, respectively.

Original: m is the number of neurons in the hidden layer.
Revised: m is the number of neurons in the hidden layer.

Original: Let ψ =
Revised: Let ψ =

Original: ( A , W O , W V , W K , W Q ) denote the set of parameters to train.
Revised: ( A , W O , W V , W K , W Q ) denote the set of parameters to train.

Original: The training problem minimizes the empirical risk f N ( ψ ) ,
Revised: The training problem minimizes the empirical risk f N ( ψ ) ,

Original: [Equation]
Revised: [Equation]

Original: where ℓ ( X n , y n ; ψ ) is the Hinge loss function, i [Equation] e.,
Revised: where ℓ ( X n , y n ; ψ ) is the Hinge loss function, i.e.,

Original: [Equation]
Revised: [Equation]

Original: The generalization performance of a learned model ψ is evaluated by the population risk f ( ψ ) , where
Revised: The generalization performance of a learned model ψ is evaluated by the population risk f ( ψ ) , where

Original: [Equation]
Revised: [Equation]

Original: The training problem (2) is solved via a mini-batch stochastic gradient descent (SGD), as summarized in Algorithm 1.
Revised: The training problem (2) is solved via a mini-batch stochastic gradient descent (SGD), as summarized in Algorithm 1.

Original: At iteration t , t = 0 , 1 , 2 , · · · , T − 1 , the gradient is computed using a mini-batch B t with |B t | =
Revised: At iteration t , t = 0 , 1 , 2 , · · · , T − 1 , the gradient is computed using a mini-batch B t with |B t | =

Original: B .
Revised: B .

Original: The step size is η .
Revised: The step size is η .

Original: Similar to (Dosovitskiy et al., 2020; Touvron et al., 2021;
Revised: Similar to (Dosovitskiy et al., 2020; Touvron et al., 2021;

Original: Jiang et al., 2021), W (0) V , W (0) Q , and W (0) K come from an initial model.
Revised: Jiang et al., 2021), W (0) V , W (0) Q , and W (0) K come from an initial model.

Original: Every entry of W O is generated from N (0 , ξ 2 ) .
Revised: Every entry of W O is generated from N (0 , ξ 2 ) .

Original: Every entry of a (0) l is sampled from { + 1 √ m , − 1 √ m } with equal probability.
Revised: Every entry of a (0) l is sampled from { + 1 √ m , − 1 √ m } with equal probability.

Original: A does not update during the training 2 .
Revised: A does not update during the training 2 .

Original:  T HEORETICAL R ESULTS 3. M AIN T HEORETICAL INSIGHTS
Revised:  T HEORETICAL R ESULTS 3. M AIN T HEORETICAL INSIGHTS

Original: Before formally introducing our data model and main theory, we first summarize the major insights.
Revised: Before formally introducing our data model and main theory, we first summarize the major insights.

Original: We consider a data model where tokens are noisy versions of label-relevant patterns that determine the data label and label-irrelevant patterns that do not affect the label.
Revised: We consider a data model where tokens are noisy versions of label-relevant patterns that determine the data label and label-irrelevant patterns that do not affect the label.

Original: α ∗ is the fraction of labelrelevant tokens.
Revised: α ∗ is the fraction of labelrelevant tokens.

Original: σ represents the initial model error, and τ characterizes the token noise level.
Revised: σ represents the initial model error, and τ characterizes the token noise level.

Original: (P1).
Revised: (P1).

Original: A Convergence and sample complexity analysis of SGD to achieve zero generalization error.
Revised: A Convergence and sample complexity analysis of SGD to achieve zero generalization error.

Original: We prove SGD with a proper initialization converges to a model with zero generalization error.
Revised: We prove SGD with a proper initialization converges to a model with zero generalization error.

Original: The required number of iterations is proportional to 1 /α ∗ and 1 / (Θ(1) − σ − τ ) .
Revised: The required number of iterations is proportional to 1 /α ∗ and 1 / (Θ(1) − σ − τ ) .

Original: Our sample complexity bound is linear in α − 2 ∗ and (Θ(1) − σ − τ ) − 2 .
Revised: Our sample complexity bound is linear in α − 2 ∗ and (Θ(1) − σ − τ ) − 2 .

Original: Therefore, the learning performance is improved, in the sense of a faster convergence and fewer training samples to achieve a desirable generalization, with a larger fraction of label-relevant patterns, a better initial model, and less token noise.
Revised: Therefore, the learning performance is improved, in the sense of a faster convergence and fewer training samples to achieve a desirable generalization, with a larger fraction of label-relevant patterns, a better initial model, and less token noise.

Original: (P2).
Revised: (P2).

Original: A theoretical characterization of increased sparsity of the self-attention module during training.
Revised: A theoretical characterization of increased sparsity of the self-attention module during training.

Original: We prove that the attention weights, which are softmax values of each token in the self-attention module, become increasingly sparse during the training, with non-zero weights concentrated at label-relevant tokens.
Revised: We prove that the attention weights, which are softmax values of each token in the self-attention module, become increasingly sparse during the training, with non-zero weights concentrated at label-relevant tokens.

Original: This formally justifies the general intuition that the attention layer makes the neural network focus on the most important part of data.
Revised: This formally justifies the general intuition that the attention layer makes the neural network focus on the most important part of data.

Original: (P3).
Revised: (P3).

Original: A theoretical guideline of designing token sparsification methods to reduce sample complexity.
Revised: A theoretical guideline of designing token sparsification methods to reduce sample complexity.

Original: Our sample complexity bound indicates that the required number of samples to achieve zero generalization can be reduced if a token sparsification method removes some label-irrelevant tokens (reducing α ∗ ), or tokens with large noise (reducing σ ), or both.
Revised: Our sample complexity bound indicates that the required number of samples to achieve zero generalization can be reduced if a token sparsification method removes some label-irrelevant tokens (reducing α ∗ ), or tokens with large noise (reducing σ ), or both.

Original: This insight provides a guideline to design proper token sparsification methods.
Revised: This insight provides a guideline to design proper token sparsification methods.

Original: (P4).
Revised: (P4).

Original: A new theoretical framework to analyze the nonconvex interactions in three-layer ViTs .
Revised: A new theoretical framework to analyze the nonconvex interactions in three-layer ViTs .

Original: This paper develops a new framework to analyze ViTs based on a more general data model than existing works like (Brutzkus & Globerson, 2021; Karp et al., 2021; Wen & Li, 2021).
Revised: This paper develops a new framework to analyze ViTs based on a more general data model than existing works like (Brutzkus & Globerson, 2021; Karp et al., 2021; Wen & Li, 2021).

Original: Compared with the nonconvex interactions in three-layer feedforward neural networks, analyzing ViTs has technical challenges that the softmax activation is highly non-linear, and the gradient computation on token correlations is complicated.
Revised: Compared with the nonconvex interactions in three-layer feedforward neural networks, analyzing ViTs has technical challenges that the softmax activation is highly non-linear, and the gradient computation on token correlations is complicated.

Original: We develop new tools to handle this problem by exploiting structures in the data and proving that SGD iterations increase the magnitude of label-relevant tokens only rather than label-irrelevant tokens.
Revised: We develop new tools to handle this problem by exploiting structures in the data and proving that SGD iterations increase the magnitude of label-relevant tokens only rather than label-irrelevant tokens.

Original: This theoretical framework is of independent interest and can potentially applied to analyze different variants of Transformers and attention mechanisms.
Revised: This theoretical framework is of independent interest and can potentially applied to analyze different variants of Transformers and attention mechanisms.

Original: 3. D ATA M ODEL
Revised: 3. D ATA M ODEL

Original: There are M (2 < M < m a , m b ) distinct patterns { µ 1 , µ 2 , · · · , µ M } in R d , where µ 1 , µare discriminative patterns that determine the binary labels, and the remaining M − 2 patterns µ 3 , µ 4 , · · · , µ M are non-discriminative patterns that do not affect the labels.
Revised: There are M (2 < M < m a , m b ) distinct patterns { µ 1 , µ 2 , · · · , µ M } in R d , where µ 1 , µare discriminative patterns that determine the binary labels, and the remaining M − 2 patterns µ 3 , µ 4 , · · · , µ M are non-discriminative patterns that do not affect the labels.

Revised: Let κ = fix the output layer weights as the random initialization in the theoretical analysis of neural networks, including NTK (Allen-Zhu et al., 2019a; Arora et al., 2019), model recovery (Zhong et al., 2017b), and feature learning (Karp et al., 2021; Allen-Zhu & Li, 2022) type of approaches.

Revised: The optimization problem here of W Q , W K , W V , and W O with non-linear activations is still highly non-convex and challenging.

Original: Let κ = min 1 ≤ i ̸ =
Revised: min 1 ≤ i ̸ =

Original: j ≤ M ∥ µ i − µ j ∥ > 0 denote the minimum distance between patterns.
Revised: j ≤ M ∥ µ i − µ j ∥ > 0 denote the minimum distance between patterns.

Original: Each token x nl of X n is a noisy version of one of the patterns, i.e.,
Revised: Each token x nl of X n is a noisy version of one of the patterns, i.e.,

Original: [Equation]
Revised: [Equation]

Original: output layer weights as the random initialization in the theoretical analysis of neural networks, including NTK (Allen-Zhu et al., 2019a; Arora et al., 2019), model recovery (Zhong et al., 2017b), and feature learning (Karp et al., 2021; Allen-Zhu & Li, 2022) type of approaches.

Original: The optimization problem here of W Q , W K , W V , and W O with non-linear activations is still highly non-convex and challenging.

Original: and the noise level τ ≤ κ .
Revised: and the noise level τ < κ/ 4 .

Revised: We take κ − 4 τ as Θ(1) for the simplicity of presentation.

Original: The label y n is determined by the tokens that correspond to discriminative patterns through a majority vote.
Revised: The label y n is determined by the tokens that correspond to discriminative patterns through a majority vote.

Original: If the number of tokens that are noisy versions of µ 1 is larger than the number of tokens that correspond to µ 2 in X n , then y n = 1 .
Revised: If the number of tokens that are noisy versions of µ 1 is larger than the number of tokens that correspond to µ 2 in X n , then y n = 1 .

Original: In this case that the label y n = 1 , the tokens that are noisy µ 1 are refereed to as label-relevant tokens, and the tokens that are noisy µ 2 are referred to as confusion tokens.
Revised: In this case that the label y n = 1 , the tokens that are noisy µ 1 are refereed to as label-relevant tokens, and the tokens that are noisy µ 2 are referred to as confusion tokens.

Original: Similarly, if there are more tokens that are noisy µ 2 than those that are noisy µ 1 , the former are label-relevant tokens, the latter are confusion tokens, and y n = − 1 .
Revised: Similarly, if there are more tokens that are noisy µ 2 than those that are noisy µ 1 , the former are label-relevant tokens, the latter are confusion tokens, and y n = − 1 .

Original: All other tokens that are not label-relevant are called label-irrelevant tokens.
Revised: All other tokens that are not label-relevant are called label-irrelevant tokens.

Original: Let α ∗ and α # as the average fraction of the label-relevant and the confusion tokens over the distribution D , respectively.
Revised: Let α ∗ and α # as the average fraction of the label-relevant and the confusion tokens over the distribution D , respectively.

Original: We consider a balanced dataset.
Revised: We consider a balanced dataset.

Original: Let D + = { ( X n , y n ) | y n = +1 , n ∈ [ N ] } and D − = { ( X n , y n ) | y n = − 1 , n ∈ [ N ] } denote the sets of positive and negative labels, respectively.
Revised: Let D + = { ( X n , y n ) | y n = +1 , n ∈ [ N ] } and D − = { ( X n , y n ) | y n = − 1 , n ∈ [ N ] } denote the sets of positive and negative labels, respectively.

Original: Then (cid:12)(cid:12)(cid:12) |D + | − |D − | (cid:12)(cid:12)(cid:12) =
Revised: Then (cid:12) (cid:12) (cid:12) |D + | − |D − | (cid:12) (cid:12) (cid:12) =

Original: O ( √ N ) .
Revised: O ( √ N ) .

Original: Our model is motivated by and generalized from those used in the state-of-art analysis of neural networks on structured data (Li & Liang, 2018; Brutzkus & Globerson, 2021; Karp et al., 2021).
Revised: Our model is motivated by and generalized from those used in the state-of-art analysis of neural networks on structured data (Li & Liang, 2018; Brutzkus & Globerson, 2021; Karp et al., 2021).

Original: All the existing models require that only one discriminative pattern exists in each sample, i.e., either µor µ 2 , but not both, while our model allows both patterns to appear in the same sample.
Revised: All the existing models require that only one discriminative pattern exists in each sample, i.e., either µor µ 2 , but not both, while our model allows both patterns to appear in the same sample.

Original: 3. F ORMAL T HEORETICAL R ESULTS
Revised: 3. F ORMAL T HEORETICAL R ESULTS

Original: Before presenting our main theory below, we first characterize the behavior of the initial model through Assumption 1.
Revised: Before presenting our main theory below, we first characterize the behavior of the initial model through Assumption 1.

Original: Some important notations are summarized in Table 1.notations
Revised: Some important notations are summarized in Table 1.

Original: [Equation]

Original: [Table]

Original: δ

Original: [Table]

Original: [Equation]

Original: [Table]

Original: τ

Original: [Table]

Original: M

Original: [Table]
Revised: [Table]

Original: m

Original: The

Original: [Table]

Original: in W

Original: [Table]

Original: [Equation]

Original: [Table]

Original: of

Original: [Table]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: Assumption 1.
Revised: Assumption 1.

Original: Assume max( ∥ W (0) V ∥ , ∥ W (0) K ∥ , ∥ W (0) Q ∥ ) ≤ 1 without loss of generality.
Revised: Assume max( ∥ W (0) V ∥ , ∥ W (0) K ∥ , ∥ W (0) Q ∥ ) ≤ 1 without loss of generality.

Original: There exist three (not necessarily different) sets of orthonormal bases P = { p 1 , p 2 , · · · , p M } , Q = { q 1 , q 2 , · · · , q M } , and R = { r 1 , r 2 , · · · , r M } , where p l ∈ R m a , q l , r l ∈ R m b , ∀ l ∈ [ M ] , q 1 = r 1 , and q 2 = r 2 3 such that
Revised: There exist three (not necessarily different) sets of orthonormal bases P = { p 1 , p 2 , · · · , p M } , Q = { q 1 , q 2 , · · · , q M } , and R = { r 1 , r 2 , · · · , r M } , where p l ∈ R m a , q l , r l ∈ R m b , ∀ l ∈ [ M ] , q 1 = r 1 , and q 2 = r 2 3 such that

Original: [Equation]
Revised: [Equation]

Original: hold for some σ = O (1 /M ) and δ < 1 / 2 .
Revised: hold for some σ = O (1 /M ) and δ < 1 / 2 .

Original: Assumption 1 characterizes the distance of query, key, and value vectors of patterns { µ j } M j =1 to orthonormal vectors.
Revised: Assumption 1 characterizes the distance of query, key, and value vectors of patterns { µ j } Mj =1 to orthonormal vectors.

Original: The requirement on δ is minor because δ can be in the same order as ∥ µ j ∥ .
Revised: The requirement on δ is minor because δ can be in the same order as ∥ µ j ∥ .

Original: Theorem 1 (Generalization of ViT) .
Revised: Theorem 1 (Generalization of ViT) .

Original: Suppose Assumption 1 holds; τ ≤ min( σ, δ ) ; a sufficiently large model with
Revised: Suppose Assumption 1 holds; τ ≤ min( σ, δ ) ; a sufficiently large model with

Original: [Equation]
Revised: [Equation]

Original: the average fraction of label-relevant patterns satisfies
Revised: the average fraction of label-relevant patterns satisfies

Original: [Equation]
Revised: [Equation]

Original: for some constant ϵ S ∈ (0 , 12 ) ; and the mini-batch size and the number of sampled tokens of each data X n , n ∈ [ N ] satisfy
Revised: for some constant ϵ S ∈ (0 , 12 ) ; and the mini-batch size and the number of sampled tokens of each data X n , n ∈ [ N ] satisfy

Original: [Equation]
Revised: [Equation]

Original: Then as long as the number of training samples N satisfies
Revised: Then as long as the number of training samples N satisfies

Original: [Equation]
Revised: [Equation]

Original: for some constant c ′ , c ′′ > 0 , and ζ ≳ 1 − η 10 , after T number of iterations such that
Revised: for some constant c ′ , c ′′ > 0 , and ζ ≳ 1 − η 10 , after T number of iterations such that

Original: [Equation]
Revised: [Equation]

Original: with a probability at least 0 .
Revised: with a probability at least 0 .

Original: 99 , the returned model achieves zero generalization error as
Revised: 99 , the returned model achieves zero generalization error as

Original: [Equation]
Revised: [Equation]

Original: Theorem 1 characterizes under what condition of the data the neural network with self-attention in (1) trained with Algorithm 1 can achieve zero generalization error.
Revised: Theorem 1 characterizes under what condition of the data the neural network with self-attention in (1) trained with Algorithm 1 can achieve zero generalization error.

Original: To show that the self-attention layer can improve the generalization performance by reducing the required sample complexity to achieve zero generalization error, we also quantify the sample complexity when there is no selfattention layer in the following proposition.
Revised: To show that the self-attention layer can improve the generalization performance by reducing the required sample complexity to achieve zero generalization error, we also quantify the sample complexity when there is no selfattention layer in the following proposition.

Original: Proposition 1 (Generalization without self-attention) .
Revised: Proposition 1 (Generalization without self-attention) .

Original: Suppose assumptions in Theorem 1 hold.
Revised: Suppose assumptions in Theorem 1 hold.

Original: When there is no self-attention layer, i.e., W K and W Q are not updated during the training, if N satisfies 
Revised: When there is no self-attention layer, i.e., W K and W Q are not updated during the training, if N satisfies [Equation]

Original: [Equation]
Revised: [Equation]

Original: then after T iterations with T in (11), the returned model achieves zero generalization error as
Revised: then after T iterations with T in (11), the returned model achieves zero generalization error as

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: Remark 1.
Revised: Remark 1.

Original: (Advantage of the self-attention layer)
Revised: (Advantage of the self-attention layer)

Original: Because m ≫ m a , m b , d , the number of trainable parameter remains almost the same with or without updating the attention layer.
Revised: Because m ≫ m a , m b , d , the number of trainable parameter remains almost the same with or without updating the attention layer.

Original: Combining
Revised: Combining

Original: Theorem 1 and Proposition 1, we can see that with the additional self-attention layer, the sample complexity 4 is reduced by a factor 1 /α 2 ∗ with an approximately equal number of network parameters.
Revised: Theorem 1 and Proposition 1, we can see that with the additional self-attention layer, the sample complexity 4 is reduced by a factor 1 /α 2 ∗ with an approximately equal number of network parameters.

Original: Remark 2.
Revised: Remark 2.

Original: (Generalization improvement by token sparsification).
Revised: (Generalization improvement by token sparsification).

Original: (10) and (11) show that the sample complexity N and the required number of iterations T scale with 1 /α 2 ∗ and 1 /α ∗ , respectively.
Revised: (10) and (11) show that the sample complexity N and the required number of iterations T scale with 1 /α 2 ∗ and 1 /α ∗ , respectively.

Original: Then, increasing α ∗ , the fraction of label-relevant tokens, can reduce the sample complexity and speed up the convergence.
Revised: Then, increasing α ∗ , the fraction of label-relevant tokens, can reduce the sample complexity and speed up the convergence.

Original: Similarly, N and T scale with 1 / (Θ(1) − τ ) 2 and 1 / (Θ(1) − τ ) .
Revised: Similarly, N and T scale with 1 / (Θ(1) − τ ) 2 and 1 / (Θ(1) − τ ) .

Original: Then decreasing τ , the noise in the tokens, can also improve the generalization.
Revised: Then decreasing τ , the noise in the tokens, can also improve the generalization.

Original: Note that a properly designed token sparsification method can both increase α ∗ by removing label-irrelevant tokens and decrease τ by removing noisy tokens, thus improving the generalization performance.
Revised: Note that a properly designed token sparsification method can both increase α ∗ by removing label-irrelevant tokens and decrease τ by removing noisy tokens, thus improving the generalization performance.

Original: Remark 3.
Revised: Remark 3.

Original: (Impact of the initial model)
Revised: (Impact of the initial model)

Original: The initial model W (0) V , W (0) K , W (0) Q affects the learning performance through σ and δ , both of which decrease as the initial model is improved.
Revised: The initial model W (0) V , W (0) K , W (0) Q affects the learning performance through σ and δ , both of which decrease as the initial model is improved.

Original: Then from (10) and (11), the sample complexity reduces and the convergence speeds up for a better initial model.
Revised: Then from (10) and (11), the sample complexity reduces and the convergence speeds up for a better initial model.

Original: Proposition 2 shows that the attention weights are increasingly concentrated on label-relevant tokens during the training.
Revised: Proposition 2 shows that the attention weights are increasingly concentrated on label-relevant tokens during the training.

Original: Proposition 2 is a critical component in proving Theorem 1 and is of independent interest.
Revised: Proposition 2 is a critical component in proving Theorem 1 and is of independent interest.

Original: Proposition 2.
Revised: Proposition 2.

Original: The attention weights for each token become increasingly concentrated on those correlated with tokens of the label-relevant pattern during the training, i.e., (cid:88)
Revised: The attention weights for each token become increasingly concentrated on those correlated with tokens of the label-relevant pattern during the training, i.e., (cid:88)

Original: [Equation]
Revised: [Equation]

Original: at a sublinear rate of O (1 /t ) as t increases for a large C > 0 and all l ∈ S n and n ∈ [ N ] .
Revised: at a sublinear rate of O (1 /t ) as t increases for a large C > 0 and all l ∈ S n and n ∈ [ N ] .

Revised: 4 The sample complexity bounds in (10) and (13) are sufficient but not necessary.

Revised: Thus, rigorously speaking, one can not compare two cases based on sufficient conditions only.

Revised: In our analysis, however, these two bounds are derived with exactly the same technique with the only difference in handling the self-attention layer.

Revised: Therefore, we believe it is fair to compare these two bounds to show the advantage of ViT.

Original: Proposition 2 indicates that only label-relevant tokens are highlighted by the learned attention of
Revised: Proposition 2 indicates that only label-relevant tokens are highlighted by the learned attention of

Original: ViTs, while other tokens have less weight.
Revised: ViTs, while other tokens have less weight.

Original: This provides a theoretical justification of magnitudebased token sparsification methods.
Revised: This provides a theoretical justification of magnitudebased token sparsification methods.

Original: layer.

Original: Therefore, we believe it is fair to compare these two bounds to show the advantage of ViT.

Original: Proof idea sketch : The main proof idea is to show that the SGD updates scale up value, query, and key vectors of discriminative patterns, while keeping the magnitude of the projections of nondiscriminative patterns and the initial model error almost unchanged.
Revised: Proof idea sketch : The main proof idea is to show that the SGD updates scale up value, query, and key vectors of discriminative patterns, while keeping the magnitude of the projections of nondiscriminative patterns and the initial model error almost unchanged.

Original: To be more specific, by Lemma 3, 4, we can identify two groups of neurons in the hidden layer W O , where one group only learns the positive pattern, and the other group only learns the negative pattern.
Revised: To be more specific, by Lemma 3, 4, we can identify two groups of neurons in the hidden layer W O , where one group only learns the positive pattern, and the other group only learns the negative pattern.

Original: Claim 1 of Lemmastates that during the SGD updates, the neuron weights in these two groups evolve in the direction of projected discriminative patterns, p 1 and p 2 , respectively.
Revised: Claim 1 of Lemmastates that during the SGD updates, the neuron weights in these two groups evolve in the direction of projected discriminative patterns, p 1 and p 2 , respectively.

Original: Meanwhile, Claim 2 of Lemmaindicates that W K and W Q update in the direction of increasing the magnitude of the query and key vectors of label-relevant tokens from 1 to Θ(log T ) , such that the attention weights correlated with label-relevant tokens gradually become dominant.
Revised: Meanwhile, Claim 2 of Lemmaindicates that W K and W Q update in the direction of increasing the magnitude of the query and key vectors of label-relevant tokens from 1 to Θ(log T ) , such that the attention weights correlated with label-relevant tokens gradually become dominant.

Original: Moreover, by Claim 3 of Lemma 2, the update of
Revised: Moreover, by Claim 3 of Lemma 2, the update of

Original: W V increases the magnitude of the value vectors of label-relevant tokens, by adding partial neuron weights of W O that are aligned with the value vectors to these vectors.
Revised: W V increases the magnitude of the value vectors of label-relevant tokens, by adding partial neuron weights of W O that are aligned with the value vectors to these vectors.

Original: Due to the above properties during the training, one can simplify the training process to show that the output of neural network (1) changes linearly in the iteration number t .
Revised: Due to the above properties during the training, one can simplify the training process to show that the output of neural network (1) changes linearly in the iteration number t .

Original: From the above analysis, we can develop the sample complexity and the required number of iterations for the zero generalization guarantee.
Revised: From the above analysis, we can develop the sample complexity and the required number of iterations for the zero generalization guarantee.

Original: Technical novelty :
Revised: Technical novelty :

Original: Our proof technique is inspired by the feature learning technique in analyzing fully connect networks and convolution neural networks (Shi et al., 2021; Brutzkus & Globerson, 2021).
Revised: Our proof technique is inspired by the feature learning technique in analyzing fully connect networks and convolution neural networks (Shi et al., 2021; Brutzkus & Globerson, 2021).

Original: Our paper makes new technical contributions from the following aspects.
Revised: Our paper makes new technical contributions from the following aspects.

Original: First, we provide a new framework of studying the nonconvex interactions of multiple weight matrices in a threelayer ViT while other feature learning works (Shi et al., 2021; Brutzkus & Globerson, 2021; Karp et al., 2021; Allen-Zhu & Li, 2022; Wen & Li, 2021) only study one trainable weight matrix in the hidden layer of a two-layer network.
Revised: First, we provide a new framework of studying the nonconvex interactions of multiple weight matrices in a three-layer ViT while other feature learning works (Shi et al., 2021; Brutzkus & Globerson, 2021; Karp et al., 2021; Allen-Zhu & Li, 2022; Wen & Li, 2021; Zhang et al., 2023) only study one trainable weight matrix in the hidden layer of a two-layer network.

Original: Second, we analyze the updates of the self-attention module with the softmax function during the training, while other papers either ignore this issue without exploring convergence analysis (Edelman et al., 2022) or oversimplify the analysis by applying the neural-tangent-kernel (NTK) method that considers impractical over-parameterization and updates the weights only around initialization.
Revised: Second, we analyze the updates of the selfattention module with the softmax function during the training, while other papers either ignore this issue without exploring convergence analysis (Edelman et al., 2022) or oversimplify the analysis by applying the neural-tangent-kernel (NTK) method that considers impractical over-parameterization and updates the weights only around initialization.

Original: (Hron et al., 2020; Yang, 2020; Allen-Zhu et al., 2019a; Arora et al., 2019).
Revised: (Hron et al., 2020; Yang, 2020; Allen-Zhu et al., 2019a; Arora et al., 2019).

Original: Third, we consider a more general data model, where discriminative patterns of multiple classes can exist in the same data sample, but the data models in (Brutzkus & Globerson, 2021; Karp et al., 2021) require one discriminative pattern only in each sample.
Revised: Third, we consider a more general data model, where discriminative patterns of multiple classes can exist in the same data sample, but the data models in (Brutzkus & Globerson, 2021; Karp et al., 2021) require one discriminative pattern only in each sample.

Original:  N UMERICAL EXPERIMENTS 4. E XPERIMENTS ON S YNTHETIC DATASETS
Revised:  N UMERICAL EXPERIMENTS 4. E XPERIMENTS ON S YNTHETIC DATASETS

Original: We first verify the theoretical bounds in Theorem 1 on synthetic data.
Revised: We first verify the theoretical bounds in Theorem 1 on synthetic data.

Original: We set the dimension of data and attention embeddings to be d = m a = m b = 10 .
Revised: We set the dimension of data and attention embeddings to be d = m a = m b = 10 .

Original: Let c 0 = 0 .
Revised: Let c 0 = 0 .

Original: Let the total number of patterns M = 5 , and { µ 1 , µ 2 , · · · , µ M } be a set of orthonormal bases.
Revised: Let the total number of patterns M = 5 , and { µ 1 , µ 2 , · · · , µ M } be a set of orthonormal bases.

Original: To satisfy Assumption 1, we generate every token that is a noisy version of µ i from a Gaussian distribution N ( µ i , c 20 · I ) with the mean µ i and covariance c 20 I .
Revised: To satisfy Assumption 1, we generate every token that is a noisy version of µ i from a Gaussian distribution N ( µ i , c 20 · I ) with the mean µ i and covariance c 20 I .

Original: W (0) Q = W (0) Q =
Revised: W (0) Q = W (0) Q =

Original: δ 

Revised: δ

Original: I /c 20 , W (0) V =
Revised: I /c 20 , W (0) V =

Original: σ 2 U /c 20 , and each entry of
Revised: σ 2 U /c 20 , and each entry of

Original: W (0) O follows N (0 , ξ 2 ) , where U is an m a × m a orthonormal matrix, and ξ = 0 .
Revised: W (0) O follows N (0 , ξ 2 ) , where U is an m a × m a orthonormal matrix, and ξ = 0 .

Original: The number of neurons m of W O is 1000 .
Revised: The number of neurons m of W O is 1000 .

Original: We set the ratio of different patterns the same among all the data for simplicity.
Revised: We set the ratio of different patterns the same among all the data for simplicity.

Original: Sample complexity and convergence rate : We first study the impact of the fraction of the labelrelevant patterns α ∗ on the sample complexity.
Revised: Sample complexity and convergence rate : We first study the impact of the fraction of the labelrelevant patterns α ∗ on the sample complexity.

Original: Let the number of tokens after sparsification be |S n | = 100 , the initialization error σ = 0 . 1 , and δ = 0 .
Revised: Let the number of tokens after sparsification be |S n | = 100 , the initialization error σ = 0 . 1 , and δ = 0 .

Original: The fraction of non-discriminative patterns is fixed to be 0 .
Revised: The fraction of non-discriminative patterns is fixed to be 0 .

Original: We implement 20 independent experiments with the same α ∗ and N and record the Hinge loss values of the testing data.
Revised: We implement 20 independent experiments with the same α ∗ and N and record the Hinge loss values of the testing data.

Original: An experiment is successful if the testing loss is smaller than 10 − 3 .
Revised: An experiment is successful if the testing loss is smaller than 10 − 3 .

Original: Figure 1 (a) shows the success rate of these experiments.
Revised: Figure 1 (a) shows the success rate of these experiments.

Original: A black block means that all the trials fail.
Revised: A black block means that all the trials fail.

Original: A white block means that they all succeed.
Revised: A white block means that they all succeed.

Original: The sample complexity is indeed almost linear in α − 2 ∗ , as predicted in 10.
Revised: The sample complexity is indeed almost linear in α − 2 ∗ , as predicted in 10.

Original: We next explore the impact on σ .
Revised: We next explore the impact on σ .

Original: Set α ∗ = 0 .
Revised: Set α ∗ = 0 .

Original: 3 and α # = 0 .
Revised: 3 and α # = 0 .

Original: The number of tokens after sparsification is fixed at 50 for all the data.
Revised: The number of tokens after sparsification is fixed at 50 for all the data.

Original: Figure 1 (b) shows that 1 / √ N is linear in Θ(1) − σ , matching our theoretical prediction in (10).
Revised: Figure 1 (b) shows that 1 / √ N is linear in Θ(1) − σ , matching our theoretical prediction in (10).

Original: The result on the noise level τ is similar to Figure 1 (b), and we skip it here.
Revised: The result on the noise level τ is similar to Figure 1 (b), and we skip it here.

Original: In Figure 2, we verify the number of iterations T against α − 1 ∗ in (11) where we set σ = 0 .
Revised: In Figure 2, we verify the number of iterations T against α − 1 ∗ in (11) where we set σ = 0 .

Original: 1 and δ = 0 .
Revised: 1 and δ = 0 .

Original: Advantage of self-attention : To verify Proposition 1, we compare the performance on ViT in 1 and on the same network with W K and W Q fixed during the training, i.e., a three-layer CNN.
Revised: Advantage of self-attention : To verify Proposition 1, we compare the performance on ViT in 1 and on the same network with W K and W Q fixed during the training, i.e., a three-layer CNN.

Original: Compared with ViT, the number of trainable parameters in CNN is reduced by only 1% .
Revised: Compared with ViT, the number of trainable parameters in CNN is reduced by only 1% .

Original: Figure 3 shows the sample complexity of CNN is almost linear in α − 4 ∗ as predicted in (13).
Revised: Figure 3 shows the sample complexity of CNN is almost linear in α − 4 ∗ as predicted in (13).

Original: Compared with Figure 2 (a), the sample complexity significantly increases for small α ∗ , indicating a much worse generalization of CNN.
Revised: Compared with Figure 2 (a), the sample complexity significantly increases for small α ∗ , indicating a much worse generalization of CNN.

Original: [Figure]
Revised: [Figure]

Original: Attention map : We then evaluate the evolution of the attention map during the training.
Revised: Attention map : We then evaluate the evolution of the attention map during the training.

Original: Let |S n | = 50 for all n ∈ [ N ] .
Revised: Let |S n | = 50 for all n ∈ [ N ] .

Original: The number of training samples is N = 200 .
Revised: The number of training samples is N = 200 .

Original: σ = 0 .
Revised: σ = 0 .

Original: 1 , δ = 0 .
Revised: 1 , δ = 0 .

Original: 2 , α ∗ = 0 .
Revised: 2 , α ∗ = 0 .

Original: 5 , α # = 0 .
Revised: 5 , α # = 0 .

Original: In Figure 4, the red line with asterisks shows that the sum of attention weights on label-relevant tokens, i.e., the left side of (15) averaged over all l , indeed increases to be close to 

Revised: In Figure 4, the red line with asterisks shows that the sum of attention weights on label-relevant tokens, i.e., the left side of (15) averaged over all l , indeed increases to be close to

Original: when the number of iterations increases
Revised: when the number of iterations increases

Original: Correspondingly, the sum of attention weights on other tokens decreases to be close to 0 , as shown in the blue line with squares.
Revised: Correspondingly, the sum of attention weights on other tokens decreases to be close to 0 , as shown in the blue line with squares.

Original: This verifies Lemma 2 on a sparse attention map.
Revised: This verifies Lemma 2 on a sparse attention map.

Original: Token sparsification : We verify the improvement by token sparsification in Figure 5.
Revised: Token sparsification : We verify the improvement by token sparsification in Figure 5.

Original: The experiment is duplicatedfor 20 times.
Revised: The experiment is duplicated 20 times.

Original: The number of training samples N = 80 .
Revised: The number of training samples N = 80 .

Original: Let |S n | = 50 for all n ∈ [ N ] .
Revised: Let |S n | = 50 for all n ∈ [ N ] .

Original: Set σ = 0 .
Revised: Set σ = 0 .

Original: 1 , δ = 0 .
Revised: 1 , δ = 0 .

Original: 5 , α ∗ = 0 .
Revised: 5 , α ∗ = 0 .

Original: 6 , α # = 0 .
Revised: 6 , α # = 0 .

Original: If we applythe random sampling over all tokens, the performance cannot be improved as shown in the red curve, because α ∗ and σ do not change.
Revised: If we apply random sampling over all tokens, the performance cannot be improved as shown in the red curve because α ∗ and σ do not change.

Original: If we remove either label-irrelevant tokens or tokens with significant noise, the testing loss decreases, as indicated in the blue and black curves.
Revised: If we remove either label-irrelevant tokens or tokens with significant noise, the testing loss decreases, as indicated in the blue and black curves.

Original: This justifies our insight P3 on token sparsification.
Revised: This justifies our insight P3 on token sparsification.

Original: 4. E XPERIMENTS ON I MAGE C LASSIFICATION D ATASETS
Revised: 4. E XPERIMENTS ON I MAGE C LASSIFICATION D ATASETS

Original: Dataset : To characterize the effect of label-relevant and label-irrelevant tokens on generalization, following the setup of image integration in (Karp et al., 2021), we adopt an image from CIFARdataset (Krizhevsky et al., 2010) as the label-relevant image pattern and integrate it with a noisy background image from the IMAGENET Plants synset (Karp et al., 2021; Deng et al., 2009), which plays the role of label-irrelevant feature.
Revised: Dataset : To characterize the effect of label-relevant and label-irrelevant tokens on generalization, following the setup of image integration in (Karp et al., 2021), we adopt an image from CIFARdataset (Krizhevsky et al., 2010) as the label-relevant image pattern and integrate it with a noisy background image from the IMAGENET Plants synset (Karp et al., 2021; Deng et al., 2009), which plays the role of label-irrelevant feature.

Original: Specifically, we randomly cut out a region with size 26 ×in the IMAGENET image and replace it with a resized CIFAR-10 image.
Revised: Specifically, we randomly cut out a region with size 26 ×in the IMAGENET image and replace it with a resized CIFAR-10 image.

Original: Architecture : Experiments are implemented on a deep ViT model.
Revised: Architecture : Experiments are implemented on a deep ViT model.

Original: Following (Dosovitskiy et al., 2020), the network architecture contains 5 blocks, where we have a 4 -head self-attention layer and a one-layer perceptron with skip connections and Layer-Normalization in each block.
Revised: Following (Dosovitskiy et al., 2020), the network architecture contains 5 blocks, where we have a 4 -head self-attention layer and a one-layer perceptron with skip connections and Layer-Normalization in each block.

Original: We first evaluate the impact on generalization of token sparsification that removes label-irrelevant patterns to increase α ∗ .
Revised: We first evaluate the impact on generalization of token sparsification that removes label-irrelevant patterns to increase α ∗ .

Original: We consider a ten-classification problem where in both the training and testing datasets, the images used for integration are randomly selected from CIFAR-10 and IMAGENET.
Revised: We consider a ten-classification problem where in both the training and testing datasets, the images used for integration are randomly selected from CIFAR-10 and IMAGENET.

Original: The number of samples for training and testing is 50 K and 10 K , respectively.
Revised: The number of samples for training and testing is 50 K and 10 K , respectively.

Original: A pretrained model from CIFAR-100 (Krizhevsky et al., 2010) is used as the initial model with the output layer randomly initialized.
Revised: A pretrained model from CIFAR-100 (Krizhevsky et al., 2010) is used as the initial model with the output layer randomly initialized.

Original: Without token sparsification, the fraction of class-relevant tokens is α ∗ ≈ 0 . 66 .
Revised: Without token sparsification, the fraction of class-relevant tokens is α ∗ ≈ 0 . 66 .

Original: α ∗ = 1 implies all background tokens are removed.
Revised: α ∗ = 1 implies all background tokens are removed.

Original: Figure 6 (a) indicates that a larger α ∗ by removing more label-irrelevant tokens leads to higher test accuracy.
Revised: Figure 6 (a) indicates that a larger α ∗ by removing more label-irrelevant tokens leads to higher test accuracy.

Original: Moreover, the test performance improves with more training samples.
Revised: Moreover, the test performance improves with more training samples.

Original: These are consistent with our sample complexity analysis in (10).
Revised: These are consistent with our sample complexity analysis in (10).

Original: Figure 6 (b) presents the required sample complexity to learn a model with desirable test accuracy.
Revised: Figure 6 (b) presents the required sample complexity to learn a model with desirable test accuracy.

Original: We run 10 independent experiments for each pair of α ∗ and N , and the experiment is considered a success if the learned model achieves a test accuracy of at least 77 .
Revised: We run 10 independent experiments for each pair of α ∗ and N , and the experiment is considered a success if the learned model achieves a test accuracy of at least 77 .

Original: We then evaluate the impact of token sparsification on removing spurious correlations (Sagawa et al., 2020), as well as the impact of the initial model.
Revised: We then evaluate the impact of token sparsification on removing spurious correlations (Sagawa et al., 2020), as well as the impact of the initial model.

Original: We consider a binary classification problem that differentiates “bird” and “airplane” images.
Revised: We consider a binary classification problem that differentiates “bird” and “airplane” images.

Original: To introduce spurious correlations in the training data, 90% of bird images in the training data are integrated into the IMAGENET plant background, while only 10% of airplane images have the plant background.
Revised: To introduce spurious correlations in the training data, 90% of bird images in the training data are integrated into the IMAGENET plant background, while only 10% of airplane images have the plant background.

Original: The remaining training data are integrated into a clean background by zero padding.
Revised: The remaining training data are integrated into a clean background by zero padding.

Original: Therefore, the label “bird” is spuriously correlated with the class-irrelevant plant background.
Revised: Therefore, the label “bird” is spuriously correlated with the class-irrelevant plant background.

Original: The testing data contain 50% birds and 50% airplanes, and each class has 50% plant background and 50% clean background.
Revised: The testing data contain 50% birds and 50% airplanes, and each class has 50% plant background and 50% clean background.

Original: The numbers of training and testing samples are 10 K and 2 K , respectively.
Revised: The numbers of training and testing samples are 10 K and 2 K , respectively.

Original: We initialize the ViT using two pre-trained models.
Revised: We initialize the ViT using two pre-trained models.

Original: The first one is pre-trained with CIFAR-100, which contains images of 100 classes not including birds and airplanes.
Revised: The first one is pre-trained with CIFAR-100, which contains images of 100 classes not including birds and airplanes.

Original: The other initial model is trained with a modified CIFAR-10 with 500 images per class for a total of eight classes, excluding birds and airplanes.
Revised: The other initial model is trained with a modified CIFAR-10 with 500 images per class for a total of eight classes, excluding birds and airplanes.

Original: The pre-trained model on CIFAR-100 is a better initial model because it is trained on a more diverse dataset with more samples.
Revised: The pre-trained model on CIFAR-100 is a better initial model because it is trained on a more diverse dataset with more samples.

Original: In Figure 6 (c), the token sparsification method removes the tokens of the added background, and the corresponding α ∗ increases.
Revised: In Figure 6 (c), the token sparsification method removes the tokens of the added background, and the corresponding α ∗ increases.

Original: Note that removing background in the training dataset also reduces the spurious correlations between birds and plants.
Revised: Note that removing background in the training dataset also reduces the spurious correlations between birds and plants.

Original: Figure 6 (c) shows that from both initial models, the testing accuracy increases when more background tokens are removed.
Revised: Figure 6 (c) shows that from both initial models, the testing accuracy increases when more background tokens are removed.

Original: Moreover, a better initial model leads to a better testing performance.
Revised: Moreover, a better initial model leads to a better testing performance.

Original: This is consistent with Remarks 2 and 3.
Revised: This is consistent with Remarks 2 and 3.

Original: [Figure]
Revised: [Figure]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: This paper provides a novel theoretical generalization analysis of three-layer ViTs.
Revised: This paper provides a novel theoretical generalization analysis of three-layer ViTs.

Original: Focusing on a data model with label-relevant and label-irrelevant tokens, this paper explicitly quantifies the sample complexity as a function of the fraction of label-relevant tokens and the token noise projected by the initial model.
Revised: Focusing on a data model with label-relevant and label-irrelevant tokens, this paper explicitly quantifies the sample complexity as a function of the fraction of label-relevant tokens and the token noise projected by the initial model.

Original: It proves that the learned attention map becomes increasingly sparse during the training, where the attention weights are concentrated on those of label-relevant tokens.
Revised: It proves that the learned attention map becomes increasingly sparse during the training, where the attention weights are concentrated on those of label-relevant tokens.

Original: Our theoretical results also offer a guideline on designing proper token sparsification methods to improve the test performance.
Revised: Our theoretical results also offer a guideline on designing proper token sparsification methods to improve the test performance.

Original: This paper considers a simplified but representative Transformer architecture to theoretically examine the role of self-attention layer as the first step.
Revised: This paper considers a simplified but representative Transformer architecture to theoretically examine the role of self-attention layer as the first step.

Original: One future direction is to analyze more practical architectures such as those with skip connection, local attention layers, and Transformers in other areas.
Revised: One future direction is to analyze more practical architectures such as those with skip connection, local attention layers, and Transformers in other areas.

Original: We see no ethical or immediate negative societal consequence of our work.
Revised: We see no ethical or immediate negative societal consequence of our work.

Original: A M ORE EXPERIMENTS
Revised: A CKNOWLEDGMENTS

Revised: This work was supported by AFOSR FA9550-20-1-0122, ARO W911NF-21-1-0255, NSFand the Rensselaer-IBM AI Research Collaboration (http://airc.rpi.edu), part of the IBM AI Horizons Network (http://ibm.biz/AIHorizons).

Revised: We thank Dr. Shuai Zhang at Rensselaer Polytechnic

Revised: Institute for useful discussions.

Revised: We thank all anonymous reviewers for their constructive comments

Revised: A C OMPARISON WITH TWO RELATED WORKS A. C OMPARISON WITH (A LLEN -Z HU & L I , 2023)

Revised: Allen-Zhu & Li (2023) studies ensemble learning and knowledge distillation.

Revised: Its main proof idea is that given large amounts of multi-view data, each single model learns one feature, and then ensemble learning integrates all learned features and, thus, improves over single models.

Revised: Knowledge distillation applies softmax logits to make use of information learned from the ensemble model.

Revised: It is analyzed in a similar approach to studying single models.

Revised: The single models considered in (Allen-Zhu & Li, 2023) is a two-layer Relu network.

Revised: In contrast, in this paper, we consider a two-layer Relu network with an additional self-attention layer.

Revised: The network architecture and training algorithm for the self-attention layer is completely different from those for the softmax logit in the knowledge distillation function.

Revised: In our proof, we analyze the impact of the gradient of W Q , W K , and W V on different patterns (Claims 2 andof Lemma 2), showing that the training process helps to enlarge the magnitude of label-relevant features.

Revised: We also show that neurons in W O mainly learn from discriminative patterns (Claim 1 of

Revised: Lemma 2).

Revised: Such a learning process is affected by the error in the initial model and the noise in tokens.

Revised: Please see details in “Proof idea sketch” and “Technical novelty” in Section 3.3 on Pageof the paper.

Revised: This technique we develop plays a critical role in our analysis of self-attention layers.

Revised: This technique is novel and did not appear in any existing works.

Revised: A. C OMPARISON WITH (J ELASSI ET AL ., 2022)

Revised: Jelassi et al.

Revised: (2022) is a concurrent work which theoretically studies Vision Transformers.

Revised: The major difference between (Jelassi et al., 2022) and our work is that we consider different data models and network architectures.

Revised: In (Jelassi et al., 2022), the data model requires spatial association between tokens.

Revised: The attention map is replaced with position encoding, and the training process of the attention map is simplified to train a linear layer.

Revised: Our setup models the data mainly based on the category of patterns.

Revised: We keep the classical structure and training process of self-attention, where W Q , W K , and W V are trained separately.

Revised: The required number of samples and iterations are derived as functions of the fraction of label-relevant patterns.

Revised: In addition, the non-linear activation function they consider is polynomial activation, instead of Relu or Gelu as in practice.

Revised: Based on these conditions, they are able to study a different and mroe general labelling function.

Revised: B M ORE EXPERIMENTS

Original: Following a similar setup in Section 4, we add more experiments.
Revised: Following a similar setup in Section 4, we add more experiments.

Original: For experiments on synthetic data, we set the dimension of data and attention embeddings to be d = m a = m b = 20 .
Revised: For experiments on synthetic data, we set the dimension of data and attention embeddings to be d = m a = m b = 20 .

Original: We vary the number of patterns M to be 10 , 15 , and 20 .
Revised: We vary the number of patterns M to be 10 , 15 , and 20 .

Original: Data generation and the network architecture follow the setup in Section 4.
Revised: Data generation and the network architecture follow the setup in Section 4.

Original: One can observe the same trend in
Revised: One can observe the same trend in

Original: Figure 7 and 8 as in Figure 4 and 5, respectively, indicating that our conclusion that the attention map becomes sparse during the training, and that pruning label-irrelevant tokens or noisy tokens improves the performance, both hold for different choices of M .
Revised: Figure 7 and 8 as in Figure 4 and 5, respectively, indicating that our conclusion that the attention map becomes sparse during the training, and that pruning label-irrelevant tokens or noisy tokens improves the performance, both hold for different choices of M .

Original: B P RELIMINARIES
Revised: C P RELIMINARIES

Original: We first formally restate the neural network with different notations of loss functions, and the Algorithm 1 of the training steps after token sparsification.
Revised: We first formally restate the neural network with different notations of loss functions, and the Algorithm 1 of the training steps after token sparsification.

Original: The notations used in the Appendix is summarized in Table 2.
Revised: The notations used in the Appendix is summarized in Table 2.

Original: For the network

Original: [Equation]

Original: The loss function of a single data, a mini-batch, the empirical loss, and the population loss is defined in the following.

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: The network output for X n and the loss function of a single data.
Revised: The network output for X n and the loss function of a single data.

Original: [Table]
Revised: [Table]

Original: The loss function of a mini-batch, the empirical loss, and the population loss, respectively.
Revised: The loss function of a mini-batch, the empirical loss, and the population loss, respectively.

Original: [Table]
Revised: [Table]

Original: The features in value, key, and query vectors at the iteration t for pattern j , respectively.
Revised: The features in value, key, and query vectors at the iteration t for pattern j , respectively.

Original: We have p j (0) =
Revised: We have p j (0) =

Original: p j , q j (0) =
Revised: p j , q j (0) =

Original: q j , and r j (0) =
Revised: q j , and r j (0) =

Original: r j .
Revised: r j .

Original: [Table]
Revised: [Table]

Original: nj ( t

Original: [Table]

Original: n nj ( t ) , o nj ( t )

Original: The error terms in the value, key, and query vectors of the j -th token and n -th data compared to their features at iteration t .
Revised: The error terms in the value, key, and query vectors of the j -th token and n -th data compared to their features at iteration t .

Original: [Table]
Revised: [Table]

Original: The set of lucky neurons at t -th iterations.
Revised: The set of lucky neurons at t -th iterations.

Original: [Table]
Revised: [Table]

Original: Approximate value of some attention weights at iteration t .
Revised: The bounds of value of some attention weights at iteration t .

Original: λ is the threshold between inner products of tokens from the same pattern and different patterns.
Revised: λ is the threshold between inner products of tokens from the same pattern and different patterns.

Original: [Table]
Revised: [Table]

Original: S nj is the set of sampled tokens of pattern j for the n -th data.
Revised: S nj is the set of sampled tokens of pattern j for the n -th data.

Original: S n ∗ , S n # are sets of sampled tokens of the label-relevant pattern and the confusion pattern for the n -th data, respectively.
Revised: S n ∗ , S n # are sets of sampled tokens of the label-relevant pattern and the confusion pattern for the n -th data, respectively.

Original: [Table]
Revised: [Table]

Original: The mean of fraction of label-relevant tokens, confusion tokens, and non-discriminative tokens, respectively.
Revised: The mean of fraction of label-relevant tokens, confusion tokens, and non-discriminative tokens, respectively.

Original: Loss
Revised: For the network

Original: [Equation]
Revised: [Equation]

Original: N (cid:88) n
Revised: The loss function of a single data, a mini-batch, the empirical loss, and the population loss is defined in the following.

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The formal algorithm is as follows.
Revised: The formal algorithm is as follows.

Original: We assume that each entry of W (0) O is randomly initialized from
Revised: We assume that each entry of W (0) O is randomly initialized from

Original: N (0 , ξ 2 ) where ξ = 1 √ M .
Revised: N (0 , ξ 2 ) where ξ = 1 √ M .

Original: a (0)( l ) i , i ∈
Revised: a (0)( l ) i , i ∈

Original: [ m ] , l ∈ [ L ] is uniformly initialized from + { 1 √ m , − 1 √ m } and fixed during the training.
Revised: [ m ] , l ∈ [ L ] is uniformly initialized from + { 1 √ m , − 1 √ m } and fixed during the training.

Original: W V , W K , and W Q are initialized from a good pretrained model.
Revised: W V , W K , and W Q are initialized from a good pretrained model.

Original: [Equation]
Revised: [Equation]

Original:  Assumption 1 can be interpreted as that we initialize W V , W K , and W Q to be the matrices that can map tokens to orthogonal features with added error terms.
Revised: [Equation] [Equation] [Equation] Assumption 1 can be interpreted as that we initialize W V , W K , and W Q to be the matrices that can map tokens to orthogonal features with added error terms.

Original: Assumption 2.
Revised: Assumption 2.

Original: Define P =
Revised: Define P =

Original: ( p 1 , p 2 , · · · , p M ) ∈ 
Revised: ( p 1 , p 2 , · · · , p M ) ∈ R m

Revised: [Equation]

Original: R m a × M , Q = ( q 1 , q 2 , · · · , q M ) ∈ R m b × M and R = ( r 1 , r 2 , · · · , r M ) ∈
Revised: × M , Q = ( q 1 , q 2 , · · · , q M ) ∈ R m b × M and R = ( r 1 , r 2 , · · · , r M ) ∈

Original: R m b × M as three feature matrices, where P = { p 1 , p 2 , · · · , p M } ,
Revised: R m b × M as three feature matrices, where P = { p 1 , p 2 , · · · , p M } ,

Original: Q = { q 1 , q 2 , · · · , q M } and R = { r 1 , r 2 , · · · , r M } are three sets of orthonormal bases.
Revised: Q = { q 1 , q 2 , · · · , q M } and R = { r 1 , r 2 , · · · , r M } are three sets of orthonormal bases.

Original: Define the noise terms z nj ( t ) , n nj ( t ) and o nj ( t ) with ∥ z nj (0) ∥ ≤ σ + τ and ∥ n nj (0) ∥ , ∥ o nj (0) ∥ ≤ δ + τ for j ∈ [ L ] .
Revised: Define the noise terms z nj ( t ) , n nj ( t ) and o nj ( t ) with ∥ z nj (0) ∥ ≤ σ + τ and ∥ n nj (0) ∥ , ∥ o nj (0) ∥ ≤ δ + τ for j ∈ [ L ] .

Original: q 1 = r 1 , q 2 = r 2 .
Revised: q 1 = r 1 , q 2 = r 2 .

Original: Suppose ∥ W (0) V ∥ , ∥ W (0) K ∥ , ∥ W (0) Q ∥ ≤ 1 , σ, τ < O (1 /M ) and δ < 1 / 2 .
Revised: Suppose ∥ W (0) V ∥ , ∥ W (0) K ∥ , ∥ W (0) Q ∥ ≤ 1 , σ, τ < O (1 /M ) and δ < 1 / 2 .

Original: Then, for x nl ∈ S nj
Revised: Then, for x nl ∈ S nj

Original: W (0) V x nl = p
Revised: W (0) V x nl = p

Original: j
Revised: j

Original: + z nj (0) .
Revised: + z nj (0) .

Original: W (0) K x nl = q j + n nj (0) .
Revised: W (0) K x nl = q j + n nj (0) .

Original: W (0) Q x nl = r j + o n j (0) .
Revised: W (0) Q x nl = r j + o nj (0) .

Original: Assumption 2 is a straightforward combination of Assumption 1 and (5) by applying the triangle inequality to bound the error terms for tokens.
Revised: Assumption 2 is a straightforward combination of Assumption 1 and (5) by applying the triangle inequality to bound the error terms for tokens.

Original: Definition 1.
Revised: Definition 1.

Revised: 1. ϕ

Original: [Equation]
Revised: [Equation]

Revised: n 2 , if y n =S n 1 , if y n = −5.

Revised: α ∗ =

Revised: E (cid:104) |S n ∗ | |S n | (cid:105) , α # = E (cid:104) |S

Revised: n # | |S n | (cid:105) , α nd = (cid:80) Ml =3 E (cid:104)

Revised: |S nl | |S n | (cid:105) .

Original: Definition 2.
Revised: Definition 2.

Original: Let θ i 1 be the angle between p 1 and W O
Revised: Let θ i 1 be the angle between p 1 and W O

Original: ( i, · ) .
Revised: ( i, · ) .

Original: Let θ i 2 be the angle between p 2 and
Revised: Let θ i 2 be the angle between p 2 and

Original: W O ( i, · ) .
Revised: W O ( i, · ) .

Original: Define W ( t ) , U ( t ) as the sets of lucky neurons at the t -th iteration such that
Revised: Define W ( t ) , U ( t ) as the sets of lucky neurons at the t -th iteration such that

Original: [Equation]
Revised: [Equation]

Original: Assumption 3.
Revised: Assumption 3.

Original: For one data X n , if the patch i and j correspond to the same feature k ∈ [ M ] , i.e., i ∈ S nk and j ∈ S nk , we have
Revised: For one data X n , if the patch i and j correspond to the same feature k ∈ [ M ] , i.e., i ∈ S nk and j ∈ S nk , we have

Original: [Equation]
Revised: [Equation]

Original: If the patch i and j correspond to the different feature k, l ∈ [ M ] , k ̸ = l i.e., i ∈ S nk and j ∈ S nl , k ̸
Revised: If the patch i and j correspond to the different feature k, l ∈ [ M ] , k ̸ = l i.e., i ∈ S nk and j ∈ S nl , k ̸

Original: = l , we have
Revised: = l , we have

Original: [Equation]
Revised: [Equation]

Original: This assumption is equivalent to the data model by (5) since τ < O (1 /M ) .
Revised: This assumption is equivalent to the data model by (5) since τ < O (1 /M ) .

Original: For the simplicity of presentation, we scale up all tokens a little bit to make the threshold of linear separability be 1 .
Revised: For the simplicity of presentation, we scale up all tokens a little bit to make the threshold of linear separability be 1 .

Revised: We also take 1 − λ and λ as Θ(1) for the simplicity.

Original: Definition 3.
Revised: Definition 3.

Original: (Vershynin, 2010)
Revised: (Vershynin, 2010)

Original: We say X is a sub-Gaussian random variable with sub-Gaussian norm K > 0 , if ( E | X | p ) 1 p ≤ K √ p for all p ≥ 1 .
Revised: We say X is a sub-Gaussian random variable with sub-Gaussian norm K > 0 , if ( E | X | p ) 1 p ≤ K √ p for all p ≥ 1 .

Original: In addition, the sub-Gaussian norm of X, denoted ∥ X ∥ ψ 2 , is defined as ∥ X ∥ ψ 2 = sup p ≥ 1 p − 12 ( E | X | p ) 1 p .
Revised: In addition, the sub-Gaussian norm of X, denoted ∥ X ∥ ψ 2 , is defined as ∥ X ∥ ψ 2 = sup p ≥ 1 p − 12 ( E | X | p ) 1 p .

Original: Lemma 1.
Revised: Lemma 1.

Original: (Vershynin (2010)
Revised: (Vershynin (2010)

Original: Proposition 5.1, Hoeffding’s inequality) Let X 1 , X 2 , · · · , X N be independent centered sub-gaussian random variables, and let K = max i
Revised: Proposition 5.1, Hoeffding’s inequality) Let X 1 , X 2 , · · · , X N be independent centered sub-gaussian random variables, and let K = max i

Original: ∥
Revised: ∥

Original: X i ∥
Revised: X i ∥

Original: ψ 2 .
Revised: ψ 2 .

Original: Then for every a = ( a 1 , · · · , a N ) ∈ R N and every t ≥ 0 , we have 
Revised: Then for every a = ( a 1 , · · · , a N ) ∈ R N and every t ≥ 0 , we have [Equation]

Original: [Equation]
Revised: [Equation]

Original: where c > 0 is an absolute constant.
Revised: where c > 0 is an absolute constant.

Original: C P ROOF OF THE M AIN T HEOREM AND P ROPOSITIONS
Revised: D P ROOF OF THE M AIN T HEOREM AND P ROPOSITIONS

Original: We state Lemma 2 first before we introduce the proof of main theorems.
Revised: We state Lemma 2 first before we introduce the proof of main theorems.

Original: Lemma 2 is the key lemma in our paper to show the training process of our ViT model using SGD.
Revised: Lemma 2 is the key lemma in our paper to show the training process of our ViT model using SGD.

Original: It has three major claims.
Revised: It has three major claims.

Original: Claim 1 involves the growth of W ( t ) O in terms of different directions of p l , i ∈
Revised: Claim 1 involves the growth of W ( t ) O in terms of different directions of p l , i ∈

Original: [ M ] .
Revised: [ M ] .

Original: Claim 2 describes the training dynamics of W ( t ) Q and W ( t ) K separately to show the tendency to a sparse attention map.
Revised: Claim 2 describes the training dynamics of W ( t ) Q and W ( t ) K separately to show the tendency to a sparse attention map.

Original: Claim 3 studies the gradient update process of W ( t ) V .
Revised: Claim 3 studies the gradient update process of W ( t ) V .

Original: Lemma 2.
Revised: Lemma 2.

Original: For l ∈ S n 1 for the data with y n = 1 , define
Revised: For l ∈ S n 1 for the data with y n = 1 , define

Original: [Equation]
Revised: [Equation]

Original: with
Revised: with

Original: [Equation]
Revised: [Equation]

Original: We also have the following claims:
Revised: We also have the following claims:

Original: Claim 1.
Revised: Claim 1.

Original: For the lucky neuron i ∈ W ( t ) and b ∈
Revised: For the lucky neuron i ∈ W ( t ) and b ∈

Original: [ T ] , we have
Revised: [ T ] , we have

Original: [Equation]

Original: p ∈ P

Original: [Equation]

Original: p

Original: [Equation]
Revised: [Equation]

Original: and for the noise z l ( t ) ,
Revised: and for the noise z l ( t ) ,

Original: [Equation]
Revised: [Equation]

Original: [Equation] For i ∈ U ( t ) , we also have equations as in (32) to (35), including[Equation] p
Revised: For i ∈ U ( t ) , we also have equations as in (32) to (35), including

Original: [Equation]
Revised: [Equation]

Original: P

Original: [Equation]
Revised: [Equation]

Original: and for the noise z l ( t ) ,
Revised: and for the noise z l ( t ) ,

Original: [Equation]
Revised: [Equation]

Original: For unlucky neurons, we have
Revised: For unlucky neurons, we have

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: P

Original: [Equation][Equation] [Equation]
Revised: [Equation]

Original: Claim 2.
Revised: Claim 2.

Original: Given conditions in (8), there exists K ( t ) , Q ( t ) > 0 , t = 0 , 1 , · · · , T − 1 such that for j ∈ S n ∗ and l / ∈ S ∗ ,
Revised: Given conditions in (8), there exists K ( t ) , Q ( t ) > 0 , t = 0 , 1 , · · · , T − 1 such that for j ∈ S n ∗ and l / ∈ S ∗ ,

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: and for j / ∈ S n ∗ , l / ∈ S n ∗ , we have 
Revised: and for j / ∈ S n ∗ , l / ∈ S n ∗ , we have softmax ( x nj ⊤ W ( t

Original: [Equation]
Revised: [Equation]

Original: For i = 1 , 2 ,
Revised: For i = 1 , 2 ,

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Claim 3.
Revised: Claim 3.

Original: For the update of W ( t ) V , there exists λ ≤ Θ(1) such that
Revised: For the update of W ( t ) V , there exists λ ≤ Θ(1) such that

Original: [Equation]
Revised: [Equation]

Original: To prove Theorem 1, we either show F ( X n ) > 1 for y n = 1 or show F ( X n ) < − 1 for y n = − 1 .
Revised: To prove Theorem 1, we either show F ( X n ) > 1 for y n = 1 or show F ( X n ) < − 1 for y n = − 1 .

Original: Take y n = 1 as an example, the basic idea of the proof is to make use of Lemma 2 to find a lower bound as a function of α ∗ , σ , τ , etc..
Revised: Take y n = 1 as an example, the basic idea of the proof is to make use of Lemma 2 to find a lower bound as a function of α ∗ , σ , τ , etc..

Original: The remaining step is to derive conditions on the sample complexity and the required number of iterations in terms of α ∗ , σ , and τ such that the lower bound is greater than 1 .
Revised: The remaining step is to derive conditions on the sample complexity and the required number of iterations in terms of α ∗ , σ , and τ such that the lower bound is greater than 1 .

Original: Given a balanced dataset, these conditions also ensure that F ( X n ) < − 1 for y n = − 1 .
Revised: Given a balanced dataset, these conditions also ensure that F ( X n ) < − 1 for y n = − 1 .

Original: During the proof, we may need to use some of equations as intermediate steps in the proof of Lemma 2.
Revised: During the proof, we may need to use some of equations as intermediate steps in the proof of Lemma 2.

Original: Since that these equations are not concise for presentation, we prefer not to state them formally in Lemma 2, but still refer to them as useful conclusions.
Revised: Since that these equations are not concise for presentation, we prefer not to state them formally in Lemma 2, but still refer to them as useful conclusions.

Original: The following is the details of the proof.
Revised: The following is the details of the proof.

Original: Proof of Theorem 1:
Revised: Proof of Theorem 1:

Original: For y n = 1 , define K l + = { i ∈ [ m ] : W ( t )
Revised: For y n = 1 , define K l + = { i ∈ [ m ] : W ( t )

Original: O
Revised: O

Original: ( i, · ) V nl ( t ) ≥ 0 } and K l − =
Revised: ( i, · ) V nl ( t ) ≥ 0 } and K l − =

Original: { i ∈ [ m ] : W ( t ) O
Revised: { i ∈ [ m ] : W ( t ) O

Original: ( i, · ) V nl ( t ) < 0 } .
Revised: ( i, · ) V nl ( t ) < 0 } .

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: By Lemma 2, we have
Revised: By Lemma 2, we have

Original: [Equation]

Original: |

Original: [Equation]

Original: [Equation]

Original: ≳

Original: [Equation]

Original: [Equation]

Original: ≳

Original: [Equation]
Revised: [Equation]

Original: |

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: Bt

Original: [Equation]
Revised: [Equation]

Original: n ∈B b |

Original: [Equation]

Original: n

Original: [Equation]

Original: | p n

Original: [Equation]

Original: b )

Original: [Equation]

Original: |S n | aM

Original: [Equation]

Original: [Equation]

Original: [Equation] τ ) [Equation] where the second step comes from (27) and the last step is by (135).
Revised: where the second step comes from (27) and the last step is by (136).

Original: By the definition of K l + [Equation] we have
Revised: By the definition of K l + , we have

Original: [Equation]
Revised: [Equation]

Original: Combining (135) and (137)
Revised: Combining (136) and (138)

Original: [Equation]
Revised: [Equation]

Original: we can obtain
Revised: we can obtain

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Revised: for some c 1 , c 2 ∈ (0 , 1) .

Original: Note that at the T -th iteration,
Revised: Note that at the T -th iteration,

Original: [Equation]
Revised: [Equation]

Original: Since that
Revised: ≳

Original: [Equation]
Revised: [Equation]

Revised: Since that

Original: [Equation]
Revised: [Equation]

Original: To find the order-wise lower bound of q 1 ( T )
Revised: To find the order-wise lower bound of q 1 ( T )

Original: [Equation]
Revised: [Equation]

Original: we need to check the equation
Revised: we need to check the equation

Original: [Equation]
Revised: [Equation]

Original: One can obtain
Revised: One can obtain

Original: [Equation]
Revised: [Equation]

Original: Therefore,
Revised: Therefore,

Original: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: for some large C > 0 .
Revised: for some large C > 0 .

Original: We require that
Revised: We require that

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where the first step is by letting a = √ m and m ≳
Revised: where the first step is by letting a = √ m and m ≳

Original: M 2 .
Revised: M 2 .

Original: We replace p n ( b ) with p n ( T ) because when b achieves the level of T , b o 1 p n ( b ) o 2 is close to b o 1 for o 1 , o 2 ≥ 0 by (61).
Revised: We replace p n ( b ) with p n ( T ) because when b achieves the level of T , b o 1 p n ( b ) o 2 is close to b o 1 for o 1 , o 2 ≥ 0 by (61).

Original: Thus,
Revised: Thus,

Original: T (cid:88)
Revised: T (cid:88)

Original: [Equation]
Revised: [Equation]

Original: We also require
Revised: We also require

Original: [Equation]
Revised: [Equation]

Original: for some ϵ 0 >0 .
Revised: for some ϵ 0 >

Revised: [Equation]

Original: We know that
Revised: We know that

Revised: [Equation]

Revised: N

Original: [Equation]
Revised: [Equation]

Original:  for c ′ > 0 and c ′′ > 0 .
Revised: 66) for c ′ > 0 and c ′′ > 0 .

Original: We can then have
Revised: We can then have

Original: [Equation]
Revised: [Equation]

Original: whereα ≥
Revised: where

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: by (160), as long as
Revised: by (157), as long as

Original: N ≥

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: and
Revised: and

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where ζ ≥ 1 − η 10 .
Revised: where ζ ≥ 1 − η 10 .

Original: If there is no mechanism like the self-attention to compute the weight using the correlations between tokens, we have
Revised: If there is no mechanism like the self-attention to compute the weight using the correlations between tokens, we have

Original: [Equation]
Revised: [Equation]

Original: which can scale up the sample complexity in (69) by α − 2 ∗ .
Revised: which can scale up the sample complexity in (69) by α − 2 ∗ .

Original: Therefore, we can obtain
Revised: Therefore, we can obtain

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Similarly, we can derive that for y = − 1 ,
Revised: Similarly, we can derive that for y = − 1 ,

Original: [Equation]
Revised: [Equation]

Original: Hence, for all n ∈ [ N ] ,
Revised: Hence, for all n ∈ [ N ] ,

Revised: Loss ( X n , y n

Original: [Equation]
Revised: [Equation]

Revised: =

Revised: [Equation]

Revised: [Equation]

Original: We also have
Revised: We also have

Original: [Equation]
Revised: [Equation]

Original: with the conditions of sample complexity and the number of iterations.
Revised: with the conditions of sample complexity and the number of iterations.

Original: Proof of Proposition 1:
Revised: Proposition 1:

Original: The main proof is the same as the proof of Theorem 1.
Revised: The main proof is the same as the proof of Theorem 1.

Original: The only difference is that we need to modify (66) as follows
Revised: The only difference is that we need to modify (66) as follows

Original: [Equation]
Revised: [Equation]

Original: where the first step is because p n ( T ) does not update since W ( t ) K and W ( t ) Q are fixed at initialization
Revised: where the first step is because p n ( T ) does not update since W ( t ) K and W ( t ) Q are fixed at initialization

Original: W (0) K and W (0) Q , and the second step is by p n (0 =
Revised: W (0) K and W (0) Q , and the second step is by p n (0) =

Original: Θ( α ∗ ) .
Revised: Θ( α ∗ ) .

Original: Since that
Revised: Since that

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Revised: N

Original: [Equation]
Revised: [Equation]

Original: Proof of Proposition 2:
Revised: Proof of Proposition 2:

Original: It can be easily derived from Claim 2 of Lemma 2, (60), and (61).
Revised: It can be easily derived from Claim 2 of Lemma 2, (60), and (61).

Original: D P ROOF OF L EMMA
Revised: E P ROOF OF L EMMA

Original: We prove the whole lemma by a long induction, which is the reason why we prefer to wrap three claims into one lemma.
Revised: We prove the whole lemma by a long induction, which is the reason why we prefer to wrap three claims into one lemma.

Original: To make it easier to follow, however, we break this Section into three parts to introduce the proof of three claims of Lemma 2 separately.
Revised: To make it easier to follow, however, we break this Section into three parts to introduce the proof of three claims of Lemma 2 separately.

Original: D. P ROOF OF C LAIM 1 OF L EMMA
Revised: E. P ROOF OF C LAIM 1 OF L EMMA

Original: Although it looks cumbersome, the key idea of Claim 1 is to characterize the growth of W ( t ) O in terms of p l , l ∈
Revised: Although it looks cumbersome, the key idea of Claim 1 is to characterize the growth of W ( t ) O in terms of p l , l ∈

Original: [ M ] .
Revised: [ M ] .

Original: We compare W ( t +1)
Revised: We compare W ( t +1)

Original: O ( i, · ) p l and W ( t )
Revised: O ( i, · ) p l and W ( t )

Original: O
Revised: O

Original: ( i, · ) p l to see the direction of growth by computing the gradient.
Revised: ( i, · ) p l to see the direction of growth by computing the gradient.

Original: One can eventually find that lucky neurons grow the most in directions of p 1 and p 2 , i.e., the feature of label-relevant patterns, while unlucky neurons do not change much in magnitude.
Revised: One can eventually find that lucky neurons grow the most in directions of p 1 and p 2 , i.e., the feature of label-relevant patterns, while unlucky neurons do not change much in magnitude.

Original: We start our poof.
Revised: We start our poof.

Original: At the t -th iteration, if l ∈ S n 1 , let
Revised: At the t -th iteration, if l ∈ S n 1 , let

Original: [Equation]
Revised: [Equation]

Original: , l ∈ [ M ] , where the second step comes from (27).
Revised: , l ∈ [ M ] , where the second step comes from (27).

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: Hence, by (18), ∂ Loss b
Revised: Hence, by (18), [Equation] Loss b

Original: [Equation]
Revised: [Equation]

Original: Define that for j ∈ [ M ] ,
Revised: Define that for j ∈ [ M ] ,

Original: [Equation]
Revised: [Equation]

Original: ) and we can then obtain 
Revised: (83) and we can then obtain (cid:28) [Equation] (cid:88) n

Revised: [Equation]

Revised: b

Original: [Equation]
Revised: [Equation]

Revised: n|S n |

Revised: [Equation]

Revised: l ∈S n a ( l ) i 1 [ W ( t ) O

Revised: ( i, · ) V n l ( t )

Revised: [Equation]

Revised: 0] (cid:88) s ∈S l softmax (

Revised: [Equation]

Revised: ( t )

Revised: [Equation]

Revised: t )

Revised: [Equation]

Revised: ) p

Revised: [Equation]

Revised: j

Revised: [Equation]

Revised: b

Revised: [Equation]

Revised: |

Revised: [Equation]

Revised: |

Revised: [Equation]

Revised: l

Revised: [Equation]

Revised: n

Revised: [Equation]

Revised: [Equation]

Revised: i

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: We then show the statements in different cases.
Revised: We then show the statements in different cases.

Original: (1) When j = 1 , since that Pr( y n = 1) =
Revised: (1) When j = 1 , since that Pr( y n = 1) =

Original: Pr( y n = − 1) = 1 / 2 , by Hoeffding’s inequality in (26), we can derive
Revised: Pr( y n = − 1) = 1 / 2 , by Hoeffding’s inequality in (26), we can derive

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Hence, with probability at least 1 − ( mB ) − c , we have|
Revised: Hence, with probability at least 1 − ( mB ) − c , we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: For i ∈ W ( t ) , from the derivation in (132) later, we have
Revised: For i ∈ W ( t ) , from the derivation in (133) later [Equation] we have

Original: [Equation]
Revised: [Equation]

Original: Denote p n ( t ) =
Revised: Denote p n ( t ) =

Original: |S n 1 | ν n ( t ) e ∥
Revised: |S n 1 | ν n ( t ) e ∥

Original: q 1 ( t ) ∥ 2 − 2 δ ∥ q 1 ( t ) ∥ .
Revised: q 1 ( t ) ∥ 2 − 2 δ ∥ q 1 ( t ) ∥ .

Original: Hence,
Revised: Hence,

Original: I
Revised: [Equation] Hence, combining (90), (92), (93), (94), and (95), we can obtain (cid:28) W ( t +1) O ( i, · )

Original: [Equation]
Revised: [Equation]

Revised: p

Revised: [Equation]

Revised: (cid:29)

Revised: [Equation]

Revised: (cid:28)

Revised: W ( t )

Revised: [Equation]

Revised: ( i,

Revised: [Equation]

Revised: [Equation]

Original: n
Revised: p

Original: [Equation]
Revised: [Equation]

Revised: (cid:29) ≳

Original: [Equation]
Revised: η a ·

Original: [Equation]
Revised: [Equation]

Original: Hence, combining (90), (92), (93), (94)
Revised: (cid:88) n [Equation]

Original: [Equation]
Revised: [Equation]

Original: and (95), we can obtain
Revised: Ba

Original: [Equation]
Revised: [Equation]

Original: ≳
Revised: ≳

Original: [Equation] aB 
Revised: η aB (cid:88) n ∈B b

Original: [Equation]
Revised: [Equation]

Original: Since that W (0) O ( i, · ) ∼ N (0 , ξ 

Revised: Since that W (0) O ( i, · ) ∼ N (0 , ξ

Original: I m a ) [Equation] by the standard property of Gaussian distribution
Revised: I m a ) , by the standard property of Gaussian distribution, we have

Original: [Equation]
Revised: [Equation]

Original: we have

Original: [Equation]

Original: Therefore

Original: [Equation]

Original:  with high probability for all i ∈ [ m ] [Equation] we have
Revised: Therefore, with high probability for all i ∈ [ m ] , we have

Original: [Equation]
Revised: [Equation]

Original: Therefore [Equation] we can derive
Revised: Therefore, we can derive

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: by verifying that
Revised: by verifying that

Original: [Equation]
Revised: [Equation]

Original: When ηt < a m , we have
Revised: When ηt < am [Equation] we have

Original: [Equation]
Revised: [Equation]

Original: When ηt ≥ am , we have that
Revised: When ηt ≥ am , we have that

Original: [Equation]
Revised: [Equation]

Original: since that g ( t ) is monotonically increasing.
Revised: since that g ( t ) is monotonically increasing.

Original: Hence, (99) is verified.
Revised: Hence, (99) is verified.

Original: Since that
Revised: Since that

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: to simplify the further analysis, we will use the bound
Revised: to simplify the further analysis, we will use the bound

Original: [Equation]
Revised: [Equation]

Original: Note that this bound does not order-wise affect the final result of the required number of iterations.
Revised: Note that this bound does not order-wise affect the final result of the required number of iterations.

Original: (2) When p j ∈ P /p
Revised: (2) When p j ∈ P /p

Original: + , we have
Revised: + , we have

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: with probability at least 1 − ( mB ) − c .
Revised: with probability at least 1 − ( mB ) − c .

Original: (107) comes from (34).
Revised: (107) comes from (34).

Original: Then, combining (90), (105), (106), (107) and (108), we can obtain
Revised: Then, combining (90), (105), (106), (107) and (108), we can obtain

Original: [Equation]
Revised: [Equation]

Original: Furthermore, we have
Revised: Furthermore, we have

Original: [Equation]
Revised: [Equation]

Original: where the last step is by
Revised: where the last step is by

Original: [Equation]
Revised: [Equation]

Original:  to ensure a non-zero gradient.
Revised: ) to ensure a non-zero gradient.

Original: (3) If i ∈ U ( t ) , following the derivation of (104) and (110), we can conclude that
Revised: (3) If i ∈ U ( t ) , following the derivation of (104) and (110), we can conclude that

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: (4) If i / ∈ ( W ( t ) ∪ U ( t )) ,
Revised: (4) If i / ∈ ( W ( t ) ∪ U ( t )) ,

Original: [Equation]
Revised: [Equation]

Original: Following (107) and (108), we have
Revised: Following (107) and (108), we have

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Hence, combining (114), (115), and (116), we can obtain 
Revised: Hence, combining (114), (115), and (116), we can obtain (cid:12)(cid:12)(cid:12)

Original: [Equation]
Revised: [Equation]

Original: 117) and
Revised: 117) and

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where the last step is by
Revised: where the last step is by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: (5) We finally study the bound of W ( t )
Revised: (5) We finally study the bound of W ( t )

Original: O ( i, · ) and the product with the noise term according to the analysis above.
Revised: O ( i, · ) and the product with the noise term according to the analysis above.

Original: By (52), for the lucky neuron i , since that the update of W ( t ) O
Revised: By (52), for the lucky neuron i , since that the update of W ( t ) O

Original: ( i, · ) lies in the subspace spanned by P and p 1 , p 2 , · · · , p M all have a unit norm, we can derive
Revised: ( i, · ) lies in the subspace spanned by P and p 1 , p 2 , · · · , p M all have a unit norm, we can derive

Original: [Equation]
Revised: [Equation]

Original: For the unlucky neuron i

Original: [Equation]
Revised: [Equation]

Original:  we can similarly obtain
Revised: For the unlucky neuron i , we can similarly obtain

Original: [Equation]
Revised: [Equation]

Original: We can also verify that this claim holds when t = 1 .
Revised: We can also verify that this claim holds when t = 1 .

Original: The proof of Claim 1 finishes here.
Revised: The proof of Claim 1 finishes here.

Original: D. P ROOF OF C LAIM 2 OF L EMMA
Revised: E. P ROOF OF C LAIM 2 OF L EMMA

Original: The proof of Claim 2 is one of the most challenging parts in our paper, since that we need to deal with the complicated softmax function.
Revised: The proof of Claim 2 is one of the most challenging parts in our paper, since that we need to deal with the complicated softmax function.

Original: The core idea of proof is that we pay more attention on the changes of label-relevant features in the gradient update, which should be the most crucial factor based on our data model.
Revised: The core idea of proof is that we pay more attention on the changes of label-relevant features in the gradient update, which should be the most crucial factor based on our data model.

Original: We then show the attention map converges to be sparse as long as the data model satisfies (8).
Revised: We then show the attention map converges to be sparse as long as the data model satisfies (8).

Original: We first study the gradient of W ( t +1) Q in part (a) and the gradient of W ( t +1) K in part (b).
Revised: We first study the gradient of W ( t +1) Q in part (a) and the gradient of W ( t +1) K in part (b).

Original: (a) By (227), we have
Revised: (a) By (231), we have

Original: [Equation]
Revised: [Equation]

Original: For r, l ∈ S n 1 , by (43) we have
Revised: For r, l ∈ S n 1 , by (43) we have

Original: [Equation]
Revised: [Equation]

Original: +

Original: [Equation]

Original: |S [Equation] | − |S n 1 | ) (126) For r / ∈ S n 1 and l ∈ S n 1 , we havesoftmax
Revised: For r / ∈ S n 1 and l ∈ S n 1 , we have

Original: [Equation]
Revised: [Equation]

Original: Therefore, for s, r, l ∈ S n 1 , letW
Revised: Therefore, for s, r, l ∈ S n 1 , let

Original: [Equation]
Revised: [Equation]

Original: (128) where
Revised: where

Original: [Equation]
Revised: [Equation]

Revised: where the last inequality holds when the final iteration log T ≤ Θ(1) .

Original: [Equation]
Revised: [Equation]

Original: (130) for some Q e ( t ) > 0 and γ ′ l > 0 .
Revised: for some Q e ( t ) > 0 and γ ′ l > 0 .

Original: Here
Revised: Here

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: Therefore, for i ∈ W ( t ) , [Equation] [Table] [Equation]
Revised: for l ≥ 2 .

Original: [Table] [Equation] [Table] [Equation]
Revised: Note that | ζ ′ l | = 0 if |S n | =

Original: [Table] [Equation] [Table] [Equation]
Revised: |S n 1 | , l ≥ 2 .

Original: [Table] [Equation] [Table] [Equation]
Revised: Therefore, for i ∈ W ( t ) ,

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: >

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where the first step is by (27) and the second step is a combination of (32) to (35).
Revised: where the first step is by (27) and the second step is a combination of (32) to (35).

Original: The final step holds as long as
Revised: The final step holds as long as

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 133) and
Revised: 134) and

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Then we study how large the coefficient of q 1 ( t ) in (125).
Revised: Then we study how large the coefficient of q 1 ( t ) in (125).

Original: If s ∈ S n 1 , by basic computation given (32) to (35),
Revised: If s ∈ S n 1 , by basic computation given (32) to (35),

Original: [Equation]
Revised: [Equation]

Original: [Equation] where the last step is by (133) and (134).
Revised: where the last step is by (134) and (135) [Equation]

Original: If s ∈ S n 2 , from (36) to (39)
Revised: If s ∈ S n 2 , from (36) to (39)

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: If i ∈ W ( t ) and s / ∈ ( S n 1 ∪ S n 2 ) ,
Revised: If i ∈ W ( t ) and s / ∈ ( S n 1 ∪ S n 2 ) ,

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: i
Revised: t (cid:88)

Original: [Equation]
Revised: [Equation]

Original: 137) by (40) to (42). Hence, for i ∈ W ( t ) , j ∈ S g 1 , combining (129) and (135), we have
Revised: = (cid:88)

Original: [Equation]
Revised: [Equation]

Original: For i ∈ U ( t ) and l ∈ S n 1 , j ∈ S g [Equation]
Revised: ( t ) (138) by (40) to (42).

Original: For i/ ∈( W ( t )∪ U ( t )) and l ∈ S n 1 , j ∈ S g 1 , 
Revised: Hence, for i ∈ W ( t ) , j ∈ S g 1 , combining (129) and (136), we have

Original: [Equation] To study the case when l / ∈ S n 1 for all n ∈ [ N ] , we need to check all other l ’s. Recall that we focus on the coefficient of q 1 ( t ) in this part. Based on the computation in (136) and (137), we know that the contribution of coefficient from non-discriminative patches is no more than that from discriminative patches, i.e., for l / ∈ ( S n 1 ∪ S n 2 ) , n ∈ [ N ] and k ∈ S n 1 ,
Revised: W ( t ) O ( i, · ) (cid:88)

Original: [Equation]
Revised: [Equation]

Original: ∈S 
Revised: ∈S n

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Similar to (138), we have that for l ∈ S n 2 , j ∈ S g 1 , and i ∈ U ( t) ,
Revised: t

Original: [Equation]
Revised: [Equation]

Original: ≲ Bt
Revised: ⊤ ·

Original: [Equation]
Revised: [Equation]

Original: (cid:88)
Revised: (cid:88)

Original: [Equation]
Revised: [Equation]

Original: (cid:88)
Revised: softmax

Original: [Equation]
Revised: [Equation]

Revised: ≳ (cid:16)

Original: Bt t (cid:88) b 
Revised: Bt t (cid:88) b =

Original: [Equation]
Revised: [Equation]

Original: (cid:88)
Revised: (cid:88)

Original: n ∈B b
Revised: n [Equation] ( b

Original: [Equation]
Revised: [Equation]

Revised: (cid:88)

Original: S n
Revised: n

Original: [Equation]
Revised: [Equation]

Original: p n ( b) ηtm
Revised: b

Original: [Equation]
Revised: [Equation]

Revised: ∥

Revised: [Equation]

Revised: ( t )

Revised: +

Revised: Bt t (cid:88) b

Revised: = (cid:88) n

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Revised: To study the case when l / ∈ S n 1 for all n ∈ [ N ]

Revised: [Equation]

Revised: we need to check all other l ’s. Recall that we focus on the coefficient of q 1 ( t ) in this part.

Revised: Based on the computation in (137) and (138), we know that the contribution of coefficient from non-discriminative patches is no more than that from discriminative patches,

Revised: i

Revised: [Equation]

Revised: e., for l / ∈ ( S n 1 ∪ S n 2 ) , n ∈ [ N ] and k ∈ S n 1 ,

Revised: [Equation]

Revised: [Equation]

Original: Therefore
Revised: Therefore

Original: [Equation]
Revised: [Equation]

Original: by the update rule
Revised: by the update rule

Original: [Equation]
Revised: [Equation]

Original: where the last step is by the condition that
Revised: where the last step is by the condition that

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: 145) for k 1 ( t ) > 0 and k 2 ( t ) > 0 from induction, i.e., q 1 ( t ) and r 1 ( t ) , q 1 ( t ) and r 1 ( t ) are in the same direction, respectively.
Revised: for k 1 ( t ) > 0 and k 2 ( t ) > 0 from induction, i.e., q 1 ( t ) and r 1 ( t ) , q 1 ( t ) and r 1 ( t ) are in the same direction, respectively.

Original: We also have
Revised: We also have

Original: K ( t )

Original: [Equation]
Revised: [Equation]

Original: Bt

Original: [Equation] Bt [Equation] Bt
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: aM

Original: [Equation] [Table] [Equation] [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation][Table] [Equation] [Table]
Revised: [Equation]

Original: [Equation] [Table] [Equation] [Table]
Revised: [Table]

Original: [Equation][Table] [Equation] [Table]
Revised: [Equation]

Original: [Equation] [Table] [Equation] [Table]
Revised: [Table]

Original: [Equation][Table] [Equation] [Table]
Revised: [Equation]

Original: [Equation] [Table] [Equation] [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table][Equation] [Table] [Equation]
Revised: [Table]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table][Equation] [Table] [Equation]
Revised: [Table]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: >
Revised: >

Original: [Equation]
Revised: [Equation]

Original: as long as
Revised: as long as

Original: [Equation]
Revised: [Equation]

Original: To find the sufficient condition for (149)
Revised: To find the sufficient condition for (150)

Original: [Equation]
Revised: [Equation]

Original: we first compare the first terms of both sides in (149).
Revised: we first compare the first terms of both sides in (150).

Original: Note that when
Revised: Note that when

Original: [Equation]
Revised: [Equation]

Original: O (
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: , (150) we have
Revised: 151) we have

Original: [Equation]
Revised: [Equation]

Original: When |S n | > |S n 1 | , min [Equation] [Equation] ϕ 
Revised: When |S n | > |S n 1 | , by (130), ϕ n ( t

Original: [Equation]
Revised: [Equation]

Original: S

Original: [Equation]

Original: −
Revised: −

Original: [Equation]
Revised: [Equation]

Original: ≳
Revised: From Definition 1, we know

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Since that
Revised: Moreover,

Original: [Equation]
Revised: [Equation]

Original: (153) for 0 < [Equation] [Equation] Moreover, |
Revised: ≤

Original: [Equation]
Revised: [Equation]

Original: (155) for [Equation] [Equation]
Revised: where the first inequality is by the triangle inequality and the second inequality comes from |S n 1 | b/ ( |S n | t ) ≤ 1 .

Original: Therefore, a sufficient condition for (149) is
Revised: Therefore, a sufficient condition for (150) is

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: i.e., [Equation] [Equation] by Hoeffding’s inequality in (26) given that |S n 1 | |S n | = Θ(1) .
Revised: by Hoeffding’s inequality in (26).

Original: Thus,
Revised: Thus,

Original: [Equation]

Original: and

Original: [Equation]
Revised: [Equation]

Original: if
Revised: if

Original: [Equation]
Revised: [Equation]

Original: for some ϵ 0 > 0 and
Revised: For the second terms on both sides in (150), since ( τ + σ ) ≤ O (1 /M ) , the inequality also holds with the same condition as in α ∗ and Bt . and

Original: [Equation]
Revised: [Equation]

Original: Given that |S n | = Θ(1) , (160) is equivalent to (8).
Revised: Given that |S n | = Θ(1) , (157) is equivalent to (8).

Original: For the second terms on both sides in (149), since ( σ + τ ) ≤ 1 /M , the inequality also holds with the same condition on α ∗ and Bt .
Revised: For the second terms on both sides in (150), since ( σ + τ ) ≤ 1 /M , the inequality also holds with the same condition on α ∗ and Bt .

Original: Note that if |S n | =
Revised: Note that if |S n | =

Original: |S n 1 | , we let |S nl | / ( |S n | − |S n 1 | )
Revised: |S n 1 | , we let |S nl | / ( |S n | − |S n 1 | )

Original: = 

Revised: =

Original: for l ∈ [ M ] .
Revised: for l ∈ [ M ] .

Original: We use the presentation in (147, 148) above and (173, 174) below for simplicity.
Revised: We use the presentation in (148, 149) above and (170, 171) below for simplicity.

Original: Then we give a brief derivation of W ( t +1)
Revised: Then we give a brief derivation of W ( t +1)

Original: Q x nj for j / ∈ S n 1 in the following.
Revised: Q x nj for j / ∈ S n 1 in the following.

Original: To be specific, for j ∈ S n / ( S n 1 ∪ S n 2 ) , 
Revised: To be specific, for j ∈ S n / ( S n 1 ∪ S n 2 ) , [Equation]

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: When K ( b ) is close to 0 + , we have
Revised: When K ( b ) is close to 0 + , we have

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where the first step is by log(1 + x )
Revised: where the first step is by log(1 + x )

Original: ≈ x when x → 0 + .
Revised: ≈ x when x → 0 + .

Original: Therefore, one can derive that
Revised: Therefore, one can derive that

Original: [Equation]
Revised: [Equation]

Original: Meanwhile, the value of p ′ n ( t ) will increase to 1 during training, making the component of q 1 ( t ) the major part in η 1 B (cid:80) n ∈B b ∂ Loss ( X n ,y n ) ∂ W ( t ) Q x nj .
Revised: Meanwhile, the value of p ′ n ( t ) will increase to 1 during training, making the component of q 1 ( t ) the major part in η 1 B (cid:80) n ∈B b ∂ Loss ( X n ,y n ) ∂ W ( t ) Q x nj .

Original: Hence, if j ∈ S nl for l ≥ 3 ,
Revised: Hence, if j ∈ S nl for l ≥ 3 ,

Original: [Equation]
Revised: [Equation]

Original: Similarly, for j ∈ S n 2 ,
Revised: Similarly, for j ∈ S n 2 ,

Original: [Equation]
Revised: [Equation]

Original: (b) For the gradient of W K , we have
Revised: (b) For the gradient of W K , we have

Original: [Equation]
Revised: [Equation]

Original: Loss b
Revised: Hence, for j ∈ S n

Original: [Equation]
Revised: [Equation]

Original: Hence, for j ∈ S n 1 , we can follow the derivation of (143) to obtain
Revised: we can follow the derivation of (144) to obtain

Original: [Equation]
Revised: [Equation]

Original: 171) where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: for λ < 1 introduced in Assumption3, and
Revised: for λ < 1 introduced in Assumption

Original: [Equation]
Revised: [Equation]

Revised: and

Original: [Equation]
Revised: [Equation]

Original: Similarly, for j ∈ S n 2 , we have
Revised: Similarly, for j ∈ S n 2 , we have

Original: [Equation]
Revised: [Equation]

Original:  t [Equation] = γ [Equation]
Revised: Q ( t ))

Original:  t ) , (175) For j ∈ S nl , l = 3 , 4 , · · · , M , we have W [Equation]
Revised: q 2 ( t ) + Θ(1) · o nj ( t ) ± t − 1 (cid:88) b = | Q e ( b ) | (1 − λ ) r 1 ( b )

Original: t [Equation] 1) K x [Equation]
Revised: +

Original: 176) where | Q f ( t ) | [Equation] Q ( t ) (177)
Revised: M (cid:88) l = γ ′ l r l ( t ) , (172)

Original: Therefore, for l ∈ S n 1 , if j ∈ S n 1 , x nj ⊤ W ( t +1) K ⊤ W ( t +1) Q x nl ≳ (1 + K ( t ))(1 + Q ( t ))
Revised: For j ∈ S nz , z = 3 , 4 , · · · , M , we have (cid:12)(cid:12)(cid:12) (cid:42) B (cid:88) n ∈B b ∂ Loss ( X n , y n ) ∂F

Original: ∥ q 1 ( t ) ∥ 2 − ( δ + τ ) ∥ q 1 ( t ) ∥ + K e ( t ) Q e ( t ) ∥
Revised: [Equation]

Original: q 2 ( t )∥∥ r 2 ( t ) ∥ +
Revised: X )

Original: M (cid:88) l = γ l γ ′ l ∥ q l ( t ) ∥∥
Revised: [Equation] (cid:43) (cid:12)(cid:12)(cid:12) ≲ (cid:12)(cid:12)(cid:12)

Original: r l ( t ) ∥ ≳ (1 + K ( t ))(1 + Q ( t ))
Revised: B (cid:88) n ∈B b ( − y n ) (cid:88) l ∈S nm (cid:88) i = a ( l ) i

Original: ∥
Revised: [Equation]

Original: q 1 ( t ) ∥ 2 − ( δ + τ ) ∥ q 1 ( t ) 
Revised: [ W O ( i, · ) W V X [Equation]

Original: ∥ − (cid:118)(cid:117)(cid:117)(cid:116) M (cid:88) l = ( 
Revised: ( X n ⊤ W ⊤ [Equation]

Original: B (cid:88) n ∈B b Q ( t ) |S nl| |S n | − |S n ∗ | ) 
Revised: Q x nl ) ≥ 0] · (cid:16) W O ( i, · ) ( (cid:88) [Equation]

Original: 2 ∥ r l( t )  ∥ 2 · (cid:118)(cid:117)(cid:117)(cid:116)
Revised: [Equation] [Equation] l ) (cid:17) ∥ [Equation]

Original: M (cid:88) l = ( 
Revised: ( t ) [Equation] (cid:12)(cid:12) (cid:12) ≤ λ | Q f ( t ) |∥ q 1 ( t ) ∥(173) (cid:12)(cid:12)(cid:12) (cid:42)

Original: B (cid:88) n ∈B b 
Revised: B (cid:88) n ∈B b ∂ Loss ( X n , y n ) ∂F ( X ) F

Original: K ( t ) |S nl | |S n | − |S n ∗ | ) [-
-]
Revised: [Equation] X ) [-
-]∂ W K [Equation] [Equation]

Original: ∥ q l ( t ) ∥≳ (1 + K ( t ) + Q ( t ))
Revised: ( t ) (cid:43) (cid:12)(cid:12)(cid:12) ≲ (cid:12)(cid:12)(cid:12) B (cid:88) [Equation] where 0 < c k 1 , c k 2 , c k 3 < 1 , and

Original: ∥
Revised: [Equation] Therefore, for l ∈ S n 1 , if j ∈ S n 1 , [Equation] where the second step is by Cauchy-Schwarz inequality.

Original: q 1 ( t ) ∥ 2 − ( δ + τ ) ∥
Revised: If j / ∈ S n 1 , [Equation] + τ ) ∥

Original: q 1 ( t ) ∥ (178)
Revised: q 1 ( t ) ∥ (178)

Original: If
Revised: Hence, for j, l ∈ S n 1 , softmax ( x nj ⊤ W ( t +

Original: [Equation]
Revised: [Equation]

Original: Hence, for j, l ∈ S n 1 ,
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where the second to last step is by the Mean Value Theorem with
Revised: where the second to last step is by the Mean Value Theorem with

Original: [Equation]
Revised: [Equation]

Original: The same conclusion holds if l / ∈ ( S n 1 ∪ S n 2 ) and j ∈ S n 1 .
Revised: We then need to study if l / ∈ ( S n 1 ∪ S n 2 ) and j ∈ S n 1 , i.e., x [Equation] For j, l / ∈ ( S n 1 ∪ S n 2 ) , [Equation]

Revised: [Equation]

Revised: We know that the magnitude of ∥ q 1 ( t ) ∥ increases along the training and finally reaches no larger than Θ( √ log T ) .

Revised: At the final step, we have

Revised: [Equation]

Revised: Therefore, when t is large enough during the training but before the final step of convergence

Revised: [Equation]

Revised: we have if l / ∈ ( S n 1 ∪ S n 2 ) and j, j ′ ∈ S n 1 , we can obtain

Revised: [Equation]

Original: Meanwhile, for l ∈ S n 1 and j / ∈ S n 1 ,
Revised: Meanwhile, for l ∈ S n 1 and j / ∈ S n 1 ,

Original: [Equation]
Revised: [Equation]

Original: where the second to last step is by the Mean Value Theorem with
Revised: where the second to last step is by the Mean Value Theorem with

Original: [Equation]
Revised: [Equation]

Original: The same conclusion holds if l / ∈ ( S n 1 ∪ S n 2 ) and j / ∈ S n 1 .
Revised: The same conclusion holds if l / ∈ ( S n 1 ∪ S n 2 ) and j / ∈ S n 1 .

Original: Note that
Revised: Note that

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: It can also be verified that this claim holds when t = 1 .
Revised: It can also be verified that this claim holds when t = 1 .

Original: D. P ROOF OF C LAIM 3 OF L EMMA
Revised: E.

Revised: P ROOF OF C LAIM 3 OF L EMMA

Original: The computation of the gradient of W V is straightforward.
Revised: The computation of the gradient of W V is straightforward.

Original: The gradient would be related to W O by their connections.
Revised: The gradient would be related to W O by their connections.

Original: One still need to study the influence of the gradient on different patterns, where we introduce the discussion for the term
Revised: One still need to study the influence of the gradient on different patterns, where we introduce the discussion for the term

Original: V i ( t ) ’s.
Revised: V i ( t ) ’s.

Original: For the gradient of W V , by (18) we have ∂ Loss b
Revised: For the gradient of W V , by (18) we have ∂ Loss b

Original: [Equation]
Revised: [Equation]

Original: Consider a data { X n , y n } where y n = 1 .
Revised: Consider a data { X n , y n } where y n = 1 .

Original: Let l ∈ S n(cid:88)
Revised: Let l ∈ S n

Original: [Equation]
Revised: [Equation]

Original: Then for j ∈ S g 1 , g ∈ [ N ] ,
Revised: Then for j ∈ S g 1 , g ∈ [ N ] ,

Original: B (cid:88) n
Revised: [Equation] If i ∈ W ( t )

Original: [Equation]
Revised: [Equation]

Original: If i ∈ W ( t ) , by the fact that S n # contributes more to V i ( t ) compared to S nl for l ≥ 3 and Assumption 3, we have
Revised: by the fact that S n # contributes more to V i ( t ) compared to S nl for l ≥ 3 and Assumption [Equation] we have

Original: [Equation]
Revised: [Equation]

Original: Similarly, if i ∈ U ( t ) ,
Revised: Similarly, if i ∈ U ( t ) ,

Original: [Equation]
Revised: [Equation]

Original: 194) if i is an unlucky neuron, by Hoeffding’s inequality in (26), we have
Revised: if i is an unlucky neuron, by Hoeffding’s inequality in (26), we have

Original: [Equation]
Revised: [Equation]

Original: For i ∈ W (0) , we have 
Revised: For i ∈ W (0) , we have −

Original: [Equation]
Revised: [Equation]

Original: Hence, (1) If j ∈ S n 1 for one n ∈ [ N ],
Revised: Hence, 1)If j ∈ S n 1 for one n ∈ [ N ]

Original: [Equation]
Revised: [Equation]

Original: (2) If j ∈ S n 2 , we have
Revised: (2) If j ∈ S n 2 , we have

Original: [Equation]
Revised: [Equation]

Original: (3) If j ∈ S n / ( S n 1 ∪ S n 2 ) , we have
Revised: (3) If j ∈ S n / ( S n 1 ∪ S n 2 ) , we have

Original: [Equation]
Revised: [Equation]

Original: j ( t ) (201)

Original: Here ∥z
Revised: Here ∥

Original: [Equation]
Revised: [Equation]

Original: t

Original: [Equation]
Revised: [Equation]

Original: (202) for t ≥ 1 .
Revised: for t ≥ 1 .

Original: Note that this claim also holds when t = 1 .
Revised: Note that this claim also holds when t = 1 .

Original: E O THER USEFUL LEMMAS
Revised: F O THER USEFUL LEMMAS

Original: Lemma 3.
Revised: Lemma 3.

Original: If the number of neurons m is larger enough such that
Revised: If the number of neurons m is larger enough such that

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: the number of lucky neurons at the initialization |W (0) | , |U (0) | satisfies
Revised: the number of lucky neurons at the initialization |W (0) | , |U (0) | satisfies

Original: [Equation]
Revised: [Equation]

Original: Proof:
Revised: Proof:

Original: Let θ l be the angle between  the initial weightand p l for one i ∈ [ m ] and all l ∈ [ M ] .
Revised: Let θ l be the angle between p l and the initial weight for one i ∈ [ m ] and all l ∈ [ M ] .

Original: For the lucky neuron i ∈ W (0) , θ 1 should be the smallest among { θ l } Ml =1 with noise ∆ θ .
Revised: For the lucky neuron i ∈ W (0) , θ 1 should be the smallest among { θ l } M l =1 with noise ∆ θ .

Original: Hence, the probability of the lucky neuron can be bounded as
Revised: Hence, the probability of the lucky neuron can be bounded as

Original: [Equation]
Revised: [Equation]

Original: where the first step is because the Gaussian W (0) O ( i, · ) and orthogonal p l ,
Revised: where the first step is because the Gaussian W (0) O ( i, · ) and orthogonal p l ,

Original: l ∈ [ M ] generate independent W (0)
Revised: l ∈ [ M ] generate independent W (0)

Original: O ( i, · ) p l .
Revised: O ( i, · ) p l .

Original: From the definition of W (0) , we have
Revised: From the definition of W (0) , we have

Original: [Equation]
Revised: [Equation]

Original: which implies
Revised: which implies

Original: [Equation]
Revised: [Equation]

Original: for small σ > 0 .
Revised: for small σ > 0 .

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: where the first step comes from that θ 1 follows the uniform distribution on [0 , 2 π ] due to the Gaussian initialization of W O .
Revised: where the first step comes from that θ 1 follows the uniform distribution on [0 , 2 π ] due to the Gaussian initialization of W O .

Original: We can define the random variable v i such that
Revised: We can define the random variable v i such that

Original: [Equation]
Revised: [Equation]

Original: We know that v i belongs to Bernoulli distribution with probability 1 M (1 − ( σ + τ ) M π ) .
Revised: We know that v i belongs to Bernoulli distribution with probability 1 M (1 − ( σ + τ ) M π ) .

Original: By Hoeffding’s inequality in (26), we know that with probability at least 1 − N − 10 ,
Revised: By Hoeffding’s inequality in (26), we know that with probability at least 1 − N − 10 ,

Original: [Equation]
Revised: [Equation]

Original: Let m ≥ Θ
Revised: Let m ≥ Θ

Original: ( ϵ − 2 m M 2 log B ) , we have
Revised: ( ϵ − 2 m M 2 log B ) , we have

Original: [Equation]
Revised: [Equation]

Original: where we require
Revised: where we require

Original: [Equation]
Revised: [Equation]

Original: to ensure a positive probability in (211).
Revised: to ensure a positive probability in (215).

Original: Likewise, the conclusion holds for U (0) .
Revised: Likewise, the conclusion holds for U (0) .

Original: Lemma 4.
Revised: Lemma 4.

Original: Let W ( t ) and U ( t ) be defined in Definition 2.
Revised: Let W ( t ) and U ( t ) be defined in Definition 2.

Original: We then have
Revised: We then have

Original: [Equation]
Revised: [Equation]

Original: as long as
Revised: as long as

Original: [Equation]
Revised: [Equation]

Original: Proof:
Revised: Proof:

Original: We show this lemma by induction.
Revised: We show this lemma by induction.

Original: (1) t = 0 .
Revised: (1) t = 0 .

Original: For i ∈ W (0) , by Definition 2, we know that the angle between W (0) O
Revised: For i ∈ W (0) , by Definition 2, we know that the angle between W (0) O

Original: ( i, · ) and p 1 is smaller than ( σ + τ ) .
Revised: ( i, · ) and p 1 is smaller than ( σ + τ ) .

Original: Hence, we have
Revised: Hence, we have

Original: [Equation]
Revised: [Equation]

Original: for all p ∈ P / p 1 .
Revised: for all p ∈ P / p 1 .

Original: (2) Suppose that the conclusion holds when t = s .
Revised: (2) Suppose that the conclusion holds when t = s .

Original: When t = s + 1 , from Lemma 2 Claim 1, we can obtain 
Revised: When t = s + 1 , from Lemma 2 Claim 1, we can obtain (cid:28)

Original: [Equation]
Revised: [Equation]

Original:  and 
Revised: (221) and (cid:12)(cid:12) (cid:12) (cid:28) W ( [Equation] +1)

Revised: [Equation]

Revised: ( i

Revised: [Equation]

Revised: [Equation]

Revised: p (cid:29)

Revised: [Equation]

Revised: (cid:29) (cid:12)(cid:12) (cid:12) ≲ η m ·

Revised: B (cid:88) n

Revised: [Equation]

Revised: b

Original: [Equation]
Revised: [Equation]

Revised: p

Revised: [Equation]

Revised: +

Revised: [Equation]

Revised: M )

Revised: [Equation]

Original:  Combining (217) and (218), we can approximately compute that if
Revised: B B ∥ [Equation] Combining (221) and (222), we can approximately compute that if

Original: [Equation]
Revised: [Equation]

Original: we can derive
Revised: we can derive

Original: [Equation]
Revised: [Equation]

Original: Therefore, we have
Revised: Therefore, we have

Original: [Equation]
Revised: [Equation]

Original: In conclusion, we can obtain
Revised: In conclusion, we can obtain

Original: [Equation]
Revised: [Equation]

Original: for all t ≥ 0 .
Revised: for all t ≥ 0 .

Original: One can develop the proof for U ( t ) following the above steps.
Revised: One can develop the proof for U ( t ) following the above steps.

Original: F E XTENSION TO MORE GENERAL CASES F. E XTENSION TO MULTI - CLASSIFICATION
Revised: G E XTENSION TO MORE GENERAL CASES G. E XTENSION TO MULTI - CLASSIFICATION

Original: Consider the classification problem with four classes, we use the label y ∈ { +1 , − 1 } 2 to denote the corresponding class.
Revised: Consider the classification problem with four classes, we use the label y ∈ { +1 , − 1 } 2 to denote the corresponding class.

Original: Similarly to the previous setup, there are four orthogonal discriminative patterns.
Revised: Similarly to the previous setup, there are four orthogonal discriminative patterns.

Original: In the output layer, a l ( i ) for the data ( X n , y n ) is changed into an R 2 vector a l ( i ) for l ∈ [ |S n | ] and i ∈ [ m ] .
Revised: In the output layer, a l ( i ) for the data ( X n , y n ) is changed into an R 2 vector a l ( i ) for l ∈ [ |S n | ] and i ∈ [ m ] .

Original: Hence, we define
Revised: Hence, we define

Original: [Equation]
Revised: [Equation]

Original: (cid:88) [Equation] [Equation] |
Revised: The dataset D can be divided into four groups as

Original: [Equation]
Revised: [Equation]

Original: (cid:88) [Equation] [Equation]
Revised: The hinge loss function for data ( X n , y n ) will be

Original: [Equation] (cid:88) [Equation] The dataset D can be divided into four groups as
Revised: Loss ( X n , y n

Original: [Equation]
Revised: [Equation]

Original: The hinge loss function for data ( X n , y n ) will be
Revised: = max

Original: [Equation]
Revised: [Equation]

Original: We can divide the weights
Revised: We can divide the weights

Original: W O
Revised: W O

Original: ( i, · ) ( i ∈ [ m ] ) into two groups, respectively.
Revised: ( i, · ) ( i ∈ [ m ] ) into two groups, respectively.

Original: [Equation]
Revised: [Equation]

Original: Therefore, for W O u in the network (223) [Equation] we have
Revised: Therefore, for W O u in the network (227), we have

Original: [Equation]
Revised: [Equation]

Original: where the derivation of ∂F 1 ( X n ) ∂ W O 

Revised: where the derivation of ∂F 1 ( X n ) ∂ W O

Original: ( i, · ) and ∂F 2 ( X n ) ∂ W O 

Revised: ( i, · ) and ∂F 2 ( X n ) ∂ W O

Original: ( i, · ) can be found in the analysis of binary classification above.
Revised: ( i, · ) can be found in the analysis of binary classification above.

Original: For any i ∈ W 2 , following the proof of Claim 1 of Lemma 2, if the data ( X n , y n ) ∈ D 2 , we have
Revised: For any i ∈ W 2 , following the proof of Claim 1 of Lemma 2, if the data ( X n , y n ) ∈ D 2 , we have

Original: [Equation]
Revised: [Equation]

Original: if ( X n , y n ) ∈ D 1 , we have
Revised: if ( X n , y n ) ∈ D 1 , we have

Original: [Equation]
Revised: [Equation]

Original: if ( X n , y n ) ∈ D 3 , we have 
Revised: if ( X n , y n ) ∈ D 3 , we have − ∂ Loss [Equation] X n , y n )

Original: [Equation]
Revised: [Equation]

Original: if ( X n , y n ) ∈ D 4 , we have
Revised: if ( X n , y n ) ∈ D 4 , we have

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: By the algorithm, W O ( i, · ) will update along the direction of p 2 for i ∈
Revised: By the algorithm, W O ( i, · ) will update along the direction of p 2 for i ∈

Original: W 2 .
Revised: W 2 .

Original: We can analyze W V , W
Revised: We can analyze W V , W

Original: K and W Q similarly.
Revised: K and W Q similarly.

Original: F. E XTENSION TO A MORE GENERAL DATA MODEL
Revised: G. E XTENSION TO A MORE GENERAL DATA MODEL

Original: We generalize the patterns from vectors to sets of vectors.
Revised: We generalize the patterns from vectors to sets of vectors.

Original: Consider that there are M (2 < M < m a , m b ) distinct sets {M 1 , M 2 , · · · , M M } where M l = { µ l, 1 , µ l, 2 , · · · , µ l,l m } , l m ≥ 1 .
Revised: Consider that there are M (2 < M < m a , m b ) distinct sets {M 1 , M 2 , · · · , M M } where M l = { µ l, 1 , µ l, 2 , · · · , µ l,l m } , l m ≥ 1 .

Original: M 1 ,
Revised: M 1 ,

Original: M 2 denote sets of discriminative patterns for the binary labels, and M 3 , · · · , M M are sets of nondiscriminative patterns.[Equation]
Revised: M 2 denote sets of discriminative patterns for the binary labels, and M 3 , · · · , M M are sets of nondiscriminative patterns.

Original: [Equation]
Revised: [Equation]

Original:  is the minimum distance between patterns of different sets.
Revised: 242) is the minimum distance between patterns of different sets.

Original: Each token x nl of X n is a noisy version of one pattern, i.e.,
Revised: Each token x n l of X n is a noisy version of one pattern, i.e.,

Original: [Equation]
Revised: [Equation]

Original: and τ ≤ κ .
Revised: Define that for l , s corresponding to b 1 , b 2 , respectively, [Equation] [Equation] we have 2 τ + ∆ < κ .

Original: To simplify our theoretical analysis, one can similarly rescale all tokens a little bit like in Assumption 3 such that tokens corresponding to patterns in the same pattern set has an inner product larger than 1, while tokens corresponding to patterns from different pattern sets has an inner product smaller than λ < 1 .
Revised: To simplify our theoretical analysis, one can similarly rescale all tokens a little bit like in Assumption 3 such that tokens corresponding to patterns in the same pattern set has an inner product larger than 1, while tokens corresponding to patterns from different pattern sets has an inner product smaller than λ < 1 .

Original: Assumption 1 can be modified such that
Revised: Assumption 1 can be modified such that

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where p j,b ⊥ p i,a for any i, j ∈ [ M ] , b ∈ [ j m ] , and a ∈ [ i m ] .
Revised: where p j,b ⊥ p i,a for any i, j ∈ [ M ] , b ∈ [ j m ] , and a ∈ [ i m ] .

Original: Likewise, q j,b ⊥ q i,a for any i, j ∈ [ M ] , b ∈ [ j m ] , and a ∈ [ i m ] .
Revised: Likewise, q j,b ⊥ q i,a for any i, j ∈ [ M ] , b ∈ [ j m ] , and a ∈ [ i m ] .

Original: r j,b ⊥ r i,a for any i, j ∈ [ M ] , b ∈ [ j m ] , and a ∈ [ i m ] .
Revised: r j,b ⊥ r i,a for any i, j ∈ [ M ] , b ∈ [ j m ] , and a ∈ [ i m ] .

Original: Therefore, we make sure that the initial query, key and value features from different sets of patterns are still close to be orthogonal to each other.
Revised: Therefore, we make sure that the initial query, key and value features from different sets of patterns are still close to be orthogonal to each other.

Original: Then, we can follow our main proof idea.
Revised: Then, we can follow our main proof idea.

Original: To be more specific, for label-relevant tokens x nl , by computing (125) and (170), W ( t ) K x nl , W ( t ) Q x nl will grow in the direction of a fixed linear combination of q l, 1 , · · · , q l,l m , and r l, 1 , · · · , r l,l m .
Revised: To be more specific, for label-relevant tokens x nl , by computing (125) and (167), W ( t ) K x nl , W ( t ) Q x nl will grow in the direction of a fixed linear combination of q l, 1 , · · · , q l,l m , and r l, 1 , · · · , r l,l m .

Original: The coefficient of the linear combination is a function of fractions of different pattern vectors µ l,b in M l .
Revised: The coefficient of the linear combination is a function of fractions of different pattern vectors µ l,b in M l .

Original: One can still derive a sparse attention map with weights of non-discriminative patterns decreasing to be close to zero during the training.
Revised: One can still derive a sparse attention map with weights of non-discriminative patterns decreasing to be close to zero during the training.

Original: F. E XTENSION TO MULTI - HEAD NETWORKS
Revised: G. E XTENSION TO MULTI - HEAD NETWORKS

Original: Suppose there are H heads in total.
Revised: Suppose there are H heads in total.

Original: The network is modified to
Revised: The network is modified to

Revised: F ( X n ) =

Revised: |S n | (cid:88)

Original: [Equation]
Revised: [Equation]

Original: where W V h ∈ R m a × d , W O =
Revised: where W V h ∈ R m a × d , W O =

Original: ( W O 1 , W O 2 , · · · W O H ) ∈
Revised: ( W O 1 , W O 2 , · · · W O H ) ∈

Original: R m ×
Revised: R m ×

Original: Hm a , and W O
Revised: Hm a , and W O

Original: h ∈ R m × m a for h ∈ [ H ] .
Revised: h ∈ R m × m a for h ∈ [ H ] .

Original: One can make similar assumptions for W (0) Q h , W (0) K h , and W (0) V h as in Assumption 1.
Revised: One can make similar assumptions for W (0) Q h , W (0) K h , and W (0) V h as in Assumption 1.

Original: Note that { p 1 , p 2 , · · · , p M } needs to be changed into { p h 1 , p h 2 , · · · , p h M } , and the set { p h 1 , p h 2 , · · · , p h M } can vary for different h ∈
Revised: Note that { p 1 , p 2 , · · · , p M } needs to be changed into { p h 1 , p h 2 , · · · , p h M } , and the set { p h 1 , p h 2 , · · · , p h M } can vary for different h ∈

Original: [ H ] .
Revised: [ H ] .

Original: It is also the same for { q h 1 , q h 2 , · · · , q h M } and { r h 1 , r h 2 , · · · , r h M } for different h ∈
Revised: It is also the same for { q h 1 , q h 2 , · · · , q h M } and { r h 1 , r h 2 , · · · , r h M } for different h ∈

Original: [ H ] .
Revised: [ H ] .

Original: Based on the modified assumption with H heads, the backbone of the proof remains the same.
Revised: Based on the modified assumption with H heads, the backbone of the proof remains the same.

Original: Lucky neurons in W O tend to learn ( p ⊤ 1 1 , p ⊤ 2 1 , · · · , p ⊤ H 1 )
Revised: Lucky neurons in W O tend to learn ( p ⊤ 1 1 , p ⊤ 2 1 , · · · , p ⊤ H 1 )

Original: ⊤ and ( p ⊤ 1 2 , p ⊤ 2 2 , · · · , p ⊤ H 2 ) ⊤ .
Revised: ⊤ and ( p ⊤ 1 2 , p ⊤ 2 2 , · · · , p ⊤ H 2 ) ⊤ .

Original: Hence, the properties of the Relu activation are almost the same as the single-head case because luck neurons are still activated by either of two label-relevant patterns with a high probability.
Revised: Hence, the properties of the Relu activation are almost the same as the single-head case because luck neurons are still activated by either of two label-relevant patterns with a high probability.

Original: In fact, one can expect a more stable training process by multiple heads due to a more stable Relu gate for lucky neurons.
Revised: In fact, one can expect a more stable training process by multiple heads due to a more stable Relu gate for lucky neurons.

Original: F. E XTENSION TO SKIP CONNECTIONS AND NORMALIZATION
Revised: G. E XTENSION TO SKIP CONNECTIONS AND NORMALIZATION

Original: Consider a basic case where a skip connection is added after the self-attention layer.
Revised: Consider a basic case where a skip connection is added after the self-attention layer.

Original: Let m a = d .
Revised: Let m a = d .

Original: The network is changed into
Revised: The network is changed into

Original: [Equation]
Revised: [Equation]

Original: The assumption of W (0) V in Assumption 1 should be changed into
Revised: The assumption of W (0) V in Assumption 1 should be changed into

Original: [Equation]
Revised: [Equation]

Original: while the assumption of W (0) Q and W (0) K remain the same.
Revised: while the assumption of W (0) Q and W (0) K remain the same.

Original: One can easily verify that the gradients of W K , W Q , and W V for (244) are almost the same as those for (1) except for the Relu gate.
Revised: One can easily verify that the gradients of W K , W Q , and W V for (249) are almost the same as those for (1) except for the Relu gate.

Original: The major differences come from the gradient of W O , which also helps to determine the Relu gate.
Revised: The major differences come from the gradient of W O , which also helps to determine the Relu gate.

Original: One needs to redefine
Revised: One needs to redefine

Original: [Equation]
Revised: [Equation]

Original: for l ∈ S n 1 . The inner product between the lucky neuron and the term (cid:80) j ̸ =1 W nj ( t )( x nj + x nl ) can still be upper bounded by the inner product between the lucky neuron and the term (cid:80) [Equation]
Revised: ) can still be upper bounded by the inner product between the lucky neuron and the x

Original:  ⊤ W( t ) Q x nl ) p 1 given good initialization of W K and W Q .
Revised: ) K ⊤ W t ) Q x nl ) p 1 given good initialization of W K and W Q .

Original: Therefore, (244) can be analyzed following our proof techniques.
Revised: Therefore, (249) can be analyzed following our proof techniques.

Original: For layer normalization, one usually use that approach to normalize each data.
Revised: For layer normalization, one usually use that approach to normalize each data.

Original: We think it is consistent with our normalization of x nl , which plays an important role in our proof.
Revised: It is consistent with our normalization of x nl , which plays an important role in our proof.

Original: By normalization, the training process becomes more stable because of the unified norm of all tokens.
Revised: By normalization, the training process becomes more stable because of the unified norm of all tokens.