Causal Inference from Small High-dimensional Datasets

Causal-Batle is a treatment effect estimator for observational studies. We assume two domains: the target-domain and the source-domain. The target-domain is the small high-dimensional dataset we extract the treatment effect from, and it is composed of covariates $X_{target}$, a binary treatment assignment $T\in\{0,1\}$, and continuous outcome $Y$. The source-domain contains the covariates $X_{source}$, which are from the same feature space as $X_{target}$. Here, we assume the source-domain is unlabeled (without treatment assignment or outcomes). Consider $X=[X_{target},X_{source}]$, and $n=n_{s}+n_{t}$, where $n_{s}$ and $n_{t}$ are the sample sizes of the source-domain and target-domain, respectively. We create a new binary variable $D$, which encodes the sample’s domain. We define the target-domain set as $\mathcal{D}_{S}=\{(x_{i},d_{i}=0):i=1,...,n_{s}\}$ and $\mathcal{D}_{T}=\{(x_{i},t_{i},y_{i},d_{i}=1):i=1,...,n_{t}\}$.

The individual treatment effect (ITE) for a given sample $i$ is defined as $\tau_{i}=Y_{i}(1)-Y_{i}(0)$, where $Y_{i}(t)$ is the potential outcomes under $t$. The challenge is each sample $i$ either has a $Y_{i}(1)$ or $Y_{i}(0)$ associated, never both. A solution is to work with the Average Treatment Effect (ATE), the focus of our work, defined as $\tau=\mathbb{E}[Y(1)-Y(0)]$. We define outcome models as $Y(t)=\mu_{t}(x)=\mathbb{E}[Y|X=x,T=t]$, and estimate ATE as $\tau=\mu_{1}(x)-\mu_{0}(x)$.

In this section, we list the underlying assumptions of Causal-Batle. To illustrate our methodology, we adopt a dragonnet architecture [19] as the backbone of our neural network. Other backbones architecture could also be adopted, such as convolutional neural networks for image-based applications.

Assumption 1: The target-domain covariates $X_{target}$ and source-domain covariates $X_{source}$ are in the same feature space.

Assumption 2: Stable Unit Treatment Value Assumption (SUTVA) [17] - the response of a particular unit depends only on the treatment(s) assigned, not the treatments of other units.

Assumption 3: Ignorability - the outcome is independent of the treatment given the covariates ($Y\perp T|X$).

Definition 1: Back-door Criterion [13] - Given a pair of variables $(T,Y)$ in a directed acyclic graph $G$, a set of variables $X\in G$ satisfies the backdoor criterion relative to $(T,Y)$ if no node in $X$ is a descendant of $T$, and $X$ blocks every path between $T$ and $Y$ that contains an arrow into $T$.

Theorem 1: Identifiability [13] - If a set of variables $X$ satisfies the back-door criterion relative to $(T,Y)$, then the causal effect of $T$ on $Y$ is identifiable.

Assumption 1 is related to our TL component to ensure both domains share the same features. According to Assumption 2, our target-domain is generated as $(Y_{i},T_{i},X_{i})\overset{\text{i.i.d.}}{\sim}D_{target}$, where $D_{target}$ is the observational distribution of the target-domain. Assumption 2 is also a standard causal inference assumption, which guarantees that one unit does not affect other units. According to Assumption 3, the observed covariates $X_{target}$ contain all the confounders of the treatment $T$ and outcome $Y$. Assumption 3 guarantees that $X_{target}$ will block all back-door paths and satisfy the back-door criterion (Definition 1). Thus, the identifiability of the treatment effect is guaranteed by Theorem 1. Note that we assume a setting with strong ignorability (Assumption 3). If that is not the case, one might need to choose a different backbone architecture and make extra assumptions to guarantee identifiability. However, while there are recent works exploring how to estimate treatment effects with unobserved confounders [20, 12, 11, 22], they still have several known limitations [15, 24]. Furthermore, working with unobserved confounders is not the focus of this paper.

