Learning to generate Wasserstein barycenters
Authors: Julien Lacombe
Julie Digne
Nicolas Courty
Nicolas Bonneel
Venue:  (2021)
Keywords: Wasserstein barycenters
Optimal Transport

Abstract: Optimal transport is a notoriously difficult problem to solve numerically, with current approaches often remaining intractable for very large scale applications such as those encountered in machine learning. Wasserstein barycenters -- the problem of finding measures in-between given input measures in the optimal transport sense -- is even more computationally demanding. 
By training a deep convolutional neural network, we improve by a factor of 60 the computational speed of Wasserstein barycenters over the fastest state-of-the-art approach on the GPU, resulting in milliseconds computational times on $512\times512$ regular grids.
We show that our network, trained on Wasserstein barycenters of pairs of measures, generalizes well to the problem of finding Wasserstein barycenters of more than two measures. We validate our approach on synthetic shapes generated via Constructive Solid Geometry as well as on the ``Quick, Draw'' sketches dataset.

Under review as a conference paper at ICLR 2021
LEARNING TO GENERATE WASSERSTEIN BARYCEN- TERS
Anonymous authors Paper under double-blind review
ABSTRACT
Optimal transport is a notoriously difficult problem to solve numerically, with cur- rent approaches often remaining intractable for very large scale applications such as those encountered in machine learning. Wasserstein barycenters the prob- lem of finding measures in in-berween given input measures in the optimal transport sense is even more computationally demanding. By training deep convolu- tional neural network, we improve by a factor of 60 the computational speed Wasserstein barycenters
the fastest satte-of-the-art approach on the GPU, re- sulting in milliseconds computational times on 512 x 512 regular grids. We show that our network, trained Wasserstein barycenters of pairs of measures, gen- eralizes well to the problem of finding Wasserstein barycenters of more than two measures. We validate our approach synthetic shapes generated via Construc- tive Solid Geometry as well as on the "Quick, Draw" sketches dataset.
INTRODUCTION
Optimal transport is becoming widespread in machine learning, but also in computer graphics, vision and many other disciplines. Its framework allows for comparing probability distributions, shapes or images, as well as producing interpolations of these data. As a result, it has been used in the context of machine learning as loss for training neural networks (Arjovsky al. 2017), a manifold for dictionary learning (Schmitz et al.) 2018). clustering (Mi et al. (2088) and metric learning ap- plications (Heitz et 2019), as way sample an embedding (Liutkus al.) 2019) and transfer learning (Courty et al. 2014). and many other applications (see Sec. (..3). However, despite recent progress in computational optimal transport, in many cases these applications have remained limited to small datasets due to the substantial computational cost of optimal transport, in terms of speed, but also memory.
We tackle the problem of efficiently computing Wasserstein barycenters of measures discretized on regular grids, a setting common to several of these machine learning applications. Wasserstein barycenters are interpolations of two more probability distributions under optimal transport dis- tances. As such, a common way to obtain them is to perform a minimization of a functional involving optimal transport distances transport plans, which is thus very costly process. Instead, we di- rectly predict Wasserstein barycenters by training a Deep Convolutional Neural Network (DCNN) specific this task.
An important challenge behind our work is to build an architecture that can handle a variable number of input measures with associated weights without needing to retrain specific network. To achieve that, we specify and adapt an architecture designed for and trained with two input measures, and show that we can use this modified network with no retraining to compute barycenters of more than two measures. Directly predicting Wasserstein barycenters avoids the need to compute a Wasserstein embedding (Courty al. 2017), and our experiments suggest that this results in better Wasserstein barycenters approximations. Our implementation is publicly available
Contributions This paper introduces a method to compute Wasserstein barycenters in millisec- onds. It shows that this can be done by learning Wasserstein barycenters of only two measures on a dataset of random shapes using a DCNN, and by adapting this DCNN to handle multiple input
httpa: httpg/gthhb.coo.. r2021 hutp:oouuuuuuttr//iaaaaa earn' ng- cooo-ggnerattt www.ttttssssssttsseeeeeeeee wwaaerratet--barrcaettrrr

---

Under review as a conference paper at ICLR 2021
measures without retraining. This proposed approach is 60x faster than the fastest state-of offte--ffiiee the-hrr-r GPU library, and performs better than Wasserstein embeddings.
RELATED WORK
2.1
WASSERSTEIN DISTANCES AND APPROXIMATIONS
Optimal transport seeks the best way to warp a given probability measure f40 to form another given probability measure μl by minimizing the total cost of moving individual "particles of earth", We restrict our description to discrete distributions. In this setting, finding the optimal transport between two probability measures is often achieved by solving large linear program (Kantorovich) 1942) more details on this theory and numerical tools can be found in the book Peyre et al. (2019). This minimization results in the so-called Wasserstein distance, the mathematical distance defined by the total cost of reshaping μo to M1. This distance can be used to compare probability distributions, in particular in a machine learning context. It also results in a transport plan, a matrix P(x, y) representing the amount of mass of do traveling from location x in μo towards location y in (41-
However, the Wasserstein distance is notoriously difficult to compute the corresponding linear program is huge, and dedicated solvers typically solve this problem in O(N log N). with N the size of the input measures discretization. Recently, numerous approaches have attempted to ap- proximate Wasserstein distances. One of the most efficient methods, the so-called Sinkhorn al- gorithm introduces an entropic regularization, allowing to compute such distances by iteratively performing fast matrix-vector multiplications (Cuturi) 2013) or convolutions in the case of regular grids (Solomon et al. 2013). However, this comes at the expense of smoothing the transport plan and removing guarantees regarding this mathematical distance (in particular, the regularized cost W. W. (o) * 0). These issues are addressed by Sinkhorn divergences (Feydy et al. 2018 Genevay et al. 2017). This approach symmetrizes the engroiy-regularrieed optimal transport distance, adding guarantees on this divergence (now, the cost (c /4) = O by construction, though triangular in- cquality still does not hold) but also effectively reducing blur, while maintaining a relatively fast numerical algorithm. They show that this divergence interpolates between optimal transport dis- tances and Maximum Mean Discrepancies. Sinkhorn divergences are implemented in the GeomLoss library (Feydy 2019), relying specific computational scheme the GPU (Feydy et al. 2019] 2018 Schmiizerr (019) and constitutes the state-of-the in term of speed and approximation of optimal transport- distances.
2.2 WASSERSTEIN BARYCENTERS
The Wasserstein barycenter of a set of probability measures corresponds to the Fréchet mean of these measures under the Wasserstein distance (i.e., weighted mean under the Wasserstein metric). Wasserstein barycenters allow to interpolate between two or more probability measures by warping these measures (contrarily to Euclidean barycenters that blends them). Similarly Wasserstein dis- tances. Wasserstein barycenters are very expensive to compute. An autropy-regulariy approach based on Sinkhorn-like iterations also allows to efficiently compute blurred Wasserstein barycen- ters. Reducing blur via Sinkhorn divergences is also doable, but does not benefit from very fast Sinkhorn- like algorithm: a weighted sum of Sinkhorn divergences needs to be iteratively minimized, which adds significant computational cost. In our approach, we rely on Sinkhorn divergence-basee barycenters to feed training data to a Deep Convolutional Neural Network, and aim at speeding up this approach. Other fast transport-brig based barycenters include that of sliced and Radon Wasserstein barycenters, obtained via Wasserstein barycenters on 1-d projections (Bonneel et al. 2013). which compare
recent trend seeks linearizations or Euclidean embeddings of optimal transport problems. Notably, Nader & Guennebaud (2018) approximate Wasserstein barycenters by first solving optimal trans- port map between a uniform measure towards n input measures, and then linearly combining Monge maps. This allows for efficient computations typically of the order of half a second for 512x512 12 images. A similar approach is taken within the documentation of the GeomLoss library (Feydy) where a single step of a gradient descent initialized with a uniform distribution used,
2019
See https://www kernel httpps////www..reeel-orrrrrrrr
i aawolosstooaee_eeeeeeeeoo hat_/__utt__aappppp_ pettappls//ppe oes/opiima
html
2

