Energy-Guided Continuous Entropic Barycenter Estimation for General Costs
Authors: Alexander Kolesov
Petr Mokrov
Igor Udovichenko
Milena Gazdieva
Gudmund Pammer
Anastasis Kratsios
Evgeny Burnaev
Alexander Korotin
Venue: NeurIPS 2024 spotlight (2024)
Keywords: energy-based model
generative model
optimal transport
entropic optimal transport barycenters
general optimal transport cost
Primary area: generative_models
TL;DR: We propose a new energy-based method to compute entropic optimal transport barycenters with general cost functions.

Abstract: Optimal transport (OT) barycenters are a mathematically grounded way of averaging probability distributions while capturing their geometric properties. In short, the barycenter task is to take the average of a collection of probability distributions w.r.t. given OT discrepancies. We propose a novel algorithm for approximating the continuous Entropic OT (EOT) barycenter for arbitrary OT cost functions. Our approach is built upon the dual reformulation of the EOT problem based on weak OT, which has recently gained the attention of the ML community. Beyond its novelty, our method enjoys several advantageous properties: (i) we establish quality bounds for the recovered solution; (ii) this approach seamlessly interconnects with the Energy-Based Models (EBMs) learning procedure enabling the use of well-tuned algorithms for the problem of interest; (iii) it provides an intuitive optimization scheme avoiding min-max, reinforce and other intricate technical tricks. For validation, we consider several low-dimensional scenarios and image-space setups, including *non-Euclidean* cost functions. Furthermore, we investigate the practical task of learning the barycenter on an image manifold generated by a pretrained generative model, opening up new directions for real-world applications. Our code is available at https://github.com/justkolesov/EnergyGuidedBarycenters.

