Learning Distributions on Manifolds with Free-Form Flows

We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:

No thorough runtime comparison study is performed. [...] How does concurrent work compare [...] at fixed performance [...] in terms of runtime (both training and inference)?

In a nutshell, M-FFF is over an order of magnitude faster at equal inference performance, and computing gradients is cheap even on arbitrary manifolds.

Inference

At inference, M-FFF is faster by a factor of typically 100x than competing methods: Wall-clock time is roughly proportional to the number of steps in Tables 3-6, since M-FFF uses approximately the same network sizes as previous methods.

As proposed, we are adding a comparison of M-FFF to Riemannian Flow Matching [1], see Figure 1 in the rebuttal PDF concerning the protein dataset “general”. To speed up R-FM, we reduce the number of steps to sample. We find that M-FFF is around 50x faster than R-FM at the same NLL. The single-step performance of R-FM is significantly worse than M-FFF.

In terms of wall-clock time, computing the M-FFF performance in the plot took 30 seconds, the R-FM dot with the same performance took 25 minutes (NVIDIA RTX 2080).

Training

Overall, computing gradients is cheap for M-FFF, depending on the manifold RFM can be more expensive. In detail: There are two contributions: Neural network evaluations and usage of manifold-related functions (projection, computing tangent spaces, …). In a single training step M-FFF calls the encoder and decoder once each and the projection function twice. R-FM calls its velocity network once, but also needs to integrate the user-specified conditional velocity field, which can be expensive unless closed form expressions are known (e.g. sphere is cheap, the mesh dataset requires solving an ODE at every training step).

In terms of the time to get a model running on a novel manifold, M-FFF simplifies the process as it only requires specifying a projection function. R-FM in addition requires the user to specify a conditional flow monotonously decreasing a premetric, increasing the complexity to set up the model.

even though the authors claim the performance is "mixed" or better "in some cases", the proposed method performs significantly worse than most multi-step approaches across almost all datasets [...]

M-FFF sets a new SOTA against multi-step approaches in 5 of the 11 experiments presented in Tables 3, 5 and 6 at orders of magnitude faster inference, hence better “in some cases”. With “mixed”, we mean that it outperforms non-SOTA multi-step methods in 3 out of 4 experiments in Table 4.

Among single-step methods, M-FFF is state-of-the-art in 12 out of 12 experiments. Notably, it works with little modification across manifolds, where each previous single-step method required special adaptation.

what are the numerical log likelihood values relative to the plots in Figure 6? how do they compare to other methods?

We will add NLLs to the paper. To the best of our knowledge, there is no established numerical benchmark for learning distributions on hyperbolic space and existing methods use a different embedding. Our experiments confirm that the generalization of our method to non-isometric geometry is correct and yields useful results. We will publish the code to generate our data sets upon acceptance so that future work can use our method as a baseline.

How [does] the variance scales with the dimensions in some experiments?

From a theoretical point of view, the variance scales just like that of the Riemannian flow matching loss (see below). From an experimental point of view, non-Riemannian free-form flows have been successfully scaled to a 261-dimensional molecular dataset, outperforming previous continuous normalizing flows.

In detail: From appendix D.3 in [1], the variance of the trace estimator $v^T A v$ scales roughly as $2 \\lVert A \\rVert_F^2$. If we apply this to a simple linear model $f(x) = Ax$ and $g(z) = Bz$ we find that summing up the variances of the gradient estimates wrt each element of $A$ leads to a total variance of $2 n \\operatorname{tr}(B^T B)$. A similar calculation for a flow matching loss of the form $\\frac{1}{2} \\lVert Ax - y \\rVert^2$ leads to total variance $n \\lVert x \\rVert^2 \\operatorname{tr}(\\Sigma)$ with $\\Sigma$ the covariance of $p(y|x)$. We can see that both expressions scale explicitly with $n$, with the trace terms also likely scaling with $n$. We thus expect no problems due to variance in high dimensions. We are happy to provide further details on the derivation.

Experimentally, [2] train a Free-Form Flow on a 261-dimensional molecular modeling dataset (QM9). They outperform a previous continuous normalizing flow on this data. Since structurally, M-FFF is very similar to FFF, we expect this behavior to transfer.

• [1] Sorrenson et al., Lifting Architectural Constraints [...], ICLR 2024
• [2] Draxler et al., Free-form Flows [...] AISTATS 2024

What are [...] the reconstruction error values?

They are low, see the table in the rebuttal PDF.

Limitations are not at all clearly stated as such in the paper. [...]

We will improve the presentation by adding a dedicated limitations section to the camera ready version. Here are the three main points:

1. M-FFF approximates the change of variables by replacing the exact inverse with an encoder decoder pair, see l. 190-191, l. 203. This might influence training performance. However, we find that the additional architecture flexibility gained from this approximation lets us outperform concurrent single-step methods with an exact volume change.
2. Compared to multi-step methods, M-FFF performs poorly on specific datasets, especially when the density changes sharply or very low density regions are present, see l. 243-245, caption of Figure 6.
3. Inaccurate likelihoods when reconstruction error does not drop low enough, see l. 240-242, Appendix B.1

