Optimal Transport for Treatment Effect Estimation
Authors: Hao Wang
Jiajun Fan
Zhichao Chen
Haoxuan Li
Weiming Liu
Tianqiao Liu
Quanyu Dai
Yichao Wang
Zhenhua Dong
Ruiming Tang
Venue: NeurIPS 2023 poster (2023)
Keywords: treatment effect estimation
optimal transport
wasserstein
causal inference
counterfactual
TL;DR: A new take on OT technology in the context of causal inference, alleviating two issues that impede the effectiveness of backdoor adjustment approaches.

Abstract: Estimating individual treatment effects from observational data is challenging due to treatment selection bias. Prevalent methods mainly mitigate this issue by aligning different treatment groups in the latent space, the core of which is the calculation of distribution discrepancy. However, two issues that are often overlooked can render these methods invalid:
(1) mini-batch sampling effects (MSE), where the calculated discrepancy is erroneous in non-ideal mini-batches with outcome imbalance and outliers;
(2) unobserved confounder effects (UCE), where the unobserved confounders are not considered in the discrepancy calculation.
Both of these issues invalidate the calculated discrepancy, mislead the training of estimators, and thus impede the handling of treatment selection bias.
To tackle these issues, we propose Entire Space CounterFactual Regression (ESCFR), which is a new take on optimal transport technology in the context of causality.
Specifically, based on the canonical optimal transport framework, we propose a relaxed mass-preserving regularizer to address the MSE issue and design a proximal factual outcome regularizer to handle the UCE issue.
Extensive experiments demonstrate that ESCFR estimates distribution discrepancy accurately, handles the treatment selection bias effectively, and outperforms prevalent competitors significantly.

Optimal Transport for Treatment Effect Estimation
Hao Wang'
Zhichao Chen'
Jiajun Fan2
Haoxuan Li Zhenhua Dong"
Tianqiao Liu" Ruiming Tang"
Weiming Liu
Quanyu Dai
Yichao Wang
'Zhejiang University
TTinghua University *Peking University
*aurdue University
Huawei Noah's Ark Lab
Maohaow@zju.edu en
quanyu. dai @conncct. wrwoonnntt.ollu..lii... polyu.hk hk
Abstract
Estimating conditional average treatment effect from observational data is highly challenging due to the existence of treatment selection bias. Prevalent methods mitigate this issue by aligning distributions of different treatment groups the latent space. However, there are two critical problems that these methods fail to address: (1) mini-batch sampling effects (MSE), which causes misalignment in non-ideal batihh with outcome imbalance and outliers; (2) unobserved confounder effects (UCE), which results inaccurate discrepancy calculation due to the neglect of unobserved confounders. To tackle these problems, we propose a principled approach named Entire Space CounterFactuaa Regression (ESCFR), take on optimal transport in the context of causality. Specifically,
which is based on the framework of stochastic optimal transport, we propose a relaxed mass- preserving regularizer address the MSE issue and design proximal factual outcome regularizer to handle the UCE issue. Extensive experiments demonstrate that our proposed ESCFR can successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.
Introduction
The estimation of the conditional average treatment effect (CATE) through randomized controlled trials serves comnestone in causal inference, finding applications in diverse fields such health care [63], e-commerce [2, 43, 81], and education [15]. Although randomized controlled trials are often considered the gold standard for CATE estimation [38, 53], the associated financial and ethical constraints can make them infeasible in practice [39, 79]. Therefore, alternative approaches that rely on observational data have gained prominence. For instance, drug developers could utilize most-marketing monitoring reports to assess the efficacy of new medications rather than conducting costly clinical A/B trials. With the growing access to observational data, estimating CATE from observational data has attracted intense research interest [41, 80, 84].
Estimating CATE with observational data has main challenges: (1) missing counterfactuals, i.e., only one factual outcome out of all potential outcomes can be observed; (2) treatment selection bias, i.e,, individuals have different propensities for treatment, leading non-random treatment assign- ments and a resulting covariate shift between the treated and untreated groups [77, 79]. Traditional meta-learners [36] handle the issue of missing counterfactuals by breaking down the CATE estimation into more tractable sub-proble problemss of factual outcome estimation. However, the covariate shift makes difficult generalize the factual outcome estimators trained within respective treatment groups the entire population and thus biases the derived CATE estimator.
Beginning with counterfactual regression [67] and its remarkable performance, there are various attempts that handle the selection bias by minimizing the distribution discrepancy between the
"Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Factual ouRcome
X
R
Factual outcome
de/Ri
W
Factual
dulli
X
Indicator
X
Covariste
Representation
Outcome estimate
(a) Handling treatment selection with v=()).
(b) Architecture of ESCFR.
Figure Overview of handling treatment selection bias with ESCER. The red and blue colors signify the treated and untreated groups, respectively. (a) The treatment selection bias manifests as a distribution shift between X1 and Xo- The scatters and curves represent the units and the fitted outcome mappings. (b) ESCFR mitigates the selection bias by aligning units from different treatment groups in the representation space: R (X),, which enables Φ1 and @u to generalize across groups.
treatment groups in the representation space [13, 27, 85, 86). However, critical issues with these methods have long been neglected, which impedes them from handling the treatment selection bias. The first issue the mini-batch sampling effects (MSE). Specifically, current representation- baaed methods compute the distribution discrepancy within mini-batches instead of the entire data space, making it vulnerable to bad sampling cases. For example, the presence of an outlier in a mini- batch can falsely increase the estimate of the distribution discrepancy, thereby misleading the update of the estimator. The second issue the unobserved confounder effects (UCE). Specifically, current approaches mainly operate under the unconfoundedness assumption [48] and ignore the pervasive influence of unobserved confounders, which makes the resulting estimators biased given the existence of unobserved confounders.
Contributions and outline. In this paper, we propose an effective CATE estimator based on optimal transport, namely Entire Space CounterFactual Regression (ESCFR), which handles both the MSE and UCE issues with generalized sinkhorn discrepancy. Specifically, after presenting preliminaries in Section 2, we redefine the problem of CATE estimation within the framework of stochastic optimal transport in Section 3.1. We subsequently showcase the MSE issue faced by existing approaches in Section 3.2 and introduce mass-presss pesss-ppeerriing regularizer counter this issue. In Section 3.3, we explore the UCE issue and propose a proximal factual outcome regularizer to mitigate its impact. The architecture and learning objectives of ESCFR are elaborated upon Section 3.4, followed by the presentation of our experimental results in Section 4.
2 Preliminaries
2.1
Causal inference from observational data
This section formulates preliminaries in observational causal inference. We first formalize the fundamental elements in Definition 2.1 following the general notation convention".
Definition 2.1. Let X be the random variable of covariates, with support X and distribution P(x); Let be the random variable of induced representations, with support R and distribution P(r); Let Y be the random variable of outcomes, with support y and distribution P(y); Let T be the random variable of treatment indicator, with support T {0,1) and distribution P(T).
In the potential outcome framework [61], unit characterized by the covariates has two potential outcomes, namely (z) given it is treated and Yo(x) otherwise. The CATE is defined as the conditionally expected difference of potential outcomes as follow:
T(x) E[Y Yo x
(1)
We use uppercase letters, e.g., X to denote a random variable, and lowercase letters, e.g., x to denote an associated specific value. Letters in calligraphic font, e.g., X represent the support of the corresponding random variable, and P() represents the probability distribution of the random variable, e.g., IPX)..

---

where one of these two outcomes is always unobserved in the collected data. To address such missing counterfactuals, the CATE estimation task commonly decomposed into outcome estimation subproblems that are solvable with supervised learning method [36]. For example, 'Tleerner models the factual outcomes Y for units in the treated and untreated groups separately; S-learner regards T as covariate, and then models Y for units in all groups simultaneously. The CATE estimate is then the difference of the estimated outcomes when T set to treated and untreated.
Definition 2.2. Let vb X R. be a mapping from support X to R. ie., Vx r= 1b(x)) R. Let o: R + y be a mapping from support R x 'T to y, Le., it maps the representations and treatment indicator to the corresponding factual outcome. Denote Y1= Φ1 (R) and Yo (R), where abbreviate @RR, T 1) and (R,, T 0) to $1 (R) and Φo (R), respectively, for brevity.
TARNet [67] combines T-learnrr and S-learner to achieve better performance, which consists of representation mapping & and an outcome mapping & as defined in Definition 2.2. For an unit with covariates x, TARNet estimates the CATE
f0.0(2) = - = (v())), =o)))())))
where v is trained over all individuals, Φ1 and @u are trained over the treated and untreated group, respectively. Finally, the quality of CATE estimation is evaluated with the precision estimation of heterogeneous effect (PEHE) metric:
()) Sfx (fo,o())-<x))) P(x) dr.
However, according Figure (()), the treatment selection bias causes distribution shift of covariates across groups, which misleads Φ1 and ón to overfit their respective group's s properties and generalize poorly to the entire population. Therefore, the CATE estimate by these methods would be biased.
2.2 Discrete optimal transport and Sinkhorn divergence
Optimal transport (OT) quantifies distribution discrepancy as the minimum transport cost, offering tool to quantify the selection bias in Figure 1(a). Monge [49] first formulated OT as finding an optimal mapping between two distributions. Kantorovich [34] further proposed a more applicable formulation in Definition 2.3, which can be seen as generalization of the Monge problem.
Definition 2.3. For empirical distributions a and 3 with n and m units, respectively, the Kantorovich problem aims to find feasible plan TRxx which transports
at the minimum cost:
(a, 3) min xEll(a,
(D,T), ()a, = { R 11mm a 1,, 11, =b}
(4)
where W(a, 3) E E is the Wasserstein discrepancy between and 3; D E R is the unit-wise distance between a and which is implemented with the squared Euclidean distance; and indicate the mass of units in a and 3, and II is the feasible transportation plan set which ensures the mass-preserving constraint holds.
Since exact solutions to (4) often come with high computational costs [5], researchers would com- monly add an entropic regularization to the Kantorovich problem follow:
Wh (a. B) = (D, x² := arg min (D, π) <H(T), H(T) == - Σ T (l)g(x.,))1)) (5) xe))(a.9)
which makes the problem c-convex and solvable with the Sinkhorn algorithm [19]. The Sinkhorn algorithm only consists of matrix-vector products, making it suited be accelerated with GPUs.
3 Proposed method
In this section, we present the Entire Space CounterFactual Regression (ESCFR) approach, which leverages optimal transport to tackle the treatment selection bias. We first illustrate the stochastic optimal transport framework for distribution discrepancy quantification across treatment groups, and demonstrate its efficacy for improving the performance of CATE estimators. Based on the framework, we then propose relaxed mass- mass-preserving regularizer to address the sampling effect, and proximal factual outcome regularizer to handle the unobserved confounders. We finally open a new thread to summarize the model architecture, learning objectives, and optimization algorithm.