---

Under review as a conference paper at ICLR 2021
which effectively corresponds to such linearization. We use this technique in work to train our network. Waang et al. (2013). Moosmüller & Cloninger (2020) and Mérigot et al. (2020) use similar linearization, possibly using non-umiform reference measure, with theoretical guaran- tees on the distorsion introduced by the embedding. Instead of explicitly building an embedding via Monge maps, such embedding can be learned. [Courty al. (2017) propose a siamese neu- ral network architecture to learn an embedding in which the Euclidean distance approximates the Wasserstein distance. Wasserstein barycenters can then be approximated by interpolating within the Euclidean embedding, without requiring explicit computations of transport plans. They show accurate barycenters on a number of datasets of low resolution (28 x 28). However, in general, is unclear whether Wasserstein metrics embed into Euclidean spaces. Negative results were shown for 3d optimal transport onto Euclidean space (Andoni al., 2016). Interestingly, in the reversed direction, Wasserstein spaces have been used to embed other metrics (Frogner et al. 2019))
Wasserstein barycenters can also be seen as particular instance of inverse problem. There is an im- portant literature on the resolution of inverse problems with deep learning models on instances such
as (non-exhaustive list) image denoising (Clyanov et al. (018)) (Burger et al. (0122) (Lefkimmiatis)
2017). super-resolution (Ledig et al. 2017). (Tai et al. 2017). (Lai al. 2017). inpainting (Ych)
et al. 2017) (Xie et al. 2012). (Liu et al. (0188).
Parallel to work, Fan al. (2020) propose model based input convex neural networks
(ICNN) developed by Amos et al. (2017). Their method allows a fast approximation of Wasserstein
barycenters of continuous input measures. This last work is also closely related to the semi-discrete
approach of Claici et al. (2018).
2.3 APPLICATIONS TO MACHINE LEARNING
For its ability to compare probability measures, optimal transport has seen important success in ma- chine learning. This is particularly the case of Wasserstein GANs (Arjovsky et al., 2017) that com- pute very efficient approximation of Wasserstein distances as loss for generative adversarial mod- els. The optimal transport loss has also been used the context of dictionary learning (Rolet et al.) 2016). Other fast approximations have allowed to perform domain adaptation for transfer learning of a classifier, by advecting samples via a computed transport plan (Courty et al., 2014). Among these approximations, Sliced optimal transport has been used to sample an embedding learned by an auto-encoder, by computing a flow between uniformly random samples and the image of encoded inputs (Liutkus et al., 2019).
Regarding the Wasserstein barycenters we are interested in, they have been used for the task of learning a dictionary out of a set of probability measures (Schmitz al., 2018), for computing Wasserstein barycentric coordinates of probability measures (Bonneel et al. (016) or for metric learning (Heitz et al. 2019). These have been performed by cifoomatc--ifffeenttaate
of Wasser- stein barycenters obtained through Sinkhorn iterations and non-linear optimization, and have thus been limited to small datasets, both due to speed and memory limitations. An adaptation k-means clustering for optimal transport was proposed by Mi et al. (2018) and (Domazakis et al.) 2020). Backhoff- et al. replaces maximum a posteriori (MAP) estimation or Bayesian model av- erage, by computing Wasserstein barycenters of posterior distributions (Backhoff- Veraguas et al.) 2018) using a stochastic gradient descent scheme. In the context of reinforcement learning, Wasser- stein barycenters used by Metelli et al. (219)) robustness to uncertainty. PCA in the Wasserstein space require the ability to compute Wasserstein barycenters: they have been studied by Bigot al. (2017) but could only be computed in 1-d where theory is simpler. In the work of Dognin et al. 2019)), Wasserstein barycenters are used for model ensembling, i.c., averaging the predictions of several models build more robust model.
way to regularize the update rule and offer
In this work, we do not focus on a single application but instead provide the tools to efficiently approximate Wasserstein barycenters on 2-d regular grids.
3 LEARNING WASSERSTEIN BARYCENTERS
This section describes our neural network and our proposed solution train it in scalable way.
3

