Adversarial balancing-based representation learning for causal effect inference with observational data

Consider the information loss when transforming covariates into a latent space. The non-linear statistical dependencies between variables can be acquired by mutual information (MI) (Shannon 1948). Thus we use MI between latent representations and original covariates as a measure to account for information loss:We denote the joint distribution between covariates and representations by \(\mathbb {P}_{Xh}\) and the product of marginals by \(\mathbb {P}_{X} \otimes \mathbb {P}_{h}\). Note that, consistent with Shannon’s information theory, MI can be represented as the Kullback-Leibler (KL) divergence:It is hard to compute MI in continuous and high-dimensional spaces, but one can capture a lower bound of MI with the Donsker-Varadhan representation of KL-divergence (Donsker and Varadhan 1983):

Here, \(\mathcal {C}\) denotes the set of unconstrained functions \(\varOmega \).

Inspired by the Mutual Information Neural Estimation (MINE) idea (Belghazi et al. 2018), we propose to establish a neural network estimator for MI. Specifically, let \(\varOmega \) be a function \(\mathcal {X} \times \mathcal {H} \rightarrow \mathbb {R}\) parameterized by a deep neural network. We have:By distinguishing the joint distribution and the product of marginals, the estimator \(\varOmega \) approximates the MI with arbitrary precision. In practice, as shown in Fig. 2, we concatenate the input covariates X with representations h one by one to create positive samples (as samples from the true joint distribution). Then, we randomly shuffle X on the batch axis to create fake input covariates \(\tilde{X}\). Representations h are concatenated with fake input \(\tilde{X}\) to create negative samples (as samples from the product of marginals). From Eq. (1) we can derive the loss function for the MI estimator:Information loss can be decreased by simultaneously optimizing the encoder \(\varPhi \) and the MI estimator \(\varOmega \) to minimize \(L_{\varPhi \varOmega }\) iteratively via gradient descent.

The representations of the treatment and the control groups are denoted by \(h(t=1)\) and \(h(t=0)\), respectively. The discrepancy between the covariate distributions within the treatment and the control groups is the issue to be addressed. To decrease this discrepancy, we propose an adversarial learning method to constrain the encoder to learn treatment and control representations that are balanced distributions. We build an adversarial game between a discriminator D and the encoder \(\varPhi \), in line with the framework of Generative Adversarial Networks (GAN) (Goodfellow et al. 2014). In the classical GAN framework, a source of noise is mapped to a generated image by a generator. A discriminator is trained to distinguish whether an input sample is from the true or the synthetic image distribution generated by the generator. The logic of GANs lies in training a reliable discriminator to distinguish fake and real images, and then, using the discriminator to train the generator, which in turn generates images constructed so as to try to fool the discriminator.

In our adversarial game: According to the architecture of ABCEI, the encoder is associated with the MI estimator \(\varOmega \), treatment outcome predictor \(\varPsi \), and adversarial discriminator D. This means that the training process is iteratively adjusting each of the components. The instability of GAN training will become serious in this context. To stabilize the GAN training, we propose to use the framework of Wasserstein GAN with gradient penalty (Gulrajani et al. 2017). By removing the sigmoid layer and applying the gradient penalty to the data between the distributions of the treatment and the control groups, we can find a function D which satisfies the \(\text {1-Lipschitz}\) inequality:We can write down the form of our adversarial game:where \(P_{\text {penalty}}\) is the distribution acquired by uniformly sampling along the straight lines between pairs of samples from \(P_{h(t=0)}\) and \(P_{h(t=1)}\). The adversarial learning process is depicted in Fig. 3.

The encoder \(\varPhi \) will now be smoothly trained to generate balanced representations. We can write down the training objectives for discriminator D and encoder \(\varPhi \), respectively:

The final step for CATE estimation is to predict the treatment outcomes with learned representations. We establish a neural network predictor, which takes latent representations and treatment assignments of units as the input, to conduct outcome prediction: \(\widehat{y_t} = \varPsi (h,t)\). We can write down the loss function of the training objective as:Here, R is a regularization on \(\varPsi \) for the model complexity.

