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

Original: OMPLEXITY
Revised: OMPLEXITY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: 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 three-layer 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 three-layer 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 nondiscriminative 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 athree-layer 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 three-layer 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, thethree-layer ViT model considered in this paper is more complicated to analyze than three-layer CNNs considered in AllenZhu 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 inter- actions 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) 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.e.,
Revised: where ℓ ( X n , y n ; ψ ) is the Hinge loss function, i [Equation] 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.

Original: Let κ = min 1 ≤ i ̸ =
Revised: Let κ = 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.
Revised: 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.
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: and the noise level τ ≤ κ .
Revised: and the noise level τ ≤ κ .

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.
Revised: Some important notations are summarized in Table 1.

Original: notations
Revised: notations

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: δ
Revised: δ

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: τ
Revised: τ

Original: [Table]
Revised: [Table]

Original: M
Revised: M

Original: [Table]
Revised: [Table]

Original: m
Revised: m

Original: The
Revised: The

Original: [Table]
Revised: [Table]

Original: in W
Revised: in W

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: of
Revised: of

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

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [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 } M j =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

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]

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: for all l ∈ S n and n ∈ [ N ] .
Revised: for all l ∈ S n and n ∈ [ N ] .

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: ViT.
Revised: 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 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.

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

Original: (Hron et al., 2020;
Revised: (Hron et al., 2020;

Original: Yang, 2020; Allen-Zhu et al., 2019a;
Revised: Yang, 2020; Allen-Zhu et al., 2019a;

Original: Arora et al., 2019).
Revised: 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: 35 and α # = 0 .δ = 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).
Revised: In Figure 2, we verify the number of iterations T against α − 1 ∗ in (11).

Original: Advantage of self-attention : To verify Proposition 1, we compare the performance on ViT inand on the same network with W K and W Q fixed during the training, i.e., a three-layer CNN.Set α ∗ = 0 . 3 and α # = 0 .
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: The sample complexity of CNN significantly increases, as shown in  Figure 3, indicating a much worse generalization of CNN than that of ViT.
Revised: Figure 3 shows the sample complexity of CNN is almost linear in α − 4 ∗ as predicted in (13). 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: 2 , δ = 0 . 2 , α ∗ = 0 .
Revised: 1 , δ = 0 . 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 labelrelevant tokens, i.e., the left side of (15) averaged over all l , indeed increases to 1 when the number of iterations increases.
Revised: In Figure 4, the red line with asterisks shows that the sum of attention weights on labelrelevant tokens, i.e., the left side of (15) averaged over all l , indeed increases to 1 when the number of iterations increases.

Original: Correspondingly, the sum of attention weights on other tokens decreases to 0 , as shown in the blue line with squares.
Revised: Correspondingly, the sum of attention weights on other tokens decreases 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.

Revised: The experiment is duplicated for 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: 2 , δ = 0 .
Revised: 1 , δ = 0 .

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

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

Original: If we apply the 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 the 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 24 ×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 IMA- GENET.
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 α ∗= 9 / 16 ≈ 0 . 
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 toa 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).

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

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 (b), 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 (b) 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 P RELIMINARIES
Revised: A M ORE EXPERIMENTS

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

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

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

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

Revised: One can observe the same trend in

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 .

Revised: B P RELIMINARIES

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: We first formally restate the neural network with different notations of loss functions, and the Algorithm 1 of the training steps after token sparsification. The notations used in the Appendix is summarized in Table 2. For the network [Equation]

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

Revised: [Equation]

Revised: [Table]

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

Revised: [Table]

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

Revised: [Table]

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

Revised: We have p j (0) =

Revised: p j , q j (0) =

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

Revised: r j .

Revised: [Table]

Revised: nj ( t

Revised: [Table]

Revised: n nj ( t ) , o nj ( 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 .

Revised: [Table]

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

Revised: [Table]

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

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

Revised: [Table]

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

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.

Revised: [Table]

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

Revised: Loss

Revised: [Equation]

Revised: N (cid:88) n

Revised: [Equation]

Revised: [Equation]

Revised: The formal algorithm is as follows.

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

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

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

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

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

Revised: [Equation]

Revised: 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 ) ∈

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

Original: R m b × M and R =
Revised: R m b × M and R =

Original: ( r 1 , r 2 , · · · , r M ) ∈
Revised: ( 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 n j ( t ) , n n j ( t ) and o n j ( t ) with ∥ z n j (0) ∥ ≤ σ + τ and ∥ n n j (0) ∥ , ∥ o n j (0) ∥ ≤ δ + τ[Equation] 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 + z n j (0) 
Revised: j + 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 nj (0) .
Revised: W (0) Q x nl = r j + o n j (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.

Original: [Equation]
Revised: [Equation]

Original: 2. ν

Original: [Equation]

Original: =

Original: 1 |S n 1 | e ∥ q 1( t ) ∥ 2 − ( δ + τ ) ∥ q 1( t ) ∥

Original: + |S n |−|S n 1 | .

Original: p n ( t ) =

Original: |S n 1 | e ∥ q 1 ( t ) ∥ 2 − ( δ + τ ) ∥

Original: q 1 ( t ) ∥ ν n ( t ) .

Original: S n ∗ = (cid:26) S n 1 , if y n =S n 2 , if y n = − 1 , S n # = (cid:26) S n 2 , if y n =S n 1 , if y n = −5.

Original: α ∗ =

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

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

Original: |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 .

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 − 1 2 ( 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: B P ROOF OF THE M AIN T HEOREM AND P ROPOSITIONS
Revised: C P ROOF OF THE M AIN T HEOREM AND P ROPOSITIONS

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

Revised: Lemma 2.

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

Revised: [Equation]

Original: Proof of Theorem 1:
Revised: with [Equation] We also have the following claims: Claim 1.

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

Revised: [ T ] , we have

Revised: [Equation]

Revised: for p ∈ P

Revised: [Equation]

Revised: p

Revised: [Equation]

Revised: [Equation]

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

Revised: [Equation]

Revised: [Equation]

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

Revised: [Equation]

Revised: p

Revised: [Equation]

Revised: P

Revised: [Equation]

Revised: 38) and for the noise

Revised: z l ( t ) ,

Revised: [Equation]

Revised: [Equation]

Revised: For unlucky neurons, we have

Revised: [Equation]

Revised: p

Revised: [Equation]

Revised: P

Revised: [Equation]

Revised: [Equation]

Revised: Claim 2.

Revised: There exists K ( t ) , Q ( t ) > 0 , t = 0 , 1 , · · · , T − 1 such that for j ∈ S n ∗ and l / ∈ S ∗ , softmax ( x nj ⊤ W ( t +1) K

Revised: [Equation]

Revised: 43) softmax ( x nj ⊤ W ( t +1)

Revised: [Equation]

Revised: and for j / ∈ S n ∗ , l / ∈ S n ∗ , we have

Revised: [Equation]

Revised: For i = 1 , 2 ,

Revised: [Equation]

Revised: Claim 3.

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

Revised: [Equation]

Revised: Proof of Theorem 1:

Original: For y n = 1 ,
Revised: For y n = 1 ,

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: (28) where the second step comes from (53) and the last step is by (135).
Revised: where the second step comes from (27) and the last step is by (135).

Original: By the definition of K + +, we have
Revised: By the definition of K + +

Original: [Equation]
Revised: [Equation]

Revised: we have

Original: [Equation]
Revised: [Equation]

Original: Combining (135) and (137), we can obtain
Revised: Combining (135) and (137)

Original: [Equation]
Revised: [Equation]

Original: n
Revised: we can obtain

Original: [Equation]
Revised: [Equation]

Original: (cid:88)

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

Original: 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: K

Original: [Equation]
Revised: [Equation]

Original: t )
Revised: Since that

Original: [Equation]
Revised: [Equation]

Original: ≳

Original: [Equation]
Revised: [Equation]

Original: Since that

Original: [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] [Equation]
Revised: [Equation] [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

Revised: [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 (35).
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,

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

Original: We know that
Revised: We know that

Original: [Equation]
Revised: [Equation]

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

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

Revised: [Equation]

Revised: where α ≥

Original: [Equation]
Revised: [Equation]

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

Original: N 
Revised: N ≥

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

Original: and
Revised: and

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: [Equation]

Original: which can scale up the sample complexity in (43) 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] [Equation]
Revised: [Equation] [Equation]

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

Original: [Equation]

Original: [Equation]
Revised: [Equation]

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

Original: [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: Proof of 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
Revised: The only difference is that we need to modify

Original: [Equation]

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

Original: N ≥
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, (34), and (35). C U SEFUL LEMMAS Lemma 2.
Revised: It can be easily derived from Claim 2 of Lemma 2, (60), and (61).

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

Original: [Equation]

Original: with

Original: [Equation]

Original: an

Original: [Equation]

Original: We also have the following claims:

Original: Claim 1.

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

Original: [ T ] , we have

Original: [Equation]

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

Original: [Equation]

Original: [Equation]

Original: For i ∈ U ( t ) , we also have equations as in (58) to (61), including

Original: [Equation] p [Equation] [Equation]
Revised: D P ROOF OF L EMMA

Original: [Equation]

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

Original: [Equation]

Original: [Equation]

Original: For unlucky neurons, we have

Original: [Equation] [Equation] Claim 2.
Revised: Proof of Lemma 2 :

Original: There exists K ( t ) , Q ( t ) > 0 , t = 0 , 1 , · · · , T − 1 such that for j ∈ S n ∗ and l / ∈ S ∗ , softmax

Original: [Equation]

Original: softmax

Original: [Equation] ) and for j / ∈ S n ∗ , l / ∈ S n ∗ , we have softmax ( x nj ⊤ W ( t +1) K [Equation] [Equation]
Revised: We prove the whole lemma by a long induction.

Original: For i = 1 , 2 ,

Original: [Equation]

Original: Claim 3.

Original: For the update of W ( t ) V , there exists λ ≤ Θ(1) such that [Equation] Proof of Lemma 2 : We prove the whole lemma by induction.
Revised: To make it easier to follow, we break this Section into three parts to introduce the proof of three claims of Lemma 2 separately.

Original: Proof of Claim of Lemma 2 :
Revised: D. P ROOF OF C LAIM 1 OF L EMMA

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: [Equation] , l ∈ [ M ] , where the second step comes from (53).
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),

Original: [Equation]
Revised: [Equation]

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

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

Original: and we can then obtain
Revised: and we can then obtain

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: [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 [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]
Revised: [Equation]

Original: Fori ∈ W ( t ) , from the derivation in (132) later, we have
Revised: For

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

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

Original: Hence, combining (90), (92), (93), (94)
Revised: Hence, combining (90), (92), (93), (94)

Original: [Equation]
Revised: [Equation]

Original: and (95), we can obtain(cid:28) [Equation] aB
Revised: and (95), we can obtain

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

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: by verifying that
Revised: by verifying that

Original: [Equation]
Revised: [Equation]

Original: When ηt < am , we have
Revised: When ηt < am , 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 ηt ≤ O (

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: [Equation] 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

Original: [Equation]
Revised: [Equation]

Original: 108) with probability at least 1 − ( mB ) − c .
Revised: [Equation] with probability at least 1 − ( mB ) − c .

Original: (107) comes from (60).
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 [Equation] [Equation] to ensure a non-zero gradient.
Revised: (110) where the last step is by ηt ≤ O (1) (111) to ensure a non-zero gradient.

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

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation] 113) (4) If i / ∈ (W ( t ) ∪ U ( t )) ,
Revised: (4) If i / ∈ (

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: Hence, combining (114), (115)
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: and (116), we can obtain

Original: [Equation]
Revised: [Equation]

Original: ≲ [Equation]
Revised: Hence, combining (114), (115), and (116), we can obtain

Original: [Equation]
Revised: [Equation]

Original:  andW
Revised: 117) and

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original:  (5) We finally study the bound of W ( t )
Revised: [Equation] (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 (78), 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
Revised: For the unlucky neuron i

Original: [Equation]
Revised: [Equation]

Original: we can similarly obtain
Revised: 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: Proof of Claim 2 of Lemma 2 :
Revised: D. P ROOF OF C LAIM 2 OF L EMMA

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 (227), we have

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: For r / ∈ S n 1 and l ∈ S n 1 , we have softmax 
Revised: For r / ∈ S n 1 and l ∈ S n 1 , we have softmax [Equation]

Original: [Equation]
Revised: [Equation]

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

Revised: [Equation]

Revised: − (cid:88)

Revised: [Equation]

Revised: n

Original: [Equation]
Revised: [Equation]

Original: 128) where 
Revised: (128) where β n 1 ( t ) ≳ [Equation] [Equation] (129) β n

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: l = ζ ′ l r
Revised: ≈ Θ( [Equation] [Equation] nj ( t

Original: [Equation]
Revised: [Equation]

Revised: +

Revised: [Equation]

Original:  t ) (130)
Revised: ( t ) [Equation] ( t [Equation]

Revised: +

Revised: [Equation]

Revised: (cid:88)

Revised: [Equation]

Revised: (cid:88) r ∈S nl

Revised: [Equation]

Revised: =

Revised: [Equation]

Revised: ζ ′

Original: [Equation]
Revised: [Equation]

Revised: r

Original: [Equation]
Revised: [Equation]

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

Revised: [Equation]

Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: Therefore, for i ∈ W ( t ) ,
Revised: Therefore, for i ∈ W ( t ) ,

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Revised: [Equation]

Revised: [Table]

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

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

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: >
Revised: >

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

Original: where the first step is by (53) and the second step is a combination of (58) to (61).
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: 133) and

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 (58) to (61),
Revised: If s ∈ S n 1 , by basic computation given (32) to (35),

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where the last step is by (133) and (134).
Revised: where the last step is by (133) and (134).

Original: If s ∈ S n 2 , from (62) to (65)
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] [Equation]
Revised: [Equation] [Equation]

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: 137) by (66) to (68).
Revised: 137) by (40) to (42).

Original: Hence, for i ∈ W ( t ) , j ∈ S g 1 , combining (129) and (135), we have
Revised: Hence, for i ∈ W ( t ) , j ∈ S g 1 , combining (129) and (135), we have

Original: [Equation]
Revised: [Equation]

Original: For i ∈ U ( t ) and l ∈ S n 1 , j ∈ S g
Revised: For i ∈ U ( t ) and l ∈ S n 1 , j ∈ S g

Original: [Equation]
Revised: [Equation]

Original: For i / ∈ ( W ( t ) ∪ U ( t )) and l ∈ S n 1 , j ∈ S g 1 ,
Revised: For i / ∈ ( W ( t ) ∪ U ( t )) and l ∈ S n 1 , j ∈ S g 1 ,

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: 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: 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 ,

Original: [Equation]
Revised: [Equation]

Original: ∈S
Revised: ∈S

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

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

Original: U ( t ) ,
Revised: U ( t ) ,

Original: [Equation]
Revised: [Equation]

Original: ≲
Revised: ≲

Original: Bt
Revised: Bt

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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 ∈B b

Original: [Equation]
Revised: [Equation]

Original: S n
Revised: S n

Original: [Equation]
Revised: [Equation]

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

Original: [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: 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.

Original: We also have
Revised: We also have

Original: K ( t )
Revised: K ( t )

Original: [Equation]
Revised: [Equation]

Original: Bt
Revised: Bt

Original: [Equation]
Revised: [Equation]

Original: Bt
Revised: Bt

Original: [Equation]
Revised: [Equation]

Original: Bt
Revised: Bt

Original: [Equation]
Revised: [Equation]

Original: aM
Revised: aM

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [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 (149)

Original: [Equation]
Revised: [Equation]

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

Original: Note that when
Revised: Note that when

Original: [Equation]
Revised: [Equation]

Original: ≤ O (1) , (150) we have
Revised: ≤ O (1) , (150) we have

Original: [Equation]
Revised: [Equation]

Original: When |S n | > |S n 1 | , min n ∈ [ N ] { ϕ n ( t )( |S n | − |S n 1 | ) } ≳ max n ∈
Revised: When |S n | > |S n 1 | , min n ∈ [ N ] { ϕ n ( t )( |S n | − |S n 1 | ) } ≳ max n ∈

Original: [ N ] { ϕ n ( t ) |S n 2 | ϕ n ( t
Revised: [ N ] { ϕ n ( t ) |S n 2 | ϕ n ( t

Original: [Equation]
Revised: [Equation]

Original: ( |S n | − |S n 1 | ) } (152)
Revised: ( |S n | − |S n 1 | ) } (152)

Original: Since that c 1 e
Revised: Since that c 1 e

Original: [Equation]
Revised: [Equation]

Original: (153) for 0 < c 1 ≤ min n ∈ [ N ] {
Revised: (153) for 0 < c 1 ≤ min n ∈ [ N ] {

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

Original: Moreover, |S
Revised: Moreover, |S

Original: [Equation]
Revised: [Equation]

Original: (155) for 0 <
Revised: (155) for 0 <

Original: [Equation]
Revised: [Equation]

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

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

Original: i.e.,
Revised: i.e.,

Original: [Equation]
Revised: [Equation]

Original: 158) by Hoeffding’s inequality in (26) given that |S n 1 | |S n | = Θ(1) .
Revised: 158) by Hoeffding’s inequality in (26) given that |S n 1 | |S n | = Θ(1) .

Original: Thus,
Revised: Thus,

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: if
Revised: if

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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 (149), 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 (147, 148) above and (173, 174) below for simplicity.

Original: Then we give a brief derivation of W ( t +1) Q x n j for j / ∈ S n 1 in the following.
Revised: Then we give a brief derivation of W ( t +1) Q x n j 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 ) ,

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

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] [Equation]
Revised: [Equation] [Equation]

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

Original: W
Revised: W

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: Loss b

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: 171) where
Revised: 171) where

Original: [Equation]
Revised: [Equation]

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

Original: [Equation] [Equation]
Revised: [Equation] [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
Revised: t

Original: [Equation]
Revised: [Equation]

Original: = γ
Revised: = γ

Original: [Equation]
Revised: [Equation]

Original: t ) , (175)
Revised: t ) , (175)

Original: For j ∈ S nl , l = 3 , 4 , · · · , M , we have
Revised: For j ∈ S nl , l = 3 , 4 , · · · , M , we have

Original: W
Revised: W

Original: [Equation]
Revised: [Equation]

Original: t
Revised: t

Original: [Equation]
Revised: [Equation]

Original: 1) K x
Revised: 1) K x

Original: [Equation]
Revised: [Equation]

Original: 176) where | Q f ( t )
Revised: 176) where | Q f ( t )

Original: |
Revised: |

Original: [Equation]
Revised: [Equation]

Original: Q ( t ) (177)
Revised: Q ( t ) (177)

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: 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 ))

Original: ∥
Revised: ∥

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

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

Original: + K e ( t ) Q e ( t ) ∥
Revised: + K e ( t ) Q e ( t ) ∥

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

Original: +
Revised: +

Original: M (cid:88) l = γ l γ ′ l ∥
Revised: M (cid:88) l = γ l γ ′ l ∥

Original: q l ( t ) ∥∥
Revised: q l ( t ) ∥∥

Original: r l ( t ) ∥ ≳
Revised: r l ( t ) ∥ ≳

Original: (1 + K ( t ))(1 + Q ( t ))
Revised: (1 + K ( t ))(1 + Q ( t ))

Original: ∥
Revised: ∥

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

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

Original: ∥ − (cid:118)(cid:117)(cid:117)(cid:116)
Revised: ∥ − (cid:118)(cid:117)(cid:117)(cid:116)

Original: M (cid:88) l = (
Revised: M (cid:88) l = (

Original: B (cid:88) n ∈B b
Revised: B (cid:88) n ∈B b

Original: Q ( t ) |S nl | |S n | − |S n ∗ | )
Revised: Q ( t ) |S nl | |S n | − |S n ∗ | )

Original: 2 ∥ r l ( t ) ∥ 2 · (cid:118)(cid:117)(cid:117)(cid:116)
Revised: 2 ∥ r l ( t ) ∥ 2 · (cid:118)(cid:117)(cid:117)(cid:116)

Original: M (cid:88) l = (
Revised: M (cid:88) l = (

Original: B (cid:88) n ∈B b
Revised: B (cid:88) n ∈B b

Original: K ( t ) |S nl | |S n | − |S n ∗ | ) 

Revised: K ( t ) |S nl | |S n | − |S n ∗ | )

Original: ∥ q l ( t ) ∥≳ (1 + K ( t ) + Q ( t ))
Revised: ∥ q l ( t ) ∥≳ (1 + K ( t ) + Q ( t ))

Original: ∥
Revised: ∥

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

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

Original: If
Revised: If

Original: [Equation]
Revised: [Equation]

Original: Hence, for j, l ∈ S n 1 ,
Revised: Hence, for j, l ∈ 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: 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]
Revised: [Equation]

Original: [Equation] [Equation][Equation]
Revised: [Equation] [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: Proof of Claim 3 of Lemma 2 :
Revised: D. P ROOF OF C LAIM 3 OF L EMMA

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
Revised: Consider a data { X n

Original: [Equation]
Revised: [Equation]

Original: y n } where y n = 1 .
Revised: y n } where y n = 1 .

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

Original: [Equation]
Revised: [Equation]

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

Original: B (cid:88) n ∈B b
Revised: B (cid:88) n ∈B b

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

Revised: B (cid:88)

Original: [Equation]
Revised: [Equation]

Original:  i ∈W( t ) V
Revised: = (cid:88) i ∈W

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[Equation]
Revised: 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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original:  if i is an unlucky neuron, by Hoeffding’s inequality in (26), we have
Revised: 194) 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]

Revised: j ( t ) (201)

Original: Here ∥ zj
Revised: Here ∥ z

Original: [Equation]
Revised: [Equation]

Original: t)
Revised: t

Original: [Equation]
Revised: [Equation]

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

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

Revised: E 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

Revised: [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 weight and p l for one i ∈ [ m ] and all l ∈ [ M ] .
Revised: Let θ l be the angle between the initial weight and p l 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 } Ml =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 (211).

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

Original: [Equation]
Revised: [Equation]

Original: 217) and
Revised: 217) and

Original: [Equation]
Revised: [Equation]

Original: Combining (217) and (218), we can approximately compute that if
Revised: Combining (217) and (218), 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: D E XTENSION TO MULTI - CLASSIFICATION
Revised: F E XTENSION TO MORE GENERAL CASES F. 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

Revised: [Equation]

Revised: [Equation]

Revised: (cid:88)

Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The dataset D can be divided into four groups as
Revised: The dataset D can be divided into four groups as

Original: [Equation]
Revised: [Equation]

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

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 [Equation] for W O u in the network (223)
Revised: Therefore, for W O u in the network (223)

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: 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 [Equation] if the data ( X n , y n ) ∈ D 2 , we have[Equation]
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: n [Equation] [Equation] if ( X n , y n ) ∈ D 3 , we have− ∂
Revised: if ( X n , y n ) ∈ D 3 , we have

Original: [Equation]
Revised: [Equation]

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

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.

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

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

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 .

Revised: M 1 ,

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

Revised: [Equation]

Revised: [Equation]

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

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

Revised: [Equation]

Revised: and τ ≤ κ .

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 .

Revised: Assumption 1 can be modified such that

Revised: [Equation]

Revised: [Equation]

Revised: where p j,b ⊥ p 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 ] .

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

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.

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

Revised: 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: The coefficient of the linear combination is a function of fractions of different pattern vectors µ l,b in M l .

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

Revised: F. E XTENSION TO MULTI - HEAD NETWORKS

Revised: Suppose there are H heads in total.

Revised: The network is modified to

Revised: [Equation]

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

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

Revised: R m ×

Revised: Hm a , and W O

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

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

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 ∈

Revised: [ 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 ∈

Revised: [ H ] .

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

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

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

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.

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

Revised: F. E XTENSION TO SKIP CONNECTIONS AND NORMALIZATION

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

Revised: Let m a = d .

Revised: The network is changed into

Revised: [Equation]

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

Revised: [Equation]

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

Revised: 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: The major differences come from the gradient of W O , which also helps to determine the Relu gate.

Revised: One needs to redefine

Revised: [Equation]

Revised: for l ∈ S n 1 .

Revised: 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)

Revised: [Equation]

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

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

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

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

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