Title: Unraveling the Gradient Descent Dynamics of Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: Optimization

Transformer
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

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.

1 Introduction

Transformer model architectures have become popular in machine learning, delivering remarkable performance across a wide array of tasks. From natural language processing [Vaswani et al., 2017, Beltagy et al., 2020] to computer vision [Dosovitskiy et al., 2020], these models have set new standards in performance and efficiency. Popular models include BERT [Devlin et al., 2018], RoBERTa [Liu et al., 2019], DeBERTa [He et al., 2020], GPT models [Radford et al., 2019, Brown et al., 2020] and ViT [Dosovitskiy et al., 2020]. Despite their empirical success, a comprehensive understanding of their optimization process remains elusive. As highlighted in Liu et al. [2020], the training of large Transformers can sometimes result in deteriorated performance. It is therefore critical to develop theoretical insights for researchers and practitioners to better understand the practical performance of Transformers. However, the complexity of their architectures, coupled with the non-convex nature

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 [Jain et al., 2017], [Jin 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 O ( poly ( N ) ) , where N is the sample size. Similarly, [Allen-Zhu et al., 2019] show that ReLU fully-connected networks with at least O ( poly ( N ) ) neurons can achieve global convergence using GD or SGD. From a statistical perspective, [Li et al., 2023] 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 4 ( p/δ )) can guarantee the global convergence with probability at least δ using GD. Despite this extensive body of work on 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 to traditional deep learning architectures, Transformers incorporate a unique level of intricacy through their attention kernel [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 et 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 a 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 to follow a 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 V ) 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

Anumber of research works have focused on the theoretical analysis and interpretation of Transformer models, revealing crucial insights into their practical performance.

Liu et al. [2020] showed that heavy reliance on the residual branch in 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 an 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 depth-dependent scaling of residual branches as a potential solution. They specifically investigated scenarios where token rank equals one, which can hinder Transformer training. Their findings demonstrate the occurrence of the vanishing gradient issue, however, this work does not comprehensively characterize the vanishing gradient problem throughout the entire training process.

A recent work Wu et al. [2024] analyzes the convergence behavior of shallow Transformer, which builds a convergence theory of 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 is different from Wu et al. [2024]. We not only derive the global convergence analysis (Our Theorem 2), but also investigates the role of different variables in optimization.

Some other works focus on improving the optimization of Transformers empirically. [Huang et al., 2020] have proposed an initialization strategy such that no warm-up or layer normalization is needed to train Transformers efficiently; in Shazeer [2020], the GLU variant of token embedding has been showed to be better than plain embedding in the optimization of Transformer models with Softmax attention kernel. It is worth noting that the above works all primarily focus on empirical investigations into the training of Transformer models, lacking a comprehensive theoretical analysis of the underlying mechanisms.

Some recent research has focused on the convergence analysis of Transformer-based models within the in-context learning (ICL) framework. For instance, Huang et al. [2023], Zhang et al. [2023] explores the learning dynamics of a one-layer Transformer with Softmax attention trained via gradient descent to learn linear function classes in-context. However, this line of study primarily addresses the general convergence performance of Transformers within the ICL setting and does not delve into the role of individual variables. More specifically, these works analyze the convergence of in-context training, where a prompt is constructed 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 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 i , y i ) } N i =1 , where X i ∈ R n × D , and y i ∈ R n is the label of the dataset. The output from the Transformer model is the prediction of the label. The Transformer structure is formulated as follows:

In the above notation, W Q h , W K h ∈ R D × d is the query weight matrix and key weight matrix, respectively; W V h ∈ R D × d is the value weight matrix; these matrices are the main optimization variables throughout the paper. Further W O ∈ R Hd × 1 is a fixed matrix, representing the weight of the output layer; H is the number of attention heads; S ( · ) is a kernel function of variables W Q , W K and input X i . Attention ( · ) is the attention head function; MH ( · ) represents the multi-head attention function. For example, with the Softmax attention [Vaswani et al., 2017], S ( · ) can be written as:

where for a given n × n matrix Z , Sofmax( Z ) := [Softmax( Z 1 ) , · · · , Softmax( Z n )] . Throughout, let us denote S ( · ) kj as the element of k -th row and j -th column in matrix S ( · ) . Let X ik · ∈ R D denote the embedding of the k -th token in data X i , which is the k -th row of matrix X i . The structure of Transformer model can be found in Fig 1, where we denote S ih := S ( W Q h , W K h ; X i ) .

Based on the above Transformer model, we consider minimizing the following empirical ℓ 2 loss function for the entire data set { X i , y i } N i =1 :

where M := ( W Q , W K , W V ) is the set of variables that can be optimized.

For notation simplicity, next we define the vector version of the Transformer model given in Equation (1), for the entire dataset { ( X i , y i ) } N i =1 . Towards this end, let X ∈ R Nn × D denote the columnstacked matrix of each single data X i . Similarly, define the stacked label y ∈ R Nn . Then we can define:

i = 1 , 2 , · · · , N, h = 1 , 2 , · · · , H for simplicity.

diag( W V 1 , · · · , W V H ) ∈ R HD × Hd denote the diagonalized weight matrices that include all value weight matrices for all attention heads. Using these definitions, We can simplify Equation (5) as

Thus the empirical loss function given in Equation (4) can be simplified as

For more notations in the following sections, we will use subscript t to represent the variables in t -th iteration, e.g, M t := { W Q t , W t K , W V t } . Similarly, we denote B t as the matrix B at t -th iteration.

