Year: 2021

Abstract:
Heterogeneous treatment effect (HTE) estimation with continuous treatment is essential in multiple disciplines, such as the online marketplace and pharmaceutical industry. The existing machine learning (ML) methods, like forest-based modeling, either work only for discrete treatments or make partially linear or parametric assumptions that may suffer from model misspecification. To alleviate these problems, we extend causal forest (CF) with non-parametric dose-response functions (DRFs) that can be estimated locally using kernel-based Double/Debiased ML estimators. Moreover, we propose a distance-based splitting criterion in the functional space of Partial DRFs to capture the heterogeneity for continuous treatments. We call the proposed algorithm generalized causal forest (GCF) as it generalizes the use case of CF to a much broader setup. We show the effectiveness of GCF compared to SOTA on synthetic data and proprietary real-world data sets.

text: Heterogeneous treatment effect (HTE) estimation with continuous treatment is essential in multiple disciplines, such as the online marketplace and pharmaceutical industry. The existing machine learning (ML) methods, like forest-based modeling, either work only for discrete treatments or make partially linear or parametric assumptions that may suffer from model misspecification. To alleviate these problems, we extend causal forest (CF) with non-parametric dose-response functions (DRFs) that can be estimated locally using kernel-based Double/Debiased ML estimators. Moreover, we propose a distance-based splitting criterion in the functional space of Partial DRFs to capture the heterogeneity for continuous treatments. We call the proposed algorithm generalized causal forest (GCF) as it generalizes the use case of CF to a much broader setup. We show the effectiveness of GCF compared to SOTA on synthetic data and proprietary real-world data sets.

heading: 1 INTRODUCTION
text: Heterogeneous treatment effect (HTE) estimation has been of growing interest for decision-makers in a wide spectrum of contexts. It uncovers the effect of interventions at sub-group levels, thereby providing highly tailored suggestions rather than a one-size-fits-all policy. When it comes to precision medicine, HTE provides the leveraged information that physicians need to precisely treat different patients with proper dosages of drugs depending on their genes, living habits, and EHR history. Recently, with the emerging big data that exhibits astronomical complexities, HTE has been essential for online marketplaces (Syrgkanis et al., 2021; Du et al., 2019; Ye et al., 2018). For largescale data that presents a huge challenge, algorithms that adapt the booming machine learning (ML) techniques to HTE estimation are proposed (Künzel et al., 2019; Zhao et al., 2017; Chernozhukov et al., 2018; Hill, 2011). However, these methods largely focus on binary or discrete treatments and thereby being not widely applicable, given the fact that continuous treatments are prevalent in practice. Examples include targeting customers with continuous incentives, customizing the duration of ads to improve user engagement, and design optimal interest rates to maximize revenue while control for default.
For HTE estimation with continuous treatments, it could be feasible to discretize the treatments and utilize the above approaches, but the limitations are two-phase and bring non-trivial challenges to researchers. First, ordinal treatments like different levels of education cannot be treated as categorical otherwise the order information across treatments will be lost, and secondly, discretizing or bucketing has the limitations of model inaccuracy and cannot recognize the right patterns of the outcome across treatments. To elucidate the complex relationship between treatment and outcome, researchers put forward ideas of dose-response function (DRF).
The estimation of DRF in the presence of high dimensional features can be viewed as a natural extension of regression problems with additional complexities. The success of random forest (RF) for regression problems motivates the development of causal forest (CF) (Athey et al., 2019) and orthogonal random forest (Oprescu et al., 2019) and the modification of Bayesian Additive Random Forest (Woody et al., 2020; Hahn et al., 2020). These methods partition the feature space with a splitting criterion fitted on treatments that can provide accurate estimations for HTE. But often
they posit linear or partially linear models of DRF which are vulnerable to model misspecifications. Figure 1 is an illustration of model misspecification of random forest and CF.
Doubly Robust estimators (DREs) allows for misspecification of either the treatment or outcome regression. Recent work (Kennedy et al., 2017) combines the estimation of generalized propensity score with regression models of the outcome that can further reduce the bias and use kernel regression for non-parametric estimation of DRF. However, it may suffer from the curse of dimensionality when the feature space is high-dimensional given large-scale instances. To advance, (Colangelo & Lee, 2020) provide kernel-based DML estimators that can easily deal with high-dimensional data. Nevertheless, the localization and weighting mechanism of CF is not taken into account, which is essential for capturing heterogeneity. It remains unexplored that how to effectively combine ML methods with localized estimations under no model assumptions.
In this paper, we propose generalized causal forest (GCF) by extending CF to a generalized one with a brand new splitting criterion, and herein the generalization is two-fold. First, we generalize the local linear model of CF to a non-parametric one with DRF. Meanwhile, we employ the distance metrics in the functional space of Partial DRF as the splitting criterion and use kernel-based DML estimators as the DRF approximation.
When it comes to implementations, many open source packages integrate recent research with the trending ML methods that facilitate HTE estimation, such as EconML (Oprescu et al.), CausalML (Chen et al., 2020) and GRF. EconML, CausalML offer a suite of high-performance ML algorithms like DML, while GRF specially supports CF. However, they are only compatible with single-machine systems and may not be efficient when it comes to large-scale instances. Apache Spark (Meng et al., 2016) is integrated for handling large-scale data with distributed and in-memory computing, while inclusively provides APIs of trending ML algorithms with the built-in MLlib. Therefore, as one of our contributions, we build our algorithm on Spark that enables large-scale data processing and easy use of any ML techniques. To summarize, the main contributions are:
• We are the first to generalize CF’s partially linear model to a non-parametric one with DRF and apply kernel-based DML estimators. The proposed algorithm pictures the nonparametric behaviors of how outcome varies with treatments.
• We propose a distance-based splitting criterion with various distance metrics in the functional space of Partial DRF that generalizes the differences of CAPE in CF.
• Our proposed GCF has been tested on both synthetic and large-scale real-world datasets. It improves on SOTA in terms of multiple evaluation metrics. We implement GCF on Apache Spark MLlib and achieves higher computatinal efficiency by distributed computing.
The rest of the paper is organized as follows. In Section 2, we present related works in the literature. In Section 3, we introduce some notations and backgrounds. In Section 4, we introduce the Generalized Causal Forest. In Section 5, we examine the empirical performance of GCF by applying it to both synthetic and real-world data sets. We conclude the paper with some discussions in Section 6.

