GCF: Generalized Causal Forest for Heterogeneous Treatment Effect Estimation Using Nonparametric Methods

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 . 1 INTRODUCTION . 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 can not 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 can not 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 . 2 RELATED WORKS . 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 can not 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 . 3 PREREQUISITES . 3.1 NOTATIONS AND ASSUMPTIONS . 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 .

Summary: 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.