Teaching Transformers Causal Reasoning through Axiomatic Training

We thank the reviewer for appreciating the contributions of the work. We answer specific questions below.

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.

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.

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.

Lastly, 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).

2. Could the authors clarify whether pretraining Llama-7B with this approach consistently outperforms prompting-based methods for large models in MultiEvalSLR Dataset?

Below 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.

MultiEvalSLR 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).

As 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).

https://arxiv.org/pdf/2405.15071, discusses the generalization capabilities of large models on synthetic dataset, could be relevant to this work.

Thanks for pointing out this work. We will add a discussion on this, especially in the context of composition-based reasoning.

We thank the reviewer for appreciating the contributions of the work. We answer the specific questions below:

1. Eval inputs being longer than training text inputs?

Yes 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.

2. The two assumptions employed...how restrictive these are?

We 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.

In 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.

3. Discussion on whether causal reasoning problems are a subclass of graph reasoning

That'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.

4. The connection to positional encoding .... encodings suited for causal reasoning.

Our 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. Although 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.

5. Prior work on passive learning

Lampinen 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.

6. The length generalization result ... compare to past work

Abbe 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.

We thank the reviewer for appreciating the contributions of our work. We answer specific questions below:

Response to Weaknesses

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.

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, 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.

Moreover, 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.

Lastly, 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).

2. This raises concerns about whether the proposed training approach provides a significant advantage over large-scale unsupervised learning.

Response We 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.

Also, 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.