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

Authors: ashok  vardhan makkuva
marco  epfl
  bondaschi
adway  girish
hyeji  kim
michael  gastpar
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-07 09:46:36

Abstract: In recent years, transformer-based models have revolutionized deep learning, particularly in sequence modeling. To better understand this phenomenon, there is a growing interest in using Markov input processes to study transformers. However, our current understanding in this regard remains limited with many fundamental questions about how transformers learn Markov chains still unanswered. In this paper, we address this by focusing on first-order Markov chains and single-layer transformers, providing a comprehensive characterization of the learning dynamics in this context. Specifically, we prove that transformer parameters trained on next-token prediction loss can either converge to global or local minima, contingent on the initialization and the Markovian data properties, and we characterize the precise conditions under which this occurs. To the best of our knowledge, this is the first result of its kind highlighting the role of initialization. We further demonstrate that our theoretical findings are corroborated by empirical evidence. Based on these insights, we provide guidelines for the initialization of single-layer transformers and demonstrate their effectiveness. Finally, we outline several open problems in this arena. Code is available at: \url{https://github.com/Bond1995/Markov}.

Main paper: 
 introduction: 
 Transformers have been at the forefront of recent successes across various fields including natural language processing [34] . To obtain insights into their impressive sequential modeling capabilities, a notable emerging theme among several recent works is to model the input data as a Markov process.Using this Markovian perspective, works such as [24, 14, 8] , among others, study the in-context learning capabilities of transformer. [23] analyzes the loss-landscape for the next-token prediction task, while [18] shows an equivalence between the attention mechanism and Markov models. Although these works reveal interesting insights about transformers and their capabilities, many fundamental questions about their learning dynamics remain unanswered. In particular, a comprehensive characterization of their training dynamics vis-á-vis the data distributional properties and the role of initialization is still missing.To address this gap, in this paper, we focus on the canonical setting of first-order Markov chains and single-layer transformers and analyze the learning dynamics in this context. Specifically, we prove (Thms. 2, 3, and 8) that the input data properties and the parameter initialization play a significant role in the convergence of the transformer parameters to either local or global minima on the loss surface. Further, we precisely characterize (Figs. 1 and Fig. 2 ) the specific data characteristics and the region of initialization under which this convergence occurs. Based on these insights, we provide guidelines for the initialization of transformer parameters and empirically corroborate our theoretical findings. On the theoretical front, our analysis provides a novel gradient flow analysis of the transformer parameters, capitalizing on their low-rank structure during training. Our main contributions can be summarized as follows:• Theoretical analysis: We precisely characterize the loss landscape and gradient flow dynamics for single-layer transformers with first-order Markov chains (Secs. 3 and 4) . We demonstrate that transformer parameters trained on next-token prediction loss can converge to global or local minima, depending on the initialization and the Markovian data properties, and determine the exact conditions under which this occurs (Thms. 2, 3, and 8). To the best of our knowledge, this is the first result of its kind.• Insights into initialization: Our theoretical analysis underscores the crucial role of initialization in transformer parameter training. Specifically, we demonstrate how the standard Gaussian initialization scheme can lead the convergence to local or global minima depending on the Markovian data properties (Thms. 2 and 8, Figs. 1 and 2 ).• Guidelines: Based on these insights, we provide practical guidelines for parameter initialization, corroborated by empirical evidence demonstrating their effectiveness (Sec. 5.2). Notation. We denote scalars by italic lower case letters like x, y and Euclidean vectors and matrices in bold: x, y, M , etc. ∥ • ∥ denotes the ℓ 2 -norm for Euclidean vectors and Frobenius norm for matrices.[k] ≜ {1, . . . , k}, and for a sequence (x n ) n≥1 , define x m k ≜ (x k , . . . , x m ) if k ≥ 1 and (x 1 , . . . , x m ) otherwise. For z ∈ R, the sigmoid σ(z) ≜ 1/(1+e -z ), ReLU(z) ≜ max(0, z) and the convex logistic loss ℓ log (z) ≜ log (1 + exp(-z)) ∈ (0, ∞). For events A and B, P (A) denotes the probability of A whereas P (A | B) the conditional probability. Let (x, y) be a pair of discrete random variables on [k] × [k] with the probability mass function (pmf) of x being p x = (p 1 , . . . , p k ) ∈ [0, 1] k . Then its Shannon entropy is defined as H(x) = H(p x ) ≜ -i∈[k] p i log p i . The conditional entropy is defined to be H(y|x) ≜ H(x, y) -H(x). The entropy rate of a stochastic process (x n ) n≥1 is defined as lim n→∞ H(x n 1 )/n. We simply write x = y to mean P (x = y) = 1. We also use the shorthand P (y = j | x) for P (y = j | x = x) as a function of the random variable x. For p ∈ (0, 1), the binary entropy function h(•) is defined as h(p) ≜ -p log p -(1 -p) log(1 -p). 
problem setting: 
 We formally define the problem setting for analysis of single-layer transformers with Markovian data, following [23] .Input data. We assume that the input word sequence {x n } N n=1 ∈ {0, 1} N is a first-order timehomogenous Markov chain with a fixed kernel P = (P ij ). That is, the transition probabilityP ij ≜ P (x n+1 = j | x n = i) = P x n+1 = j | x n = i, x n-1 1, for any x n-1 1 , i, j ∈ {0, 1}. In particular, we consider P = [1 -p, p; q, 1 -q] where p = P 01 = P (x n+1 = 1 | x n = 0) and q = P 10 = P (x n+1 = 0 | x n = 1) denote the switching probabilities from the states 0 and 1 respectively. We call p + q, the switching factor. We assume that the process is already mixed, i.e. x n ∼ π for all n, where π ≜ (π 0 , π 1 ) = (q, p)/(p + q) is the stationary distribution satisfying π = πP . Succinctly, (x n ) n≥1 ∼ (π, P ). For this process, the entropy rate, H(x n+1 |x n ) = 1 p+q (q h(p) + p h(q)), and the entropy of the marginal, H(x n ) = H(π), are both constant in n.Transformer architecture. We consider a single-layer transformer with a single-head attention and ReLU non-linearity. Given an input sequence {x n } N n=1 , it performs the following mathematical operations at each n ∈ [N ] to predict the next-token probability f θ (x n 1 ):x n ∈ {0, 1} Embedding -------→ x n Attention ------→ y n Feed-forward ---------→ z n Linear ----→ logit n Prediction -------→ f θ (x n 1 ), whereEQUATIONy n = x n + i∈[n]att n,i ∈(0,1)EQUATIONEQUATIONEQUATIONHere θ ≜ (e, {p n } N n=1 , . . . , W 1 , W 2 , b, a) ∈ R D denotes the full list of the transformer parameters from the embedding layer till the linear layer ( § A details them). While the underlying data {x n } N n=1 is Markovian, i.e. P (x n+1 = 1 | x n 1 ) = P (x n+1 = 1 | x n ), the transformer is agnostic to this fact and it can potentially utilize the full past x n 1 in the Attention layer, via the attention weights att n,i , to predict the next-symbol probability f θ (x n 1 ) = P θ (x n+1 = 1 | x n 1 ). Note that it suffices to estimate the symbol 1 probability as the vocabulary is binary. We also refer to the above architecture as "full model".Loss and training. The transformer parameters θ are usually initialized according to standard Gaussian distribution N (0, σ 2 I) with a small variance σ 2 [26] and are trained using gradient-based methods to minimize the cross-entropy loss on the next-token prediction, i.e.EQUATIONWhen the input sequence {x n } N n=1 ∼ (π, P ), the minimal loss equals its entropy-rate, i.e. L ⋆ ≜ min θ L(θ) = H(x n+1 |x n ) [11] .A key surprising observation in [23] is that the loss function L(•) admits both the global and local minima depending on the switching factor p + q of the Markovian data, and the weight-tying of the embedding and linear weights (e = a) of the transformer. In particular, they show that (i) for all (p, q) ∈ (0, 1) 2 , there exists a global minimum θ ⋆ for the loss L such that its prediction matches the Markov kernel, i.e. P θ⋆ (x n+1 = 1 | x n 1 ) = P (x n+1 = 1 | x n ). (ii) if p + q > 1 and the weights are tied (e = a), there exists a bad local minimum θ min for L whose prediction equals the marginal, i.e. P θπ (x n+1 = 1 | x n 1 ) = P (x n+1 = 1).In view of these results, we focus on the weight-tying scenario and hence let e = a to be a single parameter in R d . Thus, θ = (e = a, {p n } N n=1 , . . . , W 1 , W 2 , b). We interchangeably refer to θ as both the transformer and the set of parameters.Our objective. While the aforementioned results detail the static landscape of the loss, they do not characterize the learning dynamics on the loss surface and the effect of initialization, which plays a central role in training machine learning models [5] . In view of these shortcomings, the main objective of this paper is to address the following question:(Q.1): Can we explain how the initialization and learning dynamics affect the convergence of the transformer parameters θ to the local or global optima?3 Canonical Low-rank Parameterization Motivation. Given the complexity of the transformer architecture and the non-convex loss function, it is challenging to analyze the learning dynamics directly [24, 14] . To tackle this, we capitalize on the following empirical observation [23] which is the motivating idea behind our approach: when trained by gradient-based methods, the weight matrices (W V , . . . , W 1 , W 2 ) at the optima θ ⋆ and θ min exhibit rank-one structure, whose eigenvector is the same direction in which the both the token embedding e and the positional embeddings p n are all aligned in. Interestingly, such low-rank solutions can also be shown to be theoretically optimal ( § B details these structures). While these observations illustrate the implicit bias towards low-rank solutions at the final convergence, a natural question arises: if we initialize with low-rank parameters, will they remain low-rank during training? In Sec. 5.1, we affirmatively address this based on a thorough empirical evaluation for single-layer transformers and inspired by these empirical phenomena, without loss of generality, we restrict our attention to these low-rank manifolds to characterize the learning dynamics. This is similar in spirit to [24] , where they assume special attention matrix structure for learning induction heads.Parameterization. More specifically, we consider a special low-rank parameterization that is empirically observed and capitalize on it to address (Q.1). Interestingly, along this low-rank manifold, it suffices to consider a reduced set of parameters θ ∈ R 2 or θ ∈ R 3 given by:θ = (e, w) ∈ R 2 , or θ = (e, w, a) ∈ R 3 . (Reparameterization)Here e denotes the embedding scalar, w the weight, and a the attention parameter respectively. Now we describe the parameterization of the transformer vis-á-vis these scalars and refer to § C for a more detailed descripton. Let the input {x n } N n=1 be a first-order Markov chain as in Sec. 2 and let n ∈ [N ] be fixed. Then we haveEmbedding : e = e • α, p n = - e 2 • α → x n = e x n - 1 2 α, e ∈ R, α ∈ {±1} d / √ d, Attention : W V = α v ⊤ → y n = e x n - 1 2 α + ⟨v, α⟩ ∝a≈0   i∈[n] att n,i • e x i - 1 2   α, v ∈ R d .The scalar a is the product of ⟨v, α⟩ and the scaling in the attention weights att n,i (Eq. ( 39)), which is empirically close to zero for first-order Markov chains. Hence for the ease of exposition, we first let a = 0 and analyze the general case when a ∈ R in Sec. 4.1. We continue:Feed-forward : W 1 = |w| √ d 1 α ⊤ , W 2 = w √ d α 1 ⊤ → z n = e x n - 1 2 (1 + 4w|w|x n ) α, w ∈ R.1 is the all-one vector in R r with r = 4d typically in practice. Substituting this z n in the linear layer with e = a and bias b ∈ R, the logits and the probabilities simplify to:EQUATIONEQUATIONFinally, using the equivalence between the cross-entropy loss and the logistic loss ℓ log (•), the loss function in Eq. ( 1) can be compactly written as (Lemma 6):EQUATIONDue to convexity of ℓ log (•), it follows that L(θ, b) is convex in the bias b for any fixed θ, whose minimizer, b ⋆ (θ) = argmin b∈R L(θ, b), has a closed form expression (Lemma 5). Hence, without loss of generality, we consider the loss with this optimal bias b ⋆ :EQUATIONEmpirically, this roughly translates to running the gradient-based algorithm for the bias for more steps at each θ. In practice, one additional step is usually sufficient (see Sec. 5). Eq. ( 5) resembles the standard logistic regression loss [31] whose binary labels are 2x n+1 -1 ∈ {±1} and the logits given by e 2 (1 + 2w|w|)x n + b ⋆ -e 2 /2, for each n ∈ [N ].The key difference here is that the logits are a non-linear function of the parameters (e, w) unlike in the standard setting.We briefly summarize our assumptions below.Assumption 1 (Canonical parameterization). For our theoretical analysis, we assume that the effective transformer parameters are canonically parameterized as θ = (e, w, a) ∈ R 3 . First we study the scenario when a = 0 with θ = (e, w) and build upon these observations to study the general setting of θ = (e, w, a) in Sec. 4.1. 
loss landscape with canonical parameterization: 
 With the new set of parameters θ = (e, w) ∈ R 2 , we are now ready to analyze the loss L(•) in Eq. ( 5).First we recall the definition of a critical point [20] .A point θ ⋆ ∈ R 2 is a critical or a stationary point for L if ∇L(θ ⋆ ) = 0. A critical point θ ⋆ is a local minimum if there exists a neighborhood U around θ ⋆ such that L(θ ⋆ ) ≤ L(θ) for all θ ∈ U , and a local maximum if L(θ ⋆ ) ≥ L(θ).If the neighborhood U is whole of R 2 , it is a global minimum/maximum. On the other hand, a critical point is a saddle point if for all neighborhoods U around θ ⋆ , there areθ 1 , θ 2 ∈ U such that L(θ 1 ) ≤ L(θ ⋆ ) ≤ L(θ 2 ).Thm. 1 below provides a complete characterization of the loss landscape in terms of the aforementioned critical points.Theorem 1 (All critical points). Let the input sequence be {x n } N n=1 ∼ (π, P ), the transformer parameters θ = (e, w) ∈ R 2 , and the next-token prediction loss L(•) be as in Eq. (5) . Then for any (p, q) ∈ (0, 1) 2 with p + q ̸ = 1 and N ∈ N, (i) the set of all global minima is given byEQUATION(ii) the set of all local minima is given by Θ min (p, q) ≜ (e, w) ∈ R 2 : e = 0, (p + q -1)(1 + 2w|w|) > 0 ,(iii) the set of all local maxima is given byEQUATION(iv) and the set of all saddle points isEQUATIONThus the set of all critical points isEQUATIONIn addition, for any θ ⋆ ∈ Θ ⋆ , θ min ∈ Θ min , θ max ∈ Θ max , and θ sad ∈ Θ sad , the loss values satisfyH(x n+1 | x n ) = L(θ ⋆ ) < L(θ min ) = L(θ max ) = L(θ sad ) = H(x n+1 ).Proof. We refer to § E.Fig. 1 illustrates the loci of these critical points for p + q < 1 and p + q > 1. Motivated by empirical observations, while [23] characterizes local minima for p + q > 1, it is interesting to note that our Thm. 1 shows that local minima also exist for p + q < 1 (Eq. ( 7) and Fig. 3a ). So why did they find the minima only when p + q > 1? The answer to this, and more broadly to question (Q.1) lies in the learning dynamics and initialization for θ, which we study in the next section. 
learning dynamics: 
 Capitalizing on the loss landscape in terms of the critical points in Thm. 1, we now focus on the convergence of gradient-based algorithms to these points. In this regard, we analyze the dynamics of the gradient-flow (GF), which can be viewed as a continuous-time analogue of gradient-descent [6] .The gradient-flow of the parameters, (θ t ) t≥0 , on L is governed bydθ t dt = -∇L(θ t ), θ t = (e t , w t ) ∈ R 2 , t ≥ 0, (GF)where θ t ≜ θ(t) is a C 1 (continuously differentiable) curve in R 2 starting with a randomly initalized θ 0 . To characterize these trajectories, we define an energy function E(•, •), which plays a crucial in the GF dynamics. It is defined asEQUATIONwhere e-axis ≜ {(e, w = 0)}. Note that E is well-defined and finite for all the points in its domain.On the other hand, lim w→0 -E(e, w) = -∞ whereas lim w→0 + E(e, w) = ∞ for any fixed e ∈ R.Thus the e-axis corresponding to w = 0 serves as an energy barrier for the flow. Figs. 3a and 3b illustrate this by visualizing the energy contour lines. The utility of the energy function is captured in the following lemma.Lemma 1 (Constant energy along the flow). For any (p, q) ∈ (0, 1) 2 and initialization θ 0 = (e 0 , w 0 ) ∈ R 2 , let (θ t ) t≥0 be the corresponding GF trajectory starting from θ 0 . If θ 0 ∈ R 2 \ e-axis, the energy stays constant along the trajectory, i.e.EQUATIONOn the other hand, if θ 0 ∈ e-axis, we have that θ t ∈ e-axis for all t ≥ 0 with w t = w 0 = 0, i.e. if we initialize on the e-axis, the trajectory always stays there.We are now ready to present the main results of our paper. Specifically, Thm. 2 and Thm. 8 highlight the role of the switching factor of the Markovian data, p + q, and the parameter initialization, θ 0 , in deciding whether the GF converges to local optima or global optima. First we define the energy valueE sad ≜ E(e = 0, w = -1/ √ 2) = -(1 + log 2)/2.Theorem 2 (GF dynamics for p + q > 1). Let (p, q) ∈ (0, 1) 2 with p + q > 1, the input sequence be {x n } N n=1 ∼ (π, P ), and (θ t ) t≥0 be the corresponding GF trajectory starting from θ 0 . Then for all initializations θ 0 ∈ R 2 , the gradient flow converges to a critical point of the loss L. That is, there exists a θ lim ∈ R 2 such that lim t→∞ θ t = θ lim and ∇L(θ lim ) = 0. In particular, θ lim is a(i) a local minimum if θ 0 ∈ I min ≜ (e, w) : w ∈ (-1/ √ 2, 0), e ∈ (-g(w), g(w)), g(w) = w 2 -log(-w) + E sad ∪ {(e, w) : w ≥ 0} , (ii) a saddle point if θ 0 ∈ I sad ≜ (e, w) : w ∈ [-1/ √ 2, 0), e = ± w 2 -log(-w) + E sad , (iii) a local maximum if θ 0 ∈ I max ≜ (e, w) : e = 0, w < -1/ √ 2 , (iv) and a global minimum if θ 0 ∈ I ⋆ ≜ R 2 \ (I min ∪ I sad ∪ I max ).Consequently, when p + q > 1, if we use the standard initialization θ 0 ∼ N (0, σ 2 I 2 ) with σ 2 ≪ 1/ √ 2, θ lim will be a local minimum with high probability. If p + q < 1, under the same initialization scheme, θ lim will be a global minimum with high probability.Proof sketch. The main idea behind the proof is to show that if we do not initialize on the e-axis, the flows stays on the constant energy contour (Lemma 1) and hence converges to a critical point of the loss L, which is at the intersection of the contour line and the set of critical points (Lemmas. 10 and 11) . By determining where these intersections occur, the corresponding basins of convergence I min , . . . , I ⋆ are obtained by showing that an initialization in a specific set leads to the said critical point (Thm. 1). The proof for θ 0 ∈ e-axis is similar.Figs. 1b and 1d illustrate these initialization sets corresponding to the convergence basins for p = q = 0.9 and p = q = 0.1 respectively. An analogous result about GF dynamics for p + q < 1 is presented in Thm. 8 ( § F.2). Here a key difference is that small Gaussian initialization around origin leads to a global minimum θ lim with high probability (Fig. 1d ). Key insights. Together, Thm. 2 and Thm. 8 address our motivating question (Q.1) by fully characterizing the GF dynamics in terms of initialization and input data properties. Specifically, our results explain the phenomenon in [23] wherein they observe local minima for p + q > 1 more often than for p + q < 1, owing to standard Gaussian initialization around origin (Figs. 1b and 1d ). However, in practice, we often do not know the input switching factor, raising a natural questions: is there a data-agnostic initialization that always converges to global minima? Indeed, as can be seen from Figs. 1b and 1d , there is a common region of initialization in the negative half-plane above the saddleasymptotes (in yellow) that leads to the global minima convergence irrespective of the switching p + q. Mathematically, this region is given by I common ≜ {(e, w) : w < 0, |e| > w 2 -log(-w) + E sad }. We empirically corroborate this fact in Sec. 5.2. In this section, the consider the attention scalar a ∈ R (Sec. 3) and study the gradient flow dynamics with the parameters θ = (e, w, a) ∈ R 3 . The parameter a captures the overall scaling from the value, key, and query components in the attention layer. Recall that the soft-max attention weights are given by att n,i ∝ exp(⟨q n , k i ⟩/ √ d), where q n = W Q x n and k i = W K x i are the query and key embeddings for any position i ∈ [n]. Using the low-rank structure of the query and key matrices, satisfying W ⊤ Q W K = (q 2 d) αα ⊤ and the value matrix W V = αv ⊤ for some q ∈ R and v ∈ R d ( § G), and assuming linear attention att n,i ∝ ⟨q n , k i ⟩/ √ d, we define a single scalar a ≜ ⟨v, α⟩q 2 d 5/2 /4 that captures the essence of the attention layer (Eq. ( 39)). We note that linear attention weights are a standard assumption in the transformer analysis literature [3, 36] . Using this parameterization, similar to the steps in Sec. 3, we obtain the final loss function to beL(θ) = E ℓ log (2Y -1) e 2 X - 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b ⋆ ,where θ = (e, w, a) and b ⋆ is the corresponding optimal bias. L recovers the loss in Eq. ( 5) when a = 0. In Thm. 10, we determine the set of all critical points of L in terms of global minima and local optima in closed-form expressions, analogous to Thm. 1. Capitalizing on this characterization, we now shift our focus to the analysis of the gradient flow in R 3 . To this end, let (θ t ) t≥0 be a C 1 curve in R 3 governed bydθ t dt = -∇L(θ t ), θ t = (e t , w t , a t ) ∈ R 3 , t ≥ 0, (GF-attn)starting with a randomly initalized θ 0 . We define the energy function E(•, •, •) asEQUATIONwhere ea-plane ≜ {(e, w = 0, a)}. It is similar to its counterpart in Eq. (11) , except for the 2a 2 term. Fig. 2 visualizes this energy surface and the set of critical points, which reveal close resemblance to that of Fig. 1 in R 2 . Capitalizing on the energy function, we now present our main result with the attention. Theorem 3 (GF dynamics with attention). For any (p, q) ∈ (0, 1) 2 and initialization θ 0 ∈ R 3 , let (θ t ) t≥0 be the corresponding GF-attn trajectory starting from it. Then for all θ 0 ∈ R 3 , the gradient flow converges to a critical point of the loss L. That is, there exists a θ lim ∈ R 3 such that lim t→∞ θ t = θ lim and ∇L(θ lim ) = 0. Further,(i) if θ 0 ∈ R 3 \ ea-plane, we have E(θ lim ) = E(θ t ) = E(θ 0) for all t ≥ 0. Hence θ lim is at the intersection of the energy contour line E = E 0 with that of the set of critical points. (ii) if θ 0 ∈ ea-plane, we have θ t ∈ ea-plane for all t ≥ 0 and hence θ lim ∈ ea-plane.Proof. We refer to § G and § N.4.Thm. 3 shows that the learning dynamics with attention closely resemble those without it (Thms. 2 and 8). While the set of all critical points of L, and thus the limit points of the flow, has a closed-form expression (Thm. 10), deriving the same for the initialization sets I min and I ⋆ to determine the basin of convergence is technically challenging (see discussion in § G). Nonetheless, empirical observations with the standard Gaussian initialization around origin reveal a similar picture as in the two-dimensional setting for both the p + q < 1 and p + q > 1 cases (Fig. 2 ). We believe it's an interesting direction of future research to theoretically characterize this, analogous to Thms. 2 and 8. We refer to § G for additional details and proofs. 
empirical results: 
 We empirically validate our canonical parameterization θ ∈ R 3 (Sec. 3) by demonstrating full model convergence to low-rank parameters through both qualitative and quantitative evidence. Qualitatively, we visualize weight matrices across iterations; quantitatively, we plot the percentage of energy captured by the top-rank components across iterations. We then demonstrate the generalization of our theoretical findings on local optima and initialization with canonical parameters to the full model θ ∈ R D . We conclude with a discussion on higher-order and multi-state Markov chains.Full model converges to low-rank. We let the input Markov sequence to be {x n } N n=1 ∼ (π(p, q), P (p, q)) for p = 0.2, q = 0.3, N = 1024 and consider the full model as defined in Sec. 2 with embedding dimension d = 8. First, we initialize the parameters θ = (e = a, {p n } N n=1 , . . . , W 1 , W 2 , b) using the standard Gaussian initialization with standard deviation 0.001 [26] and train them using SGD on a batch size B = 16 and for t = 800 iterations. In Fig. 6 , we track the value matrix W V ∈ R d×d and the weight matrix W 1 ∈ R 4d×d across iterations. We observe that at convergence both W V and W 1 are approximately rank-one with one of their components being same as the embedding vector (the row in W V and column in W 1 ). Further, the embedding vector has all entries in {±1} up to a scaling. We observe the same conclusion for other weight matrices W K,Q , W 2 and for all values of (p, q) ∈ (0, 1) 2 . Fig. 3 also quantitatively demonstrates this.Full model initialized at low-rank remains low-rank during training. Inpsired by the low-rank structure obtained above, we randomly initialize the weight parameters as rank-one matrices and the embeddings on the hypercube {±1} d . After the initialization, we train them without any low-rank restrictions, and track them during the course of training. Interestingly, here we observe that the parameters still stay low-rank as illustrated in Fig. 7 and Fig. 3 . A similar conclusion holds for the remaining weight matrices. Together these results provide the empirical basis for our canonical parameterization analysis in Sec. 3.Now we investigate the findings of Sec. for p + q > 1, any small initialization around zero would lead a local minima convergence. To test this hypothesis, we compare the standard initialization where all the transformer parameters θ = (e = a, {p n } N n=1 , . . . , W 1 , W 2 , b) are randomly chosen around zero with small variance σ = 0.02, with a new initialization based on our results, where we initialize the embedding vector e such that all cordinates are equal to e = 0.5, W 1 to be constant with the scalar w 1 = 1 and W 2 constant with w 2 = -1 (corresponding to I ⋆ in Fig. 1d ). We indeed observe that the final test loss matches the unigram loss for the standard initialization, while it converges to the optimal bigram loss for our initialization (see Fig. 4 ). Together these results indicate that though our analysis used canonical parameterization, the corresponding insights are more general and apply more broadly to the general full model. In a similar spirit, analysis of initialization effects for deeper architectures is an interesting avenue of future research.While the primary focus of this paper has been on binary first-order Markov chains, we believe it's possible to extend our analysis to both multi-state and higher-order settings. On the multi-state front, akin to the binary case, [23] already demonstrates the effect of switching probability and weight-tying on the final model convergence. Here, first characterizing the loss landscape and then the associated learning dynamics in line with our approach is an interesting direction. On the other hand, a recent work [28] establishes a surprising result that any k th -order Markov chain can be represented by a three layer transformer with just one head per layer, relying on induction head mechanism. Analyzing gradient flow dynamics using appropriate canonical parameterization (cf. [24] ) in this scenario is also a fruitful direction of research. 
related works: 
 The recent success of transformer models in deep learning has sparked significant interest and active research in understanding them [38, 25, 16, 27, 15, 37, 40, 32] . In relation to our paper, they can be broadly classified into two topics: (i) In-context learning (ICL): ICL refers to the ability of transformers learn and reason from information present in their context [10, 13, 4, 35, 39, 7, 21, 17] . Along this thread, the works most relevant to ours are [8, 14, 24] , which use Markovian input data to understand the ICL mechanism. [8, 14] heuristically show how gradient-based updates can learn an induction-head mechanism using a simplified transformer architecture with frozen encodings, query matrices and linear activations. On the other hand, we consider the canonical parameterization, capitalizing on inherent low-rank parameters, to provide a full characterization of the learning dynamics. [24] demonstrates how two-layer transformers with GD learn induction head mechanism when the input has a causal tree dependency, such as in Markov chains. In this work, we focus on the GF dynamics for single-layer transformers and show how they can also converge to local optima, further highlighting the role of initialization. (ii) Training dynamics: On the other hand, numerous works have investigated the training dynamics of transformers. For instance, [9] examines the gradient flow in a simplified single-layer transformer, while [33] studies the process by which self-attention integrates input tokens, assuming the decoder learns faster than the attention layer. Unlike these settings, our focus is on understanding the training dynamics of the full transformer model end-to-end. Other related works include [30] , which analyzes gradient dynamics in LSTM Seq2seq models, [19] , which shows how Vision Transformers learn spatial structures, and [22] , which demonstrates that a single-layer transformer can learn a constrained topic model. A closely related work is [18] , which shows that self-attention has a Markovian structure, but our focus is on self-attention's capability in modeling Markov chains and the associated training dynamics. 
conclusion: 
 In this work, we present a novel characterization of gradient flow dynamics for (weight-tied) singlelayer transformers with first-order Markov chains. Specifically, we highlight the significant role of the parameter initialization and inherent properties of the Markovian data in determining the parameter convergence to either global minima or local optima. Drawing upon these insights, we offer practical guidelines for parameter initialization, corroborated by empirical results demonstrating their effectiveness. While our current analysis is limited to single-layer models, uncovering similar results with gradient flow analysis for deeper architectures and higher order Markov chains is open and an interesting avenue for future research.  
 
 Appendix: We first describe the transformer architecture from Sec. 2:EQUATIONEQUATIONEQUATIONEQUATIONHere θ ≜ (e, {p n } N n=1 , . . . , W 1 , W 2 , b, a) ∈ R D denotes the full list of the transformer parameters from the embedding layer till the linear layer. In the attention layer, the weight assigned to each value, att n,i , is computed by a compatibility function of the query vector q n ≜ W Q x n and the corresponding key vectorsk i ≜ W K x i for all i ∈ [n]. More precisely, att n,i ≜ softmax((⟨q n , k 1 ⟩, . . . , ⟨q n , k n ⟩)/ √ d) i . W K,Q,V ∈ R d×dare the respective key, query, and value matrices. For multi-headed attention, the same operation is performed on multiple parallel heads, whose outputs are additively combined.Finally, the transformer parameters θ ≜ (e, {p n } N n=1 , . . . , b, a) are trained via the cross-entropy loss on the next-token prediction:EQUATIONIn this paper, we focus on the weight-tied scenario where e = a. Hence we let them be a single parameter with θ = (e = a, {p n } N n=1 , . . . , b) ∈ R D , where D is the total parameter dimensionality.Now we recall the theoretical results from [23] about the loss landscape of L in the form of global and local minima. Theorem 4 (Global minimum). Let the input sequence be {x n } N n=1 ∼ (π(p, q), P (p, q)) for some fixed (p, q) ∈ (0, 1) 2 . Then for all (p, q), there exists a θ ⋆ ∈ R D with an explicit construction such that it is a global minimum for the population loss L(•) in Eq. ( 14), i.e.(i) L(θ) ≥ L(θ ⋆ ) for all θ ∈ R D .Further, θ ⋆ satisfies:(ii) P θ⋆ (x n+1 = 1 | x n 1 ) = P (x n+1 = 1 | x n ), the Markov kernel. (iii) L(θ ⋆ ) = H(x n+1 |x n ), the entropy rate of the Markov chain. (iv) ∇L(θ ⋆ ) = 0, i.e. θ ⋆ is a stationary point.Let L ⋆ ≜ L(θ ⋆ ) be the global minimal loss from Thm. 4. Now we recall the result on the bad local minimum. Theorem 5 (Bad local minimum). Let the input sequence be {x n } N n=1 ∼ (π(p, q), P (p, q)) for some fixed (p, q) ∈ (0, 1) 2 . If p + q > 1, there exists an explicit θ min ∈ R D such that it is a bad local minimum for the loss L(•), i.e.(i) there exists a neighborhood B(θ min , r) with r > 0 such that L(θ) ≥ L(θ min ) for all θ ∈ B(θ min , r), with L(θ min ) > L ⋆ .Further, θ min satisfies:(ii) P θπ (x n+1 = 1 | x n 1 ) = P (x n+1 = 1) = π 1 , the marginal distribution. (iii) L(θ min ) = H(x n+1 ) = H(π), the entropy of the marginal. (iv) ∇L(θ min ) = 0, i.e. θ min is a stationary point.Here we recall the low-rank structure for the global minima found by SGD consistently across multiple runs when p + q < 1 [23, Appendix C.2]. In particular, it is observed that the token and positional encodings point in the same direction α, which is a low-rank factor for the weight matrices in the attention and the feedforward layers, which in turn are all rank-one. Mathematically, Embedding. The embedding vector e obeys e = e • α for some e > 0 and α ∈ {±1} d . Further, the positional embeddings p n are constant across positions n pointing in the same direction albeit with a negative scalar, i.e.p n = -p • α for p > 0 and p ≈ e2 such that e > p . Thus from Embedding layer,EQUATIONwhich ensures that the respective embeddings for the bit x n = 0 and x n = 1 are x n = -p • α and x n = (e -p) • α, which are roughly anti-podal.Attention. Recall from the Attention layer that the output y n is given byy n = x n + W O i∈[n] att n,i • W V x i, where the attention weights att n,i are computed according toatt n,i = exp ⟨q n , k i ⟩/ √ d / j∈[n] exp ⟨q n , k j ⟩/ √ d with q n = W Q x n and k i = W K x i .Here it is observed that the matrix products are all rank-one with α being a factor, i.e.W O W V = α • v ⊤ ∈ R d×d , for some v ∈ R d , W ⊤ Q W K = (q 2 d) α • α ⊤ ∈ R d×d , for some q ∈ R. Hence, W V x i = ⟨v, α⟩(ex i -p) α,and ⟨q n , k i ⟩ √ d = 1 √ d • x ⊤ n W ⊤ Q W K x n = q 2 d √ d • (x ⊤ n α)(x ⊤ i α) (∥α∥ 2 =d) = q 2 d 3 √ d • (ex n -p)(ex i -p) = q 2 d 5/2 • (ex n -p)(ex i -p).Thus,EQUATION)It is further noticed that ⟨v, α⟩ ≈ 0 and hencey n = (ex n -p) α = x n .Feed-forward. For the Feed-forward layer, both the matrices W 1 ∈ R r×d and W 2 ∈ R d×r exhibit rank-one structure with α being one of the factors,EQUATIONEQUATIONThus W 1 y n = dw(ex n -p) w. Since -p < 0 and e -p > 0, corresponding to x n = 0 and x n = 1 respectively, we obtain ReLU(W1 y n ) = dw ((1 -x n )p • ReLU(-w) + x n (e -p) • ReLU(w)).Denoting the number of ones in w as β, i.e. β = r i=1 1(w i = 1), we further simplify:W 2 ReLU(W 1 y n ) = W ⊤ 1 ReLU(W 1 y n ) = w 2 d ((1 -x n )p • ⟨w, ReLU(-w)⟩ + x n (e -p) • ⟨w, ReLU(w)⟩) α = w 2 d ((1 -x n )p • (β -r) + x n (e -p) • β) α = w 2 d(ex n -p) ((2β -r)x n + r -β) α.HenceEQUATION)Linear. Using the fact that e = a = e • α due to weight-tying, we obtain from Linear layer thatEQUATIONPrediction. We finally obtain that the prediction probabilityf θ (x n 1 ) = σ(logit n ) = x n • σ ed(e -p) 1 + βw 2 d + b + (1 -x n ) • σ -edp 1 + (r -β)w 2 d + b .Thus we see that the prediction probability and hence the loss function L(•) in Eq. ( 14) is influenced only by the scalars e, p, w, b and β.Building on the low-rank strucutre of the transformer parameters described above, we consider a special parameterization for them. A key property of this parameterization is that it covers both the global and local minima from Thm. 4 and Thm. 5 for all (p, q) ∈ (0, 1) 2 . Recall that Thm. 5 characterizes local minima only for p + q > 1 whereas our special parameterization allows to discover local minima even for p + q < 1. Our construction follows the same outline as in Eqs. ( 15)-( 20). First we start with the embedding layer.Embedding. We let e = e • α and p n = -p • α for all n where e > 0, p = e 2 and α ∈ {±1} d / √ d. Thus the embedding x n from Eq. ( 15) simplifies tox n = e x n - 1 2 α ∈ {± e 2 } α.Attention. Substituting this x n in Eq. ( 16), we haveEQUATIONwhere the attention weightsEQUATION)) ∈ (0, 1) for some q ∈ R. Since ⟨v, α⟩ ≈ 0, we let v = 0 and obtainy n = x n = e x n - 1 2 α.Feed-forward. For the feed-forward layer, we observe from Eq. ( 17) and Eq. ( 19) that for any w ∈ {±1} r , only the number of 1's in w, β, matters for the final vector z n which further interacts with the weight scalar w. Hence without loss of generality, we set w to be the all-ones vector: w = 1 ∈ R r and hence β = r = 4d. While we observe from Eq. ( 17) that W 2 = W ⊤ 1 for p + q < 1, we observe from the proof of the Thm. 4 for p + q > 1 in [23, Appendix B.2] that we need W 2 = -W ⊤ 1 in this scenario. Hence we consider the following parameterization that covers both these scenarios:W 1 = |w| √ d 1 • α ⊤ ∈ R 4d×d , W 2 = w √ d α • 1 ⊤ ∈ R d×4d .Here w > 0 ensuresW 2 = W ⊤ 1 whereas w < 0, W 2 = -W ⊤ 1 .Using this parameterization, substituting β = r = 4d and w → w d in Eq. ( 19), we getz n = e x n - 1 2 (1 + 4w|w|x n ) α.Linear. Since e = a = e • α due to weight-tying, Eq. ( 20) simplifies tologit n = ⟨e, z n ⟩ + b = e 2 x n - 1 2 (1 + 4w|w|x n ) + b (xn=x 2 n ) = e 2 x n + 4w|w|x n - 1 2 -2w|w|x n + b = e 2 (1 + 2w|w|) x n + b - e 2 2 .Prediction. The next-token prediction probability isEQUATION)Loss. While we assumed e > 0 in the beginning, in view of Eq. ( 22) and the fact that α ∈ {±1} d / √ d, we see that e ∈ R gives us the same expression for probability. Thus the final probability depends on just the three scalars (e, w, b) ∈ R 3 . Defining θ = (e, w) ∈ R 2 , we recall the cross-entropy loss L(•) from Eq. ( 4) in Sec. 2 for this canonical model:EQUATION)It turns out that we can further remove the bias b by minimizing the loss over it which we discuss in App. E. For now in the next section, we analyze when it's present as in Eq. ( 23).D Analysis of the loss with the bias, L(θ, b), in Eq. (4) and Eq. (23) In this section, we analyze the loss function with the bias, L(θ, b), from Eq. ( 4) and Eq. ( 23), which will later be useful for studying L(θ). First we characterize the set of its critical points in R 3 . To this end, we define the following sets of pointsEQUATIONEQUATIONEQUATION)The following result establishes that these sets exhaust all the critical points. Theorem 6 (All critical points). Let the input sequence be {x n } N n=1 ∼ (π, P ), the transformer parameters (θ, b) = (e, w, b) ∈ R 3 , and the next-token prediction loss L(•) be as in Eq. (23) . Then all the stationary points of L are either in Γ ⋆ , Γ min , or Γ sad , i.e.EQUATION)Proof. We refer to App. J.1.Recall the definitions of local minima & maxima, global minima, and that of all the saddle points from Sec. 3.1. We are now ready to present the main result about the loss landscape of L(•). Theorem 7 (Loss landscape with bias). Let the input sequence be {x n } N n=1 ∼ (π, P ), the transformer parameters (e, w, b) ∈ R 3 , and the next-token prediction loss L(•) be as in Eq. (23) . Then for any (p, q) ∈ (0, 1) 2 with p + q ̸ = 1 and N ∈ N, (i) the set of all global minima of L is given by Γ ⋆ (p, q), (ii) the set of all bad local minima of L is given by Γ min (p, q), (iii) and the set of all saddle points of L is Γ sad (p, q). Furthermore, for any γ ⋆ ∈ Γ ⋆ , γ min ∈ Γ min , and γ sad ∈ Γ sad , the losses are ordered asH(x n+1 | x n ) = L(γ ⋆ ) < L(γ min ) = L(γ sad ) = H(x n+1 ).Remark 1. Note that a bad local minimum is a local minimum whose loss value is strictly less than that of the global minimum, as is the case here. Interestingly, Thm. 7 highlights that all local minima for the loss L are indeed bad local minima.Proof. We refer to App. J.2.The proofs of both Thm. 6 and Thm. 7 rely on few key lemmas that we present below. First we start with the result that rewrites the loss L(θ, b) from Eq. ( 23) in a compact manner using the logistic function ℓ log (•). Lemma 2 (Loss as a logistic function). The next-token prediction loss L(•) in Eq. (23) can be written asEQUATIONwhere (X, Y ) ∈ {0, 1} 2 are distributed according to (X, Y ) ∼ (π, P ), i.e. X is a Bernoulli random variable with X ∼ π ≡ Bern(p/(p + q)) and Y |X ∼ P (p, q), the Markov kernel.The following lemma establishes the gradients of the loss function with respect to the parameters e, w, and b. Lemma 3 (Gradient computation). For any (e, w, b) ∈ R 3 and the next-token prediction loss L(•) in Eq. ( 23), the gradients are given by∂L ∂e = E X [(f 1 X + f 2 )(2X(1 + 2w|w|) -1)] • e, ∂L ∂w = E X [(f 1 X + f 2 )X] • 4e 2 |w|, ∂L ∂b = E X [f 1 X + f 2 ] ,whereX ∈ {0, 1} is a Bernoulli random variable with X ∼ Bern(p/(p + q)), f 1 = σ 2e 2 w|w| + b + e 2 2 + q -1 -σ b -e 2 2+ p, andf 2 = σ b -e 2 2 -p.Remark 2. It is interesting to note that the gradients for both e and w are product of an expectation term and an e factor. Also, except for scaling factors in terms of (e, w, b), all the gradients are governed by the two expectation terms E[(f1 X + f 2 )X] and E[f 1 X + f 2 ]. This observation plays a key role in obtaining an ordinary differential equation which yields the energy function E, defined in Eq. ( 11). Now we characterize the Hessian at both local-minima and saddle points. Lemma 4 (Hessian at local-minima and saddle points). For the canonical parameterization γ = (b, e, w) ∈ R 3 and the next-token prediction loss L(•) in Eq. (23) , the Hessian at any γ min ∈ Γ min or γ sad ∈ Γ sad is given by∇ 2 L(γ) γ=γ min ,γ sad = π 0 π 1 1 0 0 0 2(p + q -1)(1 + 2w|w|) 0 0 0 0, where π 0 = q p+q and π 1 = p p+q . Remark 3. We note that the Hessian is computed with the parameter ordering (b, e, w).The proofs of the lemmas are presented in App. K.E Analysis of the loss without bias, L(θ), and proof of Thm. 1The proof of Thm. 1, concerning the loss L(θ) in Eq. ( 5), is similar to that of Thm. 7 which studies the loss L(θ, b) with the bias present. The main idea is to establish the analogous set of lemmas, as in App. D, when the bias is substituted with its optimal choice. First we recall the loss function EQUATIONEQUATION)Consequently, if e 2 (1 + 2w|w|) = log (1-p)(1-q) pq , then b ⋆ -e 2 2 = log p 1-p . If e = 0, then b ⋆ = log p q .Now we rewrite the loss in terms of the logistic function. Lemma 6 (Loss as a logistic function). For any θ ∈ R 2 , the next-token prediction loss L(θ) in Eq. ( 29) can be written asEQUATION)where b ⋆ follows from Eq. (30), (X, Y ) ∈ {0, 1} 2 are distributed according to (X, Y ) ∼ (π, P ), i.e. X is a Bernoulli random variable with X ∼ π ≡ Bern(p/(p + q)) and Y |X ∼ P (p, q), the Markov kernel.The following lemma establishes the gradients of the loss. Lemma 7 (Gradient computation). For any θ = (e, w) ∈ R 2 and the next-token prediction loss L(θ) in Eq. ( 29), the gradients are given by∂L ∂e = E X [(f 1 X + f 2 )X] • 2(1 + 2w|w|)e, ∂L ∂w = E X [(f 1 X + f 2 )X] • 4e 2 |w|, where X ∈ {0, 1} is a Bernoulli random variable with X ∼ Bern(p/(p + q)), f 1 = σ 2e 2 w|w| + b ⋆ + e 2 2 +q-1-σ b ⋆ -e 22 +p, andf 2 = σ b ⋆ -e 2 2 -p. Further, π 1 f 1 +f 2 = 0.Remark 4. We observe above that the gradients for both e and w are proportional to each other, except for the scaling factors in terms of e and w. This forms the basis for the derivation of the energy function discussed in App. F.The following lemma characterizes the Hessian. Lemma 8 (Hessian computation). Let γ = (b, θ) ∈ R 3 with θ = (e, w) ∈ R 2 , and L(γ) be the next-token prediction loss in Eq. (23) and L(θ) be the one in Eq. (29) . Let the Hessian of L at γ beH(γ) ≜ ∇ 2 γγ L = H bb H bθ H ⊤ bθ H θθ = ∇ 2 bb L ∇ 2 bθ L (∇ 2 bθ L) ⊤ ∇ 2 θθ L ∈ R 3×3 .Then the Hessian of L at θ ∈ R 2 is given byEQUATIONConsequently, for any γ = (b, e, w) ∈ Γ min ∪ Γ sad , the Hessian H(θ) at θ = (e, w) is given byEQUATIONwhere π 0 = q p+q and π 1 = p p+q .The proofs of the above lemmas are deferred to App. L. We are now ready to present the proof of Thm. 1.E.1 Proof of Thm. 1Proof. Let θ ∈ R 2 and γ(θ) = (θ, b ⋆ (θ) ∈ R 3 be its embedding in R 3 with the optimal bias b ⋆ (θ) = argmin b∈R L(θ, b) from Lemma 5. Define the following four sets of points:Θ ⋆ (p, q) ≜ (e, w) ∈ R 2 : e 2 (1 + 2w|w|) = log (1 -p)(1 -q) pq , Θ min (p, q) ≜ (e, w) ∈ R 2 : e = 0, (p + q -1)(1 + 2w|w|) > 0 , Θ max (p, q) ≜ (e, w) ∈ R 2 : e = 0, (p + q -1)(1 + 2w|w|) < 0 , Θ sad (p, q) ≜ (e, w) : e = 0, w = -1/ √ 2 .First we show that any critical point of L : R 2 → R has to lie in one of these sets. Then we characterize that they correspond to the set of all global minima, local minima & maxima, and saddle points respectively. pq from Eq. ( 24) and hence θ ∈ Θ ⋆ . Further, by Lemma 5, we have that the optimal bias for this θ satisfies b ⋆ -e 2 2 = log p 1-p , which is precisely the characterization of the bias b for γ = (e, w, b) in Eq. (24) . Likewise, if γ ∈ Γ min ∪ Γ sad , we have e = 0 and hence θ ∈ Θ min ∪ Θ max ∪ Θ sad . Hence by Lemma 5, b ⋆ = log p q , matching that of Eq. ( 25) and Eq. (26) . Thus the set of all critical points of L in R 3 are of the form (θ, b ⋆ (θ)) with where θ ∈ Θ ⋆ ∪ Θ min ∪ Θ max ∪ Θ sad . Since Γ ⋆ ∪ Γ min ∪ Γ sad covers the entirety of stationary points of L in R 3 , it follows that the set of all stationary points in R 2 is precisely Θ ⋆ ∪ Θ min ∪ Θ max ∪ Θ sad . Also, the ordering of losses directly follows from the aformentioned observation. Now we characterize these critical points in terms of the extrema.(ii) Set of global and local minima: From Eq. ( 24), for any global minimumγ ⋆ = (θ ⋆ , b ⋆ (θ ⋆ )) of L in R 3 , we have θ ⋆ ∈ Θ ⋆ ⊆ R 2 .Hence by definition, Θ ⋆ is the set of all global minima in R 2 . A similar argument holds for Θ min , which establishes that it is a set of all local minima.(iii) Set of local maxima and saddle points: From Eq. ( 26), for any saddle point γ = (e, w, b ⋆ (e, w)) of L in R 3 , we have that e = 0 and (p + q -1)(1 + 2w|w|) ≤ 0. Hence θ = (e, w) ∈ Θ max ∪ Θ sad . Suppose θ ∈ Θ max which implies e = 0, (p + q -1)(1 + 2w|w|) < 0. By Lemma 8, the Hessian at θ (upto a positive scale) is a diagonal matrix with the entries (p + q -1)(1 + 2w|w|) < 0 and 0, corresponding to the directions of e and w respectively. Though one of the eigenvalue here is zero, using a continuity argument as in the proof of Thm. 7 for local minima, we can establish that θ is indeed a local maximum. Thus Θ max is a set of local minima. Now suppose (e, w) ∈ Θ sad . Thus e = 0 and w = -1 √ 2 . Since it lies at the intersection of Θ min and Θ max , using a neighborhood argument, it's straightforward to see that Θ sad is indeed a set of saddle points.Finally it follows that Θ min , Θ max , Θ sad are the only set of local minima, maxima, and saddle points from the above fact about the characterization of the set of all critical points in terms of these sets and Θ ⋆ , the ordering of the losses, and using the same argument as in the final steps of the proof of Thm. 7 with the bias. This concludes the proof.In this section, we analyze the learning dynamics of the transformer parameters θ = (e, w) ∈ R 2 without the attention scalar. First, we present few important lemmas regarding the same, useful for the proofs of Thm. 2 and Thm. 8 later. Recall from Sec. 4 that the trajectory (θ t ) t≥0 is governed byEQUATIONstarting with a randomly initalized θ 0 . The energy function E(•, •) is defined asEQUATIONwhere e-axis ≜ {(e, w = 0)} and w-axis ≜ {(e = 0, w)}. Note thatE sad = E(0, -1 √ 2 ) = -1+log 2 2. We re-present the Lemma 1 from Sec. 4 below for the sake of completeness. Lemma 9 (Constant energy along the flow). For any (p, q) ∈ (0, 1) 2 and initialization θ 0 = (e 0 , w 0 ) ∈ R 2 , let (θ t ) t≥0 be the corresponding GF trajectory starting from θ 0 . If w 0 ̸ = 0, then the energy stays constant along the trajectory, i.e.EQUATION)On the other hand, if w 0 = 0, w t = 0 for all t ≥ 0. Hence, if we initialize on e-axis the trajectory always stays on the e-axis.Now we establish that the GF trajectories always converge. Lemma 10 (GF convergence). Let (θ t ) t≥0 be a continuously diferentiable GF trajectory starting from θ 0 . Then for all initializations θ 0 ∈ R 2 ,(i) (θ t ) t≥0 is bounded, (ii) there exists a θ lim ∈ R 2 such that lim t→∞ θ t = θ lim and (iii) lim t→∞ ∥∇L(θ t )∥ = ∥∇L(θ lim )∥ = 0.Hence θ lim is a critical point of L.The following result characterizes the energy of the limit point. Lemma 11 (Energy at the limit point). Consider the same setting as in Lemma 10. If θ 0 ∈ R 2 \e-axis, then E(θ lim ) = E(θ 0 ). Hence θ lim lies at the intersection of the contour line E(e, w) = E 0 with the set of critical points of L in R 2 .On the other hand, if θ 0 ∈ e-axis, then θ lim ∈ e-axis.We now study the energy function on the w-axis which plays a key role in the GF analysis. Lemma 12 (Analysis of the energy function). Let E(•, •) be the energy function defined in Eq. (48) and f (w) ≜ E(e = 0, w) = -(w 2 + sign(w) • log |w|) be the energy evaluated on w-axis for w ∈ R \ {0}. Then(i) f : (-∞, -1/ √ 2] → (-∞, E sad ] is monotonically increasing with lim w→-∞ f (w) = -∞ and the maximum being f (-1/ √ 2) = E sad , (ii) f : [-1/ √ 2, 0) → [E sad , -∞) is monotonically decreasing with lim w→0 -f (w) = -∞, (iii) f ′ (-1 √ 2 ) = 0, and (iv) f : (0, ∞) → (-∞, ∞) is monotonically decreasing with lim w→0 + f (w) = ∞ and lim w→∞ f (w) = -∞.We are now ready to prove Thm. 2 corresponding to p + q > 1.F.1 Proof of Thm. 2Proof. Let θ 0 = (e 0 , w 0 ) ∈ R 2 be the initialization for the GF trajectory (θ t ) t≥0 . Recall that II ⋆ ≜ R 2 \ (I min ∪ I sad ∪ I max ) .We consider the cases θ 0 ∈ e-axis and θ 0 ∈ R 2 \ e-axis separately. First recall from Thm. 1 and Eq. ( 6) that for p + q > 1, the loci of the global minima, e 2 (1 + 2w|w|) = log (1-p)(1-q) pq < 0, lies entirely in the negative half-plane corresponding to w < -1 √ 2 . On the other hand, all the local minima, maxima and the saddle points span the w-axis corresponding to e = 0. 9), (10) , and (11) , we have that the trajectory (θ t ) t≥0 always stays on the contour line E(e, w) = E 0 and converges to the limit θ lim which is an intersection of this contour line with the set of critical points of L. Hence the crux of the proof is to establish where these intersections occur based on the initialization θ 0 and the initial energy E 0 . This gives rise to the set of initializations I min , I max , I sad , and I ⋆ that correspond to the limit being a local minimum/maximum, a saddle point, or a global minimum.(i) θ 0 ∈ R 2 \ e-axis: Let E 0 = E(θ 0 ) ∈ R. By Lemmas. (We characterize them individually below starting with I min .Initializations for local minima, I min . For θ 0 = (e 0 , w 0 ) ∈ R 2 \ e-axis, assume that w 0 > 0. Since E 0 ∈ R, there exists an unique w ⋆ > 0 such that f (w ⋆ ) = E(0, w ⋆ ) = E 0 by Lemma 12, (iv). Further using the fact that the energy contour lines do not cross each other (by definition of a contour line) and the fact they do not intersect the e-axis (it's an energy barrier as discussed in Sec. 4), it follows that the contour line E(e, w) = E 0 stays entirely in the positive half-plane corresponding to w > 0 and w ⋆ > 0 is the unique (and only) intersection of this line with the w-axis, and hence the set of critical points. Since the w-axis corresponding to w > 0 is a set of a local minima (Eq. ( 7)), it follows that any initialization (e 0 , w 0 ) with w 0 > 0 converges to a local minimum. Now suppose -1 √ 2 < w 0 < 0 and e 0 ∈ (-g(w 0 ), g(w 0 )), where g(w 0 ) = w 2 -log(-w) + E sad . Thus |e 0 | < g(w 0 ) and hence e 2 0 -(w 2 0 -log(-w 0 )) = E(e 0 , w 0 ) = E 0 < E sad . Hence by Lemma 12, (iii), there is a unique intersection of the contour line E(e, w) = E 0 with the w-axis, which lies in the region -1 √ 2 , 0 . Further note that this contour line cannot intersect with the global minima loci as it lies in the half-plane w < -1 √ 2 , and hence its only intersection with the set of critical points is this segment of w-axis, which is precisely the set of local minima the GF initialized on this line would converge to.Thus we have shown that any initialization in I min \∪ {(e, w) : w = 0} converges to a local minimum, the set of which exhausts all the set of local minima Θ min except for the origin. Below we will estbalish that any initialization on e-axis = {(e, w) : w = 0} converges to the origin, implying I min is the full set of initializations for which the limit is a local minimum.Initializations for saddle points, I sad . It's straightforward to see that for any θ 0 ∈ I sad , e 2 0 -(w 2 0 - log(-w 0 )) = E(0, -1 √ 2 ) = E sad . Since -1 √ 2 ≤ w 0 < 0, the point (w, e) = (-1 √ 2 , 0)is the only intersection of the contour line with the set of critical points, any initialization in I sad converges to the saddle point. On the other hand, there also exists a contour line e 2 0 -( 2 and further intersecting with the global minima loci Θ ⋆ . However, if we initialize on this line the flow escapes away from the saddle point and converges instead to a global minimum. To show this, it suffices to prove that det dt > 0 and dwt dt < 0 if e 0 > 0 and w 0 < -1 √ 2 , such that (w 0 , e 0 ) is close to the saddle point (-1 √ 2 , 0) (the case for e 0 < 0 is similar as the flow is symmetric in e ∈ R). From Lemma 7 and the definition of the GF, we have thatw 2 0 -log(-w 0 )) = E sad for w 0 < -1 √ 2 that passes through (-1 √ 2 , 0) ∈ Rde t dt = - ∂L ∂e (e 0 , w 0 ) = 2E X [(f 1 X + f 2 )X] • (1 -2w 2 0 ))e 0 dw t dt = - ∂L ∂w (e 0 , w 0 ) = 4E X [(f 1 X + f 2 )X] • (-e 2 0 w 0 ). So it suffices to show that E X [(f 1 X + f 2 )X] > 0.To establish this, we have from Lemma 5 thatE X [(f 1 X + f 2 )X] = E[X](f 1 + f 2 ) = π 1 - f 2 π 1 + f 2 = -π 0 • f 2 .From the defintion of f 2 and the optimal bias b ⋆ in Lemma 7 and Lemma 5 respectively, we obtainf 2 = σ b ⋆ - e 2 0 2 -p = 1 + exp -b ⋆ + e 2 0 2 -1 -p =     1 + 2A p q -1 + p q -1 2 + 4 • p q • A     -1 -p, A ≜ exp(e 2 0 (1 -2w 2 0 )).When e 0 = 0, we have A = 1 and henceEQUATIONwhere we used the fact that p + q > 1. Hence by continuity of f 2 in e 0 , for e 0 sufficiently close to 0, f 2 < 0 which proves our claim about the direction of the flow close to the saddle point. By using the continuity of the flow, it follows that GF cannot converge to saddle point when initialized on this contour line for w 0 < -1 √ 2 . Thus I sad is the only set of initializations for convergence to Θ sad . Initializations for local maxima, I max . If p + q > 1, we have from Thm. 1 that Θ max = (e, w) ∈ R 2 : e = 0, (1 + 2w|w|) < 0 = (e, w) ∈ R 2 : e = 0, w < -1 √ 2 . Thus for any θ 0 ∈ Θ max , dθt ddt = 0 for all t ≥ 0 and hence θ lim = θ 0 . Further if we slightly perturb away from this set, from Eq. ( 36) it follows that the flow diverges and hence it's an unstable set of critical points (they are local maxima indeed). Thus the only set of initializations leading to local maxima are I max = Θ max .Initializations for the global minima, I ⋆ . Since the set of all critical points of L is Θ ⋆ ∪ Θ min ∪ Θ max ∪ Θ sad , and the initializations in I min , I sad , and I max converge to Θ min , Θ sad , and Θ max respectively, it follows that the set of initializations for which the GF converges to global minima isI ⋆ = R 2 \ (I min ∪ I sad ∪ I max ).In fact, since the loci of the global minima lies in the half-plane correspondint to w < -1 √ 2when p + q > 1, we can precisely determine the location of the global minimum for which the intersection occurs for any θ 0 ∈ I ⋆ . Specifically, we can solve the pair of equations E(e, w) = e 2 -w 2 + log(-w) = E 0 and e 2 (1 -2w 2 ) = log (1-p)(1-q) pq which has a unique solution for w < 0 (upto a sign flip in e).(ii) θ 0 ∈ e-axis ⇒ θ 0 ∈ I min : If θ 0 = (e 0 , w 0 ) ∈ e-axis, we have that w 0 = 0 and hence w t = 0 for all t ≥ 0 (Lemma 9). Lemma 10-(i) also establishes that the iterates (θ t = (e t , 0)) t≥0 stay bounded on the e-axis and monotonically decrease. Since the origin is the only critical point of L on the e-axis, and lim t→∞ θ t = θ lim exists, it follows that θ lim = (0, 0), a local minima. Thus θ 0 ∈ I min . This concludes the proof for all the initializations θ 0 ∈ R 2 .Gaussian initialization N (0, σ 2 I 2 ). When θ 0 is initialized according to the standard Gaussian distribution N (0, σ 2 I 2 ) with σ 2 ≪ 1 √ 2 , we note that θ 0 lands in the set I min with high probability. In fact, this probability can be made arbitrarily close to 1 depending on σ 2 . Thus this initialization will lead to a local minimum convergence on the w-axis.F.2 Gradient flow dynamics for p + q < 1 Theorem 8 (GF dynamics for p + q < 1). Under the same setting as in Thm. 2 with p + q < 1, and any initialization θ 0 ∈ R 2 , the GF trajectory always converges to a θ lim ∈ R 2 which is a critical point of the loss L. More specifically, θ lim is(i) a local minimum if θ 0 ∈ I min ≜ (e, w) : w < -1/ √ 2, e ∈ (-g(w), g(w)), g(w) = w 2 -log(-w) + E sad , (ii) a saddle point if θ 0 ∈ I sad ≜ (e, w) : w ≤ -1/ √ 2, e = ± w 2 -log(-w) + E sad , (iii) a local maximum if θ 0 ∈ I max ≜ (e, w) : e = 0, w > -1/ √ 2 , (iv) and a global minimum if θ 0 ∈ R 2 \ (I min ∪ I sad ∪ I max ).Consequentely, if we use the standard initialization θ 0 ∼ N (0, σ 2 I 2 ) with σ 2 ≪ 1/ √ 2, θ lim will be a global minimum.Proof. The proof for the case of p + q < 1 essentially follows the same steps as that of p + q > 1. If the initialization is not on the e-axis we use the energy equation to establish the convergence to the critical point at the intersection of the energy contour line with the critical set and if it starts on the e-axis, the only change is that it now converges to the global minimum instead of the origin as in the earlier case. This is due to the fact that origin turns out to be a local maximum when p + q < 1 and hence it's an unstable critical point (which can be established as in the proof of Thm. 2 for I max ).In this section, we analyze the learning dynamics of the transformer parameters θ ∈ R 3 with the attention scalar a ∈ R, i.e. θ = (e, w, a) ∈ R 3 . Similar to the analysis for θ = (e, w) ∈ R 2 , we first introduce the canonical parameterization including a ∈ R, then analyze the corresponding loss function L(•) in terms of its gradients and critical points, and capitalize on it to study the gradient flow dynamics using the energy. We first start with the parameterization. d. This results in the embeddingx n = e x n - 1 2 α.Attention. Similarly, we recall from Eq. ( 21) that the attention output y n is given byEQUATIONwhereatt n,i ≜ exp ⟨q n , k i ⟩/ √ d /   j∈[n] exp ⟨q n , k j ⟩/ √ d   , q n = W Q x n , k i = W K x i , W ⊤ Q W K = (q 2 d) α • α ⊤ ∈ R d×d ,for some q ∈ R. Instead of the softmax, now we assume that the attention weights are linear in the scaled dot product, i.e.EQUATION(EQUATIONNote that the 1/n factor is to ensure normalization for the attention weights in Eq. (37) . Now substituting Eq. ( 38) in Eq. ( 37), we obtainy n = e x n - 1 2 α + ⟨v, α⟩   i∈[n] att n,i • e x i - 1 2   α =   e x n - 1 2 + ⟨v, α⟩   i∈[n] 1 n q 2 d 5/2 e 2 (x n - 1 2 )(x i - 1 2 )   • e x i - 1 2   α = e x n - 1 2 1 + ⟨v, α⟩q 2 d 5/2 e 2 x i - 1 2 2 α =   e x n - 1 2   1 + ⟨v, α⟩q 2 d 5/2 1 4 a •e 2 •       α = e x n - 1 2 1 + ae 2 α,where we used the fact that (x i -1 2 ) 2 = 1 4 since x i ∈ {0, 1}, andEQUATION)is the attention scalar. Note that this includes the scaling ⟨v, α⟩ from the value matrix W V and q 2 from the query-key dot product. Thus we succinctly haveEQUATIONFeed-forward. For the feed-forward layer, we have thatW 1 = |w| √ d 1 • α ⊤ ∈ R 4d×d , W 2 = w √ d α • 1 ⊤ ∈ R d×4d. Hence Eq. ( 40) impliesW 1 y n = |w| √ d 1 • α ⊤ e x n - 1 2 1 + ae 2 α = |w| √ d e x n - 1 2 1 + ae 2 1.Thus,ReLU(W 1 y n ) = |w| √ d 1 • ReLU e x n - 1 2 1 + ae 2 = |w| √ d 1 • e ReLU x n - 1 2 1 + ae 2 = |w| √ d 1 • e x n 2 ReLU 1 + ae 2 + 1 -x n 2 ReLU -1 -ae 2 = |w| 2 √ d 1 • e x n ReLU 1 + ae 2 -ReLU -1 -ae 2 + ReLU -1 -ae 2 .Using ReLU(x) -ReLU(-x) = x above,ReLU(W 1 y n ) = |w| 2 √ d 1 • e x n 1 + ae 2 + ReLU -1 -ae 2 .Hence,W 2 ReLU(W 1 y n ) = w √ d α • 1 ⊤ |w| 2 √ d 1 • e x n 1 + ae 2 + ReLU -1 -ae 2 = 2w|w|e x n 1 + ae 2 + ReLU -1 -ae 2 α = 2w|w|e x n 1 + ae 2 + -1 -ae 2 2 + |1 + ae 2 | 2 α = 2w|w|e x n - 1 2 1 + ae 2 + |1 + ae 2 | 2 α.Thus the embedding z n is given byz n = y n + W 2 ReLU(W 1 y n ) = e x n - 1 2 1 + ae 2 α + 2w|w|e x n - 1 2 1 + ae 2 + |1 + ae 2 | 2 α = e x n - 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| α.Linear. Since e = a = e • α due to weight-tying, the logits are given bylogit n (e, w, a, b) = ⟨a, z n ⟩ + b = e 2 x n - 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b.Loss. Denote θ ≜ (e, w, a) ∈ R 3 . Similar to the case without a (Eq. ( 14) and Lemma 2), the cross-entropy loss in our setting can be compactly written asEQUATIONwhere logit X (θ, b) ≜ e 2 X -1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| +b, (X, Y ) ∈ {0, 1} 2 are distributed according to (X, Y ) ∼ (π, P ), i.e. X is a Bernoulli random variable with X ∼ π ≡ Bern(p/(p + q)) and Y |X ∼ P (p, q), the Markov kernel. Further, using the convexity of b in L(•, b), we can consider the optimal bias b ⋆ (θ) = argmin b∈R L(θ, b) in Eq. ( 41) to obtain the loss L(θ):EQUATIONWe derive the expression for the optimal bias b ⋆ in the proof of Lemma 13 below in App. N.G.2 Analysis of the loss function L(θ) from Eq. (42)Now we establish the gradients of the loss function.Lemma 13 (Gradient computation and optimal bias). For any θ = (e, w, a) ∈ R 3 and the next-token prediction loss L(θ) in Eq. (42), the gradients are given by∂L ∂e = -E (f 1 X + f 2 ) X - 1 2 • 2e 1 + ae 2 (1 + 2w|w|) -E (f 1 X + f 2 ) X - 1 2 • 2e 3 a (1 + 2w|w|) , ∂L ∂w = -E (f 1 X + f 2 ) X - 1 2 • 2e 2 1 + ae 2 (|w| + sign (w) w) , ∂L ∂a = -E (f 1 X + f 2 ) X - 1 2 • e 4 (1 + 2w|w|) ,where X ∈ {0, 1} is a Bernoulli random variable with X ∼ Bern(p/(p + q)), andf 1 ≜ 1 -p -q -ϕ 1 + ϕ 0 , f 2 ≜ p -ϕ 0 , ϕ 1 ≜ σ e 2 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b ⋆ , ϕ 0 ≜ σ e 2 -1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b ⋆ ,where the optimal bias b ⋆ is obtained by solving π 1 f 1 + f 2 = 0.Proof. We defer to App. N.Theorem 9 (All critical points for linear attention in R 4 ). Let the input sequence be {x n } N n=1 ∼ (π, P ), the transformer parameters θ = (e, w, b, a) ∈ R 4 , and the next-token prediction loss L(•) be as in Eq. (41). Then for any (p, q) ∈ (0, 1) 2 with p + q ̸ = 1 and N ∈ N, (i) the set of all global minima is given byEQUATIONEQUATION(ii) a set of local minima is given by Γ min (p, q) ≜ γ min = (e, w, b, a) ∈ R 4 : e = 0, (p + q -1)(1 + 2w|w|) > 0, b = log p q , (iii) a set of saddle points is Γ sad (p, q) ≜ γ sad = (e, w, b, a) ∈ R 4 : e = 0, (p + q -1)(1 + 2w|w|) ≤ 0, b = log p q .(iv) a set of stationary points isΓ station (p, q) ≜ γ station = (e, w, b, a) ∈ R 4 : e ̸ = 0, 1 + ae 2 = 0, 1 + 2w|w| = 0, b = log p q ,Thus the set of all critical points isθ ∈ R 2 : ∇L(θ) = 0 = Γ ⋆ ∪ Γ min ∪ Γ sad ∪ Γ station .(45) In addition, for any θ ⋆ ∈ Γ ⋆ , θ min ∈ Γ min and θ sad ∈ Γ sad , the loss values satisfyH(x n+1 | x n ) = L(θ ⋆ ) < L(θ min ) = L(θ max ) = L(θ sad ) = H(x n+1 ).Theorem 10 (All critical points in R 3 ). Let the input sequence be {x n } N n=1 ∼ (π, P ), the transformer parameters θ = (e, w, a) ∈ R 3 , and the next-token prediction loss L(•) be as in Eq. (42). Then for any (p, q) ∈ (0, 1) 2 with p + q ̸ = 1 and N ∈ N, (i) the set of all global minima is given byEQUATION(ii) a set of local minima is given by Θ min (p, q) ≜ (e, w, a) ∈ R 3 : e = 0, (p + q -1)(1 + 2w|w|) > 0 , (iii) a set of local maxima is given by Θ min (p, q) ≜ (e, w, a) ∈ R 3 : e = 0, (p + q -1)(1 + 2w|w|) < 0 , (iv) a set of saddle points isΘ sad (p, q) ≜ (e, w, a) ∈ R 3 : 0, -1/ √ 2, a .Defining a set of stationary points Θ station (p, q) ≜ (e, w, a) ∈ R 3 : e ̸ = 0, 1 + ae 2 = 0, 1 + 2w|w| = 0 , the set of all critical points isθ ∈ R 2 : ∇L(θ) = 0 = Θ ⋆ ∪ Θ min ∪ Θ max ∪ Θ sad ∪ Θ station .(47) In addition, for any θ ⋆ ∈ Θ ⋆ , θ min ∈ Θ min , and θ sad ∈ Θ sad , the loss values satisfyH(x n+1 | x n ) = L(θ ⋆ ) < L(θ min ) = L(θ max ) = L(θ sad ) = H(x n+1). Remark 5. While the remaining set of stationary points could be classified to global minima, local minima, etc., it's technically unclear what category the set of critical points Θ station (p, q) belong to, as the Hessian is undefined here. We would need to rely on local perturbation analysis to characterize these class of points.We defer the proofs of the theorems to App. N.2.Analogous to the gradient flow analysis for θ = (e, w) ∈ R 2 in App. F, we now study its countepart toegether with the attention scalar, i.e. θ = (e, w, a) ∈ R 3 . To this end, let (θ t ) t≥0 be a C 1 curve in R 3 governed bydθ t dt = -∇L(θ t ), θ t = (e t , w t , a t ) ∈ R 3 , t ≥ 0, (GF-Attention)starting with a randomly initalized θ 0 . We define the energy function E(•, •, •) asEQUATIONwhere ea-plane ≜ {(e, w = 0, a)}. The following lemma presents the crucial result that the energy is constant along the flow in GF-Attention.Lemma 14 (Constant energy along the flow). For any (p, q) ∈ (0, 1) 2 and initialization θ 0 = (e 0 , w 0 , a 0 ) ∈ R 3 , let (θ t ) t≥0 be the corresponding GF-Attention trajectory starting from θ 0 . If w 0 ̸ = 0, then the energy stays constant along the trajectory, i.e.E(θ t ) = e 2 t -(w 2 t + sign(w t ) • log |w t |) -2a 2 t = E(θ 0 ), ∀t ≥ 0.(49) On the other hand, if w 0 = 0, w t = 0 for all t ≥ 0. Hence, if we initialize on ea-plane the trajectory always stays on the ea-plane. Now we characterize the convergence of the gradient flow. Lemma 15 (GF convergence). Let (θ t ) t≥0 be a continuously diferentiable GF-Attention trajectory starting from θ 0 . Then for all initializations θ 0 ∈ R 3 , (i) (θ t ) t≥0 is bounded, (ii) there exists a θ lim ∈ R 3 such that lim t→∞ θ t = θ lim and (iii) lim t→∞ ∥∇L(θ t )∥ = ∥∇L(θ lim )∥ = 0.Hence θ lim is a critical point of L.The following result characterizes the energy of the limit point. Lemma 16 (Energy at the limit point). Consider the same setting as in Lemma 15. If θ 0 ∈ R 3 \ ea-plane, then E(θ lim ) = E(θ 0 ). Hence θ lim lies at the intersection of the contour line E(e, w) = E 0 with the set of critical points of L in R 3 .On the other hand, if θ 0 ∈ ea-plane, then θ lim ∈ ea-plane.We defer the proofs of the lemmas to App. N.The following picture enhances the behavior of the GF dynamics near the origin.(a) Gradient flow around origin (p + q < 1) (b) Gradient flow around origin (p + q > 1)Figure 5 : Gradient flow dynamics in R 3 , near the origin, for the transformer parameters with attention scalar a (Sec. 4.1). The local minima are repellors for p + q < 1, while attracting for p + q > 1.Based on the above theoretical results and empirical evidence, we conjecture the following theorem, whose informal proof we defer to App. N.5. Theorem 11 ( [Informal] Role of standard initialization for p + q > 1). If we use the standard initialization θ 0 ∼ N (0, σ 2 I 3 ) with σ 2 ≪ 1/ √ 2 for the GF-Attention, θ lim will be a local minimum with high probability. Proof. We characterize the set of global minima, local minima, and that of the saddle points individually.(i) Set of all global minima. Let γ ⋆ ∈ R 3 be arbitrary. From [23, Lemma 1], we have that γ ⋆ is a global minimum for the loss L(•) in Eq. ( 23) if and only if its prediction probability satisfies f γ ⋆ (x n 1 ) = P (x n+1 = 1 | x n ), the Markov kernel. Since the input {x n } N n=1 ∼ (π(p, q), P (p, q)), we have thatEQUATION)On the other hand, by definition, from Eq. ( 3),f γ ⋆ (x n 1 ) = σ e 2 (1 + 2w|w|) x n + b -e 2 2, where γ ⋆ = (e, w, b). Since x n ∈ {0, 1}, this can be further simplified tof γ ⋆ (x n 1 ) = σ e 2 (1 + 2w|w|) x n + b - e 2 2 = x n • σ e 2 (1 + 2w|w|) + b - e 2 2 + (1 -x n ) • σ b - e 2 2 = x n σ 2e 2 w|w|) + b + e 2 2 -σ b - e 2 2 + σ b - e 2 2 .(51)Since both f γ ⋆ (x n 1 ) and P (x n+1 = 1 | x n ) are linear functions of x n , equating them for all vallues of x n ∈ {0, 1} implies that the respective coeffecients in these functions in Eq. ( 50) and Eq. ( 51) are also equal, i.e.σ b - e 2 2 = p, σ 2e 2 w|w| + b + e 2 2 -σ b - e 2 2 = 1 -p -q,and henceEQUATION)Since σ(z) = y for y ∈ (0, 1) implies z = log y 1-y , Eq. ( 52) can be rewritten asEQUATION)Thus γ ⋆ ∈ R 3 is a global minimum for L(•) if and only if it satisfies Eq. ( 53) (note that it's already a critical point, as established in Thm. 6). Thus, the set of all global minimum Γ ⋆ (p, q) is given byΓ ⋆ (p, q) ≜ γ ⋆ = (e, w, b) ∈ R 3 : e 2 (1 + 2w|w|) = log (1 -p)(1 -q) pq , b - e 2 2 = log p 1 -p .Since the prediction f γ ⋆ (•) equals the Markov kernel for any γ ⋆ ∈ Γ ⋆ , it follows from Thm. 4 (or [23, Lemma 1] ) that L(γ ⋆ ) = H(x n+1 | x n ), the entropy rate of the Markov chain.(ii) Set of local minima and saddle points.Define Γ min (p, q) ⊆ R 3 and Γ sad ⊆ R 3 as follows:Γ min (p, q) ≜ γ min = (e, w, b) ∈ R 3 : e = 0, (p + q -1)(1 + 2w|w|) > 0, b = log p q , Γ sad (p, q) ≜ γ sad = (e, w, b) ∈ R 3 : e = 0, (p + q -1)(1 + 2w|w|) ≤ 0, b = log p q .To show that Γ min is the set of all bad local minima for L(•), we first show that any γ min ∈ Γ min is a bad local minimum and then show that every bad local minimum should belong to Γ min . Similarly for Γ sad . We start with the local minima.Let γ min = (e, w, b) ∈ Γ min . Recall that γ min is a stationary point (Thm. 6), i.e.∇L(γ min ) = 0.Rearragning the order of scalars and writing γ min = (b, e, w), from Lemma 4, the Hessian of the loss at γ min isEQUATION)By definition, γ min = (b, e, w) satisfies (p + q -1)(1 + 2w|w|) > 0. Thus its Hessian in Eq. ( 54) has a block diagonal structure of the form H b,e 0 0 0 where H b,e has both the eigen values positive, and hence positive-definite. In other words, γ min is a local minimum for L(•) in the (b, e) ∈ R 2 space for any fixed w in the set. Interestingly, using the continuity argument and the fact that L(b = log p q , e = 0, w) is constant in w ∈ R, we can essentially follow the same steps as in proof of Theorem 2 in [23, Appendix B.3] (Thm. 5 above) to show that γ min = (b, e, w) is a also local minimum for L(•) in the full parameter space R 3 . This establishes that γ min is a local minimum for L(•).For the saddle points, let γ sad = (e, w, b) ∈ Γ sad . We have that γ sad is a stationary point (Thm. 6) and Lemma 4 implies its Hessian (after rearraging the order of scalars as above with γ sad = (b, e, w)) is:EQUATIONIf w ̸ = -1 √ 2 , (p + q -1)(1 + 2w|w|) < 0 for any γ sad ∈ Γ sad , and hence the Hessian ∇ 2 L(γ sad ) in Eq. (55) as both positive, negative, and zero eigen values. Thus γ sad is a saddle point for L(•). Using a neighborhood argument, we can similarly argue for w = 1 √ 2 to establish that it's also a saddle point. Now we compute the loss value.For any γ min = (e, w, b) ∈ Γ min or γ sad = (e, w, b) ∈ Γ sad , we have that e = 0 and b = log p q . Thus for γ = γ min or Γ sad , the prediction probability in view of Eq. ( 3) isP γ (x n+1 = 1 | x n 1 ) = σ e 2 (1 + 2w|w|) x n + b - e 2 2 = σ(b) = p p + q = P (x n+1 = 1) ,the marginal distribution. Substituting this equality in the definition of cross-entropy loss L(•) in Eq. ( 1) and the fact that P (x n+1 = 1) = p p+q = π 1 , following the same steps as in [23, Appendix B.3], we obtainL(γ) = - 1 N n∈[N ] E x n+1 1 [x n+1 • log f γ (x n 1 ) + (1 -x n+1 ) • log(1 -f γ (x n 1 ))] = - 1 N n∈[N ] E x n+1 1 [x n+1 • log π 1 + (1 -x n+1 ) • log π 0 ] = 1 N n∈[N ] [-π 1 log π 1 -π 0 log π 0 ] = H(π) = H(x n+1 ). Thus L(γ min ) = L(γ sad ) = H(x n+1 ). To see that H(x n+1 | x n ) = L(γ ⋆ ) < L(γ min ) = L(γ sad ) = H(x n+1) for any global minimum γ ⋆ , observe that the gapL(γ min ) -L ⋆ = H(x n+1 ) -H(x n+1 |x n ) = I(x n ; x n+1 ) ≥ 0,where I(x n ; x n+1 ) is the mutual information between x n and x n+1 [11] . Hence the optimality gap equals zero if and only if the mutual information equals zero, which happens when x n and x n+1 are independent, i.e. P (x n+1 = 1 | x n ) doesn't depend on x n . Since P (x n+1 = 1 | x n ) = (1 -p -q)x n + p from Eq. ( 50), this happens only when p + q = 1 which contradicts the theorem assumption that p+ q ̸ = 1. Hence H(x n+1 | x n ) = L(γ ⋆ ) < L(γ min ) = L(γ sad ) = H(x n+1 ).Now we finally show that Γ min and Γ sad are the only set of bad local minima and saddle points respectively. Let γ is a bad local minimum for L(•). By definition, it's also a critical point. Recall from Thm. 6 that any stationary point γ = (e, w, b) for the loss L(•) satisfies that either γ ∈ Γ ⋆ , γ ∈ Γ min , or γ ∈ Γ sad . Clearly γ ̸ ∈ Γ ⋆ , as Γ ⋆ is the set of all global minima. Similarly, γ ̸ ∈ Γ sad as every point in Γ sad is a saddle point for the loss L(•) as established above. Hence γ ∈ Γ min . Thus every bad local minimum in R 3 belongs to Γ min . This coupled with the fact above that Γ min is a set of bad local minima implies Γ min is indeed the set of all bad local minima. The proof for Γ sad is similar.J.2 Proof of Thm. 6EQUATIONEQUATIONEQUATIONwhereX ∼ Bern(p/(p + q)), f 1 = σ 2e 2 w|w| + b + e 2 2 + q -1 -σ b -e 2 2+ p, andf 2 = σ b -e 2 2-p. Our goal is to now show that Eqs. ( 56)-(58) hold only if either (e = 0, b = log p q ) or (f 1 = 0, f 2 = 0). We consider two cases corresponding to e = 0 and e ̸ = 0.(i): e = 0. If e = 0, we readily see that ∂L ∂e = ∂L ∂e = 0. Further,f 1 = p + q -1 and f 2 = σ(b) -p. Hence, Eq. (58) implies that E X [f 1 X + f 2 ] = (p + q -1)E[X] + σ(b) -p = (p + q -1) p p + q + σ(b) -p = σ(b) - p p + q = 0,which implies that b = log p q . Since w ∈ R is arbitrary, we see in this case thatγ ∈ (e, w, b) ∈ R 3 : e = 0, (p + q -1)(1 + 2w|w|) > 0, b = log p q (e, w, b) ∈ R 3 : e = 0, (p + q -1)(1 + 2w|w|) < 0, b = log p q , = Γ min ∪ Γ sad .(ii): e ̸ = 0. Suppose e ̸ = 0. Here we show that f 1 = f 2 = 0 and hence γ ∈ Γ ⋆ . We consider two cases corresponding to w ̸ = 0 and w = 0. Let w ̸ = 0. Since both e ̸ = 0 and w ̸ = 0, and X = X 2 in Eq. ( 57) with E[X] = π 1 = p p+q > 0, we obtain thatE X [f 1 X 2 + f 2 X] = (f 1 + f 2 )E[X] = (f 1 + f 2 )π 1 = 0,and hence f 1 + f 2 = 0. Further Eq. (58) implies that f 1 π 1 + f 2 = 0. Together, this impliesf 1 = f 2 = 0.Now suppose w = 0. From Eq. ( 58), we have thatE[f 1 X + f 2 ] = f 1 π 1 + f 2 = 0.Since e ̸ = 0, Eq. ( 56) yieldsE X [(f 1 X + f 2 )(2X -1))] = 2E X [(f 1 X + f 2 )X] = 2π 1 (f 1 + f 2 ) = 0.So f 1 = f 2 = 0 in this case too. Thus we have showed that whenever e ̸ = 0, we havef 1 = f 2 = 0.Recalling the expressions for f 1 and f 2 ,EQUATIONf 1 = σ 2e 2 w|w| + b + e 2 2 + q -1 -σ b - e 2 2 + p = σ 2e 2 w|w| + b + e 2 2 + q -1 = 0,and hence,2e 2 w|w| + b + e 2 2 = log 1 -q q .Substituting b -e 2 2 = log p 1-p in the above equation,2e 2 w|w| + e 2 + b - e 2 2 = 2e 2 w|w| + e 2 + log p 1 -p = log 1 -q q ,and thus,EQUATIONIn view of Eq. ( 59) and Eq. ( 60), we have that γ = (e, w, b) ∈ Γ ⋆ .Together, we have shown that whenever ∇L(γ) = 0, we have γ ∈ Γ ⋆ ∪ Γ min ∪ Γ sad . Since Γ ⋆ ∪ Γ min ∪ Γ sad ⊆ {γ : ∇L(γ) = 0}, we are done.Proof. Recall from Eq. ( 14) that the cross-entropy loss L(•) is defined asEQUATIONwheref γ (x n 1 ) = σ(logit n ) = σ e 2 (1 + 2w|w|) x n + b -e 2 2from Eq. ( 3). For any Y ∈ {0, 1}, Z ∈ R, using the fact that 1 -σ(Z) = σ(-Z), we haveEQUATIONwhere (a) and (b) follow from the definitions of sigmoid and the logistic functions: 62), Eq. (61) simplifies toσ(z) = 1 1+exp(-z) , ℓ log (z) = log(1 + exp(-z)) for z ∈ R. Substituting Y = x n+1 ∈ {0, 1} and Z = logit n ∈ R in Eq. (L(γ) = 1 N n∈[N ] E[ℓ log ((2x n+1 -1) • logit n )].Since logit n is only a function of x n , the above expectation is over the distribution of the pairs (x n , x n+1 ), which for all n ∈ [N ] have the same law as a pair of random variables (X, Y ) with X ∼ π ≡ Bern(p/(p + q)) and Y |X ∼ P (p, q), the Markov kernel. Hence the above equality can be rewritten using the definition of logit n asL(γ) = 1 N n∈[N ] E[ℓ log ((2x n+1 -1) • logit n )] = E X,Y ℓ log (2Y -1) e 2 (1 + 2w|w|)X + b - e 2 2 .Proof. With γ = (e, w, b) and θ denoting either of the scalars e, w, or b, we have from [23, Lemma 2] that the gradient of the loss L(•) is given byEQUATIONwheref γ (x n 1 ) = σ(logit n ) = σ e 2 (1 + 2w|w|) x n + b -e 2 2 .Using the same argument as in the proof of Lemma 6, we can replace the expecatations in Eq. ( 63) with that of a pair of random variables (X, Y ) with X ∼ π ≡ Bern(p/(p + q)) and Y |X ∼ P (p, q), the Markov kernel. That is,EQUATIONNow we define the error termE(X, Y ) ≜ -Y -σ e 2 (1 + 2w|w|)X + b -e 2 2. Our goal is toshow that E[E(X, Y ) | X] = f 1 X +f 2 , where f 1 ≜ σ 2e 2 w|w| + b + e 2 2 +q-1-σ b -e 2 2 +p, and f 2 ≜ σ b -e 2 2-p, which suffices to prove the lemma. To this end, using the fact that X ∈ {0, 1}, we haveE(X, Y ) = -Y -σ e 2 (1 + 2w|w|)X + b - e 2 2 = -Y -X • σ 2e 2 w|w| + b + e 2 2 -(1 -X) • σ b - e 2 2 = -Y + X σ 2e 2 w|w| + b + e 2 2 -σ b - e 2 2 + σ b - e 2 2 .Now taking the conditional expectation with respect to X and using the fact thatE[Y |X] = P (Y = 1 | X) = (1 -p -q)X + p (since Y |X ∼ P (p, q)), we have E[E(X, Y ) | X] = -(1 -p -q)X -p + X σ 2e 2 w|w| + b + e 2 2 -σ b - e 2 2 + σ b - e 2 2 = X σ 2e 2 w|w| + b + e 2 2 + q -1 -σ b - e 2 2 + p + σ b - e 2 2 -p (a) = f 1 X + f 2 ,where (a) follows from the definition of f 1 and f 2 above. Thus Eq. (64) simplifies to∇ θ L(γ) = E X (f 1 X + f 2 ) • ∇ θ e 2 (1 + 2w|w|)X + b - e 2 2 .Letting θ = e, w, and b in the above equation, we finally obtain the individual gradients:∂L ∂e = E X [(f 1 X + f 2 )(2X(1 + 2w|w| -1))] • e, ∂L ∂w = E X [(f 1 X + f 2 )X] • 4e 2 |w|, ∂L ∂b = E X [f 1 X + f 2 ] . K.3 Proof of Lemma 4Proof. Slightly changing the variable order, for any γ = (b, e, w) ∈ R 3 , we defineEQUATIONRecall that for any γ min ∈ Γ min and γ sad ∈ Γ sad , we have e = 0 and b = log p q . Now we compute the second derivatives of L with respect to any γ = (b = log p q , e = 0, w). We start with the first derivatives. By Lemma 7, the gradients areEQUATIONwheref 1 ≜ σ 2e 2 w|w| + b + e 2 2 + q -1 -σ b - e 2 2 + p, f 2 ≜ σ b - e 2 2 -p.From Eq. (66), we see that the second derivaties of L depend on the first-derivatives of f 1 and f 2 , which we now compute. Recall that the derivative of the sigmoid function obeys σ ′ (z) = σ(z)(1 -σ(z)) = σ(z)σ(-z) for any z ∈ R. Now the gradients of f 1 and f 2 with respect to b, e, and w are Using the fact that σ log p q = p p+q = π 1 and σ -log p q = q p+q = π 0 , the above gradients evaluated for any γ = (b = log p q , e = 0, w) further reduce to∂f 1 ∂b = σ 2e 2 w|w| + b + e 2 2 σ -2e 2 w|w| -b - e 2 2 -σ b - e 2 2 σ -b + e 2 2 , ∂f 2 ∂b = σ b - e 2 2 σ -b + e 2 2 ,EQUATIONNow substituting Eq. ( 67) when computing the second-derivatives of L in Eq. (66), we obtainEQUATIONwhere (a) follows from the fact that f 1 | γ = p+q-1, f 2 | γ = σ(b)-p = p p+q -p = -p p+q (p+q-1) = -π 1 (p + q -1) and (b) from 1 -π 1 = π 0 . Returning to the remaining second derivatives,EQUATIONCongregating all the second derivatives from Eq. ( 68) and Eq. ( 69) into the Hessian H(γ) in Eq. ( 65), we finally obtainH(γ) = π 0 π 1 1 0 0 0 2(p + q -1)(1 + 2w|w|) 0 0 0 0 . L Proofs of lemmas in App. E L.1 Proof of Lemma 5Proof. Recall from Lemma 2 that for any θ = (e, w) ∈ R 2 and b ∈ R, we haveL(θ, b) = E X,Y ℓ log (2Y -1) e 2 (1 + 2w|w|)X + b - e 2 2 .Since ℓ log (•) is a convex function, Y ∈ {0, 1} and thus 2Y -1 ∈ {±1}, the convexity of L in b follows from the following fact:∂ 2 L ∂b 2 = E X,Y ℓ ′′ log (2Y -1) e 2 (1 + 2w|w|)X + b - e 2 2 ≥ 0.To find the optimal b ⋆ , we set the gradient ∂L ∂b = 0. Thus from Lemma 3, we obtain∂L ∂b = E X [f 1 X + f 2 ] = 0, f 1 = σ 2e 2 w|w| + b + e 2 2 + q -1 -σ b - e 2 2 + p, f 2 = σ b - e 2 2 -p. Substituting E ≜ e 2 (1 + 2w|w|), B ≜ b -e 2 2, and E[X] = π 1 = p/(p + q) in the above equations,π 1 (σ(E + B) + q -1 -σ(B) + p) + σ(B) -p = π 1 • σ(E + B) + π 0 • σ(B) + p(p + q -1) p + q -p = 0.Further simplifying,π 0 • σ(B) = π 1 • (1 -σ(E + B)) = π 1 • σ(-E -B). In other words, (1 + exp(-B)) -1 (1 + exp(E + B)) -1 = π 1 π 0 = p q ⇒ 1 + exp(E) exp(B) 1 + exp(-B) = p q .Defining x ≜ exp(B) and A ≜ exp(E), we thus obtain the following quadratic equation in x and its corresponding roots:Ax 2 -x p q -1 - p q = 0 ⇒ x = 1 2A   p q -1 ± p q -1 2 + 4 • p q • A   .Since x > 0, we take the root corresponding to the addition choice above and resubstituting x = exp(b -e 2 2 ) and A = exp(e 2 (1 + 2w|w|)), we obtain the final expression for b ⋆ . In particular, if e = 0, it is easy to see that A = 1 and hence x = exp(b ⋆ ) = p q , implying b ⋆ = log p q . Similarly, if A = (1-p) (1-q) pq , it's straightforward to see that exp(b ⋆ -e 2 /2) = p 1-p and hence b ⋆ -e 2 /2 = log p 1-p .Proof. The proof directly follows from Lemma 2 by substituting b = b ⋆ .Proof. By Danskin's theorem [12] , it follows that for b ⋆ = argmin Proof. Denote (e t , w t ) = (e, w) with the dependence on time implicitly assumed. Then by the definition of GF and the gradient expressions in Lemma 7, we have thatEQUATIONEQUATIONDividing Eq. ( 70) by (1 + 2w|w|)e and Eq. ( 71) by 4e 2 |w|, we have Thus defining E(e, w) = e 2 -w 2 -sign(w) • log |w| for w ̸ = 0, the above equation implies E(θ t ) = E(θ 0 ) for θ 0 = (e 0 , w 0 ) with w 0 ̸ = 0. On the other hand, it's easy to see that if w 0 = 0, Eq. ( 71) implies dw dt = 0 at t = 0 and hence w t = 0 for all t ≥ 0.Proof. To prove the convergence of the trajectory (θ t ) t≥0 , we use the classical result due to Łojasiewicz [1, Theorem 2.2] which gurantees the convergence of gradient flow for real analytic functions, as long as the trajectory is bounded. Hence we first show the boundedness of the trajectory.(i) (θ t ) t≥0 is bounded.We consider the cases θ 0 ∈ e-axis and θ 0 ∈ R 2 \ e-axis separately.Let's suppose θ 0 = (e 0 , w 0 ) ∈ e-axis, i.e. w 0 = 0. Thus it follows from Lemma 9 that w t = 0 for all t ≥ 0. That is, the trajectory always stays on the e-axis and it suffices to track (e t ) t≥0 and show that they are bounded. To this end, we show that if e 0 > 0, we have de dt < 0 and if e 0 < 0, we have de dt > 0 for all t ≥ 0, which establishes our claim. We have from the GF and Lemma 7 thatEQUATIONandEQUATIONand substituting Eq. (74) and Eq. (73) in Eq. ( 72), we obtainEQUATIONSince x t = exp(e 2 t ) = exp(-e 2 t ), in view of Eq. ( 75), with out loss of generality, we can assume that e 0 > 0 and show that det dt < 0 for all t ≥ 0. That is, the RHS Eq. ( 75) is negative. Note thatx t ≥ 1 since x t = exp(e 2 t ) and g(x t ) > 0 since the denominator p q -1 +p q -1 2 + 4 • p q • x t > p q -1 + p q + 1 = 2 • p q > 0.Further g(1) = q p and lim x→∞ g(x) = ∞. If we show that g(x) is increasing in x for x ≥ 1, it implies 1 1+g(x) -p < 1 1+g(1) -p = 1 1+ q p -p = -p p+q (p + q -1) < 0. Thus the gradient in Eq. (75) remains negative starting at t = 0 and hence the sequence (e t ) t≥0 will be bounded. Now it remains to show g(•) is increasing, i.e. g ′ (•) > 0. Defining C = p q -1 / 2 p q and D = C 2 , we have that g(x) upto a postive scaling isg(x) = x C + √ x + D . Hence g ′ (x) = C + √ x + D - x 2 √ x+D (C + √ x + D) 2 .Thus it suffices to show that h 1 (x)≜ C + √ x + D > h 2 (x) ≜ x 2 √ x+D for x ≥ 1. Note that h 1 (1) -h 2 (1) is given by h 1 (1) -h 2 (1) = p q -1 2 p q + 1 +   p q -1 2 p q   2 - 1 2 1 + p q -1 2 √ p q 2 = p q - p q 1 + p q > 0.Now we show that h ′ (x) > h ′ 2 (x) for all x ≥ 1 which implies that h 1 (x) > h 2 (x) for all x ≥ 1, thus establishing our claim. To this end, we have thath ′ 1 (x) -h ′ 2 (x) = 1 2 √ x + D - √ x + D - x 2 √ x+D 2(x + D) = x 2 √ x + D(x + D) > 0.This proves our claim that g(•) is increasing and hence (e t ) t≥0 , and consequently (θ t ) t≥0 , is bounded when θ 0 ∈ e-axis. Now let's assume that θ 0 = (e 0 , w 0 ) ∈ R 2 \ e-axis. Since (θ t ) t≥0 ⊆ R 2 \ e-axis, it follows that the loss L(•) is analytic on the trajectory (since the logistic function is analytic), and hence by [1, Theorem 2.2], it follows that lim t→∞ ∥θ t ∥ exists. Now we show that lim t→∞ ∥θ t ∥ ̸ = ∞, which implies the desired result about boundedness. To show lim t→∞ ∥θ t ∥ ̸ = ∞, we show that there exists a large B > 0 such that for any θ t = (e, w) ∈ R 2 with ∥(e, w)∥ ≥ B , the velocity vector dθ dt points inwards into the ball of radius B and thus the trajectory always stays inside this ball, and hence bounded. To establish this, let's denote (e t , w t ) = (e, w) with the dependence on time implicitly assumed. Then by the definition of GF and the gradient expressions in Lemma 7, we have thatEQUATIONEQUATIONwheref 1 = σ 2e 2 w|w| + b ⋆ + e 2 2 + q -1 -σ b ⋆ -e 2 2+ p, and f 2 = σ b ⋆ -e 2-p withπ 1 f 1 + f 2 = 0.Given that only de dt flips in sign under the transformation (e, w) → (-e, w), with out loss of generality we can assume e > 0. Now let's also assume w > 0. Thus, in view of Eq. ( 76) and GF, to show that the derivative points inwards, it suffices to show that E[(f 1 X + f 2 )X] > 0 for reasonably large B with ∥(e, w)∥ = B. Similar to Eq. (73) and Eq. (74) above, using the relationEQUATION)Using the fact that g(x) is increasing for x ≥ 1 with lim x→∞ g(x) = ∞, and |w| = w > 0, we can chose a B > 0 such that for any ∥(e, w)∥ ≥ B, in Eq. ( 78) we have 1/(1 + g(x)) < p and hence E[(f 1 X + f 2 )X] > 0. This finishes the proof of our claim. The proof for w < 0 is similar, where we make use of the fact that lim x→0 g(x) = 0 to show E[(f 1 X + f 2 )X] < 0 for e, w reasonably large.(ii) lim t→∞ θ t = θ lim . Since the logistic function ℓ log (•) is analytic, it follows from Lemma 6 that the loss L(θ) is analytic too whenever w ̸ = 0. On the other hand, when w = 0, it's easy to see that L is an analytic function of e ∈ R. By Lemma 9, we know that if w 0 ̸ = 0, w t ̸ = 0 and if w 0 = 0, w t = 0 for all t ≥ 0. Thus the loss is analytic on the trajectory for all t ≥ 0. Since the trajectory is bounded, it follows from Łojasiewicz's theorem [1, Theorem 2.2] that there exists a θ lim ∈ R 2 such that lim t→∞ θ t = θ lim .(iii) lim t→∞ ∥∇L(θ t )∥ = ∥∇L(θ lim )∥ = 0. Since the trajectory is bounded, it follows from [2, Theorem 2] that the gradient converges to zero, i.e. lim t→∞ ∥∇L(θ t )∥ = 0. Since ∇L(•) is a continuous function and lim t→∞ θ t = θ lim , we have lim t→∞ ∥∇L(θ t )∥ = ∥∇L(θ lim )∥ = 0.Proof. Since the energy function E(•, •) in Eq. ( 48) is a continuous function in R 2 \ e-axis, and any trajectory (θ t ) t≥0 with intialization θ ∈ R 2 \ e-axis stays in R 2 \ e-axis for all t ≥ 0 (Lemma 9), it follows that lim t→∞ E(θ t ) = E(θ lim ) = E(θ 0 ). As ∇L(θ lim ) = 0 from Lemma 10, it follows that θ lim lies at the intersection of the contour line E(e, w) = E 0 with the set of critical points of L in R 2 .On the other hand, if θ 0 ∈ e-axis, we have θ t ∈ e-axis from Lemma 9 for all t ≥ 0. Hence θ lim ∈ e-axis.Proof. Recall that f : R \ {0} → R, defined as f (w) ≜ E(e = 0, w) = -(w 2 + sign(w) • log |w|).If w < 0, we havef (w) = -(w 2 -log(-w)), f ′ (w) = -2w + 1 w . Hence f ′ (w) ≥ 0 for w ∈ (-∞, -1/ √ 2] and f ′ (w) ≤ 0 for w ∈ [-1/ √ 2, 0) with f ′ (-1 √ 2 ) = 0. It's also straightforward to see that lim w→-∞ f (w) = -∞, lim w→0 -f (w) = -∞, and f (-1/ √ 2) = E sad (by the definition of f ). This establishes (i), (ii), and (iii).On the other hand, for w > 0, we have f (w) = -(w 2 + log w) and f ′ (w) = -(2w + 1/w). Hence f is monotonically decreasing for w > 0 with lim w→0 + f (w) = ∞ and lim w→∞ f (w) = -∞. Note that w = 0 acts as an energy barrier since lim w→0 -f (w) = -∞ whereas lim w→0 + f (w) = ∞.Proof. First we recall the loss with the bias L(θ, b) from Eq. (41):L(θ, b) = E X,Y [ℓ log ((2Y -1) • logit X (θ, b))] ,where logit X (θ, b) = e 2 X -1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b with θ = (e, w, a). Using the fact that ℓ ′ log (z) = σ(z)-1 and 2Y -1 ∈ {±1}, we have for any θ ∈ {e, w, a, b} thatEQUATIONNow we simplify σ(logit X ) using the fact that X ∈ {0, 1}:σ(logit X ) = σ e 2 X - 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b = X σ e 2 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b ≜ϕ1 + (1 -X) σ e 2 -1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b ≜ϕ0 = Xϕ 1 + (1 -X)ϕ 0 = X(ϕ 1 -ϕ 0 ) + ϕ 0 .Thus the gradients in Eq. ( 79) are given byEQUATIONNow we compute the individual gradients with respect to e, w, a and b. Recall thatlogit X = e 2 X - 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b.Thus,∇ e logit X = 2e X - 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + e 2 2ae X - 1 2 (1 + 2w|w|) + wsign w 1 + ae 2 (2ae) = 2e X - 1 2 1 + ae 2 (1 + 2w|w|) + 2ew|w(1 + ae 2 )| + 2e 3 a X - 1 2 (1 + 2w|w|) + 2e 3 awsign w 1 + ae 2 = 2e X - 1 2 1 + ae 2 (1 + 2w|w|) + 2e 3 a X - 1 2 (1 + 2w|w|) + 2ew|w(1 + ae 2 )| + 2e 3 awsign w 1 + ae 2 .Substituting the above equation in Eq. (80), we obtain∇ e L = -E (f 1 X + f 2 ) X - 1 2 • 2e 1 + ae 2 (1 + 2w|w|) -E (f 1 X + f 2 ) X - 1 2 • 2e 3 a (1 + 2w|w|) -(E [(f 1 X + f 2 )]) • 2ew|w(1 + ae 2 )| -(E [(f 1 X + f 2 )]) • 2e 3 aw sign w 1 + ae 2 .Now we compute the derivative with respect to w.logit X = e 2 X - 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b ⇒ ∇ w logit X = 2e 2 X - 1 2 1 + ae 2 (|w| + sign (w) w) + e 2 |w(1 + ae 2 )| + w 1 + ae 2 sign w 1 + ae 2 ⇒ ∇ w L = -E (f 1 X + f 2 ) X - 1 2 2e 2 1 + ae 2 (|w| + sign (w) w) -(E [(f 1 X + f 2 )]) e 2 |w(1 + ae 2 )| + w 1 + ae 2 sign w 1 + ae 2 .Similarly, for a:logit X = e 2 X - 1 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b ⇒ ∇ a logit X = e 4 X - 1 2 1 + ae 2 (1 + 2w|w|) + e 4 w 2 sign w(1 + ae 2 ) ⇒ ∇ a L = -E (f 1 X + f 2 ) X - 1 2 e 4 (1 + 2w|w|) -(E [(f 1 X + f 2 )]) e 4 w 2 sign w(1 + ae 2 ) .Finally, since ∇ b logit X = 1, it follows from Eq. (80) thatEQUATIONFor the optimal b ⋆ , we have ∇ b L = 0 and hence E [f 1 X + f 2 ] = 0, simplifying the expressions for the gradients of e, w, and a above. In fact, there exists a closed form expression for b ⋆ in terms of e, w, and a. Recall from Eq. (80) thatEQUATIONSubstituting Eq. (82) in Eq. ( 81) and setting the gradient to zero there, we have-E X [f 1 X + f 2 ] = (σ(z 1 ) -σ(z 2 ) + p + q -1) E[X] + σ(z 2 ) -p = (σ(z 1 ) -σ(z 2 ) + p + q -1)π 1 + σ(z 2 ) -p = 0.Simplifying,(σ(z 1 ) -σ(z 2 ) + p + q -1)π 1 = p -σ(z 2 ) ⇒ (σ(z 1 ) -1) π 1 -σ(z 2 )π 1 + p = p -σ(z 2 ) ⇒ (σ(z 1 ) -1) π 1 = σ(z 2 )(π 1 -1) ⇒ σ(z 2 ) 1 -σ(z 1 ) = π 1 1 -π 1 = p q ,Using the definition of the sigmoid function and rearranging,EQUATIONBy definitions of z 1 and z 2 in Eq. ( 82), we havez 1 -z 2 = e 2 (1 + ae 2 )(1 + 2w|w|) and thus A ≜ exp(z 1 -z 2 ) = exp(e 2 (1 + ae 2 )(1 + 2w|w|)).Thus the quadratic equation in Eq. ( 83) simplifies to(exp(z 1 )) 2 + exp(z 1 )(1 - p q ) - p q • A = 0.On solving the quadratic equation for exp(z 1 ):exp(z 1 ) = 1 2   p q -1 + p q -1 2 + 4 • p q • A   ⇒ z 1 = log   1 2   p q -1 + p q -1 2 + 4 • p q • A     ⇒ b ⋆ = log   1 2   p q -1 + p q -1 2 + 4 • p q • A     -e 2 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| .Note that when a = 0 above, we recover the expression for the optimal bias in Lemma 5. This concludes the proof.N.2 Proofs of Thm. 9 and Thm. 10We prove Thm. 9 and Thm. 10 below. Note that Thm. 10 directly follows from the former by removing the bias b, since for any critical point γ = (e, w, b, a) ∈ R 4 with ∇L(γ) = 0, the bias b is already the optimal one corresponding to θ = (e, w, a) ∈ R 3 , i.e. b = b ⋆ (θ) = argmin b∈R L(θ, b). This is similar to the proof of Thm. 1, which follows from Thm. 7. Now we prove Thm. 9.Proof. We characterize the set of global minima first.Set of all global minima. Let γ ⋆ ∈ R 4 be arbitrary. From [23, Lemma 1], we have that γ ⋆ is a global minimum for the loss L(•) in Eq. ( 23) if and only if its prediction probability satisfies f γ ⋆ (x n 1 ) = P (x n+1 = 1 | x n ), the Markov kernel. Since the input {x n } N n=1 ∼ (π(p, q), P (p, q)), we have thatEQUATIONOn the other hand, by definition, from Eq.(3),f γ ⋆ (x n 1 ) = σ e 2 x n -1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b, where γ ⋆ = (e, w, b, a). Since x n ∈ {0, 1}, this can be further simplified toEQUATIONSince both f γ ⋆ (x n 1 ) and P (x n+1 = 1 | x n ) are linear functions of x n , equating them for all vallues of x n ∈ {0, 1} implies that the respective coeffecients in these functions in Eq. (84) and Eq. (85) are also equal, i.e.EQUATIONand henceEQUATIONSince σ(z) = y for y ∈ (0, 1) implies z = log y 1-y , Eq. (87) can be rewritten ase 2 1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b = log 1 -q q , e 2 -1 2 1 + ae 2 (1 + 2w|w|) + w|w(1 + ae 2 )| + b = log p 1 -p .Adding and subtracting the above two equations, we obtainEQUATIONThus γ ⋆ ∈ R 4 is a global minimum for L(•) if and only if it satisfies Eq. ( 88). It's easy to see that γ ⋆ is already a critical point for L as Eq. ( 86) is equivalent to f 1 = f 2 = 0 in Lemma 13. Thus, the set of all global minimum Γ ⋆ (p, q) is given byΓ ⋆ (p, q) ≜ {γ ⋆ = (e, w, b, a) ∈ R 4 : e 2 w|w(1 + ae 2 )| + b = 1 2 log p(1 -q) q(1 -p) , e 2 1 + ae 2 (1 + 2w|w|) = log (1 -q)(1 -p) pq }.Since the prediction f γ ⋆ (•) equals the Markov kernel for any γ ⋆ ∈ Γ ⋆ , it follows from Thm. 4 (or [23, Lemma 1] ) that L(γ ⋆ ) = H(x n+1 | x n ), the entropy rate of the Markov chain. Now we characterize the remaining set of stationary points.Non-global-min critical points. For any critical point γ = (e, w, a, b) ∈ R 4 , we have from the gradient expressions in Lemma 13 that (denoting -f 1 and -f 2 from the lemma as f 1 and f 2 ) respectively)∂L ∂b = E X [f 1 X + f 2 ] = 0, ∂L ∂e = E X (f 1 X + f 2 ) X - 1 2 2e 1 + ae 2 (1 + 2w|w|) + E X (f 1 X + f 2 ) X - 1 2 2e 3 a (1 + 2w|w|) = 0, ∂L ∂w = E X (f 1 X + f 2 ) X - 1 22e 2 1 + ae 2 (|w| + sign (w) w) = 0,EQUATIONFrom Eq. (89), we have that EX [f 1 X + f 2 ] = 0. If E X (f 1 X + f 2 ) X -1 2= E[(f 1 X + f 2 )X] = 0, we have that f 1 = f 2 = 0, implying γ is a global minimum. Hence without loss of generality, assume that E X (f 1 X + f 2 ) X -1 2 ̸ = 0. Then we can partition the above set of equations into the following regions of stationarity:(i) E X [f 1 X + f 2 ] = 0, e = 0, (ii) E X [f 1 X + f 2 ]= 0, e ̸ = 0, 1 + ae 2 = 0, 1 + 2w|w| = 0.Slightly changing the variable order, for any γ = (b, e, w, a) ∈ R 4 , we defineEQUATION)Recall that f = σ(z 1 ) -σ(z 2 ) + p + q -1, f = σ(z 2 ) -p, z ≜ e (91)Now we characterize the first set of critical points.(i) Stationary points with E X [f 1 X + f 2 ] = 0, e = 0. When e = 0, we have that z 1 = z 2 = b. Hence,f 1 = σ(b) + p + q -1 -σ(b) = p + q -1, f 2 = σ(b) -p.Thus, E X [f 1 X + f 2 ] = (p + q -1) E X [X] + σ(b) -p = (p + q -1) π 1 + σ(b) -p = 0. Rearranging and simplifying, σ(b) = p p+q and hence b = log p q . Using the fact that σ log p q = p p+q = π 1 and σ -log p q = q p+q = π 0 , the above gradients evaluated for any γ = (b = log p q , e = 0, w, a) further reduce to ∂f 1 ∂b γ = 0, ∂f 2 ∂b γ = π 0 π 1 , ∂f 1 ∂e γ = 0, ∂f 2 ∂e γ = 0,EQUATIONNow substituting Eq. ( 92) when computing the second-derivatives of L in Eq. ( 89), we obtain∂ 2 L ∂b 2 γ = E X ∂f 1 ∂b γ X + ∂f 2 ∂b γ = π 0 π 1 , ∂ 2 L ∂b∂e γ = E X ∂f 1 ∂e γ X + ∂f 2 ∂e γ = 0, ∂ 2 L ∂b∂w γ = E X ∂f 1 ∂w γ X + ∂f 2 ∂w γ = 0, ∂ 2 L ∂b∂a γ = E X ∂f 1 ∂a γ X + ∂f 2 ∂a γ = 0, ∂ 2 L ∂e 2 γ = E X (f 1 X + f 2 ) X - 1 2 2 1 + ae 2 (1 + 2w|w|) γ + E X (f 1 X + f 2 ) X - 1 2 4e 2 a (1 + 2w|w|) γ 0 + E X ∂f 1 ∂e γ X + ∂f 2 ∂e γ X - 1 2 2 1 + ae 2 (1 + 2w|w|) 0 = E X [(2f 1 (1 + 2w|w|) -f 1 + 2f 2 (1 + 2w|w|)X -f 2 ] γ = E X [(f 1 (1 + 4w|w|) + f 2 (2 + 4w|w|)X -f 2 ] γ = (f 1 (1 + 4w|w|) + f 2 (2 + 4w|w|))π 1 -f 2 γ (a)= ((p + q -1)(1 + 4w|w|) -π 1 (p + q -1)(2 + 4w|w|))π 1 + π 1 (p + q -1) = π 1 (p + q -1) (1 + 4w|w| -π 1 (2 + 4w|w|) + 1)(b)= 2π 1 π 0 (p + q -1)(1 + 2w|w|),where (a) follows from the fact that f 1 | γ = p+q-1, f 2 | γ = σ(b)-p = p p+q -p = -p p+q (p+q-1) = -π 1 (p + q -1) and (b) from 1 -π 1 = π 0 . Returning to the remaining second derivatives,EQUATIONCongregating all the second derivatives from Eq. (93) and Eq. (94) into the Hessian H(γ) in Eq. (90), we finally obtainH(γ) = π 0 π 1    1 0 0 0 0 2(p + q -1)(1 + 2w|w|) 0 0 0 0 0 0 0 0 0 0    ,which is identical to the Hessian obtained in the proof of Thm. 7 (App. J.2) for e = 0. Thus it follows that Γ min (p, q) ⊆ R 4 and Γ sad ⊆ R 4 defined below are a set of local minima and saddle points respectively:Γ min (p, q) ≜ γ min = (e, w, b, a) ∈ R 4 : e = 0, (p + q -1)(1 + 2w|w|) > 0, b = log p q , Γ sad (p, q) ≜ γ sad = (e, w, b, a) ∈ R 4 : e = 0, (p + q -1)(1 + 2w|w|) ≤ 0, b = log p q .Now we focus on the remaining set of critical points.(ii) Stationary points with E X [f 1 X + f 2 ] = 0, e ̸ = 0, 1 + ae 2 = 0, 1 + 2w|w| = 0. For this set of points, the Hessian remains undefined because ∂f1 ∂e , ∂f2 ∂e , ∂f1 ∂a , ∂f2 ∂a do not exist (Eq. ( 91)). This nonexistence arises since sign 1 + ae 2 is undefine 1 + ae 2 = 0. However, even in this scenario,when e ̸ = 0, 1 + ae 2 = 0, 1 + 2w|w| = 0, we have z 1 = z 2 = b. Hence,f 1 = σ(b) + p + q -1 -σ(b) = p + q -1, f 2 = σ(b) -p.Thus the expectation term E X [f 1 X + f 2 ] = (p + q -1) E X [X] + σ(b) -p = (p + q -1) π 1 + σ(b) -p = 0. Simplifying, σ(b) = p p+q , which implies b = log p q . We could attempt to understand the characterization of the points on this manifold through local perturbation analysis. However, in this work, we classify them as stationary points and leave the comprehensive characterization for future research. This set of points Γ station (p, q) ⊆ R 4 is defined as Γ station (p, q) ≜ γ min = (e, w, b, a) ∈ R 4 : e ̸ = 0, 1 + ae 2 = 0, 1 + 2w|w| = 0, b = log p q .This concludes the proof.Proof. Recall from Lemma 13 that for θ = (e, w, a) ∈ R 3 , we have∂L ∂e = -E (f 1 X + f 2 ) X - 1 2 • 2e 1 + ae 2 (1 + 2w|w|) -E (f 1 X + f 2 ) X - 1 2 • 2e 3 a (1 + 2w|w|) , ∂L ∂w = -E (f 1 X + f 2 ) X - 1 2 • 2e 2 1 + ae 2 (|w| + sign (w) w) , ∂L ∂a = -E (f 1 X + f 2 ) X - 1 2 • e 4 (1 + 2w|w|) ,where X ∈ {0, 1} is a Bernoulli random variable with X ∼ Bern(p/(p + q)), and (95)f 1 ≜ 1 -p -q -ϕ 1 + ϕ 0 , f 2 ≜ p -ϕ 0 , ϕ 1 ≜ σ eSimilarly for a,EQUATIONLikewise, for e,EQUATION(1 + 2w|w|) 2e 3 a.(97)By substituting the expressions of Eq. ( 95) and Eq. (96) into Eq. ( 97): Note that here c ∈ R, is a constant that depends on the initial conditions. Thus the energy E(θ t ) = E(θ 0 ) for w 0 ̸ = 0. Lemma 16, Lemma 15, and Thm. 3 Proof. We note that the proofs of Lemma 16, Lemma 15 directly follow from that of their counterparts Lemma 11 and Lemma 10 using Łojasiewicz's theorem to characterize the convergence of the gradient flow. Thm. 3 is a direct consequence of Lemma 16, Lemma 15.EQUATIONN.5 Informal proof of Thm. 11Proof.[Informal] For θ = (e, w, a), recall from Lemma 13 that the derivative of the loss L with respect to e is∂L ∂e = E X (f 1 X + f 2 ) X - 1 2 2e 1 + ae 2 (1 + 2w|w|) + E X (f 1 X + f 2 ) X - 1 2 2e 3 a (1 + 2w|w|) ,where Hence the gradient with respect to e isf 1 = σ(z 1 ) -σ(z 2 ) + p + q -1, f 2 = σ(z 2 ) -p,EQUATION)Simplifying the expectation term, E X [(f1 X + f 2 ) X] = (f 1 + f 2 )π 1 = f 1 π 1 -f 1 π 2 = (p + q - 1)(π 1 -π 2 1), where we used the fact that b is optimal in the above equations, specifically where f 1 π 1 + f 2 = 0. Thus the gradient flow for the parameter e is governed by ė = -∂L ∂e = -(p + q -1)(π 1 -π 2 1 )e ⇒ e = e 0 exp(-(p + q -1)(π 1 -π 2 1 )t).Since (p + q -1)(π 1 -π 2 1 ) > 0, e → 0, which denotes it converges to the local minima. 


