Unraveling the Gradient Descent Dynamics of Transformers

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 $\mathcal{O}\big{(}\text{poly}(N)\big{)}$, where $N$ is the sample size. Similarly, (Allen-Zhu et al., 2019) show that ReLU fully-connected networks with at least $\mathcal{O}\big{(}\text{poly}(N)\big{)}$ 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/\delta))$ can guarantee the global convergence with probability at least $\delta$ 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.

A number 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 behaviour of shallow Transformer, which shows the global convergence can be achieved with GD algorithm. However, the focus of our paper is different from Wu et al. (2024). We not only derive the global convergence analysis (Our Theorem 2), but also investigates the role of different variables in optimization.

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

Some recent research has focused on the convergence analysis of Transformer-based models within the in-context learning (ICL) framework. For instance, Huang et al. (2023); Zhang et al. (2023) explores the learning dynamics of a one-layer Transformer with Softmax attention trained via gradient descent to learn linear function classes in-context. However, this line of study primarily addresses the general convergence performance of Transformers within the ICL setting and does not delve into the role of individual variables.

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})\}_{i=1}^{N}$, where $X_{i}\in\mathbb{R}^{n\times D}$, and $y_{i}\in\mathbb{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}\in\mathbb{R}^{D\times d}$ is the query weight matrix and key weight matrix, respectively; $W_{h}^{V}\in\mathbb{R}^{D\times d}$ is the value weight matrix; these matrices are the main optimization variables throughout the paper. Further $W^{O}\in\mathbb{R}^{Hd\times 1}$ is a fixed matrix, representing the weight of the output layer; $H$ is the number of attention heads; $S(\cdot)$ is a kernel function of variables $W^{Q},W^{K}$ and input $X_{i}$. Attention$(\cdot)$ is the attention head function; MH$(\cdot)$ represents the multi-head attention function. For example, with the Softmax attention (Vaswani et al., 2017), $S(\cdot)$ can be written as:

where for a given $n\times n$ matrix $Z$, $\operatorname{Sofmax}(Z):=[\operatorname{Softmax}(Z_{1}),\cdots,\operatorname{Softmax}(Z_{n})]$. Throughout, let us denote $S(\cdot)_{kj}$ as the element of $k$-th row and $j$-th column in matrix $S(\cdot)$. Let $X_{ik\cdot}\in\mathbb{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\left(W_{h}^{Q},W_{h}^{K};X_{i}\right)$.

Based on the above Transformer model, we consider minimizing the following empirical $\ell_{2}$ loss function for the entire data set $\{X_{i},y_{i}\}_{i=1}^{N}$:

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})\}_{i=1}^{N}$. Towards this end, let $X\in\mathbb{R}^{Nn\times D}$ denote the column-stacked matrix of each single data $X_{i}$. Similarly, define the stacked label $y\in\mathbb{R}^{Nn}$. Then we can define:

$i=1,2,\cdots,N,h=1,2,\cdots,H$ for simplicity.

Let $B:=\begin{pmatrix}S_{11}X_{1}&\cdots&S_{1H}X_{1}\\ \cdots&\cdots&\cdots\\ S_{N1}X_{N}&\cdots&S_{NH}X_{N}\end{pmatrix}$, and $W^{V}:=\operatorname{diag}(W^{V}_{1},\cdots,W^{V}_{H})\in\mathbb{R}^{HD\times 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^{K}_{t},W^{V}_{t}\}$. Similarly, we denote $B_{t}$ as the matrix $B$ at $t$-th iteration.

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

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

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

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

Common convergence conditions with Softmax Attention: When the activation function $S(\cdot)$ is either the Softmax or Gaussian function, and the embedding dimension $D$ is at least $\mathcal{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(\cdot)$ is Gaussian and the embedding dimension $D$ is at least $\mathcal{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(\cdot)$ 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 $\underline{\lambda}^{V}$ as the smallest eigenvalue of $W^{V}_{0}$, $\underline{\lambda}^{B}$ as the smallest eigenvalue of $B_{0}$, $\bar{\lambda}_{h}^{Q},\bar{\lambda}_{h}^{K},\bar{\lambda}^{V}$ as the largest singular value of matrix $W_{h,0}^{Q},W_{h,0}^{K},W^{V}$, respectively. We denote $\|\cdot\|_{2}$ as $\ell_{2}$ norm and $\|\cdot\|_{F}$ as Frobenius norm. Further, we denote $\sigma_{\max}(\cdot)$ and $\sigma_{\min}(\cdot)$ 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$.

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

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\geq 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.

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.