Title: Towards Understanding Transformers in Learning Random Walks
Venue: NeurIPS 2025
Accepted: yes
Keywords: transformer

self-attention

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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nAs 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.\n\nLIMITATIONS:\nPlease check the previous sections.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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).\n\nQUESTIONS:\n1. **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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n- 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'$)?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-07-05"}