Unraveling the Gradient Descent Dynamics of Transformers

Authors: bingqing  song
boran  han
shuai  zhang
jie  ding
mingyi  hong
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:38:27

Abstract: While the Transformer architecture has achieved remarkable success across various domains, a thorough theoretical foundation explaining its optimization dynamics is yet to be fully developed. In this study, we aim to bridge this understanding gap by answering the following two core questions: (1) Which types of Transformer architectures allow Gradient Descent (GD) to achieve guaranteed convergence? and (2) Under what initial conditions and architectural specifics does the Transformer achieve rapid convergence during training? By analyzing the loss landscape of a single Transformer layer using Softmax and Gaussian attention kernels, our work provides concrete answers to these questions. Our findings demonstrate that, with appropriate weight initialization, GD can train a Transformer model (with either kernel type) to achieve a global optimal solution, especially when the input embedding dimension is large. Nonetheless, certain scenarios highlight potential pitfalls: training a Transformer using the Softmax attention kernel may sometimes lead to suboptimal local solutions. In contrast,  the Gaussian attention kernel exhibits a much favorable behavior. Our empirical study further validate the theoretical findings.

Main paper: 
 introduction: 
 Transformer model architectures have become popular in machine learning, delivering remarkable performance across a wide array of tasks. From natural language processing [Vaswani et al., 2017 , Beltagy et al., 2020] to computer vision [Dosovitskiy et al., 2020] , these models have set new standards in performance and efficiency. Popular models include BERT [Devlin et al., 2018] , RoBERTa [Liu et al., 2019] , DeBERTa [He et al., 2020] , GPT models [Radford et al., 2019 , Brown et al., 2020] and ViT [Dosovitskiy et al., 2020] . Despite their empirical success, a comprehensive understanding of their optimization process remains elusive. As highlighted in Liu et al. [2020] , the training of large Transformers can sometimes result in deteriorated performance. It is therefore critical to develop theoretical insights for researchers and practitioners to better understand the practical performance of Transformers. However, the complexity of their architectures, coupled with the non-convex nature the authors explored rank collapse issues in token representations and their impact on training. The authors discussed the origin of the phenomenon of rank collapse and proposed depth-dependent scaling of residual branches as a potential solution. They specifically investigated scenarios where token rank equals one, which can hinder Transformer training. Their findings demonstrate the occurrence of the vanishing gradient issue, however, this work does not comprehensively characterize the vanishing gradient problem throughout the entire training process.A recent work Wu et al. [2024] analyzes the convergence behaviour of shallow Transformer, which shows the global convergence can be achieved with GD algorithm. However, the focus of our paper is different from Wu et al. [2024] . We not only derive the global convergence analysis (Our Theorem 2), but also investigates the role of different variables in optimization. Some other works focus on improving the optimization of Transformers empirically. [Huang et al., 2020] have proposed an initialization strategy such that no warm-up or layer normalization is needed to train Transformers efficiently; in Shazeer [2020] , the GLU variant of token embedding has been showed to be better than plain embedding in the optimization of Transformer models with Softmax attention kernel. It is worth noting that the above works all primarily focus on empirical investigations into the training of Transformer models, lacking a comprehensive theoretical analysis of the underlying mechanisms.Some recent research has focused on the convergence analysis of Transformer-based models within the in-context learning (ICL) framework. For instance, Huang et al. [2023] , Zhang et al. [2023] explores the learning dynamics of a one-layer Transformer with Softmax attention trained via gradient descent to learn linear function classes in-context. However, this line of study primarily addresses the general convergence performance of Transformers within the ICL setting and does not delve into the role of individual variables.Recall that we have defined the structure of a single Transformer model in Equation ( 1) and Equation (5). We will further define a few notations before we introduce a few useful lemmas that are needed in our proof.(1) Operator: Denote vec(•) as the vectorization operator on a matrix; ⊗ as Kronecker product operator; ⊙ as the element product. Denote Υ(•) as a matrix operator, such that for any matrix X without zero elementEQUATION(2) Matrix: Denote I as the identity matrix. Define matrix E and E as following:E =    E . . . E    Hn×Hn , E =    1 • • • 1 . . . . . . . . . 1 • • • 1    n×n .Define matrix P h as following:P h = . . . , E h n×n , . . . , h = 1, • • • , H.(3) Matrix in Transformer: Define the following matrix C related to the attention layer Softmax kernel:EQUATIONGaussian kernel:EQUATIONEQUATIONNext, let us introduce several useful lemma which leads to Theorem 2:1.2 Lemmas of Theorem 2 Lemma 1.(1)EQUATIONRemark: The above lemma derives the closed form of the gradient of objective over W V , W Q . Notice that we can derive the derivative of W K in the same way as W Q due to symmetry, so we do not include the derivation here. Some of the lemmas here refers https://say-hello2y.github.io/2022-09-07/attention-gradient Lemma 2. Consider updatingW Q , W K , W V at iteration t. Suppose σ max (W Q ), σ max (W K ), σ max (W V) are bounded during in the optimization phase, then we have the following conclusion:EQUATIONEQUATION) (4) ∥ ∂f (M ; X i ) ∂W Q ∥ F ≤ Q i ∥MH (M ; X i ) -y i ∥ F ,whereEQUATION) (5) ∥ ∂f (M ; X i ) ∂W K ∥ F ≤ K i ∥MH (M ; X i ) -y i ∥ F ,whereEQUATION) Lemma 3. Consider updating W Q , W K , W V at iteration t. Suppose σ max (W Q ), σ max (W K ), σ max (W V) are bounded during in the optimization phase, then we have the following conclusion:EQUATIONEQUATION)Lemma 4. Let f : R n → R be a second order differentiable function. Let x, y ∈ R n be given, and assume that ∥∇f(z) -∇f (x)∥ 2 ≤ C∥z -x∥ 2 for every z = x + t(y -x) with t ∈ [0, 1]. Then, f (y) ≤ f (x) + ⟨∇f (x), y -x⟩ + C ′ 2 ∥x -y∥ 2 . Lemma 5. For matrix A ∈ R k×l , B ∈ R l×m , C ∈ R m×n . vec(ABC) = (I n ⊗ AB) vec(C) = C T B T ⊗ I k vec(A) vec(AB) = (I m ⊗ A) vec(B) = B T ⊗ I k vec(A) vec(A ⊙ B) = vec(A) ⊙ vec(B).Proof Sketch of Theorem 2:The main idea of the proof follows from [Nguyen and Mondelli, 2020] . Let us first recall a few notations. λV := 2 3 1 + σ max (W V 0 ) , λ B := σ min (B 0 ). Using GD update rule, we aim to iteratively show            σ max (W V r ) ≤ 3 2 λV , r ∈ {0, . . . , t}, σ max (W Q r ) ≤ 3 2 λQ , r ∈ {0, . . . , t}, σ max (W K r ) ≤ 3 2 λK , r ∈ {0, . . . , t}, σ min (B r ) ≥ 1 2 λ B , r ∈ {0, . . . , t}, f (M r ; X) ≤ (1 -ηµ) r f (M 0 , X) , r ∈ {0, . . . , t} (34) Denote µ := 1 4 (λ B ) 2 ∥W O ∥ 2 2 .Let us discuss about the value of µ. We know W O ∈ R Hd×1 , B ⊤ 0 ∈ R HD×N n . We require µ > 0, i.e, λ B > 0, which implies B has full row rank. For simplicity, let us consider the H = 1 case. Recall the definition of B:B := S 11 X 1 • • • S N 1 X N Suppose we initialize W Q 1 , W K1 such that each S i1 ∈ R n×n is full rank, then we can easily show that rank(S i1 X i ) = n if X i has full row rank. Suppose embedding dimension D is large, with certain assumption on X, we can show B has full row rank. For example, if each X i follows standard Gaussian distribution with D >> N , then rank(B) = N n with probability 1 if we initializeW Q 1 , W K 1 such that S i1 is full rank.Further, let us assume thatH h=1 ( λQ h ) 2 > 1, H h=1( λK h ) 2 > 1, and initialization condition satisfies:54n 2 √ N H∥X∥ 6 F λV H h=1 ( λQ h ) 2 + ( λK h ) 2 (λ B ) 2 ∥W O ∥ 2 min λQ h , λK h , 1, λ B ≤ 1 (35)Remark 5. The initialization condition can be satisfied if ∥W O ∥ 2 is large and σ max (W V ) is small. ν in Equation ( 8) is 1 54 .It is clear that Equation (34) holds when t = 0. Suppose it holds at iteration t, we prove it holds at iteration t + 1.W V r+1 -W V 0 F (i) ≤ r s=0 W V r+1 -W V r F = η r s=0 ∥∇ W V f (M t ; X)∥ F (ii) ≤ η r s=0 ∥B r ∥ F ∥W O ∥ 2 ∥MH(M r ; X) -y∥ 2 (iii) ≤ η∥B r ∥ F ∥W O ∥ 2 r s=0 (1 -ηµ) s/2 ∥MH(M 0 ; X) -y∥ 2 ,where (i) uses the triangle inequality; (ii) plugs in the expression of ∇ W V f (M t ; X) and uses the Cauchy-Schwartz inequality; (iii) is because we assume the loss function f (•) linearly decreases until t-th iteration. Let u = √ 1 -ηµ. So we haveEQUATIONEQUATIONwhere (i) is because each element in S r,i has magnitude at most 1 and ∥S r,i ∥ F ≤ n √ H, then by Cuachy-Schwartz inequality, we have ∥B∥ r ≤ √ HN ∥X∥ F ; (ii) is due to the initialization condition. Then by Weyl's inequality, there isσ max W V r+1 ≤ σ max (W V 0 ) + 1 = 3 2 λV .Similarly, let us derive the upper bound for σ max (W Q h,r ).W Q h,r+1 -W Q h,0 F (i) ≤ r s=0 W Q h,r+1 -W Q h,r F = η r s=0 ∇ W Q h f (M t ; X) F (ii) ≤ η r s=0 N i=1 Q 2 i ∥MH (M r ; X) -y∥ 2 (iii) ≤ η N i=1 Q 2 i r s=0 (1 -ηµ) s/2 ∥MH (M 0 ; X) -y∥ 2 ≤ N i=1 Q 2 i µ 1 -u r+1 1 -u 1 -u 2 ∥MH (M 0 ; X) -y∥ 2 ≤ 2 N i=1 Q 2 i µ ∥MH (M 0 ; X) -y∥ 2 (iv) ≤ 1 2 λQ h, where (i) uses triangle inequality; (ii) uses Lemma 2 (4); (iii) comes from the assumption that loss function f (•) linearly decreases until t-th iteration; (iv) is due to the initialization condition Equation ( 35). Similarly, we can showEQUATIONThen by Weyl's inequality, there isσ max W Q h,t+1 ≤ σ max (W Q h,0 ) + 1 2 λQ h = 3 2 λQ h ; σ max W K h,t+1 ≤ σ max (W K h,0 ) + 1 2 λK h = 3 2 λK h .Now we aim to bound the eigenvalues of B r+1 .∥B r+1 -B 0 ∥ F ≤ r s=0 ∥B s+1 -B s ∥ F = N i=1 r s=0 ∥S i,s+1 X i -S i,s X i ∥ F (i) ≤ N i=1 r s=0 ∥X i ∥ F ∥S i,s+1 -S i,s ∥ F (ii) ≤ N i=1 r s=0 ∥X i ∥ F H h=1 ∥S ih,s+1 -S ih,s ∥ F (iii) ≤ η N i=1 r s=0 ∥X i ∥ F • ϕ 2 i + ψ 2 i • ∥ (∇ W Q f (M s ; X i ), ∇ W K f (M s ; X i )) ∥ F (iv) ≤ N i=1 r s=0 ∥X i ∥ F • ϕ 2 i + ψ 2 2 • Q 2 i + K 2 i • ∥MH (M s , X i ) -y i ∥ F (v) ≤ η r s=0 N i=1 ∥X i ∥ 2 F (ϕ 2 i + ψ 2 i )(Q 2 i + K 2 i )(1 -ηµ) s/2 ∥MH (M 0 ; X) -y∥ 2, where (i) and (ii) uses triangle inequality and Cauchy-Schwartz inequality; (iii) comes from Lemma 2 (5); (iv) uses Lemma 2 and Cauchy-Schwartz inequality; (v) comes from Cauchy-Schwartz inequality.Together with our initialization condition, we have∥B r+1 -B 0 ∥ F ≤ 1 µ N i=1 ∥X i ∥ 2 F (ϕ 2 i + ψ 2 i )(Q 2 i + K 2 i ) • 1 -u r+1 1 -u ∥MH (M 0 ; X) -y∥ 2 ≤ 2 µ N i=1 ∥X i ∥ 2 F (ϕ 2 i + ψ 2 i )(Q 2 i + K 2 i ) ∥MH (M 0 ; X) -y∥ 2 (i) ≤ 1 2 λ B ,where (i) comes from the initialization condition 35. By Weyl's inequality, we can derive the bound for the singular values of B t :σ min (B r+1 ) ≥ σ min (B 0 ) -∥B r+1 -B 0 ∥ F ≥ 1 2 λ B .The final step is to show the last inequality holds. Since we have already showedσ max (W Q h ), σ max (W K h ), σ max (W V h )are bounded, by Lemma 3 (2) we can conclude that:∥∇f (M t+1 ; X) -∇f (M t )∥ 2 ≤ G∥M t+1 -M t ∥ FThus by Lemma 4, we choose η < 1 2G , then the following hold true:f (M t+1 ; X) = f (M t -η∇f (M t ; X); X) (i) ≤ f (M t ; X) -η ∥∇f (M t ; X)∥ 2 + G 2 η 2 ∥∇f (M t ; X)∥ 2 (ii) ≤ f (M t ; X) - 1 2 η ∥∇f (M t ; X)∥ 2 (iii) ≤ f (M t ; X) - 1 2 η ∂f (M t ; X) ∂W V 2 (iv) ≤ f (M t ; X) - 1 2 η∥W O ⊗ B ⊤ t (vec(MH(M t ; X) -y)) ∥ 2 (v) ≤ f (M t ; X) - 1 8 η∥W O ∥ 2 2 (λ B ) 2 • f (M t ; X) = (1 - 1 4 ∥W O ∥ 2 2 (λ B ) 2 ) • f (M t ; X) (vi) = (1 -ηµ)f (M t ; X),where (i) uses Lemma 4; (ii) is because we set η < 1 2G ; (iii) only considers the gradient over W V ; (iv) plugs in the closed form gradient in Lemma 1; (v) uses the property of smallest singular value and induction assumption; (vi) comes from the definition of µ.The following lemmas all consider the Transformers with Gaussian kernel attention 11. Lemma 6.(1)∂f (M ; X) ∂W V = B ⊤ (MH(M ; X) -y) W O ⊤ (39) (2) vec ∂f (M ; X) ∂W V = (W O ) ⊤ ⊗ B, vec(MH(M ; X) -y) = I Hd ⊗ B ⊤ • W O ⊗ I N • (MH(M ; X) -y) (40) (3) ∂f (M ; X) ∂W Q h = ∂f (M ; X) ∂C • ∂C ∂W Q h = N i=1 ∂f (M ; X i ) ∂C i • ∂C i ∂W Q h (41) (4) ∂f (M ; X i ) ∂C i = (MH(M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ ⊙ S i (42) (5) ∂f (M ; X) ∂C = ∂f (M ; X 1 ) ∂C 1 , • • • , ∂f (M ; X N ) ∂C N ⊤ (43) Lemma 7. Consider updating W Q , W K , W V iteration t. Suppose σ max (W Q ), σ max (W K ), σ max (W V) are bounded during in the optimization phase, then we have the following conclusion:(EQUATIONEQUATION)whereEQUATIONEQUATIONEQUATION)where min |V ′ W O | is the smallest absolute value of each element in vector V ′ W O ; min S is the smallest element in matrix S.Lemma 8. Consider updating W Q , W K , W V at iteration t. Suppose σ max (W Q ), σ max (W K ), σ max (W V) are bounded during in the optimization phase, then we have the following conclusion:∥MH(M t+1 ; X) -MH(M t ; X)∥ F ≤ Z ′ ∥M t+1 -M t ∥ F , where Z ′ is some positive constant. (50) ∥∇f (M t+1 ; X) -∇f (M t ; X)∥ 2 ≤ G ′ ∥M t+1 -M t ∥ F ,where G ′ is some positive constant.(51)1.5 Proof Sketch of Theorem 3.(1)Using GD update rule, we aim to iteratively show       σ max (W Q r ) ≤ 3 2 λQ , r ∈ {0, . . . , t}, σ min ∂C h (Mr) ∂W Q h ≥ 1 2 δ, r ∈ {0, . . . , t}, min S r ≥ κ, r ∈ {0, . . . , t}, f (M r ; X) ≤ (1 -ηγ) r f (M 0 , X) , r ∈ {0, . . . , t} (52) Denote γ := 1 2 δ 2 κ 2 min V ′ W O 2 .Let us discuss about the value of γ. We knowW O ∈ R Hd V ×1 , B ⊤ 0 ∈ R HD×N n, where Hd > 1, HD > N n. We require γ > 0, i.e, δ > 0, κ > 0, minV ′ W O > 0. It is clear that κ > 0 can hold as long as W Q h is bounded. And it is easy to show that if X i ̸ = 0, we can always choose W V and W O , such that min V ′ W O > 0. Since ∂C h (M ) ∂W Q h ∈ R N n 2 ×Dd , suppose we initialize W Q h , W K h such that rank( ∂C h (M0) ∂W Q h ) = N n 2 , then we have σ min ∂C h (M0) ∂W Q h≥ δ for some positive constant δ. Further, we assume the initialization condition satisfies:8n∥X∥ 5 F ∥W O ∥ 2 λV ( λQ h + λK h ) exp 9 4 ∥X∥ 2 F λQ h 2 + λK h 2 δ 2 (min (|V ′ W O |)) 2 • min δ, λQ h ∥MH (M 0 ; X) -y∥ 2 ≤ 1 (53)Remark 6. The initialization condition can be satisfied if ∥W O ∥ 2 is large and σ max (W V ) is small. ν ′ in Equation ( 12) is 1 8 .Similar to the proof of Theorem 2, we use induction to prove the theorem. Equation ( 52) holds when t = 0. Suppose it holds at iteration t, we prove it holds at iteration t + 1.W Q h,r+1 -W Q h,0 F (i) ≤ r s=0 W Q h,r+1 -W Q h,r F = η r s=0 ∇ W Q h f (M t ; X) F (ii) ≤ η r s=0 N i=1 Q ′ i 2 ∥MH(M r ; X) -y∥ 2 (iii) ≤ η N i=1 Q ′ i 2 r s=0 (1 -ηγ) s/2 ∥MH(M 0 ; X) -y∥ 2 ,where (i) uses triangle inequality; (ii) comes from Lemma 7 and Cauchy-Schwartz inequality; (iii) is from the induction assumption that loss function f (•) linearly decreases until t-th iteration. Let u = √ 1 -ηγ. So we haveEQUATIONwhere (i) comes from the initialization condition. Then by Weyl's inequality, there isσ max W Q h,t+1 ≤ σ max (W Q h,0 ) + 1 2 λQ h = 3 2 λQ h . ∂C h (M r+1 ) ∂W Q h - ∂C h (M 0 ) ∂W Q h F (i) ≤ r s=0 ∂C h (M s+1 ) ∂W Q h - ∂C h (M s ) ∂W Q h F (ii) ≤ η √ n∥X∥ 2 F r s=0 ∇ W Q h f (M s ; X) F (iii) ≤ η √ n∥X∥ 2 F r s=0 N i=1 Q ′ i 2 ∥MH (M s ; X) -y∥ 2 (iv) ≤ η √ n∥X∥ 2 F N i=1 Q ′ i 2 r s=0 (1 -ηγ) s/2 ∥MH (M 0 ; X) -y∥ 2 , ≤ 2 γ √ n∥X∥ 2 F N i=1 Q ′ i 2 ∥MH (M 0 ; X) -y∥ 2 (v) ≤ 1 2 δ,where (i) uses triangle inequality; (ii) applies Lemma 7 (2) and Cauchy-Schwartz inequality; (iii) uses Lemma 7 (4); (iv) applies the induction assumption that the loss function f (•) linearly decreases until t-th iteration; (v) comes from the initialization condition. Then by Weyl's inequality, there isσ max ∂C h (M t+1 ) ∂W Q h ≥ σ max ∂C h (M 0 ) ∂W Q h - 1 2 δ = 1 2 δ.For each element in S ih , we have close formS W Q h , W K h ; X i kj = exp - 1 2 √ d X ik • W Q h -X ij • W K h 2Since we have already showed that σ max W Q h,r ≤ 3 2 λQ h , it follows directly each element in matrix S t is lower bounded by some constant κ for any t. Now we derive the expression of κ:exp - 1 2 √ d X ik • W Q h,t -X ij • W K h 2 (i) ≥ exp - 1 √ d ∥X ik • W Q h,t ∥ 2 + ∥X ij • W K h ∥ 2 (ii) ≥ exp - 1 √ d 9 4 ( λQ h ) 2 ∥X ik• ∥ 2 + ( λK h ) 2 ∥X ij• ∥ 2 (iii) ≥ exp - 9 4 ∥X∥ 2 F ( λQ h ) 2 + ( λK h ) 2 := κ,where (i) uses Cauchy-Schwartz inequality; (ii) applies the induction assumptionσ max (W Q h,t ) ≤ 3 2 λQ hand property of singular value; (iii) is because d ≥ 1. Thus, we have min S t ≥ κ. Finally, we aim to show f (M t+1 ; X) ≤ (1 -ηγ)f (M t , X). By Lemma 8, since we have showed thatσ max (W Q h ) is bounded, we can directly derive that ∥∇f (M t+1 ; X) -∇f (M t ; X)∥ 2 = ∇ W Q h f (M t+1 ; X) -∇ W Q h f (M t ; X) 2 ≤ G ′ ∥M t+1 -M t ∥ FFinally, by Lemma 4, choose η < 1 2G ′ , we have the following holds:f (M t+1 ; X) = f (M t -η∇f (M t ; X); X) (i) ≤ f (M t ; X) -η ∥∇ W Q f (M t ; X)∥ 2 + G ′ 2 η 2 ∥∇ W Q f (M t ; X)∥ 2 (ii) ≤ f (M t ; X) - 1 2 η ∇ W Q h f (M t ; X) 2 (iii) = f (M t ; X) - 1 2 η ∂f (M t ; X) ∂C(M t ) • ∂C(M t ) ∂W Q h 2 F (iv) ≤ f (M t ; X) - 1 4 ηδ 2 ∂f (M t ; X) ∂C (M t ) 2 F (v) ≤ f (M t ; X) - 1 4 ηδ 2 (MH (M ; X) -y) W O ⊤ (V ′ ) ⊤ ⊙ S 2 F (vi) ≤ f (M t ; X) - 1 4 ηδ 2 κ 2 • (min |V ′ W O |) 2 ∥MH (M 0 ; X) -y∥ 2 2 (vii) = (1 -ηγ)f (M t ; X),where (i) uses Lemma 4 (2); (ii) is because we choose η < 1 2G ′ ; (iii) writes down the expression of gradient according to chain rule in Lemma 6; (iv) uses the induction assumption (2)Next, we show the convergence result for Transformer with Softmax kernel with only W Q updated. Since we assume parameters are all bounded during optimization phase, by Lemma 8, we can easily show that there exists constant G ′ (see xx for details), such thatσ max ∂C h (Mt+1) ∂W Q h ≥ 1 2 δ∇ W Q h f (M t+1 ; X) -∇ W Q h f (M t ; X) 2 ≤ G ′ ∥M t+1 -M t ∥ F (56) Then by Lemma 4, choose η ′ < 1 2G ′ we have f (M t+1 ; X) = f (M t -η∇f (M t ; X) ; X) ≤f (M t ; X) -η ′ ∥∇f (M t ; X)∥ 2 + G ′ 2 η ′2 ∥∇f (M t ; X)∥ 2 ≤ f (M t ; X) - 1 2 η ′ ∥∇f (M t ; X)∥ 2 1.6 Proof of Lemma in Section 1.2Proof of Lemma 2 (1).Proof.Step 1: When W Q , W K are updated, we aim to prove∥d(S i )∥ F ≤ n∥d(C i )∥ F .Step 2: We aim to show ∥d(C i )∥ F ≤ n √ d ∥X i ∥ 2 F H h=1 σ 2 max W K h • d W Q F .Combine the above two steps, we can derive the bound in Equation ( 27). Proof of Step 1: First, we can write down the closed form of the differential of S i :∥d(S i )∥ F = ∥S i ⊙ d(C i ) -S i ⊙ Υ((exp C i )E) ⊙ d(exp(C i )E))∥ F (57)We reorganize the terms on the right side of Equation ( 57), we have the following equation: 58). We write down the close form of the element in the k-th row and j-th column:∥d(S i )∥ F = ∥S i ⊙ d(C i ) -Υ((exp C i )E) ⊙ d((exp C i )E) ∥ F = ∥S i ⊙ d(C i ) -Υ((exp C i )E) ⊙ ((exp C i ) ⊙ d(C i )) E ∥ F (58) Since C i = [C i1 , • • • , C iH ], we will investigate each C ih , h = 1, 2, • • • , H. We focus on the term d(C i ) -Υ((exp C i )E) ⊙ (exp(C i ) ⊙ d(C i )) E in Equation (EQUATIONEQUATIONwhere (i) is expand the closed form of Equation ( 59); (ii) uses the Cauchy-Schwartz inequality; (iii) is because each element in the square root in (ii) is upper bounded by 1. With Equation ( 61), we can easily show∥d(C ih ) -Υ((exp C ih )E) ⊙ ((exp C ih ) ⊙ d(C ih )) E∥ F ≤ √ n n k=1 n j=1 ∥d(C ihk )∥ 2 F ≤ n∥d(C ih )∥ F (63)Since every element in S i has magnitude less than 1, we haveEQUATIONEQUATIONwhere (i) is from Cauchy-Schawatz inequality. 
notations and problem description: 
 In this section, we define the structure of the Transformer model and describe the training problem. We consider a one-layer attention Transformer model with multiple heads and a dataset with N samples. Each data sample consists of n discrete tokens, each with embedding dimension D. We denote the dataset as {(X i , y i )} N i=1 , where X i ∈ R n×D , and y i ∈ R n is the label of the dataset. The output from the Transformer model is the prediction of the label. The Transformer structure is formulated as follows:Attention(W Q h , W K h , W V h ; Xi) := S(W Q h , W K h ; Xi)XiW V h (1) MH(W Q , W K , W V ; Xi) := Concat (head1, . . . , headH) • W O ,whereEQUATIONIn the above notation, W Q h , W K h ∈ R D×d is the query weight matrix and key weight matrix, respectively; W V h ∈ R D×d is the value weight matrix; these matrices are the main optimization variables throughout the paper. Further W O ∈ R Hd×1 is a fixed matrix, representing the weight of the output layer; H is the number of attention heads; S(•) is a kernel function of variables W Q , W K and input X i . Attention(•) is the attention head function; MH(•) represents the multi-head attention function. For example, with the Softmax attention [Vaswani et al., 2017] , S(•) can be written as:EQUATIONwhere for a given n × n matrix Z, Sofmax(Z) := [Softmax(Z 1 ), • • • , Softmax(Z n )]. Throughout, let us denote S(•) kj as the element of k-th row and j-th column in matrix S(•). Let X ik• ∈ R D denote the embedding of the k-th token in data X i , which is the k-th row of matrix X i . The structure of Transformer model can be found in Fig 1 , where we denote S ih := S W Q h , W K h ; X i .Based on the above Transformer model, we consider minimizing the following empirical ℓ 2 loss function for the entire data set {X i , y i } N i=1 :EQUATIONwhere M := (W Q , W K , W V ) is the set of variables that can be optimized.For notation simplicity, next we define the vector version of the Transformer model given in Equation (1), for the entire dataset {(X i , y i )} N i=1 . Towards this end, let X ∈ R N n×D denote the columnstacked matrix of each single data X i . Similarly, define the stacked label y ∈ R N n . Then we can define: Let B :=EQUATIONS 11 X 1 • • • S 1H X 1 • • • • • • • • • S N 1 X N • • • S N H X N, andW V := diag(W V 1 , • • • , W V H ) ∈ RHD×Hd denote the diagonalized weight matrices that include all value weight matrices for all attention heads. Using these definitions, We can simplify Equation (5) asMH(M ; X) = B • W V • W OThus the empirical loss function given in Equation ( 4) can be simplified asEQUATION)For more notations in the following sections, we will use subscript t to represent the variables in t-th iteration, e.g,M t := {W Q t , W K t , W V t }.Similarly, we denote B t as the matrix B at t-th iteration.It is important to note that, in the above description and throughout the paper, we model the Transformer training problem by using a single-layer Transformer, with a regression loss. In practice Transformer models can exhibit greater complexity (different loss functions, multiple layers, etc). For example, the text classification task has an additional mean pooling layer followed by the output of the Transformer structure. Further, they usually contain downstream MLP modules. However, we choose to use the simplified version due to the following reasons:First, the primary objective of this work is to understand how different attention kernels affect the training dynamics of the Transformers, so we do not include the layer normalization in our model. In fact, in the literature, many works that analyze popular network structures also do not consider layer normalization. For example, in [Huang et al., 2023 , Zhang et al., 2023] , both analyze the convergence performance of Transformers but normalization is not considered. Second, we do not include the downstream MLP module in our work since we are interested in the role of self-attention layer in convergence analysis, and the single-attention model is also the standard model used in [Huang et al., 2023 , Zhang et al., 2023] . Further, the analysis of MLP is standard in literature [Allen-Zhu et al., 2019 , Du et al., 2019 , Nguyen and Mondelli, 2020] . And it is worth noting that our choice to focus on a one-layer Transformer is consistent with other works that similarly aim to investigate the core training dynamics of Transformers, e.g, in [Tian et al., 2023] , a single-layer Transformer is considered as a basic model. 
