Generative modeling typically concerns the transport of a single source distribution to a single target distribution by learning (i.e., regressing onto) simple probability flows. However, in modern data-driven fields such as computer graphics and single-cell genomics, samples (say, point-clouds) from datasets can themselves be viewed as distributions (as, say, discrete measures). In these settings, the standard generative modeling paradigm of flow matching would ignore the relevant geometry of the samples. To remedy this, we propose \emph{Wasserstein flow matching} (WFM), which appropriately lifts flow matching onto families of distributions by appealing to the Riemannian nature of the Wasserstein geometry. Our algorithm leverages theoretical and computational advances in (entropic) optimal transport, as well as the attention mechanism in our neural network architecture. We present two novel algorithmic contributions. First, we demonstrate how to perform generative modeling over Gaussian distributions, where we generate representations of granular cell states from single-cell genomics data. Secondly, we show that WFM can learn flows between high-dimensional and variable sized point-clouds and synthesize cellular microenvironments from spatial transcriptomics datasets. Code is available at [WassersteinFlowMatching](https://github.com/DoronHav/WassersteinFlowMatching).

本研究解决了传统生成建模在处理现代数据驱动领域（如计算机图形学和单细胞基因组学）中的局限性，特别是在样本视为分布时忽略其几何特征的问题。通过提出Wasserstein流匹配（WFM），本论文创新性地将流匹配扩展至分布族，并利用Wasserstein几何的Riemannian特性，显著提升了生成建模的能力，尤其是在生成高维点云和合成细胞微环境方面的应用。

Wasserstein流匹配：分布族上的生成建模