Title: Treatment Effect Estimation for Optimal Decision-Making
Venue: NeurIPS 2025
Accepted: yes
Keywords: treatment effect estimation

policy learning

causal inference
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

Abstract

Decision-making in various fields, such as 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-making performance of modern CATE estimators (meta-learners) is poorly understood. In this paper, we 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 can be suboptimal when used for decision-making in common settings where the second-stage regression is over a restricted function class (e.g., when using regularization or employing fairness/interpretability constraints). 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 re-targets 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.

1 Introduction

Data-driven decision-making across various fields, such as medicine [19], public policy [1, 39], and marketing [58], relies on understanding how treatments affect different individuals and groups. This heterogeneity in the treatment effect across individuals is typically quantified through the conditional average treatment effect (CATE) . Then, a common approach from practice to obtain decisions from a CATE is thresholding : individuals with positive CATE receive treatment, while those with negative CATE do not [13]. For example, in medicine, clinicians typically administer treatments to 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 a thresholding approach for decision-making are unclear. We argue that minimizing estimation error and optimizing

Valentyn Melnychuk

LMU Munich

Munich Center for Machine Learning

for decision-making 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, 13, 12]. Nevertheless, the optimality of decision-making based on modern CATE estimators remains unclear, and approaches are missing for how to improve thresholding-based 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 a prespecified function class. They are widely used in practice and include orthogonal learners such as the DR-learner [55, 33]. We show theoretically that, while these methods may be optimal for CATE estimation, they can lead to suboptimal decisions when combined with a 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 [18, 35].

Intuition: Why are two-stage learners suboptimal for decision-making? Fig. 1 shows the results of one of our experimental setups (details in Sec. 5), where the groundtruth CATE (red line) is positive everywhere, except for a small region of the covariate space. An optimal policy is 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 a 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 of CATE accuracy to yield a better downstream decision performance, so that the decisions coincide with thresholding the ground-truth CATE.

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 retarget the CATE to a new estimand which we call policytargeted CATE (PT-CATE) . Our PT-CATE can still be approximately interpreted as a CATE, while leading to









































=0



=0.8



=0.95

superior policies. We further propose a 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 are made; and finally, we adjust the first neural network to improve 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-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 a neural method for our learning objective with theoretical guarantees and empirically demonstrate its ability to effectively trade-off CATE estimation error and decision-making performance.

2 Related work

Our work connects to several literature streams below (see Appendix A 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] or tailored neural networks [23, 50] for CATE estimation. (ii) Model-

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-learner [55, 33]. The DR-learner has the additional advantage of being Neyman-orthogonal and thus being provably robust against estimation errors from the first-stage regression [8, 16]. In this paper, we show that such two-stage learners can be suboptimal for decision-making when using restricted second-stage model classes and propose a method to improve them.

(Direct) Off-policy learning (OPL). The goal in OPL is to directly learn an 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) [47] leverages estimates of the response functions; (ii) inverse propensity weighting (IPW) [51] re-weights the data such that they resemble samples under the evaluation policy; and (iii) the doubly robust method (DR) [11, 2] combines both. Recent work has focused on enhancing finite-sample performance through techniques such as reweighting [24, 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, 49].

OPL is different from our work as follows: OPL aims to directly learn a policy, thus bypassing the need to estimate a CATE for decision-making. While this approach can optimize decisionmaking performance, it often does so at the expense of making black-box decisions that are not based on treatment effects. In contrast, our work prioritizes the interpretability inherent in CATEbased methods while leveraging insights from OPL to improve decision-making 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., 13, 37]. 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 [4, 52]. However, these works do not focus two-stage meta-learners under model class constraints.

Relatedly, Bonvini et al. [7] establish minimax-optimality results on the OPL performance of thresholded two-stage meta-learners, but only under certain assumptions (e.g., assuming that the secondstage model is well-specified). 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 we consider thresholding, i.e., treating everyone that benefits from treatment.

Research gap: 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 a critical gap between the theoretical understanding and the practical application of CATE estimation for decision-making.

3 Problem setup

3.1 Setting

Data: We consider a standard causal inference setting with a population Z = ( X,A,Y ) ∼ P , where X ∈ X ⊆ R d are observed pre-treatment covariates, A ∈ { 0 , 1 } is a binary treatment (or action), and Y ∈ R 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) D = { ( x i , a i , y i ) } n i =1 of size n ∈ N sampled i.i.d. from P . For example, in a medical setting, X are patient covariates, A is a medical

drug, and Y is a health outcome (e.g., blood pressure). Another example is logged data of A/B tests in marketing, where X are user demographics, A is a binary decision of whether a coupon was given, and Y is some reward such as user engagement.

Notation. We define the response functions as µ a ( x ) = E [ Y | X = x, A = a ] for a ∈ { 0 , 1 } and the propensity score (behavioral policy) as π b ( x ) = P ( A = 1 | X = x ) . We refer to these functions as nuisance functions , denoted by η = ( µ 1 , µ 0 , π b ) . A policy is any function π : X → [0 , 1] that maps an individual with covariates X ∈ X to a probability π ( X ) of receiving treatment.

Identifiability: We use the potential outcomes framework [48] and denote Y ( a ) as the potential outcome corresponding to a treatment intervention A = a . The potential outcomes are not directly observed, which means that we have to impose assumptions to identify any estimands from data.

Assumption 3.1 (Standard causal inference assumptions) . For all a ∈ { 0 , 1 } and x ∈ X it holds: (i) consistency : Y ( a ) = Y whenever A = a ; (ii) overlap : 0 < π b ( x ) < 1 whenever P ( X = x ) > 0 ; and (iii) ignorability : A ⊥ Y (1) , Y (0) | X = x .

