Energy-Guided Continuous Entropic Barycenter Estimation for General Costs

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 17:22:19

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.

Main paper: 
 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 (Villani et al., 2009) . They have clear geometrical meaning and practical interpretation (Santambrogio, 2015; Solomon, 2018) . The corresponding problem of averaging probability distributions using OT discrepancies is known as the OT barycenter problem (Agueh & Carlier, 2011) . OT-based barycenters find application in various practical domains: domain adaptation (Montesuma & Mboula, 2021b;a), shape interpolation (Solomon et al., 2015) , Bayesian inference (Srivastava et al., 2015; 2018) , text scoring (Colombo et al., 2021) , style transfer (Mroueh, 2020) , reinforcement learning (Metelli et al., 2019) .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 (Li et al., 2020; Cohen et al., 2020; Korotin et al., 2021c; 2022a; Fan et al., 2021; Noble et al., 2023; Chi et al., 2023) . Most of these works are devoted to specific OT cost functions, e.g., deal with ℓ 2 2 barycenters (Korotin et al., 2021c; 2022a; Fan et al., 2021; Noble et al., 2023) ; while others require non-trivial a priori selections (Li et al., 2020) and have limiting expressivity and generative ability (Cohen et al., 2020; Chi et al., 2023) , see §3 for a detailed discussion.Contribution. 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 (EBMs) ( §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 dP(x) dx . 
background: 
 First, we recall the formulations of EOT ( §2.1) and the barycenter problem ( §2.2). Subsequently, we clarify the computational setup of the considered EOT barycenter problem ( §2.3).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 (Cuturi, 2013; Peyré et al., 2019; Genevay, 2019) consists of finding a minimizer of EOT (2) c,ε (P, Q) def = min π∈Π(P,Q) X ×Y c(x, y)dπ(x, y) -εH(π),where H(π) is the differential entropy of plan π. The case ε = 0 corresponds to classical OT, also known as the Kantorovich problem (Kantorovich, 1942) , and falls of the scope of this paper.Since by the chain rule of entropy we have H(π) = X H(π(•|x))dP(x) + H(P), where π(•|x) are conditional distributions on Y, equation (1) permits the following equivalent reformulation:EQUATIONA minimizer π * ∈ Π(P, Q) of ( 2) is called the EOT plan. Thanks to the equivalence of ( 1) and (2), its existence and uniqueness are guaranteed, see, e.g., (Clason et al., 2021, 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 (Gushchin et al., 2023b, §2) . We refer to π * (•|x) as the conditional plans (for x ∈ X ).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, see (Gozlan et al., 2017, Theorem 9.5) or (Backhoff-Veraguas et al., 2019, Theorem 1.3) :EQUATIONwhere f C : X → R is the so-called weak entropic c-transform (Backhoff-Veraguas et al., 2019, Eq. 1. 2) of the function (potential) f . The transform is defined byEQUATIONWe use the capital C in f C to distinguish the weak transform from the classic c-transform (Santambrogio, 2015, §1.6) or (c, ϵ)-transform (Marino & Gerolin, 2020, §2). In particular, formulation (3) is not to be confused with the conventional dual, see (Mokrov et al., 2023, Appendix A) .For each x ∈ X , the minimizer µ f x ∈ P(Y) of the weak c-transform (4) exists and is unique. Its density is of the form (Mokrov et al., 2023, Theorem 1) : Set Z c (f, x) def = Y exp f (y)-c(x,y) ε dy, dµ f x (y) dy def = 1 Z c (f, x) exp f (y) -c(x, y) ε .(5)By substituting ( 5) into (4) and carrying out straightforward mathematical manipulations, we arrive at an explicit formula f C (x) = -ε log Z c (f, x), see (Mokrov et al., 2023, Equation (14) ).Not only does maximizing the dual objective in (3) allow us to estimate the actual value of the primal objective EOT c,ε (P, Q), but it also provides an approximation of the EOT plan π * . Consider the distribution dπ f (x, y) = dµ f x (y)dP(x). According to (Mokrov et al., 2023, Thm. 2), we haveEQUATION)This means that the smaller the error in solving the dual problem with f ∈ C(Y), the closer the distribution π f to π * . This useful property appears only when ϵ > 0, which is one of many reasons why EOT (ϵ > 0) is often favoured over the unregularized problem (ϵ = 0).Consider distributions P k ∈ P ac (X k ) and continuous cost functions (Cuturi & Doucet, 2014; Cuturi & Peyré, 2018; Dvurechenskii et al., 2018; del Barrio & Loubes, 2020) , (Le et al., 2021; 2022) c k (•, •) : X k × Y → R for k ∈ K. Given positive weights λ k > 0 with K k=1 λ k = 1, the EOT Barycenter problemEQUATIONThe case where ε = 0 corresponds to the classical OT barycenter problem (Agueh & Carlier, 2011) and falls of the scope of this paper. By substituting EOT (2) c k ,ε , with EOT c k ,ε we deriveEQUATIONwhich differs from (7) only by the additive constant K k=1 λ k H(P k ) and has the same minimizers. It is worth noting 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 (Backhoff-Veraguas et al., 2019, Th. 2.9 ) by noting that on P(Y) the map µ → Y c k (x, y)dµ(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., (Santambrogio, 2015, Box 1.4 )), it holds that B(Q) admits at least one minimizer due to the Weierstrass theorem (Santambrogio, 2015, 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 as consequence of the strict convexity ofB.Barycenter problems, such as ( 7) or (8), are known to be challenging in practice (Altschuler & Boix-Adsera, 2022) . 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 Appendix C.4, nor for the unregularized case ( Álvarez-Esteban et al., 2016, ϵ = 0) . Furthermore, 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 following learning setup.We assume that each P k is accessible only by a limited number of i.i.d. empirical samplesX 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 (8). 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 (Li et al., 2020; Korotin et al., 2021c) . It differs from the discrete OT setup (Cuturi, 2013; Cuturi & Doucet, 2014) which aims to solve the barycenter problem for empirical distributions, e.g.,P k def = 1 N k N k n=1 δ x nk . Discrete methods are not well-suited for the out-of-sample estimation required in the continuous OT setup.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 EOT problem or (E-)OT barycenter problem. These methods try to 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 and continuous OT solvers is left to Appendix B.1. Continuous EOT solvers aim to recover the optimal EOT plan π * in (2), (1) between unknown distributions P and Q which are only accessible through a limited number of samples. One group of methods (Genevay et al., 2016; Seguy et al., 2018; Daniels et al., 2021) is based on the dual formulation of OT problem regularized with KL divergences (Genevay et al., 2016, Eq . (P ε )) which is equivalent to (2). Another group of methods (Vargas et al., 2021; De Bortoli et al., 2021; Chen et al., 2021; Gushchin et al., 2023a; Tong et al., 2023; Shi et al., 2023) takes advantage of the dynamic reformulation of (1) via Schrödinger bridges (Léonard, 2013; Marino & Gerolin, 2020) . In (Mokrov et al., 2023) , the authors propose an approach to tackle (2) by means of Energy-Based models (LeCun et al., 2006; Song & Kingma, 2021, EBM) . They develop an optimization procedure resembling standard EBM training which retrieves the optimal dual potential f * appearing in (3). As a byproduct, they recover the optimal conditional plans µ f * x = π * (•|x). Our approach for solving the EOT barycenter ( 8) 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. The works (Korotin et al., 2021c; Fan et al., 2021; Noble et al., 2023; Korotin et al., 2022a) 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 (Ambrosio & Gigli, 2013, §2) which are exploited by the aforementioned articles to build efficient procedures for barycenter estimation algorithms. In particular, (Korotin et al., 2021c; Fan et al., 2021) utilize ICNNs (Amos et al., 2017) which parameterize convex functions, and (Noble et al., 2023) 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 (Li et al., 2020) , 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 (Chi et al., 2023) 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 (Cohen et al., 2020) , 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 proposed approach overcomes many imperfections which are inherent in competing methods. 
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.In what follows, we analyze (8) from the dual perspectives. We introduce L : C(Y) K → R:L(f 1 , . . . , f K ) def = K k=1 λ k X k f C k k (x k )dP k (x k ) = -ϵ K k=1 λ k X k log Z c k (f k , x k )dP k (x k ) .Here f C k k (x k ) denotes the weak entropic c k -transform (4) 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 probablity measures on Y. Theorem 1 (Dual formulation of the EOT barycenter problem). The EOT barycenter problem (8) permits the following dual formulation: infEQUATIONWe refer to the constraintK k=1 λ k f k = 0 as the congruence condition w.r.t. weights {λ k } K k=1 .The potentials f k appearing in (9) play the same role as in (3). Notably, whenL(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 2 below. First, for f k ∈ C(Y), we define dπ f k (x k , y) def = dµ f k x k (y)dP k (x k ) and dQ f k (y) def = X k dπ f k (x k , y).Theorem 2 (Quality bound of plans recovered from dual potentials).EQUATIONwhereπ * k ∈ Π(P k , Q * ), k ∈ K are the EOT plans between P k and the barycenter distribution Q * .According to Theorem 2, an approximate solution {f k } K k=1 of (9) 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. ( 5), with closed-form energy function, i.e., the unnormalized log-likelihood. Consequently, one can generate samples from µ f kx k using standard MCMC techniques (Andrieu et al., 2003) . In the next subsection, we stick to the practical aspects of optimization of (9), which bears certain similarities to the training of Energy-Based models (LeCun et al., 2006; Song & Kingma, 2021, EBM) .Relation to prior works. Works (Li et al., 2020; Korotin et al., 2021c ) also aim to first get the dual potentials and then recover the barycenter, see the detailed discussion in §3.To maximize the dual EOT barycenter objective (9), we replace the potentialsf k ∈ C(Y) for k ∈ K with neural networks f θ,k , θ ∈ Θ. In order to eliminate the constraint in (9), we propose parametriz- ing f θ,k as g θ k - K k ′ =1 λ k ′ g θ k ′ , where {g θ k : R D → R, θ k ∈ Θ k } Kk=1 are neural networks. This parameterization automatically ensures the congruence conditionK k=1 λ k f θ,k ≡ 0. Note that Θ = Θ 1 × • • • × Θ K and θ = (θ 1 , . . . , θ K ) ∈ Θ.Our objective function for ( 9) is defined asEQUATION)While direct computation of the normalizing constant Z c k may, in general, be infeasible, the gradient of L with respect to θ can be derived similarly to (Mokrov et al., 2023, Theorem 3 ): Theorem 3 (Gradient of the dual EOT barycenter objective). The gradient of L satisfiesEQUATION)With this result, we can describe our proposed algorithm which maximizes L using (12).TRAINING. To perform stochastic gradient ascent step w.r.t. θ, we approximate (12) with Monte-Carlo by drawing samples from dπ f θ,k (x k , y) = dµ Algorithm 1: Entropic OT barycenters via Energy-Based ModellingInput : Source distributions P k , k ∈ K accessible by samples; cost functions c k (x k , y) : X k × Y → R; the regularization coefficient ε > 0; barycenter averaging coefficients λ k > 0 such that K k=1 λ k = 1; MCMC procedure MCMC procedure; batch size S > 0; potential NNs f θ,k : Y → R, such that 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 * . for iter = 1, 2, . . .do for k = 1, 2, . . . , K do Sample batch {x s k } S s=1 ∼ P k ; Draw Y k = {y s k } S s=1 with MCMC: y s k = MCMC procedure f θ,k (•)-c(x s k ,•) ε ; L k ← -λ k 1 S S s=1 f θ,k (y s k ) ; L ← K k=1 L k ; Perform a gradient step over θ by using ∂ L ∂θ ;X k . Then, we need to draw a sample from the distribution µ f θ,kx k . Since we know the negative energy (unnormalized log density) of µf θ,kx k by ( 5), 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.In all our experiments, we use ULA (Roberts & Tweedie, 1996, §1.4 .1) as a MCMC procedure. It is the simplest MCMC algorithm. Specifically, in order to draw a sampley k ∼ µ f θ,k x k , where x k ∈ X k , we initialize y (0)k from the standard Normal D-dimensional distribution N (0, I D ) and then iterate the discretized Langevin dynamics:y (l+1) k ← y (l) k + η 2ε ∇ y f θ,k (y) -c(x k , y) y=y (l) k + √ ηξ l , ξ l ∼ N (0, I D ),where 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 samplesY (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 (12) is presented in Appendix D. INFERENCE. At the inference stage, 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 (LeCun et al., 2006; Xie et al., 2016; Du & Mordatch, 2019; Song & Kingma, 2021) . 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., (Du et al., 2021; Gao et al., 2021; Zhao et al., 2021) . 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 MCMC algorithm, even without the replay buffer (Hinton, 2002) , as it serves our current objectives.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 samplesX 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 havef 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 (11):max (f1,...,f K )∈F L(f 1 , . . . , f K ) def = max (f1,...,f K )∈F K k=1 λ k N k N k n=1 f C k k (x n k ) (13) and recover ( f 1 , . . . f K ) def = arg max (f1,...,f K )∈F L(f 1 , . . . , f k ).A natural question arises: How close are the recovered plans π f k to the EOT plans π * k between P k and Q * ? Since the objective ( 11) is a sum of integrals over distributions P k , we can derive generic finite learning guarantees by using Theorem 2, see the next theorem below. Theorem 4 (Finite sample learning guarantees). The following generalization bound holds true:EQUATION)whereF C k k def = {f C k k | (f 1 , . . . , f K ) ∈ F}, and the expectation is taken w.r.t. the random realization of the datasetsX 1 ∼P 1 , . . . , X K ∼P K . Here R N k (F C k , P k ) is the Rademacher complexity (Shalev- Shwartz & Ben-David, 2014, §26) of the functional class F C k k w.r.t. samples of size N k from P k .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 a neural network parameterization is used for the potentials (as we do in §4.2). Theorem 5 (Universal approximation with neural networks). For every δ > 0 there exist K neural networks g k : R D k → R such that the congruent functionsf k = g k - K k=1 λ k g k satisfy ϵ K k=1 λ k KL π * k ∥π f k = L * -L(f 1 , . . . , f K ) < δ.Relation to prior works. To our knowledge, the generalization and the universal approximation are novel results with no analogs established for any other continuous barycenter solver. Our analysis shows that the EOT barycenter objective ( 11) is well-suited for statistical learning and approximation theory tools. This aspect distinguishes our work from the preceding works, where more complex optimization objectives may not be as amenable to rigorous study. 
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 available in the supplementary material and written in the PyTorch framework. The code will be made publicly available. 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 and extended experimental results are in Appendix C. 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 barycenter. Therefore, in most cases, our evaluation strategy is to compare the computed EOT barycenter (for small ϵ) with the unregularized one. In particular, we follow this strategy to quantitatively evaluate our solver in the Gaussian case, see Appendix C.4.2D Twister. Consider a map u : R 2 → R 2 which, in the polar coordinate system, is represented by R + ×[0, 2π) ∋ (r, θ) → r, (θ-r) mod 2π . Let P 1 , P 2 , P 3 be 2-dimensional distributions as shown in Fig. 1a . 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. 1a . We run the experiment for this cost with ϵ = 10 -2 , and the results are recorded in Fig. 1b . 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. 1c ) and the same regularization parameter ϵ. The true ℓ 2 barycenter is estimated by using the free support barycenter solver from POT package (Flamary et al., 2021) . We stress that the twisted cost barycenter and ℓ 2 barycenter differ, and so do the learned conditional plans. To be precise, the ℓ 2 EOT plan (Fig. 1d ) expectedly looks more well-structured while for the twisted cost (Fig. 1b ) it becomes more chaotic due to non-trivial structure of this cost. A classic experiment considered in the continuous barycenter literature (Fan et al., 2021; Korotin et al., 2022a; Noble et al., 2023; Cohen et al., 2020) 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 x1+x2 2 of pairs of images x 1 and x 2 coming from the ℓ 2 OT plan between 0's (P 1 ) and 1's (P 2 ). In Fig. 2 , we show the unregularized ℓ 2 barycenter computed by (Fan et al., 2021, SCWB) , (Korotin et al., 2022a, 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. 2 , we see that our solver (data space) recovers the noisy barycenter images exactly as expected.Manifold-constrained EOT barycenter. Averaging image distributions directly in the data space can be challenging. Our experiment below shows that if we a priori have some manifold M where we want the barycenter to be concentrated on, our solver can restrict the search space to it.As discussed earlier, the support of the image-space unregularized ℓ 2 -barycenter is a certain subset of M = { x1+x2 2 | x 1 ∈ Supp(P 1 ), x 2 ∈ Supp(P 1 )}. To achieve this, we train a StyleGAN model G : Z → Y (Karras et al., 2019) with Z = R 512 to generate an approximate manifold G(Z) ≈ M. 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 (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 StyleGANparameterized manifold G(Z). In this case, the barycenter Q * is the distribution of the latent variables z, which can be pushed to the manifold G(Z) ⊂ Y via G(z). The results are also in Fig. 2 . There is no noise compared to the data-space EOT barycenter because of the manifold constraint.We emphasize that the costs c k (x k , z) used here are general (not ℓ 2 cost!) because G is a non-trivial StyleGAN generator. Hence, this experimental setup with the manifold-constrained barycenters is not feasible for most other barycenter solvers as they work exclusively with ℓ 2 . In (Korotin et al., 2022a) , 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 (Liu et al., 2015) , 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.FID↓ of plans to the barycenter k = 1 k = 2 k = 3 (Fan et al., 2021) 56.7 53.2 58.8 (Korotin et al., 2022a) 49.3 46.9 61.5 Ours 8.4 8.7 10.2Table 2 : FID scores of images mapped from inputs P k to the barycenter.We follow the EOT manifold-constrained setup as in §5.2 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. 3 , 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 (Fan et al., 2021, SCWB) , (Korotin et al., 2022a, WIN) . Additionally, we report the FID score (Heusel et al., 2017) for images mapped to the barycenter in Table 2 . Owing to the manifold-constrained setup, the FID score of our solver is significantly smaller. 
discussion: 
 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.Limitations. The 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 µ fx . In addition, MCMC sampling is time-consuming. We leave the search of more efficient sampling procedures for our solver, e.g., (Levy et al., 2017; Song et al., 2017; Habib & Barber, 2018; Neklyudov et al., 2020; Hoffman et al., 2019; Turitsyn et al., 2011; Lawson et al., 2019; Du et al., 2023) , for future research.We start by showing some basic properties of the C-transform which will be used in the main proofs.Proposition 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 1. We recall the definition of the C-transformf C 1 (x) = inf µ∈P(Y) C(x, µ) - Y f 1 (y)dµ(y) ,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 thatEQUATIONfor any µ ∈ P(Y). For fixed x ∈ X we havef C 1 (x) -f C 2 (x) = inf µ∈P(Y) C(x, μ) - Y f 1 (y)dμ(y) -inf µ∈P(Y) C(x, µ) - Y f 2 (y)dµ(y) = sup µ∈P(Y) inf μ∈P(Y) C(x, μ) -C(x, µ) - Y f 1 (y)dμ(y) + Y f 2 (y)dµ(y) .By setting μ = µ we increase the value and obtainEQUATIONwhere the last inequality follows from (15). For symmetry reasons, we can swap the roles of f 1 and f 2 in ( 16), which yields the claim.Proof. By substituting in (8) the primal EOT problems (2) with their dual counterparts (3), we obtain a dual formulation, which is the starting point of our analysis:EQUATION)Here, we replaced inf with min because of the existence of the barycenter ( §2.2). Moreover, we refer to the entire expression under the min sup as a functional L : P(Y) × C(Y) K → R. For brevity, we introduce, for (f 1 , . . . , f K ) ∈ C(Y) K , the notationEQUATIONwhere the equality follows from two elementary observations that (a) M ≤ f (y)dQ(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 1. These observations allow us to apply Sion's minimax theorem (Sion, 1958, Theorem 3.4) to swap min and inf in ( 17) and obtain using ( 18)EQUATION)Next, we show that the sup in ( 19) can be restricted to tuplets satisfying the congruence conditionK 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 ( f1 , . . . , fK ) ∈ C(Y) K such that L( f1 , . . . , fK ) ≥ L(f 1 , . . . , f K ).To this end, fix (f 1 , . . . , f K ) and define the congruent tupletEQUATION)We find M def = inf y∈Y K k=1 λ k fk = 0 by the congruence and deriveL( f1 , . . . , fK ) -L(f 1 , . . . , f K ) = λ K X K f C K K (x K ) -f C K K (x K ) dP K (x K ) -M ≥ λ K X K f K - M λ K C K (x K ) -f C K K (x K ) dP(x K ) -M = λ K X K M λ K dP(x K ) -M = 0,where the first inequality follows from fK = f K -f λ K ≤ f K -M λ K combined with monotonicity of the C-transform, see (i) in Proposition 1. The second to last equality follow from constant additivity, see (ii) in Proposition 1.In summary, we obtainEQUATIONFinally, 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 (21), which yields (9).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 distributiondπ f k (x k , y) def = dµ f k x k (y) dP k (x k ),wheredµ f k x k (y) dy def = 1 Z c k (f k , x k ) exp f k (y) -c k (x k , y) ϵ , Z c k (f k , x k ) def = log Y e f k (y)-c k (x k ,y) ϵdy .Then we have by ( 6)EQUATIONMultiplying ( 22) by λ k and summing over k yieldsϵ K k=1 λ k KL π * k ∥π f k = k=1 λ k EOT c k ,ϵ (P k , Q * ) - X k f C k k (x k )dP k (x k ) - Y K k=1 λ k f (y) =0 dQ * (y) = L * -L(f 1 , . . . , f k ),where the last equality follows by congruence, i.e., K k=1 λ k f k ≡ 0. The remaining inequality in ( 10) is a consequence of the data processing inequality for f -divergences which we invoke here to getKL π * k ∥π f k ≥ KL Q * ∥Q f k ,where Q * and Q f k are the second marginals of π * k and π f k , respectively.A.4 PROOF OF THEOREM 3Proof. The desired equation ( 12) could be derived exactly the same way as it is done in (Mokrov et al., 2023, Theorem 3) .Before proving Theorem 4, we recall the required basic concepts of statistical learning theory (Shalev-Shwartz & Ben-David, 2014, §26) . Consider some class H of functions h : 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 H and the distribution P is defined byEQUATIONThe Rademacher complexity of the class H w.r.t. the distribution P and sample size N is given byEQUATIONwhere {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 ( 24) and ( 23) isEQUATIONwhere the expectation is taken w.r.t. random i.i.d. sample X ∼ µ of size N .A.5 PROOF OF THEOREM 4Proof.EQUATIONLet f be an arg max f∈F L(f). We conduct the standard analysis for the estimation error in (26):EQUATIONTo upper bound the central term in transition to line ( 27), we use maximality of f, that is, L( f) = max f∈F L(f) ≥ L(f). We plug (28) into (26), take the expectation and obtainϵE K k=1 λ k KL π * k ∥π f k ≤ L * -max f∈F L(f) + 2E K k=1 λ k Rep X k (F C k k , P n ) ≤ L * -max f∈F L(f) + 4E K k=1 λ k R N k (F C k , P k ),where the last inequality is the usual Rademacher bound (25).A.6 PROOF OF THEOREM 5Proof.EQUATION)These are continuous functions defined on a compact subset of R D k . From the general approximation theory it follows that for every δ 1 , . . . , δ K > 0 there exist neural networksg k : R D k → R (k ∈ K) which satisfy ∥g k -f ′ k ∥ ∞ def = sup y∈Y |g k (y) -f ′ k (y)| < δ k .For this, we may use (Kidger & Lyons, 2020, Theorem 3.2) or any other known neural universal approximation theorem. Pick δ k = δ 4 for all k ∈ K and suitable neural networks g 1 , . . . , g K . Next, we define the congruent sums of neural networksEQUATIONUsing ( 30) we obtain for fixed k ∈ KEQUATION)By (iv) in Proposition 1 together with (31) we findEQUATION)Now we use (32) to deriveEQUATION)Next we combine ( 29) with ( 33) to getEQUATIONBy using (34) together with Theorem 2 we obtainϵ K k=1 λ k KL π * k ∥π f k = L * -L(f 1 , . . . , f K ) < δwhich finishes the proof.Discrete OT-based solvers provide solutions to OT-related problems between discrete distributions. A comprehensive overview can be found in (Peyré et al., 2019) . The discrete OT methods for EOT barycenter estimation are (Cuturi & Doucet, 2014; Solomon et al., 2015; Benamou et al., 2015; Cuturi & Peyré, 2016; Cuturi & Peyré, 2018; Dvurechenskii et al., 2018; Krawtschenko et al., 2020) , (Le et al., 2021) . In spite of sound theoretical foundations and established convergence guarantees (Kroshnin et al., 2019) , 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 (Henry-Labordere, 2019; Makkuva et al., 2020; Korotin et al., 2021a; b; 2022b; Fan et al., 2021; Gazdieva et al., 2023; Rout et al., 2021; Amos, 2022) , are based on optimizing the dual form, similar to our (3), with neural networks. We mention these methods because they serve as the basis for certain continuous unregularized barycenter solvers. For example, ideas of (Korotin et al., 2022a) are employed in the barycenter method presented in (Korotin et al., 2021c) ; the solver from (Makkuva et al., 2020) is applied in (Fan et al., 2021) ; max-min solver introduced in (Korotin et al., 2021b) is used in (Korotin et al., 2022a) . It is also worth noting that there exist several neural solvers that cater to more general OT problem formulations (Korotin et al., 2023b; a; Asadulaev et al., 2022) . These can even be adapted to the EOT case (Gushchin et al., 2023b) but require substantial technical effort and the usage of restrictive neural architectures.Other related works. Another relevant work is (Simon & Aberdam, 2020) , where the authors study the barycenter problem and restrict the search space to a manifold produced by a GAN. This idea is Figure 4 : A schematical presentation of potential applications of barycenter solvers. 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 2D distribution via its intensity histogram and search for a single image on the GAN manifold whose density is the barycenter of the input images. To compute the barycenter, they use discrete OT solver. In summary, neither our barycenter solver is intended to be used in their setup, nor their method is targeted to solve the problems considered in our paper.It is not a secret that despite considerable efforts in developing continuous barycenter solvers (Li et al., 2020; Korotin et al., 2021c; 2022a; Fan et al., 2021; Noble et al., 2023; Chi et al., 2023) , these solvers have not found yet a real working practical application. The reasons for this are two-fold:1. 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. 2. Potential applications of barycenter solvers are too technical and, unfortunately, require substantial efforts (challenging and costly data collection, non-trivial design of task-specific cost functions, unclear performance metrics, etc.) to be implemented in practice.Despite these challenges, there exist rather inspiring practical problem formulations where the continuous barycenter solvers may potentially shine and we name a few below. These potential applications motivate the research in the area. More generally, we hope that our developed solver could be a step towards applying continuous barycenters to practical tasks that benefit humanity.1. Solving domain shift problems in medicine (Fig. 4a ). 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 (Guan & Liu, 2021; Kushol et al., 2023; Kondrateva et al., 2021; Stacke et al., 2020; Yan et al., 2019) . 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 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 (Ramesh et al., 2022) , StableDiffusion (Rombach et al., 2022) or StyleGAN-T (Sauer et al., 2023) for general image synthesis. These future models could potentially parameterize the medical data manifolds of interest, opening new possibilities for medical data analysis.2. Mixing geological simulators (Fig. 4b ). 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.Flumy1 is a process-based simulator that uses hydraulic theory (Ikeda et al., 1981) 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.Petrel2 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 .The hyper-parameters of our solver are summarized in Table 3 . Working with the manifoldconstraint setup, we parameterize each g θ k (z) in our sover 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). In this subsection we provide a mathematical derivation that the true unregularized barycenter of the distributions P 1 , P 2 , P 3 in Fig. 1a coincides with a Gaussian. This problem differs from our objective (7) with c(x, y) = 1 2 ∥x -y∥ 2 by the non-constant Qdependent term ϵH(Q); this problem yields a different solution. The difference of other mentioned approaches can be shown in the same way. In particular, (Le et al., 2022) tackles the barycenter for inner product Gromov-Wasserstein problem with entropic regularization which is not relevant for us. To our knowledge, the Gaussian ground-truth solution for our problem setup (7) is not yet known, although some of its properties seem to be established (del Barrio & Loubes, 2020) .EQUATIONStill 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 ( Álvarez-Esteban et al., 2016, Eq. (19) ). This gives us an opportunity to compare our results. 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.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 (3). To overcome the difficulty, in the main part of our manuscript, we utilize MCMC sampling from conditional distributions µ f θ,kx k and estimate the loss with Monte-Carlo. Here we discuss a potential alternative approach based on importance sampling (IS) (Tokdar & Kass, 2010) . That is, we evaluate the internal integral over Y in (3):EQUATIONwith 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:ω k (x k , y q p ) def = exp f θ,k (y q p ) -c(x k , y q p ) ε q(y q p ).Then (36) permits the following stochastic estimate: Concluding remarks. We note that IS-based methods requires accurate selection of the proposal distribution q to reduce the variance of the estimator (Tokdar & Kass, 2010). 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.f θ,kx k (y)dP k (x k ). Analogously to(Mokrov et al., 2023, §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(a) Maps from P1 to the barycenter. (b) Maps from P2 to the barycenter. (c) Maps from P3 to the barycenter.https://flumy.minesparis.psl.euhttps://www.software.slb.com/products/petrel 
 
 Appendix:  


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

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

Reviews: review: paper_summary: paper_summary: This paper presents a new entropically regularized algorithm for identifying the Wasserstein barycenter (distribution on average closest to the reference distributions as measured in Wasserstein distance) among K distributions. 

The 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). 

The authors proceed to provide a theoretical analysis of the proposed algorithm, providing KL and universal approximation bounds. 

They then follow with a number of experimental demonstrations of the method.
strength_weakness: strength_weakness: - 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. 
- 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.

Some small remarks:

- Line 208: ULA is surely not the simplest MCMC algorithm. 
- personal complaint: the notation feels a bit heavyhanded and can be distracting from the main story of the paper.
questions: questions: It is often the case in energy-based modeling that evaluating the loss (12) can be very high variance because of the necessity of the MCMC procedure and can have negative effect on the performance of the algorithm. Do the authors notice any such difficulty in doing so? How in general do you see benchmarking performance, which is usually important when an algorithm relies on stochastic gradient descent and a neural network, rather than stricter algorithms usually associated to EOT (like sinkhorn). I don't mean to be critical here, just reflective. The toy distributions example is nice.
limitations: limitations: n/a
paper_summary: paper_summary: The paper proposes a methodology to approximate (regularised) barycenters of continuous-support distributions for arbitrary transport costs. To do so, the authors combine strategies and results from weak OT and transport map approximation via neural networks. The proposal is well connected with the existing (ever-growing) literature on barycentric approximation and provides a set of theoretical results as well as numerical examples of real-world data.
strength_weakness: strength_weakness: - The use of the Style GAN seems to be pivotal in the experimental validation, yet it is only explained briefly while reporting the experiments. 

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

- There are some formatting issues: 
  - line 255: _Substituting (13) or (14) to (13) ..._
  - lines 301-302 have spacing problems
  - the top of fig 5 is cropped 
  - many formatting problems with the references
questions: questions: Please refer to my comments on computational complexity and the role of the Style GAN. 

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

Please elaborate on this point and, if possible, include it in the paper
limitations: limitations: Though the paper features a section called limitations, the authors only state that the one limitation of their method was alleviated by using Style GANs. A deeper identification and discussion of the limitations of the proposed method is lacking
paper_summary: paper_summary: The authors focus on the Wasserstein barycenter problem; an average between distributions given a Wasserstein cost. Specifically, they consider arbitrary cost functions to define the Wasserstein distance and they approximate the continuous barycenters. They further relate their method with energy based models, and show experiments and synthetic and image datasets.
strength_weakness: strength_weakness: - 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. 

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

- 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]).

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

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

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

