Attention with Markov: A Curious Case of Single-layer Transformers
Authors: Ashok Vardhan Makkuva
Marco Bondaschi
Adway Girish
Alliot Nagle
Martin Jaggi
Hyeji Kim
Michael Gastpar
Venue: ICLR 2025 Spotlight (2025)
Keywords: Markov chains
Transformers
Optimization
Landscape
Primary area: interpretability and explainable AI
TL;DR: We theoretically characterize the loss landscape of single-layer transformers and show the existence of global minima (bigram) and bad local minima (unigram) contingent upon the specific data characteristics and the transformer architecture.

Abstract: Attention-based transformers have achieved tremendous success across a variety of disciplines including natural languages. To deepen our understanding of their sequential modeling capabilities, there is a growing interest in using Markov input processes to study them. A key finding is that when trained on first-order Markov chains, transformers with two or more layers consistently develop an induction head mechanism to estimate the in-context bigram conditional distribution. In contrast, single-layer transformers, unable to form an induction head, directly learn the Markov kernel but often face a surprising challenge: they become trapped in local minima representing the unigram distribution, whereas deeper models reliably converge to the ground-truth bigram. While single-layer transformers can theoretically model first-order Markov chains, their empirical failure to learn this simple kernel in practice remains a curious phenomenon. To explain this contrasting behavior of single-layer models, in this paper we introduce a new framework for a principled analysis of transformers via Markov chains. Leveraging our framework, we theoretically characterize the loss landscape of single-layer transformers and show the existence of global minima (bigram) and bad local minima (unigram) contingent on data properties and model architecture. We precisely delineate the regimes under which these local optima occur. Backed by experiments, we demonstrate that our theoretical findings are in congruence with the empirical results. Finally, we outline several open problems in this arena.

Published as a conference paper at ICLR 2025
ATTENTION WITH MARKOV:
A CURIOUS CASE OF
SINGLE-LAYER TRANSFORMERS
Ashok Vardhan Makkuva' EPFL
Marco Bondaschi* EPFL
Adway Girish EPFL
Alliot Nagle UT Austin
Martin Jaggi
Hyeji Kim
Michael Gastpar*
EPFL
UT Austin
EPFL
ABSTRACT
Attention-bast based transformers have achieved tremendous success across a variety of disciplines including natural languages. To deepen our understanding of their sequential modeling capabilities, there is growing interest in using Markov input processes study them. A key finding is that when trained first-order Markov chains, transformers with two or more layers consistently develop an induction head mechanism estimate the in-context bigram conditional distribution. In contrast, single-layer transformers, unable to form an induction head, directly learn the Markov kernel but often face a surprising challenge: they become trapped in local minima representing the unigram distribution, whereas deeper models reliably converge to the ground-truth bigram. While single-layer transformers can theoretically model first-order Markov chains, their empirical failure to learn this simple kernel in practice remains a curious phenomenon. explain this contrasting framework for
behavior of single-layet models, in this paper introduce principled analysis of transformers via Markov chains. Leveraging our framework, we theoretically characterize the loss landscape of single-layer transformers and show the existence of global minima (bigram) and bad local minima (unigram) contingent on data properties and model architecture. We precisely delineate the regimes under which these local optima occur. Backed by experiments, we demonstrate that our theoretical findings are in congruence with the empirical results. Finally, we outline several open problems in this arena. Code is available https: /gi thub. com/ Bond1 995 /arrovv.
INTRODUCTION
Attention-based transformers have been at the forefront of recent breakthroughs in variety of disciplines, including natural language processing (Vaswani al., 2017; Radford and Narasimhan, 2018; Devlin et al., 2018). One of the key workhorses behind this success is the attention mechanism, which allows transformers to capture complex causal structures in the data, thus rendering them with impressive sequential modeling capabilities.
Given their success, there tremendous interest in understanding the sequential modeling abilities of transformers. Notably, a growing body of research explores transformers through Markov input processes investigate their in-context learning capabilities (Rajaraman et al., 2024a; Nichani et al., 2024; Edelman et al., 2024; Bietti et al., 2023). These studies reveal an interesting insight that transformers with two or more layers develop an induction head to estimate the in-context bigram conditional distribution when trained on first-order Markov chains. In contrast, single-layer transformers, unable to form an induction head (Olsson et al., 2022), directly learn the Markov kernel. Surprisingly, we empirically find that while deeper models reliably converge the ground- truth bigram, regardless of initialization, single-layer transformers often get stuck in local minima corresponding the unigram distribution (Fig. 1). Despite their theoretical ability model first-order
"Fqual contribution. ashok. www.kma@spkl.aee.eom
'Equal contribution.
Correspondence to: Ashok Vardhan Makkuva,

---

Published as a conference paper at ICLR 2025
Markov chains, they sometimes fail to learn this simple kernel in practice. Motivated by this stark contrast in behavior based on depth, and our limited understanding of we ask: Can we systematically characterize the learning capabilities of single-layer transformers with Markovian inputs?
To address this, in this paper we introduce a new framework for a principled theoretical and empirical analysis of transformers via Markov chains. Leveraging framework, characterize the loss landscape of single-layer transformers and prove the existence bad local minima and global minima corresponding the unigram and bigram, respectively. We further demonstrate that the presence of these local optima depends on the Markov state switching probabilities and the transformer's weight-tying and we precisely delineate the regimes under which this happens. Together, our analysis reveals a complex interplay between the data aattadistributionall properties, the transformer architecture, and the final model performance for single-layer transformers with Markov chains, explaining the aforementioned empirical phenomena.
In summary, we make the following contributions:
We provide novel framework for precise theoretical and empirical study of transformers via Markov chains (Sec. 3).
We characterize the loss landscape of single-layer transformers with first-order Markov chains, highlighting the effect of the data distribution and the model architecture (Sec. 4),
We show that the Markov switching probabilities and weight-tying play a crucial role in the presence of local optima on loss surface and precisely characterize the said conditions (Thms. 2 and 3).
L-layer transformer
0.775
l-yye transformer
4-layer transformer
0.750
8-layer transformer
0.725
Unigram loss (local minimum)
Test loss
Bigram loss (global minimum)
0.700
0.675
0.650
0.625
0.600
100
200
300
400
500
Iteration
Figure 1: Single-layer transformers get stuck at local minima, corresponding to the unigram model, when the input first-order Markov chain with switching probabilities p 0.5 and q 0.8 (Fig. 2b). However, deeper models escape global minima corresponding to the bigram model.
Our main findings and observations are:
We prove that weight tying can introduce bad local minima for single-layer transformers when Markovian switching is greater than one (Thm. 2). Removing the tying, however, solves the issue (Thm. 3).
When the Markovian switching is less than one, empirically observe the model always converges to the global minima irrespective of the weight tying (Fig. 3).
Interestingly, transformers with depth two and beyond always converge to global minima irrespective of the weight tying and switching (Fig. 1).
Notation. Scalars are denoted by italic lower case letters like x, and Euclidean vectors and matrices denoted by bold ones æ, y. M,, etc. We use denote the forn for Fuclidean
2

---

