Local to Global: Learning Dynamics and Effect of Initialization for Transformers

We thank the reviewer for the helpful feedback and insightful comments. We address the individual questions below:

• Generality of the insights: For the two-state Markov chain, while our analysis capitalizes on the canonical parameters and linear attention, our guidelines and insights from this setting readily generalize to non-linear soft-max attention. For instance, Figure 6 in the paper (Figure 1, 3 and 4 in the rebuttal) demonstrates the effectiveness of our initialization scheme for the full one-layer transformer model with soft-max attention (Lines 60-64). Similarly, the low-rank structure of the parameters, as shown in Figure 5 (Figure 2 in the rebuttal), is concerning the full general model. The lines 280-281 and 255 were meant to convey this message which we are happy to rewrite to make it more clear.

  For multi-state Markov chains, we observe a similar phenomenon where, under standard Gaussian initialization, a single-layer transformer can converge to either local or global minima depending on the Markovian switching. Specifically, in Figure 5 of the rebuttal, we consider a first-order Markov chain with state space $S = \\{ 0, 1, 2, 3, 4 \\}$ where any state $s \\in S$ stays in the same state with probability $1-p$ and switches uniformly random to either of the remaining states with probability $p/4$. This generalizes the binary-state Markov chain in the paper with $p = q$. As illustrated in Figure 5 of the rebuttal, depending on the values of the switching probability $p$, the transformer parameters either converge to local-minima (unigram) or global minima (bi-gram). Further, Figure 6 of the rebuttal also reveals the low-rank structure of the parameters during training for the multi-state setting too. These interesting empirical observations highlight an exciting research opportunity to mathematically characterize the gradient dynamics for the multi-state Markov chains, similar to our two state analysis.

• Low rank manifold: Indeed, using canonical parameterization to characterize the gradient descent dynamics of transformers is a standard tool in the literature. For instance, a recent work [1] employs a similar approach to show that two-layer transformers learn induction heads, assuming specific structure for the attention matrices in both layers, inspired by empirical observations (Figure 2, [1]).

• Limitations: Please note that the attention parameter $a$ is also trained alongside the embedding and weight parameters $(e,w)$, and the corresponding analysis is highlighted in Section 4.1. All the figures in the paper reflect this setting, with no parameters being frozen or omitted.

References

[1] How Transformers Learn Causal Structure with Gradient Descent? https://arxiv.org/pdf/2402.14735

We thank the reviewer for interesting questions and apologize for not being able to accommodate the requested figures in the first page of the PDF. We address the concerns below.

Attention features are indeed learnt: Note that with canonical parameterization, the attention score is given by $`attn`_{n,i} = a \cdot (x_n - \frac{1}{2}) (x_i - \frac{1}{2})$ (Eq. 45 in the Appendix). Since $a$ is a trainable parameter, the attention scores are also learned.

Linear vs. soft-max attention: We are afraid that there has been a misunderstanding about our intialization and the transformer architecture. For all the experiments in the paper, we use the full general 1-layer transformer with non-linear soft-max attention, as described in lines 63-64 of the paper. For the theoretical analysis, we use the linear attention mechanism but with rest of the architecture intact. Using these insights derived from the linear setting, we demonstrate that our initialization scheme performs significantly better than standard Gaussian initialization even on the full soft-max model. This is not due to linear vs. soft-max but rather the actual loss landscape as illustrated in Figure 4 in the paper and theoretically characterized in Theorem 9 of Appendix G. Hence we respectully disagree with the claim that "...our intialization scheme in linear attention prevents the softmax transformers from converging to local minima...".

We do not limit low-rank parameters: Please note that the low-rank parameters in the attention layer are not artificially restricted to be same as $\alpha$. The reason why they are also parameterized as $\alpha$ is because empirical evidence suggests that the transformer parameters always converge to this specific structure. Furthermore, this $\alpha$ turns out to be all $\pm 1$ vector. Hence, even after parameterizing $\alpha$ as a trainable vector in our setting, it turns out that the final loss function $L$ (in line 203) is independent of $\alpha$ (since $\|\alpha \|$ is constant, cf. Appendix B) and only depends on three scalar parameters $(e,w,a)$. Thus without loss of generality, it suffices to analyze these scalars.

