Adjusting Prediction Model Through Wasserstein Geodesic for Causal Inference
Authors: Yuguang Yan
Haolin Yang
Zecong Chen
Weilin Chen
Ruichu Cai
Zhifeng Hao
Venue: ICLR 2026 Poster (2026)
Keywords: causal inference
Primary area: causal reasoning

Abstract: Causal inference estimates the treatment effect by comparing the potential outcomes of the treated and control groups. Due to the existence of confounders, the distributions of treated and control groups are imbalanced, resulting in limited generalization ability of the outcome prediction model, \ie, the prediction model trained on one group cannot perform well on the other group. To tackle this, existing methods usually adjust confounders to learn balanced representations for aligning the distributions. However, these methods could suffer from the over-balancing issue that predictive information about outcomes is removed during adjustment. In this paper, we propose to adjust the outcome prediction model to improve its generalization ability on both groups simultaneously, so that the over-balancing issue caused by confounder adjustment can be avoided. To address the challenge of large distribution discrepancy between groups during model adjustment, we propose to generate intermediate groups through the Wasserstein geodesic, which smoothly connects the control and treated groups. Based on this, we gradually adjust the outcome prediction model between consecutive groups by a self-training paradigm. To further enhance the performance of the model, we filter the generated samples to select high-quality samples for learning. We provide the theoretical analysis regarding our method, and demonstrate the effectiveness of our method on several benchmark datasets in terms of multiple evaluation metrics.

Published as a conference paper at ICLR 2026
ADJUSTING PREDICTION MODEL THROUGH WASSER- STEIN GEODESIC FOR CAUSAL INFERENCE
Yuguang Yan', Haolin Yang's Zecong Chen', Weilin Chen', Ruichu : Zhifeng Hao
1 School of Computer Science, Guangdong University of Technology, Guangzhou, China,
2 Pazhou Laboratory (Huangpu). Guangzhou, China,
3 College of Science, Shantou University, Shantou, China.
ABSTRACT
Causal inference estimates the treatment effect by comparing the potential out- comes of the treated and control groups. Due to the existence of confounders, the distributions of treated and control groups are imbalanced, resulting in lim- ited generalization ability of the outcome prediction model, ie., the prediction model trained on one group cannot perform well on the other group. To tackle this, existing methods usually adjust confounders to learn balanced representa- tions for aligning the distributions. However, these methods could suffer from the over-balancing issue that predictive information about outcomes is removed during adjustment. In this paper, we propose to adjust the outcome prediction model to improve its generalization ability on both groups simultaneously, so that the over- werr-balancing issue caused by confounder adjustment can be avoided. To ad- dress the challenge of large distribution discrepancy between groups during model adjustment, we propose to generate intermediate groups through the Wasserstein geodesic, which smoothly connects the control and treated groups. Based on this, we gradually adjust the outcome prediction model between consecutive groups by self-trainin self-training paradigm. To further enhance the performance of the model, we fil- ter the generated samples to select high-quality samples for learning. We provide the thcoretical analysis regarding our method, and demonstrate the effectiveness of our method on several benchmark datasets in terms of multiple evaluation met- rics.
INTRODUCTION
Causal inference from observational data has been widely used in many real-world applications to
evaluate the effects of treatment (Yazdani Boerwinkle 2014, Varian 2016). such as healthcare
(Foster et al.] 2011)) and recommendation (Sato et al. 2020 Luo et al. 2024 Gao et al.
2024
Wang et al. 2024b). Based on the Rubin-Neyman potential outcome framework (Rubin 1974
Splawa-Neyman et al. 1990), the treatment effect can be estimated by comparing the potential outcomes under treatment and no treatment, where the outcome models are independently trained on the treated and control groups estimate the respective potential outcomes (Künzel al. 2019).
Nevertheless, due to confounders that affect both treatment assignment and potential outcomes, the covariates of the control and treated groups follow significantly different distributions, thereby limiting the generalization ability of the outcome prediction models. For instance, patients receiving surgery (i.e., the treated group) typically have more severe conditions than patients not receiving surgery (i.e., the control group). As a result, the prediction model trained on the treated group fails to accurately predict outcomes for the control group.
To tackle the above confounding bias, existing methods usually adjust and align confounders by learning latent representations, so that the distributions of the control and treated groups are bal- anced, and the outcome prediction models can be generalized between groups (Robins et al. 20000] Shalit et al., 2017]] Johansson et al.] 2018). However, these methods could excessively balance the distributions of two groups while ignoring discriminative information which are vital for outcome prediction, suffering from the over-balancing issue (Yao et al., 2018 Du et al. 202T). In the extreme
"Corresponding author; cairuichu @gmail.com

---

Published as a conference paper at ICLR 2026
situation where two distributions collapse to single point, the distributions are perfectly balanced, while the predictive information for outcomes is completely eliminated. Although a compromise between distribution balancing and outcome prediction can be considered (Shalit et al. 2017), it is heuristic and remains under-explored how to achieve a promising trade-off.
In this paper, seek to adjust outcome prediction models rather than covariates, so that the over- non-trivial to adjust the outcome prediction model from one
balancing issue can be avoided. group the other group due the significant distribution shift between them. To address this, we propose a method named G-learner, which generates a sequence of intermediate groups that drift gradually from the control group to the treated group, and then gradually adjusts the outcome predic- tion model trained on one group to the other. Specifically, we establish an optimal transport model between the control and treated groups, which induces geodesic path between these two distribu- tions in the sense of the Wasserstein metric (Villani 2021 Such a Wasserstein geodesic is derived from the theory of optimal transport (Monge, 1781 Kantorovitch, 1958), which has demonstrated powerful ability to characterize data distribution (Courty et al. 2016 2017). Based this, we exploit data distributions of the control and treated groups to generate intermediate groups, which smoothly connect the two groups and are beneficial for model adjustment.
Through the generated intermediate group sequence derived from the Wasserstein geodesic, we em- play a self- self-training strategy to transfer the prediction model trained on the control group towards the treated group and vice versa. Since the model adjustment is performed between two consecu- tive distributions, it is feasible effectively transfer the prediction model trained the previous distribution to the next distribution. To further improve the generalization ability of our model on different groups, design data filtering method select generated samples according to the confidence of the predicted results, so that reliable generated samples are leveraged to transfer the prediction models.
We summarize our major contributions as follows.
We propose intermediate group generation method based on the Wasserstein geodesic to smoothly connect the control and treated groups, so that one group can be gradually shifted the other group.
We gradually adjust the outcome prediction model between the control and treated groups through the generated intermediate groups, which improves the generalization ability of the outcome prediction model.
We theoretically analyze the effect estimation error of our proposed method, and empiri- cally demonstrate the effectiveness of our method on several benehmark datasets.
2 RELATED WORKS
Over the past decades, variety of methods have been proposed to address the issue of confound- ing bias. Reweighting aims to reweight samples to reduce the distribution shift between the groups (Hainmueller 2012). For example, [Rosenbaum & Rubin (1983) adopt the inverse of propensity scores as the sample weights, and Kuang et al. (2017) learn sample weights by aligning the mo- ments between the control and treated groups Representation learning methods adjust covariates by learning balanced representations, so that the distributions of the control and treated groups in the embedding space can be aligned (Johansson et al. 2016 Shalit et al. (2017) trains neural network to learn balanced representations, where the distribution shift is measured by the Integral Probability Metric (IPM). In (Kazemi & Ester, 2024), the Kulbback-Leiblrr divergence is adopted measure the distribution shift. However, these methods could excessively balance the covariates while ignor- ing the discriminative information for potential outcome prediction, which known as the issue of over-balancing and hampers the performance of causal inference (Johansson et al.) 2018 Zhao et al.) 2019). Different from them that adjust the covariates by learning sample weights or representations, adjust the outcome prediction model to improve the generalization ability on different groups, avoiding the over- balancinnn issue in existing methods.
Recently, optimal transport has also been employed for causal inference (Yan et al., 2024 Wang al., 2026). Optimal transport studies how to transport masses from one distribution to another dis- tribution with minimized cost (Monge 1781 Kantorovitch 1958 Villani 2021). Existing studies
2

---

Published as a conference paper at ICLR 2026
have demonstrated the powerful ability of optimal transport model data distribution (Courty al. 2016 2017; Adler & Lunz 2018). Inspired by this, optimal transport is applied in causal inference to reduce the confounding bias (Li et al. 202I). In (Yan et al., 2024), sample weights are learned based on a semi-relaxed optimal transport model between control and treated groups. In (Shalit et at... 2017 Wang al. ). IPM is implemented by the Wasserstein distance measure the dis- tribution shift between groups, and balanced representations are learned to minimize the Wasserstein distance. Wang et al. (2025b) utilizes Wasserstein distance to minimize the distribution gap between biased observations and the target population for effective debiasing. These methods are still under the paradigm of reweighting or representation learning, in which the optimal transport cost mini- mized to align the distributions of the control and treated groups. Different from them, we generate intermediate groups between distributions through the Wasserstein geodesic, which smoothly con- nects the control and treated groups in the sense of the metric defined by optimal transport. Based on the generated intermediate groups, we can gradually adjust the outcome prediction model from one group another.
Self-training methods (Grandvalet & Bengio 2004 Zou et al. 2019 Gao et al. (02T) have been ap- plied to the problem of sumi-supervised learning and domain adaptation (Nigam et al. 2000 Grand- valet & Bengio 2004| |Han et al. 2019). AAmini & Gallinari (2002) presents a semi-supervised algo- rithm that improves classification by iteratively training on unlabeled data with pseudo labels. Han eet al. (2019) calibrates predictive uncertainties between source and target domains using Bayesian neural networks, enabling reliable pseudo-labeling and effective unsupervised domain adaptation. He et al. (2024) proposes to generate intermediate domains and applies gradual self-training, signif- icantly improving domain adaptation when intermediate domains are scarce. Compared with them, we leverage self-training to learn counterfactual outcome prediction models, and design data fil- tering method to select reliable samples for training. Moreover, we provide a theoretical analysis regarding the effect estimation error of our self-training methods.
3 PROBLEM STATEMENT
We adopt the Neyman-Rubin potential outcomes framework (Rubin) 1974 Splawa-Neyman et al. 1990). Let {(x, t= denote samples drawn from the joint distribution of covariates X, treated assignment T,, and outcome Y, where X; E X' denotes the covariates of the i-th sample, t, E T = [0, indicates the treatment assignment, and y; E R denotes the observed factual outcome under the treatment 4, The observed outcome Y is the potential outcome (() corresponding to the actually received treatment t. The control group (t 0) received treatment is denoted as 18gg {22, and the treated group (t = 1) received the treatment is denoted as { where no and nı are the numbers of the samples in two groups, respectively. The first subscript 0/1 is the group index, and the second subscript i/j is the sample index.
For given x E X, our objective is to estimate the conditional average treatment effect (CATE) as follows:
(x) (1) Y(O)X = x].
(1)
To facilitate the identification and estimation of causal effects within the framework outlined above, typically invoke the following key assumptions:
Assumption (Stable Unit Treatment Value Assumption). 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.
Assumption 2 (Ignorability). Conditional on covariates, the treatment assignment is independent of potential outcomes: T TL Y (f)]x.
Assumption 3 (Positivity). Conditional on covariates, the treatment assignment is not deterministic: < TT = x|X = x) 1.
Under these assumptions, 7 (xx) can be represented by f(x, 1) as follow:
7xx) (1)
-xxx)(xxx
which quantifies how the effects vary with the treatment received by an individual. To predict potential outcomes, we define a function l(x, to estimate the ground-true ITE as follow:
+(x) h(x, 1) h(x, 0).
(3)
3

---

Published as a conference paper at ICLR 2026
Control group
Intermediate group
Treated group
hy
PaLaK-r"
[xg]
Dul
Nucsing (c))
bu Pecgoral
Outcome under control
Outcome under
Space of Covariates
Figure 1: Overview of G-learner. In the figure, the yellow and blue boxes represent the outcomes of the corresponding groups. Here, Yo and /I denote the observed factual outcomes of the control group and treated group, respectively, while the remaining outcomes are pseudo outcomes generated by N or ,,, The observed covariates of the control and treated groups are also shown as xo and x1; the remaining covariates Ix are generated via optimal transport.
For convenience, we define f(x, fx) and (x) h(x, t).
Throughout the paper, 1n E R" is vector with all the entry being 1. For the matrix A, Aij is the (i, j-th entry, A is the transpose of A. ||o lo is the fo-norm that counts the non-zero entries of A. The probability simplex En defined as En {v IR" Σn
EE 1 1, Vi).
METHODOLOGY
Figure [ fillustrates the main idea of our proposed method G-learner. Given the control group (t 0) and the treated group (t = 1) with covariate shift, we generate intermediate groups indexed by x E (0, 1) between the control and treated groups based on optimal transport. After that, gradually adjust outcome prediction models between groups with the help of generated data, during which the models trained on filtered samples of high quality. In the following, we present the technical details of our method.
4.1 INTERMEDIATE GROUP GENERATION
Due to the existence of confounders, the control and treated groups follow different distributions. As result, the outcome prediction model trained on one group cannot perform well the opposite group, suffering from limited generalization ability (Shalit et al. 2017). To address this, we propose to connect the control and treated groups by generating series of intermediate groups, which is derived by the Wasserstein geodesic from the optimal transport theory.
Optimal transport seeks to find the optimal plan to move mass from one distribution to another with a minimized transport cost (Villani 2008 Peyré et al. 2019). Formally, given the feature space X, let M(X') be the set of Radom measures, E M and 3 _M(X) be two distributions, whose corresponding samples are and 6, respectively. E [( 3) is a transport plan, and [( (8) the of all joint probability couplings whose marginal distributions indicates how many masses are transported from the sample to the sample 6, and the transport cost between them measured by cost function c(a, b), where e X X I can be implemented by distance metric. Optimal transport minimizes the total transport cost between the distributions and defines the following p- Wasserstein distance
and {3 respectively. 7(a, b)
Wp(a, (8) ( inf
(yEtto..)) / XxxX
cla, ()db))))(a,bb)
(4)
4

