Title: On the Convergence of Encoder-only Shallow Transformers
Venue: NeurIPS 2023
Accepted: yes
Keywords: Transformer

convergence

scaling

initialization

over-parameterization
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

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.

1 Introduction

Transformers [Vaswani et al., 2017] have demonstrated unparalleled success in influential applications [Devlin et al., 2019, Brown et al., 2020, Wang et al., 2018, Dosovitskiy et al., 2021, Liu et al., 2022b]. A fundamental theoretical topic concerns the global convergence, i.e., the training dynamics of Transformers, which would be helpful for further analysis, e.g., in-context learning [von Oswald et al., 2022, Akyürek et al., 2023], generalization [Li et al., 2023]. In fact, even within a 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 ∈ R d s × d ( d s is the number of tokens and d is the feature dimension of each token), a self-attention mechanism is defined as:

where σ s is the row-wise softmax function, τ 0 ∈ R + is the scaling factor, and the learnable weights are W Q , W K , W V ∈ R d m × d with the width d m . Given X , the input of softmax depends on τ 0 W ⊤ Q W K , including the scaling factor τ 0 and initialization schemes for learnable parameters, and thus determines the output of softmax and then affects the performance of Transformers in both

Fanghui Liu ∗

Volkan Cevher LIONS, EPFL

theory and practice. There are several scaling schemes in previous literature. For instance, given W Q and W K initialized by standard Gaussian, the scaling factor τ 0 is chosen by

τ 0 = d -1 / 2 m in the original Transformer [Vaswani et al., 2017]: each element in τ 0 W ⊤ Q W K is a random variable with mean 0 and variance 1 . This scaling avoids the blow-up of value inside softmax as d m increases [Hron et al., 2020].

τ 0 = d -1 m : 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 convolutional) networks under an infinite-width setting d m → ∞ . However, for Transformer, if one uses this scaling under the infinite-width setting, then by the law of large numbers, we have lim d m →∞ τ 0 [ W ⊤ Q W K ] ( ij ) = 0 . As a result, the input of softmax is zero and the softmax degenerates to a pooling layer. That means, the non-linearity is missing, which motivates researchers to carefully rethink this setting.

For instance, under the τ 0 = d -1 m 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 τ 0 = d -1 / 2 m and infinite width setting for simplicity.

Clearly, there exists a gap between theoretical analysis and practical architectures on the use of softmax, and accordingly, this leads to the following open question:

How can we ensure the global convergence of Transformers under a realistic setting?

The primary contribution of this work is to establish the convergence theory of shallow Transformer under a realistic setting. Despite its shallow and encoder-only architecture, our 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 a linear output layer, cf. Eq. (1.2). We adopt the τ 0 = d -1 / 2 m 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 τ 0 = d -1 m setting (as well as the NTK based analysis) is also studied, which facilitates a comprehensive comparison. Our theoretical results demonstrate notable separations among scaling settings, initializations, and architectures as below:

Scaling: The global convergence can be achieved under both τ 0 = d -1 / 2 m and τ 0 = d -1 m . Nevertheless, as suggested by our theory: for a small d m , there is no significant difference for these two scaling schemes on the convergence; but for a large enough d m , the τ 0 = d -1 / 2 m scaling admits a faster convergence rate of Transformers than that with τ 0 = d -1 m . Interestingly, under this τ 0 = d -1 m setting, our theory also demonstrates the separation on the convergence result, depending on whether the input is formed along sequence dimension ( d = 1 ) or embedding dimension ( d s = 1 ).

Initialization: Under LeCun and He initialization, our shallow Transformer admits a faster convergence rate than the NTK initialization. This could be an explanation for the seldom usage of NTK initialization for Transformer training in practice.

Architecture: Quadratic over-parameterization is enough to ensure the global convergence of our shallow Transformer. As a comparison, if the self-attention mechanism is substituted by a feed-forward ReLU layer, our shallow Transformer is close to a three-layer fully-connected ReLU neural networks to some extent, requiring cubic over-parameterization for global convergence.

We firmly believe that our theoretical analysis takes a 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-based models [Radford et al., 2018, Brown et al., 2020].

2 Related work

Self-attention, Transformer: Regarding training dynamics, Snell et al. [2021] explain why singlehead 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. [2023] provide sample complexity of shallow Transformer to study its generalization property under a good initialization from pretrained model. The separation between the Transformer and CNN is recently explored. Jelassi et al. [2022] provably demonstrate that Vision Transformer (ViT) has the ability to learn spatial structure without additional inductive bias such as the spatial locality in CNN. Chen et al. [2022] study the loss landscape of ViT and find that ViT converges at sharper local minima than ResNet. Park and Kim [2022] show that ViT is a low-pass filter while CNN is a high-pass filter, thus, these two models can be complementary.

NTK, lazy training, Hessian: The NTK was introduced by Jacot et al. [2018] to connect the infinite-width neural network trained by gradient descent and the kernel regression. The roles of such kernel include analysis of the training dynamics of the neural network in the over-parameterization regime [Allen-Zhu et al., 2019a, Chizat et al., 2019, Du et al., 2019a,b, 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 is extended from FCNN to several typical networks including Transformer [Tirer et al., 2020, Huang et al., 2020, Arora et al., 2019b, Alemohammad et al., 2021, 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 a 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].

Over-parameterization for convergence analysis: Due to over-parameterization, neural networks (NNs) can fit arbitrary labels with zero training loss when trained with (stochastic) gradient descent (SGD), both theoretically Li and Liang [2018], Du et al. [2019b] and empirically [Zhang et al., 2017]. This leads to an interesting question in theory: how much overparameterization is enough to ensure global convergence of NNs? A common recipe for the proof of global convergence relies on the variant of Polyak-Lojasiewicz condition [Polyak, 1963, Liu et al., 2022a], NTK [Du et al., 2019b,a, Zou and Gu, 2019, Allen-Zhu et al., 2019a], or the minimum eigenvalue of the gram matrix [Nguyen, 2021, Bombari et al., 2022]. In Appendix B.3, we provide a comprehensive overview of a recent line of work that improves the over-parameterization condition for ensuring the convergence of NNs. However, the over-parameterization condition for Transformer to achieve global convergence remains elusive from existing literature and we make an initial step towards this question.

3 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 ∥ · ∥ F and ∥ · ∥ 2 to represent the Frobenius norm and the spectral norm of a matrix, respectively. The Euclidean norm of a vector is symbolized by ∥ · ∥ 2 . The superscript with brackets is used to represent the element of a vector/matrix, e.g., w ( i ) is the i th element of w . The superscript without brackets symbolizes the parameters at different training step, e.g., θ t . We denote by [ N ] = { 1 , . . . , N } for short. We use σ min ( · ) and λ min ( · ) to represent the minimum singular value and minimum eigenvalue of a matrix. The NTK matrix and hessian matrix of the network are denoted by K and H , respectively. The order notation, e.g., ˜ O , ˜ Ω , omits the logarithmic factor. More detailed notation can be found in Table 2 of the appendix.

Let X ⊆ R d s × d be a compact metric space and Y ⊆ R , where d is the dimension of each token, d s is the total sequence length of the input. The training set { ( X n , y n ) } N n =1 is assumed to be iid sampled from an unknown probability measure on 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 f : X → Y in our work, such that f ( X ; θ ) parameterized by θ is a good approximation of the label y ∈ Y corresponding to a new sample X ∈ X . We use a vector θ to denote the collection of all learnable parameters.

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 et al., 2019, Dosovitskiy et al., 2021]. Given an input X ∈ R d s × d , the model is defined as below:

where the output is f ( X ; θ ) ∈ R , τ 0 and τ 1 are two scaling factors. The ingredients of a Transformer with width d m are defined as follows:

A self-attention mechanism (Eq. (1.1)): σ s is the row-wise softmax function; the learnable parameters are W Q , W K , W V ∈ R d m × d . We employ Gaussian initialization W ( ij ) Q ∼ N (0 , η Q ) , W ( ij ) K ∼ N (0 , η K ) , W ( ij ) V ∼ N (0 , η V ) with i ∈ [ d m ] and j ∈ [ d ] . Refer to Table 1 for typical initialization examples.

A feed-forward ReLU layer (in Eq. (1.2)): σ r is the ReLU activation function; the learnable parameter is W H ∈ R d m × d m . Following Yang et al. [2022], we combine W V and W H together (by setting W H = I ) 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. (1.2)): φ indicates the column-wise 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 or theoretical analysis [Dosovitskiy et al., 2021, Yang, 2020].

An output layer (in Eq. (1.2)) with learnable parameter w O ∈ R d m , initialized by w ( i ) O ∼ N (0 , η O ) .

Remarks: 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., 2019]. 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 1, we detail the formula of LeCun initialization, He initialization, and NTK initialization.

Given N input samples { X n } N n =1 , the corresponding ground truth label, the final output of network, and the output of the last hidden layer, are denoted by:

We consider standard gradient descent (GD) training of Transformer, as illustrated in Algorithm 1. Here the squared loss is expressed as ℓ ( θ ) = 1 2 ∥ f ( θ ) -y ∥ 2 2 .

4 Main results

In this section, we study the convergence guarantee of Transformer training by GD under the squared loss. Firstly, we provide a 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 τ 0 = d -1 / 2 m in Section 4.2 and τ 0 = d -1 m 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 ∥ F ≤ √ d s C x with some positive constant C x .

Remark: This assumption is standard as we can always scale the input. The embedding of token is usually assumed to have a unit norm [Li et al., 2023], which is unrelated to d .

Now we are ready to present our 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 to a unifying analysis framework under various initialization schemes, we do not include specific probabilities here. The validation of the required conditions and probability is deferred to Sections 4.2 and 4.3, respectively.

Proposition 1. Consider the Transformer with d m ≥ N . Let C Q , C K , C V , C O be some positive constants, and define the following quantities at initialization for simplification:

The norm of the parameters:

Under Assumption 1, we additionally assume that the minimum singular value of F 0 pre , i.e., α ≜ σ min ( F 0 pre ) satisfies the following condition at initialization:

If the step size satisfies γ ≤ 1 /C with a constant C depending on ( ¯ λ Q , ¯ λ K , ¯ λ V , ¯ λ O , ℓ ( θ 0 ) , ρ, τ 0 ), then GD converges to a global minimum as follows:

Remark: The parameter α in Eqs. (2) and (3) controls the convergence rate of global convergence, and the condition will be verified in the next subsection. The step-size γ is inversely proportional to N 1 / 2 due to the construction of C in Appendix C.2.

Proof sketch: The main idea of our convergence analysis is based on the variant of PolyakLojasiewicz (PL) inequality [Polyak, 1963, Nguyen, 2021], i.e., ||∇ ℓ ( θ ) || 2 2 ≥ 2 λ min ( K ) ℓ ( θ ) ≥ 2 λ min ( F pre F ⊤ pre ) ℓ ( θ ) . 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, Proposition. 1 can be split into two parts. First, by induction, at every time step, each parameter in the 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. 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 to validate Eqs. (2) and (3) via the spectral norm estimation of ¯ λ Q , ¯ λ K , ¯ λ V , ¯ λ O and a positive lower bound for F pre in the following section.

4.2 LeCun and He initialization under the d -1 / 2 m setting

Here we aim to show that, under the LeCun/He initialization with the setting of d -1 / 2 m , the conditions in Eqs. (2) and (3) will be satisfied with high probability and scaling schemes in Table 1 and hence

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.

Remark: This assumption requires that each row X ( i, :) is linearly independent for any i ∈ [ d s ] , 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.

̸

Assumption 3. For any data pair ( X n , X n ′ ) , with n = n ′ and n, n ′ ∈ [ N ] , then we assume that P (∣ ∣ 〈 X ⊤ n X n , X ⊤ n ′ X n ′ 〉∣ ∣ ≥ t ) ≤ exp( -t ˆ c ) with some constant ˆ c > 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 n and X n ′ with n = n ′ , their inner product reflects the similarity of their respective (empirical) covariance matrix. We expect that this similarity is small with a high probability. The spirit of this assumption also exists in previous literature, e.g., Nguyen et al. [2021]. The constant ˆ c determines the decay of data dissimilarity. A larger ˆ c results in less separable data. Our assumption has no requirement on ˆ c such that ˆ c can be small enough, which allows for a general data distribution.

Verification of assumption. Here we experimentally validate this assumption under a standard language IMDB dataset [Maas

0.00

0.25

0.50

0.75

0.0

0.2

0.4

0.6

0.8

1.0

Probability

et al., 2011]. We randomly sample 100 (normalized) sentences with embedding and plot the probability (frequency) of P (∣ ∣ 〈 X ⊤ n X n , X ⊤ n ′ X n ′ 〉∣ ∣ ≥ t ) as t increases. We repeat it over 10 runs and plot the result in Figure 1. We can observe an exponential decay as t increases, which implies that our assumption is fair. Besides, the validation of Assumption 3 on image data, e.g., MNIST by ViT, is deferred to Appendix E.2.

Now we are ready to present our main theorem under the d -1 / 2 m setting. The proof is deferred to Appendix C.3.

Theorem 1. Under the setting of LeCun/He initialization in Table 1 and Assumptions 1 to 3, suppose d m ≥ d , and given τ 0 = d -1 / 2 m , d m = ˜ Ω( N 3 ) , then with probability at least 1 -8 e -d m / 2 -δ -exp( -Ω(( N -1) -ˆ c d -1 s )) for proper δ , the GD training of Transformer converges to a global minimum with sufficiently small step size γ as in Eq. (4) .

Remark: The probability relates to several randomness sources, e.g., data sampling, dissimilarity of data, and parameter initialization. The quantity exp( -Ω(( N -1) -ˆ c d -1 s )) can be small for a small enough ˆ c as discussed in Assumption 3. Further discussion on our result refers to Section 4.4 for details. Besides, our proof framework is also valid for the τ 0 = d -1 m setting. We demonstrate that d m = ˜ Ω( N 2 ) suffices to achieve global convergence, see Appendix C.3 for details.

Proof sketch: To check whether the conditions in Eqs. (2) and (3) hold, the key idea is to provide the lower bound of α . Then we upper bound ¯ λ Q , ¯ λ K , ¯ λ V , ¯ λ O 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 α . Lastly, we plug these bound into the condition Eqs. (2) and (3) in order to obtain the requirement for the width d m .

4.3 NTK initialization under the d -1 m setting

The NTK theory, as a representative application of the τ 0 = d -1 m setting, can be also used for analysis of training dynamics. We also include the NTK results in this section for self-completeness. In this section, we first derive the limiting NTK formulation of Transformer under the d -1 m scaling scheme, and then show the global convergence of Transformers.

Lemma 1. Denote Φ ⋆ =: [ 1 d s X ⊤ 1 1 d s , ..., 1 d s X ⊤ N 1 d s ] ⊤ ∈ R N × d , then the limiting NTK matrix K ∈ R N × N of Transformer under the NTK initialization with τ 0 = d -1 m has the following form:

Remark: The formulation of Φ ⋆ implies that the self-attention layer degenerates as 1 d s 1 d s × d s XW ⊤ V , i.e., the dimension missing effect as mentioned before.

Now we are ready to present our convergence result under the d -1 m scaling with the proof deferred to Appendix C.4.

Theorem 2. Under the setting of NTK initialization in Table 1 and Assumptions 1 to 3 , suppose d m = ˜ Ω( N ) , then with probability at least 1 -8 e -d m / 2 -δ -exp( -Ω(( N -1) -ˆ c d -1 s )) , the GD training of Transformer converges to a global minimum with sufficiently small γ as in Eq. (4) .

Proof sketch: The overall idea is the same as the proof of the previous theorem, i.e., we need to provide the lower bound of α . However, in this case, the limit for the output of softmax exists so that we can apply concentration inequality to lower bound the α . Lastly, we plug these bound into the condition Eqs. (2) and (3) in order to obtain the requirement for the width d m .

Besides, the stability of NTK during training allows us to build a connection on training dynamics between the Transformer (assuming a squared loss) and the kernel regression predictor. Next, in order to show that the NTK is stable during GD training, below we prove that the spectral norm of Hessian is controlled by the width.

Theorem 3 (Hessian norm is controlled by the width) . Under Assumption 1 and scaling τ 0 = d -1 m , given any fixed R > 0 , and any θ t ∈ B ( θ , R ) := { θ : ∥ θ -θ 0 ∥ 2 ≤ R } , θ 0 as the weight at initialization, then with probability at least 1 -8 e -d m / 2 , the Hessian spectral norm of Transformer satisfies: ∥ H ( θ t ) ∥ 2 ≤ O ( d -1 / 2 m ) .

Remark: By [Liu et al., 2020, Proposition 2.3], the small Hessian norm is 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 on 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 τ 0 = d -1 / 2 m and τ 0 = d -1 m . Nevertheless, as suggested by our theory, for a small d m , there is no significant difference between these two scaling schemes on the convergence; but for a large enough d m , the τ 0 = d -1 / 2 m scaling admits a faster convergence rate of Transformers than that of τ 0 = d -1 m due to a more tight estimation of the lower bound of α , see Appendix C.6 for details. The intuition behind the lower convergence rate under the setting of large width and τ 0 = d -1 m 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 1 from the perspective of dimension missing : self-attention ( X ) degenerates as 1 d s 1 d s × d s XW ⊤ V .

The result under the τ 0 = d -1 m setting requires weaker over-parameterization than the τ 0 = d -1 / 2 m setting. Nevertheless, we do not claim that τ 0 = d -1 m is better than τ 0 = d -1 / 2 m . This is because, under the over-parameterization regime, the scaling τ 0 = d -1 m makes the self-attention layer close to a pooling layer. This analysis loses the ability to 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 d m and λ 0 , we can prove that quadratic (cubic) over-parameterization is required for d -1 m ( d -1 / 2 m ) scaling. This quadratic over-parameterization requirement could be relaxed if a better relationship is given while it is still unclear and beyond our proof technique.

Comparison on initializations: Though our results achieve the same convergence rate under these initialization schemes, we can still show the separation on α that affects the convergence in Eq. (4) among these initialization schemes. To be specific, we verify that under LeCun and He initialization, the lower bound of α 2 is 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 α is the same in these three initialization schemes, we usually choose a large step size under the NTK initialization scheme, see Appendix E.1.

Comparison on architectures: Note that the Transformer defined in Eq. (1.2) includes a selfattention layer, a feedforward ReLU layer, and an output layer. Our result proves that a cubic (quadratic) over-parameterization condition is required for d -1 / 2 m ( d -1 m ) under LeCun initialization. As a comparison, a three-layer fully-connected ReLU network under LeCun initialization requires d m = ˜ Ω( N 3 ) [Nguyen, 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 s = 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 τ 0 = d -1 m scaling, given vector input x ∈ R ˜ d , if one feeds the input to Transformer by setting d s = 1 , d = ˜ d , then training with GD can converge to a global minimum. On the contrary, if one sets d s = ˜ 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 ∈ R N × 1 , then the NTK XX ⊤ is a rank-one matrix. As a 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 3.

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, 2021] cannot be applied to our setting because of 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 XX ⊤ to disentangle the nonlinear softmax function. By doing so, a lower bound on the minimum singular value of F pre in Proposition. 1 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 -1 m setting (the softmax degenerates to an all-one vector) but is inapplicable to the realistic d -1 / 2 m 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., 2021] on the standard MNIST dataset in Section 5.2. Additional results and detail on the experimental setup are deferred to Appendix E.

5.1 Fitting synthetic data

In this section, we corroborate our theoretical findings on the synthetic data. We generate 100 data points where the input X ∈ R 10 × 100 is sampled from standard Gaussian distribution. The corresponding label y ∈ 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. (1.2) with LeCun initialization and τ 0 = d -1 / 2 m for 400 epochs with a fixed step size γ = 1 . We test different widths of the network including d m = { 10 , 100 , 1000 , 4000 } and plot the training progress in

0

50

100

150

200

250

300

350

400

Epochs

10

8

10

6

10

4

10

2

10 0

Training loss

dm

=10

dm

=100

dm

=1000

dm

=4000

(a) Convergence curve.

0

50

100

150

200

250

300

350

400

Epochs

10

2

10

1

10 0

Change of weights:

||

t

0

||2

||

0

||2

dm

=10

dm

=100

dm

=1000

dm

=4000

(b) Weight movement.

0

50

100

150

200

Epochs

0

25

50

75

100

125

150

175

200

Epochs

dm

=10

0

50

100

150

200

Epochs

0

25

50

75

100

125

150

175

200

Epochs

dm

=100

0

50

100

150

200

Epochs

0

25

50

75

100

125

150

175

200

Epochs

dm

=1000

0

50

100

150

200

Epochs

0

25

50

75

100

125

150

175

200

Epochs

dm

=4000

0.000

0.005

0.010

0.015

0.020

0.025

0.030

(c) Kernel distance.

Figure 2. The result shows that the training can converge except for the case with sub-linear width, see the linear convergence rate in Figure 2a and the small movement of the weight in Figure 2b. 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:

which quantitatively compares two kernels by the relative rotation, as used in Fort et al. [2020]. Figure 2c shows that the kernel changes rapidly at the beginning of training while being approximately constant at later stages. The experiment with τ 0 = d -1 m , which is deferred to Appendix E.1, obtains similar results.

Additionally, the experiment with other initialization schemes is deferred to Appendix E.1, where we observe that NTK initialization indeed yields slower convergence than LeCun/He initialization, which is consistent with our theoretical finding.

Next, in order to verify corollary 1, we feed vector input x n ∈ R 100 into Transformer by setting either d s = 1 , d = 100 or d s = 100 , d = 1 under LeCun initialization with τ 0 = d -1 m . Figure 3 shows that the training of Transformer with d s = 1 , d = 100 is similar to that of two-layer FCNN with the same width d m = 100 . However, the case of d s = 100 , d = 1 fails to converge, which is consistent with our theoretical finding.

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 a single layer and single head ViT. The dimension of d is 64. We change the dimension of 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 3 runs in Figure 4a, which shows that when the width is smaller, e.g., 16 , both τ 0 = d -1 m and τ 0 = d -1 / 2 m scaling have similar

0

50

100

150

200

250

300

350

400

Epochs

10

14

10

12

10

10

10

8

10

6

10

4

10

2

10 0

Training loss

Transformer,

d

=1,

ds

=100

Transformer,

d

=100,

ds

=1

FCNN,

d

=100

0

10

20

30

40

50

Epochs

10

2

Training loss

dm

=16,

0=

d

1

m

dm

=16,

0=

d

1/2

m

dm

=1024,

0=

d

1

m

dm

=1024,

0=

d

1/2

m

dm

=16384,

0=

d

1

m

dm

=16384,

0=

d

1/2

m

(a) Convergence curve.

Before training

After training

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.0

0.2

0.4

0.6

0.8

(b) Attention map,

d

m

= 16384

.

convergence results. However, as the width increases to 1024 and 16384 , the τ 0 = d -1 m setting admits a slower rate than that of τ 0 = d -1 / 2 m , especially a extremely slow rate under d m = 16384 . This is consistent with our theoretical result on the dimension missing effect such that the τ 0 = d -1 m setting makes Transformer difficult to fit data. Additionally, we visualize the attention map in Figure 4b, i.e., the output of softmax in the self-attention layer before training and after training under τ 0 = d -1 / 2 m setting. We could see that the self-attention layer changes a lot during training.

6 Conclusion

We present a comprehensive analysis of the global convergence properties of shallow Transformer under various scaling and initialization schemes. Regarding scaling schemes, for a large width setting, the difference on convergence between τ 0 = d -1 / 2 m and τ 0 = d -1 m can be demonstrated both theoretically and empirically. Our theory is able to explain this in a dimension missing view. Regarding initialization schemes, our theory prefers to 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 is indeed an intriguing avenue for future research.

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 on Learning Representations (ICLR) , 2021.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In International Conference on Machine Learning (ICML) , pages 242-252. PMLR, 2019a.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of 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 on 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 over-parameterization. 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-parameterization is sufficient to learn deep re{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 of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-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 at 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.

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 a matrix. Proc. Russian Acad. Sci , page 749-754, 1931.

Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J Tibshirani. Surprises in highdimensional 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?-a neural tangent kernel perspective. Advances in neural information processing systems (NeurIPS) , 33:2698-2709, 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. CesaBianchi, 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, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in neural information processing systems (NeurIPS) , 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 SohlDickstein, 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-Yu 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.

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. Association for Computational Linguistics.

Leon Mirsky. A 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: Memorization 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 overparameterization: Global convergence 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'noi Matematiki i Matematicheskoi Fiziki , 3(4):643-653, 1963.

Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding with unsupervised learning. 2018.

Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601 , 2021.

Chaehwan Song, Ali Ramezani-Kebrya, 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 preprint arXiv:2009.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 N Gomez, Łukasz 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.1017/CBO9780511794308.006.

Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science . Cambridge University Press, 2018.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge University Press, 2019.

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.14548 , 2020.

Yongyi Yang, Zengfeng Huang, and David Wipf. Transformers from an optimization perspective. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in neural information processing systems (NeurIPS) , 2022.

Atıf Emre Yüksel, Ya¸ sar Alim Türkmen, Arzucan Özgür, and Berna Altınel. Turkish tweet classification 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. Understanding 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 (NeurIPS) , 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 overparameterized deep relu networks. Machine learning , 2020.

Contents of the Appendix

The Appendix is organized as follows:

In Appendix A, we summarize the symbols and notation used in this work.

In Appendix B.1, we provide a theoretical overview for neural tangent kernel (NTK). The background 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.

The proofs for the convergence of 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 E.

Limitations and societal impact of this work are discussed in Appendix F and Appendix G, respectively.

A Symbols and Notation

We include the core symbols and notation in Table 2 for facilitating the understanding of our work.

B Theoretical background

B.1 Preliminary on NTK

In this section, we summarize how training a neural network by minimizing squared loss, i.e., ℓ ( θ t ) = 1 2 ∑ N n =1 ( f ( X n ; θ t ) -y n ) 2 , 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:

By substituting the loss into the above equation and using the chain rule, we can find that the network outputs f ( θ t ) ∈ R N admit the following dynamics:

where K t = ( ∂f ( θ t ) ∂ θ )( ∂f ( θ t ) ∂ θ ) ⊤ ∈ R N × N . Jacot et al. [2018], Arora et al. [2019b] have shown that for fully-connected neural networks, under the infinite-width setting and proper initialization, K t will be stable during training and K 0 will converge to a fixed matrix K ⋆ ∈ R N × N , where ( K ⋆ ) ( ij ) = K ⋆ ( X i , X j ) is the NTK value for the inputs X i and X j . Then, we rewrite Eq. (5) as:

This implies the network output for any X ∈ R d s × d can be calculated by the kernel regression predictor with the associated NTK:

where K ⋆ ( X , X n ) is the kernel value between test data X and training data X n .

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]. A random variable X with mean µ is called sub-exponential random variable if there exist non-negative parameters ( ν, α ) such that

and we denote by X ∼ SE ( ν, α ) .

Lemma 2. The multiplication of a scalar s ∈ R + and a sub-exponential random variable X ∼ SE ( ν, α ) is still a sub-exponential random variable: sX ∼ SE ( sν, sα ) .

Lemma 3. Given a set of independent sub-exponential random variables X i ∼ ( ν i , α i ) for i ∈ 1 , ..., N , then ∑ N i =1 X i ∼ ( √ ∑ N i =1 ν 2 i , max i α i ) .

Lemma 4. Given a sub-exponential random variable X ∼ SE ( ν, α ) with mean µ , the following inequality holds with probability at least 1 -δ :

B.3 Related work on over-parameterization 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, a 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 over-parameterized ReLU neural network with cross-entropy loss [Li and Liang, 2018]. Du et al. [2019b] show that only training the first layer of two-layer ReLU 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 Ω( N 4 ) , 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 Ω( N 8 ) , which improves upon Allen-Zhu et al. [2019a] that show the result of Ω( N 24 ) . Allen-Zhu 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. [2021]. This facilitates Nguyen [2021] to demonstrate that deep ReLU network under LeCUN initialization with width in the order Ω( N 3 ) 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.

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.1, we provide some auxiliary lemmas. Lemma 5 and Lemma 6 show that the norm of the parameters can be bounded with high probability at initialization. Lemmas 7, 9 and 10 present that the network output and the output of softmax between two adjacent time steps can be upper bounded. The Lipschitzness of network gradient and its norm is bounded in Lemmas 13 and 14. In Appendix C.2, we prove the convergence of the general cases, i.e., Proposition. 1. In Appendix C.3 and C.4 we present the proof for d -1 / 2 m and d -1 m scaling.

C.1 Auxiliary lemmas

Lemma 5 (Corollary 5.35 of Vershynin [2012]) . For a weight matrix W ∈ R d 1 × d 2 where each element is sampled independently from N (0 , 1) , for every ζ ≥ 0 , with probability at least 1 -2exp( -ζ 2 / 2) one has:

where σ max ( W ) and σ min ( W ) represents the maximum and minimum singular value of W , respectively.

Lemma 6 (Upper bound of spectral norms of initial weight) . For a weight matrix W ∈ R d m × d where d m > d , each element is sampled independently from N (0 , 1) , with probability at least 1 -2 exp( -d m / 2) , one has:

Proof of Lemma 6. Following Lemma 5, one has:

Letting ζ = d and using the fact that d > d complete the proof.

√ m m

Lemma 7. Recall from Table 2, β i = σ s ( τ 0 X ( i, :) W ⊤ Q W K X ⊤ ) ⊤ is the i -th row of the output of Softmax, if max( ∥ ∥ W t Q ∥ ∥ 2 , ∥ ∥ ∥ W t ′ Q ∥ ∥ ∥ 2 ) ≤ ¯ λ Q , max( ∥ W t K ∥ 2 , ∥ ∥ ∥ W t ′ K ∥ ∥ ∥ 2 ) ≤ ¯ λ K , then its difference

between t ′ step and t step has the following upper bound:

Proof.

Lemma 8 (Upper bound the Euclidean norm of the output of softmax) . Suppose v = Softmax ( u ) ∈ R d s , then one has: 1 / √ d s ≤ ∥ v ∥ 2 ≤ 1 .

Proof. By the inequality of arithmetic and geometric mean:

Lemma 9. If at t ′ and t step, max( ∥ ∥ W t Q ∥ ∥ 2 , ∥ ∥ ∥ W t ′ Q ∥ ∥ ∥ 2 ) ≤ ¯ λ Q , max( ∥ W t K ∥ 2 , ∥ ∥ ∥ W t ′ K ∥ ∥ ∥ 2 ) ≤ ¯ λ K , max( ∥ W t V ∥ 2 , ∥ ∥ ∥ W t ′ V ∥ ∥ ∥ 2 ) ≤ ¯ λ V , then the difference between the output of the last hidden layer at t ′ step and t can be upper bounded by:

Proof. Note that F pre ∈ R N × d m is the output of the last hidden layer given { X n } N n =1 sample as defined in Table 2. We firstly analyze each sample, i.e., f pre ∈ R d m and we drop the index n for simplification.

Next, we bound the difference given N data sample:

Lemma 10. If at t ′ and t step, max( ∥ ∥ W t Q ∥ ∥ 2 , ∥ ∥ ∥ W t ′ Q ∥ ∥ ∥ 2 ) ≤ ¯ λ Q , max( ∥ W t K ∥ 2 , ∥ ∥ ∥ W t ′ K ∥ ∥ ∥ 2 ) ≤ ¯ λ K , max( ∥ W t V ∥ 2 , ∥ ∥ ∥ W t ′ V ∥ ∥ ∥ 2 ) ≤ ¯ λ V , max( ∥ w t O ∥ 2 , ∥ ∥ ∥ w t ′ O ∥ ∥ ∥ 2 ) ≤ ¯ λ O , then the difference between the network output at t ′ step and t step can be upper bounded by:

Proof.

where we use triangle inequality. Then, the first term of the RHS can be bounded by:

where the last equality is by the definition of ρ in Proposition. 1.

Lemma 11 (Jacobian of Softmax) . Suppose v = Softmax ( u ) ∈ R d s , then ∂ v = diag ( v ) -vv ⊤ .

Proof. We can reformulate v as:

v

=

∂ u

Then, we have

Thus

Lemma 12 (Upper bound for the loss gradient norm) . If ∥ ∥ W t Q ∥ ∥ 2 ≤ ¯ λ Q , ∥ W t K ∥ 2 ≤ ¯ λ K , ∥ W t V ∥ 2 ≤ ¯ λ V , ∥ w t O ∥ 2 ≤ ¯ λ O , then the gradient norm with respect to W Q , W K , W V , w O can be upper bounded by:

Proof. To simplify the notation, in the proof below, we hide the index t . Firstly, consider the gradient w.r.t W V ,

where the last equality is by the definition of ρ in Proposition. 1. Next, consider the gradient w.r.t w O ,

̸

̸

Next, consider the gradient w.r.t W Q ,

Similarly, for the gradient w.r.t W K , we have:

Lemma 13 (Upper bound for the network function gradient norm) . If ∥ ∥ W t Q ∥ ∥ 2 ≤ ¯ λ Q , ∥ W t K ∥ 2 ≤ ¯ λ K , ∥ W t V ∥ 2 ≤ ¯ λ V , ∥ w t O ∥ 2 ≤ ¯ λ O , then one has:

where

Proof. To simplify the notation, in the proof below, we hide the index t . Firstly, note that:

Then following the step as in Eqs. (7) to (10), each term can be bounded as follows:

Plugging these bounds back Eq. (12) finishes the proof.

Lemma 14 (Upper bound for the Lipschitzness of the network gradient) . Suppose max( ∥ ∥ W t Q ∥ ∥ 2 , ∥ ∥ ∥ W t ′ Q ∥ ∥ ∥ 2 ) ≤ ¯ λ Q , max( ∥ W t K ∥ 2 , ∥ ∥ ∥ W t ′ K ∥ ∥ ∥ 2 ) ≤ ¯ λ K , max( ∥ W t V ∥ 2 , ∥ ∥ ∥ W t ′ V ∥ ∥ ∥ 2 ) ≤ ¯ λ V , max( ∥ w t O ∥ 2 , ∥ ∥ ∥ w t ′ O ∥ ∥ ∥ 2 ) ≤ ¯ λ O , and define z ≜ 2 τ 0 C 2 x d s ( ¯ λ Q + ¯ λ K ) , then one has

where

Proof. Firstly, note that:

Then, we will prove each term separately. Regarding the first term in Eq. (15), we have

where the first inequality is by Eq. (6), and in the last equality is by the definition of z in the lemma. Regarding the second term in Eq. (15), we have:

The term ∥ ∥ ∥ β t ′ i,n -β t i,n ∥ ∥ ∥ 2 can be bounded by Lemma 7 as follows:

The term ∥ ∥ ∥ w t ′ O ◦ ˙ σ r ( W t ′ V X ⊤ n β t ′ i,n ) -w t O ◦ ˙ σ r ( W t V X ⊤ n β t i,n ) ∥ ∥ ∥ 2 can be bounded by the triangle inequality, Cauchy-Schwarz inequality, and the same method in Eq. (6) as follows:

Plugging Eq. (18) and Eq. (19) back Eq. (17), we obtain:

Regarding the third term in Eq. (15), let us denote by:

Then:

We bound each term separately:

where the second inequality uses the result in Eq. (19).

where the last second inequality is by Lemma 7. Plugging back Eq. (21), we obtain:

Similarly, the fourth term in Eq. (15) can be bounded by:

Plugging the upper bound for these four terms back Eq. (15) finishes the proof.

C.2 Proof of Proposition. 1

Proof. We can reformulate Eq. (1.2) as:

It is clear that Eq. (25) holds for t = 0 . Assume that Eq. (25) holds up to iteration t . By the triangle inequality, we have

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 := √ 1 -γα 2 / 2 , we bound the RHS of the previous expression:

By Weyl's inequality, this implies:

Similarly,

By Weyl's inequality, this implies

Similarly, we can obtain that

Next, we will prove the fifth inequality in Eq. (25):

where the first inequality holds by Lemma 9. This result further implies that σ min ( F t +1 pre ) ≥ 1 2 α 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 [ θ ⊤ , θ t +1 ] . We define θ t + ϕ := θ ⊤ + ϕ ( θ t +1 -θ ⊤ ) , for ϕ ∈ [0 , 1] . Then, by triangle inequality and Cauchy-Schwarz inequality, we have:

where the last inequality is by induction rule. Then, we bound each term separately. Note that:

Then following exact the same step as in Eqs. (26), (27) and (30), we have: ∥ ∥ ∥ W t + ϕ Q ∥ ∥ ∥ 2 ≤ ¯ λ Q . By the same method, we have: ∥ ∥ ∥ W t + ϕ K ∥ ∥ ∥ 2 ≤ ¯ λ K , ∥ ∥ ∥ W t + ϕ V ∥ ∥ ∥ 2 ≤ ¯ λ V , ∥ ∥ ∥ w t + ϕ O ∥ ∥ ∥ 2 ≤ ¯ λ O . Now we proceed to bound the first term in Eq. (35).

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 13. The third term in Eq. (35) can be bounded by Lemma 14. As a result, Eq. (35) has the following upper bound:

where c 1 , c 2 , c 3 are defined at Eqs. (11), (14) and (36), and we further define C ≜ c 1 c 2 +2 c 3 ℓ ( θ 0 ) . Lastly, by applying Lemma 4.3 in Nguyen and Mondelli [2020] and Eq. (37), we have:

which concludes the proof.

C.3 Proof of Theorem 1

Lemma 15. Let Φ = [ X ⊤ 1 β 1 , 1 , ..., X ⊤ N β 1 ,N ] ⊤ ∈ R N × d , where β 1 ,n = σ s ( τ 0 X (1 , :) n W ⊤ Q W K X ⊤ n ) ⊤ , n ∈ [ N ] , then under Assumptions 2 and 3, with probability at

Proof. Due to Assumption 2, for any data X , the matrix XX ⊤ is positive definite and thus has positive minimum eigenvalue. We denote it as λ min ( XX ⊤ ) ≥ C λ .

According to [Nguyen et al., 2021, Lemma 5.3], using the Hermite expansion of σ r , one has:

where µ ( σ r ) is the 1-st Hermite coefficient of ReLU satisfying µ ( σ r ) > 0 .

Now we proceed to provide a lower bound for λ min ( ΦΦ ⊤ ) . For notational simplicity, define:

Then we can rewrite ΦΦ ⊤ as follows:

By Gershgorin circle theorem [Gershgorin, 1931], there exists k ∈ [ N ] such that:

̸

Using Von Neumann's trace inequality [Mirsky, 1975] and noting that B kj is a rank one matrix, one has:

where we use the definition of the inner product between two matrices and ∥ β 1 ,j ∥ 2 ≤ 1 . By Assumption 2, we have λ min ( X k X ⊤ k ) ≥ C λ , where C λ is some positive constant. By setting t := ∥ β 1 ,k ∥ 2 2 C 2 λ / ( N -1) 2 in Assumption 3, with probability at least 1 -exp( -Ω(( N -1) -ˆ c )) , one has

̸

Plugging back Eq. (39) and Eq. (40), we obtain:

where the last inequality is by the lower bound of β 1 ,k in Lemma 8. Lastly, plugging the lower bound of λ min ( ΦΦ ⊤ ) back Eq. (38) finishes the proof.

Remark: The estimation of λ 0 is actually tight because its upper bound is also in a constant order. To be specific, denote G := σ r ( Φ w ) σ r ( Φ w ) ⊤ , we have

where Φ ( n, :) = β ⊤ 1 ,n X n . Next, by Liao and Couillet [2018] (Sec.A in Supplementary Material):

Combine Eq. (42) and Eq. (43), we have

That means, our estimation on λ 0 is tight as its upper and lower bounds match with each other.

Now we are ready to present the proof for LeCun initialization under τ 0 = d -1 / 2 m scaling.

Proof. We select C Q = C K = 1 = C V = C O = 1 , then by Lemma 5, with probability at least 1 -8 e -d m / 2 , we have:

Plugging Eq. (44) into Eqs. (2) and (3) , it suffices to prove the following equations.

