Title: FAST: a Fused and Accurate Shrinkage Tree for Heterogeneous Treatment Effects Estimation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Data fusion

heterogeneous treatment effects estimation

shrinkage estimation

tree-based method
Abstract: This paper proposes a novel strategy for estimating the heterogeneous treatment effect  called the Fused and Accurate Shrinkage Tree ($\mathrm{FAST}$). Our approach utilizes both trial and observational data to improve the accuracy and robustness of the estimator. Inspired by the concept of shrinkage estimation in statistics, we develop an optimal weighting scheme and a corresponding estimator that balances the unbiased estimator based on the trial data with the potentially biased estimator based on the observational data. Specifically, combined with tree-based techniques, we introduce a new split criterion that utilizes both trial data and observational data to more accurately estimate the treatment effect. Furthermore, we confirm the consistency of our proposed tree-based estimator and demonstrate the effectiveness of our criterion in reducing prediction error through theoretical analysis.  The advantageous  finite sample performance of the $\mathrm{FAST}$ and its ensemble version over existing methods is demonstrated via  simulations and real data analysis.

FAST: a Fused and Accurate Shrinkage Tree for Heterogeneous Treatment Effects Estimation

Abstract

This paper proposes a novel strategy for estimating the heterogeneous treatment effect called the Fused and Accurate Shrinkage Tree ( FAST ). Our approach utilizes both trial and observational data to improve the accuracy and robustness of the estimator. Inspired by the concept of shrinkage estimation in statistics, we develop an optimal weighting scheme and a corresponding estimator that balances the unbiased estimator based on the trial data with the potentially biased estimator based on the observational data. Specifically, combined with tree-based techniques, we introduce a new split criterion that utilizes both trial data and observational data to more accurately estimate the treatment effect. Furthermore, we confirm the consistency of our proposed tree-based estimator and demonstrate the effectiveness of our criterion in reducing prediction error through theoretical analysis. The advantageous finite sample performance of the FAST and its ensemble version over existing methods is demonstrated via simulations and real data analysis.

1 Introduction

Causal effects are the magnitude of the response of an effect variable (also called outcome) caused by the effect variable (also called treatment), which is a fundamental and essential issue in the field of casual inference (Imbens and Rubin, 2016). And the heterogeneous treatment effect (abbr. HTE) is usually used to characterize the heterogeneity of causal effects across different subgroups of the population. In recent years, heterogeneous treatment effect estimation has been successfully applied in various fields such as epidemiology, medicine, and social sciences (Glass et al., 2013; Kosorok and Laber, 2019; Turney and Wildeman, 2015; Taddy et al., 2016).

In general, the causal problems can be studied through both experimental studies (also known as randomized control trials, RCTs) and observational studies. Experimental studies are widely regarded as the gold standard for assessing causal effects since the randomization process eliminates the possibility of confounding bias. However, large-scale RCTs can be challenging due to issues related to cost, time, and ethics (Edwards et al., 1999). On the other hand, observational data are often readily available with an adequate sample size. Under certain fairly strong assumptions, such as unconfoundedness assumption, there is a rich literature regarding the estimation of HTE in observational studies, such as tree-based methods (Athey and Imbens, 2016; Wager and Athey, 2018; Athey et al., 2019; Hahn et al., 2020), boosting (Powers et al., 2017), meta learners (Künzel et al.,

2019) and R -learner (Nie and Wager, 2020). However, the unconfoundedness assumption, which requires measuring all confounders, is in general untestable unless extra information such as the existing of instrumental variables is available (de Luna and Johansson, 2014). And any violations of this assumption may result in seriously invalid causal statements. Various methods have been proposed to mitigate the unmeasured confounding in observational studies, such as the sensitivity analysis (Rosenbaum and Rubin, 1983; Zhang and Tchetgen Tchetgen, 2022), the instrumental variables (IV) approach (Angrist et al., 1996) and the proximal causal inference (Kuroki and Pearl, 2014; Miao et al., 2018; Shi et al., 2020; Cui et al., 2023). However, the validity of these procedures also relies crucially on assumptions that are often difficult to verify in practice.

Given the limitations of relying on individual data sources, data fusion, as a branch of causal inference strategies that integrates both the trial and the observational data, has gained significant interest in the literature (Bareinboim and Pearl, 2016; Colnet et al., 2020; Shi et al., 2022). Existing data fusion methods for estimating the HTE include the KPS estimator obtained by modeling the confounding function parametrically (Kallus et al., 2018), the semi-parametric integrative estimator under the parametric structural models (Yang et al., 2020) and the integrative R-learner (Wu and Yang, 2022). Besides, (Tang et al., 2022) proposed the Gradient Boosting Causal Tree ( GBCT ), which integrated the current observational data and their historical controls for estimating the conditional average treatment effect on the treated group (CATT).

This paper presents a novel approach for estimating heterogeneous treatment effects (HTE) in the context of causal data fusion. The proposed method, named Fused and Accurate Shrinkage Tree ( FAST ), avoids the need for a two-stage estimation process required in conventional data fusion strategies, which involves modeling and estimating the nuisance confounding bias function. The main contributions of this work can be summarized as follows (i) The authors propose a novel shrinkage method for combining an unbiased and biased estimator, which effectively reduces the mean square error of the unbiased estimator, and provides an easy implementation of the method tailored for the HTE estimation; (ii) The authors extend the conventional node split criterion via a re-scaling technique, which automatically penalizes the use of the observational data with low quality (namely large confounding bias); (iii) The authors also provide a theoretical analysis to explain the advantages of our splitting criterion.

2 Background and motivation

2.1 Notations

Let X ∈ X = [ -1 , 1] p be a p -dimensional vector of pre-treatment covariates, U ∈ R q ( q ≥ 0 ) be a possibly unmeasured random vector consisting of the confounding variables, D be a binary treatment variable ( D = 0 denotes the control and D = 1 denotes the treated) and let Y ( d ) be the potential outcome that would be observed when the treatment had been set to d ∈ { 0 , 1 } . We follow the potential outcome framework (Rubin, 1974) to define the heterogeneous treatment effect τ ( x ) = E ( Y (1) -Y (0) | X = x ) .

Suppose that we can collect two kinds of data: trial data and observational data, and they are described by n + m quadruples, { Y i , D i , X i , S i } n + m i =1 , where S i indicates if the i -th individual would have been recruited ( S = 1 ) or not ( S = 0 ) in the trial. We also denote R = { 1 , 2 , · · · , n } the set of indices of observations in the RCT study, and O = { n +1 , n +2 , · · · , n + m } the set of indices of observations in the observational study. We define e ( X , U , S ) = P ( D = 1 | X , U , S ) as the propensity score of the trial and observational population, respectively. In practice, due to U being unknown, we usually use ˆ e ( X , S ) to estimate e ( X , U , S ) . In addition, ˆ e ( X , 1) is unbiased, but ˆ e ( X , 0) is biased because the unmeasured confounder U in the observational study can be related to the assignment of treatment D . Let τ 1 ( x ) = E ( Y (1) -Y (0) | X = x , S = 1) be the HTE on the trial population. We then make the following fundamental assumption on the trial and observational studies, which facilitates the potential for causal data fusion:

Assumption 1. (i) For any x ∈ X , τ 1 ( x ) = τ ( x ) ; (ii) Y ( d ) ⊥ D | ( X , S = 1) for d ∈ { 0 , 1 } and (iii) the propensity score δ < e ( X , S ) < 1 -δ almost surely for some constant 0 < δ < 1 / 2 .

Assumption 1 (i) states that the HTE function is transportable from the trial population to the target population. Stronger versions of Assumption 1 include the ignorability of study participation (Buchanan et al., 2018) and the mean exchangeability (Dahabreh et al., 2019). In the following of this

paper, we use | Λ | to denote the number of elements for any set Λ , ⌊ c ⌋ to denote the biggest integer less than or equal to the constant c , and [ p ] to denote the index set { 1 , 2 , · · · , ⌊ p ⌋} . For two positive sequences { a n } n ≥ 1 and { b n } n ≥ 1 , we write a n = O( b n ) if | a n /b n | is bounded.

