Energy balancing of covariate distributions

2.1 Setup

Consider a sample { ( Y i , A i , X i ) } i = 1 n of size n from a population, where Y i is the outcome of the i th unit, A i ∈ { 0 , 1 } is a binary indicator of receiving a treatment, and X i ∈ X = R p is a p -dimensional vector of covariates. Further denote n 1 = ∑ i = 1 n A i and n 0 = n − n 1 . In this article, we are interested in estimating the average causal effect of the treatment on the outcome. A formal definition of a causal effect often involves the use of so-called potential outcomes [26–29]. The potential outcome Y ( a ) is the outcome that would have been observed under level a of the treatment. As each individual only receives one level of the treatment at a given time, only one potential outcome, either Y i ( 0 ) or Y i ( 1 ) , for each individual is observable. We assume the standard stable unit treatment value assumption (SUTVA), which posits that the potential outcomes for each unit are unaffected by the potential outcomes of other units and that only one version of the treatment exists. Under SUTVA, the observed outcome is consistent with the potential outcomes in that Y i = Y i ( A i ) . We further assume the assignment mechanism is strongly unconfounded in the sense that { Y ( 0 ) , Y ( 1 ) } ⊥ ⊥ A ∣ X , which requires that there are no unmeasured confounders. The ⊥ ⊥ notation of Dawid [30] denotes (conditional) independence. We further assume positivity (or probabilistic assignment) in that the propensity score π ( x ) ≡ P ( A = 1 ∣ X = x ) [10] satisfies 0 < π ( x ) < 1 , so that everyone has a chance of receiving the treatment. Positivity, together with no unmeasured confounders, constitute strong ignorability [10].

Let us denote μ a ( X i ) ≡ E ( Y ( a ) ∣ X i ) as the conditional mean function and σ a 2 ( X i ) ≡ V ( Y i ( a ) ∣ X i ) as the conditional variance function of the response for a ∈ { 0 , 1 } . We consider scenarios where data have been collected from an observational study, and thus, the treatment groups are not comparable due to imbalances in the distributions of their baseline covariates. These differences can be characterized through π ( x ) or through F a ( x ) ≡ P ( X ≤ x ∣ A = a ) , the cumulative distribution function (CDF) of covariates X conditional on treatment level a .

Using the notation of potential outcomes, the (population) ATE is defined as τ ≡ E ( Y ( 1 ) − Y ( 0 ) ) . This can be rewritten as follows:

where F ( x ) = P r ( X ≤ x ) is the CDF of covariates X marginalized over the treatment groups. In other words, F ( x ) = F 1 ( x ) P 1 + F 0 ( x ) P 0 , where P a ≡ P ( A = a ) ∈ ( 0 , 1 ) is the probability of being assigned treatment level a ∈ { 0 , 1 } .

2.2 Weighted average estimates and distributional balance

We restrict our focus to weighted averages as estimates of the ATE. Given a vector of weights w = ( w 1 , … , w n ) , we study estimators of the form:

The most commonly used example of (2), inverse propensity score weighting, uses w i = ( A i π ( X i ) n 1 ∕ n + ( 1 − A i ) ( 1 − π ( X i ) ) n 0 ∕ n ) − 1 . As presented in Imai and Ratkovic [12] and Li et al. [20], the weights in (2) are often normalized by treatment group, i.e., ∑ i = 1 n w i I ( A i = a ) = n a for a ∈ { 0 , 1 } , to improve precision [11,16] at the cost of a small bias. When these weights are constructed as mentioned earlier and normalized by treatment group, the resulting estimator is called the Hájek estimator [31] and may yield reduced mean squared error over its nonnormalized version [32].

Given any weight vector w , we can express the error of τ ^ w as follows:

