MOCA: A Transformer-based Modular Causal Inference Framework with One-way Cross-attention and Cutting Feedback

Abstract:Causal effect estimation from observational data requires careful adjustment for confounding. Classical estimators such as inverse probability weighting and augmented inverse probability weighting can perform well under favorable model specification but may become unstable in complex settings. Machine-learning and representation-learning methods provide greater flexibility, but joint optimization may allow outcome information to alter treatment representations and compromise the intended causal structure. We propose MOCA (Modular One-way Causal Attention), a transformer-based framework that separates treatment and outcome modeling through a modular architecture. This design preserves directional information flow while retaining the flexibility of transformer architectures. Using a Markov-kernel formulation, we show that MOCA learns an autonomous treatment representation and is predictively KL-optimal within its Gaussian model class. Under correct specification, MOCA recovers the true average treatment effect, while additional two-way feedback does not further reduce average treatment effect estimation error. We also propose a conformal inference procedure for individual treatment effects. Across multiple simulation scenarios, MOCA achieved competitive or improved average treatment effect estimation compared with IPW, AIPW, the X-learner, TARNet, DragonNet, BART, Causal Forest, and Do-PFN. Ablation studies supported the contributions of the proposed architectural components. We further evaluated MOCA on the Infant Health and Development Program benchmark and the observational Dehejia-Wahba dataset. Overall, modular attention with one-way information flow provides an effective and interpretable framework for causal inference using modern deep-learning models.

41 pages, 6 figures, 6 tables. Preprint

Machine Learning (stat.ML); Machine Learning (cs.LG); Methodology (stat.ME)

arXiv:2604.23107 [stat.ML]

(or arXiv:2604.23107v2 [stat.ML] for this version)

https://doi.org/10.48550/arXiv.2604.23107 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.