Original: E ND - TO -E ND B ALANCING FOR C AUSAL C ONTINUOUS T REATMENT -E FFECT E STIMATION
Revised: E ND - TO -E ND B ALANCING FOR C AUSAL C ONTINUOUS T REATMENT -E FFECT E STIMATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We study the problem of observational causal inference with continuous treatment.
Revised: We study the problem of observational causal inference with continuous treatment.

Original: We focus on the challenge of estimating the causal response curve for infrequentlyobserved treatment values.
Revised: We focus on the challenge of estimating the causal response curve for infrequentlyobserved treatment values.

Original: 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.
Revised: 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.

Original: Our weights can be customized for different datasets and causal inference algorithms.
Revised: Our weights can be customized for different datasets and causal inference algorithms.

Original: We propose a new theory for consistency of entropy balancing for continuous treatments.
Revised: We propose a new theory for consistency of entropy balancing for continuous treatments.

Original: Using synthetic and real-world data, we show that our proposed algorithm outperforms the entropy balancing in accuracy of treatment effect estimation
Revised: Using synthetic and real-world data, we show that our proposed algorithm outperforms the entropy balancing in accuracy of treatment effect estimation

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: Often, several confounding factors are common factors of inﬂuencing both treatment and response variable, therefore for accurate causal estimation of the treatment in view, we must appropriately account for their potential impact.
Revised: Often, several confounding factors are common factors of inﬂuencing both treatment and response variable, therefore for accurate causal estimation of the treatment in view, we must appropriately account for their potential impact.

Original: Unlike binary treatments, causal inference with continuous treatments is largely understudied and far more challenging than binary treatments.
Revised: Unlike binary treatments, causal inference with continuous treatments is largely understudied and far more challenging than binary treatments.

Original: (Galagate, 2016;
Revised: (Galagate, 2016;

Original: Ai et al., 2021).
Revised: Ai et al., 2021).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: However, this approach faces several challenges:
Revised: However, this approach faces several challenges:

Original: (1) The weights only balance the confounders in expectation, not necessarily in the given data (Zubizarreta
Revised: (1) The weights only balance the confounders in expectation, not necessarily in the given data (Zubizarreta

Original: et al., 2011).
Revised: et al., 2011).

Original: The weights can be very large for some of units, leading to unstable estimation and uncertain inference.
Revised: The weights can be very large for some of units, leading to unstable estimation and uncertain inference.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We propose End-to-End
Revised: We propose End-to-End

Original: Balancing (E2B) to improve the accuracy of the weighted regression used for causal inference.
Revised: Balancing (E2B) to improve the accuracy of the weighted regression used for causal inference.

Original: E2B uses end-to-end training to estimate the base weights in the entropy balancing framework.
Revised: E2B uses end-to-end training to estimate the base weights in the entropy balancing framework.

Original: The E2B weights are thus customized for different datasets and causal inference algorithms that are based on weighting.
Revised: The E2B weights are thus customized for different datasets and causal inference algorithms that are based on weighting.

Original: 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.
Revised: 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.

Original: To theoretically analyze end-to-end balancing, we deﬁne Generalized Stable Weights (GSW) for causal inference as a generalization of the stable weights proposed by Robins et al.
Revised: To theoretically analyze end-to-end balancing, we deﬁne Generalized Stable Weights (GSW) for causal inference as a generalization of the stable weights proposed by Robins et al.

Original: We prove that weights learned by entropy balancing for continuous treatments, including E2B weights, are unbiased estimators of generalized stable weights.
Revised: We prove that weights learned by entropy balancing for continuous treatments, including E2B weights, are unbiased estimators of generalized stable weights.

Original: We also show that E2B weights are asymptotically consistent and efﬁcient estimators of the population weights.
Revised: We also show that E2B weights are asymptotically consistent and efﬁcient estimators of the population weights.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original:  P ROBLEM D EFINITION AND R ELATED W ORK
Revised:  P ROBLEM D EFINITION AND R ELATED W ORK

Original: Problem Statement.
Revised: Problem Statement.

Original: Suppose we have the triplet of ( x , a , y) , where x 2 X ⇢ R r , a 2 A ⇢ R and y 2 R denote the confounders, treatments, and response variables, respectively, from an observational causal study.
Revised: Suppose we have the triplet of ( x , a , y) , where x 2 X ⇢ R r , a 2 A ⇢ R and y 2 R denote the confounders, treatments, and response variables, respectively, from an observational causal study.

Original: 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 .
Revised: 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 .

Original: Given an i.i.d.
Revised: Given an i.i.d.

Original: sample of size n , { ( x i , a i , y i ) } ni =1 , our objective is to eliminate the impact of the confounders and identify the average treatment effect function µ ( a ) =
Revised: sample of size n , { ( x i , a i , y i ) } ni =1 , our objective is to eliminate the impact of the confounders and identify the average treatment effect function µ ( a ) =

Original: E [y ( a ) ] , which is also called the response function.
Revised: E [y ( a ) ] , which is also called the response function.

Original: We make the two classic assumptions: (1) Strong ignorability: y (a) ??
Revised: We make the two classic assumptions: (1) Strong ignorability: y (a) ??

Original: a | x .
Revised: a | x .

Original: (i.e., no hidden confounders) and (2) Positivity: 0 < P (a | x ) < 1 .
Revised: (i.e., no hidden confounders) and (2) Positivity: 0 < P (a | x ) < 1 .

Original: General Causal Inference Literature.
Revised: General Causal Inference Literature.

Original: The literature on causal inference is vast and we refer the reader to the books for the general inquiry (Pearl, 2009;
Revised: The literature on causal inference is vast and we refer the reader to the books for the general inquiry (Pearl, 2009;

Original: Imbens and Rubin, 2015; Spirtes et al., 2000; Peters et al., 2017).
Revised: Imbens and Rubin, 2015; Spirtes et al., 2000; Peters et al., 2017).

Original: Instead, we focus on reviewing the inference techniques for continuous treatments.
Revised: Instead, we focus on reviewing the inference techniques for continuous treatments.

Original: 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.
Revised: 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.

Original: Causal Inference via Weighting.
Revised: Causal Inference via Weighting.

Original: 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.
Revised: 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.

Original: Thus, regular regression algorithms can estimate the causal response curve using the pseudo-population, which resembles data from randomized trials.
Revised: Thus, regular regression algorithms can estimate the causal response curve using the pseudo-population, which resembles data from randomized trials.

