TransTEE

This is an official PyTorch implementation of the paper Can Transformers be Strong Treatment Effect Estimators?

Abstract

In this paper, we develop a general framework based on the Transformer architecture to address a variety of challenging treatment effect estimation (TEE) problems. Our methods are applicable both when covariates are tabular and when they consist of sequences (e.g., in text), and can handle discrete, continuous, structured, or dosage-associated treatments. While Transformers have already emerged as dominant methods for diverse domains, including natural language and computer vision, our experiments with Transformers as Treatment Effect Estimators (TransTEE) demonstrate that these inductive biases are also effective on the sorts of estimation problems and datasets that arise in research aimed at estimating causal effects. Moreover, we propose a propensity score network that is trained with TransTEE in an adversarial manner to promote independence between covariates and treatments to further address selection bias. Through extensive experiments, we show that TransTEE significantly outperforms competitive baselines with greater parameter efficiency over a wide range of benchmarks and settings.

Dependencies

Dependencies of different settings are listed in each subfolder.

Continuous Treatments

./Continuous

Continuous Dosage

./Dosage

Structured Treatments

./Structured

Empirical Study on Pre-trained Language Models

./CausalM

Citation

The repository is under construction. If you find this repo useful, please cite:

@misc{zhang2022transformers,
      title={Can Transformers be Strong Treatment Effect Estimators?}, 
      author={Yi-Fan Zhang and Hanlin Zhang and Zachary C. Lipton and Li Erran Li and Eric P. Xing},
      year={2022},
      eprint={2202.01336},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

nktoan / TransTEE