Q1: Generalization Beyond Random Walks May Not Be Representative + Pretrained Initialization for Downstream Tasks

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.

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

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

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

Q2: Readability of Figures

A2: Thank you for your suggestion. We will enlarge the font size for better readability in the revised version.

Q3: Figure 2 and Circularity of Random Walks

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

Q4: Role of Position Embedding

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.

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

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

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.

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

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?

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

[1] Zhang et al. Transformer learns optimal variable selection in group-sparse classification. ICLR 2025.

[2] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot. ICML 2024.

[3] Eshaan Nichani et al. How Transformers Learn Causal Structure with Gradient Descent. ICML 2024.

[4] Bietti et al. Birth of a transformer: A memory viewpoint. NeurIPS 2024.

[5] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.

[6] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv preprint arXiv:2411.05189 (2024).

[7] Huang et al. In-context convergence of transformers. ICML 2024.

[8] Ildiz et al. From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformers. ICML 2024.

[9] Jelassi et al. Vision transformers provably learn spatial structure. NeurIPS 2022.

[10] Li et al. Mechanics of next token prediction with self-attention. AISTATS 2024.

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.

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

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

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

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

Q2: The results in Section 6 appear at best weakly related to the preceding content.

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.

Q3: 2-layer attention models?

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.

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

Q4: the impact of their work

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.

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

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

[1] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot. ICML 2024.

[2] Zhang et al. Transformer learns optimal variable selection in group-sparse classification. ICLR 2025.

[3] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.

[4] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv (2024).

[5] Tarzanagh et al. Transformers as support vector machines. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning.

[6] Li et al. How do transformers learn topic structure: Towards a mechanistic understanding. ICML 2023.

[7] Makkuva et al. Attention with markov: A curious case of single-layer transformers. ICLR 2024.

Q1: Simplified Model

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.

Q2: Interpretability Contribution

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

Q3: Zero Initialization

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

Q4: Task Specificity

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.

Q5: Circle Constraint Motivation

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.

Q6: Absence of Softmax Scaling

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.

Q7: Residual Connection Impact

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.

Q8: Alternative Token Dependencies+Higher-Order Markov Chains

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.

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

[1] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.

[2] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv (2024).

[3] Tarzanagh et al. Transformers as support vector machines. In NeurIPS 2023 Workshop.

[4] Li et al. How do transformers learn topic structure: Towards a mechanistic understanding. ICML 2023.

[5] Li et al. One-Layer Transformer Provably Learns One-Nearest Neighbor In Context. NeurIPS 2024.

[6] Anej et al. Transformers Can Represent n-gram Language Models. NAACL 2024.

[7] D'Angelo et al. Selective induction heads: How transformers select causal structures in context. ICLR 2025.

[8] Nichani et al. How transformers learn causal structure with gradient descent. ICML 2024.

[9] Edelman et al. The evolution of statistical induction heads: In-context learning markov chains. NeurIPS 2024.

[10] Makkuva et al. Attention with Markov: A curious case of single-layer transformers. ICLR 2024.

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.

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.

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.

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

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

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?

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.

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

We will add these visualizations in the revision.

Moreover, 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].

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?

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.

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?

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.

Q6: Do similar conclusions about interpretability hold for the case of Markov chains with more complex transition matrices?

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.

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

[1] Zhang et al. Trained transformers learn linear models in-context. JMLR 2024.

[2] Usman et al. Adversarial Robustness of In-Context Learning in Transformers for Linear Regression. arXiv preprint arXiv:2411.05189 (2024).

[3] Li et al. How do transformers learn topic structure: Towards a mechanistic understanding. ICML 2023.

[4] Li et al. One-Layer Transformer Provably Learns One-Nearest Neighbor In Context. NeurIPS 2024.

[5] Tarzanagh et al. Transformers as support vector machines. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning.

[6] Zhang et al. Anchor function: a type of benchmark functions for studying language models. arXiv preprint arXiv:2401.08309 (2024).

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

[8] D'Angelo et al. Selective induction heads: How transformers select causal structures in context. ICLR 2025.

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.

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.

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

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

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.

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.

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.

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

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

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’$?

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.

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

[1] Makkuva et al. Attention with Markov: A curious case of single-layer transformers. ICLR 2024.

[2] Rajaraman et al. An analysis of tokenization: Transformers under markov data. NeurIPS 2024.

[3] Ildiz et al. From self-attention to Markov models: Unveiling the dynamics of generative transformers. ICML 2024.

[4] Li et al. Mechanics of next token prediction with self-attention. International Conference on Artificial Intelligence and Statistics. AISTATS, 2024.

[5] Nichani et al. How transformers learn causal structure with gradient descent. ICML 2024.

[6] Edelman, Ezra, et al. The evolution of statistical induction heads: In-context learning Markov chains. NeurIPS 2024.

Dear Reviewer ki3b,

As the author-reviewer discussion period will end soon, we would like to kindly follow up to see if our responses have adequately addressed your concerns.

Regarding your concern about the residual connection, we have clarified in our previous reply that this is a misunderstanding. Our setting already covers scenarios involving the residual connection, and even with the residual connection, the model still relies on softmax attention to select the direct parent token for optimal prediction. For more details, please refer to our previous response.

Thank you for your time and effort in reviewing our work. If your concerns are addressed, we would sincerely appreciate your reconsideration of the evaluation.

Best regards,

Authors of Submission20613

Dear Reviewer ki3b,

We would like to clarify again that your concern about the residual connection is actually a misunderstanding. The transformer model with the residual connection cannot handle the task more easily.

Additionally, we would like to reemphasize that our goal is to understand the capability and interpretability of transformers when learning classic statistical tasks. Although other models (especially those designed with prior knowledge of Markovian property) can also deal with the random walk tasks, this cannot deny the importance of studying the capability of transformers in learning such tasks. We believe our findings can inspire a more comprehensive understanding of when and how transformers can succeed.

We sincerely appreciate your time and effort in reviewing our work. If your concerns are addressed, we would be grateful for your reconsideration of the evaluation.

Best regards,

Authors of Submission20613