Title: Energy-Guided Continuous Entropic Barycenter Estimation for General Costs
Venue: NeurIPS 2024
Accepted: yes
Keywords: energy-based model

generative model

optimal transport

entropic optimal transport barycenters

general optimal transport cost
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

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 .

1 Introduction

Averaging is a fundamental concept in mathematics and plays a central role in numerous applications. While it is a 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 a different way of taking averages. To address this issue, one may carefully select a measure of distance that properly captures similarity in the space of probabilities. Then, the task is to find a procedure which identifies a 'center' that, on average, is closest to the reference distributions. One good choice for comparing and averaging probability distributions is provided by the family of Optimal Transport (OT) discrepancies [110]. They have clear geometrical meaning and practical

X1~P1

X2~P2

X3~P3

X4~P4

Y~Q

3.0

2.5

Zenith

2.0

1.5

1.0

0.5

0.0

-3-2-101

23

(a) The unfolded sphere.

(b) The sphere viewed from different viewpoints.

interpretation [89, 97]. The corresponding problem of averaging probability distributions using OT discrepancies is known as the OT barycenter problem [1]. OT-based barycenters find application in various practical domains: domain adaptation [79], shape interpolation [98], Bayesian inference [101, 102], 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, see Appendix B.1 for more details. The continuous setting turns out to be even more challenging, with only a handful of recent works devoted to this setup [72, 17, 59, 55, 32, 82, 14]. Most of these works are devoted to specific OT cost functions, e.g., deal with ℓ 2 2 barycenters [59, 55, 32, 82]; while others require non-trivial a priori selections [72] and have limiting expressivity and generative ability [17, 14], see §3 for a detailed discussion.

Contributions . We propose a novel approach for solving Entropy-regularized OT (EOT) barycenter problems, which alleviates the aforementioned limitations of existing continuous OT solvers.

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 a simple optimization procedure which closely relates to the standard training algorithm of Energy-Based models (§4.2).

We establish the generalization bounds as well as the universal approximation guarantees for our recovered EOT plans, which push the reference distributions to the barycenter (§4.3).

We validate the applicability of our approach on various toy and large-scale setups, including the RGB image domain (§5). In contrast to previous works, we also pay attention to non-Euclidean OT costs. Specifically, we conduct a series of experiments looking for a 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 , 2 , . . . , K } . Throughout the paper X ⊂ R D ′ , Y ⊂ R D and X k ⊂ R D k ( k ∈ K ) are compact subsets of Euclidean spaces. Continuous functions on X are denoted as C ( X ) . Probability distributions on X are P ( X ) . Absolutely continuous probability distributions on X are denoted by P ac ( X ) ⊂ P ( X ) . Given P ∈ P ( X ) , Q ∈ P ( Y ) , we use Π( P , Q ) to designate the set of transport plans , i.e., probability distributions on X × Y with the first and second marginals given by P and Q , respectively. The density of P ∈ P ac ( X ) w.r.t. the Lebesgue measure is denoted by d P ( x ) d x .

2 Background

First, we recall the formulations of EOT (§2.1) and the barycenter problem (§2.2). Next, we clarify the computational setup of the considered EOT barycenter task (§2.3).

2.1 Entropic Optimal Transport

Consider distributions P ∈ P ac ( X ) , Q ∈ P ac ( Y ) , a continuous cost function c : X × Y → R and a regularization parameter ϵ > 0 . The entropic optimal transportation (EOT) problem between P and Q [15, 78] consists of finding a minimizer of

Note that (1) is not the only way to formulate EOT. One more popular and equivalent formulation [19, 83, 36] substitutes the conditional entropy term E x ∼ P H ( π ( ·| x )) in (1) with full entropy H ( π ) .

Aminimizer π ∗ ∈ Π( P , Q ) of (1) is called the EOT plan; its existence and uniqueness are guaranteed, see, e.g., [16, Th. 3.3]. In practice, we usually do not require the EOT plan π ∗ but its conditional distributions π ∗ ( ·| x ) ∈ P ( Y ) as they prescribe how points x ∈ X are stochastically mapped to Y [42, §2]. We refer to π ∗ ( ·| x ) as the conditional plans .

Weak OT dual formulation of the EOT problem . The EOT problem permits several dual formulations. In our paper, we use the one derived from the weak OT theory [39, Theorem 9.5]:

where f C : X → R is the so-called weak entropic c -transform [9, Eq. 1.2] of the function ( potential ) f . The transform is defined by

We use the capital C in f C to distinguish the weak transform from the classic c -transform [89, §1.6] or ( c, ϵ ) -transform [76, §2]. In particular, formulation (2) is not to be confused with the conventional EOT dual, see [78, Appendix A].

For each x ∈ X , the minimizer µ f x ∈ P ( Y ) of the weak c -transform (3) exists and is unique. Its density has particular form [78, Theorem 1]. Let Z c ( f, x ) def = ∫ Y exp ( f ( y ) -c ( x,y ) ϵ ) d y , then

By substituting (4) into (3) and carrying out straightforward manipulations, we arrive at an explicit formula f C ( x ) = -ϵ log Z c ( f, x ) , see [78, Equation (14)].

2.2 Entropic OT Barycenter

Consider distributions P k ∈ P ac ( X k ) and continuous cost functions c k ( · , · ) : X k ×Y → R for k ∈ K . Given weights λ k > 0 with ∑ K k =1 λ k = 1 , the EOT Barycenter problem is:

The case where ϵ = 0 corresponds to the classical OT barycenter [1] and falls out 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.e., which has the same minimizers . The objective (5) is known as Schrödinger barycenter problem [15, Table 1], see the extended discussion in Appendix B.3. It is worth noting that under mild assumptions the barycenter Q ∗ which delivers optimal value of (5) exists and is unique, see Appendix A.7.

2.3 Computational aspects of the EOT barycenter task

Barycenter problems, such as (5), are known to be challenging in practice [2]. To our knowledge, even when P 1 , . . . , P K are Gaussian distributions, there is no direct analytical solution neither for our entropic case ( ϵ > 0 ), see the additional discussion in App. C.4, nor for the unregularized case [3, ϵ = 0 ]. Moreover, in real-world scenarios, the distributions P k ( k ∈ K ) are typically not available explicitly but only through empirical samples (datasets). This aspect leads to the next learning setup .

We assume that each P k is accessible only by a limited number of i.i.d. empirical samples X k = { x 1 k , x 2 k , . . . x N k k } ∼ P k . Our aim is to approximate the optimal conditional plans π ∗ k ( ·| x k ) between the entire source distributions P k and the entire (unknown) barycenter Q ∗ solving (5). The recovered plans should provide the out-of-sample estimation, i.e., allow generating samples from π ∗ k ( ·| x new k ) , where x new k is a new sample from P k which is not necessarily present in 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-sample estimation required in the continuous OT setup.

3 Related works

The taxonomy of OT solvers is large. Due to space constraints, we discuss here only methods within the continuous OT learning setup that solve the (E-)OT barycenter problem. These methods approximate OT maps or plans between the distributions P k and the barycenter Q ∗ rather than just their empirical counterparts that are available from the training samples. A broader discussion of general-purpose discrete/continuous (E-)OT solvers is in Appendix B.1.

Continuous OT barycenter solvers are based on the unregularized or regularized OT barycenter problem within the continuous OT learning setup. The works [59, 32, 82, 55] are designed exclusively for the quadratic Euclidean cost ℓ 2 ( x, y ) def = 1 2 ∥ x -y ∥ 2 2 . The OT problem with this particular cost exhibits several advantageous theoretical properties [4, §2] which are exploited by the aforementioned articles to build efficient pro-

cedures for barycenter estimation algorithms. In particular, [59, 32] utilize ICNNs [6] which parameterize 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 c 1 , . . . , c K . In [72], they also consider regularized OT with non-Euclidean 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 a 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 1, 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 it from our comparisons.

4 Proposed Barycenter Solver

In the first two subsections, we work out our optimization objective (§4.1) and its practical implementation (§4.2). In §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 ( Y ) K → R :

