1 of 41

Shantanu Ghosh,

University of Pittsburgh

Zheng Feng, Jiang Bian, Kevin Butler, Mattia Prosperi

University of Florida

DR-VIDAL - Doubly Robust Variational Information theoretic Deep Adversarial Learning for Counterfactual Prediction and Treatment Effect Estimation

S90

2 of 41

Nancy is having fever

AMIA 2022 Annual Symposium | amia.org

2

I am having a mild fever. Shall I take medicine or not?

3 of 41

Nancy is having fever

AMIA 2022 Annual Symposium | amia.org

3

I am having a mild fever. Shall I take medicine or not?

Features

Age: 32

Sex: F

Blood group: A+

Race: Asian

Blood Sugar: High

Temperature: 100°F

….

4 of 41

Nancy is having fever

AMIA 2022 Annual Symposium | amia.org

4

I am having a mild fever. Shall I take medicine or not?

Features

Age: 32

Sex: F

Blood group: A+

Race: Asian

Blood Sugar: High

Temperature: 100°F

….

Medication A

Control

T=0

Medication B

Treated

T=1

Temperature = ?

5 of 41

Nancy is having fever

AMIA 2022 Annual Symposium | amia.org

5

I am having a mild fever. Shall I take medicine or not?

Features

Age: 32

Sex: F

Blood group: A+

Race: Asian

Blood Sugar: High

Temperature: 100°F

….

Medication A

Control

T=0

Medication B

Treated

T=1

Temperature = ?

X

T

Y

6 of 41

Nancy is having fever

AMIA 2022 Annual Symposium | amia.org

6

I am having a mild fever. Shall I take medicine or not?

Features

Age: 32

Sex: F

Blood group: A+

Race: Asian

Blood Sugar: High

Temperature: 100°F

….

Medication A

Control

T=0

Medication B

Treated

T=1

Temperature = ?

X

T

Y

7 of 41

Fundamentally challenging problem

AMIA 2022 Annual Symposium | amia.org

7

I am having a mild fever. Shall I take medicine or not?

Features

Age: 32

Sex: F

Blood group: A+

Race: Asian

Blood Sugar: High

Temperature: 100°F

….

Medication A

Control

T=0

Medication B

Treated

T=1

Temperature = ?

X

T

Y

Only one outcome is observed

8 of 41

Fundamentally challenging problem

AMIA 2022 Annual Symposium | amia.org

8

I am having a mild fever. Shall I take medicine or not?

Features

Age: 32

Sex: F

Blood group: A+

Race: Asian

Blood Sugar: High

Temperature: 100°F

….

Medication A

Control

T=0

Medication B

Treated

T=1

Temperature = ?

X

T

Y

Only one outcome is observed

Counterfactual

Factual

9 of 41

We are in Big data era

AMIA 2022 Annual Symposium | amia.org

9

Prosperi et al.[Nature Machine Intelligence 2019]

10 of 41

But data can be problematic

AMIA 2022 Annual Symposium | amia.org

10

Forgetting confounders or including colliders produces biased models: predictions can still be good, but models are not causal

11 of 41

Causal Inference is challenging due to counterfactuals

AMIA 2022 Annual Symposium | amia.org

11

Randomized control trail
Propensity score matching
Inverse probability treatment weighting

Solutions

Drawbacks

Unethical, not feasible
Rest are mostly linear following

logistic regression for estimation

Treated group

P(cured|T=1) = 0.7

Control group

P(cured|T=0) = 0.2

Propensity score

P(T|X)=0.6

P(T|X)=0.9

P(T|X)=0.3

Data rebalancing via

matching

propensity scores

P(cured|T=1) = ?

P(cured|T=0) = ?

12 of 41

It is a deep learning era

AMIA 2022 Annual Symposium | amia.org

12

Ghosh et al. [Computer methods and programs in biomedicine update, 2021]

Ghosh et al. [JAMIA, 2021]

Model non-linear treatment assignment
Handle high dimensional inputs
Provide non-linear counterfactual predictions

13 of 41

Our contribution

AMIA 2022 Annual Symposium | amia.org

13

Incorporate of an underlying causal structure to approximate the data generation

process.

2. Infer the latent variables, responsible for generating data using a VAE.

Generate counterfactual outcomes using a GAN with variational information

theoretic regularization.

Estimate the individual treatment effect by minimizing the factual and

counterfactual outcomes both.

5. Utilize doubly robust regularizer for faster convergence.

14 of 41

GANITE generates counterfactuals

AMIA 2022 Annual Symposium | amia.org

14

15 of 41

Can we use latent codes?

AMIA 2022 Annual Symposium | amia.org

15

Dumoulin et al. [ICLR, 2017]

Larsen et al. [ICML, 2016]

16 of 41

Problem formulation

AMIA 2022 Annual Symposium | amia.org

16

Used potential outcome framework by Rubin.

The complete tuple {X_i, T_i, Y_i}, for i=1...N

Y_i⁰and Y_i¹are the potential outcomes for treatment T_i=0 and T_i=1

The ITE for the subject i with covariates X_i = x is defined as,

Assumption:

Followed strongly ignorable treatment assignment assumption (SITA), defined as,

17 of 41

Contribution #1: Causal structure

AMIA 2022 Annual Symposium | amia.org

17

Z: Latent space
X: Observed covariates
t: Treatment indicator
Y: Potential outcome

18 of 41

Contribution #2: Infer the latent variables

AMIA 2022 Annual Symposium | amia.org

18

Back Propagation

Forward Propagation

19 of 41

Contribution #2: Infer the latent variables

AMIA 2022 Annual Symposium | amia.org

19

The overall ELBO function to optimize