Attention features are indeed learned: The definition of $a$ in the paper is somewhat confusing, as there are two different definitions provided. Between lines 111-112, $a=\langle v, \alpha\rangle$ is stated, while on line 200, $a=q^2d^{5/2}/4\langle v, \alpha\rangle$ is given. I believe the definition on lines 111-112 is a typo, but it seems to suggest that only the $V$ layer is being trained. This would imply that $`attn`$ is not learned and that the trainable parameter $a$ comes from $v$. It is a minor issue, since it is equivalent to put this scalar with $V$ or $`attn`$.

Linear vs. Softmax Attention: I am puzzled by the authors' disagreement with my claim. My statement was intended to convey that "your initialization scheme, inspired by the linear model, indeed helps generalize to softmax attention and prevents convergence to local minima." But my concern remains: does this low-rank initialization (not necessarily rank-2) generalize to multi-state Markov models? Even within the extended rebuttal PDF, I don't see this experiment included. Furthermore, according to Figure 5, it is not strictly low-rank (and definitely not rank-1); the strict low-rank property only holds for the binary-state case in the PDF. I suggest that the authors include these details in future versions, and hopefully, all previous insights will generalize to the multi-state settings.

We do not limit low-rank parameters: I fully understand the authors need to make these simplification. But I respectively disagree with the claim that "even after parameterizing $\alpha$ as a trainable vector in our setting, the loss L is constant", since the after training the norm of $\alpha$ will change. Also, it does not necessarily true that $x_n, W_Q,W_K$ all align with $\alpha$ after you train $\alpha$ with some random initialization. And the authors limit low-rank parameters, just not artificially. It is natural and common to simplify the network as a theoretical work but it is indeed a limitation.

I will maintain my score and cannot strongly recommend acceptance.

We thank the reviewer for the helpful feedback and suggestions to improve the paper and the code which we are planning to incorporate in the revised version. We refer to the global response regarding our results across repeated trials for various $(p,q)$ and the measure of low-rankness. We address the individual questions below.

• Our insights carry over to the full general model: While our theoretical analysis uses canonical parameterization, our insights and guidelines extend to the general single-layer transformer with soft-max attention (Lines 60-64). For example, Figure 6 demonstrates the effectiveness of our initialization scheme for the full one-layer transformer model. Similarly, the low-rank structure of the parameters, shown in Figure 5 in the paper (Figure 2 and 6 in the rebuttal PDF), pertains to this general model.

• Canonical parameters exhibit the same dynamics as the full model: We strongly believe that the training dynamics of the canonical model closely resemble those of a full single-layer transformer with rank-1 initialization. Due to the time constraints, we currently lack empirical evidence for this. However, we will share this comparison during the discussion phase.

• $W_1$ and $W_2$ are not constants: We are afraid that there has been a misunderstanding about the initialization in Figure 6. While $W_1$ and $W_2$ are both initialized constant, they are trained alongside all other parameters $\theta$ in the single-layer transformer with non-linear soft-max attention (Lines 60-64). Thus, Figure 6 of the paper (Figure 1, 3 and 4 in the rebuttal PDF) effectively demonstrates the success of our initialization scheme and the broad applicability of our insights.

• Q.1: In Figure 6 in the paper, our standard initialization scheme is to let $e_0=1$ and $w_0 = -1$ which falls in the region $\mathcal{I}_\ast$ (see Figure 1d) and converges to the global minimum as predicted by our theory. To test our hypothesis on a wider range of initializations, we try the following set of initializations which concur with our theoretical predictions (Figure 3 of the rebuttal): we let $(e_0,w_0) \in \mathcal{N}(\mu, 0.02)$ with $\mu \in \{ (-1,-1), (0, -0.5), (0,0.5), (0, 0) \}$, with $p=0.5$ and $q=0.8$. As highlighted in Figure 3 of the rebuttal, we observe that the initializations around $(-1,-1)$ leads to global minimum convergence whereas the rest converge to local minima, as predicted by Figure 1d in the paper.

• Q.2: Yes indeed. We wanted to highlight the main ideas behind our gradient-flow analysis using canonical parameterization and hence chose to omit the attention parameters in the beginning for the ease of exposition and clarity.

