← All publications
BETA This research page is an early version. Explanatory text, figure selection, and presentation are still being revised.
ICLR 2026

CaTs and DAGs: Integrating Directed Acyclic Graphs with Transformers for Causally Constrained Predictions

Matthew James Vowels · Mathieu Rochat · Sina Akbari
The Fourteenth International Conference on Learning Representations · 2026
Paper GitHub ICLR Talk
Keywords.
causal machine learning transformers DAGs robustness

TL;DR.

Standard neural networks can ignore causal structure that we already know. CaTs and causal fully connected networks bake a user-specified DAG into the architecture, so predictions respect the allowed causal pathways. The aim is better robustness under distribution shift and clearer interpretation without giving up expressive models.

Why this matters.

It lets neural predictors respect causal structure that is known in advance. That can improve robustness to covariate shift while making the allowed information flow much easier to interpret.

Abstract

Artificial Neural Networks (ANNs), including fully-connected networks and transformers, are highly flexible and powerful function approximators, widely applied in fields like computer vision and natural language processing. However, their inability to inherently respect causal structures can limit their robustness, making them vulnerable to covariate shift and difficult to interpret/explain. This poses significant challenges for their reliability in real-world applications. In this paper, we introduce Causal Transformers (CaTs), a general model class designed to operate under predefined causal constraints, as specified by a Directed Acyclic Graph (DAG). CaTs retain the powerful function approximation abilities of traditional neural networks while adhering to the underlying structural constraints, improving robustness, reliability, and interpretability at inference time. This approach opens new avenues for deploying neural networks in more demanding, real-world scenarios where robustness and explainability is critical.

Main message. The central idea is to make predictive models obey the causal structure we already trust. Instead of hoping a black-box model learns the right inductive bias, we build that bias in.

The paper at a glance

An informal guide to the problem, the idea, and the main results.

causal machine learning DAG-informed prediction transformers robustness under shift

1. The problem

Standard neural networks optimize predictive fit, but they do not automatically respect known causal structure. That can hurt interpretability and robustness, especially under interventions or distribution shift.

2. The key idea

We use the DAG as an architectural constraint. The network only allows information flow that is compatible with the graph, so prediction is tied directly to the assumed causal structure.

3. Structural guarantee

Theory

We prove that the CaT predictive distribution factorizes according to the supplied DAG. The same masked architecture implements the truncated factorization used by the g-formula, so intervention queries identified by the DAG can be computed directly through the network.

4. Robustness guarantee

Theory

We prove structural robustness to covariate shift: when the causal mechanisms are shared across domains, shifts in variables that are not ancestors of the target do not change the CaT predictor. The corresponding identified causal queries are transportable in the large-sample limit.

5. How it works

  1. Start from a known or trusted DAG.
  2. Translate the graph into architectural constraints on information flow.
  3. Train the resulting neural predictor with the same data as a standard baseline.
  4. Evaluate not only predictive performance but also robustness under interventions or covariate shifts.

What the experiments show

Main empirical message.

The experiments make the value of structural constraints visible. In the motivating simulation, the true-DAG CaT attains an ATE estimation error of 0.058 versus 2.379 for a standard transformer, and CaT/CFCN remain stable under the induced covariate shift while unconstrained baselines deteriorate sharply. On the Twins and Jobs benchmarks, CaT and CFCN remain competitive with methods specialized for causal inference, and the real psychology application shows that CaT can work directly with the full multidimensional questionnaire representation.

0.058eATE for true-DAG CaT
2.379eATE for a standard transformer
895participants in the psychology application
Test mean squared error under no shift and covariate shift for causal and non-causal models.
Figure 1 from the paper. CaT and CFCN remain stable under the induced shift, while the unconstrained transformer, MLP, and random forest incur much larger test error.
Absolute ATE estimation error across different models and graph specifications.
Figure 6 from the supplementary material. The true-DAG CaT and CFCN achieve much smaller ATE estimation error than the false-DAG variants and unconstrained baselines.

Where this helps

Good fit

  • You already know or trust a causal graph and want your predictor to use it.
  • You care about robustness to intervention or covariate shift rather than only i.i.d. test accuracy.
  • You want a predictive model that is easier to interpret in causal terms.

Keep in mind

  • CaT is designed for settings where structural knowledge is available and can be used both for robust prediction and for causal queries identified by the supplied DAG.
  • The same masking principle also gives CFCNs, showing that the structural idea is broader than the transformer architecture itself.

Cite this paper

@inproceedings{vowels2026cats,
  title     = {CaTs and DAGs: Integrating Directed Acyclic Graphs with Transformers for Causally Constrained Predictions},
  author    = {Matthew James Vowels and Mathieu Rochat and Sina Akbari},
  booktitle = {The Fourteenth International Conference on Learning Representations},
  year      = {2026},
  url       = {https://openreview.net/forum?id=ZIQactmQxb}
}