{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/how-transformers-learn-causal-structure-with","title":"How Transformers Learn Causal Structure with Gradient Descent","arxiv_id":"2402.14735","date":"2024-02-22","proceeding":null,"authors":["Eshaan Nichani","Alex Damian","Jason D. Lee"],"abstract":"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.","url_abs":"https://arxiv.org/abs/2402.14735v2","url_pdf":"https://arxiv.org/pdf/2402.14735v2.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"how-transformers-learn-causal-structure-with","repo_url":"https://github.com/eshnich/transformers-learn-causal-structure","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok"}}],"tasks":[{"task_slug":"in-context-learning","task_name":"In-Context Learning"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2402.14735","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2402.14735"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/eshnich/transformers-learn-causal-structure","reach":{"status":"ok"}}],"summary":{"ran":4},"by_repo_kind":{"official":{"samples":4,"ran":4,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":4,"samples":[{"code_sha256_prefix":"cf90fc4b68f425dc","entry":"get_stationary","repo":"eshnich/transformers-learn-causal-structure","repo_kind":"official","path":"problems.py","file_url":"https://github.com/eshnich/transformers-learn-causal-structure/blob/HEAD/problems.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"cf90fc4b68f425dc"}},{"code_sha256_prefix":"72cb679283a62f39","entry":"plot_A","repo":"eshnich/transformers-learn-causal-structure","repo_kind":"official","path":"plots.py","file_url":"https://github.com/eshnich/transformers-learn-causal-structure/blob/HEAD/plots.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"72cb679283a62f39"}},{"code_sha256_prefix":"ac819defc5e7d02b","entry":"plot_A1","repo":"eshnich/transformers-learn-causal-structure","repo_kind":"official","path":"plots.py","file_url":"https://github.com/eshnich/transformers-learn-causal-structure/blob/HEAD/plots.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"ac819defc5e7d02b"}},{"code_sha256_prefix":"aa51a3b782a251f4","entry":"plot_losses","repo":"eshnich/transformers-learn-causal-structure","repo_kind":"official","path":"plots.py","file_url":"https://github.com/eshnich/transformers-learn-causal-structure/blob/HEAD/plots.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"aa51a3b782a251f4"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}