MOCA: A Transformer-based Modular Causal Inference Framework with One-way Cross-attention and Cutting Feedback
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 15 of its 39 references that are in the map. First posted on arXiv 25 Apr 2026.
Builds on
15 of 39 references.
References · 15
- 2017Attention Is All You Needcited by 60,361
- 2015Estimation and Inference of Heterogeneous Treatment Effects using Random Forestscited by 600
- 2017Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learningcited by 478
- 2017Causal Effect Inference with Deep Latent-Variable Modelscited by 410
- 2016Learning Representations for Counterfactual Inferencecited by 329
- 2019Adapting Neural Networks for the Estimation of Treatment Effectscited by 287
- 2020Conformal Inference of Counterfactuals and Individual Treatment Effectscited by 151
- 2022Causal Transformer for Estimating Counterfactual Outcomescited by 123
- 2019A synthetic approach to Markov kernels, conditional independence and theorems on sufficient statisticscited by 56
- 2025Do-PFN: In-Context Learning for Causal Effect Estimationcited by 51
- 2016A Shrinkage Principle for Heavy-Tailed Data: High-Dimensional Robust Low-Rank Matrix Recoverycited by 23
- 2020Targeted VAE: Variational and Targeted Learning for Causal Inferencecited by 7
3 more
- 2020Semi-Modular Inference: enhanced learning in multi-modular models by tempering the influence of componentscited by 4
- 2022A General Framework for Cutting Feedback within Modularised Bayesian Inferencecited by 2
- 2022Cutting feedback and modularized analyses in generalized Bayesian inferencecited by 1
Built on by
Not cited yet.