# Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation

arXiv: https://arxiv.org/abs/2510.04602 · HTML: https://mlsanity.com/paper/2510.04602/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 72 references that are in the map.
- First posted on arXiv 6 Oct 2025.

## Builds on (most cited first, 30 of 36)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Deep Hashing Network for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1706.07522/) (2017)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [CBraMod: A Criss-Cross Brain Foundation Model for EEG Decoding](https://mlsanity.com/paper/2412.07236/) (2024)
- [A review of domain adaptation without target labels](https://mlsanity.com/paper/1901.05335/) (2019)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [An elementary introduction to information geometry](https://mlsanity.com/paper/1808.08271/) (2018)
- [Optimal Transport for Multi-source Domain Adaptation under Target Shift](https://mlsanity.com/paper/1803.04899/) (2018)
- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Minibatch optimal transport distances; analysis and applications](https://mlsanity.com/paper/2101.01792/) (2021)
- [Rates of Estimation of Optimal Transport Maps using Plug-in Estimators via Barycentric Projections](https://mlsanity.com/paper/2107.01718/) (2021)
- [Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization](https://mlsanity.com/paper/2102.01752/) (2021)
- [Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks](https://mlsanity.com/paper/2007.04462/) (2020)
- [Statistical optimal transport](https://mlsanity.com/paper/2407.18163/) (2024)
- [Scalable Bayes via Barycenter in Wasserstein Space](https://mlsanity.com/paper/1508.05880/) (2015)
- [Multi-source Domain Adaptation via Weighted Joint Distributions Optimal Transport](https://mlsanity.com/paper/2006.12938/) (2020)
- [Wasserstein barycenters are NP-hard to compute](https://mlsanity.com/paper/2101.01100/) (2021)
- [Wasserstein Iterative Networks for Barycenter Estimation](https://mlsanity.com/paper/2201.12245/) (2022)
- [Existence, Duality, and Cyclical monotonicity for weak transport costs](https://mlsanity.com/paper/1809.05893/) (2018)
- [Estimating Barycenters of Measures in High Dimensions](https://mlsanity.com/paper/2007.07105/) (2020)
- [Dataset Dynamics via Gradient Flows in Probability Space](https://mlsanity.com/paper/2010.12760/) (2020)
- [Estimating Barycenters of Distributions with Neural Optimal Transport](https://mlsanity.com/paper/2402.03828/) (2024)
- [Doubly Regularized Entropic Wasserstein Barycenters](https://mlsanity.com/paper/2303.11844/) (2023)
- [Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport](https://mlsanity.com/paper/2410.03974/) (2024)

## Built on by (most cited first, 1 of 1)

- [Multi-Domain Clustering via Measure Quantization](https://mlsanity.com/paper/2609.21664/) (2026)

## Abstract (arXiv)

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