Original: Throughout this paper, we will denote the parameters of the pseudo-population with a tilde mark.
Revised: Throughout this paper, we will denote the parameters of the pseudo-population with a tilde mark.

Original: Multiple forms of propensity scores have been proposed for continuous treatments (Hirano and Imbens, 2004; Imai and Van Dyk, 2004).
Revised: Multiple forms of propensity scores have been proposed for continuous treatments (Hirano and Imbens, 2004; Imai and Van Dyk, 2004).

Original: The commonly-used stablized weights (Robins et al., 2000; Zhu et al., 2015) are deﬁned as the ratio of marginal density over the conditional density of the treatments:
Revised: The commonly-used stablized weights (Robins et al., 2000; Zhu et al., 2015) are deﬁned as the ratio of marginal density over the conditional density of the treatments:

Original: sw = f
Revised: sw = f

Original: ( a ) / f
Revised: ( a ) / f

Original: ( a | x ) .
Revised: ( a | x ) .

Original: Problems with Propensity Scores.
Revised: Problems with Propensity Scores.

Original: Zubizarreta et al.
Revised: Zubizarreta et al.

Original: (2011) list two challenges with the propensity scores: (1) The weights only balance the confounders in expectation, not necessarily in the given data.
Revised: (2011) list two challenges with the propensity scores: (1) The weights only balance the confounders in expectation, not necessarily in the given data.

Original: The weights can be very large for some of the data points, leading to unstable estimations.
Revised: The weights can be very large for some of the data points, leading to unstable estimations.

Original: The challenges are ampliﬁed 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.
Revised: The challenges are ampliﬁed 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.

Original: Kang et al.
Revised: Kang et al.

Original: and Smith and Todd (2005) provide multiple evidence that the propensity score methods can lead to large biases in the estimations.
Revised: and Smith and Todd (2005) provide multiple evidence that the propensity score methods can lead to large biases in the estimations.

Original: While Robins et al.
Revised: While Robins et al.

Original: (2007) propose techniques to ﬁx the large weights problems in the binary treatment examples discussed by
Revised: (2007) propose techniques to ﬁx the large weights problems in the binary treatment examples discussed by

Original: Kang et al.
Revised: Kang et al.

Original: (2007), learning more accurate, bounded, and stable weights has been an active research area.
Revised: (2007), learning more accurate, bounded, and stable weights has been an active research area.

Original: 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.
Revised: 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.

Original: Entropy Balancing.
Revised: Entropy Balancing.

Original: 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.
Revised: 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.

Original: Other loss functions using different dispersion metrics have been proposed for balancing (Zubizarreta, 2015; Chan et al., 2016).
Revised: Other loss functions using different dispersion metrics have been proposed for balancing (Zubizarreta, 2015; Chan et al., 2016).

Original: Zhao and Percival (2016) show that the entropy balancing is double robust.
Revised: Zhao and Percival (2016) show that the entropy balancing is double robust.

Original: 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.
Revised: 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.

Original: Ai et al.
Revised: Ai et al.

Original: (2021) propose a method for estimating the counterfactual distribution in the continuous treatment setting.
Revised: (2021) propose a method for estimating the counterfactual distribution in the continuous treatment setting.

Original:  M ETHODOLOGY
Revised:  M ETHODOLOGY

Original: To describe our end-to-end balancing algorithm, we ﬁrst need to describe entropy balancing for continuous treatments with base weights.
Revised: To describe our end-to-end balancing algorithm, we ﬁrst need to describe entropy balancing for continuous treatments with base weights.

Original: 3. E NTROPY B ALANCING FOR C ONTINUOUS T REATMENTS
Revised: 3. E NTROPY B ALANCING FOR C ONTINUOUS T REATMENTS

Original: Causal Inference via Entropy Balancing.
Revised: Causal Inference via Entropy Balancing.

Original: Entropy balancing creates a pseudo-population using instance weights w i , i = 1 , . . . , n , in which the treatment a
Revised: Entropy balancing creates a pseudo-population using instance weights w i , i = 1 , . . . , n , in which the treatment a

Original: and the confounders x are independent from each other.
Revised: and the confounders x are independent from each other.

Original: The independence is enforced by ﬁrst selecting a set of functions on the confounders  k ( · ) :
Revised: The independence is enforced by ﬁrst selecting a set of functions on the confounders  k ( · ) :

Original: X 7!
Revised: X 7!

Original: R , for k = 1 , . . .
Revised: R , for k = 1 , . . .

Original: , K , that are dense and complete in L 2 space.
Revised: , K , that are dense and complete in L 2 space.

Original: Given the  functions, we approximate the independence relationship by b
Revised: Given the  functions, we approximate the independence relationship by b

Original: E n [a  k ( x )] = 0 , for k = 1 , . .
Revised: E n [a  k ( x )] = 0 , for k = 1 , . .

Original: , K , where the empirical expectation b
Revised: , K , where the empirical expectation b

Original: E n is performed on the pseudo population.
Revised: E n is performed on the pseudo population.

Original: Hereafter, we will denote the mapped data points as  ( x i ) =
Revised: Hereafter, we will denote the mapped data points as  ( x i ) =

Original: [  1 ( x i ) , . . .
Revised: [  1 ( x i ) , . . .

Original: ,  K ( x i )] .
Revised: ,  K ( x i )] .

Original: The  k ( · ) functions can be chosen based on prior knowledge or deﬁned by the penultimate layer of a neural network that predicts (a , y) from x .
Revised: The  k ( · ) functions can be chosen based on prior knowledge or deﬁned by the penultimate layer of a neural network that predicts (a , y) from x .

Original: Our contributions in this paper are orthogonal to the choice of the  k ( · ) functions and can beneﬁt from ideas on learning these functions (Zeng et al., 2020).
Revised: Our contributions in this paper are orthogonal to the choice of the  k ( · ) functions and can beneﬁt from ideas on learning these functions (Zeng et al., 2020).

Original: The data-driven choice of the number of bases K is beyond the scope of current paper and left to future work.
Revised: The data-driven choice of the number of bases K is beyond the scope of current paper and left to future work.

Original: Balancing Constraint for the Continuous Treatments.
Revised: Balancing Constraint for the Continuous Treatments.

Original: Following Fong et al.
Revised: Following Fong et al.

