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

TL;DR

MOCA is a transformer-based causal inference framework that uses modular design and one-way attention to improve confounder adjustment, demonstrating strong performance across various simulated and real-world datasets.

Contribution

It introduces MOCA, a novel modular transformer architecture with one-way cross-attention and gradient detachment, enhancing causal effect estimation from observational data.

Findings

MOCA outperforms traditional estimators like IPW and AIPW in multiple simulated scenarios.

MOCA achieves competitive results on real-world datasets such as the Infant Health and Development Program.

The modular design with one-way attention preserves causal directionality and interpretability.

Abstract

Causal effect estimation from observational data requires careful adjustment for confounding. Classical estimators such as inverse probability weighting and augmented inverse probability weighting are effective under favorable model specification, but may become unstable when treatment assignment and outcome mechanisms are complex, non-linear, and high-dimensional. Machine learning and representation learning approaches improve flexibility, yet joint training can allow outcome-related information to influence treatment-side representations, which is undesirable from a causal perspective. We propose MOCA (Modular One-way Causal Attention), a transformer-based framework that separates treatment and outcome modeling through a modular design, and performs confounder adjustment using a one-way attention mechanism. A cutting-feedback strategy, implemented via gradient detachment, prevents the outcome loss from updating the treatment module. This design preserves directional information flow while retaining the representational power of transformer architectures for causal inference. Across multiple simulated scenarios, including linear, nonlinear, heavy-tailed, hidden confounding, and high-dimensional settings, MOCA shows competitive or improved performance relative to IPW, AIPW, X-learner, TARNet, and DragonNet. We further illustrate the method on the Infant Health and Development Program dataset and the Dehejia-Wahba dataset as real-world benchmarks. These results suggest that modular attention with one-way information flow provides a promising and interpretable direction for causal inference with modern deep learning models.