Dear Reviewer Fte8,

We sincerely appreciate the time you have taken to provide valuable feedback for our work. As we are getting closer to the end of the discussion period, could you let us know if our responses above have adequately addressed your concerns? We remain at your disposal for any further questions.

If you agree that our responses to your reviews have addressed the concerns you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments!

Sincerely, The Authors

We thank the reviewer for the helpful feedback and insightful comments. We will update the revised version with information about the attention in the main text. We address the individual questions below.

• Rank-1 input sequence……: Please note that the input is a first-order Markov chain, without the assumption of it being rank-one. We use linear attention for ease of theoretical analysis, but our guidelines and insights generalize to non-linear soft-max attention. For instance, Figure 6 in the paper demonstrates the effectiveness of our initialization scheme for the full one-layer transformer model (Lines 60-64). Similarly, the low-rank structure of the parameters, as shown in Figure 5 in the paper (Figure 2 and 6 in the rebuttal PDF), applies to the full general model.

• Initialization guidelines: As our analysis shows, understanding the theoretical foundations of transformer training and initialization, even for shallow models, is challenging. Our paper, focusing on shallow transformers, is a crucial step toward optimizing deeper models. Recent research [1] employs a similar reparameterization technique to study gradient descent dynamics in a simplified two-layer attention-only transformer, elucidating the induction head mechanism. However, unlike our work, they do not explore how initialization impacts convergence to local or global minima. Given the nascent and evolving understanding of deeper transformers and higher-order Markov chains [1,2], exploring similar initialization guidelines for large-scale models is an exciting avenue for future research.

• Low rank parameters: Recent evidence [3,4] suggests that SGD with weight decay inherently leads to rank minimization, resulting in low-rank parameters at convergence. Although these studies focus on feed-forward neural networks, we believe a similar phenomenon occurs in transformers as well..

References

• [1] How Transformers Learn Causal Structure with Gradient Descent? https://arxiv.org/pdf/2402.14735
• [2] Transformers on Markov Data: Constant Depth Suffices, https://arxiv.org/abs/2407.17686v1
• [3] Characterizing the Implicit Bias of Regularized SGD in Rank Minimization, https://arxiv.org/pdf/2206.05794
• [4] Rank Diminishing in Deep Neural Networks: https://openreview.net/forum?id=tIqzLFf3kk

We thank the reviewer for the constructive feedback and helpful suggestions to improve the paper. We will add a separate section with prerequisite background details in the revised version. We refer to the common response for experiments on multi-state Markov chains. We address the individual concerns below.

Scalability: Understanding the theoretical foundations of transformer training and initialization, even for shallow models, remains a challenging task, which our work aims to address. While this paper focuses on shallow transformers, it represents a crucial step toward understanding optimization in deeper models. For instance, recent research [1] employs a similar reparameterization technique to study gradient descent dynamics in a simplified two-layer attention-only transformer, elucidating how they learn the induction head mechanism. However, unlike our work, they do not investigate how initialization impacts convergence to local or global minima. Given the nascent and evolving understanding of deeper transformers and higher-order Markov chains [1,2], exploring similar initialization guidelines for large-scale models is an exciting avenue for future research.

References

• [1] How Transformers Learn Causal Structure with Gradient Descent? https://arxiv.org/pdf/2402.14735
• [2] Transformers on Markov Data: Constant Depth Suffices, https://arxiv.org/abs/2407.17686v1

Dear Reviewer DieV,

Thanks for the insightful questions and keeping the score. We are afraid that there has been a misunderstanding about the generality of the insights, which we addressed in the global response and also individually to the Reviewer er5c just now. We are reposting them here for the sake of completeness.

Generality of the insights: For all the experiments in the paper, we use the full general 1-layer transformer with non-linear soft-max attention, as described in lines 63-64 of the paper. For the theoretical analysis, we use the linear attention mechanism but with rest of the architecture intact. Using these insights derived from the linear setting, we demonstrate that our initialization scheme performs significantly better than standard Gaussian initialization even on the full soft-max model. Since the same soft-max attention used for both schemes empirically, our superiority is not due to linear vs. soft-max but rather due to capitalizing on the actual loss landscape as illustrated in Figure 4 in the paper and theoretically characterized in Theorem 9 of Appendix G. Hence we respectully disagree with the Reviewer er5c's claim that "...our intialization scheme in linear attention prevents the softmax transformers from converging to local minima...".