2.2 Tree-based methods

To estimate the HTE, it is reasonable to perform subgroup analysis by appropriately stratifying or matching (Frangakis and Rubin, 2002) the samples into multiple subgroups that differ in the magnitude of treatment effects. In machine learning, tree-based methods (Breiman et al., 1984; Breiman, 2001; Friedman, 2001) are usually used for such stratification tasks, which greedily optimize the loss function, also called splitting criterion, via recursively partitioning feature space. In fact, many treebased causal methods designed for the HTE estimation were also proposed (Radcliffe and Surry, 2012; Athey and Imbens, 2016; Athey et al., 2019). Along with the development of the tree-based methods, various regularization strategies, either implicit or explicit, have been proposed to mitigate overfitting (Mentch and Zhou, 2020; Agarwal et al., 2022). Recently, tree-based methods have been generalized to address the heterogeneous data from diverse data sources(Nasseri et al., 2022). For convenience, in the following we define a regression tree by two components: a set of leaves Q = { Q j } J j =1 and the associated parameter τ . We can denote a causal tree by T ( X ; Q , τ ) = ∑ J j =1 τ ( Q j )I { x ∈ Q j } , where I {·} denotes the indicator function and τ ( Q j ) denotes the casual effect of sub-area indexed by Q j .

GLYPH(cmap:d835)GLYPH(cmap:dc49)GLYPH(cmap:d835)(GLYPH(cmap:d835)

GLYPH(cmap:d835)) = 5

GLYPH(cmap:d835)GLYPH(cmap:dc49)GLYPH(cmap:d835)(GLYPH(cmap:d835)

GLYPH(cmap:d835)

) = 0.5,  GLYPH(cmap:d835) GLYPH(cmap:d835)

GLYPH(cmap:d835)  = 2

GLYPH(cmap:d835) = 0

GLYPH(cmap:d835)GLYPH(cmap:dc49)GLYPH(cmap:d835)(GLYPH(cmap:d835)

GLYPH(cmap:d835)

∗ )

= 1.26,  GLYPH(cmap:d835) GLYPH(cmap:d835)

GLYPH(cmap:d835) ∗   = 1.05

GLYPH(cmap:d835) GLYPH(cmap:d835)

GLYPH(cmap:d835)

∗

-GLYPH(cmap:d835)

2

= 2.37 <  GLYPH(cmap:d835)GLYPH(cmap:dc5a)GLYPH(cmap:d835){GLYPH(cmap:d835) GLYPH(cmap:d835)

GLYPH(cmap:d835)

-GLYPH(cmap:d835)

2

, GLYPH(cmap:d835) GLYPH(cmap:d835)

GLYPH(cmap:d835)

-GLYPH(cmap:d835)

2

}

Shrinkage

Estimation

Optimal weight   GLYPH(cmap:d835)

∗

= 0.53

一

0.2

0.0

-5.0

-2.5

0.0

2.5

5.0

0.2

0.0

-5.0

-2.5

0.0

2.5

5.0

2.3 Shrinkage estimation

It is important to note that applying conventional methods, such as the generalized random forest (Athey et al., 2019), separately to trial data and observational data can readily lead to two estimators: the first is unbiased but may exhibit large variability, while the second is potentially biased but usually has a smaller variance due to the much larger amount of observational data. Therefore, the challenge becomes finding the optimal combination of an unbiased estimator and a biased estimator in the data fusion problem. To see this, suppose we have a parameter of interest θ ∈ R , an unbiased estimator ˆ θ u , and a (potentially) biased estimator ˆ θ b of θ , such that E ( ˆ θ u ) = θ , E ( ˆ θ b ) = θ + b ( θ ) , Var( ˆ θ u ) = σ 2 u , Var( ˆ θ b ) = σ 2 b and Cov( ˆ θ u , ˆ θ b ) = 0 . Consider the family of estimators Λ w = { ˆ θ w | ˆ θ w = w ˆ θ b +(1 -w ) ˆ θ u , 0 ≤ w ≤ 1 } , then the mean square error (MSE) of its elements admits the following expansion:

Minimizing (1) with respect to w , we can obtain the unique minimizer w ∗ = σ 2 u / ( σ 2 b + b 2 ( θ ) + σ 2 u ) and the gain of the optimal weighting over the single estimators ˆ θ u and ˆ θ b can be characterized by the following formula:

Comment The weighting strategy is akin to the classical James-Stein shrinkage estimation (Efron and Morris, 1973; Green and Strawderman, 1991) method, in which it is shown that a multivariate normal vector Z ( p ≥ 3 ), as a maximum likelihood estimator (MLE) of its population mean µ = E ( Z ) , is not minimax optimal, and the MSE of the estimator Z can be reduced by shrinking it towards the zero vector 0 by some factor 0 < w < 1 . The zero vector can be viewed as a biased estimator of µ with zero variance in their setting. In comparison, we replace the deterministic estimator with a (potentially) biased estimator ˆ θ b : The larger the variance σ 2 u of the unbiased estimator is compared to b 2 ( θ ) + σ 2 b , the more the fused estimator ˆ θ w ∗ will be shrunk towards the biased estimator that is less fluctuating. By doing so, one can efficiently mitigate the occurrence of significant estimation error in the unbiased estimator caused by its high variance, as unbiasedness alone does not guarantee reliable estimation performance with a limited sample size. Figure 1 illustrates a concrete example of the benefit provided by the shrinkage estimation, where θ = 0 , ˆ θ u ∼ N( θ, 5) and ˆ θ b ∼ N ( θ +2 , 0 . 5) . The fused estimator ˆ θ w ∗ reduces over 50% of the MSE compared with the unbiased estimator ˆ θ u .

3 Methodology

In this section, we propose a new data fusion strategy, referred to as the Fused and Accurate Shrinkage Tree (FAST). We proceed in a bottom-up manner to provide a clear and intuitive illustration of the entire estimation: we will begin by applying the shrinkage estimation strategy for local data fusion within each sub-region of the feature space given by a pre-specified partition. Then, we propose a fused criterion that incorporates the information contained in the observational data via a simple re-scaling of the conventional criterion. Theoretical guarantees are established in Section 4.

3.1 Local fusion for the HTE estimation

Under a pre-specified partition Q = { Q j } J j =1 of the feature space, let R j = { i | i ∈ R , X i ∈ Q j } and O j = { i | i ∈ O , X i ∈ Q j } represent the sets of indices of the trial and observational sub-samples, respectively, that fall within the region Q j . Let

be transformed outcomes of all data, e.g., the transformed outcomes of i -th sample can be denoted by ˜ Y i . Then under Assumption 1, one can immediately show for the trial population with S = 1 :

Thus, ˆ τ u ( Q j ) = (1 / |R j | ) ∑ i ∈R j ˜ Y i is an unbiased estimator of E ( Y (1) -Y (0) | X ∈ Q j , S = 1) , which can be seen as a reasonable approximation of τ ( Q j ) if Q segments the feature space properly such that τ ( x ) varies slowly in each sub-region Q j . An estimator of Var(ˆ τ u ( Q j )) is given by ˆ σ 2 u ( Q j ) = (1 / | ( R j | ( |R j |-1))) ∑ i ∈R j ( ˜ Y i -ˆ τ u ( Q j )) 2 . In contrast, for the observational population, the conditional independence no longer holds and ˆ τ b ( Q j ) = (1 / |O j | ) ∑ i ∈O j ˜ Y i is a biased estimator concerning τ ( Q j ) , due to the presence of unmeasured confounding ( U ) on the observational data.

