Unraveling the Gradient Descent Dynamics of Transformers
Authors: Bingqing Song
Boran Han
Shuai Zhang
Jie Ding
Mingyi Hong
Venue: NeurIPS 2024 poster (2024)
Keywords: Optimization
Transformer
Primary area: optimization_for_deep_networks

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.

Unraveling the Gradient Descent Dynamics of Transformers
Bingqing Song
Boran Han Amazon Web Services boranhan@amazon, com
University of Minnesota, Twin Cities song04090umm edu
Shuai Zhang Amazon Web Services
Jie Ding
University of Minnesota, Twin Cities dingj @ummn.ed
shuaizs@amazon.com com
Mingyi Hong
University of Minnesota, Twin Cities mhong@umn. edu
Abstract
While the Transformer architecture has achieved remarkable success across various domains, 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 Transformer architectures allow Gradient Descent (GD) to achieve guaranteed convergence? and (2) Under what initial conditions and architectural specifies does the Transformer achieve rapid convergence during training? By analyzing the loss landscape of 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 Transformer model (with either kernel type) to achieve 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 much favorable behavior. Our empirical study further validate the theoretical findings.
Introduction
Transformer model architectures have become popular in machine learning, delivering remarkable performance across wide array of tasks. From natural language processing (Vaswani et al. 2017
standards in performance and efficiency. Popular models include BERT [Devlin et al. 2018), RoBERTa Liu
BBeltgy et al. 2020) computer vision (Dosovitskiy et al. 22020) these models have
al. 2019]. DeBERTa He al. 2020]. GPT models (Radford al. 2019 Brown al. 2020 and
ViT Dosovitskiy et al. 2020]. Despite their empirical success, a comprehensive understanding of
their optimization process remains clusive. As highlighted in [iu 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 work of B. Song was partially done while interning at Amazon Web Services.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

of the associated optimization problems, has made the theoretical analysis of these models very challenging.
The optimization landscape can be pivotal for understanding a certain type of neural network and providing the practical guidance [Liu et al., 2020]. Existing literature offers numerous studies on achieving zero-loss solutions in networks with ReLU activation. These studies encompass various network structures, including fully-connected, convolutional, and residual networks, as explored in (Jinn et al. 2017]. Jiin et al., 2021], and (Danilova et al., 2022]. They delve into the analysis of network optimization landscapes and provide assurances of rapid global convergence when using gradient descent (GD) or stochastic gradient descent (SGD) algorithms. For instance, in Du et al. [2019), the authors focus on fully-connected networks and ResNets with smooth activation functions, and they have demonstrated that global convergence can be achieved using GD with a network size proportional to (ppoly(N)) where N is the sample size. Similarly, [Allen-Zhu et al. 2019]
show that ReLU fully-connected networks with at least (ppoly(N)) (N)) neurons can achieve global convergence using GD or SGD. From statistical perspective, (Li et al., [203] have shown that for two-layer ReLU neural networks (with input dimension p) that admit a sparse subnetwork representation, a sample size of O(log" (p/o)) can guarantee the global convergence with probability at least ó using GD. Despite this extensive body of work traditional architectures, it is not clear what conditions we need (e.g. network size, optimizer, initialization) to ensure training Transformer models to find high-quality solutions.
Compared traditional deep learning architectures, Transformers incorporate a unique level of intricacy through their attention kemel [Vaswani et al.) 2017]. which is designed to effectively handle sequence inputs. This mechanism incorporates Softmax activation to the inner products of query and key vectors, and this inherently non-convex operation poses considerable challenges to theoretical analysis. Consequently, existing frameworks for analyzing the convergence of classical deep learning models are not directly applicable to Transformers. Further, many recent works have pointed out that the performance of Transformers depends on a number of factors such as the choice of kernel function, initialization, choice of optimizers, and forms of token embeddings (Huang al. 2020 Pan
and Li 2023 Shazeer 2020 Li et al. 2018 Tian et al. 2023]. In deep learning, these factors have been studied in line works. For example, Li et al. 2018 show that the good training performance is not universal: skip connections have the effect of smoothing the training landscape, and the
Adam algorithm tends follow more direct trajectory towards optimal solutions compared to SGD,
Therefore, it is imperative to understand what kind of conditions, including initialization, network structure, data properties, and optimizer choices, will lead to high-performing Transformers.
In this work, we will delve into the intricacies of attention kernels, discussing both their advantages and limitations in the context of model optimization. The main contributions of this work are threefold.
We derive the conditions that will make the one-layer Softmax attention Transformer reach global optimality with vanilla gradient descent. The convergence guarantee is largely attributed to the linear layer (W in the attention mechanism.
We investigate the attention kernel's effectiveness, revealing Gaussian attention achieves zero training loss, while Softmax can lead to non-optimal stationary points.
Our experiments validate that Softmax attention Transformers converge slower and present more challenging training landscapes than Gaussian counterparts, potentially leading to more local optimal solutions.
2 Related Work
A number of research works have focused on the theoretical analysis and interpretation of Transformerr models, revealing crucial insights into their practical performance.
Liu et al. 2020) showed that heavy reliance on the residual branch in multi- multi-layer Transformer models can lead to training instability, which amplifies small parameter perturbations, causing significant disturbances in the model's output. In Bhojanapalli et al. [2020), the authors illustrated the existence of a low-rank bottleneck in Transformer models with sufficiently large embedding and hidden size
(D = d). However, this work focuses on the representation ability of large size attention, while
falling short of analyzing Transformer models from optimization perspective. In Noci et al. 2022

---

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 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 behavior of shallow Transformer, which builds a convergence theory shallow Transformer with realistic structure and initialization, but they do not provide roles for different matrices in the convergence. However, the focus of our paper 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 foeus 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 taassfoomerrbbssee models within the in-context learning (ICL.) framework. For instance, [Huang et al. 2023 Zhang et al. (023) 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. More specifically, these works analyze the convergence of in-context training, where a prompt is constructed with all the training samples and single test sample. The goal of these works is achieve the zero test loss (in expectation) by optimizing over the loss function modeled by the prompt. On the other hand, our analysis is based on standard empirical loss minimization, which does not involve any prompt construction.
3 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,, W where X, E RTxD and y; E I" 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:
(eention(WR W W.: X,) = S(W?. W X.)XWWX
(1)
MH(WR, w WV; X,) = Concat (head heada) w. where head, = Attention R h W wX: X.), h = 1,..., H.
(2)
In the above notation, WºQ W.K E Rxd is the query weight matrix and key weight matrix, respectively; WY E IR Dxd is the value weight matrix; these matrices are the main optimization variables throughout the paper. Further WOO E R Háx1 1 is a fixed matrix, representing the weight of the output layer: is the number of attention heads; S(-) is a kernel function of variables W',, W K and input X, 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:
h.
(X.WX) X.W K
s (W W x.) = Softmax
(3)
Va
where for a given n X n matrix Z, Sofmax Z)) == [Softmax(Z1); ... Softmax( Zn.) Throughout, let us denote 5()), as the element of k-th row and j-th column in matrix S(.). Let X RD denote the embedding of the k-th token in data X,, which is the k-th row of matrix X₁. The structure of Transformer model can be found in Fig
where we denote S., := S (W W X.).
3

---