It is important to note that, in the above description and throughout the paper, we model the Transformer training problem by using a single-layer Transformer, with a regression loss. In practice Transformer models can exhibit greater complexity (different loss functions, multiple layers, etc). For example, the text classification task has an additional mean pooling layer followed by the output of the Transformer structure. Further, they usually contain downstream MLP modules. However, we choose to use the simplified version due to the following reasons:

First, the primary objective of this work is to understand how different attention kernels affect the training dynamics of the Transformers, so we do not include the layer normalization in our model. In fact, in the literature, many works that analyze popular network structures also do not consider layer normalization. For example, in [Huang et al., 2023, Zhang et al., 2023], both analyze the convergence performance of Transformers but normalization is not considered.

Second, we do not include the downstream MLP module in our work since we are interested in the role of self-attention layer in convergence analysis, and the single-attention model is also the standard model used in [Huang et al., 2023, Zhang et al., 2023]. Further, the analysis of MLP is standard in literature [Allen-Zhu et al., 2019, Du et al., 2019, Nguyen and Mondelli, 2020]. And it is worth noting that our choice to focus on a one-layer Transformer is consistent with other works that similarly aim to investigate the core training dynamics of Transformers, e.g, in [Tian et al., 2023], a single-layer Transformer is considered as a basic model.

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 ⊂ { W Q , W K , W V } . Below we summarize our results.

Common convergence conditions with Softmax Attention : When the activation function S ( · ) is either the Softmax or Gaussian function, and the embedding dimension D is at least O ( Nn ) , optimizing Equation (6) can achieve a global optimal solution when M = { W V } and M = { W Q , W K , W V } .

Different behavior between Softmax and Gaussian Kernel Attention . When S ( · ) is Gaussian and the embedding dimension D is at least O ( Nn ) , convergence to global optimal is also ensured for M = { W Q } . Interestingly, under the same conditions of large D , convergence to global optimal is not guaranteed when S ( · ) is Softmax.

In the subsequent sections, we will elaborate on these convergence results in detail, providing a deeper understanding of the nuances in Transformer behavior under varying configurations. To set up our analysis, we introduce λ V as the smallest eigenvalue of W V 0 , λ B as the smallest eigenvalue of B 0 , ¯ λ Q h , ¯ λ K h , ¯ λ V as the largest singular value of matrix W Q h, 0 , W K h, 0 , W V , respectively. We denote ∥ · ∥ 2 as ℓ 2 norm and ∥ · ∥ F as Frobenius norm. Further, we denote σ max ( · ) and σ min ( · ) as the largest and smallest singular value of a matrix, respectively. For any vector v , let min( | v | ) denote the smallest absolute value of vector v .

4.1 Convergence to global optimal

First, we examine the role of W V in the optimization of multi-head attention network structure. Our analysis demonstrates that with the hidden dimension HD ≥ Nn and proper initialization, the global optimal solution of (6) can be found using a vanilla gradient descent algorithm. The initialization requires that the matrix B 0 has full rank. Our first result shows that, overparameterized Transformer can be trained to global optimal solution.

Theorem 1. Consider problem (4) with S ( · ) being instantiated as the Softmax kernel given in (3). Consider the following update for the variable M = { W V } : W V t +1 = W V t -η ∇ W V f ( M t ; X ) , where η > 0 is the stepsize.

Suppose W Q 0 and W K 0 are initialized such that λ B > 0 . Then we have:

where α := ∥ W O ∥ 2 ( λ B ) 2 > 0 ; η > 0 is defined in Appendix 1.3 , and chosen such as ηα < 1 .

Remark 1. The aforementioned theorem focuses on the convergence behavior when only W V is being updated. We further elaborate on the initial conditions ensuring λ B > 0 .

Note that λ B > 0 implies that the objective function f exhibits a landscape that is nearly convex, which is crucial for optimization. By definition, this condition implies that B 0 has full rank, which can be fulfilled by selecting appropriate W Q 0 and W K 0 , plus having large enough embedding size, satisfying D ≥ Nn/H . We refer the readers to Appendix 1 . 3 for the derivation of this condition, which can be guaranteed by random initialization with high probability.

Furthermore, it is important to note that our work aligns with existing literature on the subject of embedding size in Transformer models. For example, in [Bhojanapalli et al., 2020], the authors restrict their focus to the simplified case of N = 1 , H = 1 . They establish the necessary condition for Softmax attention to overcome its low-rank bottleneck, which requires D ≥ n . In our analysis, we derive a similar necessary condition on Transformer model size ( D ≥ n × ( N/H ) ) to guarantee the global convergence when a Transformer model is trained with GD.

In Theorem 1, we have illustrated the case where only updating W V already leads to global convergence. However, in practice, all parameters W V , W Q , W K are updated. This case is more challenging to analyze due to the non-linearity introduced by the Softmax function. Next, we show that a similar result in Theorem 1 still holds when all the parameters are updated simultaneously.

Theorem 2. Consider problem (4), with S ( · ) being instantiated as the Softmax kernel. Consider the GD update where M = { W Q , W K , W V } : Suppose λ B > 0 , and the initialization M 0 satisfy

Then there exists stepsize η > 0 , such that

where β := ∥ W O ∥ 2 ( λ B ) 2 > 0 , and the constants η, ν are defined in Appendix 1.3 .

Remark 2. In the stated theorem, we simplify our analysis by excluding the downstream MLP module in the typical Transformer model, since it is easy to combine the model in Equation (2) with downstream MLP layers. Further, it can be directly showed that the Transformer with MLP will lead to the same convergence rate of the optimization problem as updating W Q , W K , W V only. To illustrate this, consider the following Transformer model:

