mlsanityData to 7 Oct 2026

How Transformers Learn Causal Structure with Gradient Descent

arXiv 2402.14735 ↗First posted 22 Feb 2024cs.LG cs.IT math.ITcited by 114

Not on the map. Why?

Cited by 24 in LLM reasoning, but placed by its own abstract, which didn’t fit.

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.

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

Builds on

24 of 29 references.

References · 24

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
  5. 2022In-context Learning and Induction Headscited by 930
  6. 2021An Explanation of In-context Learning as Implicit Bayesian Inferencecited by 915
  7. 2022Transformers learn in-context by gradient descentcited by 714
  8. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  9. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  10. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
  11. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 291
  12. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 274
12 more
  1. 2023Looped Transformers as Programmable Computerscited by 224
  2. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 168
  3. 2023Birth of a Transformer: A Memory Viewpointcited by 154
  4. 2023In-Context Convergence of Transformerscited by 112
  5. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 101
  6. 2022Vision Transformers provably learn spatial structurecited by 92
  7. 2023Transformers as Support Vector Machinescited by 67
  8. 2023Max-Margin Token Selection in Attention Mechanismcited by 57
  9. 2023Transformers learn through gradual rank increasecited by 52
  10. 2023JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attentioncited by 52
  11. 2023Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regressioncited by 47
  12. 2023Learning Transformer Programscited by 40

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 38; 2025: 46; 2026: 30.

114 papers.

Sort:

By area:

  1. 2024Mechanistic Interpretability for AI Safety -- A Reviewcited by 415
  2. 2024Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMscited by 54
  3. 2025Training Dynamics of In-Context Learning in Linear Attentioncited by 44
  4. 2025When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformerscited by 38
  5. 2024In-Context Learning with Representations: Contextual Generalization of Trained Transformerscited by 36
  6. 2024Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformerscited by 32
  7. 2025Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thoughtcited by 32
  8. 2024Iteration Head: A Mechanistic Study of Chain-of-Thoughtcited by 31
  9. 2024Pretrained transformer efficiently learns low-dimensional target functions in-contextcited by 29
  10. 2024Out-of-distribution generalization via composition: a lens through induction heads in Transformerscited by 28
  11. 2024From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiencycited by 28
  12. 2024Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamicscited by 27
38 more
  1. 2024Transformers on Markov Data: Constant Depth Sufficescited by 26
  2. 2025Metastable Dynamics of Chain-of-Thought Reasoning: Provable Benefits of Search, RL and Distillationcited by 25
  3. 2025The emergence of sparse attention: impact of data distribution and benefits of repetitioncited by 25
  4. 2024Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannotcited by 21
  5. 2026There Will Be a Scientific Theory of Deep Learningcited by 21
  6. 2025How Do LLMs Perform Two-Hop Reasoning in Context?cited by 20
  7. 2024Toward Understanding In-context vs. In-weight Learningcited by 19
  8. 2024Compact Proofs of Model Performance via Mechanistic Interpretabilitycited by 17
  9. 2024Implicit Regularization of Gradient Flow on One-Layer Softmax Attentioncited by 16
  10. 2025Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thoughtcited by 16
  11. 2024Understanding the Language Model to Solve the Symbolic Multi-Step Reasoning Problem from the Perspective of Buffer Mechanismcited by 15
  12. 2024Attention layers provably solve single-location regressioncited by 15
  13. 2024On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixturescited by 15
  14. 2025Learning In-context n-grams with Transformers: Sub-n-grams Are Near-stationary Pointscited by 15
  15. 2024Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformerscited by 14
  16. 2024Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysiscited by 14
  17. 2025Why and How LLMs Hallucinate: Connecting the Dots with Subsequence Associationscited by 14
  18. 2024How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regressioncited by 13
  19. 2024Non-asymptotic Convergence of Training Transformers for Next-token Predictioncited by 13
  20. 2025Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalizationcited by 12
  21. 2025How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Biascited by 10
  22. 2025Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformerscited by 10
  23. 2024Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformerscited by 9
  24. 2025From Markov to Laplace: How Mamba In-Context Learns Markov Chainscited by 9
  25. 2025From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamicscited by 9
  26. 2024Transformers Handle Endogeneity in In-Context Linear Regressioncited by 8
  27. 2024Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysiscited by 8
  28. 2025Understanding In-context Learning of Addition via Activation Subspacescited by 8
  29. 2025Quantitative Bounds for Length Generalization in Transformerscited by 8
  30. 2024How Transformers Get Rich: Approximation and Dynamics Analysiscited by 7
  31. 2025On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learningcited by 7
  32. 2025Selective Induction Heads: How Transformers Select Causal Structures In Contextcited by 7
  33. 2024On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recoverycited by 6
  34. 2025From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformerscited by 6
  35. 2024On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capabilitycited by 5
  36. 2024Explain the Black Box for the Sake of Science: the Scientific Method in the Era of Generative Artificial Intelligencecited by 5
  37. 2024On the Power of Convolution Augmented Transformercited by 5
  38. 2024In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learningcited by 5

50 most cited of 114 All 114 (JSON)