Energy-Guided Continuous Entropic Barycenter Estimation for General Costs
Alexander Kolesov
Petr Mokrov & Igor Udovichenko
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
Skolkovo Institute of Science and Technology
Moscow, Russia
Moscow, Russia kolesov@skoltech.c.
{p.mo Ap.mokrov, i. luudviicenkoo..oo @skoltech.
Milena Gazdieva
Gudmund Pammer
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute Moscow, Russia
Graz University of Technology Graz, Austria gudmund, pammerttugraz at
m. www.iiii@asoooteecchcooooo ru
Anastasis Kratsios
Evgeny Burnaev & Alexander Korotin Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
Vector Institute, McMaster University Ontario, Canada
kratssia@@mmmaseer..co
Moscow, Russia
{e. burnaev, a. .kkoootin@sskolteeeeeeeoo ru
Abstract
Optimal transport (OT) barycenters are mathematically grounded way of averag- ing probability distributions while capturing their geometric properties. In short, the barycenter task is to take the average of collection of probability distributions w.r.t. given OT discrepancies. We propose a novel algorithm for approximating the continuous Entropic OT (EOT) barycenter for arbitrary OT cost functions. Our ap- proach built upon the dual reformulation of the EOT problem based on weak ОT, which has recently gained the attention of the ML community. Beyond its novelty, our method enjoys several advantageous properties: (i) we establish quality bounds for the recovered solution; (ii) this approach seamlessly interconnects with the Energy-Based Models (EBMs) learning procedure enabling the use of well-tuned al algorithms for the problem of interest; (iii) it provides an intuitive optimization scheme avoiding min-max, reinforce and other intricate technical tricks. For valida- tion, we consider several low- oineeinenioonal scenarios and image-space setups, includ- ing Eun--dildd cost functions. Furthermore, we investigate the practical task learning the barycenter model, opening up new directions for real world application.. Our code is available at /https://ggthhie comp/ http:///gg....oooo/ooo//o//////e /wwp://oooeoeeeeeeeeeeeeeeaaaaaa
image manifold generated by pretrained generative
Introduction
Averaging is fundamental concept in mathematics and plays central role in numerous applications. While straightforward operation when applied to scalars or vectors in a linear space, the situation complicates when working in the space of probability distributions. Here, simple convex combinations can be inadequate or even compromise essential geometric features, which necessitates different way of taking averages. To address this issue, one may carefully select measure of distance that properly captures similarity in the space of probabilities. Then, the task is find procedure which identifies 'center' that, on average, closest to the reference distributions. One good choice for comparing and averaging probability distributions is provided by the family of Optimal Transport (OT) discrepancies [10]. They have clear geometrical meaning and practical
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

P
P2
P3 P4 Y-0"
3.0
2.0
1.5
0.5
Aztmutth
(a) The unfolded sphere.
(b) The sphere viewed from different viewpoints.
Figure I: Entropic barycenter Q" of N 4 von Mises distributions Fn the sphere (see $ 5.1 estimated
with our barycenter solver (Algorithm 1.. The used transport costs are (TK, y) 1/4 (arccos (Ia v))²
interpretation
89 197. The corresponding problem of averaging probability distributions using OT
discrepancies is known the OT barycenter problem M. OT-based barycenters find application in various practical domains: domain adaptation [79], shape interpolation [98]. Bayesian inference
[01 [02], text scoring [18] style transfer (80), reinforcement learning (77].
Over the past decade, the substantial demand from practitioners sparked the development of various methods tackling the barycenter problem. The research community 's initial efforts were focused on the discrete OT barycenter setting, Appendix [or more details. The continuous setting turns out to be even more challenging, with only a handful of recent works devoted to this setup Most of these works are devoted to specific OT cost functions, e.g., deal while others require non-trivial a priori selections [72 and have see 3 for detailed discussion.
17 59 55 32
82
4
with & barycenters 59
55
32
82
limiting expressivity and generative ability 117 14]
Contributions. We propose a novel approach for solving Entropy-regularized OT (EOT) barycenter problems, which alleviates the aforementioned limitations of existing continuous OT solvers.
1. We reveal an elegant reformulation of the EOT barycenter problem by combining weak dual form of EOT with the congruence condition ($4.1): we derive simple optimization procedure which closely relates to the standard training algorithm of Energy-Based models ($4.2).
2. We establish the generalization bounds as well the universal approximation guarantees for our recovered EOT plans, which push the reference distributions to the barycenter ($4.3).
3. We validate the applicability of our approach
various toy and large-scale setups, including the RGB image domain ($5). In contrast to previous works, we also pay attention to non-Fuclidean OT costs. Specifically, we conduct a series of experiments looking for barycenter on an image manifold of a pretrained GAN. In principle, the image manifold support may contribute to the interpretability and plausibility of the resulting barycenter distribution in downstream tasks.
Notations. We write K = {1, ..,.. K}. Throughout the paper X' C CR'' , C RD and X₂ C RD
(k: E K) are compact subsets of Euclidean spaces. Continuous functions on are denoted as C(Y). Probability distributions on X' are ((x). Absolutely continuous probability distributions on X' are denoted by (X) C P(X). Given P E P(X), Q E P()), we use [I( Q) to designate the set of transport plans, i.c., probability distributions on X' x y with the first and second marginals given by and Q respectively. The density of P E Pac (X) w.r.t. the Lebesgue measure is denoted by dP(x)
2
Background
First, we recall the formulations of EOT ($2.) 12.1 and the barycenter problem the computational setup of the considered HOT barycenter task (2.))
2..2 Next, clarify
2.1
Entropic Optimal Transport
Consider distributions P e (X), Q e P)c(), continuous cost function e X x y - R and regularization parameter >> O. The entropic optimal transportation (EOT) problem between IP and consists of finding a minimizer of
Q 15 78
EOT
del min =TTT(P
E c( v) - (E (na))= (r
{
INE
Note that T is not the only way to formulate EOT. One more popular and equivalent formulation 19 83 36 substitutes the conditional entropy term (x= in I with full entropy H (T).

---

A minimizer x E [I, Q) of (I) is called the EOT plan; its existence and uniqueness are guaranteed, see, e.g., [6 Th. 3.3]. In practice, we usually do not require the EOT plan ' but its conditional distributions H" (-[xx) E P()') as they prescribe how points x E X are stochastically mapped to y
42 $2]. We refer to n²" (+(+))) as the conditional plans.
Weak OT dual formulation of the EOT problem. The EOT problem permits several dual formula-
tions, In our paper, we use the one derived from the weak OT theory [39 Theorem 9.5]:
EOT (IP, Q)
sup E (x) SEcty)
EE
f(m)).
G
you [9]
where fO : + R is the so-called weak entropic e- c-transform f. The transform is defined by
Eq. 1.2] of the function (potential)
(x)
min { E cx, ( (EP()) ymuu
We the capital C' in le to distinguish the weak transform from the classic c-transform (8 11.6] or (c, --ttanssform (76] $2]. In particular, formulation (2) is not to be confused with the conventional
EOT dual, see 778 Appendix A].
For each æ E X, the minimizer pd E P()) of the weak c-transform (3) exists and is unique. Its density has particular form 78 Theorem 1]. Let Ze(f, x) def ʃn exp
[()) dy. then
duf del dy
1
f(v) - c(2, y)
zz(,
exp
By substituting into (3) and carrying out straightforward manipulations, we arrive at an explicit formula f (x) -e log Zo(f,2) x). see 778 Equation (14)].
2.2 Entropic OT Barycenter
Consider distributions P& Pac (x) and continuous cost functions cx(-.-) X2 x IR for k E K.
Given weights d₁ > with Σ d 1, the EOT Barycenter problem is:
K
L def
inf QEP()))
EOT cz.. (PA, Q),
(5)
The case where O corresponds to the classical OT barycenter 1 and falls of the scope of this paper. Note that the majority of previous research [20 22 30 25 68 67 consider a bit different but equivalent EOT barycenter formulation, i.c., which has the same minimizers. The objective (5) is known as Schrödinger barycenter problem [5] Table 1]. see the extended discussion in Appendix B.3 It worth noting that under mild assumptions the barycenter Q" which delivers optimal value of exists and is unique, see Appendix A.7
2.3
Computational aspects of the EOT barycenter task
Barycenter problems, such as (5). are known be challenging in practice (2). To our knowledge, even when I...... Pxx are Gaussian distributions, there is no direct analytical solution neither for our entropic case (e > 0), see the additional discussion App. C.4 nor for the unregularized case
3 0]. Moreover, in real-world scenarios, the distributions PR (k E K) are typically available explicitly but only through empirical samples (datasets). This aspect leads to the next learning setup.
We assume that each PA is accessible only by limited number of i.i.d. empirical samples
} ~ Pk. Our aim is to approximate the optimal conditional plans T (-(2%) X1 = = ......2...22 Nu
between the entire source distributions PA and the entire (unknown) barycenter Q" solving (5). The recovered plans should provide the out-of-sampll estimation, i.e., allow generating samples from T) where new is a new sample from IP, which not necessarily present the train sample.
This setup corresponds to continuous OT 72 59 It differs from the discrete OT setup 19 20 which aims to solve the barycenter task for discrete empirical distributions. Discrete OT are not well-suited for the out-of out-of-sample estimation required in the continuous OT setup.
Related works
The taxonomy of OT solvers is large. Due to space constraints, we discuss here only methods within the contimuous OT learning setup that solve the (E-)OT barycenter problem. These methods approximate OT maps plans between the distributions P, and the barycenter Q" rather than just their empirical counterparts that are available from the training samples. broader discussion of general-purport discrete/continuouss (E-)OT solvers is in Appendix B.1
3

---

Continuous OT barycenter solvers are based on the unregularized or regular- ized OT barycenter problem within the continuous OT learning setup. The works [59
Method
Admissible
Learns
Max considered
Regularization
costs
OT plans
data dim
EntropioQqudratic
general
yes
SD, images
with fixed prior
17
general
no
1x32x32 (MNIST)
Entropic (Sinkhom)
32 82, 55] are designed ex-
only 12
yes
256D, images (MINSST)
requires fixd prior
clusively for the quadratic Euclidean cost
32
only 1
yes
1x28&28
e ((x, def
2 || v1l-- The OT prob-
55]
only 12
only 1.
yes
yes
3x64x64 (CelebA, etc.)
1x28x28 (MINSSI)
Entropic
lem with this particular cost exhibits sev-
general
yes 256D. Gaussians only EntropickQuarratic
eral advantageous theoretical properties 4 $2] which are exploited by the afore- mentioned articles build efficient pro-
Ours
general
yes
3x64x64 (CelebA)
Entropic
Table 1: Comparison of continuous OT bary solvers
cedures for barycenter estimation algorithms. In particular, [39 32] utilize ICNNS 6 which parame- terize convex functions, and [82 relies on a specific tree-based Schrödinger Bridge reformulation. In contrast, our proposed approach is designed to handle the EOT problem with arbitrary cost functions CI,...,CK: Cx- In [72). they also consider regularized OT with -uclliiid costs. Similar to our method, they take advantage of the dual formulation and exploit the so-called congruence condition ($4) However, their optimization procedure substantially differs. It necessitates selecting a fixed prior for the barycenter, which can be non-trivial. The work [14] takes step further by directly optimizing the barycenter distribution in a variational manner, eliminating the need for a fixed prior. This modification increases the complexity of optimization and requires specific parametrization of the variational barycenter. In [17,, the authors also parameterize the barycenter as a generative model. Their approach does not recover the OT plans, which differs from our learning setup ($2.3) A summary of the key properties is provided in Table [I] highlighting the fact that our approach overcomes many imperfections of competing methods. We are also aware of the novel continuous OT barycenter solver [52]. This approach is more recent than ours and is significantly based on the ideas from our article. Because of this, we exclude from our comparisons.
4 Proposed Barycenter Solver
In the first two subsections, we work out our optimization objective ($411 and its practical implemen- tation (44.)) $4.3] we alleviate the gap between the theory and practice by offering finite sample approximation guarantees and universality of NNs to approximate the solution.
4.1 Deriving the optimization objective
In what follows, we analyze
(5) from the dual perspectives. We introduce L C())) - R:
K
C(Ji,... SK) del
E
(xx)) A6 FA-PA
Ae
E log Zoo (Jk, ark) P
Here Je denotes the weak entropic ca-transform (3) of Sk. Following 2
we that it coincides with -e log Zo Tx). Below we formulate our main theoretical result, which will allow us to solve the EOT barycenter task without optimization over all distributions on y.
Theorem 4.1 (Dual formulation of the EOT barycenter problem (proof ref ). Problem ( the following dual formulation:
permits
L
sup fk).
(6)
fu.....Sww Sw E C(V); Ex, ===0
We refer to the constraint Σ K Apfa = 0 as the congruence condition w.r.t. weights { The
potentials fa appearing in (6) play the same role as in
(2). Notably, when L(ft,.., fx) is close
to L" the conditional optimal transport plans n (-(xx), th E Xr, between Ps and the barycenter distribution O can be approximately recovered through the potentials fx. This intuition is formalized in Theorem 4.2 below. First, for fo E C()), we define
dx (xx. v) dudn (v)dPx((T))
and set Q E P(Y) to be the second marginal of f/ Theorem 4.2 (Quality bound of plans recovered from dual potentials (proof ref. Let (f Sa C()) be congruent potentials. Then we have
K
K
L L(hi. .......... fK) (.fi....,Kk) = <<)==>>> (mEI)()) 2 Σ XLK (Q'1()())
(7)
k=1
k=1
where के I1 k K are the EOT plans between Fx and the barycenter distribution Q"

---

Algorithm 1: BOT barycenters via Energy-Based Modelling
Input: Distributions Pk. E K accessible by samples; cost functions eu (x2, y) : Nr x y + R; the regularization coeff. < barycenter averaging coeff. A 0: Σ A MCMC procedure McMc_prro;; batch size
NNs fo.s: y R, s.t. Σ fo 0 (see $4.2
Output: Trained NNs fe- recovering the conditional OT plans between P& and barycenter Q*
for ster = 1, for = 1,2..... : K do Sample hatch {ar [ Pk:
Draw Y {yk with MCMC:
V = XCMC_proc ((22)))
Lu < - Au E f (v
L Σ k=
Le: Update 0 by using oi. 26
According to Theorem 4.2 an approximate solution {fe) of (6) yields distributions x which are close to the optimal plans T.. Each π is formed by conditional distributions pfg c.f. 4 with closed-form energy function, i.e., the unnormalized log-likelihood. Consequently, one can generate samples from /e using standard MCMC techniques (7]. In the next subsection, we stick the practical aspects of optimization of (6 which bears certain similarities to the training of Energy-Bass Based models [69 100 EBM].
Relation to prior works. Works 52 also aim to first get the dual potentials and then recover the
barycenter, see the discussion in 3 for more details.
4.2 Practical Optimization Algorithm
To maximize the dual EOT barycenter objective (6). we replace the potentials fa E C(V) for k E K with neural networks fø.k E O. In order to eliminate the constraint (6). we parametrize fo.k as 9o, parameterization automatically ensures the congruence condition -k=1 AnSa, 0. Note that θ= x ... x θạ and 0= (01....,0 E O. Our objective function for (6) is
EE= 90. where {go, TR - R, E Ox) are neural networks. This K
L(
Σ Au {
log Zo, (x)
(8)
The direct computation of the normalizing constant Zey may be infeasible. Still, the gradient of with respect to 0 can be derived similarly to (78) Theorem 3]:
Theorem 4.3 (Gradient of the dual EOT barycenter objective proof ref ). The gradient of L satisfies
K
0
(( = -
A E
0g fo.
00
y-u fea
With this result, we can describe our proposed algorithm which maximizes L using
TRAINING. To perform stochastic gradient ascent step w.r.t. 6, we approximate (9) with Monte-Cario by drawing samples from fo.x y) = du (y) & (Ix). Analogously to 78 $3.21, this can be achieved by simple two-stage procedure. At first, draw random vector 24 from P.., This is done by picking random empirical sample from the available dataset X₁- Then, we need to draw sample from the distribution f.r. fo.* Since know the negative energy (unnormalized log density) of fo.. by 4 we can sample from this distribution by applying an MCMC procedure which uses the negative energy function (fo,((y cs (*x, y)) as the input. Our findings are summarized in Algorithm [] Note that typical MCMC needs the energy functions, particular, costs Ck. to be differentiable. Otherwise, one can consider gradient-free procedures, c.g., [04].
In all our experiments, we use ULA 86, $1.4.1] as MCMC_proc. It is conventional MCMC fe.k
algorithm. Specifically, in order to draw sample ye pazz where He Xx. we initialize (() from the D-dimensional distribution (00 In) and then iterate the discretized Langevin dynamics:
(1+1) + (1) + 7 V (fo c(xb. v)) y= 2k
+ +nEi,
2
5

---

where & ~ ID), f0, .). L is number of steps, and > 0 is a step size. Note that the iteration procedure above could be straightforwardly adapted to a batch scenario, i.e., we can
simultaneously simulate the whole batch of samples Yx conditioned on x The particular values of number of steps L and step size y) are reported in the details of the experiments, see Appendix C An alternative importance sampling- based approach for optimizing (9) is presented in Appendix D
INFERENCE. We use the same ULA procedure for sampling from the recovered optimal conditional plans z > (-Iar)), see the details on the hyperparameters 1, in $
Relation to prior works. Learning distribution of interest via its energy function (EBMs) is
well-established direction in generative modelling research [69] 111 [28] [1O0]. Similar to ours, the key step in most energy- baargy approaches is the MCMC procedure which recovers samples from a distribution accessible only by an unnormalized log density. Typically, various techniques are employed improve the stability and convergence speed of MCMC, see, e.g., [27] 34 113]. The majority of these techniques can be readily adapted to complement our approach. At the same time, the primary goal of this study to introduce and validate the methodology for computing EOT barycenters in an energy-guided manner. Therefore, we opt for the simplest MCMC algorithm, even without the replay buffer 466, as it serves our current objectives.
4.3 Generalization Bounds and Universal Approximation with Neural Nets
In this subsection, we answer the question of how far the recovered plans are from the EOT plan T between P2 and Q. In practice, for each distribution I&; we know only the empirical samples X₁ = {z .k.-. ... 2x } ~ Pk. i.e., finite datasets. Besides, the available potentials fk. kt E K come from restricted classes of functions and satisfy the congruence condition. More precisely, we have fe 9k -k=1 Auga (54.2) where each gk is picked from some class G3 of neural networks. Formally, write (f1,... fk) F to denote the congruent potentials constructed this way from the functional classes gI....,9 9K. Hence, in practice, we optimize the empirical version of (8)
K
N
max E( .Jk....]] del (fi..... fx x)F
A
CL
(x)));
max (x) k=1
N
n=1
(h Jk) del
arg max (J)E...IREEF
C(f fRk)
A natural question arises: How close are the recovered plans The to the EOT plans गं between IPx and Q'? Since our objective (8) is a sum integrals over distributions Pg, the generalization error can be straightforwardly decomposed into the estimation and approximation parts.
Proposition 4.4 (Decomposition of the generalization error
ppooo ref.) ). The following bound holds:
Estimation error (upper bound)
Approstmation erior
cE XL (T((lT()) k=1 where J= def U1 (fi...., fk) E F}, and the expectations are taken w.r.t. the random real- ization of the datasets XI-MP,.... Xx.TRx. Here Repx, (JCr (FCt,PP)) is the standard notion of the representativeness of the sample Xx w.r.l. functional class JCa and distribution Pt, see App. A.5
2 k1 AgERepx, (J Pt,P))
K 22
L' max Ci Sk) (J...nx))
(10)
To bound the estimation eтror, we need to further bound the expected representativeness ERepx, x, ( Doing preliminary analysis, we found that it does not depend that much on the complexity of the functional class Fg. However, it seems to heavily depend on the properties of the cost Cj. We derive two bounds: a general one for Lipschitz (in .)) costs and better one for the feature-basee quadratic costs.
Theorem 4.5 (Bound on ERep w.r.t. Ca-trrnsform classes (proof ref.)) (a) Let F& C CC)). Assume that ct (x.y) is Lipschitz in x with the same Lipschitz constant for all yey. Then
ERepx, (FF .P)) S (NN,
(11)
(b) Let cx(zk, v) = ((())(( - v(2))12. F& be bounded (w.r.. the supremum norm) subset of C()). us Xx + RR 1" and y + D" be continuous functions. Then
ERepx, (FC+, (Pp) S &ONNE
(12)

---

Substituting
to 10 immediately provides the statistical consistency when N....., NK oc, i.e. vanishing of the estimation enror when the sample size grows.
11 or
The case (a) here is not very practically useful as the rate suffers from the curse of dimensionality. Still, this result points one intriguing property of our solver. Namely, we may take arbitrarily large set Fi (even Fs = !()))) and still have the guarantees of learning the barycenter. This happens because of Ca-transforms: informally, they make functions fü Fk smoother and "simplify" the set Fu. In our experiments, we always work with the costs as in (b). As result, our estimation error is O N₁ N==1/2); this is standard fast and dimension usually neural nets. They are indeed bounded, as required in (b), their weights are constrained.
convergence rale. In practice, Fk
While the estimation enror usually decreases when the sample sizes tend to infinity, it is natural to wonder whether the approximation error can be also made arbitrarily small. We positively answer this question when the standard fully-connected neural nets (multi-layer perceptrons) are used.
Theorem 4.6 (Vanishing Approximation Error proof ref Assume that it is non-affine and there is an E E R at which a is differentiable and o' & O. Then for every >> there exist K multi-layer perceptrons 9k TRD -> R. with activations a for which the congruent functions fk = gk -
Let
R be an activation function.
Auggs satisfy
K
Σ KL (TÉIIT(E) = - fx)e
Furthermore, each 9k has width at-most D 4.
Importantly, our Theorem 4.6 more than just result on universal approximation since it deals with (i) congruent potentials and (ii) entropic Ca-transforms. In particular, only specific properties of the entropic Ce-transforms allow deriving the desired universal approximation statement, see the proof. Summary. Our results of this section show that both the estimation and approximation errors can be made arbitrarily small given sufficient amount of data and large neural nets, allowing to perfectly recover the EOT plans z..
Relation to prior works. To our knowledge, the generalization and the universal approximation are novel results with no analogues established for any other continuous barycenter solver. Our analysis shows that the EOT barycenter objective (8) is well-suited for statistical learning and approximation theory tools. This aspect distinguishes our work from the predecessors, where complex optimization objectives may not be as amenable to rigorous study.
4.4 Learning EOT barycenter on data manifold
Averaging complex data distributions by means of EOT barycenter directly in the data space may be undesirable. In particular, for image data domain:
the entropic barycenter contains noisy images, see, e.g., our MNIST 0/1 experiment, $5.2 This is due to the "blurring bias" bias [15 49 of our entropic barycenter setup and reliance on MCMC.
searching for (entropic) barycenter is not very practical for standard OT cost functions like E is known that the true unregularized (e 0) (e-barycenter of several image domains consists of just some pixel-wise averages of images from these source domains, which is not practically useful.
To alleviate the problem, we propose solving the (entropic) barycenter problem on some a priori known data manifold M, where we want the barycenter to be concentrated on. In our experiments
(5.2) $533 these manifolds are given by pre-trained StyleGAN [50) generator models G Z + y; 2 is the latent space, M = G(Z). Technically speaking, adapt our Alogithm for manifold- constrained setup, we propose solving the barycenter problem in latent space Z with modified cost functions Ga.G (Mi,, z) == cn (æi,, G(z)). We emphasize that such costs are general (not L² cost!) because G is non-trivial StyleGAN generator. Hence, while our proposed manifold-constrained barycenter learning setup could be used on par with other OT barycenter solvers, these barycenter solvers should support general costs. In particular, the majority of competitive methods from Table [ are not adjustable to the manifold setup as they work exclusively with l².
Relation to prior works. While the utilization of data manifolds given by pre-trained (foundational) models ubiquitous in generative modeling, the adaptation of this technique for Optimal Transport barycenter is novel idea. Apart from work, this idea is exploited in follow-up paper (52).
7

---

5 Experiments
We assess the performance of our barycenter solver on small-dimensional illustrative setups ($5.1) and in image spaces ($5.2) (5.3). The source code for solver is written in the PyTorch framework and available at https: btaps://github issued in the form of convenient * *.ipynb notebooks. Reproduc- ing the most challenging experiments ($5.2) (5.3) requires less than 12 hours on a single 100 GPU. The details of the experiments. extended experimental results are in Appendix additional experiments with single-cell data are given in Appendix C.55
The experiments
Disclaimer. Evaluating how well our solver recovers the EOT barycenter is challenging because the ground truth barycenter is typically unknown. In some cases, the true unregularized barycenter (e = 0) can be derived (see below). The EOT barycenter for sufficiently small e > O expected to be close to the unregularized one. Therefore, in most cases, our evaluation strategy is to compare the computed EOT barycenter (for small e) with the unregularized onc. In particular, we use this strategy to quantitatively evaluate our solver in the Gaussian case, see Appendix C.4
5.1
Barycenters of Toy Distributions
A-h
y "(: [xi]
y y Kp
f-- x-
X) Pr
y-
6-
(a) The true unrogularized baryoenter for the twisted cost.
(b) Our EUT burycenter For the twisted cost (map from F1).
(c) unregularized baryeenter for E" cost.
(d) Our BOT barycenter for g²
cost (map from F1).
Figure 2: 2D swister example: The true barycenter 3 comets vs. one computed by our solver with
c 10 Two costs C& are considered: the twisted cost 2a 2b and
2c 2d).
2D Twister. Consider the map u : I2_ + R22 which, the polar coordinate system, is represented by
R- x [0, 2T) (r, 0) + (r; (0 r) mod 2T). The cartesian version of u is presented in Appendix Let PI, P2, PP be 2-dimensional distributions as shown in Fig. 2a For these distributions and
uniform weights A }. the unregularized barycenter (e 0) for the twisted cost ea (Z y)
[]lu(xx)= m(y) || can be derived analytically, see Appendix C. The barycenter is the centered Gaussian distribution which is also shown in Fig.
We run the experiment for this cost with We see that it qualitatively coincides with the true barycenter. For completeness, we also show the EOT barycenter computed with our solver
2a
= 10-2 and the results are recorded in Fig. 2b
for (Ix, }|1= y/l costs (Fig. (cc and the same regularization The true 82 barycenter is
estimated by using the
solver from POT package [33. We stress that
the twisted cost barycenter and E barycenter differ, and so do the learned conditional plans: the
e² EOT plan (Fig. 2d) expectedly looks more well-structured while for the twisted cost (Fig. 2b) it becomes more chaotic due to non-trivial structure of this cost.
Sphere. In this experiment, we look for the barycenter of four von Mises distributions In, supported on 3D sphere, see Figure [ The cost functions are eu (2x2,)) y) - arccos" (Jrk, v). the regularization is 10 2 The learned potentials fe. operate with ambient R33 vectors. When performing MCMC, we project each Langevin step to the sphere. Our qualitative results are shown on Figure T While the ground truth solution to the considered problem is unknown, the learned barycenter looks reasonable. This showcases the applicability of our approach to non- standard qua-quacaatic experimental setups.
5.2 Barycenters of MNIST Classes 0 and 1
Figure 3: Samples from the StyleGAN G defining the polluted manifold M.
A classic experiment considered in the continuous barycenter literature
32 55 82 17 is averaging of distributions of MNIST 0/1 digits with
weights (3. 2 in the grayscale image space X₂ = Xr = Y 1-11 10._1]22200 1]32332 The true unregularized
(e = 0) e -abarcyceterr images y are direct pixel-wise averages
z t2z of pairs of images I and 22
8

---

INPUT
P2-P2
X1-P3
TRUE y-Tplo
y-T26)
r=T2e)
WIN
SONI
2- 5=11
OURS
y-hilel y-1ijk y-00
y-
P-njj y 022i
OURS
OURS
(a) Maps from P1 to the barycenter. (b) Maps from P2 to the barycenter. (c) Maps from IP to barycenter.
Figure 4: Experiment on the Ave, celeba! barycenter dataset. The plots compare the transported inputs æ IP the barycenter learned by various solvers, The true unregularized 1 barycenter of Pi, P2, Ps are the clean
celebrity faces, see 555 $51.
SCWB K1 Tyllly 7.7Zg)
WIN
OURS (Dara Space) Bicke
OURS Matniiald Conserained y- B. CI
WIN
SCWB 1
OURS IDuta Spacel y- D
OURS IKandfoi Constra nus
Xy P Flbel
y- 0,0lo
1
4
4
/
(a) Learned plans from F, (zeros) to barycenter.
(b) Learned plans from P2 (ones) to the barycenter.
Figure 5: Qualitative comparison of barycenters of MNIST 0/1 digit classes computed with barycenter solvers in the image space w.r.t. the pixel-wise L Solvers SCWB and WIN only learn the unregularized barycenter (c 0) directly in the data space. In turn, our solver learns the BOT baryeenter in data space as well as it learn EOT barycenter restricted to the StyleGAN manifold (e 1))
coming from the 2² OT plan between O's (P:)) and I's (P2). In Fig. 5 we show the unregularized & barycenter computed by (32 SCWB], 55 WINJ.
Data space EOT barycenter. To begin with, we employ our solver to compute the e-regularized EOT bb-arcceeteer directly in the image space y for 10-2, We emphasize that the true entropic barycenter slightly differs from the unregularized one. To be precise, it is expected that regularized barycenter images close to the unregularized barycenter images but with additional noise. In Fig. 5) we see that our solver (data space) recovers the noisy barycenter images exactly as expected.
Manifold-constriined EOT barycenter: Following the reasoning from $4.4 we propose to restrict the search space for our algorithm to some pre-defined manifold M. discussed earlier, the
support of the image-space unregularized e--barycenter is certain subset of AM' def { =2 (upp )) x2 E (_ To achieve this, we train a StyleGAN (50) model G y with Z = 512 to generate some larger manifold M G(Z) which is expected to contain M'. Namely, we use all possible pixel-wise half-sums 21 2iz of digits as Φ and {1, 4, 7) as 22, see Figure with the trained StyleGAN samples. That is, our final constructed manifold M is polluted with additional samples (e.g., averages of digits O and 7) which should not to lie in the support of the barycenter. Then, we our solver with
2
10-2 to search for the barycenter of 0/1 digit
distributions X1, X2 which lies in the latent space Z w.r.t. costs (x,
del
/ | - [2=
This can be interpreted as learning the EOT barycceete in the ambient space but constrained to

---

the StyleGANN-arrmmeeriiee parameterized manifold GZ The barycenter Q' is some distribution of the latent variables 2, which can be pushed to the manifold (( C y via G(z).
The results are in Fig. 5 There is (a) no noise compared to the data-space EOT barycenter because of the manifold constraint, and (b) our solvers correctly ignores polluted samples from M.
5.3 Evaluation on the Ave, celeba! Dataset
In [55]. the authors developed a theoretically grounded methodology for finding probability distributions whose un-
Solver
FIDL
plans barycenter
k=1
k=
k=
regularized e barycenter is known by construction. Based on the CelebA faces dataset [73), they constructed an Ave, celeba! dataset containing 3 degraded subsets of faces. The
SCWB 52]
56.7
53.2
58.8
WIN 5S
49.3
46.9
GL.5
Ours
1.4..)
(.3)
10.2
true e barycenter w.r.t. the weights (1 4) is the distri-
Table 2:FID scores images mapped from inputs P& to the barycenter.
bution of Celeba faces itself. This dataset is used to test how well our approach recovers the barycenter.
We follow the EOT manifold-constrained setup ($..4) and train the StyleGAN on unperturbed celeha faces. This might sound little bit unfair, but our goal to demonstrate the learned transport plan to the constrained barycenter rather than unconditional barycenter samples (recall the setup in 2.3 Hence, we learn the constrained EOT barycenter with e = 10 4, In Fig. we present the results, depicting samples from the learned plans from each IPx to the barycenter. Overall, the map is qualitatively good, although sometimes failures in preserving the image content may occur. This is presumably due to MCMC inference getting stuck in local minima of the energy landscape. For comparison, we also show the results of the solvers by [32 SCWB], [55] WIN]. Additionally, report the FID score [45] for images mapped to the barycenter in Table 2 (std. deviations for our method correspond to running the inference with different random seeds). Owing to the manifold-constrainee setup, the FID score of our solver is significantly smaller.
Potential Impact, Limitations and Broader Impact
Potential impact. In our work, we propose a novel approach for solving EOT barycenter problems which is applicable to general OT costs. From the practical viewpoint, we demonstrate the ability to restrict the sought-for barycenter the image manifold by utilizing a pretrained generative model. Our findings may be applicable to list of important real-world applications. see Appendix B B.2 We believe that our large-scale barycenter solver will leverage industrial & socially-importaan problems.
Methodological limitations. The methodological limitations of our approach are mostly the same as those of EBMs. It is worth mentioning the usage of MCMC during the training/inference. The basic ULA algorithm which we use in $4.2 may poorly converge to the desired distribution n/_- In addition, MCMC sampling is time-consuming We leave the search for more efficient sampling procedures for our solver, e.g., [71 99 43 81 47 1088 66 26). for future research. We also note that our theoretical analysis in 44. does not take into the account the optimization errors appearing due to the gradient descent and MCMC. The analysis of these quantities is a completely different domain in machine learning and out of the scope of our work. As the most generative modelling research, we do not attempt to analyse these errors.
Problem setup limitations. Our paper aims at solving Entropic OT barycenter problem. In the image data space, due to utilization of the Entropy, the learned barycenter distribution may contain noisy images. However, the utilization of our proposed StyleGAN-inspiree manifold technique entirely alleviates the problem with the noise. This is demonstrated by our latent-space experiments with MNIST 0/1 (manifold space) and Ave Celeba! dataset.
Broader impact. This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
7 Acknowledgements
Skoltech was supported by the Analytical center under the RF Government (subsidy agreement (00000D733311P50000 Grant No. 70-2021-0015 02. 1.2021). AK acknowledges financial support from the NSERC Discovery Grant No. RGPIN-2023--4482 We would like to express special thanks to Vladimir Vanovskiy from Skoltech for the insightful discussions and details geological modelling (Appendix B.2
10

---

References
[1] Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-95) 924, 2011.
[2] Jason M Altschuler and Enric Boix-Adsera. Wasserstein barycenters are np-hard to compute. SIAM Journal on Mathematics of Data Science, 4(1):179-2000 2022.
[3] Pedro C Álvarez-Esteban, E Del Barrio, JA Cuesta-Albertos, and C Matrán. A fixed-point ap- proach to barycenters in wasserstein space. Journal of Mathematical Analysis and Applications, 441(2):744-76) 2016.
[4] Luigi Ambrosio and Nicola Gigli. A User's Guide to Optimal Transport, pages 1-155. Springer Berlin Heidelberg, Berlin, Heidelberg. 2013.
[5] Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representationss 2022.
[6] Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pages 146-155. PMLR, 2017.
[7] Christophe Andrieu, Nando De Freitas, Amaud Doucet, and Michael I Jordan. An introduction to meme for machine learning. Machine learning, 50:5 43, 2003.
[8] Arip Asadulaev, Alexander Korotin, Vage Egiazarian, and Evgeny Burnaev. Neural optimal transport with general cost functionals. The Twelfth International Conference on Learning Representations, 2024.
[9] Julio Backhoff- Backhoff-Veraguas,, Mathias Beiglböck, and Gudmund Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):203, 2019.
[10] Peter L. Bartlett and Shahar Mendelson. Rademacher and Gaussian complexities: risk bounds and structural results. J. Mach. Learn. Res., 3:463-482, 2002.
[1I] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Iterative bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111AA1 138, 2015.
[12] Elsa Cazelles, Felipe Tobar, and Joaquin Fontbona. A novel notion of barycenter for proba- bility distributions based on optimal weak mass transport. Advances in Neural Information Processing Systems, 34:1357: 14411355-135666. 2021.
[13] Tianrong Chen, Guan-Homn Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backwarrd SDEs theory. In International Conference Learning Representations, 2022.
[14] Jinjin Chi, Zhiyao Yang, Ximing Li, Jihong Ouyang, and Renchu Guan. Variational wasser- stein barycenters with c-cyclical monotonicity regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 7157-7165, 2023.
[15] Lénaic Chizat. Doubly regularized entropic wasserstein barycenters. arXiv preprint
SCIi:22303.18444 2023.
[16] Christian Clason, Dirk A Lorenz, Hinrich Mahler, and Benedikt Wirth. Entropic regularization
of continuous optimal transport problems. Journal of Mathematical Analysis and Applications, 494(1):124432, 2021.
[17] Samuel Cohen, Michael Arbel, and Mare Peter Deisenroth. Estimating barycenters of measures in high dimensions. arXiv preprint arXiv:2007.07105. 2020.
[18] Pierre Colombo, Guillaume Staerman, Chloé Clavel. and Pablo Piantanida. Automatic text evaluation through the lens of Wasserstein barycenters. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10450-10466, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
[19] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C.J. Burges, L. Bottou, M. Welling, 7.. Ghahramani, and K. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
[20] Marco Cuturi and Amaud Doucet. Fast computation of wasserstein barycenters. In Interna- tional conference machine learning, pages 685-693, PMLR, 2014.
11

---

[21] Marco Cuturi and Gabriel Peyré. A smoothed dual approach for variational wasserstein problems. SIAM Journal on Imaging Sciences, (1):320-343, 2016.
[22] Marco Cuturi and Gabriel Peyré. Semidual regularized optimal transport. SIAM Review, 60(4):941-965, 2018.
[23] 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-12965. 2021.
[24] Valentin De Bortoli, James Thorton,, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:1 144117655-17700 2021.
[25] Eustasio del Barrio and Jean-Michel Loubes. The statistical effect of entropic regularization in optimal transportation, 2020.
[26] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl- Arnad Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and meme. In International Conference on Machine Learning, pages
PMLR, 2023.
[27] Yilun Du, Shuang Li, B. Joshua Tenenbaum, and Igor Mordatch. Improved contrastive diver- gence training of energy based models. In Proceedings of the 38th International Conference Machine Learning (ICML-21), 2021.
[28] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models.
Advances in Neural Information Processing Systems, 32, 2019.
[29] J. Dugundji. An extension of Tietze's 's theorem. Pacific J. Math., 1:353-367, 1951.
[30] Pavel Dvurechenskii, Darina Dvinskikh, Alexander Gasnikov, Cesar Uribe, and Angelia Nedich. Decentralize and randomize: Faster algorithm for wasserstein barycenters. Advances in Neural Information Processing Systems, 31, 2018.
[31] Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions Machine Learning Research, 2023. Featured Certification.
[32] Jiaojiao Fan, Amirhossein Taghvaci, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 1571-1581. PMLR, 18 24 Jul 2021.
[33] 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 Sutherland, Romain Tavenard, Alexan- der Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
[34] Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energy- based models by diffusion recovery likelihood. In International Conference on Learning Representations, 2021.
[35] Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Extremal domain translation with neural optimal transport. In Advances in Neural Information Process- ing Systems, 2023.
[36] Aude Genevay. Fedlaiy-regularizee optimal transport for machine learning. PhD thesis, Paris Sciences et Lettres (ComUE), 2019.
[37] Aude Genevay, Marco Cuturi, Gabriel Peyré. and Francis Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems, 29, 2016.
[38] Lee-Ad Gottlieb, Arych Kontorovich, and Robert Krauthgamer. Adaptive metric dimensional- ity reduction. volume 620, pages 105-118. 2016.
[39] 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-3445 2017.
12

---

[40] Hao Guan and Mingxia Liu. Domain adaptation for medical image analysis: survey. IEEE Transactions on Biomedical Engineering, 69(3):1173-1185, 2021.
[41] Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems, 2023. 2023.
[42] Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrey Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the bridge of schr) odinger: continuous entropic optimal transport benchmark. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023.
[43] Raza Habib and David Barber. Auxiliary variational memc. In International Conference on Learning Representations, 2018.
[44] Pierre Henry-Laborders (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. arXiv preprint arXiv:1904 1904.0446,, 2019.
[45] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-seale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[46] Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 147()1177 2002.
[47] Matthew Hoffman, Pavel Sountsov, Joshua V Dillon, Ian Langmore, Dustin Tran, and Srinivas Vasudevan. Neutra-lizing bad geometry in hamiltonian monte carlo using neural transport. arXiv preprint arXiv: 1903. 2019.
[48] Syunsuke Ikeda, Gary Parker, and Kenji Sawai. Bend theory of river meanders, part 1. linear development. Journal of Fluid Mechanics, 12:363-377, 1981.
[49] Hicham Janati, Marco Cuturi, and Alexandre Gramfort. Debiased sinkhorn barycenters. In International Conference on Machine Learning, pages 4692-4701 PMLR, 2020.
[50] Tero Karras, Samuli Laine, and Timo Aila. style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401-4410, 2019.
[51] Dominik Klein, Théo Uscidda. Fabian Theis, and Marco Cuturi. Generative entropic neural optimal transport to map within and across spaces. arXiv preprint arXiv:2310.09254, 2023.
[52] Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin. Estimating barycenters of distributions with neural optimal transport. In Forty-first International Conference on Machine Learning, 2024.
[53] Ekaterina Kondrateva, Marina Pominova, Elena Popova, Maxim Sharaev, Alexander Bernstein, and Evgeny Burnaev. Domain shift in computer vision models for mri data analysis: an overview. In Thirteenth International Conference on Machine Vision, volume I 16055, pages 126-133. SPIE, 2021.
[54] Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. W generative networks. In International Conference on Learning Representations, 2021.
[55] Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[56] Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light schrödinger bridge. In The Twelfih International Conference on Learning Representations.
[57] Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein gans are optimal transport? Advances in Neural Information Processing Systems, 35:13333-3394, 2022.
[58] 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:14593-14605, 2021.
13

---

[59] Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Confer- ence on Learning Representations, 2021.
[60] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. In The Eleventh International Conference on Learning Representations, 2023.
[61] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In The Eleventh International Conference on Learning Representations, 2023.
[62]
Anastasis Kratsios and Léonie Papon. Universal approximation theorems for differentiable
geometric deep learning. J. Mach. Learn. Res., 23:Paper No. [196], 73, 2022.
[63] Roman Krawtschenko, César A Uribe, Alexander Gasnikov, and Pavel Dvurechensky. Dis- tributed optimization with quantization for computing wasserstein barycenters. arXiv preprint 1xXi:20010.14325 2020.
[64] Alexey Kroshnin, Nazarii Tupitsa, Darina Dvinskikh, Pavel Dvurechensky, Alexander Gas- nikov, and Cesar Uribe. On the complexity of approximating wasserstein barycenters. In International conference machine learning, pages 3530-3540. PMLR, 2019.
[65] Rafsanjany Kushol, Alan H Wilman, Sanjay Kalra, and Yee-Hong Yang. Dsmri: Domain shift analyzer for multi-center mri datasets. Diagnostics, 13(18)22947, 2023.
[66] John Lawson, George Tucker, Bo Dai, and Rajesh Ranganath. Energy-inspired models: Learning with sampler-induced distributions, Advances in Neural Information Processing Systems, 32, 2019.
[67] Khang Le, Dung Q Le, Huy Nguyen, Dat Do, Tung Pham, and Nhat Ho. Entropic gromov- wasserstein between gaussian distributions. In International Conference on Machine Learning, pages 12164-12203 PMLR, 2022.
[68] Khang Le, Huy Nguyen, Quang Nguyen, Tung Pham, Hung Bui, and Nhat Ho. robust optimal transport: Computational complexity and barycenter computation. Advances in Neural Information Processing Systems, 34:21947-21959 2021.
[69] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang, A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
[70] Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems, (4(44:1533-1577 2013.
[71] Daniel Levy, Matt D. Hoffiman, and Jascha Sohl-Dickstein. Generalizing hamiltonian monte carlo with neural networks. In International Conference on Learning Representations, 2018.
[72] Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin M Solomon. Contimuous regular- ized wasserstein barycenters. Advances in Neural Information Processing Systems, 33:17755- 17765, 2020.
[73] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
[74] Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference Machine Learning, pages 6672-6681. PMLR, 2020.
[75] Anton Mallasto, Augusto Gerolin, and Hà Quang Minh. reguaie-regullrrized 2- wasserstein distance between gaussian measures. Information Geometry, 5(1):289-333 223,,999,,, 2022.
[76] Simone Di Marino and Augusto Gerolin. An optimal transport approach for the schrödinger bridge problem and convergence of sinkhorn algorithm. Journal of Scientific Computing, 85(2):27, 2020.
[77] Alberto Maria Metelli, Amarildo Likmeta, and Marcello Restelli. Propagating uncertainty reinforcement learning via wasserstein barycenters. Advances in Neural Information Processing Systems, 32, 2019.
[78] Petr Mokrov, Alexander Korotin, Alexander Kolesov, Nikita Gushchin, and Exgeny Burnaex. Energy-guided entropic neural optimal transport. In The Twelfth International Conference on Learning Representations, 2024.
14

---

[79] Eduardo Fernandes Montesuma and Fred Maurice Ngole Mboula. Wasserstein barycenter for multi-soure source domain adaptation. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pages 16785-16799 2021.
[80] Youssef Mrouch. Wasserstein style transfer. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statisties, volume 108 of Proceedings of Machine Learning Research, pages 842- 852. PMLR, 26-28 Aug 2020.
[81] Kirill Neklyudov, Max Welling, Evgenii Egorov, and Dmitry Vetrov. Involutive meme: unifying framework. In International Conference on Machine Learning, pages 7273-7282.
PMLR, 2020.
[82] Maxence Noble, Valentin De Bortoli, Arnaud Doucet, and Alain Durmus. Tree-based diffusion schrödinger bridge with applications to wasserstein barycenters. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
[83] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11(-6))335-660 2019.
[84] Aram-Alexandre Pooladian, Carles Domingo-Enrich, Ricky T. Q. Chen, and Brandon Amos. Neural optimal transport with lagrangian costs. In ICML Workshop on New Frontiers in Learning. Control, and Dynamical Systems, 2023.
[85] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 1(2):3, 2022.
[86] Gareth O Roberts and Richard L Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli, pages 341-363, 1996.
[87] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[88] Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference Learning Representations, 2021. 2021.
[89] Filippo Santambrogio. Optimal transport for applied mathermaticians. Birkäuser, NY, 55(58- 63):94, 2015.
[90] Axel Sauer, Tero Karras, Samuli Laine, Andreas Geiger, and Timo Aila. Stylegan-t: Unlocking the power of gans for fast large-scale text-to-image synthesis. International conference on muchine learning, pages 30105-301 18. PMLR, 2023.
[91] Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. International Conference on Learning Representations, 2018.
[92] Dino Sejdinovic, Heiko Strathmann, Maria Lomeli Garcia, Christophe Andrieu, and Arthur Gretton. Kerel adaptive metropolis-hasting. In International conference on machine learning. pages 1665-1673. PMLR, 2014.
[93] Shai Shalev-Shwartz and Shai Ben- David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[94] Yuyang Shi. Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusson schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[95] Dror Simon and Aviad Aberdam. Barycenters of natural images constrained wasserstein barycenters for image morphing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7910-7919, 2020.
[96] Maurice Sion. On general minimax theorems. Pacific Journal of Mathematics, 8(1::171-176. 1958.
[97] Justin Solomon. Optimal transport on discrete domains. AMS Short Course on Discrete Differential Geometry, 2018
15

---

[98] Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 1414):1-11, 2015.
[99] Jiaming Song, Shengjia Zhao, and Stefano Ermon. A-nice-me: Adversarial training for meme. Advances in Neural Information Processing Systems, 30, 2017.
[100] Yang Song and Diederik Kingma. How to train your energy-based models. arXiv preprint (rXiv:2101.03260 2021.
[101] Sanvesh Srivastava, Volkan Cevher, Quoe Dinh, and David Dunson. Wasp: Scalable bayes via barycenters of subset posteriors. Artificial Intelligence and Statistics, pages 912-920. PMLR, 2015.
[102] Sanvesh Srivastava, Cheng Li, and David B Dunson. Scalable bayes via baryeenter in wasser- stein space. The Journal of Machine Learning Research, 19(1)311-3366, 2018.
[103] Karin Stacke, Gabriel Eilertsen, Jonas Unger, and Claes Lundström. Measuring domain shift for deep learning in histopathology. IEEE journal of biomedical and health informatics, 25(2):325-336, 2020.
[104) Heiko Strathmann, Dino Sejdinovic, Samuel I Livingstone, Zoltan Szabo, and Arthur Gretton. Gradient-free free hamiltonian monte carlo with efficient kernel exponential families. Advances in Neural Information Processing Systems, 28, 2015.
[105] Surya T Tokdar and Robert E Kass. Importance sampling: a review. Wiley Interdisciplinary
Reviews: Computational Statistics, 2(1):54 60, 2010.
[106] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024, Expert Certification.
[107] Alexander Y Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guil- laume Huguet, Guy Wolf, and Yoshua Bengio. Simulation schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics, pages 1279-1287. PMLR, 2024.
[108] Konstantin S Turitsyn, Michael Chertkov, and Marija Vucelja. Irreversible monte carlo algorithms for efficient sampling. Physica D: Nonlinear Phenomena, 240(4-5):410414,, 2011.
[109] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving
schrödinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021.
[110] Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
[111] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Yingnian Wu. A theory of generative convnet. In International Conference on Machine Learning, pages 2635-2644, PMLR, 2016.
[112] Wenjun Yan, Yuanyuan Wang, Shengjia Gu, Lu Huang, Fuhua Yan, Liming Xia, and Qian Tao. The domain shift problem of medical image segmentation and vendor-adaptation by unet-gan. In Medical Image Computing and Computer Assisted Intervention 2019: 22nd International Conference, Shenzhen, China, October 13-17, 2019, Proceedings, Part
22, pages 623-631. Springer, 2019.
[113] Yang Zhao, Jianwen Xie, and Ping Li. Learning energy-based generative models via coarse- to-fine expanding and sampling. In International Conference Learning Representations, 2021.
16

---

A Proofs
A.1 Auxiliary Statements
We start by showing some basic properties of the C-transform which will be used in the main proofs. Proposition A.1 (Properties of the C-transform). Let J1, J2: y IR be measurable functions which are bounded from below. It holds that
(i) Monotonicity: f1 S implies 18 2 f2:
(ii) Constant additivity: (/1 11 aa) = 12 for all E R;
(iii) Concavity: (Af1 + (1- X)52) (1- 2
ApC + (1 A)fE for all X E [0,
(iv) Continuity: f1 bounded implies supzex 1 (x)) & supyey (f((v) - f2(1)-
Proof of Proposition A.7 We recall the definition of the C-transfomm
{Clor.) 1000J15 [Lwddwmm))
{(x)= = inf MEP())
where C(., p)) def Jy c(x, y)))()) ell(x). Monotonicity i and constant additivity (ii) immediate from the definition.
To see (iii)
observe that the dependence of fy. fi (v)du(y) on fi is linear. Thus, 14 the pointwise infimum family linear functionals and thus concave.
Finally, to show (iv) we have by monotonicity of the integral that
J, f((w)d(()) /,. f2(u)du(v)
I sup (f((t) - f2(1)
(13)
DE
for any E P(Y). For fixed x E X' we have
((u,))) / [(w)daats) MEP())
((E,M) -11 Salw))m(w)
(1(x) (z) 12)))
inf nEPO)
inf
C(x, C(z,p) / fi f(())d(()) / folw)ddle)
sup inf mEP()) REP())
By setting 1 increase the value and obtain
sup / (fa(v) [i()) du(v) s sup (v( f2(v)), #EP())
(x) -
(x)) ((()
(14)
yEy
where the last inequality follows from fa in (I4), which yields the claim.
133 For symmetry reasons, swap the roles of fi and
A.2 Proof of Theorem
Proof. By substituting in (5) the primal EOT problems ( with their dual counterparts dual formulation, which the starting point our analysis:
we obtain
C (Ix)dPP(a)) + I, Sa(u)(Q(U))
L
min QEPO)) f SRC(
sup
(15)
(2
Here, we replaced inf with min because of the existence of the barycenter 2.2 the entire expression under the min sup functional E: P(Y) C())* E. For brevity, introduce, for (fi,...,ft fK) E C(Y) the notation
Moreover, refer
K
del
Aafk
and
M
del inf f(y) yEY
inf QEP(Y)
/ F())aw))
(16)
17