Theorem 2: Sufficiency of Propensity Score [16, 19] - If the average treatment effect is identifiable from observational data by adjusting for $X$, i.e., $ATE=E[E[Y|X,T=1]-E[Y|X,T=0]]$, then adjusting for the propensity score also suffices: $ATE=E[E[Y|f(X),T=1]-E[Y|f(X),T=0]]$

According to the Sufficiency of Propensity Score, it suffices to use only the information in $X$ relevant to predicting the propensity score, $P(T=1|X)$. The propensity score is predicted by $g(.)$, and all the relevant information to predict $P(T=1|X)$ is the output of $f(.)$. For the proof, please refer to the original publication Rosenbaum and Rubin [16].

Causal-Batle builds upon existing neural network methods that estimate treatment effects. While not tied to any specific architecture, we will explain our methodology using a Dragonnet [19] as a backbone. The Dragonnet is a neural network with three heads: one to predict the treatment assignment given the input features, and the other two to predict the outcome if treated and untreated. It assumes no hidden confounders and relies on the Sufficiency of Propensity Score (Theorem 1) to estimate its outcomes, later used to estimate the treatment effects. The Dragonnet architecture has five key components: a shared representation of the input covariates $f(.)$, treatment assignment prediction $g(.)$, a targeted regularization, and the prediction of the outcome if treated and untreated. Jesson et al. [7] showed that for high-dimensional datasets, Bayesian neural networks tend to yield better results than their non-Bayesian neural network counterpart. They proposed a modification that suits most treatment effect estimators based on neural networks: a Gaussian mixture on the last layer.

where $Y$, $T$, and $X$ are the observed outcome, the treatment assigned, and the covariates, respectively, and $w$ are the neural network weights.

Up to this point, all components of the architecture address only the estimation of the outcome in a high-dimensional setting. Note that these components assume the sample size $n$ is sufficiently large to train such a neural network. It is also worth noticing that, regardless of the backbone adopted (in our case, Dragonnet), training the shared representation $f(.)$ is the most expensive step: this is the component with the largest number of weights, which receives as input the high-dimensional data and outputs a lower-dimensional representation of that data, used as input by the other components.

Causal-Batle focuses on improving the training of $f(.)$, the hardest part of the architecture. The idea is to use the source-domain to improve the representation of $f(.)$ (hence, transfer knowledge from source-domain to target-domain). With a better trained $f(.)$, the estimation of the outcome models and the propensity score would yield better results.

Causal-Batle adds two new components to the network architecture: a discriminator $d(.)$ and a reconstruction layer $r(.)$, and their corresponding terms in the loss function. The overall architecture is shown in Figure 2, with the two new components highlighted with blue-traced boards. The discriminator component $d(.)$ is a binary classifier defined as $d(f(X))=P(D=1|X)$. Therefore, for a given sample $i$, the discriminator will predict whether it belongs to the source or target-domain. The adversarial loss incentives $f(.)$ to extract patterns common to both domains and fool $d(.)$. The adversarial loss also helps remove spurious correlations and potential biases present in only one of the domains, improving $f(.)$’s representation. The reconstruction component $r(.)$ is to ensure $f(.)$’s representation is meaningful. For the reconstruction component $r(.)$, in our work, we adopt an autoencoder for simplicity, but we point out that one could replace it with more sophisticated methods, such as variational autoencoders.

Similarly to Jesson et al. [7], Causal-Batle adopts a Bayesian neural network with a Gaussian Mixture for the outcomes (See Equation 1), and for the propensity score, with a Bernoulli Distribution:

Causal-Batle does not assume a distribution for the discriminator $d(D|X)$. Instead, it adopts a neural network classification node, which allows it to use traditional losses for adversarial learning.

The exact loss function will depend on the backbone architecture adopted. Here, we will present the loss function for the example given in Figure 2, which is a bayesian dragonnet architecture. Note that we have $n_{t}$ samples in the target-domain, $n_{s}$ samples in the source-domain, and $n=n_{t}+n_{s}$