Original: (2018), in the case of continuous treatments, we ﬁrst de-mean the confounders  ( x i ) and treatments a such that without loss of generality they are taken to have mean zero.
Revised: (2018), in the case of continuous treatments, we ﬁrst de-mean the confounders  ( x i ) and treatments a such that without loss of generality they are taken to have mean zero.

Original: The balancing objective is to learn a set of weights w i , i = 1 , . . . , n that satisfy
Revised: The balancing objective is to learn a set of weights w i , i = 1 , . . . , n that satisfy

Original: P ni =1 w i  ( x i )
Revised: P ni =1 w i  ( x i )

Original: = 0 ,
Revised: = 0 ,

Original: P ni =

Revised: P ni =

Original: w i a i = 0 , and
Revised: w i a i = 0 , and

Original: P ni =

Revised: P ni =

Original: w i a i  ( x i )
Revised: w i a i  ( x i )

Original: = 0 .
Revised: = 0 .

Original: We can write these three constraints in a compact form by deﬁning a (2 K + 1) –dimensional vector g i =
Revised: We can write these three constraints in a compact form by deﬁning a (2 K + 1) –dimensional vector g i =

Original: [  ( x i ) , a i , a i  ( x i )] .
Revised: [  ( x i ) , a i , a i  ( x i )] .

Original: The constraints become
Revised: The constraints become

Original: P ni =1 w
Revised: P ni =1 w

Original: i
Revised: i

Original: g i = 0 .
Revised: g i = 0 .

Original: We stack the g vectors in a (2 K + 1) ⇥ n dimensional matrix G for compact notation.
Revised: We stack the g vectors in a (2 K + 1) ⇥ n dimensional matrix G for compact notation.

Original: In this work, without loss of generality, we will present our idea with the ﬁrst order balancing, without higher order moments (Galagate, 2016; Wong and Chan, 2018; Hazlett, 2020).
Revised: In this work, without loss of generality, we will present our idea with the ﬁrst order balancing, without higher order moments (Galagate, 2016; Wong and Chan, 2018; Hazlett, 2020).

Original: Primal and Dual EB.
Revised: Primal and Dual EB.

Original: 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).
Revised: 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).

Original: Hainmueller (2012) originally proposed minimizing the KL-divergence between the weights and a set of base weights
Revised: Hainmueller (2012) originally proposed minimizing the KL-divergence between the weights and a set of base weights

Original: q i , i = 1 , . . . , n .
Revised: q i , i = 1 , . . . , n .

Original: Details on choice of base weights is discussed below, however, we note that q i = const .
Revised: Details on choice of base weights is discussed below, however, we note that q i = const .

Original: leads to minimization of the entropy of weights.
Revised: leads to minimization of the entropy of weights.

Original: Using this dispersion function and the balancing constraints, entropy balancing optimization is as follows:
Revised: Using this dispersion function and the balancing constraints, entropy balancing optimization is as follows:

Original: [Equation]
Revised: [Equation]

Original: The above optimization problem can be solved efﬁciently using its Lagrangian dual:
Revised: The above optimization problem can be solved efﬁciently using its Lagrangian dual:

Original: [Equation]
Revised: [Equation]

Original: where ` i = log q
Revised: where ` i = log q

Original: i are the log-base-weights .
Revised: i are the log-base-weights .

Original: Given the solution b
Revised: Given the solution b

Original:  , the balancing weights can be computed as w = softmax
Revised:  , the balancing weights can be computed as w = softmax

Original: [Equation]
Revised: [Equation]

Original: ⌘ .
Revised: ⌘ .

Original: The softmax function is deﬁned as softmax( v ) =
Revised: The softmax function is deﬁned as softmax( v ) =

Original: exp v / ⇣ 1 > exp v ⌘ for any vector v .
Revised: exp v / ⇣ 1 > exp v ⌘ for any vector v .

Original: The log base weight is a degree of freedom that we have in the
Revised: The log base weight is a degree of freedom that we have in the

Original: Eq.
Revised: Eq.

Original: (2) to improve the quality of causal estimation.
Revised: (2) to improve the quality of causal estimation.

Original: 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.
Revised: 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.

Original: We can also add an L 1 penalty term to the dual objective in Eq.
Revised: We can also add an L 1 penalty term to the dual objective in Eq.

Original: (2), which corresponds to approximate balancing (Wang and Zubizarreta, 2020).
Revised: (2), which corresponds to approximate balancing (Wang and Zubizarreta, 2020).

Original: In the next section we propose to parameterize the log-base-weights and learn them.
Revised: In the next section we propose to parameterize the log-base-weights and learn them.

Original: Our analysis in
Revised: Our analysis in

Original: Section 4 shows that with any arbitrary base weights, causal estimation using the weights learned in
Revised: Section 4 shows that with any arbitrary base weights, causal estimation using the weights learned in

Original: Eq.
Revised: Eq.

Original: (2) will be consistent.
Revised: (2) will be consistent.

Original: 3. L EARNING THE B ASE W EIGHTS
Revised: 3. L EARNING THE B ASE W EIGHTS

Original: Hainmueller (2012) suggests two approaches for choosing base weights: (1) weights obtained from a conventional propensity score model and (2), in the context of survey design, using knowledge about the sampling design.
Revised: Hainmueller (2012) suggests two approaches for choosing base weights: (1) weights obtained from a conventional propensity score model and (2), in the context of survey design, using knowledge about the sampling design.

Original: We argue that a data-driven approach that learns customized base weights for each pair of dataset and weighted causal regression algorithm can further improve performance.
Revised: We argue that a data-driven approach that learns customized base weights for each pair of dataset and weighted causal regression algorithm can further improve performance.

Original: To address this problem, we deﬁne the log-base-weights ` as a parametric function (e.g., a neural network) of the treatment variable; i.e. ` ✓ ( · ) .
Revised: To address this problem, we deﬁne the log-base-weights ` as a parametric function (e.g., a neural network) of the treatment variable; i.e. ` ✓ ( · ) .

Original: We learn the base weights with the goal of improving the accuracy of the subsequent weighted regression.
Revised: We learn the base weights with the goal of improving the accuracy of the subsequent weighted regression.

Original: This is challenging because simply optimizing the weighted regression loss (e.g., weighted MSE) leads to degenerate results.
Revised: This is challenging because simply optimizing the weighted regression loss (e.g., weighted MSE) leads to degenerate results.

Original: That is, learning ` to minimize the regression loss will lead to exclusion of the difﬁcult-to-predict data points from the regression, which is undesirable.
Revised: That is, learning ` to minimize the regression loss will lead to exclusion of the difﬁcult-to-predict data points from the regression, which is undesirable.

