Research map / Architectures and efficiency / Attention and transformers
How Transformers Learn Causal Structure with 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 that a two-layer attention-only (disentangled) transformer trained by stage-wise gradient descent learns the causal (induction-head-style) structure of Markov chain sequences, approximately minimizing loss.
Problem
How transformers learn causal structure in sequences via gradient-based training, i.e., the optimization dynamics that lead attention to recover the parent (previous-token / DAG) structure needed for in-context next-token prediction.
Approach
Introduces a disentangled transformer (attention outputs appended to the residual stream) and a reduced two-layer model, then proves that a three-stage gradient descent schedule on population loss drives the first-layer position attention to the shift (induction-head) matrix, with the second layer growing to reduce loss; also generalizes to random DAG causal structures and multiple parents.
Key results
- Theorem: with stage-wise GD, the learned loss satisfies L(theta_hat) - L* <= S/T^{1/4}, and the first-layer attention satisfies s(A1)_{i,i-1} >= 1 - O(T^{-3}) (shift/induction structure learned).
- Two-layer disentangled transformer with 2k heads can exactly implement the empirical multi-parent transition estimator in the limit, reaching near-optimal loss (o_T(1) gap).
- Experiments section is empty (placeholder 'todo'), so no empirical validation is reported.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train the reduced two-layer transformer with three-stage gradient descent | Reduced two-layer transformer · Random sequences with causal structure | excess cross-entropy loss: L(\hat \theta) - L^* \lesssim_{\gamma,\lambda} \frac{S}{T^{1/4}} | Helped |
| Represent each node's parent in the first attention layer | Depth-two multiple-head transformer · Random sequences with a general causal DAG | excess cross-entropy loss: L(\hat \theta) - L^* = o_T(1) | Helped |
| Use 2k first-layer heads to copy the k parent tokens and match parent contexts | Two-layer transformer · Sequences with k-parent causal structure | transition probability estimate | Helped |
Limitations
- Experiments are not provided in the text (section marked todo)
- Analysis restricted to a reduced two-layer model with one parent per position in main theorem
- Many proofs and assumptions are incomplete (placeholders noted)
Experimental setup
Datasets: synthetic Markov chain sequences with causal DAG structure
Metrics: population cross-entropy loss gap to optimal
Builds on: induction head tracr
Closest papers by meaning
- Selective Induction Heads: How Transformers Select Causal Structures in Context (ICLR 2025 · similarity 0.92)
- Towards Understanding the Universality of Transformers for Next-Token Prediction (ICLR 2025 · similarity 0.91)
- How Transformers Learn Structured Data: Insights From Hierarchical Filtering (ICML 2025 · similarity 0.88)
- A Fixed-Point Approach for Causal Generative Modeling (ICML 2024 · similarity 0.88)
- Linear Attention for Efficient Bidirectional Sequence Modeling (NeurIPS 2025 · similarity 0.87)
- Causality-Induced Positional Encoding for Transformer-Based Representation Learning of Non-Sequential Features (NeurIPS 2025 · similarity 0.87)
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.