where ε i ≡ Y i ( A i ) − μ A i ( X i ) , F n ( x ) = ∑ i = 1 n I ( X i ≤ x ) ∕ n is the empirical CDF (ECDF) of the combined sample { X i } i = 1 n , and F n , a , w ( x ) = ∑ i = 1 n w i I ( X i ≤ x , A i = a ) ∕ n a is the weighted ECDF for treatment level a ∈ { 0 , 1 } . In observational studies, F n is not impacted by the weights, so the error term (4) is irreducible. However, this term goes to 0 as long as the sample is representative of the desired population. Since (5) always has mean 0, the bias of τ ^ w in essence depends on the properties of (3): the difference of integrals with respect to F n , a , w − F n . Thus, the systematic source of error from the weighted estimator τ ^ w can be completely controlled by reducing the imbalance between the weighted ECDFs F n , a , w and the ECDF F n . Unlike the decomposition in Zhao [16], the decomposition of τ ^ w − τ above holds even if the treatment effect is not constant over x . Further, none of the methodology or results of this article require such a constant treatment effect to justify the validity of our methods. We note that the systematic source of bias in equation (3) can be zero without F n , a , w being balanced to F n , depending on μ 0 ( x ) and μ 1 ( x ) . However, knowledge of such an event a priori is impossible without knowledge of the mean potential outcome functions; as such, a measure of distributional imbalance is critical to characterize the degree of bias for a given dataset. The terms in (5) drives the variability of τ ˆ w and can be straightforwardly mitigated with a measure of dispersion of the weights [33]; we discuss this more at the end of Section 3; however, this is not the focus of this article.

The notion that the covariate distributions should be balanced to obtain a good estimate of τ is not new (see Imai and Ratkovic [12] and Li et al. [20], among many others). In fact, it is well understood that the weights resulting from correctly specified propensity score models asymptotically balance covariate distributions. However, slight model misspecifications of the propensity score may result in poor performance of (2) [11]. Further, two units with the same propensity score do not necessarily have the same covariate values, so in finite samples, propensity score weights may not be optimal. Other approaches that aim to estimate weights by balancing prespecified moments of the covariates [12,14] tend to be more robust than directly modeling π ( x ) . Yet, the term (3) makes it explicit that bias is directly related to distributional imbalance and that, holding μ a fixed, weights that yield less distributional imbalance will in general result in less bias. A natural conclusion is that w should be estimated to directly balance each F n , a , w to F n . In the following, we introduce a new distance metric between F n , a , w and F n , which enables one to characterize how well a set of weights re-balances the weighted distribution F n , a , w to F n .

2.3 Weighted energy distance

We introduce next a new measure of the distributional balance induced by a set of weights. This measure is based on the energy distance, which is a metric on distributions [34]. Due to the link between estimation bias and distributional imbalance, our measure can be used to evaluate the degree of bias one expects from a given set of weights and a given dataset. We will later leverage this measure to construct distributional balance weights that minimize this metric for a given dataset.

The energy distance (as surveyed in Székely and Rizzo [34]) is defined as follows. Let G and H be two finite-mean distribution functions on X , and let Z , Z ′ ∼ i . i . d G and V , V ′ ∼ i . i . d H . The energy distance between distributions G and H is defined as follows:

where ‖ ⋅ ‖ 2 is the Euclidean norm. When both G and H are ECDFs, i.e., G n is the ECDF of { Z i } i = 1 n ⊆ X and H m is the ECDF of { V i } i = 1 m ⊆ X , the energy distance ℰ ( G n , H m ) can be expressed as follows:

The energy distance has been used within a wide variety of statistical methods, e.g., for testing equivalence of distributions, for testing statistical independence [35], and for generating samples from a target distribution [36]. The energy distance is simple to compute as it only involves sums of Euclidean distances, it is nonnegative, and it can be shown to be equivalent to a distance between the characteristic functions of G and H , indicating that a value of 0 occurs if and only if G and H are the same distribution and a positive value indicates how similar they are. The name “energy” comes from the fact that the population energy distance (or limit of the empirical energy distance) has an intricate connection with gravitational potential energy (which depends on the distance between two bodies of mass). Two key attractive properties of the energy distance are that it tends to perform well in measuring the distance between two distributions in moderately high-dimensional settings, and it does not require the choice of a kernel or any tuning parameters.

We propose a weighted modification of this energy distance, which measures the distance between a weighted distribution, i.e., the weighted covariate ECDFs F n , 0 , w and F n , 1 , w for the control and treated, and a target distribution, i.e., the combined covariate ECDF F n . The weighted energy distance between F n , a , w and F n is defined as follows:

In other words, ℰ ( F n , a , w , F n ) is the energy distance between the ECDF of a sample { X i } i = 1 n and a weighted ECDF of a subsample { X i } i : A i = a .

We show this new weighted energy distance is indeed a distance between F n , a , w and F n . Let ⟨ t , s ⟩ be the inner product of vectors t and s and define the empirical characteristic function (ECHF) of { X i } i = 1 n and the weighted ECHFs of { X i } i : A i = a as follows:

respectively. The following proposition establishes the distance property of the weighted energy distance.

Thus, the weighted energy distance is a distance between the weighted distribution of interest and the target distribution, and is thus a bona fide measure of distributional balance of covariates induced by a set of weights. This proportion extends the duality results of Proposition 1 from Székely and Rizzo [34] and Theorem 1 from Székely et al. [35] for the weighted energy distance at hand. A subtle point is that Proposition 2.1, combined with the decomposition presented in Section 2.2, makes clear that ℰ ( F n , a , w , F n ) more closely aligns with evaluating how well a set of weights estimates the sample average treatment effect [32] than the population ATE τ that is our main focus. Yet, as long as F n is representative of F , the weighted energy distance still aligns well with the population ATE.

We now show that the weighted energy distance converges to the energy distance when the weights yield a well-defined limiting distribution.

Theorem 2.2 shows that the weighted energy distance converges to the limiting energy distance. If the limiting distribution implied by a set of weights is F ˜ 0 = F ˜ 1 = F , then ℰ ( F ˜ 0 , F ) + ℰ ( F ˜ 1 , F ) = 0 . Proposition 2.1 and Theorem 2.2 together imply that weights with smaller values of the sum ℰ ( F n , 1 , w , F n ) + ℰ ( F n , 0 , w , F n ) yield better distributional balance of covariates. Due to the link between imbalance and bias (see the following section), this also implies that better balance will yield estimates with smaller values for the term (3).

2.4 Bias and distributional imbalance

We now demonstrate this connection between bias and distributional imbalance (as measured by the weighted energy distance) using two illustrative examples. A more formal presentation is provided later in Section 3.2 when proving asymptotic properties of EBWs.

In the first illustrative example, we generate a one-dimensional covariate of sample size 250, which impacts treatment assignment via a logistic model under each of three scenarios: (1) logit ( π ( X ) ) = − 1 + X , (2) logit ( π ( X ) ) = − 1 + X + 2 X 2 ∕ 3 , and (3) logit ( π ( X ) ) = − 1 + X + 2 X 2 ∕ 3 − X 3 ∕ 3 . In each scenario, the response is generated as Y = X + X 3 − 1 ∕ ( 0.1 + 0.1 X 2 ) + ε , where ε ∼ N ( 0 , 2 ) . For each scenario, we construct IPWs based on three logistic regression models, which consider only a linear term in X (denoted as “IPW (1)”), a linear plus quadratic term (“IPW (2)”), and up to the cubic term (“IPW (3)”), respectively. Thus, for Scenarios 2 and 3, at least one of the fitted models is misspecified. For each set of weights w , we compute ℰ ( F n , 0 , w , F n ) + ℰ ( F n , 1 , w , F n ) , i.e., the sum of the energy distances between each treatment group and the combined sample, and compute the error τ ^ w − τ of (2) for τ .

Figure 1(a) displays the energy distances and biases over 1,000 replications of the experiment. We see that the energy distances are the largest in all scenarios for the unweighted estimator ((2) with all weights equal to 1). For scenarios where the weights are estimated using a misspecified model (IPW (1) in Scenario 2 and IPW (1) and (2) in Scenario 3), the energy distances are much larger than for weights based on correctly (or over-specified) models. Correspondingly, the bias is pronounced for the misspecified models. Thus, the weighted energy distance can be a useful tool to compare between different models, as weights with smaller weighted energy distances tend to yield estimates with smaller error.

Figure 1

(a, left) Energy distances and biases for IPW estimates based on weights from the three fitted logistic regression models; (b, right) Boxplots of the biases for IPW estimates versus weighted energy distance based on weights estimated by several methods, each with different combinations of moments included for balancing or estimation.

