Estimation of conditional average treatment effects on distributed confidential data

The Neyman–Rubin model (Imbens and Rubin, 2015) or potential outcomes framework is one of the main causal inference methods for estimating average treatment (intervention) effects and has been developed and applied in many studies since its establishment by Rubin. In recent years, there have been innovations that enable the estimation of individual or conditional average treatment effects (CATEs) by adapting not only statistical but also machine learning methods (e.g., Künzel et al. (2019); Athey et al. (2019)). Most methods for estimating CATEs assume that the data can be centralized in one place. However, if the distributed data contains confidential or private information, it is difficult to centralized them in one place.

Conversely, in the field of machine learning, a privacy-preserving analysis method called federated learning has been developed in recent years (Konečnỳ et al., 2016; McMahan et al., 2017; Rodríguez-Barroso et al., 2020; Criado et al., 2022; Rodríguez-Barroso et al., 2023; Rafi et al., 2024). Data collaboration (DC) analysis (Imakura and Sakurai, 2020; Bogdanova et al., 2020; Imakura et al., 2021b, 2023b, 2023a), one of federated learning systems, is a method that enables collaborative analysis from privacy-preserving fusion data. The privacy-preserving fusion data are constructed by sharing dimensionality-reduced intermediate representations instead of raw data. DC analysis was originally proposed to address regression and classification problems in distributed data. Recently, Kawamata et al. (Kawamata et al., 2024) proposed data collaboration quasi-experiment (DC-QE), that extends DC analysis to allow estimation of average treatment effect (ATE). It is important to note that although those studies often focus their discussions on privacy data (personal information), the technologies can also be used to preserve the management information held by companies. Therefore, in this paper, we collectively refer to such data that are difficult to disclose to the outside as confidential information. Below we discuss techniques for preserving confidential information.

In this paper, we propose data collaboration double machine learning (DC-DML) as a method to estimate CATEs while preserving the confidentiality of distributed data by incorporating double machine learning (DML) (Chernozhukov et al., 2018), a machine learning-based treatment effect estimation method, into DC-QE. Moreover, through numerical simulations, we evaluate the performance of our method compared to existing methods.

This paper relates partially linear models, one of the semi-parametric models. Since Engle et al. (Engle et al., 1986), partially linear models have been used in empirical studies where assuming a linear model between the covariates and the outcome is not appropriate, i.e., to avoid model mis-specification owing to linear assumptions. A significant advance about partially linear models in recent years is DML (Chernozhukov et al., 2018), which brought an least square estimator that can adapt common machine learning models. Various causal inference methods have been developed by developing DML (Fan et al., 2022; Bia et al., 2023). In this paper, we develop DML into a method that enables CATE estimation through distributed data with privacy-preserving.

The effects of policy interventions by governments or medical treatments by hospitals are likely to differ across people. It is possible to improve the overall performance of the treatment if one knows which subjects should be treated. For example, by concentrating scholarship support on those whose future income increases will be higher owing to the treatment, the limited budget can be used more efficiently. In addition, by concentrating medication on those who are more likely to improved or have fewer side effects from the treatment, medication can be more effective and safer to use. One approach to estimating CATEs for these situations with a high accuracy is to centralize data from many parties, but this is difficult owing to confidential or privacy concerns. Our method enables to estimate CATEs with high accuracy while maintaining the confidentiality of distributed data across many parties.

Our contributions are summarized in the following three points. First, our method enables estimation and testing of semi-parametric CATE models without iterative communication on distributed data. Semi-parametric (or non-parametric) CATE models enable estimation and testing that is more robust to model mis-specification than parametric models. However, to our knowledge, no communication-efficient method has been proposed for estimating and testing semi-parametric (or non-parametric) CATE models on distributed data. Second, our method enables collaborative estimation between different parties as well as multiple time points because the dimensionality-reduced intermediate representations can be accumulated as a knowledge base. Third, our method performed as well or better than other methods in evaluation simulations using synthetic, semi-synthetic and real-world datasets. Fig. 1 shows the summary of our study.

The remainder of this paper is organized as follows. In Section 2, we present related works on causal inference and machine learning. We describe the basics of causal inference and distributed data in Section 3. We describe our method in Section 4. In Section 5, we report the numerical simulations. In Section 6, we discuss the results. Finally, we conclude in Section 7. Acronyms and their definitions used in this paper are summarized in A.

The field of treatment effect estimation has advanced significantly in recent years through the incorporation of machine learning methods as well as statistics. However, studies that take into account confidentiality preservation, which is our aim, are limited. Here, we briefly review existing methods.

