Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport
Authors: Milena Gazdieva
Jaemoo Choi
Alexander Kolesov
Jaewoong Choi
Petr Mokrov
Alexander Korotin
Venue: ICLR 2025 Poster (2025)
Keywords: unbalanced optimal transport
barycenter
generative modeling
Primary area: generative models

Abstract: Aggregating data from multiple sources can be formalized as an *Optimal Transport* (OT) barycenter problem, which seeks to compute the average of probability distributions with respect to OT discrepancies. However, in real-world scenarios, the presence of outliers and noise in the data measures can significantly hinder the performance of traditional statistical methods for estimating OT barycenters. To address this issue, we propose a novel scalable approach for estimating the *robust* continuous barycenter, leveraging the dual formulation of the *(semi-)unbalanced* OT problem. To the best of our knowledge, this paper is the first attempt to develop an algorithm for robust barycenters under the continuous distribution setup. Our method is framed as a $\min$-$\max$ optimization problem and is adaptable to *general* cost functions. We rigorously establish the theoretical underpinnings of the proposed method and demonstrate its robustness to outliers and class imbalance through a number of illustrative experiments. Our source code is publicly available at https://github.com/milenagazdieva/U-NOTBarycenters.

Published as a conference paper at ICLR 2025
ROBUST BARYCENTER ESTIMATION USING NEURAL OPTIMAL TRANSPORT
Milena Gazdieva"
Jaemoo Choi*
Skolkovo Institute of Science and Technology
Georgia Institute of Technology Atlanta, GA, USA j jchei843@gatech.
Artificial Intelligence Research Institute
Moscow, Russia
mi milena. www.dieaassooooooooae
Alexander Kolesov
Jaewoong Choi
Skolkovo Institute of Science and Technology
Sungkyunkwan University Seoul, Korea j sewwonncchiisskkee
Artificial Intelligence Research Institute Moscow, Russia
a. ..klessv@skoltechhh. ru
Petr Mokrov
Alexander Korotin
Skolkovo Institute of Science and Technology Moscow, Russia
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute Moscow, Russia a. k.korotin@kklteehh
p. mmokrov@skoltech
ABSTRACT
Aggregating data from multiple sources can formalized an Optimal Transport (OT) barycenter problem, which seeks to compute the average of probability distributions with respect to OT discrepancies. However, in real-world scenarios, the presence of outliers and noise in the data measures can significantly hinder the performance of traditional statistical methods for estimating OT barycenters. To address this issue, we propose a novel scalable approach for estimating the robust continuous barycenter, leveraging the dual formulation the (semi- Junbalanced OT problem. To the best of our knowledge, this paper is the first attempt to develop an algorithm for robust barycenters under the continuous distribution setup. Our method is framed as a min-max optimization problem and adaptable general cost functions. We rigorously establish the theoretical underpinnings of the proposed method and demonstrate its robustness to outliers and class imbalance through number of illustrative experiments. Our source code is publicly available at https::: / //git//iit thub. com/mi Tenaga zdi eva http:/UOONooaaaaceeeee..
StyleGAN imago monifold
M
A
ondiers
ontliers
G(a))
7
3
G(
C1 color
2 2 3 3
72
((2))
3
2
2
IP2
23
C2
22
sbape
Intent space
Z
Figure 1: The semi-unbalanced barycenter of distributions of colors (IP) and digits (P₂) computed by our U-NOTB solver in the latent space of a StyleGAN model pretrained on colored MNIST images of digits '2, '3'. Our solver allows for successful elimination of outliers in the input distributions.
"Equal contribution

---

Published as a conference paper at ICLR 2025
INTRODUCTION
The world is full of data, and accurate analysis of this data allows for solving number of problems in the world related to science, medicine, engineering, etc. A common challenge in data analysis arises from two factors: (i) a huge quantity of data, and (ii) the data originating from multiple sources, such as different scientific experiments, hospitals, or engineering trials. One promising way to address these difficulties to perform data aggregation. This means reducing the partieular source-specific characteristics and leaving only source-agnostic information for solving practical case on hand.
Recently, fruitful branch of research (Li et al. 2020 Fan et al. 2021: Korotin et al. 2022a Kolesov
et al. 2024a b) has emerged that focuses on addressing the challenge of data aggregation through
the framework of the Optimal Transport (OT) barycenter problem. Given a number of reference distributions representing data from different sources, the problem is to discover a geometrically meaningful average of these distributions by minimizing the average OT costs from the reference distributions to the desired one.
Originally introduced for rather specific setup (quadratic Euclidean OT cost functions) (Agueh & Carlier ). the subfield of OT barycenter has evolved tremendously over the last decade. At present, it is the subject of deep learning and can be adapted to nigh-dimensional data setups, e.g., images. However, current barycenter research tends to ignore a typical and important property of data appearing in real-world applications: the presence of undesirable noise and outliers (Le et al. 2021). This motivates the development of robust OT barycenter framework which (i) inherits all the advances of current barycenter researches, e.g., permits deep learning and could be adapted high dims; (ii) wisely deals with aforementioned inconveniences of real-world data. To satisfy both aims,
we propose to opt for recent continuous unbalanced OT techniques (Gazdieva al. [023] Choi et al., 2024b a) which are naturally aligned with data imperfection.
Our contributions are follows:
We propose a novel semi-unbalanced OT (SUOT) barycenter algorithm that robust to outliers and class imbalance in the reference distributions ($4)
2. We develop a solid theoretical foundation of our proposed approach including the quality bounds for the recovered transport plans ($4.1)
3. We conduct number of experiments toy and image data setups to demonstrate the performance of our method and showcase its robustness to outliers and imbalance of classes ($5).
To the best of our knowledge, our work is the first attempt to construct robust OT barycenters under continuous computational setup (2.3) Existing works addressing similar problems consider
exclusively discrete settings (Le et al. 2021 Wang et al. 2024).
Notations. Throughout the paper, we denote K (1, K} for K E N and use sp:k] denote the tuplet of objects ($1, sk). We define X' RR' 24 CRRR,, Euclidian spaces. We use C()") to denote the set of continuous functions V. The sets of absolutely continuous Borel probability (or non-negative) measures on X' are denoted as P(X') (or M+ (X)). For two measures , p2 E M. we write li < μ12 to denote that the measure μii is absolutely continuous w.r.t. the measure μ2z The joint probability distributions on X'x) with the marginals P and Q are denoted [I, Q) and usually called the transport plans. All joint non- non-negative measures on x/ are referred to as M, (X' x x)). For given measure (zz, y) E P(X' x ))) (or M,, M+ (X))))), we use (r), Yy (v) to denote its marginals. All probability measures (x, y) E P(X' x ))) s.t. 7 = Q for a given probability measure Q are denoted as II ()). We use >u(z) ) to denote the conditional probability measure. For a measurable map T. T4 denotes the corresponding push-forward operator.
as the compact subsets of
The Fenchel conjugate of a function (u)) is denoted by (t) supwER (ut (.))..
def
2 BACKGROUND
In this section, we give an overview of the theoretical concepts related our paper. In 2..1 we state the SUOT problems and its semi-dual formulation. In $2.2 we formulate the OT/UOT barycenter problem. Then, in 2.3 we describe our computational setup. For details on OT, we refer to
(Santambrogio 2015 Villani et al. 2009) Peyré et al. 2019). unbalanced OT - CCiiat 2017
Liero
ct al. 2018 Séjouréé et al. 2022). OT barycenters (Agueh & Carlier 2011 Chizat 2023).
To begin with, we introduce the concept of vi-iieegeences needed for our further derivations.
divivergeeces for neg-negative measures. Let H2 E M,(X')) be two non-negative measures. Then, for function d R R2 U (00), ,, the vi-divergence between fi, #2 is given by:
2

---

Published as a conference paper at ICLR 2025
De (en l122)) def Je v
pi (z) 12(x))
dpz(x) if i < p2 and + oo otherwise.
Here, the generator function (b(()) ) is assumed to be convex, non-negative, lower semi- comi-continuous, to attain zero uniquely when u = and to satisfy the property V(w) " "=Se cc. Then, the conjugate (u) is also known be non-negative, and its domain is R. see (Séjouréé et al. 2019). Under these assumptions. D& is valid measure of dissimilarity between two non-negative measures. Well-known examples of u-divergences include Kullback-Leibler divergence DKL, (Chizat 2017 Séjouréé et al., 2022) and x² <i-diverge D22. For the extended discussion on admissible divergences, we refer to (Gazdieva et al. 2025] Appendix C).
2.1
OPTIMAL TRANSPORT
Classic OT problem. Consider two probability measures Pe P(2), QEP()) and the cost function c(2, y) E C(X x y). Then, the OT problem (Kantorovich, 1942) between IP and Q is given by
(F, def
inf xEII(P.Q) xx
c(1, (I,,))T(z, y)drdy.
Under the mild assumptions on the distributions P., Q and cost function c, the minimizer x² of (I) always exists but is not guaranteed to be unique, see (Villani et al.) 2009). In the special case of the
quadratic cost c(x, y) = 2 problem ① reduces to the well-known (squared) Wasserstein-2
|r-vll
distance (W2(()) Q)) Problem admits a semi-dual reformulation:
1/54 f (x) + f(y)d[(v)]
OT Q)
sup fEC(Y)
(2)
where (x) def inf merr) J3 (c(x,y) )) f()))dply) is the c-transform. Note that the standard c-transform in the OT duality usually defined as f*(x) inf y) f(y) Here, for the future needs, we substitute the classic c-transform by the weak c-Iransform (Backhotf-Verague Veraguuss ct al. 2019] Theorem 1.3), (Gozlan et al., 2017). This transition is valid since the infimum weak transform is anyway attained at any μ E P(Y) supported on the arg inf
f(y) set. Unbalanced & semi-unbalanced OT problems. The standard formulation of the OT problem entails various issues (Balaji et al. 2020; Séjouré et al. (2022), such as sensitivity to outliers, inability to deal with class imbalance the source and target measures and inapplicability to general non-negative measures. Still, we consider the mass transportation between probability measures and focus on the first two issues. Fortunately, they can be overcome by relaxing the hard marginal constraints of the OT problem which yields the Unbalanced OT (Chizat 2017 Liero et al. (2018) UOT) problem. Formally, the UOT problem between the probability measures PE P(Z). Q (()) is
Jeeg ()),)) y)drdy + Di (1x (P) + D (7)11))
inf UOT op (P, Q (EM-(XXX)) xxy
(3)
where c(x, y) E C(X y) is continuous cost function and Dv,, D& are didivergesces over X' and V. respectively. We can tune the degree of penalizing the marginal distributions mismatch by introducing the unbalancedness parameter т > 0 and considering De = TD Dr. De = rDor. Informally, when +00, the corresponding UOT problem tends to classic OT I. More precisely, the classic OT problem (T) is a special case of the UOT problem (3) where the --divergences Dy. D are defined such that Dar (T' 0 if TP' Q and +oo otherwise. Equivalently, it means that the generator functions v, Φ are the convex indicators of {1} and their conjugates are defined as (t)) d(t)
In this paper, we focus on the unmilanbalnnced OT (SUOT) problem, i.e., the case when only the first ①1) is softened:
marginal constraint in the OT problem
SUOT a. (P, Q) = inf
((x,y)((, y)drdy + Do (P=
(4)
7Ell(Q) xxy
Note that the minimizer "* of (4) always exists thanks to the compactness of the space X' x y (yielding the compactness of IT ()), continuity of the cost function c(x, y) and lower semi-continuity
of the function v (yielding the lower-semi-conting
of the optimized functional w.r.t. 7).
The semi-dual SUOT problem is given by
sup SUOT c. (P. fEc())
(5)
Note that under mild assumptions on the cost function c(x, v), optimal potential I of (5) always exists, see (Liero et al. 2018 Theorem 4.)4). When (+)) = L, problem (5) reduces to classic dual problem 2 for balanced OT problem
3

---

Published as a conference paper at ICLR 2025
2.2 USEMI-JUNBALANCED OPTIMAL TRANSPORT BARYCENTER
Consider probability measures P& E P(Xx) and continuous cost functions cx Xx x R. k E K. Given weights A 2 0 s.t. Σx= A = 1, the classic OT barycenter problem consists in finding minimizer of the sum of OT problems with fixed first marginals PI:K]:
K
inf QEP())
OT (IRR,
To ensure the robustness of the barycenter estimation, it is natural to consider the relaxation of the marginals Pl::K), ie., substitute the OT problems in (6) with the SUOT problems. Let M/t:k] be the set of u-diverggeness Then, the SUOT barycenter problem is
L def inf (Q) def inf QEP())
Σ
SUOT Cv. (P&s Q)-
QEP())
Clearly, the SUOT barycenter problem 7 subsumes the conventional OT one (6) We also note that in principle one may substitute OT with UOT problem B) unbalanced from the both sides (Friesecke) ett al.] [ Chung & Phung 2021 Bonafiii et al.) 2023). However, this is not practically meaningful.
Indeed, the ultimate goal is get rid of the potential outliers and noises in the input data, not in the barycenter. So there is no need for unbalancedness in the barycenter.
Thanks to (Liero et al. 2018 Corollary 2.9), the functionals Q + aUOTa..ea (Px, Q) (k K) are convex and lower semi- continonniuuuus. Thus, the functional Q + B, (Q) is convex and lower semi- continuous itself. We note that P())) is weakly compact. Thus, thanks to the Weierstrass theorem (Santambrogio) 2015 Box 1. B.. (Q) admits at least one minimizer meaning that the barycenter Q exists. Still, we can not claim that is unique since the functional Q + (Q) is not proved to be strictly convex.
2.3 COMPUTATIONAL SETUP
In real-world scenario, the distributions Pss are not available explicitly and can be assessed only via empirical samples. Assume that we are given N₂ empirical samples = =fi.N.] Px, h E K. Our goal to find approximations 5 of the SUOT plans ア between the measures Ps and unknown SUOT barycenter Q' for the given cost functions and ei-divergences Ck2 1/., k E K. After that, we can use the learned plans to perform the conditional sampling. i.e., derive new points y 7r (xx) from the (approximate) barycenter taking samples Zg ~ (fx)zz as inputs. Actually, sampling from the left marginals of the learned plans is not an easy task since the marginal can not be easily assessed using the learned plans or potentials. Fortunately, we can deal with this issue using the rejection sampling procedure (Forsythe, 1972). see our $4.2 for the additional details.
It is important to note that the learned plans should admit the new input samples, i.e., those which are not necessarily present in the training datasets. This setup is typically called continuous (Li et al., 2020; Kolesov et al., 2024a; Korotin et al. 2022a) and significantly differs from the discrete one (Peyré et 20,1 & Doucet 2014). The latter is aimed at solving the barycenter problem between the empirical measures which makes its application new data samples challenging.
3 RELATED WORKS
Due to the space constraints, below we give an overview of the most relevant OT barycenter solvers and leave the discussion of the related continuous OT solvers to Appendix B
Continuous OT barycenter solvers. We start with brief overview of continuous OT barycenter solvers and then proceed to the current state of robust (unbalanced) barycenter research. The continuous OT barycenter methods could be categorized as follows: (Fan et al. [02 Korotin et al., 202 stick to ICNN (Amos et al. 2017) parameterization and work only with quadratic Euclidean costs (L2); Korotin et al. 2022a) also deals with f but utilizes fixed point algorithm (Alvarez-Fsteban al. al. 2016);; (Noble et 2003) builds upon Schrödinger bridge; (Li et al.) 20200: Chi et al. 2023)
take the advantage of congruence condition; (Kolesov et al.) 2024a b) empower the congruence
condition with Neural OT (Korotin et al.] 2023b) and Energy-guides Neural OT (Mokrov et al. 2024)
correspondingly. The latter works are the current SOTA in the continuous OT barycenter domain.
The area of robust barycenter computation is much less explored and, to the best of our knowledge, limited exclusively to discrete ($2.3) solvers (Chizat al., 2018 Le et al. 202 Beier al.) 2023 Séjourné et al., 2023) Wang et al.) 2024, Manupriya et al. 2024 Yang & Ding 2024) based on wataaaaatttiiinanaaa
algorithm/Sliedd OT etc.). In contrast, we take a significant step forward and propose the first continuous robust barycenter approach with proper
different principles (MMD
theoretical support and practical validation ($4 5
4

---

we propose the algorithm to solve this problem. The proof' for all theoretical results are Published as a conference paper at ICLR 2025
4 PROPOSED METHOD
In
$4.1
we derive and investigate our novel optimization objective for learning SUOT barycenters.
In
3
4.2
given II Appendix A
4.1
DERIVING THE OPTIMIZATION OBJECTIVE
To develop procedure for estimating the barycenter, one may simply substitute the dual form of SUOT problem (5) in the barycenter problem 7 and get min-max problem. Actually, without additional modifications, it will lead to the min-max-min problem since the c-transforms f& also need to be optimized. In Theorem |
below, we present the way for solving the optimization problem 7 in max-min manner and without the optimization over distributions Q E P(Y)).
Theorem (Semi-dual form of SUOT barycenter problem), The dual form of SUOT barycenter problem is given by
Nol J/x - (2L)) +m]
C
sup
(8)
MER.SUNJEC (2) k=1
Σ d
Hereafter, we say that the potentials f[1:K] are m-congruent or satisfy the m-congruence condition
Σ Axfx = m (for some m E R). Substituting the definition of the c-transform in (8) we get our final optimization objective.
Corollary (Maximin reformulation for the semi-dual problem (). It holds:
K
(-/, (cK(Ix, v) fu(v)) drn (wlmx) dIP (xx) +m
L = sup inf mER, ((((x))EP()) fIAJEE (2) Ex-1 Aef
A
(9)
Xx
where the sup is taken over m E and m-congruent potentials fu.i] E CK), inf over conditional plans (u(-(rx),
Interestingly, in the special
when at least of the SUOT aU terms in the right-hand-
side of 77) reduces into OT the m-congruence condition in (9) turns into congruence condition
EX fr = O, which typically appears in balanced settings Kolesov et al.
2024a b).
Corollary 2 (Congruence Condition of the Special SUOT barycenter problem). Suppose that (t( t, i.e., SUOT = OT,, in (.. Then, dual form (9) turns into the following dual formulation:
K
/
Jn (-/f Vl(FF..))-f400 d do d dPP((T)) (10) Xxx
L = sup
inf fuuKjEC (2) ((12)))()))
A
E f=
Theorem
2 below demonstrates the two important properties of the true plans 7" solving the SUOT barycenter problem (7). First, the left marginals of these plans (7x can be characterized via the optimal potentials fū which allows us to find procedure for sampling from the marginal density during the inference, see 84..2 Second, the Theorem guarantees that the family of optimal plans is indeed one of the solutions to (9.
Theorem 2 (SUOT barycenter conditional plans are contained in optimal saddle points). Assume that there exist m and (S' (y) E C()) which deliver maximum to the problem ( Assume that og (xx y) are continuous cost functions and ty are continuously differentiable functions. Let ( K be a family of optimal SUOT plans between Pk and some barycenter Q*. Then for every plan 7 (a) the marginal of the plan can be represented as:
(x)) (
(11)
(b) the corresponding conditional plan 7i. (-1xx) satisfies
)=(=))))) arg m.) C (([ii (7x(-1==))
(12)
Theorem 2 shows that for some optimal saddle points {(EE k of (9. it holds that (nz k=1 the family of true SUOT plans between Pp:kk] and @* Meanwhile, for any family of optimal ffi:K]'
the arg min,, (-lag) sets might contain not only optimal SUOT plans K k but other functions as well which a known fake solutions issue of neural OT solvers (Korotin et al. 2023a).
5

---

Published as a conference paper at ICLR 2025
Still, in Theorem 3 below, we show that if the triple (m, [fk]k=1
1. 15k fk-1) solves the problem !
relatively well and some additional assumptions are satisfied, then the recovered conditional plans (((+)=))))) are close to those of the true SUOT plans (- [ax)) } k== Thanks to this result, we deduce that when Algorithm| 1 (34.22 optimizing 9 converged nearly to the optimum, its solutions close to the true conditional plans.
Theorem 3 (Quality bounds for recovered plans), Consider recovered m-congruent potentials FLLKK
s.t. FR E C())), k E K and recovered conditional plans (fw(()ax)) Ik ~ PR s.t. (-(aaxx) e P()), E K. Assume that the conditions of Theorem 2 are satisfied, y is convex set, and maps y + e (T u) - (v) are 8-strongly convex on V, k E K. Also, ler 0 < def (bb(-b) where is the upper bound b = max max max {cx v) fu(v)). Introduce the duality gaps: ".ENN.VEY
kEK fnE(fe.fE)
81 81(
def [(f(:K): M..K), m) 7x(2)))))))))
inf [[ (,,K, m))
52 b2(f(18)) def -- inf
[(fuxK) ((( m).
T(H-)))PP))) Then, 2 (6) + 82) 2 Ex= AE Sx, (e) (x).
1/3
Theorem 3 establishes quality bounds for the recovered plans {i (+xx))))) investigating the duality gaps, i.e., the errors for solving inner and outer optimization problems in (9) Note that the strong convexity conditions, in general, are not guaranteed in practice, but are commonly used in the related
papers (aan et al.) 2023 Rout et al. 2022), see (Kolesov et al. 2024a 84) for additional details.
4.2 PARAMETRIZATION AND PRACTICAL OPTIMIZATION PROCEDURE
Parametrization. To implement the optimization over distributions 7-1lau) E P(Y) (k E K) in p we consider parametrizing them with stochastic or deterministic functions, using the strategy defined in (Kolesov et al. 2 $4.1) and (Korotin et al. 2023a 64.1).
To realize the optimization over conditional plans >x (- (xx) in ( we define an auxiliary space S C RDD, atomless distribution 3 E P(S) and measurable maps Tx Xr x S + V. For every plan Tke we consider the representation using the map T& s.t. (- (z) = (x, S Using the stochastic
parametrization of the conditional plans, we can reformulate the optimization objective (g) as:
K
(
(faCCTN 8)) cn (xx, (Tk s)) dS(s) dP dF (13)
L. sup inf TILA)
A
mER, fnxijEK(()) (D)
k=
Xx
EK_ Aef==m
c0 [u.0q m))
We denote the expression under sup inf in 133 as [[[L.K] T[.K]. m). In some of the setups, the stochasticity of the conditional plans 7-11x) is not needed. Then we can consider a measurable map Th: : Xx + y which specifies the deterministic conditional plans as (-((- (.),
For solving (3), we parametrize the maps Th and potentials fx (k E K) as neural networks
:R Dx x R. + R and fe.0 : R Pk + R with weights def W[I:A] 22 e def 01. E O. Here Q= Ω1 x f2 x x 22x and θ = θ , x O₂ x... x Ox are the parameter spaces for the maps and potentials respectively. Note that IR D. denotes the stochastic dimension which should be omitted in the case of deterministic maps. In order to ensure the m-congruence condition for the potentials, we parametrize them as fx.0 = 0 Engk XaIk-1]9n. 9n.0 + Kbe
m where Ik.0 denote the auxiliary neural
nets. Here we take inspiration from similar strategies in (Li al. 2020 Kolesov et al. 2024b a).
Training. Recall that we work in the continuous setup of the barycenter problem, i.e., the distri- butions P[I.K] are accessible only through the empirical samples of data. Thus, we opt to estimate the objective (3) from samples using the Monte-Carlo method. Specifically, the objective op- timized using stochastic gradient descent-ascent algorithm over random batches of samples from the distributions Ps, and auxiliary distribution T For simplicity, we replace the minimization of the objective [(fo Tw.(I=k)) w.n.t. Tu.[1:k] with the direct minimization of the c-transform
Sx. Js (cu(I,, T2 s)) fa (Zz,, s)) dS(s)dPAx which has the same set of minimizers. We detail our optimization procedure in Algorithm|] Inference. Our approach for SUOT barycenter estimation relaxes the marginal constraints related to
Sampling from the stochastic distribution is omitted when we deal with the deterministic maps Tw.[1-K]-

---

Published as a conference paper at ICLR 2025
Algorithm 1: SUOT Barycenter with Semi-Unbalanced Neural Optimal Transport
Input: Distributions FI,K. accessible by samples; transport costs C : X3 D + R: maps Th and m--conguent potentials fu. k E K; number Nr of inner iterations; batch sizes. Output: Trained (stochastic) maps [l:K],. approximating SUOT plans between Py: and barycenter Q
repeat
Sample batches X ~ PR. k E K; For each 2k E Xx sample an auxiliary batch Stas] ~
f EKER
Σ
[(lb(a.))))<(( (T To (mx..)) m
NE
[S/aa]]
ILEX2
#LES(EL]
oc, to maximize Cy;
Update by using "9 for nt 1,,,..... .....::................... :
Sample batches X ~ PA: For each E Xg: sample an auxiliary batch S(xx)] S;
LT
X E
Σ
[S[ze]]
BLEX
Update w by using de
to minimize Lr;
until not converged;
the input distributions Px. It means that during the inference, the input points Ty. should be sampled not from the distributions P3 but from the left marginals of the learned plans (fk)z. However, sampling from these marginals is not available explicitly and requires the usage of the established techniques for sampling from the reweighted distributions, e.g., the rejection sampling technique. We apply this technique in the following manner.
(a), we are able to approximate the fraction of the densities of measures dF)()) (+(((f)) where fo denotes the potential learned during training. Thus, we can (1) generate new samples X Pa. and samples of the same shape from the uniform
Thanks to our Theorem
(r)) and PA:
(()))(()))
distribution ug ~ U; (2) calculate the constant c del max (V (V+(((f)) (3) accept the samples
X which satisfy the inequality uy S VV)(-(fa) and reject the others.
5 EXPERIMENTS
In this section, evaluate model through various experiments. In 5.11 compare the numerical accuracy of our method against baseline methods. In $5.22, we investigate two key properties of our U-NOTB method: robustness class imbalance and robustness outliers. Specifically, we show that these properties can be controlled by adjusting the unbalanced parameter т. In $5.3,, we evaluate our model for the general costs, leveraging shape and iovarainvariant cost on the image dataset, other than the quadratic cost, We provide the details about settings and baselines in Appendix C In Appendices D.I D.2 we test our solver in the balanced/semi unbalanced OT barycenter problem for Gaussian distributions with computable ground-truft solutions. Additionally, in Appendix D.3] we illustrate another interesting application of U-NOTB in higher dimensions showing ability to manipulate images through interpolating its distributions on the image manifolds. The code for our model written using PyTorch framework and publicly available at
ht
com http://hwww...ooooooooooeeeeeaea wtt.:e/wwwwwooootteeeeeeeeee
5.1 BARYCENTER EVALUATION ON SYNTHETIC DATASETS
Baselines. We aim to evaluate whether our model accurately estimates the optimal barycenter Q* and the corresponding transport maps T []: Because our work is the first attempt to address the continuous SUOT barycenters, there are no existing baselines for direct comparison. Therefore, we carefully designed two baseline models for comparison with our approach. These baseline models motivated by the equivalence between the OT barycenter of two distributions and the interpolation between them (Villani et al. 2009). Each baseline model is derived from two approaches for learning the unbalanced transport map Tuot between P, and Py: (i) the semi-dual UOT model (UOTM) (Choi jet al. 20244) and (ii) the OT Map model (OTM) (Rout et al.) 2022) combined with the mini-batch UOT sampling (Eyring et al. 2024). Because the SUOT problem is equivalent the OT problem
7

---

Published as a conference paper at ICLR 2025
A
4-
1-0
112
12
(a) Sctup
(b) Discrete UOT
(c) UOTM
(d) Mini-Batch UOT (e) U-NOTB (ours)
Figure 2: Transport map Ti (11), Ty ~ PI and UOT barycenter distribution Q Ti P1 obtained by
Discrete UOT, UOTM, Mini-batch UOT, and our method in Spiral - Gaussian Mixture.
Data
Metric
UOTM
Mini-Batch UOT
U-NOTB (ours)
L2
0.65
0.31
0.25
Sprial -
Gaussian Mixture
W2
0.29
0.15
0.12
C2
0.89
0.64
0.52
Moon + Spiral
Wa
0.59
0.93
0.81
Table 1: Comparison on Benchmarks on Toy datasets.
between the rescaled distributions, we conduct the displacement interpolation using this Tuot to approximate the barycenter. Specifically, we consider the following SUOT barycenter problem:
inf (A, (IP,, Q) A2SUOT, Q}},
(14)
QEP(Y) where C1, 02 are the quadratic cost functions (3r, y)= - \1 21= Let Te del X1 Id++ X2 Then, Ti and 71 approximate the barycenter distribution Q and the transport map from P2 to Q*, respectively. Note that we should set the OT distance for the first distribution IP, ensure that the estimated barycenter Q = Ti # P1 has a unit mass in two baselines. Since the OT problem is a specific case of the SUOT one, our method can be naturally applied to this case. We test U-NOTB on the synthetic dataset of Spiral-to- to--aussinnn Mixture (S (S-+GM) (Fig. 2a).
Experimental Results. Fig. 2 presents qualitative visualization of the results from our model and two baseline methods. Since there is no closed-form solution for the continuous SUOT barycenter problems, we consider the SUOT barycenter between discrete empirical distributions, obtained from the POT (Flamary et al. (021) library, as the ground-trute solution Q and T (Fig, 2b).. In Fig. 2c we visualize the spiral distribution P1, the estimated transport map Ti and barycenter Note that, because we set OT,, for IP; in Eq. (14), the rejection sampling is not applied in this case. Fig. 2 shows that our model and other baseline methods present decent results.
#
Table provides a quantitative evaluation of our model for more detailed assessment. The approxi- mate barycenter Qu- and the transport map Ti... from each method are evaluated from two perspec- tives. First, the distribution error between Qu- and Q* is evaluated by measuring the Wasserstein-2 distance "Q (W₂ metric). Second, the optimality error of the transport map is assessed by measuring C2-norm between the transport maps (C2 metric), i.e., Sx, ||Ti, (z1)- -Ti (x1) 1|2 (z 1) As shown in Table our model achieves the most accurate barycenter estimation in three out of four scores, evaluated using two metries (W2 and L2) across two synthetic datasets.
5.2 ROBUST BARYCENTER ESTIMATION UNDER OUTLIER AND CLASS IMBALANCE
In this section, we examine two key properties of our U-NOTB model: (i) robustness to class imbalancedness, (ii) robustness to outliers. The value of OT functional O. is known to be sensitive to outliers and class imbalance problems. Hence, the OT harycenter is also largely affected by these factors. To address this sensitivity, the UOT problem extends the traditional OT by relaxing the constraint on marginal densities. This relaxation introduces two notable characteristics. First, even in the presence of class imbalances, the reweighting process allows the model to appropriately align with the relevant modes in a reasonable manner (Eyring et al. 2024). Second, the model exhibits robustness to outliers, maintaining reliable transport despite their presence (Balaji et al. 2020 Choi al. 2024a). The goal of this section is twofold: (1) to investigate whether these two properties are also observed in our U-NOTB model by comparing with its OT counterpart (Kolesov et al.) 2024a) and (2) to demonstrate that this robustness can be controlled through the unbalancedness parameter
Robustness to Class Imbalance. The class imbalance issue refers to the case where each class has a different proportion across Px. This class imbalance problem can lead to undesirable behaviour in the standard OT barycenter. For example, in Fig. B we consider the upper-modes modes and bottom-modes for each distribution as class I and 2, respectively. In Pl, 75% of the data is concentrated in class 1
8

---

Published as a conference paper at ICLR 2025
Po
Po ma 2-4 ralsted
ve A y-o Reooadd
Kyming R
y-o
(a) NOTB
(b) 200
(c) T = 20
(d) T =
Figure 3: Conditional plans (v< and barycenter Q obtained by NOTB and our method in Gaussian
Mixture barycenter experiment. We evaluate various unbalancedness parameter T E {1, 20, 200).
300
-008
-3
-25
304
304
(a) Marginal distributions
(b) NOTB with outlier
(c) NOTB without outlier
(f) Figure 4: Learned barycenter Q obtained from x ~ P1 by NOTB and our method on Gaussian Mixture with 5% outliers. The evaluation is conducted for various unbalancedness parameter r {1, 20, 200}. For comparison, we additionally trained NOTB without outliers, shown in subfigure (c).
200
(e) 20
(upper-left white dots) while only 25% concentrated class for P1 (upper-right yellow dots). Therefore, the balanced OT barycenter (Kolesov et al. [024a] NOTB) generates a mode between two majority modes of k. which can an undesirable phenomenon for the barycenter (Fig. Ba).
As shown in Fig. 3 evaluated our model for diverse unbalancedness parameter {1, 20, 100)
(Eq. D.. Additionally, for comparison, we conducted experiments under balanced setting using the important
that as increases, the UOT problem converges the OT problem, allowing for less flexibility in marginal distributions (Choi et al. (024b)). When = 1. our unbalanced barycenter demonstrates robust results by generating barycenter modes between the pair of the closest modes of the P1 and P2. Our model rejects contour samples from majority modes, reweighting them to better align the corresponding modes. Moreover, our unbalanced barycenter offers controllability of this robustness through т. As 7 increases, the flexibility on marginal distributions decreases in the UOT problem. Thus, when т = 200, our model accepts nearly all samples, making the UOT barycenter closely resemble the OT barycenter. We believe this adaptability of the UOT barycenter offers practical benefits world scenarios.
OT counterpart.
Robustness to Outliers. To evaluate the robustness to outliers, we conducted experiments on synthetic dataset that included small proportion of outliers. As depicted Fig. 4. the dataset comprises three marginal distributions: IP, (beige dots), P2 (gray dots), and IP3 (orange dots). Outliers, constituting 5% of the marginal P1 and Pz, were added to the outermost point of each marginal. We evaluated our model for various unbalancedness parameters and compared it with the balanced OT counterpart. Similar to the robustness to class imbalance experiments, we expect that smaller values of will provide more robustness outliers by allowing higher flexibility in marginal distributions.
For the comparative analysis, we also tested the NOTB approach for OT barycenter estimation under the data without outliers. This experiment aims to demonstrate the desired behaviour of barycenter when there are no outliers in the dataset. In Fig. 4 we see that for 7 = 1, our U-NOTB closely aligns with the OT barycenter in this outlier-free scenario. This result shows that model offers robustness

---

Published as a conference paper at ICLR 2025
2
2
2
2
2
3
3
3
3
3
3
7 7 7
7
7
2
2 2 2
2
2
3
3
3
3
3
3
3
3
3
3
2
3
(a) In-distribution Samples from P1. Acceptance Rate: 63%
(b) Outlier Samples from Pl.
Acceptance Rate: 19%
2
2
2
3
3
3
3
2
3
2 3
(c) In In-distribution Samples from 122.
(d) Outlier Samples from P2. Acceptance Rate: 5.5%
Acceptance Rate: 71.25%
Figure 5: Examples of 21 P1 (grayscale digits) and 22 P2 (color images) and its corresponding barycenter y ~ Q samples (colored-digits) in chape-color experiment.
to outliers by rejecting outlier samples at smaller As increases, proportion of outliers. Hence, the outliers on the left and right begin to influence the barycenter, making the modes of barycenter lie between the majority modes of distributions PPx- In summary, reducing т enhances outlier robustness, while increasing yields a more precise barycenter by incorporating all data points. We believe that the flexibility adjust 7 offers promising potential for broader applications, where different levels of tolerance to outliers may be required.
model accepts higher
5.3 SHAPE-COLOR EXPERIMENT
In this section, we illustrate one interesting example demonstrating how the UOT barycenter problem can be applied using general cost. designed the problem at general costs when there are some undesirable outliers in the marginal distributions, see Fig. [ for visualization of the setup.
Shape distribution PI The first marginal distribution consists of grayscaled images of digits '2' (49% of training dataset), '3' (50%) and '7'(1% outliers) of MNIST data.
Color distribution P2. The second marginal distribution consists of three color points: red (probability mass po 0.495), green (p1 0,495). white (p2 0.01 outliers).
Manifold. We pretrained the StyleGAN Karras al. 2019) generator G 2 3x332x33 on the colored MNIST dataset of digits '2' and '3'. Let EY be the network which encodes each marginal Py the latent Z. Throughout the experiment, we transport 23 P& to the barycenter point yis by Go o F.ty (Ixa) thus confining the transformed images to the colored MNIST of digits '2' and '3'.
General Transport costs. For the first marginal sample 21 P1 and its corresponding barycenter point y1, we use the following shape-preserte decolorization operator. Moreover, for the second marginal sample 2r2 P2 and the corresponding
cost: Cj (I), /1 = | H2 (v))) where H, is
barycenter point y2, we use the following color-preserving cost: "2(222,/1 y2) _ H.(32)1
where He color projection operator defined in (Kolesov al. 2024a).
10 are presented in Figure 5 We demonstrate the examples of learned mapping (æk, (x s)) from in-distribution points Tr, see Fig. 5a |cc and from outliers, see Fig. 5b 5d Expectedly, the learned in-distribution mappings preserve the shape (P, - Q*) and color (P2
Our results for unbalancedness
turn, the outlier mappings have no reasonable interpretation.
Along with the qualitative performance, report the acceptance rates of input points 24. the explanation of our inference procedure in $4.2 Importantly, the acceptance outliers is much smaller compared to in-distribution samples. In particular, if 21 is grayscaled digit '7', it will be accepted (and processed through the mapping Ti) only with probability 0.19, while for grayscaled digits '2' and '3' the same figure reads 0.63. Thus, the experiment showcases the applicability of our robust barycenter methodology for non-trivial cases with non-Euclidean OT costs.
6 DISCUSSION
Our work continues the recent and fruitful branch of OT barycenter research. We present the first attempt to build robus! OT barycenters under continuous setup based on unbalanced OT formulation. From now the barycenter researchers have them the ability to deal with imperfect data (with outliers/nois//llass imbalance) at large scale. We believe that our proposed method will further expand the toolbox of deep learning practitioners. We discuss the limitations of our method and future research directions in Appendix
deep learning tool their disposal that gives
Reproducibility. For transparency and reproducibility, make source code publicly available at https: /ts::://gtttbt con/ mi Lenagazdi wwwaaiiiiiaa/aaae NOTBarycenters The code contains a README file with further instructions reproduce experiments.
10

---

Published as a conference paper at ICLR 2025
Code of Ethics. This paper presents work whose goal is advance the field of Machine Learn- ing. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
ACKNOWLEDGEMENTS.
The work Skoltech was supported by the Analytical center under Grant No. 70-2021-00145
the RF Government (subsidy agreement 0000000730321150000 02.11.2021). Jaemoo Choi acknowledges the support by National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) [RS-2024-00410661]
REFERENCES
Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904 924, 2011.
Pedro C Alvarez-Esteban Del Barrio, JA Deesta-Albertos, and C Matrán. A fixed-point approach to barycenters in wasserstein space. Journal of Mathematical Analysis and Applications, (2): 744-762, 2016.
Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference Learning Representations, 2023. URL http:pts///ppeeeeeeeee.....o hwwp:/ewwwweeeeeo.eeeeeeeeeeee forum?
Td=TQ5WUwS 4ai
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155, PMLR, 2017.
Arip Asadulaev, Alexander Korotin, Vage Egiazarian, and Evgeny Burnaev. Neural optimal transport with general cost functionals. In The Twelfth International Conference Learning Representations, 2024.
Julio Backhoff- Backhoff-Veraguas, monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):203, 2019.
Veraguas, Mathias Beiglböck, and Gudmund Pammer. Existence, duality, and cyclical
Yogesh Balaji, Rama Chellappa, and Soheil Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. Advances Neural Information Processing Systems, 10.12934-12944, 2020.
Florian Beier, Johannes Lindheim, Sebastian Neumayer, and Gabriele Steidl. Unbalanced multi-marginal optimal transport. Journal of Mathematical Imaging and Vision, 65(3):394 2023.
Dimitri Bertsekas and Steven Shreve. Stochastic optimal control: the discrete-time case, volume Athena Scientific, 1996.
Mauro Bonafini, Olga Minevich, and Bernhard Schmitzer. Bellinger-kantoroviin barycenter between dirac measures. ESAIM: Control, Optimisation and Calculus of Variations, 29:19, 2023.
Maciej Buze. Constrained hellinger-kantorrvich barycenters: least-cost soft and conic multi-marginal formulations. arXiv preprint arXiv:2402. 1268, 2024.
Jinjin Chi, Zhiyao Yang. Ximing Li, Jihong Ouyang, and Renchu Guan. Variational wasserstein barycenters with c-cyclical monotonicity regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7157-7165, 2023.
Lénaïc Chizat. Doubly regularized entropic wasserstein barycenters. arXiv preprint arXiv:2303.118... 2023.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Mathematics of Computation, 87(314):2563-269) 2018.
Lénaïc Chizat. Unbalanced optimal transport: Models, numerical methods, applications. PhD thesis, Université Paris sciences lettres, 2017.
11

---

Published as a conference paper at ICLR 2025
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. Advances in Neural Information Processing Systems, 36, 2024a.
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Analyzing and improving based adversarial networks. In The Twelfth International Conference on Learning Representations, 2024b.
Nhan-Phu Chung and Minh-Nhat Phung. Barycenters in the hellinger- kantorovich space. Applied Mathematics & Optimization, 84(2):1791-1820) 2021.
Marco Cuturi and Armad Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning, pp. 685-693. PMLR, 2014.
Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems, 34:12955-12966 2021.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrüdinger bridge with applications to score-based generative modeling. Advances Neural Information Processing Systems, 34:17695-17709, 2021.
Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural monge maps improves unpaired domain translation. In The Tweifth International Conference on Learning Representations, 2024.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of wasserstein barycenter via imput convex neural networks. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 1571-1581. PMLR, 18-24 Jul 2021. URL htttps::/ppooceennngggg/// mlr. wtw.ppess/vv3///eee.
html
Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL Featured Certification.
https:://penreeview wtwp:ervwwwww.eee/ouureeeee wot.://oooommmm2222m11eeeee
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy. levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8. 2021. URL http::/tmmml.........
.html
George E Forsythe. Von neumann's 's comparison method for random sampling from the normal and other distributions. Mathematics of Computation,
1972.
Gero Friesecke, Daniel Matthes, and Berharrd Schmitzer. Barycenters for the -antineer-anttrovichh distance over r'd. SIAM Journal on Mathematical Analysis, 53(1):62-110, 2021.
Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Extremal domain translation with neural optimal transport. In Advances in Neural Information Processing Systems, 2023.
Milena Gazdieva, Arip Asadulaev, Evgeny Burnaev, and Aleksandr Korotin. Light unbalanced optimal transport. Advances in Neural Information Processing Systems, 37:93907-93938. 2025.
Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 273(11):3327-3405. 2017.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnev.. Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems, 2023.
12

---

Published as a conference paper at ICLR 2025
Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024a.
Nikita Gushchin, Daniil Selikhanovych, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Adversarial schrödinger bridge matching. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024b.
Leonid V Kantorovich. On the translocation of masses, In Dokl. Akad. Nauk. USSR (NS), volume 37, pp. 199-201, 1942.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. In Proc. NeurIPS, 2020.
Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Genot: A neural optimal transport framework for single cell genomics. In The Tiirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin. Estimating barycenters distributions with neural optimal transport. In Forty-first International Conference on Machine Learning, 2024a.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Anastasis Kratsios, Evgeny Burnaev, and Alexander Korotin. Energy-guided continuous entropic barycenter estimation for general costs. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b.
Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, and Alexander Korotin. Optimal flow matching: Learning straight trajectories in just one step. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations, 2021a.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34: 184:14593-1444 14600, 2021b.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021c. URL https:://ppereeriiiee
n ftt
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022a. URL https //openreview net www./eeeutttttee forum? iidiidd--i Enzx TTaM
Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein gans are not optimal transport? Advances in Neural Information Processing Systems, 35:1 13333- 13946, 2022b.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. In The Eleventh International Conference on Learning Representations, 2023a.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In The Eleventh International Conference Learning Representations, 2023b.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light schrödinger bridge. In The Twelfth International Conference on Learning Representations, 2024, CtLpttps. /ttp:////pppeeiee/eeeee net o-t.
13

---

Published as a conference paper at ICLR 2025
Khang Le, Huy Nguyen, Quang M Nguyen, Tung Pham, Hung Bui, and Nhat Ho. On robust optimal transport: Computational complexity and barycenter computation. Advances in Neural Information Processing Systems, 34:21947-21955 2021.
Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin M Solomon. Continuous regularized wasserstein barycenters. Advances in Neural Information Processing Systems, 33:1 33:17755-17765, 2020.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a new Dellinger-kantorovve distance between positive measures. Inventiones mathematicae, 21 (3); 969-1117, 2018.
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023. URL/ttts:: httpo///AA.tss///eeeeeeeeeeeeeee http:///pweeeeee.eeeeeetttataaaa
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672 6672-6681. PMLR, 2020.
Piyushi Manupriya, SakethaNath Jagarlapudi, and Pratik Jawanpuria. MMD-regularized unbalanced optimal transport. Transactions on Machine Learning Research, 2024, ISSN 2835-8856. URL
net www.oroioeeeeiiiie
Petr Mokrov, Alexander Korotin, Alexander Kolesov, Nikita Gushchin, and Evgeny Burnaev. Energy- guided entropic neural optimal transport. In The Twelfth International Conference on Learning Representations, 2024. https:: http://www.....ooro.oedeodddaaa.
Ngoc-Hai Nguyen, Dung Le, Hoang-Phi Nguyen, Tung Pham, and Nhat Ho. On barycenter computation: Semi-unbalanced optimal transport-based method on gaussians. arXiv preprint arXiv:2410.08117. 2024.
Sloan Nietert, Ziv Goldfeld, and Rachel Cummings. Outlier-robust optimal transport: Duality, structure, and statistical analysis. In International Conference on Artificial Intelligence and Statistics, pp. 11691-11719. PMLR, 2022.
Maxence Noble, Valentin De Bortoli, Arnaud Doucet, and Alain Durmus. Tree-based diffusion schr\ odinger bridge with applications to wasserstein barycenters. arXiv preprint arXiv:2305. 2023.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications data science. Foundations and Trends in Machine Learning, 11(5-6)355-607) 2019.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2022. URL hattps:///genrrevveew
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94 2015.
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Confer- ence on Learning Representations, 2018. URL https: /wwp//oppeeeowooeeooee wet.eeewwweeeeeeeeeeeeeeeeeeeeee forummiroooo
B1zlp1bRM
Thibault Séjourné, Jean Feydy, François-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport. arXiv preprint /9/0v:/910.1 /2958, 2019.
Thibault Séjourné, Gabriel Peyré, and François-Xavier Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint w:Xiv:2211.087/ 2022.
Thibault Séjouréé, Clément Bonet, Kilian Fatras, Kimia Nadjahi, and Nicolas Courty. Unbalanced optimal transport meets wessedd-wssesttein. arXiv preprint arXiv:2306.07777 176, 2023.
14

---

Published as a conference paper at ICLR 2025
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching, Advances in Neural Information Processing Systems, 36, 2024.
Frode Terkelsen. Some minimax theorems. Mathematica Scandinavica, 31(2):005-433, 1972.
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation -fee schr\ odinger bridges via score and flow matching. arXiv preprint 2rXiv:2307.03672 2023.
Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions Machine Learning Research, 2024. ISSN 2835- Expert Certification.
8856. URL wtt.:ttpss//ggeeeee...eeeeeeeeee dc= 3AW
Omer Tov, Yuval Alaluf, Yotam Nitzan, Or Patashnik, and Daniel Cohen-Or. Designing an encoder for stylegan image manipulation. ACM Transactions on Graphics (TOG), 40(4): 2021.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Xu Wang, Jiawei Huang, Qingyuan Yang, and Jinpeng Zhang. On robust wasserstein barycenter: The model and algorithm. In Proceedings of the 2024 SIAM International Conference on Data Mining (SDM), pp. 235-243. SIAM, 2024.
Qingyuan Yang and Hu Ding. Approximate algorithms for k-sparse wasserstein barycenter with outliers. In Kate Larson (ed.). Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence, IJCAI-24, pp. 5316-5325, International Joint Conferences on Artificial Intelligence Organization, 8 2024. doi: 10. 00.24963/jcai.22244588. URL https://doi. org/
10.24 20.22963/1jcaa..222//5333 Main Track.
15

---

Published as a conference paper at ICLR 2025
A PROOFS
A.1 PROOFS OF THEOREM AND COROLLARY
Proof of Theorem lem finding
Substituting the dual form
into
13p. we can formulate the barycenter prob-
K
/ Nx
Va((f((=))dPAaa)) + +1, fu(u)d(())
J
min QEP())
sup EC())
Ax
-Va(-
(15)
f
k=1
where we replace inf with min thanks to the existence SUOT barycenter, see 32.
Note that the compactness of the space y yields the weak compactness of P()).. At the same time, the functional Q + F(Q. fn:k]) continuous and linear in Q. At the same time, is concave in
fp=:K] Indeed, the convex conjugate of the generator function v is convex and non-decreasing by definition, the c-transform operation convex, see (Kolesov et al.) 2024b Proposition A.1 (iii)). Thus, the superposition
concave, non-increasing function (-) and concave function - --)())
ie., 1hh (-f()) is a concave function as well (for every k E K). Thus, we can apply the minimax
theorem (Terkelsen 1972 Corollary and swap min and sup in 15
Au / 1=1 N/
- I fu(e))((v)]] =
sup frum EC()) GEP())
min
K
sup J (EC())
AE
-hal- va((fj((s))dPPP(t)) +
min QEP())
Au Su (n(())) dQ(m)
1=1
X2
k=1
T(w) / F())d(()) ].
K Ak
sup f
min @EP(O)
(16)
Nx
9(h.....xx
Now the following fact: infgep(y) fy T()) inf T(y). Assume that the true value m def
mingey f()) is known. Then we can restrict sup in (16) to the potentials fn:a] satisfying the m-congruence condition: Σk AafE = m. Indeed, for each tuple (f1...,, fx), ) let us consider
the tuple (F1,... ..., Fx) def (f1..., fo m-I) satisfying the congruence condition. For this tuple, infyey Ex=1 Aafr Besides, we get:
((h, AK
(-0x(-((((x))
AK f
[ (xx(- - - m
dF )F 0. (17)
Ax
Xx
Here the inequality in line follows from the properties of the c-transform function and convex conjugate of the divergence" generator function vb. First, c-transform is a decreasing function of its argument, see (Kolesov et al. 2024b Proposition A.1 (i)). Thus,
m - F s fK fK
m m- fK -- (fK
m
JON fk
AK
Ax
Second, the convex conjugate Vx is non-decreasing function which yields
- 2 -fkx + m - -
( + m yox).
Ax
16

---

Published as a conference paper at ICLR 2025
From this, the inequality (17 becomes obvious. It means that the transition to congruent tuple is fair for the known value m. It is important to note that in practice this value not given and should be optimized. Thus, the final optimization objective can be formalized as
K
du / - - + m (18)
= sup mER.Jux ECK (D) - Aam
g(f fk)=
sup
mE fu. E A.f.m
(2)
1mm1
Xx
where (88 coincides with
Proof of Corollary
7
By substituting the definition of the c-transform in we getget
L-=
sup MER.JLLKEEC (2) k=1 Σt AntiEm
Σ A
| -v ( inf Nx
Moseppo / (
dlP d + (19)
inf mivieoon {-va (-//, (c<(TA (cal )al (r)
sup fu EC (2) k=1 Σzs: fm
(20)
In transition from 19 to 20). we interchange the operation of taking inf over (u() P())) and vt-
This is possible since 1 is monotone non-increasing and infimum in (I9) is attained.
Now we note that the map
(IR, (1) + ( // (- (a (w
(21)
is measurable and bounded from below. Indeed, recall that the spaces Xx. y and P())) are com- pact, the potentials fx(2) and cost functions e (Jr..) y) are continuous. Therefore, the function (Ik, 1) + (y(Cc(TR. v) fu(v))du()) is bounded from above, i.e. 3B₁ E:
V(zk, p) E X x P(Y) - [ - fa(v))du(v) s Bk-
vb(u)
Since - monotone non-increasing and dom(vx) (-00, lim
R. then:
=+oc
V(TK, u) E X x P(Y) - (CL(E..))) fw())dpp(w)
2 > DC.
The latter justifies that (21) is indeed bounded from below; its measurability is obvious. Thanks to the properties of this map, we the rule of interchange between integral and inf (Bertsekas & Shreve 1996. Propositions 7.27 and 7.50) and get
20
K
(-
-
sup mER. E (2))
Av
inf (
dI (xx) +m
(22)
Σ A
where the inf is taken over families of measurable stochastic maps ()e((1=>)))p- Note that (22 equal to (9) after the trivial interchange between the summation and inf operations. This completes the proof.
Proof of Corollary 2 substituting fr = f1 X
Since SUOT,, OT vi Id by the definition of convex conjugate. By m
and 1 = Id into (9). we can obtain the congruence condition as follows:
L = sup
inf ((()))))
d I Xx
Je (-/
(co(Fn. v) - 100-Jaale (x()) (wll)))))))))))))))))))))))))d ( + m
m.fuit ΣR AaJm
loml
17

---

Published as a conference paper at ICLR 2025
A1 Jxn / (a(m).1)- (Rww)+ x)) an )d
sup M. m..Kk A A
inf JEPEE))))
U ( (co(Tx.V)- lu (n(olluaa dlP + t
A6
k=2
Xx
d1 L Φ1 (5(-1 (9((),))- - dl (x1)+
sup
inf
A As/eE " (((())))(()))
-Va( (-(--/, (a((E ff) dPP)x)))
Σ Av
Nxx
A.2 PROOF OF THEOREM 2
Below we recall an auxiliary thereotical result which will be used in our proof of Theorem
Theorem 4 (Connection between solutions of dual OT and UOT problems (Choi et al. 2024a Assume that ve is continuously differentiable function. Then if the optimal potential f delivering maximum to the dual UOT problem exists, f is solution of the following objective
OT PP. = ssup/ / (x))d((x) + / f(y)d(()))
p
where dx(x) V (=)(((()))))))P and dā(w) = V
This Theorem highlights the intriguing connection between the solutions the classic OT problem and its unbalanced counterpart. Note that in the case of SUOT problem, v22(y) = y. thus, the measure Q reduces to Q. Now we are ready to prove our Theorem
Proof of Theorem Below we show that optimal potentials fiik] delivering maximum to the dual SUOT barycenter problem (8) are also optimal for each of the underlying SUOT(F,,, Q) problems. Recall that the SUOT barycenter problem (7) admits minimizer which denote by Q*. Then:
(2x) (xx)+m]= Au / - fh (
SUOT che (IPR,
sup mER. / EKEKKKO (V) >* Ae/
1x
(23)
K
sup mER Spijee (2) k= ΣK Au mm
A
(xxx (==)((()))d
Σ (afa(o) S (24)
(2)
Xx
K
K
sup MIAREECK
A - for (Z X1
(o y==1
SUSU .ia (Pk, Q*). (25)
(3))==
Here the transition from line 23 to 24 follows from the fact that the optimization is performed over m-congruence potentials. If then remove this restriction on potentials, the value of sup can become bigger which explains the transition between (24) (5)). Importantly, we get that the final
objective line (25) is equal to the initial objective in line (23). Thus, the inequality in line (24) turns into equality. At the same time, it justifies that the potentials fiiik] delivering maximum to the
SUOT barycenter problem coincide with the optimal potentials for each of the SUOT problems.
Then, thanks to Theorem the SUOT barycenter problem
25 admits a reformulation:
K
K
SUOT( Q*) =
Σ X (TTTP,
(26)
18

---

Published as a conference paper at ICLR 2025
where each of the distributions Pk (k E K) is specified via the optimal potential fž delivering maximum to the problem OT, Q*) in its dual form (2) dlxx ( V (fz )xx (mx))ddPP:ar
Thus, the optimal SUOT plans now correspond to the optimal OT plans N% (z. y) belonging to (P Q7), i.e., having marginals (1%) = Pr. (nk)y Q*.
Now we fix & K. We aim to show that (-(xx) E arg (((x)))))) ((((((((R))) y) (-)) where inf is taken over measurable maps 7/ (-(2)))) Then by repeating the derivations of Corollary I we easily get that 7 (-(aa)) E arg ( (-Irx,)) [(f (=<(((x)))). To prove the former, we again use the fact that %x is a solution of the balanced OT problem between the rescaled marginal Pk and Q" Note that Vik is a continuous function since 4 is assumed to be continu- ously differentiable function. This yields the fact that I& is absolutely continuous distribu- tion w.r.t. Px.- Thus, we can use (Villani et al., 2009, Remark 5.13) which states that for each y E (1pp( (=(a))) it holds that y era arg inf (Cu y/) fv()))- The latter statement is equiva- lent to > (-|aræ) ) arg inf (tlme) (e (FR, -\ー (v))dr((1azaa which completes the proof. Note that since the infimum is attained least for one map 7/ (- (lax),, the arg (-(a))) in the latter equation is actually arg m (+In))
A.3 PROOF OF THEOREM 3
The theorem provides duality gap analysis. We are given approximations: m. Pl.x] which are mi-congruet; .p..x) Recall that the gaps are defined follows:
81 (..K))
[(PI:K; 71:K1; <(..KK):77:K11 m))
inf
D(Pu.K) (1:K], m);
(27)
(a((())))P()))
82(F(u11))
L
inf [(]u.x] 71:K], m).
(28)
Proof of Theorem 3 Our proof considerably extends the proof of (Kolesov et al. 2024a Theorem 4.2) for the case of the classical cost function.
Part 1. Several notations.
To begin with, we introduce several notations (Q* a SUOT barycenter):
(c(f.(7k) def f (-/ (CC(FA.)) dP(()))
(29)
Xx
() def Valfkr mk) + fu(v)d((
(30)
In(Kk) del inf Ve(Jk.mh) P((A))))
(31)
Ja(fk) def Julfk)+ [ fely)d(0'
(32)
Note that the following holds:
AeVieffa 7k) + m
L
sup mER fp. fpIKEEK (3) ΣK= AuEm
AeJT((i) + +
inf meR, ((22)))P())) fuKK (2) Σk faEm
sup
Also, if f(xx]] are m-congruent, then Σk d Syy fa(dd) = m and therefore:
K
K
Σ (o Σ XeValfes mx) + m = (()))))))).)))mmm
19

---

Published as a conference paper at ICLR 2025
K
K Σ k=1
Σ Ak(h)
Ix mm= inf [([1:K), M1-A], m). (((()))P)))
(33)
Now we start the proof.
Part 2. Analysis of gap 81-
At first, we express gap 61 in terms of newly introduced V 29 Tu (31 follows:
81
[([u:K) M. m) - inf C [(p.x] [u:k; m) >u(-1x)EP())
K
d (2 [&aa -
Σ 1 (fu. Tn),
where
( 7v) del Va( 7n)) Ia(Fi)::
(34)
We analyze gaps 61,1 for each particular & K] separately. For convenience, introduce the function ( y) def cx( v) (v), D& Nx. Y. Note that Dk is B-strongly convex in the second argument by the Theorem assumption. Define the maps Tę. Xx y follows:
T (mx)
def
arg inl Oo (F((,,))) y).
yey
Note that Ty is correctly defined (since ÿk is strongly convex) and Borel (cf. the proof of Kolesov et al. 2024a Theorem 4.2). It induces the family of conditional distributions:
+ Th (- E P(D);
Nh (.(2a) del (1)( (-)-
(35)
For any family ne(+1=x) E P())) we have:
f - (- (-f Te (... v)drc(llzz) dP
Ve(Fermn
(36)
X1
Vx (-
(x To (xx)) dl (Ix)
(37)
X1
(-f ( w) (ulea)) di (xx) Va((FinnE)
vhu
Xx
we utilize the properties of vn and T.C
where in transition from (36
37
Ty (IV, ()dr(()
(x Ty. (xx))
S -ğ Ty. (Ix)) Th non non-decreasinng non-
Tn (Ix, y) (v)
Je - To I((a)))
20

---

Published as a conference paper at ICLR 2025
-Tx (-) (o 2 - ( (mxx(xx (cx)))
From the above, we conclude that 7 minimizes the functional 7 ++ Ve 7k). In particular, delivers the value of Js(f)):
Ju(F Vv Vi m
(38)
From the perspective of 38 gap d1 could analyzed as follows:
(1, 7x))
Va FR)) Ve((Fu:Th
/ @ Th (2k: (xx))) (1x))))- $6
( ( (
DulOTA s)dFh(llox))) ) (aru(as))
(39)
2/x
/ {
Ia (Z () - In (2xx. Te (zx))) jdPp(oo))
2
(40)
Nx
{ (e((F - (1x Tá. (mx)) ( (dsioulla)dPalln
(41)
X, y
n3
lly (1
(42)
2
Transition from 39
40
Note that:
- T/ (xx)) 2 J Ih ( v)dr
max E.a (x(,,,)) 2 -b.
def
The first-order properties of convex differentiable function tha read as:
ViLa T
2
Veh( = n
J ( 9N (Zx, T& (xx)))
Th
0h + ( (x TV (#x)
Based on them, we derive:
VE - (Tr, Ti (xx)))
1/
2 n T
(N (sx)))
which validates the transition from from 39 to
Transition from 41 to Note that T. (Irx.) is the minimizer of + (T v). Therefore, due to B strong convexity of 9t on the second argument it holds:
v12 s (I v) (x, T/ (xx)).
which validates the transition.
Part Analysis of gap 82.
Note that we work under the conditions of Theorem 2 In particular, we assume that there exist mn and m=-congruent potentials {fr C())) which deliver maximum to D). Also, we denote ( as the family of optimal SUOT plans between Ps and Q*
Under these assumptions, we note that:
L
Av - Vh (- (fz)
(Zx)) dPp((x)) + m²]
(43)
k=1
21

---

Published as a conference paper at ICLR 2025
K
Vx
(-
x + m²)
(44)
k=1
A (efk 710
(45)
where the transition from (43 to 44
follows from the proof of Theorem In particular,
J (cu(
(1))) (xx)
Thanks to 45 and (33. one can decompose gap 82 as follows:
52
inf ((()))))))))
EUPUKK Nu. m)
Mu.KI,
A 12) -
K
cf. 38
A (Va 1k) (2((FenTe))
K
(F) def Va 760 Das
where
In what follows, we analyze gap 62 for each particular h E [1 K]:
/ { vh ( (F Th - TH (- (dF))
82 (Fa) =
(46)
X
- [ (f2(() - Fo()))
+
/ (xx) - ( (z(w))
Vina (ft) (ax)) dl (FK)
(47)
X
(x)) ot [II
+ +/ʃ (f2(g) fo())dd011
/ (Ce(FA.0) v) (1(1)) d( + (1)
X1 3
ToalERsTT ((w J, Talw)aor()))
/ (ca(In,)) (rtt I Sileddtttn
X y
-/ / (((o))) - / F(lu)dr((ma
N4.
X,.
/ (0L(Z..)) - (ax (xx))) (xx.
R/
TTL v)d ā (xx, (- (48)
Xx
22

---

Published as a conference paper at ICLR 2025
Tel (ddn -rh (2x, Ty (xx)) daPalaw)
2
7
(49)
2,
/ (Dx - Dr (rL, (xx)) (
X
11.3
/ |1 Th
(50)
X4 y
Transition from (46 to
47
follows from the first-order property of convex differentiable function
Th:
T. 2 2b -V -u
(2 Ty (rx))]
Transition from
48
49
Since T
minimizes y Ih (ak: ( v) then:
Dr (1x, ()d1)())z) (1( Te (mx)) 2 0.
(51)
Also note that:
-(f2) (xL) -- (cn ( =1EN2,DE)
- -b,
max
3
and, therefore,
(fE) 2 Vet((-b 17.
(52)
The combination of 51 and 522) validates the transition.
Part 4. Bringing it all together.
Summarizing inequalities for 61,2 and 82 50
yields:
(Fe, 71)) + 82 (fx)
nB
l1 T (ow)) + / [1w (M--T(())1BB/ (ww)llccr((llloa))))))d))))
NEXY
Nxx
118
where the last inequality obtained analogously the similar one in the proof of Kolesov et 2024a Theorem 4.2).
Aggregating the inequalities for 61,6 + 62, with weights d finishes the proof of the Theorem:
K
m
K
61 + 82 =
Σ
A6 + -2,6))
A
WF(Fx
((
2
Xx
B EXTENDED DISCUSSION OF RELATED WORKS
Below give overview of the continuous OT solvers with a specific focus unbalanced setup.
Neural OT/UOT solvers. Neural network-based continuous OT is a popular and fruitful area of recent generative modelling research. Some keynote solvers include: (Makkuva et al. 2020; Korotin ct al.] 2021a] Amos 2023) (ICNN-based, quadratic cost); (Seguy et al.] Daniels et al, [202]
Mokrov et al.
2024) (Entropic OT); (Vargas et al. 2021
De Bortoli al. 2021 Gushchin al.
23

---

Published as a conference paper at ICLR 2025
2023
Tong et al.
2023 Shi et al. 2024 Gushchin al. 2024b
Korotin al. 2024 Gushchin et al. Kornlovv et al. 20244] Klein et al.
2024a). (Schrödinger bridge); (Liu et al. 2023; Tong et al.
2024.
2024
(Flow matching). Of special importance for our developed method are max-min -mn adversarial)
OT solvers based on (semi-) dual OT formulation (Rout et al.
2022 Korotin et al. 2023b a Fan et al.
2023 Asadulaev al. 2024 Korotin et al. 2022b 2021b). Recently, the adversarial methodology
has been extended to unbalanced setup (Choi et al. 2024a; Gazdieva et al. 2023] Choi et al. 2024b
opening up new intriguing research direction in the field of robust continuous OT.
C IMPLEMENTATION DETAILS
In this section, we provide additional details about experiments in our paper.
C.I DESCRIPTION OF TOY EXPERIMENTS
Moon, Spiral, and G-Gasssian Datasets. The datasets used in Section 5.1 were implemented by following (Choi et al. 2024b).
Datasets for Class Imbalance Experiments. For the distribution IP and P2, we employ the Gaussian mixture of 4))4 + (4),,,,,,,,.... and ( 4), 0.4%) + [ 0.4.. respectively.
Datasets for Outlier Experiments. We consider three marginal distributions, denoted as PI, 12,, and P3. For P1 and P2,, generate datasets consisting of 95% in-distribution data and 5% outliers. In contrast, Py consists solely of in-distribution data. The in-distribution data for each marginal follows a Gaussian mixture model with four modes and uniform weights, i.e., Ek=1 W The means and variances for P1. F2, and P3 are defined follows:
P1: m (-5, 11), m2 1), my (1.- m4 (-1,5), with o 0.1, P2: m1 =--5, 1), m2 (5, my -11,5), m4 = (-11-- with 0.1, * IPg: m1 = (-5,,)), m2 (5. 0). mg = (0. 5), m4 (0, -5), with o 0.1,
respectively. For the outlier distributions of IP, and P2, we again Gaussian mixture model with four modes, Σ [M o²1), with the means and variances defined as:
* PP outliers:
(10, 2), m2 (10, 1), my = (10, 0), = (10,-1), with 0.02, PP outliers: m. = (-10, 1), m2 = (-10, 0), ma = (10,-1), m4 (10, -2), with
σ 0.02,
respectively.
Implementation Details. For the experiments in Section and the class imbalanced experiments, the number of iterations of 10K. For the outlier experiments in Section 5.2 we employ the number of iterations of 20K. We use the batch size of 1024. We update m with the Adam Optimizer with learning rate of 10-3 and (3 (2) (0, 0.9). For other hyperparameters, follow
the experimental settings of (Kolesov et al. 2024a).
Evaluation metric. As discussed in Section 5.1 utilize the Wa and L2 metrics, defined as (T P1, Q*) and 1|7, - I1 respectively. Here, Q" and 7i represent the true barycenter and transport maps. To approximate (Tī Q*), we solve the discrete Unbalanced Barycenter problem. Specifically, we sample 2000 points and solve the Unbalanced Optimal Transport (UOT) problem from P1 to P2 using the Python Optimal Transport (POT) library (Flamary et al. discrete UOT map Ti is then used to define T X₁ ld + X27 TeTiot and Q² 7f P.. Using these
2021). The obtained
discrete samples and transport maps, we compute the approximation of the metrics.
Training/Inference Time. The training and inference times of our model, along with comparisons to baseline methods, are presented in Table and Table 3 In the case of toy experiments, the training time of U-NOTB is comparable to that of NOTB (Kolesov et al.) 2024a). The inference time of U-NOTB is typically 2- 10 times slower than NOTB, This gap is due to the additional computational complexity introduced by rejection sampling, which requires calculating the c-transform of the
24

---

Published as a conference paper at ICLR 2025
potential function fk. i.e., [x( R c(z, ((t)) Thus, this process involves forward pass through the learned potential function. Moreover, requires to compute the cost functional, which is an additional burden for StyleGAN experiments. However, would like to highlight that thanks to our proposed sampling method, our solver gains robustness to outliers and class imbalancedness.
Experiment
2D Toy ($5.1)
Class Imbalance ($5.
Outlier ($5.2
Shape-Color Color ($5.3
UOTM
4-5 mins
Minibatch UOT
2h
NOTB
6-7 mins
6-7 mins
42-45 mins
U-NOTB
6-7 mins
6-7 mins
42-45 mins
3h 20 mins RTX 3090Ti
GPU
RTX 2080Ti
RTX 2080Ti
RTX 2080Ti
Table 2: Training time for various experiments.
Experiment
2D Toy ($5.1)
Class Imbalance ($5.2
Outlier ($5.2
Shape-Color ($5.3
UOTM
0.001 sec
Minibatch UOT
0.001 sec
NOTB
0.001 sec
0.001 sec
0.001 sec
0.02 sec
U-NOTB
0.001 sec
0.003 sec
0.003 sec
0.18 sec
GPU
RTX 2080'I
RTX 2080Ti
RTX 2080T1
RTX 3090TI
Table 3: Inference time for various experiments. We report the inference time of 1000 samples for
synthetic experiments (55.T. $5.2). For StyleGAN experiment ($5.3), we evaluated with 220 samples.
C.2 STYLEGAN EXPERIMENTS
Unless otherwise stated, we follow the implementation (Kolesov et al. Adam Optimizers with (91. 32) 0.9).
2024a). Note that we use
The exact hyperparameters for all experiments are described in Table
Experiment
2D Toy 5.51
Class Imbalance ($5.)
Outlier ( 3.2
Shape- ($5.3
D
372
K
2
1,20,200
1,20,200
10
batch size
1024
256
256
64
Epochs
10K
10K
20K
3K
NT
10
1/2
1/2
1/3
1/2
A2
1/2
1/2
1/3
1/3
Às
1/3
Divergence 4/1
balanced
KL
KL
Softplus
Divergence 2/2
KL.
KL
KL
Softplus
Divergence 4'a
KL
MLP
MLP
MLP
ResNet
Tkw
MLP
MLP
MLP
ResNet
le-3
le-3
1e-3
le4
ITT
le-3
le-3
le-3
le-7
Irm
le-3
le-3
le-3
le-2
Table 4:
settings of Algorithm for various experiments.
D EXTENDED EXPERIMENTS
In this section, provide additional experimental results. In Appendices and D.2 and demon-
strate the performance of our solver and baselines in OT/SUOT barycenter problem for Gaussian distributions with computable ground-truth solutions. In Appendix D.3 we conduct high-dimensional
25

---

Published as a conference paper at ICLR 2025
experiment demonstrating the practical advantages of our solver, i.e., its ability manipulate images through interpolating image distributions on the image manifolds.
D.1 BALANCED OT BARYCENTERS FOR GAUSSIAN DISTRIBUTIONS
In the balanced OT barycenter problem for Gaussian distributions, the ground-truth barycenter is known to be Gaussian and can be estimated using the fixed point iteration procedure (Alvarez-Esterl et al. 2016). In this section, we tested our solver in this balanced setup keeping in mind that for large unbalancedness parameters т, should provide good approximation of balanced OT barycenter problem solutions. Ultimately, we show that while our solver is designed to tackle an unbalanced OT barycenter problem, its performance in the different balanced OT barycenter problem is comparable to the current SOTA solvers.
We consider K 33 Gaussian input distributions with weights ' = = = 0.25, À = 0.5, and quadratic
OT cost functions c,(r, y) | -/1² following the experimental setup which firstly introduced in (Kolesov et al.] 2024bb and also used in (Kolesov et al.) 2024a] Appendix B.I). We perform comparison with three recent approaches for continuous barycenter estimation (Kolesov al.) 2024a NOTB), (Kolesov et al. 2024b EgBary) and (Korotin et al., 2022a WIN). For EgBary approach, which solves the Entropy-resuguularieed OT (EOT) barycenter problem, we consider small entropy regularization parameter E = 0.01. For completeness, we also tested the performance of the classic approach (Cuturi & Doucet 2014 FCWB) which approximates the barycenter by a discrete distribution on fixed number of free-support points.
For our U-NOTB solver, we consider KL divergencies and varying unbalancedness parameter E [1, 10 102, 103, 10 10 expecting that for large т, our solver will provide good results in solving the balanced problem. We assessed the performance solvers using the weighted unexplained
f- (- 1% where Q" is given ground-truth
variance percentage metrics L2-UVP(T)) = 100
Var(C"")
OT barycenter. For the continuous baseline solvers (NOTB, EgBary, WIN), report the results given in (Kolesov et al. 2024a| Table 3). Specifically, for our solver and EgBary which learn the optimal plans, we consider their barycentric projections. For the assessment of FCWB, we first calculate the discrete barycenter using large number of samples from input distributions IPg. Then we compute the optimal plans between the input and this barycenter distributions (Flamary et al. [02T]. and consider their barycentric projections. Table 5 presents the weighted sum of LV-UVP values w.r.t. the barycenter weights λg.
Method/Dim
2
16
Ours (7 1)
3.10
2.39
1.64
1.44
1.39
Ours (r = 10²)
0.03
0.09
0.06
0.08
0.09
Ours (r = 1
0.02
0.03
0.03
0.05
0.07
Ours (r = 10")
0.02
0.03
0.04
0.05
0.08
Ours (T = 10t)
0.02
0.03
0.04
0.05
0.08
NOTB
0.01
0.02
0.04
0.04
0.08
EgBary
0.02
0.05
0.06
0.09
0.84
WIN
0.03
0.08
0.13
0.25
0.75
FCWB
2.17
6.51
18.68
35.80
100.47
Table 5:
L--UVP for our method, NOTB, EgBary (c 0.01),
WIN and FCWB, D = 2, 4, 8, 16, 64.
The Table shows that for T = 102, our approach gives the results comparable with the current SOTA solver for continuous balanced barycenter estimation NOTB. Further increase does not help to improve the results which can be explained by related numerical instabilities. At the same time, discrete OT barycenter solver provides the worst results and the CuV metric increases drastically with the increase of dimension D. It is an expected behaviour, since the discrete distributions poorly approximate the continuous ones. This aspect was previously investigated and justified in (Korotin et al. 2022a 1).
26

---

Published as a conference paper at ICLR 2025
D.2 SEMI-UNBALANCED OT BARYCENTERS FOR GAUSSIAN DISTRIBUTIONS
A recent preprint Nguuen et al., 2024) provides an iteration procedure for calculating unbalanced SUOT barycenter for Gaussian distributions using the quadratic cost and KL divergences. We test our solver in this unbalanced setup for different parameters and show that it consistently outperforms the SOTA balanced solver (Kolesov al., 2024a NOTB).
Method
7 = 10
T = 10
7
T= 10
T= 10
T
10
7 7= 10
Ours
0.31
0.20
0.15
0.17
0.17
0.17
0.17
NOTB
1.71
1.22
0,47
0.21
0.20
0.19
0.19
Table 6: BW3-UVP between the learned and ground-truth unbalanced barycenters for our method and NOTB. The results are given for varying parameters 7 = 10-2, 10-1, ,,, 10, 102, 109, 10
We consider the same set of Gaussians distributions, weights and quadratic cost function as in the experiment with balanced Gaussians, see Appendix [.I] As the baseline solver, we consider SOTA balanced solver for continuous barycenter estimation (Kolesov al. 2024a] NOTB). To calculate the ground-truth barycenter, we run the Hybrid Bures- Wasserstein Gradient Descent algorithm (Nguyen jet al. 2024 84.3) for computation of SUOT barycenter problem between Gaussians. The algorithm designed for the case of KL divergencies and is endowed with the unbalancedness parameter which we denote as т.
We consider parameters T E [10-3 2 1001 1, 10, 10², 103, 10*]. For each parameter, calculate the corresponding ground-truth SUOT barycenter and train our U-NOTB solver with KL diver- gencies. Then we assess the performance of our solver and NOTB by measuring the weighted Bures- Wasserstein unexplained variance percentage metric for an implicitly given barycenters Q:
BW Q')
BW2-UVP(C
Q def
100
QQr)
where BW3 Q2) W3(N((a, Eg, (H Eo. the Wassereasnn metric, M4o Eg denote means and covariances of corresponding distribution Q. see (Korotin et al., 2021c). To get
the samples from learned barycenter Q. we first sample the input points 24. either form the input distributions P& (for NOTB) or from the left marginals the learned plans (for our U-NOTB). Here sampling from the left marginals of the plans can be done using the rejection sampling procedure, (4.2) Then we sample new points from the learned barycenter by passing the points 2y through the learned stochastic or deterministic maps T3 which parametrize the learned plans. Table presents the weighted sum BW3 values w.r.t. the barycenter weights A
We that our unbalanced U-NOTB solver outperforms NOTB for all considered unbalancedness parameters т. The difference is especially visible for small parameter т. It is expected since for small the solutions of the unbalanced barycenter problem significantly differ from the balanced one. For large this difference is getting smaller, thus, balanced NOTB solver manages to approximate the barycenter quite well but still worse than ours.
D.3 MANIPULATING IMAGES THROUGH INTERPOLATING IMAGE DISTRIBUTIONS ON THE IMAGE MANIFOLD
In this section, we present a novel real-world application of U-NOTB for high- -igh-dimensional image manipulation. Specifically, we address scenarios characterized by class imbalances in the marginal distributions. Ultimately, our experiment highlights the potential of our approach to manipulate images by learning barycenters between distinct image distributions, enabling controlled transitions across semantic attributes. Moreover, our model also demonstrates robustness to class imbalancedness in this practical task compared to other baselines.
Given multiple image distributions {P we aim to simultaneously learn all UOT barycenters Q for arbitrary tuple (A satisfying ΣK d = 0 < A 1. Specifically, we learn all intermediate barycenters simultaneously by parametrizing the transport maps {Tk K functions of the condition variable >i:K. expressed as T, (A).K. +). Note that a similar conditioning on barycenter weights is also employed in Algorithm 2 of (Fan et al. 2021
27

---

Published as a conference paper at ICLR 2025
With the learned conditional transport plan, we perform image manipulation through transforming given image x' ~ Pş to the barycenter point ~ Qx.x. In particular, we consider P1 and P2 as collections of images of young individuals (age between 5 to 20) and elderly individuals (age of over 50), respectively, from the FFHQ (Karras et al. 2019) dataset. It is important to note that the ratios of females to males in Pi and 2 are approximately 1:2 and 3:1, respectively. Thus, the data inherently exhibits class imbalance across the marginals. Here, we provide the precise experimental settings:
Marginal distributions P2. The marginal distributions Pi, I2 consist of FFHQ images of young individuals (5 age 20) and elderly individuals (age > 50). We randomly partition each distribution into 90% for the training set and 10% for the test set.
StyleGAN Manifold. Let G : W + 1332556256 be the pretrained StyleGAN2-ada
Karras et al. 2020) generator that maps W space 256-dimensional images. Let W be the pretrained FFHQ e4e encoder (Tov et al. 2021 Note that E maps image x E R8x2566256 the corresponding vector w E W i.e. G(w) æ. Note
E: 321xx200x22 250
that W+ C I 18185122
To restrict our transformed images to FFHQ images, we parametrize our transport map by follows: (z) == G Skw E(x) where æ E
and
Sk R 18x512 - R18x112
General Transport Cost. To guide image manipulation in meaningful direction, we adopt a quadratic cost function in the W latent space, defined as:
c(z, y) = al| E(x)11
(53)
Here, we set a = 10-2
Network Parametrization. Let A2) = (1 11-1,,
1, t). To generate QAna from x. Pk. we condition the transport map Tkn with t as follows:
(t, G (Sk,o (()))
(54)
We parametrize (t, + (t, (t,,)) z) and S2. (t. (1-()))) 2). Note that S1 and S2., transport the latent vector of young individuals to barycenter point, latent vector of elderly individuals to barycenter, respectively. Thus, we selected this heuristic parametrization of Sk. to ensure that (1, x:) R when z 0 and S2 (4, v) & y when Av Furthermore, the potential functions fI, J2 are also conditioned by variable / as follows:
f10(t, = E(x) mo(t), f2,0(t,) Vo(t,E(x)) + mt) 1-t
(55)
t
where Vo : R x R18x512 + R. Note that (1 t) + tf2,0 = mt) Here, mo(t) can be regarded as the parametrization of the congruence constant at time To avoid the numerical instability of (55) with respect to the time variable f, we sample from a uniform distribution within the interval [0.05, 0.95].
Results. As discussed, the two marginals PI, P1, exhibit a class imbalance with respect to gender. P1 contains more than 60% of female individuals, while I2 consists of less than 30% of female individuals. Thus, we evaluate the quantitative performance of our transport map by measuring how well the transport map preserves gender in such a class imbalanced case. Specifically, we measure the gender preservation accuracy when transforming a young individual to barycenter Ti (1) at t t= 0.9. For our model, we report the gender preservation accuracy for the accepted images. As comparison, we include the results from the NOTB method applied to the entire test dataset. For the baseline model, we precompute the global latent direction Egiobal := Eyy-,, [E()) - Ez-F, (E(x))
and measure the accuracy of the gender alignment between P1 and x + Uiglbbll-
Method
Baseline
NOTB
U-NOTB
Acceptance Rate Accuracy
67.9% (592/ 872)
69.8% (609 / 872) 63.4% (553 / 872)
88.0% (521 592)
Table 7: Gender preservation accuracy for x ~ Pı and Ti (t, at t = 0.9. U-NOTB can perform rejection sampling to address class imbalancedness.
As shown in Table 7 our model achieves a significantly high accuracy of 88%, highly outperforming other methods, which achieve less than 70%. This demonstrates that our model performs robustly
28

---

Published as a conference paper at ICLR 2025
2
DCO
(a) Accepted samples from P1.
(b) Corresponding barycenter samples.
Figure 6: Examples of randomly selected (a) accepted samples of PI and (b) its corresponding barycenter at & 0.9. Note that the acceptance rate of U-NOTB is 67.9%. Approximately, 36% of the accepted samples female individuals.
Figure 7: Examples rejected samples of U-NOTB. Approximately 80% rejected samples are female images.
under class-imbalanced conditions. Moreover, only 36. 1% of the accepted samples were female. Given that over 60% of the data in IP consisted of female samples, this result indicates that significant proportion of female samples were rejected to address the class imbalance. For qualitative examples, please refer to Figure and Figure 7
Evaluation Metric. classify the gender of images, we use the open python library called OpenCV. Specifically, adjusted the code from the following github address:
com http:///ttt..a..ooo.//eeeeeeee/a
Implementation Details. We set (w,, )3) (Softplus, Softplus) and 0.1. We employ the learning rate of 10 -5 batch size of 4, and the total number iterations of 5K. For all the networks
N NN Vo employ the same embedding approach: the time variable & is embedded
29

---

Published as a conference paper at ICLR 2025
D
512x18
K
2
0.1
batch size
Iterations
5000
Nr
A
1/2
X2
1/2
Divergence 1/1
Softplus
Divergence 4/2
Softplus
fE.0
Transformer
Thw
Transformer
fr.
le-5
le-5
Irm
le-5
Training Time
4-5
GPU
RTX 3090Ti
Table 8: Hyper-parameter settings and training time for FFHQ experiment.
into 512-dimension vector temb E IR 1x522 through sinusoidal embedding, and concatenated with E 182 For N these embeddings are processed through two transformer layers, each with four attention heads. Then, we pass through the linear layer with both input and output channels of 512, resulting the output of size IR 18x512 For the potential Vg. we use pass through four transformer layers, each with four attention heads. Then, we aggregate features by taking mean of the tokens, which is then passed through the final linear layer. Additionally, for the network mg we embed by sinusoidal embedding, and then pass it through 2-layered MLP. The outline of the hyperparameter settings and the training time are described in Table
E LIMITATIONS AND FUTURE RESEARCH DIRECTIONS
In practice, we found that the practical optimization procedure (Algorithm may work unstably. In particular, the training may diverge under improper hyper-parameters selection or the resulting quality might depend on the random seed. We hypothesise that this behaviour is due to the reliance adversarial training and leave its thorough investigation to future research. From the theoretical side, the recovered semi-unbalanced OT barycenter is not guaranteed to be unique ($2.2). Resolving these limitations paves way for future work. Another appealing direction for future research is
to adopt alternative robust OT/barycenter variants, e.g., (Nietert et al. 2022 Buze 2024). for the continuous barycenter setup.
30


Reviews:
[{"review_id": "Ji13YodqdH", "reviewer": "Reviewer_nQfu", "summary": "This paper proposes a neural network-based method to estimate continuous barycenter via the dual formulation of the semi-unbalanced OT.", "questions": "Q1. How does the time efficiency of U-NOTB compare with other baselines?\n\nQ2. The experiments seem to lack details about the size of the training/testing sets and some training specifics, such as learning rate and other hyperparameters.\n\nQ3. Since the overall training is similar to GANs with a max-min approach, is the training of U-NOTB stable?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "1. The proposed method may be the first continuous robust barycenter estimation approach with proper theoretical support and practical validation.\n\n2. The proposed method is robust and can estimate the barycenter based on the data containing the outliers.", "weaknesses": "My main concern is the novelty of the paper. The primary method seems to be a combination of barycenter calculation, neural optimal transport, and semi-unbalanced optimal transport. In my view, simply combining methods may not be sufficient for publication at ICLR, so I give a negative score. I will consider raising my score if the authors can demonstrate unique contributions that go beyond a straightforward combination of these methods.", "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": "lGP1prHkXR", "reviewer": "Reviewer_MXNH", "summary": "This paper proposes a scalable approach for estimating the robust continuous barycenter by using the dual formulation of the (semi-)unbalanced OT problem. They first attempt to develop an algorithm for robust barycenters under continuous distribution setup. They model this problem as a min-max optimization problem and provide theoretical underpinnings, along with experimental results on synthetic and real-world datasets to demonstrate the robustness and adaptability of their method.", "questions": "- Given that the Wasserstein barycenter of Gaussian distributions has a closed-form solution, could the authors provide experimental verification for the Gaussian distribution case?\n- Can the robustness parameter $\\tau$ be automatically optimized during training rather than manually tuned? If not, for different proportions or types of noise, different $tau$ is usually required. Will this limit the practicality of the scheme?\n- How does the method work with high-dimensional data? In other words, the support size of the measures may be large in practical applications. While the support size of the measures in your experiments seems rather small.\n- Why don't you compare the methods [1,2] based on robust OT? Please add corresponding experiments.\n\n[1] Nietert S, Goldfeld Z, Cummings R. Outlier-robust optimal transport: Duality, structure, and statistical analysis[C]//International Conference on Artificial Intelligence and Statistics. PMLR, 2022: 11691-11719.\n\n[2] Wang X, Huang J, Yang Q, et al. On Robust Wasserstein Barycenter: The Model and Algorithm[C]//Proceedings of the 2024 SIAM International Conference on Data Mining (SDM). Society for Industrial and Applied Mathematics, 2024: 235-243.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper introduces a continuous SUOT barycenter estimation method that addresses the issue of robustness against outliers and imbalances in real-world datasets.\n- Rigorous derivation of the SUOT-based framework and solid theoretical support for the proposed model give the approach a solid foundation.", "weaknesses": "Experiments are not sufficient and the baseline method is not comprehensive.", "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": "4eg3BLndAA", "reviewer": "Reviewer_hyxj", "summary": "In this paper, the authors investigate the problem of finding the barycenter with semi-unbalance optimal cost (SUOT) from perspective of dual theory. In particular:\n\n(1) They derive the theory for the dual form of SUOT barycenter with reformulation, and necessary condition based on marginal plan and conditional plan\n\n(2) Then, from the necessary condition,  they propose the an algorithm (Algorithm 1) to approximate the optimal transport plan from the solution. This algorithm is based on training the neural network of optimal plan.  \n\n(3) Then, they empirically validate the efficacy of the algorithm on both synthetic and real dataset.", "questions": "See Weaknesses section.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper provides an important and novel dual theory for the SUOT barycenter problem.\n\n2. The proposed algorthm works well and persistent to the imbalance and outliers. \n\n3. The paper is well-structured and easy to follow.", "weaknesses": "1. The necessary condition seems to be not difficult to obtain, which makes it not convincing enough. So, would it be possible to obtain some sufficient conditions? \n\n2. It would be better if the authors can highlight the advantages of finding the optimal plan over the optimal solution. \n\n3. As stated in the limitations, the authors did not provide a rigorous analysis for the convergence of the algorithm.", "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": "tfPTTl2Ozo", "reviewer": "Reviewer_niQu", "summary": "In this work, authors propose to solve the continuous unbalanced optimal transport barycenter. To do that, they derive a dual formulation to the problem, which is a min-max problem over potentials and conditional distributions. They solve the problem by parametrizing the potential and conditional distributions with neural networks. Finally, they show that their method works on three experiments, demonstrating that they recover the right barycenters, that the barycenter is robust to outliers and to class imbalance, and that it can handle using different costs.", "questions": "A term seems to be missing in the definition of $\\psi$-divergence, see Definition 1 in [2].\n\nThe sentence line 227 is not clear to me. The $m$-congruence is a constraint of the problem. So I do not understand why it is written \"if the potentials satisfy the m-congruence, then the optimal value of the SUOT barycenter can be derived by solving (8)\". The SUOT barycenter can always be derived by solving (8) and the potentials necessarily satisfy this constraint when solving (8).\n\nIn Corollary 1, it is stated that the sup is taken over $f_{[1,K]}$. Isn't it also taken over $m$?\n\nIn Theorem 2, equation (12), shouldn't it be an argmin?\n\nIn Section 5.1, it is stated that the UOT problem is equivalent with the OT problem between rescaled distribution. Is it truly equivalent? And could we solve the OT barycenter between the rescaled distributions instead of solving the SUOT barycenter?\n\nIs it proved that $T_{1\\\\#}\\mathbb{P}_1$ gives the UOT barycenter?\n\nI think that the reference [3] is missing. In their experiments, they compute robust OT barycenters via unbalancedness.\n\nTypos:\n- Line 11: The first sentence of the abstract feels weird: I don't see what is the \"common challenge\"\n- Line 107: \"This transition is valid since we work the infimum in weak\"\n- Line 142: \"wights\"\n- Line 365: what is $T$ in $T_1 = \\lambda_1 Id + \\lambda_2 T$\n- Legend of Figure 4: $\\mathbb{P}_0$ -> $\\mathbb{P}_1$ and $\\mathbb{P}_1$ -> $\\mathbb{P}_2$\n\n[1] Kolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., & Korotin, A. (2024). Estimating Barycenters of Distributions with Neural Optimal Transport. arXiv preprint arXiv:2402.03828.\n\n[2] Séjourné, T., Peyré, G., & Vialard, F. X. (2023). Unbalanced optimal transport, from theory to numerics. Handbook of Numerical Analysis, 24, 407-471.\n\n[3] Séjourné, T., Bonet, C., Fatras, K., Nadjahi, K., & Courty, N. (2023). Unbalanced optimal transport meets sliced-Wasserstein. arXiv preprint arXiv:2306.07176.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is overall well written and nice to read.\n\nIt provides one of the first method to solve continuous UOT barycenters leveraging a dual formulation and neural networks parametrizations, which is an interesting contribution.\n\nSeveral convincing experiments are done demonstrating that the method learns well the UOT barycenter, that it is robust to outliers and class imbalance, and that it works with several costs.", "weaknesses": "The main weakness in my opinion is that the method feels really incremental compared to [1]. The main difference is that it it adapted to the UOT problem, which just changes slightly the formulation of the dual, and how to sample from the barycenter for inference.\n\nAnother weakness, which is classical with UOT, is that the choice of the unbalancedness parameters does not seem easy.\n\nThe method also needs to solve a min-max problem, which is probably unstable and very costly.", "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)