where W l ∈ R n l -1 × n l , and n 0 = d O . Based on the Transformer model defined in Equation (10), we have the following corollary.

Corollary 1. Consider problem min M 1 2 ∥ G ( M ; X ) -y ∥ 2 , with G ( · ) being defined in Equation (10) and S ( · ) being instantiated as the Softmax kernel. Suppose that the MLP module satisfies:

Consider the following GD update (where M = { W Q , W K , W V , W 1 , · · · , W L } ): Suppose λ B > 0 . Then, there exists a step size η > 0 and initialization weight M 0 , such that the loss function linearly converges to 0 .

Remark 3. The above theorem and corollary describe the global convergence guarantee when W Q , W K and W V are updated. This is in line with the insights gained from Theorem 1. However, the conditions for initialization are more stringent, and the optimization landscape becomes inherently more complex due to the involvement of the Softmax attention through W Q and W K .

To ensure the initial condition 8, we have two options: 1) Initializing M 0 such that ∥ MH ( M 0 ; X ) -y ∥ F is small, which implies that the optimization starts in a region close to the global optimal solution and that the initial weight is close to the global optimal solution; 2) Balancing between W O and W V , in the sense that ∥ W O ∥ 2 is large and ¯ λ V is small. For a detailed account of these initialization strategies, please refer to Appendix 1.3 .

Finally, we need to point out that for Transformers with Gaussian kernel attention, we can derive similar convergence results as long as the attention kernel maintains full rank and weights are initialized appropriately. Here we do not include the theoretical statement since it is similar to the result for Softmax attention.

4.2 Softmax vs Gaussian kernel: Softmax attention Transformers may exhibit slower convergence.

In the previous section, we explored the global convergence of training Transformer models. However, from Theorem 2, it was not clear what roles do matrices W Q and W K play in the entire convergence process, since Theorem 1 indicates that optimizing W V alone already ensures the desired convergence. Nevertheless, it is the matrices W K and W Q that truly represent the power of a Transformer model, because they are used to extract token correlations.

To study how well a Transformer model can extract the token correlation, in this section, we will study the GD dynamics for Transformer models, where only W K and W Q are optimized (while fixing W V ). If optimizing these two parameters alone can still achieve zero training loss, then we claim that the input token correlation can be optimally extracted by the Transformer model.

4.2.1 Notations

To begin our study, let us define that Gaussian kernel to be an n × n matrix, where its k -th row and j -th column of is given by:

Since the training dynamics/gradients of variables W Q and W K have the same property in (3) and (11), we will only concentrate on optimizing W Q .

With some abuse of notation, define a matrix C for Softmax attention and Gaussian kernel attention, respectively. Softmax attention: C ih := X i W Q h ( X i W K h ) ⊤ √ d ∈ R n × n .

Gaussian kernel attention: C ih ∈ R n × n ; ( C ih ) kj = - ∥ X ik · W Q h -X ij · W K h ∥ 2 2 √ d .

For both Softmax attention and Gaussian kernel attention:

Using the above notation, the activation function S ( · ) in (3) and (11) can be related to the matrices C 's in the following manner:

Softmax attention : S ih = Softmax( C ih ) , Gaussian attention : ( S ih ) kj = exp ( ( C ih ) kj ) .

Additionally, note that C is a function of variables M . Therefore we will sometimes use C ( M ) when we need to emphasize the dependency of C on M .

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 W Q is updated), while under the same set of conditions, but with Softmax kernel attention, GD fails.

Theorem 3. Solve problem (4) with the following GD update (with M = { W Q } ): W Q t +1 = W Q t -η ∇ W Q f ( M t ; X ) . Suppose δ h := σ min ( ∂C ( M 0 ) ∂W Q h ) > 0 , ∀ h ∈ [1 , 2 , · · · , H ] , and the initialization condition further satisfies

(1) When S ( · ) is a Gaussian kernel function, there exists a stepsize η and a positive constant γ , such

ν ′ is defined in Appendix 1 . 5 . that where γ, η are defined in Appendix 1.5 .

(2) When S ( · ) is a Softmax function, suppose W Q t is bounded during the training phase, then there exists stepsize η , such that

where η ′ is defined in Appendix 1.5 .

Remark 4. First, it's important to note that the parameter size must satisfy Dd ≥ Nn 2 for δ > 0 to hold. It is crucial to emphasize the fundamental distinction in convergence outcomes between Transformers employing Gaussian kernel attention and those utilizing Softmax attention under these conditions. With equivalent initialization conditions, training Transformers equipped with Gaussian kernel attention achieves global convergence using gradient descent (GD). Second, it is essential to emphasize that the dimension size Dd ≥ Nn 2 is similar to the findings of works that have analyzed the convergence performance of over-parameterized neural networks Allen-Zhu et al. [2019], Du et al. [2019]. The total number of samples, consisting of N samples each with n tokens, can be calculated as Nn . Meanwhile, the total feature dimension is Dd . The inequality implies that the width of the

parameters is at least O ( N ) , a relationship also illustrated in Nguyen and Mondelli [2020]. The proof consists of two basic steps. The first step is to derive the closed form gradient of loss function over variable W Q . 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 Q 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 no 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 3, yet fails to satisfy the PL condition. Therefore, in this case, GD leads to vanishing gradients without being able to find a global optimal solution. The details of this specific example are provided below.

Example: Consider Transformer with Softmax attention, and N = 1 , n = 2 , H = 1 . Let us first write down the close form of the gradient over W Q 1 :

Next, we show there exists W O , W V , X 1 , W Q 1 , 0 , W K 1 , 0 such that the loss function is non-zero with Equation (12) satisfied, while