https://github.com/Shantanu48114860/DR-VIDAL-AMIA-22/blob/main/DR_VIDAL_AMIA-Supp.pdf

Back Propagation

Forward Propagation

The variational posteriors of the inference model is defined as,

All the latent factors - z, are assumed to have a prior gaussian distributions defined as,

20 of 41

Contribution #3: Generate counterfactuals

AMIA 2022 Annual Symposium | amia.org

20

Back Propagation

Forward Propagation

Following GANITE, the optimization function of the GAN block,

The supervised loss is defined as,

The complete loss of the counterfactual GAN block is defined as,

21 of 41

Contribution #3: Generate counterfactuals

AMIA 2022 Annual Symposium | amia.org

21

Back Propagation

Forward Propagation

Following GANITE, the optimization function of the GAN block,

The supervised loss is defined as,

The complete loss of the counterfactual GAN block is defined as,

Maximize I(z_c; G(z_G, z_c)) to solve the optimization function

22 of 41

Contribution #3: Generate counterfactuals

AMIA 2022 Annual Symposium | amia.org

22

Back Propagation

Forward Propagation

I(z_c; G(z_G, z_c)) is harder to optimize due to p(z_c|x), so using a variational distribution Q(z_c|x) to approximate the posterior p(z_c|x).

The optimal discriminator and generator will be obtained by solving the following objectives

https://github.com/Shantanu48114860/DR-VIDAL-AMIA-22/blob/main/DR_VIDAL_AMIA-Supp.pdf

23 of 41

Contribution #4: Estimate ITE by minimizing the factuals and counterfactuals

AMIA 2022 Annual Symposium | amia.org

23

The factual and the counterfactual outcomes were estimated as,

Shalit et al. [ICML, 2017]

24 of 41

Contribution #5: Doubly robust ITE Estimation

AMIA 2022 Annual Symposium | amia.org

24

Using propensity score, the doubly robust estimation of causal effect is defined as,

Jonsson et al. [American Journal of epidemiology, 2011]

where,

The propensity score is defined as,

25 of 41

Contribution #5: Doubly robust ITE Estimation

AMIA 2022 Annual Symposium | amia.org

25

Using propensity score, the doubly robust estimation of causal effect is defined as,

Jonsson et al. [American Journal of epidemiology, 2011]

where,

The predicted loss to be optimized as,

The propensity score is defined as,

26 of 41

Contribution #5: Doubly robust ITE Estimation

AMIA 2022 Annual Symposium | amia.org

26

Jonsson et al. [American Journal of epidemiology, 2011]

The factual and the counterfactual doubly robust outcomes were estimated as,

27 of 41

Contribution #5: Doubly robust ITE Estimation

AMIA 2022 Annual Symposium | amia.org

27

Jonsson et al. [American Journal of epidemiology, 2011]

The factual and the counterfactual doubly robust outcomes were estimated as,

The doubly robust loss is optimized as ,

28 of 41

Contribution #5: Doubly robust ITE Estimation

AMIA 2022 Annual Symposium | amia.org

28

Jonsson et al. [American Journal of epidemiology, 2011]

The factual and the counterfactual doubly robust outcomes were estimated as,

The doubly robust loss is optimized as ,

The complete loss to estimate ITE to be optimized as,

29 of 41

Performance metrics

AMIA 2022 Annual Symposium | amia.org

29

30 of 41

Results – Synthetic datasets 1

AMIA 2022 Annual Symposium | amia.org

30

Comparison of the performance (ATE) of DR- VIDAL vs. all other models on samples from the generative process of Synthetic dataset 1 - sample size {1000, 3000, 5000, 10000, 30000}

Louizos et al. [Neurips, 2017]

31 of 41

Results – Synthetic datasets 2

AMIA 2022 Annual Symposium | amia.org

31

Performance comparison (PEHE) of GANITE vs. DR-VIDAL

Sample sizes: {1000, 3000, 5000, 10000, 30000}

32 of 41

Results – Synthetic datasets 2

AMIA 2022 Annual Symposium | amia.org

32

Performance comparison (PEHE) of GANITE vs. DR-VIDAL

Sample sizes: {1000, 3000, 5000, 10000, 30000}

33 of 41

Results – Real world datasets

AMIA 2022 Annual Symposium | amia.org

33

IHDP
Jobs
Twins

34 of 41

Results – Ablation study

AMIA 2022 Annual Symposium | amia.org

34

35 of 41

Results – correct classification of factual outcomes

AMIA 2022 Annual Symposium | amia.org

35

36 of 41

Results – PEHE and policy risk �(mean ± st.dev)

AMIA 2022 Annual Symposium | amia.org

36

37 of 41

Results – ATE and ATT�(mean ± st.dev)

AMIA 2022 Annual Symposium | amia.org

37

38 of 41

Conclusion

AMIA 2022 Annual Symposium | amia.org

38

Beats the performance of previous generative models on synthetic datasets
Comparable performance on real-world datasets.

Code: https://github.com/Shantanu48114860/DR-VIDAL-AMIA-22/

39 of 41

Future directions

AMIA 2022 Annual Symposium | amia.org

39

1. Causal Graph is too simple.

2. How to enforce strict disentanglement?

3. Instead of Variational Information Maximization, what about Variational Information Bottleneck Layer?

4. TARNET, DRAGONNET, DCN-PD, ITE block as a downstream model instead of the Doubly robust treatment estimator.

6. More realistic data from News-8 or MIMIC-II(EBB).

40 of 41

Acknowledgement

AMIA 2022 Annual Symposium | amia.org

40

Funded by NIH awards R21CA245858, R01CA246418, R56AG069880, R01AG076234, R01AI145552,�R01AI141810, and NSF 2028221

41 of 41

Thank you!

Email me at: shg121@pitt.edu