Wasserstein Flow Matching: Generative modeling over families of distributions

## Abstract

Abstract 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 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 .

## 1 Introduction

Today’s abundance of data and scalability of training massive neural
networks has made it possible to generate hyper-realistic images on
the basis of training examples (OpenAI,, ), as well as video and
audio clips (Vyas et al.,, ; Xing et al.,, ), and, of
course, text (Bubeck et al.,, ).
All of these are instances of generative modeling:
given access to finitely many samples from a distribution,
devise a scheme which generates new samples from the same distribution.
Generative modeling has also been revolutionary in the biomedical sciences,
for drug design (Jumper et al.,, ) and single-cell genomics (Lopez et al.,, ).
Nearly all frameworks exploit the notion that datasets (of, say, genomic profiles of cells,
images, videos, or corpora of text documents) are instantiations of
probability measures, and the task is to transform a point sampled
from random noise to generate a data point that obeys the
distribution of interest.

Among the zoo of available generative models, one approach noted for its flexibility and simplicity is Flow Matching (FM) (Albergo and Vanden-Eijnden,, ; Lipman et al.,, ; Liu et al.,, ).
For a fixed target probability measure, FM learns an implicitly
defined vector field that can transform a source measure (e.g., the
standard Gaussian) to the target. Unlike discrete time and
probabilistic generative models (such as Denoising Diffusion Models by
Song et al., ()), FM learns a deterministic, continuous
normalizing flow by regressing onto a simple conditional probability
flow. This approach, while originally designed for Euclidean domains,
can be readily adopted to Riemannian geometries
(Chen and Lipman,, ).
Riemannian flow matching is widely used for generating samples over
geometries such as spheres, tori, translation/rotation groups,
simplices, triangular meshes, mazes, and molecular positions and structures.

The *Wasserstein geometry*, a canonical geometry over distributions, does not easily fit into any of
these existing frameworks and has not been successfully
adapted for flow matching.
This geometry is useful, for example, in computational
graphics where
collections of 3D shapes are represented
as empirical distributions (point-clouds).
Likewise, recent developments in single-cell genomics analysis have
demonstrated that gene-expression profiles from groups of cells
aggregated via their mean and covariance can capture cellular
microenvironments or highlight fine-grain clusters
(; Persad et al.,, ). For both point-cloud
and Gaussian settings, it is natural to search for a unified
generative model that respects the underlying geometry of the data,
namely, treating each sample as itself a probability distribution.

##### Contributions.

We introduce *Wasserstein Flow Matching* (WFM), a principled extension of the FM framework lifted to the space of probability distributions. As illustrated in Figure 1, a single point in our source and target datasets is itself a distribution (e.g., a single discrete measure or a single Gaussian), and our aim is to learn vector fields acting on the space of probability distributions and match the optimal transport map, which is the geodesic in Wasserstein space. WFM is an instantiation of Riemannian FM  (Chen and Lipman,, ), where we train a neural model to learn a continuous normalizing flow (CNF) between distributions over distributions.

Figure: Figure 1: *Left:* Table contrasting FM methods over $\mathbb{R}^{d}$, general manifolds $\mathcal{M}$, categorical and Dirichlet distributions on the $d$-simplex $\Delta_{d}$, and finally, our approach, FM problems defined over $\mathcal{P}(\mathbb{R}^{d})$. *Right:* WFM overview, which learns flows over distributions over distributions.
Refer to caption: x1.png

We demonstrate the effectiveness of our approach for generative modeling between distributions over Gaussian distributions and distributions over point-clouds.
The former task is motivated by recent directions in single-cell and spatial transcriptomics (; Persad et al.,, ), where we consider matching problems over the *Bures–Wasserstein space* (BW), the Gaussian submanifold of the Wasserstein space. In this case, we show that WFM can be further modified, resulting in the Bures–Wasserstein FM (BW-FM) algorithm. We validate BW-FM on a variety of Gaussian-based datasets, where we observe that samples generated by our algorithm are significantly more robust than naïve approaches which do not fully exploit the underlying geometry of the data. In turn, we present a generative model for cell states and niches from single-cell genomics data.

Point-cloud generation is made possible by two distinct, yet crucial, algorithmic primitives: (1) incorporating transformers in our neural network architecture (Vaswani,, ; Lee et al.,, ), and (2) recent algorithmic advances in entropic optimal transport (Pooladian and Niles-Weed,, ). Indeed, our WFM algorithm performs generative modeling in the Wasserstein space, where geodesics are given by pushforwards of optimal transport (OT) maps; see Section 2.3 for more information. Both the transformer architecture and entropic optimal transport are crucial to approximating the OT map between independent point-clouds. Indeed, the permutation equivariance of attention makes the transformer a natural basis for our model, inherently modeling the equivariance feature of the Wasserstein geometry while maintaining scalability in high-dimensions.

For datasets of 3D point-clouds with uniform sizes, the performance of WFM is comparable to other current generative models. However, due to their particular training paradigms (namely the voxelization of 3D spaces), contemporary approaches cannot scale to high-dimensional point-clouds and fail on datasets with variable sized examples. Conversely, WFM succeeds in the high-dimensional and inhomogeneous settings, unlocking generative modeling to new, previously uncharted domains such as synthesizing niches from spatial genomics data. The ability to model tissue biology in this generative manner could enhance our understanding of how environment is associated with cell state. In the context of many diseases, most notably cancer and its tumor-immune microenvironment, these insights are critical for developing novel therapeutics (Binnewies et al.,, ).

## 2 Background and related work

We let $\mathcal{P}_{2}(\mathbb{R}^{d})$ denote the set of probability distributions over $\mathbb{R}^{d}$ with finite second moment, and write $\mathcal{P}_{2,\text{ac}}(\mathbb{R}^{d})$ to be those with densities. For a probability measure $\mu$ and (vector-valued) function $f$, we interchangeably write $\int\|f(x)\|^{2}\,\mathrm{d}\mu(x)$ and $\|f\|^{2}_{L^{2}(\mu)}$. Let $\mathcal{M}$ be a Riemannian manifold, with $\mathcal{P}(\mathcal{M})$ defining the space of probability measures over said manifold. For $x\in\mathcal{M}$, we write $\mathcal{T}_{x}\mathcal{M}$ to mean the tangent space of the manifold at $x$, and write the metric on the tangent space (at $x$) as $g(x)$.
For $x_{0}\in\mathcal{M}$ with initial velocity $v\in\mathcal{T}_{x_{0}}\mathcal{M}$, the terminal location of the resulting geodesic is expressed as the output of the exponential map $v\mapsto\exp_{x_{0}}(v)\in\mathcal{M}$. Similarly, for an initial point $x_{0}$ and terminal location $x_{1}$, the logarithmic map defines the tangent vector, denoted $x_{1}\mapsto\log_{x_{0}}(x_{1})$, such that $\exp_{x_{0}}(\log_{x_{0}}(x_{1}))=x_{1}$. The set of symmetric matrices (resp. positive definite matrices) over $\mathbb{R}^{d}$ are denoted by $\mathbb{S}^{d}$ (resp. $\mathbb{S}^{d}_{++}$).

### 2.1 Riemannian flow matching

We first briefly discuss the Riemannian flow matching (RFM) framework of Chen and Lipman, (). Let $\mathfrak{p}_{0}$ be the source distribution and $\mathfrak{p}_{1}$ be the target distribution over a Riemannian manifold $\mathcal{M}$, and let $(\gamma_{t})_{t\in[0,1]}$ be a curve of probability measures satisfying $\gamma_{0}=\mathfrak{p}_{0}$ and $\gamma_{1}=\mathfrak{p}_{1}$. Letting $(w_{t})_{t\in[0,1]}$ denote a family of vector fields, we say that the pair $(\gamma_{t},w_{t})_{t\in[0,1]}$ satisfy the *continuity equation* with respect to the metric $g$, abbreviated to $(\gamma_{t},w_{t})\in\mathfrak{C}_{g}$ if

$$ $\displaystyle\partial_{t}\gamma_{t}+\nabla_{g}\!\cdot\!(\gamma_{t}w_{t})=0\,,$ (1) $$

where $\nabla_{\!g}\cdot$ is the Riemannian divergence operator.

The goal of RFM is to regress a parameterized vector field (e.g., a neural network), written $f_{\theta}(x,t)\in\mathcal{T}_{x}\mathcal{M}$ for $t\in[0,1]$, onto the family $w_{t}$ by minimizing