Denote L := ∂f ( M 0 ; X 1 ) ∂ MH ( M 0 ; X 1 ) ( W O ) ⊤ ( X 1 W V 0 ) ⊤ ∈ R 2 × 2 . ∂f ( M 0 ; X 1 ) ∂C 11 can be expressed as follows:/

Next, we will give the value of W O , W V 0 to show the case where GD leads to vanishing gradient.

easy to show that there exists W Q 1 and W K 1 such that Equation (12) holds. Further, it is easy to verify that for this scenario, the following holds:

(

M

0

∂C

Next, we can easily deduce that

;

X

11

=

(

M

0

∂C

;

X

11

= 0

. Similarly, we can demon-

̸

11

12

if y 1 satisfies that ∂f ( M 0 ; X 1 ) ∂ MH ( M 0 ; X 1 ) = 0 , it follows f ( M 0 ; X 1 ) = 0 , which means M 0 is not global optimal solution.

5 Experiment: Softmax v.s. Gaussian

In this section, we present numerical results to illustrate the behaviors of Transformers models with Softmax attention and Gaussian kernel attention across various tasks.

5.1 Dataset

We investigate two distinct tasks: Text Classification using the IMDb review dataset [Maas et al., 2011] and Pathfinder [Linsley et al., 2018]. While both tasks involve processing

TestAccGaussianvsSoftmax

TestLossGaussianvsSoftmax

0.750

0.650

0.725

0.625

0.700

Test Acc

0.600

SS07

0675

Test I

0990

0.550

0.625

0.525

0.600

0.500

Softmax

0.575

Softmax

0

2500

5000

7500

10000

Epoch

12500

15000

17500

20000

22500

25000

0.550

0

2500

5000

7500

10000

Epoch

12500

15000

17500

20000

22500

25000

long sequences, they exhibit different characteristics. Text Classification is a well-known NLP task that focuses on discerning relationships among token embeddings, while the Pathfinder task prioritizes capturing spatial information within the input pixels.

̸

1

)

)

(

∂f

1

)

)