With respect to the architecture in Fig. 1, we minimize \(L_{\varPhi \varOmega }\), \(L_{\varPhi }\), and \(L_{\varPhi \varPsi }\), respectively, to iteratively optimize parameters in the global model. The optimization steps are handled with the stochastic method due to Adam (Kingma and Ba 2015), and the training of the model is done per Algorithm 1.

Metadata on the employed datasets can be found in Table 1.

ACIC (Dorie et al. 2019) The Atlantic Causal Inference Conference (ACIC) dataset is derived from real-world data with 4802 observations with 58 covariates. There are 77 datasets which are simulated with different treatment selection and outcome functions. Each dataset is generated in 100 random replications independently. In this benchmark, different settings like degrees of non-linearity, treatment selection bias and magnitude of treatment outcome are considered.

IHDP (Hill 2011) The Infant Health and Development Program (IHDP) studies the impact of specialist home visits on future cognitive test scores. Covariates in the semi-simulated dataset are collected from a real-world randomized experiment. The treatment selection bias is created by removing a subset of the treatment group. We use the setting ‘A’ in (Dorie 2016) to simulate treatment outcomes. This dataset includes 747 units (608 control and 139 treated) with 25 covariates associated with each unit.

Twins (Louizos et al. 2017) The Twins dataset is created based on the “Linked Birth / Infant Death Cohort Data” by NBER.³ Inspired by Almond et al. (2005), we employ a matching algorithm to select twin births in the USA between 1989 and 1991. By doing this, we get units associated with 43 covariates including education, age, race of parents, birth place, marital status of mother, the month in which pregnancy prenatal care began, total number of prenatal visits, and other variables indicating demographic and health conditions. We only select twins that have the same gender who both weigh less than 2000g. For the treatment variable, we use \(t=0\) indicating the lighter twin and \(t=1\) indicating the heavier twin. We take the mortality of each twin in their first year of life as the treatment outcome, inspired by Louizos et al. (2017). Finally, we have a dataset consisting of 12,828 pairs of twins whose mortality rate is \(19.02\%\) for the lighter twin and \(16.54\%\) for the heavier twin. Hence, we have observational treatment outcomes for both treatments. In order to simulate the selection bias, we selectively choose one of the twins to observe with regard to the covariates associated with each unit as follows: \(t|x \sim \text {Bernoulli}(\sigma (w^T x + n))\), where \(w^T \sim \mathcal {N}(0,0.1\cdot I)\) and \(n \sim \mathcal {N}(1,0.1)\).

MIMIC-III (Johnson et al. 2016, 2019) This benchmark is created based on MIMIC-III, a database comprised of de-identified profile and health outcome data for critical care unit patients. We select patient samples with their demographic information as well as various observed laboratory measurements by chemistry or hematology. After filtering samples with missing values, the benchmark consists of 7413 samples with 25 covariates. We investigate the effect of prescription amount in the first day of critical care unit on the length of stay in the ICU: for the binary treatment, we let 0 represent a small prescription amount and 1 a large prescription amount. The treatment outcomes are simulated by \(y|x,t \sim (w^Tx + \beta t + n)\), where \(n \sim \mathcal {N}(0, 1)\), \(w \sim \mathcal {N}(0^{25}, 0.5 \cdot (\varSigma + \varSigma ^T))\), and \(\varSigma \sim \mathcal {U}((-1,1)^{25 \times 25})\). The treatment assignments are simulated as \(t|x \sim Bernoulli(\sigma (s^Tx + m))\), where \(m \sim \mathcal {N}(0, 0.1)\) and \(s \sim \mathcal {N}(0^{25},0.1\cdot I)\).

Jobs (LaLonde 1986; Smith and Todd 2005) The Jobs dataset studies the effect of job training on employment status. It consists of a non-randomized component from observational studies and a randomized component based on the National Supported Work program. The randomized component includes 722 units (425 control and 297 treated) with seven covariates, and the non-randomized component (PSID comparison group) includes 2490 control units.