There is a loss associated with each grey/blue square. Starting with the outcome model loss: we only observe the outcome $Y$ for either $q(0,.)$ or $q(1,.)$, with $q()$ representing a distribution. Therefore, we define the loss as the log probability:

The propensity score loss is also calculated using log probability.

To capture the losses associated with the transfer learning components, we adopt a second binary cross entropy loss for the discriminator. Defining $\hat{D}_{i}=\hat{d}(X_{i})$:

and to compete with the discriminator loss, we have the adversarial loss:

Finally, the reconstruction loss ensures the features extracted by $f(.)$ are meaningful. Denoting by $\hat{X}$ the reconstruction of input $X$, we have:

The final loss is defined as:

where the array $\alpha=[\alpha_{0},...,\alpha_{4}]$ are the losses’ weights.

Equation 1 shows how the treatment effect is estimated from a Gaussian Mixture. To calculate the Average Treatment Effect (ATE), we use the mean value $\mu_{t}(x,w)$ of the Normal Distribution (see Equation 1). The ATE can then be estimated as:

This section discusses the relationship between the source and target-domain. Causal-Batle aims to use the unlabeled source-domain to learn a better representation of the labeled target-domain. There are three losses associated with the transfer learning: the discriminator loss $\ell_{d}$, the adversarial loss $\ell_{a}$, and the reconstruction loss $\ell_{r}$. Here, we discuss the impact of $P(X_{target})$ (target-domain distribution) and $P(X_{source})$ (source-domain distribution) in our proposed methodology.

• •

  $P(X_{target})=P(X_{source})$: $\ell_{d}$ will be large, as the discriminator will have a hard time trying to distinguish elements from two identical distributions while working against the adversarial loss. The $\ell_{r}$ loss, on the other hand, will thrive. A large collection of samples from the same distribution will help $f()$ to find a good representation of the input data, lowering the $\ell_{r}$.

• •

  $P(X_{target})\neq P(X_{source})$: $\ell_{d}$ will be low, as $d(.)$ would be able to identify these differences and use them. Still, $\ell_{a}$ would attempt to fool $d(.)$, which would push $f(.)$ to learn a domain agnostic representation.

This means that, in practice, regardless of the $P(X_{target})$ and $P(X_{source})$ distributions, the losses would push $f(.)$ to have a meaningful representation of $X$. The key is to find a balance between the adversarial loss, the discriminator loss, and the reconstruction loss, which, as shown in Equation 8 is controlled by the weights (hyper-parameters) $\alpha$.

We adopted the three datasets to validate our proposed method. Please check our Appendix for an extended description of the datasets:

GWAS [22, 1]: The Genome-Wide Association Study (GWAS) dataset is a semi-simulated dataset that explores sparse settings. Proposed initially to handle multiple treatments, we adapt it to have one binary treatment $T$ and a continuous outcome $Y$. We split the data into two datasets, as explained in Setting 1. Note that $X_{source}$ and $X_{target}$ have the same distribution. We simulated a small high-dimensional dataset with 2k samples and 10k covariates.

IHDP[6]: The Infant Health and Development Program (IHDP) is a popular causal inference semi-synthetic benchmark dataset with binary treatment and continuous outcome. This dataset has ten repetitions, with 23 covariates and approximately 747 samples. We adopted Setting 1. Note that this dataset is not high-dimensional; still, it is a classic dataset in Causal Inference.

CMNIST [8]: A semi-synthetic dataset that adapts the MNIST [10] dataset to causal inference. It contains a binary treatment and continuous outcome. We follow the setup from Jesson et al. [8] with small adaptations to contain a source and a target-domain. Jesson et al. [8] generates the treatment assignments $T$ and the outcomes $Y$ partially based on the images’ representation. In our adaption, we also follow this scheme, but we only used two randomly selected digits $c_{i},c_{j}\in\{0,..,9\},c_{i}\neq c_{j}$. Therefore, in our target-domain, we have the images of the digits $c_{i}$ and $c_{j}$. We map the high-dimensional images to a one-dimensional array $\phi$ (See Appendix for more details), which is used to simulate the treatments $T$ and the outcome $Y$. For the source-domain, we use the images of the other digits. As the source-domain is unlabeled, there is no treatment $T$ or outcome $Y$ produced, and we use the images as they are observed. We follow Setting 2.