Decision: This submission presents a theoretical analysis of the learning dynamics of a simple transformer architecture on a simple Markov chain data model. Reviewers praised the technical correctness of the work and the clarity. Most concerns from reviewers are about the relevance of assumptions for more practical settings. However, these assumptions have been clearly explained in the submission and responses, and theoretical work in this area is difficult and needed, so I recommend acceptance.

Reviews: review: paper_summary: paper_summary: The paper theoretically studies the effect of initialization on the gradient dynamics of a single-layer transformer trained on Markov data. It considers a simplified model by: (1) using a binary input alphabet and (2) reducing the many parameters of the single-layer transformer to just two or three scalar parameters. The authors identify the critical points of this reduced set of parameters and compare them to the stable point under its gradient flow dynamics. This is used to determine whether an initialization will converge to global or local optima. This informs choices of initialization that will reliably converge to global optima. This is validated in a single-layer transformer where the reduction of weights is not performed. This experiment is also used to validate the reductions made for theoretical analysis.
strength_weakness: strength_weakness: It is not clear how much results vary across repeated trials. The authors state that they repeated experiments 3 to 5 times, but the results across replications are not presented in the figures or the text. For instance, in Figure 6, did every repeated experiment follow the same blue and red curves? Similarly, while the authors show the weights from a single run on Figure 5, they do not provide a numerical assessment of the validity of their low-rank approximations. Another weakness is the simplified data model, which is just a binary alphabet. The authors seem to inherit this from previous analytical work, but a discussion of how this restricts results and real-world applicability is missing.
questions: questions: 1. How much do results vary across replications? The authors could easily answer this by showing different runs in the same figure or stating measures of uncertainty. If possible, the authors could also consider more values of p and q.
2. How reliable is the low-rank approximation? It would be beneficial to quantify the low-rank structure of W1 and Wv so that it is easy to understand how consistent this approximation is in general.
3. Do similar empirical results hold with larger alphabets? More specifically, do similar low-rank approximations still hold, and do these initializations still work better?
4. The authors remark that they observe $a \approx 0$, so the attention term is often not relevant. Is this a consequence of the simplified data model? It would be nice to comment more on this observation.
5. Figures 2 and 3 should be subfigures.
limitations: limitations: The authors are fair in the scope of their claims, though it would be good to emphasize the use of reduction/reparameterization earlier in the paper or abstract (this is not made clear until the end of page 3, and this reduces the theoretical scope of their claims). Alternatively, the authors could provide stronger evidence to justify the universality of these approximations (precise quantification of low-rank structure, and more replications).
paper_summary: paper_summary: This paper investigates the learning dynamics of transformer models, specifically focusing on first-order Markov chains and single-layer transformers. The authors aim to understand how transformers learn Markov chains and the impact of initialization on their training outcomes. They provide a comprehensive theoretical analysis, demonstrating that transformer parameters trained on next-token prediction loss can converge to either global or local minima, depending on the initialization and properties of the Markovian data. The paper also offers empirical evidence to support their theoretical findings and proposes guidelines for initializing transformer parameters effectively.
strength_weakness: strength_weakness: - One important weakness of this study is its limited scope. The study focuses on single-layer transformers and first-order Markov chains, which may limit the generalizability of the findings to more complex models and data structures.
- While thorough, the theoretical analysis may be difficult for practitioners without a strong mathematical background to fully grasp and apply. Maybe the presentation can be changed and made more intuitive to also target the audience who does not have a strong mathematical background.
- The empirical validation could benefit from a broader range of experiments, including different types of datasets and more complex transformer architectures.
questions: questions: The paper does not address the scalability of the proposed initialization guidelines for large-scale transformer models, which are commonly used in practice. Can the authors provide an insight on this?
limitations: limitations: As I mentioned in weaknesses section, the most important limitation of this study is its generalizability. The study focuses on single-layer transformers and first-order Markov chains, and it is unknown if this method is generalizable to more complex models and data structures.
paper_summary: paper_summary: The paper investigates how transformers learn first-order Markov chains, focusing on the role of parameter initialization and providing a comprehensive analysis of the learning dynamics. It proves that transformer parameters can converge to global or local minima based on initialization and data properties, and supports these theoretical findings with empirical evidence.
strength_weakness: strength_weakness: 1. The paper examines first-order Markov chains with a rank-1 input sequence and considers single-layer linear attention. This general setting might be overly simplified.
2. Can the initialization guidelines provided be extended to more complex transformer models, such as those with additional layers?
3. While the empirical results support the findings, it is not intuitively obvious why the optimal parameters exhibit low rank. Could you offer a more detailed explanation?
4. A minor suggestion: the formulation $att_{n,i}= . . .$ should be included in the main body of the paper rather than just in the appendix, as attention is a crucial component of the transformer structure.
questions: questions: Please see the weakness.
limitations: limitations: Yes.
paper_summary: paper_summary: The paper seeks to characterize how single layer transformers learn (two symbol) markov chains. The analysis relies on a reparametrization of the transformer that assumes the weight matrices are rank 1. The primary results are characterizing when gradient flow on this reparameterized transformer leads to global minimums in loss or other critical points, showing how in some cases a standard gaussian initialization will lead to a local minimum instead of the global minimum.
strength_weakness: strength_weakness: 1. The canonical parameterization is a rather strong assumption (reducing the model to just three parameters) and so it needs strong results to justify it. For me, the results in this paper are close to strong enough, but some strengthening of them would be appreciated (short of removing assumptions).

