# Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation

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

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 85 references that are in the map.
- First posted on arXiv 2 Feb 2022.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Multimodal Transformer for Unaligned Multimodal Language Sequences](https://mlsanity.com/paper/1906.00295/) (2019)
- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://mlsanity.com/paper/2107.14795/) (2021)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Towards Faithfully Interpretable NLP Systems: How should we define and evaluate faithfulness?](https://mlsanity.com/paper/2004.03685/) (2020)
- [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)
- [Do Transformers Really Perform Bad for Graph Representation?](https://mlsanity.com/paper/2106.05234/) (2021)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Datamodels: Predicting Predictions from Training Data](https://mlsanity.com/paper/2202.00622/) (2022)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Estimating Counterfactual Treatment Outcomes over Time Through Adversarially Balanced Representations](https://mlsanity.com/paper/2002.04083/) (2020)
- [Orthogonal Statistical Learning](https://mlsanity.com/paper/1901.09036/) (2019)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Learning in Implicit Generative Models](https://mlsanity.com/paper/1610.03483/) (2016)
- [Kernel Instrumental Variable Regression](https://mlsanity.com/paper/1906.00232/) (2019)
- [CausaLM: Causal Model Explanation Through Counterfactual Language Models](https://mlsanity.com/paper/2005.13407/) (2020)
- [Towards Principled Disentanglement for Domain Generalization](https://mlsanity.com/paper/2111.13839/) (2021)
- [Learning Counterfactual Representations for Estimating Individual Dose-Response Curves](https://mlsanity.com/paper/1902.00981/) (2019)
- [Estimating the Effects of Continuous-valued Interventions using Generative Adversarial Networks](https://mlsanity.com/paper/2002.12326/) (2020)
- [Continuously Indexed Domain Adaptation](https://mlsanity.com/paper/2007.01807/) (2020)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects](https://mlsanity.com/paper/2001.07426/) (2020)

## Built on by (most cited first, 25 of 25)

- [AdaNPC: Exploring Non-Parametric Classifier for Test-Time Adaptation](https://mlsanity.com/paper/2304.12566/) (2023)
- [Counterfactual Generative Models for Time-Varying Treatments](https://mlsanity.com/paper/2305.15742/) (2023)
- [Benchmarking Heterogeneous Treatment Effect Models through the Lens of Interpretability](https://mlsanity.com/paper/2206.08363/) (2022)
- [Domain-Specific Risk Minimization for Out-of-Distribution Generalization](https://mlsanity.com/paper/2208.08661/) (2022)
- [Deep Causal Learning: Representation, Discovery and Inference](https://mlsanity.com/paper/2211.03374/) (2022)
- [Adversarially Balanced Representation for Continuous Treatment Effect Estimation](https://mlsanity.com/paper/2312.10570/) (2023)
- [Continuous Treatment Effect Estimation Using Gradient Interpolation and Kernel Smoothing](https://mlsanity.com/paper/2401.15447/) (2024)
- [Time Series Causal Link Estimation under Hidden Confounding using Knockoff Interventions](https://mlsanity.com/paper/2209.11497/) (2022)
- [TemporAI: Facilitating Machine Learning Innovation in Time Domain Tasks for Medicine](https://mlsanity.com/paper/2301.12260/) (2023)
- [Heterogeneous Treatment Effect with Trained Kernels of the Nadaraya-Watson Regression](https://mlsanity.com/paper/2207.09139/) (2022)
- [Uplift modeling with continuous treatments: A predict-then-optimize approach](https://mlsanity.com/paper/2412.09232/) (2024)
- [PairNet: Training with Observed Pairs to Estimate Individual Treatment Effect](https://mlsanity.com/paper/2406.03864/) (2024)
- [DISCRET: Synthesizing Faithful Explanations For Treatment Effect Estimation](https://mlsanity.com/paper/2406.00611/) (2024)
- [Sources of Gain: Decomposing Performance in Conditional Average Dose Response Estimation](https://mlsanity.com/paper/2406.08206/) (2024)
- [DAG-aware Transformer for Causal Effect Estimation](https://mlsanity.com/paper/2410.10044/) (2024)
- [CausalKANs: interpretable treatment effect estimation with Kolmogorov-Arnold networks](https://mlsanity.com/paper/2509.22467/) (2025)
- [BENK: The Beran Estimator with Neural Kernels for Estimating the Heterogeneous Treatment Effect](https://mlsanity.com/paper/2211.10793/) (2022)
- [SubgroupTE: Advancing Treatment Effect Estimation with Subgroup Identification](https://mlsanity.com/paper/2401.12369/) (2024)
- [Heterogeneous treatment effect estimation with subpopulation identification for personalized medicine in opioid use disorder](https://mlsanity.com/paper/2401.17027/) (2024)
- [Identification of Causal Relationship between Amyloid-beta Accumulation and Alzheimer's Disease Progression via Counterfactual Inference](https://mlsanity.com/paper/2307.01389/) (2023)
- [Using representation balancing to learn conditional-average dose responses from clustered data](https://mlsanity.com/paper/2309.03731/) (2023)
- [Disentangled Representation via Variational AutoEncoder for Continuous Treatment Effect Estimation](https://mlsanity.com/paper/2406.02310/) (2024)
- [Leveraging a Simulator for Learning Causal Representations from Post-Treatment Covariates for CATE](https://mlsanity.com/paper/2502.05037/) (2025)
- [Advancing Edge Classification through High-Dimensional Causal Modeling of Node-Edge Interplay](https://mlsanity.com/paper/2605.00374/) (2026)
- [Surv-IPTB: An Attention-Based Model for Estimating Individual Probability of Treatment Benefit with Survival Data](https://mlsanity.com/paper/2608.06288/) (2026)

## Abstract (arXiv)

Previous works on Treatment Effect Estimation (TEE) are not in widespread use because they are predominantly theoretical, where strong parametric assumptions are made but untractable for practical application. Recent work uses multilayer perceptron (MLP) for modeling casual relationships, however, MLPs lag far behind recent advances in ML methodology, which limits their applicability and generalizability. To extend beyond the single domain formulation and towards more realistic learning scenarios, we explore model design spaces beyond MLPs, i.e., transformer backbones, which provide flexibility where attention layers govern interactions among treatments and covariates to exploit structural similarities of potential outcomes for confounding control. Through careful model design, Transformers as Treatment Effect Estimators (TransTEE) is proposed. We show empirically that TransTEE can: (1) serve as a general purpose treatment effect estimator that significantly outperforms competitive baselines in a variety of challenging TEE problems (e.g., discrete, continuous, structured, or dosage-associated treatments) and is applicable to both when covariates are tabular and when they consist of structural data (e.g., texts, graphs); (2) yield multiple advantages: compatibility with propensity score modeling, parameter efficiency, robustness to continuous treatment value distribution shifts, explainable in covariate adjustment, and real-world utility in auditing pre-trained language models