It remains to estimate the weight w ∗ ( Q j ) composed of the tuple ( σ 2 u ( Q j ) , σ 2 b ( Q j ) , b 2 ( Q j )) . The first term σ 2 u ( Q j ) can be estimated by ˆ σ 2 u ( Q j ) . To bypass the unmeasured confounding issue of the observational population, re-sampling techniques, such as the Bootstrap (Efron, 1979; Hall, 1992), can be applied to consistently estimate σ 2 b ( Q j ) . However, in the causal data fusion setting, σ 2 b ( Q j ) = O( |O j | -1 ) is expected to be of a smaller order term compared to σ 2 u ( Q j ) = O( |R j | -1 ) ,

which is a consequence of the relative sample size between the trial and the observational data. Thus, one can just avoid estimating the negligible term σ 2 b ( Q j ) . For the last term, ̂ b ( Q j ) = ˆ τ b ( Q j ) -ˆ τ u ( Q j ) serves as a natural estimator of the bias b ( Q j ) . This leads to the following estimator of w ∗ ( Q j ) and the corresponding fused estimator

A fused estimator of the HTE function τ ( · ) under the partition Q can thus be defined as ˆ τ Q ( x ) = ∑ J j =1 ˆ τ of ( Q j )I { x ∈ Q j } .

3.2 Adaptive fusion for segmentation

In order to obtain a tree-based partition Q designed for the fusion strategy (6), a split criterion is required, which is sufficient to be defined only at the root node given the recursive nature of the partitioning. We follow the honest estimation approach (Athey and Imbens, 2016) to prevent overfitting. Specifically, given a fraction 0 < r < 1 (typically r = 0 . 5 ), ⌊ rn ⌋ observations are sampled without replacement from the trial data of sample size n for the tree structure estimation, while the rest of observations are used for local estimation of the HTE in each leaf node. Let the index sets of the trial data used for the partition and the HTE estimation be R t and R e , respectively. We do not further split the observational data to reduce uncertainty, since we have already partitioned the trial data to avoid overfitting.

The conventional criterion for growing a regression tree chooses the index of the split variable and its split value at the root node by minimizing the following goodness-of-fit criterion

where ̂ Q L = { x | x ˆ q ≤ ˆ c } , ̂ R t L = { i | i ∈ R t , X i ∈ ̂ Q L } and ˆ τ u ( ̂ Q L , R t ) = (1 / |{ i | i ∈ R t , X i ∈ ̂ Q L }| ) ∑ i ∈{ i | i ∈R t , X i ∈ ̂ Q L } ˜ Y i , and ̂ Q R , ̂ R t R and ˆ τ u ( ̂ Q R , R t ) can be defined correspondingly. Given a tree grown under (7), we fuse the trial data indexed by R e and the observational data indexed by O at each leaf node according to (6) and refer to the resulting tree estimator as a Shrinkage Tree ( ST ) . A direct modification of (7), which aligns more with the fused estimator at the leaf nodes, should be

where ˆ τ of ( ̂ Q L ) = ˆ w of ( ̂ Q L )ˆ τ b ( ̂ Q L ) + (1 -ˆ w of ( ̂ Q L ))ˆ τ u ( ̂ Q L , R t ) and ˆ τ of ( ̂ Q R ) is defined correspondingly. The replacement of the unbiased estimators in (7) with the fused estimators in (8) facilitates a goodness-of-fit criterion of the proposed fusion strategy.

Alternatively, (7) can be interpreted as minimizing the sum of the estimated MSEs of the unbiased estimators at the child nodes, if the two terms on the right-hand side of (7) are divided by the square of their respective sample sizes. By contrast, since the fused estimator ˆ τ of reduces variance by shrinking the original unbiased estimator to a potentially biased estimator, simply comparing the fused estimators with the outcomes of the trial data as in (8) fails to capture the variability at the child nodes. Instead, an appropriate criterion shall respect the MSE of the fused estimator. To this end, we introduce the following split criterion

where (1 -ˆ w of ( ̂ Q L ))ˆ σ 2 u ( ̂ Q L , R t ) and (1 -ˆ w of ( ̂ Q R ))ˆ σ 2 u ( ̂ Q R , R t ) estimate the MSE of ˆ τ of ( ̂ Q L ) and ˆ τ of ( ̂ Q R ) , respectively, according to formula (2). Compared to (7), the proposed criterion incorporates the additional information from the observational data into each node split in an adaptive manner by simply re-scaling the estimated MSE of the unbiased estimator.

Comment The criterion (9) offers the benefit of local adjustment, which can be intuitively justified. In sub-regions where the observational data exhibit moderate confounding biases, this criterion improves tree building by providing a sharper assessment of the variability of the fused estimator. On the other hand, for sub-regions where the observational data exhibit substantial confounding biases, the estimated weights of those sub-regions approach zero according to (5). In such cases, the criterion reduces to the conventional criterion (7), except for the standardization of the square of the sample size. It is worth mentioning that all the local adjustments achieved by applying this adaptive fusion strategy are data-driven, namely one can just avoid global modeling of the confounding bias function, which requires domain-specific knowledge of the observational studies. Additionally, it also enables the exclusion of the global impact of extremely large confounding biases of the observational data that only exist in certain sub-regions of the feature space.

We denote the partition obtained under criterion (9) as ̂ Q of = { ̂ Q of, 1 , ̂ Q of, 2 , · · · , ̂ Q of, | ̂ Q of | } , and the corresponding tree-based estimator of the HTE is defined as

where the superscript 'e' is to show that the RCT data used to construct the fused estimator at the leaf node is indexed by R e and 'fast' is an acronym for the name Fused and Accurate Shrinkage Tree, which is due to the data fusion nature of the criterion (9), the shrinkage-type leaf node estimator (6) and its accuracy in terms of the MSE.

3.3 Ensemble fusion

To reduce overfitting, improve robustness against outliers, and enhance generalization, we introduce the bagged version (Hastie et al., 2009) of the FAST , referred to as the rfFAST , as follows: We randomly draw index sets R ∗ of size n and O ∗ of size m , separately from R and O with replacement. We repeat the process B times, resulting in {R ∗ , ( b ) , O ∗ , ( b ) } B b =1 . Then, B estimators ˆ τ ∗ , ( b ) fast ( x ) can be calculated based on the trial data indexed by R ∗ , ( b ) and the observational data index by O ∗ , ( b ) . We then define ˆ τ rffast ( x ) = (1 /B ) ∑ B b =1 ˆ τ ∗ , ( b ) fast ( x ) . For the construction of the prediction intervals, see Zhang et al. (2020).

4 Theoretical guarantee

In this section, we formally establish the benefits of the proposed split criterion (9) compared with the conventional criterion (7). To present the theoretical result, we first pose the following regularity conditions that are standard in literature (see e.g., Györfi et al., 2002 and Scornet et al., 2015).

Assumption 2. (i) There exists a positive constant λ < ∞ such that E { exp( λ ˜ Y 2 ) | S = i } < ∞ for i = 0 , 1 . (ii) There exists positive constants σ min < ∞ such that σ 2 min < Var( ˜ Y | X = x , S = 0) for any x ∈ X .

Theorem 1 (MSE reduction of the proposed split criterion) . Let θ = ( q, c ) and Θ = [ p ] × R . Suppose the node that needs to be partitioned is Q j , under which the sample sizes of the trial data and observational data are n j and m j , respectively. Let M ( θ ) and M of ( θ ) be the sum of MSEs of the conventional HTE estimator and the fused HTE estimator on the two child nodes of Q j split by θ , respectively. Denote b max = sup x ∈ Q j |{ E ( ˜ Y | X = x , S = 0) -E ( ˜ Y | X = x , S = 1) }| . Let ˆ θ be the solution of the conventional split criterion (7) and ˆ θ of be the solution of the proposed split criterion (9). Under Assumptions 1-2, we have

(i) For any θ ∈ Θ ,

(ii) With probability at least 1 -C 1 e -t for some positive constant C 1 < ∞ , it holds that

for some positive constant C 2 , C 3 < ∞ , where θ ∗ and θ ∗ of are oracle splits defined as