---

where the equality follows from two elementary observations that (a) M S f(y)d(()) for any Q E P()) and (b) f(y) (()))d where by denotes a Dirac mass at E y.
On the hand, P()) is compact w.r.t the weak topology because y is compact, and for fixed potentials (Si....; fx) P())) we have that Z(:, (fz) continuous and linear. particular, K km is convex and l.s.c. On the other hand, for fixed Q, the functional Z(Q,-) is concave
C(-, (fxe (fie)
by (1ii) in Proposition [A.I These observations allow us apply Sion's minimax theorem 96
Theorem 3.4] to swap min and inf in (I5) and obtain using (16)
K
1/ee Xx
(FW f (awaaa
L
min f (xC0)) QEP())
sup
K
/ Fw)aa(()
sup fx.. EC())
A
f (xx))dPPP + min
QEP())
K
sup fi....xECC)))
A / N4.
fx (x + min f(v) vEy
(17)
_(ff......kk)
Next, we show that the sup in can be restricted to tuplets satisfying the congruence condition E c(y)* there exists a
Ek= Arfk = 0. It remains to show that for every tuplet
congruent tuplet (fi,...,fr fk) E c(y) such that E(f To this end, fix (fi...., fk) and define the congruent tuplet
J
(f.....,fx fxx) f1.... fk....ff.1.fa
Ax
(18)
We find M def inf yny Ek=1 Sa by the congruence and derive
Z(h, fk) i(J = Ax {
(TK) - [C (xx) dPx(ak) M
-s
Nx
(
M (xx) - - (xx)) dP(xk) - M
2 dx
Nx
M
Ak
(Il((x))
-M = =0,
Nx
Aσ
where the first inequality follows from fx = fK
{ S fK Ax
M combined with monotonicity
of the C-transform, see in Proposition A.1 The second to last equality follow from constant additivity, see (ii in Proposition A.I
In summary, we obtain
-= sup E(fi,..., (1C.0))
(19)
ΣK fu=0
Finally, observe that for congruent (fi...., fK) we have E(fi...., C(f fk). Hence, we can replace by (19), which yields
A.3 Proof of Theorem 4.2
Proof. Write Q" for the barycenter and nद for the optimizer of EO Q*). Consider congruent potentials fi,..,,K fx E C())) and define the probability distribution
dm In (IK, y) del dy (v) dT)x(xx),
where
dutr (v) def
exp fa(y) ck (x(x
(20)
dy
Zo, (Jh, xxx)
18

---

Zoe (fk,:H)) def log (fs ( dy)
(21)
Then we have by 78 Thm. 2]:
(Se
(oa())))P(a)) + SS, s(o)d((()))
EOT o PR, Q*)
-cKL (Tāl(T/))
(22)
Multiplying 22 by A5 and summing over yields
K
KL (T) =
Σ
EOT of (Px,
/ x 1 (xw))dPN(())) -
Σ
(v)
dQ'(w)
k=1
k=
- fx),
(23)
where the last equality follows by congruence, i.e., Σ A= =0
The remaining inequality in is consequence of the data processing inequality for f-divergences which we invoke here to get
f
KL ( 2 KL (Q'1)))),
where Q* and Q+k are the second marginals x and fk respectively.
A.4 Proof of Theorem 4.3
Proof. The desired equation could be derived exactly the same way as in (78) Theorem 3].
A.5 Proof of Proposition 4.4 and Theorem 4.5
The derivations of the quantitative bound for Proposition 4.4
and and Theorem 4.5 relies on the following
standard definitions from learning theory, which we now recall for convenience (see, e.g. 993 $26)).
Consider some class S of functions X' + R and a distribution μ on X. Let X be a sample of N points in X'
The representativeness of the sample X w.r.t. the class S and the distribution j is defined by
1/ s(x)dμ(x)
Rep (S, ())
sup
s(x")]
(24)
.ES
N n=1
The Rademacher complexity of the class w.r.t. the distribution μ and sample size is given by
(S, def
E
sup SES
Σ
s ((x" Jan
Jan
(25)
where {a:") x==1 μ are mutually independent, {a") j mutually independent Rademacher random variables, i.e., Prob(o" = 1) Prob(o" -1) 0.5, and the expectation is taken with respect both 1N I.N The well-celebrated relation between (25) and (4). shown in
993 Lemma 26.2], is
ERepx (8, s (x,
(26)
where the expectation is taken w.r.t. random i.i.d. sample X ~ of size N.
Proposition 4.4
Observe that by (23) we may write
AL (T)l(T(e) - L(f)
max L(f) L
max L(1 L()) FEF
(27)
FEF
Approximation emmr
Estimation emvor
Let i be a maximizer of L(f) within F. Analysing the estimation error in 27
yields
max C(T) - C(P) = L(I) - -((
19

---

[c(t) - [(F] + [2(1) - -00]+ - -
(28)
suu c(() E(n).
(29)
FEF
where central term in line 28 is bounded above by due the maximality of f, that is, 2(0 can bound the estimation error using the Rademacher
max4=FF (()) [(f). Due to 29
complexity
K Σ U f X
N
K
sup [(f) Z(0) S fe
A sup
[ (xa)) 1
Ok (1+)
ReRe (F Px).
N a
k=1
Proof of Theorem 4.3 Case (a) Lipschitz costs: Assume that, for k K, (2, y) Lipschitz with constant La 2 for every y Recall that fC is defined as the pointwise supremum of Ly- Li-Lipschitz functions and, therefore, Lipschitz continuous with the same constant. Since the value of the representativennss of sample w.r.t. a function class is invariant under translating individual elements of said class, we have that Rep (FF .P) coincides with Repx (9k, Fk) where
Gk def ( - JCh (mx) (#x) C()))
for some fixed āy. Nr. All the functions in this class are La-Lipschitz and, therefore, bounded by LA di) We may therefore apply [38 Theorem 4.3] to the class Gs and obtain
Ex Repx (FC,, Pk)) E Ex Repx (9k, PP)) S 2R (9k, PPA) S ONN, D
Case (b) Feature-based quadratic costs 1 (lua() (+) - rl)) Alternatively, consider the case where ok (Tx, y) = 11 (3x) - vly)l1 and Fx S C()) is bounded (w.r.t. the supremum norm).
To this end, recall that for measurable and bounded function f y + R. the weak entropic c-transform satisfies
1o (Ix) -c coogg (e, (f,xk))
where Zon (f, 2x) = fy exp (f()) (o))))))))))) dy. Recall that R D" xR 3 (a, b) exp Iaa-bll bll"
2e
positive definite kernel which is widely known as the Gaussian kernel. This means that there exists a Hilbert space H. and a feature map á RD" + H such that exp Due to the particular form of Ck2 we may write
||a-M] 2c
(á(a), d(b))H
exp ca(Th, cxp ( [[u(()) -(w))l?) (6((((()))
(30)
Notice that {o(ux(Iu)) X& S {U H. ||l|| 1} since (1((a) = (o(a), = for every DD"
Using the identity
30 we can express Ze (f, xxx) by
Zee (f,
J (o(un o(v())))
fiv) dy (d(un(z))), / ((v())) flui
)
where the last equality is justified as y is compact; furthermore, we note the integrals are well-defined by the measurability of Φ. us and v. Moreover, using the boundedness of each Fx and compactness of )2 we get
def
11 d(((w))e
(iv) dy|ln < ao.
R
max sup R I...,K K SEFi
(31)
Define GR del (f. f Fi) and observe that Ge S 0 def ((d(ux))) w)N h ||l] S R). This implies that RN, (G Pk) RN, Px) by the properties of the Rademacher complexity.
20

---

In turn, the latter Rademacher complexity can be bounded by [T0 Lemma 22]. Indeed, write def
Pk (ux)+Px such that
RN (Pk, RN, (((o(-), whn : ||l|x - R),
which can be directly seen from the definition of the Rademacher complexity, Thus, we can apply 10 Lemma 22] to the right-hand side, and summarizing obtain
R
RN,, (9k, S
VNk
Since the functions in 9K are bounded uniformly away from zero by some constant &: > 0 (depending on the bound of F& and and since the logarithm restricted to [K, oo) is 1 //s-Lipschitz, have that
Rv, (L) é R R
RN, (F
P&.
V
We can now bound the expected representativeness of F with the Rademacher complexity by (26 yielding the claim.
A.6 Proof of Theorem
Let a ((R) be non-affine activation function which is differentiable at some point Zo and for which (=()) + 0. 6 > A1...., Ag > O, and K E N be given. By Theorem 4.1 there exist K congruent continuous functions fi...., such that
Proof of Theorem 4.6
L - fk) 2
(32)
Applying 229
Theorem 4.1] we deduce that for each k l,.... K, there exist continuous extension fk RDD R for Sa all of R i.e. M(() [u( for each V. In particular, 32 can be rewritten as
L- - -(S, .Kk
(33)
Set dR def 8/(42x) for each I,..., Since y is a compact subset of D, and since we have assumed that E C(R) is non-affine activation function which is differentiable some point ā E RR and for which o' (3) j then the special case of [62, Theorem 9] given in [62] Proposition 53).
K.
implies that for any there exist feedforward neural networks gs RD + R (k e K) with activation function such that
1g falloo
def sup (9)(v) fk(v)l = sup (9) fc(U)1 < Sz. yEy
yEy
each Ok width at-most D + 4. Pick d, for all k E K and suitable neural networks ...... ,9.
Next, we define the congruent sums of neural networks f& 9k
Ek=l Angk. We derive
K
K
K
K
K
A
=
Σ
AE
A
- f/l
5
(34)
4
4
km
km
Using 34 we obtain for fixed & E
K
[ fallo 11. 9k +
< <11f% 1l
9kll
+ Σ Σx=
K
Au
1lo
<
5
(35)
k=
2.
By (iv) in Proposition A.1 together with (35p we find
[ (ff)
s llfa falloo <
(36)
21

---

Now we use 36 to derive
k
Ak /
Sk (o))) / (
fk) -US.....,.
S
Σ
k=
A/z
(1x)) - (
Nx
km
K
8 8
dvels (f2) <
A
(37)
2
=1
Next we combine 33 with 377 to get
C
L(Si...., JK) L CU
[k)
E(1 JK) ,.
... [k]]] <8
(38)
<6/2
<6/2
By using (38) together with Theorem
4.2
we obtain
KL (TE](T()) = - <
which completes the proof.
A.7 Existence and uniqueness of the barycenter distribution which solves
We introduce an auxiliary functional which is the argument of minimization problem
del
AREOT Q),
(39)
i.e. the optimal value of 5 could be defined as L
inf B(Q).
QEP(Y)
Note that the functional ((Q) is strictly convex and lower semicontinuous (w.rt. the weak topol-
ogy) each component EOT ct (P&k Q) is strictly convex and l.s.c. (lower semi-continuous) itself. The latter follows from 9 Th. 2.9] by noting that on P(Y) the map μr> u--Ey-aChC y) H (n)
I.s.c, bounded from below and strictly convex thanks the entropy term. Since P(Y) is weakly compact (as ) is compact due to Prokhorov's 's theorem, see, e.g., 89 Box 1.4)), it holds that B(Q)
admits at least one minimizer due to the Weierstrass theorem 89 Box 1.1], i.e., a barycenter Q" exists. In the paper, we work under the reasonable assumption that there exists least one Q which ((Q)).. this case, the barycenter Q" unique due the strict convexity of B.
B Extended discussions
B.1 Extended discussion of related works
Discrete OT-bbaed solvers provide solutions to OT-related problems between discrete distributions.
A comprehensive overview can be found in 83
The discrete OT methods for EOT barycenter
estimation are (20] 98 I 21 30
63
49
68 An alternative formulation of the barycenter
problem based weak mass transport and corresponding discrete solver could be found in [12). In spite sound thcoretical foundations and established convergence guarantees (44 these approaches can not be directly adapted to our learning setup, see $2.3
Continuous OT solvers. Beside the continuous EOT solvers discussed in 3 there exist a variety of neural OT solver for the non-entropic (unregularized, c 0) case. For example, solvers such as
dual form, similar to our (2), with neural networks. We mention these methods because they serve the basis for certain continuous unregularized barycenter solvers. For example, ideas of [55] are employed in the barycenter method
44
57 32
88 31, are based on optimizing
presented in [59]; the solver from [74) is applied in 332 maxx-in solver introduced in [58] is used
22

---

in [55]. It is also worth noting that there exist several neural solvers that cater to more general OT problem formulations 66 60 31
g 84 These can even be adapted to the EOT case 42 but require substantial technical effort and the usage of restrictive neural architectures.
Continuous EOT solvers aim recover the optimal EOT plan x* in EOT problems like ! between unknown distributions IP and Q which are only accessible through limited number of samples. One group of methods [37 91 [3] is based on the dual formulation of OT problem
regularized with KL divergences [37 Eq. (Pt)] which is equivalent to (T). Another group of methods
109 24
13
106
94] takes advantage of the dynamic reformulation of EOT via Schrödinger
bridges
70 76
In turn, [51] solve EOT with conditional flow matching [106].
In [78], the authors propose an approach tackle ①I by means of Energy- -aargd models 69 100
EBMJ. They develop an optimization procedure resembling standard EBM training which retrieves the optimal dual potential f appearing in (2). As a byproduct, they recover the optimal conditional
plans _ = x² (-Iz). Our approach for solving the EOT barycenter (5) is primarily inspired by this work. In fact, we manage to overcome the theoretical and practical difficulties that arise when moving from the EOT problem guided with EBMs to the EOT barycenter problem (multiple marginals, optimization with respect to an wnfixed marginal distribution Q), see $4 for details of our method.
Other related works. Another relevant work is [95], where the authors study the barycenter problem and restrict the search space to a manifold produced by GAN. This idea is also utilized in $5.2 and |5.3] but their overall setting drastically differs from our setup and actually is not applicable. We search for a barycenter of K high-dimensional image distributions represented by their random samples (datasets). In contrast, they consider K images, represent each image 2D distribution via its intensity histogram and search for a single image on the GAN manifold whose density is the barycenter of the input images. To compute the barycenter, they use discrete OT solver. In summary, neither our barycenter solver is intended to be used in their setup, nor their method is targeted to solve the problems considered in our paper.
B.2
Extended discussion of potential applications
is not a secret that despite considcrable efforts in developing continuous barycenter solvers 59 55 32 [4, these solvers have not found yet real working practical application. The reasons for this are two-fold:
772
1. Existing continuous barycenter solvers (Table are yet not scalable enough and/or work exclu- sively with the quadratic cost (E²), which might be not sufficient for the practical needs.
2. Potential applications of barycenter solvers are too technical and, unfortunately, require substantial efforts (challenging and costly data collection, non-trivial design of task-specific cost functions, unclear performance metrics, etc.) to be implemented in practice.
Despite these challenges, there exist rather inspiring practical problem formulations where the continuous barycenter solvers may potentially shine and we name a few below. These potential applications motivate the research in the area. More generally, we hope that our developed solver could be a step towards applying continuous barycenters to practical tasks that benefit humanity.
1. Solving domain shift problems in medicine (Fig. 6a) In medicine, it is common that the data is collected from multiple sources (laboratories, clinics) and using different equipment from various vendors, each with varying technical characteristics 40 65
53 103 112]. Moreover, the
data coming from each source may be of limited size. These issues complicate building robust and reliable machine learning models by using such datasets, e.g., learning MRI segmentation models to assist doctors.
A potential way to overcome the above-mentioned limitations is to find common representation of the data coming from multiple sources. This representation would require translation maps that can transform the new (previously unseen) data from each of the sources to this shared representation. This formulation closely aligns with the continuous barycenter learning setup ($2.3 studied in our paper. In this context, the barycenter could play the role of the shared representation.
To apply barycenters effectively to such domain averaging problems, two crucial ingredients are likely required: appropriate cost functions and a suitable data manifold M in which to search for the barycenters. The design of the cost itself may be a challenge requiring certain domain-specine knowledge that necessitates involving experts in the field. Mcanwhile, the manifold constraint is
23

---

Source
Souron 2
C1
?
C2
?
C2
Petrel
Flumy
C
Source
C3
(a) Domain averaging of MRI scans' sources.
(b) Mixing geological simulators.
Figure 6: A schematical presentation of potential applications of barycenter solvers.
required to avoid practically meaningless barycenters such as those considered in $5.2 Nowadays, with the rapid growth of the field of generative models, it reasonable expect that the new large models targeted for medical data may appear, analogously to DAILLEE [5]. StableDiifussion [87 or StyleGAN-T [90] for general image synthesis. These future models could potentially parameterize the medical data manifolds of interest, opening new possibilities for medical data analysis.
2. Mixing geological simulators (Fig. 6b) In geological modeling, variuos simulators exist to model different aspects of underground deposits. Sometimes one needs to build a generic tool which take into account several desired geological factors which successfully modeled by separate simulators.
Flumy is a process-based simulator that uses hydraulic theory 48 to model specific channel depositional processes returning detailed three-dimensional geomodel informed with deposit lithotype, age and grain size. However, its result is a 3D) segmentation field of facies (rock types) and it does not produce the real valued porosity field needed for hydrodynamical modeling.
Petre software is the other popular simulator in the oil and gas industry. It is able to model complex real-valuee geological maps such as the distribution of porosity. The produced porosity fields may not be realistic enough due to paying limited attention to the geological formation physics.
Both Flumy and Petrel simulators contain some level of stochasticity and hard Even when conditioned on common prior information about the deposit, they may produce maps of facies and permeability which do not meaningfully correspond each other. This limitation provides potential prospects for barycenter solvers which could be used to get the best from both simulators by mixing the distributions produced by each of them.
in conjunction.
From our personal discussions with the experts in the field geology, such task formulations are of considerable interest both for scientific community well as industry. Applying our barycenter solver in this context is challenge for future research. We acknowledge that this would also require overcoming considerable technical and domain-specifie issues, including the data collection and the choice of costs Ck-
B.3 Extended discussion on doubly-regularized OT barycenters
The objective 5 is not the only way to formulate Entropic OT barycenter problem. Recent studies 15 82 consider the so-called doubly-regularized Entropic OT barycenters. Following the notations
from 15] for À 2 0 and r 2 ( we define:
EOT c.A,7 (P. =EIT(P.)) (a.y)=
min E c(x,v) + AKL & Q) +
The corresponding
EOT barycenter problem is follows:
K
L dr del
inf B (Q) QEP(Y)
def
inf GEP())
BOT C.A.r (Pk, Q)
(40)
eul
https://ww.softuaat ww..rreesbb......... www..cc//ppoottutt///tttteeaeee
24

---

It turns out that our considered EOT barycenter (5) with regularization strength is the particular According to the classification [ Table 1]. this case is known
case of (40p with as Schrüdingge barycenter. The natural question arizes: Is it possible to adapt our energy-gun energy-guided methodology for the case À + r?
To answer this question, in what follows, we have a closer look at the differences between general doubly-regularized formulations and our particular Schrödinger specification. The important property of our case is that the entropy term H (Q) completely disappears from the barycenter objective. In all the other regularized cases, when À o т, this entropy term immediately reappears (either with the plus or minus sign). The presence of (Q) term notably differs from ours and seems like to be not suitable for our solver. In what follows, we explain the reasons.
In the Schrödinger case, the barycenter problem can be solved via optimizing conditional distributions T& (y)l xx)). Namely, use potentials f& combined with costs c to approximate the energy functions (unnormalized log-densities) of these conditionals. We employ EBM-based techniques to compute the gradient of the learning objective which avoids the direct computation of the entropy terms H (TT appearing in the C₂ Cg-transforms.
If we further add the non-zero term (QQ to the barycenter objective, this will presumably require (in the dual objective) separate computation of the entropy terms (T&(H)) T& (y) of each тTK- This is highly non-trivial and seems like our solver does not easily generalize this case. In our framework, we get samples of TA (y) by MCMC (via sampling ~ Ixx and then running MCMC for (z)(z))) However, esti- mation of entropy of (y) from raw samples still remains infeasible. In particular, EBM-like tech- niques (which we employ) can not be used here to
training time
Experiment
time
Toy (Ours)
3b
20s
second marginals
(Ours)
Sphere (Ours)
20s
MNIST 0/1 manifold (Ours)
10h
Im
MNIST 0/1 data (Ours)
20th
1m
Ave Celcha manifold (Ours)
Ave Celoba data (Ours)
608
2m
Ave Celcha data (WIN)
160h
10s
Ave Celeba data (SCWB)
200h
10s
Table 3: Computational complexity for Ours (all
experiments) and baselines (Ave Celeba).
derive the gradient of the objective. This because the required unnormalized density of ι (g) is itself unknown (we only know it for conditional distributions T (v(zz))
C Experimental Details
The hyperparameters of our solver are summarized in Table Some hyperparameters, e.g., L, S, iler, are chosen primarily from time complexity reasons. Typically, the increase in these numbers positively affects the quality of the recovered solutions, sec, c.g., 442 Appendix E, Table 16]. However, to reduce the computational burden, report the figures which we found to be reasonable. Working with the manifold-constraint setup, we parameterize each ge, (z) in our solver as ha, o G(zz, where G' is pre-trained (frozen) StyleGAN and he, is neural network with the ResNet architecture. We empirically found that this strategy provides better results than a direct MLP parameterization for the function go, (2).
Experiment
D
A
Xa
Aa
A4
fo.k
Iran
ster
Toy 2D
3
10-2
1/3
1/3
1/3
MLP
10
200
1.0
300
256
MNIST a/1
1024
2
10-2
10
0.5
0.5
ResNet
10
1000
0.1
500
32
MNIST 0/1
512
2
10
0.5
0.5
ResNet
10
1000
0.1
250
32
Ave, celebal
3964"
3
10-2
0.25
0.5
0.25
ResNet
10-4
5000
0.1
500
64
(Data)
Ave, celeba!
(Manifold)
512
3
10-1
0.25
0.5
0.25
ResNet
10-1
1000
0.1
250
128
Gaussions
10-#.1
0.5
MLP
10-3
50000
1024
Sphere
MLP
x10
256
Singlle-sell
1000
0.25
0.75
ML.P
x
1000
0.05
1000
1024
Table 4: Hyperparameters that we use in the experiments with our Algorithm
|
To train the StyleGAN for MNSITOI & Ave, celeba! experiments, employ the official code from
25

---

Computational complexity. We report the (approximate) time for training and inference (batch size S 128) of our method on different experimental setups, see Table (the hardware is single V100 gpu). For Ave Celeha experiment, we additionally report the computational complexity of the competitors. As we can see, all the methods in this experiment (Ave Celeha) require comparable amount of time for training. The inference with our approach expectedly costlier than competitors due the reliance MCMC.
C.1 Barycenters of 2D/3D Distributions
Cartesian representation of twister map. In 5.1 define twister map using polar coordinate system. For clearness, we give its form on cartesian coordinates. Let æ , r = (Z(1): (2))). that ||>|||| = V ( + 2(2) Then,
Note
u(x)= = u(I F(2)) cos ([|x])) sin (1|=|l) (|==I)
2(11) x1(2)
sin
cos (1|==)
i.c. the twister map rotates input points æ by angles equal to ||=|||-
Analytical barycenter distribution for 2D Twister experiment. Below, we provide mathematical derivation that the true unregularized barycenter of the distributions P1,P2,P5 Pa in Fig. 2a coincides with Gaussian.
We begin with rather general observation. Consider X = = RDD (k E K) and let OT, EOT c. denote the unregularized OT problem = 0) for given continuous cost function Let u R + R be measurable bijection and consider IP, E P(RP) for & E K. By using the change of variables formula, we have for Q' P(T.P) that
OT OT co (u (P'), " (Q') = (I', QP),Q,)),
(41)
where # denotes the pushforward operator of distributions and (u u)] y) = c(u(z), u(y)). Note that by (4T) the barycenter of P'r for the unregularized problem with cost coincides with the result of applying the pushforward operator u to the barycenter of the very same problem but with cost (u x u).
Next, we fix u to be the twister map (55. In Fig. 2a we plot the distributions P1 def # Iy, # P1, u Ps which are obtained from Gaussian distributions Pí = ((0,, 4), 12,
(((-2, ((2,,223), P/ (22 2v3), I2) by the pushforward. Here is the 2-dimensional iden- tity matrix. For the unregularized e² barycenter problem, the barycenter of such shifted Gaussians can be derived analytically 3. The solution coincides with a zero-centered standard Gaussian
Q N (0, 12). Hence, the barycenter of Pi,..., PP w.r.t. the cost f2 (u x u) is given by Q" u V From the particular choice of u it not hard to see that Q = Q NO, ( I2) as well.
C.2 Barycenters of MNIST Classes
Additional qualitative examples of our solver's results are given in Figure 7
Details of the baseline solvers. For the solver by 32 SCWB), run their publicly available code from the official repository
https: pettss:/gittane http://gtthuuu..coobbbbeeeeeeeee
com wtt..//bbbbbbs//esaaabooooo htt.aaabbeewwaserrreeeeee-aaaae- Barycenter
The authors do provide checkpoints, and train their barycenter model from scratch. In turn, for the solver by [55] WINJ, we also employ the publicly available code
Here do not train their models but just use the checkpoint available in their repo
26

---

(a) Learned plans from F1 (zeros, 1st column) to the barycenter.
1
/
1
I
1
/
1
1
{
(b) Learned plans from P2 (ones, first column) to the barycenter.
Figure 7: Experiment with averaging MNIST 0/1 digit classes. The plot shows additional examples of samples transported with our solver to the barycenter.
C.3 Barycenters of the Ave, Celeba! Dataset
Additional qualitative examples of our solver's results are given in Figure X
Extra experiment in Data Space. Our method in the manifold constrained setup on MNIST dataset performs better than in the data space setup (see Figure 5) It generates less noised images and demonstrates better perceptual quality. For this reason, we provide only Manifold-constrained setting for experiment with Ave, Celeba! dataset in $5.3
For completeness, here we also test our method directly in the data space setup for Ave, Celeba! dataset. Hyperparameters of our solver are listed in Table 4 In Figure I0, we show images obtained in Data space setup. As expected, the FID scores in Data space are not that good as the images are more noised since solve the entropy-regularize
problem. But we stress one more time that our method permits StyleGAN manifold trick, which greatly improves the performance, see the images in Figure 4 for the manifold-constrainee setup.
Convergence at training. We provide the behaviour of -V-UVP metric (between the unregularized ground truth barycenter mapping and the entropic barycenter mapping for our learned T (w)))))) for Ave Celeba experiment $5.3) see Figure We emphasize that L2-UVP directly measures (by computing pointwise MSE) how the learned mapping differs from the true mapping.
27

---

r-bi-e y betre
y 1--AANN
r-
y tero
1-
y tetho
(a) Maps from Pj the barycenter.
(b) Maps from Pa to the barycenter. (c) Maps from Ps to the barycenter.
Figure 8: Experiment on the Ave, celeba! barycenter dataset. The plots show additional examples of samples transported with our solver to the barycenter.
100
Figure 9: Training curves of L2-UVP vs. time for OUR proposed method in Manifold-constrained setup with Style-GAN on Ave. Celeba! dataset. The duration of the training is 100 h (1 GPU V100).
Convergence at inference. Since Langevin Dynamies at the heart of our method for sampling from a conditional plan, it is important to demonstrate the dependence of the inference times/quality on the number L Langevin steps in Manifold constrained as well as the Data space setup. We demonstrate Tables 5 and 6 where we show the trade-off between number of Langevin steps L and the obtained quality. To provide the comprehensive analysis, we report FID scores as well as sampling time for different number of steps L. Expectedly, inference time linearly depends L both setups.
Our results testify that (in Manifold constrained setting) our method achieves good quality even with rather small number of Langevin steps, i.e., the computational burden of our approach could be considerably reduced with rather little trade-off in quality.
Details of the baseline solvers. For the (55 WINJ solver, we use their pre-trained checkpoints provided by the authors in the above-mentioned repository. Note that the authors of [32] SCWB] do not consider such high-dimensional setup with RGB images. Hence, to implement their approach in this setting, we follow [55] Appendix B.4].
28

---

FID
FID
k=0
k=l
k=2
L
I, sec
k=0
k=
k=2
L
sec
15.8
15.3
18.3
15
125.3
122.3
187.6
10
11.3
11.2
14.3
150
45
119.4
120.0
169.7
150
105
8.4
8.7
10.2
250
75
118.5
120,4
168.8
250
175
8.3
8.2
9.9
500
150
118.3
120.1
168.9
500
350
8.2
8.2
9.8
1000
300
118.8
121.2
170.2
1000
700
Table 5: The dependence of running time of ULA number of Langevin steps during in- ference mode and the trade-off between the steps and the quality of obtained images in Ave, Celeba! dataset in Manifold constrained setting.
Table 6: The dependence of running time of
ULA on number of Langevin steps during inference mode and the trade-off between the steps and the quality obtained images in Ave, Celeba! dataset in Data space setup.
at 1y-ioM
(a) Maps from P1 to the barycenter. (b) Maps from P2 to the barycenter. (c) Maps from P3 to the barycenter.
Figure 10: Experiment on the Ave, celeba! barycenter dataset. The plots represent transported inputs 2& IP' to the barycenter learned by our algorithm in Data space. The true unregularized t barycenter of IP, P2, P3 are situated in figure 4a 4b and 4c correspondingly.
C.4 Barycenters of Gaussian Distributions
We note that there exist many ways incorporate the entropic regularization for barycenters [5 Table 1]: these problems do not coincide and yield different solutions. For some of them, the ground-truth solutions known for specific cases, such the Gaussian case. For example, 75 Theorem 3] examines barycenters for OT regularized with KL. divergence. They consider the task
I - y112 d=s(x, CKL(TT x Q))
inf QEP(Y)
Xxy
kml
f JxXy
| y12
d d v) - Jf H(TA(()))dP(())
inf QEP()))
d
A=1
K A (P +c
inf GEP(O)
objective inside
This problem differs from our objective ( with c(z, y) | y by the non-constant Q- dependent term (Q);; this problem yields different solution. The difference with the majority of other mentioned approaches can be shown the same way. particular, [67] tackles the barycenter for inner product Gromov-Wo Gromov-Wassersten problem with entropic regularization, which is not relevant for us. In turn, the authors [49] demonstrate significant progress towards GT results for barycenter problem with considered regularization, but only for univariate (ID) case. To our knowledge, the general Gaussian ground-trui solution for our problem setup (5) is not yet known, although some of its properties seem to be established 25
Still, when As 0, our entropy- reguuaa-eee barycenter expected be close to the unregularized one (c 0). In the Gaussian case, it is known that the unregularized OT barycenter for c₁ (x.x1)) y) Ø y/ itself Gaussian and can be computed using the well-celebrated fixed point iterations of 3 Eq. (19)). This gives us opportunity to compare our results with the ground-truth unregularized barycenter in the Gaussian case. As the baseline, we include the results of [55] WIN] solver which learss the unregularized barycenter (e = 0).
29

---

We consider 3 Gaussian distributions P2, P3 in dimensions D = 2,4,, 16, 64 and compute the
approximate EOT barycenters < for c= 0.01, 1 w.r.t. weights (A1. )g = C_ 1: 2/
1) with our solver. To initialize these distributions, follow the strategy of [55 Appendix C.2]. The ground truth unregularized barycenter Q" is estimated via the above-mentioned iterative procedure. We use the code from WIN repository available via the link mentioned in Appendix C.2 To assess the WIN solver, we use the unexplained variance percentage metrics defined as L2-UVP
100- - where T* denotes the optimal transport map 7", see 54 $5.1]. Since our solver computes EOT plans but not maps, we evaluate the barycentric projections of the learned plans, i.e., <x = fyy yT (u|z)) and
calculate C2-UVP(TT,, T.). We evaluate this metric using 10* samples. To estimate the barycentric [h projection in our solver, we use 10 samples y Th (2)2 for each z To keep the table with the results simple, in each case we report the average of this metric for &o = 1, 2, 3 w.r.t. the weights Ak-
Dim / Method
Ours (c
Ours (c 0.01)
55 WIN]
1.12
0.02
0.03
4
1.6
0.05
0.08
8
1.85
0.06
0.13
1.32
0.09
0.25
64
1.83
0.84
0.75
Table 7: LV-UVP for our method with 0.01, 1 and WIN, D =0 2, 2,, 16,
We see that for small c = 0.01 and dimension up to = 16, our algorithm gives the results even better than WIN solver designed specifically for the unregularized case (e 0). As was expected, larger e = leads to the increased bias in the solutions of our algorithm and L2-UVP metric increases.
Effect of the batch size. In order to test how our method is affected by batch size, we conduct an additional empirical study with varying batch size. We follow our Gaussian experimental setup from above and pick (D, E) (2, 0.1), (D,e) (16, 0.01). As the batch sizes, we consider 2, 2,4,8 8, 16, 32, 64, 128, 256, 512, 1024. As we can see in Table 9 the quality of the recovered plans T between reference and barycenter distributions naturally grows with the batch size. In all our other experiments, we typically choose the batch size to be a reasonable number which, on the one hand, allows achieving sufficient quality and, the other hand, provides reasonable computational burden. In the image experiments, we use batch size < 128 as already provides reasonable performance. Going beyond that is challenging due to the computational restrictions.
Effect of sampling steps number at training. We conducted an ablation study testing how our method affected by chosen number of discretized Langevin dynamic steps (L from (4.2) at training, the results are presented in Figure IT this experiment, we try to learn the barycenter of Gaussian distributions; pick the dimensionality D = 64 (which is the highest- dimensional case among the considered); As we can see, performance drops when the number of steps is insufficient. Overall, in all our experiments, the number of Langevin steps is chosen to achieve reasonable qualitative/quaniiative results.
SrUVP
46
128
256
10
20
60
100
10"
Figure 11: The training curves of Cv-UVP vs. iterations for OUR proposed method for the barycenter of Gaussian distributions depending on number of Langevin steps L.
C,5 Single-cell experiment
We consider the problem of predicting the interpo- lation between single cell populations at multiple
timepoints from [107] The objective is to interpolate
the distribution of cell population at any intermediate point in time, call it t, given the cell population distributions at past and future time-points 4++11 and 1,-1. Since this is an interpolation problem, is
natural to expect that the intermediate population is (entropy-regularized) barycentric average (with 2 cost) of both the population distributions available at the nearest preceding and future times. We
leverage the data pre- propesoces and metrics from paper (56), wherein the authors provide complete
30

---

Solver/DIM
50
100
1000
OUR
3322:011
2.26 + 0.09
1.34 = 0.12
LightSB-M [2]
2.33 + 0.09
2.172 + 0.08
1.33 1 0.05
SFM-sink [3]
2.66 0.18
2.52 + 0.17
1.38 = 0.05
EGNOT [1]
2.39 + 0.06
2.32 0.15
1.46 = 0.20
Table 8: Energy distance (averaged for two setups and random seeds) on the MSCI dataset along with 95%--confidence interval (+--itervalss) The best solver according to the mean value is bolded.
())c)
Batch size
2
4
8
16
32
64
128
256
512
1024
(2,0.1)
0.20
0.14
0.07
0.05
0.05
0.04
0.02
0.03
0.02
0.02
C2-UVP
(16,0.01)
0.56
0.40
0.37
0.30
0.25
0.22
0.20
0.17
0.15
0.13
Table 9: Influence of the batch size on the performance of our approach in the case of computing the barycenters of Gaussian distributions.
notebook with the code to launch the setup similar to [107] There are experimental settings with dimensions D 50, 100, 1000, each setting contains two setups: predicting day given days 2 and [107] or Section 5.3 from [56] for additional experimental details. We report the result in Table where we find that in most cases, our general-pune purpose entropic barycenter approach nearly matches the performance of leading baselines. This underscores the scope of problems in which our barycentric optimal transport technology can viable foundation model, directly out-of-the-box.
and predicting day 4 given days and 7. The metric is MMD;
D Alternative EBM training procedure
20
f (xs] Pp
X3 f
y-
P
-2
(a) trie unregularized baryoenter for the twisted cost.
(h) Our barycenter for twisted cost (map from Pt).
(c) unregularized baryeenter for l cost.
(d) BOT harycenter e cost (map from P1).
Figure 12: 2D twister example. Trained with importance sampling: The true barycenter of 3 comets vs. the one computed by our solver with = 10 10-2 2 Two costs cj are considered: the twisted cost We employ the simulation-free importance sampling procedure for training.
12a 12b and 12 12c 12d).
In this section, describe an alternative simulation-free training procedure for learning EOT barycenter distribution via our proposed methodology. The key challenge behind our approach is to estimate the gradient of the dual objective (4.)) To overcome the difficulty, in the main part of our manuscript, we utilize MCMC sampling from conditional distributions μiz with Monte-Carlo. Here we discuss potential alternative approach based on importance sampling (IS) [05). That is, we evaluate the internal integral over y in 4.3h.
fo.t and estimate the loss
I(xx) def / [ dudr Sa (v)
(42)
with help of an auxiliary proposal (continuous) distribution accessible by samples with the known density q(y). Let Y {vf..... y be a sample from q(y). Define the weights:
[So. (w/) - c( ug)
(x( 0g)
exp
1(1g).
Then 42
permits the following stochastic estimate:
31

---

Ep= Σp Wk (
89 fo.k [y)) Wk 83)
I(rx)= ~
(43)
Experimental illustration. To demonstrate the applicability of IS-based training procedure to our barycenter setting, we conduct the experiment following our 2D Twister setup. see $5.1 We employ zero-mean 161-covariance Gaussian distribution are shown in Figure 12 As we can see, the alternative training procedure yields similar results as Figund but converges faster (=s 1 min. VS 18 min. of the original MCMC-based training).
and pick the batch size P 1024. Our results
Concluding remarks. We note that IS-based methods requires accurate selection of the proposal distribution to reduce the variance of the estimator [105]. It may be challenging in real-world scenarios. We leave the detailed study of more advanced approaches the context energy-based models and their applicability to our EOT barycentric setup to follow-up research.
32

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In the Introduction section, completely describe contributions. For every contribution, we provide a link to the section about
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations in Section 6
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Proofs provided in Appendix
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Experimental details are discussed in Appendix C Code for the experiments is available publicly accessible github. All the datasets are available in public.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:Code is publicly available github repository. Experimental details provided Appendix C
All the datasets are publicly available.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Experimental details are discussed in Appendix C the hyperparameters and peculiarities are hard-coded in our publicly available source code.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Not all of the experiments conducted following the statistical significance protocol due to computational restrictions.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
33

---

Justification: In Appendix C we mention time and resources.
9. Code Of Ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of
Answer: [Yes]
Justification: Research conforms with NeurlPS Code of Ethics.
Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification:
We discuss broader impact in Appendix
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The research does not need special safeguards.
Licenses existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite each used assets.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: New code is available at public github repository, the code is nof-ddocueeted and follows our experimental protocol from the manuscript. The license MIT.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The research does not engage with Crowdsourcing or Human subjects.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human
Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The research does not engage with Crowdsourcing or Human subjects.
34


Reviews:
[{"review_id": "tPxKzvnuSt", "reviewer": "Reviewer_xS43", "summary": "This paper presents a new entropically regularized algorithm for identifying the Wasserstein barycenter (distribution on average closest to the reference distributions as measured in Wasserstein distance) among K distributions. \n\nThe proposed algorithm relies on a dual form of the entropic OT. Namely the entropic OT can be formulated as a supremum over the potential f and its entropic c-transform, the latter of which is defined by a measure whose partition function will be exploited in a learning algorithm defined by the objective function given in (12). \n\nThe authors proceed to provide a theoretical analysis of the proposed algorithm, providing KL and universal approximation bounds. \n\nThey then follow with a number of experimental demonstrations of the method.", "questions": "It is often the case in energy-based modeling that evaluating the loss (12) can be very high variance because of the necessity of the MCMC procedure and can have negative effect on the performance of the algorithm. Do the authors notice any such difficulty in doing so? How in general do you see benchmarking performance, which is usually important when an algorithm relies on stochastic gradient descent and a neural network, rather than stricter algorithms usually associated to EOT (like sinkhorn). I don't mean to be critical here, just reflective. The toy distributions example is nice.", "limitations": "n/a", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper nicely exposits how arrive at the primary objective function for the entropic algorithm, and along the way provides good intuition for where the EBM picture emerges (labeling of the \"partition function\" (5), etc...) and its multimarginal generalization. \n\nThe experiments seem pretty thorough, though I must say it is quite hard to test how well a Wasserstein barycenter algo works, as we have few ground truth examples to compare to, and the high-dimensional examples are always quite subjective.", "weaknesses": "- Please correct me if I'm wrong, but the discussion about convergence bounds seems a bit distracting because this is not in some sense a measurable convergence diagnostic (you do not know necessarily what $\\mathcal L^*$ in the second half of Prop 4.4). While it is good to show that the loss controls these KLs, Theorem 4.6 to me is what seems distracting unless it actually informs the method. We know NNs are universal function approximators. \n- Because multimarginal problems are usually quite hard to benchmark (as we have few known solutions), such discussions would indeed be appealing if they were practically informative, as we also know that the multimarginal problem is np-hard.\n\nSome small remarks:\n\n- Line 208: ULA is surely not the simplest MCMC algorithm. \n- personal complaint: the notation feels a bit heavyhanded and can be distracting from the main story of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for spending time reviewing our paper and for a positive evaluation of the work. We are happy that you found our work to be interesting and well-presented. Below we answer your questions and comments. \n\n**(W1) Convergence bounds discussions. \"[...] it is not [...] a measurable convergence diagnostic\". Significance of Theorem 4.6.**\n\nIt seems that there is a tiny misunderstanding. The results from Section 4.3 (Prop. 4.4, Theorem 4.5, Theorem 4.6) exhibit the generalization and approximation properties of our proposed solver. Namely, we give a complete (statistical) picture of what happens when we move from real reference distributions $\\mathbb{P}_k$ to their empirical counterparts. To the best of our knowledge, the ability to carry out such the analysis is exceptional in the existing neural OT literature, which strengthens the contribution of our work. While our Prop. 4.4 indeed does not give particular figures on how different the true $\\pi_k^*$ and recovered $\\pi^{\\widehat{f}_k}$ plans are, it provides the rates of convergence (w.r.t. number of empirical samples). Indeed, if we assume the approximation error to be zero (note Theorem 4.6), then Theorem 4.5. gives us the desired rates.\n\nWhile our Theorem 4.6. indeed looks simply like \"we have a loss function, and can approximate it arbitrarily good with NNs\", this result is not that easy. Indeed, the considered loss $\\mathcal{L}$ operates with (a) *congruent* potentials and (b) deals with entropic $C_k$-transforms of the potentials, which complicate the overall analysis. In particular, only specific properties of the entropic $C_k$-transforms allow deriving the desired universal approximation statement, see the proof of the Theorem. \n\n**(W2) Discussions on intrinsic difficulty of multimarginal problems.**\n\nThe problem we aim to tackle is indeed significantly complex by nature. The fact that our method theoretically (and practically) recovers entropic barycenters is quite encouraging for us. Having in mind the complexity of the barycenter problem, in our paper we pay special attention to the experimental setups where our method could be evaluated qualitatively or even quantitatively, see, e.g.:  Toy 2D Twister (Section 5.1), Ave Celeba  (Section 5.3), Gaussian case (Appendix C.4). \n\nAt the same time, we humbly think the discussions on the NP-hardness of the multimarginal problems are out of the scope of our paper, since our approach is based on parametric models (Neural Networks) learned with gradient descent/ascent, which eliminates such questions coming from the Theory of Algorithms.\n\n**(Q1) Practical aspects of optimizing Energy-based training objectives. Variance of the loss estimate. Reliance on MCMC, stochastic gradient descent, and neural networks.**\n\nEven in the presence of variance due to stochasticity induced by MCMC and gradient procedures, the proper tuning of the hyperparameters makes the training procedure to be successful. To practically demonstrate it, we provide the behaviour of $\\mathcal{L}_2$-UVP metric (between the unregularized ground truth barycenter mapping and our learned entropic barycenter mapping $\\pi^{f_k}(y \\vert x_k)$) for our Ave Celeba experiment, see Figure 3 in the attached pdf. We emphasize that $\\mathcal{L}_2$-UVP directly measures (by computing pointwise MSE) how the learned mapping differs from the true mapping. So, the results showcase that the presence of stochasticity in the loss does not invalidate the optimization.\n\n**(Q1.1) Benchmarking the performance**\n\nIn order to benchmark the performance, one either needs to know the true barycenter (it holds true for the majority of experiments we consider) or utilize some problem-specific metrics to check whether the recovered barycenter satisfies some expected problem-specific properties. The particular loss curves during the training are indeed not that informative.\n\n**Remark 1. \"ULA is surely not the simplest MCMC algorithm.\"**\n\nWe agree that we were a little hasty with this statement. Will be fixed in the revised version. Thank you!\n\n**Remark 2. \"[...] the notation feels a bit heavyhanded and can be distracting from the main story of the paper.\"**\n\nWe will simplify the background section in the final version of our manuscript by removing/moving to the appendix some not vital definitions/results. Please see our first answer to reviewer **unyK**.\n\n**Concluding remarks.** Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.", "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": "Qi2O9YQ5QX", "reviewer": "Reviewer_kw6j", "summary": "The paper proposes a methodology to approximate (regularised) barycenters of continuous-support distributions for arbitrary transport costs. To do so, the authors combine strategies and results from weak OT and transport map approximation via neural networks. The proposal is well connected with the existing (ever-growing) literature on barycentric approximation and provides a set of theoretical results as well as numerical examples of real-world data.", "questions": "Please refer to my comments on computational complexity and the role of the Style GAN. \n\nIn particular, after reading the paper, it seems to me that the use of the Style GAN was adopted during the final stages of the paper production, therefore, there's only brief mentions of it in the experiments but without the required /deserved attention, despite being a fundamental part of the success of the experiments.\n\nPlease elaborate on this point and, if possible, include it in the paper", "limitations": "Though the paper features a section called limitations, the authors only state that the one limitation of their method was alleviated by using Style GANs. A deeper identification and discussion of the limitations of the proposed method is lacking", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and well-positioned in the literature. The contribution features theoretical guarantees with their respective proofs and multiple experiments with different toy and benchmark datasets.", "weaknesses": "- The use of the Style GAN seems to be pivotal in the experimental validation, yet it is only explained briefly while reporting the experiments. \n\n- There is little mention of the computational complexity (either conceptually or empirical) related to the method and how it can affect its adoption by others. \n\n- There are some formatting issues: \n  - line 255: _Substituting (13) or (14) to (13) ..._\n  - lines 301-302 have spacing problems\n  - the top of fig 5 is cropped \n  - many formatting problems with the references", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer. Thank you for your efforts in reviewing our paper and for your valuable comments. We are delighted that you appreciated both the theoretical and practical contributions of our paper. Below you can find the answers to your questions and comments:\n\n**(W1) \"The use of the Style GAN seems to be pivotal in the experimental validation, yet it is only explained briefly while reporting the experiments.\" Proper attention to the idea with StyleGAN manifold.**\n\nIn the revised version of our manuscript, we will devote a specific subsection to the idea with StyleGAN manifold. We will explain the technique itself and clarify why it is significant and valuable in the context of entropic barycenters. Our reasoning here is approximately as follows:\n\n* At first, solving the entropic barycenter problem directly in the image data space leads to noisy images, see, e.g., our MNIST 0/1 experiment. Such \"deterioration\" follows from the theory - in fact, the Schrödinger barycenter problem we consider (see Appendix B.3) is known to have blurring bias, see work [15] from the references in our manuscript. Our practical reliance on MCMC also makes it obvious that the generated images from the barycenter are noisy.\n* Secondly, solving barycenter problems in image data space **is not very practical** (for standard costs $c_k$ like $\\ell^2$ or around). Indeed, our paper (MNIST 0/1 experiment in data space) and previous research (see, e.g., [1, Figure 7]) witness that $\\ell^2$ barycenter of image distributions $A$, $B$, $C$ etc.\\ is a bizarre distribution which includes \"averaged\" images of $a \\in A$, $b \\in B$, $c \\in C$ etc. So, we think that direct manipulation with image data space is more about benchmarking barycenter solvers rather than real-world applicability.\n* In contrast, operating with images with the assistance of our proposed StyleGAN-inspired manifold technique seems to be **much more practically appealing**, see the examples of potential practical applications in our Appendix B.2.\n\n**(W2) Computational complexity of the method.**\n\nAs per request, we report the (approximate) time for training and inference (batch size = $128$) of our method on different experimental setups, see the table below (the hardware is a single V100 gpu):\n\n|Phase\\Experiment   | Toy $2D$  | Gaussians  | Sphere   | MNIST(0/1) (data)  | MNIST(0/1) (manifold)  |  Ave Celeba (data) |  Ave Celeba (manifold) |\n|---|---|---|---|---|---|---|---|\n| Training  | 3h   | 6h  | 3h  | 20h  | 10h  | 60h  | 60h  | \n|  Inference  | 20s   | 40s  | 20s  | 1m  | 1m   | 2m  | 2m  |\n\nFor Ave Celeba experiment, we additionally report the computational complexity of the competitors (partially grabbed from the original papers, partially obtained through our own evaluations):\n\n| Phase\\Experiment   | WIN  | SCWB  |\n|---|---|---|\n| Training  | 160h  | 200h  |\n|  Inference |  10s | 10s  |\n\nAs we can see, all the methods in this experiment require a comparable amount of time for training. The inference with our approach is expectedly costlier than competitors due to the reliance on MCMC. In the revised version of our paper, we will add the information about the computational complexity.\n\n**(L1) Identification and discussion of the limitations apart those alleviated by StyleGAN.**\n\nIn Appendix E we have a paragraph called \"**Methodological** limitations\" which discusses the limitations not related to StyleGAN. When preparing the manuscript, we accidentally leave the word \"limitations\" in the \"**Methodological** limitations\" to be normal, not bold. It is our bug, sorry for this. Will be fixed.\n\n**Formatting issues.**\n\nThank you for pointing them out! We will go through the text carefully and do our best to fix the formatting problems in the final revision.\n\n**Concluding remarks.** Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.", "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": "e1j9wSwbfF", "reviewer": "Reviewer_unyK", "summary": "The authors focus on the Wasserstein barycenter problem; an average between distributions given a Wasserstein cost. Specifically, they consider arbitrary cost functions to define the Wasserstein distance and they approximate the continuous barycenters. They further relate their method with energy based models, and show experiments and synthetic and image datasets.", "questions": "- What is $\\mathcal{B}(\\mathbb{Q})$ in eq. 8 ?\n- Section 2.2 could be shortened, I believe we only really need to know that the barycenter is unique if it exists. \n\n## Minor comments and potential typos\n- Maybe a missing word on line 86 \"has particular\".\n- \"falls of the scope\" $\\to$ falls within the scope of the paper.\n- The authors should refer to the proofs of the theorems in the main body of the text, even if the proof is in appendix.\n- In Tab.2, could you add the standard deviation and the names of the other two solvers.", "limitations": "The authors address the limitations of their work in the manuscript.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The method is theoretically well-founded and has many applications. I particularly found interesting that it can work for multiple different ground cost functions.", "weaknesses": "- The background section may have too many details, which undermine readability and makes it harder to understand the main assumptions and challenges. I think some of it could be moved to the appendix. In particular, it is not clear to me that all definitions are useful, as some of them do not seem to be used in the subsequent sections (e.g. $\\hat{\\mathbb{P}}$). I believe the authors could improve readability by keeping only the main assumptions and definitions (hence notations) in the main body of the text, and moving other results, albeit very interesting, to the appendix. \n\n- It is a bit unclear how this work is different from the work [1]. Could the authors elaborate on this? Further, I believe it is worth comparing with [2], as it is now published at ICML 24.\n\n- Most of the experimental results are qualitative, which makes it difficult to attest if the method improves on existing ones. I understand that for most distributions the ground truth barycenter is unknown, but it could be approximated by the unregularized barycenter. Otherwise, the barycenter between two distributions relates to the McCann interpolation, a potential task could be to predict the interpolation between distributions at multiple timepoints (e.g. Tab.4, Tab.5 in [3]). A task on treatment effect using barycenters could also be used (similar to [4]).\n\n[1] 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.\n\n[2] Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin. Estimating barycenters of distributions with neural optimal transport. ICML, 2024.\n\n[3] Tong, A., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., ... & Bengio, Y. (2023). Simulation-free schr\\\" odinger bridges via score and flow matching. arXiv preprint arXiv:2307.03672.\n\n[4] Huguet, G., Tong, A., Zapatero, M. R., Tape, C. J., Wolf, G., & Krishnaswamy, S. (2023, September). Geodesic Sinkhorn for fast and accurate optimal transport on manifolds. In 2023 IEEE 33rd International Workshop on Machine Learning for Signal Processing (MLSP) (pp. 1-6). IEEE.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer. Thank you for your thoughtful review. We are pleased that you noted the theoretical and practical strengths of our work. Below we answer your comments and questions.\n\n**(W1)+(Q2). Too many details in the background section. Improving the readability by moving some secondary assumptions/definitions to the Appendix. Shortening Section 2.2.**\n\nFollowing the reviewer's request, in the revised version of our manuscript, we will shorten the Background section by leaving only vital definitions/assumptions in the main part. In particular, we will remove/move to the appendix Eq. (1) (definition of $\\text{EOT}^{(2)}$). Indeed, only Eq. (2) actually needed. Correspondingly, Eq. (7) will also disappear from the main part. Following your request, we will also shorten the explanations after Eq. (8) (Section 2.2) and eliminate introducing empirical measure $\\widehat{\\mathbb{P}}$, line 121. \n\n**(W2. Part 1) The difference between this work and [1].**\n\nPlease, take a look at our General response (1) to all the reviewers, where we get into details of the connection with [1] and our novelty over [1].\n\n**(W2. Part 2) Comparison with [2].**\n\nWe highlight that the authors of [2] acknowledge that *their work is significantly based on our paper*, see the Contribution subsection in [2]. Interestingly, they also seem to conduct a comparison of their approach with ours on Ave Celeba! barycenter benchmark. According to their results, on the manifold-constrained data setup, our method achieves FID scores $\\approx 9$, while their approach yields a much bigger number $\\approx 30$. Therefore, our method demonstrates a significant quantitative boost. This is a clear indicator of the strength of our energy-guided methodology.\n\n**(W3) \"Most of the experimental results are qualitative, which makes it difficult to attest if the method improves on existing ones.\" Call for more experiments with quantitative outcomes.**\n\nFollowing your kind suggestion, we conducted a novel experiment with predicting the interpolation between distributions at multiple timepoints following the setup identical to [3]. Please refer to the general answer to all the reviewers. According to the results, our method works on par with competitive methods.\n\nFor completeness, we also note that we have a purely quantitative experiment in the Gaussian case in Appendix C.4.\n\n**(Q1) \"What is $\\mathcal{B}(\\mathbb{Q})$ in eq. 8?\"**\n\nThe term $\\mathcal{B}(\\mathbb{Q})$ is defined directly in equation (8), i.e., $\\mathcal{B}(\\mathbb{Q}):= \\sum\\limits_{k = 1}^{K} \\lambda_k \\text{EOT}_{c_k, \\varepsilon}(\\mathbb{P}_k, \\mathbb{Q}).$ We agree that this makes (8) to be a bit overcomplicated, which was done due to space constraints. Sorry for this. In the final version, we will have one extra page and we will separately define $\\mathcal{B}(\\mathbb{Q})$ or even get rid of this notion completely. \n\n**Minor comments and typos**\n\nThank you for pointing them out! We will fix them in the revised version of the manuscript. The FIDs for our method with std. deviations (running the inference with different seeds) look as:\n\n| k=1 | k=1  | k=2  |\n|------|------|------|\n| 8.4 (0.3)  | 8.7 (0.3)  | 10.2 (0.7) |\n\nFor the competitors, we grab the results from WIN paper where no std. deviations are provided.\n\n**Concluding remarks.** Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.\n\n**References**\n\n[1] Mokrov et. al., Energy-guided entropic neural optimal transport. ICLR'2024.\n\n[2] Kolesov et. al., Estimating barycenters of distributions with neural optimal transport. ICML'2024. \n\n[3] Tong et. al. Simulation-free Schrödinger bridges via score and flow matching. AISTATS'2024", "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": "1jT9pWW8mA", "reviewer": "Reviewer_wGnY", "summary": "The authors propose a new algorithm to estimate entropic barycenters in this submission. \nIn particular, leveraging the c-transform based duality of the entropic OT problem, the authors reformulate the entropic barycenter problem as a constrained functional optimization problem, in which a set of dual potential functions are optimized to minimize the energy function corresponding to the c-transform. \nThe potential functions are parametrized by neural networks so that the constrained functional optimization problem can be implemented as an unconstrained and parametric optimization problem, which can be solved efficiently by the gradient descent associated with MCMC sampling.\nCompared to existing OT-based barycenter estimation methods, the proposed algorithm can be applied to general cost functions and can explicitly estimate OT plans.", "questions": "Please see above.", "limitations": "Not applicable.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The motivation is clear, and the derivations of key formulas are clear and easy to follow.\n- The theoretical part is solid, which supports the rationality of the proposed method, including its sample complexity, generalization power, and so on.", "weaknesses": "- It has been well-known that entropic OT suffers from the sensitivity of the entropic regularizer’s weight. It is not clear to me why the proposed method can support high-dimensional applications. Especially, in Lines 208-212, epsilon impacts a lot on the MCMC sampling step. The robustness test of the proposed method to the regularizer’s weight should be considered in the experimental part.\n- Parameterizing continuous potential functions by neural networks means shrinking the feasible space of the problem. How to design the neural networks? Again, the impacts of different NN models and the robustness of the method to the model selection should be considered.\n- The proposed algorithm involves MCMC sampling and gradient descent, so the convergence and the computational efficiency of the proposed algorithm should be analyzed quantitatively. In particular, in high-dimensional cases, do we need to sample a lot?\n\nMinors:\n- It would be nice if the authors could apply the proposed method in a practical task, e.g., point cloud clustering.\n- For the completeness of the main paper, a conclusion section should be added.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer. Thank you for spending time reading and evaluating our work.  We were greatly encouraged by your positive feedback on our theory and the readability of our manuscript.  Below you can find the answers to your questions and comments.\n\n**(W1) The role and impact of entropic regularizer's weight. Robustness test of the method to the regularizer's weight.**\n\nWe agree that considering too small $\\epsilon$ within our proposed method may cause some difficulties with optimizing our energy-based barycenter objective. The heart of the matter is the answer to: \n\"can our approch permit *practically-reasonable* choices for $\\epsilon$\"? \nWe believe that our considered practical use cases (Figure 5, Figure 7, Figure 8) testify that this is the case. The images sampled from our learned conditional plans $\\pi^k(y \\vert x)$ are of reasonable diversity. We want to emphasize that for Ave Celeba! experiment (manifold constrained) the considered value of $\\epsilon$ was $10^{-4}$. \nAlso, as per request, we additionally deployed our methodology on MNIST 0/1, manifold setup, with an $\\epsilon$ of $10^{-4}$ (in addition to the previously considered value of $\\epsilon=10^{2}$).  The results are demonstrated in **Figure 1 in the attached pdf**. \n\nRegarding the high-dimensional applications, our comment is as follows. It seems that learning an (especially entropic) barycenter directly in the image data space is, in most cases, not very meaningful - we will obtain just blurred pixel-wise averages of input images. This is exactly where our StyleGAN trick comes into play. In turn, the latent space of StyleGAN is moderate dimensional and allows learning meaningful barycenters in high-dimensions. And our experimental results demonstrate that our approach successfully tackles this setup.\n\n**(W2) NNs architectures to parameterize the potential functions.**\n\nOne of the key insights of our method in practice is the reliance of our approach on already well-established, built-on efficient architectures (and hyper-parameter selections) constructed for EBMs and EOT problems [1]. Since we trust the architectural design and hyperparameter specifications of those authors, we perform no additional find-tuning, using their neural networks as pre-trained foundation models. We think that doing this, in addition, is out of the scope of the current work. Please note that our primal contribution is methodological, not architectural, but has a primarily theoretical intent.\n\n**(W3) Reliance on MCMC and gradient descent. Quantitative analysis of convergence and computational efficiency in, e.g., high-dimensional case.**\n\nWe adopt MCMC and (stochastic) gradient descent as the standard tools for Energy-based Models. It is a generic knowledge that these tools allow for generating high-dimensional images (when carefully tuned); see references [34], [112] from our paper.\n\nFollowing your request, we quantitatively analyze the performance of our method as a function of the number of discretized Langevin dynamic steps, see the charts (Figure 4) in our attached pdf. \nIn this experiment, we try to learn the barycenter of Gaussian distributions; which, as explained in Appendix C.4, is known in-closed form for the unregularized case, which gives us access to a theoretical object by which we may compare and validate our computational pipeline. We pick the dimensionality $D = 64$ (which is the highest-dimensional case among the considered). As we can see, performance drops when the number of steps is insufficient. Overall, in all our experiments, the number of Langevin steps is chosen to achieve reasonable qualitative/quantitative results.\n\n**(Minor W4) Practical application.**\n\nFor an intriguing practical application, we pick an intriguing Single-Cell Data analysis experiment, where we need to predict the distribution of cell population at some intermediate time point $t_i$ given cell distributions at $t_{i + 1}$ and $t_{i - 1}$. Please see a brief explanation (and the relevant links) of the problem in our General response, our results are in Table 1 in the attached pdf. Our approach achieves competitive performance.\n\n**(Minor W5) Call for conclusion section.**\n\nSection E contains our conclusions and summary. Currently, this section is located in the Appendix. However, following your suggestion, we will move it to the main part of our manuscript in the final version.\n\n**Concluding remarks.** Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.\n\n**References**\n\n[1] Mokrov et. al., Energy-guided entropic neural optimal transport. ICLR'2024.", "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": "jctsgNZJ2q", "reviewer": "Reviewer_LoLb", "summary": "This paper deals with the question of estimating Wasserstein(-like) barycenters in the _continuous_ case based on _samples_. That is, one actually observes (iid) samples $X^k$ made of $N_k$ points from (unknown) probability distributions $\\mathbb{P}_1,\\dots,\\mathbb{P}_K$, and the goal is to estimate $\\mathbb{Q}^*$ the barycenters of the $(\\mathbb{P}_k)_k$ (and not simply the discrete barycenter of the $X^k$). \n\nThe idea developed in this work is to maximize a loss (arising from a dual problem of an _entropy-regularized_ OT formulation) $L(f_1,\\dots,f_K)$ over $K$ potential $(f_k)_k$ (themselves parametrized by neural networks). A given $f_k$ induces a measure $\\pi^{f_k}$ whose first marginal is $\\mathbb{P}_k$, and the second marginal is some $\\mathbb{Q}_k$ that aims at approximating $\\mathbb{Q}^*$. \nOf interest, $\\pi^{f_k}$ admits a disintegration and its second marginal can be sampled by the following scheme :\n- sample $x_k$ from $\\mathbb{P}_k$ (i.e. take an observation in the training set),\n- Then sample $y \\sim \\mu_{x_k}^{f_k} \\propto e^{-V}$ with potential $V(y) \\coloneqq c(x_k, y) - f_k(y)$, typically using a Langevin-type approach. \nThe objective value $L$ itself may not be accessible in close form (because the normalization constant of $\\mu_{x_k}^{f_k}$ is unknown), but its gradient---as usual in score-based problem---is accessible and optimization is thus possible though a practical algorithm. \n\nFrom a theoretical standpoint, authors prove that the error gap $(\\max L) - L(f_1,\\dots,f_K)$ controls the distances $\\mathrm{KL}(\\mathbb{Q}^* | \\mathbb{Q_k})$, i.e. a small error gap ensures that the second marginals of the $\\mathbb{Q_k}$ should be close to the target $\\mathbb{Q}^*$. They also establish a statistical consistency result, in that if the number of observations $(N_k)_k$ goes to $\\infty$ (for all $k$), the estimation error of their model goes to zero. \n\nThe approach is showcases on toy and images datasets.", "questions": "1. The approach relies on a Langevin (or other sampling) scheme with potential energy given by $f_k - c(x_k, \\dot)$. Given that $f_k$ is encoded by a neural network which is typically neither convex not smooth in the space variable $y$ (nor in the parameters $\\theta$ but that does not matter here), it's dubious to assume that the Langevin scheme converge toward the target distribution. This is only very briefly mentioned in lines 363-364 and in line 1051 and I believe this should be discussed more extensively (i.e. did you notice it during your experiments? how impactful is it?)\n\n2. While the choice the entropic regularization ($H(\\pi)$, $\\mathrm{KL}(\\pi | \\mathbb{P} \\otimes \\mathbb{Q})$...) is not crucial _when the marginal are fixed_, it starts to play a role when optimizing over a given marginal, as in the barycenter problem. This work consider the Schrodinger barycenter (using the terminology of Chizat [15, Table 1]) as discussed in appendix B.3---which also discusses the importance of this choice from a practical viewpoint. However, this modelling induces a \"blurring\" bias which is never mentioned in the work as far as I can tell. Can you briefly elaborate on this? For instance, if the input distribution are actually given by $\\delta_{x_k}$, what would be the output of the algorithm?\n\n3. As far as I can tell, the barycenter that is targeted here similar (if not exactly the same) as the $\\alpha_{\\mathrm{OT}^\\mathcal{U}}$ of Janati et al., _Debiased Sinkhorn barycenters_ (note: this work should be cited/discussed at some point). is that correct? If so, note that in that case an  _closed form_ expression for the entropic barycenter of (univariate) Gaussian distributions is accessible (showcasing formally the blurring effect); and section C.4. may be revisited under this lens.", "limitations": "Yes, the authors discuss the limitations of their work in a convincing way (in the appendix; ).", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "First, the paper is well-tailored for NeurIPS in my opinion: it tackles a problem already studied in the ML community, propose a clear methodology that yields a practical algorithm and that is supported by reasonable theoretical results, and conduct both toy and demanding experiments. The main paper is mostly self-contained, while the appendix provides insightful complementary comments. \n\nThe article provides a fairly good report of sota models/related works. \n\nMoreover, the approach can handle general cost functions (and not only the widely used squared Euclidean cost), and this advantage is nicely showcased in the experiment by enabling to _parametrize_ the barycenter (e.g. by leveraging a prior) on a latent space $Z$, i.e. the cost becomes $c(x,z) = |x - G(z)|^2$ for instance. This perspective is satisfactory as it opens the way for interesting applications of OT-based barycenters.", "weaknesses": "1. The main weakness (and for which I unfortunately do not have a specific solution) of the approach is the difficulty to properly evaluate it (note: this is not specific to the method but rather to the problem considered), due to:\n- the actual (regularized, or not) barycenter $\\mathbb{Q}^*$ being usually unknown,\n- the value of the loss $L$ (a fortiori the optimal value $L^*$ and the error gap) cannot be computed,\n- even if we could compute it, it is natural to run experiments with a small regularization parameter $\\epsilon$, but in this regime the approximation bound provided by Theorem becomes useless.\nIt thus makes very hard to \"monitor\" the algorithm and to have a practical clue on whether its output can be trusted or not. This does not invalidate the methodology in my opinion, but remains an important weakness that must be stressed. \n\n2. In terms of clarity, while I consider that the \"methodology\" part of the work (section 4) is reasonably well-written, the experimental section is much harder to follow in my opinion: I do not exactly understand if \"$\\ell^2$-barycenters\" is supposed to mean \"Wasserstein barycenters for the quadratic cost\" or \"barycenter between images seen as vector (standard means)\". Explicitly, lines 317--320 claim that those are \"piwel-wise average\" (so, in my understanding, standard means) \"coming from $\\ell^2$ OT plan\" (so, in my understanding, geodesics/McCann interpolation; from which you have access to the Wasserstein barycenter). \n\n3. This work takes many ideas from [77] (energy-based optimization, etc.), which to some extend diminishes its originality. This is discussed in lines 813--819. I believe that this should belong to the main paper and be slightly extended, explaining more precisely what are _the theoretical and practical difficulties_ that had to be overcome to go from [77] to the actual work (the discussion refers to \"multiple marginals, optimization with an unfixed marginal $\\mathbb{Q}$\", and I agree with that, but in some sense the proposed approach seems to adapt so naturally (which is good!) from [77] that I do not see \"what could have gone wrong\").", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for spending time reviewing our paper. Your kind words that our paper well suits NeurIPS are encouraging for us. Below you can find the answers to your comments and questions.\n\n**(W1) Difficulties with proper evaluation.**\n\nWe agree that the out-of-the-box evaluation is typically unavailable for the problem we aim to tackle (entropic barycenter). This is due to the inherent *complexity* of the barycenter problem. In particular, there seem to be no easy-to-calculate statistics to monitor to control the quality of the recovered barycenter (like the loss gap $L - L^*$). Surprisingly, we were able to derive a solution to the notoriously challenging entropic barycenter problem. On the other hand, we did our best to validate our methodology as accurately as possible. In particular, we consider several experimental setups where the unregularized barycenter is **known**: (a) Toy 2D Twister (Section 5.1), (b) Ave Celeba experiment (Section 5.3), Gaussian case (Appendix C.4). These setups allow us quantitatively/qualitatively estimate the recovered EOT barycenter. \n\n*A comment on small $\\epsilon$ regime.* From the theoretical point of view (approximation bounds etc.), the regime $\\epsilon \\rightarrow 0$ is indeed not so favourable. However, in terms of method evaluation, it is the regime that allows more-or-less accurate ways to \"monitor\" the quality. This is thanks to the aforementioned practical setups, where the true *unregularized* barycenters are known, and we can directly compare them with the generated ones. Practically speaking, in our paper, we consider rather small values of $\\epsilon$ (up to $10^{-4}$), and our approach treats such choices.\n\n**(W2) Questions about the experimental section and $\\ell^2$ barycenters.**\n\nYour understanding is correct: when we consider $\\ell^{2}$ barycenters of images, the $\\ell^{2}$ between two particular images is computed by flattening each image into a vector (e.g., $32\\times 32\\rightarrow 1024$). In turn, when there are **just two** distributions ($K=2$) of images considered, their $\\ell^{2}$ barycenter is indeed a certain McCahn interpolant. However, one needs to recover somehow the true $\\ell^{2}$ OT map (plan) between the distributions to access the barycenter through this interpolant.\n\n**(W3) Our work vs. [77]. Theoretical and practical efforts upon [77] which are required to end up with our approach.**\n\nPlease find the answer to your comment in the general answer to all the reviewers.\n\n**(Q1) Non-convexity of the energy potentials. Extending the discussion and practical observations.**\n\nThe problem with the non-convex energy function and its effect on MCMC is a hallmark of energy-based methods (and even score-based approaches). The non-convexity may cause some MCMC chains to become trapped by local minima of the energy landscape, making the corresponding samples somewhat different from the expected samples from $\\pi^f_k(y \\vert x_k)$. By taking a closer look at our results, e.g., Figure 7 (conditional samples from MNIST 0/1 barycenter), one may distinguish that some generated pictures \"fall out of the line\", i.e., noticeably differ from the others. We think this phenomenon is inherently caused by the non-convexity. We underline that non-convexity does not stop EBMs from successfully modelling high-dimensional distributions (see, e.g., references [34], [112] from our paper), which makes our methodology practically appealing. In the revised version of our manuscript, we will expand the discussion on the non-convexity.\n\n**(Q2) The type of the considered Entropic Barycenter problem. \"Blurring\" effect. Entropic barycenter of $\\delta$-distributions.**\n\nPlease note that we directly mention the blurring bias in our experiment with MNIST01 in data space (section 5.2, Figure 5 and lines 323-325) in our paper. However, we do not focus on it too much because in the manifold-constrained setups (MNIST01 in Section 5.2, Ave, Celeba in Section 5.3) this bias in fact disappears. More generally, we believe that manifold-constrained setups are generally more interesting and probably even necessary to obtain meaningful barycenters for downstream practical problems (see Appendix B.2 for a discussion).\n\nRegarding the $\\epsilon$-barycenters of $\\mathbb{P}\\_{k}=\\delta_{x_{k}}$, it seems like it is some Gaussian centered at $\\sum_{k=1}^{K}\\lambda_k x_{k}$ with variance proportional to $\\epsilon$ (in the case of $\\ell^{2}$ barycenters). Namely, the higher $\\epsilon$ is, the more \"blurry\" (noisy) the $\\epsilon$-barycenter is compared to the unregularized one. Following your question, we will add an extended discussion of this aspect to B.3.\n\n**(Q3) Missed citation. (Possible) closed-form expression for entropic barycenters in 1D.**\n\nThanks for noting, we were not aware of this paper. It indeed seems like their $\\alpha_{\\text{OT}^{\\mathcal{L}}}$ barycenter coincides with our considered Schrodinger barycenter. We will surely include the citation and proper discussion, including their formula for the closed-form barycenter for 1D Gaussians and $\\ell^{2}$ cost.\n\n**Concluding remarks.** Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. Thank you one more time for your efforts.", "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:
Dear reviewers, thank you for taking the time to review our paper! It is a great pleasure for us that you positively evaluated our paper, emphasizing the importance of the considered problem, our solid theoretical contribution and well-tailored experimental validation. Please find the answers to your shared questions below and answers to your non-intersecting questions in the corresponding replies to your reviews.

**(1) Our work vs. Energy-guided EOT paper [1] ([77] from the original manuscript). (reviewers unyK and LoLb)**

Following your request, we explain the theoretical, methodological and practical efforts needed to develop the ideas from [1] and, eventually, come up with our work. 

In [1], the authors solve the Entropic *Optimal Transport* (EOT) problem, while here, we built atop their methodology to solve the Entropic *Barycenter problem*. These two problems are linked through their use of optimal transport tools, but they are otherwise unrelated.  In particular, the theoretical guarantees for the EOT problem do not imply guarantees for the entropic barycenter problem, and the latter requires theoretical guarantees, built on distinct tools.  Our paper fills the theoretical and implication gaps for the entropic barycenter problem. 

We would like to begin by underscoring the **novelty** of our key theoretical result, namely Theorem 4.1, whose derivation was highly non-trivial. The beauty of our Theorem is that it allows solving the barycenter problem without explicit knowledge of the barycenter distribution itself. This *drastically differs* from the classical EOT problem considered in [1], where both marginal distributions are known through their samples.  Unlike the EOT problem, the nature barycenter problem studied in our paper requires novel technical insights, such as being able to accommodate multiple potentials and observing that these necessarily have to obey a certain congruence condition, see Theorem 4.1 and line 167. These, and several other technical obstructions, imply that there is a lengthy mathematical and technical leap from the theoretical guarantees/implementation peculiarities in [1] to those which we derived in our manuscript.  

In addition, our work is not limited by the proposed methodology for solving the entropic barycenter problem.  The technology, both theoretical and code, which we developed for solving the entropic barycenter problem, has a number of non-trivial future implications for the solvability of other problems in the broader area of computational optimal transport.  For instance, we highlight our concept with learning barycenter on an image manifold of a pre-trained StyleGAN; which allows for *meaningful* recovery of the barycenter distribution $\mathbb{Q}^*$ and would lead to new applications of our technology (see Section B.2 in the Appendix). Secondly, we have pioneering statistical learning theoretic results for the Entropic barycenter problem in (Section 4.3.), which may be of independent interest in the learning theory sub-community. Specifically, our Theorem 4.5 shows that the representativeness $\mathbb{E} Rep_{X_k}(\mathcal{F}_k^{C_k}, \mathbb{P}_k)$ with respect to the class of entropic $C_k$-transforms can be bounded by a rate independent of the functional class $\mathcal{F}$ and the cost $c_k$ functions' properties, which seems to be a bit counterintuitively.

**(2) New experiment with the biological data (reviewers unyK, wGnY)**

Following the suggestion of reviewer unyK and the request of reviewer wGnY to consider more practical tasks, we considered the problem setup of predicting the interpolation between single cell populations at multiple timepoints from [3].  The objective is to *interpolate* the distribution of cell population at any intermediate point in time, call it $t_i$, the cell population distributions at past and future time-points $t_{i + 1}$ and $t_{i - 1}$.  Since this is an interpolation problem, it is natural to expect that the intermediate population is a (entropy-regularized) barycentric average (with $\ell^{2}$ cost) of both the population distributions available at the nearest preceding and future times. We leverage the data pre-processing and metrics from paper [2], wherein the authors provide a complete notebook with the code to launch the setup similar to [3].  There are 3 experimental subsetups with dimensions $D=50,100,1000$, and the metric is MMD;  see [3] or [2, Section 5.3] for additional experimental details. We report the result in **Table 1 in the attached PDF** where we find that in most cases, our general-purpose entropic barycenter approach nearly matches the performance of the leading computational in computational biology, which are specifically designed to solve only this problem. This underscores the scope of problems in which our barycentric optimal transport technology can act as a viable foundation model, directly out-of-the-box.

**PDF content**.

We attach **PDF** file that contains additional experiments for the rebuttal. A brief explanation of the content:

* We provide evidence of our method's robustness for different entropic regularization setups on the MNIST 0/1 dataset (**wGnY**). This experiment demonstrates that the smaller the regularization term $\epsilon$, the better preserving the image content.
* We demonstrate convergence $\mathcal{L}$2-UVP of our method during the training phase on the Ave, Celeba! dataset (**kw6j, xs43**). 
* We show the performance of our method depending on the number of Langevin steps in the case of Gaussian distributions' barycenter (**wGnY**).
* We provide the practical application of our method with Single-cell experiment, comparing our approach with existing solvers (**wGnY, unyK**).

**References**

[1] Mokrov et. al., Energy-guided entropic neural optimal transport. ICLR'2024. 

[2] Korotin et. al., Light Schrödinger Bridge. ICLR'2023

[3] Tong et. al. Simulation-free schrödinger bridges via score and flow matching. AISTATS'2024

Meta-review: The paper proposes a new methodology for estimating entropic regularized Wasserstein barycenters of $K$ distributions, for arbitrary transport cost function $c$. This approach is based on a dual formulation of the original problem that the authors establish, leveraging the weak entropic $c$-transform. This dual problem consists in maximizing a convex combination of functionals, over $K$ continuous functions under some constraints. Based on this result, the authors provide an algorithm to solve this dual problem, using stochastic gradient descent and MCMC samplers. The proposed approach is well contextualized in relation to existing work and comes with both theoretical results and interesting numerical illustrations with real-world data. 

All the reviewers and I found that the paper introduces interesting and novel ideas for solving the problem at hand, which are well-developed. Therefore, we think that the paper is worth accepting.
Decision: Accept (spotlight)