Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Authors: Yuguang Yan
Zongyu Li
Haolin Yang
Zeqin Yang
Hao Zhou
Ruichu Cai
Zhifeng Hao
Venue: ICML 2025 poster (2025)
Keywords: Causal inference
continuous treatment
optimal transport
Primary area: general_machine_learning->causality

Abstract: Causal inference seeks to estimate the effect given a treatment such as a medicine or the dosage of a medication. To reduce the confounding bias caused by the non-randomized treatment assignment, most existing methods reduce the shift between subpopulations receiving different treatments. However, these methods split limited training samples into smaller groups, which cuts down the number of samples in each group, while precise distribution estimation and alignment highly rely on a sufficient number of training samples. In this paper, we propose a distribution alignment paradigm without data splitting, which can be naturally applied in the settings of binary and continuous treatments. To this end, we characterize the confounding bias by considering different probability measures of the same set including all the training samples, and exploit the optimal transport theory to analyze the confounding bias and outcome estimation error. Based on this, we propose to learn balanced representations by reducing the bias between the marginal distribution and the conditional distribution of a treatment. As a result, data reduction caused by splitting is avoided, and the outcome prediction model trained on one treatment group can be generalized to the entire population. The experiments on both binary and continuous treatment settings demonstrate the effectiveness of our method.

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Yuguang Yan
Zongyu Li!
Haolin Yang
Zeqin Yang'
Hao Zhou'
Ruichu Cai Zhifeng Hao
Abstract
1. Introduction
Causal inference aims to estimate the causal effects of treatments for supporting decision-making, where the treat- ments are usually binary (Shalit et al., 2017) or continuous (Schwab et al., 2020). The gold standard for estimating causal effects is to conduct randomized control trials (RCTs) (Fisher, I 1936), in which the assignment of treatment for sam- ples is completely random without relying on the covariates of samples. However, it is usually infeasible to conduct RCTs, and the effects are estimated from observational data involving confounding bias, which means that the data dis- tribution of subpopulation receiving one value of treatment differs from the distribution of the entire population (Ham- merton & Munafö, 2021), i.e., p(zlt) * p(z), where x is the covariates and is the treatment value.
Causal inference seeks to estimate the effect given treatment such as a medicine or the dosage of medication. To reduce the confounding bias caused by the non-randomized treatment assign- ment, most existing methods reduce the shift be- tween subpopulations receiving different treat- ments. However, these methods split limited train- ing samples into smaller groups, which cuts down the number of samples in each group, while pre- cise distribution estimation and alignment highly rely a sufficient number of training samples. In this paper, we propose distribution alignment paradigm without data splitting, which can be nat- urally applied in the settings of binary and contin- uous treatments. To this end, we characterize the confounding bias by considering different proba- bility measures of the same set including all the training samples, and exploit the optimal transport theory to analyze the confounding bias and out- come estimation error. Based this, we propose to learn balanced representations by reducing the bias between the marginal distribution and the conditional distribution of a treatment. As a re- sult, data reduction caused by splitting avoided, and the outcome prediction model trained on one treatment group can be generalized to the entire population. The experiments on both binary and continuous treatment settings demonstrate the ef- fectiveness of our method.
To address the confounding bias, most existing machine learning methods adopt a data-splitting strategy to split sam- ples into smaller subpopulations according to the treatment values, and then reduce the distribution shift between dif- ferent subpopulations. For binary treatments, one usually splits training samples to the treated group receiving treat- ment and the control group without receiving treatment, and then reduces the distribution shift between the two groups (Kuang et al., 2017; Shalit et al., 2017). For continuous treatments, the natural and widely used strategy is to split samples into multiple groups based on their received treat- ments. After that, the distribution shift reduction approach for binary treatments can be applied by considering the shift between each pair of groups (Wang et al., 2022). However, data splitting cuts down the number of samples in each sub- population, and only part of the samples are leveraged in distribution estimation and alignment. This decreases the performance of distribution estimation and confounding bias reduction, which highly relies on a sufficient number of training samples (Wang et al., 2022).
'School of Computer Science, Guangdong University of Technology, Guangzhou, China Pazhou Laboratory (Huangpu), Guangzhou, China 'College of Science, Shantou Univer- sity, Shantou, China. Correspondence to: Ruichu Cai <cairuichu gom..
In this paper, we propose a distribution alignment paradigm involving all the training samples without data splitting, which can be naturally applied to effect estimation of bi- nary and continuous treatments. Rather than reducing the distribution shift between subpopulations receiving differ- ent treatment values in existing methods, we characterize the distribution shift by different probability measures of
Proceedings of the 4rnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
the same set including all the samples. In other words, we model the conditional distribution p(arlt) by all the sam- ples, instead of only subpopulation receiving which is widely used in existing works (Shalit et al., 2017; Wang et al., 2022). By doing this, data splitring is avoided and all the samples can be leveraged to improve the performance of distribution alignment.
and advertising (Chen al., 2023; Liu al., 2021; Wei et al., 2021). In the last decades, various machine learning methods have been applied to address the problem of causal inference. Due to the confounding bias, the data distribution of a subpopulation receiving one value of treatment differs from the distribution of the entire population (Hammerton & Munafò, 2021). For example, in the treatment of dis- ease, the group receiving surgery usually has more severe conditions compared with the group receiving medication, the patients receiving higher doses of drugs usually have more severe conditions compared with the patients receiving lower doses, resulting in a distribution discrepancy between a subpopulation and the entire population.
Specifically, we establish the connection between the treat- ment effect estimation and optimal transport built on prob- ability measures involving all the samples (Villani, 2008; Peyré & Cuturi, 2017). We show that for the marginal co- variate distribution and the conditional covariate distribution given a treatment value, both the bias of covariates and the bias of outcome estimation errors can be upper bounded by the Wasserstein distances between these two distributions. Motivated by our theoretical results, we propose method named Optimal transport for Reducing blas in Causal in- ference (ORIC), which learns balanced representations to reduce the confounding bias and outcome estimation error jointly. As a result, the outcome prediction model trained on samples receiving one treatment value can be generalized to the entire population. Our theoretical results and algo- rithm can be naturally applied to both binary and continuous treatments. We conduct experiments synthetic and semi- synthetic datasets under the binary and continuous treatment settings, and the results demonstrate the effectiveness of our proposed method compared with existing methods.
Most existing works consider the binary and continuous treatment settings. The binary setting only considers whether the treatment is conducted not (Shalit et al., 2017; Shi et al., 2019; Zhang et al., 2020), and the continuous treatment setting considers the outcome of the dosage of the treatment to estimate the dose-response function (Schwab et al., 2020; Nie et al., 2021; Wang et al., 2022).
Binary Treatment. Causal effect estimation of binary treat- ments considers only two groups, i.e., the one receiving the treatment and the one not receiving the treatment (Chip- man et al., 2010; Dismuke & Lindrooth, 2006; Yoon et al., 2018; Zhang et al., 2020). To address the confounding bias between the two groups, one class of methods is to create a useudo-balanced group by lcarning weights for samples. Kuang et al. (2017) proposed to reweight samples by re- ducing the distribution discrepancy between the two groups, where the discrepancy is measured by the difference of the moments. The other class of methods is to learn balanced representations for the two groups (Johansson et al., 2016), 6). Shalit et al. (2017) and Johansson et al. (2022) proposed to learn representations with the minimized distribution dis- crepancy between two groups, where the discrepancy is measured by the integral probability metric, and a theoreti- cal analysis regarding the effect estimation error is provided.
The principal contributions are summarized as follows:
To address the confounding bias in causal inference, we propose to characterize the distribution shift by considering different probability measures of all the training samples without data splitting.
We construct the theoretical connection between the es- timation error of treatment outcomes and optimal trans- port, which measures the distribution shift between the marginal covariate distribution and the conditional covariate distribution given treatment value.
Our proposed learning model can be naturally applied in the binary treatment setting. Actually, distribution alignment between two groups splits training samples into two subsets, also cutting down the number of samples in each group. By modeling a distribution as a probability measure of all the samples, we avoid data splitting and obtain more samples for learning.
Motivated by our theoretical results, we propose a bal- anced representation learning algorithm to reduce con- founding bias and outcome estimation error jointly, and conduct experiments under different settings to demonstrate the effectiveness of the method.
Continuous Treatment. Causal effect estimation of con- tinuous treatments considers that the treatment lies in an interval, e.g., the dosage of a medication (Imbens, 2000). The natural strategy is to partition training samples into mul- tiple groups, each of which receives a similar dose of the treatment. By doing this, the existing methods for binary treatments can be applied. Schwab et al. (2020) adopted multi-head architecture to deal with multiple intervals of
2. Related Works
2.1. Causal Effect Estimation
Causal inference has been widely used in real-world appli- cations, such as economics (Davis & Heller, 2020; Kreif et al., 2021; Cockx et al., 2023), healthcare (Sanchez et al., 2022; Karboub & Tabaa, 2022; Van Goethem et al., 2021),
2

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
treatment separately. Wang et al. (2022) calculated the dis- crepancy between each pair of two groups and reduced the largest discrepancy to learn balanced representations. The strategy of data splitting cuts down the training samples in each group, highly affecting the performance of distribu- tion estimation and alignment. Different from them, we characterize the distribution discrepancy by different proba- bility measures of all the samples, avoiding data reduction in splitting.
usually considers only subpopulation, while our model represents distribution by involving all the training sam- ples and probability measure, improving the number of training samples for distribution estimation and alignment.
3. Problem Statement
We assume dataset of the form { t4, (et,,,t,)))) H=] where (x, t, y) is realization of random vector (X,T,Y Here 24 E X denotes the covariates of the i-th sample, t, E T is the treatment value that the sample i received which can be binary or continuous, and y; E ) denotes the outcome of interest for the sample i after receiving treatment t, Under the Neyman- Rubin potential outcome framework (Rubin, 1974; Rosenbaum & Rubin, 1983), the observed outcome Y is the potential outcome Y (t) corresponding to the actually received treatment T = t.
There are also few works of continuous treatments without data splitting. Nie et al. (2021) proposed a varying coeffi- cient model to estimate the effects of continuous treatment and apply a targeted regularization paradigm for estimation. Different with it, we explicitly reduce the confounding bias and theoretically reveal the connection between the con- founding bias and the generalization error of the outcome estimation, which are missing in (Nie et al., 2021). Kazemi & Ester (2024) measured the distribution discrepancy based on the Kullback-Leible (KL) divergence and employed an adversarial learning paradigm to learn the representations. However, the KI. divergence suffers from the issue of gra- dient vanish when the distribution discrepancy is too large (Arjovsky al., 2017), and the adversarial architecture is usually difficult to train (Gulrajani et al., 2017). Different from it, we measure the discrepancy by the Wasserstein distance to avoid the issue of gradient vanish, which can be casily estimated by the Sinkhorn algorithm (Cuturi, 2013).
Given the input covariates X = and the treatment 7 = t, our goal is to derive an estimator l0( 1) for the ground-truth individual response function μızz, t) as follow:
,(x, t) = EE (0) = x].
(1)
For simplicity, we will use the shorthand pz(x) = ,l t) and hx(x)) = h(x, l). The following assumptions have been made to ensure that (x)) is identifiable from observational data.
Assumption 3.1 (Stable Unit Treatment Value Assumption).
2.2. Optimal Transport
The potential outcomes for any sample do not vary with the treatments assigned to other samples, and for each sample, there are no different forms or versions of each treatment value which leads to different potential outcomes.
Optimal transport studies how to move mass from one dis- tribution to another with a minimal transport cost (Monge, 1781; Kantorovitch, 1958; Villani, 2008). Beneficial from the powerful ability to model probability distributions and exploit geometry, optimal transport has been widely ap- plied in many applications (Peyré& & Cuturi, 2017), such computer vision (Rubner et al., 2000), domain adaptation (Courty et al., 2014; 2017), data generation (Arjovsky et al., 2017; Tolstikhin et al., 201 8), graph data analysis (Peyré et al., 2016; Titouan et al., 2019), etc.
Assumption 3.2 (Ignorability). Conditional on covariates, the treatment assignment is independent of potential out- comes: T 11 X
Assumption 3.3 (Positivity). Conditional on covariates, the treatment assignment is not deterministic: 0 < p(TT-
x|x = x) < 1.
With these assumptions, pe(x)) can be rewritten as follows, and we can estimate as
Optimal transport has also been introduced into causal ef- feet estimation of binary treatments recently (Yan et al.,, 2024a;b; Wang et al., 2024). Li et al. (2021) proposed to transport the factual distribution to the counterfactual distri- bution for estimating counterfactual outcomes. Dunipace (2021) employed optimal transport to learn an intermediate distribution by reweighting samples. Shalit et al. (2017) and Wang et al. (2024) learned balanced representations between the control and treated groups by measuring the discrepancy via the Wasserstein distance. Different from the above studies that only consider binary treatments, we address the confounding bias in the setting of continuous treatments. Besides, in the above methods, a distribution
(x(x) = EY = x]] Y = xx
Without ambiguity, we omit the random variables to write p(X- = .)) as p(x) for simplicity.
4. Methodology
In this section, we first characterize the confounding bias by considering different probability measures of all the sam- ples, in which data will not be split into subpopulations. After that, we provide thcoretical results regarding the con- founding bias and the generalization error of the outcome
3

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
estimation from the perspective of optimal transport, which is built probability measures of all the samples. Based the theoretical analysis, we propose balanced representa- tion learning algorithm to reduce the confounding bias and outcome estimation error jointly.
We consider all the possible treatment values & E T, and define the total balancing error follows
E(m) - J 5(m, (p(())l
J 11 (0())(())dr S q(z))((z)d p(t)dt.
4.1. Confounding Bias in Causal Effect Estimation
Given the set of Radon measures M(X), ) let the marginal co- variate distribution be the probability measure E M())), and the conditional covariate distribution given a treatment value t T be the probability measure qr E M()). The corresponding probability density functions can be written as q(x:) p(x)), ge (a:) p(a](t) According Assump- tion 3.3, for each sample and treatment value f, we have qo (z) = p(zlt) = p( ((z)) pt)))(t)))/p( > 0, which means all the samples could be drawn from the distribution Moti- vated by this, we model 9r as probability measure involv- ing all the training samples, which is different from data splitting that only samples receiving t are involved (Shalit al., 2017; Wang et al., 2022).
We do not restrict the specific form of the function m(-) (-) as long as it can capture information from samples, enabling the balancing error E(+) to characterize the degree of con- founding bias.
In the following, we establish the connection between the treatment effect estimation and optimal transport, which motivates us to propose a balanced representation learn- ing algorithm for reducing confounding bias and outcome estimation error.
4.2. Theoretical Analysis
To analyze the confounding bias and outcome estimation error, we exploit the theory of optimal transport built on probability measures. Optimal transport aims to find the optimal plan to move mass from one distribution to another with a minimal transport cost (Villani, 2008; Peyré & Cuturi, 2017). Formally, for the samples from two spaces a E A, E B, let M(A) and M(B) be the sets of Radon measures. Consider two distributions a E M(A), 3 E M(B), and a distance function e A x B + IR + with the corresponding norm || |]. the Wasserstein distance between two distribu- tions W(c, a, B3) is defined by the following Kantorovich Problem
Specifically, for the treatment value I, based on the loss function f: Vxy + R we aim to minimize the following estimation error on the marginal distribution g(z)
(h,) = =Eq Ek(h), = (h.))
= 1 (((a((),,
(3)
and achieve small average outcome error considering all the possible values of treatment which is defined as
f Eq (((h)))(t)dt.
ع = EE
(4)
W(c, a, 3) KP 3)
inf { cla, (), b),
Nevertheless, given the observational data, we can only minimize the following factual error on the conditional dis- tribution (x)
=Ell(a,,)) AxE
(8)
where π is transport plan, and II(a, is the set of all joint probability couplings whose marginal distributions are a and 33, respectively. क (a, b) indicates how many masses are moved from to and the transport cost between them measured by the distance c(a, b). The minimized transport cost calculated by the optimal plan is the Wasserstein dis- tance to measure the discrepancy between two distributions.
Sq, (h+) = Sq. (he, Pt) = Exngi ((hx pr()))
= f -fx
(4
(5)
The principal challenge in causal effect estimation comes from the confounding bias, i.e,, q(x) ((x),, Vt T.
Given the pair of continuous functions (f,g) satisfying the constraint f(a) + g(b) S c(a, b), the above Kantorovich problem admits the following Dual Problem (Villani, 2021)
As result, the model trained to minimize Eg, cannot be well generalized to minimize Eq. To measure the level of confounding bias between (x) and q(x), given a function (e.g., balancing score) m(-) and a norm || ||. we define the balancing error between these two distributions as
DP(a, sup / f(a)da(a)+ + 9(b))3(())
1.g
s.t. f(a) + g(b) s c(a,b).
{(m, t) = [ [e)=a Ex-a(e)((z))
ll (=())m(z)dx - / m(d)m())
(6)
The following theorem shows that the confounding bias can be upper bounded by the Wasserstein distances between the

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
marginal covariate distribution and the conditional covariate distributions given value of treatment.
Eq. (4), propose to minimize the estimation error on the conditional distributions qe and the Wasserstein distances between q and d: simultaneously, as shown in the following
Theorem 4.1. Let q be the marginal covariate distribution, and qt be the conditional covariate distribution given the trealment value I, i.e,, q(z) = p(x)) and (x(x) = p(x](t. Given pair of the functions (m, e) satisfying the condition m(x;)= - -m(x;) S clzi, Tz). We have the following result
E E= { Egg((ii)tt) - /, Eg (h)(h()(t))tt
+ I W( ut DD((tdd.
(12)
5(m) & f [Wc,, 9t, d)p(t)dt.
(10)
which can be obtained from Eq. (11) immediately.
For the probability measures qi and q. a convenient property of optimal transport is that either continuous or discrete measures can be handled within the same framework, and the probabilities (z(z)) and g(z) can be easily represented as the sample weights for empirical distributions (Peyré & Cuturi, 2017). practice, given training samples {=:)==1> let 6z, be the Dirac function at the location I,, 9. (1,)) and ((x;) are the probability masses of the sample 2; in the distributions qi and g, respectively, which satisfy the simplex constraints
This theorem presents that the confounding bias character- ized by the balancing error can be upper bounded by the Wasserstein distances based on an underlying cost function .) and the probability measures di and g. where the cost function -):-) can be implemented by a distance measured on a representation space.
However, only focusing on confounding bias reduction may lead to trivial solution that loses outcome information, ie., mapping all samples to single point, which hampers the performance of outcome prediction. For the outcome estimation error, we can only train a prediction model ht on the training data to minimize Eg. (h₂) in Eq. (5), while the objective to minimize (he) in Eq. (3). The bias of the outcome estimation errors E (he) and (fz) character- ized by the following theorem.
n
(x,))) = 1,
Σ ((x))= =1
(13)
The corresponding empirical distributions de and can be represented as
Theorem 4.2. Assume that the cost function c(x, x²) [(((z)- - -)(2")llz where H. is Reproducing Kernel Hilbert Space (RKHS) induced by 6 X + H. Assume further that hi, Pt E F where F is a unit ball in the RKHS H. and the loss function ((he (2)), (x())) is convex, symmetric, bounded, obeys the triangular inequality and has the parametric form The (z) )xx for some x > O. Assume also that kernel in the RKHS H. is square-root integrable with respect to X and 0 S k(xxx²') (Ø(x), a(x!)) S K. Then the following holds
n i = Σ (())))o.
n q Σ ((,)),...
(14)
Here, all the training samples are involved in the empirical distributions, which avoids the issue of data splitting and enhances the performance of distribution estimation.
Based on this, the relation between the outcome estimation error and the Wasserstein distances measured on the em- pirical discrete distributions is provided in the following theorem.
f Eg(h))(t)dtt e (hep((t))
Theorem 4.3. Let n. be the number of samples, 9.9 be the empirical distributions of q, 9t. respectively. With the probability of least 1 8, we have:
< / M(0,0090000
(11)
This theorem shows that given an outcome prediction model he, the Wasserstein distances between the distributions q and ge provide an upper bound for the bias between the outcome estimation errors of h on q and qu- The theorem also indicates that it is not sufficient reduce W(c, qe- 9) only, since small W(c, qt, 9) cannot guarantee to obtain model h, with good performance. Even a model ht with poor prediction performance can perform similarly on and q, which happens when the information about the outcome is missing during distribution alignment. Therefore, in order to minimize& E that is the estimation error on q defined in
+ (a
E E Eg. (ho)((t)))
/ W(c, (p((t)t
+ (1/Von))
(15)
Note that in (Shalit et al., 2017), the discrepancy between the subpopulations of treated and control groups is reduced. Different from it, we reduce the discrepancy between the marginal distribution & and the conditional distribution i., both of which are modeled by all the samples equipped with the propensity scores.
5

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
4.3. Algorithm
p(],,), we approximate p((]z,) by ((4))) ((((;))) which is estimated by the generalized propensity score (GPS) (Imbens, 2000) based on the model 0(-). As result, (o(Zzg) is approximated by the normalized value de (2x) = 2 (((())) where Z = 1 (()())) is the normalized factor, so that the simplex constraint in Eq. (13) is satisfied. Please refer to Appendix for model details of the implementation of ((-).
According to the above theoretical analysis, we propose to minimize the outcome prediction on the observational distribution 4t and the Wasserstein distances between the & E TT The first part of the right side of Inequality (15) is defined as
empirical distributions & and / with
c = | En (ts))(t)dt
In practice, similar q: and & that only consider the empiri- cal discrete samples, we consider set 7 including discrete values of the treatment. For binary treatments, we have T = {0, 1). For continuous treatments, it brings a high computational cost to consider all the discrete treatments received by the samples. To alleviate this, we adopt some sampled values evenly distributed in T to construct the set T. It is worth mentioning that for each E T, all the samples are assigned by the weights á (2) and taken into consider- ation for distribution alignment, avoiding the issue of data splitting, Finally, we achieve the following optimization problem
S 1xXT ((()))))((
- JXXT (e()) ()dad..
(16)
By implementing the hypothesis as he(2) = ((((), t).
where <(-) is a model for representation learning, and vb(-)
for outcome prediction, the above loss can be written based on the empirical distribution of training samples by the following
n Σ (y: (o(())).)))
C
(17)
The second part of the right side of Inequality (15) is to minimize the Wasserstein distances on the empirical dis- tributions W(c, fe: q). where the cost function is measured in the embedding space, i.e., 2zz1, ca(2s, zz) [l(a;;) o(1,))l, and the Wasserstein distance is estimated by the following
min E+> A gcc
(21)
IET
where À is the trade-off hyperparameter between the out- come prediction loss and the distribution discrepancies, & and &b are implemented by neural networks. Figure 1 il- lustrates the framework of our propose method ORIC, and Algorithm summarizes the major procedure of ORIC.
() ===0) ΣΣ
(18)
i=1
For each / 7, we apply the Sinkhorn algorithm to com- pute the Wasserstein distance. Let and d be the numbers of samples and features, the time complexity is in On and the space complexity is in O + nd). We evaluate the time efficiency of our method in Section 5.3.
where 5' is the solution of the following optimization prob- lem
#' = arg min If ETT
co(m,,))) + 77((7) (19)
Loss Function
the set II is defined
n
IT's ( E R
Σ T (((x;) Vi,
j=l
Representation Leaming
Outcome Prediction
Σ nz = ví},
(20)
is the trade-off parameter, the entropic regularization ((x*)= = Σm Σÿ= n/3 log =f, is the negative entropy, and the Sinkhorn algorithm can be applied to solve the prob- lem efficiently (Cuturi, 2013).
The probability mass q(x;;) is approximated as to avoid density estimation (Courty et al., 2017). For the probabil- 2(x.) ((t)a)())))) x ity mass qe (z;), since (x.)) p(z/(t) p()
Figure 1. Overview of our proposed method ORIC.
6

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Synthetic
Methods
IHDP
News
B 3= 0.25
8 = 0.5
8 0.75
B =
KNN
0.2339 = 0.0294
0.2234 = 0.0296
0.2211 = 0.0235
0.2361 = 0.0209
0.8364 = 0.0917
0.6104 0.4117
BART
0.2205 + 0.0248
0.2108 + 0.0312
0.2177 + 0.0259
0.2238
0.0212
0.6825 + 0.0713
0.5639 + 0.3123
GPS
0.2103 = 0.0319
0.2056 = 0.0345
0.2063 = 0.0264
0.2219 = 0.0238
0.7247 = 0.0582
0.4422 = 0.2033
MI.P
0.2083 + 0.0275
0.2042 + 0.0311
0.2044 + 0.0252
0.2185
0.0202
0.6566 + 0.0710
0.4355 + 0.2098
MLP+GPS
0.2077 + 0.0238
0.2028 + 0.0203
0.2022 = 0.0210
0.2161
0.0157
0.6303 = 0.0826
0.4255 + 0.2113
DRNel
0.1992 1: 0.0303
0.2033 1: 0.0226
0.1967 = 0.0172
0.2046 + 0.0195
0.5714 1 0.0211
0.2380 1 0.0141
ADMIT
0.1542 + 0.0325
0.1729 + 0.0467
0.1856 + 0.0345
0.1645
0.0279
0.5222 + 0.0873
0.1832 + 0.0394
ACFR
0.1428 1 0.0250
0.1651 1 0.0325
0.1654 1 0.0334
0.1567 1 0.0248
0.5134 T 0.0523
0.1719 1 0.0767
VCNet
0.1233 + 0.0328
0.1577 + 0.0460
0.1543 + 0.0536
0.1395 + 0.0369
0.4656 + 0.0476
0.1905 + 0.1072
VCNe++TR
0.1155 1: 0.0305
0.1361 + 0.0430
0.1442 + 0.0512
0.1257 + 0.0381
0.3712 1 0.0465
0,1675 1 0.0566
ORIC
0.1098 + 0.0273
0.1234 + 0.0388
0.1313 = 0.0464
0.1168 = 0.0316
0.3595 = 0.0304
0.1507 + 0.0406
Table 1. Comparison of ORIC with baseline algorithms of related networks. The + denotes the mean and standard deviation of v VAMSE.
Algorithm 1 Optimal transport for Reducing blas in Causal
samples, with the parameter 3 adjusted to simulate various confounding biases. IHDP contains 747 subjects, with 25 covariates for each sample to capture the aspects of chil- dren and their mothers. News contains 3,000 news items randomly sampled from Newman (2008), which simulates the opinions of a media consumer when exposed to multi- ple news items. We follow similar approach in Nie et al. (2021) to generate continuous treatments and outcomes, and randomly divide the samples into training set (67%) and a testing set (33%). The detailed synthesis protocols can be found in Appendix E.
inference (ORIC).
Input:
Training samples {a;, 4, 9: }
Initialize: Representation learning model d. potential out- come prediction model vb, generalized propensity score estimator A.
1: repeat
2: Calculate the cost xz) = [0((,)) ((z;))
3: for all / f do
4: Calculate the outcome prediction loss according to Eq. (17).
5: Estimate (z();) based on the normalized general-
Compared methods. We conduct comparison of our ORIC model with several compared methods, including the traditional statistical and machine learning method BART (Chipman et al., 2010), KNN (Peterson, 2009), GPS (Imbens, 2000), and modern neural network based meth- ods MLP, DRNet (Schwab et al., 2020), ADMIT (Wang et al., 2022), ACFR (Kazemi & Ester, 2024). and VCNet (Nie et al., 2021). Specifically, for GPS, in order to en- hance the traditional statistical learning approach, we incor- porate Multilayer Perceptron Network for optimization (GPS+MLP), For VCNet, we consider the naive version of VCNet (VCNet) and VCNet with the target regularization (VC
ized propensity scores (()(z)))
6: Obtain the optimal transport plans 5' by solving Problem (19).
7: Calculate the Wasserstein discrepancies based on #' according to Eq. (18).
8: end for
9: Update and based on the gradient in Problem
(21).
10: until Convergence.
5. Experiments
Evaluation Metries. Following Nie et al. (2021), we adopt the Average Dose DessoRes Function (ADRF) curve and / AMSE as metrics. ADRF curve is the expected potential outcome under the treatment value t, which is defined as # = [t). We repeatedly carry out 100 trials on the simulated and the IHDP datasets, 20 trials on the News dataset, and report the mean and standard deviation of the results on the test set. The computational details are given in Appendix D.
In this section, we present experimental settings and results of continuous and binary treatments. The detailed experi- ments are provided in Appendix E and F.
5.1. Continuous Treatments
Dataset. For the experiments of continuous treatments, we evaluate the performance of the proposed method using one synthetic dataset and two semi-synthetic datasets: IHDP (Hill, 2011) and News (Newman, 2008). The synthetic dataset consists of 500 training samples and 200 testing
Results and Discussions. Table presents the results of
7

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
IHDP
News
Methods
VPEHE
MAE
VAMSE
<PEHE
MAE
VAMSE
BART
13.8853 = 9.3630
9.1204
10.0374 # 7.2281
7.3663 = 2.2189
5.6858 # 1.7925
5.6355 = 1.6635
OLS
14.3736 11.3114
8.8191 +
9.7246 6.9604
8.0871 + 2.3580
5.7820 + 1.6172
6.3790 + 1.865
KNN
3.1108 3.8114
0.4104 0.6477
9.7638 7.4574
7.0048 = 2.3408
5.1976 # 2.0301
5.5409 = 1.7343
ML.P
11.2789 15.3081
8.9105 + 3.1171
11.0619 + 8.5434
8.2535 + 2.4681
5.3473 + 1.6470
6.0092 + 1.7761
CFRNet
1.2809
1.7304
0.1582 0.1986
1.2739 1.7038
2.0527 + 0.6464
0.3080 + 0.2224
2.4187 + 0.6538
Dragonnet
1.8883 1.4305
0.2672 0.4576
1.3229
1.7893
1.7916 1 0.5652
0.3531 1: 0.1724
3.8169 1 1.6722
GANITE
5.0500
1.3205
4.2490 0.6251
13.4438 + 6.7216
2.6473 + 0.6873
2.6375 + 0.6867
6.1070 + 1.1409
DKLite
5.3315
7.0602
0.5472 0.7026
5.7984
7.1115
1.8172 1 0.5182
0,2328 1 0.1272
1.9610 1 0.5701
ESCFR
1.2443
2.1300
0.4112 0.5902
1.3498
2.1298
2.7671 + 0.8924
0.8651 + 0.6514
2.9547 + 0.8822
CausalOT
13.8269 13.5417
2.4498 0.8065
7.3281
6.2416
9.1213 1 2.0043
2.3308 1 0.4832
4.1533 1: 1.0084
ORIC
1.1129 + 1.4290
0.2134 + 0.3488
1.1976 + 1,3822
1.7183 = 0,5488
0.1624 = 0.1587
2.3972 + 0.5678
Table 2. Comparison of ORIC with baseline algorithms on the semi-synthetic dataset. Specifically, we perform over 100 trials the IHDP dataset, and 50 trials on the News dataset.
ORIC and the compared algorithms. Overall, the results indicate that ORIC' consistently outperforms other methods on both synthetic and semi-synthetic datasets, showing the effectiveness of the proposed method. Typically, compared with traditional statistical methods (i.e., KNN, BART, GPS), neural sewwork-based methods usually achieve performance improvement across a variety of datasets. Among the neural network methods, we observe that VC outperforms other methods, showing the advantage the doubly robust property obtained by the targeted regularization. However, lacks an explicit mechanism of distribution alignment to address confounding bias. ADMIT and DRNet split training samples into multiple smaller groups for training, suffering from the issue of data reduction for distribution alignment. Compared with them, ORIC involves all the training sam- ples without data splitting for distribution alignment, and reduces the confounding bias and the outcome estimation error jointly, achieving the best performance in different kinds of datasets. In addition, ORIC obtains promising per- formance with different values of 8, which demonstrates the robustness of the proposed method for different levels of confounding bias. Furthermore, from the ADRF curve in Figure 2, we observe that compared to VCNet, which achieves the best V AAMSF performance among other mod- els, ORIC exhibits a performance improvement on synthetic (8 0.25), IHDP, and News datasets.
Truth Venet
Truth ORIC Voniet
Responsa
Response
Truth
Venet
Prestment
Trestment
Treatment
(a) Synthetic
(b) IHDP
(c) News
Figure 2. The figures from left to right are the ADRF results for the Synthetic, IHDP, and News datasets. The yellow line illustrates the ground-truth results, the blue points represent the predicted results of VCNet, and the red points correspond to the predicted results produced by ORIC.
training set and 1,000 samples to the test set (Johansson et al., 2016). Furthermore, experiments on synthetic data are provided in Appendix F.
Compared methods. We evaluate the proposed method in the binary treatment setting with several baselines, including non nou-neural network methods BART, OLS, KNN, and neural network methods MLP, CFR (Shalit et al., 2017), Dragonnet (Shi et al., 2019), GANITE (Yoon et al., 2018), DKLite (Zhang et al., 2020), ESCFR (Wang et al., 2024), CausalOT (Li et al., 2021).
Evaluation Metrics. For the synthetic dataset, we adopt mean absolute errors(MAB) (Dehejia & Wahba, 1999) as the metric. For semi-synthetic datasets, besides MAE, we adopt VPEHE (Hill, 201 and VAMSE to evaluate the conducted methods. The computational details are given in Appendix D.
5.2. Binary Treatments
Dataset. We conduct experiments on two semi-synthetic datasets, IHDP (Brooks- Gunn et al., 1992) and News (Jo- hansson et al., 2016). For the IHDP dataset, we randomly select 100 datasets from the IHDP-1000 version and follow (Shalit et al., 2017) to split training and testing sets. In the News dataset, we assign the first 3,500 samples to the
Results and Discussions. Tables 2 demonstrate the result across two semi-synthetic datasets in the binary setting. We draw similar observations from the results of the binary
8

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Method
KNN
BART
GPS
MLP
MLP+GPS
DRNet
ADMIT
ACFR
VCNet
VCNe++R
ORIC
Time
8
18
25
26
47
24
17
23
135
Table 3. Running time results (in seconds) of different methods in the continuous treatment setting,
Method
BART
OLS
KNN
MLP
CFRNet
DragonNct
GANITE
DKLITE
ESCFR
CausalOT
ORIC
Time
0.2
0.2
0.3
47
41
4s
165
76
Table 4. Running time results (in seconds) of different methods in the binary treatment setting.
treatment setting to the continuous treatment setting. Ben- efit from the mechanism that involves all the samples for training to avoid data splitting, ORIC achives the best highly competitive performance compared with other meth- ods. This observation demonstrates that ORIC not only can handle continuous treatment, but also obtain promising per- formance in binary treatment, indicating the capability of generalization in different kinds of treatment settings.
porating additional information or assumptions, such the existence of proxy variables.
Impact Statement
This paper presents a balanced representation learning method without data splitting to reduce confounding bias in causal inference. We exploit the tool of optimal transport to analyze the confounding bias and outcome estimation error, and propose novel method based our theoretical analysis. The proposed method can be applied to wide range of applications, such as decision-making in marketing and healthcare.
5.3. Running Time Results
Table 3 reports the running time results on the continuous treatment setting (synthetic data with 3 = 0.25), and Table 4 reports the running time results on the binary treatment setting (the IHDP dataset). We observe that our method achieves moderate time efficiency.
Acknowledgments
This research was supported in part by National Science and Technology Major Project (20212D01 11500) Science Foundation of China (62206061).
6. Conclusion
1501), Natural
In this paper, we estimate the effect of binary and contin- uous treatments by reducing the confounding bias from non-RCTs. We characterize the confounding bias by differ- ent probability measures of the same set of all the samples, and analyze the confounding bias and outcome prediction error based on optimal transport built on probability mea- sures. Motivated by this, we propose to learn balanced representations to reduce the outcome estimation error and the confounding bias simultaneously. By doing this, we avoid data splitting commonly used in existing methods and enhance the generalization ability of the model. We con- duct experiments in both binary and continuous settings to evaluate the performance of our method.
References
Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein gen- erative adversarial networks. In International Conference on Machine Learning, 214-223, 2017.
Bica, I., Jordon, J., and van der Schaar, M. Estimating the ef- fects of continuous- valuee interventions using generative adversarial networks. Advances Neural Information Processing Systems, 33:16434-16445 2020.
Bolley, F., Guillin, A., and Villani, C. Quantitative con- centration inequalities for empirical measures non- compact spaces. Probability Theory and Related Fields, 137:541-59: 2007.
In the future, we plan to investigate the situations involy- ing complex treatments or unobserved confounders. For complex treatments such as bundle or graph treatments, is feasible aggregate the treatment into a continuous treat- ment value, that the proposed method can be employed. The situation with unobserved confounders is more chal- lenging since the ignorability assumption does not hold, and the confounding bias cannot be fully captured by the observed covariates. This issue may be addressed by incor-
Brooks-Gunn. J., Liaw, F-r., and Klebanov, P. K. Efffets of carly intervention on cognitive function of low birth weight preterm infants. The Journal of Pediatrics, 120 (3):350-359, 1992.
Chen, J., Dong. H., Wang, X., Feng, F., Wang, M., and He, Bias and debias in recommender system: survey

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
and future directions. ACM Transactions on Information Systems, 41():11-39, 2023.
Hill, J. L. Bayesian nonparametric modeling for causal infer-
ence. Journal of Computational and Graphical Statistics, 20(1):217-240, 2011.
Chipman, H. A., George, E. L., and McCulloch, R. E. Bart: Bayesian additive regression trees. 2010.
Hirano, K. and Imbens, G. W. The Propensity Score with Continuous Treatments, chapter 7, pp. 73-84. John Wiley & Sons, Ltd, 2004.
Cockx, B., Lechner, M., and Bollens, Priority to unem- ployed immigrants? causal machine learning evaluation of training in belgium. Labour Economics, 80:102306, 2023.
Imbens, G. W. The role of the propensity score in estimating dose-response functions. Biometrika, 87(3):706-710, 2000.
Courty, N., Flamary, R., and Tuia, D. Domain adaptation with regularized optimal transport. In European Confer- ence on Machine Learning and Principles and Practice of Knowledge Discovery in Databases, pp.
Johansson, F., Shalit, U., and Sontag, D. Learning represen- tations for counterfactual inference. In International Con- ference on Machine Learning, pp. 3020-3029. PMLR, 2016.
2014.
Courty, N., Flamary, R., Tuia, D., and Rakotomamonjy, A. Optimal transport for domain adaptation. IEEE Transac- tions on Puttern Analysis and Machine Intelligence, (9):1853-1865, 2017.
Johansson, F. D., Shalit, U., Kallus, N., and Sontag, D.
Generalization bounds and representation learning for es- timation of potential outcomes and causal effects. Journal of Machine Learning Research, 23(166):1-50, 2022.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. In Annual Conference on Neural Information Processing Systems, pp. 2292-2300, 2013.
Kantorovitch. L. On the translocation of masses. Manage- ment Science, 5(1):1-4, 1958.
Davis, J. M. and Heller, B. Rethinking the benefits of youth employment programs: The heterogeneous effects of summer jobs. Review of economics and statistics, 102 (4):664-677) 2020.
Karboub, K. and Tabaa, M. A machine learning based discharge prediction of cardiovascular diseases patients in intensive care units. In Healthcare, volume 10, pp. 966. Multidisciplinary Digital Publishing Institute, 2022.
Dehejia, R. H. and Wahba, S. Causal effects in nonexper- imental studies: Reevaluating the evaluation of training programs. Journal of the American Statistical Associa- fion, 94(448):1053-11 1999.
Kazemi, A. and Ester, M. Adversarially balanced repre- sentation for continuous treatment effect estimation. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 38, pp. 13085-13093, 2024.
Dismuke, C. and Lindrooth, R. Ordinary least squares. Methods and designs for outcomes research, 93(1) (3(3) 104, 2006.
Kreif, N., DiazOrdaz, K., Moreno-Serra, R., Mirelman, A., Hidayat, T., and Suhrcke, M. Estimating heterogeneous policy impacts using causal machine learning: a case study of health insurance reform in indonesia. Health Services and Outcomes Research Methodology, pp. 1-36. 2021.
Dunipace, F. Optimal transport weights for causal inference. arXiv preprint arXiv:2109.01991, 2021.
Fisher, R. A. Design of experiments. British Medical Journal, 1(3923):554, 1936.
Kuang, K., Cui, P., Li, B., Jiang, M., and Yang, S. Es- timating treatment effect in the wild via differentiated confounder balancing. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Dis- covery and Data Mining, pp. 265-274, 2017.
Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., and Courville, A. C. Improved training of wasserstein gans. Advances in Neural Information Processing Systems, 30, 2017.
Hammerton, G. and Munafö, M. R. Causal inference with observational data: the need for triangulation of evidence. Psychological Medicine, 51 (4):563-578. 2021.
Li, Q. Wang, Z., Liu, S., Li, G., and Xu, G. Causal optimal transport for treatment effect estimation. IEEE Transac- tions on Newral Networks and Learning Systems, 2021.
Hatt, Т. and Feuerriegel, S. Estimating average treatment effects via orthogonal regularization. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 680-689, 2021.
Liu, D., Cheng, P., Zhu, H., Dong, Z., He, X., Pan, W., and Ming, Z. Mitigating confounding bias in recommendation via information bottleneck. In Fifteenth ACM Conference on Recommender Systems, pp. 351-360, 2021.
10

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Monge, G. Mémoire sur la théorie des déblais et des rem- blais. Histoire de l'Académie Royale des Sciences de Paris, 1781.
Shi, C., Blei, D., and Veitch, V. Adapting neural networks for the estimation of treatment effects. Advances in Neu- ral Information Processing Systems, 32, 2019.
Newman, D. Bag of words data set. UCI Machine Learning Respository, 289, 2008.
Titouan, V., Courty, N., Tavenard, R., and Flamary, R. Op- timal transport for structured data with application on graphs. In International Conference on Machine Learn- ing. pp. 6275 62844 PMLR, 2019.
Nie, L., Ye, M., Nicolae, D., et al. Venet and functional tar- geted regularization for learning causal effects of continu- ous treatments. In International Conference on Learning Representations, 2021.
Tolstikhin, I., Bousquet, O., Gelly, S., and Schölkopf, B. Wasserstein outo-encoders. In International Conference on Learning Representations, 2018
Peterson, I. E. K-nearest neighbor. Scholarpedia, 4(2): 1883, 2009.
Van Goethem, N., Serrien, B., Vandromme, M., Wyndham- Thomas, C.,, Catteau, L., Brondeel, R., Klamer, S., Meurisse, M., Cuypers, L.., André, E., et al. Concep- tual causal framework to assess the effect of sars-cov-2 variants covid-19 disease severity among hospitalized patients. Archives of Public Health, (9(1): 2021.
Peyré, G. and Cuturi, M. Computational optimal transport, 2017.
Peyré, G., Cuturi, M., and Solomon, J. Gromov-wasserstein averaging of kernel and distance matrices. In Interna- tional Conference on Machine Learning, pp. 2664-2672, 2016.
Villani, C. Optimal transport: old and new, volume 338.
Springer Science & Business Media, 2008
Redko, I., Habrard, A., and Sebban, M. Theoretical analysis of domain adaptation with optimal transport. In European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases. pp. 737- 753, 2017.
Villani, C. Topics in optimal transportation, volume 58.
American Mathematical Soc., 2021.
Wang. H., Fan, J., Chen, Z., Li, H., Liu, W., Liu, T., Dai, Q.. Wang, Y., Dong, Z., and Tang, R. Optimal trans- port for treatment effect estimation. Advances in Neural Information Processing Systems, 36, 2024.
Rosenbaum, P. R. and Rubin, D. B. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55 1983.
Wang, X., Lyu, S., Wu, X., Wu, T., and Chen, H. General- ization bounds for estimating causal effects of continuous treatments. Advances in Neural Information Processing Systems, 35:8605 8617, 2022.
Rubin, D. B. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of edu- cational Psychology, 66(5))688, 1974.
Rubner, Y., Tomasi, C., and Guibas, L. J. The earth mover's distance as a metric for image retrieval. International Journal of Computer Vision, 40(2):99-121, 2000.
Wei, T., Feng. F., Chen, J., Wu, Z., Yi, and He, X. Model- agnostic counterfactual reasoning for eliminating popu- larity bias in recommender system. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pp. 1791-1800, 2021.
Saitoh, S. Integral transforms, reproducing kernels and their applications. CRC Press, 2020.
Sanchez, P., Voisey, J. P., Xia, T. Watson, H. I., O'Neil, A. Q. and Tsaftaris, S. A. Causal machine learning for healthcare and precision medicine. Royal Society Open Science, 98):220638, 2022.
Yan, Y., Yang, Z., Chen, W., Cai, R., Hao, Z., and Ng, M. K.- P. Exploiting geometry for treatment effect estimation via optimal transport. In Proceedings of the AAAI Conference on A Artificial Intelligence, volume 38, pp. 16290-16298, 2024a.
Schwab, P., Linhardt, L., Bauer, S., Buhmann, J. M., and Karlen, W. Learning counterfactual representations for estimating individual dose-nes curves. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 5612-5619, 2020.
Yan, Y., Zhou, H.,, Yang, Z.,, Chen, W., Cai, R., and Hao, Z. Reducing balancing error for causal inference via optimal transport. In Forty-first International Conference Machine Learning. 2024b.
Shalit, U., Johansson, F. D., and Sontag, D. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp. 3076-3085. PMLR, 2017.
Yao, L., Li, S., Li, Y., Huai, M., Gao, J., and Zhang, A. Representation learning for treatment effect estimation from observational data. Advances in Neural Information Processing Systems, 31, 2018.

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Yoon, J., Jordon, J., and Van Der Schaar, M. Ganite: Estima-
tion of individualized treatment effects using generative adversarial nets. In International Conference on Learning Representations. 2018.
Zhang, Y., Bellot, A., and Schaar, M. Learning overlapping representations for the estimation of individualized treat- ment effects. In International Conference on Artificial Intelligence and Statistics, pp. 1005-1014. PMLR, 2020.
12

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
A. Implementation of 0
The implementation of is based on (Hirano Imbens, 2004) and described as follows, Assuming that the conditional distribution of treatment given covariates Gaussian, i.e., P(t 23) N(((o(z))) We can estimate the parameters by maximizing the likelihood:
72 exp
1 (-
20 (-20 (t, 0(((a)))))
max L(0, 8,0
(22)
V2T22
After that, the estimated generalized propensity score is given by:
1
1 28 (1 0()(<;)))3
ptt ar)) = V2T22
exp
(23)
B. Theoretical Analysis regarding Effect Estimation Error
The effect estimation eгrоr captures the bias between the ground-truth causal effect and the estimated causal effect. The following theorem shows that our theoretical results can be applied not only to the outcome estimation error &, but also to the effect estimation error the binary and continuous treatment settings. Let q be the marginal covariate distribution, and ae be the conditional covariate distribution given the treatment value f, i.e., q(x:) = p(x) and (x) = p(x[t), we have the following results.
Theorem B.1 (Binary Treatment Error Bound). Define the binary trealment effect estimation error g(
Ex-q(z) [((h1 ho(x), (1 - -o(2))] Under the assumptions of Theorem 4.2, this error satisfies
8(h) s Em (h1) + Syo (ho) + W(c, Qı q) + W(c, qo 9).
(24)
Proof. Based on the assumptions Theorem 4.2, we first decompose the effect estimation error for the true causal effect T(x)= = M1 (xx) (x) as follows:
c = Exrg(a) (((h((z) ho(z), - mo(x))]
< E Eerglt) (2), (+)) + [eng(e) (((ha(z))
= (h1) + (ho);
(25)
where l is the Ly-norm based loss function and has the triangle inequality property.
We define the estimation eπror of the potential outcome function fl (z) and #z(z) in treatment and control groups, respectively:
Eq (h1) E (x) ((1 (1:), (z)),
(26)
Eg (ho)
(27)
According to Eq. (12), we have
E(1) < (h1)) Eg(h)) Eqi (h1) Ego (ho) W(c, q. 9) + W(c,
(28)
Theorem B.2 (Continuous Treatment Error Bound). Define the continuous treatment effect estimation error g(c) Epsp(titt))) Exzg(a) (((b(()) ha(z), Me(x)) Ho(z))] Under the same assumptions of Theorem 4.2, we have
e(e) < eq (ht) dt + f W(c, des 9) p(t) dt.
(29)
13

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Proof. Based the triangle inequality property, we have:
Els) / Bz Be) - ho(z)) [o(x))]
s (2) ((a)) +
(0 + [
(Eg
(30)
Then according Eq. (12), we have:
E/ s Ee-p(o)E((he))
$ 1ff en (ha)p(()d + J W(c,
W)lldd A)l())..
(31)
C. Proofs of Theorems
C.1. Proof of Theorem 4.1
Theorem 4.1. Let be the marginal covariate distribution, and be the conditional covariate distribution given the treatment value I, i.e. q(x)) p(()) and (e(x)) p(=:(). Given a pair of the functions (m, c) satisfying the condition m(x;) - -m(x,) S c();, We have the following result
E(m) f/ W(Cc
(32)
Proof. According the definition of {(m, t), we have:
5(m, t) (e - me-o(s)
11 m(x)dq((x) m(x)dg(x)
(33)
m(x)dg((a) - -/r m(x)dq(x)
sup m(x) m (clt)a)
(34)
inf TEll(qu.a) xxx
cla, (x)
(35)
W(c, 9)
(36)
Under the assumption of Theorem 4.1, Eq. (34) is the the worst-case of Eq. (33), and Eq. (35) holds because of the property of the dual problem, which just corresponds to the definition of the Wasserstein distance. As a result, we obtain E(m, t) < W(c, qt, q), which finishes the proof by integrating p(t) on both sides of the inequality.
C.2. Proof of Theorem 4.2
Theorem 4.2. Assume that the cost function
6(x')ln where H. is Reproducing Kernel Hilbert Space (RKHS) induced by & X - H. Assume further that his Pt E where F is a unit ball in the RKHS H, and the loss function E(hy(a (x), He(x)) ) convex, symmetric, bounded, obeys the triangular inequality and has the parametric form hy (2) xx(x){ for some X 0. Assume also that kernel ls in the RKHS H is square-roo integrable with respect to X and
=
0 < k(,,2') = (6(x), ((x1)) S K. Then the following holds
[ e (h)h))((())t
f c. (hx)p(t)dt
W(c, qu ())(t)
(37)
14

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Proof. According (Saitoh, 2020), (z), uu (1)) also belongs to the RKHS H. since it is convex loss-function defined on he, ft E F. As result, & has the reproducing property and the norn ||e]] is bounded. For simplicity, assume that ||4] is bounded by 1, which is easily extendable to the case when ||e|| - M by scaling (Redko et al., 2017). Now, the estimation error can be expressed in terms of the inner product in the corresponding Hilbert space,
eg (h,) Ex (() B (()()))) Eqe (hx) = (x) (1))))))))))))))))))))) (e E. Bx (z))))()
(38)
(39)
With (h(h) = (bebs) + Eqe (hr) - Eye (1+) and the above definitions, we have:
Eq(he) - Eq. (h+) = (0)=q())((((())) f) En (m) [(((x) Eydl
(ol - (x [0(x), H)H I (())) Be (x) [x)
- |1 / od(a(( {x)))]l
(40)
The first line is obtained by the reproducing property of and the last line is due ||e|| S 1. Now using the definition of joint distribution, we have:
1/x x
(dl()(+)) gxx)))lxx = /
xxxx (0(x) - (d (xx, x Il
x f xxxx 1x(x)) 6(x') (lzedt (x, x')
/ JxxX
inf
110(z) 6(24) (lad)
(41)
=EIl(uea)
= =W(c, qt; 9t:4),
(42)
where x e ge(2) (x) and u' ~ q(x).. As a result, we get (xe) - Ege (1+) S W(c, gt, q), which finishes the proof by integrating p(f) on both sides of the inequality.
C.3. Proof of Theorem 4.3
Theorem 4.3, Let n be the number of samples, q. qe be the empirical distributions of q, 9s, respectively. With the probability of at least 1 8, we have:
8 f En (h))p(((t) + J W(c, (p((t)tt O (0(1//0))
(43)
Proof. With the triangular inequality of the Wasserstein metric, we have:
W(c 9) S W(c, gt, (t) + ((c,
I W(c, qt q)) + W(c, q. + W(c, W(c. qt â) + W(c. W(c, q.
(44)
Next, we present Lemma showing the convergence of the empirical measure j to its true μ w.r.t. the Wasserstein metric, which allows us propose generalization bound based on the Wasserstein distance for finite samples rather than true population measures:
Lemma 1. ((Bolley et al., 2007), Theorem 1.1). Let be a probability measure in Bdd satisfying Ti (zzet)) inequality, and A Em= 8z, be its associated empirical measure with units. Then for any d' >> and Ç' < Ç. there exists some constant no depending on d' and some square exponential moment of μ. such that for any and n 2 no (d'+2), 1)
PP (M q exp (- (_'ne²)
(45)
where d', "''' can be calculated explicitly.
15

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
The Hoeffding inequality in Lemma gives the following inequality which holds with the probability at least - å:
W(c, 9c 9:) S V 2 log (3) /C'n,
W(c,(.0) S 2200 (
/c'n.
(46)
Combining Eq. (44) and Eq. (46) together, we have:
1 15'n+ +
(7)
1
W(c < / 22og
22
/en + Nccc
1
W(C, 12 +2/ 28/22og log
IC'n
W(e, + O +0(1/6mm)
(47)
which finishes the proof.
D. Evaluation Metrics
In the case continuous treatment, we calculated AMSE as the evaluation metric:
N
AMSE = N
1
Σ 1=1
(iee (x:) - Pe (xx))² p(t)dt
(48)
In the case of binary treatment, besides AMSE, we also calculated PEHE and MAE as the evaluation metrics:
MAE = ATE - ATEI
Σ (x:( ho(x;)) - n
> (m)(s)) - #o(x;))
PEHE 1 n Σ (1:1 ho(x:)) (n1 (2rz) mo(xx))1
(49)
E. Experiments of Continuous Treatments
E.1. Experimental Settings
Synthetic. We synthesize data as follows: 2, iid. Unif [0, 1], where æj is the j-th dimension of æ e R6, and generate treatment and outcome as:
i x sin (max (I1, 22, 23)) + max 24, 26) sin (Bxa) (1 + exp (x4 - Bxx))
1 1 (x. x26)
+ 23 2si (I4)
2xz 6.5 (O,
4 max (1))
y =o,,t=cos(2nr) (3))
1
sin (Ira) + (..(0.0.55)
2n
where = (1+exp(-t))-15 (0025,, 0.5,0..... 0.75, 1). is noteworthy that T(t ()) only is contingent upon 11, 22, 03, 24,
while (t, x) only is contingent upon Ii, T3, R44 26-
16

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
IHD. The original semi-synthetic IHDP dataset from Hill (2011) includes binary treatments, comprising 747 observations across 25 covariates. To facilitate comparisons using continuous treatments, randomly synthesize both treatment and response variables as follows:
2x1
max (I3, Is, 466
5 -iEssan (2g -c2)
+
+ 2 tanh
4 +++ (0,
(1 + 22)
2 + min (Ii3, 255 (13, vin)
Sdis,2
y v1 xr = sin(3r()
(tanh 5 Diesenn (I: (1)-c))
exp (2 (I1 - 260)
+
+ +0
1.2 -
[Sdis
0.5 + min (I22, T3 23))
(1,2,3,5,6} is the index set of continuous features, Sais ,1 {4, 7, 8,8, 10, 11, 12, 13, 14, 15), Sdis 22 = {16, 17, 18, 19, 20, 21, 22, 23, 24, 25} and Sais, U Sdis = [25] -Scon Here C1 E- 1aa. - Z4 C2 _ -ieem
where t = (1 + exp( (sp(-t)) -1) S
It is noteworthy that all continuous features are advantageous for T(( | 2)) and (tt, 2²) but only Sais.1 is advantageous for Q and only Sais.2 is advantageous for æ. Following Hill (2011), covariates are standardized to have a mean of O and a standard deviation of 1, while the synthesized treatment values are normalized to the range [0, I]. Furthermore, we applied denoising techniques to the error data produced during the construction of the IHDP dataset.
1St..
News. The News dataset comprises 3,000 randomly sampled news items from the NY Times corpus (?), originally introduced as a benchmark for binary treatment settings (Johansson et al., 2016). We synthesize the treatment and outcome variables similarly to the method outlined in Bica et al. (2020). We first synthesize v1, 122 and v/ from Ao 1) and then set [ 1) And we synthesize the outcome by 201
U/2 = / 1| for = (1, Given x, we synthesize t from Beta (2,
U2 x 0g x
v x, t = exp
0.3
y xzt= 20022mmmxx (max( (-2, min (2, + 200 2)) (4(t= 0.3)² sin
) + 2
E.2. Sensitivity Analysis
To empirically study the effect of the hyper-parameter À in Eq. (21) which trades off between the outcome prediction loss and the Wasserstein discrepancies, we conduct experiments on the synthetic dataset (B = 0.25) with varying values of in the range [0.5, 1.3), and present the results of v VAMSE in Figure 3(a). We observe that ORIC is able to achieve good performance with a wide range of the values of X, which verifies the sensitivity of ORIC with respect to X.
Figure 3(b) illustrates the impact of varying the entropy regularization hyperparameter y within the range [0.0001, 0.1] on the model's performance, as measured by v VAMSE. The results demonstrate the trade-off associated with different values of 7, highlighting how the choice of this hyperparameter influences the balance between exploration and exploitation in the model.
Besides, we conduct experiments on the synthetic dataset with different numbers of sampled treatment values in the discrete set T, and report the results of / VAMSE in Figure 3(c). We observe that ORIC stably achieves promising performance when the number of discrete values of the treatment is greater than 50, since more values of the treatment provide finer-grained estimation for the conditional marginal distribution qr (z).
E.3. Ablation Study
Table 5 illustrates the ablation study on the loss function involving Wasserstein distances. ORIC outperforms the version without the Wasserstein distances, which demonstrates the effect of the distance for balanced representation learning.
F. Experiments of Binary Treatments
F.1. Experiment settings
Synthetic. Following the similar protocols in (Yao et al., 2018; Hatt & Feuerriegel, 2021), we generate a synthetic dataset in the binary treatment setting as follows:
17

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
026
020
0.84
611
6.110
0:2-
6.00
0.1
0.7
12 25
6.000
1,0022 0004 23335 6000 1,0000 o 6632 6.024 60036 8000
610 30
40
50
30
(a) byper-parameter A
(b) hyper-parameter 7
(c) number of treatments T
Figure Figure (a) demonstrates the trade-off of the hyperparameter À between outcome prediction loss and the Wasserstein discrepancies with the variation of values ranging from [0.5, 1.3), and presents the results of VAMSE. Figure (b) demonstrates the trade-off the entropy regularization hyperparameter y values range from [0.0001, 0.1]. and presents the results of /AM SE. Figure (c) illustrates the experiment the synthetic dataset with different numbers of sampled treatment values in the discrete T, and reports the results of VAMSE.
Synthetic
Methods
IHDP
News
8 == 0.25
/ = 0.5
3 0.75
ORIC without
0.2077 = 0.0238
0.2028 + 0.0203
0.2022 1 0.0210
0.2161 + 0.0137
0.6303 1 0.0826
0.4255 1 0.2115
ORIC
0.1098
0.0273
0.1234 1 0.0388
0.1313 1 0.0464
0.1168 1 0.0316
0.3595 1 0.0304
0.1507 1 0.0406
Table Ablation study the loss function involving Wasserstein distances. The results of the mean and standard deviation of VAMSE are reported.
We employ Gaussian mixture model consisting of two distributions: Ni N (0.510x1 0.5 x E ,N2 N 11 0.5x EgΣT), where Σn U ((0, (00) En ~ u ((0, 1)10x10). We then synthesize 1,500 treated and control samples from x ~ a.Ni + (1 N2, xe NNNi + (1 ax) N2. fix ar to 0.5 and vary the value of a, to simulate different confounding bias. The outcomes are defined as y sin [w] æ) + cos (w) (a @ x)) + + where
w. ~ U ()3x3) (NO, 0.1).
E22 Results and Discussions
Tables 6 illustrate the results of synthetic data in different bias situations. We draw similar observation in the continuous setting. ORIC outperforms other methods and achieves the best results in different levels of confounding bias, indicating the superior performance of robustness.
Besides, we compare our method with baselines of the in-sample experiments on the IHDP dataset in Table 7. Our method also achieves promising performance in the in-sample setting.
We also compare the results using the ground-truth and estimated propensity scores on the IHDP dataset. Since the ground-truth propensity score is unknown, adjust the IHDP dataset assign treatment accordingly. Table shows the results in terms of different metries. Our method achieve comparable results even with the estimation error in the propensity score model, which demonstrates the robustness with respect to propensity scores.
E.3. Visualization Results
We conduct experiments on simulation dataset visualize the embeddings before and after representation learning based t-SNE. We generate covariates of the treated and control groups that are drawn from multivariate Gaussians: x' M1 = N (0.520x1, 0.5 x ΣgΣT) xx ~ N₂ = A' (120x1 0.5x ΣgΣT), where Σ₁ ~ U ((O, 0.8)20x00), Σg ~ u ((0, 1.2) 20x20) The visualization results are shown in Figure 4. Our method can learn balanced representations for the treated and control groups.
18

---

Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport
Synthetic
Methods
ag 0.2
0, 0.4
ag 0.6
ag 0.8
BART
0.0622 + 0.0374
0.0484 + 0.0194
0.0255 + 0.0206
0.0397 = 0.0207
OLS
0.0568 + 0.0420
0.0471 + 0.0361
0.0387 + 0.0234
0.0412 = 0.0259
MLP
0.0862 + 0.0813
0.0803 + 0.0600
0.4992 + 0.0422
0.0621 0.0388
KNN
0.0229 + 0.0196
0.0276 + 0.0198
0.0306 + 0.0184
0.0296 = 0.0259
CFRNet
0.0328 + 0.0063
0.0326 + 0.0065
0.0383 + 0.0326
0.0475 = 0.0345
Dragonnet
0.0351 + 0.0104
0.0323 + 0.0092
0.04778 + 0.0061
0.0482 0.0067
GANITE
0.1883 + 0.0530
0.1779 + 0.0672
0.3219 + 0.0574
0.3916 = 0.0581
DKLite
0.0599 + 0.0338
0.0432 + 0.0158
0.0302 + 0.0344
0.0753 0.0463
ORIC
0.0052 0.0089
0.0282 0.0048
0.0235 0.0166
0.0291 + 0.0186
Table Comparison ORIC with baseline algorithms on synthetic dataset. Specifically, we conducted over 10 trials on synthetic datasct, adopting M AE the evaluation metric.
Methods
/PEHE
MAE
VAMSE
CFR
1.0462 + 1.0905
0.4966 - 0.4711
1.0437 + 1.0215
GANITE
8.0017 + 5.3730
5.3945 1.0848
13.3972 10.6536
DKLite
5.0756 + 6.0795
0.2252
0.2440
5.5372 6.1226
CausalOT
10.2003 = 4.5611
2.7824
1.4760
8.2140 9.1621
ESCER
1.0019
1.6507
0.4434 0.5371
2.0842 1. 1.6892
ORIC
0.8463 1 0.7730
0.3539 0.3996
0.8173 0.7391
Table Performance different methods on the IHDP (binary) dataset under in-sample setting.
Methods
/PEHE
MAE
VAMSE
Ground-truth PS
1.4624 0.1222
0.1662 1 0.1255
2.1058 0.1526
Estimate PS
1.3 + 0.0800
0.2012 + 0.1477
1.9891 + 0.1329
Table 8, Results using ground-truth and estimated propensity scores.
Treated Group
Treated Group
Group
Control Group
(a) before balance
(b) after balance
Figure 4. (a) shows the distribution of original covariates, while Figure (b) displays the distribution of the learned embeddings, where the I-SNE used for projection. Samples receiving treatment are marked in red, and those control group are shown in blue.
19


Reviews:
[{"review_id": "8rrg4Fg70b", "reviewer": "Reviewer_6NpC", "summary": "This paper extends CFRNet to use all samples for each treatment when computing loss functions. Thus, it fits the continuous treatment setting better as it does not suffer from the sample splitting problem. However, this is at the cost of modeling propensity scores and density estimation for $q(x)$ and $q_t(x)$. The theoretical results extend those in CFRNet to their setting without sample splitting. Experiments show the proposed method leads to smaller errors in PEHE, MAE and AMSE.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We appreciate the reviewer for the valuable comments. We will revise the submission according to the comments and responses.\n\n**Q1**\nThe time and space complexity for computing the Wasserstein distance.\n\n**A1**\n1) In practice, to avoid heavy computation, we consider a set $\\widehat{\\mathcal{T}}$ including sampled treatment values evenly distributed in the continuous treatment space $\\mathcal{T}$. The number of treatment values is set as $[20, 100]$, as shown in Figure 3.\n\n2) For each $t \\in \\widehat{\\mathcal{T}}$, we apply the Sinkhorn algorithm to compute the Wasserstein distance. \nLet $n$ and $d$ be the numbers of samples and features,\nthe time complexity is in $O(n^2d)$, \nand the space complexity is in $O(n^2 + nd)$.\n\n3) We report the running time results in **A3** to Reviewer L5M9.\nOur method has moderate time efficiency.\n\n**Q2**\nRegarding $q_t(x)>0$.\n\n**A2**\nThe condition $p(x) > 0$ is implicitly embedded in Asm. 3.3.\nThis is because the conditional density could be rewritten as $p(t|x) = \\frac{p(x,t)}{p(x)}$, where we assume $0 < p(t|x) < 1$. \nIf $p(x) = 0$, then $p(t|x)$ would be undefined. \n\n**Q3**\nRegarding the claim that a small $\\mathcal{W}(c,q_t,q)$ leads to poor outcome prediction performance.\n\n**A3**\nOnly optimizing the Wasserstein distance without incorporating outcome prediction loss can easily lead to balanced but trivial latent representations, such as mapping all samples to a single point. This is known as the over-balancing issue, as stated in [a][b].\n\n**Q4**\nDifference between Theorem 4.3 and Theorem 1 in [c].\n\n**A4**\n1) [c] reduces the discrepancy between the subpopulations $q_1(x)$ and $q_0(x)$, each of which is modeled by the samples in one group.\nDifferent from it, we reduce the discrepancy between the marginal distribution $q(x)$ and the conditional distribution $q_t(x)$,\nwhich is modeled by all the samples equipped with the propensity scores.\nThis difference is significant since data splitting is avoided and more data are leveraged for conditional distribution modeling.\nThe data splitting issue becomes even more severe in the continuous treatment setting since multiple groups are considered, and each group receives only a part of the samples,\nhampering the performance of distribution estimation and confounding bias reduction.\n\n2) [c] applies IPM to measure the confounding bias,\nwhile our analysis applies the Wasserstein discrepancy to measure the confounding bias.\nAlthough IPM can be implemented as the Wasserstein-1 distance, the Wasserstein discrepancy based on different underlying cost functions cannot be represented as IPM.\n\n**Q5**\nAdvantage of the proposed method compared with CFRNet.\n\n**A5**\n1) CFRNet splits training data into two groups to estimate $q_1(x)$ and $q_0(x)$.\nDifferent from it, our method leverages all the training data to estimate $q_1(x)$ and $q_0(x)$.\nTherefore, more training data are involved in conditional distribution modeling.\nTheorem 4.3 also shows that more training samples, i.e., a large $n$, helps to achieve a better error bound.\n\n2) CFRNet reduces the discrepancy between $q_1(x)$ and $q_0(x)$.\nDifferent from it,\nwe aim to reduce the discrepancy between $q_1(x)$ and $q(x)$,\nand the discrepancy between $q_0(x)$ and $q(x)$,\nwhich can be naturally applied into the continuous treatment setting without considering pairs of different treatments and data splitting.\n\n3) Existing studies have demonstrated that propensity scores are helpful for representation learning [d][e].\nWe design a different approach to leverage propensity scores, i.e., to reduce the discrepancy between conditional and marginal distributions based on propensity scores.\n\n**Q6**\nHow the error in propensity score models and density estimation models impact the final performance of the proposed method.\n\n**A6**\nFollowing existing works of optimal transport such as [f], we adopt $q(x)= \\frac{1}{n}$ to avoid density estimation.\n\nWe compare the results using the ground truth and predicted propensity scores.\nSince the ground-truth propensity score is unknown, we modified the IHDP dataset to assign treatment accordingly it.\nOur method can achieve comparable results even with the errors in propensity score models. \n\n|               | PEHE            | MAE             | AMSE            |\n|---------------|-----------------|-----------------|-----------------|\n| ground-truth ps | 1.4624 ± 0.1222 | 0.1662 ± 0.1255 | 2.1058 ± 0.1526 |\n| Estimate ps   | 1.3400 ± 0.0800 | 0.2012 ± 0.1477 | 1.9894 ± 0.1329 |\n\n[a] On learning invariant representations for domain adaptation, ICML 2019.\n\n[b] Counterfactual representation learning with balancing weights, AISTATS 2021.\n\n[c] Estimating individual treatment effect: generalization bounds and algorithms, ICML 2017.\n\n[d] CounterFactual Regression with Importance Sampling Weights, IJCAI 2019.\n\n[e] Counterfactual representation learning with balancing weights, AISTATS 2021.\n\n[f] Optimal transport for domain adaptation, TPAMI 2017.", "claims_and_evidence": "See below.", "methods_and_evaluation_criteria": "The method is an extension of the loss function proposed by CFRNet to continuous treatment. \n\n1. The main concern of the proposed method is the computational cost of the wasserstein distance $\\mathcal{W}(c_{\\phi},\\hat{q}_t, \\hat{q})$ and the overhead of density estimation to obtain $\\hat{q}_t$ and $\\hat{q}$. For the Wasserstein distance, basically for each treatment $t$, each plan $\\pi^t$ is an n by n kernel. I am not sure how many values of $t$ and plans are there. Can the authors discuss the time and space complexity for computing the Wasserstein distance?", "theoretical_claims": "1. The authors claim $q_t(x) > 0$ for all $x$ given Asm 3.3, which seems not true. If $p(x)=0$, then no $q_t(x)=0$, regardless of $p(t|x)$,\n2. Although it can be true intuitively, the claim that a small $\\mathcal{W}(c,q_t,q)$ leads to poor outcome prediction performance due to losing information for outcome prediction is not well supported by theoretical analysis in this work.\n3. Theorem 4.3 is an extension of Theorem 1 in [1] to continuous treatment, which says the outcome estimation error is upper bounded by the sum of the estimation error on factuals and the distance between the treated and controlled group in the representation space. I guess the main difference is that Theorem 4.3 allows soft $p(t)$, that is, for each $x$ you estimate the propensity score $p(t)$ and plug it into the upper bound. Could the authors clarify the difference between their contribution and existing work in terms of Theorem 4.3?\n\n[1] Shalit, Uri, Fredrik D. Johansson, and David Sontag. \"Estimating individual treatment effect: generalization bounds and algorithms.\" International conference on machine learning. PMLR, 2017.", "experimental_designs_or_analyses": "The experiments include both continuous and binary treatment.\n\n1. For the binary treatment case, I think the proposed method is very similar to the CFRNet with Wasserstein distance, except it uses a soft propensity score to compute the error and the Wasserstein distance. I wonder how it can be much better than CFRNet in this case, especially considering the error of propensity score models and density estimation models can lead to error in counterfactual outcome prediction.\n\n2. It would be better to add an experiment to show how the error in propensity score models and density estimation models impact the final performance of the proposed method ORIC.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aJnEsLgYk8", "reviewer": "Reviewer_L5M9", "summary": "This paper proposes a novel method for causal effect estimation based on optimal transport. The method reduces the confounding bias without data splitting, which is different from the existing methods that partition training into multiple groups according to treatments. Theoretical and empirical results are provided to evaluate the performance of the proposed method. Extensive experiments on both binary and continuous treatment settings are conducted.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We appreciate the reviewer for the valuable comments. We will revise the submission according to the comments and responses.\n\n**Q1** \nRegarding the effect estimation error of the continuous treatment setting.\n\n**A1**\nWe analyze the effect estimation error of the continuous treatment setting below.\nFollowing [a], we define the effect estimation error $e_{\\tau }^{G}$ in the continuous treatment setting:\n\\begin{align}\ne_{\\tau}^{G} = E_{t \\sim p(t | t \\neq 0)} E_{x \\sim q(x)}[ l( h_{t}(x) - h_{0}(x), \\mu_{t}(x) -\\mu_{0}(x))]\n\\end{align}\nWith the similar procedure in Appendix B, we have:\n\n$e_{\\tau}^{G}= E_{t \\sim p(t | t \\neq 0)} E _ {x \\sim q(x)}[ l( h_{t}(x) - h_{0}(x), \\mu _ {t}(x) - \\mu _ {0}(x))]$\n\n$\\leq E _ {t \\sim p(t | t \\neq 0)} [E _ {x \\sim q(x)}[ l( h _ {t}(x), \\mu _ {t}(x))] + E _ {x \\sim q(x)} [ l( h _ {0}(x), \\mu _ {0}(x))]]$\n\n$= E _ {t\\sim p( t|t\\neq 0)}[ \\varepsilon _ {q}( h _ {t}) +\\varepsilon _{q}( h _ {0})]$\n\n$= E _ {t\\sim p( t)}[ \\varepsilon _ {q}( h _ {t})]$\n\n$\\leq \\int_{\\mathcal{T}} \\varepsilon_{q_{t}}( h_{t}) p(t) dt+\\int_{\\mathcal{T}}\\mathcal{W}(c,q_{t} ,q) p(t) dt.$\n\nwhere the first inequality holds due to the triangle inequality property, and the second inequality holds because of Eq. (12) in the paper.\n\n**Q2**\nReplace $\\mathcal{R}$ with $\\mathbb{R}$ in Section E.1.\n\n**A2**\nThank you for the valuable suggestion.\nWe will revise the submission accordingly.\n\n**Q3**\nRegarding the running time results.\n\n**A3**\nThe following table presents the running time results.\nOur method achieves moderate time efficiency.\n1) Continuous treatment setting on synthetic data $(\\beta = 0.25)$ \n\n| Methods  | Times |\n|----------|-------|\n| ORIC     | 135s  |\n| VCNet+TR | 23s   |\n| VCNet    | 17s   |\n| ADMIT    | 47s   |\n| ACFR     | 24s   |\n| DRNet    | 26s   |\n| GPS+MLP  | 25s   |\n| MLP      | 18s   |\n| GPS      | 9s    |\n| BART     | 7s    |\n| KNN      | 8s    |\n\n2) Binary treatment setting on the IHDP-1000 data \n\n| Methods   | Times |\n|-----------|-------|\n| ORIC      | 76s   |\n| CFRNet    | 47s   |\n| DragonNet | 41s   |\n| DKLITE    | 4s    |\n| ESCFR     | 165s  |\n| CausalOT  | 4s    |\n| GANITE    | 4s    |\n| BART      | 0.2s  |\n| OLS       | 0.2s  |\n| KNN       | 0.3s  |\n\n\n**Q4**\nIs it possible to extend the proposed method to more complex settings, such as bundle or graph treatments?\n\n**A4**\nThank you for the valuable comments.\nIn general, the bundle or graph treatments remain opening problems.\nInspired by the assumption in networked interference [b][c],\nit is feasible to aggregate the information of bundle or graph treatments into a continuous treatment value,\nso that the proposed method can be employed.\nWe will investigate this challenging problem in the future.\n\n[a] Estimating heterogeneous treatment effects: Mutual information bounds and learning algorithms, ICML 2023.\n\n[b] Identification and estimation of treatment and interference effects in observational studies on networks, JASA 2021.\n\n[c] Learning individual treatment effects under heterogeneous interference in networks, TKDD 2024.", "claims_and_evidence": "The claims are well supported by the theoretical analysis and extensive experiments.", "methods_and_evaluation_criteria": "Benchmark datasets of both binary and continuous treatment settings are used in the experiments, and multiple evaluation metrics are adopted.", "theoretical_claims": "The theoretical claims are correct and well presented.", "experimental_designs_or_analyses": "The experimental designs and analyses are sound", "relation_to_broader_scientific_literature": "The paper makes contributions to causal effect estimation, especially confounding bias reduction in binary and continuous treatment settings. The proposed method can be applied in different areas, such as policy decision and healthcare.", "essential_references_not_discussed": "As  far as I know, the essential references are well discussed and cited.", "other_strengths_and_weaknesses": "Strengths\n\n1.     The paper proposes a novel method to reduce confounding bias without data splitting, which is an under-explored and interesting topic.\n\n2.     The extensive theoretical analysis regarding confounding bias, outcome estimation error, and effect estimation error are provided.\n\n3.     The experiments are sufficient. Both binary and continuous treatment settings are considered, multiple evaluation metrics are adopted, and many compared methods are conducted.\n\n4.     The paper is well organized.\n\nWeaknesses\n\n1.     The balanced representation learning relies on the Wasserstein distance. It may bring more computation to solve the optimal transport problem.\n\n2.     The analysis of effect estimation error in Section B only considers the binary treatment setting. Although I understand that the studies of continuous treatment usually consider potential outcome estimation instead of effect estimation, it would be better to analyze the effect estimation error of the continuous treatment setting.", "other_comments_or_suggestions": "1. It seems that the analysis in Section B is related to the binary treatment setting. Is it feasible to derive similar results regarding the continuous treatment setting?\n\n2. In Section E.1, I suggest to replace the notation $\\mathcal{R}$ by $\\mathbb{R}$, which is consistent with the main part of the submission.", "questions_for_authors": "1. It would be better to evaluate the efficiency in terms of the running time results.\n\n2. Is it possible to extend the proposed method to more complex settings, such as bundle or graph treatments?\n\n3. Based on Section B, is it feasible to derive some theoretical results regarding the effect estimation error of the continuous treatment setting", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g66reSz6bu", "reviewer": "Reviewer_qAwp", "summary": "This paper proposes an effective algorithm for estimating treatment effects while reducing confounding bias, applicable to both binary and continuous treatments. It employs optimal transport methods to utilize all available samples for estimating confounding bias, thereby mitigating bias and avoiding the decrease in estimation accuracy associated with splitting training samples into smaller groups for distribution alignment.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We appreciate the reviewer for the valuable comments. We will revise the submission according to the comments and responses.\n\n**Q1**\nThe performance advantage of the proposed method diminishes in binary cases and limited continuous treatment cases.\n\n**A1**\nThis observation is reasonable.\nCompared with the setting with more treatment values,\nfor the setting of binary treatments or limited continuous treatments, more training samples fall into each group to achieve better performance of distribution modeling and alignment.\nFor example, for the binary treatment setting, each group has around half of the samples,\nthus can achieve a good performance.\nHowever, with more number of treatments,\neach group receives fewer samples,\nwhich tends to decrease the performance.\n\n**Q2** \nThe advantage of the algorithm over distribution alignment methods using different treatment groups.\n\n**A2**\nDistribution alignment across different treatment groups splits training data into different subpopulations,\nwhich cuts down the numbers of training data,\nhampering the performance of distribution modeling and confounding bias reduction.\nThe issue becomes more severe in the continuous treatment setting since many different treatment values are involved,\nand each group receives only a small part of the training samples.\nDifferent from them, our method considers all the training samples in each treatment group,\nwhich leverages more training data for distribution modeling and alignment. \nThe experimental results demonstrate the advantage of our method,\nand Theorem 4.3 also shows that more training data, i.e., a large $n$, helps to achieve a better error bound.\n\n**Q3** \nRegarding the computational complexity of the Wasserstein distance\nand training on large-scale data.\n\n**A3**\nThe computational complexity of the Sinkhorn algorithm is in $O(n^2d)$,\nwhere $n$ and $d$ are the numbers of the samples and features.\n\nThe following table presents the running time results.\nOur method achieves moderate time efficiency.\nFor large-scale data, it is feasible to consider optimal transport on mini-batch samples, as shown in [a].\n\n1) Continuous treatment setting on synthetic data $(\\beta = 0.25)$ \n\n| Methods  | Times |\n|----------|-------|\n| ORIC     | 135s  |\n| VCNet+TR | 23s   |\n| VCNet    | 17s   |\n| ADMIT    | 47s   |\n| ACFR     | 24s   |\n| DRNet    | 26s   |\n| GPS+MLP  | 25s   |\n| MLP      | 18s   |\n| GPS      | 9s    |\n| BART     | 7s    |\n| KNN      | 8s    |\n\n2) Binary treatment setting on the IHDP-1000 data \n\n| Methods   | Times |\n|-----------|-------|\n| ORIC      | 76s   |\n| CFRNet    | 47s   |\n| DragonNet | 41s   |\n| DKLITE    | 4s    |\n| ESCFR     | 165s  |\n| CausalOT  | 4s    |\n| GANITE    | 4s    |\n| BART      | 0.2s  |\n| OLS       | 0.2s  |\n| KNN       | 0.3s  |\n\n**Q4**\nRegarding the inconsistency in the full name of the proposed method.\n\n**A4**\nWe are sorry for the confusion.\nWe will revise the submission accordingly.\n\n**Q5**\nRegarding the existence of unobserved confounders.\n\n**A5**\nSince we characterize the confounding bias by measuring the discrepancy between $q_t(x)$ and $q(x)$, the ignorability assumption is required. If unobserved confounders exist, the confounding bias can not be fully captured by considering $q_t(x)$ and $q(x)$ only. We will investigate the situation with unobserved confounders in the future.\n\n\n[a] Improving mini-batch optimal transport via partial transportation, ICML 2022.", "claims_and_evidence": "The proposed theorem on the confounding bias bound strongly supports the algorithm's design and its proof is rather valid.\n\nExperiments compare the performance of the proposed algorithm with baseline methods in both binary and continuous treatment scenarios. Note that the algorithm's performance advantage diminishes in binary cases and in the limited continuous treatment cases shown in Appendix E. The authors may consider providing further explanations regarding this observation.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-suited to address causal inference under confounding bias. ORIC leverages optimal transport to align distributions without splitting data, overcoming the limitations of traditional methods due to reduced sample sizes—a reasonable approach since distribution estimation relies on sufficient data. The use of neural networks for representation learning and the Sinkhorn algorithm for Wasserstein distance computation is theoretically grounded.", "theoretical_claims": "The theoretical framework provides an upper bound on confounding error related to distribution differences characterized by optimal transport to support the algorithm's design. Additionally, experiments compare the proposed algorithm with multiple baseline methods, demonstrating its effectiveness.", "experimental_designs_or_analyses": "I examined the experimental designs and analyses in Section 5 and Appendices E and F. For continuous treatments, Synthetic/IHDP experiments ran 100 trials, and News ran 20, sufficient for statistical reliability. Binary treatment experiments on IHDP (100 trials) and News (50 trials) followed similar standards. Comparisons with multiple baselines (e.g., KNN, BART, VCNet, CFRNet) are comprehensive, with appropriate metrics. Sensitivity analysis in Appendix E.2 tests hyperparameters, demonstrating robustness. The ablation study (Table 3) isolates the impact of the Wasserstein component. The design is sound, data generation protocols are clear, and no obvious methodological flaws are apparent.", "relation_to_broader_scientific_literature": "This article falls within the field of causal effect estimation, particularly regarding the use of distribution alignment to reduce confounding bias.", "essential_references_not_discussed": "The related work and citations in this paper is comprehensive.", "other_strengths_and_weaknesses": "Strengths:\nS1: The article is well-organized, and the mathematical expressions are precise.\n\nS2: The proposed theorem on the confounding bias bound strongly supports the algorithm's design and is proven to be valid.\n\nS3: The paper conducts experiments comparing the proposed algorithm with multiple baseline methods, demonstrating its effectiveness.\n\nWeaknesses:\nW1: The advantages of the algorithm over distribution alignment methods using different treatment groups need further comparison and explanation.\n\nW2: The effectiveness of the algorithm concerning data splitting issues should be further demonstrated. Experimental results indicate that the performance advantage decreases in binary cases and in limited continuous treatment scenarios.\n\nW3: The proposed algorithm introduces significant computational complexity through nested loops to calculate the optimal solution for the Wasserstein distance. The paper should clarify this point.", "other_comments_or_suggestions": "Page 2, Line 67, and Page 6, Line 291: There is a inconsistency in the full name of the proposed algorithm ORIC.", "questions_for_authors": "1. I understand that the key to this paper is linking quantified confounding bias to the Wasserstein distance, which relies on several assumptions, such as the ignorability assumption for quantifying confounding bias. How does the algorithm perform when unobserved confounders exist, and the difference between $q_t$ and $q$ cannot quantify confounding bias?  \n\n2. The paper uses the Sinkhorn algorithm, which may slow down training. How does it perform on large-scale datasets?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8G73rTuQcz", "reviewer": "Reviewer_UGhX", "summary": "This paper focuses on the causal inference task, specifically in the binary and continuous treatment settings. The authors argue that data sparsity can hinder covariate distribution alignment across different treatment groups, leading to biased outcome predictions. Instead, they push all conditional marginals forward to the marginal distribution, where the formers are implemented by the generalized propensity score. Theoretically, they provide a bound for optimization. Extensive experiments are conducted to evaluate their method.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "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": "We appreciate the reviewer for the valuable comments. We will revise the submission according to the comments and responses.\n\n**Q1** \nA visualization of the balance results is missing.\n\n**A1**\nThank you for the valuable comments.\nWe have further conducted experiments to visualize the embeddings before and after representation learning based on t-SNE.\nThe results shows that our method can reduce the distribution discrepancy caused by the confounding bias.\nPlease kindly refer to https://anonymous.4open.science/r/ICML_figure_commit-3B64. \n\n**Q2**\nRegarding the conditional distribution estimation based on the generalized propensity score.\n\n**A2**\n1) We exploit the generalized propensity scores to estimate the conditional distribution,\nwhich leverages all the training samples for distribution modeling and alignment without data splitting.\nThe data splitting issue becomes even more severe in the continuous treatment setting since multiple groups are considered, and each group receives only a part of the samples,\nhampering the performance of distribution estimation and confounding bias reduction.\nTherefore, our method can achieve better performance.\nThe experimental results demonstrate the effectiveness of our method.\nIn addition,\nTheorem 4.3 also shows that more training samples, i.e., a large $n$, helps to achieve a better error bound.\n\n2) Our experimental results demonstrate the effectiveness of our method involving propensity score estimation.\nIn addition, existing studies have demonstrated that propensity scores are helpful for representation learning [a][b].\nWe design a different approach to leverage propensity scores, i.e., to reduce the discrepancy between conditional and marginal distributions based on propensity scores.\n\n**Q3**\nThe computation of the generalized propensity score is missing in Figure 1.\n\n**A3**\nDue to space limitations,\nwe present the details of the computation of the generalized propensity score in Section A of the appendix.\nWe will revise Figure 1 to highlight the computation of the generalized propensity score.\n\n**Q4**\nRegarding the in-sample results.\n\n**A4**\nThank you for the valuable comments.\nWe have added the results of the in-sample experiments on the binary IHDP dataset in the following table.\nOur method achieves promising performance on the in-sample setting.\n\n|          | PEHE             | MAE             | AMSE              |\n|----------|------------------|-----------------|-------------------|\n| CFR      | 1.0462 ± 1.0905  | 0.4966 ± 0.4711 | 1.0437 ± 1.0215   |\n| GANITE   | 8.0017 ± 5.3730  | 5.3945 ± 1.0848 | 13.3972 ± 10.6536 |\n| DKLite   | 5.0756 ± 6.0795  | 0.2252 ± 0.2440 | 5.5372 ± 6.1226   |\n| CausalOT | 10.2003 ± 4.5611 | 2.7824 ± 1.4760 | 8.2140 ± 9.1621   |\n| ESCFR    | 1.0019 ± 1.6507  | 0.4434 ± 0.5371 | 2.0842 ± 1.6892   |\n| ORIC     | 0.8463 ± 0.7730  | 0.3539 ± 0.3996 | 0.8173 ± 0.7391   |\n\n**Q5**\nThe differences of the models are minor in Table 4 while larger in Tables 1 and 2.\n\n**A5**\nIn general, the outcome values of the real-world data IHDP and News in Tables 1 and 2 are large,\nwhile the outcome values of the simulation data are small.\nAs a result, the differences of the models on real-world data are larger compared with those on simulation data.\nSimilar observations can be drawn from existing studies [c][d].\n\n[a] Counterfactual Regression with Importance Sampling Weights, IJCAI 2019.\n\n[b] Counterfactual representation learning with balancing weights, AISTATS 2021.\n\n[c] Perfect match: A simple method for learning representations for counterfactual inference with neural networks. arXiv:1810.00656\n\n[d] GANITE: Estimation of individualized treatment effects using generative adversarial nets. ICLR 2018.", "claims_and_evidence": "N/A", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes, the proof for the Theorem 4.2 is correct.", "experimental_designs_or_analyses": "1. In Tables 1 and 2, the large performance gaps between different baselines seem strange.\n2. A visualization of the balance results for mitigating the selection bias is missing.", "relation_to_broader_scientific_literature": "This paper propose a new method to help address the selection bias in treatment effect estimation, which has wide applications in healthcare, ecomony, biology, and so on.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths:**\n\n1. The proposed method is straightforward and easy to understand.\n2. The writing in this paper is well-done, and the prerequisite knowledge is comprehensively provided.\n\n**Weaknesses:**\n\n1. Instead of optimizing the traditional balance term between $q_0(x)$ and $q_1(x)$, this paper enforces all conditional marginal distributions $q_t(x)$ to be aligned with the marginal distribution $q(x)$. However, $q_t(x)$ is approximated by the generalized propensity score, which is generally difficult to predict and has high variance. Thus, the effectiveness of this method is questionable.\n\n2. Figure 1 does not highlight the unique technical contribution of this paper, such as the computation of the generalized propensity score.\n\n3. In Tables 1 and 2, the performance gaps between different baselines are too large—they are on different scales. Additionally, the in-sample experiment is missing.\n\n4. Why do the results of different models in Table 4 show only minor differences, while in Tables 1 and 2, the differences between models are much larger?\n\n5. A visualization of the balance results is missing, which is crucial because it demonstrates that balance can be achieved without data splitting.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "See weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a method called ORIC for causal effect estimation of continuous or binary treatment variables. It can use all samples to compute the loss function without the sample splitting process dividing the training samples into multiple groups based on the treatment value, which is commonly relied upon by existing methods. Getting rid of sample splitting mitigate the difficulty to alignment due to insufficient sample size within a single group. Theoretically, this article gives the upper bound of the proposed method for covariate balance error and outcome estimation error. Empirically, a large number of experiments have been conducted in continuous and binary cases, verifying the superiority of the proposed method over the baselines in multiple metrics such as PEHE, MAE, and AMSE.

Overall, the reviews for this paper is positive. On the one hand, this paper is widely recognized for its good organization and writing, solid theoretical results, and extensive experimental results. On the other hand, the reviewers are generally concerned about the complexity of ORIC, and the author added the algorithm complexity analysis and experimental runtime results in the rebuttal, showing several times the time cost compared to the baselines with continuous treatment variables, and slightly smaller increase with binary variables. The author also responded to other questions from the reviewers, including questions about theoretical details, sensitivity analysis, differences and advantages compared to existing methods, experimental results in the sample, visualized experimental results, etc. Overall, considering the review comments, rebuttal, and the paper itself, I recommend that this paper be accepted.
Decision: Accept (poster)