We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:

Is it realistic to generate data in high-dimensional spaces? [...] There is no mention of limitations related to dimensionality

Yes, it is realistic. From a theoretical point of view, the trace estimator scales similarly to Riemannian flow matching (see below). From an experimental point of view, non-Riemannian free-form flows have been successfully scaled to a 261-dimensional molecular dataset, outperforming previous continuous normalizing flows.

In detail: One metric to assess scaling behavior is to consider the variance of the trace estimator. From appendix D.3 in [1], the variance of the trace estimator $v^T A v$ scales roughly as $2 \\lVert A \\rVert_F^2$. If we apply this to a simple linear model $f(x) = Ax$ and $g(z) = Bz$ we find that summing up the variances of the gradient estimates wrt each element of $A$ leads to a total variance of $2 n \\operatorname{tr}(B^T B)$. A similar calculation for a flow matching loss of the form $\\frac{1}{2} \\lVert Ax - y \\rVert^2$ leads to total variance $n \\lVert x \\rVert^2 \\operatorname{tr}(\\Sigma)$ with $\\Sigma$ the covariance of $p(y|x)$. We can see that both expressions scale explicitly with $n$, with the trace terms also likely scaling with $n$. Since the number of parameters (elements of $A$) scales as $n^2$, the variance per parameter is constant. We thus expect no problems due to variance in high dimensions. We are happy to provide further details on the derivation of these expressions if the reviewer is interested.

Experimentally, [2] train a Free-Form Flow jointly modeling the position and atomic properties (3+6 dimensions) of up to 29 atoms, making this a 261-dimensional dataset (QM9). They outperform a previous continuous normalizing flow on this data. Since structurally, M-FFF is very similar to FFF, we expect this behavior to transfer.

• [1] Sorrenson et al., Lifting Architectural Constraints of Injective Flows, ICLR 2024
• [2] Draxler et al., Free-form Flows: Make Any Architecture a Normalizing Flow, AISTATS 2024

Do you plan to release the source code?

Yes, we will release the source code upon acceptance.

While authors of M-FFF claim their theoretical results operate on arbitrary Riemann manifolds, can the same be said for quotient manifolds?

We identify two potential pitfalls when applying M-FFF to quotient manifolds:

1. Not every quotient manifold is Riemannian. For example the metric must be constant within each equivalence class.
2. The topology might change when applying the quotient to an existing manifold, such that potentially a different embedding space and projection needs to be chosen. If the embedding space is the same, the projection must differentiably choose a representative from each equivalence class.

For the Grassman manifold $G(k, D)$, M-FFF can be applied by parameterizing it through idempotent symmetric $k \\times k$ matrices. Projecting to the manifold is then achieved via a singular value decomposition, see e.g. [1].

We think that an extension of M-FFF to arbitrary quotient manifolds is an interesting extension of our work. We thank the reviewer for the interesting pointer to Grassmann Manifold Flows, which we propose to add to our related work.

[1] Implementation of geomstats.geometry.grassmannian.Grassmannian

Isn't the calculation of R and Q dependent on the manifold necessary for computing the negative log-likelihood? If so, how are these obtained?

Yes, that is correct. Given a projection $\\pi$, we use that its Jacobian $\\pi’(x) \\in \\mathbb R^m$ spans the tangent space $\\mathcal T_x \\mathcal M$. Then, $Q$ can be obtained by taking any basis of the column space of this Jacobian $\\pi’(x)$. We will make this explicit below Theorem 2.

The authors describe that evaluating the loss in Euclidean space (Eq. (22)) is practically acceptable because "we find that ambient Euclidean distance works well in." Theoretically, wouldn’t a formulation using the Riemannian manifold's metric be necessary?

It was not necessary for the cases we considered, but the framework is readily compatible with using geodesic distances. Simply write the reconstruction loss as \\mathcal L_R = \\mathbb E[d(g_\\phi(f_\\theta(x)), x)^2]; in the end, both have the same minimum x =g_\\phi(f_\\theta(x)). Using Euclidean distance makes computation of the loss cheaper and did not reduce the performance, so we opted for this variant.

Particularly in spaces with negative curvature, such as hyperbolic space, I am concerned about the appropriateness of using Euclidean distance

We tried both hyperbolic and the Euclidean distance (in the embedding space). The hyperbolic distance did not give an improvement, so we opted for the Euclidean distance to reduce the computational effort. We find that a low Euclidean reconstruction loss also means a low hyperbolic reconstruction loss (e.g. for the swish experiment, the reconstructions have a Euclidean squared distance of $8 \\times 10^{-6}$ and a squared geodesic distance of $2 \\times 10^{-4}$ from the original samples).

We hope that this addresses the reviewer’s concerns and look forward to the subsequent discussion.

We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:

I'm not fully convinced by the likelihood based evaluation. [...] Since free form flows are not invertible, it is possible that the flow does not map to all of $M$

