Teaching Transformers Causal Reasoning through Axiomatic Training

## Abstract

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.

## 1 Introduction

Causal reasoning can be defined as a set of reasoning procedures consistent with pre-defined axioms or rules that are specific to causality []. For instance, d-separation and rules of do-calculus can be considered as axioms and specifications of a collider or a backdoor set can be considered as rules that can be derived from axioms. Typically, causal reasoning is done over data corresponding to variables in a system. Axioms or rules are incorporated as inductive biases in a machine learning (ML) model, through regularization, model architecture, or the choice of variables for a particular analysis. Depending on the kind of available data—observational, interventional, or counterfactual—Pearl’s ladder of causation [] defines the kinds of causal reasoning that is possible.

As axioms are the building blocks of causality, we study whether it is possible to directly learn the axioms using ML models. That is, rather than learning from data that is the result of axioms followed by a data-generating process, what if a model can learn an axiom (and thus causal reasoning) directly from symbolic demonstrations of the axiom? Such a model has the advantage that it can be applied for causal reasoning in diverse downstream scenarios, compared to task-specific causal models built using specific data distributions. This question gains relevance as language models make it possible to learn over symbolic data expressed in natural language. In fact, recent studies have evaluated whether large language models (LLMs) can do causal reasoning by creating benchmarks that encode causal reasoning problems in natural language [, , ].

Specifically, we propose a new way of learning causal reasoning through axiomatic training. We posit that causal axioms can be expressed as the following symbolic tuple, $\langle$premise, hypothesis, result$\rangle$ where hypothesis refers to a causal claim and premise refers to any relevant information to decide whether the claim is true or not (conclusion). The conclusion could simply be “Yes” or “No”. For example, the collider axiom from  [] can be expressed as, premise: “$A\perp\!\!\!\perp B,B\not\perp\!\!\!\perp C,A\not\perp\!\!\!\perp C$”; hypothesis: “Does $A$ cause $C?$”; and the conclusion as “Yes”. Based on this template, a 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 key question is: if a model is trained on such data, would it learn to apply the axiom to new scenarios?

To answer this question, we train a transformer model from scratch on symbolic demonstrations of the causal irrelevance axiom []. To evaluate generalizability, we train on simple chains of the causal irrelevance axiom of size 3-5 nodes and test on multiple different aspects of generalization, including length generalization (chains of size 7-15), name generalization (longer variable names), order generalization (chains with reversed edges or shuffled nodes), and structure generalization (graphs with branching). We find that a model trained on simple chains generalizes to applying the axiom multiple times over larger chains, but it is unable to generalize to the more complex scenarios like order or structure generalization. However, when we train a model on a combined dataset of simple chains and chains with some edges randomly reversed, we find that the model generalizes well across all kinds of evaluation scenarios. Extending the findings on length generalization for NLP tasks [, , , ], we find a critical role of positional embedding in ensuring causal generalization across length and other aspects. Our best model has no positional encoding, although we find that sinusoidal encoding also works well for some scenarions.

The axiomatic training approach also generalizes to a harder problem proposed in []. The task is to distinguish correlation from causation given a premise containing statistical independence statements. Solving this task requires knowledge of multiple axioms, including d-separation and Markov property. Using the same method to generate synthetic training data and train the model as above, we find that a transformer trained on task demonstrations over 3-4 variables learns to solve this task for graphs with 5 variables. On this task, our model’s accuracy is higher than larger LLMs such as GPT-4 and Gemini Pro.

Our work provides a new paradigm of teaching models causal reasoning through symbolic demonstrations of axioms, which we call axiomatic training. The data generation and training procedure is general and can be applied to learn any new axiom, as long as it can be expressed in the symbolic tuple format. More generally, our results contribute to the literature on causal learning from passive data [], showing a general way to learn any causal axiom through passive demonstrations.

## 2 Related Work

LLMs for Knowledge-Driven Causal Reasoning: Recent developments in Large Language Models (LLMs) have highlighted their potential for knowledge-driven causal discovery. Unlike traditional methods which focus on statistical patterns or correlations, LLMs utilize knowledge acquired through their pretraining to reason about and identify causal structures based on metadata of variables [, , , , ]. However, possibility of memorization of existing benchmarks in the pretraining of these LLMs has been a major criticism. As a result, recent work [] argues that LLMs 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 []. To evaluate causal reasoning capabilities of LLMs, [] and [] propose formal causal inference evaluation benchmarks to infer direct and indirect causal relationships, and highlight the failure of LLMs in performing accurate causal reasoning.

Impact of Positional Encoding on Generalization: Length generalization capabilities of transformers has been studied in the past to better understand their different failure modes across various settings [, , ]. Previous work [, , , ] emphasizes the impact of positional encoding in length generalization capability of transformers. To understand how transformers can be optimized for learning 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 [] and Sinusoidal PEs [].

Synthetic data generation for teaching transformers reasoning: Synthetic data generation has been explored for optimising model training for reasoning. For example, [, ] use LLM-generated synthetic text for training Phi-1 and Phi-1.5 models and show impressive performance for reasoning-based tasks. [] introduce a novel neuro-symbolic framework to pre-train a transformer model for Olympiad-level math problems. Their method involves generating synthetic geometric data for pre-training, followed by fine-tuning on auxiliary constructions to solve specific theorems. Building on this stream of work, we apply synthetic data generation for teaching causal reasoning.

## 3 Learning Causal Axioms using Transformers

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 axioms. We begin by asking the question “are there any minimal sufficient characterization of causal principles?”.
There has been a fundamental work from Galles and Pearl [] where they axiomatize causal relevance (or equivalently irrelevance). They show that for a given *stable probabilistic* causal model (defined below), there exists a finite set of axioms that completely characterized by axioms of path interception in corresponding directed graphs. We now study how such causal relevance statements can be incorporated into transformer models.

Let $\mathcal{M}=(\bm{X},\bm{U},\mathcal{F})$ be a causal model defined over a set of endogenous variables $\bm{X}$, exogenous variables $\bm{U}$ and the causal relationship between then defined by set of structural equations $\mathcal{F}$ []. Let $\mathcal{G}$ be the causal graph associated with the causal model $\mathcal{M}$ where the nodes $\bm{V}$ in $\mathcal{G}$ correspond to the variables in $\mathcal{M}$ and an edge $V_{i}\rightarrow V_{j}$ between any two nodes $V_{i},V_{j}$ denote the causal relationship between them.

###### Definition 3.1 ( Causal Irrelevance , Defn. 7 in [ 11 ] ) .

$X$ is probabilistically causally irrelevant to Y given Z, written ($X\nrightarrow Y|Z$) iff:
$\mathbb{P}(y|z,do(X)=x)=\mathbb{P}(y|z,do(X)=x^{\prime})\>,\forall x,x^{\prime
},y,z$
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 [] that gives a richer set of finite axiomatization for probabilistic causal irrelevance.

###### Assumption 3.1 (Stability, Definition 9 in [ 11 ] ) .

