Local to Global: Learning Dynamics and Effect of Initialization for Transformers
Authors: Ashok Vardhan Makkuva
Marco Bondaschi
Adway Girish
Alliot Nagle
Hyeji Kim
Michael Gastpar
Chanakya Ekbote
Venue: NeurIPS 2024 poster (2024)
Keywords: gradient descent dynamics
transformers
markov chains
Primary area: interpretability_and_explainability

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
Ashok Vardhan Makkuva
Marco Bondaschi
Chanakya Ekbote EPFL
Adway Girish EPFL
EPFL
EPFL
Alliot Nagle UT Austin
Hyeji Kim
Michael Gastpar
UT Austin
EPFL
Abstract
In recent years, transformer-basse based models have revolutionized deep learning, par- ticularly 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 comprehensive characterization of the learning dynamics in this context. Specifically, we prove that transformer parameters trained 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 our knowledge, this is the first result of its kind highlighting the role of initialization. We further demonstrate that our thcoretical findings are corroborated by empirical evidence. Based on these insights, provide guidelines for the initialization of single-layer transformers and demonstrate their effectiveness. Finally, we outline several open problems this arena. Code is available https: // thub. Bt////onnn99/aalln
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, notable emerging theme among several recent works is model the input data as a Markov process.
Using this Markovian perspective, works such as [24, 14, 8J. among others, study in-context learn- ing capabilities of transformer. [23] analyzes the loss- 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 char- acterization of their training dynamics vis-á-vis the data distributional properties and the role of initialization still missing.
To address this gap, in this paper, focus 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 significant role in the convergence the transformer parameters to either local or global minima on the loss surface. Further, precisely characterize (Figs. and Fig. 2) the specific data characteristics and the region of initialization under which this convergence occurs. Based on these insights, we provide guidelines
"quual contribution. Correspondence ashok. www.knakerlll..ctt.....
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 dy- namics 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 initial- ization in transformer parameter training. Specifically, we demonstrate how the standard Gaussian initialization scheme can lead the convergence to local global minima depending on the Markovian data properties (Thms. and 8, Figs. and 2).
Guidelines: Based on these insights, we provide practical guidelines for parameter initial- ization, corroborated by empirical evidence demonstrating their effectiveness (Sec. 5.2).
Siba nOs mivi
(OCN PErIOA
Local medea
Locd THOTTE
Sandle poire
1
(a) Gradicnt flow (p 4 < 1) (b) Convergence basin (p+ <
Gradient flow (p- 1) (d) Convergence basin (p+
Figure I: Gradient flow dynamics and initialization effect for single-layer transformers. (p, are Markov switching probabilities, and w) are the embedding and weight parameters (Sec. 2). (a). (c): The flow is aligned along energy contour lines, converging to local or global optima. (b), (d): I, is the basin of convergence for global minima, Tmni for the local minima, and yellow asymptotes for the saddle point. Notice the contrasting behavior for Gaussian initialization around origin for + 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 &n-norm for Euclidean vectors and Frobenius norm for matrices. [k] {1,..., }., and for sequence (2r))21, define xt & .., if k 2 1 and (#1....,.25) otherwise. For 2 IR, the sigmoid (2) e 1/(1+e ", Rz)U(z) & max(0, 2) and the convex logistic loss flog(2 (2) A log (1 exp(-2)) E (0, 20). For events A and B, (A) denotes the probability of A whereas P (A B) the conditional probability. Let (r, y) be a pair of discrete random variables on [k] [A] with the probability mass function (pmf) of æ being Pa = (PI,..., Pp)) E [0, 1jkt Then its Shannon entropy is defined as (x) = H(p)) - -Eie P. log P:. The conditional
entropy is defined to be H (y(xz) & (x,, y) HI). The entropy rate of a stochastic process (I )221 y mean P (xr y) = 1. We also the shorthand P (y = æ) for TP (y = x == x) as a function of the random variable x. For p E (0, 1), the binary entropy function h(-)) is defined as h(p)) A -p log (1 p) log(1 p).
is defined 1im H (xT)
) nn. We simply write
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,, IN n=1 E (0. 1] is a first-order time- homogenous Markov chain with a fixed kernel P (Pt)). That is, the transition probability

---