---

Published as a conference paper at ICLR 2026
In the remainder, adopt p 2 and implement the cost function by the Euclidean distance, i.e., c(a, b) = a 222, which derives the 2- W-Wasersstinn distance W2-
The Wasserstein distance not only measures the discrepancy between two data distributions, but also induces a geodesic structure of the space of probability measures. The geodesic path can be naturally constructed via optimal transport-based interpolation between two distributions (Villani, 2008), and the distributions along the geodesic path provide a smooth connection between two distributions. Mathematically, the distributions on the Wasserstein geodesic can be obtained by solving the following problem parameterized by K E [0,
arg min (1 - (a, KW2(3,
(5)
The constructed distribution Mx is also known as the Wasserstein barycenter of two distributions with the parameter K, which is the mean of two probability measures under the optimal transport metric (Bonneel et al., 2013). The barycentric projection considers the special case with Courty et al. (2016), which obtains the representations of samples in μo by weighted samples in the distribution f61-
Our target is to adjust the outcome prediction model trained on the control group to the treated group and vice versa, so that the outcome prediction models l(x,, 1) and hxx, 0) can perform well on both groups. To this end, we establish an optimal transport problem between the control and treated groups, and generate the intermediate groups between them based on the Wasserstein geodesic. Through the generated intermediate groups, we can gradually adjust the outcome prediction models from one group to another, which avoids the over-balancing issue in balanced representation learning and achieves prediction models with promising generalization ability.
Specifically, for the control samples {x and the treated samples {xir) let { and {PI.; 1 jj= be the corresponding probability mass, and 6(x) be the Dirac function at the location x. The empirical distributions of the control and treated groups can be represented as
R2
1o =
Po. (xx,)) p1 = =E P1.6((x1.)) j=1
(6)
where the first subscript 0/1 indicates the group index while the second subscript corresponds to the sample index. The probability masses can be denoted as a probability simplex, i.e., Po = [po.1..... mp.
Eno. P1 = Ip., PI.n.. E Σn,, =Without additional prior information of the underlying distribution, we usually implement po and P1 as uniform distributions (Courty et al., 2016), i.e., Po,; no Vi, P1.j zi of maximum entropy, which states that the probability distribution that best represents the current observations the one with the largest entropy. The equal weights in the uniform distribution also reflect that each sample has the same importance. We also try obtain marginal distributions based on kernel density estimation, and report the results in Appendix B.8 Based on this, the transport plan y defined in the following set
1 Vj. The uniform distribution can be derived from the principle
N[Ho P1) = {> "OX0
X
111,, Po. Ino = PIl:
(7)
and the optimal transport plan is obtained by solving the following optimization problem, which is the discrete form Problem
= arg min
7as Xa. x1.il22
s.t. YE TETOH,
(8)
This is constrained linear programming problem that can be efficiently solved by the Earth Mover Distance solver (Flamary et al.) 202T). Optimal transport requires a well-defined cost function, which is implemented by the squared Euclidean distance here. For complex types such as discrete or categorical data, one can define a distance or learn a mapping function to project them into a representation space, that optimal transport can be applied.
Following Remark 7.1 in (Peyré et al. 2019) and Theorem 5.27 in (Santambrogio 2015). given an optimal coupling n, the barycenter is obtained by the push-forward interpolation ps. P PR.g where Pt (No, 21) ) E I Rd + (1 x)00 + kzg For the discrete setup, once the optimal transport plan
5

---

Published as a conference paper at ICLR 2026
" is obtained. the intermediate group derived by Eq. (5) boils down to the interpolation defined as the following empirical distribution (Villani 2008 He et al. 2024
Ro
n1
Hx
1 7)(((1 X)x i + xx1.)).
(9)
We generate a sequence of intermediate groups with different values of x to connect the control and treated groups. For given E [0, 1/, Me involves |/7| generated samples, in which each non-zero entry Mij > O is the probability mass of the corresponding sample. As result, we obtain weighted generated samples {(x,,,j. 24)) m. > 0} that can be used to gradually adjust outcome models by
self-training. According (Peyré et 2019), the total sum of nx for the set X5. is ng + n1 1.
4.2 PREDICTION MODEL ADJUSTMENT
To smoothly connect μi and Mi, generate K intermediate groups indexed by A ( ) and then adjust the outcome prediction models between two consecutive groups by a self-training paradigm (He et al. 2024). In the following, we take the model ha(-) predicting the potential outcome of non-treatment as the example to present method of model adjustment. The adjustment for the model h1 (-) to predict the treated potential outcome is similar.
Specifically, we transfer (-))) from the group &T - x the group K, and define ho,x- (-) as the non-treatment outcome prediction model have been trained on the group HT. the initial step with <T 0, the model h- is trained on the control group 14 with ground-truth non-treatment factual outcomes. Although ho,o(-) (-) cannot perform well on the treated group (41, is expected to perform much better on the distribution M1/(K+1) compared with the distribution 11, since the covariate shift between μo and P1/(K+11 is much smaller than that between (oo and μ41. Therefore, (x) can be used as supervised information for the sample x E #1/(K+1) to train ho, which can be leveraged as the annotation function to train the model (K,2/(K++))( (-) in the next step. Finally, we can obtain the prediction model ho, (-) predict non-treatment potential outcomes for the treated group fr. (-) is gradually adjusted from the model ho.u(-) through the intermediate groups derived from the Wasserstein geodesic.
Formally. model adjustment from the group R" to the group can be achieved by solving the following problem
ho.x arg min
Σ ((h(xx))hho. (xx)),
(10)
Xx Epe
where the sample X, in the distribution (xx is annotated by the model hn (.), and then is used to train the model h(-). The loss function implemented by the squared loss based on the sample weight obtained from optimal transport, ie.,
((h(xxa)); bave (xm.is)) = v (I(xx.))) - ha.x (Xx,,s))²
(11)
where Xn.ij is sample the distribution μ4x with the weight 7i.. which is obtained according to Eq. 9
To improve the performance of potential outcome prediction, we further refine our model by filtering generated data, which is described by the following.
4.2.1 GENERATED DATA FILTERING.
During the procedure of adjusting ho(-) from the distribution μx- to frx, we train the model on the samples x, E fsc with the pseudo labels ho,x- (xx). In order to leverage high-quality samples to train the model, we filter generated data to remove samples with low prediction certainty. Specifi- cally, for the sample Xx, we conduct M times dropout-enable forward passes to obtain M predicted outcomes, which are denoted as { (x,,,xx)))M, M- After that, we calculate the standard deviation of the multiple predicted results a (xx)) to measure the uncertainty of prediction. To improve reliabil-
ity, we select the r' proportion of the generated samples with the lowest standard deviations, which enjoy high prediction certainty and are beneficial for model training (Gal & Ghahramani 2016). As

---

Published as a conference paper at ICLR 2026
result, construct a filtered group Ao with high quality generated samples, and refine Problem as follows
I0
ho.s arg min Σ
((h(xa),ho..- (xx)).
(12)
Obeying a similar approach above, the model to predict treated potential outcomes h1(+) is adjusted from the group x = = R+T to the group x by the following
h1 arg min
Σ E(h(xx), h1, (x,))(
(13)
X. EAn
where (-) (-) trained on the distribution μ1 with ground-truth treated factual outcomes, and h1,0(-) is gradually adjusted from (-) through the intermediate groups.
We summarize our proposed method G-learner in Algorithm 1 in the appendix.
4.3 THEORETICAL ANALYSIS
In this part, we analyze the estimation eirors of outcomes and effects of our method. We de- fine K = K + for notational simplicity, where K is the number of intermediate groups. Let be the average Wasserstein distance between two consecutive distributions, E (ho, fo(x)))x(x(( group, Similarly, E( -fx (x),,((() fo(x))"p(x{)
1)dx be the expected loss of the prediction model ho, the treated
0)dx is the expected loss of the predic- tion model ho.o on the control group. We have the following results
Lemma 1. Following (He al. 2024), consider arbitrary measures from group and HT, the error of | (ho,,) E(ho,x-) upper bounded with probability at least r as:
[E(ho,x) - E(ho,x )) OWW Mx-) pC Vlog(1/r)
(14)
Vn
Under Lemma I it can be readily observed that the error between group and K as measured by the Wasserstein distance Wp(pee: Mie ), depends on the sample size n and the probability r.
Lemma 2. For filtering ratio E (0,1). the loss E(ho,1) is upper bounded with probability at least 1 - as
+
/lognk)sL nk
Vn + K log(1/r)
1
log(1/r),
1(ho,1) S (((ho,))+((K) +
Vnk +
+
nK
(15)
where n. is the number of training samples, L is the model depth, and R& is the accumulated distribution shift between the control and treated groups through intermediate groups.
In addition, let [ - ffik)--f(o (x))² = 1)dx the expected loss the prediction model h1 1 the treated group. The effeet estimation error is measured by the pairwise precision in the estimation of heterogeneous effect (PEHE) is defined as 6PEHE E (+())( - (x dxx
(Shalit et al. 2017). According to Lemma 2 the effect estimation error is upper bounded by the following theorem.
Theorem 1, The effect estimation error EPEHE upper bounded by:
CPEHE & 28 (ho,o) 28(h, 20 KAA B(n, K K,,,,)
(16)
where
K L,r) Vn
log(1/r)
1
(log 23)2 + +
log(1/r) nk
+ K
+
(17)
Vnk
nk
The proof of Lemma 2 and Theorem is given in the appendix. Theorem indicates that the сpEhE upper bound by the prediction error of the treated and control group. with an additional term depending on n, K, r and L. particular, with fixed K, r, and L, the additional term approaches zero as 2o.

---

Published as a conference paper at ICLR 2026
Table 1: Results for out-sample performance on real-world datasets in terms of mean and standard deviation. A lower metric indicates better performance. We highlight the best results in bold and underline the second best results.
Method
News
Twins
Jobs
VEPEHE
EATE
VEPEBE
FATE
Rpoz
ÉATT
OLS
4,0051 + 1.6369
0.4777 + 0.2632
0.5212 + 0,4070
0.0128 0.0150
0.2612 + 0.0632
0.1825 0.1224
BART
5.8950
0.0082
0.0305
0.0075
0.1133
0.1092
T-leamer
2,9004 + 0.9077
0.5526 + 0.4227
0.3385 + 0.0087
0.0224 + 0.0126
0.2769 + 0.0180
0,1409
0.0962
K-NN
13,0410 = 11.2283
9,5025 = 5.2287
0,3735
0.0087
0.0330 = 0.0089
0.2642 # 0.0686
0.1731 0.1494
PSM
7,0947 + 2.1623
5,3695 = 1,4663
0.3954
0,0072
0.0168 = 0.004
0.2667 + 0.0746
0,1495
0.1190
GANITE
3.7838 1. 1.2759
1,7325 0.6216
0.3202
0.0085
0.0130 1 0.0055
0.2811 1 0.1308
0.1648 0.1083
DKI ite
3,4381 + 1.3630
0,8782
0.5693
0.3207
0,0086
0.0086 0.044
0.1730 + 0.0001
0.1470
0.1099
TARNet
4.2959 = 0,3899
0.3413 = 0.0085
0.0129 = 0.0054
0.2364 + 0.0672
0.0920 = 0.0810
DragonNet
3.2137 0.8048
0,7673 0.5775
0.4454
0,0110 0.0101
0.0810
0,2516 0.0506
BNN
4,2182 1 1.2550
2 1 0.6801
0,3202 1. 0.0085
0.0132 1. 0.0046
0.2576 1 0.1275
0.0782 0.0755
3.0044
1.049 0.7090
0.0095
0,0220 0.0254
0.0703
0.0920
0.0808
CFRMND
3,0688 1 1.2633
1.0373
1. 0.7943
0.3233 1 0.008S
0.0279 1. 0.0185
0.2390 1 0.0760
0.0938
0.0813
CFR-Pro
3.3837 1.2978
0.5963
0.3206 0.0088
0.0161 0.0133
0.0623
0.0937
0.0863
CE-RCFR
3.2956 + 1.0875
1.2729 = 0,4131
0.3214 + 0,0095
0.0244 + 0.0202
0.1781 + 0.0401
0,1265
0. 1012
DESCN
4,3308 = 1.2460
2.6741 = 0.6953
0,3244 = 0.0081
0.0168 = 0.0217
0.3011 = 0.0509
0,1920 00,0088
ESCFR
2.7435
0,9110
0.4255
1 0.3082
0.3209 0,0085
0.0147 1 0.0083
0.239
1 0.0438
0,0893
0.0801
G-learner
2.8681
0.8696
0.2451 0.1955
0.3200 0.0086
0.0084 0.0060
0.1691 0.0622
0.0596
0.0740
Table 2: Results on out-sample simulated dataset in terms of mean and standard deviation. A lower metric indicates better performance and we highlight the best results in bold and underline the second best results.
Method
The 0.5
ThL 0.8
TT, 1.1
mi 1.4
VEPEHE
CAvE
V VEPEEE
CATE
(PEEH
CATE
V
VFPEBE
CATE
OLS
1.12 L: 0.50
0.30
0.22
1.13 1 0.50
0.30 1 0.22
1.16 0.50
0.30 4 0.24
1.31 0.59
0.35 1. 0.28
BART
4.07 0.66
0,98
008-0.3. 0.31
4.03
0.59
0.41
0.28
4.28
0.53
4.71
0.50
+
T-leamer
8.23 + 0.78
0.41
0.34
8.93
+ 0.78
2.79
= 0.51
10.20
= 0.70
5.40 - 0.50
12.10
= 0.63
8.07
0.53
KNN
4.49 : 0.84
1
3.76
1 0.91
0.85
0.24
3.69
0.71
0.37 1 0.22
4.19 0.56
0.21 0.11
PSM
1,98 1. 0.59
0.51
151-1.55 0.53
2.68
0.69
0.61
0.41
4.16
0.99
0.86 1 0.61
5.76 1761110 1.24
1.96 1. 0.88
GANITE
0.85 + 0.08
0.28
0.22
1.49
+ 0.03
1.47 0.06
1.50
0.01
1.50 0.01
1.51 +001
1.50 + 0.01
DKLite
0.59 1. 0.11
0.07
0.05
2.73
173.1...6 0.69
2.47 0.78
5.84
0.66
5.73 1 0.65
8.96 1.63
8.89 1 0.64
TARNet
0.55 +- 0.05
0.53 0.05
0.55
+ 0.04
0.53 + 0.05
0.54 0.06
0.52 + 0.06
0.55 0.05
0.53 + 0.05
DragonNet
0.27 0.04
0.05
0.25
0.05
0.38 + 0.15
0.06 0,05
0.10 0.22
BNN
0.42 0.01
0.40
0.8
0.46=0.00
0.44 1 0.01
0.41 = 0.01
0.01
0.42 1 0.01
CFRarasa
0.52 0.03
0,50
0.03
0.05
0.51
0.54 0.07
0,55
L 0.05
CFRADD
0.51 + 0.03
0.49 - 0.03
0.51 0.03
0.49 0.03
0.50 0.08
0.48 + 0.08
0,96 + 1.38
0.56 + 0.27
CFR-Pro
0.54 0.44
0.43
0.55 0.5
0.56 0.45
0.55
0.75 0.14
0.46 0.10
CE-RCFR
0,47 = 0.02
0.03
0.03
0.09
0.52 1 0.02
0.50 0.01
0,52 0.02
0.49 1 0.02
DESCN
0.52 1: 0.05
0.49
009:000 0.06
0.52
1 0.05
0.49
1 0.06
0.51 0.06
0.49 1 0.01
0.52 0.01
0.49 1 0.01
ESCFR
0.59 4 0.15
0.12
0.07
0.75
1 0.18
0.24
+ 0.11
0.96 0.21
0.40 + 0.12
0.96 0.21
0.40
4 0.12
G-learer
0.20 1: 0.03
0.05
0.04
0.23
1. 0.04
0.05 1 0.05
0.29 0.07
0.08 1. 0.05
0.36 0.14
0.09 1 0.11
5 EXPERIMENTS
In this section, we first describe the experimental settings, including the compared methods and evaluation metrics. After that, we present experimental results and discussions on real-word and simulation datasets.
5.1
EXPERIMENTAL SETTINGS
5.1.1
COMPARED METHODS.
We compare the performance of G-learner with classical and state-of-the-art methods, including
OLS, BART (Chipman et al. 22100] Hill 2011), PSM (Rosenbaum & Rubin 1983), k-NN (Crump)
et al. 2008). T-learner (Künzel et al. | 2019). TA RNet (Shalit et al. (2017). 7). BNNN (Johansson et al.
2016), CFRwass Wass and CFR (Shalit et al.]
2017). GANITE (Yoon al. 2018), DragonNet
(Shi et al. 2019), DKLite (Zhang et al. 2020). CFR-Pro (Wang et al., 2025a
CE-RCFR (Wang
ee al. 2002a). DESCN (Zhong et al. 2022 and ESCFR (Wang et al. 2023).
8

---

Published as a conference paper at ICLR 2026
5.1.2 EVALUATION METRICS.
Following (Shalit et al. 2017 we evaluate the performance of different methods based on the Precision in Estimation of Heterogeneous Effect error (CPEhE) and the absolute error in estimating the Average Treatment Effect (CATR). Additionally, we also assess performance using the policy risk (Rpot) and the error in the Average Treatment Effect on the Treated group (EATT). Complete formulas of these metrics are given in the Appendix.
5.2 EXPERIMENTS ON REAL-WORLD DATA
We compare the performance of different methods on three real-world datasets: News, Twins, and Jobs. (i) The News dataset includes 5,000 news articles, which represented by embedding each article in a topic space and defining two device centroids (desktop and mobile). (Johansson et al., 2016) (ii) The Twins dataset is collected from U.S. twins born in 1989-199) (Almond et al. 2005) For each twin pair, we record both treatment conditions, including lighter (t = 0) and heavier = 1), and the one-year mortality outcome. To simulate the confounding bias, choose one of the twins as follows: & Bernoulli and b from (0,00.1) and a is the sigmoid function. (iii) The Jobs dataset includes experimental samples (297 treated samples and 425 control samples) and the Panel Study of Income Dynamics (PSID) comparison group (2490 control), which not only capitalizes on the internal validity of the Randomized Controlled Trial (RCT) but also draws on large-scale observational data boost in- ferential precision and enhance generalizability (LaLonde 1986). Following the experiments in (Johnnsoon et al. 2016 Machlanski et al. 2023). the News and Twins frameworks provide both factual and counterfactual outcomes for two-troaae treatmentt causal inference, so we assess methods using <PEHE and CATE In contrast, Jobs supplies only factual outcomes, which we evaluate via Roo and CATT. More details of the datasets can be found in the appendix.
+ b)), where w is drawn uniformly from (-0.1, 1) )30
Table I shows the results of different methods on the real-world datasets under the out-sample set- ting. The results under the in-sample setting are shown the appendix. We draw several interesting observations as follows.
Overall, our proposed method achieves best or highly competitive performance on all the datasets in terms of both evaluation metrics, which clearly demonstrates the effectiveness and robustness of it. These results verify the practical value and generalizability of our method in different scenarios.
Among the compared methods, OLS, BART, and T-learne train separate neural regressors for each treatment group, TARNet and DragonNet train shared representation layers and separate outcome heads. These methods suffer from the issue of covariate distribution shifts. Compared with them, G-learner demonstrates superior performance, which verifies that our method can learn effective prediction models with promising generalization ability across groups.
Compared with the generative method GANITE, G-learner leverages optimal transport to generate intermediate data. Beneficial from the geometric information involved in data extracted by optimal transport, our method achieves better performance.
The balanced representation learning method, such as BNN, CFR, and ESCFR, adjusts confounders to reduce the confounding bias, which suffers from the overbalancing issue and could discard discriminative features for potential outcome prediction. During the balancing procedure, some important information could be removed hampering the causal inference performance. Rather than adjusting the confounders, G-learner adjusts the out- come prediction models by gradually transferring the models between groups with the aid of intermediate groups. By doing this, G-learner leverages all the features to train the pre- diction model without information loss, avoiding the overbalancing issue and achieving better performance.
5.3 EXPERIMENTS ON SIMULATION DATA
Table 2 shows the results of different methods on simulation data under the out-sample setting. The results under the in-sample setting are shown in the appendix. We evaluate the robustness

---

Published as a conference paper at ICLR 2026
of each method by incrementally widening the group discrepancies, thereby emulating range of confounding biases. As the degree of confounding bias increases, all the methods show worse performance, which indicates that the distribution shift between groups affects the performance of causal inference. Nevertheless, compared to other methods, G-learner consistently achieves the best performance, demonstrating its ability preserve robust predictive accuracy and stability even as the severity of confounding bias increases.
6 CONCLUSION
In this paper, we propose a model adjustment method to lear outcome prediction models that can perform well on both control and treated groups. To achieve this, we generate a series of inter- mediate groups between the control and treated groups through the Wasserstein geodesic, which is derived from the optimal transport theory. By doing this, the control and treated groups with significant distribution shift be smoothly connected, and the outcome prediction model can be gradually adjusted between them. We further propose generated data filtering method to refine prediction model. We analyze the upper bound of the effect estimation error of our method, and conduct experiments on simulation and real-world datasets to evaluate the performance of our method. We provide alternative approach in addition covariate adjustment for addressing the issue of confounding bias in causal inference.
7 ACKNOWLEDGMENTS
This research was supported in part by National Science and Technology Major Project (0012D0111501), National Natural Science Foundation of China (62206061, U24A20233), Na- tional Science Fund for Excellent Young Scholars (621 (62122022), Guangdong Basic and Applied Basic Research Foundation (2024A151501 1901), Guangzhou Basic and Applied Basic Research Founda- tion (2023A04)1000), CCF-DiDi GAIA Collaborative Research Funds (CCF-DiDi GAIA 202521).
REFERENCES
Jonas Adler and Sebastian Lunz, Banach wasserstein gan. Advances neural information process- ing systems, 31, 2018.
Douglas Almond, Kenneth Y Chay, and David $ Lee. The costs of low birth weight. The Quarterly Journal of Economics, 120(3):0031-1083 2005.
Massih-Reza Amini and Patrick Gallinari. Semi-supervised logistic regression. In ECAI, volume 2, pp. 390-394, 2002.
Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of machine learning research, 3(Nov):463-482 2002.
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51 2015.
Jeanne Brooks-Gunn, Fong-ruey Liaw, and Pamela Kato Klebanov. Effects of early intervention cognitive function of low birth weight preterm infants. The Journal of pediatrics, 120(3): 350-359, 1992.
Hugh A Chipman, Edward I George, and Robert E McCulloch. Bart: Bayesian additive regression trees. 2010.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamoniy. Optimal transport for do- main adaptation. IEEE transactions on pattern analysis and machine intelligence, 19533 1865, 2016.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. Advances in neural information processing systems, 30, 2017.
10

---

Published as a conference paper at ICLR 2026
Richard K Crump, V Joseph Hotz, Guido W Imbens, and Oscar A Mitnik. Nonparametrie tests for treatment effect heterogeneity. The Review of Economics and Statistics, 90(3):389 405, 2008.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural
information processing systems, 26, 2013.
Rajeev H Dehejia and Sadek Wahba. Propensity score-matching methods for nonexperimental causal studies. Review of Economics and statistics, 84(1):151-161, 2002.
Vincent Dorie, Jennifer Hill, Uri Shalit, Mare Scott, and Dan Cervone. Automated versus do-it- yourself methods for causal inference: Lessons learned from a data analysis competition. 2019.
Xin Du, Lei Sun, Wouter Duivesteijn, Alexander Nikolaev, and Mykola Pechenizkiy. Adversarial brauncing-baee representation learning for causal effect inference with observational data. Data Mining and Knowledge Discovery, 35(4):1713-1738, 2021.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
Jared C Foster, Jeremy MG Taylor, and Stephen Ruberg. Subgroup identification from randomized clinical trial data. Statistics in medicine, 30(24))2867-2888) 2011.
Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learming, pp. 1050-1059. PMLR, 2016.
Chen Gao, Yu Zheng, Wenjie Wang, Fuli Feng, Xiangnan He, and Yong Li. Causal inference in recommender systems: A survey and future directions. ACM Transactions Information Systems, 42(4):1-32, 2024.
Ruijiang Gao, Max Biggs, Wei Sun, and Ligong Han. Enhancing counterfactual classification via self self-training. arXiv preprint arXiv:2112.04461, 2021.
Yves Grandvalet and Yoshua Bengio. Semi-supervised learning by entropy minimization. Advances in neural information processing systems, 17, 2004.
Jens Hainmueller. Entropy balancing for causal effects: A multivariate reweighting method to pro-
duce balanced samples in observational studies. Political analysis, 20(1):25-46 2012.
Ligong Han, Yang Zou, Ruijiang Gao, Lezi Wang, and Dimitris Metaxas. Unsupervised domain adaptation via calibrating uncertaintics. In CVPR Workshops, volume 9, 2019.
Yifei He, Haoxiang Wang, Bo Li, and Han Zhao. Gradual domain adaptation: Theory and algo- rithms. Journal of Machine Learning Research, 25(361):1-40, 2024.
Jennifer Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1):217-240, 2011.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual infer- ence. In International conference on machine learning, pp. 3020-3029, PMLR, 2016.
Fredrik D Johansson, Nathan Kallus, Uri Shalit, and David Sontag. Learning weighted representa- tions for generalization across designs. arXiv preprint arXiv:/802.088 2018.
Leonid Kantorovitch. On the translocation of masses. Management Science, 5(1):1-4, 1958.
Amirreza Kazemi and Martin Ester. Adversarially balanced representation for continuous treatment effect estimation. In Proceedings of the AAAÍ Conference on Artificial Intelligence, volume 38, pp. 13085-13093, 2024.
Kun Kuang, Peng Cui. Bo Li, Meng Jiang, and Shiqiang Yang. Estimating treatment effect in the wild via differentiated confounder balancing. In Proceedings of the 23rd ACM SIGKDD international conference knowledge discovery and data mining, pp. 265-274, 2017.
11