DML is an innovative method proposed by Chernozhukov et al. Chernozhukov et al. (2018) that allows estimation of ATE using any machine learning methods. DML is a semi-parametric method that uses machine learning to estimate nuisance parameters for estimating target parameters. More specifically, DML first estimates nuisance parameters using machine learning, and then estimates the target parameters using the estimated nuisance parameters. DML resolves influences of regularization bias and over-fitting on estimates of target parameters through Neyman–orthogonal scores and cross-fitting, respectively. As shown in Section 3, DML can be extended easily to allow estimation of CATE. DML is a useful method for estimating treatment effects, but using centralized data. To achieve our goal, we need to develop DML to address distributed data. In Section 3, we describe DML in detail.

As representative methods that utilize machine learning for CATE estimation, generalized random forest (GRF) (Athey et al., 2019) and X-Learner (Künzel et al., 2019) have been proposed, but these methods cannot be applied to distributed data. ¹¹1 Many other methods for estimating CATEs have been proposed so far, such as non-parametric Bayesian (Hill, 2011; Taddy et al., 2016) and model averaging (Rolling et al., 2019; Shi et al., 2024). However, these methods also cannot be applied to distributed data. GRF is a non-parametric method for estimating treatment effects by solving local moment equations using random forests. GRF first constructs a random forest, then calculates weights for each sample from the random forest, and finally derives the treatment effect by solving a moment equations using the calculated weights. X-Learner is a non-parametric method that can estimate treatment effects from any machine learning model. X-Learner uses propensity scores (Rosenbaum and Rubin, 1983) to resolve performance degradation that occurs when controlled or treated group is much larger than the other group.

Conversely, Secure regression (SR) (Karr et al., 2005) and FedCI (Vo et al., 2022) are methods that can estimate CATEs from distributed data. SR is a parametric method for estimating linear regression models. In SR, statistics necessary to compute least squares estimates are calculated in each local party, and these statistics are shared. SR is communication-efficient because it enables estimation of a linear model by sharing statistics only once. However, in SR, bias owing to mis-specification can occur if the assumptions of the linear model are incorrect. FedCI is a non-parametric method that constructs a potential outcome model based on Gaussian processes through federated learning approach. In FedCI, the communication between the clients and the server is iterated: each client sends the model gradient calculated from its local data to the server, then the server averages the aggregated gradients, updates its own model parameters, and returns those parameters to the clients. In the case where local data is stored in a location isolated from the Internet or synchronization of communications is not possible, FedCI is difficult to execute.

While the purpose of our method is to estimate CATEs, ifedtree (Tan et al., 2022), Federated MLE, IPW-MLE and AIPW (Xiong et al., 2023), and DC-QE (Kawamata et al., 2024) are methods that have been proposed for different estimation targets. Those methods can address distributed data. ifedtree is a tree-based model averaging approach to improve the CATE estimation at a specific party by sharing models derived from other parties instead of raw data. Our method aims to estimate CATEs on aggregated whole data, rather than on specific party data. Federated MLE, IPW-MLE and AIPW are a parametric methods that aggregates individually estimated gradients for estimating propensity scores and treatment effects. Federated MLE, IPW-MLE and AIPW estimate the ATE not conditioned on covariates.

DC-QE is a semi-parametric method for estimating ATE based on a DC analysis framework. DC-QE first construct collaborative representations based on dimensionality-reduced intermediate representations collected from all parties instead of raw data, second estimates propensity scores, finally estimates ATE using the estimated propensity scores (see B for the pseudo-code of DC-QE). Our method extends DC-QE to address DML for CATE estimation.

Table 1 summarizes five of those described in this section that are particularly relevant to our method. In Section 4, we propose a method for estimating CATEs in a semi-parametric model while preserving confidentiality of distributed data. Although DML, GRF and X-Learner are methods for centralized data, our method address distributed data. Moreover, our method is more robust than SR for model misspecification and more communication-efficient than FedCI.

In this section, we describe the DML method with the linear CATE model. Under the Neyman–Rubin causal model, let variables for subject $i$ as $y_{i}(0)$, $y_{i}(1)$, $z_{i}$ and $\boldsymbol{x}_{i}$, where $\boldsymbol{x}_{i}=(x_{i}^{1},\cdots,x_{i}^{m})^{\top}\in\mathbb{R}^{m}$ is the $m$-dimensional covariates vector, $z_{i}\in\{0,1\}$ is the treatment, $y_{i}(0)$ is the outcome that subject $i$ would have if they were treated ($z_{i}=1$), and $y_{i}(0)$ is that if not (controlled; $z_{i}=0$). Moreover, we denote $y_{i}=y_{i}(z_{i})$. With this definition, the CATE is defined as

In DML, we make the following structural equation assumptions on the data generating process as

$\theta$ is the function that represents the effect of the treatment conditional on the covariate $\boldsymbol{x}_{i}$. The covariates $\boldsymbol{x}_{i}$ affect the treatment $z_{i}$ via the function $h$ and the outcome via the function $u$, where $h$ and $u$ are defined as functions belong to convex subsets $\mathbb{H}$ and $\mathbb{U}$ with some normed vector space. $\varepsilon_{i}$ and $\eta_{i}$ are disturbances such that $\mathbb{E}[\eta_{i}|\boldsymbol{x}_{i}]=0$, $\mathbb{E}[\varepsilon_{i}|\boldsymbol{x}_{i}]=0$ and $\mathbb{E}[\eta_{i}\varepsilon_{i}|\boldsymbol{x}_{i}]=0$. In this data generating process, the function of our interest is $\theta$ because obviously $\tau(\boldsymbol{x}_{i})=\theta(\boldsymbol{x}_{i})$, while $h$ and $u$ are not functions of our interest. Note that, Chernozhukov et al. (Chernozhukov et al., 2018) assumed that $\theta$ is constant, i.e., independent of $\boldsymbol{x}_{i}$. However, as we discuss below, we can estimate $\theta$ similarly to Chernozhukov et al. (Chernozhukov et al., 2018) even when that $\theta$ linearly depends on $\boldsymbol{x}_{i}$.

In this paper, we derive $\theta$ by a partialling-out approach as follows: subtracting the conditional expectation of (1) from (1), we can re-write the structural equations as

Here, define $q(\boldsymbol{x}_{i})=\mathbb{E}[y_{i}|\boldsymbol{x}_{i}]$, where $q$ belongs a convex subset $\mathbb{Q}$ with some normed vector space. Then, the residuals for $q$ and $h$ are

So we can obtain

Note that $\mathbb{E}[\zeta_{i}|\boldsymbol{x}_{i}]=0$ because $\mathbb{E}[\eta_{i}|\boldsymbol{x}_{i}]=0$ and $\mathbb{E}[\varepsilon_{i}|\boldsymbol{x}_{i}]=0$.

We assume $\theta$ as a linear function with constant term such that

where $\boldsymbol{\beta}=[\beta^{0},\cdots,\beta^{m}]^{\top}\in\mathbb{R}^{m+1}$. Re-write (3) as

To derive $\boldsymbol{\beta}$, consider the following score vector:

Define the moment condition

and Neyman–orthogonality (Chernozhukov et al., 2018)

where $r\in[0,1)$, $q^{\prime}\in\mathbb{Q}$ and $h^{\prime}\in\mathbb{H}$. The Neyman–orthogonality implies first-order insensitivity to perturbations of the nuisance functions $q$ and $h$. Then, we have Theorem 1, whose proof is in C.

Given $\hat{q}$ and $\hat{h}$, $\hat{\beta}$ can be obtained by solving (5) empirically from dataset as:

(7) indicates that $\hat{\boldsymbol{\beta}}$ is the solution of the regression problem from $\hat{\zeta_{i}}$ to $\hat{\eta_{i}}[1,\boldsymbol{x}_{i}^{\top}]$. That is

The DML uses machine learning models to estimate $q$ and $h$, and cross-fitting to reduce the effect of over-fitting on the estimation of parameter $\theta$. In this paper, the data is split into two folds in cross-fitting. To derive the variance of $\hat{\boldsymbol{\beta}}$, re-write (4) as

where

From Theorem 3.2 in Chernozhukov et al. (Chernozhukov et al., 2018), the asymptotic variance of $\sqrt{n}(\hat{\boldsymbol{\beta}}-\boldsymbol{\beta})$ is

Here, $\hat{q}_{l}$ and $\hat{h}_{l}$ are the functions, and $N_{l}$ is the set of subjects for fold $l$. Moreover,

Then, the variance of $\hat{\boldsymbol{\beta}}$ is

In summary, the DML procedure can be expressed as the following two steps. First, $q$ and $h$ are estimated from the two folds using machine learning models to obtain the estimated residuals: $\hat{\zeta}_{i}$ and $\hat{\eta}_{i}$. Second, $\boldsymbol{\beta}$ is estimated by regressing $\hat{\zeta}_{i}$ to $\hat{\eta}_{i}[1,\boldsymbol{x}_{i}^{\top}]$.

We describe our method in this section. Prior to that, we explain the data distribution that our method focuses on. Let $n$ be the number of subjects (sample size) in a dataset. Moreover, let $X=[\boldsymbol{x}_{1},\cdots,\boldsymbol{x}_{n}]^{\top}\in\mathbb{R}^{n\times m}$, $Z=[z_{1},\cdots,z_{n}]^{\top}\in\{0,1\}^{n}$ and $Y=[y_{1},\cdots,y_{n}]^{\top}\in\mathbb{R}^{n}$ be the datasets of covariates, treatments, and outcomes, respectively. We consider that $n$ subjects are partitioned horizontally into $c$ parties, as follows:

The $k$th party has a partial dataset and corresponding treatments and outcomes, which are $X_{k}\in\mathbb{R}^{n_{k}\times m},Z_{k}\in\{0,1\}^{n_{k}},Y_{k}\in\mathbb{R}^{n_{k}}$, where $n_{k}$ is the number of subjects for each party in the $k$th row ($n=\sum_{k=1}^{c}n_{k}$).

This section proceeds as: we describe our method in Section 4.1, analyze the correctness of our method in Section 4.2, propose a dimension reduction method for DC-DML in Section 4.3, discuss the confidentiality preservation of our method in Section 4.4, and finally discuss the advantages and disadvantages in Section 4.5.

In this section we describe three simulations we conducted to demonstrate the effectiveness of DC-DML. ²²2We performed the simulations using Python codes, which are available from the corresponding author by reasonable request. In Simulation I, that uses synthetic data, we show that DC-DML is useful in the setting of data distributions that would be incorrectly estimated by individual analysis (IA), the analysis performed by the user using only their own dataset. In Simulation II, that uses semi-synthetic data generated from the infant dataset, we compare the performances between DC-DML and other existing methods. In Simulation III, that uses real-world data from the financial assets and jobs datasets, we investigate the robustness of DC-DML performance in the presence of heterogeneous numbers of subjects or treatment group proportions per party.

In Simulation II, the methods to be compared to DC-DML are as follows. First are the DMLs in IA and CA, that we denote for convenience as IA-DML and CA-DML, respectively. As described above, IA-DML and CA-DML are, respectively, targets that DC-DML should perform better and closer to. Second, GRF and X-Learner in IA and CA, denoted for convenience with the prefixes IA- and CA-: IA-GRF, CA-GRF, IA-XL, and CA-XL. Finally, the distributed data analysis methods SR and FedCI.

In the simulations, we set up the machine learning models used in DML and X-learner as follows. First, both machine learning models used to estimate $q$ and $h$ in the DML methods (IA-DML, CA-DML and DC-DML) in the Simulation I are random forests. Second, the settings of the machine learning models used in Simulations II and III are as follows. We selected the machine learning models used to estimate $q$ in the DML methods by our preliminary simulations from among five regression methods: multiple regression, random forests, K-nearest neighbor, light gradient boosting machine (LGBM) and support vector machine (SVM) with standardization. In the preliminary simulations for $q$, we repeated the regression of the outcomes on the covariates for the candidate methods in CA setting 50 times and set the criterion of selecting the one with the smallest root mean squared errors (RMSE) average. For $h$, similarly, we selected by preliminary simulations among five classification methods: logistic regression, random forests, K-nearest neighbor, LGBM and SVR. In our preliminary simulations for $h$, we followed the similar procedure as for $q$, where we set the criterion of selecting the one with the smallest average Brier score. Here, Brier Score is a mean squared error measure of probability forecasting (Brier, 1950). The hyperparameters for the candidate methods were the default values shown in E. As a result of our preliminary simulations, we chose SVM as the machine learning models to be used in the estimation of $q$ and $h$, both in Simulation II. For the financial data set in Simulation III, we set $q$ and $h$ to multiple regression and logistic regression, respectively. For the jobs dataset, we set $q$ and $h$ to multiple regression and random forest, respectively. We use the same models for the outcome and propensity score in X-Learner (IA-XL and CA-XL) as in DML in all simulations. The results of the preliminary simulations are in F.

The settings used in DC-DML are as follows. As the intermediate representation construction methods, we used six: PCA, LPP, FA, PCA+B, LPP+B, and FA+B. Here, +B represents the combination with the bootstrap-based dimensionality reduction method. The intermediate representation dimension and collaborative representation dimension were set as $\widetilde{m}_{k}=m-1$ for all $k$ and $\check{m}=m$, respectively. In combination with the bootstrap-based dimensionality reduction method, the BS dimension was set to three in Simulation I and $\lceil 0.1m\rceil$ in Simulation II and III. Here,$\lceil\cdot\rceil$ is the ceiling function. Anchor data were generated from uniform random numbers of maximum and minimum values for each covariate and thier subject number $r$ was the same of the original dataset $X$.

As described below, in Simulation II, we used the five methods described in Section 2 as comparison methods: DML, GRF, X-Learner, SR and FedCI. Also, in Simulation III, we used DML and SR as comparison methods. See G for the detailed settings of those methods.

In Simulations II and III, the evaluation measures used were how close the estimates were to the benchmarks. The benchmarks in Simulations II and III were the true values and the average values of 50 trials of CA-DML, respectively. In this study, we consider four evaluation measures: the RMSE of CATE, the consistency rate of significance of CATE, the RMSE of coefficients, and the consistency rate of significance of coefficients. Note that, for simplicity of explanation, results other than the RMSE of CATE in Simulation III were shown in H.