Attention features are indeed learnt: Note that with canonical parameterization, the attention score is given by $`attn`_{n,i} = a \cdot (x_n - \frac{1}{2}) (x_i - \frac{1}{2})$ (Eq. 45 in the Appendix). Since $a$ is a trainable parameter, the attention scores are also learned.

We do not limit low-rank parameters: Please note that the low-rank parameters in the attention layer are not artificially restricted to be same as $\alpha$. The reason why they are also parameterized as $\alpha$ is because empirical evidence suggests that the transformer parameters always converge to this specific structure. Furthermore, this $\alpha$ turns out to be all $\pm 1$ vector. Hence, even after parameterizing $\alpha$ as a trainable vector in our setting, it turns out that the final loss function $L$ (in line 203) is independent of $\alpha$ (since $\|\alpha \|$ is constant, cf. Appendix B) and only depends on three scalar parameters $(e,w,a)$. Thus without loss of generality, it suffices to analyze these scalars.

We thank the reviewer for the helpful feedback and insightful comments. We refer to the global response regarding our results across repeated trials for various $(p,q)$ and the measure of low-rankness. We address the individual concerns below.

• Binary and large alphabet: This paper primarily focuses on the binary alphabet to analytically characterize the phenomena reported in [1], specific to the binary setting. However, we observe a similar low-rank structure for transformer parameters in a first-order Markov chain on a larger alphabet (see Figure 6 in the rebuttal PDF). Consistent with the binary setting, parameters initialized as low-rank remain low-rank throughout training, as shown in Figure 6 . Theoretically, while our gradient-flow analysis can be generalized to the multi-alphabet setting, it is challenging to precisely characterize the phase portraits, as in Figure 2 of the paper, due to the increased number of parameters. This presents an intriguing avenue for future research.

• $a \approx 0$: This results from the fact that for first-order Markov chains, the Markov kernel $\mathbb{P}(x\_{n+1} = 1 \mid x_1^n) = x_n (1-p-q) + p$ is a simple linear function of the last symbol x_n​. Consequently, this function can be represented by a single-layer transformer using only the skip connection in the attention layer, without relying on the attention mechanism. Therefore, $a \approx 0$ becomes a viable solution, which we leverage for the gradient-flow analysis without attention. Figure 7 in the rebuttal empirically corrobarates this fact. As a side note, there are also non-zero attention coefficients that represent global minima, as shown in Figure 4 and discussed in Section 4.1 of the paper.

References:

[1] Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains: https://arxiv.org/abs/2402.04161

Dear Authors,

Thank you for your responses to the other reviewers and myself, and for responding

Low-rank structure: Thank you for simulating Markov chains with more data. Looking at Figure 6 (despite exceeding the page limit), it appears to be the case that the rank of $W_1$ and $W_V$ is |S| - 1, based on the binary and 5-state examples given (I assume Figure 6 used the same alphabet size as Figure 5, but this does not appear to have been explicitly stated). This further seems to imply that for more plausible settings, where alphabets are very large, we should initialize the weight matrices to be full rank. While I am not as concerned as other reviewers about the analytical simplifications of reducing assuming parameters to be low-rank at initialization for the sake of tractability, I find it hard to see how the implications, as presented in section 5, are very impactful.

Nevertheless, I appreciate that the paper is technically strong and precise. If the assumptions and limitations are made clearer throughout the text (which the authors have promised to do), I believe this paper could be a nice, small contribution towards analyzing transformers, but I cannot presently strongly recommend acceptance. So, I maintain my score.

The PDF in this global response is longer than the 1-page limit. Reviewers are requested to look only at the first page per the review policy.

Dear Area Chair Pe5t,

We apologize for overflowing the 1-page figure limit. Unfortunately, despite our best efforts we couldn't contain the experimental results requested by reviewers in less than a page.

Once again, we apologize for the inconvenience.

Regards, The Authors