Here f C k k denotes the weak entropic c k -transform (3) of f k . Following §2.1, we see that it coincides with -ϵ log Z c k ( f k , x k ) . 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 (5) permits the following dual formulation:

We refer to the constraint ∑ K k =1 λ k f k = 0 as the congruence condition w.r.t. weights { λ k } K k =1 . The potentials f k appearing in (6) play the same role as in (2). Notably, when L ( f 1 , . . . , f K ) is close to L ∗ , the conditional optimal transport plans π ∗ k ( ·| x k ) , x k ∈ X k , between P k and the barycenter distribution Q ∗ can be approximately recovered through the potentials f k . This intuition is formalized in Theorem 4.2 below. First, for f k ∈ C ( Y ) , we define

and set Q f k ∈ P ( Y ) to be the second marginal of π f k .

Theorem 4.2 (Quality bound of plans recovered from dual potentials [proof ref.]) . Let { f k } K k =1 , f k ∈ C ( Y ) be congruent potentials. Then we have

where π ∗ k ∈ Π( P k , Q ∗ ) , k ∈ K are the EOT plans between P k and the barycenter distribution Q ∗ .

Algorithm 1: EOT barycenters via Energy-Based Modelling

Input: Distributions P k , k ∈ K accessible by samples; cost functions c k ( x k , y ) : X k ×Y → R ; the regularization coeff. ϵ > 0 ; barycenter averaging coeff. λ k > 0 : ∑ K k =1 λ k = 1 ; MCMCprocedure MCMC_proc ; batch size S > 0 ; NNs f θ,k : Y → R , s.t. ∑ K k =1 λ k f θ,k ≡ 0 (see §4.2).

Output: Trained NNs f θ ∗ ,k recovering the conditional OT plans between P k and barycenter Q ∗ .

According to Theorem 4.2, an approximate solution { f k } K k =1 of (6) yields distributions π f k which are close to the optimal plans π ∗ k . Each π f k is formed by conditional distributions µ f k x k , c.f. (4), with closed-form energy function, i.e., the unnormalized log-likelihood. Consequently, one can generate samples from µ f k x k using standard MCMC techniques [7]. In the next subsection, we stick to the practical aspects of optimization of (6), which bears certain similarities to the training of Energy-Based models [69, 100, EBM].

Relation to prior works. Works [72, 59] 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 f k ∈ C ( Y ) for k ∈ K with neural networks f θ,k , θ ∈ Θ . In order to eliminate the constraint in (6), we parametrize f θ,k as g θ k -∑ K k ′ =1 λ k ′ g θ k ′ , where { g θ k : R D → R , θ k ∈ Θ k } K k =1 are neural networks. This parameterization automatically ensures the congruence condition ∑ K k =1 λ k f θ,k ≡ 0 . Note that Θ=Θ 1 ×··· × Θ K and θ =( θ 1 , . . . , θ K ) ∈ Θ . Our objective function for (6) is

The direct computation of the normalizing constant Z c k may be infeasible. Still, the gradient of L with respect to θ 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

With this result, we can describe our proposed algorithm which maximizes L using (9).

TRAINING. To perform stochastic gradient ascent step w.r.t. θ , we approximate (9) with Monte-Carlo by drawing samples from d π f θ,k ( x k , y ) = d µ f θ,k x k ( y )d P k ( x k ) . Analogously to [78, §3.2], this can be achieved by a simple two-stage procedure. At first, we draw a random vector x k from P k . This is done by picking a random empirical sample from the available dataset X k . Then, we need to draw a sample from the distribution µ f θ,k x k . Since we know the negative energy (unnormalized log density) of µ f θ,k x k by (4), we can sample from this distribution by applying an MCMC procedure which uses the negative energy function ϵ -1 ( f θ,k ( y ) -c k ( x k , y )) as the input. Our findings are summarized in Algorithm 1. Note that typical MCMC needs the energy functions, in particular, costs c k , to be differentiable. Otherwise, one can consider gradient-free procedures, e.g., [92, 104].

In all our experiments, we use ULA [86, §1.4.1] as a MCMC_proc . It is a conventional MCMC algorithm. Specifically, in order to draw a sample y k ∼ µ f θ,k x k , where x k ∈ X k , we initialize y (0) k from the D -dimensional distribution N (0 , I D ) and then iterate the discretized Langevin dynamics:

where ξ l ∼ N (0 , I D ) , l ∈ { 0 , 1 , 2 , . . . , L } , L is a 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 Y ( l ) k conditioned on X ( l ) k . The particular values of number of steps L and step size η 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 π f θ ∗ ,k ( ·| x k ) , see the details on the hyperparameters L, η in §5.

Relation to prior works. Learning a distribution of interest via its energy function (EBMs) is a well-established direction in generative modelling research [69, 111, 28, 100]. Similar to ours, the key step in most energy-based approaches is the MCMC procedure which recovers samples from a distribution accessible only by an unnormalized log density. Typically, various techniques are employed to 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 is to introduce and validate the methodology for computing EOT barycenters in an energy-guided manner. Therefore, we opt for the simplest MCMCalgorithm, even without the replay buffer [46], 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 π ∗ k between P k and Q . In practice, for each distribution P k we know only the empirical samples X k = { x 1 k , x 2 k , . . . x N k k } ∼ P k , i.e., finite datasets. Besides, the available potentials f k , k ∈ K come from restricted classes of functions and satisfy the congruence condition. More precisely, we have f k = g k -∑ K k =1 λ k g k (§4.2), where each g k is picked from some class G k of neural networks. Formally, we write ( f 1 , . . . , f K ) ∈ F to denote the congruent potentials constructed this way from the functional classes G 1 , . . . , G K . Hence, in practice, we optimize the empirical version of (8):

A natural question arises: How close are the recovered plans π ̂ f k to the EOT plans π ∗ k between P k and Q ∗ ? Since our objective (8) is a sum of integrals over distributions P k , the generalization error can be straightforwardly decomposed into the estimation and approximation parts.

Proposition 4.4 (Decomposition of the generalization error [proof ref.]) . The following bound holds:

where F C k k def = { f C k k | ( f 1 , . . . , f K ) ∈ F} , and the expectations are taken w.r.t. the random realization of the datasets X 1 ∼ P 1 , . . . , X K ∼ P K . Here Rep X k ( F C k k , P k ) is the standard notion of the representativeness of the sample X k w.r.t. functional class F C k k and distribution P k , see App. A.5.

To bound the estimation error, we need to further bound the expected representativeness E Rep X k ( F C k k , P k ) . Doing preliminary analysis, we found that it does not depend that much on the complexity of the functional class F k . However, it seems to heavily depend on the properties of the cost c k . We derive two bounds: a general one for Lipschitz (in x ) costs and a better one for the feature-based quadratic costs.

Theorem 4.5 (Bound on E Rep w.r.t. C k -transform classes [proof ref.]) . (a) Let F k ⊂ C ( Y ) . Assume that c k ( x, y ) is Lipschitz in x with the same Lipschitz constant for all y ∈Y . Then

(b) Let c k ( x k , y ) = 1 2 ∥ u k ( x k ) -v ( y ) ∥ 2 , F k be a bounded (w.r.t. the supremum norm) subset of C ( Y ) , u k : X k → R D ′′ and v : Y → R D ′′ be continuous functions. Then

Substituting (11) or (12) to (10) immediately provides the statistical consistency when N 1 , . . . , N K →∞ , i.e., vanishing of the estimation error when the sample size grows.

The case (a) here is not very practically useful as the rate suffers from the curse of dimensionality. Still, this result points to one intriguing property of our solver. Namely, we may take arbitrarily large set F k (even F k = C ( Y ) !) and still have the guarantees of learning the barycenter. This happens because of C k -transforms: informally, they make functions f k ∈ F k smoother and "simplify" the set F k . In our experiments, we always work with the costs as in (b) . As a result, our estimation error is O ( ∑ K k =1 N -1 / 2 k ) ; this is a standard fast and dimension-free convergence rate . In practice, F k are usually neural nets. They are indeed bounded, as required in (b) , if their weights are constrained.