Original: Thus, we need to ﬁnd another loss function to optimize, ideally a loss function that directly minimizes the error in estimation of the response function µ ( a ) .
Revised: Thus, we need to ﬁnd another loss function to optimize, ideally a loss function that directly minimizes the error in estimation of the response function µ ( a ) .

Original: Our idea for learning the parameters of the base weights is to generate multiple pseudo-responses y with randomly generated response functions µ ( a ) .
Revised: Our idea for learning the parameters of the base weights is to generate multiple pseudo-responses y with randomly generated response functions µ ( a ) .

Original: Now that we know the true response function µ ( a ) in the randomly generated data, we can perform causal inference and obtain the estimation of the known response curve b µ ( a ) using our weights.
Revised: Now that we know the true response function µ ( a ) in the randomly generated data, we can perform causal inference and obtain the estimation of the known response curve b µ ( a ) using our weights.

Original: Algorithm 1 outlines our stochastic training of the base weight function.
Revised: Algorithm 1 outlines our stochastic training of the base weight function.

Original: First, in Step 2, we estimate the distribution of noise using the residuals of regressing y over (a , x ) , capturing the possible heteroskedasticity in the noise.
Revised: First, in Step 2, we estimate the distribution of noise using the residuals of regressing y over (a , x ) , capturing the possible heteroskedasticity in the noise.

Original: Then, in each iteration, we draw a batch of possible datasets.
Revised: Then, in each iteration, we draw a batch of possible datasets.

Original: To generate each dataset, we randomly choose a response function µ ( a ) and use it to generate the entire dataset (see Section 5.1 for examples of random functions).
Revised: To generate each dataset, we randomly choose a response function µ ( a ) and use it to generate the entire dataset (see Section 5.1 for examples of random functions).

Original: For the entire batch, we use ` ✓ to learn the log-base-weights, and subsequently learn the weights in lines 7–8.
Revised: For the entire batch, we use ` ✓ to learn the log-base-weights, and subsequently learn the weights in lines 7–8.

Original: In line 9 we use a weighted regression algorithm to ﬁnd our estimation b µ ( a ) of the randomly-generated µ ( a ) .
Revised: In line 9 we use a weighted regression algorithm to ﬁnd our estimation b µ ( a ) of the randomly-generated µ ( a ) .

Original: Our loss function is the mean squared error between the latter quantities.
Revised: Our loss function is the mean squared error between the latter quantities.

Original: While we call our algorithm End-to-End Balancing (E2B) because of our end-to-end optimization.
Revised: While we call our algorithm End-to-End Balancing (E2B) because of our end-to-end optimization.

Original: Sample Splitting.
Revised: Sample Splitting.

Original: The E2B procedure involves estimation of two sets of parameters ✓ in the ` ✓ and  for entropy balancing.
Revised: The E2B procedure involves estimation of two sets of parameters ✓ in the ` ✓ and  for entropy balancing.

Original: The joint estimation of ✓ ,  on a single sample will result in bias (Chernozhukov et al., 2018).
Revised: The joint estimation of ✓ ,  on a single sample will result in bias (Chernozhukov et al., 2018).

Original: Thus, we split the sample to two mutually exclusive parts and perform the optimizations on separate partitions of data.
Revised: Thus, we split the sample to two mutually exclusive parts and perform the optimizations on separate partitions of data.

Original: Choice of Random Response Functions.
Revised: Choice of Random Response Functions.

Original: Ideally, we should rely on domain experts for choosing the random set of response functions µ ( a ) that includes the true response function.
Revised: Ideally, we should rely on domain experts for choosing the random set of response functions µ ( a ) that includes the true response function.

Original: Alternatively, we can choose broad function classes such as random piecewise smooth functions or polynomial functions with random coefﬁcients.
Revised: Alternatively, we can choose broad function classes such as random piecewise smooth functions or polynomial functions with random coefﬁcients.

Original: We can also use generative adversarial networks to generate data that is more similar to our sample (Athey et al., 2021).
Revised: We can also use generative adversarial networks to generate data that is more similar to our sample (Athey et al., 2021).

Original: Features Fed to ` ✓ .
Revised: Features Fed to ` ✓ .

Original: We can feed the raw values of the treatments and any handcrafted features, denoted by ( a i , x i ) .
Revised: We can feed the raw values of the treatments and any handcrafted features, denoted by ( a i , x i ) .

Original: We empirically ﬁnd that ( a i , x i ) =
Revised: We empirically ﬁnd that ( a i , x i ) =

Original: (log p ( a i ) , log p ( a i | x i )) makes training the ` ✓ easier.
Revised: (log p ( a i ) , log p ( a i | x i )) makes training the ` ✓ easier.

Original: We describe the details of our neural network model for ` ✓ and our techniques for training in Appendix B.
Revised: We describe the details of our neural network model for ` ✓ and our techniques for training in Appendix B.

Original: Weighted Regression Algorithms.
Revised: Weighted Regression Algorithms.

Original: To be able to differentiate the loss function with respect to ✓ , we need weighted regression algorithms whose estimates are differentiable with respect to the weights.
Revised: To be able to differentiate the loss function with respect to ✓ , we need weighted regression algorithms whose estimates are differentiable with respect to the weights.

Original: In the linear average treatment effect function we choose weighted linear regression and in the non-linear setting we use the weighted polynomial regression and the local kernel regression, as used by Flores et al.
Revised: In the linear average treatment effect function we choose weighted linear regression and in the non-linear setting we use the weighted polynomial regression and the local kernel regression, as used by Flores et al.

Original: Double Robustness.
Revised: Double Robustness.

Original: Zhao and Percival (2016) show that in the binary case, the entropy balancing is double robust.
Revised: Zhao and Percival (2016) show that in the binary case, the entropy balancing is double robust.

Original: We do not attempt to show double robustness for E2B because we see E2B as a meta algorithm that learns customized weights for each dataset and algorithm.
Revised: We do not attempt to show double robustness for E2B because we see E2B as a meta algorithm that learns customized weights for each dataset and algorithm.

Original: We can either (1) plug-in the E2B weights in the double robust algorithm and expect improved accuracy, or (2) learn weights that directly minimize the error of double robust algorithms such as (Kennedy et al., 2017).
Revised: We can either (1) plug-in the E2B weights in the double robust algorithm and expect improved accuracy, or (2) learn weights that directly minimize the error of double robust algorithms such as (Kennedy et al., 2017).

