Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows

Gabriele Visentin, Patrick Cheridito
Proceedings of the 42nd International Conference on Machine Learning, PMLR 267:61564-61581, 2025.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v267-visentin25a, title = {Computing Optimal Transport Maps and {W}asserstein Barycenters Using Conditional Normalizing Flows}, author = {Visentin, Gabriele and Cheridito, Patrick}, booktitle = {Proceedings of the 42nd International Conference on Machine Learning}, pages = {61564--61581}, year = {2025}, editor = {Singh, Aarti and Fazel, Maryam and Hsu, Daniel and Lacoste-Julien, Simon and Berkenkamp, Felix and Maharaj, Tegan and Wagstaff, Kiri and Zhu, Jerry}, volume = {267}, series = {Proceedings of Machine Learning Research}, month = {13--19 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v267/main/assets/visentin25a/visentin25a.pdf}, url = {https://proceedings.mlr.press/v267/visentin25a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows %A Gabriele Visentin %A Patrick Cheridito %B Proceedings of the 42nd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2025 %E Aarti Singh %E Maryam Fazel %E Daniel Hsu %E Simon Lacoste-Julien %E Felix Berkenkamp %E Tegan Maharaj %E Kiri Wagstaff %E Jerry Zhu %F pmlr-v267-visentin25a %I PMLR %P 61564--61581 %U https://proceedings.mlr.press/v267/visentin25a.html %V 267 %X 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.
APA
Visentin, G. & Cheridito, P.. (2025). Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows. Proceedings of the 42nd International Conference on Machine Learning, in Proceedings of Machine Learning Research 267:61564-61581 Available from https://proceedings.mlr.press/v267/visentin25a.html.

Related Material