It would be nice more evidence showing how the analysis on the canonical parameterization could extend to more general settings. Ideally, theoretical analyses of toy systems should be paired with empirical and heuristic evidence that they either a) approximately explain the dynamics of more realistic systems, or b) yield new insights that can be empirically verified on more realistic systems.

I believe that the paper is primarily trying to show the first claim (a), with the experiments showing the transformers initialized with rank 1 weights stay rank 1, but I think this can benefit from some additional evidence. For example, it would be great to show that the learning dynamics of the canonical parameterization are similar to that of a single layer transformer (with softmax attention and no constant parameters) with a rank one initialization.

2. Figure 6 might be trying to address the second claim (b), but I think it can be improved. Having W_1 and W_2 be constant implies that “our initialization” is more of a reparametrization than an initialization, making the comparison not feel apples to apples.

3. For figures 5 and 7, while I agree that qualitatively the matrices look rank 1, some quantitative measure of rank (like nuclear norm) would be good to graph (especially with error bars for multiple runs).

4. While it is mentioned that experiments were repeated 5 times, it might be nice to put into the paper some of these results, Figure 6 stood out as a place where it would be good and not too difficult to add error bars or additional runs. Additionally, details of how replications were done would be good to mention (such as how random seeds were chosen). The code as of now is a bit messy, and could do with some reorganizing and removing unnecessary files (like untitled.ipynb).
questions: questions: 1. How is figure 6 testing the hypothesis that “for p + q > 1, any small initialization around zero would lead to a local minima” (line 247). I feel like it's more of comparing a small initialization to your parametrization (which is still a worthy endeavor). I’d imagine that testing the stated hypothesis would require trying many different small initializations and showing that they lead to local minima (Given the nature of the problem, I imagine it could be feasible to test initializations in a grid and create something like an empirical analogue to corroborate figure 1d?). Also, a natural follow-up question (or null hypothesis) is can some standard initialization that isn’t small and/or isn’t around zero work just as well as your initialization.

