How Transformers Learn Causal Structure with Gradient Descent
summary
Gradient descent on a two-layer disentangled transformer with one-hot tokens and softmax attention provably encodes a latent DAG: the first attention layer converges to the graph adjacency and the second implements copying (induction heads) for Markov chains. The attention gradient matches a chi-square mutual information measure so the largest entries identify parents; analysis assumes single-parent sequences from Dirichlet-sampled transitions with mixing and gives a polynomial-in-S rate (informal exponent ≈1/50), and experiments recover many DAGs.
AI-written summary · about
chapters (11)
- 0:00Introduction
- 0:22Induction heads in transformers
- 6:00Causal structure data
- 10:15Transformer architecture
- 14:31Disentangled transformers
- 17:01Training two attention blocks
- 21:04Learning causal graphs
- 27:00Gradient dynamics
- 30:17Multiple parents in n-grams
- 33:11Limitations and future
- 37:00Q&A
related talks
01h 10mIntroduction to Transformers
01h 16mHow Do Transformers Work?
01h 04mTransformers can learn compositional function
46mLearning Theory of Transformers: Generalization and Optimization of In-Context Learning
01h 30mToward Understanding In-context Learning
28mEverything you need to know about Transformer Models
24mRevealing the Unseen: Visualizing Attention in Transformer Models
26mA Shallow Introduction to Self-Attention