---

Published as a conference paper at ICLR 2026
Sören R Künzel, Jasjeet S Sekhon, Peter J Bickel, and Bin Yu. Metalearners for estimating heteroge- neous treatment effects using machine learning. Proceedings of the national academy of sciences, 116(10):4156-416) 2019.
Vitaly Kuznetsov and Mehryar Mohri. Discrepaney-baset theory and algorithms for forecasting non-stationary time series. Annals of Mathematics and Artificial Intelligence, 88(4):367-399, 2020.
Robert LaLonde. Evaluating the econometric evaluations of training programs with experimental data. The American economic review, pp. 604 620, 1986.
Qian Li, Zhichao Wang, Shaowu Liu, Gang Li, and Guandong Xu. Causal optimal transport for treatment effect estimation. IEEE transactions on neural networks and learning systems, 34(8): 4083 4095, 2021.
Huishi Luo, Fuzhen Zhuang, Ruobing Xie, Hengshu Zhu, Deqing Wang. Zhulin An, and Yongjun Xu. A survey causal inference for recommendation. The Innovation, 5(2), 2024.
Damian Machlanski, Spyridon Samothrakis, and Paul Clarke. Hyperparameter tuning and model evaluation in causal effect estimation. arXiv preprint 2rXiv:2303.01412, 2023.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Histoire de I' l'Académie Royale des Sciences de Paris, 1781.
Kamal Nigam, Andrew Kachites McCallum, Sebastian Thrun, and Tom Mitchell. Text classification from labeled and unlabeled documents using em. Machine learning. 99(2):003-133 2000.
Gabriel Peyré, Marco Cuturi, et Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 116-66):355-077, 2019 2019.
James M Robins, Miguel Angel Hernan, and Babette Brumback. Marginal structural models and causal inference in epidemiology, 2000.
Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
Donald B Rubin. Estimating causal effects of treatments randomized and nonrandomized studies.
Journal of educational Psychology, 66(5):688, 1974.
Filippo Santambrogio. Optimal transport for applied mathematicians. 2015.
Masahiro Sato, Sho Takemori, Janmajay Singh, and Tomoko Ohkuma. Unbiased learning for the causal effect of recommendation. In Proceedings of the 14th ACM conference on recommender systems, pp. 378-387, 2020.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: general- ization bounds and algorithms. International conference on machine learning, pp. 3076-3085. PMLR, 2017.
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in neural information processing systems, 32, 2019.
Jeffrey A Smith and Petra E Todd. Does matching overcome lalonde's critique of nonexperimental estimators? Journal of econometrics, 125(1-2):305-355 -353, 2005.
Jerzy Splawa-Neyman, Dorota M Dabrowska, and Terrence Speed. On the application of prob- ability theory to agricultural experiments. essay on principles. section 9. Statistical Science, pp- 465-472, 1990.
Hal R Varian. Causal inference in economics and marketing. Proceedings of the National Academy of Sciences, 113(27)7710-7315, 2016.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
12

---

Published as a conference paper at ICLR 2026
Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.
Eric Wang, Hao, Licheng Pan, Yuan Lu, Zhixuan Chu, Xiaoxi Li, Shuting He, Zhichao Chen, Haoxuan Li, Qingsong Wen, and Zhouchen Lin. Distdf: Time-series forecasting needs joint- distribution wasserstein alignment. In Proc. Int. Conf. Learn. Represent, pp. 1-24, 2026.
Fan Wang, Chaochao Chen, Weiming Liu, Tianhao Fan, Xinting Liao, Yanchao Tan, Lianyong Qi, and Xiaolin Zheng. Ce-refr: Robust counterfactual regression for consensus-enabled treatment effect estimation. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 3013-3023, 2024a.
Hao Wang, Jiajun Fan, Zhichao Chen, Haoxuan Li, Weiming Liu, Tianqiao Liu, Quanyu Dai, Yichao Wang, Zhenhua Dong, and Ruiming Tang. Optimal transport for treatment effect estimation. Advances in Neural Information Processing Systems, 36:5404 5418, 2023.
Hao Wang. Zhichao Chen, Zhaoran Liu, Haozhe Li, Degui Yang, Xinggao Liu, and Haoxuan Li. Entire space counterfactual learning for reliable content recommendations. IEEE Transactions on Information Forensics and Security, 2024b.
Hao Wang, Zhichao Chen, Zhaoran Liu, Xu Chen, Haoxuan Li, and Zhouchen Lin. Proximity matters: Local proximity enhanced balancing for treatment effect estimation. In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining 2, pp. 2927- 2937, 2025a.
Hao Wang, Zhichao Chen, Honglei Zhang, Zhengnan Li, Licheng Pan, Haoxuan Li, and Ming- ming Gong. Debiased recommendation via wasserstein causal balancing. ACM Transactions on Information Systems, 2025b.
Yuguang Yan, Hao Zhou, Zeqin Yang, Weilin Chen, Ruichu Cai, and Zhifeng Hao. Reducing balancing error for causal inference via optimal transport. In Forty-first International Conference on Machine Learning. 2024.
Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. Advances in neural information processing systems, 31, 2018.
Afsaneh Yazdani and Eric Boerwinkle. Causal inference in the age of decision medicine. Journal of data mining genomics & proteomics. 6(1):163, 2014.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International conference on learning rep- resentations, 2018
Yao Zhang, Alexis Bellot, and Mihaela Schaar. Learning overlapping representations for the esti- mation of individualized treatment effects. In International Conference on Artificial Intelligence and Statistics, pp. 1005-1014. PMLR, 2020.
Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant representations for domain adaptation. International conference on machine learning, pp- 7523-7532. PMLR, 2019.
Kailiang Zhong, Fengtong Xiao, Yan Ren, Yaorong Liang, Wenqing Yao, Xiaofeng Yang, and Ling Cen. Descn: Deep entire space cross networks for individual treatment effect estimation. In Proceedings of the 28th ACM SIGKDD conference knowledge discovery and data mining, pp- 2022.
Yang Zou, Zhiding Yu, Xiaofeng Liu, BVK Kumar, and Jinsong Wang. Confidence regularized self-training. In Proceedings of the IEEE/CVF international conference on computer vision, pp.
5982-5991,2011
13

---

Published as a conference paper at ICLR 2026
A PSEUDO-CODE OF G-LEARNER
Algorithm presents the pseudo-code of our method G-learner.
Algorithm G-learner.
Require:
factual samples (x,, f, Vi 1 number of intermediate groups K, proportion parameter
r.
1: Train the model looo predict non-treatment outcomes on the control group 140-
2: Train the model 11,1 to predict treated outcomes on the treated group [41-
3: Obtain the optimal transport plan y by solving Problem (8)
4: Generate K intermediate groups {1x) according to Eq.
5: for A. KTT K+T do
6: Select the proportion of generated samples to obtain j.--
7: Train the model (.) via Problem (12).
8: Select the r proportion of generated samples to obtain Āi1xx
9: Train the model (-))-((-) via Problem 13).
10: end for
11: Output: Outcome prediction models h-) and h1,0(-)-
B EXPERIMENT
All the experiments are run on a single 24GB GPU of NVIDIA GeForce RTX 4090 GPU. Following, we introduce more details about the experiment, including more details the datasets, formulations of evaluation metrics, results of the in-sample setting, and visualization of intermediate group gen- eration.
B.1 MORE DETAILS OF DATASETS
News The News dataset is first proposed benchmark for counterfactual inference by (Johans- son et al. 2016). The News dataset simulates counterfactual inference by modeling news articles topic distributions =(x), derived from topic model trained the NY Times corpus. Multiple centroids are randomly chosen in the topic space, where one centroid represents the control group, while the other centroids represent treated groups viewing devices (treatments). Each centroid zÿ associated with a Gaussian outcome distribution: my ~ (0.45, 0.15), oj ~ (0.1,, 0.05), from which ideal potential outcomes are sampled as Ùn [[m;,, .T) ) C, where
NO, 0.15). The unscaled potential outcomes are computed as Øj = Ūj [D(2(x), 2j) + D(2(x), 2c)], where D(.,-) the Euclidean distance, ze represents the control centroid. The treatment assignment follows t/r Bernoulli( (ernouuli(sottaxx( no bias). The true observed outcomes are scaled by a constant D = 50: y = D Øj. The dataset can simulate k treatments with v 10, enabling flexible modeling of counterfactual inference scenarios.
)), with controlling the strength of assignment bias (v = 0 implies
Twins The Twins Dataset comprises twins born in the United States between 1989 and 1991, with 30 covariates related to pregnancy, birth, and parental characteristics. In causal inference studies, the treatment defined as the heavier twin, and the outcome one-year mortality. This dataset is widely used to assess the causal effect of birth weight differences infant survival. (Almond et al.) 2005) For each twin pair, we record both treatment conditions, including lighter (t = 0) and heavier = 1), and the one-year mortality outcome. To simulate the confounding bias, choose one of the twins as follows: t ~ Bernoulli(oWw x + b)). where w drawn uniformly from (-0.1, 1)30 and from (0, 0.1), and is the sigmoid function.
The study by (LaLonde 1986) is widely used benchmark in the causal inference commu- nity, where the treatment is job training and the outcomes are income and employment status after training. 1. This dataset combines randomized study based the National Supported Work pro- gram with observational data to form a larger dataset (Smith & Todd 2005). The presence of the randomized subgroup gives way estimate the "ground truth" causal effect. The study includes
Jobs
14

---

Published as a conference paper at ICLR 2026
covariates, such age and education, as well as previous earnings. We construct binary classifi- cation task, called Jobs, where the goal is to predict unemployment, using the feature set of (Dehejia & Wahba 2002). Following (Smith & Todd 2005), use the LaLonde experimental sample (297 treated, 425 control) and the PSID 2. This dataset combines a randomized study based on the Na- tional Supported Work program with observational data form a larger dataset (Smith & Todd 2005). The presence of the randomized subgroup gives a way to estimate the "ground truth" causal effect. The study includes covariates such as age and education,
well as previous earnings. We construct binary classification task, called Jobs, where the goal is to predict unemployment, using the feature set of (Dehejia Wahba 2002). Following (Smith & Todd 2005), we use the LaLonde experimental sample (297 treated, 425 control) and the PSID comparison group (2490 control). There were 482 (15%) subjects unemployed by the end of the study.
Simulation We generate 1,500 treated samples from a multivariate normal distribution x; ~- (mmm4x1, 0.5 Σ,-,), ), and 1,500 control samples from x, ~ [(m 0.5 A.E.E where each 2 is drawn from uniform distribution Σ. ((0, m.)10x10). We fix m, = 0.5 and vary me to simulate different levels of confounding bias. The potential outcomes are generated as Y 0.5 + w (11, 0.012 r w E R 10, denotes a length-10 10 vector of ones, and 100 is the 10 x 10 identity matrix. The noise term is generated as N',, 0.1²).
{ and Yo w x ç. where the weight vector is drawn from
B.2 COMPARE METHODS
Following is the details of compare methods:
Statistics-based methods: OLS trains separate linear regression models for the control and treated groups to predict potential outcomes. BART (Chipman et al., 2010; Hill 2011) yields posterior estimates of treatment effects, supporting uncertainty quantification.
Matching Methods: PSM (Rosenbaum & Rubin 1983) estimates treatment effects by
matching treated and control units using propensity scores derived from logistic regres- sion. k-NN (Crump et al.] 2008) predicts potential outcomes using factual outcomes of the & nearest neighbors in the opposite group.
Neural Network-based methods: T-learner (Künzel et al. 2019) trains a separate neural network for each treatment group. TARNet (Shalit al. 2017) reduces distribution im- balance by learning shared latent representations of covariates. BNN (Johansson et al.) 2016) combines domain adaptation and representation learning to minimize the discrep- ancy distance in the hypothesis space. CFRwass and CFRMMD (Shalit et al.) (017) align the distributions of treated and control groups in a latent space by minimizing the Integral Probability Metric, which is implemented by the Wasserstein distance and Maximum Mean Discrepancy (MMD), respectively. GANITE (Yoon et al. 2018) leverages a generative ad- versarial network output counterfactual outcomes for ITE estimation. DragonNet (Shil el al., 2019) leverages propensity scores and targeted regularization to improve outcome prediction and stabilize treatment effect estimation. DKLite (Zhang al., 2020) learns invertible representations with overlapping support and standardizes using counterfactual variance to estimate ITE. CFR-Pro (Wang al.) (025a) enhances HTE estimation from observational data by incorporating local proximity via an optimal transport regularization and a dimension-reducinn subspace projector. CE-RCFR (Wang et al. (024a) boosts ITE estimation with relaxed optimal transport alignment and consensus aggregation, taming mini- batch noise and gradient conflict in one shot. DESCN (Zhong et al. 2022) jointly models treatment and response in the full sample space to curb bias and imbalance. ES- CFR (Wang et al., (023) leverages optimal transport to align the distributions of treated and control groups and address the unobserved confounder issue.
15

---

Published as a conference paper at ICLR 2026
Table 3: Results for in-sample performance real-world datasets in terms of mean and standard deviation. A lower metric indicates better performance. We highlight the best results in bold and underline the second best results
Mothod
VEPEНE
VEPENE
RPOL
ÉAT
ULS
4.0069 # 1.6977
0.409
+: 0.2580
BA229
+ 0.1670
0.0061
0.0099
0,2380 d: 0.0427
0.3015 + 0.0417
BART
26705 = 0.8685
0.5464
+ 0.2877
0.3208
+ 0.022
0.0034
0,0028
0.2614 # 0.030
0.0939 + 0.0452
2.7336 0.7579
0.3335
+ 0.0388
0,0221
0.0108
0.0491
0.1532
0.1382
0,0154
0,0242
0.1256
0.1263
PSM
29110 = 0.7843
0.333
+ 0.2515
0.3693
+ 000028
0.0057 2 0.0044
02246 2 0.0231
0.1183 : 0.0929
GANITE
33045 1,2562
1.7301 1 0.6204
0.3191 1 0.0211
0.0169
0.0011
0,2058
0.0516
0.1861 1 0.0804
0.1502 0.0344
TARNe
1.367 0.3708
0.3455
0.2771
0.3191
01002
0,0168
0.0011
0.2498
0.0398
0.849 :: 0.0250
DragonNer
2.641 0,7849
0.800 + 0.7010
0.3614 + 010553
0,00K
0.0064
0.1088 + 0,0923
0 + 0.0244
BNN
4.2649
1.3215
0.3194
01021
0.0433
00714 0,0229
Raa
2,892
8.7105
0,4813
0,5205
0.048
0,0400
- 0.0250
CF CFRAMME RA
218800 = 0.9814
0.9752 1 0.7153
0.5224 + 0.0038
0,0302 1 0.0183
0.2503 : 0.0418
0.0865 + 0.0253
CFR-Pro
5,3859
1,6643
0,5562 : 0.4728
0.3197 1 010023
0,0183 0.0156
0.3007 - 0.0216
0/038 : 0.01
CE-RCTR
0.1132
DESCN
4.3748 = 1.3147
2.6966
+ 0.7238
0.3174 + 0.0031
0.0149 0.0105
0,2811 + 0.0349
0.1620 0000118
ESCER
1.72M = 0.3845
0.3678
: 0.2559
0.3101
+ 0.002
0.0199 + 0.0085
0.2976
2 0.0285
0 + 0.0253
G- G-leamer
227500
0,7415
0,2258 1 0.1472
0,3191 1 0.0020
0,0107 0.0099
0,2053 0.0151
04552 : 0,0306
B.3 FORMULATIONS EVALUATION METRICS
Following (Shalit et al.
2017), we evaluate Roor. and EAT on Jobs dataset:
RPoL- = [ (x) = 1, 1]. 1(T = 1) + E| 0, = 0] p( 0)),
EATT = |([7)- Σ
- Σ vi) 171 Σ (x;) ho((a)))
iE
iECNE
SET
where for model h, if h (x)
hu(x) > X, (x))) else (x)x) 0. Besides, evaluate
ÊPEHE and CATE on other datasets:
1 E (1h))))) hh(xx))) (Ji(xs) fo(xx)))
EPEHE
n
CATE Σ (hi(1;)- ho(x:)) (((()) -fo(s:))
1
B.4 EXTERNAL. RESULT OF IN-SAMPLE SETTING
We give the results of real-world and simulation datasets under the in-sample setting in Table 3 and Table 4
B.5 VISUALLIZATION OF INTERMEDIATE. GROUP GENERATION
In this part, we conduct experiments on a toy 2D dataset to visualize the generated intermediate data. We construct a control group and treated group with means of -1 and 1, respectively, both having a standard deviation of 0.1 and covariate dimensionality of 2. The visualization results are shown in Figure 2 In Figure 2(a) we present the original covariates separately for the control and treated groups. In Figure (b) eight intermediate groups are generated via optimal transport between the control and treated distributions. These intermediate nodes lie precisely along the geodesic connecting the two groups.
B.6 SENSITIVITY ANALYSIS
In order to assess the impact of hyperparameter configurations on model performance, we take sim- ulation data where m. 1.4 example to evaluate the performance of our method with varying values of hyperparameters. Figure 3(a) and Figure ((b show the results different numbers of in- termediate groups. We observe that when the number of intermediate groups increases from 0 to 4,
16

---

Published as a conference paper at ICLR 2026
Table 4: Result simulated dataset in terms of and standard deviation. A lower metric indicates better performance. We highlight the best results in bold and underline the second best results.
Method
mhe
0.5
0.8
1.4
VEPERE
CAYE
V
VEPERK
CAPE
VEPEE
CATH
V
VEPEEE
CAY
OLS
1.13 1: 0.49
0.30
0.24
1.14
1 0.50
0.30 = 0.24
1.17 1 0.50
0.30 - 0.26
1.36
0.64
0.35 1 0.29
BART
4.18
0.58
0.20
0.18
4.30
0.50
4.67
0.43
5,25
0.43
0.44
T-leamer
8.32 + 0.82
0.40
0000-0.:0 0.32
8.88
0.80
2.77 + 0.49
10.18
+ 0.66
5.44 + 0.51
12.14 + 0.58
8.16 + 0.54
4.82 0.55
0.41
0.20
3.96
0.1
3,94
0.47
4.49
0.50
0.23
PSM
2,05 + 0.64
0153=0.00 0.60
2.76
+
0.61 0.44
4.25 + 0.97
0.92 + 0.62
5,91 + 1.11
1.99 + 0.888
GANITE
0.88 = 0.08
0.23
+.09
0.06
1,50
:: 0.01
1.51 = 0.01
0.02
1,50 0.01
DKLite
0.63 + 0.12
0.08
0.04
2.78
0.69
2.51
0.78
5.87 0.68
5.76 = 0.66
8,99 0.59
8.92 0.60
TARNeI
0.55 1. 0.05
0.54
0.05
0.55
1. 0.04
0.53
0.05
0.54 0.06
0.52 1 0.06
0.55 1.05
0.54 1 0.05
DragonNet
0.29 + 0.08
0.06
0.04
0.51
0.21
0.06
0.04
0.39 0.13
0.07 0.05
0.47 0.22
0.11-00.09
BNN
0.43 = 0.00
0.00
0.01
0.46
0.44 + 0.01
0.41 = 0.01
0.4
0.42 + 0.01
0.53 0.03
0.00
0.04
0.05
0.54 0.07
0.52
0,55
0.53 0.05
CFRMDD
0.51 1. 0.03
0.49
0.03
0.51
1 0.03
0.09
0.50 1 0.07
0.48 1 0.08
1,05 1.68
0.59 1. 0.37
CFR.Pro
0.50 0.40
0.49
0.39
0.51
0.22
0.57
0.55
0,73
CE- RCFR
0.49 1. 0.01
0.45
0.01
0.52 1 0.01
0,50
0.56 1 0.43
0.55 1 0.45
0,51 0.02
0.49 1 0.02
DESCN
0.51 0.01
0.49-0.00 0.04
0.52 0.05
0.49
0.52 0.06
0,51
0.001
ESCER
0,63 + 0.16
0.07
0.78
0.20
0.24
0.09
0.99 + 0.25
0.41 0.12
0.99 0,25
0.41 0.12
G-learer
0.21 = 0.03
0.05
0.04
0.23 + 0.05
0.05 0.05
0.30
0.07
0.09 - 0.05
0.36 0.15
0.10 0.11
(a) Original covariates
(b) Intermediate group
Figure 2: Illustration results of the original covariates (left) and generated intermediate groups (right). Blue nodes in the lower-left are control samples, red nodes in the upper-right are treated samples.
the performance of our method improves in terms of both PEHE and ATE, which indicates that the prediction models can be smoothly adjusted from one group to another, enhancing the performance of causal inference. As the number of intermediate groups increases, much more samples are gen- erated, which could amplify the influence of label noise, thereby resulting in slight performance degradation. Nevertheless, the results remain relatively stable across different values of K. Figure 3(C)and Figure ((d) present the results with different filter ratios for the intermediate data. In gen- eral, our method consistently achieves promising performance in terms of both evaluation metrics with respect to varying values of r.
In addition to simulation data, we also report the results of EPEHE on News dataset under different parameter in Figure First, in figure 4[a] we observe that when K 0. the model fails achieve satisfactory performance. In contrast, the model attains its best performance when K = 2 and K on EPEHE in in-sample setting and EPEHE in out-sample setting, as the discrepancy between the two groups is not excessively large. Therefore, relatively small value of K is sufficient to yield good results. As K continues to increase, the performance tends to stabilize; however, an overly large K may amplify the influence of noise, leading to slight degradation in the final performance. Overall, the results remain relatively stable across different values of K. Second, in figure 4b) can be observed that both EÈEEE and ÉEEНE achieve their best performance when r gradually decreases, the number of samples in the intermediate group is reduced, and the model may even be trained only real samples, resulting in suboptimal performance. On the other hand, r is too large, an excessive number of low-confidence samples will be introduced, which can also impair the model's 's predictive ability and thus degrade the performance.
0.6. As
17

---

Published as a conference paper at ICLR 2026
620
034
030
634
(a) parameter: K
(b) parameter: K
(c) parameter:
(d) parameter:
Figure 3: Results including error bar of varying values of the hyperparameters on simulation data
(m, 1.4).
(a) parameter: K
(b) parameter:
Figure 4: Results including error bar of varying values of the hyperparameters on News data.
B.7 ABLATION STUDY
The ablation study results about G-learner without OT geodesic (w/o OT) or without data filtering
(w/o filtering) are given as 3 The results demonstrate the effectiveness of the modules.
B.8 OPTIMAL TRANSPORT WITH ESTIMATED MARGINAL DISTRIBUTIONS
In this experiment, we modify the marginal distribution from uniform distributions to estimated distributions based on kernel density estimation, and report the results in Table Overall, the two versions have similar performance. Due to the estimation difficulty of the marginal distribution, the version with uniform distributions performs slightly better.
B.9 COMPUTATIONAL COST
Let no. n1, and d be the numbers of control samples, treated samples, and features, respec- tively. According to Cuturi 2003)), the complexity of optimal transport in Eq. ( is in O( + ni (onnnnnoogg + n1) ) + d)). d). Fortunately, we only need to solve one optimal transport problem regardless of the number of intermediate groups. In addition, optimal transport can be approximately solved by the Sinkhorn algorithm (Cuturi 2013). whose complexity is in O(Trini + nonid), where 7' is the number of iterations. For each intermediate group, the complexity of calculating the barycenter in Eq. (9) is O((no + nı )d). We also compare the rumning time with other eepresentation-based methods simulation data in Table Our method exhibits comparable run- time to classical representation-Ie based approaches, and outperforms certain baselines by a slight margin.
1. MORE DATASETS
We also compare our method with IHDP and ACIC Table The two datasets are widely used in causal inference (Brooks-Gum et al. 1992 Dorie et al. 2019 Our method achieves highly competitive performance on these datasets.
B.11 VISUALIZATION OF PREDICTION MODEL ADJUSTMENT
We also visualize the adjustment procedure of the outcome prediction model from the control group to the treated group in Figure We generate synthetic data one-dimensional (1D) dataset,
18

---

Published as a conference paper at ICLR 2026
Table 5: Out-sample performance comparison of ablation methods News, Twins, and Jobs
datasets (mean + standard deviation). Lower is better.
Method
News (<PEnE)
News (FA7B)
Twins (<pgnE)
Twins (CATE)
Jobs (Rpol)
Jobs (EATT)
w/o OT
0.83
0.25
0.01
0.02
0.26
0.09 0.08
w/o filtering
2.89 0.83
0.29 0.24
0.32 0.01
0.01 0.01
0.19 0.07
0.19 0.26
G-learner
2.87 + 0.87
0.25 0.20
0.32 + 0.01
0.01 0.01
0.17 + 0.06
0.06 0.07
Table 6: Out-sample performance comparison of OT with KDE on News, Twins, and Jobs datasets (mean - standard deviation). Lower better.
Method
News (tpang)
News (carg)
Twins (+pgnc)
Twins (FArg)
Jobs (Rpol)
Johs (eATT)
OT with KDB
2.87 0.78
0.24 4: 0.30
0.32 = 0.01
0.02 + 0.02
0.17 = 0.10
0.23 1: 0.05
G-learner
2.87 0.87
0.25 0.20
0.32 0.01
0.01 0.01
0.17 0.06
0.06
where the means of the control group and the treated group are set to 0 and 0.5, respectively. The outcome is generated =-² + Wiz + & where Wi (1, 0.1) and No 0.1)). We observe that the outcomes are transported from the control group to the treated group smoothly.
28 DaBalll Bornn
Figure 5: Visualization of the outcome model adjustment.
C
PROOF OF LEMMA I
Lemmal Following He et al. 2024), consider two arbitrary measures from group and x=, the error of E (ho,:) E(ho,- upper bounded with probability least r as:
E(ho.e 1 OWWp Mx ++ PC + Vlo(1/T)
[E(ho.e)
(17)
Vn
Firstly, we have the following assumption:
Assumption (R-Lipschitz). We assume H be class of regressors h X' [a, is R-Lipschitz in & norm,i.e., V h E H, Vx, x' E X:
[hx)) h(2²)l < R|2 x'112.
(18)
Assumption (p-Lipschitz Loss). We assumption the loss function & p-Lipschitz each argu- ment, i.e. Vy. yf E y:
1e(y, f(y. .)) s ply v/.
(19)
(((:)) f(sw')l plu v'1
(20)
19

---

Published as a conference paper at ICLR 2026
Table 7: In-sample performance comparison of OT with KDE News, Twins, and Jobs datasets
(mean + standard deviation). Lower is better.
Method
News (sppn))
News (cars)
Twins (epgng)
Twins (EArg)
Jobs (Rpos)
Johs (eATT)
with KDE
2.29 0.57
0.35 0.26
0.32 = 0.00
0.02 + 0.02
0.21 0.06
0.17 =: 0.08
G-learner
2.28 0.74
0.23 + 0.15
0.32 0.00
0.01 + 0.00
0.21 0.02
0.06 + 0.03
Table 8: Running time comparison of different methods.
Method
Running Time (s)
R.wwasss
280,447
Fmmmd
250.137
EsCER
261.413
GANITE
38.8825
DKLite
20.369
DragonNet
254.232
BNN
32.709
G-learner
227.049
Table 9: Results comparison the IHDP and ACIC datasets (mean + standard deviation). The best results are in bold, and the second- best results are underlined.
LHDP
ACIC
Model
Vive
EsCER
1.6676
1.1347
1,2581
1.3166
0.25564
+ 0L3886
0.4348
0.5601
B.1639 + 0.0122
0.1654 + 0.0283
0.8420
300079
01574
=: D.0076
1,2879
0.2550
1.3431
0.2154
0.2155
+ 0.156
0.3082
0.3883
0.5033
0.0544
0.5033
0.0155
0.0370
300500
0U561
=: 0.0063
CP CaRaran
1.2075 4 0.0277
1.3100
1,0927
0.2104
1: 0.1267
0.3749
0.2941
0.5034
0.0139
0.5002
0.0137
0
0J077
0J0188
0.0066
DmgonNer
3.991 4-5.20 5.20
3.2636
4.8384
1.3069
2.2908
19523
1.5499
D,4000
0.4419 + +0.1192
0.1466
0.102
0.140
0.0859
GANITE
8:2263 + 7.4778
5.5415
1,2962
7.3728
+ 1049
6.4080
3.7900
0.3276
0.0287
0,2995
0.0236
0.2366
0.020
02364
0.0029
DKLac
4.0072
# 7.440
3.728%
5.7008
0.2053
+ 0.2497
0.948
1.8252
0.1809
0.0166
0.1996
0.0084
0.0331
0.112
0.299
0.0110
G-leamer
1.8888 0.71
L1982
1,6799
0.3108
1: 0.2978
0.3738
0.3059
0.1569
0.0186
0.1899
: 0.0117
0.0087
0.0887
00298
: 0.0007
Assumption 6 (Bounded Model Complexity). We assume the Rademacher complexity (Bartlett & Mendelson 2002). R, of the hypothesis H. is bounded for any distribution j considered in this paper. That is, for some constant C 2 0,
R2 (HH:) = E
sup NEW
1 n
Σ oxx((z)
C
(21)
Vn
where the expectation is w.r.t. X; fx and a; Uniform ((-1,1)) for 1,...,n.
Proof. For any predictor h, the risk difference is
(E)(h)- - Ex (h)] - (xe(x.) (h(x) - v)2x(x)dxx - - (h(x') - s/ 2x)xp(x)dx)) (22)
Introduce coupling 7 of Hkx and Me--
- Ex-- ( 5 /1((k(=).0) e(b(2z),
< p( x'1l + 1111 v'll) (by Assumption 4 and Assumption 5 - (HxsW)(
(23)
So that we can get,
((hh.,)) - E(ho,e- )1 OWW,(H.ss PC + Vlo(1/r))
(24)
Vn
Specific details of the proof above is similar to the proof of proposition in (He et al. 2024).
20

---

Published as a conference paper at ICLR 2026
D PROOF OF LEMMA 2
Lemma
2
For a filtering ratio (). the loss E(ha,1) is upper bounded with probability at least
1-rau as
K vn
(or(1/r)
(log 1oyan
log(1 /r) nk
E(ho,1) E(h,,)) + +OKK++
K
Vnk +V
nkK
where n. is the number of training samples, is the model depth, and K&A. is the accumulated distribution shift between the control and treated groups through intermediate groups
Proof. Treat gradual self-training over K + 1 domains (each with n points) as running an online learner for n sequential examples. Directly invoke Corollary of (Kuznetsov Mohri 2020 to get, with probability r, following (He et al. 2024).
k
E(bo,1) S
gkn E x) disc(qn(RR))
k=
8 In(1/r) nk
+ +n l12 V 4n ( (e H.) + +
(25)
K+T. where k = 0....., K 1. H. denotes the hypothesis class (the set of all predictors your algorithm may choose from), and C? is the uniform upper bound on the loss function é, i.c. {(h(z).y) [0, CJ for all h H and all (x, y). disc(-) denotes the
where we index each group by &
discrepancy term as defined in Kuznetsov & Mohri 2020 and (n(k+1) is taken as 9n(K+1)
(((+11
T(K+T)
m(K+1)
From Lemma from (He al. 2024 and Lemma we know:
K
log n
qen+ie (ho E(ha,)) O(KA) + O
+
K
(26)
Vn
From Lemma 2 from He et al.
2024) and Lemma 1 we know disc(In( can be bound. From
Lemma 4 and Example 2 on neural-network sequence complexity from (He et al. 2024) we know:
nk 60 V 4T (og(nn))( (f H) S ooV (log nK)31-2
(27)
In conclusion,
+ + K log(1/r)
K Vn
(logn
log(1/r)
((ho,) s (((ho)))
V nK +1
+
nk
nk
E
PROOF OF THEOREM 1
Theorem The effect estimation error ÉPEHE is upper bounded by:
СPEHE < 2h + 2E 20(KA B(n, K,L,)))
where
K Tn
log(1/r) 1
(logmK)sLs /
log(1/r) nk
B(n, K, L,r)
+ +K
Vnk
nk
21

---

Published as a conference paper at ICLR 2026
Proof.
(PEHE = J - (xx)))
((1(()) - ho(x)) - (fi(x)
f (1h1(x) f(()) (fo(x) no (x))) (x)dxx
1 x ((11(x) - f1())² + (fo(x) hx()))² )p ( xdxx
[ (111(x) - [1(x)) (xx|] = 1)dx
/ ((h1(x)) f1(x)) p(x 0)dx
+
+ / ((1o(x) [o(x)))p(xx It- = 1)dx
X
J ((ho(x) fo(x))) (6(x)))p(x(+>) p(x 0)dx (h1,1) + E(b1,0) + E(ho,)) + E(ho,o)- For a filtering ratio r E (0, 1), the loss (ho,,) is upper bounded with probability at least
(28)
Lemma 2
as
K log(1/r)
(log (Rogn
log(1 //r(
E(ho,1) S E(ho,,)) + +OKA+ +
+ K
Vnk +V
+
vn
nk
nk
(29)
where n. is the number of training samples, T, is the model depth, and K the accumulated distribution shift between the control and treated groups through intermediate groups According to Lemma 1, it is easy to get:
K
log(1/r)
(lognR)33-22 nk
log(1 /r) nk
E(ho,1) s E(ho,,) + +OKA +
+
Vn
V Vnk
K
log(1/r)
1
(lognK)3r..2 nK
log(1/r) n
E(h1,0) + +
+
K
nk
Then, combine Eq. we can obtain the upper bound of EPERE:
ЕPEHE S E(hi,, + E(b1,)) + (bo..) E(ho,))
2h1,1 + (ho.o)
K
log(1/r)
(log K3a3-2
/lo((11))
+ +20(KA +
+ K
+
+
)
Vn
nk
n.
nk
(30)
22


Reviews:
[{"review_id": "aSj1UMn6Gv", "reviewer": "Reviewer_ttMD", "summary": "In this paper, the authors address an important limitation of existing causal inference methods that rely on balanced representation learning: the over-balancing issue that arises in outcome prediction models due to imbalances in the treatment and control groups. They propose generating intermediate groups through the Wasserstein geodesic to adjust the outcome prediction model between consecutive groups via a self-training paradigm. This approach is novel and well-motivated, and supported by both theoretical and experimental evidence.", "questions": "Please see Weaknesses. It is a very interesting work, and I look forward to the authors' responses and further discussions.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The problem addressed in this paper, the over-balancing issue that arises in outcome prediction models due to imbalances in the treatment and control groups, is indeed a significant and common challenge in the real-world application of causal inference methods based on balanced representation learning. It has not been extensively studied in existing studies, and the authors' focus on this gap is highly relevant.\n\n2. The authors' idea of incorporating optimal transport theory to address this challenge is both novel and well-motivated. By leveraging the elegant properties of the Wasserstein geodesic, the approach enables the learning of intermediate distributions along the optimal transport paths between treatment and control groups. This, in turn, allows the outcome prediction model to be adjusted between consecutive groups via a self-training paradigm. Its ability to incorporate the control group distribution into the treatment group model (and vice versa) is an excellent solution to the imbalance problem.\n\n3. The theoretical proof of the proposed method is quite thorough, providing a solid foundation for the approach.\n\n4. The authors have conducted experiments on many benchmarks and included hyperparameter analysis, effectively demonstrating the effectiveness of the proposed method.", "weaknesses": "1. The current assumption that $p_0$ and $p_1$ follow uniform (or other prior) distributions seems challenging to satisfy in real-world scenarios, where such prior knowledge may not be available. I am curious whether it would be possible to directly estimate $p_0$ and $p_1$ from the data, thereby relaxing this assumption.\n\n2. In the generated data filtering process, the authors select only the samples with high confidence. However, low-confidence samples might not only arise from model errors, but could also be inherently more difficult to learn due to their scarcity, such as in the case of long-tail distributions. I wonder whether the cumulative removal of these samples in each step could lead to the situation where such samples are never adequately learned by the model.\n\n3. The paper aims to address the imbalances in the treatment and control groups, but the baselines compared do not specifically address this issue. I suggest comparing the proposed method with other baselines that also aim to resolve such imbalances, such as DESCN [1], to provide a more relevant evaluation.\n\n[1] Zhong, K., Xiao, F., Ren, Y., Liang, Y., Yao, W., Yang, X., & Cen, L. (2022, August). Descn: Deep entire space cross networks for individual treatment effect estimation. In Proceedings of the 28th ACM SIGKDD conference on knowledge discovery and data mining (pp. 4612-4620).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VPAIUXswNs", "reviewer": "Reviewer_z8oy", "summary": "To balance the distributions between treated and control groups while reducing the loss of predictive information, the authors propose to generate intermediate groups through the Wasserstein geodesic, which smoothly connects the control and treated groups. The authors present a theoretical analysis of this method and demonstrate its effectiveness across several benchmark datasets.", "questions": "1. The paper's writing style needs improvement, especially in the abstract and introduction, which fail to clearly explain the research motivation. \n\n2. The novelty of this paper should be more clearly highlighted, enabling readers to easily understand the authors' contributions.\n\n3. The authors should compare their work with some relatively new algorithms and conduct experiments on benchmark datasets, such as IHDP, for a more comprehensive evaluation.\n\n4. This paper lacks a sensitivity analysis of the parameters to demonstrate the robustness of the proposed method.\n\n5. The authors discuss the Optimal Transport in the related work section but fail to compare with it in the experiments.\n\n6. Lacking ablation experiments to verify the effectiveness of each module.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The authors provide a theoretical analysis of the estimation error for their proposed method.\n\n2. The authors focus an important issue in causal learning, treatment effect, and provide a potential solution to address confounding bias.", "weaknesses": "1. The language quality of the submission is not very good, especially in the abstract, which fails to clearly explain the motivation behind the proposed method.\n\n2. Lacking ablation experiments to verify the effectiveness of each module.\n\n3. The authors should compare one or some relatively new algorithms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CwoPHPkxo8", "reviewer": "Reviewer_c11B", "summary": "This paper proposes a novel method, `G-learner`, to address the over-balancing issue during alignment of treated and control groups in causal inference. The authors leverage the Wasserstein space and the concept of the Wasserstein barycenter to generate intermediate domains. Models are then progressively fine-tuned on these intermediate domains. Additionally, an uncertainty quantification technique is introduced to filter data during the fine-tuning stage. Theoretical analyses are provided, and extensive experiments are conducted to demonstrate the efficacy of the proposed approach.", "questions": "Please see the weaknesses", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The use of intermediate domains to generate auxiliary samples for alleviating the over-balancing issue is both creative and effective.\n2. The research topic—improving causal inference—is highly relevant to the ICLR community.\n3. The experimental results are clearly quantified and support the claims of the proposed approach.", "weaknesses": "1. To the best of the reviewer’s knowledge, some related works have explored using the Gromov-Wasserstein distance and unbalanced Wasserstein distance to account for local similarity and mini-batch sampling effect during alignment. Including such baselines (e.g., ESCFR-Pro [1] and CE-RCFR [2]) would strengthen the comparison and evaluation.\n\n2. The paper states that the tables underline the second-best results, but these underlined results are not clearly visible.\n\n3. Figure 3 should include uncertainty/error bars to better demonstrate the robustness of the proposed approach.\n\n4. An ablation study evaluating the effectiveness of the data filtering component is recommended.\n\n5. Is Eq. (9) a variant of barycentric projection? Please clarify how it is derived from Eq. (5).\n\n6. In Figure 3, why does performance decrease as the number of intermediate groups increases? \n\n--- \nReferences:  \n[1]. Proximity Matters: Local Proximity Enhanced Balancing for Treatment Effect Estimation  \n[2]. CE-RCFR: Robust Counterfactual Regression for Consensus-Enabled Treatment Effect Estimation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2uej18sFDN", "reviewer": "Reviewer_JEDq", "summary": "The paper addresses covariate shift between treated and control groups in observational causal inference. Instead of balancing covariates via reweighting/representation learning (which can cause over-balancing and remove predictive information), the authors propose to adjust the outcome prediction models themselves. They:\n\n- Build an optimal transport plan between the empirical control and treated distributions and use the Wasserstein geodesic (displacement interpolation) to generate K intermediate “groups” µκ that smoothly connect the two.\n- Perform gradual self-training along this path: starting from a supervised model trained on one endpoint (e.g., h0,0 on control), iteratively pseudo-label the next intermediate group and fit a new model h0,κ, until reaching the other endpoint (yielding h0,1). A symmetric process is used for h1.\n- Use MC-Dropout uncertainty to filter the generated samples (keep the r fraction with lowest predictive std) to mitigate pseudo-label noise.\n- Provide generalization bounds that adapt gradual domain adaptation theory (He et al., 2024) to their regression setting, bounding E(h0,1) and PEHE in terms of endpoint errors, the average inter-step Wasserstein distance Δ and finite-sample/model-complexity terms.\n- Empirically evaluate on News, Twins, Jobs, and a synthetic suite with increasing group shift, showing competitive or state-of-the-art performance on most metrics, especially as confounding grows.", "questions": "- **Necessity of OT geodesics:**\n  How much of the gain comes from using the OT coupling versus any random or nearest-neighbor pairing for interpolation? Would performing OT in a learned latent space (e.g., via an autoencoder or TARNet trunk) further improve the realism of intermediate groups? Any preliminary results?\n\n- **Computational Cost:**\n  Could the authors please comment on the computational complexity of the proposed G-learner, particularly the optimal transport step?\n\n- **Visualization of Model Adjustment:**\n  The visualization of the generated intermediate data in Figure 2 is very effective. It would be highly insightful to also visualize the function learned by the outcome models (e.g., h_0,k) at each step k. For a simple 1D or 2D covariate space, plotting how the prediction surface smoothly deforms from the control group towards the treated group would provide a powerful visual confirmation of the \"gradual model adjustment\" concept. Is this something the authors have considered or could provide?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **Originality and framing**\n  Conceptual shift from \"adjust covariates\" to \"adjust the prediction model\" via a smooth geodesic path is a clear and appealing alternative to conventional balancing. This directly targets the over-balancing concern. Using OT-induced displacement interpolation to construct intermediate \"domains\" is principled and leverages known geometry of probability measures.\n\n- **Methodological quality**\n  The algorithmic pipeline is simple, modular, and practical: compute OT; generate intermediate groups; gradual self-training with uncertainty filtering. It can wrap around standard regressors. The uncertainty-based filtering is a reasonable, low-overhead safeguard against error amplification in self-training.\n\n- **Theoretical and Empirical Support**\n  The claims are well-supported by both theoretical analysis and extensive empirical evidence. The authors provide a theoretical upper bound on the treatment effect estimation error (Theorem 1), which solidifies the method's foundations. The experimental evaluation is comprehensive, covering multiple datasets, evaluation metrics, and a large number of strong baselines. The consistently superior performance of G-learner across different scenarios robustly demonstrates its effectiveness.", "weaknesses": "- **Scope and assumptions**\n  The approach still relies on ignorability/overlap and does not address unobserved confounding. While this is standard for meta-learners, the paper’s emphasis on avoiding balancing could be read as implying robustness to confounding; a clearer articulation of identifiability limits would help.\n  \n- **Discrete vs. Continuous Covariates**\n  The use of Wasserstein distance with an L2 cost function seems most natural for continuous covariates. It is not immediately clear how the method would perform with datasets dominated by discrete or categorical features. While the method appears to work well on the presented real-world datasets which likely contain mixed data types, a brief discussion on the applicability or potential adaptations for high-dimensional discrete covariate spaces would be beneficial.\n\n- **Experimental gaps**\n  - On News, G-learner's PEHE is worse than TARNet; the paper states best or competitive, but this discrepancy merits discussion and an ablation (e.g., how K, r, or representation choice affects PEHE on News).\n  - Dataset breadth: adding IHDP/ACIC or a modern semi-synthetic benchmark would strengthen generality claims.\n  - Computational cost: solving the OT problem can be heavy for large n.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)