Teaching Transformers Causal Reasoning through Axiomatic Training
Authors: Aniket Vashishtha
Abhinav Kumar
Atharva Pandey
Abbavaram Gowtham Reddy
Kabir Ahuja
Vineeth N. Balasubramanian
Amit Sharma
Venue: ICML 2025 poster (2025)
Keywords: Causality
Axioms
Transformers
Reasoning
Generalization
Primary area: deep_learning->large_language_models

Abstract: For text-based AI systems to interact in the real world, causal reasoning is an essential skill. Since interventional data is costly to generate, we study to what extent an agent can learn  causal reasoning from passive data. Specifically, we consider an axiomatic training setup where an agent learns from multiple demonstrations of a causal axiom (or rule), rather than incorporating the axiom as an inductive bias or inferring it from data values. A key question is whether the agent would learn to generalize from the axiom demonstrations to new scenarios. For example, if a transformer model is trained on demonstrations of the causal transitivity axiom over small graphs, would it generalize to applying the transitivity axiom over large graphs? 
Our results, based on a novel axiomatic training scheme, indicate that such generalization is possible. 
 We consider the task of inferring whether a variable causes another variable, given a causal graph structure. We find that a 67 million parameter transformer model, when trained on linear causal chains (along with some noisy variations) can generalize well to new kinds of graphs, including longer causal chains, causal chains with reversed order, and graphs with branching; even when it is not explicitly trained for such settings. Our model performs at par (or even better) than many larger language models  such as GPT-4, Gemini Pro, and Phi-3. Overall, our axiomatic training framework provides a new paradigm of learning causal reasoning from passive data that can be used to learn arbitrary axioms, as long as sufficient demonstrations can be generated.

Teaching Transformers Causal Reasoning Through Axiomatic Training
Aniket Vashishtha' Abhinav Kumar? Atharva Pandey? 3 Abbavaram Gowtham Reddy 4 Kabir Ahuja
Vineeth N Balasubramanian 6
Amit Sharma
Abstract
1. Introduction
Causal reasoning can be defined as a set of reasoning proce- dures consistent with pre-defined axioms or rules that are specific to causality (Galles Pearl, 1997). For instance, under stable causal models, the transitivity axiom ("if A causes B and B causes C. then A causes C") helps answer questions of cause and effect between pairs of variables in a system. Similarly, the d-separation rule (Pearl, 2009) connects independence of variables and their causal graph structure, and forms the basis of many graph discovery and effect identification algorithms. Given a causal task and data observations from a system, axioms or rules are typically incorporated as inductive biases in machine learning (ML. algorithm, through regularization, model architecture, the choice of variables. Depending on the kind of avail- able data observational, interventional, or counterfactual Pearl's ladder of causation (Bareinhoim et al., 2022) defines the kinds of causal reasoning that is possible.
For text-based AI systems to interact in the real world, causal reasoning is an essential skill. Since active interventions are costly, we study to what extent a system can learn causal reason- ing from symbolic demonstrations of causal ax- ioms. Specifically, we present axiomatic train- ing method where the system Iearns from multi- ple demonstrations of a causal axiom (or rule), rather than incorporating the axiom as an induc- tive bias or inferring it from data values. A key question is whether the system would learn to gen- eralize from the axiom demonstrations more complex scenarios. Our results, based applying axiomatic training to learn the transitivity axiom and d-separation rule, indicate that such general- ization is possible. To avoid data contamination issues, we start with a 67 million parameter trans- former model and train it from scratch. On both tasks, we find that a model trained on linear causal chains (along with some noisy variations) can gen- eralize well to complex graphs, including longer causal chains, causal chains with reversed order, and graphs with branching. To handle diverse text inputs, the same method is extended to fine- tune language models. Finetuning 1 8B model on our axiomatic data leads to significant gains on causal benchmarks such as Corr2Cause and CLEAR, in cases providing state-of-the- art performance surpassing GPT-4.
As axioms are the building blocks of causality, we study whether it is possible to directly learn the axioms or rules using ML models. That is, rather than learning from dataset sampled from a data-generating process that obeys causal axioms, what if a model can learn an axiom (and thus causal reasoning) directly from symbolic demonstrations of the axiom? This question gains relevance as language models make it possible to learn over symbolic data ex- pressed in natural language. In fact, recent studies have evaluated causal reasoning capabilities of large language models (LLMs) by encoding causal reasoning problems in natural language (Kıcıman et al., 2023; Jin et al., 2023; 2024a). We study whether directly teaching the axioms can a viable way to improve causal reasoning of language models.
'oor primarily done as a Research Fellow at Microsoft Re- scarch India, with additional contributions made at UTUC, USA "Msssahuusetss Institute of Technology, USA 'iirooooft Research, India *CISPA Helmholtz Center for Information Security Ger- many. Part of the work done at ITT Hyderabad University of Washington, USA WWok primarily done at IIT Hyderabad, with additional contributions made MSR India. Correspondence to: Aniket Vashishtha <aniketv2 @ illinois ilinnois.eduu, Amit Sharma <amshar @ comm.
Specifically, we propose new way to learn causal rea- soning through axiomatic training. We posit that causal axioms or rules can be expressed as the following symbolic tuple, (premise, hypothesis, conclusion) where hypothesis refers to a causal claim and premise refers to any relevant information to decide whether the claim is true or not (con- clusion). The conclusion could simply be Yes or No. For example, consider the task of inferring causal relationships from correlational statements in the Corr2Cause dataset (Jin
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Evaluation Set
Longer clase
Training Set
Transformer (Trained scralch)
Sequential chains
Roveraed Cholte
Lenger Nods rarwa
Corr2Cause Benchmark
Prensee
Sapate assiaiiall rekaboro anonn thrwe 2 rataaiia A comelatea int
ajobev of nanalisa dudd Ala
Liama 3-Ob- Instruct
Hypodhoalss A alresdly affecte
itando Plppiin
O
CLEAR Benchmark
and diaecled E-vn)) plary cood, goret Zoo..
(Finstuned)
Htypothest: "Whtch the nellwing nodeaset can dissparsts node and rode
V?
Answer
Figure I: Axiomatic Training for imparting causal reasoning to language models. Given an axiom, we construct a training dataset comprising <premise, hypothesis, conclusion> triplets based on simple chain-like graphs of 3-6 nodes. A transformer model trained from scratch on such instances generalizes to much more complex graphs, including longer causal chains with >6 nodes, branched networks with higher average in-degree and out-degree, complete reversals, shuffled statements, and longer node names. Moreover, when an existing model such as LLama-3-8h model is trained on the same dataset, it leads to significant (up to 20 p.p.) improvement in accuracy on causal reasoning benchmarks such as CLEAR and Corr2Cause.
et al., 2024a), which we empirically study in this paper. The premise can be statements about statistical (in dependence: Prreisse T causes TI. 72 causes T. 72 causes T2 causes T3. causes T. causes TI. causes T3"; the hypothesis can be question about cause-and-effect, "Are T3 and T d-separated given (t, ?2) and the conclusion would be "Yes". This tuple is a demonstration of the d-separation rule (see Section 3 for definition). Based on this template, our key insight is that large number of synthetic tuples can be generated, e.g., by changing the variable names, changing the number of variables, changing the order, and so on. The question is: if a model is trained on such data, would it learn to apply the axiom to new, more complex scenarios?
reversed, generalizes well across all kinds of evaluation sce- narios. However, a model trained only on simple directed chains fails to generalize to graphs with branching or edge direction perturbations. Moreover, for d-separation, our 67 million parameter model outperforms billion-scall models such as GPT-4 under both zero-shot and multi-shot settings. Extending the findings on positional embedding for length generalization in NLP tasks (Kazemnejad et al., 2023; Bhat- tamishra et al., 2020; Haviv et al., 2022), we find that rotary positional encoding works the best for causal generalization, closely followed by no positional encoding.
Next, we study whether the same axiomatic training dataset can also help to improve causal reasoning of pre- trained large language models. We fine-tune Llama-3-8B. Instruct model over axiomatic datasets for transitivity and d-separation. We evaluate on two benchmarks of causal rea- soning in natural language: CLEAR (Chen et al., 2024) and Corr2Cause (Jin al., 2024a). The CLEAR dataset tests language models' understanding of causal graphs through wide range of binary and multi-choice tasks, including d-separation, Markov equivalence, and backdoor adjust- ment. Even though we don't finetune on their dataset or question type, we find significant performance gains on the d-separation task after axiomatic fine-tuning: accuracy in- creases from 30 to 70% on Yes/No task, and goes from 33 to 50% Multi- Choice Questions.
To answer this question, we consider setup where a model is trained on axiomatic demonstrations over simple chain- like graphs with 3-6 nodes and evaluated on more complex graphs, including longer chain- graphs with 7-15 nodes, graphs with branching, longer variable names, and edge direction perturbations (see Figure 1). To avoid any contam- ination concerns with the pre-training data of an existing language model, we first train a transformer model from scratch. For both transitivity and d-separation, we find that a model trained on axiomatic demonstrations learns to apply the axiom multiple times to answer questions over more complex graphs. In particular, diversity in the training data plays crucial role in enabling this generalization. For transitivity, a model trained on a combined dataset of sim- ple directed chains and chains with some edges randomly
Corr2Cause is another challenging dataset to assess lan-
2

---

Teaching Transformers Causal Reasoning through Axiomatic Training
guage model's ability to infer causal relationships (e-g., ancestor, collider, fork) given correlational statements expressed in natural language. Solving this task re- quires knowledge of multiple causal concepts including d-separation, transitivity, Markov property and the Markov equivalence class. Jin et al. (2024a) show how frontier models such as GPT-4 also fail at this task. We show that axiomatic finetuning leads to a substantial F1 improvement of nearly 20 percentage points over the base Llama3 model and even surpasses GPT-4 highlighting the effectiveness of axiomatic training for causal reasoning on complex datasets.
encoding in length generalization capability of transformers. To understand how transformers can be optimized for leam- ing through axiomatic training and generalizing to unseen larger causal structures, we also examine different types of positional encoding such as no positional encoding (PE), Learnable PEs (Radford et al., 2018) and Sinusoidal PEs (Vaswani et al., 2017).
Synthetic data generation for teaching transformers rea- soning: Synthetic data generation has been explored for optimising model training for reasoning. For example, (Li et al., 2023; Gunasekar et al., 2023) use LLM-generated syn- thetic text for training Phi-1 and Phi-1.5 models and show impressive performance for reasoning-basee tasks. (Trinh et al., 2024) introduce a novel neuro-symbolis framework to pre-train a transformer model for Olympiad-level math prob- lems. (Morishita et al., 2024) construct synthetic training data to improve language models' performance on logical reasoning tasks. Building on this stream of work, we apply synthetic data generation for teaching causal reasoning.
Overall, our work provides a new paradigm of teach- ing models causal reasoning through symbolic demon- strations of axioms or rules, which we call axiomatic training. Such symbolic data can be cheaply generated for multiple axioms and added to the finetuning data of language models. More generally, our results using ax- iomatic data contribute to the literature on causal learn- ing from passive data (Lampinen et al., 2023). Our code repository can be accessed at: https: / ///g/tttuuus com/ Aniket Vashi shtha /Cassal _Axiiox.s...
3. Preliminaries: Causal Axioms and Rules
2. Related Work
Instead of performing causal reasoning using observational or interventional data, we study whether it is possible to learn general rules of causality directly from symbolic ax- ioms. There has been fundamental work from Galles & Pearl (1997) where they axiomatize causal relevance (or equivalently irrelevance). They show that for given stable probabilistic causal model (defined below), there exists finite set of axioms that are completely characterized by ax- ioms of path interception in corresponding directed graphs. Additionally, causal inference in practice depends on a few key rules, such as d-separation and do-calculus rules. Learn- ing these rules can have a tangible impact on practical causal tasks such graph discovery and effect inference. While we call the method axiomatic training, we consider learning both causal axioms and rules. Throughout this work, we assume no unobserved confounders.
LLMs for Knowledge-Driven Causal Reasoning: Recent developments in Large Language Models (LLMs) have high- lighted their potential for knowledge-driven causal discov- ery. Unlike traditional methods which focus on statistical patterns or correlations, LLMs utilize knowledge acquired through their pretraining reason about and identify causal structures based on metadata of variables (Kıcıman al., 2023; Ban et al., 2023; Long et al., 2023; Willig et al., 2022: Vashishtha et al., 2025). However, possibility of memoriza- tion of existing benchmarks in the pretraining of these LLMSs has been a major criticism. As a result, recent work (Zeče- vić et al.) argues that L.MS are not actually performing causal reasoning, but simply learning correlations about causal facts. In addition, there are critical failure modes of using LLMs for causal discovery due to hallucinations or not obeying the acyclic constraint when generating graph edges (Vashishtha et al., 2025). To evaluate causal reason- ing capabilities of LLMs, (Jin et al., 2024a) and (Jin et al., 2023) propose formal causal inference evaluation bench- marks infer direct and indirect causal relationships, and highlight the failure of LLMs in performing accurate causal reasoning.
Notation. We denote a random variable with an uppercase letter (e.g., X, Y, Z) and use lowercase letters (e.g., 2; y, 2) to denote the values taken by the corresponding random variable, written as X = x., y, Z = We represent the probability of a random variable X, as P(X.). Let G(X, E) be a directed acyelic graph (DAG) consisting of a set of variables X- = ['....., } and a set of directed edges E among variables in X. Let (X,) x = [X - X,), de().)= = (X&|Xr + < X,), ch(XX)) = (X&|X; + X denote the set of parents, descendants and children of X; respectively. Given two nodes X,, X, we call third node X collider if both X, and X₃ are parents of X
Impact of Positional Encoding on Generalization: Length generalization capabilities of transformers has been stud- ied in the past to better understand their different failure modes across various settings (Hupkes et al., 2020; Zhang et al., 2023; Furrer et al., 2021). Previous work (Kazem- nejad et al., 2023; Bhattamishra et al., 2020; Haviv et al., 2022; Shen et al., 2023) emphasizes the impact of positional
3

---

Teaching Transformers Causal Reasoning through Axiomatic Training
3.1. Axioms of Causality: Transitivity
Definition 3.3 (Definition 1.2.3 in Pearl (2009)). Given a DAG ( E), two sets of random variables X; and X; are said to be d-separated by third set X, if all the paths between X; and X, in G are blocked by X. A path p between X; and X is said to be blocked by a set of nodes X; iff 1) p contains a fork (i.c., : + A + -) or a chain (i.c., A ++) .) such that the middle node A is X.. or 2) p contains a collider (. (+ A <.) )) such that the middle node A is not in X, and no descendant of A is in X..
Definition 3.1 (Causal Irrelevance, adapted from Defi. 7 in (Galles & Pearl, 1997)). X is probabilistically causally irrelevant to Y given Z, written (X Z iff: P(ylz, do) = x) P(ul2, (X) x'), Vx, x', y, i.e., once we hold Z fixed at z, intervening on X will not change the probability of Y.
Next, we restate the stability assumption for a causal model from (Galles& & Pearl, 1997) that gives a richer set of finite axiomatization for probabilistic causal irrelevance.
Given P(X) is Markov with respect to g, if two sets of random variable X, and X; are d-separated by Xz, then they are conditionally independent of each other given X..
Assumption 3.2 (Stability, Definition 9 in (Galles & Pearl, 1997)). Let M be a causal model. Then an irrelevance (X + 2| in M is stable if it is shared by all possible probability distribution over M.. The causal model M is stable if all of the irrelevances in M are stable.
4. Axiomatic Training for Transformers
Given an axiom, our key idea is to generate thousands of synthetic symbolic expressions that can be used to train a transformer on how to use the axiom. The trained model is then evaluated on whether it can apply these axioms to new causal structures that were not available in the training set. Below we describe how we generate the training data and the model architecture details.
Under the stability assumption (see Assumption 3.2), Galles & Pearl (1997) states six axioms that completely character- ize causal irrelevance (Definition 3.1) via axioms of path interception in the directed graphs. An axiom of causal ir- relevance is of the form (given in conjunctive normal form):
AV(X X
AV
x
1"
La
4.1. Training Data: Diversity is key
As mentioned above, an axiom consists of tuple, (premise, hypothesis, conclusion). Based on the specific axiom, we can map hypothesis given the premise to its correct label ('Yes' or 'No'). To create a training dataset, we randomly sample a causal DAG G and enumerate N random tuples of {(P, H, L)}N where P is the premise, H is the hypothesis and L is the label (Yes/No). The premise describes the edges of the graph and is expressed in natural language, e.g., "X causes Y. Z causes Y.". Given a premise P based on the causal graph's 's edges, if the hypothesis can be derived by applying the specified axiom (once or multiple times), then label I. is Yes; otherwise, No. For example, for the transitivity axiom, suppose the underlying true causal graph of a system is a chain, X1 - X2 + X3- Then, the premise will be X1 X2 X2 + X3- A corresponding hypothesis for the transitivity axiom could be Xi X3 will have label Yes whereas another hypothesis X3 - Xi will have label No. The former would create training data instance with the following text, "X1 causes X2. X2 causes X3. Does X1 cause X3? Yes.". Note that the axiom can be inductively applied multiple times to generate more complex training tuples. Another possible hypothesis for the d-separation rule could be "Are X1 and X₂ d-separated given (Xa)?" ) and the label will be No.
where A is "logical and", V is "logical or" and for given (s,t) or n) pair, X4, Xx Xx are disjoint subsets of ob- served variables X. In the above causal irrelevance state- ment, if the antecedent is true, the consequent is also true.
Transitivity Axiom. We illustrate our axiomatic training procedure through the transitivity axiom. Following the stability assumption above, we consider the class of in- terventional distributions in which the transitivity causal irrelevance axiom holds (Sadeghi & Soo, 2024). Formally, for a stable probabilistic causal model ($3), given variables x, Y, Z in the system, the transitivity axiom is:
(X + YIZ) = (A + Y|Z) V X + A|Z)
AA $ XUZUY
which can be simplified using the contrapositive. We call the LHS as Premise and the RHS as Hypothesis.
AE XUYUZ s.t. (X + A(2) A (X)
P:premise
(X + XI2)
H: hyppthssssss
We train the model on data from simple causal graphs such as sequential chains with 3-6 nodes and evaluate its perfor- mance on more complex graphs (refer fig. 1). To enhance generalization, we introduce structured perturbations in the
3.2. d-separation rule
The d-separation rule connects causal graph structure with conditional independence in TP( XX)

---

Teaching Transformers Causal Reasoning through Axiomatic Training
training data across three axes: node names, causal structure types, and the number of nodes in the causal graph.
et al., 2018), sinusoidal (SPE) (Vaswani al., 2017), rotary (RoPE) position encodings (Su et al., 2024), and no posi- tional encoding (NoPE) (Kazemnejad et al., 2023; Haviv et al., 2022). See Appendix F for details.
1. Node names: an alphanumeric name comprising 1-3 characters. The length of a name and the specific characters are randomly selected during data generation.
Each node in the graph is represented by
Finetuning. Apart from training a transformer from scratch, we fine-tune a pre-trained language model (Llama-3-8b- Instruct (gra, 2024)) on our axiomatic training data. Fine- tuning a pretrained model like Llama, allows us to evaluate causal reasoning in wider variety of contexts, leverag- ing its expansive vocabulary and the versatility provided by pretraining. While a model trained from scratch on axiomatic data is limited by its narrow vocabulary and struc- ture, finetuning a pretrained model allows direct evaluation on benchmarks like Corr2Cause, which implicitly require the application of the target axioms, thus enabling us to stress test framework in more diverse settings. Refer J for more details about the finetuning setup.
2. Causal Graph Topology: We consider two main types of causal graphs in the training set.
(a) Sequential: All causal edges are directed forward, forming typical chain DAG, e.g., X + Y + Z. (b) Random Flipping: Given chain of sequential nodes, we randomly reverse some edges eg. X ++ Y f Z. This can be expressed simply through natural language like: "X causes Y. Z causes Y.". This intro- duces forks and colliders that help add complexity to model training, thus aiding generalization across larger space of graphs.
3. Number of nodes in graph: To facilitate the general- ization of transformers over graphs of different sizes we incorporate chains of varying lengths, ranging from 3 to nodes in our training set.
4.3. Evaluation setup: Assessing Axiomatic Learning
We consider two types of evaluation: 1) on synthetic datasets where we directly test the models on axioms and, 2) on existing benchmarks corresponding to different high-level causal tasks where we expect the axioms to be helpful.
4.2. Tokenization, Training Loss & Architecture
We train a decoder-based 67 million parameter model based on GPT-2's architecture. The model has 12 attention layers, 8 attention heads and 512 embedding dimensions. We train the transformer model from scratch to ensure that the model has not seen such axioms in the pertaining step and thus requires true correct understanding of axioms to perform well. Later we also tested on a pre-trained model which was fine-tuned on our dataset.
Synthetic evaluation. To evaluate if a trained model has learned the correct understanding of an axiom instead of shortcuts or correlation- features, designing an out-of- distribution (OOD) evaluation set is important. We evaluate our model on multiple types of complex graphs that are unseen during training.
1. Length:
Evaluating whether our model accurately infers causal relationships for chain graphs (both sequential and ones with random flipping) longer than those in the training set. Specifically, we train the model on chains with size 3-6 and evaluate on chains of size 7-15.
Since the training dataset follows a specific structure, we develop a custom tokenizer. Alphanumeric node names are tokenized at a character level, while special terms such as 'causes', 'Does', 'cause", 'Yes', and 'No' are tokenized at the word level. Such an approach avoids out-of- vocabulary (OOV) tokens at test time since the alphanumeric node names in the test set can be different than those in the training set. Following this approach, the vocabulary size of our transformer model is 69.
Tokenization.
2. Node Name Shift: Testing the model's performance on longer node names, from 1-3 characters in the training set to 8-10 characters, motivated by Jin et al. (2024a) who show how change in node names can lead to gener- alization failure causal tasks for language models.
3. Order of Chains: a) Completely reversed chains: This evaluation is inspired by the reversal curse (Berglund et al.. 2024) that revealed generalization failure of LLMs in answering questions in reversed sequences despite knowing the answers in the original order. We evaluate our model on completely reversed chains, a structure that was not in the training data. A completely reversed chain will be of the form X <- Y + Z, written in natural lan- guage as: "Y causes X. 7. causes Y.", where X, Y. Z are replaced by random alphanumeric names. b) Shuffling of Sequences: Here we shuffle the order of causal edges presented in each training row to add complexity and break sequential order.
Loss function. Given dataset, the loss function is defined based on the ground truth label for each tuple, represented as E - log IP PP, H). A preliminary analysis indi- (P,.H.)>>PPann
cated promising results with this loss formulation compared to next token prediction loss.
Positional Encoding. In addition to tokenization and loss function, recent work (Kazemnejad et al., 2023) shows that the choice of positional encoding is important for general- izing a transformer to longer or complex inputs. Therefore, we evaluate different positional encodings on their impact on generalization causal tasks: learnable (LPF) (Radford
5

---

Teaching Transformers Causal Reasoning through Axiomatic Training
4. Branching Factor: One of the most complex evaluation setups, with DAGS containing multiple branches, collid- ers, and forks. Let the branching factor be defined as the ratio between a number of edges and number of nodes. Thus, the branching factor for the training set is I 1. Then, we create a different evaluation set with multiple densely branched networks constructed using the Erdös-Rényí model, with different branching factors.
Model/Nodes
3
8
Baselines (Zero Shot)
GPT-4
0.99
0.97
0.89
0.95
0.90
0.90
Gemini Pro Phi-3
0.75
0.73
0.72
0.76
0.71
0.68
0.74
0.88
0.86
0.82
0.79
0.76
0.73
0.79
Buselines (Multi Shof)
GPT-4
1.00
0.99
0.97
0.95
0.94
0.90
0.92
Gemini Pro
0.95
0.85
0.83
0.79
0.79
0.73
0.75
Phi 3
0.88
0.83
0.82
0.80
0.83
0.76
0.78
Axiomatic Training
4.4. Benchmark evaluation
TSI w NoPE
1.00
0.93
0.85
0.83
0.78
0.73
0.73
To test whether such simple axiomatic training is helpful in more complex scenarios, we evaluate our models causal reasoning benchmarks, Corr2Cause (Jin et al., 2024a) and CLEAR (Chen et al., 2024). Unlike past methods as in Jin et al. (2024a), we do not finetune on these datasets directly. Instead, we finetune language model on synthetic axiomatic data and evaluate whether the model generalizes to these external benchmarks. We discuss the results on these benchmarks in Section 7.
TS1 w LPE
1.00
0.93
0.87
0.83
0.79
0.74
0.73
TSI w SPE
0.99
0.92
0.85
0.81
0.76
0.74
0.61
TSI w RoPF
1.0
0.93
0.87
0.85
0.81
0.78
0.76
TS2 w NoPE
0.99
0.93
0.79
0.74
0.73
TS2 w LPE
1.00
0.92
0.85
0.83
0.77
0.74
0.71
TS2 w SPE
0.99
0.94
0.86
0.81
0.76
0.72
0.64
T52 w RoPF
1.0
0.95
0.89
0.86
0.82
0.79
0.76
OCC w RoPE
0.78
0.64
0.63 0.61
0.61
Table 1: Evaluation on MultiEvalgLn dataset. Accu- racy of axiomatically trained models with another baseline on the most complicated setups. For OCC we only report performance with RoPE encodings, which is the best per- forming setup for this dataset. See See 5.2 for details. Bold numbers denote the highest value on test set, while the underlined ones denote the second best.
4.5. Baselines Using Existing LLMs
Given recent work on how LLMs can be leveraged for causal reasoning (Kiciman et al., 2023; Vashishtha et al., 2025; Ban et al., 2023), we include language models such as GPT-4 (gpt-4-32k) (ope, 2024), Gemini (gemini-pro) (gem, 2024) and Phi-3 (Phi-3-mini-1) 128k- (iskiinstucc)) (abd, 2024) as base- lines. Each model is significantly larger than our model and known to perform well on reasoning tasks, with the smallest model Phi-3 having 3.8B parameters (Li al., 2023). Refer App. B and C. for zero and multi-shot prompts used. On evaluating the baselines for the transitivity task, we find that GPT-4 is significantly better. Hence, for the more complex d-separation task, we evaluate only GPT-4.
and 101k causal graphs where the underlying graphs has sequential chains. Since flipping is done randomly across all consecutive pairs of nodes in the given chain, some complete reversals are also formed. In this training set, around 12k completely reversed chains are present. 3. Training Setup 2 (TS2): This setup comprises more sim- ple sequential chains (132k), while we decrease chains with random flipping (42k) to keep the overall size around 175k to see the effect of adding examples with complicated graphs model's performance.
5. Axiomatic Training For Transitivity Axiom
In all our experiments, we consider an empty conditioning set Z for simplicity.
Evaluation Datasets. In addition to the evaluation sets described carlier in Sec 4.3 (length generalization, node name shift, order of chains, and branching), we add another evaluation set that is a combination of three shifts.
5.1. Training and Evaluation Datasets
Training Datasets. The training data consists of sequen- tial chains of lengths from [3,6]. In addition to sequential chains, random flipping of edges is done with 0.5 probabil- ity. See Appendix G for details on these hyperparameters. Our training data consists of 175k axiom demonstrations. We use three versions of training data to evaluate the impact of different data perturbations.
MultiEvalsLR (Shuffling + Random Flipping + Length Sequence): This setup merges three kinds of complexities shuffling of sentence for representing the sequences, each se- quence having random flipping, and some sequences having longer length than sequences in training set (upto 9).
5.2. Results:
Generalizing to Complex Causal Graphs
1. Only Causal Chains (OCC): This set comprises of graphs with only sequential chains (see causal graph topology in Sec 4.1 for details).
We train our model using axiomatic training on different kinds of datasets, TS1, TS2, and OCC, with different po- sitional encoding (NoPE, LPE, and SPE) as described in See 5. Results on all evaluation settings are in Appendix Tables A3, A4 and A5.
2. Training Setup 1 (TS1): This setup comprises of 73k ex- amples where the underlying graphs has random flipping
6

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Model/Nodes
10
12
Generalization Performances on Longer Node Names
iZero Shat)
Baselines (Zero shot)
Accuracy
GPT-4 (Multi Shot)
IS2 No PE
GPT-4
0.95
0.90
0.88
0.86
Gemini Pro
0.74
0.76
0.73
0.71
T52 Sin PE
Phi-3
0.83
0.79
0.77
0.80
Baselines (Multi shot)
Chain Length (Number Nodes)
GPT-4
0.97
0.93
0.94
0.93
Figure 2: Evaluating generalization on causal sequences (with- out random flipping) with longer node names (than the ones used in sequences in train set). TS-2 training set with no positional encoding leads to the best performance. Refer table A5 for complete results.
Gemini Pro
0.76
0.79
0.77
0.79
Phi-3
0.78
0.82
0.79
0.79
Axiomatic Training
OCC w RoPE
0.72
0.68
0.66
0.62
TSI w LPE
0.82
0.72
0.69
0.68
Generalization Performance on Longer Chains with Random Flipping
TSI w SPE
0.78
0.61
0.59
0.57
1o
GPT-4 Shot)
TSI w NoPE
0.82
0.74
0.69
0.66
GPT-4 (Mult: Shot)
TSI w RoPE
0.86
0.79
0.74
0.71
Accuracy
T52 No PE
TS2
TS2 w LPE
0.80
0.72
0.67
0.64
T52 Sin PE
TS2 w SPE
0.79
0,59
0.54
0.52
TS2 w NoPE
0.82
0.73
0.68
0.64
TS2 w RoPE
0.88
0.79
0.72
0.68
Chain Length (Numbee Nodes)
Table Evaluation with branching factor 1.4. Accuracy of axiomatically trained models with LM baselines on the causal graphs with branching factor 1.4.
Figure 3: Generalizing to longer unseen causal sequences (>6 nodes) with random flipping on TS2 and OCC (with NoPE) train sets. OCC-trained models struggle due to limited edge-level variability, while TS2 NoPE consistently performs well. Refer table A4 for complete results
MultiEvalsLR Dataset. We evaluate our models and other baselines on the challenging MultiEvallLı dataset since it includes example that are different from training dataset in terms of size and complexity of causal graph. Ta- ble 1 summarizes the results for this dataset. While GPT-4 performs best, models trained with RoPE position encod- ings still achieve strong results, surpassing Gemini Pro and Phi-3 in both zero-shot and multi-shot settings for majority of node lengths.
curacy as graph sizes increase, TS1 (RoPE) and TS2 (RoPE) model outperforms Gemini Pro (branching factor 1.4) in for graphs with size 5 and 8 under zero-shot set- tings. On graphs with 12 nodes and a 1.4 branching factor, TS2 (RoPE) achieves 68% accuracy, far better than random (50%), despite training only on graphs with branching fac- tors S 1. Although LLMs excel in multi-shot settings, our model's performance is comparable even on more compli- cated causal graphs than the ones they were trained on.
A similar trend is seen for completely reversed sequences (Table A3). This task presents extreme out-of-distribution data, as the training data contains left-to-right edges, while the test data has only right-to-left edges. TS2 (NoPE) con- sistently outperforms Gemini-Pro and Phi-3, and remains competitive with GPT-4 (zero shot). In particular, its curacy (0.94 for chains of length 6) is substantially higher than Gemini Pro and Phi-3 (0.71 and 0.75 respectively).
Further Ablations. We run further ablations to understand the generalization behavior of our models. In particular, we experiment with generalization to unseen node names (Appendix Table A5), generalization to unseen lengths for graphs with linear chains and random flipping of edges (Appendix Table A4), and generalization to graphs with branching factor of 2 (Tab. A6).
Branched Causal Graphs. Even though our models are trained on simpler graphs like sequential and random- flipping, we want to test our models on structurally harder graphs not considered in MultiEvalsLR. To do so, we in- troduce general Erdos-Rent Renyi graphs as the causal sequences while the training data contains only linear chains. We vary the branching factor of the graph as defined in Sec 4.3 be- tween 1.4 and for all our experiments on graphs with different numbers of nodes. Table summarizes the results of this experiment. While GPT-4 achieves the highest ac-
6. Axiomatic Training for D-Separation Rule
Similar to the transitivity axiom, we now train our model on instances of the d-separation rule.
6.1. Training Dataset and Setup
We follow a similar strategy to generate the training dataset as for the transitivity axiom. The training dataset consists
7

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Branching (Bfactor = 1.4)
7. Axiomatic finetuning for complex tasks
Models
5
8
10
12
Baselines
We now focus on causal reasoning benchmarks such as Corr2Cause and CLEAR and show how axiomatic finetun- ing can help improve causal reasoning of pretrained lan- guage models. We finetune -lama-3.1-8B-Insttee supervised fine-tuning on our axiomatic training data (de- scribed in Sec. 4.1). Finetuning an existing model avoids any restriction on tokens and language style; the axiomat- ically finetuned model can be applied to any general text. We use 1. -n8--rutttutt as the baseline to assess the improvement through axiomatic finetuning. We also in- clude GPT-4 baseline, given its demonstrated strength in reasoning tasks.
GPT-4
0.53
0.544
0.62
0.52
Finetumed Results
-lama-3-8h-Instmun
0.474
0.490
0.470
0.482
Finetuned LJama
0.796
0.738
0.718
0.670
using
Models with different PEs
SPE
0.67
0.59
0.56
0.55
LPE
0.67
0.61
0.57
0.56
NoPE RoPE
0.63
0.58
0.53
0.53
0.70
0.58
0.54
0.52
Table 3: Effectiveness of axiomatic training for d-separation under two training paradigms: training model from scratch and fine-tuning pretrained Llama model.
7.1. Evaluation on CLEAR Benchmark
of graphs with nodes from [3. 6] with branching factor in range [0.6, 0.8] and uses the same premises from the TS2 training set (see Sec 5. Given a causal graph premise, we create the hypotheses as follows: First, select all pairs of nodes (2) in the graph, then select the conditioning set C from the remaining sets of nodes of size up nodes. The ground truth labels denote whether x1 and 22 are d- separated given the conditioning set C. From this exhaustive set of hypotheses, we randomly subsample 175k instances.
Dataset and Evaluation Setup. CLEAR (Chen et al., 2024) provides a comprehensive benchmark assessing language models' understanding of causal graphs through 20 diverse tasks, including d-separation. It has different question types like Aulti-Choice and Yes/No type questions (Refer Fig. I for an example instance).
We test our model, fine-tuned on axiomatic d-separation instances, in zero-shot setting on CLEAR's Yes/No (YN) and Multi-Choice (MC) questions. Note that there are multi- ple differences, including in structure and wording, between the finetuning data and the test benchmark. In addition, our finetuning data does not include any MC questions, only YN questions. Our goal is to check whether a axiomatically finetuned Llama model would generalize zero-shot to new kinds of causal questions.
6.2. Evaluation
Unlike transitivity that primarily involves reasoning over linear chains, d-separation is more challenging due to its dependence on various structural patterns, including collid- ers, chains, and forks. Given the difficulty of the task, we evaluate two models: 1) 67M model trained from scratch on the d-separation axiomatic data; 2) Llama-3-8B model finetuned on the same axiomatic data. We consider two evaluation settings: longer sequences with random flipping (refer Table A7) and branching (refer Table 3).
Results. Table 5 shows the results. Our axiomatically finetuned Llama model provides a significant improvement in accuracy over the base model, for both YN and MC questions. Specifically, on MC questions, fine-tuning on ax- iomatic instances of d-separation yields a 1 17% improvement over the base model, even surpassing GPT-4. These strong zero-shot improvements across different question types and graph structures-notabe outperforming larger models like GPT-4, indicate the effectiveness of axiomatic finetuning in enhancing causal reasoning.
Overall, the model trained with RoPE achieves the highest performance among all models trained from scratch, with the NoPE-based model ranking a close second across both complex graph structures. GPT-4, in contrast, struggles to outperform random baselines in both settings. For com- plex linear graphs, our RoPE-trained model significantly exceeds random baseline performance and clearly outper- forms both GPT-4 and -lama-3-8B-- Instructt The axiomati- cally finetuned model performs comparably and maintains its effectiveness even as graph length increases, whereas the -orinnr model's performance declines for longer sequences. For branched causal graphs, although the RoPE- trained model initially performs well, its accuracy quickly drops toward the random baseline as complexity grows. In these more challenging scenarios, the Alama-finetunel model remains consistently strong and surpasses GPT-4 across all node sizes, thus highlighting the impact of ax- iomatic training in this setup.
7.2. Evaluation on Corr2Cause Dataset
Dataset details. Corr2Cause (Jin et al., 2024b) provides a more complex dataset to evaluate models on different causal tasks. Each data instance in the benchmark includes cor- relational relationships described in natural language for graphs with nodes; the goal is to infer the truth value of hypothesis. The hypothesis consists of six different kinds of graphical relationships between pairs of variables: Parent, Ancestor, Child, Descendant, Collider, and Con- founder. This task is significantly harder than applying a
8

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Model
Prec
Rec
Acc
20 p.p. (Table 4). Fine-tuning on both transitivity and d-separation improves performance. Notably, transitivity fine-tuning led to the largest gains, even surpassing GPT-4, likely due to its direct relevance in inferring graph relation- ships required for the hypothesis covered in the Corr2Cause benchmark (such as identifying child, ancestor, parent, de- scendant). The significant gains in Fl scores under an out- of-distribttion setting demonstrate the effectiveness of ax- iomatic finetuning.
Llama-3-88
0.11
0.15
0.08
0.77
Alamm-3-8b-Instruue
0.23
0.16
0.46
0.54
finetuned Dsep (175k)
Alama-3-8b-Instruue
0.32
0.20
0.93
0.39
finetuned Transitivity (175k)
GPT-4 (from (Jin et al., 2024a))
0.29
0.21
0.48
0.64
Table 4: Evaluation on the Corr2Cause Task from (Jin et al., 2024a). We finetune Llama-3-8B our transitivity and d-separation datasets and evaluate on the Corr2Cause dataset in a zero-shot setting. We observe significant increase in performance (over 20 p.p. improvement in Fl
8. Discussion and Conclusion
score compared to the base Llama model).
We introduced a general framework, axiomatic training, to add axioms and simple rules of causality as inductive priors in ML. models and showed its effectiveness in improving performance on downstream causal reasoning tasks. For applicability to real-world tasks, we also showed that it can be used to finetune pre-trained language models. Our study shows the importance of diversity in axiomatic data to enhance generalizability.
CLEAR D-Separation task (YN)
Models
Accuracy
Llama-3 3-8b-Instruct
60.0
-lama-3-Sb-Instrue Finetuned
70.0
GPT-4
63.33
CLEAR D-Separation task (MC)
Models
Accuracy
Llama-3 3-8b-Instruee
33.0
Generalization to Logical Reasoning. While our ax- iomatic training approach focuses on causal reasoning, can be applied any formal system such as deductive logi- cal reasoning. Recent work (Saparov et al., 2023) highlights LLMs' deterioration in performance as reasoning depth in- creases. Given the similarity between our setup and such tasks, it would be interesting to explore if axiomatic training can improve deductive reasoning in language models. For example, transformers may be trained on logical axioms to improve language model's reasoning; Morishita et al. (2024) provide promising evidence in this direction.
Llama-3-8b 3-8---nsruu-I Finetuned
50.0
GPT-4
36.67
Table 5: Evaluation on CLEAR dataset (Chen et al., 2024) We finetune the LMs on our d-separation dataset and evaluate on the CLEAR dataset. We observe a significant increase in performance compared to the baseline, which highlights the efficacy of axiomatic training.
single axiom. First, one needs to infer the causal graph from the correlation statements. This requires knowledge of d-separation and Markov condition (Appendix E) Then, one needs to the transitivity axiom identify the direct and indirect relationships to identify the children, ancestors, colliders, and confounders in the causal graph.
Verifiers for Causal Reasoning. Another application of axiomatic training may be in building verifiers that operate over language models' output. For instance, our axiomat- ically finetuned models can be used to detect violations of causal rules such as transitivity and d-separation a given text, and give inference-time feedback the language model to improve its generation.
It was found that popular LLMs fail to perform well on this benchmark. While finetuning on this dataset shows promise, the finetuned language models also fail in out-of- distribution settings generated by simply perturbing these questions changing node names. As for the CLEAR benchmark, we will test our axiomatically finetuned models exactly in the out-of-distribution setting: the models have only been trained on synthetic axiomatic data (described in Section 4.1) and not seen any instance of Corr2Cause. Specifically, we finetune two models, one based on transi- tivity data and one based on d-separation data, given that the questions involves multiple axioms. Due to the high imbalance between labels in the dataset, we use Fl score to measure performance.
Implications for Training Language Models. Incorporat- ing causal axiom demonstrations into pretraining or finetun- ing can boost reasoning in small language models, enabling models like Phi-3 to reach GPT-4-level accuracy on causal tasks. Papadimitriou & Jurafsky (2023) show that pretrain- ing on synthetic formal languages introduces beneficial in- ductive biases. Similarly, pretraining synthetic axiomatic data may enhance language models" reasoning abilities. An- other promising direction is to use axiomatic training to en- hance base language models before applying post-training with reinforcement learning (RL). Since RL-based align- ment typically amplifies the capabilities of a base model, in- corporating causal understanding through axiomatic training could provide stronger foundation (Gandhi et al., 2025).
Results. Our results highlight Llama-3- -eeaa----b--nnntee 8b-Insttrut Base model's 's poor performance on this task, while axiomatic fine- tnning results in significant performance gains of around

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Impact Statement
Chen, S., Xu, M., Wang, K., Zeng, X., Zhao, R., Zhao, S., and Lu, C. CLEAR: Can language models really understand causal graphs? In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 6247- 6265, Miami, Florida, USA, November 2024, Association for Computational Linguistics. doi: 10. 18653/v 1/2024. findings-emnlpp emnlp.363.
This paper presents work whose goal is to advance the field of Machine Learning. by proposing effective training strate- gies to improve causal reasoning capabilities of language models. Robust causal reasoning has clear upside: decision- support tools in medicine, public policy, and the social sciences may become more trustworthy under distribution shifts and easier for experts to scrutinize. We also recog- nize possible downsides, flawed or adversarially fine-tuned models could generate spurious causal claims that reinforee bias or justify harmful interventions. To mitigate these risks, release code and synthetic training data openly, detail known limitations, and encourage independent replication and auditing. On balance, we believe that equipping lan- guage models with transparent, principled causal reasoning capabilities will yield net benefits for socicty while enabling informed oversight.
Csordás, R., Irie, K., and Schmidhuber, J. The devil is in the detail: Simple tricks improve systematic generalization of transformers. In Moens, M.-F., Huang, X., Specia, L.,, and Yih, S. W.-t. (eds.). Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 619-634, Online and Punta Cana, Dominican Re- public, November 2021. Association for Computational Linguistics. doi: 10. 00.16633v1/202/ 10.1865//112222.cmmllmp
Furrer, D., van Zee, M., Scales, N., and Schärli, N. Compo- sitional generalization in semantic parsing: Pre-training vs. specialized architectures, 2021.
References
Transformers: State-of-the art natural language processing. In Liu, Q. and Schlangen, D. (eds.), Proceedings of the 2020 Conference on Empirical Methods in Natural Lan- guage Processing: System Demonstrations, pp. 38-45, Online, October 2020. Association for Computational Linguistics. doi: 10.18653/11/222.ccomm http://0000mm//dddmdmdoo.om demos.6.
Galles, D. and Pearl, Axioms of causal relevance. Artifi- cial Intelligence, 97(1):9-43, 1997. ISSN 0004-3702. doi: https://oo.lon///000000000000000 vance.
Rele-
Gandhi, K., Chakravarthy, A., Singh, A., Lile, N., and Goodman, N. D. Cognitive behaviors that enable self- improving reasoners, or, four habits of highly effective stars, 2025.
Phi-3 technical report: A highly capable language model locally on your phone, 2024.
Gemini:
A family of highly capable multimodal models, 2024.
Gunasekar, S., Zhang, Y., Aneja, J., Mendes, C. T. Giorno, A. D., Gopi, S., Javaheripi, M., Kauffmann, P., de Rosa, G., Saarikivi, O., Salim, A., Shah, S., Behl, H. S., Wang. X.,, Bubeck, S., Eldan, R., Kalai, A. T., Lee, T.T. and Li, Y. Textbooks are all you need, 2023.
The Ilama 3 herd of models, 2024.
Gpt-4 technical report, 2024.
Haviv, A., Ram, O., Press, O., Izsak, P., and Levy, O. Transformer language models without positional encod- ings still learn positional information. In Goldberg, Y., Kozareva, 7., and Zhang, Y. (eds.), Findings of the As- sociation for Computational Linguistics: EMNLP 2022, pp. 1382-1390, Abu Dhabi, United Arab Emirates, De- cember 2022. Association for Computational Linguistics. doi: 10. 81538661202222 h2w.222..nnnnnnnemmm.......
Ban, T., Chen, L.. Wang, X., and Chen, H. From query tools to causal architects: Harnessing large language models for advanced causal discovery from data, arXiv preprint arXiv:2306. 16902, 2023.
Bareinboim, E., Correa, J., theling, D., and Icard, T. On pearl's hierarchy and the foundations of causal inference (Ist edition). In Geffner, H., Dechter, R., and Halpern, J. (eds.), Probabilistic and Causal Inference: the Works of Judea Pearl, pp. 507-556. ACM Books, 2022.
Hupkes, D., Dankers, V., Mul, M., and Bruni, F. Compo- sitionality decomposed: how do neural networks gener- alise? 2020.
Berglund, L., Tong, M., Kaufmann, M., Balesni, M., Stick- land, A. C., Korbak, T., and Evans, O. The reversal curse: Lims trained on "a is b" fail to learn "b is a", 2024.
Jin, Z., Chen, Y., Leeb, F., Gresele, L.., Kamal, O., Lyu, Z., Blin, K., Adauto, F. G., Kleiman-Weinet M., Sachan, M., and Schölkopf, B. Cladder: Assessing causal reasoning in language models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.). Advances Neural Information Processing Systems 36 (NeurIPS
Bhattamishra, S., Ahuja, K., and Goyal, N. On the abil- ity and limitations of transformers to recognize formal languages, 2020.
10

---

Teaching Transformers Causal Reasoning through Axiomatic Training
2023), volume 36, pp. 31038-31065, New Orleans, USA, Dec 2023. Curran Associates, Inc.
December 2023. Association for Computational Linguis- ties.
Jin, Z.,, Liu, J., Lyu, Z.,, Poff, S., Sachan, M., Mihalcea, R., Diab, M., and Schölkopf, B. Can large language models infer causation from correlation? In Proceed- ings of the Twelfth International Conference on Learning Representations (ICLR 2024), Vienna, Austria, 2024a. International Conference on Learning Representations.
Pearl, J. Causality: Models, Reasoning and Inference. Cam-
bridge University Press, USA, 2nd edition, 2009. ISBN 052189560X.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. 2018
Jin, Z., Liu, J., LYU, Z., Poff, S., Sachan, M., Mihalcea, R., Diab, M. T., and Schölkopf, B. Can large language models infer causation from correlation? In The Twelfth International Conference on Learning Representations, 2024b.
Sadeghi, K. and Soo, T. Axiomatization of interventional probability distributions. Biometrika, pp. asac043, 08 2024. ISSN 1464-3510. doi: 10.109//biomet/sscc43.
Saparov, A., Pang, R. Y., Padmakumar, V., Joshi, N., Kazemi, S. M., Kim, N., and He, H. Testing the general deductive reasoning capacity of large language models using ood examples, 2023.
Kazemnejad, A., Padhi, I., Natesan Ramamurthy, K., Das, P., and Reddy, S. The impact of positional encoding on length generalization in transformers. In Oh, A., Nau- mann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Pro- cessing Systems, volume 36, pp. 24892-24928. Curran Associates, Inc., 2023.
Shen, R., Bubeck, S., Eldan, R., Lee, Y. T., Li, Y., and Zhang, Y. Positional description matters for transformers arithmetic, 2023.
Su, J., Ahmed, M., Lu, Y., Pan. S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomput, 568(C), February 2024. ISSN 2312.. doi:
KIeıman, E., Ness, R., Sharma, A., and Tan, C. Causal reasoning and large language models: Opening a new frontier for causality. arXiv preprint arX1v:2305.00050, 2023.
Trinh, T. H., Wu, Y., Le, Q. V., He, H., and Luong, T. Solving olympiad geometry without human demonstra- tions. Nature, 625(7995):476-482 2024. doi: 10.1038/ s4 441586-023-06747-55
Lampinen, A., Chan, S., Dasgupta, I., Nam, A., and Wang, Passive learning of active causal strategies in agents and language models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 1283-1297. Curran Associates, Inc., 2023.
Vashishtha, A., Reddy, A. G., Kumar, A., Bachu, S., Bala- subramanian, V. N., and Sharma, A. Causal order: The key to leveraging imperfect experts in causal inference. In Proceedings of the Thirteenth International Conference on Learning Representations (ICLR 2025), Singapore, 2025. International Conference on Learning Representa- tions. OpenReview ID: 9juyeCqLOu.
Li, Y., Bubeck, S., Eldan, R., Giorno, A. D., Gunasekar, S., and Lee, Y. T. Textbooks are all you need ii: phi-1.5 technical report, 2023.
Long, S., Piché, A., Zantedeschi, V., Schuster, T., and Drouin, A. Causal discovery with language models as imperfect experts. In ICMI. 2023 Workshop on Structured Probabilistic Inference & Generative Modeling, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit. J., Jones, L.,, Gomez, A. N., Kaiser, E.., and Polosukhin, I. Atten- tion is all you need. In Guyon, I., von Luxburg, U., Ben- gio, S., Wallach, H., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30 (NeurIPS 2017), pp. 5998-6008, Long Beach, CA, USA, 2017. Curran Associates, Inc.
Morishita, T., Morio, G., Yamaguchi, A., and Sogawa, Y. Enhancing reasoning capabilities of Ilms via principled synthetic logic corpus. arXiv preprint 1r12422411.12498, 2024.
Ontañión, S., Ainslie, J., Cvicek, V., and Fisher, Z. Making transformers solve compositional tasks, 2022.
Willig, M., Zečević, M., Dhami, D. S., and Kersting, K. Probing for correlations of causal facts: Large language models and causality. 2022.
Papadimitriou, I. and Jurafsky, D. Injecting structural hints: Using language models to study inductive biases in lan- guage learning. In Bouamor, H., Pino, J., and Bali, K. (eds.). Findings of the Association for Computational Linguistics: EMNLP 2023, pp.
Zečević, M., Willig, M., Dhami, D. S., and Kersting, K. Causal parrots: Large language models may talk causality but are not causal. Transactions on Machine Learning Research, 2023.
Singapore,

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Zhang. Y., Backurs, A., Bubeck, S., Eldan, R., Gunasekar, S., and Wagner, T. Unveiling transformers with lego: synthetic reasoning task, 2023,
12

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Appendix
Query Type
Data Instance Example (Prenise-lypodhesss-aaeee
Structure Type
Network Size (number of nodes)
(Train/
Eval)
Transitivity
Train
Mhb causes iqB. iqB causes G. Does cause iqB?: No
Short Linear
3-6
Sequence
Train
N5w causes 6D causes s. Does N5w cause s?: Yes
Short Sequence with Random Flipping
3-6
Eval
w3 causes ROv, w3 causes IQC. H causes ROv, H causes 1QC. causes ROv, causes w3, b causes H. Does iQC cause ROv?: No
Branching
5,8,10,12
Eval
LKk causes 50v. Kk causes LO. LO causes 500. causes c. Does KWO cause L.0: No
Shuffled Sequences
3-9
Eval
FDAH26m V7 causes HtzalHjIY. 7tzalHjlY causes OkspcX95Im. 0kspcX951m causes 7rhFSIx209. 7rhFSIx209 causes IPIGSLHVqp. Does FDAH26mV7
Sequences with Longer Node Names
3.9
cause 7tzalHjiY?? Yes
Eval
I causes 1Z. rZ. causes L. L. causes bUx. bUx causes Phr.
Long Linear Sequences
7-15
Pbr causes Iw. Iw causes c3. c3 causes yBQ. yBQ
causes yk. yK causes w. w causes causes kH. kH causes lu. causes jV7. jV7 causes i. Does cause rZ?: Yes
Eval
rU6 causes eF, eF causes ivC. 3R causes ivC. 3R causes A8. 2 causes A8. 2 causes causes a03. causes a03. causes y. b causes h. h causes yN. ic0 causes yN. ic0 causes Hd. Hd causes U. Does rU6 cause eF?: Yes
Long Sequences with Random
7-15
Flipping
Table A1: Table with examples of transitivity data instances of different causal structural networks used for training and evaluating models.
A. Transitivity Axioms
Length Generalization: Table A4 shows the accuracy different models when evaluated on larger causal chains that were seen during training. Among the baseline pre-trained LMs, GPT-4 obtains the highest accuracy on both sequential and randomly flipped chains for the multi-shot setting. It is remarkable that our TS2 (NoPE) model obtains competitive performance the trillion-scale GPT4 model. In particular, for chains of size 7-12, TS2 (NcPE) obtains higher comparable accuracy than GPT-4 on both sequential and randomly flipped chains. Similar trends are observed for chains of size 7-13 when compared to GPT-4 in the zero-shot setting. Our model's accuracy decreases for chains of length 14-15 (0.85 for sequential chains and 0.78 for randomly flipped chains) but is still significantly higher than that of LMs like Gemini-Pri and Phi-3. Although in-context examples multi shot setting improve the performance of baseline LLMs, TS2 (NoPE) still outperforms both Gemini Pro and Phi-3 the multi-sho setting. Note that random prediction would yield a 50%
13

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Query Type
Data Instance Example (Premise-lypottesss-aakd)
Structure Type
Network Size (number of nodes)
(Train/
Eval)
D-Separation
Train
Iel causes kT. KT causes t4d. t4d causes zW. zW causes Z4P. Z4P causes pij. Are zW and pij d-separated given [t4d, kT, Z4P)?: Yes
Short Linear
6
Sequence with Non Empty Conditioning Set
Train
nI. causes A. A causes xx. xx causes 5Cg. Are xx and nL d-separated?: No
Short Linear
4
Sequence with Empty Conditioning Set
Train
ZWn causes P9. u causes P9, B causes NS causes B. Are P9 and d-separated given [B]?: No
Short Sequence with Empty Conditioning Set
5
Train
ZWn causes P9. u causes P9. B causes u. NS causes B. Are P9 and d-separated given (B, ZWn)?: No
Short Sequence with Empty Conditioning Set
5
Eval
FZg causes I. FZg causes Y. causes vEU. a causes VBU. 15 causes a. f5 causes R. R causes O. 2WJ causes O. 2WJ causes TDA. TDA causes 9d. Are 2WJ and d-separated given (FZg)?
Long Sequence
10
with Random Flipping
Eval
I causes I causes OP. causes WT. faG causes t. faG causes Z. pK causes OP. pK causes 0K3. pK causes yUa. 0K3 causes a. 0K3 causes OP. 0K3 causes wT. Z causes yUa. rY6 causes faG. rY6 causes wT.. Are t and yUa d-separated given (faG)?: Yes
Branching
12
Table A2: Table with examples of d-separation data instances of different causal structural networks used for training and evaluating models. We have instances with null and empty conditioning set for d-separation in train and evaluation set.
accuracy, indicating that the ddiomatically-trainee TS2 (NoPE) model can generalize its reasoning to causal chains much longer than even though was trained only on chains up to length 6.
Node Name Shift: For models trained on TS2 dataset, we also evaluate generalization to changes in variable names (Figure 2). We find that TS2 (NoPE) robust to node name changes and retains its high accuracy as new, longer names are introduced. It also retains its generalizability to longer sequences with new node names, performing similarly to GPT-4.
Summary: Across all evaluation setups, our axiomatically trained model TS2 (NoPE) performs significantly better than random baselines even as chain lengths are increased beyond its training data. In particular, even though our model was not trained on fully reversed chains, it performs at par with the significantly larger GPT-4 model, while easily outperforming other billion scale models even under multi-shot settings. For other tasks, it often outperforms or matches the accuracy of billion-scale models like Gemini Pro and Phi-3. These results indicate that a model trained axiomatically can learn to reason about more complex causal structures from demonstrations of simple causal sequences.
14

---

Teaching Transformers Causal Reasoning through Axiomatic Training
Model
3
5
Baselines
Zero Shot
GPT-4
0.97
0.99
0.98
0.92
Gemini Pro
0.61
0.59
0.66
0.62
Phi-3
0.80
0.69
0.73
0.69
Multi Shot
GPT-4
1.00
1.00
1.00
0.99
Gemini Pro
0.95
0.87
0.77
0.71
Phi-3
0.93
0.89
0.75
0.75
Axiomatic Training
TSI NoPE
0.99
0.97
0.90
0.91
TS1 LPE
0.99
0.98
0.95
0.93
TSI I w SPE
1.00
0.98
0.95
0.96
TS1 w RoPE
0.97
0.97
0.96
0.98
TS2 NoPE
0.98
0.96
0.90
0.91
TS2 LPE
0.99
0.97
0.92
0.96
TS2 SPE
0.99
0.97
0.93
0.94
TS2 w RoPE
0.99
0.98
0.97
0.98
OCC NoPE
0.41
0.24
0.18
0.13
OCC RoPE
0.59
0.26
0.22
0.20
Table A3: Following (Berglund et al., 2024), evaluate models inferring cause-and-effeet from fully reversed sequences absent in training data. Models trained on OCC perform worse, highlighting the importance of edge-level perturbations for generalization. Accuracy metric is reported, with random baseline 0.5. Best performance is bolded, while second best is underlined.
15

---

Teaching Transformers Causal Reasoning through Axiomatic Training
A.1. Additional Results: Role of Data Diversity and Positional Encoding
Importance of Data Perturbations. We find that diversity of the sequences in train data plays an important role. Model trained only causal chains (OCC) generalize to longer chains (Table A4) but not to other DAG structures (see Figure 3 for edge flip, Table for branching). Models trained TSI or TS2 generalize across all scenarios, including random flip, order permutations, and branching; thus highlighting the impact of incorporating variability at the edge level through random flipping. However, across tasks, we find that TS2 yields higher accuracy than TS1, even as TS1 has more variations due to random flipping. This suggests that while perturbations aid structural generalization, excessive perturbations can hinder it (in particular, random flipping may decrease the length of available causal paths during training).
Role of Positional Encodings. When comparing models based on positional encoding, we find that models without positional encoding generalize well to both longer chains (up to length 15) and unseen complex graph structures, despite being trained only on linear chains with 3-6 nodes. Models with SPE and LPE perform well on longer chains but struggle with longer node names, even in smaller graphs (Figure 2). highlighting their sensitivity to minor perturbations. SPE also underperforms in branching and order-based settings like shuffling and reversal. Learnable PE works up to 9-1ength chains but drops afterward. Overall, our results extend earlier work on the utility of NoPE (Kazemnejad et al., 2023; Haviv et al., 2022) to the task of understanding causal sequences and generalizing to both longer length and complex structure at test time. Interestingly, all PEs perform well in randomly flipped sequences, likely due to the short effective path lengths caused by the 0.5 probability of forward-directed edges.
B. Zero Shot Setting
To evaluate the baseline models, we follow a simple zero-shot prompting strategy. For each tuple, we provide the natural language expression of the causal graph (Premise) followed by the question (Hypothexis) and prompt the LM to answer in either YYes' or 'No' (Label). Here is an example prompt: "EX causes T. T causes 9. 9 causes W. causes 7. 7 causes M. M causes a. Does EX cause T? Answer 'Yes' or 'No' only". See Table ?? contains examples of prompts used.
C. Multi-Shot Prompt
C.1. Cause-Effect Inference Task
Chain lengths of the in context examples ranged from 3 to 6 to maintains consistency with the training and testing paradigm used for our 67-million million--paammeee model.
The following multi-shot prompt was used to evaluate the baselines and models across different test sets, assessing their generalization based on length, order, and branching.
Following the given examples answer the question regarding causal relationship between two variables: '5e0 causes vAf. vAf causes VO. Does vAf cause VO?: Yes'
causes 5ee causes vAf. vAf causes VO. Does vAf cause Se0?: No'
"eOF causes Z. Z causes OU. OU causes mR. mR causes 1L. Does mR cause IL?: Yes'
"eOF causes? Z. Z causes OU. OU causes mR. mR causes IL. Does Z cause e0F?: No' "b causes K. K causes qPv. 5 causes qPv. Does b cause qPv?: Yes'
"b causes K. K causes qPv. 5 causes qPv. Does b cause 5?: No'
Mhh causes 10a. 6Eh causes Mhh. NS causes 6Eh. causes NS. causes Xu. Does xu cause 6Eh?: No'
Mhb causes 10a. 6Eh causes Mhb. NS causes 6Eh. causes NS. causes XU. Does cause NS?: Yes
16

---

Teaching Transformers Causal Reasoning through Axiomatic Training
D. Results and Analysis
Model
RF
RF
RF
FS
FS
FS
RF
FS
RF
RF
FS
RF
Baselines
Single Shot
GPT4
0.95
0,98
0.97
0.93
0.87
0.94
0.91
0.87
0.90
0.95
0.92
0.92
0.85
0.93
0.93
0.93
0.89
0.86
Gem-Pro
0.53
0.73
0.69
0.74
0.64
0.75
0.65
0.81
0.72
0.78
0.60
0.80
0.59
0.68
0.67
0.64
0.61
0.66
Multi Shot
GPT-4
0.97
0.99
0.93
0.99
0.92
0.96
0.88
0.94
0.89
0.97
0.89
0.93
0.88
0.95
0.93
0.94
0.86
0.94
Gem-Pro
0.82
0.73
Phi-3
0.83
0.92
0.89
0.88
0.75
0.86
0.66
0.87
0.80
0.90
0.80
0.85
0.79
0.82
0.71
0.81
0.72
0.82
Axiomatic Training
TSI NoPE
0.96
0.97
0.95
0.86
0.92
0.78
0.87
0.77
0.90
0.76
0.82
0.82
0.75
0.83
0.70
TSI W LPE
0.98
0.96
0.89
0.94
0.81
0.91
0.61
0.86
0.64
0.87
0.64
0.79
0.60
0.80
0.59
0.81
0.57
0.73
TSI w RoPE
0.99
0.96
0.97
0.95
0.89
0.90
0.82
0.84
0.81
0.84
0.86
0.76
0.76
0.81
0.82
0.70
0.78
0.75
TS2 NoPE
0.98
0,93
0.93
0.92
0.82
0.88
0.74
0.84
0.70
0.85
0.70
0.80
0.71
0.76
0.71
0.77
0.66
0.73
TS2 w LPE
0.99
0.95
0.96
0.94
0.86
0.90
0.72
0.86
0.69
0.85
0.80
0.78
0.73
0.78
0.75
0.80
0.68
0.77
1S2
TS2 w RoPE
0.99
0.97
0.98
0.96
0.90
0.89
0.85
0.87
0.84
0.82
0.87
0.74
0.78
0.80
0.86
0.69
0.78
0.71
OCC NoPE
0.99
0.61
0.98
0.62
0.89
0.62
0,90
0.57
0.90
0.57
0.93
0.52
0.87
0.55
0.93
0.50
0.87
0.53
OCC w RoPE
0.96
0.65
0.98
0.71
0.84
0.68
0.84
0.64
0.80
0.65
0.88
0.56
0.76
0.60
0.84
0.60
0.79
0.55
Table A4: Accuracy of different models on Transitivity axioms. this table, we show accuracy different models the transitivity axioms. The rows shows different models considered for comparison. The top rows denote the performance baseline models with different prompting strategies i.e. single shot and multi-shot prompt (see Sec 4.5 for details). The models listed after axiomatic training shows the performance of transformer models trained from scratch SPE: Sinusoidal Positional Encoding (PE), IPE:: Learnable PE,
axiomatic dataset. TS1 1 and TS2 denote pretraining
data setups and 2 described in Sec 5.1 and di fifferent modifiers w/o PE: No PE, RoPE: Rotary Position Emhedding. For axiomatic training, the model remains the same across setups (67 Million parameters based). The training dataset contain graphs of size 3-6 however the models are tested on graphs of size 7-15 (as mentioned in different columns). FS denotes the graphs that only contain chains that are oriented in forward direction and RF contains the graphs that training set.
also includes random flipping (see Sec 4.1 for details)
Model
Baselines
Single Shot
GPT-4
1,00
1.00
1.00
1,00
1.00
1.00
1.00
Gemini Pro
0.96
0.94
0.86
0.81
0.76
0.73
0.71
Phi-3
Shot
GPT-4
1.00
1.00
0.98
0.98
0.98
0.98
0.97
Gemini Pro
1.00
0.90
0.86
0.88
0.84
Phi-3
0.93
0.89
0.89
0.84
0.82
0.77
0.79
Axlomatic Training
TSI NoPF
1.00
1.00
1.00
0.99
0.98
0.92
0.88
TSI w: LPE
1.00
1.00
0.99
0.97
0.92
0.83
0.74
0.57
TS1 w. RoPE
0.65
0.56
0.56
0.49
0.45
0.49
0.50
TS2 NoPE
1.00
0.99
0.92
0.84
0.76
0.71
0.69
TS2 w LPE
1.00
0.99
0.96
0.90
0.86
0.76
0.74
TS2
0.58
TS2 w RoPE
0.51
0.48
0.48
0.50
0.46
0.46
0.48
0CC NoPE
1,00
0.99
0.98
0.96
0.96
0.91
0.93
0CC w RoPE
0.90
0.77
0.67
0.64
0.65
0.59
0.62
Table A5: Results on node name length generalization. TS1 and TS2 denote Training Data setup 2 from Section ??. OCC the third data setup comprising of sequential causal chains. SPE: Sinusoidal PF, LPE: Learnable PE, w/o PE: No PE, RoPE: Rotary Position Embedding. Model remains the same across all setups (67 Million parameter based). For longer node names, NoPE performs best on sequential linear setup. Accuracy metric is used.
17

---

Teaching Transformers Causal Reasoning through Axiomatic Training
10
Model
BF=1.4
BFa1.4
Baselines
Zero shot
GPT-4
0.98
0.95
0.91
0.90
0,84
0.88
0.82
0.86
Gemini Pro
0.72
0.76
0.71
0.73
0.73
0.71
Phi-3
0.87
0.83
0.82
0.79
0.77
0.77
0.75
0.80
Midti shot
GPT-4
0.94
0.93
Gemini Pro
0.81
0.76
0.77
0.79
0.75
0.77
0.78
0.79
Phi-3
0.77
0.78
0.79
0.82
0.78
0.79
0.80
0.79
Axiomatic Training
OCC w RoPE
0.74
0.72
0.70
0.68
0.66
0.66
0.65
0.62
TSI w LPE
0.76
0.82
0.70
0.72
0.67
0.69
0.63
0.68
TSI w SPE
0.65
0.78
0.57
0.61
0.55
0.59
0.53
0.57
TSI NaPE
0.78
0.82
0.70
0.74
0.67
0.69
0.62
0.66
TSJ w RoPE
0.81
0.86
0.75
0.79
0.73
0.74
0.69
0.71
TS2 w LPP
0.73
0.80
0.68
0.72
0.65
0.67
0.61
0.64
TS2
0.52
0.54
0.52
0.52
TS2 NoPE
0.75
0.82
0.68
0.73
0.67
0.68
0.62
0.64
TS2 w RoPE
0.81
0.88
0.74
0.79
0.70
0.72
0.68
0.68
Table A6: Evaluation with variation in branching factor. Accuracy of axiomatically trained models with LM baselines on the causal graphs with higher branching factor than that in training. See Sec 5.2 for details.
E. Preliminaries and Notations
Causal Models Let M = (X, U,F) be causal model defined over a set of endogenous variables X, exogenous variables U and the causal relationship between then defined by set of structural equations F (Galles & Pearl, 1997). Let be the causal graph associated with the causal model M where the nodes V in G correspond to the variables in M and edge V; + Vj between any nodes V,, Vj denote the causal relationship between them. The causal relationship of node X; is characterized by the functional relationship fi E F s.t., Z; fi u.). Here pa; the parent of the node X; is the corresponding causal graph G and u, U are set of exogenous variables influencing the exogenous variable X, In our work, we assume that there are no hidden confounders so we have one exogenous variable corresponding to every endogenous variable i.e. u, = u, Each exogenous variable has an associated probability distribution which quantifies the uncertainty the system i.e. u; ~ P(uu). Thus the joint distribution endogenous variable is deterministic function of other endogenous and exogenous variables the probability distribution corresponding to the endogenous variable is the push-forward of the exogenous variable i.e (P P## (U).
exogenous variable given by IUU)). Since any
E.1. Definitions
Following the formal definitions provided by (Jin et al., 2024a), we explain the following terminologies:
Markov Property In directed acyclic graph (DAG) G, the Markov property asserts that each node X; is conditionally independent of its non-descendants given its parents. This can be written as X; JL NonDe(X;) PP((X)), where NonDe(X,) represents the set of non-descendants of X;, excluding the node itself, and Pa(X;) denotes its parents. Leveraging the Markov property, the joint distribution over all the nodes can be factorized as:
N II P(X, Pa(X.)).
P( XN)
Markov Equivalence Class Two directed acyclic graphs (DAGs) are considered Markov equivalent if they induce the same joint distribution P(X). The collection of DAGs that are Markov equivalent is referred to as a Markov equivalence class (MEC). Causal graphs within the same MEC can be easily recognized as they share the same skeleton (i.c., undirected edges) and V-structures (i.e., configurations of the form A +- C, where and are not directly connected).
18

---

Teaching Transformers Causal Reasoning through Axiomatic Training
F. Positional Encodings and their Role in Generalization
Positional Encoding (PE) play crucial role of providing information about absolute and relative position of tokens in a sequence (Vaswani al., 2017). (Vaswani al., 2017) propose an absolute positional encoding strategy using periodic functions (e.g., sinusoidal cosine) to initialize these encodings. Absolute positional encoding provides definite values for all positions across any sequence length. However, studies (Ontañón al., 2022: Csordás et al., 2021) show absolute positional encoding fails in length generalization tasks for transformers. In the learnable APE variant (Radford et al., 201 8), each positional embedding is randomly initialized and trained with the model. This approach falters with sequences longer than those seen in training, as the new positional embeddings remain untrained and randomized. Interestingly, recent findings (Kazemnejad et al., 2023; Haviv et al., 2022) indicate that removal of PEs in auto-regressive models can improve model's length generalization capabilities, wherein the attention mechanism during auto-regressive decoding sufficient to encode positional information. We also experiment with Rotary Position Encodings (Su al., 2024), which have shown superior length generalization. We use 10000.0 for the base period of RoPE embeddings.
G. Formalising training and Evaluation Setup
Let Saim represent the maximum value for given perturbation dimension dim, along which we construct train and evaluation sets for our axiomatic framework. For each dimension, we choose a threshold Tdem L, such that faim < Tdim forms our training set and fdim Tdim forms the evaluation set. So, fdim E {fiens foranch, fnodeien: frevfactors fahuf fie] where:
fien maxx (len Vi)), gives the maximum number of nodes across all causal sequences. Teen for length is set at 6, with fien 6].
foranch = max ( (X.l gives the maximum branching factor a dataset, with Toranch = 0.8 (for 6 node linear sequences). For sequences the train set, the branching factor ranges from 0.6 to 0.8 for to 6 length sequences.
Let Li be the length of the name of the node Xi,, then 4.j = (I) Thereefre, the maximum length of node names across all nodes in all causal sequences can be represented as: Inodenamelen = maxi 15/5m lij- We set Trodelen for train set as 3, with Inodelen E [1, 3].
Given any causal sequence X; and function N, where N(Xi,; Xi.+1) returns natural language representation of a directed edge between j and j + 1 node in the causal chain X,. fshuffie = Dts.pPerm/ (N [N X1)+1)), where represents deviation from original sequential order of natural language sentences to represent X,-
Given causal sequence X2 and let R(X,, frevfactor) be operation on the causal chain that flips the direction of every edge in the sequence with probability Jreufactor- In the training set, there is directed edge between every sequential pair of nodes Xi,,, Xis+1 with frewfactor = 0 (for linear sequence, Xi, + x1)++1) or 0.5 (for sequence with random flipping, X., + Xij++ or Xi, t Xi)+1) In the evaluation set Sreufactor = i.e., all sequences for reversal evaluation setup are completely reversed unlike in train set where no sequence is present where all edges are completely reversed.
H. Results of Dsep on CLEAR dataset
Models
7
8
10
11
12
13
14
GPT-4
0.57
0.64
0.52
0.50
0.53
0.52
0.51
0.50
Hlma-3-8b-Inntruu
0.51
0.50
0.54
0.51
0.48
0.51
0.53
0.51
0.952
0.948
0.954
0.850
0.87
0.88
0.73
0.66
Models with different PEs trained from scratch
SPE
0.93
0.95
0.97
0.95
0.71
0.61
0.80
0.44
LPE
0.96
0.93
0.99
0.92
0.68
0.71
0.62
0.47
NoPE
0.89
0.93
0.85
0.94
0.68
0.65
0.60
0.5
RoPE
0.96
0.91
0.96
0.92
0.70
0.69
0.54
0.58
Table A7: Performance DSEP of longer chains with random flipping
19