In the second example, we consider a two-dimensional example where the true assignment mechanism depends on first and second moments of the covariates. We consider several methods for estimate weights: logistic regression, the method of Imai and Ratkovic [12], and the method of Chan et al. [14], each with different moments included for balancing or estimation. We then compare their weighted energy distances and absolute errors of (2) over 1,000 replications. Figure 1(b) displays the distances and errors for each dataset and method. We see that, in general, weights with lower energy distance have a much smaller magnitude of bias in estimating the ATE.

2.5 Other measures of distributional imbalance

There are, of course, many ways of measuring distance between distributions in the literature, including the Kolmogorov-Smirnov statistic, f -divergences (e.g., the Kullback-Leibler divergence and the Hellinger distance), the Wasserstein distance, and the maximum mean discrepancy (MMD). The energy distance has several advantages for characterizing distributional imbalance. First, while in general there is no uniformly most powerful nonparametric test for the difference between two distributions, the energy distance is often sensitive to differences in distributions, unlike the Kolmogorov-Smirnov statistic. Second, unlike the energy distance, f -divergences such as the Kullback–Leibler divergence and Hellinger distance do not metrize weak convergence; this property ensures that the distance remains stable under perturbations of the support of the distributions being measured [37]. Third, the energy distance is easy and efficient to compute, unlike the Wasserstein distance. It also works reasonably well in moderately high dimensions [36], whereas the Wasserstein distance suffers from the curse of dimensionality, in the sense that empirical Wassterstein distances converge to the true Wasserstein distance at rate O ( n − 1 ∕ p ) [37,38]. Finally, while there is a direct link between MMD and the energy distance [39], the energy distance does not require careful tuning of hyperparameters and tends to work well across a wide variety of scenarios.

3.1 Definition

We will now use the proposed weighted energy distance to estimate weights which (i) match the distribution of covariates of the treated group to the distribution of covariates of the full population, and (ii) match the distribution of covariates of the control group to the distribution of covariates of the full population.

To achieve this, we define the EBWs to be

Thus, the EBWs w n e minimize the statistical energy between the treatment and control groups to the full population. Due to the duality result of Proposition 2.1, minimizing statistical energy is directly equivalent to balancing covariate distributions. The constraints ∑ i = 1 n w i I ( A i = a ) = n a serve several purposes: they (i) preserve the sample size of the weighted pseudo-population to that of the study population, (ii) stabilize the estimated weights, and (iii) ensure that F n , 0 , w and F n , 1 , w are valid distribution functions. Also note that, due to the bias decomposition in (3) and the duality result in Proposition (2.1), the weights w n e are explicitly designed to minimize the key component of the finite-sample bias of τ ^ w n e , in a manner agnostic to the functional forms of μ 0 ( x ) and μ 1 ( x ) . If the functional forms of μ 0 ( x ) and μ 1 ( x ) are known, it is certainly possible to construct a different set of weights to better reduce bias, by emphasizing balance in regions of X where either μ 0 ( X ) or μ 1 ( X ) are pronounced. However, this information is rarely (if ever) known, and hence this is difficult to achieve in practice except by good fortune.

To illustrate the effectiveness of EBWs for distributional balance, we consider data generated under Scenario 3 of the toy example in Figure 1(a). Figure 2 shows the difference between the weighted ECDF using EBWs of the covariate in the treatment group and the ECDF of the combined (i.e., treated and untreated) sample, for varying sample sizes n . This is compared with the weighted ECDF using weights from the true data-generating propensity score, the estimated propensity score under the correct model, and the unweighted ECDF of the treatment group. As sample size increases, the difference vanishes for EBWs, but much more slowly for both the true and estimated propensity score weights. This demonstrates the improved distributional balance provided by the proposed EBWs, which should then translate to a greater reduction of bias.

Figure 2

Displayed are the absolute differences between the ECDF of the combined sample and the weighted ECDF of the covariate in the treatment group based on energy weights. Also displayed the same for the unweighted ECDF of the treated group and the weighted ECDF based on the true and estimated propensity scores.