Let $\mathcal{M}$ be a causal model. Then an irrelevance ($X\nrightarrow Y|Z$) in $\mathcal{M}$ is stable it is shared by all possible probability distribution over $\mathcal{M}$. The causal model $\mathcal{M}$ is stable if all of the irrelevances in $\mathcal{M}$ are stable.

Under the stability assumption (see Assumption 3.1), Galles and Pearl [] characterizes six axioms that completely characterize causal irrelevance (Definition 3.1) or equivalent causal relevance statements after using the corresponding contrapositive statements. An axiom of causal irrelevance is of the form (given in conjunctive normal form):

$$ $\bigwedge_{s}\bigvee_{t}(\bm{X}^{s,t}_{i}\nrightarrow\bm{X}^{s,t}_{j}|\bm{X}^{ s,t}_{k})\implies\bigwedge_{l}\bigvee_{n}(\bm{X}^{l,n}_{i}\nrightarrow\bm{X}^{ l,n}_{j}|\bm{X}^{l,n}_{k})$ $$

where $\wedge$ is “logical and", $\vee$ is “logical or" and for a given $(s,t)$ or $(l,n)$ pair, $\bm{X}_{i},\bm{X}_{j},\bm{X}_{k}$ are disjoint subsets of observed variables $\bm{X}$. In the above causal irrelevance statement, if the antecedent is true, the consequent is also true.

###### Definition 3.1 ( Causal Irrelevance , Defn. 7 in [ 11 ] ) .

###### Assumption 3.1 (Stability, Definition 9 in [ 11 ] ) .

### 3.1 Transitivity Axiom for Causality

Transitivity axiom is intuitive in nature and presented for causal irrelevance in []. In our study, we focus on the transitivity axiom for its simplicity and ease of expressing causal chains and other graph structures for building diverse evaluation sets.
For a stable probabilistic causal model (§3), given variables $X$, $Y$, $Z$ in the system, the transitivity axiom is given by:

$$ $(X\nrightarrow Y|Z)\Rightarrow(A\nrightarrow Y|Z)\vee(X\nrightarrow A|Z) \forall A\notin X\cup Z\cup Y$ (1) $$

which could be further simplified by taking the contrapositive of the expression to derive its causal relevance version which could be given as follows:

$$ $\exists A\notin X\cup Y\cup Z\>\>s.t.\underbrace{(X\rightarrow A|Z)\wedge(A \rightarrow Y|Z)}_{P:\text{premise}}\implies\underbrace{(X\rightarrow Y|Z)}_{H :hypothesis}$ $$

We call the LHS as Premise and the RHS as Hypothesis. Our key idea is that we can use such an axiom to generate thousands of synthetic symbolic expressions that can be used to teach a transformer the specific 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.

### 3.2 Axiomatic Training: Dataset, Loss Function, and Positional Encoding

Training data. Based on a specific axiom, we can map a hypothesis given the premise to its correct label (‘Yes’ or ‘No’). To create a training dataset, we enumerate all possible tuples of $\{(P,H,L)\}_{N}$ where $P$ is the premise, $H$ is the hypothesis and $L$ is the label *(Yes/No)* for a particular setting of the variables $X,Y,Z,A$. Given a premise $P$ based on a given causal graph, if the hypothesis can be derived by applying the specified axiom (once or multiple times), then label $L$ is *Yes*; otherwise, *No*.
For example, suppose the underlying true causal graph of a system has the topology of a chain, $X_{1}\rightarrow X_{2}\rightarrow X_{3}\rightarrow\dots\rightarrow X_{n}$. Then, a possible premise could be $X_{1}\rightarrow X_{2}\wedge X_{2}\rightarrow X_{3}$, and the corresponding hypothesis $X_{1}\rightarrow X_{3}$ will have label *Yes* whereas another hypothesis $X_{3}\rightarrow X_{1}$ will have label *No*. The above axiom could be inductively applied multiple times to generate more complex training tuples.

For our training setup, a synthetic dataset $D$ is constructed with $N$ axiomatic instances generated using the transitivity axiom. Each instance in $D$ is structured in the form of a premise $P$, which is the natural language expression of a causal structure (e.g., “$X$ causes $Y$. $Y$ causes $Z$”), followed by the hypothesis in the form of a question $H_{q}$ (e.g., “Does $X$ cause $Y$?”), which is then followed by the final label $L$ (e.g., “Yes” or “No”).
Each instance in $D$ is structured as $(P_{i},H_{ij},L_{ij})$; $j\in\{1,\dots,\binom{n}{2}\}$where $n$ is the number of nodes in each $i$th premise, thus effectively covering all pairs of nodes in each unique chain of a given causal graph.

Loss function. Given a dataset, the loss function is defined based on the ground truth label for each tuple, represented as $\underset{P,H,L\sim P_{\text{train}}}{\mathbb{E}}-\log(P(L|P,H))$. A preliminary analysis indicated promising results with this loss formulation compared to next token prediction loss.

Positional Encoding. In addition to the training data and loss function, another important factor is the choice of positional encoding. Positional Encoding (PE) play a crucial role of providing information about the absolute and relative position of tokens in a sequence []. [] propose an absolute positional encoding strategy using periodic functions (e.g., sinusoidal or cosine) to initialize these encodings. Absolute positional encoding provides definite values for all positions across any sequence length. However, studies [, ] show absolute positional encoding fails in length generalization tasks for transformers. In the learnable APE variant [], 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 [, ] indicate that removal of PEs in auto-regressive models can improve model’s length generalization capabilities, wherein the attention mechanism during auto-regressive decoding is sufficient to encode positional information. We experiment with different positional encoding to understand their impact on generalization in causal tasks: learnable position encoding (LPE), sinusoidal positional encoding (SPE), no positional encoding (NoPE).

### 3.3 Data Perturbation: A Key to Model Generalization

Variability or diversity in training data in the form of perturbation helps aid model generalization []. For axiomatic training, we provide structured perturbation to extend our transformer model across complex structures for which it was not explicitly trained on. We introduce perturbations at multiple levels in the training data to maximize diversity in the distribution of the training set, as explained below.

- 1.
Node names: Each node in the transitivity chain is represented by an alphanumeric name comprising 1-3 characters. The length of a name and the specific characters are randomly selected during data generation.
- 2.
Causal Graph Topology: We consider two main types of causal graphs for the training set.
2.1.
Sequential: All causal edges are directed forward, thus forming a typical transitivity chain, e.g. X $\rightarrow$ Y $\rightarrow$ Z.
2.2.
Random Flipping: Given a chain of sequential nodes, we randomly reverse some edges creating complexity by disrupting direct paths between subsequent nodes (eg. X $\rightarrow$ Y $\leftarrow$ Z). This can be expressed simply through natural language like: “X causes Y. Z causes Y."
- 3.
Length level: To facilitate transformers understanding of the axiom, we incorporate chains of varying lengths, ranging from 3 to 6 nodes in our training set.

Random flipping introduces forks and colliders, which form the building blocks of any causal DAG. This helps incorporate complexity in model training, thus aiding its capability to generalize across multiple structures.

### 3.4 Assessing Axiomatic Learning in Transformers

