Towards Understanding Transformers in Learning Random Walks
Authors: Wei Shi
Yuan Cao
Venue: NeurIPS 2025 poster ()
Keywords: transformer
self-attention
random walk
Primary area: theory

Abstract: Transformers have proven highly effective across various applications, especially in handling sequential data such as natural languages and time series. However, transformer models often lack clear interpretability, and the success of transformers has not been well understood in theory. In this paper, we study the capability and interpretability of transformers in learning a family of classic statistical models, namely random walks on circles. We theoretically demonstrate that, after training with gradient descent, a one-layer transformer model can achieve optimal accuracy in predicting random walks. Importantly, our analysis reveals that the trained model is interpretable: the trained softmax attention serves as a token selector, focusing on the direct parent state; subsequently, the value matrix executes a one-step probability transition to predict the location of the next state based on this parent state. We also show that certain edge cases not covered by our theory are indeed failure cases, demonstrating that our theoretical conditions are tight. By investigating these success and failure cases, it is revealed that gradient descent with small initialization may fail or struggle to converge to a good solution in certain simple tasks even beyond random walks. Experiments are conducted to support our theoretical findings.

Towards Understanding Transformers in Learning Random Walks
Wei Shi
Yuan Cao
School of Computing and Data Science
School of Computing and Data Science The University of Hong Kong yuancao0hku. hk
The University of Hong Kong hku. hk
Abstract
Transformers have proven highly effective across various applications, especially in handling sequential data such natural languages and time series. However, transformer models often lack clear interpretability, and the success of transformers has not been well understood in theory. In this paper, study the capability and interpretability of transformers in learning a family of classic statistical models, namely random walks circles. We theoretically demonstrate that, after training with gradient descent, a one-layer transformer model can achieve optimal accuracy in predicting random walks. Importantly, analysis reveals that the trained model is interpretable: the trained softmax attention serves as a token selector, focusing on the direct parent state; subsequently, the value matrix executes a one- step probability transition to predict the location of the next state based on this parent state. We also show that certain edge cases covered by our theory are indeed failure cases, demonstrating that our theoretical conditions are tight. By investigating these success and failure cases, revealed that gradient descent with small initialization may fail or struggle to converge to a good solution in certain simple tasks even beyond random walks. Experiments are conducted support our theoretical findings.
1 Introduction
In recent years, transformers 31 have revolutionized many fields such as natural language processing
2 23 B1,, computer vision 8 24, reinforcement learning [ [5], and have rapidly emerged as a key component in state-of-the-arr deep learning models due to their ability to capture complex dependencies in data. While transformers exhibit remarkable practical performance, the underlying mechanisms of transformers are still not well understood due to their complex architecture.
In order to theoretically understand transformers, a number of recent works have investigated their capability in learning from sequential data that follows certain classic statistical models. Specifically, 4 [2 studied sequential data with an underlying causal graph, and theoretically showed how transformers can encode the causal structure with the self-attention mechanism for in-context learning. 66 considered the task of in-context learning of Markov chains, and investigated how two-layer transformers can make predictions for Markov chains according to the context.
[22] considered
Markov chain forecasting tasks and reveals connection between context-conditioned Markov chain and the self-attention mechanism. (21 characterized the loss landscape of a one-layer transformer and demonstrated the existence of global minima and bad local minima in learning Markovian data with vocabularies of size two. Although these works provided valuable insights, there remain many open questions that require further exploration. Notably, the Markov chain data studied in [21] can be essentially understood as a random walk over a space of only two states. Therefore, given the results in 2T natural question is when and how transformers can learn more general random walks over larger state space.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

In this paper, we consider classic statistical model for random sequences, namely random walks on circles, and study the capability of one-layer transformers to learn from such data and make predictions. We consider a general setting that allows a large number of nodes (possible lo- cations of the walker) on the circle, which is an extension to the two-state setting studied in [2I]. We also consider the general setting where a walker moves on the circle clockwisely with probability p, and counter-clockwise with probability 1 1-2 p. and we build theories that cover the full range of p [0, 1]. For such a classic, yet relatively general class of random sequences, our goal is to theoret- ically study the performance of one-layer transformer trained by gradient descent in predicting the next location, and reveal the interpretability of the obtained transformer model.
1
K-1
K nodes
Figure 1: Illustration of random walks on circles with K nodes and transition probability p.
The main contributions of this paper are as follows:
We theoretically demonstrate that one-layer transformer can be trained to optimally predict the next location of a random walk with p E (0, 1). Despite non-convexity, we prove that the trained model converges to the optimal prediction function with a rate O(T-1/2), where 7 is the iteration show that the trained transformer achieves optimal
number of gradient descent. Furthermore, prediction accuracy max (p. y} after a constant number of iterations.
Our analysis reveals that the trained transformer model is interpretable, as we precisely delineate the role of each component in the model. First, the trained softmax attention can select the "direct parent" token by assigning it a near-one score. Second, the trained value matrix serves as a one-step transition model that recovers the ground-trute probability transition matrix, which is applied the "direct parent" token to make an optimal prediction.
We also identify failure cases when = 0 or 1. In these cases, we show that starting from zero initialization, the training of the one-layer transformer model with any loss function and any learning rate will always fail, resulting in a transformer model whose performance is no better than a random guess. This negative result is complementary random walks with E (0, 1), and they together give a comprehensive characterization of the capability of transformers learning random walks with p E [0. 1].
positive guarantees for learning
We provide intuitive explanations that the failure cases with p = 0 or are optimization failures caused by zero initialization. Notably, similar optimization failures may also happen beyond the cases of random walks, similar issues in optimization. We also empirically demonstrate that although training may still take longer, these failure cases can be resolved a certain extent with random initialization.
construct simple question answering tasks that also suffer from
Problem Setup
In this section, we present our problem formulations, including the random walk prediction task, the one-layer transformer model, and the training algorithm.
We study random walks on circles. Specifically, consider K nodes (possible locations) that are arranged on a circle so that each node has two neighbors. Without loss of generality, we suppose that the nodes are assigned with node IDs 1, 2,2,... ... K in clockwise manner. A "walk" on the circle refers to the process where a "walker" moves step-by-step among the nodes of the circle. We suppose that, starting from random initial location, each step, the walker moves either clockwise with probability p or counterclockwise with probability - p. to a neighboring node its current position, where p 10, is a fixed probability. In this way, random walk of length N generates sequence of "states" $1....&. SN, where s; E [K] denotes the location (node ID) of the walker at the i-th step. We aim to address the problem of predicting the walker's next location sy based on the historical locations $1 S1-1.
To better formulate this random walk prediction task, we map 81,..., 8N-1 to embeddings ær,.... -N K Our goal then to train a model to predict the target y = 8x based on Fi...., -N-1. In the following, we give detailed definition and discuss some basic properties.

---

Random walk on circles. Suppose that there are K nodes on circle and the transition probability is p. Bi.... : NN-1. y are generated as follows:
1. Draw Uni
2. For i = .,,..,........................ N, sample either 8; ((1--1 1)x with probability or 8; (8;-1 1)x with probability p. 3. Set z; Cs,, [N 11. and y SN.
Here, ( k is defined as the integer satisfying (5) E [K] and (8) = (mod K). It is clear that the sequence æ ....... NN 1 By form Markov chain, and P(y ",,
= =Pyy 1).
Moreover, the transition matrix of the Markov model is II (Tig) (xxj)xx KxK. where Ti.j p- f
j [(mod K)} + (1 p) 1 = + 1(mod K)}. visualization
II is given in Figure The
Markov property indicates that the optimal predictor of y is given by
J (x ( IN
-N-1:
and the optimal prediction accuracy any predictor can achieve is OPT = max p, - -p)- Based on the formulation of fort ). is clear that simple autoregressive model f(X) VæN-1 can already solve this random walk prediction task. However, the goal of this work is not to identify the optimal model for solving random walk tasks. Instead, aim understand and analyze the capability and interpretability of transformers when learning such classic statistical tasks from scratch. Therefore, in the following, we introduce a simple transformer model.
(c) II with p Figure 2: Visualization of the transition matrices II with p = 0.5, 0.7, 1.0 respectively.
(a) II with p 0.5
(b) II with p 0.7
We consider learning the random walk prediction task with simple one-layer transformer model. By naturally treating the one-hot vectors 2i...., EN-1 token embeddings, the task to predict the next position is a problem of next token prediction. Therefore, we define the data matrix
X [21, 22 ................................ ØN-1. 0] KK*,, We also employ positional embedding matrix P
[pL. P2..... PN] E R MXN, where M is the embedding dimension with M 2(N3/22 and p. RM is defined as
in
2i M
Min M +1
P: = sin
.sin
sin
M
+
for i = 1, N. The positional embeddings above are inspired by the fact that (p:, for all i j. which significantly helps simplify our theoretical analysis (see Lemma R.5 the appendix). Additionally, orthogonal positional embeddings are commonly considered in existing theoretical studies [3] 12 32
Then, we define the matrix X by concatenating the input matrix X and the positional embedding matrix P as
33
X
X P
32
ØN-1
12. Eyl R(K+M)XN
P1
P2
PN-1
PN
We consider one-layer transformer model to make prediction given input matrix X. The transformer defined as follows:
(x X W
(2.1)
where V E x W E (K++)) are the trainable parameter matrices, S R + R is exp(x.)
the softmax function defined by [S(=)], Σグ of all the trainable parameters. In this definition, consider reparameterization where
" and 0 (V, W) denotes the collection

---

use a single matrix W to denote the product of the key and query paraneter matrices in self- attention. Such kind of reparameterizaations are widely considered in theoretical studies of transformer models [1] 12 14.
32 34 In addition, we omit the softmax scaling factor, which is mainly for simplicity. Such omission has also been considered in most of the theoretical studies 25 [6]]
16,
22 29
17 18 22
Note that by (..1). given any input matrix X, the transformer model outputs a K--dimeniional vector. This follows the standard practice of K-class classification: for i E [K], [fa (X)]; can be treated as a predicted "score" of the i-th class. More specifically, we define the prediction rule as follows.
Definition 2.1. For any predictor f(x): RXXN RK the predicted label is given as
Px)d(f(XX) == min {
E [K] [f(),, max (I(())1 KE[K]
The definition above matches the common practice to predict the label that corresponds to the entry in f(X) with the maximum function value. It also gives a naive way to handle ties when f(X) contains multiple dimensions with the same (and maximum) function value, we always predict the dimension corresponding to the smallest label. We remark that this definition to handle ties is just to exclude ambiguity, and the detailed rule to handle ties is not essential. Our result works for all reasonable methods to handle ties.
We train the transformer model defined in 2.1 by gradient descent, minimizing the loss function
L(
E ((ee]Jo())))
(2.2)
where E(-) is a loss function. In terms of the specific choice of f(-), our analysis will cover learning random walks on circle by minimizing the log-loss {(2) = log(2 + c), which has been considered
in series of recent works [12 16 21 28
We consider gradient descent with zero initialization ( = 0 Kx W = (K++))((++) to train the model. The update rule for the parameter matrices W, V follows:
( = ((t) - V W L(0()),
v(++)) = ( L(B1))
(2.3)
where 17 > is the learning rate and t 2 0 is the iteration number. Note that the log-loss {(2) = - -log(2 & is well-defined and does not blow up zero initialization due the stability constant e > O. Gradient descent with appropriate learning rate further ensures that it does not blow up during training.
3 Main Results
In this section, present our main theoretical results on learning random walks with self-attention layer. In our result, we can choose any T = poly(n, c K, N, .M) as the maximum admissible number of iterations, and only consider the training period 0 S t T". This technical assumption regarding a polynomially large maximum admissible number of iterations prevents training from becoming exponentially long and is mild assumption since exponentially long training is impractical.
Our main results are given the following theorem.
Theorem 3.1. Suppose that < P 1, K > 4, and N 2 C2 poly(K) for some constant Cp that only depends on Further suppose that the transformer is trained by gradient descent (2.3) to minimize the loss 2.2 with {(z) = log(= c), and 7), 6 = (11). Then there exists To = O1),
such that for all To < <TT it holds that:
1. The trained transformer achieves optimal prediction accuracy:
P Pred(feery (X) = yl = OPT = max(p, 1 pp.
The transformer converges to the optimal predictor:
SOiT) (X)
OPT X
=
[ (X)ll2
(TT
2

---

3. The value matrix converges to the true transition matrix in direction:
V
II
- O
II
[F
4. Softmax attention selects the "direct parent" of
[S(X ( N 2 100-eexp(-(NN
W ww), exp(-s((N)) for j * N [S(X W
The first result in Theorem 3.] states that the transformer trained by gradient descent for a constant number of iterations can achieve prediction accuracy max ph. which matches the optimal accuracy OPT. The second result in Theorem 3.1 further gives more detailed characterization of the trained transformer, and demonstrates that the normalized model converges to the optimal
prediction model fOPT (X) II IN-1-
Token selector
One-step probability transition
1
2
1
0
0
0
O
0
0
0
0
0
V II
X
V II
xyy-1 Poy|xy-
Figure 3: Illustration of how the trained one-layer transformer performs optimal prediction in an example with 1/2, K 6. Here denote $ S(X W(T)
The third and fourth results in Theorem 3. further back up the first results by a precise characterization on how the self-attention mechanism works in predicting random walks. Specifically, the third result demonstrates that direction, the value matrix V(T) serves as a one-step transition model by recovering the ground-truth one-step transition matrix, and the last result indicates that softmax attention performs optimal token selection by assigning near-one score to the (N - 11) token (the direct parent of y). These results show that, when learning to predict random walks, the one-layer transformet obtained through gradient descent is interpretable: the trained one-layer transformer model makes predictions by selecting the correct parent token øN- of y by assigning a softmax weighting close to 1 to it, and (ii) predicting by applying a one-step transition model to N through the linear mapping defined by V. An illustration given Figure
recent related work [2 studied how one-layer transformers can learn Markov chains with a state space of size two, and analyzed the loss landscape by identifying global minima and bad local minima. We remark that we consider different parameterization of one-layer transformers, which makes our results not rigorously comparable. However, by studying random walks state spaces of arbitrary sizes and establishing theoretical guarantees directly on transformers trained by gradient descent, our work explores a setting that is not covered in (1
Notably, Theorem i.T based on the assumption that the transition probability satisfies < < 1, excluding the edge cases when 0, 1. In Section 4 we will show negative result showing that when p = 0, 1, gradient descent with zero initialization fails to achieve good prediction accuracy. Therefore, the assumption 0
in Theorem 3. inevitable.
Here we informally explain how Theorem 3.1 proved. The idea of the proof can be shown by investigating the first several gradient descent steps:
Step After the first gradient descent step, due the zero initialization and the Toeplitz property of the transition matrix II, can be shown that (1) is also Toeplitz matrix whose largest entries appear exactly the locations of the largest entries of II W(! is still zero matrix due to the fact that (0) = 0.
5

---

Step 2. With the same analysis, we can also show that ( is a Toeplitz matrix whose largest entries appear exactly on the locations of the largest entries of II" Moreover, the locations of the largest entries in V encourage W to be updated so that higher softmax weightings are put upon the "direct parent" token xx-1 (see Lemma C.5 in the appendix).
Step 3. The higher weighting on æN-1 given by W further encourages V be updated towards II in direction. And V obtained in Step 2 continues encourage W(3) to place even higher weighting on -N-1 1 (see Lemma C.7 in the appendix).
From the three gradient descent steps discussed above, it is clear that V will converge to the direction of IIT and ( will consistently place a high weighting on the direct parent token æy This is our key intuition for proving Theorem B.J and in our formal proof (given in the appendix). we use an induction to characterize the whole training procedure.
4 "Deterministic Walks" with p 0,
we demonstrate that one-layer transformer can be trained to optimally predict random walks under the assumption 0 < p < 1. In this section, we justify the necessity this assumption and analyze the edge cases 0, 1, which lead to failure for learning random walks. Random walks with 0, are essentially "deterministic walks", and for them we have the following theorem.
In Section 3
Theorem 4.1. Suppose that N = rK + with r 2 1, and p = or 1. Further suppose that the transformer is trained by gradient descent (2.3) to minimize the loss (.. Then for any loss function
{()), any learning rate v7 > O, and any 2 0, it holds that
1 ((Pred(focc) (X)) = =v) K
Moreover, with probability 1, for all T 2 0, it holds that
(T) 1KxK,
[((X w w(T =...= |S W (()))))
N-1
Theorem 4.1 | shows that the prediction accuracy of the trained transformer on deterministic walks 1/ K, equal to the accuracy of a random guess. Moreover, the characterizations of the softmax scores and the value matrix (T) further demonstrate that the transformer takes average over all tokens, and then gives the same prediction scores for all possible values of y. Notably, these results hold for any choice of the loss function and any learning rate, indicating that this failure case of the transformer cannot be resolved by simply adjusting these training setups.
Theorems 3.1 and 4.1 together provide comprehensive characterization of the transformer's perfor- mance when learning random walks with transition probabilities p [0, 1]. Based on Theorem 4.1 it is clear that the assumption in Theorem 3.1 that 0 < p < is necessary and cannot be further relaxed. Theorems and also point out an interesting fact, that compared to random walks with O < p < 1, the seemingly easier task of predicting deterministic walks with p= 0. may be more challenging for a transformer to learn. Here, we would like to remark that the failure of learning deterministic walks is optimization failure, and is highly due to zero initialization. To see the reason, we can consider the two initial gradient descent steps:
Step 1. Since the initial softmax weightings on all tokens are the same, (1) is essentially trained based on the averaged token z N-1 ΣN-1 x;- Importantly,
that
is a constant vector that does not depend on y.
This means that z does not provide any helpful information in predicting and is therefore "uninformative". As result, it can be shown that all entries in ( are equal (see Lemma D.I in the appendix). Additionally, (1) is still a zero matrix since (())= = O.
Step 2. With the same analysis Step we can show that all entries in (2) are equal. Moreover, due to the fact that (!) is proportional to the all-one matrix, it can be further shown that
(2) is updated so that the softmax weightings on all tokens æ1, NN-1 remain equal.

---

In our formal proof, we inductively show that throughout training, the value matrix ( is always proportional the all-one matrix, and the softmax weights on all tokens are always the same.
From the discussion above, we can observe that transformers struggle to learn deterministic walks due to the unbreakable symmetry in the training dynamics, arising from the "uninformative" token average a given by the zero initialization. We remark that if random initialization is used, the symmetry will be broken, and transformers may succeed in learning the optimal predictor. However, if the random initialization is too small in scale, we can still expect that the optimization for learning deterministic walks to be more challenging compared to that for random walks. While our theoretical analysis does not cover random initialization, we present empirical studies in Sections [ and 6 verify this claim. We believe that theoretically analyzing the impact of random initializations can be an important future work direction.
Experiment
In this section, we present experimental results to support theoretical analysis. We consider two cases: the first one is the zero initialization case that aligns with the setting used in our theoretical analysis, and the second one is the random initialization case, which helps verify our discussion about the optimization failure caused by zero initialization. In all experiments introduced in this section, we set the number of nodes K 6 and the length of each sequence N 97. We utilize the transformer model introduced in Section 2 and utilize the gradient method to train the model. The prediction accuracy is calculated based on 1000 test data.
Zero initialization. this case, we set the length of the positional embedding M 1000, the initialization ((()) Og*K, ( (K and the learning rate 1) The constant in the log-loss is set as c = 0.1. For both tasks, we generate 1000 sequences to train the model.
Figure 4 and Figure illustrate the results of the experiment for p 0.5 and p respectively: Figure 4(a) and Figure 5(a) present the prediction accuracy; Figure ((b) and Figure 5(b) visualize
the value matrix V after 50 iterations; Figure 4(c) and Figure 5(c) display the attention scores attached to each token after 50 iterations. To clearly observe the results, we also provide Figure 4(d)
that represents the part of Figure 4(c)
Mudef Maretrs
(a) Accuracy
(b) (T)
(c) Average attention
(d) Part the attention
Figure 4: The results of the experiments for 0.5 with initialization: (a) the test accuracy;
(b) is the visualization of V; (c) and (d) present the average attention of the test data with æ-axis representing the position of the token and y-axis representing the attention score.
We can observe that these experimental results for p 0.5 provide strong support for Theorem 3.1 Figure ((a shows that the prediction accuracy is close to the optimal accuracy (50%) within constant iterations. Figure 4(b) indicates that (T) can recover the transition matrix II" shown in Figure 2(a) Figure 4(c) presents that the (N 1)-th attention score is the highest and close to 1,
indicating that the self-attention layer is able select the true parent token. All of these experimental results demonstrate the performance of transformers in learning random walks.
In addition, we can find that the experimental outcomes for p 1 match the theoretical results stated in Theorem 4.1 We obtain an accuracy close to 0.167 from ((a) which suggests that the prediction accuracy for learning deterministic walks is approximately equal to 1/K, far away from the optimal accuracy (100%) and no better than a random guess. Figure 5(b) indicates that (T) approximately proportional to 1x*x. Figure 5(c) shows that the attention scores attached to all tokens are identical, which proves that the self-attention layer cannot select any of the tokens when learning deterministic walks. These experimental results demonstrate that the self-attention mechanism struggles in learning deterministic walks with p 0, 1.
7

---

saleddor
64100
644
BCO3
ac100
44
33
3.2
(a) Accuracy
(b) V
(c) Average attention
Figure 5: Experiments for = with zero initialization. (a) gives the prediction accuracy along
training. (b) is the visualization of V. (c) is the average attention of the test data with æ-axis representing the position the token and y-axis representing the attention score.
Random initialization, In this case, we set the length of the positional embedding M = 1000, the initialization Vi(0) W (O0, x²) with 0.01, and the learning rate 0.01. The constant in the log-loss is set as = 0.1. For both tasks, we generate 1000 sequences to train the model. To ensure numerical stability, we normalize H, 's to unit length in the softmax attention.
the results of the experiment for p 0.5 and p 1. Figure 6(a) and Figure ((c) show the prediction accuracy within 1000 iterations, respectively. In Figure 6(b) and Figure 6(d) we first normalize the output of the trained transformer model to get a K dimeensiion vector, which can be regarded as the prediction distribution of K locations. The KL-divergence between this prediction distribution and the true distribution of N- illustrated in Figure 6(b) and Figure 6(d)
(a) Accuracy
(b) BL-Divergence
(c) Accuracy
(d) KL- KL-Divergence
Figure 6: The results of the synthetic experiment with random initialization: (a) and (b) correspond to the experiment for 0.5; (c) and (d) correspond to the experiment for 1. (a) and (c) present the prediction accuracy. In (b) and (d), we first normalize the output of the trained transformer model to get K K<ddimesiional vector, representing the prediction distribution of K locations. Then, we display the KL-divergence between this prediction distribution and the distribution of y N-
Figure 6(a) clearly shows that in the experiment for p 0.5, the accuracy is close to the optimal accuracy (50%) after around 400 iterations. However, as shown in Figure 66c) for
1. the prediction accuracy cannot reach the optimal accuracy (100%) within 1000 iterations. Based on the plots of Kl-diiergence, we can also see that the transformer learns the true prediction distribution of random walks much faster than learning that of deterministic walks. Note that these results are for training with random initialization, and hence the results do not perfectly match theory for zero initialization in Section However, the experiment results still show that the optimization task is more challenging, even with small random initialization.
6 Beyond Random Walks
In Section 4 we intuitively explain that one-layer transformers may suffer from optimization issues when learning deterministic walks with p = 1 due the fact that zero initialization produces token average which is "uninformative". In this section, we briefly discuss other tasks beyond random walks where transformers with zero/small initialization may face similar optimization challenges.
We construct two question answering tasks. The detailed descriptions are given as follows. Task 1. We consider simple question answering task. Possible input questions are of the form:
Based the list 'apple, orange, apple, apple, orange', which rype of fruit appears most frequently?

---

Here, the list stated in the question can be any combination of "apple" and 'orange' with a fixed length of 5. Therefore, there are total of 32 possible questions the model may see, and each of these questions occurs with probability 1 1/32. Ignoring punctuation marks, each input sample is assumed to be 16 words involving the list and other words in the inquiry sentence. The correct response (the 'label" for classification) is the fruit that appears most frequently in the list. For example, for the question "Based on the list 'apple, orange, apple, apple, orange', which type of fruit appears most frequently?", the correct response is apple.
Task 2. We again consider simple question answering task with only two possible questions:
Based on the sentence prefer an apple to an orange', which type of fruit prefer? Based on the sentence 'I prefer an orange to an apple', which type of fruit do prefer?
Here, each of the two questions above occurs with probability 1/2. Similar to Task 1, we ignore the punctuation marks, and the input is the 18 words in the sentence. The correct response (the "label" for classification) is apple for the first question above, and orange for the second question above.
Combining all the words appearing in two tasks, we attain a vocabulary with a length of 19: ('apple'. 'orange', 'Based', 'on', 'the', 'which', 'type', 'of", 'fruit', 'list', 'appears', 'most', 'frequently' 'sentence', 'T, 'prefer', 'an', 'to', 'do"). We embed this sequence as a matrix E [el, E2 e19l E 1199199 Then, we know that the length of the vocabulary K and the length of each input sequence N are set as N) = (19, 17), (19, 19) for Task and Task 2 respectively.
In the experiments, we consider a similar transformer model as we introduced in our theoretical analysis. To train the model, we consider Gaussian random initialization V (0) w NO, o²) with σ = 0.01, and we use gradient descent with learning rate 1) 0.1 to train the model. The constant in the log-loss is set as 0.1. Both the training and test datasets contain 1000 samples.
15-
(a) Accuracy
(b) KL-Divergence
(c) Accuracy
(d) KL-Divergence
Figure 7: The results of the experiment for Task and Task 2: (a) and (b) correspond to the experiment
for Task I; (c) and (d) correspond to the experiment for Task 2.
Figure 7 shows the experiment results for Task 1 and Task Figure ((a) and Figure [cc present
the test accuracy. In Figure 7(b) and Figure 7(d) we first normalize the output of the trained transformer model to get a K -dimensioonl vector, representing the prediction distribution of K words. Then, we report the Kll-divereence between this prediction distribution and the true distribution of y|æ,, 222 ..., N in Figure 7(b) and Figure 7 The experiment results show clear difference between the performances of the transformer model in the two tasks. In Task I, the trained transformer model can successfully approach the optimal accuracy (100%) within 100 iterations. However, in Task 2, the test accuracy remains around 50% within 100 iterations.
The results can be explained following the discussion in Section Specifically, in Task 1, the average of the word embeddings E in question can help the model to find the correct response, while in Task 2, the two questions give the same average of word embeddings T., and therefore, it causes inefficient optimization.
The results for these two tasks demonstrate that our theories and explanations for random walks can also guide the construction of various other learning tasks and predict the performance of transformer model in these tasks. This confirms the validity of our theories and explanations and highlights the insights provided by our study.
7 Conclusion
This paper investigates the capability of one-layer transformer model to learn random walks on circles. We demonstrate that transformers successfully learn such walks for the transition probability 0 p 1, achieving the optimal prediction accuracy. We also show that the trained model is

---

interpretable: the softmax attention mechanism effectively selects the correct "parent" token, while the value matrix recovers one-step transition model that applies to the selected "parent" token for optimal prediction. In addition, we identify that the edge cases (p 0, 1) are failure cases, thereby proving the necessity of the assumption < p < 1. Motivated by the analysis of success and failure in learning random walks, we design simple question answering tasks that exhibit similar optimization challenges, showing the broader applicability of our analysis to other tasks beyond random walks. We also provide experimental results to validate our theoretical findings.
In future works, it important to theoretically study the impact of random initialization in learning random walks. Moreover, an interesting future work direction is to extend the results and study the performance of deeper transformer architectures, which may require more advanced theoretical tools. Moreover, extending the finding to more complicated learning tasks, such as random sequences generated by general Markov chains or Bayesian networks, is also an important future work direction.
Acknowledgments
We thank the anonymous reviewers and area chairs for their valuable and constructive comments. Yuan Cao is partially supported by NSFC 12301657 and Hong Kong ECS award 27308624.
References
[1] ABBE, E., BENGIO, S., BOIX- E., LITTWIN, and SUSSKIND, J. (2024). Trans- formers learn through gradual rank increase. Advances in Neural Information Processing Systems 36.
[2] ACHIAM, J., ADLER, S., AGARWAL, S., AHMAD, L., AKKAYA, I., ALEMAN, F. L., ALMEIDA, D., ALTENSCHMIDT, J., ALTMAN. S., ANADKAT, S. ET AL. (2023). Gpt-4 technical report. arXiv preprint arXiv:2303.08774.
[3] BIETTI, A., CABANNES, V., BOUCHACOURT, D.,, JEGOU, H. and Bortou, L. (2024). Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems 36.
[4] CAO, Y., HE, Y., WU, D., CHEN, H.-Y., FAN, J, and LIU, H. (2025). Transformers simulate mle for sequence generation in bayesian networks. arXiv preprint arXiv:2501.02547.
[5] CHEN, L., LU, K., RAJESWARAN, A., LEE, K., GROVER, A., LASKIN, M., ABBEEL, P., SRINIVAS, A. and MORDATCH, I. (2021). Decision transformer: Reinforcement learning via sequence modeling. Advances neural information processing systems 34 15084-15097.
[6] CHEN, S., SHEEN, H., WANG, T. and YANG, Z. (2024). Unveiling induction heads: Provable training dynamics and feature learning in transformers. Advances in Neural Information Processing Systems 37 66479-66567.
[7] D' ANGELO, F., CROCE F. and FLAMMARION, N. (2025). Selective induction heads: How transformers select causal structures in context. In The Thirteenth International Conference on Learning Representations.
[8] DOSOVITSKIY, A., BEYER, L., KOLESNIKOV, A., WEISSENBORN, D., ZIAII, X., UN- TERTIINER, T., DENGHANII M., MINDERER, M., HEIGOLD, G., GELLY, S. ET AL. (2020). An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint 1mXiv:2010.11929
[9] EDELMAN, E., TSILIVIS, N., EDELMAN, B., MALACH, E. and GOEL, S. (2024). The evolution of statistical induction heads: In-context learning markov chains. Advances Neural Information Processing Systems 37
[10] HE, H. and Su, W. (2025). A law of next-token prediction in large language models. Physical Review 112 035317.
[11] HUANG, Y., CHENG, Y. and LIANG, Y. (2024). In-context convergence of transformers. In Forty-first International Conference Machine Learning.
10

---

[12] ILDIZ, M. E., HUANG, Y., LI, Y., RAWAT, A. S. and OYMAK, S. (2024). From self- attention markov models: Unveiling the dynamics generative transformers. In Forty-first International Conference on Machine Learning.
[13] JANNER, M., LI, Q. and LEVINE, S. (2021). Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems 34 1273-1286.
[14] JELASSI, S., SANDER, M. and LI, Y. (2022). Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems 35 37822-378366
[15] JUMPER, J., EVANS, R., PRITZEL, A., GREEN, T., FIGURNOV, M., RONNEBERGER, O.,
TUNYASUVUNAKOOU K., BATES, R., ZIDEK, A., POTAPENKO, A. ET AL. (2021). Highly accurate protein structure prediction with alphafold. nature 596 583-589.
[16] LI. Y., HUANG, Y., ILDIZ, M. E., RAWAT, A. S. and OYMAK, S. (2024). Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics. PMLR.
[17] LI, Y., LI, Y. and RISTESKI, A. (2023). How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning. PMLR.
[18] LI. Z., CAO, Y., GAO, C., HE, Y., LIU, H., JASON, K., FAN, J. and WANG, M. (2024). One-layer transformer provably learns one-nearest neighbor in context, In Advances Neural Information Processing Systems.
[19] LU, C., SHI, R., LIU, Y., HU, K., DU. and Xu, H. (2024). Rethinking transformers in solving pomdps. In Forty-first International Conference on Machine Learning.
[20] MAHANKALI, A. V., HASHIMOTO, and MA, T. (2024). One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations.
[21] MAKKUVA, A. V., BONDASCIII, M., GIRISH, A., NAGLE, A., JAGGI, M., KIM, H. and GASTPAR, M. (2025). Attention with markov: A curious case single-layer transformers. In The Thirleenth International Conference on Learning Representations.
[22] NICHANI, E., DAMIAN, A. and LEE, J. D. (2024). How transformers learn causal structure with gradient descent. In Forty-first International Conference Machine Learning.
[23] RADFORD, A., WU, J., CHILD, R., LUAN, D., AMODEI, D., SUTSKEVER, I. ET AL. (2019).
Language models are unsupervised multitask learners. OpenAl blog 19.
[24] RAO, Y., ZHAO, W., LIU, B., LU, J., ZHOU, J. and HSIEH, C.-J. (2021). Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems 34 13937-13949.
[25] SVETE, A. and COTTERELL, R. (2024). Transformers can represent n-gram language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I: Long Papers).
[26] TARZANAGH, A., LI, Y., THRAMPOULIDIS, C. and OYMAK, S. (2023). Transformers as support vector machines. arXiv preprint arXiv:2308.1692
[27] TARZANAGH, D. A., LI, Y., ZHANG, X. and OYMAK. S. (2023). Max-margis token selection in attention mechanism. Advances in Neural Information Processing Systems 36 48314-48362
[28] THRAMPOULIDIS, C. (2024). Implicit optimization bias of next-token prediction in linear models. Advances in Neural Information Processing Systems
[29] TIAN, Y., WANG, Y., CHEN, B. and DU, S. S. (2023). Scan and snap: Understanding training dynamies and token composition in 1-layer transformer. Advances in Neural Information Processing Systems 36 71911-71947.
11

---

[30] TIAN, Y., WANG, Y., ZHANG, Z., CHEN, B. and DU, S. S. (2024). Joma: Demystifying multilayer transfomerss via joint dynamics of mlp and attention. In The Twelfth International Conference on Learning Representations.
[31] VASWANI, A., SHIAZEER, N., PARMAR, N., USZKOREIT, J., JONES, L., GOMEZ, A. N., KAISER, L. and POLOSUKHIN, (2017). Attention all you need. Advances in Neural Information Processing Systems.
[32] WANG, Z., WEI, S., Hsu, D. and LEE, J. D. (2024). Transformers provably learn sparse token selection while fully-connected nets cannot. In Forty-first International Conference on Machine Learning.
[33] ZHANG, C., MENG, X. and CAO, Y. (2024). Transformer learns optimal variable selec- tion in group-sparse classification. In The Thirteenth International Conference on Learning Representations.
[34] ZHANG, R., FREI, S. and BARTLETT, L. (2024). Trained transformers learn linear models in-context. Journal of Machine Learning Research 25 1-55.
12

---

A Additional Related Work
In this section, we give overview of some additional related works.
Token selection. Our work reveals that a one-layer transformer can learn to perform the optimal token selection by focusing on the direct parent in random walk. A line of recent works has studied the token selection of the self-attention mechanism from different perspectives. [26] 27 propose an equivalence between the optimization dynamics of one self-attention layer and an SVM problem and prove the global convergence under certain assumptions. [16] shows that when training self-attention layer, the priority in token selection is determined by directed graph extracted from the training data. [32] demonstrates that transformer models can learn the sparse token selection task effectively while fully connected networks fail in the worst case. [18] shows that self-attention layer can be trained to perform proper token selection so that the model acts as a one-nearest neighbor classifier in context.
Next-token prediction. 28 explores the implicit bias of next-token prediction employing a related SVM formulation. [99] demonstrates that transformers fail solve the Partially Observable Markov Decision Processes problem (POMDP) even with sufficient data. [10] observes a phenomenon of next-token prediction in LLM that each layer contributes equally enhancing the prediction accuracy. 29 studies the SGD training dynamics of a transformer with one self-attention layer and one decoder layer for next-token prediction, restricted to some specific assumptions like positional encoding, long input sequences, and the fact that the decoder layer learns faster than the self-attention layer.
Training dynamics of transformers. [20 34 investigate the training dynamics of in-context learning in transformers with single self-attention layer trained through gradient flow on linear regression tasks. [I]] solves in-context linear regression with the orthogonal input data by gradient descent single softmax attention layer. 14] demonstrates that the position-position block of single attention layer in vision transformer
encode spatial structure by dealing with binary classification task. [0] delves into the training process of transformers with multi-layers by analyzing the dynamics of the MLP layers. 33] analyzes synthetic in-context learning task and emphasizes the significance of weight matrices as associative memories. [ shows incremental learning dynamics in transformers with diagonal attention matrices.
B Gradient Calculation
Recall that the population loss is
L(0)- [] = lo-, (X)
Ee] + c)] = B[- log(e, VXX(X (iXS(X"Wxx)) + e]]-
The following lemma calculates the gradients of the population function.
Lemma B.1. The gradients regarding V and W are
N-1
1 VXS +E E, Σ
N-1
Vvl(0) = = -
S
S
e,
1=1
1=1
N-1
S
N-1 i==
S Em= Si
VTepP
i==
V Vw((0) = l.
NN Sip.r
Eim S:p; NN S
VTepP
EN'T
-EN S.w. E S
V VT Tepp
e, VXS +
EN
N-1
S.p.aT
-E S.p. EN S V vee
where S = S(X Wãn), and S; is the i-th element of S.
Proof of Lemma B.l For V. we have
1 de ; Vvl(0) e, so(
VXS
+
OV
e X
e VxSx+
13

---

N-1 S
e eVVXS+ + ey
For W, we have
de, V VXS(X WāN)
(l((0)
ey (x()) +E
ow
X W X V CYIN
e VXS+ +c
X(diag(S) SS 1 V V CyIN
VXS+ +e
N-1 S
-N S.z; 2 N
N
S
1
0 V Teeppl
e, VXS++ te
S.p.a
-2 Sip: 2i=1 S
N-1
S, T
Σ
N
S.m. Σ N
S T V
CyPN
VXS
-==1 S.p:a. Li=1
Sip: 2i=1
S VV eyP,
where the fact that S'(X'
WwN) (diag(S) SS and
S1
S2
diag(S)= =
SN_
To simplify the notation, we denote
N-
--1
N-1
N-1
=
S
S
Σ
S T V TeePp
1=1
1=1
N-1
N
N-1
B
S.p.a.
S.p.
S V eyPN- VTeppk.
i=1
i=1
With this notations, by Lemma B.1
we have
0
0
A
Ww e,
(B.1)
VXS
0
B
We also denote A (() the corresponding matrices at the 1-th iteration of gradient descent.
Moreover, we can observe that W
W11
W12
where Wn E I KxK Wi2 E RKXM
W21 21
W22
W21 MxK and W22 E MMMM By B.T we know that w 0кжк and w 0мжк
for all 2 1.
By the definition of the random walks circles, write the transition matrix II as II pll, + - (1-p)In, where
0
0
1 0
1 0
C Proof of Theorem 3.1
In this section, we analyze random walks with the transition probability p satisfying 0 < < 1. Without loss of generality, we assume Sp 1. We consider gradient descent starting from zero initialization. The following lemma presents the result of the first iteration.
14

---

Lemma C.I. Under the same conditions as Theorem 3.1 it holds that
N-1
V
T
Σ
N-1 and W OKK+M)(K++)-
cNK
1=1
Proof of Lemma C.1
By Lemma
B.1
we have
1
E
[[v(((0))
Ble 1
e
N
BI
2,,0 N=,00,00
N-
eN 1=1
N-1
y
NNK
where the first equation is by the initialization ( and ( the second equation is by the sampling method, and the third equation is by E[ R Lx for i [N - since æ; is uniformly
distributed in E. Thus, by the update, we can get
V (0)
n )((()))
N-1
II NN
NNK
Since V(0) 0K and W O(K+M)) we can get E [[ ((X+M)((K+M) Thus,
( W 1 w(( = = ((K+M))(K-M)
Lemma C.1 gives explicit calculations of (1), Based on these, we can further derive some properties of ( given in Lemma C.2 below. In this lemma, for all matrix indices, we consider simplified notations following the rule that an index i is not in the {1....,K} K), we treat as (l...., K} with i' = (mod K).
Lemma C.2. Under the same conditions Theorem 3.1 it holds that (V [V(1) i - i2 j2 (mod K).
Jiajz for
Proof of Lemma C.2
We induction prove that for any R N. II 1i.
R ja for
f1 12 j2 (mod K). The result is obvious R = 1. Suppose that the results hold for II
We aim to prove the results hold for I By the definition of II, for any i,j. we obtain that
R+1 li. p- II 16--1 + (1 - p)) J6,541
(C.1)
By induction, for any i1, i2. j2 satisfying i - is (mod K), we have I 1i
12..2 [1++1] Jajz. which completes the induction.
I Ja-a-1 and I b Thus, by C.I we can easily get [II
R
From Lemma C.I know (V())T= 11 DE has the same property.
N-1
holds the result. Therefore, v also
Lemma C.3. Under the same conditions as Theorem 3.1 it holds that (V = ( maxx
Proof of Lemma C.3
Without loss of generality, we assume < p < 1. From Lemma C.1 we N-1 (II I We can observe that since p 1, with increasing R, (II )R
know (1)
will be closer to + 11
as stated in Lemma E.2 and E.3 Thus, the location of the largest entries in (1) is mainly determined by the first several terms. In II there are two locations with non-zero
15

---

values 1 In (IT )² there are three locations with non-zero values p², 1p(2p(1 p), (1-p)- p)?. We can easily check that p is larger than p² and 2p(1 - pp). Thus, know the location of the value p in II" is also the largest in V i.e. |2,1 detail, we can first get that
N
II m
N-i
(IT -m
matrix of interest:
N-2
6, i=0
(e))..
Ie))
N-2
e
Σ
II
y e)). (e)
Denote
N-2 T(N)
Σ
1II
8=0
Claim: for all N, the diagonal entries in ((N) are larger than or equal to the other entries.
T(N) = ITN +)+ I.
Induction hypothesis:
T(N),1 2 T(N)as) + C
Suppose the induction hypothesis holds for I(N - 1).
For i with j - jl 2 1,
I(N), = p- I(N +1 (T(N 1)) S (T(N-1) + Cp..
T(N),1 T(N 1)1,1
E 1= k (II )² where Cp is a positive constant p
N-1 (ITT) (IT
- -
regarding Considering the terms IT (III) )k with 0 S A: S K - 1, we can easily observe that the coefficient of II (II,) ) is much larger than that of other terms. Since II is a cyclic shift matrix and the entries in (II N are almost same (Lemma E.2 and E.3 the location of the largest entry in II is also that in (1), Therefore, holds that [ 1|2.1
max-
Further, the following two lemmas provide some properties of the weights for the second iteration.
Lemma C.4. Under the same conditions Theorem 3.1 it holds that (| [Iaxx
2cK2
Proof of Lemma C.4 First, we have
ey Σ N-1 x== S E
( (1)
(() + +n
- (1) Σ S 2
Thus,
ey Σ ==
(2)
max
|
max
nE
e
V N m;
x
max
16

---

N
N-1
7
B
max
N-1
N
miniy
[(1)
==
N-1
N
1
II
ENK
N-1 min.j
y
2i
N
II
N
1=1
max
<NK
1.
11
+
+
eK
where the second inequality is by e, 2. 2 m and the last inequality is by
Lemma E.2 and E.3
Lemma C.5. Under the same conditions as Theorem 3.1 it holds that sk= 2 S.e exp(I((V)) for j N 1. Further, S 2 -
and s I exp( (2()) for j * N -
Proof of Lemma C.5
By Lemma
B.I
we have
E E[A
[1
N
N-IN-1
E
S
a
(V(1)
By
Σ
s( S T -
(V(U)) ey
PNN
1=1
4=16=1
N-1
N-1
N-IN-1
N
11
Σ
II
By
T
Bi T 42
Σ
I1
By
PN
CN K
cN3K 11 i==1
N-1
N-1
E
D
III
NN
PN T
EN2K
t=1
2=1
N-1N-1
N-1
11
II
III' ) N
PN
EN
12
N- N-1
E
B. II
J 0i PN
62K
N-1N-1N-1
E
"
Σ
Σ
T 2
II N- -1
(II N- N xis
PN
CN3K
N-1 N-1
tr
nI
(II INi
1KPN
P=1
N-IN-IN-1
7
ΣΣΣ
II N- (II N 1
where the second equation is by Lemma C.I the third equation by the sampling method, and the fifth equation is by the fact that all the x; is uniformly distributed in E for E [N - 1). Then,
w (2)
( nE w ((())) x T. Thus, We also have
(1)
N
N-1
E
(V(1))
Σ
St
T
(V(W)) ey
PN
1=1
t=1
N-1
N-1
N-1
E
II
N-
n
Σ
ps Σ
x
Σ
II
-P ey
PN
CN2K
eN3K
#=1
i=1
2=1
N
N-1
II
N-P (II
P
EN2K
Y=1
17

---

N
N-1
N-1
E
EN
17
K
pi
Σ
II (II N x, PN
1=1
F=1
N- N-1 1
P: tr
II
N-P
)mー)
PN
Tr
<N2K²
1=1 P=1
N-1N-1
7 <
P ΣΣΣN
II (IT)) ) pÑ:
where the second equation is by Lemma C.1 the third equation is by the sampling method, and the last equation is by the fact that all the x; is uniformly distributed in E for i E [N 1].
Since
X W ēN N PN W (2) PN and [ = w PN + P. w (2) 22 PN for j e
N - 1}. we can obtain that {1, 2,
X (2) X
= PN 22 PN
T PN ey VJVSS+ +
B B
E
PN
N-1
E
1
-S
PNPN
Σ
s 0.
(Va)))ee By DNNN
VXS+6
<0,
where the third equation is by P. P3 = 0 for i * And for j E (1, N 2]. we can get
XW
(2)
XW (2)
N
N-1
12 PN + +pT
(2)
W (2)
22 PN xj W PN P3 w PN PN
PN - -P, W W2 pn
PN-
(i)
B.y V X 46" N (PN-IPN-- (V()) e, PD (,,x,(((l))) ev) pkPN
(10)
c, (S()) +E N PN DD-IPN--IX- PI-- xN- - PP 2 (V()) (II -1 PNPN
7
7
(pNDN)2
((V(0)TITT 1.1 - [( 1.1
e, ( XS()) +E N
IV3
n
(DNPN)2
[(V)) - ((()) (1 11
E
max, (V)) Jis +E N
11
eN2K O
ck 6
M Ω
N2
N ((N),
where (i) is by w x I. (i) is by p. pr for * 7. (tii) is by the sampling methods,
(iv) is by the fact that all the x; is uniformly distributed in E for E [N - 11, (v) and (vi) are
by Lemma E.4 Therefore, we have s // = exp [X (2) aN
XW 2
N-1
exp(S)(N)) for j F -
1. Further,
(22)) NN =1 Σ 8 > 1 - (N- 1) exp ))xx)(-((N))))
-
18

---

which implies that
N 1
(2) NN-1
(xp((((N) + N -1 exp( (P))))
I+ - cxp( (9(())
Then, have s/00 1
5 N-1 S (xp(<x((N) for j # N --1. 1.
Then, we can derive the bounds of entries in (
Lemma C.6. Under the same conditions as Theorem 3.1 holds for 2 3 that
min ]. 2 2cK²
and [V(1) max
ek + (t - 2) Kc
Proof of Lemma C.6 First, have
min
vin[
2
min w
V
li
N-1
n
E
min w
cNK
π
N
<NK
2K
2
where the third inequality is by Lemma
F.2
and
E..3
Then, we can get that
N-1
max
(t-11)
E
(t--1) 1) Ei=1 S:
mey 1= S 27
s((-1)
Ry + 6
max
By E 5, (((-1) T
17
(t-1)
Imax +E
min
By
Σ 5
(t-1)) Imax
1)
min:,
V(t-1) li,
max + 22eK2
(2) maxx + (1-3) K2
1 + (t 2) 2cK
eK
where the third inequality is by ey Væ; 2 min; [V Visg. and the last inequality is by Lemma C.4
Next, we analyze the training dynamics over multiple iterations.
Lemma C.7. Assume the same conditions as Theorem 3.1
For 2 < t 5 T", it holds that S 2
1 - exp(-S((N) and ( ((*) II + (9) where
(+)
I (t), Here, 8(t) vnt cK
2n
max
and 7()) < ck + 2(t 1)cK²N N exp(
(NN)).
Proof of Lemma
C.7
We use induction to prove the results that
g(0) 2 Vnl KK'
< 2n
(9)
2((
1)eK exp(-S((N)),
EK
X W (1) (9)
[XW[OxN] 2 S(N).
N-1
(()) N-1
2 1 - -exp(-
19

---

It can be easily checked that the results hold for = 2. Suppose that the results hold for the t-th iteration. We aim to prove that the results hold for + 1.
For V(++1, we get
v(++) = V(O) B v (((()))
By EE ii 5 D2
N-1 (1)
V nB
e (1) Ei=1
N (
0i EE_
= V +E
n N-1 e N
(t)
+
new
Σ N-2 (1)
Si x
e V (4) N-1
e (4)
N-1 8.00 04 +
1=
1==
S CN 1 1
nyy E 5/00 x
V +E
II
+
E
B.y V(t)
Σ N
6 V(O)
Σ i=
N-1 5 m; +
i=
Then, we have
(1)
(1)
N-N-INN N
S N-1
e (+) Σm 5(t) x3 +
N-1
(1) max +E
111
300 + 7()) +e
2 (300 + ck 2cK2 + 2tck exppexxpp(pp00 +
1
(C.2)
2 (3tt) +
2n ck
where the first inequality is by e, V x; - | mmax: the second inequality is by induction, and the third inequality is by the assumption of And, we have
mey Σ S!00 x
nexp( (XxP(-2(N)) min,.
N-
E
E
Σ
eyat
e
N-1 s(t) V Σm 1=1 S
+
6. 22 exp((S((N)) .N,
[
max
max
(C.3)
where the first inequality is by induction and by Lemma C.6 Thus, we can get that
v(t 2 min,, (V kig, and the second inequality e T
n N- LN N
g(++1) 2 900
+
E
ee (())
EN S:00 x. + 1.1
8(0)
11
2,300 +
47
r
27
VTI
KK
2/
2n
VT
Vnt
eK
Vt V++I+ Vt
21
n(1 + 1)
cK
where the second incquality is by C.2 and the third incquality is by induction and the fact that
2x+ 47 is monotonically increasing for x >
7
o 2n And, we can get
/2
Rk
Σ N (t) 1=
(1+11
< <7(0) +
E
ne,
x
( e. (1) ΣN=1 8 x; +
max
<,10 + 2eK2N N
20

---

2n
- eK
+ + N exp(-9(N)))
where the second inequality is by (C.3) and the third inequality is by induction.
Next, we consider 1(++1), Recall that
A (1)
B (1)
w = w +
+n Eyy
and W (1(++)_ = W
(2) 22
+
+n
12
VXS
+e
eyy
VXS +
where
N-1
N-IN-1
10
Σ
(V(1)) ey - Σ x
Σ
s(t)
( Iii 2 (V(0)
ey
pkv
12
N-1
B
Σ
S D
T (Vili)Te,, Ey Σ
S
Pj
Σ
S vi [V( eyy
PN-
We also have
X WWO NN P W 22 pn and
XW =
w (2)
12 PN +P
W (8)
22 PN for
j E (1,2.....
N 1--1 13. Then, for j = N, we have
p B PN
WW Px 22 PN = PN W PN + B
e (())
N-1 5/ 2j +
Σ s! x e
= PN 22 PN (
T W 10)
e (+) Em= 5,00 x, +
S p with PN.
For j E (1, - -2], we have
x: A 10 PN
3 W ((11)) (1+1)
12 PN x; w 12 PN + B
e V())
E S x. +
V e, -N-1 ((1)
12 x vitee
= = w
12 PN /
e, (()) E (2)
Di +
( max
$ x w 12 PN + + MS
minig [V(') bn
2K + 4(t 2) K
'. T W (1)
12 PN // M exp(-x((N))
"
where the first inequality is by e, V and Lemma C.6 And,
100NN11 = (
maxx and the second inequality is by induction
p3 B PN
T W P.
22
(t+1)
PN
D 22 PN + +B
W (t)
y.y (())
Σ Ei=l 8/0 +
N-1
Viee, - -E N-1 5(00 x2 WWeey
= =P 22 pN + M
w (t)
e, V()) N 5:00 C; +
|
< p 22 PN + nM
T
minej(V())
2K 4(t 2)e² K
Sp T w (0)
pN + M (.MMex(p((())
"
where the first inequality is by e, V (0 -I-1= [|((t) max. and the second inequality is by induction
and Lemma C.6 For j = N 1. we have
T A (1)) PN
x W (++1)
PN - W (t)
N
PN + 1
e, (t) E S: x + C
21

---

IN Viye, By - - N-1 5(0) T V
Liz=1 vig
= = W 12 PN + T/ M
E
V Ei S:¹² 0i +E
S 2, V ey
( 2 2 12 PN + ME
2iz=1
v (+) Ei= N-1 5 a, +
N-2 (5)
V Imlma Imax
2 x W PN T/ M
11=1
12
min:. [V()]
2 .N PN 1/V NN ((NNexp(-S(()) 2K + 4(t - 2) K*
"
where the second inequality is by ey V | maax and the third inequality is by induction and Lemma
C.6
And,
P T
B (4)
PN
W 22
PN = PN W 22 PN + T/EE
PN
e (()
2i=1
N 8 a; +E
2 FN v(0e, ey DN-11 S: x
Viee,
( 22 PN + nM E
PN
e. (()) Σ N ((t)
2i +
EiS 5:00 xz Viyey
2 P W
22 PN + ME
e.y V
N-1
-== S xi +
N-2 (!()
T P W PN nM
i= i/
Imax
min;,1 [V
2 PN PN MMN exp(- 2K + 4(t 2) Kd
where the second inequality is by e, ZN-1 |V(0 maxs and the third inequality is by induction Therefore, we can get that for j F N 1,
and Lemma C.6
(W++1)
( T.N
EN
N-1
(1+1) ((+1))
PN + pT. 1 (1+
PN x3 w (++1)
PN -P, w- (++1)
EN
22
PN
NN PN + - w Px az W
(( (t)
(1) PN P) w 22 PN
A(t 2)e² K
4 (P(() 2K +
XW EN
XW eexp(-S(1 N-1
2 2(N),
where the last inequality is by induction. Thus, s(t /S > exp(S((N)) for j / N - 11 which
implies that s > 1 - exp(- Therefore, we prove that the results hold for t + 1, which completes the proof.
The next two lemmas show the convergence rates of Lemma C.8. Assume the same conditions as Theorem 3.1 For 2(n 2 2 -2- < T < T holds that
Ir and for (X) (X) 1|2-
(
S O
V
VT
Proof of Lemma C.8
By Lemma
we can get that
(T)
II
3(T)
(T)
II
(T))
F
F
22

---

3n
v
<1
II
+
|
For the first part, we have
B(T)
3
|
II
(
+ 1|
||
|1A111l ||
KK (1 p)))3(()) + ||
Ky(T)
(2² + (11 - p)2) + Ky(T)
21 + +2Tek exp( (())
V2K
(VTT
2n
2 + 2Tck3 N exp(-2((N))
2
s
where (i) is by have
|I
|| = / k(p² + (1 2)) and (ii) by Lemma C.7
For the second part, we
VTT
V ||
V
II
F
I13 || + |
KT(T)
v2K ((T) Ky(T)
(5)
2n + 2T2233 N exp(-
/2K
(VTT-
2 + 2n + 3 exp(-S((N))
ck
1
O
where (i) is by we can obtain that
|
||
V k(p2 + (1 - p)²) 2 /2K /2, and (ii) is by Lemma
C.7
Therefore,
V(T)
II
0
V
F
VT
Lemma C.9. Assume the same conditions as Theorem that
3.1
ST ST. it holds For 2
farry (X)
II -N-I
5
[ (X)l12
Proof of Lemma C.9 The output with 0 = OC0) fXcTi()
V(T) w( (
(T) EN-11 S m; Then, we can get that
Jocr (X)
II EN
[ (X)12
23

---

(8n (1L7))) + +V(T) Σ 1 S 7
EN
v(n) Ei=l 5/¹7
2
II N-
V(T) 1- 5: 32
B T Σ iis1 5 2, + (TT)
N-2
())
Σ 11
(
21
+
V(T) Σ N-1
5 3
i=
For the first part, we have
-1
II NN
V(
N
S!T
0;
3(T) $
11 EN
BITIII Em= 5 (
VIT) Σ
N-1
5 Σ; (T)
Zi
1=1
B(T) exp(- [((N)
x2 ((T) + / (T((T)
(V1jI 27 [1 exp( (
'V12 + (1 p)
(VijI おt + K K 1 + 2T'KKNN exp(-
K (ch 2Te N Nexp<<(exx(((
0
Vπj + 2TekK exp(-s((N))
where the first inequality is by Lemma C.7 the second inequality is by Lemma C.7 and | æ.1l2
2 + (1 v)2 and the third inequality is by Lemma C.7
For the second part, we have
B(T) 1-1
S.7 0. + (T)
Σ N-1
((T)
x;
V(T)
Σ 11 S x;
(T)
ΣmT 1
S (T))
+
(T) >
(TT) x:
300 N- S.T)
0;
exp( ((N)) max v (Ky(TT
V12 (1 B())
exp( (P()) (.kk +2TcK2) V
2n
+TTc exp(-x(())))))))e
lek
V0² + (1- p) (VTI
2n e
where the second inequality is by Lemma C.7 and | æ1l12 = VD2 (1- and the third
inequality is by Lemma C.7 Therefore, we can obtain that
Sacr (X) [|ferry (X)ll2
1
II -N
50
(TT
2
24

---

D Proof of Theorem 4.1
In this section, analyze "deterministic walks" with p 0, Without loss of generality, we set p == 1, which means II HT The following lemma shows the results of the first iteration.
Lemma D.I. Under the same condition as Theorem 4.1 for any loss function E(.), it holds that
V (0
T NK Kxx and w (KK+))((K+M))
Proof of Lemma
D.1
By Lemma B.1 have
N-1 E(eyar.] 1
E. [V = (0
N 1
= =
B[ 100000000
20:00
N
N-1
=(
(o)
NK
= =
1 1xxK,
NK
where the first equation is by the initialization ( and (0),, the second equation is by the sampling method, the third equation is by [[ & Ix for E [N - 1J. and the last equation is by Lemma E. Thus, by the update, we can get
( ( T [(0(0))] (((0)))
Tpr
1
NK
Since V = 0x and (0 :::0==0(x++M0 ((K+M))(K*M)) we can get E V ((()))) = O(K+M)((K+M)) Thus,
W w( nE [V (K---)
The following lemma states the results of the second iteration. Lemma D.2. II II,, then it holds that
(2) = - ((((0 + ((O(1)))
N 1xxK,
NK
172,2
w DE je
IKPN:
T
N3
n²²r
N-1
172
w ((0
N
Σ
Pt
N3
PN
Pi--
1=1
Proof of Lemma
D.2 By Lemma B.I have
N
=
E V = (0())
E(eyæ,] !
N
= =((o
NNN
B[III
y N x
N-1
= =)(O())
NK Σ (Il) N-i
= =(()())
NK 1
25

---

where the second equation is by the sampling method, the third equation is by E[ = K Ix for i [ 1], and the last equation is by Lemma E.1 Thus, we can get
v v v
= -
1 - (( 1XKK
TT
NK
NK
By Lemma B.I
we have
N-1
N-1N-1
1 =E E
Σ
DE, (V())
Cy
Σ
Σ
5(1)
s(l) 12
2i, Di
(V()))e
PN
i 1==11=1 $1
E
((
T
2 1K
1p
Σ
Σ
Tiz 2 1k
T
PN
N2K
N3K
N-1
N-1N-1
= E
)((O0)
THe
Di
qpr
Di
PN
T
N2K
N
==1
11=1
1pp
mp (N - 1)
(((1))
1K
1x
PN
N2K
N
?p?
(0(0)
N 1KPN.
where the second equation is by Lemma D.I. and the fourth equation is by the fact that all the æ; is uniformly distributed in E. We also have
-1
N
N-1 1
E B
(1) =
(V(1))
T
((1)
p:
S 2.
T
(V())
ey
Piv
i=1
1=1
N-1
N-1
-e(orm)
= =E
N2K
p 1x
N3K
pi
x
T 1
PNN
t=1
1=1
i=1
Tp
gr(N- -11 1)
Σ
N2K
Pi
N3K
Pi
PN
N
(0())
N3K
qP
Σ
pi
N3PN
PN-
where the second equation is by Lemma D.I Thus, we can get that
w
w
nE V (((
12
E Ine (B(1))
A
12222 N3K
(0
1KPT.
and
w
n w (((((()))) 122
- E..
nE (0(1)). B(1)
n
N-1
1222
-((0
Σ pt
N3K
N3
PN
Dir.
Next, we can analyze the gradient descent dynamics over multiple iterations.
26

---

Lemma D.3. I[II= = I2,, then for any t 2 0 and any sequence of learning rates } it holds that v(t 1xxK, and 5 = !
= S N-1-
(()
Proof of Lemma D.3 We use induction to prove that for some scalar ai .2 : 9g 04 it holds
(2) (8)
that for L 2 2, V =a
(t) 1xxk, 12 a (1) kpp.. and W - (a3 (1) Σ P: as PN PN-
By Lemma D22 we know that the hypothesis holds for t 2. Suppose that the hypothesis holds for
t = t'. We aim to prove that the hypothesis holds for t = t' + 1. We have
W
12 PN Pi W PN
222
12
PN
+ P2 W 22 PN
XW
W( (t))
12 PN PN- W (e)
PN
Px W
22 PN
(t')
PNPN + 033 PIPNPN
(t')
02 PNPN + 033 D2 P2PNPN
(D.1)
(t') T PPN PPN-IPPPN
02 PNPN + ag
-a4 (PNDN)?
Since pi Pı = =PIP2 = = PNPN, we have
[W (1)()) Ex]1 = [ (() =
(WW')x))-- Thus, we can get that s(1)) have
S2
S/22 = s(()) Then, we
N-1
V ((((((")))] =E
((0())) ey Σ
s(O)
T
=
Ble
=((O)))
E y N x,,0] 1
E
s(())
= =((()))
(II) N-i
K
1=1
s((')r 1KxK,
=(((()))
K
where the second equation is by the induction, the third equation is by the sampling method, the fourth equation is by the fact that 2; is uniformly distributed in E, and the last equation is by Lemma E. Thus, we can get V (('++1 +1) ((() n/ V v v((o(t)) x 1xx K. We also have
N
N-IN-1
E PIACO (00)
E
(V(')) By - Σ Σ
S mi Th Tr,, (V(l)) By
PN T
N-1
N-IN-1
x 1K ai (s()))
E
Izs x to 1k
p.v
1=1
11=1
N-1
W-1N-1
E
a
(1) ((())
Σ
x; 0
(8) (8
Σ Σ Di
PN
11 1=
(1)
(t)) - 1)
(s(')) 1)2
(1))
a
k
1x
pN-
K
K
27

---

where the second equation is by the induction, and the fourth equation is by the fact that all the æ; is uniformly distributed in E. And,
N-
N
N
B
(1))
pia; (v())
-
Σ
S(U
pj
S:C
zi
(Vic))
PN
i=1
1=1
N-1
N-1
IF
a1
(1) ((()
pi 1x 01
(1) (s(())
Pi
x;
T
1
Piv
i=1
1=1
i=1
N-1
(1) a(l))
0i (1')
(s(1))
Σ
pi
Pi
PN-
where the second equation is by the induction. Therefore, we can get
w/2
(t't +1)
w n/ V w ((0(())) 12
12
=a2 (t') (t') 1KPN (0(*) (B(A
[a. (1) (())(()) - 1)
(1) ( -1)2)
=a (1)) 1xp (0(()))))
1 1KP
K
K
(1++1)
KPN,
= 02
and
W (((+)) (f++)
+1)
w -VE/V w ((((((())) 12
N-1
a
(1)
Σ
pi aj PN
(0(1) (B(
N-1
(1)
Σ
(e)
p.
ai PN
PNN
1=1
N-1
N 1) Σ
(0(())
a
(0))(())
Σ
P: ai
(c)()))))
(8))
P: P
1m1
2=
(t'+1)
Σ P:
(1'+1) as
PN
px-
Therefore, by induction, we can conclude that for all 2 2. V ai (d 1xxk. W 02 (1) Ikpy and wi(t)
03 (1)
Di=1 p: - a
(((''x') PN)) pl. Similar (D.I) we have [X W (000
[ ØNl2 =... (X which implies that 5 S NN
(t)
E Auxiliary Lemmas
In this section, we present some auxiliary lemmas. The following lemma states the properties of IIo Lemma E.1. By the definition of Il, holds that II IK. H. IK, and ΣK K 1x-K.
II
Proof of Lemma E.1 In this proof, the index i larger than K represents i - K. For Ho, only [[o]i+ 1 for E [K] and other elements are O. We can get that for ПŽ. only [1] = for i E [K] and other elements are O. By this observation, we can derive that II I, and ΣK I 1K*K. Also, we have II II so we can get IIII = II = Ik.
Further, the following three lemmas show some properties of II with transition probability p satisfying 0 0<p<< 1. 1.
28

---

Lemma E.2. Assume that K is odd. It holds that
8p(1 p)R
R
exp
1.
K
K²
Proof of Lemma E.2
has eigenvalues Ao Ax- where
2nth
2uit
2
2
A pe
+ (1 - p)e
cos
+i(1 2p) sin
K
K
with
2
16p(1 - p)
8p(1 p)
Au = 1 4p(1 p) sin2
I V
K
K2
K
for & # O, where the first inequality is by sin(2 K) 2 sin(Ti 2/K. The eigendecomposi- tion of each entry in II can be written as
K Σ
1
II
Aki
= do + Σ Ak.
K
K=
Kę0
where C,.i
1 2mi(i- VR
1)b/K
vR
2 11)/KK
k
1
2)///K Then, we can get that
K-1
1
by
=
C
+
CR
K
k=0
Thus,
1
I1
R5
Σ
Ck.i.j
A/
K
k
[ck. M1 R
k
8p(1 - p)
R
<
ck.a.f
K2
8p(1 - p)R
<exp
K2
where the second inequality is by the bound of the absolute values the eigenvalues, and the last inequality is by 1) Ri for any 0 < and |ck.i.l for all k,ū,j.
Lemma F.3. Assume that K is even. For the case that R is even,
for odd (j - i);
2
8p(1 p)R
< exp
for even (j - i).
K
K2
For the case that R is odd,
8p(1 R)R
< exp
for odd (j - i);
K
K2
f = 0 for even (j i).
Proof of Lemma E.3
II has eigenvalues Ao, Ax A1-1, -1, where
d
2atk
+ (1 - p)e
2etk k
= cos
2mr
+ i(1
2p) sin
2mm
K
K
29

---

with
2 16p(1 -p)
8p(1 - p)
Axl 1- 4p(1 - p) sin?
V
K
K²
K²
for 1: + 0, K/2, where the first inequality is by sin (2 K) z sin( K) 2 2/K. The eigendecom- position of each entry in II can be written
Σ ChA
1
1o +
(-1)*t)
AK/2
Σ
K
K
k/0,K/2
where Ck,A, = 1
VK e
2 (acff-1160 1)K/K
VK
2 22Ti(j-1)R/K
1 -e
K
2 2))/K
Then, we can get that
n kij =
K-1
Σ
Chii.
1
Nu +
(-1)it)
n
K/22
+
E
CRAJSE
K
K
k=0
K/O
1
(-1)
Σ
Ck.
K
K
K/2
When + + R odd, it easy to obtain that
a,is + Ce+ 3.
(ri()+j-2)1/1
{ 2 2 ( (-) (-At)
X,¹ + (-1)
K
K
2x1(++5-211/3
X (-1)+t++ 2 ix
K
0.
which indicates that III "kk 0.
When ++j+ is even,
2
Ji.j
Σ
Ckiig A
K
k20,K/2
Σ
lekii
K/2
8p(1
p)
R
K2
K/2
8p(1 R)
S exp
K2
where the second inequality is by the bound of the absolute values of the eigenvalues, and the last
inequality is by (1 - t) < e Ri for any 0 < t < and [ck.a.j = K for all k, j,
D
w(p '11 p) -)- and (2,..., N 1), it holds that Lemma E.4. For II with 0 < <1, N
N-1
N-1
Σ TT'S
ITT
Σ T (ITT))
+ Cpp
1,1
1,1
where Cp is positive constant only depending p.
Proof of Lemma E.4
Without loss of generality, we assume
S p 1. We observe that
N
N-1 II'I e = e Σ m c1
N-2
=
1=0
30

---

N-1 Σ
N-1
N-2
= =
Σ ( = oj Σ Ir ee 1=0
11
Denote T(N) = - II'. We use induction to prove that
1(N),11 2 ((N)a, + (1 p) for for $
can be easily obtained that for N = 2. ((2)1.1 = r(2),.; + for any * Suppose that the induction hypothesis holds for T(N 1). Then, for [(N), by the definition of F(N), we can get T(N) = II-T(N - 1) + I. Thus, based on the definition of II, we obtain that for + j.
T(N), = p- T(N 1) (1 p) T(N 1)i.t1.
When / - 1ー 2, by induction, we have
T(N)a p p. T(N 1) + p): T(N 1)1st1
S P(T(NN - 1)1,1 - (1 -1p) N- 3-3) + (1 P)( - -1-
= =T(N- - 1)11-11111 (1-p) N-3
$ T(N- 1)1,1))11111111 -1 p) N-2
When =, we have
T(N) p- P(N - 1) (1 T(N 1)ift1
pTT(N - 11),1- (1.1--1111 N1- + (1 p))((N- 1)1,1
T(N - -p(1p)
S T(N- - - -1)1,1(1-11) -) p) N-2
When i =-- 1, we have
+ (1-p) p) T(N - 1)1.t1 T(N), = I(N - +
S - 11),1 + (1- 1)1,1)111111 (1 N33)
-
T(N 1)1, -(1 (1-p)
Obviously, T(N 1)1,1 S I(N)1,1 as I(N) = I(N 1) + N-2 Therefore, we prove that the induction hypothesis holds for T(N), which completes the induction. Thus, for all N 2 2. the diagonal element is the largest entry I(V). We analyze this result from the perspective of random walks. Given that I is the k-step transition matrix in the random walk task, the entry I(N),, NN= IT' can be regarded as the expected number of visits to state j within N- 1 steps, starting from state i. Since the largest entry in I(N) is found at the diagonal, we can conclude that the expected number of visits back state i within N 1 steps, starting from state i, is the largest. we know that for any i,j.
From Lemma F.2 and F.3
2 K
8p(1 - p)R (-8011
+ +exp (- 8p(1 - p) + +1)
(-B0(1
S
exp
K2
K2
Hence, we get that for any i,
8p(1 - P)R
8p(1 - (P) +
[II
1.1 II R++ 11,1
[IT
Jas
[I] R2+ Ja.5 2
exp
+ exp
K2
12
Thus, for constant integer No w(p w0p-1(11 -p) and N > Nu. we get for i +
T(N)1.1 I(N)s 2 (1 p)) No-2 Cpo exp(-No).
where too are constants related to p. In addition, we can get [(N),, - T(N)1,2 S Cpo cxp(-NM for any # 1, + 1.
Next, we focus on II with k 1, 2, 2,... analysis of induction above, we can use induction to prove that for + j.
Denote T2(k) IIIII Tt. Similar to
T2 1)1,1 + (2k)) 2 [2 1)5, + T2(22) - (1 -py
31

---

The proof can be directly extended from the proof for ((N) above. Then, we use this result to prove
(T2()) [2) 2 [2k 1), + [T2(2))]1,1
for any k: 2 2. We also induction to prove this. is easily checked that [T, (T
p² + - p)² 2 3p p) + 3p(1 p) [3)]3] (4)(4)]..1 Suppose that 2
+ (T2 1.1 2 (2k + 1)I,,1 + [T2 2) Then, by T2(2k + T2 (2k +
(22(2k-1) T2 (2h)) (IT )², we can get
(2k + 1)1,1 + (2k + 2)1,1 =p² 1 [T 1) [2
+ 2p(1 - (2 - + (2k))11
+ (12k)
- -) [,
- (2k 1 122kk)
where the inequality is by the result demonstrated before. Therefore, we complete the induction and get that for k 2 1, T2(1),1 + (2()),,1 (2 + 1)1, (2(2) + 2).1 2 To()), - (2(3) - (2(
=(p² + (1- (13p(1111111---- 3p)1-p)).
Since (2(2),, 0, we obtain for h: 2
(2(1)1,1 (2( 2 + (1-)(11)))11 3p(1 p)) == :
We provide an intuitive explanation from the perspective of random walks. II represents first taking a step according to the transition II, followed by steps according to the transition II With increasing k, the distribution of the possible state after le + 1 steps is closer to uniform across all states, regardless of the starting state chosen, which can also be shown by Lemma E.2 and E.3 Thus, I ei will converge the most sparse case with the largest value concentrating on the first entry.
vector corresponding uniform distribution, and IIII represents
Combining all the results obtained above, we can conclude that
N
N-11
Σ II
II
Σ II
(IIT)
e, T(N) Pa(1)e1 - e T(N): (2(k)e)
K
K Σ (2))).
F(N),, + Σ
- I2()),,
2=
x=
2
- -(2(k)))))))))) + min
(1-2))))
max
2gisk
2SiSK
min
(2(())
max (T(N),,)
25iSK
2515K
+ min (T(N),,}
max
1T,
25iSK
25iŠK
(T2(1)1,1 - -I2( - min (T(N), + min (T(N),, 251SK
max (T(N),.) 25iSk
23iSk
z 9p1 (1I-))No-2 No-2 cp0 exxeexxp))) No)) Goo exp(-No)
2 Cp.
where Cp is a positive constant related to p.
The following lemma shows the basic property of the positional embedding. Lemma E.5. Assume that
2i
Min M + 1)
p; sin
sin
sin
M +
M
for E [M]. It holds that
M+1 for i d23
T
P. Pig
Oor in # iz.
32

---

Proof of Lemma E.5
When i1 * i2 and 11 + 22 are even, we have
M
jig
jU2 M +1
P. Pig
Σ
sin
sin
M + 1
j=1
M >
jij
ji2 M +1
sin
sin
M +
j=0
1 M
Σ
exp
im 11 i2
M + +1² 1
exp
im 11 12
M + 1 )
4
1=0
iι 11 + i2
in i1 + i2
]
-exp
M + 1
exp
M + +11
1 exp (iT(1i
12)) - 1
1
exp
in (01 i2) )
+
4
exp iι 12 M++1
-1
4
exp
11 in
M+1
1
1
exp (ir((i) + 02))
-1 1
1
exp
(-in((i + i2]) - 1
4 exp
in futti u++1
4
exp
in M+1
1
0,
where the third equation is by sin(2) exp(iz)==xplt in) and the last inequality is by exp(ink) = 1 for even k. When i # i2 and i + i2 are odd, we have
M
ju
j22 M + 1
P. Pia =
sin
sin
M + +1
j=1
M
ji
ji2 M +
sin
sin
M + 1
j=0
1 M
im 11 122
i 01 122
exp
M + +11) 3
+ exp
M + )
4
j=
iл 81 + i2
- in 11 + +12
]
exp
M + (
exp
M + +1* )
1
exp (it (i1
i2)) 1
exp
in - 12))
+
4
exp in i1-12
4
in i1-da
- 1
exp
1
1
exp (in((i)
+ 12)) - 1
1
exp
im + +2))
4
iæ tit
4
inr +12
exp
M+1
exp
M+1
exp
in d
M+1
exp
-
in M+1
1
1
+
+
2
exp
in 44
M+1
1
exp
in عايي
M+1
0,
where the third equation is by sin = exp(ia) (iciai--xxpp exp(-iz) iz) for odd hs and the last equation is by exp(x)) exp(-x))-1
the fifth inequality is by exp(ink) = -1
2i
1
-1. When é 12, we have
M
jigh
P; Pig
sin
sin
M +
M +1
33

---

M
ji
jigh
sin
sin
M +
M +
==
Σ
in 11
in 11 82
exp
+ 1
+ +exp
M +
im 11 + 12
M + -
11 + 82 M +
]
eep
exp
M+1
1
in 11 22
- 11 + ( M + 15)
exp
+ exp
2
4
M +
M 1
exp
(iT (ii
i2))
exp
(- + i2))
2
exp
in 1et
M++1
-1
exp
- -
M + 1
2
where the third equation by sin(x) exp(ix) exp(ix)-exp(innn and the last inequality is by exp(irk) = for even k.
21
F
Additional Experiments
F1 Additional Experiments on Random/Deterministe Walks
In this subsection, we provide additional experiments on synthetic data with (K, N)) (20, 101). We consider the transformer model introduced in Section with the length of the position embedding M 1000. train the model, we utilize gradient descent starting with zero initialization, where the learning rate 7 1 and the constant in the log-loss is set descent algorithm for = 50 training epochs. Figure and Figure iillustrate the experiments for p 0.5 and p 1 respectively. These experimental results match Theorem and Theorem 4.1
0.1. And, we run the gradient
which also strongly supports our theoretical results.
(a) Accuracy
(b) Visualization of V (T) (c) The average attention
(d) Part of the attention
Figure 8: The results of the experiment for = 0.5 with NK = (20, 101): (a) is the test accuracy;
(b) is the visualization of (T); (c) and (d) present the average attention of the test data with x-axis representing the position the token and y-axis representing the attention score.
F.2 Additional Experiments on the Question Answering Tasks in Section 6
In this section, we conduct some additional experiments for Task 1 and Task 2 discussed in Section O We conduct some additional experiments extending the one-layer transformer model more complicated model by adding fully connected layer with ReLU activation to the transformer model. The new model has the form
fo(X) A ReLU (V VV X WTWNN))
(F.1)
where A RKxmm V E m W E ((K+M) are the trainable parameter matrices, and is the number of neurons in the fully connected layer. For Task I and Task 2, the length of
34

---

Takes Sootaor
0.002
0.015
015
0.7
650
622
600
(a) Accuracy
(b) Visualization of V
(c) Average attention of test data
Figure 9: The results of the experiment for p with K 20, 101. (a) is the prediction accuracy with r-axis representing the iteration and y-axis representing the accuracy. (b) is the visualization of V. (c) is the average attention of the test data with z-axis representing the position of the token and y-axis representing the attention score.
the vocabulary K and the length of each input sequence N are set as (K, N) (19, 17), (19, 19) respectively. In addition, we set the positional embedding M 1000 and the number of neurons m = 19. To train the model, we consider the Gaussian random initialization A A ( w N ( o²) with log-loss set as c = 0.1. Both experiments are conducted on 1024 training data and 1024 test data. Here, most of the settings remain the same in the previous experiments in Section 6
0.01, and use gradient descent with learning rate 0.1. The constant in the
/ M
Figure 10: The results (b) KL-Divergence
(a) Accuracy
(c) Accuracy
(d) KL-Divergence
experiment conducted using a more complicated transformer for Task 1 and Task 2: (a) and (b) correspond the experiment for Task 1; (c) and (d) correspond to the experiment for Task 2.
Figure TO shows the experiment results using the more complicated transformer in M learn Task 1 and Task 2. In Figure TO(a) and Figure T0(c) we present the test accuracy achieved by the transformer model learning Task and Task 2 respectively. In Figure [0bb and Figure [Ord) we first normalize the output of the trained transformer model to get a K <'-dimensional vector, representing the prediction distribution of K words. Then, we report the KL-divergence between this prediction distribution and the true distribution of y| Ø2, ..... 2N The experiment results show clear difference between the performances of the transformer model in the two tasks. In Task 1, the trained transformer model can successfully approach the optimal accuracy (100%) within 100 iterations. However, in Task 2, the test accuracy always remains around 50%, which is the accuracy a random guess.
Despite using a more complicated transformer model with an additional feedforward layer of non- linearities compared to the one considered in theoretical analysis and previous experiments, the experimental results are still similar to those reported in Section 6 These results demonstrate that more complex transformer models may still struggle with the relatively 'simple' Task 2 but excel at the relatively 'difficult' Task This indicates that our findings can be applied to cases involving additional nonlincarities, implying their applicability more complex and general conditions.
F33 Visualizations of the value matrix and Softmax scores corresponding to Figure
In this section, we present visualizations of the value matrix and softmax scores corresponding to
Figure 6
35

---

Toxen Seletion
10
02-
00
Piction
(a) Value matrix
(b) Softmax scores
Figure 11: Visualizations of the trained value matrix and the softmax scores corresponding to Figure 6(a) and Figure 6(b)
02525
0.200
03475
0.358
02400
(a) 100 training steps
(b) 400 training steps
0.7
0.40
0.16
036
033
636
(c) 700 training steps
(d) 1000 training steps
Figure 12: Visualizations of the value matrix corresponding to Figure 6(c) and Figure 6(d) at 100, 400, 700, and 1000 training steps, respectively.
Figure visualizes the trained value matrix and the softmax scores attached to all tokens, corre- sponding Figure 6(a) and Figure 6(b) shows that V recover the transition matrix II and the softmax score attached to the direct parent token is the largest and close to These findings demonstrate that our positive results for learning random walks in Theorem 3.1 are fairly robust to random initialization.
Figure 12 and Figure 13 display the value matrix and the weighted average token X S(X W x corresponding to Figure b(c) and Figure ( at 100, 400, 700, and 1000 training steps, respectively. In Figure 13 the weighted average token embedding is calculated based one sequence X as an example, while the variance each dimension of the weighted average token is calculated across 1000 sequences. We can observe that when considering the case for with random initialization.
36

---

100 Traininng Step
400 Trairing Stop
Warghted Avg Toten
Weighted Aug Bean
0.223
0.223
0325
04
02
0.332
0.000
Inde
Indox
0.534
-0532
00321
0.20
0.030
0.35
00032
025
028
0.32
00331
0.000
0:5-
0.504
0.25
0.0
Isdec
Indea
1sdec
Iices
Figure 13: Visualizations of the weighted average token X S(X Wix) corresponding to Fig- ure 66c) and Figure b(d) at 100, 400, 700, and 1000 training steps, respectively.
the value matrix and the weighted average token embedding remain approximately proportional to the all-one matrix and the all-one vector within 700 iterations, aligning with the results stated in Theorem 4.1. Additionally, the variances of all dimensions are close to O, which shows that this result general for different sequences. However, after 1000 iterations, V is no longer proportional to an all-one matrix, and the softmax score for one arbitrary token becomes much larger than the others. As discussed Section 4 the failure case for p 1 is indeed due unbreakable symmetry caused by zero initialization. The random initialization may break this symmetry, enabling the transformer successfully learn the optimal predictor.
37

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: This paper mainly focus on theory, and our main results are stated in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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 long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have discussed limitations on the studies of the impact of random initial- ization and the relatively simple setting of one-layer transformers.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiitiinn 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 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-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't '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 complete (and correct) proof?
38

---

Answer: [Yes]
Justification: We have clearly stated all theoretical assumptions in the problem setting section and in the formal theorems.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced 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 encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results 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)?
Question:
Answer: [Yes]
Justification: Experiments in this paper are mainly to back up the theoretical results and are not too complicated. We have provided detailed descriptions the experimental setups.
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 to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility 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, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS 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 a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
39

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification: Experiments in this paper synthetic data and relatively simple. The purpose is mainly to back up our theory.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://niis.cc.
aublic/guudees/Coeettbbmsss
wwtpssiiinooooyyyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https wtteee////oSSSemmsss.eeeee..
/ //nips. wttps/cccppclcccuui.eeeeeecccaea
www.ssssiicggcooo 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 proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experiments are relatively simple and the purpose is mainly to back up our theory. We have clearly explained all the experiment setups.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core 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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: The experiments in this paper are to verify our theory, and there is no need to analyze statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
40

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification: The experiments in this paper are very simple and are mainly to verify our theory. There is no need specify compute resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs 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 NeurIPS Code of Ethics http:///wwrppp...copubiii///eaaa
Answer:
[Yes]
Justification:
We confirm that this research conforms with the NeurIPS Code of Ethics.
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: This paper mainly focuses on theoretical analyses on simple one-layer trans- former models. We do not see any potential negative social impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
41

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess 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 tied to particular applications, let alone deployments. However, if there is 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, 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 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 are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: This paper mainly focuses on theoretical analyses on simple one-layer trans- former models. We do not any potential risk for misuse.
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 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 we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [NA]
Justification: This paper mainly focuses theoretical analyses simple one-layer trans- former models. We have provided appropriate citations and there is no other relevant assets credit.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
42

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.coo//aaaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
43

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [No]
Justification: We have only used LLMSs for editing purposes.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy for what should or should not be described.
44


Reviews:
[{"review_id": "k6rUkkSy0g", "reviewer": "Reviewer_ki3b", "summary": "This paper studies the dynamics of gradient descent on a 1 layer Transformer trained on a specific Markov chain arising from a biased random walk. It shows that gradient descent successfully solves this task in a constant number of iterations.", "questions": "- Would it be possible to extend the proof to the empirical loss rather than the population loss? It would be interesting to compare the sample complexity of the simple autoregressive model mentioned above with the depth 1 transformer considered in the paper. In particular, unlike the simple autoregressive model, the sample complexity of the depth 1 transformer would probably scale with the sequence length as it is forced to identify the parent tokens to emulate this autoregressive process.\n- Why is it important for the analysis to specialize to this specific class of Markov chains (random walks on a circle) rather than allowing for arbitrary Markov chains (under some non-degeneracy assumption like $\\pi(s|s') > 0$ for all $s,s'$)?", "limitations": "yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The analysis successfully handles Softmax attention rather than studying the simplified linear attention.\n- The paper is generally well written and both the proof sketch and the discussion for the optimization failure when $p \\in \\{0,1\\}$ are very clear.\n\nWeaknesses:\n- The setting is very simple and can be solved by much simpler models. For example, a simple autoregressive model (a \"depth 0\" Transformer) $f_\\theta(\\boldsymbol{X}) = \\boldsymbol{V} \\boldsymbol{x}_N$ would trivially solve this task with $\\boldsymbol{V} \\to \\boldsymbol{\\Pi}$ and the corresponding optimization problem would be convex. This could also be achieved by adding a residual connection to the transformer. The difficulty, both in the analysis and in the optimization dynamics, therefore arise from the fact that the architecture is written in such a way that the attention layer is forced to simulate an autoregressive process.\n- A minor comment is that while the use of log-loss is now fairly common in theoretical analyses of Transformers, it should still be noted that this is non-standard as the final softmax after the value matrix is omitted. The paper therefore relies on the fact that $\\mathcal{S}(\\ldots) \\ge 0$, $X \\ge 0$ (because of the one-hot positional embeddings), and $V \\ge 0$ (from the dynamics) to ensure that $f \\ge 0$ so that the log-loss doesn't blow up.", "was_revised": true, "final_justification": "After the discussion, my main concern remains that this paper takes a very simple task (learning a Markov chain) which a standard Transformer can very easily learn, and then uses an unnatural input encoding and non-standard architecture which forces the model to learn a more complicated solution using the self-attention block.", "rebuttal": ">**Q1**: The setting is very simple and can be solved by much simpler models. For example, a simple autoregressive model (a \"depth 0\" Transformer) $f(X)=Vx_N$ would trivially solve this task with $V \\to \\Pi$ and the corresponding optimization problem would be convex.\n\n**A1**: We agree that a simple autoregressive model $f(X)=Vx_N$ can also solve the classic random walk task. However, please note that our goal is not to identify the optimal model for solving random walk tasks but to understand and analyze the capability and interpretability of transformers when learning such classic statistical tasks from scratch. \n\nImportantly, autoregressive models may have an advantage in learning random walks because they assume Markovian property by definition. In comparison, the one-layer transformer model we consider follows the standard definition and is not designed with any prior knowledge of the Markovian property of the data. Therefore, investigating whether transformers can also learn random walks is more challenging. Moreover, by showing that self-attention can successfully be trained to learn the Markovian property by attending to the direct parent token, our results demonstrate how a simple transformer model can capture underlying data dependency structures and interpretably learn classic statistical tasks.\n\nWe would like to mention that there exists a line of work studying the capabilities of transformers in learning Markov chains [1,2]. The tasks considered in these studies could also be addressed by some similar autoregressive models. [1] characterizes the loss landscape of a one-layer transformer and identifies global minima and bad local minima based on Markovian data. [2] shows that transformers with tokenization can efficiently learn k-th order Markov chains, and demonstrates that the failure cases observed in [1] can be overcome with appropriate tokenization. However, these studies only focus on the Markov chains with only two states. In contrast, our framework allows a large number of states, suggesting that our setting is already more general and complex.\n\n>**Q2**: A minor comment is that while the use of log-loss is now fairly common in theoretical analyses of Transformers, it should still be noted that this is non-standard as the final softmax after the value matrix is omitted. The paper therefore relies on the fact that $S(...) \\geq 0, X \\geq 0$ (because of the one-hot positional embeddings), and $V \\geq 0$ (from the dynamics) to ensure that $f \\geq 0$ so that the log-loss doesn't blow up.\n\n**A2**: Thank you for your comment. As you mentioned, log-loss is fairly common in theoretical analyses [1,3,4], and the setting with log-loss can already provide valuable insights into the underlying mechanisms and behaviors of transformers. In the revision, we will add discussions on the reasons that the loss does not blow up.\n\n>**Q3**: Would it be possible to extend the proof to the empirical loss rather than the population loss? It would be interesting to compare the sample complexity of the simple autoregressive model mentioned above with the depth 1 transformer considered in the paper.\n\n**A3**: We appreciate your suggestion regarding the comparison of sample complexity between transformers and the simple autoregressive model. Our current analysis focuses on the population loss, as this setting can help provide clean theoretical insights and guarantees about the transformer’s capabilities and interpretability. Training on the empirical loss may require significantly more complicated perturbation analyses to handle the gaps between population and empirical gradients. Please note that most of the existing theoretical studies on the training dynamics of transformers are based on the population loss [1,3,5,6].\n\nIt is also likely that more carefully designed autoregressive models may have better sample complexity. However, as discussed in **A1**, the transformer model is designed with no prior knowledge of the Markov property, and therefore, such a comparison may not be fair. Our goal is not to optimally learn random works with the best-designed model, but to understand and demonstrate the capabilities of transformers.\n\n>**Q4**: Why is it important for the analysis to specialize to this specific class of Markov chains (random walks on a circle) rather than allowing for arbitrary Markov chains (under some non-degeneracy assumption like $\\pi(s|s’)>0$ for all $s,s’$?\n\n**A4**: Our result demonstrates different performances of transformers when learning specific random or deterministic walks. Existing work [1] also points out the existence of success and failure cases in learning Markov chains with two states. These results together indicate that it may be very challenging to develop a general and unified theory for arbitrary transition matrices. We believe that such challenges in developing a more unified theory also, to some extent, demonstrate the difficulty and importance of obtaining our results.\n\nIn fact, our setting is already more general than the previous work [1]. As mentioned in **A1**, [1] focuses on the first-order Markovian chains with vocabularies of size two, which can also be interpreted as a random walk over a space of only two states. In comparison, we consider random walk tasks on circles that allow a large number of states, indicating that our theory represents a step further towards a unified theory for more general cases.\n\n[1] Makkuva et al. Attention with Markov: A curious case of single-layer transformers. ICLR 2024.\n\n[2] Rajaraman et al. An analysis of tokenization: Transformers under markov data. NeurIPS 2024.\n\n[3] Ildiz et al. From self-attention to Markov models: Unveiling the dynamics of generative transformers. ICML 2024.\n\n[4] Li et al. Mechanics of next token prediction with self-attention. International Conference on Artificial Intelligence and Statistics. AISTATS, 2024.\n\n[5] Nichani et al. How transformers learn causal structure with gradient descent. ICML 2024.\n\n[6] Edelman, Ezra, et al. The evolution of statistical induction heads: In-context learning Markov chains. NeurIPS 2024."}, {"review_id": "SFNmajmAhm", "reviewer": "Reviewer_5AGb", "summary": "This paper theoretically characterizes the learning dynamics of a one-layer transformer on the task of predicting a circular random walk.\nThe authors prove that for transition probabilities p in (0,1), gradient descent from zero initialization enables the model to achieve optimal prediction accuracy. \nThis is accomplished via an interpretable mechanism where the self-attention layer learns to act as a token selector, while the value matrix converges to the transition matrix. \nIn contrast, for deterministic walks (p=0,1), the paper demonstrates a failure case, caused by initial symmetries, where the model's performance is comparable to random guess.", "questions": "- The theory relies heavily on zero initialization. The experiments with random initialization are interesting but could be explored further. Could the authors comment on whether the interpretability of the V and S matrices is preserved (numerically) with larger, non-zero random initializations? What are the key obstacles to extending the theoretical analysis beyond the zero-init case?\n- How would the failure mechanism for deterministic walks (p=0,1) be affected by the presence of multiple transformer layers? Would deeper models be able to break the symmetry, or would the 'uninformative' signal propagate and cause a similar failure?\n- The appendix includes an experiment with an additional fully connected layer (F.2), which is a valuable addition. It would strengthen the paper to reference and briefly discuss this result in the main text. On a related note, what would be the expected effect of adding bias terms to the model's layers?\n- Do similar conclusions about interpretability hold for the case of Markov chains with more complex transition matrices?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n- The paper's main strength is its well-chosen toy problem that allows for a clear and mathematically tractable analysis.\n- The work presents both a success case and a failure case within the same framework. The authors identify precisely which conditions lead to successful learning and which lead to optimization failure, and why.\n- The findings offer a clear, interpretable story for how the parameters of a transformer can learn that specific simple task and the experiments are convincing.\n\n**Weaknesses:**\n- The primary weakness is that the paper's analysis relies on a heavily simplified model and task. The theoretical analysis is limited to a single attention layer and omits standard transformer components like MLPs, residual connections, and layer normalization. \n- Because the proofs rely on explicit calculations made possible by these simplifications, it is  unclear whether the core insights or mathematical techniques would generalize to more realistic transformer architectures. This raises questions about the potential impact of these findings for understanding the transformers used in practice.", "was_revised": true, "final_justification": "The paper presents a clear analysis of a simplified model and task. While I still believe that the architecture and setting are too idealized to have immediate practical relevance, and that the extendibility of the results to more realistic scenarios may be limited, the theoretical analysis is well written and instructive. The authors’ rebuttal addressed most of my concerns regarding the modeling assumptions. For these reasons, I consider this a borderline accept.", "rebuttal": ">**Q1**: The primary weakness is that the paper's analysis relies on a heavily simplified model and task. The theoretical analysis is limited to a single attention layer and omits standard transformer components like MLPs, residual connections, and layer normalization.\n\n**A1**: Thank you for your suggestion. We agree that considering more complex models is important and aligns better with the practical setting. However, providing precise theoretical analysis for such more complicated architectures can be extremely challenging. Most existing theoretical studies on the training dynamics of transformers focus on simple one-layer transformer models [1,2,3,4,5]. In fact, to our knowledge, the setting considered in our paper is already more aligned with the practical transformer architecture compared to many of these recent theoretical studies. For example, [1] and [2] analyze transformers with linear attention instead of softmax attention. And, in [3], [4], and [5], the value matrix $V$ is held constant during training, with only parameters inside the softmax matrix being trained. In comparison, our work constructs a transformer architecture with nonlinear softmax attention and analyzes the training of the matrices $V$ and $W$ simultaneously, which is a step toward a more practical study.\n\n>**Q2**: Because the proofs rely on explicit calculations made possible by these simplifications, it is unclear whether the core insights or mathematical techniques would generalize to more realistic transformer architectures. This raises questions about the potential impact of these findings for understanding the transformers used in practice.\n\n**A2**: As mentioned in **A1**, our work has already analyzed a setting that is closer to practice scenarios than many existing theoretical studies. In addition, while the precise training procedure may differ, we believe that our study can still provide insights into more practical settings. Specifically, our study proves that a one-layer transformer can learn random walks by ‘selecting the parent token’ with softmax attention. The insight that softmax attention can serve as token selectors aligns with many existing theoretical and empirical studies [6,7,8].\n\nAdditionally, we believe the fact that our theory can successfully motivate us to predict the performance of transformers in two question-answering tasks demonstrates the practical value of our theory. We believe that conducting rigorous theoretical analysis in a clean setting and providing insights applicable to other settings is precisely what a theory paper should aim to do, and our paper accomplishes this.\n\n>**Q3**: The theory relies heavily on zero initialization. The experiments with random initialization are interesting but could be explored further. Could the authors comment on whether the interpretability of the V and S matrices is preserved (numerically) with larger, non-zero random initializations? What are the key obstacles to extending the theoretical analysis beyond the zero-init case?\n\n**A3**: Thanks for your question. In our experiments with random initialization, the visualizations of V and S corresponding to Figure 6(a) and 6(b) show that the trained value matrix $V$ can recover the transition matrix $\\Pi^{\\top}$ and the softmax score attached to the direct parent token is the largest and close to 1, demonstrating that our positive results for learning random walks in Theorem 3.1 are fairly robust to random initialization.\n\nIn the case for p=0,1, with random initialization, the visualizations of V and S corresponding to Figure 6(c) and 6(d) show that within 600 iterations, the value matrix is approximately proportional to an all-one matrix, and the softmax attention fails to select any token, aligning with the negative results stated in Theorem 4.1. However, after 1000 iterations, V is no longer proportional to an all-one matrix, and the softmax score for one arbitrary token becomes much larger than the others. As we have discussed in the paper, the failure case for p=0,1 is indeed due to an unbreakable symmetry caused by zero initialization. The random initialization may break this symmetry, enabling the transformer to successfully learn the optimal predictor.\n\nWe will add these visualizations in the revision.\n\nMoreover, the use of zero initialization in our theoretical analysis is primarily to simplify the dynamics of the model, which allows us to derive rigorous insights and obtain clear theoretical results. Our proof is based on the analyses of the softmax scores on each token during training. The setting with zero initialization ensures a fair start for all tokens, and we can prove that after training, the largest softmax score will be assigned to the direct parent token. When considering non-zero initialization, the initial scores assigned to all the tokens differ, leading to an unfair starting condition for training. This difference can largely influence and complicate the optimization analysis. If the direct parent token receives a relatively small initial score, the larger initial scores assigned to other tokens may cause the value matrix $V$ to learn a k-step transition matrix (other than the one-step transition), which can not achieve the optimal predictor. Additionally, in failure cases, the symmetry preserved under zero initialization is broken with non-zero initialization. This causes the largest softmax score to be assigned to an arbitrary token, which can enable the model to make an optimal prediction. Therefore, zero initialization serves as a useful idealization to facilitate theoretical understanding, while extending the results to non-zero initializations remains an interesting direction for future work. Please note that zero initialization or specific initialization has been considered in most theoretical studies [1,2,3,4].\n\n>**Q4**: How would the failure mechanism for deterministic walks (p=0,1) be affected by the presence of multiple transformer layers? Would deeper models be able to break the symmetry, or would the 'uninformative' signal propagate and cause a similar failure?\n\n**A4**: We demonstrate that one-layer transformers fail to learn deterministic walks due to the fact that the average of tokens does not provide useful information for prediction. We believe this insight can also be applied to learning the same task using multi-layer transformers. When the token average is not informative, multi-layer transformers may struggle since it always takes an average over all tokens in each layer. As a result, the output of multiple layers may lack the necessary information for prediction.\n\n>**Q5**: The appendix includes an experiment with an additional fully connected layer (F.2), which is a valuable addition. It would strengthen the paper to reference and briefly discuss this result in the main text. On a related note, what would be the expected effect of adding bias terms to the model's layers?\n\n**A5**: Thank you for your suggestion. We will include discussions about this experiment in our revised version. Regarding the bias terms, we conduct the experiment with the following model: $f(X) = A\\cdot\\text{ReLU}((VX+b_V)\\text{Softmax}(\\tilde{X}^{\\top} W \\tilde{x}_{N})) + b_A$. The learning tasks and the settings are the same as the experiments in Section F.2. The results are similar to Figure 10 in the paper. In task 1, the transformer achieves the optimal accuracy (100%) after 20 iterations. In task 2, the test accuracy remains around 50%, which is just the accuracy of a random guess. This indicates that adding bias terms may not significantly influence the results.\n\n>**Q6**: Do similar conclusions about interpretability hold for the case of Markov chains with more complex transition matrices?\n\n**A6**: Our results demonstrate different performances of transformers when learning specific random or deterministic walks. Existing work [5] also points out the existence of success and failure cases in learning Markov chains with two states. These results together indicate that it may be very challenging to develop a general and unified theory for more complex transition matrices. We believe that such challenges in developing a more unified theory also, to some extent, demonstrate the difficulty and importance of obtaining our results.\n\nIn fact, our setting is already more general than the previous work [5]. [5] characterizes the loss landscape of a one-layer transformer and identifies global minima and bad local minima based on Markovian data with vocabularies of size two, which can also be interpreted as a random walk over a space of only two states. In contrast, our framework accommodates a large number of states, suggesting that our theory represents a step further towards a unified theory for more general cases.\n\n[1] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.\n\n[2] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv preprint arXiv:2411.05189 (2024).\n\n[3] Li et al. How do transformers learn topic structure: Towards a mechanistic understanding. ICML 2023.\n\n[4] Li et al. One-Layer Transformer Provably Learns One-Nearest Neighbor In Context. NeurIPS 2024.\n\n[5] Tarzanagh et al. Transformers as support vector machines. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning.\n\n[6] Zhang et al. Anchor function: a type of benchmark functions for studying language models. arXiv preprint arXiv:2401.08309 (2024).\n\n[7] Song et al. Out-of-distribution generalization via composition: a lens through induction heads in transformers. Proceedings of the National Academy of Sciences 122.6 (2025).\n\n[8] D'Angelo et al. Selective induction heads: How transformers select causal structures in context. ICLR 2025."}, {"review_id": "AH3pRm0mcA", "reviewer": "Reviewer_Z1xn", "summary": "This paper provides a theoretical and empirical analysis of a one-layer transformer's learning random walks on a circle. The authors prove that for non-deterministic walks (where the transition probability p is between 0 and 1), a transformer trained with gradient descent from zero initialization can achieve optimal prediction accuracy. Their analysis reveals an interpretable mechanism: the self-attention layer learns to function as a token selector, focusing almost exclusively on the direct parent state, while the value matrix learns to approximate the transition matrix. Conversely, the authors demonstrate that for deterministic walks (p=0 or p=1), the same training setup fails, resulting in performance no better than a random guess. They attribute this to an optimization failure caused by unbreakable symmetries under zero initialization. The findings are supported by experiments on the random walk task and extended to simple question-answering tasks to show the broader relevance of the identified failure mode.", "questions": "1. **On the Role of Random Initialization:** The analysis of the failure case for `p=0,1` is nice, but relies heavily on zero initialization. Your experiments with random initialization (Figure 6) are interesting; they show that learning remains much harder for `p=1` than for `p=0.5`. Could you elaborate on the theoretical challenges of analyzing dynamics under random initialization? Does it simply break the symmetry, allowing gradients to find useful directions, or is there a more complex dynamic? \n    \n2. **Generalizing Beyond Circle Graphs:** The random walk on a circle possesses a convenient Toeplitz/circulant structure. How do you expect results to change for random walks on general (undirected) graphs? Specifically, would the core mechanism of \"attention selects parent, value matrix learns transition\" still hold? I suspect the attention component might remain robust, but the clean convergence of `V(t)` to `Π^T` might be complicated by the absence of structural regularity.\n    \n3. **Absence of Softmax Scaling:** The standard transformer architecture includes a `1/sqrt(d_k)` scaling factor inside the softmax, which your model in Eq. (2.1) omits. Was this omission for analytical convenience? How would its inclusion affect the dynamics? This scaling factor could influence the softmax temperature, potentially impacting the speed and sharpness with which attention concentrates on the parent token, particularly during early training stages. \n    \n4. **Circle Structure Justification:** Could you elaborate on the choice to constrain the random walk to a circle? Specifically, could you clarify in the main text how the circle structure is fundamental to your analysis and why this regular structure is necessary for proof tractability? This would help readers better understand the scope and assumptions of your theoretical contribution early on. \n    \n5. **Residual Connection Impact:** A residual connection would pass token embeddings directly to the output, effectively 'giving away' the parent token x_{N-1} to the prediction head, making parent token selection trivial. How would your analysis change with the inclusion of a residual stream? Would the failure case for deterministic walks (p=0,1) still occur, or would direct access to the parent token resolve the 'uninformative average' issue even under zero initialization?\n    \n6. **Alternative Token Dependencies:** To address the triviality introduced by skip connections, it would be interesting to see generalization to Markov chains that still depend on a single token but at different positions rather than the previous one (e.g., x_{N-2}), as explored in [2,5].\n    \n7. **Higher-Order Markov Chains:** Have you considered extensions to higher-order Markov chains where the attention mechanism must learn to attend to multiple tokens simultaneously? Even if multiple token attention is possible using multiple heads (see [1] above), do you have insights into how a mechanism similar to your Value matrix implementation would function in such cases?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\n- **Presentation:** The paper is well written, I enjoyed reading it, and I learned something from it. \n\n- **Theoretical Rigor:** The paper's primary strength is its rigorous theoretical analysis. The authors prove that a one-layer transformer trained with gradient descent learns the optimal predictor (for p∈(0,1)) and its convergence rate. They also show how in the deterministic cases (p=0,1) learning fails. This complete success/failure analysis is insightful and provides a clear picture of the model's capabilities and its training under the specified conditions.\n\n- **Interpretability of the Learned Mechanism:** The paper provides a clear and interpretable explanation of how the transformer solves the task. While the finding that attention mechanisms learn to select parent tokens and value matrices encode transition probabilities is not entirely novel, seeing how this structured solution emerges naturally through gradient descent training also in this setting offers valuable insights.\n    \n- **Empirical Validation:** The experiments, while simple, are well-designed to directly validate the theoretical claims even beyond Markov Chains. \n    \n\t\n**Weaknesses:**\n\n- **Simplified Model:** The analysis is restricted to a one-layer transformer without an MLP block, layer normalization, or residual connections. While this is a standard and necessary simplification for theoretical tractability, it represents a significant departure from practical, multi-layer transformer architectures. Although the authors are transparent about this limitation, it naturally constrains the generalization of their results.\n    \n- **Interpretability Contribution:** The novelty of the interpretable construction remains unclear. Previous works relating to n-gram learning and Markov chains have already established how attention is responsible for learning to attend to the correct token (implementing a token selector) while the transition matrix is stored in linear or MLP layers (see, for example, [1,2,3,4]).\n    \n- **Dependence on Zero Initialization:** The sharp failure case for deterministic walks critically depends on zero initialization. As the authors' own experiments in Figure 6 demonstrate, this failure can be partially mitigated through random initialization. While the zero-initialization analysis is theoretically clean and insightful, it raises questions about the practical relevance of this specific failure mode, given that random initialization is standard practice. Notably, the theory does not cover the random initialization case.\n    \n- **Task Specificity:** The random walk on a circle exhibits a highly regular structure, resulting in a Toeplitz transition matrix. This structure is heavily exploited throughout the proofs (e.g., Lemma C.2). It remains unclear how the analysis would extend to random walks on more general, less symmetric graphs, where the properties of the transition matrix might differ substantially.\n    \n- **Circle Constraint Motivation:**  The motivation for constraining the random walk to a circle is not immediately clear from the main text and could be made more explicit. This choice is crucial, as it induces a 'Toeplitz property' in the transition matrix. \n\n[1] Anej Svete and Ryan Cotterell. 2024. Transformers Can Represent n-gram Language Models. In _Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)_, pages 6845–6881, Mexico City, Mexico. Association for Computational Linguistics. \n\n[2] D'Angelo, F., Croce, F., & Flammarion, N. (2025). Selective induction heads: How transformers select causal structures in context. In _The Thirteenth International Conference on Learning Representations_. \n\n[3] Nichani, Eshaan, Alex Damian, and Jason D. Lee. \"How transformers learn causal structure with gradient descent.\" arXiv preprint arXiv:2402.14735 (2024). \n\n[4] Edelman, Ezra, et al. \"The evolution of statistical induction heads: In-context learning markov chains.\" Advances in neural information processing systems 37 (2024): 64273-64311. \n\n[5] Makkuva, Ashok Vardhan, et al. \"Attention with markov: A framework for principled analysis of transformers via markov chains.\" arXiv preprint arXiv:2402.04161 (2024).", "was_revised": true, "final_justification": "I recommend acceptance of this work. This paper presents a theoretically rigorous and insightful analysis of a one-layer transformer learning a specific, structured task. The authors' rebuttal was thorough and successfully clarified my questions. The paper is interesting and technically solid, but its core findings rely on a simplified and highly structured setting, which limits their generality. My final score reflects this balance between the paper's technical strength and its structured setup.\n\n**Resolved Issues:**\n\nThe authors' rebuttal addressed all my specific questions. \n*   **Impact of Residual Connections:** The authors corrected my misunderstanding, explaining why, in their specific prediction setup, a residual connection would not alter the model's output.\n*   **Justification for Simplifications:** The authors provided justifications for their modeling choices, i.e., the use of a circle, zero initialization, and the omission of softmax scaling.\n*   **Contribution:** They effectively positioned their work relative to prior related papers, highlighting that their contribution lies in providing rigorous guarantees for the training dynamics, not just the expressive power of the learned mechanism.\n\n**Unresolved Issues:**\n\nWhile my questions were answered, the core limitations that I raised in my initial review remain. These are not unresolved issues in the sense of a flawed rebuttal, but rather inherent constraints on the paper that the authors acknowledged:\n\n*   **Limited Generality:** The analysis is fundamentally tied to the highly symmetric structure of a random walk on a circle and the use of a simplified one-layer transformer. How these results generalize to less structured tasks or more complex architectures remains a significant open question.\n*   **Dependence on Zero Initialization:** The studied failure case for walks is an artifact of zero initialization. As the authors' own experiments show, this specific failure mode is less pronounced with random initialization, which is standard practice. The theory does not cover this more practical case.", "rebuttal": ">**Q1**: Simplified Model \n\n**A1**: We agree that considering more complex models is important and aligns better with the practical setting. However, providing precise theoretical analysis for such more complicated architectures can be extremely challenging. Most existing theoretical studies on the training dynamics of transformers focus on simple one-layer transformer models [1,2,3,4,5]. In fact, to our knowledge, the setting considered in our paper is already more aligned with the practical transformer architecture compared to many of these recent theoretical studies. For example, [1,2] analyze transformers with linear attention instead of softmax attention. And, in [3,4,5], the value matrix is held constant during training, with only parameters inside the softmax matrix being trained. In comparison, our work constructs a transformer architecture with nonlinear softmax attention and analyzes the training of V and W simultaneously, which is a step toward a more practical study.\n>**Q2**: Interpretability Contribution\n\n**A2**: We acknowledge that previous works have explored how attention mechanisms facilitate token selection. We believe that one of our contributions is to provide further theoretical evidence on the successful cases where transformers can be trained to perform token selection. Below, we give detailed comparisons between our work and existing results in [6,7,8,9]. [5] and [6] show that transformers with specific constructions can represent any n-gram language model, but they primarily focus on the expressive power of transformers and do not delve into the training dynamics. [7,8] investigate transformers' in-context learning capabilities with Markovian data, both considering two-layer transformers. [7] sets a specific initialization designed to assign distinct roles to each layer, thereby facilitating the analysis of how different layers contribute during training. Meanwhile, [8] simplifies the model to a single linear attention layer making the analysis of training dynamics more tractable. In contrast, our work focuses on directly learning Markov chains through a single-layer transformer model. We provide theoretical guarantees that standard gradient descent, without any specialized initialization or assumptions, naturally converges to an interpretable solution capturing the inherent sequential structure (Theorem 3.1).\n>**Q3**: Zero Initialization\n\n**A3**: The use of zero initialization in our theoretical analysis is primarily to simplify the dynamics of the model, which allows us to derive rigorous insights and obtain clear theoretical results. Our proof is based on the analyses of the softmax scores on each token during training. The setting with zero initialization ensures a fair start for all tokens, and we can prove that after training, the largest softmax score will be assigned to the direct parent token. When considering non-zero initialization, the initial scores assigned to all the tokens differ, leading to an unfair starting condition for training. This difference can largely influence and complicate the optimization analysis. If the direct parent token receives a relatively small initial score, the larger initial scores assigned to other tokens may cause the value matrix V to learn a k-step transition matrix (other than the one-step transition), which can not achieve the optimal predictor. Additionally, in failure cases, the symmetry preserved under zero initialization is broken with non-zero initialization. This causes the largest softmax score to be assigned to an arbitrary token, which can enable the model to make an optimal prediction. Therefore, zero initialization serves as a useful idealization to facilitate theoretical understanding, while extending the results to non-zero initializations remains an interesting direction for future work. Please note that zero initialization or specific initialization has been considered in most theoretical studies [1,2,3,4].\n>**Q4**: Task Specificity\n\n**A4**: Our result demonstrates different performances of transformers when learning specific random or deterministic walks. Existing work [10] also points out the existence of success and failure cases in learning Markov chains with two states. These results together indicate that it may be very challenging to develop a general and unified theory for arbitrary transition matrices. We believe that such challenges in developing a more unified theory also, to some extent, demonstrate the difficulty and importance of obtaining our results. In fact, our setting is already more general than the previous work [10]. [10] characterizes the loss landscape of a one-layer transformer and identifies global minima and bad local minima based on Markovian data with vocabularies of size two, which can also be interpreted as a random walk over a space of only two states. In contrast, our framework accommodates a large number of states, suggesting that our theory represents a step further towards a unified theory for more general cases.\n>**Q5**: Circle Constraint Motivation\n\n**A5**: Thank you for your suggestion. Random walks in circles enjoy better symmetry properties that greatly facilitate the calculations of gradients, which are fundamental to the training analysis. For instance, due to the symmetry, the stationary distribution is uniform across all nodes, which significantly simplifies the calculation of some expected values like the term $E[x_ix_i^T]$ during the calculation of expected gradients. While it is possible to implement our proof techniques to study other types of random walks, the calculations may be significantly more complicated and tedious. This is the main reason that we choose to study random walks in circles. We will add some discussions in the revision.\n>**Q6**: Absence of Softmax Scaling\n\n**A6**: Thanks for your question. The reason we omit the softmax scaling factor is mainly for simplicity, as such omission has been considered in most of the theoretical studies [3,4,5,6,7,8]. If we add the scaling factor, we expect that the main theoretical conclusions will remain unchanged. Intuitively, even if the value matrix V is trained much faster than the matrix W, the softmax function will predict a mixture distribution before W is sufficiently trained to assign an almost-one softmax score to the direct parent token, and the predicted distribution cannot be close to optimal. Therefore, in such cases, the loss will not be sufficiently minimized, and the training will not stop (gradients of W will not be small). Therefore, we can still expect that the value matrix and the softmax matrix will eventually be trained as the transition matrix and the direct parent selector, respectively. However, when such a scaling factor is added, it is likely that the training will take more iterations to converge to the desired model. We will add a detailed discussion in the revision.\n>**Q7**: Residual Connection Impact\n\n**A7**: Thank you for your question. We believe this is a misunderstanding. We consider the setting where the prediction is made by reading the output of the last column of the transformer output, while the corresponding input part of this column is set to zero (see equation (2.1)). When using the residual connection, the transformer model can be expressed as $f(X)=X+VXS(...)$. Since the last column of X is zero, the last column of the output is determined solely by the second term. Thus, when reading the last column as the prediction, this model's output is the same as ours. Therefore, adding the residual connection will not make parent token selection trivial. In fact, our current result can already be equivalently applied to the setting with residual connection, because the effective model output, the gradients of parameters, and thus the training dynamics, are actually the same with or without residual connection. We will clarify this point in the revision.\n>**Q8**: Alternative Token Dependencies+Higher-Order Markov Chains\n\n**A8**: Thank you for your suggestions. We believe that our theory can be extended to Markov chains that depend on a different position. We take the case of $x_{N-2}$ as an example. At the first step, $V^{(1)}$ is proportional to $\\sum_{i=1}^{[N/2]}(\\Pi^T)^i$, similar to Lemma C.1. Due to the Toeplitz property of the transition matrix, we can find that $V^{(1)}$ is also a Toeplitz matrix whose largest entries appear on the locations of the largest entries of $P_i^T$. At the next step, $V^{(1)}$ will encourage $W$ to assign the highest score to the parent $x_{N-2}$. The remaining is the same as that stated in our work. Therefore, we can extend our results to this specific task.\n\nRegarding high-order models, it may require multiple layers or multiple heads, since it is likely that one-layer transformers may not have enough expressive power to handle such cases. The softmax attention may only select one token, and the value matrix can not store too much information about the transition probabilities.\n\n[1] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.\n\n[2] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv (2024).\n\n[3] Tarzanagh et al. Transformers as support vector machines. In NeurIPS 2023 Workshop.\n\n[4] Li et al. How do transformers learn topic structure: Towards a mechanistic understanding. ICML 2023.\n\n[5] Li et al. One-Layer Transformer Provably Learns One-Nearest Neighbor In Context. NeurIPS 2024.\n\n[6] Anej et al. Transformers Can Represent n-gram Language Models. NAACL 2024.\n\n[7] D'Angelo et al. Selective induction heads: How transformers select causal structures in context. ICLR 2025.\n\n[8] Nichani et al. How transformers learn causal structure with gradient descent. ICML 2024.\n\n[9] Edelman et al. The evolution of statistical induction heads: In-context learning markov chains. NeurIPS 2024.\n\n[10] Makkuva et al. Attention with Markov: A curious case of single-layer transformers. ICLR 2024."}, {"review_id": "y1JLpHDpsc", "reviewer": "Reviewer_bJ9o", "summary": "The paper investigates how Transformers, specifically softmax attention without the feedforward network, perform on random walk problems. It addresses an interesting and relevant question, offering insightful analysis and well-motivated theoretical contributions. The paper presents clear motivations for the study and provides rigorous proofs regarding both the functionality of attention and its optimization dynamics in the context of learning random walks. The experimental results are supportive of the theoretical claims, and the paper thoughtfully explores various aspects such as initialization, positional embeddings, and architectural design. While I am not an expert in the theoretical analysis of random walks or Transformers in this specific context, I found the paper well-written and engaging in my opinion.", "questions": "As I am not very familiar with random walks, I hope these questions are not too basic. Some of them relate to Transformers and would be helpful to clarify for a broader audience.\n\n1. **Figure 2 and Circularity of Random Walks**  \n   In Figure 2, is the reason that both the top-right and bottom-left corners have probabilities `p` and `1 - p` due to the **circularity** of the random walk? That is, does the random walk wrap around the boundaries in this setup?\n\n2. **Role of Position Embedding**  \n   The position embedding used in this study appears to be specifically designed to help the Transformer solve the random walk task, effectively acting as an **inductive bias**. I’m curious: how would the **theoretical analysis** hold up **without** this guided position embedding? In practice, position embeddings like **RoPE** are often used, but may not be ideal for this type of structured task. It would be valuable to understand how the analysis extends to the case **without position embeddings**.\n\n3. **Initialization of Attention Weights**  \n   As far as I understand, there is **no clear theoretical result** regarding random initialization of attention weights in this setting. While the attention computation can be expressed as $S(XᵀWX)$, in practice, for interpretability and implementation, the weight matrix $W$ is typically decomposed as $W_Q @ W_K^T$. Could this decomposition affect or enhance the theoretical results?\n\n4. **Pretrained Initialization for Downstream Tasks**  \n   Following up on the weakness mentioned earlier, how would a Transformer perform on the tasks in Section 6 if it were first pretrained to **predict the next token** and then fine-tuned on tasks like identifying the preferable fruit? I would expect performance to improve significantly, as **random or zero initialization** is known to be insufficient for downstream tasks in Transformer-based models.", "limitations": "Please check the previous sections.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n- The paper clearly states the problem it is aiming to address and understand and the theoretical steps towards it are clear.\n- The experimental setup for the problem in the random walks are supporting the claims and well motivated.\n- Theoretical Proofs are solid and supporting the claims.\n\n## Weaknesses \n\nOverall, the paper is well-written and presents its ideas clearly. However, there are some presentation and formatting issues in certain sections that could be improved for better readability. In addition, I see one major weakness related to the initialization of the Transformer, and a conceptual gap in the tasks discussed in Section 6 (Beyond Random Walks), which seem to diverge from the core focus of the paper. These are, in my view, the most notable limitations. That said, I acknowledge that I may not be fully familiar with the specific literature in this area.\n\n- **Generalization Beyond Random Walks May Not Be Representative**  \nSection 6 attempts to extend the findings to more applied settings, such as language-related tasks. However, there is a major conceptual gap, particularly regarding initialization. It is well-established that Transformers — when evaluated on downstream tasks like those mentioned (e.g., identifying a preferred fruit or detecting repetition, as part of copying tasks) — do not perform well with random or zero initialization. Instead, they require pretraining, typically via next-token prediction, followed by fine-tuning for classification or multiple-choice tasks, as highlighted in [1].\n\n   This is consistent across all large language models: they are pretrained on large-scale language modeling tasks before being capable of handling such downstream problems. The use of zero initialization in the current study is, therefore, quite far from how Transformers are applied in practice.\n\n   In this sense, random walks are more analogous to pretraining tasks (e.g., next-token prediction), whereas the tasks in Section 6 resemble downstream evaluation tasks, which inherently require pretraining to be solvable. Thus, the theoretical results derived under the zero initialization regime may not generalize well to the tasks presented in Section 6. As such, the section’s title, \"Beyond Random Walks,\" may be misleading — it is unclear how far beyond and in what direction these extensions go, given the fundamental differences in task structure and training setup.\n\n- **Readability of Figures** Figures in general specifically from 4 to 7 are very hard to read and very small in font it would be better to enlarge font size ore resize for better readability.", "was_revised": true, "final_justification": "The authors' responses to my questions were convincing and satisfactorily addressed my concerns.", "rebuttal": ">**Q1**: Generalization Beyond Random Walks May Not Be Representative + Pretrained Initialization for Downstream Tasks\n\n**A1**: Thank you for your detailed feedback. We agree that modern large language models are typically trained on large-scale datasets, and the simple question-answering tasks we consider in Section 6 can definitely be easily handled by well-pretrained language models.\n\nHowever, we believe that this cannot deny the importance of studying the capability of transformers in learning simpler or more specific tasks, as such findings can still help us gain a more comprehensive understanding of when and how transformers can succeed. Our results show that direct training of transformers fails in some of the seemingly simpler tasks while succeeds in other tasks, which we believe can contribute to a better understanding of the nature of transformer models and their training processes.\n\nIn addition, we would like to clarify that the purpose of Section 6 is to demonstrate the value of our findings. In particular, Section 6 demonstrates that our theoretical insights derived from studying random walks can be extended to other learning tasks and help us to successfully identify success and failure cases in training simple transformer models. While it is an important question to theoretically study the performance of transformers under the setting of pretraining and fine-tuning, we feel that this is beyond the scope of our work.\n\nWe will change the title of Section 6 to “Similar Observations in Two Simple Question Answering Tasks” to clarify the goal of this section and be more precise about the type of tasks we are going to consider.\n\n>**Q2**: Readability of Figures\n\n**A2**: Thank you for your suggestion. We will enlarge the font size for better readability in the revised version.\n\n>**Q3**: Figure 2 and Circularity of Random Walks\n\n**A3**: Yes, you are right. Since we study random walks on a circle with a finite number of nodes, the boundary states 1 and $K$ are actually neighbours. This means the walker can move from node 1 to node $K$ with probability $1-p$, and from node $K$ to node 1 with probability $p$. This circular structure is why the top-right and bottom-left corners in the transition matrix have value $p$ and $1 - p$.\n\n>**Q4**: Role of Position Embedding\n\n**A4**: Thanks for your comment. We would like to clarify that the specific definitions of the position embeddings are not important. The definition in the equation between lines 108 and 109 is mainly to mimic the classic positional embeddings based on sine and cosine functions. Our analysis only relies on the property that embeddings of different positions are orthogonal to each other. As long as orthogonality is maintained, changes to the embedding matrix will not impact the validity of our analysis. Please note that orthogonal positional embeddings are common in existing theoretical studies: [1] employs positional embeddings based on orthogonal sine basis functions; [2] and [3] adopt one-hot positional embeddings that are inherently orthogonal; [4] considers near-orthogonal positional embeddings.\n\nIn our analysis, positional embeddings play a crucial role. In learning random walks, the key for a transformer model is to successfully identify the “direct parent token”. Positional embeddings can serve as identifiers of the direct parent. If positional embeddings are removed, the model may lack crucial information to solve the task optimally. In fact, we can show that if the input matrix does not consist of positional embeddings, any one-layer transformer cannot achieve a prediction accuracy better than a random guess. We will add this result and discussions about it in the revision.\n\nRegarding RoPE, since our current analysis relies on orthogonal positional embeddings, our analysis may not be directly applicable to RoPE. However, as we have mentioned, the setting with orthogonal positional embeddings is commonly considered in existing theoretical studies on transformers, and we believe that our current setting helps us to perform relatively clean theoretical analysis and therefore serves our purpose well. We will add comments that studying the case with RoPE is an interesting future work direction. \n\n>**Q5.1**: Initialization of Attention Weights: As far as I understand, there is no clear theoretical result regarding random initialization of attention weights in this setting. \n\n**A5.1**: The use of zero initialization in our theoretical analysis is primarily to simplify the dynamics of the model, which allows us to derive rigorous insights and obtain clear theoretical results. Our proof is based on the analyses of the softmax scores on each token during training. The setting with zero initialization ensures a fair start for all tokens, and we can prove that after training, the largest softmax score will be assigned to the direct parent token. When considering non-zero initialization, the initial scores assigned to all the tokens differ, leading to an unfair starting condition for training. This difference can largely influence and complicate the optimization analysis. If the direct parent token receives a relatively small initial score, the larger initial scores assigned to other tokens may cause the value matrix $V$ to learn a k-step transition matrix (other than the one-step transition), which can not achieve the optimal predictor. Additionally, in failure cases, the symmetry preserved under zero initialization is broken with non-zero initialization. This causes the largest softmax score to be assigned to an arbitrary token, which can enable the model to make an optimal prediction. Therefore, zero initialization serves as a useful idealization to facilitate theoretical understanding, while extending the results to non-zero initializations remains an interesting direction for future work. Please note that zero initialization or specific initialization has been considered in most theoretical studies [1,2,3,4,5,6].\n\n>**Q5.2**: While the attention computation can be expressed as $S(X^TWX)$, in practice, for interpretability and implementation, the weight matrix $W$ is typically decomposed as $W_QW_K^T$. Could this decomposition affect or enhance the theoretical results?\n\n**A5.2**: In our theoretical analysis, we consider a single weight matrix $W$ to denote $W_QW_K^T$ in self-attention, mainly for the simplicity of the theoretical analysis of the training dynamics. If we utilize the decomposition $W_QW_K^T$, it would require a more complex and careful analysis, as we need to deal with the simultaneous training of three parameter matrices $V$, $W_Q$, and $W_K$. In fact, we can no longer use zero initialization in this setting due to the interdependence of the training dynamics of $W_Q$ and $W_K$. The gradient of $W_Q$ involves $W_K$, and the gradient of $W_K$ also involves $W_Q$, which results in both gradients remaining zero with zero initialization. In this case, we have to consider random initialization. However, as discussed in **A5.1**, random initialization can largely complicate the analysis and may not yield clear theoretical results. There, we consider the reparameterization of a single weight matrix in our work. Please note that such reparameterization is widely considered in theoretical studies of transformer models [5,6,7,8]. \n\n[1] Zhang et al. Transformer learns optimal variable selection in group-sparse classification. ICLR 2025.\n\n[2] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot. ICML 2024.\n\n[3] Eshaan Nichani et al. How Transformers Learn Causal Structure with Gradient Descent. ICML 2024.\n\n[4] Bietti et al. Birth of a transformer: A memory viewpoint. NeurIPS 2024.\n\n[5] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.\n\n[6] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv preprint arXiv:2411.05189 (2024).\n\n[7] Huang et al. In-context convergence of transformers. ICML 2024.\n\n[8] Ildiz et al. From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformers. ICML 2024.\n\n[9] Jelassi et al. Vision transformers provably learn spatial structure. NeurIPS 2022.\n\n[10] Li et al. Mechanics of next token prediction with self-attention. AISTATS 2024."}, {"review_id": "8ZkI098wOC", "reviewer": "Reviewer_yFLc", "summary": "The paper considers single-layer 0-initialised attention mechanisms applied to the task of learning random walks on segmented circles. The authors present theoretical results about the ability of such a model to learn random walks when left/right transition probabilities are in $p \\\\in (0, 1)$. They further present theoretical results arguing that such a model fails to learn when $p \\\\in \\\\{0, 1\\\\}$. Both claims are empirically supported. They present 2 further tasks with small vocabulary and fixed grammar and argue that zero-initialised single-layer attention mechanisms can only learn when the average of tokens is informative for the task output.", "questions": "- In the zero initialisation setting, do the results extend to 2-layer attention models? Perhaps with nonlinear activations?\n- Could the authors comment further on the impact of their work, and how it can inform design of models / initalisation beyond the synthetic task considered?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- Convincing results for the zero-init setting for a single attention layer applied to learning random walks on circles\n- Interpretable / intuitive explanation of the mechanism the model learns, particularly in figure 3.\n- Thorough experiments in the setting considered, and intriguing extensions beyond this.\n\n\nWeaknesses:\n- It is not clear that the zero-initialisation setting is very useful in general. This initialisation is not common when using transformers in practice, for long-established empirical and theoretical reasons. Section 5 extends experiments to the random initialisation setting, and slow convergence is observed to the p=1 solution. It is not immediately apparent to me that this is a consqeuence of the author's prior analysis. I believe the argument here could be strengthened, by: \n    1. Showing empirically that even in the presence of random initialisation, the attention model is trained approximately on the \"uninformative\" averaged token.\n    2. Discussing the sudden jumps in accuracy. Are these e.g., because of random noise happening upon an output that is sufficiently far away from the uniform averaging? Though if that were the case, would the zero-init model also stumble on such a point?\n- The results in Section 6 appear at best weakly related to the preceding content. The conclusion appears to be that the model is only able to learn when the mean of the tokens is informative for the label. Stated differently, a conjecture this paper seems to make is \"Zero initialised models only learn when token averages have discriminative value\".  This is an interesting result, which could be strengthened with further experimental support.\n\n\nOverall, I think this is a good paper that makes a claim backed up by reasonable results. My primary concern is that the impact of the impact of the claim appears limited. I would look for further applications of this result to strengthen this paper.", "was_revised": true, "final_justification": "The paper presents an interesting method and result for analysing the behavior of single-layer transformers when learning Markov dynamics under 0-init conditions. While open questions and future directions are apparent, it is not necessary for these to be answered in this one paper, which is valuable without these.", "rebuttal": ">**Q1**: It is not clear that the zero-initialization setting is very useful in general. Section 5 extends experiments to the random initialisation setting. It is not apparent to me that this is a consequence of the author's prior analysis. Suggestions: Show empirically + Discuss sudden jumps in accuracy.\n\n**A1**: The use of zero initialization in our theoretical analysis is primarily to simplify the dynamics of the model, which allows us to derive rigorous insights and obtain clear theoretical results. Our proof is based on the analyses of the softmax scores on each token during training. The setting with zero initialization ensures a fair start for all tokens, and we can prove that after training, the largest softmax score will be assigned to the direct parent token. When considering non-zero initialization, the initial scores assigned to all the tokens differ, leading to an unfair starting condition for training. This difference can largely influence and complicate the optimization analysis. If the direct parent token receives a relatively small initial score, the larger initial scores assigned to other tokens may cause the value matrix $V$ to learn a k-step transition matrix (other than the one-step transition), which can not achieve the optimal predictor. Additionally, in failure cases, the symmetry preserved under zero initialization is broken with non-zero initialization. This causes the largest softmax score to be assigned to an arbitrary token, which can enable the model to make an optimal prediction. Therefore, zero initialization serves as a useful idealization to facilitate theoretical understanding, while extending the results to non-zero initializations remains an interesting direction for future work. Please note that zero initialization or specific initialization has been considered in most theoretical studies [1,2,3,4].\n\nFor the first suggestion, we provide the following results about the weighted average token ($X \\cdot \\text{softmax}(\\tilde{X}W\\tilde{x})$) corresponding to Figure 6(c) and 6(d).\n\n|Training steps|Weighted average token $X\\text{softmax}(\\tilde{X}W\\tilde{x})$|Variance|\n|----|----|----|\n|100 steps| [0.1648, 0.1650, 0.1653, 0.1752, 0.1646, 0.1652] |[1.499e-5, 1.522e-5, 1.404e-5, 1.529e-5, 1.516e-5, 1.281e-5]| \n| 300 steps|[0.1633, 0.1649, 0.1676, 0.1753, 0.1617, 0.1672] |[1.912e-5, 2.006e-5, 1.856e-5, 1.974e-5, 2.007e-5, 1.741e-5] |\n| 500 steps|[0.1564, 0.1638, 0.1775, 0.1769, 0.1499, 0.1756] |[1.120e-4, 1.164e-4, 1.166e-4, 1.167e-4, 1.168e-4, 1.162e-4] |\n| 700 steps|[0.1328, 0.1538, 0.2147, 0.1766, 0.1183, 0.2037] |[1.218e-3, 1.209e-3, 1.261e-3, 1.251e-3, 1.201e-3, 1.274e-3] |\n| 900 steps|[0.0766, 0.1001, 0.3765, 0.1300, 0.0642, 0.2526] |[1.284e-2, 1.223e-2, 1.303e-2, 1.288e-2, 1.164e-2, 1.300e-2] |\n\nIn the table above, the weighted average token embedding is calculated based on one sequence X as an example, while the variance of the weighted average token is calculated across 1000 sequences. We can observe that when considering the case for p=1 with random initialization, the weighted average token embedding remains approximately proportional to the all-one vector within 500 iterations. This indicates that the average token provides very limited useful information for training V, aligning with the results stated in Theorem 4.1. Additionally, the variances of all dimensions are close to 0, which shows that this result is general for different sequences. Although after 900 steps, some dimensions in the weighted average token embedding are relatively larger than others, it is still not enough to enable the transformer to achieve the optimal prediction, as illustrated in Figure 6(c).\n\nFor the second suggestion, as mentioned in the first paragraph, the starting point of the softmax score attached to each token is different. This means that sufficient training steps are needed to update the weight matrix $W$ so that the largest softmax score can be assigned to the direct parent token. Before that, the attention may select another token, which can lead to a lower accuracy than the optimal one. In addition, since we consider 6 nodes in our experiments, the accuracy can only take values from a finite set. Thus, due to the long training steps needed and the discrete nature of accuracy, sudden jumps may occur when using random initialization. In contrast, when starting from zero initialization, the model begins fairly, and the softmax score for the direct parent token can quickly become the highest within the first few steps, which facilitates rapid achievement of the optimal accuracy.\n\n>**Q2**: The results in Section 6 appear at best weakly related to the preceding content. \n\n**A2**: Thank you for your comment. We would like to clarify that the purpose of Section 6 is to demonstrate the value of our findings. In particular, Section 6 demonstrates that our theoretical insights derived from studying random walks can be extended to other learning tasks and help us to successfully identify success and failure cases in training simple transformer models. Specifically, in Task 2, the data has the property that the token averages remain the same across all sequences; thus, we expect that Task 2 is a failure case for the simple transformer model, which is verified in Figure 7.\n\n>**Q3**: 2-layer attention models?\n\n**A3**: We agree that considering more complex models is important and aligns better with the practical setting. However, precise theoretical analysis of more complicated architectures can be extremely complicated. While our theoretical analysis can not be directly extended to 2-layer models, we believe that the theoretical insights derived from studying one-layer transformers may be extended to understand deeper models. In Theorem 3.1, we demonstrate that simple one-layer transformers can already successfully learn random walks. This may indicate the capability of multi-layer transformers in learning the same task as well. In Theorem 4.1, we demonstrate that one-layer transformers fail to learn deterministic walks due to the fact that the average of tokens does not provide useful information for prediction. We believe this insight can also be applied to learning the same task using multi-layer transformers. When the token average is not informative, multi-layer transformers may also struggle with zero initialization since it always takes an average over all tokens in each layer. As a result, the output of multiple layers may lack the necessary information for prediction.\n\nPlease note that many recent theoretical studies on the training dynamics of transformers focus on simple one-layer transformer models [3,4,5,6]. In fact, to our knowledge, the setting considered in our paper is already more aligned with the practical transformer architecture compared to many of these recent theoretical studies. For example, [3,4] conduct the theoretical analysis on the transformer with a linear attention (instead of softmax attention). And, in [5,6], the value matrix $V$ is not involved in the training process; instead, the value matrix is held constant while other parameters are trained. In comparison, we construct a transformer architecture with nonlinear softmax attention and analyze the training of the matrices $V$ and $W$ simultaneously, which is a step toward a more practical study.\n\n>**Q4**: the impact of their work\n\n**A4**: To our knowledge, the most related work to ours is [7]. [7] characterizes the loss landscape of a one-layer transformer and identifies global minima and bad local minima based on Markovian data with vocabularies of size two, which can also be interpreted as a random walk over a space of only two states. In contrast, we consider random walk tasks on a circle that allows a large number of states, which can be seen as an extension to the two-state setting studied in [7]. This suggests that our theory represents a step further towards a unified theory for more general cases.\n\nIn addition, our work provides a rigorous theoretical framework explaining how transformers learn Markovian data. We demonstrate that the trained transformer model is interpretable and can be decomposed into two parts: the token selector (attention focusing on the direct parent) and the one-step transition model (value matrix recovering $\\Pi^T$). This reveals the token selection ability of softmax attention, enhancing our understanding of how softmax attention layers work in more complicated models. Moreover, this capability also addresses the reason why transformers can efficiently handle sequential data in NLP tasks, as softmax attention can capture the underlying structure of the data.\n\nWhile our work investigates the capability of transformers in learning simple or specific tasks, the findings contribute to a more comprehensive understanding of when and how transformers can succeed. This understanding can guide the design of transformer models for more complicated tasks. We demonstrate that the failure cases for $p=0\\text{or}1$ are optimization failures, and are highly due to the zero initialization. This highlights the importance of avoiding unbreakable symmetry in optimization and emphasizes the need to select appropriate initialization strategies. \n\n[1] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot. ICML 2024.\n\n[2] Zhang et al. Transformer learns optimal variable selection in group-sparse classification. ICLR 2025.\n\n[3] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.\n\n[4] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv (2024).\n\n[5] Tarzanagh et al. Transformers as support vector machines. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning.\n \n[6] Li et al. How do transformers learn topic structure: Towards a mechanistic understanding. ICML 2023.\n\n[7] Makkuva et al. Attention with markov: A curious case of single-layer transformers. ICLR 2024."}]

Author rebuttal:

Meta-review: This paper studies the dynamics of gradient descent on a one-layer Transformer trained on a specific Markov chain arising from a biased random walk. The authors show that gradient descent successfully solves this task in a constant number of iterations. Specifically, the authors prove that for transition probabilities p∈(0,1), gradient descent from zero initialization achieves optimal prediction accuracy, through an interpretable mechanism in which the self-attention layer learns to act as a token selector, while the value matrix converges to the transition matrix. In contrast, for deterministic walks (p=0,1), the paper demonstrates a failure case caused by initial symmetries, where the model’s performance is no better than random guessing.

The reviewers found the paper’s results to be an interesting and valid contribution. They appreciated that the analysis is conducted on a real softmax Transformer and noted that the paper is well written, with clear explanations of the proof ideas and optimization bottlenecks. However, as highlighted in the reviews, the analysis relies on a simplified setting with strong structural assumptions that are fundamental to the results, which limits the generality of the findings.

Overall, I recommend acceptance and encourage the authors to incorporate the reviewers’ feedback in the final version.
Decision: Accept (poster)