EBWs can be naturally extended to handle a wide variety of scenarios, such as for treatments with more than two levels, estimation of the ATT, and for the estimation of optimal ITRs. In the Supplementary Material, we show how these extensions are manifested and empirically demonstrate the benefit of using EBWs for ITR estimation.

A few key distinguishing features of our proposal from the works of Wong and Chan [18] and Kallus [19] are (1) its direct focus on distributional balance rather than moment balance. Despite the connection between balancing moments of an infinite dimensional class of functions and distributional balance, this feature helps alleviate modeling decisions about what moments to balance of what space of moments to balance and has advantages in terms of interpretability and (2) by focusing on a measure that does not require tuning parameters to characterize distributional differences, our approach can be applied broadly by practitioners of varying degrees of statistical sophistication.

3.2 Asymptotic properties

Next, we show two desirable properties of the proposed EBWs. We first show that the weighted ECDFs based on EBWs indeed converge to the population CDF F of X .

Thus, EBWs result in the almost sure convergence of the weighted ECDFs of the treated (and untreated) group to the underlying covariate distribution F .

The consistency of τ ^ w n e follows immediately from Theorem 3.1:

Next, we show that the ATE estimator τ ^ w n e is root- n consistent. To do so, we make use of the following lemma, which provides a connection between bias and distributional balance:

Lemma 3.3 provides a connection between the systematic bias in (3) and the weighted energy distance ℰ ( F n , a , w , F n ) . Under the mild condition that the conditional mean function μ a ( x ) is found in ℋ (this is discussed further at the end of the section), this lemma shows that the weighted energy distance is a key component in an upper bound on the systematic bias in (3). We note that Lemma 3.3 applies to any set of auto-normalized weights, justifying the use of the weighted energy distance to compare between different sets of weights. The proposed EBWs w n e , which minimize ℰ ( F n , a , w , F n ) , may therefore yield lower estimation bias via this upper bound, which is in line with the empirical observations in Section 2.4. With this, we now state the result on root- n consistency:

We give a brief discussion of Assumptions (A1)–(A5). Assumption (A1) concerns the regularity of the conditional mean functions μ 0 and μ 1 . It can be shown that the Sobolev space W ⌈ ( p + 1 ) ∕ 2 ⌉ , 2 ( X ) – the space of functions with square-integrable r < ⌈ ( p + 1 ) ∕ 2 ⌉ -th differentials – is contained within the native space ℋ ( X ) (Theorem 10.42 of [40]), so (A1) can be viewed as a smoothness assumption on the conditional ATE μ 1 − μ 0 (a similar assumption is made in [18]). Assumptions (A2) and (A3) require the conditional mean functions to have finite variance and the conditional variance function to be bounded, respectively. Assumption (A4) require the kernels h 0 and h 1 to have finite second moments; these kernels can be seen as a “modified” Euclidean kernel weighted by the true propensity scores. Assumption (A5) assumes all EBWs to be bounded above by C n 1 ∕ 3 for some positive constant C > 0 ; this assumption has been used in several weighting-based covariate balancing methods [18,25]. In practice, (A5) can always be checked after optimizing for the EBWs in (10). One can change the optimization procedure (10) to explicitly enforce (A5), but we have never encountered “exploding weights,” which violate (A5) in practice; EBWs in our simulations and data analysis typically satisfy (A5) with C = 1 . As we discuss in Section 3.4, if one is willing to add a penalty on the dispersion of the weights, the explicit assumption on the maximum of the weights is unnecessary and root- n consistency still holds, as shown in the Supplement.

While Theorem 3.4 proves the desired root- n consistency of the proposed EBW estimator τ ^ w n e , it unfortunately does not shed light on its asymptotic variance, which is useful for constructing confidence intervals on τ . We recommend the use of bootstrapped confidence intervals for the EBW estimator. We explore a connection to importance sampling in the Supplementary Material.

3.3 Optimization

The optimization problem (10) for computing EBWs (and the later optimization problem (16) for obtaining three-way EBWs) can be viewed as quadratic programs with linear (in)equality constraints. There has been much work on efficient algorithms for solving such programs [41], including interior point methods [42], augmented Lagrangian techniques [43], and extensions of the simplex algorithm [44]. A recent development is the operator splitting approach in Stellato et al. [45], which provides a reliable alternative for nonpositive definite quadratic programs.