2. Also, it took me a bit to understand the definition of the attention matrix (attn) (along with Wq and Wk, they seem to only only be mentioned in section 4.1 in the main body, though I may have missed something). Are they abstracted away from earlier in the paper for simplicities sake, or some other reason?
limitations: limitations: The discussion of limitations in the conclusion section (which I mentioned specifically because of the response given in the checklist, lines 1200-1201) is a bit too brief. It would be good to have all of the assumptions made for the main theorems in one place, especially those implicit and not stated in the theorem statements, including the rank one and related assumptions, linear attention, and gradient flow (which is in the theorem statements, but good to mention because transformers are not normally trained with gradient flow).
paper_summary: paper_summary: This paper focuses on training a single-layer linear attention transformer with low-rank parameterization on first-order Markov chain data. By reparameterization, they can reduce this problem to a 3-variable learning dynamics and comprehensively characterize the trajectory and local/global minimizers. This paper precisely characterizes the different convergence conditions for local/global convergence, highlighting the role of initialization. Empirically, on this specific data distribution, they verified the superiority of their special initialization compared to standard initialization.
strength_weakness: strength_weakness: 1. I am concerned with the generality of the insights on the role initialization. I appreciate that the authors make a great effort to theoretically analyze the Markov Chain with 2 states (i.e. $P\in R^{2\times 2}$) and do extensive experiments on this particular dataset with linear attention. However, the authors claim that ``we offer practical guidelines for initialization of transformer parameters and demonstrate their effectiveness."(Line 280-281) in both the abstract and the conclusion. Line 255 also claims "the corresponding insights are more general and apply more broadly to the general architecture.''. Since all the conclusions are based on this simplified setting, I don't think it can be characterized as a realistic guide for initialization. I think further experiments on the Markov data with more states and non-linear attention transformers can benefit this paper and make the results more general. I will increase the score if the results can be verified in a more general setting.

