An effective framework for estimating individualized treatment rules
Authors: Joowon Lee
Jared Davis Huling
Guanhua Chen
Venue: NeurIPS 2024 poster (2024)
Keywords: Covariate Balancing
Multicategory Treatment
Precision Medicine
Projected Gradient Descent
Convergence Analysis
Constrained Optimization
Primary area: causal_inference

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
Joowon Lee', Jared D. Huling", Guanhua Chen'
University of Lisconsin-Madisonn University of Minnesota j joowon. lee@wisc, edu, huling@umn. edu, gchen250visc.ent edu
Abstract
Estimating individualized treatment rules (ITRs) is fundamental in causal inference, particularly for precision medicine applications. Traditional ITR estimation meth- ods rely on inverse probability weighting (IPW) to address confounding factors and L1 Li-penalization for simplicity and interpretability. However, IPW can introduce statistical bias without precise propensity score modeling, while L1 -penallization makes the objective non-smooth, leading computational bias and requiring sub- gradient methods. In this paper, we propose a unified ITR estimation framework formulated as 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 '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.
Introduction
Traditional medicine often uses a "one-size-fitsall" approach where the same treatment is applied to all patients regardless of their unique attributes, aiming to find single optimal treatment that may be most effective for 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 's unique characteristics [1]].
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 socio- psychological aspects 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-depte 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 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 optimization maximum likelihood estimation (MLE) problem, where the choice of weights directly impacts both the optimization and likelihood landscapes. This, in
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