[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.
questions: questions: - What is $\mathcal{B}(\mathbb{Q})$ in eq. 8 ?
- Section 2.2 could be shortened, I believe we only really need to know that the barycenter is unique if it exists. 

## Minor comments and potential typos
- Maybe a missing word on line 86 "has particular".
- "falls of the scope" $\to$ falls within the scope of the paper.
- The authors should refer to the proofs of the theorems in the main body of the text, even if the proof is in appendix.
- In Tab.2, could you add the standard deviation and the names of the other two solvers.
limitations: limitations: The authors address the limitations of their work in the manuscript.
paper_summary: paper_summary: The authors propose a new algorithm to estimate entropic barycenters in this submission. 
In 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. 
The 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.
Compared to existing OT-based barycenter estimation methods, the proposed algorithm can be applied to general cost functions and can explicitly estimate OT plans.
strength_weakness: strength_weakness: - 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.
- 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.
- 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?

Minors:
- It would be nice if the authors could apply the proposed method in a practical task, e.g., point cloud clustering.
- For the completeness of the main paper, a conclusion section should be added.
questions: questions: Please see above.
limitations: limitations: Not applicable.
paper_summary: paper_summary: This paper deals with the question of estimating Wasserstein(-like) barycenters in the _continuous_ case based on _samples_. That is, one actually observes (iid) samples $X^k$ made of $N_k$ points from (unknown) probability distributions $\mathbb{P}_1,\dots,\mathbb{P}_K$, and the goal is to estimate $\mathbb{Q}^*$ the barycenters of the $(\mathbb{P}_k)_k$ (and not simply the discrete barycenter of the $X^k$). 

The 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}^*$. 
Of interest, $\pi^{f_k}$ admits a disintegration and its second marginal can be sampled by the following scheme :
- sample $x_k$ from $\mathbb{P}_k$ (i.e. take an observation in the training set),
- 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. 
The 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. 

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