$$ $\displaystyle\min_{\theta}\int_{0}^{1}\!\!\int\|f_{\theta}(z_{t},t)-w_{t}(z_{t })\|_{g(z_{t})}^{2}\,\mathrm{d}\gamma_{t}(z_{t})\,\mathrm{d}t\,,$ $$

assuming access to a pair $(\gamma_{t},w_{t})_{t\in[0,1]}$ that satisfies equation 1, which is not possible in many scenarios. Borrowing insights from recent work (e.g., Albergo and Vanden-Eijnden, (); Lipman et al., (); Liu et al., ()), the authors construct a simple vector field that satisfies the continuity equation, resulting in the tractable objective

$$ $\displaystyle\min_{\theta}\int_{0}^{1}\!\!\iint\|f_{\theta}(x_{t},t)-\dot{x}_{ t}\|_{g(x_{t})}^{2}\,\mathrm{d}\mathfrak{p}_{0}(x)\,\mathrm{d}\mathfrak{p}_{1} (y)\,\mathrm{d}t\,,$ (2) $$

where, for example, $x_{t}=\exp_{x}((1-t)\log_{x}(y))\in\mathcal{M}$, and $\dot{x}_{t}\in\mathcal{T}_{x_{t}}\mathcal{M}$. For complete discussions and proofs, see (Chen and Lipman,, ). Once $f_{\theta}$ is appropriately fit using equation 2, we can generate new samples from $\mathfrak{p}_{1}$:
start by sampling $X_{0}\sim\mathfrak{p}_{0}$, then follow $\dot{X}_{t}=f_{\theta}(X_{t},t)$ numerically by discretizing the dynamics given by the exponential map, resulting in $X_{1}\sim\mathfrak{p}_{1}$. We emphasize that the dynamics are only simulated at inference time and not when training $f_{\theta}$, which is commonly known as a *simulation-free* training paradigm.

### 2.2 Related work

##### Generative models for point-clouds.

Paralleling the progress in generative models for natural images, the field of point-cloud generation is rapidly expanding. Many different models have been used from this task, namely generative-adversarial-nets (Achlioptas et al.,, ), variational autoencoders (Gadelha et al.,, ), normalizing flows (Yang et al.,, ; Kim et al.,, ; Klokov et al.,, ), diffusion (Zhou et al.,, ; Cai et al.,, ) and even euclidean FM (Wu et al.,, ). Thus far, these approaches are limited to uniformly sized point-clouds in 2D & 3D, and fail on high-dimensional spaces which cannot be voxelized.

##### Generative models over families of distributions.

Our work is not the first to instantiate Riemannian FM with a manifold of probability measures. Two notable works are Fisher FM (Davis et al.,, ) and
Categorical FM (Cheng et al.,, ), which consider the FM algorithm with respect to the Fisher–Rao geometry Amari, (); Nielsen, () over the $d$-dimensional simplex $\Delta_{d}$. The work of Stark et al., () is similar in spirit, where they focus on the Dirichlet distribution for generation of discrete data. Another related work is that of Atanackovic et al., (), called Meta FM. Their approach requires pairs of distributions which are already coupled, with the goal of solving FM between a distribution over pairs. In contrast, we emphasize that our proposed Wasserstein FM applies between two separate *uncoupled* distributions over distributions.

### 2.3 Wasserstein geometry

The (squared) *$2$-Wasserstein distance* between two probability measures $\mu,\nu\in\mathcal{P}_{2,\text{ac}}(\mathbb{R}^{d})$ is given by the non-convex optimization problem over vector-valued maps $T:\mathbb{R}^{d}\to\mathbb{R}^{d}$

$$ $\displaystyle W_{2}^{2}(\mu,\nu)\coloneqq\min_{T:T_{\sharp}\mu=\nu}\|{\rm{id}} -T\|_{L^{2}(\mu)}^{2}\,,$ (3) $$

where the pushforward constraint, written $T_{\sharp}\mu=\nu$, means that, for $X\sim\mu$, the image follows $T(X)\sim\nu$. The minimizer to equation 3 is called the *optimal transport (OT) map*, denoted $T^{\mu\to\nu}_{\star}$ (we abbreviate this to $T_{\star}$ when it is clear from context).
The existence and uniqueness of the optimal transport map under the stated regularity conditions is due to Brenier, ().

The *Wasserstein space* is the space of probability densities with finite second moment endowed with the Wasserstein distance; this space is known to be a metric space (Villani,, ). Following the celebrated work of Otto, (), the Wasserstein space can be formally (meaning, non-rigorously) viewed as a Riemannian manifold, whose properties we now describe in brief; see e.g., Ambrosio et al., () for a rigorous treatment.

Following the definition in
Theorem 8.5.1 from Ambrosio et al., (), the tangent space at a point $\mu\in\mathcal{P}_{2,\text{ac}}(\mathbb{R}^{d})$ consists of all possible tangent vectors that emanate from $\mu$, written formally as

$$ $\displaystyle\mathcal{T}_{\mu}\mathcal{P}_{2,\text{ac}}(\mathbb{R}^{d}) \coloneqq\overline{\{\lambda(T_{\star}^{\mu\to\nu}-\mathrm{id})\ :\ \lambda>0, \nu\in\mathcal{P}_{2}(\mathbb{R}^{d})\}}^{L^{2}(\mu)}\,,$ $$

where the overline denotes the closure of the set in $L^{2}(\mu)$, and the norm on the tangent space is also $L^{2}(\mu)$. The exponential and logarithmic maps read

$$ $\displaystyle v\mapsto\exp_{\mu}(v)\coloneqq(\mathrm{id}+v)_{\sharp}\mu\,, \quad\nu\mapsto\log_{\mu}(\nu)\coloneqq T_{\star}^{\mu\to\nu}-\mathrm{id}\,,$ $$

where id is the identity map. Consequently, the (constant-speed) geodesic, or *McCann interpolation*, between two measures $\mu$ and $\nu$ is given by the curve $(\mu_{t})_{t\in[0,1]}$ where

$$ $\displaystyle\mu_{t}\coloneqq(T_{t}^{\mu\to\nu})_{\sharp}\mu\coloneqq((1-t) \text{id}+tT_{\star}^{\mu\to\nu})_{\sharp}\mu\equiv\exp_{\mu}((1-t)\log_{\mu}( \nu))\,,$ (4) $$

where the last expression writes the pushforward in terms of the exponential and logarithmic maps. Equivalently, at the level of the random variables, one can write $X_{t}=(1-t)X_{0}+tT_{\star}^{\mu\to\nu}(X_{0})$, where $X_{0}\sim\mu$ and $X_{t}\sim\mu_{t}$ for any $t\in[0,1]$. Combined with $(v_{t})_{t\in[0,1]}$ a suitable family of vector fields, the McCann interpolation satisfies the continuity equation equation 1 over $\mathbb{R}^{d}$, re-written as

$$ $\displaystyle\partial_{t}\mu_{t}+\nabla\cdot(\mu_{t}v_{t})=0\,,\quad\text{s.t. }\quad\mu_{0}=\mu\,,\mu_{1}=\nu\,,$ (5) $$

where the divergence operator is the usual Euclidean one over $\mathbb{R}^{d}$, thus we write $(\mu_{t},v_{t})\in\mathfrak{C}$. The link between the constant speed geodesics and the $2$-Wasserstein distance can be viewed from the celebrated Benamou–Brenier formulation of optimal transport (Benamou and Brenier,, ):

$$ $\displaystyle W_{2}^{2}(\mu,\nu)=\inf_{(\mu_{t},v_{t})\in\mathfrak{C}}\int_{0} ^{1}\|v_{t}\|^{2}_{L^{2}(\mu_{t})}\,\mathrm{d}t\,.$ (6) $$

The optimal curve of measures is given by the constant-speed geodesics described above, and the optimal velocity field is given by

$$ $\displaystyle v_{t}=(T_{\star}^{\mu\to\nu}-\text{id})\circ(T_{t}^{\mu\to\nu})^ {-1}\,.$ (7) $$

The vector field equation 7 should be interpreted as the time-derivative of the McCann interpolation:

$$ $\displaystyle\dot{X}_{t}=(T_{\star}^{\mu\to\nu}-\text{id})(X_{0})=(T_{\star}^{ \mu\to\nu}-\text{id})\circ(T_{t}^{\mu\to\nu})^{-1}(X_{t})\,,\qquad X_{0}\sim \mu\,.$ $$

#### 2.3.1 Bures–Wasserstein (BW) space

A known special case of the Wasserstein space is the *Bures–Wasserstein* space, which consists of the submanifold of non-degenerate Gaussians parameterized by means and covariances $\{(m,\Sigma):m\in\mathbb{R}^{d},\Sigma\in\mathbb{S}^{d}_{++}\}$, endowed with the Wasserstein metric. We provide a brief exposition on the geometry of the Bures–Wasserstein space and refer the interested reader to Lambert et al., () for detailed calculations and explanations, as we follow their notation conventions.

The OT map between $\mu=\mathcal{N}(m_{\mu},\Sigma_{\mu})$ and $\nu=\mathcal{N}(m_{\nu},\Sigma_{\nu})$ has a closed-form (Gelbrich,, ):

$$ $\displaystyle T_{\star}(x)\coloneqq m_{\nu}+C^{\mu\to\nu}(x-m_{\mu})\coloneqq b +\Sigma_{\mu}^{-\tfrac{1}{2}}(\Sigma_{\mu}^{\tfrac{1}{2}}\Sigma_{\nu}\Sigma_{ \mu}^{\tfrac{1}{2}})^{\tfrac{1}{2}}\Sigma_{\mu}^{-\tfrac{1}{2}}(x-m_{\mu})\,.$ $$

As this map is affine, it is clear that the McCann interpolation between two Gaussians is always Gaussian (indeed, Gaussians undergoing affine transformations remain Gaussian).
More generally, we have the succinct representation of the tangent space at a point in the Bures–Wasserstein space

$$ $\displaystyle\mathcal{T}_{\mu}{\rm{BW}}(\mathbb{R}^{d})\coloneqq\{a+S({\rm{id} }-m_{\mu})\,:\,a\in\mathbb{R}^{d},S\in\mathbb{S}^{d}\}\,,$ $$

and the exponential and logarithmic maps between two non-degenerate Gaussians are

$$ $\displaystyle(a,S)\mapsto\exp_{\mu}((a,S))\coloneqq\mathcal{N}(m_{\mu}+a,(S+I) \Sigma_{\mu}(S+I))\,,$ $\displaystyle\nu\mapsto\log_{\mu}(\nu)\coloneqq(m_{\nu}-m_{\mu},\Sigma_{\mu}^{ -\tfrac{1}{2}}(\Sigma_{\mu}^{\tfrac{1}{2}}\Sigma_{\nu}\Sigma_{\mu}^{\tfrac{1}{ 2}})^{\tfrac{1}{2}}\Sigma_{\mu}^{-\tfrac{1}{2}}-I)\,,$ $$

where the exponential map requires $S\succ-I$. We also mention that the norm on the tangent space at $\mu$ in the Bures–Wasserstein space can be written as

$$ $\displaystyle\|(a,S)\|^{2}_{{\rm{BW}}(\mu)}\coloneqq\|a-m_{\mu}\|^{2}+{\rm{Tr} }(S^{2}\Sigma_{\mu})$ $$

With the above, it is easy to compute the closed-form solutions for the mean and covariance of the McCann interpolation $\mu_{t}=(T_{t})_{\sharp}\mu=\mathcal{N}(m_{t},\Sigma_{t})$, given by

$$ $\displaystyle m_{t}\coloneqq(1-t)a+tb\,,\quad\Sigma_{t}\coloneqq T_{t}AT_{t} \coloneqq((1-t)I+tC^{A\to B})A((1-t)I+tC^{A\to B})\,.$ (8) $$

We can relate the Euclidean and Riemannian time-derivatives of $\Sigma_{t}$ through the following manipulation (the latter of which respects the exponential and logarithmic maps above):

$$ $\displaystyle\dot{\Sigma}_{t}^{\rm E}=\dot{T}_{t}AT_{t}+T_{t}A\dot{T}_{t}=\dot {T}_{t}(T_{t})^{-1}T_{t}AT_{t}+T_{t}AT_{t}(T_{t})^{-1}\dot{T}_{t}=\dot{\Sigma} _{t}^{\rm{BW}}\Sigma_{t}+\Sigma_{t}\dot{\Sigma}_{t}^{\rm{BW}}\,.$ $$

To this end, we can draw parallels to equation 7 by writing

$$ $\displaystyle\dot{m}_{t}=b-a\,,\quad\dot{\Sigma}_{t}^{\rm{BW}}=(C^{A\to B}-I)( (1-t)I+tC^{A\to B})^{-1}\,.$ (9) $$

## 3 Flow matching over the Wasserstein space

### 3.1 Training

Let $\mathfrak{p}_{0}$ and $\mathfrak{p}_{1}$ denote probability measures over the Wasserstein space.(^1^11This implies that $\mu\sim\mathfrak{p}_{0}$ is itself a distribution (e.g., a Gaussian or a point-cloud), *not* a random variable.) Our goal is to learn a vector field that transports the family of measures $\mathfrak{p}_{0}$ to the family $\mathfrak{p}_{1}$. To accomplish this, we pass in the McCann interpolation $\mu_{t}$ and optimal velocity field $v_{t}$ in the Riemannian FM objective equation 2, resulting in our Wasserstein FM (WFM) objective:

$$ $\displaystyle\min_{\theta}\int_{0}^{1}\!\!\iint\|f_{\theta}^{\texttt{geo}}(\mu _{t},t)-v_{t}\|_{L^{2}(\mu_{t})}^{2}\,\mathrm{d}\mathfrak{p}_{0}(\mu)\,\mathrm {d}\mathfrak{p}_{1}(\nu)\,\mathrm{d}t\,.$ (10) $$

Figure: Figure 2: When the number of training examples is too few, all methods collapse on the training data, though our Riemannian instantiation of BW-FM captures the covariances perfectly. In the presence of sufficiently many samples, all methods generate Gaussians along the whole spiral, and our Riemannian BW-FM algorithm produces the most consistent samples. Other methods produce Gaussians with *degenerate* covariance, highlighting their inability to effectively capture the geometry of the data.
Refer to caption: extracted/5971869/Figures/bw_spirals.png

As mentioned in the introduction, our two use-cases of interest are flow matching over (1) families of point-clouds, and (2) families of Gaussian distributions. While the theory outlined in Section 2.3 explicitly requires continuous distributions to ensure all objects are well-defined, the approximation of measures by point-clouds is reasonable for our applications and can be made computationally efficient courtesy of existing open-source packages (Flamary et al.,, ; Cuturi et al.,, ). In the case of Gaussian measures, the theory as described in Section 2.3.1 holds in full force. Our training algorithm is described in Algorithm 1, and Appendix D contains precise details regarding our neural network parameterization.

Figure: Algorithm 1 Wasserstein FM Training^d

Finally, we mention that both frameworks can be modified by training via the *multisample FM* algorithm (; Tong et al.,, ). In brief, the idea is to augment the training regime by pairing the source and target minibatch samples according to some prescribed matching rule (instead of independent draws from both $\mathfrak{p}_{0}$ and $\mathfrak{p}_{1}$). We employ this augmentation during training, which we detail in Appendix B.

#### 3.1.1 WFM over the Bures–Wasserstein space

First suppose $\mathfrak{p}_{0}$ and $\mathfrak{p}_{1}$ are distributions over Gaussians, meaning that a batch of samples drawn from $\mathfrak{p}_{0}$ and $\mathfrak{p}_{1}$ consists of a batch of mean-covariance pairs. Here, the dynamics are straightforward: the interpolant is the McCann interpolation (recall equation 8), and the velocity field over the Bures–Wasserstein manifold is also known (see equation 9). Since $\mu_{t}$ is parameterized by $(m_{t},\Sigma_{t})$, the neural network is parameterized as $f_{\theta}^{\text{BW}}:\mathbb{R}^{d}\times\mathbb{S}^{d}_{++}\to\mathbb{R}^{d
}\times\mathbb{S}^{d}$, and the norm on the tangent space simplifies the computations considerably. Our final training objective becomes