Since the ground truth CATE for the IHDP dataset and MIMIC-III benchmark is known, we can employ Precision in Estimation of Heterogeneous Effect (PEHE) (Hill 2011), as the evaluation metric of CATE estimation:Subsequently, we can evaluate the precision of ATE estimation based on the estimated CATE.

For the Jobs dataset, because we only know parts of the ground truth (the randomized component), we cannot evaluate the performance of ATE estimation. Following Shalit et al. (2017), we evaluate the precision of ATT estimation and policy risk estimation, whereIn this paper, we consider \(\pi (x^u) = 1\) when \(f(x^u,1) - f(x^u,0) > 0\).

For the Twins dataset, because we only know the observed treatment outcome for each unit, we follow Louizos et al. (2017) in using the Area Under the ROC Curve (AUC) as the evaluation metric. For the ACIC dataset, we follow Ozery-Flato et al. (2018) in using the RMSE ATE as performance metric.

We perform the comparisons with the following baselines: least square regression using treatment as a feature (OLS/\({LR_1}\)); separate least square regressions for each treatment (OLS/\({LR_2}\)); balancing linear regression (BLR) and balancing neural network (BNN) (Johansson et al. 2016); k-nearest neighbor (k-NN) as suggested by Crump et al. (2008); Bayesian additive regression trees (BART) (Sparapani et al. 2016); random forests (RF) (Breiman 2001); causal forests (CF) (Wager and Athey 2018); treatment-agnostic representation networks (TARNet) and counterfactual regression with Wasserstein distance (CFR-Wass) (Shalit et al. 2017); causal effect variational autoencoders (CEVAE) (Louizos et al. 2017); local similarity preserved individual treatment effect (SITE) (Yao et al. 2018); MMD measure using RBF kernel (MMD-V1, MMD-V2) (Kallus 2020, 2018); and adversarial balancing with cross-validation procedure (ADV-LR/SVM/MLP) (Ozery-Flato et al. 2018). We show a quantitative comparison between our method and the state-of-the-art baselines. In the comparison, we include two variants of ABCEI: by \(\text {ABCEI}^*\) we denote ABCEI without the mutual information estimation component, and by \(\text {ABCEI}^{**}\) we denote ABCEI without the adversarial learning component. All baseline methods are parameterized according to the recommended settings in the original papers.

Figure 4 displays the relative performance of ABCEI and recent balancing methods, on the ACIC benchmark dataset. As we can see, representation learning methods display a lower variance than methods based on reweighing samples on covariate space. Also, adversarial balancing methods have a lower (= better) mean RMSE in ATE estimation. ABCEI combines the benefits of both: it has the advantage of adversarial balancing as well as preserving predictive information in latent space. As a consequence, it outperforms counterfactual regression with Wasserstein distance and MMD-V1, and compared to the other baselines, performs similarly in mean but better in variance.

Experimental results on the other four datasets are shown in Tables 2, 3, 4, and 5. It would be unsound to report aggregated statistical test results over the results reported in these tables. Due to varying (un-)availability of ground truth, we must resort to reporting varying evaluation measures per dataset. It would not be appropriate to aggregate over these measures in a single statistical hypothesis test. However, one can see that ABCEI performs best in fourteen out of sixteen cases, not only by the best number in the column, but often also by a non-overlapping empirical confidence interval (\(\mu \pm \sigma \), so \(68\%\) confidence intervals) with that of the best competitor (cf. reported standard deviations). This provides evidence that ABCEI is a substantial improvement over the state of the art.

Additionally, we observe that the two out of the sixteen cases, in which ABCEI does not perform best, have several similar characteristics. They are based on in-sample measurements on the Jobs and IHDP datasets, and as Table 1 shows, these are the datasets with the smallest numbers of observations, the (joint) smallest numbers of covariates, and a relative imbalance between control and treatment group size. Among these five datasets, ABCEI always performs better on datasets with more observations, with more covariates, and with more balance between control and treatment group size, although the number of datasets is too small to claim the significance of these effects.

