# Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks

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

## Key facts

- Cited by 50 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 8 Jul 2020.

## Built on by (most cited first, 30 of 50)

- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Optimal Flow Matching: Learning Straight Trajectories in Just One Step](https://mlsanity.com/paper/2403.13117/) (2024)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Score-based Generative Neural Networks for Large-Scale Optimal Transport](https://mlsanity.com/paper/2110.03237/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization](https://mlsanity.com/paper/2102.01752/) (2021)
- [Neural Monge Map estimation and its applications](https://mlsanity.com/paper/2106.03812/) (2021)
- [Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent](https://mlsanity.com/paper/2106.08502/) (2021)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Dataset Distillation via the Wasserstein Metric](https://mlsanity.com/paper/2311.18531/) (2023)
- [Curriculum Reinforcement Learning using Optimal Transport via Gradual Domain Adaptation](https://mlsanity.com/paper/2210.10195/) (2022)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Energy-guided Entropic Neural Optimal Transport](https://mlsanity.com/paper/2304.06094/) (2023)
- [Wasserstein Iterative Networks for Barycenter Estimation](https://mlsanity.com/paper/2201.12245/) (2022)
- [Building the Bridge of Schr"odinger: A Continuous Entropic Optimal Transport Benchmark](https://mlsanity.com/paper/2306.10161/) (2023)
- [Input Convex Gradient Networks](https://mlsanity.com/paper/2111.12187/) (2021)
- [Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings](https://mlsanity.com/paper/2209.15621/) (2022)
- [Tree-Based Diffusion Schr"odinger Bridge with Applications to Wasserstein Barycenters](https://mlsanity.com/paper/2305.16557/) (2023)
- [Estimating Barycenters of Distributions with Neural Optimal Transport](https://mlsanity.com/paper/2402.03828/) (2024)
- [Energy-Guided Continuous Entropic Barycenter Estimation for General Costs](https://mlsanity.com/paper/2310.01105/) (2023)
- [Nonlinear Filtering with Brenier Optimal Transport Maps](https://mlsanity.com/paper/2310.13886/) (2023)
- [Learning High Dimensional Wasserstein Geodesics](https://mlsanity.com/paper/2102.02992/) (2021)
- [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)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Light Unbalanced Optimal Transport](https://mlsanity.com/paper/2303.07988/) (2023)

## Abstract (arXiv)

Wasserstein Barycenter is a principled approach to represent the weighted mean of a given set of probability distributions, utilizing the geometry induced by optimal transport. In this work, we present a novel scalable algorithm to approximate the Wasserstein Barycenters aiming at high-dimensional applications in machine learning. Our proposed algorithm is based on the Kantorovich dual formulation of the Wasserstein-2 distance as well as a recent neural network architecture, input convex neural network, that is known to parametrize convex functions. The distinguishing features of our method are: i) it only requires samples from the marginal distributions; ii) unlike the existing approaches, it represents the Barycenter with a generative model and can thus generate infinite samples from the barycenter without querying the marginal distributions; iii) it works similar to Generative Adversarial Model in one marginal case. We demonstrate the efficacy of our algorithm by comparing it with the state-of-art methods in multiple experiments.