On the Convergence of Encoder-only Shallow Transformers
Authors: Yongtao Wu
Fanghui Liu
Grigorios Chrysos
Volkan Cevher
Venue: NeurIPS 2023 poster (2023)
Keywords: Transformer
convergence
scaling
initialization
over-parameterization
TL;DR: We study the convergence of encoder-only shallow Transformers, and provide theoretical insight on the choice of initializations as well as scalings.

Abstract: In this paper, we aim to build the global convergence theory of encoder-only shallow Transformers under a realistic setting from the perspective of architectures, initialization, and scaling under a finite width regime. The difficulty lies in how to tackle the softmax in self-attention mechanism, the core ingredient of Transformer. In particular, we diagnose the scaling scheme, carefully tackle the input/output of softmax, and prove that quadratic overparameterization is sufficient for global convergence of our shallow Transformers under commonly-used He/LeCun initialization in practice. Besides, neural tangent kernel (NTK) based analysis is also given, which facilitates a comprehensive comparison. Our theory demonstrates the separation on the importance of different scaling schemes and initialization. We believe our results can pave the way for a better understanding of modern Transformers, particularly on training dynamics.

On the Convergence of Encoder-only Shallow Transformers
Yongtao Wu LIONS, EPFL
Fanghui Liu" University of Warwick fanghui wwwwwarwick.ccmm uk
yongtao. wu@epfl.ch
Grigorios G Chrysos' LIONS, EPFL University of Wisconsin-Madison chrysos@wisc.edd edu
Volkan Cevher LIONS, EPFL volkan cewher@epfl.cc
Abstract
In this paper, aim to build the global convergence theory of encoder-only shal- low Transformers under realistic setting from the perspective architectures, initialization, and scaling under finite width regime. The difficulty lies in how tackle the softmax in self-attention mechanism, the core ingredient of Transformer. In particular, diagnose the scaling scheme, carefully tackle the input/output softmax, and prove that quadratic overparamcterization is sufficient for global convergence of our shallow Transformers under commonly-used He/LeCun initial- ization in practice. Besides, neural tangent kernel (NTK) based analysis also given, which facilitates comprehensive comparison. Our theory demonstrates the separation on the importance of different scaling schemes and initialization. We believe our results can pave the way for better understanding of modern Transformers, particularly on training dynamics.
Introduction
Transformers Vaswani al. 2017 have demonstrated unparalleled success in influential applica- tions Devlin et al. 2019] Brown et al. 2020) Wang et al.) 2018 Dosovitskiy et al. 2021 Liu et al., 2022b]. fundamental theoretical topic concerns the global convergence, i.e., the traiiing dynamics of Transformers, which would be helpful for further analysis. e.g., in-context learning [von Oswald al. 2022, Akyüüek et al. 2023). generalization (Li et al. 2023]. fact, even within simplified Transformer framework under certain specific regimes, the global convergence guarantees still remain an elusive challenge.
To theoretically understand this, let us first recall the exact format of the self-attention mechanism, the core ingredient of the Transformer. Given the input X E Rd, (d, is the number of tokens and d is the feature dimension of each token), a self-attention mechanism is defined as:
Self-attention (X) A a. (o (XW)) (XWx) ") (xW,) = a X Wx (XW.),
where is the row-wise softmax function, E RR+ is the scaling factor, and the learnable weights are Wo, WK. Wy E Idd.. with the width dm. Given X, the input of softmax depends on To W WK. k. including the scaling factor To and initialization schemes for learnable parameters, and thus determines the output of softmax and then affects the performance of Transformers in both
Work done at LIONS, EPFL.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

theory and practice. There are several scaling schemes in previous literature. For instance, given Wo and Wx k initialized by standard Gaussian, the scaling factor To is chosen by
To = 1/2 in the original Transformer Vaswani et al. 2017 each element in To W WK k is a random variable with mean and variance This scaling avoids the blow-up of value inside softmax as dm increases [Hron et al. 2020).
Q
TO dm This scaling stems from the neural tangent kernel (NTK) analysis [Jacot et al. 2018 a commonly used technical tool for convergence analysis of fully-connected (or con- volutional) networks under an infinite-width setting dm - Do. However, for Transformer, one uses this scaling under the infinite-width setting, then by the law of large numbers, we have lima,,-oo W Wx] = O. As result, the input of softmax is zero and the softmax degener- ates to a pooling layer. That means, the non-linearity is missing, which motivates researchers to carefully rethink this setting.
For instance, under the TD = dmm setting, Yang 2020) use the same query and key matrices to prevent the softmax from degenerating into a pooling layer. Besides, to avoid the analytic difficulty of softmax due to the fact that each element of the output depends on all inputs, [Hron et al. [2020)
substitute softmax with ReLU under the To = dm 1/2 and infinite width setting for simplicity.
Clearly, there exists gap between theoretical analysis and practical architectures the of softmax, and accordingly, this leads to the following open question:
How ensure the global convergence of Transformers under realistic setting?
The primary contribution of this work is to establish the convergence theory of shallow Transformer under realistic setting. Despite its shallow and encoder-only architecture, Transformer model captures all the fundamental components found on typical Transformers, including the self-attention mechanism with the softmax activation function, one feedforward ReLU layer, one average pooling layer, and linear output layer, cf. Eq. 1.2). We adopt the To = dm2/2 scaling under the finite-width setting and compare the results of LeCun/He initializations, which are commonly used in practical applications. Besides, the convergence result under the To = d setting (as well as the NTK based analysis) is also studied, which facilitates comprehensive comparison. Our theoretical results demonstrate notable separations among scaling settings, initializations, and architectures as below:
1/22
The global convergence can be achieved under both less, as suggested by our theory: for small dmm, there no significant difference for these two scaling schemes on the convergence; but for a large enough dm. the TO 1/21/2 scaling admits faster convergence rate of Transformers than that with To d Interestingly, under this To d setting, our theory also demonstrates the separation on the convergence result, depending on whether the input formed along sequence dimension
Scaling:
1/2 and To d Neverthe-
1) or embedding dimension (d, 1).
Initialization: Under LeCun and He initialization, our shallow Transformer admits a faster conver- gence rate than the NTK initialization. This could be an explanation for the seldom usage of NTK initialization for Transformer training in practice.
Quadratic over- parameterization is enough to ensure the global convergence of our shallow Transformer. As comparison, if the self-attention mechanism substituted by feed-forward ReLU layer, our shallow Transformer is close to three-layer fully-connected ReLU neural networks to some extent, requiring cubic over-paramettrization for global convergence.
Architecture:
We firmly believe that our theoretical analysis takes significant step towards unraveling the mysteries behind Transformers from the perspective of global convergence. We hope that our analytical framework and insights on various initialization and scaling techniques would be helpful in training modern, large-scale Transformer-basee models [Radford et al. 2018 Brown et al. 2020
2 Related work
Self-attention, Transformer: Regarding training dynamics, Snell et al. 2021 explain why single-
head attention focuses on salient words by analyzing the evolution throughout training. Hron et al.)
2020 show that the output of Transformer converges to Gaussian process kernel and provide the

---

NTK formulation of Transformer. Recently, Li et al. Transformer study its generalization property under good initialization from pretrained model. The separation between the Transformer and CNN is recently explored. Jelassi et al. [022] provably demonstrate that Vision Transformer (ViT) has the ability to learn spatial structure without additional inductive bias such as the spatial locality in CNN. [Chn et al.) [2022] study the loss landscape of ViT and find that ViT converges at sharper local minima than ResNet. Park and Kim [022] show that ViT is a low-pass filter while CNN is a high-pass filter, thus, these two models can be complementary.
[2023] provide sample complexity of shallow
The NTK was introduced by Jacot et al. 2018 connect the infinite-widde width neural network trained by gradient descent and the kernel regression. The roles of such kernel include analysis of the training dynamies of the neural network in the over-parameterization regime [Allen-Zhu et al.] 2019aa Chizat et al. 2019 Du et al.) [ Zou et al. 2020]. The global convergence, generalization bound, and memorization capacity largely depend on the minimum eigenvalue of the NTK (Cao and Gu 2019] Zhu et al.) 2022 Nguyen et al., 2021, [Bombari et al.) 2022]. Even though the NTK extended from FCNN to several typical networks including Trans- former [Trer et al. 20200 Huang al., 2020 Arora et al. 2019b Alemohammad et al. 2 Nguyen and Mondelli 2020 it has not been used to analyze the global convergence of Transformer. On the other hand, the stability of the tangent kernel during training is required when connecting to kernel regression, but such stability can not be explained by the phenomenon of lazy training [Chizat et al., 2019), which indicates small change of the parameters from initialization. The hessian spectral bound is the main reason for the stability of kernel, as mentioned in Liu et al. (2020].
NTK, lazy training, Hessian:
pherpparameteiiztionn for convergence analysis: Due to peeeraaramttrizaiionn neural networks (NNs) can fit arbitrary labels with zero training loss when trained with (stochastic) gradient descent (SGD), both theoretically [i and Liang] 2018). Du et (2019b) and empirically (Zhang et al.) [017]. This leads to an interesting question in theory: how much overparameterization
enough to ensure global convergence of NNs? A common recipe for the proof of global convergence relies on the
variant
Lojaieewwii condition [Polyak] 1963 Liu et al. 2022a], NTK [Du et al., 2019b a
Zou and Gu. 2019 Allen-Zhu et al. 201 the minimum eigenvalue of the gram matrix Nguyen
2021 Bombari et al.] 2022). J. In Appendix
B.3 we provide a comprehensive overview of a recent line of work that improves the over- paramctenizaiiionn condition for ensuring the convergence of NNs.
However, the over over-parameterization condition for Transformer to achieve global convergence remains elusive from existing literature and we make an initial step towards this question.
Problem setting
This section includes the problem setting with notations and model formulation of the shallow Transformer that is studied in this paper.
3.1 Notation
Vectors (matrices) are symbolized by lowercase (uppercase) boldface letters, e.g., w. W. We use || 11I and || lla to represent the Frobenius norm and the spectral norm of matrix, respectively. The Euclidean norm of vector symbolized by || ||2.. The superscript with brackets is used to represent the element of a vector/matrix, e.g., ( is the 7th element of w. The superscript without brackets symbolizes the parameters at different training step, e.g., 0'. We denote by [N] = {1, .N for short. We use mmin(.) (-) and ) to represent the minimum singular value and minimum eigenvalue of matrix. The NTK matrix and hessian matrix of the network are denoted by K and H. respectively.
The order notation, e.g., O, á2, omits the logarithmic factor. More detailed notation can be found in Table of the appendix.
Let X C Rd, xd be compact metric space and Y C R, where d is the dimension of each token, d, is the total sequence length of the input. The training set {(X,, yn) from an unknown probability measure on X x Y. In this paper, we focus the regression task by employing the squared loss. The goal of our regression task is to find a hypothesis, i.e., a Transformer
) assumed to be iid sampled
f: X -> Y in our work, such that f(X; 0) parameterized by 0 is a good approximation of the label y E Y corresponding to a new sample X E X. We use a vector to denote the collection of all learnalle parameters.
3

---

Algorithm 1: Gradient descent training
Table I: Common initialization methods with their variances of Gaussian distribution and scal- ing factors. The choice of Tr dm 1/2 is based on standard NTK initialization on prior litera- ture (Du et al. 2019b].
Input: data (Xn, Yn ba)n- step size yy Initialize weights follows:
0 = {WQ. WX
K WV, W
Q.
for = 0 to t' 1 do
W
Q W . VWo ((0),
Init.
1/v
1QQ
TI
W wi++ = WK K -7 W {(0),
K
LeCun
dm
d-1
d-1
d-1
1
WT+1 W. -7 Vw, {(0'),
He
2dm
2d-1
2d
2d-1
1
W W 7 Vw 1
NTK
1
1
dm
dm
1/2
end for
Output: the model based on 0
3.2 Model formulation of shallow Transformer
Throughout this work, we consider the encoder of Transformer, which can be applied to both
regression and classification tasks (Yüksel al., 2019 Dosovitskiy et al. 202T]. Given an input X E Rd xd the model is defined as below:
1 Self-attention X) 0s (m) ( (XWx) (XWT).
(1.1)
A T, WH),
0g - 4A22,
[(x;)) - ag wo,
(1.2)
where the output fX::) E R, To and TI are two scaling factors. The ingredients of Transformer with width dm are defined as follows:
A self-attention mechanism (Eq. (1.I) )) og is the row-wise softmax function; the learnable pa- rameters are Wo, WK. Wy E Rdm xd, We employ Gaussian initialization W (if)
Q ~O, mo)..
w ~ N(0 nk), w(()) ~N(0, ny) with i E [dm] and j E [d]. Refer to Table for typical
initialization examples.
A feed-forward ReLU layer (in Eq. 1.2b): a. is the ReLU activation function; the learnable parameter is WH
H E Rdm.dd., Following Yang et al. [2022]. we combine Wy and WH и together (by setting WH II) for ease of the analysis. Note that it does not mean its training dynamics are the same as the joint-training of these two adjacent matrices.
An average pooling layer (in Eq. (.2b): ;p indicates the column-wiss average pooling. Note that the average pooling layer is applied along the sequence length dimension to ensure the final output is a scalar, which is commonly used in practical Vision Transformer theoretical analysis (Dosovitskiy
et al. 2021
Yang 2020
An output layer (in Eq. 1.2 with learnable parameter wo E Rdm. initialized by w (O no).
Proper initialization and scaling are required to ensure the convergence and learnability, as seen in previous work (Jacot et al. 2018 Tirer et al., 2022 Lee et al. [019]]. For our convergence analysis, we consider standard Gaussian initialization with different variances and different scaling factor that includes three typical initialization schemes in practice. In Table we detail the formula of LeCun initialization, He initialization, and NTK initialization.
Remarks:
Given N input samples (X,, J the corresponding ground truth label, the final output of network, and the output of the last hidden layer, are denoted by:
Un I RA f(0)= (f(x,; 0)1 E RRN
(0) A (as(X,, 0) E RNxx
We consider standard gradient descent (GD) training of Transformer, as illustrated in Algorithm Here the squared loss is expressed E(0) } [0] y
4 Main results
In this section, we study the convergence guarantee of Transformer training by GD under the squared loss. Firstly, provide general analytical framework in Section 4.1 covering different initialization

---

schemes, where we identify the condition for achieving global convergence. Next, we validate these
conditions for several practical initialization schemes under To dm in Section 4.2 and To d, in Section 4.3 respectively. We include NTK-based results for self-completeness. Discussion on the convergence results with different scalings, initializations, and architectures is present in Section 4.4
4.1 General framework for convergence analysis
Before presenting our result, we introduce a basic assumption.
Assumption 1. The input data is bounded X S Vd,CC, with some positive constant Oz-
Remark:
This assumption is standard as we can always scale the input. The embedding of token is usually assumed to have unit norm [Li et al. 2023), which is unrelated to d.
Now ready present proposition. with the proof deferred to Appendix C.2
Notice
that our proposition holds with high probability under some conditions, which depend on certain initialization schemes and scaling factors. Since our proposition is devoted unifying analysis framework under various initialization schemes, we do not include specific probabilities here. The validation of the required conditions and probability is deferred Sections and and 4.3 respectively.
Proposition 1. Consider the Transformer with dm 2 N. Let Co. Ck. constants, and define the following quantities at initialization for simplification:
CQ.Ck. Cv, Co be some positive
The norm of the parameters:
dy - ||@ll, 1 + Cq, dK A
[ +CK, dv - | Ov, do - 111188 02 Co.
Two auxiliary terms: p a 2/2
T C2 = 2 (1 47/C (Xg xxx))
Under Assumption
we additionally assume that the minimum singular value of F iee, a 2
[FO pre) satisfies the following condition at initialization:
mmin
a 2 ma
( 2m0t,,2moo.dddee 4 2 2 2000
V2((0%) (2)
a 2 If the step size satisfies 7 S. L/C with a constant C depending on lλq, Ax Av, ((0%) then GD converges to a global minimum as follows:
(3)
E(O") - (1 (1-7-²) a²
2
{(0)),
AI 0.
The parameter in Eqs. (2) and B controls the convergence rate of global convergence, and the condition will be verified in the next subsection. The step-size y is inversely proportional to N1/2 due to the construction of C Appendix C.2
Remark:
Proof sketch: The main idea of our convergence analysis is based on the variant of Polyak- Lojasiewicz (PL) inequality (Polyak] 1963, Nguyen| 2021]. i.c., [[V)(0)] 2Amin F ))))) Thus, if the minimum singular value of Fpre is strictly greater than then minimizing the gradient on the LHS will drive the loss to zero. To this end, Proposition. [can be split into two parts. First, by induction, every time step, each parameter in the Transformer can be bounded w.h.p: the minimum singular value of Fpre is bounded away for some positive quality at the initialization point. Secondly, we prove that the Lipschitzness of the network gradient, which means the loss function is almost smooth. Combining the above two results, the global convergence can be achieved. Furthermore, based on different initialization and scaling schemes, we are able validate Eqs. (2) and (3) via the spectral norm estimation of λq. Ax. Xv. to and positive lower bound for Fe the following section.
2Amin (K) (K)(() Z
4.2 LeCun and He initialization under the d setting
Here we aim to show that, under the LeCun/He initialization with the setting of d/, the conditions in Eqs. (2) and (3) will be satisfied with high probability and scaling schemes in Table I and hence
5

---

lead to global convergence. To derive our result, we need the following assumptions on the input data regarding the rank and dissimilarity of data.
Assumption 2. We assume that the input data X has full row rank.
This assumption requires that each row X is linearly independent for any i E [da]. which is fair and attainable in practice. For example, for language tasks, even though there might be some repeated token in X, each row in X can be uncorrelated when added with positional embedding. Similarly, in image tasks with Visual Transformer, the raw image is grouped by patch and mapped via a linear layer to construct X, thus each row in X can be uncorrelated.
Remark:
Assumption 3. For any data pair (X X.)), with # n' and n'' E [M] then we assume that P (I(X,, X,, xx) 2 t) I exp(- exp(-t") with some constant > 0.
Remark: We discuss the rationale behind this assumption:
The idea behind Assumption 3 is that different data points admit a small similarity. To be specific, for two data points X,, and Xn, with * n', their inner product reflects the similarity of their respective (empirical) covariance matrix. We expect that
1.0
0.8
Probability
this similarity is small with high probability. The spirit of this assumption also exists in previous literature, e.g., Nguyen et al. [.02]] The constant determines the decay of data dissimilarity. larger results in less separable data. Our assumption has no requirement on such that can be small enough, which allows for general data distribution.
0.4
0.2
0.0
0.00
0.25
0.50
0.75
Verification of assumption. Here we experimentally validate this assumption under a standard language IMDB dataset [Maas] Figure 1: Validation of Asm.
et 2011]. We randomly sample 100 (normalized) sentences with embedding and plot the probabil- ity (Irequency) of XX, X,,) 2 t) as increases. We repeat it over 10 runs and plot the result in Figure 1 We can observe exponential decay as t increases, which implies that assumption is fair. Besides, the validation of Assumption on image data, e.g., MNIST by ViT, is deferred to Appendix E.2
Now we are ready present our main theorem under the dm setting. The proof is deferred to Appendix C.3
Theorem 1. Under the setting of LeCun/He initialization in Table Aand Assumptions to suppose
dm 2 d, and given To = 1/2 1/2 dm = (((N'), then with probability at least
Se dm//2
8- exp (-S((N - 1) 1)= for proper the GD training of Transformer converges global minimum with sufficiently small step size as in Eq.
The probability relates to several randomness sources, e.g., data sampling, dissimilarity of data, and parameter initialization. The quantity exp (-N((N - 1) "d )) can be small for a small enough è as discussed in Assumption 3 Further discussion on our result refers to Section 4. for details. Besides, our proof framework is also valid for the To = d setting. We demonstrate that dm = ((N2) suffices to achieve global convergence, see Appendix C.3 for details.
Remark:
Proof sketch: To check whether the conditions in Eqs. (2) and (3) hold, the key idea is to provide the lower bound a. Then we upper bound λQ. Ax. ^v, based on concentration inequalities to upper bound the initial loss, one key step is to utilize Gershgorin circle theorem [Gershgorin] 1931 to provide a lower bound for a. Lastly, we plug these bound into the condition Eqs. and in order to obtain the requirement for the width dm.
4.3
NTK initialization under the setting
The NTK theory, as a representative application of the To d=11 setting, can be also used for analysis of training dynamics. We also include the NTK results in this section for self-completeness. this section, we first derive the limiting NTK formulation of Transformer under the d scaling scheme, and then show the global convergence of Transformers.

---

Lemma 1. Denote O' =: = 1
d. "d.
Id.....
N.,, E RNxd then the limiting NTK matrix
K E RRNN of TTTransfommer under the NTK initialization with To = =d has the following form:
K - d. (a, (D"")) or (4'' + df (o. (v' ($'w) (w))') ( (-'
Remark: The formulation of 4' implies that the self-attention layer degenerates as 1 d. -d,xx, xd, XWV. ie., the dimension missing effect as mentioned before.
Now we are ready to present our convergence result under the dm scaling with the proof deferred to Appendix C.4
Theorem 2. Under the setting of NTK initialization Table and Assumptions 7 to suppose dm (,)), then with probability least -emth. d exp (-S((IV - 1) d the GD training of Transformer converges to global minimum with sufficiently small as in Eq.
The overall idea is the same as the proof of the previous theorem, i.e., we need to provide the lower bound of a. However, in this case, the limit for the output of softmax exists so that we can apply concentration inequality lower bound the a. Lastly, we plug these bound into the condition Eqs. 2 and (3) in order to obtain the requirement for the width dmn-
Proof sketch:
Besides, the stability of NTK during training allows to build a connection on training dynamics between the Transformer (assuming a squared loss) and the kernel regression predictor. Next, in order show that the NTK is stable during GD training. below we prove that the spectral norm of Hessian is controlled by the width.
Theorem (Hessian norm is controlled by the width). Under Assumption I and scaling To = d given any fixed R O, and any O B(O, R) {0 |10 00 I R), O the weight at initialization, then with probability at least 1 Sc dmm/ the Hessian spectral norm of Transformer
satisfies: (H'0)) 12 S O (mmm" 1/2
Remark: By Liu et al., 2020 Proposition 2.3]. the small Hessian norm a sufficient condition for small change of NTK. Thus, Theorem 3] can be an indicator for the stability of NTK. Besides, Theorem 3 supplements the result in [Park and Kim 2022 which exhibits empirically a relationship
between the Hessian norm and the width but lacks theoretical proof.
4.4 Discussion convergence results
We compare the derived results under different scaling schemes, initializations, and architectures.
Comparison of scaling schemes: The global convergence can be achieved under both TO dm 1/2 and To d Nevertheless, as suggested by our theory, for small dm, there is no significant difference between these two scaling schemes on the convergence; but for a large enough dm. the
To d,, scaling admits a faster convergence rate of Transformers than that of To d due to more tight estimation of the lower bound of a, see Appendix C.6 details. The intuition behind the lower convergence rate under the setting of large width and To d is that the input of softmax is close to zero such that softmax roughly degenerates as a pooling layer, losing the ability to fit data. This can be also explained by Lemma T from the perspective of dimension missing: self-attention (X) degenerates as d, ld, xd, XW..
The result under the To = d setting requires weaker over-parameeterizatonn than the TO = dm 1/2 setting. Nevertheless, we do not claim that To d is better than To dd,, This is because, under the pver-parameterization regime, the scaling To = d makes the self-attention layer close to pooling layer. This analysis loses the ability capture the key characteristics of the self-attention mechanism in Transformers.
Note that the lower bound of the minimum eigenvalue is in the constant order, which is tight (since it matches the upper bound). Based on this, by studying the relationship between dm and Ag. we can prove that quadratic (cubic) over- prraattraettiiiitoonnn is required for d (dm scaling. This quadratic over-parameterization requirement could be relaxed if a better relationship is given while it still unclear and beyond our proof technique.
7

---

Comparison on initializations: Though our results achieve the same convergence rate under these initialization schemes, we can still show the separation on a that affects the convergence Eq. ( among these initialization schemes. To be specific, we verify that under LeCun and He initialization, the lower bound of a tighter than that of NTK initialization, and hence admits faster convergence, see Appendix C.5 for further details. This can be an explanation of the seldom usage of NTK initialization in practice. Besides, the NTK initialization scheme allows for a larger step size than LeCun/He initialization for training. That means, if a is the same in these three initialization schemes, we usually choose large step size under the NTK initialization scheme, see Appendix E.T
Comparison on architectures: Note that the Transformer defined in Eq. 1.2 includes a self- attention layer, feedforward ReLU layer. and an output layer. Our result proves that cubic (quadratic) over-parametterization condition is required for dm (dm=)) under LeCun initialization. As comparison, three-layer fully-connected ReLU network under LeCun initialization requires dm (2(N3) [gguyen 2021
The aforementioned result holds for matrix input. Although not as frequent, some data inputs are naturally in vector form. Two ways to feed the input into Transformer are either formulating along sequence dimension (d = 1) or along embedding dimension (d, = 1). The following result shows the separation under these two settings to understand when the Transformer works well or not regarding the input.
Corollary 1 (Convergence of vector input). Consider LeCun initialization with To d scaling,
given vector input æ E Rd, if one feeds the input to Transformer by setting dg = 1, = d, then training with GD can converge to a global minimum. On the contrary, if one sets d, = d, d = 1, the conditions in Eqs. 2 and 3 do not hold, the convergence can not be guaranteed by our theory.
Remark: Such a result is motivated by considering least squares. Specifically, given input X E RNXI, then the NTK XX is rank-one matrix. As result, when the augmented matrix [X, y] is not rank-one, then there is no solution to the linear system so that GD training can not converge to zero loss. The empirical validation can be found in Figure
Technical difficulty. The technical difficulty of our analysis includes handling the softmax function and scaling schemes beyond NTK initialization. Previous convergence analysis, e.g., (Du et al. 2019b Nguyen| [02T] cannot be applied to our setting because the following issues. First, different from classical activation functions, e.g., ReLU, in softmax each element of the output depends on all input. To tackle the interplay between dimensions, we build the connection between the output of softmax and X to disentangle the nonlinear softmax function. By doing so, lower bound on the minimum singular value of Fpre in Proposition. can be well controlled by XX and the output of softmax.
Regarding different initializations and scaling, previous NTK-based analysis is only valid under the d setting (the softmax degenerates to an all-one vector) but is inapplicable to the realistic 1/,1/22 setting, as discussed in the introduction. To tackle this issue, we analyze the input/output of softmax under LeCun/He initialization and identify the optimization properties of the loss function for global convergence under the finite-width setting.
5 Experimental validations
Our experiments are organized as follows: In Section 5.1 we conduct experiments with the model Eq. 1.2) on synthetic data and study the training dynamics. Next, we show convergence results on ViT Dosovitskiy et al. [202 on the standard MNIST dataset in Section 5.2 Additional results and detail on the experimental setup are deferred to Appendix D
5.1 Fitting synthetic data
In this section, we corroborate our thcoretical findings on the synthetic data. We generate 100 data points where the input X R010X100 is sampled from standard Gaussian distribution. The corresponding label y E R is generated from standard Gaussian distribution. Squared loss is selected as the criterion. We apply gradient descent on the shallow Transformer defined in Eq. (I.2) with LeCun initialization and To = 1/2 for 400 epochs with fixed step size 7 = 1. We test different widths of the network including dm = {10, 100, 1000, 4000} and plot the training progress in
8

---

4,-10
dc-100
woyhhs-
Bocls
dy-10
1,0
1000
4000
Epects
Spochs
(a) Convergence curve.
(b) Weight movement.
(c) Kernel distance.
Figure 2: Visualization of the training process of Transformers with LeCun initialization and To dm 1/2 scaling synthetic data. (a) Linear convergence. (b) Rate of change of the weights during training. Observe that the weights change very slowly after the 50th epoch. (c) Evolution of the NTK during the training. The result mirrors the plot (b) and demonstrates how the kernel varies
significantly at the beginning of the training and remains approximately constant later. As the width increases, the empirical NTK becomes more stable.
Figure The result shows that the training can converge except for the case with sub-linear width, see the Tinear convergence rate in Vigure[alann and the small movement of the weight in Figure For these cases, the weights do not change much after 50 epochs while the losses are still decreasing. In Figure 2c] we keep track of the evolution of NTK along each epoch. Specifically, the kernel distance is given by:
Tr KKT
Distance K, K
=
Tr (KK ) Tr KK
which quantitatively compares two kernels by the relative rotation, used in Fort jet al.] [2020]. Figure 2c] shows that the kernel changes rapidly at the begin- ning of training while being approximately constant at later stages. The experi- ment with TO Additionally, the experiment with other initialization schemes deferred to Appendix [.. where we observe that NTK initialization indeed yields slower convergence than LeCun/He initialization, which is consistent with our theoretical finding.
d= which is deferred to Appendix E. obtains similar results.
Transtomer de d. 100 Transtormer. di 100 w,
roat 105
Next, in order to verify corollary T we feed vector input In E R00 into Transformer by setting either under LeCun
de = d, = 100 or d, = 100, initialization with To d= Figure shows that the training of Transformer with d, = 1, 100 is similar to that of two-layer FCNN with the same width d,, 100. However, the case of d, 100, d = fails to converge, which is consistent with our theoretical finding.
Epocts
Figure Convergence result on syn- thetic data with vector input.
5.2 Fitting real-world dataset
Beyond synthetic data, in this section, we examine the convergence performance of Vision Transformer (VIT) on classification task on MNIST dataset [LeCun et al. 1998), which includes 10 classes of images with size 28 28. We use single layer and single head VīT. The dimension of d 64, We change the dimension the query, key, and value from 16 to 1024 and 16384. The network is optimized with SGD with step size 0.1, and momentum 0.9 for 50 epochs. We repeat the experiment for three runs. We present the convergence results over runs in Figure 4a which
shows that when the width is smaller, e.g., 16, both dim and To d -1/2 scaling have similar

---

Before training
After training
Training
Epochs
(a) Convergence
(b) Attention map, dm 16384.
Figure 4: Convergence curve on MNIST dataset with different scaling schemes and different widths (a). Visualization of attention map in (b).
convergence results. However, as the width increases to 1024 and 16384, the To dm setting admits a slower rate than that тy = dm 1/2 especially a extremely slow rate under d,, 16384. This is consistent with our theoretical result the dimension missing effect such that the Tu dm setting makes Transformer difficult to fit data. Additionally, we visualize the attention map in Figure 4b) i.c.
the output of softmax in the self-attention layer before training and after training under TO dm 1/2 setting. We could see that the self-attention layer changes a lot during training.
Conclusion
We present comprehensive analysis of the global convergence properties of shallow Transformer under various scaling and initialization schemes. Regarding scaling schemes, for large width setting, the difference convergence between To d/2 and To d can be demonstrated both theoretically and empirically, Our theory is able to explain this in dimension missing view. Regarding initialization schemes, our theory prefers using LeCun and He initialization in Transformer training, which allows for a faster convergence rate than NTK initialization. Though our analysis is limited to shallow Transformers, we believe our framework can be extended to deep Transformers. We provide further discussion into this extension in Appendix C.7 Exploring the convergence properties of deep Transformers indeed an intriguing avenue for future research.
Acknowledgements
We thank the reviewers for their constructive feedback. We thank Zhenyu Zhu for the discussion and help this work. This work has received funding from the Swiss National Science Foundation (SNSF) under grant number 200021 205011. This work supported by Hasler Foundation Program: Hasler Responsible Al (project number 21043). Corresponding author: Fanghui Liu.
10

---

References
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations (ICLR), 2023.
Sina Alemohammad, Zichao Wang, Randall Balestriero, and Richard Baraniuk. The recurrent neural tangent kernel. In International Conference Learning Representations (ICLR), 2021.
Zeyuan
Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning (ICML), pages 242-252.
PMLR, 2019a.
Zeyuan Allen- Yuanzhi Li, and Zhao Song. On the convergence rate training recurrent neural networks. In Advances in neural information processing systems (NeurIPS), volume 32, 2019b.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference Machine Learning (ICML), 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems (NeurIPS), 32, 2019b.
Simone Bombari, Mohammad Hossein Amani, and Marco Mondelli. Memorization and optimization in deep neural networks with minimum paeeeerramttiiiiionnn. In Advances in neural information processing systems (NeurIPS), 2022.
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 (NeurIPS), 2020.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in neural information processing systems (NeurIPS), 2019.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In International Conference on Learning Representations (ICLR), 2022.
Zixiang Chen, Yuan Cao. Difan Zou, and Quanquan Gu. How much over- preeeerramtttiinatoon is sufficient to learn deep (lu) networks? In International Conference on Learning Representations (ICLR), 2021.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems (NeurIPS), 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguisties: Human Language Technologies, Volume (Long and Short Pupers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 00.18653/v//119 1423.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference on Learning Representations (ICLR), 2021.
Simon Du Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning (ICML), 2019a.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations (ICLR), 2019b.
11

---

Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self- attention mechanisms. In International Conference on Machine Learning (ICML), pages 5793-5831, 2022.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Advances in neural information processing systems (NeurIPS), 2020.
Semyon Aranovich Gershgorin. About delimiting the intrinsic values of matrix. Proc. Russian Acad. Sci, page 749-754, 1931.
Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. Annals of Statistics, 2022.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning (ICML), 2020.
Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks generalize better than deep feedforward networks? neural tangent kernel perspective. Advances in neural information processing systems (NeurIPS), 33:2698-27000 2020.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances in neural information processing systems (NeurIPS), 2018.
Samy Jelassi, Michael Eli Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Alice H. Oh, Alckh Agarwal, Daniclle Belgrave, and Kyunghyun Cho, editors, Advances in neural information processing systems (NeurIPSS, 2022.
Yann LeCun, Léon Bottou, Yoshua Bengio, Patrick Haffner, et al. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2324, 1998.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems (NeurIPS), 2019.
Hongkang Li, Meng Wang. Sijia Liu, and Pin- Yun-Yuee Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In International Conference on Learning Representations (ICLR), 2023.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems (NeurIPS), 2018.
Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel "ridgeless" regression can generalize. The Annals of Statistics, 2020.
Zhenyu Liao and Romain Couillet. On the spectrum of random features maps of high dimensional data. In International Conference on Machine Learning, pages 3063-3071. PMLR, 2018
Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances in neural information processing systems (NeurIPS), 2020.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 2022a.
Fanghui Liu, Zhenyu Liao, and Johan Suykens. Kernel regression in high dimensions: Refined analysis beyond double descent. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2021.
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Conference on Computer Vision and Pattern Recognition (CVPR), 2022b.
12

---

Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142-150, Portland, Oregon, USA, June 2011. 1. Association for Computational Linguisties.
Leon Mirsky, trace inequality of john von neumann. Monatshefte für mathematik, 79:303-306, 1975.
Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Mem- orization and generalization under lazy training. The Annals of Statistics, 50(5):2816 2847, 2022.
Quynh Nguyen. On the proof of global convergence of gradient descent for deep relu networks with linear widths. In International Conference on Machine Learning (ICML), 2021.
Quynh Nguyen, Marco Mondelli, and Guido F Montufar. Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep relu networks. In International Conference on Machine Learning (ICML), 2021.
Quynh N Nguyen and Marco Mondelli. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in neural information processing systems (NeurIPS), 2020.
Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameteriiatton: Global con- vergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 1(1):84-105, 2020.
Namuk Park and Songkuk Kim. How do vision transformers work? In International Conference on Learning Representations (ICLR), 2022.
Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal Vychislitel'nu noi Matematiki Matematicheskoi Fiziki, 3(4):643-653, 1963.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing with unsupervised learning. 2018.
Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiw:2103.07601 2021.
Chaehwan Song, Ali Ramezani- Thomas Pethick, Armin Eftekhari, and Volkan Cevher.
Subquadratic overparameterization for shallow neural networks. In Advances in neural information processing systems (NeurIPS), 2021.
Tom Tirer, Joan Bruna, and Raja Giryes. Kernel-based smoothness analysis of residual networks. arXiv preprini 2009 1000.000000 10008, 2020.
Tom Tirer, Joan Bruna, and Raja Giryes. Kernel- based smoothness analysis of residual networks. In Mathematical and Scientific Machine Learning, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems (NeurIPS), 2017.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices, page 210-268.
Cambridge University Press, 2012. doi: 10. 10.1017/CB09780S 11779/000119448880000
Roman Vershynin. Hiigh-Dieensionl Probability: An Introduction with Applications in Data Science. Cambridge University Press, 2018
Johannes von Oswald, Eyvind Niklasson, Eutore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint 172.77221, 2022.
Martin J Wainwright. High-dimenssonal statistics: non-asymptotic viewpoint, volume 48, Cam- bridge University Press, 2019.
13

---

Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Conference on Computer Vision and Pattern Recognition (CVPR). 2018.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14542 2020.
Yongyi Yang, Zengfeng Huang, and David Wipf. Transformers from optimization perspective. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in neural information processing systems (NeurIPS), 2022.
Auf Emre Yüksel, Yaşar Alim Türkmen, Arzucan Özgür, and Berna Alunel. Turkish tweet classifica- tion with transformer encoder. In Proceedings of the International Conference on Recent Advances in Natural Language Processing (RANLP 2019), pages 1380-1387, 2019.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations (ICLR), 2017.
Zhenyu Zhu, Fanghui Liu, Grigorios G Chrysos, and Volkan Cevher. Generalization properties of nas under activation and skip connection search. In Advances in neural information processing systems (NeurlPS), 2022.
Difan Zou and Quanquan Gu. An improved analysis of training over-parameterized deep neural networks. In Advances in neural information processing systems (NeurIPS), volume 32, 2019.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over-
parameterized deep relu networks. Machine learning, 2020.
14

---

Contents of the Appendix
The Appendix is organized follows:
In Appendix A we summarize the symbols and notation used in this work.
Appendix B. we provide a theoretical overview for neural tangent kernel (NTK). The back-
ground in Sub-Exponential random variables is elaborated in Appendix B.2 More detailed related work on the convergence analysis of nerual networks can be found in Appendix B.3
[.22
The proofs for convergence Transformer are included in Appendix C
The derivations and the proofs for the NTK are further elaborated in Appendix D including the formulation of NTK, minimum eigenvalue of NTK, and the relationship between the Hessian and the width.
Further details on the experiments are developed in Appendix B
Limitations and societal impact of this work are discussed in Appendix and Appendix respectively.
15

---

A Symbols and Notation
We include the core symbols and notation in Table 2 for facilitating the understanding of our work.
Table 2: Core symbols and notations used in this paper.
Symbol
Dimension(s)
Definition
N(poo)
Gaussian distribution of mean and variance
Faclidean norms of vectors un Spectral norms of matrix W Frohenius norms of matrix W Nuclear norms matrix W
W
W
W
Amine W
W
Minimum and maximum cigenvalues of matrix W Minimum and Maximum singular values of matrix W
miin W Omao
W)
w
i-th element of vectors w (i. (-t))-h element of matrix W W time step
W
W
Hadamard product
o.( (W)
Row-wise Softmax activation for matrix Element wise ReL.U activation for matrix W
d, (:(W)
Wi
[event]
Indicator function for event
N
Size of the dataset
d
Width intermediate layer Sequence length of the input Dimension each token
92 TI. 1/ W0 Th. Ti
Variance of Gaussian initialization W q. Wi. Wy. mo
Scaling factor Step size
X,
gduxdd
The n-th data point
go
The target
Dx
Input data distribution Target data distribution
Dr
8r. 3o
To X.! W X
Rd
The row output softmax the n-th data point
W MN
RVV
Ground truth label of the data samples (X,
f(0) = [[ 0) (xx;0)]
RV
Network output given data samples (X,
(0)
(as(X1; 0) xx(
Output of the last hidden layer given [X., 1x
(10) 5(x:0)
Network output given
samples X
fm (x(x0))
P
Output of the last hidden layer given
Wo. WK
R xd Rr.cd
Learnable parameters
Wr. wo
Blaxd Rd.
Leurnable parameters
W
+ Cg, Ax [[ 1 Cx
Purameters norm
Ay A
W
Cy. No & |wblla Co
Parameters norm
a N3/2
C...
Va (1+ + 4 CEEE (X8 Xi))
Auxiliary
Output network for input
o, Ω and e
Standard Bachmann- - Landau order notation
* The superscript with bracket represents element of vector/matrix, e.g., (0 the 7th element ur. The superseript without bracket symbolizes the parameters at different training steps, c.g., 0'. The subscript without bracket symbolizes the variable associated to the n-th data sample, e.g., X,,
B Theoretical background
B.1 Preliminary on NTK
In this section, we summarize how training a neural network by minimizing squared loss, i.., ((0))- "EM,, (f(",,,'' 0') - yn via gradient descent can be characterized by the kernel regression predictor with NTK.
By choosing an infinitesimally small learning rate, we can obtain the following gradient flow:
d
-e(("))
16

---

By substituting the loss into the above equation and using the chain rule, we can find that the network outputs f(0") E RR admit the following dynamics:
df
(f(0') - y)
(5)
dt
where K = (
a E RNNN
Jacot et al. [2018] Arora et al. [2019b] have shown
36
that for fully-connected neural networks, under the infimit- infmite-width setting and proper initialization, Kt will be stable during training and K will converge to a fixed matrix K E RNX' where
(K*) (5)) = (X, xj)
) is the NTK value for the inputs X; and Xj. Then, we rewrite Eq. (5 as:
df(0') -KK((((0") y)
dt
This implies the network output for any X E Rdd, xd can be calculated by the kernel regression predictor with the associated NTK:
f(X) = (X,
K (x,, (x)))
where K"XX,X, is the kernel value between test data X and training data Xn-
B.2 Preliminary on Sub-Exponential random variables
Below, we overview the definition of a sub-exponential random variable and few related lemma based on Wainwright [2019] random variable X with mean j called sub-exponential random variable if there exist non-negative parameters (n, a) such that
EfeA(x [e I
1
2222 for all |
< a
and we denote by X ~ a).
Lemma 2. The multiplication of a scalar s E RR and a sub-exponential random variable X SE(V, a)) is still a sub-exponentidl random variable: sX ~ Ssv sa.). Lemma Given a set of independent sub-exxonential random variables X; (v)) for v22 1 max; oi).
1 1,...,N, N, then ENI X,>
Lemma 4. Given sub-exponent sub-exponential random variable X SE(V, a) with mean p. the following inequality holds with probability at least 8:
X ul < max
V/ 2 log z: 2a log 8
B.3 Related work on over-parameterizztion for convergence analysis
Recent empirical observation shows that neural networks can fit arbitrary labels with zero training loss when applying the first-order methods, e.g., gradient descent (GD) [Zhang et al. 2017]. Due to the highly non-convex and non-smooth instinct of the neural network, large body of work have attempted to explain such a phenomenon. Early work studied the convergence of stochastic gradient descent (SGD) for training two-layer verrpaarameterized ReLU neural network with cross-entropy loss Li and Liang 2018). Du et al. (009b) show that only training the first layer of two-layer ReL.U network with square loss by GD can lead to global convergence under the assumption that the Gram matrix is positive definite. Du et al. 2019a] extend the result to deep neural network with smooth activation function and shows that the convergence is guaranteed when the widths of all the hidden layers scale in ()NV),, where N is the number of data points. Meanwhile, Zou and Gu (2019) prove that the condition for the convergence of GD for deep ReLU network is 2 which improves upon Allen-Zhu et al. [2019a] that show the result of 52 (2(22)) [ (AllenZZhu et al. [2019a] also provide several convergence analyses under the setting of SGD and various loss functions. The assumption regarding the positive definiteness of the Gram matrix made in Du et al. [2019b] has
been rigorously proved in Nguyen et al. [222] This facilitates Nguyen 2002 to demonstrate that deep ReLU network under LeCUN initialization with width in the order 12(N3) is enough for global convergence. Recent breakthrough (Bombari et al. [2022] improves previous results by showing that
sub-linear layer widths suffice for deep neural network with smooth activation function.
17

---

Table 3: preeeaaraettiiittionn conditions for the convergence analysis of neural network under
gradient descent training with squared loss. L is the depth of the network.
Model
Depth
Initialization
Activation
Width
Allen-Zhu al.
2019a
FCNN/CNN
Deep
NTK
ReLU
Du et al. 2019a
FCNN/CNN
Deep
NTK
Smooth
((N42001)
Oymak Soltanolkotabi
2020
FONN
Shallow Standard Gaussian
ReLU
((N2)
Zou and Gu
2019
FCNN
Deep
Hc
ReLU
R(NS
Du ct al.
2019b
FCNN
Shallow
NTK
ReLU
9(N9)
Nguyen
2021
FCNN
Deep
LeCun
ReLU
S((N))
Chen et al.
2021
FCNN
Deep
NTK
ReLU
S2(17)
Song et al.
2021
FCNN
Shallow
He/Lecun
Smooth
2(N3/2)
Bombari
FONN
Deep
He/LeCun
Smooth
((VN)
Allen-Zhu al.
2019b
RNN
NTK
ReLU
9(N"),
Hron et al.
2020
Transformer
Deep
NTK
ReLU
Yang 2020
Transformer
Deep
NTK
Softmax+Rel.U
Our
Transformer Shallow
Table
Softmax+ReLU
9(N)
C
Proof for convergence analysis
This section is developed for the proof of the convergence result and we outline the flowchart below:
Specifically, in Appendix C.I1 we provide some auxiliary lemmas. Lemma and Lemma 0 show that the norm of the parameters can be bounded with high probability initialization. Lemmas 7 y
and present that the network output and the output of softmax between two adjacent time steps can be upper bounded. The Lipschitzness of network gradient
its norm is bounded Lemmas 13
and 14 In Appendix
C.2
we prove the convergence of the general cases, i.e., Proposition,
In
Appendix
C.3
and
C.4
we present the proof for 1/2 and d, scaling.
C.1 Auxiliary lemmas
Lemma 5 (Corollary 5.35 of Vershynin (0022): For a weight matrix W E Rdi xd2 where each
element is sampled independently from NO0,1 1), for every Ç with probability least 1
2exp( /2) one has:
ddn
dz S miin (W) S mmxx (W) S Vd +2
+
where maxx (W) and Juinn (W) represents the maximum and minimum singular value of W. respec- tively.
Lemma 6 (Upper bound of spectral norms of initial weight). For weight matrix W E Rdm xd where dm > d, each element is sampled independently from (00, 1), with probability at least
1 2exp( /2), one has:
|| S $3 dm
Proof of Lemma Following Lemma one has:
| S dm + Vd+ C.
Letting /d and using the fact that d,,, complete the proof.
Lemma 7. Recall from Table 2 B, (Tx (1:) W WRX is the i-th row of the output
of Softmax, if max (|) Whllp w s Xq. max (| W
S λx. then its difference
18

---

between t' step and t step has the following upper bound:
| - -8:11 s 2no
2mm00 C2ds
(1 | W
Wi
+ +ix W - -wbl))
Proof.
185 18/11
$ 188 8
$ 2 10 (6.) W W - TO TO W W X
(By Corollary A.7 in Edelman et al. [202
max
To (1. We T W
((.)
W X
2max (TO
x (1.11) 1/
W W
W W
| |,)
(By Cauchy-Schwarz inequality)
2r002dd
T
W
W
w
(By Assumption D
2
W
W
W
+ +
W -W 1,) (By Cauchy-Schwand inequality, Triangle inequality)
2To
W
Wx
+
d
W W/ 1.)
Lemma 8 (Upper bound the Euclidean norm of the output of softmax), Suppose Softmax(u) Rd, then one has: 1/ Vd, S ||la|| S 1.
Proof. By the inequality of arithmetic and geometric mean:
E (1) S ||0||2
d E(((00) =1.
((v()))) S
Vd,
Vd,
Lemma 9. at t' and step, max ( wwll)) s λy. max (II 112 | s
ÅR. max W 11 | < dv. then the difference between the output of the last hidden layer at t' step and can be upper bounded by:
FC prc
Fprc 1p s
w 11 + +V (Na W wxll, + +x |ws's w
Note that Fore E Xddm is the output of the last hidden layer given {X,) > sample as defined in Table We firstly analyze each sample, i.e., Spre Rdóm and we drop the index n for simplification.
Proof.
19

---

d. Ti
d.
- 2
Σ or
(w X (x)
Σ or (wixT
i=1
ST
W x Bi WiX x T 0:
(By Lipschitz continuity of o,)
Σ
(l(w3' wx)xx / 1| + | (8)(,))
- (By Triangle inequality)
de
ST) Σ
W - w3 || / + -|| || 0. - 3:11) (By Cauchy-Schwarz inequality)
1=1
S ST de) Σ (11m - -wil
+ +dv 1|05 - -B8 [8(()) (By Lemma and Assumption T
S
(lws -wil 1, 2x (3o
W W | + Ax We' W: () (By Lemma 7
(6)
Next, we bound the difference given N data sample:
Ft.
S SVNTICa
6/2
(|ws -
w + Av-2TO (Te W W/ 1 + dx W - -was ())
= -w will, CZ (ne) W Wk 2
W
wall
+
where the last equality is by the definition of p in Proposition.
Lemma 10. If at t' and step, max ( [ S 'v, max (1lwg wg
W 2 W
S Xa. max (l W
AK, max (I) 2:
I No. then the difference between
the network output at t' step and t step can be upper bounded by:
|| |rd 1²1
s pdv Iw0 -woll +
F F
Proof.
||5" -
5"11
F w6
F Frreo w/b
<
Frd pre
w/b
w/b 2
+
1 l1z
F
pro
2
2
where we use triangle inequality. Then, the first term of the RHS can be bounded by:
T d. or (WEXTB")
Ft' pre
S
Fr' pre
STVNa&
W
|| = ////l/l
Ti Σie or (WX x X
where the last equality is by the definition of p in Proposition.
Lemma 11 (Jacobian of Softmax). Suppose v Softmax(u4) E RRd, then de Tu
= diag(v) v0
exp (1)
E exp (u(1))
Proof. We can reformulate v as:
exp (())
Σー exp (u())
20

---

Then, we have
① exp (weli) Eis exp (ea(1) (k)
eex (u --ox (u(²) Ciny exp (u(
Oyli)
if + k
@u(k)
exp (u ) اis exp exp("'t))) (xxp) (x))
j = k
(I exp exp (u()1))
- (A)) v())
j サ k
if j
= k
Thus
du = diag(t) 00
Lemma 12 (Upper bound for the loss gradient norm). If [Wl,, S λq. [klll AK. Wx;ll2 S
Xv, | S No then the gradient norm with respect to We. WK, Wv, wo can be upper bounded
by:
V Wo E( ((0T s
2112,
V W {(0 >>e I 2proda | 1'--
V Vw, ((0 )11p S puoll1 - y111
V Ww ((0T)]12 S pdv ||5* -1112-
Proof. To simplify the notation, in the proof below, we hide the index t. Firstly, consider the gradient w.I.t Wy.
N | Σ ( vn) 0)(Xm))
| w. ((0))
OW.
1=1
N
de
= T1
Σ (f(x)) - 3n)) Σ
(wo or (WyX, Bi,)) B Xn
ST
(f(xx)) - Unn) 2//2//20 2 Co
F
n=1
N
< T
N Σ
(S(X.)
o223 = No 2ll -oo Il y112
where the last equality is by the definition of p in Proposition. 1 Next, consider the gradient w.r.t wo
N | Wwo& (00)2 = - Σ (F(..) -Lw)
vn) vn) - 0)(Xm) awo
[V(X))) - Vnn) Σ or (Wv.x, Binn)
=
1=1
2
n=1
2
N N STy (S(X,) -ym))yy))))))))) yn)) 8/21y STy
(f(X,) 22.))/28/2 XvCr pdv |l y11z
< n
(8)
21

---

Next, consider the gradient w.r.t Wq.
V ((0) - -
N
- 4n) xx
@
Σ (J(X,) y/n) Σ WxX, (dia)(Be..) Bl,nBun) Xn W. (wo oō, (Wy Bi,)) x
TiI
NV.
- 2To71
(f((x,) Un)) -n))@AAkxvoo(CCV4)
n=1
s <TOTT
N Σ (ff(x,)
2o |ff -yll2
(9)
Similarly, for the gradient w.r.t WK. we have:
Vw (0))l s 2pmaAoAvdod..321 lf yll2
(10)
Lemma 13 (Upper bound for the network function gradient norm). If #/ll s 1l
Ax, [W s \v, Av,llulll2 S No, then one has:
Vof c2. c2
where
PV X3 + X2 + (2ndAAv-ood,CC) + AZ
(11)
Proof. simplify the notation, the proof below, we hide the index L. Firstly, note that:
IVoell2 < |/olllr
I(X,;;
V s(x,; 0)11 + [(,,, 0)] 0)
(12)
Then following the step as in Eqs. to 10), each term can be bounded as follows:
N (I wo ((X, (pAv)²,
[V s (oto))
n=1
N (I J(Xx; < Avod..3
(l/ 0)1pp) < (2pmAAAA
Plugging these bounds back Eq. finishes the proof.
Lemma 14 (Upper bound for the Lipschitzness of the network gradient). Suppose max (|) W/llz w | I dQ max (lIWk l2 [ s dx, max (I) 112 |
w s
Xv, max (lumbll2. (ub/l)) s do, and define & 2no Cd, (dg + Axx), then one has
Vof" Vof*
S c3
10""
0:
(13)
22

---

where
(cs) A N(n 2/2 (1+
N
doz foz+ãoCo do(l )v2) + ]]
N C2 122
+
+ N Con (2 dy
(v + +Xo + +vvoo + +3xx3
}}
N Con (2 Tv (IoCoVa +vz)+11 10 + vvvo C2d + 3d 2
(14)
Proof. Firstly, note that:
Vot" - Votil
-wo V wo T - 0²)
2 112 Fwrf(X,n 0²)) Wwyf(X,n 0¹) |
+ VwoJ 0²) Vwaf 0¹)
||: V 00)
Vwx (Xx; 04) 00/10
(15)
Then, we will prove each term separately, Regarding the first term Eq.
15 we have
V woo S (X,n;'')) V w (Xxn; 0*)
= S $ime] 1
I Crrd/// (/wW - Wi/ 1 + +Xv25002d, (πo W Wk|
+ +
We - Wa w)
S S 8/2 (1 CTd, (Te + +xx) 10" -0"-601 00/
C/ Avz) 0r 0¹
(16)
where the first inequality is by Eq. 66 and in the last equality is by the definition of a in the lemma.
Regarding the second term in Eq. 15
we have:
V wy s(x.; 0') - Vw. f(X.;
de
wo dr (w:' (wi'x. si. )) B: x. - Σ (wó od, (wixx,. Bi, T x_.
=
< SnC.vd
Σ
ww o (we' x. Bl. Bi.n. (wo d (w; B(... Bl.n
i-1
d.
SnCovd. 2
|| W a or (wt'
x Nyt'Xx.i
3/ 11, B Binll +|litnll w6 or (wi (WiXX Bl,) w6 o (wixx 11}
SnCavd.
Σ
No |.0 - B || + wó oo (wi' (T wā od. (wi Bl.m) Bim |1.]
(17)
10% Bi.nl can be bounded by Lemma 7 as follows:
The term
13 B/r s Cdd (No W Wx
|| + Ax W/ -W/ S s=||0r 0' 2
(18)
23

---

The term |ooy o o, ( (( - w -w' of (WiX,,) (,,)) can be bounded by the triangle
inequality, Cauchy-Sch Schwarz inequality, and the same method in Eq. (6) follows:
w6 oo (W x ( w o. (W; x
$ w66 or ( X ((n)) or (Wr x s
+ or (WiX o (wo - -wooll
2
do wi X, BYT WiX x T Bi 1| + wn0 w/
s SooCovdy de (w -w |, + +xx220022, (Λo
W Wx
+ +dx We wall))
w w/o
$ [TCCo 2vz)= 1
00" 0 l
(19)
Plugging Eq.
18 and Eq. 19 back Eq.
17
we obtain:
Vur Vwr TuxJ((x,:0)) 0x)
(20)
SmCanpoe R//2 doz + ÅoCo v Th(1 + λvz) + 00" - 0 1
Regarding the third term in Eq. (15), let us denote by:
U W X (diag(O(..) - B Br 7) Xn, 1S = wwy" (wo oi (
U W W T (diag(B)..) Bl. 7) Xns
h = wt (wó où, ( x ([
Then:
de
Vwo f(X::
(xx:(Xxx: TOTI
Σ
U h (1.))
Σ
U h x((.))
i=1
de
de
de
$ ToTC
Σ
U.
Σ
Ul
U
h
S Co E (|0%_1
h nim
+ ||
U., (11
(21)
We bound each term separately:
|
5 W | | diage( Bin 7 22nnd,
| 112 S | S AvÑo
In.... -hin - W | w/ oor (WFxx. xTT((.))
w/ ooi (W. Bl. 11, + |w - |w'
wy'-w; |wol|
& [iv (AoCovdd(1 Avz) + 1) + 2o] 10" 0¹
24

---

where the second inequality uses the result in Eq. 19.
U Uii
W x
(diag((B).)) - B ') x (diag((( -i 1) |
S <Cavd
W Wk
||, |x 11 diag(( - B
3 71.
W
2
Xx,12
diag(S - B diag - Bi, in
[12]]
T
im
C.
3
0¹
AM (1 [diag((((,)) diag(S |1, 1 0 8 Bi T (1)) (22)
P.
10
8
+ +A (1 B/.n
+ (115 + |13)mm11) 1151 Bl.m
C.d.
rr
00
+ 3ス
B -B.
C.d.
00"
00
+ 2TO (no
W WX || + dx
W will)]
[C अ
0- - 0¹
where the last second inequality by Lemma 7 Plugging back Eq.
we obtain:
Vwo )²)) VwoJ(Xn: 03)
Cod, f2 o j ( (v2)) + So} +AvdolC2d + + 1100 08
(23)
Similarly, the fourth term in Eq. 15 can be bounded by:
Vw. 0²)) J 13)
$ Cad C2 Xy
Xvz) + Xo} AvdolCEd, +Kkx2]
0 0¹
(24)
Plugging the upper bound for these four terms back Eq. (5) finishes the proof.
C.2 Proof of Proposition.
Proof. We can reformulate Eq. 1.2 as:
f(X) Ti w
or (W B.) fore
with B = o. (ToX (1:
W Q
WK x 1)
Rds We show by induction for every 2 0
WW s MellWWll, I dk.
s E t)
[Wvll, 1 - Xv, lolll2 No
s [0,t].
Omin (Fr: 2 $a se [0,t],
(25)
((0%) < (1 7 1(0%),
8 s [0.4] 4]
It is clear that Eq. 255 holds for / = 0. Assume that Eq. 25 inequality, we have
holds up to iteration f. By the triangle
We wal
W W =
V Wo e(0'))lF
==0
)2)0² )
$ Soprodxddoodeee
Σ 22 s 21prorrdr Σ
(1-9
220(00)
==0
==0
(26)
25

---

where the second inequality follows from the upper bound of the gradient norm in Lemma 12 and the last inequality follows from induction assumption. Let u= / - h/11ya2/3 /2, we bound the RHS of the previous expression:
P 2 2011 (1-2²) u²)
1 11 u+
1-u u 220
(27)
4 a² V2((00), (since E (0)
S potodddydod
(28)
CQ. (by Eq. 2
(29)
By Weyl's inequality, this implies:
W | - >W@ll2 + CQ = πλ.
(30)
Similarly,
| - w/ll 5 Σ1
|0021 -wåll2 = Σ
V wo ((W')ll S <лdλv Σ 20()))
5=0
s=0
s=0
S <ypāv Σ
11 1 ?
a² 8/2
v 22((0)) S pdv
2
(1 u²
1 - vt
22((0)) S pdv 4
022 V 200) Co
a²
u
(31)
By Weyl's inequality, this implies
| 112 s $1 1|2 + +Co = =AO-
(32)
Similarly, we can obtain that
W w S 2pTododv Cod, 4 V2E(00). || 112 S Ak,
(33)
a²
W wollp S pāo 4
a² (0))
| 1|2 S πv.
(34)
Next, we will prove the fifth inequality in Eq.
25
p
F Pells
pre
|| - -WPll + || well 2TOC2 ( | || || - -W2l2++ ||| ||2ll W - wa - w
Σ 7p {pToV/2((00) C.d 2p Rxx +
(by Lemma 12)
16
S 'oo [I + C2 (Ag + xx)] /ه
V 20(
= p² 2 16
Ao a² 20)
1 2aa
(by Eq. 3
where the first inequality holds by Lemma 9. This result further implies that min (Fit') 1) 2 a by Weyl's inequality. It remains to prove the last inequality in Eq. (25) holds for step t + 1.
We start by proving the Lipschitz constant for the gradient of the loss restricted to the interval (0 06+1). We define Gi++ := O + 000++1 - TT), for & E [o, 1]. Then, by triangle inequality and Cauchy-Schwar inequality, we have:
V Vol(o - vol(o }|1
Voo''to (flt 4) Vof* (ft --)l1
(35)
< l| - f4112 Vof و|2 + Vefito - vot |||||| || -yll2
< f|tto 8011| Vo f | || + 2
Vofc+o vol 2 l=(("),
26

---

where the last inequality is by induction rule. Then, we bound each term separately. Note that:
Wh w
W
Σ Wat
W
= vw, 1|p +2
Σ
V W. ((0*)) 1| S 2
V Vw ((
we have: | S λq. By the
Then following exact the same step as in Eqs.
(6)) 27 and (30,
same method, we have:
W +
W
Xv,
1+
S 'o. Now we proceed to
2
bound the first term in Eq. (35
g ge+0 f
2
S /
W
w
w
+ +11bbble
Fpt
Fi
2
Wite
w/ 1
0
2
+ +
222 (11 || will + // 2no (1 1100011111 W - - + +| w - -wål w
s pdy 10 1000 1 -
+ pdo + Cv2 + ix))
l 0
=p(xv= do + 2d + xx)) 10 00 00112
ci 0 0 =0'
(36)
where the first inequality is by Lemma 10 the second inequality is by Lemma
9 Next, the second
term in Eq. (35) can be bounded by Lemma [3 The third term in Eq. (35 can be bounded by
Lemma 14 As result, Eq. (35) has the following upper bound:
Vol vel(O) ll2 S c102 1
0
l1 + 2cz((0") 110
11
2 C C -0
11 (37)
where C1, C2, 03 are defined at Eqs.
11 14) and (36). and we further define C c1 + 20g 200e((%)
Lastly, by applying Lemma 4.3 in guyen and Mondelli [2020) and Eq. (37). have:
C
((0r+)) < ((0))
(Vo((o))
10000+10 0
-
0
0
-00''
Vol(O) ||p
C
2
< (00))
7 Vo ))
|
(By the condition on 7)
S (00' 27 V.,, 1(0 103)
=(0))
) T(f" 3)
< <E(OT) 2 (a (Fo )) ||st - 2112
2 ge((
< 7
(By induction assumption),
which concludes the proof.
C.3 Proof of Theorem
Lemma 15. Let +
[X B1.1 TT E RNxd where B
0, (Tox( T WxXX)) n E [N], then under Assumptions and 3 with probability at
27

---

least 1 - exp (-9(( - 1) 1)=dd=))) one has: Ao == Amin 100 2 ((Tv/da)
Proof. Due to Assumption 2 for any data X, the matrix XX is positive definite and thus has positive minimum eigenvalue.
We denote Amin (XXT) 2 Cx-
According to Nguyen et al. 2021 Lemma 5.3], using the Hermite expansion of or one has:
Ao 2 nvu(o,)² Aminn (
(38)
where ((a,) the 1-st Hermite coefficient of ReLU satisfying u(a,) > 0. Now proceed to provide lower bound for Amiu ($ ) For notational simplicity, define:
B₁ = 31 Rdd xd, Cis x,X] E Rd.
Then rewrite $ as follows:
Trace B Cn) Trace B Cz)
Traco B C12)
Trace B Tracc ( C2N)
Traco B2 C2a)
Trace B Cw1)
Tracc N2 Cw2)
Trace ( CNN)
By Gershgorin circle theorem Gershgorii 1931). there exists k E [N] such that:
Amiu
Trace B kb Cka)
Tracol B Ckj))
(39)
Using Von Neumann's trace inequality Mirsky 1975 and noting that Bıy is rank one matrix, one has:
Trace( Trace(BB,,Co Ckj) & max Bky O max (Cky) = || ||Br,112 v Amax (C2y City
(40)
11,1,,lll2 / ( X X x),,
11111,1112 V (ra
where the definition of the inner product between two matrices and || S 1. By Assumption 2 we have Amin A (X) > Cx. where Cx is some positive constant. By setting
t : 1| CR/(N- - 1) in Assumption 3. with probability at least exp (-Q((N 1)) one has
x X
X x,) < 11, CN - 1 k
Plugging back Eq.
39 and Eq. 440). we obtain:
Amin Φ
Trace B T kAh C.a) C2 |1 2 Amin (Cka) (raco C2 |1,
(41)
Ammin (xx
11 Cx || 2 @(1B1, 2 O(1/ds),
where the last inequality is by the lower bound of 31, in Lemma of Amin [ T) back Eq. (38) finishes the proof.
Lastly, plugging the lower bound
Remark: The estimation of An is actually tight because upper bound is also a constant order. To be specific, denote G o, a (Drr) (Apr)) we have
Gr.
ΣN
Em ( (n.))
Ao = Amin G) S
(42)
N
N
where (n.:) = X,, Next, by Liao and Couillet
2018 (Sec.A in Supplementary Material):
| (n.))
Em (o(()
(n.:)
arccos( (-1)
(43)
2π
2
28

---

Combine Eq.
42
and Eq. 443 we have
V o (m.))
A S n=1
S nv - 0(1).
2N
That means, our estimation on Ay is tight as its upper and lower bounds match with each other.
Now we are ready to present the proof for LeCun initialization under = dm scaling.
Proof. We select Cq Cк = =1 = Cv = Co = 1, then by Lemma 5 with probability at least
1 8e ddm/2 we have:
dy (0Vdm///
= =1)1
(44)
dy OVV
AK= = O( (lVdm/d).
Plugging Eq.
into Eqs.
and suffices to prove the following equations.
2 2 N NNd V 22(
(45)
a > O (Nd 4C, dt (2)
4C,
(46)
Next, we will provide the lower bound for a² = Amin (1 pre (F ). In the following context, we hide the index 0 for simplification. One can note that Pr F. is the summation of PSD matrices, thus, suffices to lower bound: Amin (Fpre Fpte), where we introduce the following notations:
FT = -f (WW, (DWJ) pre
pre
Φ [X] BI.,.... x GIN
(47)
Bnn = -o (0 (1:)) W W n E [N].
By Matrix-Chernon inequality, we can obtain that (e.g. Lemma 5.2 of Nguyen et al. [001] w.p at least 1 81,
Amia (rpro (TT pre) 2 A0/4
(48)
long as it holds d.,, (((N/Ao), where do Amin (F (0 far (o (bun)o,((an (dw))) and Ω
hides logarithmic factors depending on 81. Lastly, w.p. at least 82, one has /20((0)) < (V.N). Plugging back Eqs. ( and 33. suffices to prove the following inequality.
dm 4 O(Nd C. C. dm/d).
(49)
(dm>>/4))9
O 4C (d
(50)
By Lemma 15 with probability at least - exp -)
1d 1))
one has Ao 2 ((vv /dz)
= (d d-1 Thus, when d,, 2 (2(N3)) all of the above conditions hold. As a result, the conditions in Eqs. (45) and (46) satisfied and the convergence of training Transformer is guaranteed in
Eq. 4 Note that one can achieve the same width requirement and probability for He initialization, and the proof bears resemblance the LeCun initialization.
For the proof under LeCun initialization and To = d scaling, we follow the same strategy. Specif-
ically: As the same in the proof with To d d1/2 scaling, select Co Ck Cv = Co = 1,
then plugging Eq. (44) into Eqs. @ and 3p. it suffices to prove the following equations.
C2 C2Vdn// / 2E()))
(51)
3 2 O (Nd
4C v 20(00).
(52)
Next, we will provide the lower bound for a = Amin ((F)))))))) (F ) ( In the following context, we hide the index 0 for simplification. One note that Fpre FT is the summation of PSD matrices,
29

---

thus, it suffices to lower bound: Amin (F Fr FT pre), where we introduce the following notations:
F pr F FT = -,o, (DW, Jo, (DW.))
pre
Φ [ $1,1, ... X NI
(53)
3 0. (1:)) W Wo wxX n [N].
By Eq. (48
w.p at least 81, Amin (Fpre pre) 2 Am 4,, as long as it holds dm 2 (NN//oo)
where λo Amin (Ew~~N(t.. L4) L4) (0r (f
and f2 hides logarithmic factors depending on 81. Lastly, note that the activation function in the output layer o, is I-Lipschitz and is applied
to (Wv X [3), where X B is bounded due to the softmax's property in Lemma 8 then by
Lemma C.I of Nguyen and Mondelli 22020]. w.p. at least 63. one has /20(0") (V(N).
Plugging back Eqs. 2 and
D suffices to prove the following inequality.
dm 2 O(N C. dd/
(54)
(d (dm>o/4) 3/2 2 O(N (C
(55)
By Lemma 15 with probability at least - exp exp 1) one has Ao 2 O(nv/ /d.)
(d d Thus, when dm (2(N2), all of the above conditions hold. As a result, the conditions in Eqs. (45) and (46) are satisfied and the convergence of training Transformer is guaranteed as in
Eq. 4.. Note that one can achieve the same width requirement and probability for He initialization, and the proof bears resemblance to the LeCun initialization.
C.4 Proof of Theorem
2 (NTK analysis)
Proof. Below, we present the proof for NTK initialization. We select Co CK Cv Co
then by Lemma 5 with probability at least beddn/ we have:
Āv =O(v dm + vd), λo = =((Vdm))
(56)
λa = O dm + vd), AR = OVVdn++n
When dnm 2 d, plugging Eq. 566 into Eqs. and 33. it suffices to prove the following equations.
O O /2 /20
(57)
O(N +41++4 LC) 20(
(58)
Next, we will provide the lower bound for a² Amin (( (FPre) hide the index 0 for simplification. One can note that Fo FT is the summation of PSD matrices, thus, suffices to lower bound: Amin (F FT pre)), where we introduce the following notations:
the following context, we
F F pre
ofa, @WW )o,((,,),))))) (T,.))
pre
① = [X] BI.1.... x N.B.NIT
(59)
Bn, = ( (1.) W W.WXX.). WxXX)) n E [N].
By Eq. 48). w.p at least 1 - 5, Amin F (Frro FT 2 dm 4, as long as it holds dm 2 (N(NNo).
pre
where Ao Amm (E (o. [o, Fer) or ( { TT) and á hides logarithmic factors depending on
8j. Lastly, w.p. at least - dg. one has / 220) < Ō(NN). Plugging back Eqs. 2 and 3 it suffices to prove the following inequality.
dm-oo/4 ((N2)/2
(60)
(dm33//433/2 > ((/3/2 C 4C12))
(61)
By Lemma 15
with probability least 1 exp (- (N)-)) )), we have Ao > 2(1). Thus.
all of the above conditions hold when dm, = ((N). As result, the conditions in Eqs. 57 and 58 are satisfied and the convergence of training Transformer is guaranteed.
30

---

C,5 Discussion for different initialization schemes
Recall that the convergence result in Theorem 2 shows:
a² 2 e(0%).
E(O) S
Thus, to discuss the convergence speed for different initialization, we need to check the lower bound for a², From the proofs for different initialization schemes above, we have the following lower bound for a i.e.
a² Amiu (F pre F
pro) 2 ridmma/44 2 dma
with high probability. Plugging the value of Ti and fjy, observe that for LeCun initialization and He initialization: a² 2 N(dmn/d) while for NTK initialization: a² 2 ((N/d). Thus, the
convergence speed of LeCun initialization and He initialization is faster than NTK initialization. As result, faster step-size is required for NTK intialization.
C.6 Discussion for TO = dm and To = dm. 1/2
Eq. 4 indicates that the convergence speed is affected by a there we compare the lower bound for for these two scaling. Appendix C.3 we have proved that under the LeCun initialization, one
has a² 2 dm / 2 d ( (1.+1)) Note that this bound holds for these two different scaling, which is inside G.. Thus in the next, we need to the difference between the lower bound of [131, in the case these two scalings. Specifically, for TO = d/2 scaling, we have proved that | 2 1/d, by Lemma 8 However, for the case of To dm scaling, when the width is large enough, the value inside the softmax tends zero, as result, ||, 1 1/d.. Thus, we can see that the width increases, the convergence speed of d could be faster. Lastly, remark on the difference in the step size for these two scales, which can be seen from the definition of C' and its corresponding C1, 02 Cg in Appendix C.2
C.7 Discussion for extension to deep Transformer and residual Transformer
Extension from our shallow Transformer to deep Transformer is not technically difficult as they share the same analysis framework. Nevertheless, the extension requires several tedious derivations and calculations involving the query, key, and value matrices along different layers. Here we point out the proof roadmap for this extension. The first step is following Proposition. to provide the sufficient condition for the convergence guarantee, e.g., Eqs. (2) and (3). The second step is similar to Theorem [ where we need to verify the aforementioned assumptions for different initialization. In the second part, the main task is to prove the lower bound of a == min (F)ce) where FpRo is the output of the last hidden layer. One can apply concentration inequality to bound the difference between Jwin [FO ) and miin (Fple). where the latter is the corresponding limit in infinite width. Lastly, one needs to plug the lower bound into the assumptions in order to obtain the width requirement.
Our proof framework general and can handle the following residual Transformer. Here we give proof sketch to show how to achieve this. Specifically, consider the residual block in the self self-attention layer:
A1 = Self-attention X) & os (To(XWQ) (XWk) (XWT)
(owk)) +x.
As result, the output becomes
f(x) = m Σ or (WiXT&+(X())))
To prove the convergence of the above residual Transformer, the first part that will be modified the proof for Proposition 1. The formula for fpz becomes as follows:
d, Sprc = T or (x
31

---

Lemmas 7 and 8 remain unchanged. In Lemma 9 only the first step in the proof changes while the remaining part does not change because the term (i:) in two adjacent time steps cancels out. In Lemma TO we have
Σ or (W (,, 1l2 S rıd, (IWill2d + Czz)
|
Similarly, in the remaining lemmas, we need to add the additional term C2 for the upper bound of ||x I|2-
The second part is the proof for Theorem regarding the lower bound for og. By Weyl's inequality:
T Tmin (Free) 2 Tminl pre ||F - Fprell2.
where we denote by Fpre- = prre (XI,, f (Xx)], )], and Fpre the corresponding one without the residual connection. Then we can upper bound the second term can be bounded as follows:
F pre
pre l12 S
|pre pre
Firalle
de [NITT
de
I VN
T
or
( B: Be+ (x(+))) Tw Σo or (WvX 3:) ll2
(62)
S VVNnds Wvll2Ca
The remains step follows the same as previous analysis.
C.8 Linear wver-parametrization and attention module behaving as pooling layer
In this section, we discuss the link between the linear over- verrpaarametrization and attention module behaving as a pooling layer under the d scaling. First, due to the d scaling, the attention module degenerates a pooling layer according the law of large numbers. In this case, the nonlinearity on X disappears and thus the minimum eigenvalue of θ can be estimated via X X Accordingly, this leads the minimum eigenvalue in the order of ((N/d). and thus linear over over-parameterization is enough.
D Proof for NTK
In this section, elaborate the proof for Lemma | Appendix D.3. respectively.
Appendix D.I the proof for Theorem 3 in
D.1 Proof of Lemma
Proof. We will compute the inner product of the Jacobian of each weight separately. Firstly,
analyze wo- Let us denote by Br = as ((.)) ( WKX E Rd, Then:
(f(x)
or (Wv.X
Dwo
E=1
The inner product of the gradient is:
0f(x)
0f(X')
lim
dm +000
Owo
Owo
lim (ar (TTX) (a, (Wv.X33.))
(63)
dm +00
1=1,==1
d (o.) (a, (w X 12.)) (ar (w X 1d.))
where the second equality uses the law of large numbers. Secondly, we analyze Wo:
0f(X)
de
(wo o or (WvxTT))) Wv (diag(B.) -,,3.) X XWT e,e, x
ow(R.4)
TOTI
32

---

Thus:
x
d, = TOT1
W
(diag(3.) - (, XW, (wo oo (WiX B:)) x1..).
oWe
1=1
The inner product of the gradient is:
0f(x) 0f(x') @
lim
dm 100
aW
lim
Trace W (diag(B))) - ),,3,) XW, (wo c o, (Wv (3))) x
400
==1Ø==
X (wo c o, (WvX')))) 33)) WvXx (diag(8)) )3,3,) x'W
lim
Σ
(1 (i:T (TOTI (x (diag((:)) B:3,) XWV (wo orr (Wv B:))
dm +00 1=1j=1
TOTI W XT (diag(8)) - (8)) B3,35)
x'W. (wo oa ( X ())))
For the first term of the inner product, we have:
lim X (diag(3.) - 3,B, XW (wo o (Wvx33.)
dm +00
TITI dm (k.1)) (k)
-k= To or (Wy X Bi) (0)
lim
X (diag(B.)) - B,B,)) x
dm +00
TOTI Edm k= Wit u or (WvX ( (1)
(. (k)
= =X
= (diag(1a,) la, 12,) x 0,
where the second equality is by the law of large numbers and Ew 0 for random variable O. Similarly, lima oW oWk 0.
w -~~ ~NO, 1). Thus: lima.. +000
af(x) aw
af(x² aw
af(x) 0f(x²
Lastly, we analyze Wv:
0f(x)
Σ
d.
(wo or 3i)) T epey TTB. Bi.
aw(r-4)
i=1
Thus:
(x
d wo or (WvX (T)x B, x.
Ti
oWr
The inner product of the gradient is:
Of(x) 0f(x')
lim dm
Ow
dwv
lim 2
Trace (wo o (WyX Bi)) B. xX" B3 (wo oō, (" xT)) ")
i
dm +00
imlj=l
d
lim
2
(wo orr (B)) (wo od (WrX B.)) B. XX 33
dm +00
i=1j==
de
dm
mw Porr (wit BT) or (w/h (k:)) x TT..) BT Xx"33 B3
(K)
lim dm
1=1,=11==
=fEL--No (ōr (w" T x 1,.)) (ō, (w x T 1a)) [1 14,)
(64)
where the last equality uses the law of large numbers.
33

---

D.2 NTK minimum eigenvalue
Lemma 16. Given defined in Eq. (53) and 4- defined in Lemma denote X, >min ( (T*
and suppose the width satisfies dm = -2( No
/log2dd Na/8) then with probability at least 1 - ö, one
x2
has || [ T Tb S Ar and Amin $
3A_ 4
Proof. In the following content, the variable with indicates the corresponding variable with infinite width dm. According to the definition of o in Eq. 53) and 4- in Lemma for each entry in p and & T we have
1 T x )(1r)) 1I X, 31, Bi X,, Bi,1
< X ( - Bi,a)l2 + x (3,, - Bi. S C2de(11B1.n -Bi +11 (irrl2))
(65)
Next, we will bound ||,,n - Bi,al12
B1,, Bi.n 0s (ToX W wxx
wWT -s (xx We W WxX x )
(66)
$ 22Cdds @ Wx - To W Wxll
( (By Lemma
7
Let first consider the absolute value of each element of To W WK To Wo Wk. i.e.,
dm w
W WK W Wx (65) = TT w wll
W wal
Since (l( 19.) ~ SE((SEE(S)))) Molxx.vmmme (mQK) is sub sub-exponential random variable, by Lemma and
dm ((a SE(TT 2dm, v2ronQu). by Bernstein's inequality, when Lemma 3 Σ g
dm Tog the following inequality holds with probability at least 0: 6:
am (a
wh (a To Σ
w
70
W
(i < 2TOTI v dm log (2/0).
(67)
8' Substituting 8 da and applying union bound, we can obtain that when dm 2 2log 2d with probability least 1 - 8':
BI, Bi,ul11 I 4TONOR CZddd dm (2d"(82)))))
(68)
Substituting back Eq. 665. the following inequality holds with the same width requirement and probability. 1( T x ((nr)) S STOTI Czd,d dm log((dd (8d) Applying the union bound over the index (n, r) for E [N] and r E [N], and substituting 8' = NT. we obtain that when the width dm 2 2og 8
5"
N²
the following inequality holds with probability least 1 8":
N
N
ΦΦ - x.
M
-...+TT(ar)p C/N dm log (2d2N222
r=1
N² (1og((2dNN30/0
In the case of LeCun initialization when dm 22 Lastly, one has:
|
F S d=
A
32
Amiin + ) X, -1 $
|
112 2 h
क.
where the first inequality is by Weyl's inequality. One can easily check that the same result also holds for He initialization and NTK initialization.
Lemma 17. Civen the xx defined in Lemmal 7 then when N 2(d)), with probability at least
1 - c d one has: Amia TT ((N/d).
34

---

Proof.
Firstly, let us define = - 1/2
1/2
then:
Amin ( " )) = Amin s Φ²) = Amin 121/2 T 0 2'E1/2T)
2 Amin (Z* Z ( [ 2 Xmin ( (z"T 1'
where the last inequality is by Assumption follows:
Note that reformulate Z by plugging $ as
1/2
Z 1->>1/12
Σ d
(8:3) Σ 1/2
Recall that in Assumption 4 define the random vector X n ( and write the covariance matrix for each feature of æ as Σ = E Jara: J. Here we further define the random vector 1/1/2x
Clearly, the covariance matrix for each feature of z is Σg [[[z>|] 1 la. Therefore, by Proposition 4 of Zhu et al. 20222). ). we have Amin (Z* Z*) 2 (N/d), which finishes the proof.
Lemma 18, Suppose the number of samples N 2 Q(d') and the width d.. = ((N2)_ where
A= Amin ( ( Tb then under Assumption 4 with probability at least 8 - E e one has
Ap Amin (F w0 La) (lor((wwoor (ww) )r
S(nvN/d)
Proof. By the Hermite expansion of or one has:
Ao 2 nv ($ )
where p(a,) is the 1-st Hermite coefficient of ReLU satisfying p(o,)) > By Lemma 16 when the width satisfies dm 22N22, ?, then with probability at least 5, one has Amin p 2 3 X, Furthermore, by Lemma 17 when 2 ((d") w.p. at least - en one has A, = ((N/d) /d) Thus, with probability at least 8 ed, one has:
No 2 1 n N(d))/
Below, provide lower bound for the minimum eigenvalue of NTK, which plays key role in analyzing convergence, generalization bounds, and memorization capacity [Arora et al.) 2019a Nguyen al. 202T Montanari and Zhong 2022 Bombari al. 2022
To prove this, we need the following assumption.
Assumption 4. Let x = Díe (x()))) Σ = E[xx JJ then we assume that Σ is positive definite, i.e., Amin (Σ) Czz for some positive constant Cz.
Remark: This assumption implies that the covariance matrix along each feature dimension is positive definite. In statistics and machine learning (Liu et al. 2021 Liang and Rakhlin 2020 Vershynin
2018 Hastie et al.] 22022. the covariance of æ is frequently assumed to be an identity matrix or positive definite matrix.
Lemma 19 (Minimum eigenvalue of limiting NTK). Under Assumptions 7 and 4 and scaling Tu d=m when N Z 2dd"), the minimum eigenvalue of can be lower bounded with probability at least 1 e as: Amiu Amin(K) ' (((N)))((Nd)) where u(a,) represents the first Hermite coefficient of the ReLU activation function.
Proof. By the Hermite expansion of or have:
Amiu Timin(Ewwshh (Ew-N(o.)) (a, ( o, (#'w)
Amiu ())
x Σ
-miin
po(a,)2 Oi= (
Nguyen and Mondelli. 2020
Lemma D.3]
2 ((o,))2 Amin (minn)))min(01 ( TT.
35

---

Note that when N 2 S2(d*), based on Lemma 17 and the fact that -- and T d share the same non-zero eigenvalues, with probability at least e-d one has:
N
Amin K) 2 1(
=
which completes the proof.
The minimum eigenvalue of the NTK plays an important role in the global convergence, similar to a in Proposition. By defining a - min (Frre) under the over-parameterized regime with dm 2 N, we have a² Amin ( pre) F which is the exact minimum eigenvalue of the empirical NTK with respect to the weight in the output layer.
Remark:
D.3 Proof of Theorem 3
Proof. Before starting the proof, we introduce some useful lemmas that are used to analyze the randomness of initialized weight.
Lemma 20. Given an initial weight vector u Rdm where each element is sampled independently
from N (0, 1), for any vp such that | w has:
12 R. with probability at least 1 - 2 eexp(-dm/ 22, one
u la I R 3
(69)
Proof of Lemma 20
By triangle inequality and Lemma 6
2 I ||00 w|l2 +
||l||2 S R+ 3
(70)
Now we are ready to start our proof, in the following content, we avoid the tilde symbol (-) over the weight for simplicity. Because we aim to study the effect of width (dm), to simplify the proof, we set the embedding dimension of the input as one, we can write the output layer of the network into the following form:
f(x) To Σ o, (ae 1=1
(Tr""wwwaa)) (wyx)) (wyxz))) wo
(71)
The Hessian matrix H of the network be written the following structure:
HQ.
HQ.K
HQ.
Hq.0
HK
HK.
HK.v
H
H
(72)
Hv.Q
Hv.
Hvv
Hv.0
Ho.
Ho.
Hov
Ho.0
where partition the Hessian H of the network each Hessian block e.g., HQQ 0²g f
5²5
Binguin HQ.K = Bwgguug etc. Based on the triangle incquality of the norm, the spectral norm of the Hessian H can be upper bounded by the sum of the spectral norm of each Hessian block. We start by analyzing HQ.Q. The time step is hidden for simplification.
0f(x)
de
TOTI Σ Ow ())
(1) mi (wo o or (wvx (.)) x (di -3,3, x,
(73)
where the Jacobian of Softmax is given by Lemma 11 Next, we calculate the second-order derivative. Each element of the Hessian Ho. is
8 f2
(jx)
Σ
( (p) w w k
(z (20 (wo or (wvx (B)) x Tı
H
To Ti
2uly) 0m:
36

---

where we denote by r, diag (B, 3,g, x) (diag (G.) - 3,3,") xB,
3,x" (diag (B.) 3,3, T) symbolizes Hadamard product. According initialization of
wo, wv wQ. wx and Lemma 20 with probability at least 1 - Se dmm we have
| 2 S 3V mvd + R
wwoll2 S 33/nodm + R,
|well22 S 3/ + R || S 33 + R.
The spectral norm of HQ.Q can be bounded with probability at least 1 - Se den/2
d.
HQ.0ll2
=1
(a (0) 2wo (wo orr (wv T Bi)) x T
i=1
d,
) (wo o or (wvx B()) T F w
(Homogeneity of norm)
d2
| 1
(1) )xx r lo (Cauchy-Schwarz Schwarz inequality)
i=1
T1
2
Σ
(x 2 ri (3vnodm + B) (3 + R) (3/nqdm + R)(33/ nkdm R)
(Lemmal 20
where the last equality is by the bound of [ by triangle inequality:
[Tll2
diag (3, B,B," T =))l2 + (d(aa (G,) B,B. 1) xB; 112 + ||B (diag (8.) - -B,B, 7) æ=|2
(74)
For the first part Eq.
74 we have
diag (B: B,3,
112 [diag (B, o (,3;;T))) - (1 -3,3, x2)x2) 1Ia
< (B3 o x) 1l)l + |,3, x ==I.c 00 I ||~||<c + +3, [B,TTol xa| s ||ا|| + ||3:11, ||m||| S 2C
For the second part of Eq. 74 have
(diag (B:) - 3,3, 23, T
2 - (113.l + 113,,,
|x3," T
s xB 1| S ||=|| 3| 1|, S 2C2.
For the third part of Eq. 744. we have
|,x T (diag (B)) - 3,,, T) 20|2 - 3,x T ||, (diag (B.) - (,B, ||
Next, we analyze the Hessian HQ.K, where each element is:
f2f
0 f2f(x)
(.P)
Σ
O (p) w K w
(ww (wo o or (urv (B))) Taa
Ouly ( owo
(p)
aug () Orol))
Due to the symmetry, similar to Hq the spectral norm of HQ,x can be bounded with probability at least 1 - ddm/ |H 222 O(1/vdm.)
Next, we analyze the Hessian HQ.v, where each element is:
8
(d.p)) H
(
(p) or (w TB.) x (diag(B.) -BB3.) x. wo
TOT
w
37

---

The spectral norm of H can be bounded with probability at least 1 - 8E Am-
|
(1) or
(p) w x (.) x (d.ag(3.) -B,3,) xwkwo '1l2
==1
z (10 [u("2x' x (3))) x= (diag((,) - -B,3.) x wo 1122
TOT
(Homogeneity of norm)
S ToTi
xi [ 3.33 laa ~ W 1 (Cauchy-Schwarz inequality)
S ToT1
b Frll|diag((3) 3,33 |2 3 R) (3v nodm (3/nodm
+ R)
(Lemma 20
1=1
=((1/dm)
where the last step is by Weyl's inequality and the range of the output of Softmax (from zero to one):
1nin(3!7)) + 113.112) I (ddia((,)) B,B, 12 -
max( + ||8 O(1.
It suffices to bound ||3 i.e., I/ (zf q (/x) 112- Since the range of each element of the output of Softmax is from to 1 and the sum of is one, || S 1.
Next, we analyze the Hessian Ho,o, where each element is:
8J
TOT Σ ww
(0) u or (w(f) 2 x (diag(3.) -B,3,) x.
0ug ow
The spectral norm of HQ.v can be bounded with probability at least -dm/22
H HQ.
||OT|
Σ
x (
( x B:) TT (d.a - 3,,35) wwx T122
i=1
x (6) dr (uV) x 3.) x= (diag(B:) - B,B,) 2: x wy T12
(p)
(Homogeneity of norm)
1=1
ToT1
| (ril(ddag((3.) 3,,, ||2||= wo 1l2lll/lvll (Cauchy-Schwarz Schwarz inequality)
d,
S TOT1
bz (edilidag(()) - 11,,3,,,21lll ||2| (3 ///oodnm + R) (3 + R)
(Lemmal 20
1=1 (1/dm))
Next, we analyze the Hessian Ho,v. Each element of Ho,v is:
8 = Owo aw
de Σ
H
or uo
B x
T [ (1=))
i=1
Note that Ho.v is diagonal matrix. Consequently, the spectral norm of Ho.v is:
Ho, = max sE(dm)
(i)) 0.V
< T Σ x C.B.l ST
= 0(1ddm).
Next, since each element of Ho.0 and Hv.v is zero, the corresponding spectral norm is zero. Lastly, due to the symmetry of wx and WQ. we can obtain the same bound for the corresponding Hessian block. Thus, the spectral norm Hessian H is upper bounded by (l) which completes the proof.
38

---

Thanting tock
$am10
Ram 10
100
Epacts
Epacta
LeCun initialization, 7
(b) He initialization,
Thshiigg boak
(c) NTK initialization,
1. (d) NTK initialization, 7 = 10.
Figure 5: Comparison of convergence curve for different initialization schemes. The convergence speed under Lecun/He initialization is generally faster than that of NTK initialization with the same step size.
E
Additional details on experiments
E.1 Additional result in Section 5.1
In this section, we present the convergence curve for different initialization schemes following the setup in Section 5.1. The results in Figures [a to 5c show that when using the same step size = for LeCun and He initialization exhibits sumilar ar behavior while NTK initialization leads to slower convergence. This is consistent with our theoretical finding. We also depict the result with larger step size = 10 for NTK initialization in Figure 5d following our analysis in Appendix
E.2 Additional result and set-up in Section 5.2
We have validated Assumption on language data in the main paper. Here will validate it image dataset. Specifically, we choose MNIST dataset and consider the embedding in the
Probabiii
0.6
ViT as X, then plot P { X, X x.x) 2 t increases in Figure 66] where we could see exponential decay.
0.4
The architecture of ViT the same as in Section 5.2
0.2
Additionally, in Figure we can see that 1/2 setting achieves faster convergence in training loss and higher test accuracy. In Figure 8 we conduct the experiment on MNIST with regres- sion task by using the same architecture. We take two classes 10 in MNIST for binary classification, the output of our regression is taken by sign( f'), and the training error is given by [sinns - y2. can see that the convergence speed of d,m 1/2 faster than that of 1 Overall, these separa- tion results on both classification and regression tasks provide good justification for the d d1/2 scaling.
0.0
0.00
0.25
0.50
0.75
Figure Verification of Assump- tion on MNIST dataset.
39

---

16, 15.5""""
de 1824, 1oddd
6. 1836 75
Pansng
de 1024, 0,
18386,
Test accuracy
Epochs
Figure Test accuracy classification task on MNIST dataset.
4. 16, ma
de 16 m
de 75000, 75
3. 1824, m.
Training loss
Iterations
Figure 8: Regression task (with MSE loss) MNIST dataset.
6,110
du
Ouigg wepho
d--1000
dc-40000
dum10
1020
4030
Epocos
Spochs
(a) Convergence curve.
(b) Weight movement.
(c) Kernel distance.
Figure 9: Experimental validation of the theoretical results on Transformers with To = d scaling trained synthetic data. (a) Linear convergence. (b) Rate of change of the weights during training. Observe that the weights change very slowly after the 50th epoch. (c) Evolution the NTK during the training. The result mirrors the plot (b) and demonstrates how the kernel varies significantly at the beginning of the training and remains approximately constant later. As the width increases, the NTK becomes more stable.
Limitation
Firstly, in this work, we prove the global convergence guarantee for Transformer under different initialization schemes. Nevertheless, we only consider the gradient descent training under squared loss.
40

---

Our result does not cover SGD training and other loss functions, e.g., hinge loss and cross- eross-yntropy loss. Secondly, the model architecture being analyzed in this work the encoder of Transformer which is widely used for regression problems or image classification problems. We do not consider modeling.
the entire Transformer including the decoder, which designed for Lastly, our model does not cover the deep Transformer. We believe our analytic framework paves way for further analyzing large-scale Transformer.
G Societal impact
This work offers convergence analysis for the Transformer, which is element within con- temporary large language models. Our theoretical analysis with realistic initializations and scaling schemes lays theoretical foundation for the interested practitioner in the ML community to further study other priorities of Transformer such as generalization and in-context learning. We do not expect any negative societal bias from this work.
41


Reviews:
[{"review_id": "naqrb93CJo", "reviewer": "Reviewer_YNcT", "summary": "The paper presents convergence results for shallow transformer networks for commonly used Gaussian-based initialization schemes in deep learning (LeCun, He) and different scaling regimes under finite overparameterization. In addition, the authors provide limiting neural tangent kernel (NTK) analysis for its minimum Eigenvalue and stability during training. They consider the encoder part of the transformer for the analysis which consists of a single-self attention layer, followed by ReLU, average pooling and a linear layer which maps to a scalar. The analysis is presented for gradient descent (GD) training with square loss applicable to classification/regression, where all parameters apart from the ReLU layer are trainable. The convergence framework for $1\\sqrt{d_m}$ scaling uses three assumptions: bounded data, linearly independent tokens in each input example and dissimilarity of two different data examples. The $1/d_m$ scaling regime additionally uses a PDness assumption on the covariance matrix. The paper also provides experiments on synthetic gaussian data and on MNIST data for a Vision Transformer (ViT) tracking convergence under scaling regimes, width and NTK evolution.", "questions": "**Some more suggestions on intuitive remarks for proof sketch**: \n1. It'd be helpful to mention the condition on gradient that makes the variant of PL and is used to prove convergence. \n2. Breaking down (briefly) how the local Lipschitz property of loss is proved.\n\n**Typos/minor bugs**:\n\n1. $ \\gamma <=1/C$ in prop. Suggestion: It’d be helpful to specify that $C$ signifies local Lipschitz constant for the gradient.\n\n2. Around line 653 in the appendix: Should it be $\\theta^{t+1}, \\theta^{t}$ instead of $\\mathbf{W}^{t+1}, \\mathbf{W}^{t}$? Also, in the same proof, second last step should be $\\sigma_{min}^2$ and $||f^t-\\mathbf{y}||^2$ - there might be a 2 factor missing in the result as well.\n\n3. Lemma 15 proof, line 661: I don't think $\\mathbf{XX}^T$ (correctly mentioned in Assm. 3 - $\\mathbf{X}^T\\mathbf{X}$) should be called the covariance matrix as the data is placed along rows.\n\n**Results**:\n\n\n1. For thm. 1 remark: For larger $\\hat{c}$, the exp quantity grows - which seems a bit counterintuitive in terms of separability of data. \n\n2. In the $1/d_m$ regime there’s an additional req. On $N$, but not in the $1\\sqrt{d_m}$ regime - is there any obvious reasoning for this or is it just a proof artifact\n\n3. Is there any step-size $\\gamma$ (or $C$, same thing) difference in $1/d_m$ vs $1/\\sqrt{d_m}$? As far as I know, typically in NNs in the latter it is $O(1)$ [5] and $O(d_m)$ in the former [4]. \n\n4. It’d be nice to know the authors' thoughts about the tightness on the width condition, especially in the $\\sqrt{1/d_m}$ regime.\n\nLimitations section says the result doesn’t cover cross-entropy loss, can such a PL variant + local Lipschitz technique be applied to it? I’ve mostly seen such analysis with square loss [1, 5].  \n\n\n**References**\n\n[1]. Nguyen, Q.N. and Mondelli, M., 2020. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33, pp.11961-11972.\n\n[2]. Li, H., Wang, M., Liu, S. and Chen, P.Y., 2023. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015.\n\n[3]. Oymak, S., Rawat, A.S., Soltanolkotabi, M. and Thrampoulidis, C., 2023. On the Role of Attention in Prompt-tuning. arXiv preprint arXiv:2306.03435.\n\n[4]. Mousavi-Hosseini, A., Park, S., Girotti, M., Mitliagkas, I. and Erdogdu, M.A., 2022. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint arXiv:2209.14863.\n\n[5]. Du, S.S., Zhai, X., Poczos, B. and Singh, A., 2018. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054.", "limitations": "I couldn’t find the code in the supplementary or any anonymous link to reproduce the results.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "**Importance and motivation**: Very interesting work, I enjoyed reading it. The paper certainly takes a step and provides novel results towards understanding the training dynamics of shallow transformers and the overparameterization required for global convergence - an important research direction considering their prevalent usage. The authors do a good job in motivating the need to study different scaling regimes in the introduction section as well. \n\n**Organization and writing**: The paper is organized, well-written and easy to follow in general.  \n\n**Assumption verification**: I like the important remarks and the experimental evidence authors use to back the assumptions, particularly for assumption 3 where they provide experiments on IMDB and MNIST dataset.  \n\nConvergence results discussion in sec 4.5 on slower convergence of $1/d_m$ scaling and NTK initialization being slower than He/LeCun provide good insights into the theoretical results. \n", "weaknesses": "**Missing intuitive explanations**: Standard convex analysis requires PL inequality + smoothness of the loss function to get a linear rate. Per my knowledge, the proof follows a similar(ish) recipe to [1], where they provide linear rates for NNs using the local gradient-Lipschitz property of the loss + a variant of PL, which gives a good analogue to convex analysis (as briefly mentioned already in the related work sec). I think it’d be nice for the reader if the paper contains a longer proof sketch with such intuitive comments/remarks for the proof steps. Some more suggestions on this can be found in the questions section. \n\n**Minor bugs and typos**: There are some minor issues in the proofs and some typos in the paper: please refer to questions. \n\n**Regarding assumption 3**: In its current form, assumption 3 forms a tail bound on the similarity of the respective empirical covariance matrices. This covariance inner product equals similarity of $\\mathbf{X}_i, \\mathbf{X}_j$ + some cross terms - how well does assumption 3 reflect dissimilarity as mentioned in the paper? It would have been nice had there been plots for similarity of $\\mathbf{X}_i, \\mathbf{X}_j$ along with the current ones.\n\n**Concerns/questions regarding the results**: I have several queries regarding separability constant $\\hat{c}$, lower bound on $N$, etc. Please see the questions section. \n\n**Applicability on contextual token data**: It might be interesting to see if/how the results apply to the commonly used contextual token data setting [2, 3] for theoretical analysis of transformers. Eg: Assm. 2 fails if sparsity [3] isn’t small enough, a small discussion might be something to consider.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer YNcT for the insightful feedback and for appreciating our theoretical analysis of Transformer and its value for the theoretical community. We address the concerns below.\n\n---\n\n> **Q1:** [Suggestions on intuitive remarks for proof sketch.]\n\n**A1:**  Thanks for the suggestion, we have already discussed about the proof sketch, see the [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA) for details. \nWe will add the proof sketch as well as ituitive remarks for better understanding.\n\n---\n\n> **Q2:** [Typos/minor bugs.]\n\n**A2:** We are thankful to the reviewer for spotting the typos. We have already corrected them in the revised version. For example, in proposition 1, it should be $\\gamma \\le 1/C$. Around line 653, it should be $\\bf \\theta^{t+1}$, $\\bf \\theta^{t}$, $\\sigma_{min}^2$, $||f^t-\\mathbf{y}||^2$, and a factor of 2. In line 661, we will remove the phase ‘covariance matrix’ to avoid confusion.\n\n---\n\n> **Q3:** [Assumption 3.]\n\n**A3:**  We plot the figure for $<X_i, X_j>$ in the same MNIST dataset as shown in **Fig.1** in the one-page pdf. \nWe can see a similar exponential decay for the inner product of two data points. \n\n---\n\n> **Q4:** [For thm. 1 remark: For larger \\hat{c}, the exp quantity grows - which seems a bit counterintuitive in terms of separability of data.]\n\n**A4:** As $\\hat{c}$ increases, the probability in theorem 1 decreases. Meanwhile, from theorem 3, we can see that a larger $\\hat{c}$ results in less separable data (line 199 should be corrected). It makes sense as less separable data leads to a lower probability. \n\nSorry for the confusion because of the typo in line 199 and we have already changed it to \"A larger $\\hat{c}$ results in less separable data\" in the revised version.\n\n---\n\n> **Q5:** [Applicability on contextual token data.]\n\n**A5:** We are thankful to the reviewer for pointing out these interesting works on the analysis of contextual token data and we believe it would be of great interest to generalize our result to the classification task with contextual token data. We have already included a discussion in our revised version.\n\nWhen extending our current analysis to contextual token data, we need a refined analysis framework to distinguish that label-relevant (irrelevant) token. Besides, Assumption 2 can be slightly relaxed to the case that $\\lambda_{min}(X_k X_k^{T}) > 0$ for any $k$ holds with a high probability. Nevertheless, a new proof framework on the minimum eigenvalue under contextual data is required if we want to get rid of Assumption 2 and its variants.\n\n---\n\n> **Q6:** [In the regime $d_m^{-1}$ there’s an additional req. On $N$, but not in the regime $d_m^{-1/2}$- is there any obvious reasoning for this or is it just a proof artifact.]\n\n**A6:** This is a proof artifact because the requirement $N > \\Omega(d^4)$ is required to provide a lower for the minimum eigenvalue of a specific matrix in the case of  $d_m^{-1}$. \n\nTo be specific, we need this assumption to ensure $\\lambda\\_{\\min} ( {\\boldsymbol{Z^*}}^\\top \\boldsymbol{Z^*} )   \\ge N/d -9N^{2/3}d^{1/3} = \\Theta (N/d)$.\n\n---\n\n> **Q7:** [Is there any step-size (or $C$ , same thing) difference in $d_m^{-1}$ vs $d_m^{-1/2}$?]\n\n**A7:** Yes, these two different scaling schemes have different $c_1$, $c_2$, $c_3$, see line 652. This leads to different C. We will make it clear in our final version.\n\n---\n\n> **Q8:** [It’d be nice to know the authors' thoughts about the tightness on the width condition, especially in the regime $d_m^{-1/2}$.]\n\n**A8:** In the regime $d_m^{-1/2}$, the lower bound of the minimum eigenvalue $\\lambda_0$ is in the constant order, which is tight.  \nBased on this, by studying the relationship between $d_m$ and $\\lambda_0$, see Eq. (48) and (49) for details, we can prove that quadratic over-parameterization is required.\n\nThis quadratic over-parameterization requirement could be relaxed if a better relationship is given. However this is unclear beyond our proof technique and we will add a remark for this in the revised version.\n\n---\n\n> **Q9:** [Limitations section says the result doesn’t cover cross-entropy loss, can such a PL variant + local Lipschitz technique be applied to it? I’ve mostly seen such analysis with square loss.]\n\n**A9:** It appears feasible to extend our current analysis from the squared loss to the cross-entropy (CE) loss. The proof idea is still based on the PL variant + local Lipschitz technique but an extra data distance assumption is needed to ensure the global convergence, refer to assumption 4.2 in [2] and assumption 2.1 in [3] for details.\n\nNote that under the CE loss, [2] and [3] deal with fully-connected networks with much stronger over-parameterization than that of the squared loss setting, therefore we think the CE loss requires a stronger quadratic over-parameterization condition for Transformer. We leave this as future work.\n\n---\n\n> **Q10:** [Code]\n\n**A10:** According to the rebuttal rule, we have sent the code to AC.\n\n---\n\n### References\n[1] Nguyen, et al. \"Global convergence of deep networks with one wide layer followed by pyramidal topology.\" NeurIPS, 2020.\n\n[2] Zou, et al. \"Gradient descent optimizes over-parameterized deep ReLU networks.\" Machine learning, 2020\n\n[3] Allen-Zhu, et al, \"A convergence theory for deep learning via over-parameterization.\" ICML, 2019."}, {"review_id": "75gj9Xh76E", "reviewer": "Reviewer_n1Tb", "summary": "This paper proves the global convergence of the shallow transformer under a more realistic setting. While I don't often read theoretical machine learning papers, I found some of the analysis interesting. However, the paper does feel a bit incremental compared to previous work.\n", "questions": "In what way does the proposed analysis different from prior analysis, aparting from a more realistic setting? To achieve the more realist setting, do the authors have to come up with a different analysis technique to do that? It is not clear to me how novel the analysis techniques are.\n", "limitations": "Most of the analysis is done on shallow Transformers as pointed out by the authors. Moreover, the practical implication of the analysis is not yet clear.\n", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I would like to commend the authors for their efforts in conducting a more realistic analysis. I also appreciate their inclusion of analysis, discussion, experiments on artificial data, and experiments on real data, which makes the paper more complete.\n", "weaknesses": "The paper does feel a bit incremental at times. It would be helpful to see more discussion of how the analysis could be applied to practical experiments.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer n1Tb for the insightful feedback. We address the concerns below.\n\n---\n\n> **Q1:** [Most of the analysis is done on shallow Transformer. The practical implication of the analysis is not yet clear.]\n\n**A1:** We agree with the reviewer that this setting is a bit far away from practical Transformers. However, from a theoretical side, even for a shallow Transformer (with a realistic setting), the global convergence analysis is unknown and is still an open question in the deep learning theory community. The contribution on the convergence analysis for Transformer has already been recognized by the rest reviewers. \n\nIn our revised version, we have already included the residual block into our analysis, to capture the fundamental components of Transformer as much as possible, see [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA).\nOur theoretical results are still able to provide some guidelines for practical settings:\n\n- Initialization: our theoretical result suggests using He/LeCun initialization instead of NTK initialization because the former initializations can lead to faster convergence.\n- Input format: We show how the data formulation has an impact on the performance of the Transformer. Specifically, given a vector input, if we formulate it along the sequence dimension, we can not guarantee the training converges from theory and it does not converge empirically. If we formulate along the embedding dimension, then the training can converge.\n- Scaling: Our theory identifies why the self-attention layer degenerates as a pooling layer under the $d_m^{-1}$ setting with an infinite-width setting. For small width $d_m$, there is no significant difference for these two scaling schemes on the convergence; But for a large enough $d_m$, then the scaling $d_m^{−1/2}$ admits a faster convergence rate than that of $d_m^{-1}$ because in this case, the self-attention layer degenerates as a pooling layer more or less.\n\n---\n\n> **Q2:** [In what way does the proposed analysis different from prior analysis, apart from a more realistic setting? To achieve the more realistic setting, do the authors have to come up with a different analysis technique to do that?]\n\n**A2:** We will add a paragraph about the technical difference from prior work in the final version, see the [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA) for details.\n\n---\n\nWe hope that our responses have addressed the concerns of the reviewer. If there are any remaining questions, we are happy to discuss them further."}, {"review_id": "pbIW0JyqWA", "reviewer": "Reviewer_9ERM", "summary": "The transformer network is a popular but complicated architecture, which are widely used in NLP and computer vision. The contribution of this paper is two-folded:\n1. it aims to give convergence guarantee for a shallow transformer network (encoder part) under different weight initialisations. Also, the paper concludes that the convergence rate of the NTK initialisation is slower than the original (LeCun/He) initialisation, which is often used in realistic setting;\n2.  it shows that a quadratic over-parametrisation guarantees global convergence of shallow transformer in usual setting.\n\nAssumptions are verified by realistic examples. Convergence and training dynamics are also tracked in both synthetic and real-world data sets. ", "questions": "The quantity $\\alpha$ in Proposition 1seems to be sub-optimal. Is it the reason why you do not plot the theoretical decay with the empirical one? Could a log-scale plot be convincing?\n\nAlso, in the proof of Lemma 1, you say by LLN we have the second equality in line (67). But the entries in $\\beta'_j$ is correlated to these in $\\beta_i$. How could you replace each vector respectively with $\\frac{1}{d_s}\\bm{1}_{d_s}$  by taking the limit before the product? ", "limitations": "Yes, the authors include a section (Section F in appendix) to address the limitation of the paper. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper gives a standard approach to proof global convergence by bounding the smallest singular value of the weight matrix in the last linear layer in the transformer network from above. The proof is delivered in a clear way, where notations and motivations are clearly stated. The lines seem to be flawless as I cannot find any typos or mistakes. \n\nThis paper clearly has its significance in showing global convergence of transformer network in the original (LeCun/He) initialisation for the first time. Also, the over-parametrisation guarantee for global convergence is original, as far as I know. It also serves as a framework for later analysis in transformer.\n", "weaknesses": "The paper only considers the shallow (encoder part of the) transformer network and training in gradient descent, which is still far away from practical architecture. Also, it would be good if I could see the comparison between theoretical and empirical convergence in a plot. Sadly, the plots only support the theorems qualitatively but not quantitatively.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer 9ERM for the insightful feedback and for appreciating our theoretical analysis of Transformer and its value for the theoretical community. We address the concerns below.\n\n---\n> **Q1:** [Only the shallow (encoder part of the) Transformer network is considered. ]\n\n**A1:** We agree with the reviewer that our results focus on a relatively simple architecture of Transformers. Nevertheless, even for such a simple architecture, convergence analysis from a theoretical standpoint is still unclear. Our results provide global convergence, as a good starting point, which would pave the way for analyzing practical architectures, as discussed in Appendix C.7.  \n\nBesides, in our revised version, we have already included the residual block into our analysis, to capture the fundamental components of Transformer as much as possible, see [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA).\n\nFurthermore, our results on shallow Transformers can still provide some findings and guidelines for practical Transformers, e.g., the impact of initialization schemes, input format, and scaling, see section 4.5 for details.\n\n---\n> **Q2:** [The comparison between theoretical and empirical convergence. The quantity $\\alpha$ in Proposition 1 seems to be sub-optimal. ]\n\n**A2:** We agree with the reviewer that the lower bound $\\alpha$ could be loose though the order of $\\alpha$ is tight. This is because our derivation involves some constants in the lower bound of $\\alpha$. We believe that these constants can be improved under a refined analysis. \n\nAccording to your suggestions, in Fig.5 of the one-page pdf, we follow the experiment at sec.5.1 and plot the comparison between the empirical convergence and the quality $\\left(1-\\frac{\\gamma \\alpha^2}{2}\\right)^t \\ell (\\theta^0)$. We can see that Proposition 1 provides an upper bound for the empirical loss: the order is the same (both of them are straight lines) but the slope is different. A refined analysis (e.g., smaller constant) on the lower bound of $\\alpha$ leads to a better estimation.\n\n\n---\n> **Q3:** [In the proof of Lemma 1, you say by LLN we have the second equality in line (67). But the entries in $\\beta'_j \\beta_i \\frac{1}{d_s}\\bf {1}\\_{d_s}$ by taking the limit before the product?]\n\n**A3:** The limit is taken after calculating the inner product of the Jacobian, as can be seen in the proof of Lemma 1, e.g., line 787.\n\n---\n\nWe hope that our responses have alrady fixed the issues that the reviewer concerns. We are happy to discuss further questions if any.\n"}, {"review_id": "Is0PzFZW5y", "reviewer": "Reviewer_p3n5", "summary": "This paper presents global convergence results of shallow transformers (one self-attention layer followed by an MLP layer). The authors consider two different scalings for the factor $\\tau_0$ used in the attention matrix as a function of $d_m$, the number of rows of $W_Q$ and $W_K$. Specifically, they prove that under LeCun/He/NTK initializations for the parameters of the Transformer, scaling  $\\tau_0$ as $d_m^{-1/2}$ requires quadratic over parametrization ($d_m = \\Omega(N^2) $) to achieve convergence, while the $d_m^{-1}$ requires only linear over parametrization ($d_m = \\Omega(N) $), though the later has a worse convergence rate. The paper claims that it is due the the lack of ability to capture pairwise interaction in the  $\\tau_0 = d_m^{-1}$ scaling. Numerical experiments are presented on synthetic data to corroborate their theoretical findings for the $d_m^{-1/2}$  scaling, as well as on MNIST, where the impact of choosing $\\tau_0 = d_m^{-1/2}$ or $\\tau_0 = d_m^{-1}$ on the training dynamics is discussed.", "questions": "1. The quartic dependency of $N$ in the dimension $d$ is not discussed in Theorem 2. Is this a realistic assumption ? \n\n2. Could the authors provide insights on why the $d_m^{-1}$ scaling, which leads the Attention module to behave as a pooling layer, requires only linear over parametrization ? \n\n3. What happens if the Attention module is explicitly replaced by a pooling layer, that is $A_1 = ( 1 / d_s ) \\times 1_{d_s \\times d_s}X W_V^T$? \n\n4. Could the authors comment the results of the synthetic data experiment for the $d_m^{-1}$ scaling ? And compare with the $d_m^{-1/2}$ scaling ?\n\n5. Could the authors provide additional details for the MNIST experiment ?  Is it a classification task ? What is the test-accuracy of the model ? What is the training loss when setting the attention module to the identity ?\n\n6. Is there a setup where we should chose the $d_m^{-1}$ scaling ?", "limitations": "Yes (Appendix F).", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. This papers tackles the significant question of global convergence of Transformers.\n\n2. The model considered is a realistic one, where both the attention module (involving pair-wise interaction between the $X^{i}$), and the MLP (applied independently on each Self-attention$(X)^{i}$ are considered (Eqs (1.1) and (1.2)). \n\n3. The assumptions on the data-distribution are clearly discussed and Assumption 3. is verified in practice. Likewise, the assumptions of the parameters of the Transformer are clearly discussed.\n\n4. Theorems 1 et 2 are strong results proving the global convergence of gradient descent of the shallow Transformer model for the general regression task studied in the paper.\n\n5. The paper is focussed on the choice of the scaling factor inside the attention module, which is a forward steps towards understanding its impact. The theoretical findings as well as the experiments of the paper underlines the advantage of using the $d_m^{-1/2}$ scaling.\n\n6. Numerical experiments are conducted to support the theoretical claims.", "weaknesses": "1. Presentation:\n\na. In my opinion the biggest weakness of the paper is its lack of clarity. The paper is not really well written, hard to read, with typos and missing words (e.g. l. 37 -42, 129, 152, 222). The paper would benefit from being rewritten in a clearer and more pedagogical manner. \n\n2. Theoretical results:\n\na. I read in Table 3 that SoftMax + ReLU leads to linear over parametrization. From the paper, we understand that this requires the scaling $d_m^{-1}$. However, as the authors remark it, this implies that the SoftMax operator will behave as a pooling layer. The link between the linear over-parametrization and the attention module behaving as a pooling layer is not clear. \n\nb. Lines 68 to 71. It is not clear whether the mentioned ReLU acts on each element of the sequence separately or mixes them together. \n\n3. Experiments:\n\na. The experiment on the synthetic data in the main paper only focuses on the $d_m^{-1/2}$ scaling. I see that  the authors provide the results for the $d_m^{-1}$ scaling in the appendix. A discussion on the impact of the chosen scaling on the performance of the model on the synthetic data experiment is missing in the main paper.\n\nb. The paper focuses on regression while the experiment on MNIST seems to be a classification task. What plays the role of $y_n$ in this experiment ?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer p3n5 for the insightful feedback and for appreciating the significance of the paper. We address the concerns below.\n\n---\n> **Q1:** [The biggest weakness of the paper is its lack of clarity.]\t\n\n**A1:** We appreciate the reviewer pointing out some typos. According to your suggestions, we have already corrected typos and rephrased the expression of this paper in the final version for better understanding. For example, line 129 is changed as follows: \n\n\"We consider the encoder of Transformer, which can be applied to regression or classification ~~task~~ **tasks**.\"\n\nApart from this, we may kindly remind the reviewer that\n- Reviewer YNcT agrees that ‘The paper is organized, well-written, and easy to follow in general.’\n- Reviewer n1Tb and 9ERM give the score of 4 (excellent) for the presentation.\n\nWe hope the revised version would have a good presentation from the reviewer's side and will appreciate it if the reviewer can reconsider this “**biggest weakness**” for evaluation.\n\n---\n> **Q2:** [The link between linear over-parametrization and attention module behaving as a pooling layer is not clear.]\n\n**A2:** Due to the $d_m^{-1}$ scaling, the attention module behaves similarly to a pooling layer according to the law of large numbers, as mentioned in the introduction. In this case, the nonlinearity on $X$ disappears and thus the minimum eigenvalue of ${\\Phi} {\\Phi}^\\top$ can be estimated via $XX^{T}$, see lines 698 and 704 for detail.\nAccordingly, this leads to the minimum eigenvalue in the order of $\\Theta(N/d)$, and thus linear over-parameterization is enough, see Lemma 17 for detail. We will make it clear in our final version.\n\n---\n> **Q3:** [Not clear: ReLU acts on each element of the sequence separately or mixes them together (Lines 68 to 71.)]\n\n**A3:** Throughout the paper, ReLU is applied for each element of the sequence. In line 68, the self-attention layer is changed to one ReLU layer, i.e., in Eq. (1.1) we have ${A}_{1} = max(0, {X} {W}_1^{T})$ in an element-wise way.\n\n---\n> **Q4:** [The quartic dependency of $N$ in the dimension $d$  is not discussed in Theorem 2. Is this a realistic assumption?]\n\n**A4:** It can be achieved in practice in some cases. For example, the embedding of the tokens is usually chosen from 8-D to 1024-D [3]. Therefore, it is realistic when the token embedding is small. More importantly, for $d_m^{-1/2}$ scaling, our analysis only requires $N>d$, which is more general and valid for practical datasets.\n\n---\n> **Q5:** [What happens if the attention module is replaced by a pooling layer?]\n\n**A5:** We answer this question both empirically and theoretically:\n- In practice: According to the reviewer’s suggestion, we also conduct the classification task (using cross-entropy loss) on MNIST of the shallow Transformer as well as the setting that the attention module is identity. As shown in Fig. 2 (training loss) and Fig. 3 (test accuracy) in the one-page pdf,  training loss increases, and test accuracy drops.\n- In theory: On one hand, the network after replacing the self-attention can be considered as a 2-layer fully-connected ReLU network that still requires linear quadratic-parametrization under LeCUN initialization. On the other hand, we should note that the considered architecture includes a softmax layer and a 2-layer fully-connected ReLU network,  if the self-attention mechanism is substituted by a fully-connected ReLU layer, it will require cubic over-parameterization for global convergence, as discussed in the paper.\n\n\n---\n> **Q6:** [More experimental details: regression and classification, synthetic data and MNIST, different scalings.]\n\n**A6:** In our work, $y_n$ is the classification label of MNIST but regression on MNIST is still doable via the MSE loss. Taking two classes of MNIST for binary classification as an example, the output of our regression is taken by $sgn(f)$ and the training error is given by $\\sum_{n=1}^N [sgn(f(x_n)) - y_n]^2$. This is commonly used in theory via the MSE loss, e.g., [4,5]. \n- Regression results: there is no significant difference between these two scaling schemes under this synthetic data. The task is too simple so the optimization process is easy. But for a complex regression task, e.g., regression on MNIST, as shown in Fig.4 in the one-page pdf, the convergence speed of $d_m^{-1/2}$  is faster than $d_m^{-1}$.\n\n- Classification results: According to the reviewer’s suggestion, we also conduct the classification task (using cross-entropy loss) on MNIST of the shallow Transformer. As shown in Fig. 2 (training loss) and Fig. 3 (test accuracy) in the one-page pdf, the $d_m^{-1/2}$ setting achieves faster convergence in training loss and higher test accuracy.\n\nThese separation results on both classification and regression tasks provide good justification for the $d_m^{-1/2}$ setting. We will add these details about the experiments in our revised version. The source code has already been sent to AC.\n\n---\n> **Q7:** [Is there a setup where we should choose $d_m^{-1}$?]\n\n**A7:** Yes. The analysis for the scaling $d_m^{-1}$ in Transformer arises from NTK literature [1,2] from the theoretical side. As mentioned in A6, the $d_m^{-1}$ scaling is a feasible choice under a small $d_m$ setting (or some simple tasks as mentioned in A6) in practice.\n\nWe hope that our responses have clarified the questions of the reviewer. If there are any remaining questions, we are happy to discuss them further. \n\n---\n\n### References\n[1] Yang, et al. \"Feature learning in infinite-width neural networks.\" ICML, 2021.\n\n[2] Jacot, et al. \"Neural tangent kernel: Convergence and generalization in neural networks.\" NeurIPS, 2018.\n\n[3] see tensorflow website: the guide of **word_embeddings**\n\n[4] Rahaman, et al. \"On the spectral bias of neural networks.\" ICML, 2019.\n\n[5] Liao, et al. \"A random matrix analysis of random fourier features: beyond the gaussian kernel, a precise phase transition, and the corresponding double descent.\" NeurIPS, 2020."}, {"review_id": "MQV5lclcxP", "reviewer": "Reviewer_sZxG", "summary": "This paper studies transformer networks consisting of one layer, with average pooling and a scalar output.  The authors consider encoder type softmax attention (unmasked) and show convergence to a global minimizer of the loss.\nThe results hold for the cases that the temperature/scaling inside the softmax $\\tau_0$ is either inversely proportional to the overparameterization of the weight matrices, or the square root of it. The results also hold for different types of initialization He/Le-Cun. \nFinally, the authors also provide an NTK based analysis.", "questions": "1. In lines 635,637, do the authors mean the triangle inequality? \n2. Why the authors study these two specific scaling schemes of $d_m^{-1/2}$ and $d_m^{-1}$? It would be interesting to have an analysis dependent on the scaling parameter $\\tau_0$.\n3. It is not very clear also what step-sizes could be permitted or constants $\\hat{c}$. To my understanding the step-size needs to be inversely proportional to the number of data samples? It would be great if the authors could include some examples, in which the probability bounds are not vacuous and the order of the step-size chosen.\n4. It would be great if the authors could include also the order of the constants $C_Q,C_K,C_V,C_O$ (if they cannot be chosen arbitrarily) or if they could point out the line in which they are defined.\n5. In line 640, how is it implied that $\\sigma_{min}(F_{pre}) \\geq \\alpha/2$?\n\n", "limitations": "This work has no potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "To the best of my knowledge this is the first theoretical work on the convergence of transformers and could potentially lead to further investigate the convergence properties of transformers and bring insights into the way they work.  It is also interesting that the authors study different initialization schemes and provide comparisons between them. ", "weaknesses": "Except for some questions I have (see below), my main concern is the formulation of the transformer model in this analysis. Specifically,\n\n1. The paper does not study sequence to sequence models as encoder transformers normally are. Thus, I find the formulation of outputting a scalar inconsistent with practice. \n2. The ReLU layers do not have any matrix after the application of the ReLU. There are also no residual connections. \n3. The pooling layer to the best of my knowledge is used only in ViT models. \n\nI also think that it would be beneficial for the paper, to include more details about proof techniques for one of the cases (this could be possible by reducing some of the text in the rest of the sections). \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful to the reviewer sZxG for appreciating that this is the first theoretical work on the convergence of Transformers and its significance for practice.\n\n---\n\n> **Q1:** [No result on sequence to sequence (seq2seq) model. The formulation of outputting a scalar inconsistent with practice. The pooling layer to the best of my knowledge is used only in ViT models. No residual connections.]\n\n**A1:** We agree with the reviewer that the current formulation does not consider seq2seq Transformer. Our model formulation as well as theoretical analysis is motivated by ViT models used for image tasks. Besides, according to the reviewer’s suggestions, our result can be extended to a shallow Transformer with residual connection, see [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA).  We hope our analysis of the self-attention module (e.g., scaling, softmax) will lay a foundation for the analysis of more general Transformer architectures.\n\n---\n\n> **Q2:** [Include more details about proof techniques for one of the cases (this could be possible by reducing some of the text in the rest of the sections).]\n\n**A2:** We are thankful to the reviewer for the suggestion, we will add the proof sketch, as present in the [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA), into the final version.\n\n---\n\n> **Q3:** [In lines 635,637, do the authors mean the triangle inequality? In line 640, how is it implied that $\\sigma_{min}(F_{pre}) \\geq \\alpha/2$]\n\n**A3:**  We use Weyl’s inequality, (a general version of triangle inequality) and the relationship between spectral norm and Frobenius norm. To be specific,\n\n\n[lemma] [Weyl’s inequality [6]]\n\n\nGiven two matrices $X, Y \\in \\mathbb{R}^{p\\times q}$ with the respective singular value ordered as follows: $\\sigma_1(X)\\geq\\ldots\\geq\\sigma_r(X)$ and $\\sigma_1(Y)\\geq\\ldots\\geq\\sigma_r(Y)$, where $r=\\min(p,q).$ Then we have $\\max_{i\\in[r]} |\\sigma_i(X)-\\sigma_i(Y)|\\leq ||X-Y||_2.$ \n\n\n[end lemma]\n\n\nIn lines 635 and 637, we choose $i = 1$ in Weyl’s inequality and use the inequality between spectral norm and Frobenius norm to obtain $||X||_2 \\le ||Y||_2+ ||X-Y||_2 \\le ||Y||_2 + ||X-Y||_F$. Thus, this can be also considered as a triangle inequality for the spectral norm. \n\n\nIn line 640, we choose $i = r$ in Weyl’s inequality and use the inequality between spectral norm and Frobenius norm to obtain $\\sigma\\_{\\min}(X) \\ge \\sigma\\_{\\min}(Y) - ||X-Y||_2 \\ge \\sigma\\_{\\min}(Y) - ||X-Y||_F$.\n\n---\n\n> **Q4:** [Why the authors study these two specific scaling schemes of $d_m^{-1/2}$ and $d_m^{-1}$? It would be interesting to have an analysis dependent on the scaling parameter$\\tau_0$ .]\n\n**A4:** These two scaling schemes are widely used in practice ($d_m^{-1/2}$ in [4,5]), and theoretical analysis for the neural network limiting ($d_m^{-1}$ in [1,2,3,7]). Our proof framework can facilitate the analysis of general $\\tau_0$, e.g., $\\tau_0 = d_m^{-a}$ with $a \\in [0,1]$.\n\n---\n\n> **Q5:** [It is not very clear also what step-sizes could be permitted or constants . To my understanding the step-size needs to be inversely proportional to the number of data samples? It would be great if the authors could include some examples, in which the probability bounds are not vacuous and the order of the step-size chosen.]\n\n**A5:** The step size is inversely proportional to $N^{1/2}$, where N is the number of data samples. To see why, in line 652, we define $C =  c_1 c_2 + 2 c_3 \\ell(\\theta^0)$. According to line 648 , we can see that $c_1$ is proportional to $\\rho$, which is defined as $\\rho = N^{1/2} d_s^{3/2} \\tau_1 C_x $ in line 170.  Therefore, C is proportional to N^{½}. Since in Proposition 1, we require the step-size $\\gamma < 1/C$, therefore, the step-size is inversely proportional to $N^{1/2}$. We will add a remark for this in our final version.\nThe probability is non-vacuous since in Lemma 6 we select $\\zeta = d_m^{-1/2}$ so that the probability is at least $1-2\\exp^{-d_m/2}$. In practice, even for $d_m = 16$, then the probability is at least 0.99$. Larger width leads to a higher probability.\n\n---\n\n> **Q6:** [It would be great if the authors could include also the order of the constants $C_Q, C_K, C_V, C_O $ (if they cannot be chosen arbitrarily) or if they could point out the line in which they are defined.]\n\n**A6:** $C_Q = C_K = C_V = C_O$ is in constant order, as can be seen in line 673 (appendix). We directly choose one for simplicity. We will make it clear in our final version.\n\n---\n\nWe hope that our responses have addressed the concerns of the reviewer. If there are any remaining questions, we are happy to discuss them further. \n\n---\n\n### References\n[1] Du, et al. \"Gradient Descent Provably Optimizes Over-parameterized Neural Networks.\" ICLR, 2019.\n\n[2] Jacot, et al. \"Neural tangent kernel: Convergence and generalization in neural networks.\" NeurIPS, 2018.\n\n[3] Hron, et al. \"Infinite attention: NNGP and NTK for deep attention networks.\" ICML, 2020.\n\n[4] Vaswani, et al. \"Attention is all you need.\" NeurIPS, 2017.\n\n[5] Dosovitskiy, et al. \"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.\" ICLR, 2021.\n\n[6] Stewart, Gilbert W. Perturbation theory for the singular value decomposition. 1998.\n\n[7] Yang, et al. \"Feature learning in infinite-width neural networks.\" ICML, 2021.\n"}, {"review_id": "TBucojKiY3", "reviewer": "Reviewer_UvAq", "summary": "This paper proves the convergence rate of GD on one-layer transformer training. In their setting, the transformer is composed by a single attention layer, followed by relu unit and a fixed fully connected layer. The linear convergence rate is achieved. The proof technique follows the standard neural network proof [Nguyen 2021], where under some conditions of minimum singualr value of gram matrix, they inductively show that the parameter drift is controlled, and hence the GD dynamic is easy to study. The main meat is to show that the minimum singular value of gram matrix is lower bounded.\n\n \n[Nguyen 2021] Nguyen, Quynh. \"On the proof of global convergence of gradient descent for deep relu networks with linear widths.\" In International Conference on Machine Learning, pp. 8056-8062. PMLR, 2021.", "questions": "1. Since original framework in [Nguyen 2021] works when all layers are jointly trained, is it possible to extend your proof to that regime as well?\n\n2. In Corollary 1, you said that when d_s and d is chosen to be some value, the conditions 2 and 3 cannot satisfied, and the convergence will fail. I think this is not a good reasoning. It could be that your framework cannot prove the convergence, but does not mean it cannot actually converge. A more convincing reasoning will be showing the lower bound of the loss, to show the failure of convergence.", "limitations": "The technical limitations are stated in weakness part. I do not see any negative social impact of this work.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "To my best knowledge, this is the first convergence proof of transformer. The rate makes sense and proof roadmap is clear to me. The first part of the proof idea is standard, that under proper conditions for initialization and over-parameterization, the model parameter will stay in linear regime, and hence the convergence is easy to control. Then the second part is that they show that the conditions can be satisfied for some initialization scheme, by showing that the minimum singular value of gram matrix is bounded. The most interesting and novel part is to derive such lower bound.", "weaknesses": "1. As I mentioned above, the first part of the proof (proposition 1) pretty much follows [Nguyen 2021]. \n\n2. The model studied is bit toy. In practice, the residual block of transformer is necessary, but it is missing in this paper.\n\n3. They assume the fully connected layer W_H is fixed to identical matrix, and argue that since W_V and W_H are adjacent, this is equavilent to W_H is trainable, However, in GD dynamic, I do not think they are equivalent. For example, in deep linear neural network, training all layer by GD is not equavilent to only training one matrix. Or if they are equivalent, there should be rigorous proof.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful to the reviewer UvAq for appreciating that this is the first theoretical work on the convergence of Transformers and the insightful feedback. We address the concerns below.\n\n--- \n> **Q1:** [The first part of the proof (proposition 1) pretty much follows [3].]\n\n**A1:** The framework of the proof for proposition 1 follows the variant of Polyak-Lojasiewicz (PL) inequality in traditional convergence analysis. But the studied structure is generally more complicated than [3] due to the softmax function, and the proof in Lemma 7-11 tailored for Transformer. \n\nMore importantly, to give the lower bound for $\\alpha_0$ in proposition 1 is another technical difficulty, which was later addressed in Theorem 1 when considering several particular initializations. See the [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA) for detailed clarification about this issue.\n\n\n---\n> **Q2:** [The model studied is bit toy. In practice, the residual block of Transformer is necessary, but it is missing in this paper.]\n\n**A2:** We are thankful for the suggestion of the reviewer. According to the reviewer’s suggestions, we have already extended our result to a shallow Transformer with residual connections, see [general response](https://openreview.net/forum?id=8ZveVHfmIE&noteId=hT7qrLWaaA). \n\nWe hope our analysis of the self-attention module (e.g., scaling, softmax) will lay a foundation for the analysis of more general Transformer architectures.\n\n\n---\n> **Q3:** [They assume the fully connected layer W_H is fixed to the identical matrix, and argue that since W_V and W_H are adjacent, this is equivalent to W_H is trainable, However, in GD dynamic, I do not think they are equivalent. Since original framework in [3] works when all layers are jointly trained, is it possible to extend your proof to that regime as well?]\n\n**A3:** We agree with the reviewer that these two cases are not equivalent under gradient descent. We adopt one learnable parameter for training just for ease of analysis, as mentioned in line 138. The same technique is also used in [1] for analyzing Transformer. According to your suggestions, we will make it clear in our final version by adding the following sentence in line 139.\n\n\"Note that we mix $W_V$ and $W_H$ together for ease of the analysis but it does not mean its training dynamics is the same as we jointly train these two adjacent matrices.\"\n\nBesides, we think it is possible to extend the proof and we illustrate the high-level idea here. If we consider all layers have been trained. Then we need to consider one more parameter $\\boldsymbol{W}_H$.  Two main parts of the analysis have to be changed. The first part is Proposition 1. Note that in line 628, $\\boldsymbol{f}$ becomes as follows: \n\n$\\boldsymbol{f}= \\tau\\_1\n\\boldsymbol{w}\\_O^\\top\n\\sum_{i=1}^{d\\_s}\\sigma\\_r\n\\left(\\boldsymbol{W}\\_H \\boldsymbol{W}\\_V\\boldsymbol{X}^\\top \\boldsymbol{\\beta}\\_i\\right)$. \n\nThen lemma 7- 8, 10, and 11 remain unchanged.\n In Lemma 9, one more step of triangle inequality is required to decouple $\\boldsymbol{W}\\_H^{t’}$ and  $\\boldsymbol{W}\\_H^{t}$. \nIn Lemma 12, 13, 14, we need an additional bound of the gradient norm and the Lipschitz constant for $\\boldsymbol{W}\\_H$ using the same technique as we bound the others. Then the proof for Proposition 1 can be finished. The second part that needs to be changed is the minimum singular value of $\\boldsymbol{F}\\_{pre}$. One can analyze the centering features w.r.t. $\\boldsymbol{W}\\_H$ following the strategy in [2] and give a lower bound for the minimum singular value of $\\boldsymbol{F}\\_{pre}$.\n\n---\n\n> **Q4:** [Fail of convergence in Corollary 1.]\n\n**A4:** Our current proof provides an upper bound of the loss function value. To show “fail of convergence”, we agree with the reviewer that the lower bound of loss function is needed. This is a quite difficult problem, and accordingly, based on the reviewer’s suggestions, we modify the corollary 1 for rigorous illustration as follows:\n\n---\n\n**Corollary** (Convergence of vector input)\n\nConsidering LeCun initialization with $\\tau_0 = d_m^{-1}$ scaling,\ngiven vector input $\\boldsymbol{x}\\in \\mathbb{R} ^{\\tilde{d}}$, \nif one feeds the input to Transformer by setting $d_s = 1, d=\\tilde{d}$, then \ntraining with GD can converge to a global minimum. However, if one sets $d_s =\\tilde{d}, d=1$, \nthe conditions in Eq.2-3 do not hold, **its convergence cannot be guranteed by our theory**.\n\n---\n\nThough the estimation (lower bound of the loss function) under Transformer is difficult, we find that it is possible to consider a simplified case, i.e., linear regression, which still shares a similar spirit to our problem.\n\n---\n\nConsider the loss $\\ell = 1/2 || \\boldsymbol{y} - \\boldsymbol{X} \\boldsymbol{w} ||_2^2 $, where $\\boldsymbol{X} \\in \\mathbb{R}^{N \\times d}$. In this case, the neural tangent kernel (NTK) is $\\boldsymbol{X} \\boldsymbol{X}^\\top$. If the NTK is a rank-one matrix, then the minimum eigenvalue is zero. In this case one can see $\\boldsymbol{X} \\in \\mathbb{R}^{N \\times 1}$. Therefore, if the augmented matrix $[\\boldsymbol{X}, \\boldsymbol{y}]$ is not rank-one, which is standard in practice, then GD training can not converge to zero loss since there is no solution.\n\n---\n\nThis would motivate us to rethink Corollary 1 under the Transformer setting as future work. We will add a remark about this in our revised version.\n\n--- \nWe hope that our responses have addressed the concerns of the reviewer. If there are any remaining questions, we are happy to discuss them further.\n\n--- \n### References\n[1] Yang, et al. \"Transformers from an optimization perspective.\" Neurips 2022.\n\n[2] Nguyen, et al. \"Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep relu networks.\" ICML, 2021.\n\n[3] Nguyen, et al. \"On the proof of global convergence of gradient descent for deep relu networks with linear widths.\" ICML, 2021.\n"}]

Author rebuttal:
**General Response**

Dear reviewers, 

We appreciate your insightful comments. Below, we address three core topics raised by reviewers and then answer individually to the questions of each reviewer. 

Concretely, below we make the following responses:
- discuss the extension to the residual Transformer, as pointed out by reviewers sZxG and UvAq 
- highlight the technique difficulty when compared to the prior analysis, as requested by reviewers n1Tb and UvAq
- add the proof sketch for better understanding, following the suggestion of reviewers YNcT and sZxG

---

> **Q1:** Extension to residual Transformer (sZxG and UvAq)

**A1:** Our proof framework is general and can handle the following residual Transformer. Here we give a proof sketch to show how to achieve this.

We consider the residual block in the self-attention layer: $ \bf{A}_1= \text{Self-attention}(\bf{X}) \triangleq \sigma_s \left( \tau_0 (\bf{X}{\bf{W}}_Q^\top) \left(\bf{X} \bf{W}_K ^\top\right)^\top \right) \left( \bf{X} \bf{W}_V^\top \right) + \bf{X}$,

which leads to the output $\bf{f}(\bf X)$ in line 628 as $\bf{f}(\bf X)= \tau\_1 \bf{w}\_O^\top \sum\_{i=1}^{d\_s} \sigma\_r
\left(\bf{W}\_V\bf{X}^\top \bf{\beta}\_i + {(\bf{X}^{(i:)})}^\top \right)$.

To prove the convergence of the above residual Transformer, we need:

- To prove the optimization properties of the loss function (e.g., almost smooth), the proof is nearly the same as the current version apart from the additional term to upper bound $||{\bf{X}^{(i:)}}||_2$ in Lemma 9.
- To prove the lower bound of $\alpha_0$, we need to estimate the output difference on the shallow Transformer with/without residual block in Lemma 15. Using the Lipschitz continuity of softmax, ReLU, this difference can be well controlled, and thus the lower bound of $\alpha_0$ can be still achieved with an extra constant. 

In the final version, we will detail this convergence result after Sec.C.7 in the appendix. We hope our analysis will lay a foundation for the analysis of more general Transformer architectures.


---

> **Q2:** Technical difficulty (n1Tb and UvAq).

**A2:** Handling the softmax function and the $d_m^{-1}, d_m^{-1/2}$ scaling schemes beyond NTK initialization in Transformer are precisely **the two challenges**. Previous convergence analysis, e.g., [1-2],  cannot be applied to our setting because of the following two issues:
 - Different from classical activation functions, e.g., ReLU, in softmax each element of the output depends on all input. To tackle the interplay between dimensions, we build the connection between $\Phi \Phi^{T}$ and $XX^T$ (see Lemma 15-17) to disentangle the nonlinear softmax function, where $\Phi$ contains information on the output of the softmax. By doing so, the lower bound of the minimum eigenvalue of  $\Phi \Phi^{T}$  can be well controlled by $XX^T$ and the output of softmax. Accordingly, a lower bound on the minimum singular value of ${F}\_{pre}$ in Proposition 1 can be obtained. 
- Regarding different initializations and scaling, previous NTK-based analysis is only valid under the $d_m^{-1}$ setting (the softmax degenerates to an all-one vector) but is inapplicable to the realistic $d_m^{-1/2}$ setting, as discussed in the introduction. To tackle this issue, we analyze the input/output of softmax under LeCun/He initialization (see Lemma 15) and identify the optimization properties of the loss function (see the proof of Proposition 1) for global convergence under the finite-width setting.

We will include this discussion in our final version and hope this clarification would address the reviewers’ concerns.

--- 

> **Q3:** Proof sketch (YNcT and sZxG).

**A3:** Here we present the proof sketch of our global convergence results under various initialization schemes.

[**Proof sketch**] The main idea of our convergence analysis is based on the variant of Polyak-Lojasiewicz (PL) inequality, i.e., (under simple calculation)

$|| \nabla \ell({\theta}) ||\_2^2 \ge 2 \lambda\_{\min}({K}) \ell({\theta}) \ge 2 \lambda\_{\min}({F}\_{pre} {F}_{pre}^{\top}) \ell({\theta})$.

 Thus if the minimum singular value of ${F}_{pre}$ is strictly greater than 0, then minimizing the gradient on the LHS will drive the loss to zero. 

To this end, we take Proposition 1 as an example. It can be split into two parts:
- By induction, at every time step, each parameter in Transformer can be bounded w.h.p;  the minimum singular value of ${F}\_{pre}$ is bounded away for some positive quality at the initialization point.
- We prove that the Lipschitzness of the network gradient (Lemma 14), which means the loss function is almost smooth.
By doing so, combining the above two results, the loss function global convergence can be achieved.

Furthermore, in Theorem 1 and Theorem 2, we consider several particular initializations and scaling and aim to provide a positive lower bound for ${F_{pre}}$, satisfying the conditions in Proposition 1. 

We will include it in our revised version for better understanding.

---

### References
[1] Du, et al. "Gradient Descent Provably Optimizes Over-parameterized Neural Networks." ICLR, 2018.

[2] Nguyen, et al. "On the proof of global convergence of gradient descent for deep relu networks with linear widths." ICML, 2021.

Meta-review: The paper derives convergence guarantees for a shallow Transformer network (only the encoder part) under different weight initializations/scalings. The analysis relies on the NTK matrix of Transformers derived in prior work and it derives a bound on the minimum eigenvalue of the NTK, similar to the work of Nguyen and Mondelli, 2020. The contribution of this work is mostly theoretical and might not have important practical consequences but I think it's a welcomed contribution for the community.

The reviewers are somewhat positive about the paper. Some concerns have been addressed in the rebuttal, leading some reviewers to increase their scores. However, the paper needs several improvements including adding more intuitive explanations of the analysis and emphasizing more the limitations. For instance, the analysis applies to a shallow encoder only (several theoretical works have studied similar models, it might be worth extending the prior work on this as well). The reviewers also made several criticisms regarding the experimental part. Finally, the title is in my view too general and should be updated to more closely reflect the setting addressed in the paper.

Overall, the paper has some merits so I recommend acceptance, but I strongly encourage the authors to make a somewhat major revision of the paper to address the problems raised by the reviewers. This will undoubtedly increase the impact of the paper so that's also in the interest of the authors.
Decision: Accept (poster)