(

∂f

5.2 Model and Experiment Method

We follow the experiment setting in [Chen et al., 2021]. For both tasks, we employ a 2-layer Transformer model with the following specifications: embedding dimension D = 64 , hidden dimension d = 128 , and number of attention heads H = 2 . To align the model with the classification task, we use an additional mean pooling layer as the final layer. We determine the batch size based on available memory constraints. Specifically, we set a batch size of 16 for the Text Classifi-

TestAccGaussianvsSoftmax

TestLossGaussianvsSoftmax

0.75

0.675

0.650

0.70

Test Acc

0.625

Test Loss

0.65

0.600

0.575

0.60

0.550

0.525

0.55

Softmax

Softmax

0.500

2500

5000

7500

10000

Epoch

12500

15000

17500

20000

0.50

2500

5000

7500

10000

Epoch

12500

15000

17500

20000

cation task with a learning rate of 1 × 10 -4 , and a batch size of 128 for the Pathfinder task with a learning rate of 2 × 10 -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.

Text Classfication, Softmax

Text Classfication, Gaussian kernel

Pathfinder,Softmax

Pathfinder, Gaussian kernel

50

50

50

50

40

40

40

40

0

10

20

30

40

50

0

10

20

30

0

10

20

30

40

50

0

10

20

30

0

10

20

30

40

50

0

10

20

30

0

10

20

30

40

50

0

10

20

30

In our second experiment, the training process consists of two stages: In the first stage, we train the Transformer model equipped with Softmax attention (defined in Equation (3)) for 8,000 steps. In the second stage, we continue training from the pre-trained model for an additional 500 steps, with the option of using either Softmax or Gaussian kernel. To explore the optimization landscape around the trained model, we employed a technique inspired by Li et al. [2018]. We select two parameter directions, specifically the W Q and W K matrices in the first Transformer layer. These two directions, denoted as d 1 , d 2 , are centered at the trained model M , and represent the parameter space of W Q , W K , respectively. We evaluate the loss function on the set { M +0 . 02( r -25) d 1 + 0 . 02( s -25) d 2 } , where r, s ∈ [1 , 2 , · · · , 50] . The above set is the neighborhood of the trained model M , and we chose the evaluation stepsize as 0 . 02 along the two directions d 1 , d 2 , with the total steps limit as 100 . Within this parameter space, we plot a 3-D surface representing the landscape around the trained model.

5.3 Results

5.3.1 Test Loss & Accuracy Curve comparison

To begin with, we present some observations in our first experiment. We plot the test performance of these two tasks on Transformers with two different types of attention. From Fig 2 and Fig 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 in Fig 3. The test accuracy has a significantly higher variance at the same training epoch. Further, the worst test accuracy after 20 , 000 epochs is around 0 . 58 for the Softmax attention Transformer, compared with 0 . 62 for the Gaussian kernel Transformer. These observations align with the experiment results in [Chen et al., 2021] and [Tay et al., 2020], where Transformers with different attention kernels are trained with the same model size and learning rate, while Softmax attention Transformers show instability in a few tasks.

5.3.2 Optimization Landscape Comparison

In Figure 4, we present a comparison of the optimization landscape between Transformers with Softmax and Gaussian kernel attention. Notably, we observe distinct differences in the training landscapes of these two attention types for both tasks. We follow the visualization method described in Section 5.2. We conduct a visualization of the optimization landscape around the trained models after a two-stage training process, with identical learning rates, network sizes, and training epochs. Keeping all other factors consistent, the disparity in the landscape provides a direct representation of the difference in the attention structure during the optimization procedure. With Softmax attention, the landscape appears more complicated compared with Gaussian kernel attention. This complexity can be interpreted as the presence of a greater number of local optima in the optimization landscape, suggesting that Transformers utilizing Softmax attention may encounter more challenges in reaching global optimal solutions. In contrast, the landscape with the Gaussian kernel is flatter. This observation aligns with our earlier findings in Figure 2 and Figure 3, where Softmax attention exhibited certain convergence issues. These observations also provide empirical evidence supporting our Theorem 3, which reflects in a slightly different perspective the complicated optimization landscape within the Softmax kernel.

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 a nuanced understanding of the advantages and disadvantages of using classical Softmax attention in Transformers. We find that while shallow Softmax attention Transformers can achieve global convergence with overparameterization, there are scenarios where this attention structure can lead to local solutions. However, those issues can be mitigated by the Gaussian kernel-based attention. In our work, we need strong initialization and large embedding size, i.e, HD ≥ Nn to obtain the global convergence, which exhibits a gap towards real case. In the future work, we will investigate how to relax the assumptions.

7 Acknowledgment

The work of B. Song was partially done while interning at Amazon Web Services. M. Hong holds concurrent appointments as an Amazon Scholar and as a faculty at the University of 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 W911NF2310315.

References

Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via over-parameterization. In International conference on machine learning , pages 242-252. PMLR, 2019.

Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

Bhojanapalli, C. Yun, A. S. Rawat, S. Reddi, and S. Kumar. Low-rank bottleneck in multi-head attention models. In International conference on machine learning , pages 864-873. PMLR, 2020.

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 in neural information processing systems , 33:1877-1901, 2020.

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 , 34:2122-2135, 2021.

Danilova, P. Dvurechensky, A. Gasnikov, E. Gorbunov, S. Guminov, D. Kamzolov, and I. Shibaev. Recent theoretical advances in non-convex 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:1810.04805 , 2018.

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

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-enhanced bert with disentangled attention. arXiv preprint arXiv:2006.03654 , 2020.

S. Huang, F. Perez, J. Ba, and M. Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learning , pages 4475-4483. PMLR, 2020.

Huang, Y. Cheng, and Y. Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249 , 2023.

Jain, P. Kar, et al. Non-convex optimization for machine learning. Foundations and Trends® in Machine Learning , 10(3-4):142-363, 2017.

Jin, P. 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 (ICML 2000) , pages 1207-1216, Stanford, CA, 2000. Morgan Kaufmann.

Li, G. Wang, and J. Ding. Provable identifiability of two-layer relu neural networks via lasso regularization. IEEE Transactions on Information Theory , 2023.

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.

Linsley, J. Kim, V. Veerabadran, C. Windolf, and T. Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. Advances in neural information processing systems , 31, 2018.

Liu, X. Liu, J. Gao, W. Chen, and J. Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08249 , 2020.

Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

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.

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:11961-11972, 2020.

Noci, S. Anagnostidis, L. 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 , 35:27198-27211, 2022.

Pan and Y. Li. Toward understanding why adam converges faster than sgd for transformers. arXiv preprint arXiv:2306.00204 , 2023.

Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 , 2020.

Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:2011.04006 , 2020.

Tian, Y. Wang, B. Chen, and S. Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv preprint arXiv:2305.16380 , 2023.

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Wu, F. Liu, G. Chrysos, and V. Cevher. On the convergence of encoder-only shallow transformers. Advances in Neural Information Processing Systems , 36, 2024.

Zhang, S. Frei, and P. L. Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

1 Appendix

1.1 Notations

Recall that we have defined the structure of a single Transformer model in Equation (1) and Equation (5). We will further define a few notations before we introduce a few useful lemmas that are needed in our proof.

(1) Operator: Denote vec( · ) as the vectorization operator on a matrix; ⊗ as Kronecker product operator; ⊙ as the element product. Denote Υ( · ) as a matrix operator, such that for any matrix X without zero element

(2) Matrix: Denote I as the identity matrix. Define matrix E and E as following:

Define matrix P h as following: P h = ( . . . , E h n × n , . . . ) , h = 1 , · · · , H .

(3) Matrix in Transformer: Define the following matrix C related to the attention layer

Softmax kernel:

Gaussian kernel:

Define matrix V ′ i for each data X i :

Next, let us introduce several useful lemma which leads to Theorem 2:

1.2 Lemmas of Theorem 2

Lemma 1.

Remark: The above lemma derives the closed form of the gradient of objective over W V , W Q . Notice that we can derive the derivative of W K in the same way as W Q due to symmetry, so we do not include the derivation here. Some of the lemmas here refers https://say-hello2y.github.io/202209-07/attention-gradient

Lemma 2. Consider updating W Q , W K , W V at iteration t . Suppose σ max ( W Q ) , σ max ( W K ) , σ max ( W V ) are bounded during in the optimization phase, then we have the following conclusion:

Lemma 3. Consider updating W Q , W K , W V at iteration t . Suppose σ max ( W Q ) , σ max ( W K ) , σ max ( W V ) are bounded during in the optimization phase, then we have the following conclusion:

Lemma 4. Let f : R n → R be a second order differentiable function. Let x, y ∈ R n be given, and assume that ∥∇ f ( z ) -∇ f ( x ) ∥ 2 ≤ C ∥ z -x ∥ 2 for every z = x + t ( y -x ) with t ∈ [0 , 1] . Then,

Lemma 5. For matrix A ∈ R k × l , B ∈ R l × m , C ∈ R m × n .

1.3 Proof of Theorem 2

Proof Sketch of Theorem 2:

The main idea of the proof follows from [Nguyen and Mondelli, 2020]. Let us first recall a few notations. ¯ λ V := 2 3 ( 1+ σ max ( W V 0 ) ) , λ B := σ min ( B 0 ) . Using GD update rule, we aim to iteratively show

Denote µ := 1 4 ( λ B ) 2 ∥ W O ∥ 2 2 . Let us discuss about the value of µ . We know W O ∈ R Hd × 1 , B ⊤ 0 ∈ R HD × Nn . We require µ > 0 , i.e, λ B > 0 , which implies B has full row rank. For simplicity, let us consider the H = 1 case. Recall the definition of B :

Suppose we initialize W Q 1 , W K 1 such that each S i 1 ∈ R n × n is full rank, then we can easily show that rank( S i 1 X i ) = n if X i has full row rank. Suppose embedding dimension D is large, with certain assumption on X , we can show B has full row rank. For example, if each X i follows standard Gaussian distribution with D >> N , then rank( B ) = Nn with probability 1 if we initialize W Q 1 , W K 1 such that S i 1 is full rank.

Further, let us assume that H ∑ h =1 ( ¯ λ Q h ) 2 > 1 , H ∑ h =1 ( ¯ λ K h ) 2 > 1 , and initialization condition satisfies:

Remark 5. The initialization condition can be satisfied if ∥ W O ∥ 2 is large and σ max ( W V ) is small. ν in Equation (8) is 1 54 .

It is clear that Equation (34) holds when t = 0 . Suppose it holds at iteration t , we prove it holds at iteration t +1 .

where (i) uses the triangle inequality; (ii) plugs in the expression of ∇ W V f ( M t ; X ) and uses the Cauchy-Schwartz inequality; (iii) is because we assume the loss function f ( · ) linearly decreases until t -th iteration. Let u = √ 1 -ηµ . So we have

where (i) is because each element in S r,i has magnitude at most 1 and ∥ S r,i ∥ F ≤ n √ H , then by Cuachy-Schwartz inequality, we have ∥ B ∥ r ≤ √ HN ∥ X ∥ F ; (ii) is due to the initialization condition. Then by Weyl's inequality, there is

Similarly, let us derive the upper bound for σ max ( W Q h,r ) .

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

Then by Weyl's inequality, there is

Now we aim to bound the eigenvalues of B r +1 .

where (i) and (ii) uses triangle inequality and Cauchy-Schwartz inequality; (iii) comes from Lemma 2 (5); (iv) uses Lemma 2 and Cauchy-Schwartz inequality; (v) comes from Cauchy-Schwartz inequality. Together with our initialization condition, we have

where (i) comes from the initialization condition 35. By Weyl's inequality, we can derive the bound for the singular values of B t :

The final step is to show the last inequality holds. Since we have already showed σ max ( W Q h ) , σ max ( W K h ) , σ max ( W V h ) are bounded, by Lemma 3 (2) we can conclude that:

Thus by Lemma 4, we choose η < 1 2 G , then the following hold true:

where (i) uses Lemma 4; (ii) is because we set η < 1 2 G ; (iii) only considers the gradient over W V ; (iv) plugs in the closed form gradient in Lemma 1; (v) uses the property of smallest singular value and induction assumption; (vi) comes from the definition of µ .

1.4 Lemma for Theorem 3

The following lemmas all consider the Transformers with Gaussian kernel attention 11.

Lemma 6.

Lemma 7. Consider updating W Q , W K , W V at iteration t . Suppose σ max ( W Q ) , σ max ( W K ) , σ max ( W V ) are bounded during in the optimization phase, then we have the following conclusion:

where min | V ′ W O | is the smallest absolute value of each element in vector V ′ W O ; min S is the smallest element in matrix S .

Lemma 8. Consider updating W Q , W K , W V at iteration t . Suppose σ max ( W Q ) , σ max ( W K ) , σ max ( W V ) are bounded during in the optimization phase, then we have the following conclusion:

1.5 Proof Sketch of Theorem 3.

(1)Using GD update rule, we aim to iteratively show

̸

Denote γ := 1 2 δ 2 κ 2 ( min ∣ ∣ V ′ W O ∣ ∣ ) 2 . Let us discuss about the value of γ . We know W O ∈ R Hd V × 1 , B ⊤ 0 ∈ R HD × Nn , where Hd > 1 , HD > Nn . We require γ > 0 , i.e, δ > 0 , κ > 0 , min ∣ ∣ V ′ W O ∣ ∣ > 0 . It is clear that κ > 0 can hold as long as W Q h is bounded. And it is easy to show that if X i = 0 , we can always choose W V and W O , such that min ∣ ∣ V ′ W O ∣ ∣ > 0 . Since ∂C h ( M ) ∂W Q h ∈ R Nn 2 × Dd , suppose we initialize W Q h , W K h such that rank( ∂C h ( M 0 ) ∂W Q h ) = Nn 2 , then we have σ min ( ∂C h ( M 0 ) ∂W Q h ) ≥ δ for some positive constant δ . Further, we assume the initialization condition satisfies:

Remark 6. The initialization condition can be satisfied if ∥ W O ∥ 2 is large and σ max ( W V ) is small. ν ′ in Equation (12) is 1 8 .

Similar to the proof of Theorem 2, we use induction to prove the theorem. Equation (52) holds when t = 0 . Suppose it holds at iteration t , we prove it holds at iteration t +1 .

where (i) uses triangle inequality; (ii) comes from Lemma 7 and Cauchy-Schwartz inequality; (iii) is from the induction assumption that loss function f ( · ) linearly decreases until t -th iteration. Let u = √ 1 -ηγ . So we have

where (i) comes from the initialization condition. Then by Weyl's inequality, there is

where (i) uses triangle inequality; (ii) applies Lemma 7 (2) and Cauchy-Schwartz inequality; (iii) uses Lemma 7 (4); (iv) applies the induction assumption that the loss function f ( · ) linearly decreases until t -th iteration; (v) comes from the initialization condition. Then by Weyl's inequality, there is

For each element in S ih , we have close form

Since we have already showed that σ max ( W Q h,r ) ≤ 3 2 ¯ λ Q h , it follows directly each element in matrix S t is lower bounded by some constant κ for any t . Now we derive the expression of κ :

where (i) uses Cauchy-Schwartz inequality; (ii) applies the induction assumption σ max ( W Q h,t ) ≤ 3 2 ¯ λ Q h and property of singular value; (iii) is because d ≥ 1 . Thus, we have min S t ≥ κ . Finally, we aim to show f ( M t +1 ; X ) ≤ (1 -ηγ ) f ( M t , X ) . By Lemma 8, since we have showed that σ max ( W Q h ) is bounded, we can directly derive that

Finally, by Lemma 4, choose η < 1 2 G ′ , we have the following holds:

where (i) uses Lemma 4 (2); (ii) is because we choose η < 1 2 G ′ ; (iii) writes down the expression of gradient according to chain rule in Lemma 6; (iv) uses the induction assumption σ max ( ∂C h ( M t +1 ) ∂W Q h ) ≥ 1 2 δ and property of singular value; (v) uses Lemma 6 (4); (vi) comes from Lemma 7 (3); (vii) uses the definition of γ .

(2)Next, we show the convergence result for Transformer with Softmax kernel with only W Q updated. Since we assume parameters are all bounded during optimization phase, by Lemma 8, we can easily show that there exists constant G ′ (see xx for details), such that

Then by Lemma 4, choose η ′ < 1 2 G ′ we have

1.6 Proof of Lemma in Section 1.2

Proof of Lemma 2 (1).

Proof. Step 1: When W Q , W K are updated, we aim to prove

above two steps, we can derive the bound in Equation (27).

Proof of Step 1: First, we can write down the closed form of the differential of S i :

We reorganize the terms on the right side of Equation (57), we have the following equation:

Since C i = [ C i 1 , · · · , C iH ] , we will investigate each C ih , h = 1 , 2 , · · · , H . We focus on the term d ( C i ) -Υ((exp C i ) E ) ⊙ (exp( C i ) ⊙ d ( C i )) E in Equation (58). We write down the close form of the element in the k -th row and j -th column:

̸

̸

where (i) is expand the closed form of Equation (59); (ii) uses the Cauchy-Schwartz inequality; (iii) is because each element in the square root in (ii) is upper bounded by 1 . With Equation (61), we can easily show

Since every element in S i has magnitude less than 1 , we have

where (i) is from Cauchy-Schawatz inequality.

Proof of Step 2: We aim to show ∥ d ( C i ) ∥ F ≤ n √ d ∥ X i ∥ 2 F √ ∑ H h =1 σ 2 max ( W K h ) · ∥ ∥ d ( W Q )∥ ∥ F . Similarly, we investigate ∥ d ( C ih ) ∥ F , h = 1 , 2 , · · · , H. We have

Then plug the above inequality to Equation (63), we can derive

Thus by Cauchy-Schwartz inequality, it is easy to show

Proof of Lemma 2 (4).

Proof. We first write down the close form of gradient of f ( · ) over W Q h by Lemma 1, and derive the upper bound of the norm of the gradient.

Denote R i = ( MH ( M ; X i ) -y i ) ( W O ) ⊤ ( V ′ i ) ⊤ , R i = [ R i 1 , · · · , R iH ] . Write down the close form of the element in the k -th row and j -th column:

̸

̸

where (1) is due to the Cauchy-Schwartz inequality; (ii) is because each element within the squre root term in (i) has magnitude at most 1 . Thus, we can further derive

where (i) if from the bound in Equation (71); (ii) comes from Cauchy-Schwatz inwquality; (iii) uses the property of Frobenious norm. Thus, by Cauchy-Schwartz inequality, we can derive the upper bound for ∥ ∥

So plug the above inequality into Equation (69), we can derive the upper bound for ∥ ∥ ∥ ∂f ( M ; X i ) ∂W Q ∥ ∥ ∥ :

h F

Proof of Lemma 3 (1). By Mean Value Theorem and Cauchy-Schwartz inequality,

where M ′ t is between M t and M t +1 . We can derive the upper bound of the norm of ∇ W V f ( M ; X i ) :

Therefore, together with Equation (73), we have

Proof of Lemma 3 (2).

Proof. By triangle inequality, we have

Step 1: Derive upper bound for

∥∇ W Q f ( M t +1 ; X i )) -∇ W Q f ( M t ; X i )) ∥ F = ∥ vec( ∇ W Q f ( M t +1 ; X i )) -vec( ∇ W Q f ( M t ; X i )) ∥ 2 . First, we give the vectorized expression of ∇ W Q f ( M t ; X i ) . Recall we denote U i = ( ( MH ( M ; X i ) -y i ) ( W O ) ⊤ ( V ′ i ) ⊤ ) ⊙ S i . By Lemma 1, we can derive the close form of vec( ∇ W Q f ( M t ; X i )) :

