Wasserstein Flow Matching: Generative modeling over families of distributions

Thank you for your insightful remarks and thoughtful questions.

The paper presents a new concept of Wasserstein Flow Matching, which extends the Flow Matching framework to probability distribution spaces. It effectively incorporates Gaussian cases through McCann’s (displacement) interpolant. The method is applied to a variety of data types, including genomic data and point cloud data. Overall, the paper is clear, well-structured, and easy to follow.

We are grateful for the reviewer’s positive remarks of our work, specifically regarding the conceptual novelty of WFM and the manuscript's clarity.

The motivation for learning the mean field on the Wasserstein space is not entirely clear, as the mean field may not capture meaningful structures like Wasserstein geodesics.

There is no mention of “mean field” in our paper — would the reviewer be able to clarify this point so we can appropriately address it?

The reviewer is correct that flowing along the Wasserstein geodesic is a critical component of our algorithm. Indeed, WFM does exactly that as our method explicitly learns

1. For Gaussians: The vector fields in equations 8-9 exactly represent the Wasserstein geodesic paths.

2. For point-clouds: The full OT solution yields the true Wasserstein geodesic, which we efficiently approximate using entropic regularization with minimal loss in accuracy.

These theoretical aspects are now discussed in depth in our revision (see Appendix A.3 Appendix B).

Additionally, the approach first estimates McCann’s interpolation using traditional methods, followed by the FM algorithm. This sequential combination of the two algorithms feels somewhat incremental in terms of contribution.

Along with other impactful generative models for statistical domains, such as Fisher, Dirichlet, Categorical and SE(3) Flow Matching [1,2,3,4], our work is also a variant of Riemannian Flow Matching (RFM). Our contribution is specifically lifting RFM to the Wasserstein space, which seamlessly models both families of Gaussian distributions and families of point-clouds.

Working in the space of distributions required several crucial innovations, such as modeling point-clouds with a transformer to capture the permutation equivariance of the OT map. We believe this is a surprising novelty that has yet to be explored in many works at the intersection of machine learning and optimal transport, with future applications abundant. The transformer combined with the entropic OT estimate allows for generation of variably-sized and high-dimensional point-clouds, which to the best of our knowledge makes WFM the first example of a model capable of this.

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

We acknowledge this deviation from the theoretical framework. Our use of entropic OT is a practical choice that enables GPU training, unlike unregularized OT. When combined with rounded matching (A.1 in the appendix), the entropic estimate still produces valid vector fields and adheres to flow-matching theory. We discuss additional details of this in Appendix A and Appendix B, and our general response above. These sections in the appendix discuss how the regularization parameter offers a trade-off between theoretical exactness and computational efficiency, as smaller values require more Sinkhorn iterations.

For completeness, we've added support for full, unregularized OT via POT, but it is significantly slower and cannot handle point-clouds of different sizes.

Even with reduced computational overhead through estimating entropic OT, I believe there would be still huge computational burden when estimating μt in point cloud application. It might be helpful if the authors could specify the computational costs associated with solving entropic OT during training.

We agree with the reviewer that there is a non-trivial computational overhead required to estimate the OT map, even in the entropic setting. As noted, WFM provides users with the ability to tune between theoretical adherence and computational efficiency via the entropic regularization parameter.

There are two sources of computational burden in our model. The first, as the reviewer noted, is the computational estimation of the OT map, and the second is the transformer and its backward pass. Surprisingly, we found that it is the transformer which consumes more resources, not the OT estimation. Both accelerating the computation of OT maps and increasing the efficiency of transformers are highly active areas of research within ML. We envision future work incorporating some of these solutions, namely MetaOT and SetTransformers++ [5,6] for streamlined training of WFM.

For the ShapeNet results in Table 3 and Figure 3, training consisted of 500,000 gradient descent steps with a batch size of 64 which required about 3 days on a single A100 GPU. WFM on datasets with fewer points per sample, such as those in Figures 5-6 and Table 2, trains much faster, requiring only 3-4 hours. We note these training parameters were not optimized, and it is entirely possible that comparable performance can be reached in fewer iterations or a smaller batch size. We have added an extended discussion on computational requirements in the revised manuscript (Appendix E).

As a reference for the reviewer, we have personally reached out to the authors of point-voxel-diffusion and point-straight-flow [7,8] (PVD and PSF in Table 3) to inquire about training time. In both cases, the authors indicated training lasted several days (3-4 for PVD and 2-3 for PSF), and required significant compute (4 GPUs for PVD and 8 for PSF). WFM achieves comparable quality with fewer resources.

I believe there would be scalability issues when the number of particles becomes more than 2000.

The reviewer is correct that WFM suffers as the number of particles increases. Again, due to the transformer and estimate of the OT map. While computational costs do increase with point cloud size, this scaling remains practical for most applications.

We refer the reviewer to the github where we added an example of training WFM to generate point-clouds with 2000 particles. In a modest 15 hours of training, (50,000 training steps and batch size of 32), WFM produces high fidelity point-clouds.

Thank you for engaging with our paper! We are grateful the reviewer found our work conceptually novel and noted its computational efficiency. We thank the reviewer for clearly outlining their remaining theoretical concerns, which now we aim to resolve.

Response to General Statement:

Do you have a reference for where the velocity $v_t(\mu_t)=\mathbb{E}[\mu_1|\mu_t]$ comes from? We believe we have a minor correction here: given the conditional velocities and paths $v_t(\cdot | x_1)$ and $p_t(\cdot|x_1)$, the marginal velocity is $v_t(x)=\mathbb{E}_{x_1}[v_t(x|x_1)p_t(x|x_1)]$. This averaging produces a valid path between from source to target $p_1\rightarrow p_0$ (see Theorem 1 and equation 8-9 in FM paper [1], or equation 7 in the RFM [2] paper), which as the reviewer points out is not the Wasserstein geodesics between $p_1$ and $p_0$.

From here, it may be helpful to clarify that our setting involves two levels of distributions:

1. The outer level: $\mathfrak{p}_0$ and $\mathfrak{p}_1$ are distributions over the Wasserstein space

2. The inner level: Random draws $\mu \sim \mathfrak{p}_0$ and $\nu \sim \mathfrak{p}_1$ are themselves distributions (point-clouds or Gaussians)

We learn conditional vector fields between samples by regressing onto the Wasserstein geodesic connecting $\mu \rightarrow \nu$. While the FM objective still holds (see derivation in Appendix Section B), producing a map $\mathfrak{p}_0 \rightarrow \mathfrak{p}_1$, we do not attempt to learn the Wasserstein geodesic between $\mathfrak{p}_0$ and $\mathfrak{p}_1$.

We have clarified this nuanced point in our re-revision (blue text in pdf).

Response to Q1:

This has been an illuminating example for us to consider, and has led to us realizing the following connection: flow matching is a special case of WFM when the measures are Diracs. This is because the OT map between $\mu = \delta_{x_{0}}$ and $\nu = \delta_{x_{1}}$ is the straightforward linear path $x_{1} - x_{0}$, so the WFM objective turns into the standard FM one. Thus, the reviewer's example, while insightful, falls outside both FM assumptions and our generative setting.

Given this, the reviewer is correct that in their described scenario, where the source measure consists of only a single point-cloud ($p_0: \mu=\delta_{0}$ with probability $1$) while target has two ($p_1: \nu = \delta_{-1}$ or $\nu = \delta_{1}$, each with probability $0.5$), there is no map from source to target. So, in the standard FM case, if the source is a single point, there is no deterministic function which splits it into 2 to draw from the target.

This has helped us clarify the following two key assumptions under which WFM operates (all of our experimental settings fit into these):

1. Outer continuity: the source $\mathfrak{p}_0$ is a continuous measure over distributions (e.g., the spiral of Gaussians in Figure 2 or shapes of all cars for point-clouds). Similar regularities assumptions for the source distribution are also required in standard FM (see Section 2 in FM paper [1])

2. Inner continuity: samples $\mu \sim \mathfrak{p}_0$ and $\nu \sim \mathfrak{p}_1$ are continuous distributions:

• Natural when $\mu$ and $\nu$ for Gaussians

• For point-clouds, $\mu$ and $\nu$ are interpreted as continuous shapes where point-clouds are drawn from them.

These assumptions ensure:

1. A transport path exists between $\mathfrak{p}_0$ and $\mathfrak{p}_1$

2. Deterministic OT maps exist between $\mu \sim \mathfrak{p}_0$ to $\nu \sim \mathfrak{p}_1$

Does it make sense? Do you agree how the example connects WFM and FM, and is outside of the assumptions of both of them?

The reviewer correctly points out that the infinite dimensional Wasserstein space is not a Riemannian manifold. While we make note of this in our manuscript (line 186), we do not fully discuss this matter for brevity. Indeed, the relevant features of Riemannian manifolds for Flow Matching, which are the existence of a tangent space and exponential/logarithmic maps, are still true in Wasserstein space

Response to Q2:

We did not mean to imply that we learn the OT map — we agree this would be misleading. Instead, WFM learns to map source to target by regressing onto OT maps (Wasserstein geodesics) between samples $\mu \sim p_0$ and $\nu \sim p_1$, which are themselves distributions. We have clarified this in lines 266-268 of the latest revision. Can you please let us know if it reads okay now?

“WFM learns to map source to target by regressing onto Wasserstein geodesics between samples $\mu \sim \mathfrak{p}_0$ and $\nu \sim \mathfrak{p}_1$, rather than learning the OT map between $\mathfrak{p}_0$ and $\mathfrak{p}_1$.”

Response to Q3:

While point-clouds are discrete distributions, we interpret them as samples from underlying continuous shapes. This continuity assumption, now explicitly stated in our re-revision, ensures the existence of OT maps between them. Do you agree this justifies the experimental settings?

We again thank the reviewer for their insightful discussion.

[1] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.

[2] Ricky TQ Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660, 2023.

Thank you for your clarification—I believe the explanations provided for Q2 and Q3 are clear. Additionally, I understood the explanation of FM; FM assumes positive density throughout the domain, and with your inner/outer assumption, this counter-example can be easily eliminated.

Turning back to the paper, while the concept is very novel, I still find both the theoretical and experimental contributions lacking. Regarding the experiments, although the paper does not include many comparisons, it does demonstrate potential applicability. Therefore, this is not my primary concern. However, my primary concern lies with the lack of rigorous theoretical grounding. As with FM or Riemannian FM, I believe a precise theoretical discussion on the matching between $p_0$ and $p_1$ is crucial.

Could the authors include a concrete theorem in the manuscript under the assumption that $p_0$ is a continuous distribution on the Wasserstein space? At present, this is not evident, as we cannot rely on the smoothness assumptions of the manifold or vector field typically used in the Riemannian setting.

We are grateful to the reviewer for their continued engagement with our work. We are pleased they were satisfied with our explanations for Q2 and Q3 and for noting the conceptual novelty and potential applicability of our work.

Regarding the experiments, although the paper does not include many comparisons...

It is worth mentioning that over the Bures--Wasserstein manifold, there are no other tractable, large-scale procedures available for learning curves of measures on this space. For high-dimensional point-clouds, other methods do not scale when learning generative models due to the fact that they employ voxelization as a key aspect of their algorithms.

The benchmarking in our manuscript follows established standards for point-cloud generative modeling. Following point-voxel-diffusion and point-straight-flow [1,2] (PVD and PSF in Table 3), we also measure the 1-Nearest-Neighbour (1-NN) accuracy for generation of planes, cars and chairs to assess the quality of our model.

In fact, our work goes beyond previous efforts, as we can generate high-dimensional and variable sized-point clouds. No previous work is capable of this, and here too we validate the fidelity of our model using the same 1-NN metric.

However, my primary concern lies with the lack of rigorous theoretical grounding… cannot rely on the smoothness assumptions of the manifold or vector field typically used in the Riemannian setting.

Regarding your theoretical concerns, we note that the regularity of the metric $g$ and measures $\mathfrak{p}_0,\mathfrak{p}_1 \in \mathcal{P}(\mathcal{M})$ are not rigorously discussed in the original FM or RFM papers. Indeed, Appendix A in RFM [3] simply assumes any and all conditions required for their proofs to work; we eagerly make these same assumptions and we now clearly mention this. Their key assumption is the positivity of the probability path, which is covered by our inner- and outer-continuity, as you have agreed.

With this in mind, we included a concrete proposition and corollary in the newly added APPENDIX B.1 (purple text in our revised manuscript, see github), which we briefly describe here:

Recall that our probability paths are of the same form as Eq (7) of the RFM paper, which for $\mu \sim \mathfrak{p}_0$ and $\nu \sim \mathfrak{p}_1$, is

$\mu_t = \exp_{\mu}((1-t)\log_{\mu}(\nu)) = ((1-t){\mathrm{id}} + tT^{\mu\to\nu})_\sharp\mu\,.$

We prove in Proposition B.1 that the squared $2$-Wasserstein distance is a valid pre-metric in the sense described by the RFM paper. From this, we compute explicitly the following (conditional) vector field (following Eq (13) in the RFM paper):

${\mathfrak{u}}_t(\mu_t|\nu) = -\frac{1}{1-t}{\mathrm{grad}}(W_2^2(\cdot, \nu))(\mu_t) = v_t,$

where $v_t$ is the optimal transport vector field that we regress onto. As a corollary (Corollary B.2 in our revised manuscript), the conditional paths and conditional vector fields are well-defined, and thus results in a valid FM objective function. We note that in all instances in our manuscript, the OT map is well-defined and non-singular, as we assume all measures are densities with respect to Lebesgue measure.

All that said, we again thank the reviewer for their detailed assessment of our work, and hope we have answered any remaining concerns.

[1] Zhou, Linqi, Yilun Du, and Jiajun Wu. "3d shape generation and completion through point-voxel diffusion." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[2] Wu, Lemeng, et al. "Fast point cloud generation with straight flows." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023.

[3] Ricky TQ Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660, 2023.

The paper tackles a novel and important problem within generative modelling of the generation of probability measures, with particular relevance for single-cell genomics data. The authors present a sound methodology in the particular cases of modelling distributions of Gaussian distributions and point-cloud data, with the notable attractive feature of being able to model point-cloud data of non-uniform sizes which is to the best of my knowledge, the first instance of this being tractable; this is backed up with convincing experimental validation. In addition, the paper is well-written and presented.”

We thank the reviewer for their positive assessment of our work. We agree with the reviewer that the ability to model point-clouds data with inhomogeneous sizes is a particularly exciting feature of WFM.

We further stress WFM’s scalability to high-dimensional data through the Transformer architecture. This unlocks key application to single-cell genomics, namely the ability to generate microenvironments from spatial transcriptomics data, which are point-cloud in gene-expression space. This ability to model cellular structure in a generative manner is an entirely novel concept made possible by WFM, and brings to the promise of generative models whose impact in single-cell has been immense (see Virtual cell efforts [1]), to the spatial domains.

A potential criticism of the work is that the main components of the methodology such as Riemannian Flow Matching, closed-form solutions for geodesics etc. on the Bures–Wasserstein space, the entropic estimation of optimal transport maps etc. have been introduced elsewhere, and the main contribution of the paper is simply the original combination of these elements applied to their problem.”

As noted in our response to reviewer R81K, WFM is indeed an instantiation of Riemannian Flow Matching, as are many other works for generative modeling on statistical domains (Fisher, Dirichlet, and Categorical Flow Matching [2,3,4]).

Our innovation lies in harnessing the Wasserstein manifold, particularly suited for Gaussians (Bures--Wasserstein) and discrete measures (point-clouds). By combining OT map estimation with flow matching via transformers, WFM is, to our knowledge, the first generative model for variable-sized, high-dimensional point-clouds, addressing critical needs in single-cell genomics and 3D modeling.

In addition, the authors could have further discussed other settings in which generative modelling on the Wassertein manifold could be useful, such as amortising Bayesian inference, beyond the two cases considered in the paper.

Thank you for pointing this out; this is a great suggestion. We did not think of this application while preparing our manuscript, but it would certainly be worthwhile to explore this in the future.

[1] Bunne, Charlotte, et al. "How to build the virtual cell with artificial intelligence: Priorities and opportunities." arXiv preprint arXiv:2409.11654 (2024).

[2] Davis, Oscar, et al. "Fisher flow matching for generative modeling over discrete data." arXiv preprint arXiv:2405.14664 (2024).

[3] Stark, Hannes, et al. "Dirichlet flow matching with applications to dna sequence design." arXiv preprint arXiv:2402.05841 (2024).

[4] Cheng, Chaoran, et al. "Categorical Flow Matching on Statistical Manifolds." arXiv preprint arXiv:2405.16441 (2024).

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

We are familiar with both of these works, however, to our knowledge, their aims are different from ours. Specifically, Klein et al. (2024) and Eyring et al. (2024) [9, 10] do not explore flow matching on the Wasserstein space; rather, they focus on flows between samples which are individual data points rather than point clouds. In contrast, WFM operates on the space of distributions, with each sample representing a complete point-cloud. These are incomparable tasks, though we have added a detailed discussion in Section 2.2: Related Work.

Additionally, WFM introduces an entirely novel concept of generating cellular environments—specifically, point-clouds derived from spatial transcriptomics data. This capability is highly relevant, as it provides unprecedented insight into spatial microenvironments, a critical aspect of cellular biology that no previous work has addressed in this way.

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

The closure of a set is the set and all its cluster points with respect to a topology (you can think of it as the boundary of a set). This is a minor technical point required in the definition of the tangent space, and we included a brief sentence in the main text to aid the reader.

[1] Davis, Oscar, et al. "Fisher flow matching for generative modeling over discrete data." arXiv preprint arXiv:2405.14664 (2024).

[2] Stark, Hannes, et al. "Dirichlet flow matching with applications to dna sequence design." arXiv preprint arXiv:2402.05841 (2024).

[3] Cheng, Chaoran, et al. "Categorical Flow Matching on Statistical Manifolds." arXiv preprint arXiv:2405.16441 (2024).

[4] Bose, Avishek Joey, et al. "Se (3)-stochastic flow matching for protein backbone generation." arXiv preprint arXiv:2310.02391 (2023).

[5] Cuturi, Marco. "Sinkhorn distances: Lightspeed computation of optimal transport." Advances in neural information processing systems 26 (2013).

[6] Pooladian, Aram-Alexandre, and Jonathan Niles-Weed. "Entropic estimation of optimal transport maps." arXiv preprint arXiv:2109.12004 (2021).

[7] Zhou, Linqi, Yilun Du, and Jiajun Wu. "3d shape generation and completion through point-voxel diffusion." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[8] Wu, Lemeng, et al. "Fast point cloud generation with straight flows." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023.

[9] Klein, Dominik, et al. "Generative Entropic Neural Optimal Transport To Map Within and Across Spaces." arXiv preprint arXiv:2310.09254 (2023).

[10] Eyring, Luca, et al. "Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation." arXiv preprint arXiv:2311.15100 (2023).

Reviewer R81K:

The proposed method can be used to deal with families with distributions, i.e., distributions over distributions. It was tested on a variety of Gaussian-based where it shows better performance than naive baselines.

We are grateful for this constructive review of our work. We appreciate that the reviewer noted the capability of our method to model families of distributions, which is an essential contribution of WFM. In particular, we aimed to address the challenge of generative modeling in the space of distributions over distributions by leveraging the Wasserstein geometry, allowing our method to capture point-clouds and Gaussians effectively.

My main concern is related to the limited significance of the developed methodology. As the authors themselves admit (lines 235-236), the main contribution of this article is the developed algorithm.

Our work is indeed an instantiation of Riemannian Flow Matching, which we emphasized throughout the manuscript. We take this moment to point out that many previous successful applications of flow matching to non-trivial domains, such as Fisher, Dirichlet, Categorical and SE(3) Flow Matching [1,2,3,4] are also specific instantiations of Riemannian Flow Matching.

We stress three novelties of our work:

1. We formally lift the flow matching problem to the Wasserstein manifold, enabling unified treatment of both Gaussians and point-clouds, both cases being novel application domains for spatial transcriptomics.
2. We combine the transformer architecture with entropic OT maps for learnable point-cloud flows
3. Combining (1) + (2), we have the unique ability to generate high-dimensional and variable-sized point-clouds, a capability not found in the existing literature (to the best of our knowledge).

From a theoretical point of view, the method is based on the results proposed in previous works (Ambrosio et al., 2008, Lambert et al., 2022). At the same time, the authors note that these theoretical results are correct if the algorithm is applied to absolutely continuous measures, which is not true in the case when the measures are approximated as point clouds. Thus, it is not absolutely fair to apply the algorithm to any measures other than Gaussians.

We understand the concerns of the reviewer, however, it is important to note some key aspects:

• The space of Gaussians is effectively the only non-trivial space of probability measures for which we have closed-form geodesics.
• Point-clouds can be readily interpreted as samples from a population distribution (consider a shape as a continuous distribution, and point-clouds constitute samples from said distribution). This is the more common regime where OT is applicable in practice. Here, practitioners who are interested in the OT cost are forced to estimate it, often using entropic optimal transport, as introduced by Cuturi (2013) [5]. This has received immense attention in the literature, and recently, Pooladian and Niles-Weed (2021) [6] proposed an estimator for the OT map based on entropic optimal transport, which has found its way to the widely used open-source toolbox OTT-JAX.

To this end, we refer the reviewer to the newly written Appendix A.3, where we discuss the delicate point of estimating geodesics using entropic optimal transport.

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

For the low-dimensional and uniformly sized point-cloud experiments, we agree that our algorithm does not achieve state-of-the-art results. However, as shown in Table 2, the use of transformers and the (entropic) OT map makes WFM the first method (to our knowledge) we can generate non-homogenous and high dimensional point-clouds; please see our general response.

Finally, as we note in our response to reviewer UUcS, WFM requires significantly less resources to train compared to other approaches. Per our personal correspondence with the authors of these papers, point-voxel-diffusion and point-straight-flow [7,8] (PVD and PSF in Table 3), require 4-8 GPUs for several days. In stark contrast, our method needs only a single GPU for 3 days to produce competitive results.

We now compute the conditional vector field as described by Eq. (13) by Chen and Lipman (2024) as a function of our choice of pre-metric ${\rm{d}}$, which we recall here, is a conditional vector field written as $\mathfrak u_t(x_t|y) = (\tfrac{\rm{d}}{\rm{d} t}\ln(1-t)) {\rm{d}}(x_t,y) \nabla_x {\rm{d}}(x_t,y)/\|\nabla_x {\rm{d}}(x_t,y)\|^2_{g(x_t)} \quad \quad (23)$

where $g$ is the metric on the tangent space at $x_t$, and $x_t$ is a geodesic in the form

$x_t = \exp_{x_0}((1-t)\log_{x_0}(x_1))\,.$

Note that our geodesics are precisely of the form

$\mu_t = \exp_\mu((1-t)\log_\mu(\nu))\,$

Where we recall $\exp_{\mu}(v) = (\mathrm{id}+v)_{\sharp}\mu$ and $\log\mu(\nu) = T^{\mu \to \nu} - \mathrm{id}$.

Replacing the general pre-metric with the $2$-Wasserstein distance in (23), as well as the metric over the tangent space $\cal T_{\mu_t}\cal P_{2,\rm{ac}}(\mathbb R^d)$ (which is now $L^2(\mu_t)$), we compute

where we make use of the fact that $T^{\mu_t\to\nu} = T^{\mu\to\nu}\circ (T^{\mu_t\to\nu})^{-1}$ (see e.g., Section 2 of Carlen and Gangbo (2003)), and that $T_t = (1-t){\rm{id}} + t{T^{\mu\to\nu}}$. Simplifying, we obtain $\mathfrak u_t(\mu_t|\nu) = \frac{1}{1-t} (T^{\mu\to\nu} - (1-t){\rm{id}} - t T^{\mu\to\nu} )\circ (T_t)^{-1} = (T^{\mu\to\nu} - {\rm{id}} )\circ (T_t)^{-1} \quad (24)\,$ which is precisely the optimal transport vector field $v_t$ stated in (7). Together, we arrive at the following corollary, which validates the probability paths generated by our approach.

Corollary B.2 The conditional flow $\mathfrak p_t(\cdot|\nu)$ generated by the conditional vector fields defined by (24) satisfy $\mathfrak p_1(\cdot|\nu) = \nu$.

Proof This follows immediately from Theorem 3.1 of Chen and Lipman (2024) by passing through the pre-metric construction above.