Wasserstein Flow Matching: Generative Modeling Over Families of Distributions

Doron Haviv, Aram-Alexandre Pooladian, Dana Pe’Er, Brandon Amos
Proceedings of the 42nd International Conference on Machine Learning, PMLR 267:22238-22258, 2025.

Abstract

Generative modeling typically concerns transporting a single source distribution to a target distribution via simple probability flows. However, in fields like computer graphics and single-cell genomics, samples themselves can be viewed as distributions, where standard flow matching ignores their inherent geometry. We propose Wasserstein flow matching (WFM), which lifts flow matching onto families of distributions using the Wasserstein geometry. Notably, WFM is the first algorithm capable of generating distributions in high dimensions, whether represented analytically (as Gaussians) or empirically (as point-clouds). Our theoretical analysis establishes that Wasserstein geodesics constitute proper conditional flows over the space of distributions, making for a valid FM objective. Our algorithm leverages optimal transport theory and the attention mechanism, demonstrating versatility across computational regimes: exploiting closed-form optimal transport paths for Gaussian families, while using entropic estimates on point-clouds for general distributions. WFM successfully generates both 2D & 3D shapes and high-dimensional cellular microenvironments from spatial transcriptomics data. Code is available at WassersteinFlowMatching.

Cite this Paper

BibTeX
@InProceedings{pmlr-v267-haviv25a, title = {{W}asserstein Flow Matching: Generative Modeling Over Families of Distributions}, author = {Haviv, Doron and Pooladian, Aram-Alexandre and Pe'Er, Dana and Amos, Brandon}, booktitle = {Proceedings of the 42nd International Conference on Machine Learning}, pages = {22238--22258}, 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/haviv25a/haviv25a.pdf}, url = {https://proceedings.mlr.press/v267/haviv25a.html}, abstract = {Generative modeling typically concerns transporting a single source distribution to a target distribution via simple probability flows. However, in fields like computer graphics and single-cell genomics, samples themselves can be viewed as distributions, where standard flow matching ignores their inherent geometry. We propose Wasserstein flow matching (WFM), which lifts flow matching onto families of distributions using the Wasserstein geometry. Notably, WFM is the first algorithm capable of generating distributions in high dimensions, whether represented analytically (as Gaussians) or empirically (as point-clouds). Our theoretical analysis establishes that Wasserstein geodesics constitute proper conditional flows over the space of distributions, making for a valid FM objective. Our algorithm leverages optimal transport theory and the attention mechanism, demonstrating versatility across computational regimes: exploiting closed-form optimal transport paths for Gaussian families, while using entropic estimates on point-clouds for general distributions. WFM successfully generates both 2D & 3D shapes and high-dimensional cellular microenvironments from spatial transcriptomics data. Code is available at WassersteinFlowMatching.} }
Endnote
%0 Conference Paper %T Wasserstein Flow Matching: Generative Modeling Over Families of Distributions %A Doron Haviv %A Aram-Alexandre Pooladian %A Dana Pe’Er %A Brandon Amos %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-haviv25a %I PMLR %P 22238--22258 %U https://proceedings.mlr.press/v267/haviv25a.html %V 267 %X Generative modeling typically concerns transporting a single source distribution to a target distribution via simple probability flows. However, in fields like computer graphics and single-cell genomics, samples themselves can be viewed as distributions, where standard flow matching ignores their inherent geometry. We propose Wasserstein flow matching (WFM), which lifts flow matching onto families of distributions using the Wasserstein geometry. Notably, WFM is the first algorithm capable of generating distributions in high dimensions, whether represented analytically (as Gaussians) or empirically (as point-clouds). Our theoretical analysis establishes that Wasserstein geodesics constitute proper conditional flows over the space of distributions, making for a valid FM objective. Our algorithm leverages optimal transport theory and the attention mechanism, demonstrating versatility across computational regimes: exploiting closed-form optimal transport paths for Gaussian families, while using entropic estimates on point-clouds for general distributions. WFM successfully generates both 2D & 3D shapes and high-dimensional cellular microenvironments from spatial transcriptomics data. Code is available at WassersteinFlowMatching.
APA
Haviv, D., Pooladian, A., Pe’Er, D. & Amos, B.. (2025). Wasserstein Flow Matching: Generative Modeling Over Families of Distributions. Proceedings of the 42nd International Conference on Machine Learning, in Proceedings of Machine Learning Research 267:22238-22258 Available from https://proceedings.mlr.press/v267/haviv25a.html.

Related Material