---

Teaching Transformers Causal Reasoning through Axiomatic Training
I. Examples of instances from each benchmark: CLEAR and Corr2Cause
Premise: Given DAG (directed acyelic graph) with nodes Z, P. V and directed edges C->V, P-SV, P.
Hypothesis: "Which of following nodesets can d E-separate node C node
A. ['22, )
B. ('V')
C. ('Z")
D. set ()
Answer: C
Figure Al: Example instance of Multiple Choice (MC) question type from (Chen et al., 2024) dataset describing d d-separation rule problem defined with different hypothesis type and semantic structure then the one our models are finetuned on.
Premise: Suppose there is closed system of 4 variables, A, B, and D. All the statistical relations among these 4 variables are as follows: A correlates with A correlates with C. A correlates with D. B correlates with C. B correlates with D. correlates with D. However, and D are independent given A, B and D are independent given A and C. C and D are independent given A, C and D are independent given A and B.
Hypothesis: There exists at least one callider (i.e., common effect) of A and B. Label: No.
Figure A2: Example instance from the Corr2Cause dataset, where the model must infer the presence of collider between variables given only correlational and conditional independence statements.
J. Implementation Details
We used a learning rate of lc-4 with linear scheduling and 3% warmup ratio, training for 4102 max steps on axiomatic instance samples with sequences of maximum length 4096 tokens. We employed mixed precision (bfloat16) training with flash attention for efficiency. After training, the LoRA weights were merged with the base model for inference. We used Huggingface (wol, 2020) for implementation. The fine-tuning used LoRA with rank 64, alpha 16, and dropout 0.1. Training was performed 3 GPUs using SpepSpeed Stage with total batch size of 128 (16 samples per GPU with gradient accumulation). A40 and A100 GPUs were used for training the transformer model from scratch for all Positional encodings based experiments.
20


Reviews:
[{"review_id": "exCC4FBNjd", "reviewer": "Reviewer_DrTc", "summary": "This paper propose an approach where the model learns symbolic axioms through demonstrations rather than directly infer causal relationships from data. And then, they investigate whether this approach allows the model to generalize from learning simple causal structures to more complex causal relationships.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating the contributions of our work. We answer specific questions below:\n\n**Response to Weaknesses**  \n> **1. GPT-4 achieves the best performance, suggesting that causal axioms relation can be learned from unstructured and massive datasets without requiring complex data preprocessing or specialized positional encoding.**\n\n**Response**:\nWhile GPT-4 performs well on transitivity-based tasks, its accuracy drops to near-random on complex causal rules like d-separation (see Tables A6, A7). Despite its scale, GPT-4 struggles with d-separation, a key rule for inferring conditional independence in causal inference. In contrast, our 67M-parameter model achieves higher accuracy on d-separation, for complex graphs unseen during its training. \n\nMoreover, finetuning Llama-3-8B-Instruct with axiomatic training further improves performance on complex graphs and the CLEAR benchmark (Table 3). On CLEAR’s multiple-choice task, the finetuned model achieves 50% accuracy, outperforming GPT-4’s 36.67% (with a random baseline of around 25%). The CLEAR benchmark has evaluation instances for d-separation that are semantically different from instances in the axiomatic training data. Also, CLEAR benchmark has different types of hypothesis (multiple choice question answer, yes/no type questions), whereas axiomatic instances only contained yes/no hypothesis instances, thus showing potential of the model to generalize for diverse problems beyond simple binary classification. \n\nLastly, Table 4 shows results on the Corr2Cause benchmark, that evaluates the model's reasoning capability to infer causal relationships from correlational statements, shows the potential of axiomatic training on real world problems which require application of different causality rules (d-separation, transitivity, markov property, markov equivalence class). A Llama-3 model finetuned on axiomatic data obtains 0.37 F1 score and outperforms GPT-4 which struggles to perform well on the task (with 0.29 F1 score). \n\n> **2. This raises concerns about whether the proposed training approach provides a significant advantage over large-scale unsupervised learning.**\n\n**Response**\nWe thank the reviewer for this point. While GPT-4 does show potential on simple graph reachability problems requiring application of transitivity axiom, we found that other billion-scale language models (Phi-3 and Gemini Pro) struggle with these problems. Keeping efficiency in mind, we believe our axiomatic framework provides a potential way to improve causal reasoning for smaller models that can balance correctness and efficiency.\n\nAlso, as stated above, while GPT-4 model does perform well on transitivity causal axiom, it struggles on applying the d-separation rule of causal inference and performs close to random baseline. This result shows that current billion scale models are unable to reason on fundamental rules of causal inference despite their large scale unsupervised training.", "claims_and_evidence": "N/A", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "I believe this paper makes a meaningful contribution to the scientific leterature.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strength\n- This paper introduces axiomatic training as a novel approach for training transformers, a field that has not been extensively explored.\n- To the best of my knowledge, this is the first work to train transformers specifically to learn causal axioms.\n- The proposed method demonstrates generalization to more complex causal relationships beyond the training set.\n\nWeakness\n\n- GPT-4 achieves the best performance, suggesting that causal axioms relation can be learned from unstructured and massive datasets without requiring complex data preprocessing or specialized positional encoding.\n- This raises concerns about whether the proposed training approach provides a significant advantage over large-scale unsupervised learning.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a6UNVQRw51", "reviewer": "Reviewer_KAqL", "summary": "This paper studies a new method for improving the causal reasoning capabilities of autoregressive transformer text models by training on synthetically generated data containing demonstrations of causal axioms or rules. Specifically, the authors consider the expressions of the form <premise, hypothesis, result> generated using random causal DAGs with structured perturbations. The authors carefully construct the training and test datasets such that the testing measures generalization to unseen causal graphs. By training (or fine-tuning, in case of natural language pre-trained) on transitivity axiom and d-separation rule, and testing on unseen causal graphs as well as causality benchmarks, the authors show that training on linear causal chains and their structural perturbations induces meaningful generalization.\n\n## update after rebuttal\n\nGiven the clarifications by the authors, I would like to keep my supportive rating.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating the contributions of the work. We answer the specific questions below:\n\n**1. Eval inputs being longer than training text inputs?**\n\nYes the evaluation set consists of chains longer than the ones in the training setup, and this translates to having longer text input than the text inputs in the training set. Table A1 consists of sample instances from our training and evaluation set, showing that the text instances in the evaluation set typically tend to be longer than the ones the model was trained on.\n\n**2. The two assumptions employed...how restrictive these are?**\n\nWe assume an empty conditioning set only for the transitivity axiom; for the d-separation rule (Section 3.2), we include conditioning sets of various sizes, as detailed in Section 6. \n\nIn both cases, no unobserved confounders were assumed to simplify our setup. If variables' data values are unobserved but the edge structure is known, our framework can readily incorporate unobserved variables symbolically; however, if the structure of unobserved variables is unknown, the problem becomes more complex—a challenge we leave for future work.\n\n**3. Discussion on whether causal reasoning problems are a subclass of graph reasoning**\n\nThat's a great point. Our work can be considered as studying a subset of graph algorithms that are relevant for causality. Our first task, transitivity axiom, can be seen as a special case of the graph reachability problem studied in Sanford et al. However, the d-separation task involves a specialized definition for causality that is usually not studied in graph algorithms literature. In other words, our work focuses on the intersection of causality and graph reasoning, specifically on the Pearlian framework of causality which focuses on DAGs. That said, the axiomatic training framework is general and can potentially be extended for other graph reasoning problems, even beyond causality.\n\n**4. The connection to positional encoding .... encodings suited for causal reasoning.**\n\nOur work builds on analyses by Kazemnejad et al. and Shen et al. (references in paper) regarding positional encodings (PEs) and length generalization. While their focus was on tasks like copying and addition, we investigate how different encodings affect length generalization for causal reasoning problems. We corroborate past findings that using NoPE outperforms absolute methods (sinusoidal and learned) because these methods make sequence length explicit during training, leading to poor performance on unseen lengths. In contrast, RoPE improves generalization and addresses these limitations. \nAlthough designing encodings specifically for causal reasoning would be interesting, our focus is on how axiomatic training aids language models in causal reasoning. Further, we used existing PEs so that our work is applicable to practical models such as Llama. We will clarify these points in the final version.\n\n**5. Prior work on passive learning**\n\nLampinen et al. examined if agents can develop causal understanding by passively observing interventions on synthetic tasks. Building on their work, our study explores whether causal reasoning axioms can be learned without active interventions by generating passive data from simulated axiom inferences on diverse synthetic graphs. Unlike Lampinen et al., who focus on observational learning for test-time interventions, our approach offers a practical training method—axiomatic training—for language models. This method enhances a transformer's ability to generalize complex causal rules over unseen networks and potentially apply interventional reasoning despite being trained only on passive data. We will add further details in the final draft. \n\n**6. The length generalization result ... compare to past work**\n\nAbbe et al. highlight the “globality barrier,” showing that high-globality tasks like syllogisms require many tokens to capture nontrivial correlations and suggesting scratchpad techniques to break them into manageable subtasks, while Prystawski et al. demonstrate that prompting intermediate steps in locally structured, chain-of-thought reasoning improves performance by decomposing complex inferences into sequential computations. In contrast, our axiomatic training framework constructs a dataset that enables compositional reasoning—such as applying the transitivity axiom repeatedly—by offering **diverse demonstrations**. In line with the claims above, when we only provide sequential causal chains of length 3 as training data, we found that the model did not generalize. We obtain generalizability only when we provide diverse training instances: chains of varying lengths (3–6) and random edge flips, which encourages the model to learn the abstract rule directly rather than relying on intermediate signals. Positional encodings like RoPE further help preserve relative information beyond the training lengths.", "claims_and_evidence": "The claims made in the submission is supported by clear and convincing evidence. In particular, the authors have made effort in ensuring that genuine generalization capability is measured at test-time, by carefully ensuring that the training data does not contain the test inputs (also for Llama). The evaluation does not only involve the datasets generated by the authors, but also includes two existing causal reasoning benchmarks, which is a good supporting evidence.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria makes sense for the problem of training a text model to have causal reasoning capabilities as far as I can confirm.", "theoretical_claims": "The authors do not make theoretical claims in this work.", "experimental_designs_or_analyses": "The experimental design is sound as far as I can verify, especially in the design of the evaluation dataset to be OOD. I have some minor clarifying questions:\n\n- In the beginning of page 5, the authors mention that the training set involves chains with size 3-6 and the evaluation set involves chains of length 5-15. Does this translate to evaluation text inputs (<premise, hypothesis, result>) being longer than the training text inputs?", "relation_to_broader_scientific_literature": "The contributions in this paper are related to causal reasoning of text models, more generally the out-of-distribution generalization of these models. In the particular context of this paper, the tasks in consideration are essentially reasoning problems on graphs, so there is also a close relation to the capability of text models in learning and executing graph algorithms (please see the Other Strengths And Weaknesses section on this).", "essential_references_not_discussed": "I am not aware of essential references not discussed in the paper.", "other_strengths_and_weaknesses": "Strengths\n\n- The paper is well-written and easy to follow.\n- The presented evidences on generalization when trained (or fine-tuned) on the proposed synthetic data are strong and significant.\n\nFor the weaknesses, I have no major concerns, but would like to hear the authors' response to the following minor concerns and questions.\n\n- The two assumptions employed in this work are the absence of unobserved confounders and the empty conditioning set $Z$. It is relatively less discussed how restrictive these assumptions are.\n- While presented in the context of causal reasoning, the axiomatic training and evaluation tasks considered in this paper are essentially same as learning and executing certain graph algorithms (Luca & Fountoulakis, 2024; Sanford et al. 2024; Wang et al. 2023). Some discussion on whether causal reasoning problems are in general a subclass of graph reasoning, or this is a particular characteristic of the problem setup considered in the paper, would be informative.\n- The connection to positional encoding and length generalization of transformer in general is interesting, but in the end the paper has resorted to trying out a range of established positional encodings empirically, rather than contributing, e.g., in-depth analysis on why certain positional encodings work better, or a new construction of positional encodings suited for causal reasoning.\n- In page 2, the authors mention that the results contribute to the literature on causal learning from passive data; it would be nice to have some discussion on how the paper differs from the prior work in the domain (it is currently missing in Section 2, as far as I can verify).\n- The length generalization result was in particular surprising to me, given that some prior work showed that such generalization typically requires some form of process-based training signal (Abbe et al. 2024; Prystawski et al. 2023). Can the authors provide some insights on why the models considered in this work do generalize in length?\n\nLuca & Fountoulakis, Simulation of Graph Algorithms with Looped Transformers (2024)\n\nSanford et al. Understanding Transformer Reasoning Capabilities via Graph Algorithms (2024)\n\nWang et al. Can Language Models Solve Graph Problems in Natural Language? (2023)\n\nAbbe et al. How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad (2024)\n\nPrystawski et al. Why think step by step? Reasoning emerges from the locality of experience (2023)", "other_comments_or_suggestions": "Minor typos:\n- In Lines 60-61, there seems like an un-erased memo.\n- In Line 153, \"adapted from\" is colored blue\n- Above Section 4, markov -> Markov\n- In Section 7, Refer I -> Refer to I\n- In Section 7.1, Refer Table 3 -> Refer to Table 3", "questions_for_authors": "Please see the weaknesses section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eWEG4fAohc", "reviewer": "Reviewer_HBJz", "summary": "This paper proposes axiomatic training, which leverages synthetic data to train small models from scratch. The authors observed that their approach enables models to generalize from small-node to large-node causal structures when evaluated on transitivity axioms and d-separation rules. Moreover, fine-tuning Llama-7B with axiomatic training improved performance on the CLEAR and Corr2Cause benchmarks, surpassing the performance of larger models like GPT-4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating the contributions of the work. We answer specific questions below. \n\n>**1. I have some concerns regarding the performance of the small models. Although the authors explored different training strategies and positional embeddings, the results appear to fall short of the prompting-based performance of existing large models.**\n\n**Response**:  While GPT-4 performs well on transitivity-based tasks, its accuracy drops to near-random on complex causal rules like d-separation (see Tables A6, A7). Despite its scale, prompting-based GPT-4 model struggles with d-separation, a key rule for inferring conditional independence in causal inference. In contrast, our 67M-parameter model achieves higher accuracy on d-separation, for complex graphs unseen during its training. This shows the potential of axiomatic training for causal tasks. Below we provide further practical evidence comparing prompting-based GPT-4 to an axiomatically finetuned Llama model.\n\nFinetuning Llama-3-8B-Instruct with axiomatic training further improves performance on complex graphs and the CLEAR benchmark (Table 3). On CLEAR’s multiple-choice task, the finetuned model achieves 50% accuracy, outperforming GPT-4’s 36.67% (with a random baseline of around 25%). The CLEAR benchmark has evaluation instances for d-separation that are semantically different from instances in the axiomatic training data. Also, CLEAR benchmark has different types of hypothesis (multiple choice question answer, yes/no type questions), whereas axiomatic instances only contained yes/no hypothesis instances, thus showing potential of the model to generalize for diverse problems beyond simple binary classification. \n\nLastly, Table 4 shows results on the Corr2Cause benchmark, that evaluates the model's reasoning capability to infer causal relationships from correlational statements, shows the potential of axiomatic training on real world problems which require application of different causality rules (d-separation, transitivity, markov property, markov equivalence class). A Llama-3 model finetuned on axiomatic data obtains 0.37 F1 score and outperforms GPT-4 which struggles to perform well on the task (with 0.29 F1 score). \n\n>**2. Could the authors clarify whether pretraining Llama-7B with this approach consistently outperforms prompting-based methods for large models in MultiEvalSLR Dataset?**\n\nBelow we present the results of applying the Llama-3-8b-Instruct model (finetuned with transitivity axiom data) to the MultiEvalSLR task.  We find that the Llama model performs at par with large prompt-based models like GPT-4. \n\nMultiEvalSLR is a complex evaluation set where each input premise contains different types of complexities compared to sequential causal chain: shuffling of premise, chains with randomly flipped edges, and longer chain lengths than the training instances. We train the Llama model on an axiomatic training dataset (same as in paper) that does not include any instances with shuffled premise. Up until length 6, its accuracy is better than GPT-4. However, the GPT-4 still obtains higher accuracy when the length of the chain exceeds the chain length the fine-tuned model saw during its training (>6). \n\nAs stated above, however, we believe the strong performance of GPT-4 is due to the simplicity of the axiom. For the d-separation rule, both axiomatically trained small model (67M) and LLama model (8B) show significant improvements over prompting-based GPT-4 (see Tables 3, 4, A6, and A7).\n\n| Model                                   | Length 3 | Length 4 | Length 5 | Length 6 | (OOD-Length) 7 | (OOD-Length) 8 | (OOD-Length) 9 |\n|-----------------------------------------|----------|----------|----------|----------|----------------|----------------|----------------|\n| GPT-4                                   | 0.99     | 0.97     | 0.89     | 0.85     | 0.95           | 0.90           | 0.90           |\n| Gemini Pro                              | 0.75     | 0.73     | 0.72     | 0.76     | 0.71           | 0.68           | 0.74           |\n| Phi-3                                   | 0.88     | 0.86     | 0.82     | 0.79     | 0.76           | 0.73           | 0.79           |\n| Llama-3-8b-Instruct-Finetuned           | 1.00     | 0.98     | 0.93     | 0.86     | 0.85           | 0.76           | 0.71           |\n\n> https://arxiv.org/pdf/2405.15071, discusses the generalization capabilities of large models on synthetic dataset, could be relevant to this work.\n\nThanks for pointing out this work. We will add a discussion on this, especially in the context of composition-based reasoning.", "claims_and_evidence": "The authors argue that axiomatic pretraining on synthetic data helps models better understand causality, and that this capability generalizes effectively to other tasks. Improvements in two experimental settings support the claim, demonstrating the versatility of this approach.", "methods_and_evaluation_criteria": "The authors constructed a synthetic NLI dataset to pretrain models. They trained a 67M parameter small model from scratch and fine-tuned Llama-7B. The evaluation was conducted on the training task as well as on CLEAR and Corr2Cause datasets. These benchmarks are reliable indicators of causal reasoning performance.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "* CLEAR dataset: Accuracy on the d-separation task (binary classification) improved from 30% to 70%, while multiple-choice accuracy improved from 33% to 50%.\n* Corr2Cause dataset: F1 scores improved by 20%, outperforming GPT-4.\n* For small models, the authors explored different positional embeddings and training strategies, providing insights for future pretraining approaches.", "relation_to_broader_scientific_literature": "This is an intriguing study. As natural datasets become increasingly scarce, the use of synthetic data as a supplement for training is an exciting direction for future research.", "essential_references_not_discussed": "https://arxiv.org/pdf/2405.15071, discusses the generalization capabilities of large models on synthetic dataset, could be relevant to this work. It may provide insights into the sources of generalization.", "other_strengths_and_weaknesses": "See additional comments.", "other_comments_or_suggestions": "Line 60: The authors left a comment that violates the anonymity principle.\nLine 138: A formula exceeds the page margin and needs to be adjusted.", "questions_for_authors": "* I have some concerns regarding the performance of the small models. Although the authors explored different training strategies and positional embeddings, the results appear to fall short of the prompting-based performance of existing large models.\n* Could the authors clarify whether pretraining Llama-7B with this approach consistently outperforms prompting-based methods for large models in MultiEvalSLR Dataset?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents a novel training method designed to integrate the transitivity axiom and d-separation rules into autoregressive transformer text models. The authors generate random directed acyclic causal graphs and use these graphs to synthetically create text data in the format of <premise, hypothesis, result>. This approach enables the development of training and testing datasets that allow for the evaluation of the models' generalization capabilities with respect to unseen causal graphs. By training on the transitivity axiom and d-separation rule, and testing on novel causal graphs and benchmarks, the authors demonstrate that training with linear causal chains and their variations results in meaningful generalization.

Overall, the method proposed in this paper is technically sound and reasonable. The simulation results indicate that this approach could enhance the causal reasoning abilities of language models in validating causal hypotheses based on provided causal knowledge. This work has the potential to impact researchers and practitioners in natural language processing and the development of pre-trained text models. 

However, the comparison baselines may be biased, as they only include transformer-based text models. There are established proof systems and efficient causal discovery algorithms within the causal inference literature that can infer causal relationships from measured correlations. It would be valuable to compare the proposed text model with these traditional causal inference methods to determine whether high-level causal reasoning capabilities could emerge by incorporating selected causal axioms and d-separation rules.
Decision: Accept (poster)