Research map / Probabilistic and scientific modeling / Causal discovery
Teaching Transformers Causal Reasoning through Axiomatic Training
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.
Trains a 67M GPT-2-style transformer from scratch on synthetic axiomatic demonstrations of transitivity and d-separation, then fine-tunes Llama-3-8B-Instruct on the same data to improve causal benchmarks (Corr2Cause, CLEAR).
Problem
Whether causal reasoning (transitivity axiom and d-separation rule) can be learned by language models directly from symbolic demonstrations of causal axioms, and whether this generalizes to more complex causal graphs and to existing causal reasoning benchmarks.
Approach
Axiomatic training: generate large synthetic <premise, hypothesis, conclusion> tuples from simple chain-like causal graphs and train models (from scratch or via fine-tuning) to apply causal axioms, relying on data diversity (random edge flipping) and positional encoding choice for out-of-distribution generalization.
Key results
- On Corr2Cause, axiomatically fine-tuned Llama-3-8B (transitivity) reaches F1 0.32 vs 0.11 for base Llama-3-8B-Instruct and 0.29 for GPT-4 (zero-shot).
- On CLEAR d-separation, fine-tuned Llama-3-8B improves Yes/No accuracy from 60.0 to 70.0 and multi-choice accuracy from 33.0 to 50.0, above GPT-4 (63.33 and 36.67).
- The 67M RoPE/NoPE transformers trained from scratch generalize to chains of 7-15 nodes, with TS2 NoPE reaching 0.94 on reversed 6-node chains versus 0.71 for Gemini Pro.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| fine-tune Llama-3-8B-Instruct on d-separation axiomatic data | Llama-3-8B-Instruct · CLEAR D-Separation YN | accuracy: 60.0 -> 70.0 | Helped |
| fine-tune Llama-3-8B-Instruct on d-separation axiomatic data | Llama-3-8B-Instruct · CLEAR D-Separation MC | accuracy: 33.0 -> 50.0 | Helped |
| fine-tune Llama-3-8B-Instruct on transitivity axiomatic data | Llama-3-8B-Instruct · Corr2Cause | F1: 0.11 -> 0.32 | Helped |
| add random edge flipping to sequential-chain training data (TS1/TS2 instead of OCC) (ablation) | 67M transformer trained for transitivity · Transitivity, reversed chains | accuracy: OCC w NoPE: 0.41, 0.24, 0.18, 0.13; OCC w RoPE: 0.59, 0.26, 0.22, 0.20 | Helped |
| increase the proportion of sequential chains and reduce random-flipping examples (TS2 versus TS1) (ablation) | 67M transformer trained for transitivity · Transitivity, generalization across evaluation setups | accuracy | Helped |
| use sinusoidal or learnable positional encoding instead of no positional encoding (ablation) | 67M transformer trained for transitivity · Transitivity, longer node names | accuracy | Hurt |
| use sinusoidal positional encoding instead of rotary positional encoding (ablation) | 67M transformer trained for transitivity · Transitivity, branching and order-based settings | accuracy | Hurt |
| use rotary positional encoding rather than NoPE, LPE, or SPE (ablation) | 67M transformer trained for transitivity · Transitivity, branching factor 1.4 | accuracy: TS2 w RoPE: 0.88; TS2 w LPE: 0.80; TS2 w SPE: 0.79; TS2 w NoPE: 0.82 | Helped |
| use rotary positional encoding instead of NoPE (ablation) | 67M transformer trained for d-separation · d-separation, branching factor 1.4 | accuracy: 0.70, 0.58, 0.54, 0.52 | Mixed |
Surprising findings
- A 67M model trained from scratch outperforms or matches billion-scale models such as GPT-4 on several length and order generalization settings.
- Training only on pure causal chains (OCC) generalizes to longer chains but fails on random-flip and branched graphs, showing data diversity is needed.
Limitations
- Evaluation assumes no unobserved confounders and only empty or limited conditioning sets for transitivity.
- Performance degrades on longer chains (14-15 nodes) and on dense branched graphs for d-separation, with RoPE model approaching random.
- Benchmark gains are reported on a small CLEAR subset and Corr2Cause zero-shot with high label imbalance.
Experimental setup
Datasets: synthetic transitivity demonstrations (chains of 3-6 nodes with random flipping) synthetic d-separation demonstrations (3-6 nodes)
Benchmarks: Corr2Cause CLEAR synthetic transitivity evaluation (MultiEval_SLR, reversed chains, branched ER graphs) synthetic d-separation evaluation
Metrics: accuracy F1
Baselines: GPT-4 Gemini Pro Phi-3 Llama-3-8B-Instruct
Builds on: Llama-3-8B-Instruct GPT-2 Corr2Cause CLEAR
Closest papers by meaning
- Unveiling Causal Reasoning in Large Language Models: Reality or Mirage? (NeurIPS 2024 · similarity 0.87)
- Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data? (ICLR 2025 · similarity 0.87)
- Disentangled Representation Learning in Non-Markovian Causal Systems (NeurIPS 2024 · similarity 0.86)
- Robust agents learn causal world models (ICLR 2024 · similarity 0.86)
- Language Agents Meet Causality -- Bridging LLMs and Causal World Models (ICLR 2025 · similarity 0.86)
- Causal Action Influence Aware Counterfactual Data Augmentation (ICML 2024 · similarity 0.86)
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.