Towards Understanding Transformers in Learning Random Walks

TL;DR

This work theoretically and empirically analyzes a one-layer transformer trained by gradient descent on the task of predicting the next state in random walks on circles. It demonstrates that for 0<p<1, the model attains optimal prediction accuracy and, importantly, is interpretable: softmax attention nearly always selects the direct parent token and the value matrix implements the one-step transition. The paper also proves that deterministic walks (p=0 or 1) are failure cases under zero initialization, and supports these findings with comprehensive experiments, including extensions to simple QA tasks. Together, these results provide sharp insight into the learnability and interpretability of transformers for classical stochastic processes and point to how initialization influences optimization in even simple tasks.

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.

Figures (13)

• Figure 1: Illustration of random walks on circles with $K$ nodes and transition probability $p$.
• Figure 2: Visualization of the transition matrices $\bm{\Pi}$ with $p=0.5, 0.7, 1.0$ respectively.
• Figure 3: Illustration of how the trained one-layer transformer performs optimal prediction in an example with $p=1/2$, $K=6$. Here we denote $\boldsymbol{S} = {\mathcal{S}}(\Tilde{\bm{X}}^{\top}\bm{W}^{(T)}\Tilde{\bm{x}}_N )$.
• Figure 4: The results of the experiments for $p=0.5$ with zero initialization: (a) is the test accuracy; (b) is the visualization of $\bm{V}$; (c) and (d) present the average attention of the test data with $x$-axis representing the position of the token and $y$-axis representing the attention score.
• Figure 5: Experiments for $p=1$ with zero initialization. (a) gives the prediction accuracy along training. (b) is the visualization of $\bm{V}$. (c) is the average attention of the test data with $x$-axis representing the position of the token and $y$-axis representing the attention score.

Theorems & Definitions (21)

• Definition 2.1
• Theorem 3.1
• Theorem 4.1
• Lemma B.1
• Lemma C.1
• Lemma C.2
• Lemma C.3
• Lemma C.4
• Lemma C.5
• Lemma C.6