2. As a minor point, the dynamics analysis is restricted to a low-rank manifold. Though it is partially corroborated by the experiments that GD finally converges to the low-rank solution, it cannot guarantee that without this restriction, the parameter can stay on the manifold once initialized on the manifold. Nevertheless, the analysis is already complicated and this simplification is acceptable.
questions: questions: 1. Does the low-rank initialization scheme insight transfer to more general data distribution or architecture?
limitations: limitations: The dynamics analysis is restricted to certain subspaces. The attention pattern does not train (which is simplified as a trainable scaling factor $a$). The experiments cannot corroborate the general claim.

Comments: title: Summary of review and discussion
Thank you for the rebuttal.
Generality of the insights
Attention features are also learnt. Insights for linear-attention generalize to soft-max. Low-rank vectors turn out not be relevant.
Discussion period ending soon
Discussion period ending soon: call for action
Apologies
Overlong PDF
Limitation section and new figures
Question
comment: Sorry for the delay, I have updated my score.
We thank all the reviewers for their constructive feedback, engaging discussion during the rebuttal, and valuable time. We are pleased that all reviewers found our paper to be well-written, insightful, clear, and precise, with solid theoretical contributions supported by robust experimental results and visualizations. We are also delighted that they found our initialization analysis to be insightful and impactful (_reviewers DieV and er5c_)

