Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization
Authors: Alexander Korotin
Lingxiao Li
Justin Solomon
Evgeny Burnaev
Venue: ICLR 2021 Poster (2021)
Keywords: wasserstein-2 barycenters
non-minimax optimization
cycle-consistency regularizer
input convex neural networks
continuous case

Abstract: Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport. In this paper, we present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures, which are not restricted to being discrete. While past approaches rely on entropic or quadratic regularization, we employ input convex neural networks and cycle-consistency regularization to avoid introducing bias. As a result, our approach does not resort to minimax optimization. We provide theoretical analysis on error bounds as well as empirical evidence of the effectiveness of the proposed approach in low-dimensional qualitative scenarios and high-dimensional quantitative experiments.

Published as a conference paper at ICLR 2021
CONTINUOUS WASSERSTEIN-2 BARYCENTER ESTIMATION WITHOUT MINIMAX OPTIMIZATION
Alexander Korotin
Lingxiao Li
Skolkovo Institute of Science and Technology Advanced Data Analytics in Science and Engineering Group
Massachusetts Institute of Technology Geometric Data Processing Group Cambridge, Massachusetts, USA lingxi wwwpxiao@mit...m
Moscow, Russia
a. korootn
Justin Solomon
Evgeny Burnaey
Massachusetts Institute of Technology Geometric Data Processing Group Cambridge, Massachusetts, USA jsolomonémit.ed
Skolkovo Institute of Science and Technology Advanced Data Analytics in Science and Engineering Group Moscow, Russia e wwwnaee@@kklltehh.cr ru
ABSTRACT
Wasserstein barycenters provide geometric notion of the weighted average of probability measures based on optimal transport. In this paper, we present a scal- able algorithm to compute Wasserstein-2 barycenters given sample access the input measures, which are not restricted to being discrete. While past approaches rely on entropic or quadratic regularization, we employ input convex neural net- works and cycle-consistency regularization to avoid introducing bias. As result, our approach does not resort minimax optimization. We provide theoretical analysis on error bounds as well as empirical evidence of the effectiveness of the proposed approach in oiwdiiensional qualitative scenarios and high-dimensional quantitative experiments.
INTRODUCTION
Wasserstein barycenters have become popular due to their ability to represent the average of probabil- ity measures a geometrically meaningful way. Techniques for computing Wasserstein barycenters have been successfully applied to many computational problems. In image processing, Wasserstein barycenters are used for color and style transfer (Rabin et al.] 2014 Mroueh] 2019). and texture synthesis (Rabin al.] 2011). In geometry processing, shape interpolation can be done by computing barycenters (Solomon et al. 2015). In online machine learning. barycenters are used for aggregating probabilistic predictions of experts (Korotin al.) 2019b). Within the context Bayesian infer- ence, the barycenter of subset posteriors converges to the full data posterior, thus enabling efficient computational methods based on finding the barycenters (Srivastava et al.) 2015 2018).
Fast and accurate barycenter algorithms exist for discrete distributions (see Peyré et al. (2019) for survey), while for continuous distributions the situation is more difficult and remains unexplored until recently (L.i et al. (2200; Fan et al., 2020 Cohen et al. 2020). The discrete methods scale poorly with the number of support points of the barycenter and thus cannot approximate continuous barycenters well, especially in high dimensions.
In this paper, we present a method to compute Wasserstein-2 barycenters continuous distributions based on novel regularized dual formulation where the convex potentials are parameterized by input convex neural networks (Amos et al. 2017). Our algorithm is straightforward without introducing bias (e.g. Li et al. (2020)) requiring minimax optimization (e.g. Fan et al. (2020)). This is made possible by combining a new congruence regularizing term combined with cycle-consistent regularization (Korotin et al. 2019a). As will show in the analysis, thanks to the properties of

---

Published as a conference paper at ICLR 2021
Wasserstein-2 distances, the gradients of the resulting convex potentials "push" the input distributions close to the true barycenter, allowing good approximation of the barycenter.
PRELIMINARIES
We denote the of all Borel probability measures R- with finite second moment by P2(RD). We use (R) C (P((IP) to denote the subset of all absolutely continuous measures (w.r.t. the Lebesgue measure).
Wasserstein-2 distance.
For P. F,Q P2(P), the Wasserstein-2 distance is defined by
| 11yl dr(z.
def min TELL
(1)
RD XBD
2
where [I, Q) is the of probability measures on RD D whose marginals are IP, respectively, This definition is known as Kantorovich's 's primal form of transport distance (Kantorovitchi 1958).
The Wasserstein-2 distance W2 is well-studied in the theory of optimal transport (Brenier 1991. McCann et al.) 1995). In particular, it has a dual formulation (Villani 2003):
| d((x))
1 RDD 2 dQ(u)
V/ (x))dd())) + )())d(()
WG(P, Q)
+
min
(2)
Ecoon
RDD
where the minimum is taken over all the convex functions (potentials) &b R RURUU and (v) max ((x, y) - (xx)) RDD + R [xc) is the convex conjugate of de (Fenchel 1949
which is also a convex function. The optimal potential o is defined up to an additive constant.
Brenier 1991) shows that if TP does not give mass sets of dimensions at most D then the optimal plan is uniquely determined by π² [idgo, T where T RRD + RD is the unique solution to the Monge's 's problem
| +(=)l1 dF(x).
T arg min
=P= RD
The connection between T and the dual formulation (2 is that T = Ve" where 10* is the optimal solution of (2). Additionally, if Q does not give mass sets of dimensions at most D then T is invertible and
('(x) V (x),
(w) (V
In partieular, the above discussion applies the case where IP, Q E (RD).
Wasserstein-2 barycenter. Let PI,...... ...... Px E (R. Then, their barycenter w.r.t. weights 01 ONN (an, and Σ On = i)
N an (PP,,,
def arg min PEP.
(4)
Throughout this paper, assume that least of P1,..., PN P22 (Tk: (TR has bounded density. Under this assumption, is unique and absolutely continuous, i.e., P P2 (RR and it has bounded density (Agueh & Carlier 2011 Definition 3.6 & Theorem 5.1).
For n E (1,2, N). let (vht, vh*) be the optimal pair of (mutually) conjugate potentials that transport P, to P, i.e., VI = - and VisedP P,, Then (10) satisfy
N
|
(
and
Or (,))(())
2 +c.
(5)
m=1
for all x E R (Agueh & Carlier 2011 Alvarez-Estebann
al. 2016). Since optimal potentials are
defined up to constant, for convenience, we set e = O. The condition (5) serves as the basis for our
algorithm for computing Wasserstein-22 barycenters. say that potentials th,,... ..., t/N are congruent
w.r.t. weights al,,.r On if their conjugate potentials satisfy D
( i.e. Σm= an (n)
||>|| for all
2

---

Published as a conference paper at ICLR 2021
3 RELATED WORK
Most algorithms in the field of computational optimal transport are designed for the discrete setting where the input distributions have finite support; see the recent survey by Peyré et al. (2019) for discussion. A particular popular line of algorithms are based on entropic regularization that gives rise to the famous Sinkhorn iteration (Cutuni 2013 Cuturi & Doucet 2014). These methods are typically limited to support of 10 10 points before the problem becomes computationally infeasible. Similarly, discrete barycenter methods (Cuturi & Doucet 2014), particularly the ones that rely on fixed support for the barycenter (Dvurechenskii et al. 2018 Staib et al., 2017), cannot provide precise approximation of continuous barycenters in high dimensions, since a large number of samples needed; see experiments in Fan et al. (2020) $4.3) for example. Thus we focus on the existing literature in the continuous setting.
Computation of Wasserstein-2 distances and maps. Genevay et al. 2016) demonstrate the pos- sibility of computing Wasserstein distances given only sample access to the distributions by param- eterizing the dual potentials functions in the reproducing kernel Hilbert spaces. Based on this realization, Seguy et al. (017)) propose a similar method but use neural networks to parameterize the potentials, using entropic or L regularization w.r.t. IP Q to keep the potentials approximately conjugate. The transport map is recovered from optimized potentials via barycentric projection.
As we note in 2 W2 enjoys many useful thcoretical properties. For example, the optimal potential zi* is convex, and the corresponding optimal transport map is given by Vib". By exploiting these properties, Makkuva al. (2019) propose a minimax optimization algorithm for recovering transport maps, using input convex neural networks (ICNNs) (Amos et al. 2017) to approximate the potentials.
An alternative to entropic regularization is the cycle-consistency regularization proposed by Korotin et al. (2019a). It uses the property that the gradients of optimal dual potentials are inverses of each other. The imposed regularizer requires integration only over the marginal measures IP and Q, instead of over IP x Q as required by entropy-based alternatives. Their method converges faster than the minimax method since does not have an inner optimization cycle.
Xie et al. (2019) propose using two generative models with shared latent space to implicitly compute the optimal transport correspondence between F and Q. Based on the obtained correspondence, the authors are able to compute the optimal transport distance between the distributions.
Computation of Wasserstein-2 barycenters. few recent techniques tackle the barycenter prob- lem 4 using continuous rather than discrete approximations of the barycenter:
MEASURE-BASHE BASED (GENERATIVE) OPTIMIZATION: Problem (4) optimizes over probability mea- sures. This can be done using the generic algorithm by Cohen al. (2020) who employ generative networks to compute barycenters w.r.t. arbitrary discrepancies. They test their method with the maximum mean discrepancy (MMD) and Sinkhorn divergence. This approach suffers from the usual limitations of generative models such as mode collapse. Applying it to W2 barycenters requires estimation of W2 (IP n. P). Fan et al. (2020) test this approach using the minimax method by Makkuva et al. (2019), but they end up with challenging min-max-min problem.
POTENTIAL- BASED OPTIMIZATION: Li al. (2020) recover the optimal potentials {v:} } via non-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.
METHODS
Inspired by et al. (020 we use a potential- approach and recover the barycenter by using
gradients of the potentials as pushforward maps. The main differences are: (1) we restrict the potentials to be convex, (2) we enforce congruence via a regularizing term, and (3) our formulation
does not introduce bias, meaning the optimal solution of our formulation gives the true barycenter.
3

---

Published as a conference paper at ICLR 2021
4.1 DERIVING THE DUAL PROBLEM
Let P be the true barycenter. Our goal is to recover the optimal potentials {v v input measures F,, into
mapping the
To start, express the barycenter objective
4
after substituting the dual formulation 2
Σ
O
|
2
dIP
dIP ( +
J00
/
Mlll dl(())
N
N
min JCony
an
Un (r)
(2)
Σ
an /
in (n(u)d(()))
(6)
The minimum is attained not just among convex potentials (wn but among congruent potentials (see discussion under )); thus, we can add the constraint that {inn are congruent to (6). Hence,
| dP, nn (x)
Σ On / Un (n) JRD
on W2 ((,,,PP) P)
Σ
On
min (wh) congruent
(7)
MultiCon(HanBB)lC
To transition from 6
to
we used the fact that for congruent {un we have
Σ (n
Ex=1 SRo (nVa(y)dP(v) SEp wvll dP
We call the value inside the minimum in 7 the multiple correlation of [Fn} with weights {a, w.r.t potentials {1h.). Notice that the true barycenter P appears nowhere the right side of (7) Thus the optimal potentials { can be recovered by solving the following
N
min (Wo congruent
Pa =
min {va) congruent
Σ 1=
an
/ (n (n())dPp())) JRD
(8)
4.2 IMPOSING THE CONGRUENCE CONDITION
challenging to impose the congruence condition on convex potentials. What if we relax the congruence condition? The following theorem bounds how close set of convex potentials {th, is to { } in terms of the difference of multiple correlation.
Theorem 4.1. Let P E P2. (RD) be the barycenter P1.... ,NN E (TR w.r.t weights O1....,aN ON. Let {1b)) } be the optimal congruent potentials of the barycenter problem. Sup- pose have B-smoouh convex potentials {mn) } for some B (o,,, and denote A MultiCorr( (am, Pn, ( 1)- MultiCorr( (Cin Pn] (Vi)). Then,
((
N
an (1) (1) |
N
A +
dP
dP
2
Σ an (V)((=) (e)n())
(9)
JRD
28
R=1
1=1
Congruence mismontch
Here 11 denotes the norm induced by inner product in Hilbert space C2 + -R., ,n). We call the second term on the left of (9) the congruence mismatch.
We prove this in Appendix Note that if the congruence mismatch non- nos-ivsitive, then
1
1 ar (a)(() - Z Σ an W
A
(10)
2B
where the last inequality of (T)) follows from (Korotin et al. 2019a conclude that for all E [L....,N} NN. have ( P.,.. P))
Lemma A.2). From (10).
BA This shows that if
the congruence mismatch is non-positive, then A, the difference in multiple correlation, provides
We say that diffirentiable function f P R is B-smooth if its gradient V is B-Liipchiizz
4

---

Published as a conference paper at ICLR 2021
upper bound for the Wasserstein-22 - distance between the true barycenter and each pushforward Ve. This justifies the use V to recover the barycenter. Notice for optimal potentials, the congruence mismatch is zero.
Thus to penalize positive congruence mismatch, we introduce a regularizing term
Ri((an), (va))
def
Σ On (v)
2 df((v).
(11)
JDD
Because take the positive part of the integrand of to get non-negative, we have
11 and that the right side of ( is
MultiConn((a,, ] (tnn) + 1 (1ann (wn)))
MultiCor(((a,,))) Pn (())) 2
for all convex potentials J. On the other hand, for optimal potentials [tn} = {w J. the inequality turns into equality, implying that adding the regularizing term 1 (f ), (v,, }) to (8) will not introduce bias the optimal solution still yields {v/))-
However, evaluating [T exactly requires knowing the true barycenter P a priori. To remedy this issue, one may replace I with another absolutely continuous measure PP (T 2 and is probability measure) whose density bounds that of FP from above almost everywhere. In this case,
N
1ull2 2
T R-((((a)) (wn)) =
an (
1. dP
(. (wr)).
(12)
Hence we obtain the following regularized version of 8 where (w)) is the optimal solution:
min ECony
MMutiCorr((a,,, P.) (v) + R(((an) [vm))]]
(13)
Selecting measure т . is not obvious. Consider the case when {Pr} are supported on compact sets . Ny..... Xy C R and P, has density upper bounded by h 06. In this scenario, the barycenter density upper bounded by h a, (Alvarez-Estebar T P supported on ConvexHull(A;, ...., .,,,,,
al. 2016 Remark 3.2). Thus, the measure with this density is an upper bound for P. We will address the question of how to choose properly in practice in $4.4
4.3 ENFORCING CONJUGACY POTENTIALS PAIRS
Throughout this subsection, we assume the upper bound finite measure T I of the P is known. The optimization problem (13) involves only the potentials {va,,, but also their conjugates {vr, }.
This brings practical difficulty since evaluating conjugate potentials is hard (Korotin et al. 2019a).
Instead we parameterize potentials v and y separately using input convex neural networks (ICNN)
as vst and vit respectively. We add an additional cycle-consistency regularizer enfore the conjugacy of eift and wh as Korotin al. (019a). This regularizer defined as
R w def
Vith (e)( - x=12 dP, (z) V Vel idg i
Note that R,"- (1tt ( this condition is necessary for vbt and with be conjugate with each other. Also, it sufficient condition for convex functions to be conjugates up to an additive constant.
use one-sided regularization. In our case, computing the regularizer of the other direction || V idgo || is infeasible, since F unknown. If fact, Korotin al. (019a) demonstrates that such one-sided condition is sufficient.
5

---

Published as a conference paper at ICLR 2021
In this way we use 2N imput convex neural networks for { vn.. By adding the new cycle consistency regularizer into 13). obtain final objective:
Approximate multiple correlation
Σ an (a.
an R (ol w.). (14) (a ((
min
(v% v.
sorp (ar)
Congruence reg.
n=1
Cycle regularizer
Note that we express the aproximate multiple correlation by using both potentials {wt and {v This is done to eliminate the freedom of an additive constant { that is not addressed by cycle regularization. We denote the entire objective MultiCorr ( (IP,, [ 1 ); IP, A). Analogous
to Theorem 4.1 we have following result showing that this new objective enjoys the same properties the unregularized version from (8)-
Theorem 4.2. Let P E P2,ac (R.a be the barycenter of Fi...., PN a2 (TR.) w.r. weights
01,...,0NN. Let {177 be the optimal congruent potentials of the barycenter problem. Suppose we have T such that convex and B smooohh convex potentials {w5)] with 0 < 3t S B < 30 and X >> 2(3
and T P. Suppose we have convex potentials {vh1) and 3i-strongly
gt Then
MultiCorr((an:: Pn (vi) (wA)iT, P,A) 2 MultiCorr( P. (v%))
(15)
Denote A MultiCorr (ar; Pn 1, (vA); P. ^) MultiCor((aan Pn] (vAa)) Then for all ,I,..., NJ. we have
24 an
(V
2
wa Velt HP P)
O(A).
(16)
A(312²
Bt 2
Informally, Theorem 4.2 states that the better we solve the regularized dual problem, (1 the closer
expect each V #P to be to the barycenter TP in W2. It follows from 15 that final
objective (14) is umbiased: the optimal solution is obtained by {v/ )
4.4 PRACTICAL. ASPECTS AND OPTIMIZATION PROCEDURE
In practice, even if the choice of т, P does not satisfy 7 P 2 P. we observe the pushforward measures dP, often converge to P.. To partially bridge the gap between theory and practice, we dynamically update the measure so that after each optimization step set (for ry [0, 1))
P : (1-111))
Σ On [PP..
ie., the probability measure for is mixture of the given initial measure P and the current barycenter estimates (V HP,,) For the initial I one may use the barycenter of (WN(p, Egg It can be efficiently computed via an iterative fixed point algorithm Jivarez-Estebar 2020). During the optimization, these estimates become closer to the true barycenter and can thus improve the congruence regularizer 12
al. 2016 Chewi et al.,
We use mini-batch stochastic gradient descent to solve 14 where the integration is done by Monte- Carlo sampling from input measures {PP,) and regularization measure IP, similar to Li et al. (2020). We provide the detailed optimization procedure (Algorithm[1 I and discuss its computational complex- ity in Appendix|3. AA] In Appendix C.3 demonstrate that the impact of the considered regularization on our model: we show that cycle consistency and the congruence condition of the potentials are well satisfied.
EXPERIMENTS
The code is written PyTorch framework and is publicly available

---

Published as a conference paper at ICLR 2021
We compare method [CW,B] with the potential-based method [CRWB] by Li al. (2020) (with Wasserstein-2 distance and [ ('Ereuularization) and with the measure-based generative method
[SCW;B] by Fan al. (2020). All considered methods recover 2NV potentials { v {vh ei and approximate the barycenter as pushforward measures PP Regularization in [CRWB) allows access to the joint density of the transport plan, feature of their method that we not consider here. The method [SCW,B] additionally outputs a generated barycenter giß P where is the generative network and S is the input noise distribution.
To assess the quality of the computed barycenter, we consider the unexplained variance percentage defined as UVP(PP 100
%. When UVP 0%, IP is good approximation of P. For values 100%, the distribution Í is undesirable: a trivial baseline Ipo бzę(د) achieves UVP(P") 100%. Evaluating UVP in high dimensions is infeasible: empirical estimates of W3 are unreliable due high sample complexity (Weed al. 2019). To overcome this issue, for barycenters given by V dP, we use L2-UVP defined by
(F(F.P (P)
ot V %
CU-UVP(VO,, P.)
def
100
UVP( Vi 1P
(17)
Var
where the inequality in brackets follows from (Krotin et al. 2019a Lemma A.2). We report the weighted average of C2-UVP of all pushforward measures w.r.t. the weights On- For barycenters given an implicit form ydS, we compute the Bures- Wasserstein UVP defined by
BW2 2 (ott, P)
BW? BW3
def
100
UVP(g!S)
(18)
2 Var (P)
where (P Q) WW(N(pp, N En)) is the Bures- Wasserstein metric and we use μe, Ep to denote the mean and the covariance of a distribution (Chewi et al. 2202).. is known that BW2 lower-bounds w3 DDwsson & Landau 1982), so the inequality in the brackets of (18) follows.
detailed discussion of the adopted metrics is given in Appendix C.2
5.1 MIGH-DIMNSIOONLL LOCATION SACATION-SCATTER EXPERIMENTS
Input distributions (
(b) True burycenter
(c) SCW, generated distribation g
(d) SCW, B, distributions Vyll
CRWB, distribations Volie,
CW3B, distributions Vol
Figure Barycenter of location scatter Swiss roll population computed by three methods.

---

Published as a conference paper at ICLR 2021
In this section, we consider N 4 with (a1 (a).......)) 04) (0.1, 0.2, 0.4) weights. We consider
the location-scatter family of distributions Alltern et al. 2016 $4) whose true barycenter can be computed. Let Po E P2..c and define the following location-scatter family of distributions F(PP) = (fs.utPu Sa + with positive definite matrix S E &M When {IP,) C F(Po), their barycenter I is also an element of F(Fo) and can be computed via fixed-point iterations (Altarez-Estthaan et al. 2016).
E MDDD. u E }R"), where fs,u : IR - R. is a linear map fs, =
Figure Ta shows 2-dimensional location- lecation-ccatter family generated by using the Swiss roll distribution
as Po. The true barycenter shown Figure Tb The generated barycenter giS of [SCW, [C,B] is given
in Figure lc The pushforward measures V VIEPP of each method are provided in Figures ld le
1f respectively. In this example, the pushforward measures V all reasonably approximate I. whereas the generated barycenter gtS of [SCW,B] (Figure Ic] visibly underfits.
For quantitative comparison, we consider two choices for Po: the D-dimensional standard Gaussian distribution and the uniform distribution on /3, +/3) Each P,, is constructed as fsçAs.. E F(FP)). where S, is random rotation matrix and A is diagonal with entries [/ 3,6 ......... 2] where b D VI. We consider only centered distributions (i.e. zero mean) because the barycenter of non-centered [ E P2... (RR ) is the barycenter of [1P,, shifted by Σ an/MP, where {P", are centered copies of { (Alarrez-Esteban et al. 2016 Results are shown in Table and
In these experiments, our method outperforms [CRWB] and [SCW,B]. For [CRWB], dimension ~ 16 is the breakpoint: the method does not scale well to higher dimensions. [SCW2B] scales with the increasing dimension better, but its errors [2-UVP and BW3-UVP are twice as high as ours. This is likely due to the generative approximation and the difficult min-max-min optimization in [SCW2B]. For completeness, we also compare our algorithm to the proposed in Cuturi & Doucet (2014) which approximates the barycenter by discrete distribution points. In our experiment, similar to Li et al.] (2020). we set 5000 as the support size. As expected, the BW2-UVP error of the method increases drastically as the dimension grows and the method is outperformed by our approach.
fixed number of free-support
To show the scalability of our method with the number of input distributions N, we conduct analo- gous experiment with high-dimeensional location-scatter family for N = 20.. We set N1
2n
for n = 1, 2, ..... 20 and choose the uniform distribution on [-v3. + 3) D as Po and construct distri- butions F. F(P)) as before. The results for dimensions 32, 64 and 128 are provided in Table 3
Similar to the results from Tables I and 2
we see that our method outperforms the alternatives.
Metric
Method
D=2
16
128
256
BW3-UVP, %
JFCRBI, Cuturi & Doucet (2014)
0.7
0.68
1.41
3.87
8.85
14.08
18.11
21.33
[SCW,,B, B). (Fan et 2020
0.16
2.85
C2-UVP, & (potentials)
0.08
0.10
0.17
0.29
0.47
0.63
1.14
1.50
[CRWB], (Liet al
2020
0.99
232
8.2
2223
6701
>100
[C
0.66
003
0.0
0.11
0.19
0.24
0.42
0.83
Table 1: Comparison of UVP for the case (P,,) C F(Po).
Po
NO ID), N
4.
Metric
Method
D=2
16
32
64
128
256
BW UVP
[FCWB]] Cuturi & Doncet
2014)
0.64
0.77
122
3,75
8.92
14.3
18.46
21.64
2)CW2B), FFan al.
2020
0.12
0.10
0.19
0.29
0.46
0.6
1.38
2.9
UVP % (potentials)
0.17
0.12
0.2
0.31
0.47
0.62
1.21
1.52
[CRWB] [[ al 2020
0.58
1.83
8.00
21.23
55.17
> 100
[CW2B] WV,
0.17
0.08
0.06
0.1
0.2
0.25
0.42
0.82
Table 2: Comparison of UVP for the case {P C ((Po). Po = Uniform (1 (1 v 3,
+
+33 P).
N
Metrie
Method
D=32
64
128
[CWBL
[FCrBL Cuturi
Doncer
2014)
14.09
2622
38.43
BW .VE, %
UV-
[SCW BB] (Fan al.
0.60
0.62
0.86
0.93
1.83
1.22
(potentials)
[CWBB]. ours
0.31
0.58
1.15
Table 3: Comparison of UVP for the case (Pn) F(Po), Po = Uniform (1-v3, +V3
20.
5.2 SUBSET POSTERIOR AGGREGATION
We apply our method to aggregate subset posterior distributions. The barycenter of subset posteriors
converges the true posterior (Srivastava al. 2018). Thus, computing the barycenter of subset
8

---

Published as a conference paper at ICLR 2021
posteriors is efficient alternative to obtaining a full posterior in the big data setting (Srivastaval ct al. 2013 Staib et al.,
2017 Li et al. 2020).
Analogous to (Li ct al. 2020), consider Poisson and negative binomial regressions for predicting the hourly number of bike rentals using features such as the day of the week and weather conditions We consider the posterior the 8-dimensional regression coefficients for both Poisson and negative binomial regressions. We randomly split the data into = 5 equally-sized subsets and obtain 10 samples from each subset posterior using the Stan library (Carpenter al., (017). This gives the discrete uniform distributions {Pn} supported on the samples. As the ground truth barycenter P,, consider the full dataset posterior also consisting of 10 points.
We use BW3 BW3-UVP(TT P) to compare the estimated barycenter P (pushforward measure Ve or generated measure gtt) with the true barycenter. The results are in Table 4 All considered methods perform well (UVP< 2%), but our method outperforms the alternatives.
SCW,B,
[CRWB]
2020
CW,B.
Regression
P gr3
W..
BVV-
Puissum
0.67
0.41
1.53
0.1
aegative binomial
0.15
0.15
1.26
Table Comparison of UVP for recovered barycenters in our subset posterior aggregation task.
5.3 COLOR PALETTE AVERAGING
For qualitative study, apply method to aggregating color palettes of images. For an RGB image I. its color palette is defined by the discrete uniform distribution PIZ) of all its pixels E [0, 1| For images } we compute the barycenter
each color palette IP, = (P( w.rt. uniform weights an 2. We apply each computed potential V pixel-wise to I, to obtain the "pushforward" image Veli #In, These "pushforward" images should be close the barycenter P of ).
(a) Original images {Zu }
(b) Color palettes {P.. of original images.
(c) Images with averaged color paletle iZ.).
(d) Barycenter palettes (P..)
Figure Results our method applied to averaging color palettes of images.
The results are provided in Figure Note that the image V zz inherits certain attributes images In and Iy: the sky becomes bluer and the trees becomes greener. On the other hand, the sunlight in images Ville T has acquired orange tint, thanks to the dominance of orange in Ii.
ACKNOWLEDGMENTS
The Skoltech Advanced Data Analytics in Science and Engineering Group acknowledges the support of Russian Foundation for Basic Research grant 20-01 000003, Skoltech-MIT NGP initiative and thanks the Skoltech CDISE HPC Zhores cluster staff for computing cluster provision.
The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grant W91 INF2010168, of Air Force Office of Scientific Research award FA9550-19-1-031 of National Science Foundation grant I1S-1838071, from the CSAIL Systems that Learn program, from the MIT-IBM Watson AI Laboratory, from the Toyota-CSAII Joint Research Center, from gift from Adobe Systems, from an nIT. Immersion LaWNCSOFF Gaming Program seed grant, and from the Skoltech-MIT Next Generation Program.
http://archi uci http:///e.....eeeeeoooooaaaaaaaa

---

Published as a conference paper at ICLR 2021
REFERENCES
Martial Agueh and Guillaume Carlier. Barycenters the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-924, 2011.
Pedro C Alvareez-Esteban, Del Barrio, JA Cuesta- Albertos, and C Matrán. A fixed-point approach to barycenters in wasserstein space. Journal of Mathematical Analysis and Applications, (2): 744-762, 2016.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Bearning-Voluume 70, pp. 146-155. JMLR. org, 2017.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Commut nications on pure and applied muthematics, 44(4):375-417 1991.
Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Be- tancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76(1), 2017.
Sinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin Stromme. Gradient descent algorithms for wassersenn barycenters. arXiv preprint arXiv:200..0/70, 2020.
Samuel Cohen, Michael Arbel, and Mare Peter Deisenroth. Estimating barycenters of measures in high dimensions. arXiv preprint 1XXi:20077.07033 2020.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013.
Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. 2014.
DC Dowson and BV Landau. The fréchet distance between multivariate normal distributions. Journal of multivariate analysis, 12(3)450-455, 1982.
Pavel Dvurechenskii, Darina Dvinskikh, Alexander Gasnikov, Cesar Uribe, and Angelia Nedich. Decentralize and randomize: Faster algorithm for wasserstein barycenters. In Advances Neural Information Processing Systems, pp. 10760-10770, 2018.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. arXiv preprint #rXii:2007.04462 2020.
Werner Fenchel. On conjugate convex functions. Canadian Journal of Mathematics, 1(1):73-77, 1949.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. In Advances in neural information processing systems, pp. 3440-3448, 2016.
Sham Kakade, Shai Shalev-Shwartz and Ambuj Tewari. On the duality of strong convexity and strong smoothness: Learning applications and matrix regularization. Unpublished Manuscript, http://ttic. uchicago. http://hhhaappekaaaaaaaaaaalleee
pdf. 2(1), 2009.
Leonid Kantorovitch. On the translocation of masses. Management Science, 5(1):1-4, 1958.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. arXiv preprint 1909, 13999999.20,,,, 2019a.
Alexander Korotin, Vladimir V'yugin, and Evgeny Burnaev. Integral mixability: a tool for efficient online aggregation of functional and probabilistic forecasts. arXiv preprint 1912v:1912.0704m 2019b.
Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin Solomon. Continuous regularized wasserstein barycenters. arXiv preprint arXiv:2008. 12534, 2020.
10

---

Published as a conference paper at ICLR 2021
Ashok Vardhan Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason D Lee. Optimal transport mapping via input convex neural networks. arXiv preprint arXiv: 190::00008.00 2019.
Robert J McCann et al. Existence and uniqueness of monotone measure-preserrvng maps. Duke Mathematical Journal, 80(2):30--224 1995.
Youssef Mroueh. Wasserstein style transfer. arXiv preprint arXiv: /28288, 2019.
Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 6(1):147-160, 1994.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport. Foundations and Trends@ in Machine Learning, 1(5-66)3355607, 2019.
Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot. Wasserstein barycenter and its application to texture mixing. In International Conference on Scale Space and Variational Methods in Computer Vision, pp. 435 446. Springer, 2011.
Julien Rabin, Sira Ferradans, and Nicolas Papadakis. Adaptive color transfer with relaxed optimal transport. 2014 IEEE International Conference on Image Processing (ICIP), pp. 4852-4856. IEEE, 2014.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathicu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv:1711.02288. 2017.
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.
Sanvesh Srivastava, Volkan Cevher, Quoc Dinh, and David Dunson. Wasp: Scalable bayes via
barycenters of subset posteriors. In Artificial Intelligence and Statistics, pp. 912-920, 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.
Matthew Staib, Sebastian Claici, Justin M Solomon, and Stefanie Jegelka. Parallel streaming Wasserstein barycenters. In Advances in Neural Information Processing Systems, pp. 2647-2658,
2017.
Cédric Villani. Topics in optimal transportation. Number 58. American Mathematical Soc., 2003.
Jonathan Weed, Francis Bach, al. Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance. Bernoulli, 25644))2620-2668 2019.
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, pp. 6882 6892, Long Beach, California, USA, 09-15 Jun 2019. PMLR. URL http: //proceedings mir press wwtp//////////..

---

Published as a conference paper at ICLR 2021
A THE ALGORITHM
The numerical procedure for solving our final objective
14
is given below.
Algorithm 1: Numerical Procedure for Optimizing Multiple Correlations (14)
Input: Distributions ....... PN with sample access; Weights ................................ ax with Σ On = 1;
Regularization distribution IP' given by sampler;
Congruence regularizer coefficient Balancing coefficient y E [0, 1];
1;
Cycle-consistency regularizer coefficient > 0:
2NV ICNNs {1b9., k:
Batch size K > 0:
for / 1, 2.... do
1. Sample batches X,, ~ P,, for all N; Compute the pushforwards Y, VX, for all n 1, I,
3. Sample batch Yo P:
4. Compute the Monte-Carlo estimate of the congruence regularizer:
1 K
N
11
Loongrience ==
One Ww.. (v)
2
YEY,
n/=1
where To = 7 and in = On," (1- for n 1,,,=,......... 5. Compute the Monte-Carlo estimate of the cycle-consistency regularizer:
N: ,,.....,y
Lcyele :- KK-
1
1 Σ an
Σ
Vieua (v). xx13
6. Compute the Monte-Carlo estimate of multiple correlations:
N
Σ (xc, Vioo, - then (Voo.
Muultioor :
On K
zEX,
7. Compute the total loss:
Lioial == LMultiCorr + À LCyele + T LCongruence:
8. Perform gradient step over (6n, by using ((R.mm)))
end
Parametrization the potentials. To parametrize potentials {v,,, tw. }. we use DenseICNI (dense input convex neural network) with quadratic skip connections; see (Korotin al. 2019a Appendix B.2). As an initialization step, we pre-train the potentials to satisfy
the (x) ~ |0-12
18/1 2
and
(w.(2)) ~
Such pre-training provides a good start for the networks: each vhe is approximately conjugate to the corresponding Vh... On the other hand, the initial networks {vo... are approximate congruent according to (5).
Computational Complexity. For a single training iteration, the time complexity of both forward (evaluation) and backward (computing the gradient with respect to the parameters) passes through the objective function (14) is ((NT). Here N is the number of input distributions and T is the time taken by evaluating each individual potential (parameterized as neural network) on a batch points sampled from either P,, or P. This claim follows from the well-known fact that gradient evaluation (aho( of hg RDD - R. when parameterized neural network, requires time proportional
12

---

Published as a conference paper at ICLR 2021
to the size of the computational graph. Hence, gradient computation requires computational time proportional to the time for evaluating the function hz(2) itself. The same holds when computing the
derivative with respect to æ. Then, for instance, computing the term V V (4) in (14) takes O(T) time. The gradient of this term with respect to 0 also takes (T)) time: ressian-vecttor products that appear can be calculated in O(T) time using the famous Hessian trick, see Pearlmutter (1994).
In practice, compute all the gradients using automatic differentiation. We empirically measured that for our DenseICNN potentials, the computation of their gradient w.r.t. input æ, i.c., V requires roughly 3-4x more time than the computation of (x)).
B PROOFS
In this section, we prove our main Theorems 4.1 and 4.2
We use (R. P + RR, (i) denote the Hilbert space of functions f PR. -R" with integrable square w.r.. a probability measure The corresponding inner product for f1, E (R + 18 denoted by
(2 def [ (f((())
where (f), (2()) is the Euclidean dot product. We use || lle Vis denote the norm induced by the inner product in (R R.P, .)).
We also recall a useful property of lower semi-continuous convex function &i RR + R:
Vi((2) arg max (v,,x) v(v)),
(19)
yERP
which follows from the fact that
arg max
Ky. - {))
( 0.
VER
We begin with the proof of Theorem 4.1
Proof. We consider the difference between the estimated correlations and true ones:
N
(i
Σ an (x)))P,,))))
an
R
JRD
[(Vvn(z), To M. (V)
On
R
N
On
(V)) - (z)
(20)
RDD
==1
where we twice use 119 for f
= thi and = 2 We note that
N
Σ
On
V
(x)
x
(x) = Σ an RRD
(v. (eT()))(F(() =
JDD
n=1
a. Ven = /
(y. v))d(((y) = |lidg
(21)
BB
RD
n=1
where we use of change-of-variable formula for HP and Analogously,
N
On
Σ an
Vi (v)dP(y) =
N
llvl
On ti (v) =
dP(y)
ild 112
(22)
JBD
2
2
n=
13

---

Published as a conference paper at ICLR 2021
Since each Vn is B-smohth, we conclude that vn is -sttonngl convex, (Kakade al. 2009).
Thus, we have
An (vi(())))) z
Vinl o (x)) VoN(2) - Von(z)) + 1
4n (Ven(z))) +
(x) - Vıhn
2B
2h (VHn(()) (2, (.)())) +
V (x) [ (23)
2B
or equivalently
- (VVn(())))) 2 - S. (z)) + (2. (o (x)(
- (24)
28
We integrate
24
w.r.t. Ir, and sum over 1, ,..,, N with weights an: an:
in (v)(())))P) 2
On
JBD
N
On
Vn
(VW) +
Σ
an (n (o)()))p
Σ an Vtin +
J
an 2B
V
Onp Fn (V)dP(y) + Σ (2 (oh(s)),
RDD n=1
an (v)()))
On 2B
(vi(() (.) (25)
We note that
{
|1y1
|1y/|
On Vn (y)dT(())
110
On vn dP(()
dT(())
2
R=1
n=1
|
N
an an (v)) dP(())
|idg 1
(26)
RD
2
Now we substitute 25
26
21 and
22
into
20
to obtain
Next, we prove Theorem 4.2
Proof. Since vi is 3 strongly convex, its conjugate vi is | smooth, i.e. has 1 Lipschitz gradient
Vel. Kakade et al. 2009). Thus, for all x, x' e ::
[VVV(()) - vet = 3t
)2 ||x-~111-
We substitute x' Veh (v)(v) = (Vw!) Vesi (v) and obtain:
(Vol()) ( - Vol Vo((2))
(27)
Since the function vă is Bt smooth, we have for all x E RP:
(VV((()) & vl(V)(()) + Virl (x), (o)( - vod(2)) + B
(V))(())) -
14

---

Published as a conference paper at ICLR 2021
that is equivalent to:
B -oll()- (28)
(a, - 2 (z, (V -
-
-2
(1)
We combine
28 with 27
to obtain
B 2(
(2, (x)) (Vvl w ())
lidgo
V V
(29)
For every n 2,2..... N integrate 29 w.r.t. P, and sum up the corresponding cycle-consistency regularization term:
100 - A
V
idR
B
N (x)dd + (A
id 0
(30)
2(31
R Te (vA1.
We sum 300 for -1,2, ..,,, N w.r.t weights at, to obtain:
N
N U dF) +
an RS F
(v wb)
an
RD
n=1
n=1
Bt
an
2' (x)dPP
on
2(3122
R (vt 25)..
MultiCo(((a, Pu
We add 7 R (( to both sides of 31 get
MultiCorr ((an) Pnd (xtt) J:,, IP, A) 2 MultiCorr((a,, Pn (1th)) +
Bt 2(33)2)
an
Ry (vit v.)).
(31)
We substract MultiCor(((a,, P.) {(+t)) from both sides and use Theorem to obtain
|
3
A
a (3)
dF
a
Vei(()) vo)())
(32)
2
N
Bt
T. T
+
an () 2(31)2) R (vi., w 2
(33)
n=1
N
Bt
3
N
On (A
R.P (wit ( +
(x) (z) (x)
(34)
2(31
2
1=
In transition from
33 to
34
we explot the fact that the sum of the first term of 32 with the g²
regularizer 7 ((v.)) Since A
218rjT from (34) we immediately conclude A > 0; i.c. the
multiple correlations upper bound holds true. the other hand, for every n 1, 1,,, N we have
2A
2A
(x) (x)
< ai
and V v o Vist idg 113, S an () 2(,81a²)
(35)
We combine the second part of
35 with
27
integrated w.r.t. P.::
2A On (A(St)2
V
Vot
(36)
B1
15

---

Published as a conference paper at ICLR 2021
Finally, use the triangle inequality for || lle, and conclude
| - Vel
S | th V
+
| Visht
Venile.
2A
(37)
On
Bt
((31)2
B1
i.e.
24
1
P)
Ve
Vol
+
O(A).
an
3t
{(3F)2
where the first inequality follows from Korotin et al.
2019a
Lemma A.2).
C EXPERIMENTAL DETAILS AND EXTRA RESULTS
In this section, provide experimental details and additional results. In Subsection demonstrate qualitative results of computed barycenters in the 2-dimensional space. In Subsection C.2 discuss used metries more detail. In Subsection C..4 list the used hyperparameters of our method (CW,B) and methods [SCW2B], [CRWB].
C. ADDITIONAL TOY EXPERIMENTS IN 2D
We provide additional qualitative examples of computed barycenters of probability measures R.2.
In Figure consider the location-scanl scatter family F(Fo) with Po Uniform[-Vi v In principle, all the methods capture the true barycenter. However, the generated distribution giS of (SCW,B (Figure 3c) provides samples that lies outside of the actual barycenter's support (Figure 3b Also, [CRWB] method, one the potentials' pushforward measure (top-right in Figure 3el has visual artifacts.
Input distributions (P.))
(h) True barycenter
(c) SCW.R penerated distribution
(d) ,B distributions Vel.p,
(c) CRWB, distributions V
( C3 distributions Vel. HP,
Figure 3: Barycenter of a random location-scatter population computed by different methods.
In Figure we consider the Gaussian Mixture example by (Fan et al. 2020). The barycenter computed (SCW,B) method (Figure 4b) suffers from the behavior similan lar mode collapse.
16

---

Published as a conference paper at ICLR 2021
Inputs (P.))
(b) 2CW
SCW VoldF.
(d) CRWB VV
CW Vel :P.
Figure Barycenter of a two 2D Gaussian mixtures.
C.2 METRICS
The unexplained variance percentage (UVP) (introduced in Section is natural and straightforward metric to assess the quality of the computed barycenter. However, it is difficult to compute high dimensions: it requires computation of the Wasserstein-2 distance. Thus, use different but highly related metrics -VV and -VV.
To access the quality of the recovered potentials { we use 22-UVP defined in (17). CVV compares not just pushforward distribution Vr n. with the barycenter P. but also the resulting transport map with the optimal transport map V It bounds UVP(V from above, thanks to (Korotin et al. 2019a Lemma A.2). Besides, LVU naturally admits unbiased Monte Carlo estimates using random samples from P..
For meesure-based optimization method, we also evaluate the quality of the generated measure gtS using Bures Wassessnsser UVP defined in (18). For measures PP, Q whose covariance matrices are not degenerate, BWy given by
BW3 (P, Q
1pp uell
Tr Ep Tr F2 Tr 2
Bures- Wassessesnrrtnin metric compares PP, Q by considering only their first and second moments. It is known that B BW3(FP Q) is a lower bound for W2(P, .). see (Dowson & Landau 1982). Thus, we estimate means
have BW3-UVP UV)((E)) S UVP(gtS). In practice, to compute BW/-UVP(gtSS and covariance matrices distributions by using 10 random samples.
C.3 CYCLE CONSISTENCY AND CONGRUENCE IN PRACTICE
To assess the effect the regularization cycle consistency and the congruence condition in practice, we run the following sanity checks.
For cycle consistency, for each input distribution P. we estimate (by drawing samples from P,,) the value V o (v) that we used for assessing the resulting transport maps. In all the experiments, this value does not exceed 2%, which means that cycle consistency and hence conjugacy are satisfied well.
/Va((/....... This metric can be viewed analog of the C²-UVP
For congruence condition, we need to check that >.Y Oi (xr) = |2/2. However, we do not know any straightforward metric to check this exact condition that is scaled properly by the variance of the distributions. Thus, we propose to use an alternative metric to check a slightly weaker condition
on gradients, e.g., that V an Voli (x) = T. This is weaker due to the ambiguity of the additive constants. For this compute D Or, V (2) /Va(())))) where the denominator is
17

---

Published as a conference paper at ICLR 2021
the variance of the true barycenter. We computed this metric and found that it is also less than 2% in all the cases, which means that congruence condition mostly satisfied.
C.4 TRAINING HYPERPARAMETERS
The code is written using the PyTorch framework. The networks are trained single GTX 1080Ti.
C.4.1 WASSERSTEIN-2 CONTINUOUS BARYCENTERS (CW2B, OUR METHOD)
Regularization. We use and P (0, Ip) our congruence regularizer T R.. We use 10 for the cycle regularization À R5 for all n = 1,2........
Neural Networks (Potentials). To approximate potentials fel, vū} } in dimension D, we use DenseICNN[2: max(64, 2D), max(64, max D)]
with CELU activation function. DenselCNN an input-convex dense architecture with additional convex quadratic skip connections. Here the rank of each input-quadratic skip-connection' Hessian matrix. Each following number max(:,) -) represents the size of hidden dense layer in the sequantial part of the network. For detailed discussion of the architecture (Korotin al. 2019a Section B.2).
Training process. We perform training according A of Appendix A We set batch size K = 1024 and balancing coefficient = 0.2. We use Adam optimizer by (Kingma & Ba 2014) with a fixed learning rate 10-3 The total number of iterations set to 50000.
C.4.2 SCALABLE COMPUTATION OF WASSERSTEIN BARYCENTERS (SCW,B)
Generator Neural Network. For the input noise distribution of the generative model we use S (0, For the generative network g RD + RD we use a fully-connected sequential ReLU network with hidden layer sizes
[max 2D), max 100, 2D), max [000, 2D)].
Before the main optimization, we pre-train the network to satisfy g(z) for all E R.D. This has been empirically verified as better option than random initialization of network's 's weights.
Neural Networks (Potentials). We used exactly the same networks as Subsection C.4.1
Training process. We perform training according to the min-max-min procedure described by (Fan al. 2020 Algorithm 1). The batch size set to 1024. We use Adam optimizer by (Kingma Ba 2014) with fixed learning rate 10-3 for potentials and 10-4 for generative network g. The number of iterations of the outer cycle (min-max-min -mi) number of iterations is set to 15000. Following (Fan) al. 2020), we use 10 iterations per the middle cycle (min-max-min) and 6 iterations per the inner cycle mmin-max-mint.
C.4.3 CONTINUOUS REGULARIZED WASSERSTEIN BARYCENTERS (CRWB)
Regularization. [CRWB] method uses regularization to keep the potentials conjugate. The authors impose entropy or L regularization w.r.t. some proposal measure IP; see (Li al. 2020 Section 3) for more details. Following the source code provided by the authors, we use La regularization (empirically shown as a more stable option than entropic regularization). The regularization measure set to be the uniform measure on box containing the support of all the source distributions, estimated by sampling. The regularization parameter is set to 10 14-4.
Neural Networks (Potentials). To approximate potentials {wi ) in dimension D, we use fully- connected sequential ReLU neural networks with layer sizes given by
(maax128 4D), max(128, 4D), www.maaxmm. (128, 4D)].
We have also tried using DenselCNN architecture, but did not experience any performance gain.
Training process. We perform training according to Li et al. 2020 Algorithm 1). We set batch size to 1024. We use Adam optimizer by (Kingma & Ba 2014) with fixed learning rate 10 10-3, The total number of iterations set to 50000.
18


Reviews:
[{"review_id": "tTB4109eyKo", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "Summary:\n\nThe paper considers the Wasserstein Barycenter problems in the continuous setting. In particular, the authors propose an algorithm to compute the Wasserstein-2 barycenter when only samples from the marginals are accessible. Some theoretical analysis of this method is presented. Several numerical examples are carried out to compare this method with two other recently proposed methods.\n\nReasons for score:\n\nThe proposed algorithm utilizes an interesting regularization of the dual formulation of Wasserstein-2 Barycenter, resulting in a single minimization problem instead of a min-max problem. This algorithm is properly justified by theoretical results as well as numerical experiments. \n\nPros:\n\n1. The paper provides theoretical results on the consistency of the proposed algorithm.\n\n2. The experiments are overall good and clear. \n\n2. The paper is well-written and easy to follow.\n\nCons:\n\n1. High-dimensional examples other than the simple Gaussian setting are missing.\n\n2. There is no analysis of computational complexity of the proposed algorithm. Also, the training expense is not reported.\n \n3. The double gradient in the second regularization term could be expensive to evaluate.\n\nQuestions: \n\n1. In both Theorem 4.1 and 4.2, the smoothness of the potentials is crucial. If the smoothness B is too large, then the bound presented is essentially useless. Please comment on it.", "extra_scores": {}, "extra_text": {}}, {"review_id": "k0xNV0z_iUi", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "This work introduces a new Wasserstein-2 barycenter computation method. The authors first derive the dual formulation of the Wasserstein-2 barycenter problem, and then parametrize the convex potentials by ICNNs. The congruent and conjugacy conditions are enforced by regularization terms, respectively. They then show that the algorithm can find a good barycenter if the objective function is properly minimized.\n\nPros:\n1. The algorithm does not introduce bias.\n2. The algorithm does not require minimax, which is efficient.\n3. The empirical performance is much better than existing methods, probably due to the above two reasons.\n\nAreas to improve:\n1. It is good that the empirical analysis include how the performance change w.r.t. D. It would be better if there is a similar analysis to N. Furthermore, since 2N ICNNs are needed to be trained, it would be better if the training time is also reported, so that we can have a more comprehensive understanding of the method. Will there be a setting that discrete method can be faster than the proposed method to enforce comparable approximation error (say, large N for 3D applications)?\n2. Since the congruent and conjugacy conditions are enforced by regularizations, they are not guaranteed to be satisfied. Therefore, it would be better if there is an experiment showing that how the conditions are satisfied. \n3. The first section of related work should also briefly include https://arxiv.org/abs/1605.08527 and https://arxiv.org/abs/1905.00158.\n\nAfter rebuttal:\n\nThe additional experiment results provided in the rebuttal stage suggests the efficiency of the proposed method, as well as the congruent and conjugacy conditions are approximately satisfied. I therefore believe this paper should be accepted.", "extra_scores": {}, "extra_text": {}}, {"review_id": "e406EvNzDIN", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "This paper proposes a method to scalably compute Wasserstein-2 barycenters given samples from input measures. In general the authors also allow for continuous measure settings. Inspired by Li et al. (2020) the paper uses a potential-based approach and recovers the barycenter by using gradients of the potentials as pushforward maps.\n\nIn general, I feel this paper is well-written and provides a fast solution to a meaningful problem, thereby supporting the claim of novelty. The theoretical developments in the paper are reasonable and the experiments carried out are quite decent, both in simulation and real-data settings.\n\nThe only point that bothers me is the approximation used. It would be great if the authors could give an extensive and detailed understanding of settings where the upper bound in Eq.(10) in the main text is small thereby leading to a good approximation. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "WO0Wjhtxges", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "**Summary**\nThe paper derives the barycenter mapping problem as an optimization over *congruent* convex functions---each convex potential corresponding to a component distribution.  Congruency is a property on the set of optimal potential functions that ties them together.  However, this optimization is quite challenging and so the paper derives an principled objective function that includes two regularization terms.  The first regularization term encourages congruency of the set of convex functions and can be seen as a variational bound on an ideal congruency regularization.  The second regularization term encourages the pairs of convex functions to be conjugate.  The paper proves that the optimal solution of this objective is the true potentials and thus no bias is introduced.  The proposed approach is demonstrated on the tasks of generative modeling (2-256 dimensions), posterior inference, and color pallete barycenters (3D)\n\n**Strengths:**\n- Nice problem formulation and setup with respect to prior methods.\n- The derivation of the final objective function is clearly laid out and well-motivated.  Each problem that is encountered is explained and then a solution or approximation is introduced.\n- The theoretical results give appropriate grounding for the approach.\n- The empirical results outperform prior potential-based methods for barycenters.\n\n**Weaknesses:**\n- It is unclear if this method can scale in terms of samples and dimensions.  What is the computational cost of estimating these input convex neural networks?  Can you provide (approximate) wall-clock times for the various methods and dimensionalities? What are the key computational bottlenecks either memory-wise or computation-wise?\n- The experiments seem small scale with a max dimension of 256.  Barycenters for high dimensional real-world data (e.g., even MNIST (784D)) or some other high-dimensional real-world dataset would improve the paper.\n- The paper lacks comparison to methods that do not recover 2N potential functions. What are the closest methods for barycenter that do not use potential functions?  For example, could the algorithms be compared to discretized barycenter algorithms to show the breakdown in higher dimensions?\n\n**Other comments or questions**\n- Is $D$ in equation 5 supposed to be $N$?\n- Some typos above Eqn. 12.\n\n**Update after author response**\n\nI appreciated the authors response to the scalability and raw computation times.  Thank you also for the additional comparison to a non-potential function method.  This will be a good comparison.  My main concerns were answered, and I still think this is a good paper.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose the 2-Wasserstein barycenter problem between measures. The authors propose a novel formulation that leverages a condition (congruence) that the optimal transport (Monge) maps, here parameterized as potentials, must obey at optimality. The introduce various regularizers to encourage that property. The idea is demonstrated on convincing synthetic experiments and on a simple color transfer problem. Although experiments are a bit limited, I do believe, and follow here the opinion of all reviewers, that there is novelty in this approach, and that this paper is a worthy addition to the recent line of work trying to leverage ICNNs/Brenier's theorem to solve OT problems.
Decision: Accept (Poster)