Research map / Architectures and efficiency / Attention and transformers
Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent
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 theoretically and empirically that one-layer (one- or two-head) transformers trained with CoT by gradient descent provably learn backward and forward path-finding in trees, with generalization to unseen trees.
Problem
How shallow transformers learn symbolic multi-step reasoning (path-finding in trees) through chain-of-thought during gradient-based training, including autonomous switching between two reasoning stages.
Approach
Explicit one-layer constructions plus gradient-descent training dynamics showing a two-head transformer learns to output the goal-to-root path and then reverse it, with one head traversing the path and the other head controlling the stage switch at the root; the extended CoT replaces the depth that prior work used.
Key results
- Backward reasoning: gradient descent from zero initialization drives training loss below epsilon in T=O(mS/(eta*eps)) iterations, with test loss bounded in terms of epsilon for unseen trees.
- Forward reasoning: a two-head one-layer transformer converges with T=O(S/eta (m/eps)^{3/2} log(N/eps)) iterations, with a generalization bound on unseen trees.
- Experiments with one-hot embeddings show training and test losses converging to zero for both tasks, matching the theoretical attention-parameter dynamics.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train a one-layer, single-head transformer with autoregressive chain-of-thought for backward path finding | One-layer transformer · Randomly generated trees, backward reasoning | training loss | Helped |
| Train a one-layer, single-head transformer with autoregressive chain-of-thought for backward path finding | One-layer transformer · Unseen trees, backward reasoning | test loss | Helped |
| Use two attention heads to traverse the path and control the reasoning stage | One-layer transformer · Randomly generated trees, forward reasoning | training loss | Helped |
| Use two attention heads to traverse the path and control the reasoning stage | One-layer transformer · Unseen trees, forward reasoning | test loss | Helped |
Surprising findings
- A one-layer model can solve forward reasoning, which prior work solved with deep multi-layer transformers, by using the extended CoT to generate the backward path as a scratchpad.
Limitations
- Training analysis restricted to perfect binary trees and orthonormal (one-hot) embeddings
- Analysis limited to one-layer transformers with specific parameterization and zero initialization
- Experiments are small synthetic settings only
Experimental setup
Datasets: synthetic random perfect binary trees randomly generated test trees (1024)
Benchmarks: synthetic random tree test set (1024 trees)
Metrics: squared error loss on test trees training loss entries of attention parameters
Builds on: chain-of-thought parity learning with CoT single-head dynamics
Closest papers by meaning
- Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization (NeurIPS 2025 · similarity 0.92)
- Transformers Provably Solve Parity Efficiently with Chain of Thought (ICLR 2025 · similarity 0.89)
- How do Transformers Learn Implicit Reasoning? (NeurIPS 2025 · similarity 0.89)
- Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization (NeurIPS 2024 · similarity 0.89)
- Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data? (ICLR 2025 · similarity 0.88)
- Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thought (ICLR 2026 · similarity 0.88)
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.