convergence analysis: 
 In this section, we present our theoretical analysis for solving problem (6). We focus on the behavior of the vanilla GD algorithm for optimizing the variable set M , where M ⊂ {W Q , W K , W V }. Below we summarize our results. = {W V } and M = {W Q , W K , W V }.Different behavior between Softmax and Gaussian Kernel Attention. When S(•) is Gaussian and the embedding dimension D is at least O(N n), convergence to global optimal is also ensured for M = {W Q }. Interestingly, under the same conditions of large D, convergence to global optimal is not guaranteed when S(•) is Softmax.In the subsequent sections, we will elaborate on these convergence results in detail, providing a deeper understanding of the nuances in Transformer behavior under varying configurations. To set up our analysis, we introduce λ V as the smallest eigenvalue of W V 0 , λ B as the smallest eigenvalue of B 0 , λQ h , λK h , λV as the largest singular value of matrix W Q h,0 , W K h,0 , W V , respectively. We denote ∥ • ∥ 2 as ℓ 2 norm and ∥ • ∥ F as Frobenius norm. Further, we denote σ max (•) and σ min (•) as the largest and smallest singular value of a matrix, respectively. For any vector v, let min(|v|) denote the smallest absolute value of vector v.First, we examine the role of W V in the optimization of multi-head attention network structure. Our analysis demonstrates that with the hidden dimension HD ≥ N n and proper initialization, the global optimal solution of (6) can be found using a vanilla gradient descent algorithm. The initialization requires that the matrix B 0 has full rank. Our first result shows that, overparameterized Transformer can be trained to global optimal solution. Theorem 1. Consider problem (4) with S(•) being instantiated as the Softmax kernel given in (3). Consider the following update for the variable M = {W V }:EQUATIONwhere α := ∥W O ∥ 2 (λ B ) 2 > 0; η > 0 is defined in Appendix 1.3, and chosen such as ηα < 1.Remark 1. The aforementioned theorem focuses on the convergence behavior when only W V is being updated. We further elaborate on the initial conditions ensuring λ B > 0.Note that λ B > 0 implies that the objective function f exhibits a landscape that is nearly convex, which is crucial for optimization. By definition, this condition implies that B 0 has full rank, which can be fulfilled by selecting appropriate W Q 0 and W K 0 , plus having large enough embedding size, satisfying D ≥ N n/H. We refer the readers to Appendix 1.3 for the derivation of this condition, which can be guaranteed by random initialization with high probability.Furthermore, it is important to note that our work aligns with existing literature on the subject of embedding size in Transformer models. For example, in [Bhojanapalli et al., 2020] , the authors restrict their focus to the simplified case of N = 1, H = 1. They establish the necessary condition for Softmax attention to overcome its low-rank bottleneck, which requires D ≥ n . In our analysis, we derive a similar necessary condition on Transformer model size (D ≥ n × (N/H)) to guarantee the global convergence when a Transformer model is trained with GD.In Theorem 1, we have illustrated the case where only updating W V already leads to global convergence. However, in practice, all parameters W V , W Q , W K are updated. This case is more challenging to analyze due to the non-linearity introduced by the Softmax function. Next, we show that a similar result in Theorem 1 still holds when all the parameters are updated simultaneously. Theorem 2. Consider problem (4), with S(•) being instantiated as the Softmax kernel. Consider the GD update where M = {W Q , W K , W V }: Suppose λ B > 0, and the initialization M 0 satisfyEQUATION)Then there exists stepsize η > 0, such thatEQUATIONwhere β := ∥W O ∥ 2 (λ B ) 2 > 0, and the constants η, ν are defined in Appendix 1.3. Remark 2. In the stated theorem, we simplify our analysis by excluding the downstream MLP module in the typical Transformer model, since it is easy to combine the model in Equation (2) with downstream MLP layers. Further, it can be directly showed that the Transformer with MLP will lead to the same convergence rate of the optimization problem as updating W Q , W K , W V only. To illustrate this, consider the following Transformer model:EQUATION)where W l ∈ R n l-1 ×n l , and n 0 = d O . Based on the Transformer model defined in Equation ( 10), we have the following corollary.Corollary 1. Consider problem min M 1 2 ∥G(M ; X) -y∥ 2 , with G(•) being defined in Equation (10) and S(•) being instantiated as the Softmax kernel. Suppose that the MLP module satisfies:n 1 ≥ n 2 • • • ≥ n L .Consider the following GD update (whereM = {W Q , W K , W V , W 1 , • • • , W L }): Suppose λ B > 0.Then, there exists a step size η > 0 and initialization weight M 0 , such that the loss function linearly converges to 0. Remark 3. The above theorem and corollary describe the global convergence guarantee when W Q , W K and W V are updated. This is in line with the insights gained from Theorem 1. However, the conditions for initialization are more stringent, and the optimization landscape becomes inherently more complex due to the involvement of the Softmax attention through W Q and W K .To ensure the initial condition 8, we have two options: 1) Initializing M 0 such that ∥MH(M 0 ; X) -y∥ F is small, which implies that the optimization starts in a region close to the global optimal solution and that the initial weight is close to the global optimal solution; 2) Balancing between W O and W V , in the sense that ∥W O ∥ 2 is large and λV is small. For a detailed account of these initialization strategies, please refer to Appendix 1.3.Finally, we need to point out that for Transformers with Gaussian kernel attention, we can derive similar convergence results as long as the attention kernel maintains full rank and weights are initialized appropriately. Here we do not include the theoretical statement since it is similar to the result for Softmax attention.4.2 Softmax vs Gaussian kernel: Softmax attention Transformers may exhibit slower convergence.In the previous section, we explored the global convergence of training Transformer models. However, from Theorem 2, it was not clear what roles do matrices W Q and W K play in the entire convergence process, since Theorem 1 indicates that optimizing W V alone already ensures the desired convergence. Nevertheless, it is the matrices W K and W Q that truly represent the power of a Transformer model, because they are used to extract token correlations.To study how well a Transformer model can extract the token correlation, in this section, we will study the GD dynamics for Transformer models, where only W K and W Q are optimized (while fixing W V ). If optimizing these two parameters alone can still achieve zero training loss, then we claim that the input token correlation can be optimally extracted by the Transformer model.To begin our study, let us define that Gaussian kernel to be an n × n matrix, where its k-th row and j-th column of is given by:S W Q h , W K h ; Xi kj = exp - 1 √ d X ik• W Q h -Xij•W K h 2 (11)Since the training dynamics/gradients of variables W Q and W K have the same property in (3) and (11), we will only concentrate on optimizing W Q .With some abuse of notation, define a matrix C for Softmax attention and Gaussian kernel attention, respectively. Softmax attention:C ih := XiW Q h (XiW K h ) ⊤ √ d ∈ R n×n .Gaussian kernel attention:C ih ∈ R n×n ; (C ih ) kj = - ∥Xik•W Q h -Xij•W K h ∥ 2 2 √ d. For both Softmax attention and Gaussian kernel attention:C i ∈ R n×Hn = [C i1 , C i2 , • • • , C iH ] ; C ∈ R N n×Hn = C ⊤ 1 , C ⊤ 2 , • • • , C ⊤ N ⊤ .Using the above notation, the activation function S(•) in ( 3) and ( 11) can be related to the matrices C's in the following manner: Softmax attention : S ih = Softmax (C ih ) , Gaussian attention : (S ih ) kj = exp (C ih ) kj .Additionally, note that C is a function of variables M . Therefore we will sometimes use C(M ) when we need to emphasize the dependency of C on M .Next, we will outline the conditions under which GD can still successfully find global optimal solutions for Transformers with Gaussian kernel attention (when only W Q is updated), while under the same set of conditions, but with Softmax kernel attention, GD fails.Theorem 3. Solve problem (4) with the following GD update (withM = {W Q }): W Q t+1 = W Q t -η∇ W Q f (M t ; X). Suppose δ h := σ min ( ∂C(M0) ∂W Q h ) > 0, ∀ h ∈ [1, 2, • • • , H],and the initialization condition further satisfiesEQUATIONν ′ is defined in Appendix 1.5.(1) When S(•) is a Gaussian kernel function, there exists a stepsize η and a positive constant γ, such that fEQUATIONwhere γ, η are defined in Appendix 1.5.(2) When S(•) is a Softmax function, suppose W Q t is bounded during the training phase, then there exists stepsize η, such thatEQUATION)where η ′ is defined in Appendix 1.5. Remark 4. First, it's important to note that the parameter size must satisfy Dd ≥ N n 2 for δ > 0 to hold. It is crucial to emphasize the fundamental distinction in convergence outcomes between Transformers employing Gaussian kernel attention and those utilizing Softmax attention under these conditions. With equivalent initialization conditions, training Transformers equipped with Gaussian kernel attention achieves global convergence using gradient descent (GD). Second, it is essential to emphasize that the dimension size Dd ≥ N n 2 is similar to the findings of works that have analyzed the convergence performance of over-parameterized neural networks Allen-Zhu et al.[2019], Du et al. [2019] . The total number of samples, consisting of N samples each with n tokens, can be calculated as N n. Meanwhile, the total feature dimension is Dd. The inequality implies that the width of the parameters is at least O(N ), a relationship also illustrated in Nguyen and Mondelli [2020] .In part (2), we demonstrate that the PL condition does not hold. In particular, we identify an initial solution that satisfies all the conditions given in Theorem 3, yet fails to satisfy the PL condition. Therefore, in this case, GD leads to vanishing gradients without being able to find a global optimal solution. The details of this specific example are provided below.Example: Consider Transformer with Softmax attention, and N = 1, n = 2, H = 1. Let us first write down the close form of the gradient overW Q 1 : ∂f (M0; X1) ∂W Q 1 = 1 √ d X ⊤ 1 ∂f (M0; X1) ∂C11 X1W K 1,0 Next, we show there exists W O , W V , X 1 , W Q 1,0 , W K1,0 such that the loss function is non-zero with Equation ( 12) satisfied, while∂f (M0; X1) ∂C11 = 0 ∈ R 2×2 . Denote L := ∂f (M0;X1) ∂MH(M0;X1) W O ⊤ X 1 W V 0 ⊤ ∈ R 2×2 . ∂f (M0;X1)∂C11 can be expressed as follows:/∂f (M0; X1) ∂C11 11 = δ • (L11 -L12), ∂f (M0; X1) ∂C11 12 = δ • (L12 -L11), δ is some constant.Next, we will give the value of W O , W V 0 to show the case where GD leads to vanishing gradient.Let D = d = 2, W O = ( 1 a , 1 a ), X 1 = 1 0 0 1 , and W V 0 = 2a a a 2a, where a is a constant. It is easy to show that there exists W Q 1 and W K 1 such that Equation ( 12) holds. Further, it is easy to verify that for this scenario, the following holds:EQUATIONNext, we can easily deduce that ∂f (M0;X1)∂C11 11 = ∂f (M0;X1) ∂C11 12= 0. Similarly, we can demonstrate that ∂f (M0;X1)∂C11 21 = ∂f (M0;X1) ∂C11 22= 0. Consequently, we have ∂f (M0;X1)∂W Q 1 = 0. However, if y 1 satisfies that ∂f (M0;X1) ∂MH(M0;X1) ̸ = 0, it follows f (M 0 ; X 1 ) ̸ = 0, which means M 0 is not global optimal solution.5 Experiment: Softmax v.s. Gaussian In this section, we present numerical results to illustrate the behaviors of Transformers models with Softmax attention and Gaussian kernel attention across various tasks. 