$$ $\displaystyle\min_{\theta}\int_{0}^{1}\!\!\iint\|f_{\theta}^{\text{BW}}((m_{t} ,\Sigma_{t}),t)-(\dot{m}_{t},\dot{\Sigma}_{t}^{\rm{BW}})\|_{{\rm{BW}}(\mu_{t}) }^{2}\,\mathrm{d}\mathfrak{p}_{0}(\mu)\,\mathrm{d}\mathfrak{p}_{1}(\nu)\, \mathrm{d}t\,.$ (11) $$

Figure: Figure 3: Environment of cells in spatial transcriptomics data can be effectively captured via the mean and covariance of gene expression from their niche (). Applying BW to realize environments of microglia from MERFISH data (Zhang et al.,, ), the composition of generated niches matches the real data along inferred cortical depth; see also Figure S2.
Refer to caption: extracted/5971869/Figures/bw_glia.png

#### 3.1.2 WFM over distributions of point-clouds

In the case of point-clouds, we lose closed-form interpolations. However, we can hope to proceed so long as we have an *approximation* of the optimal transport map between the point-clouds, written $\hat{T}$. There are many works on the approximation of these maps on the basis of samples; see Hütter and Rigollet, (); Divol et al., (); Manole et al., (); Pooladian and Niles-Weed, (). We stress that the goal of our approach is to have a methodology that holds for families of point-clouds of non-uniform size.

We consider two approximations of optimal transport maps, both of which are based on entropic optimal transport Cuturi, (), and are computationally efficient on GPUs due to Sinkhorn’s algorithm (Sinkhorn,, ). One approach is to round the optimal coupling to a permutation and perform the resulting interpolation. Another approach, which allows for inhomogeneous pairs of points, is to approximate $T_{\star}^{\mu\to\nu}$ using the entropic map (Pooladian and Niles-Weed,, ); we provide extensive background on these objects in Appendix A. To this end, let $\bm{X}$ (resp. $\bm{Y}$) represent the locations of the point-cloud $\mu\sim\mathfrak{p}_{0}$ (resp. $\nu\sim\mathfrak{p}_{1}$), and let $\hat{T}^{\mu\to\nu}$ denote the approximation of the optimal transport map.
The objective equation 10 can be approximated by

$$ $\displaystyle\min_{\theta}\int_{0}^{1}\!\!\iint\sum_{i}\|[f_{\theta}^{\text{PC }}(\hat{\bm{X}}_{t},t)]_{i}-[(\hat{T}^{\mu\to\nu}(\bm{X})-\bm{X})]_{i}\|_{2}^{ 2}\,\mathrm{d}\mathfrak{p}_{0}(\mu)\,\mathrm{d}\mathfrak{p}_{1}(\nu)\,\mathrm{ d}t\,,$ (12) $$

where $\hat{\bm{X}}_{t}=(1-t)\bm{X}+t\hat{T}^{\mu\to\nu}(\bm{X})$. Here, we stress that $\hat{\bm{X}}_{t}$ plays the role of a discretized McCann interpolation $\mu_{t}$.

We parameterize $f_{\theta}^{PC}$ with a transformer
and Appendix D provides further details.
As both the OT map and self-attention are permutation equivariant, transformers are an organic backbone for OT based models (). Moreover, unlike other point-cloud neural models such as PVCNN (Liu et al.,, ), transformers do not rely on voxelization and are not hindered by the curse-of-dimensionality.

### 3.2 Generation

Once $f_{\theta}^{\texttt{geo}}$ is trained, we can generate new samples in a simulation-free manner as in Riemannian FM. For the Bures–Wasserstein space, we appeal to the exponential and logarithmic maps; Section 2.3.1 and Algorithm 2. We emphasize that the appropriate Riemannian updates are crucial to obtain non-degenerate final samples. To generate point-clouds, we perform a standard Euler discretization of the learned flow over the whole collection of point-clouds; see Algorithm 3.

## 4 Results

### 4.1 Flow matching between families of Gaussians

We first demonstrate our flow matching framework between measures of Gaussian distributions on a synthetic and real datasets. For comparable baselines in each scenario, we construct two simpler flow matching approaches for Gaussian generation: (1)
*Frobenius FM*, which concatenates the mean and covariances values, and trains on $\dot{\Sigma}_{t}^{{\rm{E}}}$ with respect to the squared-Frobenius norm, and (2) BW-FM (Euclidean), which tries to match $\dot{\Sigma}^{{\rm E}}_{t}$ but still under the BW geometry (and not a concatenated geometry).

In all cases, we assess the quality of the learned flows by computing the the minimum distance between each generated Gaussian and the dataset using the (squared) $2$-Wasserstein distance. Notably, the flows generated by Frobenius FM and BW-FM (Euclidean) do not strictly adhere to the geometry of the Bures-Wasserstein manifold, requiring synthesized covariance matrices to be artificially projected onto the space of positive semi-definite (PSD) matrices via eigenvalue truncation. In contrast, the Riemannian BW-FM algorithm consistently produces valid and accurate results across all dimensions and datasets.

Figure: Figure 4: *Left.* Synthesized samples from WFM trained on the cars, planes or chairs datasets. *Right.* Examples generated conditionally from the *same initial noise* via a WFM model trained on the complete 40-class ModelNet dataset.
Refer to caption: x2.png

#### 4.1.1 Toy Datasets

As a first test, we design a dataset of Gaussians centered on a spiral (Figure 2). When there were only few samples, only Riemannian BW-FM reconstructed the data, despite other benchmark methods following identical training regimes. On the complete 128-sample dataset, BW-FM not only reconstructs the training data, but generalizes and synthesizes novel Gaussians whose means lie on the spiral with the correct covariance profile. BW-FM shares this generalization feature with standard FM, and is able to learn the structure underlying the measure from the training data.

**Table 1: Average minimum $W^{2}_{2}$ distance between each generated Gaussian and the reference datasets. We benchmark our BW-FM algorithm with Riemannian updates against (1) BW-FM with Euclidean updates, and (2) Frobenius FM; see Figure 2. Despite identical training schemes, the BW-FM (R) vastly outperforms other approaches on both synthetic and real data.**
|  | BW-FM (R) | BW-FM (E) | Frobenius FM |
| --- | --- | --- | --- |
| Spiral - 16 (2D) | $\mathbf{2.98\cdot 10^{-4}}$ | $4.00\cdot 10^{-4}$ | $1.03\cdot 10^{-3}$ |
| Spirals - 128 (2D) | $\mathbf{1.28\cdot 10^{-3}}$ | $1.70\cdot 10^{-3}$ | $2.69\cdot 10^{-3}$ |
| Two Moons (2D) | $\mathbf{1.84\cdot 10^{-4}}$ | $8.96\cdot 10^{-4}$ | $1.30\cdot 10^{-3}$ |
| Sphere (3D) | $\mathbf{6.65\cdot 10^{-4}}$ | $2.14\cdot 10^{-3}$ | $2.25\cdot 10^{-3}$ |
| Cities (2D) | $\mathbf{1.88\cdot 10^{-4}}$ | $7.26\cdot 10^{-3}$ | $1.75\cdot 10^{-3}$ |
| ECG (15D) | $\mathbf{9.24\cdot 10^{-2}}$ | $3.26\cdot 10^{-1}$ | $3.98\cdot 10^{-1}$ |
| MERFISH (16D) | $\mathbf{1.90}$ | $1.98$ | $2.06$ |
| scRNA-seq (32D) | $\mathbf{1.31}$ | $2.74$ | $3.21$ |

#### 4.1.2 Single-Cell Genomics

Spatial transcriptomics are a set of techniques which build on single-cell genomics and preserve physical information of cells’ location in tissues, while assaying their gene expression. demonstrated that a cell’s microenvironment can be effectively characterized using the mean and covariance of the surrounding cells gene expression. This statistical representation efficiently captures the key features of cellular neighborhoods and transforms each spatial transcriptomics dataset into a measure within the Bures–Wasserstein space, highlighting the value of generative modeling in this context.

In the motor cortex, excitatory neurons form phenotypically distinct and highly specialized cortical layer (Zeng and Sanes,, ). From the 254 gene MERFISH atlas (Zhang et al.,, ), we compute the mean and covariance of the top $16$ principal components of gene expression from all cells within an $80$ micron radius around each microglia. Despite the dimensionality of this data, BW-FM synthesizes Gaussians which are highly congruent with the real data (Figure 3 and Figure S2).

