Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation

Eduardo Fernandes Montesuma, Yassir Bendou, Mike Gartrell
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:4595-4606, 2026.

Abstract

{Wasserstein} barycenters provide a principled approach for aggregating probability measures, while preserving the geometry of their ambient space. Existing discrete methods are not because as they assume access to the complete set of samples from the input measures. Meanwhile, neural network approaches do scale well, but rely on complex optimization problems and cannot easily incorporate label information. We address these limitations through gradient flows in the space of probability measures. Through time discretization, we achieve a scalable algorithm that i) relies on mini-batch optimal transport, ii) accepts modular regularization through task-aware functions, and iii) seamlessly integrates supervised information into the ground-cost. We empirically validate our approach on domain adaptation benchmarks that span computer vision, neuroscience, and chemical engineering. Our method establishes a new state-of-the-art {Wasserstein} barycenter solver, with labeled barycenters consistently outperforming unlabeled ones. Our code at https://github.com/SigmaNova/barycentric-gradient-flows

Cite this Paper

BibTeX
@InProceedings{pmlr-v337-fernandes-montesuma26a, title = {{Wasserstein} Gradient Flows for Scalable and Regularized Barycenter Computation}, author = {Fernandes Montesuma, Eduardo and Bendou, Yassir and Gartrell, Mike}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {4595--4606}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/fernandes-montesuma26a/fernandes-montesuma26a.pdf}, url = {https://proceedings.mlr.press/v337/fernandes-montesuma26a.html}, abstract = {{Wasserstein} barycenters provide a principled approach for aggregating probability measures, while preserving the geometry of their ambient space. Existing discrete methods are not because as they assume access to the complete set of samples from the input measures. Meanwhile, neural network approaches do scale well, but rely on complex optimization problems and cannot easily incorporate label information. We address these limitations through gradient flows in the space of probability measures. Through time discretization, we achieve a scalable algorithm that i) relies on mini-batch optimal transport, ii) accepts modular regularization through task-aware functions, and iii) seamlessly integrates supervised information into the ground-cost. We empirically validate our approach on domain adaptation benchmarks that span computer vision, neuroscience, and chemical engineering. Our method establishes a new state-of-the-art {Wasserstein} barycenter solver, with labeled barycenters consistently outperforming unlabeled ones. Our code at https://github.com/SigmaNova/barycentric-gradient-flows} }
Endnote
%0 Conference Paper %T Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation %A Eduardo Fernandes Montesuma %A Yassir Bendou %A Mike Gartrell %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-fernandes-montesuma26a %I PMLR %P 4595--4606 %U https://proceedings.mlr.press/v337/fernandes-montesuma26a.html %V 337 %X {Wasserstein} barycenters provide a principled approach for aggregating probability measures, while preserving the geometry of their ambient space. Existing discrete methods are not because as they assume access to the complete set of samples from the input measures. Meanwhile, neural network approaches do scale well, but rely on complex optimization problems and cannot easily incorporate label information. We address these limitations through gradient flows in the space of probability measures. Through time discretization, we achieve a scalable algorithm that i) relies on mini-batch optimal transport, ii) accepts modular regularization through task-aware functions, and iii) seamlessly integrates supervised information into the ground-cost. We empirically validate our approach on domain adaptation benchmarks that span computer vision, neuroscience, and chemical engineering. Our method establishes a new state-of-the-art {Wasserstein} barycenter solver, with labeled barycenters consistently outperforming unlabeled ones. Our code at https://github.com/SigmaNova/barycentric-gradient-flows
APA
Fernandes Montesuma, E., Bendou, Y. & Gartrell, M.. (2026). Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:4595-4606 Available from https://proceedings.mlr.press/v337/fernandes-montesuma26a.html.

Related Material