In our implementation, we made use of well-maintained interior-point cone programming solvers in the R package cccp [46] for optimizing the EBW formulations (10) and (16). Such solvers follow a two-stage procedure: finding an initial feasible solution for w , then iteratively refining this solution by traversing the interior of the feasible region. Detailed algorithmic steps can be found in Wright et al. [47]. Interior-point algorithms enjoy empirical and theoretical advantages for solving large-scale quadratic programs (see further discussion in Gondzio [48] on its scalability and complexity), and we have observed excellent performance of such algorithms for EBW optimization. In practice, we recommend that the covariates should be normalized so that each covariate has mean zero and unit variance, before being used as inputs for the optimization problem. An efficient implementation of these methods for EBW optimization is provided in our R package ebw, which will be released on comprehensive R archive network in the near future.

3.4 Controlling weight variability

In our experience, the weights resulting from the optimization criterion (10) rarely, if ever, result in large weights; however, in the literature there has been an emphasis on methods that afford explicit control on the variability of weights [33]. To allow for such within our framework, one can simply add a penalty λ n − 2 ∑ i = 1 n w i 2 to the criterion (10) for some fixed λ > 0 . This penalty can be thought of as a penalty on the inverse of the effective sample size of the weights [33]. In our experience, since the energy distance criterion is not nearly as variable as using a kernelized distance using a universal kernel such as the Gaussian kernel, careful tuning of λ is not at all critical and we advocate for a simple choice such as λ = 1 or not using any penalty at all. On the other hand, if one were to replace the energy distance with a kernelized distance using a universal kernel, the choice of λ becomes critical. In the Supplement, we show analogs of Theorems 3.1 and 3.4 hold when using a penalized energy distance criterion, although the conditions regarding the magnitude of the weights in the latter theorem can be relaxed.

3.5 Three-way EBWs

The EBWs in (10) are designed to balance the distributions of covariates between each treatment group to that of the combined sample { X } i = 1 n . As such, the treatment group should be asymptotically balanced to the control group. Yet for finite samples, EBWs do not necessarily guarantee good distributional balance between the treatment and control arms, which can be important in practice. Consider the following re-expression of the two terms in (3):

Terms (14) and (15) shed light on how the choice of w impacts estimation error for the ATE. In particular, this error term depends not only on (i) how close the weighted ECDFs F n , 0 , w and F n , 1 , w are to that of the combined sample F n but also (ii) how close the weighted ECDF for the control group F n , 0 , w is to the weighted ECDF of the treated group F n , 1 , w . The EBWs in Section 3 take care of the imbalance in (i), but not the imbalance in (ii) between treatment and control. With finite samples, imbalance in (ii) can result in (14) dominating (15), depending on the properties of μ 0 ( x ) , μ 1 ( x ) , and μ 0 ( x ) + μ 1 ( x ) . Indeed, the importance of this three-way balance is recognized in the literature, and is emphasized in Chan et al. [14] as an important component in constructing globally efficient estimators based on moment balancing. In our case, the target is the three-way distributional balance, i.e., balance in (i) and (ii).

We propose an extension of EBWs, the improved EBWs (iEBWs), to help improve covariate balance between the treatment and control groups. These improved weights are defined as follows:

where

is the energy distance between the weighted ECDFs for treated and control. Thus, the iEBWs w n e i aim to minimize imbalance not only between treatment arms and the full population but also between the treatment arm and the control arm. Note that ℰ ( F n , 1 , w , F n , 0 , w ) still retains the properties of a weighted energy distance, in the sense that ℰ ( F n , 0 , w , F n , 1 , w ) = ∫ R p ∣ φ n , 1 , w ( t ) − φ n , 0 , w ( t ) ∣ 2 ω ( t ) d t , as Proposition 2.1 can be trivially extended to the case where both arguments of the energy distance are weighted.

4.1 Estimation of the ATT

A common target of estimation is the (population) ATT, τ ( 1 ) ≡ E ( Y ( 1 ) − Y ( 0 ) ∣ A = 1 ) , which is the mean difference in potential outcomes among those who are actually treated. Due to the unconfoundedness assumption, we can write

