Title: An effective framework for estimating individualized treatment rules
Venue: NeurIPS 2024
Accepted: yes
Keywords: Covariate Balancing

Multicategory Treatment

Precision Medicine

Projected Gradient Descent

Convergence Analysis

Constrained Optimization
Abstract: Estimating individualized treatment rules (ITRs) is fundamental in causal inference, particularly for precision medicine applications. Traditional ITR estimation methods rely on inverse probability weighting (IPW) to address confounding factors and $L_{1}$-penalization for simplicity and interpretability. However, IPW can introduce statistical bias without precise propensity score modeling, while $L_1$-penalization makes the objective non-smooth, leading to computational bias and requiring subgradient methods. In this paper, we propose a unified ITR estimation framework formulated as a constrained, weighted, and smooth convex optimization problem. The optimal ITR can be robustly and effectively computed by projected gradient descent. Our comprehensive theoretical analysis reveals that weights that balance the spectrum of a `weighted design matrix' improve both the optimization and likelihood landscapes, yielding improved computational and statistical estimation guarantees. In particular, this is achieved by distributional covariate balancing weights, which are model-free alternatives to IPW. Extensive simulations and applications demonstrate that our framework achieves significant gains in both robustness and effectiveness for ITR learning against existing methods.

An effective framework for estimating individualized treatment rules

Abstract

Estimating individualized treatment rules (ITRs) is fundamental in causal inference, particularly for precision medicine applications. Traditional ITR estimation methods rely on inverse probability weighting (IPW) to address confounding factors and L 1 -penalization for simplicity and interpretability. However, IPW can introduce statistical bias without precise propensity score modeling, while L 1 -penalization makes the objective non-smooth, leading to computational bias and requiring subgradient methods. In this paper, we propose a unified ITR estimation framework formulated as a constrained, weighted, and smooth convex optimization problem. The optimal ITR can be robustly and effectively computed by projected gradient descent. Our comprehensive theoretical analysis reveals that weights that balance the spectrum of a 'weighted design matrix' improve both the optimization and likelihood landscapes, yielding improved computational and statistical estimation guarantees. In particular, this is achieved by distributional covariate balancing weights, which are model-free alternatives to IPW. Extensive simulations and applications demonstrate that our framework achieves significant gains in both robustness and effectiveness for ITR learning against existing methods.

1 Introduction

Traditional medicine often uses a "one-size-fits-all" approach where the same treatment is applied to all patients regardless of their unique attributes, aiming to find a single optimal treatment that may be most effective for a broad population. However, since not everyone fits the mold, precision medicine has been popular in medical research, emphasizing personalized treatment based on a patient's unique characteristics [11].

One key element of precision medicine is the estimation of individualized treatment regime (ITR), also known as treatment policy or policy. ITRs consider various factors like demographics and sociopsychological aspects to optimize treatment decisions, thereby maximizing individual outcomes. Many frameworks have been proposed for learning ITRs [41, 54, 57, 7, 58, 60, 56] and the related conditional average treatment effect (CATE) estimation problem [3, 53, 46, 13, 36]. For model interpretation and other practical concerns, researchers often focus on pre-specified rule classes. For example, the method proposed in [2] focuses on estimating ITRs using shallow-depth decision trees.

The central task in the ITR estimation is controlling confounding factors to isolate the causal effect of treatment from other factors that may influence the outcome. A standard approach is by weighting each sample using their inverse probability weights (IPWs), which requires specifying a propensity score model. However, propensity score methods have long been known to be highly sensitive to model misspecification, which yields biased estimates of causal effects [25, 30]. Distributional covariate balancing weights (DCBWs) [24, 27] are modern alternatives to IPWs, directly minimizing the distance between empirical covariate distributions induced by the weights.

ITR estimation can be framed as an optimization or maximum likelihood estimation (MLE) problem, where the choice of weights directly impacts both the optimization and likelihood landscapes. This, in

turn, affects the computational and statistical errors in the ITR estimation, as well as the performance of the optimization algorithm. The key contribution of this work is to bridge the gap between causal inference and optimization perspectives in ITR estimation. Our main finding is both simple and surprising: Covariate balancing weights not only address confounding but also yield near-optimal optimization and likelihood landscapes for the ITR estimation. We show this claim by formulating the ITR estimation problem as a constrained, weighted, and smooth convex optimization problem. Our formulation is grounded in angle-based direct learning (AD-Learning) framework [39], which is a recent framework for multi-category treatments with a linear decision function class assumption. While we focus on multi-category treatments-a relatively underexplored area-we believe that the principles are equally applicable to binary treatments with a linear decision function.

Besides advancing our theoretical understanding of the impact of covariate balancing weights, we propose a unified computational and statistical framework for the ITR estimation. Notably, we introduce a hard L 1 -ball constraint to promote sparsity in regression coefficients, replacing the traditional soft L 1 -penalization. This approach keeps the optimization objective smooth, enabling fast and accurate ITR estimation using projected gradient descent (PGD). Furthermore, we integrate techniques such as (1) variable screening, (2) outcome augmentation, and (3) inverse variance weighting to account for heteroscedastic errors, which together show a synergistic effect.

Related Work. Most existing ITR approaches use IPW, which requires specifying a propensity score model for confounding control. Recent developments in causal inference have also introduced robust approaches that directly estimate weights based on balance-seeking objectives rather than relying solely on propensity scores. These approaches include entropy balancing weights [20], stable balancing weights [61] as well as DCBWs [24, 27]. However, despite the weights serving a critical role in the estimation process, these advancements have mainly focused on improving average treatment effect (ATE) and are largely under-explored in the ITR literature. To the best of our knowledge, this is the first work to apply DCBWs for ITR-Learning under a weighted optimization framework. While recent works have taken a similar approach by utilizing weighting schemes and directly optimizing the weights rather than relying on estimated IPW [28, 26], our method differs in key aspects. Unlike [26], which involves the evaluation of policy effects, our approach directly learns the optimal decision function, bypassing the need for intermediate policy effect evaluations. Although [28] introduced the concept of retargeting the population covariate distribution, which is conceptually related to our analysis, their approach relies on a reference policy, whereas our method does not. We also explore various statistical techniques to improve ITR-Learning. First, we perform variable screening using the distance covariance test [48] to retain key effect modifiers and precision variables that directly impact the outcome. Second, [50, 9] introduced outcome augmentation to minimize estimator variance in clinical trials and observational studies by IPW involving binary treatments. Building on these foundations, our work extends these approaches to a broader range of weighting schemes, such as DCBW particularly for multi-category treatments. Lastly, with inverse variance weighting, as introduced in stabilized angle-based direct learning (SABD-Learning) to address heteroscedasticity [44], this combined strategy enhances estimator precision and robustness.

Contributions. We summarize our contributions below.

Unified Framework of ITR-Learning : We introduce a novel framework that addresses the limitations of existing ITR estimation methods by formulating the problem as a constrained, weighted, and smooth convex optimization problem. Under the framework, we propose a PGD algorithm for ITR-Learning under a hard L 1 -constraint.

Improved Computational and Statistical Guarantees : We establish convergence guarantees (Theorem 3.3). Furthermore, under mild assumptions, we demonstrate that the parameter of ITR-Learning can be consistently estimated with high probability (Theorem 3.5) and provide computational and sample complexity. In both cases, we demonstrate that using covariate balancing weights controls confounding factors and leads to better optimization and likelihood landscapes, resulting in improved computational and statistical performance.

Statistical Framework of ITR-Learning : Our main contribution lies in providing a unified framework that combines DCBWs, variable screening, outcome augmentation, and inverse variance weighting in a synergistic and effective way. We demonstrated theoretical justifications for the combined approach, which, to our knowledge, have not been previously established.

2 Methods for ITR Estimation

2.1 Model Setup

Individualized Treatment Rule and Optimal Decision Function. We observe a random vector ( X , A, Y ) where X = ( X 1 , X 2 , . . . , X p ) ∈ X denotes the p -dimensional vector of pre-treatment covariates, a received treatment A ∈ A = { 1 , 2 , . . . , K } , and the corresponding outcome Y . Without loss of generality, we assume higher Y values are more favorable. Let Y ( a ) denote the random variable that describes the potential outcome that would have been observed were the individual assigned to treatment a ∈ A . We make the following standard assumptions using the potential outcome framework [43]: (i) Stable Unit Treatment Value Assumption (SUTVA): Y = Y ( A ) , (ii) positivity: 0 < π ( a, X ) := P ( A = a | X ) < 1 for all a ∈ A and all X ∈ X , (iii) no unmeasured confounding: Y ( a ) ⊥ A | X for any a ∈ A .

An ITR, d ( x ) , is a function mapping each covariate x to one of the K treatments. According to [41], the ITR d can be measured by the value function V ( d ) :

Then the optimal ITR d opt is defined as a function that maximizes the expected potential outcome d opt ( x ) ∈ arg max d ∈D V ( d ) among all functions belonging to the treatment rule class ( D ). Building on this, we adopt the following working model:

where µ ( X ) is treatment-free effect, δ k ( X ) is interaction effect between covariates and k th treatment, and ϵ is a random noise with mean zero and variance σ 2 ( A, X ) . For model identifiability, we assume ∑ K k =1 δ k ( x ) = 0 for all x . Under this model, δ k ( X ) determines the optimal ITR, while µ ( x ) has no impact on the ITR. One approach to identifying the optimal ITR is to use K separate decision functions for each treatment with the sum-to-zero constraint, but this approach can be computationally inefficient [55].

Instead, one can opt for simplex coding as an alternative, which inherently satisfies the sum-to-zero constraint. In AD-Learning [39], each of the K treatments is represented as a vertex simplex on R K -1 , denoted as u k , k = 1 , . . . , K . In particular, u k is defined as

where 1 K -1 is a K -1 dimensional vector with entries 1, and e k -1 is a K -1 dimensional vector with entries 0 except its k -th entry being 1. This vertex simplex has K vertices with equal angles between them and an origin at the center of the simplex. All u k have the same Euclidean norm 1 . Then the optimal ITR is reformulated as follows:

where u is a random treatment vector in R K , corresponding to the random treatment A . Then estimating the optimal ITR can be converted into estimating the decision function f ( x ) = ( f 1 ( x ) , . . . , f K -1 ( x )) T , assigning a K -1 dimensional vector to each covariate x . Here, while the optimal decision function f opt ( · ) can take a generic form, we assume the linear decision function within the linear treatment rule class D , induced by the linear decision function class F = { f ( X ) = B T X : B ∈ R p × ( K -1) } .

Reduction to Weighted Convex Optimization. Suppose we have a sample ( x i , a i , y i ) n i =1 of size n from the joint population distribution for ( X , A, Y ) . Estimating the optimal decision function f opt from the observed finite sample can be reduced to solving a weighted convex optimization problem:

There are three key components in the problem (3): the per-sample objective function ℓ , the persample weights w , and the regularization R ( B ) . First, ℓ denotes a model-dependent per-sample convex objective function given by

An important characteristic of causal inference problems is that only one of the possible treatments can be observed for each subject. Thus, it is critical to control confounding variables to isolate the causal effect of the treatment from other factors that might influence the outcome. This effectively is done by choosing the appropriate weight w i = w ( a i , x i ) for the i th subject. A classic approach is to use IPW: w ( a, x ) = 1 /π ( a, x ) [39].

2.2 Proposed ITR Estimation Framework

Reducing Finite Sample Bias. The classic IPW approach requires specifying a propensity score model for confounding control. However, it is well-known that propensity score methods are highly sensitive to model misspecification, which yields biased estimates of causal effects [25, 30]. To address this issue, we investigate incorporating alternative, model-free weighting schemes such as energy balancing weights (EBWs) [23, 24], which are a type of DCBWs. Other options for DCBWs include maximum mean discrepancy (MMD) balancing weights [19, 27, 8] and Wasserstein distance-based balancing weights [52, 1]. Specifically, under the working model assumption (1), the optimal decision function f opt in (2) can be decomposed as

Since E [ u π ( A, x ) ∣ ∣ x ] = 0 and E [ ϵ | x ] = 0 , it follows that f opt ( x ) = ∑ K k =1 δ k ( x ) u k , depending only on the interaction term. However, suppose the true propensity score is not correctly estimated. In that case, the estimated decision function may be biased and include additional factors including the treatment-free effect, leading to sub-optimal treatment decisions. There are two possible sources for this estimation error. First, an incorrectly specified propensity score model might lead to an estimated propensity score ˆ P n ( A = a | x ) that deviates from the true propensity score π ( A = a | x ) . Second, even with correct specification, insufficient sample size may lead to inaccurate finite-sample approximation ˆ P n ( A = a | x ) of the true propensity score π ( A = a | x ) [31]. It results in systematic bias for the estimated coefficient of the treatment-free term:

̸

Instead of using IPW, we use DCBWs such as EBWs. Such weights are data-driven and do not rely on specific model assumptions or large sample approximation. Furthermore, EBWs are known to promote independence between treatment and confounders [23], so ˆ E n [ u w ( A, x ) | x ] ≈ ˆ E n [ u ] ˆ E n [ w ( A, x ) | x ] = 0 since ˆ E n [ u ] = 0 . Thus, EBWs may reduce the unfavorable impacts of bias on the finite-sample decision function.

Improving Optimization Guarantees. To our best knowledge, the impact of the choice of the regularization R ( B ) on the model parameter B on the detailed optimization landscape, especially in terms of the convergence rate and statistical estimation error bounds has been under-investigated in the literature. Contrary to the literature, we propose to use the combination of soft L 2 -regularization and a hard L 1 -ball constraint: .

The above formulation has a smooth objective function that is unaffected by the hard L 1 -ball constraint. Also, L 2 -regularization ensures that the objective is at least λ 2 -strongly convex. Thus, the

PGD algorithm (with suitable stepsize) enjoys a strong exponential convergence guarantee toward the global optimum. In contrast, standard approaches with L 1 -penalization loss need sub-gradient methods since it is non-smooth. It is known that the convergence rate of the sub-gradient methods is much slower than gradient methods [6, 18].

We use the classical PGD to solve the convex-constrained optimization problem in (6). In each iteration, it involves conducting gradient descent with a suitable stepsize α t , followed by a projection Π B onto the L 1 -ball B := { B ∈ R p × ( K -1) : ∥ B ∥ ≤ λ 1 } :

For the projection onto the L 1 -projection, we use the algorithm in [14]. Detailed implementation of this PGD algorithm is discussed in Algorithm 1 in the appendix.

Additional Improvements by Variance and Dimension Reduction. To ensure the stability of estimates for both balancing weights and outcome regression models, a variable screening step is recommended. We propose adapting distance covariance test [48] for its computational efficiency and its ability to detect complex dependencies between treatment and outcome. By focusing on screened variables that significantly impact the outcome, we can improve the estimation of outcome regression models. Additionally, estimating balancing weights based on effect modifiers rather than all variables provides more reliable weighting estimates. Hence, the screening approach serves to reduce bias and enhance the effectiveness of ITR-Learning. See Algorithm 2 in the appendix.

Lastly, we propose using both augmented outcomes and inverse variance weighting for variance reduction in learning ITRs. Under the working model assumption in (1), the interaction effect determines optimal treatment assignments, while the treatment-free effect acts as additional noise. Outcome augmentation helps reduce variance by mitigating residual variability due to the treatment-free effect. Nevertheless, misspecification of either the outcome or treatment-free effect models can induce heteroscedastic errors due to residual treatment-free effects [35]. In such scenarios, using inverse variance weighting, introduced in SABD-Learning, enhances robustness against misspecification [44]. Consequently, the combination of outcome augmentation and inverse variance weighting offers a robust and efficient approach to ITR-Learning, as they are complementary methods. Details are discussed in the appendix C.3 and C.4.

3 Statement of results

Our theoretical analysis shows that a 'weighted design matrix', specifically a p ( K -1) × p ( K -1) positive semi-definite matrix, plays a central role in the local landscape analysis of the ITR estimation problem:

where ⊗ denotes the Kronecker product. Without loss of generality, we assume ∑ i w ( a i , x i ) = n . Note that the weights w i = w ( a i , x i ) , i = 1 , . . . , n affect the eigenvalues of the weighted design matrix Ψ . Our analysis reveals that, for the improved convergence rate of the PGD algorithm (Theorem 3.3) as well as a smaller statistical estimation error (Theorem 3.5), we need to choose the weights w i such that the minimum eigenvalue of Ψ is as large as possible. Roughly speaking, this is achieved when w 1 , . . . , w n are 'covariate balancing weights' where the weighted sample covariance matrices conditional for each treatment are approximately the same. In the special case when the covariates are discrete and one-hot encoded, the optimal weights are the ones that exactly balance the covariate distributions given treatments. Thus this provides an explicit spectral characterization of the DCBWs in causal inference literature. See the appendix D for more discussions.

3.1 Computational Guarantees

In order to control the strong convexity and the smoothness parameter of the ITR objective (6), we require that the eigenvalues of Ψ are uniformly bounded.

Assumption 3.1 (Eigenvalue bounds on the weighted design matrix) . There are constants 0 ≤ λ -≤ λ + such that the eigenvalues of Ψ in (8) is between λ -and λ + ,

For each subject i = 1 , . . . , n , let z i and p i denote its 'activation' and 'predictive probability', where

We need the activation z i to be uniformly bounded for our theoretical analysis of the binary outcome. Since the simplex vectors u a i have unit length, it is enough to require the following assumption.

Assumption 3.2 (Bounded activation) . There exists a constant M > 0 such that for all model parameter B ∈ R p × ( K -1) in the constraint set B and observed sample i = 1 , . . . , n , ∥ B T x i ∥ ≤ M . Consequently, there exist constants 0 < α -≤ α + ≤ 1 / 4 such that α -≤ p i (1 -p i ) ≤ α + for all i .

With our hard L 1 -norm constraint on the model parameter B , Assumption 3.2 can be easily verified whenever the covariate vectors x i are uniformly bounded, which is a standard in the literature.

We establish the convergence rate of Algorithm 1 for ITR-Learning with stepsizes that are fixed but sufficiently small and diminishing rate. An informal statement is given in Theorem 3.3 below, and a full statement is provided in Theorem E.1 in the appendix. Define

We say a function f : R p → R is µ -strongly convex if f ( x ) -µ 2 ∥ x ∥ 2 is convex and L -smooth if f is differentiable and ∇ f is L -Lipschitz continuous.

Theorem 3.3 (Convergence rate of PGD for ITR-Learning) . Let ( B t ) t ≥ 0 denote the sequence of parameters obtained by the PGD algorithm (7) for ITR-Learning problem (6) with arbitrary initialization B 0 . Suppose Assumptions 3.1 and 3.2 hold. Let B ⋆ denote the unique global optimum of (6) . Then the following hold:

(i) (Optimization landscape) The ITR objective L ( B ) in (6) , is µ -strongly convex and L -smooth, where µ, L are as in (10) .

(ii) (Linear convergence with fixed stepsize) Assume constant stepsize α t ≡ α < 2 /L . Denote the contraction constant ρ ( α ) := max {| 1 -αL | , | 1 -αµ |} ∈ (0 , 1) . Then ∥ B t -B ⋆ ∥ 2 F ≤ ρ ( α ) t ∥ B 0 -B ⋆ ∥ 2 F for all t ≥ 1 .

(iii) (Sublinear convergence with diminishing stepsize) Assume diminishing stepsize α t = β γ + t , where β = β ( µ, L ) > 0 is a sufficiently large constant and γ > 0 is an arbitrary constant. Then ∥ B t -B ⋆ ∥ 2 ≤ ν γ + t for all t ≥ 1 for some constant ν > 0 .

Our algorithm 1 with soft L 2 -penalization and hard constraint with L 1 -ball is guaranteed to converge to the true solution B ⋆ exponentially fast provided stepsizes α t are fixed and sufficiently small.

3.2 Statistical Guarantees

Next, we introduce generative models for the ITR estimation problem, for which we will establish statistical estimation guarantees in Theorem 3.5. Fix a joint distribution π for the pair ( X , A ) of covariates X ∈ R p and treatment A ∈ { 1 , . . . , K } . Fix a true parameter B ⋆ ∈ R p × ( K -1) . Then we assume n i.i.d. samples ( x i , a i , y i ) are drawn as ( x i , a i ) ∼ π and

where N ( µ, σ 2 ) denotes a normal distribution with mean µ and variance σ 2 , whereas Ber ( p ) denotes a Bernoulli distribution with mean p . We assume that π does not depend on B ⋆ and the noise variance σ 2 for the continuous case is known. We then seek to estimate the true parameter B ⋆ from the observed samples. It is easy to check that the negative log-likelihood function (up to additive constants) coincides with the per-sample loss function ℓ in (4). Therefore, the weighted convex optimization problem (3) corresponds to weighted MLE under the generative models in (11) (see Section H in the appendix for details).

We assume that the weighted gradient of the per-sample loss ℓ in (4) has uniformly bounded third moments for our statistical analysis of the generative ITR model above.

Assumption 3.4 (Bounded moments at B ⋆ and weights) . Suppose ( x i , a i , y i ) follows the generative model above. Denote U i := w i ∇ B ℓ ( y i , x i , a i ; B ⋆ ) ∈ R p × ( K -1) for w i := w ( a i , x i ) and U i := U i -E [ U i ] . Suppose there are constants D 1 , d 1 ∈ (0 , ∞ ) such that E [ ∥ U i ∥ 3 F ] < D 1 and min k,l Var ( U i ( k, l )) > d 1 . Also, ∑ n i =1 w 3 i / (∑ n i =1 w 2 i ) 3 / 2 = O ( n -1 / 2 ) .

Now we state our main result regarding the statistical estimation guarantee for the true generative model parameter B ⋆ . With high probability, Algorithm 1 can recover B ⋆ up to a statistical error O ( n -1 / 2 ) .

Theorem 3.5 (Statistical estimation guarantee) . Let ( x i , a i , y i ) n i =1 be i.i.d. observations from the generative models in (11) with true parameter B ⋆ ∈ R p × ( K -1) such that ∥ B ⋆ ∥ 1 ≤ λ 1 for some λ 1 ≥ 0 . Suppose Assumptions 3.1, 3.2, and 3.4 hold. Let ̂ B T denote the weighted MLE obtained after T iterations of PGD algorithm (7) for (3) . Fix a constant ε > 0 .

Then there exists a constant C = C ( ε ) > 0 such that with probability at least 1 -ε and for µ as in (10) but with λ ± = λ ± ( E [ Ψ ]) ,

provided n, T are large enough. More specifically, the above holds when

(i) (Sample complexity) n ≥ 1 large s.t. ( ∑ n i =1 w 2 i ) 3 ( ∑ n i =1 w 3 i ) 2 ≥ C 1 ε -2 for some explicit constant C 1 > 0 ;

(ii) (Computational complexity) T = O (log n ) and O ( n -1 ) when constant or diminishing stepsize as in Theorem 3.3 is used, respectively.

Furthermore, we get √ n -consistent estimation whenever λ 2 = O ( n -1 / 2 ∥ B ⋆ ∥ -1 F ) .

Notice that both terms in the error bound (12) are proportional to 1 /µ , which depends both on the minimum eigenvalue λ -of E [ Ψ ] as well as the L 2 -regularization parameter λ 2 . Therefore, choosing the weighting function ω ( A, X ) to maximize λ -helps minimize the overall statistical estimation error. (Detailed analysis and discussion on this point can be found in the appendix H.) Another easy way to increase µ is to use large L 2 -regularization, but as in the second term in (12), using L 2 -regularization gives estimation bias of O ( λ 2 ∥ B ⋆ ∥ F /µ ) . Hence, we need to choose λ 2 is small enough so that this bias term is also of O ( n -1 / 2 ) . Then we obtain an overall √ n -consistent estimator.

The key insight in our proof of Theorem 3.5 is that we can decompose the total estimation error ∥ B ⋆ -̂ B T ∥ F into computational and statistical parts:

where ̂ B denotes the exact MLE. The computational error vanishes as the number T of PGD iterations tends to infinity according to Theorem 3.3. For the statistical error, we show

Thus, the 'skewness' of the weights w i measured by the ratio √ n ∑ n i =1 w 3 i / (∑ n i =1 w 2 i ) 3 / 2 (which arises from Berry-Esseen bound) acts as a scalar multiple to the standard statistical error of O ( n -1 / 2 ) per the central limit theorem. Hence, more balanced weights lead to fewer statistical estimation errors. If this skewness ratio is uniformly bounded, we can ensure a small statistical error with probability at least 1 -ε , provided we have at least O ( ε -2 ) samples. Then, we simply need to ensure the PGD iteration T is large enough so that the computational error is at most C/ √ n to achieve the high-probability total estimation error bound in Theorem 3.5. Another insight from the theorem is that maximizing λ min ( E [ Ψ ]) promotes heterogeneous balancing weights w i , and we get increased sample complexity through the bound ( ∑ n i =1 w 2 i ) 3 ( ∑ n i =1 w 3 i ) 2 ≥ C 1 ϵ -2 (e.g., if w i ≡ 1 , this is n ⪆ ϵ -2 , but if w 1 = n, w 2 = · · · = w n = 0 , this may not be satisfied for any n ).

4 Simulation Studies

We investigate four simulation settings, each designed to evaluate various factors influencing the performance of different ITR estimators. For each setting, we mimic a randomized trial (no-confounding) and an observational study (with confounding) as per [44]. We consider 4 treatments with sample sizes n = 200 , 600 , 1000 , and covariate dimensions p = 20 , 40 , 60 . The outcome Y follows the model (1) with ϵ | A, X ∼ N ( 0 , σ 2 ( A, X ) ) . We evaluate method performance using accuracy (i.e., correct identification of the optimal treatment for each observation) and empirical value on a test dataset of 10 , 000 observations. Each simulation setting is replicated independently 100 times.

In this section, we provide partial experiment results according to the proposed estimation framework in Section 2.2. Specifically, we focus on the results from two scenarios: one mimicking randomized trials with linear ITR as the optimal rule and the other mimicking observational studies with nonlinear ITR as the optimal rule for continuous outcome. Since our proposed method assumes linear treatment rule class, the latter involves a situation where our model's prespecified treatment rule class is misspecified. Specifically, we use following treatment-free effect function µ ( X ) and interaction effect function δ ( X ) for each scenario:

Randomized Trial: Linear ITR as the true optimal

Observational Study: Nonlinear ITR as the true optimal

in the working model (1). Further details about simulation settings, other scenarios, and corresponding results are available in the appendix I.

Randomized Trial: Linear ITR as the true optimal

p=60

p=60

p=60

100%

100%

100%

80%

80%

80%

60%

%09

60%

200

600

1000

200

600

1000

200

600

1000

Sample Size

Sample Size

Sample Size

Observational Study:NonlinearITR as the true optimal

p=60

p=60

p=60

80%

80%

80%

60%

%09

60%

40%

40%

40%

20%

20%

20%

200

600

1000

200

600

1000

200

600

1000

SampleSize

Sample Size

SampleSize

AD

AD+Constrained

AD

causalDMLProposed+Penalized

Linear

AD+energy

AD+Penalized

SABD

■

policytree

Proposed+Constrained

Using Distributional Covariate Balancing Weights. In the left subplots of Figure 1, the incorporation of DCBWs (e.g., EBWs) shows improvements in classification rates for ITR-Learning. While EBWs indeed provide more effective balancing weights in observational studies, as depicted in the lower left subplot in Figure 1, it also results in improvements even in randomized trial settings where the true propensity score is known. This is achieved by reducing finite sample imbalances, as shown in the upper left subplot in Figure 1. Moreover, when the underlying true treatment rule class is nonlinear, the performance of AD-Learning suffers due to a misalignment between the assumed linear treatment rule class and the actual class. In this scenario, using EBWs yields accuracy improvements compared to the conventional IPW approach. Further details can be found in the appendix I.

Using PGD for Constrained Optimization. The middle subplots of Figure 1 show accuracy comparisons between standard penalized approach and constrained optimization using PGD. Both optimization algorithms are based on the same proposed statistical approach, integrating EBWs with SABD-Learning, outcome augmentation, and additional variable screening. For solving the L 1 -penalized regression problems, the R package glmnet is a standard choice, which is based on cyclic coordinate descent with simulated-annealing-style hyperparameter tuning. However, glmnet does not handle hard L 1 -ball constraints required for our method. Therefore, to ensure a fair comparison, we directly implemented standard (projected) (sub)gradient descent to solve optimization problems for both penalized and constrained problems. In our experiment reported in Figure 1, we observe significant improvements using our PGD method, especially in the second scenario. One possible explanation for this improvement is that large penalization may yield a perturbed solution to the ITR problem. If the minimizer of the ITR problem lies on the boundary of the L 1 -ball, the solutions of constrained optimization and penalized optimization coincide. However, if the minimizer is the interior point of L 1 -ball, constrained optimization can find the exact solution, while the penalization may induce a perturbed solution. Furthermore, the trajectory of PGD for the constraint approach appears to be significantly more stable than that of the subgradient descent in the penalization approach (see Figure 2).

scale

Constrained

Penalized

Constrained

Penalized

10

80%

601)

入1

60%

0.1

W

Loss

7

1

5

Training

Training

40%

10

5

0

10

20

30

0

10

20

30

0

10

20

300

10

20

30

Performance Evaluations between the Proposed Method and Benchmark Methods. In the right subplots of Figure 1, we evaluate accuracies by comparing benchmark methods to our proposed approach. The benchmarks include AD-Learning, SABD-Learning, a treatment-covariate interaction model with L 1 -regularization 'linear', and two popular tree-based methods:'policytree' [47, 59] and double-machine learning-based 'causalDML' [10, 29]. Our proposed method integrates EBWs with SABD-Learning, outcome augmentation, and additional variable screening for variance and dimension reduction, implemented via PGD. However, causalDML requires a sufficiently large sample size, thereby making its evaluation unfeasible under a sample size of 200. Our numerical analysis provides evidence that our proposed approach outperforms existing methods in both scenarios. 'linear' performs well in the upper subplot, which simulates a randomized trial with a true optimal linear rule, as covariates are balanced and the model accurately captures the decision function. However, in the lower subplot, simulating an observational study with a nonlinear rule, the model performs poorly because it lacks balancing weights to address confounding and cannot fully capture nonlinear relationships. Further detailed analysis, including the effectiveness of additional techniques such as variable screening for performance improvement, is provided in the appendix I.

Additionally, our proposed method processes each dataset (p = 60, n = 1000) in an average of 4.40 seconds, compared to 2.01 seconds for the penalization method using glmnet in R. Other methods exhibit scalability challenges, with policytree averaging 12.02 seconds and causalDML taking 46.75 seconds. This demonstrates that computational efficiency is a notable advantage of our approach.

5 Applications

We apply the proposed methods to two datasets from AIDS Clinical Trials Group (ACTG) 175 [21] and email marketing [22]. ACTG is a randomized trial for 2,139 patients with HIV infection who were randomly assigned to four different treatments. Twelve covariates including age and gender were used for the analysis. The outcome is the change in CD4 cell count from baseline to 20 weeks, where larger values are preferable. Email marketing dataset has 64,000 customers who were randomly chosen to receive an e-mail campaign among three different marketing methods. Eight covariates with historical customer attributes and their pairwise interactions were used. The outcome is whether customers visited the website in the following two weeks. The goal of the analysis using these two datasets is to find the best treatment/marketing methods based on individuals' attributes. Similar to [44], we randomly split the data into a training set of { 200 , 400 , 800 , 1000 , 1200 } observations for the ACTG dataset and { 1000 , 3000 , 5000 } observations for the email dataset. The remaining observations were used for test data with 10 iterations. We evaluate method performance using the empirical value function on a test dataset. The benchmarking methods include AD-Learning, SABD-Learning, and two tree-based methods (policytree [47] and causalDML [29]). However, causalDML requires a sufficiently large sample size, thereby making its evaluation unfeasible under a sample size of 200. For the binary outcome in Email dataset, SABD-Learning is not used, since the homoscedastic assumption is unnecessary for logistic regression. The results are shown in Table 1.

6 Discussion and Limitations

Our method introduces a unified, robust framework for estimating ITR. By formulating the problem as a weighted, constrained optimization problem, we incorporate DCBWs to control confounding and propose the PGD algorithm for sparse ITRs with computational and statistical guarantees. Additionally, we propose variable screening and efficient augmentation, which together show synergistic effects. We demonstrate our method with continuous and binary outcomes, and it can be readily extended to other outcomes, such as censored outcomes. Although we focus on learning linear ITR to facilitate the demonstration of our idea, it can be extended to nonlinear ITR using the basis function [34]. For example, by taking the covariate powers up to M, the linear function class becomes the class of degree-M polynomial functions. The linear function class can also serve as a useful approximation to nonlinear decision functions.

We acknowledge that although most, if not all, types of ITR-Learning frameworks can benefit from using DCBWs, variable screening, and augmentation, our proposed framework for theoretical analysis relies on the specific assumption that the model can be formulated as a constrained, weighted, and convex optimization problem. Furthermore, we restricted the decision function to the linear function class. Generalizing our theoretical analysis to include nonlinear function classes with nonconvex problems, such as boosting or deep neural networks, is an interesting direction for future investigation. Another limitation of our method is its reliance on standard assumptions to identify optimal ITRs using observational data, including the assumption of no unmeasured confounding. An interesting future direction would be to extend our framework to cases where this assumption does not hold, utilizing instrumental variables [12, 38] and proximal causal learning [40, 45].

Our work significantly impacts society by improving patient health outcomes through personalized care and advancing medical research in disease mechanisms and drug development. Beyond healthcare, these methods can be applied in marketing strategies based on customer characteristics.

References

Serge Assaad, Shuxi Zeng, Chenyang Tao, Shounak Datta, Nikhil Mehta, Ricardo Henao, Fan Li, and Lawrence Carin. Counterfactual representation learning with balancing weights. In International Conference on Artificial Intelligence and Statistics , pages 1972-1980. PMLR, 2021.

Susan Athey and Stefan Wager. Policy learning with observational data. Econometrica , 89(1):133-161, 2021.

Laura B Balzer, Maya L Petersen, Mark J van der Laan, and Search Collaboration. Targeted estimation and inference for the sample average treatment effect in trials with and without pair-matching. Statistics in medicine , 35(21):3717-3732, 2016.

Léon Bottou. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT'2010 , pages 177-186. Springer, 2010.

Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization . Cambridge university press, 2004.

Stephen Boyd, Lin Xiao, and Almir Mutapcic. Subgradient methods. lecture notes of EE392o, Stanford University, Autumn Quarter , 2004(01), 2003.

Guanhua Chen, Donglin Zeng, and Michael R Kosorok. Personalized dose finding using outcome weighted learning. J. Amer. Statist. Assoc. , 111(516):1509-1521, 2016.

Rui Chen, Jared D Huling, Guanhua Chen, and Menggang Yu. Robust sample weighting to facilitate individualized treatment rule learning for a target population. Biometrika , 111(1):309329, 2024.

Shuai Chen, Lu Tian, Tianxi Cai, and Menggang Yu. A general statistical framework for subgroup identification and comparative treatment scoring. Biometrics , 73(4):1199-1209, 2017.

Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, and Whitney Newey. Double/debiased/neyman machine learning of treatment effects. Am. Econ. Rev. , 107(5):261-65, 2017.

Francis S Collins and Harold Varmus. A new initiative on precision medicine. New England journal of medicine , 372(9):793-795, 2015.

Yifan Cui and Eric Tchetgen Tchetgen. A semiparametric instrumental variable approach to optimal treatment regimes under endogeneity. Journal of the American Statistical Association , 116(533):162-173, 2021.

Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statistics , pages 1810-1818. PMLR, 2021.

John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra. Efficient projections onto the l 1-ball for learning in high dimensions. In Proceedings of the 25th international conference on Machine learning , pages 272-279, 2008.

Jianqing Fan and Runze Li. Variable selection via nonconcave penalized likelihood and its oracle properties. Journal of the American statistical Association , 96(456):1348-1360, 2001.

Jianqing Fan and Jinchi Lv. Sure independence screening for ultrahigh dimensional feature space. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 70(5):849-911, 2008.

William Feller. An introduction to probability theory and its applications, Volume 2 , volume 81. John Wiley & Sons, 1991.

Jean-Louis Goffin. On convergence rates of subgradient optimization methods. Mathematical programming , 13:329-347, 1977.

Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research , 13(1):723-773, 2012.

Jens Hainmueller. Entropy balancing for causal effects: A multivariate reweighting method to produce balanced samples in observational studies. Political analysis , 20(1):25-46, 2012.

Scott M Hammer, David A Katzenstein, Michael D Hughes, Holly Gundacker, Robert T Schooley, Richard H Haubrich, W Keith Henry, Michael M Lederman, John P Phair, Manette Niu, et al. A trial comparing nucleoside monotherapy with combination therapy in hiv-infected adults with cd4 cell counts from 200 to 500 per cubic millimeter. New England Journal of Medicine , 335(15):1081-1090, 1996.

Kevin Hillstrom. Minethatdata e-mail analytics and data mining challenge. https://www. uplift-modeling.com/en/v0.3.1/api/datasets/fetch_hillstrom.html , 2008. Accessed: 2024-04-28.

Jared D Huling, Noah Greifer, and Guanhua Chen. Independence weights for causal inference with continuous treatments. Journal of the American Statistical Association , 119(546):16571670, 2024.

Jared D Huling and Simon Mak. Energy balancing of covariate distributions. Journal of Causal Inference , 12(1):20220029, 2024.

Kosuke Imai and Marc Ratkovic. Covariate balancing propensity score. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 76(1):243-263, 2014.

Nathan Kallus. Balanced policy evaluation and learning. Advances in neural information processing systems , 31, 2018.

Nathan Kallus. Generalized optimal matching methods for causal inference. The Journal of Machine Learning Research , 21(1):2300-2353, 2020.

Nathan Kallus. More efficient policy learning via optimal retargeting. Journal of the American Statistical Association , 116(534):646-658, 2021.

Michael C Knaus. Double machine learning-based programme evaluation under unconfoundedness. The Econometrics Journal , 25(3):602-627, 2022.

Brian K Lee, Justin Lessler, and Elizabeth A Stuart. Weight trimming and propensity score weighting. PloS one , 6(3):e18174, 2011.

Fan Li, Kari Lock Morgan, and Alan M Zaslavsky. Balancing covariates via propensity score weighting. Journal of the American Statistical Association , 113(521):390-400, 2018.

Runze Li, Wei Zhong, and Liping Zhu. Feature screening via distance correlation learning. Journal of the American Statistical Association , 107(499):1129-1139, 2012.

Haixu Ma, Donglin Zeng, and Yufeng Liu. Learning individualized treatment rules with many treatments: A supervised clustering approach using adaptive fusion. Advances in Neural Information Processing Systems , 35:15956-15969, 2022.

Jacob M Maronge, Jared D Huling, and Guanhua Chen. A reluctant additive model framework for interpretable nonlinear individualized treatment rules. The Annals of Applied Statistics , 17(4):3384-3402, 2023.

Weibin Mo and Yufeng Liu. Efficient learning of optimal individualized treatment rules for heteroscedastic or misspecified treatment-free effect models. Journal of the Royal Statistical Society Series B: Statistical Methodology , 84(2):440-472, 2022.

Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika , 108(2):299-319, 2021.

Wenliang Pan, Xueqin Wang, Heping Zhang, Hongtu Zhu, and Jin Zhu. Ball covariance: A generic measure of dependence in banach space. Journal of the American Statistical Association , 115(529):307-317, 2019.

Hongming Pu and Bo Zhang. Estimating optimal treatment rules with an instrumental variable: A partial identification learning approach. Journal of the Royal Statistical Society Series B: Statistical Methodology , 83(2):318-345, 2021.

Zhengling Qi, Dacheng Liu, Haoda Fu, and Yufeng Liu. Multi-armed angle-based direct learning for estimating optimal individualized treatment rules with various outcomes. Journal of the American Statistical Association , 115(530):678-691, 2020.

Zhengling Qi, Rui Miao, and Xiaoke Zhang. Proximal learning for individualized treatment regimes under unmeasured confounding. Journal of the American Statistical Association , pages 1-14, 2023.

Min Qian and Susan A Murphy. Performance guarantees for individualized treatment rules. Annals of statistics , 39(2):1180, 2011.

Brian C Ross. Mutual information between discrete and continuous data sets. PloS one , 9(2):e87357, 2014.

Donald B Rubin. Randomization analysis of experimental data: The fisher randomization test comment. Journal of the American statistical association , 75(371):591-593, 1980.

Kushal S Shah, Haoda Fu, and Michael R Kosorok. Stabilized direct learning for efficient estimation of individualized treatment rules. Biometrics , 2022. doi: 10.1111/biom.13818.

Tao Shen and Yifan Cui. Optimal treatment regimes for proximal causal learning. Advances in Neural Information Processing Systems , 36, 2024.

Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in neural information processing systems , 32, 2019.

Erik Sverdrup, Ayush Kanodia, Zhengyuan Zhou, Susan Athey, and Stefan Wager. policytree: Policy learning via doubly robust empirical welfare maximization over trees. Journal of Open Source Software , 5(50):2232, 2020.

Gábor J Székely, Maria L Rizzo, and Nail K Bakirov. Measuring and testing dependence by correlation of distances. The Annals of Statistics , 35(6):2769-2794, 2007.

Dingke Tang, Dehan Kong, Wenliang Pan, and Linbo Wang. Ultra-high dimensional variable selection for doubly robust causal inference. Biometrics , 79(2):903-914, 2023.

Lu Tian, Ash A Alizadeh, Andrew J Gentles, and Robert Tibshirani. A simple method for estimating interactions between a treatment and a large number of covariates. Journal of the American Statistical Association , 109(508):1517-1532, 2014.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association , 113(523):1228-1242, 2018.

Baqun Zhang, Anastasios A Tsiatis, Eric B Laber, and Marie Davidian. A robust method for estimating optimal treatment regimes. Biometrics , 68(4):1010-1018, 2012.

Chong Zhang and Yufeng Liu. Multicategory angle-based large-margin classification. Biometrika , 101(3):625-640, 2014.

Qingyuan Zhao, Dylan S Small, and Ashkan Ertefaie. Selective inference for effect modification via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology , 84(2):382-413, 2022.

Yingqi Zhao, Donglin Zeng, A John Rush, and Michael R Kosorok. Estimating individualized treatment rules using outcome weighted learning. Journal of the American Statistical Association , 107(499):1106-1118, 2012.

Xin Zhou, Nicole Mayer-Hamblett, Umer Khan, and Michael R Kosorok. Residual weighted learning for estimating individualized treatment rules. Journal of the American Statistical Association , 112(517):169-187, 2017.

Zhengyuan Zhou, Susan Athey, and Stefan Wager. Offline multi-action policy learning: Generalization and optimization. Operations Research , 71(1):148-183, 2023.

Ruoqing Zhu, Ying-Qi Zhao, Guanhua Chen, Shuangge Ma, and Hongyu Zhao. Greedy outcome weighted tree learning of optimal personalized treatment rules. Biometrics , 73(2):391-400, 2017.

José R Zubizarreta. Stable weights that balance covariates for estimation with incomplete outcome data. Journal of the American Statistical Association , 110(511):910-922, 2015.

A Preliminaries

In this section, we review several notations and basic facts on linear algebra and matrix calculus.

The Kronecker product, denoted by ⊗ , is a binary operation that combines two matrices to create a new matrix. Given two matrices A = ( a ij ) ∈ R m × n and B ∈ R p × q , then A ⊗ B is an mp × nq matrix defined as

If A ′ ∈ R m × n ′ , then for the horizontally stacked matrix [ A , A ′ ] ∈ R m × ( n + n ′ ) ,

For each A = [ a 1 , . . . , a n ] ∈ R m × n , let vec( A ) := [ a T 1 , . . . , a T n ] T ∈ R mn denote its vectorization. The commutation matrix C ( a,b ) , which is a special instance of ab × ab permutation matrix. Namely, for each integers a, b ≥ 1 , there exists a unique matrix C ( a,b ) ∈ { 0 , 1 } ab × ab such that for all A ∈ R a × b , we have C ( a,b ) vec( A ) = vec( A T ) . Note that ( C ( a,b ) ) T = C ( b,a ) . Furthermore, ( C ( a,b ) ) T C ( a,b ) = I ab since ( C ( a,b ) ) T C ( a,b ) vec( A ) = C ( b,a ) vec( A T ) = vec( A ) . Hence C ( a,b ) is positive semi-definite. Note that C ( a, 1) = I a = C (1 ,a ) since if A ∈ R a × 1 or A ∈ R 1 × a , vec( A ) = vec( A T ) . Also, for every matrices A ∈ R m × n and B ∈ R r × q ,

Also recall the relations for vectorizing product of matrices: for A ∈ R a × b , B ∈ R b × c , and C ∈ R c × d ,

For A , B , C , and D are matrices of compatible sizes allowing the matrix products AC and BD ,

Next, recall that if f : R a × 1 → R b × 1 and g : R b × 1 → R c × 1 are differentiable functions, then the Jacobian J f can be written by using gradients as J f ( x ) = ( ∇ x f T ) T . By chain rule J g ◦ f ( x ) = J g ( f ( x )) J f ( x ) , we have the following chain rule for gradients

We will also frequently use the following fact: For A ∈ R a × b ,

For two matrices A , B of the same size, we write A ⪯ B if B -A is positive semi-definite. The partial ordering ⪯ is called the Loewner ordering.

B Details of model setup

B.1 Details of angle-based approach

Consider the working model given by Equation (1) in the main text:

To make the optimal ITR identifiable, we impose the constraint ∑ K k =1 δ k ( x ) = 0 for all covariates x . Under this model, δ k ( x ) determines the optimal ITR for an individual with covariate x , while µ ( x ) has no impact on the ITR. However, using K functions with the sum-to-zero constraint can be computationally inefficient [55]. Instead, one can opt for simplex coding as an alternative, which inherently satisfies the sum-to-zero constraint.

In AD-Learning, proposed by [39], each of the K treatments is represented as a vertex simplex on R K -1 , denoted as u k , k = 1 , . . . , K . In particular, u k is defined as

where 1 K -1 is a K -1 dimensional vector with entries 1, and e k -1 is a K -1 dimensional vector with entries 0 except its k -th entry being 1. This vertex simplex has K vertices with equal angles between them and an origin at the center of the simplex. All u k have the same Euclidean norm 1 . To estimate the optimal ITR, a decision function f ( x ) = ( f 1 ( x ) , . . . , f K -1 ( x )) T is constructed, assigning a K -1 dimensional vector to each covariate x . Here, f ( · ) can take a generic form, while we assume the linear decision function in the main text. The optimal ITR is reformulated by the authors as follows:

where ∠ ( · , · ) represents the angle between two vectors. Note that due to ∑ K k =1 u k = 0 , ∑ K k =1 u T k f ( x ) = 0 for any given function f and covariate value x . Then we can use

instead of ∑ K k =1 δ k ( X ) I ( A = k ) to encode the treatment and covariate interaction. Specifically, we use u A to denote ( ∑ K k =1 u k I ( A = k ) ) T and assume f ( x ) to be linear in x (with the coefficient matrix B ). Therefore, terms pertaining to ITRs can be expressed as u T a i B T x i .

B.2 Details of binary outcome

Again, under prespecified treatment rule class D , the optimal ITR is given by

For the binary response setting, we assume the following logistic model

Under this model, f opt is indeed the optimal decision function under the angle-based direct learning framework since

Furthermore, it can be checked that the function f opt is an optimal solution to maximizing the IPW-weighted log-likelihood of the logistic regression model:

under the assumption that the expectation and derivative w.r.t. f can be exchanged ([39]).

For the effective and robust ITR estimation, replacing IPW with DCBW leads to

Now consider restricting it to the linear decision function class F = { f ( X ) = B T X : B ∈ R p × ( K -1) } . Then the above (25) becomes

Then, we can estimate the optimal parameter B by solving the following empirical loss minimization

where we imposed additionally the L 1 -penalization ∥ B ∥ 1 ≤ z, on B to promote a sparse solution. By the proposed algorithm 1, we can efficiently solve the problem (26).

C Details of proposed approaches

To keep the appendix as self-contained as possible, we provide descriptions of the building blocks of the proposed methods.

C.1 Use of distributional covariate balancing weights

In this section, we provide the background about the use of DCBWs for constructing effective and robust ITRs. For additional details, please refer to the original papers.

Many existing ITR-Learning approaches such as AD-Learning and SABD-Learning use inverse probability (of treatment) weighting (IPW), which requires specifying a propensity score model for confounding control. However, propensity score methods have long been known to be highly sensitive to model misspecification, which yields biased estimates of causal effects [25, 30]. For this reason, methods that directly estimate weights based on balance-seeking objectives instead of using IPW have been proposed. These approaches include entropy balancing weights [20], stable balancing weights [61] as well as DCBWs [24, 27]. However, despite the weights serving a critical role in the estimation process, these advancements have mainly focused on improving average treatment effect (ATE) and are largely under-explored in the ITR literature. Significant improvements in the performance of ITRs are often left on the table when the weights are not treated with the same level of emphasis and focus as the other aspects of ITR learning.

To tackle these challenges, we propose the use of DCBWs, such as energy balancing weights (EBWs) [24] and maximum mean discrepancy (MMD) balancing weights [19, 27, 8], as an alternative to IPW for ITR-Learning. Unlike traditional approaches, these methods do not rely on any pre-specified functional forms for constructing weights. Instead, they focus on minimizing the distance of the weighted empirical distributions of covariates across treatment groups to reduce potential confounding. Even for data from randomized trials, they can help reduce finite sample covariate imbalance. The emphasis on finite sample balance and the model-independent nature of distributional weighting methods effectively mitigate biases due to model misspecification, resulting in more precise and reliable ITR estimation.

In the remainder of this section, we provide additional details about EBWs. The EBWs are derived from the concept of 'weighted energy distance" which modifies the energy distance measure to allow for measurement of the distance between a weighted empirical distribution and a target empirical distribution. This measure captures the distance between a weighted empirical cumulative distribution function (ECDF) of covariates for a specific treatment group a , denoted as { X i } i : A i = a , and the combined ECDF of covariates, denoted as { X i } n i =1 . Specifically, the weighted energy distance is defined as:

where n a = ∑ n i =1 I ( A i = a ) . The weighted energy distance measures the distance between the weighted distribution of covariates among those with A i = a and the empirical distribution of covariates in the full sample; it takes value 0 if and only if these two distributions are identical.

The optimal weights are obtained by minimizing the sum of all the weighted energy distances of all treatment groups:

Thus, the EBWs encourage the distributions of covariates for each treatment group to look like that of the target population, i.e. the full sample. Since (27) is a quadratic objective function along with linear constraints, it can be efficiently solved through interior point methods. Another advantage lies in the fact that it does not require parameter tuning. The estimated weights are effective in balancing

all covariates of the covariate distributions, not just limited to the first or second moments. Lastly, the data-driven characteristic of EBWs that do not rely on specific model assumptions contributes to more robust ITR estimates.

C.2 Implementation of the PGD algorithm

We provide Algorithm 1 to solve the problem (6) using a projected gradient descent (PGD). The algorithm applies to either continuous or binary outcomes. For the projection onto the L 1 -ball, we use the algorithm in [14].

Algorithm 1

To help the convenient implementation of the above PGD algorithm, we provide the gradient of the ITR objective below. For the continuous outcome, we have (see Proposition F.1)

For the binary outcome, we have (see Proposition F.2)

C.3 Additional improvements by dimension reduction: variable screening procedure

High dimensional covariates can significantly affect the performance of methods for estimating weights, outcome regression models, and the final ITR. Specifically, using screened variables is beneficial in estimating outcome regression models by focusing on fewer variables that significantly impact the outcome. Additionally, estimating balancing weights based on effect modifiers rather than all variables provides more reliable weighting estimates. Hence, the screening approach serves to reduce bias and enhance the efficiency of ITR-Learning.

Several techniques have been developed to choose appropriate variables in causal inference. One popular approach in the ITR setting is the group-lasso method introduced in [33] which filters out variables that only affect the treatment-free effect. For the ATE estimation, sure independence screening (SIS) method [16, 32] is widely used, which involves an initial screening step based on Pearson correlations between each covariate and the outcome. The covariates are then ranked based on these correlations, and only the top-ranked ones are used in the downstream analysis. While one could also use similar variable screening methods in the ITR setting, the use of Pearson correlation may limit their effectiveness in capturing nonlinear relationships, potentially leading to biased estimates of the ITR.

To address this issue, alternative approaches that employ non-linear dependence tests are explored, including the ball covariance, mutual information test, and distance covariance [37, 42, 48]. Ball covariance, previously applied to variable selection in ATE estimation, has been found to be effective [49]. Distance covariance is another way to capture the nonlinear dependence between two random

vectors/variables [48]. The test compares the observed distances between pairs of observations to the expected distances under the assumption of independence. A zero value of distance covariance indicates independence between two random vectors. If the observed distances are significantly far from the expected distances, it indicates statistical dependence between the random vectors. The p -value is calculated by comparing the observed empirical distance with the distribution of the same test statistic obtained from random permutations of the data under the assumption of independence. We propose adapting the distance covariance due to its computational efficiency and empirical performance for the settings we explored.

Now, we introduce our proposed variable screening procedure. In a setting with K treatments, we aim to find which covariates have conditional dependence on the outcome given the treatment. Doing so allows the identification of covariates that either modify the treatment effect or otherwise affect the outcome (i.e. are precision variables). To do so, for each covariate X j where j ∈ { 1 , . . . , p } , conduct the distance covariance test between covariate and outcome K times, conditioning on each treatment value. This conditioning on the treatment creates subgroups based on treatment assignments, allowing the assessment of the dependence between the covariate and the outcome within each subgroup. Then we set the p -value of the covariate X j as the minimum p-value from these K tests. This approach ensures that we retain as many relevant variables as possible during this preliminary variable screening stage. For the same reason, we avoid the use of multiple comparison adjustments such as Bonferroni correction. Instead, we opt to select variables with a significance level of 5%. We perform the same steps for all covariates in the dataset. These procedures are summarized in Algorithm 2.

Algorithm 2 Variable Screening using Distance Covariance

This conditional independence test is crucial to retain effect modifiers and precision variables while filtering out instrumental variables associated with treatment but not the outcome [49]. While acknowledging that the procedure is ad-hoc, the simulation results show that it effectively reduces the number of covariates to estimate covariate balancing weights and outcome regression models, which improves efficiency and helps avoid overfitting.

C.4 Additional improvements by variance reduction: combining inverse variance weighting with augmented outcome

Our additional recommendation includes dually using both augmented outcomes and inverse variance weighting for variance reduction. Under our working model (1), the outcome can be explained by the sum of the treatment-free effects of covariates, the interaction effect between covariates and each treatment, and random noise. As the interaction effect alone determines optimal treatment assignments, the treatment-free effect acts as extra noise. This noise can be mitigated using outcome model augmentation. One way to implement outcome augmentation is to use the augmented outcomes ˇ Y := Y -̂ µ ( X ) in place of the original outcomes for continuous outcomes when applying an ITR estimation method. Flexible techniques such as random forest and neural networks can be utilized to estimate the treatment-free effect. This is the first implementation for multi-category treatments with DCBWs, whereas previous work focused on binary treatments with IPW [44].

Augmented outcomes can refine ITR estimates when correctly specified. Yet, misspecification of either outcome or treatment-free effect models can lead to heteroscedastic errors due to residual treatment-free effects. This is where SABD-Learning's inverse variance weighting helps, providing

robustness against misspecification. SABD-Learning extends AD-Learning designed for homoscedasticity [39] to account for heteroscedasticity, i.e. when the conditional variance of the (potential) outcome given the covariates is not constant [44]. This method estimates the conditional variance of outcomes and reweights the final objective function by adjusting each sample's contribution based on the inverse of the estimated variance. Therefore, we propose to combine outcome augmentation and inverse variance weighting for robust and efficient ITR learning, as they are complementary methods.

The key idea behind outcome augmentation is to add an asymptotically unbiased quantity to the objective function (negative log-likelihood function) so that the new optimizer for the 'augmented objective function' has a provably smaller variance than the original estimator. To do this, let h be some unknown (to be determined) augmentation function depending on the covariate. The key property of asymptotic unbiasedness is

We will justify this property in Corollary C.3 toward the end of this section. Given (30), we have

where ≈ here means asymptotic equality in probability as n →∞ . Denote the 'augmented objective function' on the right-hand side as ℓ aug ( B ) .

In the remainder of this section, we introduce the optimal augmented function h opt ( X ) that minimizes the variance of the estimator, while maintaining asymptotic unbiasedness. We demonstrate that outcome augmentation remains valid when IPW is replaced with DCBWs, such as EBWs and MMD balancing weights. Our approach extends the work of [50, 9] to a more general setting, incorporating other types of covariate balancing weights that have asymptotic consistency (38) in a multi-category treatment setting.

Let S ( B ) be the derivative of the original objective function ℓ ( B ) with respect to B . Then setting the optimal h ( X ) is equivalent to minimizing the conditional variance of

where ̂ B a is the minimizer of ℓ aug ( B ) , and B ⋆ is the minimizer of E [ ℓ ( B )] . Using the property of asymptotic unbiasedness Theorem C.2, which is E [ Z ] = 0 , we need to minimize E [ ∥ Z ∥ 2 F ] with respect to h to minimize the conditional variance of Z . Using the property of matrix trace,

and setting ∇ h E [ ∥ Z ∥ 2 F ] = 0 . Thus, h opt ( X ) satisfies

Since u T A u A = 1 , we have

Now, we determine the functional form of h opt ( X ) for both continuous and binary outcomes. For simplicity, let us define

Then for continuous outcomes, using (28) without regularization term for S ( B ) and (32), we have ℓ aug ( B )

assuming ̂ R 1 ( x i ) / ̂ c 1 = 0 . Therefore, we obtain the optimal augmented function with the minimum estimation variance by properly defining R 2 ( X ) . Our simulation studies demonstrate that setting R 2 ( X ) = E [ Y | X = x ] leads to significant improvement.

Similarly, for binary outcomes, let us define

Then using (29) without regularization term for S ( B ) and (32), we have

Therefore, by selecting proper Q 1 ( X ) and Q 2 ( X ) , the minimizer of ℓ aug ( B ) has smaller variance than that of the original estimator.

Now we prove asymptotic unbiasedness of 1 n ∑ n i =1 w ( a i , x i ) u T a i B T h ( x i ) . First, we consider IPWbased conditions for multi-category treatments.

Lemma C.1. Let A = { 1 , . . . , K } be the possible treatment sets. Assume we have i.i.d. observations ( x i , a i ) from a population distribution of ( X , A ) such that P ( A = a | x ) is uniformly positive for all a, x . Define

If g is any function with ∑ j ∈A g ( j ) = 0 , then for any function h ( X ) : R p → R p , then T → 0 as n →∞ in probability.

Proof. First note that by iterated expectation,

Then by the standard weak law of large numbers for i.i.d. samples, the empirical mean T converges in probability to its population mean E [ T ] = 0 as n →∞ .

Note that in the statement of the Lemma C.1, the choice of such g ( j ) can be each treatment vertex under the AD-Learning [39]. Now we replace IPW P ( A = a i | x i ) with covariate balancing weights w ( a i , x i ) that has the following asymptotic consistency property:

For instance, maximum mean discrepancy (MMD)balancing weights, a type of distributional covariate energy balancing weight, satisfy the property (38) [8].

Theorem C.2 (Asymptotic unbiasedness of augmented function) . Assume we have i.i.d. observations ( x i , a i ) from population distribution of ( X , A ) . If g : A → R is any bounded function with ∑ j ∈A g ( j ) = 0 , then for any bounded function h ( x ) : R p → R p , define

where the weights w ( a i , x i ) satisfy (38) . Then T 0 → 0 as n →∞ in probability. That is, T 0 is an asymptotically mean zero, making it suitable to be used as the augmented function.

Proof. By the weak law of large numbers for i.i.d. samples, the empirical mean T 0 converges in probability to its population mean E [ T 0 ] = E [ w ( A, X ) h ( X ) g ( A )] . By continuous mapping theorem, ∥ T 0 ∥ 2 converges to ∥ E [ T 0 ] ∥ 2 as n →∞ in probability. Then we have

where C is any constant such that ∥ h ( X ) g ( A ) ∥ 2 ≤ C . For (a), we use Lemma C.1; for (b) and (c), we use Jensen's inequality and the law of large numbers, respectively. The last term converges to 0 due to the property (38). Lastly, since ∥ T 0 ∥ 2 converges to 0 in probability, we have T 0 converges to 0 in probability by continuous mapping theorem.

The following corollary is useful in deriving outcome augmentation under the AD-Learning framework.

Corollary C.3. Assume the AD-learning setting. For any function h ( X ) depending only on the covariate X , and for each parameter matrix B ∈ R p × ( K -1) , denoting,

we have E [ H 0 ( B )] = 0 . Furthermore, assuming the weights w ( A, X ) satisfying the consistency property (38) , in probability as n →∞ ,

Proof. Follows immediately from Theorem C.2 by noting that ∑ K a =1 u a = 0 .

D Spectral characterization of optimal covariate balancing weight

Recall the definition of the weighted design matrix:

This is the p ( K -1) × p ( K -1) design matrix encoding the sample of patients with covariate x i and observed treatment a i , for i = 1 , . . . , n . We are free to choose the weight w i = ω ( a i , x i ) for the i th patient as a function as long as they sum to one. What is the best way to choose such weights? Our main results (Theorems 3.3, 3.5) suggest that it is best to choose them in a way that the minimum eigenvalue of Ψ is as large as possible.

To illustrate this point, suppose for simplicity the covariates are categorical and assume values in a finite set X . We decompose the summation in (41) according to the treatment and covariate as follows:

Note that n ( a, x ) is the number of patients with covariate x and treatment a , while Σ ( a, x ) is the p ( K -1) × p ( K -1) covariance matrix representing the treatment-covariate pair ( a, x ) .

Under IPW w ( a, x ) = 1 π ( a | x ) , since n ( a, x ) ∼ Binom ( n, π ( a | x )) (under suitable assumption), using a law of large numbers and central limit theorem, we have

Consequently, the covariance matrix Σ ( a, x ) of every treatment-covariate pair ( a, x ) contributes nearly equally to the weighted design matrix Ψ .

If distinct covariates are encoded as one-hot encoding vectors, then IPW maximizes the minimum eigenvalue of Ψ . To see this, for simplicity, assume we have binary treatments ( K = 2 and u 1 = 1 , u 2 = -1 ) and two discrete covariates (e.g., 'Male' and 'Female') encoded as one-hot vectors (1 , 0) T and (0 , 1) T . Then we can write Ψ as

Notice that the minimum eigenvalue of the last 2 × 2 matrix is the minimum of its two diagonal entries. With IPW, each weighted count w ( a, x ) n ( a, x ) becomes a constant, and Ψ above becomes the identity matrix, achieving the largest possible minimum eigenvalue. Also, in this case, Ψ has condition number 1 so the ITR optimization landscape becomes well-conditioned.

Using the fact that the minimum eigenvalue is a positive homogeneous and concave function on the space of symmetric matrices, we have

Here Σ w ( a ) is the aggregated weighted covariance matrix for sub-population with treatment a . In the above derivation, we used that Ψ = ∑ a Σ w ( a ) by (42) and definition and the following computation using Jensen's inequality:

If we choose the weights so that these matrices are identical, then each treatment group contributes equally to the lower bound above for λ min ( Ψ ) . This condition is satisfied, for example, when w is distributional covariate balancing weights, which ensure that the covariate distributions across all treatment groups are identical. There is a remaining degree of freedom for specifying the common covariate distribution, denoted by µ , for the treatment groups. To maximize the lower bound for λ min ( Ψ ) , µ should be chosen to maximize λ min ( Σ w ( · )) . The specific solution to this optimization depends on the structure of Ψ . In three-way balancing methods [8], µ is typically chosen to be close to the target population distribution. Another common approach, as discussed earlier, is IPW, which sets µ as the uniform distribution over the covariate space X .

E Convergence analysis of PGD for ITR estimation

In this section, we analyze the convergence properties of Algorithm 1 for ITR-Learning problem (6). Below we state the full version of Theorem 3.3 in the main text.

Theorem E.1 (Convergence rate of PGD for ITR-Learning) . Let ( B t ) t ≥ 0 denote the sequence of parameters obtained by the PGD algorithm (7) for ITR-Learning problem (6) with arbitrary initialization B 0 . Suppose Assumptions 3.1 and 3.2 hold. Let B ⋆ denote the unique global optimum of (6) . Then the following hold:

(i) (Optimization landscape) The ITR objective L ( B ) in (6) , is µ -strongly convex and L -smooth, where µ, L are as in (10) .

(ii) (Linear convergence with fixed stepsize) Assume constant stepsize α t ≡ α < 2 /L . Denote the contraction constant ρ ( α ) := max {| 1 -αL | , | 1 -αµ |} ∈ (0 , 1) . Then for all t ≥ 1 ,

In particular, if α = 2 µ + L , then ρ ( α ) = L -µ L + µ ∈ [0 , 1)

(iii) (Sublinear convergence with diminishing stepsize) Assume diminishing stepsize α t = β γ + t , where ρ := 2 µL µ + L and β > 1 /ρ , γ > 0 are constants. Then for all t ≥ 1 ,

where ν := max { 2 β 2 ∥∇ f ( B ⋆ ) ∥ 2 F βρ -1 , 4( γ +1) λ 2 1 } . In particular, the assertion holds for β = 2 /ρ = µ + L µL and ν = max { 2( µ + L ) 2 ∥∇ f ( B ⋆ ) ∥ 2 F µ 2 L 2 , 4( γ +1) λ 2 1 } .

The proof of Theorem E.1 relies on two ingredients: a general result on the convergence of PGD (Lemma E.2) and a local landscape analysis of the ITR objective (Lemma E.3).

We first establish a general result on the convergence rate of PGD to find the minimizer of a strongly convex and smooth objective f within a convex constraint B . Our argument is mostly based on the standard convergence analysis of PGD in the optimization literature [5, 4]. A new element we add is that for analyzing constrained optimization problems, the gradient of the objective f at the minimizer B ⋆ does not need to be zero, which only needs to be in the negative normal cone at the constraint set B . In our analysis, we treat the possibly nonzero gradient norm ∥∇ f ( B ⋆ ) ∥ F as 'noise' (variance) of a stochastic gradient oracle in standard stochastic gradient descent analysis [4].

Lemma E.2 (Convergence rate of PGD for strongly convex and smooth objectives) . Let f : R p × ( K -1) → R be a L -smooth and µ -strongly convex function for some µ, L > 0 . Suppose f ⋆ := inf B ∈B f ( B ) > -∞ and denote B ⋆ := arg min B ∈B f ( B ) , where B is a convex constraint set of R p × ( K -1) . Consider the following PGD iterates:

Then the following hold:

(i) (Linear convergence with fixed stepsize) Suppose f is twice differentiable and α t ≡ α < 2 /L . Let contraction constant ρ ( α ) := max {| 1 -αL | , | 1 -αµ |} ∈ (0 , 1) . Then for all t ≥ 1 ,

The bound ρ ( α ) on the linear convergence rate is minimized at α = 2 µ + L with minimum value L -µ L + µ .

(ii) (Sublinear convergence with diminishing stepsize) Denote ρ := 2 µL µ + L . Suppose α t = β γ + t , where β > 1 /ρ , γ > 0 are constants. Then for all t ≥ 1 ,

where ν := max { 2 β 2 ∥∇ f ( B ⋆ ) ∥ 2 F βρ -1 , ( γ +1) ∥ B 1 -B ⋆ ∥ 2 F } . In particular, the assertion holds for β = 2 /ρ = µ + L µL and ν = max { 2( µ + L ) 2 ∥∇ f ( B ⋆ ) ∥ 2 F µ 2 L 2 , ( γ +1) ∥ B 1 -B ⋆ ∥ 2 F } .

¯

using that ⟨∇ f ( B ⋆ ) , B t -B ⋆ ⟩ ≥ 0 since B ⋆ is a stationary point of f over B . Using the co-coercivity of µ -strongly convex and L -smooth functions and, we have

By Young's inequality,

Thus the above yields, given that α t ≤ 1 / ( µ + L ) ,

Since the projection onto the convex constraint set B , Π B , is non-expansive,

Now, we show (i) (47). Suppose α t ≡ α > 0 . Note that B ∗ is a unique minimizer of f over B because of strong convexity. Since -∇ f ( B ⋆ ) lies in the normal cone of B at B ⋆ , we have B ⋆ = Π B ( B ⋆ -α ∇ f ( B ⋆ )) . Using non-expansiveness of convex projection Π B ,

For the last inequality, since the eigenvalues of ∇ 2 f ( ˜ B ) are contained in the interval [ µ, L ] , the eigenvalues of I -α ∇ 2 f ( z ) are between min(1 -αL, 1 -αµ ) and max(1 -αL, 1 -αµ ) . Note that ρ ( α ) is minimized at α = 2 µ + L with minimum value L -µ L + µ . It follows that for any t ≥ 0 ,

Next, we show (ii) (48) by induction. Indeed, it holds for t = 1 by the choice of ν . Denoting ˆ t := γ + t , by the induction hypothesis and by the choices of α t = β/ ˆ t and ν ≥ β 2 σ 2 βρ -1 ,

This shows (ii) .

Next, we establish the following global landscape result of the ITR learning problem (6) by computing the Hessian of the objective L ( B ) in (6).

Lemma E.3 (Local landscape of the ITR objective) . Let L ( B ) denote the ITR objective in (6) .

(i) Suppose Assumption 3.1 holds. For the continuous outcome,

where the constants λ ± are defined in Assumption 3.1.

(ii) Further Assumptions 3.1 and 3.2 hold. For the binary outcome,

where the constants λ ± and α ± are defined in Assumptions 3.1 and 3.2, respectively.

To maintain the flow, we write down the proof of Lemma E.3 to Section F.

Now we can easily deduce Theorem E.1.

Proof of Theorem E.1 . Lemma E.3 shows that the eigenvalues of the Hessian ∇ vec( B ) ∇ vec( B ) T L of the ITR objective in (6) are uniformly bounded between constants µ and L as defined in (10). Under Assumptions 3.1 and 3.2, we have 0 < µ ≤ L < ∞ . This shows that L is µ -strongly convex and L -smooth over B , as claimed in Theorem E.1 (i) .

Next, noting that ∥·∥ F ≤ ∥·∥ 1 , triangle inequality and the fact that ˆ B 0 , B ⋆ ∈ B yields that, for any t ≥ 0 ,

Hence the assertions (i) -(ii) follow directly from combining Lemmas E.2 and E.3.

F Proof of Lemma E.3

In this section, we prove Lemma E.3. We will first compute the gradient and the Hessian of the ITR objective in (6) separately for the continuous and binary outcome. We will then prove Lemma E.3 at the end of the section.

F.1 Continuous outcome

Proposition F.1. Let L ( B ) denote the ITR objective in (6) with continuous outcome. Then

where Ψ is the weighted design matrix defined at (8) .

Proof. Recall that the ITR objective for continuous outcome is given by

It is trivial to show (54). For the Hessian, using (16),

This shows the assertion.

F.2 Binary outcome

Given the current model parameter B ∈ R p × ( K -1) , define the scalars z i , p i , w i for the i -th sample by

For each binary label y ∈ { 0 , 1 } and activation z ∈ R , recall the negative log-likelihood of observing y under the logistic model Ber ( p i ) is given by

An easy computation shows

For the forthcoming computations, define matrices

Proposition F.2. Let L ( B ) denote the ITR objective in (6) with binary outcome. Suppose Assumption 3.1 holds. Let z i := u T a i B T x i for i = 1 , . . . , n . Then

Proof. We first claim that

Observe that by using (18), we can write

Noting that vec(

B

T

)

T

= (

(

C

p,K

-

1)

B

))

T

= vec(

B

)

T

C

(

K

-

1

,p

)

, Hence we get

This verifies (65). Then by using the chain rule (20), we get

which verifies (66).

Next, we compute the gradients of K in (63). First, using (65) and the chain rule (20),

where the last equality uses commutation matrix relation (17). Then noting that vec( K ) T = K = [ ˙ h ( y 1 , z 1 ) , . . . , ˙ h ( y n , z n )] , it follows that

where (a) follows from (65) and the chain rule, (b) is an algebra, and (c) follows from the definition. This shows (67).

Now we compute the gradient:

where (a) follows from (66), (b) uses the commutation matrix relation (17), and (c) follows from the definition. Then by using (69) and (18) with the fact that ˙ h is a scalar, we deduce (63).

Next, we compute the Hessian. From (71),

Then using (66) with (72), we get

This shows the assertion.

F.3 Proof of Lemma E.3

We are now ready to give proof of Lemma E.3.

Proof of Lemma E.3 . Recall that Φ := [ u a 1 ⊗ x 1 , . . . , u a n ⊗ x n ] . Note that

where Ψ is the weighted design matrix defined at (8). (In particular, since W is a diagonal matrix of nonnegative entries, this shows that Ψ is positive semi-definite.)

According to Propositions F.1 and F.2, we have

Thus the assertion for the continuous outcome follows immediately. For the case of a binary outcome, by Assumption 3.2 we have

Then

It follows that

This shows the assertion for the binary outcome holds.

G A non-asymptotic consistency of weighted, constrained, and regularized MLE

In this section, we provide a general result on the non-asymptotic consistency of MLE in a general setting. Here, observations are assumed to be i.i.d. from a generative model, and the unknown true parameter of the generative model may lie on the boundary of the parameter space.

A generic sample ( X,Y ) consists of a random covariate X and outcome Y . (In the ITR setting, we will specialize X to be the pair of the patient's covariate and the corresponding treatment.) We assume the covariate X follows a probability distribution µ ( · ) and the conditional distribution of Y given X = x is parameterized as π θ ⋆ ( · | x ) for an unknown parameter θ ⋆ . Hence the generic random sample ( X,Y ) follows the parameterized distribution π θ ⋆ ( y | x ) µ ( x ) . Assume we know a smooth parametric family of conditional distributions θ ↦→ π θ ( · | x ) , where the true parameter θ ⋆ lies in a convex parameter space Θ ⊆ R d . Our goal is to estimate θ ⋆ from i.i.d. samples ( x i , y i ) , i = 1 , . . . , n drawn from π θ ⋆ ( · | · ) µ ( · ) .

To estimate θ ⋆ , a natural approach is to maximize the empirical joint likelihood function. Since the marginal distribution of the covariate x is independent of the parameter θ , we can apply the method of 'reweighting' the covariate distribution by introducing an additional random weight ω ( x ) that we can control. Namely, fix a weighting function ω ( · ) : X → [0 , ∞ ) , where X is covariate space, that assigns a nonnegative weight to each covariate value x . Consider then the following weighted expected log-likelihood maximization problem:

where we omit the marginal distribution µ ( X ) as it is independent of θ .

Recall that we get to design the weighting function ω ( · ) . We first observe that the true parameter θ ⋆ is a stationary point of the weighted expected log-likelihood function E [ -ω ( X ) log π θ ⋆ ( Y | X )] for arbitrary weighting function ω . This can be justified under the mild assumption of exchangeability of the expectation and derivatives (which holds under twice continuous differentiability of the log-likelihood function by the dominated convergence theorem). Indeed, note that

where the last equality follows from

Therefore, θ ⋆ is a critical point of E [ ω ( X ) log π θ ( Y | X )] . In particular, if we assume that the log-likelihood function log π θ ( y | x ) µ ( x ) is convex in θ , then θ ⋆ is a global maximizer of E [ ω ( X ) log π θ ( Y | X )] .

Next, we introduce the finite-sample approximation of (82). Denote w i := ω ( x i ) for the weight of the i th sample ( x i , y i ) . Then the weighted negative log-likelihood of the observed samples under the model parameter θ is

where R ( θ ) is a suitable choice of regularizer for parameter θ . The regularization term can be considered as prior knowledge about the model parameter θ .

Let ˆ θ n denote a (possibly non-unique) minimizer of the above function over Θ . This is a minimizer of the loss function L over the constraint set Θ , which we call the weighted, constrained, and regularized MLE of θ ⋆ .

As in the standard MLE analysis, we assume that the negative log-likelihood function θ ↦→L 0 ( θ ) is differentiable and strictly convex for all x ∈ R p . Statistical literature typically assumes that the true parameter lies in the interior of the constrained parameter space, Θ , a convex subset of R d . In contrast, we allow the true parameter θ ⋆ to lie on the boundary of Θ . Regardless of that, we have seen above that the gradient of the weighted expected log-likelihood function vanishes at the true parameter θ ⋆ .

In this setting, we aim to provide a high-probability guarantee that there exists a global minimizer of (83) that is close to the true parameter θ ⋆ for an arbitrary weighting function ω ( · ) , which we will optimize later for the best statistical guarantee. In Theorem G.1 below, we obtain such a result in the non-asymptotic, weighted, constrained, and regularized setting. For its proof, we combine a classical approach from [15] with concentration inequalities, namely, a classical Berry-Esseen bound for deviations from the standard normal distribution for independent but non-identically distributed random variables and a uniform McDirmid bound (Lemma G.2). The Berry-Esseen bound controls the linear term in the second-order Taylor expansion of the log-likelihood function ( T n ( θ ) in (96)), while the McDirmid bound controls the quadratic term ( S n ( θ ) in (96)). By using an ε -net argument, the latter concentration inequality can be extended to a setting where the random variables are parameterized within a compact set.

Theorem G.1 (Non-asymptotic consistency of weighted, constrained, and regularized MLE) . Consider the constrained, regularized, and weighted MLE problem (83) with unknown parameter θ ⋆ from a convex subset Θ ⊆ R d . Fix an arbitrary weighting function ω ( · ) . Assume the following holds:

(a1) (Smoothness) For each sample ( x, y ) , the per-sample negative log-likelihood function θ ↦→ log π θ ( y | x ) µ ( x ) is strictly convex, three-times continuously differentiable, and R ( θ ) is differentiable. Furthermore, denote U := ω ( X ) ∇ θ log π θ ⋆ ( Y | X ) ∈ R d and U := U -E [ U ] . Suppose there are constants D 1 , d 1 ∈ (0 , ∞ ) such that for all i ,

(a2) (First-order optimality) The true parameter θ ⋆ is a stationary point of the expected negative weighted log-likelihood function L 0 ( θ ) := E θ ⋆ [ -ω ( X ) log π θ ( Y | X )] over Θ :

(a3) (Approximate second-order optimality) Let ¯ L ( θ ) := L 0 ( θ ) + R ( θ ) denote the expected regularized negative log- likelihood function. Then the regularized Fisher information ∇ 2 ¯ L ( θ ) is positive definite at θ = θ ⋆ with minimum eigenvalue ρ > 0 .

Fix a constant C > 0 . Let D = Cn -1 / 2 + 8 ∥∇ R ( θ ⋆ ) ∥ ρ and M = M ( D ) > 0 denote the supremum of the absolute values of all third-order partial derivatives of L over all θ with ∥ θ -θ ⋆ ∥ ≤ D . Suppose ∥∇ R ( θ ⋆ ) ∥ is small enough so that

Then there are constants c 1 , c 2 , c 3 > 0 such that

That is, with high probability explicitly depending on C , ρ , d , and n , there exists a global minimizer of θ ↦→L ( θ ) in Θ within distance D from θ ⋆ .

It is important to notice that the minimum eigenvalue ρ of the weighted regularized Fisher information ∇ 2 L ( θ ) depends on the choice of weighting function ω ( · ) . Therefore, Theorem G.1 suggests that

it is best to choose the weighting function to maximize the minimum eigenvalue of the weighted regularized Fisher information in order to minimize the statistical estimation error. We elaborate more on this point in the proof of Theorem 3.5.

We devote the rest of this section to proving Theorem G.1 using Lemma G.2 and Theorem G.3.

Lemma G.2 (A uniform McDirmid's inequality) . Let X 1 , . . . , X n be independent random vectors in R p from a joint distribution π . Fix a compact parameter space Θ ⊆ R d and f θ : R d → [ -J, J ] is a bounded functional for each θ ∈ Θ such that

for some constant L > 0 . Further assume that E [ f θ ( X k )] = 0 for all θ ∈ Θ and k = 1 , . . . , n . Then there exist constant J > 0 such that for each n ≥ 0 , and η > 0 ,

Proof. Since Θ ⊆ R d is compact, it can be covered by a finite number of L 2 -balls of any given radius ε > 0 . Let U ε be such an open cover using the least number of balls of radius ε > 0 . Let N ( ε ) = |U ε | denote the least number of such balls to cover Θ . Moreover, let diam ( Θ ) denote the diameter of Θ , which is finite since Θ is compact. Then Θ is contained in a d -dimensional box of side length diam ( Θ ) . This box can be covered by ( diam ( Θ ) /ε ) d cubes of side length ε . Covering each such cube of side length ε by a ball of radius ε , it follows that

Next, fix η > 0 , θ ∈ Θ , and ε > 0 . Let θ 1 , · · · , θ N ( ε ) be the centers of balls in the open cover U ε . Then there exists 1 ≤ j ≤ N ( ε ) such that ∥ θ -θ j ∥ < ε . By the hypothesis, f θ depends on θ uniformly continuously with respect to the supremum norm. Hence there exists δ = δ ( ε ) > 0 such that

Denote H n ( θ ) := n -1 ∑ n k =1 f θ ( X k ) . Then this yields, almost surely,

Furthermore, by the assumption, ∥ f θ ∥ ∞ is uniformly bounded by J > 0 . It follows that for each θ ∈ Θ , H n ( θ ) changes its value at most by M when one of X 1 , . . . , X n is replaced arbitrarily. Therefore, by the standard McDirmid's inequality (see Theorem 2.9.1. in [51]) and a union bound, with choosing ε = η/ (2 L ) ,

The above holds for all n ≥ 1 and η > 0 . This shows the assertion.

Next, we recall the classical Berry-Esseen theorem for the rate of convergence of normal approximation for the sum of independent random variables due to Feller.

Theorem G.3 (Berry-Esseen, Theorem 2 in Ch. XVI.5 of Feller '91[17]) . Let X 1 , X 2 , . . . , X n be independent and not necessarily identically distributed random variables with zero means and finite variances. Define W = ∑ n i =1 X i let F be the distribution function of W and Φ be the standard normal distribution function. Then

Now we prove Theorem G.1.

Proof of Theorem G.1 . Suppose we have n i.i.d. observed samples ( x i , y i ) for i = 1 , . . . , n from distribution π θ ⋆ ( y | x ) µ ( x ) . Also by the assumption, L 0 is twice continuously differentiable and n -1 ∑ n i =1 w i = 1 , so E [ ∇L 0 ] = ∇ E [ L 0 ] and E [ ∇ 2 L 0 ] = ∇ 2 E [ L 0 ] by the dominated convergence theorem.

Note that, for any r > 0 ,

The inclusion ⊆ is clear. Conversely, since L is strictly convex, the event on the right implies that its unique global minimizer is within r from θ ⋆ , as desired. Thus in order to show the main result in (87), it suffices to show that the probability of the event on the right with r = D = Cn -1 / 2 + 8 ∥∇ R ( θ ⋆ ) ∥ ρ is at least the right-hand side of (87).

Fix θ ∈ Θ such that ∥ θ -θ ⋆ ∥ = D . We introduce two random variables that we will bound to be small by using some concentration inequalities:

Roughly speaking, these quantities are of order 1 with high probability, T n by Central Limit Theorem and S n by the law of large numbers. Later in this proof, we will use non-asymptotic versions of these classical limit theorems to obtain high probability bounds for related quantities.

Since θ ↦→L ( θ ) is assumed to be three-time continuously differentiable, the quantity M ≥ 0 in the assertion is well-defined and is finite. Then, using the Taylor expansion, we may write

We will lower bound the first two terms on the right-hand side above. Note that

where for (a) we use the fact that θ ⋆ is a stationary point of E [ L 0 ( θ )] over Θ and Cauchy-Schwarz inequality; for (b) we use the definition of T n ( θ ) .

Next, we turn our attention to the second-order term in the Taylor expansion (98). Recall that from the assumption (a3) in Theorem G.1,

where ρ > 0 is a constant. It follows that

Combining the above inequalities with noting that ∥ θ -θ ⋆ ∥ 2 = D 2 , we obtain

Note that I 1 ≥ 0 if

The choice of D holds the former condition, and the latter by the assumption (86). Thus, I 1 ≥ 0 .

We now take infimum over all θ ∈ Θ such that ∥ θ -θ ⋆ ∥ = D . It suffices to show that the random variable I 2 defined above is positive with high probability, since I 1 does not depend on θ . To this end, write

Then the last expression in (103) is at least ρ/ 16 if A ≥ -ρ/ 16 and B ≥ ρ/ 8 . By the hypothesis, D = O (1) so it is uniformly bounded. Then by the uniform McDirmid's inequality in Lemma G.2, there exist constants C ′ , C ′′ > 0 such that

Next, we will show the following inequalities: For K = 6 D 1 /d 3 / 2 1 ,

where Z ∼ N (0 , 1) is an independent standard normal random variable. (d) above is a simple consequence of the standard Gaussian tail bound P ( N (0 , 1) > x ) ≤ e -x 2 / 2 x √ 2 π and (d), (e) are from the choice of D which yields nD 2 ≥ C 2 . We will show (c) at the end of this proof. Note that for any two events E 1 , E 2 defined on the same probability space, P ( E 1 ∩ E 2 ) ≥ P ( E 1 ) + P ( E 2 ) -1 . Then

by combining (103), (104), and (105), we have

with the corresponding constants c 1 , c 2 , c 3 > 0 . Then the assertion (87) will follow using (95). It remains to verify (c) in (105). To this end, write

where U i = ω ( X i ) ∇ θ log π θ ⋆ ( Y i | X i ) ∈ R d and U i = U i -E [ U i ] as in the statement. Then by using Cauchy-Schwarz inequality and noting that ∥ θ -θ ⋆ ∥ = D , we get

It is important to note that the distribution of the random variable on the last term above does not depend on θ . Note that w i U i for i = 1 , . . . , n are independent mean zero (but possibly non-identical distributions due to the weights w i ) random vectors in R d , and by the hypothesis, their respective coordinates have uniformly bounded variances. Hence by a union bound,

Let Q k n = 1 √ n ∑ n i =1 w i U i ( k ) . Then by the Berry-Esseen Theorem (Theorem G.3) and the hypothesis, for Z ∼ N (0 , 1) ,

with K = 6 D 1 /d 3 / 2 1 . Combining with (106) and using a triangle inequality, we obtain

Thus (c) in (105) follows.

H Proof of Theorem 3.5

In this section, we provide proof of the computational and statistical estimation guarantee stated in Theorem (3.5) under the generative ITR model (11). Two key ingredients in our proof are the computational guarantee (Theorem. E.1) and non-asymptotic estimation guarantee of constrained and regularized MLE (Theorem. G.1).

As we have briefly explained in the main text, the connection between the weighted convex optimization (6) for ITR learning and the statistical estimation problem under the generative model (11) is established by writing down the corresponding maximum likelihood estimation problem. To see this, recall that covariate-treatment pair ( x i , a i ) follows a joint distribution π that does not depend on the true parameter B ⋆ . Further, assume that π admits a density function f π . Then, for the continuous setting, the joint log-likelihood of observing the triple ( x i , a i , y i ) gives under the generative model with parameter B

Thus, the standard (unweighted) maximum likelihood estimation of the true parameter B ⋆ can be written as

where we imposed additional L 1 -ball constraint for the parameter B . Now, if we use covariatebalancing weight w ( x i , a i ) for each subject i , the corresponding weighted MLE problem reads as

which is exactly the weighted convex optimization problem (6) for continuous outcome. A similar discussion applies to the binary outcome simply by noting that the joint log-likelihood is given by

Proof of Theorem 3.5 . Consider the generative ITR model (11) with true parameter B ⋆ . Fix ε > 0 . Let ( ̂ B t ) t ≥ 0 denote the estimated parameters by using the PGD algorithm (7) for the weighted empirical maximum likelihood estimation problem (6). Decompose the total estimation error into computational and statistical parts as

We wish to apply our general result on weighted and constrained MLE in (Theorem G.1) with L 2 -regularization R ( B ) = λ 2 2 ∥ B ∥ 2 F in order to get a high-probability bound on the term I 2 . For this, we first verify the hypothesis in Theorem G.1 for our generative ITR model (11). Assumption (a1) in Theorem G.1 follows from Assumption 3.4. For Assumption (a3) in Theorem G.1, denoting by ρ the minimum eigenvalue of the expected regularized Fisher information, by Lemma E.3 (also see (10)),

where λ -is the minimum eigenvalue of the expected weighted design matrix E [ Ψ ] in (8) (see Assumption 3.1). The above quantity is positive under the hypothesis.

Lastly, for Assumption (a2) in Theorem G.1, we need to show that the true parameter B ⋆ is a global maximizer of the expected log-likelihood function of our generative ITR model (11). This can be shown by a standard argument under the mild assumption that expectations and derivatives are exchangeable. For simplicity, we prove this for the case when B ⋆ is in the interior of the constraint set Θ . (The general case can be shown similarly by considering the first-order optimality condition for stationary points.) First note that, if we denote the likelihood of observing the triple ( X,A,Y ) under the model parameter B as p B ( X,A,Y ) = p B ( Y | X,A ) p ( X,A ) , then log p B ( X,A,Y ) is a concave function in B . Thus we only need to show that the true parameter B ⋆ is a critical point of the expected weighted log-likelihood function E [ ω ( A,X ) log p B ( X,A,Y )] . Indeed,

where we have used iterated expectation to condition first on the covariate X and then integrate out the treatment-response pair ( A,Y ) . Now note that

Thus B ⋆ is a critical point of E [log p B ( X,A,Y )] . By the concavity of the log-likelihood function, it follows that B ⋆ is a global maximizer of E [ ω ( A,X ) log p B ( X,A,Y )] . In particular, this verifies the assumption G.1 (a2) .

Now we can apply Theorem G.1 with L 2 -regularization R ( B ) = λ 2 2 ∥ B ∥ 2 F for our generative ITR model. Denoting

for C > 0 any constant, this gives

where c 1 , c 2 , c 3 > 0 are constants and d = p ( K -1) . By Assumption 3.4, the quantities √ n ∑ n i =1 w 3 i / (∑ n i =1 w 2 i ) 3 / 2 are uniformly bounded in n . Hence we can first choose n large enough so that the last two terms combined in the right-hand side in (113) are at most ε/ 2 , for which n ≳ ε -2 is sufficient. Then we can choose C > 0 large enough so that the first term on the right-hand side of (113) is at most ε/ 2 . For this, we choose

Accordingly, we can choose T large enough so that I 1 ≤ C/ √ n . Then the above gives

This shows the assertion.

I Additional simulation details and results

The code supporting this study is available at https://github.com/ljw9510/effective-ITR, with plans for release as an R package soon.

I.1 Details of data-generating mechanism

In this section, we explain how data are generated and treatment assignments are allocated in our simulation studies. The outcome is generated under the working model assumption given by Equation 1 with heteroscedastic random noise ϵ ∼ N (0 , σ 2 ( A, X )) . We have four different settings/cases for the treatment-free effect and interaction effect functions, which are outlined below.

Case 1: Linear interaction functions with simple treatment-free effect functions

For Case 1, our randomized study aligns with Scenario 7 and our observational study mirrors Scenario 8 described in [44]. Specifically, we consider the following settings for the treatment-free effects and interaction effects:

Randomized trial:

Observational study:

Case 2: Linear interaction functions with complicated treatment-free effect functions

For Case 2, we consider the following treatment-free effect function for model assumptions of Y :

Case 3: Non-linear interaction functions with simple treatment-free effect functions

For Case 3, we consider the following non-linear interaction functions for model assumptions of Y :

Case 4: Non-linear interaction functions with a complicated treatment-free effect function

For Case 4, we consider the following complicated treatment-free effect function for model assumptions of Y :

In all four cases, we adopt the same treatment assignment principles and variance function σ 2 ( A, X ) that induces heteroscedastic error for observational studies as outlined in [44]:

For randomized trials, each individual has an equal probability of assigning each treatment, therefore, no propensity score estimation was made for AD-learning and SABD-learning. For the heteroscedastic error for randomized trials, σ 2 ( A, X ) = 0 . 25 + 0 . 2 (1 . 5 -X 2 ) 2 was used.

I.2 Implementation details

All numerical experiments were performed on a 2022 Macbook Air with M1 chip and 16 GB of RAM.

In the computation of the final ITR estimates, we use two different approaches: one with weighted multivariate regression with L 1 -regularization for the penalized approach (named 'penalized' in the main text). The other uses the proposed PGD algorithm with L 1 -ball (named 'constrained' in the main text). To ensure a fair comparison, both the penalized method and the constrained optimization algorithm were implemented using code developed by the author. We do not consider additional L 2 -regularization in this setting. To determine the regularization constant or L 1 -ball size, we use mean squared error as the criterion. The iterate number T of the PGD algorithm is 1000.

In our simulation study, we employ the default configuration of the random forest algorithm as implemented in the randomForest package in R to estimate the residual variance function in SABDLearning, the treatment-free effect, and the propensity score. This default configuration includes the construction of an ensemble of 500 decision trees, with each tree's splitting criterion determined using the Gini index. The number of variables considered for splitting at each node is set to the square root of the total number of covariates. For model performance evaluation, we utilize the out-of-bag (OOB) error estimate.

I.3 Additional simulation results

The goal of these additional simulations in the appendix is to demonstrate, under the pre-specified linear ITR, that DCBW, variable screening, and augmentation can synergistically benefit the final ITR. Methods using energy balancing weights are denoted as '_e' and those employing the distance covariance test for variable screening are indicated as '_s' in the table. Note that the augmented 'SABD_e_s' corresponds to ' proposed+penalized' in the main text. In this section, we include the performance of various methods implemented only by the penalized approach to compare the effectiveness of the proposed statistical methods. The simulation was implemented using package glmnet in R.

I.3.1 Treatment Decision Accuracy

Here, we show the averages and standard deviations of treatment decision accuracy from the additional simulation studies. The first four tables present the results of experiments that mimic randomized trials, and the remaining four tables show the results of observational studies.

The simulation results indicate that individual components alone do not significantly improve performance. Instead, the synergistic effect of combining these components is essential for the ITR-Learning. For example, in high-dimensional settings, using DCBW alone balances the empirical distributions of all covariates (up to 60), including a significant portion of irrelevant ones. Additionally, estimated decision functions may include irrelevant variables, reducing the impact of DCBW. Similarly, using variable screening alone remains challenging due to model misspecification or highly nonlinear outcomes. Combining other methods, such as outcome augmentation, helps mitigate model misspecification effects. This combined application ensures efficient and robust optimal ITR estimation across diverse scenarios.

Below is a detailed investigation of each scenario.

First, we examine Case 1, where the treatment rule class is correctly specified (Tables 1, 5). With a relatively simple treatment-free effect, AD-Learning's performance is comparable to SABD-Learning, particularly following outcome augmentation. When combined with EBWs as AD_e, SABD_e, these methods exhibit improved performance over their original versions (AD, SABD). In randomized trials, EBWs effectively reduce finite sample imbalance compared to IPW with true propensity scores. In observational studies, energy balancing yields more effective balancing weights, resulting in higher accuracy.

Next, we delve into Case 2, involving a more complicated treatment-free effect function compared to Case 1 (Tables 2, 6). Notably, since the treatment-free effect does not affect ITR learning, a misspecified treatment-free effect function introduces heterogeneity and increases noise variance. Then we can see the synergistic effectiveness of outcome augmentation and inverse variance weighting in SABD-Learning, effectively handling additional heterogeneity in ITR estimation (See 'augmented' panel in Tables 2, 6). However, it's important to note that all variants of AD-Learning show subpar performance, even with augmented outcomes. AD-Learning results in an intercept-only model frequently as its estimated decision function, which suggests that the mean squared error criterion in LASSO regularization may not be effective for selecting models. In such instances, different penalty types or performance metrics could prove beneficial.

Case 3 and Case 4 handle situations where the underlying treatment rule class is misspecified (Tables 3, 4, 7, and 8). Particularly, in Case 4, estimating ITR becomes even more challenging due to a highly non-linear covariate-outcome relationship, which leads to suboptimal results of the variants of ADLearning as in Case 2. Our results demonstrate that the more robust EBWs outperform standard IPW, particularly in observational studies. Additionally, variable screening enhances ITR learning in a highdimensional setting. Employing variable screening (SABD_s) enhances the performance of SABD, especially with limited sample sizes, as it aids in selecting effect modifiers and precision variables, thereby facilitating decision function estimation. However, even with perfect screening, estimating the decision function remains challenging due to model misspecification or highly nonlinear outcomes. Simultaneous use of EBWs and screening for SABD (SABD_e_s) further elevates performance, likely influenced by dimensionality in nonparametric balancing weight methods.

The significant enhancements from integrating DCBWs, variable screening, outcome augmentation, and inverse variance weighting underscore their collective importance. Their combined application ensures efficient and robust optimal ITR estimation across diverse scenarios.

I.3.2 Empirical values

In this subsection, we show the averages and standard deviations of empirical values from the additional simulation studies. The empirical values of the estimated rule E [ Y ( ˆ d ( X ))] are computed by utilizing the generated potential outcomes on the test dataset.

We include the empirical value of the optimal rule E [ Y ( d opt ( X ))] as well as the values obtained by assigning everyone to each of the four classes E [ Y ( K )] for K = 1 , . . . , 4 as benchmarks, assuming the higher value of empirical value is better. Again, the first four tables present the results of experiments that mimic randomized trials, and the remaining four tables show the results of observational studies.

Here, we can observe the benefits of the three components that we found from the treatment decision accuracy results. In addition, the benchmarks of the empirical values reveal some intriguing insights. Firstly, the optimal empirical value provides information on how challenging the current scenario is. For instance, in the randomized trial of Case 1 (Table 9), the optimal empirical value is not significantly different from the mean empirical value evaluated by current methods. Especially with augmented outcomes using a sample size of 1000, it yields a value quite similar to the optimal. On the other hand, in the observational study of Case 4 (Table 16), the obtained empirical values show significant differences from the optimal value. Furthermore, the mean empirical value evaluated from AD-Learning with a sample size of 200 is not markedly different from the empirical value obtained by assigning all individuals to treatment 1. This implies that the estimated ITR provides results akin to a 'one-size-fits-all' outcome in small sample sizes, and this finding is not irrelevant to the fact that a majority of AD-Learning results in an intercept-only model as the final decision function under conditions of highly non-linear covariate-outcome relationships. In such instances, we can explore alternative methods to further advance our current results such as other penalty types or performance metrics.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Our paper discusses the limitations of the work in the section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide the full set of assumptions in the main text and complete proof in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Our paper fully discloses all the information needed to reproduce all simulation results and data applications of the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide open access to the code, with sufficient instructions to faithfully reproduce the main experimental results.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We specify all the training and test details including data splits, hyperparameters, how they were chosen, and type of optimizer for the benchmark.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report the means and standard deviations with 100 (10) iterations in our simulation studies (applications), respectively.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide sufficient information on the computer resources in the appendix. Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We strictly follow the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss both potential positive societal impacts and negative societal impacts of our work in the section 6.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cite the original papers and include the URL for the webpage.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We introduce our proposed with sufficient details in the paper. The documentation is provided alongside our code.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We used open-source datasets and our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a framework for estimating individualized treatment rules (ITRs) in precision medicine applications. ​ The traditional methods for ITR estimation rely on inverse probability weighting (IPW) and L1-penalization, but these methods may have limitations such as statistical bias and computational bias. ​ The proposed framework uses model-free distributional covariate balancing and a hard L1-ball constraint to address these issues. ​ The optimal ITR is computed using projected gradient descent (PGD). ​ The paper provides a comprehensive analysis of the framework, including convergence guarantees and statistical estimation guarantees. ​ Simulation studies demonstrate that the proposed framework outperforms existing methods in terms of robustness and effectiveness for ITR learning.\n\nSTRENGTHS:\n1. The focus on weight optimization to enhance the efficiency of Individualized Treatment Rules (ITR) introduces a useful approach. The paper proposes an algorithm using Proximal Gradient Descent (PGD), providing a new angle on optimizing treatment protocols.\n\n2. The paper is supported by a solid analytical framework, combining theoretical guarantees with simulation studies and real data analysis. This mix of theoretical and empirical evidence effectively demonstrates the method's reliability.\n\n3. The research contributes to the field of personalized medicine by potentially improving the efficiency of treatment customization. Its impact may extend to influencing healthcare practices and patient management.\n\nWEAKNESSES:\n1. Previous papers like \"Balanced policy evaluation and learning\" from NIPS 2018 and \"More efficient policy learning via optimal retargeting\" in JASA 2021 also emphasized weight adjustments, focusing on minimizing the MSE of the value estimate. How does this paper differ from those approaches in its methodology and outcomes?\n2. The author claims that using distributional covariate balancing weights rather than Inverse Probability Weighting (IPW) can reduce the finite sample bias in ITR-Learning. Can the author theoretically connect this reduction to the findings in Theorem 3.5?\n3. This paper applies Proximal Gradient Descent (PGD) with constrained optimization to address the ITR problem. Is there a comparison available regarding computation speed between this method and others used in similar contexts?\n4. The simulations describe two versions of the proposed method, labeled \"penalized\" and \"constrained,\" as detailed in Appendix H.2. What are the main differences in tuning strategies between these versions? For a fair comparison, should the same tuning strategy be applied to both methods, or is it justifiable to use different strategies such as one focusing on MSE and the other on empirical value? Clarifying these differences could be crucial, especially if the paper aims to highlight the computational efficiency of the PGD over the penalization method.\n5. Some typos: In equation (2), it should be $u_k$ rather than $u$?\n6. The title of the paper with \"multi-category trts\" can be confusing. The paper focus on adjusting the weight for robustness and efficiency, rather than extending binary trts to multiple ones. The benefit of dealing with multiple trts may only come from Angle-based direct learning method.\n\nQUESTIONS:\n1. The paper provides a theoretical guarantee for the convergence of estimation parameters $B$. Does this focus on scenarios with fixed covariates $p$ and a fixed number of treatments $k$? Is it feasible to extend this to high-dimensional cases or scenarios where $p$ and \n$k$ depend on the sample size $n$?\n\n2. How does the paper estimate the treatment-free effect in the regression model? Could potential model misspecification for the treatment-free effect impact the performance of the ITR?\n\n3. The paper mentions the use of an ad-hoc Bonferroni correction for variable selection. What is the accuracy of this method for ensuring variable selection consistency? In the NIPS 2022 paper \"Learning Individualized Treatment Rules with Many Treatments: A Supervised Clustering Approach Using Adaptive Fusion,\" a group-lasso method is employed to filter out variables that only contribute to the treatment-free effect, thereby reducing the covariate dimension in $B$. Could further discussion and investigation into this method enhance the performance?\n\nLIMITATIONS:\nPlease see it in above weakness and questions section.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an approach for estimating individualized treatment rules (ITRs) for linear-in-feature decisions involving multi-category treatments. Unlike previous methods, the authors utilize distributional covariate balancing instead of inverse propensity weighting, and apply a combination of $L_1$ and $L_2$ penalties for regularization. The method is supported by rigorous computational and statistical analysis. Simulations on both synthetic and real datasets demonstrate its effectiveness, particularly in small sample, high-dimensional data settings.\n\nSTRENGTHS:\nThe paper is well written and proposes a more robust estimator for linear ITRs in small data settings by avoiding the use of inverse-propensity weighting, which can yield biased estimates in finite samples. The paper is introduces a somewhat novel approach in that utilizes distributional weighting and incorporates a hybrid $L_1$-$L_2$ regularization scheme. This approach is supported by rigorous theoretical analysis and thorough experimental evaluation.\n\nWEAKNESSES:\n*  The paper focuses mostly on linear decision rules for small, high-dimensional datasets which should be clearly stated somewhere in the abstract and emphasized in the introduction. Thus, the scope of the paper is more narrow than advertised.\n*  The approach is somewhat incremental and not necessarily novel as covariate balancing via distributional weighting has been used in treatment effect estimation tasks before. Furthermore, while the theoretical analysis is robust and valuable, it primarily mirrors that of linear regression with mixed $L_1$-$L_2$ regularization.\n*  The paper needs a thorough literature review and comparisons with literature on policy learning and treatment effect estimation (see, for e.g. [1]). \n* Missing bias-variance tradeoff curves between the methods. The main appeal of this method is that it reduces the finite sample bias by not using IPW. It would be insightful to see how this reduction in bias plays out in the finite sample scenarios within the simulations.\n* Missing baselines: Additional baselines should be considered. Eq. 2 suggests that using a linear regression model with k as a one-hot-encoded feature, along with regularization (either $L_1$ or a mixture of $L_1$ and $L_2$) applied solely to the non-k coefficients, could serve as a viable alternative.\n\nOverall, I believe this paper is strong from a technical standpoint but has limited impact compared to existing literature. It is correct and technically meets the criteria for acceptance. However, I remain undecided and am open to revising my score based on further discussion.\n\n[1] Nekipelov, Denis, Vira Semenova, and Vasilis Syrgkanis. \"Regularised orthogonal machine learning for nonlinear semiparametric models.\" The Econometrics Journal 25.1 (2022): 233-255.\n\nQUESTIONS:\nSee points 4 and 5 above. See points 4 and 5 above. Could you provide a limited evaluation to demonstrate the improvements of your method (in both bias and accuracy) over the IPW method and the proposed outcome-based baseline?\n\nLIMITATIONS:\nThe authors have aqequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a framework for estimating individualized treatment rules (ITRs) with multi-category treatments to address the problem of misspecified propensity score in inverse probability weighting (IPW) and the computation bias of L1 penalization. The authors propose using energy balancing weights (EBWs) for weighting and a hard L1-ball constraint to maintain objective smoothness, with projected gradient descent (PGD) for optimization. Theoretical results have been provided regarding the convergence rate of PGD and the estimated parameters in the linear ITR model. A simulation study and two real datasets are used to compare the proposed method with baseline methods.\n\nSTRENGTHS:\n-\tThe paper presents theoretical guarantees for the proposed method in terms of both the convergence rate of the computation algorithm and the statistical properties of the estimated ITR parameters.\n-\tExtensive simulation and real data analysis have been conducted to test the proposed method.\n\nWEAKNESSES:\n-\tFigure 1 only shows the average accuracy of different methods, although the standard deviations are reported in the appendix. It would be more informative to also plot the 95% confidence intervals to demonstrate whether the differences between methods are significant.\n-\tTable 1 suggests that the empirical value function of the proposed method has a large standard deviation. The difference between the value function of the proposed method and other baseline methods is smaller than one standard deviation (except when the training size is 3000 or 5000 in the Email dataset), providing insufficient evidence to demonstrate the superiority of the proposed method.\n-\tThe tables in the appendix may overwhelm readers. Replacing them with figures that highlight the differences between methods could be more beneficial.\n\nQUESTIONS:\n-\tIn Figure 1, why is the “AD + Energy” line in the left column different from the “Proposed + Constrained” line in the middle column? Have the methods in the left column used additional techniques like outcome augmentation?\n-\tIt would be helpful for readers’ comprehension if the authors could add a high-level summary and a pointer to the appendix in Section 2.2, explaining how the EBW is constructed and used.\n-\tThe vector $\\mathbf{u}$ in equation (2) has not been defined.\n\nLIMITATIONS:\nThe authors have discussed the limitations of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an algorithm for treatment rule estimation under standard no unmeansured confounding + SUTVA assumption. The algorithm builds on the AD-learning approach but changes to energy balancing weights and different regularizations. The paper then concludes with both simulations and real data example.\n\nSTRENGTHS:\nThe paper combines AD-learning with other useful techniques to improve treatment rule estimation and have theoretical guarantees.\n\nWEAKNESSES:\nI think the main weakness comes from two aspects: presentation and originality. \nPresentation:\n1. There is no seperate related work section. \n2. A lot of the main contexts (for example the algorithm 1, details on augmentation and inverse variance weighting) are put into the appendix. If the paper is too long to make it self-contained in 9 pages, I feel submitting to a journal is better suited. \n3. The theoretical sections seem just there to have some resultsm. The convergence of the algorithm is not really algorithm 1 specific and the statistical guarantees rely on strong generative model assumptions. \n4. Some detailed questions in the questions section. \n\nOriginality:\n1. From my understanding (please correct me if I am wrong), all the crucial pieces of the algorithm are already in the literature and the paper simply just combines them together. For example, balancing weights, AD-learning all are not original.\n\nQUESTIONS:\n1. Can you explain how you get the second \"=\" in equation (2)? Is it by assumption?\n2. What is the definition of $u$ in equation (2)? It took me to the original AD-learning paper to figure out so you should be clear here. \n3. On line 101, 'which inherently satisfies sum-to-zero constraint': I understand the $u_k$'s sum to zero but I thought sum-to-zero means $\\delta_k(x)$?\n4. In figure 1, are there error bars? \n5. In figure 1 the leftmost subplot, it seems in RCT energy weights improve more, which seems counterintuitive. And also, would it make sense to also compare observational + linear with RCT + linear just to see whether the weights help? Similarly we can also compare observational + linear vs observational non-linear. \n6. In table 1, do you know why the SE decreases then increases again?\n\nLIMITATIONS:\nAddressed in the paper", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-21"}