End-to-End Balancing for Causal Continuous Treatment-Effect Estimation

We study the problem of observational causal inference with continuous treatment . We focus on the challenge of estimating the causal response curve for infrequentlyobserved treatment values . We design a new algorithm based on the framework of entropy balancing which learns weights that directly maximize causal inference accuracy using end-to-end optimization . Our weights can be customized for different datasets and causal inference algorithms . We propose a new theory for consistency of entropy balancing for continuous treatments . Using synthetic and real-world data , we show that our proposed algorithm outperforms the entropy balancing in accuracy of treatment effect estimation . 1 INTRODUCTION . In many applications in business , social , and health sciences , we wish to infer the effect of a continuous treatment such as drug dosage or administration duration on a health outcome variable . Often , several confounding factors are common factors of influencing both treatment and response variable , therefore for accurate causal estimation of the treatment in view , we must appropriately account for their potential impact . Unlike binary treatments , causal inference with continuous treatments is largely understudied and far more challenging than binary treatments . ( Galagate , 2016 ; Ai et al. , 2021 ) . This is primarily because continuous treatments induce uncountably many potential outcomes per unit , only one of which is observed for each unit and across units , a sparse coarsening of the underlying information needed to infer causal effects without uncertainty . Propensity score weighting ( Robins et al. , 2000 ; Imai and Van Dyk , 2004 ) , stand-alone or combined with regression-based models to achieve double robustness ( Kennedy et al. , 2017 ) , has quickly become the state of the art for causal inference . If the weights , inversely proportional to the conditional distribution of the treatment given the confounders , are correctly modeled , the weighted population will appear to come from a randomized study . However , this approach faces several challenges : ( 1 ) The weights only balance the confounders in expectation , not necessarily in the given data ( Zubizarreta et al. , 2011 ) . ( 2 ) The weights can be very large for some of units , leading to unstable estimation and uncertain inference . As a possible remedy , entropy balancing ( Hainmueller , 2012 ) estimates the weights such that they balance confounders subject to a measure of dispersion on the weights to prevent extreme weights . In this work , we note that low-entropy weights do not directly optimize the quality of subsequent weighted regression , and we introduce an alternative approach that does . We propose End-to-End Balancing ( E2B ) to improve the accuracy of the weighted regression used for causal inference . E2B uses end-to-end training to estimate the base weights in the entropy balancing framework . The E2B weights are thus customized for different datasets and causal inference algorithms that are based on weighting . Because we do not know the true treatment response function in real data , we propose a new approach to generate synthetic training datasets for end-to-end training . To theoretically analyze end-to-end balancing , we define Generalized Stable Weights ( GSW ) for causal inference as a generalization of the stable weights proposed by Robins et al . ( 2000 ) . We prove that weights learned by entropy balancing for continuous treatments , including E2B weights , are unbiased estimators of generalized stable weights . We also show that E2B weights are asymptotically consistent and efficient estimators of the population weights . We perform three sets of experiments to demonstrate accuracy improvements by E2B . Two experiments with synthetic data , one with linear and another with non-linear response functions show that the E2B is more accurate than the baseline entropy balancing and inverse propensity score techniques . In the experiments on real-world data , we qualitatively evaluate the average treatment effect function learned by E2B . We also show that the base weights learned by E2B follow our intuition about up-weighting low frequency treatments . 2 PROBLEM DEFINITION AND RELATED WORK . Problem Statement . Suppose we have the triplet of ( x , a , y ) , where x 2 X ⇢ Rr , a 2 A ⇢ R and y 2 R denote the confounders , treatments , and response variables , respectively , from an observational causal study . In our continuous treatment setting ( Galagate , 2016 , Ch . 1.2.6 ) , we denote potential outcomes as y ( a ) , which means the value of y after intervention in the treatment a and setting its value to a . Given an i.i.d . sample of size n , { ( xi , ai , yi ) } ni=1 , our objective is to eliminate the impact of the confounders and identify the average treatment effect function µ ( a ) = E [ y ( a ) ] , which is also called the response function . We make the two classic assumptions : ( 1 ) Strong ignorability : y ( a ) ? ? a | x . ( i.e. , no hidden confounders ) and ( 2 ) Positivity : 0 < P ( a|x ) < 1 . General Causal Inference Literature . The literature on causal inference is vast and we refer the reader to the books for the general inquiry ( Pearl , 2009 ; Imbens and Rubin , 2015 ; Spirtes et al. , 2000 ; Peters et al. , 2017 ) . Instead , we focus on reviewing the inference techniques for continuous treatments . In particular , we narrow down our focus on propensity score weighting approaches ( Robins et al. , 2000 ; Imai and Van Dyk , 2004 ) , because they can either be used alone or combined with the regression algorithms to create double robust algorithms . Causal Inference via Weighting . A popular approach for causal inference is to create a pseudopopulation by weighting data points such that in the pseudo-population the confounders and treatments are independent . Thus , regular regression algorithms can estimate the causal response curve using the pseudo-population , which resembles data from randomized trials . Throughout this paper , we will denote the parameters of the pseudo-population with a tilde mark . Multiple forms of propensity scores have been proposed for continuous treatments ( Hirano and Imbens , 2004 ; Imai and Van Dyk , 2004 ) . The commonly-used stablized weights ( Robins et al. , 2000 ; Zhu et al. , 2015 ) are defined as the ratio of marginal density over the conditional density of the treatments : sw = f ( a ) /f ( a|x ) . Problems with Propensity Scores . Zubizarreta et al . ( 2011 ) list two challenges with the propensity scores : ( 1 ) The weights only balance the confounders in expectation , not necessarily in the given data . ( 2 ) The weights can be very large for some of the data points , leading to unstable estimations . The challenges are amplified in the continuous setting because computing the stabilized weights requires correctly choosing two models , one for the marginal and one for the conditional distributions of the treatments . Kang et al . ( 2007 ) and Smith and Todd ( 2005 ) provide multiple evidence that the propensity score methods can lead to large biases in the estimations . While Robins et al . ( 2007 ) propose techniques to fix the large weights problems in the binary treatment examples discussed by Kang et al . ( 2007 ) , learning more accurate , bounded , and stable weights has been an active research area . Further techniques have proposed techniques to learn more robust propensity scores for binary treatments ( Li et al. , 2018 ; Zhao , 2019 ) too , however , the case of continuous treatments have received considerably less attention . Entropy Balancing . To address the problem of extreme weights , Entropy Balancing ( EB ) ( Hainmueller , 2012 ) estimates weights such that they balance the confounders subject to a measure of dispersion on the weights to prevent extremely large weights . Other loss functions using different dispersion metrics have been proposed for balancing ( Zubizarreta , 2015 ; Chan et al. , 2016 ) . Zhao and Percival ( 2016 ) show that the entropy balancing is double robust . Entropy balancing has been extended to the continuous treatment setting ( Fong et al. , 2018 ; Vegetabile et al. , 2021 ) , where the balancing condition ensures that the weighted correlation between the confounders and the treatment is zero . Ai et al . ( 2021 ) propose a method for estimating the counterfactual distribution in the continuous treatment setting . 3 METHODOLOGY . To describe our end-to-end balancing algorithm , we first need to describe entropy balancing for continuous treatments with base weights . 3.1 ENTROPY BALANCING FOR CONTINUOUS TREATMENTS . Causal Inference via Entropy Balancing . Entropy balancing creates a pseudo-population using instance weights wi , i = 1 , . . . , n , in which the treatment a and the confounders x are independent from each other . The independence is enforced by first selecting a set of functions on the confounders k ( · ) : X 7 ! R , for k = 1 , . . . , K , that are dense and complete in L2 space . Given the functions , we approximate the independence relationship by bEn [ a k ( x ) ] = 0 , for k = 1 , . . . , K , where the empirical expectation bEn is performed on the pseudo population . Hereafter , we will denote the mapped data points as ( xi ) = [ 1 ( xi ) , . . . , K ( xi ) ] . The k ( · ) functions can be chosen based on prior knowledge or defined by the penultimate layer of a neural network that predicts ( a , y ) from x . Our contributions in this paper are orthogonal to the choice of the k ( · ) functions and can benefit from ideas on learning these functions ( Zeng et al. , 2020 ) . The data-driven choice of the number of bases K is beyond the scope of current paper and left to future work . Balancing Constraint for the Continuous Treatments . Following ( Fong et al. , 2018 ; Vegetabile et al. , 2021 ) , in the case of continuous treatments , we first de-mean the confounders ( xi ) and treatments a such that without loss of generality they are taken to have mean zero . The balancing objective is to learn a set of weights wi , i = 1 , . . . , n that satisfy Pn i=1 wi ( xi ) = 0 , Pn i=1 wiai = 0 , and Pn i=1 wiai ( xi ) = 0 . We can write these three constraints in a compact form by defining a ( 2K + 1 ) –dimensional vector gi = [ ( xi ) , ai , ai ( xi ) ] . The constraints become Pn i=1 wigi = 0 . We stack the g vectors in a ( 2K + 1 ) ⇥ n dimensional matrix G for compact notation . In this work , without loss of generality , we will present our idea with the first order balancing , without higher order moments ( Galagate , 2016 ; Wong and Chan , 2018 ; Hazlett , 2020 ) . Primal and Dual EB . A variety of dispersion metrics have been proposed as objective function for minimization such as entropy or variance of the weights ( Wang and Zubizarreta , 2020 ) . Hainmueller ( 2012 ) originally proposed minimizing the KL-divergence between the weights and a set of base weights qi , i = 1 , . . . , n. Details on choice of base weights is discussed below , however , we note that qi = const . leads to minimization of the entropy of weights . Using this dispersion function and the balancing constraints , entropy balancing optimization is as follows : bw = argmin w nX i=1 wi log ✓ wi qi ◆ , ( 1 ) s.t . ( i ) Gw = 0 , ( ii ) 1 > w = 1 , ( iii ) wi 0 for i = 1 , . . . , n. The above optimization problem can be solved efficiently using its Lagrangian dual : b = argmin log 1 > exp > G + ` , ( 2 ) where ` i = log qi are the log-base-weights . Given the solution b , the balancing weights can be computed as w = softmax ⇣ b > G + ` ⌘ . The softmax function is defined as softmax ( v ) = exp v/ ⇣ 1 > exp v ⌘ for any vector v. The log base weight is a degree of freedom that we have in the Eq . ( 2 ) to improve the quality of causal estimation . We select the mapping dimension K such that problem ( 2 ) is well-conditioned and leave the analysis of the high dimensional setting K ⇡ n to future work . We can also add an L1 penalty term to the dual objective in Eq . ( 2 ) , which corresponds to approximate balancing ( Wang and Zubizarreta , 2020 ) . In the next section we propose to parameterize the log-base-weights and learn them . Our analysis in Section 4 shows that with any arbitrary base weights , causal estimation using the weights learned in Eq . ( 2 ) will be consistent . Algorithm 1 Stochastic Training of ` ✓ for End-to-End Balancing Require : Data tuples ( xi , ai , yi ) for i = 1 , . . . , n with an unknown response function µ ( a ) . Require : Representation functions ( · ) and ( · ) , split size n1 < n and batch size B . 1 : Generate a random set of indexes I , |I| = n1 and its complement Ic and split the data to S and Sc using them . 2 : Estimate the distribution of noise in y given ( a , x ) as bF '' . 3 : Compute G by stacking gi = [ ( xi ) , ai , ai ( xi ) ] , for i = 1 , . . . , n. 4 : for Number of Iterations do 5 : Generate B datasets { ( xi , ai , yi , b ) } ni=1 for b = 1 , . . . , B using `` ⇠ bF '' , and randomly selected µ ( a ) b response functions . 6 : ` i ` ✓ ( ( ai , xi ) ) . 7 : b argmin log 1 > exp > G + ` using only S data . 8 : w softmax ⇣ b > G + ` ⌘ using only Sc data . 9 : bµ ( a ) b weighting-based causal estimates using ( ai , yi , b , wi ) in Sc for b = 1 , . . . , B . 10 : Take a step in ✓ to minimize 1B PB b=1 ( bµ ( a ) b µ ( a ) b ) 2 . 11 : end for 12 : return The ` b✓ function .

Summary: The authors propose a way to intelligently choose the base weights in an entropy balancing framework for estimating treatment effects from observational data. Typically, the weights are chosen arbitrarily (e.g. to be uniform) or based off of domain expertise, which is not always available. The weights are chosen by positing forms for the response function, generating pseudo-responses, and then finding weights that maximize treatment effect estimation accuracy. In this sense, the procedure is end-to-end because the weighting is done specifically with the subsequent estimation in mind. The authors provide theoretical results showing that the weights are asymptotically normal and that they converge to the population minimizers of the entropy balancing optimization problem. Experimental results on synthetic datasets show that the proposed method is superior to entropy balancing with naive weights with respect to treatment effect estimation accuracy.