Matching without Group Barrier for Heterogeneous Treatment Effect Estimation
Authors: Yuguang Yan
Haolin Yang
Shihao Zhang
Weilin Chen
Ruichu Cai
Zhifeng Hao
Venue: ICLR 2026 Poster (2026)
Keywords: causal inference
matching
Primary area: causal reasoning

Abstract: In heterogeneous treatment effect estimation from observational data, the fundamental challenge is that only the factual outcome under the received treatment is observable, while the potential outcomes under other treatments or no treatment can never be observed. As a simple and effective approach, matching aims to predict counterfactual outcomes of the target treatment by leveraging the nearest neighbors within the target group. However, due to limited observational data and the distribution shifts between groups, one cannot always find sufficiently close neighbors in the target group, resulting in inaccurate counterfactual prediction because of the manifold structure of data. To address this, we remove group barriers and propose a matching method that selects neighbors from all samples, not just the target group. This helps find closer neighbors and improves counterfactual prediction. Specifically, we analyze the effect estimation error in matching, which motivates us to propose a self optimal transport model for matching. Based on this, we employ an outcome propagation mechanism via the transport plan for counterfactual prediction, and exploit factual outcomes to learn a distance as the transport cost. The experiments are conducted on both binary and multiple treatment settings to evaluate our method.

Published as a conference paper at ICLR 2026
MATCHING WITHOUT GROUP BARRIER FOR HETERO- GENEOUS TREATMENT EFFECT ESTIMATION
Yuguang Yan', Haolin Yang', Shihao Zhang Weilin Chen', Ruichu Zh Hao
1 School of Computer Science, Guangdong University of Technology, Guangzhou, China,
2 School of Computing. National University of Singapore, Singapore,
3 Pazhou Laboratory (Huangpu), Guangzhou, China,
4 College of Science, Shantou University, Shantou, China.
ABSTRACT
In heterogeneous treatment effect estimation from observational data, the funda- mental challenge is that only the factual outcome under the received treatment is observable, while the potential outcomes under other treatments or no treatment can never be observed. As a simple and effective approach, matching aims to predict counterfactual outcomes of the target treatment by leveraging the nearest neighbors within the target group. However, due to limited observational data and the distribution shifts between groups, one cannot always find sufficiently close neighbors in the target group, resulting in inaccurate counterfactual prediction be- cause of the manifold structure of data. To address this, we remove group barriers and propose matching method that selects neighbors from all samples, not just the target group. This helps find closer neighbors and improves counterfactual prediction. Specifically,
analyze the effect estimation error in matching, which motivates us to propose a self optimal transport model for matching. Based on this, employ an outcome propagation mechanism via the transport plan for counterfactual prediction, and exploit factual outcomes to learn a distance as the transport cost. The experiments are conducted both binary and multiple treat- ment settings to evaluate our method.
INTRODUCTION
Estimating heterogeneous treatment effects from observational data has been widely applied in many real-world data applications (Hitsch et al., 2024), such as healthcare (Foster et al. (011), economics
(Heckman 2000). and recommendation systems (Sato et al. 2020: .uo et al. [024]] Gao et al.) 2024). Based on the framework of the Neyman- (tymann-Ruin potential outcome model, the treatment effect can be estimated by comparing the potential outcomes of different treatments (Splawa-Neymanl
al. 1990.
Rubin 2005). Nevertheless, we can only observe the factual outcome of the received treatment, while counterfactual outcomes under other treatments
treatment can never be obtained.
To predict counterfactual outcomes, a variety machine learning methods have been proposed (Johansson et al.) 2016; Feuerriegel et al. 2024). Among them, matching has attracted significant attention because of its simplicity and interpretability (Stuart 2010 Kallus 2020). To predict the counterfactual outcome of target treatment, classical matching identifies the nearest neighbors in the group receiving the target treatment, and then aggregates their factual outcomes for prediction [Kallus 2020). The cornerstone underlying matching the assumption that samples close in dis- tance tend to have similar potential outcomes.
However, in practice, due to limited observational data and distribution discrepancies between groups caused by the confounding bias (Greenland et al.] [999] Shalit et al.] [2017). there exist regions where samples under the target treatment are scarce or even absent, making it difficult to find sufficiently close samples within the target group. Consequently, the matched samples may suffer from large distances. Since data samples typically lie on an intrinsic manifold, where the Euclidean distance is meaningful only locally, large distances between matched samples may not
"Corresponding author; cairuichu gmail.com

---

Published as a conference paper at ICLR 2026
capture true relationships. This inconsistency weakens counterfactual prediction. In other words, matching performs well only when samples are close enough.
To address the above challenge, matching method to find neighbors from all the samples regardless of their received treatments. By be matched, which is beneficial for capturing
propose to remove the barriers between groups and design
doing this, closer samples with small distances relations between samples for counterfactual prediction. Specifically, we analyze the outcome esti- mation error of our matching method and provide an error bound in terms of the sample distances. Our theoretical result enjoys an explanation from the perspective of optimal transport, which studies how to move masses from group of samples to another group with the minimal total transport cost VVillani et al. 2009 Peyré et al. 2019). Motivated by this explanation, we propose a self optimal transport model to select neighbors from all the samples for matching.
Nevertheless, for the matched samples not coming from the target group, their potential outcomes under the target treatment are unknown, bringing challenge to counterfactual prediction. To allevi- ate this, inspired by the information propagation mechanism used in semi-supervised learning (Zhu and Ghahramani] 2002), we construct transition probability matrix based on the optimal transport plan, allowing us to employ a random walk algorithm (Xia et al. 2019) for counterfactual prediction. predicctioooo.
To preserve the relations between factual outcomes in the transport cost of our model, we introduce factual outcomes to learn distance as the transport cost within the optimal transport framework, in which the transport cost measured on covariates consistent with the optimal transport plan of factual outcomes. To evaluate the performance of our method, we conduct experiments both semi-synthetic data and simulation datasets, including both binary and multiple treatment settings. We name our method as Matching OithOut Group bArrier (MOGA), and summarize the major contributions as follows.
We propose a matching method to select neighbors from all the samples, which is formu- lated as self optimal transport model, allowing closer samples to be matched for better capturing sample relationships.
We propose a counterfactual prediction approach for estimating heterogeneous treatment effects, using outcome propagation mechanism and the optimal transport plan modeled as a transition probability matrix.
We propose distance learning method that improves causal effect estimation by leveraging factual outcomes within the optimal transport framework.
BACKGROUNDS
In this section, we first present the notations used in the paper, and then provide the background of optimal transport and heterogeneous treatment effect estimation. The comprehensive review of related work on causal effect estimation and optimal transport is provided in the Appendix AA]
Given a vector q RRN, 4 is the i-th entry. represents vector or matrix with all the entries being 1. The probability simplex EN is defined as Σx {q (R+) Σ q =}. For matrix A, AT is the transpose of A, and Aij is the (i, entry. For the probability distribution A E (R )NN, the entropy is defined as H(A) = - - Ex Asg( A15-1)
-
2.1
OPTIMAL TRANSPORT
Given the sets of probability measures P(U) and P(V) on the spaces 24 and V, respectively, and a cost function e U x V + RR. Let a P(U) and B P(V) be two distributions with the samples E l and E V. The Kantorovich problem of optimal transport aims to find the optimal probabilistic coupling 7 E P(U x V) by solving the following problem
min /
c(u, v)dr(u, s.t. 7 E T(a,8).
Juxy
where I(a, B) C. P(U x V) is the set of probabilistic couplings with marginal distributions a and 3.
2

---

Published as a conference paper at ICLR 2026
One of the advantages of optimal transport is that it can be performed without knowing the underly- ing distribution. Optimal transport can work even on discrete distributions represented by empirical samples. Specifically, for the discrete situation, given the observed samples (u and {v: he with Ra and ni, being the numbers of samples, respectively, let d(w;) (resp., 8(vx)) be the Dirac function the location u, (resp., 5(1:)). The vectors a Σ,,, and b Eng are the probability simplexes, and the i-th entry a; (resp., 6;) is the probability masses associated with the sample u, (resp., vz). Based on the above notations, the empirical distributions can be written as
á = Σ (us) B= 2=20 b.ō(vz).
(2)
im
im1
Let C be the cost matrix with the entry Ciş = clui, U2), and y the transport matrix belonging to the set
r () = f> E (R+) y1 = a,y" In. = =b), the discrete form of optimal transport reads
(3)
min (C,7) Y E (). Y
(4)
2.2 CAUSAL EFFECT ESTIMATION
Our analysis follows the Noyman-Rubin potential outcomes framework (Rubin 1974; Splawa-
Neyman et al. 1990). We denote t; the treatment received by the i-th sample, and as treatment value in the space T = (0, 1...., 7}. where 7 is the number the different treatment values, and O indicates the control group received no treatment. The samples are represented as {(x;, Vs f,) )= where is the number of samples, X; E Rd is the covariate vector with d being the number of covariates, Y/ E R is the observed factual outcome and t, E T is the received treatment. For the treatment group I, the samples are represented {(x', v/) } with n, being the number of sam- ples in the treatment group L. Further, we denote Yx(x) as the potential outcome for the specific individual given its covariates under the treatment
Our task is to estimate the heterogencous treatment effect (HTE), which captures how the impact of a treatment differs based on individual characteristies. In this paper, focus on the multiple treatment setting, thus the task is to estimate all HTES under all possible treatments. Formally, for given treatment and f', HTE defined as:
TE. -x(Y((x,) - Ye (xx))x]
=}((x;) - (xs);
where we denote nuisance function under treatment /. as fx == [x;))x Following (Yan al, |cotiia and Gutman 2019: Schwab et al.] 2018), we make the following assumptions to ensure
the identification:
Assumption (Stable Unit Treatment Value Assumption). The potential outcome of unit is unaf- fected by the treatment status of other units, and there is no variation in the treatment levels.
Assumption 2 (Unconfoundedness)) For the i-the sample, the received treatment is independent of the potential outcomes Y; (x(x) conditioned on the covariates x,, Formally, Vt E I, TT,, (x.x)) IL 4||x;.
Assumption (Overlap). For each sample, there is a non-zero probability of being assigned either treatment control, conditional on the covariates x; Formally, Vi E 0 < P(t x < 1.
T) and x,, we have
Assumption 4 (Lipschitz Continuity). Given the mapping function Φ Rdd - Rd and the norm ||: Xyا1 = [|x(x;) d(x;)ll- For any treatment & T, the nuisance function fe(.) is Lipschitz continuous with the constant L, > 0, i.e. [x,, - fx(x;) < L, x,l1a-
Assumptions 1 and are common and standard assumptions in causal inference (Rubin) 1974] Hill 2011 Johansson et al. 2016), which guarantee the identification of HTE. Here we adopt the classic assumptions under the standard identification framework. Additionally, we make assump- tion in Assumption on the function fx),, which ensures that the function ft does not change too rapidly and provides bound on how much the potential outcome varies x changes. Intu- itively, Assumption means that two close samples usually have similar potential outcomes. This assumption is reasonable and easy be satisfied in practice (Kallus 2020).
3

---

Published as a conference paper at ICLR 2026
3 METHODOLOGY
3.1 MATCHING WITHOUT GROUP BARRIER
Given a sample x" with t' * t, in order predict its counterfactual outcome under the treatment t, classical matching finds the nearest neighbors from the treatment group & {x,, ) based on distance d(x", x(), and then combines the factual outcomes of the neighbors to predict the coun- terfactual outcome g¹. The assumption underlying matching that if d(x"", xx) is small, then they have similar potential outcomes. However, this approach faces a challenge in practice that one can- not always find sufficiently close neighbors the treatment group t. This occurs when the number of observational samples is limited, or the groups t and t' suffer from a large distribution discrep- ancy because of the confounding bias. Although one can find a neighbor with a large distance, the manifold structure of data makes large distances unreliable to accurately characterize the structure of potential outcomes, resulting in inaccurate counterfactual prediction.
To address this, break the barriers between different groups and propose matching method to find neighbors from all the samples regardless of their received treatments, so that smaller dis- tances between matched samples are expected, improving the reliability of counterfactual prediction. Specifically, to predict the counterfactual outcome of the treatment t for the sample X,, we find near- est neighbors of x; from all the samples rather than the treatment group t, and then combine the potential outcomes of the matched neighbors for prediction. Without loss of generality, let Wij be the matching degree between x, and Xj, the counterfactual outcome can be estimated by the following
{x(x,) = Σ Wxy))
where E/=1 Wi, 1. Ideally, Wi, of the matched neighbors should be large, and Way of the non- matched samples should be small or close to O. In addition, Wa: is expected to be zero since the outcome estimation of z; relies on its neighbors rather than itself. The consistency analysis of the estimator discussed in Appendix B The outcome estimation error of the treatment is analyzed by the following theorem with an upper error bound:
Theorem 1. Let cy, (f=()((x)) (x,)))
be the estimation error of the potential outcomes under the trealment I, if the assumption Vi, V Var(y.xx;,1) = 1/² holds, then the error is upper bounded by the following:
ey: S 2L
Wyllo(x)) 6(x))1 2nn
W 1j
(8)
==1
==
The proof is in Appendix Kallus 2020).
The homogeneity assumption of the variance used in (Kuang et al.,
2019
Based on Theorem 1 we discuss the difference between our matching method and the classical matching method from the perspective of outcome estimation error. To predict counterfactual out- come (x,), no matter which treatment/ is considered, our matching method is able to find neigh- bors from all the samples {xy}ー=1 regardless of their received treatment {t,} On the contrary, classical matching only considers the target group {x, by = t} the candidates, which highly restricts the search space for matching. As a result, the neighbors in the other group are excluded, which suffers from larger distances between x; and the matched samples, resulting in a looser upper bound.
The upper bound in Theorem enjoys a clear explanation from the perspective of optimal transport. The first term can be modeled as the total transport cost, and the second term can be modeled as the Frobenius norm of the transport matrix, motivating us to propose a regularized optimal transport model for learning the matching degree matrix. Specifically, we define the empirical distribution of all the samples as μ = Σ px5(x;), p. Vi = 1,.... ..... n, where the uniform probability mass p: indicates that all the samples contribute equally to the estimation error. The uniform masses also prevent large subgroups from dominating small subgroups, which ensures that internal subgroups maintain influence in the matching process. By setting the probabilistic coupling as Tig 1 Wij.
4

---

Published as a conference paper at ICLR 2026
minimize the upper bound of the potential outcome estimation error in Theorem by the following optimal transport problem
m 7) + Xy s.t. r E L((,H) u)) Yi = 0, = ,,...,
(9)
7
where Cº is the cost matrix with the (i.j)-th entry being Cify lxx - o(x) 112 s277) - 2 is the square of the Frobenius norm, and Ay is the trade-off hyperparameter.
Different from the classical optimal transport model involving distributions discussed in Section
2.1 Problem (9) is self optimal transport model that considers transport from the set of samples this set while excluding moving one sample to itself (Landa al.] 202T; Yan et al. 2024). As result, for the sample x,, the neighbors found from all the samples are matched with a large weight Wiy. while the samples far away from x; are assigned with a weight Wi close to O.
The following theorem further shows that by minimizing the upper bound in Theorem the effect estimation error is also minimized
Theorem 2. Let (x;) denote the predicted outcome for the i-th sample under the treatment t.
The effect estimation error is measured by the pairwise precision in the estimation of heterogeneous
effect (mPEHE) defined as cm PEHE +T(T+1 Eose CIST Ei)(((((xx))) - (x.)) (f((x.))
fe (xx))) (Schwab al. 2018: Guo al. 2023). The effect estimation error is upper bounded by the outcome estimation error as follows
EMPEHE S
EY
(10)
n(T
t=0
The proof is given in Appendix D We observe that the bound of in Theorem is upper bounded in Theorem
Consequently, we establish theoretical connection between our matching method and optimal transport. In practice, to remove the constraints "/s 0, we follow (Yan et al., 2024) to con- struct a cost matrix C = C + LI,, where L is sufficiently large value and I, R is the identity matrix. By doing this, the diagonal entries C will induce % to be close to 0, avoiding tackling the constraints % H(y) to minimize the negative entropy of 7, so that the Sinkhorn algorithm (Cuturi 2013) can be applicd to efficiently solve the optimal transport problem. Finally, we achieve the following optimal transport problem:
explicitly. In addition, we borrow entropic regularization term
min (C,,) +y)((7)) (r) s.t. I(,
(11)
where A the hyperparameter.
Based on optimal transport model, we leverage the results of our model for counterfactual pre- diction Section 3.2, and incorporate factual outcomes to learn a distance as the transport cost in Section 3
3.2
COUNTERFACTUAL PREDICTION
Given the optimal transport plan 7 obtained by solving Problem 1T we can find the matched sam- ples and predict the counterfactual outcome Y-(-) according to Eq. 7 The optimal transport plan reflects the matching degrees used for counterfactual outcome estimation. Optimal transport will adaptively assign larger values Yi between close sample pairs, and pair far away from each other will receive quite small Vij. For outliers that are far away from most samples, optimal transport will assign small weights to them. As a result, the estimated outcomes are basically determined by close samples with large weights, and the outliers will make limited contribution in counterfac- tual outcome estimation. However, for the matched samples not come from the treatment group t, their potential outcomes under the treatment are unknown. To tackle this, we consider an informa- tion propagation mechanism to iteratively update the counterfactual predictions by a random walk method (Xia al. 2019)
5

---

Published as a conference paper at ICLR 2026
Recall that 7 [(H,, (i) is a doubly stochastic matrix, meaning that Σ Mr 1 Vi 1..... n. We can simply construct a transition probability matrix W (I+))** by setting W ny. The (i,j)-th entry Wi indicates the probability that the i-th sample moves to the j-th sample, where the probability measured based on the transport cost between them compared with the costs between other pairs. Based on this, we develop random walk algorithm to predict potential outcomes for all the treatments over all the samples.
Specifically, let Y E [ax(T+1) be the matrix including all the factual outcomes of all the samples, which is defined as
Vi for t =t t. 4 +
Yu
(12)
0 for
and M {0, 1)xx((+1) be the factual outcome mask matrix defined as
1 for = M = 0 for t #t..
(13)
At the R-th iteration, use Y 1 to denote the predicted potential outcome matrix including all the treatments and samples. We update the predicted potential outcome matrix by SÝx where S is the affinity matrix constructed as S - oW + (1 p)1, which introduces self- connections with the coefficient p (0, 1), which balances exploration (via W) and memory (via I). After that, we replace the predicted entries Y% with the known corresponding factual outcome by Yu. In summary, we iteratively update the predicted potential outcome matrix by the following random walk rule
P SS @ (1 M) + Y ( +Y,
(14)
and the initial potential outcome matrix is set as
Y =Y.
Eq. (4) contains a diffusion term S @ (1 M), which gradually propagates outcome informa- tion along the manifold to the unmasked part denoted by 1 M. mimicking geodesic aggregation and ensuring a smooth geodesic estimation (Tenenbaum et al. 2000). By doing this, the manifold structure of data is leveraged improve the causal effect estimation. Finally, the outcome informa- tion will gradually propagate to all the samples under all the treatments. The fixed observation term Y O M remains unchanged across iterations.
3.3 DISTANCE LEARNING
Based on Theorem D the estimation error of heterogeneous treatment effects relies on the cost cx(x;, x²) [xx)x; o(x,)11 which determined by p(-). In this part, we discuss how to implement the function c(-).
The vanilla approach is the identity function d(x) x, and the cost (,,, x,) ||x; - x, 1 the squared Euclidean distance, which is commonly used in existing works of optimal transport (Courty jtt al. 2017).
The key to the success of matching is to find sample with similar potential outcomes, which means that (x, can capture the difference between their potential outcomes. However, ca((,-) does not take the outcome into consideration. To enhance the distance measurement for potential outcome prediction, we introduce the factual outcomes into distance learning. In addition, since the distance adopted as the transport cost in our optimal transport model in Problem (9). we also apply the framework of optimal transport to learn a distance. Specifically, we consider sample transport within each treatment group, which involves the transport plans {7,)=-0- with the constraint that 7 T(4, 14), where the empirical distribution of one group is defined as Pr Σn px)(x ( 1 Vi 1,...,
Based on the above discussions, we first exploit factual outcomes to learn an optimal transport plan for each group, and then enforce the learned distance covariates to admit the same optimal

---

Published as a conference paper at ICLR 2026
transport plans obtained from factual outcomes. Specifically, we learn the optimal transport plan based on factual outcomes by the following problem
4. arg min (C) .)) (r,) s.t. Y ['ve Pe),
(15)
%
where the cost matrix C measured by the factual outcomes of the treatment group / is constructed Ch,es (vt nو)². Similar to the self optimal transport model in (Yan et al. 2024), we set
Ci,s = L as a sufficiently large value to avoid the trivial solution. After that, we learn the mapping function o(.) based on the optimal transport plan % by the following
T
min 1
2(C 5.).
(16)
t=
where C is the cost matrix determined by the function (6-) on the covariates of the treatment group 4. Intuitively, for the pair of i-th and j-th samples, if their potential outcomes are similar, a large mass transport Je,y will be induced, resulting in a small cost C2as" Eq. 16 encourages C to approach C.², which improves the consistency between cx(X,, and Ci,,j As a result, the outcome information is effectively captured in the learned cost (..).)). Moreover, the ordinal relation of the potential outcomes {ví f==1 is well preserved in <(x). which has been shown to compress the manifold on which ((x) lie, leading to improved generalization ability (Zhang et al. 2024).
To further introduce oi-) and (7,)75 into a unified problem, we propose the following self op- timal transport model, which considers the transport cost on both covariates and factual outcomes collaboratively, and learn the transport cost and plans jointly
+ Ny . (nH)
min (Y.)-8
Yr ((44,)4)) t. O,...,TT (17)
where A the trade-of off hyperparameter. We initialize the optimal transport plans yr based on the solution to Problem (1). and then solve Problem (17) to refine Y, shared by the cost of covariates and factual outcomes. As result, the coupled transport cost C measured on covariates can be supervised by the factual outcomes.
Now discuss how to solve Problem 17 to obtain the optimal transport plans fy, 1T and the mapping function Φ(-) involved in the transport cost. Problem (17) contains multiple blocks of parameters. We adopt the alternate method to solve the problem, during which optimize one block of parameters with the other blocks fixed.
Specifically, given the fixed mapping function e(-) and the corresponding cost matrix C. and the cost matrix C the optimal transport problem within each group can be separated and solved indi- vidually. For the treatment group I, the subproblem with respect to Y can be formulated as follows
min (C%,7d + X, (C) m) - A H7r) %
s.t. 7 E T(u,,
(18)
which is standard self optimal transport problem with the cost matrix C + AN and can be solved by the Sinkhom algorithm (Cuturi 2013).
Given the fixed transport plans {r we optimize the mapping function s(.) to obtain the cou- pled cost function. Here, we implement o6.) as projection operation Φ(x) = PT x with P E Rdxdd being the parameters to be optimized, so that the transport cost C2,s be obtained
Ci = co X,) = |
x 11.
(19)
As a result, the transport cost is directly guided by the factual outcomes, ensuring that the learned distance reflects meaningful relations of samples in the sense that close samples with small learned distance have similar outcomes. In addition, the transport cost is calculated in the supervised sub- space, which can alleviate the issue of high-dimensional data.
7

---

Published as a conference paper at ICLR 2026
To avoid the trivial solution and induce orthogonal projected features, make P follow the constraint P M (P pd
P'P=l] Based on this, the subproblem with respect to P is
given as follows
T
min C .1)
s.t. PEM.
(20)
t=0
The following proposition provides the closed-form solution this problem. Proposition 3. Let X4 R be the matrix including all the samples in the treatment group t. Problem 20 is equivalent to the following problem
T
min tr
P 1=
P
P.)
S.L.
P
(21)
P
where the matrix O₂ is constructed as
θ = 2(X,) diag(re1 7x.
(22)
The closed-form solution to this problem is obtained by the eigenvectors associated with the d' smallest eigenvalues of the matrix E(= o
The proof is given in Appendix
The pseudo-code of our algorithm is given in Appendix G
4 EXPERIMENTS
In this section, we first describe the experimental settings including the compared methods and eval- uation metrics. After that, we present experimental results and discussion on semi-synthetic and simulation datasets. More experiments can be found in the appendix, including matching visualiza- tion results, ablation studies, and sensitivity analysis. All the experiments can be run on a single 24GB GPU of NVIDIA GeForce RTX 4090.
4.1 EXPERIMENTAL SETTINGS
Compared Methods We compare the performance of MOGA with the following methods: k- NN(Crump et al., 2008) finds & nearest neighbors from the target group and then predicts the poten- tial outcome based on the factual outcomes of the neighbors. OLS/LR-2 applies linear regression with separate regression models for each treatment group. BART (Chipman al. 2010 Hill 201T) provides a posterior distribution of the treatment effects, allowing for uncertainty quantification in causal inference tasks. TARNet (Shalit et al. 2017) learns latent representations reduce the distribution discrepancy between the treated and control groups. CFR (Shalit et al.) 2017) minimizes the distribution discrepancy between treated and control groups in the latent rep- resentation space via the Integral Probability Metric, which is implemented by the Wasserstein dis- tance. We defined regularized shared layer, extending CFR multiple treatment settings. GANITE (Yoon et al.) 2018) estimates individual treatment effects using a generative model based on Generative Adversarial Networks. PSM (Rosenbaum and Rubin 1983) estimates the treatment effect by matching individuals in the treated and control groups based on the propensity score, which is predicted by logistic regression. PM (Schwab al. 2018) enhances the matching method by learning neural network estimate propensity scores within mini-batches. CP (Harada and Kashimal (021)) constructs a graph based on similarities between samples, and then applies a graph-based semi-supervised learning method for causal inference. GOM (Kallus) 2020) unifies and extends matching, covariate balancing, and doubly-robust estimation by minimizing a bias-variance trade-off under a general function norm. KOM (Kallus] 2020) instantiates GOM with an RKHS norm to achieve robust causal estimates. CEM (lacus et al. 2012) (Coarsened Exact Matching) enhances causal inference by strategically reducing the precision of covariates through data coarsening, followed by exact matching. MitNet
covariates to
treatments to have the same activation distribution in the topmost
(Guo et al., 2023) proposes to use mutual information to characterize confounding bias in heteroge- neous treatment effect estimation.
8

---

Published as a conference paper at ICLR 2026
Table 1: Result Semi-synthetic data 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.
Daaser
Newe 2
Newe-4
Nive 4
TEGA
244
2A09:
5,4112
21:41022
3.541:
191.
1,315 1479
0.911 1.500
53228 2104
2164
0.00
33
1445-:1 144%
1.93:
12.5742
03334
11,00 0.00
18588 1472
6.98L
03811 1288
2413 1-288
11.3
A1112:
TARNG:
1334 2357
1983
$.189 1.927
1.330 3.87
0393
1.533
13,99
2891 0.011
11.23
CVA
1.311-2374
12788
02466
1-700
1-28
34.177 3.20
03772 1.26
0988
2817-
2451
3.19
03744 3
2,80
3.50A
8,033
1.304
15,70
8.55
13:2042
3.75
4.034 3.23
10-36-
25221
1915
12333 3.835
17
1.111
15:034
2.48
7.416
111834
GOM
6:51-21155
1512-1-77
134
1416.4
12.05122065
4.53341
18.445..45
1.83-1177
22086
11455-1714
1653643.90
530-0065
CTM
4.773-2441
10-1242001
304
851
12.01432065
478241
940
8.200480603
9399-3356
1.3-4 1468
64016
2-53
2761
37864 3713
1.66
13:1-
3185
1915 1120
13-77-0344 366
1830-8
6311
4.N0--1 D:365
M0GH 3385-3550
1401--300
1739-
50004
3-00- 811
7-46- 7430
0-20-1130 155
198
10715-0102
1676-8
1441
8333V
1342
1368 177
7.09-
0033
2751--83
Table 2: Result on synthetic data 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.
Dataue
JLL 5.00)
m JLL 1.111
06.5 6-2
1
01230
0:318 01134
1 2154
1,4%5
3.55
01235 0-11
1 01925
1510 0.104
0.354.5 0:366
1351 2-1.92
1.900 -00333
01200 $.0.99
13449 -2 EX02
1,05
0.102
1.496
0.104
0.002
$111/2
TARNi
0U779
EUT7
1.33
5.655
3.565
01146:1 43550
1:300- m1366
1.566
5.14
5.44
1.380
1.505
B.164
1.273
CPK
1398-2 00102
0184:2 036
1192 9305
1441 6-339
0U005 03335
1.30 1001
1
0.000
2108
1.34 0.000
140% 2-8.54
1106 10323
1300 = EU1C4
P534
3164-1 01143
0.58 0.33
1 8
3::94-3:3
1957 800
33165
1172--33 8.177
0.335
5313
1.633 0.133
1664 5110
KOM
1/27 (1566
0185 0/66
8105
0U99
1223 01221
1902
5700
0.194 5323
1.34 2.052
2-0.4
10090
CEN
1 01:04
9304
0.55 01.111
13301
1516
8.140
0311 5135
1.353
5.0.36
11,00
0-084-1 8-10
1972
1.6
155
191 0011
8101 1403
0114
0089 1811
100 6017
5105
1:01 1100
16825
3004
1316 0004
BJ46 BJOS
160
135
115
0845 0.15
1381 2001T
1308 1002
8843 23 8019
1182 130
1376 01BI
0UM4 8 0003
1706 -1 BX2
Evaluation Metries Following (Guo et al. 2023), we adopt multiple metries to evaluate the per- formance of the conducted methods, including Precision in Estimation of Heterogeneous Effect (PEHE), Average Treatment Effect (ATE), and Average Mean Squared Error (AMSE). In partic- ular, for the setting of multiple treatments, we consider the pair-wise version of ATE and PEHE denoted as mPEHE and mATE, respectively. The computational details of the metries are presented in Appendix H
4.2 RESULTS ON SEMI-SYNTHETIC AND SIMULATION DATA
In this section, we present the experimental results on both semi-synthetic and simulated datasets. More details about the dataset settings can be found in the Appendix
Semi-synthetic data. As shown in Table in both binary treatments (News-2) and multiple treat- ments (News-4/8, TCGA), MOGA achieves promising performance and reliable results across dif- ferent treatment scenarios. Specifically, compared to traditional matching methods such as PSM and KNN, MOGA achieves significant improvement. This suggests that our approach effectively leverages information from all samples, leading to more accurate predictions of potential outcomes. In comparison to match- methods like PM, MOGA shows superior performance. This is be- cause MOGA simultaneously accounts for relations between neighbors and distance learning based on factual outcomes, which further enhances the quality of matching. In comparison to CP, which also employs a semi-supervised graph learning algorithm. MOGA demonstrates better performance. This be attributed to the incorporation of outcome information during the distance learning in MOGA. Compared to other methods such as CFR and MitNet, MOGA also achieves highly com- petitive performance, which benefits from the usage of the underlying manifold structure of data and information from all groups to find close neighbors.
Simulation data. The results are shown in Table 2 To verify the robustness of different strengths of confounding bias, progressively increase the mean differences between the groups simulate different intensities of confounding bias. Overall, MOGA consistently outperforms other methods in terms of /EPEHE. CATE and v VAMSE with different levels of confounding biases, demon- strating superior effectiveness and stability. With the increase of strengths of confounding biases,

---

Published as a conference paper at ICLR 2026
all methods perform worse, which is reasonable since confounding factors affect the performance of bias reduction and outcome prediction. Nevertheless, MOGA still achieves competitive perfor- mance compared with the others, which demonstrates the robustness of our method. Additionally, m controls the distribution means, and the more spread out m is, the less the group distributions overlap and the worse all methods perform. Our method can expand the matching pool by consid- cring all groups rather than only the target group as candidate samples, As result, our method still remains competitive with different values of m, which demonstrates the robustness our method.
More experimental results can be found in Appendices J K and L including visualization results, effects of distance functions and hyperparameters.
CONCLUSION
In this paper. we propose a matching method without group barriers for estimating heterogencous treatment effects. Different from existing matching that finds neighbors from only the target group, our method considers neighbors from all the samples, so that closer samples can be matched to enhance counterfactual prediction. We analyze the estimation error of matching method and propose self optimal transport model based on our analysis. We further leverage the transport plan design outcome propagation method for counterfactual prediction, and incorporate fac- tual outcomes to learn a distance as the transport cost. We conduct experiments on both binary and multiple treatment settings, and the experimental results demonstrate the effectiveness of our proposed method.
6 ACKNOWLEDGMENTS
This research was supported in part by National Science and Technology Major Project 10212D0111501), 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 (2024A 151111100111 1901), Guangzhou Basic and Applied Basic Research Founda- tion (2023A0411700), CCF-DiDi GAIA Collaborative Research Funds (CCF-DiDi GAIA 202521).
REFERENCES
Günter Hitsch, Sanjog Misra, and Walter W Zhang. Heterogeneous treatment effects and optimal targeting policy evaluation. Quantitative Marketing and Economics, 22(2):115-168, 2024.
Jared C Foster, Jeremy MG Taylor, and Stephen Ruberg. Subgroup identification from randomized clinical trial data. Statistics in medicine, 30624)) 2011.
James J Heckman. Causal parameters and policy analysis in economics: A twentieth century retro- spective. The Quarterly Journal of Economics, 1 115(1):45-97, 2000.
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 reconmender systems, pages 378-387, 2020.
Huishi Luo, Fuzhen Zhuang, Ruobing Xie, Hengshu Zhu, Deqing Wang, Zhulin An, and Yongjun Xu. A survey on causal inference recommendation. The Innovation, 2024.
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 on Information Systems, 4234):1-32, 2024.
Jerzy Splawa-Neyman, Dorota M Dabrowska, and Terrence P Speed. On the application of proba- bility theory to agricultural experiments. essay on principles. section 9. Statistical Science, pages 465-472, 1990.
Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-331 2005.
10

---

Published as a conference paper at ICLR 2026
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual infer- ence. In International conference on machine learning. pages 3020- 3020-3029, PMLR, 2016.
Stefan Feuerriegel, Dennis Frauen, Valentyn Melnychuk, Jonas Schweisthal, Konstantin Hess, Ali- cia Curth, Stefan Bauer, Niki Kilbertus, Isaac S Kohane, and Mihaela der Schaar. Causal machine learning for predicting treatment outcomes. Nature Medicine, 30(4))958-966 2024.
Elizabeth A Stuart. Matching methods for causal inference: A review and a look forward. Statistical science: a review journal of the Institute of Mathematical Statistics, 25(1):1, 2010.
Nathan Kallus. Generalized optimal matching methods for causal inference. Journal of Machine Learning Research, 21(62):1-54, 2020.
Sander Greenland. Judea Pearl, and James M Robins. Confounding and collapsibility in causal inference. Statistical science, 14(1):29 1999.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: gener- alization bounds and algorithms. In International conference on machine learning, pages 3076- 3085. PMLR. 2017.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Gabriel Peyré, Mareo Cuturi, et Computational optimal transport: With applications data science. Foundations and Trends in Machine Learning, 11(5-6):355-677, 2019.
Xiaojin Zhu and Zoubin Ghahramani. Towards semi-supervised classification with markov random fields. Technical Report CMU-CALD-02-106, Carnegie Melton University, 2002.
Feng Xia, Jiaying Liu, Hansong Nie, Yonghao Fu, Liangtian Wan, and Xiangjie Kong. Random walks: A review of algorithms and applications. IEEE Transactions on Emerging Topics in Com- putational Intelligence, 4(2):95-107, 2019.
Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of educational Psychology, 66(5):688, 1974.
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.
Anthony D Scotina and Roee Gutman. Matching algorithms for causal inference with multiple treatments. Statisties in medicine, 38(17)3113-3166 2019.
Patrick Schwab, Lorenz Linhardt, and Walter Karlen. Perfect match: A simple method for learning representations for counterfactual inference with neural networks. arXiv preprint 1rXiv:1810.00656, 2018
Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1):217-240, 2011.
Kun Kuang, Peng Cui, Bo Li, Meng Jiang, Yashen Wang, Fei Wu, and Shiqiang Yang. Treatment effect estimation via differentiated confounder balancing and regression. ACM Transactions on Knowledge Discovery from Data (TKDD), 14(1):1-25, 2019.
Boris Landa, Ronald R Coifman, and Yuval Kluger. Doubly stochastic normalization of the gaussian kernel robust to heteroskedastic noise. SIAM journal mathematics of data science, 3(1):388 413, 2021.
Yuguang Yan, Zhihao Xu, Canlin Yang, Jie Zhang, Ruichu Cai, and Michael Kwok-Po Ng. An optimal transport view for subspace clustering and spectral clustering. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pages 16281-16289, 2024.
Xingzhuo Guo, Yuchen Zhang, Jianmin Wang, and Mingsheng Long. Estimating heterogeneous treatment effects: Mutual information bounds and learning algorithms. In International Confer- on Machine Learning, pages 12108 112211. PMLR, 2023.

---

Published as a conference paper at ICLR 2026
Marco Cuturi. Sinkhom distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science, 290(5500):2319-2255 2322333,,,,,2,,,,. 2000.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamoniy. Optimal transport for do- main adaptation. IEEE Transactions on Puttern Analysis and Machine Intelligence, 39(9): 1853- 1865, 2017.
Shihao Zhang, Kenji Kawaguchi, and Angela Yao. Deep regression representation learning with topology. In International Conference Machine Learning (ICML., 2024.
Richard K Crump, Joseph Hotz, Guido W Imbens, and Oscar A Mitnik. Nonparametric tests for
treatment effect heterogeneity. The Review of Economics and Statistics, 90(3):389-405, 2008
Hugh A Chipman, Edward I George, and Robert E McCulloch. Bart: Bayesian additive regression trees. 2010.
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
Paul R Rosenbaum and Donald Rubin. The central role the propensity score in observational studies for causal effects. Biometrika, 70(1): 14-15,, 1983.
Shonosuke Harada and Hisashi Kashima. Counterfactual propagation for semi-supervised individual treatment effect estimation. In Machine Learning and Knowledge Discovery in Databases: Eu- ropean Conference, ECML PKDD 2020, Ghent, Belgium, September 14-18, 2020, Proceedings, Part I, pages
Springer, 2021.
Stefano M Iacus, Gary King, and Giuseppe Porro. Coarsened exact matching. Political analysis, 20(1):1-24 2012.
Causal inference without balance checking:
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.
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 on knowledge discovery and data mining, pages 265-274, 2017.
Insung Kong, Yuha Park, Joonhyuk Jung, Kwonsang Lee, and Yongdai Kim. Covariate balancing using the integral probability metric for causal inference. In International Conference on Machine Learning. pages 17430-1 14430-1461.. PMLR, 2023.
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances neural information processing systems, 32, 2019.
Fredrik D Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and rep- resentation learning for estimation potential outcomes and causal effects. Journal of Machine Learning Research, 23(166): 2022.
Xin Du, Lei Sun, Wouter Duivesteijn, Alexander Nikolaev, and Mykola Pechenizkiy. Adversarial balancing-basser representation learning for causal effect inference with observational data. Data Mining and Knowledge Discovery, 35(4):1713-173 1738, 2021.
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.
Nathan Kallus, Aahlad Manas Puli, and Uri Shalit. Removing hidden confounding by experimental grounding. Advances in neural information processing systems, 31, 2018.
12

---

Published as a conference paper at ICLR 2026
Haotian Wang, Wenjing Yang. Longgi Yang, Anpeng Wu, Liyang Xu, Jing Ren, Fei Wu, and Kun Kuang. Estimating individualized causal effect with confounded instruments. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 1857- 1867, 2022.
Sheng Li and Yun Fu. Matching on balanced nonlinear representations for treatment effects estima- tion. Advances in Neural Information Processing Systems, 30, 2017.
Yale Chang and Jennifer Dy. Informative subspace learning for counterfactual inference. In Pro- ceedings of the AAAI Conference Artificial Intelligence, volume 31, 2017.
Zhixuan Chu, Stephen L. Rathbun, and Sheng Li. Matching in selective and balanced representation space for treatment effects estimation. In Proceedings of the 29th ACM International Conference on Information & Knowledge Management, pages 205-214, 2020.
Donald B Rubin. Matching to remove bias in observational studies. Biometrics, pages 159-183, 1973.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pages 666-704, 1781.
Leonid Kantorovitch. On the translocation of masses. Management science, 5(1):1-4, 1958.
Leonid V Kantorovich. On the translocation of masses. Journal of mathematical sciences, 133(4): 1381-1382, 2006.
levgen Redko, Amaury Habrard, and Marc Sebban. Thcoretical analysis of domain adaptation with optimal transport. In European Conference Machine Learning and Principles and Practice of Knowledge Discovery in Databases, pages 737-753, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pages 214-223, 2017.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schölkopf. Wasserstein auto- encoders. In International Conference Learning Representations, 2018.
Gabriel Peyré, Marco Cuturi, and Justin Solomon. Gromov-wasserstein averaging of kernel and distance matrices. In International Conference on Machine Learning, pages 2664-2672, 2016.
Vayer Titouan, Nicolas Courty, Romain Tavenard, and Rémi Flamary. Optimal transport for struc- tured data with application on graphs. International Conference on Machine Learning, pages 6275-6284 PMLR, 2019.
Hongteng Xu, Dixin Luo, Hongyuan Zha, and Lawrence Carin Duke. Gromov-wasserstiin learning for graph matching and node embedding. In International conference on machine learning, pages 6932-692 6941. PMLR, 2019.
Florian Gunsilius and Yuliang Xu. Matching for causal effects via multimarginal unbalanced optimal transport. arXiv preprint arXiv:211..00
2021.
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, 2024.
Eric Dunipace. Optimal transport weights for causal inference. arXiv preprint arXiv:2109.0199. 2021.
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.
John N Weinstein, Erie A Collisson, Gordon B Mills, Kenna Shaw, Brad A Ozenberger, Kyle Ellrott, Ilya Shmulevich, Chris Sander, and Joshua M Stuart. The cancer genome atlas pan-cancer analysis project. Nature genetics, 45(10):1113-1123 2013.
13

---

Published as a conference paper at ICLR 2026
Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pages 5612-5619, 2020.
Tobias Hatt and Stefan Feuerriegel. Estimating average treatment effects via orthogonal regulariza- tion. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pages 680-689, 2021.
Robert LaLonde. Evaluating the econometric evaluations of training programs with experimental data. The American economic review, pages 604-620, 1986.
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)44156-4DD 2019.
Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects.
Biometrika, 108(2):299-319, 2021.
Kosuke Imai and Marc Ratkovic. Covariate balancing propensity score. Journal of the Royal Statis- fical Society Series Statistical Methodology, 76(1):243-266 263, 2014.
Harsh Parikh, Cynthia Rudin, and Alexander Volfovsky, Malts: Matching after learning to stretch. Journal of Machine Learning Research, 23(240):1-42, 2022.
Johannes Gasteiger, Marten Lienen, and Stephan Günnemann. Scalable optimal transport in high
dimensions for graph distances, embedding alignment, and more. In International Conference on Machine Learning. pages 5616-5627. PMLR, 2021.
Khai Nguyen, Dang Nguyen, Tung Pham, Nhat Ho, et al. Improving mini-batch optimal transport via partial transportation. In International conference on machine learning, pages 16656-16690. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2026
A RELATED WORKS
A.1 CAUSAL EFFECT ESTIMATION
In the last decades, various methods based on machine learning have been proposed for causal effect estimatiin. Most the existing methods do not consider unobserved confounders and can be categorized into three classes: reweighting, representation learning, and matching. The reweighting approach aims to construct pseudo-balanced groups by reweighting samples. Rosenbaum and Rubin (1983) estimate the propensity scores for samples and take the inverse of the propensity scores the weights. To avoid estimating propensity scores, some methods learn sample weights to minimize the distribution shift between groups, in which the shift is measured by some predefined metric, such as the difference of moment (Hainmueller) 2012) Kuang et al. (017) or the Integral Probability Metric (IPM) (Kngg et al. 2023). The representation learning approach is devoted to learning balanced representations to reduce the distribution shift between groups (Johansson et al. 2016 Shi et al,)
2019 Johansson et al., 2022). Shalit et al. (2017) trains a neural network to learn representations for minimizing the IPM between treated and control groups. [uoo et al. (2023) instead leverages the mutual information to capture the distribution shift between groups the setting of multiple treatments. Nevertheless, during the learning of balanced representations, some information highly
related to potential outcomes could be lost, suffering from the over- wver-balancing issue (Du et al. 2021
Yao al. 2018).
There are also some studies considering unobserved confounding (Kallus et al. 2018 Wang et al., 2022), which poses difficulties for all the causal inference methods relying standard assumptions. In this study, we still consider the standard assumption, including the unconfoundedness assumption. If the unconfoundedness assumption is violated, additional assumptions required in existing studies, such as the presence of instrumental variables (Wang et al. 2022), or auxiliary random controlled trial data (Kallus et al. 2018).
Matching assumes that two samples with similar covariates usually have similar potential outcomes, based on this, predict the counterfactual outcome of treatment, classical matching finds the nearest neighbors in the target treatment group and predict counterfactual outcomes based on the matched neighbors and and Fu 2017; Chang and Dy 2017 Chu al. 2020). The similarity between two samples is usually measured by the distance of covariates (Rubin 1973) or the difference of propensity scores estimated by logistic regression (Rosenbaum and Rubin 1983). [Schwab et al. (018)) improve the matching approach based on propensity scores by learning a neural network, and propose a matching method within mini-batches Kallus (2020) models matching as a problem of sample weight learning, and analyzes the estimation error under the framework of worst-case analysis.
Different from these matching methods that find matched samples in the target treatment group, we propose a novel matching method to find nearest neighbors from all the samples, so that more samples are involved improve the data efficiency, and closer neighbors counterfactual prediction. We further model our method as a self optimal transport model, whose transport cost is supervised by factual outcomes and the solution is leveraged for counterfactual outcome prediction.
be found to boost
A.2 OPTIMAL TRANSPORT
Optimal transport, which is originally proposed in (Monge 1781) and then extended by Kantorovich in (Kantorovitch) 1958 Kantorovich, 2006), seeks the best plan to move one probability distribution
into another distribution by minimizing the transport cost (Villani et al., 2009 Peyré et al. 2019). Recently, optimal transport has been widely applied in machine learning and data mining, inc luding domain adaptation (Courty et al. 2017 Redko al.. 2017), generative model (Arjovsky et al. 2017
Tolstikhin et al. (2018), structured data analysis (Peyré et al. 2016 Titouan et al. 2019 Xu et al., 2019). etc. Optimal transport is also introduced into causal inference, focusing on confounding bias reduction between treated and control groups (Gunsilius and Xu| 2021 Wang al.) 2024 Dunipace 2021). These methods usually learn weights or representations for samples to minimize the discrepancy measured by the theory of optimal transport (Li et al., 2021 Yan et al.). Different from them, we model our matching method as a self optimal transport model, which is able to find matched samples from all the groups rather than only the target group.
15

---

Published as a conference paper at ICLR 2026
B CONSISTENCY
Theorem 4. Under equation 7 assumption mild regularity conditions (Kong et al. (023) (i.e.
+ - 0). while - oc, the outcome estimation error f((4)) Yx(x,) - 0.
Proof. We analyze the outcome estimation error as follows
(f(())) - Fr(t,)) = (f(()))
Σ
Woy)(
j=1
= 11((2)) Σ Wey(e(tz)++E) +E)]
Σ
(o(l((;; - fe(z,)) Σ
W
Wiylf((a))) f((x))+
WiyEi
j=1
L Σ Wollo((a))-laa +(x))|+
Σ Wobh
(23)
Since the weights Wij are obtained via self optimal transport learning, where larger [l)(x;) o(x,) ||| corresponds to smaller Wij. it follows that for sufficiently large n, the product
We (xxx) o(x,)l II approaches zero (Kong et al., 2023). We next show that as n + co, the
sum Dj=1 W converges to zero under the regularity condition (Kong et al. 2023), i.e., + co,
Σラ E W O. We first show that as + co, its mean is 0:
W
Σ
EWW B
Σ
E(Wo] 0
(24)
where the first equality is based on Wey 11 E,, and also its variance is zero:
n
Var WoEs =
VarlWisEsl + Σ
Cou WaEe))
j=1
VarlWi; + (BIW,&,W.EEL - EIW,S)JE(WW&.))
Rep
War
Σ E[W? [[W3]Var(s. Var(s,)
2
E[W? 2 13
0,
(25)
where the third equality is based on zero mean of & and (W W Il. Ey and Eý II &x. and in fifth equality we set o² = V Var[K,], and the last equality holds due to the regularity condition. Eq. 24 and Eq. 25 together imply - 0c, Σj= WiE 0, which finishes the proof.
16

---

Published as a conference paper at ICLR 2026
C PROOF OF THEOREM 1
Theorem cy, = Eii ((fe(xx,) - (²² be the estimation error of the potential outcomes under the treatment t. if the assumption Vi, Var(yixx,,t;) = sf² holds, then the error upper bounded by the following:
ey S 2
Wiylo((;)- -xx;) 2n0
W 1j
(26)
/=1
Proof. Based on the assumptions Section 2.2 and the condition Σ}= Wij = the upper bound derived as follows:
(xf((x:) {x(x)))2 B(fi(x)) Σ (o)(f((x)) +
(x)() - -(((x)))((xx)))))))
Σ Wiy&,)²
j=1
j=1
$ 2
(o( - -f((x)))) + 2E Woy
(27)
==1
where the inequality holds because of the condition that (a b) 2n + 2b2. In the following, we analyze the two terms 22==1 Wi (fx(f((x.) - fx(x))))² and 2
< j=1 W, respectively.
For the first term in 27 we have
n W Wif(f(xx)) - f((x))))==
Σ v W -
j=1
S 2
W
Wo(f((x)) - fa(x)))3)
j=1
j=1
2
- - ft(x,))2
1=1
S 2L Σ Wyllo((x)) o(x,)
(28)
where the first inequality holds according to the Cauchy-Schhee inequality, the second inequality holds because of the Lipschitz continuity of the function ft()).
For the second term in 27
we have
2F
WisS)) <
2(
W
E
/=1
==1
2nn"
W
(29)
where according to the Cauchy-Schwarce inequality, then can simply rewrite E E E3) as 2(n
1)m?
Based on the conclusions above, we can derive:
-((x))-Y( < 2 Σ W.yll((x.) - - + 2n17
Σ
W
(30)
1=1
and Theorem Can be obtained by considering all the samples.
17

---

Published as a conference paper at ICLR 2026
D PROOF OF THEOREM 2
Theorem Let (x;) denote the predicted outcome for the i-th sample under the treatment The effect estimation error measured by the pairwise precision in estimation of heterogeneous effect (mPEHE) is defined as EMPEEE
2
Losreigt (i=(((,( (xx) (fe(x,) -
(1+1
fe (xx))) (Schwab al. 2018: Guo al. 2023). The effect estimation error is upper bounded by the outcome estimation error as follows
T
PMEEE S
CY.
(31)
n(T +1
t=0
Proof.
Σ
(P((x;) - (x,)) - (f((x,) - (e((xx))Je(x)) fe(())))
EMPEHE
nT(T + 1)
DSU'IITT i=1
Σ
Σ
(1M((x)) - f((x)) + +(fe (x,(x,))) 9
nT +1)
OSN'CIST iml
4
Σ
(DM(()) f((x))) (fx((x)) - (x:)))
nT + 1)
i=1
T
(di(x;)) -x(x)))
n(T + 1)
1= 1=1
n(T
1=
E THEOREMS REGARDING AMSE AND ATE
Theorem 5. Let AMSE = miFनI of (Ye(xi) fe be the average mean squared
error of the potential outcomes. It is upper bounded by the following
T ΣΣ
AMSE < n(T +
Wisll((x;) - -(xx)) m
Σ
W
(32)
t=0
3=1
Proof.
T nT +1)
T
AMSE
ΣΣ
(x((x,)) - 1(xx))1
n(T +1)
EYo
(33)
1=
1=0
By combining Theorem it follows directly that:
2
AMSE <
W -(x;)
mmn"" 2
W
nT +1 1)
1=
Theorem 6.
Let ATA T[TFI Eose H
2
Er (",x))) Yx(x:)) - (f(()))
fe (xx))) be the error of the average treatment effect. is upper bounded by the following
2
EMATE n( +1
LL
W Wi Vill((x)) xxxx)) mn Σ
WyD.
(34)
18

---

Published as a conference paper at ICLR 2026
Proof.
2
Σ OSUKIST
(M((s)) (x:)) (1(x))) - (x()))
CMATE
T(T + 1)
T
1=1
2
((P{(x;) (x,)) - -(e(xx)) - Sx.(xx)))
T(T +)
n
OSUCIST
i=l
$
2
Σ
(|V((x)) - f((x)) + (xx)) (x((x)))
(T 1)
t=1
T
2
[((x;) f(((x))
n(T + 1)
1=0
T
2
Wo)(((s)) +5,) Σ (xyf((x))
n(T +
1= i=
==1
T
Wo)l((s;) - f((x)) + EG)
n(T +
1)
2
ΣΣ
Wo) f((x)) + - n
WiEE,))
n( + 1)
j=1
5=1
2
Wold(x;) o(x;)l + n
Wyyl)
n(T +1)
F PROOF OF PROPOSITION 3
Proposition 3 Let X, E TR.x xd be the matrix including all the samples the treatment group t.
Problem 20
equivalent the following problem
min
P F
P
(35)
1=
where the matrix θ is constructed as
+ = 2(x)) diag(Tre1 - Xt,
(36)
The closed-form solution to this problem is obtained by the eigenvectors associated with the d' smallest eigenvalues of the matrix Σ O
Proof. First of all, we implement of-) as a projection operation Φ(x) = PTx x with P E Idd are
the parameters to be optimized. Beside, set Cl denotes |P X; P Xzy 11.. After that, the transport cost between and X, can be rewritten as:
Ti Ri
(CP ..)
Clay 7eu
P X; PT X 11
1=1 1=1
= 2
| x;
2
(PT (Txx,) x, TT PTx))mas
1=1
=((XP)) (XP) ddag(y1)) - (X.P)) .7))
2tr((P x, (diag(y1) - -))
-tr(P -tr(TT
(37)
19

---

Published as a conference paper at ICLR 2026
Based on this, the objective function of Problem (20 can be written
T
Σ 71) = Σ tr ( - tr
P)
P
O.)
(38)
t=0
The solution to minimize this objective is the cigenvectors associated with the d' smallest eigenval- ues of the matrix > 1 O
G PSEUDO-CODE OF MATCHING WITHOUT GROUP BARRIER (MOGA).
Algorithm D
presents the pseudo-code of our method MOGA.
Algorithm Matching without Group Barrier (MOGA)
Input: Data samples {(x;, Mr, 4)
1: Initialize Y/ by solving Problem
2: loop
3: Update according to Proposition 3
4: Update " by solving Problem (88).
5: end loop
6: Construct the matrix C based on Eq. 19
7: Obtain matching matrix by solving Problem
8: loop
9: Update outcome matrix based on Eq. 14
end loop
For Algorithm| [ n and n, be the numbers of all the samples and the samples in the treatment group f, d and d' be the numbers of the features before and after projection. For distance learning, the complexity of Step is O + n,d + ) the complexity of Step 4 (n, For optimal transport matching, the complexity of Steps 6 and 7 is O dd). For counterfactual prediction, the complexity of Step 9 is Onn?T), where T is the number of different treatment values.
The overall space complexity of the Algorithm T dominated by the storage of the N x N tran- sition probability matrix W, which involves two steps: optimal transport and label propagation. First, during the calculation of optimal transport, the primary memory requirement is for simulta- neously storing the input data and the probability matrix W, whose complexities are ((Nd) and O(N2), respectively, where is the feature dimension. Second, the total space required during label propagation accommodates the largest input matrix W E N along with the label matrix Y E (0, 1 N xT' whose complexities are O(N2) and O(NT), respectively, where < N. In summary, the algorithm's dominant space complexity is O(N2 + Nd).
H EVALUATION METRICS
To evaluate the performance of the conducted methods, we follow (Schwab al, 2018 adopt the following metrics
2
CAPEHE = nT(T +1)
(()( -fe
(39)
USUCIST'R im1
CMATE T(T +1)
2
(f(())) Fx(x:))
-
(40)
Besides, we also add a metric:
T ΣΣ (f(())) f((x)))
AMSE
(41)
n(T + 1)
1= tm1
20

---

Published as a conference paper at ICLR 2026
1 DATASET SETTING
News The News dataset is first proposed as a benchmark for counterfactual inference by byJohansson et al. (0166) and is used in the multiple treatment setting in (Shwwb al. (2018). The News dataset simulates counterfactual inference by modeling news articles as topic distributions =(x), derived from topic model trained on 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 2j associated with Gaussian outcome distribution: my (0.44, 0.15), Jj (0.1.,00 0.05), from which ideal potential outcomes are sampled as Új ~ N (mm,. (,) 0j) + 6, where ~ (O 0.15). The unscaled potential outcomes are computed as y ú3 (D(2(x), 2,) (x=(x), =e)]. where -)::-) is the Euclidean distance, and &c represents the control centroid. The treatment assignment follows t|r ~ Bermouli(softmaax// with controlling the strength of assignment bias (v = 0 implies no bias). The true observed outcomes are scaled by a constant D = 50: Vj = D Øj. The dataset can simulate = 2,4,8 treatments with = 10, enabling flexible modeling of counterfactual inference scenarios.
The Cancer Genome Atlas (TCGA) project collects gene expression data from 9,659 in- dividuals with various types of cancer, incorporating 20,531 covariates Weinstein et al. (2013). The dataset includes three clinical treatment options: medication, chemotherapy, and surgery. To esti- mate the risk of cancer recurrence following any of these treatments, a synthetic outcome function, the dose-response curve, was applied using real-world gene expression data. The modeling of out- comes follows the method described by [chwwa et al. (2020), where the treatment assignment bias coefficient is set to 10. Furthermore, to evaluate the robustness of the model, we artificially in- troduce Gaussian noise to the outcome variable to simulate random perturbations in the experiment. Specifically, for given sample i, the observed outcome y/- V/,,, & where & ~ NNO0, o²) and we set a = 5.
TCGA
Simulation data Following Yao et al. 2018 Hatt and Feuerriegel 2021
we generate synthetic
data for four treated group and one control groups by sampling features x from Gaussian mixture distribution
W1 - (m, Σ) + u2 (2mm, E). where m [m.... ml] E Rd and Σ 0.5- (2rand -. ). with Laad ~ [(0, (dund) The outcomes are modeled as y y= sin (w, x)- exp(cos( [w] (x @ x)))+ where W1, W2 ~ 2d((,, 1)dxk) are random weight metrices, k represents the total number of treatments and control groups, and & (o 0.5) represents noise. We vary the value of m across different groups to explore the performance of the conducted methods under different data conditions.
J MATCHING VISUALIZATION
We conduct an experiment on simulation data to visualize the matching results of our method. The data consist of three treated groups and one control group, each of which includes with 20 samples with 25 features. Similar to simulation data generation in Section I we generate syn- thetic data follows. Let x Wi ~(m,, 0.5E E Emma) + W2 (2m., 0.5Eram where m = [1.5, 1.75, 1.0, 0.75] and Erand ~ U((0, [1.2, 3.4, 0.6) ddxd). We learn the projection matrix P to map data into 2D space, and show the matching results Figure We find the nearest neighbors for the objective samples, the dotted lines represent matched samples, and the color depth of matched samples represents the matching degree. We observe that samples in different groups are matched, and closer neighbors have darker colors, which means that they contribute more to the prediction.
K
MORE RESULTS
K.
DIFFERENT DISTANCE LEARNING
Besides the distance learned in Section the following methods: For the squared Euclidean distance, we have ca(x;, xj) = ||x: - Xý / For the Cosine distance, we measure the distance without considering the scale covariates, which given as X,)
also calculate the distance between x; and xy using
x:
2 2-2 xUx, We take the TCGA dataset as an
[x
21

---

Published as a conference paper at ICLR 2026
0.2
0.1
-.1
A6
-22
-022 -0.1
0.0
0.1
0.2
0.3
Figure I: Visualization results of our matching method MOGA. Hollow points in different colors and shapes represent different groups, while solid black stars denote the objective nodes. The figure displays the five matched samples with the top matching degrees, with darker colors indicating higher matching degrees.
example compare different distances used in method, and report the results in Table 3 We observe that MOGA takes the outcome into consideration through optimal transport, thus achieving the best performance.
Table 3: Results of different distances on TCGA dataset.
(CEEE
CATE
VAMSE
Euclidean
12,5989 0.0263
6,0822 0.0419
9.7808 0.0203
Cosine distance
11.1576+ 0.0452
4.1404 0.0892
8.3178 0.0416
MOGA
10.5965 0.0263
2.7850 0.0752
7.7511 0.0407
0.0407
K.2 COMPARISON WITH TRADITIONAL MATCHING
We also consider a variant of our matching method, which leverages the distance learned in Section 3.3 but selecting neighbors within the target group only. We take the News-4 dataset an example and report the results in Table 4 We observe that MOGA performs better which demonstrates the advantage of our matching method considering all the samples regardless of their received treat- ments.
Table 4: Results of different matching methods on the News-4 dataset.
/EPEHE
EATE
V VAMSE
matching only within the target group MOGA
8.0039 + 2.2043
2.9041 + 1.7250
8.0388 + 2.5587
5,9601 + 1.1798
1,1551 0.7063
4.4197 + 0.9348
K.3 PERFORMANCE OF REAL-WAL WORLD DATASET
To compare the performance in the Real-world dataset, we report the results of CATE on the Lalonde dataset[LaLonde 1986) in Table The LaLonde dataset consists of experimental and observational components. The experimental part comes from the National Supported Work (NSW) randomized controlled trial, while the control group is replaced by observational data from the Panel Study of
22

---

Published as a conference paper at ICLR 2026
0.77
6.0
$6000
2470
0.00
2.00
(a) Ay
(b) parameter.A/
(c) parameter.Ap
62-
(d) parameter.p
Figure 2: Results of varying values of hyperparameters on News-4 dataset.
Table MAE results on Lalonde dataset (mean + standard deviation). Lower is better.
Method
MAE of Lalonde
MitNet
393.9513 + 87.6213
CFR
495.3089 + 59.5542
PSM
728.7922 + 500.1740
KNN
275.2712 + 148.3843
GOM
434.4803 + 231.6803
KOM
438.3496 + 255.4207
CEM
796.7605 + 610.0612
MOGA
246.1526 = 86.9842
Income Dynamics (PSID) dataset. The treatment indicates participation in a job training program, and the outcome is earmings in 1978. The dataset widely used for benchmarking causal inference methods. We observe that our proposed method still achieves promising results.
K.4 MORE BASELINES
We also compare with more basclines on the News data. S-learner is a single unified model trained with the treatment indicator input feature to directly estimate potential outcomes under differ- ent treatments and their difference. T-learner uses two separate models trained for the treated and control groups, and the individual treatment effect is obtained taking the difference between their predictions. X-learner (Künzel et al. 2019) estimates ITE via cross-fitting of imputed treatment ef- fects and propensity score-based weighting. making it especially effective under treatment-contro imbalance. R-learner (Nie and Wager 202T) formulates ITE estimation through orthogonalized residual regression by removing the effects of covariates on both the outcome and treatment, yielding robustness to confounding. DragonNet (Shi et al.) 2019) leverages propensity scores and targeted regularization improve outcome prediction and stabilize treatment effect estimation. CE-RCFR (Wang et al., 2022) enhances ITE estimation through relaxed optimal transport alignment and con- sensus aggregation, simultancously mitigating mini-batch noise and gradient conflicts. CBPS (Imai and Ratkovic 2014) estimates propensity scores while directly optimizing covariate balance be- tween treatment groups, improving causal effect estimation. MALTS (Parikh et al.] 2022) a matching method that leverages learnable distance metric optimize feature- ceature-space similarity and accurately estimate ITE. The results are listed in Table 6 We observe that our method achieves promising performance
23

---

Published as a conference paper at ICLR 2026
Table 6: Results News dataset (mean + standard deviation) on more baselines. Lower is better.
News
News
News
V.P.M.
V
Vi.PPMME
V.ASSS
VANSE
DragonNet
9,408
2.524
1.567
6.969
10.359 4.153
3.757
3,960
3,015 4: 2.611
16,904
5.657
8,503
18.367
6.896
CE RCTR
2.397
1
1.561
6.972
1.876
9.453
2.791
1.179
3.017 4: 1.736
10.545
1.429
4.057
10.506
1.855
R-lemmer
9.336
2.967
1,310
1.531
6.902
22.276
9.665
2,048
1.118
1.654
5.619 3: 2.212
10.400
1.282
3.867
00.05
11,005
1.243
S-Jeamer
8.646 = 2.574
1.533-4
1.468
6.385 6: 1.743
8.933 +: 1.545
2.591
=: 1.282
1.645 :: 1.747
10.330
1.313
3,000
00951
10.577
1.670
Tleamer
8.581 = 2.246
1.581
1,512
6.258 4: 1.725
8,800 1: 1.545
2,658
=: 1.273
1.774 4: 1.744
10.284 =: 1.297
3.914
0.927
10,502
1.645
Xelmmeer
5,577 2: 2.2770
1.574
1,510
6.329 4: 1.732
8.825 1: 1,642
2.661
=: 112755
1.789 1: 1.748
10,250
=: 1,293
3.922
0.925
10,570
1,643
CBPS
10,645 4: 2,847
2,572
1 2.439
7.525
4. 2,013
13.214
2.679
2.891
=: 1.766
10,929
2,079
15.817
1.987
3.457
0.974
14,545
2.355
MALTS
10,692
3: 2,065
1,553
1.436
7.911
1: 2,025
10,382
1.995
2,658
1.273
8,71.
2: 1.833
13.123
1,390
3.278
1.439
10.224
2.435
MOGA
5.001 1.693
0,449
0,344
3,591
1: 1.197
5.90
1,180
1.155
:: 0.706
4,820
3: 0.935
8.904
1.214
2.386
0.725
7.819
1.212
Table 7: Comparison between two-stage and MOGA on News Dataset (mean = standard deviation). Lower better.
Dunaee
News
News-4
Meeric
VALSE
VALSE
Fearr
VAMS2
Two stage
5,218
1.079
0.741
0.027
3,999
0.758
6,126
1.887
1,300
0.247
5,963
0.791
9,091
0.841
2,788 a:
0.487
8,257 0.687 7.819 1.212
MOGA
3,081
1.693
0,449
0.344
3,591
1.197
5,900
1.180
1,155 0.706
4.420
0.938
8,904
1.214
2.386
0.728
K.5 TWO-STAGE METHOD ON DISTANCE LEARNING
In this experiment, modify the distance learning method in Section 3.3 to a two-stage approach. In the first stage, only the optimal transport matrices T/2 are calculated. In the second stage, with Te fixed, the matrices θ are computed and then used to derive the mapping function parameterized by the matrix P. Table [ shows the results of the two-stage strategy and the original joint learning method. We observe that our proposed joint learning strategy achieves better performance compared with the tto-stage strategy.
K.6 PERFORMANCE OF DIFFERENT NUMBERS OF TREATMENTS
We provide the performance and running time results under different numbers of treatment T in Table g Our method maintains competitive performance as the value of increases. Although the running time grows, our method still has modest running time. If T is extremely large, we can accelerate the distance learning stage by considering a subset of samples, and accelerate optimal transport by some fast algorithms (Gasteiger al. 2021 Nguyen al. (022).
L SENSITIVITY ANALYSIS
We take News-4 an example evaluate the effects of the hyper-parameters in our model. Figure 2 shows the results in terms of mPEHE with varying values of the hyperparameters Xy. Ay, Ag- and p. From Figure ()) the performance decreases with large An since large Å will induce uniform transport plan y which cannot reflect different matching degrees based on the distances between samples. Fig
that a large X helpful to achieve better performance, which demonstrates the advantage of incorporating factual outcomes for distance learning. From Figures we observe that our method performs stably in wide range of values of Xy and p.
2(b) and 2(d)
We also conduct ablation studies by setting the value of λu. Xy or X,y 0. The results also shown in Figure 2 We observe that the performance with A = 0 and X, = 0 is worse compared with non-zero values, which verifies the effects of the entropic and Frobenius regularizations in Problem II A > 0 achieves better performance compared with that of A = O, which demonstrates the effectiveness of introducing factual outcomes for distance learning.
24

---

Published as a conference paper at ICLR 2026
Table 8: Results under different T Settings (mean + standard deviation). Lower is better.
[wup
8-003:53
C3R
4550
6.2337
2.313
694.76
0.512
0.1004
1,831
2.112
2511.334
100
133
1961
135% ma85e
11788
0395 60033
1.833 -0991
14998
16099 4-01751
0335
(7761
60M
00:90
1080
17808 -80m
9.770 04877
17115 -08700
--80175
6-430
5405
1,292
88446
24536 à 00655
1/687 0 0.267
19-77
1/085 8.415
03320 0U6SS
1,802 L 00112
312.32
13059 1 9.855
0.300 - 03485
1.335 1 (0062
350.59
Table 9: Computation Time Comparison (in milliseconds).
Method
Calculation Time (ms)
k-NN
26.61
OLS/LR-2
1528.69
BART
1806.28
TARNet
2094.19
CFR
2415.28
GANITE
20477.37
PSM
146.81
PM
5565.81
CP
1619.49
MitNet
29507.23
MOGA
4028.86
M VISUALIZATION OF OPTIMAL TRANSPORT PLAN
Figure is the visualization of the optimal transport plan in Eq. Similar to simulation data
generation in Section we generate data m 00.1, 0.2, 0.3) and Eaan ~ ~(((0, [1, 1.5, 2] ydxd).
We observe that is meaningful and far away from a uniform coupling.
Optimal Transport Plan
1406
10
Figure 3: Visualization of the optimal transport plan. Lighter colors indicate larger weights, while darker colors represent weights closer to zero.
N RUNNING TIME RESULTS
We also provide the running times of representative methods in Table 9 We observe that the calcu- lation time of our method is comparable to other methods.
25


Reviews:
[{"review_id": "SkgV2NHJDS", "reviewer": "Reviewer_VXGC", "summary": "The paper introduces a novel approach for estimating heterogeneous treatment effects (HTE) using matching without group barriers. In traditional matching methods, neighbors for a treatment group are selected solely from within that same group. However, the authors argue that such methods are limited by distribution discrepancies and scarce samples in some treatment groups. To address this, the paper proposes a method that removes the barriers between treatment groups, selecting nearest neighbors from all available samples, improving counterfactual predictions. This method is modeled as a self-optimal transport problem, where the transport cost is measured using the distance between samples and factual outcomes. This new method, called Matching withOut Group Barrier (MOGA), is demonstrated to improve prediction accuracy, particularly in both binary and multiple treatment settings, with extensive experiments on synthetic and semi-synthetic datasets.", "questions": "Refer to Weaknesses above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tInnovative Approach: The method removes the group barrier in matching, which is a common issue in traditional causal inference methods, allowing for better matching of similar samples from different treatment groups.\n\n2.\tTheoretical Foundation: The method is supported by a robust theoretical analysis, providing error bounds and drawing a connection with optimal transport.\n\n3.\tComprehensive Experiments: The paper presents experiments on a variety of datasets, including semi-synthetic data and real-world datasets like TCGA, comparing the proposed method against several traditional methods (e.g., k-NN, PSM, BART). The experimental results demonstrate that MOGA outperforms traditional matching methods, as well as other advanced models like BART and TARNet, in terms of precision in estimating heterogeneous treatment effects.", "weaknesses": "1.\tComplexity: The method involves solving optimal transport problems and learning transport plans, which may be computationally intensive, especially for large datasets. The algorithm requires significant computational resources (e.g., a high-performance GPU) .\n\n2.\tDependence on Hyperparameters: The model's performance is sensitive to hyperparameters like λh, λf, λy, and ρ. While the authors provide sensitivity analysis, the method might require careful tuning depending on the dataset .\n\n3.\tAssumptions: A key assumption of the proposed method is that samples from different treatment groups can be meaningfully matched, which may not hold in real-world scenarios. In the presence of unobserved confounding, where latent variables affect both treatment assignment and outcomes, the matching may fail to balance treatment groups, leading to biased estimates. Moreover, if treatment effects vary significantly across subgroups, the assumption of similar effects for matched samples breaks down, as heterogeneous treatment effects may not be captured accurately. Finally, substantial distributional shifts between treatment groups can further hinder meaningful matching, particularly when groups are non-exchangeable due to underlying differences. Thus, the method's effectiveness is highly contingent on these assumptions, which may not always hold in practical settings.", "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": "BBQIyO1WlU", "reviewer": "Reviewer_7WJ5", "summary": "This paper addresses a fundamental challenge in heterogeneous treatment effect estimation—the unobservability of counterfactual outcomes. Targeting the issue of inadequate matching quality caused by within-group constraints in traditional matching methods, it proposes a Matching method named as Matching withOut Group bArrier (MOGA). MOGA introduces a matching framework that transcends inter-group barriers. While conventional matching only seeks nearest neighbors within the target treatment group, MOGA innovatively allows matching across all treatment groups. By optimizing matching weights through a self-optimal transport model, it expands the effective matching pool and mitigates the problem of excessive matching distances caused by data distribution discrepancies or insufficient samples. Unlike traditional matching that relies solely on covariates, MOGA incorporates factual outcomes into distance learning, ensuring that the matching distance aligns with outcome relevance. Furthermore, it enhances the smoothness of counterfactual predictions by disseminating outcome information via random walks and leveraging the manifold structure of the data, thereby balancing covariate similarity with consistency in outcome patterns.", "questions": "1. The optimal transport models relied upon by MOGA typically assume that data distributions satisfy smoothness or convexity assumptions. If the data are high-dimensional, sparse, non-convex, or contain outliers, might the transport plans introduce biases that affect the stability of counterfactual predictions?\n\n2. The experiments are primarily based on semi-synthetic and simulated data. Has MOGA been validated on real observational data? If so, do metrics such as PEHE and ATE align with those obtained from synthetic data? If not, how can the method's generalization capability be ensured in real-world scenarios with unobservable counterfactuals?\n\n3. The paper claims MOGA exhibits enhanced stability in multi-treatment settings; however, it does not specify the exact number of treatments evaluated (e.g., K=3, 5, 10) nor the associated computational complexity as K increases. Could you provide performance comparisons across different K values and elaborate on scalability optimization strategies for scenarios where K > 10?\n\n4. Could you provide comparative results benchmarking MOGA against more advanced methodologies, such as DragonNet, R-learner, and Meta-learners?\n\n5. MOGA expands the matching pool through cross-group matching, which may implicitly impose stronger exchangeability requirements. Has sensitivity analysis  been conducted to quantify its robustness to unobserved confounding?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "In terms of originality, this paper addresses the limitations of traditional within-group matching methods by proposing the MOGA framework, which enables cross-group matching. It integrates an optimal transport model with an outcome propagation mechanism, thereby forming a methodological adjustment to the existing matching paradigm. Regarding quality, the method is designed based on theoretical foundations such as the error bound of optimal transport optimization. Experiments conducted on semi-synthetic and simulated datasets, benchmarking against baseline methods like k-NN and PSM, demonstrate its relative advantages in metrics such as PEHE and ATE. The contribution of individual modules is further validated through ablation studies. In terms of clarity, the paper exhibits a logically coherent structure, with comprehensive explanations of technical details (e.g., the optimal transport formulation) and well-defined key concepts (e.g., cross-group matching), ensuring strong comprehensibility. As for significance, the study addresses the inadequacy of within-group matching in estimating heterogeneous treatment effects from observational data by offering a solution that broadens the matching pool. It enhances counterfactual prediction accuracy in multi-treatment scenarios, thereby holding referential value for methodological applications in the field of causal inference.", "weaknesses": "1. Insufficient Verification of Theoretical Assumptions.\nThe core of MOGA relies on the \"self-optimal transport model\" and the \"outcome propagation mechanism.\" However, the paper does not adequately justify the rationality and applicability boundaries of its key assumptions. Optimal transport (OT) models require the transport cost (i.e., the distance between samples) to satisfy mathematical properties such as convexity and lower semicontinuity, while also assuming a certain level of smoothness in the data distribution. The paper does not explicitly clarify whether real-world data meet these conditions, nor does it verify whether the OT model may fail (e.g., by concentrating matching weights on a few samples) when the data exhibit high-dimensional sparsity, non-convex distributions, or outliers, thereby compromising the stability of counterfactual predictions.\n2. Practical application is questionable. \nThe computational complexity of optimal transport models is typically . Although the paper mentions the use of approximation algorithms for optimization, it does not report specific computational time or memory usage. Furthermore, in multi-treatment scenarios, the dimensionality of the transport matrix expands from  to , leading to an exponential increase in complexity. The paper fails to provide optimization strategies or experimental validation for such scenarios, raising doubts about the scalability of the method in practical applications.\n3. Insufficient Experimental Validation.\nThe paper only compares MOGA with k-NN, PSM, and TARNet, without including more advanced methods in the current HTE field (e.g., DragonNet, R-learner, or the Meta-learner series) or improved matching-based methods (e.g., Covariate Balancing Propensity Score). This makes it difficult to demonstrate whether the \"significant advantages\" of MOGA stem from methodological innovation or baseline selection bias. Additionally, the experiments primarily rely on semi-synthetic and simulated data. Semi-synthetic data, often generated based on known models, may fail to capture real-world challenges such as high-dimensional nonlinear confounding and sample selection bias.", "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": "8k19x0xcXS", "reviewer": "Reviewer_MkPW", "summary": "This paper proposes MOGA (Matching without Group Barrier), a novel method for estimating heterogeneous treatment effects (HTEs) under the potential outcomes framework.\nThe key innovation lies in establishing a theoretical and algorithmic connection between Sinkhorn-regularized optimal transport (OT) and counterfactual effect estimation.\nInstead of restricting matching to samples within the same treatment group, MOGA finds neighbors across all groups, enabling more reliable counterfactual estimation even when group distributions differ.\n\nThe method further introduces a random walk–based outcome propagation procedure to iteratively diffuse factual outcomes over the learned sample graph, achieving global consistency of counterfactual prediction.\nAdditionally, MOGA learns a distance function via self-supervised optimal transport driven by factual outcomes, aligning the feature space with outcome similarity.\n\nExperiments on semi-synthetic and simulated datasets show consistent improvements in PEHE, ATE, and AMSE metrics over a broad range of baselines.", "questions": "- How would the method perform on real-world biased observational datasets, or what challenges prevent such evaluation?\n- How does the approach address intra-group bias or unequal propensity distributions that may exist within treatment groups?\n- The Sinkhorn regularization introduces an entropy term that can make the transport plan overly smooth, potentially blurring the matching structure. How sensitive is the method to the entropy coefficient $\\lambda_h$? Have the authors considered strategies to prevent the transport plan from collapsing into a nearly uniform coupling?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper formulates causal matching as a Sinkhorn-regularized optimal transport problem, and provides theoretical error bounds which explicitly motivate the design of the objective.\n- The proposed framework integrates a novel random-walk mechanism for counterfactual propagation, combining matching, representation-learning and outcome-prediction into one unified pipeline. This global propagation over the sample manifold is inventive and extends beyond traditional nearest-neighbour or single‐group matching methods.", "weaknesses": "- The empirical evaluation is restricted to semi-synthetic and simulation datasets. While these are common in causal inference research, including at least one real-world biased observational dataset would significantly strengthen the practical claims of the method, ensuring it generalises beyond controlled or synthetic settings.\n- The experimental comparisons omit recent OT-based causal estimators. \n- The current formulation appears to assume relatively homogeneous treatment-group distributions and does not explicitly model or correct for internal biases within treatment groups or highly unequal propensity distributions.", "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": "O8eIJ4u3Cb", "reviewer": "Reviewer_dnf7", "summary": "This paper investigate the limitation of existing matching approaches in causal inference, which select neighbors only from the target treatment group, leading to suboptimal counterfactual predictions due to data sparsity and distribution shifts. Hence, this paper proposes a novel matching framework: The core innovation is removing group barriers entirely, allowing neighbor selection from all samples regardless of treatment assignment. Then it develops a self optimal transport model and information propagation mechanism to realize the algorithm. The method is rigorously evaluated across both synthetic and semi-synthetic datasets.", "questions": "1. I am curious about the role of $Y \\odot M$ in equation (14). I think $Y \\odot M = Y$. Why does this term is incorporated?\n\n2. In the current framework, the mapping function $\\phi$ and the transport cost matrix are learned simultaneously. I am interested in exploring the effect of adopting a two-stage learning process where these two components are learned separately.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper identifies a clear limitation in conventional matching methods-the data sparsity in the specific treatment group can induce the difficulty in finding the near neighborhood. \n\n2. It provides a solid theoretical motivation by analyzing the estimation error in matching, which naturally leads to the formulation of the self optimal transport model. \n\n3. The experimental design is thorough. A lot of typical baselines are incoporated into the experiments. This shows that the approach is effective.", "weaknesses": "1. This paper propose the algorithm based on the matching method. Although I believe the proposed method is superior to the vanilla matching method, it is still unclear why the matching-based method is superior to the other classes of methods, such as regression-based method (e.g. CFRNet and GANITE).\n\n2. In the simulation experiments, only the statistics over all the setup $m=0.1, 0.4, 0.7, 1.0, 1.3$ is calculated. I suggest to reveal the results on each setup respectively. This is beneficial for understanding the behavior of the method under different scenarios.", "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)