Yes, this could happen in principle if the reconstruction loss is not fulfilled. To prevent this pathology, we monitor the reconstruction loss of validation data during training. In the rebuttal PDF, we report the reconstruction loss after training for all datasets, which are consistently low.

We will add these numbers and the above explanation to the main paper.

There is no error bound on the likelihood gradient

Good point! We find that the same derivation applies to M-FFF, only that the Jacobians are projected to the tangent space. Here’s the bound we find:

\\left| \\text{tr}(R^T (\\nabla_{\\theta} J_{f_\\theta}) J_{g_\\phi} R) - \\nabla_{\\theta} \\log |R^T J_{f_\\theta} Q| \\right| \\leq \\lVert R^T (\\nabla_{\\theta} J_{f_\\theta}) J_{f_\\theta^{-1}} R \\rVert_F^2 \\lVert R^T J_{f_\\theta} J_{g_\\phi} R - I_n \\rVert_F^2

Notation: J_{f_\\theta} = f_\\theta'(x), J_{g_\\phi} = g_\\phi'(z) and J_{f_\\theta^{-1}} = f_\\theta^{-1\\prime}(z). We provide a formal statement in the rebuttal PDF, and will add a proof to the camera ready version (if the reviewer is interested, we can also post it here).

The $m$ and $n$ dimensions in the proof of theorem 2 seem to be flipped in the appendix. [...] It would be nice to clarify in the text that all of the quantities considered are written in terms of the Euclidean ambient space coordinates.

The dimensions are not flipped, but we did not make the use of internal/ambient coordinates explicit enough. We will clarify this as needed for the camera ready version.

We hope that this addresses the reviewer’s concerns and look forward to the subsequent discussion.

We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:

The geometry may be respected very approximately

No, this is a misunderstanding. M-FFFs always respect the geometry: In Eq. (17), we exactly project the output of an unconstrained neural network to the manifold. Thus, encoder and decoder are guaranteed to be manifold-to-manifold functions via the projection $\\pi : \\mathbb R^m \\to \\mathcal M$ (see Fig 1, right).

The projection loss encourages the unconstrained neural network to produce points close to the manifold before they are projected to the manifold. Therefore, the projection only leads to small (and thus numerically stable) corrections, and topologically problematic mappings (e.g. to the center of a Euclidean embedded sphere) are avoided. The uniformity loss encourages encoder and decoder to become inverses of each other even at locations where only few training points are available, i.e. where the NLL training signal is weak.

Thus, projection and uniformity loss mainly complement the NLL loss, and we generally find that M-FFFs are not very sensitive to the relative strength of these losses, see the ablation in Figure 1b provided in the rebuttal pdf.

learning symmetric densities over $SU(n)$

We need to disambiguate two requirements:

1. Distributions that are invariant under the action of $SU(n)$.
2. Distributions that assign a probability to the elements of $SU(n)$.

The reviewer’s remark refers to requirement 1, whereas our paper addresses requirement 2: M-FFFs can exactly learn distributions over the elements of $SU(n)$ by embedding $U \\in SU(n)$ into $\\mathbb C^{n \\times n}$, and associating this space with $\\mathbb R^{2 \\times n \\times n}$. We can project from $\\mathbb C^{n \\times n}$ to the $SU(n)$ manifold by various methods, e.g., via polar decomposition or via the constrained Procrustes algorithm.

We agree that combining density estimation on manifolds with group invariance (requirement 1) is an interesting future direction.

There are clearly contexts in which this method is suboptimal relative to existing multi-step methods [...] This should also be well noted

Multi-step and single-step methods occupy different spots in the trade-off between competing modeling goals: Multi-step methods have slightly higher generative quality, especially in the presence of sharp boundaries in the density. Single-step methods are $\approx 100\times$ faster and allow for much more efficient explicit calculation of the probability density (expensive integration of the instantaneous change-of-variables is avoided). Note also that our method outperforms previous single-step approaches, even when they are specialized to a specific type of manifold. See our response to all reviewers above for more details.

As an example, Figure 1 in the rebuttal PDF compares the performance of M-FFF to Riemannian Flow Matching: R-FM requires 50x the compute to achieve the quality of M-FFF on the protein dataset “General”.

In Section 5, the authors mention that the "reconstruction error sometimes does not drop to a satisfactory level." My question is, when does this usually happen, and can we more precisely characterize the situations in which one of the non-NLL losses fails to give a satisfactory result?

Yes, we are happy to elaborate. In general the likelihoods we measure with M-FFF are trustworthy as long as the validation reconstruction errors are low. This is the case for all datasets as shown in our rebuttal pdf. However, if there are sharp edges present in the probability density, a small error in reconstruction near the edge can have a large effect on the measured likelihood. We identify this problem in the earth datasets (e.g. many floods at coastlines, but none at sea) and present a diagnostic to verify the measured likelihoods in appendix B.1.

We will add this explanation to Section 5 and will add it to an explicit limitation section in the camera ready version.

We hope that this addresses the reviewer’s concerns and look forward to the subsequent discussion. We thank the reviewer for the list of corrections which we will incorporate into the camera ready version.