---

3.1 Stochastic optimal transport for counterfactual regression
Representation-bun approaches mitigate the treatment selection bias by calculating distribution discrepancy in the representation space and then minimizing it. OT is preferred method to quantify the discrepancy due its numerical advantages and flexibility over competitors. It is numerically stable in the cases where other methods, such as o-divergence (e.g., Kullback-Leibler divergence), fails [64]. Compared with adversarial discrepancy measures [7, 33, 87], the can calculated efficiently and integrated naturally with the traditional supervised learning framework.
where (r) (r) are the distributions of representations in treated and untreated groups, respec- tively, induced by the mapping
Formally, we denote the OT discrepancy between treatment groups as W (TT-(() and
(æ:). The discrepancy is differentiable with respect [24], and thus can be minimized by updating the representation mapping t with gradient-basee optimizers.
Definition 3.1. Let (x) {=:-->}²= and (x)= { {x:]-=) J
be the empirical distributions of covariates at mini-batch level, which contain n treated units and m untreated units, respectively; 1BT (r) and fsT=0 (r) be that of representations induced by the mapping = (xxx) Definition 2.2.
However, since prevalent neural estimators mainly update parameters with stochastic gradient meth- ods, only a fraction of the units is accessible within each iteration. A shortcut in this context is to calculate the discrepancy stochastic mini-batel level:
W (r)
The effectiveness of this shorteut is investigated by Theorem 3.1 (refer to Appendix A.4 for proof), which demonstrates that the PEHE be optimized by iteratively minimizing the estimation error of factual outcomes and the mini-batch group discrepancy in (6).
Theorem 3.1. Let vp and d be the representation mupping and factual outcome mapping, respectively; W be the group discrepancy at a mini-batch level. With the probability of at least 1 we have:
6PEHE 21cF--((0) (v, + ep T=0 (vb, B 20 O( )
(7)
8N
where T-1 and c are the expected errors of factual outcome estimation, N the batch size, dy is the variance of outcomes, B& is constant term, and O()) is sampling complexity term.
3.2 Relaxed mass-preserving regularizer for sampling effect
Although Theorem 3.1 guarantees that the empirical OT discrepancy (6) bounds the PEHE, the sampling complexity term O()) inspires to investigate the potential risks of bad cases caused by stochastic sampling. Precisely, the term (-)) results from the discrepancy between the entire population and the sampled mini-batch units (see (30) and (32) in Appendix A) which is highly dependent on the un- controllable sampling quality. Therefore, reliable discrepancy measure should be robust to bad sam- pling cases, otherwise the resulting vulnerability would impede us from quantifying and minimizing the actual discrepancy.
(a) Ideal
(b) Imbalance
(c) Outlier
Figure 2: Optimal transport plan (upper) and its geometric interpretation (down) in three cases, where the connection strength represents the transported mass. Different colors and vertical positions indicate different treatments and out- comes, respectively.
We consider three sampling cases in Figure 2, where the transport strategy is reasonable and applicable in the ideal sampling case in Figure 2(a). However, the transport strategy calculated with (6) is vulner- able to disturbances caused by non-ideal sampling conditions. For instance, Figure 2(b) reveals false matching of units with disparate outcomes in the sampled mini-batch where the outcomes the two groups are imbalanced; Figure 2(c) showcases false matching of the
outliers with normal units, causing substantial disruption of the

---

transport strategy for other units. Therefore, the vanilla OT in (6) fails to quantify the group discrep- ancy for producing erroneous transport strategies non-ideal bini-batches, thereby misleading the update of the representation mapping vb. We term this issue the mini-batch sampling effect (MSE).
Furthermore, this MSE issue is not exclusive OT- methods but is also prevalent in other representation-base approaches. Despite this, OT offers distinct ad advantage: it allows the formaliza- tion of the MSE issue through its nass-preservation constraint, as indicated in (5). This constraint mandates a match for every unit in both groups, regardless of the real-world scenarios, complicating the transport of normal units and the computation of true group discrepancies. This issue is further exacerbated by small batch sizes. Such formalizability through mass prisssseesevvttoon provides lever for handling the MSE issue, setting OT apart from other representation based methods [85, 86).
An intuitive approach to mitigate the MSE issue is to relax the marginal constraint, i.e., to allow for the creation and destruction of the mass of each unit. To this end, inspired by unbalanced and weak transport theories [14, 65], a relaxed mass- mass-pres regularizer (RMPR) is devised in Definition 3.2. The core technical point is to replace the hard marginal constraint in (4) with soft penalty in (9) for constraining transport strategy. On the basis, the stochastic discrepancy is calculated as
W Wee (TTT-(()) (+), (r))
(8)
where the hard mass muass-peservation constraint is removed to mitigate the MSE issue. Building on Lemma in Fatras et al. [23]. we derive Corollary 3.1 to investigate the robustness of RMPR to sampling effects, showcasing that the effect of mini-batch outliers is upper bounded by a constant.
Definition 3.2. For empirical distributions a and 3 with n and m units, respectively, optimal transport with relaxed mass-preserving constraint seeks the transport strategy πI at the minimum cost:
WF (a, B) = (D, क arg (D.x) eH((T) + (<DKL (T a) + ( b))
where De E RR.. is the unit-wise distance, and a, indicate the mass units in a and B, respectively. Corollary 3.1. For empirical distributions a, 3 with n and m units, respectively, adding an outlier a' to a and denoting the disturbed distribution as d', we have
w0.0 (a', -..> (a. s 2x(1 - Eucold'th) 2 )(( 1),
(10)
which is upper bounded by 2x/ (n= + 1). vy0.s is the unbalanced discrepancy as per Definition 3.2.
In comparison with alternative methods [9, 83] to relax the marginal constraint, the RMPR imple- mentation in Definition 3.2 enjoys a collection of theoretical properties [65] and can be calculated via the generalized Sinkhorn algorithm [14]. The calculated discrepancy differentiable w.r.t. va and thus can be minimized via stochastic gradient methods in end-to-end manner.
3.3 Proximal factual outcome regularizer for unobserved confounders
Existing representation-buse methods fail to clim- inate the treatment sclection bias due to the unob- served confounder effects (UCE). Beginning with CFR [67], the unconfoundedness assumption A.1 (see Appendix A) is often taken to circumvent the UCE issue [75]. In this context, given two units ri E P7-1 P.PT7-1 (r) and rj E (r), for instance, OT discrepancy in Definition 3.2 calculates the unit- wise distance as D., ||r: - 121 If Assump- tion A. I holds, it climinates the treatment selection bias since it blocks the backdoor path X - T in Figure 3(a) by balancing confounders in the latent space. However, Assumption A. is usually violated due to the existence of unobserved confounders as per Figure 3(b,, which hinders existing methods from handling treatment selection bias since the backdoor path X' + T is not blocked. The existence of unobserved covariates X' also makes the transport strategy in vanilla OT unidentifiable, which invalidates the calculated discrepancy.
x
X
(a)
(b)
Figure 3: Causal graphs with (a) and w/o (b) the unconfoundedness assumption. The shaded node indicates the hidden confounder X".
To mitigate the effect of X', we introduce a modification to the unit-wise distance. Specifically, we observe that given balanced X and identical T,, the only variable reflecting the variation of X' is the

---

outcome Y. As such, resonating with Courty et al. [16], we calibrate the unit-wise distance D with potential outcomes as follow:
-m + T0 - - T 1lu y "13,
(11)
where controls the strength of regularization. The underlying regularization is straightforward: units with similar (observed and unobserved) confounders should also have similar potential outcomes. As such, for pair of units with similar observed covariates, i.e,, Ir= Il2 a O, if their potential outcomes given the same treatment / = {0,1) differ greatly, i.e, llv1 - wÿl >> O, their unobserved confounders should likewise differ significantly. The vanilla OT technique in (6) where Dij Ir; rrl
I2 would incorrectly match this pair, generate a false transport strategy. and consequently misguide the update of the representation mapping vh. In contrast, OT based on D would not match this pair as the difference of unobserved confounders is compensated with that of potential outcomes.
Moving forward, since - and yf in (11) are unavailable due to the missing counterfactual outcomes, the proposed proximal factual outcome regularizer (PFOR) uses their estimates instead. Specifically, let ğ and 0ū be the estimates of y y and 1/'-1, respectively, PFOR refines (11) as
D Ir. -ry1² 7 [1 y5 ||ū y F],
i (o()),, 01
(12)
Additional justifications, assumptions and limitations of PFOR are discussed in Appendix D.3.
3.4 Architecture of entire space counterfactual regression
The architecture ESCER is presented in Figure 1(b), where the covariate X is first mapped to the representations R with v-(), and then the potential outcomes with e(-). The group discrepancy W calculated with the optimal transport equipping with the RMPR in (9) and PFOR in (12).
The learning objective is to minimize the risk of factual outcome estimation and the group discrepancy.
Given mini-batch distributions (x=) and (x)==(x) in Definition 3.1, the risk of factual outcome estimation following [68] can be formulated as
LF o) = E (1) I6 (i (v(z;)) vil 2 +
=, (0)(v(;))) - y3
(13)
where yi and yj the factual outcomes for the corresponding treatment groups. The discrepancy is:
(v²) := W (IT (r) (r),
(14)
which in general the optimal transport with RMPR in Definition 3.2, except for the unit-wise distance D calculated with the PFOR in (12). Finally, the overall learning objective of ESCFR is
ESCER == (F o) A (v),
(15)
where \ controls the strength of distribution alignment, controls the entropic regularization in (5), controls RMPR in (9), and controls PFOR in (12). The learning objective (15) mitigates the selection bias following Theorem 3.1 and handles the MSE and UCE issues.
The optimization procedure consists of three steps as summarized in Algorithm 3 (see Appendix B). First, compute 5.K.7 by solving the OT problem in Definition 3.2 with Algorithm (see Appendix B), where the unit-wise distance is calculated with D Second, compute the discrepancy in (14) as (mr- D%), which is differentiable v and é. Finally, caleulate the overall loss in (15) and update u and & with stochastic gradient methods.
Experiments
4.1 Experimental setup
Missing counterfactuals impede the evaluation of PEHE using observational benchmarks. Therefore, experiments are conducted on two semi-synthetic benchmarks [67, 85), i.e,, IHDP and ACIC. Specifically, IHDP is designed to estimate the effect of specialist home visits on infants' po- tential cognitive scores, with 747 observations and covariates. ACIC comes from the collaborative perinatal project [51], and includes 4802 observations and 58 covariates.
Datasets.

---

Table 1: Performance (meanistd) on the PEHE and AUUC metries. marks the baseline estimators that ESCFR outperforms significantly at p-value < 0.05 over paired samples t-test.
Datieet
(PERE)
HHDP
ACIC (AUUC)
IHDP (AUUC)
Model
In-sample
Out-sample
In-sample
Out-sample
In-sample
Out-sample
In-sample
Out-sample
OLS
3.749=0.080
4 4.34000.177
3.85626.015
5.67449.026
0.84320.007
0.49620.017
0.652s0.050
0.492=0.032
R.Forest
3.597=0.064"
165"
2.635+3.598
4.67++9.291
0.902-9.016
0.702+0.026'
0.736+0.142
0.661=0.259
S. S.Leamer
3,572=0.269'
3 3.63620.25 254"
1,70611 1000
3.038:5.31 319
0.90520.641
0.62740.014'
0.63550.118
0.702:0.330
T.Leamee
3.429 142"
1.6660.244
1.96741.136
2.73023.627
0.846:20019
0.632+0.020'
0.65140.179
0.707.0.335
TARNet
3.236=00.266'
3.254+0. 150"
0.749+0.29
1.788+2.812
0.886=0.046
0.662+0.014
0.654+0.184
0.71 120.29
C.Fortst
3.44920. 101"
3. 3.19620.177" 177"
4.0181 602"
4. 18618.6777
0.717:0.005
0.70910.018
0.64310.11
0.69510. 294
K-NN
5.605=0. 168"
5.892=0.138
222042.233'
4.319
0.892=0.007
0.507+0.034'
0.725
0.668:0.299
O.Forest
8.094-4.669'
4.148+2.223
2.605+2.418*
3.136+5.644
0.744=0.013
0.69940.022'
0.662+0.157
0.702+0.325
PSM
5.228=0,154*
5.09420,301"
3,21924,3352
4.63428.374
0.89-100.100
0.745=0.021
0.74020.149
0.681
3.345=0,233*
3.368:0. 176"
0. 0,70940.330
1.06622.897
0.89200.03
0.45:00013"'
0.654x0.184
0,711:0.329
CFR-MMD
3.18220,174"
3.35740.321
0.77740.327
1.791s2.741
0.871
0.69910.017
0.65510.183
0.71010.329
CER WASS
3.128=0,263"
3.207=0. 169
0.657+0.673
115
0.873=00.029
0.66900.018
0.656*0.188 187
0.715=0.311
ESCER
2.252:0.297
2.31620.613
0.502:0.252
1.282:22312
0.79620.030
0.75420.021
0.66510.166
0.719:0.311
Baselines. We consider three groups of baselines. (1) Statistical estimators: least square regression with treatment covariates (OLS), random forest with treatment covariates Forest),) a single network with treatment as covariates (S. learerr [36]), separate neural regressors for each treatment group (T.learnrr [36]). TARNet [67]: (2) Matching estimators: propensity score match with logistic regression (PSM [60D, k-nearest neighbor (k-NN [18]), causal forest Fores [73]), orthogonal forest (O.Forest [73]; (3) ippresentation-based estimators: balancing neural network (BNN (31)) D, counterfactual regression with MMD (CFR-MMD) and Wasserstein discrepancy (CFR-WASS) [67].
Training protocol. A fully connected neural network with two 60-dimensional hidden layers is selected to instantiate the representation mapping v and the factual outcome mapping é for ESCFR and other neural network based baselines. To ensure fair comparison, all neural models are trained for a maximum of 400 epochs using the Adam optimizer, with the patience of early stopping being 30. The learning rate and weight decay are set to le-3 and le respectively. Other settings of optimizers follow Kingma and Ba [35]. We fine-tune hyperparameters within the range in Figure 5, validate performance every two epochs, and save the optimal model for test.
Evaluation protocol. The PEHE in (3) the primary metric for performance evaluation [67, 85]. However, it is unavailable in the model selection phase due to missing counterfactuals. As such, we use the area under the uplift curve (AUUC) [2] to guide model selection, which evaluates the ranking performance of the CATE estimator and can be computed without counterfactual outcomes. Although AUUC is not commonly used in treatment effect estimation, we report it as an auxiliary metric for reference. The within-sample and out-of- results are computed on the training and test set, respectively, following the common settings [44, 45, 67, 85).
4.2
Overall performance
Table
1 compares ESCER and its competitors. Main observations are noted as follows.
Statistical estimators demonstrate competitive performance on the PEHE metric, with neural estimators outperforming linear and random forest methods due to their superior ability to capture non-linearit, In particular, TARNet, which combines the advantages of T-leamnr and S-learner achieves the best overall performance among statistical estimators. However, the circumvention to treatment selection bias results in inferior performance.
Matching methods such as PSM exhibit strong ranking performance, which explains their popularity in counterfactual ranking practice. However, their relatively poor performance on the PEHE metric limits their applicability in counterfactual estimation applications where accuracy of CATE estimation is prioritized, such as advertising systems.
Representation- based methods mitigate the treatment selection bias and enhance overall perfor- mance. In particular, CFR-WASS reaches an out-of-sample PEHE of 3.207 on ACIC, advancing most statistical methods. However, it utilizes the vanilla Wasserstein discrepancy, wherein the MSE and UCE issues impede it from solving the treatment selection bias. The proposed ESCFR

---

Table 2: Ablation study (meanistd) the ACIC benchmark. "*" marks the variants that ESCER outperforms significantly at p-vuuee < 0.01 over paired samples t-test.
In-samplc
Out-sample
SOT
RMPR
PTOR
PEHE
AUUC
PEHE
AUUC
3.2367
3,254240.1505
0.6624+0.00
X
0.6698-0.01877
2.645900.2777
0.836640.0266
2.768840 4009
0.7099+0.0157
2.5705.0.3403*
0.8270a0.0541
2.633000 4672
0.7110:0.0257
2.2550=0.29955
0.7968#0.0307
2.3165+0.6136
0.7542+0.0202
achieves significant improvement over most metrics compared with various prevalent baselines" Combined with the comparisons above, attribute its superiority the proposed RMPR and PFOR regularizers, which accommodate ESCFR to the situations where MSE and UCE exist.
4.3 Ablation study
In this section, to further verify the effectiveness of individual components, an ablation study is conducted on the ACIC benchmark in Table 2. Specifically. ESCFR first augments TARNet with the stochastic optimal transport to align the confounders in the representation space, described in Section 3.1, which reduces the out-of-sample PEHE from 3.254 to 3.207. Subsequently, it mitigates the MSE issue with RMPR as Section 3.2 and the UCE issue with PFOR per Section 3.3, which reduces the out-of-sample PEHE to 2.768 and 2.633, respectively. Finally, ESCFR combines the RMPR and PFOR in unified framework as detailed in Section 3.4, which further reduces the value PEHE and advances the best performance other variants.
4.4 Analysis of relaxed mass- mass-preserving regularizer
Most prevalent methods fail to handle the label imbalance and mini-batch outliers Figure 2(b- c). Figure 4 shows the transport plan generated with RMPR in the same situations, where RMPR alleviates the MSE issue in both bad cases. Ini- tially, RMPR with x 50 presents similar matching scheme with Figure 2, since in this setting the loss of marginal constraint strong, and the solution thus similar to that of the vanilla OT problem in Def- inition 2.3; RMPR with k 10 further looses the marginal constraint and avoids the incorrect match- ing of units with different outcomes; RMPR with 2 further gets robust to the outlier's interference and correctly matches the remaining units. This success is attributed to the relaxation of the mass- preserving constraint according to Section 3.2.
(a) 50
(b) 10
(c) K
Figure 4: Geometric interpretation of OT plan with RMPR under the outcome imbalance (up- per) and outlier (down) settings. The dark area indicates the transported mass of a unit, i.e.. marginal of the transport matrix æ. The light area indicates the total mass.
Notably, RMPR does not transport all mass of a unit. The closer the unit is to the overlapping zone in batch, the greater the mass is transferred. That is, RMPR adaptively matches and pulls closer units that are close to the overlapping region, ignoring outliers, which mitigates the bias of causal inference methods in cases where the positivity assumption does not strictly hold. Current approaches mainly achieve by manually cleaning the data or
dynamically weighting the units [32], while RMPR naturally implements it via the soft penalty in (9).
We further investigate the performance RMPR under different batch sizes and in Appendix D.2 to verify the effectiveness of RMPR more thoroughly.
3 An exception would be the within-sample sample AUUC, which is reported over training data and thus easy to be overfitted, This metric is not critical as the factual outcomes are typically unavailable in the inference phase. We mainly rely out-of-sample AUUC instead evaluate the ranking performance and perform model selection.

---

3.0
3.2-
In-sample Out-sammpe
In-sample Sumps
In-sample Sum-lempll
In-sample
2.9
Out-sample
PEHE
3.0
PEHE
2.8
2.8
2.7
2.6
2.6
0.5
1.5 2.5 3.0 (b)
0.1 0.5 1.
1..0
5.0
10.0 15.0 20.0 25.0
(a)
(c)
(d)
Figure 5: Parameter sensitivity of ESCFR, where the lines and error bars indicate the mean values and 90% confidence intervals, respectively. (a) Impact of alignment strength (A). (b) Impact of entropic regularization strength e. (c) Impact of PFOR strength y (x10*). (d) Impact of RMPR strength k.
4.5 Parameter sensitivity study
In this section, we investigate the role of four critical hyperparameters within the ESCFR frame- work-namely, X, e, K, and 7. These hyperparameters have a profound impact on learning objectives and the performance, as substantiated by the experimental results presented in Figure 5.
First, vary À investigate the efficacy of stochastic optimal transport. Our findings suggest that gradual increase in the value of À leads to consistent improvements. For instance, the out-of uut-of-sample PEНE diminishes from 3.22 distribution balancing within a multi-task learning framework can result in compromised factual outcome estimation and, consequently, suboptimal CATE estimates. The next parameter scrutinized is which serves as an entropic regularizer. Our study shows that a larger accelerates the computation of optimal transport discrepancy [24]. However, this computational advantage comes at the cost of skewed transport plan, as evidenced by an increase in out-of-sample PEHE with oscillations.
0 roughly 2.85 at À 1.5. Nonetheless, overly emphasizing
Following this, explore the hyperparameters and K, associated with PFOR and RMPR, respec- tively. Extending upon vanilla optimal transport, we incorporate unit-wise distance via (12) and assess the impact of PFOR. Generally, the inclusion of PFOR positively influences CATE estimation. However, allocating excessive weight to proximal outcome distance compromises the performance, it dilutes the role of the unit-wise distance in the representation space. Concurrently, modify the transport problem with (9) to examine the implications of RMPR. Relieving the mass-preserving constraints through RMPR significantly enhances CATE estimation, but excessively low value of 6 hampers performance as fails to ensure that the representations across treatment groups are pulled closer together in the optimal transport paradigm.
Related works
Current research aims to alleviate treatment selection bias through the balancing of distributions between treated and untreated groups. These balancing techniques can broadly categorized into three categories: reweighting- based, matching- based, and popresentationnbasee methods.
Reweighting- based methods mainly employ propensity scores achieve global balance between groups. The core procedure comprises two steps: the estimation of propensity scores and the construction unbiased estimators. Propensity scores are commonly estimated using logistic regression [11, 20, 37, 88]. To enhance the precision of these estimates, techniques such as feature selection [69, 76, 77]. joint optimization [88, 89], and alternative training techniques [92] have been adopted. The unbiased estimator is exemplified by the inverse propensity score method [59], which inversely re- re-weights each units with the estimated propensity scores. While theoretically unbiased, it suffers from high variance with low propensity and bias with incorrect propensity estimates [39, 40, 74]. To alleviate these drawbacks, doubly robust estimators and variance reduction techniques have been introduced [26, 42, 58]. Nevertheless, these methods remain constrained by their reliance on propensity scores, affecting their efficacy in world applications.
Matching- based methods aim to match comparable units from different groups to construct locally balanced distributions. The key distinctions between representative techniques [8, 46, 60, 79] lie in their similarity measures. A notable exemplar is the propensity matching approach [60], which

---

computes unit (dis)similarity based on estimated propensity scores. Notably, Tree- methods [73] can also be categorized as matching approaches. but use adaptive similarity measures. However, these techniques are computationally intensive, limiting their deployment in large-scale operations.
Representation- methods aim construct mapping feature space where distributional dis- crepancies are minimized. The central challenge lies in the accurate calculation of such discrepancies. Initial investigations focused on maximum mean discrepancy and vanilla Wasserstein discrepancy [31, 67], later supplemented by local similarity preservation [85, 86]. feature selection [13, 27], representation decomposition [27, 80] and adversarial training [87] mechanisms, Despite their suc- cess, they fail under specific but prevalent conditions, such as outlier fluctuations [23] and unlabeled confounders [90], undermining the reliability of calculated discrepancies.
The recent exploration of optimal transport in causality [71] has spawned innovative reweighting- based [21], matching-based [ [10] and representation--asee methods [44, 45]. For example, Li et al. [45] utilize OT to align factual and counterfactual distributions; Li et al. [44] and Ma et al. [48] use OT to reduce confounding bias. Despite these advancements, they largely adhere the vanilla Kantorovich problem that corresponds to the canonical Wasserstein discrepancy, akin to Shalit et al. [67]. Adapting OT problems meet the unique needs of treatment effect estimation remains an open area for further research.
Conclusion
Due to the effectiveness mitigating treatment selection bias, representation learning has been primary approach estimating individual treatment effect. However, existing methods neglect the mini-batch sampling effects and unobserved confounders, which hinders them from handling the treatment selection bias. A principled approach named ESCFR is devised based on generalized OT problem. Extensive experiments demonstrate that ESCFR can mitigate MSE and UCE issues, and achieve better performance compared with prevalent baseline models.
Looking ahead, two research avenues hold promise for further exploration. The first avenue explores the use of normalizing flows for representation mapping, which offers the benefit of invertibility [12] and thus aligns with the assumptions set forth by Shalit et [67]. The second avenue aims apply our methodology to real-world industrial applications, such as bias mitigation in recommendation systems [74], an area currently dominated by high-variance reweighting methods.
Acknowledgements
This work is supported by National Key R&D Program of China (Grant No. 2021YFC2101000, National Natural Science Foundation of China (62073288, 12075212, 12105246, 11975207) and Zhejiang University NGICS Platform.

---

References
[1] Jason M. Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. In NeurlPS, pages 1964-19744 2017.
[2] Artem Betlei, Eustache Diemert, and Massih-Reza Amini, Uplift modeling with generalization guarantees. In SIGKDD, pages 55-65, 2021.
[3] Jose H. Blanchet, Arun Jambulapati, Carson Kent, and Aaron Sidford. Towards optimal running times for optimal transport. CoRR, foss/1810.07717 2018.
[4] François Bolley, Arnaud Guillin, and Cédric Villani. Quantitative concentration inequalities for empirical measures non-compact spaces. Probability Theory and Related Fields, 137(3): 541-593, 2007.
[5] Nicolas Bonneel, Michiel van de Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. ACM Trans. Graph., 30(6):158, 2011.
[6] M Alan Brookhart, Sebastian Schneeweiss, Kenneth Rothman, Robert Glynn, Jerry Avorn, and Til Stürmer. Variable selection for propensity score models. American journal of epidemi- ology, 163(12): 149-1156, 2006.
[7] Aaron Chadha and Yiannis Andreopoulos. Improved techniques for adversarial discriminative domain adaptation.
IEEE Trans. Image Process., 29:2622-2637, 2020.
[8] Yale Chang and Jennifer G. Dy. Informative subspace learning for counterfactual inference. In AAAI, pages 1770-1776, 2017.
[9] Lactitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced optimal transport through non-negative penalized linear regression. In NeurlPS, pages 23270- 23282, 2021.
[10] Arthur Charpentier, Emmanuel Flachaire, and Ewen Gallic. Optimal transport for counterfactual estimation: A method for causal inference. arXiv preprint vrXiv:2301.07755 2023.
[11] Jiawei Chen, Hande Dong, Yang Qiu, Xiangnan He, Xin Xin, Liang Chen, Guli Lin, and Keping Yang. Autodebias: Learning to debias for recommendation. In SIGIR, pages 21-30, 2021.
[12] Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In NewrIPS, pages 6572-6583, 2018.
[13] Mingyuan Cheng, Xinru Liao, Quan Liu, Bin Ma, Jian Xu, and Bo Zheng. Learning disentangled representations for counterfactual regression via mutual information minimization. In SIGIR, pages 1802-1806, 2022.
[14] Lenaic Chizat, Gabriel Peyré, Berhard Schmitzer, and François-Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Math. Comput., 876314):2533-2669 2018.
[15] José M Cordero, Víctor Cristóbal, and Daniel Santín, Causal inference on education policies: A survey of empirical studies using pisa, timss and pirls. J. Econ. Surv., 32(3):878-99) 995,, 2018.
[16] Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In NeurlPS, pages 3730-3739, 2017.
[17] Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Trans. Pattern Anal. Mach. Intell., 390)): 2017.
[18] Richard Crump, V Joseph Hotz, Guido W Imbens, and Oscar A Mitnik. Nonparametric tests for treatment effect heterogeneity. Rev. Econ. Stat., 90(3):389-445, 2008.
[19] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. NeurIPS, pages 2292-2300, 2013.

---

[20] Quanyu Dai, Haoxuan Li, Peng Wu, Zhenhua Dong, Xiao-Hua Zhou, Rui Zhang, Rui Zhang, and Jie Sun. generalized doubly robust learning framework for debiasing post-click conversion rate prediction. In SIGKDD, pages 252-262, 2022.
[21] Eric Dunipace. Optimal transport weights for causal inference. CoRR, abs/2109.01199, 2021.
[22] Pavel E. Dvurechensky, Alexander V. Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent better than by sinkhorn's 's algorithm. In
ICML, volume 80, pages 1366-13113 1375, 2018.
[23] Kilian Fatras, Thibault Séjouréé, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport: applications domain adaptation. In ICML, volume 139, pages 3186-3197, 2021.
[24] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélic Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T. H. Gayraud, Hicham Janati, Alain Rakotomamoniy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: python optimal transport. J. Mach. Learn, Res., 22:78:1-78:88 2021.
[25] Sergio Garrido, Stanislav Borysov, Jeppe Rich, and Francisco Pereira. Estimating causal effects with the neural autoregressive density estimator. Journal of Causal Inference, 91):211-2228 2021.
[26] Siyuan Guo, Lixin Zou, Yiding Liu, Wenwen Ye, Suqi Cheng, Shuaiqiang Wang, Hechang Chen, Dawei Yin, and Yi Chang. Enhanced doubly robust learning for debiasing post-click conversion rate estimation. In SIGIR, pages 275-284. 2021.
[27] Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In ICLR. OpenReview.net, 2020.
[28] Guido W Imbens and Donald Rubin. Causal inference statistics, social, and biomedical sciences. Cambridge University Press, 2015.
[29] Arun Jambulapati, Aaron Sidford, and Kevin Tian. direct o11 /epsilln) iteration parallel algorithm for optimal transport. In NeurlPS, pages 11355-11366 2019.
[30] Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How to escape saddle points efficiently. In ICML, pages 1724-1732, 2017.
[31] Fredrik D. Johansson, Uri Shalit, and David A. Sontag. Learning representations for counterfac- tual inference. In /CML,, pages 3020-3029, 2016.
[32] Fredrik D. Johansson, Uri Shalit, Nathan Kallus, and David A. Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. CoRR, Abs/2001.074222 2020.
[33] Nathan Kallus. Deepmatch: Balancing deep covariate representations for causal inference using adversarial training. In ICML, volume 119, pages 5067-5077, 2020.
[34] Leonid V Kantorovich. On the translocation of masses. J. Math. Sci., 133(4):1381-11) 2006.
[35] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[36] Sören Reinhold Künzel, Jasject Sekhon, Peter Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proc. Natl. Acad. Sei. U. S. A., 116 (10):45664116, 2019.
[37] Jae-woong Lee, Seongmin Park, and Jongwuk Lee. Dual unbiased recommender learning for implicit feedback. In SIGIR, pages 1647-1651, 2021.

