Computing Wasserstein Barycenters via Operator Splitting: The Method of Averaged Marginals
Abstract.
The Wasserstein barycenter (WB) is an important tool for summarizing sets of probability measures. It finds applications in applied probability, clustering, image processing, etc. When the measures’ supports are finite, computing a (balanced) WB can be done by solving a linear optimization problem whose dimensions generally exceed standard solvers’ capabilities. In the more general setting where measures have different total masses, we propose a convex nonsmooth optimization formulation for the so-called unbalanced WB problem. Due to their colossal dimensions, we introduce a decomposition scheme based on the Douglas–Rachford splitting method that can be applied to both balanced and unbalanced WB problem variants. Our algorithm, which has the interesting interpretation of being built upon averaging marginals, operates a series of simple (and exact) projections that can be parallelized and even randomized, making it suitable for large-scale datasets. Numerical comparisons against state-of-the-art methods on several data sets from the literature illustrate the method’s performance.
Keywords
MSC codes
Get full access to this article
View all available purchase options and get full access to this article.
References
1.
M. Agueh and G. Carlier, Barycenters in the Wasserstein space, SIAM J. Math. Anal., 43 (2011), pp. 904–924, https://doi.org/10.1137/100805741.
2.
J. M. Altschuler and E. Boix-Adsera, Wasserstein barycenters can be computed in polynomial time in fixed dimension, J. Mach. Learn. Res., 22 (2021), 44.
3.
J. M. Altschuler and E. Boix-Adserà, Wasserstein barycenters are NP-Hard to compute, SIAM J. Math. Data Sci., 4 (2022), pp. 179–203, https://doi.org/10.1137/21M1390062.
4.
E. Anderes, S. Borgwardt, and J. Miller, Discrete wasserstein barycenters: Optimal transport for discrete data, Math. Methods Oper. Res., 84 (2016), pp. 389–409, https://doi.org/10.1007/s00186-016-0549-x.
5.
H. H. Bauschke and P. L. Combettes, Convex Analysis and Monotone Operator Theory in Hilbert Spaces, CMS Books Math. 2nd ed., Springer, Cham, 2017, https://doi.org/10.1007/978-3-319-48311-5.
6.
J.-D. Benamou, G. Carlier, M. Cuturi, L. Nenna, and G. Peyré, Iterative Bregman projections for regularized transportation problems, SIAM J. Sci. Comput., 37 (2015), pp. 1111–1138, https://doi.org/10.1137/141000439.
7.
D. P. Bertsekas, Convex Optimization Algorithms, Vol. 1, Athena Scientific, 2015.
8.
S. Borgwardt, An LP-based, strongly-polynomial 2-approximation algorithm for sparse Wasserstein barycenters, Oper. Res., 22 (2022), pp. 1511–1551, https://doi.org/10.1007/s12351-020-00589-z.
9.
S. Borgwardt and S. Patterson, Improved linear programs for discrete barycenters, INFORMS J. Optim., 2 (2020), pp. 14–33, https://doi.org/10.1287/ijoo.2019.0020.
10.
S. Borgwardt and S. Patterson, On the computational complexity of finding a sparse Wasserstein barycenter, J. Comb. Optim., 41 (2021), pp. 736–761.
11.
G. Carlier, A. Oberman, and E. Oudet, Numerical methods for matching for teams and Wasserstein barycenters, ESAIM Math. Model. Numer. Anal., 49 (2015), pp. 1621–1642, https://doi.org/10.1051/m2an/2015033.
12.
P. L. Combettes and J. Eckstein, Asynchronous block-iterative primal-dual decomposition methods for monotone inclusions, Math. Program., 168 (2018), pp. 645–672, https://doi.org/10.1007/s10107-016-1044-0.
13.
P. L. Combettes and J.-C. Pesquet, Stochastic quasi-Fejér block-coordinate fixed point iterations with random sweeping, SIAM J. Optim., 25 (2015), pp. 1221–1248, https://doi.org/10.1137/140971233.
14.
L. Condat, Fast projection onto the simplex and the \(\boldsymbol{{l}_{\mathbf{1}}}\)ball, Math. Program., 158 (2016), pp. 575–585.
15.
M. Cuturi, Sinkhorn distances: Lightspeed computation of optimal transport, in Advances in Neural Information Processing Systems, Vol. 26, C. J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, eds., Curran Associates, 2013, pp. 1–9.
16.
M. Cuturi and A. Doucet, Fast computation of Wasserstein barycenters, in Proceedings of the 31st International Conference on Machine Learning, Proceedings of Machine Learning Research 32, E. P. Xing and T. Jebara, eds., PMLR, Bejing, China, 2014, pp. 685–693.
17.
M. Cuturi and G. Peyré, A smoothed dual approach for variational Wasserstein problems, SIAM J. Imaging Sci., 9 (2016), pp. 320–343, https://doi.org/10.1137/15M1032600.
18.
W. de Oliveira, C. Sagastizábal, D. D. J. Penna, M. E. P. Maceira, and J. M. Damázio, Optimal scenario tree reduction for stochastic streamflows in power generation planning problems, Optim. Methods Software, 25 (2010), pp. 917–936.
19.
J. Douglas and H. H. Rachford, On the numerical solution of heat conduction problems in two and three space variables, Trans. Amer. Math. Soc., 82 (1956), pp. 421–439.
20.
J. Eckstein and D. P. Bertsekas, On the Douglas—Rachford splitting method and the proximal point algorithm for maximal monotone operators, Math. Program., 55 (1992), pp. 293–318, https://doi.org/10.1007/bf01581204.
21.
A. Fu, J. Zhang, and S. Boyd, Anderson accelerated Douglas–Rachford splitting, SIAM J. Sci. Comput., 42 (2020), pp. A3560–A3583, https://doi.org/10.1137/19m1290097.
22.
A. Gramfort, G. Peyré, and M. Cuturi, Fast optimal transport averaging of neuroimaging data, in Information Processing in Medical Imaging, S. Ourselin, D. C. Alexander, C.-F. Westin, and M. J. Cardoso, eds., Springer, Cham, 2015, pp. 261–272.
23.
F. Heinemann, M. Klatt, and A. Munk, Kantorovich–Rubinstein distance and barycenter for finitely supported measures: Foundations and algorithms, Appl. Math. Optim., 87 (2022), 4, https://doi.org/10.1007/s00245-022-09911-x.
24.
F. Iutzeler, P. Bianchi, P. Ciblat, and W. Hachem, Asynchronous distributed optimization using a randomized alternating direction method of multipliers, in Proceedings of the 52nd IEEE Conference on Decision and Control, IEEE, 2013, pp. 3671–3676, https://doi.org/10.1109/cdc.2013.6760448.
25.
G. Peyré, BregmanOT, 2014, https://github.com/gpeyre/2014-SISC-BregmanOT.
26.
G. Peyré and M. Cuturi, Computational optimal transport: With applications to data science, Found. Trends Mach. Learn., 11 (2019), pp. 355–607, https://doi.org/10.1561/2200000073.
27.
G. C. Pflug and A. Pichler, Multistage Stochastic Optimization, Springer Ser. Oper. Res. Financ. Eng., Springer, Cham, 2014, https://doi.org/10.1007/978-3-319-08843-3.
28.
G. Puccetti, L. Rüschendorf, and S. Vanduffel, On the computation of Wasserstein barycenters, J. Multivariate Anal., 176 (2020), 104581, https://doi.org/10.1016/j.jmva.2019.104581.
29.
Y. Rubner, C. Tomasi, and L. J. Guibas, The earth mover’s distance as a metric for image retrieval, Int. J. Comput. Vision, 40 (2000), pp. 99–121, https://doi.org/10.1023/A:1026543900054.
30.
T. Sejourne, G. Peyre, and F.-X. Vialard, Unbalanced optimal transport, from theory to numerics, Handb. Numer. Anal., 24 (2023), pp. 407–471, https://doi.org/10.1016/bs.hna.2022.11.003.
31.
D. Simon and A. Aberdam, Barycenters of natural images constrained Wasserstein barycenters for image morphing, in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020, pp. 7910–7919.
32.
R. Sinkhorn, Diagonal equivalence to matrices with prescribed row and column sums. ii, Proc. Amer. Math. Soc., 45 (1974), pp. 195–198.
33.
G. Tartavel, G. Peyré, and Y. Gousseau, Wasserstein loss for image synthesis and restoration, SIAM J. Imaging Sci., 9 (2016), pp. 1726–1755, https://doi.org/10.1137/16M1067494.
34.
Tijmen, affnist, 2013, https://www.cs.toronto.edu/∼tijmen/affNIST/.
35.
J. v. Lindheim, Simple approximative algorithms for free-support Wasserstein barycenters, Comput. Optim. Appl., 85 (2023), pp. 213–246, https://doi.org/10.1007/s10589-023-00458-3.
36.
C. Villani, Optimal Transport: Old and New, Grundlehren Math. Wiss. 338, Springer-Verlag, Berlin, 2009.
37.
H. Wang and A. Banerjee, Bregman alternating direction method of multipliers, in Advances in Neural Information Processing Systems, Vol. 27, Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger, eds., Curran Associates, 2014, pp. 1–9.
38.
J.-P. Watson and D. L. Woodruff, Progressive hedging innovations for a class of stochastic mixed-integer resource allocation problems, Comput. Manag. Sci., 8 (2011), pp. 355–370, https://doi.org/10.1007/s10287-010-0125-4.
39.
Z. Xu, M. Figueiredo, and T. Goldstein, Adaptive ADMM with spectral penalty parameter selection, in Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, Proceedings of Machine Learning Research 54, PMLR, A. Singh and J. Zhu, eds., 2017, pp. 718–727.
40.
J. Ye and J. Li, Scaling up discrete distribution clustering using ADMM, in Proceedings of the 2014 IEEE International Conference on Image Processing (ICIP), 2014, pp. 5267–5271, https://doi.org/10.1109/ICIP.2014.7026066.
41.
J. Ye, P. Wu, J. Z. Wang, and J. Li, Fast discrete distribution clustering using Wasserstein barycenter with sparse support, IEEE Trans. Signal Process., 65 (2017), pp. 2317–2332, https://doi.org/10.1109/TSP.2017.2659647.
Information & Authors
Information
Published In
SIAM Journal on Mathematics of Data Science
Pages: 1000 - 1026
DOI: 10.1137/23M1584228
ISSN (online): 2577-0187
Copyright
© 2024 Society for Industrial and Applied Mathematics.
History
Submitted: 5 July 2023
Accepted: 29 July 2024
Published online: 18 October 2024
Keywords
MSC codes
Authors
Metrics & Citations
Metrics
Citations
If you have the appropriate software installed, you can download article citation data to the citation manager of your choice. Simply select your manager software from the list below and click Download.
Cited By
There are no citations for this item
View Options
- Access via your Institution
- Questions about how to access this content? Contact SIAM at [email protected].