heading: 2 RELATED WORKS
text: A growing amount of literature has been devoted to address the problem of HTE estimation with continuous treatments. The algorithms in the context of DRF include methods targeting at confounding bias, kernel-based or ML-based methods for regression bias reduction, and techniques of DRE or DML that balance the trade-off between the two biases. DRF combined with IPW (Zhu et al., 2015; Graham et al., 2015) can achieve consistent estimators by weighting the estimation with probability density of treatment. For the regression bias, (Galagate & Schafer, 2015) employ parametric estimations and (Flores et al., 2007) model DRF as non-parametric functions using kernel regression. To advance, (Kennedy et al., 2017) propose a doubly robust estimators for DRF by combine the estimation of generalized propensity score with the estimation of outcome using kernel regression. Thus far, these approaches only provide global estimations and cannot handle massive amounts of high-dimensional data.
Recently, the great efficiency of ML methods motivates their generalization to the problem of HTE estimation. Towards that end, estimating DRF with ML-based algorithms is developed and among which tree-based models are a great candidate since it partitions the feature space for dimensionality reduction and maximizes the heterogeneity as well. Causal Forest (CF) proposed in (Athey et al., 2019) utilizes a subset of training samples for growing trees by recursively partitioning via a splitting criterion. Then HTE estimation is given by weighted average over the outcomes of the remaining training samples, known as honesty principle. The final estimator obtained from CF further exhibits a lower degree of bias by tree ensembles. However, their splitting criterion relies on a linear model assumption and is formulated as the difference between the slopes of linear models. The same limitation applies to Orthogonal Random Forest (Oprescu et al., 2019) and Bayesian Additive Regression Tress (Hill, 2011; Woody et al., 2020; Hahn et al., 2020). Nevertheless, the complexity of HTE estimation may not be fully captured by linear or general parametric models, which brings the necessity of non-parametric ML models.
Kernel-based DML (Colangelo & Lee, 2020) estimates the nuisance functions with cross-fitting, constructs a non-parametric DML estimator by Gaeutax Derivative. This motivates the utilization of DML and non-parametric estimation in our work, though it only provide global estimators with the limited capacity of localization. A fully non-parametric DML with locally weighted estimations for continuous treatments has not been considered yet.
Building on prior art, we point out that our integration of DRF estimation into CF overcome the challenges aforementioned by using local non-parametric DRF to constructing splitting criterion. We introduce the conceptual partial DRF as a component of the splitting criterion and employ the distance in the functional space of Partial DRF as a proxy for heterogeneity instead of the difference of slopes. Moreover, the partial DRF can be estimated precisely by the robust kernel-based DML estimators in our splitting criterion.

