Deep Multi-Modal Structural Equations For Causal Effect Estimation With Unstructured Proxies
Abstract
Estimating the effect of an intervention from observational data while accounting for confounding variables is a key task in causal inference.
Oftentimes, the confounders are unobserved, but we have access to large amounts of additional unstructured data (images, text) that contain valuable proxy signal about the missing confounders.
This paper argues that leveraging this unstructured data can greatly improve the accuracy of causal effect estimation.
Specifically, we introduce deep multi-modal structural equations, a generative model for causal effect estimation in which confounders are latent variables and unstructured data are proxy variables.
This model supports multiple multi-modal proxies (images, text) as well as missing data.
We empirically demonstrate that our approach outperforms existing methods based on propensity scores and corrects for confounding using unstructured inputs on tasks in genomics and healthcare. Our methods can potentially support the use of large amounts of data that were previously not used in causal inference.
Introduction
An important goal of causal inference is to understand from observational data the causal effect of performing an intervention—e.g., the effect of a behavioral choice on an individual’s health [36].
As an initial motivating example for this work, consider the problem of determining the effect of smoking on an individual’s risk of heart disease.
This problem is complicated by the presence of confounders: e.g., it is possible that individuals who smoke have a higher likelihood to be sedentary, which is a lifestyle choice that also negatively impacts their heart disease risk.
If the individual’s lifestyle is available to us as a well-defined feature, we may adjust for this factor while computing treatment effects [36].
However, confounders are often not observed and not available in the form of features, making accurate causal inference challenging.
Oftentimes, datasets in domains such as medicine or genomics come with large amounts of unstructured data —e.g., medical images, clinical notes, wearable sensor measurements [7].
This data contains strong proxy signal about unobserved confounding factors—e.g., wearable sensor measurements can help reveal individuals who are sedentary.
However, existing causal inference methods are often not able to leverage this “dark data” for causal effect estimation [18].
The goal of this paper is to develop methods in causal inference that improve the estimation of causal effects in the presence of unobserved confounders by leveraging additional sources of unstructured multi-modal data, such as images and text.
For example, given time series from patients’ wearables, our methods may disentangle the effects of being sedentary from the effects of smoking by using clusters in the sensor measurements (which would correspond to groups of active and sedentary individuals) as a proxy for the patients’ lifestyles and without requiring explicit lifestyle features.
Concretely, our paper formalizes the task of estimating causal effects using rich, unstructured, multimodal proxy variables (e.g. images, text, time series) and introduces deep multi-modal structural equations, a generative model in which confounders are latent variables.
This model can perform causal effect estimation with missing data by leveraging approximate variational inference and learning algorithms [6] developed for this task. Previous methods relied on propensity scoring with neural approximators [41, 50, 51], which may output volatile probabilistic outputs that lead to unreliable effect estimates [18]. Our methods are generative, and thus naturally sidestep some of the aforementioned instabilities, support high-dimensional treatment variables and incomplete data, and as a result are applicable to a broader class of causal inference problems.
We evaluate our methods on an important real-world causal inference task—estimating the effects of genetic mutations in genome-wide association studies (GWASs)—as well as on benchmarks derived from popular causal inference datasets.
The intervention variables in a GWAS are high-dimensional genomic sequences, hence existing methods based on propensity scoring are not easily applicable. In contrast, we demonstrate that our algorithms naturally leverage high-dimensional genetic and environmental data (e.g., historical weather time series) and can discover causal genetic factors in plants and humans more accurately than existing GWAS analysis methods.
Contributions In summary, this paper makes three contributions: (1) we define the task of estimating causal effects using rich, unstructured multi-modal proxy variables; (2) we introduce deep multi-modal structural equations, a generative model tailored to this problem, and we describe associated variational learning and inference algorithms; (3) we demonstrate on an important real-world problem (GWAS) that unstructured data can improve causal effect estimation, enabling the use of large amounts of “dark data” that were previously not used in causal inference.
Background
Notation Formally, we are given an observational dataset D = { ( x ( i ) , y ( i ) , t ( i ) ) }
ni =1 consisting of n individuals, each characterized by features x ( i )
∈ X ⊆ R d , a binary treatment t ( i ) ∈ { 0 , 1 } , and a scalar outcome y ( i ) ∈ R .
We initially assume binary treatments and scalar outcomes, and later discuss how our approach naturally extends beyond this setting.
We also use z ( i ) ∈
R p to model latent confounding factors that influence both the treatment and the outcome [28].
We are interested in recovering the true effect of T = t in terms of its conditional average treatment effect (CATE), also known as the individual treatment effect (ITE) and average treatment effect (ATE).
[Equation]
where do ( · ) denotes an intervention [37].
Many methods for this task rely on propensity scoring [41, 50, 51], which uses a model p ( t | x ) to assign weights to individual datapoints; however, when x is high-dimensional and unstructured, a neural approximator for p ( t | x ) may output volatile and miscalibrated probabilities close to { 0 , 1 } that lead to unreliable effect estimates [18].
Structural Equations An alternative approach are structural equation models of the form
[Equation]
where Z ∼ p
( Z ) is drawn from a prior and the ε i are noise variables drawn independently from their distributions [10].
Structural equations define a generative model p ( x, y, z, t ) of the data.
When this model encodes the true dependency structure of the data distribution, we can estimate the true effect of an intervention by clamping t to its desired value and drawing samples.
Deep Structural Equations Equations 2 can be parameterized with deep neural networks, which yields deep structural equation models [49, 28].
Expressive neural networks may learn a more accurate model of the true data distribution on large datasets, which improves causal effect estimation.
Such models have been used for GWAS analysis [49] and to correct for proxy variables [28].
Causal Effect Estimation With Unstructured Proxy Variables
Oftentimes, datasets in domains such as medicine or genomics come with large amounts of unstructured data (medical images, clinical notes), which contains strong proxy signal about unobserved confounding factors.
Our paper seeks to develop methods that leverage unstructured data within causal inference.
We start by formalizing this task as causal effect estimation with unstructured proxy variables; these proxies may come from multiple diverse modalities (images, text).
3.1 Task Definition
Formally, consider a causal inference dataset D = { ( x ( i ) , y ( i ) , t ( i ) ) } ni =1 in which x ( i ) =
( x ( i ) 1 , x ( i ) 2 , . . .
, x ( i ) m ) is a vector of m distinct input modalities x
( i )
j ∈ X j
(e.g., images, text, time series, etc.).
In other words, X = X 1 × . . .
×
X m , where each X j corresponds to a space of images, time series, or other unstructured modalities.
Here, t ( i ) ∈ T (binary or continuous) is the treatment and y ( i
) ∈
Y is the output.
Some modalities may also be missing at training or inference time.
We are interested in recovering the true effect of t in terms of the individual and average treatment effects.
We are specifically interested in estimating the individual treatment effect (ITE) from arbitrary subsets of modalities M ⊆ { 1 , 2 , ..., m } , indicating that certain inputs may be missing at test time.
Y [ x, t, M ] =
E [ Y | do ( T = t ) , X j = x j for j in M ] ITE ( x, M )
=
Y [ x, t = 1 , M ] − Y [ x, t = 0 , M ] (3)
To help make this setup more concrete, we define two motivating applications.
Healthcare Consider the task of determining the effect of smoking on heart disease from an observational dataset of patients.
The observational study may contain additional unstructured data about individuals, e.g., clinician notes, medical images, wearable sensor data, etc.
This data may hold information about hidden confounders: for example, raw wearable sensor data can be clustered to uncover sedentary and active indivudals, revealing a latent confounding factor, sedentary lifestyle.
Genomics Consider the problem of estimating the effects of genetic variants via a genome-wide association study (GWAS).
Modern GWAS datasets in plants or humans feature large amounts of unstructured inputs [7]: clinical notes, medical records, meteorological time series.
For example, historical weather data (e.g., precipitation, wind strength, etc.) can reveal distinct climatic regions that affect plant phenotypes and whose confounding effects should be corrected for in a GWAS [55].
Deep Structural Equations for Causal Effect Estimation
Next, we derive models and inference algorithms for the task of causal effect estimation with unstructured proxy variables.
Our approach uses deep structural equations to extract confounding signal from the multi-modal proxies x ( i ) j .
We use neural networks because they naturally handle unstructured modalities via specialized architectures (e.g., convolutions for images) that can learn high-level representations over raw unstructured inputs (e.g., pixels).
Parameterizing structural equations with neural networks also presents challenges: they induce complex latent variable models that require the development of efficient approximate inference algorithms [10, 6].
We present an instantiation of Equations 2 that admits such efficient algorithms.
4.1 Deep Multi-Modal Structural Equations
We start by introducing deep multi-modal structural equations (DMSEs), a generative model for estimating causal effects in which confounders are latent variables and unstructured data are proxy variables.
We define a DMSE model as follows:
[Equation]
where p x j , p y are probability distributions with a tractable density over x j and y , respectively, and the θ x j , θ y are the parameters of p x j , p y . The θ x j , θ y are themselves functions of z, t parameterized by neural networks—e.g., when x j is Gaussian, the θ x j ( z ) are a mean and a covariance matrix µ x j ( z ) , Σ x j ( z ) that are parameterized by a neural network as a function of z (see [21, 28]).
Note that other modeling choices (e.g., Bernoulli distributions for discrete variables) are also possible.
Note that the models for θ x j ( z ) can benefit from domain-specific neural architectures—e.g., a convolutional parameterization for µ x j ( z ) , σ x j ( z ) as a function of z is more appropriate when the x j are images.
See Appendix B for details on recommended architectures. While Section 3.1 defines y, t as scalars following existing literature [28, 54, 50], DMSEs can also define a model with high-dimensional y, t —we simply choose the distributions over y, t to be multivariate. Our inference and learning algorithms will remain unchanged, except for the parameterization of specific approximate posteriors (e.g., q ( z | y, t ) ). In fact, we apply DMSEs to high-dimensional t in our GWAS experiments in Section 5.
Note that this is a setting where existing propensity scoring methods (which learn a model of p ( t | x ) ) are not directly applicable [41, 50]—our approach, on the other hand, can easily be used with high-dimensional t on tasks like GWAS analysis. Dependency Structure Equations (4) define a density p ( z ) p ( t | z ) p ( y | z, t ) (cid:81)
mj =1 p ( x j | z ) . Note that each proxy x j is independent of the others conditioned on z . Figure 1 shows these dependencies as solid lines. In our setting, the x j represent image pixels, waveform measurements, etc; thus, they should not directly influence each other or y, t .
For example, it would not make sense for the pixels of an image to causally influence the samples of a waveform measurement—they influence each other only through a latent confounder z
(e.g., patient health status), and the x j are therefore assumed to be conditionally independent given z .
These independence assumptions will also enable us to derive efficient stochastic variational inference algorithms, as we show below.
[Figure]
In certain settings, we may also want to model observed confounders as well as additional structured proxies x j that have direct causal effects on each other and on y, t .
Figure 1 shows these additional dependencies as dashed lines.
Our model admits simple generalizations that support these modeling assumptions.
In brief, our variational inference algorithms can be conditioned on the observed variables, and mutually dependent sets of proxies can be treated as a group represented by one high-dimensional variable.
See Appendix H for details.
4.2 Approximate Inference and Learning Algorithms
In the full multi-modal setting, the x j are conditionally independent given z (Figure 1), which enables us to apply efficient algorithms inspired by Wu & Goodman [56].
These algorithms offer the following improvements: (1) we may perform learning, inference, and causal estimation with missing modalities x j ; (2) the process for performing causal inference does not require training auxiliary inference networks as in previous work [28].
The DMSE model induces a tractable joint density p ( x, y, t, z ) , which allows us to fit its parameters using stochastic variational inference by optimizing the evidence lower bound (ELBO):
[Equation]
where p ( y ( i ) , t ( i ) , z )
=
p ( y ( i ) | t ( i ) , z ) p ( t ( i )
| z ) p ( z ) and q ( z | x ( i ) , y ( i ) , t ( i ) ) is the approximate variational posterior.
We assume a total of m modalities.
Structured Multi-Modal Variational Inference
We may use the independence structure of p (Figure 1) to derive an efficient structured form for q .
First, observe that because the true posterior factorizes as p ( z | x, t, y ) ∝ ( p ( z | t, y ) (cid:81) mj
=1 p ( z | x j )) / (cid:81) m − 1 j =1 p ( z ) , the optimal approximate posterior q must also factorize as q ( z | x, t, y )
∝ ( q ( z | t, y ) (cid:81) mj
=1 q ( z | x j )) / (cid:81) m − 1 j =1 p
( z ) .
This decomposition implies that we can maintain the optimal structure of q by training modality-specific inference networks ˜ q ( z | t, y ) and ˜ q ( z | x j ) such that q ( z | x j ) =
˜ q
( z | x j ) p ( z ) and q ( z | t, y ) =
˜ q
( z | t, y ) p ( z ) and by defining a joint posterior as
[Equation]
This network can be seen as a product of experts (PoE) [56].
Computing the density q is in general not possible.
However, because p ( z ) , q ( z | t, y ) and q ( z | x j ) are Gaussians, we may use the fact that a product of Gaussians with means µ i and covariances V i is µ = ( (cid:80) µ i T i ) /
( (cid:80) T i ) and V =
( (cid:80) i T i ) − 1 , where T i =
1 /V
i .
Thus, computing q ( z | x, t, y ) for any subset of modalities is possible without having to train an inference network for each subset of modalities separately.
A Multi-Modal ELBO Training this architecture with ELBO X will not necessarily yield good single-modality inference networks ˜ q ( z | x j ) , while training with each modality separately will prevent the network from learning how the modalities are related to each other.
Hence, we train our model with a sub-sampled ELBO objective [56] that is computed on the full set of modalities, each individual modality, and a few subsets of modalities together.
For this, we randomly pick at each gradient step s non-empty subsets { M k } sk =1 of the set of modalities M k ⊆ { 1 , 2 , ..., m } .
The final objective is
[Equation]
, where
[Equation]
Our model doesn’t train extra auxiliary inference networks, unlike that of Louizos et al.
4.3 Deep Gaussian Structural Equations
The DMSE model can be simplified in settings in which there is only one type of proxy x (i.e., m = 1 ).
This simplified model, which we call deep Gaussian structural equations (DGSEs), has a tractable joint density p ( x, y, t, z ) = p ( z ) p ( x | z ) p ( t | z ) p ( y | z, t ) , where the latent z is Gaussian.
The DGSE model can also be fit using stochastic variational inference by optimizing the ELBO objective (cid:80)
[Equation] ( x ( i ) , y ( i ) , t ( i ) , z ) − log q ( z | x ( i ) , y ( i ) , t ( i ) ) (cid:3) , where q ( z | x, y, t ) is an approximate variational posterior.
We optimize the above objective using gradient descent, applying the reparameterization trick to estimate the gradient.
We compute the counterfactual Y [ x, t ] using auxiliary inference networks as in earlier work [28].
See Appendix B for the full derivation.
4.4 Properties of Deep Multi-Modal Structural Equations
Recovering Causal Effects The DMSE and DGSE models determine the true causal effect when their causal graph is correct and they recover the true data distribution.
The following argument is analogous to that made in most previous works on causal deep generative models [28, 57, 35].
Theorem 1.
The DMSE and DGSE models recover the true ITE ( x, M ) for any subset M ⊆ { 1 , 2 , ..., m } of observed modalities whenever they represent the true data distribution p ( x, y, t, z ) .
Proof: We establish the theorem for DMSEs; the proof for DGSEs is analogous with m = 1 .
Let x M = { x j | j ∈ M } be the data from the observed subset of modalities.
We need to show that p ( y | x M , do ( t = t ′ )) is identifiable for any t ′ .
Observe that
[Equation]
where the second equality follows from the rule of do-calculus (applying backdoor adjustment).
Since our proof holds for any t ′ and all elements on the right-hand side are identifiable, the claim follows.
Note that in practice our assumption may not hold (e.g., neural network optimization is non-convex and may fail), but there is evidence of both failure modes [44] as well as successful settings in which deep latent variable models provide useful causal estimates [57, 35, 30, 48].
See our Discussion section for additional details.
Identifiability in Linear Models Structural equations parameterized by non-convex neural networks are less amenable to analysis than simpler model classes.
However, we may provide theoretical guarantees in the special case where a DMSE model (Equations 4) is linear, i.e., each equation with input variables u ∈ R d 1 has the form A · u + b for some A ∈ R d 2 × d 1 , b ∈ R d 1 .
Specifically, we establish in Appendix I the following result.
Theorem 2. Given a binary treatment t , a univariate outcome y , confounder z and proxy variables u, v, w , the causal effect P ( y | do ( t )) is identifiable if The structural equations follow a DMSE model (Figure 1, solid edges) and are linear. Three independent views of z are available in the form of proxies u, v, w such that u ⊥ v ⊥ w | z and the equations between z and u, v, w are parameterized by matrices of rank dim ( z )
Our proof extends techniques developed by Kuroki and Pearl [22] to high-dimensional proxy variables.
Interestingly, our result crucially relies on an independence structure specified by Figure 1 (specifically, the existence of three independent proxy variables), which lends additional support for this modeling assumption and for the development of variational techniques specialized to this model family.
Computing Causal Effects Given a subset of modalities M , we can compute the ATE & ITE as E ( ITE ( x, M )) , where ITE ( x, M ) =
Y [ x, t = 1 , M ] − Y [ x, t = 0 , M ] and p ( y | x,
[Equation]
= [Equation] where we use our variational posterior formulation from Equation 6 to approximate the true posterior p ( z | x ) .
Experimental Results 5.1 Synthetic Demonstration Dataset
We start with a demonstration that provides intuition for why proxy variables are important, and how unstructured proxies can serve in place of featurized (structured) ones.
The following small-scale synthetic setup (Louizos et al.
[28]) involves a data distribution P over binary variables
y, t, z, x :
P (
[Equation]
=
[Equation]
( z =
[Equation]
= 0 .
[Equation]
where 0 < ρ x 1 , ρ x 0 , ρ t 1 , ρ t 0 < 1 are parameters.
We also introduce an unstructured proxy variable X that represents an “image version” of x .
The variable X will be a random MNIST image of a zero or one, depending on whether x = 0 or x = 1 .
Formally, X is distributed as follows:
P ( X | x = 1) is unif.
over MNIST images of ‘1’
[Equation]
[Table]
First, this is a setup that requires us to model proxies : treating X as a confounder as using a model of P ( y | X , t ) recovers the true ATE only when ρ t 1 = 1 − ρ t 0 and ρ x 1 = 1 − ρ x 0 (i.e., when X is perfectly informative of z ), otherwise it fails (see also Appendix A).
We also show that structural equations solve this task.
We sample 3000 data points from P and fit DGSE models to 80% of the data points { x, y, t } (the B INARY setting) as well as on { X , y, t } (the I MAGE setting).
We note the Average Treatment Effect (ATE) on the training and test sets, and we report results in Table 1.
We compare DGSE with the Inverse Probability of Treatment Weighted estimator (IPTW) [29] and the doubly robust Augmented-IPTW [45]—in each case the propensity score model is an MLP trained to predict t from either x or X . We found that replacing x with an image X causes the model to output highly miscalibrated probabilities close to 0 , 1 (while maintaining good accuracy), which results in large and volatile inverse propensity weights and in poor ATE estimates.
5.2 Benchmark Datasets for Causal Effect Estimation
IHDP The Infant Health and Development Project (IHDP) is a popular benchmark for causal inference algorithms [13] that contains the outcomes of comprehensive early interventions for premature, low birth weight infants.
We create a benchmark for multi-modal causal inference based on IHDP in which we can replace certain features with their ”unstructured version”.
We chooseof the 25 features available in IHDP in order to magnify their relative importance and accurately measure the effects of their removal.
Please refer to Appendix D and E for detailed setup.
[Table] STAR The Student-Teacher Achievement Ratio (STAR) experiment [1] studied the effect of class size on the performance of students.
We consider small class size as treatment; the outcome is the sum of the reading and math scores of a student.
We ‘derandomize’ this dataset by removing 80% of the data corresponding to white students in the treated population.
Similarly to IHDP, we selectattributes for the multi-modal experiment.
Further details can be found in the Appendix D and E.
Adding Unstructured Modalities We create a benchmark for multi-modal causal inference derived from IHDP and STAR in which we replace features with unstructured inputs that contain the same information as their featurized versions.
On
IHDP, we replace the attribute ‘baby’s gender’ with the CLIP embedding [24] of an image of a child between ages 3 to 8 years, drawn from the UTK dataset [61].
On STAR, we replace the attributes corresponding to the student’s ethnicity and gender by selecting an image of a child with the same ethnicity and gender from the UTK dataset.
We train and evaluate models on datasets where the image ‘replaces’ the attribute.
(e.g., 8 ATTRS + IMAGE ).
We also consider two other settings for comparison: a) the original attribute is included (e.g., 9 ATTRS ) and b) the attribute is dropped from the reduced set of input features (8 ATTRS ).
Results As seen in Table 2, the degradation in ATE error from replacing the baby’s gender by a photograph is lower as compared to removing the attribute entirely.
This shows our models leverage signal found in the unstructured image modality with the help of deep neural networks.
We compare these results with a simple Ordinary Least Squares model (OLS) baseline as described by Shalit et al.
[47] to predict treatment effect.
OLS shows a similar behavior when replacing baby’s gender with corresponding image, however ATE errors are generally worse as compared to ATE error produced by DGSE.
We also compare this with Counterfactual Regression Network (CFRNet) [16] baseline.
However CFRNet did not show benefits of using image modality unlike our approach.
In Table 6, replacing gender and ethnicity attributes on STAR with the corresponding image improves
ATE errors as compared to dropping these two attributes entirely.
This shows that we can use an image to extract multiple attributes while doing causal inference.
The CFRNet baseline shows a similar behavior, but the difference between average ATE errors across different setups is small.
5.3 Genome-Wide Association Studies
We evaluate our methods on an important real-world causal inference problem—genome-wide association study analysis (GWASs).
A GWAS is a large observational study that seeks to determine the causal effects of genetic markers (or genotypes) on specific traits (known as phenotypes).
In this setting, treatment variables are high-dimensional genomic sequences, and existing propensity scoring methods (which learn a model of p ( t | x ) for a binary t ) are not easily applicable [41, 50]—they may require training an impractical number of models. Our approach, on the other hand, can easily be used with high-dimensional t . Background and Notation As motivation
[Table]
consider the problem of linking a plant’s genetic variants t ∈ { 0 , 1 } d with nutritional yield , which we model via a variable y ∈ R .
Our goal is to determine if each variant t j is causal for yield, meaning that it influences biological mechanisms which affect this phenotype [8].
We also want to leverage large amounts of unstructured data x (e [Table] g., health records, physiological data) that are often available in modern datasets [7, 25] [Table]
[Table] A key challenge in finding causal variants is ancestry-based confounding [2, 52].
Suppose that we are doing a GWAS of plants from Countries A and B; plants in Country A get more rainfall, and thus grow faster and are more nutritious.
A simple linear model of y and t will find that any variant that is characteristic of plants in
Country A (e.g., bigger leaves to capture rain) is causal for nutritional yield.
Most existing GWAS analysis methods for estimating the effect of a variant t j rely on latent variable models: (1) they treat all remaining variants x as proxies and obtain z via a linear projection (e.g., PCA [38, 39] or LMM [59, 26]) of x into a lower dimensional space where genomes from Countries A and B tend to form distinct clusters (because plants from the same country breed and are similar); (2) we assume a linear model β ⊤ t of y and add z into it, which effectively adds the country as a feature ( z reveals the cluster for each country); this allows the model to regress out the effects of ancestry and assign the correct effect to variants t (one at a time).
Baselines.
Our main baselines are Principal Component Analysis (PCA) and Linear Mixed Model (LMM), as described above and implemented via the popular LIMIX library [27].
We also compare against Factor Analysis (FA), a standard linear technique for deriving latent variables, Uniform Manifold Approximation and Projection (UMAP), a manifold learning technique for dimensionality reduction [31] and a linear model with no correction for confounding.
DSE Models.
We compare against deep structural equation models that leverage one or more sets of proxy variables coming from the following unstructured modalities: genomic sequences, weather time series, simulated physiological time series.
We fit DSE models via a stagewise strategy analogous to how classical GWAS models are fit: (1) we fit the component p ( z k | x k ) for each proxy x k ; (2) we fit a linear model of y given t and the z to estimate causal effects.
Thus, our p ( x k | z ) components are deep, while p ( y | x, z ) is shallow (following standard assumptions on epistasis in GWAS).
5.3.1 Simulated Human GWAS
We used the 1000 Human Genomes [3] dataset to generate a simulated multi-modal GWAS dataset, following the biologically-inspired “Spatial” strategy studied in [49].
In addition to genotypes, we generated random physiological time series by sampling Fourier series conditioned on the confounders.
Please see Appendix G for details.
Table 3 shows our results.
Amongst the non-oracle baselines, multi-modal DSE has the smallest error in estimating causal effect, as well as the highest precision and recall at identifying causal SNPs.
The uni-modal DSE, while worse than LMM, still beats PCA and FA.
Note that the LMM model is not compatible with multiple unstructured proxies.
In general, we see that precision starts to deteriorate faster than recall, suggesting that false positives are more likely from the weaker linear deconfounding methods such as PCA/FA.
Our results here further support that additional sources of unstructured multi-modal data can improve GWAS.
5.3.2 Real-World Genomic Prediction and GWAS in Plants
[Table]
We also tested our methods on a realworld plant GWAS dataset from the 1001 Genomes Project for Arabidopsis
Thaliana plants [19, 55].
The challenge here is that we don’t know true causal effects; therefore, we define a phenotype y for which the true causal effect is zero—specifically, we set y to the GDP of the country where each plant grows.
Because the genomes of plants from the same country are similar, there exist spurious correlations caused by latent subpopulation groups z .
Our goal is to detect and correct for these confounding effects.
In addition to genomic data, we gather historical weather time series from the location of each plant (see Appendix F) and use both modalities to correct for confounding.
data to identify the latent confounding variable better.
[Table]
We evaluate whether each model learned to correctly account for latent confounding effects by measuring the predictive power of β ⊤ t for y , where β are the causal effects and t is the vector of variants.
Specifically, we compare the Rcorrelation between β ⊤ t and y —here, lower is better , since a model that has learned the causal effects should not be predictive of the phenotype.
In Table
[Table]
we see the effect of extracting confounding variables using DSEs as opposed to using the standard PCA technique.
We can see that the R 2 values produced using DSEs are closer to 0 as compared to using PCA.
This experiment shows that neural network architectures are effective in dealing with unstructured genomic and weather data while correcting for confounding.
Please refer to Appendix F for details.
5.4 Multimodal Experiments With Missing Modalities
We demonstrate the ability of the DMSE model to handle missing data.
Our IHDP and STAR benchmarks involve two modalities: images (e.g., baby’s gender in IHDP) and tabular data (e.g., the remaining features).
We compare DMSE and DGSE models on these datasets when some of the modalities may be missing.
For DMSE, we define two different modalities X 1 and X 2 for the tabular and image modalities respectively.
For DGSE, we concatenate the image embedding to the tabular modality while training the network.
We evaluate ATE while randomly dropping 50% of the images.
DMSE utilizes its product-of-experts inference network to approximate the posterior distribution when modalities are missing.
DGSE cannot do this, and we resort to feeding it a vector of zeros when an image is missing.
Table 4 shows that DMSE produces improved ATE estimates as compared to DGSE.
Related Work
Multi-Modal Causal Inference While previous work analyzed unstructured interventions t consisting of natural language [41–43] (e.g., determining the effect of a polite vs. a rude response) as well as unstructured y [5, 35] (e.g., MRI images), our work proposes methods to handle unstructured x .
Veitch et al. [50] developed models that correct for confounders from a single unstructured proxy x derived from text [33, 51] or a graph [50].
These approaches rely on a propensity scoring framework—they train a discriminative model of p ( t | x ) —hence do not support proxy variables or missing data, and require pre-trained text embeddings.
Additionally, propensity scoring methods rely on neural approximators for p ( t | x ) , which may output volatile probabilistic outputs that lead to unreliable effect estimates [18].
Our method (i) works across all modalities (beyond text or graphs), (ii) supports arbitrary numbers of proxies, (iii) supports missing data by virtue of being generative.
Deep Latent Variable Models Representation learning in causal inference has been studied by
Johansson et al.
[16, 15, 17] and Sch¨olkopf et al.
Deep latent variable models find applications throughout causal inference [28, 30, 57, 35, 60, 53, 20].
Pawlowski et al.
[35] study unstructured outcomes y (MRI scans), but do not support proxies.
Louizos et al.
[28] use variational auto-encoders to estimate confounders from proxies; we introduce a more structured model that handles multiple proxies that can be missing, and obviates the need for auxiliary modules.
Tran & Blei [49] propose implicit deep structural equations for GWAS; ours are explicit and thus easier to train.
Discussion
Identifiability Approaches to causal effect estimation that rely on deep learning [28, 47, 35, 48, 30, 11, 60, 58] can never guarantee the recovery of causal effects—neural network optimization is itself non-convex and has no guarantees.
Other failure modes of deep latent variable models (DLVMs) include potentially not having a sufficiently expressive model, not having enough data to learn the model, as well as shortcomings of approximate inference algorithms.
That said, there is ample evidence of both failures [44] and successes of DLVMs [57, 35, 30]. The DLVM approach is appealing over existing propensity scoring methods [41, 50] because: (i) it naturally handles unstructured proxies that may be missing at random; (ii) it supports high-dimensional treatment variables in settings like GWAS, where propensity scoring algorithms are not easily applicable.
Rissanen & Marttinen [44] empirically identify multiple failure modes of DLVMs; our work and that of others identifies success cases (particularly in GWAS [49, 54]), and ultimately the validation of DLVM methods needs to be empirical [48, 58, 50, 11, 20].
Missing Data We make the common assumption that data is missing at random (MAR).
This poses challenges if, for example, patients missing outcomes are ones that are more likely to be sick.
When two modalities and their missingness are correlated, their x j , x k nodes in Figure 1 could be merged, somewhat addessing the issue.
We leave the full exploration of non-MAR models to future work.  Conclusion
In conclusion, we proposed an approach based on deep structural equations that can leverage useful signal present in unstructured data to improve the accuracy of causal effect estimation.
Unlike previous methods that relied on propensity scores [41, 50], ours does not suffer from instabilities caused by volatile predictive probabilities coming out of neural networks, naturally handle missing data, and are applicable in settings in which the treatment variable is high-dimensional (such as in genome-wide association studies). Our work highlights the benefits of using large amounts of “dark” data that were previously left unused by existing methods to improve the accuracy of causal effect estimation.
A Details of the toy experiment.
Dataset.
In this experiment, we generate synthetic dataset consisting of 5 variables: a latent binary variable Z , a binary variable X ∗ , an MNIST image X , and two other binary variables y, t standing for outcome and treatment respectively.
[Equation]
These variables are sampled from the following distribution:
[Equation]
Accounting for Confounding in ATE Computation.
As mentioned before, the definition of the Average Treatment Effect (ATE) is as follows:
[Equation]
Now we consider the first term within the outer expectation:
[Equation]
[Equation]
[Equation]
[Equation]
z
[Equation] [Equation]
z
[Equation]
[Equation]
[Equation]
[Equation] [Equation] [Equation] where the second last equation is due to the fact that y = t ⊕ Z .
Similarly, we can compute the second term as
[Equation]
Using the predefined generative process for this dataset, we can also write
[Equation]
where x is an image of binary variable x ∗ .
Plugging in the previous results, we can compute the individual treatment effect (ITE):
[Equation]
Therefore we can plug in the previous equation and get the final result of ATE:
[Equation]
ATE Computation with Non-Causal Model. ATE computation goes wrong when X is taken to be the only confounder. In the following computation of E [ y | X = x, do ( t = 1)] , the second step goes wrong as adjustment is done over X instead of Z .
[Equation]
of 0) · P [ y = 1 | X = image of 0 , t = 1]
= P ( X ∗ = 1) · P [ y = 1 | X ∗ = 1 , t = 1] + P ( X ∗ = 0) · P [ y = 1 | X ∗ = 0 , t = 1] = P ( X ∗ = 1) · P [ y = 1 , X ∗ = 1 , t = 1]
P [ X ∗ = 1 , t = 1] + P ( X ∗ = 0) · P [ y = 1 , X ∗ = 0 , t = 1]
P [ X ∗ = 0 , t = 1] = P ( X ∗ = 1) · (cid:80)
z ( P [ y = 1 , X ∗ = 1 , t = 1 | Z
= z ] P ( Z = z ))
(cid:80) z
( P [ X ∗ = 1 , t = 1 | Z
=
z ] P ( Z = z ))
+ P ( X ∗ = 0) · (cid:80) z ( P [ y = 1 , X ∗ = 0 , t = 1 | Z
= z ] P ( Z = z ))
(cid:80) z
( P [ X ∗ = 0 , t = 1 | Z
=
z ] P ( Z = z ))
=
P ( X ∗ = 1) · ( P [ y = 1 , X ∗ = 1 , t = 1 | Z = 0] · 0 . 5) (cid:80) z ( P [ X ∗ = 1 , t = 1 | Z
=
z ] · 0 . 5)
+ P ( X ∗ = 0) · ( P [ y = 1 , X ∗ = 0 , t = 1 | Z = 0] · 0 . 5) (cid:80) z
( P [ X ∗ = 0 , t = 1 | Z
=
z ] · 0 . 5)
=
{+P ( X ∗ = 1) · ρ x 0 ρ tρ x 1 ρ t 1 + ρ x 0 ρ t 
+}
+
P ( X ∗ = 0) · ((1 − ρ x 0 ) ρ t 0 ) (1 − ρ x 1 ) ρ t 1 + (1 − ρ x 0 ) ρ t Similarly, we get wrong expectation when performing intervention on variable t to set it to zero. [Equation]
do
[Equation]
Also,
[Equation] With the above expressions, we can check that setting ρ x 0 = 0 . 1 , ρ t 0 = 0 . 2 , ρ x 1 = 0 .
3 , ρ t 1 = 0 . 4 gives a non-zero ATE when computed using non-causal methods (i.e. without accounting for hidden confounder). This corresponds to the non-causal baseline in Table 1.
To get ATE of zero using non-causal baseline, we need to set ρ t 1 = 1 − ρ t 0 and ρ x 1 = 1 − ρ x 0 . B Neural Architecture of Deep Structural Equations and Approximate Inference Networks
Our Architecture In this section, we add the details of the DGSE and DMSE architecture that we used.
X i denotes an input datapoint, i.e. the feature vector (possibly containing multiple modalities), t i is the treatment assignment, y i denotes the corresponding outcome and Z i is the latent hidden confounder.
Within DGSE and DMSE, the latent variable is modeled as a Gaussian.
For DGSE, we write (similar to Louizos et al.
[Equation]
where σ ( · ) is the sigmoid function, Bern is the Bernoulli distribution, D x , D z are the dimensions of
X and Z respectively, and p [ X ij | Z i ] is an appropriate probability distribution for the covariate j .
If the treatment variable is not binary, we can modify the distribution appropriately.
Within DMSE, it is possible to further factorize the distribution p ( X i | Z i ) into product of distributions over component modalities owing to the conditional independence.
If the outcome y is discrete, we parameterize its probability distribution as a Bernoulli distribution:
[Equation]
and if it is continuous, we parameterize its distribution as a Gaussian with a fixed variance ˆ v , defined as:
[Equation]
Here each of the NN i ( · ) is a neural network.
The posterior distribution for DGSE is approximated as
[Equation]
where
[Equation]
For DMSE, the posterior distribution is computed differently using Product-of-Experts (PoE) [56] formulation, due to which it can handle missing modalities during training and inference gracefully.
The objective of DGSE model is the variational lower bound defined as:
[Equation]
The DMSE, on the other hand, requires a sub-sampled training objective to ensure that the modality specific posterior networks are trained and the relationships between individual modalities is captured.
For DGSE, we also define the auxiliary encoders and the extra term in the variational lower bound following Louizos et al.
Auxiliary Encoders:
[Equation]
For discrete y i , we have
[Equation]
For continuous y i , we write
[Equation]
This introduces the following extra term in the variational lower bound:
[Equation]
DMSE does not involve these extra terms within its ELBO objective.
Compared with Louizos et al.
[28], we can extend DGSE to different types of architectures for the posterior distribution
q [ Z i | X i , t i , y i ] .
When X is an image (e.g. medical scans, patient photos), we can use a suitable Convolutional Neural Network (CNN) architecture for extracting information effectively [23].
In our experiments with image modality, we used pretrained CLIP embeddings [24] in the first layer to extract relevant features from the images.
To avoid the overwhelming difference between the image and two binary variables t, y , we also apply dimension reduction techniques such as Principle Component Analysis to the embeddings of the image before feeding it into the network that is shared with t, y .
When X is time-series data, (e.g. text, recording), we can change the architecture to recurrent neural networks such as Long Short Term Memory [14].
More generally, we can choose modality specific architectures and make appropriate design choices to perform learning and inference over unstructured modalities as inputs.
DMSE can handle different types and lengths of modalities gracefully and also work with missing modalities owing the specialized variational learning and inference procedures.
C Comparing Our Methods with Other VAE- Based Estimators
While our method is an instance of generative models, we identify the following key differences: We propose new generative model architectures that extend existing models (e g., DSE, CEVAE) to multiple proxies
X i , each possibly coming from a different modality.
We derive novel inference algorithms for these extended models, which have the following benefits:
(a) Our algorithms scale better to large sets of modalities by leveraging the independence structure of the X i .
(b) Our inference algorithms naturally handle missing X i .
(c) They are also simpler: they don’t require auxiliary networks (e.g., like in CEVAE [28]).
Lastly, our key contribution is that we demonstrate the effectiveness of generative models at modeling unstructured proxies (many previous methods instead relied on propensity scoring).
Appendix H.5.1 empirically shows that DMSE model compares favorably against CEVAE on synthetic datasets.
D Setups used for IHDP and STAR Dataset experiments D.1 IHDP Experiments
The data corresponding to non-white mothers in the treated set of children is removed so that causal effect of the intervention cannot be estimated directly.
The column corresponding to mother’s race is removed so that this confounder cannot be obtained directly from the input.
We considerreplicates of this dataset, where the output is simulated according to setting ’A’ of NPCI package [9].
The true treatment effect is known as the simulation provides expected output values for both values of binary treatment variable.
We train a DGSE model on each replicate with a 63/27/10 ratio of training, validation and test dataset size.
We set the latent dimension to be 20 units and the number of hidden layers to be 2.
The hidden layers have size of 20 units.
The IHDP-Full Setting There are 25 input features in this experimental setting.
We report the absolute error in ATE produced by DGSE and OLS for this setting in Table 7.
The IHDP-Mini Setting Here, we choose 9 features from the 25 input features so that removal of the feature ‘baby’s gender’ produces statistically significant treatment effect.
We used mutual information and F-statistics between each of the original 25 features and the target variable y to assess the importance of each feature in the initial 100 replicates of IHDP.
While making sure that the absolute ATE errors don’t deviate too much from the corresponding errors produced by IHDPFull setting, we experimented with several combinations of the high ranking features to select the following 9 features in the IHDP-Mini setting.
Feature 6: ‘sex of baby’ 2.
Feature 0: ‘birth-weight’ 3.
Feature 1: ‘b.head’ 4.
Feature 2: ‘preterm’ 5.
Feature 3: ‘birth.o’ 6.
Feature 8: ‘mom married?’
Feature 9: ‘mom’s education lower than high school?’
Feature 12: ‘Smoked cig during pregnancy?’
Feature 20: ‘harlem’
Table 7 shows a comparison of absolute ATE errors between the IHDP-Full setting and IHDP-Mini setting.
Table 2 shows the comparison of our approach with CFRNet [16] and Ordinary Least Squares (OLS) approach.
OLS takes the concatenation of the covariates and treatment variable value as input to produce output.
We see that the CFRNet baseline does not utilize the image effectively while
OLS shows small difference between the setting where baby’s gender was removed (8 attrs) and the setting where baby’s gender was retained (9 attrs) setting.
Hence the replacement of image ( attrs+image) produces a small average improvement as compared to the setting where baby’s gender was dropped (8 attrs).
We also note that we have dropped 6 replicates from the 100 IHDP replicates under consideration.
These 6 replicates showed a large degradation in ATE estimates by adding baby’s gender (9 attrs) as compared to removing it (8 attrs).
D.2 STAR Experiments
We ‘derandomize’ this dataset by removing 80% of the data corresponding to white students in the treated population.
The dataset [1] has 15 input attributes.
The true treatment effect can be estimated directly since the original dataset corresponds to a randomized controlled trial.
Hence, it is possible to compute ATE error as the absolute difference between true ATE estimate and the ATE as predicted by the model.
Similar to IHDP [Table] we choose the following subset of attributes for the STAR experiment
Feature 2: ‘Student grade’ 2.
Feature 3: ‘Student class-type’ 3.
Feature 4: ‘Highest degree obtained by teacher’ 4.
Feature 5: ‘Career ladder position of teacher’ 5.
Feature 6: ‘Number of years of experience of teacher’ 6.
Feature 7: ‘Teacher’s race’ 7.
Feature 10: ‘Student’s gender’ 8.
Feature 11: ‘Student’s ethnicity’
[Table]
In Table
[Table]
we repeat the experiment 100 times and report average ATE errors along with standard error.
We removed 8 repetitions in the experiment where DGSE or CFRNet showed lack of convergence as evidenced by very high validation loss on any setting of input attributes.
[Table]
E Evaluating Quality of Pre-Trained Embeddings
We demonstrate that our pre-trained embeddings contain useful signal by building a neural network model that predicts the gender, age and ethnicity from the CLIP embedding [24] of the corresponding image.
We build a simple neural network model that takes the CLIP embedding of an image as input and predicts the age of the person in that image.
We use 5-dimensional PCA embeddings of 500 randomly chosen images of people aged 10-45 years (corresponding to the age-group in the IHDP experiment).
We have an independent test dataset corresponding to 100 images chosen in a similar way.
We see that the R 2 value for the age prediction on test dataset is 0.45.
If we increase the size of PCA embedding to 50, this R 2 value increases to 0.58.
Thus, it is possible to extract the age information from randomly chosen images using a simple neural network.
In the above setting, we also studied the classification accuracy of separate neural networks that predict gender and ethnicity from the CLIP embeddings.
We saw that gender was predicted with 94% accuracy and ethnicity was predicted with 58% accuracy using a 5-dimensional PCA embedding.
After increasing the size of PCA embedding to 50 dimensions, the gender prediction accuracy increased to 95% and ethnicity prediction accuracy increased to 77%.
This further supports our idea of replacing the attribute corresponding to ’s baby’s gender or student ethnicity/gender with an appropriate image.
F Plant GWAS
Setup We apply our deep structural equations framework for correcting the effects of confounding.
We fit a DGSE model in two stages: (1) first, we only fit the model of p ( T | Z ) using the DGSE
ELBO objective; (2) then we fit p ( Y | Z, T ) with a fixed Z produced by the auxiliary model q ( Z | T ) .
We found this two-step procedure to produce best results.
The subsampled SNPs corresponding to each genome are taken as input X .
The encoder and decoder use a single hidden layer of 256 units while a 10-dimensional latent variable Z is used.
This network is optimized using ClippedAdam with learning rate of 0.01, further reduced exponentially over 20 training epochs.
The confounding variable for each genome can now be computed as latent representation produced by the DSE.
To measure the success of confounding correction, we compute the R 2 values between the true GDP of the region and the GDP output as predicted via our model and the baselines.
If we have corrected for confounding, then we should get low R 2 values.
Historical Weather Data We used historical weather data collected from Menne et al.
[32] to add a new modality while performing plant GWAS.
We use per day precipitation data from yearcollected by weather station closest in distance to the latitude/longitude coordinates of the location from which the SNPs of plant were collected.
For the locations where weather data was missing, we replaced those entries with zeros.
G Simulated GWAS Experiments
We provide additional details on this experiment here.
G.1 Data generating process
To simulate the confounders, SNPs (genotypes), and the outcomes (phenotypes), we follow the “Spatial” simulation setup from Appendix D.1 & D.2 of [49].
Specifically, we generate random lowrank factorization of the allele frequency logits F = σ − 1 (Γ S ) , where σ is sigmoid, as is common in the literature [4, 40]. In the “Spatial” Setting [49], the S matrix is interpreted as geographic spatial positions of the individuals.
For the m -th SNP of the n -th individual, we generate the SNP
X nm ∼ Bin(3 , σ ( F nm )) .
In this simulation, we considered M = 10 SNPs with N =individuals.
Individuals are clustered into K = 3 groups based on their locations, and the individual’s cluster is the unobserved confounder.
Then, the outcomes are calculated from both the SNPs, where only c = 2 SNPs have a non-zero causal effect, plus a confounding term that is a function of the cluster, and some i.i.d.
Gaussian noise.
One small deviation from Tran & Blei [49] is that our noise is i.i.d.
Gaussian (hence, the variance does not depend on the confounder).
We further augmented the dataset to include a time-series proxy that can help identify the unobserved spatial position of individuals.
For each cluster, we came up with some Fourier coefficients,
Cluster 0: a (0)0 = 0 . 0 , a (0)1 = − 1 . 0 , a (0)
[Equation]
= 1 . 0 , b (0)1 = − 1 . 0 , b (0)2 = 1 . Cluster 1: a (1)0 = 1 . 0 , a (1)1 = − 5 . 0 , a (1)2 = 2 . 0 , b (1)1 = − 5 . 0 , b (1)2 = 2 . Cluster 2: a (2)0 = − 1 . 0 , a (2)1 = − 2 . 0 , a (2)2 = 5 . 0 , b (2)1 = − 2 . 0 , b
(2)2 = 5 . 0 .
For each individual in the k -th cluster, we sampled at N samples =
50 uniformly spaced times, across N periods = 2 periods of length T = 5 .
That is, our time series proxy consists of points { x i } i ∈
[ N samples ] ,
[Equation]
Please see Fig.
3 for a visualization of the time-series generated for each cluster.
G.2 DSE Modeling and evaluation setup
To handle the SNP data, we trained a Multi-Layer Perceptron (MLP) based β -variational autoencoder [12] with ReLU activations, which output parameters of a Gaussian latent distribution.
The MLPs consisted of 3 hidden layers, each with 64 units, and our latent dimension z was 2.
To handle time series data (new modality), we trained a β -VAE that encoded and decoded with 1D convolutions, as 1D convolutions have been shown to strongly approximate fast Fourier transforms [34].
There werehidden convolutional layers, each of which had 32 output channels, with a kernel of length 3 and a stride of 1.
Following the convolution was an MLP with 1 hidden layer consisting of 64 units.
ReLU activations were used throughout the architecture.
In our experiments, we swept over β ∈ { 0 . 1 , 0 . 2 , 0 .
Our general hypothesis was that lower β values would make our VAE perform better, as unlike in the standard VAE evaluation scheme, we sample from the posterior to generate latents for individuals, and not the prior, so the KL divergence term in the ELBO doesn’t matter as much.
When evaluating on a set of SNPs and time series data corresponding to a set of individuals, we generate latents by passing the data of each respective type into the appropriate encoder and sampling from the resulting distributions.
The latent z for each individual is the concatenation of the SNP latent z snp and the time series latent z ts .
Once we have the latent confounder z for an individual, we calculate the causal effect of the m th SNP by performing the linear regression,
[Equation]
where ϵ i are i.i.d.
standard normal.
We solved this using the closed form solution for least-squares linear regression, given by
[Equation]
where A i =
[ z i , 1 , x m ] is the i th row in the matrix A .
Then (cid:98) γ m is our estimated causal effect for the m th SNP.
We let (cid:98) γ be a vector of length M to denote our estimates for each SNP’s causal effect.
G.3 Results
We tested our single modality DSE model and multi-modal DSE model against several linear latent model baselines.
Two baselines, Principle Components Analysis (PCA) and Factor Analysis (FA), also generate a latent z , which we then use within our linear regression approach to calculate (cid:98) γ .
We also ran the Linear Mixed Model (LMM) implementation by Limix [27].
We also plot two oracle baselines.
The first is when (cid:98) γ
= γ ⋆ , labeled “truth”.
The second is when z is the true confounder used when generating the data, according to “Spatial”.
Let γ ⋆ denote the ground truth causal effect vector.
We report ℓ 1 and ℓ 2 norms to measure ∥ (cid:98) γ − γ ⋆ ∥ .
We also report true/false positive/negatives, which we define as follows.
For an individual, let τ = min i : γ i ̸ {+=
+} | γ i | / 2 .
Then, SNP m is a
True Positive (tp): if (cid:98)
γ m ≥ τ and γ ⋆m ≥ τ or (cid:98) γ m ≤ − τ and γ ⋆m ≤ − τ .
True Negative (tn): if | (cid:98) γ m | < τ and | γ ⋆m | < τ . 3.
False Negative (fn): if | (cid:98) γ m | < τ but | γ ⋆m | > τ . 4.
False Positive (fp): if none of the above hold.
Concretely, there are two cases.
First, if | (cid:98) γ m | > τ but | γ ⋆m | < τ .
Second, if the direction is wrong, i e. (cid:98) γ m > τ but γ ⋆m < − τ , or (cid:98) γ m < − τ but γ ⋆m > τ .
Finally, recall that precision is tp/ ( tp + fp ) and recall is tp/ ( tp + fn ) , and higher is better for both.
In the following table, we show the mean and standard error of the mean (sem) over 10 seeds at evaluation time.
The unimodal DSE only encodes the SNPs vector into a 2 dimensional embedding.
The multimodal DSE also encodes the time-series vector into a 2 dimensional embedding, and we concatenate this to the 2 -dimensional embedding of SNPs to form a 4 dimensional embedding that can be used in our linear regression.
PCA k denotes Principal Components Analysisn(PCA) with k components. Similarly, FA k denotes Factor Analysis (FA) with k components.
Amongst the non-oracle baselines, multi-modal DSE has the smallest errors and highest precision/recall. It beats all linear baselines including LMM.
Unimodal DSE, while worse than LMM, still beats PCA and FA.
In general, we see that precision starts to deteriorate faster than recall, suggesting that false positives are more likely from the weaker linear deconfounding techniques such as PCA/FA.
It’s also interesting that when PCA/FA used more components, we see worse performance across the board.
This illustrates the importance of picking the right size of latent dimension, which is an open question.
If the latent dimension is too small then we may not capture all the confounders.
If the latent dimension is too large, the latent encoding may also end up capturing spurious correlations.
H Alternative Causal Graph Structures
This section derives estimators of ATE and ITE under assumptions on the causal graph that are different from the ones described in Figure 1.
The key takeaway of this section is that ATE and ITE can be computed under arbitrary causal graphs using minor modifications of our algorithms, which we describe in this section.
H.1 Alternative Causal Graph Structures
We start by formally defining the class of causal graphs that will be studied in this section.
Formally, we identify two sets of modification to Figure 1.
We will derive small modifications to our algorithms to allow the new types of causal graphs implied by these modifications.
[Table] [Table] Observed confounders We consider an expanded set of graphs over a space of random variables ( X, Y, T, Z, V ) , where V represents observed (non-proxy) confounders and the other random variables are associated with the observed data x, y, t, z .
We look at causal graphs implied by structural equations of the form:
[Equation]
where P X j , P Y are probability distributions with a tractable density and the µ, σ, π, θ are functions parameterized by neural networks that output the parameters of their respective probability distribution as a function of ancestor variables in the causal graph.
In the above equations, V is variable that is assumed to be always observed (just like y ( i ) and t ( i ) ).
All the other technical terms are defined as in the main body of the paper.
The above equations result in a causal graph with edges between V and Y, T and define a distribution p ( x, y, t, z, v ) .
Dependent proxies Another possible set of modifications to Figure 1 is the presence of edges between proxies
X i , X j , which can be denoted as
[Equation]
where pa ( X j )) denotes the set of parents of X j among the other unstructured proxy variables X i .
H.2 A General Estimator Class
The following Theorem shows that we can estimate ATE and ITE when the data distribution follows the structure in Figure 1, plus the two types of modifications outlined above (observed confounders and dependent proxies).
Theorem 3.
The true ITE ( x, M ) for any subset M ⊆ { 1 , 2 , ..., m } of observed modalities is identifiable when the true data distribution p ( x, y, t, z, v ) has the causal graph structure of a DMSE model (Figure 1) in addition to having observed confounders v and possibly dependent confounders.
Proof: Let x M = { x j | j ∈ M } be the data from the observed subset of modalities.
Let v be the observed proxy variable.
We need to show that p ( y | x M , v, do ( t = t ′ )) is identifiable for any t ′ .
Observe that
[Equation]
where the second equality follows from the rule of do-calculus (applying backdoor adjustment).
Since our proof holds for any t ′ and all elements on the right-hand side are identifiable, the claim follows.
H.3 Observed Confounders
Next, we derive an extension of the DMSE model to the setting in which we have observed confounders V .
We refer to this modified model as DMSE-V .
As earlier, the DMSE-V model induces a tractable joint density p ( X, Y, T, Z, V ) , which allows us to fit its parameters using stochastic variational inference by optimizing the evidence lower bound (ELBO) on the marginal log-likelihood p ( y ( i ) , x ( i ) , t ( i ) , v ( i ) ) defined over an expanded dataset { y ( i ) , x ( i ) , t ( i ) , v ( i ) } ni =1 :
[Equation]
i ) j | z ) + log p ( y ( i ) , t ( i ) , v ( i ) , z ) − log q ( z | x ( i ) , y ( i ) , t ( i ) , v ( i ) )   , (9) where p ( y ( i ) , t ( i ) , v ( i ) , z ) =
p ( y ( i ) | t ( i ) , v ( i ) , z ) p ( t ( i ) | z, v ( i ) )
p ( z ) p ( v ( i ) ) and q ( z | x ( i ) , y ( i ) , t ( i ) , v ( i ) ) is the approximate variational posterior.
Note that since v ( i ) is always observed (just like y ( i ) and t ( i ) ), the p ( v ( i ) ) term can be ignored.
In practice, this reduces to the vanilla DMSE model with the following modifications:
• The log p ( y ( i ) | t ( i ) , v ( i ) , z ) term becomes additionally conditioned on v ( i ) .
• The log p ( t ( i ) | v ( i ) , z ) term becomes additionally conditioned on v ( i ) .
• The approximate posterior q ( z | x ( i ) , y ( i ) , t ( i ) , v ( i ) ) becomes additionally conditioned on v ( i ) .
Crucially, the specialized inference algorithms derived for the DMSE model remain unchanged.
Since
V is always observed, learning a DMSE-V model is equivalent to learning a model p ( y, x, t, z | v ) , which has the same structure as a DMSE model.
In particular, all the X i are conditionally independent given Z .
Hence, the same learning and inference algorithms apply.
H.4 Causal Links Among Proxies
Another type of causal graph that we consider is one in which proxies
X i are connect by causal edges.
First, we note that when proxies are unstructured, such causal edges are expected to be rare, i.e., we do not expect the pixels of an image X i to have a direct influence on other variables.
When X i takes on a structured form and directly influences other proxies, our strategy is to “collapse” any sets of variables X j that have edges among them, until we have the conditionally independent structure in Figure 1.
In the extreme case, we might need to collapse all proxies X i into a single proxy X that consists of their concatenation.
The result is a model that has the same structure as DMSE, and that can be learned using the same set of inference and learning algorithms.
The only drawback is increased computational efficiency.
H.5 Additional Synthetic Data Experiment Details
We generate synthetic data to simulate a process with modified causal links as covered in subsection H.3 and H.4.
Specifically, we simulate four different datasets with generative process detailed in the following section.
We generate { X i , t i , y i } i
= m i =0 where m = 10000 .
We use train/val/test split of 63/27/10.
Here, ⊕ stands for logical XOR.
1. Case of Original Causal Graph (Dataset A)
Variables: X 1 is unstructured input, X 2 is structured input, T is treatment, Y is output, Z is confounder
Edges in graph: Z → X 1 , Z → X 2 , { Z } → T, { T, Z } → Y
Thus we added no extra edges.
Generative process:
(a) P ( z = 1) = 0 .(b) P ( x 1 ′ =
| z = 1)
=
P ( x 1 ′ =
| z = 0)
= 0 .
1 ( x 1 ′ is an intermediate variable) (c) P ( x 2 = 1 | z = 1) =
P ( x 2 =
| z = 0)
= 0 .(d)
P ( t = 1 | z = 1) =
P ( t = 0 | z = 0)
= 0 .(e) y
=
( z ⊕ t ) (f) P ( x 1 | x 1 ′ = 1) is unif.
over MNIST images of ‘1’, P ( x 1 | x 1 ′ = 0) is unif.
over MNIST images of ‘0’
2. Case of the Observed Confounder (Dataset B)
Variables: X 1 is unstructured input, X 2 is structured input, T is treatment, Y is output, Z is confounder
Edges in graph: Z → X 1 , { X 2 , Z } → T, { X 2 , T, Z } → Y
Thus we added extra edges X 2 → T and X 2 → Y .
Generative process:
(a) P ( z = 1) = 0 .
5 , P ( x 2 = 1) = 0 .(b) P ( x 1 ′ =
| z = 1)
=
P ( x 1 ′ =
| z = 0)
= 0 .
1 ( x 1 ′ is an intermediate variable)
(c) P ( t = 1 | z
= 1 , x 2 = 1) =
P ( t = 0 | z = 0 , x 2 = 1) = 0 .
2; P ( t = 1 | z
= 1 , x 2 = 0) =
P ( t = 0 | z = 0 , x 2 = 0)
= 0 .(d)
y = x 2 AND ( z ⊕ t ) (e) P ( x 1 | x 1 ′ = 1) is unif.
over MNIST images of ‘1’, P ( x 1 | x 1 ′ = 0) is unif.
over MNIST images of ‘0’
3. Case of the Observed Confounder (Dataset C)
Variables: X 1 is unstructured input, X 2 is structured input, T is treatment, Y is output, confounder Z
=
( Z 1 , Z 2)
Edges in graph: Z → X 1 , Z → X 1 , Z − > X 2 , { X 2 , Z } → T, { X 2 , T, Z } → Y
Thus we added extra edges Z → X 2 , X 2 → T and X 2 → Y .
Generative process:
(a) P ( z 1 = 1) = 0 .
5 , P ( z 2 = 1) = 0 .(b)
P ( x 1 ′ = [-
-]1 | z 1 = 1) =
P ( x 1 ′ =
| z 1 = 0)
= 0 .
1 ( x 1 ′ is an intermediate variable)
(c) P ( x 2 = 1 | z 2 = 1) =
P ( x 2 = 0 | z 2 = 0)
= 0 .(d)
[-X 1 → X 
 Data generation:(a) P ( z 1 = 1) = P (-] z = z 1 ⊕ z(e) P ( t = 1 |
z
= 1 , x 2 = 1) =
P ( t = 0 | z = 0 , x 2 = 1) = 0 .
9; P ( t = 1 | z = 1 , x 2 = 0)
=
P ( t = 0 | z = 0 , x 2 = 0)
= 0 .(f)
[-P ( x 2 = 1 | z 2 = 1 , x 1 = 1) = P ( x 2 = 
 | z 2 = 0 , x 1 = 1) = 0 .-]y
= x 2 ⊕ t ⊕ ( z 1 ⊕ z 2) (g) P ( x 1 | x 1 ′ =[-
-]
) is
unif
over
MNIST images of ‘1’, P ( x 1 | x 1 ′ = 0)
is
unif.
over MNIST images of ‘0’
4. Case of Causal Links Among Proxies (Dataset D)
Variables: X 1 is structured input, X 2 is structured input, X 3 is unstructured input, T is treatment, Y is output, confounder Z
= ( Z 1 , Z 2 , Z 3)
Edges in graph: Z → X 1 , Z → X 2 , Z → X 3 , X 1 → X 2 , { Z } → T, { T, Z } → Y
[-P ( x ′ i = 
 | z i = 0) =-]Thus we added extra edge: X 1 → X
Data generation:
Case
3 z + 2) , P ( y = 1 | t = 0) = sigmoid (3 z − 2) (f) P ( x i | x ′ i = 1) is unif.
over MNIST images of ‘1’, P ( x i | x ′ i = 0) is unif.
over MNIST images of ‘0’
[Table]
[Table] H.5.1 Experimental Results DMSE under alternative graph structures and comparison with CEVAE :We generate synthetic data corresponding to Datasets A, B, C and D. Table 13 demonstrates that DMSE recovers ATE under modified causal graph structures. We demonstrate that with increasing dataset size, the ATE error on test dataset continues to fall.
Thus, with dataset size approaching infinity, we can recover the true ATE as long as the model class contains true distribution and our optimizer can find the minimum.
In Table 12, DMSE also compares favorably with CEVAE and recovers the ATE in this extended setting.
In this extended setting, we have a combination of structured and unstructured input modalities.
CEVAE takes a concatenation of these modalities as its input while DMSE has a separate model and inference network for each modality.
Hence DMSE can handle diverse modality types gracefully as compared to CEVAE.
[Table]
[Table]
DMSE under increasing number of modalities and comparison with CEVAE : We generate synthetic data corresponding to Dataset E with varying number of modalities (i.e. proxies).
Tablecontains the results of experiment with just one modality; the key difference between the two models is the inference procedure.
Table 11 shows that even in this setting, DMSE recovers the ATE more accurately than CEVAE.
Next, we compare CEVAE vs. DMSE when the data has many unstructured modalities in Table 10.
We generate synthetic data from K modalities (Dataset E).
The CEVAE model treats them as one concatenated vector; DMSE models them as separate vectors.
As expected, DMSE handles large numbers of modalities better than CEVAE.
I Identifiability of Causal Effects in the Presence of Proxies to Hidden Confounder
If the structural equations have linear dependencies, we can establish the identifiability of total treatment effect.
Kuroki & Pearl [22] define the total treatment effect of X on Y as ‘the total sum of the products of the path coefficients on the sequence of arrows along all directed paths from X to Y’.
We extend the identifiability result of Kuroki & Pearl [22] to handle vector-valued confounders ¯ U , and hence we require one additional view of the confounder for a total of three independent views.
Concretely, we consider the following setup.
¯ U : hidden confounder X : binary treatment variable
¯ W, ¯ Z, ¯ V : proxies for ¯ U
Assume causal graph as
[Equation]
U
Assume linear structural equations
[Equation]
Note that the three views of ¯ U are independent given ¯ U , i.e. ¯
W ⊥ ¯ Z ⊥ ¯ V | ¯ U .
Theorem 4.
Given the above setup, the causal effect of X on Y , i.e. P ( y | do ( x )) , is identifiable if the dependency matrices β ⋄ ¯ U for ⋄ ∈ { ¯ V , ¯ W, ¯ Z } have rank | ¯ U | .
Proof Sketch First, we analyze what happens when we know the β matrices a priori, i.e. treatment effect estimation with external information.
Then, we analyze how to leverage the three independent views of the confounder to recover enough information about the matrices for treatment effect estimation, i.e. treatment effect estimation without external information.
In this second step, we use the rank constraint so that each proxy has enough information about the confounder.
Treatment effect estimation with external information :
If any one β matrix corresponding to a proxy is known (i.e. we know the causal mechanism ¯ U → ¯ W , ¯
U → ¯ Z or ¯ U → ¯ V ), then the treatment effect is identifiable given that the β matrix is invertible.
In principle, we need any one proxy and the β matrix corresponding to it for computing the treatment effect.
The following derivation is done using the matrix β ¯ W ¯ U .
We define σ AB = Covariance ( A, B ) where A and B are univariate.
We define Σ ¯ A ¯ B to be the covariance matrix between sets of variables in the vectors ¯ A and ¯ B .
For the above causal graph Figure 4 with linear dependencies, the total treatment effect τ Y X of variable X on output Y can be derived using backdoor adjustment [22] as
[Equation]
Using formula of covariance and owing to the linear nature of structural equations, we know
that Σ X
¯ W
=
Σ
X
[Equation]
U
[Equation]
T ¯ W ¯ U , Σ Y
¯ W = Σ
Y ¯ U β T ¯ W ¯ U .
Thus, if the matrix β ¯ W ¯ U is a square and invertible, we get
[Equation]
Note : If the matrix β ¯ W ¯ U is not square, it needs to have rank | ¯ U | at the least.
In case | ¯ W | > | ¯ U | , but rank( β ¯ W ¯ U ) =
| ¯ U | , we can still work with a modified system of equations such that ¯ W is replaced with ¯ W ′ where | ¯ W ′ | = | ¯ U | = rank ( β ¯ W ′ ¯ U ) .
To obtain this modified system, we apply elementary row transformations on the equation ¯ W = β
¯ W ¯
U
¯ U
+ ¯ ϵ W such that β ¯
W ¯ U is an upper triangular matrix.
After that, we drop the bottom | ¯ W | − | ¯ U | rows from matrices on both sides of equation to obtain ¯ W ′ and
β ¯ W ′
¯ U .
Consider the case where rank( β ¯ W ¯ U ) <
| ¯ U | and rank( β ¯ Z ¯ U ) <
| ¯ U | .
In this case, if we can concatenate the two proxies ¯ W and ¯ Z to a new proxy ¯ V
=
¯
W : ¯
Z such that rank( β ¯ V ¯ U ) =
| ¯ U | , we can estimate treatment effects using β ¯ V ¯ U .
Treatment effect estimation without external information :
In this case, we need three independent views of hidden confounder U.
As we are working with multivariate confounders, the univariate treatment variable X cannot serve as the third view of U anymore.
Now using the properties of covariance matrix and the linearity of structural equations, we can write the following
[Equation]
Assuming that the matrices Σ ¯ U ¯ U , β ¯ V ¯ U , β ¯ Z ¯ U are square and invertible, we get Σ ¯ W ¯ Z Σ − 1¯ V ¯ Z Σ ¯ V ¯ W
=
β
¯ W ¯ U Σ ¯
U ¯ U β T
¯ W ¯
U
Hence, if β ¯ W ¯ U is invertible additionally, then the treatment effect can be identified using the following equation
[Equation]
What happens when the dimensionality of proxies is greater than dimensionality of true confounder?
We are interested in high-dimensional, unstructured proxies.
In this case, our matrices β ¯ V ¯ U , β ¯ Z ¯ U , β ¯ W ¯ U will need to have rank | ¯ U | .
In essence, we still need three views for our vector ¯ U .
When we have access to external information in the form of β matrices, verifying the rank is possible.
We can then apply appropriate row transformations on the corresponding structural equations (as mentioned in I) to obtain modified proxies with same length as the confounder ¯ U .
In absence of external information, rank of β matrices is not verifiable.
Assuming that we have β ¯ V ¯ U , β ¯ Z ¯ U , β ¯ W ¯ U with rank | ¯ U | , we apply a dimensionality reduction procedure to map the proxies ¯ V , ¯ Z, ¯ W to modified proxies ¯ V ′ , ¯ Z ′ , ¯ W ′ such that | ¯
V
| =
| ¯ Z ′ | =
| ¯ W ′ | = | ¯ U | and β ¯ V ′ ¯ U , β ¯ Z ′ ¯ U , β ¯ W ′ ¯
U have rank | ¯ U | . In practice, this dimensionality reduction can be a technique like PCA or the result of applying pre-trained neural network layer to obtain an embedding of unstructured data. J Additional Details on Mathematical Proofs J.1 Evidence Lower Bound for Deep Structural Equations We discuss the evidence lower bound (ELBO), presented in Section 4.2.
The left-hand-side of the ELBO can be written as
[Equation]
Due to the causal graph structure we can factorize the following distribution as
[Equation]
Thus, we can rewrite (10) as n (cid:88)
[Equation]
)  
[Equation]
Note that the first four terms in the expectation over the posterior distribution q make up a reconstruction loss of the original data x , y, t in terms of the latent variable z , while the last two terms form the negative KL divergence − D KL ( q ( z | x i , t i , y i )
∥ p ( z )) between the posterior q and the prior p ( z ) .
Thus, we can write our ELBO as
[Equation]
which is exactly our variational objective.
We can sum or take the average over the number of datapoints n to form our empirical objective.
J.2 Derivation of posterior
Here we show why the true posterior factorizes as [Equation] =
[Equation]
y ) (by cond.
indep.
of causal graph)
=
[Equation]
(removing terms independent of z )