Original:  A NALYSIS
Revised:  A NALYSIS

Original: We prove that for any arbitrary choice of the log-base-weight function ` ✓ , our approach consistently estimates causal effects.
Revised: We prove that for any arbitrary choice of the log-base-weight function ` ✓ , our approach consistently estimates causal effects.

Original: Before proving the consistency results, we characterize the quantity that our solution converges to.
Revised: Before proving the consistency results, we characterize the quantity that our solution converges to.

Original: All long proofs are relegated to Appendix A.
Revised: All long proofs are relegated to Appendix A.

Original: Deﬁnition 1.
Revised: Deﬁnition 1.

Original: Generalized Stable Weights.
Revised: Generalized Stable Weights.

Original: Suppose f ( a, x ) denote the joint probability density function of treatments and confounders in a population.
Revised: Suppose f ( a, x ) denote the joint probability density function of treatments and confounders in a population.

Original: Suppose e f ( a ) and e f ( x ) denote two arbitrary density functions, possibly different with the marginal density functions in our population, that satisfy
Revised: Suppose e f ( a ) and e f ( x ) denote two arbitrary density functions, possibly different with the marginal density functions in our population, that satisfy

Original: E x ⇠ e f ( x ) [ x ] = 0 and E a ⇠
Revised: E x ⇠ e f ( x ) [ x ] = 0 and E a ⇠

Original: [Equation]
Revised: [Equation]

Original: We deﬁne the Generalized Stable Weights as follows
Revised: We deﬁne the Generalized Stable Weights as follows

Original: [Equation]
Revised: [Equation]

Original: Remark.
Revised: Remark.

Original: Our deﬁnition generalizes the stabilized weights deﬁned by Robins et al.
Revised: Our deﬁnition generalizes the stabilized weights deﬁned by Robins et al.

Original: (2000), where e f ( a ) and e f ( x ) match the marginal probability density functions in the original population.
Revised: (2000), where e f ( a ) and e f ( x ) match the marginal probability density functions in the original population.

Original: Proposition 1.
Revised: Proposition 1.

Original: The generalized stable weights w GSW satisfy E [ w GSW a x ] = 0 .
Revised: The generalized stable weights w GSW satisfy E [ w GSW a x ] = 0 .

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: where the last equation is because of zero mean assumption for the e f ( a ) and e f ( x ) distributions.
Revised: where the last equation is because of zero mean assumption for the e f ( a ) and e f ( x ) distributions.

Original: Now, we can show that with an appropriate choice of the  functions, the solution of Eq.
Revised: Now, we can show that with an appropriate choice of the  functions, the solution of Eq.

Original: (2) approximates the generalized stable weights.
Revised: (2) approximates the generalized stable weights.

Original: Consider the population version of Eq.
Revised: Consider the population version of Eq.

Original: [Equation]
Revised: [Equation]

Original: The weights corresponding to  ? can be calculated as w ?
Revised: The weights corresponding to  ? can be calculated as w ?

Original: =
Revised: =