---

Under review as a conference paper at ICLR 2021
3.1
PROPOSED MODEL
Contractive path qi (xn)
Expansive path 4
Diinscompingg
Upsanprrr - Seftma
20 Input
512
Predicted
Depth
Level
(1) (2) (3)
(6)
(5)
(4)
(1)
toput
(1)
(6)
Details of the connections between and
enput 2
Baryoortaa
(6)
(6)
(5)
(2)
Iegut m
Copy Weaghood
Concetoration
(2)
(6)
Figure Our model divided into n contractive paths Pi, sharing the same architecture and th. Blue rectangles represent feature maps and arrows denote the different operations we use (see legend). training time n 2, but by duplicating the contractive paths, we can adapt to the measures barycenter problem at test time, without needing to retrain the network.
weights, and one expansive path
Our model aims at obtaining approximations of Wasserstein barycenters from > 2 probability measures {p:limm.... discretized on 512 x 512 regular grids, and their corresponding barycentric weights {A:)i=1..n- Based on the observation that the Sinkhorn algorithm is mainly made of suc- cessive convolutions, we propose to directly predict a Wasserstein barycenter through an end-to-end neural network approach, using a Deep Convolutional Neural Network (DCNN) architecture. This DCNN should be deep enough to allow accurate approximations but shallow enough to reduce its computational requirements.
We propose network consisting of n contractive paths {pp))= and one expansive path v (see Fig. [T.. Importantly enough, n is not fixed and can vary test time. In fact, the same path with the same architecture and sharing the same weights. The contractive paths are made of successive blocks, each block consisting of two convolutional layers followed by ReLU activation. We further add average pooling layers between each block in order to decrease the dimensionality. The expansive path is symetrically constructed, each block also being made of 2 convolutional layers with ReLU activations. To better invert average poolings, layers with nearest-neighbor interpolation. Finally, to recover an output probability distribution, we softmax activation the end of the expansive path. All the 2D convolutions of model use 3 x kernels with a stride and a padding equal to 1. The architecture might look similar to the U-Net architecture introduced by RRonneberger al. (0015), because of the nature of the contractive and expansive paths. However the similarities end here, since our architecture uses a variable number of contractive paths to handle multiple inputs. The connections we use from the contractive paths to the expansive path also highly differ: first, we take all the feature maps from each contractive path and not only part of as is done U-Net, and, second, we compute weighted sum of all these activations using barycentric weights which results in a weighted feature map which is then symmetrically concatenate to the corresponding activations in the expansive path. Our network is deeper than U-Net and we do not use the same succession of layers nor the same downsampling and upsampling methods which are respectively max- maxx-ooolig and up-convolutions in the case of U- Net. We also use Instance Normalization (Ulyanov et al. 2016) which has empirically shown better results than Batch Normalization for our model. These normalization layers are placed before each ReLU activation.
duplicates of
upsampling
4

---

Under review as a conference paper at ICLR 2021
The connections going from the contractive paths to the expansive path defined as follows: after each block in a contractive path (i at depth level j, we take the resulting activations h=1..s compute their linear combination Fj Lien A,F.y, and concatenate it symmetrically to the corre- sponding activations in the expansive path (see figure
3.2 TRAINING
Our solution allows to generalize network trained for computing the barycenter of two measures to an arbitrary number input measures while remaining fast to train.
Variable number of inputs. We expect network to produce accurate results without con- structing an explicit embedding whose existence remains uncertain (Andoni et al. 2008). However, Euclidean embedding trivially generalizes an arbitrary number of input measures. A key insight to our work is that, since contractive paths weights are shared, our network can be trained using only two contractive paths for the task of predicting Wasserstein barycenters of two probability mea- sures. Once trained, contractive paths can be duplicated to the desired number of input measures. In practice, we found this procedure yield accurate barycenters (see Sec.
Loss function. Training the network requires comparing the predicted Wasserstein barycenter to a groundtruth Wasserstein barycenter. Ideally, such comparison should be performed via an optimal transport cost - those are ideal to compare probability distributions. However, computing optimal transport costs on large training datasets would be intractable. Instead. we resort Kullback- Leibler divergence between the output distribution and the desired barycenter.
Optimizer. To optimize the model parameters, we use a stochastic gradient descent with warm restarts (SGDR) (Loshchilov & Hutter 2016). The exact learning rate schedule we used for our models shown in appendix A Fig. 10
Training data. We strive to train our network with datasets that would cover a wide range of input sketches. To achieve this, we built dataset made 100k pairs of 512 x 512 random shape contours with random barycentric weights and their corresponding 2D Wasserstein barycenter. These 2D shapes are generated Constructive Solid Geometry fashion: we randomly assemble primitives shapes using logical operators and detect contours in post-processing. primitive corresponds to a filled ellipse, triangle, rectangle or a line. We assemble these primitives together by using the classical boolean operators OR, AND, XOR, NOT. To generate shape, we initialize it with random primitive. Then we combine it with another random primitive using a randomly chosen operator, and repeat this operation d times (0 < d < 50) where follows the probability distribution d ~ 1 (0, 50) + (N 2.5) + (50, 2.5)) which promotes simple (d close to 0) and complex (d close to 50) shapes. Finally, we apply a Sobel filter to create contours. We thus create 10k random 2D shapes from which we build 100k Wasserstein barycenters.
3
We then use the GoomLoss library (Feydy) 2019) to build good approximations of Wasserstein barycenters in reasonable time, with random pairs of inputs sampled from the set of gen- crated shape contours. Given two 2D input distributions li and μl with their corresponding barycentric weights A and X2 1 Å₁. their barycenter b" can be found by minimizing: 6" = arg min, (b, p1) + (2, p2) where S. corresponds to the Sinkhorn divergence with quadratic ground metric, and the regularization parameter (we use e le-4). We use a La- grangian gradient descent scheme that first samples the distributions as b Σÿ byōs, and then perform descent using 2 (k+1) x + An + >2n where bith is the displacement vector. This vector is computed as the gradient of the Sinkhorn divergence: 1/5h = - V., 2, (b, 14). These successive updates can be computationally expensive when inputs are large. To speed up compu- tations, we use a linearized approach that performs a single descent step, starting from a uniform distribution. In practice, this allows to precompute distribution and each of the input measures in the database, and obtain approximate Wasserstein barycenters by using weighted average of these transport maps.
optimal transport map between uniform
5

---

Under review as a conference paper at ICLR 2021
4 EXPERIMENTAL RESULTS
While our model is exclusively trained on our synthetic dataset, at test time we also consider the Quick, Draw! dataset from Google (2020) and the Coil20 dataset (Nane et al. 1996). The first dataset contains million grayscale drawings divided multiple classes and has been created by asking users to draw with mouse given object in limited time. The second one is made of images of 20 objects rotating on a black background and contains 72 images per object for a total of 1440 images. We rasterized these two datasets to 512 X 512 images.
4.1 TWO-WAY INTERPOLATION RESULTS
In Fig. 2 we show a visual comparison between barycenters obtained with Geomloss and our method. Wassertein barycenters taken from the test dataset and the corresponding predictions shown. We also compare these results to classical approaches (linear program, regularized barycen- ters) and to another approximation method known as Radon barycenters (Bonneel et al. 2015
inpot
A
0.4586
0.5618
A 0.4137
0.5414
0.3427
0.7433
Geamloss
Our Model
Linor Progrem
G
Regnlbrried
4R
Radon
Figure 2: We illustrate typical results and comparisons to GeomLoss Feydy 2019 a linear pro-
gram via network simplex (Bonneel al. 2011), regularized barycenters computed log domain (see for instance |Peyré et al. (2019)) with regularization parameter of 1e-3 and Radon barycen- ters (Bonneel al. 2015).
To further visually assess that the barycenters approximating are close the exact ones, also present comparison with the method of Claici et al. (2018) in Fig. 3 Input distributions are taken from the Quick, Draw! dataset.

---

Under review as a conference paper at ICLR 2021
Inpuls
GeomLoss
Our Model
Inputs
Geomiouss
Our Model
Figure 3: We superimpose the centroids (A1 = A2 = 0.5) found by the method of Claici et al. (0188) (in red) using 100 Dirac masses over the ones computed by GeomLoss and by our model, on images from the Quick, Draw! dataset. The solution of Claici et al. (2018) was found within 37 hours of computation.
We compare method with the Deep Wasserstein Embedding (DWE) model developed by Courty jet al. (2017) on Quick, Draw! images. We propose two versions of DWE. The first version relies on the exact original architecture which only process 28 28 images, retrained version of our shape contours dataset see Fig. 5 for this comparison. In the second version, we adapt their network to process 512x 512 imputs. The encoder and decoder of this second version have the same architecture as the contractive and expansive paths that we use in our model without our skip connections, but is used to compute the embedding rather than directly predicting baryeenters see Fig.
a downsampled
(a)
(b)
13.--
125
11.-
Narbe
08
23
30.2 350
Figure 4: Approximation error of our model compared to the ones of DWE (version adapted to handle 512 x 512 images), respectively measured in terms of (a) KL-Divergence and (b) LI distance, on images coming from our synthetic test dataset. Fach one of the 1000 x 4 points corresponds to barycenter. The x-axis represents the error measured between the GeomLoss barycenter and the barycenter predicted by our model while the y-axis represents the one between the GeomLoss barycenter and the barycenter predicted by DWE. The color of a point associated to a barycenter represents its number of inputs.
In Fig 4 we show a numerical comparison of approximation errors between our model and DWE adapted to 512 512 images of shape contours dataset, in terms of KL-divergence and LI distance. Our results clearly show that our method is able to approximate more accurately the Wasserstein barycenter 512 512 input measures.
Finally, we study the limitations of the generalization of our network on the Coil20 dataset (Nanc) al. 1996), which consists of images of objects on a black background. In figure 7 we show the interpolation of 2 cars; ; additional results are available in appendix B figure
4.2 N- WAY BARYCENTERS
Even if our model has been trained using only barycenters computed from pairs of inputs, we can apply it to predict barycenters of more than two measures. We display interpolations between spectively three and five input measures in Fig which surprisingly tends to show that our model generalize what it learned on pairs of inputs, at least partially. Additional results Quick, Draw! are also shown in appendix B Fig. [2 A 100-way barycenter comparison can be found in appendix Sec. Fig. 13

---

Under review as a conference paper at ICLR 2021
Geoml Loss
Our Model
DIV
Figure 5: Interpolations between two 28 x 28 images from the Quick, Draw! dataset using Geomloss, model and the original Deep Wasserstein Embedding (DWE) method from (Courty al. 2017 Our model directly considers 512 x 512 inputs and its results are downsampled from 512 x 512 to
28 28.
Our Model GoomLoc oes
DWE
Figure 6: Interpolations between two 512 x 512 images from the Quick, Draw! dataset using Geomloss, our model and the Deep Wasserstein Embedding (DWE) method from (Courty et al., 2017) adapted to handle 512 x 512 images.
Numerically when the number of inputs is greater than 2, our model also achieve to find better approximations than the ones obtained with DWE, as shown in Fig. 4
4.3 SPEED
In order to assess computational times, obtain average running time over 1000 barycenter compu- tations on average, our model predicts barycenters of two images 0.0092 seconds. We compare the average speed of our model with GeomLoss in different settings. The first one considers the full 512 x 512 images GeomL Loss computes such barycenters setting takes advantage the sparsity of our images and only uses the 2D coordinates of the points with non-zero mass in this case, GeomLoss computes barycenters in 0.589 seconds. Our method provides nearly 64x speedup compared with this last approach. comparison, an exact barycenter computation of two (sparse) measures using network simplex (Bonneel et al. 2011) ranges from 4 80 seconds for typical shape contours images that contains few thousands pixels carrying mass. The time required to compute barycenters using the method of (Claici et al. 2018) 1) depends on the number iterations, our setting 100 iterations with the inputs shown figure 3 require 37 hours while 50 iterations are achieved in 14 hours. A 512 x 512 Radon barycenter (Bonneel et al.) 2015) requires 0.2 seconds for 720 projection directions, but remains far from the expected barycenter.
1.41 seconds. The second
DISCUSSION AND CONCLUSION
While our method produces good approximation of Wasserstein barycenters n inputs, some shapes are surprisingly difficult to handle. The barycenter of simple translated and scaled shapes such as lines or ellipses should theoretically also be lines or ellipses, but are failure cases for our model (Fig. 9 while more complex shapes well handled (Fig. 8). In addition, we rely on lin- earized barycenter to train our network (Nader & Guennebaud 2018 Wang et al. 2013)] Moosmüller & Cloninger 2020; Mérigot al. 2020). which incurs some error. This can be seen in appendix Sec. C. Fig. 14 While using more iterations of gradient descent yields more accurate results and re- moves this linearity, it also prevents easy combination and makes the dataset generation intractable.
8

---

Under review as a conference paper at ICLR 2021
GeomL Loss
Our Model
Figure 7: Interpolations between two 512 x 512 images from the Coil20 dataset (Nane et al. 1996 using GeomLoss and model trained with synthetic shape contours. In order to perform compu- tations on the shapes and not on the background, mass has been inverted.
Mode GoomLocr3)
CommLos
Our Model
Figure 8: Wasserstein barycenters of three inputs (top rows) and five inputs (bottom rows) from
Quick, Draw!, respectively computed with Geomloss and with model trained with only pairs from our synthetic training dataset. Barycentric weights are randomly chosen.
Grent oos
Ours
2-ciroles
2- lines
S-cincles
5-lines
Figure 9: Wasserstein barycenters of sets of lines or ellipses should result in lines (resp. ellipses). Our prediction for two-way barycenters (here, with equal weights) of such shapes remains correct (left). However, the predicted barycenter is highly distorted for 5-way barycenters of simple shapes (right) although it remains plausible for more complex shapes (see Fig. 88.
Nevertheless, in many cases our DCNN able to synthesize a barycenter from an arbitrary number of inputs. The main strength of our approach lies in its capacity be trained from only 2-inputs barycenters examples and to generalize to any number of inputs. We showed that the results ex- ceeded the ones obtained by explicit Wasserstein Embedding computation while having very low computation time. We hope our fast approach will accelerate the adoption of optimal transport in machine learning applications.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International Confer- ence on Machine Learning, pp. 146-155, 2017.
Alexandr Andoni, Piotr Indyk, and Robert Krauthgamer. Earth mover distance over high- dimensional spaces. In SODA, volume 8, pp. 342, 2008.
Alexandr Andoni, Assaf Naor, and Ofer Neiman. Impossibility of sketching of the 3d transportation metric with quadratic cost. 43rd International Colloquium on Automata, Languages, and Programming (ICALP 2016). Schloss Dagstuhl-Lcibniz-ZZennni fuer Informatik, 2016.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan, 2017.
Julio Backhoff- Backhoff-Veraguas, Joaquin Fontbona, Gonzalo Rios, and Felipe Tobar. Bayesian learning with wasserstein barycenters. arXiv preprint arXiv:1805.1083: 2018.
Jérémie Bigot, Raúl Gouet, Thierry Klein, Alfredo López, et al. Geodesic pea in the wasserstein space by convex pea. In Annales de I' l'Institut Henri Poincaré, Probabilités Statistiques, vol- ume 53, pp. 1-26. Institut Henri Poincaré, 2017.
Nicolas Bonneel, Michiel van de Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement Inter- polation Using Lagrangian Mass Transport. ACM Transactions Graphics (SIGGRAPII ASIA 201/), 30(6), 2011.
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51(1):22-45, 2015.
Nicolas Bonneel, Gabriel Peyré, and Marco Cuturi. Wasserstein Barycentric Coordinates: His- togram Regression Using Optimal Transport. ACM Transactions on Graphics (SIGGRAPH 2016), 35(4), 2016.
Harold Burger, Christian Schuler, and Stefan Harmeling. Image denoising: Can plain neural networks compete with bm3d? In 2012 IEEE conference on computer vision and pattern recog- nition, pp. 2392-2399. IEEE, 2012.
Sebastian Claici, Edward Chien, and Justin Solomon. Stochastic wasserstein barycenters. arXiv preprint (tX.Xv:8802.057577 2018.
Nicolas Courty, Rémi Flamary, and Devis Tuia. Domain adaptation with regularized optimal transport. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 274-289 Springer, 2014.
Nicolas Courty, Rémi Flamary, and Mélanie Ducoffe. Learning wasserstein embeddings. arXiv preprint arXiv:1710.07457, 2017.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013.
Pierre Dognin, Igor Melnyk, Youssef Mroueh, Jerret Ross, Cicero Dos Santos, and Tom Sercu. Wasserstein barycenter model ensembling. arXiv preprint arXiv:1902.0440 2019.
G Domazakis, D Drivaliaris, S Koukoulas, G Papayiannis, A Tsekrekos, and A Yannacopoulos. Clustering measure-valuse valuud data with wasserstein barycenters. arXiv preprint 1rXiv:1912.11801, 2020.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462, 2020.
Jean Feydy. Geometric loss functions between sampled measures, images and volumes, 2019. URL
https: w/www kernel www.etterittitnee io/geomloss/
Jean Feydy, Thibault Séjouréé, Langois-Xavier Vialard, Shun-Ichi Amari, Alain Trouvé, and
Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. arXiv preprint arXiv: 1810 2018
10

---

Under review as a conference paper at ICLR 2021
Jean Feydy, Pieree Roussillon, Alain Trouvé, and Pietro Gori. Fast and scalable optimal transport for brain tractograms. In International Conference on Medical Image Computing and Computer- Assisted Intervention, pp. 636-644. Springer, 2019.
Charlie Frogner, Farzaneh Mirzazadeh, and Justin Solomon. Learning embeddings into entropic wasserstein spaces. arXiv preprint arXiv:/905.03322 2019.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn diver- gences. arXiv preprint 1oXxvv/706.002220 2017.
Inc. Google. The quick, draw! dataset, 2020.
URL https: com/
Matthieu Heitz, Nicolas Bonneel, David Coeurjolly, Marco Cuturi, and Gabriel Peyré. Ground Metric Learning on Graphs. Technical Report 1rXiv:1911.03111 November 2019.
L Kantorovich. On the transfer of masses (in russian). In Doklady Akademii Nauk, volume 37, pp. 227-229, 1942.
Wei-Sheng Lai, Jia-Bin Huang, Narendra Ahuja, and Ming-Hsuan Yang. Deep laplacian pyramid networks for fast and accurate super-resolution. In Proceedings of the IEEE conference com- puter vision and pattern recognition, pp. 624-632, 2017.
Christian Ledig Lucas Theis, Ferenc Huszár, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, et al. Photo-realistic sin- gle image super-resolution using a generative adversarial network. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4681-4690, 2017
Stamatios Lefkimmiatis. Non-local color image denoising with convolutional neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3587- 3596, 2017.
Guilin Liu, Fitsum A Reda, Kevin J Shih, Ting-Chun Wang, Andrew Tao, and Bryan Catanzaro. Image inpainting for irregular holes using partial convolutions. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 85-100, 2018
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter. Wlisee--assersteinn flows: Nonparametric generative modeling via optimal transport and diffu- sions. International Conference on Machine Learning, pp. 4104-4113. PMLR, 2019.
Iya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv
preprint 1608. 03988, 2016
Quentin Mérigot, Alex Delalande, and Frederic Chazal. Quantitative stability of optimal transport maps and linearization of the 2-wasserstein space. volume 108 of Proceedings of Machine Learn- ing Research, pp. 3186-3196 2020.
Alberto Maria Metelli, Amarildo Likmeta, and Marcello Restelli, Propagating uncertainty in rein- forcement learning via wasserstein barycenters. In Advances in Neural Information Processing Systems, pp. 4333-4345, 2019.
Liang Mi, Wen Zhang, Xianfeng Gu, and Yalin Wang. Variational Wasserstein clustering. In Pro- ceedings of the European Conference on Computer Vision (ECCV), pp. 322-337, 2018.
Caroline Moosmüller and Alexander Cloninger. Linear optimal transport embedding: Provable fast wasserstein distance computation and classification for nonlinear problems, 2020.
Georges Nader and Gael Guennebaud. Instant transport maps on 2d grids. ACM Trans. Graph.,, 37 (6), 2018. ISSN 0730-0301.
SA Nane, SK Nayar, and H Murase. Columbia object image library: Coil-20. Dept. Comp. Sci., Columbia University, New York, Tech. Rep, 1996.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5-6)335--61) 2019.
11

---

Under review as a conference paper at ICLR 2021
Antoine Rolet, Marco Cuturi, and Gabriel Peyré. Fast dictionary learning with smoothed wasser- stein loss. In Artificial Intellgence and Statistics, pp. 330-638,70 2016.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference Medical image computing and computer- assisted intervention, pp. 234-241. Springer, 2015.
Morgan A. Schmitz, Matthieu Heitz, Nicolas Bonneel, Fred Maurice Ngolè Mboula, David Coeur- jolly, Marco Cuturi, Gabriel Peyré, and Jean-Luc Starck. Wasserstein dictionary learning: Op- timal transport-based unsupervised non-linear dictionary learning. SIAM Journal on Imaging Sciences, (1)). 2018.
Bernhard Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing. A1(3):A1443AA144 2019.
Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transporta- tion on geometric domains. ACM Transactions on Graphics (TOG), 34(4):1-11, 2015.
Ying Tai, Jian Yang. and Xiaoming Liu. Image super-resolution via deep recursive residual network. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3147- 3155, 2017.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing in- gredient for fast stylization. arXiv preprint arXiv:1607.08022, 2016.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Deep image prior. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 9446-9454, 2018.
Wei Wang, Dejan Slepčev, Saurav Basu, John A Ozolek, and Gustavo K Rohde. A linear optimal transportation framework for quantifying and visualizing variations in sets of images. Interna- tional journal of computer vision, 101(2):554-20) 2013.
Junyuan Xie, Linli Xu, and Enhong Chen. Image denoising and inpainting with deep neural net- works. In Advances in neural information processing systems, pp. 341-349, 2012.
Raymond A Ych, Chen Chen, Teck Yian Lim, Alexander G Schwing, Mark Hasegawa-Johnsen and Minh N Do. Semantic image inpainting with deep generative models. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5485-54449 2017.
A LEARNING STRATEGY
Instead of using fixed learning rate decreasing learning rate, choose learning rate schedule with warm restart as proposed by Loshchilov & Hutter (2016). The learning schedule is shown in Figure TO the learning rate decreased and is periodically restarted to its initial value, the period increasing as the number of epochs grows. This schedule was chosen after comparing with stepwise schedules or constant learning rates and yielded better convergence in practice.
B ADDITIONAL RESULTS
To better show the limitations of the generalization of our network when the number of inputs is 2, we show additional interpolations between 2 objects from the Coil20 in figure [T] There are two reasons for these bad results: first, our model trained synthetic shape contours and do not look at all like these images. Furthermore, the cup image seem to be even more challenging than the car image for our network, and our best explanation for this failure that the cup covers almost the whole image. We provide additional experiments showing barycenters of 5 sketches Figure I2 The weights evolve linearly inside the pentagon. As a stress test, we also show a barycenter of 100 cats with equal weights in Fig. 13 and compare it with barycenter computed with GeomLoss. While both results recover more or less the global shape of the cat, details are clearly lost and our result looks much smoother.
12

