Wasserstein Iterative Networks for Barycenter Estimation
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.
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 [61]).
Methods for computing barycenters have been successfully applied to various practical problems.
In geometry processing, shape interpolation can be performed by barycenters [55].
In image processing, barycenters are used for color and style translation [46, 41], texture mixing [47] and image interpolation [29, 53].
In language processing, barycenters can be applied to text evaluation [13].
In online learning, barycenters are used for aggregating probabilistic forecasts of experts [28, 44].
In Bayesian inference, the barycenter of subset posteriors converges to the full data posterior [56, 57] allowing efﬁcient computation of full posterior based on barycenters.
In reinforcement learning, barycenters are used for uncertainty propagation [38].
Other applications are data augmentation
[6], multivariate density registration [7], distributions alignment [22], domain generalization [36] and adaptation [40], model ensembling [15].
The bottleneck of obtaining barycenters is the computational complexity.
For discrete measures, fast and accurate barycenter algorithms exist for low-dimensional problems; see [45] for a survey.
However, 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 [31, 17, 27] are mostly based on entropic/quadratic regularization or parametrization of Brenier potentials with input-convex neural networks (ICNNs, see [4]).
The regularization-based [ CR W B ] algorithm by [31] recovers a barycenter biased from the true one.
Algorithms [ C W 2 B ] by [27] and [ SC W 2 B ] by [17] based on ICNNs resolve this issue, see [27, Tables 1-3].
However, despite the growing popularity of ICNNs in OT applications [37, 25, 39], they could be suboptimal architectures according to a recent study [26].
According to 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-scatter (e.g. Gaussians)
[3, (cid:77) 4] or 1 -dimensional [8, (cid:77) 2.3].
Recent works [31, 27, 17] 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 ( (cid:77) 4) for estimating Wasserstein-2 barycenters based on the ﬁxed point approach by [3] combined with a neural solver for optimal transport [26].
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 RGB images 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 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 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 R D unless stated otherwise.
We denote the set of all Borel probability measures on R D with ﬁnite second moment by P 2 ( R D ) .
We use P 2 , ac ( R D ) ⊂
P 2 ( R D ) to denote the subset of absolutely continuous 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 Q by Π( P , Q ) .
For a measurable map T : R D
→ R D , we denote the associated push-forward operator by T(cid:93) .
For φ : R D → R , we denote by φ its Legendre-Fenchel transform [18] 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]
where the minimum is taken over measurable functions (transport maps) T : R D → R D mapping P to Q .
The optimal T ∗ is called the optimal transport map .
Note that (1) is not symmetric, and this formulation 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, [23] proposed the following relaxation:
[Equation]
where the minimum is taken over all transport plans π , i.e., measures on R D × R D whose marginals are 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 [60] 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 for all x, y ∈ R D .
The functions u and v are called potentials .
There exist optimal u ∗ , v ∗ satisfying u ∗ = ( v ∗ ) c , where f c ( y ) def =
min x
[Equation]
(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 [61, Cases 5.3 & 5.17] to deﬁne u ( 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 be recovered from the dual solution u ∗ (or ψ ∗ ) of (3): T ∗ ( x )
= x − ∇ u ∗ ( x ) = ∇ ψ ∗ ( x ) [51, Theorem 1.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 bounded density.
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ﬁnition 3.6 & Theorem 5.1].
For n ∈ { 1 , 2 , . . . , N } , let T P → P n = ∇ ψ ∗ n be the OT maps from P to P n .
The following holds P -almost everywhere:
[Equation]
see [3, (cid:77) 3].
If P ∈ P +2 , ac ( R D ) , then (6) holds for every x ∈ R 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
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., [63, 34], parameterize T ∗ using complicated generative modeling techniques with adversarial losses to handle the pushforward constraint T(cid:93)
P = Q in the primal form (1).
They depend on careful hyperparameter search and complex optimization [35].
Dual-form continuous solvers [19, 52, 43, 59, 25] based on (3) or (4) have straightforward optimization 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 [26].
According to the evaluation, the best performing OT solver is reversed maximin solver (cid:98) MM:R (cid:101) , a modiﬁcation of the idea proposed by [43] in the context of Wasserstein-1 GANs [5].
In this paper, we employ this solver as 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 [48, 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 [26, (cid:77) 2] or [50, (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 , one may employ a generator G ξ : R H → R D with a latent measure S on R H and train ξ by minimizing
N (cid:88)
[Equation] [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 W2 terms.
To our knowledge, only ICNN-based minimax solver [37] has been used to compute W 2in (8) yielding [ SC W 2 B ] algorithm [17].
Alternatively,
Potential-based optimization.
[31, 27] 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 is recovered 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 [31] use entropic or quadratic regularization and [ C W 2 B ] algorithm by [27] uses ICNNs, congruence and cycle-consistency [25] regularization.
Other methods.
Recent work [12] combines the variational (8) and potential-based optimization via the c -cyclical monotonity regularization.
In [14], an algorithm to sample from the continuous
Wasserstein barycenter via the gradient ﬂows is proposed.
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 our algorithm 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 indeed absolutely 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 3.1].
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 in P
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 also makes parameterization (c) simple since both measures P n and H k ( P ) can be parametrized by means and covariance matrices, and the maps T H k
( P ) → P n are linear with closed form.
[Figure]
For general continuous measures P n , it remains an open problem to ﬁnd sharp conditions on inputs P n and the initial measure P of the ﬁxed-point iteration for the sequence {H k ( P ) } k to converge to the barycenter.
In this work, we empirically verify that the ﬁxed point approach works well for the input measures that we consider and for a randomly initialized generative model representing the evolving barycenter ( (cid:77) 4.2).
We tackle challenge (c) and develop a scalable optimization procedure that requires only sample access to P n ∈P 2 ,ac ( R D ) .
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 is a latent measure, e.g., S = N (0 , I H ) , and G ξ is a neural network R H
→ R D with parameters ξ .
Our approach 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 = α n T θ n (cid:0) G ξ 0 ( z ) (cid:1)(cid:19)
d S ( z ) → 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 uses regression to update the generator (Figure 2b).
The optimization procedure is detailed in Algorithm 1.
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 generator updates, K G , is equal to 1 .
More speciﬁcally, we show the equivalence of the gradient update w.r.t.
Input : latent S and input P
P measures; weights
>
; number of per network
K
; generator G
R D ; mapping networks T θ 1 , . .
, T θ N : R D
→ R D ; potentials v ω 1 , . .
. , v ω N : R D → R ; regression loss (cid:96) :
R D × R D
→ R + ;
Output : generator satisfying G ξ (cid:93) S ≈ P ; OT maps satisfying T θ n (cid:93) ( G ξ (cid:93) S )
≈ P n ; repeat # OT solvers update for n = 1 , 2 , . . .
, N do for k v = 1 , 2 , . .
, K v do
Sample batches Z ∼ S , Y ∼ P n ; X ← G ξ ( Z ) ;
L v ← 1 | X | (cid:80) x ∈ X v ω n (cid:0) T θ n ( x ) (cid:1) − 1 | Y | (cid:80) y ∈ Y v ω n (cid:0) y (cid:1) ;
Update ω n by using ∂ L v ∂ω n ; for k T = 1 , 2 , . .
, K T do
Sample batch Z ∼ S ; X ← G ξ ( Z ) ;
L T = 1 | X | (cid:80) x ∈ X (cid:2) 12 (cid:107) x − T θ n ( x ) (cid:107) 2 − v ω n (cid:0) T θ n ( x ) (cid:1)(cid:3) ;
Update θ n by using ∂ L T ∂θ n ; # Generator update (regression)
G ξ 0 ← copy (cid:0) G ξ (cid:1) ; for k G = 1 , 2 , . . .
, K G do
Sample batch Z ∼ S ;
L G ← 1 | Z | (cid:80) z ∈ Z (cid:96) (cid:18)
[Equation]
[Equation]
; parameters ξ of the generator in our iterative Algorithm 1 and that of (8).
We assume that W 22 terms are computed exactly in (8) regardless of the particular OT solver.
Similarly, in Algorithm 1, we assume 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 is the unique (up to a constant) dual potential for ( P , P n ) [58].
Since potentials ψ n ’s are congruent, the barycenter 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 are absolutely continuous, then P
=
P is the unique barycenter ( (cid:77) 2).
[Figure]
and production of measures with known W 2 barycenter ( (cid:77) 5).
If one obtains N congruent ψ n , then for any P ∈ P 2 ,ac ( R D ) , pushforward measures P n = ∇ ψ n (cid:93) P can be used as the input measures for the barycenter task.
For P accessible by samples, measures P n are 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 provide a novel method to transform a single convex function ψ into a pair ( ψ l , ψ r ) of convex functions satisfying α ∇ ψ l ( x ) + (1 − α ) ∇ ψ r ( x ) =
x for all x ∈ R D (Lemma 2).
Next, we extend the method to 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 3a.
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 tools with ∇ ψ 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 in each 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] [Equation]
Then ψ 1 , . .
. , ψ N are congruent w.r.t.
weights α n def = (cid:80)
Mm =1 w m (cid:2) β m γ l nm + (1 − β m ) γ r nm (cid:3) .
We prove Lemma 3 in Appendix A.
We visualize the idea of our Lemma 3 in Figure 3b.
The lemma provides 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 linear combinations of gradients ∇ ψ lm and ∇ ψ rm .
Dataset creation.
We use CelebA 64 × 64 faces dataset [33] 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 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 to make 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 [26, Appendix B.1], σ 1 , σ 2 are random permutations of pixels and channels, s 1 , s 2 are axis-wise random reﬂections, λ = 1100 .
In both functions, d m is a de-colorization transform which sets R, G, B channels of each pixel to ( 710 R + 125 G + 1350 B ) for ψ 01 and 13 ( R + G + B ) for ψ 02 .
The weights of ICNNs are initialized by the pre-trained potentials of W 22 ”Early” transport benchmark which map blurry faces to the clean ones [26, (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 parts containing ≈ 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 4.
The samples from the respective parts are in green boxes.
Evaluation
The code 1 is written on the PyTorch and includes the script for producing Ave, celeba!
dataset.
The experiments 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 [17] on
Ave, celeba!
dataset.
Both algorithms use a generative model P ξ =
G ξ (cid:93) S for the barycenter and yield approximate 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 5a and 5b respectively.
Recall the ground truth barycenter is P Celeba .
Thus, for quantitative evaluation we use FID score [21] computed on 200K generated samples w.r.t.
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.
[Figure]
[Figure]
[Figure]
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 during training.
Our algorithm does not compute them.
Thus, we separately ﬁt the inverse maps after main training by using (cid:98) MM:R (cid:101) solver between each input P n and learned P ξ (Algorithm 2 of Appendix 2).
The inverse maps are given in Figure 6; their FID scores – in Table 2.
Here we add an additional constant 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 ) =
( 1 3 , 1 3 , 1 3 ) of notably different datasets: 50K Shoes [64], 138K Amazon Handbags and 90K Fruits [42].
All the images are rescaled to 64 × 64 .
The ground gruth barycenter is unknown , but one may imagine what it looks like.
Due to (6), each barycenter image is a pixel-wise average of a shoe, a handbag and a fruit, which is supported by our result shown in Figure 7.
In the same ﬁgure, we also show the maps between datasets through the barycenter (Figures 7b, 7c, 7d): this allows generation of images in other categories with styles similar to a given image.
For instance, in Figure 7d, we
[Figure]
We competitive generate an orange bag and an orange shoe by pushing the image of an orange through the barycenter.
We provide more examples in Figure 15 of Appendix C.5.
Extra results.
In Appendix C.1, we compute the barycenters of toy 2D dsitributions.
In Appendix
C.2, we provide quantitative results for computing barycenters in the Gaussian case.
In Appendix
C.3, we test how our algorithm works as a generative model on the original 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.4, similar to [17], we compute barycenters of digit classes 0/1 of 32 × 32 grayscale MNIST [30].
We also test our algorithm on FashionMNIST [62] 10 classes dataset.
Discussion
Potential impact (algorithm).
We present a scalable barycenter algorithm based on ﬁxed-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-dimensional dataset for the barycenter problem except for location-scattered cases (e.g. Gaussians), where the transport maps are always linear.
Hence our proposed dataset ﬁlls 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 ( (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 continuity of P ξ = G ξ (cid:93) S , one may use an invertible network [16] for G ξ and an absolutely continuous latent measure S in a latent space of dimension H
= D .
However, our results suggest this is unnecessary in practice — common GANs approaches also assume H (cid:28) D .
During the ﬁxed-point iterations, the barycenter objective (5) decreases.
However, there is no guarantee that the sequence of measures converges to a ﬁxed 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 T P n → P ; we compute them with an OT solver as a follow-up.
To avoid this second step, one may consider using invertible neural nets [16] 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 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.
A CKNOWLEDGEMENTS .
Skoltech 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 September 1, 2022. Neither Mr.