Bayesian Federated Estimation of Causal Effects from Observational Data
Abstract
We propose a Bayesian framework for estimating causal effects from federated observational data sources.
Bayesian causal inference is an important approach to learning the distribution of the causal estimands and understanding the uncertainty of causal effects.
Our framework estimates the posterior distributions of the causal effects to compute the higher-order statistics that capture the uncertainty.
We integrate local causal effects from different data sources without centralizing them.
We then estimate the treatment effects from observational data using a non-parametric reformulation of the classical potential outcomes framework.
We model the potential outcomes as a random function distributed by Gaussian processes, with deﬁning parameters that can be efﬁciently learned from multiple data sources.
Our method avoids exchanging raw data among the sources, thus contributing towards privacy-preserving causal learning.
The promise of our approach is demonstrated through a set of simulated and real-world examples.
INTRODUCTION
Causal effect estimation is important in many real-life situations.
For example: What is the effect of war in a speciﬁc region on world food supply?
How would the blood pressure of a patient change if that patient took a new drug?
How does coronary heart disease affect age- and gender-speciﬁc mortality rates?
These questions are common in many areas, including personalized medicine [Powers et al., 2018], digital experiments [Taddy et al., 2016], political science [Green and Kern, 2012], etc., and especially recent events in the
Covid-19 pandemic and the war of Ukraine.
In practice, the relevant data essential for accurate and meaningful causal inference may reside in multiple, decentralized data sources which cannot be shared or combined due to geographical, organizational, process, and/or privacy constraints.
Some alternative solutions such as establishing data use agreements or creating secure data environments may not be possible and are often not easily implemented.
In addition, it is important to know whether the causal estimands are reliable.
Thus, estimating a conﬁdence interval of the relevant causal effect together with its point estimates would give helpful insights into the uncertainty of the causal estimand.
For example, a narrow conﬁdence interval for individual treatment effect of smoking on lung cancer, where zero falls outside the conﬁdence interval, means that the patient is at a higher risk of getting cancer.
Most of the recent causal effect estimators, e.g., Louizos et al.
[2017], Shalit et al.
[2017], Madras et al.
[2019], are point estimates without considering the uncertainty of the causal estimands.
Bayesian approaches, e.g., Imbens and
Rubin [1997], Daniels et al.
[2012], Talbot et al.
[2015], Gutman et al.
[2018], Ning et al.
[2019], on the other hand, aim to learn the posterior distributions of the causal estimands to obtain higher-order statistics that capture the uncertainty.
To derive these posterior distributions of the causal estimands, however, most, if not all of the existing efforts involve pooling the distributed data from multiple sources centrally to compute the model marginal likelihood, thus violating the privacy constraints mentioned above.
We propose a Bayesian framework that can learn the causal effects of interest without combining data sources to a central site, and, at the same time, learn higher-order statistics of the causal effects to understand their uncertainty.
This federated learning approach [McMahan et al., 2017] has not been well studied for causal inference.
Our contributions are summarized as follows:
• We propose the Federated Causal Inference (FedCI)framework that fuses federated learning and causal inference to incorporate multiple data sources while maintain-
ing the sources at their local sites.
– FedCI generalizes the Bayesian imputation approach [Imbens and Rubin, 2015] to a more generic model based on Gaussian processes (GPs); the resulting model is decomposed into multiple components, each of which handles a distinct data source.
– FedCI minimizes information transmitted among the sources, thus enabling privacy-preserving causal inference.
The framework could support multiparty computation and differential privacy in future.
• We propose a variational approximation scheme for the proposed model, whose evidence lower bound can be decomposed additively across different data sources.
This allows the parameters to be optimized via federated gradient averaging.
We then leverage the computed predictive distribution to efﬁciently estimate the desired treatment effect quantities.
• We empirically evaluate the proposed framework on benchmark datasets, and show its competitive performance as compared to the recent baseline approaches trained on the combined datasets.
RELATED WORK
Causal inference.
In most causal inference literature, the estimation of causal effects is performed directly on accessible local data sources.
Hill [2011], Alaa and van der
Schaar [2017, 2018] proposed nonparametric approaches to estimate causal effects.
A growing literature, including
Shalit et al.
[2017], Yoon et al.
[2018], Yao et al.
Künzel et al.
[2019], Nie and Wager [2021], used parametric methods to model the potential outcomes.
Louizos et al.
[2017], Madras et al.
[2019] used the formulation of Pearl [1995] to estimate causal effects under the existence of latent confounding variables.
Bica et al.
[2020a,b] formalized potential outcomes for temporal data with observed and unobserved confounding variables to estimate counterfactual outcomes for treatment plans.
Imbens and Rubin [1997],
Daniels et al.
[2012], Talbot et al.
[2015], Gutman et al.
[2018], Ning et al.
[2019] are typical Bayesian methods that learn posterior distributions of the causal estimands.
All these works were not proposed for the context of multisource data which cannot be shared and combined as a uniﬁed dataset.
Our model, in contrast, learns treatment effects while preserving the source data at their local sites.
It is different from the problem of transportability of causal relations [e.g., Pearl and Bareinboim, 2011, Bareinboim and Pearl, 2013b,a, 2016, Lee et al., 2020], where theoretical tools were developed to transport causal effects from a source population to a target population, which does not take into account the above data privacy constraint.
Federated learning.
Federated learning aims to train an algorithm across multiple decentralized clients, thus respect the privacy information of the data [McMahan et al., 2017].
Federated stochastic gradient descent [Shokri and
Shmatikov, 2015] and federated averaging [McMahan et al., 2017] are two variations of federated learning.
Recent developments of federated learning, e.g., Álvarez et al.
Zhe et al.
[2019], de Wolff et al.
[2020], Joukov and Kuli´c [2020], Sattler et al.
[2019], Mohri et al.
[2019] are formalized for a typical classiﬁcation or regression problem.
Also, it has recently been applied in facilitating multi-institutional collaborations without sharing patient data [Rieke et al., 2020, Sheller et al., 2020] and healthcare informatics [Lee and Shin, 2020, Xu et al., 2021].
Other biomedical applications of federated learning include predicting adverse drug reactions [Choudhury et al., 2019], stroke prevention [Ju et al., 2020], mortality prediction [Vaid et al., 2020], predicting outcomes in SARS-COV-2 patients [Flores et al., 2021], etc.
Ng and Zhang [2022], Gao et al.
[2021] are noticeable works that estimate causal graphs in federated setting, which is different from our work in estimating casual effects.
Following some recent works [e.g., Shalit et al., 2017, Yao et al., 2018, Oprescu et al., 2019, Künzel et al., 2019, Nie and Wager, 2021], we develop a federated causal inference algorithm based on the potential outcomes framework.
We summarize the related models in the subsequent sections.
2.1 POTENTIAL OUTCOMES AND THE BAYESIAN IMPUTATION MODEL
The concept of potential outcomes was proposed in Neyman [1923] for randomized trial experiments.
Rubin [1975, 1976, 1977, 1978] re-formalized the framework for observational studies.
We consider the causal effects of a binary treatment w , with w = 1 indicating assignment to ‘treatment’ and w = 0 indicating assignment to ‘control’.
Following the literature, the causal effect for individual i is deﬁned as a comparison of the two potential outcomes, y i (0) and y i (1) , where these are the outcomes that would be observed under w i = 0 and w i = 1 , respectively.
We can never observe both y i (0) and y i (1) for any individual i , because it is not possible to go back in time and expose the i –th individual to the other treatment.
In this work, we generalize the
Bayesian imputation model of Imbens and Rubin [2015] since it captures uncertainty of the causal estimands in a
Bayesian setting:
[Equation]
where (cid:15)
i and (cid:15)
i are the Gaussian noises.
The key to compute treatment effects is y i (0) and y i (1) .
So we need to impute one of the two outcomes.
Let y i, obs , y i, mis be the observed and unobserved (or missing) outcome.
The idea is to ﬁnd the marginal distribution p ( y i, mis | y obs , X , w ) .
Once the missing outcomes are imputed, the treatment effects can be estimated.
To proceed, Imbens and Rubin [2015] suggested four steps based on the following equation p
[Equation]
(cid:82) p ( y i, mis |
y obs , X , w , θ ) p ( θ | y obs , X , w ) dθ , where θ = { β 0 , β 1 } .
The aim is to ﬁnd p ( y i, mis | y obs , X , w , θ ) and p ( θ | y obs , X , w ) , and then compute the integral to obtain p ( y i, mis | y obs , X , w ) , which is a non-parametric prediction.
The above procedure shows that learning the distribution p ( y i, mis | y obs , X , w ) would require data from all sources since it is conditional on y obs , X , and w .
Thus, it violates the data privacy constraint.
In Sections 3.3, 3.4 and 3.5, we generalize this model with Gaussian processes and decompose it into multiple components to perform federated inference of the causal effects, which minimizes the risk of privacy leak of the data.
OUR APPROACH
We generalize the Bayesian imputation model presented in
Section 2.1 to a generic model based on Gaussian Processes (GPs).
We introduce the Federated Causal Inference (FedCI) method to decompose the model into multiple components, each associated with a data source, to estimate causal effects under a federated setting.
3.1 PROBLEM FORMULATION
Problem setting & notations.
Suppose we have m data sources that are organized and curated at their local sites.
Each source is denoted by D s = { ( w s i , y s i, obs , x s i ) } n s i =1 , where s = 1 , 2 , . . .
, m , and the quantities
w s i , y s i, obs and x s i are the treatment assignment, observed outcome associated with the treatment, and covariates of individual i in source s , respectively.
In this work, we focus on binary treatment w s i ∈ { 0 , 1 } , thus y s i, obs can be either of the potential outcomes y s i (0) or y s i (1) , i.e., for each individual i , we can only observe either y s i (0) or y s i (1) , but not both of them.
We further denote the unobserved or missing outcome as y s i, mis .
These variables are related to each other through the following equations:
[Equation]
Thus, y s i (1) =
y s i, obs when w s i =
and y s i (1) =
y s i, mis when w s i = 0 , and similarly for y s i (0) .
For notational convenience, we further denote y s (0) =
[ y s 1 (0) ,..., y s n s (0)] (cid:62) , y s obs =
[ y s 1 , obs ,..., y s n s , obs ] (cid:62) , and similarly for y s (1) , y s mis ,
X s and w s .
Causal effects of interest.
We estimate the individual treatment effect (ITE) 2 and the average treatment effect (ATE) deﬁned as follows:
[Equation]
where y s i (1) , y s i (0) are realization outcomes of the corresponding random variables, and n = (cid:80) m s =1 n s .
The causal estimands.
Inserting Eq.
(2) and (3) into (4), we obtain the estimate of ITE:
[Equation]
where y obs , X , w denote the vectors/matrices of the observed outcomes, covariates and treatments concatenated from all the sources.
The estimate of ATE is as follows:
[Equation]
where 1 is a vector of ones.
Hence, the remaining task is to learn the posterior p ( y mis (cid:12)(cid:12) y obs , X , w ) , which is the predictive distribution of y mis given all the covariates, treatments and observed outcomes from all sources.
3.2 ASSUMPTIONS
The following assumptions are made to enable federated causal estimations:
Assumption 1 (Strong Ignorability) .
( i )
The potential outcomes are independent of the treatment assignment conditional on the covariates ( unconfoundedness ) , i.e., y s i (1) , y s i (0) ⊥⊥
w s i | x s i , and ( ii ) every individual has some positive probability to be assigned to every treatment ( positivity ) , i.e., 0 < p ( w s i = 1 | x s i ) < 1 .
[Rosenbaum and
Rubin, 1983]
Assumption 2 (Stable Unit Treatment Value Assumption or SUTVA) .
( i )
The potential outcomes for any individual do not vary with the treatments assigned to other individuals, and ( ii ) there are no different forms or versions of each treatment level, which would lead to different potential outcomes.
[Imbens and Rubin, 2015]
Assumption 3.
The individuals from all sources share the same set of common covariates.
Assumption 4.
There exists a set of features such that any individual is uniquely identiﬁed across different sources.
We refer to this set as ‘primary key’.
Assumption 5.
Data in different sources are drawn from parts of the population.
The multi-source data, which may be homogeneous or heterogeneous in nature, together reﬂect the characteristics of the population.
Assumption 1 and 2 are standards in causal inference, as discussed in, e.g., Imbens and Rubin [2015], Shalit et al.
Assumption 3 is reasonable, e.g., decentralized data in Choudhury et al.
[2019], Vaid et al.
[2020], Flores et al.
[2021] (to name a few) satisfy this assumption for federated learning.
In Assumption 4, a ‘primary key’ is not limited to the observed data used for inference as described in Section 3.1, but it can include any features to uniquely identify an individual, such as { nationality, national id } of a patient.
Assumption 4 allows a secure preprocessing procedure to remove repeated individual records in different sources, if necessary, without sharing raw data among the sources (see
Appendix A for details).
Assumption 5 ensures that there is no imbalanced data bias across the sources.
In the subsequent sections, we assume that all of the above assumptions are satisﬁed, and the preprocessing procedure is already performed if necessary.
3.3 GP-BASED IMPUTATION
The model presented in Eq.
(1) is a simple Bayesian linear model.
In this section, we present a more general nonlinear
Bayesian model.
In particular, since β 0 and β 1 in Eq.
(1) follow multivariate normal distributions, the two components β (cid:62)
x i and β (cid:62) 1 x i also follow multivariate normal distributions.
The generalisation of these two components are f 0 ( x i ) =
β (cid:62) 0 ω ( x i ) and f 1 ( x i ) =
β (cid:62) 1 ω ( x i ) , where ω ( x i ) is a vector of basis functions with input x i .
This formulation would lead to the fact that the marginal of f 0 ( x ) and f 1 ( x ) are Gaussian processes (GPs).
Thus, we propose:
[Equation]
where f 0 ( x i ) and f 1 ( x i ) are two random functions evaluated at x i , i.e., f 0 ( x i ) ∼ GP ( µ 0 ( X ) , K ) and
f 1 ( x i ) ∼
GP ( µ 1 ( X ) , K ) , where K denotes the covariance matrix computed with a kernel function k ( x , x (cid:48) ) .
Similar to the imputation model of Imbens and Rubin [2015], this model also requires ﬁnding the marginal distribution p ( y i, mis |
y obs , X , w ) , through accessing the observed data from all the sources.
Similarly, although this model is generic, it requires access to all the observed data to compute K , which is impossible without violating the privacy constraints mentioned above.
In the subsequent sections, we propose a federated model to address this problem.
3.4 THE PROPOSED MODEL
Recall that the aim is to ﬁnd p ( y mis | y obs , X , w ) so that we may in turn compute Eqs.
(5)-(8) to arrive at the quantities of interest.
To that end, we propose to model the joint distribution of the potential outcomes as follows:
[Equation]
where ε s i ∼ N ( 0 , I 2 ) is to handle the noise of the outcomes.
As mentioned in Section 2.1 and 3.3, all the outcomes from all sources are interdependent in the Bayesian imputation approach, which is problematic for federated learning.
This dependency is handled via f s j ( x i ) and g s j ( j ∈ { 0 , 1 } ), which enable federated learning for the proposed model.
We refer to the dependency handled by f s j ( x i ) as intra-dependency and the one captured by g s j as inter-dependency.
Intra-dependency.
f s 0 ( x i ) and f s 1 ( x i ) are GP-distributed functions, which allows us to model each source dataset simultaneously along with its heterogeneous correlations.
Speciﬁcally, we model f s 0 ( x i ) ∼ GP ( µ 0 ( X s ) , K s ) and
[Equation]
∼
[Equation]
1 ( X s ) , K s ) , where K s is a covariance matrix computed by a kernel function k
( x s i , x s j ) , and µ 0 ( · ) , µ 1 ( · ) are functions modelling the mean of these GPs.
Parameters of these functions and hyperparameters in the kernel function are shared across multiple sources.
The above GPs handle the correlations within one source only.
Inter-dependency.
To capture dependency among the sources, we introduce variable g =
[ g 0 , g 1 ] , where g 0 =
[ g 10 ,..., g m 0 ] (cid:62) ∼ N ( r 0 , M ) and g 1 =
[ g 11 ,..., g m 1 ] (cid:62) ∼
N ( r 1 , M ) .
Both g s
and g s 1 are shared within the source s , and they are correlated across multiple sources s ∈ { 1 ,..., m } .
The correlations among the sources are modelled via the covariance matrix M which is computed with a kernel function.
The inputs to the kernel function are the sufﬁcient statistics (we used mean, variance, skewness, and kurtosis) of each covariate x s within the source s .
We denote the ﬁrst four moments of covariates as (cid:101) x s ∈ R 4 d x ×and the kernel function as γ ( (cid:101) x s , (cid:101) x s (cid:48) ) , which evaluates the correlation of two sources s and s (cid:48) .
This formulation implies that g 0 and g 1 are GPs.
The elements of r 0 and rare computed with the mean functions r 0 ( (cid:101) x s ) and r 1 ( (cid:101) x s ) , respectively.
Herein, we only share the sufﬁcient statistics of covariates, but not covariates of a speciﬁc individual.
The two variables Ψ and Σ .
These are positive semideﬁnite matrices capturing the correlations between the two possible outcomes y s i (0) and y s i (1) , Ψ 12 and Σ 12 are their Cholesky decomposition matrices.
Note that Ψ and Σ are also random variables.
Since these are positive semideﬁnite matrices, we model their priors using Wishart distribution Ψ ∼ Wishart ( V 0 , d 0 ) , Σ ∼ Wishart ( S 0 , n 0 ) , where
V 0 , S 0 ∈ R 2 × 2 are predeﬁned positive semi-deﬁnite matrices and d 0 , n 0 ≥ 2 are predeﬁned degrees of freedom.
3.5 THE PROPOSED ALGORITHM
Based on some results on the joint distribution of potential outcomes, we construct a federated objective function for the proposed federated causal inference algorithm (FedCI).
3.5.1 Joint Distribution of the Outcomes
We ﬁrst present some results that are helpful in constructing the federated objective function in Section 3.5.2.
The proofs of these results are in the appendices.
To simplify the exposition, we denote g s =
[ g s 0 , g s 1 ] , where g s
= [ g s 0 ,..., g s 0 ] (cid:62) and g s 1 = [ g s 1 ,..., g s 1 ] (cid:62) .
Lemma 1.
Let Ψ , Σ , K , µ 0 ( X s ) , µ 1 ( X s ) , and g s satisfy the model in Eq.
Then,
[Equation]
where ⊗ is the Kronecker product.
The proof of Lemma 1 is presented in Appendix C. From
Lemma 1, we observe that Ψ , K s , Σ , and I n s are positive semi-deﬁnite, thus the covariance matrix Ψ ⊗ K s + Σ ⊗ I n s is positive semi-deﬁnite due to the fundamental property of Kronecker product.
This is why we choose Ψ and Σ to be positive semi-deﬁnite in our model; otherwise, the covariance matrix is invalid.
From Lemma 1, we can obtain the result in Lemma 2 as follows:
Lemma 2.
Let Ψ , Σ , K , µ
( X s ) , µ 1 ( X s ) , and g s satisfy the model in Eq.
Then,
[Equation]
The mean functions µ obs ( X s ) and µ mis ( X s ) are:
[Equation]
where we denote m
= ψ ∗ 11 ( µ 0 ( X s ) + g s 0 ) and m 1 = ψ ∗ 21 ( µ 0 ( X s ) + g s 0 ) + ψ ∗ 22 ( µ 1 ( X s ) + g s 1 ) with ψ ∗ ij is the ( i, j ) –th element of Cholesky decomposition matrix of Ψ ,is a vector ones
[Equation]
and (cid:12) is the element-wise product.
The covariance matrices K s obs , K s mis , and K s om are computed by kernel functions:
[Equation]
where ψ ab and σ ab are the ( a, b ) –th elements of Ψ and Σ , respectively.
The proof of Lemma 2 is in Appendix D. Lemma 2 has two important roles in our work.
First, we can obtain the conditional likelihood to help infer the parameters and hyperparameters of our proposed model.
Second, we can also obtain the posterior of y s mis to help us estimate ITE and ATE.
3.5.2 Federated Objective Function
The proposed model in Eq.
(10) would lead to an objective function that can be decomposed into m components, each associated with a data source.
Since estimating p ( y s mis (cid:12)(cid:12) y s obs , X s , w s ) exactly is intractable, we sidestep this intractability via a variational approximation.
To achieve this, we maximize the following evidence lower bound
[Equation]
where
[Equation]
Herein, D KL [ · ] is the Kullback–Leibler divergence.
Details of the ELBO are presented in Appendix B.
The conditional likelihood p ( y s obs |· ) is obtained from Lemma 2 by marginalizing out y s mis , i.e.,
[Equation]
The above conditional likelihood is free of σ 21 and σ 12 , which capture the correlation of two potential outcomes.
Thus the posteriors of these variables would coincide with their priors, i.e., the correlation cannot be learned but set as a prior.
This is well-known as one of the potential outcome cannot be observed [Imbens and Rubin, 2015].
In
Eq.
(11), the ELBO L is derived from the of joint marginal likelihood of all m sources, and it is factorized into m components L s , each component corresponds to a source.
This enables federated optimization of L .
The ﬁrst term of L s is expectation of the conditional likelihood with respect to the variational posterior q ( g , Ψ , Σ) , thus this distribution is learned from data of all the sources.
In the following, we present its factorization.
Variational posterior distributions.
We use the typical mean-ﬁeld approximation to factorize among the variational posteriors q (Ψ , Σ , g ) =
q (Ψ) q (Σ) q ( g ) .
Let (cid:101) y s obs (0) , (cid:101) y s obs (1) , (cid:101) x s , and (cid:101) w s ( s = 1 , 2 ,..., m ) be the ﬁrst four moments of the observed outcomes, covariates, and treatment of the s –th source.
Let (cid:101) X
=
[ (cid:101) x 1 ,..., (cid:101)
x m ] (cid:62) , (cid:101) y obs (0) =
[ (cid:101) y 1 obs (0) ,..., (cid:101)
y m obs (0)] (cid:62) , (cid:101) y obs (1) =
[ (cid:101) y 1 obs (1) ,..., (cid:101)
y m obs (1)] (cid:62) , and (cid:101) w =
[ (cid:101) w 1 ,..., (cid:101) w m ] (cid:62) .
We parameterize q ( g ) =
(cid:81) j ∈{ 0 , 1 } N ( g j ; h j ( (cid:101)
y obs (0) , (cid:101) y obs (1) , (cid:101) X , (cid:101) w ) , U ) , where h 0 ( · ) and h 1 ( · ) are the mean functions, U is the covariance matrix computed with a kernel function κ ( u s , u s (cid:48) ) , where u s : =
[ (cid:101)
y s obs (0) , (cid:101) y s obs (1) , (cid:101)
x s , (cid:101) w s ] .
Since Ψ and Σ are positive semi-deﬁnite matrices, we model their variational posterior as Wishart distribution:
[Equation]
where d q , n q are degrees of freedom, V q , S q are the scale matrices.
We set the form of these scale matrices as follows
[Equation]
where ν i , ρ, δ i , η are parameters to be learned and ρ, η ∈ [0 , 1] .
Reparameterization.
To maximize the ELBO, we approximate the expectation in L s with Monte Carlo integration, which requires drawing samples of g , Ψ and Σ from their variational distributions.
This requires a reparameterization to allow the gradients to pass through the random variables g , Ψ and Σ .
The reparameterization trick for g are: g j =
[Equation]
y obs
[Equation] [Equation]
y
[Equation] [Equation]
} , where ξ j ∼ N ( 0 , I m ) and U 12 is the Cholesky decomposition matrix of U .
Since q (Ψ) is a Wishart distribution, we introduce the following procedure to draw Ψ : Ψ
[Equation]
q )
[Equation]
, ζ ∼ Wishart ( I 2 , d q ) , where V 12 q is the
Choleskey decomposition matrix of V q .
Likewise, we also apply this procedure to draw Σ .
Federated optimization algorithm.
With the above model and its objective function, we can compute gradients of the learnable parameters separately in each source without sharing data to a central server.
We summarize our procedure in Algorithm 1.
3.5.3 Predicting Causal Effects from Multiple Sources
To understand why data from all the sources can help predict causal effects in a source s , we observe that
[Equation]
Eq.
(13) is an approximation of the predictive distribution of the missing outcomes y s mis
and it depends on the following three components:
(i).
The observed outcomes, covariates and treatment assignments from the same source s .
(ii).
The shared parameters Θ learned from data of all the sources.
(iii).
Sufﬁcient statistics of the observed data from all the sources.
The two last components (ii) and (iii) indicate that the predictive distribution in source s utilizes knowledge from all the sources through Θ and the sufﬁcient statistics [ (cid:101) y obs (0) , (cid:101)
y obs
[Equation]
1) , (cid:101) X , (cid:101) w ] .
This explain why data from all of the sources help predict missing outcomes in source s .
Parameters: Let Θ be set of parameters 1 begin
Initialize Θ and send to all source machines;repeatfor source machine s ∈ { 1 , 2 , . .
. , m } do
Compute ∇ Θ L s and
EXPERIMENTS
The baselines and experimental objectives.
We ﬁrst examine the performance of FedCI.
We then compare the performance of FedCI against recent causal inference methods, such as BART [Hill, 2011], TARNet,
CFR Wass (CFRNet with Wasserstein distance), CFR MMD (CFRNet with maximum mean discrepancy distance)
[Shalit et al., 2017],
CEVAE [Louizos et al., 2017], OrthoRF [Oprescu et al., 2019], X-learner
[Künzel et al., 2019], and R-learner [Nie and Wager, 2021].
All these methods do not consider learning causal effects in a federated setting.
This analysis aims to show the efﬁcacy of FedCI as compared with the baselines trained in three different cases: ( 1 ) training a local model on each source data, ( 2 ) training a global model with the combined data of all sources, ( 3 ) using bootstrap aggregating (also known as bagging, which is an ensemble learning method) of Breiman [1996] where m models are locally trained on each source data; then taking average of the predicted treatment effects of each model.
Although case ( 2 ) violates the privacy constraint of federated data, we use it for comparison purposes.
In general, we would like to assess the performance of the federated causal inference approach against the baselines using combined data in case ( 2 ).
We use publicly available libraries and source codes to implement the baseline methods.
In particular, CEVAE,
TARNet, CFR Wass, and CFR MMD are readily available on github.
We use the online packages BartPy for
BART, causalml [Chen et al., 2020b] for X-learner and R-learner, and econml [Microsoft Research, 2019] for OrthoRF.
For all the methods, we ﬁne-tune the learning rate in { 10 − 1 , 10 − 2 , 10 − 3 , 10 − 4 } and regularizers in { 10 1 , 10 0 , 10 − 1 , 10 − 2 , 10 − 3 } .
Evaluation metrics.
We report two evaluation metrics: (i) precision in estimation of heterogeneous effects (PEHE)
[Hill, 2011] for evaluating ITE, and (ii) absolute error for evaluating ATE.
Details are presented in Appendix E. These metrics are for point estimates, which are the mean of ITE and ATE in their estimated distributions.
We also report the estimated distribution of ATE in our model.
4.1 SYNTHETIC DATA
We analyses FedCI in terms of three types of outcomes: (1) real-value, (2) binary, and (3) count.
While (1) is examined in a well-speciﬁed case for the outcomes, (2) and (3) are studied in misspeciﬁed cases.
4.1.1 Real-value Outcomes
Data.
Obtaining ground truth for evaluating causal inference algorithm is challenging.
Thus, most methods are evaluated using synthetic or semi-synthetic datasets.
In this experiment, we simulate the data with the following distributions:
[Equation]
where ϕ ( · ) is the sigmoid function, λ ( · ) is the softplus function, and x i =
[ x i 1 ,..., x id x ] (cid:62) ∈ R d x with d x = 20 .
We simulate two synthetic datasets: DATA-1 and DATA-2.
For
DATA-1, the ground truth parameters are randomly set as follows: σ 0 = σ 1 = 1 , ( a 0 , b 0 , c 0 ) =
, a 1 ∼
N ( 0 , 2 · I d x ) , b 1 ∼ N ( 0 , 2 · I d x ) , c 1 ∼ N ( 1 , 2 · I d x ) .
For
DATA-2, we set ( b 0 , c 0 ) =
(6 , 30) , b 1 ∼ N (10 · 1 , 2 · I d x ) , c
1 ∼ N (15 · 1 , 2 · I d x ) , and the other parameters are set similar to that of DATA-1.
The purpose is to make two different scales of the outcomes for the two datasets.
For each dataset, we simulate 10 replications with n = 5000 records.
We only keep { ( y i , w i , x i ) } ni =1 as the observed data, where y i =
y i (0)
if w i = 0 and y i
=
y i (1) if w i
= 1 .
We divide the data into ﬁve sources, each consists of n s =records.
In each source, we use 50 records for training,for testing and 400 for validation.
We report the evaluation metrics and their standard errors over the 10 replications.
The parameters chosen for this simulation study satisfy Assumption 1 since y i (0) and y i (1) are independent of w i given x i .
Assumption 2 is respected as the treatment on an individual i does not effect the outcome of another individual j ( i (cid:54) = j ).
Since we ﬁxed the dimension of x i and draw it from the same distribution, Assumption 3 is implicitly satisﬁed.
Assumption 4 holds true since each record drawn from the above distributions is attributed to one individual.
This means that there are no duplicates of individuals in more than one source.
Assumption 5 is also satisﬁed since we have divided the data equally from one dataset.
FedCI vs. training on combined data.
Figure 1 reports the three evaluation metrics of FedCI compared with two data source settings: training on combined data and training locally on each data source.
As expected, the ﬁgures show that the errors of FedCI are as low as those of training on
[Table]
PE
H E
The error of ITE, √ (cid:15) PEHE
[Table] [Table]
[Table] [Table]
(cid:15)
A T E
The error of ATE, (cid:15) ATE
[Table]
[Figure] [Figure]
[Table]
the combined data.
This result veriﬁes the efﬁcacy of the proposed federated algorithm
[Table]
Inter-dependency component analysis.
We study the impact of the inter-dependency component (see Section 3.4)
[Table]
by removing it from the model.
Figure 2 presents the errors of FedCI compared with ‘no inter-dependency’ (FedCI without inter-dependency).
The ﬁgures show that the errors in predicting ITE and ATE of ‘no inter-dependency’ seem to be higher than those of FedCI.
This result showcases the importance of our proposed inter-dependency component.
In Figure 1, the error (cid:15) ATE of FedCI increases as the number of sources increases from 1 to 2.
In Figure 2, (cid:15) ATE of FedCI is larger than that of without inter-dependency.
These results might be due to the non-convex optimisation which could lead to a local minima.
A potential direction to improve is to use a minibatch stochastic gradient descent for GPs [Chen et al., 2020a].
Contrasting with existing baselines.
In this experiment, we compare FedCI with the existing causal inference methods.
All these baseline methods do not consider estimating causal effects on multiple sources.
Thus, we train them in three cases as explained earlier: (1) train locally ( loc ), (2) train with combined data ( com ), and (3) train with bootstrap aggregating ( agg ).
Note that case ( 2 ) violates constraint that data are stored at their local sites.
We expect that the error of FedCI to be close to case (2) of the baselines.
Table 1 and 2 report the performance of each method in estimating ATE and ITE.
Regardless of different scales on the two synthetic datasets, the ﬁgure shows that FedCI achieves competitive results as compared with all the baselines.
FedCI is in the top-3 performances among all the methods.
Importantly, FedCI obtains lower errors than those of BART com , X-Learner com , R-Learner com ,
OthoRF com , TARNet com , CFR Wass com , and CFR MMD com , which were trained on combined data and thus violate constraint of federated data setting.
Compared with CEVAE com ,
FedCI is better than this method in predicting ITE and comparable with this method in predicting ATE (slightly higher errors).
However, we emphasize again that this result is expected since FedCI is a federated learning algorithm
while
CEVAE com works directly on combined data.
The estimated distribution of ATE.
To analyse uncertainty, we present in Figure 3 the estimated distribution of ATE in the ﬁrst source ( s = 1 ).
The ﬁgures show that the true 1 source 2 sources 3 sources
A T E
E [ τ ]
[Figure]
ATE 1 source 2 sources 3 sources
A T E
[Figure]
ATE is covered by the estimated interval and the estimated mean ATE shifts towards its true value (dotted lines) when more data sources are used.
This result might provide useful information about the application in practice
4.1.2 Misspeciﬁcation Analysis: Binary and Count Outcomes
Data.
In this experiment, we analyse the performance when the model is misspeciﬁed.
We compare FedCI with the baselines in two cases: binary outcomes and count outcomes.
We reuse the ground truth distributions of x ij and w i as in
Section 4.1.1.
For the outcomes, we simulate them with the following distributions:
Binary outcomes: y i (0) ∼ Bern ( ϕ ( b 0 + x (cid:62) i b 1 )) , y i (1) ∼ Bern ( ϕ ( c 0 + x (cid:62) i c 1 )) .
Count outcomes: y i (0) ∼ Poisson (exp( b
+ x (cid:62) i b 1 )) , y i (1) ∼ Poisson (exp( c
+ x (cid:62)
i c 1 )) .
Results and discussion.
From Table 3 and 4, FedCI gives competitive results compared with the baselines trained on combined data.
The reason for the good performance for
FedCI and some baselines in these misspeciﬁcation cases is because they provide good estimates for the mean of the missing outcomes.
This might in turn be due to the mean estimation of Gaussian distribution in FedCI coincides with the mean estimation of the other distributions.
Nevertheless, since these are misspeciﬁed cases, the continuous posterior distribution is not a good estimation.
To obtain better posterior distributions of the missing outcomes and the causal estimands, we would need to consider some other appropriate distributions in our model.
4.2 IHDP DATA
Data.
The Infant Health and Development Program (IHDP)
[Hill, 2011] is a dataset with 747 data points, each has 25 covariates.
These data are obtained from a randomized study on the impact of specialist visits to children’s cognitive
[Table] [Table]
development.
Herein, specialist visit is the treatment and children’s cognitive development is the outcome.
We use the NPCI package [Dorie, 2016] to simulate two potential outcomes for the treatment (with or without specialist visit) of each child.
Hence, the true individual treatment effect can be computed for evaluation purposes.
There are 10 replicates of the dataset
[Table]
and each of them is divided into three sources of size 249.
For each source, we then split it into three equal sets for the purpose of training, testing
[Table]
and validating the models.
The mean and standard error of the aforementioned evaluation metrics are reported over the above 10 replicates of the data.
Results and discussion.
Similar to the experiment for synthetic datasets, here we also train the baselines in three cases as explained earlier.
We also expect that the errors of
FedCI to be close to the baselines trained with combined data ( com ).
The results reported in Table 5 show that FedCI achieves competitive results compared to the baselines (we skipped the ﬁrst and second cases ( loc and agg ), please see Appendix F.2 for the full table).
Indeed, FedCI is in the top-3 performances among all the methods.
This result again veriﬁes that FedCI can be used to estimate causal effects effectively under privacy-perserving, federated data settings.
The estimated distribution of ATE is presented in
Appendix F.2 due to limited space
[Table]
CONCLUSION
We have introduced FedCI, a Bayesian causal inference paradigm via a reformulation of multi-output GPs to learn
Please see the table in
[Table]
causal effects, while keeping data at their local sites.
An inference method involving the decomposition of ELBO is presented, allowing the model to be trained in a federated setting.
This work is an important step towards a privacy-preserving causal learning model.
One interesting future research direction is to combine FedCI with differential privacy to give a stronger privacy guarantee.
This new direction would require adding an appropriate noise component, such as
Laplace noise, to the parameters while training the model.
Our future research would also involve further exploration on combining FedCI with differential privacy for Gaussian processes [Smith et al., 2018] and multiparty differential privacy algorithms [Pathak et al., 2010, Rajkumar and Agarwal, 2012, Hamm et al., 2016].
The inherent use of GPs in our approach would incur computational time of inverse covariance matrix in each source of cubic time complexity.
Another possible future research direction is to reformulate this in terms of sparse Gaussian process models [Hensman et al., 2013, Hoang et al., 2017, 2020].
Acknowledgements
This research/project is supported by the National Research Foundation Singapore and DSO National Laboratories under the AI Singapore Programme (AISG Award
No: AISG2-RP2020-016).