which suggests that a plug-in estimator can be obtained by replacing F 1 ( x ) with a suitable energy-weighted ECDF. To do so, we define the new weights w n e 1 = ( w 1 e 1 , … , w n e 1 ) , where w i e 1 = 1 for { i : A i = 1 } and

where F n , 1 ( x ) = 1 n 1 ∑ i = 1 n A i I ( X i ≤ x ) . Thus, w n e 1 balances the covariate distribution for the control group to the covariate distribution of the treated group. Similar to Theorem 3.1, we have lim n → ∞ F n , 0 , w n e 1 ( x ) = F 1 ( x ) a.s. for every continuity point x ∈ X . With the new weights w n e 1 , the natural plug-in estimate for the ATT τ ( 1 ) is τ ˆ w n e 1 (i.e., the weighted estimator in (2) with w = w n e 1 ).

4.2 Multi-category treatments

EBWs can also be constructed for multi-category treatments. When the treatment A i takes multiple values, i.e. A i ∈ A = { a 1 , … , a K } . Denote n a = ∑ i = 1 n I ( A i = a ) , for a ∈ A . Each unit has K potential outcomes ( Y i ( a 1 ) , … , Y i ( a K ) ) , one for each level. The unconfoundedness assumption in this setting becomes { Y i ( a 1 ) , … , Y i ( a K ) } ⊥ ⊥ A ∣ X . Following Lopez and Gutman [49], the positivity assumption required is now 0 < π ( a , x ) = P ( A = a ∣ X = x ) < 1 for all a ∈ A and all possible x ∈ X . The standard IPW estimator for E [ Y ( a ) ] involves inverse weighting each sample i by π ( A i , X i ) . Instead, we propose to estimate E [ Y ( a ) ] or any causal contrast τ ( a − a ′ ) ≡ E [ Y ( a ) − Y ( a ′ ) ] for a , a ′ ∈ A with EBWs.

We define the EBWs for the multiple treatment case as follows:

Improved EBWs, which encourage covariate balance between all pairs of treatment options, can be defined similarly as (16), where an additional weighted energy distance between each pair of treatment options is added to the objective. Given any two treatment options a , a ′ ∈ A , we can then estimate the causal contrast τ ( a − a ′ ) with

4.3 Estimation of ITRs

As many treatments exhibit heterogeneous effects for different patients, there is great interest in tailoring treatment decisions to patients. A main line of work in this area is the development of statistical methods aimed at finding an optimal ITR, which maps patient characteristics to treatment decisions. Thus, the immediate goal is to estimate a mapping d : X ↦ { 0 , 1 } which optimizes the expected potential outcomes under the distribution induced by d . Following Qian and Murphy [23], and Zhao et al. [24] and assuming that larger values of the outcome are preferred, the optimal ITR is defined as

where π ( a , X ) = P ( A = a ∣ X ) and the second equality holds due to the causal assumptions outlined in Section 2.1. The optimal ITR d * has the property that d * ( x ) = a ⇒ μ a ( x ) > μ 1 − a ( x ) . The last term in (19) appears as a weighted classification task due to the weighted 0–1 loss. With observed data, the objective becomes to minimize

Due to the nonconvexity of (20), in practice I ( A i ≠ d ( X i ) ) is replaced with a surrogate, such as the hinge function ( 1 − ( 2 A i − 1 ) d ( X i ) ) + (see, e.g., the outcome weighted learning (OWL) method of Zhao et al. [24]). Yet, this objective still requires estimation of the propensity score π ( X ) and involves plugging in this estimate to (20), which subjects it to the same issues of propensity score weighted estimates of τ . To see how EBWs can be used in place of π ( A , X ) , we can express (19) as argmin d ∫ x ∈ X ∑ a ∈ { 0 , 1 } μ a ( x ) I ( d ( x ) ≠ a ) d F ( x ) . Thus, d * ( x ) is a functional of F ( x ) . This suggests a plug-in estimator by replacing F ( x ) with energy-weighted ECDFs (either F n , a , w n e ( x ) or F n , a , w n e i ( x ) ). Thus, we propose to estimate the optimal ITR as follows: