Treatment Effect Estimation for Optimal Decision-Making
Authors: Dennis Frauen
Valentyn Melnychuk
Jonas Schweisthal
Mihaela van der Schaar
Stefan Feuerriegel
Venue: NeurIPS 2025 poster ()
Keywords: treatment effect estimation
policy learning
causal inference
Primary area: probabilistic_methods
TL;DR: We show that two-stage CATE estimators can be suboptimal for decision-making and propose a novel method to improve them.

Abstract: Decision-making across various fields, such as medicine, heavily relies on conditional average treatment effects (CATEs). Practitioners commonly make decisions by checking whether the estimated CATE is positive, even though the decision-making performance of modern CATE estimators is poorly understood from a theoretical perspective. In this paper, we study optimal decision-making based on two-stage CATE estimators (e.g., DR-learner), which are considered state-of-the-art and widely used in practice. We prove that, while such estimators may be optimal for estimating CATE, they can be suboptimal when used for decision-making. Intuitively, this occurs because such estimators prioritize CATE accuracy in regions far away from the decision boundary, which is ultimately irrelevant to decision-making. As a remedy, we propose a novel two-stage learning objective that retargets the CATE to balance CATE estimation error and decision performance. We then propose a neural method that optimizes an adaptively-smoothed approximation of our learning objective. Finally, we confirm the effectiveness of our method both empirically and theoretically. In sum, our work is the first to show how state-of-the-art CATE estimators can be adapted for optimal decision-making.

Treatment Effect Estimation for Optimal Decision-Making
Dennis Frauen"
Valentyn Melnychuk
LMU Munich
LMU Munich
Munich Center for Machine Learning
Munich Center for Machine Learning
Jonas Schweisthal
Mihaela van der Schaar University of Cambridge
LMU Munich
Munich Center for Machine Learning
Stefan Feuerriegel
LMU Munich
Munich Center for Machine Learning
Abstract
Decision-making in various fields, such medicine, is heavily based on conditional average treatment effects (CATEs). Practitioners commonly make decisions by checking whether the estimated CATE is positive, even though the decision decision-making performance modern CATE estimators (meta-learners) is poorly understood. this paper, study optimal decision-making based on two-stage meta-learners (e.g., DR-learner), which estimate CATE via a second-stage regression. We show that these meta-learners common settings where the second-stage regression is over a restricted function class (e.g., when using regularization or employing fairnessinterpretabillty con- straints). Intuitively, this occurs because such estimators prioritize CATE accuracy in regions far away from the decision boundary, which is ultimately irrelevant decision- decision-making. As a remedy, we propose a novel two-stage learning objective that re-targets the CATE balance CATE estimation error and decision perfor- mance. We then propose a neural method that optimizes an adlptively-smoonee approximation of our learning objective. Finally, confirm the effectiveness of our method both empirically and theoretically.
be suboptimal when used for decision-making in
Introduction
Data-driven decision- across various fields, such as medicine 19), public policy [ 39 and marketing ( relies understanding how treatments affect different individuals and groups. This heterogeneity in the treatment effect across individuals typically quantified through the conditional average treatment effect (CATE). Then, common approach from practice obtain decisions from CATE is thresholding: individuals with positive CATE receive freafment, while those with negative CATE do not (13). For example, in medicine, clinicians typically administer treatments the subset of patients who are expected to benefit from the intervention 37
However, despite being widely used in practice, the optimality properties of such thresholding approach for decision-making are unclear. We argue that minimizing estimation error and optimizing
"Correspondence to: frauen Imu.de
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

for decision-makee performance are inherently different objectives. Existing literature acknowledges the distinction between CATE estimation error and decision making performance (14) and draws connections between the two in specific situations 7
Nevertheless, the optimality of makiing based on modern CATE estimators remains unclear, and approaches are missing for how to improve thresholding-
13 12
decision rules.
In this paper, we study optimal decision-making based on two-stage meta-learners. Two-stage meta-learners estimate CATE by employing a second-stage regression over prespecified function class. They are widely used in practice and include orthogonal learners such as the DR-leamnr (55) 33). We show theoretically that, while these methods may be optimal for CATE estimation, they can lead to suboptimal decisions when combined with thresholding approach, particularly when the second-stage model class is restricted. This is crucial in various real-world applications, in which the CATE model class is often restricted, e-g., due to fairness constraints
35
Why are two-stage learners suboptimal for decision-making?? ? Fig. [gg.[showw the results of one of our experimental setups (details in Sec. D). where the ground- truth CATE (red line) is positive everywhere, except for a small region of the covariate space. An optimal policy one that administers the treatment everywhere, except for that region (because the treatment is harmful in that region). Further, we show three two-stage learners with regularized model class: the estimator in blue (- 0) achieves the lowest CATE estimation error, but yields the wrong policy in the region of negative CATE, In contrast, the estimators shown in violet (generated by our method that we propose later) is preferred for decision-making it sacrifice a small amount CATE accuracy to yield better downstream decision performance, so that the decisions coincide with thresholding the ground-trun CATE.
Intuition:
Lagand
Deciue bunday
PTCATT yn PT/AATC 2-0
Figure 1: Illustrative example show- ing the suboptimality of CATE estima- tion for decision-making. The dotted lines show regularized two-stage CATE estimators. The blue line corresponds standard two-stage CATE estimation, while the green and violet lines are gen- erated by our method. The parameter yy quantifies the trade-ofi off between CATE estimation error and decision-making performance. Details are in Sec. 5
To address the shortcomings of the thresholding approach from above, we propose a novel second-stage learning objective that re-targets CATE to balance CATE estimation error and decision performance. By doing so, we re- target the CATE to new estimand which we call policy- targeted CATE (PT-CATE). Our PT-CATE can still approximately interpreted as CATE, while leading to superior policies. We further propose neural method to optimize our objective and estimate the corresponding PT-CATE. Our method follows a three-step procedure: first, we learn a neural network to estimate CATE; second, we learn a separate neural network to identify regions where incorrect decisions made; and finally, we adjust the first neural network to improve decision decision-making in these problematic regions.
Our contributions 2 are: (i) We show when two-stage CATE estimators with restricted model-classes can be suboptimal for decision-maie making. (ii) We develop a novel two-stage learning objective that effectively balances CATE accuracy and decision performance. For this, introduce a new estimand for policy-targeted CATE estimation (PT-CATE), (iii) We propose neural method for our learning objective with theoretical guarantees and empirically demonstrate its ability to effectively trade-off CATE estimation error and decision makinng performance.
2 Related work
Our work connects to several literature streams below (see Appendix for an extended version).
CATE estimation. Methods for CATE estimation can be broadly categorized into (i) model-based and (ii) model-agnostic approaches. (i) Model-based methods propose specific models, such as
regression forests 21 59 tailored neural networks 23 50 for CATE estimation. (ii) Model-
"Code is available at huttss//ggtttt..cmmmmmmm/s//tsee

---

agnostic methods (also called meta-learners) are "recipes" for constructing CATE estimators that can be combined with arbitrary machine learning models (e.g., neural networks) 38 9.
Meta-learners often follow a two-stage approach to account for the fact that the CATE is often structurally simpler than its components (response functions) [33] Prominent examples of two-stage meta-learners include the RA/X-learner (38), IPW-learner [9], and the DR-Iearner [55] 33 The DR-learner has the additional advantage of being Neyman-orthooonall and thus being provably robust against estimation errors from the first-stage regression [8 [6]. In this paper, we show that such two-stage learners can be suboptimal for decision-male making when using restricted second-stage model classes and propose method to improve them.
(Direct) Off-policy learning (OPL). The goal in OPL is directly learn optimal policy by maximizing the so-called policy value. Approaches for estimating the policy value from data follow three primary approaches: (i) the direct method (DM) 477 leverages estimates of the response functions; (ii) inverse propensity weighting (IPW) (51] re- re-weights the data such that they resemble
samples under the evaluation policy; and (iii) the doubly robust method (DR) [I] 2] combines both. Recent work has focused on enhancing finite-sample performance through techniques such as
reweighting [2 25 and targeted maximum likelihood estimation (5). Further, extended versions have
been developed for specific scenarios, including distributional robustness 31 fairness considerations
18]. and continuous treatments 30
OPL is different from our work as follows: OPL aims to directly learn a policy, thus bypassing the need to estimate CATE for decision-making While this approach can optimize decision- making performance, it often does so at the expense of making black-box decisions that are not based treatment effects. In contrast, our work prioritizes the interpretability inherent in CATE- based methods while leveraging insights from OPL to improve deciisionmmking performance. This distinction is particularly relevant in fields like medicine, where the effectiveness of treatments is often evaluated using CATES, and the CATEs are then used to guide interpretable clinical decisions [15].
CATE estimation vs. decision-making In practice, it is common to use CATE estimators for Decision-making through thresholding [e.g.. B 37j. Yet, few works have formally studied the effectiveness of this approach. One literature stream discusses the suboptimality of estimating CATE for decision-making as compared with outcome/ response function modeling [14] 13 12
In this context, Zou et al. (63) study a continuous treatment setting for counterfactual prediction and propose to re-weight the loss with the inverse magnitude of the treatment effect. Additional works propose to
adjust the objective to tailor the learning process towards decision-making works do not focus two-stage meta-learners under model class constraints.
[2 However, these
Relatedly, Bonvini et al. 7 establish minimax- -ottimllttp results on the OPL performance of thresh- olded two-stage meta-learners, but only under certain assumptions (e.g., assuming that the second- stage model is well-speeified). In contrast, we allow for the misspecification of second-stage models (e.g., by incorporating fairness or interpretability constraints) and instead show in such cases that two-stage learners may be suboptimal. Finally, (32) considers learning optimal treatment effect rankings under possible resource constraints while consider thresholding, i.e., treating everyone that benefits from treatment.
To the best of our knowledge, we are the first to show the suboptimality of two-stage meta-learners for decision-making under model class restrictions along with proposing novel methods for improving decision performance. Our work thus bridges critical gap between the theoretical understanding and the practical application of CATE estimation for decision -maiin...
Research gap:
3
Problem setup
3.1 Setting
Data: We consider a standard causal inference setting with a population Z = 2X(XA,,Y) ~ IP, where
X E X C Rd are observed pre-treatment covariates, A E {0, 1) a binary treatment (or action), and
Y E IR is an outcome (or reward) of interest that is observed after the treatment A. We assume that we have access to a dataset (either randomized or observational) = {(x,, ai, b2 of size n E N sampled i.i.d. from IP. For example, in medical setting, X are patient covariates, A is medical
3

---

drug, and Y is health outcome (e.g., blood pressure). Another example is logged data of A/B tests in marketing, where X are user demographics, a binary decision of whether a coupon was given, and Y is some reward such as user engagement.
Notation, We define the response functions paa (z) Y X = A= a] for {0, 1} and the propensity score (behavioral policy) as Tb( = (A 1 z). We refer to these functions as nuisance functions, denoted by r7 (1), Ho: mb). A policy any function æ: X + [0, 1] that maps an individual with covariates X E X' to a probability π (XX) of receiving treatment.
Identifiability: We the potential outcomes framework 48 and denote Y(a) the potential outcome corresponding to a treatment intervention A = a. The potential outcomes are not directly observed, which means that we have impose assumptions to identify any estimands from data.
Assumption 3.1 (Standard causal inference assumptions). For all a E {0, and x E X' it holds: (i) consistency: Y(a) Y whenever A = at (ii) overlap: 0 < TT(2)) and (iii) ignorability: A L (O)) X x.
whenever PP(X x) > 0:
Assumption 3.1 is standard in the causal inference literature [56] 99. (i) Consistency prohibits interference between individuals; (ii) overlap ensures that both treatments are observed for each covariate value; and (iii) ignorability excludes unobserved confounders that affect both the treatment A and the outcome Y. Note that (ii) and (iii) are usually fulfilled in randomized experiments, which fall within our setting.
3.2 Mathematical preliminaries
Policy value. The decision-making performance of a policy a is usually quantified via its policy value. Formally, the policy value is defined via
((=)= = EE) (xx))) = (1=( (X)) BIx + (1 T (X)) (X)):
(1)
Under Assumption 3.1 it is identified via V(T) (X (X) 1000-x(330 (X))) and thus can be estimated from the available data.
CATE. We define the conditional average treatment effect (CATE) as
T(x) (LY (0) x = x] (1(x) #o(x),
(2)
where identifiability in terms of response functions (x) follows again from Assumption 3.1 The CATE captures heterogeneity in the treatment effect across individuals characterized by X.
CATE estimation. A straightforward approach estimating the CATE is the so-called plug-in approach. Here, we first obtain estimators for the response functions ر and Pu (which are standard regression tasks) and then obtain a CATE estimator via fP)(z) = Pr (xx) fz)(z).
However, it is well known that the plug-in approach is suboptimal as it suffers from so-called plug-in bias (34). In contrast, state-of-the-art approaches for CATE estimation are based the following two-stage principle: in stage 1, one obtains estimators i of the nuisance functions 17, and, in stage 2, one estimates the CATE directly via second-stae stage regression
fg arg min Lg(9). gEG
where G is some function class and (g) is a (population) second-stage loss.
Two-stage CATE meta-learners offer two key advantages: First, they enable to direct incorporationcof
onstraints on the CATE estimate, such as fairness requirements (45 (35) or interpretability conditions
[54). Second, it leverages the inductive bias that the CATE structure is typically simpler than its constituent response functions, making direct estimation more effective (33)
One approach for obtaining an optimal policy x² is direct off-policy learning (OPL): here, we directly maximize the estimated policy value and solve TÔPL arg max+en (TT) for some estimator (Tz) of (TT) over a prespecified class of policies II. An advantage of the OPL approach is that it directly optimizes for decision decision-making performance. However, the obtained x² is a black-box policy and may be hard to provide with meaningful interpretation.
Direct OPL:
CATE-baase OPI. Another common approach, which we focus on this paper, is to use the CATE 72x) for decision-makin by thresholding [1] The approach has two steps. First, the CATE

---

f is estimated via e.g., a two-stage meta-learner. Second, the CATE- based policy is obtained via T (z) 1(f(z)= > 0). The treatment is thus only applied individuals with a positive CATE (= individuals for which the treatment helps on average). To see why this is a valid approach note that we can write the policy value as
(T))
(o( [o (X)()),(X)),
where x denotes equivalence up constant (irrelevant to maximization). Hence, the (ground-truth) thresholded CATE policy T (z) maximizes the policy value if ,. E 11.
A benefit of CATE-based policy learning is that the CATE provides individualized estimates of the incremental benefits from treatment. Unlike direct OPL methods, which yield black-box policies optimized solely for overall performance, the CATE rX)) explicitly allows to quantify the net gain from treatment. As a result, CATE-based policy allows to compare the estimated treatment effects against domain knowledge. Further, CATE-basee policy learning enables practitioners to weigh the benefits against potential side effects when making treatment decisions, a critical consideration in domains like personalized medicine.
3.3 Research questions
In this paper, we study the optimality of CATE-based policy learning when the CATE estimator îg is obtained via second-stage regression function class g (as in Eq. 33). More formally:
① Do two-stage estimators fig yield policies n. that maximize the policy value V(T) among thresholded policies π E 1Ig = (1gg > 0) g e 9}?
If E G (i.e., 9 contains the ground-truth truth CATE), optimality (in population) of T is guaranteed by Eq. 44). However, in two-stage CATE estimation, G is often restricted such that $ G. This occurs, for instance, when fairness or interpretability constraints are imposed (55) or when regularization applied to smooth the second-stage model (33). In this setting, later show in Sec. [4. that there can exist policies т E Ilg with (T) > ((T+). In other words, thresholding a two-stage CATE estimator may not yield an optimal policy, even when the policy class is restricted in an analogous manner. This leads to our second research question, where we seek policy that achieves (i) a low CATE estimation error and (ii) good decision performance:
2 How can we learn a function E G that satisfies two key properties: (i) s T (g is good approximation of the CATE), and (ii) (T(x)) (g)(x) > 0) is approximately optimal, that is, V (Tg)) R (TT)), where oद is an optimal policy among the class llg?
4 Re-targeting CATE for decision-making
We now answer both research questions from Sec. 3.3 First, in Sec. 4.1 we show the suboptimality of two-stage CATE estimators for Decision-making when G. Then, Sec. 4.2 we propose new learning objective that balances CATE estimation error and policy value. Finally, in Sec. 4.3 and we propose a two-stage learning algorithm and provide theoretical guarantees.
Sec. 4.4
4.1
Suboptimality of CATE for decision-nit decision-making
To provide intuition why two-stage CATE estimators can be suboptimal for decision-making we first consider a toy example illustrated Fig. 2 (left). Here, we examine a two-stage CATE estimator with one-dimensional covariates X and linear functions. Hence, the policy class we consider is IIg = (1 + b > 0)), which represents the class of thresholded linear policies. The ground- truth CATE is nonlinear that & G.
{g(x) az + b 6 E R} being the class of
We make two key observations: (i) The optimal policy n² arg max V(z) assigns a treatment in the region of the covariate space where ground-truth CATE is positive, but no treatment in the region where negative. This is equivalent to thresholding the ground-truth CATE (represented by the red line). (ii) The optimal linear approximation to the CATE is g" arg mingeg [[(r(X) =
5