In the above theorem, the (i) part establishes a uniform MSE reduction result for any split choice θ ∈ Θ of the proposed split criterion (9). As revealed in (11), the criterion (9) leads to larger MSE reduction on the nodes with a larger variance of ˜ Y and less bias of the observational data. In addition, the upper bound in (11) decreases as the node sample size n j decreases, implying that our proposed criterion leads to increasing relative benefits as the tree grows deeper. Besides, in the (ii) part we present non-asymptotic bounds for the discrepancies between the MSEs under the empirically estimated splits and the oracle splits, showing that the MSEs under the estimated splits can achieve a fast convergence rate. As a direct consequence of Theorem 1, the consistency of our final HTE estimator (10) can be established, since it is known from Scornet et al. (2015) and Athey et al. (2019) that the conventional tree-based estimator using only the trial data is mean-squared consistent, and our proposed method leads to a reduced MSE.

Proposition 1 (Consistency of ˆ τ fast ) . For almost every x ∈ [ -1 , 1] p , we have ˆ τ fast ( x ) → τ ( x ) in probability as n, m →∞ .

5 Experiments

In this section, we demonstrate the results of a series of experiments to answer the following two questions: (i) Whether the proposed method can effectively alleviate the impact of confounding bias of observational data and limited sample size of trial data; (ii) Whether the techniques we proposed including local fusion in tree leaves and adaptive fusion in partitioning are valid, respectively.

In consequence, we conducted experiments on both simulated and real-world datasets to verify the effectiveness of our method. We evaluated our method against both traditional tree-based and data fusion-based casual methods. The former includes the classical Transformed Outcome Honest Tree ( HT ) Athey and Imbens (2016) and its ensemble version Generalized Random Forest ( GRF ) Athey et al. (2019). The latter includes the simplest fusion estimator ( SF ) training both trial data and observational data together without distinction and the KPS estimators Kallus et al. (2018). In order to facilitate better comparison and understanding of our proposed method, we demonstrate three versions: the simple implementation, Shrinkage Tree ( ST ), described in Section 3.1; the improved version, Fused and Accurate Shrinkage Tree ( FAST ), described in Section 3.2; and its final ensemble version rfFAST described in Section 3.3. The results of each simulation experiment were based on B = 100 replications. The ensemble size for all the ensemble estimators was set to 100 . For the tree estimators, the minimum number of observations required to be at a leaf node was set to 5 and the maximum depth of the tree was set to 10 .

5.1 Simulation

We conducted two sets of simulation experiments to evaluate the finite sample performance of the fused estimator and various baseline estimators. In both experiments, we first generated the pretreatment covariates X = ( X 1 , X 2 , · · · , X p ) T from Uniform[ -1 , 1] p and the unobserved variable U from N(0 , 1) . Then, we generated the potential outcomes by Y ( d ) = dτ ( X )+ ∑ p j =1 X j +1 . 5 U + ϵ ( d ) , where τ ( X ) = 1 + X 1 + X 2 1 + X 2 + X 2 2 and ϵ ( d ) ∼ N(0 , 1) for d = 0 , 1 . Thus The treatment assignments for the trial sample of size n and the observational sample of size m were generated as follows: D | ( X , U, S = 1) ∼ Ber(0 . 5) and D | ( X , U, S = 0) ∼ Ber(1 / (1 + exp( -βU -0 . 5 X 1 ))) . Thus, the magnitude of β controls the strength of the unmeasured confounding: a larger β leads to a larger confounding bias. The test data X test,j for 1 ≤ j ≤ p were generated from Uniform( -1 , 1) with sample size 1000 .

In the first experiment, we aim to verify the effectiveness of the proposed data fusion strategy via an ablation study. We compared the robustness of the ST and the FAST against different levels of confounding bias parameter β . Two baselines were considered: (i) the HT using only the trial data and (ii) the SF estimator obtained by directly merging all the available data and constructing a Fit-Based Causal Tree (Athey and Imbens, 2016). We set the sample sizes of the trial data and the observational data be n = 200 and m = 2000 , respectively, the dimension of covariates p = 5 and β ∈ { 0 . 1 c | c ∈ N , c ≤ 19 } . The following three conclusions could be drawn from Figure 2: (1) When confounding bias in observational data was small, the simple fusion ( SF ) strategy can effectively improve the model performance. But when it became large, the SF was very vulnerable to confounding bias in observational data; (2) Even with the increase of β , both ST and FAST consistently showed resistance to confounding bias; (3) FAST was significantly better than other methods including ST , which verified the effectiveness of our proposed split criterion (9) numerically.

Estimators 一 HT 一 SF - ST - FAST

2.2

RMSE

2.0

1.8

0.0

0.5

1.0

1.5

2.0

β

In the second experiment, we evaluated the RMSEs with respect to different n and β . We set m = 2000 and p = 5 . We included seven estimators in the analysis: The first two estimators were calculated purely based on the trial data: (i) the Transformed Outcome Honest Tree ( HT ) (Athey and Imbens, 2016) and (ii) the Generalized Random Forest ( GRF ) (Athey et al., 2019). The rest estimators were calculated using different data fusion strategies: (iii) the Shrinkage Tree ( ST ) estimator,(iv) the Fused and Accurate Shrinkage Tree ( FAST ) estimator, (v& vi) the KPS estimators (Kallus et al., 2018) with a parametric (OLS) estimator and a non-parametric (Random Forest) specification of the confounding function, respectively and (vii) the bagged FAST estimator ( rfFAST ).

Table 1 reports the RMSEs of the seven estimators, conveying a good estimation accuracy of both the FAST and its ensemble version rfFAST . Among the three individual estimators, the ST and FAST , exhibited superior performance compared to the HT , and the FAST outperformed the ST . These relative performances provided support for the FAST approach compared to the classical honest regression tree, the proposed split criterion (9), and the shrinkage estimation strategy (6), which are implemented progressively. Among the three ensemble estimators, the rfFAST estimator demonstrated the best performance among all the six combinations of the trial sample size n and the confounding bias parameter β . On the other hand, the performance of the KPS estimators appeared to be less stable. The KPS ols outperformed the GRF only when the trial sample size was relatively large ( n = 200 ). Under the non-parametric specification of the confounding function, the KPS RF did not gain benefit from incorporating the observational data and was consistently inferior to the baseline estimator GRF .

5.2 Real-world data

In this sub-section, we report an analysis of the Tennessee Student/Teacher Achievement Ratio (STAR) Experiment (Krueger, 1999) to demonstrate the proposed FAST for the HTE estimation. We aim at quantifying the treatment effect of the class size on the student's academic achievement.

Data description The STAR Experiment was a randomized controlled trial conducted in the late 1980s. Students were randomly assigned to one of the two types of classes during the first school year: D = 1 for small classes containing 13 -17 pupils and D = 0 for regular classes containing 22 -25 pupils. The outcome Y is the average of the listening, reading, and math standardized tests at the end of first grade. The vector of covariates X includes gender, race, birth month, birthday, birth year, free lunch given or not, and teacher id. This made a universal sample of 4218 students, among which 2413 were randomly assigned to regular-size classes ( D = 0 ) and 1805 to small classes ( D = 1 ).

Ground-truth In practice, the ground-truth τ ( · ) is not accessible, so we replaced it with an estimate calculated by a generalized random forest (Athey et al., 2019) based on all the 4218 observations.

Construction of the trial, observational and test data Following Kallus et al. (2018), we introduced confounding bias by splitting the population over a variable which is known to strongly affect the observed outcome Y (Krueger, 1999): rural or inner-city ( U = 1 , 2811 students) and urban or suburban ( U = 0 , 1407 students). The trial data were generated by randomly sampling a fraction h of the students with U = 1 , where h ranges from 0 . 1 to 0 . 5 . The observational data were constructed as follows: From students with U = 1 , we took the controls ( D = 0 ) that were not sampled in trial data, and the treated ( D = 1 ) whose outcomes were in the lower half of outcomes among students with D = 1 and U = 1 ; From students with U = 0 , we took all of the controls ( D = 0 ), and the treated ( D = 1 ) whose outcomes were in the lower half of outcomes among students with D = 1 and U = 0 . The test data consisted of a held-out sub-sample of all the observations in the universal sample excluding the trial data.

