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

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

## 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 (most cited first, 15 of 15)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Causal Effect Inference with Deep Latent-Variable Models](https://mlsanity.com/paper/1705.08821/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Conformal Inference of Counterfactuals and Individual Treatment Effects](https://mlsanity.com/paper/2006.06138/) (2020)
- [Causal Transformer for Estimating Counterfactual Outcomes](https://mlsanity.com/paper/2204.07258/) (2022)
- [A synthetic approach to Markov kernels, conditional independence and theorems on sufficient statistics](https://mlsanity.com/paper/1908.07021/) (2019)
- [Do-PFN: In-Context Learning for Causal Effect Estimation](https://mlsanity.com/paper/2506.06039/) (2025)
- [A Shrinkage Principle for Heavy-Tailed Data: High-Dimensional Robust Low-Rank Matrix Recovery](https://mlsanity.com/paper/1603.08315/) (2016)
- [Targeted VAE: Variational and Targeted Learning for Causal Inference](https://mlsanity.com/paper/2009.13472/) (2020)
- [Semi-Modular Inference: enhanced learning in multi-modular models by tempering the influence of components](https://mlsanity.com/paper/2003.06804/) (2020)
- [A General Framework for Cutting Feedback within Modularised Bayesian Inference](https://mlsanity.com/paper/2211.03274/) (2022)
- [Cutting feedback and modularized analyses in generalized Bayesian inference](https://mlsanity.com/paper/2202.09968/) (2022)

## 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.