An Adaptive Kernel Approach to Federated Learning of Heterogeneous Causal Effects
Abstract
We propose a new causal inference framework to learn causal effects from multiple, decentralized data sources in a federated setting.
We introduce an adaptive transfer algorithm that learns the similarities among the data sources by utilizing Random
Fourier Features to disentangle the loss function into multiple components, each of which is associated with a data source.
The data sources may have different distributions; the causal effects are independently and systematically incorporated.
The proposed method estimates the similarities among the sources through transfer coefficients, and hence requiring no prior information about the similarity measures.
The heterogeneous causal effects can be estimated with no sharing of the raw training data among the sources, thus minimizing the risk of privacy leak.
We also provide minimax lower bounds to assess the quality of the parameters learned from the disparate sources.
The proposed method is empirically shown to outperform the baselines on decentralized data sources with dissimilar distributions.
Introduction
Many important questions posed in the natural and social sciences are causal in nature:
What are the long-term effects of mild Covid-19 infection on lung and brain functions?
How is mortality rate influenced by the daily air pollution?
How would a welfare policy affect employment rate of a minority group?
Causal inference has been applied in a wide range of domains, including economics (Finkelstein and Hendren 2020), medicine (Henderson et al. 2016; Powers et al.
and social welfare (Gutman et al. 2017).
The large amount of experimental and/or observation data needed to accurately estimate the causal effects often resides across different sites.
In most cases, the data sources cannot be combined to support centralized processing due to some inherent organizational or policy constraints.
For example, in many countries, medical or health records of cancer patients are kept strictly confidential at local hospitals; direct exchange or sharing of the records among hospitals, especially for research purposes, are not allowed (Gostin et al. 2009).
The main research question is:
How to securely access these diverse data sources to build an effective global causal effect estimator, while balancing the risk of breaching data privacy and confidentiality?
Current causal inference approaches (e.g., Shalit et al.
; Yao et al. 2018) require the shared data to be put in one place for processing.
Current federated learning algorithms (e.g., Sattler et al.
Wang et al. 2020) allow collaborative learning of joint models based on non-independent and identically distributed (non iid) data; they cannot, however, directly support causal inference as the different data sources might have disimilar distributions that would lead to biased causal effect estimation.
For example, the demographic profile and average age for cancer patients from two different hospitals may be drastically different.
If the two data sets are combined to support causal inference, one distribution may dominate over the other, leading to biased causal effect estimation.
We introduce a new approach to federated causal inference from multiple, decentralized, and disimilarly distributed data sources.
Our contributions are summarized as follows:
• We propose a new federated causal inference algorithm, called CausalRFF 1 , based on the structural causal model (SCM) (Pearl 2009a), leveraging the Random Fourier Features (Rahimi and Recht 2007) for federated estimation of causal effects.
The Random Fourier Features allow the objective function to be divided into multiple components to support federated training of the model.
• We perform federated causal inference with CausalRFF from data sources with different distributions through the adaptive kernel functions ; the inference is carried out without sharing raw data among the sources, hence minimizing the risk of privacy leak.
• We provide the minimax lower bounds to explicate the limits of estimation and optimization procedures in our federated causal inference framework.
Our work is an important step toward privacy-preserving causal inference.
We explore the possibility of combining CausalRFF with multiparty differential privacy at the end of the paper.
Related Work
Little work has been done on combining causal inference with federated learning in a privacy preserving manner.
On causal inference: The authors Hill (2011); Alaa and van der Schaar (2017, 2018); Shalit et al.
(2017); Yoon et al.
; Yao et al.
; Künzel et al.
; Nie and Wager (2020) proposed learning causal effects directly from local data sources; these methods adopt the standard ignorability assumption (Rosenbaum and Rubin 1983).
Louizos et al.
; Madras et al.
(2019) adapted the structural causal model (SCM) of Pearl (1995) to estimate the causal effects with the existence of latent confounding variables.
Our work is closely related to and extends the notion of transportability , where Pearl and Bareinboim (2011); Bareinboim and Pearl (2016); Lee et al.
and
related work formulated and provided theoretical analysis of intervention tools on one population to compute causal effects on another population.
Lee et al.
(2020) generalized transportability to support identification of causal effects in the target domain from the observational and interventional distributions on subsets of observable variables, forming a foundation for drawing conclusions for observational and experimental data (Tsamardinos et al. 2012; Bareinboim and Pearl 2016).
Causal inference from multiple, decentralized, dissimilarly distributed sources that cannot be combined or processed in a central site is not addressed.
Recently, Aglietti et al.
(2020), conducted randomized experiments on the source to collect data and then estimated a joint model of the interventional data from source population and the observational data from target population.
Our work is different in that we do not work with randomized data; we estimate causal effects through transfers using only observational data .
This corresponds to an important setting in real-life, where only retrospective observational data are available, e.g., Covidrelated case and intervention records, bank and financial transaction records.
On federated learning: Federated learning enables collaboratively learning a shared prediction model while keeping all the training data decentralized at source (McMahan et al. 2017).
Some federated learning approaches combine federated stochastic gradient descent (Shokri and Shmatikov 2015) and federated averaging (McMahan et al.
2017) to address regression problems Álvarez et al.
; Zhe et al.
; de Wolff et al.
Joukov and Kuli´c (2020) and Hard et al.
Zhao et al.
(2018); Sattler et al.
; Mohri et al.
Recent federated learning algorithms allow collaborative learning of joint deep neural network models based on non-iid data (Sattler et al.
Wang et al. 2020).
All these algorithms, however, do not directly support causal inference as the different data sources might have dissimilar distributions that would lead to biased causal effect estimation.
Little work has been focused on federated estimation of causal effects.
Vo et al.
(2022) proposed a Bayesian approach that estimates posterior distributions of causal effects based on Gaussian processes, which does not allow dissimilar distributions of the sources.
Xiong et al.
(2021) estimated average treatment effect (ATE) and average treatment effect on the treated (ATT) and assumed that the confounders are observed .
Our work, on the other hand, estimate conditional average treatment effect (CATE) (which is also known as individual treatment effect, ITE) and average treatment effect (ATE) under the existence of latent confounders .
We utilize Random Fourier
Features to build an integrative framework of causal inference in a federated setting that allows for dissimilar data distributions .
The Proposed Model
In this section, we first detail the problem formalization.
We then present the causal effects of interest and the scheme to estimate them.
Lastly, we describe the assumptions and the structural equations.
3.1 Problem Description
Problem setting & notations.
Suppose we have m sources of data, each denoted by D s = { ( w s i ,
y s i , x s i ) } n s i =1 , where s ∈ S :
= { s 1 , s 2 , . .
, s m } , and the quantities w
s i , y s i and x s i are the treatment assignments, observed outcome associated with the treatment, and covariates of individual i in source s , respectively.
These data sources D s are located in different locations and their distributions might be completely different.
All the sources share the same causal graph as shown in Figure 1, but the data distributions may be different, e.g., p s 1 ( x , w, y ) ̸
= p s 2 ( x , w, y ) , where p s 1 ( · ) and p s 2 ( · ) denote the two distributions on two sources s 1 and s 2 , respectively.
Similarly, the marginal and the conditional distributions with respect to these variables can also be different (or similar).
The objective is to develop a global causal inference model that satisfies both of the following two conditions: (i) the causal inference model can be trained in a private setting where the data of each source are not shared to an outsider, and (ii) the causal inference model can incorporate data from multiple sources to improve causal effects estimation in each specific source.
Causal effects of interest.
Given a causal model trained under the aforementioned setting, we are interested in estimating the conditional average treatment effect (CATE) 2 and average treatment effect (ATE).
Let Y , W , X
be random variables denoting the outcome, treatment, and proxy variable, respectively.
Then, the CATE and ATE are defined as follows (Louizos et al. 2017; Madras et al.
[Equation]
Y | do ( W =0) ,
X = x (cid:3) , τ :=
E [ τ ( X )] , (1) where do ( W = w ) represents that a treatment w ∈ { 0 , 1 } is given to the individual.
This definition is followed from Louizos et al.
; Madras et al.
Given a set of n new individuals whose covariates/observed proxy variables are { x i } ni =1 , the CATE and ATE in this sub-population are obtained by τ ( x i ) and τ = (cid:80) ni
=1 τ ( x i ) /n .
[Figure] [Figure]
The central task to estimate CATE and ATE is to find E [ Y | do ( W = w ) , X = x ] .
Since the data distribution of each source might be different from (or similar to) each other, we use the notation
E [ Y | do ( W = w s , X
= x s ] to denote the expectation of the outcome Y under an intervention on W of an individual in source s .
With the existence of the latent confounder Z , we can further expand this quantity using do -calculus
(Pearl 1995).
In particular, from the backdoor adjustment formula, we have
[Equation]
Eq.
(2) shows that the causal effect is identifiable if we can find the conditional distributions p s ( y s | w s , z s ) and p s ( z s | x s ) for each source s .
The second distribution can be further expanded by p s ( z s | x s ) =
(cid:80) w s (cid:82) p s ( z | x s , y s i , w s ) p s ( y s | x s , w s )
p s ( w s | x s )
d y s .
Following the forward sampling strategy, the remaining is to find the following distributions
[Equation]
called and then systematically draw samples from these estimated distributions to obtain the empirical expectation of Y given do ( W = w s ) and X = x s .
Identification.
The CATE and ATE are identifiable if we are able to learn the distributions in Eq.
(3), which involve latent confounder Z .
Louizos et al.
(2017) showed that this is possible if Z has a relationship to the observed variables X , and there are many cases that it is identifiable such as:
Z is categorical and X is a Gaussian mixture model (Anandkumar et al. 2014), X includes three independent views of Z (Goodman 1974; Allman et al. 2009; Anandkumar et al. 2012), Z is a multivariate binary and X are noisy functions of Z (Jernite et al.
2013; Arora et al. 2017)
, to name a few.
Following the works by Louizos et al.
; Madras et al.
(2019), we use variational inference in the spirit of the variational auto-encoder (VAE) to recover the latent confounders, since it can learn a rich class of latent-variable models, and thus recovering the causal effects.
Identification of our work follows closely from the literature, however our main contribution is in the federated setting of the model.
Please refer to Appendix for the proof of identifiability.
3.2 The Causal Graph and Assumptions
Since our method adopts the SCM approach with the causal graph in Figure 1, there are some implicit assumptions that follow from the axioms and properties of SCM: (A1) Consistency : W = w = ⇒
Y ( w ) =
Y , this follows from the axioms of SCM.
(A2) No interference : the treatment on one subject does not affect the outcomes of another one.
This is because the outcome has only a single treatment node as its parent.
(A3) Positivity : every subject has some positive probability to be assigned to every treatment.
These assumptions are standard in any causal inference algorithm.
One can find further discussion in Pearl (2009a,b); Morgan and Winship (2015).
For our proposed federated setting, we make two additional assumptions as follows:
(A4)
The individuals in all sources have the same set of common covariates
(A5)
Any individual does not exist in more than one source
Assumption (A4) has been implicitly shown in our setup since all the sources would share the same causal graph.
This is a reasonable assumption as we intend to build a unified model on all of the data sources, e.g., decentralized data in Choudhury et al.
; Vaid et al.
Flores et al.
(2020) satisfy this assumption for federated learning.
Assumption (A5) is to ensure that no individuals would dominate the other individuals when training the model.
For example, if an individual appears in all of the sources, the trained model would be biased by data of this individual (there is imbalance caused by the use of more data from this particular individual than the others).
Hence, this condition would ensure that such bias does not exist.
In practice, Assumption (A5) sometimes does not hold.
To address such a problem, we perform a pre-training step to exclude such duplicated individuals.
This step would use a one-way hash function to perform a secured matching procedure that identifies duplicated individuals.
Details of the pre-training step are presented in Appendix.
3.3 The Structural Equations
This section presents how the causal relations are modeled.
Since Z is the root node in the causal graph, we model it as a multivariate normal distribution: Z ∼ N ( µ , σ 2 z I d z ) for the all sources.
We now detail the structural equations of Y , W and X .
Let V be a univariate variable that represents a node or a dimension of a node in the causal graph (Figure 1), i.e., V can be Y , W or a dimension of
X .
Let pa ( V ) be set of V ’s parent variables in the causal graph, i.e, the nodes with directed edges to
V .
We model the structural equation of V in two cases as follows: if
[Equation]
is continuous:
[Equation]
where ϵ v ∼ N (0 , σ 2 v ) for the former case and ϵ v ∼ U [0 , 1] for the latter case, φ ( · ) is the logistic function and 1 ( · ) is the indicator function.
The latter case implies that V given pa ( V ) follows
Bernoulli distribution with p ( V = 1 | pa ( V )) =
φ ( f v ( pa ( V ))) .
Furthermore, if W ∈ pa ( V ) , then we further model
[Equation]
Example.
If Y ∈ R , W ∈ { 0 , 1 } and X k ∈ R ( X k is the k –th dimension of X ), then the structural equations are as follows:
[Equation]
In the subsequent sections, we present how to learn the functions f v ( v ∈ { y 0 , y 1 , w, x } ) in a federated setting and then use them to estimate the causal effects of interest.
CausalRFF: An Adaptive Federated Inference Algorithm
This section presents a new federated algorithm to learn the distributions in Eq.
The central task is to decompose the objective function into multiple components, each associated with a source.
4.1 Learning Distributions Involving Latent Confounder
To estimate causal effects, we need to estimate the four quantities detailed in Eq.
This section presents how to learn p s ( z s | x s , y s , w s ) and p s ( y s | w s , z s ) .
Since the marginal likelihood has no analytical form, we learn the above distributions using variational inference which maximizes the evidence lower bound (ELBO)
[Equation]
where q ( z s ) =
N ( z s ; f q ( y s , w s , x s ) , σ 2 q I ) is the variational posterior distribution.
The function f q ( · ) is modeled as follows:
f q ( y s , w s , x s ) =
(1 − w s )
f q 0 ( y s , x s ) + w s f q 1 ( y s , x s ) , where f q 0 and f qare two functions to be learned.
The density functions p s ( y s | w s , z s ) , p s ( w s | z s ) and p s ( x s | z s ) are obtained from the structural equations as described in Section 3.3.
Please refer to Appendix for details on derivation of the ELBO.
Adaptive modeling.
Since the observed data from each source might come from different (or similar) distributions, we would model them separately and adaptively learn their similarities.
In particular, we propose a kernel-based approach to learn these distributions.
To proceed, we first obtain the empirical loss function (cid:98) L from negative of the ELBO L by generating M samples of each latent confounder
Z using the reparameterization trick (Kingma and Welling 2013): z s i
[ l
[Equation]
+
[Equation] [Equation]
] , where ϵ s i [ l ] is drawn from the standard normal distribution.
We obtain a complete dataset
[Equation]
Using this complete dataset, we minimize the following objective function
[Equation]
with respect to f c , where A = { y 0 , y 1 , w, x, q 0 , q 1 } , and R ( · ) denotes a regularizer.
The minimizer of J would result in the following form of f c
[Equation]
where u v j is obtained from the j –th tuple of the dataset ˜ D v .
Details are presented in Appendix.
Since data from the sources might come from a completely different (or similar) distribution, we would use an adaptive kernel to measure their similarity.
In particular, let k ( u s , u v ) be typical kernel function such as squared exponential kernel, rational quadratic kernel, or Matérn kernel.
The kernel used in
Eq.
(9) is as follows: κ ( u s , u v )
=
λ s , v k ( u s , u v ) , if s ̸ = v ; otherwise, κ ( u s , u v ) =
k ( u s , u v ) , where λ s , v ∈ [0 , 1] is the adaptive factor and it is learned from the observed data.
Remark.
Eq.
(9) indicates that computing f c ( u s ) requires collecting all data points from all sources, and so the objective function in Eq.
(8) cannot be optimized in a federated setting.
Next, we present a method known as Random Fourier Features to address the problem.
Random Fourier Features.
We show how to adapt Random Fourier Features (Rahimi and Recht 2007) into our model.
Let k ( u , u ′ ) be any translation-invariant kernel (e.g., squared exponential kernel, rational quadratic kernel, or Matérn kernel).
Then, by Bochner’s theorem (Wendland 2004,
Theorem 6.6), it can be written in the following form:
[Equation]
where s ( ω ) is a spectral density function associated with the kernel (please refer to Appendix for spectral density of some popular kernels).
The last equality follows from the fact that the kernel function is real-valued and symmetric.
This type of kernel can be approximated by
[Equation]
where ϕ ( u ) =
B − 12 [cos( ω ⊤ 1 u ) ,..., cos( ω ⊤ B u ) , sin( ω ⊤ 1 u ) ,..., sin( ω ⊤ B u )] ⊤ .
The last equality follows from the trigonometric identity: cos( u − v ) = cos u cos v + sin u sin v .
Substituting the above random Fourier Features into Eq.
(9), we obtain
[Equation]
where θ s c = (cid:80) n s i =1 ϕ ( u s )
α s i and λ s , v ( s , v ∈ S ).
While optimizing the objective function J , instead of learning α s i , we can directly consider θ s as parameter to be optimized.
This has been used in several works such as Rahimi and Recht (2007); Chaudhuri et al.
(2011); Rajkumar and Agarwal (2012).
This approximation allows us to rewrite the objective function J as a summation of local objective functions in each source:
[Equation]
where J
[Equation]
=
[Equation]
where ζ ∈ R
+ is a regularizer factor.
Each component J ( s ) is associated with the source s and it can be computed with the local data in this source.
Hence, it enables federated optimization for the objective function J .
Figure 2 illustrates our proposed federated causal learning algorithm with three sources, where θ denotes the set of all parameters to be learned including θ s and λ s , v from all the sources.
The federated learning algorithm can be summarized as follows:
First, each source computes the local gradient, ∇ θ J ( s ) , using its own data and sends to the server.
The server, then, collects these gradients from all sources and subsequently updates the model.
Next, the server broadcasts the new model to all the sources.
Minimax lower bound.
We now compute the minimax lower bound of the proposed model, which gives the rate at which our estimator can converge to the population quantity of interest as the sample size increases.
We first state the following result that concerns the last two terms in Eq.
Lemma 1 (With presence of latent variables) .
Let θ = { θ s c : c ∈ { y 0 , y 1 , x, w } , s ∈ S } and ˆ θ be its estimate.
Let y s i ∈ R and x s
i ∈ R
d x .
Let S \ s = S \ { s } .
Then,
[Equation]
The LHS of Eq.
(14) can be seen as the worst case of the best estimator, whereas the RHS depicts the behavior of the convergence.
The bounds do not only depend on the number of samples ( n s , training size) of each source but also the adaptive factors λ s , v .
When the adaptive factors are small, the lower bounds are large since data from a source s are only used to learn its own parameter θ s .
When the adaptive factors are large, the lower bounds are smaller, which suggests that data from a source would help infer parameters associated with the other sources.
This bound gives a guarantee on how data from all the sources impact the learned parameters that modulate the two distributions p s ( z s | x s , y s , w s ) and p s ( y s | w s , z s ) .
The proof of Lemma 1 can be found in Appendix.
4.2 Learning Auxiliary Distributions
The previous section has shown how to learn p s ( z s | x s , y s , w s ) and p s ( y s | w s , z s ) .
To compute treatment effects, we need to learn two more conditional distributions, namely p s ( w s | x s ) and p s ( y s | x s , w s ) .
Since all the variables in these two distributions are observed, we estimate them using maximum likelihood estimation.
In the following, we present a federated setting to learn p s ( w s | x s ) .
Similar to the previous section, the objective function here can also be decomposed into m components as follows:
J w ≃ (cid:80) s ∈ S J ( s ) w , where J ( s ) w = (cid:80) n s i =1 ℓ ( w s i , φ ( g ( x s i )))
+ m − 1 (cid:80) v ∈ S ζ w ∥ ψ v ∥ 22 and g ( x s i ) =
(cid:80) v ∈ S ϕ ( x s i ) ⊤ ( ψ s + γ s , v ψ v ) , γ s , v ∈ [0 , 1] is the adaptive factor, ψ s is the parameter associated with source s , and ℓ ( · ) denotes the cross-entropy loss function since w s i is a binary value.
The first component of J ( s ) w is obtained from the negative log-likelihood.
Learning of p s ( y s | x s , w s ) is similar.
For convenience, in the subsequent analyses, we denote the parameters and adaptive factors of this distribution as β s and η s , v , where s , v ∈ S and s ̸ = v .
The next lemma shows the minimax lower bound for the first two sets of parameters ψ and β in Eq.
(3), but this time without involving the latent variables:
Lemma 2 (Without the presence of latent variables) .
Let ψ = { ψ s } m s =1 , β = { β s }
m s =
and ˆ ψ , ˆ β be their estimates, respectively.
Let y s i ∈ R .
Then,
[Equation]
The proof of Lemma 2 can be found in Appendix.
The bounds presented in Lemma 1 and 2 give helpful information about the number of samples to be observed and the cooperation of multiple sources of data through the transfer factors.
Since we used variational inference and maximum likelihood to learn the parameters in our model, these methods give consistent estimation as shown in
Kiefer and Wolfowitz (1956); Van der Vaart (2000); Wang and Blei (2019); Yang et al.
4.3 Computing Causal Effects
The key to estimate causal effects in our model is to compute the outcome in Eq.
We proceed by drawing samples from the distributions in Eq.
Generating samples from the conditional distributions p s ( w s | x s ) , p s ( y s | x s , w s ) , and p s ( y s | w s , z s ) is straightforward since they are readily available as shown in either Section 4.1 or 4.2.
There are two options to draw samples from the posterior distribution of confounder p s ( z s | x s , y s , w s ) .
The first one is to draw from its approximation, q ( z s ) , since maximizing the ELBO in Section 4.1 is equivalent to minimizing
KL ( q ( z s ) ∥
p s ( z s | x s , y s , w s )) .
As a second option, we note that the exact posterior of confounder can be rewritten as p s ( z s | x s , y s , w s )
∝ p s ( y s | z s , w s ) p s ( w s |
z s ) p s ( x s | z s ) p
z s ) , whose components on the right hand side are also available in Section 4.1.
Thus, we can draw from this distribution using the
Metropolis-Hastings (MH) algorithm.
Since Z is a multidimensional random variable, the traditional
MH algorithm would require a long chain to converge.
We overcome this problem by using the
MH with independent sampler (Liu 1996) where the proposal distribution is the variational posterior distribution q ( z s ) learned in Section 4.1.
The second approach would give more accurate samples since we select the samples based on exact acceptance probability of the posterior p s ( z s | x s , y s , w s ) .
This would help estimate the CATE given x s i .
The local ATE is the average of CATE of individuals in a source s .
These quantities can be estimated in a local source machine.
To compute a global ATE, the server would collect all the local ATE in each source and then compute their weighted average.
Further details are in Appendix.
Experiments
The baselines.
In this section, we first carry out the experiments to examine the performance of
CausalRFF against standard baselines such as BART (Hill 2011), TARNet (Shalit et al. 2017), CFRwass (CFRNet with Wasserstein distance) (Shalit et al. 2017),
CFR-mmd (CFRNet with maximum mean discrepancy distance) (Shalit et al. 2017), CEVAE (Louizos et al. 2017), OrthoRF (Oprescu et al.
, X-learner (Künzel et al. 2019), R-learner (Nie and Wager 2020), and FedCI (Vo et al. 2022).
In contrast to CausalRFF, these methods (except FedCI) do not consider causal inference within a federated setting.
We compare our method to these baselines trained in two ways: ( a ) training a global model with the combined data from all the sources, ( b ) using bootstrap aggregating of Breiman (1996) where m models are trained separately on each source data and then averaging the predicted treatment effects based on each trained model.
Note that case ( a ) violates federated data setting and is only used for comparison purposes.
In general, we expect that the performance of CausalRFF to be close to that of the performance of the baselines in case ( a ) when the data distribution of all the sources are the same.
In addition, we also show that the performance of CausalRFF is better than that of the baselines in case ( a ) when the data distribution of all the source are different.
Implementation of the baselines.
The implementation of CEVAE is from Louizos et al.
Implementation of TARNet, CFR-wass, and CFR-mmd are from Shalit et al.
For these
[Figure]
1 .1 .1 .1 .√
(cid:15)
PE H E
The error of CATE    5 Numberofsources, m 0 .0 .0 .0 .1 .(cid:15)
A T E
The error of ATE
[Figure]
[Figure] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
methods, we use Exponential Linear Unit (ELU) activation function and fine-tune the number of nodes in each hidden later from 10 to 200 with step size of addition by 10.
For BART, we use package
BartPy , which is readily available.
For X-learner and R-learner, we use the package causalml (Chen et al. 2020).
For OrthoRF, we use the package econml (Microsoft Research 2019).
For FedCI, we use the code from Vo et al.
For all methods, the learning rate is fine-tuned from 10 − 4 to 10 − 1 with step size of multiplication by 10 .
Similarly, the regularizer factors are also fine-tuned from 10 − 4 to 10 0 with step size of multiplication by 10 .
We report two error metrics: ϵ PEHE (precision in estimation of heterogeneous effects) and ϵ ATE (absolute error) to compare the methods.
We report the mean and standard error over 10 replicates of the data.
Further details are presented in Appendix.
5.1 Synthetic Data
Data description.
Obtaining ground truth for evaluating causal inference algorithm is a challenging task.
Thus, most of the state-of-the-art methods are evaluated using synthetic or semi-synthetic datasets.
In this experiment, the synthetic data is simulated with the following distributions:
[Equation]
where Cat ( · ) , N ( · ) , and Bern ( · ) denote the categorical distribution, normal distribution, and Bernoulli distribution, respectively.
φ ( · ) denotes the sigmoid function, sp ( · ) denotes the softplus function, and x i =
[ x i 1 ,..., x id x ] ⊤ ∈ R d x with d x = 30 .
Herein, we convert z s i to a one-hot vector.
To simulate data, we randomly set the ground truth parameters as follows:
ρ = [ . 11 , .
12] ⊤ , ( c 0 , d 0 ) =
(0 . 9 , 7 . 9) , ( c 1 , d 1 , d 1 ) are drawn i.i.d from N ( 0 , 2 I 5 ) , a j 0 and elements of a j 1 are drawn i.i.d from N (0 , 2) .
For each source, we simulate 10 replications with n s = 1000 records.
We only keep { ( y s i , w s i , x s i ) } n s i =1 as the observed data, where y s i =
y s i (0)
if w s i =
and y s i
=
y s i (1) if w s i = 1 .
In each source, we use 50 data points for training, 450 for testing and 400 for validating.
We report the evaluation metrics and their standard errors over the 10 replications.
Result and discussion (I).
In the first experiment, we study the performance of CausalRFF on multiple sources whose data distributions are the same.
To do that, we simulate m = 5 sources from the same distribution, i.e., we set the ground truth ∆ = 0 .
for all the sources.
We refer to this dataset as DATA same .
In this experiment, we expect that the result of CausalRFF, which is trained in federated setting, is as good as training on combined data.
The results in Figure 3 show that the error in two cases seem to move together in a correlated fashion, which verifies our hypothesis.
In addition, to study the performance of CausalRFF on the sources whose data distributions are different, we also simulate m = 5 sources.
However, the first source is with ∆ = 0 .
and the other
[Figure]
[Table]
four sources are with ∆ = 4 . 0 .
We refer to this dataset as DATA diff .
We test the error of CATE and
ATE on the first source.
In this case, we expect that the errors of CausalRFF to be lower than that of training on combined data since CausalRFF learns the adaptive factors which prevent negative impact of the other four sources to the first source.
The results in Figure 4 show that CausalRFF achieves lower errors compared to training on combined data (there are two cases of combining: stacking data, and adding one-hot vectors to indicate the source of each data point), which is as expected.
In the third experiment, we study the effect of ∆ on the performance of CausalRFF.
We simulate m = 2 sources with different values of ∆ .
In particular, the first source is with ∆ = 0 .
and the second source is with ∆ varying from 0.0 to 8.0.
We compare our CausalRFF method with that of training on combined data.
Again, Figure 5 shows that CausalRFF achieves lower errors as expected.
Result and discussion (II).
This section aims to compare CausalRFF with the baselines on both datasets: DATA same and DATA diff .
Except FedCI (which is a Bayesian federated method), the other baselines are trained on two cases: combined data (cb) and bootstrap aggregating (ag) as mentioned earlier.
On DATA same , we expect that the performance of the proposed method is as good as the baselines trained on combined data.
The results in Table 1 show that the performance of CausalRFF is as expected.
For DATA diff , we report the results on Table 2.
The figures reveal that the performance of CausalRFF is as good as the baselines in predicting ATE.
In terms of predicting CATE, the performance of the baselines significantly reduces as we add more data sources whose distribution are different from the first source.
Meanwhile, the performance of CausalRFF in predicting CATE is slightly reduced, but it is still much better than those of the baselines.
The reason of this is because we used adaptive factors to learn for the similarity of data distributions among the sources.
5.2 Large-scale Synthetic Data
Data description.
In this section, we conduct experiments on a large number of sources.
The set up in this section is similar to that of Section 5.1.
We simulate two cases: (1) DATA-LARGE same : a dataset of 100 sources, where we set ∆ = 0 for all sources so that their distributions are the same.
(2) DATA-LARGE diff : a dataset of 100 sources, where we draw uniformly the discrepancy factor ∆ ∼ U [0 , 8] for each source so that their distributions are different.
In both cases, we use test set from the first 20 sources for evaluation.
Result and discussion.
Table 3 shows that CausalRFF achieves competitive results in estimating ATE and CATE when the sources have the same distribution.
Table 4 shows that CausalRFF outperforms the baselines when the sources have different distributions.
These results are consistent with our discussions in Section 5.1.
5.3 A Real World Dataset
Data description.
The Infant Health and Development Program (IHDP)
(Hill 2011) is a randomized study on the impact of specialist visits (the treatment) on the cognitive development of children (the
[Table] [Table]
outcome).
The dataset consists of 747 records with 25 covariates describing properties of the children and their mothers.
The treatment group includes children who received specialist visits and control group includes children who did not receive.
This dataset was ‘de-randomized’ by removing from the treated set children with non-white mothers.
For each child, a treated and a control outcome are then simulated, thus allowing us to know the ‘true’ individual causal effects of the treatment.
Further details are presented in Appendix.
[Table]
Result and discussion.
Table 5 reports the experimental results on IHDP dataset.
Again, we see that the proposed method gives competitive results compared to the baselines.
In particular, the error of CausalRFF in predicting ATE is as low as that of the baselines, which is as we expected.
In addition, the errors of CausalRFF in predicting CATE are lower than those of the baselines, which verifies the efficacy of the proposed method.
Most importantly, CausalRFF is trained in a federated setting which minimizes the risk of privacy breach for the individuals stored in the local dataset.
Conclusion
We have proposed a new method to learn causal effects from federated, observational data sources with dissimilar distributions.
Our method utilizes Random Fourier Features that naturally induce the decomposition of the loss function to individual components.
Our method allows for each component data group to inherit different distributions, and requires no prior knowledge on data discrepancy among the sources.
We have also proved statistical guarantees which show how multiple data sources are effectively incorporated in our causal model.
Our work is an important step toward privacy-preserving causal inference.
Future work may include combining the proposed method with a multiparty differential privacy technique (e.g., Pathak et al. 2010; Rajkumar and Agarwal 2012;
Pettai and Laud 2015; Hamm et al. 2016), which might lead to a stronger privacy guarantee model.
Another direction is to extend the proposed method with some recent ideas (e.g., Khemakhem et al.
2020; Sun et al.
to study the identifiability of the model.
Acknowledgments and Disclosure of Funding
This research/project is supported by the National Research Foundation Singapore and DSO National
Laboratories under the AI Singapore Programme (AISG Award
No: AISG2-RP-2020-016).
AB was partially supported by an NRF Fellowship for AI (NRF-NRFAI11-2019-0002) and an
Amazon Research Award.
This work was conducted while YL was at Harvard University and the views expressed here do not necessarily reflect the position of Roche AG.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] Please see our experiments in Section 5.
(b) Did you describe the limitations of your work?
[Yes] Please see Section 6.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] Please refer to our Appendix. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Please see our experiments in Section 5 and supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Please see our experiments in Section 5.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] Please see our experiments in Section(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Please see our Appendix. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] Please see our experiments in Section 5.
(b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] Please see our supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]