Wasserstein Iterative Networks for Barycenter Estimation
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Wasserstein barycenters have become popular due to their ability to representthe average of probability measures in a geometrically meaningful way.
In thispaper, we present an algorithm to approximate the Wasserstein-2 barycentersof continuous measures via a generative model.
Previous approaches rely onregularization (entropic/quadratic) which introduces bias or on input convex neuralnetworks which are not expressive enough for large-scale tasks.
In contrast, ouralgorithm does not introduce bias and allows using arbitrary neural networks.
Inaddition, based on the celebrity faces dataset, we construct Ave, celeba!
datasetwhich can be used for quantitative evaluation of barycenter algorithms by usingstandard metrics of generative models such as FID.
Introduction
The ﬁgure shows random samples from the input subsets and generated images from the barycenter.
Wasserstein barycenters [1] provide a geometrically meaningful notion of the average of probability measures based on optimal transport (OT,see [59]).
Methods for computing barycentershave been successfully applied to various practical problems.
In geometry processing, shapeinterpolation can be performed by barycenters[53].
In image processing, barycenters are usedfor color and style translation [44, 39], texturemixing [45] and image interpolation [27, 51].
In language processing, barycenters can be applied to text evaluation [12].
In online learning,barycenters are used for aggregating probabilistic forecasts of experts [26, 42].
In Bayesianinference, the barycenter of subset posteriorsconverges to the full data posterior [54, 55] allowing efﬁcient computation of full posteriorbased on barycenters.
In reinforcement learning,barycenters are used for uncertainty propagation[36].
Other applications are data augmentation[6], multivariate density registration [7], distributions alignment [20], domain generalization [34]and adaptation [38], model ensembling [13].
The bottleneck of obtaining barycenters is the computational complexity.
For discrete measures, fastand accurate barycenter algorithms exist for low-dimensional problems; see [43] for a 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 [29, 15, 25] are based on entropic/quadratic regularizationor parametrization of Brenier potentials with input-convex neural networks (ICNNs, see [4]).
Theregularization-based [ CR W B ] algorithm by [29] recovers a barycenter biased from the true one.
Algorithms [ C W 2 B ] by [25] and [ SC W 2 B ] by [15] based on ICNNs resolve this issue, see [25,
Tables 1-3].
However, despite the growing popularity of ICNNs in OT applications [35, 23, 37], theycould be suboptimal architectures according to a recent study [24].
According to the authors, moreexpressive networks without the convexity constraint outperform ICNNs in practical OT problems.
Furthermore, evaluation of barycenter algorithms is challenging due to the limited number ofcontinuous measures with explicitly known barycenters.
It can be computed when the input measuresare location-scatter (e.g. Gaussians)
[3, (cid:77) 4] or 1 -dimensional [8, (cid:77) 2.3].
Recent works [29, 25, 15]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 matrixinversion and square root extraction [2, Algorithm 1] with the cubic complexity in the dimension.
Contributions.
• We develop a novel iterative algorithm ( (cid:77) 4) for estimating Wasserstein-2 barycenters based on
the ﬁxed point approach by [3] combined with a neural solver for optimal transport [24].
Unlike
predecessors, our algorithm does not introduce bias and allows arbitrary network architectures.
• We construct the Ave, celeba!
( ave raging celeb rity faces,
(cid:77)
5) dataset
consisting of 64 × 64 RGBimages for large-scale quantitative evaluation of continuous Wasserstein-2 barycenter algorithms.
The dataset includes 3 subsets of degraded images of faces (Figure 1).
The barycenter of thesesubsets corresponds to the original clean faces.
Our algorithm is suitable for large-scale Wasserstein-2 barycenters applications.
The developeddataset will allow quantitative evaluation of barycenter algorithms at a large scale improving transparency and allowing healthy competition in the optimal transport research.
Notation.
We work in a Euclidean space R D for some D .
All the integrals are computed over 63 R D unless stated otherwise.
We denote the set of all Borel probability measures on R D with ﬁnitesecond moment by P 2 ( R D ) .
We use P 2 , ac ( R D ) ⊂
P 2 ( R D ) to denote the subset of absolutelycontinuous measures.
We denote its subset of measures with positive density by P +2 , ac ( R D ) ⊂
P 2 , ac ( R D ) .
We denote the set of probability measures on R D × R D with marginals P and Qby Π( P , Q ) .
For a measurable map T : R D
→ R D , we denote the associated push-forwardoperator by T(cid:93) .
For φ : R D → R , we denote by φ its Legendre-Fenchel transform [16] deﬁned byφ ( y ) =
max x ∈ R D [ (cid:104) x, y (cid:105) − φ ( x )] .
Recall that φ is a convex function, even when φ is not.
Preliminaries
Wasserstein-2 distance .
For P , Q ∈P 2 ( R D ) , Monge’s primal formulation of the squared
Wasserstein-2 distance, i.e., OT with quadratic cost , is
[Equation]
(1) where the minimum is taken over measurable functions (transport maps) T : R D
→ R D mapping Pto Q .
The optimal T ∗ is called the optimal transport map .
Note that (1) is not symmetric, and thisformulation does not allow mass splitting.
That is, for some P , Q ∈ P 2 ( R D ) , there might be no map
T that satisﬁes T(cid:93) P = Q .
Thus, [21] proposed the following relaxation:
[Equation]
where the minimum is taken over all transport plans π , i.e., measures on R D × R D whose marginalsare P and Q .
The optimal π ∗ ∈ Π( P , Q ) is called the optimal transport plan .
If π ∗ is of the form[ id , T ∗ ] (cid:93) P ∈ Π( P , Q ) for some T ∗ , then T ∗ minimizes (1).
The dual form [58] of W 22 is:
[Equation]
where the maximum is taken over u ∈ L 1 ( P ) , v ∈ L 1 ( Q ) satisfying u ( x ) + v ( y ) ≤ 12 (cid:107) x − y (cid:107) 2 forall x, y ∈ R D .
The functions u and v are called potentials .
There exist optimal u ∗ , v ∗ satisfyingu ∗ =
( v ∗ ) c , where f c ( y ) def = min x ∈ R D (cid:2) 12 (cid:107) x − y (cid:107) 2 − f ( x ) (cid:3) is the c -transform of f .
We rewrite (3) as
[Equation]
where the maximum is taken over all v ∈ L 1 ( Q ) .
It is customary [59, Cases 5.3 & 5.17] to deﬁneu ( x )= 12 (cid:107) x (cid:107) 2 − ψ ( x ) and v ( y )
=
12 (cid:107) y (cid:107) 2 − φ ( x ) .
There exist convex optimal ψ ∗ and φ ∗ satisfyingψ ∗ =
φ ∗ and φ ∗ = ψ ∗ .
If P ∈ P 2 ,ac ( R D ) , then the optimal T ∗ of (1) always exists and can berecovered from the dual solution u ∗ (or ψ ∗ ) of (3): T ∗ ( x )
= x − ∇ u ∗ ( x ) = ∇ ψ ∗ ( x ) [49, Theorem1.17].
The map T ∗ is a gradient of a convex function, see the Brenier Theorem [9].
Wasserstein-2 barycenter .
Let P 1 , . .
. , P N ∈ P 2 ,ac ( R D ) such that at least one of them has boundeddensity.
Their barycenter w.r.t.
weights α 1 , . .
. , α N ( α n > 0 ; (cid:80)
Nn =1 α n = 1 ) is given by [1]:
[Equation]
The barycenter P exists uniquely and P ∈ P 2 ,ac ( R D ) .
Moreover, its density is bounded [1, Deﬁnition3.6 & Theorem 5.1].
For n ∈ { 1 , 2 , . . . , N } , let T P → P n = ∇ ψ ∗ n be the OT maps from P to P n .
Thefollowing holds P -almost everywhere:
[Equation]
see [3, (cid:77) 3].
If P ∈ P +2 , ac ( R D ) , then (6) holds for every x ∈
[Equation]
D , i.e.
[Equation]
We call such convex potentials ψ ∗ n congruent .
Related Work
Below we review existing continuous methods for OT.
In (cid:77) 3.1, we discuss methods for OT problems(1), (2), (3).
In (cid:77) 3.2, we review algorithms that compute barycenters (5).
3.1 Continuous OT Solvers for W 2
We use the phrase OT solver to denote any method capable of recovering T ∗ or u ∗ (or ψ ∗ ).
Primal-form solvers based on (1) or (2), e.g., [61, 32], parameterize T ∗ using complicated generativemodeling techniques with adversarial losses to handle the pushforward constraint T(cid:93) P = Q in theprimal form (1).
They depend on careful hyperparameter search and complex optimization [33].
Dual-form continuous solvers [17, 50, 41, 57, 23] based on (3) or (4) have straightforward optimiza-
[Figure]
tion procedures and can be adapted to various tasks without extensive hyperparameter search.
A comprehensive overview and a benchmark of dual-form solvers are given in [24].
According to theevaluation, the best performing OT solver is reversed maximin solver (cid:98) MM:R (cid:101) , a modiﬁcation of theidea proposed by [41] in the context of Wasserstein-1 GANs [5].
In this paper, we employ this solveras a part of our algorithm.
We review it below.
Reversed Maximin Solver .
In (4), v c ( x ) can be expanded through v via the deﬁnition of c -transform:
[Equation]
In (7), the optimization over y ∈ R D is replaced by the equivalent optimization over functions
T : R D → R D .
This is done by the interchanging of integral and minimum, see [46, Theorem 3A].
The key point of this reformulation is that the optimal solution of this maximin problem is given by( v ∗ , T ∗ ) , where T ∗ is the OT map from P to Q , see discussion in [24, (cid:77) 2] or [48, (cid:77) 4.1].
In practice,the potential v : R D → R and the map T : R D → R D are parametrized by neural networks v ω , T θ .
To train θ and ω , stochastic gradient ascent/descent (SGAD) over mini-batches from P , Q is used.3.2 Algorithms for Continuous W 2 Barycenters
Variational optimization.
Problem (5) is optimization over probability measures.
To estimate P , onemay employ a generator G ξ : R H → R D with a latent measure S on R H and train ξ by minimizing
[Equation]
Optimization (8) can be performed by using SGD on random mini-batches from measures P n and S .
The difference between possible variational algorithms lies in the particular estimation method for 121 W 22 terms.
To our knowledge, only ICNN-based minimax solver [35] has been used to compute W 22in (8) yielding [ SC W 2 B ] algorithm [15].
Potential-based optimization.
[29, 25] recover the optimal potentials { ψ ∗ n , φ ∗ n } for each pair ( P , P n )via a non-minimax regularized dual formulation.
No generative model is needed: the barycenter isrecovered by pushing forward measures using gradients of potentials or by barycentric projection.
However, the non-trivial choice of the prior barycenter distribution is required.
Algorithm [ CR W B ] by[29] uses entropic or quadratic regularization and [ C W 2 B ] algorithm by [25] uses ICNNs, congruenceand cycle-consistency [23] regularization.
Iterative W 2 -Barycenter Algorithm
Our proposed algorithm is based on the ﬁxed point approach by [3] which we recall in (cid:77) 4.1.
In (cid:77) 4.2,we formulate our algorithm for computing Wasserstein-2 barycenters.
In (cid:77) 4.3, we show that ouralgorithm generalizes the variational barycenter approach.
4.1 Theoretical Fixed Point Approach
Following [3], we deﬁne an operator H : P 2 ,ac ( R D ) →P
2 ,ac ( R D ) by H ( P ) =
[ (cid:80) Nn
=
α n T P
→ P n ] (cid:93) P ,where T P → P n denotes the OT map from P to P n .
The measure H ( P ) obtained by the operator is indeedabsolutely continuous, see [3, Theorem 3.1].
According to (6), the barycenter P deﬁned by (5) is aﬁxed point of H , i.e., H ( P )
=
P .
This suggests a way to compute P by picking some P ∈ P 2 ,ac ( R D )and recursively applying H until convergence.
However, there are several challenges :
(a) A ﬁxed point P ∈ P 2 , ac ( R D ) satisfying H ( P ) =
P may be not the barycenter [3, Example
The situation
is
analogous to that
of the
iterative
k -means
algorithm
for a
different problem
clustering
There may be ﬁxed points which are not globally optimal.
(b) The sequence {H k ( P ) } k is tight [3, Theorem 3.6] so it has a subsequence converging inP 2 ,ac ( R D ) , but the entire sequence may not converge.
Nevertheless, the value of the objective
(5) decreases for H k ( P ) as k → ∞ [3, Prop. 3.3].
(c) Efﬁcient parametrization of the evolving measure H k ( P ) is required.
Moreover, to get H k +1 ( P )
from H k ( P ) , one needs to compute N optimal transport maps T H k
( P ) →
P n which can be costly.
In [11] and [2], the ﬁxed point approach is considered in the Gaussian case where the sequence
H k ( P ) is guaranteed to converge to the unique ﬁxed point – the barycenter.
The Gaussian case alsomakes parameterization (c) simple since both measures P n and H k ( P ) can be parametrized by meansand covariance matrices, and the maps T H k
( P ) → P n are linear with closed form.
For general continuous measures P n , it remains an open problem to ﬁnd sharp conditions on inputs 152 P n and the initial measure P of the ﬁxed-point iteration for the sequence {H k ( P ) } k to converge tothe barycenter.
In this work, we empirically verify that the ﬁxed point approach works well for theinput measures that we consider and for a randomly initialized generative model representing theevolving barycenter ( (cid:77) 4.2).
We tackle challenge (c) and develop a scalable optimization procedurethat requires only sample access to P n ∈P 2 ,ac ( R D ) .
[Figure]
4.2 Practical Iterative Optimization Procedure
We employ a generative model to parametrize the evolving measure, i.e., put P ξ =
G ξ (cid:93) S , where S isa latent measure, e.g., S = N
(0 , I H ) , and G ξ is a neural network R H
→ R D with parameters ξ .
Ourapproach to compute the operator H ( P ξ ) and update G ξ consists of two steps.
First, we approximately recover N maps T P ξ → P n via (cid:100) MM:R (cid:99) solver, i.e., we use N pairs or networks{ T θ n , v ω n } and train them by optimizing (7) with P ← P ξ and Q ← P n .
For each n = 1 , 2 , . . .
, N ,we perform SGAD by using batches from G ξ (cid:93) S and P n and get T θ n ≈ T P ξ → P n .
Second, we update G ξ to represent H ( P ξ ) instead of P ξ .
Inspired by [10], we do this via regression.
We introduce G ξ 0 , a ﬁxed copy of G ξ .
Next, we regress G ξ ( · ) onto (cid:80) Nn =
α n T θ n (cid:0) G ξ 0 ( · ) (cid:1) (cid:90) z (cid:96) (cid:18)
G ξ ( z ) ,
N (cid:88) n =
[Equation]
→ min ξ by performing SGD on random batches from S , e.g., by using squared error (cid:96) ( x, x (cid:48) ) def = 12 (cid:107) x − x (cid:48) (cid:107) 2 .
Thus, generator G ξ ( · ) becomes close to (cid:80) Nn =
α n T θ n (cid:0) G ξ 0 ( · ) (cid:1) as a function R H → R D .
We get
P ξ =
G ξ (cid:93) S ≈ (cid:2)
N (cid:88) n = α n T θ n (cid:3) (cid:93) (cid:2) G ξ 0 (cid:93) S (cid:3) = (cid:2) N (cid:88) n = α n T θ n (cid:3) (cid:93) P ξ 0 ≈ (cid:2)
N (cid:88) n = α n T P ξ 0 → P n (cid:3) (cid:93) P ξ
= H ( P ξ 0 ) , i.e., the new generated G ξ (cid:93) S measure approximates H ( P ξ 0 ) .
Summary.
Our two-step approach iteratively recomputes OT maps T P ξ → P n (Figure 2a) and then usesregression to update the generator (Figure 2b).
The optimization procedure is detailed in Algorithm 1(Appendix B.2).
Note that when ﬁtting OT maps T P ξ → P n , we start from previously used { T θ n , v ω n }rather than re-initialize them.
Empirically, this works better.
4.3 Relation to Variational Barycenter Algorithms
We show that our Algorithm 1 reduces to variational approach ( (cid:77) 3.2) when the number of generatorupdates, K G , is equal to 1 .
More speciﬁcally, we show the equivalence of the gradient update w.r.t.parameters ξ of the generator in our iterative Algorithm 1 and that of (8).
We assume that W 22 termsare computed exactly in (8) regardless of the particular OT solver.
Similarly, in Algorithm 1, weassume that maps G ξ 0 (cid:93) S → P n before the generator update are always exact, i.e., T θ n = T P ξ 0 → P n .
Lemma 1.
Assume that P ξ =
G ξ (cid:93) S ∈ P 2 ,ac ( R D ) .
Consider K G = 1 for the iterative Algorithm 1,i.e.
, we do a single gradient step regression update per OT solvers’ update.
Assume that (cid:96) ( x, x (cid:48) )
=12 (cid:107) x − x (cid:48) (cid:107) 2 , i.e., the squared loss is used for regression.
Then the generator’s gradient update in
Algorithm 1 is the same as in the variational algorithm:
[Equation]
where the derivatives are evaluated at ξ = ξ 0 .
We prove the lemma in Appendix A.
In practice, we choose K G = 50 as it empirically works better.
Ave, celeba! Images Dataset
In this section, we develop a generic methodology for building measures with known W 2 barycenter.
We then use it to construct Ave, celeba!
dataset for quantitative evaluation of barycenter algorithms.
Key idea.
Consider α 1 , . .
. , α N > 0 with (cid:80) Nn =1 α n = 1 , congruent convex functions ψ 1 , . . .
ψ N :
R D → R , and a measure P ∈ P
+2 , ac ( R D ) with positive density.
Deﬁne P n = ∇ ψ n (cid:93) P .
Thanks to
Brenier’s theorem [9], ∇ ψ n is the unique OT map from P to P n .
Since the support of P is R D , ψ n isthe unique (up to a constant) dual potential for ( P , P n ) [56].
Since potentials ψ n ’s are congruent, thebarycenter P of P n w.r.t.
weights α 1 , . .
, α N is P itself [11, C.2].
If ψ n ’s are such that all P n areabsolutely continuous, then P
=
P is the unique barycenter ( (cid:77) 2).
If one obtains N congruent ψ n , then for any P ∈ P 2 ,ac ( R D ) , pushforward measures P n = ∇ ψ n (cid:93)
Pcan be used as the input measures for the barycenter task.
For P accessible by samples, measures P nare also accessible by samples: one may sample x ∼ P and push samples forward by ∇ ψ n .
The challenging part is to construct non-trivial congruent convex functions ψ n .
First, we providea novel method to transform a single convex function ψ into a pair ( ψ l , ψ r ) of convex functionssatisfying α ∇ ψ l ( x ) + (1 − α ) ∇ ψ r ( x ) =
x for all x ∈ R D (Lemma 2).
Next, we extend the methodto generate congruent N - tuples (Lemma 3).
Lemma 2 (Constructing congruent pairs) .
Let ψ be a strongly convex and L -smooth (for some L > 0 )function.
Let β ∈
Deﬁne β -left and β -right functions of ψ by
[Equation]
and
[Equation]
Then βψ l ( x ) + (1 − β ) ψ r ( x ) = (cid:107) x (cid:107) 2 2 for x ∈ R D , i.e., convex functions ψ l , ψ r are congruent w.r.t.weights ( β, 1 − β ) .
Besides, for all x ∈ R D the gradient y l def = ∇ ψ l ( x ) can be computed via solvingβ -strongly concave optimization:
[Equation]
In turn, the value y r def = ∇ ψ r ( x ) is given by y r = ∇ ψ
( y l ) .
The proof is given in Appendix A.
We visualize the idea of our Lemma 2 in Figure 7a.
Thanks to
Lemma 2, any analytically known convex ψ , e.g., an ICNN, can be used to produce a congruent pairψ l , ψ r .
To compute the gradient maps, optimization (11) can be solved by convex optimization toolswith ∇ ψ computed by automatic differentiation.
Lemma 3 (Constructing N congruent functions.) .
Let ψ 01 , . . .
, ψ 0 M be convex functions,β 1 , . .
β M ∈ (0 , 1) and ψ lm , ψ rm be β m -left, β m -right functions for ψ 0 m respectively.
Letγ l , γ r ∈ R N
×
M be two rectangular matrices with non-negative elements and the sum of elements ineach column equals to 1 .
Let w 1 , . .
. , w M > 0 satisfy (cid:80) Mm =1 w m = 1 .
For n = 1 , . . .
, N deﬁne
[Equation]
Then ψ 1 , . .
. , ψ N are congruent w.r.t.
weights α n def = (cid:80)
Mm =1 w m (cid:2) β m γ lnm + (1 − β m
[Equation]
γ rnm (cid:3) .
We prove Lemma 3 in Appendix A.
We visualize the idea of our Lemma 3 in Figure 7b.
The lemmaprovides an elegant way to create N ≥ 2 congruent functions from convex linear combinations of functions in given congruent pairs ( ψ lm , ψ rm ) .
Gradients ∇ ψ n of these functions are respective linearcombinations of gradients ∇ ψ lm and ∇ ψ rm .
Dataset creation.
We use CelebA 64 × 64 faces dataset [31] as the basis for our Ave,celeba!
dataset.
We assume that CelebA dataset is an empirical sample from the continuous measure P Celeba ∈ P +2 , ac ( R 3 × 64 × 64 ) which we put to be the barycenter in our design, i.e.,
P = P Celeba .
We construct difﬁrentiable congruent ψ n with bijective gradients that produce 222 P n = ∇ ψ n (cid:93) P ∈ P +2 , ac ( R 3 × 64 × 64 ) whose unique barycenter is P Celeba .
In Lemma 3, we set N = 3 ,
M = 2 , β 1 = β 2 = 12 , w 1 = w 2 = 12 and
[Equation]
which yields weights ( α 1 , α 2 , α 3 )
= ( 14 , 12 , 14 )
We choose the constants above manually tomake sure the ﬁnal produced measures P n are visually distinguishable.
We use ψ 0 m ( x ) =
ICNN m (cid:0) s m ( σ m ( d m ( x )))
(cid:1) + λ (cid:107) x (cid:107) 2 2 as convex functions, where ICNNs have ConvICNN64 architecture [24, Appendix B.1], σ 1 , σ 2 are random permutations of pixels and channels, s 1 , s 2 are axis-wiserandom reﬂections, λ = 1100 .
In both functions, d m is a de-colorization transform which sets R, G, Bchannels of each pixel to ( 710 R + 125 G + 1350 B ) for ψ 01 and 13 ( R + G + B ) for ψ 02 .
The weights of ICNNsare initialized by the pre-trained potentials of W 22 ”Early” transport benchmark which map blurryfaces to the clean ones [24, (cid:77) 4.1].
All the implementation details are given in Appendix B.1.
Finally, to create Ave, celeba!
dataset, we randomly split the images dataset into 3 equal partscontaining ≈ 67 K samples, and map each part to respective measure P n = ∇ ψ n (cid:93) P Celeba by ∇ ψ n .
Resulting 3 × 67 K samples form the dataset consisting of 3 parts.
We show the samples in Figure 3.
The samples from the respective parts are in green boxes.
Evaluation
The code is written on the PyTorch and includes the script for producing Ave, celeba!
dataset.
Theexperiments are conducted on 4 × GPU GTX 1080ti.
The details are given in Appendix B.
6.1 Evaluation on Ave, celeba! Dataset
We evaluate our iterative algorithm 1 and a recent state-of-the-art variational [ SC W 2 B ] by [15] on
Ave, celeba!
dataset.
Both algorithms use a generative model P ξ =
G ξ (cid:93) S for the barycenter and yieldapproximate maps (cid:98) T P ξ → P n to input measures.
In our case, the maps are neural networks T θ n , while in[ SC W 2 B ] they are gradients of ICNNs.
The barycenters of Ave, celeba! ﬁtted by our algorithm and[ SC W 2 B ] are shown in Figures 4a and 4b respectively.
Recall the ground truth barycenter is P Celeba .
Thus, for quantitative evaluation we use FID score [19] computed on 200K generated samples w.r.t.
[Figure]
[Figure] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
the original CelebA dataset, see Table 1.
Our method drastically outperforms [ SC W 2 B ] .
Presumably,this is due to the latter using ICNNs which do not provide sufﬁcient performance.
Additionally, we evaluate to which extent the algorithms allow to recover the inverse OT maps
T P n → P ξ from inputs P n to the barycenter P ξ ≈ P .
In [ SC W 2 B ] , these maps are computed duringtraining.
Our algorithm does not compute them.
Thus, we separately ﬁt the inverse maps after maintraining by using (cid:98) MM:R (cid:101) solver between each input P n and learned P ξ (Algorithm 2 of Appendix2).
The inverse maps are given in Figure 5; their FID scores – in Table 2.
Here we add an additionalconstant shift (cid:100) CS (cid:99) baseline which simply shifts the mean of input P n to the mean µ of P .
The vectorµ is given by (cid:80) Nn
=1 α n µ n , where µ n is the mean of P n [3].
We estimate µ from samples y ∼ P n .
6.2 Additional Experimental Results
Different domains.
To stress-test our algorithm 1, we compute the barycenters w.r.t.
( α 1 , α 2 , α 3 ) =( 13 , 13 , 13 ) of notably different datasets: 50K Shoes [62], 138K Amazon Handbags and 90K Fruits[40].
All the images are rescaled to 64 × 64 .
The ground gruth barycenter is unknown , but one mayimagine what it looks like.
Due to (6), each barycenter image is a pixel-wise average of a shoe, a
[Figure]
handbag and a fruit, which is supported by our result shown in Figure 6.
In the same ﬁgure, we alsoshow the maps between datasets through the barycenter (Figures 6b, 6c, 6d): this allows generationof images in other categories with styles similar to a given image.
For instance, in Figure 6d, wegenerate an orange bag and an orange shoe by pushing the image of an orange through the barycenter.
We provide more examples in Figure 13 of Appendix C.4.
Extra results.
In Appendix C.1, we provide quantitative results for computing barycenters in the Gaussian case.
In Appendix C.2, we test how our algorithm works as a generative model on theoriginal CelebA dataset, i.e., when N = 1 and P 1 = P Celeba .
We show that in this case it achieves FID scores comparable to recent WGAN models.
In Appendix C.3, similar to [15], we computebarycenters of digit classes 0/1 of 32 × 32 grayscale MNIST [28].
We also test our algorithm on FashionMNIST [60] 10 classes dataset.
Discussion
Potential impact (algorithm).
We present a scalable barycenter algorithm based on ﬁxed-pointiterations with many application prospects.
For instance, in medical imaging, MRI is often acquiredat multiple sites where the overlap of information (imaging, genetic, diagnosis) between any twosites is limited.
Consequently, the data on each site may be biased and can cause generalizability androbustness issues when training models.
The developed algorithm could help aggregate data frommultiple sites and overcome the distributional shift issue across sites.
Potential impact (dataset).
There is no high-dimensional dataset for the barycenter problem exceptfor location-scattered cases (e.g. Gaussians), where the transport maps are always linear.
Hence ourproposed dataset ﬁlls an important gap, thereby allowing quantitative evaluation of future relatedmethods.
We expect our Ave, celeba!
to become a standard dataset for evaluating continuousbarycenter algorithms.
In addition, we describe a generic recipe ( (cid:77) 5) to produce new datasets.
Limitations (algorithm).
In our algorithm, the evolving measure P ξ is not guaranteed to be continuous, while it is continuous in the underlying ﬁxed point approach.
To enforce the absolute continuityof P ξ =
G ξ (cid:93) S , one may use an invertible network [14] for G ξ and an absolutely continuous latentmeasure S in a latent space of dimension H
= D .
However, our results suggest this is unnecessary inpractice — common GANs approaches also assume H (cid:28) D .
During the ﬁxed-point iterations, thebarycenter objective (5) decreases.
However, there is no guarantee that the sequence of measuresconverges to a ﬁxed point which is the barycenter.
Identifying the precise conditions on the inputmeasures and the initial point is an important future direction.
Besides, our algorithm does not recoverinverse OT maps T P n → P ; we compute them with an OT solver as a follow-up.
To avoid this secondstep, one may consider using invertible neural nets [14] to parametrize maps T θ n in our Algorithm 1.
Limitations (dataset).
To create Ave, celeba!
dataset ( (cid:77) 5), we compose ICNNs with decolorization,random reﬂections and permutations to simulate degraded images.
It is unclear how to produce otherpractically interesting effects via ICNNs.
It also remains an open question on how to better hide theinformation of the barycenter image in the constructed marginal measures.
Studying these questionsis an interesting future direction that can inspire benchmarking other OT problems.