BayesIMP: Uncertainty Quantification for Causal Data Fusion

1 Introduction . Causal inference has seen a significant surge of research interest in areas such as healthcare [ 1 ] , ecology [ 2 ] , and optimisation [ 3 ] . However , data fusion , the problem of merging information from multiple data sources , has received limited attention in the context of causal modelling , yet presents significant potential benefits for practical situations [ 4 , 5 ] . In this work , we consider a causal data fusion problem where two causal graphs are combined for the purposes of inference of a target variable ( see Fig.1 ) . In particular , our goal is to quantify the uncertainty under such a setup and determine the level of confidence in our treatment effect estimates . Let us consider the motivating example in Fig.1 , where a medical practitioner is investigating how prostate cancer volume is affected by a statin drug dosage . We consider the case where the doctor 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ∗Denotes equal contribution with alphabetical ordering only has access to two separate medical studies describing the quantities of interest . On one hand we have observational data , from one medical study D1 [ 1 ] , describing the causal relationship between statin level and prostate specific antigen ( PSA ) , and on the other hand we have observational data , from a second study D2 [ 7 ] , that looked into the link between PSA level and prostate cancer volume . The goal is to model the interventional effect between our target variable ( cancer volume ) and the treatment variable ( statin ) . This problem setting is different from the standard observational scenario as it comes with the following challenges : • Unmatched data : Our goal is to estimate E [ cancer volume|do ( statin ) ] but the observed cancer volume is not paired with statin dosage . Instead , they are related via a mediating variable PSA . • Uncertainty quantification : The two studies may be of different data quantity/quality . Furthermore , a covariate shift in the mediating variable , i.e . a difference between its distributions in two datasets , may cause inaccurate extrapolation . Hence , we need to account for uncertainty in both datasets . Formally , let X be the treatment ( Statin ) , Y be the mediating variable ( PSA ) and T our target ( cancer volume ) , and our aim is to estimate E [ T |do ( X ) ] . The problem of unmatched data in a similar context has been previously considered by [ 5 ] using a two-staged regression approach ( X → Y and Y → T ) . However , uncertainty quantification , despite being essential if our estimates of interventional effects will guide decision-making , has not been previously explored . In particular , it is crucial to quantify the uncertainty in both stages as this takes into account the lack of data in specific parts of the space . Given that we are using different datasets for each stage , there are also two sources of epistemic uncertainties ( due to lack of data ) as well as two sources of aleatoric uncertainties ( due to inherent randomness in Y and T ) [ 8 ] . It is thus natural to consider regression models based on Gaussian Processes ( GP ) [ 9 ] , as they are able to model both types of uncertainties . However , as GPs , or any other standard regression models , are designed to model conditional expectations only and will fail to capture the underlying distributions of interest ( e.g . if there is multimodality in Y as discussed in [ 10 ] ) . This is undesirable since , as we will see , interventional effect estimation requires accurate estimates of distributions . While one could in principle resort to density estimation methods , this becomes challenging since we typically deal with a number of conditional/ interventional densities . In this paper , we introduce the framework of Bayesian Interventional Mean Processes ( BAYESIMP ) to circumvent the challenges in the causal data fusion setting described above . BAYESIMP considers kernel mean embeddings [ 11 ] for representing distributions in a reproducing kernel Hilbert space ( RKHS ) , in which the whole arsenal of kernel methods can be extended to probabilistic inference ( e.g . kernel Bayes rule [ 12 ] , hypothesis testing [ 13 ] , distribution regression [ 14 ] ) . Specifically , BAYESIMP uses kernel mean embeddings to represent the interventional distributions and to analytically marginalise out Y , hence accounting for aleatoric uncertainties . Further , BAYESIMP uses GPs to estimate the required kernel mean embeddings from data in a Bayesian manner , which allows to quantify the epistemic uncertainties when representing the interventional distributions . To illustrate the quality of our uncertainty estimates , we apply BAYESIMP to Causal Bayesian Optimisation [ 15 ] , an efficient heuristic to optimise objective functions of the form x∗ = arg minx∈X E [ T |do ( X ) = x ] . Our contributions are summarised below : 1 . We propose a novel Bayesian Learning of Conditional Mean Embedding ( BAYESCME ) that allows us to estimate conditional mean embeddings in a Bayesian framework . 2 . Using BAYESCME , we propose a novel Bayesian Interventional Mean Process ( BAYESIMP ) that allows us to model interventional effect across causal graphs without explicit density estimation , while obtaining uncertainty estimates for E [ T |do ( X ) = x ] . 3 . We apply BAYESIMP to Causal Bayesian Optimisation , a problem introduced in [ 15 ] and show significant improvements over existing state-of-the-art methods . Note that [ 16 ] also considered a causal fusion problem but with a different objective . They focused on extrapolating experimental findings across treatment domains , i.e . inferring E [ Y |do ( X ) ] when only data from p ( Y |do ( S ) ) is observed , where S is some other treatment variable . In contrast , we focus on modelling combined causal graphs , with a strong emphasis on uncertainty quantification . While [ 17 ] considered mapping interventional distributions in the RKHS to model quantities such as E [ T |do ( X ) ] , they only considered a frequentist approach , which does not account for epistemic uncertainties . Notations . We denote X , Y , Z as random variables taking values in the non-empty sets X , Y and Z respectively . Let kx : X × X → R be positive definite kernels on X with an associated RKHS Hkx . The corresponding canonical feature map kx ( x′ , · ) is denoted as φx ( x′ ) . Analogously for Y and Z . In the simplest setting , we observe i.i.d samples D1 = { xi , yi , zi } Ni=1 from joint distribution PXY Z which we con- catenate into vectors x : = [ x1 , ... , xN ] > . Similarly for y and z . For this work , X is referred as treatment variable , Y as mediating variable and Z as adjustment variables accounting for confounding effects . With an abuse of notation , features matrices are defined by stacking feature maps along the columns , i.e Φx : = [ φx ( x1 ) , ... , φx ( xN ) ] . We denote the Gram matrix as Kxx : = Φ > x Φx and the vector of evaluations kxx as [ kx ( x , x1 ) , ... , kx ( x , xN ) ] . We define Φy , Φz analogously for y and z. Lastly , we denote T = f ( Y ) + as our target variable , which is modelled as some noisy evaluation of a function f : Y → T on Y while being some random noise . For our problem setup we observe a second dataset of i.i.d realisations D2 = { ỹj , tj } Mj=1 from the joint PY T independent of D1 . Again , we define ỹ : = [ ỹ1 , ... , ỹM ] > and t : = [ t1 , ... , tM ] > just like for D1 . See Fig.2 for illustration . 2 Background . Representing interventional distributions in an RKHS1 has been explored in different contexts [ 18 , 17 , 19 ] . In particular , when the treatment is continuous , [ 17 ] introduced the Interventional Mean Embeddings ( IMEs ) to model densities in an RKHS by utilising smoothness across treatments . Given that IME is an important building block to our contribution , we give it a detailed review by first introducing the key concepts of do-calculus [ 20 ] and conditional mean embeddings [ 21 ] . 2.1 Interventional distribution and do-calculus In this work , we consider the structural causal model [ 20 ] ( SCM ) framework , where a causal directed acyclic graph ( DAG ) G is given and encodes knowledge of existing causal mechanisms amongst the variables in terms of conditional independencies . Given random variables X and Y , a central question in interventional inference [ 20 ] is to estimate the distribution p ( Y |do ( X ) = x ) , where { do ( X ) = x } represents an intervention on X whose value is set to x . Note that this quantity is not directly observed given that we are usually only given observational data , i.e , data sampled from the conditional p ( Y |X ) but not from the interventional density p ( Y |do ( X ) ) . However , Pearl [ 20 ] developed do-calculus which allows us to estimate interventional distributions from purely observational distributions under the identifiability assumption . Here we present the backdoor and front-door adjustments , which are the fundamental components of DAG based causal inference . The backdoor adjustment is applicable when there are observed confounding variables Z between the cause X and the effect Y ( see Fig . 3 ( Top ) ) . In order to correct for this confounding bias we can use the following equation , adjusting for Z as p ( Y |do ( X ) = x ) = ∫ Z p ( Y |X = x , z ) p ( z ) dz . The front-door adjustment applies to cases when confounders are unobserved ( see Fig . 3 ( Bottom ) ) . Given a set of front-door adjustment variables Z , we can again correct the estimate for the causal effect fromX to Y with p ( Y |do ( X ) = x ) = ∫ Z ∫ X p ( Y |x ′ , z ) p ( z|X = x ) p ( x′ ) dx′dz . We rewrite the above formulae in a more general form as we show below . For the remainder of the paper we will opt for this notation : p ( Y |do ( X ) = x ) = EΩx [ p ( Y |Ωx ) ] = ∫ p ( Y |Ωx ) p ( Ωx ) dΩx ( 1 ) For backdoor we have Ωx = { X = x , Z } and p ( Ωx ) = δxp ( Z ) where δx is the Dirac measure at X = x . For front-door , Ωx = { X ′ , Z } and p ( Ωx ) = p ( X ′ ) p ( Z|X = x ) . 1We refer the reader to a detailed review of RKHS methods provided in the Appendix of [ 5 ]

Summary: The paper proposes a framework for embedding interventional distributions into RKHS that accounts for uncertainty in causal datasets especially when the confounders, treatments, mediators, and outcomes are not all observed in a single observational dataset. The main problem that is addressed here is quantifying the uncertainty in the estimation of causal effects when the treatment and the outcome are measured in two separate sources (potentially under different data-generating conditions) and the aim is to estimate the effect from a causal data fusion perspective. The paper introduces Bayesian Interventional Mean Processes for representing interventional distributions in the reproducing kernel Hilbert space. The method is applied to the task of Causal Bayesian Optimization where it is shown that the method converges faster as compared to other baseline methods. The main contribution of the paper is obtaining the uncertainty estimates for $E\[T\mid do(X) = x\]$.