Title: Local to Global: Learning Dynamics and Effect of Initialization for Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: gradient descent dynamics

transformers

markov chains
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}.

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

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 nexttoken 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: https://github.com/Bond1995/Markov .

1 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).

3

2

1

0

1

2

3

2

1

0

1

2

Global minima

Local minima

Local maxima

Saddle point

9

7

5

3

1

1

3

5

7

9

11

w

→

e

↑

(a) Gradient flow

(

p

+

q <

1)

3

2

1

0

1

2

3

3

2

1

0

1

2

3

min

w

→

e

↑

(b) Convergence basin

(

p

+

q<

1)

3

2

1

0

1

2

3

2

1

0

1

2

Global minima

Local minima

Local maxima

Saddle point

9

7

5

3

1

1

3

5

7

9

11

w

→

e

↑

(c) Gradient flow

(

p

+

q >

1)

3

2

1

0

1

2

3

3

2

1

0

1

2

3

min

w

→

e

↑

(d) Convergence basin

(

p

+

q>

1)

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

2 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 probability

P 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 ) :

Here θ ≜ ( 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.

When 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].

Loss landscape. 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:

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 have

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:

1 is the all-one vector in R r with r = 4 d 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:

Finally, 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):

Due 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 ⋆ :

Empirically, 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 2 x n +1 -1 ∈ {± 1 } and the logits given by e 2 (1 + 2 w | 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.

3.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 by

(ii) the set of all local minima is given by

(iii) the set of all local maxima is given by

(iv) and the set of all saddle points is

Thus the set of all critical points is

In addition, for any θ ⋆ ∈ Θ ⋆ , θ min ∈ Θ min , θ max ∈ Θ max , and θ sad ∈ Θ sad , the loss values satisfy

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.

4 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 by

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 as

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

On 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 value E 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

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

4.1 Gradient Flow with Attention

(a)

p

+

q <

1

: Gradient flow

(b) Energy manifold and minima

(c)

p

+

q >

1

: Gradient flow

(d) Energy manifold and minima

Global minima Local minima Ball around origin

Global minima Local minima  Constant energy manifolds

Global minima Local minima Ball around origin

Global minima Local minima  Constant energy manifolds

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 be

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 by

starting with a randomly initalized θ 0 . We define the energy function E ( · , · , · ) as

where ea-plane ≜ { ( e, w = 0 , a ) } . It is similar to its counterpart in Eq. (11), except for the 2 a 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 .

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.

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

5.1 Low-rank Parameters

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 4 d × 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.

5.2 Effect of Initialization: Broader Implications

Now we investigate the findings of Sec. 3 and Sec. 4, derived for the canonical low-rank model, more broadly in the context of full model in Sec. 2. In particular, as shown in Thm. 2 and Fig. 1d

0

20

40

60

80

100

Iteration (multiples of 100)

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Fraction of energy

Random initialization

Rank-one initialization

(a)

W

V

0

10

20

30

40

50

60

70

80

Iteration (multiples of 100)

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Fraction of energy

Random initialization

Rank-one initialization

(b)

W

1

0

100

200

300

400

500

600

Iteration

0.60

0.65

0.70

0.75

0.80

0.85

0.90

Test loss

Standard initialization

Our initialization

Unigram loss

Bigram loss

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.

5.3 Higher-Order and Multi-State Markov Chains

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.

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

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

Acknowledgments and Disclosure of Funding

Ashok would like to thank Aditya Vardhan Varre for many helpful discussions about the project. This work was supported in part by the Swiss National Science Foundation under Grant 200364.

References

Pierre-Antoine Absil, Robert Mahony, and Ben Andrews. Convergence of the iterates of descent methods for analytic cost functions. SIAM Journal on Optimization , 16(2):531-547, 2005.

Arzu Ahmadova. Convergence results for gradient flow and gradient descent systems in the artificial neural network training. arXiv preprint arXiv:2306.13086 , 2023.

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/forum? id=LziniAXEI9 .

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? Investigations with linear models. In The Eleventh International Conference on Learning Representations , 2023.

Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks, 2019.

Francis Bach. Effortless optimization through gradient flows. https://francisbach.com/gradientflows/ , 2020.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Workshop on Efficient Systems for Foundation Models @ ICML2023 , 2023.

Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Pritam Chandra, Tanmay Kumar Sinha, Kabir Ahuja, Ankit Garg, and Navin Goyal. Towards analyzing self-attention via linear neural network, 2024. URL https://openreview.net/ forum?id=4fVuBf5HE9 .

Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality, 2024.

Thomas M Cover and Joy A Thomas. Elements of information theory . John Wiley & Sons, 2nd edition, 2006.

John M Danskin. The theory of max-min, with applications. SIAM Journal on Applied Mathematics , 14(4):641-664, 1966.

Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, Lei Li, and Zhifang Sui. A survey on in-context learning, 2023. URL https://arxiv. org/abs/2301.00234 .

Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains, 2024.

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread , 2021. URL https://transformer-circuits.pub/2021/framework/index.html .

Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? A study through the random features lens. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

MEmrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self-attention to markov models: Unveiling the dynamics of generative transformers. arXiv preprint arXiv:2402.13512 , 2024.

Samy Jelassi, Michael Eli Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Jason D Lee, Max Simchowitz, Michael I Jordan, and Benjamin Recht. Gradient descent only converges to minimizers. In Conference on learning theory , pages 1246-1257. PMLR, 2016.

Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: generalization and stability in in-context learning. In Proceedings of the 40th International Conference on Machine Learning , 2023.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding, 2023.

Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with Markov: A framework for principled analysis of transformers via Markov chains. arXiv preprint arXiv:2402.04161 , 2024.

Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 , 2024.

Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In Proceedings of the 40th International Conference on Machine Learning , 2023.

Matteo Pagliardini. GPT-2 modular codebase implementation. https://github.com/epfml/llmbaselines , 2023.

Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is Turing-complete. Journal of Machine Learning Research , 22(75):1-35, 2021.

Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vardhan Makkuva. Transformers on markov data: Constant depth suffices, 2024. URL https: //arxiv.org/abs/2407.17686 .

Alexander Shapiro. Second-order derivatives of extremal-value functions and optimality conditions for semi-infinite programs. Mathematics of Operations Research , 10(2):207-219, 1985.

Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. CoRR , abs/2103.07601, 2021. URL https://arxiv.org/abs/2103.07601 .

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research , 19 (70):1-57, 2018.

Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Shaolei Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Conference on Parsimony and Learning (Recent Spotlight Track) , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , pages 5998-6008, 2017.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174, 2023.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2023.

Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating Turing machines with transformers. In Advances in Neural Information Processing Systems , volume 35, pages 12071-12083, 2022.

Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning , pages 11080-11090, 2021.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of incontext learning as implicit Bayesian inference. arXiv preprint arXiv:2111.02080 , 2021. URL https://arxiv.org/abs/2111.02080 .

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations , 2020.

Contents

A Single-layer transformer: architecture and results

We first describe the transformer architecture from Sec. 2:

Here θ ≜ ( 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 vectors k 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 × d are 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:

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

A.1 Loss landscape results

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.

B Low-rank structure of the optima

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

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.

for p > 0 and p ≈ e 2 such that e > p . Thus from Embedding layer,

which 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 by y n = x n W O ∑ i ∈ [ n ] att n,i · W V x i , where the attention weights att n,i are computed according to att 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.

Hence, and

Thus,

It is further noticed that ⟨ v , α ⟩ ≈ 0 and hence y 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,

Thus 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( W 1 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 ✶ ( w i = 1) , we further simplify: W 2 ReLU( W 1 y n ) = W ⊤ 1 ReLU( W 1 y n )

Hence

Linear. Using the fact that e = a = e · α due to weight-tying, we obtain from Linear layer that

Prediction. We finally obtain that the prediction probability

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

C Canonical reparameterization

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 to

Attention. Substituting this x n in Eq. (16), we have

where the attention weights att n,i = exp ( e 2 q 2 d 5 / 2 ( x n -1 2 )( x i -1 2 ) ) ∑ j ∈ [ n ] exp ( e 2 q 2 d 5 / 2 ( x n -1 2 )( x j -1 2 ) ) ∈ (0 , 1) for some q ∈ R . Since ⟨ v , α ⟩ ≈ 0 , we let v = 0 and obtain

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 = 4 d . 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:

Here w > 0 ensures W 2 = W ⊤ 1 whereas w < 0 , W 2 = -W ⊤ 1 . Using this parameterization, substituting β = r = 4 d and w ↦→ w d in Eq. (19), we get

Linear. Since e = a = e · α due to weight-tying, Eq. (20) simplifies to

Prediction. The next-token prediction probability is

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:

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 points

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.

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 as

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.

D.1 Technical lemmas

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 as

where ( 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

where X ∈ { 0 , 1 } is a Bernoulli random variable with X ∼ Bern( p/ ( p + q )) , f 1 = σ ( 2 e 2 w | w | + b + e 2 2 ) + q -1 -σ ( b -e 2 2 ) + p , and f 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 [( f 1 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

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

The 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

We start with the result that establishes a closed form expression for b ⋆ .

Lemma 5 (Optimal bias) . For θ = ( e, w ) ∈ R 2 and b ∈ R , let L ( θ , b ) be the next-token prediction loss defined in Eq. (23) . Then, for any θ ∈ R 2 , L ( θ , b ) is convex in b and the minimizer b ⋆ ≜ argmin b ∈ R L ( θ , b ) is given by

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 as

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

where X ∈ { 0 , 1 } is a Bernoulli random variable with X ∼ Bern( p/ ( p + q )) , f 1 = σ ( 2 e 2 w | w | + b ⋆ + e 2 2 ) + q -1 -σ ( b ⋆ -e 2 2 ) + p , and f 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 γ be

Then the Hessian of L at θ ∈ R 2 is given by

Consequently, for any γ = ( b, e, w ) ∈ Γ min ∪ Γ sad , the Hessian H ( θ ) at θ = ( e, w ) is given by

where π 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. 1

Proof. 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:

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.

(i) Set of all critical points: Recall from Thm. 7 that for any critical point γ = ( θ , b ) = ( e, w, b ) ∈ R 3 of L , γ ∈ Γ ⋆ ∪ Γ min ∪ Γ sad . Here the main observation is that all these critical points are of the form ( θ , b ⋆ ( θ )) where θ ∈ Θ ⋆ ∪ Θ min ∪ Θ max ∪ Θ sad . To see this, let γ ∈ Γ ⋆ . Here we have e 2 (1 + 2 w | w | ) = log (1 -p )(1 -q ) 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 + 2 w | w | ) ≤ 0 . Hence θ = ( e, w ) ∈ Θ max ∪ Θ sad . Suppose θ ∈ Θ max which implies e = 0 , ( p + q -1)(1 + 2 w | w | ) < 0 . By Lemma 8, the Hessian at θ (upto a positive scale) is a diagonal matrix with the entries ( p + q -1)(1 + 2 w | 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.

F Gradient flow analysis without attention

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 by

starting with a randomly initalized θ 0 . The energy function E ( · , · ) is defined as

where e-axis ≜ { ( e, w = 0) } and w-axis ≜ { ( e = 0 , w ) } . Note that E 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.

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.

Lemma11 (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. 2

Proof. Let θ 0 = ( e 0 , w 0 ) ∈ R 2 be the initialization for the GF trajectory ( θ t ) t ≥ 0 . Recall that

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 + 2 w | 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 .

(i) θ 0 ∈ R 2 \ e-axis : Let E 0 = E ( θ 0 ) ∈ R . By Lemmas. (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.

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 -( w 2 0 -log( -w 0 )) = E sad for w 0 < -1 √ 2 that passes through ( -1 √ 2 , 0) ∈ R 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 d e t d t > 0 and d w t d t < 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 that

So it suffices to show that E X [( f 1 X + f 2 ) X ] > 0 . To establish this, we have from Lemma 5 that

From the defintion of f 2 and the optimal bias b ⋆ in Lemma 7 and Lemma 5 respectively, we obtain

When e 0 = 0 , we have A = 1 and hence

where 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 + 2 w | w | ) < 0 } = { ( e, w ) ∈ R 2 : e = 0 , w < -1 √ 2 } . Thus for any θ 0 ∈

Θ max , d θ t d dt = 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 is I ⋆ = 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 √ 2 when 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 -2 w 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

(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 ).

G Gradient flow analysis with attention

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.

G.1 Canonical parameterization with attention

Embedding. Recall from App. C that we let e = e · α and p n = -p · α for all n where e > 0 , p = e 2 and α ∈ {± 1 } d / √ d . This results in the embedding

Attention. Similarly, we recall from Eq. (21) that the attention output y n is given by

where

Instead of the softmax, now we assume that the attention weights are linear in the scaled dot product, i.e.

Note that the 1 /n factor is to ensure normalization for the attention weights in Eq. (37). Now substituting Eq. (38) in Eq. (37), we obtain

where we used the fact that ( x i -1 2 ) 2 = 1 4 since x i ∈ { 0 , 1 } , and

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 have

Feed-forward . For the feed-forward layer, we have that W 1 = | w | √ d 1 · α ⊤ ∈ R 4 d × d , W 2 = w √ d α · 1 ⊤ ∈ R d × 4 d . Hence Eq. (40) implies

Thus,

Using ReLU( x ) -ReLU( -x ) = x above,

Hence,

Thus the embedding z n is given by

Linear. Since e = a = e · α due to weight-tying, the logits are given by

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 as

where logit X ( θ , b ) ≜ e 2 [( X -1 2 ) ( 1 + ae 2 ) (1 + 2 w | 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 ( θ ) :

We 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

where X ∈ { 0 , 1 } is a Bernoulli random variable with X ∼ Bern( p/ ( p + q )) , and

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 by

(ii) a set of local minima is given by

̸

(iii) a set of saddle points is

(iv) a set of stationary points is

Thus the set of all critical points is

In addition, for any θ ⋆ ∈ Γ ⋆ , θ min ∈ Γ min and θ sad ∈ Γ sad , the loss values satisfy

̸

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 by

(ii) a set of local minima is given by

(iii) a set of local maxima is given by

(iv) a set of saddle points is

̸

Defining a set of stationary points Θ station ( p, q ) ≜ { ( e, w, a ) ∈ R 3 : e = 0 , 1 + ae 2 = 0 , 1 + 2 w | w | = 0 } , the set of all critical points is

In addition, for any θ ⋆ ∈ Θ ⋆ , θ min ∈ Θ min , and θ sad ∈ Θ sad , the loss values satisfy

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.

G.3 Gradient flow analysis

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 by

starting with a randomly initalized θ 0 . We define the energy function E ( · , · , · ) as

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

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

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.

G.4 Role of standard initialization

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)

0.5

0.5

e

0.0

e

0.0

5

-0.5Q0

5

0.5

ILocal Minima

Local Minima

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.

H Additional empirical results

H.1 Gaussian initialization converges to low-rank

0

5

0

5

10

15

20

25

30

0.04

0.02

0.00

0.02

0.04

(a)

W

1

at initialization

0

5

0

5

10

15

20

25

30

0.10

0.05

0.00

0.05

0.10

(b)

W

1

after 50 iterations

0

5

0

5

10

15

20

25

30

0.2

0.1

0.0

0.1

0.2

(c)

W

1

at convergence

0

2

4

6

0

1

2

3

4

5

6

7

0.02

0.00

0.02

0.04

(d)

W

V

at initialization

0

2

4

6

0

1

2

3

4

5

6

7

0.10

0.05

0.00

0.05

0.10

(e)

W

V

after 50 iterations

0

2

4

6

0

1

2

3

4

5

6

7

0.1

0.0

0.1

(f)

W

V

at convergence

H.2 Low-rank initialization stays low-rank

0

5

0

5

10

15

20

25

30

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

(a)

W

1

at initialization

0

5

0

5

10

15

20

25

30

0.08

0.06

0.04

0.02

0.00

0.02

0.04

0.06

0.08

(b)

W

1

after 50 iterations

0

5

0

5

10

15

20

25

30

0.10

0.05

0.00

0.05

0.10

(c)

W

1

at convergence

0

2

4

6

0

1

2

3

4

5

6

7

0.075 0.050 0.0250.000 0.025 0.050 0.075

(d)

W

V

at initialization

0

2

4

6

0

1

2

3

4

5

6

7

0.06

0.04

0.02

0.00

0.02

0.04

0.06

(e)

W

V

after 50 iterations

0

2

4

6

0

1

2

3

4

5

6

7

0.075 0.050 0.0250.000 0.025 0.050 0.075

(f)

W

V

at convergence

I Model architecture and hyper-parameters

J Proofs of theorems in App. D

J.1 Proof of Thm. 7

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 that

On the other hand, by definition, from Eq. (3), f γ ⋆ ( x n 1 ) = σ ( e 2 (1 + 2 w | w | ) x n + b -e 2 2 ) , where γ ⋆ = ( e, w, b ) . Since x n ∈ { 0 , 1 } , this can be further simplified to

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.

and hence

Since σ ( z ) = y for y ∈ (0 , 1) implies z = log y 1 -y , Eq. (52) can be rewritten as

Using 2 e 2 w | w | + b + e 2 2 = e 2 (1 + 2 w | w | ) + b -e 2 2 = e 2 (1 + 2 w | w | ) + log p 1 -p in the second equality above, we obtain

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

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:

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.

Rearragning the order of scalars and writing γ min = ( b, e, w ) , from Lemma 4, the Hessian of the loss at γ min is

By definition, γ min = ( b, e, w ) satisfies ( p + q -1)(1 + 2 w | 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:

̸

If w = -1 √ 2 , ( p + q -1)(1 + 2 w | 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) is

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 obtain

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 gap

̸

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

Proof. Let γ = ( e, w, b ) ∈ R 3 be such that ∇ L ( γ ) = ( ∂L ∂e , ∂L ∂w , ∂L ∂b ) ⊤ = 0 . By Lemma 3, we have

̸

where X ∼ Bern( p/ ( p + q )) , f 1 = σ ( 2 e 2 w | w | + b + e 2 2 ) + q -1 -σ ( b -e 2 2 ) + p , and f 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

which implies that b = log p q . Since w ∈ R is arbitrary, we see in this case that

̸

̸

(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 that

̸

̸

̸

̸

and hence f 1 + f 2 = 0 . Further Eq. (58) implies that f 1 π 1 + f 2 = 0 . Together, this implies f 1 = f 2 = 0 .

̸

Now suppose w = 0 . From Eq. (58), we have that E [ f 1 X + f 2 ] = f 1 π 1 + f 2 = 0 . Since e = 0 , Eq. (56) yields

̸

So f 1 = f 2 = 0 in this case too. Thus we have showed that whenever e = 0 , we have f 1 = f 2 = 0 . Recalling the expressions for f 1 and f 2 ,

and hence,

Substituting b -e 2 2 = log p 1 -p in the above equation,

and thus,

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

K Proofs of technical lemmas in App. D

K.1 Proof of Lemma 2

Proof. Recall from Eq. (14) that the cross-entropy loss L ( · ) is defined as

where f γ ( x n 1 ) = σ (logit n ) = σ ( e 2 (1 + 2 w | w | ) x n + b -e 2 2 ) from Eq. (3). For any Y ∈ { 0 , 1 } , Z ∈ R , using the fact that 1 -σ ( Z ) = σ ( -Z ) , we have -[ Y log σ ( Z ) + (1 -Y ) log(1 -σ ( Z ))] = -[ Y log σ ( Z ) + (1 -Y ) log σ ( -Z )] ( a ) = Y log(1 + exp( -Z )) + (1 -Y ) log(1 + exp( Z )) = log (1 + exp( -(2 Y -1) Z )) ( b ) = ℓ log ((2 Y -1) Z ) , (62)

where ( a ) and ( b ) follow from the definitions of sigmoid and the logistic functions: σ ( 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. (62), Eq. (61) simplifies to

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 as

K.2 Proof of Lemma 3

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 by

where f γ ( x n 1 ) = σ (logit n ) = σ ( e 2 (1 + 2 w | 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,

Now we define the error term E ( X,Y ) ≜ -( Y -σ ( e 2 (1 + 2 w | w | ) X + b -e 2 2 )) . Our goal is to show that E [ E ( X,Y ) | X ] = f 1 X + f 2 , where f 1 ≜ σ ( 2 e 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 have

Now taking the conditional expectation with respect to X and using the fact that E [ Y | X ] = P ( Y = 1 | X ) = (1 -p -q ) X + p (since Y | X ∼ P ( p, q ) ), we have

where ( a ) follows from the definition of f 1 and f 2 above. Thus Eq. (64) simplifies to

Letting θ = e, w , and b in the above equation, we finally obtain the individual gradients:

where

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

K.3 Proof of Lemma 4

Proof. Slightly changing the variable order, for any γ = ( b, e, w ) ∈ R 3 , we define

Recall 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 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

Now substituting Eq. (67) when computing the second-derivatives of L in Eq. (66), we obtain

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,

Congregating all the second derivatives from Eq. (68) and Eq. (69) into the Hessian H ( γ ) in Eq. (65), we finally obtain

L Proofs of lemmas in App. E

L.1 Proof of Lemma 5

Proof. Recall from Lemma 2 that for any θ = ( e, w ) ∈ R 2 and b ∈ R , we have

Since ℓ log ( · ) is a convex function, Y ∈ { 0 , 1 } and thus 2 Y -1 ∈ {± 1 } , the convexity of L in b follows from the following fact:

To find the optimal b ⋆ , we set the gradient ∂L ∂b = 0 . Thus from Lemma 3, we obtain

Substituting E ≜ e 2 (1 + 2 w | w | ) , B ≜ b -e 2 2 , and E [ X ] = π 1 = p/ ( p + q ) in the above equations,

Further simplifying,

In other words,

Defining x ≜ exp( B ) and A ≜ exp( E ) , we thus obtain the following quadratic equation in x and its corresponding roots:

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 + 2 w | 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 .

L.2 Proof of Lemma 6

Proof. The proof directly follows from Lemma 2 by substituting b = b ⋆ .

L.3 Proof of Lemma 7

Proof. By Danskin's theorem [12], it follows that for b ⋆ = argmin b ∈ R L ( θ , b ) and L ( θ ) = L ( θ , b ⋆ ) , we have ∇ θ L ( θ ) = ∇ θ L ( θ , b ⋆ ) . Using the gradient expressions of L ( θ , b ) w.r.t θ from Lemma 3, and using the fact that ∂L ∂b = E [ f 1 X + f 2 ] at b = b ⋆ , the claim follows.

L.4 Proof of Lemma 8

Proof. Since L ( θ ) = L ( θ , b ⋆ ) where b ⋆ = argmin b ∈ R L ( θ , b ) , the identity in Eq. (32) about the Hessian of the loss L with respect to θ follows from the classical result of [29, Lemma 2.2] about second-derivatives of extremal-value functions. Finally Eq. (33) follows from substituting the full Hessian in R 3 × 3 from Lemma 4 in this identity.

M Proofs of lemmas in App. F

M.1 Proof of Lemma 9

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 that

Dividing Eq. (70) by (1 + 2 w | w | ) e and Eq. (71) by 4 e 2 | w | , we have

Noting that d d t (sign( w ) · log | w | ) = 1 | w | d w d t , the above equation can be rewritten as

̸

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 d w d t = 0 at t = 0 and hence w t = 0 for all t ≥ 0 .

M.2 Proof of Lemma 10

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 d e d t < 0 and if e 0 < 0 , we have d e d t > 0 for all t ≥ 0 , which establishes our claim. We have from the GF and Lemma 7 that

and

Defining

̸

and substituting Eq. (74) and Eq. (73) in Eq. (72), we obtain

Since 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 d e t d t < 0 for all t ≥ 0 . That is, the RHS Eq. (75) is negative. Note that x 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 is

Hence

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

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 that

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 θ d t 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 that

where f 1 = σ ( 2 e 2 w | w | + b ⋆ + e 2 2 ) + q -1 -σ ( b ⋆ -e 2 2 ) + p , and f 2 = σ ( b ⋆ -e 2 2 ) -p with π 1 f 1 + f 2 = 0 . Given that only d e d t 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 relation π 1 f 1 + f 2 = 0 , we obtain

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 .

M.3 Proof of Lemma 11

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 .

M.4 Proof of Lemma 12

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 have

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 ) = -(2 w +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 ) = ∞ .

N Proofs of lemmas in App. G

N.1 Proof of Lemma 13

Proof. First we recall the loss with the bias L ( θ , b ) from Eq. (41):

where logit X ( θ , b ) = e 2 [( X -1 2 ) ( 1 + ae 2 ) (1 + 2 w | w | ) + w | w (1 + ae 2 ) | ] + b with θ = ( e, w, a ) . Using the fact that ℓ ′ log ( z ) = σ ( z ) -1 and 2 Y -1 ∈ {± 1 } , we have for any θ ∈ { e, w, a, b } that

Now we simplify σ (logit X ) using the fact that X ∈ { 0 , 1 } :

Thus the gradients in Eq. (79) are given by

Now we compute the individual gradients with respect to e, w, a and b . Recall that

Thus,

Substituting the above equation in Eq. (80), we obtain

Now we compute the derivative with respect to w .

Similarly, for a :

Finally, since ∇ b logit X = 1 , it follows from Eq. (80) that

For 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) that

Substituting Eq. (82) in Eq. (81) and setting the gradient to zero there, we have

Simplifying,

Using the definition of the sigmoid function and rearranging,

By definitions of z 1 and z 2 in Eq. (82), we have z 1 -z 2 = e 2 (1 + ae 2 )(1 + 2 w | w | ) and thus A ≜ exp( z 1 -z 2 ) = exp( e 2 (1 + ae 2 )(1 + 2 w | w | )) . Thus the quadratic equation in Eq. (83) simplifies to

On solving the quadratic equation for exp( z 1 ) :

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

We 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 that

On the other hand, by definition, from Eq. (3), f γ ⋆ ( x n 1 ) = σ ( e 2 [( x n -1 2 ) ( 1 + ae 2 ) (1 + 2 w | w | ) + w | w (1 + ae 2 ) | ] + b ) , where γ ⋆ = ( e, w, b, a ) . Since x n ∈ { 0 , 1 } , this can be further simplified to

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. (84) and Eq. (85) are also equal, i.e.

and hence

Since σ ( z ) = y for y ∈ (0 , 1) implies z = log y 1 -y , Eq. (87) can be rewritten as

Adding and subtracting the above two equations, we obtain

Thus γ ⋆ ∈ 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

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)

̸

From Eq. (89), we have that E X [ 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:

̸

Slightly changing the variable order, for any γ = ( b, e, w, a ) ∈ R 4 , we define

Recall that

From Eq. (89), 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, w and a are

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,

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

Now substituting Eq. (92) when computing the second-derivatives of L in Eq. (89), we obtain

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,

Congregating all the second derivatives from Eq. (93) and Eq. (94) into the Hessian H ( γ ) in Eq. (90), we finally obtain

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:

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 + 2 w | w | = 0 . For this set of points, the Hessian remains undefined because ∂f 1 ∂e , ∂f 2 ∂e , ∂f 1 ∂a , ∂f 2 ∂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 + 2 w | w | = 0 , we have z 1 = z 2 = b . Hence,

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

̸

This concludes the proof.

N.3 Proof of Lemma 14

Proof. Recall from Lemma 13 that for θ = ( e, w, a ) ∈ R 3 , we have

where X ∈ { 0 , 1 } is a Bernoulli random variable with X ∼ Bern( p/ ( p + q )) , and

̸

where the optimal bias b ⋆ is obtained by solving π 1 f 1 + f 2 = 0 . Using the definition of the gradient flow that ˙ θ = -∇ L ( θ ) for θ = θ t , we have

Similarly for a ,

Likewise, for e ,

By substituting the expressions of Eq. (95) and Eq. (96) into Eq. (97):

Thus, we obtain

On rearranging and simplifying:

Integrating the above equation on both sides:

̸

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 .

N.4 Proofs of Lemma 16, Lemma 15, and Thm. 3

Proof. Wenote 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.

N.5 Informal proof of Thm. 11

Proof. [Informal] For θ = ( e, w, a ) , recall from Lemma 13 that the derivative of the loss L with respect to e is

where

Assuming the initialization is very small, making any product of quantities in θ = ( e, w, a, b ) much smaller than the individual quantities. Therefore, we can consider these products to be approximately zero. That is, ∀ x, y ∈ θ , x ≥ xy & y ≥ xy & xy ≈ 0 . Hence,

Hence the gradient with respect to e is

Simplifying the expectation term, E X [( f 1 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

Since ( p + q -1)( π 1 -π 2 1 ) > 0 , e → 0 , which denotes it converges to the local minima.

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

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

Answer: [Yes]

Justification: All claims are in the form of theorems that are proved in Sections 3 and 4 and experimental results in Sec. 5.2, for which our code is made available for reproducibility.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

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

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Though we do not have an explicit section for limitations, we outline the shortcomings in our approach in the conclusion and list them as future direction.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theorems are stated with required assumptions and proofs are provided either in the appendix or the main paper.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

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

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide all our code as open access and describe the experimental setup in App. I.

Guidelines:

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

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

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

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

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide open access to the code, available at https://anonymous.4open. science/r/Local-to-Global-C70B/ .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All experiment details are in App. I.

Guidelines:

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

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All experiments are repeated 5 times.

Guidelines:

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

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

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

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

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

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

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

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

8. Experiments Compute Resources

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

Answer: [No]

Justification: These are unrelated to the results of the paper.

Guidelines:

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

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

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

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

9. Code Of Ethics

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

Answer: [Yes]

Justification: All ethical practices were followed.

Guidelines:

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

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

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

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: No direct societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

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

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

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

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

11. Safeguards

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

Answer: [NA]

Justification: No risk.

Guidelines:

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

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

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

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

12. Licenses for existing assets

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

Answer: [Yes]

Justification: We use code from [26] which is cited.

Guidelines:

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

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

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

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

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

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

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

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No assets are releasd, but code is made publicly available.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: No human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper 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.\n\nSTRENGTHS:\nThis paper is the first result highlighting the role of initialization. Theoretically, they simplify the model to a low-rank parameterization with a fixed subspace for each parameter and rigorously characterize the landscape and trajectory. With this simplified model, all the possible stationary points can be found, and the dynamics can be calculated given different initializations. The simplification is somehow corroborated by the empirical evidence, and the analysis is complicated. \n\nSurprisingly, empirically a random low-rank initialization can outperform the standard initialization on a linear transformer on the Markovian task. It implies the theoretical insights from the low-rank parameterization provide some guidelines for transformer initialization.\n\nWEAKNESSES:\n1. 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.\n\n2. 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.\n\nQUESTIONS:\n1. Does the low-rank initialization scheme insight transfer to more general data distribution or architecture?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe paper is generally well written and readable, with well explained theory (in the main body at least, I have not checked the appendix proofs carefully) and high quality visualizations. I like the step by step presentation of the architecture.\n\nTheoretically analyzing the training dynamics and initializations of transformers is definitely significant, and currently somewhat rare to see even in toy settings such as this one.\n\nPost Rebuttal: I believe the paper is technically very solid, with my main concern still being the contribution. I did not raise my confidence only because I am less confident that the contribution quite reaches \"moderate to high\", but I am more confident (4/5) in the other aspects of the paper.\n\nWEAKNESSES:\n1. 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).\n\nIt 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.\n\nI 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.\n\n2. 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.\n\n3. 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).\n\n4. 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).\n\nQUESTIONS:\n1. 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.\n\n2. 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?\n\nLIMITATIONS:\nThe 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).", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: 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.

STRENGTHS:
1. This paper studies the training dynamics of transformers in learning Markov chains and study the effects of initialization of parameters, which is a meaningful direction which have not been fully explored.
2. The paper is well-structured and well-written, offering clear explanations and visualizations that enhance comprehension.
3. The empirical studies are well-designed and the results convincingly support the theoretical findings.

WEAKNESSES:
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:
Please see the weakness.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n-  The paper provides a detailed theoretical framework that explains the conditions under which transformer parameters converge to global or local minima, filling a gap in the understanding of transformer learning dynamics.\n- The theoretical findings are supported by empirical experiments, which strengthens the credibility and applicability of the results.\n- The authors highlight the critical role of initialization in the training process and offer practical guidelines for initializing transformer parameters, which can be valuable for practitioners looking to optimize transformer training.\n\nWEAKNESSES:\n- 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.\n- 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.\n- The empirical validation could benefit from a broader range of experiments, including different types of datasets and more complex transformer architectures.\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThis paper studies the dynamics of a simplified transformer model under next-token prediction loss. This extends previous work, which just studied the static landscape of the loss. In its study of dynamics, this paper also newly considers questions about initialization and its role in convergence toward global minima.\n\nThe paper is written clearly and generally supports its claims, highlighting the questions asked and key contributions.\n\nWEAKNESSES:\nIt 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.\n\nQUESTIONS:\n1. 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.\n2. 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.\n3. Do similar empirical results hold with larger alphabets? More specifically, do similar low-rank approximations still hold, and do these initializations still work better?\n4. 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.\n5. Figures 2 and 3 should be subfigures.\n\nLIMITATIONS:\nThe 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).", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-25"}