Another common instance of BW manifolds arising in single-cell genomics is through aggregating cells into common states. These clusters can be summarized by their mean gene expression and its covariance. On a scRNA-seq atlas elucidating human immune response to COVID (Stephenson et al.,, ), we combine cells into MetaCells (Persad et al.,, ), and quantify the gene expression mean and covariance for each. Here too we apply BW-FM conditioned on cell-state, which encompass the heterogeneity of immune profiles appearing as response to COVID infection. Despite the plurality of labels, BW-FM can still synthesize appropriate examples for each condition (see Figure S1).

### 4.2 Flow matching between families of point-clouds

When the two measures are point-clouds, we turn to entropic optimal transport to estimate the McCann interpolation (Cuturi,, ). When the dataset consists of source and target point-clouds of the same size, we use a GPU-efficient rounding scheme to approximate the OT map using the entropic OT coupling; see Appendix A. When the support size of the source and target point-clouds vary, it is worth mentioning that an OT map may not even exist. Nevertheless, we approximate these curves using the entropic transport map of Pooladian and Niles-Weed, (). Together, these two approaches method offer a computationally feasible solution while maintaining accuracy in transport map estimation.

**Table 2: 1-Nearest-Neighbour Accuracy for high-dimensional or variable size point-clouds. WFM employs a transformer backbone and relies on the efficient computation of the entropic transport map, allowing it to scale to arbitrary dimensions and learn flows between point-cloud of variable sizes, key features all previous point-cloud generation approaches lack.**
|  | MNIST ($4$) | Letters ($A$) | seqFISH | XENIUM |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | CD $\downarrow$ | EMD $\downarrow$ | CD $\downarrow$ | EMD $\downarrow$ | CD $\downarrow$ | EMD $\downarrow$ | CD $\downarrow$ | EMD $\downarrow$ |
| Current methods | NA | NA | NA | NA | NA | NA | NA | NA |
| WFM (ours) | 63.34 | 59.97 | 62.12 | 58.68 | 61.79 | 64.34 | 60.69 | 64.20 |

We compare WFM to many other point-cloud generation algorithms. Following in their footsteps, we measure generation quality based on the 1-Nearest-Neighbour accuracy metric between generated and test-set point-clouds. On uniform, 3D datasets, WFM is competitive with current approaches (Table 3), but exemplifies itself with its unique ability to generate point-clouds with varying sizes and in high-dimensions (Table 2).

#### 4.2.1 2D & 3D Point-Clouds

Figure: Figure 5: Generated point-clouds from MNIST and EMNIST datasets using WFM, where $n$ denotes the number of points in each cloud.
Refer to caption: x3.png

Derived from 3D CAD designs, ShapeNet & ModelNet (Wu et al.,, ; Chang et al.,, ) are touchstone point-cloud datasets in computational geometry. Trained individually on samples from the chair, car and plane classes of ShapeNet, WFM synthesized high quality point-clouds with diverse profiles and matches the performance of previous 3D generation algorithms; see Figure 4 and Table 3.
Our framework’s versatility allows for seamless integration of label information during training, enabling the synthesis of point-clouds conditioned on specific classes. On the full 40-class ModelNet dataset, WFM learned condition dependent flows, allowing for the same initial point-cloud to generate a diverse cohort of shapes based on the desired label; see Figure 4. We stress that WFM is not restricted to only noisy source measures but can generate transformations between any two collections of point-clouds. To this end, we demonstrate that WFM can interpolate between two arbitrary elements in the dataset (e.g., between a lamp and a handbag) and complete the point-clouds based on partial profiles (e.g., generate the remaining parts of a plane); see Figure S3.

Another novel facet of WFM is its ability to perform generative modeling from inhomogeneous datasets, where the number of points varies between independent samples. This happens in the MNIST or Letters datasets, where data is generated by thresholding grayscale numerical values. In this setting, we use the entropic transport map to approximate the objective; see equation 14. WFM sets itself apart from other methods, which are restricted to uniform datasets, by leveraging the entropic OT map’s ability to compute feasible transformations between point-clouds of different sizes. Our experiments in Figure 5 demonstrate that WFM generates high-quality & diverse samples, despite large variability in the number of points per sample, which is itself a novel contribution.

**Table 3: Using 1-Nearest-Neighbour Accuracy based on Earth Mover’s Distance (EMD) and Chamfer’s Distance (CD). Wasserstein Flow Matching is competitive with existing approaches (data from Wu et al., ()), while producing diverse samples, see Figure 4.**
|  | Airplane | Chair | Car |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | CD $\downarrow$ | EMD $\downarrow$ | CD $\downarrow$ | EMD $\downarrow$ | CD $\downarrow$ | EMD $\downarrow$ |
| PointFlow | 75.68 | 70.74 | 62.84 | 60.57 | 58.10 | 56.25 |
| SoftFlow | 76.05 | 65.80 | 59.21 | 60.05 | 64.77 | 60.09 |
| DPF-Net | 75.18 | 65.55 | 62.00 | 58.53 | 62.35 | 54.48 |
| Shape-GF | 80.00 | 76.17 | 68.96 | 65.48 | 63.20 | 56.53 |
| PVD | 73.82 | 64.81 | 56.26 | 53.32 | 54.55 | 53.83 |
| PSF | 71.11 | 61.09 | 58.92 | 54.45 | 57.19 | 56.07 |
| WFM (ours) | 73.45 | 71.72 | 58.98 | 57.77 | 56.53 | 57.95 |

#### 4.2.2 Spatial Transcriptomics (High Dimensional Point Clouds)

In spatial transcriptomics, the niche of a cell is the point-cloud in high-dimensional gene-expression space of its immediate nearest neighbours. This approach is complementary to the BW representation of a niche (recall Section 4.1.2), and serves as a more high fidelity view suited for fine-grain interactions. Due to their high dimensionality, cellular microenvironments have remained beyond the reach of point-cloud-based generative models that depend on voxel-based neural networks. In contrast, WFM leverages transformers, which due to their permutation equivariance and indifference to dimensionality, have already emerged as possible architecture to model spatial transcriptomics based point-clouds ().

During embryogenesis, specific regions within the primitive gut tube differentiate into organs such as the liver or lungs based on interactions between the gut and surrounding mesenchyme (Nowotschin et al.,, ). Applied on environments of gut-tube cells from a seqFISH dataset of mouse embryogenesis (Lohoff et al.,, ), WFM synthesized cellular niches conditioned on organ labels, thus demanding an understanding of the interplay between spatial context and phenotype. Despite the intricate nature of the gastrulation process, compunded by the dataset’s dimensionality, WFM can accurately generate organ-specific niches; see Figure 6 and Table 2.

Figure: Figure 6: The niche of each cell is the point-cloud from gene-expression profiles for cells in its environments. Using the top 16 principal components of gene expression, WFM learns to generate high-dimensional cellular microenvironments of gut-tube cells conditioned on the gastrulating organs.
Refer to caption: extracted/5971869/Figures/wfm_seqfish_gut.png

## 5 Conclusion and outlook

This work shows how to appropriately *lift* the Riemannian flow matching paradigm of Chen and Lipman, () to the Wasserstein space, resulting in Wasserstein flow matching. Our motivations stem from modern datasets, where each sample of data can itself be viewed as a probability distribution, necessitating this extension for generative modeling purposes. Our contributions are algorithmic in nature, which incorporate various elements, such as estimating optimal transport maps via entropic optimal transport, closed-form expressions over the Bures–Wasserstein space, and attention mechanisms in neural network architectures. Our algorithm is capable of generating realistic data from Gaussian and variable-size or high-dimensional point-clouds. Both contexts are highly relevant in single-cell and spatial transcriptomics for synthesizing of microenvironments and cellular states.

## Acknowledgements

AAP acknowledges support from NSF grant DMS-1922658. DP is an Investigator at the Howard Hughes Medical Institute and is supported by National Cancer Institute grants
P30 CA08748 and U54 CA209975.

## Appendix A Entropic estimation of OT maps

We briefly discuss how to estimate optimal transport maps between point-clouds using entropic optimal transport. We refer the interested reader to Pooladian and Niles-Weed, () for more information on this approach.

