Learning Distributions on Manifolds with Free-Form Flows

Summary

In this paper the authors propose to adapt the formalism in Draxtler et al. (2024) and Sorrensen et al. (2024) to Riemannian geometry. The idea they build upon is called Free-form flows and consists of: (i) relaxing the bijectivity requirement of Normalizing Flows by approximating the inverse transformation with a learnable function and (ii) approximating the gradient of the Jacobian determinant with the Hutchinson's trace estimator (and in particular to use the approximate inverse at this point). The contribution consists in generalizing such framework to accommodate for Riemannian manifolds (and specifically to accommodate for the further multiplication by the tangent space orthonormal basis).

Strengths

- the proposed Theorem 3 is elegant and simple, and allows to easily adapt the framework to the manifold under study - even if not discussed in these terms, the proposed method seem to work particularly well in settings where the underlying density is highly multi modal (e.g. Table 3 and Table 6)

Weaknesses

- 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 (see Table 4). In almost all other experiments (Table 3, 5, 6) the proposed method performs equal or worse than other multi-step approaches. - I would then expect the proposed method to highlight a trade-off between performance and computational speed, but no thorough runtime comparison study is performed.

Questions

1. as I understand the method, the main advantage comes from fast inference. Since the performance compared to other methods are worse or equal in most settings (except highly multi modal distributions), I believe it is important to show a runtime comparison with SOTA approaches. In particular, the very recent method in [1] is particularly flexible and much faster than other concurrent work. How does concurrent work compare with the proposed method at fixed performance (e.g. for a given log likelihood value) in terms of runtime (both training and inference)? 2. what are the numerical log likelihood values relative to the plots in Figure 6? how do they compare to other methods? 3. the proposed method relies on two fundamental approximations: estimation of the trace for the Jacobian determinant and approximate inverse of the learnt transformation. In the appendix the authors argue how to reduce the variance by rescaling v and that this helps particularly in low dimensions. I would be interested to see empirically how the variance scales with the dimensions in some experiments 4. given the several losses used in optimization (eq. 22) I was wondering how well is the inverse approximated in practice, which is the second fundamental approximation of the model. Also, the additional two losses kind of suggest that it is hard in practice to learn the approximate inference, which is however fundamental for the method to work. The authors also highlight that the reconstruction error plays a crucial role to learn the correct log likelihood. What are in practice the reconstruction error values (and the other loss terms) in the experiments you have run? As a general advice, I would suggest the authors to further investigate how and why the proposed approach performs better when the learnt densities are multi modal. I think this is a very promising aspect - and from experience maybe the true limitation of bijective layers - but it is not explored in the paper [1] Chen et al., Flow Matching on General Geometries, ICLR 2024 [2] Draxler et al., Free-form Flows: Make Any Architecture a Normalizing Flow, AISTATS 2024

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

Limitations are not at all clearly stated as such in the paper. Some limitations are referenced in the checklist but their description is vague and some references are missing (e.g. the checkerboard dataset. Without a comparison is also hard to grasp how much of a limitation this is). I would suggest the authors to include a paragraph concerning limitations, where limitations are mentioned and discussed.

Summary

This paper introduces M-FFF, a single-step generative model that operates on Riemann manifolds. The authors extend the existing FFF, which functions in Euclidean space, establishing the theoretical foundation for learning distributions on various Riemann manifolds and enabling faster generation than existing multi-step and single-step methods. The effectiveness of the proposed approach is evaluated on several datasets.

Strengths

1. The authors have extended the single-step generator FFF to arbitrary Riemann manifolds, making it possible to generate data in non-Euclidean spaces in a single step with simulator-free which developed in recent flow matching. 2. M-FFF demonstrated superior performance compared to existing flow models in experiments with different datasets. It not only generated high-quality samples but also effectively learned the distribution on each manifold.

Weaknesses

1. Extending existing generative models operating in Euclidean space to Riemannian manifold is a well-trodden path in recent research, so this extension alone lacks novelty. However, the theoretical contributions specific to the extension are solid. 2. The authors evaluated the method's effectiveness only under limited conditions where the dimensionality n is ~10 or less. Given that the theory and method are highly general without specific applications, there should be a discussion on the potential limitations of the algorithm.

Questions

1. Is it realistic to generate data in high-dimensional spaces (e.g., 100 ~ 1000 dimensions), or is it inherently challenging? Understanding the short rebuttal period, I don't expect new experiments, but could you provide insights on the expected impacts of increased dimensionality? 2. Do you plan to release the source code, including reproduction experiments? Considering the scarcity of source code for generative models on Riemann manifolds, releasing the code would significantly contribute to the field. 3. Regarding Specialized architectures in related work, refer to, for example, the paper for quotient manifolds by Yataka et al., Grassmann manifold flows for stable shape generation, NeurIPS2023. While authors of M-FFF claim their theoretical results operate on arbitrary Riemann manifolds, can the same be said for quotient manifolds? 4. M-FFF is described as requiring "requiring only a projection function from an embedding space." However, 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? At least, it seems there is a discrepancy between the claims in the Introduction and the necessity of Eq. (9) in Section 5 (Toy distributions on hyperbolic space). 5. 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." However, considering that Riemannian flow matching is rigorously formulated based on the Riemannian metric, this approach seems inconsistent. Theoretically, wouldn’t a formulation using the Riemannian manifold's metric be necessary? 6. Particularly in spaces with negative curvature, such as hyperbolic space, I am concerned about the appropriateness of using Euclidean distance. If some measures are needed, please explain why the formulation in Eq. (22) is appropriate as a general evaluation metric for M-FFF.

Rating

7

Confidence

5

Soundness

4

Presentation

3

Contribution

3

Limitations

1. There is no mention of limitations related to dimensionality.