The following list summarizes the additions that we will make to the final vision of our manuscript. These additions are based on the comments and questions of the reviewers during the official review, and are already made available in our rebuttal.

---------------------------------------------------------------------------------------------------------------------------------
1. **Assumptions and generality of the insights:**  We believe there maybe a misunderstanding about the generality of the insights obtained through our analysis in the current paper. In the modified version, we will include a separate section highlighting the specific setting used for our theoretical analysis (_canonical_) and the architecture used for the experiments (_full general model_). We will also include additional experiments for all $(p,q)$ highlighting the generality of our results. We thank all the reviewers for this suggestion.

2. **Error-bars and low-rank metric:** We will modify all the existing curves averaged over runs with the shaded error bars, as shared in the global response below. We will also quantify and better illustrate the low-rankness of matrices using the low-rank metric we discussed in the rebuttal. We thank reviewers LVoR and Fte8 for these suggestions.

3. **Multi-state Markov chains:** We will add an additional section about multi-state Markov chains discussing how one could obtain similar initialization guidelines in this scenario. In addition to the empirical results highlighting the role of Markovian switching and the number of states  on the loss landscape, as inlcuded in the rebuttal PDF, this entails a clear theoretical analysis of the loss landscape first and the associated learning dynamics. We thank reviewers LVoR and er5c for the suggestion.