heading: 3 PREREQUISITES

heading: 3.1 NOTATIONS AND ASSUMPTIONS
text: We formally introduce the notations for HTE estimation with continuous treatments. Following the potential outcome framework in (Neyman, 1923; Rubin, 1974), we let T be the continuous treatment, X = (Xj)pXj=1 be the pX -dim confounding variables, U be the outcome-specific covariates, Z be the treatment-specific covariates that are independent of U , and Y be the outcome of interests. The population Ω = (X,U ,Z, Y, T ) ∈ RpX+pU+pZ+1+1 satisfies
Y = g(T,X,U) + ;T = f(X,Z) + ν
where , ν are noises of zero mean and g : Rp × R→ R and f : Rp → R. {(Xi,Ui,Zi, Yi, Ti), i = 1, . . . , n} are i.i.d. samples drawn from the population Ω. In practice, the decomposition ofX,U ,Z from observed covariates is difficult. Therefore, without additional specification, we use X to represent observed covariates throughout the paper to simply notations. The potential outcomes under treatment t is Y(t). Recall that Propensity Score (PS) (Rubin, 1974) for a discrete treatment is P (T = t|X), the probability for a unit receiving treatment t given the covariates X . For continuous treatments, (Rubin, 1974; Hirano & Imbens, 2004) introduce Generalized Propensity Score (GPS) which is the probability density function π(T = t|X).
The estimand of interests, CATE θ(t,X), is formally defined as
θ(t,X) = E[Y(t)|X]− E[Y(0)|X]
To identify θ(t,X) , common assumptions as in (Holland, 1986; Kennedy et al., 2017) are made throughout the paper.
Assumption 1. Consistency: E[Y |T = t] = E[Y(t)|T = t], i.e. the outcome of any sample solely depends on its treatment;
Assumption 2. Ignorability: The potential outcomes Y(T ) is independent of treatment T given covariatesX .
Assumption 3. Positivity: The GPS π(T = t|X) > pmin > 0, ∀t,X , i.e. the density for any sample receiving any treatment is bounded away from 0.
Under the above assumptions, we have
θ(t,X) = E[Y(t)|X]− E[Y(0)|X] = E[Y |T = t,X]− E[Y |T = 0,X] = E[g(t,X)|T = t,X]− E[g(t,X)|T = 0,X]
where the first equality holds by Assumption 1 and the second equality holds by Assumption 2. Positivity is indispensable for the conditional expectation in the last line to be well-defined whereas being often too strong. In practice, it can be reduced to the following.
Assumption 4. Weak Positivity: The variance of GPS σ(π) = ∫ t t2 · π(T = t|X)dt −(∫
t t · π(T = t|X)dt )2 > σmin > 0, ∀t,X .

heading: 3.2 DOSE-RESPONSE FUNCTION
text: For continuous treatments, the treatment effect can be characterized by DRF, formally defined as
µ(t) := E[Y |T = t] (1)
Generally speaking, DRF can be either parametric or non-parametric, with parametric ones being applied to CATE estimation as in (Robins, 2000; Van der Laan et al., 2003). Non-parametric ones can better explore the functional space of ADRF out of flexibility and hence being the tool of our proposed algorithm.
The target estimand of DRF heavily depends on the contexts (Galagate, 2016), such as treatment effect ( E[Y(t)]− E[Y(0)] ) , partial effect
( E[Y(t2)]− E[Y(t1)]
t2 − t1
) and the elasticity ( ∂log(E[Y(t)])
∂log(t)
) .
Specifically, conditional DRF (CDRF) µ(t,X) = E[Y |T = t,X] characterizes the treatment effect at individual levels and thereby leading to a CATE estimator as θ(t,X) = µ(t,X)− µ(0,X), which is the target estimand hereafter.