---

[38] Haoxuan Li, Yanghao Xiao, Chunyuan Zheng, and Peng Wu. Balancing unobserved con- founding with a few unbiased ratings in debiased recommendations. In WWW, pages 305-315, 2022.
[39] Haoxuan Li, Quanyu Dai, Yuru Li, Yan Lyu, Zhenhua Dong, Xuao-Hua Zhou, and Peng Wu. Multiple robust learning for recommendation. In Brian Williams, Yiling Chen, and Jennifer Neville, editors, AAAI, pages 4417-4425, 2023.
2023.
[40] Haoxuan Li, Yanghao Xiao, Chunyuan Zheng, Peng Wu, and Peng Cui. Propensity matters: Measuring and enhancing balancing for recommendation. In ICML, volume 202 of Proceedings of Machine Learning Research, pages 20182-20194, PMLR, 2023.
[41] Haoxuan Li, Chunyuan Zheng, Yixiao Cao, Zhi Geng, Yue Liu, and Peng Wu. Trustworthy policy learning under the counterfactual no-harm criterion. In ICML, volume 202 of Proceedings of Machine Learning Research, pages 10575-20599 PMLR, 2023.
[42] Haoxuan Li, Chunyuan Zheng, and Peng Wu. Stabledr: Stabilized doubly robust learning for recommendation data missing not at random. In ICLR, 2023.
[43] Haoxuan Li, Chunyuan Zheng, Peng Wu, Kun Kuang, Yue Liu, and Peng Cui. Who should be given incentives? counterfactual optimal treatment regimes learning for recommendation. In KDD, pages 1235-1247. ACM, 2023.
[44] Qian Li, Zhichao Wang, Shaowu Liu, Gang Li, and Guandong Xu. Deep treatment-adaptive network for causal inference. VLJDBJ, pages 1-16, 2022.
[45] Qian Li, Zhichao Wang. Shaowu Liu, Gang Li, and Guandong Xu. Causal optimal transport for treatment effect estimation. IEEE Trans. Neural Networks Learn. Syst., 34(8):4083-409) 2023.
[46] Sheng Li, Nikos Vlassis, Jaya Kawale, and Yun Fu. Matching via dimensionality reduction for estimation treatment effects in digital marketing campaigns. In IJCAI, pages 3768-3774, 2016.
[47] Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In ICML, volume 97, pages 3982-3991, 2019.
[48] Jing Ma, Mengting Wan, Longqi Yang. Jundong Li, Brent Hecht, and Jaime Teevan. Learning causal effects on hypergraphs. In SIGKDD, page 1202-1212, 2022.
[49] Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pages 666-704, 1781.
[50] X Nie and S Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2):299-319 2020.
[51] Kenneth R Niswander and Myron Gordon. The women and their pregnancies: the Collaborative Perinatal Study of the National Institute of Neurological Diseases and Stroke, volume 73. National Institute of Hcalth, 1972.
[52] Elizabeth L Ogburn and Tyler VanderWeele, On the nondifferential misclassification of a binary confounder. Epidemiology, 23(3):433, 2012.
[53] Judea Pearl and Dana Mackenzie. The book of why: the new science of cause and effect. Basic books, 2018.
[54] Ofir Pele and Michael Werman. Fast and robust earth mover's distances. In ICCV, pages 460 460-467,,000 2009.
[55] Gabriel Peyré and Marco Cuturi. Computational optimal transport. Found, Trends Mach. Learn., 11(5-6):355-607 2019.
[56] Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In ICML, volume 119, pages 7673-7682, 2020.