where (i) uses the Lemma 1; (ii) and (iii) comes from the property of vectorization in Lemma 5; (vi) uses the definition of S i ; (v) gives an equivalent expression of vec( U i ) ; (vi) reorganizies (v). Further, it is easy to verify that:

Next, let us derive upper bound for ∥∇ W Q f ( M t +1 ; X i ) -∇ W Q f ( M t +1 ; X i ) ∥ F .

∥∇

(

i

)

∥

=

W

∥

=

Q

(

∥

∥

I

f

n

(

I

(

ii

)

≤ ∥

(

2

n

I

2

M

H

H

2

n

H

t

+1

-

(

;

X

E

-

-

(

(

i

⊗

E

E

⊗

⊗

)

I

-∇

n

I

)

)(

n

I

)

n

)(

)

∥

F

W

Q

f

(

M

U

i,t

t

+1

+1

U

i,t

(

∥

)

+1

)

U

i,t

;

X

⊙

i

)

∥

F

⊙

+1

S

i,t

S

-

U

+1

i,t i,t

+1

)

)

-

)

∥

F

-

+

∥

U

U

U

i,t i,t

t

∥

F

)

∥

)

⊙

⊙

S

i,t

S

i,t

S

)

i,t

+1

-

+1

S

i,t

)∥