heading: 3.3 KERNEL REGRESSION AND DOUBLE/DEBIASED MACHINE LEARNING
text: For non-parametric modeling, kernel regression (Fan & Gijbels, 2018) works with solid theoretical foundations. Specifically, given a non-parametric model y = g(x) + with data (Xi, Yi), kernel regression with linear kernel smoothers gives the estimator ĝ(x) = ∑n i=1Kh(x,Xi) · yi∑n i=1Kh(x,Xi) where Kh(x,Xi) is a scaled kernel density with h being the bandwidth for smoothness. Typical choices of kernel include Uniform, Epanechnikov, Biweight, Triweight and Gaussian.
The kernel-based DML estimator proposed in (Colangelo & Lee, 2020) performs a 2-stage DRF estimation maintaining asymptotically normal. The DML estimator for the DRF µ(t) is proposed as follows. First, we estimate the CDRF µ(t,X) with µ̂(t,X) and the GPS π(T = t|X) with π̂(t|X). Then, with the nuisance functions plugged in, kernel-based double/debiased estimator gives
µ̂(t) = 1 n · n∑ i=1 µ̂(t,Xi) + Kh(Ti − t) π̂(t|Xi) · (Yi − µ̂(t,Xi))

heading: 4 GENERALIZED CAUSAL FOREST
text: In this section, we formally present the details of our algorithm, GCF. Building on CF, GCF relaxes partially linear assumption by adapting non-parametric DRF to the splitting criterion that can be estimated using kernel-based DML estimator described in the previous section. In what follows, we give it to a workflow of GCF at both the training stage and prediction stage, followed by the elaborations on the details of splitting criterion, estimations, and practical tweaks for implementation.
Algorithm 1: Generalized Causal Forest Input Data set O = (Xi, Ti, Yi), i = 1, . . . , n; number of trees B; number of features sampled
for growing a tree mtry; the minimum sample size on each leaf node min.node.size; honesty fraction α; tolerance τ ; positivity threshold ζ;
Trainings begin Data O1 and data O2 by a (α,1− α) split rule applied to O for honesty; Pre-train outcome regression model µ̂ and treatment density estimation π̂ on sample Ω1; b←− 1; while b ≤ B do
Randomly sample a feature set of size mtry XS fromX; for The stopping rule is not satisfied do
Identify each parent node P ; Compute the splitting criterion ∆(·) (3) over the samples in O1 withXS ; Grow the tree Tb by splitting at the parent node P according to the ∆(·);
end Assign samples in O2 to leaf nodes based on Tb and let b←− b+ 1;
end end Output Causal forest with B trees by recursive partitioning on X; node assignments for samples in O2; Predictions begin CDRF estimation µ̂b(t,x) by local weighted average over the outcomes of samples in Ω2
that falls into Lb(x) as ∑n i=1
1Xi∈Lb(x) · 1Ti=t · Yi |Lb(x)| ; CATE estimation
θ̂b(t,x) = µ̂b(t,x)− µ̂b(0,x) and θ̂(t,x) = 1
B
∑B b=1 θ̂b(t,x);
end
We introduce the partial derivative of DRF, PDRF at treatment twith a slightly abuse use of notations
Φ(t) := ∂µ(T )
∂T ∣∣∣∣ T=t and Φ(t,X) := ∂µ(T,X) ∂T ∣∣∣∣ T=t .
Intuitively, it tells about how µ(·) varies with treatments and characterizes the ”slope” along t. Different leaf nodes on a tree shall have different slopes as the heterogeneity of covariatesX .
Let the number of trees be B and then GCF grows trees Tb, b = 1, . . . , B, by repeating tree-growing process B times with bootstrapping. At the training stage, with PDRF we construct the tree Tb by recursive partitioning based on a splitting criterion ∆(c1, c2). It is proportional to the discrepancy between the PDRF in the left and right child nodes, to maximize the heterogeneity. We elaborate on this splitting criterion in the next part.
Our algorithm is implemented on Apache Spark for large-scale data processing and the mechanism of the tree-growing process is different from that in CF. To elucidate on that, the data is stored at the master machine and trees are copied to each branch machine. Then data is randomly distributed to branch machines for the computation and collected back to the master machine for the final criterion that is used for updating a tree. The updated tree will be copied to each branch machine again and we repeat the process until the tree stops splitting. This distributed computation and clustered storage leverage the computational power of multiple machines for speeding up the training process. We refer readers to the appendix for more information on GCF’s Spark implementation.
The splits or the tree-growing process terminate when various pre-specified stopping criteria are met, for example, the sample size of child nodes is smaller than the parametermin.node.size or the distance D between PDRF is smaller than min.info.gain τ .
For prediction, we also follow the honesty principle proposed in (Athey et al., 2019) by partitioning the training samples into two parts where one is for constructing the tree and the other is for estimating CATE on the leaf node. That is to say that each training data can either be utilized to estimate CATE or contribute to growing a tree. At the prediction stage, the leaf node of given samples x on each tree Tb is denoted by Lb(x). The estimation of CDRF µ̂b(t,x) on Lb(x) takes a local
weighted average on training samples in Lb(x), which is given by ∑n i=1
1Xi∈Lb(x) · 1Ti=t · Yi |Lb(x)| and
θ̂b(t,x) = µ̂b(t,x) − µ̂b(0,x). The final CATE estimation for X is the average of Lb(x) over B trees, as θ̂(t,x) = 1
B
∑B b=1 θ̂b(t,x).