---

[57] levgen Redko, Amaury Habrard, and Mare Sebban. Theoretical analysis of domain adaptation with optimal transport. In PKDD, volume 10535, pages 737-753, 2017.
[58] James M Robins, Andrea Romitzzky, and Lue Ping Zhao. Estimation of regression coefficients when some regressors are not always observed. J. Am. Stat. Assoc., 89(427):846 866, 1994.
[59] Paul Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-555 1983.
[60] Paul Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
[61] Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. J. Educ. Psychol., 66(5):688, 1974.
[62] Brian Sauer, M Alan Brookhart, Jason Roy, and Tyler VanderWeele. review of covariate selection for non-experimental comparative effectiveness research. Pharmaacoppidemiolggy and drug safery, 22(11) (1:1111111 139 1145, 2013.
[63] Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. In AAAI, pages 5612-5619, 2020.
[64] Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In ICIR, 2018.
[65] Thibault Séjouméé, Jean Feydy, François- Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport. CoRR, 129558 2019.
[66] Thibault Séjourné, Sranoois-aavie Vialard, and Gabriel Peyré. Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe. In AISTATS, volume 151, pages 4995-5021, 2022.
[67] Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In ICMI, pages 3076-3085, 2017.
[68] Claudia Shi, David M. Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. In NeurlPS, pages 2503-2513, 2019.
[69] Susan M Shortreed and Ashkan Ertefaie, Outcome-adaptive lasso: variable selection for causal inference. Biometrics, 73(4):1111-1122 2017.
[70] Tamar Sofer, David B Richardson, Flena Colicino, Joel Schwartz, and Eric J Tchetgen Tchetgen. On negative outcome control of unobserved confounding as a generalization of diffcrence-in- differences. Stat. Sci., 31(3):348, 2016.
[71] Ruibo Tu, Kun Zhang, Hedvig Kjellström, and Cheng Zhang. Optimal transport for causal discovery. In ICLR. OpenReview.net, 2022.
[72] Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
[73] Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. J. Am. Stat. Assoc., 1 122823):2228-224 2018.
[74] Hao Wang, Wei Chang, Tianqiao Liu, Jianmin Huang, Zhichao Chen, Chao Yu, Ruopeng Li, and Wei Chu. Esem2: Entire space counterfactual multi-task model for post-click conversion rate estimation. In SIGIR, 2022.
[75] Haotian Wang, Wenjing Yang, Longqi Yang, Anpeng Wu, Liyang Xu, Jing Ren, Fei Wu, and Kun Kuang. Estimating individualized causal effect with confounded instruments. In KDD, pages 1857-1867. ACM, 2022.
[76] Haotian Wang, Kun Kuang, Haoang Chi, Longqi Yang, Mingyang Geng, Wanrong Huang, and Wenjing Yang. Treatment effect estimation with adjustment feature selection. KDD, pages 2290-2301. ACM, 2023.

