Unraveling the Gradient Descent Dynamics of Transformers

Your comment: The new insight we can gain compared to Wu et al [2024]; the fixed $W^O$ in optimization.

Response: See the general rebuttal at the top of the page Comment 1.

Your comment: The comparison between the Gaussian and Softmax kernels only updates $W^Q$.

Response: We'd like to clarify that we also studied the role of the variable $W^V$ in Gaussian and Softmax kernel and found them to be the same. We have commented on the similarity between the two Transformers in Remark 3. Specifically, the conclusions in Theorem 1 and Theorem 2 hold for Transformers with Gaussian attention, and the linear decrease of the loss function also comes from $W^V$ when $W^Q,W^K,W^V$ are updated. Since the conclusion and proof technique are very similar to Theorem 1 and Theorem 2, we did not include the details in the paper. To avoid confusion, we will add a corollary to clearly state the similarity of $W^V$ in both Transformers in our revised version.

Your question: The comparison of assumptions in Wu et al [2024] and our paper.

Response: See the general rebuttal at the top of the page Comment 2.

Your question regarding the importance of analyzing individual variables. Can the results obtained with one variable trainable and others fixed explain their roles when they are simultaneously updated?

Response: The importance of unravelling individual variables: First, unravelling the individual role of each variable helps us understand the optimization bottlenecks in Transformers [2,3]. For example, in [2], it is shown that the rank collapse of the attention kernel can occur when the variable $W^Q,W^K$ are updated, which can lead to failure in optimization. Second, acceleration optimization algorithms can be designed based on the understanding of the individual role of each variable [3,5]. In these works, Nyström method is used in gradient computation to greatly accelerate the training. The algorithm design is specific to the training dynamics of individual variables in Transformers. Third, analyzing the training dynamics of individual variables leads to a new network design that improves the optimization performance [3,4,5]. In these works, attention structures other than Softmax attention are proposed to improve the generalization and computational efficiency.

Simultaneously update: Yes, the property obtained when analyzing a single variable updates does carry over to the situation when this variable is updated to other variables simultaneously. The reason is as follows: Our over-parameterization structure and initialization ensure all the variables always stay in a locally near-convex region. Within this region, the optimization landscape is relatively smooth, and the partial gradient of any single variable does not change much when other variables are updated, thus analyzing the gradients of different variables separately is able to explain the role of each variable when they are updating simultaneously.

Your question regarding the existence of global optimum in line 286.

Response: Yes, it is implicitly assumed that the global optimum solution has zero loss. This is true because of the over-parameterization of the network, where the network size is $\Omega(Nn)$ ($N$ is the sample size). When there are more parameters than samples, there exists a solution that can completely fit all the samples. Further, in some other works with similar settings, the existence of global optimal solution can also be derived from the limit of a Cauchy sequence (See [8] Appendix B.7). To make the proof more rigorous, we will add a similar derivation of the Cauchy sequence in our revised version.

Your question regarding the statement "transformers perform exceptionally well".

Response: We claim that our result explains why Transformers perform exceptionally well because we derive a global convergence analysis of training Transformers (Theorems 1 and 2). These two theorems show that even the plain gradient decent algorithm can find the global optimal solution with a network size of $\Omega(Nn)$, allowing the training loss to decrease to $0$. This provides theoretical proof for the outstanding training performance of Transformers.

Your question regarding the sentences like "only updating $W^V$ already leads to global convergence".

Response: No, it does not suggest that it is easier to optimize more variables. As we aim to scrutinize the role of each variable in the attention kernel, we found that the global convergence is dominated by $W^V$, which prompted us to later focus on studying the role of $W^Q$ and $W^K$ in optimization. These sentences are meant to emphasize that the role of $W^Q, W^K$ remains unclear if updating $W^V$ alone already leads to global convergence (See Theorem 1,2), i.e, does updating $W^Q,W^K$ contribute to the linear convergence of loss function? Thus, these statements lead to our discussion on $W^Q$ alone in Theorem 3.

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

[2] Noci, Lorenzo, et al. "Signal propagation in transformers: Theoretical perspectives and the role of rank collapse." Advances in Neural Information Processing Systems 35 (2022): 27198-27211.

[3] Chen, Yifan, et al. "Skyformer: Remodel self-attention with gaussian kernel and nystr" om method." Advances in Neural Information Processing Systems 34 (2021): 2122-2135.

[4] Choromanski, Krzysztof, et al. "Rethinking attention with performers." arXiv preprint arXiv:2009.14794 (2020).

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

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

