Sinkhorn Discrepancy for Counterfactual Generalization
Authors: Hao Wang
Quanyu Dai
Jiajun Fan
Weiming Liu
Zhichao Chen
Tianqiao Liu
Yichao Wang
Zhenhua Dong
Ruiming Tang
Venue: Submitted to ICLR 2023 (2023)
Keywords: causal inference
treatment selection bias

Abstract: Estimating individual treatment effects from observational data is very challenging due to the existence of treatment selection bias.
Most existing representation-based methods mitigate this issue by aligning distributions of different treatment groups in the representation space. However, they still suffer from two critical problems: (1) Mini-batch Sampling Effects (MSE), where the alignment easily fails due to the outcome imbalance or outliers in the batch; (2) Unobserved Confounder Effects (UCE), where the unobserved confounders damage the correct alignment. To tackle these problems, we propose a principled approach named Entire Space CounterFactual Regression (ESCFR) based on a generalized sinkhorn discrepancy for distribution alignment within the stochastic optimal transport framework. Based on the 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 our proposed ESCFR can successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.

SINKHORN DISCREPANCY FOR COUNTERFACTUAL GENERALIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Estimating individual treatment effects from observational data is highly challeng- ing due to the existence of treatment selection bias. Most prevalent approaches mitigate this issue by aligning distributions of different treatment groups in the representation space. However, there are two critical problems circumvented: (I) mini-batch sampling effects (MSE), where the alignment easily fails due to the outcome imbalance or outliers at mini-bbtch level; (2) unobserved confounder effects (UCE), where the unobserved confounders damage the correct alignment. To tackle these problems, we propose principled approach named Entire Space Counter Factual Regression (ESCFR) based on a generalized sinkhorn discrepancy for distribution alignment within the stochastic optimal transport framework. Based on the framework, we propose a relaxed mass pass-preserving regularizer to address the MSE issue and design a proximal factual outcome regularizer to handle the UCE issue. Extensive experiments demonstrate that our proposed ESCFR successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.
INTRODUCTION
Estimating individual treatment effect (ITE) with randomized controlled trials is a common practice in causal inference, which has been widely used in e-commerce (Betlei al.) (021), education (Cordero et al. 2018), and health care (Schwab et al. 2020). For example, drug developers would conduct clinical A/B tests to evaluate the drug effects. Although randomized controlled trials are the gold standard (Pearl & Mackenzie, 2018) for causal inference, it is often prohibitively expensive to conduct such experiments. Hence, observational data that can be acquired without intervention has been a tempting shortcut. For example, drug developers tend to assess drug effects with post-marketing monitoring reports instead of clinical A/B trials. With the growing access to observational data, estimating ITE from observational data has attracted intense research interest.
Estimating ITE with observational data has two main challenges: (1) missing counterfactuals, i.e., only one factual outcome of all potential outcomes be observed; (2) treatment selection bias, i.e. individuals have their preferences for treatment selection, making units in different treatment groups heterogeneous. To handle missing counterfactuals, meta-learners (Künzel et al. 2019) decompose the ITE estimation task into solvable factual outcome estimation subproblems. However, the treatment selection bias makes it difficult to generalize the factual outcome estimators trained within respective treaiment groups to the entire population; consequently, the derived ITE estimator is biased.
Beginning with counterfactual regression (Shalit al., 2017) and its revolutionary performance, most prevalent methods handle the selection bias by minimizing the distribution discrepancy between groups in the representation space (see Liuyi et al., 2018 Hassanpour & Greiner 2020 Cheng et al. (022). However, two critical issues with these methods have long been neglected, which significantly impedes them from handling the treatment selection bias. The first problem is the mini-batch sampling effects (MSE). Specifically, current representation-basee methods (Shalit et al.) 2017 Liuyi et al. 2018) compute distribution discrepancy within
instead of the entire data space, making vulnerable bad sampling cases. For example, given two aligned distributions, a mini- batch outlier exists in the sampled distribution, the mini-batch discrepancy will be significant, making the training process noise- filled. The second problem is the unobserved confounder effects

---

Factual outcome
Factual outcome
W
X
Treatment indicatoo
X2
X2
R
Covariate
Representation
Outcome estimate
(a) Handling treatment selection bias with <)().
(b) Architecture of ESCFR.
Figure 1: Overview of handling treatment selection bias with ESCFR. Red (blue) indicates treated (untreated) group. (a) Treatment selection bias causes the shift between X1 and Xg. impeding Φ1 and Φn to generalize beyond the respective group's properties, Scatters and curves indicate the units and fitted outcome mappings, respectively. (b) ESCFR handles this issue by mapping covariates to an overlapped representation space with R (X) where Φi and Φo are mutually generalizable.
(UCE). Specifically, current approaches directly assume unconfoundedness Ma et al. (2022), while the unobserved confounders widely exist in real scenarios and make the resulting estimators biased.
Contributions and outline. this paper, we propose an effective ITE estimator based on optimal transport, Entire Space CounterFactual Regression (ESCFR), which tackles both the MSE and UCE issues with a generalized sinkhorn discrepancy. Specifically, after preliminaries in Section 2 we first reformulate the ITE estimation problem as a stochastic optimal transport problem in Section 3.1
We next showcase the MSE issue faced by existing approaches in Section 3.2 and propose relaxed mass passerrreer regularizer mitigate this issue. We further investigate the UCE issue in Section 3.3 and propose proximal factual outcome regularizer to solve it. We finally formulate the architecture and learning objectives ESCER Section 3.4 and report the experimental results in Section
PRELIMINARIES
2.1 CAUSAL INFERENCE FROM OBSERVATIONAL. DATA
This section formulates basic definitions and models in observational causal inference. We first formalize the fundamental elements in Definition 2.1 following the general notation convention
Definition 2.1. Let X the random variable of covariates, with support and distribution P ( Let R be the random variable of induced representations, with support R and distribution P(r);; Let be the random variable of outcomes, with support y and distribution Pyy) Let T be the random variable of treatment indicator, with support T {00,1 1) and distribution IPT).
Following the potential outcome framework (Rubin 1974). an individual with covariates x has two po- tential outcomes, namely Y1 (x) given is treated and Yu(x)) otherwise. The ground-truth individual treatment effect (ITE) is always formulated as the expected difference of potential outcomes:
T (z) == E[Y Yo x].
where one of these two outcomes always unobserved. To address such missing counterfactuals, the ITE estimation task is commonly decomposed into potential outcome estimation subproblems that are solvable with any supervised learning method (Künzel et al. 2019). For example, T-learner models the factual outcomes Y for units in the treated and untreated groups separately: S-learner regards the treatment indicator T as one the covariates X, and models for all units simultaneously. The ITE estimate is then the difference of the estimated outcomes when T is set to treated and untreated.
Definition 2.2. Let & X R be mapping from support X to R, i.e., Yx é X, 3r (xx) e R. Let R T + y be mapping from support R x T Y, i.e., maps the representations and treatment indicator to the corresponding factual outcome. For example, Yi (R),,, Yo Φ (R), where we abbreviate oi R, T 1) and (R, I 0) to Φ (R) and do (R), respectively, for brevity.
'We use uppercase letters, e.g., X to denote random variable, and lowercase letters, e.g., to denote an associated specific value. Letters in calligraphic font, egg., A' represent the support of the corresponding random variable, and P() represents the probability distribution of the random variable, egg., (PX).

---