dataset: 
 We investigate two distinct tasks: Text Classification using the IMDb review dataset [Maas et al., 2011] and Pathfinder [Linsley et al., 2018] . While both tasks involve processing long sequences, they exhibit different characteristics. Text Classification is a well-known NLP task that focuses on discerning relationships among token embeddings, while the Pathfinder task prioritizes capturing spatial information within the input pixels. In the first experiment, we plot the test accuracy and test loss within the training steps with both Softmax and Gaussian kernel attention on both tasks. We repeat the training for 10 times and make the shadow plot on the test performance. In our second experiment, the training process consists of two stages: In the first stage, we train the Transformer model equipped with Softmax attention (defined in Equation ( 3)) for 8,000 steps. In the second stage, we continue training from the pre-trained model for an additional 500 steps, with the option of using either Softmax or Gaussian kernel. To explore the optimization landscape around the trained model, we employed a technique inspired by Li et al. [2018] . We select two parameter directions, specifically the W Q and W K matrices in the first Transformer layer. These two directions, denoted as d 1 , d 2 , are centered at the trained model M , and represent the parameter space of W Q , W K , respectively. We evaluate the loss function on the set {M + 0.02(r -25)d 1 + 0.02(s -25)d 2 }, where r, s ∈ [1, 2, • • • , 50].The above set is the neighborhood of the trained model M , and we chose the evaluation stepsize as 0.02 along the two directions d 1 , d 2 , with the total steps limit as 100. Within this parameter space, we plot a 3-D surface representing the landscape around the trained model.To begin with, we present some observations in our first experiment. We plot the test performance of these two tasks on Transformers with two different types of attention. From The test accuracy has a significantly higher variance at the same training epoch. Further, the worst test accuracy after 20, 000 epochs is around 0.58 for the Softmax attention Transformer, compared with 0.62 for the Gaussian kernel Transformer. These observations align with the experiment results in [Chen et al., 2021] and [Tay et al., 2020] , where Transformers with different attention kernels are trained with the same model size and learning rate, while Softmax attention Transformers show instability in a few tasks.In Figure 4 , we present a comparison of the optimization landscape between Transformers with Softmax and Gaussian kernel attention. Notably, we observe distinct differences in the training landscapes of these two attention types for both tasks. We follow the visualization method described in Section 5.2. We conduct a visualization of the optimization landscape around the trained models after a two-stage training process, with identical learning rates, network sizes, and training epochs. Keeping all other factors consistent, the disparity in the landscape provides a direct representation of the difference in the attention structure during the optimization procedure. With Softmax attention, the landscape appears more complicated compared with Gaussian kernel attention. This complexity can be interpreted as the presence of a greater number of local optima in the optimization landscape, suggesting that Transformers utilizing Softmax attention may encounter more challenges in reaching global optimal solutions. In contrast, the landscape with the Gaussian kernel is flatter. This observation aligns with our earlier findings in Figure 2 and Figure 3 , where Softmax attention exhibited certain convergence issues. These observations also provide empirical evidence supporting our Theorem 3, which reflects in a slightly different perspective the complicated optimization landscape within the Softmax kernel. 
conclusion and future work: 
 In conclusion, our study addresses critical gaps in our understanding of why Transformer models perform exceptionally well in a variety of machine learning tasks. Our work also provides a nuanced understanding of the advantages and disadvantages of using classical Softmax attention in Transformers. We find that while shallow Softmax attention Transformers can achieve global convergence with overparameterization, there are scenarios where this attention structure can lead to local solutions. However, those issues can be mitigated by the Gaussian kernel-based attention. In our work, we need strong initialization and large embedding size, i.e, HD ≥ N n to obtain the global convergence, which exhibits a gap towards real case. In the future work, we will investigate how to relax the assumptions. Y. Wu, F. Liu, G. Chrysos, and V. Cevher. On the convergence of encoder-only shallow transformers.Advances in Neural Information Processing Systems, 36, 2024.R. Zhang, S. Frei, and P. L. Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023.1 Appendix 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [No] Justification: We do not include the compute resources detail.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: The paper has no harm in the research process or negative social impact. The paper is anonymous.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification:The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We have cited the code framework we use Chen et al. [2021] .Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Step 2: We aim to show ∥d(C i )∥ F ≤ n √ d ∥X i ∥ 2 F H h=1 σ 2 max W K h • d W Q F . Sim- ilarly, we investigate ∥d(C ih )∥ F , h = 1, 2, • • • , H. We have ∥d(C ih )∥ F = X i d(W Q h ) X i W K h ⊤ √ d F ≤ 1 √ d ∥X i ∥ 2 F σ max (W K h )∥d(W Q h )∥ F (67)Then plug the above inequality to Equation ( 63), we can deriveEQUATION)Thus by Cauchy-Schwartz inequality, it is easy to show∥d(S i )∥ F ≤ n √ d ∥X i ∥ 2 F H h=1 σ 2 max W K h • d W Q F .Proof of Lemma 2 (4).Proof. We first write down the close form of gradient of f (•) over W Q h by Lemma 1, and derive the upper bound of the norm of the gradient.∂f (M ; X i ) ∂W Q h F = 1 √ d X ⊤ i ∂f (M ; X i ) ∂C i P ⊤ h X i W K h F ≤ ∥X i ∥ 2 F σ max (W K h ) ∂f (M ; X i ) ∂C i F (69) By Lemma 1, there is ∂f (M ; X i ) ∂C i = (MH(M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ ⊙ S i - (MH(M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ ⊙ S i ⊙ Υ (exp C i )E E ⊤ ⊙ exp C i (70) Denote R i = (MH (M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ , R i = [R i1 , • • • , R iH ].Write down the close form of the element in the k-th row and j-th column:R ih S ih -(R ih ⊙ C ih ⊙ Υ ((exp C ih ) E)) E ⊤ ⊙ (exp C ih ) kj = R ihkj S ihkj - exp(C ihkj ) n j=1 R ihkj S ihkj n j=1 exp(C ihkj ) =     S ihkj - (exp C ihkj )S ihkj n j=1 exp(C ihkj )     • R ihkj - p̸ =j (exp C ihkp )S ihkj n j=1 exp(C ihkp ) R ihkp (i) ≤ 1 - exp (C ihkj ) n j=1 exp (C ihkj ) 2 + p̸ =j     exp (C ihkp ) n j=1 exp (C ihkj )     2 • ∥R ihk ∥ F (ii) ≤ √ n∥R ihk ∥ Fwhere (1) is due to the Cauchy-Schwartz inequality; (ii) is because each element within the squre root term in (i) has magnitude at most 1. Thus, we can further derive ∂f (M ;X i ) ∂C ih F = R ih ⊙ S ih -(R ih ⊙ S ih ⊙ Υ((exp C ih )E)) E ⊤ ⊙ exp C ih F (i) ≤ √ n n k=1 n j=1 ∥R ihk ∥ F (ii) ≤ n∥R ih ∥ F (iii) ≤ n∥X i ∥ F ∥W O ∥ 2 σ max (W V h )∥MH(M ; X i ) -y i ∥ 2, where (i) if from the bound in Equation ( 71); (ii) comes from Cauchy-Schwatz inwquality; (iii) uses the property of Frobenious norm. Thus, by Cauchy-Schwartz inequality, we can derive the upper bound for ∂f (M ;Xi) ∂Ci F.EQUATIONSo plug the above inequality into Equation ( 69), we can derive the upper bound for ∂f (M ;Xi)∂W Q h F : ∂f (M ; X i ) ∂W Q h F ≤ ∥X i ∥ 2 F σ max W K h ∂f (M ; X) ∂C i F ≤ n √ H ∥X i ∥ 3 F W O 2 σ max W K h σ max W V h ∥MH (M ; X i ) -y i ∥ 2 ≤ n √ H ∥X i ∥ 3 F W O 2 H h=1 σ 2 max W K h σ max W V ∥MH (M ; X i ) -y i ∥ 2Proof of Lemma 3 (1). By Mean Value Theorem and Cauchy-Schwartz inequality, |f (M t+1 ;EQUATIONwhere M ′ t is between M t and M t+1 . We can derive the upper bound of the norm of∇ W V f (M ; X i ): ∂f (M t ; X i ) ∂W V F = ∥B ⊤ i (MH(M t ; X i ) -y i ) W O ⊤ ∥ F ≤ ∥B i ∥ F ∥MH(M t ; X i ) -y i ∥ F ∥W O ∥ 2 ≤ n √ H∥X i ∥ F ∥W O ∥ 2 ∥MH(M t ; X i ) -y i ∥ F (73) By Lemma 2, we know ∂f (M t ; X i ) ∂W Q F ≤ Q i ∥MH (M ; X i ) -y i ∥ 2 ; ∂f (M t ; X i ) ∂W K F ≤ K i ∥MH (M ; X i ) -y i ∥ 2 . ∥f (M t+1 ; X i ) -f (M t ; X i )∥ 2 ≤ Q 2 i + K 2 i + n 2 Hσ 2 max (X i )∥W O ∥ 2 ∥M t+1 -M t ∥ F := Z i ∥M t+1 -M t ∥ F(74) Therefore, together with Equation (73), we haveEQUATIONProof of Lemma 3 (2).Proof. By triangle inequality, we haveEQUATIONEQUATIONStep 1: Derive upper bound for∥∇ W Q f (M t+1 ; X i ))-∇ W Q f (M t ; X i ))∥ F = ∥ vec(∇ W Q f (M t+1 ; X i ))-vec(∇ W Q f (M t ; X i ))∥ 2 .First, we give the vectorized expression of∇ W Q f (M t ; X i ). Recall we denote U i = (MH (M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ ⊙ S i .By Lemma 1, we can derive the close form ofEQUATIONwhere (i) uses the Lemma 1; (ii) and (iii) comes from the property of vectorization in Lemma 5; (vi) uses the definition of S i ; (v) gives an equivalent expression of vec(U i ); (vi) reorganizies (v). Further, it is easy to verify that:EQUATIONNext, let us derive upper bound forEQUATIONwhere (i) plugs in the expression in Equation ( 79); (ii) uses the fact that each element in S i,t+1 has magnitude at most 1, and Cauchy-Schwartz inequality; (iii) comes from the definition of I, E and R;(iv) uses the definition of U i,t ; (v) is because triangle inequality; (vi) uses the fact that each element in S i,t+1 has magnitude at most 1, and Cauchy-Schwartz inequality. Next, we aim to derive upper bound of ∥R i,t+1 -R i,t ∥ F in Equation ( 81).EQUATIONwhere (i) is because of the triangle inequality; (ii) uses the definition of Z i in Equation ( 74), Cauchy-Schwartz inequality and triangle inequality; (iii) uses the Cauchy-Schwartz inequality; (iv) reorganizes the terms in (iii). Plug Equation (82) into Equation (81), we can finally derive the bound for∥∇ W Q f (M t+1 ; X i ) -∇ W Q f (M t+1 ; X i )∥ F . ∥∇ W Q f (M t+1 ; X i ) -∇ W Q f (M t+1 ; X i )∥ F (i) ≤ n √ H ∥R i,t+1 -R i,t ∥ F + ∥R i,t ∥ F ∥S i,t+1 -S i,t ∥ F + R ∥S i,t+1 -S i,t ∥ (ii) ≤ n √ HP i ∥M t+1 -M t ∥ F + 2 Rn √ H∥S i,t+1 -S i,t ∥ F (iii) ≤ n √ HP i ∥M t+1 -M t ∥ F + 2 Rn √ H ϕ 2 i + ψ 2 i ∥M t+1 -M t ∥ F := L Q i ∥M t+1 -M t ∥ F ,where (i) is from Equation ( 81); (ii) uses the definition of R in Equation ( 80); (iii) comes from Lemma 3 (3). Since W Q and W K are symmetric in the Transormer structure, similarly, we can derive L K i .Step 2: In this step, we aim to derive bound for 1); (ii) uses triangle inequality; (iii) uses Cauchy-Schwartz inequality; (iv) comes from the definition of B i,t ,, Z i (in Equation ( 74)), Cauchy-Schwartz inequality and triangle inequality; (v) comes from Lemma 2 (3) and Cauchy-Schwartz inequality; (vi) reorganizes (v).∥∇ W v f (M t+1 ; X i ) -∇ W v f (M t ; X i )∥ F . ∥∇ W V f (M t+1 ; X i ) -∇ W V f (M t ; X i )∥ F (i) = B ⊤ i,t+1 (MH (M t+1 ; X i ) -y) W O ⊤ -B ⊤ i,t (MH (M t ; X i ) -y i ) W O ⊤ F (ii) ≤ B ⊤ i,t+1 (MH (M t+1 ; X i ) -y i ) W O ⊤ -B ⊤ i,t+1 (MH (M t ; X i ) -y i ) W O ⊤ F + B ⊤ i,t+1 (MH (M t ; X i ) -y i ) W O ⊤ -B ⊤ i,t (MH (M t ; X i ) -y i ) W O ⊤ F (iii) ≤ ∥B i,t+1 ∥ F ∥ ∥MH (M t+1 ; X i ) -MH (M t ; X i )∥ F ∥W O ∥ 2 + ∥B i,t+1 -B i,t ∥ F ∥MH (M t ; X i ) -y i ∥ F ∥W O ∥ 2 (iv) ≤ n √ H∥X i ∥ F ∥W O ∥ 2 Z i ∥M t+1 -M t ∥ F + ∥S i,t+1 -S i,t ∥ F ∥X i ∥ F ∥W O ∥ 2 (∥MH (M t+1 ; X i )∥ F + ∥y i ∥ 2 ) (v) ≤ ϕ 2 i + ψ 2 i ∥X i ∥ F W O 2 n √ Hσ max W V ∥X i ∥ F W O 2 + ∥y i ∥ 2 ∥M t+1 -M t ∥ F + n √ H W O 2 ∥X i ∥ F Z i ∥M t+1 -M t ∥ F (vi) ≤ ϕ 2 i + ψ 2 i ∥X i ∥ F W O 2 n √ Hσ max W V ∥X i ∥ F W O 2 + ∥y i ∥ 2 + n √ H W O 2 Z i ∥M t+1 -M t ∥ F := L V i ∥M t+1 -M t ∥ F where (i) is from Lemma 1 (Step 1 and Step 2, and plug into Equation ( 78), we can finally deriveEQUATIONEQUATION)1.7 Proof of Lemma in Section 1.4Proof. Proof of Lemma 6 (1): We consider the differential of the element in the k-th row and j-th column. First, let us write down the closed form of each element:(C ih ) kj = -∥X ik• W Q h -X ij• W K h ∥ 2 /2 √ dNext, we consider the differential of each element over W Q h :d (C ih ) kj = - 1 2 √ d X ik• W Q h + d(W Q h ) -X ij• W K h 2 -X ik• (W Q h ) -X ij• W K h 2 = - 1 √ d ⟨X ik• d(W Q h ), X ik• W Q h -X ij• W K h ⟩ + o d(W Q h ) ,where o(d(W Q h )) denotes the higher order of d(W Q h ). Leave out the higher order differential term, we derive∥d (C ih ) kj ∥ F ≤ 1 √ d ∥X ik• ∥ 2 ∥d(W Q h )∥ F • σ max (W Q h )∥X ik• ∥ 2 + ∥d(W Q h )∥ F • σ max (W K h )∥X ik• ∥ 2 ∥X ij• ∥ 2 ≤ 1 √ d ∥X ik• ∥ 2 ∥d(W Q h )∥ F (σ max (W Q h )∥X ik• ∥ 2 + σ max (W K h )∥X ij• ∥ 2 ) ≤ 1 √ d ∥X ik• ∥ 2 σ 2 max (W Q h ) + σ 2 max (W K h ) • ∥X ik• ∥ 2 2 + ∥X ij• ∥ 2 2 ∥d(W Q h )∥ F ∥d (C ih )∥ F = n k=1 n j=1 ∥d(C ih ) kj ∥ 2 F ≤ 1 √ d n k=1 n j=1 ∥X ik• ∥ 2 σ 2 max W Q h + σ 2 max W K h • ∥X ik• ∥ 2 2 + ∥X ij• ∥ 2 2 ∥d(W Q h )∥ F ≤ 1 √ d σ 2 max W Q h + σ 2 max W K h n k=1 ∥X ik• ∥ n∥X ik• ∥ 2 2 + n j=1 ∥X ij• ∥ 2 F ∥d(W Q h )∥ F ≤ 1 √ d σ 2 max W Q h + σ 2 max W K h • n k=1 ∥X ik• ∥ 2 F • n k=1 (n∥X ik• ∥ 2 2 + n j=1 ∥X ij• ∥ 2 F )∥d(W Q h )∥ F = 1 √ d σ 2 max W Q h + σ 2 max W K h • ∥X i ∥ F • √ 2n∥X i ∥ F ∥d(W Q h )∥ F = 2n d ∥X i ∥ 2 F σ 2 max W Q h + σ 2 max W K h ∥d(W Q h )∥ FProof. Proof of Lemma 7 (2): First, let us write down the closed form ofEQUATIONThus, we can derive upper bound for dEQUATIONThus, we have the following:d ∂ (C ih ) ∂W Q h F ≤ n k=1 n j=1 d ∂ (C ih ) kj ∂W Q h F ≤ ∥d(W Q h )∥ F n k=1 n j=1 ∥X ik• ∥ 2 2 ≤ n∥X i ∥ 2 F ∥d(W Q h )∥ F Proof. Proof of Lemma 7 (3): ∂f (M ; X i ) ∂C i F = (MH (M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ ⊙ S i F ≥ (MH (M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ F • min |S i | ≥ min |V ′ i W O | • min |S i | • ∥MH (M ; X i ) -y i ∥ 2 .Proof. Proof of Lemma 7 (4):∂f (M ; X i ) ∂W Q h F = vec ∂f (M ; X i ) ∂W Q h 2 = vec ∂f (M ; X i ) ∂C i • ∂C i ∂W Q h 2 ≤ ∂f (M ; X i ) ∂C i F • ∂C i ∂W Q h 2 = (MH (M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ ⊙ S i F • 2n d ∥X i ∥ 2 F σ 2 max W Q h + σ 2 max W K h ≤ 2n d ∥X i ∥ 3 F ∥W O ∥ 2 σ max (W V ) σ 2 max W Q h + σ 2 max W K h ∥MH (M ; X i ) -y i ∥ 2Proof. Proof of Lemma 8 (1): The proof is similar to the proof of Lemma 3 (1). So we do not include the details here. We can similarly deriveEQUATIONProof. Proof of Lemma 8 (2): By triangle inequality, we haveEQUATIONStep 1: Derive upper bound for∥∇ W Q f (M t+1 ; X i ))-∇ W Q f (M t ; X i ))∥ F = ∥ vec(∇ W Q f (M t+1 ; X i ))-vec(∇ W Q f (M t ; X i ))∥ 2 .First, we give the vectorized expression of∇ W Q f (M t ; X i ). Recall we denote U i = (MH (M ; X i ) -y i ) W O ⊤ (V ′ i ) ⊤ ⊙ S i .By Lemma 6, we can derive the close form ofvec(∇ W Q f (M t ; X i )): vec(∇ W Q f (M ; X i )) (i) = vec(U i ) • vec ∂C i ∂W Q h (89)Further, recall we have defined R and the following inequality holds:EQUATIONNext, let us derive upper bound for ∥∇EQUATIONNext, we aim to derive upper bound of ∥R i,t+1 -R i,t ∥ F in Equation ( 92). Similar to the derivation in Equation ( 81), we can deriveEQUATIONPlug Equation (82) into Equation (92), we can finally derive the bound forEQUATIONplug into Equation (78), we can finally derive • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean.EQUATION• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification: There is no societal impact Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 


Decision: All reviewers agreed this paper should be accepted: it addresses an important problem, the analysis is insightful and novel, and the writing is clear. The author response addressed reviewer concerns causing a reviewer to increase their score and confidence. A clear accept. Authors: you've already indicated that you will update the submission to respond to reviewer changes, if you could double check their comments for any recommendation you may have missed on accident that would be great! The paper will make a great contribution to the conference!

Reviews: review: paper_summary: paper_summary: This paper analyzes the convergence behavior of Transformer models with different attention mechanisms, specifically comparing Softmax and Gaussian kernel attention. The authors provide theoretical results on the conditions for global convergence and empirically demonstrate differences in optimization landscapes between the two attention types. The work aims to provide insights into why Transformers perform well and the potential advantages/disadvantages of Softmax attention.
strength_weakness: strength_weakness: - The empirical evaluation is limited and doesn't fully validate the theoretical claims in practical settings. The experiments use simplified Transformer models on relatively small datasets (IMDB and Pathfinder). It would be more convincing to see results on larger, more complex Transformer architectures and standard NLP benchmarks.
- The paper's contribution relative to prior work is not clearly articulated. While the authors cite some related papers, they don't adequately explain how their results extend or differ from existing analyses of Transformer convergence, e.g., [1].
- The practical implications of the theoretical results are not sufficiently discussed. It's unclear how the insights about Gaussian attention could be applied to improve real-world Transformer models or training.

[1] Y. Huang, Y. Cheng, and Y. Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
questions: questions: See weakness above.
limitations: limitations: The authors acknowledge some limitations in Section 6, noting that they require strong initialization and a large embedding size to obtain global convergence guarantees. They also mention the gap between their analysis and real-world scenarios. However, a more thorough discussion of limitations would strengthen the paper.
paper_summary: paper_summary: The authors establish different convergence theorems on the training of a single layer transformer with different trainable weight matrices and kernel functions. They prove that, under certain conditions, a one-layer Transformer with a Gaussian kernel converges faster than one with a Softmax kernel. Finally, they conduct empirical experiments to show the effectiveness of Gaussian kernel compared with Softmax kernel, validating their theory.
strength_weakness: strength_weakness: Minor points:
1.	It may be beneficial to include insights in the main text about the critical steps in the proof that result in the different convergence rates of the Gaussian and Softmax kernels in Theorem 3.

Major points:
1.	The experimental results in Figure 2 suggest that the Transformer with a Gaussian kernel converges faster. However, the difference in convergence seems more like a difference in the constant factor rather than a difference in the convergence rate.
questions: questions: 1.	The authors construct a compelling example to illustrate the failure of a Transformer with a Gaussian kernel, but it requires a specific initialization form. To achieve a similar convergence rate for the Softmax kernel, how many additional conditions are required? 
2.	I don’t fully understand why optimizing the query matrix alone leads to similar dynamics as when optimizing both the query and the key matrix. Could you clarify this?
limitations: limitations: The author has discussed some of the limitations in the paper. Specifically, the authors explain why they choose to study a simple one-layer transformer with regression loss while in practice, transformer models are usually trained with cross-entropy loss and with multiple layers.
paper_summary: paper_summary: This paper proves that with appropriate initialization GD can train a Transformer model with either Softmax or Gaussian kernel to achieve a global optimal solution. Besides, this paper highlights the Gaussian attention kernel exhibits much favorable behavior than Softmax in certain scenarios.
strength_weakness: strength_weakness: 1. It is unclear what new insights could be gained from the global convergence result Theorem 2, given that Wu et al [2024] already have one. The weight $W^O$ is trainable in Wu et al [2024], which is in line with practice, but $W^O$ is fixed in the present paper. 

2. The comparison between the Gaussian and Softmax kernels only updates $W^Q$.
questions: questions: 1. The conditions in Theorem 2 is different from those in Wu et al [2024]. How do these conditions compare with each other? Which one is weaker? What new insights can we obtain?

2. The authors emphasize that one important feature of the present work is that they delve in the roles of individual variables. Why is this important? Can the results obtained with one variable trainable and others fixed explain their roles when they are simultaneously updated? 

3. On line 286, is it implicitly assumed that the global optimum should have zero loss? Why must this be true? 

4. Lines 376-377 claim that the results in this paper explain why transformers perform exceptionally well. Can you clarify the explanation? How does it follow from the results of the paper?

5. There are several sentences like "only updating $W^V$ already leads to global convergence". Are they suggesting that it is easier to optimize more variables? How should we understand these sentences?
limitations: limitations: Not quite. See the weakness section.

Comments: title: Thank you and raise my confidence score from 2 to 3
Reply to Reviewer cxoC's questions
Thank you. Raise score from 5 to 6. More questions.
Response to Reviewer ynwg's question
comment: Thank you for your reply. Now, I understand this paper better, and my confidence score has increased from 2 to 3.
We thank the reviewer for the feedback and raise in score!

**Comment**: In Remark 1, it seems that the only assumption for data and initialization is that $B_0$ is full rank. Then, by over-parameterization, i.e. $D=\Omega(Nn)$, the loss landscape is nearly convex around initialization. This sounds very similar to NTK or a lazy learning regime. I wonder whether there are any feature learning or rich learning regime effects in your analysis.

**Response**: Yes, our optimization analysis is similar to the NTK regime with lazy update. Our current theory does not consider feature learning or the rich learning effect. However, we believe it is possible to extend our analysis towards the feature learning regime. For example, the high-dimensional embedding can be further modeled by a function that extracts data features, e.g, networks, allowing the attention head to be interpreted as computing the "correlation" of extracted features. We plan to investigate the feature learning regime in our future work. 

**Comment**: Is there any idea about the generalization with some assumptions on data distribution?

**Response**: Yes, generalization analysis can be derived from our theory with certain assumptions on data distribution and model weights, although it is not analyzed in this paper as we primarily focus on the optimization landscape. As a concrete example, it is possible to extend our analysis by following the framework in [1], where the data is balanced and a binary classification problem is considered. We can further derive a similar generalization bound as [1] Theorem 1.


**Comment:** The formula of Gaussian kernel.

**Response**: Respectfully, we would like to clarify that the formula of Gaussian kernel in your comment is missing an additional activation function. With different activations included, the Gaussian kernels (defined in [2, 3] and equation (11)) and Softmax kernels differ significantly, as the Softmax activation considers the entire row of the attention matrix, whereas the Gaussian does not. 
Specifically, the $k$-th row and $j$-th column of Gaussian attention is given by (see equation (11)):
$$S\left(W\_h^Q, W_h^K ; X_i\right)\_{k j}=\operatorname{\exp}\left(-\frac{1}{\sqrt{d}}\left(X\_{ik\cdot} W_h^Q-X\_{ij\cdot} W\_h^K\right)^2\right)$$.
The Softmax attention is given by:
$$S\_{ih}:=\text{Softmax}\left(\frac{X\_{i} W_h^Q\left(X\_{i} W\_h^K\right)^{\top}}{\sqrt{d}} \in \mathbb{R}^{n\times n}\right).$$
Notice that the entry in Gaussian kernel is only related to the $k$-th and $j$-th token in sample $X_i$. In contrast, in Softmax attention, each entry is also related to the entries in the same row, since the Softmax activation computes the regularization based on a row. This leads to a more complex optimization landscape and makes the convergence more challenging. These two different kernels also result in different performance on some tasks [2, 3].


[1] Li, Hongkang, et al. "A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity." arXiv preprint arXiv:2302.06015 (2023).

[2] Lu, Jiachen, et al. "Soft: Softmax-free transformer with linear complexity." Advances in Neural Information Processing Systems 34 (2021): 21297-21309.

[3] Chen, Yifan, et al. "Skyformer: Remodel self-attention with gaussian kernel and nystr" om method." Advances in Neural Information Processing Systems 34 (2021): 2122-2135.
Thank you for your rebuttal. It fixed some of my concerns and I raised my score from 5 to 6. 

On the other hand, I have some other further questions. 
- In Remark 1, it seems that the only assumption for data and initialization is that $B_0$ is full rank. Then, by over-parameterization, i.e.,  $D = \Theta (N n)$, the loss landscape is nearly convex around initialization. This sounds very similar to NTK or a lazy learning regime. I wonder whether there are any feature learning or rich learning regime effects in your analysis. 
- Is there any idea about the generalization with some assumptions on data distribution? 
- Why Gaussian kernel is $-\\|XW_Q - XW_K\\|_2^2$ rather than $(XW_Q W_K^T X^T)^2$? From my perspective, $-\\|XW_Q - XW_K\\|_2^2 \approx C  + 2 XW_Q W_K^T X^T$ with some constant $C$, which is not far different from the standard kernel.
Thanks for the clarification. I'll keep my score.
I thank the authors for the response! The authors have clarified the new insights and novel contributions compared with previous works. As a result, I am raising my score to 6.
First, let us clarify that global convergence in our theorems means **convergence to global optimal solution**. The definition is also used in other paper that analyzes the optimization of networks, e.g,[1]. And this definition is different from the traditional definition of "global convergence", which means that an algorithm converges (to certain solution) regardless of initialization. To avoid confusion, we will correct it to "convergence to global optimal solution" in our revised version.

Second, with the clarification above, we need to clarify that the condtion of initialization in near-convex condition **does not conflict** with our analysis of achieving global optimal solution in Theorem 1,2 and 3. Instead, it is an **essential condition** for deriving the result in our paper. This type of claim (local initialization to achieve global optimal solution) has been common [1,2,3], though in some works it is not pointed out directly.
 


[1] Nguyen, Quynh N., and Marco Mondelli. "Global convergence of deep networks with one wide layer followed by pyramidal topology." Advances in Neural Information Processing Systems 33 (2020): 11961-11972.

[2] Wu, Yongtao, et al. "On the convergence of encoder-only shallow transformers." Advances in Neural Information Processing Systems 36 (2024).

[3] Song, Bingqing, et al. "Fedavg converges to zero training loss linearly for overparameterized multi-layer neural networks." International Conference on Machine Learning. PMLR, 2023.
Thanks for the response. The response claims global convergence on the one hand, and initialization in a locally near convex region. How do you reconcile these two claims?