We first outline the numerical aspects of the approach; we follow Peyré and Cuturi, (). Let $\mu=\sum_{i}m^{-1}\delta_{x_{i}}$ and $\nu=\sum_{j}n^{-1}\delta_{y_{j}}$, where $\bm{X}=\{x_{1},\ldots,x_{m}\}$, $\bm{Y}=\{y_{1},\ldots,y_{n}\}$. We first define the following polyhedral constraint set

$$ $\displaystyle U_{m,n}\coloneqq\bigl{\{}P\in\mathbb{R}^{m\times n}_{+}\ :\ P\bm {1}_{m}=m^{-1}\bm{1}_{m}\,,P^{\top}\bm{1}_{n}=n^{-1}\bm{1}_{n}\bigr{\}}\,,$ $$

which represents the possible couplings between the two discrete measures. The entropic optimal transport coupling between the two discrete measures $\mu$ and $\nu$ is defined as the minimizer to the following strictly convex optimization problem

$$ $\displaystyle\bm{P}^{\star}\coloneqq\operatorname*{argmin}_{P\in U_{m,n}} \langle C,P\rangle+\varepsilon H(P)\,,$ (13) $$

where $\varepsilon>0$, $H(P)\coloneqq\sum_{i,j}P_{i,j}(\log(P_{i,j})-1)$, and $C_{i,j}\coloneqq\|x_{i}-y_{j}\|^{2}_{2}$. Sinkhorn’s matrix scaling algorithm (Sinkhorn,, ) makes it possible to solve for $\bm{P}^{\star}$ with a runtime of $O(mn/\varepsilon)$ (Altschuler et al.,, ). We briefly stress two points:

- 1.
The coupling $\bm{P}^{\star}$ is *not* a permutation matrix. The coupling lies inside the polytope $U_{m,n}$ and not at the vertices, and therefore is not a permutation matrix.
- 2.
When $\varepsilon=0$, the objective becomes a standard linear program with a runtime of $\tilde{O}(mn(m+n))$ (up to log factors) (Peyré and Cuturi,, , Chapter 3). However, this approach does not come with a GPU-efficient implementation and we decide to omit it entirely.

