Sobolev Gradient Ascent for Optimal Transport: Barycenter Optimization and Convergence Analysis
Authors: Kaheon Kim
Bohan Zhou
Changbo Zhu
Xiaohui Chen
Venue: ICLR 2026 Poster (2026)
Keywords: optimal transport; Wasserstein barycenter; concave dual; gradient ascent;
Primary area: optimization
TL;DR: This paper introduces a constraint-free, concave dual formulation for the Wasserstein barycenter, and proposes a scalable Sobolev gradient ascent algorithm with convergence guarantee.

Abstract: This paper introduces a new constraint-free concave dual formulation for the Wasserstein barycenter. Tailoring the vanilla dual gradient ascent algorithm to the Sobolev geometry, we derive a scalable Sobolev gradient ascent (SGA) algorithm to compute the barycenter for input distributions supported on a regular grid. Despite the algorithmic simplicity, we provide a global convergence analysis that achieves the same rate as the classical subgradient descent methods for minimizing nonsmooth convex functions in the Euclidean space. A central feature of our SGA algorithm is that the computationally expensive $c$-concavity projection operator enforced on the Kantorovich dual potentials is unnecessary to guarantee convergence, leading to significant algorithmic and theoretical simplifications over all existing primal and dual methods for computing the exact barycenter. Our numerical experiments demonstrate the superior empirical performance of SGA over the existing optimal transport barycenter solvers.

