Wasserstein Flow Matching: Generative Modeling Over Families of Distributions
Authors: Doron Haviv
Aram-Alexandre Pooladian
Dana Pe'er
Brandon Amos
Venue: ICML 2025 poster (2025)
Keywords: Optimal Transport
Flow Matching
Generative Modeling
Single Cell
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: Flow Matching between distributions over distributions on the Wasserstein space for generation of Gaussians & general distributions.

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](https://github.com/WassersteinFlowMatching/WassersteinFlowMatching/).

Wasserstein Flow Matching: Generative Modeling Over Families of Distributions
Doron Haviv 12
Aram-Alexandre Pooladian Dana Pe'er Brandon Amos 5
Abstract
access to finitely many samples from a distribution, devise scheme which generates new samples from the same dis- tribution. Generative modeling has also been revolutionary in the biomedical sciences, for drug design (Jumper et al., 2021), and single-cell genomics (Lopez al., 2018). Nearly all frameworks exploit the notion that datasets (of, say, ge- nomic profiles of cells, images, videos, or corpora of text documents) are instantiations of probability measures, and the task to transform a point sampled from random noise to generate a data point that obeys the distribution of interest.
Generative modeling typically concerns transport- ing a single source distribution target distri- bution via simple probability flows. However, in fields like computer graphics and single-cell ge- nomics, samples themselves can be viewed as dis- tributions, where standard flow matching ignores their inherent geometry. We propose Wasserstein flow matching (WFM), which lifts flow matching onto families of distributions using the Wasser- stein geometry. Notably, WFM is the first algo- rithm 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 valid FM objective. Our algorithm leverages optimal trans- port theory and the attention mechanism, demon- strating 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 avail- able at Wasserstein Flow Matching.
Among the zoo of available generative models, one approach noted for its flexibility and simplicity is Flow Matching (FM) (Albergo & Canden-Fijnden, 2022; Lipman et al., 2022; Liu et al., 2022). For a fixed target probability mea- sure, FM learns an implicitly defined vector field that can transform a source measure (c.g., the standard Gaussian) to the target. Unlike discrete time and probabilistic generative models (such Diffusion Models by Song et al. (2020)), FM learns deterministic, continuous normalizing flow by regressing onto simple conditional probability flow. This approach, while originally designed for Euclidean domains, can be readily adopted to Riemannian geometries (Chen & Lipman, 2023). Riemannian flow matching (RFM) is widely used for generating samples over geometries such as spheres, tori, translation/rotation groups, simplices, triangu- lar meshes, mazes, and molecular positions and structures.
However, modern data can exhibit a richer structure in which the samples themselves are distributions. In computa- tional graphics, 3D models are shape distributions observed through 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 microenviron- ments or highlight fine-grain clusters (Haviv et al., 2024b; Persad et al., 2023). For both general distributions and Gaus- sian settings, it is natural to search for a unified generative model that respects the underlying geometry of the data.
1. Introduction
Today's abundance of data and scalability of training mas- sive neural networks has made it possible to generate hyper- realistic images on the basis of training examples (OpenAI, 2022), as well as video and audio clips (Vyas et al., 2023; Xing et al., 2023), and, of course, text (Bubeck et al., 2023). All of these are instances of generative modeling: given
'Equal contribution 'Weill Comell Memorial Keeterng Cancer Center *Center for Data Science, New York University 'Howard Hughes Medical Institute MMta AL Correspondence to: Doron Haviv <havivd@mskoce.rr orgg..orgss..c.......ooo.omrm... Aram-Alexande Pooladian <ap6599@myu.cdal nyu.edu>, Brandon Amos <bdda @ meta.com>>
Contributions. To this end, 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. These distributions are either represented analytically as Gaussians or realized
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Wasserstein Flow Matching
through point-cloud samples. 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& & Lipman, 2023), where we train a neural model to learn a continuous normalizing flow (CNF) between distributions over distributions.
po
p1
N( Enx)
N(my, Σn)
1 m Di das
1 =
Σ3 60z
We demonstrate the effectiveness of our approach for gen- erative modeling between distributions over Gaussian dis- tributions and distributions over point-clouds, which are realization of general distributions. The former task is mo- tivated by recent directions in single-cell and spatial tran- scriptomics (Haviv et al., 2024b; Persad et al., 2023), where we consider matching problems over the Wasserrassrsniin space (BW), the Gaussian submanifold of the Wasserstein space. In this case, we show that WFM can be further modi- fied, resulting in the Basse-Wannrn FM (BW-FM) algo- rithm. We validate BW-FM on variety of Gaussian-bar bassdd datasets, where we observe that samples generated by our al- gorithm are significantly more robust than naive approaches which do not fully exploit the underlying geometry of the data. As an application, we present generative model for cell states and niches from single-cell genomics data.
000
Figure I. WFM learns flows between distributions over distribu- tions, where distributions are either analytically represented (Gaus- sians) or empirically observed through point-clouds.
are realized as point-clouds with variable sizes. Conversely, WFM succeeds in these challenging settings, enabling gen- erative modeling in new domains like synthesizing tissue microenvironments from spatial genomics data. Modeling 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 immune microenvironment, these insights are critical for developing novel therapeutics (Binnewies et al., 2).
Method
Data type
Source
Target
2. Background and related work
FM over Rd
e Rd
I Po
y~ p
We let (R()) denote probability distributions over Rd with finite second moment, and write P2.cc for those with densities. For a probability measure l and function f. we write [=() for the squared L L²( norm. Let M be a Riemannian manifold with tangent space T&M at z. For xo E M with velocity v E M, the exponential map expzo (v) gives the terminal location, while the logarithmic map logzo (x1) gives the initial velocity for the geodesic from my to x1. The set of symmetric (resp. positive definite) matrices over IRdd are denoted by gd (resp. S4
FM over M
E M
x ~ Po
y P1
FM over Ad
EE P(Aa)
p Po
1 ~ p1
Wasserstein FM
P(Rd)
p. Po
Gaussians
W(m, x.)
(mm E.)) Po
((m,,, + 66 ~ p
P1 F..)
Point-Clouds
2 Sz, ~ Po
Table 1. Contrasting FM methods over R general manifolds M. categorical and Dirichlet distributions the d-simplex Ad. and finally, our approach, FM problems defined over P(R*).
Generation of general distributions is made possible by two distinet, yet crucial, algorithmic primitives: (1) incor- porating transformers in our neural network architecture (Vaswani, 2017: Lee et al., 2019), and (2) recent algorith- mic advances in entropic optimal transport (Pooladian & Niles-Weed, 2021). Our WFM algorithm performs gener- ative modeling in the Wasserstein space, where geodesics are given by pushforwards of optimal transport (OT) maps; see Section 2.3 for more information. Since these OT maps lack closed-form solutions for general distributions, we rep- resent distributions as point-clouds and estimate the maps using entropic optimal transport. The permutation equivari- ance of attention makes transformers a natural basis for our model, inherently capturing the equivariance of Wasserstein geometry while maintaining scalability in high dimensions.
2.1. Riemannian flow matching
We first briefly discuss the Riemannian flow matching (RFM) framework of Chen & Lipman (2023). Let po be the source distribution and p1 be the target distribution over Riemannian manifold M, and let (1t)cejo.1) be curve of probability measures satisfying To = Po and = P1. Letting (wx) (w=))ejo.1 denote a family of vector fields, we say that the pair (7t, We We)ee(u,1) satisfy the continuity equation with respect to the metric g, abbreviated to (Jt, we) E Cg if
O V g (((w)))) = =0,
(1)
For datasets of distributions in 3D, WFM's performance matches existing generative models. However, due to their reliance on voxelization, current approaches cannot scale to high-dimensional distributions and fail when distributions
where Vg is the Riemannian divergence operator.
The goal of RFM is to regress a parameterized vector field (e.g., a neural network), written fo t) TyM for e
2

---

Wasserstein Flow Matching
[0, II. onto the family we by minimizing
Generative models for single-cell genomics. Deep learn ing based generative models have transformed single-cell genomics through various approaches. Variational auto- encoders (Lopez et al., 2018; Gayoso et al., 2022) have successfully addressed technical artifacts, integrating multi- modal data and imputing missing features in scRNA-seq data. More recently, Transformer-based foundation models have been noted for their ability to integrate large atlases of data (Cui et al., 2024; Theodoris et al., 2023). Flow Match- ing has also emerged as a promising direction (Klein et al., 2024; Eyring et al., 2024) for learning both balanced and unbalanced OT maps between cell populations. While these prior FM applications focus on cell-to-cell mappings, our work introduces a new paradigm: generating entire distribu- tions representing cellular populations. This is particularly relevant for spatial genomics, where cellular microenvi- ronments are naturally represented distributions. WFM enables synthesis of whole cellular neighhorhoods, a novel approach in generative modeling for the single-cell field.
min / - we
- we(2))))ee))
dri(z))
assuming access to a pair (7e, wy we)refo.j that satisfies (I). This is not possible in many scenarios. Borrowing insights from recent work (e.g., Albergo & Vanden-Eijndee (2022); Lipman et al. (2022); Liu et al. (2022)), the authors construct a simple vector field that satisfies the continuity equation, resulting in the tractable objective
min Jf [fo(xx.)) - 2 d2o(x) dp dt, (2)
where, for example, 32 = cxx,, ((1-)) t) log, E M, and f. E M, For complete discussions and proofs, see Chen & Lipman (2023, Section 3. Once So appropriately fit using (2), we can generate new samples from P1: start by sampling Xu ~ Po, then follow X, fx(e.,)) t) numerically by discretizing the dynamics given by the exponential map, resulting in X1 ~ P1. We emphasize that the dynamics are only simulated at inference time and not when training fa, commonly known as a simulation-fre free training paradigm.
2.3. Wasserstein geometry
The (squared) 2- Wasserstein distance between two probabil- ity measures μu, V/ E a2 is given by the optimization problem over vector-valued maps T : Idd - Rd
(1, == min Iiid - T(2220)) TYTTEE
2.2. Related work
Generative models for shapes. Paralleling the progress in generative models for natural images, the field of shape gen- eration rapidly expanding. Many different models have been used from this task, namely generative-aadversaiial nets (Achlioptas et al., 2018), variational autoencoders (Gadelha et al., 2018), normalizing flows (Yang et al., 2019; Kim et al., 2020; Klokov et al., 2020), diffusion (Zhou et al., 2021; Cai et al., 2020) and even euclidean FM (Wu et al., 2023a). Thus far, these approaches are limited to uniform- sized samplings of shape distributions in 2D & 3D, and fail on datasets where realizations are variably sized.
where the pushforward constraint, written Typ = v, means that for X ~ p. T(X) ~ v. The minimizer to (3) is called the optimal transport (OT) map, denoted T (we abbre- viate this to T. when clear). The existence and uniqueness of the optimal transport map under the stated regularity conditions is due to Brenier (1991).
The Wasserstein space is the space of probability densi- ties with finite second moment endowed with the Wasser- stein distance; this space is known to be a metric space (Villani, 2009). Following the celebrated work of (Otto, 2001), the Wasserstein space can be formally (meaning, non-rigorously) viewed as Riemannian manifold, whose properties we now describe in brief; see e.g., (Ambrosio et al., 2008) for a rigorous treatment. Following the defini- tion by Ambrosio et al. (2008, Theorem 8.5.1) 1), the tangent space at a point μ E P2 (R consists of all possible vec- tors that emanate from jr, written
Generative models over families of distributions. Our work is not the first instantiate Riemannian FM with manifold of probability measures. Two notable works are Fisher FM (Davis et al., 2024) and Categorical FM (Cheng et al., 2024), which consider the FM algorithm with respect to the Fisher-Rao geometry (Amari, 2016; Nielsen, 2020) over the d-dimensional simplex A,. The work of Stark et al. (2024) 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. (2024), 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.
id): -d): 0, E P2(R')) (H)
Tr. (RC) = [X(T -
where the overline denotes the closure of the set (i.e., the set and its limit points) in L²) and the norm on the tangent space is LH))). The exponential and logarithmic maps read
v + (x) = (id + v)sM, v + log == TTh - id,
where id is the identity map. Consequently, the (constant- speed) geodesic, or McCann interpolation, between two
3

---

Wasserstein Flow Matching
Training Data
Frobenius FM
BW-FM (Euclidean)
BW-FM (Riemannian)
AE HONNCDOINH DO
Degenerate Gaussians
Figure 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, they do not model geometry of the data, When there are only few examples, BW-FM accurately reconstructs the training data; see Figure $1.
measures μ and v is given by the curve (H4)1E(o.1) where
covariances {m, Σ) m E R I E 84 endowed with the Wasserstein metric. We provide a brief exposition on the geometry of the BarresWasssessedd space and refer to Lambert et al. (2022) for detailed calculations and explanations, as we follow their notation conventions.
1 = ((1 tjid 17T""))s 2o14 exp, (1 (o)),
N4
(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
The OT map between N(m,, Σ,)) and A(m,, Σ,) has closed-form expression (Gelbrich, 1990);
X, (1 - Xo + "IT=)(XX)) where Xg f and X, jr for any [o, 1]. Combined with (1 suitable family of vector fields, the McCann interpolation satisfies the continuity equation (1) over Rd, re-written as
Tx(x) = my + +CH (x- mp))
1 = my En (E (E 2xx mm+)
Oeple V (1404) = 0, s.t. N40 , = v (5)
As this map is affine, it is clear that the McCann interpo- lation between two Gaussians is always Gaussian. More generally, we have the succinct representation of the tangent space at point in the Bures Wassersesnn space
where the divergence operator is the usual Euclidean one over Idd thus we write (Me, Ut) E C. The link between the constant speed geodesics and the 2- 2-Wasserstein distance can be viewed from the celebrated Benamou-Brent for- mulation of optimal transport (Benamou Brenier, 2000):
Ty BW() = fa + S(id - m,) a E RRR... SSS.)
and the tangent space norm at μ can be written as:
[' [Ivell [1oell22(w)) de.
Wz,
inf (miniee
(6)
(I ()l = +
The optimal curve of measures is given by the constant- speed geodesics described above, and the optimal velocity field is given by
With the above, it is straightforward to compute the McCann interpolation = (Te)ap= =N(mm, Ex), with
= (T - id) (T!
(7)
my == (1- - m) + lm
(8)
Σ, = T,AT, T, (1 - +
The vector field (7) should be interpreted as the time- derivative of the McCann interpolation, with Xo ~ μ
We can relate the Euclidean and Riemannian time- derivatives of E, through the following manipulation:
X, = (T - -Xo)
(Tm- -dd) o ((x)))
Σマ - + + = BBw + E
(9)
2.3.1. BURES- WASSERSTEIN (BW) SPACE
To this end, we can draw parallels to (7) by writing
known special case of the Wasserstein space the Wasses-esssrsteii space, which consists of the submanifold of non-degenerate Gaussians parameterized by means and
m = my -Ma
(10)
BBW (CA-B [T)((1 - t)T --1-11-1

---

Wasserstein Flow Matching
Plane
Chair
Monitor
Vase
Toilet
Unconditional
Conditional
Figure 3. Left. Synthesized samples from WFM trained on the cars, planes or chairs datascts. Right. Examples generated conditionally from the same initial noise via a WFM model trained on the complete 40-class ModelNet datasct.
3. Flow matching over the Wasserstein space
Finally, we prove that regressing onto vector fields results in valid conditional probability flows between distributions of distributions; see Appendix B.1 for a proof.
Let Po and iı denote probability measures over the Wasser- stein space. Our goal is to learn vector field that trans- ports the family of measures Po to the family P1- WFM learns to map source to target by regressing onto Wasser- stein geodesics between samples jt Po and V ~ P1.' To accomplish this, we pass in the McCann interpolation fe and optimal velocity field v, the Riemannian FM objective (2), resulting in our Wasserstein FM (WFM) objective:
Proposition 3.1. Conditional probability paths pa(-Iv) gen- erated by conditional vector fields of the form vt (recall (7)) satisfy pi (-Ip) = V.
To prove this result, we following a general recipe by Chen & Lipman (2023) with calculations specific to the Wasser- stein space. First, we prove that the Wasserstein distance is pre-metric (see Lemma B.1). Then, we prove that the conditional vector fields e are defined pre-metric (see Lemma B.2). With these calculations, the proof of Proposition 3.1 then follows immediately from Theorem 3.1 of Chen & Lipman (2023).
min [f (1 (pr. t) velli.((m)) dpo dp1 dt. (11)
function of this
See Appendix B for a derivation of this training objective. remains to ensure that this objective describes valid con- ditional probability flows between families of distributions.
3.1. WFM over the Bures- space
First and foremost, we require two levels of regularity on the distributions of distributions po and p1:
First suppose Po and P1 are distributions over Gaussians, meaning that a batch of samples drawn from Po and pr consists of a batch of mean-covariance pairs. Here, the dynamics are straightforward: the interpolant is the Mc- Cann interpolation, and the velocity field over the Bures- Wasserstein manifold is also known (see (8) & (10), respec- tively). Since be is (m, E,), the neural network is parame- terized as fBw Rdd x 5 + Ie x gd, and the norm on the tangent space simplifies the computations considerably. Our final training objective becomes
1. Outer continuity: The source measure Po satisfies smoothness conditions (as in Appendix A of Chen et al. (2024)) ensuring the conditional paths pe(-lv) can be marginalized over p1 to form a valid flow
2. Inner continuity: For any µ Po and P1. there exists an optimal transport map between μ and mak- ing the conditional flows well-defined.
The condition of "inner continuity" is fairly mild, as this ensured for any distribution μ ~ Po with density. For Gaussian distributions, inner continuity holds naturally. For general distributions, we assume continuity but work with point-clouds as empirical realizations to approximate OT maps with statistical guarantees (see Appendix A.3) and computational efficiency (Flamary et al., 2021; Cuturi et al., 2022). There exists a slight gap between our theoretical re- sults which consider classic OT and our practical implemen- tation which uses regulay-regllariied OT approximations. This aligns with other geometric methods (see Algorithm 1 by Chen & Lipman (2023)) that rely on numerical ap- proximations when exact solutions are not tractable. The "outer continuity" condition is purely technical, and it serves the same role as in prior work. Our training algorithm is described in Algorithm 1, and Appendix E contains precise details on our neural network design.
[/ LBW (a ((m,,Ee) Et). (mi,, Σ. Il dpu dp dt.
3.2. WFM over distributions of general distributions
In the of general distributions, lose closed-form interpolations. However, we can hope to proceed so long as we have an approximation of the optimal transport map between them, written T, based on empirical samples (point- clouds) drawn from the distributions. There are many works on the approximation of these maps on the basis of samples; see Hütter Rigollet (2021); Divol et al. (2022); Manole et al. (2021); Pooladian & Niles-Weed (2021).
While existing point-cloud generative models work with uniform-sio sized samples from high-resolution CAD models, many datasets provide inherently variable-sizer samples. Though we assume these samples come from underlying continuous distributions, we only observe finite, discrete
"Crucially, we are not learning the OT map between Po and P1-
5

---

Wasserstein Flow Matching
Algorithm 1 Wasserstein FM Training
Algorithm 2 (Rd) Generation
Data: base Po E ())(R))) target p1 E P(P(IR)),
Data: Trained fow, step size h 1/N Init: (mmo Eo) ~ Po
geo E (BW, ]
Init: Parameters 0 of f
for k = .,..., N
do (sx,5) + fo (BW ((MRM: Ekh), kh)
repeat
Sample time L ~ ((,) 1)
mk++1)h - my + hsp. U* + (I + hSa)
Sample source measure f ~ Po Sample target measure V ~ p1 if geo is BW then
E(+11)A -- UK
end for Return: N(MNN,
Pt 4- (m4 24) via equation 8
Ut + (m,, > via equation 9
Algorithm 3 General Distribution Generation
else
Nr Approximate via equation 4 < Approximate via equation 7 end if
Data: Trained f8 step size h = 1/NN Init: Po ~ Po. X,- {X1, :.....: ..,] /4o
for k ===0 O...., N - 1 do
((0) + 1158 (He,t) ve (2l(; 0 < optimizer_step(O, (()). Vol())
X ((+1)) + Xah + h kh)
(Xi,,
end for
until Convergence
Return: XN4
realizations. This natural variation in sample sizes necessi- tates our development of WFM to handle point-clouds of non-uniform size, making WFM applicable across both high- resolution shape modeling and naturally discrete datasets.
4. Results
4.1. Families of Gaussians
We first demonstrate our flow matching framework between measures of Gaussians 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 co- variances, and trains on Σ. (equation (9)) with respect to the squared-Froheniu norm, and (2) BW-FM (Fuclidean), which tries to match E but still under the BW geometry.
We approximate optimal transport maps using entropic OT (Cuturi, 2013), which is computationally efficient on GPUs via Sinkhorn's 's algorithm (Sinkhorn, 1964). Let X and Y represent point-clouds described by data from f Po and P1 respectively, and let πuーン denote the entropic ap- proximation of the optimal transport map (see Appendix A for full details). The objective (1I) can be approximated by
BW-FM (R)
BW-FM (E)
Frobenius FM
Spiral 16 (2D)
2.98 10 4
4.00 10
1.03 10-3
Jf (1.80((x. - ((ffm-((x) x)) dpo dp dt,
Spirals 128 (2D)
1.28 10-3
1.70- 10-3
10-3
Two Moons (2D)
1.84 10-4
8.96 10-
1.30- 10-3
Sphere (3D)
6.65 10-4
2.14 10-3
2.25 10-3
Cities (2D)
1.88
7.26 10
1.75 10
ECG (15D)
9
3.26 3.26-10"
3.98 10
where X, x, -1 (1-1) x +ITx-v(X) represents a discretized McCann interpolation. We parameterize ffC with a trans- former, leveraging the natural alignment between the per- mutation equivariance of both OT maps and self-attention. Unlike point-cloud models based on voxelization, transform- ers maintain efficiency in high dimensions.
seqFISH (16D)
3.11 10-1
5.31 10-1
6.82 10-1
scRNA- seg (32D)
1.31
2.74
3.21
Table 2. Average min. W3 distance between each generated Gaus- sian and the reference datasets. Despite identical training schemes, BW-FM (R) outperforms other approaches on both synthetic and real data across several dimensions.
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; see Table 2. Notably, the flows gen- erated by Frobenius FM and BW-FM (Euclidean) do not strictly adhere to the geometry of the Bures wasserressr manifold, requiring synthesized covariance matrices to be projected onto the space of PSD matrices. In contrast, the Riemannian BW-FM algorithm consistently produces valid and accurate results across all dimensions and datasets.
3.3. Generation
Once 1800 trained, we can generate new samples as in Riemannian FM. For the -asse-Wassnrntein space, we ap- peal to the closed-form exponential and logarithmic maps; see Algorithm 2. We emphasize that the appropriate Rie- mannian updates are crucial to obtain non-degenerate final samples. For general distributions, we perform a standard Euler discretization of the learned flow on their point-cloud realizations; see Algorithm 3.
6

---

Wasserstein Flow Matching
4. TOY DATASETS
et al., 2023), and quantify the gene expression mean and covariance for each. Here too we apply BW-FM condi- tioned on cell-state, which encompass the heterogeneity of immune profiles appearing as response to COVID infec- tion. Despite the plurality of labels, BW-FM can synthesize correct examples for each condition; see Figure S3.
As a first test, we design a dataset of Gaussians centered on spiral: See Figures 2 & S1. When there were only few samples, only Riemannian BW-FM reconstructed the data, despite other benchmark methods following identi- cal training regimes. On the complete 128-sample dataset, generalizes beyond the training data and synthesizes novel Gaussians lying on the spiral. BW-FM shares this general- ization feature with standard FM, and is able to learn the structure underlying the measure from the training data.
4.2. Families of general distributions
Recall that for general distributions, we use entropic opti- mal transport to estimate McCann interpolations and veloci- ties (Cuturi, 2013). With high-resolution CAD datasets, we samplr 2048- poontt realizations from each shape to approxi- mate OT maps during training. For spatial transcriptomics, MNIST and Letters datasets, we work directly with their naturally varying sample sizes, leveraging entropic OT's ability to handle unequal point-cloud sizes while maintain- ing accurate transport estimation. These align with our assumption that the underlying distributions are continuous, with point-clouds serving as discrete realizations.
4.1.2. SINGLE-CELI GENOMICS
Spatial transcriptomics are set of techniques which build on single-cell genomics and preserve physical information of cells' location in tissues and assay their gene expression. Haviv al. (2024b) demonstrated that a cell's microenvi- ronment can be effectively characterized using the mean and covariance of the surrounding cells gene expression. This representation captures key features of cellular neighbor- hoods and transforms spatial transcriptomics datasets into measure within the Bures-Wasserstein space, highlighting the value of generative modeling in this context.
Airplanc
Chair
Car
CD J
EMD
CD
EMD
CD
EMD
PointFlow
75.68
70.74
62.84
60.57
58.10
56.25
During embryogenesis, specific regions within the primi- tive gut tube differentiate into organs such as the liver or lungs based interactions between the gut and surrounding mesenchyme (Nowotschin al., 2019). Applied envi- ronments of gut-tube cells from a seqFISH dataset of mouse embryogenesis (Lohoff et al., 2022), BW-FM synthesized Gaussian cellular niches conditioned on organ labels, thus demanding an understanding the interplay between spatial context and phenotype. Despite the intricate nature of the gastrulation process, compounded by the dataset's dimen- sionality, BW-FM can accurately generate organ-specific niches; see Figure $2.
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)
69.88
64.44
57.62
57.93
53.41
58.10
Table 3. Performance evaluation using 1-NN Accuracy (Earth Mover's Distance and Chamfer Distance) on distributions from ShapeNet, realized as uniform 3D point-clouds. Top three results for each metric are highlighted (Ist, 2nd, 3rd). WFM achieves sfate-of-the-art CD performance on 2 datasets while maintain- ing competitive results across other metrics (data from Wu et al. (2023a)), while producing diverse samples, see Figure 3.
We compare WFM to many other generative models. Fol- lowing in their footsteps, we measure generation quality based on the 1-Nearest-Neighbour accuracy metric between generated and test-set shapes. On uniform, 3D datasets, WFM is competitive with current approaches, but exempli- fies itself with its unique ability to generate samples with varying sizes and in high-dimensions; see Tables 3 & 4.
n= 78
n= 110
511
n= 388
Figure 4. Generated distributions from MNIST and Letters datasets using WFM. Each distribution is realized as a point-cloud with naturally varying sample sizes (n²), demonstrating WFM's unique ability to learn from and generate distributions without requiring fixed-size realizations.
4.2.1. 2D & 3D SHAPES
Derived from 3D CAD designs. ShapeNet & ModelNet (Wu et al., 2015; Chang et al., 2015) are touchstone shape datasets in computational geometry. Trained individually on samples from the chair, car and plane classes of ShapeNet, WFM synthesized high quality shapes with diverse profiles and matches the performance of previous 3D generation algorithms; see Figure 3 and Table 3. Our framework's 's ver- satility allows for seamless integration of label information during training, enabling the synthesis of shapes conditioned
Another common instance of BW manifolds arising in single-cell genomics is through aggregating cells into com- mon 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 (Stephen- son et al., 2021), we combine cells into MetaCells (Persad
7

---

Wasserstein Flow Matching
Subtyde
on specific classes. On the full 40-class ModelNet dataset, WFM learned condition-dependen flows between distribu- tions, enabling generation of diverse shapes from the same initial distribution based on the desired label; see Figure 3. We stress that WFM is not restricted to only noisy source measures but can generate transformations between any two collections of distributions. To this end, we demonstrate that WFM can interpolate between two arbitrary elements in the dataset (e.g., between lamp and handhag) and complete shapes based on partial profiles (e.g.. generate the remaining parts of plane); see Figure S5.
Sist Niches
WFM
Co
Cmh2 Hppe Cm
Stynä Th
Chool
OT
OT X
Figure 5. WFM enables high-dimensional generation of cellular microenvironment. (A) Wasserstein 2D embedding of observed Sst-neuron niches in motor cortex, colored by subtype. (B) WFM generated niches faithfuly reproduce of the tissue landscape
Current methods CD 4
WFM (ours)
EMD +.
CD 4
EMD 4
Variable-sized realizations
nized into phenotypically distinct subtypes, each localized to specific cortical layers (Wu et al., 2023b). Using the 254-gene MERFISH atlas (Zhang et al., 2021), we represent each Sst neuron's 's microenvironment as a distribution in gene expression space, observed through its 16 nearest neighbors. Despite operating in this nigh-dimensional space, WFM successfully generates niche distributions that capture both global tissue organization and subtype-specific molecular signatures; see Figures 5 & 6.
MNIST (4)
NA
NA
63.34
59.97
Letters (A)
NA
NA
62.12
58.68
diih-diimensonal distributions
MERFISH
NA
NA
52.80
54.52
XENIUM
NA
NA
60.69
64.20
Table 1- N-Neerest-Neighbor Accuracy demonstrating WFM's unique capabilities. Using transformers and entropic transport maps, WFM handles both variable-sized realizations and high- dimensional distributions, unlike previous approaches limited to fixed-size, low-dimensional samples.
Calb2+ Sst Niche
Another novel facet of WFM is its ability to perform gen- erative modeling from inhomogeneous datasets, where the number of points per point-cloud realization varies between independent samples. This happens in the MNIST or Let- ters datasets, where data is generated by thresholding low- resolution grayscale images. WFM sets itself apart from other methods, which are restricted to uniform datasets, by leveraging the entropic OT map's ability to compute feasi- ble transformations between empirical samples of different sizes. Our experiments in Figure 4 demonstrate that WFM generates high-quality & diverse samples, despite large vari- ability in the number of particles, itself a novel contribution.
120
122
123
Secg
Desand
Sictaa
Th+ Sst Niche
[6
16
18
CT
CT
4.2.2. SPATIAL TRANSCRIPTOMICS
Antro
Dreamel
Cocz
Lemad
salle
Cldes Sould Opanoo
In spatial transcriptomics, cell's 's niche is represented by the distribution of gene expression in its neighboring cells, which we observe through its immediate nearest neighbors in high-dimensional space. This approach is complementary to the BW representation of niche (recall Section 4.1. and serves as a more high fidelity view suited for fine-grain interactions. Due to their high dimensionality, these cellu- lar microenvironment distributions have remained beyond the reach of generative models that depend on voxel-based neural networks. Instead, WFM uses transformers, whose permutation equivariance and indifference to dimension- ality make them natural architectures for modeling these high-dimensional distributions through their point-cloud realizations (Haviv et al., 2024b).
Figure 6. Examples of generated niches for Calb2+ and Th+ Sst- neurons (see Figure 5) display expected gene-expression signa- tures, with Calb2+ niches enriched for upper-laye markers (Cux2) and Th+ for deep-layer markers (Ning2 and Foxp2). Each niche is represented as high-dimensional distribution realized through point-clouds, demonstrating WFM's s unique capability to generate distributions in spaces where previous methods, limited by vox- elization, cannot operate.
5. Conclusion and outlook
This work shows how to appropriately lift the Riemannian flow matching paradigm of Chen & Lipman (2023) to the Wasserstein space, resulting in Wasserstein flow matching. Our motivations stem from modern datasets, where each data sample is a probability distribution, necessitating this
In the motor cortex, somatostatin (Sst) interneurons are orga-
8

---

Wasserstein Flow Matching
extension for generative modeling purposes. Our contri- butions are algorithmic in nature, which incorporate vari- ous elements, such as estimating optimal transport maps via entropic OT, closed-fome expressions over the Bures- Wasserstein space, and attention mechanisms in neural net- work architectures. Our algorithm is capable of generating realistic data from Gaussian and general distribution real- ized as variable-size or high- nigh-dimensional point-clouds. Both contexts are highly relevant in single-cell transcriptomics for synthesizing of microenvironmentt and cellular states.
Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition, 2008. ISBN 978-3--7643-8711-1 Cited page 3.
Amos, B., Cohen, S., Luise, G., and Redko, I. Meta optimal transport. arXiv preprint MtXiv:2206.0526; 2022. Cited on page 18.
Atanackovic, L.., Zhang, X., Amos, B., Blanchette, M., Lee, L. J., Bengio, Y., Tong, A., and Neklyudov, K. Meta flow matching: Integrating vector fields on the wasserstein manifold. arXiv preprint arXiv:2408. 14608, 2024. Cited on page 3.
Impact Statement
Wasserstein Flow Matching designed for generative mod- eling of distributions, with applications in computational biology. As with any generative model, care should be taken when using WFM to draw biological conclusions.
arXiv preprint Ba, J. L. Layer normalization, NTlv:1607.064500 2016. Cited on page 18.
Benamou, J.-D. and Brenier, Y. A computational fluid me- chanics solution the Monge-Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000. Cited page 4.
Acknowledgements
AAP acknowledges support from NSF grant DMS-1922658. DP an Investigator at the Howard Hughes Medical Insti- tute and is supported by National Cancer Institute grants P30 CA08748 and U54 CA209975.
Bennett, J. OpenStreetMap. Packt Publishing Ltd, 2010.
Cited on page 20.
Binnewies, M., Roberts, E. W., Kersten, K., Chan, V., Fearon, D. F., Merad, M., Coussens, L. M., Gabrilovich, D. I., derstanding the tumor immune microenvironment (time) for effective therapy. Nature medicine, 24(5):541-550, 2018. Cited on page
Meta was involved only in advisory role. All experimen- tation and data processing was conducted at MSKCC.
S., Hedrick, C. C., et al. Un-
References
Boeing, G. Osmnx:
New methods for acquiring, con- structing, analyzing, and visualizing complex street net- works. Computers, environment and urban systems, 65: 136-139, 2017. Cited page 20.
Achlioptas, P., Diamanti, O., Mitliagkas, I., and Guibas, L. Learning representations and generative models for 3d point clouds supllemtup material-. arXiv preprint arXiv: 077829 2018. Cited on page 3.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milnee S., et al. Jax: Autograd and xla.
Albergo, M. S. and Vanden-Eijnden, E. Building normal- izing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571 2022. Cited on pages 1 and 3.
Astrophysics Source Code Library, pp. ascl-2111, 2021. Cited on page 18.
Altschuler, J., Weed, J., and Rigollet, Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, 2017. Cited on page 13.
Brenier, Y. Polar factorization and monotone rearrangement vector-valued functions. Comm. Pure Appl. Math., 44 (4):375-411 1991. Cited on page 3.
Bubeck, S., Chandrasekaran, V., Eldan, R., Gehrke, J., Horvitz, E., Kamar, F.,, Lee, P., Lee, Y. T., Li, Y., Lundberg, S., et al. Sparks of artificial general intel- ligence: Early experiments with GPT-4. arXiv preprint ar arXiv:2303. 12712, 2023. Cited on page
Altschuler, J., Chewi, S., Gerber, P. R., and Stromme, A. Av- eraging on the Bures-Wasserste manifold: Dimension- free convergence of gradient descent. Advances in Neural Information Processing Systems, 14:22132-22144 2021.
Cited on page 16.
Cai, R., Yang, G., Averbuch- Elor, H., Hao, Z., Belongie, S., Snavely, N., and Hariharan, B. Learning gradient fields for shape generation. In Computer EECV 2020: 16th European Conference, Glasgow, UK, August 23 28, 2020, Proceedings, Part III 16, pp. 364-381. Springer, 2020. Cited on page 3.
Amari, S.-i. Information geometry and its applications, vol- ume 194. Springer, 2016. Cited on page 3.
Ambrosio, L., Gigli, N., and Savaré, G. Gradient flows in metric spaces and in the space of probability measures.

---

Wasserstein Flow Matching
Carlen, E. A. and Gangbo, W. Constrained steepest descent in the 2-Wassersteds metric. Annals of mathematics, pp. 807-846, 2003. Cited page 17.
Eyring, L.., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z., and Theis, F. Unbalancedness in neural monge maps improves unpaired domain translation, 2024, URL https: http:///www....oaaa//3.......... on page 3.
Cited
Chang, A. X., Funkhouser, T., Guibas, L., Hanrahan, P., Huang. Q. Li, Z., Savarese, S., Savva, M., Song, S., Su, H., et al. Shapenet: An information-rich 3d model repository. arXiv preprint arXiv: 2015.C0 Cited on page
Flamary, R., Courty, N., Gramfort, A., et al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021. URL http:// jmlr. org http//ppasss//2222/ 5 and 13.
html. Cited on pages
Chen, R. T. and Lipman, Y. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660 2023. Cited on pages 1, 3, 5, 8, 15, 16, and 17.
Gadelha, M., Wang. R., and Maji, S. Multiresolution tree networks for 3d point cloud processing. In Proceedings of the European Conference Computer Vision (ECCV), pp. 118- 2018, Cited page 3.
Chen, Y., Goldstein, M., Hua, M., Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Probabilistic forecasting with stochastic interpolants and Föllmer processes. arXiv preprint arXiv:2403.1372. 2024. Cited page 5.
Gayoso, A., Lopez, R., Xing, G., Boyeau, P., Valiol- lah Pour Amiri, V., Hong, J., Wu, K., Jayasuriya, M., Mehlman, E., Langevin, M., et al. A python library for probabilistic analysis of single-cell omics data. Nature biotechnology, 40(2):163-166, 2022. Cited on page 3.
Cheng, C., Li, J., Peng, J., and Liu, G. Categorical flow matching statistical manifolds. arXiv preprint 14X4412405.6641,, 2024. Cited page 3.
Cui, H., Wang, C., Maan, H., Pang, K., Luo, F., Duan, N., and Wang. B. segpt: toward building foundation model for single-cell multi-omics using generative ai. Nature Methods, pp. 1-11, 2024. Cited on page 3.
Gelbrich, M. On a formula for the 12 Wasserstein metric between measures on Euclidean and Hilbert spaces. Mathematische Nachrichten, 147(1):185-203, 1990. Cited on page 4.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013. Cited on pages 6 and 7.
Haviv, D., Kunes, Z., Dougherty, T., Burdziak, C., Nawy, T., Gilbert, A., and Pe Pr'Er, D. Wasserstein wormhole: Scalable optimal transport distance with transformers. ArXiv, 2024a. Cited on page 20.
Cuturi, M., Srang-Paaxxanthos L., Tian, Y., Bunne, C., Davis, G., and Teboul, O. Optimal transport tools (OTT): A JAX toolbox for all things Wasserstein. arXiv preprint arXiv:2201.12322 2022. Cited on pages 5, 13, and 18.
Haviv, D., Remšík, J., Gatie, M., Snopkowski, C., et al. The covariance environment defines cellular niches for spatial inference. Nature Biotechnology, pp. 1-12, 2024b. Cited on pages 2, 7, 8, 17, and
Davis, O., Kessler, S., Petrache, M., Bose, J., et al. Fisher flow matching for generative modeling over discrete data. arXiv preprint arXiv:2405.14666 2024. Cited on page
Hütter, J.-C. and Rigollet, P. Minimax estimation of smooth optimal transport maps. The Annals Statistics, 49(2): 1166 1194, 2021. Cited on pages 5 and 14.
Deb, N., Ghosal, P., and Sen, B. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. Advances in Neural Information Cited
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zidek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-58) 2021. Cited on page
Processing Systems, 34:29736-29753, 2021. page 14.
Divol, V., Weed, J., and Pooladian, A.-A. Opti- mal transport map estimation in general function spaces. arXiv preprint arXiv:2212.03722, 2022. Cited on pages and 14.
Kim, H., Lee, H., Kang, W. H., Lee, J. Y., and Kim, N. S. Softflow: Probabilistic framework for normaliz- ing flow on manifolds. Advances in Neural Information Processing Systems, 33: 33:16388-16397 2020. Cited on page 3.
Emami, P. and Pass, B. Optimal transport with opti- mal transport cost: the Monge-Kand Kantorovich problem on Wasserstein spaces. Calculus of Variations and Partial Differential Equations, 64(2):43, 2025. Cited on page 16.
Kingma, D. Adam: method for stochastic optimization. arXiv preprint arXiv:1412. 2014. Cited on page 18.
10

---

Wasserstein Flow Matching
Klein, D., Uscidda, T., Theis, E, and Cuturi, M. Genot:
Nielsen, F. An elementary introduction information ge- ometry. Entropy, 22(10):1100, 2020. Cited page 3.
Entropic (gromov) wasserstein flow matching with ap- plications single-cell genomies, 2024. URL https: //arxiv. org/ abbs//3000 0 09254. Cited on page 3.
Nowotschin, S., Setty, M., Kuo, Y.-Y,, Liu, V., Garg, V., Sharma, R., Simon, C. S., Saiz, N., Gardner, R., Boutet, S. C.. et al. The emergent landscape of the mouse gut endoderm at single-cell resolution. Nature, 569(7756): 361-367, 2019. Cited on page 7.
Klokov, R., Boyer, E., and Verbeek, J. Discrete point flow networks for efficient point cloud generation. In European Conference on Computer Vision, pp. 694-710. Springer, 2020. Cited on page 3.
OpenAl. Dall.e-2.
https: /ttsss//ppeai.ccom com/
Lambert, M., Chewi, S., Bach, F., Bonnabel, S., and Rigol- let, P. Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, 35: 14434-14447, 2022. Cited on page 4.
dall-e-2/, 2022. 2022..... Cited on page 1.
Otto, F. The geometry of dissipative evolution equations:
The porous medium equation. 2001. Cited on page 3.
Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh, Y. W. Set transformer: A framework for attention based intnutaationinvariian neural networks. In International conference on machine learning, pp. 3744-3753. PMLR, 2019. Cited on page 2.
Panaretos, V. M. and Zemel, Y. The Wasserstein Space, pp. 37-57. Springer International Publishing, Cham, 2020, ISBN 978-3-030-38433 38438-8. Cited on page 14.
Persad, S., Choo, Z.-.,, Dien, C., Sohail, N., Masilionis, I., et al. Seacells infers transcriptional and epigenomic cellular states from single-cell genomics data. Nature Biotcchnology, (1(1)::17111661777333 2023. Cited on pages 1, 7, and 19.
Lipman, Y., Chen, R. T., Le, M. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. Cited on pages 1, and 15.
H., Nickel, M., and
Peyré, G. and Cuturi, M. Computational optimal transport. Foundations and Trends@ in Machine Learning, 11(5-6): 357-607, 2019. Cited on page 13.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 000000, 2022. Cited on pages and 3.
Pooladian, A.A. and Weed. J. Entropic esti- mation of optimal transport maps. arXiv preprint arXiv:2109. 12004, 2021. Cited on pages 2, 5, 13, and 14.
Lohoff, T., Ghazanfar, S., Missarova, A., Koulena, N., Pier- son, N., Griffiths. J. A., Bardot, E. S., Eng, C.-H., Tyser, R. C., Argelaguet, R., et al. Integration of spatial and single-cell transcriptomic data elucidates mouse organo- genesis. Nature biotechnology, 40(1):74-85, 2022. Cited on pages and 19.
Pooladian, A.-A., Cuturi, M., and Niles Weed, J. Debi- aser beware: Pitfalls of centering regularized transport maps. arXiv preprint MiXiv:220..08919 2022. Cited on page 14.
Lopez, R., Regier, J., Cole, M. B., Jordan, M. I., and Yosef, N. Deep generative modeling for single-cell transcrip- tomics. Nature methods, 15(12):1053-1088 2018. Cited on pages and 3.
Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich C., Amos, B., Lipman, Y., and Chen, R. T. Multisample flow matching: Straightening flows with minibatch cou- plings. arXiv preprint arXiv:2304. 13472200..4472.. 2023a. Cited on page 17.
Manole, T., Balakrishnan, S., Niles-Weed J., and Wasser- man, L. Plugin estimation of smooth optimal transport maps. arXiv preprint arXiv:2107.12366 2021. Cited on page 5.
Pooladian, A.-A., Divol, V., and Niles-Weed J. Minimax estimation of discontinuous optimal transport maps: The semi-discrete case. arXiv preprint arXiv:2301. 2023b. Cited on page 14.
Mena, G. and Niles-Weed. J. Statistical bounds for entropic optimal transport:
sample complexity and the central limit theorem. Advances in Neural Information Processing Systems, 32, 2019, Cited page 13.
Rigollet, P. and Stromme, A. J. On the sample com- plexity of entropic optimal transport. arXiv preprint Ttlv:2206.134722 2022. Cited on page 14.
Muzellec, B., Vacher, A., Bach, F., Vialard, F-X., and Rudi, A. Near-optimal estimation of smooth trans- port maps with kernel sums-of of-ns-ss-quures.... arXiv preprint arXiv:2112.01907, 2021. Cited on page 14.
Santambrogio, F. Optimal transport for applied mathemati- cians. Birkäuser, NY, 55(58-63):94, 2015. Cited on page 16.

---

Wasserstein Flow Matching
Scetbon, M., Cuturi, M., and Peyré, G. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pp. 9344-9354. PMLR, 2021. Cited on page 18.
Wu, L.., Wang. D., Gong, C., Liu, X., Xiong, Y., Ranjan, R., Krishnamoorthi, R., Chandra, V., and Liu, Q. Fast point cloud generation with straight flows. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9445-9454, 2023a. Cited on pages 3, 7, and 20.
Sinkhorn, R. A relationship between arbitrary positive matrices and doubly stochastic matrices. The annals of mathematical statistics, 35(2):876-879, 1964. Cited on pages and 13.
Wu, S. J., Sevier, E., Dwivedi, D., Saldi, G.-A., Hairston, A., Yu, S., Abbott, L.,, Choi, D. H., Sherer, M., Qiu, Y., et al. Cortical somatostatin interneuron subtypes form cell-type-specific circuits. Neuron, 111((7):2675-69) 2023b. Cited on page 8.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011. 131466, 2020. Cited page 1.
Wu, Z., Song, S., Khosla, A., Yu, F., Zhang, L.,, Tang, X., and Xiao, J. 3d shapenets: A deep representation for volu- metric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912-1920, 2015. Cited on page 7.
Stark, H., Jing, B., Wang, C., Corso, G., Berger, B., Barzi- lay, R., and Jaakkola, T. Dirichlet flow matching with applications to DNA sequence design. arXiv preprint arXiv:2402.0584) 2024. Cited on page 3.
Xing, Z., Feng, Q., Chen, H., Dai, Q., Hu, H., Xu, H., Wu, Z. and Jiang, Y.-G. survey on video diffusion models. ACM Computing Surveys, 2023, Cited on page
Stephenson, E., Reynolds, G., Botting, R. A., Calero-Nieto, F. J., Morgan, M. D., Tuong, Z. K., Bach, K., Sung- nak, W., Worlock, K. B., Yoshida, M., et al. Single-cell multi-omics analysis of the immune response in covid-19 19. Nature medicine, 27(5):904-911 2021. Cited on page
Yang, G., Huang, X., Hao, Z., Liu, M.-Y., Belongie, S., and Hariharan, B. Pointflow: 3d point cloud generation with continuous normalizing flows. In Proceedings the IEEE/CVF international conference on computer vision, pp. 15411-450, 2019. Cited on page 3.
Theodoris, C. V., Xiao, L., Chopra, A., Chaffin, M. D., Al Sayed, Z. R., Hill, M. C., Mantineo, H., Brydon, E. M., Zeng, Z., Liu, X. S., al. Transfer learning enables predictions in network biology. Nature, 618(7965):616- 624, 2023. Cited on 3.
Zemel, Y. and Panaretos, V. M. Fréchet means and Pro- crustes analysis in Wasserstein space. 2019. Cited on page 16.
Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector- Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. Im- proving and generalizing flow-based generative mod- els with minibatch optimal transport. arXiv preprint arXiv:2302.00482 2023. Cited on page 17.
Zhang, L., Tozzo, V., Higgins, J., and Ranganath, R. Set norm and equivariant skip connections: Putting the deep in deep sets. In International Conference on Machine Learning, pp. 15559-265744 PMLR, 2022. Cited on page 18.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017. Cited on pages 2 and 18.
Zhang, M., Eichhorn, S. W., Zingg, B., Yao, Z., Cotter, K., Zeng, H., Dong, H., and Zhuang, X. Spatially resolved cell atlas of the mouse primary motor cortex by merfish. Nature, 19877879) 2021. Cited on page 8.
Villani, C. Optimal transport: old and new, volume 338. Springer, 2009, Cited on pages? and 16.
Zhou, L., Du, Y., and Wu, J. 3d shape generation and completion through point-voxel diffusion. In Proceedings the IEFE/CVF international conference on computer vision, pp. 5826-5835, 2021. Cited on page 3.
Vyas, A., Shi, B., Le, M., Tjandra, A., Wu, Y.-C., Guo, B., Zhang, J., Zhang, X., Adkins, R., Ngan, W., al. Au- diobox: Unified audio generation with natural language prompts. arXiv preprint 158iv22312.158211 2023. Cited on page 1.
Wang, W., Ozolek, J. A., Slepčev, D., Lee, A. B., Chen, C., and Rohde, G. K. An optimal transportation approach for nuclear structure-based pathology. IEEE transactions on medical imaging, 30(3):621-63 2010. Cited on page 14.
12

---

Wasserstein Flow Matching
A. Entropic estimation of OT maps
We briefly discuss how to estimate optimal transport maps between distributions using entropic OT. We refer the interested reader Pooladian Niles-Weed (2021) for more information on this approach.
We first outline the numerical aspects of the approach; follow Peyré & Cuturi (2019). Let μ E; m and / -ー Σ, n be two empirical distributions where X (21,..., Y , Wn We first define the following
polyhedral constraint set
Un.n = P EPPERTA RMXX : Plm = m '1m PT 1, = 1,)
which represents the possible couplings between the two discrete measures. The entropic optimal transport coupling between the two discrete measures f and is defined as the minimizer to the following strictly convex optimization problem
P = argminnC,, P) cH(P),
(12)
P@Um,
where e > 0, = Eig Pi,j(lo((P..))
-11, 1), and Crj = | - リリ 121l.. Sinkhorn's matrix scaling algorithm (Sinkhorn, 1964) makes it possible solve for ps with runtime of ((m//e) (Altschuler al., 2017). We briefly stress three points:
=
1. The coupling P is not permutation matrix. The coupling lies inside the polytope Um,, and not the vertices, and therefore is not a permutation matrix.
2. When e = 0, the objective becomes a standard linear program with a runtime of Ómmn(m + ni)) (up to log factors) (Peyré & Cuturi, 2019, Chapter 3). While we include CPU implementation (Flamary et al., 2021) in the WFM codebase, this approach lacks GPU efficiency and substantially increases training time, making it impractical for most use cases.
3. Instead, the regularization parameter e serves tunable training hyperparameter. Lower e values better approximate true the optimal transport map but require more Sinkhorn iterations for convergence, creating direct trade-off between accuracy and computational efficiency.
In all our experiments, we used the open-source package OTT-AX (See http://ott.aarraatttteoosss//ees compute the entropic coupling and the out-of-sample mapping (Cuturi et al., 2022).
A.1. Rounded matchings
Our first approach holds when n. In this case, greedily round the noisy matching matrix P to become 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
primal-dual relationship of the strictly convex program (12) shows that there exist vectors (f",g") E R x R" such that
P. = - G
These two vectors called the Kantorovich potentials, which are initially defined the support of and v, respectively. However, they can be readily extended to all of Rd (Mena & Niles Weed, 2019), resulting in two functions
f(x) - log
exp((e) - -1x -v,111//))
m
((v) ee log
exp((Si -11y-1;111 =x,13)/ss)
13

---

Wasserstein Flow Matching
Following Pooladian & Niles-Weed (2021), we can define the entropic transport map, where the last equality is simple calculation:
وو= By (g) دj=1 exp
- ;)(1))/
((=)==-v[ia
(13)
1 V/J (111)//)
This estimator was initially to provide statistical upproximations to the optimal transport map T on the basis of samples; see Pooladian Niles-Weed (2021); Rigollet Stromme (2022); Pooladian et al. (2023b; 2022). Note that T3 (x) can be interpreted as the conditional expectation of the plan P conditioned on out-of-sample inputs æ E R which is well-defined due to the relations above. Finally, stress that this estimator be adapted to settings where the point-clouds μ and not only have different numbers of points, but also non-uniform weights. As this estimator is also by-product of Sinkhorm's 's algorithm, it is also scalable and GPU-friendly.
A.3. On the (statistical) approximations of geodesics
We briefly collect basic results pertaining to the (statistical) approximation of optimal transport paths. This bound shows that the error grows along the trajectory, but is limited by the overall distance of the maps.
Proposition A.1. Let Pt, v be two probability measures and suppose y has a density, and let T be the optimal transport map from p to v. Let T be an estimator to the optimal transport map, defined with respect to data X1y...5 Xn ~ p and Yi....,Y Y, v. Then for t [0,1]
T (P2 S S 1=
where the outer expectation is taken with respect to the data, and we define
Pr = ((1 - t)id Mp P: = ((1- t)id + tf
Proof. The result follows immediately from a standard coupling argument to obtain the linearized Wasserstein distance (Wang et al., 2010; Panaretos & Zemel, 2020)
W3 (pc, Pe) s [I +)id+
+ (T*) ((1- t)id +T) (B= = 2 T 1
The entropic Brenier map is one particular estimator. We note two key properties of this map; see Pooladian Niles-Weed (2021) for in-depth discussions.
Theorem A.2. Suppose μ. have density bounded above and below, and that the optimal transport map between them, denoted T*, is such that (T*)-1 is at least twice differentiable and there exists X, A >> such that
X DT 3A.
Then, when estimated from samples from and samples from μ. the entropic Brenier map has the following error
E - T2 S
1/2 1/2 log(n)== dd/-11 = +
(14)
where we suppress constants that depend our assumptions. Performing bias-variance trade-off in the regularization parameter, one obtains (n.) x n -1((44))) and the total error becomes
Elit. - -T*112 Stog(n) n. 22(d+44)
We emphasize that the assumptions in Theorem A.2 are standard in the literature (Hütter & Rigollet, 2021; Deb et al., 2021; Muzellec et al., 2021; Divol et al., 2022). While the rate scales exponentially poorly with the dimension, we stress that existing lower bounds of estimation (see Hütter & Rigollet (2021)) also suffer from the curse of dimensionality. which unavoidable for this task. Combining these two results, we can compare the geodesic given by the OT map, and the one induced by using the entropic map, where we write
eE = (11 (1--tich 1jid + (T.).u.
Corollary A.3. Consider the same setting as Theorem A.2. Then the geodesic given by the estimated entropic Brenier map, denoted by Pe, on the basis of samples and x n 1/(444) is close to the true geodesic with the following error
E(W, (Pt. ( S ftt 2/(d+1)
14

---

Wasserstein Flow Matching
B. Derivation of the WFM objective
In this section, we give validity to the WFM for optimization purposes, and our choice curves. For instance, recall the original Flow Matching objective (Lipman al., 2022)
(F = L/
(x()))l
(15)
where fo R x [0, 1] + Rd a neural network, and (pe, 1e)ee10.1] are a density field pair that satisfy the continuity equation between two distribution μ and
Öp (pen)) 0, Po = p.
Note that, implicit (F) is the endpoint constraints μ and Now, we average over possibly choices of μ po and Pr, resulting in
Lo Ex.. -x,(X)] dpo(1) (v) dt.
(16)
As particular case, take (pt, uz) < (ju, Ue), where the first argument is the McCann interpolation between μ and v, and us is the optimal velocity field, which is a function of the optimal transport map from j to (recall Section 2.3). This yields our final objective (11), which we recall here for convenience
LWFM (0) = L.' /f Ex.xmillooxx.. 2x.)))ll dpi (v) dt.
(17)
When Po: Pi are distributions over Gaussians, we have closed-form expressions for objects of interest. When Po, P1 are distributions over general distributions, we approximate the geodesics based on point-cloud samples using entropic Brenier maps and their respective interpolations. We emphasize that the oounded-matching employ (see Appendix A.1) is also valid curve.
B.1. Validity of Conditional Flows
The key idea behind Flow Matching is that the (tractable) conditional velocity probability paths come together to form correct marginal velocities and paths from source distribution target. In Riemannian Flow Matching, given source and target ; Po E (PM) distributions over Riemannian manifold M. the conditional velocity field from Po to sample 11 is generated by flowing from each point ~ Po towards y using the geodesic path between them. This produces the conditional probability path p.(-Iv) generated by the vector field u.(-))), which obey the continuity equation for Riemannian manifolds:
(-p.(-P) + div))(P)(())))((()))
(18)
where div M is the Riemannian divergence operator. Following Theorem in Lipman al. (2022) and equation 7 in Chen Lipman (2023), the marginal vector field:
(e(e) = 1 u (Helk) Pe(Hell) (v) pr((e)
(19)
generates the marginal probability path:
pr(Ne) = S pe(uelv) dpr
(20)
as (14) and (144 together obey the continuity equation. The proof requires regularity constraints on pr (p(lp) and u, Me (V), which we assume as in Chen & Lipman (2023). A withstanding requirement that we need is that a unique optimal transport map exists between any two measures (ji, v) E Po x P1- This can be ensured if, for example, all measures Po
15

---

Wasserstein Flow Matching
have density, which is assumed throughout the text. Positivity of the path 14 can be ensured through myriad of conditions (e.g., source and target measures have positive density over the same support). In the explicit Gaussian-to- Gaaussaan case (i.e., the Busss-Wasserstiin manifold), this assumption is trivially satisfied as the manifold is genuine finite-dimensional Riemannian manifold and the Riemannian machinery of Chen & Lipman (2023) goes through with no modifications. We carry out the rest of our computations for general distributions over the Wasserstein space, and hence the computations are non-rigorous. For a fully rigorous treatment, see Emami & Pass (2025).
CONDITIONAL FLOWS AND PATHS FOR WFM
To continue, we require the following proposition. Recall that pre-metric is a function d: M M R that satisfies the following three requirements
1. d(z, v) 2 0 for all I, E M,
2. d(z, y) 0 if and only if x y,
3. V2d(I, v) * 0 if and only if + y.
Lemma B.1. The 2- 2-Wasserstein distance defines valid pre-metric over P2 (R²))
Proof. As the 2- Wasserstein distance is metric (Villani, 2009; Santambrogio, 2015), the first and second conditions of pre-metric are trivially satisfied. For the third condition, we note that for 14, (2
V V (1)) =- log
2(T - -dd)
(21)
For a proof of this computation, see Altschuler et al. (2021) or Zemel & Panaretos (2019). By chain-rule, we have that
V2 (4) W2(n, VW2( = log (v)
This latter quantity is only zero if and only if μ v. and thus the proof is complete.
We now recall how construct conditional vector fields function of pre-metrics d, described by Eq. (13) by (Chen Lipman, 2023). The construction is the following
uy (xrll) d In(1 t)) V) where is the metric on the tangent space at xz, and is a geodesic in the form
111 ((
(22)
expx, ((1 t) log.ro (21)).
Note that our geodesics are precisely of the form
exp, - t) log,
(v)),
where we recall - (x) (id v)zfa and + log,, (v) Tu- - id. We now specialize these calculations the case where the pre-metric is the 2- 2-Wasserstein distance.
Lemma B.2. Fix V replaced by the 2- 2-Wasserstein distance, evaluated at geodesic between any p ~ Po and the fixed v. Then it holds that ue ue(14 ()) = Ub where v is given by Equation (7) between I and v.
Let μe + ue(uiv)) denote the conditional vector field defined using Equation (22) with d
Proof. We compute, using Equation (21),
ue(p)))) = ( ( In(1 (WW,, (11P())))
1
(Tu - id)
t
-1
(Tu-t (T:) id)
(Tu-oV Tx) 0 (T:)-1
16

---

Wasserstein Flow Matching
where we make use of the fact that Th TW (TH+ (see e.g., Section 2 of Carlen & Gangbo (2003)), and that
T. (1 t)id LTu-V. Simplifying, obtain
ue ue(pel)) 1 (TH - (1- 1)id ) (T:)- - (TK - id) (T:)-*
(23)
which is precisely the optimal transport vector field v2 stated (7).
Together, we arrive at the following corollary, which validates the probability paths generated by our approach.
Corollary B.3. The conditional flow p+(-)) generated by the conditional vector fields defined by (23) satisfy pi (-11))
Proof. This follows immediately from Theorem 3.1 of Chen & Lipman (2023) by passing through the pre-metric construction above.
C. 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 (Pooladian et al., 2023a; Tong et al., 2023). The core technique behind MS-FM is to 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 E {l...., Bsz}. In the BW-FM setting, given two sets of Gaussians {(a;, P2) and {(b,. B,) |ei. their Frechét (W2) distance matrix is:
C₂2 = -a byl1 Tr(A, + B, - B,
(24)
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 (see Appendix A.I), ensuring the entire batch is used in training.
For general WFM, applying MS requires computing pairwise OT distance between all source and target samples within minibatch. For large point-cloud samples, this is exorbitantly expensive, even with Sinkhorn iterations. For efficient approximate, here too we rely on the Frechét distance, computed between empirical means and covariances of cach point- cloud. Computation the Frechét distance is markedly less resource-intensive than entropic OT, yet is notably correlated with EMD values (see Table in Haviv et al. (2024b)).
D. Choice of source measure
In both WFM and BW-FM, learning flows requires source measure which is straightforward to sample from. For source distribution on the space of {(m, E) m E Rd Σ E 84 J. 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 general distributions, it essential that the initial (source) measure be diverse, rather than collapsed and degenerate. Indeed, while is alluring produce noise point-clouds by sampling particles from single base distribution, i.c. X | z; ~ (0, as n grows, the Wasserstein distance between empirical distributions goes to 0. To alleviate this, we draw point-clouds from multivariate Gaussians with a stochastic covariance:
L ~ N I)
X f N(O,
where #L & σL are the average and standard deviation of the Cholesky factors from the empirical covariances of the target measure point-clouds. This ensures wider source measure, producing diverse range of noise point-clouds.
17

---

Wasserstein Flow Matching
Training Data
Frobenius FM
BW-FM (Euclidean)
BW-FM (Riemannian)
Degenerate Gaussians
Figure SI. When the number of training examples is too few (in comparison to Figure 2), all methods collapse on the training data, though only the Riemannian instantiation of BW-FM captures the covariances perfectly.
E. Neural architecture & training
E.1. BW-FM on Gaussians
The goal of BW-FM is to train neural network match the (Riemannian) time-derivative along the BW geodesics between Gaussians. The model employs standard, fully connected neural network which takes input concatenated values of (m,, Mx,t) t) based 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, 2017). To streamline training, two separate networks are employed: one to match the time derivative of the mean rh, and another for the time derivative the covariance matrix SBWW The BW tangent norm is used as the loss function for training these networks.
By default, all models use 6-layer neural network using relu non-linearity, with 1024 neurons per layer. applying skip connections and layer-norm (Ba, 2016). Training is performed for 100, 000 gradient descent steps using the Adam optimizer (Kingma, 2014) with exponential learning rate decay of 0.97 every 1000 steps and batch size of 128.
E.2. WFM general distributions
WFM estimates the optimal transport (OT) map for given pair of interpolate point-cloud and time (X, Here the time component t is first converted into Fourier features. The model's architecture begins with embedding layer, followed by series of alternating multi-head attention and fully-connected layers. Skip connections and layer-norm are applied after each operation. The final layer projects the embeddings back to X's original space using 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 constructed with regularization weight of 0.002 and 200 Sinkhorn iterations, which found to be sufficient for convergence. Whenever the dataset consists of uniformly sized point-cloud realizations, apply the out-of-sample estimator (see Appendix A.2) which can
rounded matching (see Appendix A.I), otherwise 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. Our model is optimizer with Adam (Kingma, 2014) using an exponential learning rate decay and batch size of 64.
WFM relies on JAX and OTT-JAX (Bradbury et al., 2021; Cuturi et al., 2022) and enjoys seamless optimization via end-to-end just-in-tine compilation. For the ShapeNet experiments (see Table 3), the model is trained for 500,000 training steps, totaling about days of trainings of a single A100 GPU. All other experiments (see Table 4) trained for 100,000 steps, requiring around 3-4 hours of GPU use. We note the Transformer's 's forward and backwards pass was the most significant source of computational overhead, as opposed to the Sinkhorn based approximation of OT maps. The computational complexity of both components is quadratic with point-cloud size, which limits the scalability of WFM. Making Transformers simpler to optimize and accelerating OT computations are both active areas within ML research (Amos et al., 2022; Scetbon et al., 2021; Zhang et al., 2022), and future iterations of WFM can incorporate solution from those spaces.
18

---

Wasserstein Flow Matching
Organ
BW-FM
Gut tube
Spatial Pattern
Thymus
Lung
Liver
Intesting
BW UMAP
Spatial
BW UMAP
Spatial
Figure S2. Comparing real and generated cellular microenvironmens embryonic gut tube development. Left: UMAP embedding of BW-FM generated (left panel) and observed (middle panel) microenvironments based Bures- Wasserstein distance, colored by developing organ type. Right: Spatial arrangement of cells in the seqFISH embryo showing physical organization of organ development.
F. Experiment Details
F.. Spatial Transcriptomics
In our manuscript, applied WFM and BW-FM on several spatial transcriptomics datasets, encompassing variety of technologies and tissue contexts. From the 351-gene seqFISH embryogensis dataset, (Lohofff et al., 2022), we focus on niches of gut-tube cells. We compress gene-expression profiles down to their 16 principal components (PC) and aggregate all the cells around each gut cell within an 80 micron radius, yielding on average 29 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 delineate budding organs into their appropriate regions (see Figure S2).
MetaCells
Conditional BW-FM
Cell Type
B_cell
CD16
CD14
DCs
Platelets
Plasmablast
HSC
CDB
CD4
gdT
BW UMAPY Y
RBC
Lymph_pront
pDC
Treg
MAIT
BW UMAP
Figure S3. Conditional BW-FM applied to single-cell RNA sequencing data of immune response to COVID-19 19. Large scale single-cell atlases commonly grouped into highly dedicated clusters called MetaCells (Persad et al., 2023). In this application, BW-FM 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 label accuracy 93.13%.
In complementary approach, WFM is applied directly on gene-expression based point-clouds of niches, representing realization of the continuous distribution of their environments. Uniquely suited for high-dimensional data, we apply WFM on a MERFISH atlas of the motor cortex XENIUM assay melanoma metastasis to the brain (Haviv et al., 2024b). In both dataset, we select the k 16 physical nearest neighbours of every cell, and aggregate expression profiles based on their first
19

---

Wasserstein Flow Matching
16 PCs.
We concentrated on Sst interneurons in the motor cortex datas, which divided into spatially segregated, phenotypically distinct subtypes. Applied unconditionally, WFM generated niches match the distribution of the real data based on EMD and CD mnearstt-neihhhon accuracy (see Table 4). We then assessed WFM's capability to comprehend the relationship between cell state and environment and tasked it with conditional generation based on subtype label. Based on OT distances estimated via Wormhole embeddings (Haviv et al., 2024a), distributions generated by WFM recapitulated true organ environment. The label accuracy for WFM- -eeeereed data was 63.86%, which was nearly identical the test-set real data accuracy 62.19%.
6
Figure S4. Cities Dataset. Gaussians representing the 100 populous cities in the continental US. The data obtained from (Bennett, 2010) via OSMix (Bocing, 2017). The mcan parameter is the longitude and latitude coordinate of cach city and the covariance the 2nd moment approximation their metro area.
G. 2D & 3D shapes
The ShapeNet dataset consists of 3D shapes of 55 different classes, Emulating the benchmarking effort in (Wu al., 2023a), apply WFM generate 2048 sized examples from the plane, and chair classes. At each gradient descent step, sample 64 shapes from the training set for each class, and randomly select 2048 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.
Interpolation
Completion
1=0.00
1=0.25
1=0.50
1=0.75
1=1.00
Figure Interpolation and shape completion with WFM. Top. Using the lamps and handbags as source and target measures, WFM learns transform given (unseen test-set) lamp point-cloud into valid handbag. Bottom. Trained generate full planes, WFM reconstruct complete point-clouds from partial views of test-set samples.
ModelNet has 40 classes of shapes, each consisting of 2048 particles. Conditioned on class label, WFM is trained to
20

---

Wasserstein Flow Matching
generate 1000 sized point-clouds here too. In this setting, the noise measure is the standard normal and did multi-sample matching. According 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 a widely used collection of handwritten digits, consisting of 28x28 pixel grayscale images of the numbers 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 distributions, we threshold each image and extract the coordinates of the above-threshold pixels. This produces cohort of point-clouds of variables sizes, as each image contains different number of relevant pixels. We apply the entropic OT map (see Appendix A) based WFM to synthesize distributions of the digit 4 and letter a. Despite the data heterogeneity, WFM produces realistic examples (see Figure 4), while capturing the data distribution (see Table 4).
21


Reviews:
[{"review_id": "oqJe2UXu5O", "reviewer": "Reviewer_8NGf", "summary": "This paper proposes Wasserstein flow matching, a generative model which can be treated a variant of Riemannian Flow Matching on Wasserstein manifold. This model allows for working with the families of distributions and is tested in a variety of experiments.\n\n## **After the rebuttal.**\n\nI kept my initial positive score and explained the reasons in my answer to the authors.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer 1 (**8NGf**)\n\nThank you for your thoughtful review and constructive feedback.\n\n**\"Overall, the proposed method and evaluation criteria make sense. However, the paper is lacking an experiment which shows that the proposed approach beats its competitors in the case of non-Gaussian and at least moderate-dimensional experiment.\"**\n\nOur updated **Table 3** (see response to  reviewer 3 - **ToSj**) now demonstrates that WFM outperforms competing methods on 2/3 ShapeNet datasets according to the CD metric. Importantly, WFM achieves these results with approximately 70% lower computational requirements (~120 GPU hours vs. ~400 GPU hours for PSF and PVD), making it substantially more efficient.\n\n**\"The main issue here is the gap between the established theoretical results which where classic OT problem is considered and practical implementation of the approach which utilizes the solutions of the entropy-regularized OT problem (in the experiments with general distributions).\"**\n\nWe acknowledge this important point. Computing exact OT maps between general distributions is not possible in most circumstances. Outside of univariate measures, product measures, and multivariate Gaussians, closed-form expressions of optimal transport maps are not known to us. Entropic OT provides the only feasible approach with both computational efficiency and statistical guarantees; see [Pooladian and Niles-Weed (2021)](https://arxiv.org/abs/2109.12004). We added a clarifying comment in the main text to address this approximation gap. In short, an approximation scheme (like entropic OT) is consistent with other geometric methods (including RFM) which rely on numerical approximations when exact solutions aren't available for complex geometries (see **Algorithm 1** in [Chen & Lipman, 2024](https://arxiv.org/pdf/2302.03660)).\n\nThe new text now reads (line 248, revision in **bold**):\n\nThe condition of \"inner continuity\" is fairly mild, as this is ensured for any distribution $\\mu ∼ p_0$ with density. For Gaussian distributions, inner continuity holds naturally. For general distributions, we assume continuity but work with point-clouds as empirical realizations to approximate OT maps with statistical guarantees (see Appendix A.3) and computational efficiency (Flamary et al., 2021; Cuturi et al., 2022). **We note there exists a gap between our theoretical results which consider classic OT and our practical implementation which uses entropy-regularized OT approximations. This approach aligns with other geometric methods (see Algorithm 1 in Chen & Lipman, 2024) that rely on numerical approximations when exact solutions are not tractable.** The \"outer continuity\" condition is purely technical, and it serves the same role as in prior work. Our training algorithm is described in Algorithm 1, and Appendix E contains precise details on our neural network design.\n\n**\"Is it possible to conduct the experiments with point-clouds in higher dimensions and assess the method's performance w.r.t. the other baselines? I understand that many baselines can not operate in high dimensions but maybe such kind of comparison is possible for moderate dimensions (>2d,3d)?\"**\n\nWe appreciate this thoughtful suggestion. We explored adapting other methods to moderate dimensions (4-10D), but the computational requirements proved prohibitive as even our most efficient implementations would require several weeks of computation time per experiment. Instead, we've focused on demonstrating WFM's effectiveness through **Table 4**, which shows ~60% 1-NN accuracy across all metrics (with 50% being optimal), and **Figures 4-6**, which showcase realistic generation of variable-sized and high-dimensional distributions. \n\n**\"Experiments related to biology also do not clarify this point for me since I am not an expert in this field\"**\n\nWe would like to elucidate the biological application of WFM in spatial transcriptomics, a technique that combines molecular profiling of gene-expression with spatial localization of cells within tissues. In these applications, cellular microenvironments are naturally represented as distributions in gene-expression space based on neighboring cells within a tissue. WFM enables generative modeling of cellular niches, synthesizing biologically plausible microenvironments for analysis and study.\n\nBy modeling the relationship between cell phenotypes and their environments, WFM could help researchers better understand potential tissue-level responses to cellular changes. WFM addresses the inherently distributional nature of microenvironments by lifting Flow Matching to the Wasserstein space, making it well-suited for this application.", "claims_and_evidence": "Overall, the paper is well-written and most of the claims are well-supported.", "methods_and_evaluation_criteria": "Overall, the proposed method and evaluation criteria make sense. However, the paper is lacking an experiment which shows that the proposed approach beats its competitors in the case of non-Gaussian and at least moderate-dimensional experiment.", "theoretical_claims": "I skimmed through the provided proofs. The main issue here is the gap between the established theoretical results which where classic OT problem is considered and practical implementation of the approach which utilizes the solutions of the *entropy-regularized* OT problem (in the experiments with general distributions). While the authors show that error of estimating OT maps with entropic Brenier maps converges to zero with the increase of the number of samples (Appendix A.3), this issue should be directly stated in the main text, possibly in the limitations subsection.", "experimental_designs_or_analyses": "Experiments are valid but have some issues:\n- Practically, the method demonstrates strong performance mainly for distributions which can be approximated by Gaussian families. In point cloud experiments, the method actually shows a fair performance in low dimensions (3d), while in high dimensions there is a lack of baseline comparisons which makes it difficult to assess its effectiveness. Experiments related to biology also do not clarify this point for me since I am not an expert in this field.", "relation_to_broader_scientific_literature": "The paper applies the Riemannian flow matching approach to the Wasserstein space. The theoretical results are strongly connected with the original paper (Chen, R. T. and Lipman, Y, 2023).\n\nChen, R. T. and Lipman, Y. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660, 2023.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths.** The paper is well-written, has a nice structure and combines theoretical and experimental results. It proposes an approach which can be treated as an application of Riemannian flow matching to Wasserstein space, is justified theoretically and has interesting practical properties. \n\n**Weaknesses.** I was reviewing this paper at one of the previous conferences. Overall, the paper became a lot better since many important theoretical aspects are now clarified. Previously, I was expressing my concerns regarding the lack of tractable experiments in moderate or high dimensions where the proposed approach shows superior results. These concerns still hold since new experiments were not added.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- Is it possible to conduct the experiments with point-clouds in higher dimensions and assess the method's performance w.r.t. the other baselines? I understand that many baselines can not operate in high dimensions but maybe such kind of comparison is possible for moderate dimensions (>2d,3d)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8w6S0CBQBN", "reviewer": "Reviewer_C2Yo", "summary": "This paper proposes Wasserstein Flow Matching, a method for building a generative model where the datapoints themselves are distributions. The basic idea is to treat the Wasserstein space as a manifold and use Riemannian flow matching techniques. The model is studied in two settings: (1) the Bures-Wasserstein space, and (2) empirical measures / point clouds. The method is evaluated on synthetic data, 3D point clouds, and some applications in single-cell genomics and spatial transcriptomics.\n\n## update after rebuttal\nAfter the author's rebuttal, I will maintain my positive score of weak accept. I think the core idea of the paper is interesting and novel, but stronger theoretical and/or empirical results would be necessary to further strengthen the paper and raise my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer 2 (**C2Yo**)\n\nThank you for your review and insightful comments. We address each of your concerns below:\n\n**\"The authors claim that existing approaches for 3D point cloud generation cannot handle point clouds of variable sizes. However, a method like PSF for example should in principle be able to do this, as long as a suitable architecture is selected.\"**\n\nPSF fundamentally relies on 'Euclidean' distance between point clouds ($X_1 - X_0$), which only exists for fixed-size point clouds. Such an interpolation between point-clouds of different sizes is not possible. While one might hypothetically modify PSF to handle variable sizes, such modification would likely require implementing some form of optimal transport to align points between differently realised distributions, precisely what WFM does by design. \n\n**\"The use of Equation 11 / Proposition 3.1 do not seem to be properly justified, at least in the general setting. In particular, Riemannian flow matching makes an important assumption that the manifold is finite dimensional.\"**\n\nThis is an important assumption, which we do indeed make. In the papers you mentioned (such as [Functional Flow Matching](https://arxiv.org/pdf/2305.17209)), the authors simply assume that a notion of continuity equation holds at the level of the flow defined over functionals. We readily adopt this approach and assume the continuity equation holds in the Wasserstein space as well. Another possible interpretation of can be inferred from **Appendix B** --- we are simply averaging the original flow matching loss (not the conditional flow matching loss) over pairs of measures.\n\n**\"The proposed method has a fairly high computational cost, as each training step requires solving an OT problem\"**\n\nThe reviewer is correct that there is a computational complexity incurred due to the OT solver. However, as **Table 3** shows (see response to  reviewer 3 - **ToSj**) , WFM achieves competitive performance while requiring significantly less compute than competing methods such as PVD and PSF.\n\n**\"The relationship between the submission and Meta Flow Matching could be described in more detail.\"**\n\nWe expand upon the comparison here. The main difference is that  Meta Flow Matching (MFM) requires paired couplings of distributions on the Wasserstein space while WFM operates with unpaired couplings. WFM is a general generative framework for creating new samples from a learned distribution of distributions, while MFM specifically learns transformations between paired distributions.\n\n**\"Equation 11 is critical to the proposed method; dedicating more space to the derivation and justification of this in the main paper might be useful.\"**\n\nWe provide a complete derivation of **Eq 11** in **Appendix B**. We now dedicate more background on its motivation in the main text to make it clearer to the reader.\n\n**\"There appear to be some typos in Section 2.3.1...**\n\nFixed, we thank the reviewer for spotting this typo.", "claims_and_evidence": "There are two claims that I feel are not entirely justified.\n\n- The authors claim that existing approaches for 3D point cloud generation cannot handle point clouds of variable sizes. However, a method like PSF for example (https://arxiv.org/abs/2212.01747) should in principle be able to do this, as long as a suitable architecture is selected. The ability to handle variable-size point clouds seems to have less to do with the proposed methodology than with the specific architectural choice (in this case, just a transformer).\n\n- The use of Equation 11 / Proposition 3.1 do not seem to be properly justified, at least in the general setting. In particular, Riemannian flow matching makes an important assumption that the manifold is **finite dimensional**. In this setting one is thus generally justified in using densities, e.g., those appearing in Equation 19 of this submission. However, in the setting of this paper, the Wasserstein manifold is in general an infinite dimensional space, where the derivations in Appendix B no longer work without significant extra steps.\n   - While this is fine in the Bures-Wasserstein case (as this is a finite dimensional submanifold), this paper does not seem to have solved the problem in as much generality as it claims.\n   - See, e.g., recent works on diffusions and flows in infinite-dimensional spaces [[1](https://arxiv.org/abs/2302.07400)] [[2](https://arxiv.org/abs/2305.17209)] [[3](https://arxiv.org/abs/2302.10130)], amongst others.", "methods_and_evaluation_criteria": "The methods and evaluation criteria make sense.", "theoretical_claims": "See previous comment", "experimental_designs_or_analyses": "The experimental analyses seem appropriate to me.", "relation_to_broader_scientific_literature": "- In relationship to the existing literature, this work can be viewed as an extension of [Riemannian Flow Matching](https://arxiv.org/abs/2302.03660) to the case of manifolds of distributions. This involves taking known results about the geometry of Wasserstein space and combining these tools with RFM to obtain their model. \n- Prior work on flow matching for point cloud generation (that the authors cite) develop similar approaches, but do not provide the nice formalism the authors here do. \n- The relationship between the submission and [Meta Flow Matching](https://arxiv.org/abs/2408.14608) could be described in more detail.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "# Strengths\n- The general framework of FM on Wasserstein space is very clear and compelling, and lends itself nicely to the methodology developed.\n- The synthetic experiments demonstrate the importance of appropriately accounting for the geometry of the space.\n\n# Weaknesses\n- Some of the claims are not entirely justified (see above).\n- The proposed method has a fairly high computational cost, as each training step requires solving an OT problem.\n   - While this is fine if it leads to better models, I am not convinced that it does. On the 3D point cloud experiments for instance, the proposed method seems to do a bit worse than PSF, which is essentially the same technique as the one proposed in this work, except without using an inner optimal transport solver.", "other_comments_or_suggestions": "- Equation 11 is critical to the proposed method; dedicating more space to the derivation and justification of this in the main paper might be useful.\n- There appear to be some typos in Section 2.3.1 -- namely the use of a, b, A, B to (I think?) represent means/covariances.", "questions_for_authors": "See above comments on claims", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xuznN5AB5O", "reviewer": "Reviewer_ToSj", "summary": "The paper addresses the problem of learning generative models of high-dimension distribution, i.e., where each sample from the model is a distribution itself. The authors propose Wasserstein flow matching (WFM), which builds on top of recent advances in the Riemannian flow matching  (RFM) framework (Chen & Lipman, 2023) and extend it to the Wasserstein geometry (Ambrosio et al. 2008).  The authors show that WFM can generate high-dimensional distribution either represented analytically as Gaussians or empirically as point clouds and derive valid flow-matching objectives for these cases. The authors then apply the method to single-cell and spatial transcriptomics datasets. Additionally, the method is applied to 3d shape generation on point cloud data (ShapeNet & ModelNet), on which it performs similarly, although not better than other methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer 3 (**ToSj**)\n\nWe are grateful to the reviewer for their positive assessment of our work.\n\n**\"While Tab 3 reveals that WFM performs similarly to other methods on ShapeNet and ModelNet, it never outperforms them.\"**\n\nAll reviewers expressed concerns about our performance in 3D shape generation experiments (**Table 3**). We've addressed this by making the following improvements:\n\n1.  **Fixed a minor inference bug**: We corrected a time-stepping loop that previously ran from t=0 to t=1+$\\Delta$t (i.e. from $i=0:N$) , instead of the correct t=0 to t=1 ( $i=0:N-1$). \n2.  **Standardized evaluation metrics**: We now use the same approximate EMD implementation as other benchmark methods. Our original results used true EMD calculations, creating an unintended evaluation discrepancy. Using consistent benchmarking code provides a fairer comparison.\n3.  **Matched point cloud size**: We've increased our point cloud size from 1000 to 2048 points to align with other methods' evaluations. This directly addresses reviewer concerns about scalability while enabling more direct comparisons.\n\nThese changes have significantly improved our results:\n\n| | Airplane | | Chair | | Car | |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|\n| | CD ↓ | EMD ↓ | CD ↓ | EMD ↓ | CD ↓ | EMD ↓ |\n| PointFlow | 75.68 | 70.74 | 62.84 | 60.57 | 58.10 | 56.25 |\n| SoftFlow | 76.05 | 65.80 | 59.21 | 60.05 | 64.77 | 60.09 |\n| DPF-Net | 75.18 | 65.55 | 62.00 | 58.53 | 62.35 | 54.48 |\n| Shape-GF | 80.00 | 76.17 | 68.96 | 65.48 | 63.20 | 56.53 |\n| PVD | 73.82 | 64.81 | **56.26** | **53.32** | 54.55 | **53.83** |\n| PSF | 71.11 | **61.09** | 58.92 | 54.45 | 57.19 | 56.07 |\n| WFM (ours) | **69.88** | 64.44 | 57.62 | 57.93 | **53.41** | 58.10 |\n\nWFM now achieves state-of-the-art CD performance on 2/3 datasets while maintaining competitive performance on all metrics. Importantly, WFM achieves these results with substantially lower computational requirements (~120 GPU hours on an 80GB GPU versus ~400 GPU hours for PSF and PVD).\n\n**\"Page 3, line 110 typo 'moodal data'.\"**\n\nThank you for catching this typo. It has been fixed to 'multi-modal data'.", "claims_and_evidence": "The authors claim to introduce WFM, extending the Flow-matching framework to the space of probability distributions. To my knowledge, this is a novel contribution. The authors provide a sound theoretical foundation for the method and demonstrate its effectiveness on various tasks. Empirical results suggest that it outperforms other approaches (Table 2) in its intended application and is on par with current state-of-the-art models for 3d shape generation (Table 3) while allowing flexible choice of number of particles (Table 4).", "methods_and_evaluation_criteria": "The proposed method is well-motivated, and the evaluation is thorough.", "theoretical_claims": "The authors prove the validity of the WFM objective in Appendix B. The proof is sound and largely builds upon previous work on FM (Lipman et al., 2022, Chen & Lipman, 2023).", "experimental_designs_or_analyses": "Experimental design and analysis are well thought out and the results are presented clearly and concisely. The experiments are well motivated and the results are convincing.", "relation_to_broader_scientific_literature": "Related work is discussed in Section 2.2 and is, to my knowledge, comprehensive.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n- Very well written and structured.\n- Extensive experiments on toy datasets, standard benchmarks, and fitting scientific applications in genomics.\n\nWeaknesses:\n- While Tab 3 reveals that WFM performs similarly to other methods on ShapeNet and ModelNet, it never outperforms them.", "other_comments_or_suggestions": "- Page 3, line 110 typo \"moodal data\".", "questions_for_authors": "- Can the authors discuss the results \"lower\" performance in Tab 3?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dfUUR2y9e0", "reviewer": "Reviewer_euNm", "summary": "The paper shows how Riemannian Flow Matching (RFM) (Chen & Lipman, 2023) is applied to the Wasserstein metric space, the space of distributions endowed with the Wasserstein metric. For RFM to work, a parametric vector field that such that when marginalizing it over the [0,1] time interval, turns a source, Gaussian distribution to a target distribution, an appropriate metric $g$ and a family of conditional flows. The paper shows that, for the Wassserstein metric space, we can use McCann interpolation, its time derivative naturally matches as the desired parametric vector field, and the Wasserstein metric acts as the Riemannian metric. The authors coined the approach as Wasserstein Flow Matching (WFM).\nThe paper further discusses a special subspace of the Wasserstein space, the subspace of non-degenerate Gaussian distributions, called the Bures-Wasserstein space, where the authors show that the McCann interpolation, its time derivatives and all other key entities can be computed in closed forms, exploring the analytical form of Gaussian distributions. The Bures-Wasserstein space has an application in single-cell gnenomics where cellular microenvironments or highlight fine-grain clusters can be captured with just means and covariances.\nFor the general case of general distributions though, since optimal transport (OT) maps are intractable, the authors proposed to rely on entropic OT (Cuturi, 2013), which approximates OT and can be efficiently computed with GPUs via Sinkhorn's algorithm (Sinkhorn, 1964).\nExperiments were conducted over the family of Gaussians and the family of general distributions using synthetic and real data, where WFM was shown to be competitive against existing baseline approaches.\n\n## Update after rebuttal\n\nMy review is already *positive* for the paper. The additional clarifications from the authors are helpful. I've kept my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer 4 (**euNm**)\n\nWe thank the reviewer for their insightful feedback and positive assessment of our work. We address each point below.\n\n**\"...What I have seen in the paper is that every point cloud in consideration are sampled with 1,000 points. This is too small to represent a realistic point cloud in industry.\"**\n\nIn our revised **Table 3** (see response to reviewer 3 - **ToSj**), we now apply WFM to point-clouds with 2048 particles, which is the standard size used by other generative models for shapes in the literature (see [PVD](https://arxiv.org/abs/2104.03670) and [PSF](https://arxiv.org/abs/2212.01747)).  In this regime, WFM performs on par with or better than other methods, while being significantly more computationally efficient. We note that while training is performed on 2048-sized realizations, inference can be done at any scale. For biological applications like cellular microenvironments, niche sizes are typically on the order of tens of cells, so our methodology is more than sufficient for these important use cases.\n\n**\"...In addition, WFM in this case has to rely on the approximate approach of entropic OT, in which the relationship between point cloud size and computational cost is not clear in the paper.\"**\n\nAs we discuss in **Appendix A**, the complexity of entropic OT is between point-clouds with $n$ particles is $O(n^{2}/\\varepsilon)$, where $\\varepsilon$ is the entropic regularisation strength. We also discuss the trade-offs between true and entropic OT and how the regularisation parameters $\\varepsilon$ provides a straightforward way to balance computational efficiency and OT accuracy.\n\n**\"I do not understand the need for Lemma B.1. The Wasserstein distance function is a metric (Villani, 2009). That means it satisfies the 4 metric axioms: identity of indiscernibles, positivity, symmetry and triangle inequality. A premetric only requires 2 axioms: identity of indiscernibles and positivity. By definition, every metric is a premetric. RFM only requires a premetric.\"**\n\n The reviewer correctly points out that the classical definition of a premetric requires only positivity and identity of indiscernibles, both of which are satisfied by any metric. However, the RFM framework (Chen & Lipman, 2023) uses a slightly different definition that includes a third necessary condition:\n \n1. Positivity: $d(x,y)\\geq  0$ for all $x,y$\n\n2. Identity of indiscernibles: $d(x,y)=0$ iff $x=y$\n\n3. Non-degeneracy: $\\nabla_x d(x, y) \\neq  0$ iff $x\\neq  0$\n\n\nIndeed, if a (squared-)metric is differentiable, then it is easy to sketch out that condition (3) holds. We found it pedagogically useful to explicitly write out this third condition, which introduced the logarithmic and exponential maps.", "claims_and_evidence": "The main claim, which is instantiation of RFM on the Wasserstein metric space is supported by clear and convincing evidences.\n\nThere is one small claim (end of section 1), that \"...current approaches cannot scale to high-dimensional distributions and fail when distributions are realized as point-clouds with variable sizes. Conversely WFM succeeds in these challenging settings, enabling generative modeling in new domains like synthesizing tissue microenvironments...\". I acknowledge that the authors have a very interesting application. However, I am not totally convinced by the evidences. There is nothing in the Wassertein geometry or RFM that deals with point cloud-related computational costs. What I have seen in the paper is that every point cloud in consideration are sampled with 1,000 points. This is too small to represent a realistic point cloud in industry. In addition, WFM in this case has to rely on the approximate approach of entropic OT, in which the relationship between point cloud size and computational cost is not clear in the paper.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are reasonable to me. In the experiments over the family of Gaussians, the results make sense.\n\nMy concern is at the experiments over the general family of point clouds where, if I understand correctly, all point cloud sample sizes are at 1,000 points, which in my view is too low in practice. There is no study regarding the impact of point cloud size vs the overall WFM performance in terms of both accuracy and speed. In my view point, this seems to be a substantial weak point because WFM in this case has to rely on entropic OT due to the intractability of the true OT.", "theoretical_claims": "I did check the proposed proofs. They appear valid to me and they are not that hard to derive.", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "From a theoretical point of view, the contribution here is interesting to branches dealing with synthesizing 2D and 3D point clouds. It might be good to relate to Dirichlet distributions as well for classification applications.\nFrom a practical point of view, applications that can stem from this work like modelling tissue biology and cellular microenvironment distributions can benefit from the work.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "I have some concerns above but this is a good paper overall because it has interesting applications and its theoretical contribution is sound, although not very technically difficult.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this study, the authors proposed a Wasserstein flow matching method for learning generative models of distributions/probability measures. All the reviewers, including AC, admit the solidity of the theoretical part and the novelty of the proposed method. 

However, the proposed method may inevitably suffer from the curse of dimensionality, and its computational complexity seems too high for some practical applications, which harms its feasibility in practice. 

Moreover, the are also related work on modeling exchangeable data and neural processes [1, 2, 3], which should be included. 

Therefore, I tend to weakly accept this work.

[1] Korshunova, I., Degrave, J., Huszar, F., Gal, Y., Gretton, A., and Dambre, J. BRUNO: A deep recurrent model for exchangeable data. In Advances in Neural Information Processing Systems, pp. 7188–7197, 2018.

[2] Garnelo, Marta, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J. Rezende, S. M. Eslami, and Yee Whye Teh. "Neural processes." arXiv preprint arXiv:1807.01622 (2018).

[3] Yang, M., Dai, B., Dai, H., & Schuurmans, D. (2020, November). Energy-based processes for exchangeable data. In International Conference on Machine Learning (pp. 10681-10692). PMLR.
Decision: Accept (poster)