Summary

The extends a recent class of generative models, called free-form flows, to Riemannian manifolds. These models are trained like normalizing flows, but relax the bijective constraint and instead employ separate encoder and decoder networks that are trained to be inverses of each other. The proposed method differs from existing flows for manifolds in that it constructs a map from the manifold M, back to M, using a full rank transformation plus a projection. This differs from existing manifold flows which typically use injective transformations. The authors derive a change of variables on manifolds for this setting and derive an unbiased estimator of its gradient and propose a loss function that affords stable training of these models. The method consistently outperforms existing generative models on manifolds at learning the various target distributions from samples at a significantly lower expense than previous methods.

Strengths

- The paper is well written - The idea is novel and extends an interesting line of work around a new class of generative models - The method seems to perform well in practice, although I have some concerns - Section 3 includes a loss function that mitigates potential failure cases like the networks failing in low data regions, or the ambient space encoder mapping to points away from the manifold. - The method is evaluated on the typical datasets seen in flows for manifold papers.

Weaknesses

- I'm not fully convinced by the likelihood based evaluation. The invertibility of normalizing flows ensures that the flow can produce samples anywhere on $\mathcal{M}$ and similarly, can assign a likelihood to any point on $\mathcal{M}$. Since free form flows are not invertible, it is possible that the flow does not map to all of $\mathcal{M}$. This would make the comparison between the likelihood values of normalizing flows and free form flows invalid because the likelihoods are valid on different domains, which could potentially invalidate the reported results. Instead, the experiments could report a sample based metric to measure the method's performance. - There is no error bound on the likelihood gradient like there is in Draxler et al. 2024. Although it seems like the method performs well in practice, I would like to know exactly how much is lost by using a decoder network instead of a true inverse. In injective flows, I don't think that the free form flow gradient works because $\nabla \log |f'(x)| = \text{Tr}{(\nabla f'(x){f^+}'(z))}$ where ${f^+}'$ is the pseudo inverse of $f'$, so using any other possible inverse would give a wrong gradient. I understand that the proposed method is not an injective flow and might not suffer from this, but I would like to see a gradient error bound to know for sure. - The $m$ and $n$ dimensions in the proof of theorem 2 seem to be flipped in the appendix.

Questions

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

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations.

Summary

This paper introduces Manifold Free-Form Flows (M-FFF), an extension of free-form flows to Riemannian manifolds. M-FFF is easy to adapt to an arbitrary Riemannian manifold, since the method requires only a projection function from an embedding space. Moreover, the method itself relies only on a single function evaluation during training and sampling, which speeds up inference over multi-step methods by several orders of magnitude. As demonstrated by the experimental section, M-FFF consistently matches or outperforms single-step methods on several benchmarks on spheres, tori, rotation matrices, and hyperbolic space, while being reasonably competitive with multi-step methods.

Strengths

1. This adaptation of free-form flows to manifolds is novel and allows for unconstrained flows over manifolds that are easier to learn (due to fewer geometric restrictions) and much faster as a consequence of the single-step nature of the method. 2. The results given are quite reasonable and show that M-FFF is an excellent single step method and competitive with multi-step methods, while being much faster.

Weaknesses

1. The method, in essence, trades the hard manifold constraints and theoretical cleanliness of previous Riemannian Normalizing Flows (Mathieu et al., 2020) for a projection-based approximation. Specifically, the benefit of a loose and adaptable "encoder-decoder" framework comes at the cost of having to optimize three additional losses: the reconstruction loss $\mathcal{L}_R$, the uniformity loss $\mathcal{L}_U$, and the projection loss $\mathcal{L}_P$. That is to say, unless these losses are optimized well, the geometry may be respected very approximately, which is a drawback relative to existing multi-step methods. The authors themselves admit that they have had issues with the reconstruction loss being too high (see Section 5). Additionally, such a method may be rather suboptimal for applications where there are hard physical constraints, e.g. see learning symmetric densities over $SU(n)$ [a] in the context of lattice QFT for downstream use in gauge integrals. M-FFF is nonetheless an interesting method, but these limitations should be well-noted. 2. There are clearly contexts in which this method is suboptimal relative to existing multi-step methods, e.g. see Figure 6. Due to the many approximations made in the method, it can be hard to match some of the more rigid density geometry exhibited by the checkerboard pattern, that strict on-manifold multi-step methods have no trouble learning (see Lou et al. (2020)). This should also be well noted; the speed flexibility of M-FFF come at a cost. ### References [a] Equivariant manifold flows. https://arxiv.org/abs/2107.08596

Questions

My questions to the authors are listed below: 1. 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? I believe this question is key to better determining when precisely M-FFF is a reasonable and/or good alternative to strict multi-step methods such as Lou et al. (2020) or Mathieu et al. (2020). ### Additional Comments and Minor Corrections The writing in this paper is for the most part quite clear and the presentation of material is excellent, with several expository figures given to elucidate important parts of the method. Nonetheless, I would like to give a non-exhaustive list of minor corrections below: L62: "for building normalizing flows such as SO(3)" -> "for building normalizing flows for manifolds such as SO(3)" L203: I believe the authors meant to say the substitution happens for equation (13), and the results is shown in equation (19). That is, "This allows us... in Eq. (19)" -> "This allows us.... in Eq. (13)" L228: "by rotations matrices with positive determinant" -> "by rotation matrices with positive determinant"

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Although some limitations of the method are mentioned throughout the paper, this work would benefit from an explicit section made to address limitations of the method (in particular incorporating some of what I have mentioned in the "Weaknesses" section above).

Thank you for the rebuttal

Thank you for your rebuttal comments; in particular, the comments regarding when reconstruction error does not drop to a satisfactory level. I maintain my original score and believe this paper should be accepted.