---

Figure 2: Experimental results for our proposed method with & being the class of linear models. Left: CATE estimator (blue) is the best linear approximation of the (nonlinear) ground-truth CATE (red). Center: the trained a a(X) detects the region in which the estimated CATE has the wrong sign. Right: re-targeted CATE estimators using our proposed loss with trained (X) and different y values.
(represented by the blue line). Note that the blue line does not intersect the x-axis at the same point as the true CATE. As result, there exists region where the policy Ry- (x) (gg > 0) makes the wrong treatment decision. Thus, the policy ag. is suboptimal, i.e., V(x*) > (TTg The optimal
policy is instead obtained by thresholding linear function that intersects the z-axis at the same point as the ground-truth CATE.
To generalize the above example, we derive the following theorem. We denote (= = ((T(X))((X) to note the dependency of the policy value on the underlying ground-truth т.
Theorem 4.1 (Suboptimality of CATE-based decision-making! Let G be a class of neural networks with fixed architecture. Then, there exists a CATE T" & G, so that, for any optimal CATE approxima- tion g E arg mingeg (X) - (())) it holds that Vy- (Tgt.) ) Vy- (.. policy म..- E arg maX, Vr- (T))
). for any optimal
Proof. See Appendix B
Interpretation. Theorem demonstrates that, regardless of how we choose our class 9 in the second-stage CATE regression, there always exists a ground-truth CATE for which the estimated CATE suboptimal in terms of thresholded decision decision-making (-> thus answering O). Specifically, there always exists more optimal function g E which, while not necessarily the best CATE approximation, yields improved policy value.
In general, the discrepancy between CATE-based and optimal policy value is determined by how well the CATE projection using the model class G can estimate the sign of the ground-truth CATE. This depends on several factors, including (i) the error of the optimal projection (related to model complexity), and (ii) the structure of the ground-truth CATE. As an illustration, consider Figure 2 (left), where the discrepancy in policy value is characterized by the distance of the two points where the blue (estimator) and red line (ground- gground-truth CATE) intersect zero. For (i), we could make the red line steeper (increasing the projection error), and thus arbitrarily widen this gap to obtain an arbitrary discrepancy between CATE-based and optimal policy value. For (ii), we could simply shift the red line upwards, which would also widen the gap without changing the projection error. In Appendix [B] we present an additional result bounding the discrepancy in policy values if we can bound the CATE projection error via G.
4.2 A novel learning objective for re-targeting CATE
Basic idea. Motivated by our previous analysis, we propose a learning objective that learns a re-targeted CATE, leading to an improved policy value while maintaining interpretability approximate CATE estimator (-> thus answering 2)). Our motivation comes from Fig. 2 (right), in which we observe that there exists continuous set of solutions between the optimal CATE approximation and the optimal linear function that maximizes policy value (after thresholding). The basic idea is to minimize convex combination of the CATE estimation error and the negative policy value of the thresholded policy.
Definition 4.2. We define the 7-policy-targeted CATE (T-PCAE) as any solution g" minimizing
(g(9) = (1 (1-r))((((X) 9(X)) (g() > 0))(X)).
(5)
over class of function g E G.

---

The hyperparameter y controls the trade-off between CATE accuracy and policy value optimization. For y 0, the objective reduces to standard CATE estimation, while, for 7 1. corresponds to pure policy value maximization (OPL.) and thus disregards the CATE estimation error. We discuss principled methods for selecting 7 in Section 4.4
Optimization challenges. The loss in Eq. (5) does not allow for gradient-based based optimization due to the con-differentiabiiity of the indicator function (g) > 0). A naive approach would be to use smooth approximation of the indicator via the sigmoid function a (ag(()) for a sufficiently large a. However, this introduces a challenging trade-off: large values for provide a better approximation to the indicator function but suffer from vanishing gradients, while small values for a maintain useful gradients but poorly approximate the indicator function. Furthermore, small values for a may incentivize the model to compensate by increasing gXX), thereby degrading CATE quality.
daptive indicator approximation. We address this optimization challenge by introducing ((X)> > O as a function of the covariates X. That is, we approximate (g) from Eq. (5) via
(g) = (1 -) B[(T(x) ( - - [F[
(a (a(X)s(x)).
(6)
for some fixed adaptive approximation (xX) 0. Such adaptive approach allows to be large when the sign of is correct, thereby providing an improved indicator approximation in regions where signal from the gradient from the policy value term is needed. Fig. iillustrates this concept, where we show an estimated CATE that is suboptimal for decision-making (left plot). The a(X) in Fig. 2 (center) is effective in identifying the region in the covariate space where the sign of is incorrect and provides gradients for these regions. Once we obtain a suitable (:X)), we can minimize (right plot).
C. (g) re-target the CATE estimate in regions of suboptimal decision
Learning ax (x). We obtain a loss for learning (xX) for fixed g by transforming the OPL component
in Eq. (5) into a classification problem (following approach similar to, e.g.,
3
We can write
V(Tg) x af (X) = [X[r TT (X) sgn(r(X))
(7)
By noting that maximizing 1 (g( > 0) sgg(((XX (X)) is equivalent to minimizing the binary cross- entropy loss over g with label (Tx(X) > 0), we can obtain a for fixed g by minimizing
(a)g(a) = 1P. [x(x]| ( ((X))
(8)
where l(u; v) = 1y > 0) log(o(u)) 1(y < 0) (1 - (u)). subject to axx) E [a. oo) for all X' and 0 < The scalar can be tuned by minimizing the loss from Eq. (5) validation set.
Interpretation as stochastic policy. The policy (z) (a(())((z) can be interpreted as the best stochastic policy that is achievable for fixed g. Here, the CATE approximation
determines the sign (i.e., whether to give treatment or not), while the approximation a determines the
stochasticity of the resulting policy. As shown in Fig. _ axx) will be large whenever g(x) has the correct sign, therefore providing policy To.g( (x) that closer to being deterministic.
4.3 Estimated nuisance functions
So far, we have assumed that the true CATE r(X) is known, which is not the case in practice. To address this, employ two-stage estimation procedure similar established CATE estimators
9 33). In the first stage, we obtain estimators f of the nuisance functions, standard regression or classification tasks that be solved using various model-based methods from the literature [50] [5]. In the second-stage, we substitute these first-stage estimates into a second-stage loss that coincides with Eq. (8) and Eq. (6) in expectation.
(/4, (J4,/o, 7). These
To start with, we define
()) (1 - ") 9(x)) - =E(""" (a(X)s(X)) and "y) |
(a(()) y,"). (9)
where ym is one of the following oseudo-outeomes: n
yPI (1(X)
YRA A -o(X) (x)()))))
(10)
WIPW (4-)
(A (a-))(X)) (x))
yDR
mi(X) Pg(X)
X6(X) (x()))
*s(X) (1 xx(
Theoretical analysis. We now justify our pseudo-outcome-baane loss from Eq. 9 theoretically. The first result shows that minimizing CT (T),,,(9) provides meaningful minimizer.
7

---

Theorem 4.3 (Consisteney).. If the muisance functions are perfectly estimated (i.e., n = n), the pseudo-outcome loss C, (g) has the same minimizer as C (g) w.E.t g E for all a and m E (PI, RA, DR).
Proof. See Appendix B
In practice, we use estimated nuisance functions é, which means that Theorem 4.3 may not hold for f% (g) due to possible nuisance estimation errors. However, the following results provides an upper bound on how much the minimizer can deviate in the presence of estimation errors.
Theorem 4,4 (Error rates). Let g" = arg mingeg f. and و = arg mingeg L (,))), be the minimizers of the PT-CATE loss with ground-truth and estimated nuisances. Then, under the additional assumptions listed in Appendix [
it holds
1 01| & a (9.9" Min ((1 (11-))+9) + CCa)
(11)
where || - is the -o-ormm, RT
5,0,0 (g9)) is an optimization-depeere dependent term, C. > 0 is constant depending on ax and M is the (pseudo-outcomee detcomed))pendee))) rate term, in
aa. (g.g") C. a, (g) -
defined via
MPT fi.1 x 11 -2ll2 + 11 Holl?, MEW x l Toll?
(12)
MDR || mul1 (1lmn mill3 + llão - -noll)
(13)
Proof. See Appendix B
Theorem 4.4 shows that, as long as we are able to estimate the nuisance functions rj involved in the corresponding pseudo-outcome reasonably well, we can ensure a sufficiently good second-stage learner. Importantly, the doubly robust pseudo-outcome leads to a doubly robust nuisance error rate: only either the propensity score Te the response functions Pa need to be estimated well for the second-stage stage learner to converge well.
4.4 Learning algorithm
We provide a concrete learning algorithm to obtain g(x) and a(x) from finite data. Given a dataset
D (x; az / we can define empirical versions f. 2, (g) of C (g) and i.m E (a) of C ( (a) by replacing expectations with empirical means. To minimize both empirical losses, we propose to parametrize as and go as neural networks, where & and denote their respective parameters. For training, we propose three-step iterative learning algorithm (shown in Fig, 3
Step 1 (initial CATE estimation): train go by minimiz- ing LT f. (g) over 0, using randomly initialized as with frozen. This gives an initial CATE estimator. Step (region detection): train ag by minimizing Ê. D (Cxp) over Φ, keeping 0 frozen. The objective for ag to iden- tify covariate regions where produces incorrect predic- tions of the sign. Step (CATE refinement): Retrain go by minimizing im fm (g) over O, with & frozen. This step corrects ge in regions previously identified as having incorrect sign predictions. Fig. 2 shows experimental re- sults for each of the three steps of our algorithm. Steps 2 and 3 can be repeated iteratively until convergence. The pseudocode is in Appendix C
LeNning ANg
H&H
Sep 3 Prieed
aU H.H.D
Calerres
Figure 3: Overview of our second-stage architecture and our learning algorithm.
Selecting y The parameter 7 quantifies the trade-off between CATE estimation error and decision performance. Setting = 0 corresponds to standard CATE estimation, while = 1 corresponds to OPL (ignoring a meaningful CATE estimation completely). As such, the selection of o is mainly driven by domain knowledge. In practice, we recommend plotting the estimation CATE error and policy value (e.g., as Fig. 4 and choosing 7 accordingly, or comparing trained models for multiple values of (sensitivity analysis). If the main objective is optimal decision-makenn making, practitioners may
8

---

choose = 1. The corresponding PT-CATE can be interpreted as the best CATE approximation among all functions that can be thresholded for optimal
as long as 7 1.
Additionally, the choice of y should reflect the expected misspecification the function class G. If g is overly flexible, then O may already yield near-optimal decisions, g is capable of closely approximating the true CATE. However, in many real-world scenarios G is deliberately constrained, e.g., for interpretability, fairness, or due limited sample sizes. In such cases, selecting larger ? allows to compensate for this model misspecification.
Experiments
We now confirm the effectiveness of our proposed learning algorithm empirically. As is standard in causal inference [50 33 we use data where we have access to ground-truth values of causal quantities. We also provide experimental results using real-world data. Additional experimental results and robustness checks are reported Appendix F
Implementation details. We use standard feed-forward neural networks with tanh activations for go and with ReLU activations for ag- We use p(z) + as the final activation function for ag to ensure (a())) > a, where p(x:) denotes the softplus function. We perform training using the Adam optimizer [36. Further details regarding architecture, training, and hyperparameters are Appendix
Evaluation. We evaluate a function learned by our method using two established metrics (50
24j: (i) the precision of estimating heterogeneous treatment effects (PEHE) (x) =² which quantifies the CATE estimation error, and (ii) the policy loss (negative policy value) given by (g(() > ())( For the experiments using simulated datasets, we use the known ground-truth CATE rXX) for evaluation. For the experiments using real-world data, we evaluate by using the doubly robust pseudo-outcome YPR instead, as the ground-truth CATE is not available.
Baselines. Standard two-stage CATE learners (i.e.. PI/ RA/ IPW/ DR-learner) correspond to method when setting 7 = 0. To ensure a fair comparison, we use the same neural network architecture for all values of in experiment. We refrain from benchmarking with specific model architectures as we are not claiming general state-of-the-an results using our specific implementation. Additional experiments using different model architectures are in Appendix F Of note, OPL methods can be viewed as special case of our method when setting = 1.
Simulated data. Experiments with ground-truth nuisance functions. Fig. 1 and Fig. 1as22pep (right) show the results of stage of our PT-CATE algorithm for different values of when using ground-truth nuisance functions in the first stage of Algorithm [ The results show visually that our algorithm is effective in improving the decision threshold (and, thus, the policy value) as compared to the result for 7 0, while maintaining good CATE approximations.
6.0
Figure 4: Experimental results for setting A. Shown: PEHE and policy loss over 7 (lower better). Shown: mean and standard errors over runs.
Figure 5: Experimental results for setting B. Shown: PEHE and policy loss over 7 (lower better). Shown: mean and standard errors over 5 runs.

---

Experiments with estimated nuisance functions. We now consider two settings to analyze the effectiveness of our algorithm when using estimated nuisance functions in stage In setting A, we consider a synthetic dataset with nonlinear CATE and aim to learn linear y (similar to Fig. . In setting B, we consider a non-linear but regularized g (similar to Fig. Details regarding the datasets are in Appendix D
Results. We report PEHE and policy loss for all four pseudo- outcomes (PI, R.,, IPW, and DR) in Fig. (Setting A) and Fig. (Setiing B). The results demonstrate that our algorithm is effective in decreasing the policy loss compared to the baselines (7 = 0) when increasing ?. The results for IPW and DR are more noisy as compared to the ones for PI and RA, which is likely due to higher variance as a result of divisions by propen- sity scores (a known issue for these estimation methods; see (9). Nevertheless, our PT-CATE algorithm leads to better average decision performance across all estimation methods while only minimally increasing the PEHE.
-001
4.3
-00
Fake
02
35
Figure 6: Experimental results for real-world data. Shown: PEHE and policy loss over y (lower bet- ter). Shown: Mean and 80% confi- dence intervals over 5 runs.
Real-world data. Dataset. Here, we provide additional exper- imental results using the Hillstrom Email Marketing dataser of = 64000 customers. Details regarding the dataset and our preprocessing are in Appendix
Results. Similar to the experiments with simulated data in Fig. Fand F plot the (estimated) PEHE and policy loss over different values of and compare it to the baseline CATE y = 0 (here: for the DR-learner). The results are shown in Fig The results are consistent with our experiments on synthetic data: our algorithm effective in improving the policy loss as compared to standard CATE estimation (7 10=0) 0).
We also compare against the behavioral policy that generated the data (i.e., using the propensity score T6 as a policy). For this, report the improvement over the behavioral policy for different values of y in Table As we can see, using our PT- CATE algorithm with y 0.98 can lead to 24.45% improved response probability as compared to just using a CATE-based
Table 1: Improvement of our PT- CATE-based policy over the obser- vational policy.
Policy
Improx:
Impeow,
Ubs. policy
1.450 &: 0.000
L.738 d: 0.04G
0.029
19:827
policy (7 0).
0.2
1.792 d: 0.014
01034
23.575
= 0.4
1.796. d: 0.009
0,015
23323
0.8
1. 1.803 d: 0.003
0,01
34.366
= 0.9
1.8011 4: 0.006
0.015
34423
7 7:500 0.96
1.805 3: 0.006
0.035
34-45
6 Discussion
Repoeted: policy values (meen 10
avenuge imppovement soods.
In this paper, we showed that standard two-stage CATE esti- mators can be suboptimal for decision-maker and propose a tolicy-targeted CATE (PT-CATE) to balance estimation and decision performance. Our neural algorithm improves CATE for decision -aainn while maintaining interpretability as CATE.
Limitations: If the second-stage model class is not restricted, our method will not lead to improve- ment over existing two-stage learners. However, it will also not introduce additional bias as the PT-CATE simplifies to standard CATE.
Societal risks: As with any causal inference methods, there are risks of misuse if applied without proper understanding of underlying assumptions or in contexts with significant unmeasured confound- ing. Additionally, automated decision systems based on approach could perpetuate amplify existing biases if training data reflects historical inequities.
Future work: Future directions may include extensions to other settings, such as time series and reinforcement learning (e.g., Q-learning), as well as real-world validation in healthcare and public policy. Furthermore, one may consider incorporating uncertainty into our method using e.g., Bayesian approaches. Finally, methods for variance reduction such as stabilized weighting or propensity clipping may be employed to improve performance in practice.
Conclusion: In sum, our method provides practitioners with principled tool for reliable, data- driven decision-make making by improving the decision performance of two-stage meta-learners under model-class restrictions.
10

---

Acknowledgements
This paper is supported by the DAAD program "Konrad Zuse Schools of Excellence in Artificial Intelligence", sponsored by the Federal Ministry of Education and Research. Dennis Frauen gratefully acknowledges financial support from G-Research.
References
[1] Joshua D. Angrist. "Lifetime earnings and the vietnam era draft lotter: Evidence from social security administrative records". In: The American Economic Review 80.3 (1990), pp. 313-336.
[2] Susan Athey and Stefan Wager. "Policy learning with observational data". In: Econometrica
89.1 (2021), pp. 133-161.
[3] Andrew Bennett and Nathan Kallus. "Efficient policy learning from surrogate- classification reductions". In: JCML. 2020.
[4] Omar Besbes, Robert Phillips, and Assaf Zeevi. "Testing the validity demand model: An operations perspective". In: Manufacturing & Service Operations Management 12.1 (2010),
pp. 162-183.
[5] Aurelien Bibaut et al. "More efficient off-policy evaluation through regularized targeted learning". In: ICML. 2019.
[6] Ioana Bica al. "Estimating counterfactual treatment outcomes over time through adversarially balanced representations". In: ICLR. 2020.
[7] Matteo Bonvini, Edward H. Kennedy, and Luke J. Keele. "Minimax optimal subgroup identifi-
cation". In: arXiv preprint arXiv:2306. 17464 (2023).
[8] Victor Chernozhukov et al. "Double/debiased machine learning for treatment and structural parameters". In: The Econometrics Journal 21.1 (2018), pp. C1-C68.
[9] Alicia Curth and Mihaela van der Schaar. "Nonparametric estimation of heterogeneous treat- ment effects: From theory to learning algorithms". In: AISTATS. 2021.
[10] Alicia Curth and Mihaela van der Schaar. "On inductive biases for heterogeneous treatment
effect cstimation". In: NeurIPS. 2021.
[11] Miroslav Dudik, John Langford, and Lihong Li. "Doubly robust policy evaluation and learning".
In: ICML. 2011.
[12] Carlos Fernández-Loría and Jorge Loría. "Inferring effect ordering without causal effect estimation". In: arXiv preprint arXiv:2206. 12532 (2024).
[13] Carlos Fernández-Lorte and Foster Provost. "Causal classification: Treatment effect estimation
vs. outcome prediction". ". In: Journal of Machine Learning Research 23.59 (2022), pp. 1-35.
[14] Carlos Fernández-Lorft and Foster Provost. "Causal decision making and causal effect esti- mation are not the same... and why it matters". In: INFORMS Journal on Data Science 1.1 (2022), pp. 4-16.
[15] Stefan Feuerriegel et al. "Causal machine learning for predicting treatment outcomes". In: Nature Medicine (2024).
[16] Dylan J. Foster and Vasilis Syrgkanis. "Orthogonal statistical learning". In: The Annals of Statisties 53.3 (2023), pp. 879-908.
[17] Dennis Frauen, Konstantin Hess, and Stefan Feuerriegel. "Model-agnostic meta- neta-learners for estimating heterogeneous treatment effects over time". In: ICLR. 2025.
[18] Dennis Frauen, Valentyn Melnychuk, and Stefan Feuerriegel. "Fair off aff-policy learning from observational data". In: ICML. 2024.
[19] Thomas A. Glass et al. "Causal inference in public health". In: Annual Review of Public Health 34 (2013), pp. 61-75.
[20] Negar Hassanpour and Russell Greiner. "Learning disentangled representations for counterfac- tual regression". In: ICLR. 2020.
[21] Jennifer L. Hill. "Bayesian nonparametric modeling for causal inference". In: Journal of Computational and Graphical Statistics 20.1 (201 pp. 2017-2040.
[22] Nan Jiang and Lihong Li. "Doubly robust off-policy value evaluation for reinforcement
learning".. In: ICML. 2016.

---

[23] Fredrik D. Johansson, Uri Shalit, and David Sonntag. "Learning representations for counter-
factual inference". In: ICML. 2016.
[24]
Nathan Kallus. "Balanced policy evaluation and learning". In: NeurIPS, 2018.
[25]
Nathan Kallus. "More efficient policy learning via optimal retargeting". In: Journal of the American Statistical Association 116.534 (2021), pp. 646-658.
[26] Nathan Kallus and Masatoshi Uehara. "Double reinforcement learning for efficient off-policy evaluation in markov decision processes".. In: Journal of Machine Learning Research 21 (2020), pp. 1-63.
[27] Nathan Kallus and Masatoshi Uchara. "Doubly robust off policy value and gradient estimation for deterministic policies". In: NeurlPS. 2020.
[28] Nathan Kallus and Masatoshi Uehara. "Efficiently breaking the curse of horizon in off- policy evaluation with double reinforcement learning". In: Operations Research 70.6 (2022), pp. 3282-3302.
[29] Nathan Kallus and Masatoshi Uehara. "Intrinsically efficient, stable, and bounded off-policy evaluation for reinforcement learning". In: NeurIPS. 2019.
[30] Nathan Kallus and Angela Zhou. "Policy evaluation and optimization with continuous treat- ments". In: AISTATS. 2018.
[31]
Nathan Kallus et al. "Doubly robust distributionally roust off-policy evaluation and learning". In: ICML. 2022.
[32]
Fahad Kamran, Maggie Maker, and Jenna Wiens. "Learning to rank for optimal treatment allocation under resource constraints". In: AISTATS, 2024.
[33] Edward H. Kennedy. "Towards optimal doubly robust estimation of heterogeneous causal
effects". In: Electronic Journal of Statistics 17.2 (2023), pp. 3008-3049.
[34] Edward H. Kennedy, Sivaraman Balakrishnan, and Larry Wasserman. "Semiparametric coun- terfactual density estimation". In: Biometrika (2023).
[35] Kwangho Kim and José R. Zubizarreta. "Fair and robust estimation of heterogeneous treatment
effects for policy learning". In: ICML. 2023.
[36] Diederik P. Kingma and Jimmy Ba. " "Adam: A method for stochastic optimization". In: ICLR.
2015.
[37] Mathias Kraus, Stefan Feuerriegel, and Maytal Saar-Tsechansky "Data-driven allocation of preventive care with application to diabetes mellitus type II",. In: Manufacturing Service Operations Management) 26.1 (2024), pp. 137-153,
[38] Sören R. Künzel et al. "Metalearners for estimating heterogeneous treatment effects using machine learning". In: Proceedings of the National Academy of Sciences (PNAS) 116.10 (2019), pp. 4156-4165.
[39] Milan Kuzmanovic et al. "Causal machine learning for cost-effective allocation of development Aid". In: KDD. 2024.
[40] Bryan Lim, Ahmed M. Alaa, and Mihaela van der Schaar. "Forecasting treatment responses over time using recurrent marginal structural networks". In: NeurIPS. 2018
[41] Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. "Causal transformer for estimat- ing counterfactual outcomes". In: ICML. 2022.
[42] Pawel Morzywolek, Johan Decruyenaere, and Stijn Vansteelandt. "On general class of orthogonal learners for the estimation of heterogeneous treatment effects". In: arXiv preprint arXiv:2303. 12687 (2023).
[43] Susan Murphy, Mark van der Laan, and James M. Robins. "Marginal mean models for dynamic regimes". In: Journal of the American Statistical Association 96.456 (2001), pp. 1423. [44] Susan A. Murphy. "Optimal dynamic treatment regimes". In: Journal of the Royal Statistical Society: Series B 65.2 (2003), pp. 331-355.
[45] Razieh Nabi and Ilya Shpitser. "Fair inference on outcomes". In: AAAI 2018
[46]
Xinkun Nie and Stefan Wager. "Quasi-oracle estimation of heterogeneous treatment effects".
In: Biometrika 108.2 (2021), pp. 299-319.
[47] Min Qian and Susan A. Murphy. "Performance guarantees for individualized treatment rules". In: Annals of Statistics 39.2 (2011), pp. 1180-1210.
[48] Donald B. Rubin. "Estimating causal effects of treatments in randomized and nonrandomized
studies". In: Journal of Educational Psychology 66.5 (1974), pp. 688-701.
12

---

[49]
Jonas Schweisthal et al. "Reliable off off-policy learning for dosage combinations". In: NeurlPS. 2023.
[50]
Uri Shalit, Fredrik D. Johansson, and David Sontag. "Estimating individual treatment effect: Generalization bounds and algorithms". In: JCML. 2017.
[51]
Adith Swaminathan and Thorsten Joachims. "Counterfactual risk minimization: Learning from
logged bandit feedback". In: ICML. 2015.
[52]
Akira Tanimoto et al. "Regret minimization for causal inference on large treatment space", In: 2021.
[53] Phillip S. Thomas and Emma Brunskill. "Data-efficient off-Policy policy evaluation for rein- forcement Learning".. In: ICML. 2016.
[54] Daniel Tscherutttr, Tobias Hatt, and Stefan Feuerriegel. "Interpretable off-policy learning via hyperbox search". In: JCML. 2022.
[55] Mark J. van der Laan. "Statistical inference for variable importance". In: The International Journal of Biostatistics 2.1 (2006), pp. 1-31.
[56] Mark van der Laan and Donald B. Rubin. "Targeted maximum likelihood learning". In: The International Journal of Biostatistics 2.1 (2006).
[57] Aart van der Vaart. Asymptotic statisties. Cambridge: Cambridge University Press, 1998.
[58] Hal R. Varian. "Causal inference in economics and marketing". In: Proceedings of the National Academy of Sciences (PNAS) 113.27 (2016), pp. 7310-7315.
[59] Stefan Wager and Susan Athey. "Estimation and inference of heterogeneous treatment effects using random forests". In: Journal of the American Statistical Association 113.523 (2018), pp. 1228-1242.
[60] Baqun Zhang and Min Zhang. "C-learning: A new classification framework to estimate optimal dynamic treatment regimes". In: Biometrics 74.3 (2018), pp. 891-899.
[61] Baqun Zhang al. "Estimating optimal treatment regimes from classification perspective".
In: Stat 1.1 (2012), pp. 103-112.
[62] Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. "Learning overlapping representations for the estimation of individualized treatment effects". In: A/STATS. 2020.
[63] Hao Zou et al. "Counterfactual prediction for outcome-oriented treatments". In: ICML. 2022.
13

---

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:
All claims are either backed theoretically (Sec. or empirically (Sec. 5).
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: See Appendix 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those 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 well-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 low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't '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?
Answer: [Yes]
14

---

Justification:
See Appendix B
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 the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to 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 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 of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: See Sec. provided.
Sec.
Appendix
Appendix
D and Appendix
Code
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer 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 is a specific model and empirical evaluation, it may be necessary to either make 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 can also be provided via detailed 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, or other means that are appropriate to the research performed.
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 reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
15

---

Answer: [Yes]
Code is available via anonymized GitHub and can be used to reproduce experiments.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm ec/
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 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 SonddeSubmissiooPP.
//nips. //nips.cc/ppl wwwppbblccgguiicuceccccccoe.oopp
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code 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: See Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper 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: Error bars provided over multiple runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
16

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 to 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 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://aaaaaaa..cpppaaaca/ttataa
Answer: [Yes]
Justification:
The code of ethics was respected in every step.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS 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- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
See Appendix 6
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact 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.
17

---

The conference expects that many papers will be foundational research and not tied 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 an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
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:
No such components have been used.
Guidelines:
The answer NA means that the paper poses 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 usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: No existing licenses used.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/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.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The repository comes with documentation about project structure and code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cooe/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification: No human subjects involved.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in 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 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: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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 expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
19

---

Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy for what should or should not be described.
20

---

A Extended related work
A.1 Deep learning for CATE estimation
In recent years, deep neural networks have gained considerable traction for estimating the conditional average treatment effect (CATE) due to scalability reasons and the ability extract complex features from multimodal data. Notable advances include methods for learning representations that improve
CATE estimation through balancing techniques [3 50 62]
disentanglement strategies (20). and the
incorporation of inductive biases
10
While these approaches focus on estimating the nuisance functions y, they do estimate CATE directly. Instead, they are compatible with two-stage meta-learners including the algorithm proposed in this work. Importantly, existing approaches prioritize accuracy in CATE estimation rather than exploring the interplay between CATE estimation and decision-makinn making, which remains a gap in the existing literature.
A.2 Orthogonal learning
Orthogonal learning (also called debiased learning) is rooted in semiparametric efficiency theory, which has become widespread in estimating heterogeneous treatment effects (57 56). These learners are designed to be robust against errors in nuisance function estimation and offer strong theoretical guarantees.
For CATE estimation, orthogonal learning has been proposed by [8] [6]. Specific instantiations are, for instance, the DR-learner [55 33 and the R-learner [46), where the latter can be interpreted within the broader context of overlap-weighee
DR estimators 442]. Orthogonal learners are, by construction, two-stage meta-leamers and thus applicable to proposed framework. In this work, we leverage the DR-leamer outlined Eq. (10). but extensions to the R-learnr// werlap-weighter methods could be of interest for future research.
A.3 Dynamic settings
Both off-policy learning (OPL) and CATE estimation have a well-established history in dynamic settings, where treatments administered, and outcomes are observed over time. In the context of
CATE estimation, methods have been developed for both model-based approaches [40] 6 41] and
model-agnostic techniques [17). For OPL, dynamic treatment regimes have been extensively studied
43 44 60, 28), alongside reinforcement learning approaches for Markov decision processes under
stationarity assumptions [22] 53 29 26 27). Extending approach to dynamic settings presents an
intriguing avenue for future work. Such extensions could address challenges unique to temporal data, including time-varying confounding.
21

---

B Proofs
B.1 Proof of Theorem 4.1
Proof. Let S be the set of step functiins
S = {f f - R
En E N, 3a = sg < 21 < ... < Fn b,
C11 c2. .2.............................. C.. E R such that f(x) = Ei= 1(r E (x--,, J. (14)
Let Sg denote the set of step functions T such that the optimal approximation g" gets the sign correct everywhere, i.e.,
Sg = {o ((s)n(g+(z)) sgn(r(æ)) for all and 9r arg min B((
(15)
sEG
Note that Sg is non-empty as any constant positive function is in Sg. For any T E S, define the supremum norm as
|| SuPae(a. [r()). We set
M sup ||
(16)
rESp
By definition of the supremum, there exists sequence (Tn)x21 C Sg such that || + M as n + co.
(17)
Hence, for any fixed e >> 0, one can select an index n, such that
|n. Ibo M E.
(18)
Therefore, there exists a corresponding approximation
JTne E arg min E (X) - (X)-g(X))3 {(x))
(19)
9Eg
satisfying
sgn sgn (x)) = sgn((n)) for all x [a,b]. which implies that the corresponding thresholded policy Taree optimal, i.e.
(20)
T Tgrne
E arg max V.n. (T). #Elle
(21)
Because Ti. E Sg is step function, there exists interval Trn, C [a, b] of positive measure that
1Tn, (x)) || lbx > M- for all IT..
(22)
on which (xx) is nearly ||
Let e > 0 be fixed, and let e. We define the CATE T" as
Tri (2) + sgn (x), if 1, if $ 1.
(z)
(23)
Tre (x),
Then, by definition of 7., we have
|| > M,
(24)
which implies by definition of M that 7" satisfies
r" Sg.
(25)
That is, the optimal approximation of gr- with functions in 9 must fail to preserve the sign of 7" on a subset of positive measure E, i...,
sgn (gr (z)) # sgn( (x)) for all x E E.
(26)
We now compare the plug-in policy induced by gr., i.e.,
gg (z) = 1 (g.- (x) > 0),
(27)
22

---

against the policy obtained by thresholding an approximation of Tn, (x), i.e.,
Tr (x) = 1 (gr.. (r) > 0).
(28)
On the set & the decisions made by Tg. and Tie differ in a way such that
(Tr..) (Tgrm.))
(29)
Furthermore, Tn. and 5" have the same sign by definition, which implies together with Eq. (1 that
E arg max (T)( *Ellp
(30)
grnn
because Tne Sg and thus Gine Ng. Hence,
V+- (To..)
max (T)(,,, Hellg
(31)
which completes the proof.
B.2 Proof of Theorem
Proof. One show that, for all pseudo-outcomee Y,," it holds that
[Yy" m
x] = T(X)
(32)
(see e.g.. 99] for proof). Hence, can apply the tower property and write
Cm y,0, (g) = -11 7) B[(Y," ((x))² - 7
(a(X)s -
(33)
(1 7) B (x T(X)
- -(x) -(xx)) - L (a(X)))))
(34)
(1 7) [[( -((1)))
+ g( -(X))]
(35)
y E
(1 (1-7)E E) [[(
- (X, - -x
(36)
- [ (a())(((()))
(1 - (1- ((- (F/
(37)
(g).
(38)
B.3 Theoretical result with nuisance errors (Theorem 4.4
In the following, we provide a new, slightly stronger theoretical result than in Theorem 4.4 but which guarantees that minimizing our proposed loss results in reasonable PT-CATE estimator, the nuisance functions are estimated with errors. Importantly, we upper bound of the PT-CATE error the nuisance errors of the respective adjustment method (pseudo-outcome). For the DR pseudo-outcome, we establish doubly robust convergence rate.
when
Theorem B.1. Let g" arg mingeg (T) and 9 = arg mingeg (T) the PT-CATE loss with ground-truth and estimated nuisances for fixed indicator approximation and 7 E [0, 1]. We assume the following regularity condition: there exists a constant ê > O, so that,
(g) be the minimizers of
for all g E star(G,g") {tg"
t)glg E G), it holds that
E
" (a ( |0( X)) ((XXXXX)XX (x)))
2 8,
(39)
||g* - 011
where ||9" - 01l2 (x) { denotes the squared T--norm and o o"(.) denotes the second derivative of the sigmoid function. Furthermore, assume that the propensity estimator and
23

---

ground-truth response functions are bounded via p ((x) S 1 p and |a( S e for constants p,c > 0 and for all x E X.
Then, all pI, P2 > 0 so that - 7)pp + _22 < 1 - + A it holds that
R na. (Ú. (0 + Mm 9m
(1-2) P1
+
2p2
Ca
119" -112
(40)
P1 +
+1 +
where R ( 9²) C aia (ف) L",,,(l'')) is an optimization optimizaationddeeendenn term, Co 0 is
constant depending on and Min is the (eseuuto-uuccomeeddeeeeddne))))) rate term, defined via
MPI PI 6.
MRA 2| A1 21ll 21 - moll?
(41)
9.21
2²
MPPW
9.1
=
p2
Aい - 1o11
(42)
2
MDR
l. Tbll (Ilx) -ill + 11:0 mall3)
(43)
pa
Proof. Recall that
,T (g) = (1 7) EE( ( 1 1 ( (a(X)((x))
(44)
= (1 - 7) L. Lm, - - (g).
(45)
G. () [X; 9 { We can write
(46)
L MSE (9") 11g" 011 21E [0 g"
- -xx))
(47)
For L (g), we can do functional Taylor expansion, i.e., there exists a E star(g, gg" with
C (g) = D g'l 1 2 D - 9- 9 9"),
(48)
where Dy denotes the functional derivative. For the first-order derivative, we obtain
(g 19 g"] dt +1g0)) -
(49)
1=
=E [r,"d
- - (x))]
(50)
where o'(-) denotes the derivative of the sigmoid function.
For the second-order derivative, we obtain
DyD (0)l0 ." 9"1
(51)
d²
E +()
-00--gg00 (x())))
(52)
dtdv
t=
d
dt [Y, (a()) (x)))) xx()))
(53)
1=0
= E [Y, o"
- [x))]
(54)
S <- - 2
(55)
where the last inequality follows from the regularity assumption. Putting everything together, we obtain that
T (0) 2 (1 7) (CW + l1g" - -11 2E - - -XX0))))
(56)
6 11 -
? cm (0") E [Y,To'
(57)
24

---

or equivalently
(1 (1-7+ 21 - -5112 < (00 - 9" (x))]
(58)
ト [x o' (x)X1
(59)
where R. RT,..,,, (,.0)().9') a (g) T. a.. (g") (9")
E [YY - (xxxx)x(x) -
E - -Xm))(X))
(60)
+ + 7 -xx
(61)
= (X)6)"(()((x)))91101))
(62)
(63)
(64)
with (X) = B|y,"" - Similarly,
E y o' (XXX
(65)
A EA (XA (11 -x((x))
(66)
(XX) (x-))
(67)
Putting everything together, we obtain
(1 11-18 71-7+7 2 7 -112 s
( -00 (X()
(68)
+ 2(1 (1)1)))|[((())---((()(()))))))--
(69)
+
(70)
+ (1
(71)
Note that
2(1 - ())(((((()))))))))))) - (X (X)]
(72)
+ +7B(r(X)o (X) - -
(73)
- (1 - -)) M.. (9" 16 961 -Dy (g") g"]
(74)
Dyf. (g [00 9²]
(75)
<0
(76)
because 9" a minimizer of the oracle nuisance loss L "00000" Hence,
5 0|12 S 1211 (x))) (77)
(1
TF T (X)o' ((Xx) -x
(78)
For the different pseudo outcomes, we can write
A Pl (x) = (x) - (1 (X) + (x) Ho(X)
(79)
RA (X) = (xOo - (x)) + (1 - (X)) - (x))
(80)
A IPW (X) = (x(X)
Mo(x)
(x) - -x)) 1 (X) (f6(X) - (x))
(81)
*o(X)
1
A (R)(X) xx) (X)(X) (x)) mi (X)) - (x) (#o(X) (x))
(X))
(82)
25

---

By applying the Cauchy-Schwarz inequality, we obtain
E ((PP) --(x)) < || - و" I| (llæ1 - full ll + 11ã0 Holl)
(83)
E A RA
(XX)) 9"(x)) s |1 g" (1la Hall ll + 11ã0 mall)
(84)
E (A IPW (XXX) - -"(x))] s واا 9"
1 mbll
(85)
p
E (A DR
-'(x))] s || g"ll (x) (X)
(, pp
(1 Tall (11A1 - w1ll + 11 poll))
(86)
By applying AM-GM inequality and the fact that (a+ + b)2 - 2(a² + +b2), it holds for any p > 0 that
2
2E [A (XP 9x()))] p.ll g¹11²
+
(llă1 m1ll2 + |1 - moll²)
(87)
P1
2
2E
- 9'())) S pillá- - g²1l? + (lli) - m1l1? + I1ã0 - moll?)
(88)
P
6²
21E [AA (X) - 9" -x(x)) S pill - g0||²
|| mll1
(89)
Pip
2
2E
[(4 (X) (x)) S Pılls- - 9011 +
1 1ul1 (1121 pIl12 + 11:0
moll
pip²
(90)
We can write this in generalized form via
1
2E [(> (x))] S p1l19
g²112 +
M
(91)
P1
Using the same arguments and the fact that we can upperbound of (a(X)g x(x) s C. for some constant Ca > 0, we obtain
E [AAT( (a(X)g (X)) -
P2 11g g"1 112
Ca
M
(92)
2p2
Hence, it holds that
(1 - (1-7+ + 8 2 "llo" 0112 s Ry g²) + (1 (1-7) (1llg- - g--12+ +
1
Mm
(93)
P1
P2
11 g²112 + Ca Mm
+?
(94)
2
2p2
11-71
or, equivalently, for pl, P2 so that (1 - 7)p1 + 22 7 + 7, have
I 7,a. n (0,9") + (1 - 7) h My - 202 9.9
Ca M
11g" -012
(95)
1-7+ 1 - (1 --)pp - -
R 1 (0.9") + MM ( (1--) 7
Co
P
(96)
1 1- + 7 ( - 1 + P1 2
B.4 Additional result bounding policy value discrency with projection error
If the projection error is small (i.e., our learned CATE close to the ground-truth CATE), we are able to bound the discrepancy in policy value as follows.
Lemma B.2. Let the true CATE be (xx) = ja1 (x) - - E T.2(P). For a function class & C 12((P)
define its 12 projecction of T as
g" = arg min IE (T(X)) s(x)))),
gEG
26

---

and let E2 = ||T 92112. Furthermore, introduce the corresponding thresholded treatment policies >0},
(x(x) (K(())) > 0},
T
and the (shifted) policy value
V(=) =xx))
Then it holds that
< V( V( S E2.
Proof. Define the disagreement set
M == fx : sign( + sign(g
Then
A = ((T+)) - (x = E (XX) (T (x,
(m))))
On M we have 7 gX) 0, hence [[(x] (x(x) (X)). Therefore
A E (XX) g" (X)) 1x) - [X| - g" (X)l] = ||T 11--g"li
112 S E2.
where the last inequality follows from Hölder's inequality.
27

---

C Implementation details
Estimation of nuisance functiinn We estimate all nuisance functions μ1, Po, and n6 with standard feed-forward neural networks using 4 layers with tanh activations. The response functions Pa are regression functions, which we fit by minimizing the MSE loss on the filtered datasets where we condition on A a. Estimating the propensity score is classification task so that we apply sigmoid output activation function and minimize the binary cross entropy loss. For the synthetic experiments, we mimic randomized controlled trials (RCTs) and use the ground-truth propensity score which we assume to known.
Second-stage model: For our second-stage model (Fig. 3). we model each of ge and as as feed- forward neural networks with 4 layers with tanh activations for Ge and ReLU activations for ag. For the experiments with linear ge (i.e., Fig. 2 and Fig. 4. we set to a single linear layer. For the experiments with regularized go (i.e., Fig. and Fig. we choose custom regularization parameter for each pseudo-outcome type that yields misspecified initial CATE estimate. For the synthetic experiments Fig and Fig[5] we normalize go by applying tanh output activation in step 2 of our learning algorithm (Algorithm() I as we observed that this can stabilize the optimization of Otp- We also applied a weighting scheme in step via 1/ax(X () to further encourage sharp indicator approximation of regions where the initial CATE is correct. We ran our algorithm for K = iteration.
Hyperparameters. To ensure fair comparison, we use the same hyperparameters for each second- stage learner across different and random seeds. For reproducibility purposes, we report the hyperparameters used (c.g., dimensions, learning rate) for all experiments and models (including nuisance functions) as yaml files
Runtime. For the second-stage models, training took approximately samples and a standard computer with AMD Ryzen Pro CPU and 32GB of RAM.
minutes using 2000
Full learning algorithm. The full learning algorithm is reported in Algorithm below.
Algorithm 1: Re-targeted CATE estimation (PT-CATE)
1: Input: Training data {3), 0e, ym)) pstudo-outcome type m, trade-off ? E [0, 1]. learning rates 77 Na epochs E1, E2, Es, itcrations K.
2: Stage 1: Estimate nuisance functions f) = (fi, Ao. f,); compute pseudo-outcomes {y]
3: Stage 2: Initialize parameters 0 (for g) and (for a:).
4: for epoch 1,.... Ei do
5: 0 - 0 119 V 2m
(go) (Step 1}
6: end for
7: for iter = 1,...,KK K do
8: for epoch = 1,.... Eg do Am
9: á - Φ - 17a V oo. (as) (Step 2)
10: end for
11: for epoch = 1.... Es do
12: 0 -- 0 1/8 V a (Step 3)
am
end for
14:
end for
15: Output: Go and ar.
*Code is at autttper/gggtaaa.co///mssssssssa
28

---

D Details regarding simulated data
Data-generating process. Our general data-generating process for simulating datasets is as follows: we start by simulating initial confounders X ~ 410, 1] from uniform distribution. Then, we simulate binary treatments via
A X ~ Bernoulli (Tx(X)
(97)
for some propensity score (x)). Finally, we simulate continuous outcomes via
Y ~ N (N)x(X),
(98)
where ua (X) denotes the response function and 0.01 the noise level.
Dataset for Fig. Here, we emulate an RCT and set the propensity score to Ty (X) = 0.5. We set the response function to Maa (z) a (12x(10 0.5), where a(-) denotes the sigmoid funetion. We sample a training dataset size Taain = 1000 and a test dataset of size Puat = 3000.
Dataset for Fig. 2 Here, we again emulate an RCT and set the propensity score to T = 0.5. (0.6 0.5) + (x > 0.25) sin (10(x - 0.25) + 1.5)). We sample a training dataset size Thtrain = 1000 and test dataset of size Thest 3000.
We set the response function to Ha (ar) a(1 (2r (10
0.25)) + + + (--0.5 <
Dataset for Fig. Here, we use the same propensity and response functions as in Fig. I We sample a training dataset of size Terain 2200 and test dataset of size Thest 3000.
Dataset for Fig. [ Here, we set the propensity score to (X() = (0.1z). We then define the
response function (x) a1 < - -0.25) (0....in((f.) + 0.25)) 0.3) + 1(-0.2 < <
(2a + 2)) - 0.5) + 1 (x: > 0.25) (0.5 (0.5ii0.0) - 0.25) + 1.5)). We sample training
dataset of size Plrain 2200 and test dataset of size The 3000.
29

---

E Details regarding real-world data
The data is taken from https://causeiinfrr rwwdthedocs www.cheeoo....o/e////aeeeeeeeeee http//////////tttttaaaa/aa////o/ html The dataset consists of = 64000 customers who purchased a product within the last 12 months and who were involved an email experiment: group randomly received an email advertising merchandise for men, group 2 for women, and group 3 did not receive an email (control), study the effect of receiving men's 's merchandise email (A = 1) versus receiving no email at all (A ===)) O). Covariates X include various customer features such as purchasing history. Finally, chose an indicator of whether people responded to the email (by clicking on the link to the website) our outcome Y. We split the data into training dataset with 50% of the data, validation set with 20%, and a test set with 30% of the data, All details regarding our data preprocessing are provided within codebase
*Code is available at httt:////gtttaa.commmmmm//ss/sst
30

---

Additional experiments
F1 Motivational experiments with estimated nuisance functions
Legend Coesennoaadda sday Grundhiees CATE
PT:CATE
PTICATE 0.8 PTIGATE y 0.55
38
Figure 7: Results from Figure 1 of the main paper but with using estimated nuisance functions and doubly robust pseudo-outcomes. The dotted lines show regularized two-stage CATE estimators. The blue line corresponds to standard two-stage CATE estimation, while the green and violet lines are generated by our method for different values of 7. The results remain robust.
Lagend Deceon boundryy
Onereistt GATE
RI-CATE y
PT-CATE 5.5
PT-CATE 2.0
FT-CATE 5.98
X
Figure 8: Results from Figure 2 of the main paper but with using estimated nuisance functions and doubly robust pseudo-outcomes. The dotted lines show regularized two-stage CATE estimators. The blue line corresponds to standard two-stage CATE estimation, while the other ones show the re-targeted CATE estimators using our proposed loss with different values for The results remain robust.
31

---

F2 Experiments with sample splitting
0.2
13
0339
Figure 9: Experimental results for setting A with sample splitting, We experiments from Fig. of the main paper but use sample splitting. Shown: PEHE and policy loss over (lower better) with mean and standard errors over 5 runs. Importantly, the results consistent with the results of our main paper.
210
48
RELAR
Marka
105
0.5
0.04
40%
Figure 10: Experimental results for setting B with sample splitting, We re-ran experiments from Fig. of the main paper but use sample splitting. Shown: PEHE and policy loss over 7 (lower The results consistent with the results of
better) with mean and standard errors our main paper.
32

---

E3 Experiments with different nuisance model baselines
Mry
POLERA
030
Figure 11: Experimental results for setting A with TARNet. We re-run experiments from Fig. 4 of the main paper but use now TARNet (Shalit et al. 2017) for estimating the nuisance functions. Shown: PEHE and policy loss over y (lower better) with mean and standard errors over 5 runs. The results are consistent with the results of our main paper.
0.00
480
105
100
Lagerd
380
Figure 12: Experimental results for setting B with TARNet. We re-run our experiments from Fig.
of the main paper but use now TARNet (Shalit al. 2017) for estimating the nuisance functions. Shown: PEHE and policy loss over (lower = better) with mean and standard errors over 5 runs. The results are consistent with the results of our main paper.
33

---

6.14
400
Figure 13: Experimental results for setting A with SNet. We re-run experiments from Fig. 4 of the main paper but use now SNet (9] for estimating the nuisance functions. Shown: PEHE and policy loss over y (lower better) with mean and standard errors over runs. The results are consistent with the results of our main paper.
317
230
120
400
102
100
-099
Figure 14: Experimental results for setting B with SNet. We re-run experiments from Fig. 5 of the main paper but use now SNet (9] for estimating the nuisance functions. Shown: PEHE and policy loss over (lower better) with mean and standard errors over runs. The results consistent with the results of our main paper.
34

---

30:
48
300
388
Figure 15: Experimental results for setting A with FlexNet. We re-run experiments from Fig. of the main paper but use now a version of FlexNet (9) for estimating the nuisance functions. Shown: PEHE and policy loss over 7 (lower better) with mean and standard errors over 5 runs. The results are consistent with the results of our main paper.
FOCEA
49
Adyybe
0.0
103
Lagrd
Seatl 2p12
Figure 16: Experimental results for setting B with FlexNet. We re-run our experiments from Fig. of the main paper but use now a version of FlexNet (9) for estimating the nuisance functions. Shown: PEHE and policy loss over y (lower better) with mean and standard errors over 5 runs. The results are consistent with the results of our main paper.
35

---

E4 Experiments with multivariate covariates
140
139
0.0
140
Figure 17: Experimental results for setting A with multivariate covariates. We re-run experiments from Fig. 5 of the main paper but use now sample X of dimension = 5. Shown: PEHE and policy loss over y (lower better) with mean and standard errors over 5 runs. The results are consistent with the results of our main paper:
36


Reviews:
[{"review_id": "prohcRdZck", "reviewer": "Reviewer_Yh69", "summary": "The paper revisits the crucial problem that there is a difference between optimizing CATE estimation error and actual decision-making \\when using two-stage CATE estimators (e.g., DR-learner). Interestingly, the paper proves theoretically that minimizing CATE estimation error alone can yield suboptimal treatment policies. Finally, the authors propose a novel γ-policy-targeted CATE objective, which balances CATE estimation accuracy with policy value optimization by introducing a trade-off hyperparameter $\\gamma$.", "questions": "1. Is it possible to show on what occasions CATE estimation will give a precise causal decision-making? It is believed that if the CATE estimator perfectly equals the true CATE, i.e., $\\hat{\\tau}=\\tau^*$, the causal decision making would be perfect. Is it possible to relax this condition? I guess Theorem 4.3 and Theorem 4.4 might be helpful to this question. I understand this may not be the main target of this paper, just out of curiosity, and your response to this question will not affect my score.", "limitations": "See above.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. This paper targets a very important research question: the gap between CATE estimation and causal decision making.\n2. This paper gives a theoretical explanation for the suboptimality of the two-stage CATE estimator in causal decision making.\n3. This paper proposes a very interesting and novel learning objective for retargeting CATE.\n\nWeaknesses:\n\n1. I think this paper uses too much content to state the gap between CATE estimation and causal decision making, which is actually not the novel point of this paper. I understand that maybe readers may not be very familiar with the literature, but I suggest the authors move these parts to supplementary or directly suggest refer [1]. I prefer to see more theoretical explanations and analysis on why the CATE estimator might lead to suboptimality in decision making. This is not provided in [1], so it would be helpful to strengthen your contribution, while you just leave Theorem 4.1 with a short interpretation.\n\n2. I found the new learning objective $\\gamma$-policy-targeted CATE objective is not very intuitive. It indeed balances CATE estimation accuracy and policy learning, but how can we control it theoretically? The authors provide some empirical insights on $\\gamma$ selection, but is it possible to give a specific solution? What is the definition of optimal $\\gamma$?\n\n\n[1] Causal Decision Making and Causal Effect Estimation Are Not the Same…and Why It Matters.", "was_revised": true, "final_justification": "I maintain the score 5.", "rebuttal": "Thank you for your positive review and your helpful comments. Below we have drafted careful responses to your suggestions. We will incorporate all points marked with **Action** into the camera-ready version of our paper. \n\n## Response to Weaknesses\n\n1. **Emphasis on CATE estimation vs decision-making.** Thank you for pointing this out. **Action:** We will follow your suggestions closely, and we will rewrite our introduction to put less emphasis on the general topic of “CATE estimation vs decision-making”. Instead, we will elaborate more on the specific research question we target in our paper: the decision-making capabilities of two-stage learners, and a novel two-stage learning algorithm to improve these.\n\n    **Further analysis on the suboptimality of CATE estimators.** We are happy to provide such analysis and kindly refer you to our response to “Questions” below.\n\n1. **Selection of $\\gamma$.** Thank you for raising this point. In general, there is no “mathematically optimal” way of choosing $\\gamma$ as $\\gamma$ defines the trade-off between CATE estimation performance and policy value we wish to achieve. In practice, we propose two different ways of selecting an appropriate value of $\\gamma$:\n\n    - Our PT-CATE framework can be retrained for multiple values of $\\gamma$, and then, the corresponding PEHE (CATE error) and policy value can be estimated (e.g., using doubly robust estimators). Now, a $\\gamma$ can be selected similar to an Elbow plot (e.g., for selecting the optimal number of clusters in a clustering algorithm): choose a value that results in a large policy value improvement, while resulting in a sufficiently accurate CATE estimator. For example, consider our real-world experiments (Fig. 6 in Sec. 5). Here, a reasonable value would be $\\gamma= 0.8$, which is desirable in terms of both policy value improvement and CATE accuracy.\n\n    - If we primarily care about decision-making, we suggest using a large value, such as $\\gamma=0.99$. This ensures that the resulting policy is comparable to standard off-policy learning (in terms of policy value), while still allowing for CATE interpretation. In contrast, the decision function for off-policy learning (i.e., $\\gamma= 1$) has no interpretation as a treatment effect. The exact value (e.g., $\\gamma \\in \\{0.98, 0.99\\}$) can be determined via the same approach as in bullet point A.\n\n**Action:** We will discuss methods for selecting $\\gamma$ more explicitly in the paper.\n\n\n## Response to Questions\n\n1. Thank you for this interesting question. In general, the discrepancy between CATE-based and optimal policy value is determined by how well the CATE projection using the model class $\\mathcal{G}$ can estimate the sign of the ground-truth CATE $\\tau^\\ast$. This depends on several factors, including (i) the error of the optimal projection (related to model complexity), and (ii) the structure of the ground-truth CATE $\\tau$.\n\n    As an illustration, consider Figure 2 (left) from our paper, where the discrepancy in policy value is characterized by the distance of the two points where the blue (estimator) and red line (ground-truth CATE) intersect zero. For (i), we could make the red line steeper (increasing the projection error), and thus can arbitrarily widen this gap to obtain an arbitrary discrepancy between CATE-based and optimal policy value. For (ii), we could simply shift the red line upwards, which would also widen the gap without changing the projection error.\n\n\n    If the projection error is small (i.e., our learned CATE is close to the ground-truth CATE), we are able to bound the discrepancy in policy value as follows:\n\n    **Lemma.** Let let the true CATE be $\\tau(x)=\\mu_{1}(x)-\\mu_{0}(x)\\in L^{2}(\\mathbb{P})$. For a function class $\\mathcal{G}\\subset L^{2}(\\mathbb{P})$, define its $L^{2}$‐projection of $\\tau$ as $g^{\\ast} = \\arg\\min_{g\\in\\mathcal{G}} \\mathbb{E} \\left[(\\tau(X)-g(X))^{2}\\right]$ and $\\epsilon_2 =\\|\\tau-g^{\\ast}\\|_{L^2}$. \n\n    Furthermore, we introduce the corresponding thresholded treatment policies $\\pi_{\\tau}(x)=1(\\tau(x)>0)$, $\\pi_{g^{\\ast}}(x)=\\mathbf{1}(g^{\\star}(x)>0)$ and the (shifted) policy value $V(\\pi) = \\mathbb{E} \\bigl[\\pi(X) \\tau(X)\\bigr]$. Then, it holds that $0 \\leq V \\bigl(\\pi_{\\tau}\\bigr)-V \\bigl(\\pi_{g^{\\star}}\\bigr) \\leq \\epsilon_{2}$.\n\n    **Proof (sketch).** We define the *disagreement set* $\\mathcal{M}:=\\{ x | sign(\\tau(x)) \\neq sign(g^{\\star}(x)) \\}$. Then, it holds that \n$\\Delta =V(\\pi_{\\tau})-V(\\pi_{g^{\\star}}) =\\mathbb{E}\\bigl[\\tau(X)\\bigl(\\pi_{\\tau}(X)-\\pi_{g^{\\star}}(X)\\bigr)\\bigr]\n=\\mathbb{E}\\bigl[\\tau(X)\\mathbf{1}_{\\mathcal{M}}(X)\\bigr]$. On $\\mathcal{M}$ we have $\\tau(X)g^{\\star}(X)\\le 0$, hence\n$|\\tau(X)|\\le|\\tau(X)-g^{\\star}(X)|$. \n\nThis implies $\\Delta \\leq \\mathbb{E}\\bigl[|\\tau(X)-g^{\\star}(X)|1_{\\mathcal{M}}(X)\\bigr] \\leq \\mathbb{E}\\bigl[|\\tau(X)-g^{\\star}(X)|\\bigr] =|\\tau-g^{\\star}|_{1} \\leq \\epsilon_2$., where the last inequality follows from Hölder’s inequality.\n    \nNote, however, that in many applications the projection error will be large as a result of constraints on the function class $\\mathcal{G}$ to ensure, e.g., interpretability [1] or fairness [2, 3], or to accommodate small data regimes. In such scenarios, the discrepancy in policy values may be large, and our method would provide substantial benefits over classical two-stage learners.\n    **Action:** We will add the new theoretical result to our paper along with the above discussion.\n\n## References\n\n[1] Tschernutter et al. Interpretable Off-Policy Learning via Hyperbox Search. ICML 2022.\n\n[2] Fang et al. (2021). Fairness-Oriented Learning for Optimal Individualized Treatment Rules. Journal of the American Statistical Association.\n\n[3] Kim et al. Fair and Robust Estimation of Heterogeneous Treatment Effects for Policy Learning. ICML 2023."}, {"review_id": "OY1yUsgEub", "reviewer": "Reviewer_9gJM", "summary": "This paper investigates how commonly used two-stage CATE estimators (like DR-learner), though accurate for estimating treatment effects, can be suboptimal for decision-making. The key insight is that these estimators often focus on regions far from the decision boundary, which doesn't help practical decisions. To address this, the authors propose a new two-stage learning objective that balances CATE estimation with decision performance. They develop a neural method to optimize this objective and demonstrate its effectiveness both theoretically and empirically. This is the first work to adapt state-of-the-art CATE methods specifically for decision-making.", "questions": "1. This proposed policy-targeted CATE introduced a tuning parameter \\gamma. In the paper, this parameter is proposed to to selected on domain knowledge. In practice, is there any valid way to better select it? How sensitive does it affect the performance of the estimated policy? More experiment results to justify the robustness of the method is helpful.", "limitations": "1. Though the paper provides simulation study to compare with standard two-stage CATE learners, there seems no comparison with direct policy search using OPE. Will the proposed method outperform these methods?\n2. A minor one: the research problem in this paper is novel but this is not reflected in the title of this paper. Probability adding more general description in the title? For now it seems it is too broad.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "1. The paper provide s a novel method to mitigate the gap between CATE estimators and the optimal decision policy.\n2. The paper provides solid theoretical results to justify the proposed research questions. This is also validated in simulation analysis.", "was_revised": true, "final_justification": "I will keep my positive rating after reading author's rebuttal", "rebuttal": "Thank you very much for your positive review! Below, we address your comments and suggestions in detail. We will incorporate all points marked with **Action** into the camera-ready version of our paper.  \n\n\n## Response to Questions\n\n\n\n1. **Selecting $\\gamma$.** Thank you for raising this critical point. We propose two different ways of selecting an appropriate value of $\\gamma$\n    - Our PT-CATE framework can be retrained for multiple values of $\\gamma$, and, then, the corresponding PEHE (CATE error) and policy value can be estimated (e.g., using doubly robust estimators). Now, a $\\gamma$ can be selected similar to an Elbow plot (e.g., for selecting the optimal number of clusters in a clustering algorithm): choose a value that results in a large policy value improvement, while resulting in a sufficiently accurate CATE estimator. For example, consider our real-world experiments (Fig. 6 in Sec. 5). Here, a reasonable value would be $\\gamma=0.8$, which is desirable in terms of both policy value improvement and CATE accuracy.\n    - If we primarily care about decision-making, we suggest using a large value, such as $\\gamma=0.99$. This ensures that the resulting policy is comparable to applying off-policy learning (in terms of policy value), while still allowing for CATE interpretation. In contrast, the decision function for off-policy learning (i.e., $\\gamma=1$) has no interpretation as a treatment effect. The exact value (e.g., $\\gamma \\in \\{0.98, 0.99\\}$) can be determined via the same approach as in bullet point A.\n\n    **Action:** We will discuss methods for selecting $\\gamma$ more explicitly in the paper.\n\n\n## Response to Limitations\n\n\n\n1. **Comparison to direct policy learning.** In our paper, we primarily focus on (widely applied) CATE-based thresholded decision rules over direct policy learning. However, it turns out that our method corresponds to direct policy learning/search when setting $\\gamma= 1$. As a consequence, we essentially have **direct policy learning baselines included in our experiments**, specifically the regression adjustment method [1], inverse propensity weighting (also known as importance sampling) [2], and doubly robust policy learning [3].\n\n    As indicated by our experiments, our PT-CATE framework effectively trades off the CATE estimation error (which increases with $\\gamma$) and the policy value (which decreases with $\\gamma$). For direct policy learning ($\\gamma= 1$), the policy value is the largest, as it is the only objective being optimized. However, CATE performance is the worst, and the resulting thresholding function is no longer interpretable as a CATE. This is particularly relevant in various disciplines (e.g., medicine), where CATE-based decision-making is essential due to its interpretability, and potentially harmful side effects need to be balanced.\n\n\n    **Action:** We will add a clarification to the paper.\n\n1. **Paper title.** Thank you for your suggestion. We are open to renaming the paper after the rebuttal period (which is in line with the guidelines of NeurIPS). A possible title would be “Re-targeting heterogeneous treatment effects for optimal decision-making/ policy learning”. We are happy to consider any other suggestions you may have.\n\n\n## References\n\n[1] Qian & Murphy (2011). Performance guarantees for individualized treatment rules. Annals of Statistics.\n\n[2] Swaminathan and Joachims (2015). Counterfactual risk minimization: Learning from logged bandit feedback. ICML 2015.\n\n[3] Athey and Wager (2021). Policy learning with observational data. Econometrica."}, {"review_id": "NYVwQUhTmX", "reviewer": "Reviewer_XmcU", "summary": "This paper investigates a critical disconnect between the goals of estimating CATE and using those estimates for optimal decision-making. The authors argue that state-of-the-art CATE estimators, such as the DR-learner, are optimized to minimize estimation error across the entire covariate space. However, for decision-making (e.g., assigning treatment if CATE > 0), accuracy is most crucial near the decision boundary. This misalignment can lead to suboptimal policies, especially when the model class for the CATE estimator is restricted (e.g., by regularization or for interpretability).\n\nThe authors propose a novel learning objective that balances the CATE estimation error with the policy value through a tunable hyperparameter. To overcome the challenge of optimizing the non-differentiable policy value term, they introduce an adaptively-smoothed approximation. \n\nThe authors provide theoretical proofs for the suboptimality of standard CATE estimators in this context and establish consistency and error rate guarantees for their proposed method. Through experiments on both synthetic and real-world datasets, they demonstrate that their PT-CATE approach can successfully trade a small amount of CATE estimation accuracy for a significant improvement in policy value.", "questions": "See weakness.", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n- The paper addresses a subtle but highly practical issue. The distinction between optimizing for estimation versus decision-making is a crucial insight, and the authors do an excellent job of motivating it with clear, intuitive examples and visualizations. \n\n- The proposed PT-CATE objective is an elegant solution to the identified problem. The method for handling the non-differentiable indicator function via an adaptive smoothing function ($\\alpha(X)$) is innovative and allows for gradient-based optimization. The approach is well-grounded in theory, with formal proofs backing its main claims.\n\n- The paper provides comprehensive experiments.\n\n## Weaknesses\n\n- The method's key hyper-parameter, $\\gamma$, controls the trade-off between estimation accuracy and decision value. I would suggest authors to give some data-driven heuristic approaches for practical guidance. \n\n- The method is benefit from when the CATE model class is mis-specified or restricted ($\\tau\\notin G$). The paper acknowledges that if the model class is flexible enough to perfectly capture the true CATE, the proposed method will not offer an improvement over standard learners. In practice, we often use simpler model due to the limited sample size or better interpretation. As noted by [1], a mis-specified model can result in near-optimal policy. Is there a way to balance the model complexity, estimation accuracy and policy value?\n\n\n[1] Besbes, O., Phillips, R., & Zeevi, A. (2010). Testing the validity of a demand model: An operations perspective. Manufacturing & Service Operations Management, 12(1), 162-183.", "was_revised": true, "final_justification": "I increased the score to 5.", "rebuttal": "Thank you very much for your thoughtful review! Below, we address all your comments and suggestions. We will incorporate all points marked with **Action** into the camera-ready version of our paper.  \n\n## Response to Weaknesses\n\n1. **Selecting $\\gamma$.** Thank you for raising this critical point. We propose two different ways of selecting an appropriate value of $\\gamma$\n    - Our PT-CATE framework can be retrained for multiple values of $\\gamma$, and then, the corresponding PEHE (CATE error) and policy value can be estimated (e.g., using doubly robust estimators). Now, a $\\gamma$ can be selected similar to an Elbow plot (e.g., for selecting the optimal number of clusters in a clustering algorithm): choose a value that results in a large policy value improvement, while resulting in a sufficiently accurate CATE estimator. For example, consider our real-world experiments (Fig. 6 in Sec. 5). Here, a reasonable value would be $\\gamma=0.8$, which is desirable in terms of both policy value improvement and CATE accuracy.\n    - If we primarily care about decision-making, we suggest using a large value, such as $\\gamma=0.99$. This ensures that the resulting policy is comparable to applying off-policy learning (in terms of policy value), while still allowing for CATE interpretation. In contrast, the decision function for off-policy learning (i.e., $\\gamma=1$) has no interpretation as a treatment effect. The exact value (e.g., $\\gamma\\in\\{0.98,0.99\\}$) can be determined via the same approach as in bullet point a.\n    **Action:** We will discuss methods for selecting $\\gamma$ more explicitly in the paper.\n\n\n\n1. **Model complexity, estimation accuracy, and policy value.** Thank you for this interesting question. In general, the discrepancy between CATE-based and optimal policy value is determined by how well the CATE projection using the model class $\\mathcal{G}$ can estimate the sign of the ground-truth CATE $\\tau$. This depends on several factors, including (i) the error of the optimal projection (related to model complexity), and (ii) the structure of the ground-truth CATE $\\tau$.\n\n    As an illustration, consider Figure 2 (left) from our paper, where the discrepancy in policy value is characterized by the distance of the two points where the blue (estimator) and red line (ground-truth CATE) intersect zero. For (i), we could make the red line steeper (increasing the projection error), and thus can arbitrarily widen this gap to obtain an arbitrary discrepancy between CATE-based and optimal policy value. For (ii), we could shift the red line upwards, which would also widen the gap without changing the projection error.\n\n\n    If the projection error is small (i.e., our learned CATE is close to the ground-truth CATE, we are able to bound the discrepancy in policy value as follows.\n\n    **Lemma.** Let let the true CATE be $\\tau(x)=\\mu_{1}(x)-\\mu_{0}(x)\\in L^{2}(\\mathbb{P})$. For a function class $\\mathcal{G}\\subset L^{2}(\\mathbb{P})$, define its $L^{2}$‐projection of $\\tau$ as $g^{\\ast} = \\arg\\min_{g\\in\\mathcal{G}} \\mathbb{E} \\left[(\\tau(X)-g(X))^{2}\\right]$ and $\\epsilon_2 =\\|\\tau-g^{\\ast}\\|_{L^2}$. \n\n    Furthermore, we introduce the corresponding thresholded treatment policies $\\pi_{\\tau}(x)=1(\\tau(x)>0)$, $\\pi_{g^{\\ast}}(x)=\\mathbf{1}(g^{\\star}(x)>0)$ and the (shifted) policy value $V(\\pi) = \\mathbb{E} \\bigl[\\pi(X) \\tau(X)\\bigr]$. Then, it holds that $0 \\leq V \\bigl(\\pi_{\\tau}\\bigr)-V \\bigl(\\pi_{g^{\\star}}\\bigr) \\leq \\epsilon_{2}$.\n\n    **Proof (sketch).** We define the *disagreement set* $\\mathcal{M}:=\\{ x | sign(\\tau(x)) \\neq sign(g^{\\star}(x)) \\}$. Then, it holds that \n$\\Delta =V(\\pi_{\\tau})-V(\\pi_{g^{\\star}}) =\\mathbb{E}\\bigl[\\tau(X)\\bigl(\\pi_{\\tau}(X)-\\pi_{g^{\\star}}(X)\\bigr)\\bigr]\n=\\mathbb{E}\\bigl[\\tau(X)\\mathbf{1}_{\\mathcal{M}}(X)\\bigr]$. On $\\mathcal{M}$ we have $\\tau(X)g^{\\star}(X)\\le 0$, hence\n$|\\tau(X)|\\le|\\tau(X)-g^{\\star}(X)|$. \n\nThis implies $\\Delta \\leq \\mathbb{E}\\bigl[|\\tau(X)-g^{\\star}(X)|1_{\\mathcal{M}}(X)\\bigr] \\leq \\mathbb{E}\\bigl[|\\tau(X)-g^{\\star}(X)|\\bigr] =|\\tau-g^{\\star}|_{1} \\leq \\epsilon_2$., where the last inequality follows from Hölder’s inequality.\n\nNote, however, that in many applications the projection error will be large as a result of constraints on the function class $\\mathcal{G}$ to ensure, e.g., interpretability [1] or fairness [2, 3], or to accommodate small data regimes. In such scenarios, the discrepancy in policy values may be large, and our method would provide substantial benefits over classical two-stage learners.\n\n**Action:** We will add the new theoretical result to our paper. We will also add the paper mentioned to our related work, noting, however, that it does not focus on a causal inference setting or CATE estimation directly.\n\n## References\n\n[1] Tschernutter et al. Interpretable Off-Policy Learning via Hyperbox Search. ICML 2022.\n\n[2] Fang et al. (2021). Fairness-Oriented Learning for Optimal Individualized Treatment Rules. Journal of the American Statistical Association.\n\n[3] Kim et al. Fair and Robust Estimation of Heterogeneous Treatment Effects for Policy Learning. ICML 2023."}, {"review_id": "LeXx6zHX7h", "reviewer": "Reviewer_x8UB", "summary": "Aiming at the decision-making performance, i.e., maximizing the outcome, the proposed method simultaneously learns the CATE regression and the classification of $I(Y_1\\ge Y_0|x)$.\nThe proposed methods exhibit improved decision-making performance over their base learners (transformed outcome regression methods).", "questions": "[Q1] $α$ is said to represent uncertainty, but what does uncertainty mean in this context? For example, I don't think it refers to aleatoric uncertainty such as $0.5-|p(Y_1\\ge Y_0|x)-0.5|$.\n\n[Q2] There exist methods that estimate the uncertainty, but they require privileged information (special covariates that are available only in training). Why does uncertainty estimation benefit learners without such additional assumptions? Is $\\alpha(x)$ necessary?\n\n[Q3] Why is the alternate learning of $\\phi$ and $θ$ performed a fixed number of times rather than until they converge? Is there any guarantee for this procedure?\n\n[4] Hernández-Lobato, Daniel, et al. \"Mind the nuisance: Gaussian process classification using privileged noise.\" Advances in Neural Information Processing Systems 27 (2014).", "limitations": "The proposed method is only applicable to transformed outcome-based methods.\nThe authors may want to compare with other approaches to check if this limitation is critical.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n[S1] Causal inference aimed at decision making is an important area that has been little studied. In particular, it is important to note that when inductive bias is present, the optimal CATE estimator is not the optimal decision function.\n\n[S2] The proposed approach of learning uncertainty simultaneously is interesting.\n\n## Weaknesses\n\n[W1] The experiment is limited. In particular, there is only a comparison with the basic baselines. For example, RMNet [1] proposes using the geometric mean of regression loss and classification loss, which seems to be closely related to the proposed method. There is also no comparison with architectural innovations such as FlexTENet [2] and SNet [3]. Comparison with OPE methods would also be important.\n\n[W2] The method for estimating uncertainty is somewhat heuristic, and there is no empirical evidence regarding its necessity. See questions below.\n\n[1] Tanimoto, Akira, et al. \"Regret minimization for causal inference on large treatment space.\" *International Conference on Artificial Intelligence and Statistics*. PMLR, 2021.\n[2] Curth, Alicia, and Mihaela Van der Schaar. \"On inductive biases for heterogeneous treatment effect estimation.\" *Advances in Neural Information Processing Systems* 34 (2021): 15883-15894.\n[3] Curth, Alicia, and Mihaela Van der Schaar. \"Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms.\" *International Conference on Artificial Intelligence and Statistics*. PMLR, 2021.", "was_revised": true, "final_justification": "Weak baselines (base models) are resolved by adding SNet and FlexTENet.\nHowever, these methods have strengths in their flexible architectures, and the model flexibility restriction (Settings A&B) harms their strengths.\nAlso, exiting work such as RMNet and OOSR (see below) seems to be closely related in that they pursue both the prediction performance and the decision-making performance, but they are not compared or discussed enough.\nThe authors' response to this point seems to focus on fairly subtle points that may not be of interest to most readers.", "rebuttal": "Thank you for your review and for giving us the opportunity for clarification.  Below, we address your concerns in detail. In case of acceptance, we will incorporate all points marked with **Action** into the camera-ready version of our paper.\n\n\n## Response to Weaknesses\n\n* **[W1] Baselines (first-stage learners).** Our work focuses on improving *two-stage CATE meta-learners*, such as RA-, IPW-, and DR-learners. Two-stage learners are widely considered state-of-the-art for CATE estimation [see e.g., 1] due to two properties:(i) orthogonal learners (e.g., the DR-learner) offer theoretical guarantees regarding robustness to first-stage estimation errors. We also provide similar guarantees for our PT-CATE framework (see Theorem 4.4) (ii) Two-stage learners allow for imposing constraints on the CATE directly via the function class $\\mathcal{G}$, such as fairness, interpretability, or complexity control, which is often desired in real-world applications.\n\n    The baselines you mentioned (e.g., FlexTENet, and SNet) are all *first-stage learners*, i.e., model-based estimators of the response functions. Our method is designed to improve the *second stage*, i.e., the regression step that transforms plugin estimates into CATEs for decision-making. Hence, plugin learners are not directly comparable within our evaluation framework, as they do not support imposing constraints on $\\mathcal{G}$, and thus do not address our research question.\n\n\n    Nonetheless, first-stage learners can be *combined* with two-stage learners by leveraging the first-stage predictions in the second-stage loss. To show that our method is robust w.r.t. the choice of first-stage learners, **we ran additional experiments in Appendix F.3** where we replaced the T-learner used in the main experiments with a TARNet. The results remain stable.\n\n\n    **Action**: To address your suggestion, we further ran **additional experiments** where we replaced the first stage with SNet and FlexTENet. **The results remain consistent and our algorithm effectively trades-off estimation error and policy value.** We will include them in the camera-ready version in case of acceptance. We summarize some of the results in the following table (we will add a plot in our revised paper, as sharing figures is prohibited during the rebuttal):\n\n### Policy value (×10, mean ± std)\n\n|   $\\gamma$                    | 0.00        | 0.20        | 0.40        | 0.80        | 0.90        | 1 *(Direct policy learning)*        |\n|:----------------------|:------------|:------------|:------------|:------------|:------------|:------------|\n| Setting A + SNet      | 6.55 ± 0.07 | 6.61 ± 0.08 | 6.68 ± 0.08 | 6.79 ± 0.08 | 6.80 ± 0.08 | 6.80 ± 0.08 |\n| Setting B + SNet      | 9.63 ± 0.08 | 9.63 ± 0.08 | 9.78 ± 0.09 | 9.86 ± 0.06 | 9.85 ± 0.06 | 9.83 ± 0.05 |\n| Setting A + FlexTENet | 6.54 ± 0.08 | 6.60 ± 0.08 | 6.67 ± 0.08 | 6.78 ± 0.08 | 6.79 ± 0.08 | 6.80 ± 0.08 |\n| Setting B + FlexTENet | 9.63 ± 0.08 | 9.64 ± 0.09 | 9.77 ± 0.10 | 9.86 ± 0.06 | 9.84 ± 0.06 | 9.82 ± 0.05 |\n    \n**Baselines (OPE methods).** In our paper, we primarily focus on (widely applied) CATE-based thresholded decision rules over direct policy learning. However, our method corresponds to direct policy learning/search when setting $\\gamma= 1$. As a consequence, we essentially **have direct policy learning baselines included in our experiments**, specifically the regression adjustment method [2], inverse propensity weighting (also known as importance sampling) [3], and doubly robust policy learning [4].\n\nAs shown by our experiments, our PT-CATE framework effectively trades off CATE estimation error (increasing in $\\gamma$) and policy value (decreasing in $\\gamma$). For direct policy learning ($\\gamma= 1$), the policy value is the largest, as it is the only objective being optimized. However, CATE performance is the worst, and the resulting thresholding function is not interpretable as a CATE anymore. This is particularly relevant in various disciplines (e.g., medicine), where CATE-based decision-making is important due to its interpretability, and potentially negative side effects need to be balanced.\n**Action:** We will add a clarification to the paper.\n\n**Novelty as compared to RMNet.** Thank you for bringing this paper to our attention. RMNet targets a fundamentally different setting: it is designed for a multi-action causal inference setting. In contrast, our setting considers the standard setting with binary treatments and emphasizes interpretable, thresholded policies derived from two-stage learners. Importantly, RMNet is a *first-stage learner*, while ours is a *second-stage learner*. Furthermore, unlike RMNet, our contributions are: (i) A mathematical proof of the suboptimality of standard two-stage CATE learners for decision-making under constraints (Theorem 4.1); (ii) a novel learning objective that trades off CATE accuracy and policy value in the second stage; and (iii) a principled learning algorithm with guarantees under nuisance estimation error (Theorem 4.4) and empirical validation with several nuisance models. To the best of our knowledge, **none of these contributions are addressed by RMNet or any existing related work.**\n**Action:** We will add this discussion to our related work.\n\n* **[W2] Uncertainty.** Here, we kindly refer you to our response to Q1 below.\n\n\n## Response to Questions\n\n* **Q1: Interpretation of $\\alpha$ as uncertainty.** Thank you for bringing up this point. We would like to clarify that the purpose of $\\alpha$ is **not** quantifying estimation uncertainty nor uncertainty of the data-generating process (i.e., epistemic or aleatoric uncertainty). Instead, $\\alpha$ ensures that the resulting PC-CATE policy is a **stochastic policy** that primarily ensures the **feasibility of our optimization procedure.**\n\n    Stochastic policies are common in large parts of the off-policy learning (OPL) and reinforcement learning literature (RL) (e.g., as in [5] or [6]). Even though the ground-truth policy is often deterministic, learning algorithms in OPL or RL often minimize objectives over a class of stochastic policies, as this brings benefits for, e.g., gradient-based optimization by introducing smoothness.\n\n    **Intuition in our setting:** Let us consider Figure 4.1. (left). Here, we would like to classify the covariate space at which the red and the blue lines differ in their sign. In other words, we would like to learn the deterministic function $f(x) = I(\\hat{\\tau}_\\mathcal{G} \\tau(x) &lt; 0)$. One way to think about our approach is to approximate the (non-differentiable function $f(x) \\approx \\sigma(h(X))$ with a smooth classifier, where $h(X)$ is some real-valued score. This classifier is a smooth function that can be trained with stochastic gradient descent. In our setting, we decompose $h(X) = \\alpha(X) g(X)$, where $g(X)$ should remain interpretable as CATE, and $\\alpha(X) > 0$ determines the “steepness” of the sigmoid function. This is what we refer to as “uncertainty”; however, this only refers to the stochasticity of the learned policy, **not** to quantified epistemic or aleatoric uncertainty. \n\n    **Action:** We admit that the notion of “uncertainty” is ambiguous in our setting. We will refrain from using uncertainty and highlight differences between stochastic policies in our paper.\n\n* **Q2: Methods for uncertainty quantification.** We kindly ask you to refer to our above response on why our current approach is unrelated to methods for epistemic or aleatoric uncertainty quantification. However, we agree that incorporating such methods into our approach may be an interesting future direction. For example, one could optimize for distributional versions of the policy value or combine our method with Bayesian estimation or conformal prediction. This would, however, be **unrelated to $\\alpha(X)$ in our setup**, which primarily enables gradient-based optimization.\n\n    **Action:** We will add these considerations to our outlook on future work.\n\n* **Q3: Alternating optimization.** Thank you for your thoughtful question. In principle, our alternating learning procedure can indeed be iterated over arbitrary rounds. However, we observed convergence already after a single iteration across all our experiments. This is due to the structure of our algorithm: the initial CATE estimation provides a reasonable approximation of the treatment effect; the subsequent fitting of $\\alpha$ identifies regions where decision errors (i.e., incorrect signs) occur; and the final refitting step corrects for those errors, resulting in improved downstream decision-making. We emphasize that further iterations can be applied if needed, as our framework does not preclude this. Yet, our empirical results demonstrate that a single iteration suffices to achieve substantial improvements in policy value and CATE estimation across various settings.\n\n    **Theoretical guarantees:** Alternate optimization is a standard approach in min-max and adversarial learning, and has also been used in causal inference literature [6]. While there exist works that provide convergence guarantees for specific alternate optimization setups [7], we view formal convergence guarantees of alternating optimization beyond the scope of our current contribution.\n\n## References\n\n[1] Kennedy et al. Minimax rates for heterogeneous causal effect estimation. Annals of Statistics.\n\n[2] Qian & Murphy (2011). Performance guarantees for individualized treatment rules. Annals of Statistics.\n\n[3] Swaminathan and Joachims (2015). Counterfactual risk minimization: Learning from logged bandit feedback. ICML 2015.\n\n[4] Athey and Wager (2021). Policy learning with observational data. Econometrica.\n\n[5] Kallus. Balanced policy evaluation and learning. NeurIPS 2018.\n\n[6] Schweisthal et al. Reliable Off-Policy Learning for Dosage Combinations. NeurIPS 2023.\n\n[7] Lin et al. On Gradient Descent Ascent for Nonconvex-Concave Minimax Problems. ICML 2022."}, {"review_id": "xGTKjFzzqw", "reviewer": "Reviewer_ytEC", "summary": "The paper addresses a problem in causal inference: while CATE estimators are widely used to guide treatment decisions, they are typically optimized for accurate estimation — not necessarily for decision-making performance, i.e., selection of optimal treatment.  This work studies optimal decision-making based on two-stage CATE estimators. It shows that minimizing CATE estimation error does not always lead to optimal treatment policies. In particular, thresholding estimated CATEs (e.g., at zero) can yield suboptimal decisions when the true CATE is not in the function class $\\mathcal{G}$ in the two-stage estimator. They propose a Policy-Targeted CATE (PT-CATE) estimator, which optimizes a combined loss function balancing CATE estimation accuracy and Policy value (i.e., decision performance). To handle the non-differentiability of the indicator function (used in policy decisions), PT-CATE uses an adaptive indicator approximation mechanism, which learns where estimation errors most impact decisions. Experiments on synthetic and real-world datasets are conducted to assess the performance of PT-CATE in terms of decision-making and estimation accuracy compared to existing two-stage CATE methods.", "questions": "1)\tProvide more justifications about the key premise for the proposed method, i.e., the ground truth CATE is not in $\\mathcal{G}$. When would practitioners choose $\\mathcal{G}$ that does not contain the ground truth CATE or a close approximation of the ground truth CATE? It would be helpful to present theoretical and/or empirical results that help illustrate the impact of the distance between the ground truth CATE and $\\mathcal{G}$ on the performance of the proposed approach.\n2)\tProvide more clarifications about the first term in the RHS of (11) in Theorem 4.4 or redo Theorem 4.4.\n3)\tThe synthetic experiments can be improved considerably to produce results for more realistic settings that can provide more valuable insights to practitioners.", "limitations": "There is no discussion of the limitations and/or potential negative societal impact of their work.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "$\\textbf{Strengths:}$\n\nThis works addresses an important problem: while CATE estimators are widely used to guide treatment decisions, they are typically optimized for accurate estimation — not necessarily for decision-making performance, i.e., selection of optimal treatment.  This work studies optimal decision-making based on two-stage CATE estimators. The basic idea is to minimize a weighted average of the CATE estimation error and the negative policy value of the thresholded policy where the weight $\\gamma$ is a tuning parameter. Since the indicator function in the proposed loss function is non-differentiable, an adaptive indicator approximation method is used to approximate the indicator function with a adaptively weighted sigmoid function $\\sigma(\\alpha(x)g(x))$ where $\\alpha(x)$ is learned from the data. Overall, the paper is well-written. Some theoretical results are provided to support the main claims. Experiments results are provided to support part of the main claims.\n\n\n$\\textbf{Weaknesses:}$\n\n$\\underline{Quality:}$ Theorem 4.4: It is unclear why the results in Theorem 4.4 imply “as long as we are able to estimate the nuisance functions $\\eta$ involved in the corresponding pseudo-outcome reasonably well, we can ensure a sufficiently good second-stage learner.” The first term on the RHS of (11) depends on both $g^*$ and $\\hat{g}$ both of which are included in the LHS too; this seems convoluted. Under what conditions, this first term will go to zero? Plus, the second term on the RHS of (11) suggests that in order for the proposed approach to achieve good performance, the true $\\eta$ need to be estimated accurately. This also begs the question about the key premise of the proposed method, i.e., the ground-truth CATE is not in $\\mathcal{G}$. If we can estimate $\\eta$ well, why cannot we find a sufficiently rich function class $\\mathcal{G}$ such that it contains the ground-truth CATE or a close approximation of the ground-truth CATE?\n\nSection 4.4: The proposed approach for selecting $\\gamma$ is heuristic and not data driven, and it would . Plus, the two terms on the RHS of (6) and (9) may have very different scales (e.g., one has significantly larger values than the other), so the choice of $\\gamma$ is important. \n\nSection 5: In the synthetic data experiments, the simulation settings are fairly simplistic and the results and insights gained from these experiments might not be relevant to real-world data analysis. First, in all simulations, a single X was generated and used; it would strengthen the simulations substantially to consider the case of multivariate X. Second, the true propensity score functions are also very simple for Figures 1, 2, 4 and 5. In addition, in the experiments summarized in Fig 1, Fig 2 and Fig 4, $\\mathcal{G}$ is the set of linear functions. While these simulation experiments can illustrate the power of the proposed approach in a somewhat contrived set-up, they do not reflect what would be done in analysis of real data, i.e., typically a richer and more flexible $\\mathcal{G}$ would be used in real data analysis. The synthetic data experiments can be improved considerably by using $\\mathcal{G}$ that is often used in practice.  For Fig 5, it’s stated that “in setting B, we consider a non-linear but regularized $g$; please provide more information about the formal definition of $\\mathcal{G}$ in this setting.\n\nIn the experiments using real-world data, I am unclear why it is meaningful to use the doubly robust pseudo-outcome instead, given that the ground-truth CATE is not available. It has been shown in the prior works when both working models in the doubly robust estimator are mis-specified, the resulting doubly robust estimator can be more biased than the IPW and RA/PI (e.g., Kang and Schafer, 2007, Statistical Science). In addition, how was $\\mathcal{G}$ defined in the real data experiments?\n\nTo mitigate the issue of higher variance of IPW and DR, could one use stabilized PS weights in the experiments?\n\nThere is no discussion about limitations of the proposed methods.\n\n$\\underline{Clarity:}$ Line 156-158: $\\eta$, $g$, and $L_\\eta(g)$ should be defined or at least refer to later sections where they are defined, say $\\eta$ and $g$ are defined in section 4.3.\n\nFigure 2: Combining the left and middle panels in Figure 2 could offer a better illustration that $\\alpha(x)$ detects the region in which the estimated CATE has the wrong sign.\n\nTable 1: should the numbers in the Improv. column be multiplied by 10? In addition, it is unclear what the ‘Obs. Policy’ row represents. Should the $\\gamma=0$ row be the baseline, as this row corresponds to the standard policy estimator using the estimated CATE?\n\n$\\underline{Significance:}$ Section 3.3: If the function class $\\mathcal{G}$ contains the ground truth CATE, then the existing two-stage estimator $\\hat{\\tau}\\_{\\mathcal{G}}$ yields the optimal policy. It seems that distance between the ground truth CATE and the function class $\\mathcal{G}$ plays a significant role in determining whether $\\hat{\\tau}\\_{\\mathcal{G}}$ can yield the optimal policy or a close-to-optimal policy. Particularly, if $\\mathcal{G}$ is sufficiently rich and general, would it still be necessary to use the proposed method? Also see my earlier comment.\n\n$\\underline{Originality:}$ The ideas of using a combined loss function to balance CATE estimation error and decision performance and an adaptively weighted sigmoid function to approximate the indicator function in the loss represents incremental advances. The concern about the significance further diminishes the potential impact of this work.", "was_revised": true, "final_justification": "The authors have addressed most of my comments, so I raised my score accordingly.", "rebuttal": "Thank you for your detailed feedback and insightful questions. Below, we respond to the specific weaknesses and questions raised in your review. We will incorporate all points marked with **Action** into the camera-ready version of our paper.\n\n## Response to Quality\n\n**First term in the RHS of Theorem 4.4.** Thank you for your helpful comment. We agree that the original phrasing was imprecise. In Theorem 4.4. (RHS), the first term $R^m_{\\gamma, \\alpha, \\hat{\\eta}}(\\hat{g}, g^\\ast)$ quantifies the optimization error from the second-stage regression, while the second term $M_{\\hat{\\eta}, \\eta}^m$  captures the estimation errors in the first-stage nuisance functions (response functions and propensity score). Thus, to ensure a small overall error, both components must be controlled. Importantly, **our result aligns with the broader literature on statistical orthogonal learning** (e.g., Foster & Syrgkanis 2023, Kennedy 2023), which decomposes the estimation error into optimization and nuisance terms. The key insight of Theorem 4.4 is how the nuisance errors affect the final estimation error, and that we obtain a doubly robust rate for the DR-pseudo outcome.\n\n**Restricting $\\mathcal{G}$ when nuisances are complex** Thank you for bringing up this important point, which directly relates to the key motivation of our paper. **There are various important scenarios in which the nuisance functions can be complex, yet it is desirable to restrict the CATE model class $\\mathcal{G}$.** For example:\n\n1. **Simple propensity, complex responses:** In randomized experiments or settings with simple, well-specified propensity models, the response functions can still be complex and challenging to estimate with limited sample size, requiring restricting $\\mathcal{G}$ (e.g., via regularization as in our experiments). Hence, two-stage learners, such as the IPW- or DR-learner, will only recover the best $L_2$ projection of the CATE, which can result in suboptimal decision performance, as shown in Theorem 4.1.\n2. **Domain-specific CATE constraints (e.g., interpretability or fairness):** In many applications, one deliberately restricts $\\mathcal{G}$ to satisfy domain-specific requirements (e.g., using simple models for interpretability or imposing fairness constraints). While nuisances may be estimated flexibly, the CATE model must remain constrained. \n\n**Action**: We will revise our introduction to clarify the crucial motivation for our paper.\n\n**Selecting $\\gamma$.** Thank you for raising this point. We propose two practical ways of selecting an appropriate value of $\\gamma$:\n\n1. Our PT-CATE framework can be retrained for multiple values of $\\gamma$, and then, the corresponding PEHE (CATE error) and policy value can be estimated (e.g., using doubly robust estimators). Now, a $\\gamma$ can be selected by choosing a value that results in a large policy value improvement, while resulting in a sufficiently accurate CATE estimator. For example, in Fig. 6 in Sec. 5, a reasonable value would be $\\gamma= 0.8$.\n2. If we primarily care about decision-making, we suggest using a large value, such as $\\ gamma=0.99$. This ensures that the resulting policy is comparable to applying off-policy learning (in terms of policy value), while still allowing for CATE interpretation. \n\n**Action:** We will discuss methods for selecting $\\gamma$ more explicitly in the paper.\n\n**Synthetic experiments.** Thank you for your helpful suggestions regarding our synthetic experiments.  We would like to emphasize that experiments using synthetic data (usually with somewhat simplified data-generating processes) are standard in the causal inference literature. Nevertheless, we followed your suggestion and performed **additional experiments using a DGP with multivariate covariates** (p = 10) and obtained consistent results as in Fig. 4 & 5.\n\n### PEHE (×10, mean ± std)\n\n|   $\\gamma$     | 0.00        | 0.20        | 0.40        | 0.80        | 0.90        | 0.98        |\n|:-------|:------------|:------------|:------------|:------------|:------------|:------------|\n| Plugin | 2.42 ± 0.02 | 2.43 ± 0.02 | 2.47 ± 0.03 | 2.74 ± 0.05 | 2.83 ± 0.06 | 2.90 ± 0.07 |\n| RA     | 2.42 ± 0.02 | 2.43 ± 0.02 | 2.46 ± 0.03 | 2.55 ± 0.03 | 2.58 ± 0.03 | 2.59 ± 0.04 |\n| IPW    | 2.47 ± 0.03 | 2.48 ± 0.04 | 2.48 ± 0.04 | 2.55 ± 0.07 | 2.58 ± 0.08 | 2.61 ± 0.08 |\n| DR     | 2.43 ± 0.02 | 2.44 ± 0.03 | 2.46 ± 0.03 | 2.55 ± 0.03 | 2.58 ± 0.03 | 2.59 ± 0.03 |\n\n### Policy value (×10, mean ± std)\n\n|  $\\gamma$      | 0.00        | 0.20        | 0.40        | 0.80        | 0.90        | 0.98        |\n|:-------|:------------|:------------|:------------|:------------|:------------|:------------|\n| Plugin | 6.39 ± 0.09 | 6.46 ± 0.09 | 6.56 ± 0.09 | 6.79 ± 0.09 | 6.82 ± 0.08 | 6.84 ± 0.09 |\n| RA     | 6.41 ± 0.10 | 6.45 ± 0.10 | 6.52 ± 0.10 | 6.64 ± 0.09 | 6.66 ± 0.09 | 6.67 ± 0.09 |\n| IPW    | 6.40 ± 0.12 | 6.42 ± 0.12 | 6.46 ± 0.12 | 6.59 ± 0.11 | 6.63 ± 0.11 | 6.66 ± 0.11 |\n| DR     | 6.42 ± 0.11 | 6.46 ± 0.11 | 6.52 ± 0.11 | 6.64 ± 0.10 | 6.66 ± 0.10 | 6.67 ± 0.10 |\n\nRegarding the function classes $\\mathcal{G}$: We chose linear function classes in Fig. 4 to showcase an important real-world use case in which we want to obtain simple interpretable CATE estimates but where the nuisances can be arbitrarily estimated. Furthermore, our additional experiments in Fig. 5 do not use a linear model class. Instead, here we use a forward neural network with 4 layers and 120 hidden neurons. To emulate model class restriction, we impose L2 regularization of 0.03. Intuitively, one can think about this setting similarly to the one displayed in Fig. 1.\n\n**DR outcomes for real-world experiments.** Thank you for your thoughtful question. First, DR pseudo-outcomes yield *asymptotically efficient* estimators when at least one nuisance function (outcome or propensity) is correctly specified. This makes them a standard choice in the literature for various causal inference problems. Furthermore, DR-outcomes are generally more robust to misspecification as only either the propensity or the response function needs to be specified correctly (doubly robustness properties). This is in contrast to RA/IPW, where misspecification of a single nuisance function can lead to several biases. While the result by Kang & Schafer cautions against DR estimators under *severe* dual misspecification, this finding is more relevant to small-sample behavior and extreme scenarios. In our setting, we have access to a sufficiently large dataset to use flexible nuisance models (neural networks) to benefit from the efficiency properties of the DR-estimator.\n\n**Stabilized IPW/ DR.** You are absolutely right: using stabilized versions of IPW/DR might further improve the estimation variance. In fact, stabilization is only one possible practical possibility for improving our method. Other methodological tweaks may be leveraged, such as propensity clipping/ trimming, or calibration of nuisance functions via isotonic regressions. However, we follow established causal inference literature and restrict our methodology and experiments to the standard versions of IPW/ DR. Our goal is to provide a clean, proof-of-concept for our policy-targeted CATE (PT-CATE) framework to evaluate the core methodology under controlled conditions and in the standard setting to ensure that the improvement stem from our proposed objective, rather than auxiliary refinements.\n\n**Action:** We will add an overview of possible methodological improvements to our paper.\n\n## Response to Clarity\n\nThank you for the helpful suggestions. We will update our paper accordingly. \n\nRegarding Table 1: You are correct, we multiplied the numbers by 10 for improved readability. The observational/ behavioural policy is the one generating the data, defined as $\\pi(x) = \\mathbb{P}(A=1 | X=x)$, (i.e., the propensity score which we estimate). In contrast, the row $\\gamma= 0$ corresponds to the thresholded policy using the standard DR-learner for CATE estimation. This improves policy value compared to the observational policy, but remains inferior to our PT-CATE policies.\n\n\n## Response to Significance\n\nTo avoid redundancies, we kindly refer to our “Response to Quality” above for detailed motivational examples. In particular, there are many important applications in which $\\mathcal{G}$ is restricted and does not capture the ground-truth CATE (e.g., when CATE is regularized or fairness or interpretability constraints are imposed, as in e.g., Kim and Zubizarreta et al., ICML 2023). \n\n\n## Response to Originality\n\nThank you for the opportunity to clarify our contributions. While combining CATE loss and policy value into a single loss may seem incremental in isolation, we believe the full set of contributions presents a non-trivial and novel advance in the field:\n\n1. **Theoretical insight**: We formally prove that two-stage CATE estimators can be suboptimal for decision-making under thresholding (Theorem 4.1), which is a result that, to our knowledge, has not been shown before. \n\n2. **New learning algorithm**: We propose a novel learning algorithm that iteratively corrects regions where the CATE sign is likely incorrect, using an adaptively weighted sigmoid to address optimization challenges (see Fig. 2). \n\n3. **Mathematical guarantees**: We provide consistency and error bounds (Theorems 4.3 and 4.4), including doubly robust rates when using DR pseudo-outcomes. \n\n\nIn sum, we are confident that these points present **non-trivial and novel contributions** that, to the best of our knowledge, have not been considered in previous work..\n\n\n## Response to Questions\n\nTo avoid redundancies, we kindly refer to our “Response to Clarity” above.\n\n\n## Response to Limitations\n\nOur discussion on limitations and societal implications is included in Appendix G. We apologize for not highlighting this discussion more clearly in the main paper. We will move the discussion to the main paper of the camera-ready version in case of acceptance."}]

Author rebuttal:

Meta-review: In this paper, the authors propose an approach for optimal decision-making using two-stage CATE estimators. They demonstrated that a direct application of optimal CATE estimator can lead to suboptimal decision-making. To address this, they introduced a new objective that balances the CATE error and decision performance. The authors provided solid theoretical justification and empirical support for their claims.

The paper was well-received by most reviewers. Most reviewers (Reviewers Yh69, XmcU, x8UB, ytEC) acknowledged the importance of the problem, and several praised the paper's novelty (Reviewers Yh69, 9gJM, XmcU) and meaningful theoretical results (Yh69, 9gJM, XmcU, ytEC). In their rebuttal, the authors successfully addressed some key concerns, including the selection of $\gamma$ (raised by yh69, 9gJM, XmcU, ytEC) and comparisons with direct policy learning (raised by 9gJM, x8UB), which led to some reviewers raising their scores.

The only remaining negative rating ("borderline reject") came from Reviewer x8UB despite that the reviewer had increased the rating after the rebuttal. In my final discussion with the reviewer, Reviewer x8UB indicated that the only remaining concern is that the authors do not clearly state their focus on settings with limited model complexity in the title and abstract, as well as how it contrasts with the method without such restriction. However, Reviewer x8UB also acknowledged the paper's novelty within this setting. In my opinion, the focus on settings with limited complexity is still practically relevant, given the authors' examples of interpretability and fairness constraints listed in the response.

Overall, the paper received strong feedback, and its contributions outweigh the remaining negative criticism. Hence, I recommend acceptance.
Decision: Accept (poster)