Assumption 3.1 is standard in the causal inference literature [56, 9]. (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 π is usually quantified via its policy value . Formally, the policy value is defined via

Under Assumption 3.1, it is identified via V ( π ) = E [ π ( X ) µ 1 ( X ) + (1 -π ( X )) µ 0 ( X )] , and thus can be estimated from the available data.

CATE. We define the conditional average treatment effect (CATE) as

where identifiability in terms of response functions µ a ( 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 to estimating the CATE is the so-called plug-in approach. Here, we first obtain estimators for the response functions ˆ µ 1 and ˆ µ 0 (which are standard regression tasks) and then obtain a CATE estimator via ˆ τ PI ( x ) = ˆ µ 1 ( x ) -ˆ µ 0 ( x ) .

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 on the following two-stage principle: in stage 1, one obtains estimators ˆ η of the nuisance functions η , and, in stage 2, one estimates the CATE directly via a second-stage regression

where G is some function class and L η ( 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].

Direct OPL : One approach for obtaining an optimal policy π ∗ is direct off-policy learning (OPL): here, we directly maximize the estimated policy value and solve π ∗ OPL = arg max π ∈ Π ˆ V ( π ) for some estimator ˆ V ( π ) of V ( π ) over a prespecified class of policies Π . An advantage of the OPL approach is that it directly optimizes for decision-making performance. However, the obtained π ∗ is a black-box policy and may be hard to provide with meaningful interpretation.

CATE-based OPL. Another common approach, which we focus on in this paper, is to use the CATE τ ( x ) for decision-making by thresholding [13]. The approach has two steps. First, the CATE

ˆ τ is estimated via e.g., a two-stage meta-learner. Second, the CATE-based policy is obtained via π ˆ τ ( x ) = 1 (ˆ τ ( x ) > 0) . The treatment is thus only applied to 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

where ∝ denotes equivalence up to a constant (irrelevant to maximization). Hence, the (ground-truth) thresholded CATE policy π τ ( x ) maximizes the policy value if π τ ∈ Π .

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 τ ( X ) 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-based 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 a second-stage regression over a function class G (as in Eq. (3)). More formally:

1 Do two-stage estimators ˆ τ G yield policies π ˆ τ that maximize the policy value V ( π ) among thresholded policies π ∈ Π G = { 1 ( g > 0) | g ∈ G} ?

If τ ∈ G (i.e., G contains the ground-truth CATE), optimality (in population) of π ˆ τ is guaranteed by Eq. (4). However, in two-stage CATE estimation, G is often restricted such that τ / ∈ G . This occurs, for instance, when fairness or interpretability constraints are imposed [54, 35], or when regularization is applied to smooth the second-stage model [33]. In this setting, we later show in Sec. 4.1 that there can exist policies π ∈ Π G with V ( π ) > V ( π ˆ τ ) . 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 a policy that achieves (i) a low CATE estimation error and (ii) a good decision performance:

2 How can we learn a function g ∈ G that satisfies two key properties : (i) g ≈ τ ( g is a good approximation of the CATE), and (ii) π g ( x ) = 1 ( g ( x ) > 0) is approximately optimal, that is, V ( π g ) ≈ V ( π ∗ G ) , where π ∗ G is an optimal policy among the class Π G ?

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, in Sec. 4.2, we propose a new learning objective that balances CATE estimation error and policy value. Finally, in Sec. 4.3 and Sec. 4.4, we propose a two-stage learning algorithm and provide theoretical guarantees.

4.1 Suboptimality of CATE for decision-making

To provide an intuition on why two-stage CATE estimators can be suboptimal for decision-making, we first consider a toy example illustrated in Fig. 2 (left). Here, we examine a two-stage CATE estimator with one-dimensional covariates X and G = { g ( x ) = ax + b | a, b ∈ R } being the class of linear functions. Hence, the policy class we consider is Π G = { 1 ( ax + b > 0) } , which represents the class of thresholded linear policies. The ground-truth CATE is nonlinear so that τ / ∈ G .

We make two key observations: (i) The optimal policy π ∗ = arg max π ∈ Π G V ( π ) assigns a treatment in the region of the covariate space where ground-truth CATE is positive, but no treatment in the region where it is 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 min g ∈G E [( τ ( X ) -g ( X )) 2 ]



















































































(

x

)















































=0



=0.8



=0.9



=0.98

(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 a result, there exists a region where the policy π g ∗ ( x ) = 1 ( g ∗ ( x ) > 0) makes the wrong treatment decision. Thus, the policy π g ∗ is suboptimal, i.e., V ( π ∗ ) > V ( π g ∗ ) . The optimal policy is instead obtained by thresholding a linear function that intersects the x -axis at the same point as the ground-truth CATE.

To generalize the above example, we derive the following theorem. We denote V τ ( π ) = E [ τ ( 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 τ ∗ / ∈ G , so that, for any optimal CATE approximation g ∗ τ ∗ ∈ arg min g ∈G E [( τ ∗ ( X ) -g ( X )) 2 ] , it holds that V τ ∗ ( π g ∗ τ ∗ ) < V τ ∗ ( π ∗ τ ∗ ) , for any optimal policy π ∗ τ ∗ ∈ arg max π ∈ Π G V τ ∗ ( π ) .

Interpretation. Theorem 4.1 demonstrates that, regardless of how we choose our class G in the second-stage CATE regression, there always exists a ground-truth CATE for which the estimated CATE is suboptimal in terms of thresholded decision-making ( → thus answering 1 ). Specifically, there always exists a more optimal function g ∈ G 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-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. 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 as an approximate CATE estimator ( → thus answering 2 ). Our motivation comes from Fig. 2 (right), in which we observe that there exists a 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 a convex combination of the CATE estimation error and the negative policy value of the thresholded policy.

Definition 4.2. We define the γ -policy-targeted CATE ( γ -PT-CATE) as any solution g ∗ minimizing

over a class of function g ∈ G .

The hyperparameter γ controls the trade-off between CATE accuracy and policy value optimization. For γ = 0 , the objective reduces to standard CATE estimation, while, for γ = 1 , corresponds to pure policy value maximization (OPL) and thus disregards the CATE estimation error. We discuss principled methods for selecting γ in Section 4.4.

Optimization challenges. The loss in Eq. (5) does not allow for gradient-based optimization due to the non-differentiability of the indicator function 1 ( g ( X ) > 0) . A naïve approach would be to use a smooth approximation of the indicator via the sigmoid function σ ( αg ( X )) for a sufficiently large α . 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 α maintain useful gradients but poorly approximate the indicator function. Furthermore, small values for α may incentivize the model to compensate by increasing g ( X ) , thereby degrading CATE quality.

Adaptive indicator approximation. We address this optimization challenge by introducing α ( X ) > 0 as a function of the covariates X . That is, we approximate L γ ( g ) from Eq. (5) via

for some fixed adaptive approximation α ( X ) > 0 . Such an adaptive approach allows α to be large when the sign of g is correct, thereby providing an improved indicator approximation in regions where no signal from the gradient from the policy value term is needed. Fig. 2 illustrates this concept, where we show an estimated CATE that is suboptimal for decision-making (left plot). The α ( X ) in Fig. 2 (center) is effective in identifying the region in the covariate space where the sign of g is incorrect and provides gradients for these regions. Once we obtain a suitable α ( X ) , we can minimize L γ ( g ) to re-target the CATE estimate in regions of suboptimal decision-making (right plot).

Learning α ( X ) . We obtain a loss for learning α ( X ) for fixed g by transforming the OPL component in Eq. (5) into a classification problem (following an approach similar to, e.g., [61, 3]). We can write

By noting that maximizing 1 ( g ( X ) > 0) sgn( τ ( X )) is equivalent to minimizing the binary crossentropy loss over g with label 1 ( τ ( X ) > 0) , we can obtain α for fixed g by minimizing

where ℓ ( u ; y ) = -1 ( y > 0) log( σ ( u )) -1 ( y < 0) log(1 -σ ( u )) , subject to α ( x ) ∈ [ a, ∞ ) for all x ∈ X and 0 < a . The scalar a can be tuned by minimizing the loss from Eq. (5) on a validation set.

Interpretation as stochastic policy. The policy π α,g ( x ) = σ ( α ( x ) g ( x )) can be interpreted as the best stochastic policy that is achievable for a fixed g ∈ G . Here, the CATE approximation g determines the sign (i.e., whether to give treatment or not), while the approximation α determines the stochasticity of the resulting policy. As shown in Fig. 2, α ( x ) will be large whenever g ( x ) has the correct sign, therefore providing a policy π α,g ( x ) that is closer to being deterministic.

4.3 Estimated nuisance functions

So far, we have assumed that the true CATE τ ( X ) is known, which is not the case in practice. To address this, we employ a two-stage estimation procedure similar to established CATE estimators [9, 33]. In the first stage , we obtain estimators ˆ η of the nuisance functions, η = ( µ 1 , µ 0 , π ) . These are standard regression or classification tasks that can be solved using various model-based methods from the literature [50, 59]. 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.

To start with, we define

where Y η m is one of the following pseudo-outcomes:

Theoretical analysis. We now justify our pseudo-outcome-based loss from Eq. (9) theoretically. The first result shows that minimizing L m γ,α,η ( g ) provides a meaningful minimizer.

Theorem 4.3 (Consistency) . If the nuisance functions are perfectly estimated (i.e., ˆ η = η ), the pseudo-outcome loss L m γ,α,η ( g ) has the same minimizer as L m γ,α ( g ) w.r.t. g ∈ G for all α and m ∈ { PI , RA , IPW , DR } .

Proof.

In practice, we use estimated nuisance functions ˆ η , which means that Theorem 4.3 may not hold for L m γ,α, ˆ η ( 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 min g ∈G L m γ,α,η ( g ) and ˆ g = arg min g ∈G L m γ,α, ˆ η ( g ) be the minimizers of the PT-CATE loss with ground-truth and estimated nuisances. Then, under the additional assumptions listed in Appendix B, it holds

where || · || is the L 2 -norm, R m γ,α, ˆ η (ˆ g, g ∗ ) = L m γ,α, ˆ η (ˆ g ) -L m γ,α, ˆ η ( g ∗ ) is an optimization-dependent term, C α > 0 is a constant depending on α , and M m ˆ η,η is the (pseudo-outcome-dependent) rate term, defined via

Theorem 4.4 shows that, as long as we are able to estimate the nuisance functions η 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 π b or the response functions µ a need to be estimated well for the second-stage learner to converge well.

4.4 Learning algorithm

We provide a concrete learning algorithm to obtain g ( x ) and α ( x ) from finite data. Given a dataset D = { ( x i , a i , y i ) } n i =1 , we can define empirical versions ˆ L m γ,α, ˆ η ( g ) of L m γ,α, ˆ η ( g ) and ˆ L m γ,g, ˆ η ( α ) of L m γ,g, ˆ η ( α ) by replacing expectations with empirical means. To minimize both empirical losses, we propose to parametrize α ϕ and g θ as neural networks, where ϕ and θ denote their respective parameters. For training, we propose a three-step iterative learning algorithm (shown in Fig. 3):

· Step 1 (initial CATE estimation): train g θ by minimizing ˆ L m γ =0 ,α ϕ , ˆ η ( g ) over θ , using randomly initialized α ϕ with ϕ frozen. This gives an initial CATE estimator. · Step 2 (region detection): train α ϕ by minimizing ˆ L m γ,g θ , ˆ η ( α ϕ ) over ϕ , keeping θ frozen. The objective is for α ϕ to identify covariate regions where g θ produces incorrect predictions of the sign. · Step 3 (CATE refinement): Retrain g θ by minimizing ˆ L m γ,α ϕ , ˆ η ( g ) over θ , with ϕ frozen. This step corrects g θ in regions previously identified as having incorrect sign predictions. Fig. 2 shows experimental results for each of the three steps of our algorithm. Steps 2

Legend:

Covariates

Treatments

Outcomes

MLP layer

Est. nuisances

Neural architecture

Learning steps

and 3 can be repeated iteratively until convergence. The pseudocode is in Appendix C.

Selecting γ . The parameter γ 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 γ is mainly driven by domain knowledge. In practice, we recommend plotting the estimation CATE error and policy value (e.g., as in Fig. 4) and choosing γ accordingly, or comparing trained models for multiple values of γ (sensitivity analysis). If the main objective is optimal decision-making, practitioners may

choose γ ≈ 1 . The corresponding PT-CATE can be interpreted as the best CATE approximation among all functions that can be thresholded for optimal decision-making as long as γ < 1 .

Additionally, the choice of γ should reflect the expected misspecification of the function class G . If G is overly flexible, then γ ≈ 0 may already yield near-optimal decisions, as 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 to limited sample sizes. In such cases, selecting a larger γ allows to compensate for this model misspecification.

5 Experiments

We now confirm the effectiveness of our proposed learning algorithm empirically. As is standard in causal inference [50, 9, 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 in Appendix F.

Implementation details. We use standard feed-forward neural networks with tanh activations for g θ and with ReLU activations for α ϕ . We use ρ ( x ) + a as the final activation function for α ϕ to ensure α ϕ ( x ) > a , where ρ ( x ) denotes the softplus function. We perform training using the Adam optimizer [36]. Further details regarding architecture, training, and hyperparameters are in Appendix C.

Evaluation. We evaluate a function g learned by our method using two established metrics [50, 24]: (i) the precision of estimating heterogeneous treatment effects (PEHE) ˆ E n [( g ( X ) -τ ( X )) 2 ] , which quantifies the CATE estimation error, and (ii) the policy loss (negative policy value) given by -ˆ E n [ 1 ( g ( X ) > 0) τ ( X )] . For the experiments using simulated datasets, we use the known ground-truth CATE τ ( X ) for evaluation. For the experiments using real-world data, we evaluate by using the doubly robust pseudo-outcome Y DR ˆ η instead, as the ground-truth CATE is not available.

Baselines. Standard two-stage CATE learners (i.e., PI/ RA/ IPW/ DR-learner) correspond to our method when setting γ = 0 . To ensure a fair comparison, we use the same neural network architecture for all values of γ in our experiment. We refrain from benchmarking with specific model architectures as we are not claiming general state-of-the-art results using our specific implementation. Additional experiments using different model architectures are in Appendix F. Of note, OPL methods can be viewed as a special case of our method when setting γ = 1 .

Simulated data. Experiments with ground-truth nuisance functions. Fig. 1 and Fig. 2 (right) show the results of stage 2 of our PT-CATE algorithm for different values of γ when using ground-truth nuisance functions in the first stage of Algorithm 1. 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 γ = 0 , while maintaining good CATE approximations.









































=0



















































=0

















































=0



















































=0





















































=0





















































=0

















































=0













































=0



















Experiments with estimated nuisance functions. We now consider two settings to analyze the effectiveness of our algorithm when using estimated nuisance functions in stage 1. In setting A, we consider a synthetic dataset with nonlinear CATE and aim to learn a linear g (similar to Fig. 1). In setting B, we consider a non-linear but regularized g (similar to Fig. 2). Details regarding the datasets are in Appendix D.

Results. We report PEHE and policy loss for all four pseudooutcomes (PI, RA, IPW, and DR) in Fig. 4 (Setting A) and Fig. 5 (Setting B). The results demonstrate that our algorithm is effective in decreasing the policy loss compared to the baselines ( γ = 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 propensity scores (a known issue for these estimation methods; see [9]). Nevertheless, our PT-CATE algorithm leads to a better average decision performance across all estimation methods while only minimally increasing the PEHE.

Real-world data. Dataset. Here, we provide additional experimental results using the Hillstrom Email Marketing dataset of n = 64000 customers. Details regarding the dataset and our preprocessing are in Appendix E.



































=0















Results. Similar to the experiments with simulated data in Fig. 4 and Fig. 5, we plot the (estimated) PEHE and policy loss over different values of γ and compare it to the baseline CATE γ = 0 (here: for the DR-learner). The results are shown in Fig 6. The results are consistent with our experiments on synthetic data: our algorithm is effective in improving the policy loss as compared to standard CATE estimation ( γ = 0 ).

We also compare against the behavioral policy that generated the data (i.e., using the propensity score π b as a policy). For this, we report the improvement over the behavioral policy for different values of γ in Table 1. As we can see, using our PTCATE algorithm with γ = 0 . 98 can lead to a 24.45% improved response probability as compared to just using a CATE-based policy ( γ = 0) .

6 Discussion

In this paper, we showed that standard two-stage CATE estimators can be suboptimal for decision-making and propose a policy-targeted CATE (PT-CATE) to balance estimation and decision performance. Our neural algorithm improves CATE for decision-making while maintaining interpretability as CATE.

Limitations: If the second-stage model class is not restricted, our method will not lead to improvement 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 confounding. Additionally, automated decision systems based on our approach could perpetuate or 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 a principled tool for reliable, datadriven decision-making by improving the decision performance of two-stage meta-learners under model-class restrictions.

References

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.

Susan Athey and Stefan Wager. 'Policy learning with observational data'. In: Econometrica 89.1 (2021), pp. 133-161.

Andrew Bennett and Nathan Kallus. 'Efficient policy learning from surrogate-loss classification reductions'. In: ICML . 2020.

Omar Besbes, Robert Phillips, and Assaf Zeevi. 'Testing the validity of a demand model: An operations perspective'. In: Manufacturing & Service Operations Management 12.1 (2010), pp. 162-183.

Aurelien Bibaut et al. 'More efficient off-policy evaluation through regularized targeted learning'. In: ICML . 2019.

Ioana Bica et al. 'Estimating counterfactual treatment outcomes over time through adversarially balanced representations'. In: ICLR . 2020.

Matteo Bonvini, Edward H. Kennedy, and Luke J. Keele. 'Minimax optimal subgroup identification'. In: arXiv preprint arXiv:2306.17464 (2023).

Victor Chernozhukov et al. 'Double/debiased machine learning for treatment and structural parameters'. In: The Econometrics Journal 21.1 (2018), pp. C1-C68.

Alicia Curth and Mihaela van der Schaar. 'Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms'. In: AISTATS . 2021.

Alicia Curth and Mihaela van der Schaar. 'On inductive biases for heterogeneous treatment effect estimation'. In: NeurIPS . 2021.

Miroslav Dudik, John Langford, and Lihong Li. 'Doubly robust policy evaluation and learning'. In: ICML . 2011.

Carlos Fernández-Loría and Jorge Loría. 'Inferring effect ordering without causal effect estimation'. In: arXiv preprint arXiv:2206.12532 (2024).

Carlos Fernández-Loría and Foster Provost. 'Causal classification: Treatment effect estimation vs. outcome prediction'. In: Journal of Machine Learning Research 23.59 (2022), pp. 1-35.

Carlos Fernández-Loría and Foster Provost. 'Causal decision making and causal effect estimation are not the same... and why it matters'. In: INFORMS Journal on Data Science 1.1 (2022), pp. 4-16.

Stefan Feuerriegel et al. 'Causal machine learning for predicting treatment outcomes'. In: Nature Medicine (2024).

Dylan J. Foster and Vasilis Syrgkanis. 'Orthogonal statistical learning'. In: The Annals of Statistics 53.3 (2023), pp. 879-908.

Dennis Frauen, Konstantin Hess, and Stefan Feuerriegel. 'Model-agnostic meta-learners for estimating heterogeneous treatment effects over time'. In: ICLR . 2025.

Dennis Frauen, Valentyn Melnychuk, and Stefan Feuerriegel. 'Fair off-policy learning from observational data'. In: ICML . 2024.

Thomas A. Glass et al. 'Causal inference in public health'. In: Annual Review of Public Health 34 (2013), pp. 61-75.

Negar Hassanpour and Russell Greiner. 'Learning disentangled representations for counterfactual regression'. In: ICLR . 2020.

Jennifer L. Hill. 'Bayesian nonparametric modeling for causal inference'. In: Journal of Computational and Graphical Statistics 20.1 (2011), pp. 2017-2040.

Nan Jiang and Lihong Li. 'Doubly robust off-policy value evaluation for reinforcement learning'. In: ICML . 2016.

Fredrik D. Johansson, Uri Shalit, and David Sonntag. 'Learning representations for counterfactual inference'. In: ICML . 2016.

Nathan Kallus. 'Balanced policy evaluation and learning'. In: NeurIPS . 2018.

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

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.

Nathan Kallus and Masatoshi Uehara. 'Doubly robust off policy value and gradient estimation for deterministic policies'. In: NeurIPS . 2020.

Nathan Kallus and Masatoshi Uehara. 'Efficiently breaking the curse of horizon in offpolicy evaluation with double reinforcement learning'. In: Operations Research 70.6 (2022), pp. 3282-3302.

Nathan Kallus and Masatoshi Uehara. 'Intrinsically efficient, stable, and bounded off-policy evaluation for reinforcement learning'. In: NeurIPS . 2019.

Nathan Kallus and Angela Zhou. 'Policy evaluation and optimization with continuous treatments'. In: AISTATS . 2018.

Nathan Kallus et al. 'Doubly robust distributionally roust off-policy evaluation and learning'. In: ICML . 2022.

Fahad Kamran, Maggie Maker, and Jenna Wiens. 'Learning to rank for optimal treatment allocation under resource constraints'. In: AISTATS . 2024.

Edward H. Kennedy. 'Towards optimal doubly robust estimation of heterogeneous causal effects'. In: Electronic Journal of Statistics 17.2 (2023), pp. 3008-3049.

Edward H. Kennedy, Sivaraman Balakrishnan, and Larry Wasserman. 'Semiparametric counterfactual density estimation'. In: Biometrika (2023).

Kwangho Kim and José R. Zubizarreta. 'Fair and robust estimation of heterogeneous treatment effects for policy learning'. In: ICML . 2023.

Diederik P. Kingma and Jimmy Ba. 'Adam: A method for stochastic optimization'. In: ICLR . 2015.

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.

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.

Milan Kuzmanovic et al. 'Causal machine learning for cost-effective allocation of development Aid'. In: KDD . 2024.

Bryan Lim, Ahmed M. Alaa, and Mihaela van der Schaar. 'Forecasting treatment responses over time using recurrent marginal structural networks'. In: NeurIPS . 2018.

Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. 'Causal transformer for estimating counterfactual outcomes'. In: ICML . 2022.

Pawel Morzywolek, Johan Decruyenaere, and Stijn Vansteelandt. 'On a general class of orthogonal learners for the estimation of heterogeneous treatment effects'. In: arXiv preprint arXiv:2303.12687 (2023).

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. 1410-1423.

Susan A. Murphy. 'Optimal dynamic treatment regimes'. In: Journal of the Royal Statistical Society: Series B 65.2 (2003), pp. 331-355.

Razieh Nabi and Ilya Shpitser. 'Fair inference on outcomes'. In: AAAI . 2018.

Xinkun Nie and Stefan Wager. 'Quasi-oracle estimation of heterogeneous treatment effects'. In: Biometrika 108.2 (2021), pp. 299-319.

Min Qian and Susan A. Murphy. 'Performance guarantees for individualized treatment rules'. In: Annals of Statistics 39.2 (2011), pp. 1180-1210.

Donald B. Rubin. 'Estimating causal effects of treatments in randomized and nonrandomized studies'. In: Journal of Educational Psychology 66.5 (1974), pp. 688-701.

Jonas Schweisthal et al. 'Reliable off-policy learning for dosage combinations'. In: NeurIPS . 2023.

Uri Shalit, Fredrik D. Johansson, and David Sontag. 'Estimating individual treatment effect: Generalization bounds and algorithms'. In: ICML . 2017.

Adith Swaminathan and Thorsten Joachims. 'Counterfactual risk minimization: Learning from logged bandit feedback'. In: ICML . 2015.

Akira Tanimoto et al. 'Regret minimization for causal inference on large treatment space'. In: 2021.

Phillip S. Thomas and Emma Brunskill. 'Data-efficient off-Policy policy evaluation for reinforcement Learning'. In: ICML . 2016.

Daniel Tschernutter, Tobias Hatt, and Stefan Feuerriegel. 'Interpretable off-policy learning via hyperbox search'. In: ICML . 2022.

Mark J. van der Laan. 'Statistical inference for variable importance'. In: The International Journal of Biostatistics 2.1 (2006), pp. 1-31.

Mark J. van der Laan and Donald B. Rubin. 'Targeted maximum likelihood learning'. In: The International Journal of Biostatistics 2.1 (2006).

Aart van der Vaart. Asymptotic statistics . Cambridge: Cambridge University Press, 1998.

Hal R. Varian. 'Causal inference in economics and marketing'. In: Proceedings of the National Academy of Sciences (PNAS) 113.27 (2016), pp. 7310-7315.

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.

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.

Baqun Zhang et al. 'Estimating optimal treatment regimes from a classification perspective'. In: Stat 1.1 (2012), pp. 103-112.

Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. 'Learning overlapping representations for the estimation of individualized treatment effects'. In: AISTATS . 2020.

Hao Zou et al. 'Counterfactual prediction for outcome-oriented treatments'. In: ICML . 2022.

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. 4) 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 be perceived well by the reviewers.

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

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

2. Limitations

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

Answer: [Yes]

Justification: 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 are not discussed in the paper.

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

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

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

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

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

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

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

3. Theory assumptions and proofs

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

Answer: [Yes]

Guidelines:

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

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

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

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

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

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

4. Experimental result reproducibility

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

Answer: [Yes]

Justification: See Sec. 4, Sec. 5, Appendix C, Appendix D, and Appendix E. Code is provided.

Guidelines:

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

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

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

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

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

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

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

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

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

5. Open access to data and code

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

Answer: [Yes]

Justification: Code is available via anonymized GitHub and can be used to reproduce experiments.

Guidelines:

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

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

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

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

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

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

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

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

6. Experimental setting/details

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

Answer: [Yes]

Guidelines:

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

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

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

7. Experiment statistical significance

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

Answer: [Yes]

Justification: 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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

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

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

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

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

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

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

8. Experiments compute resources

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

Answer: [Yes]

Guidelines:

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

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

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

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

9. Code of ethics

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

Answer: [Yes]

Justification: The code of ethics was respected in every step.

Guidelines:

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

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

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

10. Broader impacts

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

Answer: [Yes]

Justification: See Appendix 6.

Guidelines:

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

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

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

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

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

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

11. Safeguards

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

Answer: [NA]

Justification: No such components have been used.

Guidelines:

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

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

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

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

12. Licenses for existing assets

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

Answer: [NA]

Justification: No existing licenses used.

Guidelines:

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

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

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

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

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

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

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

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

13. New assets

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

Answer: [Yes]

Justification: 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/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

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

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

14. Crowdsourcing and research with human subjects

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

Answer: [NA]

Justification: No human subjects involved.

Guidelines:

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

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

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

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

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

Answer: [NA]

Justification: Not applicable.

Guidelines:

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

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

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

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

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is 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.

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

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 to extract complex features from multimodal data. Notable advances include methods for learning representations that improve CATE estimation through balancing techniques [23, 50, 62], disentanglement strategies [20], and the incorporation of inductive biases [9, 10].

While these approaches focus on estimating the nuisance functions η , they do not 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-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, 16]. 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-weighted DR estimators [42]. Orthogonal learners are, by construction, two-stage meta-learners and thus are applicable to our proposed framework. In this work, we leverage the DR-learner outlined in Eq. (10), but extensions to the R-learner/ overlap-weighted 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 are 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 our approach to dynamic settings presents an intriguing avenue for future work. Such extensions could address challenges unique to temporal data, including time-varying confounding.

B Proofs

B.1 Proof of Theorem 4.1

Proof. Let S be the set of step functions

Let S G denote the set of step functions τ such that the optimal approximation g τ gets the sign correct everywhere, i.e.,

Note that S G is non-empty as any constant positive function is in S G .

For any τ ∈ S , we define the supremum norm as ∥ τ ∥ ∞ = sup x ∈ [ a,b ] | τ ( x ) | . We set

By definition of the supremum, there exists a sequence { τ n } n ≥ 1 ⊂ S G such that

Hence, for any fixed ϵ > 0 , one can select an index n ϵ such that

Therefore, there exists a corresponding approximation

which implies that the corresponding thresholded policy π g τnϵ is optimal, i.e.,

Because τ n ϵ ∈ S G is a step function, there exists an interval I τ nϵ ⊆ [ a, b ] of positive measure so that

on which | τ ∗ ( x ) | is nearly ∥ τ ∗ ∥ ∞ .

Let ϵ > 0 be fixed, and let δ > ϵ . We define the CATE τ ∗ as

Then, by definition of τ ∗ , we have

which implies by definition of M that τ ∗ satisfies

That is, the optimal approximation of g τ ∗ with functions in G must fail to preserve the sign of τ ∗ on a subset of positive measure E , i.e.,

̸

We now compare the plug-in policy induced by g τ ∗ , i.e.,

satisfying

against the policy obtained by thresholding an approximation of τ n ϵ ( x ) , i.e.,

On the set E , the decisions made by π g τϵ and π ∗ τ nϵ differ in a way such that

Furthermore, τ n ϵ and τ ∗ have the same sign by definition, which implies together with Eq. (21) that

because τ n ϵ ∈ S G and thus g τ nϵ ∈ Π G . Hence,

which completes the proof.

B.2 Proof of Theorem 4.3

Proof. One can show that, for all pseudo-outcomes Y η m , it holds that

(see e.g., [9] for a proof). Hence, we can apply the tower property and write

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 a reasonable PT-CATE estimator, even when the nuisance functions are estimated with errors. Importantly, we upper bound of the PT-CATE error on the nuisance errors of the respective adjustment method (pseudo-outcome). For the DR pseudo-outcome, we establish a doubly robust convergence rate .

Theorem B.1. Let g ∗ = arg min g ∈G L m γ,α,η ( g ) and ˆ g = arg min g ∈G L m γ,α, ˆ η ( g ) be the minimizers of the PT-CATE loss with ground-truth and estimated nuisances for a fixed indicator approximation α and γ ∈ [0 , 1] . We assume the following regularity condition: there exists a constant δ > 0 , so that, for all ¯ g ∈ star( G , g ∗ ) = { tg ∗ +(1 -t ) g | g ∈ G} , it holds that

where || g ∗ -ˆ g || 2 = E [( g ∗ ( X ) -ˆ g ( X )) 2 ] denotes the squared L 2 -norm and σ ′′ ( · ) denotes the second derivative of the sigmoid function. Furthermore, assume that the propensity estimator and

ground-truth response functions are bounded via p ≤ ˆ π ( x ) ≤ 1 -p and | µ a ( x ) | ≤ c for constants p, c > 0 and for all x ∈ X .

Then, for all ρ 1 , ρ 2 > 0 so that (1 -γ ) ρ 1 + γ 2 ρ 2 < 1 -γ + δ 2 γ , it holds that

where R m γ,α, ˆ η (ˆ g, g ∗ ) = L m γ,α, ˆ η (ˆ g ) - L m γ,α, ˆ η ( g ∗ ) is an optimization-dependent term, C α > 0 is a constant depending on α , and M m ˆ η,η is the (pseudo-outcome-dependent) rate term, defined via

Proof. Recall that

We can write

For L m ˆ η,α (ˆ g ) , we can do a functional Taylor expansion, i.e., there exists a ¯ g ∈ star( G , g ∗ ) with

where D g denotes the functional derivative.

For the first-order derivative, we obtain

where σ ′ ( · ) denotes the derivative of the sigmoid function.

For the second-order derivative, we obtain

where the last inequality follows from the regularity assumption.

Putting everything together, we obtain that

or equivalently

where R m γ,α, ˆ η (ˆ g, g ∗ ) = L m γ,α, ˆ η (ˆ g ) -L m γ,α, ˆ η ( g ∗ ) .

(64)

Putting everything together, we obtain

Note that

because g ∗ is a minimizer of the oracle nuisance loss L m γ,α,η . Hence,

For the different pseudo outcomes, we can write

By applying the Cauchy-Schwarz inequality, we obtain

By applying AM-GM inequality and the fact that ( a + b ) 2 ≤ 2( a 2 + b 2 ) , it holds for any ρ > 0 that

We can write this in generalized form via

Using the same arguments and the fact that we can upperbound σ ′ ( α ( X ) g ∗ ( X )) 2 α ( X ) 2 ≤ C α for some constant C α > 0 , we obtain

Hence, it holds that

or, equivalently, for ρ 1 , ρ 2 so that (1 -γ ) ρ 1 + γ 2 ρ 2 < 1 -γ + δ 2 γ , we have

B.4 Additional result bounding policy value discrency with projection error

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.

Lemma B.2. Let the true CATE be τ ( x ) = µ 1 ( x ) -µ 0 ( x ) ∈ L 2 ( P ) . For a function class G ⊂ L 2 ( P ) , define its L 2 -projection of τ as

and let ε = ∥ τ -g ∥ . Furthermore, introduce the corresponding thresholded treatment policies

and the (shifted) policy value

Then it holds that

Then

On M we have τ ( X ) g ⋆ ( X ) ≤ 0 , hence | τ ( X ) | ≤ | τ ( X ) -g ⋆ ( X ) | . Therefore

where the last inequality follows from Hölder's inequality.

Proof. Define the disagreement set

̸

C Implementation details

Estimation of nuisance functions We estimate all nuisance functions µ 1 , µ 0 , and π b with standard feed-forward neural networks using 4 layers with tanh activations. The response functions µ a 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 π b is a classification task so that we apply a 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 be known.

Second-stage model: For our second-stage model (Fig. 3), we model each of g θ and α ϕ as feedforward neural networks with 4 layers with tanh activations for g θ and ReLU activations for α ϕ . For the experiments with linear g θ (i.e., Fig. 2 and Fig. 4), we set g θ to a single linear layer. For the experiments with regularized g θ (i.e., Fig. 1 and Fig. 5), we choose a custom regularization parameter for each pseudo-outcome type that yields a misspecified initial CATE estimate. For the synthetic experiments in Fig 4 and Fig 5, we normalize g θ by applying a tanh output activation in step 2 of our learning algorithm (Algorithm 1) as we observed that this can stabilize the optimization of α ϕ . We also applied a weighting scheme in step 3 via 1 /α ϕ ( X ) to further encourage sharp indicator approximation of regions where the initial CATE is correct. We ran our algorithm for K = 1 iteration.

Hyperparameters. To ensure a fair comparison, we use the same hyperparameters for each secondstage learner across different γ and random seeds. For reproducibility purposes, we report the hyperparameters used (e.g., dimensions, learning rate) for all experiments and models (including nuisance functions) as .yaml files. 3

Runtime. For the second-stage models, training took approximately two minutes using n = 2000 samples and a standard computer with AMD Ryzen 7 Pro CPU and 32GB of RAM.

Full learning algorithm. The full learning algorithm is reported in Algorithm 1 below.

Algorithm 1: Re-targeted CATE estimation (PT-CATE)

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 ∼ U [0 , 1] from a uniform distribution. Then, we simulate binary treatments via

for some propensity score π b ( X ) . Finally, we simulate continuous outcomes via

where µ A ( X ) denotes the response function and ε = 0 . 01 the noise level.

Dataset for Fig. 1. Here, we emulate an RCT and set the propensity score to π b ( X ) = 0 . 5 . We set the response function to µ a ( x ) = a (2 σ (10 x ) -0 . 5) , where σ ( · ) denotes the sigmoid function. We sample a training dataset of size n train = 1000 and a test dataset of size n test = 3000 .

Dataset for Fig. 2. Here, we again emulate an RCT and set the propensity score to π b ( X ) = 0 . 5 . We set the response function to µ a ( x ) = a ( 1 ( x < -0 . 25)(0 . 6 sin(8( x +0 . 25))+0 . 3)+ 1 ( -0 . 25 < x < 0 . 25)(2 σ (10( x +2)) -0 . 5) + 1 ( x > 0 . 25)(0 . 5 sin(10( x -0 . 25) + 1 . 5)) . We sample a training dataset of size n train = 1000 and a test dataset of size n test = 3000 .

Dataset for Fig. 4. Here, we use the same propensity and response functions as in Fig. 1. We sample a training dataset of size n train = 2200 and a test dataset of size n test = 3000 .

Dataset for Fig. 5. Here, we set the propensity score to π b ( X ) = σ (0 . 1 x ) . We then define the response function as µ a ( x ) = a ( 1 ( x < -0 . 25)(0 . 6 sin(8( x + 0 . 25)) + 0 . 3) + 1 ( -0 . 25 < x < 0 . 25)(2 σ (10( x + 2)) -0 . 5) + 1 ( x > 0 . 25)(0 . 5 sin(10( x -0 . 25) + 1 . 5)) . We sample a training dataset of size n train = 2200 and a test dataset of size n test = 3000 .

E Details regarding real-world data

The data is taken from https://causeinfer.readthedocs.io/en/latest/data/hillstrom. html . The dataset consists of n = 64000 customers who purchased a product within the last 12 months and who were involved in an email experiment: group 1 randomly received an email advertising merchandise for men, group 2 for women, and group 3 did not receive an email (control). We study the effect of receiving a men's merchandise email ( A = 1 ) versus receiving no email at all ( A = 0 ). Covariates X include various customer features such as purchasing history. Finally, we chose Y an indicator of whether people responded to the email (by clicking on the link to the website) as our outcome Y . We split the data into a training dataset with 50% of the data, a validation set with 20% , and a test set with 30% of the data. All details regarding our data preprocessing are provided within our codebase. 4

F Additional experiments

F.1 Motivational experiments with estimated nuisance functions









































=0



=0.8



=0.95













































=0



=0.8



=0.9



=0.98

F.2 Experiments with sample splitting





































=0



















































=0

























































=0























































=0























































=0

























































=0







































=0











































=0





































F.3 Experiments with different nuisance model baselines

































=0





















































=0



















































=0











































=0













































=0

















































=0

















































=0

























































=0















































=0





















































=0



















































=0





















































=0













































=0























































=0



















































=0

























































=0













































=0





















































=0



















































=0











































=0













































=0

















































=0



















































=0

























































=0















F.4 Experiments with multivariate covariates





































=0

















































=0



















































=0



















































=0



















Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n$\\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.\n\nQUESTIONS:\n1)\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.\n\nLIMITATIONS:\nThere is no discussion of the limitations and/or potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAiming 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).\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\n[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).\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nSee weakness.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

comments: 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.

STRENGTHS AND WEAKNESSES:
1. The paper provide s a novel method to mitigate the gap between CATE estimators and the optimal decision policy.
2. The paper provides solid theoretical results to justify the proposed research questions. This is also validated in simulation analysis.

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?
2. 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.
TITLE: Official Review
CLARITY: 2
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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$.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}