To evaluate if a trained model has learnt the correct understanding of an axiom instead of relying on shortcuts or correlation-based features, designing an out-of-distribution (OOD) evaluation set is important. To tackle this, we evaluate our model across multiple types of complex structures that are not seen during training. We divide the structural complexities as follows:

- 1.
Length: Evaluating whether our model accurately infers causal relationships for sequences or chains (both sequential and ones with random flipping) longer than those in the training set.
- 2.
Node Name Shift: Testing the model’s performance with longer node names, increasing from 1-3 characters used in the training set to 8-10 characters. Findings of [] show how shift in node names used in finetuning of language models result in generalization failure in inferring causal relationships from correlational relationships even though the model had impressive performance on in-distribution node names, further motivating this analysis.
- 3.
Order of Chains:
a) Completely reversed chains: This evaluation is inspired by the reversal curse [] that revealed generalization failure of LLMs in answering questions in reversed sequences despite knowing the answers in the original order. We evaluate the capacity of axiomatic training to enable reasoning over reversed chains even when not explicitly trained on any completely reversed chains. A completely reversed chain will be of the form X $\leftarrow$ Y $\leftarrow$ Z with its natural language representation as: “Y causes X. Z causes Y.", where $X,Y,Z$ are replaced by random alphanumeric names.
b) Shuffling of Sequences: Causal sequences with random edge flips, as defined in 3.3 represented by natural language statements sequentially (A causes B. B causes C …), are shuffled to add complexity and break sequential order. This tests transformers’ ability to infer accurate relationships regardless of sequence order of premise.
- 4.
Branching: Causal graphs with dense branching pose a challenging evaluation task. While the training set comprises simplistic linear sequences, this evaluation setup involves multiple branches, colliders, forks, and chains in one network, with significantly high complexity. We measure complexity of a graph using its branching factor: Number of edges/Number of nodes.

Unlike length and node name generalization, the reversal and branching evaluation setups change the causal structure and hence better evaluate whether the model has learnt accurate representations for causal structure. Branching is perhaps the most challenging since it contains new structures (due to a higher branching factor) that were unseen during training.

Figure: Figure 1: Evaluating structural generalization of transformers through axiomatic training. We train a transformer on two simple causal structures: chains and chains with random flipping of some edges. All training instances consist of 3-6 nodes. The trained model is evaluated on significantly more complex structures: bigger causal graphs with >6 nodes, more branched networks with higher average in-degree and out-degree, complete reversals, longer sequences, shuffled natural language statements of sequences and longer node names.
Refer to caption: extracted/5722718/figures/CausalAxioms.png

## 4 Applying Axiomatic Training to Learn Causal Transitivity Axiom

### 4.1 Training and Evaluation Datasets

Based on the above discussion, a dataset can be characterized by relative frequency of different types of causal structures such as linear, completely reversed sequences, and so on. Given a set of causal sequences $X=\{X_{1},\dots,X_{n}\}$ where each causal sequence $X_{i}$ takes the form: $X_{i}=[e_{jk}\mid j>0,k>0]$. Here $e_{jk}$ represents an edge between node $j$ and $k$ in the $ith$ causal sequence (can also be represented as $X_{i,j}$, $X_{i,k}$) such that each $e_{jk}$ translates to "$X_{i,j}$ causes $X_{i,k}$." in natural language. $V_{i}$ = $\{X_{i,1},X_{i,2},...,X_{i,m}\}$, represents set of vertices in causal sequence $X_{i}$.

Let $f_{dim}$ represent the maximum value for a given perturbation dimension $dim$, along which we construct train and evaluation sets for our axiomatic framework. For each dimension, we choose a threshold $\tau_{dim}\in L$, such that $f_{dim}<\tau_{dim}$ forms our training set and $f_{dim}\geq\tau_{dim}$ forms the evaluation set. So, $f_{dim}\in\{f_{len},\ f_{branch},\ f_{nodelen},\ f_{revfactor},\ f_{shuffle}\}$
where:

- •
$f_{len}=\max_{\forall i}(len(V_{i}))$, gives the maximum number of nodes across all causal sequences. $\tau_{len}$ for length is set at 6, with $f_{len}\in[3,6]$.
- •
$f_{branch}=\max_{\forall i}(|X_{i}|/|V_{i}|)$ gives the maximum branching factor in a dataset, with $\tau_{branch}=$ 0.8 (for 6 node linear sequences). For sequences in the train set, the branching factor ranges from 0.6 to 0.8 for 3 to 6 length sequences.
- •
Let $l_{i,j}$ be the length of the name of the node $X_{i,j}$, then $l_{i,j}=(len(X_{i,j}))$.Therefore, the maximum length of node names across all nodes in all causal sequences can be represented as: $f_{nodenamelen}=\max_{1\leq i\leq n,\,1\leq j\leq m}l_{i,j}$. We set $\tau_{nodelen}$ for train set as 3, with $f_{nodelen}\in[1,3]$.
- •
Given any causal sequence $X_{i}$ and a function $N$, where $N(X_{i,j},X_{i,j+1})$ returns natural language representation of a directed edge between $j$ and $j+1$ node in the causal chain $X_{i}$. $f_{shuffle}=\cap_{\forall i,j}\text{Perm}(N(X_{i,j},X_{i,j+1})),$ where $N(X_{i,j},X_{i,j+1})$ represents deviation from original sequential order of natural language sentences to represent $X_{i}$.
- •
Given a causal sequence $X_{i}$ and let $R(X_{i},f_{revfactor})$ be an operation on the causal chain that flips the direction of every edge in the sequence with probability $f_{revfactor}$. In the training set, there is a directed edge between every sequential pair of nodes $X_{i,j},X_{i,j+1}$ with $f_{revfactor}=0$ (for linear sequence, $X_{i,j}\rightarrow X_{i,j+1}$) or $0.5$ (for sequence with random flipping, $X_{i,j}\rightarrow X_{i,j+1}$ or $X_{i,j}\leftarrow X_{i,j+1}$) In the evaluation set $f_{revfactor}=1$ 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.

#### 4.1.1 Training Datasets

Using a straightforward transitivity chain-based training setup, we aim to understand how variability or noise, introduced through randomly flipped edges, varying sizes of node names, and different chain lengths, enhances the generalization abilities of our transformer model to handle longer, branched, and reordered complex settings. We perform multiple ablations with different training sets to understand possible contributing factors towards the model’s generalization. Our training setup consists of around 175k instances of sequential chains with size of chains ranging from 3 to 6 nodes. We use three versions of training data to evaluate the impact of different noise perturbations, and the impact it has across diverse evaluation setups when a transformer is trained from scratch using that. Each of our training set is a balanced set with equal number of instances from both label class (‘Yes’ and ‘No’).
Our three sets can be described as follows.