TARNet (Shalit et al., 2017) and S- S-leamee, consisting of a representation mapping e and an outcome mapping & as defined in Definition 2.2 For an individual with covariates X, TARNet estimates ITE as:
obtains better performance by absorbing the advantages of both T-learner
(X) Y2 YY, where Y = ===1(((())
Yo= = ()=)((((X)),
(2)
where v is trained over all individuals, Φ1 and on are trained over the treated and untreated group, respectively, to minimize the factual error. Finally, the performance of ITE estimators is mainly evaluated with the precision in estimation of heterogeneous effect (PEHE);
(PE fxx -
(())) dr.
However, according to Figure [(a) the treatment selection bias causes distribution shift of covariates across groups, which misleads @1 and do to overfit their respective group's properties and generalize poorly to the entire population. Therefore, the ITE estimate by these methods would be biased.
2.2 DISCRETE OPTIMAL TRANSPORT AND SINKHORN DIVERGENCE
Optimal transport (OT) instantiates distribution discrepancy as the minimum transport cost, which provides grip for quantifying the treatment selection bias in Figure [(a] Monge (1781) first
formulated OT as finding an optimal mapping between two distributions. However, this formulation cannot guarantee the existence and uniqueness of solutions. Kantorovich (2006) proposed a more applicable formulation in Definition 2.3 which can be seen as a generalization of 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 e Bxxm which transports to 3 at minimum cost:
W(a
(D.T), II(a, = { paxm xcI((a.))
min
'Imm = a, "1, =b.
(4)
where W(a, 3) R is the Wasserstein discrepancy between a and 3; D E RR is the unit-wise distance between a and 3; a and b indicate the mass of units in a and 3, and II is the feasible transportation plan set which ensures the mass-preserving constraint holds.
However, exact solutions (Bonneel et al. 2011) to always come with high computational costs. As such, researchers would always add an entropic regularization to the Kantorovich problem:
(a, (D, ne aig (D, π) cH(r), H(T) Σ n. (log(x..s) 1), (5)
xell(a,3)
making the problem e-convex and solvable with the Sinkhorn algorithm (Cuurii 2013 The Sinkhorn
algorithm only consists of matrix-vector products, making it suited to be accelerated with GPUs.
3 PROPOSED METHOD
In this section, we present the proposed Entire Space CounterFactual Regression (ESCFR) approach based on optimal transport to tackle the treatment selection bias. We first illustrate the stochastic optimal transport framework for distribution discrepancy minimization across treatment groups. Based on the framework, we then propose a relaxed mass- paass-preeervng regularizer to address the sampling effect, and proximal factual outcome regularizer to handle the unobserved confounders. We finally summarize the model architecture, learning objectives, and optimization algorithm.
3.1 STOCHASTIC OPTIMAL TRANSPORT FOR COUNTERFACTUAL REGRESSION
Representation based approaches mitigate the treatment selection bias by calculating the distribution discrepancy in the representation space and then minimizing it. We select optimal transport to compute the discrepancy since it has shown compelling advantages over its competitors. Specifically, it accounts for the distribution's geometry and thus works in the cases where o-divergence (e-g.. Kullback-Leibler divergence) fails (Seguy et al., 2018). addition, the calculated discrepancy can
We calculate the unit-wise distance with the squared Fuclidean metrie following Courty al. (011b))

---

be optimized with the traditional supervised learning framework instead of the adversarial learning framework, and is therefore easier to optimize than adversarial- based methods (Kallus 2020
Optimal transport calculates the group discrepancy as W (T) where PTT (r) and (r) are the distributions of representations in treated and untreated groups, respectively, induced by the mapping r + (z). The discrepancy is differentiable with respect to 1b (Flamary et al. 2021 thus can be minimized by updating the representation mapping vb with gradient-based optimizers.
Definition 3.1. Let TT (x) and (TT (x) be the empirical distributions of covariates at a mini-batch level, which contain treated units and m untreated units, respectively; (r) and PT--0 (r) be
that of representations induced by the representation mapping = (x)) defined in Definition 2.2
However, since prevalent neural estimators mainly update parameters with stochastic gradient meth- ods, only fraction of the units is accessible within each iteration. A shortcut in this context is to calculate the group discrepancy at a stochastic mini-batch level:
W M == W ( (r), PT (TTTT(())) (r)).
(6)
To further investigate the effectiveness of this shortcut, Theorem 3.1 demonstrates that PEHE can be optimized by iteratively minimizing the factual outcome estimation error and the mini-batch group discrepancy (6). The proof of the theorem can be found in Appendix A.3
Theorem 3.1. Let sp and & be the representation mapping and factual outcome mapping, respectively; We be the group discrepancy at a mini-batch level. With the probability of at least 1 we have:
EPEHE (v,, (6,,0) s {
T-1 (v, T-0 (v.0)+ B 20 0
5N
T-1 and cy are the expected errors of factual outcome estimation, N the butch size, dy is the variance of outcomes, B3 is a constant term, and (()) is sampling complexity term
where Ep
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 (.) inspires us to investigate the
potential risks of bad cases caused by stochastic sampling. Precisely, O--) results from the discrep- ancy between the entire population and the sampled mini-batch units (see (30) and 32 in Appendix A which is highly dependent sampling quality. Therefore, the discrepancy mea- sure should be robust to bad sampling cases, other- wise the resulting huge variance will impede it from computing and reducing the actual discrepancy.
the uncontrollable
(a) Ideal
(b) Imbalance
(c) Outlier
The OT discrepancy in (6) can be easily disturbed by many sampling cases, except for the ideal case in Figure 2(a) where the transport strategy is rea- sonable and applicable. For example, according to Figure 2bb it falsely matches units with unrelated factual outcomes in the mini-batch where the out- comes between groups are imbalanced; according to Figure [ it falsely matches the mini-batch outliers to normal units, causing a substantial disrup- tion of the transportation strategy. In short, the vanilla OT technique in (6) fails to quantify the group discrepancy for producing erroneous transport strategies in non-ideal mini-batches, and thus misleads the update of the representation mapping vh. We summarize this phenomenon the mini-batch sampling effect (MSE) issue
Figure 2: Optimal transport plan (upper) and its geometric interpretation (down) in three cases, where the connection strength depicts the trans- ported mass. Different colors (vertical posi- tions) indicate different treatments (outcomes).
Liuyi et al. 2018 fail to handle the treatment selection hias for neglecting the MSE issue. The ability formalize the MSE issue through the mass-preserving constraint important advantage of OT over other techniques, as it provides grip for handling it.
'Apart from OT, most prevalent methods (Ma et al. 2022

---

This issue is attributed to the mass mass-prese preservation constraint in (5), which requires that all units in both groups match each other, regardless of the actual situation. bani-batch outliers, for instance, would be compelled to be transported according to Figure 2 which impedes the transport of normal units and the computation of the actual group discrepancy. A small batch size would exacerbate this defect. Definition 3.2. For empirical distributions a and B with and m units, respectively, optimal transport with relaxed mass-preserving constraint seeks the transport strategy π at the minimum cost:
WM (a. == min (D, T) ,I:== == argmin (D, -D)(T) <(DKL((D (T1mm; a) ( b)) (8)
where D E RR is the wnit-wise distance, and a; b indicate the mass of units in a and B, respectively.
An intuitive approach to mitigate MSE is to relax the marginal constraint and allow for the creation and destruction of unit's mass. To this end, relaxed mass- muss-preserving regularizer (RMPR) is devised in Definition] 3.2] which replaces the hard marginal constraint in (4) with a soft penalty in (8) for
deriving transport strategy. In this context, the stochastic discrepancy is calculated
W W.S (r),
where the hard mass prsseerreeevvtion constraint is removed mitigate the MSE issue. Inspired by Fatras
it al. 2021). the robustness of RMPR to sampling effects can be further theoretically investigated in
Theorem 3.2, where the effect of mini-batch outliers is upper bounded by constant.
Theorem 3.2. For empirical distributions a, B with n and units, respectively, adding an outlier d' to a and denoting the disturbed distribution as a', we have
ypg0. (a',3) -0p000 3) S 2x(111 e Emn(d'-d) )/n,
(10)
which is upper bounded by 2k/nn. W00x is the unbalanced discrepancy as per Definition 3.2
In addition, compared with alternatives (Xu al., 2020 Chapel et al. 2021) relax the marginal constraint, the approach in Definition 3.2 has better metric properties (Séjourné et al. (0199) and can be accelerated via the generalized Sinkhorn algorithm (Chizat et al. 2018). is differentiable w.r.f.
ai and thus can be minimized via stochastic gradient methods in an end-to-end manner.
3.3 PROXIMAL FACTUAL OUTCOME REGULARIZER FOR UNOBSERVED CONFOUNDERS
Existing representation-basee methods fail to elim- inate the treatment selection bias due to the unob- served confounder effects (UCE). Beginming with CFR (Shalit et al.] 2017). the unconfoundedness assumption 1 (see Appendix A is often taken to
x
X
circumvent the UCE issue (Ma et al. 2022). Given two units ri E pT- (r) and Tj E P
(a)
(b)
for in- stance, optimal transport in Definition 3.2 calcu- lates the unit wise distance Dij == |r:-r Ty If Assumption A. holds, this approach mitigates the treatment selection bias since it blocks the backdoor path X - 7' in Figure ((a] by balancing the confounders across groups in a latent space. However, Assumption A. is usually violated in practice as per Figure 3(b which hinder existing methods including OT from handling treatment selection bias since the backdoor path X' + T is not blocked.
Figure 3: Causal graphs with (a) and w/o (b) the unconfoundedness assumption. The shaded node indicates the hidden confounder X'.
According to Figure 3(b) given balanced X and identical T. the only variable reflecting the variation of X' is the outcome Y. As such, inspired by the joint distribution transport technique (Courty et al. 2017a), we propose to calibrate the unit-wise distance D with the potential outcomes as follow:
D7- Ir. -y.--r1o 7 [[luT~0 v -012 + 11o]-1- 11]]
(11)
where y controls the strength of regularization. The underlying regularization is: units with similar (both observed and unobserved) confounders should have similar potential outcomes. As such, for a pair of units with similar observed covariates, iee. Ir, - ry a 0, if their potential outcomes given the same treatment t {0, 1} differ greatly, i..., lv! vIll >> 0, their unobserved confounders should
likewise differ significantly. The vanilla OT technique in 6 with Dij = ||r rryl? would incorrectly

---

match this pair because lr; 1r:-r, rryl s 0, generate false transport strategy, and consequently misguide the update of the representation mapping vb. In contrast, OT based on D would not match this pair as the difference of unobserved confounders is compensated with that of potential outcomes.
Moreover, since v, and y in [TI are unavailable due to the missing counterfactual outcomes, the proposed proximal factual outcome regularizer (PFOR) uses their estimates instead. Specifically, let ÿ and Ūy be the estimates of y and v =1 respectively, PFOR refines [1I as
D = In: - -ry1² + 7 [l Ns l2 ||ù yı 13]
ÿ = (Ti), 0 (r;))
(12)
Additional justifications, assumptions and limitations of PFOR are discussed in Appendix D.3
3.4 ARCHITECTURE OF ENTIRE SPACE COUNTERFACTUAL REGRESSION
The architecture of ESCFR is presented in Figure ((b) where the covariate X is first mapped to the representations R with (-), and then to the potential outcomes with @(.). The group discrepancy W
calculated with the optimal transport equipping with the RMPR in (8) and PFOR in (12).
The learning objective is to minimize the risk of factual outcome estimation and the group discrepancy,
Given mini-batch distributions pT-1 (xx) and (z) in Definition 3.1 the risk of factual outcome estimation following Shi et al. (2019) can be formulated as
LF (v) 6) E z (2) di (1(x;)) v:1 B o (x) (00(v(z;)) y3l
(13)
where y and y the factual outcomes for the corresponding treatment groups. The discrepancy is: SD C (v) w (r),
(14)
which is in general the optimal transport with RMPR in Definition 3.2 distance D calculated with the PFOR (12). Finally, the overall learning objective of ESCFR
except for the unit
ESCPR == (F(,, の) + X LT (bb),
(15)
where A controls the strength of distribution alignment, controls the entropic regularization in (5)
controls RMPR in (8), and 7 controls PFOR (2).. The learning objective (55) 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 First, compute " FiN by solving the linear programming problem in Definition 3.2 with Algorithm where the unit-wise distance is calculated with D", Second, compute the discrepancy in (14) as (mT* D7). where D is differentiable v making it feasible minimize this discrepancy with its gradient w.r.t. th. Finally, calculate the overall loss in (15) and update yi and & with stochastic gradient methods.
EXPERIMENTS
4.1 EXPERIMENTAL. SETUP
Datasets. Missing counterfactuals impedes the evaluation of PEHE with observational benchmarks. Following Liuyi et al. 2018); (halit et al. (2017). experiments are conducted on two semi-synthetic benchmarks. Specifically, the THDP benchmark aims to estimate the effect of specialist home visits on infants' potential cognitive scores, with 747 observations and 25 covariates; the ACIC dataset comes from the collaborative perinatal project, with 4802 observations and 58 covariates.
Baselines. We consider three groups of baselines. (1) Statistical estimators: least square regression
with the treatment as covariates (OLS), random forest with the treatment as covariates (R. Forest,, single network with the treatment as covariates (S.learner by Künzel et al. (2019)), separate neural regressors for each treatment group (T.learner by Künzel et al. (2019)), ) TARNet (Shalit et al. 2017):
(2) Matching estimators: propensity score match with logistic regression (PSM by Rosenbaum &
Rubin (1983))), k-nearest neighbor (k-NN by Crump et al. (2008) ) causal forest (C.Forest by Wager
& Athey (2018)), ). orthogonal forest (O.Forest by Wager & Athey (201)); (3) Representation basedd based
methods: balancing neural network (BNN by Johansson et al. (2016)) counterfactual regression with
MMD (CFR-MMD) and Wasserstein metric (CFR-WASS)) by Shalit et al. (2017).

---

Table I: Performance (meantstd) on the PEНE and AUUC metrics. "*" marks the baseline estimators that ESCFR outperforms significantly at p-value < 0.05 over paired samples t-test.
Dataser
(PEHE)
IHDP (PEHE)
(AULC))
IHDP (AUUC)
Model
la-sample
Out-samplc
In-sample
Out-sample
In-sample
Out-sample
In-sample
-at-sample
OLS
3.749=0.080'
4. 4,340+0.111 117"
3.85646.018
5.67449.000
0.843=0.007
0,496=0.017
0.65240.050
0.492+0.0322
R.Forest
3.597=0.004
3,999: 165
2.63513.
4.671+9.291
0.90200.016
0.702:0.026'
0,736x0.142
0.661:0
S.Learner
25992-0.269'
363620.2541
1.706=1.600
3.03825,319
0.90520.041
0.62740.014'
0.63540.183
0.70210.330
T.Leamee
3.429=0. 142"
3.566.40 248
1.967+1.136
2.730+3.627
0.846=0.019
0.632+0.020'
0.651+0.179
0.307=0.333
TAKNet
3.236=0: 266"
3.254=0. 150"
0.749H0.291
1.788822.812
0.86620:046
0.662:0.014
0.654+0.184
0.711
C.Forest
3.449=0. 101*
3. 3.196++0 177"
4. 8+5 602"
4, 4.486+88.66 677
0.717=0.005*
0170
0,643 141
0.695ed. 294
5.60520.168
5.89240.
2.20812.233'
4.31917.336
0.89220.007
0.50740.034"
0.72510.116
0.66810.299
O F.ross
8.094g4.669'
4.148=2.224'
2.60542.418'
3.13645.644 6.22
0.744=0.013
0.69940.022'
0.86440.155
0.702+0.325
PSM
5.228-0.1550 154*
5.094+0.301 301*
3,222
4.634+8.174
0.884-0.010
0.745+0.021
0.740=001199
0.681+0.25
3.34520.233"
1.36840.170
0.70940.330
1.80612.88
0.852:0.033
0.61510.013'
0.65410.118
0.7.
3.182=0.1744
3.357=0.321
0.77740.327
1.791+2.744 7441
0.87100.032
0.659=0.017
0.655+0.118
0.710=0.329
CHR- WASS
3.128=0,369"
3,207:0 169*
0.657406673
1,704+5.115
0.873=0.029
0.66-00.018'
0.656+00187
0.715e0,329
2.252=0.297
2.31620.613
0.607=0.328
1.25720.677
0.79620.030
0.75420.021
0.65910.187
0.734:0.329
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 a fair comparison, all neural models are trained for 400 epochs with Adam optimizer, with the learning rate and weight decay both set to 0.001. Other settings of optimizers follow Kingma & Ba (2015). We fine-tune hyperparameters within the range in Figure [ validate performance every two epochs, and save the optimal model for test.
Evaluation protocol. Following Liuyi et al. 2018). PEHE in (3) is used as the precision metric for performance evaluation. However, is unavailable in the model selection phase due to missing counterfactuals. As such, we use the area under the uplift curve (AUUC) (Betlei et al., 202I) to guide model selection, which measures the counterfactual ranking performance of the model and can be computed without counterfactual outcomes. The within within-sample and out-of-sample results are reported on the training and test data, respectively, following (Shalit et al. (2117).
4.2 PERFORMANCE COMPARISON
Table I reeport the performance of ESCFR and its competitors over ten runs. Statistical estimators exhibit competitive performance on the PEHE metric. Due to the superiority to depict non-linearity, neural estimators outperform linear and random forest methods. In particular, TARNet that absorbs the advantage of T-learner and S-learner achieves the best overall performance in statistic estimators. However, the circumvention to treatment selection bias leads to inferior performance.
Matching methods, e.g., PSM exhibit compelling ranking performance, which explains why they are favored in counterfactual ranking practice (Betlei et al. 202I). However, their poor performance on PEHE hinders their application in counterfactual estimation applications such as advertising systems that place more emphasis on the accuracy of treatment effect estimation.
Representation- based methods mitigate the treatment selection bias and enhance overall performance. In particular, WASS reaches an out-of-sample PEHE of 3.207 on ACIC, significantly outper- forming most statistical methods. However, the MSE and UCE issues impede these methods from solving the treatment selection bias. The proposed ESCFR achieves significant improvement over most metrics compared with various prevalent baselines Combined with the comparisons above, we attribute its superiority to the proposed RMPR and PFOR regularizers, which makes it robust to MSE and UCE. See Appendix D.5 for additional comparison results.
4.3 ABLATION STUDY
To verify the effectiveness of individual components, an ablation study is conducted on the ACIC benchmark Table Specifically, ESCFR first augments TARNet with stochastic optimal transport in Section 3.1 which effectively reduces the out-of-sample PEHE from 3.254 to 3.207. Then, it
An exception would be the within-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 on out-of-sample AUUC instead evaluate the ranking performance and perform model selection.

---

Table 2: Ablation study (mean=std) the ACIC benchmark. "** marks the variants that ESCFR outperforms significantly at p-value < 0.01 over paired samples t-test.
In-sample
Out-sample
SOT
RMPR
PFOR
PERE
AUUC
PEHE
AUUC
3.2367=0.2666
0.8862+0.0462
3.2542+0.1505
0.6624+0.0149
X
1284-002000 2638
0873440.0291
3.2073+0.
0.609%+0.01887
X
64549-0.22222 2747"
083564000286
2.768880
0.7199-0.0177
5,355000,4403
0.827010.0341
2. 2.6330x0.4667 46722
0,7110:0.0287
2.2520
0.796810.0307
2.3165+0.6136
0.7542:0.0202
mitigates the MSE issue with RMPR in Section 3.2 and the UCE issue with PFOR in Section 3.33
reducing the out-of-sample PEHE 2.768 and 2.633, respectively. Finally, ESCER combines the
RMPR and PFOR in a unified framework in Section 3.4 reducing the out-of-sample PEHE to 2.316.
4.4 ANALYSIS OF RELAXED
REGULARIZER
Most prevalent methods fail to cope with the label imbal- ance and mini-batch outliers in Figure bbb-).. Figure 4
shows the transport plan generated with RMPR in the same situations, where RMPR alleviates the MSE issue in both bad cases. RMPR with λ 10, for instance, avoids the in- correct matching of units with different outcomes; RMPR with x = 2 gets robust to the outlier's s interference and cor- rectly matches the remaining units. We attribute the success to the relaxed mass-preserrva constraint in Section 3.2
Notably, RMPR does not transport all mass of unit. The closer the unit is to the overlapping zone in a batch, the greater the mass is transferred. That is, RMPR adaptively matches and pulls closer units that are close to the overlap- ping region, ignoring outleers, which mitigates the bias of causal inference methods in cases where the positivity as- sumption does not strictly hold. Current approaches mainly achieve it by manually cleaning the data dynamically weighting the units (Johansson et al. 2020), while RMPR naturally implements it via the soft penalty in (8).
(a) 10
(b)
Figure 4: Geometric interpretation of optimal transport plan with RMPR un- der the outcome imbalance (upper) and outlier (down) settings. The dark area indicates the transported mass of unit, i.e., marginal of the transport matrix The light area indicates the total mass.
We further investigate the performance of RMPR under different batch sizes and x in Appendix D.2 to verify the effectiveness of RMPR more thoroughly.
4.5 PARAMETER SENSITIVITY STUDY
We discuss four critical hyperparameters in ESCER, i.e., X, e, Ki 7, which are the weights in the learning objective and influence the final performance significantly.
We firstly vary investigate the influence of stochastic optimal transport. Specifically, increasing consistently improves the precision of ITE estimates. For example, the out-of-sample PEHE reduces from 3.22 À = 0 to approximately 2.85 À = 1.5. However, assigning higher focus on distribution balancing in multi-task learning framework leads to difficulties in factual outcome estimation, leading to sub-optimal ITE estimates. Then, we vary c to investigate the entropic regularizer. It does accelerate the computation of optimal transport discrepancy with large corresponding to faster computation speed (Flamary et al.) 202T). However, comes with a biased transport plan, evidenced by the oft-of-sample PEHE increasing to 2.95 with fluctuations.
We further vary y and k to showcase the influence of PFOR and RMPR, respectively. PFOR benefits ITE estimation, while assigning considerable weight to the proximal outcome distance is detrimental to ITE estimates, because the unit-wise distance calculated with representations (covariates) would be neglected. Relaxing the mass preserving constraint via RMPR could significantly improve the ITE estimates; however, too-small x is always detrimental as we cannot guarantee that the representations across treatment groups are drawn closer together in the optimal transport approach.

---

3.2
3.0
In-sample Out-sample
In-sample Out-sample
In-sample Out-sample
In-sample
2.9
2.9
Out-sample
PEHE
3.0
PEHE
2.8
2.8
2.8
2.7
2.6
2.6
2.6
2.0
3.0
0.5
1.0
5.0
10.0
15.0
25.0
(c)
Figure 5: Parameter sensitivity study for critical hyper-paramettess of ESCER
5
RELATED WORKS
Current works mitigate the treatment selection bias by balancing the distributions of treated and untreated groups, which can be divided into three groups according to the methods used for balancing.
Reweighting-based methods weight individuals with balanced scores to achieve globally balanced distributions, represented by the inverse propensity score (IPS) approach (Rosenbaum & Rubin 1983a and its doubly robust variant (Robins et al. 1994). Imai & Ratkovic (2014) and and Fong et al. (2018) propose calculating the balancing score by solving optimization problem. Kuang et al. (2017b) and Kuang et al. (201a) consider additional non-confoun factors in covariates. However, these methods susceptible to non-overlapping units and suffer from high variance issue.
Matching-based methods match similar units from different groups to construct locally balanced distributions. Representative methods (Rosenbaum & Rubin 1983b Chang & Dy [017 Li et al.
2016) mainly differ in terms of similarity measures. Notably, Tree-based methods (Wager & Athey 2018) can also be considered matching methods with adaptive neighborhood metries. However, the computational cost hinders the application of these methods in large-scale scenarios.
Begining with BNN (Johansson et al., 2016) and CFR (Shalit al., (017). representation- methods minimize the group discrepancy in the latent space. Liuyi et al. (2018) and Hassanpour & Greiner 2020) further augment CFR with local similarity and non-confounding factors, respectively. Kallus (020) and Yoon et al.] (2018) propose to balance the distributions of representations with adversarial training. Due to its scalability and avoidance of the high variance issue, representation- based methods have been predominant for handling the treatment selection bias.
necessary to distinguish ourselves from emerging OT-based causal inference approaches. [Dnii pace (202T) augments the IPS method with propensity score estimator based on OT; however, it is limited by the aforementioned high variance issue. Torous et al. (022] uses the push-forward operator to improve change-in-change models; however, they are designed for multi-phase data which is not available in our case. Li et al. 2022) has a similar setup to to us, while it focuses on the decomposition of latent variables and is identical to (Shalit et al. (2017) in terms of alignment technology. Our work is a new take on OT under the CFR framework, alleviating the MSE and UCE issues that have been long neglected by the causal inference community until this year.
CONCLUSION
Due to the effectiveness of mitigating treatment selection bias, representation learning has been the primary approach to estimating individual treatment effect. However, existing methods neglect the mini-match sampling effects and unobserved confounders, which hinders them from handling the treatment selection bias. A principled approach named ESCFR is devised based on generalized Sinkhorn discrepancy. Extensive experiments demonstrate that ESCFR can largely mitigate MSE and UCE issues, and achieve better performance compared with other baseline models.
There are two directions of future work that we intend to pursue. The first direction attempts to con- struet the representation mapping with normalizing flows (Chen et al. 2018). which is invertible and thus satisfies the assumption to the representation mapping by (Shalit et al. (2017). Another direction seeks to extend methodology to industrial applications, e.g., debias in recommenders (Wang et al.) 2022), which has long been dominated by reweighting methods with the high variance issuc.

---

REFERENCES
Jason M. Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algo-
rithms for optimal transport via sinkhorn iteration. In NeurlPS, pp. 1964-1974, 2017.
Artem Betlei, Eustache Diemert, and Massih-Rees Reza Amini. Uplift modeling with generalization guarantees. In SIGKDD, pp. 55 55-65, 2021.
Jose H. Blanchet, Arun Jambulapati, Carson Kent, and Aaron Sidford. Towards optimal running times for optimal transport. CoRR, 0bs//810.077) 2018.
François Bolley, Arnaud Guillin, and Cédric Villani. Quantitative concentration inequalities for empirical measures on non-compact spaces. Probability Theory and Related Fields, 137(3): 541-593, 2007.
Nicolas Bonneel, Michiel de Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. ACM Trans. Graph., 30(6):158, 2011.
Yale Chang and Jennifer G. Dy. Informative subspace learning for counterfactual inference. In AAAI, pp. 1770-1776, 2017.
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, pp. 23270-23282, 2021.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In NeurIPS, pp. 6572-6583, 2018.
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, pp. 1802-1806, 2022.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Math. Comput., 87(314):256-225 2018.
José Cordero, Víctor Cristóbal, and Daniel Santín. Causal inference on education policies: survey of empirical studies using pisa, timss and pirls. J. Econ. Surv., 32());878911 2018.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In NeurlPS, pp. 3730-3739, 2017a.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamoniy. Optimal transport for domain adaptation. IEEE Trans. Pattern Anal. Mach. Intell., 39(9): 2017b.
Richard K Crump, V Joseph Hotz, Guido W Imbens, and Oscar A Mitnik. Nonparametric tests for treatment effect heterogeneity. Rev. Econ. Stat., 90(3):389 2008.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In NeurIPS, pp. 2292-2300, 2013.
Eric Dunipace. Optimal transport weights for causal inference. CoRR, abs/2109.09991,, 2021.
Pavel E. Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's 's algorithm. In ICML, volume 80, pp. 3366-1375. 2018.
Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In JCML, volume 139, pp. 3197, 2021.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon. Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T. H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: python optimal transport. J. Mach. Learn. Res.,
2021.

---

Christian Fong, Chad Hazlett, and Kosuke Imai. Covariate balancing propensity score for continuous treatment: Application to the efficacy of political advertisements. Ann. Appl. Stat., 12((1:156-11) 2018.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In ICLR. OpenReview.net, 2020.
Xiaofeng He, Guoqiang Xu, Cunxiang Yin, Zhongyu Wei, Yuncong Li, Yancheng He, and Jing Cai. Causal enhanced uplift model. In PAKDD, volume 13282, pp. 119-131, 2022.
Kosuke Imai and Mare Ratkovie, Covariate balancing propensity score. R. Stat. Soc. Series. Stut. Methodol., 76(1):243 2014.
Arun Jambulapati, Aaron Sidford, and Kevin Tian. direct o (1//esillo) iteration parallel algorithm for optimal transport. In NeurlPS, pp. 11355-11366, 2019.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How to escape saddle points efficiently. In ICML, pp. 1724-1732, 2017.
Fredrik D. Johansson, Uri Shalit, and David A. Sontag, Learning representations for counterfactual inference. In ICML, pp. 0020-3029, 2016.
Fredrik D. Johansson, Uri Shalit, Nathan Kallus, and David A. Sontag. Generalization bounds and rep- resentation learning for estimation of potential outcomes and causal effects. CoRR, abs/2001 .07426, 2020.
Nathan Kallus. Deepmatch: Balancing deep covariate representations for causal inference using adversarial training. In ICML, volume 119, pp. 5067-5077, 2020.
Leonid V Kantorovich. On the translocation of masses. J. Math. Sci., 133(4):1381-1322, 2006.
Wenwei Ke, Chuanren Liu, Xiangfu Shi, Yiqiao Dai, Philip S. Yu, and Xiaoqiang Zhu. Addressing exposure bias in uplift modeling for large-scale online advertising. In ICDM, pp. 1156-1161, 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015,
Kun Kuang, Peng Cui, Bo Li, Meng Jiang, and Shiqiang Yang. Estimating treatment effect in the wild via differentiated confounder balancing. In SIGKDD, pp. 265-274, 2017a.
Kun Kuang, Peng Cui, Li, Meng Jiang, Shiqiang Yang, and Fei Wang. Treatment effect estimation with data-driven variable decomposition. In AAAI, pp. 140-146, 2017b.
Sören Reinhold Künzel, Jasjeet Sekhon, Peter Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proc. Nat. Acad. Sci. U. S. A., 116(10): 4156-4166 2019.
Qian Li, Zhichao Wang, Shaowu Liu, Gang Li, and Guandong Xu. Deep treatment-adaptivee network for causal inference. VLDBJ, pp. 1-16, 2022.
Sheng Li, Nikos Vlassis, Jaya Kawale, and Yun Fu. Matching via dimensionality reduction for estimation of treatment effects in digital marketing campaigns. In IJCAI, pp. 3768-3774, 2016.
Tianyi Lin, Nhat Ho. and Michael I. Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In ICML, volume 97, pp. 3982-3991, 2019.
Yao Liuyi, Li Sheng, Li Yaliang, Huai Mengdi, Gao Jing, and Zhang Aidong. Representation learning for treatment effect estimation from observational data. In NeurlPS, pp. 2638-2648, 2018
Jing Ma, Mengting Wan, Longqi Yang, Jundong Li, Brent Hecht, and Jaime Teevan. Learning causal effects hypergraphs. In SIGKDD, pp. 1202-1212, 2022.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sei., pp. 666-704, 1781.

---

X Nie and S Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2): 299-319, 2020.
Elizabeth L Ogburn and Tyler J VanderWeele. On the nondifferential misclassification of binary confounder. Epidemiology: 23(3):433, 2012.
Judea Pearl and Dana Mackenzie. The book of why: the new science of cause and effect. Basic books, 2018.
Ofrr Pele and Michael Werman. Fast and robust earth mover's distances. In ICCV, pp. 460-467 2009.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Found. Trends Mach. Learn., 11 (5-6):355-607, 2019.
Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In ICMI., volume 119, pp. 7673-7682, 2020.
levgen Redko, Amaury Habrard, and Mare Sebban. Theoretical analysis of domain adaptation with optimal transport. In PKDD, volume 10535, pp. 737-753, 2017.
James Robins, Andrea Rotnitzky, and Lue Ping Zhao. Estimation of regression coefficients when some regressors are not always observed. J. Am. Stat. Assoc., 89(427):846-866, 1994.
Paul Rosenbaum and Donald B Rubin. The central role of the propensity score observational studies for causal effects. Biometrika, 70(1):41-55, 1983a.
Paul Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-53 1983b.
Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Educ. Psychol., 6685))688. 1974.
Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. In AAAI, pp- 5612- 5619, 2020.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In ICLR, 2018.
Thibault Séjourné, Jean Feydy, François-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport. CoRR, abs/1910.12955 2019.
Thibault Séjourné, François- Xavier Vialard, and Gabriel Peyré. Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d woafe- In AISTATS, volume 151, pp. 4995-5021, 2022.
Uri Shalit, Johansson Fredrik D, and Sontag David. Estimating individual treatment effect: general- ization bounds and algorithms. In ICML. pp. 3076-3085, 2017.
Claudia Shi, David M. Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. In NeurlPS, pp. 2503-25133 2019.
Tamar Sofer, David Richardson, Elena Colicino, Joel Schwartz, and Eric Tehetgen Tchetgen. On negative outcome control of unobserved confounding generalization of difference-in- differences. Stat. Sci., 3183):348, 2016.
William Torous, Florian Gunsilius, and Philippe Rigollet. An optimal transport approach to causal inference. CoRR, fis/2108.05855. 2021.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. J. Am. Stat. Assoc., 1 113(523): 12288-21 1242, 2018.

---

Hao Wang, Kei Chang, Tianqiao Liu, Jianmin Huang, Zhichao Chen, Chao Yu, Ruopeng Li, and Wei Chu. Esem Entire space counterfactual multi-task model for post-click conversion rate estimation. In SIGIR, 2022.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning. volume 2. MIT press Cambridge, MA, 2006.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. ICLR, 2020.
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, pp. 2540- 2546, 2020.
Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. ACE: adaptively Similarity-preserved representation learning for individual treatment effect estimation, In ICDM, pp. 1432-1437, 2019.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: estimation of individualized treatment effects using generative adversarial nets. In ICLR, 2018
Jiajing Zheng, Alexander D' Amour, and Alexander Franks. Copula-based sensitivity analysis for multi-treatment causal inference with unobserved confounding. CoRR, abs/2 122.0941,, 2021.

---

A CAUSAL INFERENCE WITH OBSERVATIONAL STUDIES
We introduce all the necessary preliminaries about causal inference for readers and reviewers that are unfamiliar with this area.
A.1 PROBLEM FORMULATION
This section formalizes the definitions, assumptions, and useful lemmas causal inference from observational data. Following the notations Section 2.1 an individual with covariates æ has two potential outcomes, namely Y1 (x) given it is treated and (z) otherwise. The ground-truth individual treatment effect (ITE) is the difference in its potential outcomes.
Definition A.1. The individual treatment effect (ITE) for unit with covariates x is
T(x) E[ Yo! =],
(16)
where we abbreviate (1(2:) to Y1 for brevity. The expectation is over the potential outcome space y.
Estimating ITE with observational data is common practice in causal inference, which has long been confronted with two primary challenges:
Missing counterfactuals: where only the factual outcome is observable. If patient is treated, for instance, we can never observe what would have happened the patient was untreated in the same situation.
Treatment selection bias, where individuals have preferences for treatment selection. For example, doctors would adapt different treatment plans for patients with different health conditions. It would make the treated and untreated populations heterogeneous. ITE estimators naïvely trained minimize the factual outcome error would overfit the respective group's s properties and thus cannot generalize well to the entire population.
Pearl & Mackenzie 2018)) suggested a two-step methodology to overcome these two challenges. The first step is identification, which aims to construct unbiased statistical estimand to identify the causal estimand (e.g., T())) based on the adjustment formula. Note that not all causal estimands are identifiable, e.g., ITE is identifiable only if Assumption A.3
Assumption A.1. (Unconfoundedness). For all covariates x in the population of interest (i.e., I with IP(X x) 0), we have conditional independence (Y,, Y1) 1 T X X x. That is, potential outcomes are conditionally independent of treatment assignment.
Assumption A.2. (Consistency). For all covariates x in the population of interest, we have = Y,. That is, the observed outcome is consistent with the potential outcome w.r.t. the assigned treatment. Assumption A.3. (Positivity). For covariates x in the population of interest, we have 0 < P(T 1 X xx) 1. That is, all individuals have chance be assigned both treatments.
The second step is estimation, which aims to estimate the derived statistical estimand with observa- tional data. Lemmal A.l iillustrates how this two-step approach can be used for ITE estimation. Lemma A.1. The ITE estimand Txx) can be identified as:
E [Y - Yo X x] E Y
X x]
E No
X = x]
X x, 1] E Yo X x. =
(17)
Y X 7 = = E X = -
where (1) stems from the unconfoundednes: assumption A. (2) stems from the consistency assump- Ai. The derived estimand is fully composed of statistical estimands, which can only be estimated under the positivity assumption A.3.. Otherwise, the positivity assumption is violated, we have:
X= T= 1]= f y- y-P(Y- = X x. T= dy
E Y
(18)
f P = P( x)
= =(==
dy.
which is not computable as there exists X which makes P(T X x)

---

X
Y
X
UP
4
%
0
(a) S-learner
(b) T-learner
(c) TARNet
Figure 6: Architecture of Meta-learner based ITE estimators, consisting of inputs (green), outputs (white), shared mappings (yellow), and mappings for treated and untreated units (red and blue, respectively).
A.2 META-LEARNERS FOR ITE ESTIMATION WITH OBSERVATIONAL DATA
In an effort solve missing counterfactuals, existing meta-learner based methods (Künzel al. 2019; Nie & Wager 2020) decompose the ITE estimation problem into several subproblems that can be solved with any supervised learning method. As depicted Figure S-learner regards the treatment indicator T as one of the covariates X, and utilizes the shared representation mapping & and outcome mapping @ to estimate the factual outcomes. However, because the network structure does not highlight the role of treatment indicator, it may be overlooked when treatment effects are minimal. T-learner models the factual outcomes for treated units X₃ and untreated units Xo separately, which highlights the treatment indicator's effect; however, it reduces the data efficiency and is therefore inapplicable when the dataset is small. Künzel et al. (2019) discuss the advantages and limitations of these approaches in more detail.
Definition A.2. Let 3p: X + R be mapping from support X' to R. That is, VI E X, 3r= =()) e R. Let o: R + y be a mapping from support T y. That is, it maps the representations and treatment indicator to the corresponding factual outcome. For example, Yi Φ1 (R), Yo do R)), where we will always abbreviate ((R, = and d R, T 0) to o1 (R) and on (R), respectively, for brevity.
Assumption A.4. 6: - R is differentiable and invertible, with its inverse h-1 defined over R.
TARNet (Shalit et al. 2017 in Figure 6 (c) obtains better results by absorbing the advantages of both T-learner and S-learner, which consists of a representation mapping i and an outcome mapping á as defined in Definition A.2 For unit with covariates TARNet estimates ITE as the difference in predicted outcomes when 7 is set to treated and untreated:
o Y1 Yo- where +xx)(v(X))
= =o=o)(((X),
(19)
where / is trained over all units, Φ1 and do are trained over the treated and untreated units, respectively, minímize the factual error 6p (0, 1h) in Definition A.33 Finally, the performance of the ITE estimator mainly evaluated with PEHE:
EPEHE (b. 6) [x (f).( P(x)
(20)
Let IL be the loss function that measures the quality of outcome estimation, e-g., the squared loss. The expected loss for the units with covariates and treatment indicator t is:
Definition A.3.
(x,((x,t)= S, S, P(Y dY,
(21)
where L is realized with the squared loss: L(Y,, (((x). -Y: (((((((()),))) in our scenario. The expected factual outcome estimation error for treated, untreated and all units are:
+ T=1 (v,,)) = S 1) ()..1)))pp))(eeee
T-o (v. f ( dr, CF
(22)
(F) Sr fxurlun (x, Ip( dxdt.

---

A.3 REPRESENTATION -BSED METHODS FOR TREATMENT SELECTION BIAS
However, the treatment selection bias makes covariate distributions across groups shift. As such, Φi and do would overfit the respective group's 's properties and thus cannot generalize well to the entire population. For example, as shown in Figure [[a (a), the potential outcome estimator Φ1 trained with treated units cannot generalize to the untreated units. Therefore, the resulting would be biased.
Definition A.4. Let (x) == P( T = 1) and (x)) == P(x T = 0) be the covariate distribution for treated and untreated groups, respectively. Let (r) and v (x) be that of representations induced by the representation mapping r (x) defined in Definition 2.2
To mitigate the effect of treatment selection bias, representation-basel approaches (Johansson et al. 2016; Shalit et al. (017)) minimize the distribution discrepancy of different groups in the representa- tion space. In particular, the integral probability metric (IPM) in Definition A..4|i A.4fis widely used metric that measures the discrepancy of two distributions. Shalit et al. (017)) propose to optimize the PEHE by minimizing the estimation error of factual outcomes OF and the IPM of learned representations between treated and untreated groups They further provide theoretical results to back up their claim as per Theorem A.T
Definition A.5. Consider two distribution functions FT= (x) and (z) supported over X, let F be a sufficiently large function family, the integral probability metric induced by F is
IPMF (1FT=1,
=sup feF
Jxx f(x) (x)-
dx
(23)
Theorem A.1. Let ab and o be the mappings in Definition 2.2 F be a predefined sufficiently large function family of é. 1 be the integral probability metric induced by F. Assume there exists 14..0 (I,, 1) E F holds. Shalit et al. (017)) demonstrate:
constant B& > O, such that for E 10
B
éPEHE 1b, () $ (
T=0 (1b, Φ) 6q
T=1
(vb + 1& ) -20ł)
(24)
where FF p0 and of-1 follow Definition
A.3
PT-1 (r) and v (xx) follow Definition A.4
v
A.4 THEORETICAL RESULTS AND EXTENSIONS
In this section, we describe the intuition of Theorem 3.1 and Theorem 3.2 and provide rigorous proofs to support our claims. Two weaknesses hinder the existing Theorem A. by (Shalit al. (2017) from supporting representation--asse causal inference approaches. Firstly, the IPM is a discrepancy metric with profound theoretical properties while being difficult to compute numerically. To counter this, note that the IPM holds for any sufficiently large function families, e.g., the 1-Lipschitz function family where the IPM is equivalent to the Wasserstein divergence as per Cantorovich-Rubinstee duality VVillani 2009). As such, shorteut IPM is the Wasserstein distance per Lemma [A.2]
Lemma A.2. Consider two distribution functions P1 (z) and P2(x)) supported over X; let be the family of "Lipschitz,function W be the Wasserstein distance, Villani (2009) demonstrates
1PMs (PI,,P2) IP2) = W (P1,P2)
(25)
Another weakness is that neglects the Mini-bater Sampling Effects (MSE). Specifically, Theo- rem A. holds only if the entire populations of treated and untreated groups are available. However, since the representation- basee approaches update parameters with stochastic gradient methods, only a subset of the population is accessible within each iteration. As such, it remains questionable whether Theorem A. holds at a mini-batch level in practice.
Lemma A.3. Let Prx)) be probubility measure supported over X R.d satisfying Ti(A) inequality. Let P(x) 2/4 EN 8z, be the corresponding empirical measure with N units. Bolley et al. (007 and Redko el al. (2017) demonstrate that for any d' d and X' < X, there exists some constant No.
such that for any > 0 and N No maxxa(=--tt (d-22) 1), we have
X Ne²
P
exp
(26)
where d','' can be calculated explicitly.
'While there is risk of symbol reuse, use e here to denote sampling emror, and to control the strength of entropic regularization in optimal transport.

---

Hoeffding's inequality is powerful statistical tool quantify such sampling effects, which is proved to be applicable for V by Bolley et al. (2007). Therefore, it is natural to expand W according to Lemma A.3 to extend Theorem A.I to mini-batch situations, in order quantify the sampling effects. Theorem A.2. Let v and e be the representation mapping and factual outcome mapping, respectively;
Wy be the discrepancy across groups at mini-batch level. With the probability of at least - ö, we have:
ÖN }
EPENE (0) s T= (v, teF (v.6)) + B - -20- O
(27)
where e and = are the expected losses factual outcome estimation over treated and untreated units, respectively. N is the batch size, oy is the variance of outcomes, By is some constant such that By (x, belongs to the family of I-Lipschitz functioon, OC) is the sampling complexity term.
1
Proof. According Theorem A.I we have:
T=0 (vb Φ) T T-1 (v) B (P _ - ÉPEHE v. s (cF
(28)
Assuming that there exists a constant B3 > O, such that for / {0,
B (3,(( belongs to the
1
family of I-Lipschitz functions. According to Lemma A.2 we have
CPEНE Φ) s (GF T-0 (v. + oF (w, () B w (PT-1 (PT-I,PT-)) -
(29)
Following Definition 3.1 let (r) and PT- be the empirical distributions of representations mini-batch level, containing Ni treated units and No untreated units, respectively. Then we have:
W (p2-1 v
pT-=)) w (PPT-1 [P=11)++ w (IP3-0 1P1-1)
s w [P T- -T-1)++ W (ITT -T W (
(30)
w ( w (TT- + W. +
because we have the triangular inequality for W. The Hoeffding inequality in Lemma gives the following inequality which holds with the probability at least 1 8
A.3 further
w
(PT T 1T S
2og 5
/NN1
(31)
w (F pT-O) s
2
/NNo-
Denote N No N₁ as the batch size, = N1/ N the ratio of treated units in the current batch.
Combining 30 and 31 we have
W (P p pT==)
s Wp V 2 (
IN
2
5
1Nooo
W. p
2log
1
INN
(32)
V
1-0
1
: W. O
5N
where O() satisfies
Vogg ( / (1+ (1+ -1)
oN 4 /V log 1 /NN,
(33)
where O (NN) 0 0.5. This corollary can be derived by differentiating the function f(x) 1//6+ 1/VI-x. Combining
( reaches its maximum when 0 1/NV 0 0 1 1/N, reaches its minimum when
29 and
we have
B 20 0 1
CPEHE (1, 6) $2 6F T
(v) cT T
(34)
5N
where we denote 00,0 x/ O ( and the proof is completed.

---

Theorem A.2 extends Theorem A.1 and derives the upper bound of PEHE in the stochastic batch form, which demonstrates that the PEHE can be optimized by iteratively minimizing the factual outcome estimation enror and the optimal transport discrepancy at mini-batch level.
Corollary A.1. The empirical variance of the PEHE estimates in (27) largely depends on the batch size and the proportion of treated and untreated units. Large batch size and balanced proportion produce low empirical variance, and vice versa.
Proof. It can be drawn directly from 27 (batch size) and (33 (treatment proportion).
Theorem A.3. For discrete measures Ei=1 a,ox, and B Ee bydgg, adding an outlier 8x to and denote the disturbed distribution as d' we have
00 (a', (3) - (a. s 2x(1 - E (x'-xxx (2
(35)
which is upper bounded by 2k/(n 1). 780, is the unbalanced discrepancy as per Definition 3.2
Proof. This a direct corollary of the Lemma by Fatras et al. (202T). under the assumption that all the units including the outlier 8xr share the same mass.
B DISCRETE OPTIMAL TRANSPORT
This section proposes the definitions and algorithms to calculate optimal transport between discrete measures. We have omitted the case of general measures, it beyond the scope this work. Readers interested in this topic should refer to Peyré & Cuturi (2019) Cuturi (2013) for details.
B. PROBLEM FORMULATION
Optimal transport derives from the formulation of Monge (1781). We provide an equivalent interpre- tation under discrete measures. Consider warehouses and m factories, where the i-th warehouse contains a, units of materials; the j-th factory needs by units of materials. Now construct a mapping from warehouses to factories, satisfying: (1) all materials of warehouses are transported; (2) all requirements factories are satisfied; (3) materials from one warehouse are transported to no more than one factory (mapping constraint). Every feasible mapping is associated with global cost, calculated by aggregating the local cost of moving a unit of material from the i-th warehouse to the j-th factory. Our objective, find feasible mapping that minimizes the transport cost, is formulated in Definition[B. [B.1
Definition B.1. For discrete measures Σi= a,dx, and B Ej=1 bj the Monge problem seeks for a mapping T; {x,), + {x, ) that associates to each point x, a single point x; and pushes the mass of a to B. That is, Vj (l...., .m} we have b, = E=T(x.)xx, a, This mass pressrving constraint is abbreviated as Tra = B. The mapping should also minimize the transportation cost denoted as c(x, y). To this end, Monge problem for discrete measures is formulated as:
[d(x..(x))))
min T-T-a-8
(36)
This problem was further utilized to compare two probability measures where Σ; a; However, Monge's formulation cannot guarantec the existence and uniqueness of solutions (Peyré Cuturi 2019). Kantorovich] (2006) relaxed the mapping constraint by allowing the transport from one warehouse to many factories and reformulated the Monge problem linear programming problem in Definition B.2]
b,
Definition B.2. For discrete measures = Σi â;&x, and 3 Ey the Kantorovich problem aims to find a feasible plan n RRax which transports a to 3 at minimum cost:
Wa, 3) = min (D,r), 3) = { paxm T = a, π 1, = b). xe)I(a.8)
(37)

---

Algorithm Sinkhorn Algorithm
Input: discrete measures Σx=1 a,ox, and B = by distance matrix Dij = | x,
Parameter: e: strength of entropic regularization: mmax maximum iterations.
Output: T=: the entropic regularized optimal transport matrix.
1: K + exp(-D/ec
2: u 1n. V + 1m
3: while & Emax do
4: u (VV).
5: v b/(Kuu)
6: f f++++
7: N + diag(u) (
where W(a, 3) R the Wasserstein discrepancy between a and 3; E IR the unit-wise distance between between a and 3; a and b indicate the mass of units in a and 3, and II is the feasible transportation plan set which ensures the mass-preserving constraint holds.
B.2 SINKHORN DISCREPANCY AND ALGORITHM
Exact solutions to the Kantorovich problem suffer from great computational costs. The interior-point and network-simpe simplex methods, for example, have complexity of O(n logn) (Pele & Werman 2009). A shortcut is to add an entropic regularizer as
(a 3) (D, π := arg min (D, eHI(T), H(T) - rell(a.))
Txi (log(TT)) - 1), (38)
which makes the problem e-convex and solvable with the Sinkhorn algorithm (Cuturi 2013), with lower complexity of Onn?/e"). Besides, the Sinkhorn algorithm consists of matrix-vecto products only, which makes it suited to be accelerated with GPUs. Specifically, let IR" and g I be lagrangian multipliers, the Lagrangian of (38) is:
$ n f. (D, R) -eH(m))) (f, (a) -gg Imm b)
(39)
According to the first-order condition of constraint optimization problem, we have:
( f.g) Diy + log (T.y)) g, 0,
(40)
equivalently, the best transport matrix ग' should satisfy:
f
Dig
g
iH exp
exp
cxp
(41)
Let u, := exp(f./c), v, exp(g)/c). Kij (pp(-Di//)) then we have x' (ddaag) The transport matrix should also satisfy the mass muss-preserving constraint, such that:
diag(u) dlia(vv]. a, diag(v) diag(u)1, = b,
(42)
or equivalently, let @ be the entry-wise multiplication vectors, we have:
u n (Kv)= and v (KT b.
(43)
known as the matrix scaling problem. An intuitive approach to solve them iteratively:
(e+1)=
a (v(()
and (++1)=
(44)
(Tu((+1)
which is the critical step of Sinkhorn algorithm in Algorithm I The optimal transport matrix n' acting constant matrix further induces the Sinkhorn discrepancy We following (38). As D is differentiable to a and B, is feasible to minimize W by adjusting the generation process of a and B, ie., the representation mapping in Definition A.2 with gradient-based optimizers.

---

Algorithm 2 Generalized Sinkhorn Algorithm for Unbalanced Optimal Transport
Input: discrete measures Σx=1 a,,x, and 3 = 2==1 b, distance matrix Dij ||x; -x-x Parameter: e: strength of entropic regularizer; k: strength of mass preserving; maxi max iterations.
X2
Output: x = the entropic regularized unbalanced optimal transport matrix.
1: K + exp(-D/e).
2: On g- 88-0m Om l -
3: while l Emax do
4:
u - exp(f:/c). V exp(8y/c)
5:
diag(u) diag(v).
6:
a' Tln, b' + - 1m-
7:
if 2/2/22 0 then
8:
f -
+ +log(a) log (a')) en
9:
else
10:
- [ + log(bb - log (b')] ex
11:
( f+1.
12:
T. diag( k) diag(v).
B.3 UNBALANCED OPTIMAL TRANSPORT AND GENERALIZED SINKHORN
We have reported the mini-batch sampling effect (MSE) issue of W in Section 3.2 and attributed it to the mass-preserving constraint in (38) An intuitive approach mitigate MSE relax the marginal constraint and allow for the creation and destruction of mass. To this end, RMPR is proposed in Definition
B.3 which replaces the hard marginal constraint with soft penalty.
Definition B.3. For empirical distributions and 3 with and m units, respectively, unbalanced optimal transport seeks a transport plan at minimum cost:
37pE.M (a. 3) == min (D, T) π == argmin (D, x) cH(x) + ( KL( 11 b)), (45)
where D 1B is the unit-wise distance, and and b indicate the mass of units in and B.
The unbalanced optimal transport problem in Definition B.3 has similar structure with (38) and thus can be solved with generalized Sinkhorn algorithm (Chizat et al. 2018). The derivation starts from Fenchel-Legendre dual form of (45):
f, gj D
F" (-f) G (-g) Σ exp
max fc R Rm
1.
e
F (f) =e
maxz RKL a) 18 (e (f/H,a))
a) a aa1
(46)
G (g) maxz g xKL(z||b) x: (cRin Badl",b, .b) 1bbimm ze Rc
where the functions F' (.) and GO are the Legendre transformation of KL divergence, Ignoring the constant terms, we can obtain the equivalent optimization problem:
f. + g;
Dij -Dis x x f/x a) + (< (e-8/1 8/x b).
min fel geRm
Σ 1J
exp
(47)
According the first-order condition, the minimizer's gradient of 47 should be zero. As such, fixing g". the updated pê-1 ought to satisfy:
8
D
( p.+1
exp
Σ
exp
exp
a;
(48)
We further multiply both sides by exp (f:
p
f
exp
a
=
exp
exp
a;
(49)
'In this work, we calculate the unit-wise distance with the squared Fuclidean metric following Courty et al.
2017b

---

Algorithm 3 ESCFR Algorithm
Input: covariates of treated units {x:)_-1 and untreated units (=) factual outcomes (1i)21 and {5>) representation mapping vi; outcome mapping ø.
Parameter: X: strength of optimal transport; strength of entropic regularizer; k: strength of RMPR; strength of PFOR; mmax: max iterations
LCR.T, Output: ESCER' the learning objective of ESCFR.
1: (riM-1 + {(((x)))=1
{ m + {1{ f-1
2: (Dit
{851 {Ø(rz, 0) m
j-1'
3: (Tili- + (,(r,,
(g 1m {Ø(T,, 1) j-1'
4: D + |x x,113 + + |1
-51l2 + 18 d 1.
5: D stop stopgradient
6: π ENNit Algorithm2
a= (r: M- 3= {ry} D D /oopp)
7: LF(D, Ø) 1 -
Vill Σz-1 1iù 3
n
8: T (1)) <- (D7, ).
9: ESCER
C.N.N.A + LF(1, 1) + +
C (+).
Table 3: Running time (mean+std) in seconds of Algorithm
1-2 with 100 runs.
Parameter
n= 32
n- 64
n 128
n 256
n = 512
n= 1024
Algorithml
0.0266=0.0102 0.005020.0004
0.0241#0.0075
0.0326=0.0068
0.0499+0.0099 0.010410.0005
0.0725=0.0128
0.1430=0.0259
Algorithm2
0.005110.0001
0.0065_0.0002
0.013820.0008
0.0256.0.000
Parameter == 0.1
<=
= 1.0
= 10.0
c= 100.0
Algorithml
0.1683=0.0038
0.120200.010
0.0699=0.0095
0.0153+0.0013
0.0097=0.0009
0.0072+0.00m
Algorithm2
0.0166=0.0019
0.006800.0010
0.0052=0.0011
0.0047+0.0010
0.0045=0.0011
0.0043+0.0009
Parameter
0.1
0.5
1.0
5.0
10.0
x= 100.0
Algorithm2
0.005010.0011
0.005910.0008
0.0060_0.0011
0.011210.0014
0.016220.0016
0.103910.0033
where a' == m1, with Tıj
exp(f. +g]- - Diy). Similarly, fixing f we have g+1 as:
gj f+1
gi
exp
b', exp
exp
by
(50)
where b' := m. 49 and (50) construct the critical iteration steps of the generalized Sinkhorn algorithm (Chizat et al. 2018), which we formulate in Algorithm 2 The transport matrix -. further induces the generalized Sinkhorn discrepancy we.. in Definition B.3 As D is differentiable with respect to a and (8, it is feasible to minimize Wr. by adjusting the generation process of a and B, ie., the representation mapping in Definition with gradient- optimizers.
B.4 OPTIMIZATION OF ENTIRE SPACE COUNTERFACTUAL REGRESSION
Algorithm 3 shows how to calculate the learning objective mini-batch level. Specifically, we first calculate the factual outcome estimates (step 2), counterfactual outcome estimates (step 3), and the unit-wise distance matrix with PFOR (step 4). Afterwards, fix the gradient of the distance matrix (step 5) and calculate the transport matrix with (egorithm@(ster 6). Finally, calculate the factual outcome estimation error (step 7) and distribution discrepancy (step 8), and aggregate them to acquire the learning objective of ESCFR (step 9). According Section B.3 the learning objective is differentiable to ve and é and thus can be optimized end-to-end with stochastic gradient methods.
B.5 COMPLEXITY ANALYSIS
One primary concern would be the overall complexity of solving discrete optimal transport problems. Exact algorithms, e.g., the interior-point method and network-simplex method, suffer from a high computational cost of O(n" log n) (Pele & Werman 2009). An entropic regularizer is thus introduced in (5, making the problem solvable by the Sinkhorn algorithm (Cuturi 2013) in Algorithm 1. The

---

complexity was shown to be O /2n) bv Altschuler al. 2017 in terms of the absolute error of the mass-press mass-preservation constraints. Dvurechensky et al. (2018) improved it to O /e). which can be
further accelerated with greedy algorithm by Lin al. (2019). Several recent explorations (Blanchet
et al. 2018 Jambulapati et al. 2019) have also attempted to further reduce the complexity to (nn*())
Entropic regularization trick is still applicable to speed up the solution of the unbalanced optimal transport problem in RMPR, represented by the Sinkhorn- algorithm in Algorithm 2 Pham et al.
2020 further proved that the complexity of Algorithm 2 is (n)
Table reports the practical running time the commonly-used batch settings. In general, the computational cost of optimal transmission is not a concern at the mini-batch level. Notice that enlarging e speeds up the computation while making the resulting transfer matrix biased, hindering the transportation performance, as per Figure 5 In addition, a large relaxation parameter k makes the computed results closer to those by Sinkhorn algorithm yet significantly contributes more iterations, which is discussed and mitigated by Séjouré et al. (2022).
C REPRODUCTION DETAILS
C.I DATASETS
We conduct experiments on two semi-synthetic benchmarks to validate our models. For the IHDF benchmark, we report the results over 10 simulation realizations following Liuyi et al. (0018) 8) However, the limited size (747 observations and covariates) makes the results highly volatile. As such, we mainly validate the models on the ACIC benchmark, which was released by the ACIC-2016 competition Since the scale of the ACIC benchmark (4802 observations and 58 covariates) is much larger than the IHDP benchmark, we mainly perform ablation studies on it to get more reliable results.
All datasets are randomly shuffled and partitioned a 0.7:0. 15:0.15:.15 15 ratio for training, validation, and test, where maintain the same ratio of treated units in all three splits to avoid numerical unreliability in the validation and test phases. We find that these datasets are overly easy to fit by the model because they are semi-synthetic. To increase the distinguishability of the results, we omit preprocessing strategies, such as min-max scaling, to increase the difficulty of the learning task.
C.2 BASELINES
We compare the proposed method with baselines based statistical estimators (Künzel al. 2019 Shalit et al. 2017), matching estimators (Rosenbaum & Rubin 1983b; Wager & Athey/ 20118] Crump al. 2008 and representation--aasee estimators (Johansson et al. 2016 Shalit et al. (017). We implement all these baselines by hand, largely built upon Pytorch for neural network models, Sklearn for statistical models, and EconML for tree and forest models. These implementations will be open-sourced.
C.3 METRICS
Following existing works Yao al., 2019 Liuyi et al.) 2018), the Precision in Estimation of Heterogeneous Effect (PEHE) is primarily used as the precision metric for performance evaluation.
(PEHE(.. (w. x) Sx (())
(51)
However, PEНE is unavailable during the model selection phase as the counterfactual outcomes are missing. shorteut adapted by Liuyi et al. (2018) for model selection would be the root mean squared error of factual outcome estimates (RMSEp). which can be evaluated in the absence of the counterfactual outcomes. However, owing to the treatment selection bias, RMSEp is not reliable as it does not consider the precision of counterfactual estimation, and thus cannot effectively evaluate the quality of treatment estimation.
'It can be downloaded from https:/wwww.fredj.comm/
"It can be downloaded from http:///wwiiiilliiiiiiiiccccoaaa

---

Out-Sample PEHE
Out-sample PEHE
Out-sample PEHE
4.0
4.0
4.0
3.5
3.5
3.5
3.0
3.0
3.0
2.5
2.5
2.5
2.0
3.0
0.675
0.700
0.85
0.90
0.95
RMSEF
Out-sample AUUC
In-sample AUUC
(a)
(b)
(c)
Figure 7: Fidelity of model-selection criterion to out-of-sample PEHE
Area Under the Uplift Curve (AUUC) (Betlei et al. 2021 evaluates the model's performance in ranking units based on potential treatment benefit. more feasible than PEHE because it can be calculated without counterfactual outcomes; more reliable than factual MSE because it partially reflects the models' counterfactual ranking ability. Therefore, it has been the primary selection standard by practitioners from Criteo (Betlei et al., 202I). Alibaba (Ke et al., 2021). and Tencent (He et al., 2022). where engineers always fine-tune their causal inference models concerning AUUC performance and decide whether deploy the models online traffic. such, we use AUUC the model selection criteria. Moreover, we report the within-sample results on the training dataset and the out-of-sample results on the test dataset, where the factual outcome is available the within-sample case following Shalit al. (2017).
compare the fidelity of RMSEp and AUUC on evaluation set to the out-of-sample PEHE on the test set in Figure 7 The first observation is the weak correlation between RMSEpp on the validation set and the out-of-sample PEHE on the test set. In contrast, the sut-ol-samplee AUUC on the evaluation set could better reflect the variation of PEHE. Another observation is that the within-sample AUUC is not good criterion for model selection, as higher within-sample AUUC corresponds to higher out-of-sample PEHE, which contrasts with reality. This reasonable as better within-sample performance does not necessarily correspond to better out-of-sample performance.
D ADDITIONAL DISCUSSIONS
D.1 ADDITIONAL DISCUSSION FOR STOCHASTIC OPTIMAL TRANSPORT
According to Theorem 3.1 one critical hyperparameter for CFR-WE WASS and ESCFR is the batch size, which directly affects the variance stochastic optimal transport in Section 3.1 3.I and thus performance of both methods. As such, it is necessary to verify whether ESCFR outperforms CFR for different batch sizes. We conduct extensive experiments and summarize the results in Table Interesting observations are noted:
Increasing batch size in wide range improves the performance of CFR-WASS and ESCFR. For example, The PEHE CFR-WASS decreases from 3. at b 32 to 2.932 == 128, and the PEHE of ESCER exhibits similar pattern. The performance gain is attributed to the decreased variance in (6. which backs up Theorem 3.1
By finetuning batch size, we can easily exceed the performance we report Table I However, did not finetune it as the PEHE is invisible during hyper-parameter tuning p
The performance drop for the overly large batch sizes comes from the sub-optimal backbone (TARNet) performance. Due to the limited number of training samples, e.g., 4.8k * 70% units for ACIC and 0.7k 70% units for IHDP, a large batch size might block the models from escaping saddle points (Jin al., 2017) and sharp minima (Xie et al. 2020), and thus deteriorating the accuracy of factual outcome estimation.
Most of the experiments in Table were performed with fixed batch size b = 32, which is selected by the factual estimation performance of TARNet.

---

Table 4: Out-of-sample PEHE of ESCFR and important baselines with different batch sizes
Model
=
b=
96
b= 128
196
b= 256
TARNet
3.3293=0.1853
3.2054
2.9262:0.310 3.0869+0.2812
3.4619+0.6652
3.6309:0.2026
CPR-WASS
3.1143-0,4578
3.0819_0.3407
1.911810.1017
2.932610:4122
4.07401 144127
3.4675.0.1552
ESCER
2.3736_0.1621
2.3082:0.4324
2.971910.2889
2.312540.1836
2.0373a0.1538
2.2777=0.4230
4.0
k=1.0
K=1.0
3.5
K=3.0
3.5
K=3.0
5.0
k=5.0
1=00.
1=10.0
PEHE
3.0
K
3.0
2.5
2.5
2.0
2.0
1.5
1.5
64
Batch 96
128
64
Batch 96
128
(a) Out-of-sample Performance
(b) In-sample Performance
Figure 8: PEHE of ESCFR and CFR-WASS (n = oo) under different batch size.
D.2 ADDITIONAL DISCUSSION FOR RELAXED MASS- AASS-PRESERIING REGULARIZER
Existing methods (Salit al., 2017 Johansson et al. 2016 Liuyi al. 2018 suffer from the mini-batch sampling effect (MSE) issue, as indicated by the two bad cases in Figure 2 RMPR mitigates the MSE issue by relaxing the mass pass-preserving constraint, the performance of which is affected by two critical hyperparameters, i.e., the batch size b and the strength of mass mass-presservin constraint On top of the ablation studies, it is necessary explore the performance of ESCFR different settings of and k, to investigate 1) how RMPR works; 2) the limitation and bottleneck of RMPR; 3) the robustness of RMPR to hyperparameter setting. The results are presented in Figure X and the observations are summarized as follows.
The optimal value of и increases with the increase of batch size. For example, the optimal is 1.0 at 6 32, and 5.0 at 6 128. This observation partially verifies how RMPR works as described in Section 3.2 Specifically, small batch sizes where sampling outliers dominate the sampled batches, a small y effectively relaxes the mass mass-preserving constraint and avoids the damage of mini-batch outliers, thus improving the performance effectively and robustly. At large batch sizes, the noise of sampling outliers is reduced, and it is reasonable to increase K to match more units and obtain more accurate wasserstein distance estimates.
Even with large batch sizes, oversized K, e.g., 2 10 does not perform well. Although the effect of sampling outliers is reduced, some patterns such as outcome imbalance are present for all batch sizes, resulting in false matching given large mass-preserving constraint strength which might be primary bottleneck of RMPR.
Hyper-parameter tuning is not necessarily the reason why ESCFR works well, since all ESCFR implementations outperform the strongest baseline CFR-WASS ( < oo) all batch sizes, most of which are statistically significant. This can be further supported by our extensive ablation study in Section 4.3 and parameter study in Section 4.5
In summary, is necessary to relax the mass-preservine constraint under all settings of batch size, which strongly verifies the effectiveness of RMPR in Section 3.2
D.3 ADDITIONAL DISCUSSION FOR PROXIMAL FACTUAL OUTCOME REGULARIZER
Existing representation based methods block the backdoor path X + T by balancing the distribution of the observed covariates in latent space. Given the unconfoundedness assumption A.1 this

---

P1st
X
x
pT
In(#)
2sin(x)
R₂
(a) Toy example
(b) Good cases
(c) Bad cases
Figure 9: A diagram showing how PFOR works and its limitations. (a) A toy example of PFOR, where and X' indicate the balanced representations and an unobserved confounder, respectively; scatters indicate the empirical distribution of units in the treated and control groups; for solid scatters with balanced R. the colored dashed line indicates the ground truth outcome y /R] R3 X in each group, the black dashed line measures the difference of unobserved X (b) Cases that satisfy Assumption D.I. where the the outcome is monotone with unobserved X' given observed confounders in R. (c) Cases that violate Assumption D.I where the Y is non-monotone with X".
approach effectively handles the treatment selection bias. However, Assumption A. is usually violated in practice, which invalidates this approach as the backdoor path from the unobserved confounder X' to T is not blocked.
According to the designed causal graph in Figure ((b) all factors associated with outcome Y include the observed confounders X, treatment T, and unobserved confounders X". Therefore, is reasonable to derive that given balanced X and identical T, the only variable reflecting the variation of X' is the outcome Y. As such, inspired by the joint distribution transport technique (see Courty et al. 2017a), PFOR calibrates the unit-wise distance D with the potential outcomes in (12). The underlying regularization is: units with similar (observed and unobserved) confounders should have similar potential outcomes. Equivalently, for pair of units with similar observed covariates, i.e., [r; -r-ra 2 x 0, if their potential outcomes under the same treatment t = {0, 1} differ significantly, i.e., lo - >> O, their unobserved confounders should also differ significantly. As such, is reasonable to utilize the difference outcomes to calibrate the unobserved confounding effect.
a in the population of interest, let T = and X' x' be the treatment assignment and unobserved confounders, respectively, have E[ X x, X' x', T t] is monotonically increasing or decreasing with respect to x'.
Assumption D.1. (Monotonicity). For all observed covariates
Ad Advantages.
The advantages of PFOR can further interpreted follows.
From statistical perspective, PFOR encourages units with similar outcomes to share similar representations. It is valid prior that inspires many learning algorithms, e.g., e. K- nearest neighbors and gaussian process (see Williams & Rasmussen 2006). As effective statistical regularizer, PFOR also works in the absence of unobserved confounders, especially small data sets.
From domain adaptation perspective, vanilla Sinkhorn aligns the distributions PT- and pT-0 (r), where is the learned representations in Definition A.4 PFOR further aligns the transition probabilities PT- (Y = +) r) and PT t) r) for t 0. 1. The discrepancy between transition probabilities can attributed unobserved confounders that can be viewed as parameters of the transition probabilities (Courty et al. (2017a). As such, it is feasible to align the unobserved confounders by aligning the transition probabilities.
Toy example. Let the ground truth Y := / R X where T is omitted one group. R, and R2 are the representations of observed confounders that have been aligned with Sinkhorn algorithm. Let the unobserved X' 0 for controlled units and X" for treated units,
only consider

---

Table 5: RMSE of ESCER and its competitors for factual and counterfactual outcome estimation.
Dataset
Outcome
TARNet
CFR- WASS
ESCFR
Train
Factual
2.2066a0.134
2.1744=0.2179
2.1608
Train
Counterfactual
3,4731 1493
3,4306= 0904
2.7025
Test
Factual
3.1924=0.0357
3.1887=0. 1493
2.6334+0.2255
Counterfactual
3.5577=0.0957
3.4163=0.0499
2.7976:0.1376
which makes X' an unobserved confounder as it is related to and different between groups. As shown in Figure 9 a). given balanced R1 and R2, the variation of Y reveals that of X'. As such, it is reasonable to employ Y calibrate the unit-wise distance that ignores X"
Synthetic labels. PFOR remains effective for semi-synthetic data, where outcomes are synthetic from the covariates and treatment assignments. One source hidden confounders in such data the information loss from the raw data space to the representation space, where not all valuable information (e.g.. confounders) is extracted and preserved. Besides, this improvement could also come from the statistical regularization, encouraging units with similar outcomes share similar representations, which is an effective prior according to the K-nearest neighhoring methods.
Limitations. PFOR fails for confounders adding constant effects all units. Specifically, for unobserved confounder X' and treatment assignment t 0, if E[Y X' =1, T t] E[y X,X' T2, = f]. PFOR fails eliminate the confounding effect of X'. Examples can be found in Figure (c). Nevertheless, in real scenarios, it is rare that different values of X' only add constant effect to the outcome (see Sooer et al., 2016, Zheng et al.) 2021 Ogburn & HanderWeels 2012). making PFOR still effective in wide range of application scenarios.
This limitation formalized as Assumption D.I where outcome should be monotonically increasing or decreasing with unobserved confounders given observed confounders and treatment assignment, as shown Figure (b). Notably, it is commonly used assumption in confounder analysis (Sofer et al. 2016 Zheng et al. 2021). Besides, this assumption is often plausible, least approximately, conditional on T'= = (Zheng et al. 202I) For example, it naturally holds for binary confounders; and generally holds in applications such epidemiology (Ogburn & VanderWeed 2012). Finally, this assumption is only imposed on the hidden confounder FollowingZZann et al. (2021), which further weakens Assumption D.I significantly.
D.4 ADDITIONAL DISCUSSION FOR PERFORMANCE IMPROVEMENT
In practice, ESCFR serves as efficient regularizer on TARNet. To further investigate its regulariza- tion effects, the RMSE for estimating factual and counterfactual outcomes on the training and test sets is reported in Table 5]
We find the effectiveness of ESCFR comes from the three sources:
Better fitting the training set, where ESCFR achieves the minimum RMSE for factual outcomes the training set. This improvement is attributed to RMPR, which mitigates the mini-batch sampling effect, thus reducing mismatching.
Effective statistical regularization, where ESCFR reduces the gap between the RMSE for factual outcomes the training set and that the test set by 52. % compared with TARNet. It is mainly attributed to PFOR, which encourages samples with similar outcomes to share similar representations.
Effective counterfactual regularization, where ESCFR reduces the gap between the RMSE on the test set for factual outcomes and that for counterfactual outcomes by 55.1% % based on TARNet; the gap between the RMSE on the training set for factual outcomes and that for counterfactual outcomes by 57.3%. makes ESCFR robust to the treatment selection bias.
D.5 ADDITIONAL COMPARISON RESULTS
In addition to the competitors in Table I add two cutting-edge representation- based ITE esti- mators to bring our scope baselines up to date. Instead of designing better alignment methods, the additional baselines introduce variable decomposition and employ the distribution alignment

---

Table 6: Additional comparison (mean=std) the ACIC benchmark.
Out-sample
In-sample
Model
PEHE
AUUC
PEHE
AUUC
TARNet
3.2054=0.
0.6683+0.0139
3,2933#0 4076
0.8597+0.0350
CFR-WASS
3.0819=0.3407
0.6693+0.0250
3.4280+0.288 2857
0.8439#0.0349
DRCER
3.0199=0.0658
8.6360=0.0088 0.000
3.4957+0.0895
0.8456+0.00773
DRCFR-MIM
2.9593=0. 1532
0.6456=0.0088
3.3875+0. 2208
0.8434+0.0093
ESCER
2.3082=0.4334
0.7311#0.0123
2.5385+0.5022
0.8416#0.0367
technology (CFR) as a component. Therefore, it would be unfair for ESCFR to compare with them, as ESCFR not equipped with the variable decomposition technology. As a result, these results are not included in Table
Baselines. The first additional baseline is DRCFR (Hassanpour & Greiner 2020). It first decom- poses the latent space into instrument variables, confounding variables, and adjustment variables, and then MMD distance align the distribution of the adjustment variables. The second additional baseline is the MIM-DRCFR (Cheng et al.) 2022). As the SOTA approach as recent as 2022, augments DRCFR with three orthogonal constraints and uses the Wasserstein distance to align the distribution of adjustment variables. We highlight that neither of the two methods considers how to mitigate the mini-batch sampling effect and unobserved confounder effect, the two key problems in representation-bast based causal inference. Therefore, their contributions are quite different from our study.
Settings. Since the full implementations of both baselines are not released, we reproduce them largely from scratch. Nevertheless, DRCFR released its model structure based on tensorflow 1.x, which gives us an important guideline. Experiments are conducted on the ACIC benchmark. In terms of parameters shared with CFR, we set them to the values of CFR and ESCFR for fairness. These shared parameters mainly include the learning rate, weight decay, batch size. strength of distribution alignment, etc. Notably, set the batch size to 64 instead of 32 in Table to highlight the advantages of these representation-batt bottom and Factor-specific Disentangled representation module consists of shared fully connected layer (60 units) cascaded by three parallel fully connected layers (30 units per layer). The outcome prediction modules consist of two fully connected layesr (60 units per layer) following TARNet, CFR, and ESCFR. Empirically, it is good setup with competitive performance and similar number of learnable parameters as ESCER, ensuring fair comparison.
baselines. In terms of model structure, the Shared-
Results. As shown in Table 6 the additional baselines show a more stable performance (less vari- ance) and do outperform CFR- WASS in terms of oft-of-sample PEHE. However, ESCFR outperforms the additional baselines significantly terms of most metrics. The only exception, in-sample AUUC, has been analyzed in Section
D.6 ADDITIONAL RESULTS FOR ABLATION STUDY
In this section, provide additional results on hyperparameter study as it reflects how different components of ESCER affect the performance. Specifically, we report the ranking performance under different settings in Figure 10 and Table 7 The out-of-sample AUUC exhibits converse variation patterns with PEHE, which is consistent with our expectations. We add the estimation error of average treatment effect estimation for the entire population (EATE) and treated population (EATT) in Figure II Figure 12 and Table g where the PFOR significantly reduces both average estimation errors. Finally, we report the root mean squared error of factual outcomes (RMSEp) and counterfactual outcomes (RMSEF). Overall, all components of ESCFR could reduce RMSE significantly. One interesting observation that the dropping of RMSF always comes with the dropping of RMSEp. It shows that distribution adjustment does not sacrifice the performance of the factual outcome estimators, but provides prompting information that improves both factual and counterfactual outcome estimators.
10
We take 40% confidence interval plot bars highlight trends. Exact values given in tables.

---

Table 7: Individual treatment effect estimation performance (mean=std) of ESCFR under different
settings.
Metric Model
PEНE
AUUC
In-sample
Out-sample
In-sample
Out-sample
A 0.0
3.2367=0.266
3.2542=0.1500
0.8862= 0.062
0.6624=0.0149
= 0.5
2.8069=0.3913
2.9891=0.4981
0.8309=0.029)
0.696520.0166
À 1.0
2.7324+0.2716
2.9053=0.3744
0,8312+ 0.022
0.7022=0.0223
A= 1.5
1.6522=02600
2.8434= 0.5605
0.8222 0.0385
0.7114=0.0187
x=
2.7580=0.300 3066
2.8741 0,4639
0.8427=0.0412
0.6978=0.0181
A= 3.0
2.6771
2.8807=0.4148
0.8305
0.7088=0.0201
0.5
2.6855+0.2577
2.8640=0.4057
0.8285=0.029)
0.7039=0.0199
2.6842=0.3529
2.9532=0.54444
0.8383=0.0223
0.7042=0.0248
== 1.5
2.7913=0.23%
3.0501
0.854330.0358
0.7042:0.0192
2=0
2.7324
2.9053=0.3744
0,8312= 0.0262
0.7022=0.0225
2.5
2.848610 O. 2229
0.3463
08521=0.0263
0.6945.0.0172
=.0
2.7025
2.9353=0.4180
0,8417
0.6946=0.0162
2.6695#0.2958
2.8612+0.5707
0.8395=0.0315
0.7096:0.0193
==
2.6598=0.266 2627
2.7782=003791
18396=0.0299
0.7085=0.0155
2.6024 25226
2.7582=0.5217
0.827=00.0219
0.7065=0.0196
K 20
2.6140=0.255 2399
2.7201=0.4249
08362=0.02777
0.7059=0.0191
25
2.8598
19556=0.3422)
0.845100.0318
0.6990=0.0167
7 0.0
2.7324
2.9053=0.3744
0.8312 0.0262
0.7022:0.0223
7 0.1
2.6120#
0.5804
0.8272= 0.0392
0.7064=0.0223
7 0.5
0.6084
26779=0.4628
0.8267
0.7095=0.0275
7 1.0
2.6067+0.3012
2.7940=0.4347
0,8317#0.0501
0.0204
7= 5.0
2.6920=0.3115
2.7865=0.4249
0.8334
17.012000034
7 10.0
2.764320.3658
2.9006:0.5537
0.8494
0.7037=0.0269
Table 8: Average treatment estimation error (tean=std) of ESCFR under different settings.
Metric
ETT
Model
In-sample
Out-sample
In-sample
Out-sample
À 0.0
0.2930=0.31176
0.3681
0.6688
0.6880=0.3808
A= 0.5
0.3108=0.2974
0.8235=0.7135
0.4830= 0.4306
0.7677=0.6753
A= 1.0
0.377
0.8992:0.6360
0.5760= 0.4385
0.7802=0.6336
x= 1.5
0.5184:0.244 1
1.047=007144
0.7806=0 .,223
0.9899=0.6105
2.0
0.2898=0.2166
0.9060=0.6212
0.
0.7868=0.5812
3.0
0.408820.3408
1.1433=0.6801
0.5424 0.5269
1.0604=0.6303
0.5
0.2561+0.2307
0.8395=0.5324
0.3989
0.7391=0.4853
0.2617+0. 1962
1.024700.4543
0.3512=0.2772
0.9203=0145155
1.5
0,3353+0. 0.2555 0,2556
1.2782=0.5477
0.4339=0.3214
1.1656=0.5019
2=0
0.3779#0.299 2954
0.8992:2.6366
0.5760= 0.485
0.7602#0.6336
2.5
0.4934+0.288 2857
.3844=0.5700
0.4625 4451
0.2784=0.5561
=.0
0.2889+0.1677 1677
1.644200.5469
0.3765=0.2717
0.9154=0.5379
0.3598
1.789220.7782
0.3969= 0.06900.344
1.007:00.779
0.2964+0.255 2338
0.8721=0.6014
0.4106=0.2210
0.7972=0.5751
15
0.3797=0.2400 2.2402
0.9187=0.6839
0.5244
0.8463=0.6493
0.2799 0,2918
0.8306
0.3740= 0.4330
0.4300
0.750
== 25
0.2469#0. 1943
0.6571=0.5344
0.548 + 4.95
0.6031
7 = 0.0
0.3779#0.20 2954
0.8992:0.6366
0.5760= 0.4385
0.7802
7 0.1.1
0.3176#0.3216
0.7454+0.7342
0.5524
0.6759=0.6982
7 0.5
0.2170 1470
0.6312=0.4785
0.3740= 0.3230
0.5505=0.4044
7 1.0
0.2371+0.178
0.8427=0.6200
0.3036= 0.2446
0.7764=0.5583
7 7= 5.0
0.3903+0.3033
0.8365=0.6857
0.6347 0.5429
0.7601=0.6364
7 10.0
0.2977+0.1908
0.8416=0.4937
0.5226=0. 0.4002
0.7335=0.4856

---

Table 9: Outcome estimation percision (mean=std) of ESCFR under different settings.
Metric
RMSEp
RMSEcP
Model
In-sample
Out-sample
In-sample
Out-samplc
= 0.0
2.7785=0.3642
3.2952=0.2815
1.2367=0.2666
3.4277=0.3103
À 0.5
2.2107=0.4151
2.6970=0.3937
2.8069=0.391
2.9530=0.4813
A 1.0
2.1674=0.3397
2.6953=0.3274
2.7324 0.2
2.8730=0.2968
A 1.5
2.1288=0.392.
2.6346=0.360
2.6522= 0.2605
2.8538
=
2.3265
2.7251=0:4254
2.7580= 0.3065
2.9292=0.3954
À 3.0
2.218220.3444
2.294
2.6771
2.8553=0.2897
=
2,1481003396
2.6331=0.3469
16855=0.2577
2.8517=0.378.
= 1.0
2.1540=0.3650
2.6842=0.3800
2.6842 0.3529
2.8743=0.40K
e 1.5
2.386610.405
2.8419=0.36 0.3468
2.7913 0.2336
2.9439=0.3116
2
2.1674
2.6953=0.3274
2.7324=0.2716
2.8730=
2
2,4194 3648
2.8643 0.2888
2.8486 0.2229
2.9970=0.2555
=
2.
2.7804=0.4149
2.7025=0.3612
2.8609#0.3312
2,2200+0.506
=.4200
2.6695=0.299 29391
2.8626=0.3787
==
2.1754
1.648=002866
2.6598 0.2627
2.802500.328
=
2.0850+0.3475
0.5494=0.3441
2.6024
2.7893=0.2970
R = 20
2.0867 3,110
2.5790= 2.7977
2.6140 O 0,2392
2.7546=0.4400
K =
2,325 +5.3592
2.7818003316
2,8598 0.3253
3.0228
7 0.0
2.1674#0.33P9
0.6973
2.7324 O 0.2
2.8730 2989
7 0.1
2.0466=0.5003
0.4424=0.4407
2.8192=0.4226
7 0.5
2.0507=0.4153
2.5509=0.4264
2.6084=0.342 0.3425
2.7283=0.4048
7 1.0
2.088820.5184
2.5788 0.4708
2.6067=0.301
2.7722=0.3143
5.0
2.1111
2..112203557
2.6920=0.3115
2.7858=0.322
10.0
2.1696#0.4620
2.6797=004161
2.3643
2.8965=0.5432
0.9
0.85
AUUC
0.8-
In-sample Out sample
Ju
In-sample Out sample
AUUC
0.80
In-sample
AUUC
0.80
In-sample
Out-sample
Out sample
0.75
0.75
0.7
0.70
0.70
0.70
0.0
0.5
1.0
1.5
3.0
0.5
1.0
1.5
2.0
2.5
3.0
0.1
0.5
1.0
10.0
5.0
10.0
15.0
20.0 25.0
(a)
Figure 10: Parameter sensitivity study for critical hyper-parameters of ESCFR (AUUC).
In-sample Out sample
1.5
In-sample
1.00
Satilample
1.0
1.0
In-sample sample
0.75
In-sample
sample
0.5
0.25
0.0
0.5
1.0 1.5 2.
0.5
1.0
1.5
2.0
2.5
3.0
0.1
0.5
1.
1..0
5.0
10.0
15.0
20.0 25.0 (d)
(a)
(b)
(c)
Figure 11: Parameter sensitivity study critical hyper-parameterd ESCFR (EATE).

---

1.25
In-sample Out sample
In-sample Out sample
1.0-
In-sample Out sample
Out Out-samplc
0.8-
1.0
1.0
Ear7
0.75
0.6
0.50
0.5
0.5
0.2
2.5
3.0
15.0
25.0
y
K
(a)
(b)
(c)
(d)
Figure 12: Parameter sensitivity study for critical hyper-parameters of ESCFR (EATT).
D.7 ADDITIONAL DISCUSSION FOR THE BASIC ASSUMPTIONS IN CAUSALITY
This section reviews the fundamental assumptions in causality and illustrates their relationship to ESCFR.
The positivity assumption A.3 impiieee that the treated and untreated groups should contain overlapping units. The stochastic optimal transport in Section [.1 seeks to achieve in the latent representation space; however, the MSE issue leads outcome imbalance and sampled outliers at mini-batch level. In this case, pulling all samples into common overlapping region leads to incorrect matches and thus misleads the update of the representation mapping v. The RMPR adaptively matches and aligns the units that are close to the overlapping region by ignoring outliers, which reduces incorrect matches and prevents biased update of vb.
The PFOR is closely related to the unconfoundedness assumption A. On the basis of the vanilla Sinkhorn in Section 3.1 further aligns the transition probabilities 1TT (Y (I'
r) and (T- (T = 1) r). That is, for a pair of units from different treatment groups, if they share similar confounders r, their potential outcomes, i.e., Y Y' should also be similar. independent of the specific treatment assignment given confounders r obtained with the assistance of PFOR, i.e., (Y", Y²) и T which is exactly the unconfoundedness assumption A.I
In other words, the potential outcomes


Reviews:
[{"review_id": "VdqCXGg0K_k", "reviewer": "Reviewer_2Gnu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper addresses the problem of inferring individual treatment effects from observational data using representation learning via a discrepancy constraint between the representations of the treatment and control groups. In particular, the authors address the problem of dealing with the mini batch sampling which occurs during the deep learning process. ", "summary_of_the_review": "As I stated above, I think this paper overall contains a very interesting idea which serves as a real contribution to the literature. Unfortunately a combination of loose claims and writing make it difficult to engage more fully with the work as it presently stands. To be entirely honest, because of this it wasn't always clear _which_ claims are being made versus are the unfortunate artifact of loose writing. Additionally, I believe that the authors would be well served to clarify exactly what they are correcting for in terms of the _consequences_ of minibatches, not just the presence of minibatching itself. With that being said, I think there is a very nice idea here! I would strongly encourage the authors to make a substantial editing pass to improve the language and narrative of the paper. ", "strength_and_weaknesses": "Strengths: \n(1) This is an interesting insight regarding minibatches and discrepancy estimation. While I find it unlikely that this effect would render existing methods unusable, the increased variance is certainly an item of concern and the authors do a nice job of addressing it.\n(2) The solution is well reasoned and motivated and elegant in execution.\n(3) There are strong empirical results. \n\nWeaknesses:\n(1) The language in this paper is _very_ difficult to follow. In particular, it is exceedingly loose in some places to the point that it is almost misleading. Some examples:\n*  it is always too expensive to conduct randomized experiments. → I think you meant to say often instead of always here?\n\n* I’m not entirely sure what happened but a number of your citations cite either the first name or middle initial instead of last name (e.g. R, et al., Uri).\n\n* “Existing representation-based methods fail to eliminate the treatment selection bias due to the Unobserved Confounder Effects (UCE). “ Is this really what you meant? If so, can you provide specific theoretical evidence of this?\n\n(2) This may partially be an artifact of (1) but the claims in this paper appear to be not terribly well stated and substantiated. This is in contrast with the proposed method which I found to be quite elegant! In particular is the claim that minibatching renders existing discrepancy based learners inconsistent? At which batch sizes? Can we quantify this? The underlying claim that minibatching and the resulting variance is problematic I think is fine, but the authors need to be more careful in the overall language and claims. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "O5yM9EzDoN", "reviewer": "Reviewer_VvUp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this work the authors propose an optimal transport framework algorithm called Entire Space Counter-Factual Regression (ESCFR). The method is interesting in itself, with the additions the authors propose (specially the PFOR) add some light on how neural representations could be better used for counterfactual estimation", "summary_of_the_review": "Overall, this paper presents a nice idea of how to use neural representations for counterfactual estimation. I really enjoyed the PFOR idea and thats part of why I'm leaning towards an accept\n\nI think more empirical studies are merited - and Id like to see a proof of what is the minimal set of assumptions needed for PFOR to work. Maybe the authors understand this, but even after careful reading - I couldn't really figure it out\n\n---\n\nMany thanks to the authors for their detailed and insightful comments - I have read them and will keep my score as is. I'm still not sure how justified the procedure is without a concrete theoretical result.\n\n---\n\nUpdate post discussion: After extensive reading and discussions with the other reviewers, I share their concerns about clarity and novelty. In particular, it isnt really clear to me what is being added over the paper by Shalit et al (2017). Based on this, I will lower my score", "strength_and_weaknesses": "Strengths:\n\n- PFOR is a nice technique to mitigate the missing confounder problem\n- Performance is good on test sets\n\nWeaknesses:\n\n- I think PFOR is the most interesting part of this paper. A lot of counterfactual analysis comes down to the issue of missing confounders. It seems assuming a monotonic effect is all PFOR needs but I didn't see any theorem proving this", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "AaS3ct03tl", "reviewer": "Reviewer_gP25", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Treatment effect estimation is a challenging problem that requires generalising to counterfactuals and involves bias and variance due to confounding and lack of overlap between treatment groups. The authors build on the extensive literature studying this problem from the representation learning perspective to propose a new regularisation function with improvements for balancing distributions of treated and untreated populations and dealing with unobserved confounders.", "summary_of_the_review": "The proposed extensions, in my view, are reasonable heuristics to improve treatment effect estimation although the formalism and theoretical understanding behind them thin (which is understandable given that this problem is challenging). I do believe however that this is a good contribution to the literature.", "strength_and_weaknesses": "**Strengths**\n- This paper is well-written. All ideas follow intuitively and all quantities are well defined. \n- Treatment effect estimation in the presence of unobserved confounding is challenging (if not impossible without strong prior assumptions). I commend the authors for studying it. \n- The proposed approach substantially outperforms in the presented experiments.\n\n**Weaknesses**\n- Given unobserved confounders, there aren't any assumptions that discuss their influence on the system. For instance, in Sec. 3.3 the authors assume that all variation in outcomes that does not come from observed confounders or treatment, must be due to unobserved confounders. Could it not be due to variables with an independent causal effect on the outcome? I take it that the authors assume the graph in Fig. 2, but still this graph hides the contribution of exogenous variables so that it is not necessarily the case that two units with similar covariates and unobserved confounders have similar outcomes.\n- The datasets used in the experiments do not have unobserved confounders (UC), right? How are you evaluating the UC regularization term?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2T-gX2m2Eo", "reviewer": "Reviewer_ECPK", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the problem of individual causal estimation under treatment selection bias or missing confounders. The authors proposed a new estimator based on representation learning and optimal transport, which aims to minimize certain regularized discrepancy measure in the representation space. The robustness of the proposed estimator under outliers, and the adjusted regularization for missing confounders are analyzed theoretically. The authors further provided empirical study on the performance of the proposed estimator on two datasets.", "summary_of_the_review": "Overall the paper studies an important and open problem. The proposed solution is intuitive and natural, while the theoretical analysis and experimental comparisons can be improved.\n\n\n\n------author rebuttal acknowledgement------\nI thank the authors for the detailed response, revising the manuscripts, and answering the questions I raised during my review. I believe this work proposed very reasonable heuristics based on optimal transport and various regularization terms to address the minibatch sampling effect and unobserved confounders problem in causal effects estimation. I also believe that the presentation and empirical evaluations could be further improved to make this work stronger. I keep my original score after the author feedback.\n", "strength_and_weaknesses": "Strength:\n- Causal effect estimation under treatment selection bias or missing confounders is an important problem, with many downstream applications. Aligning the distribution of each treatment groups in a \"learnt\" representation space is a natural solution, and has been studied in the prior works. This work builds on works in this line, and provided further adjustments to mitigate the mini-batch sampling bias and unobserved confounder problems, which are all well-motivated and significant research questions to investigate in.\n- Using optimal transport appears to be an elegant solution for aligning two distributions. The proposed estimator using optimal transport exhibits simplicity while outperforming several other baselines.\n- The proposed estimator was compared with a wide selection of baselines in the prior works\n\nWeakness:\n- I think many places in writing for this work can be improved. For example, Figure 2 was actually referenced before Figure 1, and the Figures are in lack of explanations in the caption or text about how the readers should read them. Besides, there were several places in the theoretical analysis where the authors left sentences like \"see more rigorous analysis in the appendix\" w/t further explanations on what are the exact content to look for, and where to look for in the Appendix. While deferring non-urgent details to the appendix is completely fine, such a way of writing gives the reader a sense that the paper was written maybe in a rush or so. \n- It was not very clear what is the key adjustment which avoids the overfitting issue under mini-batches, and the main novelty over (Uri et al 2017) . As the author noted, (Uri et al 2017) minimizes PEHE (eq(3)), which may overfit to the respective group’s properties and thus cannot generalize well to the entire population. However, for the optimal transport solution, suppose that the treated/untreated groups are highly imbalanced, shouldn't that give bias issue to the optimal transport optimization loss too?\n- From table 1 the proposed method not only has lower out-of-sample loss, but also the lowest in-sample loss. It was unclear if the other baselines are tuned and trained to the minimal in-sample loss? Otherwise what is the intuition that the proposed estimator also performs the best in the in-sample regime?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject