Wasserstein Iterative Networks for Barycenter Estimation
Authors: Alexander Korotin
Vage Egiazarian
Lingxiao Li
Evgeny Burnaev
Venue: NeurIPS 2022 Accept (2022)
Keywords: optimal transport
continuous barycenter
neural networks
Wasserstein-2 distance

Abstract: Wasserstein barycenters have become popular due to their ability to represent the average of probability measures in a geometrically meaningful way. In this paper, we present an algorithm to approximate the Wasserstein-2 barycenters of continuous measures via a generative model. Previous approaches rely on regularization (entropic/quadratic) which introduces bias or on input convex neural networks which are not expressive enough for large-scale tasks. In contrast, our algorithm does not introduce bias and allows using arbitrary neural networks. In addition, based on the celebrity faces dataset, we construct Ave, celeba! dataset which can be used for quantitative evaluation of barycenter algorithms by using standard metrics of generative models such as FID. 

Wasserstein Iterative Networks for Barycenter Estimation
Alexander Korotin
Vage Egiazarian
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
Skolkovo Institute of Science and Technology
Moscow, Russia
Moscow, Russia
vage. www.aaarraanaakkltccch....
a korotin@skolltech.rr
Lingxiao Massachusetts Institute of Technology Cambridge, Massachusetts, USA lingxiaodmit.edd edu
Evgeny Burnaev
Skolkovo Institute of Science and Technology
Artificial Intelligence Research Institute Moscow, Russia e. burnaer@@sklttech.coo ru
Abstract
Wasserstein barycenters have become popular due to their ability to represent the average of probability measures in geometrically meaningful way. In this paper, we present an algorithm to approximate the Wasserstein-2 barycenters of continuous measures via generative model. Previous approaches rely regularization (entropic/quadratic) which introduces bias or on input convex neural networks which are not expressive enough for large-scale tasks. In contrast, algorithm does not introduce bias and allows using arbitrary neural networks. In addition, based the celebrity faces dataset, construct Ave, celeba! dataset which can be used for quantitative evaluation of barycenter algorithms by using standard metrics of generative models such FID.
Introduction
Wasserstein barycenters [1] provide geometri- cally meaningful notion of the average proba- bility measures based on optimal transport (OT, see (65]) Methods for computing barycenters have been successfully applied to various prac- tical problems. geometry processing, shape interpolation can be performed by barycenters 58. In image processing, barycenters are used for color and style translation 44 texture mixing [50 and image interpolation [32] [56].
ELH
PP
P
P
P
Barycenter
In language processing, barycenters be ap- plied to text evaluation [15). In online learning, barycenters are used for aggregating probabilis-
P argmin
on W...P.)
P
tie forecasts of experts B1 47 27]. Bayesian inference, the barycenter of subset posteriors converges to the full data posterior [59] 60) al- lowing efficient computation of full posterior based on barycenters. In reinforcement learning, barycenters used for uncertainty propagation Other applications are data augmentation
12
Figure 1: The barycenter w.r.t. weights (a1, 02, 03)) ) = (1) 3: { of Ave, Celeba! subsets computed by our Algorithm
The figure shows random samples from the input subsets and generated images from the barycenter.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

[8] multivariate density registration 19]. distributions alignment [24]. domain generalization 39 and adaptation [33. model ensembling averaging of persistence diagrams 63
The bottleneck of obtaining barycenters is the computational complexity. For discrete measures, fast and accurate barycenter algorithms exist for low-dimensional problems; see [48] for survey. How- ever, discrete methods scale poorly with the number of support points of the barycenter. Consequently, they cannot approximate continuous barycenters well, especially in high dimensions.
Existing continuous barycenter approaches 34
30] are mostly based on entropic/quadratic
regularization parametrization of Brenier potentials with input-convex neural networks (ICNNs, see 44). The regularization-baat [CRWB] algorithm by [34] recovers a barycenter biased from the true one. Algorithms [CW2B] by 30 and [SCW₂B] by [9] based on ICNNs resolve this issue, see 30 Tables 1-3]. However, despite the growing popularity of ICNNs in OT applications
40
28
42
they could be suboptimal architectures according to recent study (29). According the
authors, more expressive networks without the convexity constraint outperform ICNNS in practical OT problems.
Furthermore, evaluation of barycenter algorithms is challenging due to the limited number of continuous measures with explicitly known barycenters. It can be computed when the input measures are location-scatte (e.g. Gaussians) 3 $4] or 1-dimensional [10, $2.3]. Recent works (34 30 [9] consider the Gaussian case in dimensions 256 for quantitative evaluation. In higher dimensions, the computation of the ground truth barycenter is hard even for the Gaussian case: it involves matrix inversion and square root extraction
2 Algorithm 1] with the cubic complexity in the dimension.
Contributions.
We develop a novel iterative algorithm ($4) for estimating Wasserstein-2 barycenters based on the fixed point approach by [] combined with neural solver for optimal transport (29]. Unlike predecessors, our algorithm does not introduce bias and allows arbitrary network architectures.
We construct the Ave, celeba! (averaging celebrity faces, $5) dataset consisting of 64 x 64 RGB images for large-scale quantitative evaluation of continuous Wasserstein-2 2 barycenter algorithms. The dataset includes 3 subsets of degraded images of faces (Figure I.. The barycenter of these subsets corresponds to the original clean faces.
Our algorithm is suitable for large-scale Wasserstein-2 barycenters applications. The developed dataset will allow quantitative evaluation of barycenter algorithms at a large scale improving trans- parency and allowing healthy competition in the optimal transport research.
Notation. We work in Euclidean space RD for some D. All the integrals are computed over R unless stated otherwise. We denote the set of all Borel probability measures on H with finite second moment by P). We use (R) C (2 to denote the subset of absolutely continuous measures. We denote its subset of measures with positive density by P2txc(R) C
P2. PRP). We denote the set of probability measures on R.D R.D. with marginals and Q by [IP Q). For measurable map T RD R,,,, we denote the associated push-forward operator by Tz. For o RD + H. we denote by Φ its Legendre- Fenchel transform (20) defined by Φ(v) max (x,,) (xx)]. Recall that Φ a convex function, even when is not.
2 Preliminaries
Wasserstein-2 distance. For P., QRP), Monge's primal formulation of the squared Wasserstein-2 distance, i.c., OT with quadratic cost, is
def min TPE
- 1 (T)
(1)
where the minimum is taken over measurable functions (transport maps) T LRP Rd mapping P to Q. The optimal T is called the optimal transport map. Note that ( is not symmetric, and this formulation does not allow mass splitting. That is, for some P, Q (2 )), there might be map
T that satisfies TVP = Q. Thus, (5 proposed the following relaxation:
W def
min Tell(F.0) (P.()) P 2 ||
{
1
vll dm (2. v).
(2)

---

where the minimum is taken over all transport plans i.e. measures on RD' x IR" whose marginals
are IP and Q. The optimal π [I, Q) called the optimal transport plan. If a² of the form [id. HPP E [T(, Q) for some 7 then T minimizes ① The dual form 64 of W2 is:
w3 (T, = max
[ (((())d((w)) f ((0))())))((()00000)
(())dQ(V))
(3)
where the maximum is taken over u E (P), E (Q) satisfying u(a) o(y) S 2 vll for all n, YE.. The functions u and called potentials. There exist optimal u", v² satisfying
un (v*) where (v)
d
min
[ yll2 f(x)] is the c-transform of f. We rewrite (3 as
zER
max / (x)d((x) + f ((y))(Q())
W2
(4)
where the maximum is taken over all E (²))) is customary [65 Cases 5.3 & 5.17] to define
u(x))= = =|= (x) and v(y) = |lyl|
((x)). There exist convex optimal 1/" and 6² satisfying
- = x0 and 0- = ~." P E (I. )) then the optimal T ( always exists and can be
recovered from the dual solution u² (or 16") of 3 (x)
Vu Vbb" [54 Theorem 1.17]. The map 7 a gradient a convex function, see the Brenier Theorem
Wasserstein-2 barycenter. Let 1P PN E (P such that at least one of them has bounded density. Their barycenter w.r.t weights an aN (an > 0; > an given by MI:
P def arg min PEP2(RD)
an WP,
(5)
n=1 The barycenter IP exists uniquely and P E PP a2. Moreover, its density is bounded 1 Definition
3.6 & Theorem 5.1J. For n E (1, (1., N). J, let Tp P-P. = V be the OT maps from P to P.. The following holds P-almost everywhere:
Σ
an T T
(x)
Σ On (x)
(6)
see 3 53]. I (R then O6) holds for every E RD, i.e, En an (x)
||x||2
We call such convex potentials vi congruent.
3 Related Work
Below we review existing continuous methods for OT. 33.1 we discuss methods for OT problems In 3.2 review algorithms that compute barycenters (5).
3.1
Continuous OT Solvers for W2
We the phrase OT solver to denote any method capable of recovering T or u² (or 16*).
Primal-form solvers based on or 2) c.g., 67 37
parameterize T using complicated generative
modeling techniques with adversarial losses to handle the pushforward constraint TiP Q in the primal form They depend on careful hyperparameter search and complex optimization (38).
Dual-form continuous solvers [1 55 46 62 28 based on 3 or have straightforward optimiza-
tion procedures and can be adapted to various tasks without extensive hyperparameter search.
comprehensive overview and benchmark of dual-form solvers are given in 29 According to the evaluation, the best performing OT solver is reversed maximin solver [M:R]] modification of the idea proposed by [46] in the context of Wasserstein-1 GANs [5. In this paper, we employ this solver part of our algorithm. We review it below.
Reversed Maximin Solver. In
(x) can expanded through via the definition of c-transform:
v1l v(v) dP(x)+ v(y)d(())=
max e
min HERD
2
||
T(x
- (T(x)) dP(x) dp(x)+/ 1001
max min
(7)
T
2
3

---

In 77. the optimization over y E RD is replaced by the equivalent optimization over functions RR". This is done by the interchanging of integral and minimum, see [5] Theorem 3A].
T E 1D-R"
The key point of this reformulation that the optimal solution of this maximin problem given by (v",77 )"), where T is the OT map from IP to Q. discussion in 29 [2] [53] 94.1]. I]. In practice, the potential v I + IR and the map I' R + IR are parametrized by neural networks tw T.. To train 0 and w, stochastic gradient ascent/descent (SGAD) over mini-batches from P, Q is used.
3.2 Algorithms for Continuous W Barycenters
Variational optimization. Problem may employ a generator Gg
5 optimization over probability measures. To estimate P. one with a latent measure $ on IR and train & by minimizing
R +
N Σ
an 2 (
min. {
(8)
Optimization (8) can be performed by using SGD on random mini-batches from measures P,, and The difference between possible variational algorithms lies in the particular estimation method for W3 terms. To our knowledge, only ICNN-based minimax solver [40] has been used to compute W3 in (8) yielding [SCW,B] algorithm
[9
Potential-bise -based optimization. 34
30 recover the optimal potentials { ot } for each pair (P, IP.)) via a non-mon mon-minimax regularized dual formulation. No generative model is needed: the barycenter is recovered by pushing forward measures using gradients of potentials or by barycentric projection. However, the non-trivial choice of the prior barycenter distribution required. Algorithm [CRWB] by 34 entropic quadratic regularization and [CW,B] algorithm by 33 uses ICNNs, congruence and consisenns (28) regularization.
Other methods. Recent work [4] combines the variational (8) and potential-based optimization via the e-cyclical monotonity regularization. In [6] an algorithm sample from the continuous Wasserstein barycenter via the gradient flows is proposed.
Iterative Wz-Barycenter Algorithm
Our proposed algorithm is based on the fixed point approach by 3 which we recall in $4.1 In $4.2 formulate our algorithm for computing Wasserstein-2 barycenters. In $4.3 show that algorithm generalizes the variational barycenter approach.
4.1 Theoretical Fixed Point Approach
Following 3 we define operator Ht P22.c (RD)) (TR (R. by H(PP) E O Tp HP, where Ti denotes the OT map from IP to In,,. The measure H(F) obtained by the operator is indeed
absolutely continuous, see 3 Theorem 3. According to (6). the barycenter P defined by (5) is
fixed point of H. i.c., HIP) = This suggests a way to compute F by picking some P P2.a [
and recursively applying H until convergence. However, there several challenges:
(a) A fixed point P E P2... (RD) satisfying H(P) = P may be not the barycenter 3 Example
3.1]. The situation analogous to that of the iterative k-means algorithm for a different problem clustering. There may be fixed points which are not globally optimal.
(b) The sequence (H*(P)) is tight 3 Theorem 3.6] so has a subsequence converging in a2 (R) but the entire sequence may not converge. Nevertheless, the value of the objective
5 decreases for (P)) as A + do 3 Prop. 3.3].
Efficient parametrization of the evolving measure H²P(P) is required. Moreover, to get 7*++1 (TP)
from Hºt (IP), one needs to compute N optimal transport maps Tx (P)-PP, which can be costly.
In 13 and 2) the fixed point approach is considered in the Gaussian H (P)) is guaranteed to converge to the unique fixed point - the barycenter. The Gaussian case also makes parameterization (c) simple since both measures F, and H² (TP) can be parametrized by means and covariance matrices, and the maps (T),,))-,,, are linear with closed form.
where the sequence

---

1
P
Te
Te
PP
Σa,
Σa, B,
m N
Gg #
# # S
N
Te.
GE
G
GE
Amiient space Latent space
Ambient space
Latent space
(a) Step 1. compute N approximate OT maps Ton. from generated measure P3 to input measures Pn.
(b) Step 2. We regress Gg(-) onto
G
En On Ton (Geo (-), where Geo is fixed copy of Gę before the update.
Figure Our proposed two-step implementation of the fixed-point operator H(-)..
For general continuous measures P.,. it remains an open problem to find sharp conditions on inputs P. and the initial measure of the fixed-point iteration for the sequence (P)) & to converge to the barycenter. In this work, we empirically verify that the fixed point approach works well for the input measures that we consider and for a randomly initialized generative model representing the evolving barycenter (84.2). We tackle challenge (c) and develop scalable optimization procedure that requires only sample access to P, E P2.00 D).
4.2 Practical Iterative Optimization Procedure
We employ generative model to parametrize the evolving measure, i.c., put PP : GGdS, where is a latent measure, e.g., S _O H). and Ge is neural network R - R2D with parameters &. Our approach compute the operator ((PE)) and update Ge consists of steps.
First, we approximately recover maps Tr.-s., via [MM:R] solver, i.e., we use pairs or networks (To.. ( }r and train them by optimizing (7) with <- Pg and Q + P,, For each n = 1,2, N, we perform SGAD by using batches from GedS and P,, and get To., a TPp-pPa.
Second, we update Gg to represent M((Pg) instead of IPg. Inspired by 12], we do this via regression. We introduce Gg. fixed copy of Gg. Next, we regress Gg(-) onto E On To., (G (-)
/ fe [e(Ge()) Can (G dS(2) + min
n=1
by performing SGD on random batches from S, e.g., by using squared error L(x, x') del 2/2 2/10 11 Thus, generator Gę(-) becomes close to En= To, (G5o as function RH + R.D. We get
N
Pg = GetS R Σ an ToToo.J = Σ To. HTPSo
12 n=1
on T
pr iP& H(TP..).
==1
i.e., the generated GetS measure approximates H((P.s))
Summary. Our two-step approach iteratively recomputes OT maps Ti.-P, (Figure 2a and then uses regression to update the generator (Figure 2b). The optimization procedure is detailed in Algorithm I Note that when fitting OT maps Te.-pe we start from previously used {To., V.,) rather than re-initialize them. Empirically, this works better.
4.3 Relation to Variational Barycenter Algorithms
We show that our Algorithm ( to variational approach (33.2) when the number of generator updates, KG, is equal to 1. More specifically, we show the equivalence of the gradient update w.r.t.
5

---

Algorithm 1:
Wasserstein Iterative Networks (WIN) for Barycenter Estimation
Input
lattent S and input ........ PN measures, weights oi,,...,
.......... .N 0
-n=1 On 1);
number of iters per network: Ka, Kr, Ki generator Gg R + R";
mapping networks To ..... .Txx RD + RP; potentials U Vay :RD - IR;
regression loss e IR" x R + R4;
satisfying GetS P: OT maps satisfying To,, To, Rs Pn
repeat
# OT solvers update
for n = 1, ........N............... .N do
for k, = 1,2, ... Kr do
Sample batches 7, 7, Y~ X Gg((Z));
L. [XT
Σ Uwn (To, (I)) TEX
Σ Uwl (v);
yEY
Update wn by using acc.
Ow.,
for kT .,., ...,, KT do
Sample batch Z ~ S; X + Gz(Z);
LT
Σ 1 To, (x)ll -ww,
Vws (To, (x))]:
X
ZEX
Update 6n by using UCT; a0,."
# Generator update (regression) G& copy (opy(Gg): for kg = 1..,.... Ke do Sample batch S;
LG 121 ZEZ
Σ l GG(z), En On Ton (Gea (G)(
(=))
Update by using
&ca. of
untll not converged;
parameters & of the generator in our iterative Algorithm 1 and that of (8) We assume that W3 terms are computed exactly in (8) regardless of the particular OT solver. Similarly, in Algorithm I we assume that maps Gg. S + P,, before the generator update are always exact, i.e., Te., = TP.o
Lemma 1. Assume that I. # P2,a Consider KG for the iterative Algorithm ie., we do a single gradient step regression update per OT solvers' update. Assume that E(x, x') 1|0 x' 2 i.e., the squared loss is used for regression. Then the generator's gradient update in Algorithm 7 is the same as in the variational algorithm:
N
ag [[ -
CanTres -F. (G5o (2)))
Σ an W.
(9)
de
n=1
where the derivatives are evaluated at ç = Eo.
We prove the lemma in Appendix A In practice, we choose K = 50 as it empirically works better.
5
Ave, celeba! Images Dataset
In this section, we develop a generic methodology for building measures with known W2 barycenter. We then use it construct Ave, celeba! dataset for quantitative evaluation of barycenter algorithms.
Key idea. Consider oi,,..., ........,a AN with Σ On 1, congruent convex functions 11,, d/ R + R. and measure PP E P DI with positive density. Define P , HP. Thanks to Brenier's theorem [1]. V is the unique OT map from IP to P,, Since the support of is D Vh. the unique (up to constant) dual potential for P.) [61] Since potentials th, are congruent, the barycenter P of P,, w.r.t. weights a1,... aN is IP itself [ C.2]. If v/,," 's are such that all P,,
absolutely continuous, then IP IP is the unique barycenter 2

---

y
V
V
P
Vul
P
VPP
Vey
Vey+P
P
PR
Vu
(a) Construction of pair of congruent functions v 2t" from convex vb, see Lemma
(b) Construction congruent the convex combinations of M congruent pairs (N% v see Lemma
Figure 3: Construction of tuples of congruent functions and production of measures with known W barycenter (5))
one obtains congruent 4h,, then for any E P2 PR pushforward measures P,, VI can be used as the input measures for the barycenter task. For I accessible by samples, measures P,, are also accessible by samples: one may sample x ~ P and push samples forward by Vii..
The challenging part is to construct non-trivial congruent convex functions thn- First, we provide novel method to transform single convex function into pair (1² v²") of convex functions satisfying Vild (x) (1 V (a: for all IR (Lemma 2). Next, extend the method
to generate congruent N-tuples (Lemma 3)
Lemma (Constructing congruent pairs), Let s be a strongly convex and L-smooth (for some > 0) function. Let B (0, 1). Define B-left and B-right functions by
1 def
12 + + -1
-1-
and
1p (1-
2 BBb.
(10)
2
Then (x) + (1 B) (x) 1=11 for x E RR,, i.e. convex functions v vh" are congruent w.r.t.
weights (3, B). Besides, for all x E DR. the gradient of def Voll (x)) can be computed via solving
3-strongly concave optimization:
|1811 (1- (1--8)8 (-B))(())
y = arg max (z. v)
(1I)
PER
2
In turn, the value yf def S (x) is given by -V((').
The proof is given Appendix A We visualize the idea of our Lemma in Figure 3a Thanks to
Lemma 2 any analytically known convex vi, e.g. an ICNN, can be used produce congruent pair dd,, 2b" To compute the gradient maps, optimization II] can be solved by convex optimization tools with V computed by automatic differentiation.
Lemma 3 (Constructing N congruent functions.). Let Vy...., V be convex functions, B1,...,Bku BM E 1) and wio vin be $m-left, Bm-rrght functions for wh respectively. Let T' E N be two rectangular matrices with non-negative elements and the sum of elements in
each column equals to 1. Let Wi...., WM satisfy Em=1 Wm 1. For 1...., N define
M Vn (x) del 2 Win 1mmT wit (2)) + (1 3mm Ban))a ) wor (x:)
(12)
Em Wm [mm 2 tha + (1 3m TTrma
Then .,,,..,, N are congruent w.r.. weights On EM wm [mm hom + Bmmm
def
We prove Lemma B Appendix A We visualize the idea of Lemmal in Figure 3b The lemma provides an elegant way to create functions in given congruent pairs (1/ e/y,). Gradients Vúù, of these functions are respective linear combinations of gradients Viid and V or
2 congruent functions from convex linear combinations of
7

---

Barycenter p= PCeeba
Viy AP = P1
Vy2AP= P2
Vyy
an HD
Figure 4: The production of Ave, celeba! dataset. The 1st line shows images Peeleta- Each of 3 next lines shows OT maps V (x) V = Pn to constructed measures Pn. Their barycenter w.r.t. (a1, 02, 03) = G 2 ) is Peekba. The last line shows congruence of thas i...
Σ On (1)
Samples in green boxes are included to dataset.
Dataset creation. We use CelebA 64 64 faces dataset [36 as the basis for our Ave. celeba! dataset. We assume that CelebA dataset is an empirical sample from the continu- ous measure I'Celeha E -2 (13366663)) which we put be the barycenter in our design, i.e.,
F- = Colcha. We construct diffirentiable congruent Uhn with bijective gradients that produce P. = Ven TT E (TR (R3x61x61) whose unique barycenter is Peteebu- In Lemma we set = 3, 2. W1
M 2, 81 B2
W2 and
(T)
1 0
T
-0:0 1010
which yields weights (a1, a2, az) ( (. We choose the constants above manually to make sure the final produced measures P, are visually distinguishable. We use (%
|=| I as convex functions, where ICNNs have ConvICNN64 architec- ture 29 Appendix B.1], 01, 02 random permutations of pixels and channels, 81, 82 axis-wise both functions, dm a de-colorization transform which sets R, G, B B) for vi and (R+G+B) for v.g. The weights of ICNNs are initialized by the pre- praitee potentials of W3 "Early" transport benchmark which map blurry faces to the clean ones 229 54.1]. All the implementation details are given in Appendix B.1
ICNN (sm (z)))) +
random reflections, À == 100
channels of each pixel to 1 R+ 25 G+ 13
50
Finally, to create Ave, celeba! dataset, we randomly split the images dataset into 3 equal parts containing 67K samples, and map each part respective measure P., Cel,t/Pcellha by V.l, Resulting 3 x 67K samples form the dataset consisting of 3 parts. We show the samples in Figure The samples from the respective parts are in green boxes.
Evaluation
The code_l is written the PyTorch and includes the script for producing Ave, celeba! dataset. The experiments are conducted on 4 4XGPU GTX 1080ti. The details are given in Appendix B
6.1 Evaluation on Ave, celeba! Dataset
We evaluate our iterative algorithm and a recent state-of-the-arr variational [SCW,B] by 19 on Ave, celeba! dataset. Both algorithms use generative model Pa = G for the barycenter and yield approximate maps TE.-RR to input measures. In our case, the maps are neural networks Te.. while in [SCW,B] they are gradients of ICNNs. The barycenters of Ave, celeba! fitted by our algorithm and [SCWBB] are shown in Figures 5a and 5b respectively Recall the ground truth barycenter is IPCeteaaa Thus, for quantitative evaluation we use FID score 22 computed on 200K generated samples w.r.t. the original CelebA dataset, Table I Our method drastically outperforms [SCW)B), Presumably, this is due to the latter using ICNNs which do not provide sufficient performance.

---

Pg=GGf5 [AID-758]
G:= [PI0=100]]
P P
Fg - P2
Pp- Pg
H(P))
H(Pp) [-0)=1000]
[AD=52285]
(a) Our Algorithm D
(b) Competitive [SCW,B] algorithm.
Figure 5: The barycenter and maps to input measures estimated by barycenter algorithms. The 1st line shows generated samples Pg Ge1S R Celebaa Lines 2-4 show maps TF.-F.. The last line shows the average map Σ On TP.
1
P2
P3
[CS]
[CS] fP FE IRID-55.5
|CS] 5 (PID 88.8)
Pe
(FID - 90.5)
[SCW/B]
[SCWBB] B] P2 62.4) 62.4]
[SCW)B] Pg - f [HDx=100)
+ Pt [HD - 67.4)
Our
Our Pg Pe
Our Pg Pe
Pj Pe
[RD = 49.3]
(FID = 46.9]
[FID= 61.5]
(a) Maps from P1 to the barycenter. (b) Maps from P2 to the barycenter.
(c) Maps from P3 to the barycenter.
Figure 6: Maps from inputs I',, to the barycenter IP estimated by the barycenter algorithms in view. For comparison with the original barycenter images, the faces are the same as in Figure
Method
F7D4
Meshod
FIDA
[SCW F3]
Gelz)
156.3
mm O. The P. (Ge(z)
152.1
CCS
(n))
90.5
75.5
88.8
Ours
Gelz)
75.8
The -p, (M)
67.4
62.4
319.62
Cn Th P. (Ge(z)
52.85
Ours
Tha - (v)
49.3
46.9
61.5
Table 1: FID scores of images from the learned barycenter.
Table 2: FID scores of images
mapped from inputs P,,.
Additionally, we evaluate to which extent the algorithms allow to recover the inverse OT maps Tr. from inputs P,, to the barycenter F.g = P. In [SCW,B], these maps are computed during training. Our algorithm does not compute them. Thus, we separately fit the inverse maps after main training by using [MM:R] solver between each input P. and learned IP& (Algorithm 2f Appendix 2. The inverse maps are given in F their FID scores in Table Here we add an additional constant shift [CS] baseline which simply shifts the mean of input P, to the mean ア7 of P. The vector is given by Σ Ont/n, where H. is the mean of P,, B] We estimate Іr from samples y P
6.2 Additional Experimental Results
Different domains. To stress-test our algorithm we compute the barycenters w.r.t. (01. 022, as))
(1 (1:3:3 3) of notably different datasets: 50K Shoes 68). 138K Amazon Handbags and 90K Fruits [5].. All the images are rescaled to 64 x 64. The ground gruth barycenter is unknown, but one may imagine what it looks like. Due to (). each barycenter image a pixel-wise average of shoe,
handbag and fruit, which is supported by our result shown in Figure 7
In the same figure, we also 7c d this allows generation
show the maps berween datasets through the barycenter (Figures 7b
of images in other categories with styles similar to given image. For instance, in Figure 7

---

Pp=GAA
P2
Pg
PI
R3 PE
P P1
P2 2t
93 1 Pt
Pt-+
PE P1
S P:
FI Pg
Py P1 P
93 PE F/ -
Pp-P3
P Sp Ps
P2 % P:
Pg + Pt Pe Ps
= ~7(P))
(a) Generated samples P3 R P., fitted maps to each P., and their average.
(b) Samples y ~ IP. mapped through Fg to each P...
(c) Samples y ~ P2 mapped through Ig to each Fn.
(d) Samples y ~ Ps mapped through Pg to each Fn.
Figure 7: The barycenter of Handbags, Shoes, Fruit (64 x 64) datasets fitted by our Algorithm We give the results of competitive [SCW2B] in Appendix B.4
generate an orange bag and an orange shoe by pushing the image of an orange through the barycenter. We provide more examples in Figure o3 Appendix C.5
Extra results. In Appendix C.I. we compute the barycenters of toy 2D dsitributions. In Appendix we provide quantitative results for computing barycenters in the Gaussian case. In Appendix we test how our algorithm works as a generative model on the original CelebA dataset, i.e., when =1 and IP PCeleha- We show that in this case it achieves FID scores comparable recent WGAN models. In Appendix C.4 similar to [1)] we compute barycenters of digit classes 0/1 of 32 32 grayscale MNIST 33
We also test our algorithm on FashionMNIST (6) 10 classes dataset.
7 Discussion
Potential impact (algorithm). We present a scalable barycenter algorithm based on fixed- fixed-point iterations with many application prospects, For instance, in medical imaging, MRI is often acquired at multiple sites where the overlap of information (imaging, genetic, diagnosis) between any two sites is limited. Consequently, the data on each site may be biased and can cause generalizability and robustness issues when training models. The developed algorithm could help aggregate data from multiple sites and overcome the distributional shift issue across sites.
Potential impact (dataset). There is no high high-dimensional dataset for the barycenter problem except for
cases (e.g. Gaussians), where the transport maps are always linear. Hence our proposed dataset fills an important gap, thereby allowing quantitative evaluation of future related methods. We expect our Ave, celeba! to become a standard dataset for evaluating continuous barycenter algorithms. In addition, we describe a generic recipe ($55 to produce new datasets.
Limitations (algorithm). In our algorithm, the evolving measure I' is not guaranteed be continu- ous, while it is continuous in the underlying fixed point approach. To enforce the absolute continuity of Ip'g = GG one may use an invertible network 18] for Gg and an absolutely continuous latent measure S in latent space of dimension H D. However, our results suggest this is unnecessary in practice - common GANs approaches also assume H & D. During the fixed-point iterations, the barycenter objective 5 decreases. However, there is no guarantee that the sequence of measures converges to fixed point which is the barycenter. Identifying the precise conditions on the input measures and the initial point is an important future direction. Besides, our algorithm does not recover inverse OT maps Tp -5 we compute them with an OT solver as a follow-up. To avoid this second step, may consider using invertible neural nets [8 to parametrize maps Te, in Algorithm [ Limitations (dataset). To create Ave, celeba! dataset (:5), we compose ICNNs with decolorization, random reflections and permutations to simulate degraded images. It is unclear how to produce other practically interesting effects via ICNNs. It also remains an open question on how to better hide the information of the barycenter image in the constructed marginal measures. Studying these questions is an interesting future direction that can inspire benchmarking other OT problems.
E. Burnaev was supported by the Russian Foundation for Basic Research grant 21-51-12005 NNIO a. A portion of this project was funded by the Skolkovo Institute of Science and Technology as part of the Skoltech NGP Program and funds were received by the Massachusetts Institute of Technology prior to Septemher 1, 2022. Neither Mr. Li, nor any other MIT personnel, contributed to any substantive or artistic alteration or enhancement of this publication after August 31, 2022.
OCKNOWLEDGEMER
10

---

References
[1] Martial Agueh and Guillaume Carlier. Barycenters the Wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-924, 2011.
[2] Jason M Altschuler, Sinho Chewi, Patrik Gerber, and Austin J Stromme. Averaging on the bures-wasserstein manifold: dimension-free convergence of gradient descent, arXiv preprint arXiv:2/06.08852,, 2021.
[3] Pedro C Alvarez-Esteban E Del Barrio, JA Cuesta-Albertos, and C Matrán. A fixed-point ap- proach to barycenters Wasserstein space. Journal of Mathematical Analysis and Applications, 441(2):744-762, 2016.
[4] Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 146-155. JMLR. org. 2017.
[5] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. arXiv preprint arXiv: 1701/008777 2017.
[6] Serguei Barannikov, Ilya Trofimov, Nikita Balabin, and Evgeny Burnaev. Representation topology divergence: A method for comparing neural network representations. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, edi- tors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 1607-1626. PMLR, 17-23 Jul 2022.
[7] Serguei Barannikov, Ilya Trofimov, Grigorii Sotníkov, Ekaterina Trimbach. Alexander Korotin, Alexander Filippov, and Evgeny Burnaev. Manifold topology divergence: framework for comparing data manifolds. Advances in Neural Information Processing Systems, 34:7294-7305. 2021.
[8] laroslav Bespalov, Nazar Buzun, Oleg Kachan, and Dmitry V Dylov. Data augmentation with manifold barycenters. arXiv preprint arXiv:2104.00925, 2021.
[9] Jérémic Bigot, Elsa Cazelles, and Nicolas Papadakis. Data-driven regularization of wasserstein barycenters with an application to multivariate density registration. Information and Inference: A Journal of the IMA, 8(4):719-755, 2019.
[10] Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasser- stein barycenters of measures. Journal of Mathematical Imaging and Vision, 51(1):22-45, 2015.
[11] Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375-417, 1991.
[12] Yucheng Chen, Matus Telgarsky, Chao Zhang, Bolton Bailey, Daniel Hsu, and Jian Peng. A gradual, semi-discrete approach to generative network training via explicit Wasserstein minimization. In International Conference on Machine Learning, pages 1071-1080. PMLR, 2019.
[13] Sinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin J Stromme. Gradient descent algorithms for sares-wasserstein barycenters. In Conference on Learning Theory, pages 1276- 1304. PMLR, 2020.
[14] Jinjin Chi, Zhiyao Yang. Jihong Ouyang, and Ximing Li. Variational wasserstein barycenters with c-cyclical monotonicity. arXiv preprint arKiv:2/ 10. /707, 2021.
[15] Pierre Colombo, Guillaume Staerman, Chloe Clavel, and Pablo Piantanida. Automatic text evaluation through the lens of wasserstein barycenters, 2021.
[16] Chiheb Daaloul, Thibaut Le Gouic, Jacques Liandrat, and Magali Tourus.. Sampling from the wasserstein barycenter. arXiv preprint arXiv:2105.01700 2021.
[17] Pierre Dognin, Igor Melnyk, Youssef Mroueh, Jerret Ross, Cicero Dos Santos, and Tom Sercu. Wasserstein barycenter model ensembling. arXiv preprint 1499999 2019.
11

---

[18] Christian Etmann, Rihuan Ke, and Carola-Bibiane Schönlieb. iunets: learnable invertible and downsampling for large-scale inverse problems. In 2020 IEEE 30th International Workshop on Machine Learning for Signal Processing ( (MISP), pages 1-6. IEEE, 2020.
[19] Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arX0::2007.04462, 2020.
[20] Werner Fenchel. On conjugate convex functions. Canadian Journal of Mathematics, 1(1):73-77, 1949.
[21] Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. In Advances in neural information processing systems, pages 1440-3448, 2016.
[22] Aude Genevay, Gabriel Peyré, and Marco Cuturi. Gan and vae from an optimal transport point of view. arXiv preprini urXiv:1706.00 1706.01807, 2017.
[23] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to local nash equilibrium. In Advances in neural information processing systems, pages 6626-6637) 2017.
[24] David I Inouye. Zeyu Zhou, Ziyu Gong, and Pradeep Ravikumar. Iterative barycenter flows. arXiv preprini arXiv:2104.022: 2021.
[25] Leonid Kantorovitch. On the translocation of masses. Management Science, 5(1):1-4, 1958.
[26] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[27] Diana Koldasbayeva, Polina Tregubova, Dmitrii Shadrin, Mikhail Gasanov, and Maria Pukalchik. Large-scale forecasting of heracleum sosnowskyi habitat suitability under the climate change on publicly available data. Scientific reports, 12(1):1-11,200 2022.
[28] Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations,
2021.
[29] Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34: 34:14593-14605, 2021.
[30] Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein- 2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021.
[31] Alexander Korotin, Vladimir V'yugin, and Evgeny Burnaev. Mixability integral losses: A key to efficient online aggregation of functional and probabilistic forecasts. Pattern Recognition, 120: 1081 17,, 2021.
[32] Julien Lacombe, Julie Digne, Nicolas Courty, and Nicolas Bonneel. Learning to generate wasserstein barycenters. arXiv preprint arXiv:2102.12272 2021.
[33] Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010.
[34] Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin Solomon. Continuous regularized Wasserstein barycenters. arXiv preprint arXiv:2008. /2534, 2020.
[35] Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein GAN with quadratic transport cost. In Proceedings of the IEEE International Conference on Computer Vision, pages 4832- 4841, 2019.
[36] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
12

---

[37] Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Large- scale optimal transport via adversarial training with cycle-consistency. arXiv preprint arXiv:2003.06633 2020.
[38] Mario Lucic, Karol Kurach, Marcin Michalski, Sylvain Gelly, and Olivier Bousquet. Are GANs created equal? a large-scale study. In Advances in neural information processing systems, pages 700-709, 2018.
[39] Boyang Lyu, Thuan Nguyen, Prakash Ishwar, Matthias Scheutz, and Shuchin Aeron. Barycen- teric distribution alignment and manifold-restricee invertibility for domain generalization, 2021.
[40] Ashok Vardhan Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason D Lee. Optimal transport mapping via input convex neural networks. arXiv preprint 1908v: 10962, 2019.
[41] Alberto Maria Metelli, Amarildo Likmeta, and Marcello Restelli. Propagating uncertainty in reinforcement learning via wasserstein barycenters. In 33rd Conference on Neural Information Processing Systems, NeurlPS 2019, pages 4335-43477 Curran Associates, Inc., 2019.
[42] Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34:15243-15256, 2021.
[43] Eduardo Fernandes Montesuma and Fred Maurice Ngole Mboula. Wasserstein barycenter for multi-source domain adaptation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16785-16793, 2021.
[44] Youssef Mroueh. Wasserstein style transfer. arXiv preprint arXiv: 19t../995.2.... 2019.
[45] Horea Mureşan and Mihai Oltean. Fruit recognition from images using deep learning, arXiv preprint arXiv: 1712.00588 2017.
[46] Quan Hoang Nhan Dam, Trung Le, Tu Dinh Nguyen, Hung Bui, and Dinh Phung. Threeplayer Wasserstein GAN via amortised duality. Proc. of the 28th Int. Joint Conf. on Artificial Intelligence (IJCAI), 2019.
[47] Quentin Paris. Online learning with exponential weights in metric spaces. arXiv preprint arXiv22103.14389, 2021.
[48] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Foundations and Trends@ in Machine Learning, 11(5-6):355-09)) 2019.
[49] Julien Rabin, Sira Ferradans, and Nicolas Papadakis. Adaptive color transfer with relaxed optimal transport. In 2014 IEEE International Conference on Image Processing (ICIP), pages 4852-4856. IEEE, 2014.
[50] Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernt. Wasserstein barycenter and its application to texture mixing. In International Conference on Scale Space and Variational Methods Computer Vision, pages 435-446. Springer, 2011.
[51] R Tyrrell Rockafellar. Integral functionals, normal integrands and measurable selections. In Nonlinear operators and the calculus of variations, pages 157-207. Springer, 1976.
[52] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pages 234-241, Springer, 2015.
[53] Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2021.
[54] Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58- 63):94, 2015.
13

