A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity
Authors: Hongkang Li
Meng Wang
Sijia Liu
Pin-Yu Chen
Venue: ICLR 2023 poster (2023)
Keywords: Vision transformer
Learning
Generalization
Sample comeplxity
Token sparsification
Theory

Abstract: Vision Transformers (ViTs) with self-attention modules have recently achieved great empirical success in many vision tasks. Due to non-convex interactions across layers, however, the theoretical learning and generalization analysis is mostly elusive. 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. We characterize the sample complexity to achieve a zero generalization error. 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. 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. 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. Empirical experiments on synthetic data and CIFAR-10 dataset justify our theoretical results and generalize to deeper ViTs. 

Published as a conference paper at ICLR 2023
A THEORETICAL UNDERSTANDING OF SHALLOW VI- SION TRANSFORMERS: LEARNING, GENERALIZA- TION, AND SAMPLE COMPLEXITY
Hongkang Li
Meng Wang
Sijia Liu
Pin-Yu Chen
ABSTRACT
Vision Transformers (ViTs) with self-attention modules have recently achieved great empirical success in many vision tasks. Due to non-convex interactions across layers, however, the theoretical learning and generalization analysis is mostly elusive. Based on a data model characterizing both label-relevant and label-irrelvvant 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 classification task. We characterize the sample complexity to achieve zero generalization error. Our sample complexity bound is positively correlated with the inverse the fraction of label-relevant tokens, the token noise level, and the initial model error. We also prove that training process using stochastic gradient descent (SGD) leads to a sparse attention map, which is a formal verification the general intuition about the success of attention. Moreover, this paper indicates that proper token sparsification can improve the test performance by removing label-irrelevant and/or noisy tokens, including spurious correlations. Empirical experiments on synthetic data and CIFAR-10 dataset justify our theoretical results and generalize to deeper ViTs.
INTRODUCTION
As the backbone of Transformers (Vaswani et al. 2017). the self-attention mechanism (Bahdanau 2014) computes the feature representation by globally modeling long-range interactions within the input. Transformers have demonstrated tremendous empirical success in numerous areas, in- cluding nature language processing (Kenton & Toutanova 2019 Radford et al. 2019 2018 Brown jet al. 2020). recommendation system (Zhou et 2018 Chen et al. 2019 Sun et al. 2011)) ) and
reinforcement learning (Chen al. Janne]aanner al. [202] 1:21:0 Zheng al. (022). Starting from the models (Touvron al. 2021 Jiang al. 2021 Wang et al. 2021| Liu et al.] 2021a) gradually replace convolutional neura network (CNN) architectures and become prevalent in vision tasks. Various techniques have been developed to train ViT efficiently. Among them, token sparsification (Pan al.) 202T Rao
advent of Vision Transformer (ViT) (Dosovitskiy et al. 2020). Transformer-
al. 2021 Liang et al. 2022 Tang et al. 20222 Yin et al. (0222) removes redundant tokens (image patches) of data to improve the computational complexity while maintaining comparable learning performance. For example, Liang et al. (2022); Tang al. 2022) prune tokens following criteria designed based on the magnitude of the attention map. Despite the remarkable empirical success, one fundamental question about training Transformers is still vastly open, which is
Under what conditions does a Transformer achieve satisfactory generalization?
Some recent works analyze Transformers theoretically from the perspective of proved Lipschitz constant of self-attention (James Vuckovic, 2020 Kim al.) (021). properties of the neural tangent
kernel (Hron et al. 2020 Yang 2020) and expressive power and
(Dehghani
ct al. 2018 Yun et al. 2019 Bhattamishra et al. (2020a b Edelman et al. 2022 Dong et al. 2021
"Rnnssllarr Polytechnic Institute. Email: Lih350rpi.c.
Rensselaer Polytechnic Institute. Email: wangm wwmmierrpiieee
Michigan State University & IBM Research. Email: Liusijii@mass jimmsuu.edds
TBM Research. Email: Pin-Yu Yu. Chen@ bm, com

---

Published as a conference paper at ICLR 2023
Likhosherstov et al. 2021: Cordonnier et al., 2019 Levine al. 2020) with statistical guarantees Snell et al, 2021: Wei et al. 2021). Likhosherstov et al. (2021) showed model complexity for the function approximation of the self-attention module. [Cordonnier al. (2019) provided sufficient and necessary conditions for multi-head self-attention structures to simulate convolution layers. None of these works, however, characterize the generalization performance of the learned model theoreti- cally. Only [delman et al. (2022) theoretically proved that a single self-attention head can represent a sparse function of the input with a sample complexity for 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.
To the best of our knowledge, this paper provides the first learning and generaliza- tion analysis of training a basic shallow Vision Transformer using stochastic gradient descent (SGD). This paper focuses on binary classification problem on structured data, where tokens with dis- criminative patterns determine the label from a majority vote, while tokens with non-discrimenative patterns do not affect the labels. We train a VIT containing self-attention layer followed by two-layer perceptron using SGD from a proper initial model. This paper explicitly characterizes the required number of training samples to achieve desirable generalization performance, referred to as the sample complexity. 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-irrelevan patterns and less noise from a better initial model. The highlights of our technical contributions include:
Contributions:
First, this paper proposes a new analytical framework to tackle the non-convex optimization and generalization for shallow ViTs. Due the more involved non-convex interactions of learn- ing parameters and diverse activation functions across layers, the ViT model, i.e., three-layer neural network with one self-attention layer, considered in this paper is more complicated to analyze than three-layer CNNs considered in AlleenZZuu et al.] (2019a); : Allen-Zhu & Li (2019), the most compli- cated neural network model that has been analyzed so far for across-layer nonconvex interactions. We consider a structured data model with relaxed assumptions from existing models and establish new analytical framework to overcome the new technical challenges to handle ViTs.
Second, this paper theoretically depicts the evolution of the attention map during the training and characterizes how "attention" paid to different tokens during the training. Specifically, 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. This insight provides a theoretical justification of the magnitude-basee token pruning methods such as (Liang et al. 2022 Tang et al. 2022) for efficient learning.
Third, we provide theoretical explanation for the improved generalization using token spar- sification. We quantitatively show that a token sparsification method can remove glass-irrelevar and/or highly noisy tokens, then the sample complexity is reduced while achieving the same test- ing accuracy. Moreover, token sparsification can also remove spurious correlations to improve the testing accuracy (Likhomanenko et al. 2021] Zhu al., [2021a). This insight provides a guideline in designing token sparsification and few-shot learning methods for Transformer (He et al. 2022
Guibas al. 2022).
1.1 BACKGROUND AND RELATED WORK
Efficient VIT learning. 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. Zhu al. (202Tb) identifies the importance of different dimensions in each layer of ViT' and then executes model pruning. [iu et al. b)) Lin et al. (2022) ; Li et al. (2022d) quantize weights and inputs to compress the learning model. Li et al. 2022a) studies automated progressive learning that automatically increases the model capacity on- the-fly. Moreover, modifications of attention modules, such as the network architecture based on local attention (Wang et al. 2021 Liu et al. 202 Chu et al. 2021), can simplify the computation of global attention for acceleration.
Theoretical analysis of learning and generalization of neural networks. One line of research
(Zhong et al.] [0177]] Fu et al.] [2200] Zhong et al., 2017a] Zhang et al.) [020a]]] Li et al.] (022c) analyzes the generalization performance when the number of neurons is smaller than the number
2

---

Published as a conference paper at ICLR 2023
of training samples. The tanaelltttgeet-eeeeene (NTK) analysis (Jacot et al. 2018 Allen- al. 2019a bj Arora et al. 2019; Cao & Gu 2019 Zou & Gu 2019; Du et al.) 2019 Chen et al., 2020 Li
2022b) considers strongly overparameterized networks and eliminates the nonconvex inter- actions across layers by linearizing the neural network around the initialization. The generalization performance is independent of the feature distribution and cannot explain the advantages of self-
attention modules.
Neural network learning on structured data. Li & Liang (2018) provide the generalization anal- ysis of a fully-connected neural network when the data comes from separated distributions. Daniely & Malach (2020); Shi et al. (202I); Karp et al. (202T); : Brutzkus & Globerson (202T); Zhang et al. (2023) study fully connected networks and convolutional neural networks assuming that data con- tains discriminative patterns and background patterns. Allen-Zhu & Li (2022) illustrates the ro- bustness of adversarial training by introducing the feature purification mechanism, in which neural networks with non-linear activation functions can memorize the data-dependent features. Wenn Li (021 extends this framework to the area of self-supervised contrastive learning. All these works consider one-hidden-laye neural networks without self-attention.
Vectors are in bold lowercase, and matrices and tensors are in bold uppercase. Scalars are in normal fonts. Sets are in calligraphy font. For instance, Z is a matrix, and x is a vector.
Notations:
2 denotes the i-th entry of x, and Z,,, denotes the (i.j)-th entry of Z. [K] (K 0) denotes the set including integers from to K. We follow the convention that f(x) = O(g())) (or ((g(z)), (g(())) means that f(x) increases at most, at least, or in the order of g(r), respectively.
2 PROBLEM FORMULATION AND LEARNING ALGORITHM
We study a binary classification problem following the common setup in (Dosovisskiy et al. 2020 Touvron et al., 2021: Jiang et al. 2021). Given N training samples {(X", [ generated from an unknown distribution D and fair initial model, the goal is to find an improved model that maps X to y for any (X,y) ~ D. Here each data point contains tokens 2f, 2 ,.,,.., ,T. i.e.
X [ .... ,. E RddL where each token dd-dimensional and unit-norm. y² E (+1,-1] is a scalar. A token can be an image patch (Dosovitskiy et al. (020). We consider a general setup that also applics to token sparsification, where some tokens are set to zero to reduce the computational time. Let 8 C [ denote the set of indices of remaining tokens in X after sparsification. Then |S" S L. and 8 = [L] without token sparsification.
Learning is performed over basic three-layer Vision Transformer, a neural network with a single- head self-attention layer and a two-layer fully connected network, as shown in ( This is a simpli- fied model of practical Vision Transformers (Dosovitskiy et al. (2220) to avoid unnecessary compli- cations in analyzing the most critical component of ViT', the self-attention.
F(X")
Σ LES
T Relu Wo Wy X softmax X T W Woz?)), any
(1)
SD
where the queue weights We in RT the key weights WK in Rd and the value weights Wv in Rm xd in the attention unit are multiplied with X to obtain the queue vector WoX", the key vector Wx K X", and the value vector X" respectively (Vaswani et al. 2017p. Wo is in Rm and A = (a(1), a(2).... ... aL) where am E Rmm Le [L] are the hidden-layer and output-layer weights of the two-layer perceptron, respectively. m is the number of neurons in the hidden layer. Relu Rm + T where Relu(æ)) = max(ar, 0}. softmax RL + R&Z where
softmax(x)) (e _........ ........................... "/",,,""/// Et e Let v (A, Wo, Wv, WK. K, We) denote the set of parameters to train. The training problem minimizes the empirical risk (x(1).
min &
fN(1)
1 N
Σ
{(X" v")))
(2)
where (("",""0)) is the Hinge loss function, i.e.,
((X", v²: - max(1 y² F(.
(3)
Exteension multi-classification is briefly discussed in Section G.1
3

---

Published as a conference paper at ICLR 2023
The generalization performance of learned model v is evaluated by the population risk f(e), where
f(t) f(A Wo. Wr. WK. Wo) B(x (x) F(X),
(4)
The training problem 22 is solved via mini-batch stochastic gradient descent (SGD), sum- At iteration I, 1 = 0,1,2,-.. T 1, the gradient is computed using a
marized in Algorithm
mini-batch B, with [3 B. The step size is 1/-
Similar to Dosovitskiy et al. 2020 Touvron et al. 2021 Jiang et al. 2021). W. V / and W
K (0) come from an initial model. Every entry of Wo is generated from (0, <²). Every entry of aj is sampled from {+ vm
vm } with equal probability. A does not update during the training
THEORETICAL RESULTS
3.1 MAIN THEORETICAL INSIGHTS
Before formally introducing our data model and main theory, we first summarize the major insights. We consider a data model where tokens are noisy versions of label-relevant patterns that determine the data label and label-irrelevant patterss that do not affect the label. is the fraction of label- relevant tokens. o represents the initial model error, and t characterizes the token noise level.
(PI). A Convergence and sample complexity analysis of SGD to achieve zero generalization error. We prove SGD with a proper initialization converges to a model with zero generalization error. The required number of iterations is proportional to 1/a., and 1/(
T). Our sample complexity bound is linear in o_² and (0(1) - a -)= 2. Therefore. the learning performance is improved, in the sense of 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.
(P2). A theoretical characterization of increased sparsity of the self-attention module dur- ing training. 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 con- centrated at label-relevant tokens. This formally justifies the general intuition that the attention layer makes the neural network focus on the most important part of data.
(P3). A theoretical guideline designing token sparsification methods to reduce sample com- plexity. Our sample complexity bound indicates that the required number of samples to achieve zero generalization can be reduced a token sparsification method removes some label-irrelevant tokens (reducing a.),, or tokens with large noise (reducing o), or both. This insight provides guideline to design proper token sparsification methods.
(P4). A new theoretical framework to analyze the nonconvex interactions in three-layer ViT's. This paper develops a new framework to analyze ViTs based on a more general data model than existing works like (Brutzkus & Globerson 202T Karp et al. 202 Wen & Li 202T). 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. 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 tokens. This theoretical framework is of independent interest and
only rather than can potentially applied to analyze different variants of Transformers and attention mechanisms.
3.2 DATA MODEL
There are M (2 < M < ma, my) distinct patterns {1L #2. pM} in Rd, where μ1, P2 are discriminative patterns that determine the binary labels, and the remaining M 2 pat- terns P3. p4, MM are non-discriminative patterns that do not affect the labels. Let K =
I is common to fix the output layer weights as the random initialization in the theoretical analysis of neural networks, including NTK (Allen Zhu et al. [201a]] Arora et al. 2019), model recovery (Zhong et al. 2017b). and feature learning (Karp et al. [202]] Allen-Zhu & (022) type of approaches. The optimization problem here of Wa. WK. Wv, and Wo with non-linear activations is still highly non-convex and challenging.
4

---

Published as a conference paper at ICLR 2023
min l|u ply 0 denote the minimum distance between patterns. of X is noisy version of one of the patterns, i.e.,
Each token zi
min
| 1ll S T,
(5)
SE[M]]
and the noise level 7 < k/4. We take K. 4T (1) for the simplicity of presentation.
The label y" is determined by the tokens that correspond discriminative patterns through ma- jority vote. If the number of tokens that are noisy versions of μl is larger than the number of tokens that correspond to μ12 in X", then yy" 1. In this case that the label y" 1, the tokens that are noisy ft are refereed to label-relevant tokens, and the tokens that are noisy l are referred to as confusion tokens. Similarly, there are more tokens that are noisy μl than those that are noisy P1. the former are label-relevant tokens, the latter are confusion tokens, and y" - All other tokens that are not label-relevant are called label-irrelevan tokens.
Let a and Ox the average fraction of the label-relevant and the confusion tokens over the distri- bution D, respectively. We consider balanced dataset. Let D+ = {(X", y") = +1, E [N] and D_ {Xx ly" 11, [N]] denote the sets of positive and negative labels, respec- tively. Then D D O(vN).
Our model is motivated by and generalized from those used the state-of-art analysis of neural networks structured data (Li & Liang 2018 Brutzkus & Globerson 2021 Karp et al. (021)) All the existing models require that only one discriminative pattern exists in each sample, i.c., either μ4 or u2g but not both, while model allows both patterns to appear in the same sample.
3.3 FORMAL THEORETICAL. RESULTS
Before presenting main theory below, first characterize the behavior of the initial model through Assumption T Some important notations are summarized in Table
Table 1: Some important notations
Initialization error for value vectors
Initialization error for query and key vectors
Minimum μa - p4y I| any j E [M], #
M
Total number of patterns
m
The number of neurons Wo
0+
Average fraction of label-relevant tokens
Ox#f
Average fraction of confusion tokens
V exist three (not necessarily different) sets of orthonormal bases P {PI,P2,- P2. .MM) J. Q = {qi, ... .Mr), and R = {ri, r2, ... ,TM,, where pt E Rma, qi E Rmm Vl E [M].
Assumption 1. Assume max(|) W
(0) IW
1 without loss of generality. There
q1 = Tı, and 92 = r such that
|| (0) pj P/ll 1[ (0)
Hy glll S 8, and WWo H, r, 5 8.
(6)
for some a = (1/M)) and 8 < 1/2.
Assumption characterizes the distance of query, key, and value vectors of patterns 1 to orthonormal vectors. The requirement on is minor because 5 can be in the same order as ||||||
Theorem (Generalization of ViT), Suppose Assumption 7 holds; T & min(a, d); sufficiently large model with
ス E 2M² log N for > 0,
the average fraction of label-relevant patterns satisfies
a#
(8)
ESE (6+T)(1 (o
Thhe condition Zi and 92 is to eliminate the trivial case that the initial attention value is very small. This condition can be relaxed but keep this form simplify the representation.
5

---

Published as a conference paper at ICLR 2023
for some constant cs E (0, 2) and the mini-bated size and the number of sampled tokens of each data X", n E [M] satisfy
B 2 Q((ar, - e ( 0)) ?)
|S"| 2 2((1)
(9)
Then as long as the number of training samples N satisfies
1
N 1 2
(10)
(a. - cc(cc () di T))²
for some constant d, d''> O, and Ç 2 1 - 11, after number of iterations such that
10
T = 0 (1
(11)
(a+T)M yna,
with probability least 0.99, the returned model achieves zero generalization error as f(A (0) W ( (
(O)) ) =0
(12)
Theorem [characterizes under what condition of the data the neural network with self-attention in ①1 trained with Algorithm can achieve zero generalization error. 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 self- attention layer in the following proposition.
Proposition (Generalization without self-attention). Suppose assumptions in Theorem ! hold. When there is no self-attention layer, i.e., WK and Wq are not updated during the training, if N satisfies
1
N
(a, (a. (a.aaaa-7)) --T)
(13)
then after T iterations with T in 411. the returned model achieves zero generalization error as f(A (0) W w w ( =
(14)
Remark 1. (Advantage of the self-attention layer) Because m > may mi, d. the number of train- able parameter remains almost the same with without updating the attention layer. Combining Theorem I and Proposition [ we can see that with the additional self-attention layer, the sample complexity reduced by a factor 1/a² with an approximately equal number of network parameters. Remark 2. (Generalization improvement by token sparsification). (10) and (77) show that the sam- ple complexity N and the required number of iterations T scale with 1/a², and 1/a.,, respectively. Then, increasing a.., the fraction of label-relevant tokens, can reduce the sample complexity and speed up the convergence. Similarly, N and T scale with 1/(@(1) -)2 and 1/(O(1) r). Then decreasing т, the noise in the tokens, can also improve the generulization. Note that properly designed token sparsification method can both increase a. by removing label-irrelevant tokens and decrease by removing noisy tokens, thus improving the generalization performance.
Remark 3. (Impact of the initial model) The initial model W w(O) W(O) affects the learning performance through and 5, both of which decrease as the initial model is improved. Then from To and TTy. the sample complexity reduces and the convergence speeds up for a better initial model.
Proposition&ssaen that the attention weights are increasingly concentrated label-relevant tokens during the training. Proposition iis critical component in proving Theorem I and is of independent interest.
Proposition 2. The attention weights for each token become increasingly concentrated on those correlated with tokens of the label-relevant pattern during the training, i.e.,
exp(xx T w(O)
w Q
softmax(X x
w K
(1
x
Σ.
+ +1- (15)
REST Ereso exp(27 (
w
REST
K
at sublinear rate of O1/(t) when I. is large for a large C > 0 and all I E S and n E [N].
"The sample complexity bounds in 10) and 13) are sufficient but not necessary. Thus, rigorously speak- ing, one can not compare two cases based on sufficient conditions only. In our analysis, however, these two bounds are derived with exactly the same technique with the only difference in handling the self-attention layer. Therefore, we believe is fair compare these hounds show the advantage of VIT.

---

Published as a conference paper at ICLR 2023
Proposition 2 indicates that only label-relevant tokens highlighted by the learned attention of ViTs, while other tokens have less weight. This provides a theoretical justification of magnitude- based token sparsification methods. softmaax-); in (15) denotes the i-th entry of softmax()).
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 non- discriminative patterns and the initial model error almost unchanged. To be more specific, by Lemma о we can identify two groups of neurons in the hidden layer Wo. where one group only learns the positive pattern, and the other group only learns the negative pattern. Claim | of Lemma 2 states that during the SGD updates, the neuron weights in these two groups evolve in the direction of projected discriminative patterns, P1 and P2. respectively. Meanwhile, Claim 2 of Lemma 2 indicates that WK and Wo update in the direction of increasing the magnitude of the query and key vectors of label-relevant tokens from 1 to (log 7'), such that the attention weights correlated with label-relevant tokens gradually become dominant. Moreover, by Claim of Lemma the update of Wy increases the magnitude of the value vectors of label-relevant tokens, by adding partial neuron weights of Wo that are aligned with the value vectors to these vectors. Due to the above properties during the training, one can simplify the training process to show that the output of neural network (T) changes linearly in the iteration number t. From the above analysis, we can develop the sample complexity and the required number of iterations for the zero generalization guarantee.
Our proof technique is inspired by the feature learning technique in analyzing fully connect networks and convolution neural networks (Shi et al.] [021 Brutzkus & Globerson
Technical novelty:
2021]. Our paper makes new technical contributions from the following aspects, First, we provide a
new framework of studying the nonconvex interactions of multiple weight matrices in a three-layer
ViT while other feature learning works (Shi et al.] [02T]] BButzkus & Globerson [02T] Karp et al.
2021 Allen-Zhr & Li 2022; Wen & Li, 2021 Zhang et al. 2023) only study one trainable weight matnx in the hidden layer of a two-layer network. 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. (202)) or oversimplify the analysis by applying the neural- suuralltangnttkeende (NTK) method that considers impractical over- wwer-prameterization and updates the weights only around initialization. (Hron et al. 22020; Yang 2020 Allen-Zhu al.)
2019a; AArora et al. 2019). 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 & Karp et al. (02T) require one discriminative pattern only in each sample.
Globerson 2021
4 NUMERICAL EXPERIMENTS
4.1 EXPERIMENTS ON SYNTHETICC DATASETS
We first verify the theoretical bounds in Theorem I on synthetic data. We set the dimension of data
and attention embeddings to be d ma my 10. Let co = 0.01. Let the total number of
patterns M- 5, and {M1 .4...... MMM} } be set of orthonormal bases. To satisfy Assumption I we generate every token that is a noisy version of μu from a Gaussian distribution N NH,, oB 11) with
the mean μe and covariance ciI, where I E Rd is the identity matrix. w(0)) w 1 w = 0 /ej, and each entry of w follows (0. g²), where U is m, m. orthonormal matrix, and g = 0.01. The number of neurons m of Wo is 1000. We set the ratio of different patterns the same among all the data for simplicity.
Sample complexity and convergence rate: We first study the impact of the fraction of the label- relevant patterns a, the sample complexity. Let the number of tokens after sparsification be |S 1000 the initialization error a = 0.1, and g = 0.2. The fraction of non-discriminative patterns is fixed to be 0.5. We implement 20 independent experiments with the same a, and N and record the Hinge loss values of the testing data. An experiment is successful if the testing loss is smaller than 10-3 3, Figure [ (a) shows the success rate of these experiments. A black block means that all the trials fail, A white block means that they all succeed. The sample complexity is indeed almost linear in a, as predicted in 10 We next explore the impact on a. Set a. = 0.3 and a# 0.2. The number of tokens after sparsification is fixed at 50 for all the data. Figure (b)) shows that 1/ 1/VN is linear in O(1) a, matching our theoretical prediction in (10). The result on
7

---

Published as a conference paper at ICLR 2023
the noise level is similar to Figure 1 (b), and we skip it here. In Figure 2 verify the number of iterations T against a_ in (11 where we set o = 0.1 and ò = 0.4.
Ad Advataae of self-attention: To verify Proposition [ compare the performance on ViT in and on the same network with Wx and We fixed during the training, i.c., a three-layer CNN. Compared with ViT, the number of trainable parameters in CNN is reduced by only 1%. Figure 3 shows the sample complexity of CNN is almost linear in o_ as predicted in 13). Compared with Figure (a), the sample complexity significantly increases for small Q., indicating a much worse generalization of CNN.
400
300
-15
40
0.8
02
0.4
0.5
0.5
3.5
1.5
(a)
(b)
Figure 2: The number of iterations against a,""
Figure 1: The impact of a, and a on sample complexity.
Attention map: We then evaluate the evolution of the attention map during the training. Let S 50 for all n E [N]. The number of training samples N 200. a 0.1, 8 0.2, az = 0.5, O# 0.05. In Figure 4 the red line with asterisks shows that the sum of attention weights on label-relevant tokens, i.c., the left side of (15) averaged over all 4, indeed increases to be close to 1 when the number of iterations increases. Correspondingly, the sum of attention weights on other tokens decreases to be close to 0, as shown in the blue line with squares. This verifies Lemma on sparse attention map.
Sum of stbcation wesights
Random spancfication
0.8
Noiey tokuns
a.
-1 Lbcl-eelson fenture
Test
-8 Other faruss
0.05
Q2
The number removed tekenss
Figure 3: Comparison of VIT and CNN
Figure Concentration of at- Figure 5:
Impact of token
tention weights
sparsification testing loss
Token sparsification: We verify the improvement by token sparsífication in Figure 5 The exper- iment is duplicated 20 times. The number of training samples N 80. Let S 50 for all E [N] Set o = 0.1, 5 = 0.5, a, = 0.6, a = 0.05. If we apply random sampling over all tokens, the performance cannot be improved as shown in the red curve because a, and do not change. we remove either label-irrelevant tokens or tokens with significant noise, the testing loss decreases, as indicated in the blue and black curves. This justifies our insight P3 on token sparsification.
4.2 EXPERIMENTS ON IMAGE CLASSIFICATION DATASETS
Dataset: To characterize the effect of label-relevant and label-imelevant tokens on generalization, following the setup of image integration in (Karp et al. 202T), we adopt an image from CIFAR- 10 dataset (Krizhevsky background image from the IMAGENET Plants synset (Karp et al. 2021 Deng et al. 22009). which plays the role of label-irrelevant feature. Specifically, we randomly cut out region with size 26 26 in the IMAGENET image and replace it with a resized CIFAR-10 10 image.
al.) 2010) as the label-relevant image pattern and integrate it with noisy
Experiments implemented 2020), the network architecture contains blocks, where we have a 4-head self atlf-attention layer and one-layer perceptron with skip connections and Layer-Normalization in each block.
Architecture:
deep ViT model. Following (Dosovitskiy al.
8

---

Published as a conference paper at ICLR 2023
We first evaluate the impact generalization of token sparsification that removes label-irrelevar patterns to increase dx. We consider 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. The number of samples for training and testing is 50K and 10K, respectively. A pre- trained model from CIFAR-100 (Krizhevsky al. (010) is used the initial model with the output layer randomly initialized. Without token sparsification, the fraction of class-relevant tokens is 0+ 0.66. a+ 1 implies all background tokens are removed. Figure 6 (a) indicates that larger a, by removing more label-irrelevant tokens leads to higher test accuracy. Moreover, the test per- formance improves with more training samples. These are consistent with our sample complexity analysis in 10). Figure( presents the required sample complexity to learn model with desirable test accuracy. We run 10 independent experiments for each pair of a, and N, and the experiment considered a success if the learned model achieves a test accuracy of at least 77.5%.
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. We consider binary classification problem that difterentiates "bird" and "airplane" images. 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. The remaining training data are integrated into clean background by zero padding. Therefore, the label "bird" is spuriously correlated with the class-irrelevant plant background. The testing data contain 50% birds and 50% airplanes, and each class has 50% plant background and 50% clean background. The numbers of training and testing samples are 10K and 2K, respectively. We initialize the ViT using two pre-trained models. The first one is pre-trained with CIFAR- 100, which contains images of 100 classes not including birds and airplanes. The other initial model is trained with a modified CIFAR-10 with 500 images per class for total of eight classes, excluding birds and airplanes. The pre- pre-trained model on CIFAR- 100 is better initial model because it is trained on more diverse dataset with more samples.
In Figure 6 (c), the token sparsification method removes the tokens of the added background, and the corresponding a, increases. Note that removing background in the training dataset also reduces the spurious correlations between birds and plants. Figure 6 (c) shows that from both initial models, the testing accuracy increases when more background tokens are removed. Moreover, better initial model leads to a better testing performance. This is consistent with Remarks and
5000
Tou OCORIN
Testassrayy
Santhrr of
2000
- exaifiid CIFAR-I GELTIARR
0.9
0.56
0.78
Narber of cku
(a)
(c)
Figure 6: (a) Test accuracy when N and a. change. (b) Relationship of sample complexity against a.. (c) Test accuracy when token sparsification removes spurious correlations.
5 CONCLUSION
This paper provides a novel theoretical generalization analysis of three-layer ViTs. Focusing on data model with label-relevant and label-irreeevant tokens, this paper explicitly quantifies the sample complexity as function of the fraction of label-relevant tokens and the token noise projected by the initial model. proves that the learned attention map becomes increasingly sparse during the train- ing, where the attention weights are concentrated those of label-relevant tokens. Our theoretical results also offer a guideline on designing proper token sparsification methods to improve the test performance.
This paper considers a simplified but representative Transformer architecture to theoretically exam- ine the role of self-atted seff-attention layer as the first step. One future direction is to analyze more practical architectures such as those with skip connection, local attention layers, and Transformers other areas. We see ethical immediate negative societal consequence of our work.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This work was supported by AFOSR FA9550-20-1-0112,, ARO 191 NSF 1932196 and the Rensselaer-IBM AI Research Collaboration (bttp:/airccrpii.dd part of the IBM AT Hori- zons Network (http://ibm.iiz/ oitp//ommm We thank Dr. Shuai Zhang at Rensselaer Polytechnic Institute for useful discussions. We thank all anonymous reviewers for their constructive comments.
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Ad- vances Neural Information Processing Systems, 32, 2019.
Zeyuan Allen-Zhu Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium Foundations of Computer Science (FOCS), pp. 977-988. IEEE, 2022.
Zeyuan Allen-Zhu and Yuanzhi Li. Understanding ensemble, knowledge distillation and self- distillation in deep learning. In International Conference on Learning Representations, 2023. URL httpttt::p/ppeeeeetaaettee wet.ieeww..eeeeeeeeeeeeeeeeeeeee www.eeewwee//rrrrrrteeeeee wot.orrrridduuuuatrtXX
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameter- ized neural networks, going beyond two layers. Advances in neural information processing systems, pp. 6155-6166 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp. 242-252. PMLR, 2019b.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterrize two-layer neural networks. In 36th In- termational Conference on Machine Learning, ICML. 2019, pp. 477-502. International Machine Learning Society (IMLS), 2019.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. arXiv preprint arXiv:/409) 449.00433 2014.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009.1 1264, 2020a.
Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.09288 2020b.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 1377-1111 2020.
Alon Brutzkus and Amir Globerson. An optimization and generalization analysis for max-pooling networks. In Uncertainty Artificial Intelligence, pp- 1650-1660. PMLR, 2021.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, pp. 10836-10846, 2019.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 15845997 2021.
Qiwei Chen, Huan Zhao, Wei Li, Pipei Huang, and Wenwu Ou. Behavior sequence transformer for e-commerce recommendation alibaba. In Proceedings of the International Workshop on Deep Learning Practice for High-Dimensione Sparse Data, pp. 1 1-4, 2019.
Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. A generalized neural tangent kernll analysis for twoel neural networks. Advances in Neural Information Processing Systems, 33, 2020.
10

---

Published as a conference paper at ICLR 2023
Xiangxiang Chu, Zhi Tian, Yuqing Wang. Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Twins: Revisiting the design of spatial attention in vision transformers. Advances in Neural Information Processing Systems, 34:9355-995 1-9355-9366, 2021.
Cean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self- attention and convolutional layers. In International Conference on Learning Representations, 2019.
Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Infor- mation Processing Systems, 33:20356-20366 2020.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2018.
J. Deng, W. Dong, R. Socher, L. J. Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248 255, 2009.
Yihe Dong, Jean- Sann-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference Machine Learning, pp. 2793-2803. PMLR, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An im- age is worth 16x16 words: Transformers for image recognition at scale. International Confer- ence on Learning Reprexentations, 2020.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations, 2019. URL
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference Machine Learning, pp- 5793-5831. PMLR, 2022.
Haoyu Fu, Yuejie Chi, and Yingbin Liang. Guaranteed recovery one-hidden-llaer neural networks via cross entropy. IEEF Transactions on Signal Processing, 68:3225-3222 3235, 2020.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Efficient token mixing for transformers via adaptive fourier neural operators. In International Conference on Learning Representations, 2022.
Yangji He, Weihan Liang, Dongyang Zhao, Hong- Yu Zhou, Weifeng Ge, Yizhou Yu, and Wenqiang Zhang. Attribute surrogates learning and spectral tokens pooling in transformers for few-shot learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recogni- tion, pp. 9119-9129, 2022.
Jiri Hron, Yasaman Bahri, Jascha Sohl- Sohl-Dickstein, and Roman Novak. Infinite attention: Nugp and ntk for deep attention networks. In International Conference on Machine Learning, pp. 4376 4386. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, pp. 8571- 8580, 2018
Remi Tachet des Combes James Vuckovic, Baratin Aristide. A mathematical theory of attention. arXiv preprint arXiv:2007.0288 2020.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence mod- cling problem. In ICML 2021 Workshop on Unsupervised Reinforcement Learning, 2021.
Samy Jelassi, Michael Eli Sander, and Yuanzhi Li. Vision transformers provably learn spatial struc- ture. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho Advances in Neural Information Processing Systems, 2022. URL Id=eMM9AkXaRRs
net forum?

---

Published as a conference paper at ICLR 2023
Zi-Hang Jiang, Qibin Hou, Li Yuan, Daquan Zhou, Yujun Shi, Xiaojie Jin, Anran Wang, and Jiashi Feng. All tokens matter: Token labeling for training better vision transformers. Advances in Neural Information Processing Systems, 34:18590-1860,, 2021.
Stefani Karp, Ezra Winston, Yuanzhi Li, and Aarti Singh. Local signal adaptivity: Provable feature learning in neural networks beyond kernels. Advances in Neural Information Processing Systems, 34:24883-24897, 2021.
Jacob Devlin Ming-We Beig Chang Kenton and Lee Kristina Toutanova. Bert:
Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACI-HET, pp. 4186, 2019.
Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschítz constant self-attention. In International Conference Machine Learning, pp. 5562-5571. PMLR, 2021.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced re- search). URL http://www. cs. toronto. edu/kriz/cifar: html, 5(4):1, 2010.
Yoav Levine, Noam Wies, Or Sharir, Hofit Bata, and Amnon Shashua. Limits to depth efficiencies of self-attention. Advances Neural Information Processing Systems, 33:22640-22651, 2020.
Changlin Li, Bohan Zhuang, Guangrun Wang, Xiaodan Liang, Xiaojun Chang, and Yi Yang. Au- tomated progressive learning for efficient training of vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12486-12496, 2022a.
Hongkang Li, Meng Wang, Sijia Liu, Pin- Yu Chen, and Jinjun Xiong. Generalization guarantec of training graph convolutional networks with graph topology sampling. In International Conference on Machine Learning, pp. 13014-13051. PMLR, 2022b.
Hongkang Li, Shuai Zhang, and Meng Wang. Learning and generalization of one-hidden-layer neural networks, going beyond standard gaussian data. In 2022 56th Annual Conference on formation Sciences and Systems (CISS), pp. 37-42. IEEE, 2022c.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31, pp. 8157-8166. Curran Associates, Inc., 2018.
Zhexin Li, Tong Yang, Peisong Wang, and Jian Cheng. Q-vit: Fully differentiable quantization for vision transformer. arXiv preprint arXiv:2201.00 07700, 2022d.
Youwei Liang, GE Chongjian, Zhan Tong, Yibing Song, Jue Wang, and Pengtao Xie. Not all patches are what you need: Expediting vision transformers via token reorganizations. In International Conference on Learning Representations, 2022.
Tatiana Likhomanenko, Qiantong Xu, Gabriel Synnaeve, Ronan Collobert, and Alex Rogozhnikov. Cape: Encoding relative positions with continuous augmented positional embeddings. Advances in Neural Information Processing Systems, 34:16079- 16092, 2021.
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. the expressive power of self-attention matrices. arXiv preprint arXiv:2106.03764 2021.
Yang Lin, Tianyu Zhang, Peiqin Sun, Zheng Li, and Shuchang Zhou. Fq-vit: Post-training quanti- zation for fully quantized vision transformer. 2022.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer. Hierarchical vision transformer using shifted windows. In Proceedings the IEEE/CVF International Conference on Computer Vision, pp. 10012-10022, 2021a.
Zhenhua Liu, Yunhe Wang, Kai Han, Wei Zhang, Siwei Ma, and Wen Gao. Post-training quanti- zation for vision transformer. Advances in Neural Information Processing Systems, 34:28092- 28103, 2021b.
12

---

Published as a conference paper at ICLR 2023
Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, and Jianfei Cai. Scalable vision transformers with hierarchical pooling. In Proceedings of the IEEE/cvf international conference on computer vision, pp. 377-386, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, al. Improving language under- standing by generative pre-training. 2018 8.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. 2019.
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems, 34:13937-13999, 2021.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning. pp. 8346 8356. PMLR, 2020.
Zhenmei Shi, Junyi Wei, and Yingyu Liang. A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations, 2021.
Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.0760 2021.
Fei Sun, Jun Liu, Jian Wu, Changhua Pei, Xiao Lin, Wenwu Ou, and Peng Jiang. Bert4rec: Sequen- tial recommendation with bidirectional encoder representations from transformer. In Proceedings of the 28th ACM international conference on information and knowledge management, pp. 1441- 1450, 2019.
Ychui Tang, Kai Han, Yunhe Wang, Chang Xu, Jianyuan Guo, Chao Xu, and Dacheng Tao. Patch slimming for efficient vision transformers. In Proceedings the IEEE/CVF Conference Com- puter Vision and Pattern Recognition, pp. 12165-12177 2022.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference Machine Learning, pp. 10347-10357. PMLR, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint 1rX1v:1011.3027, 2010.
Wenhai Wang. Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp- 568-578 508,578,70 2021.
Ziyu Wang, Wenhao Jiang, Yiming M Zhu, Li Yuan, Yibing Song, and Wei Liu. Dynamixer: vision mlp architecture with dynamic mixing. In International Conference on Machine Learning, pp. 12271-2 PMLR, 2022.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. arXiv preprint arXiv:2107.13163 2021.
Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning, pp. 11112-11122. PMLR, 2021.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
13

---

Published as a conference paper at ICLR 2023
Hongxu Yin, Arash Vahdat, Jose M Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. A-vit: Adaptive tokens for efficient vision transformer. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 10809-10818, 2022.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeene functions? In International Con- ference on Learning Representations, 2019.
Shuai Zhang, Meng Wang. Sijia Liu, Pin- Yu Chen, and Jinjun Xiong. Fast learning of graph neural networks with guaranteed generalizability: One-hidden-layer case. In International Conference on Machine Learning, pp. 11268-11222 11277. PMLR, 2020a.
Shuai Zhang, Meng Wang. Jinjun Xiong, Sijia Liu, and Pin- Yu Chen. Improved linear convergence of training with generalizability guarantees: A one-hidden-layer case. JEEE Transactions on Neural Networks and Learning Systems, 2020b.
Shuai Zhang, Meng Wang. Pin-Yu Chen, Sijia Liu, Songtao Lu, and Miao Liu. Joint edge-model sparse learning provably efficient for graph neural networks. The Eleventh International Conference on Learning Representations, 2023. URI/httpp::/opeeeeviie net net.rorumm
1d=40laFt2_CV2
Qinging Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. Proceedings of the 39th International Conference on Machine Learning, 2022.
Kai Zhong, Zhao Song, and Inderjit S Dhillon. Learning non-overlapping convolutional neural networks with multiple kernels. arXiv preprint arXiv:1711.03440 2017a
Kai Zhong, Zhao Song, Prateek Jain, Peter L Bartlett, and Inderjit S Dhillon. Recovery guarantees for onyy neural networks. Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 4140-4149, 2017b. URL btttps://exxiv.org//poo
1706. 03175.117.... pdf
Guorui Zhou, Xiaoqiang Zhu, Chenru Song. Ying Fan, Han Zhu, Xiao Ma. Yanghui Yan, Junqi Jin, Han Li, and Kun Gai. Deep interest network for click-through rate prediction. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp- 1059-1068, 2018
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in Neural Information Processing Systems, 14:17723-17736 2021a.
Mingjian Zhu, Yehui Tang, and Kai Han. Vision transformer pruning. KDD 2021 Workshop on Model Mining, 2021b.
Difan Zou and Quanquan Gu. An improved analysis of training over-parameterized deep neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, 2019.
14

---

Published as a conference paper at ICLR 2023
The appendix contains sections. We first provides a brief discussion about comparisons be- tween our works and other two related works in Section A] In Section B we add additional experiments for the verification of our theory. In Section C we introduce some definitions and assumptions in accordance with the main paper for the ease of the proof in the following. Section D first states a core lemma for the proof, based on which we provide the proof of The-
orem I and Proposition and Section E gives the proof of Lemma 2 with three subsections to prove its three main claims. Section shows key lemmas and the proof of lemmas for this paper. We finally discuss the extension of our analysis in Section G. including extension to multi-clasiffcation cases, general data model cases, multi-head attention cases, and cases with skip connections in Section G.1
G.3 and G.4 respectively.
A COMPARISON WITH TWO RELATED WORKS
A.1 COMPARISON WITH (LLEN--HU & LI. 2023
Allen-Zhu & Li 2023) studies ensemble learning and knowledge distillation. Its main proof idea that given large amounts of multi-view data, each single model learns one feature, and then en- semble learning integrates all learned features and, thus, improves over single models. Knowledge distillation applies softmax logits to make use of information learned from the ensemble model. It is analyzed in similar approach to studying single models. The single models considered in (Allen-Zhne & Li 2023 is a two-layer Relu network.
In contrast, in this paper, we consider two-layer Relu network with an additional self-attention layer. The network architecture and training algorithm for the self-attention layer is completely different from those for the softmax logit in the knowledge distillation function. In our proof, analyze the impact of the gradient of Wo. WK. and Wy on different patterns (Claims 2 and of Lemma 2), showing that the training process helps to enlarge the magnitude of label-relevant features. We also show that neurons in Wo mainly learn from discriminative patterns (Claim I of Lemma 2). Such a learning process is affected by the error in the initial model and the noise in tokens. Please see details in "Proof idea sketch" and "Technical novelty" in Section 3.3 on Page 7 of the paper. This technique we develop plays critical role in our analysis of self-attention layers. This technique is novel and did not appear in any existing works.
A.2 COMPARISON WITH JELASSI ET AL. 2022
Jelassi et al. 2022) is a concurrent work which theoretically studies Vision Transformers. The major difference between (Telassi al. 2022) and our work is that we consider different data models and network architectures. In (Jelassi et al. 2022), the data model requires spatial association between tokens. The attention map is replaced with position encoding, and the training process of the atten- tion map is simplified to train linear layer. Our setup models the data mainly based on the category of patterns. We keep the classical structure and training process of self-attention, where Wo. WK. and Wy are trained separately. The required number of samples and iterations are derived as func- tions of the fraction of label-relevant patterns. addition, the non-linear activation function they consider is polynomial activation, instead of Relu or Gelu as in practice. Based on these conditions, they are able to study different and mroe general labelling function.
B MORE EXPERIMENTS
Following a similar setup in Section we add more experiments.
For experiments on synthetic data, we set the dimension of data and attention embeddings to be d = m, = m, 20. We vary the number of patterns M to be 10, 15, and 20. Data generation and the network architecture follow the setup in Section One can observe the same trend in Figure 7 and 8 as in Figure 4 and S. respectively, indicating that our conclusion that the attention map becomes sparse during the traming, and that pruning label-irrelevant tokens or noisy tokens improves the performance, both hold for different choices of M.
15

---

Published as a conference paper at ICLR 2023
Sum of attenttoon weights
Sun of attention weights
Sum of attention weights
0.8
0.8
0.8
0.6
+ Labrl-rlecant feaure
0.6g
Lubelrrlevant feature
0.6g
Label-releaan festure
0.4
- Other fostrn
0.40
Other
0.40
O-Other features
02
0.2
0.2
50
100
50 Iterations
100
Iterations
Figure Concentration of attention weights when (a) M 10 (b) M 15 (c) M 20.
Random spersiftcation cabeeerrreeaattt
Random sparsificaiin
Randam sparsification
Ei Labcl-inelevatt likers
-B Labet sabet-inclevaal tokens
ei
Naney tnkens
Nousy
1est loss
o.
toes
0.1
Test loss
0.05
0.05
0.05
0
10
20
numbar removed tokens
The number of removed tokens
Figure 8: Impact of token sparsification on testing data when (a) M 10 (b) M (c) M 20.
C PRELIMINARIES
We first formally restate the neural network with different notations loss functions, and the Al- gorithm of the training steps after token sparsification. The notations used in the Appendix is summarized in Table
Table 2: Summary of notations
F(X") ") LosssX", .y")
The network output for X n and the loss function single data. The loss function of mini-batch, the empirical loss, and the population loss, respectively.
Loss,, Loss, Loss
P,(). (t).
The features in value, key, and query vectors at the iteration for pattern j. respectively. We have (0) Pj. q, (0) = q;, and (0) Tz The error terms in the value, key, and query vectors the j-th token and n-th data compared to their features at iteration t.
x (1). ((),,,(()) (t), (t)
W(t),U(())
The set of lucky neurons at L-th iterations.
Φ (f), vn((),, (t,
The bounds of value of some attention weights iteration threshold between inner products of tokens from the same pattern and different patterns.
the
SJ.S., S", S
S is the set of sampled tokens of pattern for the n-th data. S., S are sets of sampled tokens of the label-relevant pattern and the confusion pattern for the n-th data, respectively.
Oia Orip, Ond
The mean of fraction of label-relevant tokens, confusion tokens, and -on-discriminative tokens, respectively.
For the n
F(X")
1
Σ
aTyy Relu Wo Wy X" softmaxx X
W Wq
(16)
Sh
IES"
The loss function of single data, mini-batch, the empirical loss, and the population loss is defined in the following.
Loss( X" y²) max(1 F(X"),
(17)
Note that in our proof the Appendix, we aften use the notation softmax(x," ( W- xf
which
the meaning softmax T w!(!" W)(0 xx))
16

---

Published as a conference paper at ICLR 2023
Loss:
1
Σ Loss(X",;0 .1"))
(18)
B
REB.
N
Loss
Loss(X" ,"*)
(19)
N
n=1
Loss = E w) Losss]
(20)
The formal algorithm is as follows. We assume that each entry of W is randomly initialized from (O, {?) where & vMT +1 mm } and fixed during the training. Wv, WK, and Wq are initialized from a good pretrained model.
1 Define that (0) i E [m], I [Z] is uniformly initialized from
Algorithm Training with SGD
1: Input: Training data {(X", ") size B.
n==1 the step size 1), the total number of iterations T, batch
2: Initialization: Every entry of w from (o ((0.5)). and every entry of (())
(0)
from
Uniform([++)).
vm. - - w
w K
and W
Q from a pre-trained model.
3: Stochastic Gradient Descent: fort = 0, T and ((() E (W O' w W (0)
w
( (v(() B
V Vwo W w w K w())
(21)
RER.
4: Output: W (T) W (T) w
Assumption can be interpreted as that initialize Wv, WK, and Wo to be the matrices that map tokens to orthogonal features with added error terms.
Assumption 2. Define P (Pi, P2. PM) E RM Q = (qi, Q q.. .M)) E m2 XM and R (T1, "2, TMS) E R Rm as three feature matrices, where P (PI. p2. ... Pad } Q {qn 42. ................................ .qM} and R = (21,,2.... ,rM } are three sets of orthonormal bases. Define the noise terms (1), (1) and (1) with (0)]l S a + and [|m]]( (0)ll. [10] S
for j E [L]. q1 P1, q2 T2. Suppose ( ||| W W001, [|. [ s 1, 0,7 O(1/M) and < 1/2. Then, for x/ E ST
1. O of Pj + 2 (0).
(0)
2. W K ox 4i + =. (0).
3. W o/ = r, + (0).
Assumption 2 is straightforward combination of Assumption and (5) by applying the triangle inequality to bound the error terms for tokens.
Definition 1.
on
ssp jella Coll +:0 +1511
2. (1) = |ST|ella 0110+ S51-1ST1
3. (t) = | lan (8+))==)))))) (t).
(S yf
(Sg
ify"
4. ST
S
SSy y = -1'
S
if = -1
5. 0.. = EEE
[
[3
aw = E
sm -
Ond
= Σ E
17

---

Published as a conference paper at ICLR 2023
Definition 2. Let 02 be the angle between p1 and Wo Let 82 be the angle between Pa and
Woay Define W(t), U(t) the sets of lucky neurons at the t-th iteration such that
W(t) {i: 0f < T, E [m]]
(22)
u(t) = {i: : 02 S [m]]
(23)
Assumption 3. For one data X", if the patch and j correspond to the same feature ks E [M], i.e. i E St and E Sg. we have
x; 23 If the patch i and j correspond to the different feature k, I E [M]. k * li.e,, i i St and j S k + we have
(24)
x S
(25)
This assumption is equivalent to the data model by (5) since < O(1/M). For the simplicity of presentation, we scale up all tokens a little bit make the threshold of linear separability be 1. We also take - À and À as (1) for the simplicity.
Definition 3. (Vershynin 2010) We say X is a sub-Gaussian random variable with sub- nab-Gaussian
norm K > O, if PP) S K /Pfor all p 2 1, In addition, the sub-Gaussian norm of X, denoted 11 defined as [|X 10/z supp21 p (B)X
Lemma 1. (Vershynin (2010) Proposition 5.1, Hoeffding's 's inequality) Let X1, X2..... XN be in- dependent centered sub gaubsaaasiinn random variables, and let K max; |X (a1,'" ,UN) E RN and every 1 2 O, we have
Then for every
P{
Σ a:Xi
2 Se.expl-
expl
et
(26)
K2||a||2
where с is an absolute constant.
D PROOF OF THE MAIN THEOREM AND PROPOSITIONS
We state Lemma 2 first before introduce the proof of main theorems. Lemma 2 is the key lemma in our paper to show the training process of our ViT model using SGD. It has three major claims. Claim involves the growth of w-( in terms of different directions of pi, [M]. Claim
describes the training dynamics of (i(t) and W K separately to show the tendency to a sparse attention map. Claim studies the gradient update process of (!tt).
Lemma 2. For E ST for the data with y" 1, define
Vi" = ( (t)
v x" softmax X "
( K
w x x/)
softmax(x.. a
w xt) JP + +(())= + Σ CWT(t)ps
K
SES
J#1
(27)
Σ
V.(W)i (
T
Σ v(())( (b)
On.) )
7
(EW(L)
(EW(M)
with
WP(1) (1)SST
(28)
SS
vi(t) <
Σ
(t),),
ie
(29)
2B
mL
n
S
Vi(t)= 2 2B neB
Σ
(t),
iEU())
(30)
mL
Vi(I) 2
if is an unlucky neuron.
(31)
/ Bm
We also have the following claims:
18

---

Published as a conference paper at ICLR 2023
Claim 1. For the lucky neuron E W(t) and b E [T] we have
t/
w (0)
Pi
BL
SmIaa
+ E(1 Si Pi pm (b)
(a+
(32)
n
W OG..) p S Ellp|l, for p E P/Pl.
(33)
t2tbm|
2
/pllpx(b)) - W 112 1pl2
(t)
||p
(34)
a² u²1Sm|
nEB,
and the noise (t).
n72t m
W Oa.
(1)l S ((o + +T))
)|pl
(35)
a²
For i E U(t), we also have equations as in 32 to 335
including
1f Pllm STI
O
||2ll (b) + 5(1 (ar +)))
(36)
Bt
S SS"|a la?
b= nEBs
w p Elep||, for p E P\p2,
(37)
S m a² S²
I (b))) = || 1|2 Me + t m
(38)
Bt
a
1=1 nEB, for the noise (I).
1ppmm 21)l|| ((0 ((+++))(VM
1 10 a..)
)|pl
(39)
a
For unlucky neurons, we have
w(l o p S E11pp1, for E (/(p),
(40)
|| Oa.) (1)l S ((o + MElpp
(41)
(t) 11 S |pl2
(42)
Given conditions in there exists ((1,, Q(1) 0, where i. large enough before the end of training, such that for E ST.
Claim 2.
softmax(a)) (a)) w(( (++1)
W
((+
at) >
(43)
Q
(1) - + | - (ST
sofimax(x,) w (++1)
W w(t+)) (++1)
m!") softmax(a.) 1 w Q x
Q
S ST
(44)
ell (6)1-((++))ll))))) (11l K(t),
= (157) and j $ SR. have
+ +[S01 -(55)|-1S())
softmax(a") (+++1) w (1+1) x.1) S (00 + ([8*) - (S
(45)
softmax(x7)
w (++1 x softmax(z))
w (-(t)
Q x"")
S (1)())
S
|9)0 (1)l
K(t)
(46)
(0 51 (0 + ([S")
18
For 1,2,
(t)
K()))q:
(47)
(t())
(1 + Q())::
(48)
1=0
19

---

Published as a conference paper at ICLR 2023
Claim 3. For the update of W V. there exists ((1) such that
W V 23j Pi
Σ
vi(b
Oa.)
(0)
+
Xv((b) (b)
oi..)
))+ + 21),,, j (49)
b=1 iEW(b)
W = P1 - "i Σ
Σ
Vi(b (8)
+ Σ
Av (b) T 2z((t,, jE (50)
(G))
o (1,)
b=1 (EU(b)
squ((b
WV (((+1)
J = P1
Wl(l))
(b) T
+ +2)(t), jE|5"1/((S) US!)
(51)
(z) S (0
(52)
To prove Theorem [ we either show F(X") > 1 for y" = 1 or show F(X") < <-1 for y" = -1 Take y" an example, the basic idea of the proof is make use of Lemma 2 to find lower bound as a function of axx, complexity and the required number of iterations in terms of a., 0, and such that the lower bound is greater than 1. Given balanced dataset, these conditions also ensure that F(X") < -1 for y" 1. During the proof, we may need to use some of equations as intermediate steps in the proof of Lemma 2 Since that these equations are not concise for presentation, we prefer not to state them formally in Lemma? 2 but still refer to them as useful conclusions. The following is the details of the proof.
etc.. The remaining step is to derive conditions on the sample
Proof of Theorem T
For yd 1, define C. {i E [m] W (t" 0) and K!_ {i E [m]
w
(t) (G.)
0}. We have
Relu(W (0)
(t)) + S
Σ
Relu( (4)
(t)(
STh
m
Oa,)
m
(6.)
JES
sEW())
IES (W(1)
Σ
Relu (1) (8)
(t)
SSn
m
(6.)
IES"
iEK!
(53)
By Lemma
we have
Relu ((
(t))
S
(G.)
lES" ((t)))
Σ
Relu (e
(t) (G.)
(t) +
Σ Σ
1
Relu (tttt Vi"
|S"|
m
IEST SEW()) m
lES
(EW())
S|ST
(t)
Σ p.otmaax (x" (( wl + + + Σ
Wi(u)p:
a
Oaa
BEST
1+8
V3(1WO
4
Σ V3(DA ()
Oo,, 3))IN(01+0
OGn
jEW())
jEW(
1ST mm
TM+
2222m
|p p1)))p)()) (a + (t)
18 18"|a.M
Cm
Bt
a
t|Sa
b= nEB,
ST po (b)mm
(a + -TM
+ Bt
S |SI|aM aM
Em
π
b=1 nEB
17 tbm ST a² S
Pn 16)) 11pppp11122 l1
Bt
(54)
20

---

Published as a conference paper at ICLR 2023
where the second step comes from 27 and the last step is by (136). By the definition of K2,, have
m Relu( (1) (c.))
()) 2
Sth
(55)
LES" tExt,W(0)
Combining
(136)
and 138.. can obtain
Relu(W Ou,) (t"(t)
8"
IES iEK!
18% Im $ S
(a+ T)M TM (Elppl + Bt nEB,
(1 cm
cm
Σ Σ
π
SMIaM
(o+T)M
pfttm (0 + T)1l
(1 6m
a²
1 AVM
7
S Pm (b))tt
|p11 + ((o +
VBa²
a²
Bt
S aMM
nEB
(o +T 1p 111p2113) (t) +
M S
"p Av MER m
(1 6m
(Sl|pll
|pl
a
ST
v /Ba²
n2tmm m a²
(S + | PP (t)mmm
(o +
+ +(o+T))(V
+
)lpl + Bt
18
Cm
b= nEB
1ppm
|SP|
2bm ST
-1sT)) +
M
| (pt)
a²
S
Bt
a² Sm
h nEB.
ST P (f)mm
tbtmn||
110 T) - pm((t) +
Bt
S aMMaM aaM
02
M
Bt
a² S*
b=1 nE
b=1nEh
Pr. (p1())))lp1))ppm((()
(56)
for some C1, E (0, 1).
Note that at the T-th iteration, (1))
Bt
Σ
S Pa (1)mtm
S laM
(1 cm (T)
)
Bt
Σ Σ
ttbm|
a?|sa
(n)) ppm)))
Pr.
b=
NEB
1=1 nEB,
2 a²
SST]
S" S (n( (b) (a + (2))))/:1(())))))00(0()(1111 187))ll)))))))))))))0))))))
n Bt
17
(t))
(57)
Since that
41 (T) 2 (1 + min {K()))
(58)
2 (1)
To find the wrder-wise lower bound of qı (T), we need to check the equation
91 (1 + ellq: (T)) 2 -- (6+7) q1(()]] T
(59)
One can obtain
Olog T) = (T) o(T)
(60)
21

---

Published as a conference paper at ICLR 2023
Therefore,
T
Pa (T) 2
21
2 1 - -(19) (1-1)
(61)
TO
1-a
a
1-a
&n (1()(((511)) - (s S Sn
(62)
for some large C O.
We require that
|ST m
(o +M
ST (b)))(
(o T)M
S"h aM
Fm
π
BT
S
cm
π
b=1 NEB,
T
n tbbm ST
(b)) || (m))))
BT
a² ST
1=1 REB,
pm D
|P1| (p) - (ar + =))
BT
b=1 nEBs
|
(o TM
IS (T)mTmm S IM
6m TT
1 N
T2 ST a²|s
Sth aM
cm
1 N
T
S
(T)( (T) +
Ppi 10 Pm (T) - (o + +)))
a²
St
=00(1))) + >1,
(1 aa
(63)
where the first step is by letting a a= /m and m 2 We replace pr (b) with pa (T)) because when b achieves the level of T, b Pm. (b) is close to 6 for 01, 02 2 by (61. Thus,
T
8 pm (b) 26))= T (p (0(1) T) 2 T pm. - (T)
(64)
We also require
17² Xv M E0:
(65)
/Ba
for some 60 O.
We know that
N
ST
|S
Pa
- (T) - EE
N
SD
||S0|
n=1
ST
(T) (T) (a+ + T)) - E
S
Pm (p) - (a +r)+ (a
N
S"
|Sa
(66)
E
(P) (T) (o + 7))
S
log N
+ c( () + c" ((o +))
N
for d' > and c" > 0. We can then have
12 T 1)
1) 1
(a+r)
15 [
pp (12 - (a + T))px(() (1)
cm
N
17
=0(
(1
cm
(o+T)dd
[ST (
1
log N
c(1 () - (a + T))
(67)
17
=0(
(1 - -mmm
(
ST
3
22

---

Published as a conference paper at ICLR 2023
where
1 and
a 2
(68)
E5e E5e (6+7) (1- TT)
by [17p. as long as
1
N 2 M1
(69)
(a cd((1 - c'( + +)))²
and
1 B 2 (T o)]e (6+r) 2a, (r (6) = 0 e (T + +o))
(1)
(70)
where ( 1
10_ If there is no mechanism like the self- compute the weight using the correlations between tokens, we have
()--() (1a
(71)
which can scale up the sample complexity in (69 by o 2 Therefore, can obtain
F(X") >
(72)
Similarly, we can derive that for y/ -1,,
F(X)= -1
(73)
Hence, for all [N].
Loss ,y") =0
(74)
We also have
Loss (x)) Loss X" y""]] 0 with the conditions of sample complexity and the number of iterations.
(75)
Proof of Proposition [
The main proof is the same as the proof of Theorem I The only difference is that we need to modify
66 as follows
ST
S
Pn (T) - (a +))
S
S
S
S
Pr (0)(()) (0) +))
Pa (a
(a
N
Sm
S
(76)
ST
(0)p.))) - (a T)) - -1)]
E
S
log
+ 1 - 0(x)) + (a-))( + +)
where the first step is because P.. (T) does not update since W (1) and w are fixed at initialization ((()) (0) and w((0) and the second step is by pm (0) O(a.). Since that
log N + +11- O(a + { a.
(77)
N
we have
N (O(u:.) + ( (c.)) T))2
(78)
= U2 (a, - -))2 -
Proof of Propusition] 2 It can be easily derived from Claim of Lemma 2 (60). and 661
23

---

Published as a conference paper at ICLR 2023
E PROOF OF LEMMA 2
We prove the whole lemma by a long induction, which is the reason why we prefer to wrap three claims into one lemma. To make easier to follow, however, we break this Section into three parts to introduce the proof of three claims of Lemma separately.
E. PROOF OF CLAIM OF LEMMA
Although looks cumbersome, the key idea of Claim to characterize the growth of W in terms of pt, I [M]] We compare W Pt and W (t) Ou.., Pt see the direction of growth by computing the gradient. One can eventually find that lucky neurons grow the most in directions of Pi and P2, i.e., the feature of label-relevant patterns, while unlucky neurons do not change much in magnitude. We start our poof. At the i-th iteration, if SY. let
(t) W X "softmax X
W oh x/²)
Σ
softmax T w
w ot pi + z(1) + Σ (T(0p)
(79)
eESi
j#1
Σ
Vi(W
(
Σ Vi(O) (1)
O )
Oa.)
(EW())
(EW())
. E [M], where the second step comes from
27
Then we have
S (91 (i1l
(T((1) S (1S*)
+ -00 = (t))ST)
(80)
18
(e
Hence, by (18).
Loss
1
1d
Σ LES
ai, WWo,.. (7"(1) (1)
(81)
dW WWo
B
nEB
Define that for j [M].
I B
1
Σ
my"
S
Σ
am) 1 (0)
(t) Σ Vi(t) OG.., Pj O..)
(82)
nEB,
LES"
kEW())
16
my
S
am, (w ((")) E
Ou.)
Va(t) (1)
Oce,, Pj.
(83)
RER,
LES"
kEW(L)
and we can then obtain
W (t-1)
witt) T
O6.
PPj
PPj
aa1) (w((9 (1)
(t) 2 (t) p
7y
B
8
IES"
1
(( (t)
(t" 01zi(t) p
B
S
O
W
lESn
1
[w(( (t)
(1" 0 Σ softmax (wo. wi!)
B
T
S
a())
(w))
p3
nEB
LEs"
1
1w ( (n.)
Vi" wi(l) P5 + T4 Is
T
am
B
S
DEB
LES"
1 + lg + I3
1 + Is.
(84)
24

---

Published as a conference paper at ICLR 2023
where
1 =
Σ
ny"
Σ
(W OR.) (1 2 (}z((t)
(85)
B
S
nE B
LES"
B nEB
1
Σ
mV S
(()). 1 (W
(
(0)
(t" 0)
Σ
softmax (2r w
w m)pp
Pj
LES"
SESt
(86)
13
nV" n S
Σ am)) 1 (c.)) (t"
Σ
Wi(t)pkP
(87)
B
REB,
JES"
Reghl
We then show the statements in different cases.
(I) When j = L, since that Pr(y" can derive
= = Pr(y" = -1) = 1/2, by Hoeffding's incquality in 26
log B
Pr
Σ yrh N
S B
(88)
B
B
nEB,
Pr ([=(1) T P1 2 v ((o + r) logm) m)
S m
(89)
Hence, with probability at least (mB)) - we have
|1 ?1(0 + r)) log m log
(90)
B
For i E W(t), from the derivation in 133 later, we have
w v x softmax(x"" w T w xt) > 0
(91)
K
Denote (t)
|8T (t)e q (2)1l 26 (+)
(, Hence,
12 2
SST|
S
P 11 (t)
2 11 B |
Σ
STT|
||1 112 (t)))
(92)
B
S
NEB,
NEB,
13 =0
(93)
SUSTT
SST m S JaM
(
14 B
Pr (t)llp (1 cm
Wo P1
W
SSa| 2B
(94)
1=1 NEB
REB,
(a + +TM
S m
[6
(1 cm
la Pm. Po(())))))) (1]pp) 2
Wo P2
B
S
2B
S"
b=
neBs
nEB
(95)
17 2m
+
W W Pas (1 (11+000 (a (T++))
VBa²
Hence, combining
90
92
933, (94),
and 95
we can obtain
( Oun
(1))
P1
Oc..) PP
1
ST
(t) ((0 + r) + nt ST
SS m
(a + TT
B
Sn
|
2B nEB,
S I''
Pm (t)
Em
neB,
Wo Pi (1 (T) n/ S
|S Im
(a (o+
S aaM Pr (1) cm
SnI
2B nEB,
π
n/m W Wo... PM (1 + (a + 7))
Wo
T)) (1
(a
)|(pi|
/Ba
17 Σ aB
SS
nt/S
1
SS m
Pt) - ((o +7)
S aMM Pro (t)
S
St
2B
BE,
n
TM
(1 6m
Wo Pi) Pi)lpp 12
π
(96)
25

---

Published as a conference paper at ICLR 2023
Since that W (0) (s.)) No( ) ma
by the standard property of Gaussian distribution, we have
Pr( Pr[ (0)
(...) | 5) 5
(97)
Therefore, with high probability for all i [m]] we have
W 00)
6.) 2
(98)
Therefore, we can derive
H+1
17 bb 1)m
(++1) O(yy
Pı exp(
Smlaa |ST]ll + (a+
B(t + 1)
REB. 1
(99)
expl
1
Σ
+ 1)
SS P Pm (1)) + E(1 -
(a (a +))
B
18"1a2
BEB
by verifying that 17 2 tm a
1
n
m
1
12 1)m
exp(
1 +
2 exp(
(exp(
1)
9(1)
a2
(1)
a²
O(1) a²
1
1 m
1 mm
2 exp(
(1))
a²
2a
(100)
When mjt < 4 have
+ tm
()
(101)
When npt 2 m. we have that
s(
17 tm 1
1 exp(
1 4)
17
g( 17m
> 0
(102)
a
a²0(1)
since that gt) is monotonically increasing. Hence, (99 is verified. Since that
nyt nt 0(1,
(103)
simplify the further analysis, we will use the bound
1 [+1 (t 1)bm |S"|a2
w (++1) (
P1
ΣΣ
|S] [ell|p111 (b) + 5(1 (a + +))
(104)
Bt
neb,
Note that this bound does not order-wise affect the final result of the required number iterations. (2) When Pj E P/pp., we have
I2 = 0
(105)
S
log log
Tal
S
(1)
||||
(106)
a
B
REB
77
log m log
ST
bybm
(nt)²
144 <
aaM Pn (b)
+ Ellppl
(107)
a
B
2B
Sh
a²
b=1 nEB,
17 122mm
1
log m log B
S m S SS Pr. (t)
|sl <
Ellp||
(108)
BBa²
B
2B
6=1
nEB,
with probability at least (mB)-
107
comes from
34
Then, combining
90
105
106
107 and (08). we can obtain
(++1) T
-p:)
w
Pg
Wo,,
7
ST
St
vn(() + ((o +))
a
B
S"
(109)
n
S Pm (b)mm 1f1mm SS'laM
log m log
tm
+ +())
|| +
/Ba² Ellp|
a²
B
26

---

Published as a conference paper at ICLR 2023
Furthermore, we have
(++1)
St
S) (m()) + ((o + +T)
S"
BEB
S pm (6)mm
n1²t2
logm log B
1
(110)
V)
|pl|
+
Elpp|
+
Elpll
|Sm|a aM
B
v BBa
b=
<11p
where the last step is by
nt O(1)
(111)
to ensure a non-zero gradient.
(3) if E U(t). following the derivation of (104 and 110).
conclude that
(t 1)mm S
w
P2 2 Oun)
| pall'pa(b) + E(1 (o + 7))
(112)
B(t
1)
STIaa
n
wiit)
S Ellp|l, for p E P/p2, Oo..
(113)
(4) if & (W(L) UU()).
1/2 + 1al S
VV
?
log m log
|pl
(114)
B
Following 107 and 108
we have
log log B
1
S m
1f
[44|
S Ja Pr (b)
(b)
+ 1ppl
(115)
a
B
2B
a²
b=1
REB,
17 1m tm
logm log B
S m
$lpp|
8 aM Pro (b)
(116)
v 2
B
2B
nEB,
Hence, combining
114)
115
and
116
we can obtain
(1+1) T Oc.,)
w
(t)
PP
o
P
5! (lpl) + ((a +
T)) +(a+))))+ +
S Pa (b)mm
22
+ E))
log m log B
|pll +
7m
Elpp!
S Ma
a²
B
VB
(117)
and
w Ou.) P
(lpll (T))
ST p. (6)mm
+ E))
SS
a²
log log B
12 m 2
(118)
|pl|
Ellp|1 Ellell
B
<El|p|I
where the last step is by
11 0()
(119)
(5) We finally study the bound of W Oc.. and the product with the noise term according to the analysis above.
By 52) for the lucky neuron i, since that the update of W w (1)
lies in the subspace spanned by P
and Pl "2 PM all have a unit norm, we can derive Oa..
W Oa.
(2
(wittD)
(++1) Oci pi) (1)+ > (wit+!)
(1+1) Oc. P1)²
(120)
?t + 1)b|ST m a S²
2(
2 (b)))
B(t neB,
27

---

Published as a conference paper at ICLR 2023
(P(++1)2m) 1111p
w (
ppl
(121)
a²
(++1) o On.
21(1)| (()+))))
Σ
W (t)
(i.)) :P)
(122)
(2 (t+
1)²m
<((o+ (V)
1pp
a²
For the unlucky neuron i, we can similarly obtain
w )
(1+1) Wo 20(1) s ((a + r))
Σ
((a VV)
(a.) .P
(123)
PEP
(1+1
1 < |p|²
(124)
(
We can also verify that this claim holds when = 1. The proof of Claim | finishes here.
E.2 PROOF OF CLAIM OF LEMMA 2
The proof of Claim is one the most challenging parts in our paper, since that we need to deal with the complicated softmax function. The core idea of proof is that pay more attention on the changes of label-relevant features in the gradient update, which should be the most crucial factor based data model. We then show the attention map converges to be sparse long the data model satisfies (8)
We first study the gradient of W-(++1) in part (a) and the gradient of W (t++) in part (b).
K
(a) By 232). we have
Loss X )
B
OW
BEB
1
Σ
Loss X yyy F(X
-1
B
OF(X")
aWq
NEB
-1
1
Σ
m
a(), Wo Wv X xsotmaxx W Wqar?) 2 0]
Sn
NEBs
(Wows
Wvz
Wx Woæ?))
#ES"
(125)
softmax(x""
(xt T W Wok")Wk(x"" - x²)0""
TES
1
Σ
a()) 1o Wr X" softmaxx (X W Woa?) 2 0]
B
Sn LES" i=1
NEBs
(Wons
Wvx" "oftmasftmax(x"" T W Woze))
(Wxœ"" Σ
softmax(x" W Woo")
TESt
Fo ST. by 43 we have
el -
softmax(x)
w (t) K
w
>
(126)
(y (1) (1S") (
For & ST and / E ST, we have
softmax(2;) n w (t+1) w (1+
x/") S
(127)
(()]e(+ (()))) -[S" - (ST
28

---

Published as a conference paper at ICLR 2023
Therefore, for r, E ST. let
W K
Σ
softmax (x,
w K
w of W x. : (0
(128)
rES"
where
S" |
2 2
S elle (+)1*+(6++) (x)) +88 ST
(129)
on (1) 18
(t) S (t((
pn ( (l
& ( (130)
where the last inequality holds when the final iteration log T ((1)
(t =0(1) ()(() + Qe(t)r2(t)
((((t) Σ Σ
softmax (x, W (1)
W!(2z))r(() x1)ra(t)
a=1rEsT
=0(1) Σ (m
(131)
for some Qa(t) > and o > Here
|S
K S (1)
(132)
SSt ST
for 2 Note that |C/| = |f* Therefore, for i E W(t).
= S 1 2.
w-l0)
l x softmaxx"""
w w xi))
Oa.)
*
-w!l Oa,
Σ p:softmax(x")
w (1) w x))) + x(t)
Σ
Wi(w)pi
sES1
I7s
(2(b) (6)
Σ W((b)
og.)) (g.) 7)
(6)
jEW())
JEW(b)
Stm
17 p2mm
(pn(6))311p)) 11l ((o + T))
||p1| + {(1 (a - ((o + +))
Bt
S
a²
b=1 B
ST S0 (b)m SSth laM
(o+T)M
12 y²1 | 1|
Mellppll Elppill
+ T
(1 Cm
π
a²
nEB,
ST P. (b)m
(a + T)M
2222m
17 Ag 2
||
1
B
Sth aaM
v Ba
b=1 nEB
>0
(133) The final step
where the first step by 27 and the second step is combination of holds as long as
32
35
OL1)
(134)
and
XEV
M
B >
2
(135)
1
S
(t)(2))2
B
ME
S
29

---

Published as a conference paper at ICLR 2023
Then we study how large the coefficient of qı (t) in 125
If s E ST,, by basic computation given (32 to 35
woo
w x softmax(x""
(t) W wx
x
)
Pm
ST
S Pn (t) (o + T)) Ip 12-
Pil - Ellpp11 +
d²
Bt b=1 BEB,
B
b=1
|ST Pn (t)m Stt laM
(a TM (1 - cm
17A MPm
(b)
Bt.
m Sm
BBa²
nEB,
b=1 NEB,
ST Pa (t(m)
(a TMM
qp 12 m 8(1o + 1))11111))
(t)
1
B
SS MMM
(1 6m
a²
|SP
b=1
NEB,
|S t2
Pr (t)
Sm Pn (t) (a + +T))[1P
112
BI
a
|ST
b=
nEB
ST Im (6)m
(o TT
tbm |
Pn(()
(1 Em
a² (s Pn (b)) (p)
B
Sth JaM
B1
|ST
BEB
b= neB,
(136)
where the last step is by 134 and (35 If i E ST, from (36) to 39 we have
will w x "oftmax(x" Wx
w (4)
Q x7")
|S Pn (b)mbb
) p
(0 y + Tllpp
S(El|pl|
+ Bt
(1 Cm
18 laM
π
a²
b=1 nEB
(137)
122m
S po (f)mm
1|p1122 + ((a++)))))))))) (V +
+
VBa2
a²
Bt
SS aaM
b=1 nEBs
(a (
(1 -1-6m 6m
a²
π if W(I) and 8 & (ST U S
w
WooD wh x softmax (x
Wx (t) w x²")
1/ Mgt m
7m2m
S(El1p]|+ +
||||2 + ((o + (v)((V
1lpl
VBa
a²
(138)
1 + Bt
(IST) + SST (T (6)mm)
(a (T+
172t'm
(1 -1 Cm
=1 (t)
S aM
a²
b= n
by 40) to Hence, for i
E W(I.j E St. combining
129 and 136), we have
(10 (0)
w x. softmaxx(x" Wx w 2 (()
"E
-(W K x
Σ softmax(x" (x;; T w( wi x²" ) W (2) x²) T ag
S pm (b)mmm
(o+
1 ST
(1 cm
Pn (b)) |p pl
Bt
SM|aM aM
π
Bt
az|s
b= nEB.
b= n
|ST|b 22
SS Pa (b)
(a + (T))))llp 11 (t) o, (t) (1 (t)l
Bt
a²
b=1
(139)
30

---

Published as a conference paper at ICLR 2023
For i u(t) and I E S j E S
w (ee)
Σ
w x softmax((""
T w w =f))(t(((t)
K
SES"
(W K x.
softmax
l ee oi
W K
F. )2f x
rEs"
ST Pm (b)mttm Sthh aM
11 f (
)
(1))
Bt
cm
a²
b= nEBs
S Pm (6)mt
(o T
1 pm 2 + 1)11
BI
S
cm
7
a²
NEB.
(1))
(140)
Fori & (W(L) UU())) and E j E
Σ
w x softmax(x"" Wx ( a")(1)
Wo.).
BES"
(w xy
Σ
softmax (x, W Wo o. Jar )as 2;
rES"
(141)
pptm
(SS
+ |S Pr (b)mtm
S(Ellpll + ((o + T))(
+ (ME))p 21 + Bt
a²
aM|Sa
6= nEB,
(a + T)M
1 22mm
26 Mm Mmlppl? (1(((a)
(1 - cm
|pil| +
π
a
VBa²
To study the case when 1 & ST for all n E [N]. we need to check all other l's. Recall that we focus on the coefficient of qı (f) in this part. Based on the computation in (137) and (138), we know that the contribution of coefficient from non-discriminative patches is no more than that from discriminative patches, i.e., for & (S U UST), n E [N] and ST.
( x. softmax(x"" w (T) who xt (1))((11))))))))))))))))))))))))
Woo..
"ES"
(W(() (1) xx"
Σ
softmax (W (1) xf ( who x: w (t) x jor agl
rES"
(142)
(w(0 (t)
Σ ( x. softmaax( T W
w x (t))
Woa.)
K
SES
(W x
softmax W xx w w x; w(( x; xx x
rESH
31

---

Published as a conference paper at ICLR 2023
Similar to 139
have that for / E Sg. J E Si, and e U(t),
Σ
who x. softmax(x"" w W x (t)
SESM
(W K 2.
softmax W zF
w K
W K
x,," xx; x
Q
rES
S pm (t)mttm
lo + T)M
11 p
|S
(1 cm
)²
3 :B)(t))
#
B1
SmIaM
a²
St ST
b= n
ST Pm (b)nttm
(o + T)M
n2
(1)11))1 +
(1 - Em
( + +1)111133100)
B
Sth aM
a²
b=1 nEB,
S
(t)l
S S
(143)
Therefore, by the update rule,
w((+!) Q
x, - -W! æ; ( OT.
w
(owo
Q
M
+ +
+ + (t,(t) +
K. (b)b) (b))g2(()) (2()) + + (t)
(144)
b=0
1=3
(1 + (t)) 0(1) (t) +
(b)lgz(2) + Σ (t)
b=
1=3
where the last step is by the condition that
(1) (1) -r)
(145)
and
(2(1) kz(() rl)
(146)
32

---

Published as a conference paper at ICLR 2023
for (t) and (t) >> 0 from induction, i.., q1(() and r rt()), (1()) and (1( are in the same direction, respectively. We also have
(t)
S Pn (b)mlm SSIaM
(a (M+TM
1f ST
20
(1 6m
(b)))
Bt
Bt
a² a²|sa
6= NEB
b= nEBs
f'fm
bJST!
(4)
t Pn(t) (1) - (a + (1)))p::11pp..
11 (1) on (100(((1--(5)))))))
Bt
a²
nEB,
S Pr (b)mt Sth laM
(o + T)
p )² , (1)0)))))))
(1 6m
Bt
a
b=1 REB
SS pn (b)mtm
lo (
1j 1ttm m
(a + +) 2d (0)
11 Bi.
SS
- cm
a²
NEB
p2
(0))tt
(t)1 n(Elppll + ((or + +T))(
a
Bt
aMIS"
hm nc
(o +T
1/2t2
ntAE Amlplll!
(11 -m
(t)ll
/Ba²
S Pr (b)mt
(a + T)
17 122
|S
>²|1p2|1? B1(()) |Sa [S
1) Bi
Sth MM
6m
π
a
b=1 mEB.
ST Pm
(a T)
17 pp ) + 11)1p1/331()) )
((1(())12
(1 - Em
Bt
Shn aM
a²
nEB,
S
| (1)
S
1
SS pn (b)mbm
(a (o+T)M
1
1 Sb I
27
(1 Cy
Pa (b))
Bt
S"
π
Bt
a?|sa
b=1 nEB
b=1 n
12t2m m
b|ST|
(t)
(t) - (a + -))]lpp11 (/²m(t)))))((()))) on (t)(ss) (
Bt
a²
1=1 n
>0
(147)
MI S 1 Σ K(I)
S
Ml
(148)
B
SSa
18
REB
1 (K(0) B
Σ
K(t)
SS
(149)
[S"|
ST
REB.
33

---

Published as a conference paper at ICLR 2023
as long
L Σ
L' m
|
SS pm (t) |8m|a.0
ES
Pn
pn(t)(
[ Pn (b) - (o + [))[[ppp
+
Bt
d²
Bt
b= nEB
b= neB,
(o + M)M
1
nf tbm | az|sma|
(1 6m
Pn (pplt)) (m)(t))(((t))
π
Bt
b=1 neB,
(18 -15 (t)1
S pp (b)mtm
(o T)+ cm
p
)(
pp2 (1) 21
#
S laM
a²
|S S
b=1 nEB
(o + M)
11² f
||1(1011+ + Bt
(1 cm
( + ( 11p1113310)
SMIaM
a²
b=1
NEB.
S Shh
[q (1)l
S
(150)
To find the sufficient condition for [50), we first compare the first terms of both sides in [50). that when
Note
n s 0(),
(151)
we have
n2²2 ス 1/²th
(152)
When |S
>>> |ST by (30).
2n (t)(1 (S 2 (t)
(153)
From Definition I we know
S
1 12 (t)) 2 pa = 0
1SP1e + | 1879 ) 2 0(e (65-r)
(154)
Moreover,
1
ST
( oo) (6+T) (r)
(1 (r o))e (6+T)
EDl
Bt 1=1 nEB
|
|S |S
ST b |Sa t
- (r + (e (6+T) - c -6+++) ((
(r + σ )]
Bt b=1 BEB
Bt
b=1 nEB,
+ +/e (6+ D E
ST
(6+ (T + ()
1
ST
Bt
St
2
2
5
REB,
<e- (6++)
log Bt
21 6 (6+))
1
-
(r + o)e (6++)
Bt.
log Bt.
<e- (6+r)
+ (T + toe (6+7)
((+t))
Bt
(155)
where the first inequality is by the triangle inequality and the second inequality comes from |ST b/((S < 1. Therefore, a sufficient condition for [150 is
1 Epl Ese (6+z)) (T + a))
S S
S
log Bi
ED
(156)
2
S"
Bt
S
Bt
b=1 nEB
by Hoeffding's inequality in (6). Thus,
ST |Sa
1- and
0. =: Epl
>
(157)
+ Ege (6+r)
34

---

Published as a conference paper at ICLR 2023
if
Bt 2 (T (5 (t-T) za, -r)²
(158)
For the second terms on both sides in 150p. since (T + σ) S O(1/M). the inequality also holds with the same condition as in and Bi.
and
S
a#
E
|ST
(159)
M
|
and
S
(160)
1=3
Given that [S*| (1), 157 is equivalent to 8
For the second terms on both sides in 150 since (o + 7) $ 1/M, the inequality also holds with the same condition on a, and Bt.
Note that |S |. we let |511 (ST)
for E [M]. We use the presentation in 149) above and 170 171) below for simplicity.
148
Then we give brief derivation of (t+1)
w 2y for j & ST in the following.
To be specific, for E S S
Loss X" y") 2. ((()))
11
owll) (0)
BEB,
Q
(o + M
1y tbm ST
2n
(1 Em
Pr (b)) | (1)
8 (S
n
Bt
a2 d
b=1
BEB
b= nEB
|5 7
(t))) S Pr (t) (o +1)))lpp,1ppp.. on -1000
B1
a²
nEB
(161)
where
ST (e (0000
(t)
(162)
(e9 Eh,, ((0):())-00+7)))0000) S - S
When ((b) is close to 0+, we have
- K()1000000 2 Ei (()]
K0)
(163)
1=1
where the first step is by log(1 + 2) xe x when + 0+. Therefore, one can derive that
1
Loss X ,y")
2/. (1(()) 2 0 K(t) ((t))
(164)
B
aw Q
nEB,
Meanwhile, the value of p(, (1) will increase to 1 during training, making the component of qı(() the ÖLoss( x"
major part in 17 1/3 ΣnEB, Hence, if j S/ for 2 3,
aw x..
W w((--) (1-11 2j ())) +01)
Σ ((b)q((b) + Σ (1((1)
(165)
Similarly, for E E
H-1
M
σ; = (1 K(t)g +01) (t) + (t-1) Q
Σ
(b)lgn1)))) + Σ (t)
(166)
b=
1=2
35

---

Published as a conference paper at ICLR 2023
(b) For the gradient of WK. K, have
Loss
1
Loss X" 3²) F(X)
OWK
B
(F(X)
OWk
nEB
1
Σ an) ( Wy xootmaxx (X"
W Woxo)) 20
B
BEB
LES" 2m1
(167)
Woon
Wv Wva"
W Woze
SES"
(x
softmax (.:; WK Womf)a;;)
rES"
Hence, for j E ST, we can follow the derivation of
144
to obtain
(((+1) 2. = + Q()) (1) + 0(1) (0()) +
(Q.())((1 2))2(0) + Σ m/(()); (168)
where
@(t) 2 - X) >>>0 0
(169)
for x < 1 introduced in Assumption
3 and
|St| |8|
hul S
Q(t)
(170)
ST
NEB
|S
1Q((11
Σ Q(1)
(171)
B
S
BEB,
Similarly, for j E Sg. we have
M
-h (++1) 2j Rd (1 + Q(1))2((() + 01)
()))
Σ
10.00)(1- ()r Σ
km((t), (172)
6=0
1=3
For j E Sn = 3,4,... M, have
1 Moss X R 1n F(X)
OW 2/ a101)
B
(F(X)
REB,
(w)) 1 (s.)) Wy X softmax (X W Woaf) 2 0]
B
(173)
BEB
LEST i=1
(Woa))
Σ
+1
Wv x softmax(x" T Wx WKW2)z)) (t)
SES
SEST
Σ B REB
Loss vnt
F(X) dW ". 20())
OF(X)
am: LWo Wy X xsotmaax (x" W. Wox)) 2 0]
B
(174)
BEB
LEST 1m1
(Woan Σ +1 Σ
Wv sottaaxaa T Wx Wozz)) (1:
BEST
=ES
(t))(())
((+1) x.j x=1 + Ck, (t)(()))((())+))))))))))))) (t)
4-1
M
(175)
C
Σ ()))
cc.
Σ (b)(()))2(0)) + Σ kn(()))
0=00
b=0
36

---

Published as a conference paper at ICLR 2023
where 0 < Cki Ckg Cks 1, and
20)0))0000
(176)
Therefore, for I Sp. if E St.
3y (T(tt)) (1+1) W (t+ zi
+ 2(++<(0)) - Q(0)) (0))))
Σ
Ka(b) Σ (011)())))
b=0
mill (
211 +
(8 (t)l
M
|
M
S
Σ
Q(t)
(t)1.
Ktt
) (1)(1))
8th STh
Sth S
DEB
nEB,
211+ K(t) + Q() - (6 +)ql
(177)
where the second step is by Gauchy-Sehward inequality. If j & ST:
" I Wo whl of (0)(t)ll + (6+))llan(0))
(178)
Hence, for j. E ST.
" softmax(x") T w (1+ wh (++1) x²)) 2 10 + ( (179)
softmax(x"" w (++1)
w a²") softmax w (t)
w (t) xx')
(I+*()))l (0)))
e (007
SP|e(L+K(0))l (1)]l (6++) (1)) + ([Sa) - (S |0 ([S*) 1 SS
1 |S
( (6+7))=0)))
-1)
(1S le + ([S*) (S
S
ST
K(I)
((67)) (6++) (1)l - (IS* - 1)
(180)
where the second to last step is by the Mean Value Theorem with
x E [llg] (1)11 - (5 + (1Il (1 + (0)) (6 (6
(181)
We then need to study if 1 (ST U and i.e.
2 w (t+1) T (++1)
W
x 2 (1 + (()))
(())))(())))))))) - 1 (+)]
(182)
6=0
For j,l $ (S S
2. T w((++) T of S Ciay X- Σ (b) ca,A Σ (Q(())r2(b)
(c+1) ((+
(183)
+ (1 + cr (Qr(0)))
We know that the magnitude of || (t)] increases along the training and finally reaches no larger than O @(Vlo T'. At the final step, we have
T
((b) 2 (6+7)))m))) 2 e(Viogg)
(184)
37

---

Published as a conference paper at ICLR 2023
Therefore, when is large enough during the training but before the final step of convergence, have if j' & (S) USE) and j E ST, we can obtain
- x;)) wk+ (++1)
(1+ We oj 2 (1) ((1 (()))l - +))
(185)
We can derive the same conclusion for j E S in 185 [)] in [177, we can obtain
Therefore, by |S2] - ege ege (6+ - (o
softmax(x)) W 11+1
W
(++1)
Q
(L) (6++)) (1))
(SS + S (0)
+ (15*) - |S (S
(186)
+ ([S*1-10) [S
softmax(x.) ((+1)
W (1+11) at)) - softmax wlo w (t) xx))
S |S
(K))
(107 (10 ((5*) 1
18" - |S
K(())
(1S7)((LKK())ll (6+7) ((S") [S,
(187)
Meanwhile, for E EST and j $ ST.
softmax (2)] T WK W
(++1) ((th") x²)) X (T()(++(()))((0)) (4)))) (an - ((S") - -ST))
(188)
softmax(x") W (++ W
(++1)
xt) softmax(a") T w
w x")
(1) (6+7) (e)
(ill + -1S*1 (ST
S (0 (t)il + -
|8T1
-1)
(187le* (18") -1 1ST
S
(0000-00))) K(L)
S ([S]
(|5"|-13 ))
(189)
where the second to last step is by the Mean Value Theorem with
x E [||q (t) (8 + +))l (+) (1 (t)) (6))))((())))
(190)
The same conclusion holds & (SF U S and j $ ST. Note that
9i (1 v(1 K()))g (t)
(191)
(t(t 1) v(1 K(())a(t))
(192)
(t + 1) = v(1 + Q(1))
(193)
r2(t + 1)
+ (())(()))2()
(194)
It can also be verified that this claim holds when / = 1.
E.3 PROOF OF CLAIM 3 OF LEMMA
The computation of the gradient of Wy is straightforward. The gradient would be related to Wo by their connections. One still need to study the influence of the gradient on different patterns, where introduce the discussion for the term (i('s..
38

---

Published as a conference paper at ICLR 2023
For the gradient of Wv. by 18 have
Loss
Loss X", v²) OF(X")
dWv
(F
dwv
nEB,
(1) WW
Wv X softtmax X W Woz 2 0]
B
Sh nEB,
LES"
Wo
softmax
[X"
W (ox x
(195)
Consider data {x",'' 3² where y" 1. Letle ST
softmax(x"" w W w x(¹) 2 (t)
(196)
SEST
Then for j E St, g [M].
1
Loss X" ))
W Zj
B
@w!()
R
1
(1)) (W (1)
Σ
softmax(:""
T w w at) W V 2 0)
B
S"
Oc.))
LES"
1=1
BES"
(1)
softmax(z" w
(t) Wh xx)x" 2
O
SES
Σ vi(t)Wo.s.
Σ (t(t))on..) T
iEW(()
SEW(()
(197)
lf E W(t). by the fact that S contributes more to Vi(t) compared to S/ for 2 3 and Assumption 3 we have
Vi(t) <
Σ
S
(f)
+
S
- (S
2B
S
nEB
(198)
S
Σ
Pa (1)
2B
a S
Similarly, if u(t).
S
Vi(t) R Σ 2B n
alSu Pr (t)
(199)
if i is unlucky neuron, by Hoeffding's inequality in 26 we have
1 Vi(t) > VB a
VM
(200)
1
> BBa
39

---

Published as a conference paper at ICLR 2023
For i W(O), we have
(b)
Σ
(0)
(.)
(s.)
b=1
jEW(b)
M
(o + T)M
S (b) - (r)
M
Ey
2BI
a|Sn
b= nEB
7 m2 ST
||p|| (r.))
Bt
a²|sTM|
nEB.
(0 + T)M
1
(1)pp(b)m
1
ny tbm SP
211
cm
aM (b(
pm. (b)))
2Bt
Bt
a² a²|sa1
1= nEB
nEB,
(201)
W WORL)
Σ
V
T
((.))
jEu(()
(202)
|S Pm (b)m
(o ()
p + T)lp
pf (o
6m
)
Bt
S aM
a²
nEB,
ptmmX / /
Wo
Σ JE(W())LE(())
vs(t
(203)
v 2
Hence,
If) E 8 for one n E [N].
w (((11) x/ w. x.j
ØL
WW
11
(204)
= - iEW(b)
Vi(i
Ois
(6)
Σ
Σ
Av((6)
00.
+ 21(
b=1 (gW(()
(2) jf E S2 we have
0L w 2j w x. (
w( 2
aww
(205)
=P Σ Σ W(()
(8)
Σ Σ b=1 igu((b)
Oie. 2,((t)
(b)
Oc.)
6=1 (EL(()
(3) jf 87/(5))) U U ) we have
W (++1) 2n
w. 10) x3
T(
OL awv
w((0) x"
4+1
(206)
Pk 17
W (b) Oc.))
+ zt(()
6=1 i=
Here
|=,(()) (a
(207)
for Note that this claim also holds when 1.
F OTHER USEFUL LEMMAS
Lemma 3. If the number of neurons m is larger enough such that m 2 M2 log N, the number of lucky neurons at the initialization [W(O)), [Z((0) satisfies (a M
(208)
[WO)), [2(0)] 2
M (1 Cm
(209)
π
40

---

Published as a conference paper at ICLR 2023
Proof:
Let fr be the angle between pt and the initial weight for one E [m] and all E [M.] For the lucky neuron i W(O), 0 should be the smallest among {8:)/1 with noise A6. Hence, the probability
of the lucky neuron can be bounded as
Pr 01
+ A
0 A0 S 2x, 2 251SM) M)
Pr
(1 + A0 < 01 - AB s 2m
(210)
1=2
2T 01 - 220 M-
2
where the first step is because the Gaussian W (0)
(s,) and orthogonal pt, E [M]] generate indepen-
dent ( (u,,, Pt. From the definition of W(0), we have
2 sin - & (a+
(211)
which implies
A0 1o +T)
(212)
for small a > 0. Therefore,
Pr (1E
2 1
2T
01
220
MM11d0 d0,
=
2T
2T
2T
2A8
2x
M
M
2m
1
(213)
1
A
yM
M
π
+)M (ar
M
π
where the first step comes from that 0, follows the uniform distribution on [0, 2T] due to the Gaussian initialization of Wo. We can define the random variable v such that
1, i E W(O),
U.
lo else
(214)
We know that v belongs to Bernoulli distribution with probability 1 (1 (a+7)M) By Hoeffding's inequality in 26 we know that with probability at least N-10
(1 (a T)M
log N
(a + T)M
log N m
+
(215)
M
M
Let m 2 0 log B), we have
[ =
2 M
m
(1 Em
(a +M
(216)
x
where we require
(a + T) S
(217)
M
to ensure positive probability in (11 Likewise, the conclusion holds for (((0).
Lemma 4. Let W(L) and U(I) be defined in Definition 2
We then have
W(0) S W(t)
(218)
u(0)) C (())
(219)
41

---

Published as a conference paper at ICLR 2023
long
B ((1)
(220)
Proof:
We show this lemma by induction. (1) For E W(0), by Definition than (a + r). Hence, we have
we know that the angle between W (0) (v.))
and p: is smaller
w (0) Oisy Pi( (a + r)) 2 w O p(1 + (a++))
(221)
for all p P/P1- (2) Suppose that the conclusion holds when / == 8. When / = + 1, from Lemma 2 Claim obtain
we
W T ((.)) -P
w(s) T (a.) :P
|S
((o (T)ss))) VNmi S
N Pn(s) 5
(222)
Su Pn(s)
B
nEB,
and
w O6.)
Wo... T .P)
.
7
|
((or + T))ns|(s) SP
ST (P)n(s))pp)) Ipl
Pn(s)
St1
Tm|sa
(223)
m
BEB,
S (s)
log m log [pl| +
y Elpel Bm
S M
B
Combining 222 and
223
we can approximately compute that if
B 1+ (a +T) 3² O(1,,
(a+ (
(224)
we can derive
w pi +)) w ( p(1 (T) Ou.)
(225)
Therefore, we have
W(0) S W(s 1)
(226)
In conclusion, obtain
W(0) S W(t)
(227)
for all 2 One can develop the proof for U(t) following the above steps.
G EXTENSION TO MORE GENERAL CASES
G.1 EXTENSION TO MULTII-CLASSIFICATEE
Consider the classification problem with four classes, we use the label y E [+1,-1]² to denote the corresponding class. Similarly to the previous setup, there four orthogonal discriminative patterns. In the output layer, 4io for the data (X".y") is changed into an R² vector ale for I E [| and i E [m]. Hence, we define
F(X")=
n a Relu Wo Wr softtaaxx T W Wkæx))
(228)
S"
Fi(X")
al Relu( X' W Wk*7))
(229)
Smh
1ES"
Fx(X")= S"1
Σ
Mary Relu Wo XX W W7TKaam x
(230)
1E8"
42

---

Published as a conference paper at ICLR 2023
The dataset D can be divided into four groups
D1 = =(x", yy" = (1,1))
D2 =(x", y²) = (1,-1)
(231)
D3 -(x", yy" "ly" (-1,
D4 =((",, "y" """""" = (-1,
The hinge loss function for data (X" y will be
Loss X" y) max (1- n" F ((
(232)
We can divide the weights Wo... (i [m]) into two groups respectively.
1
W1 = [il (1
(1,1)}
Vm
1
W2 = [1
(1,
11)}
Vm
(233)
W33 = (1/0
(-1,1)}
Vm
1
W4 (i) 040
(-1,-1)3
m
Therefore, for Wo, in the network 228 we have
ÖLoss X" yn)
OF (X"))
(X)
v
V/2
(234)
aWom.
OW WWo..
Won
where the derivation of (x."
3Wome.
and BWoza.)
(X")
can be found in the analysis of binary classification
above. For any i E W2, following the proof of Claim I of Lemma if the data (X" y²) E D2, have
ØLoss X", yy)
(X
OF
1 P2 -2-11(--1
2
Vi OW (Woa.))
2/2
Rix 1.
P2 mm
P2 (235)
aWou.
W Woza.)
vm
\m.
(wilt))
will) Oa,
P2
20
(236)
if (X", D1, we have
Loss X y")
2sx 1.
P1 + +1. (-
PP 0
(237)
aWo W
mm
vm
(w Oa.)
w((
pp ===
(238)
(B.)
if (x"."") (P)e D3. we have
ØLoss XX,, X",""" yn)
1
2SX -1
P3 + +1.(-
Pa ps
(239)
aWo
mm
mm
mm
(Wit+1)_ O
((() Jps 50 Ou,
(240)
(X","" E D4, we have
àLoss(X )
1
1
2X -1
ps 1. (-
)p4 =0 0 Vm
(241)
aWo...
vm
(W11+)) (
(t+1)
w
JP4 0
(242)
By the algorithm, Wo will update along the direction of p2 for i E W3. We can analyze Wv, Wx and Wq similarly.
43

---

Based on the modified assumption with H heads, the backbone the proof remains the same. Lucky neurons in Wo tend to learn (pI, PP. PDH)) and Pi.......... P.......PPP..... .P,,)) )T Hence, the properties of the Relu activation are almost the same as the single-head case because luck neurons Published as a conference paper at ICLR 2023
G.2 EXTENSION TO MORE GENERAL DATA MODEL
We generalize the patterns from vectors to sets of vectors. Consider that there are M (2 < M <
May m2) distinct sets (Mi, "Mr) where M₁ [uL,1 |4,,," ",ii Im 2 1. Mi. are sets of non-
M2 denote sets of discriminative patterns for the binary labels, and
discriminative patterns.
min KiSJSM, StStm.lshii
I| M3,0l 0
(243)
is the minimum distance between patterns of different sets. Each token ax" of X is a noisy version of one pattern, i.c.,
min JE[M]]bElmm Define that for I, s corresponding bi. b2, respectively,
| - Me, S
(244)
min je (e
||=|
x"'l S A,
(245)
we have 27 + A <
To simplify our theoretical analysis, one can similarly rescale all tokens little bit like in Assumption
3 such that tokens corresponding to patterns in the same pattern set has an inner product larger than
I. while tokens corresponding to patterns from different pattern sets has an inner product smaller
than À < 1.
Assumption can be modified such that
(0)
Mi P3 5
(246)
W (0)
14 93.bll 5
(247)
W (0)
Q 43, -Tybll 5
(248)
where Pisb 1 Pia for any i.j [M].
[im], and E [im]]. Likewise, Fj 1 qia for any
i.j [M]. b E [i], and a E [8m]. Ty.b 1 Ti,a for any i.j [M]. [5m], and E [im].
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. Then, we can follow our main proof idea. To be more
specific, for label-relevant tokens 27" by computing 125 and (167) W 27 w xc will grow in the direction of a fixed linear combination of It,1," It.... and 1t,,," Pr,,,.. The coefficient of the linear combination is a function of fractions of different pattern vectors JAL in Mr. One can still derive sparse attention map with weights of non-discriminative patterns decreasing to be close to zero during the training.
G.3 EXTENSION TO MULTI-HEAD NETWORKS
Suppose there are H heads in total. The network is modified
F(X")
Sm
1
Σ
(1)
Relu
Wo
|I h=1 SES*
Σ
Wv, x. softmax(x"" WT W Q. x())
Kr.
IES"
(249)
H
1
T Relu in
Wo
Σ Wv, x" softmax(x" WT, Kr Wo, xf))
Sm
TES"
h=l
S
where Wv, E Rm xd Wo = (Wo,, Wo.." h E [H].
Won) E RmxHMm, and Wo,, E Raxma for
One can make similar assumptions for W (0)
Note that (Pi.P2.... .... } needs to be changed into {Pr, Phay Phas J. and the set (Ph,, Phot Phss can vary for different h E [H]. is also the same for (qh, -ha:" -hh and (Th, "h That for different h [H].
W(0)) K,, and ((()) Vi as in Assumption
Qr
still activated by either of two label-relevant patterns with a high probability. In fact, can expect a more stable training process by multiple heads due to a more stable Relu gate for lucky neurons.
44

---

Published as a conference paper at ICLR 2023
G.4 EXTENSION TO SKIP CONNECTIONS AND NORMALIZATION
Consider a basic case where a skip connection is added after the self-attention layer. Let ma d. The network changed into
F(X")
Σ
alyRell
Σ
Wof"ssotmaxx""""
W Wo + xi))
(250)
SHh
LES"
SES
The assumption of w/ in Assumption I should be changed into
I[ + ΠH; - -Pll S
(251)
while the assumption of W and W K remain the same.
One can easily verify that the gradients of WK, We. and Wv for 250) are almost the same as those for (I) except for the Relu gate. The major differences come from the gradient of Wo. which also helps to determine the Relu gate. One needs to redefine
(1) W (1) X "oottmaa" ( w T W x0)) +
Σ
softmax(æ"
T W (1)
w 2t a())p + x(1) + Σ [W]() + +
(252)
sES:
n
Σ V.(t
Σ V((t)aW (n) (EW())
T
(0) Oc.
(EW()
for 1 E Sp. The inner product between the lucky neuron and the term Ejj1 (1) + xt) can still be upper bounded by the inner product between the lucky neuron and the term EaEs, softmax(>" w (
Q az Pr given good initialization of WK and Wo. Therefore, 250 can be analyzed following our proof techniques. For layer normalization, one usually use that approach to normalize each data. It is consistent with our normalization of az", which plays an important role our proof. By normalization, the training process becomes more stable because of the unified norm of all tokens.
45


Reviews:
[{"review_id": "xhm5_P1q28", "reviewer": "Reviewer_7JzE", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper gives a theoretical analysis of a simple (three layer) vision transformer network. The analysis is focused on training a ViT on a simple distribution of images constructed from label-relevant and label-irrelevant patches. Namely, there is a dictionary of M patches, and each patch in the input image is a noisy version of one of the patches. Only 2 out of the M patches are important for determining the label, which is the majority over the number of patches from these two classes of patches. The authors show that when training a ViT using SGD on this distributions, SGD converges to a solution with low error on the distribution. Furthermore, they analyze the sample and runtime complexity of SGD and show their dependence on different parameters of the distribution. The authors also show that the self-attention maps converge to sparse maps, which depend on the relevant/irrelevant patches. Finally, the authors complement their finding with different synthetic experiments.", "summary_of_the_review": "The paper is overall good, with novel contributions, but I still find some problems with the paper that could be improved (mentioned above).", "strength_and_weaknesses": "To my knowledge, this is the first work on theoretical analysis of learning with vision transformers, and one of the few works studying the theory of transformers in general. The fact that the authors analyze both optimization and generalization of transformers is unique, and allows the authors to draw novel conclusions on the theoretical properties of transformers, e.g. showing that the attention maps converge to sparse maps.\n\nThere are a few problems I still find with the paper:\n1. The theoretical results comparing training with and without self-attention maps are not convincing in my view. To establish a real \"separation\" between training with and without the self-attention map, showing that the former is indeed better, the authors need to show that training without self-attention requires more samples or training steps (i.e., giving a lower-bound on the sample/computational complexity). Instead, the authors show that using the same technical analysis leads to an inferior upper-bound when not using self-attention. While the authors also admit this in a comment, I think the presentation of the result still suggests that such a separation is established.\n2. I do not understand the results on token sparsification: How is the sparsification done? Do you assume prior knowledge about the distribution (namely, which patches are relevant/irrelevant)? Is the sparsification done before or after training? How do you make sure that the relevant patches are not removed?\n3. The theoretical results are limited to a very specific data distribution. While I do not necessarily see this as a major problem, since it allows the authors to show stronger results in this context, it would be good to understand whether the results can be generalized beyond this specific case. For example, can this result be extended to the case where there is more than one relevant basis patch for each of the two classes? Is the result on achieving sparse maps in the presence of irrelevant patches true in more general cases? Can we show, for a general distribution, that patches which do not affect the label get zero weight in the self-attention layer?\n4. It seems to me that Assumption 1 should hold with high-probability over the initialization, and also that \\sigma (the error of the initial model) can be bounded with high-probability. Why is this given as a separate assumption? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "bwtufS2t397", "reviewer": "Reviewer_sr6i", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper provides theoretical results on training three-layer ViTs for classification tasks. The authors quantify the importance of self-attention on sample complexity for zero generalization error, as well as the sparsity of attention maps when being trained by SGD. The authors then also show that token sparsification can improve generalization performance by removing class-irrelevant tokens and noisy tokens.  ", "summary_of_the_review": "Ultimately, this work provides valuable contributions towards understanding transformer architectures, for which there is not enough theory. While there are some concerns around reproducibility, the contributions of this work outweighs this minor issue. ", "strength_and_weaknesses": "Strengths\n- A valuable problem to study and valuable theoretical results.  \n- Experiments around token sparsification are compelling (Fig 6). \n\nWeaknesses\n- Only single head attention is studied, even though multi-head attention is typically used. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QjGvYSgLq3", "reviewer": "Reviewer_8Nb6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper theoretically studies the generalization of a multi-layer neural network containing a self-attention module (resembling a simplified vision transformer), trained with SGD on a dataset consists of a mixture of label-relevant and irrelevant tokens. This paper is the first to formally study the generalization of ViT-like neural networks. The analysis characterize how attention map evolves during training and the effects on sample complexity and reduced spurious correlation. ", "summary_of_the_review": "This paper is a theoretical study of the learning and generalization of a 3-layer networks inspired by vision transformer (ViT) models. This is one step towards, but far from a full theoretical understanding of ViTs. That being said, this seems to be one of the first papers making such attempts. I'm willing to raise my rating if the authors clearly address these limitations n the title and abstract, and if other reviewers who are more familiar with the theories acknowledge that there are novelties in the proof techniques and analysis. \n\n------\nThanks for the clarifications, and modification to the paper to make the topic of study more precise. I increased my rating.", "strength_and_weaknesses": "**Strength**\n\n1. This paper is among the first to formally study the generalization of ViT like neural network models.\n\n2. This paper characterizes sample complexity and evolution of the attention map for such models.\n\n**Weakness**\n\n1. The architecture studied in this paper contains a self attention module, but is still quite different from ViT models, even when restricting to one encoder/decoder block. In particular, there is no normalization layer, and the skip connections are not present. A closer simplified model would be four layers (1 attention layer, 2 fully connected layers for the MLP, 1 final frozen random prediction layer) with skip connections (and better, also with normalization layers). \n\n2. This paper, while making one step closer towards understanding ViTs comparing to previous theoretical papers studying general DNNs and CNNs, is still a big overclaim. The theoretical model is just one encoder block (which is also drastically simplified, see the previous point), and it relied on the inputs already having clearly separating class relevant and irrelevant token embeddings (the experiments also uses ViTs that are already pre-trained on other datasets). I think the title should at least say this is an understanding of *shallow* ViTs.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "5SwGjZ5fwg", "reviewer": "Reviewer_3WM1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper provides a sample complexity bound for Transformers, as well as the required number of SGD steps. These bounds are for the model to achieve 0 generalization loss (using the hinge loss).\n\nThe task is binary classification, where the label depends  only on the \"label-relevant patterns\" in the input. The bound suggests that the sample complexity is improved as we 1) increase the fraction of label-irrelevant patterns, or 2) reduce the initial error of the model parameters (how close are initial encoding to an orthogonal basis; see Def 1), or 3) reduce the noise in the tokens.\n\nThe paper also provides empirical results to corroborate the theory.", "summary_of_the_review": "This paper studies an interesting and challenging problem: the optimization of Transformers.\nOn the plus side, the paper is among the first to provide such results (there is a concurrent work: https://arxiv.org/abs/2210.09221) and such type of results is of great interests to the community.\n\nHowever, I'm currently concerned about the novelty of the techniques and the presentation.\nI'm willing to update my review after the rebuttal period & revisions.\n\n===========\n\nUpdate post-rebuttal: my concerns have been addressed by the author responses & the revision. I have raised my score accordingly.", "strength_and_weaknesses": "To my knowledge, this paper is the among the first to analyze the training dynamics of Transformers on binary classification. The data setup (e.g. patterns for multiple classes can co-exist in the same input sample) and proof ideas seem quite similar to those of Allen-Zhu and Li (2020), which is however not cited in the paper.\n\nThe main proof idea is to track the growth of a set of lucky neurons, and show that the attention weights concentrates on the label-relevant patterns. I didn't check the proof details too carefully since I find the current appendix hard to track; please see my comments on writing.\n\n The paper provides empirical results to corroborate the theory which is appreciated, though I think the empirical section needs some improvement; please see my comments below.\n\nQuestions:\n- \"token sparsification method\": what sparsification is used in Theorem 1 (i.e. how is $S^n$ defined)? Regarding designing better token sparsification method, how do we know what tokens are label-irrelevant in practice?\n- Prop 2: is there a finite step result, i..e at what rate (in terms of $T$) does the sum of the weights goes to 1?\n- Sec 4.1\n    - How are $\\delta, \\sigma$ enforced in practice?\n    - Sample complexity: why $10^{-3}$ as the threshold? Fig 3 seems to track till reaching a much lower loss.\n    - Fig 3-5: are these over 20 duplicates as well?\n    - Fig 5: is the shadow area standard deviation or standard error?\n    - For Fig 4 and 5, using a larger $M$ (currently $M = 5$) will make the trends clearer and the results more interesting.\n- Sec 4.2\n    - Fig 6: are all models trained till convergence? Will architectures other than Transformers follow the same trend as $\\alpha_*$ varies?\n    - Fig 6(a): currently the number of samples are not sufficient to get (near) perfect accuracy; what if we further increase the number of samples? Will the lines for different $\\alpha_*$ cross?\n        - A more direct comparison may be to plot the number of samples required to reach a certain accuracy (e.g. 99%), as we vary $\\alpha_*$, e.g. for $\\alpha_* \\in \\{0.1, 0.2, \\cdots, 0.9, 1\\}$; the choices of $\\alpha_*$ currently seems a bit arbitrary.\n\nComments:\n- The counting of layers is a bit non-standard.\n  - The model uses 1 attention layer w/ a 2-layer MLP, which is usually called a 1-layer Transformer (since a Transformer layer includes attention + MLP). Moreover, $W_O$ is usually considered as part of the attention layer, rather than an extra layer.\n  - The paper considers the model as 3-layer possibly by counting the weight matrices, however, there is no nonlinearity between $W_O$  and $W_V$ so they can be considered as 1 layer?\n- There are some sloppiness in the proof; for instance, $\\ll$ in eq (34) and $\\approx$ in eq (35) are not precisely quantified.\n\nReference: Allen-Zhu and Li 2020: https://arxiv.org/abs/2012.09816", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster