Response to Reviewer 4 (euNm)

We thank the reviewer for their insightful feedback and positive assessment of our work. We address each point below.

"...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 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 and PSF). 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.

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

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

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

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:

1. Positivity: $d(x,y)\geq 0$ for all $x,y$

2. Identity of indiscernibles: $d(x,y)=0$ iff $x=y$

3. Non-degeneracy: $\nabla_x d(x, y) \neq 0$ iff $x\neq 0$

Indeed, 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.

Response to Reviewer 3 (ToSj)

We are grateful to the reviewer for their positive assessment of our work.

"While Tab 3 reveals that WFM performs similarly to other methods on ShapeNet and ModelNet, it never outperforms them."

All reviewers expressed concerns about our performance in 3D shape generation experiments (Table 3). We've addressed this by making the following improvements:

1. 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$).
2. 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.
3. 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.

These changes have significantly improved our results:

WFM 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).

"Page 3, line 110 typo 'moodal data'."

Thank you for catching this typo. It has been fixed to 'multi-modal data'.

Response to Reviewer 2 (C2Yo)

Thank you for your review and insightful comments. We address each of your concerns below:

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

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

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

This is an important assumption, which we do indeed make. In the papers you mentioned (such as Functional Flow Matching), 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.

"The proposed method has a fairly high computational cost, as each training step requires solving an OT problem"

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

"The relationship between the submission and Meta Flow Matching could be described in more detail."

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

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

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

"There appear to be some typos in Section 2.3.1...

Fixed, we thank the reviewer for spotting this typo.

Response to Reviewer 1 (8NGf)

Thank you for your thoughtful review and constructive feedback.

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

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

"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)."

We 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). 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).

The new text now reads (line 248, revision in bold):

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

"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)?"

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

"Experiments related to biology also do not clarify this point for me since I am not an expert in this field"

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

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