# DAG-aware Transformer for Causal Effect Estimation

arXiv: https://arxiv.org/abs/2410.10044 · HTML: https://mlsanity.com/paper/2410.10044/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 28 references that are in the map.
- First posted on arXiv 13 Oct 2024.

## Builds on (most cited first, 12 of 12)

- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Recursive Partitioning for Heterogeneous Causal Effects](https://mlsanity.com/paper/1504.01132/) (2015)
- [Causal Transformer for Estimating Counterfactual Outcomes](https://mlsanity.com/paper/2204.07258/) (2022)
- [Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition](https://mlsanity.com/paper/1707.02641/) (2017)
- [RealCause: Realistic Causal Inference Benchmarking](https://mlsanity.com/paper/2011.15007/) (2020)
- [Causal Inference under Networked Interference and Intervention Policy Enhancement](https://mlsanity.com/paper/2002.08506/) (2020)
- [Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2202.01336/) (2022)
- [Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation](https://mlsanity.com/paper/2211.01939/) (2022)
- [Bounds on Representation-Induced Confounding Bias for Treatment Effect Estimation](https://mlsanity.com/paper/2311.11321/) (2023)
- [Towards Causal Foundation Model: on Duality between Causal Inference and Attention](https://mlsanity.com/paper/2310.00809/) (2023)
- [Deep Learning Methods for Proximal Inference via Maximum Moment Restriction](https://mlsanity.com/paper/2205.09824/) (2022)

## Built on by (most cited first, 2 of 2)

- [CausalPFN: Amortized Causal Effect Estimation via In-Context Learning](https://mlsanity.com/paper/2506.07918/) (2025)
- [CoT-ICL Lab: A Synthetic Framework for Studying Chain-of-Thought Learning from In-Context Demonstrations](https://mlsanity.com/paper/2502.15132/) (2025)

## Abstract (arXiv)

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.