4. **Limitations and scalability:** Being one of the first works to precisely characterzie the role of intialization in learning dynamics of transformers, our work opens new doors for similar analyses for deeper architectures and multi-state settings. We will add a separate section outlining the current limitations of our approach and outline some important research directions to extend to more complicated settings. We thank reviewers LVoR, DieV, and JG7D for this suggestion.
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.
Dear Reviewer er5c,

Thanks for your prompt response. We address the individual comments below.

**Attention features:** As you already noted, indeed the attention scalar $a$ captures the essence of both the attention mechanism and the value matrix and hence can be thought of as a trainable parameter that takes both of them into account, as also implied by its definition in Line 200. Sorry for the confusion through the typo in Line 112. Here we just wanted to motivate the definition of $a$ without getting into too much notation-heavy details. We will correct this.

**Linear vs. soft-max attention:** Sorry if we misunderstood your claim. Our earlier disagreement was with the initialization scheme and not regarding the multi-state Markov chains. As the Figure 6 in the rebuttal suggests, while it's true that the matrices are not exactly rank-1, nonetheless they still converge to a relativley low-rank corresponding to rank-4. And as also highlighted in Figure 5 of the rebuttal, the input Markovian switching $p$ has an effect on model converging to local or global minima, albeit with a threshold different than $0.5$, which is the threshold for the binary case. This is not surpising given that this corresponds to multi-state setting. Together these observations suggest that while similar phenomena as in binary-setting happen also in multi-state setting, this warrants a more detailed study and corresponding gradient flow analysis. This is considerably more challenging than the binary case, but nonetheless an important and exciting venue of future research and currently outside the scope of our current paper. We will include a discussion and corresponding results for multi-state case in the revised version to provide more context to our work.

