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 Bayesian approach for estimating the treatment effect by combining two data sources. The main advantage of the Bayesian approach over the frequentist counterpart [17] is its ability to quantify uncertainty. The main machinery was GP which replaces KRR used in existing works.