Flow Matching on General Geometries

Authors response

We thank the reviewer for their commendations and concise review.

Authors response

We thank the reviewer for their commendations and meticulous comments. We address them below. > How fast is the training of the proposed approach compared to the training of previous diffusion models? As RSGM (Riemannian Score-based Generative Model) observes that they achieve similar test NLL with only 5 in-training simulation steps (Section O.1), I think that the simulation-free training does not provide significant time efficiency. Note that in Appendix I, we provide wallclock time comparison between simulating the path and our simulation-free approach. We see a 17x improvement in training speed even after taking into account neural network evaluation and gradient descent. This is compared to a 200 step baseline. While it may be true that fewer steps are needed for some manifolds, we note that this can be a tricky hyperparameter to tune. To the best of our knowledge, in the existing Riemannian diffusion model works [1, 2], the number of steps is tuned on each manifold and can be very different depending on the manifold (ranging from 100 to 1000). On the other hand, being able to perform exact sampling with a simulation-free approach completely foregoes the need to consider such a hyperparameter. Also note that the likelihood values shown in O.1 in the RSGM paper [2] (with up to 200 steps) are still quite inferior to the final results reported in Table 4 (using 1000 steps) in the main part of the paper [2]. > Why does the proposed method outperform the Riemannian Diffusion model? A potential reason could be that we have zero numerical & approximation errors. Please see Appendix D, where we highlighted, in red, sources of potential errors during training between our framework and Riemannian diffusion. Namely, for diffusion models, both the sample $x_t$ and the “supervision signal” (either denoising or implicit score matching) must be approximated. This is not the case for RFM, where we have both in closed form for simple manifolds. Also, the noising processes in diffusion models do not reach the uniform distribution until infinite time, requiring yet another hyperparameter to tune which is the time horizon of the simulation. > Further, why does CNF Matching show superior performance on Earthquake and Flood datasets over the proposed method? Whereas RFM prescribes both the marginal probability path and the marginal vector field (which is being regressed onto), there is no unique vector field for a chosen probability path. CNF Matching takes advantage of this flexibility and prescribes only the marginal probability path and gives the model the freedom to choose any vector field that matches this probability path. It may be this extra freedom that is allowing CNF Matching to perform better. Note however, that CNF Matching has a biased loss that cannot scale as well as RFM. > Why is the result of the proposed method for Flood dataset in Table 2 highlighted in bold even though CNF Matching shows better result? Apologies, this is a mistake on our part. It was not our intention to mislead, and we have fixed this. Thanks for catching this. We have also weakened some of our statements surrounding the experiments to reflect that we only achieve state-of-the-art on some but not all data sets that we consider. > Can the other diffusion models applied to general manifolds such as the triangular mesh? For example, RSGM seems to be applicable when using the denoised score matching with the Varadhan approximation. There are two approaches for approximating the conditional score function that RSGM [2] proposed: (1) based on a truncated heat kernel expansion using eigendecomposition, and (2) based on the Varadhan approximation. The second option (2) actually will not work on general manifolds as it (i) requires the logarithm map globally and (ii) is only justified for small t values. For approach (1), the truncated heat kernel expansion using eigendecomposition is possible, but the number of eigenfunctions required for reasonable approximation of the score is very high for RSGM. It is much easier to satisfy the properties of a premetric in our RFM framework---possible with a small finite number—than it is to approximate the heat kernel using eigenfunctions. Indeed, we **expanded the discussion on k and added extra illustrative examples in Appendix G.1** on comparing the number of eigenfunctions required by RSGM and RFM, where we see very large score approximation errors even on a simple 2D-sphere using 100s of eigenfunctions, while spectral distance already provides a valid premetric using 3 eigenfunctions. We believe this property is further exacerbated on more general complex manifolds. We hope this answers the reviewer’s concerns. If we missed anything, please feel free to let us know. [1] Riemannian Diffusion Models. https://arxiv.org/abs/2208.07949. [2] Riemannian Score-Based Generative Modelling. https://arxiv.org/abs/2202.02763.

Authors response

We thank the reviewer for their commendations and insightful questions. Below we address their main points: > the effect of k (the number of eigenfunctions) on the learned model. Theoretically we know there exists a finite (bounded) value of k that allows the spectral distances to satisfy our premetric conditions. To help provide intuition regarding the effect of $k$, we have **expanded the discussion and added a simple experiment in Appendix G.1** where we compare spectral distance and diffusion model’s conditional probability on the 2D-sphere for different choices of $k$. The spectral distance changes only slightly for different $k$ and already separates all points for $k=3$ and is a valid premetric. Alongside this, we compare to diffusion models, which in contrast, require many more eigenfunctions to be able to concentrate probability mass at small time values $t$ and to get an accurate estimate of the conditional score function. In fact, regardless of how many eigenfunctions are used for diffusion the score approximation is always biased. > choice of scheduler. We believe the optimal choice of scheduler depends on the data distribution itself, and it is difficult to infer at the time of training how to best choose a scheduler. For instance, [1] chooses a scheduler that minimizes the kinetic energy of the learned model (for the Euclidean case). It could be interesting to see how that translates to manifolds, but it does not seem trivial. > Is there any benefit to considering non-uniform base distributions? While the theory easily admits arbitrary base distributions, are there practical challenges in using non-uniform distributions? We note first that the maze manifolds use Gaussian mixture models as base and target distributions, and the high-dimensional tori experiment uses a wrapped Gaussian. It is actually often harder to use a uniform distribution because in high dimensions, the volume grows exponentially (so density decreases exponentially). For manifolds with boundaries, ideally the base distribution does not have significant mass on the boundary itself. Overall, we aren’t aware of any practical challenges of using non-uniform distributions. However, choosing a good base distribution is also unclear a priori; a good base distribution for generative modeling is perhaps one that is already close to the data distribution. > For instance, can we say anything about what makes a particular premetric "good"? We agree that being able to answer such questions can bring many insights into this and similar frameworks. However, we don’t think there is a clear cut answer. Similar to the scheduler, it may very well be that a “good” premetric depends on the data distribution. > Are there principled reasons someone would choose the diffusion distance over the biharmonic distance? Diffusion distances have a time parameter ($\tau$) that can be tweaked to achieve certain properties. For example, as can be seen in Figure 2 in [1], that long time parameter provides a smoother overall distance field, while short time parameters is more isotropic near the source point at the cost of a more wiggly distance landscape. Practically, we note that the biharmonic is often easier to use as a first choice because of the lack of hyperparameters, where diffusion distance requires strict tuning of the time parameter which we found it is quite sensitive to. > Are spectral distances chosen in Section 3.2 primarily for practical reasons, or are there other aspects that make the corresponding premetrics desirable? Spectral distances have a certain computational benefit over geodesic distances as they can be computed on the fly between any pair of points on the manifold by simply computing the euclidean distance following a weighted eigenfunction embedding of these points. Apart from that, diffusion distances are intuitively constructed by averaging many random walks, so they end up being smoother and more robust to topological noise (e.g., holes, narrow passages). See https://www.pnas.org/doi/abs/10.1073/pnas.0500334102, e.g., Figure 2, which initially motivated diffusion distances and diffusion maps. We see the benefit of spectral distances in Figure 1 (bottom row), where the geodesic (bottom-left) wants to always transport along the boundary, but this is not ideal as the marginal distribution ends up being a Dirac on the boundary. Spectral distances (bottom-right) avoid the boundary and allow for a diffeomorphic flow of mass within the manifold. > Is there a compelling reason spectral distances use the spectrum of Laplace-Beltrami operator, rather than some other operator? The eigenfunctions of the Laplace-Beltrami (LB) operator are considered to be the most natural (or even, the “correct”) generalization of the Fourier basis on the unit circle (arguably the simplest 1D compact manifold) to general manifolds. [1] On Kinetic Optimal Probability Paths for Generative Models. https://arxiv.org/abs/2306.06626.