Results We compared the performance of the rfFAST with various baseline estimators. In particular, the NF and the SF estimators were constructed using the Random Forest regressor. The NF estimator utilized only trial data, while the SF estimator utilized both trial data and observational data together without distinction. As shown in Figure 3, the proposed rfFAST method consistently outperformed other estimators.

6 Discussion

This paper explores the estimation of heterogeneous treatment effects (HTE) within the framework of causal data fusion. Drawing inspiration from the classical James-Stein shrinkage estimation (Green and Strawderman, 1991) approach, the authors introduce a new method called Fused and Accurate Shrinkage Tree ( FAST ) that effectively incorporates observational data in both feature

Estimators → NF + SF + KPS(ridge)→ KPS(RF)+ rfFAST

40

30

RMSE

20-

10

0.1

0.2

0.3

0.4

0.5

h

space segmentation and leaf node value estimation. This new approach is shown to outperform existing data fusion methods via numerical experiments.

The above estimation framework can be generalized to any data fusion problem if there exists an unbiased estimator and a biased estimator of some functions of interest. It would be worthwhile to explore the combination of the FAST method with other ensemble methods, such as the boosting and the grf-style (Athey et al., 2019) bagging, in addition to Breiman-style (Breiman, 2001) bagging used in rfFAST . Moreover, extending the framework to handle time-series observational data would be an interesting direction for future research. Additionally, investigating statistical inference under the proposed fusion framework would also be valuable.

References

Agarwal, A., Tan, Y. S., Ronen, O., Singh, C., and Yu, B. (2022). Hierarchical Shrinkage: Improving the accuracy and interpretability of tree-based models. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S., editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 111-135. PMLR.

Angrist, J. D., Imbens, G. W., and Rubin, D. B. (1996). Identification of causal effects using instrumental variables. Journal of the American Statistical Association , 91(434):444-455.

Athey, S. and Imbens, G. (2016). Recursive partitioning for heterogeneous causal effects: Table 1. Proceedings of the National Academy of Sciences , 113:7353-7360.

Athey, S., Tibshirani, J., and Wager, S. (2019). Generalized random forests. The Annals of Statistics , 47(2):1148 - 1178.

Bareinboim, E. and Pearl, J. (2016). Causal inference and the data-fusion problem. Proceedings of the National Academy of Sciences , 113:7345-7352.

Breiman, L. (2001). Random Forests. Machine Learning , 45:5-32.

Breiman, L., Friedman, J., Olshen, R., and Stone, C. (1984). Classification And Regression Trees . Chapman & Hall/CRC.

Buchanan, A. L., Hudgens, M. G., Cole, S. R., Mollan, K. R., Sax, P. E., Daar, E. S., Adimora, A. A., Eron, J. J., and Mugavero, M. J. (2018). Generalizing Evidence from Randomized Trials Using Inverse Probability of Sampling Weights. Journal of the Royal Statistical Society Series A: Statistics in Society , 181(4):1193-1209.

Colnet, B., Mayer, I., Chen, G., Dieng, A., Li, R., Varoquaux, G., Vert, J.-P., Josse, J., and Yang, S. (2020). Causal inference methods for combining randomized trials and observational studies: a review. arXiv preprint arXiv:2011.08047 .

Cui, Y., Pu, H., Shi, X., Miao, W., and Tchetgen, E. T. (2023). Semiparametric Proximal Causal Inference. Journal of the American Statistical Association , 0(0):1-12.

Dahabreh, I. J., Robertson, S. E., Tchetgen, E. J., Stuart, E. A., and Hernán, M. A. (2019). Generalizing causal inferences from individuals in randomized trials to all trial-eligible individuals. Biometrics , 75(2):685-694.

de Luna, X. and Johansson, P. (2014). Testing for the Unconfoundedness Assumption Using an Instrumental Assumption. Journal of Causal Inference , 2(2):187-199.

Edwards, S., Lilford, R., Braunholtz, D., Jackson, J., Hewison, J., and Thornton, J. (1999). Ethical issues in the design and conduct of Randomised Controlled Trials. Health technology assessment (Winchester, England) , 2:i-vi, 1.

Efron, B. (1979). Bootstrap Methods: Another Look at the Jackknife. The Annals of Statistics , 7(1):1 - 26.

Efron, B. and Morris, C. (1973). Stein's estimation rule and its competitors-an empirical bayes approach. Journal of the American Statistical Association , 68(341):117-130.

Frangakis, C. and Rubin, D. B. (2002). Principal stratification in causal inference. Biometrics , 58.

Friedman, J. H. (2001). Greedy function approximation: A gradient boosting machine. The Annals of Statistics , 29(5):1189 - 1232.

Glass, T. A., Goodman, S. N., Hernán, M. A., and Samet, J. M. (2013). Causal inference in public health. Annual review of public health , 34:61-75.

Green, E. and Strawderman, W. (1991). A James-Stein type estimator for combining unbiased and possibly biased estimators. Journal of The American Statistical Association - J AMER STATIST ASSN , 86:1001-1006.

Györfi, L., Köhler, M., Krzy˙ zak, A., and Walk, H. (2002). A distribution-free theory of nonparametric regression , volume 1. Springer.

Hahn, P. R., Murray, J. S., and Carvalho, C. M. (2020). Bayesian Regression Tree Models for Causal Inference: Regularization, Confounding, and Heterogeneous Effects (with Discussion). Bayesian Analysis , 15(3):965 - 2020.

Hall, P. (1992). The bootstrap and Edgeworth expansion . Springer-Verlag New York.

Hastie, T., Tibshirani, R., and Friedman, J. (2009). The Elements of Statistical Learning : Data Mining, Inference, and Prediction . Springer series in statistics. Springer, 2nd edition.

Imbens, G. W. and Rubin, D. B. (2016). Causal inference for statistics, social, and biomedical sciences: An introduction . Taylor & Francis.

Kallus, N., Puli, A. M., and Shalit, U. (2018). Removing hidden confounding by experimental grounding. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc.

Kosorok, M. and Laber, E. (2019). Precision medicine. Annual Review of Statistics and Its Application , 6:263-286.

Krueger, A. (1999). Experimental Estimates Of Education Production Functions. The Quarterly Journal of Economics , 114:497-532.

Kuroki, M. and Pearl, J. (2014). Measurement bias and effect restoration in causal inference. Biometrika , 101(2):423-437.

Künzel, S., Sekhon, J., Bickel, P., and Yu, B. (2019). Meta-learners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences , 116:4156-4165.

Mentch, L. and Zhou, S. (2020). Randomization as regularization: A degrees of freedom explanation for random forest success. Journal of Machine Learning Research , 21(171):1-36.

Miao, W., Geng, Z., and Tchetgen Tchetgen, E. J. (2018). Identifying causal effects with proxy variables of an unmeasured confounder. Biometrika , 105(4):987-993.

Nasseri, K., Singh, C., Duncan, J., Kornblith, A., and Yu, B. (2022). Group probability-weighted tree sums for interpretable modeling of heterogeneous data.

Nie, X. and Wager, S. (2020). Quasi-oracle estimation of heterogeneous treatment effects. Biometrika , 108(2):299-319.

Powers, S., Qian, J., Jung, K., Schuler, A., Shah, N., Hastie, T., and Tibshirani, R. (2017). Some methods for heterogeneous treatment effect estimation in high-dimensions. Statistics in Medicine , 37.

Radcliffe, N. J. and Surry, P. D. (2012). Real-world uplift modelling with significance-based uplift trees.

Rosenbaum, P. and Rubin, D. (1983). Assessing Sensitivity to an Unobserved Binary Covariate in an Observational Study with Binary Outcome. Journal of the Royal Statistical Society. Series B (Methodological) , 45(2):212-218.

