Wasserstein Flow Matching: Generative modeling over families of distributions

生成模型通常关注于通过学习（即回归）简单的概率流，将单一的源分布传输到单一的目标分布。然而，在现代数据驱动领域，如计算机图形学和单细胞基因组学中，数据集中的样本（例如点云）本身可以被视为分布（例如离散测度）。在这种情况下，标准的生成模型范式——流匹配，会忽略样本的相关几何结构。为了解决这一问题，我们提出了*Wasserstein流匹配*（WFM），该方法通过利用Wasserstein几何的黎曼性质，将流匹配适当地扩展到分布族上。我们的算法利用了（熵正则化）最优传输理论和计算方面的最新进展，以及神经网络架构中的注意力机制。我们提出了两个新颖的算法贡献。首先，我们展示了如何对高斯分布进行生成建模，从单细胞基因组学数据生成颗粒状细胞状态的表示。其次，我们展示了WFM可以学习高维和可变大小点云之间的流，并从空间转录组学数据集中合成细胞微环境。代码可在 [WassersteinFlowMatching](https://github.com/DoronHav/WassersteinFlowMatching) 获取。

该论文旨在解决标准生成模型在处理现代数据驱动领域（如计算机图形学和单细胞基因组学）中的点云样本时，忽略了样本本身作为分布所具有的几何特性的问题。这是一个相对较新的问题，因为随着数据类型的多样化，如何有效利用样本的内在结构成为了一个重要的研究方向。

论文提出了一种名为Wasserstein流匹配（WFM）的新方法，通过利用Wasserstein几何的Riemannian性质，将流匹配提升到分布族上。这种方法不仅考虑了样本之间的距离，还考虑了样本分布的几何结构，从而更好地捕捉数据的本质特征。与传统的生成模型相比，WFM能够更准确地建模复杂的数据分布。

论文的主要亮点包括：1) 在高维和可变大小的点云之间学习流，能够合成细胞微环境；2) 通过生成高斯分布来表示单细胞基因组数据中的细胞状态；3) 实验设计合理，使用了单细胞基因组学和空间转录组学数据集进行验证；4) 开源代码已发布，便于其他研究人员复现和扩展研究。未来可以进一步探索WFM在更多类型数据上的应用，以及与其他生成模型的结合。

近期在这个领域的一些相关研究包括：1) 'Sinkhorn Divergences for Unsupervised Point Cloud Registration'，提出了利用Sinkhorn散度进行无监督点云配准的方法；2) 'Optimal Transport for Domain Adaptation'，探讨了最优传输在域适应中的应用；3) 'Generative Adversarial Networks for Point Clouds'，研究了生成对抗网络在点云生成中的应用。这些研究都为WFM的发展提供了理论基础和技术支持。