heading: 4.1 SPLITTING CRITERION
text: Given a parent node P and training samples with covariates Xω , our splitting criterion ∆(c1, c2) is proposed as
∆(c1, c2) = nc1nc2 n2P D(Φ̂c1(t,Xω), Φ̂c2(t,Xω)), (2)
where Φ̂c1 , Φ̂c2 are the PDRF estimators, respectively. The sample sizes of parent node, left child node and right child node are nP , nc1 and nc2 , respectively. The ratio
nc1nc2 n2P in ∆(c1, c2)
is to balance the sample sizes between the two child nodes. Distance metric D measures the distance between Φ̂c1(t,Xω) and Φ̂c2(t,Xω) in the functional space of PDRF, which is to be maximized in the splitting so as to optimize the heterogeneity. Some commonly used metrics are L1, L2 and L∞ as in (Dette et al., 2018) D1 = ∫ t |Φ̂c1(t,Xω) − Φ̂c2(t,Xω)|dt,D2 =∫
t |Φ̂c1(t,Xω) − Φ̂c2(t,Xω)|2dt,D∞ = maxt|Φ̂c1(t,Xω) − Φ̂c2(t,Xω)|. We next show how
to estimate PDRF Φc1 and Φc2 using kernel-based DML estimator (Colangelo & Lee, 2020).

heading: 4.2 PDRF ESTIMATION
text: The estimation for DRF is a precondition for the estimation of PDRF. In this paper, we point out kernel-based DML estimator in (Colangelo & Lee, 2020) can be adapted to the DRF estimation in our GCF and guides the tree splitting.
Recall that, the DML estimation for DRF is given by µ̂(t) = ∑ i∈C ( µ̂(t,Xi) + Kh(Ti − t) π̂(t|Xi) · (Yi − µ̂(t,Xi)) )
where µ̂(t,Xi) is the pretrained estimator for µ(t,X) atXi, π̂(t|Xi) is the pretrained estimator for GPS atXi and C denotes the child node.
When adapted to our framework, we additionally derive the following estimator for CDRF by innovatively combine kernel regression with the DML estimator.
CDRF estimator The estimator is given by
µ̃(t,Xi) = µ̂(t,Xi) + Kh(Ti − t) π̂(t|Xi) · (Yi − µ̂(t,Xi))
PDRF estimator The PDRF estimator is given by Φ̂(t) = ∂µ̃(t)∂t and Φ̂(t,x) = ∂µ̃(t,x) ∂t .
In estimating PDRF, the smoothness condition of DRF plays an important role. For smooth DRF, we can explicitly derive the closed-form PDRF. For smooth Gaussian kernels, Φ(t) is
Φ̂(t) = ∂µ̃(t)
∂t =
∑n i=1K
′ h(Ti − t)Yi∑n i=1Kh(Ti − t) − [ n∑ i=1 Kh(Ti − t)Yi ] ∑n i=1K ′ h(Ti − t) [ ∑n i=1Kh(Ti − t)]2
However, for general DRF, the PDRF may be discontinuous or even ill-defined. For example, nondifferentiable kernel functions like uniform, epanechnikov, biweight and triweight, Φ can only be estimated via numerical approximations of Φ(t) = E[µ(t+ δ)− µ(t)]
δ .

heading: 4.3 PRACTICAL CONSIDERATIONS
text: In the work of (Kennedy et al., 2017), they claim that the choice of bandwidth h for kernels weighs more than the choice of density functions, since h deals with the trade off between bias and variance. Usually, a small h avoids a large variance while a large h reduces the bias. The empirical ways of choosing an optimal h are Cross Validation and Rule-of-thumb (Dehnad, 1987)
When utilizing kernel densities for weighting, we often suffer from the boundary bias if without normalization. To this end, we normalize estimators with the cumulative density function of range [Tmin, Tmax] of treatments. More specifically, the estimators are divided by ∫ Tmax Tmin Kh(Ti − t)dt.
Regarding the positivity assumption in practice, we use a hyperparameter ζ to regularize for the boundary of σ(π) that is the variance of GPS. This is also a guarantee on the weak assumption as stated in Assumption 4. Formally, our proposed splitting criterion will be a regularized one as
∆(c1, c2) = nc1nc2 n2P D(Φ̂c1(t,Xω), Φ̂c2(t,Xω)) + ζ · σ(π̂), (3)
where π̂ is the estimated GPS. With this additional term, the splitting will tend to maximize the variance of GPS and hence taking the variety of treatments across samples into account.

heading: 4.4 ASYMPTOTIC PROPERTY
text: Let Assumptions 1-3 in (Colangelo & Lee, 2020) hold and thusfar the convergence of kernel-based DML estimator is naturally guaranteed at a rate of √ n as long as the outcome is unbiased or the generalized propensity score is unbiased. That is to say that we can utilize a subsample s of order√ n to achieve a converged kernel DML estimator. Then s satisfies that s/n → 0 and s → ∞ which is stated as Specification 1 in (Athey et al., 2019). Under the above assumptions and suppose one additional condition be that the score function is lipschitz with a bounded support of treatment, then Assumption 1-6 in (Athey et al., 2019) automatically hold for the score function. Throughout our algorithm, we also follow the honesty principle. Then by Theorem 3 (Athey et al., 2019), our final estimator is asymptotically unbiased. This shows the doubly robustness of our estimator by combining kernel DML and Causal Forest under slightly stronger assumptions on the score function.

heading: 5 EXPERIMENT
text: In this section, we validate GCF by showing that 1) GCF can provide CATE estimations in a variety of contexts such as simulated data with true labels and real-world data without the ground truth 2) GCF is flexible yet effective compared to SOTA which provides enough evidence for its practical use. More specifically, we present the detailed experimental results of our proposed algorithm and demonstrate its effectiveness compared to the prior art on synthetic data and real-world data sets.

heading: 5.1 EVALUATION
text: In what follows we introduce the commonly used metrics for model evaluation on both synthetic data and real-world data sets, named as model errors and Qini score and Qini curve, respectively. For real-world data sets, the model errors are not available since the counterfactual outcomes in the potential outcomes cannot be observed.
Synthetic Data We adopt PEHE and RMSE as the measure for model precision. Formally, let n be the number of treatment groups, θ̂ti and θ t i be the predicted and true treatment effect of ith
sample at treatment t, θ̂t = ∑t i=1 ŷ t i − ŷ0i , ŷ0i be the reference outcome that is pre-determined.
PEHE,RMSE,and ADRF are ∑n i=1(|θ̂t − θt|))
n ,
√∑n t=1(θ̂ t − θt)2
n , ˆy(t) =
∫ x µ̂(t, x)dx.
Real-world data sets We refer readers to (Gutierrez & Gérardy, 2017) for the formal definition of Qini Score and Qini Curve. In practice, a larger Qini Score indicates that model has better performance in identifying HTE.

heading: 5.2 MODEL SETUPS
text: The baselines for GCF include RF, CF, and Kennedy (Kennedy et al., 2017). Both RF and CF use num.trees equals to 500, min.node.size to 50. For Kennedy (Kennedy et al., 2017), we use the code of function cts.eff from R package npcausal available in (ehkennedy, 2017). SuperLearner library is set to SL.ranger and SL.glm. We made small tweaks to SL.ranger to make share same hyperparaters with RF and CF for comparability.
GCF We implement the proposed GCF on Apache Spark MLlib (Meng et al., 2016) with kernelbased DML estimator, denoted by GCF. For the implementation of GCF, we fix the distance metric D to be d2 out of simplicity. The number of trees, the minimum node size and other hyperparameters are set to match those in RF and CF.
RF is a forest-ensemble method without controlling confoundedness. CF makes the partially linear assumption on DRF. Kennedy is a non-parametric doubly robust algorithm but uses a normal splitting criterion. The comparisons inclusively validate our algorithm for both non-parametric DRF estimation and specially designed splitting criterion.

heading: 5.3 SIMULATION
text: Here we show a comprehensive simulation study on synthetic data. First, we elaborate on the data generating process. It is followed by the details of model implementations. Finally we illustrate the numerical results that consistently show the superior performance of our GCF.

heading: 5.3.1 DATA GENERATING PROCESS
text: Let n denote the number of samples and p = pX + pU + pZ be the dimension of covariates. The covariate matrix is (Xji ) j=1,...,pX i=1,...,n = (X1, . . . ,Xn) ∈ Rn×pX . Data generating process (DGP) is
Y = µ(T ) + 0.2(X1 ·X1 +X4)T +X · βX +U · βU + T = 20 ·Ψ(φ(X · βX +Z · βZ)) + ν
where φ is the sigmoid function and Ψ is the pdf of Beta distribution. Here we set DRF µ to be polynomial(Poly), exponential(Exp), and sinusoidal functions(Sinus).
The above DGP gives us a total of 6 datasets by taking the possible combinations of 3 DRFs µ(t) and 2 setups of covariate Σ. Noises , ν follow Unif(−1, 1). Covariates X and coefficients β follow N(0, IpX ). We also allow sparsity in covariates by randomly setting some of the coefficients to 0. Testing data sets are generated with randomly assigned treatments to allow unbiased evaluation.

heading: 5.3.2 SIMULATION RESULTS
text: Since some models output the final prediction of ˆY (t) and some output the treatment effect ˆθ(t), we deliberately compute the pseudo outcome ˜Y (t) = Y (t) − Y (0) as the label in the purpose of the comparability across baselines. The simulation results pf 3 datasets are summarized in Table 1 and those of the other 3 datasets are in Appendix. Here PEHE and RMSE are averaged and standard deviations of the metrics across simulation runs are also reported. Overall, GCF significantly outperforms baseline methods, which again validates the superiority of the newly proposed algorithm. More specifically, GCF exhibits the smallest biases across multiple DGP setups which we defer details to the supplementary.
Additionally, we compare the ADRF curve given by different models to the ground truth. An example of ADRF under a DGP specification is shown in Fig. 2.

heading: 5.4 REAL-WORLD DATA SETS
text: Our method is also implemented on real-world data sets with size of 10,698,884 which sourced from a leading ride-hailing company. A discount that randomly sampled (one of 6 different values {d0, d1, d2, d3, d4, d5 : d0 < d1 < . . . < d5}) is assigned to all trips in a given origin-destinationtime tuple. The effect of discount on the trip complete rate is of primary interests. We apply GCF to this data set and evaluate the performance of models by calculating Qini Scores.
The Qini scores of different models are summarized in Table 2. The performance of GCF is superior than the company’s current models as GCF has the highest Qini scores across all levels of discounts. We further create an incentive policy based on GCF output and deploy it for online purpose. Online A/B testing shows a significant lift in KPI for GCF-driven policy.

heading: 6 CONCLUSION
text: In this paper, we propose a novel forest-based non-parametric algorithm to address the problem of HTE estimation with continuous treatments. Under the situation that a fully non-parametric and localized ML-based algorithm has not been proposed yet, we extend Causal Forest with a DRFbased splitting criterion computed by the distance in the functional space of PDRF. To estimate this DRF, We use the kernel-based DML estimator to guarantee double robustness that mitigates model misspecifications. We implement GCF on Spark to leverage the computational efficiency and test it on both synthetic data and real-world data sets and compare it with the prior art approaches. The numerical results demonstrate that our method significantly outperform competing methods.

heading: 7 APPENDIX
text: The appendix runs as follows. First, we list all the notations used in the appendix. Then we dive into the details of the data generating process for the simulation study and present additional simulation results which further show the effectiveness of our proposed method. Subsequently, we briefly outline GCF’s Spark framework, which enables large-scale instances processing and faster model training. Finally, we demonstrate how to apply our method to an online schema and achieve superior performance as its another validation.

heading: 7.1 NOTATIONS
text: We restate the notations that are consistent with the main paper.
Following the potential outcome framework in (Neyman, 1923; Rubin, 1974), we let T be the continuous treatments, X = (Xj)pj=1 be the pX -dim confounder, U be the pU -dim outcome-specific adjustment variable, Z be the pZ-dim treatment-specific adjustment variable, and Y be the observed outcome.The potential outcomes under treatment t is Y(t). The population Ω = (Σ = (X,U ,Z), Y, T ) ∈ RpX+pU+pZ+1 satisfies
Y = g(T,X,U) +
T = f(X,Z) + ν
where , ν are noises of zero mean,
and g : RpX+pU × R→ R and f : RpX+pZ → R.
{(Σi, Yi, Ti), i = 1, . . . , n} are i.i.d. samples drawn from the population Ω. Then the covariate matrix Σ = (Σji ) 1≤j≤p 1≤i≤n = (Σ1, . . . ,Σn) ∈ Rn×p where p = pX + pU + pZ . The Generalized Propensity Score is π(T = t|X), which is the probability density for a unit receiving treatment t given the covariateX .

heading: 7.2 DATA GENERATING PROCESS
text: Recall that the covariate matrix Σ = (Xji , U j i , Z j i ) = (X1, . . . ,Xn,U1, . . . ,Un,Z1, . . . ,Zn) ∈ Rn×(pX+pU+pZ) and treatment T ∈ R where n denotes the number of observations and p is the dimension of covariates. Data generating process (DGP) works as
Y = µ(T ) + 0.2(X21 +X4)T +X · βX +U · βU + T = 20 ·Ψ(φ(X · β∗X +Z · βZ)) + ν
where φ is the sigmoid function and Ψ is the pdf of Beta distribution with shape parameters set to 2 and 3. Here we set DRF µ as polinomial(Poly), exponential(Exp), and sinusoidal functions(Sinus).
µ(t) = 0.2 ∗ (t− 5)2 − t− 5, Polynomial µ(t) = log(1 + exp(t)/(t+ 0.1))− log(11), Exponential µ(t) = 5 ∗ sin(x) + x, Sinusoidal
The above DGP gives us a variety of data sets by taking the combination of DRF µ(t) and covariate Σ that are specified by n,p and DRF . Parameter , ν are noises following N(0, 1). Covariates X,U,Z follows N(0, I) and coefficient vector β follow [Unif(−1, 1)]. Following this rule, we generate 100 rounds. Meanwhile, in test data, we randomly assigned the treatments to make sure unbiased evaluations.

heading: 7.3 SIMULATION RESULTS
text: In addition to the results provided in the main paper, here we provide the inclusive results of our simulation study, which covers all 6 data sets from the Data Generating process.
The results are summarized in Table 3. Compared to CF (Athey et al., 2019) that is the benchmark of our generalization, our GCF outperforms it with significantly smaller errors. With the baseine fixed to be Kennedy (Kennedy et al., 2017), our algorithm almost enjoys smaller errors except one dataset
where DRF is sin. Overall, these experimental evidence demonstrates that our GCF is superior than the algorithms for continuous treatments.
Meanwhile, ADRF curves are also inclusively presented here as validations of our proposed GCF.
Figure 3−8 are ADRF curves of baseline methods and GCF on different synthetic data sets. Across all the data sets, our proposed GCF performs the best by being the closest to the ground truth.

heading: 7.4 SPARK IMPLEMENTATION
text: Apache Spark (Meng et al., 2016) has the power of large-scale data processing and provides APIs of any machine learning algorithms. Consequently, we build our proposed algorithm on Spark to do parallel computing and distributed model training that leverage the resources of multiple machines simultaneously.
The workflow of Spark with GCF is shown in Figure 9, which include integrated data and distributed computation. More specifically, with the parallel structure of spark, the tree-growing process runs and the process of GCF is depicted in Figure 10, which is significantly different from that in CF (Athey et al., 2019) by distributing tasks for computations

heading: 7.5 ONLINE DEPLOYMENT AND EXPERIMENTS
text: The deployment of models in the online marketplace is twofold as illustrated in Figure 11. On the one hand, the offline model that is trained on historical data provides HTE estimation that decisionmakers need to customize incentives at the origin-destination-time level. On the other hand, the outcome in response to the customized incentives contributes to the data prepared for the subsequent model training. The following is an example of an online experiment with the GCF-driven policy.
We conduct online A/B testings with the model on 2 equal-size non-overlapping groups, i.e. the cohort. The evaluation metric of the online experiment is the finish order (FO) increment. The numerical results of GCF and CF are shown in Figure 12. GCF improves on CF by 15.1%, benefiting from the relaxation of the assumptions on the model specification. It shows that the more complicated or more general the treatment is, the more efficient the model is, which again shows the superiority of our novel combination of DRF and CF for continuous treatment.

Review summaries:
The paper proposes a generalized version of the causal forest for heterogeneous treatment effect estimation for continuous treatment by non-parametric modeling of the dose-response function. To provide non-parametric modeling, the proposed method makes use of kernel-based double/debiased estimators. Experiments of the proposed methods are run on assorted synthetic datasets as well as real-world datasets to demonstrate the effectiveness of the proposed method compared to competing methods.

The authors extend the generalized random forest by Athey et al. (2019) from partially linear models to nonparametric ones and combine it with a doubly robust estimation step. The paper contains a small scale simulation study and a small real data analysis.