F

∥

) + vec(

U

)

∥

F

i,t

)

⊙

S

i,t

+1

)

-

U

i,t

)

⊙

where (i) plugs in the expression in Equation (79); (ii) uses the fact that each element in S i,t +1 has magnitude at most 1 , and Cauchy-Schwartz inequality; (iii) comes from the definition of I , E and ¯ R ; (iv) uses the definition of U i,t ; (v) is because triangle inequality; (vi) uses the fact that each element in S i,t +1 has magnitude at most 1 , and Cauchy-Schwartz inequality. Next, we aim to derive upper bound of ∥ R i,t +1 -R i,t ∥ F in Equation (81).

where (i) is because of the triangle inequality; (ii) uses the definition of Z i in Equation (74), CauchySchwartz inequality and triangle inequality; (iii) uses the Cauchy-Schwartz inequality; (iv) reorganizes the terms in (iii). Plug Equation (82) into Equation (81), we can finally derive the bound for ∥∇ W Q f ( M t +1 ; X i ) -∇ W Q f ( M t +1 ; X i ) ∥ F .

where (i) is from Equation (81); (ii) uses the definition of ¯ R in Equation (80); (iii) comes from Lemma 3 (3). Since W Q and W K are symmetric in the Transormer structure, similarly, we can derive L K i .