---

[77] Haotian Wang, Kun Kuang. Long Lan, Zige Wang, Wanrong Huang, Fei Wu, and Wenjing generalization with causal feature separation. IEEE Trans. Knowl.
Yang. Data Eng., 2023.
[78]
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processss for machine learning, volume 2. MIT press Cambridge, MA, 2006.
[79] Anpeng Wu, Kun Kuang, Ruoxuan Xiong, Bo Li, and Fei Wu. Stable estimation of heteroge- neous treatment effects. In ICML, volume 202 of Proceedings of Machine Learning Research, pages 37496-37510. PMLR, 2023.
[80] Anpeng Wu, Junkun Yuan, Kun Kuang, Bo Li, Runze Wu, Qiang Zhu, Yucting Zhuang, and Fei Wu. Learning decomposed representations for treatment effect estimation. IEEE Trans. Knowl. Data Eng., 35()):498-5) 5001, 2023.
[81] Peng Wu, Haoxuan Li, Yuhao Deng, Wenjie Hu, Quanyu Dai, Zhenhua Dong, Jie Sun, Rui Zhang, and Xiao-Hua Zhou. On the opportunity of causal learning in recommendation systems: Foundation, estimation, prediction and challenges. In IJCAI, pages 23-29, 2022.
[82] Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In ICLR, 2020.
[83] Renjun Xu, Pelen Liu, Yin Zhang, Fang Cai, Jindong Wang, Shuoying Liang, Heting Ying, and Jianwei Yin. Joint partial optimal transport for open set domain adaptation. In IJCAI, pages 2540-2546, 2020.
[84] Mengyue Yang. Xinyu Cai, Furui Liu, Weinan Zhang, and Jun Wang. Specify robust causal representation from mixed observations. In KDD, pages 1978-2987.. ACM, 2023.
[85] Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. In NeurlPS, pages 2638-2648, 2018.
[86] Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. ACE: adaptively similarity- representation learning for individual treatment effect estimation. In ICDM, pages 1437, 2019.
[87] Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: estimation of individualized treatment effects using generative adversarial nets. In ICLR, 2018.
[88] Wenhao Zhang, Wentian Bao, Xiao- Yang Liu, Keping Yang, Quan Lin, Hong Wen, and Ramin Ramezani. Large-scale causal approaches to debiasing post-click conversion rate estimation with multi-task learning. WWW, pages 2775-2783 2020.
[89] Yang Zhang, Dong Wang, Qiang Li, Yue Shen, Ziqi Liu, Xiaodong Zeng, Zhiqiang Zhang, Jinjie Gu, and Derek F Wong. User retention: A causal approach with triple task modeling, In IJCAI, 2021.
[90] Jiajing Zheng, Sensitivity Analysis for Causal Inference with Unobserved Confounding. Univer- sity of California, Santa Barbara, 2021.
[91] Jiajing Zheng, Alexander D' Amour, and Alexander Franks. Copula-based sensitivity analysis for multi-treetmmnt causal inference with unobserved confounding. CoRR, abs/2102.09412, 2021.
[92] Ziwei Zhu, Yun He, Yin Zhang, and James Caverlee. Unbiased implicit recommendation and propensity estimation via combinational joint learning. In RecSys, pages 551-556, 2020.