Next, we will provide the lower bound for α 2 = λ min (( F 0 pre )( F 0 pre ) ⊤ ) . In the following context, we hide the index 0 for simplification. One can note that F pre F ⊤ pre is the summation of PSD matrices, thus, it suffices to lower bound: λ min ( ˆ F pre ˆ F ⊤ pre ) , where we introduce the following notations:

By Matrix-Chernoff inequality, we can obtain that (e.g. Lemma 5.2 of Nguyen et al. [2021]) w.p at least 1 -δ 1 ,

as long as it holds d m ≥ ˜ Ω( N/λ 0 ) , where λ 0 = λ min ( E w ∼N (0 ,η V I d ) [ σ r ( Φ w ) σ r ( Φ w ) ⊤ ] ) , and ˜ Ω hides logarithmic factors depending on δ 1 . Lastly, w.p. at least 1 -δ 2 , one has √ 2 ℓ ( θ 0 ) ≤ ˜ O ( √ N ) . Plugging back Eqs. (2) and (3), it suffices to prove the following inequality.

By Lemma 15, with probability at least 1 -exp( -Ω(( N -1) -ˆ c d -1 s )) , one has λ 0 ≥ Θ( η V /d s ) = Θ( d -1 d -1 s ) . Thus, when d m ≥ ˜ Ω( N 3 ) , 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.

For the proof under LeCun initialization and τ 0 = d -1 m scaling, we follow the same strategy. Specifically: As the same in the proof with τ 0 = d -1 / 2 m scaling, we select C Q = C K = C V = C O = 1 , then plugging Eq. (44) into Eqs. (2) and (3), it suffices to prove the following equations.

Next, we will provide the lower bound for α 2 = λ min (( F 0 pre )( F 0 pre ) ⊤ ) . In the following context, we hide the index 0 for simplification. One can note that F pre F ⊤ pre is the summation of PSD matrices,

thus, it suffices to lower bound: λ min ( ˆ F pre ˆ F ⊤ pre ) , where we introduce the following notations:

By Eq. (48) w.p at least 1 -δ 1 , λ min ( ˆ F pre ˆ F ⊤ pre ) ≥ d m λ 0 / 4 , , as long as it holds d m ≥ ˜ Ω( N/λ 0 ) , where λ 0 = λ min ( E w ∼N (0 ,η V I d ) [ σ r ( Φ w ) σ r ( Φ w ) ⊤ ] ) , and ˜ Ω hides logarithmic factors depending on δ 1 . Lastly, note that the activation function in the output layer σ r is 1-Lipschitz and is applied to σ r ( W V X ⊤ β i ) , where X ⊤ β i is bounded due to the softmax's property in Lemma 8, then by Lemma C.1 of Nguyen and Mondelli [2020], w.p. at least 1 -δ 3 , one has √ 2 ℓ ( θ 0 ) ≤ ˜ O ( √ N ) . Plugging back Eqs. (2) and (3), it suffices to prove the following inequality.

By Lemma 15, with probability at least 1 -exp( -Ω(( N -1) -ˆ c d -1 s )) , one has λ 0 ≥ Θ( η V /d s ) = Θ( d -1 d -1 s ) . Thus, when d m ≥ ˜ Ω( N 2 ) , 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 C Q = C K = C V = C O = 1 , then by Lemma 5, with probability at least 1 -8 e -d m / 2 , we have:

When d m ≥ d , plugging Eq. (56) into Eqs. (2) and (3), it suffices to prove the following equations.

Next, we will provide the lower bound for α 2 = λ min (( F 0 pre )( F 0 pre ) ⊤ ) . In the following context, we hide the index 0 for simplification. One can note that F pre F ⊤ pre is the summation of PSD matrices, thus, it suffices to lower bound: λ min ( ˆ F pre ˆ F ⊤ pre ) , where we introduce the following notations:

By Eq. (48), w.p at least 1 -δ , λ min ( ˆ F pre ˆ F ⊤ pre ) ≥ d m λ 0 / 4 , as long as it holds d m ≥ ˜ Ω( N/λ 0 ) , where λ 0 = λ min ( E w ∼N (0 , I d ) [ σ r ( Φ w ) σ r ( Φ w ) ⊤ ] ) , and ˜ Ω hides logarithmic factors depending on δ 1 . Lastly, w.p. at least 1 -δ 3 , one has √ 2 ℓ ( θ 0 ) ≤ ˜ O ( √ N ) . Plugging back Eqs. (2) and (3), it suffices to prove the following inequality.

By Lemma 15, with probability at least 1 -exp( -Ω(( N -1) -ˆ c d -1 s )) , we have λ 0 ≥ Ω(1) . Thus, all of the above conditions hold when d m = ˜ Ω( N ) . As a result, the conditions in Eqs. (57) and (58) are satisfied and the convergence of training Transformer is guaranteed.

C.5 Discussion for different initialization schemes

Recall that the convergence result in Theorem 2 shows:

Thus, to discuss the convergence speed for different initialization, we need to check the lower bound for α 2 . From the proofs for different initialization schemes above, we have the following lower bound for α 2 , i.e.,

with high probability. Plugging the value of τ 1 and η V , we observe that for LeCun initialization and He initialization: α 2 ≥ Ω( d m N/d ) while for NTK initialization: α 2 ≥ Ω( N/d ) . Thus, the convergence speed of LeCun initialization and He initialization is faster than NTK initialization. As a result, faster step-size is required for NTK intialization.

Eq. (4) indicates that the convergence speed is affected by α , there we compare the lower bound for α for these two scaling. In Appendix C.3, we have proved that under the LeCun initialization, one has α 2 ≥ d m λ 0 / 4 ≥ d m η V µ ( σ r ) 2 Θ( ∥ β 1 ,k ∥ 2 2 ) . Note that this bound holds for these two different scaling, which is inside β . Thus in the next, we need to see the difference between the lower bound of ∥ β 1 ,k ∥ 2 2 ) in the case of these two scalings. Specifically, for τ 0 = d -1 / 2 m scaling, we have proved that ∥ β 1 ,k ∥ 2 2 ≥ 1 /d s by Lemma 8. However, for the case of τ 0 = d -1 m scaling, when the width is large enough, the value inside the softmax tends to zero, as a result, ∥ β 1 ,k ∥ 2 2 ≈ 1 /d s . Thus, we can see that as the width increases, the convergence speed of τ 0 = d -1 / 2 m could be faster. Lastly, we remark on the difference in the step size for these two scales, which can be seen from the definition of C and its corresponding c 1 , c 2 , c 3 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. 1 to provide the sufficient condition for the convergence guarantee, e.g., Eqs. (2) and (3). The second step is similar to Theorem 2, 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 α := σ min ( F 0 pre ) , where F pre is the output of the last hidden layer. One can apply concentration inequality to bound the difference between σ min ( F 0 pre ) and σ min ( F ⋆ 0 pre ) , 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 is general and can handle the following residual Transformer. Here we give a proof sketch to show how to achieve this. Specifically, we consider the residual block in the self-attention layer:

As a result, the output becomes

To prove the convergence of the above residual Transformer, the first part that will be modified is the proof for Proposition 1. The formula for f pre becomes as follows:

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 X ( i :) ⊤ in two adjacent time steps cancels out. In Lemma 10, we have

Similarly, in the remaining lemmas, we need to add the additional term C x for the upper bound of || X ( i :) || 2 .

The second part is the proof for Theorem 1 regarding the lower bound for α 0 . By Weyl's inequality:

where we denote by F pre = [ f pre ( X 1 ) , · · · , f pre ( X N )] , and F ∗ pre is the corresponding one without the residual connection. Then we can upper bound the second term can be bounded as follows:

The remains step follows the same as previous analysis.

C.8 Linear over-parametrization and attention module behaving as a pooling layer

In this section, we discuss the link between the linear over-parametrization and attention module behaving as a pooling layer under the d -1 m scaling. First, due to the d -1 m scaling, the attention module degenerates to a pooling layer according to the law of large numbers. In this case, the nonlinearity on X disappears and thus the minimum eigenvalue of ΘΘ ⊤ can be estimated via XX ⊤ . Accordingly, this leads to the minimum eigenvalue in the order of Ω( N/d ) , and thus linear over-parameterization is enough.

