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

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.

💡 Research Summary

The paper tackles a fundamental limitation of modern deep learning models: their reliance on statistical correlations without respecting known causal structures, which leads to vulnerability under covariate shift and poor interpretability. To address this, the authors introduce two model families that embed a pre‑specified Directed Acyclic Graph (DAG) directly into the architecture. The primary contribution is the Causal Transformer (CaT), a modification of the standard transformer that enforces DAG‑constrained attention, and a baseline Causal Fully Connected Network (CFCN) that applies analogous masking to a multilayer perceptron.

Model design
Both models assume a DAG G = (V,E) over |Z| variables (nodes). Input data are organized as a batch B of tensors with shape (B × |Z| × C), where C is the feature dimension of each variable. In CaT each variable is first projected by an independent linear layer into a common embedding space of dimension d_E (d_E ≥ C). A learnable global embedding γ ∈ ℝ^{|Z|×d_E} is initialized randomly and serves as the query source for the first attention block. For each block the usual key, query, value matrices are computed as

 K = X_E W_K + b_K, Q = γ W_Q + b_Q, V = X_E W_V + b_V

where X_E is the embedded input. The attention scores QK^T are element‑wise multiplied (Hadamard product) with the transposed, topologically‑sorted adjacency matrix A^T, thereby zeroing out any attention from a node to a non‑parent. After the softmax, the masked scores weight V to produce the block output O. This masked cross‑attention is performed in parallel across H heads and stacked across L transformer blocks; each block’s output is passed through a residual connection, layer normalization, and then becomes the new γ for the next block. Notably, CaT does not add self‑connections to the mask after the first layer, relying on the DAG’s acyclicity to preserve information flow.

CFCN follows the MADE approach: each dense layer’s weight matrix W is multiplied by a binary mask M derived directly from the DAG, ensuring that node j can influence node i only if (j→i)∈E. This yields a simple, fully‑connected architecture that respects the same conditional independencies as CaT but lacks the ability to handle high‑dimensional embeddings without additional preprocessing.

Key advantages

1. Robustness to distribution shift – By restricting learning to causal pathways, the models ignore spurious correlations that may change between training and deployment. Empirical illustrations (e.g., camel vs. cow classification with background changes) show that CaT and CFCN maintain accuracy where conventional CNNs, MLPs, or vanilla transformers degrade.
2. Interpretability – The attention mask aligns one‑to‑one with the DAG, making it straightforward to trace which parents drive each prediction. This is valuable for domains such as medicine or policy where causal effect attribution is required.
3. Scalability – The masking operation does not increase computational complexity; CaT retains the O(B·|Z|·d_E·h) cost of standard transformers. Multi‑head and multi‑block parallelism are unchanged.

Technical observations

• The dimensionality d_E must be sufficiently larger than the original feature size C; otherwise the network struggles to disentangle variables, especially when C = 1.
• γ acts as a learnable “global token” that aggregates information across blocks; its random initialization works in practice, but sensitivity to the initial seed was not deeply explored.
• The method assumes the DAG is known a priori. In realistic settings, DAG discovery or expert elicitation will be necessary, and the impact of misspecified edges on performance remains an open question.
• The authors provide a formal proof (Supplement S13) that the DAG‑induced conditional independencies are preserved through the transformer layers.

Related work positioning
Previous attempts to make transformers “causal” have focused on autoregressive masking (restricting attention to past tokens) which only captures Granger‑type temporal causality. Other works embed DAGs into normalizing flows or structured neural networks, but often require complex training objectives or are limited to scalar inputs. CaT distinguishes itself by (i) supporting arbitrary high‑dimensional embeddings, (ii) applying the DAG mask directly to the attention matrix rather than to weight parameters, and (iii) offering a unified architecture that does not need separate subnetworks for treatment, covariates, and outcomes.

Experiments & results (as described in the paper excerpt)

• Synthetic image classification tasks demonstrating covariate shift robustness.
• Medical causal inference scenarios (e.g., treatment → mediator → outcome) where CaT yields unbiased average treatment effect estimates under the back‑door criterion, while unconstrained models exhibit bias.
• Ablation studies comparing CaT, CFCN, and baseline transformers/MLPs, highlighting the trade‑off between predictive power on a stable distribution and robustness under shift.

Limitations & future directions

• Reliance on a known DAG limits immediate applicability; integrating structure learning (e.g., score‑based or constraint‑based methods) with CaT is a natural next step.
• The approach has not been evaluated on large‑scale language models or graph‑structured data where the number of nodes can be in the thousands; scalability of the mask generation and storage may become an issue.
• Sensitivity to incorrect edges is not quantified; robustness to misspecification would be crucial for real‑world deployment.
• The current formulation focuses on single‑step interventions; extending to dynamic or sequential decision‑making (e.g., reinforcement learning) would broaden impact.

Conclusion
The authors present a principled way to inject causal knowledge into transformer architectures via DAG‑based attention masking, and a simpler fully‑connected counterpart. By aligning model inductive bias with known causal structure, CaT achieves improved robustness to covariate shift and enhanced interpretability without sacrificing the expressive power of modern deep networks. The work opens a promising research avenue at the intersection of causal inference and deep learning, inviting further exploration of automated DAG discovery, scalability to massive models, and robustness analyses under structural misspecification.