Due to the existence of treatment selection bias, regression based methods suffer from a high generalization error. Nearest neighbor based methods consider unit similarity to overcome selection bias, but cannot achieve balance globally. Recent advances in representation learning bring improvements in causal effect estimation. Unlike CFR-Wass, BNN, and SITE, ABCEI considers information loss and balancing problems. The mutual information estimator ensures that the encoder learns representations preserving useful information from the original covariate space. The adversarial learning component constrains the encoder to learn balanced representations. This causes ABCEI to achieve better performance than the baselines. We also report the performance of our model without mutual information estimator or adversarial learning, respectively, as \(\text {ABCEI}^*\), \(\text {ABCEI}^{**}\). From the results we can see that performance suffers when either of these components is left out, which demonstrates the importance of combining adversarial learning and mutual information estimation in ABCEI.

We adopt ELU (Clevert et al. 2016) as the non-linear activation function if there is no specification. We employ various numbers of fully-connected hidden layers with various sizes across networks: four layers with size 200 for the encoder network; two layers with size 200 for the mutual information estimator network; three layers with size 200 for the discriminator network; and finally, three layers with size 100 for the predictor network, following the structure of TARnet (Shalit et al. 2017). The gradient penalty weight \(\beta \) is set to 10.0, and the regularization weight is set to 0.0001.

In the training step, firstly we minimize \(L_{\varPhi \varOmega }\) by simultaneously optimizing \(\varPhi \) and \(\varOmega \) with one-step gradient descent. Then the representations h are passed to the discriminator to minimize \(L_D\) by optimizing D with 3-step gradient descent, in order to find a stable discriminator. Next, we use discriminator D to train encoder \(\varPhi \) by minimizing \(L_{\varPhi }\) with one-step gradient descent. Finally, encoder \(\varPhi \) and predictor \(\varPsi \) are optimized simultaneously by minimizing \(L_{\varPhi \varPsi }\).

Due to the fact that we cannot observe counterfactuals in observational datasets, standard cross-validation methods are not feasible. We follow the hyper-parameter optimization criterion in (Shalit et al. 2017), with an early stopping with regard to the lower bound on the validation set. Hyper-parameter search space details are displayed in Table 6. The optimal hyper-parameter settings for each benchmark dataset are reported in Table 7.

Assuming that the size of mini-batch is n, and the number of epochs is m, the computational complexity of our model is \(\mathcal {O}(n \cdot m \cdot ((\varPhi _h-1)\varPhi _w^2 + (\varOmega _h-1)\varOmega _w^2 + (D_h-1)D_w^2 + (\varPsi _h-1)\varPsi _w^2))\). Here \(\varPhi _h, \varOmega _h, D_h, \varPsi _h\) indicate the numbers of layers, and \(\varPhi _w, \varOmega _w, D_w, \varPsi _w\) indicate the numbers of neurons in each layer, in Neural Networks \(\varPhi , \varOmega , D, \varPsi \).

To investigate the performance of our model when varying the level of treatment selection bias, we generate toy datasets by varying the discrepancy between the treatment and control group. We draw \(8 000\) samples with ten covariates \(x \sim \mathcal {N}(\mu _0, 0.5 \cdot (\varSigma + \varSigma ^T))\) as control group, where \(\varSigma \sim \mathcal {U}((-1,1)^{10 \times 10})\). Then we draw \(2 000\) samples from \(x \sim \mathcal {N}(\mu _1, 0.5 \cdot (\varSigma + \varSigma ^T))\). By adjusting \(\mu _1\), we generate treatment groups with varying treatment selection bias, which can be measured by KL-divergence. For the outcomes, we generate \(y|x \sim (w^Tx + n)\), where \(n \sim \mathcal {N}(0^{2 \times 1}, 0.1 \cdot I^{2 \times 2})\) and \(w \sim \mathcal {U}((-1,1)^{10 \times 2})\).