Based on the above Transformer model, we consider minimizing the following empirical & loss function for the entire data (X,, yi yili=1
N
minM Σ [MH(M;;;) -vill",
where M == (WQ,WK w WWV
is the set of variables that can be optimized.
For notation simplicity, next we define the vector version of the Transformer model given in Equa- tion I for the entire dataset {(X,, H Towards this end, let X E IR Nx denote the column- stacked matrix of each single data X;. Similarly, define the stacked label y E RNV, Then we can define:
S
SIHXI SNHXN/
MH(M; ==
diag(WY, ,WH w
(5)
SNIXN
i == 1,2, .N, 1,2,
H for simplicity.
S11
S1HX1
Let B
and WV
SNIXN diag(Wy Bo. .... WH) E RD denote the diagonalized weight matrices that include all value weight matrices for all attention heads. Using these definitions, We can simplify Equation (5) as
XNHXN
MLP Layers
WO
MH(M: x) = B W. .
Sin
X, X,WY
Thus the empirical loss function given in Equation can be simplified as
X,W h
X X,W h
miny |MH(M;X) - -
(6)
For more notations in the following sections, will use sub- script to represent the variables in f-th iteration, c.g, M, := {W? : W WV). Similarly, we denote B, as the matrix B at 1-th iteration.
X;
Figure/ 1: One head in Trans- former architecture with Soft- max Attention.
important to note that, in the above description and throughout the paper, 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 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 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 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., 22023. Further, the analysis of MLP is standard in literature (Allen-Zhu et al. 2019 Du et al. 2019 Nguyen and Mondelli 2020). And is worth noting that our choice to focus on lne-l Transformer is consistent with other works that similarly aim to investigate the core training dynamics of Transformers, e.g, in (Tian et al. 2023), layer Transformer is considered basic model.

---

4 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 C (WQ, W WV) ). Below we summarize our results.
Common convergence conditions with Softmax Attention: When the activation function S(-) either the Softmax or Gaussian function, and the embedding dimension D is at least O (NN). optimizing Equation (6) can achieve global optimal solution when M (WV) and M W (WR, w J.
Different behavior between Softmax and Gaussian Kernel Attention. When S(-) is Gaussian and the embedding dimension D is at least O(Nn), convergence to global optimal also ensured for M = { Interestingly, under the same conditions of large D, convergence to global optimal is not guaranteed when 5(-) is Softmax.
In the subsequent sections. we will elaborate on these convergence results in detail, providing deeper understanding of the nuances in Transformer behavior under varying configurations. To set up our analysis, we introduce A as the smallest eigenvalue of W A as the smallest eigenvalue of Bo, X Λ X as the largest singular value of matrix Wh K W respectively We denote || ll2 as l norm and || || as Frobenius norm. Further, we denote o max (-) and min( as the largest and smallest singular value of matrix, respectively. For any vector v, let min((u]) denote the smallest absolute value of vector v.
4.1 Convergence global optimal
First, we examine the role of WV in the optimization of multi-head attention network structure. Our analysis demonstrates that with the hidden dimension HD 2 Nn and proper initialization, the global optimal solution of (6) can be found using vanilla gradient descent algorithm. The initialization requires that the matrix Bg has full rank. Our first result shows that, overparameterized Transformer can be trained global optimal solution.
Theorem 1. Consider problem { with S(-) being instantiated as the Softmax kernel given in (3) Consider the following update for the variable M = }: WX, where n) > 0 is the stepsize.
= W.Y - V f(M; X).
Suppose Wo and W are initialized such that A > O. Then we have:
f (M,; X) s (1 - noo) (Mo; X),
where a: ||00 0: 1 > 0 is defined in Appendix 1.3, and chosen such as na 1. Remark 1. The aforementioned theorem focuses on the convergence behavior when only W is being updated. We further elaborate on the initial conditions ensuring A 00
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 Bo has full rank, which can be fulfilled by selecting appropriate W and W plus having large enough embedding size, satisfying D 2 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, 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 = 1, Hl 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 2 n x (N/H) to guarantee the global convergence when Transformer model is trained with GD.
In Theorem 1 we have illustrated the case where only updating WV already leads to global con- vergence. However, in practice, all parameters WV W W are updated. This case is more challenging to analyze due to the non-linearity introduced by the Softmax function. Next, we show that similar result in Theorem still holds when all the parameters are updated simultancously.
5

---

Theorem 2. Consider problem 4 with S(.) being instantiated as the Softmax kernel. Consider the GD update where M (WQ, WK WVY Suppose A= > O, and the initialization Mo satisfy
H n² D'VNH XII Σ hm1
(A + X
||MH(Mo;; -112 S
(8)
|Woll22 )2 min (2 A/ 29)
Then there exists stepsize oj > such that
f ()
(Mo:
(9)
where B = -|W012229)>>> O, 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 WQ, W WV only. To illustrate this, consider the following Transformer model:
G a W WV; x.) -MH(WO, w WY, X.) w WL
(10)
where W R XMu, and no = do Based on the Transformer model defined Equation we have the following corollary.
10
Corollary 1. Consider problem min 1 |IIG(M; y1l². with G(-) being defined in Equation 10 and 5()) being instantiated as the Softmax kernel. Suppose that the MLP module satisfies:
11 2 n2 2 TLL.
Consider the following GD update (where M = (WQ, W WV W "..,,,, Suppose X O. Then, there exists step size 17 >> O and initialization weight Mo. such that the loss function Tinearly converges to O.
Remark 3. The above theorem and corollary describe the global convergence guarantee when WQ, W K and W. are updated. This in line with the insights gained from Theorem | 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 WQ and W
To ensure the initial condition 8 we have two options: 1) Initializing My such that [MH(Mo; X):
glle small, which implies that the optimization starts in a region close the global optimal solution and that the initial weight is close to the global optimal solution; 2) Balancing between W and WV, in the sense that ||WO|| [|2 is large and X is small. For 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 include the theoretical statement since it is similar the result for Softmax attention.
4.2 Softmax vs Gaussian kernel: Softmax attention Transformers may exhibit slower convergence.
In the previous section, explored the global convergence of training Transformer models. However, from Theorem was not clear what roles do matrices WQQ and W play in the entire convergence process, since Theorem indicates that optimizing W alone already ensures the desired convergence. Nevertheless, is the matrices W and Wº that truly represent the power of a Transformer model, because they used extract token correlations.
To study how well a Transformer model can extract token correlation, in this section, we will study the GD dynamics for Transformer models, where only W and WQ are optimized (while fixing WV), 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.

---

4.2.1 Notations
To begin our study, let us define that Gaussian kernel to be an n2 x 7u matrix, where its k-thh row and j-th column of is given by:
(WR W xx:) exp
(-v/2
Va (X..W - -x,,,
(11)
Since the training dynamics//radients of variables W and W have the same property in (1I. we will only concentrate on optimizing WQ
3 and
With some abuse of notation, define matrix C for Softmax attention and Gaussian kernel attention, respectively. Softmax attention: CsA =
X; X.W (x, (x,
vd
Xi.. W X. W
Gaussian kernel attention: Cia E Rux", (Cih)ay For both Softmax attention and Gaussian kernel attention:
2vd
C I R [Ca, Ci., Cuul; C [CT C2 ...... .N.
Using the above notation, the activation function S(-) in 3 and I1 can be related the matrices C's in the following manner:
Softmax attention Shh = Softmax (Can), Gaussian attention (Sin)a = exp ((Cin)
Additionally, note that is a function of variables M. Therefore we will sometimes use ((M) when we need to emphasize the dependency of C on M.
4.2.2 Main Results
Next, we will outline the conditions under which GD can still successfully find global optimal solutions for Transformers with Gaussian kernel attention (when only WQ 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 (with M = =WW 1h = @C(Mo)
W.e w J(Me; X). Suppose 8, == Tmin
O, A h E [1,2,... HI and the
ow
initialization condition further satisfies
n/ xI (5% + xx) exp (2 12 ((X9)2 (xK)3) x X 10 |MH (M;; - =ll2 S v/
(min(V -minn (6n, (0)
(12)
V is defined in Appendix 1.5.
(1) When S(-) is a Gaussian kernel function, there exists a stepsize nj and a positive constant 7, such that
I [M,; X) s (7-) (M;;
(13)
where Y, n are defined in Appendix 1.5.
(2) When S(.) is a functiin suppose W is bounded during the training phuse, then there exists stepsize T), such that
f ((M;X) S (Moo
Σ Wo (M;X)]I
(14)
where of is defined in Appendix 1.5.
Remark 4. First, it's important to note that the parameter size must satisfy Dd > Nn² for ò > 0 to hold. It is crucial 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 2 Nn² is similar to the findings of works that have analyzed the convergence performance of aver-parameterized neural networks Allen-Zhu et al.) (2019). Du et al. [2019]. The total number of samples, consisting of N samples each with tokens, can be calculated as N n. Meanwhile, the total feature dimension is Dd. The inequality implies that the width of the
7

---

parameters is at least (NN, a relationship also illustrated in Nguyen and Mondelli [2020]. The proof consists of two basic steps. The first step is to derive the closed gradient of loss function over variable WQ, Intuitively, the gradient of Softmax attention is much more complicated than the Gaussian attention Transformer, which will lead to a more complicated landscape and more local solutions. The second step is to analyze the gradient of Transformers with both kernels and the same initialization Equation (12) For Gaussian attention Transformer, it can be iteratively shown during the gradient descent training: 1) The variable W is bounded; 2) The PL condition holds (i.e, the optimization landscape remains near-convex); 3) The loss function decreases linearly. For the Softmax attention Transformer, there is guarantee that the PL condition holds during iterative gradient descent update.
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 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, 2, H 1. Let us first write down the close form of the gradient over WP:
O (Ma; ow
1 Vd Xi Os X1) Xi WIK
a011
X1, W 1,: W such that the loss function is non-zero with Next, we show there exists W0 WV Equation [12) satisfied, while
(Mo: X1) E R2x2
80u1
Denote I. ==
(WO) (X
E R222
aS(Mo:X.)
can be expressed as follows:/
3MH(M.:X.)
80n
Of (Mo; X1)
[af (Mo; X1)
(M):
= (L11 - L12),
= 8. (L12 - Ln1), ō is some constant.
00n1
8011
11
12
Next, we will give the value of W Wy show the case where GD leads to vanishing gradient.
Let D 2, w ( f), X1
[
1) and Wy
0)
[2a a
where a is constant. is
0
2a
easy to show that there exists W.O and W such that Equation that for this scenario, the following holds:
12 holds. Further, it is easy to verify
Lu = L12, L21
L2z.
(15)
Next, we can easily deduce that (
[M((M))
= 0. Similarly, we can demon-
8011
11
8011
12
strate that (DfiMoix.)) 8C11 1)(
[M((MoXX])
= 0. Consequently, we have ([(Mo = 0. However,
aw?
if y satisfies that JMH(Mg:X.)
+ O, follows (MM X1) * O, which means Mo 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 at- tention and Gaussian kernel attention across various tasks.
Acc Gaussian vs Softmax
Test Loss Gaussian Softmax
5.1 Dataset
Ceusstar
We investigate two distinct tasks: Text Classification using the IMDb review dataset Maas et al. and Pathfinder [Linsley et al. While both tasks involve processing long sequences, they exhibit different characteristics. Text Classification is well-known NLP task that focuses on discerning relationships among token embeddings, while the Pathfinder task prioritizes capturing spatial information within the input pixels.
2011
Figure 2: Test performance on text classification task with different attention kernels
2018
8

---

5.2 Model and Experiment Method
We follow the experiment setting in [Chen et al. 202T]. For both tasks, we employ a 2-layer Transformer model with the following specifications: em- bedding dimension D = 64, hidden dimension d 128, and number of attention heads H = 2. To align the model with the classification task, use an additional mean pooling layer as the final layer. We determine the batch size based on available mem- ory constraints. Specifically, we batch size of 16 for the Text Classifi-
Test Softmex
Test Loss Geussian Softmax
Gemrn
Colavi
Figure 3: Test performance on pathfinder task with different attention kernels
cation task with a learning rate of 1 x 10-4, and batch size of 128 for the Pathfinder task with learning rate of 2 x 10-4. 4, For optimization, we use Stochastic Gradient Descent (SGD) for the Text Classification task and Adam for the Pathfinder task. We conduct two types of experiments.
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.
Yot Suinee
farthoor Sulteel
huriieer Goumirr
40
6
40
20
to
Figure 4: The loss landscapes on text classification task and Pathfinder task. For both tasks, we use the two-stage training in Section 5.2 with the same training hyperparameters, while the only difference is the attention structure in the second training stage. The two axes represent the two directions dj and d2 as defined in Section 5.2
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 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 matrices in the first Transformer layer. These two directions, denoted as di, d2, are centered at the trained model M, and represent the parameter space of Wº. W respectively. We evaluate the loss function on the set (M + 0.02(r 25)d, 0.02(s }. where r, E 2, 50). The above set is the neighborhood of the trained model M, and we chose the evaluation stepsize 0.02 along the two directions d1, d2, with the total steps limit as 100. Within this parameter space, we plot 3-D surface representing the landscape around the trained model.
5.3 Results
5.3.1
Test Loss & Accuracy Curve comparison
To begin with, we present some observations our first experiment. We plot the test performance of these two tasks on Transformers with two different types of attention. From Fig and 3 we can conclude that in both tasks, Transformers with Gaussian kernel attention exhibit faster convergence and higher test accuracy than Softmax attention with the same model size and learning rate. Especially,

---

training Transformers with Softmax attention in the Pathfinder task can lead to unstable performance as indicated Fig 3 The test accuracy has 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. 202T] and (Tay et al., [020]. where Transformers with different attention kernels are trained with the same model size and learning rate, while Softmax attention Transformers show instability in few tasks.
5.3.2 Optimization Landscape Comparison
In Figure we present comparison of the optimization landscape between Transformers with Softmax and Gaussian kernel attention. Notably, we observe distinct differences in the training landscapes 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 direct representation of the difference in the attention structure during the optimization procedure. With Softmax attention, complicated compared with Gaussian kernel attention. This complexity can be interpreted as the presence of 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 Baaureeandd and Figure 3 where Softmax attention exhibited certain convergence issues. These observations also provide empirical evidence supporting our Theorem 3 which reflects in slightly different perspective the complicated optimization landscape within the Softmax kernel.
the landscape appears
6 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 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 overparameterizztton,, 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 Z Nn to obtain the global convergence, which exhibits gap towards real case. In the future work, we will investigate how to relax the assumptions.
Acknowledgment
The work of B. Song was partially done while interning Amazon Web Services. M. Hong holds concurrent appointments as an Amazon Scholar and as a faculty at the University Minnesota. This paper describes their work performed at Amazon. The work of Jie Ding was supported in part by the Army Research Office Early Career Program Award under grant number 191IN12310315.
References
Z. Allen-Zhu, Y. Li, and Song. convergence theory for deep learning via over-parameteriizaion.. In International conference machine learning, pages 242-252. PMLR, 2019.
I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long docu-entt transformer. arXiv preprint 005500 2020.
Bhojanapalli, C. Yun, S. Rawat, Reddi, and S. Kumar. Low-rank bottleneck in multi-head attention models. In International conference on machine learning, pages 863- PMLR, 2020.
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances neural information processing systems, 33:1877-1901, 2020.
10

---

Y. Chen, Q. Zeng, H. Ji, and Y. Yang. Skyformer: Remodel self-attention with gaussian kernel and nystr om method. Advances in Neural Information Processing Systems, 0-212222113, 2021.
M. Danilova, P. Dvurechensky, A. Gasnikov, E. Gorbunov, Guminov, D. Kamzolov, and Shibaev. Recent theoretical advances in non-conver optimization. In High-Dimensional Optimization and Probability: With a View Towards Data Science, pages 79-163. Springer, 2022.
Devlin, M.-W. Chang. K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 2018.
A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pages 1675-1685. PMLR, 2019.
He, X. Liu, J. Gao, and W. Chen. Deberta: Decoding-enhaneee bert with disentangled attention. arXiv preprint arXiv:2006.03664 2020.
X. S. Huang, F. Perez, J. Ba, and M. Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learning, pages 4483. PMLR, 2020.
Y. Huang, Y. Cheng, and Y. Liang. In-context convergence of transformers. arXiv preprint arXiv:231000524,, 2023.
P. Jain, P. Kar, et al. Non-convex optimization for machine learning. Foundations and Trends in Machine Learning, 10(3-4):142-366 2017.
C. Jin, Netrapalli, R. Ge. S. M. Kakade, and M. I. Jordan. On nonconvex optimization for machine learning: Gradients, stochasticity, and saddle points. Journal of the ACM (JACM), 68(2):1-29, 2021.
Langley, Crafting papers on machine learning. In P. Langley, editor, Proceedings of the 17th International Conference on Machine Learning (ICMI. 2000), pages 1207-1216, Stanford, CA, 2000. Morgan Kaufmann.
G. Li, G. Wang, and Ding. Provable identifiability of two-layer relu neural networks via lasso regularization. IEEE Transactions Information Theory, 2023.
H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
D. Linsley, J. Kim, V. Veerabadran, C. Windolf, and T. Serre. Learning long-range spatial dependen- cies with horizontal gated recurrent units. Advances in neural information processing systems, 31, 2018
L. Liu, X. Liu, Gao, W. Chen, and J. Han. Understanding the difficulty of training transformers. arXiv preprini 20Xi::2004408249, 2020.
Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoy- anov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 16692. 2019.
A. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. 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, 2011.
Q N. Nguyen and M. Mondelli. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:1 1961- 2020.
L. Noci, S. Anagnostidis. I.. Biggio, A. Orvieto, S. P. Singh, and A. Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 15:27198-27211 2022.

---

Y. Pan and Y. Li. Toward understanding why adam converges faster than sgd for transformers, arXiv preprint XrXiv:2306.00 00204, 2023.
A. Radford, J. Wu, R. Child, D. Luan, Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8)99 2019.
N. Shazeer. Glu variants improve transformer. arXiv preprint arXi::2002.052022, 2020.
Y. Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L.. Yang. S. Ruder, and D. Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint arXiv:201 1.04066, 2020.
Y. Tian, Y. Wang. Chen, and S. Du Scan and snap: Understanding training dynamics and token composition in I-layer transformer. arXiv preprint arXiv:2305. 16380, 2023.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
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.09990 2023.
12

---

Appendix
1.1 Notations
Recall that we have defined the structure of a single Transformer model in Equation 1 and Equa- tion (5) We will further define a few notations before introduce few useful lemmas that needed in our proof.
(1) Operator: Denote vec(.) the vectorization operator matrix; & Kronecker product operator; O as the element product. Denote T(-) as a matrix operator, such that for any matrix X without zero element
1/211
1/21n
T (Xmxa)=
(16)
1/uml
1/rmn
mxn (2) Matrix: Denote as the identity matrix. Define matrix E and E as following:
E
E
E
E
HnxHn
nXn
Define matrix PA as following: Ir, = (..., Eh Paxm:...) h = 1.... H. (3) Matrix Transformer: Define the following matrix related to the attention layer Softmax kernel:
Ca X.Wo (X.WR)
Sth Softmax(Caa)
(17)
Vd
Gaussian kernel:
X W Xij WK
(Can)ss
(Sen)s exp ((Cr))
(18)
2vd
C = =C C.Hl, S.= [St, SSHl SS
(19)
Define matrix V/ for each data X₁:
X,WY
V
V V V........ .VI..
(20)
X.WY Hoxd
Next, let us introduce several useful lemma which leads to Theorem 2
1.2 Lemmas of Theorem 2
Lemma 1.
(1) 0f(M:X) BT (MH(M;X) - (WO)
(21)
OWV
(Of(M:X) ((WO)T vec(M(M((MMM))) X) n))
(2) vec
OWV
(Ina & B (wo & IN). (MH(M; - y)
(22)
(3) 8f(M:X)
TP 0f(M;X)
N
1 Of(M:X,)
XW = Σ
x Ph
X,W
(23)
aw Vd"
0C
Vd
0c,
[=1
0f(M: X,) = (IMH(M;X.) - vi) (WO) (V1)) o S
(4)
(24)
OC:
(((MM((M;X))- - -n) (WO) (V1)
3 S 3 T( C.) E BT) o exp C;
(25)
Df(M:X) diag
(Of(M;X))
Of(M;
(5) OC
(26)
00,
dC
13

---

Remark:
The above lemma derives the closed form of the gradient of objective over W W Notice that we can derive the derivative of W in the same way as W due symmetry, so do not include the derivation here. Some of the lemmas here refers 09-07/attention-ee gradient
Lemma 2. Consider updating Wº, W WV at iteration L. Suppose max max (W4), max (WK)
Fmax V) are bounded during in the optimization phase, then we have the following conclusion:
n (1) so where Φ; vd
|
d max
(WA)
(27)
(2) $ vslld )l., where v/. = vd
|
Σ o max
(w²)
(28)
(3) [[d(S))]l= S of + oz ((d(W)), d(W*
(29)
(4) (S(M:X.) ||5 S Q.IIMH (M;
ville,
OWQ
" where Qi = Hn X 1 3 ||
Σ h=1
2 (W max (WV).
(30)
XI) OWK
(5) 11
|| S (M; willF,
H Σ of (w²) mmax (W
where K₁ = = H I [w
(31)
Lemma 3. Consider updating W W WV at iteration t. Suppose maax (WQ), m (W*), Fmax (WV) are bounded during in the optimization phase, then we have the following conclusion:
(1) [IMH(M(+1:X)) X) MH(M;; X))le & M|| - Melle, where Z is some positive constant.
(32)
(2) I (M+11: - V (M: 1)11 G|| Molle, where G is some positive constant.
(33)
Lemma 4. Let f R + RR. be second order differentiable Let xr, y E R. be given, and assume that [V) fxx))l2 S C||2 for every x t(y x) with E [0, 1). Then,
C' || y
f(v) < f(x) + (VI(x), - x) +
2
Lemma S. For matrix A E Rixl B E R lxm C E Ruxn
vec(ABC) (Tn & AB) veo(C) (CT BT s Tx) vec(A)
(Tm. @ A) vec B) (BT s It) vec(A)
vec(A A B) vec A) 2 vec(B).
1.3 Proof of Theorem 2
Proof Sketch of Theorem 2:
The main idea of the proof follows from
Nguyen and Mondelli 2020]. Let first recall a few notations, X 3 (1 + maax (W))) )), X Omiu (By). Using GD update rule, we aim to iteratively show
max max(WV)
3 TV r {0, (0.....t.. t)
maxx (W?)
1 r {0,
.o.....t]. t
amax (W:)) < 3T r (0,
(),
(34)
mmin (B,)
11 ,re {0, (0....,th.
f (M;; X)
I (1 mH)" (Mo. X), r ...
14

---

Denote P = 1 Let us discuss about the value of We know wo E R R By RHDxNN We require μ O, i.e, A > O, which implies B has full row rank. For simplicity, let us consider the H = case. Recall the definition of B:
SI1
B
SNIXN
Suppose we initialize W w such that each Sn1 E F is full rank, then we can easily show that (an
X; has full row rank. Suppose embedding dimension is large, with certain assumption on X. we can show B has full row rank. For example, if each X, follows standard Gaussian distribution with D >> N, then rank(B) Nn with probability if we initialize W W K such that S21 is full rank.
Further, let us assume that Σ (2 > 1, Σ (x
1, and initialization condition satisfies:
5 X Σ -
h=1
S
(35)
A 1, (A >>22|/o0ll min (X)
hh
Remark 5. The initialization condition can be satisfied if W ||2 is large and a (WV) is small. Equation 8 54'
is clear that Equation iteration t + 1.
holds when Suppose holds at iteration t, we prove holds
| W. Ir
S
WV r+ (463) B|IBAII=| 2 Σ (1-1p)) |MH(M;; X) - -1l2.
W = Σ Vwv of (M:: X)Ile
(1)) 7
BxIle||Wo [MH(M,; -lll
where (i) uses the triangle inequality; (ii) plugs in the expression of Vwy f (M;X) and uses the Cauchy-Schwand inequality; (iii) because we assume the loss function f'(-) linearly decreases until I-th iteration. Let VI m/.. So we have
n 1 F
W
Σ (1
12)0/2 JIMH (Mo; X) 2112
2
1
u-+1
$
|BBl
(1 u²2) [|MH(Mo; - 8112
1 u
1 II
1- ++ (1 u²) ||HH(M;: -
(36)
HN μ
u
2n
(1))
|X]|i|/l MH(M;; X) y|l I 1,
(37)
where (i) is because each element in Sr has magnitude at most 1 and || < nv H, then by Cuachy-Schwartz inequality, we have ||B||, < VHN [[; (ii) is due to the initialization condition. Then by Weyl's 's inequality, there is
T max (WV)) (W+1) S max (W,)) + TV.
15

---

Similarly, let us derive the upper bound for mmax (W
h.r
W h,++1
Σ
=1
V VwwJ (M;; [
=
(66)
((0)
|H
17
Σ
Q
(1 m/)
MH (Aoo X)
-112
i=1
N
1 - u r+1
(1 u1) [IMH (Mo; - v112
1-u u
2
Q
IM (Mo; X) vll2 ((v)
12 A/
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 show
nn)a/ |M (Mo: 5ll2
i=1
2
Σ
K?
1 MH (Mo: - 5/12 S
2 1/
(38)
Then by Weyl's inequality, there is
Tmax ( (1
S maxx (W
1
00-
3 3.0 2
maax hee+))) S mmax (W h.o)
1 TK
3 A
2
2 h
Now we aim to bound the eigenvalues of Br+1.
[B,41 Boll r
Bs+1 Bzll ΣE
|S - SexXielF
i=
(1) Shallr
N
H
S
|
Σ 4=1
|| Sst
==0
(0z)
N
| F- V V + 0g f(Me: Vwx f(Ms; llf
Xelle V0 +/2 0 + +3 (M,, - ville
(o + (Q) + (1)(111))*12 n/)/// |M (Mo: - -yl1
where (i) and (ii) uses triangle inequality and Cauuch-SShwaad inequality; (iii) comes from Lemma 2 (5); (iv) uses Lemma and Cauchy-Schie Schwartz inequality; (v) comes from Sauchy-Schwarr inequality. Together with our initialization condition, we have
|I + +P) + K7)- 1 - u+1
|B,11 Bulls S
MH (Mo;
v112
u
(4) |x,1I + (Q) + +3) ||H (Mo; x)-xx)-- -vl1
<
S
B
2
16

---

where (i) comes from the initialization condition 35 By Weyl's inequality, we can derive the bound for the singular values of By:
0 min Br++) 2 Jmin (Bo) - |B,+1 Bull F 2 10.
The final step is to show the last inequality holds. Since we have already showed m (W,)) O mmax
O (W ) mmax ( are bounded, by Lemma (2) we can conclude that:
|V/ (M+1;; - (S(M)ll S G - MiIlr
Thus by Lemma 4 we choose
<
20 then the following hold true:
f (M++1:X) = (M 7/ (MM;;)
(9)
G " IV (M; XMI2 - (M, X) 7 IVV (M: (M;X)I|++
(9)
$ *M,; X) 2" /VV (M;
(08)
05 (M; x)
S / (M, 2"
OWV
3 f (M,x)
(2)
1
W @ B, (vec(MH((M; X) v)) 1|?
2
(9)
$ (M,x)
8
1
1 f(MciX)
(1 1 13(A2)))))) ((M;X)
(vd) (1 -1-mi [
1 (iii) only considers the gradient over WV; where (i) uses Lemma (ii) is because we set < 2G (iv) plugs in the closed form gradient in Lemmal) [ (V) uses the property of smallest singular value and induction assumption; (vi) comes from the definition of
1.4
Lemma for Theorem 3
The following lemmas all consider the Transformers with Gaussian kernel attention 11
Lemma 6.
(1) as(M:X)
OWV = BT BB (MH(M;X))- - () (WO)
(39)
(fflM:X) = (T @ (eN - 4))
(2) vec
OWV
= =(Laa @ BT) (wo @ IN) (MH(M;X) -
(40)
(3) 0S(M:X) 0s(M:X) OC
Ow OC aw = as(M:X.)
OC;
(41)
0C:
aw h
(4) 0s(M:X.) ((MH(M)X.) wi) (WO) (v:) o S
(42)
OC;
Of(M:X)
[OS(M;X])
0f(M;XN))
(5) 0C
(43)
aC1
8CN
17

---

Lemma 7. Consider updating W w WV at iteration t. Suppose o m (WQ), mma (W*) max (WV) are bounded during in the optimization phase, then we have the following conclusion:
()))ld( S
|
Thax (w:) + Thax (ld))d
(44)
d
|CCn
(2)
d
$ vn X; 1 )l(
(45)
(M;
(3)
F min |, min |IMH (M; -v1l2. ll2
(46)
ac;
where R4 (MH (; 21) (WO)² (V.").
(47)
(4)
((IMM((M;X)) 34112.
(48)
aw
2n
||x.1I [ 0 (WV) V Pmnxx (W:)) + o (WX)
Q
V d
(49)
where min |V'WO| is the smallest absolute value of each element in vector V'WO; min the smallest element in matrix S.
Lemma 8. Consider updating Wº, W.K. WV at iteration t. Suppose m (WQ), max (WK),
maax (WY) are bounded during in the optimization phase, then have the following conclusion:
XX MH(M;;X)lle s 1' Melle, where Z' is some positive constant.
(50)
IV (M++11 (M, G' Melle, where G' is some positive constant.
(51)
1.5 Proof Sketch of Theorem 3.
(I)Using GD update rule, we aim iteratively show
mmxx (W?)
2 1 E t),
min (CC.
2 f0r E (0......).
aw
(52)
min S, 2 K, r E {O......). 1), (M,i (1 m)) (Mo, X), r (0.
Denote : y= 162n2 (min VWW ) Let us discuss about the value of We know Wo E R Hdd x1 BJ E "N where Hd > 1, HD > Nn. We require 0, i.,, 8 > 0, min V'Wol O. It clear that 6 > can hold as long as Wh is bounded. And it is easy to show that if X, * 0, we
always choose V and Wo, such that min V'Wo 0. Since (C*(M) ERNn'x DDd suppose
aw
we initialize w w such that rank OCa (Ma)) = Nn², then we have Jmin aw some positive constant Further, we assume the initialization condition satisfies:
(OC,(Mo) aw
> 8 for
8n|| 8n||X|]//|[/>>22 (A) + A exp 1IIXI ((x)) + (AA))
|MM (Mo;
=112 S
82 (min (IV')11)2 min (6, 14
(53)
Remark The initialization condition can be satisfied if [ large and a (WV) small. v/ in Equation 12 is
18

---

Similar to the proof of Theorem 2 t O. Suppose it holds iteration t, we prove it holds at iteration t ++ 1.
we use induction to prove the theorem. Equation (52 holds when
[ - who 1.0
w - W hrr
| =
w. ʃ /Mri x)
1
(1))
N 2 [IMH(M;; x) - vllp S "
(181
Σ
Q
m 8/2 |MH X) - -vl12>
1=1
s=0
where (i) uses triangle inequality; (ii) comes from Lemma from the induction assumption that loss function f(-) linearly decreases until t-th iteration. Let vI 77y So we have
7 and Cauchy-Schwand inequality; (iii)
N
W hr++1 - Who
S
Σ
Q
2
11 m) ||MH (Mo; X) - vll2
(54)
A.0
1=1
==0
N
0
1
+1
(1 (1-u²) u²) (M((M;;)) - ull2
1=1
N
2
Σ
2
(0) 1 10_ (M; X) - ylll I 2
2^h h:
(55)
where (i) comes from the initialization condition. Then by Weyl's inequality, there is
Tmax (W
S Tmax (W
h,0)
+
1 5/ 3
2 %
2
aCn (Mr+1)
8Ch (Mo)
(3)
8Ch (M++1)
a0n (M.)
5
aw
aw
ow
ow
F
F
7V
n
| Σ
V w (M,; x
(8))
N
7V n
X 2
Σ
|M (Ms;)) - -1l2
N
X
Σ 0/ 2 Σ (1 7/) |MH (Mo; - v/122
s=0
X
Σ
Q MH (Mo;
(Mo; X) - v112
1 2
where (i) uses triangle inequality; (ii) applies Lemma 7 (2) and Cauchy-Schwartz inequality; (iii) uses Lemma] 7 (4); (iv) applics the induction assumption that the loss function f(-) linearly decreases until t-th iteration: (v) comes from the initialization condition. Then by Weyl's incquality, there is
aCp (Me+1) aw? h
aChh (Mp) aw?
5 2
o Omax
Tmax
For each element in Sih, we have close form
S ( W WK; x.)
exp
Xa
W.
X,, W k (?2
kj
2vd
19

---

Since we have already showed that max (W
h,T S - it follows directly each element in matrix
St is lower bounded by some constant for any L. Now we derive the expression of k:
Xik WhO - -X, wx'l("
exp
2
1
(1 W2,011 2 1x W Wx
exp
(1))
2 exp
1
7
9
(A%
21 11 (A P11x. 113)
Vi
(16)
IIXTI (12 (889))
exp
K2
where (i) uses Sauchy-Schwarce inequality; (ii) applies the induction assumption m (W m.e) S 2 and property of singular value; (iii is because d 1. Thus, we have min S; k. Finally, we aim to show f (M++1; X) - (1 - my)f (M, X). By Lemma since we have showed that m (W is
bounded, we can directly derive that
IV f (M+11; - Vf (M,;
= V was (Mc+1; V wwg (M: (Mr X)
- G - Melle
Finally, by Lemma choose 20 we have the following holds:
f (M MM+1;X7
= (M, VM(;
(0)
- /M, - - VV (M; XII2 + _
(0)
< *(M;X)
-
1
V (Mg (M/x)|
(Ms
(00)
X(Mx)- 1
Xf(M;
(C(M)
2"
8C(M))
ow
(1)
Of (M;
S (Max)) 4 n88
F aC (M.)
(9)
& f (M:x) X) 1 n [((MH (M;X) - 3) (Wo) T (v') @ $1
(v8)
$ XMi:X) 4 n n22 (min 1' |IMH (Mo: X)- - y112
(v1) (1 (1=m))(M;x))
where (i) uses Lemma (2); (ii) is because we choose " < 2 (iii) writes down the ex-
pression of gradient according to chain rule in Lemma (iv) uses the induction assumption
> _ and property of singular value; (v) uses Lemma (4); (vi) comes from
Tmax
aw (3); (vii) uses the definition of 7-
Lemma
Ne)t, 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
Vwyy (M++1; - WW
(Ms x S G G11 - Mell
(56)
20

---

Then by Lemma 4 choose 1/' < 1
20' we have
f XM = (M, (M;;)))):
Sf (M; - 1/ |V (M;; G (M;
Sf(M; - /" [M,
1.6 Proof Lemma in Section 1.2
Proof of (1).
Proof. Step 1: When Wº, W are updated, we aim to prove
|d((s)]le & mlld(C))llF
Step 2:
We aim show ||d(C:)]l= < Ja |
Σ taa (WK) ||d (WQ) || Combine the
above two steps, we can derive the bound in Equation (7
Proof of Step 1: First, we can write down the closed form of the differential of S,:
||d(S))l || = || @ d(Ci) S, @ T(xxxp C,)E) o d(exp(C:)E))]Iz
(57)
We reorganize the terms the right side of Equation (57 have the following equation:
|d((S))lr = || @ (d(C;) - ((cxp C E) O d((exp C)) ||
= || @ (d(C:) - ((exp C E) @ ((exp C:) O d(C:))
(58)
Since C, = [Ca..... ... C. we will investigate each Ca., h 1,2, ... H. We focus on the term d(C;) - ((exp C))E (C,) d(C.)) E in Equation 58) We write down the close form of the element the k-th row and j-th column:
[d(Ca.) T(exp(C.n)
(exp (C.)) O d(CCn)) B]s
(59)
Σ p#j
exp (Cinkp) d(Cinkp)
i
exp (Canks)
d(Cinky)
(60)
Σ exp (Conk))
cxp (Cinkj)
2 exp (Cinkj)
(68)
exp(Crrkkp)
1
(d(Cens)))
(61)
exp (Canks)
P#j
exp (Cihkz)
(156)
(62)
where (i) is expand the closed form of Equation (59): (ii) uses the Cauchy-Schwart inequality; (iii) is because each element in the square root in (ii) is upper bounded by 1. With Equation (61). we casily show
[|d(Cin.) ((expC ((exp Can) d(Ca)) Vn
[(((Csn)) S nnlld(Gan)l
(63)
Since every element in S; has magnitude less than 1, we have
[d(S)))lr ||s @ (d,),) Y ((exp C,) o ((exp C.) 6 d (C..) B)l,
(64)
W Wn
[ld(Cax)) T(exp C., EE) @ ((exp C.a) o d(CGn)) E
(65)
ny nVH)dd(C))
(66)
21

---

where (i) is from Sauchy-Schaweer inequality.
Proof of Step 2: We aim to show ||d (C:))l I va Xill f V E max (W/)): l|d (W?) ||
1- Sim-
ilarly, we investigate Ild(Cat) MI. 1.2, H. We have
x (X,WW)
||d(Cin)ll==
h vd
| (Wx) )|(
(67)
Then plug the above inequality to Equation (63 we can derive
n
|x max (wx) |a(
[ld(Sa)lle -
(68)
Vd
Thus by Schwar-Schwartz inequality, it is easy to show
H
|d((S))le S Vi
of
(Wx) 9
Proof of (4)
Proof. We first write down the close form of gradient of f(-) over W by Lemma and derive the upper bound of the norm of the gradient.
0f(M: x:) a h
1
X
((M:X.)
[x max (W
0f (M;
Vd
ac,
00,
(69)
By Lemma 1, there is
0f(M;X,) = ((MH(M;))) - Ui) v.) (WO) (V.) )) o S:
ac,
- -(((MM((M)))) (b) (WO) (v:) S o ((exp E BT) @ exp C.
(70)
Denote R, = (MH XM; - vi) (Wo) (V') R = Ra,... RaH). Write down the close form of the element in the k-th row and j-th column:
San ((Ran @ Ca @ Y (exp E)) ET) @ (exp Ca))as
exp(Conkk) Σ
Sinky
SinkkSShkkj
exp(Cinkj)
(exp Cinkk Sanki
(exp Dahkk))sshs Piikkp
Sinks
Ranks
Σ exp(Cinns)
pt1
exp(Cinia)
1=1
exp (Cihkj)
exp (Cinke)
+
Σ pン3
j= exp (Canks)
exp (Cenks)
Rinall
22

---

where (1) is due to the Sauchy-Schwand inequality; (ii) is because each element within the squre root term in (i) has magnitude at most 1. Thus, we can further derive
f X,)
Rur @ Sth ((Ran @ Sah Ca)) B & exp Ciall
ah
F
Ramell Ralle
(161)
S 11ll (W%) [M;H(M;;,) - Vill2,
max
where (i) if from the bound in Equation 7T]: (ii) comes from Cauchy-Schward inwquality; (ii) uses the property of Frobenious norm. Thus, by Cauchy-Schwartz inequality, can derive the upper
bound for [)()
8C
0f (M; X.)
1I 12 max [MM(MM
-vill2
(71)
Oc;
So plug the above inequality into Equation (69)
we can derive the upper bound for
[ 8W2
0f XM:
Of (M;X)
x++| max (W,x)
aw
aC2
mV H IXel
[ 11e max (W
mmax (W |IMH (M; X,) - Vill2
H max (W max (WV) |IMH (M; - 2/1l2
S <n/H||X
112
Proof of Lemma 3 (1). By Mean Value Theorem and Cauchy- inequality,
If (M++1: X,) - J (M; X:)
Of X.) MM+1 - M,
OW
0f(M;
OF(M) X.)
2
@f X.)
+
| -MelF,
(72)
OWQ
OWK
OWV
where Mi is between M, and M4+1. We can derive the upper bound of the norm of Vwv X.M;
0f(M; X.)
IIB (MH(M; X₂) - vi) (Wo) I|
awv
< ||B/||=||MH(MM Ville|1o111
< I | vellF
(73)
By Lemma 2
we know
Of(M; X,)
0f(M;; dWk
s Q: |MH (M; X²) valle:
< K; |MH (M; vell2.
DWQ
F
IIf (M++1; X,) - (M; X.)I1, S Q + K? + n² Ho² max (X.) (X) -Melrr
= -1 - M
(74)
Therefore, together with Equation
73 we have
Ilf (Mc+1;X) (Mr; X)ll, S N max Q? max K? n²H max X| | Mell
= M1| - -
(75)
Proof of Lemma 3 (2).
23

---

Proof. By triangle inequality, we have
Wf X)- V (MM;:)) I
/Vx XMM
+ Vv f(M++1i - Wwv
(76)
(IV (M (M VV w (M++1:X) - V W (Me+1i X)Il
(77)
+ Vw
(Mc+1; ff
- Vw (Mc-1: 11F)
(78)
Step 1: Derive upper bound for W (M; (M) = (M) ((M;; X.)
First, give the vectorized expression of Wwef (Mr; Recall we denote Uj (MMHMM; (M; vi) (WO)T (V?)") o S.. By Lemma 1 we can derive the close form of vec(VooJ(MM) X))):
veo( V w (M;; @ vec(U.)) vee (UU @ C.)R) BT) & veolexp Ci)
(8)
vec(U,) (E I..) vee (U: @ T ((xxxp Ci BB) @ vec(exp
(61)
veo(U,) (E vec U.)) 2 vec (T((exp C)) @ vec(exp
vec(U,) [E X In) vee U?) @ vec(S:)
I vec (U:) C vec I HH (E @ In) vec(U:) @ vec(S:)
(L2 (E @ L..) vee(U.) vec(1.1 S;),
(79)
where (i) uses the Lemma (ii) and (iii) comes from the property of vectorization in Lemma 5 (vi) uses the definition S; : (V) gives an equivalent expression of vec(U:); (vi) reorganizies (v). Further, it is easy to verify that:
|WllE - (M)) (M; v:) (Wo)² T (V:)")
St s:
S ||
(IIMH (M; X.) 2l (1 (1b:ll))l |/2 X || o max (WV)
(nv Ho (W || 112 + llvx | [l || e max (WV)
nn Homaa (WV) 1|
| + llvll:) || || p σ max (W.)
R
(80)
Next, let us derive upper bound for // (M++1; X.) - (Meri ( ||V We X,) [wee[Meer;
[[(E22 - (E @ I.,)) vec(U @ Ove(SS) -vo(Ui) o vec(St.)))l
vee(U.) @ vec(S.,++) 1) + vec(Us.)) @ veo(Si,41) vec(Us,) -vo((U)) vec(S [T - (B 60 Tw.)) vec(U @ vec(Si,t1) -vec(t))
(16)
S - / U1,e))l [|UcelF - StrllF))
(00)
nVH (livec (U,,411 -2,,)l R S - Ss,,ll)))
(w0)
nv H ([ 1 St,+1 RR,, 2 1 + R||St,11 - St,llF)
n. HT ( ① S1,441 -Ra S4,41 Ri, S RA, St.)l R||S S..llr))
HVH(([(R - Ri,e) 0 -i.--111 || + Ra, @ Si, - Ri, o St)l F + R|| -S.ll)
ny (1 Ra,el | S Stall + ReliSt.t11 Sa,
(81)
24

---

where (i) plugs in the expression in Equation 79 (ii) uses the fact that each element in St,++1 has magnitude at most 1, and Sauchy-Schwartz inequality; (ii) comes from the definition of II, E and R: (iv) uses the definition of Ur (v) is because triangle inequality; (vi) uses the fact that each element in S has magnitude at most 1, and Cauchy-Schwartz inequality. Next, we aim to derive upper bound of [2,,,+11 Rr,ell in Equation (8I)
[R,,41 R || ([MH(MM+)) 110-v)))1 v))
(MH(M;;X,) -MMM (w)
= -[(MH(MM-i:X,) - (
- (MH((,;;;) - (++++
(MH(M X,, - (VM()))))MM(M( (V
19 II(MH(M++ 11 X.) MH(M;; (V) + II(MH((H X.) - v.) V(o)o%Il
(K)
Z.IlM++1 Mellelll
(W 111
+IIMH(M X.)ll= 11 X: ||W W. 11=lw/0
(20)
Z X FA (W : 12 |1 - -
(nv Hamax (WV) || | 10 11uslla) I|xelle | w ||
(10)
(Ze |Xall maa (WV) 1| + (no Homax (WV) |Xillr || + (1111112) || | 1|
x M - MellF
P. |M M.llF,
(82)
where (i) is because of the triangle inequality; (ii) uses the definition of Z, 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 (8I). we can finally derive the bound for V f (M++1; X,) - Vwo. Xi)l.
J(M+H1 - Vwo I(Me+1i;
RVH (1 R1 + || || - Sr.ll + R R1 - St.el)
(1)
3 nVHHP||Me+I MIlFF + 2Rov S1++1 -Seell
(18)
- nVHP,I|Mi41 Mile= + 2Rn H V 0f + + - MillF
19 MillF,
where (i) from Equation (1) ; (ii) uses the definition of R in Equation (80;; (iii) comes from Lemma 3 (3). Since WQ and W are symmetric in the Transormer structure, similarly, we can
derive I.A
Step 2: In this step, aim to derive bound for V (Me+1:X) Vw- (M):
25

---

VV wv JMMM(+1;;; Nw f(Me; x;)
Bilet1 (MH (Me+1; X;) u) (Wo)² - -B, (MH (M; vb) (Wo)
(16)
T (MH (M++1; - -1)) (WO) B (MH (M; X;) - and (oo)
F
+ B (MH X,) v.) (WO)* -B,, (MH (M; - -w) (W0) T1,
(18)
$
|| |MH (M+1;; MH (M; X,)I || || Beelle || (M;; -. -Vll=ll
nV "V 12 Me+ Melle + |S St,el P || (IIHH (M++1: x:)Il +
(1)
Vo8 + of + 112 (NVV L.VHomma. (W ||Xill | + +1lssll)) (oMell))1[M+4>> -Mllle
+ H | Xx 11 Z4
- M
(18)
V of + /ेद |Xillr 10 (NVToomax (WV) ||Xel| [W0l|2 + +1bslll)) + ny [ Zi
| Mellr
== I. M - F
where (i) is from Lemma (1): (ii) uses triangle inequality: (iii) uses Cauchy-Schwan inequality; (iv) comes from the definition of Bi.... Za(in Equation (74), Cauchy-Schwartz inequality and triangle inequality; (v) comes from Lemma 2 (3) and Cauchy- Bhuchy-Schwartz inequality; (vi) reorganizes (v).
Now we combine the result in Step and Step 2. and plug into Equation 78
we can finally derive
N I + M - MF
Vww (M++1: W (M: X)Il $
+
1=1
s <Nmmx L + max L + max LV M| Mell
(83)
G|]MM+++ - 11MM+11-Mell-
(84)
1.7 Proof of Lemma in Section 1.4
Proof. Proof of Lemma (1): We consider the differential of the element in the k-tth row and j-th column. First, let us write down the closed form of each element:
(Can)s -| W Xıj. W 11/22Vd
Next, we consider the differential of each element over W
d (Csh)a 2vd (f)xae (W + a(W()) X,j W ||2 -
-x,,
Vi Vid" (Xix X.x. X.y, w + (d((W)))
where o(d(W))) denotes the higher order of d d(w ). Leave out the higher order differential term, we derive
I|d (Can)as |1. S 1/2 (1xxxelllld(( max |1 112 + (dd(w Il T (W ||Ix Iaall (1-l11))
vJ
1
$ X [
[I Fmax (W l112)
1
Vd Xxaxlll 11a (W 02 T
(W
| + 1
26

---

I|d [dn)ll = k=1j=1
Id(Canaaill
||Xat-l22 112 v Imax (Wx)) + mmx (Wx)
v /1xxxx + ||x (1 )l
mux (W(²) + muxx (Wx)
||Xix-ll, 1l
n 11 11 + Σ
|| 11
j=1
max (wx²) + +maxx (W/+)-
Xix 11
(nllIXix-113 +
|x 1x ll )ll
Vd
Vd. V maxx (w/)) + max (W/). ||x:|le v2n x2
2m
| |xell V 0 (W/)) + o ll
d
Proof. Proof of Lemma 7 (2): First, let us write down the closed form of @(Cin)e aw We have
((Ca))a (Xx. W )la @ X
(85)
aw
Thus, we can derive upper bound for d (crinn )
0 (CaNas)) aw?
)l = (x (W + ((W?)) X1, Kg,,Wh Ta Nu + (x X,y, Tia 00 X Ivvll./Vaa
d(W 3 La & Iiall
$ a| BILalellddw
(86)
Thus, have the following:
0 (Ca.) d
("aw
(C(Cm)ss
Σ
d
aw
F
j=1
F
< |
1x
S |m1xx(11111(((3)1111>
Proof. Proof of Lemma (3):
of (M;X,)
= -I((MH(M: v.) v.) (WO) T (V7)7) o. F
a0,
|F
[(MH (M; - wi (WO)T T (V:) )11- min SS
2 mimmiI//v 1 min |S ||MH XM; - -vill2.
27

---

Proof. Proof of Lemmar (4):
Xf,
(f(M:;))
(f((M:
ac;
O
=
vec
aw
vec
80;
aw
F
1
Xf(M;X;)
ac;
8Ci
aw
= (I( (M; X.) - vi) (Wo) T (V!) .
2n
|xl|
Pmax (Wo))
+ o (Wx)
F
2m V d
||X 120 (W /
o max
( + o (W || (M; 36ll2
Proof. Proof of Lemma 8 (1): The proof is similar to the proof Lemma (1). So we do not include the details here. We can similarly derive
Ilf (M++; (M;; X)112 S N max Q? + max K;a n² H max | [Wol
M,
Mell
Z' [|M++++ -
(87)
Proof. Proof of Lemma 8 (2): By triangle inequality, we have
f Vw XMI:
+ f(Me11xX) Vwx XMM + VV f(M++1; - ww (M
V Wwo V
N
(Il (M++15X) - V w (M+11;
V (M+Hix)) - V (Me+i;
V (M++1; X) - Vw-f (M X) l[F)
(88)
Step 1: Derive upper bound for
W/ (M{+1 we f(M;; Xi))ll Mex we f X,))
First, we give the vectorized expression of Wwof X,). Recall we denote U, = (IMH (M; - vi (Wo) (v:))) o S. By Lemma 6 we can derive the close form of veo(Vwe X;)):
vec (WW ( vec(U.) vec
ac
(89)
aw
Further, recall we have defined R and the following inequality holds:
| < R
(90)
28

---

Next, let us derive upper bound for wew (M++1:X;) - Vwef (M
X,)
X,)
aC;(M++1)
(C.(M))
vec(U.,++1)
vec(Upa))
ow
aw
(C,(M++1)
(CC((
vec(UL,++)
vec(Ui,+1)
(91)
ow?
OW
(M:)
)
(Mt(
+ veo(Us++1)
veo(Ui,))
ow
aw
(0)
ac; (M++1
ac; (M))
+ veo Usellla ac; (Mr) Ow
aw
aw
12
(00)
x RVm X
a (w:)
||vec (U1++1
ULD)ll Villl (Amax (W:2)
+ max (W
(w
+ + (amax (WR)) + Tmax (Wx+) [| O S4,41 - Rag O St
F
S RVTIIX.
(W:)
+ |
Tmaxx
(w2) + Tmax (Wx+)
RA,)) S1.t1 ||a,, O S4,41
Rig 3 Sz) (IF)
X
(d (W:) + Vr |X (omux (W.))
mmax (W:K)
(I) R + ||R.,ell r ||S,,41 st,
(92)
Next, we aim to derive upper bound of || |Ri,+1 RR.elF in Equation 92 Equation (8T). we can derive
Similar to the derivation
(iv)
RR,41 RRlll S [Z Xiile Tmax (WV) || + (nv Homax (WV) |X ||lo|| ll2 + (1vill2) |x,11
|woll2
x || - MellF
== M - Mel r.
(93)
Plug Equation
82
into Equation
992.
we can finally derive the bound for
/ (M++1; X3) Vee (Me+i; X:) F-
IV I X2)
Vwo (ovoJMMen [MM++1 X,)lf
< R/Tl X // | (W |1 FF
+ v Xx Xx11 ( (amax (w:)) + mmx (W
x (1|R4411 - RR + |R || - S),llF))
< R 11 -Mel + vn Xall (omax (W,) + Tmax (Wx))
x (P/II Mille + VnR |xll (omaax (w²) + m (Wx)) |1 - Mellr)
= 1
and plug into Equation 78 we can finally derive
N Σ (I) + IK + 1 - Mallr
Ww (M++1; Vw (Mr X)Il &
1=1
< <N(max I + max LK + max LV (1 M| Mille G|IM4+ 1111111-M,1-- M,IFF
(94)
(95)
29

---

2
NeurIPS paper checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:[Yes Yes]
Justification: See Theorem 123
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please see our conclusion 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
30

---

Justification:
See Appendix, which provides proof for each Theorem.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See experiment setting in Section 5.2
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
31

---

Answer:
[No]
Justification:
We do not include the open access to code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm
Ahbllii/guidde//oeeeubibiiii
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ lttipp.cc/pullic//uiice gwccc///wuee//eee /oode/umddeeemmii cerbbbmsssioon awiionooliccoo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See Section 5.2 for experiment setting.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Please see Fig 2 and Fig 3 We have 1-a eтror bar.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
32

---

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 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful 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.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification:
We do not include the compute resources detail.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well 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 failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
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 NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:[NA]
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 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.
33

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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 use Chen et al. [021
Guidelines:
The answer NA means that the paper does not 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, papprswiiiiiie thcode com datastts has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
34

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
35


Reviews:
[{"review_id": "pnfxOTEisz", "reviewer": "Reviewer_cxoC", "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.", "questions": "See weakness above.", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- The analysis of Gaussian kernel attention is insightful and novel. The authors show theoretically and empirically that Gaussian attention can lead to better convergence properties compared to Softmax attention. This provides valuable understanding of alternative attention mechanisms.\n- The paper addresses an important question about the optimization dynamics of Transformers, which is crucial given their widespread use. The motivation to understand why Transformers work well and potential limitations of - Softmax attention is timely and relevant.\n- The theoretical analysis is reasonably thorough, with the authors deriving conditions for global convergence for both Softmax and Gaussian attention under different update scenarios (Theorems 1-3). This helps formalize the intuitions about attention behavior.", "weaknesses": "- 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.\n- 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].\n- 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.\n\n[1] Y. Huang, Y. Cheng, and Y. Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">>**Your comment**: 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.\n\n**Response**: Thanks the reviewer for the comment. First, we need to clarify that the goal of our paper is to investigate the **training dynamics** and **landscape** of Transformer model. The main purpose of our experiment is to **provide a concrete example** to illustrate the existence of local solution in Softmax attention. Similar experiments of **shallow** networks can be found in [1,2,9], which have settings similar to ours. \n\nWe agree that additional experiments on larger models could potentially provide further support for our theory. However, due to computational budget constraints, we are unable to complete these experiments during the rebuttal period. We will add results on these larger models at a later time.\n\n>>**Your comment**: 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., [4].\n\n**Response**: We thank the reviewer for the comment. We will add the following detailed comparison with the related works in the revised version:\n1. Comparison with [3,4]: These works analyze the convergence of in-context training, where a prompt is contructed with all the training samples and a single test sample. The goal of these works is to achieve the zero test loss (in expectation) by optimizing over the loss function modeled by the propmpt. However, our analysis is based on standard empirical loss minimization, which does not involve any prompt construction.\n2. Comparison with [5,6,7,8]: This line of works analyze the global convergence of training over-parameterized fully connected networks. However, Transformer structure and the optimization landscape in our study are more complicated.\n3. Comparison with [9]: This work is closely related to our work. However, the primary goal of [9] and our work is **different**. In [9], the main purpose is to **build a convergence theory** of shallow Transformer with realistic structure and initialization, but they do not provide roles for different matrices in the convergence. While in our paper, we focus on **investigating the role of each variables** within the attention mechanism in optimization. Further, we analyze **the convergence behavior of different attention kernels**, i.e, Softmax, Gaussian, which is not covered by [9].\n\n\n\n>>**Your comment**: 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.\n\n**Response**: We thank the reviewer for the comment. Here we discuss the insights about the benefits of using a Gaussian kernel in training Transformer models. As discussed in [1], Gaussian kernels have a natural interpretation of assigning “attention” to different tokens, similar to Softmax attention. However, the Gaussian kernel performs intrinsic normalization, which allows the attention mechanism to have a more reasonable condition number than Softmax attention. Therefore, designing specific Transformers with Gaussian attention can greatly improve the stability of model training. For example, in [1] Table 1, the kernelized attention (Gaussian kernel) and Skyformer (modified Gaussian kernel) show better performance than Softmax attention. In [10], it is showed that the Gaussian Transfomer is more efficient and generalizes better than Softmax Transfomer.\n\n[1] 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.\n\n[2] Tay, Yi, et al. \"Long range arena: A benchmark for efficient transformers.\" arXiv preprint arXiv:2011.04006 (2020).\n\n[3] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. \"Trained transformers learn linear models in-context.\" arXiv preprint arXiv:2306.09927 (2023).\n\n[4] Huang, Yu, Yuan Cheng, and Yingbin Liang. \"In-context convergence of transformers.\" arXiv preprint arXiv:2310.05249 (2023).\n\n[5] Allen-Zhu, Zeyuan, Yuanzhi Li, and Zhao Song. \"A convergence theory for deep learning via over-parameterization.\" International conference on machine learning. PMLR, 2019.\n\n[6] Du, Simon, et al. \"Gradient descent finds global minima of deep neural networks.\" International conference on machine learning. PMLR, 2019.\n\n[7] 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.\n\n[8] Allen-Zhu, Zeyuan, Yuanzhi Li, and Zhao Song. \"On the convergence rate of training recurrent neural networks.\" Advances in neural information processing systems 32 (2019).\n\n[9] Wu, Yongtao, et al. \"On the convergence of encoder-only shallow transformers.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[10] Lu, Jiachen, et al. \"Soft: Softmax-free transformer with linear complexity.\" Advances in Neural Information Processing Systems 34 (2021): 21297-21309.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z2vPUog0MI", "reviewer": "Reviewer_mp2C", "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.", "questions": "1.\tThe 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? \n2.\tI 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": "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.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is overall clearly written. The authors delve into the role of different weight matrices in Transformer optimization and establish different convergence theorems with regard to different trainable weight matrices. This aspect is novel. Their theory also suggests that Gaussian kernel may be more effective than the widely used softmax kernel, which if validated in larger models, can have a great impact on the field.", "weaknesses": "Minor points:\n1.\tIt 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.\n\nMajor points:\n1.\tThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">>**Your comment**:  The difference in convergence speed in Fig. 2 seems like constant level with same rate.\n\n**Response**: Thank the reviewer for the comment. We need to clarify that it is **reasonable** that the difference in convergence rate in Fig. 2 is constant level. Further, we provide **empirical reason** for the different theoretical convergence rates in Theorem 3 by landscape visualization in Fig. 4. \n\nFrist, Transformers with Gaussian and Softmax attentions theoretically should have **the same convergence rate** in our experiment setting in Fig. 2. Notice that the algorithm we use in Fig. 2 is updating **all the variables** in the attention ($W^Q,W^K,W^V$). As discussed in Remark 3, when all the variables are updated, both Transformers could achieve linear convergence rate with appropriate network size and initialization. Thus, the theoretical difference in convergence rate is constant level. However, due to the bottleneck of optimizing $W^Q$ (or $W^K$) in Softmax Transformer as illustrated in Theorem 3, Gaussian Transformer still converges **faster** than Softmax Transformer when $W^Q, W^K$ are included in optimization.\n\nSecond, in terms of the different theoretical convergence rate in Theorem 3 when **only $W^Q$ (or $W^K$)** is updated in both Transformers, we  provide **empirical reason** for the different convergence rates in Theorem 3 by landscape visualization in Fig. 4. By showing the existence of **more local solutions** in Transformers with Softmax attention, we provide evidence for the comparison of the convergence rates: Compared to Transformer with Gaussian attention, a Transformer with Softmax attention is more likely to **obtain a sublinear convergence to stationary points**, not global optimum.\n\n\n>>**Your comment**: The insights of critical proof steps that cause different convergence rates in Theorem 3.\n\n**Response**: We thank the reviewer for the suggestion. We will add the proof sketch of Theorem 3 in the revised version. Here we will briefly summarize the critical steps.\n\n**Step 1**: Derive the **closed form gradient** of loss function over variable $W^Q$. Please see Lemma 1 equation (23),(24) in line 778 for Softmax attention Transformer, and Lemma 6 equation (41), (42) in line 821 for Gaussian attention Transformer. Intuitively, the gradient of Softmax attention is much more complicated than Gaussian attention Transformer, which will lead to a more complicated landscape and more local solutions.\n\n**Step 2**: **Analyze the gradient** of Transformers with both kernels with the **same initialization** Equation (12). For Gaussian attention Transformer, it can be iteratively shown during the gradient descent training: 1) The variable $W^Q$ is bounded; 2) The PL condition holds (i.e, the optimization landscape remains near-convex); 3) The loss function decreases linearly. (Please see Equation (52) in line 829.) For the Softmax attention Transformer, there is no guarantee that the PL condition holds during iterative gradient descent update. To illustrate this claim, we construct a concrete counterexample to show that the PL condition does not hold. (Please see line 284).\n\n>>**Your question**: 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?\n\n**Response**: Respectfully, we need to clarify that we **do not** construct an example to illustrate the failure of the Transformer with a **Gaussian kernel**. Instead, we provide a detailed example to illustrate that the Transformer with Softmax attention converges to a **stationary point** with a sublinear rate, but fails to achieve linear global convergence. (Please see Equation (14) in Theorem 2 and the failure to satisfy the PL condition in lines 284 and 285). \n\n**Additional condition for Softmax Transformer**: For a Transformer with Softmax attention, the **PL condition or convexity of loss function** should be assumed during the **whole training phase** to achieve the same linear convergence rate as a Transformer with Gaussian attention. This condition is much more stringent than the condition for Gaussian attention, where the PL condition is assumed **only at initialization**.\n>>**Your question**: 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?\n\n**Response**: The similarity comes from the symmetry of $W^Q$ and $W^K$ in attention mechanism. The symmetry means the training dynamics/gradients of  $W^Q$ and $W^K$ have **exactly the same** properties. Please refer to Appendix 1.2 Lemma 1 (3). We write down the closed form gradient of $W_h^Q$ over $f$ for the Softmax Transformer. Similarly, we can derive the gradient over $W_h^K$, with the only difference from Lemma 1 (3) is that, the $W_h^K$ term on the right side of the equation is replaced with $W_h^Q$. Thus, the gradients of $W^Q$ and $W^K$ have the **same** structure. As long as we initialize $W^Q$ and $W^K$ with the same properties, the training dynamics are also **the same** for both gradients. Intuitively, this stems from the fact that $W^Q$ and $W^K$ are in symmetric positions with respect to calculating their gradients in the attention head, see equation (3).\n\n[1] Huang, Baihe, et al. \"Fl-ntk: A neural tangent kernel-based framework for federated learning analysis.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] Gao, Tianxiang, et al. \"A global convergence theory for deep relu implicit networks via over-parameterization.\" arXiv preprint arXiv:2110.05645 (2021).\n\n[3] Wu, Yongtao, et al. \"On the convergence of encoder-only shallow transformers.\" Advances in Neural Information Processing Systems 36 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qdTJ0sK4Le", "reviewer": "Reviewer_ynwg", "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.", "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?\n\n2. 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? \n\n3. On line 286, is it implicitly assumed that the global optimum should have zero loss? Why must this be true? \n\n4. 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?\n\n5. 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": "Not quite. See the weakness section.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper gives a theoretical analysis of training dynamics of Transformers, and provides conditions to guarantee global convergence.\n\n2. The comparison between Gaussian and Softmax kernels is interesting.", "weaknesses": "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. \n\n2. The comparison between the Gaussian and Softmax kernels only updates $W^Q$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">>**Your comment**: The new insight we can gain compared to Wu et al [2024]; the fixed $W^O$ in optimization.\n\n**Response**: See the general rebuttal at the top of the page **Comment 1**.\n\n>> **Your comment**: The comparison between the Gaussian and Softmax kernels only updates $W^Q$.\n\n**Response**: We'd like to clarify that we also studied the role of the variable $W^V$ in Gaussian and Softmax kernel and found them to be the same. We have commented on the similarity between the two Transformers in **Remark 3**. Specifically, the conclusions in Theorem 1 and Theorem 2 hold for Transformers with Gaussian attention, and the linear decrease of the loss function also comes from $W^V$ when $W^Q,W^K,W^V$ are updated.  Since the conclusion and proof technique are very similar to Theorem 1 and Theorem 2, we did not include the details in the paper. To avoid confusion, we will add a corollary to clearly state the similarity of $W^V$ in both Transformers in our revised version.\n\n>>**Your question**: The comparison of assumptions in Wu et al [2024] and our paper.\n\n**Response**:  See the general rebuttal at the top of the page **Comment 2**.\n\n>>Your question regarding the importance of analyzing individual variables. Can the results obtained with one variable trainable and others fixed explain their roles when they are simultaneously updated?\n\n**Response**:\n**The importance of unravelling individual variables**: First, unravelling the individual role of each variable helps us understand the optimization bottlenecks in Transformers [2,3]. For example, in [2], it is shown that the rank collapse of the attention kernel can occur when the variable $W^Q,W^K$ are updated, which can lead to failure in optimization. Second, acceleration optimization algorithms can be designed based on the understanding of the individual role of each variable [3,5]. In these works, Nyström method is used in gradient computation to greatly accelerate the training. The algorithm design is specific to the training dynamics of individual variables in Transformers. Third, analyzing the training dynamics of individual variables leads to a new network design that improves the optimization performance [3,4,5]. In these works, attention structures other than Softmax attention are proposed to improve the generalization and computational efficiency.\n\n**Simultaneously update**: Yes, the property obtained when analyzing a single variable updates does carry over to the situation when this variable is updated to other variables simultaneously. The reason is as follows: Our over-parameterization structure and initialization ensure all the variables **always stay in a locally near-convex region**. Within this region, the optimization landscape is relatively smooth, and the partial gradient of any single variable does not change much when other variables are updated, thus analyzing the gradients of different variables separately is able to explain the role of each variable when they are updating simultaneously.\n>>Your question regarding the existence of global optimum in line 286.\n\n**Response**: Yes, it is implicitly assumed that the global optimum solution has zero loss. This is true because of the over-parameterization of the network, where the network size is $\\Omega(Nn)$ ($N$ is the sample size). When there are more parameters than samples, there exists a solution that can completely fit all the samples. Further, in some other works with similar settings, the existence of global optimal solution can also be derived from the limit of a  Cauchy sequence (See [8] Appendix B.7). To make the proof more rigorous, we will add a similar derivation of the Cauchy sequence in our revised version.\n\n>>Your question regarding the statement \"transformers perform exceptionally well\".\n\n**Response**: We claim that our result explains why Transformers perform exceptionally well because we derive a **global convergence** analysis of training Transformers (Theorems 1 and 2). These two theorems show that even the plain gradient decent algorithm can find the global optimal solution with a network size of $\\Omega(Nn)$, allowing the training loss to decrease to $0$. This provides theoretical proof for the outstanding training performance of Transformers.\n>>Your question regarding the sentences like \"only updating $W^V$ already leads to global convergence\".\n\n**Response**: No, it does not suggest that it is easier to optimize more variables. As we aim to scrutinize the role of each variable in the attention kernel, we found that the global convergence is dominated by $W^V$, which prompted us to later focus on studying the **role of $W^Q$ and $W^K$** in optimization. These sentences are meant to emphasize that the role of $W^Q, W^K$ remains unclear if updating $W^V$ alone already leads to global convergence (See Theorem 1,2), i.e, does updating $W^Q,W^K$ contribute to the linear convergence of loss function? Thus, these statements lead to our discussion on $W^Q$ alone in Theorem 3.\n\n[1] Wu, Yongtao, et al. \"On the convergence of encoder-only shallow transformers.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[2] Noci, Lorenzo, et al. \"Signal propagation in transformers: Theoretical perspectives and the role of rank collapse.\" Advances in Neural Information Processing Systems 35 (2022): 27198-27211.\n\n[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.\n\n[4] Choromanski, Krzysztof, et al. \"Rethinking attention with performers.\" arXiv preprint arXiv:2009.14794 (2020).\n\n[5] Lu, Jiachen, et al. \"Soft: Softmax-free transformer with linear complexity.\" Advances in Neural Information Processing Systems 34 (2021): 21297-21309.\n\n[6] 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for the comments and suggestions. We will summarize the strength of our work from reviewers as following:
1. The analysis of Transformer is under realistic setting, and requires no strict data assumption.
2. We delve into each variables within attention kernels. The analysis on Transformer landscape and comparison between different attention kernels are novel and significant.
4. The analysis is non-trivial and requires sufficient techniques.

Regarding the weakness and questions, we will address the several comments about the comparison with [1] here. Please see our detailed response individually for other comments. We are glad to answer your further questions.

**Comment 1**: The new insight of our Theorem 2 compared to [1]; the fixed $W^O$ in our paper.

**Response**: We thank the reviewer for the comment. We state the differences between [1] and our work in terms of optimization algorithm and discuss the new insight we can gain from our setting.

**Difference in optimization algorithm**: In [1], all the variables ($W^Q,W^K,W^V,W^O$) are updated to build a convergence theory, while in our work, we consider updating $W^Q,W^K,W^V$ in Theorem 2. We find that including $W^O$ in the analysis will **hide the true role** of $W^Q,W^K,W^V$ in optimization (see the explanation in the following paragraph). Given that we aim to unravel the role of different variables **within the attention head** in the convergence analysis, we leave out $W^O$. However, we need to point out that if we include $W^O$ as a variable, we can still achieve the same convergence rate as in Theorem 2.

**Difference in insight**: In Proposition 1 of [1], the linear decrease of the loss function **only** comes from updating the **output layer $W^O$**, In our work, however, the linear decrease **comes from $W^V$ in attention head**. Our result provides a **different** insight from [1]: Optimization on the **attention head alone** can achieve linear convergence rate. In our analysis, the same linear convergence rate can be derived if we include $W^O$ in the optimization, but it remain unclear whether $W^V$, or even the attention head contributes to the linear decrease since updating $W^O$ alone can already lead to the same result. Thus, leaving out $W^O$ makes our insight **more clear** regarding the role of $W^Q,W^K,W^V$. 

**Comment 2**: Comparison of assumptions in our Theorem 2 and [1]. New insight from the conditions.

**Response**:

**Comparison of conditions**: The conditions in [1] and our Theorem both consist of two parts: the network size and the scale of initial weights. We will compare the two conditions below:
(1) Network size: [1] considers the single head attention mechanism, which requires $D\geq n,d\geq N$. It means the embedding dimension and model size are lower bounded by sequence length, sample size, respectively. In our work, we require $Nn\leq HD$, while $d$ can be small. If we compute the total number of neurons in each variable of the attention mechanism in $H=1$ case, then the lower bound of total neuron number in our result is $\Omega(Nn)$, which is **exactly the same** as [1]! This result shows that our network size lower bound is **consistent** with the literature.
(2) Scale of initial weights: The initialization  conditions in two papers are very similar; it is **almost equally difficult** to satisfy the initialization conditions in two papers. The initialization condition in (8) implies that: 1. Both works require the smallest singular value of the attention head $B_0$ ($\underline{\lambda}^B$ in our paper) not to be too small; 2. Both works require the scale of each $W_h^Q, W_h^K$, and $W^V$ not to be too large; 3. Both works require the initial weight not to be too far from the global optimal solution; 4. Our Equation (8) requires the scale of $W^O$ to be large, while in Wu et al. [2024], the scale of $W^O$ has an **upper bound**. From above comparison, we can conclude that the initial conditions in two papers are **very similar**, except for the requirement for the scale of $W^O$, which is opposite. Further, we can verify that the gap between lower bounds for the smallest singular value of attention head ($\bar{\lambda}^B$ in our paper) is constant level, as well as the gap between the upper bounds for the scale of $W^Q,W^K,W^K$. The only difference is the scale of $W^O$, which results from the difference technique between papers: The linear decrease of loss function in [1] comes from updating $W^O$, while in our paper, the linear decrease of loss function in our paper comes from updating $W^V$. Please refer to the response to the previous question (**Difference in insight**).

**Regarding the new insight**: As we discussed in the previous point, the new insight we can derive from Theorem 2 is that: Optimization on the **attention head alone** can achieve a linear convergence rate. Please refer to the response to the previous question (**Difference in insight**). 

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

Meta-review: 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!
Decision: Accept (poster)