In all our experiments, we used the open-source package OTT-JAX(^2^22See [https://ott-jax.readthedocs.io/en/latest/](https://ott-jax.readthedocs.io/en/latest/).) to compute the entropic coupling and the out-of-sample mapping (Cuturi et al.,, ).

### A.1 Rounded matchings

Our first approach holds when $m=n$. In this case, we can greedily *round* the noisy matching matrix $\bm{P}^{\star}$ to become a permutation. This is achieved through an iterative process of selecting the maximum value (argmax) and zeroing out corresponding rows and columns. This method repeatedly identifies the largest remaining probability, sets it to 1, and eliminates other entries in its row and column, ultimately resulting in a permutation matrix that preserves the probabilistic assignment implied by the original doubly stochastic matrix. This is merely a GPU-friendly heuristic approximation to the true optimal permutation matrix between the two point-clouds.

### A.2 Entropic transport map: An out-of-sample estimator

A primal-dual relationship of the strictly convex program equation 13 shows that there exist vectors $(\bm{f}^{\star},\bm{g}^{\star})\in\mathbb{R}^{m}\times\mathbb{R}^{n}$ such that

$$ $\displaystyle\bm{P}_{i,j}^{\star}=e^{\bm{f}_{i}^{\star}/\varepsilon}e^{-\bm{C} _{i,j}/\varepsilon}e^{\bm{g}_{j}^{\star}/\varepsilon}$ $$

These two vectors are called the Kantorovich potentials, which are initially defined on the support of $\mu$ and $\nu$, respectively. However, they can be readily extended to all of $\mathbb{R}^{d}$ (Mena and Niles-Weed,, ), resulting in two functions

$$ $\displaystyle\hat{f}(x)=-\varepsilon\log\Bigl{(}\sum_{j=1}^{n}n^{-1}\exp((\bm{ g}_{j}^{\star}-\|x-y_{j}\|^{2})/\varepsilon)\Bigr{)}\,,$ $\displaystyle\hat{g}(y)=-\varepsilon\log\Bigl{(}\sum_{i=1}^{m}m^{-1}\exp((\bm{ f}_{i}^{\star}-\|y-x_{i}\|^{2})/\varepsilon)\Bigr{)}\,.$ $$

Following Pooladian and Niles-Weed, (), we can define the *entropic transport map*, where the last equality is a simple calculation:

$$ $\displaystyle\hat{T}_{\varepsilon}(x)$ $\displaystyle\coloneqq x-\nabla\hat{f}(x)=\frac{\sum_{j=1}^{n}y_{j}\exp((\bm{g }_{j}^{\star}-\|x-y_{j}\|^{2})/\varepsilon)}{\sum_{j=1}^{n}\exp((\bm{g}_{j}^{ \star}-\|x-y_{j}\|^{2})/\varepsilon)}\,.$ (14) $$

This estimator was initially to provide statistical approximations to the optimal transport map $T_{\star}^{\mu\to\nu}$ on the basis of samples; see Pooladian and Niles-Weed, (); ; Pooladian et al., ().
Note that $\hat{T}_{\varepsilon}(x)$ can be interpreted as the conditional expectation of the plan $\bm{P}^{\star}$ conditioned on out-of-sample inputs $x\in\mathbb{R}^{d}$, which is well-defined due to the relations above. Finally, we stress that this estimator can be adapted to settings where the point-clouds $\mu$ and $\nu$ not only have different numbers of points, but also non-uniform weights. As this estimator is also a by-product of Sinkhorn’s algorithm, it is also scalable and GPU-friendly.

## Appendix B Multisample Wasserstein Flow Matching

Since optimal transport can be applied on the Wasserstein manifold itself, both WFM and BW-FM can be seamlessly integrated with the multisample FM (MS-FM) framework (; Tong et al.,, ). The core technique behind MS-FM is to use OT to match minibatches from source and target measures during training, rather than relying on random pairings. This has shown to improve learned flows while requiring fewer function evaluations to synthesize new samples. Applying MS-FM requires computing the pairwise distance matrix between source and target batch samples, denoted from $i\in\{1,\ldots,\texttt{Bsz}\}$. In the BW-FM setting, given two sets of Gaussians $\{(a_{i},A_{i})\}_{i=1}^{\texttt{Bsz}}$ and $\{(b_{i},B_{i})\}_{i=1}^{\texttt{Bsz}}$, their Frechét ($W^{2}_{2}$) distance matrix is:

$$ $C_{i,j}=\|a_{i}-b_{j}\|^{2}_{2}+\text{Tr}(A_{i}+B_{j}-2(A_{i}^{1/2}B_{j}A_{i}^ {1/2})^{1/2})$ (15) $$

We then use entropic OT to approximately solve the assignment problem on $C$ and compute a transport matrix. This is the converted into a one-to-one assignment matrix via rounded matching (Section A.1), ensuring the entire batch is used in training.

For WFM point-clouds, applying MS requires computing pairwise OT distance between all source and target samples within a minibatch. For large point-clouds, this is exorbitantly expensive, even with Sinkhorn iterations. For an efficient approximate, here too we rely on the Frechét distance, computed between empirical means and covariances of each point-cloud. Computation of the Frechét distance is markedly less resource-intensive than entropic OT, yet is notably correlated with EMD values (see Table 1 in ()).

## Appendix C Sampling from source measure

In both WFM and BW-FM, learning flows requires a source measure which is straightforward to sample from. For a source distribution on the space of $\{(m,\Sigma):m\in\mathbb{R}^{d},\Sigma\in\mathbb{S}^{d}_{++}\}$, we simply sample means and covariance matrices using independent Gaussian and Wishart distributions, respectively. By default, the parameters for the Gaussian component of the source matches the average and standard deviation of the means in the target, while the scale parameter in the Wishart is the barycenter of the data covariance.

To achieve high-quality generation of point-clouds, it is essential that the initial (source) distribution be diverse, rather than collapsed and degenerate. Indeed, while it is alluring to produce noisy point-clouds by sampling points from a single base distribution, i.e. $X=\{x_{i}\}_{i=1}^{n},\,x_{i}\sim\mathcal{N}(0,I_{d})$, as $n$ grows, the Wasserstein distance between instances goes to 0 0. To alleviate this, we draw point-clouds from multivariate Gaussians with a stochastic covariance:

$$ $\displaystyle L\sim\mathcal{N}(\mu_{L},$ $\displaystyle\sigma_{L}\cdot I)$ $\displaystyle X=\{x_{i}\}_{i=1}^{n},\,$ $\displaystyle x_{i}\sim\mathcal{N}(0,LL^{T})$ $$

where $\mu_{L}$ & $\sigma_{L}$ are the average and standard deviation of the Cholesky factors from the empirical covariances of the target measure point-clouds. This ensures a wider source measure, producing a diverse range of noise point-clouds.

Figure: Figure S1: Conditional BW-FM applied to single-cell RNA sequencing data of immune response to COVID-19. Large scale single-cell atlases are commonly grouped into highly dedicated clusters called MetaCells (Persad et al.,, ). In this application, BW-FM is conditioned on cell state and trained to generate means and covariances of gene expression, focusing on the top 32 principal components, derived from aggregated cells. The model achieves high-quality sample generation, as evidenced by a label accuracy of 93.13%.
Refer to caption: extracted/5971869/Figures/bw_metacells.png

## Appendix D Neural architecture & training

### D.1 BW-FM on Gaussians

The goal of BW-FM is to train a neural network to match the (Riemannian) time-derivative along the BW geodesics between Gaussians. The model employs a standard, fully connected neural network which takes as input concatenated values of $(m_{t},\Sigma_{t},t)$ based on the McCann interpolation formula from Section 2.3.1. Since the covariance matrix is symmetric, only its lower-diagonal values are used, flattened into a vector of length $d(d+1)/2$. Time values are converted to Fourier features, an approach inspired by positional encodings in transformer literature (Vaswani,, ). To streamline training, two separate networks are employed: one to match the time derivative of the mean $\dot{m}_{t}$ and another for the time derivative of the covariance matrix $\dot{\Sigma}^{{\rm{BW}}}_{t}$. The BW tangent norm is used as the loss function for training these networks.

By default, all models use a $6$-layer neural network using relu non-linearity, with $1024$ neurons per layer, applying skip connections and layer-norm Ba, (). Training is performed for $100,000$ gradient descent steps using the Adam optimizer (Kingma,, ) with an exponential learning rate decay of $0.97$ every $1000$ steps and batch size of $128$.

Figure: Figure S2: Spatial arrangement of microglia in the motor cortex. Bures-Wasserstein distance based 2D UMAP visualization of real microglia and BW-FM synthesized niches, colored by their first diffusion component (DC). This DC corresponds to the cortical depth of the microglia across the MERFISH slices.
Refer to caption: extracted/5971869/Figures/bw_glia_supp.png

### D.2 WFM on point-clouds

WFM is designed to estimate the optimal transport (OT) map for a given pair of interpolate point-cloud and time $(\bm{X}_{t},t)$. Here too the time component $t$ is first converted into Fourier features. The model’s architechture begins with an embedding layer, followed by a series of alternating multihead attention and fully-connected layers. Skip connections and layer-norma are applied after each operation. The final layer projects the embeddings back to $X$’s original space using a dense layer with zero initialization. The model is trained by minimizing the squared distance between the predicted and true OT maps.

By default, the entropic OT map is constructed with regularization weight of $\varepsilon=0.002$ and $200$ Sinkhorn iterations, which we found to be sufficient for convergence. Whenever the dataset consists of uniformly sized point-clouds, we use rounded matching (Section A.1), otherwise we apply the out-of-sample estimator (Section A.2) which can calculate maps between point-clouds with different sizes. The transformer network is composed to $6$ multi-head attention blocks, with an embedding dimension of $512$ and $4$ heads. The model is trained for $500,000$ steps, using the Adam optimizer with exponential learning rate decay and batch size of $64$.

## Appendix E Experiment Details

### E.1 Spatial Transcriptomics

In our manuscript, we applied WFM and BW-FM on several spatial transcriptomics datasets, encompassing a variety of technologies and tissue contexts. From a 254-gene MERFISH atlas of the motor cortex (Zhang et al.,, ), we focus on niches of microglia cells. We compress gene-expression profiles down to their $16$ principal components (PC) and aggregate all the cells around each microglia within an $80$ micron radius, yielding on average $26.6$ cells per niche. We then calculated the gene-expression PC mean and covariance within each environment to produce Gaussians for BW-FM. Generated Gaussians align with real data and span for the full cortical depth of the microglia niches (Figure S2). In Figure 3, we predict the environment composition by cell type for generated Gaussians via nearest-neighbour regression in BW space using real data as supervision, demonstrating congruence between the two across cortical depth.

Figure: Figure S3: Interpolation and shape completion with WFM. Top. Using the lamps and handbags as the source and target measures, WFM learns to transform a given (unseen test-set) lamp point-cloud into a valid handbag. Bottom. Trained to generate full planes, WFM can reconstruct complete point-clouds from partial views of test-set samples.
Refer to caption: extracted/5971869/Figures/wfm_mug2bowl.png

In a complementary approach, WFM is applied directly on gene-expression based point-clouds of niches, and does not require the Gaussian representation. Uniquely suited for high-dimensional data, we apply WFM on seqFISH assay of embryogenesis (Lohoff et al.,, ) and a XENIUM experiment of melanoma metastasis to the brain (). In both dataset, we select the $k=8$ physical nearest neighbours of every cell, and aggregate their first $16$ PCs to produce environment point-clouds.

From the seqFISH dataset, we concentrated on the gut-tube region, which is divided into spatially segregated, gastrulating organs. Applied unconditionally, WFM generated niches match the distribution of the real data based on EMD and CD 1-nearest-neighbour accuracy (Table 2). We then assessed WFM’s capability to comprehend the relationship between cell state and environment and tasked it with conditional generation based on organ label. Based on OT distances estimated via Wormhole embeddings (), point-clouds from WFM recapitulated true organ environment.The label accuracy for WFM-generated data was 78.86%, which was nearly identical to the test-set real data accuracy of 79.59%.

## Appendix F 2D & 3D point-clouds

Figure: Figure S4: Cities Dataset. Gaussians representing the $100$ most populous cities in the continental US. The data was obtained from Bennett, () via OSMnx (Boeing,, ). The mean parameter is the longitude and latitude coordinate of each city and the covariance is the 2nd moment approximation of their metro area.
Refer to caption: extracted/5971869/Figures/bw_cities.png

The ShapeNet dataset consists of 3D point-clouds of $55$ different classes, each one comprised of $15,000$ points. Emulating the benchmarking effort in (Wu et al.,, ), we apply WFM to generate $n=1000$ sized examples from the plane, car and chair classes. At each gradient descent step, we sample $64$ point-clouds from the training set for each class, and randomly select $n=1000$ points from each. To evaluate generation quality, we synthesize point-clouds to much the size of the test set, and calculate the real or generated $1-NN$ accuracy based on EMD and CD metrics.

ModelNet has $40$ classes of point-clouds, with $2048$ points in each. Conditioned on class label, WFM is trained to generate $n=1000$ sized point-clouds here too. In this setting, the noise measure is the standard normal and we did not use multi-sample matching. According to nearest-neighbour classification from OT preserving Wormhole embeddings, generated samples match their class with an accuracy of $77.66\%$, approaching the $79.98\%$ purity of test set samples from real-data.

The MNIST dataset is a widely used collection of handwritten digits, consisting of 28x28 pixel grayscale images of the numbers 0 through 9. EMNIST (Extended MNIST) is an expansion of MNIST that includes handwritten letters as well as digits. To convert samples from these datasets into point-clouds, we threshold each image and extract the coordinates of the above-threshold pixels. This produces a cohort of point-clouds of variables sizes, as each image contains a different number of relevant pixel. We apply the entropic OT map (see Appendix A) based WFM to synthesize point-clouds of the digit $4$ and letter $a$. Despite the data heterogeneity, WFM produces realistic examples (Figure 5), while capturing the data distribution (Table 2). We again stress that this is a unique feature of WFM, lacking from any previous point-cloud generation algorithm

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - The notation $\overline{(\cdot)}^{L^2(\mu)}$ is not sufficiently clear in lines 182-183. Could you please write rigorously what is meant here? And I recommend moving such kind of notations in the corresponding part of Section 2.

**References.** 

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savare ́. Gradient flows in metric spaces and in the space of probability measures. Lectures in Mathematics ETH Zu ̈rich. Birkha ̈user Verlag, Basel, second edition, 2008. ISBN 978-3-7643-8721-1. 

Marc Lambert, Sinho Chewi, Francis Bach, Silve`re Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, 35:14434–14447, 2022.

Klein, D., Uscidda, T., Theis, F., \& Cuturi, M. (2024). GENOT: A Neural Optimal Transport Framework for Single Cell Genomics. NeurIPS 2024

Eyring, L., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z., \& Theis, F. J. Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation. ICLR 2024


**After rebuttal.** In the last response of the authors, they have clarified several theoretical aspects of their work. That is, they establish the conditions guaranteeing the validity of their conditional flows. I appreciate these important additions. However, overall, I still have some concerns regarding this paper, e.g., moderate performance of the approach in low dimensions with an absence of baselines in high dimensions (I found out that here the 1-NN accuracy seems to be also quite moderate ~ 60%), some gaps between theory and practical implementation. I see the potential in this paper, thank the authors for their work during the rebuttal but currently keep my score.
soundness: 2
strengths: The proposed method can be used to deal with families with distributions, i.e., distributions over distributions. It was tested on a variety of Gaussian-based where it shows better performance than naive baselines.
summary: This paper proposes a new type of generative model based on flow matching, i.e., Wasserstein flow matching. The proposed model allows for working with the *distributions of distributions* and was tested in a variety of experimental setups.
weaknesses: My main concern is related to the limited significance of the developed methodology. As the authors themselves admit (ines 235-236), the main contribution of this article is the developed algorithm. From a theoretical point of view, the method is based on the results proposed in previous works (Ambrosio et al., 2008, Lambert et al., 2022). At the same time, the authors note that these theoretical results are correct if the algorithm is applied to absolutely continuous measures, which is not true in the case when the measures are approximated as point clouds. Thus, it is not absolutely fair to apply the algorithm to any measures other than Gaussians.

Meanwhile, the theoretical groundlessness of the algorithm could be overlooked if it showed inspiring results from a practical point of view. However, the method demonstrates good results only when working with distributions that can be approximated by families of Gaussian distributions (including the spatial transcriptomics dataset, see lines 414-420). In the case of working with more general measures (point clouds), the method shows results that do not outperform those of its competitors (see Table 3 and authors comment in lines 95-96).

In addition, the experimental part could be improved by including more baselines, for example in problems from the single cell biology domain. In this regard, you might be interested in considering the methods from two recent papers that also rely on flow matching and consider problems from single cell biology (Klein et al., 2024, Eyring et al., 2024).

**Overall**, in my opinion, the presented practical results alone are not enough to publish the paper at the current conference.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: I do not have any questions.
soundness: 3
strengths: The paper tackles a novel and important problem within generative modelling of the generation of probability measures, with particular relevance for single-cell genomics data. The authors present a sound methodology in the particular cases of modelling distributions of Gaussian distributions and point-cloud data, with the notable attractive feature of being able to model point-cloud data of non-uniform sizes which is to the best of my knowledge, the first instance of this being tractable; this is backed up with convincing experimental validation. In addition, the paper is well-written and presented.
summary: The paper demonstrates an application of Riemannian Flow Matching for the novel problem of generating probability measures on the Wasserstein manifold, in particular, for the case of the Bures–Wasserstein space - i.e. distributions over Gaussian distributions, and for point-cloud data considered as empirical measures. The authors present tractable training objectives and sampling schemes for both settings, relying on closed-form solutions available in the Bures–Wasserstein space, and the approximation of the optimal transport map by entropic couplings in the point-cloud setting. Notably, their methodology allows for the generation of point-clouds with non-uniform size. They further validate their framework via experiments with single-cell genomics data and 2D and 3D point-clouds, presenting competitive results and improved scalability.
weaknesses: A potential criticism of the work is that the main components of the methodology such as Riemannian Flow Matching, closed-form solutions for geodesics etc. on the Bures–Wasserstein space, the entropic estimation of optimal transport maps etc. have been introduced elsewhere, and the main contribution of the paper is simply the original combination of these elements applied to their problem. In addition, the authors could have further discussed other settings in which generative modelling on the Wassertein manifold could be useful, such as amortising Bayesian inference, beyond the two cases considered in the paper.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - Could the authors provide training times, including times for solving entropic OT, for the point cloud generation tasks? Additionally, a comparison with other models, such as diffusion-based point cloud generation methods, would be insightful.

- Does the optimal velocity field $v_t$ that minimizes WFM objective guarantee the transportation of $\mathfrak{p}_0$ to $\mathfrak{p}_1$? Could the authors provide the theoretical proofs?

- What is the difference of this algorithm compare to mini-batch OT with Wasserstein distance cost (in point clouds)?

$ $

**Additional Comments** Thank you for your response. However, certain aspects of the authors' explanation remain unclear. First, the authors assume that the marginals must have a "pdf." Furthermore, they claim to apply this assumption to Flow Matching, which requires a positive density across the entire domain (generally, I think the domain is $\mathcal{W}_{2,ac}(M)$). I believe these two assumptions are contradictory since it is (typically) not possible to define a pdf in infinite-dimensional spaces; only distributions exist in such cases. This seems to create a fundamental inconsistency in the assumptions. 

That said, I do understand and appreciate some points the authors have discussed. For instance, the Bures-Wasserstein manifold is a finite-dimensional Riemannian manifold, which allows for the direct application of Riemannian Flow Matching. Consequently, I understood that we can define a "nice subspace" that ensures the flow matching vector field $v$ can transport one marginal to another. With this clarification, the discussion becomes much clearer. 

I believe that remaining small issues can be easily addressed, so I will raise my score to 6.
soundness: 3
strengths: - The paper presents an new concept of Wasserstein Flow Matching, which extends the Flow Matching framework to probability distribution spaces.
- It effectively incorporates Gaussian cases through McCann’s (displacement) interpolant.
- The method is applied to a variety of data types, including genomic data and point cloud data.
- Overall, the paper is clear, well-structured, and easy to follow.
summary: This paper introduces Wasserstein Flow Matching (WFM) for families of set of particles and Gaussians. In other words, this paper attempts to train the FM model on Wasserstein space. For given the pairs of probability densities $\Gamma$, the method first approximates McCann’s interpolant of $(\mu, \nu) \sim \Gamma$ through traditional OT algorithms or closed-form solutions. By this approximation, they obtain the interpolant $\mu_t$, and the velocity $v_t$. Then, the approach trains the model $F_\theta$ that learns the mean field on each $(t, \mu_t)$ through flow matching loss ($\\| F_\theta (t, \mu_t) - v_t \\|^2$). This paper evaluates their algorithm on the cell dataset and point cloud generation.
weaknesses: - The motivation for learning the mean field on the Wasserstein space is not entirely clear, as the mean field may not capture meaningful structures like Wasserstein geodesics. Additionally, the approach first estimates McCann’s interpolation using traditional methods, followed by the FM algorithm. This sequential combination of the two algorithms feels somewhat incremental in terms of contribution.

- For cases where families of particle sets (point cloud) are considered, the paper estimates entropic OT, which deviates from the theoretical framework. Calculating the McCann's interpolant $\mu_t$ requires estimating the standard OT rather than entropic OT. 

- Even with reduced computational overhead through estimating entropic OT, I believe there would be still huge computational burden when estimating $\mu_t$ in point cloud application. It might be helpful if the authors could specify the computational costs associated with solving entropic OT during training. I believe there would be scalability issues when the number of particles becomes more than 2000.

- Overall, the method involves an inaccurate approximation (e.g. standard OT -> entropic OT) and seems to be computationally demanding (e.g. sinkhorn for every iterations). Moreover, its performance results is not convincing for me, too. I believe the benefit of WFM method and the potential applications should be more carefully carried.

Meta-review: In the paper, the authors proposed Wasserstein flow matching, a type of generative model based on flow matching. While all the reviewers agree that the proposed method is novel, the paper still has some major weaknesses, including (1) the gap between theoretical justification and practical implementation of the method; (2) the limited significance of the method, i.e., the experimental results are not thorough. For instance, as pointed out by the reviewer, the experiments with point clouds in high dimensional settings are lacking; (3) the assumptions are not consistent as pointed out by the reviewers. 

I myself also read the paper and agree with the reviewers about the current weaknesses of the paper. Therefore, I recommend rejecting the paper at the current stage. The authors are encouraged to revise the paper based on the suggestions of the reviewers for future resubmission.

Decision: Reject