**Low-rank parameters:** We agree with your assessment. We initialized all the low-rank vectors to be a single $\alpha$ because that's what happens in practice at convergence (one could also experimentally check how fast this happens from the beginning). But we agree that analyzing the full dynamics with different vectors is indeed an interesting direction of future research though considerably more challenging, as already revealed in our simplified setting.
Thank you for the reply. Some of my concerns have been addressed. Hence, I will keep my original score.
**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 $\texttt{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 $\texttt{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.
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 $\texttt{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 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 $\texttt{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.
Dear Authors, I carefully read your responses and the other reviewers' reviews. I agree to the reviewer er5c on their comments on the generality of the insights. Additionally, my concerns regarding generalizability and the scalability of the proposed method sustain. That's why I will maintain my score.
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
Dear Reviewer DieV,

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.

Sincerely, 
The Authors
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
Thanks for the detailed response. After careful consideration, I will maintain my score. Below, I provide specific reasons for each of the points.

**Generality of the insights**: Regarding the softmax attention for $P\in \mathbb{R}^{2\times 2}$ mentioned in the paper, I believe that the initialization scheme in linear attention prevents the softmax transformers from converging to local minima. However, the exact architecture used in the experiments should be explicitly stated in the main paper. For the multi-state Markov Chain setting, since I was requested to look only at the first page, the current content does not sufficiently demonstrate general results that would illustrate the low-rank structure of the final minimizer or the effectiveness of your low-rank initialization. For completeness, I suggest including experiments that explore whether your low-rank initialization helps transformers not converge to local minima in multi-state Markov Chain tasks.

**Low-rank manifold**:  I fully understand the difficulty of analyzing transformer dynamics, and I agree that reparameterizing the transformer to simplify the analysis is a reasonable approach. However, I must note that restricting the analysis to the ground-truth low-rank manifold with only a trainable scalar factor $a$ may be an oversimplification. Compared with the related works, [1] and other linear regression ICL papers only limit the trainable parameter by zeroing out some irrelevant parts of the matrices. They can then prove the gradient can help learn the correct ground truth, instead of limiting the feature on the ground-truth direction. I believe that if the low-rank parameters in the attention layer were trained (initialized or parameterized in a low-rank manner, but not necessarily aligned with $\alpha$), it would significantly enhance the contribution of this paper. It may be too challenging to analyze with so many trainable parameters, so the authors are free to consider or disregard this suggestion.

**Limitations**: I have revised my previous inaccurate comment. My primary concern is that the attention features of the transformers are not learned, with the only trainable parameter being the scalar $a$. Specifically, the attention score $\texttt{attn}_{n, i}$ is not trained.
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 reviewer,

- **Canonical vs. low-rank**: Thanks for your prompt response and thought provoking questions. In our earlier response, due to time constraints , we couldn't attach the figure showcasing the similar training profiles of canonical and low-rank model but we now have the corresponding empirical results and the figure that illustrate this phenomenon. We are currently checking with the ACs about how best to share these results as the official rules forbid us from sharing links. We'll keep you updated.

- **Multi-state**: In this setting as well, there would exist analogous good and bad initialization regions as highlighted for the binary case in Figs. 1 and 4 in our paper. While empirically we can verify if our scheme still works here as is, however we have to rederive and precisely determine the corresponding basins of attraction for the local and global minima in the multi-state setting separately. Also the Markovian switching condition $p+q>1$ would also have to correspondigly modified to accommodate more states. We already see a glimpse of this in Figure 5 of the rebuttal for five-state Markov chain: $p=0.6$ still converges to global minima unlike the binary-state setting where $p+q = 0.6 + 0.6 >1$ would have led to local minima convergence.

- **Limitations and assumptions**: Indeed, in our revised version we will add two separate subsections. The first one outliining all the corresponding assumptions and the setting for our results. Then the second discusses the limitations of our results and potential future directions to address them.

Once again, we would like to thank the reviewer for their constructive efforts to improve the paper which we really appreciate. We are wondering if they are willing to change their score if their concerns are addressed. We remain at your disposal for any further questions.
I was wondering (after seeing figure 6 in the rebuttal pdf), if theres any chance your initialization scheme help in the multi state markov chain setting (I know that this might be beyond the scope of this work, but I am curious)?

I also wanted to ask if you plan for the discussion of limitations in the conclusion to be any different in the camera ready version than the current version?