---

Under review as a conference paper at ICLR 2021
10-4
Ir Value
10-5
10-6
10-7
5
10
Epochs
Figure 10: Learning rate schedule used to train our models, following the SGDR method described by Loshchilov & Hutter (2016). Our training runs for total of 31 epochs. Compared a constant learning rate to stepwise schedules, SGDR has empirically shown a better convergence in our context.
Geoml
Our Model
Figure 11: Additional interpolations between two 512 x 512 images from the Coil20 dataset using GeomLoss and our model. Mass is inverted.
C LINEARIZED BARYCENTERS
Fig. 14 shows the error introduced by using linearized version of Wasserstein barycenters (Nader
& Guennebaud 2018 Wang al. [2013] Moosmüller & Cloninger 2020 Mérigot al. 2020). Our predicted barycenters reflect this error.
13

---

Under review as a conference paper at ICLR 2021
Figure 12: Interpolations between 5 inputs from Quick, Draw!, shown as pentagons. Left pentagon corresponds to GeomLoss barycenters while the right one shows predictions of our model trained on our synthetic dataset.
14

---

Under review as a conference paper at ICLR 2021
GoomLoss
Our DCNN
Figure 13: Stress test. We predict a barycenter of 100 of the QuickDraw dataset, with equal weights.
input
input
Geomloss
(10)
0.4382
>ı 0.5618
Figure 14: Wasserstein barycenter computed from a pair of inputs respectively using Geomloss with only one descent step, Geomloss with 10 descent steps and using model trained synthetic training dataset.
15


