V ARIATIONAL C OUNTERFACTUAL P REDICTION UNDER R UNTIME D OMAIN C ORRUPTION
Anonymous authors
Paper under double-blind review
A BSTRACT
To date, various neural methods have been proposed to address the causal effect estimation based on observational data, where the counterfactual prediction commonly assumes the same distribution and availability of variables at both training and inference (i.e., runtime) stages.
In reality, covariate shift commonly happens, and the accessibility of variables is usually impaired due mainly to privacy and ethical concerns.
We term the co-occurrence of domain shift and inaccessible variables runtime domain corruption , which seriously challenges the generalizability of the trained counterfactual predictor on top of the existence of confoundedness and selection bias.
To counteract runtime domain corruption, we subsume counterfactual prediction under the notion of domain adaptation.
Specifically, we upper-bound the error w.r.t.
the target domain (i.e., runtime covariates) by the sum of source domain error and inter-domain distribution distance.
In addition, we build an adversarially unified variational causal effect model, named VEGAN, with a novel two-stage adversarial domain adaptation to implicitly reduce the distribution disparity between treated and control groups first, and between training and inference domains afterwards.
We demonstrate that VEGAN outperforms other state-of-the-art baselines on individual-level treatment effect estimation in the presence of runtime domain corruption on benchmark datasets.
I NTRODUCTION
In predictive analytics, causal inference is increasingly important in guiding decision-making in high-stake domains, such as healthcare (Glass et al., 2013), education (Cordero et al., 2018), ecommerce (Phang et al., 2019), etc.
Normally, randomized control trial (RCT) is the gold standard for estimating the causal effect.
Given that implementing RCTs is costly, time-consuming, and sometimes ethically intractable, various applications alternatively turn to use the passively collected observational data to perform causal inference in a data-driven fashion (Johansson et al., 2016; Yao et al., 2021).
Denoting input variables as x , treatment as t , outcome as y , the observational dataset with N samples { ( x i , t
i , y i ) } N i =1 commonly does not satisfy the RCT standard due to unmeasured confounders and selection bias, which correspond to the two prominent challenges below.
Challenge 1: The untestable unconfoundedness assumption assumes no unobserved confounders.
However, such an assumption cannot be satisfied in many cases, rendering the causal effect estimation erroneous (Rosenbaum, 2002; Caliendo & Kopeinig, 2008).
Challenge 2: The imbalanced covariate distribution between the treated group and the control group makes the treatment assignment t not random and subject to selection bias, which further weakens the effectiveness of causal inference.
A healthcare example to explain these two challenges is that, if only the rich can afford drug A while the poor have to use the cheaper alternative drug B, then people’s financial status could be a hidden confounder if unmeasured.
Consequently, it confounds the treatment assignment, and the effectiveness of drug A and drug B cannot be validly compared based on the rich and the poor groups due to the skewed distribution of variables.
To date, by addressing either of the two challenges or both, several neural approaches (Shalit et al., 2017; Louizos et al., 2017; Shi et al., 2019) are made available for causal effect estimation with observational data.
Notably, the majority of such counterfactual predictions ignores the domain shift situation, and also bear an assumption that all input variables used for training are also accessible during the inference stage (i.e., runtime).
However, real-world applications commonly see that the medical diagnostic models are learned with high-quality open benchmarks, but in the deployment stage, covariate shift could happen, also not all end-users are able to provide the same set of attributes due to accessibility issues (e.g., high-cost medical checks), privacy constraints (e.g., historical treatments), and ethical concerns (e.g., gender and race).
We specifically define this challenge as the runtime domain corruption as follows.
Challenge 3: We define each variable vector x
=
[ x 1 , x 2 , ..., x d ] ∈ R d as a concatenation of multihot-coded categorical features and non-zero numerical features.
During training, all d entries x s , ∀ s ≤ d , are available and assigned corresponding values.
Then, during inference, runtime domain corruption occurs when the covariate distribution shifts, and vector x contains an arbitrary number of unavailable variables which are set to x s ′ = 0 , s ′ ≤ d .
Domain Shift vs. Runtime Domain Corruption.
Runtime domain corruption can be interpreted as one step above the well-defined domain/covariate shift, where the runtime not only witnesses changed covariate distribution but also missing values.
Note that Challenge 3 specifies the input vector x only contains non-zero values for known attributes, which can be conveniently achieved via common preprocessing steps like rescaling, exponential, and normalization.
In short, in our definition, domain shift happens when covariate distribution shifts, while domain corruption is caused by the co-occurrence of domain shift and missing values.
Compared with conventional domain shift, runtime domain corruption more aggressively challenges the generalizability of the trained counterfactual prediction model, because variables deemed important in training might no longer be present during inference, and the domain-invariant patterns are unable to be mapped to those missing variables.
Furthermore, a high corruption rate of runtime variables can make the counterfactual predictor learned on full training data incur large generalization errors, and rendering data imputation impractical.
Though one can consider discarding the corrupted attributes in the training set, it may lead to an underfitting issue and also it is impractical as the corrupted attributes and differ among individuals (e.g., users may choose to withhold different personal information).
The Use of Zero-Padding.
Specifically, we do not discard inaccessible variables during runtime, which will cause dimensionality reduction.
Instead, we pad zeros to entries that correspond to missing variables such that the dimensionality is kept unchanged.
Also, zero-padding is a more feasible solution in real applications, as each runtime instance x may have an arbitrary number and combination of attributes missing, rendering it impractical to train a specific latent feature extractor for each case if the missing features are abandoned.
In contrast, zero-padding is a more flexible and scalable approach for learning domain-invariant latent representations with a shared feature extractor, where all unknown variables padded with zeros will be filtered out during projection.
This work focuses on causal inference using the Neyman-Rubin potential outcome framework (Neyman, 1923; Rubin, 2005) under the runtime domain corruption circumstance.
In this work, we aim to learn a robust, causal, and domain-invariant latent representation z of variable x , for which the latent distributions across various domains are well-balanced to counter the aforementioned three challenges simultaneously.
Our main contributions are:
• We identify an important yet largely understudied setting in causal inference, namely runtime domain corruption where each unit tends to have an incomplete set of covariates during inference besides the potentially shifted covariates.
In our paper, we propose an investigation of runtime domain corruption, confoundedness, and selection bias in a unified view.
• We derive the upper bound of the generalization error by extending the in-sample causal inference to the corrupted out-of-sample scenario.
To efficiently optimize the multiple Kullback-Leibler (KL) divergence terms, we propose a two-stage domain adaptation scheme, namely the Variational autoEncoder Generative Adversarial Network (VEGAN) for unifying multiple inter-domain distances.
• We compare VEGAN to state-of-the-art baselines for performing predictions on both insample covariates and out-of-sample, corrupted runtime covariates.
The empirical results demonstrate our model’s stronger robustness to runtime domain corruption.
R ELATED W ORK
We relate our work to the representation learning branch in causal inference, which is overlapped with the domain adaptation field due to its unique counterfactual nature.
With the strong represen- tation learning ability of deep learning (Bengio et al., 2013; LeCun et al., 2015), new works are proliferating by leveraging the deep learning framework to learn the latent representation on top of the observed covariates.
The TARNet (Shalit et al., 2017) builds a shared feature extractor followed by a two-headed neural network to model the outcome for each type of treatment separately.
It also incorporates the integral probability metric (IPM), e.g., Wasserstein distance (Vallender, 1974) or maximum mean discrepancy (MMD) (Gretton et al., 2012), to minimize the distance of the learned latent covariate distribution between treated and control groups to mitigate the selection bias.
Following that, a variational autoencoder (VAE) framed CEVAE model (Louizos et al., 2017) emphasizes handling the confounding problem by building robust latent representation, and its performance is stated to be more robust than many previous methods, and dragonnet (Shi et al., 2019) leverages the neural net-enhanced propensity estimation and the innovative targeted regularization for causal effect estimation.
In addition, other works such as GANITE (Yoon et al., 2018) and DeepMatch (Kallus, 2020) adopt generative adversarial network (GAN) (Goodfellow et al., 2014) and build their own designated GAN learning network.
Our work differs from theirs and we relate our work further to the unsupervised domain adaptation (Blitzer et al., 2007) by building a GAN-integrated
VAE model, as we additionally consider the domain corruption situation where the trained model’s performance could dramatically decline in causal effect estimation.
In addition, it is noted that Jesson et al.
(2020) propose an uncertainty estimation plug-in to the state-of-the-arts such as TARNet, CEVAE, and Dragonnet to allow these models to estimate epistemic uncertainty in high-dimensional conditional average treatment effect (CATE) estimation, thus to inform the decision maker to be vigilant when making recommendations if the high uncertainty present.
It considers the domain shift during runtime, but it emphasizes making no treatment recommendation if the epistemic uncertainty exceeds a certain threshold.
Hence, our work differs from it as we focus on more accurate treatment effect estimation.
It should also be noted that some existing works (Qu & Lipkovich, 2009; Mayer et al., 2020; Berrevoets et al., 2022) have been proposed for treatment effect estimation with missing values, where the core is to leverage imputation algorithms to handle the missing values.
Since runtime domain corruption also includes domain shift, the imputed target domain data could still deviate heavily from the source domain, rendering those methods inaccurate in such conditions.
Furthermore, imputation algorithm falls short in accurate imputing when the number of missing value is large, it even becomes useless when attributes that are completely missing at distribution level during testing.
We also relate our work to algorithmic fairness topics, e.g., disparate learning processes (DLPs), in which the ethically concerned, privacy-related features are not available or impermissible to be used during runtime (Lipton et al., 2018; Simons et al., 2021).
A similar approach to DLPs is a doubly-robust counterfactual prediction model with additional handling of the confounding problem during training (Coston et al., 2020).
However, its causal inference differs from the common causal effect estimation as it assumes that one of the potential outcomes is a known constant for a binary treatment, and is hence inapplicable to the problem studied in this paper.
P ROBLEM S ETUP AND A NALYSIS
For simplicity, we consider binary treatment t of 1 or 0 to denote the treated group and the control group, respectively.
The individual treatment effect (ITE) for a variable vector x is defined as:
[Equation]
where y ∗ 1 and y ∗ 0 are potential outcomes with treatment t = 1 and t = 0 respectively.
To validly estimate τ ( x ) , the following unconfoundedness assumption is required.
Assumption 1 (Unconfoundedness): Treatment assignment is independent to the potential outcome given the pre-treatment covariate x , i.e., t ⊥⊥ { y ∗ 0 , y ∗ 1 }| x .
Let Ψ : X × { 0 , 1 } → R parameterizes the hypothesis, our goal is to build the regression model Ψ 1 − t ( x , t ) =
E [ y | X = x , T = 1 − t ] with observed outcome y which can accurately recover the counterfactual outcome for x with treatment t , thus the causal effect can be calculated as Ψ 1 − Ψ 0 .
However, Challenges 1 and 2 arise when the observational dataset does not follow the RCT standard, making the trained models Ψ 1 and Ψ 0 unable to reflect the true treatment outcome for each x .
We perceive the observed covariates of treated and control groups from the conventional domain shift perspective, in which covariate x is a noisy measurement, normally less informative and more confounded (Griliches & Hausman, 1986; Maddala & Nimalendran, 1996), than the domaininvariant latent representation z .
Therefore, the unconfoundedness changes from t ⊥⊥ { y ∗ 0 , y ∗ 1 }| x to t ⊥⊥ { y ∗ 0 , y ∗ 1 }| z .
In addition to the treated and control groups from the in-sample set, this paper considers the runtime causal effect estimation where the out-of-sample set is affected by domain corruption (Challenge 3), where the trained model’s generalizability is significantly challenged.
3. T ARGET D OMAIN E RROR U PPER B OUND
Shalit et al.
(2017) have shown that the expected precision in estimation of heterogeneous effect (PEHE), or ϵ PEHE , is upper-bounded by both the trained model error ϵ F on factual outcomes and the distance between treated and control distributions, measured by integral probability metric (IPM).
Since their derived upper bound for ϵ PEHE does not consider runtime domain corruption on out-ofsample variables, we fill the gap by deriving the bound in Theorem 1.
Theorem 1: Let ϕ : X → Z be the invertible latent representation mapping function (a.k.a. feature extractor) with inverse Φ .
Let Ψ : Z × { 0 , 1 } → R be the updated hypothesis.
Let F be a family of functions f : Z → R , f ∈ F .
The source domain is the observational data for treated and control groups, and the target domain is the runtime test/inference set with corrupted variables.
Under the conditions of Assumption 8.1, Lemma 8.1, Theorem 8.1, and the upper breakdown of absolute value inequality, we derive the upper bound of target domain error (i.e., generalization error) as 1 :
[Equation]
where ϵ t F denotes the factual training error, P tϕ is the probability measure within treatment group t in the training set, ϵ tr PEHE and ϵ sr PEHE respectively indicate the target and source domain errors, P trϕ and
P srϕ are probability measures which denote the covariate distribution in target domain and source domain respectively, and B ϕ is a bounded constant.
We provide the proof in in Appendix 8.1.
The upper bound given in Theorem 1 suggests that, to bring down the runtime/target domain error ϵ tr PEHE , we are essentially minimizing: (1) the prediction errors on observed outcomes; (2) the imbalance between treated and control groups; (3) the discrepancy between the training and test sets altogether.
It guides our algorithm design in general for runtime causal inference.
Note that if no domain corruption exists, which means P trϕ = P srϕ and thus
IPM ( P trϕ , P srϕ ) = 0 , the runtime error becomes identical to the source domain error ϵ tr PEHE =
ϵ sr PEHE .
V ARIATIONAL I NFERENCE
Our solution is built upon the variational autoencoder (VAE).
To start with, in this section we introduce the minimization of the factual error ϵ F , the distribution disparity between treated and control groups first, and between training and runtime domains afterwards.
4. E VIDENCE L OWER B OUND
For modelling the observed treatment outcome y , we use the maximum likelihood estimation (MLE) to approximate the parameters.
For simplicity, log is commonly used to decompose the joint marginal likelihood p
( y )
into:
[Equation]
where y =
[ y 1 , y 2 , . . .
, y N ] is a vector hosting all N samples’ ground truth, N = N 1 + N 0 , Nand N 0 respectively denote the number of samples in treated and control groups.
Thus, to maximize the joint marginal log-likelihood of observing y , we can maximize the individual log-likelihood log p ( y | t ) , and the expression log p ( y | t ) itself serves as a general term to represent a distribution of all observed y .
As we assume that there exists a latent representation z and treatment t that causally determine the observed treatment response y t , i.e., y t ∼ p ( y | z , t ) in a probabilistic way, while the observed proxy x does not have any causal relations but statistical correlations with y .
Due to the potentially high dimensionality of z , the marginal likelihood p t ( y ) =
p ( y | t ) of treatment group t is intractable.
Here, we apply the variational methodology (Kingma & Welling, 2013) to our scenario to tackle p ( y | t ) by establishing an encoder network ϕ t to learn latent representation z t ∼ p ϕ t ( z | x ) , and a decoder network Ψ t to estimate treatment response y t ∼ p
Ψ t ( y | z , t ) .
The detailed probabilistic mechanism to model p Ψ t ( y | z , t ) is described in Appendix 8.2, which leads us to the evidence lower bound (ELBO t ) as follows:
[Equation]
where P ϕ t and P Ψ t are posterior and prior distributions respectively.
D KL ( · ) returns the Kullback–Leibler (KL) divergence between two distributions.
As such, the task of maximizing the intractable log p Ψ t ( y | t ) can be indirectly solved by pushing up its associated ELBO t , thus minimizing the factual error ϵ t F .
According to the decomposition in Eq. 3, our objective is to maximize the sum of two ELBOs for treated and control groups:
[Equation]
It is worth noting that, our derived bound ELBO can be easily extended from our binary treatment setting to scenarios that involve multiple treatments.
4. T REATED /C ONTROL D OMAIN A DAPTATION
According to the second term in Eq. 4, for t ∈ { 0 , 1 } , we have both KL divergence terms that regularize the posterior distribution P ϕ t and the prior distribution P Ψ t as follows :
[Equation]
By pushing up the ELBO in Eq.
5, one can notice that both posteriors P ϕ 1 and P ϕ 0 are regularized to approach the same prior distribution if we set the prior P Ψ 1 and P Ψ 0 to be the same, e.g., standard normal distribution N ( 0 , 1 ) .
Thus, the domain adaptation (DA) for both groups can be naturally achieved to balance their latent distributions and counter Challenge 2 by adjusting the priors using the VAE framework.
It is worth noting that KL divergence is an unbounded asymmetric distribution distance measure (Kullback & Leibler, 1951) which does not belong to IPM, so we replace it with a bounded symmetric distribution similarity measurement in Section 5 as a better approximation.
4. T RAINING /R UNTIME D OMAIN A DAPTATION
In addition to the DA across treated and control groups within the training set, we would also like to do DA between the entire training and runtime sets to minimize the tightness bound
B ϕ IPM F ( P trϕ , P srϕ ) given in Lemma 8.1 and thus alleviate Challenge 3.
As such, for a well-trained model, we aim to make the out-of-sample performance as good as the in-sample performance, i.e., the out-of-sample results would not deviate from the in-sample ones drastically while keeping good in-sample performance.
Intuitively, if the VAE prediction framework is applied to the full runtime test set { ( x trj , t trj , y trj ) } N ′ j = (i.e., target domain in this context), one can end up with the objective to be maximized similar to the
ELBO t presented in Eq.
4 as follows:
[Equation]
However, the label y tr and treatments t tr are apparently unknown in practice, and such an objective cannot be optimized.
Since the only available information is the runtime covariates which can be used to extract the domain-invariant representation from DA, the second term in Eq.
7 can be utilised for such purpose with a mild modification.
To be specific, we alternatively walk around to minimize the KL divergence between the runtime posterior P tr ϕ and the entire training set posterior P sr ϕ , namely D KL ( P tr ϕ || P sr ϕ ) , where ϕ is a shared
𝑿 𝒕=𝟏𝒔𝒓 𝑿 𝒕=𝟎𝒔𝒓 𝑿 𝒕𝒓
Shared Feature Extractor 𝐺 ∅
Treated, Control, Discriminator 𝐷 𝛿
Train, Test, Discriminator 𝐷 𝛽
Treated Predictor Ψ
Control Predictor Ψ𝝁 𝒕=𝟏𝒔𝒓 𝝈 𝒕=𝟏𝒔𝒓 𝝁
𝒕=𝟎 𝒔𝒓 𝝈 𝒕=𝟎𝒔𝒓 𝝁 𝒕𝒓 𝝈 𝒕𝒓  𝑁 𝑦ො𝑦ො𝑝 𝑖 𝑝 𝑗 ′
MLP 𝜇
MLP 𝜎 feature extractor.
Thus, to achieve the second-stage DA, our proposed ultimate evidence lower bound (ELBO ulti ) for the intractable joint log-likelihood p ( y ) is:
[Equation]
A DVERSARIAL L EARNING
So far, we have multiple D KL ( · ) terms: two in Eq.
6 align the posteriors to the same prior N ( 0 , 1 ) , and the one in Eq. 8 that aligns the posteriors of training set and runtime set.
As the direct calculation of KL divergence is computationally inefficient and may even be infeasible with high dimensional data (Nelken & Shieber, 2006; Moral et al., 2021), we propose to implicitly minimize them and unify these terms into a compact generative adversarial network (GAN) (Goodfellow et al., 2014) shown in Figure 1.
Apart from that, optimizing the minimax game in GAN is equivalent to minimizing the Jensen-Shannon divergence (Goodfellow et al., 2014), which is a bounded symmetric distribution similarity measurement (Nielsen, 2019).
Such technique resonates with adversarial variational
Bayes introduced in Mescheder et al.
(2017), while our motivation and implementation differ.
Here, we present our V ariational auto E ncoder G enerative A dversarial N etwork runtime counterfactual regression model, coined as VEGAN.
In what follows, we unfold the design details of VEGAN.
Firstly, we instantiate ϕ , the shared feature extractor among x srt =1 , x srt =0 and x tr .
It includes G ϕ and the following two multi-layer perceptrons (MLPs) that map all the data from the original R d into latent space R l .
Due to the variational nature of the model, the j -th latent dimension of individual i , denoted by x sri , is modelled by a Gaussian distribution with its dedicated mean µ ij and variance σ 2 ij as follows:
[Equation]
where mean µ ij and standard deviation σ ij are respectively the j -th element of latent representations µ i and σ i .
In VEGAN, µ i , σ i ∈ R l are denoted as:
[Equation]
which allows us to obtain the latent representation
z i for the subsequent DAs and inference.
Secondly, for the treated/control group DA, we propose an adversarial way to implicitly reduce inter-domain distribution distance.
In this regard, G ϕ is essentially our generator for notation simplicity, and a discriminator D δ is thus designed to pair up the generator to facilitate adversarial learning.
The minimax game is designed as: the discriminator D δ tries to differentiate the standard Gaussian sample n ∼ N ( 0 , 1 ) from z sr learned from the training sample; in the meantime, feature extractor G ϕ tries to update the latent representation z sr to make it indistinguishable from n .
When an equilibrium state is reached, the treated and control domains are well adapted because both latent representations z sr fall in the same distribution.
Thus, the two terms D KL ( P ϕ 1 || P Ψ 1 ) and
D KL ( P ϕ
|| P Ψ 0 ) , are minimized in an adversarial way.
As Figure 1 shows, the output p i
=
D δ ( w i ) is the scalar probability of being a Gaussian sample, where w i
=
η i n i +(1 − η i )
z sri with η i ∈ { 1 , 0 } labeling the i -th sample from two buckets (1 for Gaussian samples, and 0 for training samples).
Note that n is resampled for every training case in every epoch.
In our supervised learning setting, we have the cross-entropy loss for the discriminator D δ ( · ) :
[Equation]
Then, in an adversarial setting, the minimization of two KL divergence in Eq.
6 for treated/control domain adaptation is replaced by the following:
[Equation] [Equation]
Analogously, for the train/runtime domain adaptation, we design another discriminator D β ( · ) to form the second GAN system between G ϕ ( · ) and D β ( · ) , where D β ( · ) predicts the probability p ′
j of the sample from the source domain, i.e., training set.
The only difference from the first GAN system is that, it takes the training sample as real while the runtime sample is treated as fake.
Thus,
D KL ( P tr ϕ || P sr ϕ ) is replaced by the followings:
[Equation]
[Equation]
Finally, to build the probabilistic model p ( y | z , t ) , we model each of the treatment classes through two separate MLPs, namely Ψ 1 and Ψ 0 respectively, thus the general representation of modelling the observed outcome y for individual i is given as p Ψ t ( y i
|
z i , t i ) =
N (ˆ µ i , ˆ σ 2 i ) , where ˆ µ t,i =
Ψ t i ( z sri , t i ) , and we follow Louizos et al.
(2017) to set ˆ σ 2 i = 1 for simplicity.
In a nutshell, to promote a computationally efficient algorithm, we propose to minimize the following loss function L along with optimizing the minimax game together such that the ELBO ulti in Eq.
8 will be maximized:
[Equation]
where L = − (cid:16) E P sr Ψ
[log p Ψ 1 ( y | z , t =
[Equation]
] + E P sr Ψ0 [log p Ψ 0 ( y | z , t =
[Equation]
)] (cid:17) .
We summarize our VEGAN model optimization scheme in Appendix 8.11.
Please note that the notation changed accordingly as the reparameterization trick z = ω ( µ ϕ , σ 2 ϕ , ϵ ) (detailed in Appendix 8.3) is applied as a necessity to get the gradient ∇ ϕ for the feature extractor G ϕ .
Also, the original minimax game in Eq.
14 is adjusted to the double minimization tradition for gradient descent.
E XPERIMENTS
We perform evaluation on two datasets, one is the widely used semi-synthetic dataset IHDP (Infant
Health Development Program)
(Hill, 2011), and the other is the high-dimensional synthetic dataset from ACIC (Atlantic Causal Inference Conference) (ACIC, 2019).
Our models, VEGAN I (with only treated/control domain adaptation) and VEGAN II (with both adaptation steps), are implemented using Pytorch (Paszke et al., 2019), and their hyperparameter settings are left in Appendix 8.9.
We compare our models with seven state-of-the-arts: TARNet (Shalit et al., 2017), BTARNet (Jesson et al., 2020), CFR WASS (Shalit et al., 2017), SITE (Yao et al., 2018), Dragonnet Base (Shi et al., 2019),
Dragonnet TR (Shi et al., 2019), CEVAE (Louizos et al., 2017) with their provided default settings.
We adopt two well-established metrics: (1) ϵ CATE (Hill, 2011), i.e., error of conditional average treatment effect (CATE) for overall performance; and (2) √ ϵ PEHE (Shalit et al., 2017), i.e., precision in estimation of heterogeneous effect (PEHE) for individual-level performance.
Notes on Additional Results Appended.
As we are interested in causal inference under runtime domain corruption, we lay more emphasis on the out-of-sample predictions with corrupted covariates at the individual level.
The details for generating out-of-sample test data with domain corruption are in Appendix 8.4 along with the analysis on the varying magnitudes of the Gaussian noise, where we use corruption level (CL) to denote the average ratio of corruption per selected variable distribution.
[Table] [Table]
We leave CATE prediction error ϵ CATE to Appendix 8.10, where VEGAN demonstrates the superiority of its DA schemes.
In the main results, we focus on individual-level causal effect estimation measured by √ ϵ PEHE .
To demonstrate the generalizability of our proposed second stage adversarial as a plug-in, we study its compatibility with the most representative baseline TARNet in Appendix 8.5.
Additionally, the training stability of VEGAN is analyzed in Appendix 8.6.
6. P ERFORMANCE ON IHDP
The semi-synthetic dataset contains 25 real-world covariates relating to characteristics of the newborns and their mothers, and we use the non-linear response surface B for simulated outcomes designed by Hill (2011).
In IHDP, we selectively corrupt 7 privacy-related features described in
Appendix 8.8 with CL ∈ { 120 , 18 , 15 , 13 , 1 } .
Table 1 shows the in-sample predictions, for which our models perform the best in estimating CATE while staying competitive on individual-level causal effect estimation.
For out-of-sample predictions, we conduct the tests on 5 randomly corrupted test sets with ratio CL ∈ { 120 , 18 , 15 , 13 , 1 } .
Notably, CL = 1 represents an extreme case where all the sensitive features are completely inaccessible during runtime and only the remaining 18 variables are available for prediction.
As Table 2 demonstrates, VEGAN II yields the second best performance when the domain corruption is relatively restrained, but obtains the highest accuracy after the corruption ratio increases to 15 .
The best baseline is CFR WASS when CL is low, but it overfits the training set significantly and thus does not generalize to a higher domain corruption level, while VEGAN II is less sensitive to the stronger corruption of private variables.
VEGAN II also outperforms its base model VEGAN I consistently owing to the additional DA between training and runtime domains.
6. P ERFORMANCE ON ACIC
We use the fully-synthetic high-dimensional dataset with 100 simulations from ACIC (2019), which includes 200 features and 1,000 samples from each simulation.
Because there is no clear definition for sensitive features, we allow the corruption to take place for all the features in ACIC dataset with a ratio of CL ∈ { 120 , 18 , 15 , 13 } .
With this, we can mimic situations where individuals can withhold an arbitrary combination of variables in privacy-sensitive applications.
Note that, we omit CL =in ACIC dataset as it will set all variables to zero and thus make any predictions infeasible.
This setting also allows for additional comparisons with the imputation method, which is a natural choice to handle missing values.
Specifically, we implement the imputation algorithm MICE (Van Buuren & Oudshoorn, 2000) widely adopted in treatment effect estimation (Mayer et al., 2020; Berrevoets et al., 2022) to enhance all our baselines.
Given the space limit, we only include the the most performant baseline CEVAE after enhancement (denoted as CEVAE-IMP) in Tables 3 and 4.
The full experiments and analysis with imputation for all baselines are in Appendix 8.7.
As a result, VEGAN outperforms all the other models for both in-sample and out-of-sample prediction as shown in Table 3 and Table 4.
We notice that when the domain corruption level climbs, the fluctuations of prediction errors are small in magnitude.
Volatility Analysis.
To better quantify the advantage of VEGAN under domain corruption on ACIC, in Figure 2 we analyse the deviation ( ∆ ) of each model’s performance between in-sample and corrupted prediction tasks, i.e., ∆ = 100% × | ϵ in-sample − ϵ corrupted | /ϵ in-sample .
∆ quantifies the instability of the model, as we commonly rely on models obtained with the training set and prefer lower gener-
[Table]
[Table] [Table]
alization errors.
We include CEVAE-IMP’s results while leaving other imputation-enhanced baselines to Appendix 8.7.
All models become more volatile as CL increases, while VEGAN II maintains an excellent stability with only 0.22% variation at corruption level 13 and achieves the best accuracy in terms of √ ϵ PEHE .
The exact numerical values of ∆ are recorded in Appendix 8.10.
Comparison with Data Imputation.
When the corruption rate is low, using imputation is generally helpful for slightly increasing the prediction performance, but the improvements are smaller when more attributes are corrupted (e.g., at a rate of 1/5 and 1/3).
Also, as 8.7 suggests, imputation methods do not make a difference in volatility, as the imputed target domain is still subject to significant distribution shift.
In short, data imputation has very limited benefit under the domain corruption setting.
Furthermore, in scenarios where an attribute is completely missing for all instances (e.g.,
CL = 1 in IHDP), it is infeasible to impute this attribute for prediction.
6. C OMPUTATIONAL E FFICIENCY OF VEGAN
[Figure]
One core motivation of utilizing GAN to replace the straightforward
KL divergence optimization is to preserve training efficiency with high dimensionality.
Hence, we further test VEGAN’s efficiency by comparing its training time (in seconds) consumption per 100 epochs with CEVAE, which is the based method VEGAN improves upon.
The test is performed on 12th Gen Intel i7-12700K 12-Cores 20Threads CPU on Ubuntu 22.04 LTS.
Figure 3 shows that as a GANbased alternative to CEVAE, VEGAN I , has significantly faster training speed (over 6 × speedup).
Furthermore, the introduction of our second-stage adaption in VEGAN II is still able to maintain high computational efficiency, witnessed by over 4 × speedup over CEVAE.
C ONCLUSION
This paper formalizes the runtime causal inference problem under domain corruption, where novel strategies are proposed to counter the imbalance between treated and control groups and the interdomain discrepancy.
We further adopt adversarial learning to implicitly reduce the distribution disparity for the sake of efficiency.
For our proposed approach VEGAN, its performance exceeds other state-of-the-arts under the runtime domain corruption setting in semi-synthetic and full-synthetic benchmark datasets.
In addition, the second stage adversarial plug-in is demonstrated applicable to off-the-shelf models to reduce generalization errors.
A PPENDIX 8. P ROOF OF T HEOREM
To proof Theorem we need some preliminary definitions and the proof of Lemma
1 beforehand
Definition
1 Let F be a family of functions, we have f : Z → R , f ∈ F .
The distribution distance measure - integral probability metric (IPM) between two distribution P sr and P tr is defined as:
[Equation]
Definition 8.2 Let Ψ : X × { 0 , 1 } → R be the hypothesis, the estimated individual treatment effect for unit x is:
[Equation] [Equation]
Here we slightly extend the expected PEHE defined by Shalit et al.
(2017) by replacing the squared loss of causal effect estimation with the expected squared loss defined in Definition 8.3 for the proof’s sake.
Definition 8.3 Let ϕ : X → Z be the inevitable latent representation mapping function with inverse Φ .
Let Ψ : Z × { 0 , 1 } → R be the updated hypothesis.
The expected squared loss l ϕ, Ψ ( x ) of unit x is defined as:
[Equation]
where τ is the true treatment effect defined in Eq.
1 and L (ˆ τ, τ ) =
(ˆ τ − τ ) 2 .
Definition 8.4 The expected Precision in Estimation of Heterogeneous Effect (PEHE) of ϕ and Ψ with expected loss function of unit x is defined as:
[Equation]
where l ϕ, Ψ ( x , t ) is defined in Definition 8.3 .
Lemma 8.1 : Let ϕ : X → Z be the invertible latent representation mapping function with inverse Φ .
Let Ψ : Z × { 0 , 1 } → R be the updated hypothesis.
Let F be a family of functions f : Z →
R , f ∈ F .
Assume we have B ϕ > 0 s.t. 1 B ϕ l ϕ, Ψ (Φ( z ))
∈ F .
The tightness of target domain error w.r.t.
the source domain one is bounded by the distribution distance denoted by IPM:
[Equation]
where ϵ tr PEHE and ϵ sr PEHE indicate target domain error and source domain error respectively, P trϕ and
P srϕ denote covariate distribution in target domain and source domain respectively, and B ϕ is a bounded constant.
We denote the expected PEHE in Eq.
18 in target domain and source domain as ϵ tr PEHE and ϵ sr PEHE respectively, also tr and sr indicates the test set and training set in our main text where p tr ( x ) ̸ = p sr ( x ) if domain corruption exists.
We firstly present the proof for Lemma 8.1 as follows:
Proof.
[Equation]
The first equality is by Definition 8.4, the second equality is by change of variable, the first inequality is by the premise that 1 B ϕ l ϕ, Ψ belongs to the function family F , the fourth equality is by
Definition 8.1, the last equality is by the property that IPM is non-negative.
Theorem 8.1 Let ϕ : X → Z be the invertible latent representation mapping function with inverse Φ .
Let Ψ : Z × { 0 , 1 } → R be the hypothesis.
Let F be a family of functions f : Z → R , f ∈ F .
Assume we have B ϕ > 0 s.t. 1 B ϕ l ϕ, Ψ (Φ( z ) , t ) ∈ F (Shalit et al., 2017), we have:
[Equation]
21) where ϵ t F denotes the factual training error within treatment group t, P tϕ is probability measure for treatment group t, and σ 2 y ∗ is the variance of potential outcome.
Assumption 8.1 Stable Unit Treatment Value Assumption (SUTVA): a)
The potential outcomes for any unit do not vary with the treatment assigned to other units.
b) For each unit, there are no different forms or versions of each treatment level, which lead to different potential outcomes (Imbens & Rubin, 2015) .
Proof.
Under the conditions of unconfoundedness assumption t ⊥⊥ { y ∗ 0 , y ∗ 1 }| ϕ ( x ) and STUVA, the potential outcome y ∗ t is determined by x via function g , namely y ∗ t = g ( x ) , resulting the single unit expected squared loss equals the squared loss, and thus the expected PEHE defined in Definition 8.4 equals the one defined by Shalit et al.
(2017), it also results in the variance σ 2 y ∗ = 0 in Eq. 21.
To conclude, the Theorem 1 is thus proofed under the further conditions of Lemma 8.1 and Theorem 8.1 as follows:
[Equation]
We align the function family F to the one Shalit et al.
(2017) use as different choice of function family F will require different assumption about the joint distribution p ( z , t, y 1 , y 0 ) , the representation mapping function ϕ , and the hypothesis Ψ .
Thus, we share the same bounded constant B ϕ .
8. P ROOF OF E VIDENCE L OWER B OUND
We present the reasoning procedures for our evidence lower bound for log p ( y | t ) with Ψ t and ϕ t implemented for approximation as follows:
Proof.
[Equation]
The conditional prior p Ψ t ( z | t ) is actually a marginal prior p Ψ t ( z ) .
Because for each treatment group, treatment t is just a constant for all units within that group, e.g., p Ψ 1 ( z | t = 1) =
p Ψ 1 ( z ) .
8. R EPARAMETERIZATION
All the neural networks’ parameters can be directly differentiated but the shared feature extractor
G ϕ ’s due to the variational nature of the model.
For instance, to deal with the gradient from discriminator D δ ’s output:
[Equation]
Because z sr is drawn from a distribution parameterized by ϕ .
Thus, a way out is to use the “reparameterization trick” (Kingma & Welling, 2013).
The methodology is to let the generating of z be dependent of something which is independent from ϕ , preferably in the form z = ω ( µ ϕ , σ 2 ϕ , ϵ ) .
According to Eq.
9, z follows the normal distribution N ( µ ϕ , σ 2 ϕ ) parameterized by the neural nets
G ϕ , thus it can be rephrased as z = µ ϕ + σ 2 ϕ ⊙ ϵ , where ⊙ denotes the element-wise product and ϵ follows an independent standard normal distribution.
Thus,
[Equation]
Due to the reparameterization z = ω ( µ ϕ , σ 2 ϕ , ϵ ) , the expectation notation changed accordingly.
For instance:
[Equation]
8. D OMAIN C ORRUPTION S IMULATION
Our domain corruption is built one step above the conventional domain shift/covariate shift, thus it not only considers the shift in the covariate distribution but the missing value found in the spreadsheet.
Another intuitive example for our defined domain corruption scenario is that respondents 𝑥 !
𝑖 𝑥 !
𝑖 are reluctant to give privacy-related information themselves when responding to a data-collection survey, e.g., falsely report their age or just leave it blank, thus creating domain corruption in the collected dataset compared to the fine-grained one we possess in the database.
8.4. G ENERATING T EST S ET
Each of the dataset is randomly split with 3:1 ratio for train/test set, then we corrupt the dataset to simulate the domain corruption scenario following the rule: For a given domain corruption level (CL), as CL ∈ { 120 , 18 , 15 , 13 , 1 } , and chosen feature x s ∈ x
i , we firstly add a random variation drawn from Gaussian distribution N (¯ µ, 0 . 1) to x s with probability p = CL for each individual i if x s is continuous, or we flip the binary value if x s is binary, ¯ µ follows an uniform distribution U (0, 0.25).
Secondly, we randomly zero the value of x s to mimic the missing value situation with probability p =
CL for each individual
i .
The first step and second step are performed independently, thus making the randomly held-out dataset noisy.
Figure 4 depicts our domain corruption mechanism in a two-dimension spreadsheet with row representing unit i and column representing feature x s .
8.4. V ARYING THE M AGNITUDE OF A DDED G AUSSIAN N OISE
To justify the setting for generating drifted data, the Gaussian noise we used follows a normal distribution N (¯ µ, ¯ σ ) .
For each test instance, the random noise has different mean ¯ µ as it is drawn from a uniform distribution.
On the other hand, we provide additional experiments by varying the upper bound of the uniform distribution for ¯ µ (e.g., ¯ µ 1 ∼ U (0 , 0 . 25) , ¯ µ 2 ∼ U (0 , 1) , ¯ µ 3 ∼ U (0 , 1 . 5) ) and take different variance ¯ σ ∈ { 0 . 1 , 0 . 25 } at the domain corruption level 1/3.
The test results with different settings are provided in the following:
[Table] [Table]
We compare VEGAN II with TARNet, as all the other baselines built on top of TARNet are equally the same in terms of lacking the functionality to handle noised test set.
When the noise magnitude increases, the prediction error of both methods increase, but VEGAN holds a consistent and significant improvement margin compared with TARNet’s results.
8. A PPLICABILITY OF S ECOND S TAGE A DVERSARIAL P LUG - IN TO O THER B ASELINES
We build the second stage adversarial plug-in to the most representative model - TARNet.
The experiments are conducted using the ACIC dataset, and the results are presented in Tables 7 and 8:
[Table] [Table]
We denote the TARNet with adversarial plug-in as TARNet-Plugin.
As the results suggest, when the corruption level becomes higher, the benefit of the stage 2 domain adaptation will be enlarged.
When the this adversarial plug-in is in use
[Table]
it effectively helps TARNet reduce prediction risks under runtime domain corruption as the volatility of the TARNet-Plugin is stabilized at around 2%.
8. S TABILITY A NALYSIS OF VEGAN M ODELS
[Figure]
As GAN is known to be unstable, we provide stability analysis of our proposed VEGAN models in terms of prediction loss convergence and equilibrium status between feature extractor and discriminators.
Figure 5 shows the models convergence’s on root mean square error (RMSE) during training.
It is noted that a higher corruption level brings more challenges in the adversarial training, but we observe an equilibrium state in the majority of the cases.
Take the harder 1/3 corruption level in IHDP dataset as an example, both discriminators (for treated/control and training/runtime adaptations) can quickly converge to the equilibrium by returning an average binary cross-entropy loss of 0.69, which means the discriminators are completely decepted by the feature extractors, and always give 0.5 probability for the samples from each of the groups.
We occasionally have the situation where one of the discriminators outperforms the feature extractor, predominantly the discriminator for treated/control adaptation.
After repeatedly executing 100 runs for each setting, we observe that on average, this happens less frequently than 5% for ACIC dataset and 30% for IHDP dataset.
As such, training VEGAN in an adversarial setting is completely attainable.
8. B ASELINES E NHANCED WITH I MPUTATION
To compare imputation-based methods with VEGAN
[Table]
we using the MICE algorithm (Van Buuren & Oudshoorn, 2000) adopted by Mayer et al.
; Berrevoets et al.
(2022) to impute the missing values in the test set.
Notably, when the corruption level (i
[Table]
e., missing rate) reaches 1/
[Table]
the MICE algorithm melts down
[Table]
and we then turn to KNN-based imputation (Troyanskaya et al., 2001) from the corruption level of 1/3.
We provide the experiment results on the ACIC dataset with imputationenhanced baselines (noted with IMP) in Table 9 and 10:
[Table] [Table]
As shown in the tables, when the corruption rate is low, using imputation is generally helpful for slightly reduce the √ ϵ PEHE , but may provide negative impact when more attributes are corrupted (e.g., at a rate of 1/5 and 1/3).
At the same time, imputation methods do not make a difference in terms of volatility, as the imputed target domain is still subject to significant distribution shift w.r.t.
the training domain, and the volatility of the model increases rapidly as corruption level goes up.
To conclude, data imputation has very limited practicality under the domain corruption setting.
Furthermore, in scenarios where an attribute is completely missing for all instances (e.g., medical checks that are unavailable in some hospitals, which resembles the rightmost column in Table 2), it is infeasible to impute this attribute for prediction.
8. S ELECTED S ENSITIVE F EATURES IN IHDP
We selectively choose 7 features which can be considered as sensitive and thus highly possible to be corrupted: momage, sex, twin, b.marr, cig, drugs, and work dur.
The explanations are cited from
Hill (2011) as follows:
[Table]
The mother’s age when she
[Table]
birth
[Table]
the
[Table]
Baby’s
[Table]
Is
[Table] [Table]
Is the mother
[Table]
when
[Table]
birth to the baby
[Table]
cig Does the mother
[Table]
mother
[Table]
8. VEGAN A RCHITECTURE
We use the general structure for evaluating VEGAN I and VEGAN II on both IHDP and ACIC benchmarks.
G ϕ is used to represent the whole feature extractor for notation simplicity which has twoheaded output with each of them modelling the mean and variance of the latent representation respectively by multi-layer perceptrons (MLPs).
[Table]
For other hyperparameters, we set epochs = 100, convergence tolerance = 10, batch size = 100, model complexity L2 regularization = 0.01, learning rate = 0.
. 8. O THER R ESULTS
It is worth noting that 0 CL results are omitted as it resembles the in-sample predictions in Tables 1 and 3.
We leave the CATE estimation errors for IHDP and ACIC in Table 13 and Tablerespectively, and the variation ∆ (%) on ACIC in Table 15.
[Table]
[Table] [Table]
[Table]
8. O PTIMIZATION S CHEME OF VEGAN
1: Input
Train set
( x sri , t sri , y sri )
Ni runtime variable set
N ′ j optional) hyperparameters
.g learning rate initialized neural networks’ parameters
Ψ do 3:
Sample a mini-batch b m ∈ { b m } Mm =1 of size | b m | = | b t =1 m | + | b t =
m | , with | b t =1 m | = | b t =
m | , from { ( x sri , t sri , y sri ) } Ni =1 .
Sample equal size instances from standard Gaussian N ( 0 , 1 ) ; 4:
[Equation]
←
[Equation]
; 5: if runtime domain corruption exists then 6:
Randomly draw batch b ′ m from { ( x trj ) } N ′ j =1 with size | b ′ m | =
| b m | ; 7:
[Equation]
←
[Equation]
; 8: end if 9: ϕ
[Equation]
10: end while