Pej A (2 Fn = =
n- .i,j i,j E {0.1}. In = ) for any मं ++ nn = 0) and
particular, we consider P = P. P: 9, 1 where p Po1
P10 IP [In+1= 0 In. 1) denote the switching probabilities from the states 0 and respectively, We call - the switching factor. We assume that the process is already mixed, i.c. 2n ~ π for all n, where π _ (To, TT) )- = (g.p)/(p (p++)) q) is the stationary distribution satisfying ा = π TP Succinetly, - ptg hq +p plg)), and
wwitchnngfactor.
(In)n21 (T,, P). For this process, the entropy rate, H (In+11/ (T,))
the entropy of the marginal, M (Ix) = (T)) 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 {r/, 1 ,. it performs the following mathematical operations at each E [M] to predict the next-token probability fo (z]):
Embedding,
Attention,
Feed-forward,
Linear
Prediction
E 1 where
nn
yn
Zn
logit,
So fo(T))
In = In + P. E Rd
(Embedding)
y, = In + Σ att W x; E R
(Attention)
iE(n E(U,1)
2n yn + W2 W2RcLU(W yn) Rd
(Feed-forward)
logit, (a, kn) +
R.
(Linear)
fo - Fo (In+1 = = = E [0, 1].
(Prediction)
Here 0 é (e, (p W 1, W2.b, E R." denotes the full list of the transformer parameters from the embedding layer till the linear layer ($ A details them). While the underlying data {ar,, J.Y is Markovian, i.e. TP (2 = xT) = P (xn+1== =1 v..), the transformer agnostic to this fact and it can potentially utilize the full past in the Attention layer, via the attention weights att,, to predict the next-symbol probability fo(xT) ) = Pe (In+1 = =1 xT) Note that it suffices to estimate the sym- bol 1 probability as the vocabulary is binary. We also refer the above architecture as "full model".
Loss and training. The transformer parameters are usually initialized according to standard Gaussian distribution (10, o²I) with small variance o² [26] and are trained using gradient-based methods minimize the cross-entropy loss on the next-token prediction, i.e.
min L()). L(0) &
Σ DE[N]
+ [xn+1 log fo (z + In+1) log(1 - fo("))) (1)
N
When the input sequence {I,, IN ~ (T, P), the minimal loss equals its entropy-rate, i.e. L, A ming L(O)- H (zn-1122, [1]].
Loss landscape. A key surprising observation in [23] is that the loss function L(-) admits both global and local minima depending the switching factor + q of the Markovian data, and the weight-tying of the embedding and linear weights (e = a) the transformer. In particular, they show that
(i) for all (p.g) 10, there exists a global minimum 0, for the loss L such that its prediction matches the Markov kernel, i.c. Fe. (In+1 = 1 x²) = I (Ix+11 = In).
(ii) if p + g > 1 and the weights are tied (e = a), there exists bad local minimum 0min for whose prediction equals the marginal, i.e. Pe,, (#n+1
x;) = P (In+1 = 1).
In view of these results, focus on the weight-tying scenario and hence let a to be a single parameter in Rdd Thus, 0 = (e fP..... P W1, 2 b). We interchangeably refer to 0 as both the transformer and the parameters.
Our objective. While the aforementioned results detail the static landscape of the loss, they do not characterize the learning dynamics 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 0 to the local or global optima?
3

---

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 the motivating idea behind our approach: when trained by gradient-basee based methods, the weight matrices (W (Wv..... W1, 1, W2) at the optima O, 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,, are all aligned in. Interestingly, such low-rank solutions can also be shown to be thcoretically optimal ($ B details these structures). While these observations illustrate the implicit bias towards low-rank solutions at the final convergence, 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 thorough empirical evaluation for single-layer transformers and inspired by these empirical phenomena, without loss of generality, we restrict our attention these low-rank manifolds characterize the learning dynamics. This is similar in spirit to [24], where they assume special attention matrix structure for learning induction heads.
Parameterizzttinn More specifically, empirically observed and capitalize on to address (Q.1). Interestingly, along this low-rank manifold, it suffices to consider reduced set of parameters 0 R or 0 I given by:
consider special low-rank parameterization that is
0 (e,w) E IR?, or 0 (e, w, a) R
(Reparamcterizattion)
Here denotes the embedding scalar, w the weight, and the attention parameter respectively. Now we describe the parameterization of the transformer vis-á-vis these scalars and refer to ş for more detailed descripton. Let the input {In J be a first-order Markov chain in Sec. 2 and let E [N] be fixed. Then we have
Embedding e Pr = =-22) a In =
In - aa CER,OE(tI)
Attention:
WV a
1 yn = In -
a + (v, a
Σ attn,i EE|n]
e
2i (x-2) a, ER.
ocas:0
The scalar is the product of (v, a) and the scaling in the attention weights att,,; (Eq. (39)), which is empirically close to zero for first-order Markov chains. Hence for the case of exposition, we first let 0 and analyze the general case when E in Sec. 4.1. We continue:
Feed-forward Wi
a W Vd
|wl
Vd a1 + 2,, = 1 - (1 4w|w(ar) a, w E R.
1 is the all-one vector in Rr with r 4d typically in practice. Substituting this nn in the linear layer with e = a and bias b R, the logits and the probabilities simplify to:
Linear: logit, (e, w. b) = e² + 2w/w() xn + bb ERR
Prediction f(0.0) = (logit,,) E (0, 1), 0 A (e, w).
Finally, using the equivalence between the cross-entropy loss and the logistic loss bog (-), the loss function in Eq. (1) can be compactly written (Lemma 6):
L(O, b)
Σ DE[N]
Elloo ((2xn+1 - 1) logit,, (0)),, 0 RERR
(4)
N
Due to convexity of hog it follows that L(0,b) is convex in the bias b for any fixed 0, whose minimizer, (0) argmnineeR L(0,b), b). has a closed form expression (Lemma 5). Hence, without loss of generality, we consider the loss with this optimal bias b.:
(()) A I(O,bb)
1
Σ NE[N]
E flog (2 11) (2211 2w 2w/ww() Fn, + b. e²2 22) 2
(5)
N

---

Empirically, this roughly translates to running the gradient- based algorithm for the bias for more
steps at each 0. In practice, one additional step usually sufficient (see Sec. 5). Eq. (5) resembles the standard logistic regression loss [31] whose binary labels are 2 1 E (+1) and the logits
given by (1 + 2 2wrlw() zn, + b e² /2, for each n E [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 0 = (e, w, a) E RR.. First we study the scenario when of 0 = (e, w. a) in Sec. 4.1.
with (e, w) and build upon these observations to study the general setting
3.1 Loss Landscape with Canonical Parameterization
With the new set of parameters 0 = (e, w) E R², we are now ready to analyze the loss L(-) in Eq. (5). First we recall the definition of a critical point [20]. A point 0, E R2 is a critical or a stationary point for L if (L().)
0. A critical point 0, is local minimum if there exists neighborhood U around 0, such that L(O.) < L( for all 0 E U, and a local maximum if L(O.) Z L(O). If the neighborhood U is whole of TR2 it is global minimum/maximun. On the other hand, critical point is saddle point for all neighborhoods U around 0., there are 1 0₂ E U such that L(01)= I(O.)= I L(O2).
Thm. below provides a complete characterization of the loss landscape in terms of the aforemen- tioned critical points.
Theorem 1 (All critical points). Let the input sequence be {x 1NN (T,, P), the transformer parameters 0 = (e, w) E R2, and the next-token prediction loss L(.) be as in Eq. (5). Then for any (p. 9) E (0, 1) with p # and N E N.
(i) the set of all global minima is given by
(p. ((e,w) (RW))RR::0. e 2w/w() log (1 -1-93 p1 - 9)
(6)
pq
i the set of all local minima is given by
min (p. 4) & {(e, w) E (
> 0}.
(7)
(ii) the set of all local maxima is given by
+ max (p, 4) A {(e, w)) E :e= :e=
0, (p+ 1) -1 ww/w/) < 0}}
(8)
(iv) and the set of all saddle points is
snd (p. a (0.0-11V22)
(9)
Thus the set of all critical points is In addition. for any 0, E 0 E ① minn- θ maxx 2 maxx and 0 sad E ② Ssads the loss values satisfy
{0 E R2 ( 0} = θ U min max U + sad.
(10)
H(Ia+1 IT) = L(0.) < L(@min) = [(fma))= = L(9gad) = H(xn+1).
Proof. We refer to $ E.
Fig. 1 illustrates the loci of these critical points for p < 1 and p p- > 1. Motivated by empirical observations, while [23] characterizes local minima for p + g > 1, it is interesting to note that our Thm. 1 shows that local minima also exist for + y < 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. /) lies the learning dynamics and initialization for 0, which we study in the next section.
5

---

4 Learning Dynamics
Capitalizing on the loss landscape in terms of the critical points in Thm. I, 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, (0)) on L is governed by
d0, di
VL(O.), 0. = (ec, wx) E R., 02
(GF)
where 0, - (t) a C (continuously differentiable) curve in IR2 starting with randomly initalized Pu- To characterize these trajectories, we define an energy function E(;, )) which plays crucial in the GF dynamics. It is defined
E(e, w) - c² - (w² + sign(w) log////)) (w]), v(e, w) E R² c-axis,
(11)
where c-axis _ {(e, w = 0)}. 00). Note that & is well-defined and finite for all the points in its domain.
On the other hand, -im- E(e, w) = oo whereas lim {e, w) = oo for any fixed e E R. Thus the e- e-xxis 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 captured in the following lemma.
Lemma 1 (Constant energy along the flow). For any (p.9) E (0, 1)² and initialization Oo (eo, Wo) E RR, let (0z)t20 be the corresponding GF trajectory starting from Og. 1f0g E RR2 e-axis, the energy stays constant along the trajectory, i.e.
E(O,) = [w] sign(w:) log (ux,)) E 11 2 0.
(12)
On the other hand, if 0, E c-axis, have that 0 E c-axis for all 2 0 with wa We = 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. highlight the role of the switching factor of the Markovian data, + and the parameter initialization, O. in deciding whether the GF converges to local optima or global optima. First we define the energy value
Exad A E(c- 0, w = 1/V2) -1 log 2)/
Theorem 2 (GF dynamics for p 1). Let (p.9) E (0, 1) with p q > 1, the input sequence be {In | (π, P), and (0z) be the corresponding GF trajectory starting from Ou Then for all initializations 0g E RR2, the gradient flow converges to a critical point of the loss That is, there exists a Pim E RR² such that lim θ:- = Pum and V ((fmm) = 0. In particular, Oim is a
local minimum if
Oo E Immin A (e,w)) : ( (-1/22,)) e eE (-g(w). g(w) Vw2 log( + Exad
Uf(e,w) : w 2 0},
(i)) saddle point if 0g E Innd {(e,w) : E == + + llo( (og(-w) sw)++ + Exnd
(ii) local maximum if 0n Imax ((e =,=0, <
(iv) and a global minimum 0. EI. =I (Imin U U Ummax).
Consequently, when + q> if we use the standard initialization Og NO, 0212) with o² E 1/v22, Otim will be a local minimum with high probability, pf 1, under the same initialization scheme, Oim will be global minimum with high probability.
Proof sketch. The main idea behind the proof to show that if we do not initialize on the c-axis. the flows stays on the constant energy contour (Lemma 1) and hence converges to a critical point of the loss L, which 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
Imin; ..., I. are obtained by showing that an initialization a specific set leads the said critical
point (Thm. 1). The proof for 0g e-axis is similar.

---

Figs. 1b and 1d illustrate these initialization sets corresponding to the convergence basins for p= = 0.9 and p q === 0.1 respectively. An analogous result about GF dynamics for p + q is presented in Thm. 8 ($ F.2). Here a key difference is that small Gaussian initialization around origin leads to a global minimum fim with high probability (Fig. Id).
Key insights. Together, Thm. 2 and Thm. 8 address our motivating question (Q.1) by fully character- izing 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 > 1 more often than for ppt q 1, owing to standard Gaussian initialization around origin (Figs. lb and 1d). However, in practice, we often do not know the input switching factor, raising natural questions: is there data-agnostic initialization that always converges to global minima? Indeed, as can be seen from Figs. lb and 1d, there is a common region of initialization in the negative half-plane above the saddle- asymptotes (in yellow) that leads to the global minima convergence irrespective of the switching p+ g. Mathematically, this region given by Ioommon _ {(e, w) w < 0, lel > vw² log(loo[ w) + Esad We empirically corroborate this fact in Sec. 5.2.
4.1 Gradient Flow with Attention
60 Sefrmmn alcod mine (P cooudd witla winu evie worvame nerboce
Croude ara atoccl ghationue slgon
(0) p + < 1: Gradient flow
(b) Energy manifold and minima
(c) + >> 1: Gradient flow
(d) Energy manifold and minima
Figure 2: Gradient flow dynamics for the canonical parameters 0 (e, w, a) E R with the attention scalar a. Notice the contrasting behavior for Gaussian initialization around origin for p + smaller and greater than one. For an enhanced view of the flow near the origin, please refer Fig. 5.
In this section, the consider the attention scalar a E TR (Sec. 3) and study the gradient flow dynamics with the parameters 0 (e, a) E R3 The parameter captures the overall scaling from the value, key, and query components the attention layer. Recall that the soft-max attention weights are given by att,,; x exp((q,.: k,)/vd, /d,, where qn = W ax and k; W K2; are the query and key embeddings for any position i E [n:. Using the low-rank structure of the query and key matrices, satisfying W W (d aa and the value matrix Wy ov for some q IR and v E IRd (S G), and assuming linear attention att,, X (gn k:) /d, we define single scalar (v, dd5/// /4 that captures the essence of the attention layer (Eq. (39)). We note that linear attention weights are a standard assumption the transformer analysis literature [3, 36). Using this parameterization, similar to the steps in Sec. 3, we obtain the final loss function to be
T(O)) b1og (22-11)
(1 (xx(-)))
(x((x-))
ac²) (1 + 2 w)l wo/w(1 + ae?) + by ))
where 0 (e, a) and b, is the corresponding optimal bias. L recovers the loss in Eq. (5) when a O. In Thm. 10, we determine the set all critical points of in terms global minima and local optima in closed-form expressions, analogous to Thm. 1. Capitalizing we now shift our focus to the analysis the gradient flow in R To this end, let (0,) be a C'l curve in R3 governed by
this characterization.
d0. V 0. = (et, We, ar) E R t2 0,
(GF-attn)
dt
starting with randomly initalized 0g. We define the energy function E(.,,, as E(e, a. A e² - (w² + sign(w) log (w]) - -2²2, V(e, w. a) R33 ea-plane,
(13)
7

---

where ea-plane = 0, ). It is similar to its counterpart in Eq. (11), except for the 2a term. Fig. 2 visualizes this energy surface and the set of critical points, which reveal close resemblance to that of Fig. 1 in R2P Capitalizing on the energy function, we now present our main result with the attention.
Theorem 3 (GF dynamics with attention). For any (p, q) E (0, 1)2 and initialization Og E R3, let (01)c20 be the corresponding GF-attn trajectory starting from it. Then for all 0 E R², the gradient flow converges
critical point of the loss L. That is, there exists a Otm E R such that lim 0- = lim and V ((0mm)) = 0. Further,
(i) if 0o E R \ ca-plane, we have [(0 E(2,) = E(Ou) for all t 2 0. Hence Plm is at the intersection of the energy contour line Eg with that of the set of critical points.
ii) if 0u E ea- ca-plane, we have 0, E ea-plane for all t 2 and hence fumm E ea- ea-plane.
Proof. We refer to $ G and $ N.4.
Thm. 3 shows that the learning dynamics with attention closely resemble those without it (Thms. and 8). While the set of all critical points of L, and thus the limit points of the flow. has closed-form expression (Thm. 10), deriving the same for the initialization sets Iimi 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 similar picture as in the two-dimensional setting for both the p interesting direction of future research to theoretically characterize this, analogous to Thms. and 8. We refer G for additional details and proofs.
and p + q cases (Fig. 2). We believe it's an
5 Empirical Results
We empirically validate our canonical parameterization 0 E R (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 We conclude with discussion on higher-order and multi-state Markov chains.
IR
5.1 Low-rank Parameters
Full model converges to low-rank. We let the input Markov sequence to be L.W ~ P(p, g)) for p 0.2, 0.3, 1024 and consider the full model as defined in Sec. 2 with embedding dimension d = First, we initialize the parameters 0 = (e a, {p.. (pa)==I..... W,, 1. W2, 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 f. = 800 iterations. In Fig. 6, track the value matrix Wy Rddxd and the weight matrix W1 Rddddd across iterations. We observe that at convergence both Wv and W, are approximately rank-one with one of their components being same the embedding vector (the row in Wy v and column in W D. 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, W2 and for all values of (p, q) E (0, 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, restrictions, and track them during the course of training. Interestingly, here we observe that the parameters still stay low-rank 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.
train them without any low-rank
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. Id
8

---

1.2
1.0
Fraction of energy
Fracioon enerrgy
0.8
0.8
0.6
0.4
Random inftialization
0.2
Random initalizztion
Rank-one initialization
Rank-one inilialization
0.0
20
40
60
80
100
0
10
20
30
40
50
60
70
80
Iteration (mulliples 100)
(multiples 100)
(a) W
v
(b) W
Figure Convergence to rank one parameters: percentage of energy contained in the first rank component of the weight matrices W and Wr across iterations. The percentage is computed as where the o, are the singular values of the matrices in descending order.
of
0.90
Standard initalization
0.85
Our initialization
Unigram loss
Bigram
Test
0.75
0.65
0.60
100
200
300
400
500
600
Iteration
Figure 4: Comparison between the average loss curve for the standard gaussian initialization around 0 and our initialization, for p 0.5 and g 0.8. Starting from the standard initialization, the model converges to a local minimum corresponding to the unigram model. With our initialization, converges the global minimum corresponding to the bigram model.
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.. pr.......... W1. W2, b) are randomly chosen around zero with small variance = 0.02, with new initialization based on our results, where we initialize the embedding vector such that all cordinates are equal to e 0.5, W to be constant with the scalar w1 = and Wa constant with W/ -1 (corresponding to I, in Fig. Id). 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 similar spirit, analysis of initialization effects for deeper architectures is an interesting avenue of future research.
5.3 Higher-Order and Multi-State Markoy 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 line with our approach an interesting direction. On the other hand, a recent
work [28] establishes surprising result that any 4th -oddrr Markov chain be represented by

---

three layer transformer with just one head per layer, relying induction head mechanism. Analyzing gradient flow dynamics using appropriate canonical parameterization (cf. [24]) this scenario 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 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-basee an induction-heas mechanism using simplified transformer architecture with frozen encodings, query matrices and linear activations. On the other hand, we consider the canonical parameterization, capitalizing inherent low-rank parameters, to provide full characterization of the learning dynamics. [24] demonstrates how two-layer transformers with GD learn induction head mechanism when the input has causal tree dependency, such as in Markov chains. In this work, 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 simplified single-layer transformer, while [33] studies the process by which self-at 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 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 's capability in modeling Markov chains and the associated training dynamics.
updates can learn
7 Conclusion
In this work, present a novel characterization of gradient flow dynamics for (weight-tied) single- layer 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
[1] 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-54) 1-5477 2005.
[2] Arzu Ahmadova. Convergence results for gradient flow and gradient descent systems in the artificial neural network training. arXiv preprini arXiv:2306. 13086, 2023.
[3] 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 httph:t////// //tpp://///peeeiieew nwt./forumrr id=LziniAXEI9
10

---

[4] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? Investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
[5] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks, 2019.
[6] Francis Bach. Effortless optimization through gradient flows. https///ranciibaa..comoguuueeeee flows/ 2020.
[7] 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.
[8] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer. memory viewpoint. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[9] Pritam Chandra, Tanmay Kumar Sinha, Kabir Ahuja, Ankit Garg, and Navin Goyal. Towards analyzing self-attention via linear neural network, 2024. URL forum?id=4f VuBf5HE9.
net/
[10] 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.
[11] Thomas M Cover and Joy A Thomas. Elements of information theory. John Wiley & Sons, 2nd edition, 2006.
[12] John M Danskin. The theory of max-min, with applications. SIAM Journal on Applied Mathematics, 14(4):641-665 1966.
[13] Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, Lei Li, and Zhifang Sui. survey on in-context learning, 2023. URL https://arxii org/abs/ /t////2311.0222...
[14] Benjamin I.. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains, 2024.
[15] 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-Dodte Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. URL https: http//ttasssfttt ormer circuits http:tssspppp/e0e/eeee http//22222///rrooe/ waneeeioko//ie/xxxee/e index.nndxx//// html.
[16] Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can single attention layer learn? A study through the random features lens. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[17] Shivam Garg, Dimitris Tsipras, Percy 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.
[18] M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self-attention to markov models: Unveiling the dynamics generative transformers. arXiv preprint arXiv:2402. 135 2024.
[19] 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.
[20] Jason Lee, Max Simehowitz, Michael Jordan, and Benjamin Recht. Gradient descent only converges to minimizers. In Conference learning theory, pages 1246-1257. PMLR, 2016.
11

---

[21] Yingeong Li, M. Emrullah Iidiz, 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.
[22] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding, 2023.
[23] 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.
[24] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402. 14735, 2024.
[25] 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.
[26] Matteo Pagliardini. GPT-2 modular codebase implementation. https///gituub.cm//efff///lt baselines, 2023.
[27] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is Turing-complete.. Journal of Machine Learning Research, 22(75):1-35, 2021.
[28] Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vard- han Makkuva. Transformers on markov data: Constant depth suffices, 2024. URL https:
/ //arxiv. wwt.:oggg//.a.....ooo .1w.244777776.....
[29] Alexander Shapiro. Second-order derivatives of extremal-value functions and optimality condi- tions for semi-infinite programs. Mathematics of Operations Research, 10(2):207-219, 1985.
[30] Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. CORR, abs/2103.07601 2021. URL https://axxxi orgg httgg//bb//00...../
2103. 2160.07601.
07601.
[31] 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 170):1-57, 2018.
[32] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. In Thirty- Conference Neural Information Processing Systems, 2023.
[33] Yuandong Tian, Yiping Wang. Beidi Chen, and Simon Shaolei Du Scan and snap: Under- standing training dynamies and token composition in I-layer transformer. In Conference on Parsimony and Learning (Recent Spotlight Track), 2023.
[34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Informa- tion Processing Systems, pages 5998 9908,600, 2017.
[35] 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 Machine Learning, pages 35151-35174, 2023.
[36] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2023.
[37] 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
2022.
[38] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference Machine Learning, pages 11800-11000, 2021.
12

---

[39] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in- context learning as implicit Bayesian inference. arXiv preprint arXiv:2111.02080, 2021. URL http://aaxiw..ogabbbb///// http///111.0000...1111
[40] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of Conference Learning Representations, 2020.
functions? International
13

---

Contents
Introduction
Problem Setting
3
Canonical Low-rank Parameterization
3.1
Loss Landscape with Canonical Parameterization
4
Learning Dynamics
4.1
Gradient Flow with Attention
5
Empirical Results
5.1
Low-rank Parameters
5.2
Effect of Initialization: Broader Implications.
8
5.3
Higher-Order and Multi-Stare Markov Chains
6
Related Works
10
7
Conclusion
10
Single-layer transformer: architecture and results
16
A.I Loss landscape results
16
Low-rank structure of the optima
17
Canonical reparameterization
18
Analysis of the loss with the bias, L(O,b), in Eq. (4) and Eq. (23)
20
D.1 Technical lemmas
20
Analysis of the loss without bias, L()), and proof of Thm.
22
E.1 Proof of Thm.
23
Gradient flow analysis without attention
25
F.1 Proof of Thm. ................................
25
F.2 Gradient flow dynamics for q
27
G Gradient flow analysis with attention
29
G.1 Canonical parameterization with attention
29
G.2 Analysis of the loss function L(0) from Eq. (42)
................................
G.3 Gradient flow analysis
32
G.4 Role of standard initialization.
33
H Additional empirical results
34
14

---

H.1 Gaussian initialization converges to low-rank
34
H.2 Low-rank initialization stays low-rank
34
I Model architecture and hyper-parameters
36
J Proofs of theorems in App. D
37
J.1 Proof of Thm.
37
J.2 Proof of Thm.
39
K Proofs of technical lemmas in App. D
40
K.1 Proof of Lemma
40
K.2 Proof of Lemma 3
41
K.3 Proof of Lemma 4
42
Proofs of lemmas in App. E
45
L.1 Proof of Lemma 5
45
L.2 Proof of Lemma
45
L,3 Proof of Lemma
45
L,4 Proof of Lemma 8
45
M Proofs of lemmas in App.
46
M.1 Proof of Lemma
46
M.2 Proof of Lemma 10
46
M.3 Proof of Lemma 11
48
M.4 Proof of Lemma
N Proofs of lemmas in App. G
49
N.1 Proof of Lemma 13
49
N.2 Proofs of Thm. 9 and Thm. 10
51
N.3 Proof of Lemma 14
57
N.4 Proofs of Lemma 16, Lemma 15, and Thm. 3
59
N.5 Informal proof of Thm.
59
15

---

a) are trained via the cross-entropy loss A Single-layer transformer: architecture and results
We first describe the transformer architecture from Sec. 2:
Zn In e + P
R
(Embedding)
n/ nn + Σ
atbn, W V E R?,
(Attention)
ic/ml
2n yn + We (e.UeWeeWW..eeeeee vn R. E
(Feed-forward)
logit, = (a, zn) +
(Linear)
fe( & Pe (In+1 2 = o (logit,.)
(Prediction)
E(001]
Here 0 - (e, fp. W 2,, b, a) eters from the embedding layer till the linear layer. In the attention layer, the weight assigned to each value, atin,i, is computed by a compatibility function of the query vector qn - W En and the corresponding key vectors k; & W æ; for all E [n]. More precisely, attn,i & softmax(((q, k1),............... (q ka))/Vd), W K.Q.V E Rxdd 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.
R.D denotes the full list of the transformer param-
Finally, the transformer parameters 0 (e, on the next-token prediction:
L( _
Σ REIN]
B Ba [4+t1 log So(x)) + (1 1) log(1 - [o())))]
(14)
N
In this paper, focus the weight-tied scenario where Hence let them be single parameter with 0 = (e = a, (p J E RDD, 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 { (T((, q). P(v, 4)) for some fixed (p, q) E (0,1)3. Then for all (p. 9). there exists a 0, E RD with an explicit construction such that it is global minimum for the population loss L(-) in Eq. (14), i.e.
L(O)= 2 L(0-).or all O E RDD
Further, 0, satisfies:
(ii) Po. (#n+1 = x²)) P (In+1 = v.). the Markov kernel.
(ii) L(O.) (Znn11l lIn). the entropy rate of the Markov chain.
(iv) (O.)) =0, i.e. O2 is stationary point.
Let L, & L(O.) be the global minimal loss from Thm. 4. Now recall the result the bad local minimum.
Theorem 5 (Bad local minimum). Let the input sequence be )N (T(p, q), P(p,g)) q)) for some fixed (p, q) E (0, 1)2. pf local minimum for the loss L(-), i.e.
+ > 1, there exists an explicit O Omin E RDD such that it is a bad
(i there exists neighborhood B(0 ((0min))) r) with r > 0 such that L( Z L(0man) for all O B 1 r), with L(@min)) L..
Further, minn satisfies:
(ii) Po. (In+1 = a = TP (In+1 = 1) = nı, the marginal distribution.
(iii) L(0miu) = H (22+1)) = HTT), ), the entropy of the marginal.
(iv) L(( min) 0, i.e. Omin is stationary point.
16

---

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 q < [23, Appendix C.2J. In particular, it is observed that the token and positional encodings point in the same direction a, which is 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 obeys
e a
for some e > 0 and a E [=1]d. Further, the positional embeddings P,, are constant across positions pointing in the same direction albeit with negative scalar, i.e.
Pn -p-e-
for > 0 and & such that p. Thus from Embedding layer,
In = (EZn, p) ax,
(15)
which ensures that the respective embeddings for the bit In. 0 and 2,, = 1 are -p- and Rn (e (e-p)) p) a, which are roughly anti-podal.
Attention. Recall from the Attention layer that the output yn is given by yn = xn W LiE[n] att,, W v Fi where the attention weights ntt,, are computed according att,, exp (gn, k.)/vV / (Ejetn) exp (ig.. k5)////// ks)\vV) with qn We En and k; W k C. Here is observed that the matrix products are all rank-one with a being factor, i.e.
WoWy a. V. E Rdddd for some Rd
W WOWx Wxx (1² a..0.. a dddd
for some R.
Hence,
Wy x, = = a) p)
and
(qn, k.)
(IaaI[==d) g²dd
x, W W K
(z ax) a)
(ex, - p)x(e - p)
vd
Vd
Q
vd
vd
n2d 2/25
(eIn p) - p).
Thus.
yn In +
att
W WoWv xi
iE|n]
= (ex; - p) α +
attm, a) - -p)
sEin|
exp (g² (eIn p)(ex - p)) rE(m] Ejgjn] exp (92d3/2 (exn -V)(Er, p))
(er; -p) -p)
(or, ax)
(v,0)
(16)
It is further noticed that (v, a.) Ri 0 and hence y, (eIn - p) 0 Dn.
Feed-forward. For the Feed-forwail forward layer, both the matrices W E grxd and W2 E Rdxr exhibit oan structure with being one of the factors,
Wi
a for some {+1) ,
(17)
W W..
(18)
Thus Wiy, = dw(er, w. Since -p < 0 and e > 0, corresponding to x, = 0 and Fr = respectively, we obtain ReLU(W 19,)) du ((1 - Fn p ReLU(- w) + xn p) RcLU(w))..
Denoting the number of ones in as B, i.e. B = Σn= 1 (ww; = 1), we further simplify:
W2ReLU(Wing 19,) W ReLU(W 1
17

---

w2d ((1 Znip- (w, ( + 2n p) (w, ReLU(w)))
w2d ((1 xn)p (8 r) + In p) 3)
wdd(exn p) (128 r) + 3] a.
Hence
2n = yn + W2ReLU(WIsS) (eIn - p) (1 + 2 ((2 r)æn + --)--)) 3) a.
(19)
Linear. Using the fact that e = a = e a due to weight-tying, we obtain from Linear layer that logit, (e, xn) + b ed(ez, p) (1 w ((28 - r) + (r)) + b.
(20)
Prediction. We finally obtain that the prediction probability
[o(=T) - (logiit,)) In (ed(e p) (1 + Bw + b) + (1 xx,) (-edp (1 (r Byw + b)
Thus we see that the prediction probability and hence the loss function L(-) in Eq. (14) is influenced only by the scalars e, w, b and B.
C Canonical reparameterization
Building on the low-rank strucutre of the transformer parameters described above, we consider special parameterization for them. A key property of this parameterization is that it covers both the global and local minima from Thm. and Thm. for all (p, E (0. 1)2. Recall that Thm. characterizes local minima only for p g > I whereas our special parameterization allows to discover local minima even for p- + q < Our construction follows the same outline as in Eqs. (15)-(20). First we start with the embedding layer.
Embedding. We let e = e a and pn, = -p- a for all n where e >> 0,p and a E (+1)d//4 Thus the embedding æn from Eq. (15) simplities to
In = xn - 2
2/2
1 a
f e
)
2
Attention. Substituting this F, in Eq. (16), we have
1 2 Un e In - a + (v)
Σ
Iz - a. attn,i
(21)
(E(n]
exp(o? (sn (a,, 11)
where the attention weights att,, ΣjE( exp(e" (22 (=n (x,,, 2] Since (v, a) x 0, we let = 0 and obtain
(0, 1) for some E R.
yn = Fn =
I' a
Feed-forward. For the feed-forward layer, we observe from Eq, (17) and Eq. (19) that for any w E {=1 only the number of I's in w, 3, matters for the final vector zn which further interacts with the weight scalar w. Hence without loss generality, we set w to the all-ones vector: w = E I' and hence 3 = r = 4d. While we observe from Eq. (17) that W₂ = W for p + g < 1, we observe from the proof of the Thm. 4 for + q 1 in [23, Appendix B.2] that need W₂ = W in this scenario. Hence we consider the following parameterization that covers both these scenarios:
Iw Vd
Wi
1 T Radxd
W2 2 Vd
a I &Rdx4d
Here w 0 ensures Wa 2 I whereas 0. W2 W T. Using this parameterization.
substituting 3 = 4d and w + 7 in Eq. (19), we get
1 (1 + www/ a
2n = In -
18

---

Linear. Since e
a due to weight-tying, Eq. (20) simplifies to logit,, = zn) + b = In
(In
(1 + 4w ( + b
2
(x. e² + 220
1
2w 2o
+ b
2
e²
= e² +1+ 2w/w|) +
Prediction. The next-token prediction probability is
(() a o (e2 + 2w/ww) on + マ 2 0 (e,w)
(22)
Loss. While we assumed > 0 in the beginning. view Eq. (22) and fact that a E (#1)*/ Vd. we see that e E R gives the same expression for probability, Thus the final probability depends on just the three scalars (e, w, b) R$. Defining 0 = (e, w) E IR? we recall the cross-entropy loss L(.) from Eq. (4) in Sec. 2 for this canonical model:
L(O,6)
Σ NE
2T 1 log (x)) + (1 - In+1) log(1 - (1
(23)
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).
19

---

D Analysis of the loss with the bias, L(O,b) b), in Eq. (4) and Eq. (23)
In this section, we analyze the loss function with the bias, L(0, 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 To this end, define the following sets of points
(p( E
R c2 2w/w() log (1 -)pp - 9)
b
c²
log
P
pq
p
(24)
T min (p. (e, (e..,)) :, 1)(1 [00]]12wwwwuz 2w/w/l 0, log
p
(25)
q
T sa (e, :e= o, 1) + 2w/w/) 0, log
p
(26)
The following result establishes that these sets exhaust all the critical points.
Theorem 6 (All critical points). Let the input sequence be { ~ (T, P), the transformer parameters (0, (e, w. b) E R2, and the next-token prediction loss L(-) be as in Eq. (23). Then all the stationary points of I. are either in D,, I miin or I sad, i.e.
{(0,b) R VL(O, b) = 0) = r, UT min UT sad.
(27)
Proof. We refer to App. J.I.
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 I--)) Theorem 7 (Loss landscape with bias). Let the input sequence be {Ira IN former parameters (e, w. b) E TR3, and the next-token prediction loss T(,)) be as in Eq. (23). Then for any (p. q) E (0, 1)2 with p + q # 1 and N E N.
(æ. P), the trans-
the set of all global minima of L is given by (p,
(it) the set of all bad local minima of L is given by minn(p: 4).
(ii)) and the set of all saddle points of L is I Fead (p. 9)-
Furthermore, for any y, E IT,7Tin min and Yaad I sadd the losses ordered
H(E++1 2n,) L(r.) L(7min) L(Tsaa) H (xx+1).
Remark 1. Note that a bad local minimum is 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 both Thm. and Thm. 7 rely on few key lemmas that we present below. First we start with the result that rewrites the loss T.(0, b) from Eq. (23) in compact manner using the logistic function log (-).
Lemma 2 (Loss as logistic function). The next-token prediction loss I(.) in Eg. (23) can be written as
L(O,b)
Σ Blltog ([2mn41 - 1) -logit,,)l
N
REIN]
(28)
Ex foo (2Y-1) (211))221+2
+ - )
where (X, Y) E 102 are distributed according to (X, ~ (T,, P), 1),i i.e. a Bernoulli random variable with X = Bern(pp (p + 4)) and Y P(p, q), the Markov kernel.
20

---

The following lemma establishes the gradients of the loss function with respect to the parameters e, u, and b.
Lemma (Gradient computation). For any (e, w, b) R and the next-token prediction loss L(-) in
Eg. (23), the gradients are given by
aT.
de Ex [ (2) + 2w(w/) 1)] e.
AL 1(fix +22xx] 4e//w|,
dw
OL
Ex [fiX + [2],
ōb
where {0, 1 is Bernoulli random variable with X Bern(p/(p + 9)), fi a 2 -P.
122wo/wl + b + c² - o (b-
e²) + +p, and
Remark 2. It is interesting to note that the gradients for both and w are product of expectation term and an factor. Also, except for scaling factors in terms of (e, w. b). all the gradients are governed the two expectation terms B[( /2) X] and [B[ X + [2]. This observation plays 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 -irammm and saddle points). For the canonical parameterization 7y (b, e, w) R& and the next-token prediction loss T(.) in Eq. (23), the Hessian any Ymin T min or Yead E I sad is given by
VLL(M)
ToT
2(p + 1)(1 + 2w w|
Tmin Tead
where To p+g Remark 3. We note that the Hessian is computed with the parameter ordering (b, w).
9 and T1 p-t"""
P
The proofs of the lemmas are presented in App. K.
21

---

E Analysis of the loss without bias, L(0), and proof of Thm. 1
The proof of Thm. concerning the loss L( in Eq. (5), similar that Thm. 7 which studies the loss L(O, b) with the bias present. The main idea is to establish the analogous set of lemmas, as App. when the bias substituted with its optimal choice. First we recall the loss function
L( A L(O.b.)
Σ E Eegti[Zn+ log ((f)) + (1 Fn+1) -log(1 f(o. (z nE[N]
N
(29)
b.
argmin (),
bER
We start with the result that establishes closed form expression for b, Lemma 5 (Optimal bias). For 0 = (e, w) E I² and b E IR, let L(O,b) b) be the next-token prediction
loss defined in Eg. (23). Then, for any 0 E R2, L(0,b) b) is convex in b and the minimizer b. & argmingeR L(0, b) is given by
exp b+
1
+
(2-1)
P. +4.
A q
A exp(e? 2w/w))). (30)
2
2A
Consequently, c 2w(w() log
(1 (1-p)(1-g)
then b,
2²
log _. Ife ===00 0, then log P. I-p"
pq
Now rewrite the loss in terms of the logistic function.
Lemma 6 (Loss logistic function). For any 0 E IR?, the next-token prediction loss L( in Eq. (29) can be written as
L(
Σ DE[N]
Elliog ((2xn+1- 1) logit,,))
N
(31)
Ex,y
(los
(2Y 12 (e(
+ +b, (w(1+2w|wl)xx (1
62) 2)
where b. follows from Eq. (30), (X,Y) {0, 132 are distributed according to (X, ~ PY,P),, i.e. X is a Bernoulli random variable with ~ π = Bern(p/ q)) and Y|X ~ P(p, 9). the Markov kernel.
The following lemma establishes the gradients the loss.
Lemma 7 (Gradient computation). For any (e, w) R22 and the next-token prediction loss (())
in Eg. (29), the gradients are given by
aT.
= Ex [fhx + f2) 2(1 + 2w(w])e
de
Ow = Ex [ 12) de2/
where E (0, 1 is a Bernoulli random variable with X Bern(p/(p + 1)), fi 2e2w|w] w| + b. + e +q - 1-o (b. +p. and σ (b. - Further, Tifi+
We observe above that the gradients for both e and w are proportional to each other, except for the scaling factors in terms of and w. This forms the basis for the derivation of the energy function discussed App. F.
Remark 4.
The following lemma characterizes the Hessian.
Lemma 8 (Hessian computation). Let 7 = (b, 0) E R with 0 (e, w) E R2, and L(y) be the nexi-token prediction loss in Eg. (23) and L(O) be the one in Eq. (29). Let the Hessian of L at y be
H(>) A 22 22 =
Hbb
Hoo
V2LL.
V20
E 1333
H
Hoo
(V GoL)
L.
22

---

Then the Hessian of L at 0 E R2 is given by
(0) VoOL Hee Hio Hib Hoo-
(32)
Consequently, for any 7 (b, e, w) I min UT the Hessian (0) at (e, w) given by
2(p + 1)(1 + 2w(w|)
0
(0) = πoम
(33)
0
where To p+g and 1T p-t-q"
The proofs of the above lemmas are deferred App. L. We are now ready to present the proof of Thm. 1.
E.1 Proof of Thm. 1
Proof. Let E R² and 70) (0) E R be its embedding in R$ with the optimal bias
(0) = argminges b)) from Lemma Define the following four sets of points:
+ (p, A w) E ER² e + w = -log (1- 11pp((1 ②)
14
min (p. g) (e, E IR2 : 0, (p + q 1)(1 + 2w 2w/w() > 0)
max (p. _ ((e E <R² : 0, (p q - 1) + 2 2w/wwl < <0},
sad (p. { (e, : 0, w= = -1/22
First we show that any critical point of L RR2 + 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 y (0, (e, w),b) R3 of L. 7 E P. U miin U sad. Here the main observation is that these critical points are of the form (0, (0)) where 0 E O, U O min U O max U sad To see this, let y E r.. Here we have e² + 2w| wl) = log (1-p) from Eq. (24) and hence O E O.. Further, by Lemma 5, we have pg
that the optimal bias for this 0 satisfies b, - _ = log T which is precisely the characterization of the bias for Y = (e, w, in Eq. (24). Likewise, if E I miin U Fead, we have e = 0 and hence 0 Omin θ max U O Hence by Lemma 5, b, log q. matching that of Eq. (25) and Eq. (26). Thus the set of all critical points of L in IR3 are of the form (0, (0)) with where E O U U 10mn O U @ Since r, UT min U I sad covers the entirety of stationary points of L in R², it follows that the set of all stationary points in R² is precisely θ U 0 U O U + 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 7. (0., (..)) of L in RR3, we have O, E O_ R.2 Hence by definition, O, is the set of all global minima in IR. A similar argument holds for θ Omin- 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 7
(e, w. by (e, w)) of L in R we have that
and (p + q 1)(1 + 2w/ww) < 0. Hence 0 = (e, w)= @max U sadd- Suppose O θ mmax which implies = 0, g ( 2w/w)) 0. By Lemma 8, the Hessian at 0 (upto a positive scale) is a diagonal matrix with the entries (p + () + 2w/w) w/) < 0 and 0, corresponding to the directions of and w respectively. Though one of the eigenvalue here is zero, using a continuity argument as in the proof of Thm. for local minima, we can establish that 0 is indeed a local maximum. Thus O mmax is a set of local minima.
Now suppose (c, w) E aaad- Thus e 0 and -/2 Since it lies at the intersection of θ min and mmxx using a neighborhood argument, it's straightforward to see that e sdd indeed a set of saddle points.
23

---

Finally it follows that e min. 2 max ead 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.
24

---

F Gradient flow analysis without attention
In this section, we analyze the learning dynamics of the transformer parameters 0 (e, w) g without the attention scalar. First, we present few important lemmas regarding the same, useful for the proofs of Thm. 2 and Thm. later. Recall from Sec. 4 that the trajectory (0c)r20 is governed by
d0, dt = VL(O,), O, = (ee, wx) E RR, 2,,2. 0,
(GF)
starting with a randomly initalized Dg. The energy function ((; .) is defined as Ele, w) 2 e -w² + sign(w) log (wl), V(e, R² e-axis,
(34)
where e-axis - {(e, w 0}} and w-axis - {(e= 0, w)). Note that Esad = E(O, 1/2 We re-present the Lemma from Sec. 4 below for the sake of completeness.
1+log
2
Lemma 9 (Constant energy along the flow). For any (p,q) (0, 1)2 and initialization O, (eo, Wo) E R?, let (0z)e20 be the corresponding GF trajectory starting from Qu. If wo & 0, then the energy stays constant along the trajectory. i.e.
E(O.) = eτ (wî + sign(ws) log (uze) E((o),
V20 0.
(35)
On the other hand, if wi 0, wn 0 for all i 2 0. Hence, if we initialize e-axis the trajectory always stays on the e-axis.
Now we establish that the GF trajectories always converge. Lemma 10 (GF convergence). Let (0,)r20 continuously diferentiable GF trajectory starting from 0y Then for all initializations 0g E R
(0:)22 bounded,
i there exists Olmm E 2 such that lim4 O2 Olim and
(iii) lim [L(O))Il = [ = O.
Hence Oum is critical point of T..
The following result characterizes the energy of the limit point.
Lemma 11 (Energy at the limit point). Consider the same setting in Lemma 10. If0n E R (-aaxss then (00m = EOu). Hence fim lies at the intersection of the contour line E(e, w) = Ey with the set of critical points of L in R.?.
On the other hand, 0g e-axis, then fim E ee-xxis
We now study the energy function on the w-axis which plays key role in the GF analysis. Lemma 12 (Analysis of the energy function). Ler E(..) be the energy function defined in Eq. (48)
and (w) A E(e = 0, w) = -ww2 sign(w) log (w]) be the energy evaluated on w-axis for
E TR (0). Then
( f : (-o0, -1//2] + (-00, Esad] is monotonically increasing with lim (ww) = -00
and the maximum being f(-1/v2) Eead. Eead.
(i) 1=1-1/2,,0) + [Esad: -00) is monotonically decreasing with lim [(w) -00,
(iii) f'( - /2 = 0. and
(iv) f (0) + (-00, 00) is monotonically decreasing with lim f(w) = oo and
+oo fww) 00.
We are now ready to prove Thm. 2 corresponding to 1.
F1 Proof of Thm.
Proof. Let Og = (eo, wo) R22 be the initialization for the GF trajectory (0, (0.) Recall that Imin A (e, : (- e e (--((())))))) (w)), g(w) = Vw² log( w) + Exad
u((e,w) : w > 03 ((e, w) w 0}
25

---

Iead= { (e,ww) w E 1-1/v2,0), e e= ==#v=² (og(-w) + Esad
I Imux (e,w) w) : c = 0, w < -1/v2
I. a R (Imin L Isad U Imax).
We consider the cases 0n E e-axis and 0g E RR2 e-axis separately. First recall from Thm. 1 and Eq. (6) that for p + q > 1, the loci of the global minima, e²(1 + 2w(w]) = log (1=))1-0) < 0, lies entirely in the negative minima, maxima and the saddle points span the w axis corresponding to = 0.
pq
corresponding to w < - v2' ~z On the other hand, all the local
(i) 0g IR? \ ee-axis: Let E E(Oo) R. By Lemmas. (9), (10), and (11), we have that the trajectory (0:)220 always stays on the contour line E(e, w) = Eg and converges to the limit Amm 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 0g and the initial energy Eg. This gives rise to the set of initializations Imin: Tmaxx Isad, and I, that correspond to the limit being a local minimum/maxximun, saddle point, or global minimum.
We characterize them individually below starting with Imin.
Initializations for local minima, Imin- For Og (Eo, wo) E R² e-axis, assume that we > 0. Since & E I. there exists an unique w, > such that f(w.)) = EO, w., & by Lemma 12, (iv). Further using the fact that the energy contour lines do not cross each other (by definition of contour line) and the fact they do not intersect the e-axis (it's an energy barrier as discussed in Sec. 4), follows that the contour line E(e, w) Eo stays entirely in the positive half-plane corresponding to w > 0 and w. > O 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
O is set of a local minima (Eq. (7)), it follows that any initialization (Co, Wo) with Wi > O converges to a local minimum.
Now suppose v2 < wi < 0 and eq E (-g(wo), ((w))), where g(wn) = / /22 log(-w) + Eead- Thus leal < g(w)) and hence enj -wg Sooe-loog' w.)) E(e, Wo) = Eg Eaad Hence by Lemma 12, (iii), there is unique intersection of the contour line E(e, w) = Eo with the w-axis, which lies in the
region ( (-13-0) 0). Further note that this contour line cannot intersect with the global minima loci as it lies in the half-plane w < - 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 this line would converge to.
Thus we have shown that any initialization in Imin \U w) w 0} converges to a local minimum, the set of which exhausts all the set of local minima Omin except for the origin. Below we will estbalish that any initialization on e-axis = {(e, w) w = 0} converges to the origin, implying Imin the full set of initializations for which the limit a local minimum.
Initializations for saddle points, Iend- It's straightforward to see that for any On E Iand, oj (w) log(-w)) E(O, -22 Eead Since - 1 S w/ < the point (w,, e) (- 1 0) is the only y2'
V22
intersection of the contour line with the set of critical points, any initialization Izad converges to Egad for WD < V2 that passes through (- 11 0) E IR? and further intersecting with the global minima loci v2 .. 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 de, > 0 and dan, < if en > 0 dt
the saddle point. On the other hand, there also exists contour line eg (wy - llo((ww))
and Wo v such that (wo, ea) is close to the saddle point (- -.23 0) (the case for cp < 0 is similar the flow is symmetric in RR. From Lemma 7 and the definition of the GF, we have that de; ØL
(eo, wo) = 2Ex [(1x + f2) (1 2w6))eo
di de
dw;
aL
dt Ow (co, wo) = 4Ex 1(fi f2) (-cāuu)-
So it suffices to show that Ex X + x2) > 0. To establish this, we have from Lemma that
(- f2 + fz = -To f2-
Ex [(fiX + x = E + f2) = .t
26

---

From the defintion of fz and the optimal bias b, in Lemma 7 and Lemma 5 respectively, we obtain
e
1+ exp (bb + e )
f2 b.-
p
- p P.
2
2
2A
1
A - expleg(1 - 2wz)).
2
P 1 + V (
2 - -1
+
44.2.4
2
When eo = 0, we have A = 1 and hence
f2 (1 y p
p
P
p
P p
(p 4 -1)11
< <0;
(36)
p +
where we used the fact that p > 1. Hence by continuity of fz in cu, for sufficiently close to O, fe < O which proves our claim about the direction of the flow close 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 Wp < V2 Thus Isnd the only set of initializations for convergence a
Initializations for local maxima, Imax. If p + q > 1, we have from Thm. 1 that @max (1 + 2w/w]) 0} (c, w) E RR:
{(c,w) E R
52)- Thus for any 0u
+ max ddr
for all 2 O and hence fimm = 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 Imax
dØ_
maxx
Initializations for the global minima, I., Since the set of all critical points of L is O
min θ max U ② sad, and the initializations in Imin. Isad, and Imax converge O min- O sad. and O max respectively, it follows that the set of initializations for which the GF converges to global minima is I,- R2 (Imin Iaad I_max).
In fact, since the loci of the global minima lies in the half-plane correspondint to w < /2 when p + g > 1, we can precisely determine the location of the global minimum for which the intersection occurs for any 0n E I.. Specifically, we can solve the pair of equations E(c, w) Eo and e² 2 = log (1-p)(1-q) which has a unique solution for w < 0 (upto sign flip in e).
e² ---wd 22 + logoloo-o
Pq
(ii) On c-axis Ou E I If 0. (ea, wo) c-axis, have that wh = 0 and hence w:
for all t 0 (Lemma 9). Lemma 10-(i) also establishes that the iterates (0, (ee, 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 0, ii exists, it follows that O (0, 0), local minima. Thus On E Imin-
This concludes the proof for all the initializations On E TR2.
Gaussian initialization (0, o² 12). When Og is initialized according the standard Gaussian distribution N(O, ((O,,²T2) with x² < 22' we note that 0n lands in the set Imin with high probability. In fact, this probability can be made arbitrarily close to depending on o². Thus this initialization will lead to local minimum convergence on the w-axis.
F.2 Gradient flow dynamics for p + <1
Theorem 8 (GF dynamics for + < 1). Under the same setting in Thm. 2 with p p q and any initialization O, E R2, the GF trajectory always converges to Omm E IR2 which is a critical point of the loss L. More specifically, 0mm is
ii a local minimum if
0g E Immin & ((e.w)):or << e (E (-g(w).((w)) g(w) Vu² log(-w) + Eead
(ii) a saddle point if 0g Isad A (e) -1/V2, e + log(- ww) + Eead
27

---

(iii) local maximum if 0u Imax - ((e,w) : -,
(iv) and a global minimum if 9o E RR \ (Imin UIsad U (mmx))
Consequentely, if we use the standard initialization 0g ~N(0, with 0² 11/V2, Oum will be global minimum.
Proof. The proof for the case of p + < essentially follows the same steps as that p + q > 1. 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 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 carlier case. This is due to the fact that origin turns out to be local maximum when p ++ q < 1 and
hence it's unstable critical point (which can be established as in the proof of Thm. 2 for Imax). D
28

---

G Gradient flow analysis with attention
In this section, analyze the learning dynamics of the transformer parameters 0 R with the attention scalar a E IR, i.e. 0 = (e, w. a) E R Similar to the analysis for 0 = (c, w) E IR?, we first introduce the canonical parameterization including function [()) in terms its gradients and critical points, and capitalize on it to study the gradient flow dynamics using the energy. We first start with the parameterization.
R. then analyze the corresponding loss
G.1 Canonical parameterization with attention
Embedding. Recall from App. that we let e e-a a and pn. -p- a for all where > 0.p and a E {+1]%// /dd This results in the embedding
2/ a
Zn = E Dp
Attention. Similarly, recall from Eq. (21) that the attention output v, is given by
y, e In -
1 _/2 a + (0.
Σ
att e (x:- 22) a.
(37)
2
(E(m)
where
(Eesp /Σ /E|m]
a e exp
((m (i)xVa))V)))
exp (In. (Inn 9, WW Ini k; WK K Zi
W K = (q² a a E Radd, for some E R. Instead of the softmax, now assume that the attention weights are linear in the scaled dot product, ie.
(gn 1
g²d nvd Vd x W W Rmn nvd (x) a)
att,,;
(I ==d) d²dd
nyd (EIn p) p)
(38)
q² 1/2
(eI, - p)x - p)
n
/
1
In
2
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
1
Un =
In
a +
(v
Σ
attmi e
83
2
2
EE [m]
In 2
1
(0,a) a)
Σ
q 5
2
1 ))(z
e
1) 2
a
n
(E|ml
1 (v, (0)
1
2
)
sn
2
1
1 (v, 2) (4 20/2 .2.
In 2
1
In 2/
(1 + ace) (I+ae?) a
29

---

where we used the fact that (x, 2)2 = since 2, {0, 13, and
(v, d
(39)
is the attention scalar. Note that this includes the scaling (v, a) from the value matrix W v and q² from the query-key dot product. Thus we succinctly have
In 5) (1 + aeae) a.
yn
e
(40)
Ferd-forward. For the feed-forward layer, we have that W vi
[ vil
1 a E Riddd
W₂
1 E Rdxid Hence Eq. (40) implies
w 1- a
T e 1a-1 1 2 (1+ + ae²) (1+ae2)
on 2/2 (1 ae²) 1.
W Wiyn
Vd
e
2
Thus,
|w|
1) (1) (1 + ar²))
ReL.U(W 19.)
19)))
ReLU
2) (2(
dd
|w|
(I(2n-
1
1 RoLU
(1 + ac²))
vd
2
[w]
In ReLUUl + ac*)
1 In ReLUU(1- act)
1.e
+
Vd
2
|w|
(In [ReLU(L+ ae²) ReeLU-LL- ae²)] ReLU (-1-ae²))
2vd
Using ReLU(x) ReLU(--ee x above.
|wl
ReLU(W 13,)
1 e (xn (1+ ae") ReLU (-1 aE?)).
2vd
Hence,
W W2ReLU((W (1 11,) a
w
lue|
1. (In (1+ + ReLU (-1 ae?))
Va
2vd
wwwwle (In (1+ad al+) + ReLU
a
on (1+ + (-1-
2 + | + ac?
2w/wle
2
2w/wwle ((m-))) 1 (1 ae?) + 11 ac?
a
Thus the embedding nn is given by
Zn yn
+ WyReLU(( 13/n)
a + 2w 2w/wole ((.--2) (1 + ae²) +
11 + ac²
e In - 2/2 (1 + )e²)
2 a
1
(1 + ae?) (1 + 2w/w!) + w|wi(1 + ac?) a
e In 2
Linear. Since e a e a due weight- tying, the logits are given by
2 (=.. - (1 (1+ae²) ae²) (1 + 2w/w() + w w= + ae?)) })
logit, (e, w, b) (a. knd
Loss. Denote 0 (e, w. a) E R Similar to the case without (Eq. (14) and Lemma 2), the cross-entropy loss in our setting can be compactly written as
L(O,b) N
Σ
ElAoo ((2xn41 - 1) (0, Exy [4gg ((2Y 1) llgit (0, b))],
N
(41)
30

---

+ w|w(1 + ae (X, E {0,1} are distributed according to (x, v) ~ (T,, P). i.e. X a Bernoulli random variable with X ~ π = Bern(p/ (p- + q)) and y| P(p. q). the Markov kernel. Further, using the convexity of b L(- 6). we can consider the optimal bias b. (0) argminseR L(0.b) b) in Eq. (41) to obtain the loss L(0):
where logit x (0,6) c² [[X - 2) (1 + ac²) (1 2w (l) +
L( & L(O,b.) Exx [log ((2Y- - 1) logit (0,b.))]
Ex,y flog (2Y 1) e² [( - 2 (1 + ae²) (1 + 2w/w|) + w|w(( + ae2) + +b,
) (42)
We derive the expression for the optimal bias in the proof of Lemma 13 below in App. N.
G.2 Analysis the loss function L(0) from Eq. (42)
Now establish the gradients of the loss function.
Lemma 13 (Gradient computation and optimal bias). For any 0 = (e, W, a) E R and the next-token prediction loss L00) in Eq. (42), the gradients are given by
0L
=
-
(fiX + f2)
X
2e (1 + aE") (1 2w/w))
de
x x- 1
E
(fiX + fe)
2e a (1 2w/w)),
2
0L
) x 2
= E
(Лi + f2)
2ea (1 + ac?) (lwl + sign (w)w), w),
Ow
OL = E
X ,
(fix + fo)
e (1 + 2w/ww)
0a
2
where X E (0, 1} is Bernoulli random variable with X Bern(p/(p+ 4)), and
f1 - 1
q
01 + do: J2 Φu.
Φ A σ
1
[ + ae (1 + wl) + w) + ae?)) + tb,)
2
-1
Φ o
(1+ ae?) (1+ w + 1 w0 + ac + b,
2
where the optimal bias b. is obtained by solving vı + f2 = 0.
Proof. We defer to App. N.
Theorem 9 (All critical points for linear attention in TRA). Let the input sequence be (x, J. (T, P). the transformer parameters
(e, w, b. a) E IR, and the next-token prediction loss L(.) be as in Eq. (41), Then for any (p. g) E (0, 1)a with + q # 1 and N e N.
the set of all global minima is given by
p(1 q)
(.(p.)) 2 (le, e : +
log q(1 - p)'
(43)
a (1 + ac?) (1 + 2w/ww) log (1 4) p)
(44)
pq
iii set of local minima is given by
T min (p. A 1 { 7min (e, ERt 0, + -- 1)(1 2w/ww) > 0, log
q
31

---

(iii) set of saddle points is
{ (aad = (e, ER :e= 0, (r (x - 1)(1 + 2w/w/) - 0, = log P
ssaa(p: ④) -
(iv) set of stationary points is satatio (p.q) - Vstation we, R e 1.,001 0,1 ac² 0,11+ |0|
p
6 log
4
Thus the set of all critical points is
{0 IR? VL(0) - 03 - UT..Tmi
U.umm min UT sad UT statioon
(45)
In addition, for any 0, 0 E r min and 0 I sad, the loss values satisfy
H( In) = I(0+) < L(fuin)) = m(f L(f = H (In+1).
Theorem 10 (All critical points in RR²). Let the input sequence be {æn IN (IT. P), the transformer parameters 0 = (e, w, a) E R2, and the next-token prediction loss L(.) be as in Eq. (42). Thenffor any (p, q) E (0, 1)2 with p + q # 1 and N E N.
(i) the set of all global minima is given by
O. (p.g) _ {(e, ((e,w,a) E Rd :22 (1+ ac?) (1+ 2w/wwl) log (1 9) p)
(46)
pq
iii) set of local minima is given by min (p. v) A {(e,w, E ERS : : 0, (p+ + 1)(1 + 2w//w/)) > >0,
(iii) a set of local maxima is given by + Omiu (p. q) & {le, (le,waa) E R :
-= (p 1 + w|) < <0},
(iv) set of saddle points is
e (p. (le R (
Defining set of stationary points e station (p. 2) 2 ((e,w,a) E the set of all critical points is
R e 0, + ac² 0,, =w/l = 0}.
{0 E R VL(O) 0] = u min + max U @ U Staation
(47)
In addition, for any 0_ E O., Omin E θ @min and A sad E s the loss values satisfy
(In+1 IT)) = L(O)) L(0min) L(0 (0max) L(Osaa) H (In+1).
Remark 5. While the remaining set of stationary points could be classified to global minima. local
minima, etc., it's technically unclear what category the set of critical points @station (p. q) belong to, as the Hessian is undefined here. We would need to rely 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 0 = (e, w) E RR in App. F, we now study its countepart toegether with the attention scalar, i.e. 0 (e, a) E RR3 To this end, let (0.)rzo be C curve in R3 governed by
d0,
VL(O.) 0, = (et. We a;) t20
(GF-Attentio))
dt
starting with a randomly initalized 0g. We define the energy function -( as
E(e, w. a) A 2 (w² + sign(w) log (w]) 2a² V(e, w. a) E B ea-plane, where ea-plane _ ((e, w = 0,a)) The following lemma presents the crucial result that the energy constant along the flow in GF-Attention
(48)
32

---

Lemma 14 (Constant energy along the flow). For any (p. E (0. and initialization do (eo, Wo. Gg) R3, let (0): be the corresponding GF-Artention trajectory starting from Og- # O, then the energy stays constant along the trajectory, ie.
{)2) ef (wr + sign(w; wy) log [wy]] 2a E(Oo),
VI 2.0.
(49)
On the other hand, if wo = O, Wi for all 2 0. Hence, if we initialize on ea- plane the trajectory always stays on the ca-plane.
Now characterize the convergence of the gradient flow.
Lemma (GF convergence). Let (0t)r20 a continuously diferentiable GF-Attention trajectory starting from Og. Then for all initializations On E R.
(0)) bounded,
i1 there exists Oim E 3R such that lim;-+00 O, Olim and
(iii) lim (L().)Il = [V =
Hence Olim is critical point of T..
The following result characterizes the energy of the limit point.
Lemma 16 (Energy at the limit point). Consider the same setting in Lemma 15. If Oo R3 \ ea-plane, then {(0 E())). Hence Ulim lies at the intersection of the contour line E(,, w) Eo with the of critical points of in R
On the other hand, if 0g ea- cl-plane, then Oim sa-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.
0.5
0.
Local Minima
Local Minima
(a) Gradient flw around origin (p + 4
(b) Grudient Bow around origin (p + 1)
Figure 5: Gradient flow dynamics in R3, near the origin, for the transformer parameters with attention
scalar a (Sec. 4.1). The local minima are repellors for p + < 1, while attracting for p + > 1.
Based on the above theoretical results and empirical evidence, we conjecture the following theorem, whose informal proof we defer to App. N.5.
Theorem 11 ( [Informal] Role of standard initialization for > 1). If we use the standard initialization 0g ~ (0, 1²13) with o² & f/ the GF-Attention, Ori will be a local minimum with high probability.
33

---

H
Additional empirical results
H.1
Gaussian initialization converges low-rank
0.2
0.10
0.04
5
5
0.1
0.05
10
0.02
10
10
15
0.00
15
0.00
15
0.0
20
20
20
-000
-0.05
-0.1
25
25
25
-0.04
-0.10
-0.2
30
30
0
5
(a) Wi at initialization
(b) W. after 50 iterations
(c) Wi at convergence
0.02
0.00
0.02
0.04
0.10
0.05
0.00
0.05
0.10
0.1
0.0
0.1
3
3
5
5
(d) initialization
Wv after 50 iterations
(1 W convergence
Figure 6: Evolution of parameters W1 and Wy across iterations, starting from standard Gaussian initialization. At convergence, all the parameter matrices are approximately rank-one.
H.2 Low-rank initialization stays low-rank
34

---

0
0.08
0.0015
0.10
0.06
5
0.0010
0.05
10
10
0.0005
0.02
15
0.0000
0.00
20
-0,0005
20
0.02
20
-0.05
-0.04
000010
25
-0.06
-0.10
30
-000015
30
30
0.08
5
(a) W, at initialization
(b) W after 50 iterations
(c) Wi at convergence
0.075 0.050 0250.000 0.025 0.050 0.075
-0.06
-0.04 -0.02 0.00 0.02 0.04 0.06
-0.075-0.0500
0.025
1
3
3
4-
5-
6-
(d) W. at initialization
(c) Wv after 50 itcrations
(f) W at convergence
Figure 7: Evolution of parameters W and W across iterations, starting from rank-one initializa- tion. The parameters maintain rank-one structure across the entire training.
35

---

I Model architecture and
Table 1: Parameters in the transformer architecture with their shape.
Parameter
Matrix shape
transformer.wte
2 x d N x d d 1 3d x d d d dx 1 4d x d d 4d d dxl
transformer.wpe
transformer.h.n.
transformer.h.atii
trnsfoomerh..ttm..ccm
munsfomer.h.in.2.
transfommerb
munsformer...mp...mpo
transformer.n
Table 2: Settings and parameters for the transformer model used in the experiments.
Dataset
k-th order binary Markov source
Architecture
Based on the GPT-2 architecture as implemented in [26]
Batch size Accumulation steps
Grid-scarched (16, 50}
Optimizer
AdamW (31 = 0.9, 82 = 0.95)
Learning rate
0.001
Scheduler
Cosine
# Iterations Weight decay
8000
x 10-3
Dropout
0
Sequence length
Grid-searched in (512, 1024, 2048)
Embedding dimension
Grid-searched in (4, 8, 16, 32, 64)
Transformer layers
Attention heads
Repetitions
3 or
36

---

J Proofs of theorems in App. D
J.1 Proof of Thm.
Proof. We characterize the set of global minima, local minima, and that of the saddle points individu- ally.
(i) Set of all global minima. Let 7 E IR& be arbitrary. From [23, Lemma 1J, we have that Y, is a global minimum for the loss L(.) in Eq. (23) if and only if its prediction probability satisfies J.. (a+) IP (In+1 = In), the Markov kernel. Since the input (sn)m=r (T(P, q), ((p 9)),
we have that
PP (In+1 = 1 xx,) = (1 ra)P + n (1 q) = (1-p-- 11-pp- (1-p-9) -m + p.
(50)
On the other hand, by definition, from Eq. (3), Sr, (x+) (c= + 2nc (xll + b -
where
Y. = (e, w. b). Since E {0, 1}, this can be further simplified to
²
[x,(57) (1 + 2w/w() In +
e²
6² rn). b
+ c²2 + 2w/ww)) In
(1
(51)
6² b
e²
6²
In
2e wwl +
+
2
2
Since both fr, (xf) and P (#n+1 xx,)) are linear functions of Tn, equating them for all vallues of {0, 1] implies that the respective coeffecients in these functions in Eq. (50) and Eq. (51) are also equal, i.e.
2
= p.
2
e²
2
26 w|wl + bb
-
= - p 9.
2
2
and hence
c² 6
P, ((262ulul) + b
= ) =1- q.
E2 //
(52)
Since (12) y for y 10, implies log -ny Eq. (52) can be rewritten
b c²
log
p 1-p"
2e- +
e² = -log
1 -
2
q
Using 2c u/ w| + b + = e1 + 2w w
+ b
c² = = + 2w/w/) + log _ in the second
equality above, we obtain
b 2
p
log
1- p 1
(53)
(1 2w/w|) = log - q + log
1 1-p
log (1-
p) (1-p)(1 g)
q
pq
Thus 7 E IR3 is global minimum for L(-) if and only it satisfies Eq. (53) (note that it's already
critical point, as established in Thm. 6). Thus, the set of all global minimum r. (p, q) is given by
(p( & 2 {7- 1 (e, b) e R² ea + 2w(w]) = log (1 p) - 9)
2 = log
pg
2 1 1
Since the prediction f-. (.) equals the Markov kernel for any 7, T., follows from Thm. 4 (or [23, Lemma 1J) that L(,.) H (0n+1 xa), the entropy rate of the Markov chain.
37

---

(ii) Set of local minima and saddle points.
Define I Fimn (p, q) C R3 and I Iadd C IR as follows:
T min (p, _
Mmi (e,,,,) E R e 0, (p -
2w//w)) > 0, b log 1)(1 (1
saad (p, q)
A Yead (e, (e, E E e 0, (p + + 1)(1 + 2w/w/) - 0, - log
P q
To show that F min is the of all bad local minima for L(.). first show that any Mmin E T 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 Ymin = (e, w. b) E P min- Recall that Ymin is a stationary point (Thm. 6), i..
Rearragning the order of scalars and writing Yiin (b, e, w). from Lemma 4, the Hessian of the loss at Y is
I
0
0
A L2 = क
2(p + q 1)(1 + 2w 2w(w) wl)
0
(54)
0
0
0
By definition, Ymin (b, w) satisfies (p +
1)(1 + 2w/ (wl) 0. Thus its Hessian in Eq. (54)
H where H has both the eigen values positive,
has a block diagonal structure of the form
0
and hence positive-definite. In other words, Vmin is a local minimum for L(-) in the (b, e) E R² space for any fixed in the set. Interestingly, using the continuity argument and the fact that L = log q² f E = 0, w) is constant in w E E, we can essentially follow the same steps as in proof of Theorem 2 in [23, Appendix B.3] (Thm. 5 above) to show that Y (b, e, w) is also local minimum for L(-) in the full parameter space IR3, This establishes that mii is a local minimum for L(-).
For the saddle points, let Yaad = (e, w. b) E T sad- We have that 7 a stationary point (Thm. 6) and Lemma 4 implies its Hessian (after rearraging the order of scalars as above with Yaad = (b, e, w)) is:
0 -))(1 + 2w/we)
L To
0 2(p +
(55)
0
0
0
If # "22' (p 11-+q-171 1)(1 2wrlwe) < 0 for any Yead T sad, and hence the Hessian V ( sad
in Eq. (55) as both positive, negative, and zero eigen values. Thus Ypad is a saddle point for L-.).
to establish that it's also Using a neighborhood argument, we can similarly argue for /2 saddle point. Now we compute the loss value.
For any Ymin (e, w. b) E Tmin or sad = (e, w. b) E I aad, we have that 0 and b = log 4 Thus for 7 miin or I sead. the prediction probability in view of Eq. (3) is
e² (b))
P. ( = 1 x
e2 2w(ww) x. b
P
P ( 1),
2
p+c
the marginal distribution. Substituting this equality in the definition of cross-entropy loss L(-) in Eq. (1) and the fact that IP (In+1 = =1) = ptg = Ti, following the same steps as in [23, Appendix B.3], we obtain
L(r)
1
Σ
Ege+i [2n+1 log fz(2") + (1 In41) log(1 (x+)))
N
nE(N]
Σ
E [Tx+1 log γ + (1 - xn+1) log mol
N
REIN]
38

---

1-m T1 log m1 - To log To]
N
MEIN]
H(T) - H (xn+1).
Thus m(mmi)) L(Ysaa) H (In+1)) To see that (Ix+1 xx)) = L(r,)) < L(7min)
L(Ysud) = H (In+1) ) for any global minimum Y observe that the gap
[(7min L. (xn+1)- (1 (20)) = [Zns Fn+1) 0,
where [xn; 1n+1) is the mutual information between In and In+1 [1I]. Hence the optimality gap equals zero if and only if the mutual information equals zero, which happens when In and In+1 are independent, i.c. P (In+1 = 1 zn,) doesn't depend on 2n. Since P (In+1 = zx) =
(L-P-o0 q) + p from Eq. (50), this happens only when p + 9= which contradicts the theorem assumption that + q # 1. Hence H (x,+1
(n.) = L(7,) < ((7 = L(7 = H +1)
Now we finally show that mimi and sad are the only set of bad local minima and saddle points respectively. Let y is a bad local minimum for L(-). By definition, it's also a critical point. Recall from Thm. 6 that any stationary point 7 (c, w. b) for the loss L(-) satisfies that either 7 E I E T min. or y E T sad. Clearly 7 E T., as r. is the set of all global minima. Similarly, 7 $ I sad as every point in I sad is saddle point for the loss L(-) as established above. Hence y E r min- Thus every bad local minimum in Rs3 belongs to I min- This coupled with the fact above that I min is a set of bad local minima implies min is indeed the set of all bad local minima. The proof for I sad is similar.
J.2 Proof of Thm. 6
Proof. Let (e, w. b) E R be such that VL(7))
[0L aL AL T (de
-.. By Lemma 3, we have
Sa 85
0L
de Ex [IfiX + f2X (w/w|) - 1))] 0,
(56)
0L Ex [[iX + +2 46" |w| = 0,
(57)
du
OL
Ex [fiX + f2]= =0
(58)
0b
where X Bern(p/(p + 4)), fi (262 + b + c2 5 + q 11 1-0 (b-
02) ) p. and f2
b- a²
2 p. Our goal is now show that Eqs. (58) hold only if either (e 0, b log 2)
or (fi 0, f2 0). We consider two cases corresponding to 0 and # O.
(i): e = O. If e = O. we readily see that OL al Hence, Eq. (58) implies that
= 0. Further, fi = p + q 1 and f2 = =bb)- pp
Be
de
Ex [AX + fz) (p + 1)E[X] + (b) - p- = (p q 1)
P + a (b) p- o(b)
0,
p + q
p p 9
which implics that 6 = log p. Since w E R is arbitrary, we see in this case that
E E (e, E : :e= 0, 1)(1 + 2w w|) > 0, b log 1
(G..,))ER' cc= 0, + - 1)(1 2w/w/) < 0, b log p
9
= =I min UT sad.
(ii): e + O. Suppose e * 0. Here we show that f1 = f2 = and hence y E r., We consider two
cases corresponding to w + 0 and w O. Let w of 0. Since both * 0 and w of 0, and X = =x² in Eq. (57) with [[[]] ग of O, we obtain that
Ex + f2x] = (fi + 1 = (fi + 12 = O.
39

---

and hence fi + f2 = 0. Further Eq. (58) implies that fi + f2 = 0. Together, this implies
fi f2 =.
Now suppose w = 0. From Eq. (58), we have that Æ + [2] fi + f2 = 0. Since e # 0, Eq. (56) yields
Ex [ f2X 1))1 - x1xx [( + f2] = 2 (fi + f2) - 0.
So f1 - 0 in this too. Thus have showed that whenever o O. have S1 O. Recalling the expressions for fi and f2.
fe =
2 -pp= -
e²
e²
= log
p
(59)
2
1 --9
e²)
c²
e² 22 we + +
fi
22 2e/wlel + +b+ +q 1- 0
6
+p
+ - 1=0,
2
2
2
and hence,
2
1 -
2e 2e²w|uu |w| | + b+
= log
2
Substituting
2² =log P in the above equation, 1-p
6²
we2 + t² +b
=2e?wlul + e² + log
log
2
q
and thus,
e² (1 + (w/w() = log (1 - p)(1 - 9)
(60)
Pq
In view of Eq. (59) and Eq. (60), we have that 7 (e, w. r..
Together, we have shown that whenever VI(r) = 0, we have Y E r, U U r sad- Since T, UT min U ssad (y: VI(()) 0}, done.
K Proofs of technical lemmas in App. D
K.1 Proof of Lemma
Proof. Recall from Eq. (14) that the cross-entropy loss L(-) is defined as
L(>) = N nE(N]
Σ E. [[nn11 log fx + (1 - In+1) log(1 fo (xf))),
(61)
where [x(x)) (logit,,)) a (c2 2ur w/) xn + - from Eq. (3). For any Y
(0, 11,2, Z E R, using the fact that o(Z) o--Z), we have
lYg (2) + (1 - Y) log(1 (2))] |Ylg o (2) + Y) loga
(a)
Y log(1 + exp( -Z) + Y) log(1 + (Z))
log (1 + exp(-((21 - 1)2))
bog ((2Y-1)) 1)2),
(62)
where (a) and (b) follow from the definitions of sigmoid and the logistic functions: o(2) =
I+exp(-3) -3) loog (z) = log(1 + exp(-2)) for = E R. Substituting Y In+1 E {0, 1) and
= logil, E I in Eq. (62), Eq. (61) simplifies to
I(>)
1 N
Σ NE[M]
Elliog ((20n41 - 1). logit,)))
40

---

Since logit, n is only a function of In, the above expectation is over the distribution of the pairs (In: In which for all E [V] have the same law as a pair of random variables (X, Y) with X ~ = Bern(p/(p + q)) and Y|X ~ P(p, 4), the Markov kernel. Hence the above equality can be rewritten using the definition of logit, as
B[fiog ((2mn+1 - 1). logit,,)] = Ex,, blog (2Y -1) 1) (e((1+2uulu)) 2w/wo)) x + - e²
2 )
I(>) = N
Σ [E(N]
K.2 Proof of Lemma
Proof. With y (e, w, w) and 0 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
Vol(7) = N
Σ NN
Bgg+ [(In+1 - fo(x))) Vo logit,),
(63)
where fz(x)) (logit,,) (c2 + 2w/ww) x,, + : 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) Y) with X π Bern(p/ (p q)) and YIX- P(p,q), the Markov kernel. That is,
VoL(7) -Ex, [x- (00 2w/wl) +
e²) )
Vo (e w//w))x x + 6
(64)
Now we define the error term E(X,Y) - - (r-o (=² 22 )X + -
-----)) ) Our goal is to
show that (X, Y) X] = fiX f2, where fi a a 22e +b+ + 1 1-o 6 ) +p.
and f2 A o 6 ) X E (0, 1}, we have
ps which suffices to prove the lemma. To this end, using the fact that
Y - σ ( (2 + 2w w/ I) +b -
E²)
E(X,Y) =
2
6²
e² b
Y - X 202 +
-(1
(1-X) X) 0
--Y+X +X x (o 2e2 +
2² 0 b
2
2²
b-
2
Now taking the conditional expectation with respect to X and using the fact that E
P(Y = 1 X) = (1 9)xX + p (since Y ~ P(p, g)), we have
0 1202wlux + 6²
e²
6²
B[E(X, Y) X - p - 1 - p X
b
+
b
2
x o 26 2w2 b+ e²)
c
+ 1
b
+p
p
where (a) follows from the definition of fi and f2 above. Thus Eq. (64) simplifies to
Val(>) = Ex [ f) va (o(1+ + (10+2ww/)))) 2 X + b 2
- }
Letting 0 = e, w. and in the above equation, we finally obtain the individual gradients:
ØL
Ex I(fi (2X + 2w/wl| 1))] e,
de
41

---

0L Ow
Ex [fiX + [2) /48/
OL
= Ex 1fi + (2)-
0b
K.3
Proof Lemma
Proof. Slightly changing the variable order, for any = (b, w) E R., we define
8²4
8 0bde
0 L
86
86ow
H(>) A V2
8 Bedb
8 de
8 L Uedu
E 3x3
(65)
L Dwab
8 OwDe
8 L Ow²
Recall that for any Yiin miin and Yad aadd we have O and b b= log P. Now compute
q'
the second derivatives of L with respect to any 7 (b log P e 0, w). We start with the first derivatives. By Lemma the gradients are
q:
OL
Ex [/iX 12).
0b
OL
Ex 1(fx + f2)) 2w///) - 1))]
(66)
de
AL
1xf + f2) 4e2
Dw
where
e²
e²
fi
222 + +b+
+ 1 2
+ p.
2
c²
f2 4
b
p.
From Eq. (66), we see that the second derivaties of L depend on the first-derivatives of fi and f2, which compute. Recall that the derivative of the sigmoid function obeys d'(2) aI=)(1 -=))) = (=)o(==)) for any = E IR. Now the gradients of fi and f2 with respect to b, e,
and are
0f
2² 2e 2e + b
o -2e2 (e3w|w - - o²) 2 -
e² 2
o 6 e² )
o
0b
0fz = a
-b b+ 2
b
0b
2
2
0f
= (4ew|w) + +ejo
222w|ul + bb+
22
(2202uluol - b -
e2
tea (b
b - _
2)
o
-b _
de
2
2
2
0f2 (-e)a b 2²
-b 2
de
2
05 (4e² wsign(w))
2
))
(-2-0u2lll b e²
2e 2e2 + b +
dw
2
05 O.
dw
42

---

Using the fact that log
_ mi and log
_
xo. the above gradients
q) p+.
p+o
evaluated for any 7 (b log 2,
q² 0, () further reduce
0f
0f2
O,
T
0b
0b
0,
(67)
de
de 7
7
0f
f
0,
0.
Ow
dw
Now substituting Eq. (67) when computing the second-derivatives of L in Eq. (66), we obtain
L
051
0f2
062
Ex
0b
X
0b
To
0f1
0f
Ex
X
0,
0bde
de
de
7
LL
0f1
0f
Ex
X
0,
Obdw
dw
Ow
P...
Ex [(fx (2X + 2w/ww) 1))]
Dea
7
(68)
= Ex [(2/1 + 2w//vl)) f1 + 2f2(1 + 2wr/wr) X - [2]
Ex [1f1 4w/w/) f2(2 4w/w/))x fel
= (fi - 4w/w/) (2 + 4w/w))) ff
(4)
+ /o) m (p 1) + 4wllw/)) + + + 4w0 w|) 1)
1)
11
2T 2m + q 1) + 2w/ww),
43

---

where (a) follows from the fact that fil, = -- -
--
(p + q 1) and (b) from 1 Ao- Returning the remaining second derivatives,
T. θ
(El(fiX f2)x]] 4ea (wol)
dedw
de
0 (BI(/1 f2)x] 4x01
de
0
((fi fe) 4TTe?|Wll
de
0
((
c²
_mmee/wll) |
2e Re?w/wi [wl +
+ +9-1)
de
2 we + bb+ -
4 e )
(69)
de
b+ + (0 2E w/www/+ +b+
e²²
_
+ - 1 )
0
(4 e²
( (0))
2
de
0,
L
(Bl(fiX f2)x] 4601w))
Ow2
dw
B[([XX + [2x]
4|w1l 6²
aw
0.
Congregating all the second derivatives from Eq. (68) and Eq. (69) into the Hessian H(y) in Eq. (65), we finally obtain
H (>) To
0 2(p g 1)(1 2w w
0
44

---

L Proofs of lemmas in App.
L.1 Proof of Lemma
Proof. Recall from Lemma that for any 0 = (e, w) E IR² and E R, we have
200 )XX + 2 ) L(O,b) b) Ex, flog (2X-1) - ( (1012u//
<
Since (+) is a convex function, Y E {0. 1} and thus 2Y E {+1}, the convexity of in follows from the following fact:
LL 06
Ex,
&%og
(2Y 11)
(1) <1 200 ) + -
2 2 0.
To find the optimal b,, we set the gradient 0L 06
0. Thus from Lemma 3, we obtain
ab = Ex [/i +
[2] = 0,
e² e²) + 99-1-1 (b
e² )
+ p. f2 σ b E²
f1 = 2e 222wwuu w//ll/ ul + b
- p.
Substituting E' 2 22 2w/ w|))
B 2 2. and E[x] pp + 4) in the above equations,
oEE + B) + o(B)+ p(P + 1) o p π (3
T1 B)
p+q q p 0.
Further simplifying.
To oBB) =m (1 (1- o(E - BB) का a(-E - B).
In other words,
(1 + exp(-B))
T1 _ 1 exp(E) exp cxp B) P
(1 + exp(E + B))
To
1 1 exp(-B) q
Defining x & exp(B) (B) and A a exp(E), E), we thus obtain the following quadratic equation in x and its corresponding roots:
Ax² P (1-1) -
1 - +
+ 2
2A
Since x 0, we take the root corresponding to the addition choice above and resubstituting
x cxp(b )) and A = cxx((e2 + 2w(ww)), we obtain the final expression for b,. In particular, e it easy to see that A and hence exp(b.) 2. implying b. log 2 q
G. Similarly, A (1-p)(1-)) it's straightforward to see that exp(b. e² e2/2) i-p and hence pq
b. e 22- log -
L.2 Proof of Lemma 6
Proof. The proof directly follows from Lemma 2 by substituting = b..
L.3 Proof of Lemma 7
Proof. By Danskin's 's theorem [12], it follows that for b, = argming L(0, b) and L(0) = L(0,b.) b.). have (O) VoL(O, b.). Using the gradient expressions of L(O, b) w.r.t from Lemma
and using the fact that at E [f 1 X + fa] at b., the claim follows.
L.4 Proof of Lemma 8
Proof. Since L( L(O,b.) where b, argminneR L(O, b), the identity in Eq. (32) about the Hessian of the loss L with respect to 0 follows from the classical result of [29, Lemma 2.2] about grrond-derivatived of vxtremal-value functions. Finally Eq. (33) follows from substituting the full Hessian in 133x from Lemma 4 in this identity.
45

---

M Proofs of lemmas in App. F
M.1 Proof of Lemma
Denote (et, wx) = (e, w) with the dependence on time implicitly assumed. Then by the definition of GF and the gradient expressions Lemma we have that
Proof.
de OL (0.) 2 f2)X] (1 2w/wwDe,
(70)
dt de
dw ØL
(0:) - + f2) 4e ww|
(71)
df Dw
Dividing Eq. (70) by (1 + 2w w ) and
Eq. (71) by 4e |w], we have
de 1 dw
de
dw
1 2w/ua dt Noting that at (sign(w) log (wl)
2e/wl dt
dt
1o
2
dt
dw Tw di
the above equation can be rewritten as
d
(e² w² sign(w) log (w]) 0.
dl
Thus defining {ee w) = e² w² sign ww) log |w| for w * 0, the above equation implies E(O,) ((u) for 0, (eo, wo) with wy # 0. On the other hand, it's easy to that if wo 0. Eq. (71) implies du de 0 at 0 and hence wy 0 for all t 2 0.
M.2 Proof of Lemma 10
Proof. To prove the convergence of the trajectory (0:)120. we use the classical result due to Lo- jasiewicz [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 the trajectory.
(i) (0,) is bounded.
We consider the cases 0g e-axis and 0, R22 e-axis separately.
Let's suppose 0n = (eg, wo) c-axis, i.e. wo = Thus it follows from Lemma 9 that w, = 0 for all t 2 0. That is, the trajectory always stays the e-axis and it suffices track (et)r20 and show that they bounded. To this end, we show that if eo > 0, we have de 0 and if eg O, have de > 0 for all 0, which establishes our claim. We have from the GF and Lemma that de
de, AL (ct, Wi = 0) = 2Bx [fix + f2)x e
(f + f2)er.
(72)
di de
and
fa
To f2
c
f1 +
fe
+ fo
(b.);
p
π
π
2
To
1 exp
((bb): +
2 ))
p
To
1+
2x.
- P
Ft. 2 exp(e?) (73)
mi
+ ( -
2 (1-1) 1)
+ 4
4
Zt
Defining
2r
g(x)) A
(74)
1 + v ( (1-1)
1)
+
46

---

and substituting Eq. (74) and Eq. (73) in Eq. (72), we obtain
de: 2
1
P) Et.
(75)
dt
1 g(xx)
Since 2, exp(e?) exp(-e?) e)), in view of Eq. (75), with out loss of generality,
assume that eo 0 and show that des < 0 for all 2 0. That is, the RHS Eq. (75) is negative. Note that
exp(ei) and g(x;) > 0 since the denominator / - --1+1 v ( - + -4.5.00 I
1)
Ie 22 since - 1+ P + 2 q increasing in x for x 2 1, it implies I+g(2] p < 1+9(1)
O. Further g()) p and gim g(x) 30. If show that g(x) is p (p + q 1) < 0.
1+3
ptg
Thus the gradient in Eq. (75) remains negative starting at = 0 and hence the sequence (e))120 will be bounded. Now it remains to show g(-) is increasing, i.c. g'() > 0. Defining C
2 7 2 9
and D = C², we have that g(Ø) upto postive scaling is
g(x) C+V
Hence
C+V++ D 2N++D
g'(x) (C+ Vx+D)²
Thus it suffices to show that hy (x) A C + VI VD+ > 12(x)) 2 bi (1) h22(1) is given by
x
for x 2 1. Note that
2V++D
p
P 1
1
h1 - hz(1)
9
+
1 +
2
p
2
2
14
2
p
V q
0.
Now we show that h²(x)= > (x( for all 2 which implies that h₂ (x)) > h2(z) for all x 2 1, thus
establishing our claim. To this end, we have that
Vx + D
(x) 11(())==)-)))))))))())))))))11 W/+D
2V
2(x D) I
2/4 D(x+ >0. 0.
This proves our claim that g()) is increasing and hence (et)12o, and consequently (0:) is bounded when 0 E c-axis.
Now let's assume that O, (eo, wo) R2 \ e-axis. Since (0:)r20 C R2 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;- ||0,]| exists. Now we show that lim |00 A co, which implies the desired result about boundedness. To show lim ||0,|l * cc, we show that there exists de
a large B > such that for any 0, (e, w) E IR² with [[(e, w)ll 2 B the velocity vector är 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 (et, = (e, w) with the dependence on time implicitly assumed. Then by the definition of GF and the gradient expressions Lemma
have that
de aL (0)) -2 X + f2) (1 + 2w/ww]e
(76)
dt De
dw aL
(.)) -E((XX f2) 4e0
(77)
dt Ow
47

---

where (2e2w) f| ue + +b, +
+ + q------- 100----bb--- (b.-5) p. and f2 b. - -)
p with
T fi + f2 = O. Given that only de flips in sign under the transformation (e. w) + (-e, w). with out loss of generality we can assume
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[ + f2) X] > 0 for reasonably large B with (e, w)ll B. Similar to Eq. (73) and Eq. (74) above, using the relation Ti + = O, we obtain
El - = FI f2) = - ( -P), x _ exp(e" ( + 2w/w=)) (78)
Using the fact that g(x) is increasing for x 2 with lim g(x) = oo, and |u| = w > 0, we can chose B >> 0 such that for any (e, w)ll 2 B. in Eq. (78) we have 1/(1 g(x)) p and hence
E[ + f2) > 0. This finishes the proof of our claim. The proof for w < is similar, where we to show E((f + f2) < 0 for e, reasonably large.
make of the fact that lim (
(ii) lim, 0: 1im. Since the logistic function (.) is analytic, follows from Lemma 6 that the loss T.) analytic too whenever w of O. On the other hand, when ar 0, it's easy to see that I is an analytic function of e E R. By Lemma 9, we know that if o 0. ut of 0 and if wo = O, wr = 0 for all t 2 O. Thus the loss is analytic on the trajectory for all 2 O. Since the trajectory is bounded, it follows from Lojasiewicz's theorem [1, Theorem 2.2] that there exists Olim E R² such that li 0, lim-
(iii) oim VL = VL ml = 0. Since the trajectory is bounded, it follows from [2, Theorem 2] that the gradient converges zero, i... lim;-+oo VL(0,)] 0. Since VL(-) is continuous function and lim;-+00 0, = Onm. we have lim (L))) = [ (0n))] =
M.3 Proof of Lemma 11
Proof. Since the energy function (., in Eq. (48) is continuous function in IR2 e-axis, and any
trajectory (0.) with intialization 0 E TR22 e-axis stays in RR2 e-axis for all 2 0 (Lemma 9), follows that lim E(0,) E fi E(o)). As V V ( jim) 0 from Lemma 10, it follows that
Olim lies at the intersection of the contour line E(e, w) & with the set of critical points of L in R².
On the other hand, if 0g E e-axis, have 0, e-axis from Lemma 9 for all > 0. Hence Pim E c-axis.
M.4 Proof of Lemma 12
Proof. Recall that f IR [0] - R, defined f(w) E(e = w, -ww² + sign(w) log (w]). < 0, we have
f(w) = (w² -lolool- f'(w) = -2 +
Hence f'(w) 2 for w (-o0, -1/v2 and (w( S for w E [-1//2, 0) with f'( v2l It's also straightforward to see that lim (w) = -00, lim (ww) -00, and f(-1/v2) Esad (by the definition of f). This establishes (i), (ii), and (iii).
On the other hand, for is monotonically decreasing for w > O with lim ((w) oo and lim f(w) -o. Note that w acts as an energy barrier since lim f(w) oc whereas lim fww co. O
have f(w)= (w² + log w) and f'(w) = -2w 1/w). Hence
48

---

N Proofs of lemmas in App. G
N.1 Proof of Lemma 13
Proof. First we recall the loss with the bias L(O,b) from Eq. (41):
[(0,b) Ex,y 1nng ((2Y 1) logit (0. b))]
where (ogitx(0.b)) = =e [X - 2) (1 - ae²) (1 + / + w|w([ + ae²) }] + b with o (e, w. a). Using the fact that llog a(2) 11 and 2Y - E {+1], we have for any 0 E {e, w. a, bj
that
Vol. = B[ 1) loggi 1) (2Y 1)) Vologit E ((o(git) -)) Vologitx] (79)
Now we simplify o(logitx) ) using the fact that X E 13:
è [xxx (x-1) 1) (1 + ac²) (1 + 2w|w() + w|w(1 + ac?) | +b]
(1)gitx)= = a
2
[1
= =ool e²
(1 + ac²) (1 + 2w] w)) + w [0 + ac²)| +
2
Adi
-1
+ (1 X)o
c²
(1 + ae") (1 + (w|w() + w|w(1 + ae +6
2
=xo + - dx
X(Ø1 do) + du-
Thus the gradients in Eq. (79) are given by
V&L -B [Y X(1 - Φ0) do) Velogitxl
Ex E p [E]Y | x] X(o) Φ0) - -b)) V Vologit x
Exx [[(1 - p 9) X + X(01 6o) hb) V Vologitxi
(80)
1-p 1- q Φ1 + S
-Ex
X de
Vologit x
fa
Exx [Vix + f2) Vologit
Now we compute the individual gradients with respect to e, w, and b. Recall that
logitx = 62 =2 (x-1) (1 + ae²) (1 + (w/w/) + wlw[[1 ac²)|
+ b.
Thus,
V. Velogitx x = 2e
(x
1 (x-1]
2
(1 + ae?) (1 w + w/ w0(1 + ue?)
1
1
+ 3
2
X
(1 + ww/www + wsign (w (1 + )e?)) (2ae)
2
1
2e
x
(1 + ae?) (1 + ww/w/) + 2e (1 + ae?)
2
1
+ 20
x 2
(1 (1+2w|wl) + 20 awsign (w (1 + ae?))
2e
[ - 2
1
(1 + ae?) (1+ 2w / + 2 x (x-1 2 (1 + 2w 2w/w])
+ 2ew 2ew/w(1 + de²)| + 2 awsign (w (1 + ae²))
49

---

Substituting the above equation in Eq. (80), we obtain
(fix + f2) X (x-1))) 2e (1 + +c²²)) (1 + 2w w|)
V.L
EF
E (fiX + fe) xx 2) (x-2)) 2e (1 + 2w(wl)
(E + -2)l)) 2ew (1 + aE2)l
(E + f2ll)= 2e aw sign (w (1+ (1+ae))) ac"))
Now we compute the derivative with respect w.
logit x = c²
[
x
(1 + ae?) (1 + 2w/w|) w|w(1 + ac?) ac?) b
2
X 1
V logiit 20²
(1+ ac*) (1wo + sign (w) w)
2)
+ e ||uo(1
+ ae²)| + (1 + ae²) sign (w (1 + ac²))]
[B (f1x + f2)
X xx xx=b)) 20 (1 ae*) (lwi sign (w) w)
V.l. =
- [B[[iX + fa)l) e² []w([ + ae²) ) + w (1 + ae²) sign (w (1 + ae 2 ))
Similarly, for a:
logitx = 62 [ (x-
(1 (1+33 20/ww/ + wlw + ae 2 +
2
Valogitx = e
(xx (x-
(x-2)
2) (1 + ae?) (1 + 2w/w/)
1
+ e (w(1 + )))))
VaL = EE (fiX + f2) x - 1) e' (1 + +wwwwwl
[[[ + fe)l) e sunigg (w(1 + ac²)).
Finally, since Vologitx = 1, follows from Eq. (80) that
-LL=BIAA> + fel.
(81)
For the optimal 6. we have VoL= 0 and hence EE [fixX + fa] = 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 w. and a. Recall from Eq. (80) that
fi (2))
(222) +p+ q
- f2 = (22)
- P.
21 2
(1 + ac
(1+ + 2w/ww) + w|w(1 + ac²) +b,
(82)
2
[1
22 A e
(1 + ae²) (1 + 2w w|) + w w0 + ac²)
+
2
Substituting Eq. (82) in Eq. (81) and setting the gradient to zero there, have
-Ex [fiX + +2] = (o(2)) - (())))-oaa)+ + - p
= (a(2)) - + + + - p 0.
50

---

Simplifying,
(0(21) - (22)) + 1)m (22)
(a(21) 1) m +p p (=2))
(a(z)) - 1) T = (22) - 1)
(22))
" p
( - m
q
Using the definition of the sigmoid function and rearranging,
1 1 exp(2))
p
+ exp(-22)
q
= exp(21) +
p (1 + exp( -22))
q
(83)
* exp(221) + exp(21)
P exp(21) + P exp(21 22)
q
+ (exp(z1))² + exp(xi)(1 - [ L exp(22 22) = 0. q
By definitions of 21 and 22 in Eq. (82), we have 21 22 e1 + ae²) + 2w/w|) and thus ف exp(21 - 22) exp(e?(1 (1 + ae?) + 2w ())w)) simplifies to
Thus the quadratic equation in Eq. (83)
(exp(z1))² + exp(21)(1 P)
p A 0. q
On solving the quadratic equation for exp(z1):
exp(21)
p
+
V
(2.
P
2
+
P A
9
1
P
( (2-1)
2
+4. P ..P.A.
21 log
1
V
A
2
9
q
1
P
[--1
P
2
P A
b. = log
1
V
+4 4.
2
9
q
1
è²
(1 + ac²) (1 2n wl) + w ww(1 ac²)
Note that when above, we recover the expression for the optimal bias in Lemma 5. This concludes the proof.
N.2 Proofs Thm. 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 7 = (e, w. b, a) E R with VL(r) = 0, the bias b is already the optimal This is similar to the proof of Thm. 1, which follows from Thm. 7.
corresponding to 0 (e, w, E R3 i.e. b = (0) argmingR L(O, b).
Now prove Thm. 9.
51

---

Proof. We characterize the set of global minima first.
Set of all global minima. Let 7 IRA be arbitrary. From [23, Lemma 1J, have that 7 is global minimum for the loss I(.) in Eq. (23) if and only if its prediction probability satisfies f, (aT) = IP (xx+1 = 111/xn) In), the Markov kernel. Since the input {Ia]n=z (T(P, P(p, 9)).
have that
PP ((en+1= = 1 x..) = (1 Fa.p + Fn - 9) = (1 -p-pp -)5 + +p.
(84)
the
other hand, by definition, from Eq. (3), (x7)
a (e² [I 2) (1+ + ae²) (1 2w 2w/w() + w/ colw(1 + ae))[] + b), where 7+ = (e, w. b, a). Since
In E {0, 1), this can be further simplified to
Jr. (27)
c²
In
2 (2-22 (1 + ae²) (1 2w/w/) + w|w(( + aee>| +b)
1
(e2
[
(1 aE?) (1 2w/w)) + w + ac>) +1)
Ino
-1
+ (1
In)o
c²
(1+ ae?) (1+ae")
(1 + wl([ + aE
+ +b)
2
e²
(1 + ac²)
(1 ww/w)) + w w ac?) 2/ +b)
-1
E²
(1+ aea) (1 + 2w/wl) + + + ac?))
+6)
(85)
2 (1 + uE²) (1 + 2w/ + + ac + +b)
Since both fr. (rT) and P (In+1 = 1 xn,) are linear functions of In equating them for all vallues of xn {0, 1) implies that the respective coeffecients in these functions Eq. (84) and Eq. (85) are also equal, i.c.
c²
1
+ ac?) (1 + 2w/w/) w|w(1 + ac?) b)
p 1
=0
2
-1
2
(1 + ac²) (1 2n w|) + (1 ac²) +b]
(86)
2
-1
c²
(1 + ae?) (1 + 2w/w|) + wlw(1 ac2)1 +b)
p p= a
2
and hence
1
2²
2 (1 + ac?)
(1+ 2w/w/) + w|w + ac?) +
g:
(87)
c²
(1 + ac?) ) (1 + ww( + w| + ac?)) 1
+b) =
2
Since (2)) = for y 1) implies
log .ny. Eq. (87) can be rewritten
1
1 -
2
+ ac?) (1 2w/w]) w|w(( + ae?))
+ b = log
2
q
-1
c²
(1 + ac?)
(1 + 2w/w]) + w|w(1 + ac?)
+ b = log
2
1 p
52

---

Adding and subtracting the above two equations, we obtain
é² @l + ) log p(1 9)
q(1 - p)
(88)
e² (1 + ae²) 2w/w/] log (1 9) p)
pq
Thus Y. E R² is global minimum for L(-) if and only if it satisfies Eq. (88). It's easy that Y already critical point for L as Eq. (86) is equivalent to fi = f2 = 0 in Lemma 13. Thus, the set of all global minimum r, (p(p.) q) is given by
1
p(1
(. g) {y. (e, (,,, ER : wlwww( ae²) + = log
g(1 p)
e² + uE?) (1 2w/ww]) = log (1 - q) -
p4
Since the prediction f7. (.) equals the Markov kernel for any 7 E r., it follows from Thm. (or [23, Lemma I]) that [(y,) H (#n+1 Ix)), the entropy rate of the Markov chain. Now characterize the remaining set of stationary points.
Non-global-min critical points. For any critical point y = (e, w. a, b) E R., we have from the gradient expressions in Lemma 13 that (denoting fi and -f2 from the lemma as fi and f2) respectively)
0L 0b
Ex [/XX + -2]= =,
0L
(1iX +12) 12) 22(xx (x-1 1
de
Ex
2
2e (1 ae²) (1 2w/w()
Ex
(1ix + 12)
(x (x-
26aa (1 2w/w/l) 0,
(89)
0L
x 1
Ex
(fix + f2)
2ea (1 +tcc)) (lwl sign (w) w) = 0,
2
0L
X 1
Ex
(/iX 12)
6 (1 2w/ww/) 0.
Da
2
From Eq. (89), have that Ex [fiX + fel 0. Ex fxx E(
f2) = 0, we have that fi = f2 = 0, implying is global minimum. Hence without loss of
generality, assume that Ex ( + fa) (X-1) 1) * 0. Then we can partition the above set of equations into the following regions of stationarity:
Ex X + [2] = 0, =
iii Ex f + fe] = 0, + 0, ae² = 0, 0,14 + 2w/w [|| 0.
Slightly changing the variable order, for any 7 = (b, e, w. a) E R2, we define
a²
83 0hde
8 8h3m
0²L
862
L Bedb
8
8 L Bedhe
LLL deda
062
H(>) A v2
E 42x4
(90)
n Dwab
r. Dwde
8. Bu
A Dwtu
LL
L
L Budw
0²L
0uab
dade
dd²
53

---

Recall that
f1 = (21) (22)) + p +pp q 111,
f2 o (22) p.
1
21 _ c²
(1 ac²) (1+ 2w(ww) + w/w(1 + ac?)| }
+ b,
-1 (1+ ac") (1 2w wwl) w/w(1 + ac²) +b. b.
22 _ e²
2
From Eq. (89), we see that the second derivaties of L depend on the first-derivative of fi and f2, which we now compute. Recall that the derivative of the sigmoid function obeys (=)) (2) o(2) = (=)) for any & E R. Now the gradients fi and fz with respect to b, e, w and are
0f (21)0(-22) - a(2))o(-222.
06
0f (22)0(-2))
0b
9f o (2) (2)0(-21) {
[
2e
(1 ac*) (1 + 2w/w)) + w/ [[ ac²)
de
[1
(1 + 2w/w/) w|ulsign(1 2w/ 2w/wol) }
((21)0(-21))
2ac
1
((22)((-22)
2e
(1 ac?) (1 2w/w/) w|w + ac
2
1
a(22)o(-2z) (-22)
2
(1 2w(w() + w|wlsign(1 + 2w0|oo)
2
0f2 (22)0(( =2)
1
2e
(1 + ac?) (1 2w/lw() + w + ace))
de
2
(91)
1
(
2
(1
2w|w() + w|wlsign(1 + 2001oo1)
2
(2)) -21)
26
1
(1 + ac?)
|wo| +wwl| |w||| acc1
ow
2
|1 ae? c
(322)((-20 (-22)
2c
(1 + ae²) |wol + |liml
-
0f
1
|m|| |m||| ae? }
2
[
(1 ae?)
Ow
2
0f
o ( (2)0((211)
{
1
(1+ + 2w 2w/w/)
1 slwlsign (1 + ae?) }}
Da
a(2))o(-22) {
1
(1 ae?)) (1 2w/w) + w ( sign (1 ae?) }
(1 +
2
0f2 a22)) 22) { e
1
da
2
Now we characterize the first set of critical points.
(i) Stationary points with Ex [ + fal = 0, 0. When e 0, we have that 21 22 = b.
Hence,
fi = o(b) +p (b) = p+ -
f2 o(b) p.
54

---

Thus, Ex [fiX + f2] = (p ( 1x +o p q +o -p = 0. Rearrang- ing and simplifying, o(b) p+v and 0 (- ) further reduce
P and hence b = log -- Using the fact that log "
n p+4 77
_ the above gradients evaluated for any 7 (b log 2. = 0, w. a)
Qf
0f2
0,
To
Øb
Øb
7
Ofi
ff
O,
0,
de
de
(92)
0f
0f
0,
Ow
dw 7
7
051
0/2
O,
0. 7
Da
da
Now substituting Eq. (92) when computing the recond-derivativer of L in Eq. (89), we obtain
0f1
0f2
Ex
X
063
0b
06
LL
0fi
0f2
Ex
X
0,
0bde
de
De
L
0f1
0fz
Ex
X
0,
Ow
Ov
7
L
0f1
0fz
= Ex
X
+
0,
abda
da
aa 7
(fiX + f2) (x (x-2) -11
Ex
2 2(1 ae?) (1 2w/w)) 7
De2
- 2 ]
+ Ex
(n1X + -2)
4e2a (1 + 2nc/wrl)
7
(93)
05i
0f2
)
(xx (x-11) 1)) 2 2(1 (1+ + ac²) (1 + +w//w/)
+ Ex
de
+
De 1,)
Ex [12f( + 2w/w/) fi + 2f2(1 2w /w/w////2 w|)x - 12]
7
Ex [[f( + 4w/w== fz 4w/uu)x - fel
(1i Aw/w/) f2 4wl -f2
(4) (p+ 4w wl) x1 1)(2 + 4w ww/ q 1)
1) 4w/ (2 + 4m0 w:l) + 1)
2m Tolp 1) + 2w(w)),
where (a) follows from the fact that fily -ーー
=
-m (p + q 1) and (6) from - no. Returning the remaining second derivatives,
55

---

8LL
a Ex
(fix + 12) x (x-²) 1 4ea (1+ + ae²) |wo
dedm
de
4
(x-1) 2) (8e + ae*) |mm + Se a w0
-Ex (fiX + f2)
(0fi)
0f2
) (x- 4ea (1+ + ae*) (loo)
Ex
X
de
de
0,
8L
① Ex
X (x-1) ) et (1 + 2w wl)
(fiX + f2) +52)
2)
deda
de
7
(x-1) 1)
Ex
(ViX + f2)
4eah (1 + 200/wol)
7
[2f de
0,2
) (x- -1/
+
+Ex
X
e* (1+ + 2w0 w)l)
De
0,
8LL
(JiX +52) (2) X (x-1) ) 4ea (1 (1+ (1+ae²) |wo
Ex
dw
7
Ex [ f2)
(x-1) 1
4ea (1 + ae?) sign(w)
2
(94)
0.
) (x- 2) 4ea (1+ aea)) [wl
+ Ex
X +
Ow
0,
8L
(fiX+ 12) (x (x- : e* (1 + ((1+2wl
Ex
Owda
Ir
7
X xx-) 2e 2e
Ex [fiX + f2)
[(
X + 012
1 e (1 + ww/w()
+ Ex
iw (@w
Ow
7
0,
PL. =
0 Ex
(fiX + f2) (x-1) 1 e² (1 + 2w wl)
da² 7
Da
(x (x-1) 2/
1
-xx (fiX 12)
2x 2e
[Off
(
[0fi
0.
|
+ Ex
X
et (1 + 2w|w/)
Ow
Dw
7
= 0.
Congregating all the second derivatives from Eq. (93) and Eq. (94) into the Hessian H (>) in Eq. (90). we finally obtain
0
0 2(p + ạ
1)(1
+ 2w w)
H H()) = गoo
0
0
56

---

which is identical to the Hessian obtained in the proof of Thm. 7 (App. J.2) for e = 0. Thus it follows that T min (p. 9) - R$ and I sad C R+ defined below are set of local minima and saddle points
respectively:
T minn (P;n(P
Ymin = (e,w, : -
1) + 2w/w() > 0, = log L
9
I (p, _ Ysad (e, = - 2w 2w/ww) S 0, b log
Now we focus on the remaining set of critical points.
(ii) Stationary points with Ex [fiX fel 0, & 0, 1,11 + ac² = 0 = 0. For this set of points, the Hessian remains undefined because afi 0f2 0h 0f2 do not exist (Eq. (91)). This non-
de de de 8n
existence arises since sign (1 + ae?) is undefine 1 + ae = 0. However, even in this scenario,when
* 0, 1 + ac 0, + 2w/w) 0, we have 22 b. Hence,
fi o(b) +p+ o(b) = -
f2 = obb)
Thus the expectation term Ex [] + [2] (p - Ex [X] + ob(b - (p + q 1) m
(6) -p p = O. Simplifying, o(b) p. which implies b = log
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 I stationn (P, q) S I is defined as
{
P { (7min = b,, ER :e : 0,1 + ac² = -,,1+ 2w/ wl = 0, b = log
I sstation (p. 9) &
q
This concludes the proof.
N.3 Proof of Lemma 14
Proof. Recall from Lemma that for 0 = (e, w. a) E TR3, we have
0L
E
(JiX
12)
X 2
2e (1 + at?) (1 2w/w/l
de
x - !
E
(ЛiX f2)
20aa (1 + 2w/lw)),
2
0L
x (x-1) 1
E
(/iX F22
2
26 (1 + ae²) (lwl + sign (w)w),
Dw
7L
[x- !
E
(1iX + f2)
e (1 2w(ww)
da
where E {0, 1} a Bernoulli random variable with ~ Bern(p/ + q)), and
fi
A 1
11- q 61 + Φo.
f2 - p- Φo.
1
2
(1 + ae. )) (1 + 2w/ w w|w(( + ue²) + bbx) by
2
1
Φ
A
2
(1 + ac²)
(1 + w)) + wlw((1 + ac + by
2
57

---

where the optimal bias b, is obtained by solving TIfi + f2 = 0. Using the definition of the gradient flow that
(0) for 0 = fr, we have
L( E (fiX + f2)
i dw
X
2e² (1 + ae²) (12ol + sign (w) w)
2/
(95)
1)
E (fiX + f2)
X
2 2(1 + acc
e² (|w) + sign (w) w) Similarly for a,
2
I(0)
1-11
)
à
=E (fiX + f2)
X - 2
e² (1 + 2w(ww)
da
(96)
à
X 1)
E (fix+ + (2)
2)
(1 (1+2w)wl).
Likewise, for
è
L(0)
= E (fiX + f2)
X -
2)
(1 + aE²) e(1 + w
de
(97)
+ E ( + 12)
X (1+ + 2w0 (ollwl) 2
By substituting the expressions of Eq. (95) and Eq. (96) into Eq. (97):
(1iX f2) x 1x-11 2
é = E
2(1 =taee)) (1 + 2w(w))
2
Eq. (95)
(98)
1
+
(1i f2)
x -- 2
(1+ 2w/w/) 26 a.
Eq.
Thus, we obtain
2e (wl) + sign w)) e (1 + 2 (ol) + à
20aa.
(99)
On rearranging and simplifying:
eè (wwl + sign (w) w) (1 + 2w uo|) + a2a
(100)
vb
eé
(1 2wo w/ ) + 2aá.
2 ([w])
Integrating the above equation on both sides:
f cè 4 ([w]) (1 (wl + 2ad
ib
(t) sign(w(()) log (w(t) + w(t)²
+ a1)2 + e² (1) = sign(w((1) log (t) + w=1)² + 2a(1)² +
c
(101)
2
Note that here E IR. is constant that depends the initial conditions. Thus the energy 8(0,)
E E(On) for wo # 0.
58

---

N.4 Proofs of Lemma 16, Lemma 15, and Thm. 3
Proof. We that the proofs of Lemma 16, Lemma 15 directly follow from that their counterparts Lemma 11 and Lemma 10 using Lojasiewicz's 's theorem to characterize the convergence of the gradient flow. Thm. 3 is direct consequence of Lemma 16, Lemma 15.
N.5
Informal proof of Thm. 11
Proof. [Informal] For 0 = (e, w, a), recall from Lemma 13 that the derivative of the loss with respect is
0L
x (x-2) 1 2e (1+ ac²) (1+ 2w|wel
de = Ex (fiX f2)
(x- _
2 )
+ Ex (fiX + f2)
2e a (1 + 2w/w)),
where
f1 = (21) - -0(22) + +
f2 o (22) Dr
[1
21 - e
(1 ac?)) (1 2w/wwl w/w((1 ac?)
+ b,
-1
2/ _ 2
2 (1 + ac²) ac²)
(1 2w (w/w/) w w[111 + act)| b.
Assuming the initialization is very small, making any product of quantities in O (e, w, a, b) much smaller than the individual quantities. Therefore, we can consider these products to be approximately zero. That is, Vx, y E 0, x 2 xy & y 2 xy & ay s O. Hence,
21 b,
22 = b,
f1 P + 9
f2 = 0 (b) p.
Hence the gradient with respect to e is
0L de
2Ex [[ix + f2) (X)]
(102)
Simplifying the expectation term, Ex [ fz) x -
g
() Tf), where used the fact that b is optimal in the above equations, specifically where fi + f2 = 0. Thus the gradient flow for the parameter e governed by
0L
é =
(pp + 1)( - m7 ep expeeoeppppppppppp 1)(m =)))).
de
Since (p + - 1)( xf) 0, + O, which denotes it converges to the local minima.
59

---

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 short (1-2 sentence) justifiiaton 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 (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 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 binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write 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
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: All claims are in the form of theorems that are proved in Sections 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 the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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]
60

---

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.
Justification:
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 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 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-ten 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, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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. the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix 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 ex- perimental results of the paper to the extent that 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.
61

---

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 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 novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary 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 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- 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. the case of closed-source models, 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 have some path 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes)
Justification: We provide open access to the code, available at http:::tpp///// h/atttpps:///nnnnousssss 4open. ocience/r/Locall 1ooo--oobll-700-
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/
for more details.
While we encourage the release of code and data, we understand that this might 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 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: / ciitps.cc/uubbice httpppulppipuila.pppppppppppaeee /ttgg/dee///ooeuumiiiiiin htt.e//mmmmiissiooooooo..ooooooo ssi.noolicyy 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. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
62

---

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) 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, hyper- parameters, 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 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, 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, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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).
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 should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful 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 the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: These are unrelated to the results of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
63

---

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 compute required for each of the individual experimental runs well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics attps://nnet stttps://neurps.... httpp:///aerrpp....ppbl//////ii
Answer: [Yes]
Justification: All ethical practices were followed.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration 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 no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
Answer: [NA]
Justification: No risk.
64

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
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 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 used and, 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 (c.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 in the package should be provided. For popular datasets, paperswi snporswwihcode thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should 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 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 release new assets.
Researchers should communicate the details 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 details about compensation (if any)?
65

---

Answer:
[NA]
Justification:
No human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, 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 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 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 NeurlPS 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.
66


Reviews:
[{"review_id": "AgIVYeqaH2", "reviewer": "Reviewer_LVoR", "summary": "The paper theoretically studies the effect of initialization on the gradient dynamics of a single-layer transformer trained on Markov data. It considers a simplified model by: (1) using a binary input alphabet and (2) reducing the many parameters of the single-layer transformer to just two or three scalar parameters. The authors identify the critical points of this reduced set of parameters and compare them to the stable point under its gradient flow dynamics. This is used to determine whether an initialization will converge to global or local optima. This informs choices of initialization that will reliably converge to global optima. This is validated in a single-layer transformer where the reduction of weights is not performed. This experiment is also used to validate the reductions made for theoretical analysis.", "questions": "1. How much do results vary across replications? The authors could easily answer this by showing different runs in the same figure or stating measures of uncertainty. If possible, the authors could also consider more values of p and q.\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.", "limitations": "The authors are fair in the scope of their claims, though it would be good to emphasize the use of reduction/reparameterization earlier in the paper or abstract (this is not made clear until the end of page 3, and this reduces the theoretical scope of their claims). Alternatively, the authors could provide stronger evidence to justify the universality of these approximations (precise quantification of low-rank structure, and more replications).", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This 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.", "weaknesses": "It is not clear how much results vary across repeated trials. The authors state that they repeated experiments 3 to 5 times, but the results across replications are not presented in the figures or the text. For instance, in Figure 6, did every repeated experiment follow the same blue and red curves? Similarly, while the authors show the weights from a single run on Figure 5, they do not provide a numerical assessment of the validity of their low-rank approximations. Another weakness is the simplified data model, which is just a binary alphabet. The authors seem to inherit this from previous analytical work, but a discussion of how this restricts results and real-world applicability is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback and insightful comments. We refer to the global response regarding our results across repeated trials for various $(p,q)$ and the measure of low-rankness. We address the individual concerns below.\n\n- **Binary and large alphabet**: This paper primarily focuses on the binary alphabet to analytically characterize the phenomena reported in [1], specific to the binary setting. However, we observe a similar low-rank structure for transformer parameters in a first-order Markov chain on a larger alphabet (see Figure 6 in the rebuttal PDF). Consistent with the binary setting, parameters initialized as low-rank remain low-rank throughout training, as shown in Figure 6 .\nTheoretically, while our gradient-flow analysis can be generalized to the multi-alphabet setting, it is challenging to precisely characterize the phase portraits, as in Figure 2 of the paper, due to the increased number of parameters. This presents an intriguing avenue for future research.\n\n- **$ a \\approx 0 $**: This results from the fact that for first-order Markov chains, the Markov kernel $ \\mathbb{P}(x\\_{n+1} = 1  \\mid x_1^n) = x_n (1-p-q) + p $ is a simple linear function of the last symbol x_n​. Consequently, this function can be represented by a single-layer transformer using only the skip connection in the attention layer, without relying on the attention mechanism. Therefore, $ a \\approx 0 $  becomes a viable solution, which we leverage for the gradient-flow analysis without attention. Figure 7 in the rebuttal empirically corrobarates this fact. As a side note, there are also non-zero attention coefficients that represent global minima, as shown in Figure 4 and discussed in Section 4.1 of the paper.\n----------------------------------------------------------------\n**References:**\n\n[1] Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains: https://arxiv.org/abs/2402.04161", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hYgOreHxRn", "reviewer": "Reviewer_DieV", "summary": "This paper investigates the learning dynamics of transformer models, specifically focusing on first-order Markov chains and single-layer transformers. The authors aim to understand how transformers learn Markov chains and the impact of initialization on their training outcomes. They provide a comprehensive theoretical analysis, demonstrating that transformer parameters trained on next-token prediction loss can converge to either global or local minima, depending on the initialization and properties of the Markovian data. The paper also offers empirical evidence to support their theoretical findings and proposes guidelines for initializing transformer parameters effectively.", "questions": "The paper does not address the scalability of the proposed initialization guidelines for large-scale transformer models, which are commonly used in practice. Can the authors provide an insight on this?", "limitations": "As I mentioned in weaknesses section, the most important limitation of this study is its generalizability. The study focuses on single-layer transformers and first-order Markov chains, and it is unknown if this method is generalizable to more complex models and data structures.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "-  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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and helpful suggestions to improve the paper. We will add a separate section with prerequisite background details in the revised version. We refer to the common response for experiments on multi-state Markov chains. We address the individual concerns below.\n\n**Scalability**: Understanding the theoretical foundations of transformer training and initialization, even for shallow models, remains a challenging task, which our work aims to address. While this paper focuses on shallow transformers, it represents a crucial step toward understanding optimization in deeper models. For instance, recent research [1] employs a similar reparameterization technique to study gradient descent dynamics in a simplified two-layer attention-only transformer, elucidating how they learn the induction head mechanism. However, unlike our work, they do not investigate how initialization impacts convergence to local or global minima. Given the nascent and evolving understanding of deeper transformers and higher-order Markov chains [1,2], exploring similar initialization guidelines for large-scale models is an exciting avenue for future research.\n\n---------------------------------------------------------------------------\n\n**References**\n\n- [1] How Transformers Learn Causal Structure with Gradient Descent? https://arxiv.org/pdf/2402.14735\n- [2] Transformers on Markov Data: Constant Depth Suffices, https://arxiv.org/abs/2407.17686v1", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cZOweq14pQ", "reviewer": "Reviewer_JG7D", "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.", "questions": "Please see the weakness.", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "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.\n2. The paper is well-structured and well-written, offering clear explanations and visualizations that enhance comprehension.\n3. 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.\n2. Can the initialization guidelines provided be extended to more complex transformer models, such as those with additional layers?\n3. 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?\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback and insightful comments. We will update the revised version with information about the attention in the main text. We address the individual questions below.\n\n- **Rank-1 input sequence……:** Please note that the input is a first-order Markov chain, without the assumption of it being rank-one. We use linear attention for ease of theoretical analysis, but our guidelines and insights generalize to non-linear soft-max attention. For instance, Figure 6 in the paper demonstrates the effectiveness of our initialization scheme for the full one-layer transformer model (Lines 60-64). Similarly, the low-rank structure of the parameters, as shown in Figure 5 in the paper (Figure 2 and 6 in the rebuttal PDF), applies to the full general model.\n\n- **Initialization guidelines**: As our analysis shows, understanding the theoretical foundations of transformer training and initialization, even for shallow models, is challenging. Our paper, focusing on shallow transformers, is a crucial step toward optimizing deeper models. Recent research [1] employs a similar reparameterization technique to study gradient descent dynamics in a simplified two-layer attention-only transformer, elucidating the induction head mechanism. However, unlike our work, they do not explore how initialization impacts convergence to local or global minima. Given the nascent and evolving understanding of deeper transformers and higher-order Markov chains [1,2], exploring similar initialization guidelines for large-scale models is an exciting avenue for future research.\n\n- **Low rank parameters**: Recent evidence [3,4] suggests that SGD with weight decay inherently leads to rank minimization, resulting in low-rank parameters at convergence. Although these studies focus on feed-forward neural networks, we believe a similar phenomenon occurs in transformers as well..\n\n-------------------------------------------------------------------------------------------------------------------------\n**References**\n- [1] How Transformers Learn Causal Structure with Gradient Descent? https://arxiv.org/pdf/2402.14735\n- [2] Transformers on Markov Data: Constant Depth Suffices, https://arxiv.org/abs/2407.17686v1\n- [3] Characterizing the Implicit Bias of Regularized SGD in Rank Minimization, https://arxiv.org/pdf/2206.05794\n- [4] Rank Diminishing in Deep Neural Networks: https://openreview.net/forum?id=tIqzLFf3kk", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0JEsFBruKc", "reviewer": "Reviewer_Fte8", "summary": "The paper seeks to characterize how single layer transformers learn (two symbol) markov chains. The analysis relies on a reparametrization of the transformer that assumes the weight matrices are rank 1. The primary results are characterizing when gradient flow on this reparameterized transformer leads to global minimums in loss or other critical points, showing how in some cases a standard gaussian initialization will lead to a local minimum instead of the global minimum.", "questions": "1. How is figure 6 testing the hypothesis that “for p + q > 1, any small initialization around zero would lead to a local minima” (line 247). I feel like it's more of comparing a small initialization to your parametrization (which is still a worthy endeavor). I’d imagine that testing the stated hypothesis would require trying many different small initializations and showing that they lead to local minima (Given the nature of the problem, I imagine it could be feasible to test initializations in a grid and create something like an empirical analogue to corroborate figure 1d?). Also, a natural follow-up question (or null hypothesis) is can some standard initialization that isn’t small and/or isn’t around zero work just as well as your initialization.\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?", "limitations": "The discussion of limitations in the conclusion section (which I mentioned specifically because of the response given in the checklist, lines 1200-1201) is a bit too brief. It would be good to have all of the assumptions made for the main theorems in one place, especially those implicit and not stated in the theorem statements, including the rank one and related assumptions, linear attention, and gradient flow (which is in the theorem statements, but good to mention because transformers are not normally trained with gradient flow).", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The 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.", "weaknesses": "1. The canonical parameterization is a rather strong assumption (reducing the model to just three parameters) and so it needs strong results to justify it. For me, the results in this paper are close to strong enough, but some strengthening of them would be appreciated (short of removing assumptions).\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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback and suggestions to improve the paper and the code which we are planning to incorporate in the revised version. We refer to the global response regarding our results across repeated trials for various $(p,q)$ and the measure of low-rankness. We address the individual questions below.\n\n\n- **Our insights carry over to the full general model**: While our theoretical analysis uses canonical parameterization, our insights and guidelines extend to the general single-layer transformer with soft-max attention (Lines 60-64). For example, Figure 6 demonstrates the effectiveness of our initialization scheme for the full one-layer transformer model. Similarly, the low-rank structure of the parameters, shown in Figure 5 in the paper (Figure 2 and 6 in the rebuttal PDF), pertains to this general model.\n\t\n-  **Canonical parameters exhibit the same dynamics as the full model**: We strongly believe that the training dynamics of the canonical model closely resemble those of a full single-layer transformer with rank-1 initialization. Due to the time constraints, we currently lack empirical evidence for this. However, we will share this comparison during the discussion phase. \n\n- **$W_1$ and $W_2$ are not constants**: We are afraid that there has been a misunderstanding about the initialization in Figure 6. While $W_1$ and $W_2$ are both initialized constant, they are trained alongside all other parameters  $\\theta$ in the single-layer transformer with non-linear soft-max attention (Lines 60-64). Thus, Figure 6 of the paper (Figure 1, 3 and 4 in the rebuttal PDF) effectively demonstrates the success of our initialization scheme and the broad applicability of our insights.\n\n- **Q.1**: In Figure 6 in the paper, our standard initialization scheme is to let $e_0=1$ and $w_0 = -1 $ which falls in the region $ \\mathcal{I}_\\ast $ (see Figure 1d) and converges to the global minimum as predicted by our theory. To test our hypothesis on a wider range of initializations, we try the following set of initializations which concur with our theoretical predictions (Figure 3 of the rebuttal): we let $(e_0,w_0) \\in \\mathcal{N}(\\mu, 0.02) $ with $ \\mu \\in \\{ (-1,-1), (0, -0.5), (0,0.5), (0, 0)  \\} $, with $ p=0.5$ and $ q=0.8$. As highlighted in Figure 3 of the rebuttal, we observe that the initializations around  $ (-1,-1)$ leads to global minimum convergence whereas the rest  converge to local minima, as predicted by Figure 1d in the paper.\n\n- **Q.2**: Yes indeed. We wanted to highlight the main ideas behind our gradient-flow analysis using canonical parameterization and hence chose to omit the attention parameters in the beginning for the ease of exposition and clarity.\n\n-----------------------------------------", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pduCqFP8dQ", "reviewer": "Reviewer_er5c", "summary": "This paper focuses on training a single-layer linear attention transformer with low-rank parameterization on first-order Markov chain data. By reparameterization, they can reduce this problem to a 3-variable learning dynamics and comprehensively characterize the trajectory and local/global minimizers. This paper precisely characterizes the different convergence conditions for local/global convergence, highlighting the role of initialization. Empirically, on this specific data distribution, they verified the superiority of their special initialization compared to standard initialization.", "questions": "1. Does the low-rank initialization scheme insight transfer to more general data distribution or architecture?", "limitations": "The dynamics analysis is restricted to certain subspaces. The attention pattern does not train (which is simplified as a trainable scaling factor $a$). The experiments cannot corroborate the general claim.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This 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.", "weaknesses": "1. I am concerned with the generality of the insights on the role initialization. I appreciate that the authors make a great effort to theoretically analyze the Markov Chain with 2 states (i.e. $P\\in R^{2\\times 2}$) and do extensive experiments on this particular dataset with linear attention. However, the authors claim that ``we offer practical guidelines for initialization of transformer parameters and demonstrate their effectiveness.\"(Line 280-281) in both the abstract and the conclusion. Line 255 also claims \"the corresponding insights are more general and apply more broadly to the general architecture.''. Since all the conclusions are based on this simplified setting, I don't think it can be characterized as a realistic guide for initialization. I think further experiments on the Markov data with more states and non-linear attention transformers can benefit this paper and make the results more general. I will increase the score if the results can be verified in a more general setting.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback and insightful comments. We address the individual questions below:\n\n- **Generality of the insights**: For the two-state Markov chain, while our analysis capitalizes on the canonical parameters and linear attention, our guidelines and insights from this setting readily generalize to non-linear soft-max attention. For instance, Figure 6 in the paper (Figure 1, 3 and 4 in the rebuttal) demonstrates the effectiveness of our initialization scheme for the full one-layer transformer model with soft-max attention (Lines 60-64). Similarly, the low-rank structure of the parameters, as shown in Figure 5 (Figure 2 in the rebuttal), is concerning the full general model. The lines 280-281 and 255 were meant to convey this message which we are happy to rewrite to make it more clear.\n\n     For multi-state Markov chains, we observe a similar phenomenon where, under standard Gaussian initialization, a single-layer transformer can converge to either local or global minima depending on the Markovian switching.  Specifically, in Figure 5 of the rebuttal, we consider a first-order Markov chain with state space $S = \\\\{ 0, 1, 2, 3, 4 \\\\} $ where any state $s \\\\in S$ stays in the same state with probability $1-p$ and switches uniformly random to either of the remaining states with probability $p/4$. This generalizes the binary-state Markov chain in the paper with $ p = q$. As illustrated in Figure 5 of the rebuttal, depending on the values of the switching probability $p$, the transformer parameters either converge to local-minima (unigram) or global minima (bi-gram). Further, Figure 6 of the rebuttal also reveals the low-rank structure of the parameters during training for the multi-state setting too. These interesting empirical observations highlight an exciting research opportunity to mathematically characterize the gradient dynamics for the multi-state Markov chains, similar to our two state analysis. \n\n- **Low rank manifold**: Indeed, using canonical parameterization to characterize the gradient descent dynamics of transformers is a standard tool in the literature. For instance, a recent work [1] employs a similar approach to show that two-layer transformers learn induction heads, assuming specific structure for the attention matrices in both layers, inspired by empirical observations (Figure 2, [1]). \n\n- **Limitations**: Please note that the attention parameter $a$ is also trained alongside the embedding and weight parameters $(e,w)$, and the corresponding analysis is highlighted in Section 4.1. All the figures in the paper reflect this setting, with no parameters being frozen or omitted.\n----------------\n\n**References**\n\n[1] How Transformers Learn Causal Structure with Gradient Descent? https://arxiv.org/pdf/2402.14735", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## **Generality of the insights: Our insights and conclusions hold for all pairs of $(p,q)$, across repeated trials, and for non-linear soft-max attention**

We thank the reviewers for the constructive feedback. We address the common concerns here regarding the error bars across repeated trials, experimental results for various $p$ and $q$, a metric for the low-rankness, and linear vs. soft-max attention.  

### **Error bars**

In the paper, we reported results averaged across three to five trials but omitted the error bars due to an oversight. The rebuttal PDF includes all figures with error bars, confirming that our conclusions remain valid. For instance, Figures 1, 3 and 4 in the rebuttal shows that our initialization scheme consistently converges to a global minimum across multiple trials while standard Gaussian initialization gets stuck at the local minima.

### **Experimental results for various $(p,q)$**

We would like to emphasize that our theoretical results hold for  all values of $(p,q)$ with the regions $p+q <1 $ and $p+q > 1$ exhibiting fundamentally different behaviors as highlighted in respective Theorems 7 and 2, and Figure 1. For the empirical results in the paper, especially for Figure 6, we chose a single pair $(p,q)= (0.5, 0.8)$ in the region $p+q>1$ for the ease of exposition only. Complementing this result, we further empirically demonstrate in Figure 4 of the rebuttal that our initialization scheme consistently outperforms the standard Gaussian one for various pairs of $(p,q) \in  \{ (0.6, 0.6), (0.7, 0.9)   \} $ all satisfying $p+q>1$.

### **Low-rank metric**

Indeed, complementing our visual illustration of the low-rank matrices in the paper, we can also mathematically quantify this property via the following metric: we compute the relative energy contained in their top singular value compared to their overall singular values, i.e. $ \frac{\sigma\_1^2}{\sum_{i} \sigma_i^2 }  $. The closer this fraction is to one, the closer the corresponding matrix is to being rank-one. Tracking this metric training for various matrices, Figure 2 in the rebuttal highlights that while random initialization training eventually reaches rank-one parameters, initializing at rank-one always retains that structure during training. This figure corresponds to the same setting as Figure 5 in the manuscript. 

### **Linear vs. soft-max attention, canonical parameters vs. full-transformer**

We are afraid that there has been a misunderstanding about the setting in the paper regarding linear and soft-max attention. For the two-state Markov chain, while our theoretical analysis capitalizes on the canonical parameters and linear attention, our guidelines and insights from this setting readily generalize to the full single-layer transformer with non-linear soft-max attention (Lines 60-64). For instance, Figure 6 in the paper (Figure 1 in the rebuttal) demonstrates the effectiveness of our initialization scheme for the full one-layer transformer model with soft-max attention across various values of $(p,q)$ and repeated trials. Similarly, the low-rank structure of the parameters, as shown in Figure 2 and 6 of the rebuttal (Figure 5 in the paper), pertains to the full general model.

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