The approach is showcases on toy and images datasets.
strength_weakness: strength_weakness: 1. The main weakness (and for which I unfortunately do not have a specific solution) of the approach is the difficulty to properly evaluate it (note: this is not specific to the method but rather to the problem considered), due to:
- the actual (regularized, or not) barycenter $\mathbb{Q}^*$ being usually unknown,
- the value of the loss $L$ (a fortiori the optimal value $L^*$ and the error gap) cannot be computed,
- 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.
It 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. 

2. 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). 

3. 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").
questions: questions: 1. The approach relies on a Langevin (or other sampling) scheme with potential energy given by $f_k - c(x_k, \dot)$. Given that $f_k$ is encoded by a neural network which is typically neither convex not smooth in the space variable $y$ (nor in the parameters $\theta$ but that does not matter here), it's dubious to assume that the Langevin scheme converge toward the target distribution. This is only very briefly mentioned in lines 363-364 and in line 1051 and I believe this should be discussed more extensively (i.e. did you notice it during your experiments? how impactful is it?)

2. 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?

3. As far as I can tell, the barycenter that is targeted here similar (if not exactly the same) as the $\alpha_{\mathrm{OT}^\mathcal{U}}$ of Janati et al., _Debiased Sinkhorn barycenters_ (note: this work should be cited/discussed at some point). is that correct? If so, note that in that case an  _closed form_ expression for the entropic barycenter of (univariate) Gaussian distributions is accessible (showcasing formally the blurring effect); and section C.4. may be revisited under this lens.
limitations: limitations: Yes, the authors discuss the limitations of their work in a convincing way (in the appendix; ).

Comments: title: Acknowledge of rebuttal
Thanks
De-anonymization
Response to Rebuttal
Thanks
comment: Dear Authors, 

many thanks for the reply. The contribution of the paper is much clearer to me know. I appreciate the clarification about computational complexity.
Thank you for your response, I appreciate the clarifications and additional experiments. It is now clearer to me how this work differs from [1].
One final remark -- going forward, please don't reference the arxiv version of your paper, which totally de-anonymizes you. You did this in the response to reviewer  unyK. While I understand they were asking you to compare your method to another paper, and it may be hard to do that without explicit referencing, it's important to not break the "double blind" reviewing criteria.
Thanks for clarifying the nuance to Theorem 4.6. This point is quite subtle and it may be worth making a statement explicitly in the text that this more than just a result of universal approximation.

With regards to remarks about NP-hardness -- I'm not asking you to comment on this, I'm saying that in general it is hard to measure if you've succeeded in solving your barycenter problem. This is a general remark about the circumstances where you'd actually like to discover a barycenter using your algorithm, and contextualizing the experiments done in that regard. 

While it is interesting and good to see that proper hyperparameter tuning can make the training procedure successful, the reviewer definitely thinks the community reading the paper, who may want to further research it, understand some of these challenges. Indeed, more practical means of reducing the variance of the loss b/c of this reliance on MCMC would definitely be interesting, and contextualizing if this is a real challenge would be quite valuable :). 

Thanks for your detailed response. I will raise my score a point!
Thanks for taking time answering my review and providing clarification / details on the few questions I had. So far, I am satisfied with the work, I'll try to engage in discussion with other reviewers to make my mind clearer.