Reviews:
[{"review_id": "wrmkbEuxm9m", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "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 numerical method to approximate 2d Wasserstein barycenters, that relies on a convolutional network. Namely, the proposed architecture is composed of\n1.\tContractive paths, made of series of 2 convolutional layers, relu activation and a pooling step, characterized by weight sharing.\n2.\tAn expansive path, symmetrically constructed, with upsampling by NN interpolation, and a final softmax activation.\n\nThough fixed during training, the varying number of contracting paths in the test setting is able to predict the barycenter of a varying number of measures. The KL loss is used to train the network predicting the Wasserstein barycenter.\nThe method is assessed on the Quick-draw dataset, to produce barycenters of two measures and its generalization to barycenters of n measures.\n\nStrong points of the paper include:\n1.\tThe writing of the paper is clear.\n2.\tThe goal to propose a speedup for the computation of Wasserstein barycenters with a varying number of input measures is pertinent.\n\nWeak points include:\n1.\tThe paper proposes an architecture tailored for very specific setting of 2d images. Its extensions to higher dimensional or non-image settings are limited and are not discussed.\n2.\tThough mentioned as closely related works, the paper is missing important baselines, for instance (Claici et al.,2018).\n3.\tThe fact that only one dataset is considered undermines the paper, which focuses on numerics.\n4.\tThe figure at the end of page 6 (which is lacking both title and legend) does not provide a relevant metric to assess the benefits of the proposed method. \n5.\tThe paper does not provide other metrics (such as inception scores) to assess the performance of the methods other than visually.\n\nI recommend a reject on the aforementioned grounds.", "extra_scores": {}, "extra_text": {}}, {"review_id": "pVu1uO1LcpN", "reviewer": "AnonReviewer4", "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": "A modification of U-Net to compute Wasserstein barycenters\n\nThis paper presents a deep network architecture which learns to generate Wasserstein barycenters associated with a given set of input measures. The paper is clearly written and presents a fairly objective view on their work. It gives one more example of optimization task (here the convex minimization problem of Wasserstein barycenters) that can be replaced by deep neural networks with seemingly satisfactory results.\nOverall, although the method is relatively straightforward, it has the merit of showing that fast approximations of linear approximations of Wasserstein barycenters can be performed at a cheap forward pass cost of a U-Net. Note that the method has no built-in information on the structure of the problem e.g. preservation of total mass for instance or more higher-level properties.  These fast approximations could be proven to be useful in other settings outside machine learning too.\n\nThe benefit of the proposed architecture is that it can be adapted to compute the barycenters for any given number N of inputs while only trained on pairs (N =2) of input measures. Apart from minor modifications in the pipeline, the architecture remains close to that of a U-Net with the modification that the barycenters weights are explicitly used to compute averages of the activations at different layers of the Deep CNN. This work is quite weak in terms of novelty (structure of the network for instance).\n\nThe experiments are compared against and trained with the GeomLoss code which seems to be state of the art for computing entropic regularity. They also compare with Wasserstein embeddings methods favorably. \n\n— References to literature in using deep learning for solving inverse problems would be welcome. Wasserstein barycenter is a particular instance of such class.\n\n— Comments on the sensitivity of the results to training data would be welcome.\n\n— I find the description of the algorithm for building the training data quite unclear. If I understand correctly, the training data is an approximation of Wasserstein barycenters between two measures. These approximations are simply computed with one step of the gradient descent of the Sinkhorn divergence and use an average of the corresponding transport maps. The cost of computation of these approximate barycenters are essentially two sinkhorn loops up to the given tolerance. This would have been interesting to have a comparison, in the supplementary material, of the difference between these approximations and more accurate approximations of Wasserstein barycenters such as given by entropic regularization with very small regularization parameter and discuss in more details what can be the impact of such inaccuracies. The authors say « While using more iterations of gradient descent yields more accurate results and removes this linearity, it also prevents easy combination and makes the dataset generation intractable ». I do not understand this statement. Cost of the gradient descent linearly scales with the number of steps, so what do the authors mean by « easy combination », the linear approximation step ?\n\nMinor remarks: \n— the (not numbered) figure on page 6 called adjacent is almost impossible to read and it is not explained enough.\n—  A dot is missing after KL-divergence on page 6.\n— Why is DWE performing so badly in figure 3 in comparison with figure 4 ?", "extra_scores": {}, "extra_text": {}}, {"review_id": "nAWHJe14ScA", "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": "Summary:\n\nThis paper proposes a network architecture and a training model for learning Wasserstein barycenters of 512x512 discrete probability distributions. Results are compared with other approximation methods.\n\nReason for my score:\n\nThis paper is ell written, and the idea is interesting.  However, the elephant in the room is: How close are these learned barycenters from the actual barycenter? For me this is the main question, that cannot be answered with CNN approaches, and thus, makes it difficult to translate into uses for optimal transport techniques. This limits the impact that the paper can have on a broader audience.\n\nThe authors made a great effort to provide a rather complete review of the literature, and the provided experiments are sound. However.\n\n1. The comparison of the results is made against another approximating method. Again, I would like to see some evidence that the learned model is actually generating good approximations of the Wasserstein Barycenter.\n\n2. It is a less exciting if the loss function used is not OT itself but KL.\n\n3. The main advantage of entropy regularized methods is that one can actually control how close the computed WB is from the actual WB. This is a main issue. In the proposed method, we have no way to know this, and I cannot avoid thinking this is just an informed guess.\n\n4. I believe one main advantage of the proposed method would be to generate initialization points of entropy regularized approaches.\n\n5. The proposed architecture is well explained. Why not other architecture, some experiments exploring other ways would be appreciated as well.\n\n6. For some basic distributions it would have been more useful to compare the learned barycenter with the output of the linear program. At least in a couple of cases.  This can provide some evidence that the learned barycenter is actually a good approximation.\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Given the reviewer's exchange with the authors, and my own examination of the paper, I don't think that it can be accepted in the present form.

First, since this paper aims at solving an optimization problem (for which existing methods exist, with theoretical guarantees) via a NN, it is important to compare appropriately to those methods, which is not done here.

Further, there are possible issues when applying these only to 2D data, and it is possible that it would not extend appropriately to other types of geometries, and costs in OT problems.
Decision: Reject