Original: C exp( g >  ?
Revised: C exp( g >  ?

Original: + ` ) , where C = R exp( g >  ?
Revised: + ` ) , where C = R exp( g >  ?

Original: + ` )d F ( a, x )   1 is the normalization constant.
Revised: + ` )d F ( a, x )   1 is the normalization constant.

Original: Assumptions.
Revised: Assumptions.

Original: f ( a, x )  c > 0 for all ( a, x ) 2 A ⇥ X pairs, where c is a constant. 2.
Revised: f ( a, x )  c > 0 for all ( a, x ) 2 A ⇥ X pairs, where c is a constant. 2.

Original: Suppose the basis functions are dense and rich enough such for some small values of   K that they satisfy:
Revised: Suppose the basis functions are dense and rich enough such for some small values of   K that they satisfy:

Original: [Equation]
Revised: [Equation]

Original: Suppose the population problem in Eq.
Revised: Suppose the population problem in Eq.

Original: (4) has a unique solution  ?
Revised: (4) has a unique solution  ?

Original: and the corresponding weights are denoted by w ? .
Revised: and the corresponding weights are denoted by w ? .

Original: The following theorem shows that the solution to Eq.
Revised: The following theorem shows that the solution to Eq.

Original: (4) converges to w GSW :
Revised: (4) converges to w GSW :

Original: Theorem 1.
Revised: Theorem 1.

Original: Given the assumptions, the solution to the population problem satisﬁes:
Revised: Given the assumptions, the solution to the population problem satisﬁes:

Original: [Equation]
Revised: [Equation]

Original: If we select the function set  K such that   K
Revised: If we select the function set  K such that   K

Original: = o (1) , the theorem shows that w ?
Revised: = o (1) , the theorem shows that w ?

Original: ( a, x ) is an unbiased estimator of w GSW ( a, x ) .
Revised: ( a, x ) is an unbiased estimator of w GSW ( a, x ) .

Original: Notice that Assumption 1 is only slightly stronger than the common positivity assumption.
Revised: Notice that Assumption 1 is only slightly stronger than the common positivity assumption.

Original: Assumption 2 requires us to select the mapping functions such that zero the correlation between the mapped confounders and the treatment implies their independence.
Revised: Assumption 2 requires us to select the mapping functions such that zero the correlation between the mapped confounders and the treatment implies their independence.

Original: We provide the proof in Appendix A.1.
Revised: We provide the proof in Appendix A.1.

Original: Note that the quality of the features and neural network training does not affect the unbiasedness of the E2B because of the balancing constraint is still satisﬁed.
Revised: Note that the quality of the features and neural network training does not affect the unbiasedness of the E2B because of the balancing constraint is still satisﬁed.

Original: The ﬂexibility in choice of e f distributions in the deﬁnition of w GSW is due to the fact that we require only ﬁrst order balancing.
Revised: The ﬂexibility in choice of e f distributions in the deﬁnition of w GSW is due to the fact that we require only ﬁrst order balancing.

Original: If we enforce higher order balancing constraints in the form of E [w ?  1 (a)  2 ( x )]
Revised: If we enforce higher order balancing constraints in the form of E [w ?  1 (a)  2 ( x )]

Original: =
Revised: =

Original: E [  1 (a)] · E [  2 ( x )] for any suitable functions  1 and  2 , Theorem 1 in (Ai et al., 2021) shows that w ?
Revised: E [  1 (a)] · E [  2 ( x )] for any suitable functions  1 and  2 , Theorem 1 in (Ai et al., 2021) shows that w ?

Original: =
Revised: =

Original: f ( a ) / f
Revised: f ( a ) / f

Original: ( a | x ) .
Revised: ( a | x ) .

Original: The more ﬂexible form of weights in Eq.
Revised: The more ﬂexible form of weights in Eq.

Original: (3) allows us to pick the marginals e f ( a ) and e f ( x ) with more freedom.
Revised: (3) allows us to pick the marginals e f ( a ) and e f ( x ) with more freedom.

Original: In this work, we have chosen a data-driven way to learn them.
Revised: In this work, we have chosen a data-driven way to learn them.

Original: Finally, the following theorem establishes the asymptotic consistency and normality result for each individual weight estimated by E2B, under the common regularity conditions for problem (2).
Revised: Finally, the following theorem establishes the asymptotic consistency and normality result for each individual weight estimated by E2B, under the common regularity conditions for problem (2).

Original: Theorem 2.
Revised: Theorem 2.

Original: Suppose ⇤ ⇢ R 2 K +1 is an open subset of Euclidean space and the solution
Revised: Suppose ⇤ ⇢ R 2 K +1 is an open subset of Euclidean space and the solution

Original: b  n 2 ⇤ to Eq.
Revised: b  n 2 ⇤ to Eq.

Original: (2) is within the subset.
Revised: (2) is within the subset.

Original: The weights estimated by Eq.
Revised: The weights estimated by Eq.

Original: (2) are asymptotically normal for i = 1 , . .
Revised: (2) are asymptotically normal for i = 1 , . .

Original: . , n :
Revised: . , n :

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: We provide the population form of  2 ( a i , x i ) and an unbiased sample estimate for it in Appendix A.2.
Revised: We provide the population form of  2 ( a i , x i ) and an unbiased sample estimate for it in Appendix A.2.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We use two synthetic and one real-world datasets to show that E2B outperforms the baselines.
Revised: We use two synthetic and one real-world datasets to show that E2B outperforms the baselines.

Original: In the synthetic datasets, we have access to the true treatment effects; thus we measure accuracy of the algorithms in recovering the treatment effects.
Revised: In the synthetic datasets, we have access to the true treatment effects; thus we measure accuracy of the algorithms in recovering the treatment effects.

Original: In the real-world data, we qualitatively evaluate the estimated causal treatment effect curve and inspect the learned log-base-weight function.
Revised: In the real-world data, we qualitatively evaluate the estimated causal treatment effect curve and inspect the learned log-base-weight function.

Original: Baselines.
Revised: Baselines.

Original: A key baseline in our study is the Inverse Propensity score Weighting (IPW) with Stable
Revised: A key baseline in our study is the Inverse Propensity score Weighting (IPW) with Stable

Original: Weights (Robins et al., 2000) as the most commonly used technique.
Revised: Weights (Robins et al., 2000) as the most commonly used technique.

Original: To avoid extreme weights and prevent instability, we trim (Winsorize) the weights by [5 , 95] percentiles (Cole and Hernán, 2008;
Revised: To avoid extreme weights and prevent instability, we trim (Winsorize) the weights by [5 , 95] percentiles (Cole and Hernán, 2008;

Original: Chernozhukov et al., 2018).
Revised: Chernozhukov et al., 2018).

Original: However, the main baseline in our experiments is Entropy Balancing, which is equal to E2B with ` ✓ = const , corresponding to the constant base weights.
Revised: However, the main baseline in our experiments is Entropy Balancing, which is equal to E2B with ` ✓ = const , corresponding to the constant base weights.

Original: EB allows us to do an ablation study and see the exact amount of improvement by learning a customized ` ✓ function.
Revised: EB allows us to do an ablation study and see the exact amount of improvement by learning a customized ` ✓ function.

Original: We also include EB with the stabilized weights (SW) as base weights ( ` ✓ = log
Revised: We also include EB with the stabilized weights (SW) as base weights ( ` ✓ = log

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: Finally, we also include the permutation weighting algorithm (Arbour et al., 2021) that proposes to compute the weights using permutation of the treatments and a classiﬁer that predict the probability of being permuted.

Original: l og ( b p ( a | x ))
Revised: We provide further details on this algorithm in Appendix B.4.

Original: [Table]
Revised: [Table]

Original: [Figure]

Original: Training Details.
Revised: Training Details.

Original: We provide the details of the neural networks used for the ` ✓ and propensity score estimation for IPW in Appendix B. All neural networks are trained using Adam (Kingma and Ba, 2014) with early stopping based on validation error.
Revised: We provide the details of the neural networks used for the ` ✓ and propensity score estimation for IPW in Appendix B. All neural networks are trained using Adam (Kingma and Ba, 2014) with early stopping based on validation error.

Original: The learning rate and architectural parameters of the neural networks are tuned via hyperparameter search on the validation data.
Revised: The learning rate and architectural parameters of the neural networks are tuned via hyperparameter search on the validation data.

Original: 5. S YNTHETIC D ATA E XPERIMENTS
Revised: 5. S YNTHETIC D ATA E XPERIMENTS

Original: Linear.
Revised: Linear.

Original: We use the following steps to generate 100 datasets, each with 1000 data points.
Revised: We use the following steps to generate 100 datasets, each with 1000 data points.

Original: Generate confounders x 2 R 5 , x ⇠ N ( 0 , ⌃ ) , where ⌃ is a tridiagonal covariance matrix with diagonal and off-diagonal elements equal to1 . 0 and 0 . 2 , respectively.
Revised: Generate confounders x 2 R 5 , x ⇠ N ( 0 , ⌃ ) , where ⌃ is a tridiagonal covariance matrix with diagonal and off-diagonal elements equal to 0 and 0 . 2 , respectively

Original: a ⇠ N ( µ a , 0 .
Revised: a ⇠ N ( µ a , 0 .

Original: 3 2 ) , where µ a = sin(  > xa x ) and  xa,k ⇠ Unif(  1 , 1) for k = 1 , . .
Revised: 3 2 ) , where µ a = sin(  > xa x ) and  xa,k ⇠ Unif(  1 , 1) for k = 1 , . .

Original: y ⇠ N ( µ y , 0 .
Revised: y ⇠ N ( µ y , 0 .

Original: 5 2 ) , where µ y =
Revised: 5 2 ) , where µ y =

Original:  > xy x +  ay a , where  ax ,  xy,k ⇠ N (0 , 1) for k = 1 , . .
Revised:  > xy x +  ay a , where  ax ,  xy,k ⇠ N (0 , 1) for k = 1 , . .

Original: We use weighted least squares as the regression algorithm and report the average | d  ay   ay | over all 100 datasets.
Revised: We use weighted least squares as the regression algorithm and report the average | d  ay   ay | over all 100 datasets.

Original: Nonlinear We ﬁrst generate confounders x and treatments a similar to steps 1 and 2 of the linear case.
Revised: Nonlinear We ﬁrst generate confounders x and treatments a similar to steps 1 and 2 of the linear case.

Original: Then, we generate the response variable according to y ⇠ N ( µ y , 0 .
Revised: Then, we generate the response variable according to y ⇠ N ( µ y , 0 .

Original: 5 2 ) , where µ y =
Revised: 5 2 ) , where µ y =

Original:  > xy x + h  ay (a) , where  xy,k ⇠ N (0 , 1) for k = 1 , . . .
Revised:  > xy x + h  ay (a) , where  xy,k ⇠ N (0 , 1) for k = 1 , . . .

Original: The Hermit polynomials are deﬁned as h  ( z ) =
Revised: The Hermit polynomials are deﬁned as h  ( z ) =

Original:  0 +  1 z +  2 ( z 2  1) +  3 ( x 3  3 x ) .
Revised:  0 +  1 z +  2 ( z 2  1) +  3 ( x 3  3 x ) .

Original: Similar to the linear case, we generate 100 samples of size 1000.
Revised: Similar to the linear case, we generate 100 samples of size 1000.

Original: We use the weighted polynomial regression as the regression algorithm to estimate b  and report the average RMSE between true  and b  .
Revised: We use the weighted polynomial regression as the regression algorithm to estimate b  and report the average RMSE between true  and b  .

Original: We report the mean and standard error of errors on 100 datasets in Table 1.
Revised: We report the mean and standard error of errors on 100 datasets in Table 1.

Original: [Figure]

Original: C a r d i o v a s c u l a r M o rt a li t y R a t e ( p e r 100 K )

Original: [Figure] [Figure] [Figure] As seen in Table 1, in both linear and non-linear datasets, the E2B is signiﬁcantly more accurate in uncovering the true treatment response functions.
Revised: As seen in Table 1, in both linear and non-linear datasets, the E2B is signiﬁcantly more accurate in uncovering the true treatment response functions.

Original: Both constant and IPW base weights perform worse than the base-weights learned by end-to-end balancing.
Revised: Both constant and IPW base weights perform worse than the base-weights learned by end-to-end balancing.

Original: As Robins et al.
Revised: As Robins et al.

Original: (2007) caution, synthetic data evaluation might exacerbate the extreme weights issues because unlike real data, usually no manual inspection of weights are done.
Revised: (2007) caution, synthetic data evaluation might exacerbate the extreme weights issues because unlike real data, usually no manual inspection of weights are done.

Original: To gain more insights, in Figure 1, we plot the log-base weight function that we learn as a function of log( b p ( a )) and log( b p ( a | x )) .
Revised: To gain more insights, in Figure 1, we plot the log-base weight function that we learn as a function of log( b p ( a )) and log( b p ( a | x )) .

Original: We align all curves at their starting point and plot the median of 100 runs.
Revised: We align all curves at their starting point and plot the median of 100 runs.

Original: Both ﬁgures, show more variations in the log( b p ( a )) –axis, rather than the log( b p ( a | x )) –axis.
Revised: Both ﬁgures, show more variations in the log( b p ( a )) –axis, rather than the log( b p ( a | x )) –axis.

Original: Not that, especially in the linear case, the smaller conditional probability leads to larger base-weights, inline with the IPW base-weights.
Revised: Not that, especially in the linear case, the smaller conditional probability leads to larger base-weights, inline with the IPW base-weights.

Original: Finally, the complexity of the plots emphasizes the need for end-to-end methods for learning weights.
Revised: Finally, the complexity of the plots emphasizes the need for end-to-end methods for learning weights.

Original: 5. R EAL D ATA E XPERIMENTS
Revised: 5. R EAL D ATA E XPERIMENTS

Original: We study the impact of PM 2 .
Revised: We study the impact of PM 2 .

Original: 5 particle level on the cardiovascular mortality rate (CMR) incounties in the US using the data provided by the National Studies on Air Pollution and Health (Rappold, 2020).
Revised: 5 particle level on the cardiovascular mortality rate (CMR) incounties in the US using the data provided by the National Studies on Air Pollution and Health (Rappold, 2020).

Original: The data is publicly available under U.S. Public Domain license.
Revised: The data is publicly available under U.S. Public Domain license.

Original: The PM 2 .particle level and the mortality rate are measured by µg/m 3 and the number of annual deaths due to cardiovascular conditions per 100,000 people, respectively.
Revised: The PM 2 .particle level and the mortality rate are measured by µg/m 3 and the number of annual deaths due to cardiovascular conditions per 100,000 people, respectively.

Original: We use only the data for 2010 to simplify the experiment setup; thus we measure the same year impact of PM 2 .
Revised: We use only the data for 2010 to simplify the experiment setup; thus we measure the same year impact of PM 2 .

Original: 5 particle level.
Revised: 5 particle level.

Original: Other than the treatment and response variables, the data includes 10 variables such as poverty rate, population, and household income, which we use as confounders.
Revised: Other than 0 log( b p ( a ))  3 . 2 . 2 . 1 . 1 . 0 .l og ( b p ( a | x ))  ( a )  0 .0 .0 .0 .1 .1 .2 .2 .3 .3 .

Revised: [Figure]

Revised: Design 0 log( b p ( a ))  3 . 2 . 2 . 1 . 1 . 0 .l

Revised: og ( b p ( a | x ))  ( a ) 0 .0 .0 .0 .0 .1 .1 .1 .1 .

Revised: [Figure]

Revised: [Figure]

Revised: l

Revised: [Figure]

Revised: t

Revised: [Figure]

Revised: a t

Revised: [Figure]

Revised: ( b p ( a |

Revised: [Figure]

Revised: ))  (

Revised: [Figure]

Revised: )

Revised: [Figure]

Revised:

Revised: [Figure]

Revised:

Revised: [Figure]

Revised: the treatment and response variables, the data includes 10 variables such as poverty rate, population, and household income, which we use as confounders.

Original: We provide the descriptive statistics and the histograms for the treatment and effect in Appendix C.
Revised: We provide the descriptive statistics and the histograms for the treatment and effect in Appendix C.

Original: To train E2B, we create the random dataset (Line 5 in Algorithm 1) using Hermite polynomials of max degree 3, µ y =
Revised: To train E2B, we create the random dataset (Line 5 in Algorithm 1) using Hermite polynomials of max degree 3, µ y =

Original: [Equation]
Revised: [Equation]

Original:  We use absolute value to capture the positivity of our response variable.
Revised:   . We use absolute value to capture the positivity of our response variable.

Original: The data also shows heteroskedasticity; we model the noise as a zero mean
Revised: The data also shows heteroskedasticity; we model the noise as a zero mean

Original: Gaussian variable with variance  2 ( [Equation]
Revised: Gaussian variable with variance  2 ( b y ) =

Original: = 6 .b
Revised: 6 .b

Original: y .
Revised: y .

Original: For regression, we use the non-parametric local kernel regression algorithm.
Revised: For regression, we use the non-parametric local kernel regression algorithm.

Original: We measure the uncertainty in the curves using the deep ensembles technique (Lakshminarayanan et al., 2017) with 100 random ensembles.
Revised: We measure the uncertainty in the curves using the deep ensembles technique (Lakshminarayanan et al., 2017) with 100 random ensembles.

Original: Figure 2a shows the average treatment effect curve for the impact of PM 2 . 5 on CMR.
Revised: Figure 2a shows the average treatment effect curve for the impact of PM 2 . 5 on CMR.

Original: We show the one standard deviation interval using the shaded areas.
Revised: We show the one standard deviation interval using the shaded areas.

Original: Starting around PM 2 . 5 = 5 .
Revised: Starting around PM 2 . 5 = 5 .

Original: 3 µg/mthe curve increases with a steep slope; conﬁrming the previous studies that increased PM 2 .
Revised: 3 µg/mthe curve increases with a steep slope; conﬁrming the previous studies that increased PM 2 .

Original: 5 levels increase the probability of cardiovascular mortality.
Revised: 5 levels increase the probability of cardiovascular mortality.

Original: Our results are generally aligned with the results reported in (Wu et al., 2020).
Revised: Our results are generally aligned with the results reported in (Wu et al., 2020).

Original: We can see that after PM 2 . 5 = 6 . 4 µg/m 3 the curve plateaus and mortality rate stays at elevated levels.
Revised: We can see that after PM 2 . 5 = 6 . 4 µg/m 3 the curve plateaus and mortality rate stays at elevated levels.

Original: Looking at the histogram of the treatments in Figure 3a in the appendix, we observed that most counties have PM 2 .
Revised: Looking at the histogram of the treatments in Figure 3a in the appendix, we observed that most counties have PM 2 .

Original: 5 between 6 and 8.
Revised: 5 between 6 and 8.

Original: This might justify the ﬂuctuations that we see in this interval and may allude about potential unmeasured confounders.
Revised: This might justify the ﬂuctuations that we see in this interval and may allude about potential unmeasured confounders.

Original: Figure 2b shows the log-base-weight function that we learn in this data.
Revised: Figure 2b shows the log-base-weight function that we learn in this data.

Original: Similar to the synthetic experiments, we show the median of 100 runs.
Revised: Similar to the synthetic experiments, we show the median of 100 runs.

Original: While the plot shows smaller variations, it is generally inline with the observations we had in the synthetic data.
Revised: While the plot shows smaller variations, it is generally inline with the observations we had in the synthetic data.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: Causal inference is a well-studied problem; its main goal is to remove biases due to confounding by balancing the population to look similar to randomized controlled trials.
Revised: Causal inference is a well-studied problem; its main goal is to remove biases due to confounding by balancing the population to look similar to randomized controlled trials.

Original: Removing the impact of confounders can play a critical role in reducing and possibly eliminating bias in our decision making leading to potentially positive societal impacts.
Revised: Removing the impact of confounders can play a critical role in reducing and possibly eliminating bias in our decision making leading to potentially positive societal impacts.

Original: Our results rely on two classical assumptions: (1) unconfoundedness and (2) positivity.
Revised: Our results rely on two classical assumptions: (1) unconfoundedness and (2) positivity.

Original: While these assumptions are sometimes reasonable in practice, their violations might lead to biased causal inferences.
Revised: While these assumptions are sometimes reasonable in practice, their violations might lead to biased causal inferences.

Original: For example, the positivity assumption might be violated if we do not collect any data for a sub-population.
Revised: For example, the positivity assumption might be violated if we do not collect any data for a sub-population.

Original: Overall, the debiasing property of causal inference should not relieve us from rigorous data collection and analysis setup.
Revised: Overall, the debiasing property of causal inference should not relieve us from rigorous data collection and analysis setup.

Original: In our experiments, we have been careful to quantify uncertainty in our causal estimation and be wary of over-conﬁdence in our results.
Revised: In our experiments, we have been careful to quantify uncertainty in our causal estimation and be wary of over-conﬁdence in our results.

Original: We performed our experiments on a CPU machine with 16 cores from a cloud provider that uses hydroelectric power.
Revised: We performed our experiments on a CPU machine with 16 cores from a cloud provider that uses hydroelectric power.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We observed that in the entropy balancing framework, the base weights provide an extra degree of freedom to optimize the accuracy of causal inference.
Revised: We observed that in the entropy balancing framework, the base weights provide an extra degree of freedom to optimize the accuracy of causal inference.

Original: We propose end-to-end balancing (E2B) as a technique to learn the base weight such that they directly improve the accuracy of causal inference using end-to-end optimization.
Revised: We propose end-to-end balancing (E2B) as a technique to learn the base weight such that they directly improve the accuracy of causal inference using end-to-end optimization.

Original: In our theoretical analysis we ﬁnd the quantity that E2B weights are approximating and discuss E2B’s statistical consistency.
Revised: In our theoretical analysis we ﬁnd the quantity that E2B weights are approximating and discuss E2B’s statistical consistency.

Original: Using synthetic and real-world data, we show that our proposed algorithm outperforms the entropy balancing in terms of causal inference accuracy.
Revised: Using synthetic and real-world data, we show that our proposed algorithm outperforms the entropy balancing in terms of causal inference accuracy.