- •
Training Setup 1 (TS1): This setup comprises of 73k chains with random flipping and 101k sequential linear 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. However, only when evaluating for reversal chain test setup, we remove all the reversed chains from this training set and re-train the model from scratch with the remaining sequential and randomly flipped chain sequences (to ensure that the model has not seen any completely reversed chain in its training).
- •
Training Setup 2 (TS2): This setup comprises of more simple sequential chains (132k), while we decrease chains where random flipping is done (42k), to keep the overall number of training instance constant (around 175k).
- •
Only Causal Chains (OCC): This set comprises of sequential transitivity chains without any edge being randomly flipped. This helps us understand whether addition of noise in the form of randomly flipping of edges is helpful or not for generalization and to what extent.

#### 4.1.2 Evaluation Datasets

To evaluate the impact of variability in causal sequences on training transformers across different OOD structural dimensions as described in Section 3.4, we design different evaluation sets encompassing natural language representation of causal structures with varying levels of complexity. This helped us evaluate whether the model is capable of learning accurate structural understanding and application of transitivity axiom to more complex networks. Similar to train sets, the evaluation sets have equal number of instances from both label classes. Our evaluation sets to assess generalization capabilities of the model could be described as follows:

- •
Length Generalization EvalSet: Testing on causal sequences with length $>$6 upto 15, longer than any sequence encountered by the model in training set. Length generalization is evaluated for both sequential chains and chains with randomly flipped edges.
- •
Node Name EvalSet: Assessed model’s generalization capabilities to sequences with longer node names, increasing from 1-3 characters used in training set to 8-10 characters. To add onto the complexity, we also include sequences longer than any sequence in the train set ($>$6) upto 9 length.
- •
Reversal EvalSet:
Evaluated performance of our transformer model, with no completely reversed sequence in its training, on reversed causal sequences. Sequences upto 6 length were evaluated.
- •
$\mathbf{MultiEval_{SLR}}$ (Shuffling + Random Flipping + Length Sequence): This setup involves evaluation on 3 levels of complexities together: shuffling of sentence for representing the sequences, each sequence having random flipping, and some sequences having longer length than sequences in training set (upto 9).
- •
Branching EvalSet: One of the most complex evaluation setups, with dense networks containing multiple branches, colliders and forks. While each sequence in the training set had values of 1-2 for both in-degree and out-degree across all nodes, in this setting a node can have maximum value of $n-1$ for both, and minimum of 0 creating more complicated structures than the ones transformer had encountered during its training. To add onto the complexity we evaluate on structures with more nodes (8,10,12), than any unique causal sequence in the training set besides 5 node networks. We evaluate multiple densely branched networks constructed using the Erdös-Rényi model, where we provide number of edges and nodes in accordance to the values of branching factor (1.4 and 2) we use for evaluation. We implement this using igraph package in python [] to get different unique graphs with required branching factors for evaluation.

### 4.2 Implementation Details: Architecture, Tokenizer and Training Procedure

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. The model is trained from scratch on each of our training datasets. To understand the effect of Positional Encodings (PE), we consider Sinusoidal PE (SPE) [], Learnable PE (LPE) [] and having no PEs (NoPE)  [, ]. All models are trained for 100 epochs using the AdamW optimizer with 1e-4 learning rate.

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.

### 4.3 Baselines Using Existing LLMs

Given recent work on how LLMs can be leveraged for causal reasoning [, , ], we include language models such as GPT-4 (gpt-4-32k) [], Gemini (gemini-pro) [] and Phi-3 (Phi-3-mini-128k-instruct) [] as baselines. Note that each of these models is significantly larger than our model and known to perform well on reasoning tasks, with the smallest baseline model Phi-3 having 3.8 billion parameters.
We incorporate both commercial (GPT-4 and Gemini Pro) and open-source (Phi-3) models covering a range of size and capabilities.
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 (Hypothesis) and prompt the LM to answer it in either ‘Yes’ or ‘No’ (Label). Here is an example prompt: “EX causes T. T causes 9. 9 causes W. W causes 7. 7 causes M. M causes a. Does EX cause T? Answer in ‘Yes’ or ‘No’ only.” A1 contains some examples of all types of prompts and instances used for querying as well as training for different structural types.

## 5 Results

### 5.1 Generalization to Complex Causal Scenarios

We present results on how well an axiomatically trained transformer can generalize to larger and more complex causal graphs, and how it compares to pre-trained LLMs.

**Table 1: Accuracy on longer linear sequential chains with all edges in forward direction (Forward Sequence, FS) and sequences with randomly flipped edges (Random flipping, RF). Trained model remains the same across all setups. For longer chains, NoPE performs best on sequential linear setup. Random baseline accuracy = 0.5.**
| Model | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| FS | RF | FS | RF | FS | RF | FS | RF | FS | RF | FS | RF | FS | RF | FS | RF | FS | RF |  |
| Baselines |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| GPT-4 | 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.63 | 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 |
| Phi-3 | 0.81 | 0.85 | 0.96 | 0.85 | 0.85 | 0.85 | 0.87 | 0.89 | 0.90 | 0.86 | 0.84 | 0.85 | 0.91 | 0.84 | 0.90 | 0.80 | 0.78 | 0.85 |
| Axiomatic Training |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| TS1 w NoPE | 1.00 | 0.99 | 0.95 | 0.96 | 0.88 | 0.89 | 0.76 | 0.88 | 0.73 | 0.90 | 0.77 | 0.92 | 0.61 | 0.82 | 0.67 | 0.78 | 0.68 | 0.81 |
| TS1<br>w LPE | 0.98 | 0.96 | 0.92 | 0.97 | 0.77 | 0.90 | 0.59 | 0.87 | 0.57 | 0.86 | 0.57 | 0.84 | 0.55 | 0.73 | 0.51 | 0.76 | 0.50 | 0.68 |
| TS1 w<br>SPE | 0.99 | 0.95 | 0.95 | 0.94 | 0.86 | 0.76 | 0.80 | 0.75 | 0.76 | 0.79 | 0.84 | 0.68 | 0.79 | 0.63 | 0.85 | 0.65 | 0.77 | 0.69 |
| TS2 w NoPE | 1.00 | 0.98 | 0.99 | 0.97 | 0.92 | 0.91 | 0.88 | 0.90 | 0.86 | 0.92 | 0.95 | 0.90 | 0.96 | 0.83 | 0.81 | 0.84 | 0.85 | 0.78 |
| TS2 w LPE | 1.00 | 0.98 | 0.88 | 0.97 | 0.80 | 0.88 | 0.62 | 0.92 | 0.66 | 0.91 | 0.64 | 0.81 | 0.65 | 0.75 | 0.62 | 0.75 | 0.62 | 0.77 |
| TS2 w SPE | 0.95 | 0.93 | 0.81 | 0.84 | 0.56 | 0.34 | 0.50 | 0.38 | 0.50 | 0.44 | 0.51 | 0.57 | 0.46 | 0.74 | 0.52 | 0.75 | 0.50 | 0.77 |
| OCC w NoPE | 0.98 | 0.58 | 0.79 | 0.49 | 0.86 | 0.51 | 0.92 | 0.49 | 0.72 | 0.57 | 0.90 | 0.50 | 0.81 | 0.52 | 0.84 | 0.52 | 0.83 | 0.46 |

