DAG-aware Transformer for Causal Effect Estimation
Manqing Liu, David R. Bellamy, Andrew L. Beam
TL;DR
The paper tackles causal effect estimation for $\tau$ and $\tau(x)$ in high-dimensional settings with complex causal graphs. It introduces a DAG-aware Transformer that encodes causal structure directly into the attention mechanism via a DAG-derived mask, enabling flexible estimation under G-formula, IPTW, AIPW, and proximal inference. Key contributions include the DAG-aware attention design, joint vs. separate training strategies, comprehensive evaluation on Lalonde, ACIC, and proximal tasks showing improved accuracy, and extension to proximal inference with NMMR-U/NMMR-V. This approach advances robust, structure-aware causal inference with potential impact on policy, medicine, and economics.
Abstract
Causal inference is a critical task across fields such as healthcare, economics, and the social sciences. While recent advances in machine learning, especially those based on the deep-learning architectures, have shown potential in estimating causal effects, existing approaches often fall short in handling complex causal structures and lack adaptability across various causal scenarios. In this paper, we present a novel transformer-based method for causal inference that overcomes these challenges. The core innovation of our model lies in its integration of causal Directed Acyclic Graphs (DAGs) directly into the attention mechanism, enabling it to accurately model the underlying causal structure. This allows for flexible estimation of both average treatment effects (ATE) and conditional average treatment effects (CATE). Extensive experiments on both synthetic and real-world datasets demonstrate that our approach surpasses existing methods in estimating causal effects across a wide range of scenarios. The flexibility and robustness of our model make it a valuable tool for researchers and practitioners tackling complex causal inference problems.