Rubin, D. (1974). Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of Educational Psychology , 66.

Scornet, E., Biau, G., and Vert, J.-P. (2015). Consistency of random forests. The Annals of Statistics , 43(4):1716 - 1741.

Shi, X., Miao, W., Nelson, J., and Tchetgen, E. (2020). Multiply robust causal inference with double-negative control adjustment for categorical unmeasured confounding. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 82.

Shi, X., Pan, Z., and Miao, W. (2022). Data integration in causal inference. WIREs Computational Statistics , 15(1).

Taddy, M., Gardner, M., Chen, L., and Draper, D. (2016). A nonparametric bayesian analysis of heterogenous treatment effects in digital experimentation. Journal of Business & Economic Statistics , 34.

Tang, C., Wang, H., Li, X., Cui, Q., Zhang, Y.-L., Zhu, F., Li, L., Zhou, J., and Jiang, L. (2022). Debiased causal tree: Heterogeneous treatment effects estimation with unmeasured confounding. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A., editors, Advances in Neural Information Processing Systems , volume 35, pages 5628-5640. Curran Associates, Inc.

Turney, K. and Wildeman, C. (2015). Detrimental for some? heterogeneous effects of maternal incarceration on child wellbeing. Criminology & Public Policy , 14.

Wager, S. and Athey, S. (2018). Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association , 113(523):1228-1242.

Wu, L. and Yang, S. (2022). Transfer learning of individualized treatment rules from experimental to real-world data. Journal of Computational and Graphical Statistics , 0(0):1-10.

Yang, S., Zeng, D., and Wang, X. (2020). Improved inference for heterogeneous treatment effects using real-world data subject to hidden confounding. arXiv preprint arXiv:2007.12922 .

Zhang, B. and Tchetgen Tchetgen, E. J. (2022). A Semi-Parametric Approach to Model-Based Sensitivity Analysis in Observational Studies. Journal of the Royal Statistical Society Series A: Statistics in Society , 185:S668-S691.

Zhang, H., Zimmerman, J., Nettleton, D., and Nordman, D. J. (2020). Random Forest Prediction Intervals. The American Statistician , 74(4):392-406.

A Additional figures

Estimators — HT － SF - ST — FAST

4.0

3.5

RMSE

3.0

2.5

2.0

0.0

0.5

1.0

1.5

2.0

β

B Pre-processing of the real-world data

In the STAR dataset, each of the pre-treatment covariate X j ( 1 ≤ j ≤ p ) was standardized to a range of -1 to 1 , and the outcome variable Y was standardized to a range of 0 to 100 .

C Proof of Theorem 1

The proof follows the similar arguments as in Györfi et al. (2002) and Scornet et al. (2015). It is sufficient to show the result at the root node given the recursive nature of the partitioning. We will use the following notations in the sequel. We denote E T , P T and E O , P O as the expectation and probability under trial data and observational data, respectively. We let Z = ( X, ˜ Y ) . For any q ∈ [ p ] and c ∈ R , let θ = ( q, c ) and the corresponding two partitioned notes are denoted as Q L ( θ ) = { x | x q ≤ c } and Q R ( θ ) = { x | x q > c } . The parameter space of θ is denoted as Θ = [ p ] × R . Let µ L and µ R be the predictions for Y on Q L ( θ ) and Q R ( θ ) , respectively and denote τ = ( τ L , τ R ) . Let M i ( θ ) and M i of ( θ ) be the MSEs of the conventional HTE estimator and the fused HTE estimator on the child nodes of Q i ( θ ) , respectively, for i ∈ { R,L } .

(i). For any θ ∈ Θ , according to Equation (2) in the main paper we have

for i ∈ { R,L } , where the weight w i ( θ ) satisfies

by Equation (5) in the main paper, which is lower bounded by σ 2 min σ 2 min + nb 2 max , where σ 2 min < Var( ˜ Y | X = x , S = 0) and b max = sup x ∈ Q j |{ E ( ˜ Y | X = x , S = 0) -E ( ˜ Y | X = x , S = 1) }|

APPENDIX

Therefore, we conclude that

which reveals the MSE reduction effect of the proposed split criterion.

(ii). The proof includes two parts. In Part 1, we will derive the bounds for the discrepancies between the MSEs under the empirically estimated split and the oracle split under the conventional criterion, and in Part 2 the similar results under the proposed split criterion.

Part 1. We define the following criterion function:

For i ∈ { L, R } , let

Then L i ( θ, τ i ) represents the MSE of τ i on the region Q i ( θ ) . For a given split θ = ( q, c ) , it is straightforward to see that the optimal τ ( θ ) = ( τ L ( θ ) , τ R ( θ )) is given by

for i ∈ { L, R } , which is the sample mean of Y on the region Q i ( θ ) . Therefore, by the definition of M i ( θ ) , it holds that L i ( θ, τ i ( θ )) = M i ( θ ) for i ∈ { L, R } . The optimal split θ 0 = ( q 0 , c 0 ) on the population level is defined via minimizing the profiled criterion function:

Define M i n ( θ ) = ℓ i n ( θ, τ i ( θ ) , R t n ) for i ∈ { L, R } and the empirical optimal split ̂ θ = ( ̂ q, ̂ c ) is defined via minimizing the sample criterion function:

Step 1 (Main error decomposition).

Now we will bound M ( ̂ θ ) -M ( θ 0 ) , which represents the discrepancy of the MSEs of the oracle and empirical split. To apply empirical process theories for stochastic error analysis, we will use a truncation argument. We let M i n,β n ( θ, π i , R t n ) = E n ( T β n ˜ Y -T β n π i ( θ )) 2 I ( X ∈ Q i ( θ )) and M i β n ( θ ) = E T { M i n,β n ( θ, π i ( θ ) , R t n ) } , where T β n x =: ( | x | ∧ β n )sign( x ) for any β n > 0 . Correspondingly, let M β n ( θ ) = M L β n ( θ ) + M R β n ( θ ) and M n,β n ( θ ) = M L n,β n ( θ ) + M R n,β n ( θ ) . Then we have the following error decomposition:

By the definition of ̂ θ , we have S 4 ,n ≤ 0 . In following steps, we will bound S 1 ,n , S 2 ,n and S 3 ,n , respectively. The truncation level β n is chosen as β n = β 0 log( n ) for β 0 ≥ 2 σ Y .

Step 2 (Bounding S 1 ,n ).

For any θ , it holds that

For T 1 ,n , by Cauchy-Schwarz inequality we have

where the second inequality is because E T ( ˜ Y 2 ) ≤ ∞ and E T { ̂ π 2 i ( θ ) } ≤ E T ( ˜ Y 2 ) / | Q i ( θ ) | . Since

therefore,

Since E T exp( σ Y | Y | 2 ) < ∞ and β n = β 0 log( n ) , we conclude that | S 5 ,n | ≲ 1 n . With the same argument, we have S 6 ,n ≲ 1 n , implying that

for any θ ∈ Θ . Therefore, the truncation error S 1 ,n ≲ 1 n .

Step 3 (Bounding S 2 ,n ).

Let M N,of = { f = ( T β n ˜ Y -T β n π ) I ( X ∈ Q i ( θ )) : θ = ( q, c ) ∈ [ p ] × R } . By applying Lemma 2 we obtain

where z n 1 is any set { z 1 , · · · , z n } ∈ [0 , 1] p ×Y and c > 0 is a universal constant. It follows from Lemma 1 that

Taking γ = 1 /n and α = ( t +log( pn )) β 4 n /n implies that with probability at least 1 -C 1 e -t for some universal constant C 1 > 0 ,

Therefore, we conclude that with probability at least 1 -C 1 e -t , the stochastic error S 2 ,n ≲ { t +log( pn ) log 4 ( n ) } /n .

Step 4 (Bounding S 3 ,n ). According to (15), we have