Step 2: In this step, we aim to derive bound for ∥∇ W v f ( M t +1 ; X i ) -∇ W v f ( M t ; X i ) ∥ F .

where (i) is from Lemma 1 (1); (ii) uses triangle inequality; (iii) uses Cauchy-Schwartz inequality; (iv) comes from the definition of B i,t , , Z i (in Equation (74)), Cauchy-Schwartz inequality and triangle inequality; (v) comes from Lemma 2 (3) and Cauchy-Schwartz inequality; (vi) reorganizes (v).

Now we combine the result in Step 1 and Step 2 , and plug into Equation (78), we can finally derive

1.7 Proof of Lemma in Section 1.4

Proof. Proof of Lemma 6 (1) : We consider the differential of the element in the k -th row and j -th column. First, let us write down the closed form of each element:

Next, we consider the differential of each element over W Q h :

where o ( d ( W Q h )) denotes the higher order of d ( W Q h ) . Leave out the higher order differential term, we derive

Proof. Proof of Lemma 7 (2) : First, let us write down the closed form of ∂ ( C ih ) kj ∂W Q h . We have

Thus, we can derive upper bound for ∥ ∥ ∥ d ( ∂ ( C ih ) kj ∂W Q h )∥ ∥ ∥ F :

Thus, we have the following:

Proof. Proof of Lemma 7 (3) :

Proof. Proof of Lemma 7 (4) :

Proof. Proof of Lemma 8 (1): The proof is similar to the proof of Lemma 3 (1). So we do not include the details here. We can similarly derive

Proof. Proof of Lemma 8 (2) : By triangle inequality, we have

Step 1: Derive upper bound for

First, we give the vectorized expression of ∇ W Q f ( M t ; X i ) . Recall we denote U i = ( ( MH ( M ; X i ) -y i ) ( W O ) ⊤ ( V ′ i ) ⊤ ) ⊙ S i . By Lemma 6, we can derive the close form of vec( ∇ W Q f ( M t ; X i )) :

Further, recall we have defined ¯ R and the following inequality holds:

Next, we aim to derive upper bound of ∥ R i,t +1 -R i,t ∥ F in Equation (92). Similar to the derivation in Equation (81), we can derive

Plug Equation (82) into Equation (92), we can finally derive the bound for ∥∇ W Q f ( M t +1 ; X i ) -∇ W Q f ( M t +1 ; X i ) ∥ F .

and plug into Equation (78), we can finally derive

2 NeurIPS paper checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer:[Yes]

Justification: See Theorem 1,2,3

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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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 are not discussed in the paper.

The authors are encouraged to create a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

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

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of 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 a 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 it 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 to 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 a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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 reproducibility. In the case of closed-source models, it 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 to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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 a 1σ error 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, confidence intervals, or statistical significance tests, at least for the experiments that support the 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, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: We do not include the compute resources detail.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The paper has no harm in the research process or negative social impact. The paper is anonymous.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have cited the code framework we use Chen et al. [2021].

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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/model 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 as 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 crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 nor 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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$.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nNot quite. See the weakness section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nMinor 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.\n\nQUESTIONS:\n1.\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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nSee weakness above.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}