Learning to Generate Wasserstein Barycenters

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—are even more computationally demanding as it requires to solve an optimization problem involving optimal transport distances. By training a deep convolutional neural network, we improve by a factor of 80 the computational speed of Wasserstein barycenters over the fastest state-of-the-art approach on the GPU, resulting in milliseconds computational times on $512\times 512$ 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 demonstrate the efficiency of our approach for computing barycenters of sketches and transferring colors between multiple images.

Appendices

A Learning Strategy

Instead of using a fixed learning rate or a decreasing learning rate, we choose a learning rate schedule with warm restart as proposed by [38]. The learning schedule is shown in Fig. 12: 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 both Adam and SGD with stepwise schedules, and yielded better convergence in practice.

Fig. 15

Interpolations between 5 inputs from Quick, Draw!, shown as pentagons. Upper pentagon corresponds to GeomLoss barycenters while the lower one shows predictions of our model trained on our synthetic dataset

Fig. 16

Stress test. We predict a barycenter of 100 cats of the Quick, Draw! dataset, with equal weights

B Test of Equivariance

Wasserstein barycenters are equivariant under rotation, translation and scaling. This amounts to $Barycenter(\{T(\mu _i), \lambda _i\}_i) = T(Barycenter(\{\mu _i, \lambda _i\}_i))$ for T a rotation, translation or scaling. We verify this behavior qualitatively on the output of our network on two examples shown in Fig. 13.

Fig. 17

Comparisons when using a single or ten Bregman projection steps to minimize the Sinkhorn divergence functional in the debiased barycenter approach of Janati et al. [31]. While using ten iterations makes it slower to converge, this leads to sharper results and is the default approach adopted in the authors’ implementation: we thus used this value in the rest of the paper.

Fig. 18

Color grading obtained by transferring the colors of $n=3$ images onto a target image, aiming at reproducing with our method the results from Bonneel et al. [9], Fig. 12. Results are shown in triangles (Left: interpolated chrominance histograms; Right: corresponding transfer results). The images corresponding to the target chrominance histogram $\nu$ and to the histograms $\mu _{i}$—which are interpolated to obtain a barycenter—are shown in top row. Each $\mu _{i}$ corresponds to a vertex of the triangle in a clockwise order beginning with $i=1$ at the uppermost vertex. Each row presents the results for a different method, from top to bottom: GeomLoss, our model trained on synthetic shape contours (ContoursDS) and our model trained on chrominance histograms from Flickr images (HistoDS)

C Additional Results

While the model presented in this paper uses 6 depth levels (following the convention of Fig. 1), we provide additional experiments showing the differences between our model with different number of depth levels in Fig. 14. Note that while the gap in performance is particularly important between the model with 4 depth levels (DL), 5DL and 6DL, there is much less difference between 6DL and 7DL.

We provide additional experiments showing barycenters of 5 sketches in Fig. 15. The weights evolve linearly inside the pentagon. As a stress test, we also show a barycenter of 100 cats with equal weights in Fig. 16 and compare it with a 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.

In the debiasing approach of Janati et al. [31], a single iteration of Bregman projections to minimize the functional in the variable d is used in their paper. However, their available implementation uses 10 iterations. Figure 17 shows the (minor) difference in quality between these two approaches. Our comparisons were made against the original implementation.

Finally, we provide an additional color transfer experiment in Fig. 18 reproducing an experiment from [9] with our model trained with ContoursDS and HistoDS.

D Linearized Barycenters

Figure 19 shows the error introduced by using a linearized version of Wasserstein barycenters [40, 43, 44, 59]. Our predicted barycenters reflect this error.

Fig. 19

Wasserstein barycenter computed from a pair of inputs, respectively, using GeomLoss with only one descent step, GeomLoss with 10 descent steps and using our model trained on our synthetic training dataset