Published as a conference paper at ICLR 2025
vectors and Frobenius norm for matrices. [k] _ {1....,k), k}, and for sequence (#n)n21 define oi ( Fmm) if k 2 1 and (3):- (21:-..,, Fm) otherwise. For 2 R. the sigmoid d(z) - 1/(1 + e) and ReLU(2)) & max For events A and B, (A) denotes the probability of A whereas IP B) the conditional probability. Let (æ:, v) be pair of discrete random variables on [k] x [k:] with the probability mass function (pmf) of x being P= = (pi.....PPh PPs) E [0, 1]k.
Then its Shannon entropy is defined as (x) = (p.) _ EiEM
ps log Pt. and the conditional entropy (y) A Ha,, v) (x)). The entropy rate of stochastic process (Fn)n21 defined as lim, H(=))/n n. Finally, for p E (0, 1). the binary entropy function h..)) is defined as h(p) & H(p, p) - pogp (1 p) log p).
2
BACKGROUND
We describe the transformer architecture and the Markovian input process.
2.1 TRANSFORMERS
We study a single-layer transformer with a single-head softmax attention and ReLU non-linearity. We omit the layer norm since its influence marginal in the settings we consider (Sec. 4). We consider an input vocabulary X' of finite size | For the ease of exposition, in this paper we mainly focus on |X| 2, i.c. X' = {0, 1}. and outline our results for multi-state setting in Sec. 4.3. Let {In IN E 1] N be an input sequence of length N. Then for each n E [N], the transformer operations are mathematically given by (Fig. 2a):
In In G1 + (1 Inn In) Eu + P., ER.
(Embedding)
y, In + Wo Σ
attn,i Wv æ; Rd.
(Attention)
(e(m)
nn y, + W2 ReL((1, y,)) R logit, = (a, zn) +
(FF) (Linear)
E R,
fø (xr) A Pg (In+1 = 1 1x (logit,, E [0, 1].
(Prediction)
Here 0 A (e1, eo, {Pn] ..... b, denotes the full list of the transformer parameters. is the embed- ding dimension, E, and eg in Rd are the token-embeddiings corresponding to I, = 1 and Ta 0 respectively, and p,, is the (trainable) positional encoding. We have matrices Wo E Rd and Wv E Rmndd and the attention weights att,,; (0. 1) are computed using the query and key matri- ces ($ A). W₂ E ıddrr and W1 E Rrxd are the weight matrices in the FF layer, whereas E IR and b E IR are the weight and bias parameters for the linear layers. For a multi-layer transformer, we apply the successive attention and feed-fome forward layers multiple times before the final linear layer. Finally, we compute the probability for the symbol using the logits: fo (x²) A Pg (In+1 = xT) (logit,,.) E [o, 1]. Note that single symbol probability suffices as the vocabulary binary.
Loss.
The parameters I are trained using the next-token prediction loss between the predicted probability fo (x²)) and the corresponding ground truth symbol Z441 across all the positions n E [N]:
L A N
Σ NE[N]
B [u4t1 log log.o((z) fo())) + (1 2n+1) 1) logg1 - fo()))),,
(1)
where the expectation is over the data distribution of the sequence { J. In practice, it is replaced by the empirical averages across the sequences { sampled from the corpus, with stochastic optimizers like SGD or Adam (Kingma and Ba, 2015) used to update the model parameters.
1N
2.2 MARKOV CHAINS
We model the input as a first-order Markov chain, i.e. a Markov chain with (order) memory m 1. For these processes, the next state is influenced only by the current state and none of the past:
Py 2 IP (In+1 æn = i) P (xnt1 In = x =1--1)
3

---

Published as a conference paper at ICLR 2025
fa(x))
fa()))
fa(a)")
a())
((())
p
1-
1-
logit,
logit
logity
Linear
Linear
Linear
FF
FF
L
1/
y/
N
Attention
1 1-p
(In+1 Endn21 ~ P(p. _
9 1 1
Pij = PP=nt(1 =j In = =i), i.j E {0,1)
Embedding
Embedding
Embedding
NN E (0.1)
(a) The transformer model with hinary input data: (b) State transition diagram and Markov kernel for
for each zf, next-bit prediction probability is first-order Markov chain ((r.g) with flipping proba- fa())) = Po (Io41 = )
bilities Pot and Pio
Figure 2: Analysis of transformers via Markov chains.
any i1.... (1...,,,--1,,11 E Xn 2 Here the Markov kernel P= (Pis) governs the transition
dynamics of the process: if (m) E [0, 1] 1]12" denotes the probability law of In at time n, then
n (m+1) (("). Of particular interest to us in this paper is the kernel P(p. q) _ [1-p. p; 4, 1-q] on the binary state space with the switching probabilities Poi = p and Pio = q, for p. q E (0, 1). Fig. 2b illustrates the state transition diagram for this kernel. Here we refer to the sum p - as the switching factor. We denote a first-order binary Markov chain (Ix)x21 with the transition kernel (p,, 9) and starting with an initial law (()) as (In)n21 ~ (T)) (p). When the initial distribution is understood from context, we simply write (In+1 kn)n21 P(p, q). For this process, the entropy rate equals H (Zn+11 (z,.) = p+g hg)) + ph())), which is independent of n.
Stationary distribution. A stationary distribution of Markov chain is distribution π on X that invariant to the transition dynamics, i.e. (")) = ग. then we have (n+1) = =P m and consequently, (1)) π for all 2 n. Also referred to the steady-state distribution, its existence and uniqueness can be guaranteed under fairly general conditions (Norris, 1997), and in particular for P(p.g) g) when p. q * 0, For P(p. q). the stationary distribution given by π (p, v) A (To:: MI = p+q (q.p). The higher the flipping probability q, the higher the likelihood for the chain to in the state 0 at the steady state. Similarly for the state We can verify that æ indeed satisfies nt
For brevity, we drop the dependence on (p, q) and simply write π and P when the parameters are clear from context.
FRAMEWORK: TRANSFORMERS VIA MARKOV CHAINS
We present our mathematical framework for a principled analysis of transformers via Markov chains. In this paper we focus first-order binary Markovian data and single-layer transformers though our framework readily generalizes to higher orders and deeper architectures (Sec. 4.4), and multi-state Markov chains (Sec. 4.3).
Data. We assume that the vocabulary X' = {0, 1} and the input data {an N ~ (T g). P(p. q)). for some fixed sequence length N > 1 and (p. q) E (0. 1)2. Recall that p is the switching factor. The parameters and provide a tractable mechanism to control the input data, which plays a crucial role in transformer learning.
4

---

Published as a conference paper at ICLR 2025
Model. We consider single-layer transformer with single-head attention, without layer norm. As the input is binary, the Embedding layer can be simplified to
nn In e1 (1 2n) eo P., 2n, + Pn
(Uni-embeddingg)
where e a G1 Ey is the embedding vector and Pn - eo + Pn is the new positional encoding. Note
that Fn E {0. } and hence the embedding is either e + pn or just P., depending on 2n. The other layers are the same as Sec. 2.1:
In E 10
Uni-embedding
En
Attention
yn
FF
Zn
Linear logit,
Prediction
fa (z())
(2)
Let 0 (e, {Pr "nn..nl......... a) E Rº denote the joint list of the parameters from all the layers, with D being the total dimensionality. training large language models, it is a common practice to tie the embedding and linear layer weights, i.e. In this case, the list of all parameters, = (e = a, (p, IN pp........, E R-d since a is no longer a free parameter. We analyze both weight-tiee and general cases.
referred to as weight tying (Press and Wolf, 201 7).
Loss. We consider the cross-entropy loss L from Eq. (1).
Objective. Towards understanding the phenomenon in Fig. 1, we utilize the aforementioned frame- work to study single-layer transformers. In particular, our objective is to address the following question:
Can we characterize the loss landscape of singe-layer transformers when the input is Markovian?
To build intuition about the loss surface, we first examine its global minima and then provide detailed characterization of the loss landscape, focusing on local optima, in Sec.
3.1 SINGLE-LAYER TRANSFORMERS: GLOBAL MINIMA
Since the loss L in Eq. (1) is the cross-entropy loss, it achieves its minimum when the predictive probability matches the Markov kernel (Lemma 1): fa(")) IP (z,-1 x.)). In other words, this occurs when the transformer outputs the correct transition probabilities. This raises natural question: can a single-layer transformer exactly represent a first-order Markov chain? Intuitively speaking, this seems plausible since the transformer, even with access the full past information Ii at each E [N], can rely solely on the current symbol 2,, (Sec. 2). The following result confirms this intuition, showing that such a realization is indeed a global minimum for the loss L(-):
Theorem (Global minimum). Let the input sequence be { (T q). P(p. 4)) for some fixed (p, 9) E (0, 1)² and e I D-d be the transformer parameters for weight-tied case. Then for all (p, q). there exists Q, E RRDd with explicit construction such that is global minimum for the population loss L(-) in Eq. (1) and its prediction matches the Markov kernel, i.e.
(i L(0) > L(0_) for all 0 E R-d and (ii) Po. (In+1
xT) (Ia+1 wn). the Markov kernel or the bigram.
Further, O, satisfies:
(iii) L(O.)= H ((nnt1lzn)) the entropy rate of the Markov chain.
(iv) VI(0.) 0, i.e. 0, is stationary point.
In addition, the same result holds for the non-weight-tied case when the parameters are in R...
Remark 1. In fact, there exist many such global minima as highlighted in the proof ($ B).
Proof sketch. The key idea here is show that any satisfying (x P (In+1 global minimum and is stationary point with the loss being the entropy rate (Lemmas. I and 2). To construct such a f, we utilize the fact the Markov kernel only a function of r, and thus we can ignore the past information in the Attention layer using only the skip. We defer the full proof to $ B. L
xn,) is
Empirical evidence for learning the Markov kernel. As demonstrated the proof above, a canonical way to realize the Markov kernel by the single-layer transformer is to rely only on the current
5

---

Published as a conference paper at ICLR 2025
symbol in and ignore the past in the Attention layer. We now empirically confirm this fact. For our experiments, we use the single-layer transformer (Table 1) and report the results averaged across 5 runs and corresponding to the best set of haper-parameterss after grid search (Table 2). In particular, for p = 0.2, q = 0.3, and = 4, we generate sequences (In J ~ (T q). P(r. 4)) of length N 1024 and train the transformer parameters 0 (weight-tied) entropy loss in Eq. (1). At inference, we interpret the attention layer and observe that the relative magnitude of the attention contribution to the final attention output y, negligible, i.e. the ratio | o ExE(v) attn,: Wv 2;11/Ily,11 ~ 0.01. Hence, the attention contribution can be neglected compared to the skip- connecton,,,, i.c. yn, A In. Using this approximation, in $ D we derive formula for the final predicted probability Jo(=)) as it is learnt by the network. This formula reveals interesting insights about the learnt parameters of the transformer:
minimize the cross-
Constant embeddings. The positional embedding Pn constant across n, i.c. it is independent of the sequence position, reflecting the fact that it has learnt to capture the homogeneity of the Markovian chain just from the data.
Low-rank weights. The weight matrices are all approximately rank-one; while is not fully clear why the training algorithm converges to low rank solutions, they do indeed provide a canonical and simple way realize the Markov kernel, illustrated in D.
Further, we show in & D that plugging in the numerical values obtained from the average of five runs, the probability given by our formula matches the theory, i.e. the model learns correctly output the Markov kernel probabilities. Indeed, Fig. 3b shows that the test loss of the model converges to the theoretical global minimum (Thm. 1), the entropy rate of the source, corresponding to the bigram, when p 0.2 and = 0.3 (p + 1). We likewise observe a similar phenomenon without weight tying. For the prediction probability, we focus on the zero positions n n& such that Fra = 0. Fig. 3d shows that irrespective of the index it and the past x² -1 if the current bit Fng is O, the model correctly predicts the probability for the next bit In being 1, which equals p theoretically (Fig. 2b). More precisely, fo(xr"*1 In, 0) for all == and k. in line with property (ii) of Thm. 1. similar conclusion holds with In, = and prediction probability q. This indicates that the model has learned to recognize the data first-order Markovian, relying solely 2n, to predict In+11
While Thm. I and above empirical results highlight the presence of global minima on the loss surface, they does not address local optima, as empirically shown in Fig. 1. We precisely address this in the next section and analyze the loss landscape in terms of the local optima.
4 SINGLE-LAYER TRANSFORMERS: LOCAL OPTIMA
In this section we present our main results about the loss landscape of single-layer transformers in terms of local optima. In particular, we prove the existence of bad local minima and saddle points on the loss surface (Thms. and 3), in addition to the global minima discussed above (Thm. 1). Interestingly, the presence of these local optima is influenced by two key factors: switching factor of the Markov chain and the weight rying of the transformer, highlighting the intricate interplay between the input data and the model architecture. First, we present the results for the weight tying scenario.
4.1 WEIGHT TYING: BAD LOCAL MINIMA
When the embedding and linear layers are tied, i.e. E = a, our analysis reveals the following surprising fact: if the switching factor p + q is greater than one, there exist bad local minima 0x E Tdd where the prediction probability (-) is the marginal stationary distribution π (unigram), disregarding the past and the present information (Thm. 2 and Fig. 3c). Now we state the result formally. Let L. A I(O.) denote the global minimal loss from Thm. 1.
Theorem 2 (Bad local minimum). Let the input sequence be IN ~ (π g), P(p.4)) for a fixed (p. q) (0. 1)2 and the transformer parameters be weight-tied. 1f + q > 1, there exists explicit 0x E TR D-d such that it is a bad local minimum for the loss L(.), i.e.
there exists neighborhood (60,,, (T, with 0 such that L(0) L(0x) for all 0 B(O,.)). with L(0x), > La.
Further, Or satisfies:

---

Published as a conference paper at ICLR 2025
0.80
0.70
Without weight-tying
Without weight-tying
With weight-tying
With weight-tying
0.75
Bigram distribution loss
0.65
Bigram loss
Unigram distribution loss
Test loss
Test loss
0.70
0.60
0.65
0.55
0.60
0.50
20
80
100
20
40
100
Iteration
Iteration
(a) Test loss (p q 1)
(b) Test (p + 1)
0.30
Without weight-tying
Without weight-lying
With weight-tying
With weight-tying nrightttying.......
Prediction fo(a)")
Transition probability
Prediction fo(x?")
0.25
Transition probability p
Stationary probability m
0.20
0.15
0.10
20
80
100
20
40
100
Index l : Tng
Index k Ink 0
(c) Predicted probability
1)
(d) Predicted probability (p + < 1)
Figure 3: Effect of weight tying on test loss and predicted probabilities fe(x"") for zero indices (ng 100
]== such that Ini = 0. For (a),(c): p 0.5, 0.8. With weight tying. the loss converges to local minimum, and the predicted probability is Tı p/(p + q)- Without weight tying, we predict the correct probability p and converge global minimum. For (b). p 0.2, q = 0.3. The test loss always converges to global minimum, and the predicted probability is p.
(ii) Po. (In+1 = xT)) = (kn+1 1) nı, the marginal distribution or the unigram. (iii) L(O,) = H(*n+1) = H(T), the entropy of the marginal.
(iv) (L(0)) 0 i.e. 0, is stationary point.
H(Ix+1) and L. L(#x) L, = (I (I = [(x;; 1n-1) 2 O, where [( 4141) is the mutual information between J., and Fn+11 (Cover and Thomas, 2006). It equals zero if and only if x,, and Int1 are independent, which happens for + = (since (In+1 = x.) = Ix (1
Remark 2. Since L(Ox)
H HI (xx), the optimality gap
q) p).
Proof sketch. The main idea behind constructing 0 that we e a in the Linear layer, the model ignores the inputs all together and outputs constant probability, and in particular by choosing the bias appropriately, i.e. fe, (xx²) = गi for all of n. For this O, it's easy to show that L(#x)
H TT) and that it's stationary point. Further show that the Hessian e follows the
H.a 0' where Ha > 0 when + q > 1, and that it implies the local minimality of Oz.
structure 0
We defer the full proof to B.4.
Empirical evidence for bad local minima. The proof of Thm. 2 above highlights that when the linear weight a is zero in 0, serves as a bad local minimum. While this might seem as a theoretical anamoly, empirically confirm that it is not. Specifically, use the same weight-tied setting
7

---

Published as a conference paper at ICLR 2025
Prediction fo (1r
Eng.
Prediction fo (2r (t J.n,
pt
0.201
6334
0.532
6,602
74747
0,207
6/431
0.530
6.681
0.200
5,400
0.502
0.560
0.201
5.400
0.503
6.502
0.788
0.5
0.199
6.399
0.499
0.611
0,525
0.5-
0.204
5.396
0.591
0.500
0.22
0.6
0.203
5.400
0.661
0,581
0.616
0.147
5.38
0.535
0,547
0.79
0.0
0,291
0:547
0.402
0.464
0.726
0.0
0.292
5.518
0.435
0.36
0.522
0.5
(a) With weight-tying
(b) Without weight-tying
Figure 4: Average of predicted probabilities across 5 runs for different values of p and q, with and without weight-tying. In the former case, there is clear demarcation beteen the cases where < 1 and those where p + > 1. For + < 1, all runs accurately predict the correct conditional probability. For p + q 1, some of the runs predict the stationary probability instead, causing the average to diverge from the correct p. In the latter case, the model always predicts the correct probability for all p and q.
as before but with the flipping probabilities p 0.5 and y 0.8 instead, we observe that the magnitude of the vector a is approximately 0.01 whereas that of nn is 4.0 in the Linear layer. Thus, (a, 2,,) 0.04, while the bias term b x -0.4 implying ((a, x,,) + 6) as ob). Hence the final prediction returned by the network only depends on the bias of the linear layer, totally independent of the input sequence x:\. Fig. 3e illustrates this fact and shows that the model always predicts the stationary probability क at all positions in the sequence, independent of the input. Here we plot for the indices & such that Tna for all z1, thus verifying property (ii) of Thm. 2. Similarly, Fig. 3a demonstrates that the model test loss converges to the entropy of the stationary distribution HTT),, the unigram loss, instead of the global minimum L, given by the entropy rate of the source (Thm. 2 (iii) ). Fig. 4 illustrates for wider range of (p. 4) E (0, 1)2.
but we observe the same phenomenon for Ing i.e. [o(x) m
Interpreting global and local minima. Thm. 2 should be interpreted in the light that it guarantees the existence of bad local minima only for + q > 1 (in sync with the experiments, Fig. 4 While there could exist such minima even when + q < 1, we empirically observe that the model always converges to the global minimum this scenario (Fig. 3b). Indeed, Makkuva et al. (2024) shows that local minima also exist for p q < 1, however, the standard Gaussian initialization around zero falls outside their basin of attraction, and hence gradient based methods escape them. Likewise, while Thm. guarantees the existence of global minima for all (p. q) and in particular for p + q > empirically the model often converges to bad local minima as highlighted above (Fig. 3a).
4.2 WITHOUT WEIGHT TYING: SADDLE POINTS
Now we let the token-embedding e Rd and the linear weight a E 12d be independent parameters. Interestingly, in this scenario, the earlier local minimum O becomes a saddle point.
Theorem 3 (Saddle point), Consider the same setting as in Thm. 2 and for p + q >> 1, let Oxr = .......by. bx) E RD-d be the corresponding bad local minimum for the loss T(.)) in the weight-tied scenario. Then its extension 0 A (0 (m:: ax) E RDD saddle point for L(-) in IRP in the non-weight-tiee ried case. Further, 0 satisfies the same properties (ii)-(iv) as in Thm. 2.
(ex = a.,...,
Empirical evidence and interpretation. In view of the theoretical results above, removing weight tying is possibly beneficial: the bad local minimum in the weight-tied case for p q > 1 suddenly becomes a saddle point when the weight tying is removed. We observe a similar phenomenon empirically (Fig. 4): shown in Fig. 3a, when not weight-tied, the model's 's test loss converges to
8

---

Published as a conference paper at ICLR 2025
the entropy rate of the source when p + q 1, in contrast the weight-tied case, possibly escaping the saddle point (Thm. 3). The fact that the model eventually learns the correct Markovian kernel is further demonstrated by the red curve in Fig. 3c. Figs. 3b and 3a together highlight that the model always (empirically) converges to the global minimum in the non-weight-tied case irrespective of the switching factor p + q.
Key insights. Together, our theoretical and empirical results highlight that when the switching p u > 1, the weight-tied model can get stuck at bad local minima corresponding to the unigram. In contrast, the non-weight-tied model can potentially escape saddle points to converge to the global minima, corresponding to the bigram (Markov kernel). This explains the phenomenon in Fig. 1 for the single-layer transformer, where p = 0.5 and q = 0.8. When + q < 1, we empirically observe that the model always converges to a global minimum irrespective of weight-tying.
4.3 MULTI-STATE MARKOV CHAINS
Our framework focusing on the binary setting X' = {0, 1} above, readily generalizes to the multi-state setting X' {0, 1, 1), where S is the state (vocabulary) size. Let the input (In)n21 P(p) be a first-order Markov chain, where either Tx+1 = In with probability 1 or Tn+1 switches to a different state in the vocabulary uniformly randomly with probability p/(S 1), for some p E 1). The symmetric kernel P()) generalizes that of the binary Markovian case, P(p.g) in Sec. 2.2, when p = q. Fig. illustrates this for the five-state setting S 5. For the transformer, we use the same architecture as in Sec. 2.1 except for the fact that we now have S token embeddings in the embedding and linear layers.
With weight-ty ng Witter ue gitilliingg Skyion lo%s
1.600
weigtttiing 3igram loss
1
1.575
Tee
1.0
p/t
0.5
1000
2000
3002
4000
1000
2000
3500
4000
p/t
p/
(b) Loss for 0.1
(d) Loss for p 0.6
1.6
swight-yy ng
1:61
Wittor: we ghr-liing Skgan bos
1.60
Test
weigh+
Wenoot weighssyyng Un gram loso Nigram loss.
P
p
1.2
1:35
1.0
1:57
4000
2200
4000
(a) Markov chain diagram
(c) Loss for p 0.3
(e) Loss for p 0.9
Figure 5: Effect of weight-tying for the multi-state Markov chain with S states. (a) shows the symmetric multi-state Markov data model used, while (b), (c), (d) and (e) show the test loss for different values of p. Similar to the binary case, when p is large enough, a 1-layer transformer with weight-tying gets stuck in local minima, whereas it escapes to global minima without.
Results. When a single-layer transformer is trained on a multi-state Markov chain with S = 5 for various values of p E {0.1, 0.3, 0.6, 0.9). Fig. 5 (analogous to Fig. 3) illustrates that models with weight tying tend to get stuck in local minima (unigram). whereas those without converge to global minima (bigram). As discussed in Sec. 4.1 for binary case, the switching condition p + q > 1 is equivalent to p > 0.5 when p = g. where the model empirically gets trapped in local minima. Interestingly, for the five-state case here, we observe convergence to global minima even at p 0.6 > 0.5, suggesting that the switching threshold exceeds 0.5 in this scenario. In the binary case, this threshold of 0.5 was derived theoretically via Hessian computation (Eq. (13)). In a similar vein, the proofs of Thm. and Thm. 3 which extend to the multi-state setting, yield comparable condition via Hessian analysis. Together, these findings highlight that similar theoretical results could be established highlighting the role of weight tying for multi-state scenarios as well, in congruence with
9

---

Published as a conference paper at ICLR 2025
above empirical findings. Obtaining an explicit mathematical expression for the switching threshold, however, is an exciting direction for future research and beyond the current scope of this paper.
4.4 DOES DEPTH HELP ESCAPE LOCAL MINIMA?
For single-layer transformers, the aforementioned results highlight the significance of the switching factor and the weight tying on the loss landscape. In stark contrast, we empirically observe that for transformers of depth 2 and beyond. the loss curve eventually reaches the global minimum regardless of these factors, as highlighted in Fig. 1. Interestingly, we observe during training that it first reaches a plateau at loss value corresponding to ((T)) and after a few additional iterations, it further drops down until it reaches the global minimum (the entropy rate). This suggests that, while there could still be local minima, increasing the number of layers positively affects the loss curvature at the bad local minima in manner making it casier to escape and reach the global minimum. In the context of beed-forward neural networks, depth of the architecture has been shown to play a major role in terms of the representation power and learning capabilities (Telgarsky, 2016; Arora et al., 2019). Given our empirical observations, a similar analysis for transformers that demonstrates the benefits of depth is intriguing direction for future research.
RELATED WORK
There is tremendous interest and active research in understanding transformer models from various perspectives (Weiss ct al., 2021; Giannou et al., 2023; Oymak et al., 2023; Li et al., 2023a; Fu et al., 2023; Noci et al., 2023; Tian et al., 2023). Yun et al. (2020); Pérez al. (2021); Wei et al. (2022); Malach (2023); Jiang and Li (2023) demonstrate the representation capabilities of transformers and show properties such as universal approximation and Turing-completeness. Another line of inquiry (Elhage et al., 2021; Snell et al., 2021; Wang et al., 2023; Geva et al., 2023) is mechanistic interpretability. i.e. enginerenggnn transformer operations on specific synthetic tasks (e.g.. matrix inversion and eigenvalue decomposition in Charton (2022). modular addition in Nanda et al. (2023)) to understand the transformer components but they usually lack theoretical guarantees, as opposed to ours. Li et al. (2023c) studies how transformers learn semantic structures across words while we are interested in how they learn sequentiality in input data. Tarzanagh et al. (2023abb) take an optimization theooretic perspective study training dynamics and characterize implicit biases in transformer models trained with gradient descent. In contrast, we characterize the local and global minima of the loss landscape of these models under sequential input data.
On the other hand, Chen et al. (2024); Bietti et al. (2023); Dong et al. (2023); Akyürek et al. (2023); Von Oswald et al. (2023); Xie et al. (2021); Bai et al. (2023); Li et al. (2023b); Garg et al. (2022) study in-context learning capabilities of transformers. Along this line, our work is closely related to Nichani et al. (2024); Bietti et al. (2023); Edelman et al. (2024) in the Markovian setup, however, while they focus on in-context learning, which is apparent for two or more layers, we investigate the loss landscape of single-layer models. Chung et al. (2021) provide empirical evidence to suggest that weight tying has drawbacks in encoder-only models, which is in line with our observations that removing weight tying is beneficial in decoder-only models with Markovian input data. More recently, Rajaraman et al. (2024b) study the effect of tokenization on learning Markov chains. Ildiz et al. (2024) show an equivalence between the attention mechanism and Markov models, whereas we characterize the loss landscape of attention-based transformers when the input is Markovian.
6 CONCLUSION AND OPEN QUESTIONS
In this work, we provide a novel framework for systematic theoretical and empirical study of the sequential modeling capabilities of transformers through Markov chains. Leveraging this framework, we theoretically characterize the loss landscape of single-layer transformers and show the existence of global minima and bad local minima contingent upon the specific data characteristics and the transformer architecture, and independently verify them by experiments. We further reveal interesting insights for deeper architectures. We believe our framework provides a new avenue for a principled study of transformers. In particular, some interesting open questions in this direction include estab- lishing similar results for higher-order Markov chains and deeper transformer models, showcasing the interplay between the model depth and Markovian order.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
The work was supported in part by the Swiss National Science Foundation under Grant 200364.
REFERENCES
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? Investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with co-context algorithm selection. In Workshop on Efficient Systems for Foundation Models ICMI-2023, 2023.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. In Thirfy-seventh Conference on Neural Information Processing Systems, 2023.
François Charton. What is my math transformer doing? Three results on interpretability and generalization, 2022. URL /tttts::/arxiv...tt
00170.
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.
Hyung Won Chung, Thibault Fevry, Henry Tsai, Melvin Johnson, and Sebastian Ruder. Rethinking embedding coupling in pre-trained language models. In International Conference Learning Representations, 2021.
Thomas M Cover and Joy A Thomas. Elements of information theory. John Wiley & Sons, 2nd edition, 2006.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT; Pre-training of deep bidirectional transformers for language understanding, 2018. URL /tttss//s//axxiv org/ abs/ abs/1810 04805.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, Lei Li, and Zhifang Sui. A survey on in-context learning, 2023. URL https: / http///xxvvvo///// org/ abs/ /oss/2301.. 00234.
Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning Markov chains, 2024.
Nelson Elhage. Neel Nanda. Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Comerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfiel-DDodss, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. URL https: /tt::://rraaaa www.aaassooorrrroooe www.orreeeerrrriteeeeeeeeeeeeeee pub/ pubb2022 2021/ www.wwork/nndeee html.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can single attention layer learn? A study through the random features lens. seirty- Conference on Neural Information Processing Systems, 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598 2022.
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in reto-regressive language models. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
11

---

Published as a conference paper at ICLR 2025
Angeliki Giannou, Shashank Rajput, Jy- Yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers, In Proceedings of the 40th International Conference Machine Learning, pages 11398-11442, 23-29 Jul 2023.
Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge university press, 2012.
M. Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self-attention to Markov models: Unveiling the dynamics of generative transformers, 2024.
Haotian Jiang and Qianxiao Li. Approximation theory of transformer networks for sequence modeling. 2023. URL https: / arxiv org/ abs 2000 18475.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference on Learning Representations, 2023a.
Yingcong Li, M. Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algo- rithms: generalization and stability in in-context learning. In Proceedings of the 40th International Conference on Machine Learning, 2023b.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How transformers learn topic structure: towards mechanistic understanding. In Proceedings of the 40th International Conference Machine Learning, 2023c.
Ashok Vardhan Makkuva, Marco Bondaschi, Chanakya Ekbote, Adway Girish, Alliot Nagle, Hyeji Kim, and Michael Gastpar. Local to global: Learning dynamics and effect of initialization for transformers. arXiv preprint arXiv:2406.0300 03072, 2024.
Eran Malach. Auto-regressive next-token predictors are universal learners, 2023.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023.
Eshaan Nichani, Alex Damian, and Jason Lee. How transformers learn causal structure with gradient descent. arXiv preprint XXXx:2402.14735, 2024.
Lorenzo Noci, Chuning Li. Mufan Bill Li. Bobby Hc, Thomas Hofmann, Chris J. Maddison, and Daniel M. Roy. The shaped transformer: Attention models in the infinite depth-and limit. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
J. R. Norris. Markov Chains. Cambridge Series in Statistical and Probabilistic Mathematics.
Cambridge University Press, 1997.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint ar arXiv:2209.11 220.1199,9. 2022.
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 Machine Learning, 2023.
Mattco Pagliardini. GPT-2 modular codebase implementation. 2023.
Ofir Press and Lior Wolf. Using the output embedding to improve language models. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers, pages 157-163, April 2017. URL https: / a/ aeclannhology org/ E17 20255.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic, Attention is Turing- Journal of Machine Learning Research, 22(75):1-33 2021.
12

---

Published as a conference paper at ICLR 2025
Alec Radford and Karthik Narasimhan. Improving language understanding by generative pre-training. 2018. URL https: sttp::/:s/.pp..aaaaaaataaa.aaaaa
aewant icscholla.... org/CorpusID:: :ttpeco......oooccc.rrr.a::..... 493133445
Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vardhan
Makkuva. Transformers on Markov data: Constant depth suffices, arXiv preprint arXiv:2407.1766 2024a.
Nived Rajaraman, Jiantao Jiao, and Kannan Ramchandran. Toward theory of tokenization in LLMs, 2024b.
Charlic Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns, 2021. URL https:: http://awww....aa///////////se// wtt./oo/////./////// 07601.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023a.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. In Thirty-seventh Conference Neural Information Processing Systems, 2023b.
Matus Telgarsky. Benefits of depth in neural networks. In Conference on learning theory, pages 1517-1539. PMLR, 2016.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Shaolei Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Conference on Parsimony and Learning (Recent Spotlight Track), 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998-6008, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174, 2023.
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: study approximating Turing machines with transformers. In Advances in Neural Information Processing Systems, volume 35, pages 12071-12083, 2022.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 1 11800-10900 2021.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit Bayesian inference. arXiv preprint urXiv:2111.0208/ 2021. URL https: //a farxiiv. orgg abs/21 /211 02080.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are trans- formers universal approximators of eqquencetto-squeene functions? In International Conference Learning Representations, 2020.
13

---

Published as a conference paper at ICLR 2025
Organization. The appendix is organized follows:
App. A details the transformer architecture, especially that of the attention mechanism.
App. B provides the proofs for our theoretical results on first-order Markov chains.
App. contains additional experimental details and results for first-order Markov chains.
A THE TRANSFORMER ARCHITECTURE
We describe the Transformer architecture from Sec. 2.1 detail, using the embedding layer simplica- tion from Sec. 3:
Zn In e + Pr E RR
(Uni (lni-embedding)
Yn an + Wo W Σ
atln, W v as E Rd
(Attention)
iE|n
Zn = yn + W ReReLL 1y,) E R4,
(FF)
logit, (a, 2n) +
R,
(Linear)
fa(x)) e Po (In+1 = 1 ( o (logit,
(Prediction)
E(O.1]
(i) Embedding: The discrete tokens T, and F, = 0 are mapped the eoken-embeddings and 0 Rd respectively, where d is the embedding dimension. The positional embedding p,, R encodes the positional information (varies with n.). The sum of these two embeddings constitutes the final input embedding 2n E IRd.
The attention layer can be viewed mappping query and of key-value pairs to an output, which are all vectors Vaswani et al. (2017). That is, on top of the skip-connection n, the output y, E Id is computed as a weighted sum of the values V. A W v I., The weight assigned to each value, att,,i, is computed by a compatibility function the query vector In & We In and the corresponding key vectors k; & W 2; for all i [n]. More precisely, att,,; A softmax( k1),..., (qn. ka))/ W K.Q.V m the respective key, query, and value matrices, and Wo Ra is the projection matrix. For multi-headed attention, the same operation is performed on multiple parallel heads, whose outputs are additively combined.
(ii) Attention:
(iii) Fordardrward (FF): The FF transformation consists of skip-connection and single-hidden layer with ReLU activation and weight matrices W, E RRxx, and W, E RRrxd The FF layer is applied liken-wise on each y, E TR. to output 2, Rd with the same dimensionality.
(iv) Linear: The linear layer transforms the final output embedding nn to a scalar logit, E R, with the weight parameter Rd and the bias b IR
(v) Prediction: The sigmoid activation finally converts the scalar logits to probabilities for the next- token prediction. Since the vocabulary has only symbols, it suffices compute the probability for the symbol 1: (x+)) & Pa (In+1 1 =f)) a (logit,,) E [0, 1]. More generally, the logits of the same dimensionality as the vocabulary and are converted to the prediction probabilities using softmax layer, which simplifies the sigmoid for the binary case. Likewise, there are as many token-enan eekeddembeddings as the words in the vocabulary and several layers of multi-headed attention and FF operations are applied alternatively on the input embeddings to compute the final logits.
Finally, The Transformer parameters 0 A (e, I ... ., RD are trained via the cross- entropy loss on the next-token prediction:
L _
Σ R
E a4 [n+t1 log fo(z)) + (1 2n+1) log(1 fo(x)))
(3)
R
14

---

Published as a conference paper at ICLR 2025
B PROOFS OF SEC. 4
We now present our proofs for the technical results in Sec. 4. Towards this, we first establish two useful lemmas on the loss function L(-) and the corresponding gradient computation. Let
R D-d = (e. (P,, I............................... a) E R be the list of parameters in the non-weight-tied case and O (e a, (p,, 1
weight-tied case. With a slight abuse of notation, by w we mean a specific parameter among the set {e, PI. ...P....PPPT. ................................ .N...... a). Since the weight-tied scenario is a special case of the non-weight-tied one with e = a, we directly present the results for the general non-weight wen--eiihh-tieee case, but both lemmas hold for 0 PD- well. First start with the loss function.
Lemma (Loss as KL divergence). Let the input sequence be {-n]n=1 (T(p. q), P(p, (oogg)) for some fixed (p, E 1)², 0 (e, (p.).=.......................... a) ED be the full list of the transformer parameters, and L(0) be the corresponding cross-entropy loss in Eq. (1). Then the loss function L(-) equivalent the KL. divergence between the Markov kernel and the predicted distribution, i.e.
((O)
Σ NE(N]
Exp (P (In+1 = xxx)) Po (En+1 = - + + ((I+11/T)) (4)
N
where (K || Q) is the KL divergence between two distributions P and Q, and H (In+1 is the entropy rate of the Markov chain.
Remark 3. Consequently, Eq. (4) highlights that any parameter O with the predicted probability
fa("") = P (In+1 = Φn) is a global minimum for the loss L as DKL(- (- -) 2 0 Cover and Thomas (2006). We utilize this fact in the proof of Thm. below.
Proof. We defer the proof to $ B.6.
Lemma 2 (Gradient computation). Consider the same data and parameter setting as in Lemma and (() be the cross-entropy loss in Eq. (1). Then for any parameter w E
L
Σ NE[N]
E B [( -fo fz())) V (a nn +b)
(5)
Σ WE[N]
Egp [(P (#n+1= xn) - (o(x))) V. (a zn + +b]]
= N
Remark 4. Eq. (5) highlights that any parameter 0 with the predicted probability fa (xf²)
P (In+1 =1 xn) is also a stationary point for the loss L. We utilize this fact in the proof of
Thm. 1 below.
Proof. We defer the proof to $ B.7.
We now detail the proofs of theorems in Sec. 4. We prove the global minimum result in Thm. 1 in two parts, separately for the cases when p + y < and p + q > For both these cases, first we consider weight-tying and in App. B.3, the non-tied case.
B.1 PROOF OF THM. FOR p + q S 1, WEIGHT-TYING
Proof. We assume that p q 1 and that we use weight tying, i.e. the list of parameters 0 (e = a, (P, .
E RD- Thus view Lemma and Lemma 2, follows that any satisfying fo fz(x)) P (In+1 = 1 zw,) is global minimum with loss equalling the entropy rate, and is a stationary point. Hence it suffices to construct such a
To build our intuition towards designing 0., recall that the Markov kernel IP (I.+1 = 1 xx) can be succintly written as P (In+1 = xx,) = (1-p-g)) To ensure that fo(2f) = x, (1-p it suffices for the transformer utilize only the information from the current symbol F, and ignore the past z -- In view of the transformer architecture $ A, a natural way to realize this is to let W o =
+p,
and Wa in Attention and FF respectively. This implies that zn Yn Dn Fr e + P.. Hence
15

---

Published as a conference paper at ICLR 2025
the logits given by logit,, (e, xn. + nn |e]l + (e, Pn)) b. Since fo(x?) (logit,,. and it equals the Markov kernel, we have that
ologit,,,) a ||el (e, pa) + b) nn - p- 9) p.
Rewriting,
- +p
In eel + (e. Pa) log
In E {0.1}.
(1 (1-p)(1 xn) + q
Substituting 0 and Zn, = 1, we further simiplify to
p
(e, P,) + log
1 -
1 q
|||| + (e. p,) + log
Subtracting both the equations obtain that global minimum 0 should satisfy
(1 - p1 9)
||e|| = log
pq
(6)
Pm) + log 1 p
Note the above choice of e is well-defined since (1-p)(1-g) m1 when p + < and hence
log (1-p)(1-)) > 0. While there exist infinitely many solutions for (e.Pn. b) satisfying Eq. (6), a
po
canonical such solution for the global minimum 0 = 0, is
(1
(1 - ②)
0. = =====1 log
[P. = -) Wo = 0, W W..VV, W2 0, W,, b = log
pq
1 p
where 1 E RD-d denotes the all-one vector, the position embeddings P., are set to zero, W K.Q.V e Rmxdd and 1 E grxd can be set to any arbitrary value. This concludes the explicit construction of 0. and the proof.
B.2 PROOF THM. FOR p + q > 1, WEIGHT-TYING
Proof. We similar idea as in the proof for the p + q S 1 case by constructing 0 E RD- satisfying fo(xT) = IP (In+1 = 1 xx) = In x1 - q) + p. However, in this case we need to use the ReLU component of the FF mechanism unlike the earlier case where we set Wa = 0. Now we start with constructing 0,.
Let the embedding denotes the all-one vector. Thus 2, a,, with on + when In and an - when In = O. Now let W = 0 in the Attention layer. Hence y, = Cn a, For the FF layer, let W and Wa be such that (to be determined later)
1 and the positional encoding P., = for all > 1, where 1 Rd
W₂ ReRRRLLLWW,, y,) = =,,1,
and hence
2, = Fn + W, (eLU(W, vn)) = 0, + B,, = (a,, + B.
Thus the logits are given by logit,, = o( 2n) + b) = a(d(ar, + 3a) + b). Since fo(x)) (logitt,.) = P (In+1 = kn), we have that
ologitt,)) = (d(an, + Br) + b) = In (1 - p - q) + p, on (0,1).
Rewriting,
In In (1 - p 9) + -p (1
d( + 3r) + log
In E (0,
(1 - xn,) + grn
16

---

Published as a conference paper at ICLR 2025
Substituting 2n and 2n. 1, and denoting corresponding I's by B1 and 8o (with slight abuse of notation), we further simiplify to
1
+
90
+ = log
(8)
2
p
1
9
d
+
91
+ b = log
2
Subtracting both the above equations we obtain
(1 p1 9)
d(1 + 31 Bo) log
pg
(9)
<0 when p+tq211
Now it suffices to find 31 and Bo satisfying Eq. (9). Recall that Bn. obeys Wa ReLU(W, v,) 8,1. Let W = 11 and W = - for some w E Since yn az 1, we have
11 ReLU(w 11 T a = 1.
Simplifying.
an = = 1
Bn = -w ReLUUaa11)
Thus 80 0 (corresponding to In them in Eq. (9), we have
0 and an - and 31
w (otherwise). Substituting
2d loog
(1 - () - 9) (1-p)(1
2 d
pq
Let w. be solution to the above equation, i.c.
(1 p) - ④)
W.
log
d
pq
By substituting 90 in Eq. (8) we obtain the bias b. = log together, let
p T-p
d 2.
Piecing everything
0. = 1e
Wo = = W k.v.v W, = =w, 11 W. = W.. b= b.).
(10)
and we are done.
B.3 PROOF OF THM. FOR NON- WHIN-WWIGGHTTIE TIED
Proof. By extending 0, R-d 0. A a.) DD it follows from the Transformer archi- xT) = Pe. (#n+1= = [ (Æn+1 = 1 2n), the Markov B.2 and $ B.I establish, prediction probability equalling the kernel
tecture in $ A that Po (In+1 kernel. As the proof Thm. 1 is a sufficient condition for global-optimality Hence 0, is global minimum for L(-) in R.P.
B.4 PROOF OF TIM.
Proof. First construct explicit Oz E R-D such that it satifies properties (ii)-(iv of Thm.
ie. it is stationary point with loss value being the entropy of the marginal ( ) and that it captures the marginal distribution P (I 1) πi. Then compute its Hessian and show that it is local minimum for p + >> thus proving property (i). On the other hand, the same e could either be a local minimum saddle point for p + q 1. We start with the construction.
17

---

Published as a conference paper at ICLR 2025
Recall that the full set of the Transformer parameters in the weight- -ieh case is given by 0 (e a, (PnIn=) Wo, W KQQ.V 22, 1.6 E R-dd Define 7 E IR D-d to be
0x e= 0, (p by Wo 0, W kov, W2 0, W 1,
log ()
)
(1I)
where (pr In=In=l C TRd W K.Q.V E Rmxd and Wi Rexd can be set to any arbitrary value. Now we start with property (ii).
(ii): So (x+) Pe,, (In+1
x P (zn+1 = NT.
Since 0, it follows from (Linear) and (Prediction) layers that fo, (xf) o(6) = (log(p/4))
ptg = T1. In other words, the model ignores all the inputs and outputs constant probability TΙ.
(ii): L( (In11) (T)..
Since Sa (+) =
B[ it follows from Eq. (3) that
T(0)))
Σ
Bag [#a+1 log Son, (xf) + (1 - In+1) log(1 So, (af)))
[E
Σ
E Fn+11 log # + (1 - In+1) log Tol
N
REIN]
1
Σ
- log T0 log
N
M
H(T) = H (xx+1).
(iv): vI((,,)
At 0 0. the individual layer outputs of the Transformer ($ A) given by Uni-embedding, xzn = Pn.
In E 10
Attention
yn = Pr.
FF
nn P.. Linear logit,
Prediction
for (x?+ TT
In other words, none of the layer outputs depend on the input sequence { IN In view of this fact and [[ = TI, using Lemma 2 the gradient with respect to a of L at 0 = Oz is given by
Val.
Σ ME[N]
E R
[(Ent]
So, (x+)) Va (a
b)]
N
Σ
12 [[N+11 T1) (Zn + Va a)l
ME[N]
(a
1
E. Ex [xn+1 TT) P.l
NE(N
0.
Similarly, for
VoL
1 N
Σ
E. Ex [(an1 fer (ITT) Vo (a nn +b] =
Σ
E [In+1 mi] = 0.
N
RE[N]
nE[M]
For any other parameter w E 0 apart from b, from Eq. (5) that the gradient V L has the term V.. V. (a" 2n)) = (V zn) a inside the expectation E S.nta[... Since a = 0, this equals zero and
hence V L = 0.
Together VL(O)) 0.
(i): 0r is a bad local minimum for L when p + 1.
Towards establishing this, we first let = (b, a) and G = (1Pm Wo. W W.Q.V, W2, Wi) be two different sets of parameters comprising i.e. (ax, ()) and compute the Hessian H A
18

---

Published as a conference paper at ICLR 2025
( 1(0) and show that it has the following block block-diagonal structure:
Ha 0 0
H
where Ha corresponds to the Hessian with respect to the parameters a and b in a. Further we show that if p + q 1, a > 0 i... positive-definit. This helps us in establishing that 0 local minimum. Now we start with the Hessian computation.
Hessian computation. We first compute the Hessian with respect a.
From Lemma 2, we have that second derivative with respect to at 0 = 0~ is given by
L-
Vo
Σ
B [ fo(x)))
N
nE[N]
(1
V6 (a 2n + +6]
N
[N[N]
E HoT
N
T 0,
where (a) follows from the fact that Vo fo(xf) Vao(a =n +) fo(2))(1 fo(x)))
V (a 2n + b). Now we compute the second derivative with respect a. From Lemma 2, we
obtain
(VaL Va
E Eutt (In+I- - Va(aTzz)
nE(N]
Σ
E B((nnt1 fo(x))) + (Vazn) -a)]
RE/N|
(
Σ
E [fo - fo) + (Vazn) (= + (Vazz)-a)
N
N
E [[(cn41-ffl2i)110 (2(4
WE[N]
where (a) follows from the gradient of the product rule and the fact that Vafe(z)) fo(=4)(1 (o( + (Vaza) a). A = Oz. this further simplifies to
(2) L
(b)
(E[TiTo P.P... 2B 2E(41+1 - Tl)
BEIN]
Σ
((TM))
-PpD)) 2 9) p t 71)..,1 I)
N
ME[N]
Σ
((TOT)) P..P..) - 2E(2n(TT) - {)
N
N
1
Σ
((TOTT)
P.PP..) - 2 (TO 1)
N
N
Σ
PAP
2
1
I
nE[N]
N
30
19

---

Published as a conference paper at ICLR 2025
MOTI
Σ REIN
PuP
+ +2p q 1)1
N
where (b) follows from the fact that Vaz, In 1 at 0 = 0r where is the identity matrix is IRdxd (c) from the observation that [[zn+1|<n>] = nn (1 - p- -)1 + p. and (d) from the fact that xo ptg
_
Now we compute the cross-derivative of second order Va,cL. L. Again, invoking Lemma 2,
Va
B, B (ent1 - [o(zi)
N
nE(N]
[So(1 - [o) + (Vazn) a)]
[E[N]
(0=0~)
E TM P.l
N
ME
Σ
Pa
TOT1
N
EEIN]
Piecing all the results together we obtain that for a (b, a), its corresponding Hessian is given by
PaP T
Ha ( (2) I(a 3x) = T
u
2(p + q I)
V
[E
nE(N]
(12)
We now show that the Hessian H3
2 L(a Bx) 0. Recall that
((p Wo. W.Q.V. W2, W1). For any W1, W2 E B, Lemma 2 implies that
V un wg L Vwy (Vig T.) = V Tm
E [ [ So wz Zn a))
N
NE(N]
(a) 1
Σ
E
[fo(1 - En a) 2n [
N
NEIN]
(0=0x)
where (a) follows from the fact that Vw, fo(2;) = Vw, aa 2n + b) fo(x)) fo(=T)) wi 2n a). Thus H3 O. Similarly, show that H VagL L = hence HØ H aag 0. Thus,
and
H. 0 !
H = ((0
Now it remains to show that H., is positive-definite when p g 1 and it implies that O is local minimum.
1
Positive-definitenss of Ha. Recall from Eq. (12) that H.
2 where u
ENEIN V. Ene[N] P.P.T /N + 2(p + - 1)I. From the characterization of positive-
definiteness by Schur's complement Horn and Johnson (2012), we have that H. > 0 < 1 >> 0 and V uu > We have that
V uu = =2p -
Σ
P.P. N
Σ DE[N]
Pn
Σ MEIN]
P N
N
nE[N]
20

---

Published as a conference paper at ICLR 2025
= 9 1)T+
Σ
Pn Σ
Pa
Pr.
Pa N
N
N
nE[N]
NEIN
REIN]
+(p 1)I + Cov Cov((pn)a 1
(13)
T
where (ov(())))))) = EnE[N] + (pn - -
Pu.
Pn - Ene[N]] & is the covariance
N matrix of the set {P,, (P,]n=) and hence positive semi-definite. Thus if p + 9 > 1, we have that 2(p+ + q 1)I > 0 and together, we obtain that V uu 0. Hence H > 0. Now it remains to show that ,r is a local minimum.
Ha is positive-definit implies 0~ is a local minimum. Since Ha 0, let H > AI for some > (in fact À = 2(p + - 1) works). Since 0 = (a, 3) E R-d interpret [( = L( .9) as a function of two variables and B3 with appropriate dimensions. We know the following facts about L(::)):
Fact 1. a + L(a, (x) has local minimum (as function of one variable) at ax (since Ha > 0).
Fact 2. B - L( 3) is constant in 3 (since On = 0, the probability so (o(")) is constant
w.r.t. Zn and hence w.r.t (Linear)).
Ha 0 with Ho 7 X.. 0
Fact 3. VL(ar, Bx) 0 and H (2) L( 3x)
show that (ain: 3x) ) 0r is also local minimum in two-variables. We prove this by contradiction. Suppose that (0,, 33)) is not local minimum for L(-,)). Without loss of generality, by shift of cordinates treat Ox the origin, i.e. (ar 0, 0, not local minimum for ... Then there exists an unit direction d (u. v) E R Rn-d with [d 2
Using these facts
is
uu|2 2 + ||0||2 1 and an EO such that
((dd < L(0), VO<E 1.
(14)
Clearly ||u||| > 0, otherwise it will contradict Fact 2. Using the definition of directional-derrivtivve. we have that
(VIss) - (VI(0),
d (2) L(0,0) lim
c-1
E
(VL(sd), d)
lim
8-2
On the other hand, using the Hessian structure the LHS equals dl ((,) 2 > 0. Thus
lim (L(e)), K1 >0
6 a00
Thus there exists an E > and > such that
(VL(sd), d) K, 0A AO<E S1,
which implies
(VL(s d) 2 Ke, AO S..
Defining the function R_ + as g(e) = I.(sd), we obtain that g'(e) = (VI(sd),d) 2 Ke for
0 I < Ej. Using the fundamental theorem of Calculus, have that for any S1.
L g'(t)dt
((E) 9(0)
2
Ktdt
21

---

Published as a conference paper at ICLR 2025
Ke²
Thus g(e) = ((ed) 2 L(0) + K for all < whereas L(sd) < L()) for all 0 < < Eg
from Eq. (14). Choosing &, min(Eo, E1), we have contradiction for 0 E E.. Thus 0 = d = (a )x) is a local minimum.
B.5 PROOF OF THM..
.,,,),,) Proof. Since 7 A (m,,aT) ax) IR. is canonical extension of 0T (ex ax....,,
RD-d which is a local-minimum for L(.) in R RD-A, following the same-steps for the gradient computation and probability evaluation as in proof of Thm. 2 in B.4, it immediately follows that O₂ also satisfies properties (ii)-(iv), i.e. it's a stationary point, it captures the marginal, since Pe. (xn+1 = = = Po,, (In+1- = 1 xf) = P ( = 1) = म, and hence its loss equals entropy of stationary distribution H (TT). In similar fashion, the Hessian computation is essentially the same except for a slight difference the Hessian structure, i.c.
1
U T V
0 (p (p+9-1 0
Ha
H00)) V L(0x)
Ha ToT u
11
1)1
0
0
(
P.Pa In the weight-tied case, we observe that the matrix V where u - EnE[N] Pa 2 EmE(N] N
N also contains the q-+q-11 1) terms which in the non-weight-tiee case gets de-coupled (due to separate and parameters). In fact, Ha corresponds the Hessian w.r.t the parameters a (b, a, e), i.c.
Ha = v! Lla=a, Now it remains to show that H is indefinite and hence 0, saddle point.
Clearly, Ha cannot be negative definite since with d = (1, ......)). we have d H.d = T > 0
for (p. q) E (0, 1). Now we show that cannot be positive definite either. Denoting
Pa EnEIN]
EnEIM (p q
P.P. N 1)1
(0+9-10 1)1
Ha No
C
0
Using the characterization of dositive-definitenes by Schur's s complement Horn and Johnson (2012),
we have that H. > ++ 1 >>>>>>>>> and C - bb > 0. This can be further simplified to
M 2 C bbT
EnEIN (p+ q 1)1
PAR. N
(p+ (p+9-1 q 1)
Σ
Eng(N] NJ
Pu N
EmELN Fx o]
Cov((p J. (w+ 9 1)1
(p 1 1)T
T
where Cov((p, I = EnE[N] t Pn -Pn-nn
N
(P, -Ene(NN]]
Px.
is the covariance
matrix of the set (p, Now we show that M cannot be positive definite. Suppose not. Then there exists a such that v Mv 2 A||a||2 for all v (v1,02) E 2d This further implies that
v._ CoC Cov((p,)) + 2(p + q () v2) 2 \lIl
Voi, U2 Rd.
Taking 1 = 0 the above inequality imples that 2/ < for all U2 E RR, which a contradiction.
Hence M cannot be positive definite and consequently neither can H..
B.6 PROOF LEMMA
Proof. Consider the loss function L(-) given in Eq. (1). We can rewrite it as follows:
L(0) =
Σ
E Egp+1 [Ix41 log (x + (1 #a+1) 1) log(1 [e(x]))
N
RE[N]
Σ
Exg
E Eunssley [#n+1]] log [x + R Rxneslz xn+1] m]] log(1 fo(x?))
N
N
22

---

Published as a conference paper at ICLR 2025
fa(z?)
fa(2²)
Σ
Bgg (In+1 In) log P (In+11 = [nn)
(In+1 [0n))))))))))))0) log
N
P P = 0 (nn))
n REIN]
P (xn+1 = 1 Ln) log
P (xnn11 = 00 2n) IP (In+1 = 0 Inn)
IP (In+1 = 1 Inn
fo()))
fo (x
Σ Eug (P = Fxn log P P(In+1 = 1 (in)
+ PP (In+1 = 0 =0(5)) log
N
P P = 0 (in))
NE[N]
Exa
PP (I((nt1 Inn log (In+1 = Fn)) (x()
IP + = =0 Inn) log
Jan0
[In+1 0 2
[o(x))
Σ Euy (I [x)) log P (In+1 (nn))
+ +P (In+1 0 (r)) log P (In+1 0 3.p))
N
NEIN]
Eiz,
(P (P(nt11 [[a))n))))) log (In+1 = [n))
1
(In+1 0 xnn) log P ( = 0 1nn
1
Since fo(=f) = IPo (In+1 = | In), we have that the first term above is
fa(
1 fo(xf)
P P( = 1 (n) log P [xnnt1 = 1 Fn)
+ P (In+1 = -0 (n) log P P(xn+1 = 0 Fn) =-DKL (In+1 1 xn) || Pg (In+1 = 2
Further, observe that the second term is exactly the entropy rate ( [mn)). Hence, the above expression for the loss reduces to
L( N
Σ nE[N]
Exy [DKL (In+1 = (x) || P. (5nt1 = --1xi (I
and we are done.
B.7 PROOF OF LEMMA 2
Proof. It suffices to show that for any component 0, of 0 E RR,
0
L(0)
1
Σ
EE
E
(In+1- - fz(")) foo="))
00; (a nn + +b)
00;
N
DE[N]
Σ
EzT
(P (P(2nt1= -x) fo(z))) θ
(a T 2n + +b)
N
00,
BEIN]
Recall from Eq. (3) that L(-) is given by
((0)=
1
Σ
E
E [#nt1 log fo()) + (1 2n41) logt) - fo
N
nE
N
which implies that
0
à log fo(x") + (1 2n+1).
0
L(0)
Σ
E
Int1
log(1 fo(z)))
00,
N
00,
00,
NN(N]
1 à
N
E
31
o+1
Fatl
fa(TT) 00,
fo(x)) + (1 Int1)
1 fa( 00;
(1 - fo(a)))
Since [0(x]) = a(a² T in + b), we first note that the derivative of a is given by d'()) IIte-ay
(1))(1 - (2)).. Hence, the derivative
a a0, fo (2T)) can be written as
0 i o(a T 2n + b) = a(a" 2n + b) oaa T + b)]
fo (o())) =
(a 2n + b)
00
00,
00,
23

---

Published as a conference paper at ICLR 2025
a 00;
(ola))(1-faa')))
(a T on + +).
Plugging this into the above expression, have
0
1(0)
Σ
E
1+
Jati -1-18 [xxx))])]
(a T Zn +b) xn+1) Sa(x))
(a nn +b)
00,
N
21
00,
00,
N
0
Σ NE(N]
E
Intl 11-11-10 (1-f(((x)) (x(
(a T Zn + (1 xn+1) f (z
(a nn b)
N
41
00,
00,
Σ
E
n+1
(a T nn +b)
N
00,
AE(N]
Σ
Egp
[= {a()))
(a T 2n b)
N
80
N
0 (a nn + +bb
Σ
Egp
(Pn+1==111 xn) fo(x)))
00;
NE(N]
and we done.
24

---

Published as a conference paper at ICLR 2025
C ADDITIONAL RESULTS FOR FIRST-ORDER MARKOV CHAINS
C.I MODEL ARCHITECTURE AND HYPER- HYPER-PARAMENT PARAMETERS
Table I: Parameters in the transformer architecture with their shape.
Parameter
Matrix shape
transformer.wte
2
transformer.wpe
x
transformer..ln. (xt)
d x 1
transfomme.h.attum attn (x
3d x d
transformer.haate proj (x
d dx1 4d x d d d 4d
transformerh.hn (x
transformer.h. wwt.mme............mmm mmlp.c_fc (x
transformerh.mipl httpoommrrrh.ml._pprrr.
(x()
transformer.I
dx
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 Pagliardini (2023)
Batch size Accumulation steps
Grid-searched in {16, 50)
Optimizer Learning rate
AdamW (81 = 0.9, 88 = 0.95)
0.001
Scheduler # Iterations
Cosine
8000
Weight decay
1 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
Between 1 6 depending on the experiment
Attention heads
Grid-searched in 2, 4,8}
Mask window
Between 2 and full causal masking depending on the experiment
Repetitions
3 5
D EMPIRICAL FORMULA FO 2Op+0 q < 1 BASED ON LOW-RANK SOLUTIONS
In this section we compute the function (= that gives the next-symbol probability predicted by the network, using the values of the weight matrices obtained five independent experiment runs. By substituting the empirical weights into the transformer architecture from $ A, i.e.
Cn In Pa E Rdd
(Uni-embedding)
Yn I. + W Σ
allmn
W v 8, E Rdd
(Attention)
y, + W₂ ReLU(W v.,) E Rd
(FF)
logit, (a, zn) +
R,
(Linear)
fo A Po (I = 1 x = (logit
(Prediction)
We can obtain an explicit expression for (afr) as it is actually learned by the model. We now analyze each section of the model architecture separately.
25

---

Published as a conference paper at ICLR 2025
Embedding. All the five independent runs show that the word embedding vector has the structure
c.
(15)
where v = (Ui...., v2) is such that 8 E [-1,+1} for all i, i.c., v E {-1, and is some constant. Moreover, the positional embeddings approximately constant across positions and they share a similar structure to e. In particular, we always have that
Ps p p. p.u Vn
(16)
for some constant p R. Furthermore, the constants are always such that p and + p 0.
Attention. Across all the runs, we observe that the contribution of the attention mechanism is negligible compared to the skip-connection. In particular, we observe that
Io Eiejnj att W v æill
0.01
(17)
uniformly for all n. Therefore, we can use the approximation
yn ~ In Vn.
(18)
FF. For the MLP layer, we observe that W1 and W2 have a clear joint structure. In fact, we empirically see that
W W1 w TT where v is again the same vector as in Eq. (15), w E {-1,1] and E R. Hence, W, is a rank-one matrix. As customary in the GPT-2 model, for our experiments we used r 4d 16. Furthermore, we see that
(19)
W WT.
(20)
Due to this structure and the formula for yn described above, we have
W yn Wi i0, wi wid(ee, p)w Let now r ReLUUW, y,,) Due to the fact that p < and e + > O, we have that, if x, 1,
(21)
p, if wi
T1 0,
(22)
if Wi = 1.
While if O,
0,
if W: 1,
r;
(23)
p.
if Wi -1
Let 3= ET 1 [w.=1].
Since W2 2 = W = wi w", we have that, for 2n.
wfdlle + p). v, if In = wfdplr - B) u. if if 0.
(24)
Or more compactly,
T = wf + ()2 - r)x,, +r 3}
(25)
and
2n nn + F (ex,, + (1 + wfd((23 r) + B))
(26)
Linear.
Since a = due to weight-tying, we have
logit,, = ed(er, + (1 + (2 - r)æn + - B)) + b
(27)
Prediction. We can now plug in the empirical values obtained by averaging five independent runs. The numerical results that we get are
0.3618
p -0.1539
W1 0.3264
(28)
b 00.1229
26

---

Published as a conference paper at ICLR 2025
Plugging these numbers into Eq. (27), we get
(0.8191, 1.3387
if 1,
logit,
(29)
if 2n
Hence, by applying the sigmoid function to the logit values, we obtain the predicted probabilities
[0.08191)) = 0.694, (0(--1.887) 0.199, if = 0.
I = 1,
fo = Po (In-1 = 1 x.,))
(30)
The numerical results correspond almost exactly to the expected theoretical values of q 0.7 and p = 0.2.
27


Reviews:
[{"review_id": "3pTzWiXWoW", "reviewer": "Reviewer_pFAs", "summary": "The authors study a curious case of transformers via Markov Chains, in which one-layer transformer struggles to find the global minimum in the Markov task. They further provide theoretical and empirical studies, showing that weight tying can lead to bad local minima when Markovian switching is greater than one. They also extend the results to multi-state Markov Chains. Experiments in multi-layer transformers further show that they can reach the global minimum, which is an intriguing future direction.", "questions": "1. Can authors comment on the role of attention layers in the paper?\n2. Can authors present the attention patterns of the transformers trained the Markov tasks?\n3. If time permits, can authors train pure MLP layers on the Markov task?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The problem setup is interesting and concise.\n2. The presentation is clear. \n3. The analysis is solid, containing both theory and extensive experiments.", "weaknesses": "The attention structure is missing here. Specifically, the Markov chain requires no attention structure, nor does theoretical analysis include the attention layer. It is more like \"MLPs with Markov\" instead of \"attention with Markov\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "MiLIJBfu4b", "reviewer": "Reviewer_rH7j", "summary": "This paper establishes a general and novel framework for the theoretical and empirical study of the attention mechanism on Markov data. They theoretically show that for one layer of Transformers, if the data comes from a Markov chain with binary states and transition probability given by P = (1-p, p \\\\ q, 1-q), the global minima of the next-token cross-entropy loss can recover the optimal risk. They theoretically show that there exist bad local minima with weight tying if p+q > 1, and without weight tying, local minima become saddle point. Empirically, they validate the theory and show those saddle point can be avoided if without weight tying. They also present experimental results on multi-state Markov chain and mention some open problems about deep Transformers.\n\nI really like this paper. I think the framework is novel, interesting, and worth investigating. The theoretical and empirical results are solid, and the presentation is particularly clear. Although this is very initial research, they propose some future directions about why single-layer Transformers can sometimes be trapped in bad local minima when trained on the Markov data, and why deep Transformers can behave well and avoid local minima or saddle point. I think we should encourage this style of theoretical research since there are too much theoretical research with complex setups and heavy math, that explains nothing at all.", "questions": "/", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The theoretical results are solid and can be validated via experiments.\n\n2. They establish a novel framework to investigate the attention mechanism on Markov data.\n\n3. The presentation is particularly clear and comfortable.", "weaknesses": "Current, I do not see any big issues.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "SXdZ7V92qM", "reviewer": "Reviewer_dP9v", "summary": "This paper introduces a framework for learning a first-order Markov chain using a single-layer transformer. Within this framework, the authors characterize the network’s loss landscape, identifying the presence of both global and bad local minima.", "questions": "1. See weakness.\n\n2. In the proof sketch of Theorem 2 (line 363), the authors set $\\mathbf{e} = \\mathbf{a} = 0$, which implies that for the uni-embedding (line 213), $\\mathbf{x}_n$ depends solely on the position encoding and lacks information from the state $x_n$. If I understand correctly, this setup could lead to problematic solutions for the model, as it is unable to fully utilize input information—potentially independent of the effects of weight tying. Could the authors clarify this point?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Examining the limitations of a single-layer transformer in learning a simple Markov kernel is an intriguing aspect of this work.\n\n2. The authors characterize the roles of Markov switching probabilities and weight-tying in contributing to the presence of local optima.\n\n3. Extensive experimental results are provided to support the theoretical insights", "weaknesses": "The main limitation of this paper is that the current results may not fully address their central question: systematically characterizing the learning gap of a single-layer transformer. The authors show that a bad local minimum exists when $p+q > 1$, but they do not rule out the possibility that only global optima exist when $p+q < 1$. Given that the results are primarily existence arguments, this does not entirely support the empirical observation that the transformer converges to a bad minimum in the $p+q > 1$ regime while learning the global minimum in the $p+q < 1$ regime. This implicit bias would benefit from additional theoretical support, such as insights from learning dynamics.  Moreover, the data model is overly simplified, which raises expectations for a more rigorous analysis that fully characterizes the solution learned by transformers rather than relying on existence arguments alone.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "fdFpvu1G4Z", "reviewer": "Reviewer_FHWC", "summary": "This paper introduces a new framework for analyzing transformers through the lens of Markov chains, which is a novel and theoretically rich approach. This framework enhances our understanding of how transformers handle sequential data, particularly first-order Markov processes. Specifically, the theoretical characterization of the loss landscape for single-layer transformers, identifying the presence of both global minima (bigram) and bad local minima (unigram), is a good contribution. This helps explain the empirical observations and provides a solid foundation for future work.", "questions": "Could the authors clarify the difference between the following work [1] and the proposed results?\n\n[1] Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality. COLT, 2024.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ The theoretical results are novel and informative.\n\n+ The experiments are well-designed, providing solid support of the considered insights.", "weaknesses": "I found no explicit assumptions in the paper, which may imply that few assumptions are utilized in the theoretical analysis. Thus, I suggest that the authors list the assumptions that were used.\n\nThe set of words under consideration is too toyish, which is the primary concern.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Spotlight)