Teaching Transformers Causal Reasoning through Axiomatic Training

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.

Collecting the costly “what-if” data that usually teaches AI systems cause-and-effect is tricky, so we tried a simpler route: we showed a mid-sized transformer (67 million parameters) many short, synthetic examples of one basic causal rule: transitivity (“if A causes B and B causes C, then A causes C”). Remarkably, after this “axiomatic training” with variety in training data and changes in model while teaching it this rule, the model could spot causes in much larger and more complicated graphs it had never seen before, including longer chains, reversed chains, and branching structures, and it matched or even outperformed far bigger models like GPT-4, Gemini Pro, and Phi-3. The results suggest that by feeding AI enough clear demonstrations of any causal axiom, we can cheaply equip it with robust causal-reasoning skills important for real-world decision-making.