# How Transformers Learn Causal Structure with Gradient Descent

arXiv: https://arxiv.org/abs/2402.14735 · HTML: https://mlsanity.com/paper/2402.14735/

## Key facts

- Cited by 118 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 29 references that are in the map.
- First posted on arXiv 22 Feb 2024.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Vision Transformers provably learn spatial structure](https://mlsanity.com/paper/2210.09221/) (2022)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)
- [Transformers learn through gradual rank increase](https://mlsanity.com/paper/2306.07042/) (2023)
- [JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention](https://mlsanity.com/paper/2310.00535/) (2023)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)
- [Learning Transformer Programs](https://mlsanity.com/paper/2306.01128/) (2023)

## Built on by (most cited first, 30 of 118)

- [Mechanistic Interpretability for AI Safety -- A Review](https://mlsanity.com/paper/2404.14082/) (2024)
- [Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs](https://mlsanity.com/paper/2410.13835/) (2024)
- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [In-Context Learning with Representations: Contextual Generalization of Trained Transformers](https://mlsanity.com/paper/2408.10147/) (2024)
- [When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers](https://mlsanity.com/paper/2504.10957/) (2025)
- [Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers](https://mlsanity.com/paper/2405.05219/) (2024)
- [Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought](https://mlsanity.com/paper/2502.21212/) (2025)
- [Iteration Head: A Mechanistic Study of Chain-of-Thought](https://mlsanity.com/paper/2406.02128/) (2024)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency](https://mlsanity.com/paper/2410.05459/) (2024)
- [There Will Be a Scientific Theory of Deep Learning](https://mlsanity.com/paper/2604.21691/) (2026)
- [Out-of-distribution generalization via composition: a lens through induction heads in Transformers](https://mlsanity.com/paper/2408.09503/) (2024)
- [Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamics](https://mlsanity.com/paper/2405.04669/) (2024)
- [Transformers on Markov Data: Constant Depth Suffices](https://mlsanity.com/paper/2407.17686/) (2024)
- [Metastable Dynamics of Chain-of-Thought Reasoning: Provable Benefits of Search, RL and Distillation](https://mlsanity.com/paper/2502.01694/) (2025)
- [The emergence of sparse attention: impact of data distribution and benefits of repetition](https://mlsanity.com/paper/2505.17863/) (2025)
- [Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot](https://mlsanity.com/paper/2406.06893/) (2024)
- [How Do LLMs Perform Two-Hop Reasoning in Context?](https://mlsanity.com/paper/2502.13913/) (2025)
- [Implicit Regularization of Gradient Flow on One-Layer Softmax Attention](https://mlsanity.com/paper/2403.08699/) (2024)
- [Toward Understanding In-context vs. In-weight Learning](https://mlsanity.com/paper/2410.23042/) (2024)
- [Compact Proofs of Model Performance via Mechanistic Interpretability](https://mlsanity.com/paper/2406.11779/) (2024)
- [Attention layers provably solve single-location regression](https://mlsanity.com/paper/2410.01537/) (2024)
- [Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Points](https://mlsanity.com/paper/2508.12837/) (2025)
- [Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thought](https://mlsanity.com/paper/2509.23365/) (2025)
- [Understanding the Language Model to Solve the Symbolic Multi-Step Reasoning Problem from the Perspective of Buffer Mechanism](https://mlsanity.com/paper/2405.15302/) (2024)
- [Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers](https://mlsanity.com/paper/2406.18400/) (2024)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [Why and How LLMs Hallucinate: Connecting the Dots with Subsequence Associations](https://mlsanity.com/paper/2504.12691/) (2025)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (2024)
- [Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis](https://mlsanity.com/paper/2410.02167/) (2024)

## Abstract (arXiv)

The incredible success of transformers on sequence modeling tasks can be largely attributed to the self-attention mechanism, which allows information to be transferred between different parts of a sequence. Self-attention allows transformers to encode causal structure which makes them particularly suitable for sequence modeling. However, the process by which transformers learn such causal structure via gradient-based training algorithms remains poorly understood. To better understand this process, we introduce an in-context learning task that requires learning latent causal structure. We prove that gradient descent on a simplified two-layer transformer learns to solve this task by encoding the latent causal graph in the first attention layer. The key insight of our proof is that the gradient of the attention matrix encodes the mutual information between tokens. As a consequence of the data processing inequality, the largest entries of this gradient correspond to edges in the latent causal graph. As a special case, when the sequences are generated from in-context Markov chains, we prove that transformers learn an induction head (Olsson et al., 2022). We confirm our theoretical findings by showing that transformers trained on our in-context learning task are able to recover a wide variety of causal structures.