Setting 1 (Figure 3A-D): We explored the following ratios $r\in\{0.25,0.67,1.5,4\}$, with low $r$ being the ideal use-case of Causal-Batle. For the IHDP dataset, we adopted $b_{d}=10$ datasets replications with $b_{m}=30$ model repetitions, for a total of $B=10\times 30=300$ estimated values for each $r$. For the GWAS dataset, we adopted $b_{d}=10$ dataset replications with $b_{m}=10$ model repetitions for each $r$ ($B=100$). The IHDP dataset is not a small high-dimensional dataset; however, it is a classic benchmark dataset. With $r\leq 1$, Causal-Batle outperforms all the other baselines as shown in Figure 3.A. With $r>1$, the performance of all methods improves, including Causal-Batle, which stays in the top 2. Figure 3.B shows the $b_{m}=30$ model repetitions on a dataset replication, where the colors indicate the $r$ value (lighter colors $\rightarrow$ low $r$). Besides having the lowest MAE’s, Causal-Batle estimations are well distributed around the true value, while some baselines often underestimate the true values, especially for low $r$ values (light-colored dots).

The GWAS dataset fits more in our ideal use-case, containing 2k samples and 10k covariates. As a reminder, $p_{t}=0.2(20\%)$ means that all methods will receive $2k\times 0.2=400$ samples as $X_{target}$, which contains the treatment assignment and the outcome of interest; Causal-Batle, on top of the $X_{target}$, also receives $X_{source}$, which contains $2k\times(1-0.2)=1.6k$ unlabeled samples (only covariates available, without treatment assignment nor outcome). Figure 3.C shows that Causal-Batle has the lowest MAE among all methods considered, with the distance between Causal-Batle and the other methods decreasing as $r$ increases. Figure 3.D shows that Causal-Batle estimated values are well distributed around the true value. Note that the Bayesian Dragonnet and CEVAE are also well distributed around the true value; however, they underestimate $\hat{\tau}$ with low $r$ values.

Setting 2 (Figure 3E-F): The HCMNIST had $b_{d}=4$ dataset replications and $b_{m}=25$ model repetitions, for a total of $4\times 25=100$ estimated values for each $r\in\{0.25,0.5,0.75,1\}$. As Figure 3.E shows, Causal-Batle has the smallest MAE for $r\leq 0.5$, and the second-best MAE for $r>0.5$. According to the experimental results in Figure 3.F, both Causal-Batle and the Bayesian Dragonnet tend to underestimate the treatment effect with low $r$ values (light colors). However, Causal-Batle has lower variance, and its estimates are more accurate than Bayesian Dragonnet’s estimates.

Discussion: Causal-Batle has more accurate estimates (low MAE) than the baselines for low $r$ values. Note that the only difference between Causal-Batle and the Bayesian Dragonnet is the transfer learning components (discriminator loss, adversarial loss, reconstruction loss, and the use of the $X_{s}$ dataset). An explanation for these results is that transfer learning improves the estimates. Considering the IHDP dataset, when $r>1$, the simpler architecture of the Bayesian Dragonnet performs better, indicating a possible threshold for the IHDP dataset when a simpler architecture is preferable over a more complex one with transfer learning. For the GWAS and HCMNIST datasets, that threshold is above $r>4$ and $r>0.5$. Note that we did not experiment with CEVAE with the HCMNIST as that would require many changes to the CEVAE to work with image data, which was not our focus. Finally, we investigated Dragonnet’s poor performance. Our main hypothesis is that the Bayesian component included in the Bayesian Dragonnet and Causal-Batle reduces the variability and improves $\tau$ estimates. The Bayesian component predicts the outcome and the treatment assignment as a Gaussian Mixture and Bernoulli distribution, respectively, and predictions are calculated using MC-dropout (30 forward passes).