Published as a conference paper at ICLR 2026
SOBOLEV GRADIENT ASCENT FOR OPTIMAL TRANS- PORT: BARYCENTER OPTIMIZATION AND CONVER-
GENCE ANALYSIS
Kaheon Kim*
Bohan Zhou'
Department of ACMS University of Notre Dame Notre Dame, IN 46556 kk 1.mm2 wend.endd.e..
Department of Mathematics University California
Santa Barbara, CA 93106
bhzhou @ucou edu
Changbo Zhu Department of ACMS University of Notre Dame Notre Dame, IN 46556 ezhu4@nd. edu
Xiaohui Chen
Department of Mathematics University of Southern California Los Angeles, CA 90089 xi wwc.huicuuscceeecc
ABSTRACT
This paper introduces a new constraint-free concave dual formulation for the Wasserstein barycenter. Tailoring the vanilla dual gradient ascent algorithm to the Sobolev geometry, we derive a scalable Sobolev gradient ascent (SGA) algorithm compute the barycenter for input distributions discretized over regular grid, Despite the algorithmic simplicity, we provide global convergence analysis that achieves the same rate the classical subgradient descent methods for minimizing nonsmooth convex functions in the Euclidean space. A central feature of our SGA algorithm is that the computationally expensive c-concavity projection operator en- forced on the Kantorovich dual potentials is unnecessary to guarantee convergence, leading to significant algorithmic and theoretical simplifications over all existing primal and dual methods for computing the exact barycenter. Our numerical exper- iments demonstrate the superior empirical performance of SGA over the existing optimal transport barycenter solvers.
INTRODUCTION
Recent years have seen an increasing trend of interest in applying the theory of optimal transport (OT) in statistics and machine learning. these data-driven fields, a fundamental question is how to average? Averaging is useful to find common patterns in the input data and to summarize the algorithm output from multiple data sources. Unlike the tabular data, where each row a data point with features being viewed as a vector in an Euclidean space, information in data with complex and often geometric structures (such as 2D images and 3D shapes) cannot be effectively extracted with conventional techniques. A pervasive OT analog of the mean for collection of distributional input data Mi..... Pn as probability measures on IRd is the Wasserstein barycenter introduced by Agueh & Carlier (001)) which is defined any minimizer of the following barycenter functional
B(v)
A
Ci W (H4.
2
where a,'s are positive weights such that E a₂ = 1 and v) is the 2-Wasserstein distance between two probability distributions y and p. This variational formulation of "middle point" in Problem I a natural extension of the variance functional for Euclidean random variables, and is
joiint first author

---

Published as a conference paper at ICLR 2026
applicable to general metric spaces. As in the conventional setting, the Wasserstein barycenter has been used in several machine learning tasks such as modeling the centroids in clustering measure- valued data (Zhuang et al. 2022 Ho et al. 2017) and aggregating distributional outputs from subsets of a massive dataset in Bayesian methods (Srivastava et al. 2018).
It remains a major challenge compute the Wasserstein barycenter in tractable (or even scalable) way with principled theoretical guarantees in the literature. For high-dimensional distributions, computing or even approximating the barycenter is worst-case NP-hard (Altschuler & Boix-Adsera 2022). In light of the curse of dimensionality obstacle, we focus on computing the barycenter for
fixed dimension d in this paper. There are various primal-dual reformulations of the barycenter
functional in (1) that lead to different optimization methods -Alvarez-Esteban et al. 2016 Zemel &
Panaretos 2019 Zhou & Parno 2024; Kim et al. 2025) (cf. Literature review in Section
1.1 LITERATURE REVIEW
Barycenter computation can be broadly categorized into two types, depending on whether the optimization problem I is solved in the primal space or dual space. Carlier, Guillaume et al. (2015); Ge et al. (2019) noted that the Wasserstein barycenter problem can be recast as linear program (LP), which, despite being tractable, has poor scalability for large-scale problems (e.g., high-resolution grids or dense point clouds). Alvarez-Esteban et al. (2016) proposed a method directly solving the fixed point iterations on the primal variable in (I). which empirically works for the dnration-scatten family of distributions without quantitative convergence guarantee. First-order methods such as the Wasserstein gradient descent were considered in (Zemel & Panaretos) 2019; Chewi et al. 2020). In each iteration, one needs to compute m OT maps, a highly undesirable feature for problems with a large (or even moderate) number of marginal distributions. Quantitative convergence rate is only available for Warress-sssrsttinn barycenter, when all marginals are Gaussian distributions (Chewi al. 2020).
On the other hand, a dual approach based on the multimarginal OT (MOT), originated from (Gangbo & Swięch 1998), was proposed by /Zhou & Parno (024), where the MOT problem was represented by fully connected undirected graph. This approach avoids the computation of OT maps in the primal problem, however with the cost of introducing ((m²) constraints in the dual problem, thus resulting an algorithm with per-iteration time complexity of Om x log n) for input distributions observed on a grid of size n. Moreover, there is no convergence guarantee established for this approach. More minimax formulation that leverages the benefits from both the primal and dual structures. The resulting Wasserstein descent -ascent (WDHA) algorithm achieves a convergence rate of 0(2)1/2) for T iterations under a fairly expensive convex and smooth projection operation with a quadratic time complexity O(n²) in grid size n per iteration. In practice, the projection is replaced with a conjugate envelope operation, which reduces the per-iteration complexity to On n log n). Nonetheless, this substitution introduces gap between the theoretically analyzable algorithm and the one that is computationally feasible in practice.
recently, Kim et al. (2025) proposed nonconvex-come
It is worth mentioning that there is an extensive literature on regularized notions and algorithms for OT barycenters (Cuturi & Doucet 2014 Bigot et al.) 2019 Janati et al. 2020 Li al. 2020 Carlier et al.) | Fan et al. [022]] Chizat [2033] Chi et al. [023)] Noble et al. [023] Kolesov et al. 2024b Li & Chen 2025). Nevertheless, we emphasize that the current paper is mainly concerned
with the computation of the exact Wasserstein barycenter without any regularization. High-accuracy Wasserstein barycenter has seen applications in, for example, medical imaging (Gramfort et al. 2015)
and scalable Bayesian inference (Srivastava et al. 2018)
1.2 CONTRIBUTION AND LIMITATION
The goal of this paper is to compute the exact barycenter for a collection of input distributions with densities discretized over a regular grid. We propose a novel constraint-free and concave formulation of the Wasserstein barycenter optimization problem that achieves strong duality. Our barycenter formulation fully operates in the dual space, thus avoiding the OT map computations in the primal problem. This unconstrained concave structure allows to derive computationally simple and efficient gradient method in an appropriate dual geometry. Remarkably, we show that, without the computationally expensive e-concave projection steps (or equivalently, the convex conjugate
2

---

Published as a conference paper at ICLR 2026
projection) in existing algorithms, proposed Sobolev gradient ascent (SGA) algorithm retains a strong algorithmic convergence guarantee. In particular, we prove a global convergence rate that matches the rate of classical subgradient descent methods for minimizing nonsmooth convex functions in the Euclidean space (Nesterov) 2013) Lan 2020). We report via numerical studies on 2D and 3D examples demonstrate the superior performance of SGA over all existing exact Wasserstein barycenter algorithms.
On the other hand, we would like mention some computational limitation of the current approach. The SGA algorithm is mainly designed to compute the Wasserstein barycenter of 2D and 3D distributions discretized over the grid in compact domain. Since the grid size scales exponentially in the domain dimension, how to make the SGA approach viable to sigher-dimensional domains is an open problem that we leave to future research.
However, to the best of our knowledge, all existing high-accuracy exact barycenter algorithms are limited to 2D or 3D problems. The entropic methods implemented in the POT library Fllaaay et al.) (2021) are limited to the computation of barycenters for 2D distributions. For higher dimensional probles, existing literature either rely on nerual network based methods (Fan et al., (2025] Kolesov et al. 2024a) or target regularized notion of barycenter (Chi et al. 2023), which trade numerical accuracy for scalability.
To extended our approach from regular to non-uniform grids, three components require modification. First, the fast c-transform must be replaced, for example, by adapting the parametric Legendre transform alogrithm by Hiriart- Hiriart-Urruty & Lucet (2007). Second, the FFT-based Poisson solver can be substituted by multigrid methods. Third, the computation of the pushforward measure, currently obtained by pushing forward discretized marginals under a discretized map, would also need multigrid techniques to handle the Jacobian. Since multigrid methods are substantially more expensive than FFT-based Poisson solvers, higher computational costs are expected.
1.3 NOTATIONS
Given u, Rd, we denote the standard vector inner product and the Euclidean norm by (u, v) and ||l||2 = / (v, v) respectively. Let C Rd be a compact and convex set. The set of probability measures 52 with finite second moments is denoted by P2)).. For any μ E P2(S2), T4 represents the pushforward of μ by T : Ω - Q. Finally, we use Ta- to denote the optimal transport map that pushes to For any f Ω R and continuous symmetric function the o-transform of f. denoted as f Ω + R, is defined as (x) := inf y) f(v)- In this paper, we focus on the case that c(x, y) == ||z /113/2. The homogeneous Sobolev space HI! (52) (f + R| Jo fdr = and []f la < oc) is a Hilbert space equipped with the H inner product (f. == (V (x), gg(x)) dæ and Hil -norm || from the standard gradient V/, we denote the -rrddiideent of the functional H - IR as VI, which is identified by the Riesz representation thcorem via (VI(f), h) (h), where EI, (h) is the Gateaux derivative of the functional we also use μ to denote its density, and define || = maX ((z). For any two probability densities μr, M, the dual of H -oorm is defined as || - vlla == max { & ddp-v) -llllg & 1}.
Ω x Ω R.
ff To distinguish it
ff e IH' in the direction With slight abuse of notation,
Further information on HHI --orm can be found in Section 5.5.2 of Santambrogio (2015).
2 PRELIMINARY. OPTIMAL TRANSPORT COMPUTATION
A function f: + R is called c-concave if = g for some Ω + R. In this case, f = For any μl, V E P22(2) that are absolutely continuous with respect the Lebesgue measure, the optimal transport T can be obtained by solving the Kantorovich dual problem,
{I() (I()==/A fdu+ ('dd).. f
max fre-concave
where the supremum is taken over all e-concave functions. Let f be a maximizer, we have
To = T' and Ty-to = Tÿ, where Ti = id - Vh.
(3)
3

---

Published as a conference paper at ICLR 2026
The optimal transport map T N--, Tu-p also called the Brenier maps (Brenier 1991). Moreover, W2(H, 2) [2](f) For more details, we refer the reader to Santambrogio (2015).
Our barycenter algorithm in Section 4 is motivated by the exact computation of the two-marginal optimal transport problem with distributions discretized over a grid of size n. The Iii -gradien based approach proposed by Jacobs & Léger (2000) a popular algorithm due to its numerical stability and effectiveness handling large n.. In particular, the H -gradiein of I(f) is defined
VI()) 9:
where g is the solution to the Neumann problem (Santambrogio 2015 Subsection 5.5.3).
Ag
(Ty=) in Ω on O2.
0g 0,
Then, Algorithm consisting of a H' gradient ascent step and a double c-transform step, has been introduced to solve 2 (Jacobs & Léger 2020). Numerically, both VI(() and fce can computed with time complexity log n) and space complexity (n)), assuming the probability measures have known densities on a grid of size n. In contrast, the exact approach via linear programming for computing optimal transport between two discrete distributions has a time complexity of ((n²)) and
a space complexity of Onn2) (Kuhn, 1955 Bertsekas & Castanon 1989
Luenberger & Ye 2008
In practice, Algorithm[[ [anoi be enhanced by a back-and-forth approach (Jacobs & Léger 2020). where
the idea is to apply the H' gradient ascent step and the double c-transform step iteratively by hopping between Problem 2 and its twin problem, in order to bound the Hessian in one problem by its twin. The double c-transform is employed for two main reasons: (1) applying the double c-transform to ((t- does not decrease the objective value, i.e., I(J( I((
> I(f((=)) =);;;)); (2) The double c-transform step ensures that { produced by Algorithm I all c-concave. The functions f and fcr share the same regularity as c(x,, y) = || y 1//2, 2, which is generally not the case for f itself.
Algorithm 1 Constrained lill- -Grdiient Ascent Algorithm
Initialize f()); for t = 1, 1,2,
do
f( 1) = =(tt + m VI(((t-));
f(() = (f(1 c)c,
end for
Issues. Despite the empirical advantages of applying the double c-transform, it also has theoretical drawbacks, most notably, it is not an orthogonal projection, and consequently, for any c-concave func- tion g. the inequality| _e - gllia < IIلʃ - glliz does not necessarily hold. As a result, Algorithm is a constrained algorithm and the global convergence guarantee is still open, despite the fact that I: H
IR is concave functional.
3 CONSTRAINT-FREN SOBOLEV GRADIENT ASCENT
Before moving to the barycenter problem, we first discuss a simplified unconstrained H -grrddiie ascent algorithm for the two-marginal OT computation. This algorithm provides the backbone structure for developing our new coordinate-wist version of the barycenter optimization algorithm in Section
Our starting point is the key observation that the first variation of Z(f), and consequently VI()), exist even when is not c-concave (cf. Proposition 2.9 in (Gangbo 2004)). and the computation of ((((()) each iteration in Algorithm l depends solely [[(o)== Recall feee f for any continuous function f. follows that V VI((()) remains unchanged regardless of whether the double c-transform step applied. Therefore, the Kantorovich dual problem (2) can be relaxed to be optimized among continuous functions while retaining the same maximum value (Santambrogio) 2015 Section 1.2); that is.
arg max (I(J) : f c-concave) arg max max(I(f) f continuous).
4

---

Published as a conference paper at ICLR 2026
These together allow to consider the unconstrained problem maxy I(f) (beyond continuity) and to propose the following one-step Sobolev gradient ascent method.
Algorithm 2 Sobolev Gradient Ascent (SGA) Algorithm
Initialize f();;
for 1 1,2 ......... T do
(() = -(t-1) + 1 VI(((-D));
end for
The proposed one-step Algorithm 2 enjoys the following global convergence guarantec, with a rate that aligns with subgradient descent methods for minimizing nonsmooth convex functions in the Euclidean space; see Theorem 3.2.2 in (Nesterov) 2013) or Chapter 3.1 in 2020).
Theorem (Convergence rate for SGA). Let { be the sequence computed from Algorithm Assuming that and v are absolutely continuous with respect to the Lebesgue measure, and f Pe= are continuous on Q. Then,
11 1100-f11
fl + + Et=1 1/1-
I(f) (besst)))
2 2 71-1
where M max (V((((-))) lia is any c-concave maximizer and f(best) is the best solution found, i.e., I[f(bestt) 2 max I(f(9).
Remark 2. The boundedness of M plays the role the Lipschitz condition in Euclidean space when the objective function lacks stronger smoothness. note that [| = =1, computed in the implementation
(Tige-nyy)) + a- of Algorithm for the 2D densities in Section 5.7 are reported in Table 3 (Appendix)) They decrease with t, thus supporting the boundedness assumption on M and indicating convergence to the target solution. The continuity assumption the trajactory of updates [T= is mild and may be inferred from the boundedness condition M, depending on dimension d, see Subsection A.3 in the Appendix for more details.
The values Tf := 1|1 (T)ge-oy))gz11
4 WASSERSTEIN BARYCENTER COMPUTATION
Now, we are ready to introduce new constraint-free concave dual formulation for the Wasserstein barycenter problem as follows:
Σ Oi f du: + 0m / füix dpm
{D(fF...... fm-1)
max famnfm-1
50
where fmix Ei= 1- fi and the supremum over f1...., f 1 is unconstrained beyond being This formulation also facilitates the analysis of the signed Wasserstein barycenter problem (Jacobs & Zhou 2026), in which negative weights are permitted.
continuous (cf. Lemma S).
First, we show that the functional D is jointly concave.
Lemma 3. The mapping (f-....,fm-))
fm-1) is jointly concave.
The Wasserstein barycenter is unique provided with at least one marginal μ4i that is absolutely continuous (Aguch Carlier 2011 Kim Pass 2017). Our next result characterizes this unique
barycenter by the maximizer ())
Theorem 4 (Strong duality and barycenter characterization), Assume that fı. pi..... Pm absolutely continuous with respect to the Lebesgue measure. Then the followings are true.
(i) If (f....,f-11) is optimal to Difi...,, fm-1), then for any i 1,2, identitical to e-concave function pl-almost everywhere. Moreover, Ímix := Dme identical to c-concave function Paa- almost everywhere.
is
as
am
Pea

---

Published as a conference paper at ICLR 2026
(ii) The strong duality holds: min B(v) max f1.x/...... D(fi..... fm-1).
Let : fm-1) be e-concave maximizer to D(fI,... fm-1). Then the unique Wasserstein barycenter can be obtained via the formula
(T He (TPo je ) Pmn.
4.1 SOBOLEV GRADIENT ASCENT FOR BARYCENTER COMPUTATION
We propose a Sobolev gradient ascent (SGA) approach for solving Problem (6). Our algorithm to maximize D. detailed in Algorithm 3l admits a global convergence guarantee and meanwhile enjoys a time complexity Om x n log(n)) n per iteration for distributions whose densities are known on the same grid of size To describe the SGA algorithm, we first derive the formula for the Soboley gradient of D in the following lemma.
Lemma (Dual gradient). Assume for 1,2...., that fo are continuous. Then,
(i) The first variation D at fu. denoted as 8D,, can be expressed as 8D, = -aa(((T) +1 p4= (T + Hm))
(ii) The - -raddent of at Jú. denoted as V f. D can be expressed as
- -- a; ((T))) - denotes the negative inverse Laplacian operator with zero Neumann boundary condition.
where (-A)-
Algorithm 3 SGA Algorithm for Barycenter Computation
Initialize (0) i 1,2,
1;
for t = 1,,,.... ,....,Td.. do
for i 1,2,
do
((0) 1
f. D(J (t-1)
+ qh
end for
end for
We show that the SGA Algorithm| enjoys the following global convergence, which is the main theoretical result of this paper.
Theorem 6 (Main theorem:
convergence rate for barycenter optimization with SGA), Let ( :i 1...., - - be the sequence produced from Algorithm 3 Assuming that {p:)7h are absolutely continuous with respect to the Lebesgue measure, and that the functions ( : I,..., 1 are continuous on Ω. Then, we have the following two scenarios.
(i) Constant step size. If the total number of iterations T is fixed a priori, then by setting M m (E I?T (.ee My0+HHi (T (ftle 1) )+ 4mmml
1| p 1/2 and 1p-1 = (Em-) II (0)
fug /MV// we have
m
D(h....,II fm-1) ((best))))
(best)) < M
1s - AIR 1/2
VT'
where (1-11) m- ll (1-1)
Jm-1) is any c-concave maximizer of D, and
am
D( (best) ((best)
max D(f(t) ...., ...., f(ll1)
(ii) Annealing step size. If the total number of iterations T is not fixed priori, then by setting T--1 = (EmT) ||S
T.11,,) /22 (MV)), we have
D(f....,,Fm.1)) Fm-1)) -(Ff(bbes)
(best) m-1
< M(
m-
11
F.12 1/2 In(T) +++++ +2
VT

---

Published as a conference paper at ICLR 2026
Similar Remark the continuity assumption of potential functions ( is mild in the lit- erature of optimal transport theory and can be inferred from the boundedness condition that 3c. The two schemes on the step size choice in our Theorem 6 strongly resonates with the standard convergence results of subgradient methods for minimizing convex nonsmooth functions in I Specifically, it is known from Theorem 3.2.2 in (Nesterov, 2013) that the subgradient method with constant step size (equivalently, with fixed iteration number I') is optimal for such problem uniformly in all dimension d and the annealing step size scheme is sub-optimal with an extra In 7 factor. Therefore, we see that Theorem extends the Euclidean convergence rates to the Wasserstein barycenter optmization setting with essential structures.
m [
Remark 7. Yao et al. (2025) derived a proximal method for minimizing the general nonsmooth convex functionals in the space of probability measures based on the idea of discretizing the KI. divergence gradient flow. Their coordinate KI. divergence gradient descent (CKL.D) algorithm operates solely in the primal space, achieving a similar rate to our Theorem D and the classical subgradient methods in the Euclidean space (Nesterov 2013 Lan 2020). However, the computation of CKLGD relies heavily on the specialized form of the objective functional in certain entropic regularized problems by sampling. Even though the functional 4 is convex in the linear structure, it is unclear and highly nontrivial how to derive an algorithm to minimize the Wasserstein barycenter functional without entropic regularization.
5 EMPIRICAL STUDY
We compare our method with the Debiased Sinkhorn Barycenter (DSB) algorithm (Janati et al. 2020). the Convolutional Wasserstein Barycenter (CWB) (Solomon al. 2015). and the recently proposed Setersttein-Dsceent iiI AAscett (WDHA) algorithm (Kim et al. 2025). For DSB and CWB, we used the implementations provided in the "POT: Python Optimal Transport" library (Flamary et al. (02TT, and for WDHA, we used the implementation available at httpps: / /kahhenkii. github. 1o/ WDHA/] All computations are performed on Dell PowerEdge R6525 server equipped with two 32-core AMD EPYC 7543 CPUs.
5.1 2D SYNTHETIC DISTRIBUTION
In this example, the goal is to compute the weighted Wasserstein baycenter four densities supported on different shapes whose values are know a equally spaced grid of size 1024 x 1024. The stepsizes of WDHA and SGA are set to be Th 0.1 for all t, and the regularization parameter of CWB and DSB is set to be reg- 0.005. The computed barycenters using 300 iterations are shown in Figure 2 The results of WDHA
aequencal
anddom
ener
18-1
22
and SGA with annealing step sizes (17- = 0.5/ 0.5/vD) are provided in Appendix, Subsection B.I Although the barycenters computed by SGA and WDHA show no obvious visual differences between them, they exhibit clearer and sharper details compared to those from CWB and DSB. This is expected, as SGA and WDHA are exact methods based on the -gradint,,,,, whereas CWB and DSB rely on entropic approximations and are not exact. In addition, we report the computation runtime and barycenter values in Table where we apply the back-and-forth approach Jacobs & Léger (2020) to compute the Wasserstein distance between the estimated barycenter and each distribution. SGA achives the smallest barycenter value for all cases. In terms of computation time, SGA requires approximately half the time of WDHA, one-third the time of CWB, and one-seventh the time of DSB. In addition, SGA is advantageous over WDHA in terms of simplicity, faster runtime, and guaranteed global convergence.
22
100
200
250
Numbe--f Intrationn
Figure 1: Empirical convergence rates of par- allel (Algorithm 3), sequential ( (Algorithm{4) and random (Algorithm (5) SGA algorithms.
Algorithm 3 outlines a parallel scheme for computing the Wasserstein barycenter. As in Yao et al., 2024) for coordinate- wise optimization, also report the empirical convergence rate of this parallel scheme alongside two alternative variants: (1) SGA-S: sequential scheme outlined in Algorithm 4 (2) SGA-R: random scheme outlined in Algorithm 5 By default, refer SGA for
7

---

Published as a conference paper at ICLR 2026
Table 1: Barycenter functional value and runnting time for SGA, WDHA, CWB, and DSB for the 21)
synthetic distribution example in Section 5.1 [J] means the smaller number, the better performance.
Bold numbers indicate the best performer among the four methods under comparison.
Barycenter Functional Value (x 10 [4]
Time (in see) [4]
Weights
SGA
WDHA
CWB
DSB
SGA
WDHA
CWB
DSB
(2. 02,0,0, )
3.917
3.940
4.202
3.926
298
625
1346
3643
(1 0,0,0,
3.917
3.923
4.160
3.925
221
569
1328
3679
(2 3,3.0,0)
5.257
5.275
5.559
5.260
257
602
1145
3728
(3::0:13)
(1111
5.923
5.952
6.153
5.936
700
1064
2290
5003
(4:614:3)
5.299
5.313
5.513
5.314
640
1049
2464
5192
(0.0.
1.646
1.649
1.821
1.648
202
550
1262
3798
(1 (,'3.0.0)
5.257
5.268
5.513
5.263
214
537
1256
3885
(1 (1111)
(::3:4:6)
5,964
5.987
6.182
5.977
672
984
2520
5062
5.219
5.232
5.424
5.239
626
961
2515
5076
4'314
(0,0, 31
1.645
1.647
1.841
1.648
206
1255
3664
(0, 2
31 3.3000
3.812
3.833
4.006
3.814
211
550
1247
3675
(0, 3:
1 3.813
3.814
4.023
3.814
209
510
1251
3648
SGA with the parallel scheme. Using (f(hes)) ..... ....,Im-)) plbe))) as an estimation for (0h...,fm-)) we
plot (og(D()1t ((eest)
D(S() | 1)) against in Figure
Algorithm 4 SSGA-S
Initialize (00) ;i= 1,2, 1;
fort t 1,2...., do
for 1,2, . m do
(1) 1 ((1-1) + 11-
f: D ..... (0) (1-1)
(11,5(101) ..... M..,,--P))
end for end for
Algorithm 5 SGA-R
Initialize 5(0) 1,2
1;
for 1.2.... do
Sample 4 uniformly from {1, 1..,....7 m - 1):
end p!( (1-1) + T Vi, D( (1-1) ..... f (1-1))
5.2 3D INTERPOLATION
We apply Wasserstein barycenter algorithms to interpolate between 3D ball and 3D cube, dis- cretized over a grid of size 200 x 200 x 200. Currently, optimal transport computations for 3D distributions not supported by the POT library (Flamary et al. 202T). In this experiment, annealing step size of T/- = 5 10-3 /t is used for both SGA and WDHA. The results of SGA are shown in Figure with barycenter values (in 10-3) of 1.443, 1.577, and 1.455 from left to right. In- terestingly, the WDHA algorithm diverges under this setting, indicating that SGA is numerically more stable. This increased stability may stem from SGA's simpler structure and its global convergence guarantees.
8

---

Published as a conference paper at ICLR 2026
DSB
CWB
Sight Labatito
Pup CLAUI
WDHA
SGA
SigH EAAGANAS
PUP ILLAME
Figure 2: Comparison of weighted Wasserstein barycenter of densities supported different shapes.
uniform Jeneity
form
Figure Interpolation using SGA between 3D ball and 3D cube.
5.3
REAL DATA EXAMPLE
The Wasserstein barycenter can be used to compress meaningful information from video frames. Such compressed representations may be valuable for tasks like object detection, motion track- ing, and behavior analysis. Figure (left) displays 16 frames from a video the electric scooter tracking data, available from Kaggle
www. kwggle..gle............kkkk..e.. www.cotooaaaaeeeeecooo datasets/ In these frames, two prominent moving objects observed: person walking toward the surveillance camera from the top right, and car in the center making a right turn. The static elements include the grass, trees, and road. compute the Wasserstein barycenter of the frames, treating each image probability density. We
Erainingdatapro. httpa//aawweeeeeeeeeeooooooteeee

---

Published as a conference paper at ICLR 2026
7), 0.1 for WDHA and SGA, and reg 0.0005 for CWB and DSB. The resulting barycenters are shown in Figure 4 (right). Among the methods compared, the SGA algorithm appears to provide the most informative representation, clearly capturing the trajectories of the moving objects while preserving a sharper background for the static components with less distortion than the other three methods.
Figure 4: Left half displays 16 surveillance video frames. Right half displays their baryceters computed by CWB (top left), DSB (top right), WDHA (bottom left) and SGA (bottom right).
REPRODUCIBILITY STATEMENT
Detailed proofs of all theorems are provided the Appendix. The code used to produce the simulation results is available as supplementary material. The electric scooter tracking data is publicly available
on (URL (ttp:///wwww..ggggommomas/s/snne trackkngg..
ACKNOWLEDGMENTS
C. 7hu partially supported by NSF DMSS-4112822. X. Chen partially supported by NSF DMS-2413404, NSF CAREER Award DMS-2347760, and a gift from Simons Foundation.
REFERENCES
M. Agueh and G. Carlier. Barycenters the Wasserstein space. SIAM J. Math. Anal., 43(2):904-924, 2011. ISSN 1016-1410. doi: 10. 10.1137/100805411
Jason Altschuler and Enric Boix-Adsera. Wasserstein barycenters are NP-hard to compute. SIAM Journal on Mathematics of Data Science, 4(1):179-203, 2022.
Pedro C. Álvarez-Esteban, E. del Barrio, J. A. Cuesta-Albertos, and C. Matrán. A fixed-point approach barycenters in Wasserstein space. Math. Anal. Appl., 441(2):744-762. 762. 2016. 3.. doi: 10. imw..//..maa.211...4... 1016/j. jmaa. 2016 04.
ISSN
URL htttps:///.i..oottot
L. Ambrosio and N. Gigli. A user's guide to optimal transport. In Modelling and optimisation of flows on networks, volume 2062 of Lecture Notes in Math., pp. 1-155. Springer, Heidelberg, 2013. doi: 10.1007/978 3-642-32166
Cédric Beaulac and Jeffrey S Rosenthal. Introducing high-resolution handwritten digits data set with writer characteristics. SN Computer Science, 4(1):66, 2022.
Dimitri P. Bertsekas and David A. Castanon. The auction algorithm for the transportation problem. Annals of Operations Research, 20(1):67-96, 1989.
10

---

Published as a conference paper at ICLR 2026
Jérémie Bigot, Elsa Cazelles, and Nicolas Papadakis. Penalization of Barycenters in the Wasser- stein Space. SIAM Journal on Mathematical Analysis, 5163):2211-2285,, 2019. doi: 10.1137/ 18M1185065. BRL https::///etoo http:///ooooorgg./// 113701.1111111111 185065
Y. Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Comm. Pure Appl. Math., 44(4))375-411 1991. ISSN 0010-3640.
Guillaume Carlier, Katharina Eichinger, and Alexey Kroshnin. Entropic-Wasserstenn Barycenters: PDE Characterization, Regularity, and CLT. SIAM Journal on Mathematicaal Analysis, 53(5):8880- 5914, 2021. doi: 10.1 00.1137/00M1887222
Carlier, Guillaume, Oberman, Adam, and Oudet, Edouard. Numerical methods for matching for teams and Wasserstein barycenters. ESAIM: M2AN, 196))))61111144 2015, doi: 10.1051/mmmi 1aa.1////mm/221500333
URL http:: org10/.0..o.....oo.ooo 1ww.1155mmmmaa. wttp://o/...../.mmm///...aoo http./11/mm/mmmmmm.m....
Sinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin J. Stromme. Gradient descent algorithms for Bures-Wassersien barycenters. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 1276-1304, PMLR, 09-12 Jul 2020. URL https: https://proceedinngg
press/ wttss//v25//cccc http1///chww11.......mm
Jinjin Chi, Zhiyao Yang, Ximing Li, Jihong Ouyang, and Renchu Guan. Variational wasserstein barycenters with e-cyclical monotonicity regularization. Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7157-7165, 2023.
Lénaïc Chizat. Doubly Regularized Entropic Wasserstein Barycenters. arXiv:2303. 2023.
URL http:://arrav..orr/aass//// 11844
Marco Cuturi and Arnaud Doucet. Fast computation of Wasserstein barycenters. In Proceedings of the 31st International Conference on Machine Learning, pp. 685-693, 2014.
Lawrence C. Evans. Partial differential equations, volume 19 of Graduate Studies in Mathematics. American Mathematical Society, Providence, RI, second edition, 2010. ISBN 978-0-8218-4974.3. doi: 10. 10.009/////001
J. Fan, I. Haasler, J. Karlsson, and Y. Chen. On the complexity of the optimal transport problem with graph-structured cost. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, pp. 9147-9165. PMLR, 2022.
stess/v151/Ean2a..
Jiaojao Fan, Amirhossein Taghvaci, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. In Proceedings of the 42nd International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research, 2025.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021. URL /ttt://::// /jmrr./mrrrr//// htatt::://mml..or//ppppsss wvtp/2222222222/22 451 html
W. Gangbo and A. Swięch. Optimal maps for the multidimensional Konge-Kantorovier prob- lem. Comm. Pure Appl. Math., 51(1):23 1998. ISSN 0010-3640. doi: 10. (00I(C)
Wilfrid Gangbo. An Introduction to the Mass Transportation Theory and its Applications, 2004. URL https: wwwww//www/////m math. ucla. edu
www.uuaw......u.eeoooo wwwp://wwwwwpppp....oeooe.e wtw.:/aatiiooo/ooteccctotcc wwwp/mo..ummmmmmmmm pdf
Dongdong Ge, Haoyue Wang, Zikai Xiong, and Yinyu Ye. Interior-point methods strike back: Solving the wasserstein barycenter problem. Advances neural information processing systems, 32, 2019.
A. Gramfort, G. Peyré, and M. Cuturi. Fast optimal transport averaging of neuroimaging data. In Information Processing in Medical Imaging, pp. 261-272, 2015.
11

---

Published as a conference paper at ICLR 2026
J Hiriart- Hiriart-Urruty Urruty and Yves Lucet. Parametric computation of the legendre- conjugate with application to the computation of the moreau envelope. Journal of Convex Analysis, 14(3):657, 2007.
Nhat Ho, Xuan Long Nguyen, Mikhail Yurochkin, Hung Hai Bui. Viet Huynh, and Dinh Phung. Multilevel clustering via wasserstein means. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, ICML'17, pp. 1501-1509. Jor..org, 2017.
M. Jacobs and F, Léger. A fast approach to optimal transport: the back-and-forth method. Numer. Math., 146(3):513-544, 2020. ISSN 0029-599X. doi: 100000000211-00-10154444 155-88.
Matt Jacobs and Bohan Zhou. The Signed Wasserstein Barycenter Problem, 2026. URL https
Matt Jacobs, Wonjun Lee, and Flavien Léger. The back-and-forth method for wasserstein gradient flows. ESAIM: Control, Optimisation and Calculus of Variations, 27:28, 2021.
Hicham Janati, Marco Cuturi, and Alexandre Gramfort. Debiased Sinkhorn barycenters. In Proceed- ings of the 37th International Conference on Machine Learning, pp. 4692-4701, 2020.
Kaheon Kim, Rentian Yao, Changho Zhu, and Xiaohui Chen. Optimal transport barycenter via nonconvex concave minimax optimization. In International Conference on Machine Learning (ICML), July 2025.
Young-Heon Kim and Brendan Pass. Wasserstein barycenters over riemannian manifolds. Ad- vances Mathematics, 307:640-683, 2017. ISSN 0001-8708. doi: http:///bi.org/11.016ff..ii 2016.11.026. URL https::/ www.:eiceiieeeecccccoco http:///ww...eoeeeeeeeeeeaataaaa
50C01870815304641
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin. Estimating barycenters of distributions with neural optimal transport. In Proceedings of the s1 International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 25016-2500 25041. PMLR, 21-27 Jul 2024a.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Anastasis Kratsios, Evgeny Burnaev, and Aleksandr Korotin. Energy-guie guided continuous entropic barycenter estimation for general costs. Advances in Neural Information Processing Systems, 37: 17:107513- 107546, 2024b.
H. W. Kuhn. The Hungarian method for the assignment problem. Naval Research Logistics Quarterly, 2(1-2):83-97, 1955.
G. Lan. First-order and Stochastic Optimization Methods for Machine Learning Springer Series in the Data Sciences. Springer International Publishing, 2020. ISBN 9783030395681. URL
Pretps://books www.ooo.ooooo.........oo.o...... wttp://booossooddTTdoooo
Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin Solomon. Continuous Regularized Wasserstein Barycenters. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 17755-17765. Curran Asso- ciates, Inc., 2020. URL http:h/tppee//prreeeeengggrrrrrr 2020/F1le/cat 3R00F1aa.cd0034444333330ac4aa333
wwtpppperrrpppppppeppppppp.pp... fi filess/pperr
Pengtao Li and Xiaohui Chen. Multimarginal SchrongginnnnSnSaniiacScSSggiiee Barycenter. arXiv:2502.02726, February 2025. httppppp..//aaail...rrrr/a//22..
David G. Luenberger and Yinyu Ye. Linear and Nonlinear Programming. Springer New York, 2008.
V. Mikha" ilov. Partial differential equations. "Mir", Moscow; distributed by Imported Publications, Chicago, II, 1978. Translated from the Russian by P. C. Sinha.
Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2013.
12

---

Published as a conference paper at ICLR 2026
Maxence Noble, Valentin De Bortoli, Arnaud Doucet, and Alain Durmus. Tree-based diffusion schrödinger bridge with applications to wasserstein barycenters. Advances in Neural Information Processing Systems, 36:55193-55555
2023.
Santambrogio. Optimal transport for applied mathematicians, volume 87 of Progress Nonlinear Difserential Equations and their Applications. Birkhäuser, 2015. ISBN 978-3-319-20827 978-3- 319-20828-2. doi: 10. 100000097-83119-20000
Calculus of variations, PDEs, and modeling.
Justin Solomon, Fernando de Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphs, 34(4), 2015.
Sanvesh Srivastava, Cheng Li, and David B. Dunson. Scalable Bayes via barycenter in Wasserstein space. The Journal of Machine Learning Research, 19(1):312- 346, 2018.
Rentian Yao, Xiaohui Chen, and Yun Yang. Wasserstein proximal coordinate gradient algorithms. Journal of Machine Learning Research, 25(269):1-66, May 2024.
Rentian Yao, Atsushi Nitanda, Xiaohui Chen, and Yun Yang. Learning Density Evolution from Snapshot Data. arXiv:2502. 17738, February 2025. URL
wwtp://s//r/////////////o/.bbb// /2502
17738
Yoav Zemel and Victor M. Panaretos. Fréchet means and Procrustes analysis in Wasserstein space. Bernoulli, 25(2):932-976, 2019. ISSN 1350-7265. doi: 10.31 URL https
//doi.org/10.
Bohan Zhou and Matthew Parno. Efficient and exact multimarginal optimal transport with pairwise costs. Journal of Scientific Computingg 100(1):25, 2024. doi: 10. 00.0007/1109115024-2252222 URL
https: //doi .org/orroo. / 0. 2007/s1090 024-02572-8
Yubo Zhuang, Xiaohui Chen, and Yun Yang. Wasserstein K-means for clustering probability distributions. In Proceedings 36th Conference Neural Information Processing Systems, 2022.
13

---

Published as a conference paper at ICLR 2026
A TECHNICAL DETAILS
A.1 TECHNICAL LEMMAS
In this section, all lemmas are adapted from existing literature under the specific setting where the cost function c: Q x Ω + R is given by c(z, y) = 2 e112 for some convex compact subset Ω E IR'. This assumption allows us to avoid introducing concepts such as the c-superdifferential. Specifically, Lemma
adapted from Gangbo (2004); Jacobs Léger (2020); Lemma 9 from
Jacobs & Zhou (026): Lemma from Brenier 1991 Ambrosio & Gigli (2013).
Lemma 8. Let f E C(2)) be continuous function on a convex compact subset Ω E Rd Then holds
(i) (r) f(y) and fro((x) (x).
(ii) f is differentiable almost everywhere.
(ii)) For any C(9)), we have [|(S + sh) - 11 <|<|||||
(iv) If f is differentible at a, then
(f (a)) = h T (x).
lim e00
Lemma 9. For any Lipschitz continuous function o, constant >> 0, and absolutely continuous probability measure μ, we have
/ ødv + - WZl(up)) / a( 2 n2
min PEP2(()
Lemma 10. Consider E PP((2)) for some convex compuct subset Ω E Rd. If a map T is in the form Ty = =id f for some c-concave function f. then T the optimal map between p and (T # Consequently.
T du= dd / fd((T)-)
W (Tx)) 2
A.2 PROOF TO THEOREM
Since is compact, we can define :: max | oc. If f continuous, then by Proposition 2.9 (ii) in Gangbo (2004), for any 22, Ω we have
(()))) (x2)) < 2R| 22ll2-
This implies that f is 2-- continuous, and hence differentiable almost everywhere on Ω with respect the Lebesgue measure. Since ş and are absolutely continuous with respect the Lebesgue measure, it follows that fo is also differentiable y-almost everywhere. By Proposition 2.9 (iii) of Gangbo (204).. the first variation of I at f, denoted 8I,, is given by
h. / h Ty. To
8I,(1) =
(VI(J)
for any continuous function h on Q. Using the concavity of I, we have the inequality
I(f2) I(fi) (VZ(fi) f2 - fi)
Following standard subgradient descent arguments for non-smooth convex functions in Fuclidean space, cf. (Nesterov) 2013), we have
[|f(0) - -11 + +M² Et=1 nh
I(f) ((bee)))) <
2 2I
1= T7-1
where M max ||I (Tipa 1)
14

---

Published as a conference paper at ICLR 2026
A.3 SUPPLEMENT TO REMARK 2
The following lemma summarize results (Mikha" ilov 1978. Theorem 4 on page 217) and (Evans) 2010, page 366) on the regularity of solution Poisson equation with zero Neumann boundary condition.
Lemma 11. Consider bounded open subset Q2 Rd, assume (S2) with Sa p 0 and aQ E C++2 for certain k 3 then the solution g to the Poisson equation with zero Neumann boundary condition belongs H
Ag Ōg on ZΩ.
in Q
On
Lemma 12 (Morrey's inequality). Consider bounded open subset Q R&d assume g E H
d d
d $
Z;
(k > 4. then E Ck+2-151<, where 7
2
2
2
Here
[T
d
any element in (0.1),
E Z.
2
denotes ceiling integer of d 2-
In Theorem we assume that the trajactory [T] are continuous functions. Recalling the itera- tive update in Algorithm 2 we just need a simple induction to show that VI(f(t-1) is continuous whenever (r-1) is continuous. Note that (I(f(t-1) is solution to the Neumann problem (see equation and equation 5 thus we may combine the above lemmas propose alternative assumption ensuring continuity of the Poisson solution.
Specifically, let f2 C Rdd be bounded open subset. To ensure g used in above lemmas continuous, suffices to require & E ( with fg p dz = 0 and ön E CK++ for any integer k > 2.
For example,
In ID, = -1 is sufficient. The continuity assumption can be replaced by an E C² and VI((()) in
= (Figtsy))#: i remains bounded.
In 2D or 3D, h: = 0 is sufficient. The continuity assumption can be replaced by an E C2 and remains bounded. In particular, given initial distributions μ,, v with L2 densities, one just needs to check the density (Tiytoye) y v to stay La bounded.
/ - (T (())))))
A.4 PROOF TO LEMMA 3
We first show that f + fc is concave. For any f. f, and t E [0, 1].
[11 (1-0) --- + (z]((x) = inf 1
11 v11 - (1 -11 - -((()
PEn
inf (1 t) | - v1ا F(0)
+t x 1---~11 2112 F(w)
yen
> (1 (1-
15

---

Published as a conference paper at ICLR 2026
To show : is jointly concave, just need to show (fi...., (- Di=l
ai fi) concave, as the other terms are separable. am
ai ((1-t)f
(x)
am
m-
inf 2
|| y12
04 ((1 t)f((u) + +f((w))
DEN
Om
(1 --f) t)
inf
1
2
21
m
Σ
a4 fv(
+
inf
x
v11
Σ
f())
yES? 2
am
NEN 2
am
i=1
M-
m
=(1 -t) t)
a4
fs
(z) +t
a3s am
fo
(z).
am
A.5 PROOF TO THEOREM 4
(i) We first observe that c-transform is ordering reverse, i.e, suppose f(x) > g(x) for any E Ω then
1 1 11100 y12 f(v) & inf || y112 g(v) < (xx).
(x) = inf 2
VER
2
Recall for any f. one has fee 3 f and fooe f Suppose fr > f on set of positive measure,
then fmix = a1 For
note that Σ1 mmm1 as fr < fmix = - Eie on f, on a set of positive measure. then
1 Jidun = -f frem dpm:
J fin dum 3 ( finn dum
Consequently, DiJEs fm- > P.Ii...,fmm1 Repeating these steps, we may replace the maximizer (Ji...,, 1m-1) by (fr (ffe,...,, ...., fmc. which is e-concave. As result, one can always select e-concave maximizer, still denoted by (fi...., fm-1).
Assume that Jinx m=- as is not c-concave, that is, Fce. > Smix on a set of positive measure. am
Now we replace by (1) = om fec fmix) < fi, resulting in
(1 > f1:
c())
04 fo
01
(()
f1
am
O'm
am
am
am
2=
Omm
(fee
fuix)) )
siix
jee
Q'm
an
Consequently, D(F(2) (1) to (ff...,,m-11. which yields that f(²) is identitical to c-concave function fi Pr-almost everywhere, while (()) is e-concave. As result, we may select e-concave maximizer, still denoted by (/i,... fm- such that fmix - 11 as f. is identical to c-concave function fo almost everywhere.
fe..... fm-1) > Fe...f....... fm-1). Since (f...., fm-1) is maximizer
am
16

---

Published as a conference paper at ICLR 2026
i We show mi() B(v) max /1. fm....n D fm-1) by noting that
min rEPa(0)
B(v) = min (EPa(9)
03 (22,,
dp / fi dv) + am (2m,,
min (EP)(O)
ai mmx
max hrof fimm
onfi dps + min
asfil dv + aim W2 (1mm) v)
i= a
vEP2(())
i=1
J atfi du, + am /1-
ai f dpm
max
a1
where the last equality is due to Lemma
To complete the strong duality, let us choose e-concave maximizer (f1...., fm-1)
such that finix E am fi is c-concave as well. We will show that D(f1..... fm-1) = B(t), where i (T -TT # #m is defined in (7). Thanks to Lemma 10 f is the optimal map between ,4; and v = (Tzz #H4 and Tje. is the optimal map between am and (T Pm. That is,
mn
04 W (1, 1) = ail S dp d.d.
for i 1, I...., m
2
Omm W (1m, 2
(1m, D) -mml fuix dum + // smix dēj.
Consequently,
m-1
= Σ
pd + am f jem dpm
ai
i=1
1
Σ
fiddus
+ am ( frin dpm
Σ
Oi
W (1:, B(D)
2
i=1
xm1
Σ
a dī am Fiin Inix BBB)
B(D).
where the last equality is due to the fact that Eme-' anJ. amf
-mm= cx
aml Em-l or f.) = 0. am
Now, max D(Si...., fm-.)) > fis..-mm-1
=B())
min B(v), which completes MEP2(0)
the strong duality.
A.6 PROOF TO LEMMA 5
(i) We first caleulate the first variation of
(- 0i Om
Ejti
f
aj fz). am
Let - Ejjti Dem
aj fz. apply Lemma
8
(- as fi + u am
as Eh (- ai f, + w) --(-
lim
Cem
am
O4i h) T am f+tue
0i h Ty
Omm
0mm
17

---

Published as a conference paper at ICLR 2026
Apply this result, the Gateaux derivative &D (h) with fixed other inputs fozt is given by (h)
/ (Js (Ji+ -se
am /
(-an as am
(-
fo + u as ch) (- a: fs + u)
lim a;
dp,
dum
500
=o: / (-ho Ty:) duu + am / ai ho Tyan dpm
am
=- --ai (ad(T(T(t))) (T ( Mm).
Consequently, the first variation öD -a, (TTT # + #m).
(ii) a consequence from the fact that the ii' gradient is defined via I foti), h) we obtain that
jD, (h). Similarly with equation 4 and equation 5
D -- a. -aT((TT))
(T ) Mm
A.7 PROOF TO THEOREM 6
We note that D is a functional over the product space (IH-))--1 == Hi x ... x TH For any
f = ( fm-1) and (91)....mm1)) (H jm- is a Hilbert space with inner product
m-1
(f. (i == (f..))
i=1
We let V (VD (f).... VD -m (f)). Using the concavity of D and the concave inequality
D(g) D(f) - (VD(f). f) payye ym & O, it follows similarly from Theorem and standard analysis of subgradient descent that
Ei= 1
fell + M ΣT n7
D(f....., D(S1 (best)
(best)
<
2Σ 91-1
where M max (2imh
mm az (I () (F # )
(tyde- Mm 1/ 12/2 If we fix the total
number of iterations T a priori, the optimal step size that minimize the right hand side is 7h-1 = (Em-) !1!(0) - ful 3/2/(MVT) This gives
M
11:00 111
D(f..... Fm-1) D(S
(bsa))
VT
If we don't 't fix T priori, then by letting 7h-1
Ei=1
m-
||
Fil 1/2 /////M//// Vt,, we have
m-1
ll f 1 11/2 11/2 logloo((T] +
Bmf....I-11) - D(Jibess elbeat)
< M
VT
(=1
where we used inequlities ΣT 1 < In())) and ΣT
VT.
Vt
A.8 PROOF TO LEMMA 9
Given that y is absolutely continuous, for any P2()),, there exists a measurable map T that pushes μ to v (Santambrogio 2015 Corollary 1.29). Thus,
min PEP2(1)
dd + 2
W (1, =a min
dv + 2
W (M,
PE (EP2(
dl(T)) #4] + / ()(xx))12 dp du
inf
(T(()))
inf
T( du
| w1 اu) | dp = a / a (=))
=a
inf
T=)(m)
18

---

Published as a conference paper at ICLR 2026
Table 2: Barycenter functional value and runnting time for SGA and WDHA using annealing step
sizes for the example in Section 5.1
Barycenter Functional Value (x103) [4]
Time (in sec)
[4]
Weights
SGA
WDHA
CWB
DSB
SGA
WDHA
CWB
DSB
0,0,0 )
3.917
4.012
4.202
3.926
328
497
1346
3643
0,.000,
3,915
3.931
4.160
3.925
215
458
1328
3679
( (1)))
5.917
5.999
6.153
5.936
678
822
2290
5003
(2. (..3.0.0
5.256
5.340
5.559
5.260
280
477
1145
3728
(.111
5.288
5.323
5.513
5.314
639
811
2464
5192
(0,0, 1.646
1.646
1.821
1.648
201
447
1262
3798
(1 (2:2.0,0)
5.256
5.263
5.513
5.263
224
470
1256
3885
5.958
6.019
6.182
5.977
690
822
2520
5062
(1 4'31416
66 41:41344
5.205
5.231
5.424
5.239
646
814
2515
5076
(0,0, 3'
1.646
1.648
1.841
1.648
218
440
1255
3664
(0, 2 0.0) 0)
3.812
3.871
4.006
3.814
211
452
1247
3675
3 3.00
1
3.813
3.817
4.023
3.814
219
433
1251
3648
B ADDITIONAL NUMERICAL RESULTS
B. 2D SYNTHETIC EXAMPLE WITH ANNEALING STEP SIZE
We examine the annealing step sizes (7) = 0.5/ VI) for SGA and WDHA using the synthetic example introduced in Subsection 5.1 The computed barycenters are shown in Figure with corresponding barycenter values and computation times reported in Table 2 SGA continues to outperform the other methods. certain cases, such when weights = (1/4, 176, 1/4, 1/3), visual inspection suggests that the algorithm has not yet converged. This observation aligns with our expectations based on Theorem 6 which indicates that for fixed total number of iterations (T 300), constant step sizes are optimal.
WDHA
SGA
MOPLEALSANME
EHERANSE
CENSIN
Figure 5: Comparison of weighted Wasserstein barycenter of densities supported on different shapes.
19

---

Published as a conference paper at ICLR 2026
Table T x 10 computed the implementation of Algorithm 2 with 2D densities Section 5.1 The first column refers to marginals used compute OT.
(p, v)
100
200
300
(1,2)
2.459
1.943
1.813
(1,3)
1.411
0.952
0.817
(1,4)
2.348
1.915
1.820
(2,
2.320
1.357
(2,3)
0.666
0.397
0.300
(2,4) 1.169
0.737
0.573
(3,1) 2.909
1.792
1.368
(3,2)
2.061
1.236
0.924
(3,4)
1.008
0.593
0.436
(4,
2.930
1.838
1.435
(4,2)
1.715
1.155
0.958
(4,3)
0.593
0.388
0.331
B.2 BOUNDEDNESS OF M IN THEOREM
The radius of the bounded ball (i.e., M in Theorem 1) can be easily computed and the boundedness checked in the process. Further more, r; == || (T ) expected decrease as increases, as it quantifies the distance between and the pushforward measure y/ by Ty If f(t-1) is optimal for I(f)) == Jg f du + + f dı, we have m, O. Empirically, we compute m at different for pairwise 2D densities in Section 5.1, shown in Table []] The results confirm that r, decreases with I, indicating that the boundedness assumption of M true and the solution will converge to the truth.
Algorithm Adagrad-adapted SGA for barycenter computation
Require: Stepsize 0, small constant >> 0 (0)
Initialize ) and r < for i 1, I,...,
for t = 1 T do
for 1..... m
1 do
{1
(1-1))
(1-1)
(1)
(t)
r t Ti
+ 9i
g
(t)
f10) + {( ((-1)
+
(t)
ri +
end for end for
B.3 HANDWRITTEN DIGITS DATA
The handwritten digits data is commonly employed to access the performance of barycenter algo- rithms, see Cuturi Doucet 2014); Ge et al. (2019): Kim al. (2025) among others. We compare the proposed methods on hign-resolution handwritten digit data from Beaulac & Rosenthal (2022), treating each image as probability density supported on [0, 1 Specifically, we compute the Wasserstein barycenter of 10 images of the digit 2, each of size 500 x 500 pixels. For SGA and WDHA, we use a constant step size 1), = 0.1. We also adapt the SGA to the AdaGrad scheme in barycenter computation problem, yielding the SGA-AdaGrad described in Algorithm 6 For AGA-Adagrad we set = 0.1 and 10 8 The regularization parameters for CWB and DSB are both chosen as 5 10-3, The results are shown in Table 4nd Figure 6 Among all methods, SGA,
20

---

Published as a conference paper at ICLR 2026
and similarly for SGA-Adagrad, produce the barycenter with the finest visible texture, while also achieving the lowest barycenter functional value and the shortest computation time.
2
2
2
2
2
2
2
(a) 10 handwritten images of the digit "2".
2
2
2
2
10
10
(b) SGA
(c) AGAaua
(d) WDHA
(e) CWB
(f) DSB
Figure 6: Input handwritten digits (top) and barycenters computed by different algorithms (bottom).
Method
Barycenter Functional Value [4]
Time [J]
SGA
5.992
10-3
406
SGA(Adagrad)
5.992 X 10 -3
409
WDHA
6.049
x 10-3
415
CWB
6.463
10
430
DSB
6.021
10-3
670
Table 4: Comparison of barycenter functional value and computation time for different algorithms.
B.4 GAUSSIAN DISTRIBUTIONS
We consider three Gaussian probability measures f11 /1, /43 R.2, where 14 N(m;, E4),
1,...,3. 3. Here, we set my (40, 40)"" = (50, 50) m3 (20, 50)"" and
2.
0.22
Σg 0.42
Σ1
Σn
0.2
0.6
0.42
We note that methods work on compact domain, thus those distribution truncated large enough square, i.e., [0, 100]2, to make the truncation error sufficiently small. The Wasserstein barycenter jı of /1, M4g is again Gaussian, j N(m, E), with barycenter mean and standard deviation given by
110 140)
(44) 0
m =
and =
3
0.4
21

---

Published as a conference paper at ICLR 2026
We set 0.005 for SGA and WDHA, reg 0.005 for CWB and DSB. We represent all truncated Gaussian distributions on an equally spaced 1024 x 1024 grid, and evaluate the squared 2- Wassersteen distance between the computed barycenter and μ using the Back-and-Forte approach (Jacobs et al.) 2021). In this Gaussian setup, our method(SGA)) returns the smallest squared Wasserstein distance between computed barycenter and the groundtruth barycenter shown in Table The barycenter fucntional value is shown in T 6
SGA
WDHA
CWB
DSB
0.4636
20.3449
21.7867
0.4925
Table Squared 2- Wasserstein distance between the j and the barycenter computed by each method.
Groundtruth
SGA
WDHA
CWB
DSB
88.9067
89.0074
102.0016
109.4954
89.0247
Table 6: Barycenter functional value for groundtruth and different methods.
We compare the approximation error of the optimal transport maps from the barycenter to each distribution for our method and for WDHA, which also relies Kantorovich potentials compute the Wasserstein barycenter. We note that CWB and DSB are entopic regularized methods and does not produce optimal transport maps without further modifications. The approximate optimal transport
maps T; are obtained from the learned Kantorovich potentials f via [x()) = x - (x))
1, 2 and (x) x - fmix(x), with fmix = -f - f2. We evaluate the approximation error between the learned maps T; and the ground-truth optimal transport maps T; from z to 14, Tx(x)= := 1/2 (x 11/
m, + /3>,
m), using the relative L norm T] T, / | where - -= So [+]()) 1 dr. The results provided in the following Table As shown in Table 7 the approximation errors for fú, T2. and T3S under our method are smaller than those obtained with WDHA.
[ -
SGA
WDHA
[ 1 11
6.6620 10
6.7319 X 10
T 12
4.1357 x 10
4.2259 x 10-2
1773 7
4.0582 10
4.0961 x 10
Table 7: Comparison of approximation errors for the optimal transport maps learned by SGA and WDHA. Smaller values (in bold) indicate better approximation.
C THE USE OF LARGE LANGUAGE MODELS (LLMS)
A Large Language Model (LLM) was used solely to identify grammar mistakes, polish language, and detect typographical errors. No part of research ideation, discovery, or substantive content was generated by the L.M..
22


Reviews:
[{"review_id": "BFE1fzvrtQ", "reviewer": "Reviewer_DYXD", "summary": "This paper proposed a constraint-free and concave formulation of the Wasserstein barycenter optimization problem which achieves strong duality.  The paper derived a scalable Sobolev gradient ascent (SGA) algorithm without the computationally expensive c-concave projection steps. Then, it showed that the proposed SGA algorithm achieves the same rate as the classical subgradient descent methods for minimizing nonsmooth convex functions in the Euclidean space. Numerical experiments demonstrate the effectiveness of SGA over the existing optimal transport barycenter solvers.", "questions": "Is the convergence rate for SGA in theorem 6 optimal?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Overall I found this paper to be interesting and provided a good theoretical contribution in computing Wasserstein barycenter efficiently:\n-  the constraint-free concave dual formulation allows the algorithm to avoid the OT map computations in the primal problem\n-  the proposed Sobolev gradient ascent (SGA) algorithm is simple and straightforward, but matches the rate of classical subgradient descent for minimizing nonsmooth convex functions in the Euclidean space.\n- with synthetic and real world data, it was shown that SGA achieved comparable or better performance than existing baselines.", "weaknesses": "- In comparison to the theoretical results, I found the numerical experiments to be on the weaker side, with limited data involved in testing SGA. For real-world data, only one set of video frames was used to demonstrate SGA's performance, and the comparison between WDHA and SGA's performance was not very obvious. WDHA also captured the trajectory of the moving object, but for some reason the picture appeared to be darker. The paper will be strengthened if the real-world data experiments part can be more comprehensive and quantitative.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jQg0oiVxzq", "reviewer": "Reviewer_VvTD", "summary": "This paper proposes a constraint-free concave dual formulation for the Wasserstein barycenter and a dual ascent algorithm adapted to the Sobolev geometry. The algorithm matches the complexity rate of the subgradient method. A key advantage is that the new dual algorithm bypasses the expensive c-concavity projection and achieves high efficiency. The experiments demonstrate the advantages of the proposed method.", "questions": "See the weakness part", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper presents a novel concave dual formulation for the Wasserstein barycenter problem that is unconstrained yet achieves strong duality. The theoretical analysis appears reasonable and expected. Empirically, the authors demonstrate that SGA is significantly faster and more numerically stable.", "weaknesses": "The stepsize rule seems impractical—it requires an M that depends on all iterates. It's somewhat bizarre that the stepsize depends on future iterates. If the stepsize can be chosen more freely, please provide a detailed convergence rate analysis.\n\nAdditionally, would an adaptive gradient method (such as AdaGrad) perform well in your examples?\n\nAs the authors pointed out, the algorithm's complexity scales with the grid size $n$, which scales *exponentially* with the dimension $d$. This effectively limits the current method to low-dimensional problems ($d=2, 3$). This could be a practical constraint for many machine learning applications.\n\nThe real-world example using electric scooter tracking data appears quite simple. Are there more challenging datasets in this area?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OlL5ourhcv", "reviewer": "Reviewer_Vdfa", "summary": "This paper introduces a novel, constraint-free, and concave dual formulation for the Wasserstein barycenter problem. By removing the constraints present in previous dual formulations, the authors develop a scalable Sobolev Gradient Ascent (SGA) algorithm. A key theoretical contribution is a global convergence analysis for the proposed method. The efficacy of the algorithm is validated through numerical experiments that show superior performance over existing methods.", "questions": "Please fix a hanging citation \"Matthew Jacobs and Bohan Zhou. The Signed Wasserstein Barycenters, 2025\". It does not seem to exist?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper focuses on a fundamental problem of Barycenter Optimization. Theorem 4 is the paper's cornerstone. It introduces a constraint-free concave dual formulation for the Wasserstein barycenter problem. This key result makes the problem amenable to optimization via SGA and allows the authors to establish a global convergence guarantee.", "weaknesses": "1. The method's efficiency hinges on a regular grid, which allows for an efficient FFT-based solver for the inverse Laplacian. It would significantly strengthen the paper to discuss the adaptations required for irregular meshes, where one would need to employ more computationally expensive methods.\n2. The grid-based approach is subject to the curse of dimensionality, meaning its computational cost grows exponentially with the number of dimensions. To their credit, the authors explicitly acknowledge this and correctly position their method for low-dimensional (2D and 3D) problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RFiD2jfy8h", "reviewer": "Reviewer_h7bp", "summary": "In this paper,  the authors propose an algorithm for computation of the exact Wasserstein barycenter for a collection of input distributions discretized over a regular grid. The algorithm is based on the idea of Sobolev gradient descent and is tested in several experimental setups.", "questions": "- Could you perform comparison with the ground-truth barycenter and OT maps in the Gaussian setup?\n- Could your approach be extended to the case of more general cost functions?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "In general, the paper is well-written and provides an algorithm supported with theoretical analysis and experimental illustrations in several setups.", "weaknesses": "My main concerns are related to the limited practical usefulness of the developed algorithm. For example, as the authors claim themself in lines 101-105, the developed algorithm is not suitable for the distributions of the dimensions higher than 3D. Besides, the algorithm is designed for the computation of the optimal transport barycenter under the quadratic cost function assumption. This point also could be listed among the limitations of the developed methodology. While I understand that the listed limitations might correspond not to the algorithm itself, but rather to the particular  class of the algorithms for exact barycenter computation, I have other concerns regarding the demonstration of the practicality of the proposed algorithm.\n\nSpecifically, all of the experimental results only show the ability of the method to compute the barycenter of distributions, but do not assess its ability to approximate well the optimal transport (OT) maps between each of the distributions and barycenter. I think that it is an important aspect since the learned maps have more practical use cases than the barycenter itself. Meanwhile, the authors do not perform comparison with the ground-truth barycenters which are known, e.g., for Gaussian distributions. Thus, I kindly suggest the authors to perform comparison of their approach with the ground-truth barycenter and OT maps in the setting with Gaussian distributions which are known thanks to (Chewi et al., 2020). \n\nThe related work section lacks the overview of recent methods for approximating the barycenter of continuous distributions (Fan et al, 2021; Chi et al., 2023; Noble et al., 2023; Kolesov et al, 2024a,b). I think these algorithms deserve mentioning since in lines 106-107 you write that “almost all existing barycenter algorithms are limited to 2D problems” which is not true when we are talking about the methods listed above. Thus, I kindly suggest the authors add the references in the related work sections and make the phrase in lines 106-107 more clear.\n\n**Overall**, I think that the experimental evaluation of the proposed approach  should be improved since now it lacks comparison with the ground-truth barycenters and OT maps which is important for clearly revealing the performance of the developed approach.\n\n**References.** \n\nFan, J., Taghvaei, A., and Chen, Y. Scalable computations of wasserstein barycenter via input convex neu- ral networks. In Meila, M. and Zhang, T. (eds.), Pro- ceedings of the 38th International Conference on Ma- chine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 1571–1581. PMLR, 18–24 Jul 2021.\n\nChi, J., Yang, Z., Li, X., Ouyang, J., and Guan, R. Varia- tional wasserstein barycenters with c-cyclical monotonic- ity regularization. In Proceedings of the AAAI Confer- ence on Artificial Intelligence, volume 37, pp. 7157–7165, 2023.\n\nNoble, M., Bortoli, V. D., Doucet, A., and Durmus, A. Tree-based diffusion schro ̈dinger bridge with applica- tions to wasserstein barycenters. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.\n\nKolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., and Korotin, A. Energy-guided continuous entropic barycenter estimation for general costs. Advances in Neural Information Processing Systems, 2024a\n\nKolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., and Korotin, A. Estimating barycenters of distributions with neural optimal transport. In Forty-first International Conference on Machine Learning, 2024b\n\nSinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin J. Stromme. Gradient descent algorithms for Bures-Wasserstein barycenters. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 1276–1304. PMLR, 09–12 Jul 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)