Since ˜ Y is sub-Gaussian by assumption, it is straightforward to see that ( T β n ˜ Y -T β n π i ( θ )) 2 I ( X ∈ Q i ( θ )) and ( ˜ Y -π i ( θ )) 2 I ( X ∈ Q i ( θ )) are sub-exponential for i ∈ { L, R } . Suppose ∥ ∥ ∥ ( T β n ˜ Y -T β n π i ( θ )) 2 I ( X ∈ Q i ( θ )) ∥ ∥ ∥ ψ 1 ≤ σ 0 and ∥ ∥ ∥ ( ˜ Y -π i ( θ )) 2 I ( X ∈ Q i ( θ )) ∥ ∥ ∥ ψ 1 ≤ σ 0 for all θ ∈ Θ , where ∥·∥ ψ 1 is the sub-exponential norm operator. By applying Bernstein's inequality, for any s > 0 , we have

for i ∈ { R,L } , where c > 0 is a universal constant. Taking s = σ 0 t cn = C 2 t , for any t ≥ 0 we obtain

for any n > t/c . Since the above result holds for any θ ∈ Θ , we conclude that for any t > 0 , with probability at least 1 -4 e -t , we have S 3 ,n ≲ ( t +1) /n .

Combining the results on S 1 ,n , S 2 ,n and S 3 ,n , we conclude that for any t > 0 , with probability at least 1 -C 3 e -t , it holds that

for some universal constants C 3 , C 4 > 0 .

Part 2. The proposed scale criterion can reformulated as follows. For i ∈ { L, R } , let

where τ 0 ,i ( θ ) = E n ( ˜ Y 0 | X 0 ∈ Q i ( θ )) and τ 1 ,i ( θ ) = E m ( ˜ Y 1 | X 1 ∈ Q i ( θ )) , and

where σ 2 u ( Q i ( θ )) = Var n ( τ 0 ,i ( θ )) , σ 2 b ( Q i ( θ )) = Var m ( τ 1 ,i ( θ )) and b ( Q i ( θ )) = τ 1 ,i ( θ ) -τ 0 , 1 ( θ ) . Let F s,i ( θ ) = E s ( F s,i ( θ )) for s ∈ { 0 , 1 } and F s ( θ ) = F s,L ( θ ) + F s,R ( θ ) , the population criterion is defined as M of ( θ ) = F 0 ( θ ) + F 1 ( θ ) . For the empirical criterion, we first define F n,i ( θ ) = E n ( F R i ( θ )) and F m,i ( θ ) = E m ( F R i ( θ )) . Let M N,of ( θ ) = F n,L ( θ ) + F n,R ( θ ) and F m ( θ ) = F m,L ( θ ) + F m,R ( θ ) , the empirical criterion is the denoted as M N,of ( θ ) = M N,of ( θ ) + F m ( θ ) . The population and empirical optimal splits are defined by

We first have the following error decomposition:

By the definition of ̂ θ of , we have T 4 ,n ≤ 0 . In the following steps, we will bound T 1 ,n , T 2 ,n and T 3 ,n , respectively. Following the same argument as for S 1 ,n , it can be obtained that T 1 ,n ≲ 1 n .

We now bound T 2 ,n Let G n = { g : g ( y, x ) = √ 1 -w ( θ )˜ y -τ ) I ( x ∈ Q ( θ ) , θ ∈ Θ n ) } and H n = { h : h ( y, x ) = √ w ( θ )(˜ y -τ ) I ( x ∈ Q ( θ ) , θ ∈ Θ n ) } , then via Lemma 2 we have

for any δ > 0 , where z n 1 is any set { z 1 , · · · , z n } ∈ [0 , 1] p ×Y and c > 0 is a universal constant. It follows from Lemma 1 that for any α 1 , γ 1 > 0

Taking γ 1 = 1 /n and α 1 = ( t +log( pn )) β 4 n /n implies that with probability at least 1 -C 4 e -t for some universal constant C 4 > 0 ,

Similary, for any α 2 , γ 2 > 0 ,

Taking γ 2 = 1 /n and α 2 = ( t +log( pn )) β 4 n /m implies that with probability at least 1 -C 5 e -t for some universal constant C 5 > 0 ,

Combining (19) and (20) delivers that with probability at least 1 -2 C 1 e -t ,

for ant t > 0 , since ̂ θ of , θ f ∈ Θ n and M of,β n ( θ ) = F 0 ,β n ( θ ) + F 1 ,β n ( θ ) and M N,of,β n ( θ ) = F n,β n ( θ ) + F m,β n ( θ ) for any θ ∈ Θ n .

Now we turn to T 3 ,n , the truncation error for the empirical loss. With the similar argument as in (16), we have with probability at least 1 -4 e -t , it holds that T 3 ,n ≲ ( t +1) /n +( t +1) /m for any t > 0 .

Combining the results for T 1 ,n , T 2 ,n and T 3 ,n , we conclude that for any t > 0 , with probability at least 1 -C 6 e -t ,

which completes our proof.

D Supporting lemmas

The following to lemmas are from Section 11.3 and Section 13.1 of Györfi et al. (2002), which are useful for our proofs.

Lemma 1. (Deviation inequality of quadratic process). Suppose that G is a class of uniformly bounded functions G = { g : R d → R ∥ g ∥ ∞ ≤ M } . Let F = { g 2 : g ∈ G } . Then for any n ≥ 1 , it holds that

where z n 1 = ( z 1 , · · · , z n ) ∈ R d , α, γ > 0 and 0 < ε ≤ 1 / 2 .

Lemma 2. (Covering number of piece-wise constant functions). Let Π be the family of partitions of [0 , 1] p . For any set x n 1 = { x 1 , · · · , x n } ⊂ [0 , 1] p , let ∆(Π) be the maximal number of partitions of x n 1 induced by elements of Π . Let M (Π) be the maximal number of sets contained in a partition P ∈ Π . Denote the piece-wise constant functions on [0 , 1] p be F (Π) with ∥ f ∥ ∞ ≤ β n for any f ∈ F (Π) . Then using Lemma 13.31 and Theorem 9.4 of Györfi et al. (2002) we have

for any δ > 0 , where c 1 > 0 is some univiersal constant. Specifically, in each partition for a node C k of a tree, we have M (Π) = 2 and ∆ n (Π) ≤ ( pa n ) 2 , where a n is the sample size of C k .

Peer reviews:
comments: SUMMARY:
This paper proposes a novel strategy for estimating the heterogeneous treatment effect called the Fused and Accurate Shrinkage Tree (FAST). The authors confirm the consistency of the proposed tree-based estimator and demonstrate the effectiveness of their criterion in reducing prediction error through theoretical analysis. The advantages of the proposed method over existing methods are demonstrated via simulations and real data analysis. As I am not very familiar with this field, it might be better to consider my opinion less.

STRENGTHS:
1. This paper is technically sound.
2. The proposed method has better performance than the existing methods.

WEAKNESSES:
Imcomplete references: The tree-based method seems to not be a new method in this field. There might be some other references as follows.

Agarwal, Abhineet, et al. "Hierarchical Shrinkage: Improving the accuracy and interpretability of tree-based models." International Conference on Machine Learning. PMLR, 2022.

Nasseri, Keyan, et al. "Group Probability-Weighted Tree Sums for Interpretable Modeling of Heterogeneous Data." arXiv preprint arXiv:2205.15135 (2022).