---

[55] Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv: 1111111112222 1.02883, 2017.
[56] Dror Simon and Aviad Aberdam. Barycenters of natural images constrained wasserstein barycenters for image morphing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7910-7919, 2020.
[57] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409. 1556, 2014.
[58] Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (TOG), 34(4):1-11, 2015.
[59] Sanvesh Srivastava, Volkan Cevher, Quoe Dinh, and David Dunson. Wasp: Scalable bayes via barycenters of subset posteriors. Artificial Intelligence and Statistics, pages 912-920, 2015.
[60] Sanvesh Srivastava, Cheng Li, and David Dunson. Scalable bayes via barycenter Wasser- stein space. The Journal of Machine Learning Research, 19(1)3112-346, 2018.
[61] Thomas Staudt, Shayan Hundrieser, and Axel Munk. On the uniqueness of kantorovich potentials. arXiv preprini arXiv:2201.08316 2022.
[62] Amirhossein Taghvaci and Amin Jalali. 2- 2-Wasserstein approximation via restricted convex potentials with application to improved training for GANs. arXiv preprint arXiv: 2019.
[63] Jules Vidal, Joseph Budin, and Julien Tiemy. Progressive wasserstein barycenters of persistence diagrams. IEEE transactions on visualization and computer graphics, 26(1):151-61,, 2019.
[64] Cédric Villani. Topics in optimal transportation.. Number 58. American Mathematical Soc., 2003.
[65] Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
[66] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 11t,.07707777 2017.
[67] Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. volume 97 of Proceedings of Machine Learning Research, pages 6892-689,, Long Beach, California, USA, 09-15 Jun 2019, PMLR.
[68] Aron Yu and Kristen Grauman. Fine-grained visual comparisons with local learning. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 192-199, 2014.
14

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
[Yes] See g
(b) Did you describe the limitations of your work?
[Yes] See $7
(c) Did you diseuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms them?
[Yes]
2. If you are including theoretical results...
(a) Did you state the full set assumptions of all theoretical results?
[Yes] All the assumptions stated in the main text.
(b) Did you include complete proofs of all theoretical results?
[Yes] All the proofs are given in the appendices.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either the supplemental material or as a URL)?
[Yes] The code and the instructions are included in the supplementary material. The datasets
that we use are publicly available.
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)?
[Ye]] See and the supplementary material (appendices code).
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] Due to the well-known high computational complexity of learning generative models, most experiments (both with our method and alternatives) were conducted only once.
(d) Did you include the total amount compute and the type of resources used (e.g.. type of GPUs, internal cluster, or cloud provider)?
[Yes] See the discussion in $6 and Appendices.
you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work existing assets, did you cite the creators?
[Yes] See the discussion in section 5
(b) Did you mention the license of the assets?
[No] We refer to the datasets' public pages.
(c) Did you include any new assets either in the supplemental material or URL?
[Yes] See the supplementary material
(d) Did you discuss whether and how consent was obtained from people whose data you're yru're using/curating?
[N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[No] For CelebA faces dataset, refer to the original authors publication.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applica- ble? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
(IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation?
[N/A]
15


Reviews:
[{"review_id": "q0W2A-7-7PS", "reviewer": "Reviewer_NEXn", "summary": "Computing the Wasserstein barycenter between continuous measures is an important problem. In this paper, authors propose a method based on a fixed-point approach to compute the Wasserstein barycenter with arbitrary neural networks. Moreover, they introduce a way to create datasets to compare to quality of the obtained barycenter and compare their method to a ICNN based method on the \"ave, celeba!\" dataset, which they generated using their method.", "questions": "Some related works are not cited. In [1], a method is proposed to compute continuous Wasserstein barycenter for abitrary costs. I believe it would be worth also comparing the method against their method.\n\n[2,3] are also a missing reference on the methods used to approximate Wasserstein barycenters.\n\nThe method seems pretty heavily computationally (since it requires to compute N OT maps at each iteration and solve a regression problem), but I guess it is also the case of the other method which is compared with. I would be curious to see some comparison of the training time between the methods.\n\n\n[1] Chi, Jinjin, et al. \"Variational Wasserstein Barycenters with c-Cyclical Monotonicity.\" arXiv preprint arXiv:2110.11707 (2021).\n\n[2] Backhoff-Veraguas, Julio, et al. \"Stochastic gradient descent in Wasserstein space.\" arXiv preprint arXiv:2201.04232 (2022).\n\n[3] Daaloul, Chiheb, et al. \"Sampling from the wasserstein barycenter.\" arXiv preprint arXiv:2105.01706 (2021).", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nThis paper propose to use generative model into the fixed-point approach previously introduced. This allows to use NNs to parametrize the barycenter and to obtain better results than other papers. Moreover, they introduce a way to construct benchmark datasets to compare the efficiency of the method to compute Wasserstein barycenters. This is a very nice contribution which can be very helpful to compare methods.\n\nStrengths:\n- Fixed-Point approach with neural networks and the regression approach.\n- A new way to construct a benchmark dataset to compare methods on Wasserstein barycenters\n\n\nWeaknesses:\n- Comparison with only one method. A more complete comparison with the different methods on the newly introduced dataset could be nice.\n- Still some open questions on whether the fixed-point approach approximates the barycenter.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "bWmEjz3lvnv", "reviewer": "Reviewer_xsZZ", "summary": "Authors propose a methodology for estimation of wasserstein barycenter in continuous setup. Authors argue that it beats the state of the art. In the passing authors conceive a new way of transforming/splitting datasets leading to new 'ave, celeba' dataset which might be useful in the future for benchmarking in OT tasks.", "questions": "Why do you only compare against SCW2B? I understand that this is the state of the art but by only comparing against it, and given that there are multiple components to your methodology, it is hard to identify where the gains come from. I suggest to whenever further experiment that allow us to debug where the gains come from. What are the advantages to each of the components in section 4?\n\nI strongly suggests to add a toy example. Is there any reason for why this should be not necessary? For example, one could try to find the barycenter of some simple continuous gaussian data. Adding such toy example (which suitable baselines, even if they don't correspond to the SOTA SCW2B) would help 'debug' and interpret the computational gains.\n\nMethods related: on many parts it is mentioned that ICNN don't work so well but at the same time the congruent functions are computed with ICNN. Could you make this distinction more explicit in the text (why for dataset creation it makes sense to use ICNN while they don't work so well for barycenter estimation?)", "limitations": "To me understanding authors correctly discuss limitations. But lack of sufficient interpretability is a substantial limitation of this work, as it is submitted at this point.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The main strengths are that it is a thoughtful paper on a relevant topic that keeps demanding new methods, and where the proposed methodology is valuable because the 'tricks' are novel, and because they lead to better results. It is clear to me that there is a large amount of work involved on this project.\n\nThe new dataset is also a major strength of this paper, and I hope authors will work in order to make it easily publicly available for future work. Although the math is rather simple, it is a strength of the paper that the use of math leads to quite valuable methodological tricks.\n\nThe main weakness is that although there are certainly quite a few clever methodological nuggets, altogether it is very hard to identify how everything is put together to lead to a better algorithm (see for example questions below). If that is not clear enough then the community has reasons to be skeptical about the evaluation. This is mostly an interpretability issue and I hope it can be addressed.\n\nThis requires some work in two related aspects: 1) although writing is clear, what is less clear to me are the different components of the proposed algorithm.  Figure 2 is beautiful but it is not clear enough. Here I missed the usual Algorithm box. I believe authors should try to polish the methods section (4) so it is crystal clear what the contributions are. I would have love to have a more thorough discussion than 'it works better empirically'. Hopefully, we would see a clear justification to every design choice made in section 4.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "VFLqrLmqDg-", "reviewer": "Reviewer_idDE", "summary": "The article proposes an iterative procedure to compute Wasserstein barycenters, the parametrised barycenter is modelled as a latent distribution pushed forward by a neural network. \n\nThe authors study the relationship with existing methods and use their proposal to build a benchmark dataset for barycentric computation", "questions": "How can the author justify the general applicability of their method given that the experiments only consider purposed-built datasets? ", "limitations": "the limitations of this work are identified in the manuscript ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper does a great job in the literature review, where it puts its method in context and compares it with other techniques in terms of their advantages and disadvantages.\n\nThe theoretical treatment of the contribution is well explained, the algorithm consists in iteratively applying a neural operator to a (latent) density according the known fixed point barycentric method by [3]. The proposal is validated in the celebrity dataset. \n\nIn my opinion, the weakest point of the article is the experimental validation. Though the proposal is sound, only two sets of experiments are shown: i) evaluation of a synthetic dataset built upon the proposed method, and ii) other synthetically-generated dataset. These experiments hardly validate the idea of a general method for barycenter computation, since there is no evidence that the proposed method operates beyond datasets built in the specific manner mentioned by the authors. This issue is identified in the manuscript, in Sec 7 (lines >295) where the authors acknowledge that their dataset only cater for specific (dispersion) properties - given this observation, it was clear that the experimental validation should've included richer datasets. \n\nThis weakness in the contribution is further clarified when the authors identify a possible application in MRI (Sec 7), where datasets have limited overlap. Perhaps a stronger experimental validation could include real world datasets of that kind.\n\nThough well written in general, the article could benefit from a proofreader ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "N9gqJsk4z7", "reviewer": "Reviewer_3V2X", "summary": "This paper proposes to learn the Wasserstein barycenter through the fixed point iteration method. They interchangeably 1) solve the OT map between the barycenter candidate and each marginal distribution, 2) update the barycenter candidate by fixed-point iteration, then regress a generator to fit the barycenter candidate. It also develops a CelebA barycenter dataset that has a known barycenter solution.", "questions": "The paper claims they cannot obtain the inverse map: marginal distribution -> barycenter during training, but the inverse map can be recovered by $x- \\nabla v(x)$ if using eq 7 to solve the optimal map. This is discussed at the end of section 2 in [W2 solver benchmark paper](https://openreview.net/pdf?id=CI0T_3l-n1). Why not use this inverse map as a baseline?\n\nIn Figure 5, SCWB is particularly bad for the map corresponding to $P_3$. This is somehow weird. Do authors have clue why this happens?\n\nFor the ablation study, I think it would be helpful if the authors can add some experiments with $k_G=1$. Thus we can be more clear about whether SCWB has worse performance because ICNN is not expressive enough.", "limitations": "see weakness", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nIt provides a nice interpolation between the gradient updates of the least square regression to fixed point barycenter update and the variational barycenter algorithms.\n\nIt gives theoretical guidance to generating probability measures with known Wasserstein barycenter.\n\nIt obtains advantageous performance compared to the competitive SCWB algorithm.\n\nWeaknesses:\n\nAs at the beginning of section 4.1, the fixed point doesn't promise to be the barycenter, and the entire sequence may not converge, which may cause numerical issues. These two challenges would question the foundation of the proposed fixed-point iteration solver.\n\nThe paper claims to tackle challenge c, but the algorithm still needs to compute N maps. So it doesn't decrease the calculation burden, at least when compared with SCWB.\n\nIt seems to me that computational time is rarely mentioned in the paper. For example, the paper uses $k_G=50$, and also a large number of iterations $k_v=50$, while the corresponding parameter is only 1 and 10 for SCWB. I conjecture the computational time of this method is several times larger than SCWB.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept