Mutual Information Treatment Network (ICML 2023)

Estimating Heterogeneous Treatment Effects: Mutual Information Bounds and Learning Algorithms [paper]

Towards estimating heterogeneous treatment effect (HTE) for general treatment spaces, we propose:

• Theoretical bound under mutual information for estimating HTE.
• Theory-guided algorithm MitNet for practical usage.

Figure 1. Overview of MitNet.

Get Started

1. Install Python 3.8. For convenience, execute the following command.

pip install -r requirements.txt

2. Prepare Data. You can obtain experimental datasets from the following links.

Dataset	Link
TCGA	[Google Drive]

3. Train and evaluate model.

cd TCGA
bash run_mitnet.sh

Results

We experiment on three benchmark datasets. MitNet reaches remarkable performance.

Table 1. Results for HTE estimation on IHDP, News and TCGA. A lower metric indicates better performance.

Citation

If you find this repo useful, please cite our paper.

@inproceedings{guo2023MitNet,
  title={Estimating Heterogeneous Treatment Effects: Mutual Information Bounds and Learning Algorithms},
  author={Xingzhuo Guo and Yuchen Zhang and Jianmin Wang and Mingsheng Long},
  booktitle={International Conference on Machine Learning},
  year={2023}
}

Contact

If you have any questions or want to use the code, please contact gxz23@mails.tsinghua.edu.cn.

Acknowledgement

We appreciate the following github repos a lot for their valuable code base:

https://github.com/thuml/Transfer-Learning-Library