Your comment: The difference in convergence speed in Fig. 2 seems like constant level with same rate.

Response: Thank the reviewer for the comment. We need to clarify that it is reasonable that the difference in convergence rate in Fig. 2 is constant level. Further, we provide empirical reason for the different theoretical convergence rates in Theorem 3 by landscape visualization in Fig. 4.

Frist, Transformers with Gaussian and Softmax attentions theoretically should have the same convergence rate in our experiment setting in Fig. 2. Notice that the algorithm we use in Fig. 2 is updating all the variables in the attention ($W^Q,W^K,W^V$). As discussed in Remark 3, when all the variables are updated, both Transformers could achieve linear convergence rate with appropriate network size and initialization. Thus, the theoretical difference in convergence rate is constant level. However, due to the bottleneck of optimizing $W^Q$ (or $W^K$) in Softmax Transformer as illustrated in Theorem 3, Gaussian Transformer still converges faster than Softmax Transformer when $W^Q, W^K$ are included in optimization.

Second, in terms of the different theoretical convergence rate in Theorem 3 when only $W^Q$ (or $W^K$) is updated in both Transformers, we provide empirical reason for the different convergence rates in Theorem 3 by landscape visualization in Fig. 4. By showing the existence of more local solutions in Transformers with Softmax attention, we provide evidence for the comparison of the convergence rates: Compared to Transformer with Gaussian attention, a Transformer with Softmax attention is more likely to obtain a sublinear convergence to stationary points, not global optimum.

Your comment: The insights of critical proof steps that cause different convergence rates in Theorem 3.

Response: We thank the reviewer for the suggestion. We will add the proof sketch of Theorem 3 in the revised version. Here we will briefly summarize the critical steps.

Step 1: Derive the closed form gradient of loss function over variable $W^Q$. Please see Lemma 1 equation (23),(24) in line 778 for Softmax attention Transformer, and Lemma 6 equation (41), (42) in line 821 for Gaussian attention Transformer. Intuitively, the gradient of Softmax attention is much more complicated than Gaussian attention Transformer, which will lead to a more complicated landscape and more local solutions.

Step 2: Analyze the gradient of Transformers with both kernels with the same initialization Equation (12). For Gaussian attention Transformer, it can be iteratively shown during the gradient descent training: 1) The variable $W^Q$ is bounded; 2) The PL condition holds (i.e, the optimization landscape remains near-convex); 3) The loss function decreases linearly. (Please see Equation (52) in line 829.) For the Softmax attention Transformer, there is no guarantee that the PL condition holds during iterative gradient descent update. To illustrate this claim, we construct a concrete counterexample to show that the PL condition does not hold. (Please see line 284).

Your question: The authors construct a compelling example to illustrate the failure of a Transformer with a Gaussian kernel, but it requires a specific initialization form. To achieve a similar convergence rate for the Softmax kernel, how many additional conditions are required?

Response: Respectfully, we need to clarify that we do not construct an example to illustrate the failure of the Transformer with a Gaussian kernel. Instead, we provide a detailed example to illustrate that the Transformer with Softmax attention converges to a stationary point with a sublinear rate, but fails to achieve linear global convergence. (Please see Equation (14) in Theorem 2 and the failure to satisfy the PL condition in lines 284 and 285).

Additional condition for Softmax Transformer: For a Transformer with Softmax attention, the PL condition or convexity of loss function should be assumed during the whole training phase to achieve the same linear convergence rate as a Transformer with Gaussian attention. This condition is much more stringent than the condition for Gaussian attention, where the PL condition is assumed only at initialization.

Your question: I don’t fully understand why optimizing the query matrix alone leads to similar dynamics as when optimizing both the query and the key matrix. Could you clarify this?

Response: The similarity comes from the symmetry of $W^Q$ and $W^K$ in attention mechanism. The symmetry means the training dynamics/gradients of $W^Q$ and $W^K$ have exactly the same properties. Please refer to Appendix 1.2 Lemma 1 (3). We write down the closed form gradient of $W_h^Q$ over $f$ for the Softmax Transformer. Similarly, we can derive the gradient over $W_h^K$, with the only difference from Lemma 1 (3) is that, the $W_h^K$ term on the right side of the equation is replaced with $W_h^Q$. Thus, the gradients of $W^Q$ and $W^K$ have the same structure. As long as we initialize $W^Q$ and $W^K$ with the same properties, the training dynamics are also the same for both gradients. Intuitively, this stems from the fact that $W^Q$ and $W^K$ are in symmetric positions with respect to calculating their gradients in the attention head, see equation (3).