D Proof for NTK

In this section, we elaborate the proof for Lemma 1 in Appendix D.1, the proof for Theorem 3 in Appendix D.3, respectively.

D.1 Proof of Lemma 1

Proof. We will compute the inner product of the Jacobian of each weight separately. Firstly, we analyze w O . Let us denote by β i := σ s ( τ 0 X ( i, :) W ⊤ Q W K X ⊤ ) ⊤ ∈ R d s . Then:

The inner product of the gradient is:

where the second equality uses the law of large numbers. Secondly, we analyze W Q :

Thus:

The inner product of the gradient is:

τ 0 τ 1 W k W k X diag ( β j ) -β j β j X W V w O ◦ ˙ σ r W V X β j ⟩ . For the first term of the inner product, we have: lim d m →∞ τ 0 τ 1 X ⊤ ( diag ( β i ) -β i β ⊤ i ) XW ⊤ V ( w O ◦ ˙ σ r ( W V X ⊤ β i ))    

where the second equality is by the law of large numbers and E w = 0 for a random variable w ∼ N (0 , 1) . Thus: lim d m →∞ 〈 ∂f ( X ) ∂ W Q , ∂f ( X ′ ) ∂ W Q 〉 = 0 . Similarly, lim d m →∞ 〈 ∂f ( X ) ∂ W K , ∂f ( X ′ ) ∂ W K 〉 = 0 . Lastly, we analyze W V :

Thus:

The inner product of the gradient is:

where the last equality uses the law of large numbers.

D.2 NTK minimum eigenvalue

Lemma 16. Given Φ defined in Eq. (53) and Φ ⋆ defined in Lemma 1, denote λ ∗ = λ min ( Φ ⋆ Φ ⋆ ⊤ ) , and suppose the width satisfies d m = Ω( N 2 √ log(2 d 2 N 2 /δ ) λ 2 ∗ ) , then with probability at least 1 -δ , one has ∥ ∥ ΦΦ ⊤ -Φ ⋆ Φ ⋆ ⊤ ∥ ∥ F ≤ λ ∗ 4 and λ min ( ΦΦ ⊤ ) ≥ 3 λ ∗ 4 .

Proof. In the following content, the variable with ⋆ indicates the corresponding variable with infinite width d m . According to the definition of Φ in Eq. (53) and Φ ⋆ in Lemma 1, for each entry in ΦΦ ⊤ and Φ ⋆ Φ ⋆ ⊤ , we have

Next, we will bound ∥ ∥ β 1 ,n -β ⋆ 1 ,n ∥ ∥ 2 ,

Let first consider the absolute value of each element of τ 0 W ⊤ Q W K -τ 0 W ⋆ ⊤ Q W ⋆ K , i.e.,

Since W ( q,i ) Q W ( q,j ) k ∼ SE ( √ 2 η QK , √ 2 η QK ) is sub-exponential random variable, by Lemma 2 and Lemma 3, τ 0 ∑ d m q =1 W ( q,i ) Q W ( q,j ) k ∼ SE ( τ 0 η QK √ 2 d m , √ 2 τ 0 η QK ) , by Bernstein's inequality, when d m ≥ 2 log(2 /δ ) , the following inequality holds with probability at least 1 -δ :

Substituting δ = δ ′ d 2 and applying union bound, we can obtain that when d m ≥ 2 log 2 d 2 δ ′ , with probability at least 1 -δ ′ :

Substituting back to Eq. (65), the following inequality holds with the same width requirement and probability. | ( ΦΦ ⊤ -Φ ⋆ Φ ⋆ ⊤ ) ( n,r ) | ≤ 8 τ 0 η QK C 4 x d s d √ d m log(2 d 2 /δ ′ ) Applying the union bound over the index ( n, r ) for n ∈ [ N ] and r ∈ [ N ] , and substituting δ ′ = δ ′′ N 2 , we obtain that when the width d m ≥ 2 log 2 d 2 N 2 δ ′′ , the following inequality holds with probability at least 1 -δ ′′ :

In the case of LeCun initialization when d m = Ω( N 2 √ log(2 d 2 N 2 /δ ′′ ) λ 2 0 ) , ∥ ∥ ΦΦ ⊤ -Φ ⋆ Φ ⋆ ⊤ ∥ ∥ F ≤ λ ∗ 4 . Lastly, one has:

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. Given the Φ ⋆ defined in Lemma 1, then when N ≥ Ω( d 4 ) , with probability at least 1 -e -d one has: λ min ( Φ ⋆ Φ ⋆ ⊤ ) ≥ Θ( N/d ) .

Proof. Firstly, let us define Z ⋆ = Φ ⋆ Σ -1 / 2 , then:

where the last inequality is by Assumption 4). Note that we can reformulate Z ⋆ by plugging Φ ⋆ as follows:

Recall that in Assumption 4, we define the random vector x = X ( i, :) ⊤ n and write the covariance matrix for each feature of x as Σ = E [ xx ⊤ ] . Here we further define the random vector z = Σ -1 / 2 x . Clearly, the covariance matrix for each feature of z is Σ z = E [ zz ⊤ ] = I d . Therefore, by Proposition 4 of Zhu et al. [2022]), we have λ min ( Z ⋆ ⊤ Z ⋆ ) ≥ Θ( N/d ) , which finishes the proof.

Lemma 18. Suppose the number of samples N ≥ Ω( d 4 ) and the width d m = ˜ Ω( N 2 λ -2 ∗ ) , where λ ∗ = λ min ( Φ ⋆ Φ ⋆ ⊤ ) , then under Assumption 4, with probability at least 1 -δ -e -d , one has λ 0 = λ min ( E w ∼N (0 ,η V I d ) [ σ r ( Φ w ) σ r ( Φ w ) T ] ) ≥ Θ( η V N/d ) .

Proof. By the Hermite expansion of σ r , one has:

where µ ( σ r ) is the 1-st Hermite coefficient of ReLU satisfying µ ( σ r ) > 0 . By Lemma 16, , when the width satisfies d m = ˜ Ω( N 2 λ -2 ∗ ) , then with probability at least 1 -δ , one has λ min ( ΦΦ ⊤ ) ≥ 3 4 λ ∗ . Furthermore, by Lemma 17, when N ≥ Ω( d 4 ) w.p. at least 1 -e -d one has λ ∗ = Θ( N/d ) . Thus, with probability at least 1 -δ -e -d , one has:

Below, we provide a lower bound for the minimum eigenvalue of NTK, which plays a key role in analyzing convergence, generalization bounds, and memorization capacity [Arora et al., 2019a, Nguyen et al., 2021, Montanari and Zhong, 2022, Bombari et al., 2022].

To prove this, we need the following assumption.

Assumption 4. Let x = ∑ d s i =1 ( X ( i, :) ) ⊤ , Σ = E [ xx ⊤ ] , then we assume that Σ is positive definite, i.e., λ min ( Σ ) ≥ C Σ for some positive constant C Σ .

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., 2022], the covariance of x is frequently assumed to be an identity matrix or positive definite matrix.

Lemma 19 (Minimum eigenvalue of limiting NTK) . Under Assumptions 1 and 4 and scaling τ 0 = d -1 m , when N ≥ Ω( d 4 ) , the minimum eigenvalue of K can be lower bounded with probability at least 1 -e -d as: λ min ( K ) ≥ µ ( σ r ) 2 Ω( N/d ) , where µ ( σ r ) represents the first Hermite coefficient of the ReLU activation function.

Proof. By the Hermite expansion of σ r , we have:

Note that when N ≥ Ω( d 4 ) , based on Lemma 17 and the fact that Φ ⋆ Φ ⋆ ⊤ and Φ ⋆ ⊤ Φ ⋆ share the same non-zero eigenvalues, with probability at least 1 -e -d one has:

which completes the proof.

Remark: The minimum eigenvalue of the NTK plays an important role in the global convergence, similar to α in Proposition. 1. By defining α ≜ σ min ( F 0 pre ) , under the over-parameterized regime with d m ≥ N , we have α 2 = λ min (( F 0 pre ) ⊤ F 0 pre ) , which is the exact minimum eigenvalue of the empirical NTK with respect to the weight in the output layer.

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 w ∈ R d m where each element is sampled independently from N (0 , 1) , for any ˜ w such that ∥ ˜ w -w ∥ 2 ≤ R , with probability at least 1 -2 exp( -d m / 2) , one has:

Proof of Lemma 20. By triangle inequality and Lemma 6:

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 ( d m ), 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:

The Hessian matrix H of the network can be written as the following structure:

where we partition the Hessian H of the network to each Hessian block e.g., H Q,Q = ∂ 2 f ∂ w Q ∂ w Q , H Q,K = ∂ 2 f ∂ w Q ∂ w K , etc. Based on the triangle inequality 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 H Q,Q . The time step t is hidden for simplification.

where the Jacobian of Softmax is given by Lemma 11. Next, we calculate the second-order derivative. Each element of the Hessian H Q,Q is

where we denote by Γ i = { diag ( β i ◦ x -β i β i ⊤ x ) -( diag ( β i ) -β i β i ⊤ ) xβ i ⊤ -β i x ⊤ ( diag ( β i ) -β i β i ⊤ ) x } , ◦ symbolizes Hadamard product. According to initialization of w O , w V , w Q , w K and Lemma 20, with probability at least 1 -8 e -d m / 2 , we have

The spectral norm of H Q,Q can be bounded with probability at least 1 -8 e -d m / 2 :

where the last equality is by the bound of ∥ Γ i ∥ 2 , by triangle inequality:

For the first part of Eq. (74), we have

For the second part of Eq. (74), we have

For the third part of Eq. (74), we have

Next, we analyze the Hessian H Q,K , where each element is:

Due to the symmetry, similar to H Q,Q , the spectral norm of H Q,K can be bounded with probability at least 1 -8 e -d m / 2 : ∥ H Q,K ∥ 2 = O (1 / √ d m ) .

Next, we analyze the Hessian H Q,V , where each element is:

The spectral norm of H Q,V can be bounded with probability at least 1 -8 e -d m / 2 :

where the last step is by Weyl's inequality and the range of the output of Softmax (from zero to one):

It suffices to bound ∥ β i ∥ 2 , i.e., ∥ σ s ( x ( i ) xw ⊤ Q w K ) ∥ 2 . Since the range of each element of the output of Softmax is from 0 to 1 and the sum of is one, ∥ β i ∥ ≤ 1 .

Next, we analyze the Hessian H Q,O , where each element is:

The spectral norm of H Q,V can be bounded with probability at least 1 -8 e -d m / 2 :

Next, we analyze the Hessian H O,V . Each element of H O,V is:

Note that H O,V is a diagonal matrix. Consequently, the spectral norm of H O,V is:

Next, since each element of H O,O and H V,V is zero, the corresponding spectral norm is zero. Lastly, due to the symmetry of w K and w Q , we can obtain the same bound for the corresponding Hessian block. Thus, the spectral norm of Hessian H is upper bounded by O (1 / √ d m ) , which completes the proof.

0

50

100

150

200

250

300

350

400

Epochs

10

8

10

6

10

4

10

2

10 0

Training loss

dm

=10

dm

=100

dm

=1000

dm

=4000

(a) LeCun initialization,

γ

= 1

.

0

50

100

150

200

250

300

350

400

Epochs

10

10

10

8

10

6

10

4

10

2

10 0

10 2

Training loss:

dm

=10

dm

=100

dm

=1000

dm

=4000

(b) He initialization,

γ

= 1

.

0

50

100

150

200

250

300

350

400

Epochs

10

10

10

8

10

6

10

4

10

2

10 0

10 2

Training loss:

dm

=10

dm

=100

dm

=1000

dm

=4000

(c) NTK initialization,

γ

= 1

.

0

50

100

150

200

250

300

350

400

Epochs

10

10

10

8

10

6

10

4

10

2

10 0

10 2

Training loss:

dm

=10

dm

=100

dm

=1000

dm

=4000

(d) NTK initialization,

γ

= 10

.

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 5a to 5c show that when using the same step size γ = 1 for LeCun and He initialization exhibits similar 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 C.5.

E.2 Additional result and set-up in Section 5.2

We have validated Assumption 3 on language data in the main paper. Here we will validate it on image dataset. Specifically, we choose MNIST dataset and consider the embedding in the ViT as X , then we plot P (∣ ∣ 〈 X ⊤ n X n , X ⊤ n ′ X n ′ 〉∣ ∣ ≥ t ) as t increases in Figure 6, where we could see exponential decay. The architecture of ViT is the same as in Section 5.2.

Additionally, in Figure 7, we can see that d -1 / 2 m setting achieves faster convergence in training loss and higher test accuracy. In Figure 8, we conduct the experiment on MNIST with regression task by using the same architecture. We take two classes { 0 , 1 } in MNIST for binary classification, the output of our regression is taken by sign ( f ) , and the training error is given by ∑ N n =1 [ sign ( f ( X n )) -y n ] 2 . We can see that the convergence speed of d -1 / 2 m is faster than that of d -1 m . Overall, these separation results on both classification and regression tasks provide good justification for the d -1 / 2 m scaling.

0.00

0.25

0.50

0.75

0.0

0.2

0.4

0.6

0.8

1.0

Probability

0

10

20

30

40

50

Epochs

20

40

60

80

100

Test accuracy %

dm

=16,

0=

d

1

m

dm

=16,

0=

d

1/2

m

dm

=1024,

0=

d

1

m

dm

=1024,

0=

d

1/2

m

dm

=16384,

0=

d

1

m

dm

=16384,

0=

d

1/2

m

Pooling

0

500

1000

1500

2000

2500

Iterations

10

4

10

3

10

2

10

1

10 0

Training loss

dm

=16,

0=

d

1

m

dm

=16,

0=

d

1/2

m

dm

=1024,

0=

d

1

m

dm

=1024,

0=

d

1/2

m

0

50

100

150

200

250

300

350

400

Epochs

10

8

10

6

10

4

10

2

10 0

Training loss:  (

W

t

)

dm

=10

dm

=100

dm

=1000

dm

=4000

(a) Convergence curve.

0

50

100

150

200

250

300

350

400

Epochs

10

2

10

1

10 0

Change of weights:

||

W

t

W

0

||

F

||

W

0

||

F

dm

=10

dm

=100

dm

=1000

dm

=4000

(b) Weight movement.

0

100

200

300

400

Epochs

0

50

100

150

200

250

300

350

400

Epochs

dm

=10

0

100

200

300

400

Epochs

0

50

100

150

200

250

300

350

400

Epochs

dm

=100

0

100

200

300

400

Epochs

0

50

100

150

200

250

300

350

400

Epochs

dm

=1000

0

100

200

300

400

Epochs

0

50

100

150

200

250

300

350

400

Epochs

dm

=4000

0.000

0.005

0.010

0.015

0.020

0.025

0.030

(c) Kernel distance.

Figure 9: Experimental validation of the theoretical results on Transformers with τ 0 = d -1 m scaling trained on synthetic data. (a) Linear convergence. (b) Rate of change of the weights during training. Observe that the weights change very slowly after the 50 th 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 NTK becomes more stable.

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

Our result does not cover SGD training and other loss functions, e.g., hinge loss and cross-entropy loss. Secondly, the model architecture being analyzed in this work is the encoder of Transformer which is widely used for regression problems or image classification problems. We do not consider the entire Transformer including the decoder, which is designed for sequence-to-sequence modeling. Lastly, our model does not cover the deep Transformer. We believe our analytic framework paves a way for further analyzing large-scale Transformer.

G Societal impact

This work offers a convergence analysis for the Transformer, which is a core element within contemporary large language models. Our theoretical analysis with realistic initializations and scaling schemes lays a 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nTo 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe technical limitations are stated in weakness part. I do not see any negative social impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nTo 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.\n\nWEAKNESSES:\nExcept 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThis work has no potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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 ?\n\nQUESTIONS:\n1. 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 ?\n\nLIMITATIONS:\nYes (Appendix F).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nYes, the authors include a section (Section F in appendix) to address the limitation of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

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

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.

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.

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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\nWEAKNESSES:\n**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\nQUESTIONS:\n**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.\n\nLIMITATIONS:\nI couldn’t find the code in the supplementary or any anonymous link to reproduce the results.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}