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 potentialin estimating causal effects, existing approaches often fall short in handlingcomplex causal structures and lack adaptability across various causalscenarios. In this paper, we present a novel transformer-based method forcausal inference that overcomes these challenges. The core innovation of ourmodel lies in its integration of causal Directed Acyclic Graphs (DAGs) directlyinto the attention mechanism, enabling it to accurately model the underlyingcausal structure. This allows for flexible estimation of both average treatmenteffects (ATE) and conditional average treatment effects (CATE). Extensiveexperiments on both synthetic and real-world datasets demonstrate that ourapproach surpasses existing methods in estimating causal effects across a widerange of scenarios. The flexibility and robustness of our model make it avaluable tool for researchers and practitioners tackling complex causalinference problems.