[1] Huang, Baihe, et al. "Fl-ntk: A neural tangent kernel-based framework for federated learning analysis." International Conference on Machine Learning. PMLR, 2021.

[2] Gao, Tianxiang, et al. "A global convergence theory for deep relu implicit networks via over-parameterization." arXiv preprint arXiv:2110.05645 (2021).

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

Your comment: The empirical evaluation is limited and doesn't fully validate the theoretical claims in practical settings. The experiments use simplified Transformer models on relatively small datasets (IMDB and Pathfinder). It would be more convincing to see results on larger, more complex Transformer architectures and standard NLP benchmarks.

Response: Thanks the reviewer for the comment. First, we need to clarify that the goal of our paper is to investigate the training dynamics and landscape of Transformer model. The main purpose of our experiment is to provide a concrete example to illustrate the existence of local solution in Softmax attention. Similar experiments of shallow networks can be found in [1,2,9], which have settings similar to ours.

We agree that additional experiments on larger models could potentially provide further support for our theory. However, due to computational budget constraints, we are unable to complete these experiments during the rebuttal period. We will add results on these larger models at a later time.

Your comment: The paper's contribution relative to prior work is not clearly articulated. While the authors cite some related papers, they don't adequately explain how their results extend or differ from existing analyses of Transformer convergence, e.g., [4].

Response: We thank the reviewer for the comment. We will add the following detailed comparison with the related works in the revised version:

1. Comparison with [3,4]: These works analyze the convergence of in-context training, where a prompt is contructed with all the training samples and a single test sample. The goal of these works is to achieve the zero test loss (in expectation) by optimizing over the loss function modeled by the propmpt. However, our analysis is based on standard empirical loss minimization, which does not involve any prompt construction.
2. Comparison with [5,6,7,8]: This line of works analyze the global convergence of training over-parameterized fully connected networks. However, Transformer structure and the optimization landscape in our study are more complicated.
3. Comparison with [9]: This work is closely related to our work. However, the primary goal of [9] and our work is different. In [9], the main purpose is to build a convergence theory of shallow Transformer with realistic structure and initialization, but they do not provide roles for different matrices in the convergence. While in our paper, we focus on investigating the role of each variables within the attention mechanism in optimization. Further, we analyze the convergence behavior of different attention kernels, i.e, Softmax, Gaussian, which is not covered by [9].

Your comment: The practical implications of the theoretical results are not sufficiently discussed. It's unclear how the insights about Gaussian attention could be applied to improve real-world Transformer models or training.

Response: We thank the reviewer for the comment. Here we discuss the insights about the benefits of using a Gaussian kernel in training Transformer models. As discussed in [1], Gaussian kernels have a natural interpretation of assigning “attention” to different tokens, similar to Softmax attention. However, the Gaussian kernel performs intrinsic normalization, which allows the attention mechanism to have a more reasonable condition number than Softmax attention. Therefore, designing specific Transformers with Gaussian attention can greatly improve the stability of model training. For example, in [1] Table 1, the kernelized attention (Gaussian kernel) and Skyformer (modified Gaussian kernel) show better performance than Softmax attention. In [10], it is showed that the Gaussian Transfomer is more efficient and generalizes better than Softmax Transfomer.

[1] Chen, Yifan, et al. "Skyformer: Remodel self-attention with gaussian kernel and nystr" om method." Advances in Neural Information Processing Systems 34 (2021): 2122-2135.

[2] Tay, Yi, et al. "Long range arena: A benchmark for efficient transformers." arXiv preprint arXiv:2011.04006 (2020).

[3] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. "Trained transformers learn linear models in-context." arXiv preprint arXiv:2306.09927 (2023).

[4] Huang, Yu, Yuan Cheng, and Yingbin Liang. "In-context convergence of transformers." arXiv preprint arXiv:2310.05249 (2023).

[5] Allen-Zhu, Zeyuan, Yuanzhi Li, and Zhao Song. "A convergence theory for deep learning via over-parameterization." International conference on machine learning. PMLR, 2019.

[6] Du, Simon, et al. "Gradient descent finds global minima of deep neural networks." International conference on machine learning. PMLR, 2019.

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

[8] Allen-Zhu, Zeyuan, Yuanzhi Li, and Zhao Song. "On the convergence rate of training recurrent neural networks." Advances in neural information processing systems 32 (2019).

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

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