In Fig. 5, we can see the robustness of ABCEI, in comparison with CFR-Wass, BART, and SITE. The reported experimental results are averaged over 100 test sets. From the figure, we can see that with increasing KL-divergence, our method achieves the most stable performance. We do not visualize standard deviations as they are negligibly small.

To investigate how our model performs in regards to the numbers of covariates or confounders, we generate synthetic datasets in which we control the number of covariates, according to the procedure outlined in Ning et al. (2020). The outcome functions are also simulated following Ning et al. (2020). The results are displayed in Fig. 6. As one would expect, all methods experience higher errors when the number of covariates increases; the error increase in ABCEI’s performance is in line with the error increase for competing methods. ABCEI has a consistently though not necessarily significantly lower error than the competitors, but more importantly for this specific experiment: ABCEI does not suffer more or less than the competitors from the increase in the number of covariates.

As shown in Fig. 6, our method, without estimating propensity scores, can perform robustly in high-dimensional settings, while preserving the valuable predictive information during covariate balancing.

To investigate the impact of minimizing the information loss on causal effect learning, we block the adversarial learning component and train our model on the IHDP dataset. We record the values of the estimated MI and \(\epsilon _{PEHE}\) in each epoch. In Fig. 7, we report the experimental results averaged over 1 000 test sets. We can see that with increasing MI, the mean square error decreases and reaches a stable region. But without the adversarial balancing component, the \(\epsilon _{PEHE}\) cannot be further lowered due to the selection bias. This result indicates that even though the estimators benefit from highly predictive information, they will still suffer if imbalance is ignored.

In Fig. 8, we visualize the learned representations on the IHDP and Jobs datasets using t-SNE (van der Maaten and Hinton 2008). Such visualizations display high-dimensional datapoints in a two-dimensional embedding, such that similar datapoints find themselves close together and dissimilar datapoints find themselves far apart. As such, the t-SNE visualization represents natural grouping behavior within the dataset. We use this visualization to interpret how well the encoders \(\varPhi \), learned by the adversarial learning methods, are capable of generating a balanced representation between the control and the treatment groups; if the encoder works well, the t-SNE visualization of the learned representation should display a consistent mix of the control and the treatment group units. It is of no importance in these visualizations what the distribution of the full dataset looks like in particular: the encoder is free to choose a representation that naturally falls apart into an arbitrary number of clusters. Of importance, however, is that the distribution of the control group naturally spans the same locations as the distribution of the treatment group: if there are no parts of the space where the control group appears but the treatment group doesn’t, and there are no parts of the space where the treatment group appears but the control group doesn’t, then the encoder has rather successfully balanced the control and the treatment groups.

Figure 8 contains t-SNE visualizations of two datasets: IHDP in the top row, and Jobs in the bottom row. For each dataset, we show three visualizations: the left column displays the t-SNE visualizations in the original covariate space, the middle column illustrates the representations learned by ABCEI, and the right column illustrates the representations learned by CFR-Wass. In the top row we see that on the IHDP dataset, both methods achieve a reasonably good matching. The biggest unmatched area is found in the CFR-Wass figure, where one can draw a circle with midpoint (5, 17) featuring quite a few control units and no treatment units. However, in the ABCEI figure, the area around midpoint \((16,-20)\) is not much better matched; there is little difference between the two. On the Jobs dataset, however, the results are quite a bit different. The two methods choose strikingly different embeddings: the ABCEI encoder goes for one long and thick connected component, while the CFR-Wass encoder goes for multiple shorter and thinner components. In itself, this is interesting but not necessarily a mark of quality: the one is not intrinsically better than the other. However, throughout ABCEI’s connected component, we can almost always find both control and treatment units nearby. Some of the components of the CFR-Wass encoder also decently balance the control and treatment unit pockets, but there are other components featuring very few treatment units, and three components featuring no treatment units at all. These clusters of control units are badly or not at all balanced by treatment units, and hence we can conclude that ABCEI achieves a better coverage of the treatment group over the control group in the learned representation space than CFR-Wass does. This showcases the degree to which adversarial balancing improves the performance of ABCEI, especially in population causal effect (ATE, ATT) inference.