Variational Causal Inference
Abstract
Estimating an individual’s potential outcomes under counterfactual treatments is a challenging task for traditional causal inference and supervised learning approaches when the outcome is high-dimensional (e.g. gene expressions, impulse responses, human faces) and covariates are relatively limited.
In this case, to construct one’s outcome under a counterfactual treatment, it is crucial to leverage individual information contained in its observed factual outcome on top of the covariates.
We propose a deep variational Bayesian framework that rigorously integrates two main sources of information for outcome construction under a counterfactual treatment: one source is the individual features embedded in the high-dimensional factual outcome; the other source is the response distribution of similar subjects (subjects with the same covariates) that factually received this treatment of interest.
Introduction
In traditional causal inference, individualized treatment effect (ITE) is typically estimated by fitting supervised learning models on the covariate-specific efficacy p
( Y | X, T ) or E [ Y | X, T ] with observed outcome Y , covariates X and treatment T .
However, in cases such as the single-cell genetic perturbation datasets [2; 8; 12] where Y has thousands of dimensions while X has less than a dozen features and mostly categorical, such model could hardly be relied on to produce useful individualized results.
To construct outcome Y ′ under a counterfactual treatment T ′ , it is important and necessary to learn the individual-specific efficacy p
( Y ′ | Y, X, T, T ′ ) , which leverages the rich information embedded in the factual outcome Y that cannot be recovered by the handful of covariates.
For example, given a cell with type A549 ( X ) that received SAHA drug treatment ( T ), we may want to know what its gene expression profile ( Y ′ ) would look like if he had received Dex drug treatment ( T ′ ) instead.
In this case, we would want to take the profiles of other cells with type A549 that have indeed received Dex drug treatment as reference, but would also want to extract as much individual features as possible from this cell’s own expression profile ( Y ) to combine with the reference so that the reconstructed counterfactual expression profile could preserve its own identity.
Yet the lack of observability of Y ′ presents a major difficulty for supervised learning when Y is taken as an input.
In previous works that involve self-supervised counterfactual generators [7; 18; 6], there is no explicit regulation on the trade-off between covariate-specific efficacy (reference) and individual features unaccounted for by the covariates (own identity).
In Louizos et al.
[7] and Yoon et al.
[18], there is no such regulation; the trade-off is black-boxed and has to depend on network sizing.
In Lotfollahi et al.
[6], the trade-off is implicitly encouraged by the balancing between the adversarial reward and reconstruction loss of the generator.
In this work, we present a semi-autoencoding scheme and its variational Bayes [4; 11] counterpart to address this problem.
Our framework is able to encode latents that explicitly honor covariate-specific efficacy while preserving individuality.
[Figure]
Proposed Method 2.1 Semi-autoencoding Potential Outcomes
We define outcome Y : Ω → R n to be a n -dimensional random vector, X : Ω → E X to be a mix of m categorical and real-valued covariates, T : Ω → E T to be a categorical or real-valued treatment and Z : Ω → R d to be the latent feature vector on a probability space (Ω , Σ , P ) .
Suppose the causal relations between random variables (and random vectors) follow a structural causal model (SCM) [10] depicted by the causal diagram in Figure 1, where we formulate counterfactuals Y ′ and T ′ as separate variables apart from Y and T , but having a conditional outcome distribution p ( Y ′ | Z, T ′ = a ) identical to that of its factual counterpart p
( Y | Z, T = a ) on a treatment level a .
Y ′
and T ′ are parts of the full data D =
( X, Z, T, T ′ , Y, Y ′ ) , yet are never observed.
Latent Z is designed to be a feature vector of Y such that random factor U Y in structural equation Y =
f ( Z, T, U Y ) is significantly less noisy than that of the outcome equation Y =
f ( X, T, U Y ) in a traditional SCM of triplets ( Y, X, T ) .
Under this setting, Z has a posterior distribution p ( Z | Y, X, T ) given the observations; Y and Y ′ can be reconstructed by p ( Y | Z, T ) and p ( Y ′ | Z, T ′ ) respectively if Z is recovered.
Similar to prior works on autoencoder [17; 1], we estimate the latent recovery model and outcome reconstruction model with deep neural networks q ϕ and p θ , whose dimensionalities control the noise level of U Y .
Figureshows the graphical model for the encoder and decoder.
While the reconstruction of Y is self-supervised, we can only assess the construction of Y ′ by looking at its resemblance to similar individuals that indeed received treatment T ′ .
Hence we may develop a semi-autoencoding scheme that works as follows.
During training, we estimate Z given the observed triplets ( Y, X, T ) , then predict both factual and counterfactual outcomes ˜ Y θ,ϕ and ˜ Y ′ θ,ϕ .
Inputting T into decoder p θ yields ˜
Y θ,ϕ , which is then evaluated by the reconstruction loss
L ( ˜ Y θ,ϕ , Y )
; inputting
T ′ into p θ yields ˜ Y ′ θ,ϕ , which we evaluate by the negative likelihood loss −L p
( Y ′ | X,T ′ )
( ˜ Y ′ θ,ϕ ) with respect to the counterfactual outcome distribution p
( Y ′ | X, T ′ ) .
The intuition is that, if ˜ Y ′ θ,ϕ is indeed one’s outcome under T ′ , then the likelihood of ˜ Y ′ θ,ϕ coming from the outcome distribution of individuals with the same attributes that factually received treatment T ′ is supposed to be high.
The total loss of the semi-autoencoder can then be formed as a weighted combination:
[Equation]
where ω is a scaling coefficient and ˆ p is the traditional covariate-specific outcome model fit on the observed variables ( X, T, Y ) (notice that p ( Y ′ | X, T ′ = a )
= p ( Y | X, T = a ) for any a ).
If ˆ p is difficult to fit, we can train a discriminator D ( X, T, Y ) for p ( Y | X, T ) (and p ( Y ′ | X, T ′ ) ) instead with the adversarial approach [3] and use its discrimination loss over counterfactuals L D ( X, T ′ , ˜ Y ′ θ,ϕ ) in place of −L ˆ p ( Y ′ | X,T ′ )
( ˜ Y ′ θ,ϕ ) .
In cases where covariates are limited and discrete such as the single-cell perturbation datasets, ˆ p can simply be the empirical outcome distribution under T stratified by X .
[Figure]
2.2 Variational Causal Inference
Now we present our main result by rigorously formulating the objective, providing a probabilistic theoretical backing, and specifying an optimization scheme that reflects the intuition described in the previous section.
Suppose we want to optimize p ( Y ′ | Y, X, T, T ′ ) in addition to the traditional outcome model p ( Y | X, T ) .
The following theorem states the evidence lower bound (ELBO) and thus provides a roadmap for stochastic optimization:
Theorem 1.
Suppose random vector D =
( X, Z, T, T ′ , Y, Y ′ ) follows a causal structure defined by the Bayesian network in Figure 1.
Then log [ p ( Y ′ | Y, X, T, T ′ )]
+ log [ p ( Y | X, T )] has the following evidence lower bound:
[Equation]
Proof of the theorem can be found in Appendix A.1.
Theorem 1 states that, in order to maximize log [ p ( Y ′ | Y, X, T, T ′ )]
+ log [ p ( Y | X, T )] , we can instead estimate and maximize the ELBO which consists of the expected reconstruction likelihood E p ( Z | Y,X,T ) log [ p ( Y | Z, T )] and covariate-specific outcome likelihood log [ p ( Y ′ |
X, T ′ )] , echoing the intuition highlighted by Equation 1, with an additional divergence term − KL [ p ( Z | Y, X, T )
∥ p ( Z | Y ′ , X, T ′ )] that regularizes the similarity across latent distributions under different potential scenarios.
The objective after weighting and parameterizing with deep neural networks is given as
[Equation]
where ω 1 and ω 2 are scaling coefficients;
˜ Y ′ θ,ϕ ∼ E q ϕ
( Z | Y,X,T ) p
θ ( Y ′ | Z, T ′ )
and p ( ˜ Y ′ θ,ϕ | X, T ′ ) can be estimated by ˆ p ( ˜ Y ′ θ,ϕ | X, T ′ ) with the approaches described at the end of Section 2.1.
Note that this framework is fundamentally different from variational autoencoder (VAE) based frameworks such as CVAE [13] and CEVAE [7] since the variational lower bound is derived directly from the causal objective and the KL-divergence term serves a causal purpose that is completely different from the latent bounding purpose of VAEs.
An explanation of this divergence term is given below.
Divergence Interpretation
The divergence term − KL (cid:104)
q ϕ
( Z | Y, X, T ) ∥
q ϕ ( Z | ˜ Y
′ θ,ϕ , X, T ′ )
(cid:105) in
Equation 3 encourages the preservation of individuality in counterfactual outcome constructions.
In the autoencoding process, the latent variable Z is required to preserve individual information contained in Y beyond covariates X in order to obtain a quality reconstruction of Y .
Yet the counterfactual construction p θ ( Y ′ | ˜
Z ϕ , T ′ ) could discard this individuality in Z once it realizes a counterfactual treatment T ′ is in the inputs, if the outcome construction was only penalized by the covariate-specific likelihood loss − log [ˆ p ( Y ′
| X, T ′ )] .
Such behavior is regulated with the addition of the divergence term, since otherwise we would not be able to recover a latent distribution q ϕ ( Z | ˜ Y ′ θ,ϕ , X, T ′ ) close to q ϕ
( Z | Y, X, T ) .
2.3 Marginal Effect Estimation
Upon obtaining individual level results, we introduce a robust estimator for the marginal treatment effect Ψ( p ) =
E p ( Y ′ | T ′ = a ) of a treatment level a in the variational causal inference (VCI) framework.
By Van der Vaart [16], a sequence of estimators ˆΨ n is asymptotically efficient if √ n (ˆΨ n − Ψ
( p ))
{+= 1 / √ n (cid:80) nk =
+}
˜ ψ p ( D k ) + o p (1) where ˜ ψ p is the efficient influence function of Ψ( p ) and D k ∼ p ( D ) .
The theorem below gives this efficient influence function and thus provides a construction of an asymptotically efficient regular estimator for Ψ :
Theorem 2.
Suppose D : Ω → E follows a causal structure defined by the Bayesian network in
Figure 1, where the counterfactual conditional distribution p ( Y ′ , T ′ | Z, X ) is identical to that of its factual counterpart p
( Y, T | Z, X ) .
Then Ψ( p ) has the following efficient influence function:
[Equation]
Proof of the theorem can be found in Appendix A.2.
In practice, we only need p ( Y ′ |
Z, T ′ =
a ) to be identical to its factual counterpart p ( Y | Z, T = a ) given a treatment level a such that they can be estimated with the same decoder p θ .
Since Y ′ and T ′ are both unobserved in the data, it does not bother to simply acknowledge that T ′ is sampled according to p ( T | X ) .
By Theorem 2, the following estimator is asymptotically efficient under some regularity conditions [15]:
[Equation]
where ( Y k , X k , T k ) are the observed variables of the k -th individual and
˜ Z k,ϕ ∼ q ϕ
( Y k , X k , T k ) ; ˆ p is an estimation of the propensity score.
As can be seen, ˆΨ n ( θ, ϕ ) takes a very similar form as the augmented inverse propensity weighted (AIPW) estimator of the average treatment effect on treated (ATT) in traditional causal inference.
Covariate-specific Marginal Effect
The robust estimator within the VCI framework can also be employed in estimating covariate-specific marginal treatment effect Ψ( p ) =
E p ( Y ′ |
X
=
c, T ′ =
a ) for a given covariate c of interest, by applying Equation 5 on the set of observations with X = c .
Note that the regression adjustment term E p θ ( Y ′ | ˜ Z k,ϕ , T ′ k
= a ) by our deep network design can vary across different subjects within the same covariate group, which is not true for that of the AIPW estimator in traditional causal inference.
Experiments
We tested our framework on single-cell perturbation datasets along with four benchmark models —
CEVAE [7], GANITE [18], CPA [6] and a naive autoencoding method.
In single-cell perturbation datasets, each observation is an n -dimensional vector (the outcome) storing the expression counts of n selected genes in a single cell, along with the perturbation (the treatment) this cell was subjected to and the cell covariates.
We employed the publicly available sci-Plex dataset from Srivatsan et al. [14] (Sciplex) and the CRISPRa dataset from Schmidt et al.
[12] (Marson) in our experiments.
Two thousand most variable genes were selected from each dataset.
Data with certain treatment-covariate combinations are held out as the out-of-distribution (OOD) set and the rest are split into training and validation set (four-to-one).
Details of the hyperparameter settings can be found in Appendix C.
Out-of-Distribution Selections We randomly select a covariate category (e.g. a cell type) and hold out all cells in this category that received one of the twenty perturbations whose effects are the hardest to predict.
We use these held-out data to compose the out-of-distribution (OOD) set.
We computed the Euclidean distance between the pseudobulked gene expression of each perturbation against the rest of the dataset, and selected the twenty most distant ones as the hardest-to-predict perturbations.
This is the same procedure carried out by Lotfollahi et al.
Differentially-Expressed Genes
In order to evaluate the quality of the predictions on the genes that were substantially affected by the perturbations, we select sets of 50 differentially-expressed (DE) genes associated with each perturbation and separately evaluate model performance on these genes.
This is the same procedure carried out by Lotfollahi et al.
3.1 Out-of-Distribution Predictions
Same as Lotfollahi et al.
[6], for each perturbation of each covariate level (e.g. each cell type of each donor) in the OOD set, we compute the R 2 (coefficient of determination) of the average outcome predictions for all genes and DE genes using samples from the validation set against the true empirical average with samples from the OOD set.
The average R 2 over all perturbations of all covariate levels is then calculated as the evaluation metric.
Table 1 shows the best results of each model over five independent runs.
All models are trained on Amazon’s accelerated computing EC2 instance G4dn (12xlarge) which contains 4 NVIDIA T4 Tensor Core GPUs.
[Table]
In the experiments, we used the empirical outcome distribution stratified by X and T (both discrete) to estimate the covariate-specific outcome model p ( Y | X, T ) .
There are a few gradient detaching options for ˜ Y ′ θ,ϕ when inputting ˜ Y ′ θ,ϕ into the KL-divergence term KL (cid:104)
q ϕ
( Z | Y, X, T ) ∥
q ϕ ( Z | ˜ Y
′ θ,ϕ , X, T ′ )
(cid:105) : not detached; detached from encoder; detached from both encoder and decoder.
Here we chose to not detach ˜ Y ′ θ,ϕ .
Our variational Bayesian causal inference framework beat all state-of-the-art models in these experiments, with the largest fractional improvement on DE genes which are most causally affected by the perturbations.
3.2 Marginal Estimations
In this section [Table] we use the same evaluation metric, but compute each R 2 with the robust marginal estimator and compare the results to that of the empirical mean estimator.
Note that on OOD set, the robust estimator reduces to empirical mean since no perturbation-covariates combination exist in validation set.
Therefore, we compute the R 2 of the marginal estimators using samples from the training set against the true empirical average on the validation set in these experiments.
In each run, we train a VCI model for individualized outcome predictions [Table] and calculate the evaluation metric for each marginal estimator periodically during the course of training.
Table 2 shows the results on
Marson over five independent runs.
Note that the goal of robust estimation is to produce less biased estimators with tighter confidence bounds, hence here we reported the mean and standard deviation across runs instead of the maximum to reflect its effectiveness regarding this goal
[Table]
[Table] As is shown in the table [Table] the robust estimator provides a crucial adjustment to the empirical mean of model predictions especially on the hard-to-predict elements of high-dimensional vectors.
Such estimation could be valuable in many contexts involving high-dimensional predictions where deep learning models might plateau at rather low ceilings.
Conclusion
In this work we introduced a variational Bayesian causal inference framework for high-dimensional individualized treatment effect prediction.
With this framework, covariate-specific efficacy and individual identity can be explicitly balanced and optimized.
In experiments, our framework excelled the state-of-the-art algorithms at out-of-distribution predictions on single-cell perturbation datasets — a widely accepted and biologically meaningful task.
We note that although several approaches were described for the estimation of covariate-specific distribution loss, our most compelling method is for discrete treatments, where the distribution can be estimated empirically via stratification.
Acknowledgments
We thank Ivana Malenica, Dominik Janzing, David E. Heckerman for the insightful discussions and
Carlo De Donno for the guidance on data processing.
Appendix A Proof of Theorems A.1 Proof of Theorem
Proof.
By the d-separation [9] of paths on the causal graph defined in Figure we have log [ p ( Y ′ |
[Equation]
X
[Equation]
[Equation]
[Equation]
[Equation] [Equation]
Reorganizing the terms yields
[Equation]
A.2 Proof of Theorem
Proof.
Following Van der Vaart [16], we define a path p
ϵ ( D ) =
p ( D )(1 + ϵS ( D )) on density p of D as a submodel that passes through p at ϵ = 0 in the direction of the score S ( D ) =
d dϵ log [ p ϵ ( D )] (cid:12)(cid:12) (cid:12) ϵ =0 .
Following the key identity presented in Levy [5]:
[Equation]
[Table]
[Equation]
where pa ( D i ) denotes the parent nodes of variable D i and minuscule of a variable denotes the value it takes, we have
[Equation]
(cid:27) p ( d ) (24) by assumptions of Theorem 2 and factorization according to Figure 1.
Hence d dϵ Ψ( p ϵ ) (cid:12)(cid:12)(cid:12) ϵ
=0 = (cid:28)
S ( D ) , I ( T = a ) p ( T | X )
( Y − E p
[ Y | Z, T ]) + E p
[ Y ′ | Z, T ′ =
a ] − Ψ (cid:29)
L 2 (Ω; E ) (25) and we have ˜ ψ p
= I ( T = a ) /p
( T | X ) · ( Y − E p
[ Y | Z, T ]) + E p
[ Y ′ | Z, T ′ =
a ] − Ψ .
B Benchmark Adaptations B.1 Autoencoder
The adapted autoencoder reconstructs the outcome during training similar to a generic autoencoder, but takes treatment and covariates as additional inputs.
During test time, we simply plug in the counterfactual treatments along with factual outcomes and covariates to generate the counterfactual outcome predictions.
B.2 GANITE
GANITE [18]’s counterfacutal generator does not scale with a combination of high-dimensional outcome and multi-level treatment, thus here we only input one randomly sampled counterfactual treatment to the generator and correspondingly construct one counterfactual outcome for each sample.
See Figure 3 for the original and adapted structure of the model.
[Figure]
The discriminator predicts the logits l f , l cf of y f , ˜ y cf separately.
The cross entropy loss of ( l f , l cf ) against (1 , 0) is then calculated.
C Hyperparameter Settings
All common hyperparameters of all models are set to the same as the defaults of CPA: an universal number of hidden dimensions 128 ; number of layers 6 (encoder 3 , decoder 3 ); discriminator (if any) hidden dimensions 64 and number of layers 2 ; adversarial training (if any) period 3; an universal learning rate 3 − 4 , decay rate 4 − decay period 45 .
For more details regarding the hyperparameter settings of our framework, visit our code repository https://github.com/yulun-rayn/variational-causal-inference .