tur, 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 ITT 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-categon treatments- relatively underexplored area-we 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 L1 -ball constraint to promote sparsity in regression coefficients, replacing the traditional soft L -pnpaaalizttion. This approach keeps the optimization objective smooth, enabling fast and accurate ITR estimation using projected gradient descent (PGD). Furthermore, we integrate techniques such as (I) variable screening. (2) outcome augmentation, and (3) inverse variance weighting to account for heteroscedastic errors, which together show 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 sellance-sekking objectives rather than relying solely on propensity scores. These approaches include entropy balancing weights [20], stable balaneing 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 retargeting the population covariate distribution, which is conceptually related
analysis, their approach relies on a reference policy, whereas our method does not. We also explore various statistical techniques to improve LeR-near First, we perform variable screening using the distance covariance test [48] 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- categoyy treatments. Lastly, with inverse variance weighting, as introduced in stabilized angle- angle-based direct learning (SABD-Leaming! 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-Leerning under a hard L1 consttarii
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 unified framework that combines DCBWs, variable screening, outcome augmentation, and inverse variance weighting in a synergistic and effective way. We demonstrated theoretical justi- fications for the combined approach, which, to our knowledge, have not been previously established.
2

---

2 Methods for ITR Estimation
2.1 Model Setup
Individualized Treatment Rule and Optimal Decision Function. We observe random vector (X, A, Y) where X = (2,... xp) E X' denotes the p-dimensional vector of pre-treatment covariates, received treatment A E A 21. K. and the corresponding outcome Y. Without loss of generality, we assume higher Y values are more favorable. Let ((a) denote the random variable that describes the potential outcome that would have been observed were the individual assigned to treatment E A. We make the following standard assumptions using the potential outcome framework [43]: (i) Stable Unit Treatment Value Assumption (SUTVA): Y (A), (ii)
positivity: 0 < (a, X) := P(A- x < 1 for all a E A and all X X, (iii) no unmeasured confounding: (a) A A for any E A.
An ITR, d(x), 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 Vdd:
Y 1(A d(x))
V(d) E (d(x)) E| d(x) =E
xAA,
Then the optimal ITR ddps is defined as function that maximizes the expected potential outcome (x) E arg max V(d) among all functions belonging to the treatment rule class (D). Building
on this, we adopt the following working model:
K
Y=
+ Σ (x) +,
k=1
where (xX) is treatment-free effect, (X) is interaction effect between covariates and kth treatment, and is random noise with mean zero and variance (A, For model identifiability, we assume Ex= (x) = 0 for all x. Under this model, 6z (X) determines the optimal ITR, while μ(x) has no impact on the ITR. One approach to identifying the optimal ITR is 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-Leaming [39], each of the K treatments represented as a vertex simplex on RK-1 denoted as Ug. h 1,..., K. In particular, up is defined as
if k = 1,
U1
(K- 1K-1 V K
if 2 2,
where I K-1 is K - dimensional vector with entries 1, and e is a K dimensional vector with entries 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ş have the same Euclidean norm 1.
Then the optimal ITR is reformulated as follows:
PP (x)
arg max
[X A= arg max us kE(L....KK
X
RE(L.....K)
π A,
x)
=:fap
where is random treatment vector in R corresponding the random treatment A. Then estimating the optimal ITR can be converted into estimating the decision function f(x) (x), (k-,..,j--i(x))) the optimal decision function (-) 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
assigning a K 1 dimensional vector to each covariate Here, while
F (f( = B B E (K--)).
Reduction Weighted Convex Optimization. Suppose we have a sample (x;, a, a;, yt) of size n from the joint population distribution for (X, A, Y).. Estimating the optimal decision function fopt from the observed finite sample can be reduced to solving weighted convex optimization problem:
min
w(a,, x;) f(y,, x,, a; B) Σ 1m1
))+ R(B).
(3)
BE
3

---

There are three key components in the problem (3): the per-sample objective function L, the per- sample weights w, and the regularization R(B). First, l denotes model-dependent per-sample convex objective function given by
R / B" B'xx) log (1 exp(u?. B x;)) you B for binary outcome.
for continuous outcome,
E(ys, Xz. 0g; B)
(4)
An important characteristic of causal inference problems is that only one of the possible treatments can be observed for each subject. Thus, it 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 wa = w(a,, x.) for the ith subject. A classic approach is to use IPW: w(a, x) (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 (I), optimal decision function fops in (2) can be decomposed as
u/
fopr
p(x)E
6 (
+E . (A, x
E|e[x].
(5)
. (A, x
π (A, x)
treatment-frce
interaction
x] and E[e]x] O, it follows that (x) Σ (x)xux depending only on the interaction term. However, suppose the true propensity score is not correctly estimated. In that case, the estimated decision function may 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 (A = a/x) that deviates from the true propensity score ( = x)).. Second, even with correct specification, insufficient sample size may lead to inaccurate finite-sample approximation P, (A a/x) of the true propensity score <AA alx) [31]. It results in systematic bias for the estimated coefficient of the treatment-frre term:
Since E[ [(A.)
K
Ê,
Pn(A-
k(x) + O.
A. x
k=
Instead of using IPW, we use DCBWs such as EBWs. Such weights are data-driven and do not rely specific model assumptions large sample approximation. Furthermore, EBWs known promote independence between treatment and confounders [23], E. [uw(A, x R E,, [u] E.,, (w(( A, x) x] since i.,, [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 L-zregularization and a hard L, ball constraint:
min [||
C(B) = Σ wla x:) Ely), Tz. a;; B)
+
A2 2
B
B
(6)
BE
The above formulation has a smooth objective function that unaffected by the hard -all constraint. Also, Logggllatinniionn ensures that the objective is at least Ag-strongly convex. Thus, the

---

PGD algorithm (with suitable stepsize) enjoys strong exponential convergence guarantee toward the global optimum. In contrast, standard upproaches with L pennaiizatin loss need sub-grade 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-constraiie optimization problem in (6). In each iteration, involves conducting gradient descent with a suitable stepsize at, followed by a projection Bgg onto the L1-ball B == {B [
[B| I A1}:
B, f 11 (B,-1- a VBC(B,-1))
For the projection onto the Li-projection, we the algorithm in [14]. Detailed implementation of this PGD algorithm is discussed in Algorithm 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, 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-LLearning. See Algorithm in the appendix.
improve the estimation of outcome regression
Lastly, we propose using both augmented outcomes and inverse variance weighting for variance reduc- tion learning ITRs. Under the working model assumption in (1), the interaction effect determines optimal treatment assignments, while the treatment-free freee 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 robust and efficient approach to ITRRLLerning, they are complementary methods. Details discussed in the appendix C.3 and C.4.
3 Statement of results
Our theoretical analysis shows that "weighted design matrix", specifically a p((K 1) x p(K 1) positive semi-definite matrix, plays a central role in the local landscape analysis of the ITR estimation problem:
V
11 Σ
waa x:) (Un, u..) & (xix?),
(8)
where @ denotes the Kronecker product. Without loss of generality, we assume E, w(a;, x;) = ......
Note that the weights w, x4), i 1...., n affect the eigenvalues of the weighted design matrix V. 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; such that the minimum eigenvalue of & is as large as possible. Roughly speaking, this is achieved when Wi.... Wn 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 DCBW Ws 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), require that the eigenvalues of are uniformly bounded.
Assumption 3.1 (Figenvalue bounds on the weighted design matrix). There are constants 0 S A- S t such that the eigenvalues of & in (8) is between X and +
5

---

For each subject i = 1,... let 2; and denote its 'activation' and 'predictive probability', where := U B'xx X
== ex=(x))) + exp(z:)).
(9)
We need the activation 2; be uniformly bounded for our theoretical analysis of the binary outcome. Since the simplex vectors Un; 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 E [R in the constraint set B and observed sample = 1, ....,,.............. n. |B" xx S M. Consequently, there exist constants 0 < a-= S art S 1/4 such that a S ps - p.) - a for all i.
With hard Looom constraint the model parameter B, Assumption 3.2 be easily verified whenever the covariate vectors X; are uniformly bounded, which is a standard in the literature.
We establish the convergence rate of Algorithm I 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
X2 + X2 L
A2
for continuous outcome for binary outcome
(10)
a
a A2
We say a function f IRP + is justtrongly convex if f(x) - =| 2 is convex and L-smooth if f is differentiable and V is I-LLipschtz continuous.
Theorem 3.3 (Convergence rate of PGD for ITTR-Lerming). Let (Bz)ızo denote the sequence of parameters obtained by the PGD algorithm (7) for ITR-Learning problem (6) with arbitrary initialization Bo- Suppose Assumptions 3.1 and 3.2 hold. Let B denote the unique global optimum of (6). Then the following hold:
ii) (Optimization landscape) The ITR objective L(B) in (6), is p-strongly convex and L-smooth, where μ. L are as in (10).
(ii) (Linear convergence with fixed stepsize) Assume constant stepsize at a < I.. Denote the contraction constant p(oc)= = max oT,, atdo E (0, Then [B, B B p(a)"||B, - 1 for all 1 1.
(iii) (Sublinear convergence with diminishing stepsize) Assume diminishing stepsize a 7+1' where B 3(,, L) is a sufficiently large constant and ? is an arbitrary constant. Then B, - B+1 < 7++ for all t 2 1 for some constant v >> 0.
Our algorithm 1 with soft T-z-penalization and hard constraint with T,y-ball is guaranteed to converge to the true solution B exponentially fast provided stepsizes a4 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 Theorem 3.5. Fix a joint distribution π for the pair (X, A) of covariates E R2 and treatment A {l...., K..,, Fix a true parameter B, E (Px((-1), Then assume i..dd samples (X,, ai, yi) are drawn as (x,, ag) ~ π and
K-1 N
K zz, a2)
for continuous outcome for binary outcome
Ye (il(x,, a;)
Ber
expla.) [+exp(#.)
where := u x,x,
(11)
where N(,, o²) denotes a normal distribution with mean and variance o², whereas Ber(p) denotes a Bernoulli distribution with mean p. We assume that π does not depend on B, and the noise variance o² 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- liglliieldoood function (up to additive constants) coincides with the per-sample loss function f 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,, y4) follows the gener- ative model above. Denote Uj := w; V VBL(y, X,, a; B,) E Rex(k-1) for == w w(a; x,) and
U. == U, E. Suppose there are constants D1. d1 E (0, oo) such that E| < D1 and
ming.1 Var(U, (k: D) > d1. Also, Ei= / (Ei=1 wf) 3/2 O( 1/2).
Now state main result regarding the statistical estimation guarantee for the true generative model parameter B.. With high probability, Algorithm can recover B, up to a statistical error O(n 1/2)..
Theorem 3.5 (Statistical estimation guarantee). Let (Xi, dis Ui, yi) be i.i.d. observations from the generative models in (11) with true parameter B, E [Rx((K-) such that |[B,|| S X for some A 2 Suppose Assumptions 3.1, 3.2, and 3.4 hold, Let By denote the weighted MLE obtained after T iterations of PGD algorithm (7) for (3). Fix constant O.
Then there exists constant C C(E) such that with probability least 1 and μ as in (10) but with x = (E(V))
B. - BrilF Colp/n)logc-
8A2
(12)
P
provided n, T are large enough. More specifically, the above holds when
(- u²)² (E w 2 2 some explicit constant C₁ 0,
ii) (Sample complexity) 2 large s.t.
(Computational complexity) T Oloo n) and On when constant or diminishing stepsize as in Theorem 3.3 is used, respectively.
Furthermore, we get Vn-consistent estimation whenever X2 Onn01/11n)1) 1/2 [B
Notice that both terms in the error bound (12) are proportional to 1/p, which depends both on the minimum eigenvalue X of ( as well as the Læ-regularization parameter X2. Therefore, choosing the weighting function w(A, X) to maximize A 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 Ly-regularization, but as in the second term in (12), using L2 regullariizztion gives estimation bias of [..2|BBBB| (F/p). Hence, we need to choose X2 small enough so that this bias term is also of (n 10--1//1 1/2).. Then we obtain an overall Vn-consistent estimator.
The key insight in our proof of Theorem 3.5 that we can decompose the total estimation error |B. - Br into computational and statistical parts:
[B. - BBll < B-
B. B
(13)
computational emmr
error
where B denotes the exact MILE The computational error vanishes as the number of PGD iterations tends to infinity according to Theorem 3.3. For the statistical error, we show
2C
8A₂ [. 2 20-0
Σn wp
P BB. - B S
(14)
Vn
1
(E w 3/3
Thus, the 'skewness" of the weights w; measured by the ratio /n Ei w (E w 8/2 (which arises from Berry-Essead bound) acts as a scalar multiple to the standard statistical error of O(n-112 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 e, provided we have at least ( samples. Then, simply need to ensure the PGD iteration T is large enough so that the computational error is at most C/ s/vn to achieve the high-pro gighhbrbbbbility total estimation error bound in Theorem 3.5. Another insight from the theorem is that maximizing Amin (E) promotes heterogeneous balancing weights w,, and we get increased
(E =
sample complexity through the bound
(27 w 2 Cie (c.g., if w. = 1, this is n 2 =2, but
W1 n, w2 w, 0, this may not be satisfied for any n))
7

---

4 Simulation Studies
We investigate four simulation settings, each designed to evaluate various factors influencing the per- formance different ITR estimators. For each setting, we mimic 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, (A, 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 assumees linear treatment rule class, the latter involves a situation where our model's 's prespecified treatment rule class is misspecified. Specifically, we use following treatment- effect function ((X) and interaction effect function 6(X) for each scenario:
1. Randomized Trial: Linear ITR as the true optimal (X) = 1 + 2X, + 2X2,
- A 6(X)= 0.75 1.5X, - 1.5X2 - 1.5X3 1.5X A 3,0.75 1,5X, 1.5X, 1.5x, 1.5x, 14,
(0.75 + 1,5X, + 1.5X2 1.5X, + 1;0.75
Observational Study: Nonlinear ITR as the true optimal
(x) = 1 + 2X, + 2X2 + 2X, 2X3 + 2X, + 2e Xx + sin(Xs);
(0.5 + 1.0,, 2.0X, + 0.5X3 A= 1; 1.0 1.0, + 1.0X, 1.0X?, A = 2;
5(X) ( 1..5 + 2.0X, 1.0X4
1.0X A 3; 1.0 1.0X, 1.0.2 1.0X7, A 4.
the working model (1). Further details about simulation settings, other scenarios, and corresponding results are available in the appendix I.
Randoriized Trial: Linear ITR as the true optimal
60
= 60
100%
Acuuntt
Sampl 5tw
Surple Sime
Sampie Som
Observational Study: Nonlincar the true optimal
p= 60
p= 60
p
00%
805
60%
305
20%
20%
1000
1200
AD Constaliedd
assgg
Peraliadd
Arcaoeed' Cocaroleed
Figure 1: Accuracy Comparison: proposed methods vs. benchmark methods. All subplots in the same row share the same simulation setting, focusing on randomized trials with linear ITR as the true optimal rule (top) and observational studies with nonlinear ITR as the true optimal rule (bottom). Each subplot presents (Left) accuracy comparisons based on weights, illustrating the difference between the standard IPW approach of AD-Learning with the proposed approach using EBW, (Middle) accuracy comparisons based on optimization algorithms, illustrating the difference between the standard L1 Li-penalized approach against the proposed constrained optimization with PGD, (Right) evaluation of accuracies between existing standard approaches and the proposed method, which integrates EBWs, variance and dimension reduction techniques implemented through PGD. Error bars represent the standard errors of the mean (SEM) of accuracies across multiple simulations.

---

Using Distributional Covariate Balancing Weights. In the left subplots of Figure 1, the incorpora- tion 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 I, it also results in improvements even 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 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-L.eaming, outcome augmentation, and additional variable screening. For solving the I Ty-penalized regression problems, the R package gimnet is standard choice, which based on cyclic coordinate descent with not handle hard -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 the boundary of the L -bll, the solutions of constrained optimization and penalized optimization coincide. However, the minimizer is the interior point of L, -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).
hyperparameter tuning. However, gimnet does
Training Loss (log scale)
Constrained
Penalized
Constrained
Penalized
10
Training Accuracy (%)
80%
Ag
60%
5
10
20
30
10
30
10
30
10
20
Figure 2: Comparison of Optimization Methods by A1. X1 is -ball size of the constraint set in the constrained optimization and the regularization parameter with additional I=-regularization of the model parameter in the penalized regression, respectively.
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-Leaming, a treatment-covariite interaction model with Zz-regularization 'linear', and two popular bree- methods: 'policytree' [47, 59] and machiee learning-based 'causalDML' [10, 29]. Our proposed method integrates EBWs with SABD-Leaming, outcome augmentation, and additional variable screening for variance and dimension reduction, implemented via PGD. However, causalDML requires 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. 'lincar' performs well in the upper subplot, which simulates a randomized trial with a true optimal linear rule, 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 variable screening for performance improvement, is provided in the appendix I.
Additionally, our proposed method processes each dataset (p = 60, = 1000) in an average of 4.40 seconds, compared to 2.01 seconds for the penalization method using gimnet 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 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 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/marketiing methods based on individuals' attributes. Similar to [44], we randomly split the data into training set of (200, 420,,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 test dataset. The benchmarking methods include AD-Learning, SABD-Learning, and two tree-based methods (policytree [47] and causalDML [29]). However, Ma 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-Learnigg is not used, since the homoscedastic assumption is unnecessary for logistic regression. The results are shown in Table 1.
Table 1: Average empirical value functions across different approaches for ACTG/Email datasets. Mean values with the corresponding standard errors of the mean (SEM) in parentheses are provided. The highest- methods marked in bold.
Dataset
Training size
AD
SABD
policytree
causalDML
Proposed
200
37.624 (6.9)
33.645 (6.7)
36.871 (4.6)
42.301 (5.7)
400
51.329 (3.2)
53.644 (1.9)
47.191 (2)
49.552 (0.2)
54.868 (3.4)
ACTG
800
56.021 (1.5)
56,700 (1.6)
(2.6)
55,506 (2.1)
59.200 (2.5)
1000
56.825 (2.2)
56.963 (2.1)
56.005 (1.7)
56,591 (2.1)
59.552 (2.5)
1200
55.092 (2.6)
55.221 (2.8)
58.819 (3.4)
57.104 (2.8)
58.207 (2.9)
1000
0.169 (0.0)
0.171 (0.0)
0.172 (0.0)
0.175 (0.0)
Email
3000
0.176 (0.0)
0.177 (0.0)
0.177 (0.0)
0.181 (0.0)
5000
0.177 (0.0)
0.178 (0.0)
0.178 (0.0)
0.182 (0.0)
6 Discussion and Limitations
Our method introduces 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. Addi- tionally, 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 a useful approximation to nonlinear decision functions.
We acknowledge that although most, if not all, types of ITTR-Learnin 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 personal- ized care and advancing medical research disease mechanisms and drug development. Beyond healthcare, these methods can be applied in marketing strategies based on customer characteristics.
10

---

Acknowledgement
The authors would like to thank Dr. Michael Kosorok and Dr. Haoda Fu for their valuable comments and encouragement on the early version of this work. Lee and Chen G.'s efforts were partially supported by NSF Grant DMS-2005 MMS-2004336.
References
[1] 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.
[2] Susan Athey and Stefan Wager. Policy learning with observational data. Econometrica, 89(1)1333-161 2021.
[3] Laura B Balzer, Maya L Petersen, Mark 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-37322 2016.
[4] Léon Bottou. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT' pages 177-186. Springer, 2010.
[5] Stephen Boyd, Stephen Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
[6] Stephen Boyd, Lin Xiao, and Almir Mutapeic. Subgradient methods. lecture notes of EE392o, Stanford University, Autumn Quarter, 2004(01), 2003.
[7] Guanhua Chen, Donglin Zeng, and Michael R Kosorok. Personalized dose finding using outcome weighted learning. Amer. Statist, Assoc., 11 111666)):509-1551 2016.
[8] Rui Chen, Jared D Huling, Guanhua Chen, and Menggang Yu. Robust sample weighting to facilitate individualized treatment rule learning for target population. Biometrika, 11 110(())309- 329, 2024.
[9] Shuai Chen, Lu Tian, Tianxi Cai, and Menggang Yu. general statistical framework for subgroup identification and comparative treatment scoring. Biometrics, 13(4):1199-1209, 2017.
[10] Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, and Whitney Newey. Double/debiaseddeemma machine learning of treatment effects. Am. Econ. Rev., 107(5):261-65) 2017.
[11] Francis S Collins and Harold Varmus. new initiative on precision medicine. New England journal of medicine, 372(9):933-795, 2015.
[12] Yifan Cui and Eric Tchetgen Tchetgen. semiparametric instrumental variable approach to optimal treatment regimes under endogeneity, Journal of the American Statistical Association, 122- 2021.
[13] Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treat- ment effects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statisties, pages 1810-1818. PMLR, 2021.
[14] John Duchi, Shai Shaley- Shwartz, Yoram Singer, and Tushar Chandra. Efficient projections onto the 1-ball for learning in high dimensions. In Proceedings of the 25th international conference on Machine learning, pages 272-279, 2008.
[15] 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.

---

[16] 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.
[17] William Feller. An introduction to probability theory and its applications, Volume 2, volume 81.
John Wiley & Sons, 1991.
[18] Jean-Louis Goffin. On convergence rates of subgradient optimization methods. Mathematical programming, 13:329-347, 1977.
[19] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
[20] Jens Hainmueller. Entropy balancing for causal effects: A multivariate reweighting method to produce balanced samples in observational studies. Politicul analysis, 20(1):25 2012.
[21] Scott M Hammer, David Katzenstein, Michael D Hughes, Holly Gundacker, Robert 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 Journul of Medicine, 335(15); 100111010000000, 1996.
[22] Kevin Hillstrom. Minethatdata e-mail analytics and data mining challenge. attps://wws 2008. Ac-
splift-modelln lingeeningeoeeoo www/en/vO. cessed: 2024-04-28.
[23] Jared D Huling, Noah Greifer, and Guanhua Chen, Independence weights for causal inference with continuous treatments. Journal of the American Statistical Association, 119(546):1657- 1670, 2024.
[24] Jared D Huling and Simon Mak. Energy balancing of covariate distributions. Journal of Causal Inference, 12()):2022022, 2024.
[25] Kosuke Imai and Marc Ratkovic. Covariate balancing propensity score. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 76(1):243-263, 2014.
[26] Nathan Kallus. Balanced policy evaluation and learning. Advances in neural information processing systems, 31, 2018.
[27] Nathan Kallus. Generalized optimal matching methods for causal inference. The Journal of Machine Learning Research, 21(1):2300-235 2020.
[28] Nathan Kallus. More efficient policy learning via optimal retargeting. Journal of the American Statistical Association, 116(534):646-558 2021.
[29] Michael C Knaus. Double machine learning- based programme evaluation under unconfounded- ness. The Econometrics Journal, 25(3):602-627, 2022.
[30] Brian K Lee, Justin Lessler, and Elizabeth Stuart. Weight trimming and propensity score weighting. PloS one, 6(3):e18174, 2011.
[31] Fan Li, Kari Lock Morgan, and Alan M Zaslavsky. Balancing covariates via propensity score weighting. Journal of the American Statisticul Association, 1136521):39--400, 2018.
[32] Runze Li, Wei Zhong, and Liping Zhu. Feature screening via distance correlation learning. Journal of the American Statistical Association, 107(499): 129-49991119-11333 129-1119, 2012.
[33] 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: 1555661995,,,, 2022.
[34] 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, 1744):38443402, 2023.
12

---

[35] Weibin Mo and Yufeng Liu. Efficient learning of optimal individualized treatment rules for heteroscedastic misspecified treatment-fred -fre effect models. Journal of the Royal Staristical Society Series B: Statistical Methodology, 84(2)4440447 -42,, 2022.
[36] Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2):299-119 2021.
[37] 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, 1155529) 2019.
[38] 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-3444 2021.
[39] 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--5)) 1011666776111111 2020.
[40] 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.
[41] Min Qian and Susan A Murphy. Performance guarantees for individualized treatment rules. Annals of statistics, 39(2):1 180. 2011.
[42] Brian C Ross. Mutual information between discrete and continuous data sets. PloS one, 9(2)e87357, 2014.
[43] Donald B Rubin. Randomization analysis of experimental data: The fisher randomization test comment. Journal of the American statistical association, 75(371):591-595 1980.
[44] Kushal Shah, Haoda Fu, and Michael Kosorok. Stabilized direct learning for efficient estimation of individualized treatment rules. Biometrics, 2022. doi:
13818.
[45] Tao Shen and Yifan Cui. Optimal treatment regimes for proximal causal learning. Advances in Neural Information Processing Systems, 36, 2024.
[46] Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in neural information processing systems, 32, 2019.
[47] 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):223,, 2020.
[48] Gábor J SzŌkely, Maria L Rizzo, and Nail Bakirov. Mcasuring and testing dependence by correlation of distances. The Annals of Statistics, 35(6):2769-22) 2007.
[49] Dingke Tang, Dehan Kong, Wenliang Pan, and Linbo Wang. Ultra-high dimensional variable selection for doubly robust causal inference. Biometrics, 79(2):903 -944, 2023.
[50] Lu Tian, Ash Alizadeh, Andrew Gentles, and Robert Tibshirani. A simple method for estimating interactions between treatment and large number of covariates. Journal of the 2014.
American Statistical Association, 109(508):1517-15
[51] Roman Vershynin. High-dimenssioall probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018
[52] Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
[53] Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 113(623): 2018.
13

---

[54] Baqun Zhang, Anastasios A Tsiatis, Eric B Laber, and Marie Davidian. A robust method for estimating optimal treatment regimes. Biometrics, 68(4):010--1018, 2012.
[55] Chong Zhang and Yufeng Liu. Multicategory angle-based large-margin classification. Biometrika, 101(3):625-640 2014.
[56] 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-415 413, 2022.
[57] Yingqi Zhao, Donglin Zeng, John Rush, and Michael Kosorok. Estimating individual- ized treatment rules using outcome weighted learning. Journal of the American Statistical Association, (074499):106-1111 2012.
[58] Xin Zhou, Nicole Mayer-Hamblett, Umer Khan, and Michael Kosorok. Residual weighted learning for estimating individualized treatment rules. Journal of the American Statistical Association, 112(51)):169-1 2017.
[59] Zhengyuan Zhou, Susan Athey, and Stefan Wager. Offline multi-action policy learning: Gener- alization and optimization. Operations Research, 71(1):148-183 2023.
[60] Ruoqing Zhu, Ying-QU 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.
[61] José R Zubizarreta. Stable weights that balance covariates for estimation with incomplete outcome data, Journal of the American Statistical Association, 110(511):910-22) 2015.
14

---

A Preliminaries
In this section, we review several notations and basic facts on linear algebra and matrix calculus.
The Kronecker product, denoted by 8, is a binary operation that combines two matrices to create new matrix. Given two matrices = (acy) E Rmxx and B Rexq, then @ B an mp x rug matrix defined as
an B
an B
A A
(15)
Oml
Amn
If A' E Rmxx'', then for the horizontally stacked matrix [A, A E Rmx((+x'),
A A
B & BJ.
(16)
For each A [ai....,ayy an) E Rmxx let vec(A) == |af. TT E mo denote its vectorization. The commutation matrix Cla.)), which is a special instance of ab x ab permutation matrix. Namely, for each integers a, b 1, there exists a unique matrix C(a.b) E {0. 1) jabxab such that for all A E Raxb we have C(a.)) vec(A) veo(A"). Note that (C(+.))
C(b,a), Furthermore, (Cla.ii) C(a.6) = Las since (Cla.))) ((ab)) vec(A) - C(b.a) veo(A") = vec(A). Hence Cla.6) positive semi-definite. Note that C(a.1) I. C(1.a) since if A E Rax1 or A E RRxx veo (e) Also, for every matrices A E m and B E Tgrx,
Clrmm) (A
[ A)C(em)
(17)
Also recall the relations for vectorizing product of matrices: for A E Raxb B E Rèxe and C E Bexd
(AB) = (Ie @ A) vec(B) (BT Ia) vec(A),
(18)
rec(ABC) (C" & A) vec(B) (I. AB) veo(C) (C" In) vec(A).
(19)
For A, B, C, and D matrices of compatible sizes allowing the matrix products AC and BD,
(A & (C @ D) (AC @ BD).
Next, recall that if f Rax1 - bxx and Rbx1 + ex are differentiable functions, then the Jacobian Jy can be written by using gradients as (x) (V=f*)² )T By chain rule Jgof(x) (x)),, have the following chain rule for gradients
V (a(f(z))²) = V. (f(x)²) 5(a) (o(f(x))²)
(20)
We will also frequently the following fact: For A Raxb
V Aoo(A) voc(A)" Lab:
(21)
V Veo(A) (oc(AAT = V Vvec(A) (e)(A)² ((b.a) = C(b.a).
(22)
For two matrices A, B of the same size, write A 3 BffB A is positive semi-definite. The partial ordering = 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:
Y (X)
8u (X) k))
To make the optimal ITR identifiable, we impose the constraint Σ 5x(x) = 0 for all covariates X. Under this model, (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 AeD-nnnrnngg proposed by [39]. each of the К treatments represented as vertex simplex on R 1 denoted as ug, h = i...., K. In particular, is defined as
K-1
if = 1,
Up
1-VK (K-1)) 1--1
K-1
if 2 2,
where 1-1 is K dimensional vector with entries 1, and e is a K dimensional vector with entries except its k-th entry being 1. This vertex simplex has K vertices with equal angles between them and origin at the center of the simplex. All uy have the same Euclidean norm To estimate the optimal ITR, a decision function f(x) = (fi(x), (x)..... f (x)) is constructed, assigning K dimensional vector each covariate Here, f(.) can take generic form, while we assume the linear decision function in the main text. The optimal ITR is reformulated by the authors as follows:
d (x) = arg max fxx) = arg min L(ux, f(x)), (K.....K
KY
where -(-,-) represents the angle between two vectors. Note that due Ek u 0,
Σ ug f(x) for any given function f and covariate value x. Then we can use
K Σ
(u k)) f(x)
instead of Ek_ 8 (x) k) to encode the treatment and covariate interaction. Specifically, we use 4 to denote [EKn u u = k)) and assume f(x) to be linear in x (with the coefficient matrix B). Therefore, terms pertaining to ITRs can be expressed as u B
T
B.2 Details of binary outcome
Again, under prespecified treatment rule class D, the optimal ITR is given by
Yu
oP (x) = arg max E|Y A = k]
arg max u ke kE (K.....)
x
π (A, x
ke(1......
Fire (x)
For the binary response setting, we assume the following logistic model
exp(ug Fop (x)) I exp(u, (xp(
P[Y = 1|y= A=

---

Under this model, fopt is indeed the optimal decision function under the angle-based based direct learning framework since
(x) arg max P[Y 1|x, A X:] ke(L...,,K
exp(ug fop (x))
arg
(23)
kE kE(I....K) 1 ++ exp(ug fog (x))
arg max u (x).
ke(l...,.k
Furthermore, it can be checked that the function foot is an optimal solution to maximizing the IPW-weighted log- legliikl of the logistic regression model:
Yu"f TAA,))
log(exp(u + 1)
argmin arg min
(24)
(A,
under the assumption that the expectation and derivative w.r.t. can be exchanged ([39]).
For the effective and robust ITR estimation, replacing IPW with DCBW leads
arg min E wxx A) YYuuf + log(exp(u' ? + +)}
(25)
Now consider restricting it to the linear decision function class F {f(X) = B Be RPx (x-)) 1). Then the above (25) becomes
E [ A)(Yuu B log(exp(u" (u + 1)}]
arg
BEBxX(K
Then, we can estimate the optimal parameter B by solving the following empirical loss minimization
arg 1 Σ (KR w(xi, a:) (-y.u,B" BTx, + log(exp(u" T B")x;) + +))
s.t B
S (26)
where imposed additionally the L1 penaaaliiztiion |B]] S B promote sparse solution. By the proposed algorithm 1, we can efficiently solve the problem (26).
3

---

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 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 selanne--eeeingg objectives instead of using IPW have been proposed. These approaches include entropy balancing weights [20], stable balancing weights [61] as well 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 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 IPW for ITR-Leane TTR-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 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, 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 weighted empirical cumulative distribution function (ECDF) of covariates for specific treatment group a, denoted as (X,):A and the combined ECDF of covariates, denoted as {X; Specifically, the weighted energy distance is defined as:
2 nnn
E (Fn,,,,; Fn)
(A, =a) XX,12
waw I A, IX, X,
1
X: X,112
where Tha -E I(A, a). The weighted energy distance measures the distance between the weighted distribution of covariates among those with 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:
W E argmin
E Fn) s.t. Σ Ii (A: a.) n2
(27)
(M...... aEA
for all a E A and w; 2 0 for = 1,...,
Thus, the EBWs encourage distributions covariates for each treatment group to look like that of the target population, i.e. the full sample. Since (27) is 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 rely on specific model assumptions contributes to more robust ITR estimates.
C.2 Implementation of the PGD algorithm
We provide Algorithm I 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 Lyabal, we use the algorithm in [14].
Algorithm 1 Projected Gradient Descent Algorithm Estimate Decision Function for ITR-Learning
Input: Data set (x;, Qz U:) E Raxp x {1, K x R (D yi can be continuous or binary);
loss function L
Parameters: B E RPx(K-1) (model coefficient), Z 0 (Li-ball size); A 2 0 (additional N (number of iterations), az 2 0 (stepsize).
L2 uz-penalization for R(B)))
Constraints: 11-ball B = [B (Px((-1) || S =
Compute x == x,u?, for all i = 1.... n for faster computation.
Choose any point from Bo as initialization.
For = 1,2, 7 do:
B - 20-1-avVee VBC(B.))
B, + II
End for Output: By
To help the convenient implementation of the above PGD algorithm, provide the gradient of the ITR objective below. For continuous outcome, we have (see Proposition F.1)
K K - uxxuu" + B.
BB C(B) = Σ x()
x Bu,, u
(28)
For the binary outcome, we have (see Proposition F.2)
x Σ
exp(uz B B²x;) + [i+exp(u exp B"x:)
VB L(B)
w(a xx)
V/ x;u + A2B.
(29)
7
C.3 Additional improvements by dimension reduction: variable screening procedure
High dimensional covariates 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 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 sereening approach serves to reduce bias and enhance the efficiency of TTR-Learnigg
Several techniques have been developed choose appropriate variables in causal inference. One popular approach 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 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 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; where E {1.....P) p). conduct the distance covariance test between covariate and outcome K times, conditioning 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 Xj as the minimum p-value from these K tests. This approach ensures that we retain 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 significance level of 5%. We perform the same steps for all covariates in the dataset. These procedures are summarized in Algorithm 2.
each treatment
Algorithm 2 Variable Screening using Distance Covariance Input: X E R Y E R"; A {l...., K}.
Initialize X 0
for j = ..2.....7 do
for k 1,2,..., K do Perform distance covariance test(X,, Y|A = k) and obtain p-vuue,, Pj.k- end for Pj <- mmii(pp,,, Pj.2 .p..)....PPppkk if P5 < 0.05 then
end if end for Output: X. the set of screened variables containing effect modifiers and precision variables.
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 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 working model (1), the outcome can be explained by the sum of the treatment- effects 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
the augmented outcomes
Y y Y (xX 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 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-Leamin''s 'ss inverse variance weighting helps, providing

---

robustness against misspecification. SABD-Learning extends AD-Learning designed for homoscedas- ticity [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 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
a
Σ w(a, B"((x) - as n - oc in probability for all B E (Px(K-1),
(30)
We will justify this property Corollary C.3 toward the end of this section. Given (30), we have
n 71
((B) = (B)
w. w.u B h(x(),
(31)
where z here means asymptotic equality in probability as - oc. Denote the 'augmented objective function' on the right-hand side as laug (B).
In the remainder of this section, we introduce the optimal augmented function Aape (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 EBWs and MMD) balancing weights. Our approach extends the work of [50, 9] to more general setting, incorporating other types of covariate balancing weights that have asymptotic consistency (38) in multi-categorry treatment setting.
Let S(B be the derivative of the original objective function {(B) with respect to B. Then setting the optimal (XX is equivalent to minimizing the conditional variance of
B B S(B - w (A, X
Z
where B, is the minimizer of Caug (B), and B* is the minimizer of E[((B)). Using the property of asymptotic unbiasedness Theorem C.2, which is E|Z] = O. we need to minimize ||| with respect to h minimize the conditional variance of Z. Using the property of matrix trace,
h
E
Z - Vurl(rrZ'z 2E/w Z Zux) x).
and setting E
Z
0. Thus, hops (X) satisfies
E (A, ()) w(A, (x)
Since uAu = 1, we have
[(w(AA X x
Mope(x
(32)
Elw =
Now, we determine the functional form of hem (X) for both continuous and binary outcomes. For simplicity, let us define
w == w(, t; - BTx x;
R, (x) = E(w²u u XX =x], R2(x) = X =x] x],
E(w X = x].
7

---

Then for continuous outcomes, using (28) without regularization term for S(B) and (32), we have
laug (B)
K
2
W/
K V/
n1,
B²xx
u
T huppl (x:)
B
K
2
wiul T B 1
2K
Ro )
Wi
Vi
uz
BTxx
xxxx
(x:)
x,
K
1
K
K
2
L BxBBBRIxs Ri(x)))
2K 1 K
Σ
Wi
1/i
-
t
x
x
+ ti
Rxxxx))
61
K
2
WN
t K 1
81
Rxxx)) či
+ constant w.r.t B
assuming (x,)) / = 0. Therefore, we obtain the optimal augmented function with the minimum
estimation variance by properly defining Rx) Our simulation studies demonstrate that setting (X) = E X = x] leads to significant improvement.
Similarly, for binary outcomes, let us define
exp(u"B"xX) (x : E 2m + cxp( Then using (29) without regularization term for SBB) and (32), we have
x=
x
X x=x]
Q2(x) = E[w X =x].
Cang B
yyu T B x; + log(1 + exp(un 2 B xx)x)) uas T B² hopa (x:)]
- @2(x0) )
Σ
00
y;t + log(1 + exp((.)) + 2
1=1
1
Σ
[2)(x) - (2(x))) + log(1 + exp((.))
t4
y/
7
21
Therefore, by selecting proper (X) and Qx),., the minimizer of Laug ( has smaller variance than that of the original estimator.
Now we prove asymptotic unbiasedness of E wla; x, ), BT (x,)) First, we consider IPW- based conditions for multi-category treatments.
Lemma C.1. Let A = {1,..., K} be the possible treatment sets. Assume we have i.i.d. observations (xi, 04) from a population distribution of (X, A) such that TP( (A a x) is uniformly positive for all
a, X. Define
T P(A= M(x;)g(a)).
(33)
1m
If g is any function with
= O, then for any function h(() Rp + RP, then T + 0 as
n. - oo in probability.
Proof. First note that by iterated expectation,
(A))
((X)9(A)
T E
Ex
E
x
(34)
P(A
((X)g(1)
Ex
P(A x|X)
(35)
P(A = x|X)
JEA
Ex
((X)
Σ g(6)
(36)
jEA
0.
(37)
8

---

Then by the standard weak law of large numbers for i..d. samples, the empirical mean T converges in probability to its population mean [T]
+ oc.
Note that in the statement of the Lemma C.1, the choice of such g(i) can be each treatment vertex under the AD-Learnnn [39]. Now we replace IPW P(A = a; (x; with covariate balancing weights w(a;, xx) that has the following asymptotic consistency property:
lim
1
Σ t=1
w(a;,x,)
=0
(38)
P(a. (x.)
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..dd observations (x,, (4) from population distribution of (X, A). If 9 A I is any bounded function with ΣjeA9(i) = 0, then for any bounded funncin. h(x): : IP + R2, define
To Σ w(as x; xx))((x))(aa))
(39)
where the weights w(as, x:) satisfy (38). Then To - 0 as n - oc in probability. That is, To 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 To converges in probability to its population mean E[To] E[w((A A, (V(())((A)) By continuous mapping theorem, |Toll converges to [[] as n + oc probability. Then we have
((X)g(A)
||] [[ ())
E Www (A, (X)))(())())))))))))) E
P(A[X)
1
E
w(A,X)
h {())(())
P(A(X)
E
1: (1,1X) X) P(A[X)
h(X)((A)
1
< CCE
((A,
P(AX)
2
lim C 10
w(a,, x.)
- 0,
P(aa x;a]);)
where C is any constant such that | ( < C. For (a), we use Lemma C.I; for (b) and (c).
we use Jensen's ' inequality and the law large numbers, respectively. The last term converges to due to the property (38). Lastly, since || I2 converges to 0 in probability, we have To converges to probability by continuous mapping theorem.
The following corollary is useful in deriving outcome augmentation under the AD-Learning frame- work.
Corollary C.3. Assume the AD-learning setting. For any function h(X) depending only the covariate X, and for each parameter matrix B E TRx denoting,
Ho( w/ A, BB
have E[Ha 0. Furthermore, assuming the weights w(A, X) satisfying the consistency property (38), in probability as n
xc,
w(a,
x. B h(xx)) +
(40)
7
Proof. Follows immediately from Theorem C.2 by noting that ΣX_ u. = 0.

---

D Spectral characterization of optimal covariate balancing weight
Recall the definition of the weighted design matrix:
wi Ua, u (x;x?)
(41)
This is the p(K -) p(K design matrix encoding the sample of patients with covariate x; and observed treatment a; for i i= 1....,,.. n. We are free to choose the weight wy = w(a xi) for the ith patient function long they sum to one. What is the best way to choose such weights?
Our main results (Theorems 3.5, 3.5) suggest that is best to choose them in way that the minimum eigenvalue of y is large as possible.
To illustrate this point, suppose for simplicity the covariates are categorical and assume values in finite set X. We decompose the summation in (41) according the treatment and covariate as follows:
Σ
w(a, x) (ua uz & (X:X)
(a: a, x; x)
(42)
n1 (axx)xAxX
(a.x)
=:n(a.))
Σ
w(a,x) nlaa x) x(a,x).
(43)
(xx
Note that na, x) is the number of patients with covariate x and treatment a, while {(a, x) is the
p(K 1) x K 1) covariance matrix representing the treatment-covariate pair (a, x).
1. Under IPW w(a, x) F[ax} since nı(a, x) Binommn, <aaxx) (under suitable assumption), using law of large numbers and central limit theorem, we have
w(a, x) na x) 1 Or(n)1/2)
Consequently, the covariance matrix E(a, x) of every treatment-covariite pair (a, x) con- tributes nearly equally to the weighted design matrix v.
If distinct covariates are encoded as one-hot encoding vectors, then IPW maximizes the minimum eigenvalue of v. To see this, for simplicity, assume we have binary treatments
(K = 2 and #1 = 1, = -1) and two discrete covariates (e.g.. 'Male' and 'Female") encoded as one-hot vectors (1, 0) and (0, 1)T. Then can write as
!1
1 0
n (M((,, M)
(2,, M (2, M)
0 0
0
0]
[0
0
+ ww(1, F),,, F)
+ F), 1*)
1
1
[w1,, M) + (2, (2, w(1, F)n(1, F) + w(2, F)m(2, F)
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 zz)n(a, æ) becomes a constant, and V above becomes the identity matrix, achieving the largest possible minimum eigenvalue. Also, in this case, & has condition number so the ITR optimization landscape becomes well-conditioned.
2. Using the fact that the minimum eigenvalue is a positive homogeneous and concave function on the space of symmetric matrices, we have
x) [a, Σ Amin Σ w(a,x)
Amin y
n
(a)
10

---

Here [a) is the aggregated weighted covariance matrix for sub population with treatment
a. In the above derivation, we used that >., Ew(a) by (42) and definition and the following computation using Jensen's inequality:
Amain Emo(a)
= Amin
Σ A Ew((a)
A
Amin En (a)) = =E (E (a)).
If we choose the weights that these matrices are identical, then each treatment group contributes equally to the lower bound above for Amin (V.) This condition is satisfied, for example, when 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 μa, for the treatment groups. To maximize the lower bound for Amin ( should be chosen to maximize Amin (w (-)) The specific solution to this optimization depends on the structure of L. In three-way balancing methods [8], ft 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 for ITR-LLarning problem (6). Below state the full version of Theorem 3.3 in the main text.
Theorem E.1 (Convergence rate of PGD for ITR- TRR-Learning, Let (Bi)120 denote the sequence of parameters obtained by the PGD algorithm (7) for ITR- TR--Learnig problem (6) with arbitrary initialization Bg- 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 p-strongly convex and L-smooth, where μs I. are as in (10).
(ii) (Linear convergence with fixed stepsize) Assume constant stepsize ot = contraction constant p(a) == max |1 - aL), | apl E (0, Then for all t 2 1,
2/ <
Denote the
B: - B1 < 4p(a)'
(44)
In particular, if ++L then p(a)
I-= L+ッ 1).
(iii) (Sublinear convergence with diminishing stepsize) Assume diminishing stepsize a p p= 24L and 1/p. 7 > O are constants. Then for all 2 1, I+L
3 7+0'
where
BB. B 7+t'
(45)
where V = max
(1 MI
4(7 + 1)x23.
In particular, the assertion holds for 11 4(1 + 1)
Ap-1
8 = 2/p=
H+L and v = max HL
The proof of Theorem E. relies on two ingredients: general result on the convergence of PGD (Lemma E.2) and local landscape analysis of the ITR objective (Lemma E.3).
We first establish a general result on the convergence rate PGD to find the minimizer of a strongly convex and smooth objective ʃ within 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 V f [BB) [[ 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 (P + IR be L-smooth and ju-strongly convex function for some p., L > 0. Suppose f = BE ((B) > -x and denote B* := arg minBes f(B), where B is convex constraint set of (o)(K-1), Consider the following PGD) iterafes:
B 1g (B, a. (B..)
(46)
Then the following hold:
(1) (Linear convergence with fixed stepsize) Suppose f is twice differentiable and o: = < 2/T. Let contraction constant p(ax) == max aL), | anl E (0, 1). Then for all 2
[B -B.-B'lF
B+ S p(a)
B²
(47)
The bound p(a) on the linear convergence rate is minimized at a μ+L with minimum L-n
-2
value
L+p
2uL (ii) (Sublinear convergence with diminishing stepsize) Denote p:= u+I Suppose at = 747 where
3 1/p. > O are constants. Thenffor all t 2 1,
| B
(48)
7 +1
"V 3p-1
(7 + 1) Bi B 11 }. In particular, the assertion holds
where V = max
for 2/p = L+L and max pL
2(u+2)² (BB 1|
(7 + 1) 1|
12

---

Proof. Let B, := B, - 0, (B.. Then
[B. - B² = |B: B. + B. + - -B,, B, B--B²)
= a V
+ B. BB 2a (B, - B", VB.((Bi)
< Vn. + [B, - B'I 2a,, (B, - BB, VB, Vn/(B")
using that (B, of p-strongly convex and smoohhh functions and, we have
since B* is stationary point of over B. Using the co-coercivity
|B, B'IF x f II
+ | - BI 1BB
μL
2a,
BB - -
vf
1 W L
P L
20
20:
[B, 1B.-B"Ie
+ + V
[V. V
#+1 1
By Young's s inequality,
s V
ff +
Thus the above yields, given that at S 1/(u + L).
20,pl.
# H+L ) [VS(B.) V)(B))|
[B - B 13
[B. B" 1/ + 204
04
M L
20yu
[B, B + \
1 T.
Since the projection onto the convex constraint set B, 1lg, is non-expansive,
|Bi+1 - 1 < B, B" | < (1 - a.p) [B, - B² + ajo?
2ut. where we have denoted p a+t and a² V IB
Now, we show (i) (47). Suppose = a > 0. Note that B' is a unique minimizer of f over B because of strong convexity. Since (B+) lies in the normal cone of B at B*, we have B" = lB a (B*)). Using non-expansiveness of convex projection Ilg,
|B,+1 B* ||
||B aVV((B)) IIx(B* a
[B aVV( - (B" - (B*)(=))
1/'0 (I (B, (B, - B*) ds
supls - /V² B, - B
B
S max - a1 apl} [B, - B*I
=p(a)
For the last inequality, since the eigenvalues of V are contained in the interval ju, L}. the eigenvalues of / a f are between min(1 - aL, I - au) and max(1 - aL, 1 - ap). Note that p(a:) is minimized at a = + with minimum value L- L+μ follows that for any t 2 0,
BB+1 B'IlF I p(a)*+1 ++1 |Bo B*IIF.
Next, we show (ii) (48) by induction. Indeed, it holds for l 1 by the choice of v. Denoting 0²0²
i : y + 6, by the induction hypothesis and by the choices of a; = 3/8 and v > B
B B <
3p
V
f202
f2
8
1
Bp
Po²
12
50
f-1
V
11
11 1)( +1)+
7+t+1
13

---

This shows (ii).
Next, we establish the following global landscape result of the ITR learning problem (6) by computing the Hessian of the objective (B)) in (6).
Lemma E.3 (Local landscape of the ITR objective). Let LB) denote the ITR objective (6).
Suppose Assumption 3.1 holds. For the continuous outcome,
(X + Pc) vec voc
3 (X+ + TcK-
(49)
where the constants X are defined Assumption 3.1. (ii) Further Assumptions 3.1 and 3.2 hold. For the binary outcome,
(a->> + ( I V Ve V vec L(B) 3 (a+ + [
(50)
where the constants x= and a are defined in Assumptions 3.1 and 3.2, respectively.
To maintain the flow, write down the proof of Lemma E.3 Section F. Now we can easily deduce Theorem E.I.
Proof of Theorem E.I. Lemma E.3 shows that the eigenvalues of the Hessian voc V vee 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 < μ & < co. This shows that L is p-strongly convex and L-smooth over B3 as claimed in Theorem E. (i).
Next, noting that ||-1| < Il-I|1, triangle inequality and the fact that B,, B E B yields that, for any t 0.
[B, - B < B, [B,ll)²²
(51)
< B, [B]
(52)
S 4X7-
(53)
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. will first compute the gradient and the Hessian ITR objective in (6) separately for the continuous and binary outcome. We will then prove Lemma E.3 at the end of the section.
F1 Continuous outcome
Proposition F.1. Let L(B) denote the ITR objective in (6) with continuous outcome. Then
VB C(B) 1 Σ w(a;, x:) (xix Bua, u K Vxx.u% + A2B, (54)
vec(B) V Vec C(B) = 1 + X2 II(
(55)
where W is the weighted design matrix defined at (8).
Proof. Recall that the ITR objective for continuous outcome is given by
K
A |
L(B)
Σ
w(a,, x;)
u
T B
B
+
2n 1=1
K
-
2
14

---

It is trivial to show (54). For the Hessian, using (16),
VB [() = V (B) V Bec [(vec(B))
r K - vixuu.
+v )
BB) vec
Σ x(),
x;)
K xx-; Bu, u
n
VB) vec
Σ w xi)
(xix
Buo Bua un:))
1o Ip(--1)
w(a;,x) ((mm.))) (xxx)) voo(B)
17
vec
X2 I)(K-1)
w(a,, xx) vea(B)7 (ua ( - (xxx)) )
|
vB)
X2 Ip(K-1)
(u n² & (xxx2)) X2 dzlp(k-1) (um
w(a;, xx)
Σ w(a: x, (ue, uZ,) @ (x,x²) + de I)(K-1)
+ X2 (K-
This shows the assertion.
F22 Binary outcome
Given the current model parameter E RPx((-1), define the scalars 2;, Pi, wi for the i-th sample by
exp(2+)
2,
u x.
Ps
:; w(a,, xz).
(56)
+ exp(2;))
For each binary label y {0, 1} and activation R. recall the negative log- ligelike of observing y under the logistic model Ber (p:) is given by
2) = log(1 + exp(2)) y2.
An easy computation shows
exp(2)
h( =) =
0z e(y 2) 1 +xp(()) - y,
(57)
h 2) = 022 llu. 2) = (1 (1+xxp(2)
8²
exp(=)
exp(2)) ) (1- exp(=) 1
(58)
=xx(((2
1 exp(2)
4
For the forthcoming computations, define matrices
K = (1111,1), h(yn E Rxxn
(59)
M == diag [h(y) 21).... h(yn, zn) E RRxx
(60)
$ [Ma] XI,...,, Un. 8 xn) RR(K
(61)
W = diag( (W...... ..., wx)) E R=
(62)
Proposition F.2. Let C(B) denote the ITR objective in (6) with binary outcome. Suppose Assumption 3.1 holds. Let 2; = u2 B x, for i I,..., n.. Then
VB L(B)
Σ
wi h( 2) x)
+ A2B,
(63)
(B V LBB()
((MW) T >2Tp(K-1)
(64)
n
15

---

Proof. We first claim that
V Vec(BB22 = ((K-1Pp) (x: 8 Ua.).
(65)
V (B) (Vi. 2;) C(K-1P)) (x: @ un.) h();)
(66)
Vee(B) vec vec(K)*= Vvec(B) K 4M.
(67)
Observe that by using (18), we can write
2, vco(2;) = voc(u B x;) Bxxx)= Noting that vec(B")) = (C(p.K-D 1) vec (B)* = vec [B) C(K-111 Hence we get
(xF @ u vcc(B3) = voo(B")" (x; @ ua.).
v)) Ver vee """ (x; ua.)
= =C(K-1,) (x, @ ua.).
This verifies (65). Then by using the chain rule (20); get
vee (y): 2;) = (Bc V f(y 2i)
C (K- (x; @ Un: h zi),
which verifies (66).
Next, we compute the gradients of K in (63). First, using (65) and the chain rule (20),
V Veo h( 2,) = V vB 2 V., h(y, 2i)
1,p) (x, Ua, (i( 2,)
xx) 1 (yt, 22), (u,,
where the last equality uses commutation matrix relation (17). Then noting that vee(K)² K
[h(y), =1),.., (yy 2n)]. follows that
V vec(B) vec(K) T (e) (uay @ x1 xx)di((1. (ua, @ xx )n))i(y 2n)
(6) Uu x1)..... ..... (ua, & x,,]] diag (i(vı, 2...... li(ynn a.)))
M
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:
B))
wi Elyi, zx)
Wi
& un;) higy, 25)
(68)
n
Y1
(6)
1
Wi Wua ? x;) h1)),
(69)
21(1),21)
[w] (May @ ()))...... un (nnu xn)
(70)
Ih(yn 2n)
1
FW vec(K),
(71)
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 scalar, we deduce (63).
Next, we compute the Hessian. From (71),
V C(B) + W vee(K) + de vec(B).
(72)
16

---

Then using (66) with (72), we get
vec V
ven L(B)
V (B)
(ec(K) W + (2)
(73)
#(MW) T" A2Ip(K-1)
(74)
n
This shows the assertion.
F.3 Proof of Lemma E.3
are now ready to give proof of Lemma E.3.
Proof of Lemma E.3. Recall that +
Lua @ X1
Uo @ xn). Note that
WW Φ T
wy (ua; @ (u,, @ xx,)
(75)
wi (Ua ug @ x;xx²)
(76)
(77)
where & is the weighted design matrix defined at (8). (In particular, since W is diagonal matrix of nonnegative entries, this shows that v positive cemi-definite)) According to Propositions F.1 and F.2, we have
n1 " + p(K n #MW4" >2II((K1)) for binary outcome.
for continuous outcome
(B) V (B) L(B)
(78)
Thus the assertion for the continuous outcome follows immediately. For the case of binary outcome, by Assumption 3.2 we have
a' In 3 M 1 I
(79)
Then
1 (owwr)
1 (MW) zat (=ow0")
a
(80)
follows that
(" @ww>") [)(K)) ((MW) "" 3 a ( ('wwo") ((--1)
Amin
(81)
This shows the assertion for the binary outcome holds.
17

---

G A non-asymptotic consistency of weighted, constrained, and regularized MLE
In this section, provide general result on the non-asymplotic consistency of MLE in general setting. Here, observations are assumed to be ii.d. from a generative model, and the unknown true parameter of the generative model may lie on the boundary of the parameter space.
generic sample (X, Y) consists of random covariate X and outcome Y. (In the ITR setting. we will specialize X to be the pair of the patient's s covariate and the corresponding treatment.) We assume the covariate
follows a probability distribution ju()) and the conditional distribution of Y given X x is parameterized as mo, (.1z;) for an unknown parameter O,. Hence the generic random sample Y) follows the parameterized distribution Tq. (y x))(()) Assume we know smooth parametric family of conditional distributions 0 + TT I)) where the true parameter 0, lies in a convex parameter space A C Rd. Our goal is to estimate 0, from i.i.d. samples (I;, Us). i ........ drawn from veo .))(-))-
To estimate 0., a natural approach is to maximize the empirical joint likelihood function. Since the marginal distribution of the covariate independent of the parameter 0, we can apply the method of "reweighting" the covariate distribution by introducing an additional random weight w(x) that we can control. Namely, fix weighting function w(-) X' + [0, oo). where X' is covariate space. that assigns a nonnegative weight to each covariate value x. Consider then the following weighted expected likelihod maximization problem:
arg min (x.)) (-w(X) log mo. (YI
(82)
o
where we omit the marginal distribution (xX) as it is independent of 0.
Recall that we get design the weighting function w(-). We first observe that the true parameter is a stationary point of the weighted expected log- likeliood function E' [-w((X) log тe. (Y [X] for arbitrary weighting function w. This can be justified under the mild assumption of exchangeability of the expectation and derivatives (which holds under twice continuous differentiability of the log-likes likeeloo function by the dominated convergence theorem). Indeed, note that
VoTo (X)
Exx,y (x) log Na (Y (Y| X)] = E(x,
(XX)
TY X)
(X) V X)
= Ex By
To (X)
Vono XX
Ex w(X) EX (.12)
X
0,
TY X)
where the last equality follows from
Vono, (YlX) To. (Y ()
VOTO - (ylX) xy
By (-1x)
X xx
Vo no. (v) X)
Vol O.
Therefore, 0, a critical point of IE [X) log (Y X)]. In particular, if we assume that the liglliill function log No (y <))( is convex in 0, then 0, is global maximizer of (X( log (Y X)].
Next, we introduce the finite-sample approximation of (82). Denote w, := w(x;) for the weight of the ith sample (I,, yi). Then the weighted negative log-likelihood the observed samples under the model parameter 0 is
(( : 1
log To (v: xx) + R(O).
(83)
C:Co(
where R(0) suitable choice of regularizer for parameter O. The regularization term can be considered prior knowledge about the model parameter 0.
18

---

Let O, denote a (possibly non-unique) minimizer of the above function over O This is minimizer of the loss function L over the constraint set O which we call the weighted, constrained, and regularized MLE of 0..
As in the standard MLE analysis, we assume that the negative log- nigllikelihood function 0 - Lo (0) is differentiable and strictly convex for all x E IRP. Statistical literature typically assumes that the true parameter lies in the interior of the constrained parameter space, O a convex subset of IR In contrast, we allow the true parameter 0, to lie on the boundary of O Regardless of that, we have seen above that the gradient of the weighted expected log- oog-lielihood function vanishes at the true parameter 0..
In this setting, we aim provide high-probabillty guarantee that there exists global minimizer of (83) that is close to the true parameter 0. for an arbitrary weighting function w(+). which we will optimize later for the best statistical guarantee. In Theorem G.1 below, we obtain such result in the non-asymptotic, weighted, constrained, and regularized setting. For its proof, we combine a classical approach from [15] with concentration incqualities, namely, classical Berry-Esseen bound for deviations from the standard normal distribution for independent but non-identically distributed random variables and 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,, (0) in (96)), while the McDirmid bound controls the quadratic term (S,, (0) (96)). By using an E-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). Con- sider the constrained, regularized, and weighted MIE problem (83) with unknown parameter O from convex subset Q Rd Fix an arbitrary weighting function w(+) Assume the following holds:
(a1) (Smoothhess) For each sample (xr, y). the per-sample negative log-likelihood function 0 + log T z))((z) is strictly convex, three-times continuously differentiable, and R(0) is differentiable. Furthermore, denote U == w wX))
V a log ne. (Y (X) E I2d and U U
EEUU. Suppose there are constants D1, d1 E (0, 20) such that for all i,
B < D1,
min Var(U():) > d1 IskSd
(84)
(a2) (First-order optimality) The true parameter 0, stationary point of the expected negative weighted bogliikel function Lo(O) == Es. [-w(X) log To (Y [X]] over O
(Vo (O.),, 0 - 0.) 2 0 V ①.
(85)
(a3) (Approximate second-order optimality) Let Σ(0) = (0) + R(0) denote the expected regu- larized negative log- likelihood function. Then the regularized Fisher information V ((0) is positive definite 0 O, with minimum eigenvalue p >>>>>>>> O.
Fix a constant C> O. Let = Cn-11/2 + & and M = M(D) > denote the supremum of
the absolute values of all third-order partial derivatives of L over all 0 with 0 0.1l < D. Suppose || R(O.)] is small enough so that
3p D D
(86)
4M(D)
Then there are constants C1, 022 C3 > O such that
P 2100- arg min ((0) SD) 2 11-0 e1 -xeeep exp
C2p²
C2 Σi=
w
ece
2
(Ei w²) 02)337
3/2
(87)
That is, with high probability explicitly depending on C, p. d, and n, there exists a global mininizer of 0 + L(0) in e within distance D from o..
It is important to notice that the minimum eigenvalue of the weighted regularized Fisher information (²Z( depends on the choice of weighting function w(-). Therefore, Theorem G. 1 suggests that
19

---

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 s incquality). Let X....., Xn be independent random vectors in IPP from a joint distribution л. Fix a compact parameter space O S IR and fe R + [-J, J] is bounded functional for each 0 E - such that
Il 11 S L 0'1l
40, E
(88)
for some constant L > O. Further assume that [Xx) = 0 for all 0 E O and k = ,,...,
Then there exist constant J such that for each n 2 0, and rj > 0,
12 Ja (Xx)
2L (i
n 2.J²
P
sup OCO
2n
<
exp
(89)
71
Proof. Since C C Rdd is compact, it can be covered by a finite number of L2-balls of any given radius 0. Let Ug be such open cover using the least number of balls of radius 0. Let (E) 24 denote the least number of such balls to cover O. Moreover, let diam )) denote the diameter of θ, which finite since O is compact. Then θ is contained in a d-dimensional box of side length diam @). This box can be covered by (diam (=)=)* cubes of side length e. Covering
each such cube of side length e by a ball of radius e, follows that
N(c) I (diam ()) d
(90)
Next, fix 1) > 0, θ, and > 0. Let 0,. ... x(()) be the centers of balls in the open cover
Ug. Then there exists 1 < j < (s) such that 110 - 0,ll < s. By the hypothesis, Se depends on uniformly continuously with respect the supremum norm. Hence there exists ô = -(e) > 0 such that
fl5o - s0, I Le.
(91)
Denote (0) == n n fo(Xx). Then this yields, almost surely,
(H, H. (0,) S LE.
(92)
Furthermore, by the assumption, || uniformly bounded by J > 0. It follows that for each O. H₁, changes its value at most by M when one of X Xx, is replaced arbitrarily. Therefore, by the standard McDirmid's inequality (see Theorem 2.9.1. in [51]) and a union bound, with choosing e n/(22),
(w//2L)
22 diam (2Ldiam
())
n7'n 2.
P (IH,(() n)
((H),() 2 n/2) < K
exp
(93)
71
The above holds for n 2 and n > 0. This shows the assertion.
Next, we recall the classical Berry-Esseen theorem for the rate of convergence of normal approxima- tion for the sum of independent random variables due to Feller.
Theorem G.3 (Berry-Esseed Theorem in Ch. XVL.5 of Feller '91[17]). Let X1, X2..... Xn, be independent and not necessarily identically distributed random variables with zero means and finite variances. Define W = Σ X, let F be the distribution function of W and & be the standard normal distribution function. Then
finctio
6 DE
1
|F
$
(94)
E 3/2
Now prove Theorem G.I.
20

---

Proof of Theorem G.I. Suppose we have n i..d. observed samples (Zz, yi) for i - i l........ n from distribution no. (y p))(2)). Also by the assumption, Lo is twice continuously differentiable and n "'EE== so E V Co) [B[Co] and [o] = { by the dominated convergence
theorem.
Note that, for any > 0,
{
{
0. arg min (()) sr]
inf 110-0.1==r
- (10)-((0.) >0 =0)
(95)
The inclusion S is clear. Conversely, since C is strictly convex, the event on the right implies that its unique global minimizer is within from 0., as desired. Thus in order to show the main result in (87), it suffices to show that the probability of the event the right with D Cn least the right-hand side of (87).
8|| (())
Fix θ such that ||0 0+ = We introduce two random variables that we will bound to be
small by using some concentration inequalities:
Vni
(0)
(V.Lo(O.) - (x) 0x).
(96)
0 0.1
(0)
0 0.1 0 0.)² (VoVer C(0.) - voVor (B[(() (0 -O)) (97)
Roughly speaking, these quantities are of order with high probability, T, by Central Limit Theorem and Sn 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 + (() is assumed to be three-time continuously differentiable, the quantity M 2 in the assertion is well-defined and is finite. Then, using the Taylor expansion, may write
(()) - ((0.) 2 (VoC(o.), 0 0.) -0 -00) -V. C(O.) - 0.)
0-0.)
M(1l0 0.
10 0.11.
(98)
6
We will lower bound the first two terms on the right-hand side above. Note that
(VoL(O.), 0 0.) (VoC (0), 0 - 0.) - -Ei(()) 0 0.)]
+ (ToE[(())
0.) (VR().). 0.)
V (o.) 0 - 0.) ([L(X) 0 0.)
|R(0.)) 1 0.10.
0.1 (, [| 1--0111
where for (a) we use the fact that 0, is a stationary point of [() over and Cauchy-Schward inequality; for (b) we use the definition of (0).
Next, we turn our attention to the second-order term in the Taylor expansion (98). Recall that from the assumption (a3) in Theorem G.I.
C(O.) = VoVer ((o_)))))))ANT((((())(())1).))) 2 plp:
(99)
where p > a constant. It follows that
(0 - 0.) VoVor 0.)
(100)
2 (0 0.)² ((..) VOV.r (El=(x))((0)1100000)))) 0.
(101)
z ||0 110-0,11 (s,
(102)
21

---

Combining the above inequalities with noting that ||0 0.112 = D2, we obtain
((0) ((0.)
0 0+11
1
(Vo((O.); 0 -0.) + 2 ( Vo0 0+)
00 0.12
M(1 0.1l)
6 100 0
(T( |VR((+))|
M( 0x1l)
2 (5, + +
|0 0+1
00 0-1l
VT
1
-||VR(00)] P
M(D) D
(m) P)
VnD (0)
10 0.11
=1h
=112
Note that 2 if
2 8 -22IVR and L
M(D) D
8 6
The choice of holds the former condition, and the latter by the assumption (86). Thus, J 2
We now take infimum over all 0 θ such that |00 0.1 D. It suffices to show that the random variable l2 defined above positive with high probability, since I does not depend on 0. To this end, write
1 (%) +
2 2
inf 0:0 || 10-0.1I=
12 2
inf Oce
VnD Tn(
inf dEB 10-0.11=D
(103)
2
|10 II
=:B
Then the last expression in (103) least p/ 1 if A 2 -p/16 and B 2 p/8. By the hypothesis, D = O(1) so is uniformly bounded. Then by the uniform McDirmid's inequality in Lemma G.2, there exist constants C² C" such that
P B < 2)= =P
inf ( 8=0
p
DCC DD'C'exp(-C" exp(-C"nn).
(104)
|10
0.1
D
Next, we will show the following incqualities: For K = 6D1 2/22 /d//d?/?
VnDp
P A
P
inf |00 0
T.(0)) >
16
16
=D
VnDp'
K Σ
i=1 W
d
IP
Z
32vd
Vn ( i=1 w 3
(d)
32vd
nD?pp²
K Σm= wp?
CPv2T
exp
2.32dd
Vn トス w 3/2
W
32Vd.
C2p²
K 1ー wp
d
2m
exp
d2
vn (2n 12 w 3/2
(105)
where ((0 1) an independent standard normal random variable. (d) above is a simple consequence of the standard Gaussian tail bound (O, 1) > x) مجج =v2x and (d), (c) are from the
choice of D which yields n 2 C We will show (c) the end of this proof. Note that for any two events Ej, Eg defined on the same probability space, PEE,, n E2) > P(E)) + P(E)) 1. Then
22

---

by combining (103), (104), and (105), we have
[( - C(O.) C(O.)
P
inf
BEE
|0 0.11
110-0. I=D
P
+PBB2
P
P A N
16
8)
-1
32vd
C²p²
K n wp
2 1 d
2T
exp
2. 22322d
+ w,?)
3/2 3/?3
+ +(1- DDC exp(p
= 1 - Cj exp
C²p²
02
w
O(exp(-can)),
2 32
ii=
with the corresponding constants C1, C2, C3 > O. Then the assertion (87) will follow using (95).
remains verify (c) (105). To this end, write
n w;Ui:
(O.) [[](())VCo n
where U: w( (X:) log ne, (Y, (,) Rd and U. U; EEU,]] in the statement. Then by
using Cauchy-Schward inequality and noting that |00 - 0.1 = D, we get
0 D
(,( =
U.
U
w (k)
1=1
A=
important to note that the distribution of the random variable on the last term above does not depend on O. Note that w,U,, for i I,..., .,, n are independent mean zero (but possibly non-identical distributions due to the weights w;) random vectors in RR and by the hypothesis, their respective
coordinates have uniformly bounded variances. Hence by a union bound.
inf oc DD
(0) 2t
S
W (K)
2
(106)
Vd
A=1
Let th = v Em we (k.). Then by the Berry-Esseen Theorem (Theorem G.3) and the hypothesis, for 7 N(0 1),
6 DE U.
P (Q) < - P ( s 2
sup zER
ΣL w Var U (k:)) 3/2
6E || Σn
1- w
(Var (0, (a(())))) 3/2
(E w 3/2
6D1 Σn wp
3/2 (ELL wz)
K Σ w
for k = I,...,d
EE w 3/2
3/2 with K =6Di/dd Combining with (106) and using triangle inequality, we obtain
n 1 ux K
inf d
(0)
2
<
2 ZZ
+
(107)
2vd
E w:? 3/2 3/2
1 D
Thus (c) in (105) follows.
23

---

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 proof are the computational guarantee (Theorem. E.1) and non-asymptotic estimation guarantee of constrained and regularized MLE (Theorem. G.I).
As we have briefly explained in the main text, the connection between the weighted convex optimiza- tion (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, a4) follows joint distribution that does not depend on the true parameter B. Further, assume that т admits a density function fx. Then, for the continuous setting, the joint ligeliill of observing the triple (x, a, Ut) gives under the generative model with parameter B
K K
u
BT X
+
log fx (xi, a:) + constant.
20 1=1
==1
Thus, the standard (unweighted) maximum likelihood estimation of the true parameter B, can be written as
K 1-1
2 uz B x1
min (E 2. B
where we imposed additional Li-ball constraint for the parameter B. Now, if we use covariate- balancing weight w((x, a.) for each subject i, the corresponding weighted MLE problem reads as
n w(x at)
K K 1-
2 x
B u
(108)
BER
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
Σ (- (1 + exp(u? BT x:)) + you BTx;
+
log (x;;
Proof of Theorem 3.5. Consider the generative ITR model (11) with true parameter B.. Fix 0. Let (Bz)220 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
[B. Bylle S B
B Ft
B
B.
=:
72
By Theorem E.I, we have
4p(a)
with constant stepsize with diminishing stepsize.
I 5
J+I
We wish to apply general result weighted and constrained MLE in (Theorem G.1) with Le-regularization R(B) A [B ? in order get a high-probability bound on the term 1g. For this, we first verify the hypothesis in Theorem G. for our generative ITR model (11), Assumption (al) in Theorem G.I follows from Assumption 3.4. For Assumption (a3) in Theorem G. denoting by p the minimum eigenvalue of the expected regularized Fisher information, by Lemma E.3 (also see (10)),
X + X2
for continuous outcome
== μ
a ---
for binary outcome.
where A_ is the minimum eigenvalue of the expected weighted design matrix [[ in (8) (see Assumption 3.1). The above quantity is positive under the hypothesis.

---

Lastly, for Assumption (a2) in Theorem G.I. we need to show that the true parameter B. is a global maximizer of the expected log-like reeliikli function our generative ITR model (11). This can be shown by 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 PB (X, A, Y) PB(Y X, A) p(X, A), then log PB (X, A, Y) is concave function in B. Thus we only need to show that the true parameter B, is a critical point of the expected weighted log- ligelikl function E (s, X) log PB (X, A. Y)). Indeed,
Vn (A, X) log pm (X, A, Y)]
= E [A,, X V (log PB (Y) X, A) log B log PB (Y] X,X) A)]
E [A, X
V BPB Y (X, A)
W(A, X)
pB (Y| A)
(Y A)
Ex
((Y,A))
(A,
X
PB X, A)
BPB (Y|X,
Ex
(a, X))(a, X) EyppB, PPB, (11x,a)
X, A,A= a
PB (Y) x,
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
BBPB. (Y|X, A) PB. (Y A)
Ey (11..))
X,A = a = V VBPB. A)
(109)
VB / PB. (Y Ay
(110)
VBl 0.
(111)
Thus B, is a critical point of E [log PB (X, A, Y)]. By the concavity of the likeliidddood function, follows that B, is global maximizer [A,, log pn (X, A, Y)]. In particular, this verifies the assumption G.1 (a2).
Now we can apply Theorem G.I with Ly-regularizztion R(B) = A2 B for our generative ITR model. Denoting
8|VR((.))lF
822 B.
D = nn
+
+
(112)
vn
for C' > 0 any constant, this gives
C²
22 Σn -i=
wi
IP (12 > D) S c1 exp
+ +(exp(-con))
32
EE ui? 3/2 3/23/22
where C1, C2 > are constants and d = p(K 1). By Assumption 3.4, the quantities
vn Ei= / (E w,?) 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 (113) are at most e/2, for which R E 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 e/2. For this, we choose
32 2d
201 log
Vploge
C =
Accordingly, we choose T large enough that C/VTn. Then the above gives
20
822
|
[B. - Brlle
+
< P P12 D)
(113)
Vn
This shows the assertion.
25

---

I Additional simulation details and results
The code supporting this study is available at http:///ttttb..ooo/yyy/fffffffff with plans for release as an R package soon.
L1 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 I with heteroscedastic random noise e ~ (0, (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 our randomized study aligns with Scenario 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:
((X)
2X 2X2,
0.75 + + 155
0.75 + 1.5x, - - 15x2-115.
A = 2;
6(X)
0.75 1.5Xi 1.5X, 115x, 1.5x4,
0.75 1.5X, + +15X, 1.5x, 1.5.4 =4,
Observational study:
x(X) = 1 X5 + 3X6 + 2X,X2;
1+ X X3, 8(X) 11.5+ +1.5+3X1 3x, - + A= 1- - X2 X3, = 4,
+ +
X3, A 1=
Case 2: Linear interaction functions with complicated treatment-free effect functions
For Case 2, we consider the following treatment-free free effect function for model assumptions of Y:
(XX + X5 + X3 2c X2 sin(X3).
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:
(XX) 1 2X1 2X, + 2X, 2X3 22,
0.0%1100X, 2.0X. 0.5X2, == 1.0 1.0,, 1.0X, - 1.0X3, 22 8(X) [1.5 2.0X; 1.0X4 1.0X3, A = 3; 1.0 - 1.0X; - 1.0X, - -1003 ,A A = 4.
Case 4: Non-linear interaction functions with complicated treatment-free effect function
For Case 4, we consider the following complicated treatment- effect function for model assump- tions of Y:
M X 1 2x, + 2X, + 2X, - 2X + 2X,,,, + 2e -x + sin(Xa).
In all four cases, adopt the same treatment assignment principles and variance function o² that induces heteroscedastic error for observational studies as outlined in [44]:
f0.25 I(X; 0) 0.4 I(X;
> 0 1= (0.25 I(X, <) 0.2 I(2 > 4- 2,3,4.
T (A,
(2 = 0.25 + I(X, 0) + I(X, 0,4 1) X4 I(X, > 0,4 2)
26

---

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-learnig,, For the heteroscedastic error for randomized trials, (A, X) = 0.25 + 0.2 (1.5 X2) was used.
1.2
Implementation details
All numerical experiments were performed on 2022 Macbook Air with MI 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 regglalariiition for the penalized approach (named "penalized" in the main text). The other uses the proposed PGD algorithm with L,-ball (named 'constrained" in the main text). To ensure fair comparison, both the penalized method and the constrained optimization algorithm were implemented using code developed by the author. We do not consider additional L2 lg-regularization this setting. To determine the regularization constant or Ly-ball size, we use mean squared error as the criterion. The iterate number I of the PGD algorithm is 1000.
In simulation study, employ the default configuration of the random forest algorithm as implemented in the uanomForret package in R to estimate the residual variance function in SABD- Learning, 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 's splitting criterion determined using the Gini index. The number of variables considered for splitting at each node is set the square root of the total number of covariates. For model performance evaluation, we utilize the out-of-hag (OOB) error estimate.
I.3 Additional simulation results
The goal of these additional simulations in the appendix 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 _es corresponds to proposed+ 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 gimnet in R.
in the main text. In this section, we include
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 perfor- mance. 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 misspeci- fication 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 relatively simple treatment-free effect, AD-Learnin's 's performance is comparable to SABD-Leaming, particularly following outcome augmentation. When combined with EBWs as AD _e_e SABD_e _e, these methods exhibit improved performance over their original versions (AD, SABD). trials, EBW Ws effectively reduce finite sample imbalance compared to IPW with true propensity scores. In observational studies, energy balancing yields more effective balancing weights, resulting higher accuracy.
randomized
27

---

Next, we delve into Case 2, involving a more complicated treatment-free free effect function compared Case (Tables 2. 6). Notably, since the treatment-free effect does not affect ITR learning, misspecified treatment-free effect function introduces heterogeneity and increases noise variance. Then we can see the synergistic effectiveness outcome augmentation and inverse variance weighting in SABD-Learning, effectively handling additional heterogeneity in ITR estimation (See "augmented" panel 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-
relationship, which leads to suboptimal results of the variants of AD- Learning as Case Our results demonstrate that the more robust EBWs outperform standard IPW, particularly in observational studies. Additionally, variable screening enhances ITR learning in high- dimensional setting. Employing variable screening (SABD_s 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 model misspecification highly nonlinear outcomes. Simultaneous use of EBWs and screening for SABD (SABD e 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.
28

---

Table 2:
Average accuracy and their standard deviations (in parenthesis) in the randomized trial in Case I: Linear interaction functions with simple treatment-free effect functions.
AD
ADe
AD
ADe.
SABD
SABDe SABD, SABD
Original
0.679
0.734
0.703
0.813
0.720
0.745
0.734
0.804
(0.06)
(0.06)
(0.06)
(0.06)
(0.06)
(0.06)
(0.06)
(0.06)
200
0.637
0.673
0.658
0.755
0.691
0.700
0.702
0.761
(0.07)
(0.08)
(0.07)
(0.08)
(0.06)
(0.06)
(0.07)
(0.07)
0.633
0.649
0.644
0.720
0.671
0.686
0.685
0.739
(0.07)
(0.08)
(0.06)
(0.08)
(0.06)
(0.07)
(0.05)
(0.06)
20
0.816
0.872
0.828
0.925
0.832
0.867
0.818
0.875
(0.05)
(0.05)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.03)
600
40
0.814
0.863
0.808
0.895
0.840
0.867
0.821
0.871
(0.04)
(0.03)
(0.04)
(0.07)
(0.04)
(0.03)
(0.04)
(0.05)
60
0.802
0.848
0.783
0.876
0.836
0.859
0.815
0.864
(0.04)
(0.03)
(0.04)
(0.06)
(0.04)
(0.03)
(0.04)
(0.03)
20
0.865
0.902
0.867
0.953
0.866
0.882
0.846
0.875
(0.03)
(0.05)
(0.03)
(0.04)
(0.03)
(0.04)
(0.04)
(0.03)
1000
40
0.854
0.900
0.846
0.927
0.866
0.893
0.839
0.876
(0.03)
(0.03)
(0.03)
(0.07)
(0.04)
(0.03)
(0.04)
(0.04)
60
0.857
0.895
0.838
0.922
0,874
0.893
0.846
0.879
(0.03)
(0.02)
(0.03)
(0.05)
(0.03)
(0.02)
(0.03)
(0.03)
Augmented
20
0.798
0.807
0.806
0.853
0.807
0.807
0.820
0.854
(0.04)
(0.04)
(0.05)
(0.05)
(0.04)
(0.04)
(0.05)
(0.05)
200
40
0.781
0.779
0.776
0.814
0.784
0.776
0.787
0.813
(0.04)
(0.05)
(0.06)
(0.07)
(0.05)
(0.05)
(0.06)
(0.07)
60
0.769
0.764
0.759
0.789
0.772
0.762
0.776
0.795
(0.05)
(0.05)
(0.05)
(0.06)
(0.05)
(0.05)
(0.05)
(0.06)
20
0.891
0.908
0.898
0.939
0.899
0.912
0.896
0.934
(0.02)
(0.02)
(0.02)
(0.02)
(0.02)
(0.02)
(0.02)
(0.02)
600
40
0.891
0.898
0.892
0.923
0.899
0.902
0.897
0.925
(0.02)
(0.02)
(0.02)
(0.03)
(0.02)
(0.02)
(0.02)
(0.03)
60
0.883
0.889
0.873
0.910
0.893
0.893
0.886
0.915
(0.03)
(0.02)
(0.02)
(0.03)
(0.02)
(0.03)
(0.02)
(0.02)
0.921
0.929
0.924
0.957
0.927
0.931
0.921
0.948
(0.02)
(0.03)
(0.02)
(0.02)
(0.01)
(0.03)
(0.02)
(0.02)
1000
0.918
0.929
0.916
0.949
0.924
0.933
0.918
0.946
(0.02)
(0.02)
(0.01)
(0.02)
(0.01)
(0.01)
(0.02)
(0.02)
0.915
0,923
0.908
0.942
0.921
0,926
0.915
0.942
(0.02)
(0.02)
(0.02)
(0.03)
(0.02)
(0.02)
(0.02)
(0.02)
29

---

Table 3: Average accuracy and their standard deviations (in parenthesis) in the randomized trial in
Case 2: Linear interaction functions with complicated treatment-free effect functions.
AD
ADe
AD
ADe.s
SABD
SABDe
SABD,
SABD
Original
0.348
0.324
0.366
0.373
0.440
0.412
0.488
0.525
(0.12)
(0.14)
(0.14)
(0.18)
(0.13)
(0.16)
(0.14)
(0.21)
200
0.316
0.295
0.346
0.329
0.379
0.337
0.454
0.457
(0.10)
(0.10)
(0.13)
(0.14)
(0.12)
(0.12)
(0.15)
(0.19)
0.309
0.286
0.336
0.310
0.345
0.320
0.454
0,444
(0.08)
(0.08)
(0.11)
(0.12)
(0.10)
(0.10)
(0.13)
(0.18)
20
0.440
0.472
0.470
0.523
0.693
0.773
0.695
0.785
(0.15)
(0.19)
(0.17)
(0.22)
(0.08)
(0.12)
(0.09)
(0.12)
600
40
0.391
0.415
0.409
0.448
0.660
0.736
0.692
0.784
(0.16)
(0.20)
(0.17)
(0.21)
(0.11)
(0.13)
(0.10)
(0.12)
60
0.367
0.367
0.387
0.408
0.599
0.688
0.669
0.757
(0.14)
(0.16)
(0.15)
(0.18)
(0.13)
(0.15)
(0.11)
(0.12)
20
0.458
0,473
0.477
0.528
0.769
0.818
0.746
0.816
(0.16)
(0.18)
(0.17)
(0.21)
(0.07)
(0.10)
(0.08)
(0.08)
1000
40
0.408
0.410
0.436
0,480
0.741
0.826
0.744
0.827
(0.17)
(0.19)
(0.18)
(0.21)
(0.09)
(0.09)
(0.09)
(0.09)
60
0.436
0.467
0.465
0.492
0.746
0.833
0.766
0.816
(0.17)
(0.20)
(0.18)
(0.21)
(0.09)
(0.10)
(0.06)
(0.12)
Augmented
20
0.417
0.418
0.436
0.474
0.650
0.647
0.716
0.740
(0.18)
(0.18)
(0.19)
(0.20)
(0.10)
(0.11)
(0.10)
(0.10)
200
40
0.367
0.362
0.384
0.413
0.556
0.541
0.679
0.689
(0.15)
(0.15)
(0.17)
(0.19)
(0.14)
(0.15)
(0.11)
(0.12)
60
0.353
0.348
0.374
0.393
0.517
0.506
0.648
0.665
(0.14)
(0.13)
(0.15)
(0.16)
(0.15)
(0.14)
(0.11)
(0.12)
20
0.492
0.520
0.514
0.583
0.857
0.862
0.873
0.894
(0.19)
(0.20)
(0.20)
(0.21)
(0.04)
(0.04)
(0.03)
(0.06)
600
40
0.437
0.444
0.457
0.481
0.833
0.835
0.864
0.878
(0.20)
(0.20)
(0.20)
(0.22)
(0.05)
(0.04)
(0.04)
(0.04)
60
0.396
0.399
0.409
0.444
0.803
0.798
0.843
0.861
(0.17)
(0.18)
(0.17)
(0.19)
(0.06)
(0.08)
(0.04)
(0.04)
0.485
0.492
0.501
0.541
0.897
0.896
0.905
0.918
(0.18)
(0.19)
(0.19)
(0.21)
(0.03)
(0.04)
(0.02)
(0.04)
1000
0.433
0.436
0.453
0.504
0.888
0.890
0.903
0.921
(0.19)
(0.20)
(0.20)
(0.21)
(0.03)
(0.03)
(0.02)
(0.03)
0.469
0.485
0.490
0.513
0.884
0.883
0.899
0.908
(0.19)
(0.20)
(0.20)
(0.21)
(0.03)
(0.03)
(0.02)
(0.03)
30

---

Table 4: Average accuracy and their standard deviations (in parenthesis) in the randomized trial in
Case 3: Non-linear interaction functions with simple treatment-free effect functions.
AD
ADe
AD
ADe.
SABD
SABDe
SABD,
SABDe.
Original
0.385
0.353
0.411
0.404
0.476
0.438
0.521
0.550
(0.12)
(0.13)
(0.12)
(0.17)
(0.12)
(0.14)
(0.12)
(0.16)
200
0.339
0.316
0.375
0.358
0.408
0.366
0.473
0.486
(0.12)
(0.11)
(0.11)
(0.14)
(0.12)
(0.12)
(0.10)
(0.14)
0.319
0.300
0.357
0.342
0.370
0.343
0.454
0.438
(0.11)
(0.10)
(0.11)
(0.12)
(0.12)
(0.11)
(0.10)
(0.13)
20
0.534
0,561
0.558
0.630
0.678
0.705
0.659
0.746
(0.11)
(0.11)
(0.10)
(0.10)
(0.07)
(0.07)
(0.08)
(0.07)
600
40
0.490
0.486
0.529
0.576
0.651
0.653
0.651
0.723
(0.11)
(0.13)
(0.08)
(0.11)
(0.07)
(0.09)
(0.07)
(0.08)
60
0.472
0.460
0.501
0.536
0.637
0.648
0.639
0.687
(0.10)
(0.12)
(0.08)
(0.10)
(0.07)
(0.07)
(0.07)
(0.07)
20
0.577
0.598
0.599
0.690
0.740
0.755
0.689
0.768
(0.09)
(0.09)
(0.08)
(0.06)
(0.06)
(0.05)
(0.08)
(0.07)
1000
40
0.584
0.583
0.599
0.663
0.724
0.728
0.702
0.779
(0.10)
(0.09)
(0.09)
(0.07)
(0.06)
(0.08)
(0.06)
(0.06)
60
0.574
0.579
0.584
0.636
0.723
0.731
0.707
0.760
(0.09)
(0.09)
(0.06)
(0.08)
(0.05)
(0.05)
(0.05)
(0.07)
Augmented
20
0.524
0.524
0.555
0.603
0.608
0.604
0.642
0.687
(0.11)
(0.12)
(0.10)
(0.11)
(0.08)
(0.09)
(0.07)
(0.08)
200
40
0.476
0.463
0.513
0.548
0.565
0.545
0.607
0.630
(0.12)
(0.13)
(0.10)
(0.11)
(0.10)
(0.09)
(0.07)
(0.08)
60
0.443
0.431
0.484
0.508
0.525
0.505
0.601
0.601
(0.12)
(0.13)
(0.12)
(0.13)
(0.11)
(0.12)
(0.08)
(0.09)
20
0.674
0.680
0.685
0.719
0.776
0.784
0.776
0.805
(0.06)
(0.07)
(0.05)
(0.06)
(0.04)
(0.04)
(0.04)
(0.04)
600
40
0.654
0.659
0.664
0.697
0.761
0.763
0.768
0.796
(0.07)
(0.06)
(0.06)
(0.06)
(0.04)
(0.04)
(0.04)
(0.04)
60
0.643
0.645
0.652
0.676
0.750
0.749
0.756
0.780
(0.06)
(0.07)
(0.05)
(0.06)
(0.05)
(0.05)
(0.04)
(0.04)
0.706
0.708
0.715
0.746
0.805
0.805
0.787
0.817
(0.05)
(0.06)
(0.04)
(0.04)
(0.03)
(0.05)
(0.04)
(0.04)
1000
0.704
0.704
0.711
0.738
0.805
0.809
0.807
0.831
(0.06)
(0.05)
(0.05)
(0.04)
(0.03)
(0.03)
(0.04)
(0.03)
0.705
0.707
0.710
0.731
0.800
0.798
0.803
0.820
(0.05)
(0.04)
(0.04)
(0.04)
(0.03)
(0.03)
(0.04)
(0.04)
31

---

Table 5: Average accuracy and their standard deviations (in parenthesis) in the randomized trial in Case Non-linear interaction functions with complicated treatment- free effect function.
AD
ADe
AD
ADe.
SABD
SABDe
SABD, SABD
Original
0.298
0.300
0.319
0.299
0.351
0.345
0.384
0.379
(0.09)
(0.09)
(0.10)
(0.10)
(0.12)
(0.12)
(0.13)
(0.16)
200
0.273
0.269
0.283
0.279
0.303
0.293
0.350
0.343
(0.06)
(0.05)
(0.07)
(0.07)
(0.09)
(0.09)
(0.11)
(0.13)
0.279
0.275
0.299
0.289
0.298
0.287
0.359
0.340
(0.07)
(0.06)
(0.08)
(0.09)
(0.09)
(0.09)
(0.12)
(0.13)
20
0.320
0.331
0.339
0.353
0.570
0.570
0.572
0.622
(0.11)
(0.12)
(0.12)
(0.13)
(0.12)
(0.13)
(0.13)
(0.13)
600
40
0.308
0.306
0.329
0.332
0.501
0.511
0.552
0.600
(0.10)
(0.09)
(0.11)
(0.13)
(0.13)
(0.13)
(0.10)
(0.11)
60
0.297
0,293
0.308
0.307
0,463
0.473
0.529
0.578
(0.08)
(0.08)
(0.09)
(0.10)
(0.13)
(0.14)
(0.11)
(0.14)
20
0.325
0.331
0.337
0.356
0.638
0.656
0.614
0.658
(0.11)
(0.12)
(0.12)
(0.14)
(0.09)
(0.09)
(0.10)
(0.10)
1000
40
0.312
0.316
0.322
0.335
0.631
0.635
0.625
0.674
(0.10)
(0.11)
(0.11)
(0.13)
(0.09)
(0.12)
(0.09)
(0.09)
60
0.308
0.328
0.323
0.341
0.629
0.637
0.628
0.675
(0.09)
(0.12)
(0.10)
(0.13)
(0.08)
(0.10)
(0.08)
(0.10)
Augmented
20
0.347
0.355
0.355
0.380
0.500
0.481
0.558
0.595
(0.12)
(0.13)
(0.14)
(0.14)
(0.13)
(0.14)
(0.12)
(0.13)
200
40
0.288
0.287
0.299
0.314
0.376
0.377
0.506
0.499
(0.08)
(0.08)
(0.10)
(0.11)
(0.13)
(0.12)
(0.11)
(0.14)
60
0.307
0.302
0.319
0.327
0.379
0.362
0.492
0.481
(0.10)
(0.09)
(0.11)
(0.12)
(0.13)
(0.12)
(0.13)
(0.14)
20
0.348
0.369
0.361
0.401
0.725
0.724
0.749
0.764
(0.13)
(0.14)
(0.14)
(0.16)
(0.08)
(0.08)
(0.06)
(0.09)
600
40
0.349
0.353
0.360
0.382
0.667
0.653
0.729
0.747
(0.13)
(0.14)
(0.14)
(0.15)
(0.09)
(0.10)
(0.06)
(0.06)
60
0.318
0.318
0.336
0.343
0.651
0.644
0.720
0.734
(0.12)
(0.11)
(0.13)
(0.13)
(0.10)
(0.11)
(0.07)
(0.07)
0.339
0.352
0.350
0.390
0.775
0.779
0.781
0,800
(0.14)
(0.15)
(0.14)
(0.16)
(0.05)
(0.06)
(0.06)
(0.05)
1000
0.332
0.339
0.346
0.361
0.764
0.756
0.793
0.803
(0.12)
(0.14)
(0.13)
(0.14)
(0.06)
(0.07)
(0.05)
(0.05)
0.331
0.347
0,342
0.358
0.766
0.753
0.788
0.803
(0.13)
(0.13)
(0.14)
(0.14)
(0.05)
(0.05)
(0.04)
(0.04)
32

---

Table 6: Average accuracy and their standard deviations (in parenthesis) in the observational study in
Case I: Linear interaction functions with simple treatment-free effect functions.
AD
ADe
AD
ADe.
SABD
SABDe SABD, SABD
Original
0.560
0.604
0.580
0.660
0.555
0.634
0.563
0.694
(0.11)
(0.12)
(0.10)
(0.11)
(0.10)
(0.10)
(0.10)
(0.10)
200
0.494
0.499
0.532
0.589
0.517
0.550
0.542
0.624
(0.12)
(0.14)
(0.11)
(0.13)
(0.09)
(0.12)
(0.10)
(0.11)
0.498
0.494
0.520
0.546
0.487
0.533
0.512
0.583
(0.11)
(0.13)
(0.09)
(0.13)
(0.11)
(0.12)
(0.10)
(0.12)
20
0.728
0.785
0.738
0.849
0.675
0.824
0.664
0.867
(0.07)
(0.07)
(0.07)
(0.06)
(0.06)
(0.04)
(0.05)
(0.05)
600
40
0.716
0.754
0.721
0.824
0.674
0.802
0.670
0.856
(0.08)
(0.07)
(0.07)
(0.07)
(0.07)
(0.06)
(0.07)
(0.05)
60
0.704
0.756
0.693
0.795
0.661
0.778
0.653
0.829
(0.08)
(0.07)
(0.07)
(0.08)
(0.08)
(0.06)
(0.07)
(0.05)
20
0.782
0.845
0.788
0.894
0.708
0.873
0.689
0.905
(0.06)
(0.05)
(0.06)
(0.05)
(0.06)
(0.03)
(0.05)
(0.04)
1000
40
0.774
0.831
0.767
0.886
0.700
0.859
0.684
0.901
(0.06)
(0.04)
(0.06)
(0.05)
(0.06)
(0.04)
(0.05)
(0.04)
60
0.764
0.825
0.755
0.868
0.702
0.845
0.684
0.885
(0.06)
(0.05)
(0.06)
(0.06)
(0.07)
(0.04)
(0.05)
(0.04)
Augmented
20
0.692
0.699
0.690
0.727
0.712
0.722
0.702
0.741
(0.07)
(0.07)
(0.07)
(0.08)
(0.07)
(0.07)
(0.08)
(0.08)
200
40
0.660
0.662
0.661
0.687
0.676
0.678
0.677
0.707
(0.09)
(0.08)
(0.08)
(0.09)
(0.07)
(0.08)
(0.08)
(0.08)
60
0.648
0.632
0.631
0.650
0.660
0.657
0.649
0.672
(0.08)
(0.09)
(0.08)
(0.08)
(0.08)
(0.08)
(0.10)
(0.10)
20
0.840
0.861
0.841
0.889
0.841
0.881
0.836
0.904
(0.04)
(0.03)
(0.04)
(0.04)
(0.04)
(0.03)
(0.04)
(0.04)
600
40
0.829
0.841
0.828
0.872
0.831
0.870
0.832
0.898
(0.05)
(0.04)
(0.05)
(0.04)
(0.05)
(0.03)
(0.04)
(0.03)
60
0.817
0.829
0.811
0.849
0.827
0.856
0.819
0.874
(0.05)
(0.05)
(0.04)
(0.05)
(0.05)
(0.04)
(0.04)
(0.04)
0.879
0.900
0.879
0.926
0.867
0.919
0.861
0.932
(0.03)
(0.03)
(0.03)
(0.02)
(0.03)
(0.02)
(0.03)
(0.02)
1000
0.878
0.892
0.874
0.917
0.867
0.909
0.860
0.927
(0.04)
(0.03)
(0.04)
(0.03)
(0.03)
(0.03)
(0.03)
(0.02)
0.868
0.883
0.862
0.905
0.859
0,904
0.852
0.919
(0.03)
(0.03)
(0.03)
(0.04)
(0.04)
(0.02)
(0.04)
(0.03)
33

---

Table 7: Average accuracy and their standard deviations (in parenthesis) in the observational study in
Case 2: Linear interaction functions with complicated treatment-free effect functions.
AD
ADe
AD,
ADe.s
SABD
SABDe
SABD,
SABD
Original
0.321
0.327
0.328
0.356
0.292
0.361
0.319
0.394
(0.15)
(0.14)
(0.15)
(0.16)
(0.12)
(0.17)
(0.14)
(0.20)
200
0.301
0.300
0.315
0.325
0.284
0.318
0.309
0.345
(0.11)
(0.12)
(0.12)
(0.12)
(0.11)
(0.10)
(0.13)
(0.14)
0.303
0.320
0.315
0.323
0.271
0.341
0.314
0.381
(0.11)
(0.12)
(0.11)
(0.12)
(0.12)
(0.13)
(0.12)
(0.15)
20
0.364
0.383
0.373
0.413
0,403
0.659
0.414
0.689
(0.15)
(0.17)
(0.15)
(0.18)
(0.09)
(0.15)
(0.09)
(0.16)
600
40
0.355
0.379
0.368
0.388
0.360
0.571
0.404
0.662
(0.13)
(0.15)
(0.14)
(0.15)
(0.10)
(0.19)
(0.09)
(0.16)
60
0.327
0.335
0.329
0.350
0.366
0.515
0.401
0.645
10.11)
(0.14)
(0.12)
(0.14)
(0.10)
(0.19)
(0.09)
(0.17)
20
0.348
0.390
0.352
0.393
0,427
0.751
0.417
0.726
(0.14)
(0.16)
(0.14)
(0.17)
(0.06)
(0.14)
(0.07)
(0.14)
1000
40
0.364
0.382
0.374
0.420
0,426
0.712
0.430
0.767
(0.15)
(0.16)
(0.15)
(0.17)
(0.07)
(0.13)
(0.07)
(0.11)
60
0.372
0.415
0.386
0.449
0.413
0.741
0.436
0.771
(0.14)
(0.18)
(0.15)
(0.19)
(0.08)
(0.13)
(0.08)
(0.12)
Augmented
20
0.368
0.393
0.382
0.440
0.576
0.590
0.636
0.662
(0.17)
(0.18)
(0.18)
(0.20)
(0.15)
(0.15)
(0.14)
(0.14)
200
40
0.340
0.358
0.353
0.395
0,472
0.493
0.586
0.625
(0.15)
(0.15)
(0.16)
(0.18)
(0.17)
(0.16)
(0.14)
(0.15)
60
0.330
0.355
0.339
0.378
0,453
0.447
0.567
0.579
(0.12)
(0.14)
(0.12)
(0.14)
(0.13)
(0.14)
(0.14)
(0.14)
20
0.403
0.446
0.414
0.497
0.809
0.827
0.828
0.861
(0.17)
(0.18)
(0.19)
(0.22)
(0.06)
(0.05)
(0.05)
(0.05)
600
40
0.381
0.422
0.397
0.458
0.761
0.766
0.818
0.843
(0.16)
(0.18)
(0.17)
(0.18)
(0.09)
(0.08)
(0.05)
(0.06)
60
0.341
0.366
0.351
0.395
0.740
0.739
0.794
0.815
(0.12)
(0.14)
(0.13)
(0.16)
(0.10)
(0.10)
(0.06)
(0.06)
0.374
0.420
0.381
0.449
0.853
0.862
0.864
0.892
(0.17)
(0.18)
(0.18)
(0.20)
(0.07)
(0.05)
(0.04)
(0.03)
1000
0.392
0.429
0.403
0.443
0.836
0.848
0.867
0.885
(0.17)
(0.18)
(0.18)
(0.18)
(0.07)
(0.04)
(0.03)
(0.03)
0.405
0.446
0.420
0.484
0.839
0,844
0.864
0.885
(0.17)
(0.19)
(0.18)
(0.20)
(0.04)
(0.05)
(0.03)
(0.03)
34

---

Table 8: Average accuracy and their standard deviations (in parenthesis) in the observational study in
Case 3: Non-linear interaction functions with simple treatment-fres effect functions.
AD
ADe
AD
ADe.
SABD
SABDe SABD,
SABD
Original
0.293
0.301
0.326
0.336
0.382
0.393
0.423
0,474
(0.08)
(0.09)
(0.10)
(0.14)
(0.10)
(0.12)
(0.10)
(0.17)
200
0.304
0.314
0.358
0.346
0.368
0.376
0.428
0.465
(0.09)
(0.09)
(0.10)
(0.13)
(0.11)
(0.12)
(0.10)
(0.15)
0.277
0.284
0.330
0.345
0.337
0.331
0.401
0.416
(0.08)
(0.08)
(0.10)
(0.11)
(0.10)
(0.10)
(0.09)
(0.12)
20
0.405
0.491
0.458
0.587
0.547
0.685
0.529
0.747
(0.11)
(0.12)
(0.11)
(0.13)
(0.07)
(0.07)
(0.06)
(0.08)
600
40
0.382
0.462
0.447
0.549
0.523
0.639
0.531
0.729
(0.10)
(0.12)
(0.10)
(0.13)
(0.08)
(0.08)
(0.07)
(0.08)
60
0.386
0.438
0.465
0.545
0.519
0.628
0.534
0.696
(0.10)
(0.11)
(0.08)
(0.11)
(0.08)
(0.09)
(0.06)
(0.08)
20
0.471
0.572
0.518
0.671
0.587
0.739
0.548
0.779
(0.09)
(0.10)
(0.08)
(0.08)
(0.06)
(0.05)
(0.06)
(0.06)
1000
40
0.440
0.542
0.502
0.633
0.575
0.725
0.554
0.775
(0.08)
(0.09)
(0.07)
(0.08)
(0.06)
(0.05)
(0.05)
(0.07)
60
0.419
0.531
0.495
0.610
0.570
0.719
0.559
0.759
(0.09)
(0.09)
(0.07)
(0.10)
(0.06)
(0.05)
(0.05)
(0.06)
Augmented
20
0.448
0.465
0.492
0.554
0.535
0.568
0.571
0.663
(0.12)
(0.13)
(0.11)
(0.12)
(0.10)
(0.11)
(0.08)
(0.09)
200
40
0.456
0.461
0.505
0.544
0.511
0.534
0.562
0.629
(0.11)
(0.12)
(0.10)
(0.11)
(0.11)
(0.11)
(0.09)
(0.09)
60
0.394
0.394
0.445
0.476
0.449
0.458
0.532
0.576
(0.11)
(0.12)
(0.11)
(0.11)
(0.11)
(0.12)
(0.09)
(0.09)
20
0.617
0.661
0.639
0.707
0.700
0.773
0.675
0.798
(0.07)
(0.07)
(0.07)
(0.06)
(0.06)
(0.04)
(0.06)
(0.05)
600
40
0.606
0.643
0.631
0.688
0.691
0.753
0.685
0.792
(0.07)
(0.07)
(0.06)
(0.06)
(0.07)
(0.05)
(0.06)
(0.05)
60
0.612
0.639
0.634
0.678
0.692
0.748
0.693
0.776
(0.07)
(0.07)
(0.07)
(0.08)
(0.06)
(0.04)
(0.06)
(0.05)
0.672
0.707
0.686
0.743
0.743
0.811
0.708
0.820
(0.05)
(0.05)
(0.05)
(0.05)
(0.06)
(0.03)
(0.05)
(0.04)
1000
0.657
0.693
0.674
0.719
0.744
0.797
0.726
0.826
(0.05)
(0.05)
(0.05)
(0.07)
(0.06)
(0.03)
(0.05)
(0.04)
0.648
0.687
0.662
0.714
0.724
0.794
0.710
0.815
(0.05)
(0.05)
(0.05)
(0.05)
(0.06)
(0.04)
(0.05)
(0.04)
35

---

Table 9; Average accuracy and their standard deviations (in parenthesis) in the observational study in Case Non-linear interaction functions with complicated treatment-free effect function.
AD
ADe
AD
ADe.
SABD
SABDe
SABD, SABD
Original
0.274
0.279
0.281
0.283
0.302
0.322
0.325
0.368
(0.07)
(0.08)
(0.08)
(0.09)
(0.08)
(0.11)
(0.09)
(0.13)
200
0.270
0.270
0.285
0.276
0.285
0.305
0.322
0.353
(0.07)
(0.07)
(0.07)
(0.07)
(0.07)
(0.09)
(0.09)
(0.13)
0.266
0.270
0.281
0.272
0.274
0.281
0.313
0.323
(0.06)
(0.06)
(0.07)
(0.06)
(0.07)
(0.08)
(0.08)
(0.10)
20
0.294
0,309
0.304
0.338
0,409
0.555
0.422
0.619
(0.09)
(0.11)
(0.10)
(0.15)
(0.06)
(0.12)
(0.06)
(0.12)
600
40
0.283
0.286
0.300
0.314
0.369
0.489
0.412
0,580
(0.08)
(0.09)
(0.09)
(0.12)
(0.07)
(0.13)
(0.07)
(0.14)
60
0.278
0.291
0.292
0.308
0.371
0.457
0.408
0,583
(0.07)
(0.08)
(0.08)
(0.11)
(0.07)
(0.15)
(0.06)
(0.12)
20
0.288
0.314
0.296
0.324
0,441
0.643
0.445
0.663
(0.09)
(0.12)
(0.09)
(0.12)
(0.04)
(0.10)
(0.05)
(0.10)
1000
40
0.307
0.314
0.327
0.347
0,431
0.619
0.441
0.693
(0.09)
(0.11)
(0.11)
(0.13)
(0.04)
(0.09)
(0.04)
(0.10)
60
0.289
0.312
0.306
0.344
0,421
0.603
0.445
0.668
(0.07)
(0.10)
(0.09)
(0.13)
(0.05)
(0.09)
(0.04)
(0.09)
Augmented
20
0.317
0.336
0.335
0.384
0.439
0.456
0.505
0.571
(0.11)
(0.13)
(0.12)
(0.15)
(0.14)
(0.14)
(0.12)
(0.13)
200
40
0.291
0.307
0.305
0.329
0,370
0.379
0.481
0.512
(0.09)
(0.10)
(0.10)
(0.11)
(0.13)
(0.12)
(0.12)
(0.14)
60
0.287
0.289
0.302
0.328
0.333
0.329
0.446
0.461
(0.07)
(0.07)
(0.09)
(0.11)
(0.10)
(0.10)
(0.13)
(0.13)
20
0.337
0.361
0.346
0.390
0.660
0.702
0.676
0.761
(0.12)
(0.13)
(0.13)
(0.15)
(0.08)
(0.07)
(0.07)
(0.07)
600
40
0.329
0.347
0.350
0.373
0.613
0.630
0.684
0.741
(0.12)
(0.13)
(0.13)
(0.15)
(0.10)
(0.12)
(0.07)
(0.06)
60
0.307
0.322
0.326
0.347
0.612
0.635
0.679
0.726
(0.10)
(0.11)
(0.11)
(0.14)
(0.10)
(0.11)
(0.07)
(0.06)
0.318
0.342
0.325
0.369
0.734
0.769
0.730
0.801
(0.11)
(0.14)
(0.12)
(0.15)
(0.08)
(0.08)
(0.06)
(0.06)
1000
0.332
0.346
0.348
0.379
0.714
0.741
0.741
0.796
(0.12)
(0.14)
(0.13)
(0.15)
(0.09)
(0.07)
(0.06)
(0.05)
0.322
0.346
0.334
0.384
0.711
0.748
0.734
0.793
(0.11)
(0.13)
(0.12)
(0.15)
(0.07)
(0.06)
(0.06)
(0.04)
36

---

133.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 (d(() are computed by utilizzig the generated potential outcomes on the test dataset.
We include the empirical value of the optimal rule E[ (d (X)) as well as the values obtained by assigning everyone to each of the four classes E (K)] for K I,..., 4 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 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 is. For instance, in the randomized trial of Case (Table 9), the optimal empirical value is not significantly different from the mean empirical value evaluated by current methods. Especially with augmented outcomes using 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-Leaming with 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 "one-size-fits- outcome in small sample sizes, and this finding is not irrelevant to fact that majority of AD-Learning results in an intercept-only model as the final decision function under conditions of highly non-linear covariate-outcone relationships. In such instances, we can explore alternative methods to further advance our current results such as other penalty types performance metrics.
how challenging the current scenario
37

---

Table 10: Average empirical values and their standard deviations (in parenthesis) in the randomized trial Case 1: Linear interaction functions with simple treatment- free effect functions. The optimal empirical value is 5.160 and the average empirical values when assigning one treatment are 1.751, 1.752, 1.751, and 1.752, respectively.
AD
AD_
AD_s
e
SABD Original
SABD_e
SABD_8
SABD e,
20
4.554
4.725
4.637
4.934
4.701
4.769
4.739
4.925
(0.25)
(0.21)
(0.23)
(0.17)
(0.21)
(0.19)
(0.21)
(0.16)
200
40
4.381
4.505
4.465
4.778
4.591
4.616
4.626
4.798
(0.31)
(0.34)
(0.30)
(0.28)
(0.24)
(0.24)
(0.27)
(0.26)
60
4.364
4.409
4.412
4.670
4.528
4.566
4.578
4.746
(0.29)
(0.36)
(0.25)
(0.27)
(0.21)
(0.26)
(0.20)
(0.21)
20
4.963
5.053
4.985
5.117
4.992
5.050
4,956
5.063
(0.10)
(0.10)
(0.09)
(0.07)
(0.09)
(0.10)
(0.11)
(0.06)
600
40
4.964
5.049
4.950
5.071
5.013
5.059
4.972
5.056
(0.10)
(0.06)
(0.09)
(0.19)
(0.08)
(0.06)
(0.09)
(0.12)
60
4.935
5.025
4.891
5.059
5.006
5.045
4.962
5.054
(0.11)
(0.07)
(0.10)
(0.11)
(0.09)
(0.07)
(0.09)
(0.07)
20
5.054
5.092
5.057
5.140
5.053
5.073
5.012
5.060
(0.07)
(0.11) (0.07)
(0.07)
(0.07)
(0.08)
(0.09)
(0.07)
1000
40
5.034
5.098
5.020
5.101
5,051
5.090
5.000
5.060
(0.07)
(0.05)
(0.07)
(0.16)
(0.07)
(0.05)
(0.09)
(0.09)
60
5.046
5.096
5.013
5.112
5,070
5.095
5.021
5.072
(0.06)
(0.04)
(0.06)
(0.09)
(0.06)
(0.04)
(0.08)
(0.07)
Augmented
20
4.928
4.940
4.936
5.022
4.951
4.942
4.964
5.024
(0.10)
(0.10)
(0.15)
(0.14)
(0.09)
(0.11)
(0.15)
(0.14)
200
40
4.887
4.873
4.857
4.940
4.888
4.863
4.884
4.936
(0.12)
(0.13)
(0.22)
(0.20)
(0.12)
(0.13)
(0.21)
(0.22)
60
4.851
4.833
4.817
4.888
4,860
4.826
4.865
4.905
(0.12)
(0.14)
(0.17)
(0.16)
(0.13)
(0.15)
(0.16)
(0.16)
20
5.092
5.109
5.100
5.136
5.101
5.114
5.097
5.133
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
600
40
5.095
5.101
5.096
5.124
5.105
5.105
5.102
5.127
(0.05)
(0.04)
(0.04)
(0.06)
(0.04)
(0.04)
(0.04)
(0.05)
60
5.086
5.091
5.073
5.114
5.097
5.096
5.089
5.120
(0.05)
(0.05)
(0.05)
(0.05)
(0.05)
(0.05)
(0.05)
(0.04)
20
5.125
5.127
5.127
5.148
5.130
5.131
5.125
5.143
(0.04)
(0.05)
(0.04)
(0.04)
(0.04)
(0.05)
(0.04)
(0.04)
1000
40
5.121
5.129
5.118
5.140
5.126
5.132
5.120
5.139
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
60
5,121
5.128
5.115
5.140
5.127
5.131
5.120
5.141
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
(0.04)
38

---

Table 11: Average empirical values and their standard deviations (in parenthesis) in the randomized trial in Case Linear interaction functions with complicated treatment-fres effect functions. The optimal empirical value is 10.333 and the average empirical values when assigning one treatment are 6.923, 6.924, 6.923, and 6.924, respectively.
AD
AD_
AD_s
AD_e,
SABD
SABD_e
SABD_s
SABD_e
Original
20
7.704
7.474
7.820
7.764
8.360
8.073
8.677
8.736
(0.90)
(1.00)
(0.98)
(1.14)
(0.96)
(1.16)
(0.94)
(1.26)
200
40
7.430
7.262
7.641
7.482
7.907
7.529
8.415
8.339
(0.80)
(0.77)
(0.94)
(0.98)
(0.93)
(0,92)
(1.08)
(1.19)
60
7.413
7.205
7.605
7.359
7.690
7.470
8,475
8,275
(0.79)
(0.74)
(0.93)
(0.94)
(0.91)
(0.83)
(0.94)
(1.23)
20
8.334
8.437
8.488
8.719
9.791
9.966
9.803
T0.023
(1.11)
(1.27)
(1.18)
(1.29)
(0.56)
(0.76)
(0.56)
(0.63)
600
40
7.983
8.051
8.079
8.250
9.610
9.821
9.757
9.991
(1.11)
(1.25)
(1.16)
(1.33)
(0.69)
(0.73)
(0.64)
(0.69)
60
7.834
7.788
7.956
8.025
9.311
9.648
9.657
9,930
(1.06)
(1.12)
(1.14)
(1.26)
(0.85)
(0.80)
(0.66)
(0.59)
20
8.395
8.458
8.486
8.706
9.997
10.087
9.929
10.101
(1.15)
(1.20)
(1.19)
(1.27)
(0.47)
(0.54)
(0.50)
(0.50)
1000
40
8.033
7.993
8.199
8.400
9.874
10.077
9.893
10.101
(1.20)
(1.24)
(1.25)
(1.32)
(0.59)
(0.51)
(0.58)
(0.47)
60
8.275
8.403
8.439
8.506
9.933
10.122
10.021
10.053
(1.17)
(1.24)
(1.21)
(1.32)
(0.54)
(0.59)
(0.44)
(0.65)
Augmented
20
8.111
8.117
8.201
8,429
9.552
9.526
9.801
9.877
(1.19)
(1.16)
(1.24)
(1.28)
(0.64)
(0.65)
(0.53)
(0.49)
200
40
7.758
7.712
7.859
8.038
9.017
8.912
9.646
9.678
(1.09)
(1.09)
(1.17)
(1.28)
(0.97)
(1.02)
(0.62)
(0.65)
60
7.694
7.649
7.834
7.939
8.814
8.764
9.519
9.589
(1.03)
(1.01)
(1.11)
(1.15)
(0.96)
(0.95)
(0.70)
(0.67)
20
8.551
8.713
8.663
9.028
10.236
10.244
10.265
10.278
(1,25)
(1.23)
(1.28)
(1.21)
(0.36)
(0.35)
(0.35)
(0.39)
600
40
8.212
8.237
8.317
8.437
10.173
10.176
10.234
10.249
(1.25)
(1.27)
(1.27)
(1.34)
(0.37)
(0.37)
(0.36)
(0.36)
60
7.972
7.998
8.048
8.249
10.100
10.080
10.201
10.230
(1.16)
(1.16)
(1.20)
(1.27)
(0.38)
(0.44)
(0.34)
(0.33)
20
8.509
8.549
8.587
8.769
10.262
10.255
10.271
10.280
(1.22)
(1.22)
(1.24)
(1.26)
(0.41)
(0.41)
(0.41)
(0.41)
1000
40
8.158
8.154
8.265
8.543
10.230
10.232
10.252
10.266
(1.26)
(1.27)
(1.31)
(1.35)
(0.38)
(0.38)
(0.38)
(0.38)
60
8,414
8.510
8.532
8.648
10.262
10,260
10.281
10.287
(1.26)
(1.25)
(1.28)
(1.31)
(0.40)
(0.40)
(0.40)
(0.40)
39

---

Table 12: Average empirical values and their standard deviations (in parenthesis) in the randomized trial in Case 3: Non-linear interaction functions with simple fratment-free effect functions. The optimal empirical value is 1.422 and the average empirical values when assigning one treatment are 0.041, -0.919, -0.419, and -0.920, respectively.
AD
AD_e
AD_s
AD_
SABD Original
SABD_e SABD_s
SABD
20
0.210
0.157
0.294
0.382
0.434
0.275
0.546
0.592
(0.45)
(0.48)
(0.44)
(0.45)
(0.47)
(0.55)
(0.46)
(0.59)
200
40
0.023
-0.018
0.131
0.185
0.196
0.044
0.436
0.462
(0.46)
(0.46)
(0.44)
(0.45)
(0.50)
(0.54)
(0.38)
(0.50)
60
-0.051
-0.087
0.072
0.100
0.081
-0.008
0.346
0.315
(0.49)
(0.45)
(0.45)
(0.45)
(0.51)
(0.48)
(0.43)
(0.48)
20
0.733
0.826
0.788
0.983
0.963
1.023
0.931
1.095
(0.29)
(0.26)
(0.27)
(0.21)
(0.18)
(0.17)
(0.23)
(0.17)
600
40
0.631
0.646
0.715
0.857
0.910
0.915
0.923
1.025
(0.33)
(0.32)
(0.22)
(0.29)
(0.21)
(0.24)
(0.18)
(0.21)
60
0.609
0.585
0.651
0.770
0.905
0.922
0.903
0.993
(0.28)
(0.32)
(0.22)
(0.27)
(0.19)
(0.18)
(0.17)
(0.17)
20
0.879
0.930
0.920
I.1II
1.089
1.104
1.009
1.123
(0.22)
(0.22)
(0.19)
(0.12)
(0.15)
(0.14)
(0.18)
(0.14)
1000
40
0.864
0.880
0.887
1.061
1.027
1.020
0.993
1.117
(0.23)
(0.21)
(0.21)
(0.14)
(0.15)
(0.25)
(0.16)
(0.14)
60
0.870
0.890
0.878
1.002
1.050
1.064
1.030
1.102
(0.20)
(0.21)
(0.16)
(0.24)
(0.14)
(0.14)
(0.14)
(0.16)
Augmented
20
0.715
0.697
0.794
0.918
0.883
0.871
0.967
1.056
(0.32)
(0.38)
(0.24)
(0.24)
(0.19)
(0.26)
(0.17)
(0.18)
200
40
0.588
0.545
0.680
0.779
0.791
0.724
0.884
0.926
(0.30)
(0.36)
(0.27)
(0.25)
(0.26)
(0.30)
(0.20)
(0.22)
60
0.481
0.444
0.578
0.653
0.661
0.587
0.849
0.840
(0.35)
(0.38)
(0.32)
(0.35)
(0.32)
(0.42)
(0.19)
(0.23)
20
1.082
1.093
1.100
1.151
1.184
1.194
1.216
1.245
(0.11)
(0.12)
(0.10)
(0.10)
(0.10)
(0.10)
(0.07)
(0,08)
600
40
1.051
1.065
1.063
1.118
1.166
1.163
1.189
1.220
(0.12)
(0.11)
(0.11)
(0.12)
(0.11)
(0.10)
(0.10)
(0.09)
60
1.039
1.042
1.051
1.092
1.161
1.156
1.173
1.204
(0.12)
(0.13)
(0.10)
(0.12)
(0.10)
(0.10)
(0.09)
(0,09)
20
1.141
1.134
1.156
1.196
1,219
1,210
1.229
1.260
(0.08)
(0.15)
(0.08)
(0.07)
(0.08)
(0.12)
(0.08)
(0.07)
1000
40
1.133
1.137
1.143
1.183
1.207
1,213
1.231
1.258
(0.09)
(0.09)
(0.09)
(0.07)
(0.07)
(0.08)
(0.07)
(0.06)
60
1.141
1.146
1.148
1.177
1,212
1,207
1.235
1,250
(0.09)
(0.08)
(0.08)
(0.08)
(0.09)
(0.08)
(0.08)
(0.08)
40

---

Table 13: Average empirical values and their standard deviations (in parenthesis) in the randomized trial in Case Non-lineat interaction functions with a complicated treatment-free effect function. The optimal empirical value is 5.623 and the average empirical values when assigning one treatment are 4. 4.243, 3. 3.282, 3. 3.782, and 3.281 respectively.
AD
AD_
AD_s
AD_e,
SABD Original
SABD_e
SABD_8
SABD_e.
3.979
4.033
4.058
4.052
4.192
4.132
4.381
4.366
(0.55)
(0.53)
(0.57)
(0.58)
(0.58)
(0.61)
(0.58)
(0.65)
200
40
3.890
3.876
3,920
3,913
3.984
3.936
4.184
4.168
(0.55)
(0.54)
(0.57)
(0.59)
(0.65)
(0.61)
(0.63)
(0.65)
60
3,939
3.931
3.991
3,970
4.018
3.938
4.241
4.137
(0.57)
(0.56)
(0.58)
(0.59)
(0.60)
(0.59)
(0.63)
(0.70)
20
4.134
4.198
4.202
4,282
4.995
4,980
5.009
5.127
(0.60)
(0.64)
(0.64)
(0.62)
(0.49)
(0.52)
(0.46)
(0.48)
600
40
4.032
4.066
4.111
4.183
4.761
4.774
4.951
5.056
(0.64)
(0.64)
(0.68)
(0.66)
(0.57)
(0.60)
(0.48)
(0.49)
60
3.989
4.021
4.029
4.042
4.667
4.685
4.885
4.964
(0.64)
(0.62)
(0.66)
(0.67)
(0.61)
(0.62)
(0.49)
(0.60)
20
4.117
4.130
4.152
4.218
5.138
5.168
5,083
5.191
(0.66)
(0.67)
(0.68)
(0.71)
(0.43)
(0.44)
(0.46)
(0.42)
1000
40
4.047
4.061
4.093
4.188
5,075
5.078
5.070
5.188
(0.62)
(0.64)
(0.66)
(0.71)
(0.49)
(0.51)
(0.48)
(0.43)
60
4.103
4.190
4.154
4.232
5.144
5.152
5.155
5.237
(0.64)
(0.67)
(0.65)
(0.69)
(0.45)
(0.49)
(0.43)
(0.48)
Augmented
20
4.203
4.247
4.225
4.333
4.733
4.648
4.923
5.016
(0.58)
(0.58)
(0.60)
(0.61)
(0.58)
(0.66)
(0.52)
(0.55)
200
40
3.962
3.962
4.010
4.029
4.332
4.288
4.745
4.700
(0.60)
(0.62)
(0.61)
(0.65)
(0.64)
(0.64)
(0.54)
(0.59)
60
4.011
3.987
4.050
4.095
4.350
4.318
4.707
4.643
(0.64)
(0.61)
(0.67)
(0.64)
(0.61)
(0.55)
(0.56)
(0.63)
20
4.245
4.293
4.291
4.425
5,307
5.302
5.364
5.375
(0.63)
(0.68)
(0.65)
(0.68)
(0.38)
(0.39)
(0.36)
(0.39)
600
40
4.199
4.262
4.227
4.328
5.202
5.160
5.320
5,334
(0.69)
(0.67)
(0.71)
(0.73)
(0.42)
(0.39)
(0.37)
(0,36)
60
4.073
4.084
4.140
4.166
5.175
5.161
5.313
5.327
(0.71)
(0.71)
(0.74)
(0.74)
(0.41)
(0.39)
(0.34)
(0.35)
20
4.147
4.199
4.175
4.331
5,356
5.355
5,372
5,393
(0.71)
(0.70)
(0.72)
(0.75)
(0.39)
(0.41)
(0.42)
(0.42)
1000
40
4.124
4.141
4.168
4.214
5,331
5.310
5.369
5,371
(0.69)
(0.71)
(0.72)
(0.77)
(0.38)
(0.40)
(0.38)
(0.40)
60
4.183
4.264
4.213
4.298
5.377
5.362
5,418
5,433
(0.69)
(0.69)
(0.73)
(0.71)
(0.40)
(0.41)
(0.40)
(0.40)
41

---

Table 14: Average empirical values and their standard deviations (in parenthesis) in the observational study in Case Linear interaction functions with simple treatment-free free effect functions. The optimal empirical value is 4.263 and the average empirical values when assigning one treatment are 1.503, 2.003, 2.502, and 2.001, respectively.
AD
AD_
AD_s
e
SABD
SABD_e
SABD_8
SABD e,
Original
20
3.511
3.621
3.581
3.790
3.553
3.730
3.566
3.879
(0.41)
(0.46)
(0.36)
(0.36)
(0.29)
(0.33)
(0.32)
(0.30)
200
40
3.258
3.223
3,415
3.558
3.430
3.450
3.503
3,704
(0.51)
(0.62)
(0.43)
(0.51)
(0.30)
(0.50)
(0.29)
(0.38)
60
3,289
3.210
3.382
3,412
3,304
3.396
3.391
3,564
(0.47)
(0.58)
(0.38)
(0.53)
(0.44)
(0.51)
(0.38)
(0.47)
20
3.992
4.079
4.009
4.165
3.881
4.143
3.851
4.188
(0.13)
(0.11)
(0.12)
(0.09)
(0.15)
(0.07)
(0.14)
(0.08)
600
40
3,970
4.027
3.982
4.135
3.888
4.115
3.874
4.182
(0.15)
(0.23)
(0.14)
(0.12)
(0.17)
(0.12)
(0.15)
(0.08)
60
3.942
4.037
3.918
4.094
3.848
4.079
3.826
4.147
(0.16)
(0.12)
(0.15)
(0.13)
(0.18)
(0.11)
(0.16)
(0.08)
20
4.086
4.166
4.094
4.211
3.942
4.200
3,893
4.224
(0.09)
(0.08)
(0.09)
(0.07)
(0.12)
(0.05)
(0.12)
(0.06)
1000
40
4.073
4.150
4.061
4.205
3,931
4.183
3.892
4.219
(0.11)
(0.08)
(0.11)
(0.07)
(0.12)
(0.06)
(0.12)
(0.06)
60
4.069
4.150
4.050
4.190
3,941
4.177
3,900
4.212
(0.10)
(0.07)
(0.10)
(0.11)
(0.14)
(0.06)
(0.13)
(0.07)
Augmented
20
3.913
3.922
3.894
3.970
3.964
3.969
3.923
3.993
(0.17)
(0.18)
(0.18)
(0.17)
(0.16)
(0.17)
(0.20)
(0.18)
200
40
3.828
3.827
3.824
3.879
3.877
3.874
3.869
3.927
(0.22)
(0.23)
(0.21)
(0.25)
(0.17)
(0.20)
(0.21)
(0.21)
60
3.806
3.739
3.737
3.789
3,847
3.817
3.793
3.831
(0.22)
(0.35)
(0.27)
(0.24)
(0.20)
(0.24)
(0.31)
(0.11)
20
4.169
4.188
4.168
4,210
4.172
4.207
4.165
4.222
(0.06)
(0.05)
(0.06)
(0.06)
(0.06)
(0.05)
(0.06)
(0.05)
600
40
4.158
4.169
4.156
4.200
4.162
4.201
4.162
4.224
(0.07)
(0.06)
(0.07)
(0.06)
(0.06)
(0.05)
(0.06)
(0.06)
60
4.144
4.153
4.135
4.175
4.154
4.186
4.146
4.201
(0.08)
(0.08)
(0.07)
(0.08)
(0.07)
(0.06)
(0.07)
(0.07)
20
4.208
4.223
4.207
4,239
4.197
4.235
4.191
4.242
(0.05)
(0.05)
(0.05)
(0.05)
(0.05)
(0.04)
(0.05)
(0.04)
1000
40
4.205
4.215
4.201
4.232
4.196
4.228
4.188
4.238
(0.06)
(0.05)
(0.06)
(0.05)
(0.06)
(0.05)
(0.06)
(0.05)
60
4.204
4.214
4.198
4.230
4.194
4.232
4.187
4,240
(0.05)
(0.05)
(0.05)
(0.05)
(0.05)
(0.04)
(0.06)
(0.04)
42

---

Table 15: Average empirical values and their standard deviations (in parenthesis) in the observational study Case 2: Linear interaction functions with complicated treatment- feeeee effect functions. The optimal empirical value is 9,433 and the average empirical values when assigning one treatment are 6.672, 7.173, 7.672, and 7.171, respectively.
AD
AD_
AD_s
AD_e,
SABD Original
SABD_e
SABD_s
SABD_e
20
7.592
7.598
7.608
7.696
7.561
7.736
7.715
7.977
(0.73)
(0.74)
(0.75)
(0.78)
(0.67)
(0.80)
(0.72)
(0,87)
200
40
7.462
7.452
7.540
7.554
7.462
7.497
7.647
7.671
(0.60)
(0.66)
(0.67)
(0.67)
(0.64)
(0.58)
(0.68)
(0.69)
60
7.503
7.531
7.568
7.538
7,380
7.612
7.645
7.815
(0.58)
(0.61)
(0.58)
(0.65)
(0.68)
(0.66)
(0.63)
(0.76)
20
7.798
7.824
7.840
7.970
8.147
8,948
8.197
9.024
(0.70)
(0.79)
(0.72)
(0.84)
(0.46)
(0.58)
(0.45)
(0.61)
600
40
7.732
7.778
7.787
7.841
7.954
8.610
8.142
8.923
(0.70)
(0.79)
(0.75)
(0.77)
(0.54)
(0.74)
(0.47)
(0.58)
60
7.615
7.599
7.625
7.666
7.953
8.413
8.125
8.875
(0.66)
(0.71)
(0.68)
(0.73)
(0.59)
(0.85)
(0.50)
(0.70)
20
7.712
7.833
7.731
7,823
8.191
9.122
8.166
9.084
(0.77)
(0.81)
(0.79)
(0.87)
(0.47)
(0.59)
(0.48)
(0.60)
1000
40
7.748
7.787
7.797
7.958
8.163
9.047
8.206
9.180
(0.75)
(0.81)
(0.74)
(0.83)
(0.42)
(0.51)
(0.43)
(0.43)
60
7.873
7.967
7.931
8.113
8.151
9.143
8.248
9.211
(0.72)
(0.84)
(0.74)
(0.88)
(0.49)
(0.55)
(0.47)
(0.51)
Augmented
20
7.770
7.887
7.838
8.042
8.667
8.711
8,859
8.936
(0.79)
(0.81)
(0.82)
(0.90)
(0.68)
(0.67)
(0.58)
(0.58)
200
40
7.621
7.713
7.680
7,854
8.240
8.308
8.703
8.802
(0.78)
(0.77)
(0.81)
(0.88)
(0.79)
(0.78)
(0.65)
(0.68)
60
7.585
7.693
7.629
7,812
8.169
8.146
8.647
8.680
(0.65)
(0.69)
(0.65)
(0.72)
(0.68)
(0.66)
(0.65)
(0.65)
20
7.920
8.102
7.956
8.263
9.315
9.340
9.343
9.380
(0.81)
(0.84)
(0.85)
(0.91)
(0.35)
(0.34)
(0.34)
(0.33)
600
40
7.792
7.962
7.867
8.103
9.194
9.211
9.311
9.342
(0.80)
(0.85)
(0.80)
(0.86)
(0.45)
(0.41)
(0.34)
(0.34)
60
7.630
7.743
7.676
7.845
9.153
9.159
9.277
9.308
(0.68)
(0.74)
(0.74)
(0.84)
(0.43)
(0.45)
(0.34)
(0.34)
20
7.756
7.949
7.781
8.043
9,318
9.341
9.347
9,371
(0.87)
(0.91)
(0.89)
(0.92)
(0.46)
(0.41)
(0.41)
(0.41)
1000
40
7.847
7.997
7.885
8.034
9.292
9.318
9.342
9.359
(0.87)
(0.85)
(0.89)
(0.87)
(0.41)
(0.38)
(0.39)
(0.38)
60
7.948
8.076
8.006
8.246
9.343
9.346
9.373
9,393
(0.80)
(0.85)
(0.84)
(0.85)
(0.41)
(0.42)
(0.40)
(0.41)
43

---

Table 16: Average empirical values and their standard deviations (in parenthesis) in the observational study Case 3: lineir interaction functions with simple treatment- free effect functions. The optimal empirical value is 1.422 and the average empirical values when assigning one treatment are 0.042, -0.920, -0.419, and -0.920, respectively.
AD
AD_e
AD_s
AD_e,
SABD Original
SABD_e
SABD_s
SABD_e
20
0.029
-0.014
0.113
0.142
0.312
0.178
0.405
0.411
(0.36)
(0.42)
(0.41)
(0.46)
(0.35)
(0.49)
(0.38)
(0.58)
200
40
0.091
0.058
0.244
0.151
0.258
0.149
0.415
0.335
(0.33)
(0.38)
(0.28)
(0.42)
(0.40)
(0.53)
(0.36)
(0.59)
60
0.014
0.010
0.110
0.129
0.168
0.031
0.349
0,224
(0.32)
(0.32)
(0.37)
(0.38)
(0.37)
(0.45)
(0.34)
(0.56)
20
0.482
0.660
0.607
0.910
0.773
1.008
0.721
T.095
(0.29)
(0.36)
(0.27)
(0.28)
(0.16)
(0.14)
(0.15)
(0.16)
600
40
0,420
0.592
0.567
0.802
0.717
0.927
0.720
1.076
(0.25)
(0.28)
(0.23)
(0.29)
(0.18)
(0.18)
(0.16)
(0.17)
60
0.437
0.546
0.612
0.796
0.718
0.907
0.737
1.009
(0.27)
(0.28)
(0.22)
(0.29)
(0.20)
(0.20)
(0.15)
(0.19)
20
0.662
0.883
0.772
1.077
0.832
1.72
0.744
1.132
(0.22)
(0.25)
(0.19)
(0.17)
(0.15)
(0.13)
(0.13)
(0.14)
1000
40
0.577
0.805
0.721
1.008
0,822
1.070
0.764
1.104
(0.21)
(0.22)
(0.17)
(0.17)
(0.14)
(0.12)
(0.12)
(0.20)
60
0.533
0.795
0.706
0.964
0.812
1.065
0.781
1.108
(0.21)
(0.21)
(0.16)
(0.21)
(0.13)
(0.13)
(0.12)
(0.15)
Augmented
20
0.521
0.559
0.647
0.783
0.766
0.759
0.862
1.018
(0.35)
(0.36)
(0.30)
(0.35)
(0.26)
(0.36)
(0.19)
(0.19)
200
40
0.559
0.544
0.684
0.763
0.715
0.713
0.830
0.922
(0.31)
(0.38)
(0.25)
(0.32)
(0.27)
(0.30)
(0.20)
(0.24)
60
0.353
0.342
0.475
0.558
0.527
0.474
0.745
0.798
(0.37)
(0.38)
(0.35)
(0.35)
(0.35)
(0.42)
(0.24)
(0.26)
20
0.987
1.071
1.025
1.146
1.118
1.188
L.099
1.246
(0.15)
(0.12)
(0.14)
(0.10)
(0.12)
(0.09)
(0.12)
(0.07)
600
40
0.960
1.029
1.006
1.107
1,092
1.157
1.086
1.225
(0.14)
(0.14)
(0.12)
(0.11)
(0.13)
(0.10)
(0.13)
(0.09)
60
0.980
1.030
1.019
1.094
1.106
1.149
1.107
1.198
(0.14)
(0.14)
(0.13)
(0.14)
(0.11)
(0.09)
(0.11)
(0.10)
20
1,090
1.144
1.111
1.191
1,179
1.228
1.124
1.266
(0.10)
(0.11)
(0.09)
(0.08)
(0.10)
(0.07)
(0.11)
(0.07)
1000
40
1.061
1.121
1.088
1.152
1.176
1.201
1.150
1.260
(0.10)
(0.09)
(0.09)
(0.14)
(0.10)
(0.09)
(0.10)
(0.08)
60
1.046
1.122
1.072
1.159
1.151
1.212
1.133
1,257
(0.11)
(0.09)
(0.10)
(0.08)
(0.11)
(0.08)
(0.10)
(0.07)
44

---

Table 17: Average empirical values and their standard deviations (in parenthesis) in the observational study Case 4: Non-linear interaction functions with a complicated treatment- free effect function. The optimal empirical value is 5.623 and the average empirical values when assigning one treatment are 4. 4.243, 3.282, 3. 3.782, and 3.281, respectively.
AD
AD_
AD_s
AD_e,
SABD Original
SABD_e
SABD_8
SABD_e,
20
4.034
3.954
4.062
4.006
4.309
4.048
4.381
4.316
(0.54)
(0.59)
(0.56)
(0.57)
(0.55)
(0.70)
(0.53)
(0.71)
200
40
3,978
3.883
4.001
3,908
4.219
4.023
4.372
4.184
(0.58)
(0.62)
(0.59)
(0.62)
(0.52)
(0.64)
(0.51)
(0.72)
60
4.019
3,926
4.038
3,879
4.198
4.007
4.287
4.056
(0.50)
(0.53)
(0.52)
(0.53)
(0.50)
(0.59)
(0.54)
(0.65)
20
4.163
4.T03
4.191
4,215
4.723
4.992
4.745
5.132
(0.53)
(0.61)
(0.56)
(0.69)
(0.35)
(0.48)
(0.37)
(0.45)
600
40
4.096
4.060
4.131
4.132
4,606
4.785
4.703
5.020
(0.58)
(0.60)
(0.62)
(0.69)
(0.41)
(0.54)
(0.41)
(0.48)
60
4.048
3.980
4.074
4.051
4.619
4.699
4.707
5.042
(0.58)
(0.64)
(0.60)
(0.69)
(0.39)
(0.62)
(0.38)
(0.48)
20
4.047
4.086
4.071
4.089
4.762
5.150
4.764
5.180
(0.63)
(0.68)
(0.63)
(0.69)
(0.41)
(0.48)
(0.41)
(0.47)
1000
40
4.130
4.103
4.186
4.158
4.727
5.091
4.743
5.236
(0.62)
(0.70)
(0.66)
(0.77)
(0.38)
(0.44)
(0.38)
(0.44)
60
4.093
4.029
4.128
4.168
4.742
5.117
4.792
5.239
(0.62)
(0.66)
(0.64)
(0.71)
(0.42)
(0.47)
(0.40)
(0.44)
Augmented
20
4.146
4.125
4.193
4.337
4.645
4.576
4.859
4.973
(0.59)
(0.68)
(0.62)
(0.72)
(0.60)
(0.69)
(0.50)
(0.51)
200
40
3.975
4.014
4.003
4.099
4,441
4.344
4.795
4.765
(0.66)
(0.68)
(0.67)
(0.69)
(0.56)
(0.64)
(0.49)
(0.56)
60
3.976
3.957
4.019
4.056
4.276
4.160
4.657
4.577
(0.57)
(0.58)
(0.60)
(0.64)
(0.58)
(0.59)
(0.55)
(0.65)
20
4.188
4.267
4.231
4.380
5,266
5.291
5.293
5.396
(0.63)
(0.64)
(0.64)
(0.69)
(0.37)
(0.36)
(0.36)
(0.35)
600
40
4.141
4.216
4.212
4.331
5.147
5.106
5.298
5,350
(0.70)
(0.70)
(0.74)
(0.76)
(0.47)
(0.49)
(0.37)
(0.33)
60
4.061
4.063
4.124
4.186
5.135
5.124
5.297
5,333
(0.66)
(0.71)
(0.69)
(0.73)
(0.47)
(0.48)
(0.37)
(0.37)
20
4.043
4.154
4.072
4,207
5.341
5.346
5,346
5,405
(0.69)
(0.74)
(0.72)
(0.76)
(0.46)
(0.47)
(0.41)
(0.42)
1000
40
4.149
4.198
4.204
4.302
5.307
5.310
5.352
5,384
(0.71)
(0.72)
(0.71)
(0.74)
(0.40)
(0.38)
(0.39)
(0.37)
60
4.086
4.131
4.111
4.290
5.349
5.371
5.388
5,433
(0.66)
(0.71)
(0.68)
(0.72)
(0.41)
(0.41)
(0.40)
(0.41)
45

---

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 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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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 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 sell-specification, asymptotic approximations only holding locally). The authors should reflect 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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, worse outcome might 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 impor- tant 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 complete (and correct) proof?
46

---

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 cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide 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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
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 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 novel architecture, describing the architecture fully might suffice, or if the contribution 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 accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions 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) the contribution is primarily 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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing or verifying the results.
5. Open access to data and code
47

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide open access to the code, with sufficient instructions faithfully reproduce the main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr
(eglli//gidee/Codeeeee
cwteimissssooooliiii onPolicy) icy for more details.
While encourage the release of code and data, 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wctpublic//aeccccccccccccccc gwwpi/gg///////eeee// /Ce///eeeubbiiicc wetsSuuumissionn ssionPolicy) 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 proposed method and baselines. only 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 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, hyper- parameters, 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, hyperparame- ters, 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 the core the paper to 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
48

---

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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 compute required for each of the individual experimental runs as well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: : /ttp//neripppp//epee
Answer: [Yes]
Justification:
We strictly follow the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations 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.
49

---

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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include URI.
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.
50

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification: We introduce our proposed with sufficient details in the paper. The documenta- tion 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/mmodel 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included 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 crowdsoureing nor research with human subjects.
51

---

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.
52


Reviews:
[{"review_id": "AwulKagy0X", "reviewer": "Reviewer_1Znf", "summary": "The 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.", "questions": "1. 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?", "limitations": "Addressed in the paper", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper combines AD-learning with other useful techniques to improve treatment rule estimation and have theoretical guarantees.", "weaknesses": "I 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Weakness 1,2:** Thank you for the feedback. Our work views the ITR framework as a weighted convex optimization problem, focusing on using robust weights and the PGD algorithm to find sparse solutions effectively, a novel approach in this context. We believe that these contributions are significant for the NeurIPS community.\\\nDue to the comprehensive analysis and detailed presentation required, some details including related work were placed in the appendix to meet the page limit. In the revision, we will include a related work section in the main text and aim to make the paper more self-contained while emphasizing our key contributions.\n\n>**Weakness 3:** We appreciate your comments.\\\n>First, our computational guarantee of the proposed PGD algorithm does rely on a general convergence analysis for PGD (Lem.D.2). However, when applying such results to the weighted convex optimization problems that reformulate the ITR-Learning problems, our analysis reveals that the convergence rate of PGD depends crucially on the eigenvalues of the weighted design matrix $\\boldsymbol{\\Psi}$ in eq.8, a new concept introduced in this work. Our analysis shows that the Hessian of the ITR-Learning problem is given by $\\boldsymbol{\\Psi}$. Thus, choosing the weights to maximize the minimum eigenvalue of $\\boldsymbol{\\Psi}$ ensures the best possible convergence rate for the PGD algorithm. This provides a novel insight into choosing the optimal weights. Specifically, the minimum eigenvalue of the weighted design matrix is maximized when the weights are *distributional covariate balancing* where the weighted sample covariance matrices for each treatment are approximately the same. In cases where the covariates are discrete and one-hot encoded, the optimal weights exactly balance the covariate distributions given treatments. We will clarify these points in the revision.\\\n>Second, our generative models for ITR-Learning are not restrictive. It is common in statistical literature to assume a probabilistic model for the model parameter, ensuring that the model is identifiable and estimable under reasonable conditions. The resulting problems from our generative models correspond exactly to the weighted convex optimization formulation of ITR-Learning in eq.3-4.\\\n>The reviewer may find the linear decision function class restrictive. However, the linear function class, in principle, contains the polynomial function class. By taking the covariate powers up to M, the linear function class becomes the class of degree-M polynomial functions. Also, the linear function class can serve as a useful approximation to nonlinear decision functions. This linear approximation captures essential patterns in the data and makes interpretable treatment decisions, crucial for practitioners.\\\n>Additionally, our statistical guarantee introduces new techniques and results to the ITR literature. We provide non-asymptotic estimation guarantees using uniform concentration inequalities and Berry-Esseen theorem, which give explicit bounds on the estimation error for a given finite sample size. We also unify our computational and statistical guarantees to obtain joint computational and sample complexity results, demonstrating how the true model parameters can be recovered.\\\n>We will incorporate these points in the revision.\n\n>**Originality:** We kindly ask you to refer to the global response.\n\n>**Q1:** Yes, the second \"=\" in equation (2) can be derived using standard causal assumptions: positivity, conditional ignorability, and consistency.\n\n>**Q2:** Thank you for pointing that out. The vector $u$ is a random treatment vector in $\\mathbb{R}^K$, corresponding to the random treatment A. We will add the definition in the revision.\n\n>**Q3:** The sum-to-zero constraints hold for $\\delta_k(x)$ for all $x$ to ensure model identifiability. However, estimating K different decision functions that satisfy the sum-to-zero constraint is challenging. An alternative approach is to use AD-Learning framework, which assumes $\\sum_k u_k=0$. It leads to $\\sum_k u_k^T f(x)=0$ for any given decision function $f$ and covariate $x$.\n\n>**Q4:** The updated Figure 1 can be found in the response pdf.\n\n>**Q5:** Thank you for your questions and suggestions. First, the left subplot highlights the role of energy weights. Even in RCT, finite sample imbalances in covariates can occur. Energy weights help reduce these imbalances, leading to observed improvements. Second, for comparison \"observational vs RCT+linear\", the table below shows the mean accuracies (n=1000, 100 rep). \n>\n|p|AD (Obs)|AD+e(Obs)|AD(RCT)|AD+e(RCT)|\n|-|-|-|-|-|\n|20|0.844|0.905|0.865|0.895|\n|40|0.830|0.893|0.854|0.900|\n|60|0.822|0.894|0.857|0.895|\n>\n>The performance improvement with energy weights compared to IPW is greater in the observational study within the same rule. This comparison directly assesses the effect of weights in less challenging scenarios. We will include these results in the revision. Lastly, regarding the comparison \"observational+linear vs nonlinear\", we have already conducted simulation studies in Appendix Sec.H (Case 1 vs 3).\n>\n|p|AD(nonlinear)|AD+e(nonlinear)|AD(linear)|AD+e(linear)|\n|-|-|-|-|-|\n|20|0.471|0.572|0.782|0.845|\n|40|0.440|0.542|0.774|0.831|\n|60|0.419|0.531|0.764|0.825|\n>\n>The performance improvement with energy weights compared to IPW is greater under the nonlinear rule, indicating the robustness of energy weights.\n\n>**Q6:** Thank you for your question. The high SE is due to the dataset's variability. We used a small portion of the 2500 data points for training, as done in [43], mimicking clinical trials with small sample sizes. Even with 1200, it's still less than 50% of the total dataset.\\\n>Increasing the training set generally decreases SE, but if the added data has more noise, SE may increase. Since this pattern is not specific to any method, as also observed in [43], it suggests the issue is related to the data nature rather than a particular method.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9De5UbtNvs", "reviewer": "Reviewer_3MV8", "summary": "The 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.", "questions": "-\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.", "limitations": "The authors have discussed the limitations of the proposed method.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "-\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.", "weaknesses": "-\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Response for Weakness 1:**  \nThank you for the suggestion. We updated Figure 1 with an error bar with the standard error of the mean in the one-page author response.\n\n> **Response for Weakness 2:**  \nThank you for the insightful comment. Instead of reporting the standard deviation, we will report the standard error of the mean empirical value in the revised version, which is a convention in the machine learning literature (see the attached table in the one-page author response). Although some values still fall within the confidence interval, the overall mean empirical values of our method are higher than those of the baseline methods. The high standard deviation is mainly due to the variability of the dataset. In our analysis, only a small portion of the data out of 2500 was used for the training set, following the setting in [43], which mimics clinical trials (typically characterized by small sample sizes).\nWe want to emphasize the importance of computational efficiency in addition to empirical value. Our proposed methods are significantly faster than two popular methods: *policytree* and *causalDML*. Specifically, for analyzing one simulated dataset with 60 covariates and 1000 samples, our methods use only 9.4% of the computational time required by *causalDML* and 36.6% of the time needed by *policytree*.\nWe recognize that real data is complex, and it is unrealistic to expect a single method to outperform all others across every dataset. However, our extensive simulation studies clearly demonstrate the conditions under which our unified method performs well. We hope you also recognize the value of this comprehensive analysis.\n\n> **Response for Weakness 3:**  \nWe appreciate your valuable feedback. We will add figures to highlight the differences between methods.\n\n> **Response for Question 1:**  \nThank you for your question. In the left column of Figure 1, we used vanilla AD-Learning as a baseline to demonstrate that incorporating energy weights provides a performance gain in both randomized trial and observational study settings. To reduce possible confusion, we updated the middle column to compare the performance of each algorithm based on AD-Learning. \nFor your information, the \"proposed\" approach in the right column combines energy weights with additional techniques including variable screening, outcome augmentation, and inverse variance weighting. We explained it in lines 294-297, but we will further clarify it in the revised version. Thank you for your feedback.\n\n> **Response for Question 2:**  \nThank you for the comment. Energy balancing weights are derived from the concept of \"weighted energy distance\". These weights minimize the distance of the weighted empirical distributions of covariates across treatment groups, mitigating biases due to model misspecification of traditional inverse probability weighting (IPW) and resulting in more precise and reliable ITR estimation. We will include a high-level summary of section 2.2 to aid readers in understanding the key concepts.\n\n> **Response for Question 3:**  \nThank you for pointing that out. The vector $\\mathbf{u}$ is a random treatment vector in $\\mathbb{R}^{K}$, corresponding to the random treatment $A$. We will add the definition in the revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FVJxBpVeCL", "reviewer": "Reviewer_oq1w", "summary": "This 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.", "questions": "See 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?", "limitations": "The authors have aqequately addressed the limitations of their work.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The 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.", "weaknesses": "*  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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Response for Weakness 1:**\n>\n>Thank you for your valuable feedback. We agree that clarifying the scope of our work is essential.\nIn this paper, we focus on linear decision boundaries, which are a standard approach in statistical literature due to their interpretability. One significant advantage of using linear decision classes is the convexity of the associated optimization problem. This convexity allows us to obtain statistical estimation guarantees, recovering the true parameter with high probability under the MLE framework of the generative model. \n>\n>While our current focus is on linear decision boundaries, it is essential to note that our framework can be extended. For example, extending to polynomial decision boundaries would allow us to obtain similar optimization analyses, including computational and statistical guarantees. Additionally, our framework can handle more complex, non-convex decision functions, such as neural networks. In such cases, the problem becomes a non-convex constrained optimization problem, meaning finding a global minimization is not feasible. However, we can guarantee the convergence to the first-order stationary point.\n>\n>This highlights the flexibility of our framework and we can demonstrate different computational and statistical guarantees depending on various function classes. Although this was not included in the current paper, we recognize the importance of this point and will provide further discussion in the revised version. We sincerely appreciate your valuable comments and feedback.\n\n> **Response for Weakness 2:**\n>\n>We thank the reviewer for raising this point. We would like to highlight that the main innovation of our work is the significant performance improvement achieved by using an $L_1$-constraint, instead of the soft $L_1$-penalization (or mixed $L_1$ and $L_2$ regularization) commonly used in the statistical literature.\n>\n>We have theoretically justified this performance improvement. In the traditional statistical literature, incorporating hard constraints into the MLE framework was not feasible. However, we strengthened MLE analysis to enable statistical estimation even when the true parameter lies on the boundary of constraint. This is crucial because if the solution lies within the interior of $L_1$ ball, it implies that the solution is not sparse. For sparsity, the solution must lie on the boundary of $L_1$ ball. By combining probabilistic techniques with constrained optimization, we achieved significant improvements.\n>\n>Regarding the novelty (originality) of our work, we kindly ask you to refer to the global response. Unfortunately, we were unable to include all the details due to character limits. Thank you for your understanding.\n\n> **Response for Weakness 3:**\n>\n>Thank you for your valuable feedback. We appreciate your suggestion to include a more thorough literature review and comparisons with existing work on policy learning and treatment effect estimation. In the revised version of our paper, we will incorporate a comprehensive review of relevant literature, including the work by Nekipelov, Semenova, and Syrgkanis (2022), which proposed a novel approach to finding sparse solutions based on the Lasso estimator in high-dimensional settings and demonstrated that the estimator converges at the oracle rate under mild conditions. We will provide detailed comparisons between our proposed methods and those discussed in the literature to highlight our contributions.\n\n> **Response for Weakness 4:**\n>\n>Thank you for your insightful suggestion. The bias-variance tradeoff plot is given in the one-page author's response.\n>\n>Under Case 3, where the true optimal decision function is nonlinear and the covariates are 60-dimensional, we measure the reduction of bias using the difference between the true value functions of the optimal treatment decisions and the estimated treatment decisions, $V(d^\\text{opt}) - V(\\hat{d})$. The bias cannot be zero due to our model assuming a linear decision function.\n>\n>Except for the case with a sample size of 200, where additional optimization required for estimating energy balancing weights leads to an increase in bias, both models show a significant decrease in squared bias with only a minor increase in variance across both training and test sets. Notably, no tradeoff between bias and variance was observed for either model. The curves illustrate that using energy weights significantly reduces bias, which enhances the overall performance of the method. Thank you for your insight and valuable comment.\n\n> **Response for Weakness 5:**\n>\n>Thank you for this very interesting suggestion. We have included your proposed baseline model in Figure 1, labeled as \"Linear_Baseline\" in the response PDF. \n>\n>This model performs excellently in the left panel, where the scenario mimics a randomized trial with a true optimal linear rule. This is expected because, in a randomized trial, covariates are well-balanced across treatment groups, allowing for an accurate comparison between treatment levels. Additionally, the baseline model can accurately capture the true linear decision function, which explains its strong performance.\n>\n>However, in the right panel, where the simulation mimics an observational study with a true optimal nonlinear rule, the baseline model shows poor performance. This is due to two reasons: firstly, the linear model requires balancing weights, such as IPW or energy weights, to reduce confounding effects present in observational studies. Secondly, it cannot fully capture the nonlinear decision function.\n>\n>We believe this interesting observation highlights the need for appropriate weighting in ITR-Learning, especially in observational studies and the importance of model flexibility. We plan to further investigate this suggestion as it provides valuable insights into the factors that significantly impact ITR-Learning. Thank you again for your insightful suggestion.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f8ewahRYZQ", "reviewer": "Reviewer_kG5t", "summary": "This 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.", "questions": "1. 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?", "limitations": "Please see it in above weakness and questions section.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> We sincerely thank the reviewer for the extremely insightful comments. Due to space constraints, we had to cut half of our initial response. We would be happy to have a more in-depth discussion.\n\n>**W1:** We see that both of these works take a similar approach to ours by using weighting schemes and opting for directly optimizing the weights, rather than using estimated IPW. However, there are some crucial differences that we would like to point out.\\\n>[K18] has a general setup for policy evaluation and learning from historical data, where accurate policy evaluation through SAPE is needed for policy learning. The author proposes to optimize the weights directly by minimizing the worst-case posterior MSE. In contrast, our approach avoids the evaluation of policy effects. Instead, we estimate the optimal decision function directly through a weighted convex empirical loss minimization. Our approach introduces a simple spectral criterion for choosing the optimal weights: maximize the minimum eigenvalue of the weighted design matrix, which may be compared with the approach in [K18].\\\n>[K21] uses a similar approach to [K18] with the novel technique of \"retargeting\" the population covariate distribution. This idea has an interesting connection to our analysis. Under our generative model assumption, the true parameter $B_ *$ that parameterizes the optimal decision function is the global maximizer of the weighted log-likelihood function $E[w(X,A) \\log \\pi_B(Y|X,A)]$ for an *arbitrary* weighting function $w(X,A)$. So to minimize the statistical estimation error, we select the weight $w(X,A)$ such that the corresponding weighted Fisher information has the largest minimum eigenvalue. Despite this conceptual similarity, the retargeting weights in [K21] depend only on the covariates, whereas our weights also depend on the treatments.\n\n>**W2:** We kindly ask you to refer to the global response.\n\n>**W3:** Our proposed method takes 4.40 sec per dataset with p=60, n=1000 (glmnet for penalization takes average 2.01 sec). Other methods face scalability issues (12.02 sec for policytree and 46.75 sec for causalDML). We believe that computational efficiency is one of the advantages of our method.\n\n>**W4:** We fully agree that the same hyperparameter tuning strategy is crucial for a fair comparison between the PGD and the penalized method. In the revised experiments, we have implemented the penalized method using unconstrained subgradient descent to minimize the L1-penalized ITR loss function and unified the hyperparameter selection across both methods by using the value function. In Figure 1 (author response PDF), we observe that the performance of L1-penalization is improved, but PGD still shows overall better performance.\\\n>To further clarify, we have included Figure 2 (author response PDF). This figure compares the optimization trajectories for the multiple L1-penalization parameters and the L1-ball sizes using the same stepsizes $\\alpha/t$. For a large $\\lambda_1=10$ in L1-penalization, necessary for a sparse solution, the optimization trajectory becomes highly fluctuating, and the algorithm fails to achieve low values for the true (unregularized) ITR objective. This is because the large L1-penalization significantly perturbs the true ITR objective. In contrast, PGD maintains smooth optimization trajectories and effectively minimizes the true objective across a wide range of L1-ball sizes. This illustrates a key advantage of the constrained method over the penalized method.\n\n>**W5:** It is indeed $u$. The vector $u$ is a random treatment vector in $\\mathbb{R}^K$, corresponding to the random treatment $A$. We will add the definition in the revision.\n\n>**W6:** Since our framework is not restricted to binary treatments, *multi-category treatment* is included in the title. However, we appreciate your suggestion and will consider changing the title to better reflect the primary focus on effectiveness.\n\n>**Q1:** A revised form of our Thm.3.5 shows that with high probability, the statistical estimation error is $O(\\sqrt{p/n}/\\mu)$. Thus, in the high-dimensional setting $p= \\alpha n$, our bound on the statistical error is non-vanishing. As long as $\\alpha = p/n$ is small and the strong convexity parameter $\\mu$ stays bounded away from zero, our result shows that the statistical error is also small. However, we believe that a more tailored analysis in the high-dimensional setting (e.g., high-dimensional concentration inequalities) could improve our current error bound. There is some difficulty in the computational complexity of PGD since the largest eigenvalue of $\\Psi$ can be order $p$. We are happy to discuss further.\n \n>**Q2:** We use random forests to estimate the treatment-free effect. As the reviewer correctly points out, potential model misspecification can significantly impact the ITR performance. Additional simulations, case 2 and 4, in Sec H show this situation. To address this issue, we combine outcome augmentation with inverse variance weighting to reduce heteroscedastic errors, resulting in more accurate results.\n\n>**Q3:** We considered two simulations with p=60, n=200: one with linear treatment-free/interaction effects involving four covariates, and the other with nonlinear effects involving two covariates. The goal of variable screening is to retain as many relevant variables as possible since we can rely on the PGD for a sparse solution. Thus, our primary interest is in the true positive rate (TPR), the proportion of true covariates correctly identified by screening. With 100 repetitions, for case 1, the mean(sd) TPR is 0.99(0.05), FPR is 0.151(0.05), and accuracy is 0.86(0.04). For case 2, the mean(sd) TPR is 0.995(0.05), FPR is 0.178(0.05), and accuracy is 0.83(0.05). Screening performs well, retaining 99% of true covariates with low FPR. We conjecture that group-lasso may further improve our proposed method by lowering the FPR. We will include this discussion in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear Reviewers,

Thank you for reviewing our work and providing valuable comments. We appreciate the time and effort you have taken to provide feedback on our paper.

In response to your suggestions, we have made the following updates to Figure 1 in the one-page author's response:

- **Error bar**: We have updated Figure 1 to include error bars, which provide a clearer representation of the variability in the results.

- **Improved performance of penalized approach**: We identified ways to improve the performance of the penalized approach to ensure a fair comparison with our proposed method. We have implemented these improvements and updated the corresponding results. We would like to extend our special thanks to reviewer kG5t for the valuable comments regarding the performance comparison. The performance gap between the constrained and penalized approaches has decreased but remains significant, highlighting the advantages of our proposed method.

Regarding the **originality** of our work, while it is true that our work builds on existing components such as balancing weights and AD-Learning, our contribution lies in providing a unified framework that combines these approaches in an effective way. We demonstrated theoretical justifications for the combined approach, which, to our knowledge, have not been previously established.

For example, in the appendix, we prove that outcome augmentation with our suggested balancing weights yields an estimator with minimum variance. Previous research has shown that under traditional inverse probability weighting, the model parameter after outcome augmentation achieves the minimum solution, but there were no results for other types of balancing weights. Our work extends these results to a more general setting, providing new insights and theoretical guarantees for the use of balancing weights.

Additionally, our unified framework allows for a more comprehensive understanding of the interplay between different components, leading to improved performance and robustness. We believe that these theoretical contributions provide significant value and advance the field beyond simply combining existing methods.

We also appreciate Reviewer **kG5t** for raising the following insightful question: 

>* 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?

Our theoretical analysis does indeed justify the use of distributional covariate balancing weights. Our optimization landscape analysis demonstrates that the strong convexity and smoothness of the convex optimization formulation for the ITR problem depend crucially on the following weighted design matrix $\Psi$ in Eq. (8). While we are free to choose the weights, our analysis provides a guiding principle for selecting these weights. Specifically, the strong convexity parameter of the ITR landscape is the minimum eigenvalue of the expected design matrix $E[\Psi]$, which we desire to be as large as possible. Namely, our analysis shows that to achieve an improved convergence rate for the PGD algorithm (Thm. 3.3) and a smaller statistical estimation error (Thm. 3.5), we need to choose the weights to maximize the minimum eigenvalue of $\Psi$. We acknowledge that this point was not clearly conveyed in our concise statement of Thm. 3.5 in the initial submission, where many important constants were suppressed in $O(\cdot)$ notations. In the revision, the high-probability bound (12) in Thm. 3.5 will read as 
\begin{align}
\lVert  B_\star - B_T \rVert_F \le \frac{C\sqrt{(p/n)\log \epsilon^{-1}} }{\mu} + \frac{8 \lambda_2 \lVert B_\star \Vert_F  }{\mu},
\end{align}
where $\mu$ is essentially the minimum eigenvalue $\lambda_{\min}(E[\Psi])$ of the expected weighted design matrix (see Eq. (10)). The number $n$ of the required sample size must be large enough to satisfy $\frac{(\sum_{i=1}^{n} w_{i}^{2})^{3}}{(\sum_{i=1}^{n} w_{i}^{3})^{2}} \ge C_{1} \epsilon^{-2}$ for some explicit constant $C_{1}>0$. The first term in this error bound represents the statistical error and the second term represents the bias introduced by the L2-regularization. Notice that both terms are proportional to $1/\mu$, which is essentially $1/\lambda_{\min}(E[\Psi])$. Thus, choosing the weighting function $w(A,X)$ to maximize $\lambda_{\min}(E[\Psi])$ helps minimize the overall statistical estimation error. Another interesting point we would like to add is that when maximizing $\lambda_{\min}(E[\Psi])$ encourages us to choose highly heterogeneous balancing weights, we get increased sample complexity through the bound $\frac{(\sum_{i=1}^{n} w_{i}^{2})^{3}}{(\sum_{i=1}^{n} w_{i}^{3})^{2}} \ge C_{1} \epsilon^{-2}$ (e.g., if $w_{i}\equiv 1$, this is $n\gtrapprox \epsilon^{-2}$, but if $w_{1}=n,w_{2}=\dots=w_{n}=0$, this may not be satisfied for any $n$).

Very interestingly, the minimum eigenvalue of the weighted design matrix is maximized when the weights are `distributional covariate balancing', making the empirical covariate distributions conditional for each treatment approximately the same. In the special case where the covariates are discrete and one-hot encoded, the optimal weights are those that exactly balance the covariate distributions given treatments (we can provide a concrete example). Therefore, maximizing the minimum eigenvalue of the weighted design matrix $\Psi$ provides an explicit spectral condition for obtaining effective distributional covariate balancing weights. We will elaborate on this point and provide further explanations in the revision. 

We hope these updates address your concerns and further clarify the contributions and findings of our work. Thank you once again for your insightful feedback.

Meta-review: This paper presents an estimation framework for individualized treatment rules with multi-category treatments based on distributional covariate balancing weights. Most reviewers praised this work for its robust theoretical analysis and extensive numerical studies, which convincingly demonstrate the effectiveness of the proposed method.

While initial concerns were raised regarding the novelty of the submission as several building blocks of the framework are not new, the authors effectively addressed them in their rebuttal. They clarified that the key contribution lies in the innovative combination of these elements, coupled with a rigorous theoretical justification and understanding. In the discussion period, all reviewers were satisfied with the authors' responses and raised their scores. Considering the rebuttal and the discussion, I recommend acceptance.
Decision: Accept (poster)