While the estimation error 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.]) . Let σ : R → R be an activation function. Assume that it is non-affine and there is an ˜ x ∈ R at which σ is differentiable and σ ′ ( ˜ x ) = 0 . Then for every δ > 0 there exist K multi-layer perceptrons g k : R D → R with activations σ for which the congruent functions f k = g k -∑ K k =1 λ k g k satisfy

Furthermore, each g k has width at-most D +4 .

Importantly, our Theorem 4.6 is more than just result on universal approximation since it deals with (i) congruent potentials and (ii) entropic C k -transforms. In particular, only specific properties of the entropic C k -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 a sufficient amount of data and large neural nets, allowing to perfectly recover the EOT plans π ∗ k .

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 ℓ 2 . It is known that the true unregularized ( ϵ = 0 ) ℓ 2 -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, §5.3) these manifolds are given by pre-trained StyleGAN [50] generator models G : Z → Y ; Z is the latent space, M = G ( Z ) . Technically speaking, to adapt our Alogithm 1 for manifoldconstrained setup, we propose solving the barycenter problem in latent space Z with modified cost functions c k,G ( x k , z ) := c k ( x k , G ( z )) . We emphasize that such costs are general (not ℓ 2 cost!) because G is a 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 1 are not adjustable to the manifold setup as they work exclusively with ℓ 2 .

Relation to prior works. While the utilization of data manifolds given by pre-trained (foundational) models is ubiquitous in generative modeling, the adaptation of this technique for Optimal Transport barycenter is a novel idea. Apart from our work, this idea is exploited in follow-up paper [52].

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 our solver is written in the PyTorch framework and available at https://github.com/justkolesov/EnergyGuidedBarycenters . The experiments are issued in the form of convenient *.ipynb notebooks. Reproducing the most challenging experiments (§5.2, §5.3) requires less than 12 hours on a single TeslaV100 GPU. The details of the experiments, extended experimental results are in Appendix C, additional experiments with single-cell data are given in Appendix C.5.

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 ( ϵ = 0) can be derived (see below). The EOT barycenter for sufficiently small ϵ > 0 is expected to be close to the unregularized one. Therefore, in most cases, our evaluation strategy is to compare the computed EOT barycenter (for small ϵ ) with the unregularized one. 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) The true unregularized barycenter for the twisted cost.

(b) Our EOT barycenter for the twisted cost (map from P 1 ).

(c) The true unregularized barycenter for ℓ 2 cost.

(d) Our EOT barycenter for ℓ 2 cost (map from P 1 ).

6

4

X1~P1

X2~P2

X3~P3

X~Q*

X1~P1

X1~P1

X2~P2

X3~P3

y~Q*

y~n(·|x1)

X1~P1

2

0

-6

-6

4

-2

0

2

4

6

-6

-4

2

0

2

4

6

-6

4

2

0

2

4

6

-6

-4

-2

0

2

4

6

2D Twister . Consider the map u : R 2 → R 2 which, in the polar coordinate system , is represented by R + × [0 , 2 π ) ∋ ( r, θ ) ↦→ ( r, ( θ -r ) mod 2 π ) . The cartesian version of u is presented in Appendix C.1. Let P 1 , P 2 , P 3 be 2 -dimensional distributions as shown in Fig. 2a. For these distributions and uniform weights λ k = 1 3 , the unregularized barycenter ( ϵ = 0 ) for the twisted cost c k ( x k , y ) = 1 2 ∥ u ( x k ) -u ( y ) ∥ 2 can be derived analytically, see Appendix C.1. The barycenter is the centered Gaussian distribution which is also shown in Fig. 2a. We run the experiment for this cost with ϵ = 10 -2 , and the results are recorded in Fig. 2b. We see that it qualitatively coincides with the true barycenter. For completeness, we also show the EOT barycenter computed with our solver for ℓ 2 ( x, y ) = 1 2 ∥ x -y ∥ 2 costs (Fig. 2c) and the same regularization ϵ . The true ℓ 2 barycenter is estimated by using the free_support_barycenter solver from POT package [33]. We stress that the twisted cost barycenter and ℓ 2 barycenter differ, and so do the learned conditional plans: the ℓ 2 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 P n supported on 3D sphere, see Figure 1. The cost functions are c k ( x k , y ) = 1 2 arccos 2 ⟨ x k , y ⟩ , the regularization is ϵ = 10 -2 . The learned potentials f θ,k operate with ambient R 3 vectors. When performing MCMC, we project each Langevin step to the sphere. Our qualitative results are shown on Figure 1. While the ground truth solution to the considered problem is unknown, the learned barycenter looks reasonable. This showcases the applicability of our approach to nonstandard non-quadratic experimental setups.

5.2 Barycenters of MNIST Classes 0 and 1

. A classic experiment considered in the continuous barycenter literature [32, 55, 82, 17] is averaging of distributions of MNIST 0/1 digits with weights ( 1 2 , 1 2 ) in the grayscale image space X 1 = X 2 = Y =[ -1 , 1] 32 × 32 . The true unregularized ( ϵ = 0 ) ℓ 2 -barycenter images y are direct pixel-wise averages x 1 + x 2 2 of pairs of images x 1 and x 2

90

B

D

INPUT

d~tx

X2 ~ P2

TRUE

y=T（x1)

y=T²（x2)

y=T（x3)

y=(x1)

y=f2（x)

y=f3（x3)

WIN

SCWB

y=1(（x1)

(x)=

y~（1x1）y~i（1x1)y~i（1x1)

(x1.)y~

(Ex1-)∈y~(x1-)y~

OURS

y~i2（-1x2)

OURS

y~i2（-1x2)

y~i（·1x3)

OURS

WIN

SCWB

OURS(DataSpace)

WIN

SCWB

OURS(DataSpace)

x~P1y~t（x)y~f（x)

y~（·｜x）

y~（-|x1）

x2~P2y~T2（x2）y~↑2（x2)

y~i2（-｜x2)

y~2（-1x2）

ee9e

1

O

D

(b) Learned plans from

P

2

(ones) to the barycenter.

coming from the ℓ 2 OT plan between 0's ( P 1 ) and 1's ( P 2 ). In Fig. 5, we show the unregularized ℓ 2 barycenter computed by [32, SCWB], [55, WIN].

Data space EOT barycenter. To begin with, we employ our solver to compute the ϵ -regularized EOT ℓ 2 -barycenter 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 are 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-constrained EOT barycenter. Following the reasoning from §4.4, we propose to restrict the search space for our algorithm to some pre-defined manifold M . As discussed earlier, the support of the image-space unregularized ℓ 2 -barycenter is a certain subset of M ′ def = { x 1 + x 2 2 | x 1 ∈ Supp ( P 1 ) , x 2 ∈ Supp ( P 2 ) } . To achieve this, we train a StyleGAN [50] model G : Z → Y with Z = R 512 to generate some even larger manifold M = G ( Z ) which is expected to contain M ′ . Namely, we use all possible pixel-wise half-sums x 1 + x 2 2 of digits 0 as x 1 and { 1 , 4 , 7 } as x 2 , see Figure 3 with the trained StyleGAN samples. That is, our final constructed manifold M is polluted with additional samples (e.g., averages of digits 0 and 7) which should not to lie in the support of the barycenter. Then, we use our solver with ϵ = 10 -2 to search for the barycenter of 0/1 digit distributions on X 1 , X 2 which lies in the latent space Z w.r.t. costs c k,G ( x, z ) def = 1 2 ∥ x -G ( z ) ∥ 2 . This can be interpreted as learning the EOT ℓ 2 -barycenter in the ambient space but constrained to

the StyleGAN-parameterized manifold G ( Z ) . The barycenter Q ∗ is some distribution of the latent variables z , which can be pushed to the manifold G ( Z ) ⊂ 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 unregularized ℓ 2 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 true ℓ 2 barycenter w.r.t. the weights ( 1 4 , 1 2 , 1 4 ) is the distribution 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.4) and train the StyleGAN on unperturbed celeba faces. This might sound a little bit unfair, but our goal is 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 ϵ = 10 -4 . In Fig. 4, we present the results, depicting samples from the learned plans from each P k 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, we 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-constrained setup, the FID score of our solver is significantly smaller.

6 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 to the image manifold by utilizing a pretrained generative model. Our findings may be applicable to a list of important real-world applications, see Appendix B.2. We believe that our large-scale barycenter solver will leverage industrial & socially-important 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 µ f x . In addition, MCMCsampling is time-consuming. We leave the search for more efficient sampling procedures for our solver, e.g., [71, 99, 43, 81, 47, 108, 66, 26], for future research. We also note that our theoretical analysis in §4.3 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-inspired 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.

References

Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis , 43(2):904-924, 2011.

Jason M Altschuler and Enric Boix-Adsera. Wasserstein barycenters are np-hard to compute. SIAM Journal on Mathematics of Data Science , 4(1):179-203, 2022.

Pedro C Álvarez-Esteban, E Del Barrio, JA Cuesta-Albertos, and C Matrán. A fixed-point approach to barycenters in wasserstein space. Journal of Mathematical Analysis and Applications , 441(2):744-762, 2016.

Luigi Ambrosio and Nicola Gigli. A User's Guide to Optimal Transport , pages 1-155. Springer Berlin Heidelberg, Berlin, Heidelberg, 2013.

Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations , 2022.

Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning , pages 146-155. PMLR, 2017.

Christophe Andrieu, Nando De Freitas, Arnaud Doucet, and Michael I Jordan. An introduction to mcmc for machine learning. Machine learning , 50:5-43, 2003.

Arip Asadulaev, Alexander Korotin, Vage Egiazarian, and Evgeny Burnaev. Neural optimal transport with general cost functionals. In The Twelfth International Conference on Learning Representations , 2024.

Julio 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.

Peter L. Bartlett and Shahar Mendelson. Rademacher and Gaussian complexities: risk bounds and structural results. J. Mach. Learn. Res. , 3:463-482, 2002.

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):A1111-A1138, 2015.

Elsa Cazelles, Felipe Tobar, and Joaquin Fontbona. A novel notion of barycenter for probability distributions based on optimal weak mass transport. Advances in Neural Information Processing Systems , 34:13575-13586, 2021.

Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Representations , 2022.

Jinjin Chi, Zhiyao Yang, Ximing Li, Jihong Ouyang, and Renchu Guan. Variational wasserstein barycenters with c-cyclical monotonicity regularization. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 37, pages 7157-7165, 2023.

Lénaïc Chizat. Doubly regularized entropic wasserstein barycenters. arXiv preprint arXiv:2303.11844 , 2023.

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.

Samuel Cohen, Michael Arbel, and Marc Peter Deisenroth. Estimating barycenters of measures in high dimensions. arXiv preprint arXiv:2007.07105 , 2020.

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.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc., 2013.

Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning , pages 685-693. PMLR, 2014.

Marco Cuturi and Gabriel Peyré. A smoothed dual approach for variational wasserstein problems. SIAM Journal on Imaging Sciences , 9(1):320-343, 2016.

Marco Cuturi and Gabriel Peyré. Semidual regularized optimal transport. SIAM Review , 60(4):941-965, 2018.

Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for largescale optimal transport. Advances in neural information processing systems , 34:12955-12965, 2021.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Eustasio del Barrio and Jean-Michel Loubes. The statistical effect of entropic regularization in optimal transportation, 2020.

Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and mcmc. In International Conference on Machine Learning , pages 8489-8510. PMLR, 2023.

Yilun Du, Shuang Li, B. Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. In Proceedings of the 38th International Conference on Machine Learning (ICML-21) , 2021.

Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems , 32, 2019.

J. Dugundji. An extension of Tietze's theorem. Pacific J. Math. , 1:353-367, 1951.

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.

Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research , 2023. Featured Certification.

Jiaojiao Fan, Amirhossein Taghvaei, 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.

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, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research , 22(78):1-8, 2021.

Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energybased models by diffusion recovery likelihood. In International Conference on Learning Representations , 2021.

Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Extremal domain translation with neural optimal transport. In Advances in Neural Information Processing Systems , 2023.

Aude Genevay. Entropy-regularized optimal transport for machine learning . PhD thesis, Paris Sciences et Lettres (ComUE), 2019.

Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems , 29, 2016.

Lee-Ad Gottlieb, Aryeh Kontorovich, and Robert Krauthgamer. Adaptive metric dimensionality reduction. volume 620, pages 105-118. 2016.

Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis , 273(11):3327-3405, 2017.

Hao Guan and Mingxia Liu. Domain adaptation for medical image analysis: a survey. IEEE Transactions on Biomedical Engineering , 69(3):1173-1185, 2021.

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.

Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrey Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the bridge of schr \ " odinger: A continuous entropic optimal transport benchmark. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track , 2023.

Raza Habib and David Barber. Auxiliary variational mcmc. In International Conference on Learning Representations , 2018.

Pierre Henry-Labordere. (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. arXiv preprint arXiv:1904.04546 , 2019.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017.

Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation , 14(8):1771-1800, 2002.

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.03704 , 2019.

Syunsuke Ikeda, Gary Parker, and Kenji Sawai. Bend theory of river meanders. part 1. linear development. Journal of Fluid Mechanics , 112:363-377, 1981.

Hicham Janati, Marco Cuturi, and Alexandre Gramfort. Debiased sinkhorn barycenters. In International Conference on Machine Learning , pages 4692-4701. PMLR, 2020.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4401-4410, 2019.

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.

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.

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 11605, pages 126-133. SPIE, 2021.

Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations , 2021.

Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light schrödinger bridge. In The Twelfth International Conference on Learning Representations .

Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein gans are not optimal transport? Advances in Neural Information Processing Systems , 35:13933-13946, 2022.

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.

Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations , 2021.

Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. In The Eleventh International Conference on Learning Representations , 2023.

Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In The Eleventh International Conference on Learning Representations , 2023.

Anastasis Kratsios and Léonie Papon. Universal approximation theorems for differentiable geometric deep learning. J. Mach. Learn. Res. , 23:Paper No. [196], 73, 2022.

Roman Krawtschenko, César A Uribe, Alexander Gasnikov, and Pavel Dvurechensky. Distributed optimization with quantization for computing wasserstein barycenters. arXiv preprint arXiv:2010.14325 , 2020.

Alexey Kroshnin, Nazarii Tupitsa, Darina Dvinskikh, Pavel Dvurechensky, Alexander Gasnikov, and Cesar Uribe. On the complexity of approximating wasserstein barycenters. In International conference on machine learning , pages 3530-3540. PMLR, 2019.

Rafsanjany Kushol, Alan H Wilman, Sanjay Kalra, and Yee-Hong Yang. Dsmri: Domain shift analyzer for multi-center mri datasets. Diagnostics , 13(18):2947, 2023.

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.

Khang Le, Dung Q Le, Huy Nguyen, Dat Do, Tung Pham, and Nhat Ho. Entropic gromovwasserstein between gaussian distributions. In International Conference on Machine Learning , pages 12164-12203. PMLR, 2022.

Khang Le, Huy Nguyen, Quang M Nguyen, Tung Pham, Hung Bui, and Nhat Ho. On robust optimal transport: Computational complexity and barycenter computation. Advances in Neural Information Processing Systems , 34:21947-21959, 2021.

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energy-based learning. Predicting structured data , 1(0), 2006.

Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems , 34(4):1533-1574, 2013.

Daniel Levy, Matt D. Hoffman, and Jascha Sohl-Dickstein. Generalizing hamiltonian monte carlo with neural networks. In International Conference on Learning Representations , 2018.

Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin M Solomon. Continuous regularized wasserstein barycenters. Advances in Neural Information Processing Systems , 33:1775517765, 2020.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV) , December 2015.

Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning , pages 6672-6681. PMLR, 2020.

Anton Mallasto, Augusto Gerolin, and Hà Quang Minh. Entropy-regularized 2-wasserstein distance between gaussian measures. Information Geometry , 5(1):289-323, 2022.

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.

Alberto Maria Metelli, Amarildo Likmeta, and Marcello Restelli. Propagating uncertainty in reinforcement learning via wasserstein barycenters. Advances in Neural Information Processing Systems , 32, 2019.

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.

Eduardo Fernandes Montesuma and Fred Maurice Ngole Mboula. Wasserstein barycenter for multi-source domain adaptation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 16785-16793, 2021.

Youssef Mroueh. Wasserstein style transfer. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics , volume 108 of Proceedings of Machine Learning Research , pages 842-852. PMLR, 26-28 Aug 2020.

Kirill Neklyudov, Max Welling, Evgenii Egorov, and Dmitry Vetrov. Involutive mcmc: a unifying framework. In International Conference on Machine Learning , pages 7273-7282. PMLR, 2020.

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 on Neural Information Processing Systems , 2023.

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

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.

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.

Gareth O Roberts and Richard L Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli , pages 341-363, 1996.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn 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.

Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations , 2021.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(5863):94, 2015.

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. In International conference on machine learning , pages 30105-30118. PMLR, 2023.

Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Conference on Learning Representations , 2018.

Dino Sejdinovic, Heiko Strathmann, Maria Lomeli Garcia, Christophe Andrieu, and Arthur Gretton. Kernel adaptive metropolis-hastings. In International conference on machine learning , pages 1665-1673. PMLR, 2014.

Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms . Cambridge university press, 2014.

Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

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.

Maurice Sion. On general minimax theorems. Pacific Journal of Mathematics , 8(1):171-176, 1958.

Justin Solomon. Optimal transport on discrete domains. AMS Short Course on Discrete Differential Geometry , 2018.

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) , 34(4):1-11, 2015.

Jiaming Song, Shengjia Zhao, and Stefano Ermon. A-nice-mc: Adversarial training for mcmc. Advances in Neural Information Processing Systems , 30, 2017.

Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint arXiv:2101.03288 , 2021.

Sanvesh Srivastava, Volkan Cevher, Quoc Dinh, and David Dunson. Wasp: Scalable bayes via barycenters of subset posteriors. In Artificial Intelligence and Statistics , pages 912-920. PMLR, 2015.

Sanvesh Srivastava, Cheng Li, and David B Dunson. Scalable bayes via barycenter in wasserstein space. The Journal of Machine Learning Research , 19(1):312-346, 2018.

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.

Heiko Strathmann, Dino Sejdinovic, Samuel Livingstone, Zoltan Szabo, and Arthur Gretton. Gradient-free hamiltonian monte carlo with efficient kernel exponential families. Advances in Neural Information Processing Systems , 28, 2015.

Surya T Tokdar and Robert E Kass. Importance sampling: a review. Wiley Interdisciplinary Reviews: Computational Statistics , 2(1):54-60, 2010.

Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research , 2024. Expert Certification.

Alexander Y Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics , pages 1279-1287. PMLR, 2024.

Konstantin S Turitsyn, Michael Chertkov, and Marija Vucelja. Irreversible monte carlo algorithms for efficient sampling. Physica D: Nonlinear Phenomena , 240(4-5):410-414, 2011.

Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy , 23(9):1134, 2021.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

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.

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-MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13-17, 2019, Proceedings, Part II 22 , pages 623-631. Springer, 2019.

Yang Zhao, Jianwen Xie, and Ping Li. Learning energy-based generative models via coarseto-fine expanding and sampling. In International Conference on Learning Representations , 2021.

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 f 1 , f 2 : Y → R be two measurable functions which are bounded from below. It holds that

(i) Monotonicity : f 1 ≤ f 2 implies f C 1 ≥ f C 2 ;

(ii) Constant additivity : ( f 1 + a ) C = f C 1 -a for all a ∈ R ;

(iii) Concavity : ( λf 1 +(1 -λ ) f 2 ) C ≥ λf C 1 +(1 -λ ) f C 2 for all λ ∈ [0 , 1] ;

(iv) Continuity : f 1 , f 2 bounded implies sup x ∈X | f C 1 ( x ) -f C 2 ( x ) | ≤ sup y ∈Y | f 1 ( y ) -f 2 ( y ) | .

Proof of Proposition A.1. We recall the definition of the C -transform

where C ( x, µ ) def = ∫ Y c ( x, y )d µ ( y ) -ϵH ( µ ) . Monotonicity (i) and constant additivity (ii) are immediate from the definition.

To see (iii), observe that the dependence of ∫ Y f 1 ( y )d µ ( y ) on f 1 is linear. Thus, f C 1 is the pointwise infimum of a family of linear functionals and thus concave.

Finally, to show (iv) we have by monotonicity of the integral that

for any µ ∈ P ( Y ) . For fixed x ∈ X we have

By setting ˜ µ = µ we increase the value and obtain

where the last inequality follows from (13). For symmetry reasons, we can swap the roles of f 1 and f 2 in (14), which yields the claim.

A.2 Proof of Theorem 4.1

Proof. By substituting in (5) the primal EOT problems (1) with their dual counterparts (2), we obtain a dual formulation, which is the starting point of our analysis:

Here, we replaced inf with min because of the existence of the barycenter (§2.2). Moreover, we refer to the entire expression under the minsup as a functional ˜ L : P ( Y ) ×C ( Y ) K → R . For brevity, we introduce, for ( f 1 , . . . , f K ) ∈ C ( Y ) K , the notation

where the equality follows from two elementary observations that (a) M ≤ ∫ ¯ f ( y )d Q ( y ) for any Q ∈ P ( Y ) and (b) ¯ f ( y ) = ∫ ¯ f ( y ′ )d δ y ( y ′ ) where δ y denotes a Dirac mass at y ∈ Y .

On the one hand, P ( Y ) is compact w.r.t. the weak topology because Y is compact, and for fixed potentials ( f 1 , . . . , f K ) ∈ P ( Y ) K we have that ˜ L ( · , ( f k ) K k =1 ) is continuous and linear. In particular, ˜ L ( · , ( f k ) K k =1 ) is convex and l.s.c. On the other hand, for a fixed Q , the functional ˜ L ( Q , · ) is concave by (iii) in Proposition A.1. These observations allow us to apply Sion's minimax theorem [96, Theorem 3.4] to swap min and inf in (15) and obtain using (16)

Next, we show that the sup in (17) can be restricted to tuplets satisfying the congruence condition ∑ K k =1 λ k f k = 0 . It remains to show that for every tuplet ( f 1 , . . . , f K ) ∈ C ( Y ) K there exists a congruent tuplet ( ˜ f 1 , . . . , ˜ f K ) ∈ C ( Y ) K such that ˜ L ( ˜ f 1 , . . . , ˜ f K ) ≥ ˜ L ( f 1 , . . . , f K ) .

To this end, fix ( f 1 , . . . , f K ) and define the congruent tuplet

We find ˜ M def = inf y ∈Y ∑ K k =1 λ k ˜ f k = 0 by the congruence and derive

where the first inequality follows from ˜ f K = f K -¯ f λ K ≤ f K -M λ K combined with monotonicity of the C -transform, see (i) in Proposition A.1. The second to last equality follow from constant additivity, see (ii) in Proposition A.1.

In summary, we obtain where

Finally, observe that for congruent ( f 1 , . . . , f K ) we have ˜ L ( f 1 , . . . , f K ) = L ( f 1 , . . . , f K ) . Hence, we can replace ˜ L by L in (19), which yields (6).

A.3 Proof of Theorem 4.2

Proof. Write Q ∗ for the barycenter and π ∗ k for the optimizer of EOT c k ,ϵ ( P k , Q ∗ ) . Consider congruent potentials f 1 , . . . , f K ∈ C ( Y ) and define the probability distribution

Then we have by [78, Thm. 2]:

Multiplying (22) by λ k and summing over k yields

where the last equality follows by congruence, i.e., ∑ K k =1 λ k f k ≡ 0 .

The remaining inequality in (7) is a consequence of the data processing inequality for f -divergences which we invoke here to get

where Q ∗ and Q f k are the second marginals of π ∗ k and π f k , respectively.

A.4 Proof of Theorem 4.3

Proof. The desired equation (9) 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 Theorem 4.5 relies on the following standard definitions from learning theory, which we now recall for convenience (see, e.g. [93, §26]). Consider some class S of functions s : X → R and a distribution µ on X . Let X = { x 1 , . . . , x N } be a sample of N points in X .

The representativeness of the sample X w.r.t. the class S and the distribution µ is defined by

The Rademacher complexity of the class S w.r.t. the distribution µ and sample size N is given by

where { x n } N n =1 ∼ µ are mutually independent, { σ n } N n =1 are mutually independent Rademacher random variables, i.e., Prob ( σ n = 1 ) = Prob ( σ n = -1 ) = 0 . 5 , and the expectation is taken with respect to both { x n } N n =1 , { σ n } N n =1 . The well-celebrated relation between (25) and (24), as shown in [93, Lemma 26.2], is

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

Let f be a maximizer of L ( f ) within F . Analysing the estimation error in (27) yields

where central term in line (28) is bounded above by 0 due the maximality of ̂ f , that is, ̂ L ( ̂ f ) = max f ∈F ̂ L ( f ) ≥ ̂ L ( f ) . Due to (29), we can bound the estimation error using the Rademacher complexity

Proof of Theorem 4.5. Case (a) - Lipschitz costs : Assume that, for k ∈ K , x ↦→ c k ( x, y ) is Lipschitz with constant L k ≥ 0 for every y ∈ Y . Recall that f C k k is defined as the pointwise supremum of L k -Lipschitz functions and, therefore, Lipschitz continuous with the same constant. Since the value of the representativeness of a sample w.r.t. a function class is invariant under translating individual elements of said class, we have that Rep X ( F C k k , P k ) coincides with Rep X ( G k , P k ) where

for some fixed ˜ x k ∈ X k . All the functions in this class are L k -Lipschitz and, therefore, bounded by L k · diam ( X k ) . We may therefore apply [38, Theorem 4.3] to the class G k and obtain

Case (b) - Feature-based quadratic costs 1 2 ∥ u k ( · ) -v ( · ) ∥ 2 2 : Alternatively, consider the case where c k ( x k , y ) = 1 2 ∥ u k ( x k ) -v ( y ) ∥ 2 2 and F k ⊆ C ( Y ) is bounded (w.r.t. the supremum norm).

To this end, recall that for a measurable and bounded function f : Y → R , the weak entropic c k -transform satisfies

where Z c k ( f, x k ) = ∫ Y exp ( f ( y ) -c k ( x k ,y ) ϵ ) dy . Recall that R D ′′ × R D ′′ ∋ ( a, b ) ↦→ exp ( -∥ a -b ∥ 2 2 ϵ ) is a positive definite kernel which is widely known as the Gaussian kernel . This means that there exists a Hilbert space H and a feature map ϕ : R D ′′ →H such that exp ( -∥ a -b ∥ 2 2 ϵ ) = ⟨ ϕ ( a ) , ϕ ( b ) ⟩ H . Due to the particular form of c k , we may write

Notice that { ϕ ( u k ( x k )) : x k ∈ X k } ⊆ { v ∈ H : ∥ v ∥ H = 1 } since ∥ ϕ ( a ) ∥ 2 H = ⟨ ϕ ( a ) , ϕ ( a ) ⟩ H = 1 for every a ∈ R D ′′ .

Using the identity in (30), we can express Z c k ( f, x k ) by

where the last equality is justified as Y is compact; furthermore, we note the integrals are well-defined by the measurability of ϕ , u k and v . Moreover, using the boundedness of each F k and compactness of Y , we get

Define G k def = { Z c k ( f, · ) : f ∈ F k } and observe that G k ⊆ G ′ k def = {⟨ ϕ ( u k ( · ) ) , w ⟩ H : ∥ w ∥ H ≤ R } . This implies that R N k ( G k , P k ) ≤ R N k ( G ′ k , P k ) by the properties of the Rademacher complexity.

In turn, the latter Rademacher complexity can be bounded by [10, Lemma 22]. Indeed, write P ′ k def = ( u k ) # P k such that

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

Since the functions in G k are bounded uniformly away from zero by some constant κ > 0 (depending on the bound of F k and ϵ ), and since the logarithm restricted to [ κ, ∞ ) is 1 /κ -Lipschitz, we have that

We can now bound the expected representativeness of F C k k with the Rademacher complexity by (26), yielding the claim.

A.6 Proof of Theorem 4.6

Proof of Theorem 4.6. Let σ ∈ C ( R ) be a non-affine activation function which is differentiable at some point x 0 ∈ R and for which σ ′ ( x 0 ) = 0 . Let δ > 0 , λ 1 , . . . , λ K > 0 , and K ∈ N be given. By Theorem 4.1, there exist K congruent continuous functions f ′ 1 , . . . , f ′ K such that

̸

Applying [29, Theorem 4.1] we deduce that for each k = 1 , . . . , K , there exist a continuous extension f ′ k : R D → R for ˜ f k to all of R D ; i.e. f ′ k ( y ) = ˜ f k ( y ) for each y ∈ Y . In particular, (32) can be rewritten as

̸

Set δ k def = δ/ (4 λ k ) for each k = 1 , . . . , K . Since Y is a compact subset of R D , and since we have assumed that σ ∈ C ( R ) is non-affine activation function which is differentiable at some point ˜ x ∈ R and for which σ ′ ( ˜ x ) = 0 then the special case of [62, Theorem 9] given in [62, Proposition 53] , implies that for any there exist feedforward neural networks g k : R D k → R ( k ∈ K ) with activation function σ , such that

each g k width at-most D +4 . Pick δ k = δ 4 for all k ∈ K and suitable neural networks g 1 , . . . , g K . Next, we define the congruent sums of neural networks f k def = g k -∑ K k =1 λ k g k . We derive

Using (34) we obtain for fixed k ∈ K

By (iv) in Proposition A.1 together with (35) we find

Now we use (36) to derive

Next we combine (33) with (37) to get

By using (38) together with Theorem 4.2 we obtain

which completes the proof.

A.7 Existence and uniqueness of the barycenter distribution which solves (5)

We introduce an auxiliary functional which is the argument of minimization problem (5):

i.e. the optimal value of (5) could be defined as L ∗ = inf Q ∈P ( Y ) B ( Q ) .

Note that the functional Q ↦→B ( Q ) is strictly convex and lower semicontinuous (w.r.t. the weak topology) as each component Q ↦→ EOT c k ,ϵ ( 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 µ ↦→ E y ∼ µ c k ( x, y ) -H ( µ ) is l.s.c, bounded from below and strictly convex thanks to the entropy term. Since P ( Y ) is weakly compact (as Y is compact due to Prokhorov'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 at least one Q for which B ( Q ) < ∞ . In this case, the barycenter Q ∗ is unique due to the strict convexity of B .

B Extended discussions

B.1 Extended discussion of related works

Discrete OT-based 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, 11, 21, 22, 30, 63, 49, 68]. An alternative formulation of the barycenter problem based on weak mass transport and corresponding discrete solver could be found in [12]. In spite of sound theoretical foundations and established convergence guarantees [64], 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, ϵ = 0 ) case. For example, solvers such as [44, 74, 54, 58, 57, 32, 35, 88, 5, 31], are based on optimizing the dual form, similar to our (2), with neural networks. We mention these methods because they serve as the basis for certain continuous unregularized barycenter solvers. For example, ideas of [55] are employed in the barycenter method presented in [59]; the solver from [74] is applied in [32]; max-min solver introduced in [58] is used

in [55]. It is also worth noting that there exist several neural solvers that cater to more general OT problem formulations [61, 60, 31, 8, 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 to recover the optimal EOT plan π ∗ in EOT problems like (1) between unknown distributions P and Q which are only accessible through a limited number of samples. One group of methods [37, 91, 23] is based on the dual formulation of OT problem regularized with KL divergences [37, Eq. ( P ϵ )] which is equivalent to (1). Another group of methods [109, 24, 13, 41, 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 to tackle (1) by means of Energy-Based models [69, 100, EBM]. 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 µ f ∗ x = π ∗ ( ·| x ) . 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 unfixed 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 a 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 as a 2 D 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

It is not a secret that despite considerable efforts in developing continuous barycenter solvers [72, 59, 55, 32, 82, 14], these solvers have not found yet a real working practical application. The reasons for this are two-fold:

Existing continuous barycenter solvers (Table 1) are yet not scalable enough and/or work exclusively with the quadratic cost ( ℓ 2 ), which might be not sufficient for the practical needs.

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 a 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 c k 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-specific knowledge that necessitates involving experts in the field. Meanwhile, the manifold constraint is

Source 1

Source 2

C2

Petrel

Flumy

Source3

?

(a) Domain averaging of MRI scans' sources.

(b) Mixing geological simulators.

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 is reasonable to expect that soon the new large models targeted for medical data may appear, analogously to DALL-E [85], StableDiffusion [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.

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 can take into account several desired geological factors which are successfully modeled by separate simulators.

Flumy 1 is a process-based simulator that uses hydraulic theory [48] to model specific channel depositional processes returning a 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.

Petrel 2 software is the other popular simulator in the oil and gas industry. It is able to model complex real-valued 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 are hard to use in conjunction. Even when conditioned on common prior information about the deposit, they may produce maps of facies and permeability which do not meaningfully correspond to 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.

From our personal discussions with the experts in the field of geology, such task formulations are of considerable interest both for scientific community as well as industry. Applying our barycenter solver in this context is a challenge for future research. We acknowledge that this would also require overcoming considerable technical and domain-specific issues, including the data collection and the choice of costs c k .

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 λ ≥ 0 and τ ≥ 0 we define:

The corresponding doubly -regularized EOT barycenter problem is as follows:

̸

It turns out that our considered EOT barycenter (5) with regularization strength ϵ is the particular case of (40) with λ = τ = ϵ . According to the classification [15, Table 1], this case is known as Schrödinger barycenter. The natural question arizes: Is it possible to adapt our energy-guided methodology for the case λ = τ ?

̸

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 λ = τ , this entropy term immediately reappears (either with the plus or minus sign). The presence of H ( 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 π k ( y | x k ) . Namely, we use potentials f k combined with costs c k 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 ( π k ( y | x k )) appearing in the C k -transforms.

If we further add the non-zero term H ( Q ) to the barycenter objective, this will presumably require (in the dual objective) a separate computation of the entropy terms H ( π k ( y )) of second marginals π k ( y ) of each π k . This is highly non-trivial and it seems like our solver does not easily generalize to this case . In our framework, we can get samples of π k ( y ) by MCMC (via sampling x k ∼ P k and then running MCMC for π ( y | x k ) ). However, estimation of entropy of π k ( y ) from raw samples still remains infeasible. In particular, EBM-like techniques (which we employ) can not be used here to

derive the gradient of the objective. This is because the required unnormalized density of π k ( y ) is itself unknown (we only know it for conditional distributions π k ( y | x ) ).

C Experimental Details

The hyperparameters of our solver are summarized in Table 4. Some hyperparameters, e.g., L, S, iter , are chosen primarily from time complexity reasons. Typically, the increase in these numbers positively affects the quality of the recovered solutions, see, e.g., [42, Appendix E, Table 16]. However, to reduce the computational burden, we report the figures which we found to be reasonable. Working with the manifold-constraint setup, we parameterize each g θ k ( z ) in our solver as h θ k ◦ G ( z ) , where G is a pre-trained (frozen) StyleGAN and h θ k is a neural network with the ResNet architecture. We empirically found that this strategy provides better results than a direct MLP parameterization for the function g θ k ( z ) .

√

To train the StyleGAN for MNSIT01 & Ave, celeba! experiments, we employ the official code from

https://github.com/NVlabs/stylegan2-ada-pytorch

Computational complexity. We report the (approximate) time for training and inference (batch size S = 128 ) of our method on different experimental setups, see Table 3 (the hardware is a single V100 gpu). For Ave Celeba experiment, we additionally report the computational complexity of the competitors. As we can see, all the methods in this experiment (Ave Celeba) require a comparable amount of time for training. The inference with our approach is expectedly costlier than competitors due to the reliance on MCMC.

C.1 Barycenters of 2D/3D Distributions

Cartesian representation of twister map. In §5.1 we define twister map u using polar coordinate system. For clearness, we give its form on cartesian coordinates. Let x ∈ R 2 , x = ( x (1) , x (2) ) . Note that ∥ x ∥ = √ x 2 (1) + x 2 (2) Then,

i.e., the twister map rotates input points x by angles equal to ∥ x ∥ .

Analytical barycenter distribution for 2D Twister experiment. Below, we provide a mathematical derivation that the true unregularized barycenter of the distributions P 1 , P 2 , P 3 in Fig. 2a coincides with a Gaussian.

We begin with a rather general observation. Consider X k = Y = R D ( k ∈ K ) and let OT c def = EOT c, 0 denote the unregularized OT problem ( ϵ = 0 ) for a given continuous cost function c . Let u : R D → R D be a measurable bijection and consider P ′ k ∈ P ( R D ) for k ∈ K . By using the change of variables formula, we have for all Q ′ ∈ P ( R D ) that

where # denotes the pushforward operator of distributions and [ c ◦ ( u × u )]( x, y ) = c ( u ( x ) , u ( y ) ) . Note that by (41) the barycenter of P ′ 1 , . . . , P ′ K for the unregularized problem with cost c coincides with the result of applying the pushforward operator u -1 # to the barycenter of the very same problem but with cost c ◦ ( u × u ) .

Next, we fix u to be the twister map (§5.1). In Fig. 2a we plot the distributions P 1 def = u -1 # P ′ 1 , u -1 # P ′ 1 , u -1 # P ′ 3 which are obtained from Gaussian distributions P ′ 1 = N ( (0 , 4) , I 2 ) , P ′ 2 = N ( ( -2 , 2 √ 3) , I 2 ) , P ′ 3 = N ( (2 , 2 √ 3) , I 2 ) by the pushforward. Here I 2 is the 2-dimensional identity matrix. For the unregularized ℓ 2 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 , I 2 ) . Hence, the barycenter of P 1 , . . . , P K w.r.t. the cost ℓ 2 ◦ ( u × u ) is given by Q ∗ = u -1 # Q ′ . From the particular choice of u it is not hard to see that Q ∗ = Q ′ = N ( 0 , I 2 ) 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], we run their publicly available code from the official repository

The authors do no provide checkpoints, and we train their barycenter model from scratch. In turn, for the solver by [55, WIN], we also employ the publicly available code

Here we do not train their models but just use the checkpoint available in their repo.

C.3 Barycenters of the Ave, Celeba! Dataset

Additional qualitative examples of our solver's results are given in Figure 8.

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 M 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 10, 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 we solve the entropy-regularized 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-constrained setup.

Convergence at training. We provide the behaviour of L 2 -UVP metric (between the unregularized ground truth barycenter mapping and the entropic barycenter mapping for our learned π f k ( y | x k ) ) for Ave Celeba experiment §5.3, see Figure 9 . We emphasize that L 2 -UVP directly measures (by computing pointwise MSE) how the learned mapping differs from the true mapping.

X1~P1

X2~P2

(x1-)2（x1-)2(2x1)（x1.)y（2x12(x1-)（x1)2(x1-)（x1-)(2x1)

y~n（-1x3)

y~(1x1)y(1x1)

y~i3（-1x3)

y~i3（-1x3)

(x1-)y~(1x1)Ty~

y~i（-1x3)

y~ig（-1x3)

y~i（1x1)

y~i（-1x3)

(x1)~

y~is（-1x3)

y~i（-1x1)

y~i（-1x3)

y~i（1x1)y~i（1x1)

y~i（-1x3)

(Ex1.)~

>

100

100

100

L2-UVP of P_2

80

80

80

%

60

60

60

40

40

40

20

20

20

0

20

40

t.hours

60

80

100

20

40

t.hours

60

80

100

0

20

40

t.hours

60

80

100

Convergence at inference. Since Langevin Dynamics is 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 of L Langevin steps in Manifold constrained as well as in 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 on 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, WIN] 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 a high-dimensional setup with RGB images. Hence, to implement their approach in this setting, we follow [55, Appendix B.4].

(xo,z)

(x，2)

(x2，2)

C.4 Barycenters of Gaussian Distributions

We note that there exist many ways to incorporate the entropic regularization for barycenters [15, Table 1]; these problems do not coincide and yield different solutions. For some of them, the ground-truth solutions are known for specific cases, such as the Gaussian case. For example, [75, Theorem 3] examines barycenters for OT regularized with KL divergence. They consider the task

This problem differs from our objective (5) with c ( x, y ) = 1 2 ∥ x -y ∥ 2 by the non-constant Q -dependent term ϵH ( Q ) ; this problem yields a different solution. The difference with the majority of other mentioned approaches can be shown in the same way. In particular, [67] tackles the barycenter for inner product Gromov-Wasserstein problem with entropic regularization, which is not relevant for us. In turn, the authors of [49] demonstrate significant progress towards GT results for barycenter problem with our considered regularization, but only for univariate (1D) case. To our knowledge, the general Gaussian ground-truth solution for our problem setup (5) is not yet known, although some of its properties seem to be established [25].

Still, when ϵ ≈ 0 , our entropy-regularized barycenter is expected to be close to the unregularized one ( ϵ = 0 ). In the Gaussian case, it is known that the unregularized OT barycenter for c k ( x, y ) = 1 2 ∥ x -y ∥ 2 is itself Gaussian and can be computed using the well-celebrated fixed point iterations of [3, Eq. (19)]. This gives us an 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 learns the unregularized barycenter ( ϵ = 0 ).

(x)²~x

We consider 3 Gaussian distributions P 1 , P 2 , P 3 in dimensions D = 2 , 4 , 8 , 16 , 64 and compute the approximate EOT barycenters π ̂ f k k for ϵ = 0 . 01 , 1 w.r.t. weights ( λ 1 , λ 2 , λ 3 ) = ( 1 4 , 1 4 , 1 2 ) with our solver. To initialize these distributions, we 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 L 2 -UVP ( ˆ T ) = 100 · [ ∥ ˆ T -T ∗ ∥ ] 2 P where T ∗ denotes the optimal transport map T ∗ , see [54, §5.1]. Since our solver computes EOT plans but not maps, we evaluate the barycentric projections of the learned plans, i.e., ̂ T k ( x ) = ∫ Y yπ ̂ f k k ( y | x ) , and calculate L 2 -UVP ( ̂ T k , T ∗ k ) . We evaluate this metric using 10 4 samples. To estimate the barycentric projection in our solver, we use 10 3 samples y ∼ π ̂ f k k ( y | x k ) for each x k . To keep the table with the results simple, in each case we report the average of this metric for k = 1 , 2 , 3 w.r.t. the weights λ k .

We see that for small ϵ = 0 . 01 and dimension up to D = 16 , our algorithm gives the results even better than WIN solver designed specifically for the unregularized case ( ϵ = 0 ). As was expected, larger ϵ = 1 leads to the increased bias in the solutions of our algorithm and L 2 -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,ϵ ) = (2 , 0 . 1) , ( D,ϵ ) = (16 , 0 . 01) . As the batch sizes, we consider 2 , 4 , 8 , 16 , 32 , 64 , 128 , 256 , 512 , 1024 . As we can see in Table 9, the quality of the recovered plans π f k 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, on the other hand, provides reasonable computational burden. In the image experiments, we use batch size ≤ 128 as it 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 is affected by chosen number of discretized Langevin dynamic steps ( L from §4.2) at training, the results are presented in Figure 11. In this experiment, we try to learn the barycenter of Gaussian distributions; we pick the dimensionality D = 64 (which is the highestdimensional case among the considered); ϵ = 10 -1 . 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.

C.5 Single-cell experiment

We consider the problem of predicting the interpolation between single cell populations at multiple timepoints from [107]. The objective is to interpolate

6

5

L =32

C2-UVP

4

L =48

L = 64

3

L=128

L=256

2

20

40

60

80

100

iterx103

the distribution of cell population at any intermediate point in time, call it t i , given 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 ℓ 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 [56], wherein the authors provide a complete

notebook with the code to launch the setup similar to [107]. There are 3 experimental settings with dimensions D = 50 , 100 , 1000 , each setting contains two setups: predicting day 3 given days 2 and 4 and predicting day 4 given days 3 and 7. The metric is MMD; see [107] or Section 5.3 from [56] for additional experimental details. We report the result in Table 8 where we find that in most cases, our general-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 act as a viable foundation model, directly out-of-the-box.

D Alternative EBM training procedure

(a) The true unregularized

barycenter for the twisted cost.

(b)

Our

EOT barycenter for the

twisted cost (map from

P

1

).

(c) The true unregularized

barycenter for

ℓ

2

cost.

(d)

Our

EOT barycenter for

ℓ

2

cost (map from

P

1

).

6

4

X1~ P1

X2~P2

X3~P3

X~Q.

X1~P1

X1~P1

X2~P2

X3~P3

y~Q*

X1~P1

2

0

2

-4

-6

-6

4

2

0

2

4

6

-6

-4

-2

0

2

4

6

-6

-4

-2

0

2

4

6

-6

-4

-2

0

2

4

6

In this section, we 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.3). To overcome the difficulty, in the main part of our manuscript, we utilize MCMC sampling from conditional distributions µ f θ,k x k and estimate the loss with Monte-Carlo. Here we discuss a potential alternative approach based on importance sampling (IS) [105]. That is, we evaluate the internal integral over Y in (4.3):

with help of an auxiliary proposal (continuous) distribution accessible by samples with the known density q ( y ) . Let Y q = { y q 1 , . . . , y q P } be a sample from q ( y ) . Define the weights:

Then (42) permits the following stochastic estimate:

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 16 I -covariance Gaussian distribution as q and pick the batch size P = 1024 . Our results are shown in Figure 12. As we can see, the alternative training procedure yields similar results as Figure2 but converges faster ( ≈ 1 min. VS ≈ 18 min. of the original MCMC-based training).

Concluding remarks. We note that IS-based methods requires accurate selection of the proposal distribution q to reduce the variance of the estimator [105]. It may be challenging in real-world scenarios. We leave the detailed study of more advanced IS approaches in the context of energy-based models and their applicability to our EOT barycentric setup to follow-up research.

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, we completely describe our contributions. For every contribution, we provide a link to the section about it.

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 set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Proofs are provided in Appendix A.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it 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 at 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification:Code is publicly available at a github repository. Experimental details are provided in 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, hyperparameters, 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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Not all of the experiments are conducted following the statistical significance protocol due to computational restrictions.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In Appendix C, we mention time and resources.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Research conforms with NeurIPS Code of Ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss broader impact in Appendix 6.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The research does not need special safeguards.

12. Licenses for 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 self-documented and follows our experimental protocol from the manuscript. The license is 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.

15. 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nFirst, 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.\n\nWEAKNESSES:\n1. 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\").\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nYes, the authors discuss the limitations of their work in a convincing way (in the appendix; ).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nPlease see above.\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe method is theoretically well-founded and has many applications. I particularly found interesting that it can work for multiple different ground cost functions.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe authors address the limitations of their work in the manuscript.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\nPlease 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\n\nLIMITATIONS:\nThough 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", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nIt 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.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-24"}