Length Generalization:
Table 1 shows accuracy of different models when evaluated on larger causal chains that were not seen during training. Among the baseline pre-trained LMs, GPT-4 obtains the highest accuracy on both standard and randomly flipped chains. It is remarkable that our TS2 (NoPE) model obtains competitive performance to the trillion-scale GPT-4 model, even though it had never seen larger sequences during training. In particular, for chains of size 7-13, TS2 (NoPE) obtains higher or comparable accuracy than GPT-4 across the standard and randomly flipped chains. Its accuracy decreases for chains of length 14-15 (0.85 for standard chains and 0.78 for randomly flipped chains) but is still significantly higher than that of LMs like Gemini-Pro and Phi-3. Note that a random prediction would yield a 50% accuracy, indicating that the axiomatically-trained TS2 (NoPE) model can generalize its reasoning to causal chains much longer than 6 even though it was trained only on chains upto length 6.

Node Name Shift: For models trained on TS2 dataset, we also evaluate generalization to changes in variable names (Figure 4). We find that TS2 (NoPE) is 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.

Order of Causal Sequences: We now consider how variations in the causal structure impact generalization of axiomatically-trained models. In Table 2(b), we consider the complex evaluation setup $\mathbf{MultiEval_{SLR}}$ that includes shuffled order of causal sequences with random flipping for increasing length (even beyond the ones in train set). On this task, TS2 (NoPE) obtains higher accuracy than Gemini Pro and Phi-3 on chains of length up to 8. At length 9, TS2 (NoPE) obtains 0.73 accuracy which is comparable to Gemini Pro (0.74) and significantly better than random.

We observe a similar pattern for evaluation on completely reversed sequences in Table 2(a).

Figure: Figure 2: Accuracy on completely reversed sequences. Our Model trained on TS2 with NoPE performs best and at par with GPT-4, while other LLMs like Phi-3 and Gemini Pro perform poorly, even when the sequence length is small (upto 6).
Refer to caption: x1.png

This is an extreme case of out-of-distribution data since most causal edges are left-to-right in the training data whereas the test data contains all right-to-left edges.
On this task, our axiomatically trained model TS2 (NoPE) outperforms GPT-4 when restricted to chain lengths of 3-6. In particular, its accuracy ($0.94$ for chains of length $6$) is substantially higher than Gemini Pro and Phi-3 ($0.62$ and $0.69$ respectively).

Branching: Finally, we consider the hardest evaluation task involving non-linear chains where we introduce general Erdos-Renyi graphs as the causal sequences while the training data contains only linear chains. Here the length of sequence corresponds to the number of nodes in the graph and we study the performance differences as the branching factor is varied. While GPT-4 obtains the best accuracy across increasing graph sizes, our TS2 (NoPE) model obtains higher accuracy than Gemini Pro for all graph sizes except one. Even when evaluated on graphs with 12 nodes and 1.4 branching factor, the TS2 (NoPE) model obtains 70% accuracy, significantly better than random (50%). Note that the training data only included graphs with a branching factor of 1.

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 (Fig. 2). 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. This suggests the potential of axiomatic training for reasoning over causal graphs.

### 5.2 Additional Results: Role of Data Diversity and Positional Encoding

**Table 2: Reversal and Shuffling**
| Model | 3 | 4 | 5 | 6 |
| --- | --- | --- | --- | --- |
| Baselines |  |  |  |  |
| 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 |
| Axiomatic Training |  |  |  |  |
| TS1 w NoPE | 0.98 | 0.99 | 0.92 | 0.91 |
| TS1 w SPE | 1.00 | 0.99 | 0.99 | 0.97 |
| TS2 w NoPE | 0.99 | 0.99 | 0.95 | 0.94 |
| TS2 w SPE | 0.98 | 0.97 | 0.93 | 0.94 |
| TS2 w LPE | 0.99 | 0.98 | 0.95 | 0.97 |
| OCC w NoPE | 0.33 | 0.18 | 0.10 | 0.09 |

**Table 3: Evaluated on branched graphs created using Erdös-Rényi model, with varying branching factors (calculated by number of edges/number of nodes). Decoder model remains the same across all setups (67 Million parameter), accuracy metric is used, with random baseline performance = 0.5**
| Model | 5 | 8 | 10 | 12 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| BF=2 | BF=1.4 | BF=2 | BF=1.4 | BF=2 | BF=1.4 | BF=2 | BF=1.4 |  |
| Baselines |  |  |  |  |  |  |  |  |
| GPT-4 | 0.98 | 0.95 | 0.91 | 0.90 | 0.84 | 0.88 | 0.82 | 0.86 |
| Gemini Pro | 0.77 | 0.74 | 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 |
| Axiomatic Training |  |  |  |  |  |  |  |  |
| OCC w NoPE | 0.52 | 0.51 | 0.53 | 0.52 | 0.52 | 0.55 | 0.49 | 0.47 |
| TS1 w LPE | 0.79 | 0.84 | 0.71 | 0.76 | 0.68 | 0.69 | 0.65 | 0.65 |
| TS1 w SPE | 0.72 | 0.79 | 0.63 | 0.64 | 0.56 | 0.61 | 0.52 | 0.59 |
| TS1 w NoPE | 0.77 | 0.84 | 0.73 | 0.76 | 0.68 | 0.70 | 0.62 | 0.66 |
| TS2 w LPE | 0.72 | 0.80 | 0.61 | 0.71 | 0.62 | 0.63 | 0.56 | 0.63 |
| TS2 w SPE | 0.52 | 0.70 | 0.49 | 0.49 | 0.49 | 0.49 | 0.51 | 0.52 |
| TS2 w NoPE | 0.83 | 0.86 | 0.74 | 0.77 | 0.69 | 0.74 | 0.64 | 0.70 |

Figure: Figure 3: Evaluating generalization on causal sequences with longer node names (than the ones used in sequences in train set), and the impact of different PEs for TS-2 training set, which yields the best performance
Refer to caption: x2.png

Role of Positional Encodings. Comparing the performance of models wrt. choice of positional encoding, we find that models with no positional encoding generalize well to both longer lengths (upto chain length of 15) and complex, unseen graph structures, even though they are only trained on chains over 3-6 nodes. Models with SPE and LPE also perform well on longer chains but poorly when node names’ length increases, even for chains with smaller number of nodes (Figure 4). This generalization failure using SPE and LPE highlights the models’ inability to handle minor perturbations in the sequences present in their training set. Moreover, SPE does not perform well across different structural dimensions like branching, and order-based settings such as shuffling and reversal. Learnable PE performs well up to 9-length linear chains but drops sharply afterward. Overall, our results extend earlier work on the utility of NoPE [, ] to the task of understanding causal sequences and generalizing to both longer length and complex structure at test time.

Note that for the evaluation setup with randomly flipped edges, all types of PEs perform well. This may be due to the fact that the maximum effective length of a directed path between any two nodes is very small in randomly flipped sequences (much smaller than that of sequential causal chains) because the probability of an edge being forward directed is 0.5, and as we move across all sequential edges, the probability keeps on decreasing as $0.5^{l}$ where l is length of path.

