Optimal Transport Barycenter via Nonconvex-Concave Minimax Optimization
Authors: Kaheon Kim
Rentian Yao
Changbo Zhu
Xiaohui Chen
Venue: ICML 2025 poster (2025)
Keywords: Wasserstein barycenter; optimal transport; nonconvex-concave minimax optimization; convergence rate;
Primary area: optimization->nonconvex

Abstract: The optimal transport barycenter (a.k.a. Wasserstein barycenter) is a fundamental notion of averaging that extends from the Euclidean space to the Wasserstein space of probability distributions. Computation of the *unregularized* barycenter for discretized probability distributions on point clouds is a challenging task when the domain dimension $d > 1$. Most practical algorithms for approximating the barycenter problem are based on entropic regularization. In this paper, we introduce a nearly linear time $O(m \log{m})$ and linear space complexity $O(m)$ primal-dual algorithm, the *Wasserstein-Descent $\dot{\mathbb{H}}^1$-Ascent* (WDHA) algorithm, for computing the *exact* barycenter when the input probability density functions are discretized on an $m$-point grid. The key success of the WDHA algorithm hinges on alternating between two different yet closely related Wasserstein and Sobolev optimization geometries for the primal barycenter and dual Kantorovich potential subproblems. Under reasonable assumptions, we establish the convergence rate and iteration complexity of WDHA to its stationary point when the step size is appropriately chosen. Superior computational efficacy, scalability, and accuracy over the existing Sinkhorn-type algorithms are demonstrated on high-resolution (e.g., $1024 \times 1024$ images) 2D synthetic and real data.

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
Kaheon Kim
Rentian Yao?
Changbo Zhu
Xiaohui Chen
Abstract
measures. Direct applications of the Wasserstein barycen- ter include smooth interpolation between shapes (Solomon et al., 2015), texture mixing (Rabin et al., 2012), and averag- ing of neuroimaging data (Gramfort et al., 2015), among oth- ers. More importantly, the computation of the Wasserstein barycenter often serves as a key stepping stone to derive more advanced machine learning and statistical algorithms. For instance, centroid basddd methods for clustering distri- butions rely the computation of Wasserstein barycenters (Cuturi & Doucet, 2014; Zhuang et al., 2022). Addition- ally, regression models and statistical inference methods for distributional data that utilize Wasserstein geometry often employ the barycenter as an "anchor" measure to map dis- tributions to a linear tangent space (Dubey & Müller, 2020; Zhang et al., 2022; Chen et al., 2023; Zhu & Müller, 2023; Zhu & Müller, 2023; Jiang al., 2024).
The optimal transport barycenter (a.k.a. Wasser- stein barycenter) is a fundamental notion of av- eraging that extends from the Euclidean space to the Wasserstein space of probability distribu- tions. Computation of the unregularized barycen- ter for discretized probability distributions on point clouds is a challenging task when the do- main dimension d > 1. Most practical algorithms for the barycenter problem are based on entropic regularization. this paper, we introduce nearly linear time (m log m) and linear space complex- ity (m)) primal-dual algorithm, the Wasserstein- Descent H -Ascent (WDHA) algorithm, for com- puting the exact barycenter when the input proba- bility density functions are discretized on an m- point grid. The key success of the WDHA algo- rithm hinges on alternating between two different yet closely related Wasserstein and Sobolev opti- mization geometries for the primal barycenter and dual Kantorovich potential subproblems. Under reasonable assumptions, we establish the conver- gence rate and iteration complexity of WDHA to its stationary point when the step size is ap- propriately chosen. Superior computational effi- cacy, scalability, and accuracy over the existing Sinkhorn-type algorithms are demonstrated on high-resolution (c.g., 1024 1024 images) 2D synthetic and real data.
Despite the widespread applications, computationally effi- cient or even scalable algorithms of the Wasserstein barycen- ter with theoretical guarantees remain to be developed. Ex- isting approaches to compute the Wasserstein barycenter of a collection of probability density functions j4,,.... .....HH I' in Rd rely on Wasserstein analog of the gradient descent algorithm, which requires to compute n OT maps per itera- tion (Zemel & Panaretos, 2019; Chewi et al., 2020). Álvarez- Esteban et al. (2016) propose fixed point approach that is effective for any location-scatter family. Method that resorts to using first-order methods on the dual for exact barycenter calculation (Carlier et al., 2015), and classical OT solvers, such as Hungarian method (Kuhn, 1955), auction algorithm (Bertsekas & Castanon, 1989) and transportation simplex (Luenberger & Ye, 2008), scale poorly for even moderately mesh-sized problems. This presents a substan- tial computational barrier for computing the barycenter of multivariate distributions. On the other hand, various reg- ularized barycenters have been proposed to mitigate the computational difficulty (Li et al., 2020a; Janati et al., 2020; Bigot et al., 2019; Carlier et al., 2021; Chizat, 2023), and Sinkhorn's algorithm is perhaps one of the most widely used algorithms to compute the entropy-regularize, barycen- ter (Peyré & Cuturi. 2019; Lin et al., 2022; Carlier, 2022). Also, Li et al. (2020b) propose a new dual formulation for the regularized Wasserstein barycenter problem such that discretizing the support is not needed. However, the computation thrifty of these methods is subject the ap-
1. Introduction
The Wasserstein barycenter, introduced by Agueh & Carlier (2011) 1) based on the theory of optimal transport (OT), ex- tends the notion of a Euclidean average to measure-value data, thus representing the "average" of a set of probability
Department of ACMS, University of Notre Dame, IN, USA. 2 Department of Mathematics, University of British Columbia, BC, Canada. 'Department of Mathematics, University of South- em Califomia, CA, USA. Correspondence to: Changbo Zhu <czhu4 nd.cu>>>
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Optimal Transport Barycenter via Nonconvex-Coneaee Minimax Optimization
proximation accuracy trade-off (Nenna & Pegon, 2024). In low-dimensional settings such 2D images, this is often manifested in visually undesirable blurring effects the barycentric images. To break the curse of dimensionality, al- gorithms using input convex neural networks (Korotin et al., 2021) and generative models (Korotin et al., 2022) have been investigated for the Wasserstein barycenter problem.
in the Euclidean nonnonvee--ttonnyyyoooneee optimiza- tion problems.
We demonstrate superior numeric accuracy and compu- tational efficacy over Sinkhorn-tye algorithms on high- resolution 2D synthetic and real image data, where the standard Wasserstein gradient descent algorithm cannot be practically implemented on such problem size.
In this paper, we recast the unregularized optimal transport (a.k.a. Wasserstein) barycenter problem nonconvex- concave minimax optimization problem and propose a co- ordinate gradient algorithm, termed as the Wasserstein- Descens H -Ascens (WDHA), which alternates between the Wasserstein and Sobolev spaces. The key innovation of our WDHA algorithm to combine two different yet closely related primal-dual optimization geometries between the primal subproblem for updating barycenter with the Wasser- stein gradient and the Kantorovich dual formulation of the Wasserstein distance for updating the potential functions with a homogeneous lil gradient (cf. Definition in Sec- tion 2.2). In contrast with the usual L gradient V. our choice of the H gradient can be interpreted as an isometric dual embedding of the potential function & corresponding to the earthmoving effort for pushing source distribution j to a target distribution via || " f |
For limitations, the current approach mainly limited to computing the Wasserstein barycenter of 2D 3D distribu- tions supported on a compact domain.
Notations. We use 12 HI, and (5()) to represent the d-dimensional Euclidean space, a Hilbert space, and the Wasserstein space on Q respectively. Let || |/2 and (:-) al - ll and (+-))a denote the Euclidean (Hilbert) norm and inner-product. Given function i Rd + R and functionals I H + R, F: Pj - R, we use V.p, VI, and W.F to represent the standard gradient on IRd, the H- gradient, and the Wasserstein gradient respectively. Let 8e := sup,,(..) p. and 4 denote the second convex conjugate. We use id to represent the identity map and the notation [T] 21, TT. Given two probability measures and μu, let т` denote the optimal transport map that pushes μu, and let Φe be the corresponding Kantorovich potential. A more detailed list of notations is provided in the appendix.
ppy) denote the convex conjugate of
This OT perspective is critical to ensure stability for our H- gradient ascent subproblem and therefore the convergence of the overall WDHA algorithm.
2. Preliminary
Contributions. Current work is among the first to combine the Wasserstein and homogeneous Sobolev gradients de- rive a simple, scalable, and accurate primal-dual coordinate gradient algorithm for computing the exact (i.e., unregular- ized) OT barycenter. The proposed WDHA algorithm is particularly suitable for computing the barycenter for dis- cretized probability density functions on a large m-point grid such as images of 1024 x 1024 resolution. The pro- posed WDHA algorithm enjoys strong theoretical properties and empirical performance. The following summarizes our main contributions.
2.1. Monge and Kantorovich Optimal Transport Problems
Let P2(2)) be the set of probability measures convex compact set Q C Rd with finite second-orde, moments, i.c., it holds that Ja | du((z) < oo for any μ E P2)))). For v, la E (92), the Monge's optimal transport (OT) problem for the quadratic cost can be written as
(V, = inf M(T), T:TaN=H
where
Diseretizing the input probability density functions Mi...., Mn onto grid of m points, the per iteration run- time complexity of our algorithm is O(m log m) for up- dating each Kantorovich potential. This is in sharp con- trast with the time cost )mm²) for computing an OT map between two distributions supported on the same grid via linear programming (LP). In addition, the space complex- ity for the Hil gradient is (m), which also substantially reduces the LP space complexity Omm?).
M(T) = f 12((() x112 de(x)
and T is the push-forward measure of 1 by T. W2 p1) is called the 2-Wasserstein distance between and μ. Though the solution of Monge's problem may not exist, its relaxation, the Kantorovich formulation of the optimal transport problem shown below, always admits solution,
min I((A:= = Jaxa - v12d)(z,8)) xE
/ 1
Under reasonable assumptions, we provide an explicit algorithmic convergence rate and iteration complexity of the WDHA algorithm to its stationary point with appro- priately chosen step sizes for the gradient updates. In particular, WDHA achieves the same convergence rate as
where I[v, 14) is the set of probability measures on Ω x Ω with marginal distributions v and j. The optimal solution A is called the optimal transport plan. When v E PT()),
2

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
the subset of (2t) consisting of all absolutely continuous probability measures (with respect to the Lebesgue measure on S2), it is known that the solution Т of Monge's problem exists, and the optimal transport plan is À (id, TV } #V, In this work, we will utilize the following dual form of the Kantorovich'' s problem,
Σ : HI + IR is concave functional. With the definition of H gradient, the following THT -gradient ascent algorithm (Algorithm 2.2) can be applied to solve max, IE (1P), where 4 represents the convex conjugate of yp
Algorithm HH' Graddient Ascent Algorithm
Initialize yp!;
min x((A= xEll(M,))
max Pil is convex
(p))
for 1,2, ,T 17-10 do
Gt+ = 0² + 1/. VIZ((");
where [D(()) = Sa | p()))d())))))))))))()(())))))) + S | (=(x) du(x)) and 1p" Ω + R is the convex conjugate of p. Maximizers of the above Kantorovich dual prob- lem are referred to as Kantorovich potentials. Brenier's Theorem states that the Kantorovich potential tp is unique when y E P and the optimal transport map satisfies T = Vp. More details of optimal transport theory are referred to the monograph (Santambrogio, 2015).
8 = ()))
end for
For an arbitrary function ye, its second convex conjugate y is always convex and satisfies 8- S s. Consequently, the step y = (I*+1)** can be interpreted as projecting @++1 onto the space of convex functions. In addition, holds that [H S (y²"), ), indicating that applying the second convex conjugate does not reduce the functional value.
2.2. I Gradient
In this subsection, we review the concept of gradient on a Hilbert space and introduce a H1 graddiet descent approach for finding the maximizers of (xp) proposed by Jacobs & Léger (2020). Gâteaux derivative generalizes the standard notion of directional derivative to functionals. Given functional F IH - R defined on a Hilbert space HH with inner product (., >),, the Gâteaux derivative of F at o H in the direction h E IH, denoted by 5FF (h), is defined as
2.3. Wasserstein Gradient
In this subsection, we review the definition of Wasserstein gradient and a Wasserstein gradient descent approach for finding the Wasserstein barycenter of absolutely continuous probability measures. Let H. : (52) + R be a functional over the nonlinear space (9)). For any v E P22 n I, ( (S2), i.e., is absolutely continuous with an I density function, we can define the first variation of H. The map KH
d
of (n) Ω + IR is called the first variation of H at ur, if
(h) = de
F(6 + ch)
c=0
= J.
8H
Furthermore, the map VF H H is referred to as the Hi gradient of F if (V F( hyn FF holds for all h E H. When IHI is Rd, VF simplifies to the standard gradient of a function. Now, we consider the following homogeneous Sobolev space,
H(u + exx c= 8u de
(x)))x(())
for the direction X V μ.. Lemma 10.4.1 (Ambrosio et al., 2008) implies that the Wasserstein gradient of H. at μ is given by WH(p) v 8H
Ap (14) under mild conditions.
1 p(x) dx /
| dr, 2
We remark here that the Wasserstein gradient is fundamen- tally different from the gradient in a Hilbert space as defined in Subsection 2.2. The primary reason is that P5 (52) is not a linear vector space, and standard arithmetic operations such as addition and subtraction do not exist. For instance, given v, μ E PE(),, their difference ability measure and hence V μ & P5 (S). For the same reason, a different notion of convexity is appropriate for H. : PE)) + R. Specifically, H is said to be geodesically convex if, for any v4 l E (2) and c E [0, 1]. it holds that H((eTy + (1 - <) v)) & eH((p) + (1 ())().
= 1 R
[ [Vp(+) do < )
1 not valid prob-
where Vip is the weak derivative of yp (Evans, 2022). It is shown that - is a Hilbert space with the inner product (P1, P2) = (V)) (=)(x)) dx. As demonstrated by Jacobs& & Léger (2020), the liltt gradient of ZØ lāt- + R is given by
V
(V+*)
(2)
3. Mancoovve-CConcave Minimax Formulation For Optimal Transport Barycenter
where (- -1 denotes the negative inverse Laplacian op- erator with zero Neumann boundary conditions. We can always assume (p) E Hi! by noting that adding a con- stant to a function does not affect its Laplacian. Note that
In this section, we fomulate the Wasserstein barycenter problem as nonconvex-conent optimization problem. By
3

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
reviewing existing methods for computing the Wasserstein barycenter, we demonstrate that nonconvex-coneave for- mulation is more realistic and practical. We then propose gradient descent-ascent type algorithm and provide relavant convergence analysis.
3.2. Existing Approach for Wasserstein Barycenter
Given n probability measures f1, :2.....:Ar P the Wasserstein barycenter is defined as the minimizer of the barycenter functional F (I( + R. given by
F(v) = 1 n
3.1. W.N.oonvex-oonaanne Minimax Optimization in Euclidean Space
Σ W2(
(3)
Before presenting our algorithm, we first discuss nonconvex- concave optimization algorithms in Euclidean space to bet- ter understand the challenges and feasible objectives in such problems. Given smooth function Rd x E + I,, the nonconvex-concave minimax optimization problem is generally formulated
Wasserstein barycenter can be viewed as a generalization of the arithmetic mean in the Wasserstein space with met- ric W2. It is shown that the barycenter functional ad- mits a Wasserstein gradient WF(v) id - Σ.= TV; and a Wasserstein gradient based approach has been pro- posed (Zemel & Panaretos, 2019; Chewi et al., 2020) for numerically computing the Wasserstein barycenter by itera- tively updating as follows:
min max f(x y), RR VET
where V C Rd4 is convex and compact, and f(x, is concave for any fixed while (( can be nonconvex for a given y. Let ((x)= == maxyev (az v). If there ex- ists unique VI E Y that attains this maximal value, i.e., ((x) = f",, then by Danskin's s Theorem (Bernhard Rapaport, 1995; Bertsekas, 1997). ((xr) is differentiable and the gradient can be computed as (+(x) v (a,, vi), where v computes the gradient with respect to æ only.
n11 (id (id- W.F(F) 1t
However, the above algorithm implicitly assumes that the optimal transport maps (T } are known. In practice computing Th for multivariate distributions is particularly challenging and often can only be approximated to certain accuracy, for example, by using the Sinkhorn algorithm.
In this work, we reformulate the Wasserstein barycenter problem as conconvex-connan minimax problem. Rather than computing the optimal transport maps in each iteration, we propose a gradient descent-ascent algorithm to solve the associated minimax problem, where the transport maps are updated using -ascent in each iteration. Our approach alleviates the computational burden of solving optimal transport problems per iteration compared with the tradi- tional approaches.
The ultimate goal of the minimax optimization problem is to find the global minimum of कb However, such problem is NP hard due to the nonconvexity of & (Lin et al., 2020). A common surrogate in nonconvex optimization is to seek a stationary point æ of 6 where (+(x)) = 0. A simple and efficient method is the gradient descent-ascen, (GDA) algo- rithm (Algorithm 2), where Py is the projection operator onto Y.
Algorithm 2 Gradient Descent-Ascen Algorithm on Eu-
clidean Domain
3.3. Our Approach: Wasserstein-Deseen HHr -Aseent Algorithm
Initialize 31, V/ for /. 1,2, .............................. T do
Let Fa.s be a subset of H1' consisting of all functions that are a-strongly convex and 3-smooth, i.e. for every f E Fa, and x, y Ω2 it holds that f E H and 2| y/1 < f(z) f(y) (Vf y) < G vll..
x'+1 xx V V y²); n²+ = Pell' + + +VVyf(a' (f (²))
end for
With this notation, Fu,cc represents the set of all convex functions on Ω. The dual formulation of Kantorovich prob- lem implies
Despite the complex structure of the minimax problem and the nonconvexity of T the GDA algorithm remains theoreti- cally trackable. (Lin et al., 2020) proved that, if / is convex in and strongly concave in y, with suitable choices of step sizes (17, r). the following bound holds:
(v.
max IH' (p::). PFEFo,.0
Given n probability measures pl,..., pis....,/4, M E PT))), we can reformulate the Wasserstein barycenter problem as
min (2*())) [
which indicates that a good approximation of the station- ary point can be achieved with e-accuracy within the first O(1/E2) iterations.
-
n 1
I Ih (12)) ) (4)
min max vEP PLEKo,
= Σ

---

Optimal Transport Barycenter via Noneonvex-Coneave Minimax Optimization
Since the inner maximization part of (4) consists of n sepa- rable subproblems, using the notation
the - update v. However, in simulations, the Wasserstein gradient update significantly outperforms the L - -gradient update.
e (v) == max I (pi),
(5)
3.4. Convergence Analysis
(4) can be rewritten as min, max p.EF. a J (V: (p)
min, 1/2
n -i= (v). When a and B = xc, F 9 the
We now establish the notion of stationary points for Fa.p(v) := Ei= (v) and show that the output se- quence from Algorithm 3 converges to stationary point of Fv,)(v). We start from presenting several properties of the functionals I5 and L*.
set of convex functions and (2 po) max (p)). Thus, we have min, 1 = (v)()- min, F(v). The constraint set aa enforces additional regularity on the Kantorovich potentials. This technique has been frequently used in the optimal transport literature (Paty et al., 2020: Hütter& Rigollet. 2021; Manole et al., 2024). Fix v and 0 < < 3 xc, the objective functional [(s (p) is strongly concave in each fi: However, if we fix {
The first lemma demonstrates the strong concavity and smoothness of the functional IP on F F with < S B oc, when the density function of μ is bounded from below and above.
Fa Fa.9 without further assumptions, J(v (p) is not geodesically convex unless 3 S 1. Thus, (4) is a manconvex-concayye minimax optimization problem.
Lemma 3.1 (Strong concavity and smoothness of I5), If 0 < a S uuz) S b < do for all x E 22 then for any 91: 9 E IF Fa set A = aa /B and = b /a, the following inequalities hold,
We now dicuss different notions of gradients for the objec- tive functional J (12) x H' x x...X x H' + R. Given (2), the Th gradient of J with respect to yf; can be computed using (2). For fixed set {p C F. the definitions in subsection 2.3 imply that the Wasserstein gradient of J is given by W 4p) = id id-Ve VV, where p= 1 Ps. Before introducing a GDA type algorithm for solving the minimax optimization in (4), we summa- rize different notions of gradients for readers' convenience: (i) the usual gradient of yoi : _d + E denoted as Vip;; (ii) the HI' gradient of J with respect to yp; is computed as V a. (v,
A 11 +1ll 2
2
IT((2) - (11) (VI -11)
B
2 2
11
401 11
The following lemma provides an explicit form of the Wasserstein gradient of LH.
Lemma 3.2. If < a S p(z) Sb< < do for all x E D then (pp) admits a unique maximizer in Fa.g. Let Ty := arg max max a.8 (p)). Then, we have: (i) the first variation of L" at is 6C (v) (5d,id) :h: (ii) the Wasserstein gradient of L at is WCr(p) = id -
+ (-A)) + (Vo)) uu)): (iii) the Wasserstein gradient of J with respect to v is computed as [ (p) = id V where ক = 1/2 >i=1 Pi:
du
Let Ps... be the projection operator onto Fa,g. This pro- jection is well-defined and unique since Fo,s C H is a complete and convex metric space. We propose the follow- ing Wasserstein-Descenn HT -Ascent (WDHA) algorithm, of which the pseudocode is provided in Algorithm 3.
The above result directly implies that WFa,s(v) id -V where ~ Σ. Em= PM.. Consequently, it is natural to define the stationary points of Fa.g as probabil- ity measures for which the Wasserstein gradient has zero -,2-norm.
Algorithm 3 Deasserstein-eescen H -Asceet Algorithm
Definition 3.3. We call E (2) a stationary point of Fa,s if and only if Sa ll (F)a,(()) de O.
Initialize v',p''
for (. 10 1,2...., ..., T - 1 do
for = 1,2, 1,2,.... n do
We denote the dual norm of || 11 as || lla- defined by || == inf ff f.fg o dv || < 1}. For more informa- tion on || |B 1. we refer readers to Chapter 5 of (Santam- brogio, 2015). The following lemma indicates that W is Lipschitz continuous with constant 1/A A with respect to the H 'nomm where A is the constant from Lemma 3.1. -norm,
@t 4p + 17
T(r )
p PB., (
end for
yf+1_ (id - W J (vt, (p Am
end for
Lemma 3.4 (Lipschitzness of Wasserstein gradient W (.). 1ʃ0 a S m(x) < b < xo, then
Since [(. sp) is a functional on the space of absolutely continuous probability measures, it can also be viewed as a functional mapping density functions to R. By embedding all densities functions into (.2),),, we can alternatively use
W (C WC2(v2))lz 1|=>
= -1 - 6h ll S A| - -22lit11
5

---

Optimal Transport Barycenter via Noneonvex-Coneave Minimax Optimization
where
|| I|اa denote the ---norm of the function.
case where || diverges. Therefore, we hypothesize that this technical assumption can be inferred from other assump- tions, which we leave as an open problem. (iii) By definition, D is Wasserstein barycenter if id 1 Ei= V = 0, where oh is the Kantorovich potential between 17 and pai If yp E Fa,s for all i, then E/ is a stationary point of Fa,s. i.e., WFa,,s (V) = O. Reversely, we assume that the Kan- torovich potential between the true barycenter and each fai is in Fa,s, then (F 0 would mean that v' is Wasserstein barycenter.
Applying Theorem 5.34 in (Santambrogio. 2015), the above result further implies
x.
A Pth lla S 1| - Mallia 1 5
Vmmx<(lu looe (W2(1, (6)
We emphasize that the inequality above holds because Phy Pū are restricted to Fa,s- Otheewisee only a weaker bound in W1 metric is available (Theorem 1.3, Berman, 2021). l1 - Ltty Iلpn S <1 (11, where Wi 12) is the 1-Wasserstein distance between in and P2, and C₁ is constant depending on 1/1 and 1/. Follow- ing standard notations in the literature, we define ||0 = sup,, 1/1 (xx and ||ا||ع = sup, 12(2), where (x) and (2(2)) are density functions of vn and V2 evaluated at the point æ E Ω2 Following the above discussion and applying Lemma 3.4, we derive the smoothness of L with respect to the W2 metrie.
3.5. Computational Complexities
In this subsection, we discuss the implementation and computational complexity of the WDHA algorithm (Algo- rithm 3). To implement Algorithm 3, we need to numerically approximate the infinite dimensional objects v, (P1...... .P through discretization. Given V and 3p supported on a fixed grid of size m, the computation of the convex con- jugate yp", the pushforward measure (V and the neg- ative inverse Laplacian (-A)-1(v) with zero Neumann boundary conditions only requires a time complexity of (m log(m)) and space complexity of O(m), as demon- strated by (Jacobs & Léger, 2020). However, computing the projection Pe... (yp) is computationally expensive, with a time complexity (m²) (Simonetto, 2021). For more effi- cient computation, we recommend replacing the projection step Pi..o (yp) with computing the second convex conjugate (p))* in Algorithm 3 in practice. Although (p)** only en- forces the convexity and not strong convexity or smoothness, the modified algorithm performs well empirically. This ad- justed algorithm achieves a time complexity O(m log(m))
Lemma 3.5 (Smoothness of L*). Let C- = + max [|n ol 1|pz loo),
we have
C (12 (
C P
(id Vid-V TM: id) dv +
Finally, we establish the convergence of the WDHA algo- rithm 3 to a stationary point of Fa, below.
Theorem 3.6 (Convergence rate of WDHA), Assume that there are constants a and b, such that the density functions satisfy O- < a S px(x)) < b xo for all 1,2 ........ and E Ω2 Recall that A = /B and B 6,3d/ /a.
per iteration, and is provided below.
Algorithm 4 Wasserstein-Descenn I''- -Asentt Algorithm
If max, ||,'Il I V < 30 for some constant V > 0, by choosing the step sizes (T.1) satisfying 7 < 1/ B and T A2n
Initialize v² , for = 1,2, ... T - do
we have
An(Aa+A+V)+ 1V 2Am
for 1,2,.... n do
p + n V
J (v ,t);;
1
T
s (v') 11 dd
++1 (
W F
T
end for
8
T(V' (0 ( ))²²
(id T2 W
4 An
F - Fu
end for
TT/2
where 8² = (1² 10,,,p
_ ......,, pn) > 0 is constant.
Theorem 3.6 suggests that the parameters т, r) should be bounded above. Empirically, we find that the above algo- rithm works better with diminishing step sizes, potentially due to two reasons: (1) diminishing step sizes may reduce the effect of discrete approximations; and (2) diminishing step sizes may be more effective for nonsmooth convex func- tions. Since the second convex conjugate does not enforce strong convexity and smoothness, Lemma I no longer holds, and Ib is only guaranteed to be concave. In addition, dimin-
Remark 3.7. (i) The minimum value of the squared L of the Wasserstein gradient WF.,s over the first 7' iterations converges to zero at rate of O(T This convergence rate is consistent with the results obtained by using GDA to solve nonconvex-concat minimax problems in the Eu- clidean space, as demonstrated by (Lin et al., 2020). (ii) We assume that the density functions of all iterates sf are uni- formly bounded. In practice, we have not encountered any
6

---

Optimal Transport Barycenter via Noneonvex-Coneaee Minimax Optimization
Wasserstein Barycenter Uniform Densities
Uniform Densities
WDHA (Ours)
CWB (reg=0.005)
CWB (thresholded)
DSB (reg=0.005)
DSB (thresholded)
+
Figure 1. Illustration of Wasserstein barycenters computed by different methods. The goal is to compute the barycenter of four uniform densities supported the square, circle, heart, and cross, respectively, displayed the top left image. The blended shape shown in the top middle image is the barycentric density computed using our method. Barycentric densities computed using CWB and DSB with reg 0.005, their thresholded versions are shown the top right image and bottom three images.
ishing step size may speed up the learning process in the early stages and a inverse time decay for Ty, i.e., Ti = 1/t, works equally well in the simulation studies.
shapes of a square, a circle, a heart, and a cross, respec- tively. Their densities are discretized on a fixed grid of size m 1024 1024 and are displayed in Figure (top left). For this simulation, 7 exp(-t/TT and n! 0.05 for all and decrease it by factor of 0.99 if I (x(++) < I (v²))
apply Algorithm 4 and set
4. Numerical Studies
Using both synthetic and real data, we compare our approach with two existing methods applicable to dis- tributions supported on large grids: (1) Convolutional Wasserstein Barycenter (CWB) (Solomon al., 2015) and (2) Debiased Sinkhorn Barycenter (DSB) (Janati et al., 2020). Both CWB and DSB employ entropic regularization techniques and are implemented Python functions convoluti onal barycenter2d and www.onaa..arcceeteeeed.......... spectively, in the Python library "POT: Python Optimal Transport" (Flamary al., 2021). additional simulation with known ground truth is provided Appendix C.I.
The barycenters computed by our method, CWB, and DSB after 300 iterations are displayed in Figure 1. The density of the barycenter distribution generated by our method, as shown in the center the top middle image in Figure 1, is uniformly distributed over blended shape comprising square, circle, heart, and cross, with sharp edges. In contrast, the barycenters computed using CWB with regu- larization strength parameter reg = 0.005, displayed in the top right image in Figure 1, appear blurred. DSB yields a better representation than CWB but remains unclear; see the bottom middle image in Figure 1 for barycenters pro- duced by DSB with reg 0.005. Notably, encounter division by zero error if the regularization strength param- eter is set to 0.001. We also compute thresholded versions of barycenters of CWB and DSB by removing intensities
re-
Sythentic Uniform Distributions. In this example, aim to compute the barycenter of four uniform distribu- tions whose supports are contained in [0. 1|2 and take the

---

Optimal Transport Barycenter via Nonconvex-Coneaee Minimax Optimization
Wasserstein Barycenter of Handwritten Eight
Exemplary Digit 8
Exemplary Digit 8
Exemplary Digit 8
&
<
WDHA (Ours)
CWB (reg=0.005)
DSB (reg=0.005)
8
8
Figure Top row displays three exemplary digit images. Bottom displays barycenters computed by different methods using 300 iterations.
smaller than the threshold such that the removed intensities amount to 10% of the total mass. The thresholded barycen- ters are shown in the bottom left and right images in Figure 1. However, the resulting barycenters lack the inward sharp curvature inherited from the heart and cross shape. Thus, the barycenter obtained by our method offers a clearer and more representative summary of the set.
densities, we aim to compute barycenter of one hundred handwritten images of the digit 8, each with a size 500 500 pixels. Three exemplary images displayed in the top row of Figure 2. To run Algorithm 4, we set T. exp(-t/T)) and 1y! = 0.5 at iteration t = 1 and decrease it by factor of 0.95 whenever Z (gp 1) < I (sp)). The barycenters computed by our method, CWB, and DSB using T 300 iterations are displayed the bottom row of Figure 2. The barycenter computed by WDHA exhibits clearer and more detailed textures, revealing variations of the digits viewed as densities. Furthermore, method is more time- efficient, taking 3,299 seconds compared to 10,808 seconds for CWB and 1 11116 seconds for DSB.
We report the program run times for computing these barycenters and the corresponding 2-Wasserstein barycen- ter functional values F(vesr), where ysst represents the computed barycenter. All functional values reported below are estimated using the back-and-forth approach (Jacobs & Léger, 2020) and are multiplied by 10", All methods were executed on Google Colab with an L4 GPU, Our method takes 676 seconds, whereas CWB takes 3731 seconds and DSB takes 7249 seconds. Additionally, our method achieves the smallest barycenter functional value (74.5791), com- pared to CWB (75.0689) and DSB (74.5804). The func- tional values for the thresholded barycenters are 74.7346 (CWB) and 74.5921 (DSB).
5. Conclusion
In this paper, we introduced Wasserstein- Iip-- Ascent (WDHA) algorithm for computing the Wasserstein barycenter of n probability density functions supported on compact subset of Rdd. Our key technique is motivated by the recent progress in nonconvex-concave minimax op- timization problems in the Euclidean space. Compared to existing methods for high-resolution densities, the WDHA algorithm is computationally more efficient and produces
High-resolution Handwritten Digits. Here, our method applied to the high-resolution handwritten digits data (Beaulac& & Rosenthal, 2022). By treating the digit images
8

---

Optimal Transport Barycenter via Concenvex-Concave Minimax Optimization
clearer, sharper, and more detailed barycenter. We believe that work not only advances computational techniques for Wasserstein barycenters but also sheds new light optimizing nonlinear functionals using a combination of geometric structures.
Bigot, J., Cazelles, E., and Papadakis, N. Penalization of
barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 516):2221-2285, 2019.
Carlier, G. On the linear convergence of the multimarginal Sinkhorn algorithm. SIAM Journal on Optimization, 32 (2):766-794, 2022.
Acknowledgements
C. Zhu was partially supported by NSF DMS-2412832. X. Chen partially supported by NSF DMSS2413404, NSF CAREER Award DMS-2347760, and gift from Simons Foundation.
Carlier, G., Oberman, A., and Oudet, E. Numerical meth- ods for matching for teams and wasserstein barycenters. ESAIM: Mathematical Modelling and Numerical Analy- sis, 4966):1211-1642 2015.
Carlier, G., Eichinger, K., and Kroshnin, A. Entropic- Wasserstein barycenters: PDE characterization, regular- ity, and CLT. SIAM Journal on Mathematical Analysis, 53(5):5880-5914, 2021.
Impact Statement
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Chen, Y., Lin, Z., and Müller, H.-G. Wasserstein regression. Journal of the American Statistical Association, 118(542): 869-882, 2023.
References
Chewi, S., Maunu, T., Rigollet, P., and Stromme, A. J. Gradi- ent descent algorithms for Bures- Wasserstein barycenters. In Conference on Learning Theory, pp. 1276-1300 1304, 2020.
Agueh, M. and Carlier, G. Barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 43(2): 904-924, 2011.
Chizat, L. Doubly regularized entropic Wasserstein barycen- ters. arXiv preprint arXiv:2303. /844, 2023.
Álvarez-Esteban, P. C., Del Barrio, E., Cuesta-Albertos, J., and Matrán, C. A fixed-point approach to barycenters in wasserstein space. Journal of Mathematical Analysis and Applications, 441 441(2):744-762 2016.
Corrias, L. Fast Segeedrre-Fenchee transform and appli- cations Hamilton-Jccobi equations and conservation laws. SIAM Journal on Numerical Analysis, 33(():1534- 1558, 1996.
Ambrosio, L., Gigli, N., and Savaré, G. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008.
Cuturi, M. and Doucet, A. Fast computation of Wasserstein barycenters. In Proceedings of the 31st International Conference on Machine Learning, pp. 685 693, 2014.
Beaulac, C. and Rosenthal, S. Introducing new high- resolution handwritten digits data set with writer charac- teristics. SN Computer Science, 4(1):66, 2022.
Dubey, and Müller, H.-G. Fréchet change-point detection. The Annals of Statistics, 4816)3331 2020.
Berman, R. Convergence rates for discretized Monge- Ampère equations and quantitative stability of optimal transport. Foundations of Computational Mathematics, 21(4):1099-1140 2021.
Evans, L. C. Partial differential equations, volume 19,
American Mathematical Society, 2022.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. 7.,, Bois-
Bernhard, and Rapaport, A. On a theorem of Danskin with an application theorem of Von Neumann-Sion. Nonlinear Analysis: Theory, Methods & Applications, 24 (8):1163-118) 1995.
bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V., Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. POT: Python optimal transport. Journal of Ma- chine Learning Research, 22(78):1-8, 2021. URL http: / j/jmml r. o wwwpgg/appppsaaeeeem v 20--2/5 html.
Bertsekas, D. and Castanon, D. The auction algorithm for the transportation problem. Annals of Operations Research, 20(1):67-96 1989.
Gramfort, A., Peyré, G., and Cuturi, M. Fast optimal trans- port averaging of neuroimaging data. In Information Processing in Medical Imaging, pp. 261-272, 2015.
Bertsekas, D. P. Nonlinear programming. Journal of the Operational Research Society, 48(3):334-34, 1997.

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
Hütter, J.-C. and Rigollet, P. Minimax estimation of smooth optimal transport maps. The Annals of Statistics, 49(2): 1166 1194, 2021.
Nenna, L. and Pegon, P. Convergence rate of entropy- regularized multi-marginal optimal transport costs. Cana- dian Journal of Mathematics, pp. 1-21, 2024.
Jacobs, M. and Léger, F. A fast approach to optimal trans- port: The back-and-forth method. Numerische Mathe- matik, 146(3):513-544 2020.
Paty, F.PP., d' Aspremont, A., and Cuturi, M. Regularity as regularization: Smooth and strongly convex brenier po- tentials in optimal transport. In International Conference Artificial Intelligence and Statistics, pp. 1222-1232, 2020.
Janati, H., Cuturi, M., and Gramfort, A. Debiased Sinkhorn barycenters. In Proceedings of the 37th International Conference on Machine Learning, pp. 4692-4701, 2020.
Peyré, G. and Cuturi, M. Computational optimal trans- port: With applications data science. Foundations and Trends@ in Machine Learning, 11(5-6):355-607, 2019.
Jiang, F., Zhu, C., and Shao, X. Two-sample and change- point inference for non-euclidean valued time series. Elec- tronic Journal of Statistics, (8(1::448894) 2024.
Rabin, J., Peyré, G., Delon, J., and Bernot, M. Wasserstein barycenter and its application texture mixing, In Scale Space and Variational Methods in Computer Vision, pp. 435 435-446, 2012.
Korotin, A., Li, L.,, Solomon, J., and Burnaev, E. Continu- ous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021. URL https: //ppneeeeeeeeeeeee net/forume F..
Santambrogio, Optimal transport for applied mathemati- cians. Birkduser. NY, (5(58-63):94) 2015.
Simonetto, A. Smooth strongly convex regression. In 2020 28th European Signal Processing Conference, pp. 2130- 2134, 2021.
Korotin, A., Egiazarian, V., Li, L., and Burnaev, E. Wasser- stein iterative networks for barycenter estimation. In Advances in Neural Information Processing Systems, vol- ume 35, pp. 15672-15686. Curran Associates, Inc., 2022.
Solomon, J., de Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphs, 34
Kuhn, H. W. The Hungarian method for the assignment problem. Naval Research Logistics Quarterly, 2(1-2):
83-97, 1955.
(4), 2015.
Li, L., Genevay, A., Yurochkin, M., and Solomon, J. Contin- wous regularized Wasserstein barycenters. In Proceedings of the 34th International Conference on Neural Informa- tion Processing Systems, pp. 17755-17765, 2020a.
Zemel, Y. and Panaretos, V. M. Fréchet means and Pro- crustes analysis in Wasserstein space. Bernoulli, 25(2): 932 976, 2019.
Zhang. C., Kokoszka, P., and Petersen, A. Wasserstein autoregressive models for density time series. Journal of Time Series Analysis, 43(1):30-52, 2022.
Li, L., Genevay, A., Yurochkin, M., and Solomon, J. M. Con- tinuous regularized wasserstein barycenters. Advances in Neural Information Processing Systems, 33:17755- 17765, 2020b.
Zhu, C. and Müller, H.-G. Geodesic optimal transport re- gression. arXiv preprint arXiv:2312.15370 2023.
Lin, T., Jin, C., and Jordan, M. On gradient descent ascent for nonconvex-conee minimax problems. In Interna- tional Conference on Machine Learning, pp. 6083-6093, 2020.
Zhu, C. and Müller, H.-G. Autoregressive optimal transport models. Journal of the Royal Statistical Society Series Statistical Methodology, 85(3):0012-1033, 2023.
Zhuang, Y., Chen, X., and Yang, Y. Wasserstein k-means for clustering probability distributions. In Proceedings of 36th Conference on Neural Information Processing Systems, 2022.
Lin, T., Ho, N., Cuturi, M., and Jordan, M. I. On the com- plexity of approximating multimarginal optimal transport. The Journal of Machine Learning Research, 23(1):2835-
2877, 2022.
Luenberger, D. G. and Ye, Y. Linear and Nonlinear Pro- gramming. Springer New York, 2008
Manole, T., Balakrishnan, S., Niles- Weed, J., and Wasser- man, L. Plugin estimation of smooth optimal transport maps. The Annals of Statistics, 52(3):966 9988 2024.
10

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
Appendix
A. List of Notations
Notations
Meaning
H H-1
homogeneous Sobolev space dual space of Hi
Ra.s Pa(2) PI( Ta T Wy(,, v)
subset of THF consisting of a-strongly convex and B smooth functions space of probability measures with finite second order moment subset of (22) consisting of absolutely continuous probability measures
pushforward measure of under the map the optimal transport map from v to μ p- Wasserstein distance between y and Kantorovich dual functional defined in (1) first variation of the functional FF (22 + R
I
65
WF
Wasserstein gradient of the functional F
(V)
Wasserstein barycenter functional defined (4)
("(v)
the maximal functinoal defined in (5)
Fal
average of maximal functionals defined Fa,,
- L
A)-
negative inverse Laplacian operator with zero Neumann boundary conditions
Pro aai oo lp*
projection operator onto Fa.8
best a- e-sttongly convex, f-smooth Kantorovich potential defined Lemma 3.2
convex conjugate the function yp
V
(standard) gradient of the function yp identity map 1.2-norm of yp defined as (Sa, 111 dz2 L (y)))nnom of y,, defined (Jg || dz)
id
11<pllea
||<l|
B. Technical Details
B.1. Proof Lemma 3.1
Before proving the lemma, let us demonstrate a technical result for computing the A inner product first.
Lemma B.1. For any functions $0, E H and v E (S2)), we have
/ 01 122 du / ( 82 [viiddd du.
VID 122 -1))
n
Proof. Let = VI (P)). By the definition of VIJ in (2), we have
Ag = (Vvi)
Therefore, have
[ (12 /(022- dr
VIH(P)) -1)
Vg. V dz
In
n
(19 (2) - (Vzi)#m) de
n
/ p1 12 dv
91 o Vpi - 62 V du.
Here, (i) is by the definition of inner product in H. and (ii) due integration by parts.
Let us now prove the lemma.

---

Optimal Transport Barycenter via Nonconvex-Coneae Minimax Optimization
Proof of Lemma 3.1. Note that we have
I (122 (+1)
(PI) +i)
| (2(()) (x(x) + /
|
@
-((x)
2
2
12
[/A | 2) du(x)
| - du(x))
(i(
2
[L/ 82
(1 p2(x)
(11 o (v))(a)) (02 du(())
(i(x) - (z(() + (p) o vpt)(2) - (122 o (p))(x) du(z))
Here, we use the definition of Iø and Lemma B.1 to derive (i). By properties of convex conjugate, Vp*(y) arg max,en((,y) (xx), which further implies that
-²))= (Vp - ((V)'())).
Combining the above equality with (7) yields
[x( - - (VI ( pi)m
( x- (V (2((((())) 92 22 (V+i(x)) du(x))
@
p2 (v((((())) P2 (Vp2(z)) (VP2(())))) Vpi(z) Vyp(z)) du(x)
Bp (z), Vo)(1)) dm(z))
where B,, is the Bregman divergence of 92 defined
B2 (xx x²) == p2((x) (22') (Vip(t')), x').
Here, in (i), we use the fact that Vy2 V id. By properties of Bregman divergences,
B2 (z), Vx)(()) Bpg V p2 (z) Vi2 Vv)(()) = B,g (Vi2 (z)
Since u22 is a--strongly convex and f3-mmoth, we know 3 is strongly convex and 1/a-smooth. Thus, we have
1
V2 Vip(()) =112 - p3; (Vp2 o (p)(x), -
2a |V o Vpi(z) ==12.
23
Integrating the above inequality with respect to f and applying change of variable formulas entail
1
Vip
28
n
<
/
Bpp (VP2 o (p)(
1
/ ||Vp2(1) (=)
$ 20
where we used the fact that Vipi Vipi = id. The density of (Vri) = no V D2x Vipl. Since e is a-strongly convex and B-smooth, we have add < p(x) < bgd and thus
a
|1:02 11 s 1 Bpg (Vi2 V))(()),)) du S
6,3d
|| 1p 1111
2,3
2a
12

---

Optimal Transport Barycenter via Nonconvex-Coneaee Minimax Optimization
B.2. Proof of Lemma 3.2
Proof. We first show the uniqueness of maximizer. Suppose there exits two maximizers 91 + p2, then for any 7 [0, 1]. Pr 7P: (1 p2 is again maximizer. Applying Lemma 3.1,
A(1 7)2
1102 --illi 2 (pp) (101 (1 (1-7 81)
(8)
2
2
|| - 1111 2 (pr) (V
(9)
2
Adding (8) multiplied by y and (9) multiplied by 1 - Y gives
A(1 - 1)7 p2 - pi 2 - - - (11-)) (122).
2
For fixed & (0, 1), the right-hand side of above is O, while the left-hand side is strictly smaller than 0. This shows a contradiction and the uniqueness is proved.
Next, show that the first variation is given by 6C du (v) So
Ialli T.t dx. Note that
d
max Iy Ib+ex(o) max (4)
C (v + cx)
de
=
vtex (Ph+ee)
-I
(id,id)
Plttex dx
2
{id, id)
@ dx
as < 0.
2
On the other hand, we have
max (p)) max,eB.... I
In (pth
(Tt)
(id,
5t dx.
Recall that the Wasserstein gradient is just the standard gradient of the first variation. Together, the Lemma is concluded.
B.3. Proof of Lemma 3.4
The equality part is direct by applying Lemma 3.2,
WC - w (22))la = V Vig Iم = 11 -4th lla.
To prove the inequality part, note that Iz, is concave by Lemma 3.1, and Fa,s is convex set. By the optimality of G%y G for any &o E Fa we have
(p P
mm 50, (7%
(10)
(p
)) S 0. VI
(11)
Substituting 9 P in (10) and y T in (11) and summing them together results
(Fth
Piz
V (Ft,
VI (it )) S.. 0.
(12)
The following two inequalities follow from Lemma 3.1,
A
IH (7
IHt (TM
VVI (GH))
F4 5
1 11 1
A
IH (7Hh)
Ih, (5t
1
11
13

---

Optimal Transport Barycenter via Nonconvex-Coneae Minimax Optimization
Summing over the above two inequalities gives
(PtH) VI ( Ply) S A PC, llp
Then, combining the above inequality with (12) shows that
A
Pún ll S VI (Ptt) - VIth (Pth Pi pty))r
| 6 2ht d(u1 122
F - phe 111l us lle
Here, (i) is derived by applying Lemma B.I as
( ) Pi - -bb)))
V (P/)) P - (
ph du / 17t
Plzz
( )" -h (( )"] du
1 Pin 1th d22
/
1 v((tt)) - p ( []dy
82
n Piz Ph d(u - U2).
B.4. Proof of Lemma 3.5
Proof. Notice that
(id, id)
(2222) - -C"(1))-
60t dv/ - 2/1
2
(id,id)
to)te(u2-n) duz - de - 1/// (id, id)
Púù duz de
2
2
8th
-
-pi+t((22-1) dVz 11 de
J0
Ja
J0 1|74 3u
Iiin 1/ - -1 de
|10(uz - milla 1|u2 - -villa de = 2A ||4 - 122
A
(10 max
|Pn bloo 1122 W1()),)))
2A
where (i) is due Lemma 3.4, and (ii) follows from Theorem 5.34 (Santambrogio, 2015). Since (dd,d) Gh is (1 (1+0 a)- smooth, have
iid, id)
2 Gy, dv2
V1
(d,
id)
iid, id)
(t,) T
2 Gy, id du
2
+ a
<
(id
V b4
Tue id} +
| id 11 del
0
Combining the above two results would conclude the lemma.
14

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
B.5. Proof of Theorem 3.6 Lemma B.2. For any xpt E (a)) let r0tt = .ra,s ((² + ( If S I/B, then 13411-8110 (1 - An)l - 1
Proof. Recall that Φ = arg max,pax...... , (g)). We have
|| - pb
11:0" 17 (II - 13
=11 -1|Ø- @pt li + +
+ 2n (I:((')) 1 -
A
V ( || + (v0')) - -I5 (+t)
11:0'~ -
2
Here, (i) is due to the property of the projection map, and (ii) by Lemma 3.1. Since Hi is a linear space and p² + v (Ig('') E H., we have
(+t) (5' + VIE(T"))
B
+ (F') 1 (z)
1/-
B 1 - -
+
B
2
1 2B
-x'( +
(1'
Again, (i) is due to Lemma 3.1. Combining the above two inequalities yields
11 - l|
< <1100 - 1E + (T) 21 TT((')) - (x6) 2 11 - 1
= An) - G + 21y(1 Bn) (1¹)) - IT
If S 1/B, we have || - 1 1| S (1 - An) - 25112
Proof of Theorem 3.6. Since vi+1 = (id - - Vip')) where yo = n each i,
4 Ei= shi
we have from Lemma 3.5 that for
- ST Vid VE' - id) du de' + +
/
- -id dud"
2 1
where C1 =1++++ V. A. Averaging over yields
F Fa, - -Fay)(Fa,s (v²)
(id- -V Vp* id) du + -² /
ST J
id d2
2 .2
T
|V - V
T
V - id
--
-0
/
|V - id de
2 n
2
2 n
2T -C
37 7201
- de
VV - id
11 de
2
4
nn
(2T- - PGC)
37
72C
A
12 dz
id
12
2
4
3r
-2
/
STV
id 12
de
4
VV
37
7²C
id
de
(13)
n
15

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
where (1) is due to the fact & Jp. ||V - id 1/ = Sa IV
V + So V id 11 de' Set o = -In |V
Vogil dr. By Young's s inequality,
1
J
|V V [1 1 An
( 1 1)] / | 1.d
1
A 1)
An
For the first term above, applying Lemma B.2 yields
[ + 2( 1 1) An
1
| - S (1- A o
)
2
For the second term, Lemma 3.2 and Theorem 5.34 (Santambrogio, 2015) implies that
Vith 11 dr l1
1
( V -
||
I A²
Wg(e²~1,1
A2
(1))
V
(
(Id--WJ((v" 1) id =1======= V²V
A² / || 11pp--''11
d2
A2
Jn
2
VV 1 Vi- 1/ d
1| 1 VV
id id12dr
A²
2
V
n
/ |V -id 1| du
A²
n
Here, (i) is due to Lemma 3.4, (ii) is due to (6) (Theorem 5.34, Santambrogio, 2015), and (iii) is because id - WJ(''11 1 is transport map from n² to v' Combining above pieces together yields
0 S [1
An 0:
+
1 +
2
1
2
V
/
id 12 d
2
AIJ
22
Set 8' = = Ex of and
1
An -
(2 And A
Averaging the above inequality for i E {1, ,n] yields
8 S 70' 8- + 2T²V V22V(2 (2- An)
=-- - id 112 do-
VV
A3
181 + 2 An))
[-1
S 1811881
- id 1 de
(14)
A3
8
Putting all pieces together yields
T
Fa.g(uTt)) - Faa)(u')= =
Σ
[ (Fa.(utt))) - [(a)((()]
1=1
T
3 - -2C
8
VV - id 11 de
5
(8)
T
2- V2 An)
1-1
37
G2G BGG
1-8 +
|V
id 112 dv
VV
id
dv
A'n
4
n
k=1
1 1-1T T
T
V22 - An)
7
T-t
3r
=V8
Vi - id dd
A
1
4
1=1
(m)) 4
V2 - AA)
4
3r -
-2C
T
14 - id 12 de
An
A
An
4
[=1
4
37 An)
3T
7 C
T
|V
id dv
An
A4r2
4
16

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
reg
Wasserstein distance
L di-dintance
F(
WDHA
3.758 x 10-3
0.4869
9.001 X 10
0.003
2.98 x 10-4
1.321
9.031 x 10-2
CWB
0.002
2.538 x 10-2
3,041
0.1154
0.005
1.2 10
0.8219
9.013 10
DSB
0.004
1.164 x 10-2
2.281
0.1016
Table I. Simulation results for uniform distributions supported on round disks.
Here, (i) is from (13), (ii) is due to (14), and in (iii) we use the fact that
An
2r V (2
An)
An 4
2
A3
A²n 2V (2(2-Am)
when <
Therefore, we have
T T
4rV8
+ Fa.)(12) F ((vEt))
VV
id
du
du'
S
1
An
T
G,72
81222-A11)
R
1
A'11² 1/1²
4r18"
+ F Fay)(v)) -
- -Fa,((utt))
An
TT/2
A²1
where the last inequality holds when <
C Ci+2
sv2
AVV
C. Additional Empirical Studies
C.1. Uniform Distributions with Ground Truth
Here, the goal is to compute the barycenter of four uniform distributions, each supported within [0, 1]2 and taking the shape of round disks with a radius of 0.15, centered at (0.2,0.2), (0.2, 0.8), (0.8, 0.2), (0.8, 0.8), respectively. Their densities are discretized on a 1024 x 1024 grid. It's clear that the true barycenter is uniform on the disk of radius 0.15 centered at (..5,0.5). The computed barycenter densities by WDHA, CWB and DSB are shown in Figure 3. We note that regularization parameters reg= 0.003 and reg=0.005 are the optimal choices for CWB and DSB respectively. Smaller regularization parameters lead nonconvergent and worse results for both CWB and DSB. We report in Table 1 the Wasserstein distance between computed barycenter distribution and the truth, _ between computed barycenter densities and the true density, and the barycenter functional value. Our method is uniformly the best, and in particular, the improvement in the Wasserstein distance is of orders of magnitude.
C.2. Experiments on 1D Distributions
In this empirical study, we compare the performance of Algorithm 3 (with projection onto F. .. and Algorithm 4 (with double convex conjugates) on ID distributions. For each repetition and i = 1, 3, we let μ4, be the truncated version of N N(a; a,?) on the domain [0. 1]. where a; ~ uniform (0.3, 0.7] and a; ~ uniform [0. 1]. Let V,, 1/,,, 1/2, be the true barycenter, computed barycernter from Algorithm 3, computed barycernter from Algorithm respectively. We repeat the experiment 300 times and report two types of average distances between the groundtruth and each computed barycenters average W2- distance $ DT W₂ (V,, 1,,e) and average L'-distance between densities $ ΣT ( (T) )=.)(r)) ,j 1, 2.
10-3 B 10²) performs slightly better with Wi-distance 7.610 10 (standard deviation: 3.367 x 10-5) and 1/2 d.2-distance 0.608 (standard deviation: 0.194), while Algorithm 4 has Wy-distance 7.771 x 10-5 (standard deviation: 3.32 10 10-5) and L'-distance 0,6434 (standard deviation: 0.451). In addition, CWB has Wi-distance 0.0022 (standard deviation: 5.85 x 10-4) and L distance 0.082 (standard deviation: 0.039). DSB achieves Wy-distance 1.225 x 10 (standard deviation: 2.644 x 10 10-5) and L dii-tttance 0.0699 (standard deviation: 0.034). DSB performs the
Algorithm 3 (with
best in Wasserstein distance for this example.
17

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
Wasserstein Barycenter Uniform Densities
Uniform Densities
WDHA (Ours)
CWB (reg=0.003)
CWB (reg=0.002)
DSB (reg=0.005)
DSB (reg=0.004)
Figure 3. Illustration of Wasserstein barycenters computed using WDHA, CWB and DSB.
C.3. Experiments on 3D distributions
We extended our WDHA implementation to handle 3D distributions and conducted an example in such setting. The goal is to compute the barycenter of three uniform densities supported on different shapes of ellipsoid contained [o, 1]3. Our simulation setup is: (x-0
(a-0.5) < 0.82, where the parameters
(a, b, E ((0.25, 0.25, 0.5), (0.5, 0.25, 0.25), (0.25, 0.5, 0.25))
These three densities are discretized on the grid of size x 200 x 200. In addition, we include weighted barycenter setting, where the weights 5 hours compute the barycenters. Figure showw the computed barycenters by our methods. Moreover, note that current version of the Python package POT does not support the computation of barycenter for 3D distributions.
be (1/4, 1/4, 1/2). We ran our algorithm on cluster with 2000 iterations and it took
C.4. descent H -ascent algorithm
Let v(x), (xx) be density functions,
write
J(u,0) 1 n
| 2
/ (Ilrl1 po (4(x)) dr.
Σ
p((()) v(x)
v(x) dx
Let (²(A) = {h: //(x///((////4/ < x
be the space of functions that are L'a-interrable with respect to the Lebesgue measure = 13 be the set density functions. Notice that the L2-gradient of with
and D (f E (()
respect to v(z) is (v, := + | (x)), we may consider the below
- -deseett H -sscett algorithm described
18

---

Optimal Transport Barycenter via Nonconvex-Coneave Minimax Optimization
Uniform Densities
WDHA(Ours)
Weight 1/3
Weight 1/3
Weight 1/3
Weight 1/4
Weight 1/4
Weight 1/2
Figure 4. Illustration of Wasserstein barycenters computed using WDHA for 3D distributions. The results are smoothed using a Gaussian filter.
Algorithm 5 T.2- [.-Descent THI -Ascent Algorithm
Initialize u1 1,²;; T 1-110 do
for / = 1,2, 2, ... ,
for = 1,2, 1.2.....77 do
Gt+ ot + m/
p.J sp));
+ (6--1)- -
end for
jf+1 y¹ V. JV 0)::
u = Pp (it+1);
end for
Here, (D(D++)E is the projection i++1 onto D, which can be computed using Python function pyproximal Simp lex in the library PyProximal. We apply this algorithm to the uniform distributions supported on round disks, and observe that diverges leading wrong result. Nevertheless, we plot the output in Figure
19

---

Optimal Transport Baryeenter via Concevvex-Concave Minimax Optimization
Figure 5. Plot of output from Algorithm C.4 to uniform distributions round disks.
D. Additional Details
D.I. Wasserstein gradient
To compute (V,,,)) we apply the definition in subsection 2.3. Note that
d [v + Sx.
de
6==0
d
|
/ |
(2(()) d(v + Ex)+ +
dpp(x)
de
n1
1m
2
2
1 2
d
|
((())ax
da
2
1=1
d
|<012
F(x) dx
de
|
(xx) dx.
8
2
which implics that
87 ou
(x)(2) =
1n13
(2x). By the definition of Wasserstein gradient, WJ
v 0.7 (1) id Vip
2
By
D.2. Implementation
Let Ω = [0, 1J and {(x;, Mj) }r be the equally spaced grid points, we have xo O, yo = 0 and .;; = i/m, = j/ m for ?.j * 0. Given the evaluations (Pay = sp(x,,,,) 833} of function y on these grid points, we compute the gradient at point
(1,, i. # as
Piui -
Ve((,,8))
lg Pi.s 2..
where h 1/m. For the computation of convex conjugates of ip, note that for the 1D case, convex conjugate can be computed efficiently using the method in (Corrias, 1996). For the 2D case, notice that
(v) (v),
sup (m1 /1)3/2 + (x2 y2)² x2x)1
31.2
sup -01-913 ( + sup { (x2 1/2 p(2)112))
sup ((11 - 3//3/2 + (5(01, (v2))
Thus, the convex conjugate the 2D case can be computed by iteratively applying the ID solver to each row and column.
20

---

Optimal Transport Barycenter via Noneonvex-Coneaee Minimax Optimization
To discuss the implementation of (V, describe how the mass v(z My) (density value of at point (11) is splitted and mapped (Jacobs & Léger, 2020) as follows. Since yp is convex, observe that Vaslzzi Mi) Vag(Ii+1, Wj) and Vyppaz, Uj) I Vup(zz Vz+1). Let R(x,, v2) be the quadrilateral formed by 4
points Vip(zs, Vep(a++1.8)) Uz), Vipzz, U5+1), Vp(zz+11 Ms+1) and pick the mesh grids {[Zv, My R(x4, where
(Tv. = (1 av av) By Ve(rs, 1j) + az - By Ve((t+1)))
+ (1 02)33 )3, (z, U)+1) + 0:. By Ve(
with 0 = 0o S 01 < ... < 02 1,0 30 31 < ... v 3x = The mass of v(z,, Uj) is first uniformly distributed
to the meshed grids {(zv, yy ) Then, the mass at each point (ăy, Fy is distributed to nearest grid points in {(r., 9;)}, inversely proportional to their distances. ØV exceeds the grid specified, the mass will be distributed to the boundary points instead.
21


Reviews:
[{"review_id": "AsMW4ljbvY", "reviewer": "Reviewer_9GjF", "summary": "This paper introduces the Wasserstein-Descent H-Ascent (WDHA) algorithm, a primal-dual method for computing the exact Wasserstein barycenter with nearly linear time and space complexity. WDHA alternates between Wasserstein and Sobolev optimization geometries for the primal barycenter and dual Kantorovich potential subproblems. The algorithm outperforms existing Sinkhorn-type methods in terms of computational efficacy, scalability, and accuracy, as demonstrated on high-resolution 2D synthetic and real data.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments. Please find our response to your\nquestions below.\n\n> I think \\[1\\] should be discussed that one can resort to using first\n> order methods such as subgradient descent on the dual for exact\n> barycenter calculations.\n\nThanks for pointing out this relevant first-order method; we will\ndiscuss the paper \"\\[1\\] Numerical methods for matching for teams and\nWasserstein barycenters\\\" in the revised version.\n\n> The experiments seem too simplistic\\... I suggest the authors include\n> more comprehensive experiments\\...\n\nWe extended our WDHA implementation to handle 3D distributions and\nconducted an example in such setting. The goal is to compute the\nbarycenter of three uniform densities supported on different shapes of\nellipsoid contained in $[0,1]^3$. Our simulation setup is:\n$\\frac{(x-0.5)^2}{a^2} + \\frac{(y-0.5)^2}{b^2} + \\frac{(z-0.5)^2}{c^2} \\leq 0.8^2$,\nwhere the parameters are\n\n$$(a,b,c) \\in \\{ (0.25, 0.25, 0.5), (0.5, 0.25, 0.25), (0.25, 0.5, 0.25) \\}.$$\n\nThese three densities are discretized on the grid of size\n$200 \\times 200 \\times 200$. In addition, we include a weighted\nbarycenter setting, where the weights is set to be $(1/4, 1/4, 1/2)$.\n(see below for the extension to weighted barycenter). We ran our\nalgorithm on a cluster with 2000 iterations and it took 5 hours to\ncompute the barycenters. The figure showing the computed barycenters by\nour methods can be accessed using the anonymous link\n<https://drive.google.com/file/d/1k5Yg4YF9FzkEOi7PC03jep5coXKNb2Vz/view?usp=sharing>.\nMoreover, we note that current version of the Python package POT does\nnot support the computation of barycenter for 3D distributions.\n\n> Is this algorithm limited to barycenter weights set to $1/n$? How\n> would it handle other barycenter weight settings?\n\nOur algorithm can be generalized to the computation of weighted\nWasserstein barycenter with any barycentric coordinates\n$(\\alpha_1, \\dots, \\alpha_n)$ on the probability simplex\n\n$$\\overline{\\nu} = \\arg\\min_\\nu \\sum_{i=1}^n \\alpha_iW_2^2(\\nu, \\mu_i).$$\n\nPrecisely, we can replace the $\\dot{\\mathbb{H}}^1$-gradient of\n$\\mathcal{J}$ with\n$\\mathbf{\\nabla}\\_{\\varphi\\_i} \\mathcal{J}(\\nu, \\mathbf{\\varphi}) = \\alpha\\_i (-\\Delta)^{-1} (-\\nu + {(\\nabla \\varphi\\_i^{*})}{\\\\#} \\mu\\_i)$\nand replace the Wasserstein gradient of $\\mathcal{J}(\\nu,\\varphi)$ with\n$id - \\sum\\_{i=1}^n\\alpha\\_i\\varphi\\_i$.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes, the proposed algorithm seems to work according to Figure 1.", "theoretical_claims": "Yes, I have checked the proof of Theorem 3.6 in Appendix B.5.", "experimental_designs_or_analyses": "Yes, the performance is shown in Figure 1, and the running time and time complexity are given in Sections 3.5 and 4.", "relation_to_broader_scientific_literature": "This paper proposes a GPU-friendly primal-dual descent method for exact Wasserstein barycenter calculation, whereas previous works mainly focus on regularized barycenter calculations, as shown in [1,2,3].\n\n[1] Iterative Bregman projections for regularized transportation problems, 2016\n\n[2] Fast optimal transport averaging of neuroimaging data\n\n[3] Fast computation of Wasserstein barycenters", "essential_references_not_discussed": "I think [1] should be discussed that one can resort to using first order methods such as subgradient descent on the dual for exact barycenter calculations. \n\n[1] Numerical methods for matching for teams and Wasserstein barycenters", "other_strengths_and_weaknesses": "Strengths: The writing is good, the idea is interesting, and the theory is complete, with full proofs provided for every lemma and theorem. Additionally, the authors have provided a clear and complete algorithm along with its code.\n\nWeaknesses: The experiments seem too simplistic.", "other_comments_or_suggestions": "I suggest the authors include more comprehensive experiments.", "questions_for_authors": "Is this algorithm limited to barycenter weights set to $1/n$? How would it handle other barycenter weight settings?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "J7EcreVL1f", "reviewer": "Reviewer_BeVm", "summary": "This work introduces a coordinate optimization algorithm for the computation of Wasserstein barycentres for probability measures discretized on compact Euclidean domains (such as images). By reformulating the traditional barycenter problem as a nonconvex-concave minimax optimization problem and employing a Gradient Descent-Ascent (GDA) approach, authors alternate between the barycenter atoms updates in the Wasserstein space and the _n_ Wasserstein dual maximizers (Kantorovich potentials) updates in the Sobolev space. The Sobolev gradient, under reasonable assumptions, allows for faster numerical convergences and more precise barycenters. Multiple theoretical analysis highlights the properties of the proposed approach, the Wasserstein descent H1 ascent algorithm (WDHA), as well as 2-dimensional examples demonstrate that WDHA outperforms OT state-of-the-art algorithms.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments and careful reading. We will fix minor errors in the revision. Please find our response to your main questions below.\n\n> As for experiment 3, here it seems that spatial notions are introduced\n> in the OT problem. Space, understood as the coordinate (x, y) of the\n> four input distributions considered \\... the comparison with\n> techniques such as CWB and DSB is unfair.\n\nThe experiments 1-3 were arranged on the same $1024 \\times 1024$ grid, except that they are uniform distributions on different supports/shapes. We agree that the spatial notion of placing the shapes (e.g., centers of disks/circles) generally implies an embedding of distributions. However, in our experimental setting, the 4 shapes were **equilaterally** arranged in the corners from the grid center. This means that the coordinate information $(x, y)$ does not impact the barycenter. For instance, we can move all shapes towards the center of the square grid domain (squeezing), and the barycenter does not change. Effectively, all of the 4 input distributions sit in a common domain centered around $(0,0)$ and the problem can be seen as a discretization on a fixed grid. Thus the comparison with CWB and DSB is fair in this particular configuration. To avoid confusion, we will revise line 911 to ``*Here, the goal is to compute the barycenter of four uniform distributions, each supported within $[0,1]^2$ and taking the shape of round disks with a radius of 0.15, centered equilaterally at $(0.2, 0.2), (0.2, 0.8), (0.8, 0.2), (0.8, 0.8)$, respectively. Their densities are discretized on a $1024 \\times 1024$ grid.*\"\n\n> Is not straightforward for me the nonconvexity (w.r.t. the\n> barycenters) and concavity (w.r.t. the potentials) of Problem (4).\n\nWe first provide some insights about the nonconvex-concave structure of\nthe barycenter problem. As mentioned in replying to Reviewer 9GjF, the\nbarycenter problem is defined as\n$$\\min_\\nu \\sum_{i=1}^n \\alpha_iW_2^2(\\nu, \\mu_i),$$ where the objective\nfunction is **geodesically nonconvex** in $\\nu$ for $n \\geq 3$ (in fact,\nit is concave over Gaussians!) However, the barycenter objective is\nconvex in the linear structure (cf. our reply to Reviewer sfZf). This is\nconnected to the **concavity in the dual program of potentials**, where\ntheir $L_2$ gradients are key to defining the Wasserstein gradient of\n$\\mathcal{J}$. Hence, the Wasserstein barycenter functional above can be\nrecast into Problem (4) as a primal-nonconvex and dual-concave program\nin two different geometries. Below we give more rigorous details about\nthis structure.\n\n**Strong concavity in $\\varphi$.** Lemma 3.1 implies that\n$\\mathcal{I}\\_\\nu^\\mu(\\varphi)$ is (strongly) concave on the function\nspace $\\mathbb{F}\\_{\\alpha,\\beta}$ w.r.t. $\\dot{\\mathbb{H}}^1$-norm when\n$\\mu$ admits a density function bounded away from zero and uniformly\nupper bounded. As the objective functional is\n$\\mathcal{J}(\\nu, \\varphi) = \\frac{1}{n}\\sum_{i=1}^n \\mathcal{I}\\_\\nu^{\\mu\\_i}(\\varphi\\_i)$\nis separable in potentials, we know $\\mathcal{J}(\\nu,\\varphi)$ is\nstrongly concave w.r.t. $\\varphi$.\n\n**Nonconvexity in $\\nu$.** Recall that\n$$\\mathcal{I}\\_\\nu^\\mu(\\varphi) = \\int \\frac{ \\\\| x  \\\\| \\_2^2}{2} - \\varphi(x)  d \\nu + \\int \\frac{ \\\\| x \\\\|\\_2^2}{2} - \\varphi^\\ast(x)  d \\mu,$$\nwhere $\\varphi^\\ast$ is the convex conjugate of $\\varphi$. Only the\nfirst integration in the above definition depends on $\\nu$. It can be\nshown that the functional\n$\\nu\\mapsto\\int \\frac{\\rm{id}^2}{2} - \\varphi  d \\nu$ is geodesically\nconvex if and only if the function $\\frac{\\rm{id}^2}{2} - \\varphi$ is\nconvex, which cannot be guaranteed in our WDHA algorithm.\n\n> is not clear if parameters alpha and beta, defining the subset of the\n> Sobolev space, are fixed a priori; are related (in some sense) to the\n> input data; or are optimized.\n\nPrameters $\\alpha$ and $\\beta$ are fixed: they are used for the theoretical purpose for convergence analysis of WDHA algorithm. In practice, we suggested using Algorithm 4 by taking the conjugate envelope (i.e., double convex conjugates) at each iterate, which could help to reduce the computational time than projection. We verified by numeric example that practical differences between small-$\\alpha$-large-$\\beta$ and conjugate envelope are small (see also our reply to Reviewer rrvS7 on Experimental Designs).\n\n> you suggest to replace the dual potentials projections with computing\n> their second convex conjugate. Did you do some experiments to check\n> the difference of the two versions of the algorithm for 2D cases (with\n> visualizations)?\n\nReviewer rvS7 raised the same question. Please find our reply therein.\n\n> Is it possible to test the behaviour of the WDHA solution in a 3D\n> context?\n\nReview 9GjF raised the same question to give a 3D example. Please find\nour reply therein.", "claims_and_evidence": "Yes, all the claims presented in the paper are all supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The authors are clear about the context and limitations of the problem addressed and the proposed approach. They say, “the current approach is mainly limited to computing the Wasserstein barycenter of 2D or 3D distributions supported on a compact domain”. Moreover, comparison methods and evaluation criteria are coherent with the presented application (images). However, exploring the solution of such approach in 3D scenarios could be interesting.", "theoretical_claims": "No", "experimental_designs_or_analyses": "Yes, experiments 1 and 2 are consistent. As for experiment 3 (in section C.1 and Figure 3), here it seems that spatial notions are introduced in the OT problem. Space, understood as the coordinate (x, y) of the four input distributions considered. I am not convinced that this scenario is coherent with the paper and the previous applications, where the distributions are discretized on m-point grids. \nFurthermore, the notion of space implies a natural embedding of distributions as point clouds, characterising circles, and consequently the comparison with techniques such as CWB and DSB is unfair.", "relation_to_broader_scientific_literature": "The proposed approach (WAHD algorithm) extends in an innovative way the optimal transport (OT) literature and further expands the wide range of methodologies for computing barycenters of continuous probability measures. Although the Wasserstein barycenter problem has been deeply studied, under different frameworks (discrete and continuous case, addition of entropic regularisations, etc.) and across a wide range of domains (images, shapes, high-dimensional distributions, etc.), the WAHD algorithm establish a direct connection between “classical” OT optimization and differential optimization. The key step of this approach is to pass to solve _n_ independent OT problems (primal problems), where _n_ is the number of input measures, to solve the respective dual problems in a Sobolev space. The updates of the dual potentials via gradient ascent with the homogeneous Sobolev gradient, leads to a smaller computational cost per iteration, which is very interesting in the case of high-resolution images.", "essential_references_not_discussed": "To the best of my knowledge, no. The bibliography is comprehensive and sufficient for covering the different areas of interest discussed in the article. Moreover, accurate and consistent citations are provided for the general understanding.", "other_strengths_and_weaknesses": "This work incorporates modern gradient-based optimization techniques with classical optimal transport problems in an original way. The switch to the computation of gradients of the dual potential in the Sobolev space is very interesting and proves to be effective for a more precise and detailed barycenter. I am very interested to see the potential applications of this approach, as it opens up new avenues for the application of modern OT techniques in machine learning.", "other_comments_or_suggestions": "List of minor errors:\n- Throughout the paper, there are numerous references to Theorems, Lemmas and Algorithms (introduced in the paper) that are inconsistent. For example, Algorithm 4 is referred to in the experimental section as Algorithm 3.5. The same applies to the theoretical statements on pages 5 and 6.\n- I assume that the method used for the computation of the second barycenter in Figure 2 (row 2, column 2) is CWB. WSB has never been defined.\n- Page 4, Line 215, l.h.s. : Missing an \"S\" for “(S)et”\n- Page 4, Line 212, r.h.s. : First equation in Sec. 3.3, \"The dual formulation of Kantorovich problem implies\", the min sound be a max (am I right?)\n- Page 5, Line 229, r.h.s. : Error in writing the functional “_I_\\_{alpha, beta}” should be “_I_\\_{nu}^{mu}”", "questions_for_authors": "1. Is not straightforward for me the nonconvexity (w.r.t. the barycenters) and concavity (w.r.t. the potentials) of Problem (4). Could you please provide deeper details for such claims? Moreover, is not clear if parameters _alpha_ and _beta_, defining the subset of the Sobolev space, are fixed a priori; are related (in some sense) to the input data; or are optimized. Could you clarify better this point?\n2. In Sec. 3.5, to speed up the computational costs you suggest to replace the dual potentials projections with computing their second convex conjugate. Did you do some experiments to check the difference of the two versions of the algorithm for 2D cases (with visualizations)? For example, how the barycenter change in the scenario of Example 1, estimating it using Algorithm 3?\n3. Is it possible to test the behaviour of the WDHA solution in a 3D context? For example, interpolating between shapes?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "N7zvmogNEu", "reviewer": "Reviewer_sfZf", "summary": "The paper propose a novel Wasserstein Descent H1-Ascent algorithm for the barycenter on W2-space, which is an important problem without an accurate and computational efficient solution. The proposed method adopt Gradient on the functional in the dual formlation of the W2 optimazation and apply it into the min-max problem for the barycenter with a extension on the Euclidean space. The convergence is proved under the assumption for the bounded density. The computational complexity is also discussed with a 2D numerical study to show the effectiveness of the method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments. In the revision, we will cite the paper [Daaloul 2021],  which proposed an interesting notion for approximating and sampling from unregularized Wasserstein barycenter. Please find our response to your main questions below.\n\n>the authors also need to clarify which algorithm uses GPUs or all of them used...\n\nAll algorithms are executed on Google Colab using the same hardware configuration. We compared the runtimes on L4 GPU and CPU for Example 1:\n\n|                | Ours    | CWB   | DSB   |\n|:----------:|:--------:|:--------:|:-------:|\n| L4-GPU  | 674     | 3731    | 7249  |\n|  CPU      |  854     |  4671  | 8887   |\n\nNote that our WDHA algorithm can naturally be parallelized because the core operations, such as the convex conjugate for the $n$ dual programs (the primary computing bottleneck), can be solved independently on high-performance computing (HPC) clusters (CPU/GPU). Our current implementation is not optimized for HPC-acceleration since the core code is written in C++ with `for` loops. It is possible that the WDHA algorithm can fully leverage the power of HPC frameworks (e.g., CUDA) and substantially benefit from parallelization. To engineer an HPC-accelerated version of WDHA is a valuable future direction.\n\n>For the min-max problem, we can only find for the stationary point but not global opt. Therefore the convergence proof is actually not for the optimal solution for barycenter problems.\n\nIf the distributions are regular enough, then a stationary solution is indeed a global optimum because the barycenter functional is convex in the **linear structure** (Proposition 7.19 in Santambrogio. *Optimal Transport for Applied Mathematicians*) [even though not geodesically convex]. This implies that any global minimizer $\\overline{\\nu}$ is a barycenter of $\\mu_1, \\dots, \\mu_n$ iff $0=id - \\frac{1}{n} \\sum_{i=1}^n  \\nabla \\varphi_{\\overline{\\nu}}^{\\mu_i} $. Moreover, if $\\mu_i$'s have densities, the barycenter functional is strictly convex and $\\overline{\\nu}$ is unique.\n\nNow, suppose that $\\varphi\\_{\\overline{\\nu}}^{\\mu\\_i} \\in \\mathbb{F}\\_{\\alpha, \\beta}$ for all $i$ (i.e., no approximation error due to assumed regularity), then by Lemma 3.2 and Definition 3.3, any smooth stationary point $\\nu$ would satisfy that (Wasserstein gradient of $\\mathcal{F}\\_{\\alpha, \\beta}(\\nu)$)= $ id - \\frac{1}{n} \\sum\\_{i=1}^n \\nabla \\varphi\\_{\\nu}^{\\mu\\_i}$ implies that $\\nu$ is a barycenter. This justifies our convergence result to a first-order stationary point solves the global opt problem. This discussion was provided in Remark 3.7. \n\nA simple example with known ground truth is provided in Appendix, Subsection C.1, wherebarycenter of four uniform round disks is computed under the same setting as in Examples 1 and 2. The true barycenter is known as a round disk in the middle. Our method performs uniformly the best to recover the ground truth.\n\n>how does the dimension influence the algo, can we apply more numerical studies on 3D cases.\n\n Yes, you are right. A key factor impacting computation speed is the total number of grid points, which grows exponentially in dimension and results in the major computation burden for higher dimensional cases. One solution is to use parallel computing on HPC clusters to accelerate via the decoupled dual programs. Review 9GjF raised the same question to give a 3D example. Please find our reply therein.\n\n>How does the grid resolution influence the proposed algo, what's the sample complexity if we hope to extend the method to contionous ones.\n\n A larger number of grid points leads to more accurate approximations but also increases computational cost, creating a tradeoff. Comparing to existing methods in the literature, we demonstrate that a grid size of 500 or 1000 achieves the highest accuracy with the least computation time in Examples 1 and 2. Sample complexity to accommodate discretization error is a challenging question. Given the stability estimate (Lemma 3.1), we conjecture that WDHA sample complexity depends on the smoothness of the continuous dual potentials and the domain dimension $d$ in a nonparametric fashion.\n\n>Please clarify the technical contribution for the theoretical results.\n\n   We acknowledge that Lemma 3.1 is derived based on the results of [Jacobs \\& Leger]. However, the other fundational Lemmas 3.2, 3.4, 3.5 are novel and original. Notably, the bound shown in Lemma 3.4 (see eqn (6)) improves upon existing results in literature regarding the quantitative stability of OT maps. This is achieved by restricting $\\varphi$ to be $\\alpha$-strongly convex and $\\beta$-smooth. Lemma 3.2 establishes the first variation of $\\max\\_{\\varphi \\in \\mathbb{F}\\_{\\alpha, \\beta}} \\mathcal{I}\\_{\\nu}^{\\mu}(\\varphi)$, which is both original and nontrivial due to the presence of maximization. These lemmas serve as fundamental building blocks for Theorem 3.6, integrating techniques from applied mathematics, OT theory, and optimization in Euclidean space.", "claims_and_evidence": "Yes. The claim in the paper is almost self-contained except for the verification of the assumptions. But it will not influence much about its applications.", "methods_and_evaluation_criteria": "Yes. The method makes sense that the proposed method design optimization for the min-max problem with gradient in H1 space.", "theoretical_claims": "Yes. I check for the proof of the convergence of the WDHA under the assumptions and the proof for the computational complexity.", "experimental_designs_or_analyses": "Yes. An 2d barycenter problem is designed for different shapes. A table might be reported for better visulization. If possible, the authors also need to clarify which algorithm uses GPUs or all of them used. I am not sure whether GPU versions are used for the methods compared.", "relation_to_broader_scientific_literature": "The computable method for W2 barycenter is very significant. For my expertise about the differential privacy, I believe it is deeply related to the data privacy level and will better solve the problem in the field. More effective application about the data processing, contol or editing based on SDE formulation can also be reduced. More applications in 3D geometry (like basic morphing) could be improved.", "essential_references_not_discussed": "Almost all the possible and important solutions for w2 barycenter are referred, except for sampling  methods like \"Sampling From the Wasserstein Barycenter\" [Daaloul 2021]).", "other_strengths_and_weaknesses": "Pros:\n1. The barycenter problem is significant and the method is novel and computation effective in 2D. \n2. Writing is great and clear for understanding.\n\nCons:\n1. For the min-max problem, we can only find for the stationary point but not global opt. Therefore the convergence proof is actually not for the optimal solution for barycenter problems. But as mentioned in the content, it is not a problem since the complexity. But if more simple examples for the global optim equals to local should give more insights to the practitioner.\n2. I'm not sure how the dimension influence the algo. Actually, for most of the barycenter computation, high dimension will bring problems. I think the proposed method also suffers from the problem. But I think it is also not a problem. More data dimension reduction can be adopted. Actually, what I care about  is 3D cases but only 2D showed. I hope a fast 3D computation can be achieved.", "other_comments_or_suggestions": "Line 318 No Theorem 1.", "questions_for_authors": "1. The same above. how does the dimension influence the algo, can we apply more numerical studies on 3D cases.\n2. How does the grid resolution influence the proposed algo, what's the sample complexity if we hope to extend the method to contionous ones.\n3. Please clarify the technical contribution for the theoretical results. I am not whether there are novel theoretical tools or insights included. I think the main bonus is coming from the computation complexity in H1 sapce but it is referred in [Jacobs $ Leger et al.].I do not think it is necessary yo keep very big novelty in techniques for the paper to solve the barycenter problem. I just hope to justify.\n\n\n\n## Update after Rebuttal\n\nThanks for the authors' efforts. Most of my concerns are well addressed.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "TfBQnbO3bE", "reviewer": "Reviewer_rvS7", "summary": "This paper investigates the Wasserstein barycenter problem. The approach involves transforming the problem into the dual formulation of Kantorovich's optimal transport problem, which is then reformulated as a nonconvex-strongly concave minimax optimization problem. To solve this, the standard gradient descent-ascent method is employed, with theoretical guarantees provided for its convergence. Numerical experiments demonstrate the effectiveness of the proposed algorithm.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments and careful reading. We will fix minor errors in the revision. Please find our response to your main questions below.\n\n>The author should emphasize that the problem is not merely concave but strongly concave, as this distinction is crucial for ensuring the validity of the derived results. Additionally, the complexity bound introduced in line 210 does not hold for general nonconvex-concave problems and should be clarified accordingly.\n\n Thanks for pointing out that our algorithmic convergence rate and the existing Euclidean result rely on strong concavity. Indeed, Theorem 1 works for  $0 < \\alpha < \\beta < \\infty$, under which Problem (4) is a nonconvex-**strongly**-concave problem. In the revision, we will clarify in line 106 to \"as in the Euclidean **nonconvex-strongly-concave** optimization problems\", revise line 207 to  \"(Lin et al., 2020) proved that, **if $f$ is convex in $x$ and strongly concave in $y$**, with suitable choices of step sizes $(\\eta, \\tau )$, the following bound holds:\", and revise line 233 to \"Fix $\\nu$ and $0 < \\alpha < \\beta < \\infty$, the objective functional $\\mathcal{J} (\\nu, \\varphi_1, \\dots, \\varphi_n )$ is **strongly concave** in each $\\varphi_i$.\" to emphasize that the dual program is strongly concave.\n\n>The algorithm tested in the numerical section differs from the one analyzed in the theoretical parts. In addition to the discretization of the function, the practical implementation approximates the projection onto the convex function space using the convex envelope. To enhance the rigor of the analysis and the credibility of the numerical results, it would be beneficial for the paper to either provide a more precise theoretical justification for these approximations or conduct numerical experiments on the exact algorithm analyzed.\n\nResults produced by projection and conjugate envelope (i.e., double convex conjugates) are identical if $\\alpha=0$ and $\\beta=\\infty$, and are close if we set $\\alpha$ small and $\\beta$ large. In Appendix, Subsection C.2, we compared the two algorithms on 1D distributions, where we set $\\alpha=0.001$ and $\\beta=1000$ in Algorithm 3 (projection), which performs comparably well as Algorithm 4 (convex envelope) for this example, i.e., no significant difference. Specifically, projection achieves $\\mathcal{W}_2$-distance $7.610\\times 10^{-5}$ (std: $3.367\\times 10^{-5}$) and $L^2$-distance 0.608 (std: $0.194$), while convex envelope has $\\mathcal{W}_2$-distance $7.771\\times 10^{-5}$ (std: $3.32 \\times 10^{-5}$) and $L^2$-distance 0.6434 (std: $0.451$).\n\nTo compute the projection for 2D distributions, one can use the algorithm proposed by Simonetto. *Smooth Strongly Convex Regression*. However, this involves solving a large convex quadratically constrained quadratic problem. For 2D distributions discretized on the grid of size $1024^2$, this problem would have $1024^4$ quadratic constraints, potentially taking weeks to solve. Due to the rebuttal time constraint, we are unable to report numeric results of running projection in 2D.\n\n>In line 253, I would like to understand why the projection operator onto the convex gradient smooth function class is guaranteed to be unique...\n\nThis projection is defined as $\\mathcal{P}\\_{\\mathbb{F}\\_{\\alpha,\\beta}}(\\varphi) = \\arg\\min\\_{\\psi\\in\\mathbb{F}\\_{\\alpha,\\beta}} \\\\| \\varphi - \\psi  \\\\| \\_{\\dot{ \\mathbb{H}^1}}$ for every $\\varphi\\in\\dot{\\mathbb{H}}^1$. First, we note that $\\mathbb{F}\\_{\\alpha,\\beta}$ is a convex set. In fact, For every $\\varphi_0, \\varphi_1\\in\\mathbb{F}_{\\alpha,\\beta}$, by definition, we have\n\n$$\n\\frac{\\alpha}{2} \\\\| x-y \\\\|_2^2 \\leq \\varphi_i(x) - \\varphi_i(y) - \\langle\\nabla\\varphi_i(y), x-y\\rangle \\leq \\frac{\\beta}{2} \\\\| x-y \\\\|_2^2,\\quad\\forall\\,i=0,1.\n$$\n\nThis implies that for every $ \\lambda \\in [0, 1]$, with $\\varphi_\\lambda = \\lambda \\varphi_1 + (1-\\lambda) \\varphi_0$, it holds that\n$$\n\\frac{\\alpha}{2} \\\\| x-y \\\\|\\_2^2 \\leq \\varphi\\_\\lambda(x) - \\varphi\\_\\lambda(y) - \\langle \\nabla \\varphi\\_\\lambda(y), x-y\\rangle \\leq \\frac{\\beta}{2} \\\\| x-y \\\\|_2^2.\n$$\n\nSo, the convexity of the set $\\mathbb{F}\\_{\\alpha,\\beta}$ is proved. Next, we will show that the projection is unique, which follows from the standard argument for projection to convex subset in Hilbert spaces. For $\\varphi \\in \\dot{\\mathbb{H}}^1$, if both $\\psi\\_0$ and $\\psi\\_1$ in $\\mathbb{F}\\_{\\alpha,\\beta}$ minimizes $\\\\| \\cdot - \\varphi \\\\| \\_{\\dot{\\mathbb{H}}^1}$, by taking $\\widetilde \\psi = (\\psi \\_0 + \\psi\\_1)/2 \\in\\mathbb{F}\\_{\\alpha,\\beta}$, we have\n$$\n\\\\| \\widetilde\\psi - \\varphi \\\\| \\_{\\dot{\\mathbb{H}}^1}^2 = \\frac{1}{2} \\\\| \\psi\\_0 - \\varphi \\| \\_{\\dot{\\mathbb{H}}^1}^2 + \\frac{1}{2} \\\\| \\psi\\_1 - \\varphi \\\\| \\_{\\dot{\\mathbb{H}}^1}^2 - \\frac{1}{4} \\\\| \\psi\\_1 - \\psi\\_0 \\\\| \\_{\\dot{\\mathbb{H}}^1}^2.\n$$\nThe optimality of $\\psi\\_0$ and $\\psi\\_1$ implies that $\\psi\\_0 =\\psi\\_1$. Thus, the uniqueness of projection is proved.", "claims_and_evidence": "Yes, it is clear.", "methods_and_evaluation_criteria": "The algorithm is tested on standard machine learning tasks.", "theoretical_claims": "The overall theoretical claims are reasonable. The specific dual formulation of Kantorovich's problem allows the original Wasserstein barycenter problem to be verified as a strongly concave minimax problem. The author should emphasize that the problem is not merely concave but strongly concave, as this distinction is crucial for ensuring the validity of the derived results. Additionally, the complexity bound introduced in line 210 does not hold for general nonconvex-concave problems and should be clarified accordingly.", "experimental_designs_or_analyses": "The algorithm tested in the numerical section differs from the one analyzed in the theoretical parts. In addition to the discretization of the function, the practical implementation approximates the projection onto the convex function space using the convex envelope. To enhance the rigor of the analysis and the credibility of the numerical results, it would be beneficial for the paper to either provide a more precise theoretical justification for these approximations or conduct numerical experiments on the exact algorithm analyzed.", "relation_to_broader_scientific_literature": "This paper is a valuable supplement to the existing literature on computing the Wasserstein barycenter through its reformulation approach.", "essential_references_not_discussed": "Not found.", "other_strengths_and_weaknesses": "Please see above.", "other_comments_or_suggestions": "1. line 215, it should be ''Let''.\n2. Several citations in the paper, such as \"Algorithm 3.3\" and \"Theorem 1,\" are incorrect and should be properly referenced.", "questions_for_authors": "In line 253, I would like to understand why the projection operator onto the convex gradient smooth function class is guaranteed to be unique. Could the authors provide further clarification or a formal justification for this claim?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In the paper, the authors propose a Wasserstein-Descent H1-Ascent (WDHA), a near linear time and linear space complexity algorithm, for computing the exact barycenter when the input probability measures are discrete on an $m$-point grid.

All the reviewers are positive about the contributions of the papers, including: (1) the novelty of the proposed WDHA algorithm that has solid theoretical guarantee; (2) the experiments are sufficiently extensive and appropriate to evaluate the proposed algorithm; (3) the writing and presentation of the paper are good. After the rebuttal, most of the concerns of the reviewers were addressed, and all the reviewers are happy with the current stage of the paper.

In my opinion, the contributions and originality of the proposed algorithm WDHA are sufficient for acceptance at ICML. Therefore, I recommend accepting it in its current form. However, I encourage the authors to address the reviewers’ suggestions and integrate their feedback into the camera-ready version of their paper.
Decision: Accept (poster)