Learning Distributions on Manifolds with Free-Form Flows
Authors: Peter Sorrenson
Felix Draxler
Armand Rousselot
Sander Hummerich
Ullrich Koethe
Venue: NeurIPS 2024 poster (2024)
Keywords: generative model
riemannian geometry
riemannian manifolds
free-form flows
normalizing flows
Primary area: generative_models
TL;DR: We propose Manifold Free-form Flows, the first generative model for data on arbitrary manifolds that sample in a single function evaluation at high quality.

Abstract: We propose Manifold Free-Form Flows (M-FFF), a simple new generative model for data on manifolds. The existing approaches to learning a distribution on arbitrary manifolds are expensive at inference time, since sampling requires solving a differential equation. Our method overcomes this limitation by sampling in a single function evaluation. The key innovation is to optimize a neural network via maximum likelihood on the manifold, possible by adapting the free-form flow framework to Riemannian manifolds. M-FFF is straightforwardly adapted to any manifold with a known projection. It consistently matches or outperforms previous single-step methods specialized to specific manifolds. It is typically two orders of magnitude faster than multi-step methods based on diffusion or flow matching, achieving better likelihoods in several experiments. We provide our code at https://github.com/vislearn/FFF.

Learning Distributions on Manifolds with Free-Form Flows
Peter Sorrenson*, Felix Draxler*, Armand Rousselot Sander Hummerich, Ullrich Köthe Computer Vision and Learning Lab, Heidelberg University
*equal contribution, firstname. Jastname@ivz umi-heidelberg.cco
Abstract
We propose Manifold Free-Form Flows (M-FFF), simple new generative model for data manifolds. The existing approaches to learning a distribution on arbitrary manifolds are expensive inference time, since sampling requires solving a differential equation. Our method overcomes this limitation by sampling in a single function evaluation. The key innovation to optimize neural network via maximum likelihood on the manifold, possible by adapting the free-form flow framework to Riemannian manifolds. M-FFF is straightforwardly adapted to any manifold with known projection. It consistently matches or outperforms previous single-step methods specialized to specific manifolds. is typically two orders of magnitude faster than multi-step methods based on diffusion or flow matching, achieving better likelihoods in several experiments. We provide our code at hhttts://ggthiii com/ vieel/r//leeaeFFF
1 Introduction
Generative models have achieved remarkable success in various domains such as image synthesis [Rombach al. 2022]. natural language processing [Brown al., 2020), scientific applications (Noé al. [019) and more. However, the approaches are not directly applicable when dealing with data inherently structured non-Euclidean spaces, which is common in fields such the natural sciences, computer vision, and robotics. Examples include earth science data on a sphere, the orientation of real-world objects given rotation matrix in SO(3). or data on special geometries modeled by meshes or signed distance functions. Representing such data naively using internal coordinates, such as angles, can lead topological issues, causing discontinuities or artifacts.
Luckily, many generative models can be adapted to handle data on arbitrary manifolds. However, the predominant methods compatible with arbitrary Riemannian manifolds involve solving differential equations-stocaare (SDEs) or ordinary (ODEs)- (foss sampling and density estimation [Rozen et al.
2021 Mathieu and Nickel 2020 Huang et al. 2022 De Bortoli al. 2022 Chen and Lipman 2024]
These methods are computationally intensive due to the need for numerous function evaluations during integration, slowing down inference.
To address these challenges, we introduce a novel approach for modeling distributions on arbitrary Riemannian manifolds that cireumvents the computational burden of previous methods. This is achieved by using a single feed-for forwwfdr neural network on an embedding space as generator, with outputs projected the manifold (Fig. TT. We learn this network generalizing the free-form flow framework [Draxler et al., [2024] Sorrenson et al.) 2024) to Riemannian manifolds. The core innovation is estimating the gradient of the negative log- likelhhood within the langent space of the manifold.
normalizing flow, facilitated by
In particular, make the following contributions:
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Function embedding
Project to manifold
Sphere
Torus T
Hyperbolic HF
Curved surface
Figure 1: Manifold Free-Form Flows (M-FFF) learn generative models variety of manifolds. (Left) The learned distributions (colored surface) accurately match the test points (black dots). (Right) We parameterize M-FFF using neural network in embedding space, whose outputs projected to the manifold. This enables simulation-free training and inference, and naturally respects the corresponding geometry, yielding fast sampling and continuous distributions regardless of the manifold.
We extend free-form flows Riemannian manifolds, yielding manifold free-form flows (M-FFF) in Section
M-FFF can easily be adapted to arbitrary Riemannian manifolds, requiring only a projection function from an embedding space.
It only relies a single function evaluation during training and sampling, speeding up inference over multi-step methods typically by two orders of magnitude.
M-FFF consistently matches outperforms previous single-step methods several bench- marks on spheres, tori, rotation matrices, hyperbolic space and curved surfaces (see Fig. [ and Section 5). In addition, it is consistently faster than multi-step methods by two orders of magnitude, while also outperforming them in terms of likelihood in several cases.
Together, manifold free-form flows offer a novel and efficient approach for learning distributions on manifolds, applicable to any Riemannian manifold with a known embedding and projection.
2
Related work
Table 1: Feature comparison of generative models manifolds. We give "V" if any method in category meets this requirement.
Respects topology
Single step sampling
Arbitrary manifolds
Euclidean Specialized
Continuous time
X
M-FFF (ours)
Existing work on learning distributions on manifolds can broadly categorized as follows: (i) leveraging Euclidean generative models; (ii) building specialized architectures that respect one particular kind of geometry; and (iii) learning a continuous time process on the manifold. We compare method to these approaches in Table and give additional detail below.
Euclidean generative models. One approach maps the n-diimensional manifold to R and learns the resulting distribution [Gemici et al. (016)_ Another approach generalizes the reparameterization trick Lie groups by sampling on the Lie algebra which can be parameterized in Fuclidean space Falorsi et al. 2019]. These approaches come with the downside that Euclidean representation

---

may not respect the geometry of the manifold sufficiently, e.g. mapping the earth to plane causes discontinuities at the boundaries. This can be overcome by learning distributions on overlapping charts that together span the full manifold [Kalatzis et al.) [02T]. An orthogonal solution is to embed the data and add noise to it in the off-manifold directions, so that the distribution can be learnt directly in an embedding space I [Brofos et al., 2021]: this only gives access to an ELBO instead of the exact density. Our method also works in the embedding space so that it respects the geometry of the manifold, but directly optimizes the likelihood on the manifold.
Specialized architectures take advantage of the specific geometry of certain kind of manifold to come up with special coupling blocks for building normalizing flows such as SO(3) [iu et al. 2023), SU (d). (d) [Boyda et al. [2021 Kanwar et al. [020]; hyperbolic space [Bose et al., 2020];
tori and spheres [Rezende et al. 22220]. Manifold wree-form Hows are not restricted to one particular
manifold, but can be easily applied any manifold for which an embedding and a projection to the manifold is known. As such, our model is an alternative to all of the above specialized architectures.
Continuous time models build generative model based on parameterizing an ODE or SDE on any Riemannian manifold, meaning that they specify the (stochastic) differential equation in the langent space [Rozen et al. 2021 Falorsi 2021 Falorsi and Forré 2020, Huang et al.) 2022 Mathieu and Nickel 2020 De Bortoli et al. 20222 Chen and Lipman| 2024]] Lou et al. 2020 Ben-Hami et al.
2022 These methods come with the disadvantage that sampling and density evaluation integrates
the ODE or SDE, requiring many function evaluations. Our manifold free-form flows do not require
repeatedly evaluating the model, a single function call followed by a projection is sufficient.
At its core, our method generalizes the recently introduced free-form flow (FFF) framework [Draxler et al. 2024 based on an estimator for the gradient of the change of variables formula (Sorrenson
et al. 2024 We give more details in Section 3.1
3
Background
In this section, we provide the background for our method: We present an introduction to free-form flows and Riemannian manifolds.
3.1 Free-form flows
Free-form flows are class of generative models that generalize normalizing flows work with arbitrary feed-forward neural network architectures (Draxler et al.] 2024
Euclidean change-of-variables Traditionally, normalizing flows are based invertible neural networks (INNs, see Kobyzev et al. 2021] for an overview) that learn an invertible transformation
= =a) mapping from data E R to latent codes = RR.. This gives explicitly parameterized probability density po()) via the change-of-variable
log)o(x) = log log + log (f((),
where fx(x) RR*n is the Jacobian matrix of fo (x) with respect to x, evaluated at 2; 1fg (x) ] is its absolute determinant. The distribution of latent codes pz(z) is chosen such that the log-density is easy to evaluate and it is easy sample from, such standard normal. Normalizing flows can be trained by minimizing the negative log-likelihood over the training data distribution:
min LNLL = min Ex, (- log Pe (z)).
This is equivalent to minimizing the Leiiaac--eibbbe--ivvveeeee and the parameterized distribution KL lpo)). Sampling from the model is achieved by pushing samples from the latent distribution ~ pz through the inverse of the learned function: x fo Po-
between the true data distribution
Euclidean gradient estimator Naively computing the volume change log [x)) in Eq. I1 is
expensive since it contains the full Jacobian matrix (x) E Raxn, requiring ((n) automatic differ- entiation steps to compute. Normalizing flow architectures usually avoid this expensive computation by further restricting the architecture to allow fast computation. Luckily, even if such a shorteut not available, its gradient, which is all we need for training, can still be efficiently estimated as follows:
3

---

Theorem 1 (Volume change gradient estimator, Draxler et al. (2024). ). Let fe R + R" be diffeomorphism. Let I be a random variable with zero mean and unit covariance. Then, the derivative of the volume change has the following trace expression, where = fo(æ):
Vo log [f(())] tr((V tr((Vof(==))f
(3)
Eo 1. (VoS((=))
(4)
Replacing the expected value over v by a single sample, and using stop-grad (SG) operation. Theorem allows us to compute a term whose gradient is an unbiased estimator for the gradient of Eq. 2
(o) R Vo( (ogp)(z))))) - o"
(z)o)).
(S)
Comparing Eqs. and (5) reveals that log [x)( is replaced by single vector-Jacobian of (z) and Jacobian- veector product fg (z)E, each of which require only one automatic differentiation
operation. Note that while the gradient estimate is unbiased, computing the term in the brackets is
informative about LNLL. Thus, for density estimation inference, Eq. (T) explicitly evaluated using the full Jacobian.
Free-form architectures The central idea of free-form flows is soften the restriction that the learned model be invertible. Instead, they learn two separate networks, an encoder fo and decoder Is coupled by a reconstruction loss, cireumventing the need for an invertible neural network fo:
LR Epau(x) [flg)(f(())) - x=1²]-
Together, this gives the loss of free-form flows with B, the reconstruction weight as hyperparameter:
LuF = LNLL + BLR.
(7)
This allows replacing constrained invertible architectures with free-form neural networks. Since fo is not restricted to efficiently compute the volume change, free-form flows use Eq. (5) to compute the gradient of LNLL. To compute Eq. (5), free-form flows approximate fo (z) (which is not tractable) by (z) during training:
VoLNLL 2 Vo(- log v² ( ((())(((((()))
(8)
For density estimation inference, Draxler et al. [2024] recommend using the explicit decoder Jacobian for the volume change.
3.2 Riemannian manifolds
manifold is fundamental concept in mathematics, providing framework for describing and analyzing spaces that locally resemble Euclidean space, but may have different global structure. For example, small region on a sphere similar to the Euclidean plane, but walking a straight line on the sphere in any direction will return back to the starting point, unlike on a plane.
Mathematically, n-dimensional manifold, denoted as M, is space where every point has neighborhood that is topologically equivalent to IR. A Riemannian manifold (M, G) extends the concept of manifold by adding Riemannian metric G which introduces notion of distances and angles. At each point x on the manifold, there is an associated tangent space T.M which is an n--imennional Fuclidean space, characterizing the directions in which you travel and still stay on the manifold. The metric G acts in this space, defining an inner product between vectors. From this inner product, we can compute the length of paths along the manifold, distances between points as well as volumes (see next section).
In this paper, we consider Riemannian manifolds globally embedded into m-dimensional Euclidean space RTH, with n & m. Embedding means that we represent a point on the manifold x E M as vector in Rm confined to an n-dimensional subspace; π : R + M a projection from the embedding space to the manifold. A global embedding is smooth, injective mapping of the entire manifold into Rm its smoothness preserving the topology.
write M C R and denote by
In most cases, we work with, but are not limited to, isometrically embedded manifolds, meaning that the metric is inherited from the ambient Euclidean space. Intuitively, this means that the length of

---

Table 2: Manifolds, global embedding and corresponding projections considered in this paper.
Manifold
Dimension
Emiedddin
Projection
Generic
(an)()((((e))
{ar 1 m(x) m1)
x(x)
Rotations (O(d)
(d 1)d/2
(Q e Raxd : QQ" = Q=1
It at arg mingesote I1Q Rlle: see Eq.
102
Sphere g"
n
(a E R : ||] =1]
22 x/ ///||=||
Torus T = y
n
{X e R [XIl fori= fo.i=1...n
X X for 1...
Hyperbolic H
2" |
min e)/l
path on the manifold is just the length of the path in the embedding space. We note that due the Nash embedding theorem (Nash, [1566, every Riemannian manifold has a smooth isometric embedding into Euclidean space of some finite dimension, so in this sense using isometric embeddings is not a limitation. Nevertheless, for some manifolds (especially with negative curvature, e.g. hyperbolic space) there may not be sensible isometric embedding.
4 Manifold free-form flows
The free-form flow (FFF) framework allows training any pair of parameterized encoder fo(x) and decoder ga(2) as generative model, see Section 3.1 In this section, we demonstrate how to generalize the steps in Section to arbitrary Riemannian manifolds. Note that for simplicity, we choose the same manifold in data and latent spaces, i.e. .xx .Mz M, but the method readily applies to Mx * MM or Gx + Gz as long as they are topologically compatible, like a sphere and a closed 3D surface without holes. The detailed derivations in the appendix consider this generalization.
Manifold change of variables The volume change manifolds generalizes the Euclidean variant Eq. ( by (a) considering the change volume the tangent space and (b) accounting for volume change due to changes in the metric:
Theorem 2 (Manifold change of variables). Let (M, G) be n-dimensional Riemannian manifold embedded in R.T ie.,, M C Rm. Let px be probability distribution on M and let f M -> M be a diffeomorphism. Let pz be the pushforward of px under (i.e., if px is the probability density of X, then pz is the probability density of f(X)).
Let x E M. Define Q E &Rux as an orthonormal basis for T,M and R E Rmx as an orthonormal basis for M.
Then, the probability densities px and pz are related under the change of variables x +> f(x) by the following equation:
R G(f(2))
logpx(x) = logpz(((z)) + log| RRT f'(x) log
(9)
QQ=((TQQI
where Q and R depend on and f(x), respectively, although this dependency is omitted for brevity.
To give an intuition for this result, Fig. 2 shows how the volume change is computed for an isometri- cally embedded manifold, that is G to:
that RRT GR !Q" 1. This simplifies Eq. (9)
log (x) = log P2(()))) + log f()))-
(10)
This is very similar to the familiar change of variables formula in the Euclidean case in Eq. 1 the
only difference being that the determinant is evaluated the n projection of f'(x:) into the
tangent spaces. These projections are necessary as the Jacobian of f is singular in the embedding space, since its action is restricted to the local tangent spaces. See the full proof in Appendix A.I
Manifold gradient estimator We now generalize the volume change gradient estimator in The- orem I to an invertible function on the manifold fe M + M. We find that taking the gradient of the manifold change of variables in Eq. (9) results in essentially the same computation as in the Euclidean case, but the trace in is now evaluated in the local tangent space:
Theorem 3. Under the assumptions of Theorem 2 with f = So. Let v E R be a random variable with zero mean and covariance R.R... Then, the derivative of the change of variables term has the
5

---

Embedding Space
R
T.M
Tig
R
f(a
Figure 2: Computation of the volume change in the tangent space of the manifold: The manifold change of variables formula in Eq. (10) requires to compute the change of a volume element in the tangent spaces under f, which in this example is given by the ratio of lengths of dt and dt'. Since is a map in the embedding space, (z) defines a linear map between vectors from the embedding space. To correctly compute the change in volume, use and R change coordinates to the intrinsic tangent spaces, resulting in the linear map Rd (x) T.M - T( which maps di to di'.
following trace expression, where fo(x):
Ve log R² fo())Q = tr [ R))
(11)
= =E, [ " (Vor)())) (s)u
(12)
This shows that the adaptation of free-form flows for invertible function to isometrically embedded manifolds is remarkably simple (see full proof in Appendix A.2: if the manifold is not isometrically embedded, add the corresponding term in Eq. (9):
VoLMNLL = VolEz, [- log P2(2) - v ( f(())sc(f
(13)
The only change is that v must have covariance RRT rather than the identity. We achieve this by sampling standard normal vectors è E I and then projecting them into the tangent space using the Jacobian of the projection function:
n' (((o(z))).
(14)
Constructing e like this fulfills the conditions of Theorem because Eg [u] and:
Cov [c] Eg < { ((((t)))
RR.
(15)
Just like [Sorrensonn et al. 2024 Draxler al. 2024 we further normalize reduce the variance of the trace estimator. Equation (13) now allows training invertible architectures on manifolds even
the volume change log Rt [=]] is not tractable.
Despite using a stochastic estimator for the gradient, we argue in Appendix A.5 that the scaling of the estimator variance with dimension is comparable to the variance due to stochasticity in flow matching and similar methods.
Free-form manifold-to-manifii to--mnifoodd neural networks As discussed in Section 2 invertible architec- tures have to be specially constructed for each manifold. To overcome this limitation, we now soften the constraint that the learned model be analytically invertible. Instead, we learn pair of free-form coiffoldto--mainnn neural networks, an encoder fo (x) and a decoder go (2) as arbitrary functions on the manifold:
fo(x) M - M, (
(16)
We choose fulfill Eq. 160 using feed-forward neural networks fo, Go Rm RM working in an embedding space Ih of M, but ensure that their outputs lie on the manifold by appending a projection gm + M, mapping points from the embedding space mm the manifold M:
fa(x) x (fo(x)), (0(2) = π (z).
(17)

---

Figure illustrates this for an example circle M =
Just like in the Fuclidean case, employ reconstruction loss to make and de approximately inverse to one another:
LR = B [[og(()()))) -
(18)
We measure the distance the embedding space; one can modify this to use an on on-manifold distance (e.g. great circle distance for the sphere) but we find that ambient Euclidean distance works well in practice, since is almost identical for small distances and this is the regime we work in.
This allows us to substitute So" (z) 2 (=(2) in Eq. (13
VOLM.NLL R VO [- pz() o² (C (=))=]]]
(19)
In Theorem 5 we show that the error of the gradient estimator is bounded by measure of the mismatch between the encoder and decoder Jacobian matrices. When the encoder and decoder are true inverses, the error reaches zero.
Regularization and final loss We find that adding the following two regularizations to the loss improve the stability and performance of our models. Firstly, the reconstruction loss on points sampled uniformly from the data manifold:
Lu = Eu-uum))lloooootto
x=13)]
(20)
helps ensure that we have globally consistent mapping between the data and latent manifolds in low data regions. Secondly, the squared distance between the output of fo and its projection to the manifold:
Lp = (x.. [1) [a(x)]1
(21)
discourages the output of So from entering unprojectable regions, for example the origin when the manifold is S". The same regularizations can be applied starting from the latent space.
The full loss is:
== CML + SRCx + BuCu + BpCp
(22)
where the gradient of LM-NI. NII. is computed using Eq. (19). and BR, Br and Bp are hyperparameters. We give our choices Appendix
Experiments
We now demonstrate the practical performance of manifold free-form flows on various manifolds. We choose established experiments to ensure comparability with previous methods, and find:
M-FFF matches or outperforms previous single-step methods. M-FFF uses a simple ResNet architecture, whereas previous methods were specialized to the given manifolds, hindering adoption to novel manifolds.
M-FFF generates samples faster by typically two orders of magnitude than methods sampling in several steps. Despite this great reduction in compute, it achieves a higher generative quality in several cases.
In our result tables, we mark as bold (a) the best method overall (both single- and multi-step), and (b) the best single-step method. We provide reconstruction losses of method and all details necessary to reproduce the experiments in Appendix B Furthermore, our code is available at phttps //github. httetthubcooomeeeeeeeemmeeeeetee htt./viilllaeerreeeee We each experiment multiple times with different data splits and report the mean and standard deviation of those runs.
Synthetic distribution over rotation matrices The group of 3D rotations S0(3) can be represented by rotations matrices with positive determinant, i.e., all Q E R3x3 with QTQ = and det Q = 1. We choose 33x3 as embedding space and project the manifold by solving the constrained Procrustes problem via SVD [Lawrence et al. 2019] (see Appendix [..2).
We evaluate M-FFF synthetic mixture distributions proposed by De Bortoli et al. 2022 with M mixture components for M = 16, 32 and 64. Samples from one of the distributions samples from our model are depicted in Fig. g
7

---

Figure 3: Manifold free-form flows synthetic S0(3) mixture distribution with M = 64 mixture components proposed by De Bortoli et al.) [022]. (Left) 10,000 samples each from the ground truth distribution and (right) describe a rotation matrix, and then plot the first two angles on the projection sphere and the last by color [Murphy et al. (02T]. We find that our model nicely samples from the distribution with few outliers between the modes.
model. This visualization computes three Euler angles, which fully
Table shows that M-FFF outperforms the normalizing flow developed for SO(3) by Liu et al. 2023 as well as the diffusion-based approaches for the mixtures M and 64.
Table 3: Test negative log-likelihood (NLL, 1) on SO(3) for multi-step and single-step methods. M-FFF consistently outperforms the specialized normalizing flow by Liu et al. [2023) on synthetic distributions of SO(3) matrices, and outperforms multi-step methods the cases with more mixture components. Multi-step baseline values are due to De Bortoli et al. [0022].
M 16
M 64
Fast inference?
Moser flow Rozen al. 2021
-0.85
-0.17 o.
0.49
X: 1000 steps
Exp-wrapped SGM Bortoli et al.
2022 2022
-0.87 + 0.4
-0.16 0.0
0.58 am
X: steps
Riemannian SGM [D Bortoh et al.
-0.89 0203
-0.20 o.
0.49 am
X: 100 steps
SO(3)-NF [Liu et al.
2023
-0,81 = 0.01
-0.12 =: 0.004
0.61 a 0.01
M-FFF (ours)
-0.87 0u2z
-0.21 0.02
0.45 0.2
Earth data on the sphere We evaluate manifold free-form flows on spheres with datasets from the domain of earth sciences. We use four established datasets compiled by Mathieu and Nickel| (0200 for density estimation on S2: Volcanic eruptions INGCC/WDS 2022b). earthquakes [NGDC/WDS
2022a floods [Brakenridge] 2017 and wildfires [EOSDIS] 2020).
Figure T shows an example for model trained on flood data. As the reconstruction error sometimes does not drop to satisfactory level we employ the method described in Appendix B.I I to ensure that the measured likelihoods are accurate. Table (oole4]show shows that M-FFF again outperforms the specialized single-step model; the performance compared to multi-step methods is mixed. We think that multi- step models have an advantage on the considered data, as there are large regions of empty space between highly concentrated data points (see density and sample plots in Appendix B.3
Table 4: M-FFF significantly outperforms the previous single-step density estimator [Peel al.)
2001] on the sphere on real-worke earth datasets in terms of negative log- likeelild (lower is better).
Baseline values collected from De Bortoli et al. 2022 2024..
Huang et al. [022] Chen and Lipman
Volcano
Earthquake
Flood
Fire
Fast inference?
Riemannian CNF Mathicu and Nickel 2020
-6.05 Lost
0.14 an
1.11 amw
-0.80 6si
X: ~ 100 steps
Moser flow [Rozed ot al 2021 Stcrcographic scorc-based
-0.16
0.57 2310
-1.28
X: 100 steps
Bortoli
2022
-0.19
0.59
-1.28
X: 100 steps
Riemannian score-hased [De Bortol
-4,92 1os
-0.19
0.45
-1.33
X: 100 steps
Riemannian diffusion [Huang et al. 2022
-6.61
-0.40
0.43 cim
-1.38 cus
X: >> 100 steps
Riemannian flow matching Chhen and Lipman 2024
-7.93 1147
-0.28 exx
0.42 aus
-1.86 an
X: 1000 steps
Mixture Kent (Pel et al.
2001
-0.80.om
0.33 nos
0.73 sm
-1.18
M-FEF (ours)
-2.25
-0.23
0.51 am
-1.19
/
827
6120
4875
12809
8

---

Table 5: M-FFF consistently outperforms normalizing flows specialized tori (Rezende al. 2020 on torus datasets, without requiring the development of a specialized architecture. In addition. our method comes close to the performance of the multi-step methods and outperforms them on the Glyeine dataset. Baseline values are due to Huang et al. [002)] Chen and Lipman (024).
General
Glycine
Proline
Pre-Pm
RNA
inference?
Riemannian diffusion [Huang al. 2022
1.04 10m2
1.97
0.12
1.24 Luse
-3.70
~ 1000 stops
Riemanmian low matching Chen and Ligman 2024
1.11 eas
1.90 4oo
0.15 aeer
1.18 woes
-3.20 aM6T
1000 sleps
Mixture of power spherical (Hoaps
2022
1.15
2.08
0.27
1.34
4.08
Circular Spline Coupling Flows Rezende al 2020
1.03 -an
1.91 c00
0.21 oo
1.24 -a31
4.01 c03
M-FFF (ours)
1,03
1.89
0.17
1.23
-4.27
Table 6: Test NLL Stanford bunny data proposed by [Chen and Lipman 2024]. living
manifold with nontrivial curvature (see Fig. [ M-FFF outperforms the multi-step model for datasets with more modes.
k 10
I 50
A = 100
Fast inference?
Riemannian Flow Matching (w/ diffusion) [ Chen and Lipman 2024
1.16 tu
1.48 sow
1.53 1:m
X: 1000 steps
Riemannian Flow Matching (w/ biharmonic) Chen and Lipman 2024
1.06 200
1.55 =cm
1.49 =am
1000 steps
M-FFF (ours)
1.21 6M
1.34 =aM
1.28 =aM
Torsion angles of molecules on tori To benchmark manifold free-form flows on tori T", we follow (Huang et al., 2022) and evaluate our model on two datasets from structural biology. We consider the torsion (dihedral) angles of the backbone of protein and RNA substructures respectively.
We represent a tuple of angles (øI...., )n)) E [0, 2n " by mapping each angle to a position on a circle: X, (cos sin (b)) E S Then we stack all X; into matrix X E TRxx2 compare Table 2
The first dataset is comprised of 500 proteins assembled by [Lovell et al. 2003] and is located on T2. The three dimensional arrangement of protein backbone can be described by the so called Ramachandran angles |Ramachandran et al. [963] & and V,, which represent the torsion of the protein backbone around the N-C,, and C. bonds. The data is split into four distinct subsets General, Glycine, Proline and Pre-Proline depending on the residue of each substructure.
The second dataset is extracted from a subset of RNA structures introduced by Murray et al. (2003]. As the RNA backbone structure can be characterized by seven torsion angles, in this case we are dealing with data on T
We report negative ligliikkls in Table 5 finding that M-FFF outperforms circular spline coupling flow, a normalizing flow particularly developed for data on tori [Rezende et al.) 2020) as well as the multi-step methods one of the datasets. In addition to the quantitative results, we show the log densities of the M-FFF models for the four protein datasets inFig. in Appendix B.4
Toy distributions on hyperbolic space We apply M-FFF to the Poincaré ball model, which embeds the 2-dimensional hyperbolic space H2 of constant negative curvature in the 2-dimensional Euclidean space IR,2, as specified in Table 2 As this embedding is not isometric, and distances between points grow when moving away from the origin, we must include the last term of Eq. when changing variables under map this embedded manifold.
We show that M-FFF can be applied to non-isometric embeddings using Eq. (9) and visualize learned densities in Fig. and in Fig. i Appendix B.5) for several toy datasets defined on the 2-dimensional Poincaré ball model. Further details can be found in Appendix B.5
Manifold with non-trivial curvature Finally, we follow Chen and Lipman 20244 and train M-FFF given by synthetic distributions on the Stanford bunny [Turk and Levoy 1994 on the data provided with their paper, see Fig. [ The natural embedding of this mesh is R and we train a separate neural network to project from the embedding space to the mesh. This ensures that the projection is continuously differentiable, which we identify to be important for stable gradients.
Table shows that M-FFF performs well on this manifold, outperforming Riemannian flow matching in two out of three cases. This experiment underlines the flexibility of our model: We only need projection function to the manifold in order train a generative model.

---

Limitations
Manifold Free-From Flows achieve high generative quality on manifolds despite the approximations made during training: First, the exact inverse of the encoder Jacobian is approximated by the decoder Jacobian, which is implicitly regularized via the reconstruction loss (see Eq. [19 ). Second, the final gradient computation in Eq. (8) is estimated with a single v for each item in the batch, adding noise the system.
At inference time, the negative log-likelihoods we report in all tables are based on the decoder Jacobian. We choose this because even if the decoder ends up not be invertible after training (that is several latent codes z yield the same generation I go(=)), the computed densities are a conservative estimate of the true probability density. The downside is that if the reconstruction loss is high, the likelihoods become inaccurate, see Appendix B.I for details, We therefore ensure that the final reconstruction losses are vanishing in Table 8
From a high level perspective, we observe that M-FFF performs less favorable compared to multi-step methods when the density changes sharply very low density regions are present.
7 Conclusion
In this paper, we present Manifold Free-Form Flows (M-FFF), a generative model designed for manifold data. To the best of our knowledge, it is the first generative model manifolds with single-step sampling and density estimation readily applicable to arbitrary Riemannian manifolds. This significantly accelerates inference and allows for deployment on edge devices.
M-FFF matches or outperforms single-step architectures specialized to particular manifolds. also surpasses multi-step methods in several cases, despite reducing the inference compute by typically two orders of magnitude.
Adapting M-FFF to new manifolds is straightforward and only requires selecting an embedding space and projection to the manifold. contrast, competing multi-stey methods are more challenging to adapt as they require implementing diffusion process or computing distances on the manifold.
Acknowledgements
This work is supported by Deutsche Forschuugsgemmeinscaatt (DFG, German Research Foundation) under Germany's 's Excellence Strategy EXC-2181/1 -390900948 (the Heidelberg STRUCTURES Cluster Excellence). is also supported by the Vector Stiftung in the project TRINN (P2019-0022) and by Informatics for Life funded by the Klaus Tschira Foundation. AR acknowledges funding from the Carl-Zeiss-Stiftung The authors acknowledge support by the state of Baden- Württemberg through bwHPC and the German Research Foundation (DFG) through grant INST 35/1597-1 FUGG.
References
Heli Ben-Hamu Samuel Cohen, Joey Bose, Brandon Amos, Maximillian Nickel, Aditya Grover, Ricky TQ Chen, and Yaron Lipman. Matching normalizing flows and probability paths on manifolds. In International Conference on Machine Learning, 2022.
Joey Bose, Ariella Smofsky, Renjie Liao, Prakash Panangaden, and Will Hamilton. Latent variable modelling with hyperbolic normalizing flows. In International Conference on Machine Learning, 2020.
Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez Rezende, Michael S. Albergo, Kyle Cranmer, Daniel C. Hackett, and Phiala E. Shanahan. Sampling using SU(N) gauge equivariant flows. Physical Review D: Particles and Fields, 103(7):074504, April 2021. doi: 10. 1..110/PyysReee..00..... 00.44400
Brakenridge. Global active archive of large flood events, 2017.
James A Brofos, Mareus Brubaker, and Roy Lederman. Manifold density estimation via generalized dequantization. arXiv preprint arXiv:2102.07143, 2021.
10

---

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Ricky TQ Chen and Yaron Lipman. Flow Matching General Geometries. In International Conference on Learning Representations, 2024.
Tim R Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub M Tomezak. Hyperspherical variational auto-encoders. In Conference on Uncertainty in Artificial Intelligence, 2018. 8.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Amaud Doucet. Riemannian score-based generative modelling. Advances in Neural Information Processing Systems, 35:2406-244 2422, 2022.
Felix Draxler, Peter Sorrenson, Lea Zimmermann, Armand Rousselot, and Ullrich Köthe. Free-form Flows: Make Any Architecture a Normalizing Flow. In Artificial Intelligence and Statistics, 2024.
EOSDIS. Land, atmosphere near real-time capability for EOS (LANCE) system operated by NASA's earth science data and information system (ESDIS), 2020.
William Falcon and The PyTorch Lightning team. PyTorch lightning, March 2019.
Luca Falorsi. Continuous normalizing flows on manifolds. arXiv preprint arXiv:2104.14955 2021.
Luca Falorsi and Patrick Forré. Neural ordinary differential equations on manifolds. arXiv preprint (rXiv:2006.06663 2020.
Luca Falorsi, Pim de Haan, Tim R Davidson, and Patrick Forré. Reparameterizing distributions on lic groups. In The 22nd International Conference on Artificial Intelligence and Statistics, 2019.
Mevlana C Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on riemannian manifolds, arXiv: /61 1.02304, 2016.
A Girard. A fast Monte-Carlo cross-validation procedure for large least squares problems with noisy data. Numerische Mathematik, 56:1-23, 1989.
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7855):5777- 2020.
Kevin Sheppard, Tyler Reddy,
Chin-Wer Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron Courville. Riemannian diffusion models. Advances in Neural Information Processing Systems, 35:2750-2761, 2022.
J. D. Hunter. Matplotlib: A 2D graphics environment. Computing in Science & Engineering, 9()): 90-95, 2007.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications Pouristics-Siulaatiio and Computation, 1853):1059-106. 1989.
Jürgen Jost. Riemannian Geometry and Geometric Analysis, volume 42005. Springer, 2008.
Dimitris Kalatzis, Johan Ziruo Ye, Alison Pouplin, Jesper Wohlert, and Søren Hauberg. Density estimation smooth manifolds with normalizing flows. arXiv preprint 2XX::2200003500, 2021.
Gurtej Kanwar, Michael S. Albergo, Denis Boyda, Kyle Cranmer, Daniel C. Hackett, bastien Racanière, Danilo Jimenez Rezende, and Phiala E. Shanahan. Equivariant flow-based sampling for lattice gauge theory. Physical Review Letters, 125(12), September 2020. doi: 10. 10.1100 Sreppyseeeele.... 125, 12.61..

---

Ivan Kobyzev, Simon J.D. Prince, and Marcus A. Brubaker. Normalizing Flows: An Introduction and Review of Current Methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43 (1 (1):3964-3979 2021.
Lars Erik Kühmichel and Felix Draxler. Lightning trainable, 2023.
Jim Lawrence, Javier Bernal, and Christoph Witzgall. A purely algebraic justification of the Kabsch- Umeyama algorithm. Journal of research of the National Institute of Standards and Technology. 124:1, 2019.
Yulin Liu, Haoran Liu, Yingda Yin, Yang Wang, Baoquan Chen, and He Wang. Delving into discrete normalizing flows on SO (3) manifold for probabilistic rotation modeling. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 21264 21273, 2023,
Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser Nam Lim, and Christopher M De Sa. Neural manifold ordinary differential equations. Advances Neural Information Processing Systems, 33:17548 17558, 2020.
Simon C. Lovell, Ian W. Davis, W. Bryan Arendall III, Paul W. de Bakker, Michael Word, Michael G. Prisant, Jane S. Richardson, and David C. Richardson. Structure validation by geometry: o.46 and cB deviation. Proteins: Structure, Function, and Bioinformatics, 50(3): doi: 10. 10.1002//rot.10286
437-450, 2003.
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503-2515, 2020.
Wes McKinney. Data Structures for Statistical Computing in Python. In Stéfan van der Walt and Jarrod Millman, editors, 9th Python in Science Conference, 2010.
Nina Miolane, Nicolas Guigui, Alice Le Brigant, Johan Mathe, Benjamin Hou, Yann Thanwerdas, Stefan Heyder, Olivier Peltre, Niklas Koep. Hadi Zaatiti, Hatem Hajri, Yann Cabanes, Thomas Gerald, Paul Chauchat, Christian Shewmake, Daniel Brooks, Bernhard Kainz, Claire Donnat, Susan Holmes, and Xavier Pennec. Geomstats: A python package for riemannian geometry in machine learning. Journal of Machine Learning Research, 21(323): 2020.
Nina Miolane. Saiteja Utpala, Nicolas Guigui, Luis F. Pereira, Alice Le Brigant, Hzaatiti, Yann Cabanes, Johan Mathe, Niklas Koep, elodiemaignant, ythanwerdas, xpennec, tgeral68, Christian, Tra My Nguyen, Olivier Peltre, pchauchat, mules-Deschamps, John Harvey, mortenapedersen, Maya95assal, Quentin Barthélemy, Patellaoui-Souhall Adele Myers, Felix Ambellan, Florent- Michel, Shubham Talbar, Stefan Heyder, Yann de Mont-Marin, and Marius. Geomstats/geomsttas:: Geomstats v2.7.0. Zenodo, August 2023.
Kieran Murphy, Carlos Esteves, Varun Jampani, Srikumar Ramalingam, and Ameesh Makadia. Implicit-PDF. Non- parametric representation of probability distributions on the rotation manifold. In International Conference on Machine Learning, 2021.
Laura Murray, W Arendall, David Richardson, and Jane Richardson. RNA backbone is rotameric. Proceedings of the National Academy of Sciences of the United States of America, 100: 13913-99-99 9, December 2003. doi: 10.1073/pnas. 1835769100.
John Nash. The imbedding problem for Riemannian manifolds. Annals of mathematics, 63(1):20-6) 63, 1956.
National Geophysical Data Center / World Data Service NGDC/WDS. NCEI/WDS global significant carthquake database, 2022a.
National Geophysical Data Center / World Data Service NGDC/WDS. NCEI/WDS global significant volcanic cruptions database, 2022b.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, (65(6457)ccaw11444 2019.
12

---

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style. nigh-performance deep learning library. In Advances in Neural Information Processing Systems, 2019.
David Peel, William Whiten, and Geoffrey McLachlan. Fitting mixtures of Kent distributions to aid in joint set identification. Journal of the American Statistical Association, 96(453):56 63, 2001.
G. N. Ramachandran, C. Ramakrishnan, and V. Sasisekharan. Stereochemistry of polypeptide chain configurations. Journal of Molecular Biology, 7(1):95-99, 1963. doi: 10.1016/50022-2836666) 80023-6.
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racaniere, Michael Albergo, Gurtej Kan- war, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. International Conference on Machine Learning, 2020.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence- brrreence-based generative modeling on manifolds. Advances in Neural Information Processing Systems, 34: 17669-17680, 2021.
Peter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich, Lea Zimmermann, and Ullrich Köthe. Lifting architectural constraints of injective flows. In International Conference on Learning Representations, 2024.
The pandas development team. Pandas-dev/pandas:: Pandas, February 2020.
Greg Turk and Mare Levoy. Zippered polygon meshes from range images. Proceedings of the 2 annual conference on Computer graphics and interactive techniques, pages 311-318, 1994.
13

---

A Free-form flows on Riemannian manifolds
In this appendix, we will focus on intuitive definitions of concepts from topology and differential geometry. For a more rigorous treatment of these concepts, see (ost 2008).
An n-dimensional manifold M is space where every point has neighborhood which is homeo- morphic to an open subset of I=" Intuitively, this means that there is a small region of M containing which be bent and stretched in continuous way to map onto a small region in R This is what is meant when we say that the manifold locally resembles R", If all these maps from M to lk"" are also differentiable then the manifold itself is differentiable, as long as there a way connect up the local neighborhoods in a differentiable and consistent way.
The tangent space of the manifold at denoted T.M, is n-dimensional Euclidean space, which is a linearization the manifold at x: if we zoom in to a very small region around x the manifold looks flat, and this flat Euclidean space is aligned with the tangent space. Because the tangent space is linearization of the manifold, this where derivatives on the manifold live, e.g. if f MM - Mz a map between two manifolds, then the Jacobian ('(x) is linear map from Ma.Mx to Mz.
A Riemannian manifold (M, G) is a differentiable manifold which is equipped with a Riemannian metric G TMM x TaM E which defines an inner product on the tangent space, which allows us to calculate lengths and angles in this space. The length of a smooth curve y [0, 1] + .M is given by the integral of the length of its velocity vector (t) E M. This ultimately allows us define a notion of distance on the manifold, as the curve of minimal length connecting two points.
In the remainder of the appendix we only consider Riemannian manifolds.
A.1 Manifold change of variables
Embedded manifolds function
We define an n-dimensional manifold embedded in R via projection
π : R m
(23)
where P C RM is the projectable set. We require the projection to have the following properties (the first is true of all projections, the others are additional requirements):
1. π o = = π
2. क is smooth on PP
for all æ e Given such projection, we define manifold by
3. rank(T((F<)))
M = fx E Rm
n (x) = x}
(24)
with the tangent space
T-M col(x'(x
(25)
where col denotes the column space. Since the rank of (xx) with x E is n, the tangent space is n-dimensional and M is an n-dimensional manifold. To avoid clutter we denote the Riemannian metric and its m x matrix representation with G interchangeably. If M is isometrically embedded then G(x) just the identity matrix.
The Jacobian of the projection is a projection matrix, meaning x <'<<>>'()) (xx for x E M. For any v in the column space of (z).. there is a such that e x (z)u and due to the projection property, x' (z)v = N (z)u = v. Similarly, for any w in the row space of ा' r'xx), wT' = w. If π is an orthogonal projection, x is symmetric by definition and hence the row and column spaces are identical.
Integration on embedded manifolds In order to perform integration on the manifold, we cannot work directly in the m m-dimensional coordinates of the embedding space, instead we have to introduce some local ve-dimensional coordinates. This means that the domain of integration has to be diffeo- morphic to an open set in R= Since this might not be the case for the whole region of integration, we might need to partition it into such regions and perform integration on each individually (each such region, together with its map to I is known chart and a collection of charts is atlas).
14

---

For example, if we want to integrate function on the sphere, we could split the sphere into two hemispheres and integrate each separately. hemisphere can be continuously stretched and flattened into a 2-dimensional region, whereas the whole sphere cannot without creating discontinuities.
Given open set U in R", and diffeomorphie local embedding function & U -> M,, the integral of a scalar function p : M + R on o(U) C is
PdV [ (p o)
du.
(26)
Ju
((U)
The integral on the right is an ordinary integral in R". The quantity inside the determinant is known as the pullback metric.
Theorem
2 (Manifold change of variables). Let (Mx, Gx) and (Mz, Gz) be n-dimensional Riemannian manifolds embedded in R.T,, i.e., Mx S Rnm, and assume they have the same global topological structure. Let px be a probability distribution on Mx and let f Mxx + Mz be a diffeommrphism. Let pz be the pushforward of px under f (i.e., px is the probability density of X, then pz is the probability density of f(X)).
Let E Mx. Define Q Rmxo as an orthonormal basis for TaMx and R E Ruxn as an orthonormal basis for M[)
Then, the probability densities px and pz are related under the change of variables x i> f(x) by the following equation:
| G2((x))
logpx(x)= = logp2(f(x)) + log ((T)
2 log
(27)
[Q*Gx(I)QI
where Q and R depend on and f(x), respectively, although this dependency is omitted for brevity.
Below, we provide two versions of the proof, the second being less rigorous and more geometric variant of the first.
Proof. Let Ø R" Mx be defined by Φ(u) xx (2 + Qu). Let U be an open subset of IR" containing the origin which is small enough so that di is bijective. Let v/ : I - Mz be defined by (wr)= = (z( + Rwy). Define yp = wi o o and let W (UU).
Note that ø' (u) (x Qu) Q and hence (0()) = x'x (r) = Q (since each column of Q is in TgMxx = co (z))).
Similarly, (0) = R. Since & is a map from n to m dimensions, there is not a unique function from RT to R which is - on the manifold and there remaining degrees of freedom in the off- manifold behavior which can result in different Jacobians. For our purposes, we define the inverse a i- such that ei r-1 is an orthogonal projection onto Mz. This means v (t (f)=))()==))(ff))) RRT and hence (w (x)( = =R".
Since pz is the pushforward of of px under f. the amount of probability mass contained in @(U) is the same as that contained in I (&)()) = (W):
/ )a( px(x)dVx= JNen Pz(2)dVZ
(28)
and therefore:
[ px(de))yoooo ... dun
dw.. (29)
Changing variables of the RHS with w == ((u) gives us
[ (rof du"
/
iw
du ... du" (30)
ou
15

---

Since U was arbitrary, we can make it arbitrarily small, demonstrating that the integrands must be equal for u 0:
dw
px
(31)
Iu
Since w +5- (f(o()))) the Jacobian has the following form when evaluated at the origin (note ((0) = r)::
Ow -w-1)(((=)) ('x)) d'(0)
(32)
du = {()QQ
(33)
Substituting this into the equality, rearranging and taking the logarithm gives the result:
| (f((x))
logpx(x) (z) log pz (f())) + log RR" (')))) 1oo
(34)
[Q'Gx(")
Alternative proof Here is a less rigorous and more geometric proof, which may be more intuitive for some readers.
Proof. Let a be a point on MMx. Consider a small square region U S .M around x (hypercubic region higher dimensions). If the sides of the square are small enough, the square approximately tangent to the manifold since the manifold looks very flat if we zoom in. Suppose Q is basis for the tangent space at x and q'....,9" are the columns of Q. Suppose that the sides of the square (or hypercube) are spanned by u' = eq' for small The volume spanned by parallelotope (higher-dimensional analog of a parallelogram) is the square root of the determinant of the Gram matrix of inner products:
vol(u', ..... )²)= = Viku', w.
(35)
The inner product is given by G. namely (u,v) v) Gr. We can therefore write the volume of U as
vol(U) ~ c" ViQTGol
(36)
Now consider how U is transformed under f. It will be mapped to region f(U) on MM with approximately straight edges, forming an approximate parallelotope in the tangent space at = = f()).
This region will be spanned by the columns of f'(x)( (since f(x + w²) z f(x) + f' (x))'' and hence will have a volume of
vol(s(U)) ~ e" (
(37)
=c" TTTTTTTTTTRRRRTTRRRRTT ((()))
(38)
=C"IRT RTT
(39)
where is a basis for the tangent space at f(x). We can introduce RR into the expression since is a projection in the tangent space at fxx) and is essentially the identity within that space. Since the RHS of (z) and the LHS f'(x) both live in this tangent space, we can introduce R...T between them without changing the expression. Then in the last step we use that AAB| ||| for square matrices.
The probability density in U and f(U) should be roughly constant since both regions are very small. Since the probability mass both regions should be the same we can write
(x(x) vol(U) = pz(f(x)) vol()(U))
(40)
and therefore
= (z)))= (' / R
(41)
where the approximation becomes exact by taking the limit of infinitesimally small e. Taking the logarithm, arrive at the result of Theorem 2
D
16

---

A.2 Loss function
Under the assumptions of Theorem 2 with - So. Let v E In be a random variable with zero mean and covariance RR Then, the derivative of the change of variables term has the following trace expression, where = [o(x):
Theorem 3
Vo log R² =((=) (Vof(())
(42)
=E, 0² (Vol3())
(43)
Proof. For brevity, we drop the index and denote g = =f First, a reminder that (l(u)
R Q with i = x- f o. Let x = -p- ,.e. X = 0 o v. Jacobi's formula tells us that
d
log I((() - (dA(() A(0-)).
(44)
di
Note also that since x()(())= = u, therefore ((p(())) = I and x'(p(u))= = = Applying Jacobi's formula to of (u):
Ve log Tologls((u) (w))V
(45)
(o((Vo
(46)
and substituting in f and g:
Volog l'
tr tr (V (()))) {(f(
(47)
Q does not depend on 8, but R depends on fxx and hence 6, so it must be considered in the derivative. However,
Vg Votr(RRR)) = =r(()) + RV RV R²)) = (RV
(48)
and since tr(R)() tr(() is a constant, (((VVRR = O. Expanding Eq. (47)
Ve Volog RRT (x) tr(Vo (()) g (f(=)) tr (z)) of(=))
(49)
Since Q an orthonormal basis for TxMx, QQ" is a projection matrix onto T,Mx. This is because (2Q = QQ "TTQQQ=== =Q" using Q" Q = I. As result, QQT=((p) = (z) Since g can also be written inside projection: g(2) = (g()))), therefore g'(2) 7z (=)))) (z), so QQ² g' (2) g' (z). Note also that (==)g'(f(z)) I since fo g id. This simplifies the equation:
Vo log RR" (x)))1 = (r(V R) + tr R Vo (z))
(50)
and finally
V RR f(() ur Vo (f( ()
(51)
In the above proof we used the fact that gQTgg((20 = =='=))))))) g2()),, where we dropped the index 0 and
use g := f-1 for brevity. Can we use RR f'(x) = f'(x) simplify the equation further? No, we cannot, since the expression involving f' is actually its derivative with respect to parameters, which may not have the same matrix structure as f'. Is it instead possible to use g' (=) = gz(()) for simplification? If is implemented as 72 (ō(2)) this is not necessarily true, as g 0((z) might not be a map from the tangent space at 2 to the tangent space at g(2). For example, if we add small deviation u to z, where is orthogonal to the tangent space at 2, then g(2 + v) might not equal (2).
However, this would mean that derivatives in the off-manifold direction can be non-zero. meaning that g(()) * g' RR" v = 0 (since RR will project v to 0). We can change this by prepending by a projection:
g = π og o π..
(52)
If T2 is an orthogonal projection, meaning that T'Z is symmetric, the column space and row space of Tz will both be the same as those R.R... meaning (2) n'Z and hence g' RR) g (z).
This leads to the following corollary:
17

---

Corollary 4. Suppose the assumptions of Theorem hold with f = fe and the following implemen- tation:
-o e So az
(53)
where iT is orthogonal projection. Then the derivative of the change of variables term has the following trace expression, where = =
V log R R fa(x)Q1 = tr(Voo
(54)
Proof. Again, we drop the index and let g = f for brevity. Take the result of Theorem 3 and use the cyclic property of the trace and the properties of g' discussed above:
(=)
(f)) (r(Vo(f((=)) (f())).
(55)
(56)
We use Hutchinson-style trace estimators to approximate the traces given above. This uses the property that, for matrix A E Rox" and a distribution p(v) in R" with unit second moment (meaning E[vv")== I)
Ep Au] = (r((ppu)) Au))
(57)
= (t) wa [A]
(58)
tt(AA)
(59)
meaning that x Au as tr(A) an unbiased estimate of the trace of A.
We have two variants of the trace estimate derived above, one evaluated in TR.,, the other in IR.. The first can be estimated using the following equality (again dropping the index 0 and using g f-1):
tr( R
Ea(c)(u( R (f(f(()))) Rf(
(F)
(60)
Enjnjl" (z) (o) (f()))
(61)
[o (f ()())
(62)
where p(u) has unit second moment in R" and p(v) is the distribution of Ru, which lies in the tangent space at i and has unit second moment in that space by which we mean Bluo"] = RR An
example of such a distribution is the standard normal projected to the tangent space, i.e. v RR ō where ū standard normal.
In the second case, we can just sample from a distribution with unit second moment in the embedding space Rm:
tr (((oo((())) VoLEpos/ Vo
VV (o(((()))SG)) (F())
(63)
A.3 Error bound
The error bound on the gradient from Draxler et 22044 Theorem 4.2] can be readily extended to Riemannian manifolds:
Theorem 5. Carry over the assumptions of Theorem with f So and let go be a manifold-to- manifold function. Let J fz))), Jg. (2), and
1 1 = (z). Then:
|tr (VeJs,)Jg. Vo log R J1, S | (VeJ R Jo. Inll.. (64)
Proof. The proof closely follows [Draxler et al. 2024] and utilizes the Cauchy- -chhwaz inequality for the Frobenius inner product, which states that for matrices A and B, we have tr(A" B) < |||||||B||> Applying this to our case:
|r((RT (VoJs.) R) - Vo Folog|1 RT J
(65)
= -IR(RR (Vo R) tr(R (VoJ (VoJ
(66)
= =|ur(ne" WoJJo))Ja )R)
(67)
18

---

We can re-express this term by introducing the identity matrix in terms of the Jacobians of fo and its inverse:
(VoJJa)) (J Im)
RJ
(68)
By applying the Cauchy-Schwande inequality, we obtain the bound:
< =| (V :Rll |/ Jo R HnllF.
(69)
To further clarify, we recall the function yp introduced in the proof of Theorem 2 This yields:
In J 1.p Q" 1 J.Q
(70)
Thus, we can represent J9. as:
Jog QQ Jos
(71)
Q(Q" Q RR ,. Q))"")1)). J
(72)
= =3, .RRR J
(73)
where we used
Jgp QQ" Jos and J QQ" J,,1 using similar reasoning to the proof of
Theorem
1
A.4 Variance reduction
When using Hutchinson trace estimator with standard normal v E Rt we can reduce the variance of the estimate by scaling o to have length vn (see [Girard] 1989 ). The scaled variable will still have zero mean and unit covariance the estimate remains unbiased, but the variance is reduced, with the effect especially pronounced in low dimensions.
While we can take advantage of this effect in both our options for trace estimator, the effect is more pronounced in lower dimensions, so we reduce the variance more by estimating the trace in an n-dimensional space rather than an m-dimensional space. Hence the first version of the trace estimator, where is sampled from distribution in MM is preferable in this regard.
Let's provide some intuition with an example. Suppose n = 1, m = 2 and R = (1, We want to first sample v RR e with
estimate the trace of A diag(1 (1, 0). Using the first estimator, standard normal which results in [u, 0)² where E R is standard normal. Then we scale so it has length /ni 1. This results in e (r, 0)" where is a Rademacher variable (taking the value -1 and 1 with equal probability). The trace estimate is therefore r² = 1, meaning we always get the correct answer, so the variance is zero. The second estimator samples directly from 2d standard normal, then scales it to have length vm /2. Hence is sampled uniformly from the circle with radius V2. We can write
v2(cos 0, sin 0)T with sampled uniformly in [0. 2n]. The estimate v" Au = 2 cos? 0. This is a random variable whose mean is indeed I as required but has a nonzero variance, showing that the variance higher when estimating the m-dimennionall space.
For this reason, we choose the first estimator, sampling u in the tangent space at x. This also simplifies the definition of g. meaning that don't have to prepend it with projection.
A.5 Scaling behavior
Since generating in high-dimensional spaces can raise concerns about an estimator's 's scaling behavior, argue thcoretically that the free-form flow estimator scales comparably to flow matching as dimensions increase. Experimental results further support this, showing that non- non-Riemannian free- form flows exhibit strong scaling performance in spaces up to 261 dimensions [Draxler et al., (2024)
One metric assess scaling behavior to consider the variance of the gradient estimator, With standard normal noise, the variance of the trace estimator v Au is A| [Hutchinson, 1989 If we apply this to simple linear free-form flow model /(x) = Ax and y(=) Bz, we find that summing up the variances of the gradient estimates with respect to each element of A leads to total variance of 2n tr( B²B".. Note that a converged model, BB" is equal to the covariance of the data. similar calculation for flow matching loss of the form || y/1² leads to a total variance || tr(E)
19

---

with Σ the covariance of p(y(z). These results are proven below. We can see that both expressions scale as n², assuming that ||=||² and the trace terms scale as n. These assumptions are fulfilled if. for example, the data is Gaussian, with covariance that doesn't I depend parameters (elements of A) scales as n², the variance per parameter is constant. We thus expect similar scaling behavior to flow matching, with no problems due to variance in high dimensions.
Since the number of
Lemma Let Raxn be a matrix with entries Aij, and let B E Txxx be any matrix. Define
0A B for each i. Let E R be random vector with entries independently drawn from OAi,
C
the standard normal distribution. Then, the total variance of the Hutchinson estimators o² C for tr(C²) over all i,j is given by
Total Variance
Var (u"Cr)) 2m|| 2ntr(BB'),
(74)
where || || denotes the Frobenius norm.
Proof. We know from Hutchinson
1989 that
Var
(v" = 21
(75)
Note the form of a4 BAg
using the Kronecker delta:
0A aAig
(76)
kI
This implies
Akm
(C Hz = Σ
Bmt = Sxtōmi Bynt = 81g Byt-
(77)
aAi
T
Thus, O has non non-zer zero entries only in the i-th row, and that row is equal to the j-th row of B:
By
i
(C
)m
(78)
0
otherwise.
Next, calculate the Frobenius norm of Cris;
| 11 = (o (BB))² = BB
(79)
k.1
where By: denotes the j-th row of B.
Finally, sum over all and j to find the total variance:
n
Total Variance
B
B
(80)
Since
BB = =IBB2 Σ 1=1
(81)
the total variance simplifies to
Total Variance 2m| B
(82)
Noting that |B = tr(BBT) completes the proof.
Lemma 7. Let A E Rax be fixed matrix, R a fixed vector, and let y 1h be random
vector with conditional distribution p(y]z) having covariance matrix 2, Define the loss function
L(A)
ÔL AA
Then, the total variance of the gradient estimators
over all and j under p(y) x)) is given by
Total Variance
|| tr()))
20

---

Proof. The loss function is given by
1
((A)
| u112
-(4r 4) (Ax 8).
(83)
2
The derivative of L(A) with respect to Aij is
0L
(Ar - v) d(Ar)
(84)
dAg
AAis
Since Ax is vector whose k-th component is (Ax));: Ei=1 A the derivative of (Ax), with respect to Ag is
((Ar)s OR1TG,
(85)
OAi
where des is the Kronecker delta. Therefore, the derivative becomes
aL
= (Ax - 1)) zzz.
(86)
0Ai
ØL
We are interested in the variance of estimator under p(y(z)) Since and x are fixed, the aAi
only randomness comes from y. Assuming that y is a random vector with mean μ Ely [x] and covariance matrix E, we have
ÖL 0Aig
Var
Var ((Ax xg) =3 Var ((Av y)
(87)
Since (Ax- 1):
-0) = (Ax); - and (Ax); is deterministic, follows that
Var ((Ar 3),) Var(y:) Exs.
(88)
Therefore,
0L AAiz
Var
= Ei..
(89)
The total variance over all and j is
ÖL
Total Variance
>Σ Var
0Aiz
=
2 M = tr())t((E))
(90)
1=1
Therefore, we have proven that the total variance is ||a:||2 tr(E).
B Experimental details
In accordance with the details provided Section 4 our approach incorporates multiple regularization loss components in addition to the negative log- ligeliiel objective. This results in the final loss expression:
C CNLL + B
22/2 C + ////////===///////////// + + c//=
(91)
For each of the terms, there is variant in x- and in s-space, indicated by the superscript. In detail: The first loss CR represents the reconstruction loss:
L Eg--pan (dl.gooooo (fo())))
(92)
Cx = Ex- (flo
(93)
Here d(xr, v) = || 12 is the standard reconstruction loss in the embedding space. This could be replaced with distance the manifold. However, this would be more expensive compute and
21

---

since we initialize networks close to identity, the distance in the embedding space is almost equal to shortest path length on the manifold.
In order to have accurate reconstructions outside of training data, we also add reconstruction losses for data uniformly sampled on the manifold, both for and 2:
Cj E [d(z, (o(f))))) f ( (d(2, fz)))
(94)
(95)
Finally, we make sure that the function learned by the neural networks is easy to project by regularizing the distance between the raw outputs by the neural networks in the embedding space and the subsequent projection to the manifold (compare Eq. (7)):
Lp = Ex- 1l (Fo(r) (o(
(96)
Cp Pew fox)) fo(z)]19).
(97)
If these superscripts not specified explicitly in the following summary of experimental details, mean 3 = 3
In all cases, we selected hyperparameter using the performance on the validation data.
B.1 Likelihood Evaluation
Sampling from trained model can be easily achieved by sampling from the latent distribution and performing a single pass through the decoder g. However, in order to evaluate the likelihood of the
test set under our model, in [Draxler al. 2024], we need calculate the change of variable formula w.r.t. g-1
(z)
logpx(x) logpz(9- log R² -=))))()))11 log
(98)
1Q"Gx(=)Q
G2((((()))
R log gff log
(99)
[Q=Gx(")QI
Here we used the approximation f f gg While is expensive compute for training, is reasonably fast to compute during inference time. To show the validity of this approximation, we report the reconstruction losses computed on the test dataset for all experiments in table
Insufficient reconstruction losses. If the assumption of f g-1 is not sufficiently fulfilled the measured likelihoods might be inaccurate. We try to identify such cases by sampling points from M around the proposed latent point with small noise strength o
1 = (f(x) + (O o)).
(100)
As inverse of the decoder g use the sample which results in the lowest reconstruction loss
l (a) Ri arg minl (=())) x:12-
(101)
We also do this with samples drawn around ã = π( + (O, σ)) and uniformly from the manifold I ((M). We note that in most cases except for the earth datasets the likelihoods computed in this way agree with f - In case of the earth datasets, we note that the newly computed likelihoods now agree with observed model quality. Specifically, whenever reconstruction loss is low we also see agreement between the likelihoods computed via f and our approximation of g- via sampling. Otherwise, the disagreement between f and g and the resulting drop in model quality correctly diagnosed. Therefore, for the all g experiments report the likelihoods computed by our approximation.
Sampling based evaluation. In order to evaluate our models independently of their reconstruction capabilities, we propose to also report a sample- based metric. Due the curse of dimensionality, sample-ber based metrics are only tractable in low dimensions. Therefore, we report the Wasserstein-2 distance between model samples and the test dataset, including the standard deviation over multiple training runs, for all 2-dimensional manifolds in Table As competing method have neither reported a sampling based evaluation metric nor published their models, we propose to use our baseline results for future benchmarking.
22

---

Table 7: Wasserstein-2 distances for all 2-dimensional manifolds.
Volcano
0.249 1.0.
Earthquakes
0.068 0,022
Flood
0.047 0,010
Fire
0.072
0,027
General
0.21
Glycine
0.32 1 0.08
Proline
0.51
Pre-Pro
0.47 0.04
Bunny (k 10)
0.09
0.02
Bunny 50)
0.046
0,006
Bunny (k = 100)
0.026
0.007
Table 8: The reconstruction losses LR show that the reconstructed points lie close to the original data. This justifies evaluating M-FFF via negative log-likelssoods.
S0(3) (K = 16)
(2.7 (2.7 0.3) 10
SO(3) (K 32)
(1.6 10 (7.2 10-5
SO(3) (K = 64)
Volcano
(6.1 10-6
Earthquakes
(1.9 10-7
Flood
(1.9 10-5
Fire
x.7=204)x000
General
(7.2 x 10
Glycine
(1.6 10-5
Proline
01.3 x 10-5
Pre-Pro
(1.5+ x 10
RNA
8.9 0.6) 10
Bunny (k: 10)
(2.6 + 0.5)
x 10-5
Bunny (k: 50)
(22 0.5)
10-5
Bunny (k: 100)
(2.3 + 0.5)
10
Hyperparameter
Value ResNet
Layer type
Residual blocks
2
Inner depth
5
Inner width
512
Activation
ReLU
PR
500
8R
0
Bu
10
3p
10
Latent distribution
uniform
Optimizer
Adam
Learning rate
5 10-3
Scheduler
Exponential w/ y = 10 1.0
Gradient clipping
Weight decay
3 10-5
Batch size
1,024
Step count
585,600 3
#Repetitions
Table Hyperparameter choices for the rotation experiments. Bu and 33 are the same for both the sample and latent space.
23

---

Dataset
Number of instances
Noise strength
Volcano
827
0.008
Barthquake
6120
0.0015
Flood
4875
0,0015
Fire
12809
0.0015
Table 10: Dataset overview for the earth data experiments. Each dataset is split into 80% for training. 10% for validation and 10% for testing.
Hyperparameter
Value ResNet
Layer type
residual blocks
4
Inner depth
2
Inner width
256
Activation
sin
gg
10a
0
2 x 102
3p
0
Latent distribution
VMF-Mixture (Neamp 5)
Optimizer
Adam
Learning rate
2 x 10-4
Scheduler
onecyclelr
Gradient clipping
10.0
Weight decay
5 x 10-5
Batch size
32
Step count
1.2M
#Repetitions
5
Table 11: Hyperparameter choices for the earth data experiments. Bu and 3p are the same for both the sample and latent space.
B.2 Special orthogonal group
To apply manifold free-form flows, we project an output matrix R E 3ax3 from the encoder/decoder to the subspace of special orthogonal matrices by finding the matrix Q e SO(3) minimizing the Frobenius norm ||Q R|] r. This is known the constrained Procrustes problem and the solution Q can be determined via the SVD R UEVT Lawrence et al. 2019
Q = USVT.
(102)
where the diagonal entries of were sorted from largest smallest and Diag(I.... det"U.)))
The special orthogonal group data set is synthetically generated. We refer (De Bortoli al. 2022] for a description of the data generation process. They use an infinite stream of samples. To emulate this, we generate data set of 10" samples from their code, of which reserve 1,000 for validation during training and 5,000 for testing. We vectorize the 3 x 3 matrices before passing them into the fully-connected networks. All training details given in Table 9 one training takes approximately 7 hours on a NVIDIA A40.
The data set is synthetically generated; of the N 100, 000 data points, use 1% of for validation and hyperparameter selection and 5% for test NLL. Each run uses different random initialization of papers.
B.3
Earth data
We follow previous works and use a dataset split 80% for training, 10% validation and 10% for testing. For the earth datasets we use mixture of 5 learnable Von-Ntises-Fisher distributions for the target latent distribution.
base our implementation on the hyperspherical. vae library Davidson et al. 2018.. In order to stabilize training apply small amount of Gaussian noise to
24

---

carthquake ke
volcano
flood
fire
Figure Density estimates of our model on the earth datasets. Blue points show the training dataset, red points the test dataset.
Dataset
Number of instances
Noise strength
General
138208
0
Glycine
13283
0
Proline
7634
0
Pre-Proline
6910
0
RNA
9478
10
Table 12: Details on the torus datasets. Each dataset is randomly split into train dataset (80%), validation dataset (10%) and test dataset (10%). During training, we add Gaussian noise with mean zero and standard deviation given by 'noise strength' the data, counteract overfitting.
every batch (see table 10 and project the resulting data point back onto the sphere. Other training hyperparameters can be found in table IT Each model trained around 20h on compute cluster using a single NVIDIA A40.
B.4 Tori
The torus datasets are randomly split into a train dataset (80%), validation dataset (10%) and test dataset (10%). To counteract overfitting, we augment the RNA dataset with random Gaussian noise. The noise strength and total number of instances is reported in Table 12
We use a uniform latent
Hyperparameter Layer type residual blocks
Value (T
Value (T")
ResNet
ResNet
2
Inner depth
3
2
Inner width
256
256
Activation
SiLU
SiLU
B
100
1000
8
100
100
B
100
100
0
1000
3p
Latent distribution
uniform
uniform
Opumizer
Adam
Adam
Learning rate
1 1 10-3 3
1 x 10-3 3 onecyclelr
Scheduler
oneclycleIr
Gradient clipping
Weight decay
1 10-3
1 x 10-3
Batch size
512
512
Step count
120k
120k
#Repetitions
5
Table 13: Details on the model architecture, loss weights and optimizer parameters for the torus datasets. We use the same configuration for all protein datasets on T
25

---

General
Glycine
12
Prolino
Le
Pre-Proline
-6.4
-98
-08
Figure 5: Log density of M-FFF models in the (Φ. v -plane of protein backbone dihedral angles (known Ramachandran plot@amaahandran et (963)). The learned density matches the true density indicated by the test dataset (black dots) very well. Note also that the learned distribution obeys the periodic boundary conditions.
distribution. We train for 120k steps with a batch size of 512 which takes 2.5 to 3 hours on a NVIDIA GeForce RTX 2070 graphics card. Further hyperparameters used in training can be found in Table I3
B.5 Hyperbolic space
straightforward way define distributions on hyperbolic space (but also other Riemannian manifolds) is, to define probability density Prangeni in the tangent space at the origin and use the exponential map expo to pushforward this distribution onto the manifold using Eq.
G maaniiooal (e)p(()))
log Pmanifold (expo (v)) log Prangent (v) log Jeexpo (v) log
(103)
Giangem (v)
where maanifoldd denotes the metric tensor of the embedded manifold and Grangent the metric tensor of the tangent space. This is also known as a 'wrapped' distribution.
We use the Poincaré ball model, which embeds the n-dimensional hyperbolic space H" in the dimensional Euclidean space R" as defined in Table [] The exponential map at the origin of this embedding and its Jacobian determinant are simply given by:
expo (v) = tanh([]u|l)
v ||a||
and
[Jexpo (v) = ||e||| cosh cosh)((|v]]))
tanh([|u|l)
(104)
The metric tensor at some point E defined by: ( = Xooe, with An = 2/((1 11-11111larl ||x||²). The
metric tensor of the tangent space is the usual Euclidean metric tensor: G -āj-
With this at hand, we can define latent and toy distributions as wrapped distributions at the origin, as depicted in Fig. [
For training, sample 100k data points from each distribution. Hyperparameters for each model can be found in Table 14] Training takes approximately 2.5, 16, 8 and 16 hours on a NVIDIA A40 graphics card for the one Gaussian, five Gaussians, swish and checkerboard dataset respectively. The resulting model densities are shown in Fig. 6
B.6
3D mesh
We base our experiment the manifold and data provided by [Chen and Lipman| 2024 using 80% for training, 10% for validation and hyperparameter tuning. We report test NLL on the remaining 10% of the data. Each starts from different parameter initialization. They give the manifold as triangular mesh, consisting of vertices v; E IR3 1,... N, and triangular faces ʃ E (1,..., J.....,,, N₂
Since the projection to the nearest point on the mesh has zero gradient in parts of 3 we instead project to the manifold using a separately trained auto encoder with a spherical latent space. This autoencoder consists of an encoder e 3 3 consisting of five hidden layers with 256 neurons each, SiLU activations and an overall skip connection. The latent codes are computed by projecting encoder outputs e e(x:) to sphere as =(x) (xx)) so that the latent space has the same topology as the input mesh. Then, decoder d(z) with the same structure as the encoder is trained to reconstruct the original points by minimizing | - d(e(x)/ [<<))l We train for 218 262, 144
steps, with each batch consisting of all N, = 502 vertices and additional N = 2, 502 uniformly
26

---

Figure 6: Density estimation on the Poincaré ball model. As latent distribution we use wrapped normal distribution with standard deviation 0.5 (Left). As target distributions (top row) we define several toy distributions in the tangent space at the origin and use Eq. (T03) to push forward the manifold. We will reference each distribution from left to right as 'one Gaussian', 'five Gaussians' "swish' and "checkerboard', We train M-FFF on these target distributions using the full expression in Eq. (9) to compute the change in variables and evaluate the densities of the models (boltom row). M-FFF are capable to adapt to coniisometrically embedded manifolds. The learned densities on the one Gaussian, five Gaussians and swish dataset closely follow the target densities. On the checkerboard dataset, M-FFF cannot fully reproduce the sharp edges and density of the dataset.
Hyperparameter
Value (one wrapped)
Value (five gaussians)
Value (swish)
Value (checkerhoard)
Layer type
ResNet
ResNet
ResNet
ResNet
residual blocks
6
6
Inner depth
3
3
Tnner width
128
256
256
256
Activation
SiLU
SiLU
SiLU
SiLU
PR
1000
1000
1000
1000
Pg
100
100
100
80
100
100
100
100
86
0
Ap
Latent distribution
Wrapped normal
Wrapped normal
Wrapped normal
Wrapped Normal
Optimizer
Adam
Adam
Adam
Adam
Learning rate
2 10
x 10
x 10-4
x 10
Scheduler
Exponential w/ = 0.9986
onocyclelr
onccycleir
onccycleir
Gradient clipping
Weight decay
10
10-3
x 10-3
x 10-3
Batch size
4096
4096
4096
Step count
84k
485k
2400
495k
Table 14: Details on the model architecture, loss weights and optimizer parameters for the Poincaré ball experiments.
random points on the original mesh. We find that for successful training, it is helpful to filter out data with 22 > 0.5 n/10,, 000, where is the step number. This prevents the long bunny ears from collapsing as the are slowly grown, allowing the model to adapt.
We then train M-FFF with the hyperparameters given in table Table 15
using the pretrained autoen- coder as our projection to the manifold. Note that we do not train the distribution on the latent sphere of the encoder, but directly on the manifold spanned by it. Training takes approximately 14 hours a NVIDIA A40.
B.7 Libraries
We base code on PyTorch Paszke et al. 2019]. Py Torch Lightning Falcon and The Py Torch
Lightning team 2019) J, Lightning Trainable Kühhichel and Draxler 2023], Numpy [Harris et al. 2020 Matplotlib Hunter 2007 for plotting and Pandas [McKinney [2010 The pandas development team 2020] for data evaluation. We use the geomstats Miolane et al. 2020 2023 package for embeddings and projections.
27

---

Hyperparameter
Value
Layer type
ResNet
Residual blocks
2
Inner depth
5
Inner width
512
Activation
ReLU
Pg
1000
0
3u
10
100
56
10
Latent distribution
uniform
Optimizer
Adam
Learning rate
5 10-4
Scheduler
Exponential w/ y
1 0.0039
Gradient clipping
1.0
Weight decay
3 10-5
Batch size
1,024 469.199
Step count
#Repetitions
3
Table 15: Hyperparameter choices for the bunny experiments.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Manifold Free-Form Flows are introduced in Section 4 and the experimental performance is demonstrated in Section 5
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer:
[Yes]
Justification:
We provide a dedicated limitations section in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
28

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
The proof for Theorems and are given as Appendices A. and A.2 in
appendix. The necessary assumptions are listed with each theorem and we explicitly state in the main text whenever they are relaxed in practice. We intuition on the proofs upon statement in the main text.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Thcorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The gradient estimator in Theorem 3 can be directly implemented using the text below, and we give detailed instructions in Appendix for each experiment, and released all necessary code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
29

---

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide code to all experiments in Appendix
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7/
gewps///gueeeeeeee CCodeSoodieeeeeee ssi mollicyy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
30

---

Answer: [Yes]
Justification: See Section S and Appendix for these details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We caleulate the mean and standard deviation across for different data splits and initialization using pandas. std Tables to A description of how errors are evaluated is included in Section S and details are included in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the eтror bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, we provide the runtime of each experiment with the training details in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
31

---

9. Code Of Ethics
Question:
Does the research conducted in the paper conform, in every respect. with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We adhere to the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The paper improves generative models on manifolds methodological level. While generative models in general can be used for malicious purposes, we do not see direct positive negative impact for society for the case of generating data manifolds.
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The data considered are synthetic scientific, and as such do not have risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
32

---

Released models that have high risk for misuse or dual-use use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the original datasets in Appendix B and the underlying software in Appendix B.7
Guidelines:
The answer NA means that the paper does existing
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (c.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pataasess.paperswtt paperswi thcode. cow/coo////ooooo /atasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No assests are introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No crowdsourcing nor research with human subjects.
33

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
34


Reviews:
[{"review_id": "94cwAMho29", "reviewer": "Reviewer_m8ZT", "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.", "questions": "My questions to the authors are listed below:\n\n1. 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).\n\n### Additional Comments and Minor Corrections\n\nThe 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:\n\nL62: \"for building normalizing flows such as SO(3)\" -> \"for building normalizing flows for manifolds such as SO(3)\"\n\nL203: 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)\"\n\nL228: \"by rotations matrices with positive determinant\" -> \"by rotation matrices with positive determinant\"", "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).", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "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.\n\n2. 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.\n\n2. 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.\n\n### References\n\n[a] Equivariant manifold flows. https://arxiv.org/abs/2107.08596", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:\n\n### The geometry may be respected very approximately\n\nNo, 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).\n\nThe 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.  \n\nThus, 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. \n\n\n### learning symmetric densities over $SU(n)$\n\nWe need to disambiguate two requirements:\n\n1. Distributions that are *invariant* under the action of  $SU(n)$.\n2. Distributions that assign a probability to the *elements* of $SU(n)$.\n\nThe 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.\n\nWe agree that combining density estimation on manifolds with group invariance (requirement 1) is an interesting future direction.\n\n\n### There are clearly contexts in which this method is suboptimal relative to existing multi-step methods [...] This should also be well noted\n\nMulti-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.\n\nAs 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”.\n\n### 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?\n\nYes, 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.\n\nWe will add this explanation to Section 5 and will add it to an explicit limitation section in the camera ready version.\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b15hJQC0bU", "reviewer": "Reviewer_RvPv", "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.", "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.", "limitations": "The authors adequately addressed the limitations.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written\n- The idea is novel and extends an interesting line of work around a new class of generative models\n- The method seems to perform well in practice, although I have some concerns\n- 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.\n- 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.\n- 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.\n- The $m$ and $n$ dimensions in the proof of theorem 2 seem to be flipped in the appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:\n\n### 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$\n\nYes, 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.\n\nWe will add these numbers and the above explanation to the main paper.\n\n\n### There is no error bound on the likelihood gradient\n\nGood 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:\n\n$ \\\\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 $\n\nNotation: $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).\n\n### 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.\n\nThe 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.\n\nWe hope that this addresses the reviewer’s concerns and look forward to the subsequent discussion.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VAYfIyYfca", "reviewer": "Reviewer_nDWB", "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.", "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?\n2. 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.\n3. 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?\n4. 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).\n5. 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?\n6. 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.", "limitations": "1. There is no mention of limitations related to dimensionality.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n2. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:\n\n### Is it realistic to generate data in high-dimensional spaces? [...] There is no mention of limitations related to dimensionality\n\nYes, 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.\n\nIn 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.\n\nExperimentally, [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.\n\n- [1] Sorrenson et al., Lifting Architectural Constraints of Injective Flows, ICLR 2024\n- [2] Draxler et al., Free-form Flows: Make Any Architecture a Normalizing Flow, AISTATS 2024\n\n### Do you plan to release the source code?\n\nYes, we will release the source code upon acceptance.\n\n### While authors of M-FFF claim their theoretical results operate on arbitrary Riemann manifolds, can the same be said for quotient manifolds?\n\nWe identify two potential pitfalls when applying M-FFF to quotient manifolds:\n\n1. Not every quotient manifold is Riemannian. For example the metric must be constant within each equivalence class.\n2. 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.\n\nFor 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]. \n\nWe 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.\n\n[1] Implementation of geomstats.geometry.grassmannian.Grassmannian\n\n### 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?\n\nYes, 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.\n\n### 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?\n\nIt 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.\n\n### Particularly in spaces with negative curvature, such as hyperbolic space, I am concerned about the appropriateness of using Euclidean distance\n\nWe 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).\n\nWe hope that this addresses the reviewer’s concerns and look forward to the subsequent discussion.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wPxKLeo2la", "reviewer": "Reviewer_Z5qx", "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).", "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)?\n2. what are the numerical log likelihood values relative to the plots in Figure 6? how do they compare to other methods?\n3. 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\n4. 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?\n\nAs 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\n\n[1] Chen et al., Flow Matching on General Geometries, ICLR 2024\n\n[2] Draxler et al., Free-form Flows: Make Any Architecture a Normalizing Flow, AISTATS 2024", "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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- the proposed Theorem 3 is elegant and simple, and allows to easily adapt the framework to the manifold under study\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and helpful feedback! Below, we address the points mentioned:\n\n### No thorough runtime comparison study is performed. [...] How does concurrent work compare [...] at fixed performance [...] in terms of runtime (both training and inference)?\n\nIn a nutshell, M-FFF is over an order of magnitude faster at equal inference performance, and computing gradients is cheap even on arbitrary manifolds.\n\n*Inference*\n\nAt 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.\n\nAs 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.\n\nIn 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).\n\n\n*Training*\n\nOverall, 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).\n\nIn 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.\n\n\n### 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 [...]\n\nM-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.\n\nAmong 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.\n\n\n### what are the numerical log likelihood values relative to the plots in Figure 6? how do they compare to other methods?\n\nWe 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.\n\n\n### How [does] the variance scales with the dimensions in some experiments?\n\nFrom 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.\n\nIn 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.\n\nExperimentally, [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.\n\n- [1] Sorrenson et al., Lifting Architectural Constraints [...], ICLR 2024\n- [2] Draxler et al., Free-form Flows [...] AISTATS 2024\n\n### What are [...] the reconstruction error values?\n\nThey are low, see the table in the rebuttal PDF.\n\n### Limitations are not at all clearly stated as such in the paper. [...]\n\nWe will improve the presentation by adding a dedicated limitations section to the camera ready version. Here are the three main points:\n\n1. 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.\n2. 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.\n3. Inaccurate likelihoods when reconstruction error does not drop low enough, see l. 240-242, Appendix B.1", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their helpful and constructive feedback. Below, we collect and address the main concerns by the reviewers.

 

### Performance comparison (reviewers m8ZT, Z5qx)

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.

Multi-step methods occupy a different spot in the trade-off between competing modeling goals: They have slightly higher generative quality, especially in the presence of sharp boundaries in the density, which comes at the cost of $\approx 100\times$ slower sampling. Despite this, **M-FFF outperforms SOTA multi-step methods in 5 of the 15 experiments** we perform.

To evaluate the performance of multi-step methods with varying compute, we reduce the sampling steps of a Riemannian Flow Matching model trained on protein backbone angle data “general”. We find **that M-FFF is 50x faster than Riemannian Flow Matching [1] at the same quality**.

### Scaling with dimension (reviewers Z5qx, nDWB)

We do not foresee a problem with scaling M-FFF to manifolds with higher dimension given the theoretical scaling of the trace estimator with dimension and the structural similarity to free-form flows [2].

### Reconstruction losses (reviewers Z5qx, m8ZT, RvPv)

We confirm that validation reconstruction losses are low in the rebuttal pdf. We will add these numbers to the camera ready revision.

### FFF error bound transfers to M-FFF (reviewer RvPv)

We confirm that the error bound on the gradient of the likelihood as derived in [2] transfers to the surrogate as proposed for manifold free-form flows (see Theorem 1 in the rebuttal pdf). We propose to add this theorem and a proof in the appendix of 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

We give more details on all points in the respective answer to the reviewers and look forward to the subsequent discussion.

Meta-review: The provides a simple extension of free-form flows to manifolds. While the work is incremental (the driving theory was available in the recent papers by Draxler & Sorrenson) it nonetheless makes the ideas explicit and the empirical demonstrations are valuable. The authors are strongly encouraged to update the paper according to the key discussions, in particular, the low reconstruction loss was generally deemed an important concern. The main empirical observation is the improvement of single-step methods; it would perhaps be reasonable to tone down the claims regarding the comparison to multi-step methods as the submitted text could be misunderstood to imply that the proposed method is competitive with such.
Decision: Accept (poster)