Research map / Architectures and efficiency / Attention and transformers
Towards Understanding Transformers in Learning Random Walks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that a one-layer softmax transformer trained by zero-initialized gradient descent provably learns random walks on a circle for 0<p<1 (selecting the parent token and recovering the transition matrix), while p=0 or 1 deterministically fails due to symmetry.
Problem
Understanding whether and how transformers can learn random walks on circles with many states, and characterizing the interpretability of the trained model, including edge cases of deterministic walks (p=0,1).
Approach
Analyzes gradient descent dynamics from zero initialization of a one-layer transformer (reparameterized attention W, value V, log-loss) on random walks over K-node circles, showing attention concentrates on the direct parent token and V converges in direction to the transition matrix; for deterministic walks, symmetry from zero init prevents learning.
Key results
- Theory: for 0<p<1, trained transformer attains accuracy max{p,1-p} and attention on the parent token is at least 1-exp(-Omega(N)).
- Theory: for p=0 or 1 with zero initialization, prediction accuracy equals 1/K for any loss and learning rate.
- Experiments with K=6, N=97: zero-init p=1 gives accuracy about 0.167 (=1/K), while random init p=0.5 reaches the optimal 50% in around 400 iterations and p=1 does not reach 100% within 1000 iterations.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train from zero initialization on random walks with 0<p<1 | One-layer transformer · Random walks on circles | prediction accuracy: max\{p,1-p\} | Helped |
| Train from zero initialization on deterministic walks (p=0 or 1) | One-layer transformer · Deterministic walks, p=0 or 1, N=rK+1 | prediction accuracy: \frac{1}{K} | Hurt |
| Use zero initialization | One-layer transformer · Deterministic walk, p=1 (K=6, N=97) | test accuracy: 0.167 | Hurt |
| Use Gaussian random initialization (σ=0.01) instead of zero initialization | One-layer transformer · Deterministic walk, p=1 (K=6, N=97) | test accuracy: 100% | Mixed |
| Use Gaussian random initialization (σ=0.01) | One-layer transformer · Random walk, p=0.5 (K=6, N=97) | test accuracy: 50% | Helped |
| Train on questions where the word-embedding average distinguishes the answer | One-layer transformer · Question answering Task 1 (fruit frequency) | test accuracy: 100% | Helped |
| Train on questions with the same average word embedding for both answers | One-layer transformer · Question answering Task 2 (stated preference) | test accuracy: 50% | Hurt |
Surprising findings
- Deterministic walks (p=0 or 1) are harder to learn than stochastic walks under zero initialization, with failure for any loss function and learning rate.
- In the QA tasks, identical average embeddings across questions cause inefficient optimization, with accuracy stuck near 50% for Task 2.
Limitations
- Analysis restricted to a one-layer transformer with reparameterized attention and no softmax scaling.
- Theory covers zero initialization only; random initialization is studied empirically.
- Experiments are small synthetic settings (K=6, N=97).
Experimental setup
Datasets: synthetic random walks on circles (K=6, N=97) synthetic QA tasks (fruit counting, fruit preference)
Benchmarks: synthetic random walks on circles synthetic QA tasks
Metrics: prediction accuracy KL divergence to true next-state distribution attention scores normalized value matrix direction error
Baselines: optimal predictor f^OPT = Pi^T x_{N-1} random guess
Builds on: Markov chain in-context learning analysis (Huang et al.) reparameterized attention in transformer theory
Closest papers by meaning
- Non-asymptotic Convergence of Training Transformers for Next-token Prediction (NeurIPS 2024 · similarity 0.87)
- Algorithmic Capabilities of Random Transformers (NeurIPS 2024 · similarity 0.87)
- Revisiting Random Walks for Learning on Graphs (ICLR 2025 · similarity 0.87)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.86)
- Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.86)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.