Reviews:
[{"review_id": "MZzpZoaQFM", "reviewer": "Reviewer_VwsR", "summary": "The paper introduces a new approach called Entire Space CounterFactual Regression (ESCFR) to estimate individual treatment effects from observational data, addressing two critical problems that existing methods fail to solve. The first problem is mini-batch sampling effects (MSE), which lead to misalignment in non-ideal mini-batches with outcome imbalance and outliers. The second problem is unobserved confounder effects (UCE), which result in inaccurate discrepancy calculation due to the neglect of unobserved confounders. ESCFR combines the principles of optimal transport and causality to overcome these issues.", "questions": "See in Weaknesses", "limitations": "See in Weaknesses", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The theoretical evidence is solid and convincing.\n- The topic, mentioned in the paper as the mini-batch misalignment and UCE error, are interesting and practical issues of ITE estimation.\n- Experimental results are solid.", "weaknesses": "Is there some theoretical guarantee for the UCE regularizer proposed in this paper? If so, I think the quality can be further improved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe extend our sincere gratitude for your insightful comments and constructive critique. We acknowledge your point concerning the theoretical support of the UCE regularizer. We are pleased to elaborate our thoughts and methodologies to clarify the theoretical foundation.\n\n**Is there some theoretical guarantee for the UCE regularizer proposed in this paper? If so, I think the quality can be further improved.**\n\n- **Theoretical Guarantees for the UCE Regularizer.** The UCE regularizer's main theoretical backing derives from the **flexibility of optimal transport in defining distribution geometry** within latent space.\n\n  - **Problem Definition:** The core issue revolves around measuring distribution discrepancies between treated and controlled groups in the presence of unobserved confounder effects. By estimating this discrepancy and minimizing it as a loss function, we encourage the overlapping of units from different groups in latent space, thereby controlling the generalization risks of outcome mappings $\\phi$ within both treatment groups.\n\n  - **Role of OT:** OT enables us to **calculate discrepancy with unit-wise distance.** Not only does this facilitate computation, but more importantly, it enables the **editing of geometry according to task characteristics** in the latent space. Specifically, according to the Kantorovich problem, we have $\\mathbb{W}(\\alpha, \\beta):=\\min_{\\boldsymbol{\\pi} \\in \\Pi(\\alpha, \\beta)}\\langle\\mathbf{D}, \\boldsymbol{\\pi}\\rangle$. Then, we edit the geometry of latent space via unit-wise distance $\\mathbf{D}$. In general, we can use the Euclidean geometry, where the distance between the i-th treated unit and the j-th untreated unit is $\\mathbf{D}_{ij}=||x_i-x_j||_2^2$ , with $x$ denoting observed confounders. To consider unobserved confounders $x^\\prime$ , we can define the unit-wise distance as $||x_i-x_j||_2^2+||x_i^{\\prime}-x_j^{\\prime}||_2^2$ , which is so far grounded in rigorous theoretical support.\n   - **Outcome Calibration:** Given the monotonic assumption, we calibrate the distance with the outcome difference rather than unobserved covariates. Since $x^\\prime$ is not observable, a natural idea is that given similar observed covariates, i.e., $x_i=x_j$ and treatment indicator, the effect of unobserved confounder should be reflected on the outcomes given the same treatment. As such, it is reasonable to use the outcome difference to calibrate $\\mathbf{D}_{ij}$, given the assumption that there exists a positive relationship between unobserved confounders and potential outcomes. We discussed the assumption as a limitation of this work in Appendix D.3.\n\n- **Inspiration from Partial Identification.** \n  - **Partial identification:** An estimand $\\theta$ is partially identified if the observed data distribuion is compatible with multiple values of $\\theta$.  In causal inference, challenges like unobserved confoundings might prevent precise causal effect pinpointing. A weaker alternative is to obtain a range of possible causal effects, known as a \"identified set\", and reduce the size of the set using proper assumptions.  For example, given assumptions, e.g., monotone treatment selection, we can narrow the bound of treatment effect. More recently, according to Lu and Ding's formulation [1], assuming that we know two sensitivity parameters $\\frac{\\mathbb{E}(Y(1) \\mid Z=1, X)}{\\mathbb{E}(Y(1) \\mid Z=0, X)}=\\varepsilon_1(X)$, $\\frac{\\mathbb{E}(Y(0) \\mid Z=1, X)}{\\mathbb{E}(Y(0) \\mid Z=0, X)}=\\varepsilon_0(X)$, we can identify the treatment effect with unmeasured confounding.\n \n  - **Connection with partial identification:** There are significant similarities between the proposed PFOR regularizer and partial identification principles. One significant similarity is that the transport strategy, i.e., $\\pi$ in Eq.(2), is partially identified given the existence of unobserved confounders. More specifically, assuming that the unobserved confounder $x^\\prime$ has multiple candidate values, there should be multiple corresponding transport strategies, so we can only identify a stratege set. Then, making the monotonic assumption between $x^\\prime$ and outcomes, we calibrate the unit-wise distance with the outcome differences, to reduce the size of strategy set and achieve more accurate estimation among possible transport strategies.\n  - **Assumption Criticality:** In both works, whether the assumptions made fit the application is a crucial aspect. In our work, we assume that the unobserved confounders are monotone with the outcome given $x$ and $t$, as is formulated in Assumption D1. We note that this assumption is also employed in emerging partial identification work [2], and its plausibility is analyzed in Appendix D.\n\n- **Conclusion and Future Directions.**\nWe recognize that our method only mitigates unobserved confounding effects in specific cases where the assumptions are met. Nevertheless, we maintain that our approach is conceptually sound and may inspire further research into the flexibility of optimal transport in causal inference. We have detailed the assumptions, limitations, and our scope in the appendix and will emphasize them in the main paper's revision.\n\n\n[1] Lu, Sizhu and Peng Ding. \"Flexible sensitivity analysis for causal inference in observational studies subject to unmeasured confounding.\" (2023).\n\n[2] Zheng, Jia-wen et al. \"Copula-based Sensitivity Analysis for Multi-Treatment Causal Inference with Unobserved Confounding.\" (2021).\n\n\n---\n\nWe hope this response addresses your concerns and elucidates the theoretical foundation of our work. Your feedback is instrumental in refining our paper, and we are open to further dialogue to enhance our contribution."}, {"review_id": "pr0cb1Ek3Z", "reviewer": "Reviewer_hjjj", "summary": "The objective of this research paper is to tackle two primary concerns: (1) the effects of mini-batch sampling (MSE), which lead to misalignment in non-ideal mini-batches characterized by outcome imbalance and outliers, and (2) the impacts of unobserved confounders (UCE), resulting in inaccurate discrepancy calculation due to the neglect of these unobserved factors. To address these challenges, this paper presents a novel approach based on optimal transport in the context of causality.\n\nMore specifically, the proposed approach builds upon the stochastic optimal transport framework. It introduces a relaxed mass preserving regularizer to mitigate the MSE issue and devises a proximal factual outcome regularizer to address the UCE problem. By incorporating these techniques, the proposed method achieves a principled solution for treatment selection bias.\n\nThe effectiveness of the proposed method is extensively evaluated through a series of experiments. The results demonstrate that the approach successfully overcomes treatment selection bias and outperforms existing state-of-the-art methods by a significant margin.", "questions": "See above", "limitations": "See above", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and well-presented. It delves into the important question of improving the estimation of individualized treatment effects, a crucial topic within the field of causal inference. Optimal transport, an emerging field in machine learning, is effectively employed to measure treatment selection bias. Furthermore, the two proposed regularizers are technically intriguing and efficient, as supported by both theoretical and empirical evidence.", "weaknesses": "I have two main concerns: \n\n1. The potential applicability of the proposed relaxed mass-preserving regularizer extends beyond causal effect estimation and encompasses a wide range of application settings. It appears to be independent of the estimation of causal effects. The authors acknowledge that the OT discrepancy can be easily affected by various sampling cases, highlighting that it is not solely limited to the distribution discrepancy across treatment groups. Consequently, the main contribution of this regularizer may be seen as an incremental addition to the field of causal inference, raising concerns about the originality and novelty of the paper.\n\n2. The existing literature on the subject appears to be considerably understudied. For example, important references like \"Optimal Transport for Counterfactual Estimation: A Method for Causal Inference\" by Arthur Charpentier, Emmanuel Flachaire, and Ewen Gallic are missing and remain unaddressed in the related work section. It would be beneficial for the authors to conduct a more comprehensive investigation of the related literature and consider incorporating these references into their study.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1] The main contribution of the proposed relaxed mass-preserving regularizer may be seen as an incremental addition to the field of causal inference, raising concerns about the originality and novelty of the paper.**\n\n- Thank you for your sincere comments. We acknowledge your concern regarding the contribution of the proposed relaxed mass-preserving regularizer to the field of treatment effect estimation (TEE). However, it is imperative to underline the context and specific challenges our work seeks to address.\nAlthough the MSE issue is not unique to TEE, this issue takes on particular significance in this field.\nIt is specified as two issues, i.e., outcome imbalance and outliers, which directly impacts the reliability of the alignment strategy and thus the causal estimators. \n\n- Notably, this issue has been somewhat overlooked in current progress, while it is present across numerous representation-based methods.\nOur regularizer's value lies in addressing this vital yet neglected issue, specifically within the realm of current TEE strategies based on representation alignment. By relaxing the mass-preserving constraint of the vanilla OT, this challenging issue is mitigated, which improves the TEE performance.\n\n- We also wish to slightly stress that the application of optimal transport within causal inference remains a nascent area of study. Our proposed method represents an interesting attempt and trial in this direction, reflecting a thoughtful response to existing problems and a promising pathway for future exploration.\n\n**[W2] The existing literature on the subject appears to be considerably understudied. For example, important references like [1] are missing and remain unaddressed in the related work section. It would be beneficial to conduct a more comprehensive literature review and consider incorporating these references into their study.**\n\n- Thank you very much for your kind suggestion. The related work section in the main paper seems not comprehensive enough due to the limited page quota. We have included a more comprehensive literature review in appendix. According to the methods used for balancing, current solutions can be divided into matching based methods, re-weighting based methods and distribution alignment methods. We focus on the re-weighting based methods and distribution alignment methods in the related works.\n- We agree that it is essential to comprehensively review literature in the intersection of OT and causal inference, so we have opened a new thread in the related work to give a brief introduciton, and explained our difference with them.\n- Thanks for recommending an important reference. \n    - summary of this paper. Overall, this work first match each unit in the control group with another one in the treat group, and then constructs CATE estimators based on it. The core of this work, in our view, is a matching-based CATE estimator, i.e., the mutatis mutandis uantile-based CATE in Definition 3.1-3.2. The authors further specify the optimal matching strategy as an OT problem in Section 4. \n    - Similarity and dissimilarity. Both works aim to achieve CATE estimation through optimal transport. They both acknowledge the selection bias issue and use OT to construct (locally) balanced representations. However, [1] is a matching-based method, which aims to match similar units in the data space, while this work is a representation-based method, which aims to minimize the distribution discrepancy in the latent space. On top of that, this work mainly cares about the flexibility of OT variants to tackle issues, i.e., MSE and UCE, which are critical in the causal inference field. These regularizations, in our view, could inspire further research based on [1], to refine matching strategies to make it more adaptable to causal inference.\n\n- After careful reading and discussion, we appreciated the insights and contributions of this work, and decided to include it in our related work section.\n\n[1] Charpentier, Arthur et al. \"Optimal Transport for Counterfactual Estimation: A Method for Causal Inference.\" (2023)."}, {"review_id": "egwo0wM9SK", "reviewer": "Reviewer_MG3J", "summary": "The authors propose to minimize a (relaxed) Sinkhorn distance between representations under the two treatment groups in order to balance the latent space and better predict the treatment effects. They highlight two problems with existing TARNet-based approaches that do not regularize with optimal-transport cost: mini-batch sampling effects (MSE) and unobserved confounder effects (UCE). Theorem 3.1 & Corollary 3.1 formally characterize how the proposed method suffers less from MSE. The authors also introduce additional regularization terms that are meant to deal with UCE. Finally, they show improved performance with two classical semi-synthetic benchmarks.", "questions": "How do you know that you address UCE? As the paper stands, I recommend that you refrain from discussing hidden confounders in your motivation for the method and stick with what really *can* be shown for your proposed regularization. I am willing to raise my score once this is addressed.\n\nIs MSE really an issue with previous approaches? Can you show that clearly?", "limitations": "Limitations or potential societal impacts are not discussed.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": " * The problem of estimating heterogeneous treatment effects is important, the paper is written well, and the solution is creative.\n * The motivation for the relaxed Sinkhorn distance (Sec. 3.2) is great. \n * The benchmark results are solid.\n * The additional empirical explorations (Sec. 4.3 - 4.5) are interesting as well.", "weaknesses": "My biggest problem is with the UCE part of the paper.\n * Sec. 3.3: the explanation for the hidden-confounding regularizer makes intuitive sense. However, I do not see a formal connection. Once the authors substitute the actual potential outcomes with the model's predictions (Eq. 11 -> Eq. 12), one begins to wonder if this can actually help with hidden confounding.\n * A major confounder with the issue above is that the authors do not appear to include hidden confounders in their semi-synthetic benchmarks. Therefore, it is impossible to tell if this regularization approach really helps with UCE.\n * If the authors still wish to make the case for UCE mitigation, they need to reference the growing literature on partial identification / sensitivity analysis for hidden confounders and how their regularization relates to sensitivity models.\n\nBesides that, I think the authors should spend more time addressing what the theoretical results mean for the MSE issue that they highlighted in the introduction.\n\nSmall nitpicks:\n* Figure 1a, the orange annotation looks like it is referring to the X_1 part of the figure.\n * line 107-108: \"Optimal transport is a preferred method due to its advantages over competitors.\" is overly vague.\n * line 108-109, \"It accounts for the distribution’s geometry, making it effective [where the KL-divergence] fails.\" is a strange statement. Why does the geometry help here? No explanation is provided. \n * The beginning of the paper is a bit verbose.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1] About the hidden-confounding regularizer**\n- **Formal connection.** Please refer to our response to the Reviewer VwsR, section \"Theoretical Guarantees for the UCE Regularizer\".\n- **Substituting Potential Outcomes with Model Predictions.** Pseudo label (using model predictions to substitue ground-truth labels) has been shown effective in both causal inference and machine learning areas.\n  - In causal inference, doubly robust learning uses an imputation model to impute labels or errors for unobserved data (without ground-truth labels), and use the imputed values as label for model training to achieve double robustness. It is wodely used in treatment effect estimation [1] and debiasing recommendation [2].\n  - In machine learning, for example, in weakly supervised learning, wrapper methods (e.g., self-training), as one of the oldest and prevalent methods, augment the labeled training data with unlabeled data whose labels are obtained from model predictions [6]. In robust machine learning, it is shown that pseudo label magnify useful underlying information in data, and thus can be used to correct noisy labels [3-5]. \n  - Thus, using model prediction is a reasonable practice, which has also been shown effective empirically.\n\n[1] Estimation of regression coefficients when some regressors are not always observed. Journal of the American statistical Association. 1994. \\\n[2] Doubly Robust Joint Learning for Recommendation on Data Missing Not at Random. ICML. 2019. \\\n[3] Who said what: Modeling individual labelers improves classification. AAAI. 2018. \\\n[4] Noise against noise: stochastic label noise helps combat inherent label noise. ICLR. 2020. \\\n[5] A closer look at memorization in deep networks. ICML. 2017. \\\n[6] A survey on semi-supervised learning. Machine learning. 2020.\n\n**[W2] About experiments on semi-synthetic benchmarks**\n\n- We do realize this problem. Unobserved confounders are generally only present in real-world observational data. But in observational data, the performance of ITE estimators cannot be evaluated due to missing counterfactuals. As such, we followed the general settings and conduct experiments on semi-synthetic data. **For semi-synthetic data, one source of hidden confounders is the information loss from the raw data space to the representation space, where not all valuable information (e.g., confounders) is extracted and preserved in the representation space.** In this context, the validity of PFOR was verified by ablation and sensitivity studies in Section 4.3 and 4.5.\n- **We note PFOR also works without unobserved confounders.** One source of performance gain would be the statistical regularization, encouraging samples with similar outcomes to share similar representations, which is an effective prior according to nearest-neighbor approaches. Another benefit comes from the prompts that the outcomes provide to the alignment module. Specifically, vanilla Sinkhorn considers each variable equally during the matching phase; in contrast, PFOR highlights variables with large correlation with Y (i.e., more significant confounding effects), which is an effective prompt.\n\n**[W3] About connections with partial identification**\n\nIt is indeed an invaluable comment. Please refer to our response to the Reviewer VwsR, which is highlighted as **\"Inspiration from Partial Identification\"**.\n\n**[Q1] About the regularizer for UCE & the motivations**\n\nThere are other alternative explanations for the effectiveness of the proposed PFOR regularizer. To list a few:\n- The performance gain could come from the prompts given to the alignment module. Vanilla Sinkhorn considers each variable equally; in contrast, PFOR focuses on variables with great correlation with Y (i.e., significant confounding effects), providing prompts to the alignment algorithm. \n- The gain could come from the statistical regularization: samples with similar outcomes to share similar representations, which is an effective prior according to nearest-neighbor approaches.\n\nNevertheless, in this work, we suggest that PFOR help alleviate the UCE issue in very limited cases given the Monotonicity Assumption in Appendix D, instead of functioning as a general solution for addressing it. Following your kind advice, we will refine our narrative in this regard and discuss the possible alternative explanations above in the main text.\n\n**[Q2] About MSE issue with previous approaches** \n- Yes, MSE is a concern in previous representation-based approaches. \n  - **Theoretically**, the complexity term $\\mathcal{O}(\\cdot)$ in Theorem 3.1 unveils the risks induced by stochastic sampling, which exists in other representation-based methods as well. Precisely, $\\mathcal{O}(\\cdot)$ emerges from the discrepancy between the entire population and the sampled mini-batch units (refer to Eq.(30)), and is highly susceptible to the uncontrollable quality of sampling. Therefore, **the discrepancy should be robust to bad sampling cases; otherwise, the variance can hinder the computation severely.** This vulnerability is ignored by existing representation-based solutions. \n  - **Empirically**, Figure 2 illustrates that Sinkhorn discrepancy **fails to quantify the group discrepancy for producing erroneous transport strategies** in non-ideal mini-batches, and thus **misguiding the update of the representation mapping $\\psi$**. Established methods are notably affected by this issue. We have opened a new thread to discuss it with additional experiments in Appendix D2.\n- **Beyond OT**, most alignment methods, e.g., MMD and adversarial training, grapple with the MSE issue. When faced with a mini-batch with imbalanced outcomes, for example, they naively minimize the discrepancy, erasing discriminative characteristics essential for outcome estimation.\n- **Advantage of OT**, OT allows to **formalize the MSE issue** as the mass-preserving constraint, thereby giving a natural solution: relaxing the marginal constraint, which is the primary intuition behind RMPR. "}, {"review_id": "1AsjJyusNd", "reviewer": "Reviewer_jCP6", "summary": "In this paper, authors have applied optimal transport to address selection-bias issue in binary treatment setting for individualized treatment effect estimation. They have proposed a relaxed mass-preserving regularizer to address the mis-alignment generated from outcome imbalance and outliers in non-ideal mini-batches. They also proposed a proximal factual outcome regularizer to relax the unconfoundedness assumption. They have done extensive evaluation of the proposed approach on IHDP and ACIC, two semi-synthetic datasets.", "questions": "- Training all networks for 400 epochs do not look fair to me. Are all the methods getting trained before 400 epochs or not? In my experience this affects the performance in causal inference. All the networks should be trained using early stopping and by setting large number of max-epochs so that every learner gets chance to train properly.\n- Have you used two layers each in representation learning and outcome mapping or total two layers for both?\n- First line of Introduction section does not look right to me. Doesn't RCT mostly answer population-level or heterogeneous-level questions but not the individual-level treatment effects?\n- In figures 1(a) and 2, please clearly state what are the different axis, e.g., for Figure 2, it is not clear what is x-axis or what is y-axis?\n- Why does title says 'a new perspective' when the idea is very widely used, including in causal inference, as authors have themselves discussed that?", "limitations": "Authors did not discuss the assumptions and limitations of thier work.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- A novel method is proposed for treatment effect estimation in binary setting based on optimal transport (OT) that uses a relaxed mass-preserving regularizer to address the mis-alignment generated from outcome imbalance and outliers in non-ideal mini-batches, as well as a proximal factual outcome regularizer to relax the unconfoundedness assumption.\n- Extensive evaluations is performed to study the proposed model that provides necessary analysis and ablation studies.\n- Theoretical results are provided to bound the error in terms of the proposed components to show the convergence (I have not verified the proof).\n- Idea was clearly stated with sufficient details.\n", "weaknesses": "- Baselines used in the paper seem old. There are very large number of methods proposed for treatment effect (TE) estimation. It would have been nice to see comparison with recent techniques. Hoever, even if authors don't add new baselines, I think paper has sufficient novelty to be above the acceptance threshold.\n- The paper solves (equation (1)) Conditional average treatment effect estimation also called as heterogeous treatment effects estimation, as called in machine learning literature, but paper did not even mention these names. Please update the paper so that readers don't have any confusion. There is already a lot of confusion in ML regarding the terminology.\n- Causal inference methods are typically based on some assumptions -- but I did not see any. Please clearly state the assumptions required to apply these methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1] Comparison with recent techniques**\n- Thanks for your kind recognition.\nWe agree that comparison with recent techniques could improve this work. Thus, we incorporated two more recently baselines founded on representation learning, DRCFR [1] and DRCFR-MIM [2], in appendix, with the aim of modernizing our scope of baselines.\nThese two methods introduce variable decomposition and employ the distribution alignment technology (CFR) as a component.\nTherefore, *it would be somewhat unfair for ESCFR to be directly compared to them, given that ESCFR is not equipped with variable decomposition technology*. As a result, we elected to place these results in appendix rather than the main content of the paper to ensure a fair presentation while preserving completeness.\n\n- Here, we provide the comparisons (Table 6 in the appendix). DRCFR and DRCFR-MIM exhibited more stable performance (with less variance) and did surpass CFR-WASS in terms of out-of-sample PEHE. Nonetheless, ESCFR significantly outperforms them across most metrics.\n\n  | Dataset   | Out-sample    | Out-sample    | In-sample     | In-sample     |\n  |-----------|---------------|---------------|---------------|---------------|\n  | Model     | PEHE          | AUUC          | PEHE          | AUUC          |\n  | TARNet    | 3.2054±0.2676 | 0.6683±0.0139 | 3.2933±0.4076 | 0.8597±0.0350 |\n  | CFR-WASS  | 3.0819±0.3407 | 0.6693±0.0250 | 3.4280±0.2857 | 0.8439±0.0349 |\n  | DRCFR     | 3.0199±0.0658 | 0.6360±0.0080 | 3.4957±0.0895 | 0.8456±0.0073 |\n  | DRCFR-MIM | 2.9593±0.1532 | 0.6456±0.0189 | 3.3875±0.2208 | 0.8434±0.0093 |\n  | ESCFR     | 2.3082±0.4334 | 0.7311±0.0123 | 2.5385±0.5020 | 0.8416±0.0367 |\n\n[1] Negar Hassanpour, et al. Learning disentangled representations for counterfactual regression. In ICLR. 2020. \\\n[2] Mingyuan Cheng, et al. Learning disentangled representations for counterfactual regression via mutual information minimization. In SIGIR. 2022.\n\n**[W2] Terminologies: ITE, CATE and HTE**\n\n- We sincerely apologize for any confusion caused. Your observation regarding the naming convention is entirely valid, and we appreciate your suggestion. In this paper, **we delve into the individual-level causal effect estimation, where each individual is identified by its covariates $x_i$**, in line with existing works [1,2,3]. Under this circumstance, individual treatment effect estimation (concerning the effect of an individual unit with covariates $x_i$) is interchangeably referred to as CATE or HTE, as you correctly pointed out, in the machine learning literature [1,2,3]. \n- We acknowledge the importance of maintaining clarity and consistency in terminology, especially given the existing confusion in the field. Consequently, we commit to updating the paper to employ the term \"CATE\" and will elucidate the connections between these terminologies in Appendix A.\n\n[3]  Uri Shalit, et al. Estimating individual treatment effect: generalization bounds and algorithms. In ICML. 2017.\n\n**[W3 & Limiations] Assumptions & limitations**\n- Thanks for your suggestion. To maintain the concise presentation of the main paper, we have indeed delineated the necessary assumptions in **Appendix A.1**.\nAdditionally, the assumptions and limitations specific to our work are thoroughly discussed in **Appendix D.3**.\n\n**[Q1] Training epochs & early stopping**\n-  We sincerely apologize for an oversight during the proofreading stage, where we inadvertently deleted a critical technique regarding early stopping. Allow us to clarify the following points:\n- The maximum number of training epochs is 400, with early stopping implemented (tolerance=30). \n- We utilized AUUC for model selection and early stopping, as it better aligns with our objectives.\n\n**[Q2] About network structures**\n- We employed a structure consisting of 3 layers for both representation learning and outcome mapping, inspired by CFR. As articulated in CFR: `Both versions are implemented as feed-forward neural networks with 3 fully-connected exponential-linear layers for the representation and 3 for the hypothesis.`\n- In terms of key parameters, such as the number of hidden nodes and layers, we fine-tuned them based on TARNet since most representation-based methods can be viewed as a regularization on TARNet. Moreover, we ensured that all deep models possessed either the same (at least similar) number of parameters to facilitate a fair comparison.\n\n**[Q3] About the first line of Introduction**\n- We sincerely apologize for the confusion caused. In this paper, we study the individual-level causal effect estimation, **where each individual is identified by its covariates $x_i$** as [1,2,3], which is also known as CATE and HTE (see our response to W1). RCT can successfully avoid the treatment selection bias issue and the unobserved confounding issue since it cuts the edge between the covariates and treatment in the causal graph. Since in RCT, there is not selection bias, vast machine learning models (e.g., XGBoost, T-learner) can be used without adjustment, which therefore helps ITE estimation.\n\n**[Q4] Fig. 1(a) and 2: meanings of x- and y-axis**\n- Sorry for the confusion caused. In both figure 1(a) and 2, the x-axis represents the covariates, and the y-axis represents the outcomes.\n\n**[Q5] 'a new perspective' in title**\n- We note that although there are some attempts that use OT to empower causal inference, these attempts are fairly few and remain in an initial exploratory stage. More importantly, there has not been attempt to improve current **representation-based estimators** by devising OT algorithm, with the aim at mitigating the specific issues in causal inference. We hope to regard this paper as a exemplar to show the potential and flexibility of OT in this field, which we believe offers a fresh perspective for representation-based solutions. We are also pleased to receive any further comments or suggestions you may have in this regard. We look forward to hearing from you."}]

Author rebuttal:

Meta-review: Reviewers agreed the paper presents insightful ideas, giving new and deep theoretical results at the intersection of two disparate fields: optimal transport and counterfactual inference, including expanding the use of optimal transport to recently proposed proximal inference approaches. The paper is well-written and the experiments are sound, and contributes to a better understanding of this sub-field of causal effect inference.
Decision: Accept (poster)