# Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 28 May 2025.

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

- [Schr"odinger Bridge Matching for Tree-Structured Costs and Entropic Wasserstein Barycentres](https://mlsanity.com/paper/2506.17197/) (2025)
- [Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation](https://mlsanity.com/paper/2510.04602/) (2025)
- [Towards Universal Wasserstein Barycenters through Flow Matching](https://mlsanity.com/paper/2609.38547/) (2026)

## Abstract (arXiv)

We present a novel method for efficiently computing optimal transport maps and Wasserstein barycenters in high-dimensional spaces. Our approach uses conditional normalizing flows to approximate the input distributions as invertible pushforward transformations from a common latent space. This makes it possible to directly solve the primal problem using gradient-based minimization of the transport cost, unlike previous methods that rely on dual formulations and complex adversarial optimization. We show how this approach can be extended to compute Wasserstein barycenters by solving a conditional variance minimization problem. A key advantage of our conditional architecture is that it enables the computation of barycenters for hundreds of input distributions, which was computationally infeasible with previous methods. Our numerical experiments illustrate that our approach yields accurate results across various high-dimensional tasks and compares favorably with previous state-of-the-art methods.