Importance of Data Perturbations:
In addition to the positional encoding, diversity of the sequences in train data plays an important role. Model trained on only causal chains (OCC) generalize to longer chains (Table 1) but not to other DAG structures (see Figure 4 for edge flip, Figure 2 for reversal, Table 3 for branching). Models trained on TS1 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).

## 6 Inferring Causation from Correlation using Axiomatic Training

While the above study evaluated transformers’ capability to generalise the transitivity axiom from small causal chains to large graphs, we now study whether this capability transfers to other causal tasks.

To this end, we apply axiomatic training to a task on inferring causation from statements about correlation in observational data []. As Figure 5 shows, each data instance includes correlational relationships described in natural language for graphs with 3 to 6 nodes; and the goal is to infer the truth value of a hypothesis (of 6 different types: Parent, Ancestor, Child, Descendant, Collider, Confounder) on whether there exists a direct or indirect relationship between any given nodes, and possible presence of colliders and confounders.

This task is significantly harder than applying the transitivity axiom. First, there are multiple hypothesis types to evaluate: direct effect, indirect effect, children, ancestors, colliders and confounders. Second, solving the task requires an understanding of d-separation and the Markov property. Specifically, it involves mapping correlational statements to multiple possible causal graphs and determining if the query is satisfied across all graphs in the Markov Equivalence Class.

Figure: Figure 5: Example instance from [] dataset describing correlational relationships of a 4 node system followed by a hypothesis to be verified.

**Table 4: Extension to ‘Correlation to Causation’ Experiments from []. Results indicate that transformer models trained on correlational statements of smaller causal structures demonstrate some generalization capabilities when tested on more complex causal structures. These models outperform significantly larger baseline models like GPT-4**
| Model | Precision | Recall | F1 Score | Accuracy |
| --- | --- | --- | --- | --- |
| Ours | 0.72 | 0.50 | 0.59 | 0.64 |
| Phi-3 | 0.52 | 0.60 | 0.56 | 0.52 |
| Gemini pro | 0.52 | 0.59 | 0.55 | 0.52 |
| GPT-4 | 0.59 | 0.50 | 0.54 | 0.58 |

##### Model Training:

We use the same model architecture as in Section 4.2 and train our model from scratch for 100 epochs using NoPE, since it performed consistently well across diverse OOD settings in our transitivity based experiments. For creating a train set, we consider the subset of the original dataset with correlational statements for graph consisting of 3 and 4 nodes. As the test set, we evaluate the model’s performance directly on 5 node correlational statements.

To aid generalization, we take inspiration from our transitivity-based experiments and create different combinations of randomly created alphanumeric node names. We then derive a training set from the original dataset by instantiating the correlational statements with different combinations of alphanumeric node names. We balance the dataset by sampling equally from both classes to avoid bias in our transformer model to get a train set with 113k instances. Then, we create a test set with 1000 randomly sampled instances of correlational statements for 5-node graph networks. Since the correlational statements are not simplistic unlike the premise from our transitivity experiments, we tokenize at the character level for nodes. For a straightforward evaluation, we tokenize all input text at the token level and use the same node names for evaluation as in the training set to avoid potential out-of-vocabulary issues.

##### Comparison with Baselines:

As reported in [], due to the complexity of the task, we find that pre-trained LMs such as Gemini Pro and Phi-3 perform similar to a random guess (52% accuracy, see Table 4). While GPT-4 does perform slightly better, it’s performance is still low (58% accuracy). Remarkably, our small transformer model performs better than all baselines with 64% accuracy; 6% points higher than GPT-4. With further exploration of different training setups, axiomatically-trained transformer models may be optimised further for such causal reasoning tasks.

## 7 Discussion and Conclusion

In this work, we use the formal axiomatic framework proposed by [] to study whether transformers can be taught causal axioms. To this end, we provide a method to create training data containing diverse demonstrations of an axiom and explore modelling choices to learn the axiom. We find that a transformer model trained from scratch can learn how to apply an axiom based on a large training dataset containing axiomatic demonstrations. We evaluate the proposed axiomatic training approach on two causal tasks—graph traversal using the transitivity and inferring causal relationship from correlational statements. For these tasks, small 67M transformers trained using our approach can generalize over longer and complex complex graph structures that they never saw during training, often providing accuracy better than existing LLMs such as GPT-4, Phi-3 and Gemini Pro.

##### Applicability to Causal Tasks.

While our current work focuses on the transitivity axiom for causal relevance, extending the work to other causal axioms from [] is an interesting research direction. In addition, we may consider other axioms that are relevant for downstream tasks such as effect inference. For example, if a transformer model can be trained to validate the d-separation rule—given two variables X and Y, are they independent given a variable set Z?—then repeated applications of the rule can be used to derive a backdoor set.
Another interesting direction is to extend the training approach for both deterministic and probabilistic causal models.

##### Generalization to Logical Reasoning.

While we focused on causal reasoning axioms, our axiomatic training approach is general and can be applied to any formal system based on axioms. For instance, the same axiomatic training procedure can be used for teaching LMs logical reasoning tasks such as deductive reasoning. For instance, recent work [] evaluates the deductive reasoning capabilities of LLMs and measures their generalization abilities along depth, width, and compositional abilities. As the depth increases, performance of LLMs deteriorates. It will be interesting to see whether axiomatic training can be applied to learn deductive reasoning axioms and improve the reasoning abilities of LMs.

##### Implications for Training Language Models.

GPT-4 shows impressive generalization on the causal tasks we evaluated, even though it was not trained specifically for these tasks. A fundamental question is how LLMs like GPT-4 achieve the capability to reason over causal graphs. Based on our preliminary experiments, we conjecture that axiomatic training may be one explanation for GPT-4’s ability. While GPT-4 was not trained specifically for the causal tasks, the underlying axioms may have been present in diverse forms in its web-scale training data (e.g., in books, blogs, and technical forums) that helped it learn the causality rules.

At the same time, performance of other LMs such as Gemini Pro and Phi-3 can be improved. For instance, inability of Gemini Pro and Phi-3 to reason over completely reversed chains (even when their length was small) highlights how zero-shot reasoning for causal tasks is missing in these models. Incorporating causal axiom demonstrations as a part of language models’ pretraining could help improve the reasoning of these models, so that even small language models like Phi-3 can perform comparably to GPT-4 for causal tasks.

## 8 Acknowledgements

We would like to thank Kabir Ahuja for his helpful insights and suggestions for this work.

## Appendix

**Table A1: Table with examples of data instances of different causal structural networks used for training and evaluating models. Each instance is broken down into premise, hyopthesis, and label. During evaluation, only the premise followed by the corresponding hypothesis is provided, whereas during training of transformer, the model is trained on the loss of prediction of the label token.**
| Query Type (Train/<br>Eval) | Data Instance Example<br>(Premise-Hypothesis-Label) | Structure Type | Network Size (number of nodes) |
| --- | --- | --- | --- |
| Train | Mhb causes iqB. iqB causes G. Does G cause iqB?: No | Short Linear Sequence | 3-6 |
|  |  |  |  |
| Train | N5w causes s. 6D causes s. Does N5w cause s?: Yes | Short Sequence with Random Flipping | 3-6 |
|  |  |  |  |
| Eval | w3 causes ROv. w3 causes tQC. H causes ROv. H causes tQC. b causes ROv. b causes w3. b causes H. Does tQC cause ROv?: No | Branching | 5,8,10,12 |
|  |  |  |  |
| Eval | LKk causes 5Ov. Kk causes L0. L0 causes KWO. 5Ov causes c. Does KWO cause L0?: No | Shuffled Sequences | 3-9 |
|  |  |  |  |
| Eval | FDAH26mV7 causes 7tzaIHjlY. 7tzaIHjlY causes 0kspcX95Im. 0kspcX95Im causes 7rhFSlx2o9. 7rhFSlx2o9 causes 1PlG5LHVqp. Does FDAH26mV7 cause 7tzaIHjlY?: Yes | Sequences with Longer Node Names | 3-9 |
|  |  |  |  |
| Eval | r causes rZ. rZ causes L. L causes bUx. bUx causes Pbr. Pbr causes 1w. 1w causes c3. c3 causes yBQ. yBQ causes yK. yK causes w. w causes P. P causes kH. kH causes 1u. 1u causes jV7. jV7 causes i. Does r cause rZ?: Yes | Long Linear Sequences | 7-15 |
|  |  |  |  |
| Eval | rU6 causes eF. eF causes ivC. 3R causes ivC. 3R causes A8. 2 causes A8. 2 causes i. i causes a03. y causes a03. b 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 Flipping | 7-15 |
|  |  |  |  |

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: + Could the authors contextualize the uses for axiomatic training in the broader causal inference and decision making context to help shed light on its significance as a training technique?

+ Can you clarify if the transitivity axiom in (1) is linked to the previous logical rules by being one of the six axioms you reference? 

+ Can you clarify how examples are generated for section 5, given that many graphical relations in a causal model are not identified by conditional independences alone?

+ Can you provide some more plain English context about the transitivity axiom and its importance? It seems like it's a statement about the transitivity of indirect causation, where causation is codified through causal (ir)relevance. Is that fair to say?
soundness: 3
strengths: + With the exception of a few expositional sentences, the paper is clear about definitions, training, and the experimental setup.

+ To the best of my knowledge, the paper considers a novel setup for training Transformers to apply logical rules, at least in the context of causal reasoning.

+ The empirical setup is thorough, identifying four ways of generating out-of-distribution premises at inference time and teasing apart the impact of multiple aspects of the training setup.
summary: This paper introduces a way of training Transformers from scratch to learn causal reasoning axioms. The idea, called axiomatic training, is to demonstrate axioms through labeled examples that encode whether or not a hypothesis is true given a premise. For example, the premise can be conditional independences entailed by a causal graphical model, the hypothesis, whether or not some variable is a collider. To study axiomatic training, the paper considers the transitivity of causal relationships as an axiom and develops an extensive experimental design that evaluates various training strategies on different forms of generalization at test time. The paper consistently finds that their trained Transformers outperform many LLMs, while performing comparably to GPT-4. As a second application, they use axiomatic training to study performance on queries about additional graphical relationships such as colliders, direct or indirect causes and find that axiomatic training again generalized favorably compared to several LLMs.
weaknesses: + At the highest level, it would be useful to get more perspective on the significance of training Transformers from scratch to automatically learn to resolve logical causal statements. First, the empirical findings suggest that GPT-4 is pre-trained to handle these queries in generalizable way already. Second, even if axiomatic training were improved to outcompete GPT-4 on these structured queries, how can this technique be incorporated into the larger decision making context? Much of text-based decision making won't feature explicit premises like conditional independence statements but rather informal natural language.  Relatedly, although the paper does summarize related work, the authors can better contextualize the significance of their contribution in the broader context of LLM-based causal reasoning. They discuss some papers, and use the Corr2Cause benchmark -- in what way does this training procedure build on the promise of LLMs answering causal questions factually?

+ I found some sentences hard to process, potentially because they were using jargon and could have been put more simply. For example, in lines 44-45, in the context of causal models, the data-generating process arguably entail the axioms -- it's not that the axioms of d-separation entail a causal model -- and so the phrase "data that is result of axioms obeyed by a data-generating process" is hard to digest. In section 3 in line 131, it would have been useful if phrases like "causal principles" were explained better before being invoked. In lines 134-135, "finite set of axioms ... completely characterized by axioms of path interception ..." is again hard to parse. 

+ Relatedly, it's hard to link the transitivity axiom in (1) back to the logical formula given above it -- is the transitivity axiom supposed to be one of the six axioms you reference? 

+ Some technical details are missing. In assumption 3.1, $\mathcal{M}$ is conventionally used to define a structural causal model (SCM). Is that what you're referring to? If so, the paper ought to define a (probabilistic) SCM, and clarify phrases like "all possible probability distribution(s) over $\mathcal{M}$", which again don't parse. In section 5, when applying axiomatic training to identify additional graphical relationships, it's not clear how the examples are created since many graphical relationships (e.g., direct causation) are not identified from conditional independence statements. Thus, for many premises and hypothesis pairs, there is no way to arrive at a yes or no label. How is this handled?

+ Section 3 feels like it is both missing some details (see above) but also introducing a lot of notation without giving context. In particular, the build up of the conjunctive normal form culminates in Equation (1), which is never explained simply in plain English. But looking at it, it seems to just say that (indirect) causation is transitive. There is a big notation overhead to arrive at a fairly simple idea, which makes the lead-up feel like it's purposefully trying to create an air of complication that's not needed.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Could you provide sensitivity analyses for:
- Training epochs
- Model size variations
- Different positional encoding schemes
2. How do you ensure fair comparison with baselines given the synthetic nature of the training data?
3. Have you considered incorporating semantic understanding mechanisms to enhance the model's causal reasoning capabilities?
soundness: 3
strengths: 1. Novel framework for teaching causal reasoning through symbolic demonstrations
2. Comprehensive evaluation across multiple generalization dimensions
3. Strong empirical results showing competitive performance with much larger models
4. Clear potential for extending to other axioms and reasoning tasks
summary: This paper introduces an "axiomatic training" framework for teaching transformers causal reasoning through symbolic demonstrations of causal axioms. The key contributions are:

1. A method to generate synthetic training data containing demonstrations of causal axioms
2. Empirical evidence that a 67M parameter transformer trained on simple causal chains can generalize to more complex causal structures
3. Application to both transitive causal reasoning and inferring causation from correlation statements
4. The authors demonstrate that their model often outperforms larger language models like GPT-4 and Gemini Pro on these causal reasoning tasks.
weaknesses: 1. Experimental Design Limitations:
- The fixed training regime of 100 epochs lacks justification and sensitivity analysis
- Insufficient ablation studies on model size impact
- Limited exploration of different positional encoding schemes' effects on performance
2. Baseline Comparison Issues:
- The baseline methods may not adequately handle noise and "synthetic data"
- Potential fairness concerns in comparisons due to the presence of contradictory causal relationships in training data
- The baseline methods might need adjustment to better align with the causal reasoning tasks
3. Semantic Consideration:
- Insufficient attention to the role of semantic information in causal reasoning tasks
- Limited exploration of how semantic understanding could enhance model performance
- Potential unexplored opportunities in leveraging Transformer's semantic capabilities

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Questions:
- Recent research into decision transformers showed their significant issues with trajectory stitching (see e.g. Zhou et al. 2024, "Free from Bellman Completeness"). Do you think it has any relation to your problem of computing graph reachability?
 - Why do you think the model truly generalised in Sec 5, and not just remembered some patterns, given the very effective dataset size?
 - What is the definition of the correlational statement, and what is the complexity of the algorithm that answers it?
 - Do you think tokenizer issues might have impacted off-the-shelf LLMs abilities to perform on the task?

Small things:
 - Citations in line 347 and label in line 914 are not displayed correctly
 - Colours are inverted in Fig A1.
soundness: 2
strengths: - The problem of reasoning, and, specifically, causal reasoning, is important and not very well understood yet, which makes the paper topical and interesting.
- The axiomatic training method suggested in the paper seems a priori reasonable - LLMs get access to many demonstrations of an axiom, but  augmenting the data in specific ways might yield better results. The fact that the model performs better, or on par, with off-the-shelf, cutting-edge LLMs is evidence that the method has some potential to improve training.
- Empirical testing of various levels of generalisation for the transitivity axiom seems solidly done, including retraining the model on a filtered data to test specific claims.
summary: The paper investigates the problem of causal reasoning in transformer models. Specifically, authors try to come up with a different training protocol for enabling causal reasoning, which they call axiomatic training. The dataset for axiomatic training consists of a large number of demonstrations of (abstract) causal reasoning - in this work, primarily based on one specific axiom, transitivity. One such demonstration could look like "X causes Y, Z causes Y. Does Z cause X?" followed by an answer "No". After generating a large dataset of those types of prompts, varying lengths, variable names and causal graphs topology, a small transformer model is trained to predict the last token on different subsets of the dataset, to check what levels of generalisation it would reach. 
Trained model is compared with off-the-shelf LLMs (GPT-4, Phi-3, Gemini). Empirical comparison is extended further by looking at more general correlational statements from Corr2Cause dataset. The paper also investigates the impact of positional encoding on generalisation abilities. Given the results, authors suggest that similar techniques might inspire specific ways of synthetic data augmentation and help with LLM reasoning.
weaknesses: - Writing was a bit disjointed. The main thread of the narrative was to investigate generalisation of a specific training setup. But then there was also some (disjointed) discussion of positional encoding, which - although important and informative - was distracting from the main story. Some other pieces of writing were also a bit confusing, such as the transitions between sections 3.2, 3.3 and 4.1 (e.g. random flipping is described twice in 3.2 and 3.3, then branching is described twice in 3.3 and 4.1) - it might be good to consolidate those sections.
 - Constituting most of the paper - investigation into transitivity axiom alone seems underwhelming, in relation to the claims about teaching transformers causal reasoning. Transitivity check does not involve any substantial *causal* component - it is simply a question about graph reachability.
 - Thus, the most interesting section, and the only one that truly uses the heavy machinery of d-separation etc., is Section 5. But it is very short, and the full details are not even included in the appendix. Section starts by describing the task in the benchmark as "significantly harder" because of various potential questions, but this feels misleading, since the full benchmark is not used. There is no formal definition of what the "correlational statement" is, which that empirical evaluation measures - appendix A.2 simply lists all possible questions.
- The task is also quite underwhelming on a conceptual level - since the variables names are repeated, it is not clear to me to what extent the model could have just memorised some patterns, which are enough to get to 64% accuracy, as the data sizes are very small - only graphs with 3 or 4 nodes (respectively 6, 31 non-isomorphic DAGs) and testing on 5 nodes (302 non-isomorphic DAGs).
- It would be good to include a theoretical discussion as to whether we can expect transformer-based models to solve this task in principle - I don't know if the algorithm of checking d-separation can be implemented in a single pass of a bounded-depth circuit.
- In general, this sort of technique seems to quickly hit diminishing returns, as per the bitter lesson. It is not a general-purpose algorithm, as it requires very specific data augmentation for different logic and reasoning systems. I would be more interested in leaning deeper into trying to infer generalisation properties of the architecture.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: There are some minor questions that need further clarification
1. Line 317 "We train a decoder-based 67 million parameter model based on GPT-2’s architecture". Does it mean the model in the paper is trained from scratch?
soundness: 2
strengths: 1. Teaching LLM to reason is an important topic.
2. Using high-quality synthetic data for pertaining LLM is important and the explanations of the synthetic axiom generation are clear.
summary: The paper mainly presents a synthetic data generation protocol for axiom learning in transformers. Through training a GPT-2 like transformer on the synthetic demonstrations, experiments show good generalization ability when evaluating it on a larger network. The paper is well formatted and the axiom generation procedure is well explained.
weaknesses: 1. There is no proper explanation for why training the transformer on the given symbolic expressions will generalize better and even outperform the SoTA GPT-4 model. If my understanding is correct, the author only changes the dataset for pertaining and still uses negative log-likelihood (SFT) as the loss function to train the transformer.  Is it possible the model overfits the symbolic expression dataset the author has provided?
2. The author attempts to investigate how different embedding techniques, i.e., SPE, LPE, and NoPE, will affect the causal reasoning ability of LLM. As far as I know, most open-source LLMs like llama, they are all using RoPE [1] for positional embedding which is a kind of relative embedding suitable for format learning. I am curious why the author did not mention RoPE in this paper. Moreover, the evaluation of positional encoding is only limited to the final performance and I could not find any deep insight into why SPE will perform better. [2] is a good reference paper for probing different LLM layers and testing their effectiveness.

[1] Jianlin Su et al. RoFormer: Enhanced Transformer with Rotary Position Embedding

[2] Tian Ye et al. Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process

Meta-review: The paper investigates the ability of transformers to learn the axioms of causal reasoning.  Instead of directly structuring a model to follow the axioms of causal reasoning, the model is trained with data to learn the axioms.  The strengths of the paper are the interesting investigation and good write up.  The weaknesses include limited experimentation with general axioms of causality beyond transitivity and data that is restricted to abstract symbols instead of general natural language.  Overall, the reviewers found the investigation interesting, but too limited for publication at this point.  Given that GPT4 performs well in many settings while handling natural language, it is not clear what is the benefit of restricting the investigation to symbolic data.  Furthermore the paper advertises a general investigation of the axioms of causality, but most of the paper focuses exclusively on transitivity, except for Section 5.  Reviewers would like Section 5 to be expanded with more details and a broader investigation.  Overall, this is very interesting work, but perhaps too preliminary to make interesting observations that could help us understand what axioms of causality LLMs may or may not be able to generalize to in natural language.

Decision: Reject