QUESTIONS:
None

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper deals with the problem of estimating the heterogeneous treatment effects with multiple data sources. In particular, the paper aims to utilize the information from the observational data to better estimate the causal effects in the trial data. Inspired by the shrinkage estimation, a weighting scheme is developed to balance the unbiased estimator based on trial data and the potentially biased estimator based on observational data. Specifically, a tree-based algorithm with new split criterion is proposed based on above motivations. Some theoretical results about the causal effect estimation is derived. Finally, the author provides simulations and real data analysis to demonstrate the performance of the proposed method.\n\nSTRENGTHS:\nThe papers deal with an interesting problem in practice, i.e., data fusion. In particular, we may have multiple data sources. However, some sources has limited observations with unbiased causal effects and other sources have sufficient observations with biased causal effects. The paper utilizes the tree-based algorithm with a new splitting criterion to tackle this issue. In addition, some theoretical analysis are provided to prove the advantages of the method.\n\nWEAKNESSES:\n1. Although the paper considers an important problem in practice, the reason why the method chooses tree-based algorithm is not convincing. In particular, other ML methods can also achieve data fusion. The advantages of using tree-based algorithm over other methods are confusing and are not clearly discussed. \n2. In the introduction and simulation studies, the author also mentions many other methods that deal with the data fusion problem. However, the reason why the proposed tree-based method can perform better than other methods are not interpreted. \n3. Although the paper provides the theoretical analysis for the causal effect estimation, the interpretation for the theorem are not convincing. For example, many other methods also have theoretical guarantee for the causal effect estimation. Which part in the theoretical analysis can illustrate the advantages of data fusion and the tree-based algorithm?\n4. In the general picture, the idea of data fusion is very similar to that of transfer learning, i.e., we want to transfer the information from the observational studies to help the estimation of causal effects in trial data. However, the paper does not mention any related literature in transfer learning. In particular, what is the advantages and differences of the method compared with transfer learning? A more comprehensive literature review is encouraged.\n\nQUESTIONS:\n1. When the bias of the causal effect estimation in observation studies is large, will the data fusion cause any harm to the estimation of the causal effects? Can it be interpreted as a trade-off between the bias and variance?\n2. In section 2.1, the author mentions $\\hat{e}(X,1)$ is unbiased while $\\hat{e}(X,0)$ is biased. Why is that? Didn't they come from the same data source with $\\hat{e}(X,1) + \\hat{e}(X,0) = 1$?\n3. One of the main questions is that, why using tree-based methods for data fusion?\n4. In real data analysis, the author mentions that the true causal effect is unknown and hence using the estimation of the generalized random forest as the ground truth. This is confusing. If the estimation using random forest is not accurate, then the results in the real data analysis is not reliable and convincing. How do the author correct this potential bias?\n\nLIMITATIONS:\nPlease see in Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a novel shrinkage method that fuses an unbiased estimator with a biased estimator. This method effectively reduces the MSE of the unbiased estimator. The approach offers a practical and straightforward implementation specifically tailored for estimating heterogeneous treatment effects. The authors extend the conventional node split criterion to align with the fused estimator and penalizes the use of observational data with substantial confounding bias. The authors also provide a theoretical analysis that explains the advantages of the modified splitting criterion.\n\nSTRENGTHS:\n- The application of the weighting strategy from shrinkage estimation to fusing unbiased and biased estimators in order to reduce the MSE of the unbiased estimator is a great idea.\n- The modification of the node-splitting criterion that aligns with the fused estimator is an excellent enhancement to the methodology.\n- The paper is well-organized and thanks to the authors's thoughtful and consistent notation, the methodology is easy to follow.\n\nWEAKNESSES:\nI'm concerned with the omission of $\\sigma^2_b$ in practice. If we only look at the weight $w$, it make sense if $\\sigma^2_b$ is small comparing to $\\sigma^2_u$. However, in the tree building process, we need working estimates of the MSE of fused estimator, $\\frac{(\\sigma^2_b+b^2)\\sigma^2_u}{\\sigma^2_b+b^2+\\sigma^2_u}$, and I don't think omitting $\\sigma^2_b$ is justified by the same reason anymore.\n\nQUESTIONS:\nThe sample sizes considered in the experiment are always large. If a sample from RCT of size 100 or 200 is avaliable, the baseline models would do the job. I'm more curious in the senario when we have a relatively small RCT sample and that we do need more data from obervational studies to enhance the estimation.\n\nLIMITATIONS:\nNo other limitation.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a Fused and Accurate Shrinkage Tree (FAST) algorithm for heterogenous treatment effect estimation given trial and observational data. The FAST algorithm introduces (i) a shrinkage based approach that combines trial and observational data for MSE reduction in treatment effect estimation, and (ii) a split criteria which down-weights observational data with high confounding bias. Further, the paper provides theoretical analysis demonstrating the benefits of the proposed split criteria. Experimental results on synthetic and real-world data demonstrate that the proposed approach outperforms baselines per metric MSE.\n\nSTRENGTHS:\n- The paper is well written and easy to follow. The reviewer enjoyed reading this paper.\n- The proposed FAST algorithm is well-justified and the theoretical analysis might be of interest to some readers. \n- The paper tackles an important problem (fusing small RCT with large readily available observational data ) with many applications. \n- The algorithm seems simple and easy to implement\n\nWEAKNESSES:\n-  Eqn. 5:  The paper seems to have glazed over the rationale for dropping $\\sigma_b$ in the shrinkage estimator. It's unclear in what scenarios, e.g., how large the observational sample size must be for $\\sigma_b$  to become negligible.\n\nQUESTIONS:\n**Experiments**\n- Table 1: Could you include results for the NF estimator (using trial data only) and SF estimator (both trial and observation data)\n- Figure 3: Could you comment on why NF is not monotonically decreasing with sample size? Shouldn't we expect the rFAST to NF estimators to cross when $n$ is large enough?\n\n**There are some limitations with tree-based methods that the paper does not adequately address:**\n- How sensitive is the FAST algorithm to the dimension $p$ and heterogeneity of the covariates? The experiments set $p=5$, which is unrealistic in real-world scenarios.\n- What is the complexity, e.g., training time of the rFAST algorithm compared to baselines, and how easy is it to scale given high-dimensional heterogeneous covariates?\n\n\nMinor:\n- Figure 1: Could you add more details summarising the plots in the caption\n\nLIMITATIONS:\nThe limitation discussion is inadequate. I encourage the authors to add a paragraph discussing the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the problem of estimating (heterogeneous) treatment effects via both interventional and observational data. The authors proposed a new estimation, namely the Fused and Accurate Shrinkage Tree (FAST), which optimally weights the interventional and observational estimator, and combines with a new spilt criteria for tree-based heterogeneous treatment effect estimation. The authors further conducted experiments to compare FAST against existing methods.\n\nSTRENGTHS:\n- Apart from a few typos, the paper is well written and ideas are presented in a rigorous but clear also way. \n- The idea of applying shrinkage method to combine interventional and observational data for better estimator is novel and could be a nice addition to the literature.\n- Note that I have not gone through all the proofs in the appendix, the mathematical correctness might need further input from other reviewers.\n\nWEAKNESSES:\nGenerally I like this paper, but there are still a few weaknesses.\n\n- The main issue with shrinkage method is interpretability: we need to understand better how the variance-bias trade-off behaves in different regimes, especially the authors takes a more analytic way to first estimate the required quantities, then solve the optimal weights. More specifically for example, it would be beneficial to at least see different how trial mechanisms affects the estimation. For example, in a more realistic setting, one may consider non-randomized trials rather than RCT, in which treatments are assigned by a *known* true model. By adjusting the parameters of such true assignment model, the estimator variance for the trial population HTE estimator can be controlled (even with fixed N). Then the performance of FAST can be evaluated against different variance regimes of trial HTE estimator, which will help us understand the sweat spot of the method.\n\n- Regarding experimental settings. It is indeed quite standard for these type of papers to have 1 or 2 synthetic experiments and 1 real data experiment. However, in the case of this paper I found the simulation setting is a bit weak. It would be great to perform experiments on multiple data generating mechanisms with randomly sampled parameters and coefficients, allowing us to evaluate the marginal performance of the method. Otherwise the authors at most demonstrated the capability of the method on only one single data generation mechanism (which arguably is much easier to hack/cherry-pick).\n\n\n- The other potential room for improvement is the baseline. I understand that the paper mainly only compares to data fusion methods. However, due to the variance-bias trade-off of the shrinkage method, it would be natural to also expect some comparisons to variance reduction methods for trial estimators as well.\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nThe authors has somewhat discussed the limitations of the method.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-01"}