Generalised Flow Maps for Few-Step Generative Modelling on Riemannian Manifolds
Authors: Oscar Davis
Michael Samuel Albergo
Nicholas Matthew Boffi
Michael M. Bronstein
Joey Bose
Venue: ICLR 2026 Poster (2026)
Keywords: generative modelling
Riemannian geometry
few-step generative modelling
Primary area: generative models
TL;DR: New Riemannian generative models that require few (down to 1) sampling steps for SOTA sample quality.

Abstract: Geometric data 
and purpose-built generative models on them have become ubiquitous in high-impact deep learning application domains, ranging from protein backbone generation and computational chemistry to geospatial data. 
Current geometric generative models remain computationally expensive at inference---requiring many steps of complex numerical simulation---as they are derived from dynamical measure transport frameworks such as diffusion and flow-matching on Riemannian manifolds. In this paper, we propose Generalised Flow Maps (GFM), a new class of few-step generative models that generalises the Flow Map framework in Euclidean spaces
to arbitrary Riemannian manifolds. We instantiate GFMs with three self-distillation-based training methods: Generalised Lagrangian Flow Maps, Generalised Eulerian Flow Maps, and Generalised Progressive Flow Maps. We theoretically show that GFMs, under specific design decisions, unify and elevate existing Euclidean few-step generative models, such as consistency models, shortcut models, and meanflows, to the Riemannian setting. We benchmark GFMs against other geometric generative models on a suite of geometric datasets, including geospatial data, RNA torsion angles, and hyperbolic manifolds, and   
achieve state-of-the-art sample quality for single- and few-step evaluations, and superior or competitive log-likelihoods using the implicit probability flow.

Published as a conference paper at ICLR 2026
GENERALISED FLOW MAPS FOR FEW-STEP GENERA- TIVE MODELLING ON RIEMANNIAN MANIFOLDS
Oscar D Michael S. Albergo Nicholas M. Boffi', Michael M. Bronstein '6 Avishek Joey Bose 1.7.8
'University of Oxford, PHarvard University, 'Kempner Institute,
*Institute for Artificial Intelligence and Fundamental Interactions, MIT,
CCanngie Mellon University, AITHYRA, MMila, *Imperal College London
ABSTRACT
Geometric data and purpose-built generative models on them have become ubiqui- tous in high-impact deep learning application domains, ranging from protein back- bone generation and computational chemistry to geospatial data. Current geomet- ric generative models remain computationally expensive at inference-requinn many steps of complex numerical simulation-as they are derived from dynamical measure transport frameworks such as diffusion and flow- mow-migching on Rieman- nian manifolds. In this paper, we propose GENERALISED FLOW MAPS (GFM), a new class of few-step generative models that generalises the Flow Map framework in Euclidean spaces to arbitrary Riemannian manifolds. We instantiate GFMs with three self seff-ditillation-basee training methods: Generalised Lagrangian Flow Maps, Generalised Eulerian Flow Maps, and Generalised Progressive Flow Maps. We theoretically show that GFMs, under specific design decisions, unify and el- evate existing Euclidean few-step generative models, such as consistency mod- els, shortcut models, and meanflows, to the Riemannian setting. We benchmark GFMs against other geometric generative models on a suite of geometric datasets, including geospatial data, RNA torsion angles, and hyperbolic manifolds, and achieve state-of-the-art sample quality for single- and few-step evaluations, and superior or competitive log-likelihoods using the implicit probability flow.
INTRODUCTION
Dynamical measure transport offers a unifying and prescriptive framework for constructing neural network-based generative models that learn to sample desired target distribution by pushing forward a tractable prior. Numerical solution of the resulting dynamical systems has led to popular method families-ai including diffusion models (Song & Ermon, 2020), flow matching (Lipman et al., 2022; Peluchetti, 2023; Liu et al., 2022; Albergo & Vanden Vanden-Eijnden, 2022), and stochastic interpolants (Albergo et al., state-of-the--at results over continuous modalities (Karras et al., 2024; Polyak et al., 2024). While often applied to Fuclidean data such as images, this powerful paradigm naturally extends to data types that are inherently geometric and lie on a known Riemannian manifold, in which case the associated flows and diffusions are defined directly on the manifold (Huang al., 2022b; De Bortoli et al., 2022; Chen & Lipman, 2024). Such geometric generative models have found widespread application in high-impact scientific settings such as rational drug design on the SE(3)* manifold of protein backbones (Bose et al., 2024; Huguet et al., 2024; Watson et al., 2023), generative material design in computational chemistry (Miller et al., 2024), and even discrete data using the Fisher-Rao geometry on the probability simplex (Davis et al., 2024; Cheng et al., 2025a).
together have revolutionised the field and led to
The scalability of dynamical transport-based generative models arises from their use of simple regression-baset objectives that lead to rapid sinulation-free training. However, unlike training, obtaining high-quality generated samples at inference time requires numerical solution of the parametrised dynamical system, which necessitates evaluating the large learnt model numerous times. The computational complexity of inference further burdened in geometric settings, where
"Correspondence oscar daviisccs ax.ac.uE

---

Published as a conference paper at ICLR 2026
each step of simulation requires computing potentially numerically unstable operators that must guarantee manifold constraints for faithful inference. For instance, for matrix Lie groups like S3(3), computing the standard exponential and logarithmic maps requires truncating infinite matrix power series (Al-Mohy & Higham, 2012), which may incur additional sources of error beyond just step-size discretisation error compared to inference on Euclidean spaces. Therefore, reducing the number of Riemannian operations, while maintaining this powerful inductive bias, has the potential lead to higher fidelity generated samples.
The search for accelerated inference with Riemannian Probability Flow dynamical generative models defined on Eu- clidean spaces has spurred the development of fast inference techniques that fall broadly into two methodological families. The first family involves the intricate design of higher-order samplers (Dockhorn et al., 2022a; Zhang & Chen, 2022: Sabour et al., 2024; Karras et al., 2022), which can produce high-quality samples with less than ten steps. The second pursues direct learning of the flow map- the global solution of the governing differential cquation-rathi than the instantaneous vector field used in standard simulation (Boffi et al., 2024). This flow map perspective underlies distillation-baser approaches, the most popular of which include consistency models (Song et al., 2023b; Song Dhariwal, 2023; Lu Song, 2024), which attempt to learn the one-step flow map of pre-trained teacher. Distillation- based methods have also been extended to the multi-step flow map setting we consider here at scale (Sabour et al., 2025). Recent advances have further demonstrated that state-of-the-e one-step generative models can be trained from scratch without a pre-trained teacher through self-consistent self-distillation (Boffi et al., 2025; Frans et al., 2025) Monte Carlo direct training (Geng al., 2025) objectives. These exciting developments naturally raise the question:
Generalised Flow Map
X
X
1
Figure 1: Left: The instantaneous vector field ve of the Riemannian probability flow ODE, Right: The Generalised flow map X.. that jumps along the trajectory of the Riemannian probability flow ODE by steps of size t - s. Both models are depicted on S?.
Can the same flow map principles be extended to Riemannian settings, leading to a new class of few-step geometric generative models?
Here, we answer this question the affirmative by introducing GENERALISED FLOW MAPS (GFM). method that generalises and extends flow maps (Boffi et al., 2025) to arbitrary Riemannian manifolds. Overall, our main contributions can be summarised as:
1. Models. We derive the technical conditions required for flow map respect manifold straints, showing how they lead to three variants of GFM: Eulerian, Lagrangian, and Progressive. 2. Learning. We elevate the self-distillation losses presented in Boffi et al. (2025) to the Rieman- nian setting, showing how each variant can be learned efficiently, and how each recovers their Euclidean counterparts as a specific case.
3. Empirical validation. We comprehensively demonstrate the performance benefits of GFM over previous Riemannian generative models on standard geometric datasets. Specifically. we obtain state-of-thc-ant performance for single- and few-step sample quality on torsion angles found in RNA backbones and protein side-chains, geospatial data of natural disasters on the Earth's sur- face, synthetic data on (3(3) and on the Poincaré disk model of hyperbolic geometry, demon- strating our method's robustness to manifolds with non-trivial structures.
Our methodology opens the door to an under-explored design space of highly expressive generative models defined over geometric data that enjoy rapid inference.
BACKGROUND
2.1 RIEMANNIAN GEOMETRY
Informally, d-dimensional manifold M is topological space covered by family of sets and corresponding maps, called charts, that render the manifold locally homeomorphic to Rd. Requiring C differential structure the charts makes the manifold smooth, and thereby naturally defines
2

---

Published as a conference paper at ICLR 2026
the tangency of vector T.M point E M, where T&M denotes the tangent space at point x, and the disjoint union of all tangent spaces TM == LoEM T&M forms the tangent bundle.
A Riemannian manifold is pair (M,g), where M a smooth manifold, and g is a metric tensor, which smoothly assigns bilinear and symmetrie positive semi-definite inner product Or : T..M x TaM + R. This structure induces notions of curve lengths, distances, volumes, and angles. We assume MM is connected, orientable, and complete, with volume element dz. A vector field is a function from M to TMM such that æ + v(z) E T..M. For a continuously differentiable map. f M + .M, we define its differential at M. dfz : T&M + M,, M, as the map given by df. + (f 7) (0) for any continuously differentiable curve y [0,1] + M,, such that 7(0) = x and (0) = v. The resulting map between tangent spaces is independent of
Manifold operations. In the context of this work, we require few usual operations on Rie- mannian manifolds. Specifically, we require a tangential projection from the ambient space, R. to any arbitrary tangent plane T&M for x M. The tangential projection enjoys closed form: Ps P== -
| where Tr E T.M is basis vector. Roughly, the exponential map, exp, T.M + M. generalises the Euclidean notion of "x + v", and amounts to taking step in the direction of the vector v. The logarithmic map, log, :M - TaMM, generalises "q x" and provides vector pointing towards its argument. Finally, the shortest path between two points p, q E M is called a geodesic and is the curve y [0, 1] + M, with 7(0) = and r(1) = q.
n==
2.2 STOCHASTIC INTERPOLANTS AND FLOW MAPS
We are interested in the dynamical measure transport problem, which seeks to transport an easy- to-sample reference measure Py to a specified target measure f1. As Riemannian manifolds can be embedded within a larger ambient space, M C IR we consider both po; PI E P(M) C P(RR).
Stochastic interpolants in Ted In the context of generative modelling, a target distribution of where 2f id p1(21), interest is provided in the form of a dataset of samples. D { )
thus defining empirical distribution, (11) == 1 - 8(x1 z1). The goal in generative modelling is to build a model pf, with parameters 0 E Q R that approximates oi = Pi in the distributional sense (e.g., W2(P1, pf), or (pr || p while allowing us to draw new samples.
A modern and scalable approach for solving the dynamical measure transport problem-an and by extension the generative modelling problem-is to leverage the framework of stochas- tic interpolants (Albergo et al., 2023). A stochastic interpolant is a stochastic process I [0, 1] x Rd x Rd + Rd that combines, in time, samples from the reference and target measures (t, Fy. x1) + [( x1) axza + 3, for some choice of continuously differentiable functions and 8, such that ou = 1, = 0, Bg = O and B1 = 1. In addition, drawing pairs (1) ~ p( 21) from given coupling must further satisfy the following marginal constraints p(zo). Connecting the measures /n and /r, the interpolant defines a probability path, (P=)ce(0.1) which also follows the corresponding probability flow and transports particles using the following ordinary differential equation (ODE):
p( x1 )dz0 = p(x1) and f (()0,,1)d2
0,I, Ez.~pp(.)) 10,4 4 2,] vx,)). with Zo Po(xa)
(1)
Setting 0; 1 and B, = recovers well known methods such as flow-matching (Albergo & Vanden-Eijnden, 2022; Lipman et al., 2022) and rectified flow (Liu et al., 2022), while imposing optimal transport costs on the coupling results in OT-CFM (Tong et al., 2023).
Flow maps in Rd. An alternative path to generating samples is to instead learn the integrator of the probability flow ODE directly, in order to avoid, at inference time, the costly numerical integration. Introduced in Boff et al. (2024), flow maps are functions X : [0. 1 x XId - Rdd that satisfy the jump any solution of the probability flow. One can thus
condition: (x.,) = Ite where (ar))s(o.1)
sample from P1 by first sampling xo ~ Po and then applying Xx, (za) ~ P1. These are parametrised as x (x.) x + (t s)ef, (2.), for any x, s and t, so that the boundary condition, XL, = Id is automatically satisfied. There are various self self-distillation (i.e., from scratch) objectives for learning flow-map, which are used in conjunction with typical flow- fow-matching objectives (Al- bergo et al., 2023; Tong et al., 2023; Lipman et al., 2023), For brevity, the three (Euclidean) flow map self-distillation loss variations- Fulerian, Lagrangian, and progressive-ant are provided in $A.
s,
3

---

Published as a conference paper at ICLR 2026
3 METHOD
We seek to define a generative model over arbitrary Riemannian manifolds that enables accelerated inference through few-step sampling. Towards this goal, we generalise the notion of flow maps to Riemannian manifolds, yielding GENERALISED FLOW MAPS (GFM) in $3.1. We then show how train such GFM from scratch using oef-distillation losses in $3.2, defined for arbitrary Riemannian manifolds recovering the Euclidean case of Boffi et al. (2025) as special case.
3.1 GENERALISED FLOW MAPS
We begin by defining a flow map in the context of Riemannian manifolds M. To do so, we first recall that the flow map allows us to jump along the trajectory of the probability flow ODE connecting two measures Po. P1 E P(M). The transport described by this ODE can be written in terms of the corresponding Riemannian flow and continuity equations:
0,2: = (x(()) (xpe(x) = -iv, (pe(x))a(z)),
(2)
with 20 Po, and where div, denotes the Riemannian divergence induced by the metric g- As the probability flow ODE lives on M. this presents an immediate point of departure from Euclidean spaces: the interpolant must trace a curve on M. Therefore, the generalised inter- polant, j0, 1] x M + M, satisfying the same endpoint constraints, Lo(zo,
zo and (20, x1) = 31, may not any more be any arbitrary linear combination of zo and æ1, as may not belong the manifold of interest. Consequently, we parametrise the interpolant as the geodesic connecting the points Tu and ø1. I: (t, mo. Ø1) + exPr. (ay log.., (21)), with ay = 0 and a1 1 (Chen & Lipman, 2024). This chosen form of I also enables us to write the vector field of the flow in equation 2 as 0, T, x1 = of log,, (z1)/(1 o4)) E M. Note that when the manifold is an Euclidean space (i.e,, M = RR,, we recover Albergo & Vanden-Eijnden (2022).
Given the stochastic interpolant, we now define GFM that jumps along the trajectory of equation 2.
Definition 1. (GENERALISED FLOW MAPS) Let (M,g) be Riemannian manifold, and lei Po and P1 be two distributions (M,g). The generalised flow map is the unique function X [0,1]2 x M + M. such that, for any solution (11)ee of Equation (2), and any (s,t) E [0, 113, (x,((x))=1x It-
Analogous to the Euclidean case, GFM enables one-step sampling by first sampling 20 ~ po (20) and then applying (zo) ~ Pl. A natural parametrisation for constructing the GFM is (x,(xx,) = exPz, ((t ().),( with the underlying vector field v [0, 1]2 x M + T M. This automatically satisfies the boundary condition (x.) = Is. since expx, (0) F. We can thus characterise GFM in three different ways. generalising the Euclidean case.
Proposition 1. Let X. be parametrised as Xx,(xx) = exp, ((1 - (1)). Then X is the unique GFM for equation if and only if it satisfies any the following conditions:
(i) Generalised Lagrangian Condition:
V(s, 1) E [0, EMM
o.X,,((x)) = ve (x,,,)())))..
(3)
(ii) Generalised Eulerian Condition:
tss E [0, 1]2, 2, M
0xX,,((x)) +
=0,
(4)
(iii) Generalised Semigroup Condition: ((s,t,u)= E [0. 1,,,, x, E M,
Xa, (X,, = (z,(
(5)
We include the proofs of this proposition in $B.1 alongside proofs of the legality of the claims. As stated in Proposition 1, the extension to Riemannian manifolds requires the use of manifold operations such as the differential in place of the Fuclidean gradient. Moreover, instantaneously. the GFM recovers the vector field in its derivative and defines implicit flow.
4

---

Published as a conference paper at ICLR 2026
Lemma 1. (Generalised Tangent Condition) Let vx( = Be 10rl and Ie E M be the drift of equation 2. Then it holds that lim,- 8, a,, (xx) = ve(xt)
(0ll 1. xz for any t
The proof for Lemma I is provided in $B.2 and also illustrated in Figure 1. The lemma underscores the key idea that, for s (on the diagonal of the times space, [0. 1]2), the derivative of the GFM is the instantaneous vector field, U2 of equation 2. Moreover, from the parametrisation Proposition I, and from the fact that lim s.X Vt, for any E [o, 1]. (We thoroughly prove the limit in Lemma 4.) Therefore, we hereinafter use v and v, interchangeably, and we may train the GFM on the diagonal s = t using Riemannian Flow Matching (RFM) (Chen & Lipman, 2024), the loss of which is given by:
= (r,((2) for any æ, follows that
FRFM (0) = E (rro.aa)) (atp:/))))-toommmaaaea) (xe) - 0, 11)1
(6)
3.2 GENERALISED DILT-LISTILAAIONN LOSSES
Given the above characterisations of the GFM in Proposition 1, we define three new generalised self-distillation objectives, leading to Generalised Lagrangian self-distillation (G-LSD), Gener- alised Eulerian self-distillation (G-ESD), and Generalised progressive self-distillation (G-PSD). Specifically, we consider self-distillation objectives of the following form:
[(0) = (0) + LGFM (0).
To adapt to the Riemannian setting, we further constrain our vector field o to lie on the tangent plane at the point where it is evaluated. Specifically, letting f [). M Rd be our underlying neural network, then, for any 0 S S 1, p E M, uh,,
projT, (p). This ensures that all the usual Riemannian manifold operations required are not ill-defined.
Generalised Lagrangian self-distillation. As linear combination of vectors also belong to the same langent space, may freely consider the difference between (L.) and ef,, (1.). However, we further adapt the loss to incorporate the Riemannian metric in the norm of the resulting vector
so that the losses are comparable with the flow matching ones (namely in terms of magnitude).
Proposition 2. (Generalised Lagrangian self-distilllation) The GFM the unique minimiser of the objective in equation 7 for LGFM-SD (0) = LG-LSD(()) where
LG.LSD (uoos.)) [|0,,X2, (1.) (x2 (x))x
(8)
The proof for Proposition 2 is provided in $B.3.
Eulerian self-distillation. an analogous manner to the Generalised Lagrangian self-distillation loss, we can instantiate Eulerian loss to the Riemannian setting.
Proposition 3. (Generalised Eulerian self-distillation) The GFM is the unique minimiser of the objective in equation 7 for LGFM-SD (0) = LG-ESD (0), where
LG-ESD (0) = E,,,a((0,,)))
[10.x% + d(x 11,
(9)
For completeness, the proof for Proposition 3 is included in $B.3.
Progressive self-distillation. Naturally, Xu,t o Xs,. .M - M is well-defined for all (s, t) [0, 113. Thus, we may use the geodesic distance, dg. to derive the Generalised Progres- sive Self-Distillation (G-PSD) objective to enforce the semigroup condition equation 5. The G-PSD objective is the simplest to port, as it is devoid of any spatial or time derivatives of the GFM as stated in the following proposition. The proof is deferred to $B.3.
Proposition 4. (Generalised Progressive Self-Distilaatinn) The GFM is the unique minimiser over 0g of equation for LGFMSD (0) = LG-PSD (9) and (s, t) ~ U(s t), where
(G-PsD(O) = (t,,,n(( [dg (X. (I.), X (x" s. (T.)))
(10)
5

---

Published as a conference paper at ICLR 2026
Algorithm 1: GFM training, for any choice of self-distillation.
Input: Riemannian manifold, (M, g); time distributions, T. S T; coupling, pú batch size, M.
repeat
Draw batch (1z, S;, ry. (rb.=1]f a1)==1111 ~ (T, T, p(
Construct .r = [arg 2i 8*), compute nn 0g. (1 21 (1);
Estimate (R,M(0, (ar) (wt) = (0):
Estimate LGFM-SD s', st,')) (CFMss(());
Optimisation step (R (z)) ()) LaFM.SD(O, (oi) ( (1)) until converged;
Output: Flow map X
3.3 TRAINING GENERALISED FLOW MAPS
While it is possible to implement the GFM objectives naively, it may incur unstable training dynamics. Instead, and in line with the literature (Sabour et al., 2025; Boffi et al., 2025; Geng et al., 2025), we opt for relf-bootstrapped objective by placing a stop-gradient operator to only optimise parts of each objective. This, in turn, converts one term in the objective the "teacher", which is distilled to terms in the loss that have non-zero parameter gradients. We include Algorithm the pseudocode used to train GFM for any self-distillation loss.
G-LSD loss. We place the stop- aop-gradient on the second term of equation 3, yielding the objective:
LOLSD( = (x..2 (1.x),,(T.) stopgrad (x2,,
(11)
Back-propa pocckpropaggting through 0, is simple in usual modern ML libraries, as forward- mode auto- matic differentiation is typically available alongside the Jacobian- Product (JVP) operation.
G-ESD loss. Similarly, for the Eulerian loss, we apply the stop-gradient to the second term, which contains the spatial derivative- and thus avoiding higher order derivatives-leadeeeee to:
ESD (0) (xo,ar)) [|0.xx.. (1.) stopgrad (d(x 10 (4.)) []]
(12)
Note that this results in objective closely related to that of (a Riemannian generalisation of) Mean Flows (Geng et al., 2025)-a connection we prove in $B.4. We also implement the latter, and term this objective G-MF (Generalised Mean Flows), and define it fully in $B.4. This objective is trained as in Mean Flows; that is to say, without the flow matching loss.
G-PSD loss. Finally, also utilise the stop-gradient operator on the PSD objective using the two, smaller steps (s to u, u to f) as the teacher for the larger step (s to f), resulting in the following:
Ec-PDD (0) ((,,a)) [ (xx,, (La), stopgrad (Xf, (x.,.
(13)
We follow Boffi al. (2025) by setting f/s + ft,, leading to two half-steps, thus generalising shortcut models of Frans et al. (2025) to the manifolds connection we also detail in $B.4.
4 EXPERIMENTS
We test the empirical calibre of GFM on a suite of standard geometric generative modelling benchmarks. Specifically, we instantiate GFM on torsion angles (T2 S² x $1) found in protein side chains and RNA backbones T = ( (Lovell et al., 2003; Murray et al., 2003), catastrophic geospatial events on Earth ($2) as introduced in Mathieu & Nickel (2020a), synthetic dataset of the manifold of 3D rotations (SO(3)). and on the Poincaré disk (hyperbolic geometry).
Metrics. To evaluate the few-step sample quality of our models, we use the empirical MMD (Maximum Mean Discrepancy) between the test-set and the samples, with RBF kernel using the manifold's distance, d,, and a bandwidth of к 1. See $C for the calculation details. To assess the learnt vector field, also compute the negative log-likelihood (NLL) the test when

---

Published as a conference paper at ICLR 2026
Table 1: Test NLL on protein sidechain and RNA torsion angles. Standard deviation estimated over
5 runs. indicates baseline numbers taken from Huang et al. (2022b).
General (2D)
Glycine (2D)
Proline (2D)
Pre-Pro (2D)
RNA (7D)
Dataset size
138,208
13,283
7,634
6,910
9,478
MoPS1
0.002
2.08-0.009
0.27++0.008
1.34++0.019
4.08 0.368
RDM (Huang et al., 2022b)
1.0440.012
1.97-0.012
0.1240.011
1.24
-3.70=0.592
(Chen & Lipman, 2024)
1.90 0.005
0.15 0.002
1.18 0.055
G-LSD (ours)
0.99 = U.0S
1.99
0.24
1.11 0.02
-4.15 0.09
G-PSD (ours)
0.95 0.02
1.94
0.08 0.04
1.00
-4.40
G-ESD (ours)
0.99 0.04
1.95
0.19
0.00
00.07
G-MF (ours)
0.97 1 0.01
1.97
0.21 = 0.04
1.02 0.044
-3.99
Table 2: MMD for NFE with the test-set for proteins torsion angles and RNA backbones. Standard deviation estimated over seeds.
General (2D)
Glycine (2D)
Proline (2D)
Pre-Pro (2D)
RNA (7D)
RFM (Chen & Lipman, 2024)
0.45
0.006
0.27 0.008
0.52 1 0.057
0.47 0.022
0.68 1 0.011
G-L.SD (ours)
0.02 0.003
0.03 + 0.004
0.04 + 0.012
0.05 + 0.004
0.08 + 0.007
G-PSD (ours)
0.11 0.016
0.05 0.019
0.07 + 0.011
0.08 0.015
0.14 + 0.027
G-ESD (ours)
0.29 0.002
0.13 0.006
0.44 + 0.024
0.26 = 0.016
0.45 + 0.006
G-MF (ours)
0.029
0.04
0.046
0.017
0,20 1
available. We have included Appendix C.I a discussion explaining why NLL is relevant this context. Finally, we provide qualitative samples to further assert the method's soundness.
Baselines. We use RFM (Chen & Lipman, 2024). which the statc-of-the the-art method, as the main baseline. Additionally, include results for Riemannian diffusion model (RDM) (Huang al., 2022b) and Riemannian score-based generative models (RSGM) (De Bortoli et al., 2022) with test NLL results taken directly from the respective papers due lack of open source code. We also report mixture of power spherical distributions (MoPS) (De Cao & Aziz, 2020) for test NLL.
Finally, we also detailed the computational cost of having using our proposed methods against RFM in Appendix C.4.
4.1 PROTEINS TORSION ANGLES AND RNA BACKBONES ON FLAT TORI
We train model protein dataset the flat 2D and 7D (RNA) tori (S1 x 3 and (S²)7). The 2D data is from Lovell et al. (2003) and the 7D data from Murray et al. (2003), compiled together by Huang al. (2022a). We report our MMD results Tables 1 and and qualitative results in Figures 2 and 3. We observe that the test NLLS produced by GFM outperform those of RFM for the protein side chain datasets, and are marginally worse on the RNA dataset. Most importantly, we test the one-step generative capability that is unique to GFM in Table 2 and Figure 3 and find that GFM offers considerable gains, especially G-LSD, which offers an improvement of up to 22x on the MMD (on "General") for single function evaluation. We also plot Ramachandran plots on the torsion angles in Figure 2, where we observe that our methods achieve log- likeliiod landscapes comparable those of RFM. Finally, of inference steps and find that GFM consistently improve on the MMD for low NFEs, with quality consistently improving with higher steps as RFM.
ablation, also plot in Figure 3 the MMD function
Figure 2: Ramachandran plots on the General protein dataset. Test-set samples depicted in red.
porpet
(a) G-ESD
(b) G-LSD
(c) G-PSD
(d) G-MF
(e) RFM

---

Published as a conference paper at ICLR 2026
Figure 3: MMD on protein datasets against the NFE.
Groerll
Clyritn
Pregro
005
145
0.2
0.5
010
110
0.2
61100
4.2 EARTH CATASTROPHES ON THE 2-SPHERE
We evaluate GFM on collection of geospatial data, represented on the 2-sphere, S The dataset was first introduced in Mathieu & Nickel (2020b) and is curated from various sources (NOAA, 2020a;b; Brakenridge, 2017: EOSDIS, 2020). We report results in Figure 4 and Table 3. Analogous to tori, we observe that the implicit flow uf, within GFM offers log-likelihoods that outperform those of RFM and all other methods three out of four datasets, with the exception of "Volcanoes". To assess sample quality, we again plot the MMD as a function of the number of integration steps, and observe once more great improvement on MMD for all our methods at 1 or 2 NFEs, while preserving high quality for higher NFEs. Finally, we provide the densities learnt by RFM and G-LSD on 2D Earth plots in Figure (with the remaining plots in Figure 7), and remark similar likelihood landscapes between RFM and G-LSD.
Table 3: Test NLL the Earth datasets. Standard deviation is estimated over 5 runs. indicates baseline numbers taken from Huang et al. (2022b).
Volcano
Earthquake
Flood
Fire
Datasel size
827
6,120
4,875
12,809
Mixture of Kent!
-0.00 47
0.33=0.06
0.73=0.07
1.1840.0
Riemannian CNFI (Mathieu & Nickel, 2020a)
00.97
0.19:0.04
0.90:00.03
0.66
Moser Flow (Rozen et al., 2021)
2.02.10.42
-0.0910.02
0.62.10.44
-..0..0.03
Stereographic Score-Based
4.18.10.0
0.04
1.3110.16
0.2810.20
RSGM (De Bortoli al., 2022)
5.5611.26
0.2110.03
0.5210.02
1.24-0..0T
RDM (Huang al., 2022h)
-6.61+00.97
-0.40+0.05
0.43+0.07
-1.38
RFM (Chen & Lipman, 2024)
7.93 1.67
0.28 + 0.08
0.42 0.05
1..8 +: 0.11
G-LSD (ours)
0.66
-0.93
-0.38 0.33
0.42
G-PSD (ours)
+.22
00663 0.13
0.76+0.13
2.48 0.71
G-ESD (ours)
-4.49 0.20
-0.67 0.08
-0.88
-2.29 + 0.08
G-MF (ours)
-3.73 + 0.41
-1.08 = 0.09
-0.72
-2.24 0.30
Figure 4: MMD on Earth datasets against the NFE.
Earthgull
082
0.
0.3
05
G-PED
014
014
014
114
09
52
4.3 SO(3) SYNTHETIC DATA
We next instantiate GFM on the manifold of 3D rotations, SO(3), using the synthetic dataset from Huang et al. (2022a). We compute both test NLL and MMD values and report them in Table 4. We find that, test NLL, all methods perform roughly equally, and there is clear winner. On MMD. we find that versions of GFM outperform RFM, with G-LSD being again the most performant. This demonstrates the effectiveness of all methods non-trivial manifolds.
8

---

Published as a conference paper at ICLR 2026
Figure 5: Plots of densities for the various datasets and all compared methods. Depicted in red are the test-set samples. Datasets from left to right: volcano, earthquake, flood, fire.
RFM
Lattuak
G-LSD
131
-10
130
-30
130
-33
130
Longitale
Figure 6: MMD on the hyperbolic dataset against the NFE.
Table 4: Results on the SO(3) test-set with stan- dard deviation estimated over 5 seeds.
10
MAD
NNFE
CLPSD
RFM
0.147 0.807
0,083 a 0,005
0,942 0.002
0.00
09
G-LSB (ours)
0.064 007
0.009 =: 8.005
77.11 +:0.0
04
G-ESD
G- G-MM (ouma)
0.991 0.007
0.780 4 0.015
0.963 4-000
-4.85 6.000
09
01
4.4 HYPERBOLIC MANIFOLDS
We evaluate GFM on a manifold with a non-trivial metric. in particular, on the Poincaré ball for hyperbolic geometry. We draw 20,000 samples from target distribution that wrapped normal, which we then use to train all methods, including RFM. We report the MMD as function of the NFES in Figure 6, and observe that all versions GFM outperform RFM, even at high NFE, except for G-MF and G-ESD. Indeed, it seems that the variance during training was higher than for the other methods, which caused under-fit the distribution at lower time steps.
RELATED WORK
Riemannian generative models. The most related early efforts to build manifold structure into generative models come from conventional normalising flows built out of iterative coupling layers (Tabak Vanden-Eijnden, 2010; Rezende & Mohamed, 2016; Dinh et al., 2017), where each coupling layer was designed to preserve the manifold structure (Bose al., 2020; Rezende et al., 2020; Kanwar et al., 2020: Boyda et al., 2021; Bose et al., 2021). These ideas have also been extended continuous-time flows and diffusions
general Riemannian structures (Lou et al., 2020; Mathieu Nickel, 2020a; Falorsi & Forré, 2020; Bortoli et al., 2022), and their optimisation has also been made simulation-free (Rozen et al., 2021; Chen & Lipman, 2024; Bose et al., 2020; Davis et al., 2024). Here, our focus is on extending these approaches so that the more efficient any-step flow map is well-posed and learnable general geometries.
Accelerated inference in generative models. Early work on accelerated inference focused on a macher-student procedure (Song et al., 2020a: Luhman & Luhman, 2021: Salimans & Ho, 2022; Meng et al.. 2023) where an expensive inference model is distilled to produce the same output in fewer steps. Orthogonal this have been efforts to parallel diffusion inference with adaptive (Chen et al., 2024; Dockhorn et al., 2022b; Shih et al., 2023; Tang et al., 2024) speculative (Bortoli et al., 2025) schemes. Consistency models aim to directly learn the one-step map to the data distribution from any point along the trajectory (Song et al., 2023b; Kim et al., 2024; Song & Dhariwal, 2023). The flow map (Boffi et al., 2024) has emerged as a unifying picture, and recent efforts have shown how to distill it (Sabour et al., 2024) directly learn it (Boffi et al., 2025; Geng et al., 2025) using the equations that characterise We take the leap here to show how to generalise this complete class of models the Riemannian setting for performant ends.

---

Published as a conference paper at ICLR 2026
Concurrent work. Concurrent and most related to work is that of Cheng al. (2025b). The authors propose a new method, Riemannian Consistency Models (RCM), to train few-step generative models on Riemannian manifolds from scratch. Their approach directly ports the original Consistency Models (Song et al., 2023b) to Riemannian geometries, relying on more sophisticated geometric constructions. In contrast, GFM enjoys simpler practical instantiations it relies on the self-distillation of the Flow Maps framework (Boffi et al., 2025) and recovers shortcut models as special case of the PSD objective executed Riemannian manifolds.
CONCLUSION
We propose GENERALISED FLOW MAPS, new class of geometric generative models that capable of performing ftw-step inference on arbitrary Riemannian manifolds. To build GFM, we provide three equivalent theoretical conditions that characterise a flow map on manifolds, and three corresponding self-distillation objectives. We demonstrate the empirical performance of GFM the low NFE regime and achieve state-of-the-art results sample-based metrics, and competitive test likelihoods in comparison to RFM. While each GFM condition leads to a different corresponding objective, at present, the Lagrangian objective remains the most performant in generating high-quality samples, and understanding this empirical observation from a theoretical lens a natural direction for future work. Additionally, as we empirically demonstrate, the implicit flow within the flow map may lead to better NLL than the flow learned through RFM, which points to an interesting direction for gaining theoretical understanding in future work.
ACKNOWLEDGEMENTS
OD is funded by both Project CETI and Intel. AJB is partially supported by an NSERC Post-doc fellowship. This research is partially supported by the EPSRC Turing AT World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub No. EP/Y028872/1. MSA is supported by Junior Fellowship the Harvard Society of Fellows as well as the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, http:/iaifi.org/ This work has been made possible part by gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence.
ETHICS STATEMENT
We hereby acknowledge and declare to have abided by the ICLR Code of Ethics. Our current work does not utilise any sensitive data, nor does it directly enable nefarious usage, although other derivative works could use ours in unethical contexts. The authors declare no competing financial interests.
REPRODUCIBILITY STATEMENT
All of results are reproducible and have been run with the same 5 random seeds that were set in advance, and that were never changed throughout our evaluation, to provide us our standard deviations. The datasets be made public upon acceptance alongside detailed instructions for running it locally, and the exact configuration files that were used.
entirely public and freely available to all. Our code will
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Awad H Al-Mohy and Nicholas J Higham. Improved inverse scaling and squaring algorithms for the matrix logarithm. SIAM Journal on Scientific Computing. 34(4) 2012. (Cited on page 2)
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic inter- polants, arXiv preprint arXiv:2209. 2022. (Cited on pages 1, 3, and 4)
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A uni- fying framework for flows and diffusions, 2023. URL htttps://arxiv.ortt 08797. (Cited on pages and 3)
Nicholas M Boffi, Michael S Albergo, and Eric Vanden-Eijnden. Flow map matching. arXiv preprint 0rXiv:2406.07507,, 2, 2024. (Cited on pages 2, 3, 9, 16, and 21)
Nicholas M. Boffi, Michael S. Albergo, and Eric Vanden-Eijnden. How build consistency model: Learning flow maps via self-distillation, 2025. URL https: ittps://arxixs /atxivaaavvv/aa//////eee/eawaal/ org/ abs/ 2505. 18825. (Cited on pages 2, 4, 10, 16, and 17)
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnau Doucet. Riemannian score-based generative modelling, 2022. URL https: org/abs 2202 03763. (Cited on page 9)
Valentin De Bortoli, Alexandre Galashov, Arthur Gretton, and Araad Doucet. Accelerated dif- fusion models via speculative sampling, 2025. URL https: arxiv. 05370. (Cited on page 9)
Avishek Joey Bose, Ariella Smofsky, Renjie Liao, Prakash Panangaden, and William L. Hamilton. Latent variable modelling with hyperbolic normalizing flows, 2020. URL https://arxxis org/abs/ 2202/ 06336. (Cited page 9)
Avishek Jocy Bose, Marcus Brubaker, and Ivan Kobyzev. Equivariant finite normalizing flows
arXiv preprint arXiv:2110.08644 2021. (Cited page
Avishek Joey Bose, Tara Akhound Sadeeh, Guillaume Huguet, Kilian Fatras, Jarrid Rector- Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Se(3)-stochastic flow matching for protein backbone generation, 2024. URL https: // ///axiiv.rrt wwwpiv..org/ali abs/ 2310. 02391. (Cited page 1)
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 equiv- ariant flows. Phys. Rev. D, 103:074504, Apr 2021. doi: 10. 10..hs.p/ps 103.003.0000 100.0004000 URL 103 074504. (Cited on page 9)
https: aps org/ /tw/c/..111111/// 201111//PPyyaae
G. R. Brakenridge. Global active archive of large flood events. http: /ttf///ooobberrtttooo/ colorado. aacuuArccivess. index..... html, 2017. Dartmouth Flood Observatory, University of Colorado. (Cited page 8)
Haoxuan Chen, Yinuo Ren, Lexing Ying, and Grant M. Rotskoff. Accelerating diffusion models with parallel sampling: Inference at sub-linear time complexity, 2024. URL //arxiv.. org//bs//444 2405. 15986. (Cited on page 9)
Ricky T. Q Chen and Yaron Lipman. Flow matching on general geometries, 2024. URL https: // //arxiv.. org/ abs / 200. 03660. (Cited on pages 1, 4, 5, 7, 8, 9, and 21)
Chaoran Cheng, Jiahan Li, Jian Peng, and Ge Liu. Categorical flow matching statistical mani- folds, 2025a. URL https: arxivv org/abs/ 22g//bbs/24055 /205 16441. (Cited on page 1)
Chaoran Cheng, Yusong Wang, Yuxin Chen, Xiangxin Zhou, Nanning Zheng, and Ge Liu. Rieman- nian consistency model, 2025b. URL https: arxiv. org/aba/2510. 00983. (Cited on
page 10)
11

---

Published as a conference paper at ICLR 2026
Oscar Davis, Samuel Kessler, Mircea Petrache, ismail Ilkan Ceylan, Michael Bronstein, and Avishek Joey Bose. Fisher flow matching for generative modeling over discrete data, 2024. URL https: //arivvvv... org/ abs 2405. 12244.15.66.... 14664. (Cited pages and
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornon,, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. Advances in neural information processing systems, 35:2406-2422, 2022. (Cited on pages 1, 7, and 8)
Nicola De Cao and Wilker Aziz. The power spherical distribution. arXiv preprint arXiv:2006.0445) 2020. (Cited on page 7)
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp, 2017.
URL https: /taps///rree aariiviiivv.////iv///a.a///aaavv org/ atp///ass/11// 08803. (Cited on page 9)
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers.
Advances in Neural Information Processing Systems, 35:30150-30166, 2022a. (Cited on page 2)
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers, 2022b. URL httptttpp//// /ttp//aaavvv////////aa//laaa.aaa org abs/ / 2210. 054 75. (Cited on page 9)
EOSDIS. eart prrtt-obseevation-cee
nasa gov/ /neae- real-elllee -ime firmm/// /twt/vitiveeeeeee fire data, 2020. Land, Atmosphere Near real-time Capability for EOS (LANCE) system operated by NASA's Earth Science Data and Information System (ESDIS). (Cited on page 8)
Active fire data.
https://earthharr
Luca Falorsi and Patrick Forré. Neural ordinary differential equations on manifolds, 2020. URL htttps:/arxiv.orvt wwtp:///ggg//b////////oo// 2006. 0 66633. (Cited on page 9)
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models, 2025. URL www.h:ttp://// http:v///axxvvooo org/ abs 2410 12557. (Cited on pages 2, 6, and 21)
Zhengyang Geng, Mingyang Deng. Xingjian Bai. J. Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling, 2025. URL htt.t//se//aaa..ooorrr/////../.. (Cited on pages 2, 6, 9, and 18)
Chin- Wei Huang, Milad Aghajohari, Avishek Joey Bose, Prakash Panangaden, and Aaron Courville. Riemannian diffusion models, 2022a. URL htttps://are www./or/////...ooooo../.oooo..be (Cited on pages and 8)
httg/ab//22200 /2008. 07949.
Chin-Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Rie- mannian diffusion models. Advances in Neural Information Processing Systems, 35:2750-2761, 2022b. (Cited on pages 1, and
Guillaume Huguet, James Vuckovie, Kilian Fatras, Eric Thibodeau-Lauffe Pablo Lemos, Riashat Islam, Chenghao Liu, Jarrid Rector-Brooks, Tara Akhound-Sadend Sadegh, Michael Bronstein, et al. gequence-augmented se (3)-flow matching for conditional protein generation. Advances in neural information processing systems, 37:33007-33036, 2024. (Cited on page
Gurtej Kanwar, Michael S. Albergo, Denis Boyda, Kyle Cranmer, Daniel C. Hackett, Sébastien Racanière, Danilo Jimenez Rezende, and Phiala E. Shanahan. Equivariant flow-based sampling for lattice gauge theory. Phys. Rev. Lett., 125:1,, Sep 2020. doi: 10. so. 125. 125.121601. URL https: http//nnkk..o. aps org/ doi 1tt.://1111111111..1111111111111 www.1/1/yyeeeeeeettt....o 125. 121601. (Cited on page 9)
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565-266 26577, 2022. (Cited page 2)
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyz- ing and improving the training dynamics of diffusion models. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24174-241 2024. (Cited on page 1)
12

---

Published as a conference paper at ICLR 2026
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Ue- saka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion, 2024. URL htttp::/aaai...o//aabbst..e 02279. (Cited on page 9)
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. (Cited on pages and 3)
Yaron Lipman, Ricky T. Q. Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, 2023. URL attpe://arxiv.orgg abs pages 3 and 21)
wttp:ii..org/ab///211111 02747. (Cited on
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022. (Cited on pages I and
Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser-Nam Lim, and Christo-
pher De Sa. Neural manifold ordinary differential equations, 2020. URL https://arxis
org/abs/ 2006. 6. 0224. (Cited on page 9)
Simon C Lovell, Ian W Davis, W Bryan Arendall, 3rd, Paul W de Bakker, J Michael Word, Michael G Prisant, Jane Richardson, and David C Richardson. Structure validation by calpha geometry: phi.psi and cheta deviation. Proteins, 50(3):437-450, February 2003. (Cited pages and 7)
Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.1108/ 2024. (Cited on page 2)
Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-tin time consistency models, 2025. URL https:
11081. (Cited on page 18)
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed, 2021. URL
wwg. /2101.10110001 02388. (Cited on page 9)
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. arXiv preprint arXiv:2006. 10605, 2020a. (Cited on pages 6, 8, and 9)
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows, 2020b. URL https: / arxiv. www./brggsssooo.. 2006. 10605. (Cited on page 8)
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik P. Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models, 2023. URL htttpps///aaii....t// abs/2210 03142. (Cited on page
Benjamin Kurt Miller, Ricky TQ Chen, Anuroop Sriram, and Brandon M Wood. Flowmm: Gener- ating materials with riemannian flow matching. arXiv preprint arXiv:2406.04713, 2024. (Cited on page 1)
Laura J. W. Murray, W. Bryan Arendall, David C. Richardson, and Jane S. Richardson. Rna back- bone is rotameric. Proceedings of the National Academy of Sciences, 100024):13944113999 2003. doi: 10. 10.1103/pnas.188 17733/mms.18357611000 1835769100. URL https:/ /www///www/w///w/t//////////.t pnw.ppasssco wttp:/.pt/////o//////////p/ooo// 1073/pnas 1835769100. (Cited on pages 6 and 7)
NOAA. Global significant earthquake database.
/ov/ arg--binniee fiiinisoiid=goo nooa. ngdc. mgg wagg..aan www.rdazza.:aaaaa. :otpearrdd::::111000 G012153, 2020a. National Geo- physical Data Center / World Data Service (NGDC/WDS); NCEI/WDS Global Significant Earth- quake Database. NOAA National Centers for Environmental Information. (Cited on page 8)
https: http///aaaa./// ttw./.oaaaoooooaoooooooooooooo..
NOAA. Global significant volcanic eruptions database. https: : /aata.aaa.aammam/aaao www.ooo..oo.ooo noaa. gov/ egi-bin/ist /ss//ssoiii id=gov. noaa. .gdcc sogydc.mqg... mag hazards G10147, 2020b. National Geo- physical Data Center World Data Service (NGDC/WDS): NCEI/WDS Global Significant Vol- canic Eruptions Database. NOAA National Centers for Environmental Information. (Cited on page 8)
13

---

Published as a conference paper at ICLR 2026
Stefano Peluchetti. Non-denoising forward-tims diffusions. arXiv preprint arXhv:2312.145888 2023. (Cited on page 1)
Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Chih-Yao Ma, Ching-Yao Chuang, al. Movie gen: A cast of media founda- tion models. arXiv preprint arXiv:2410. 13720, 2024. (Cited on page I)
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows, 2016.
URL httph:tt///c/... pectps://aziive /ttpivvvvvvv.////e//eee/eaaaeava wrg. aba 05770. (Cited on page 9)
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racanière, Michael Albergo, Gurtej Kanwar, Phiala E. Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres, 2020.
URL httph:tp///// /tttpp::/aaxiv..oo
abs / 02 22428. (Cited on page 9)
Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence- based generative modeling on manifolds. Advances in Neural Information Processing Systems, 34, 2021. (Cited on pages and 9)
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your steps: Optimizing sampling schedules in diffusion models. arXiv preprint 145007 2024. (Cited on pages and 9)
Amimmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your flow: Scaling continuous- flow map distillation, 2025. URL /ettps://axxvv.or//aatt pages and 6)
14603. (Cited on
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models, 2022.
URL httpt::t/aaiiiiiiii org/ abs/ 2202. 00512. (Cited on page
Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffiusio models, 2023. URL https: /eatts://arxiv.... a//aarxivvv....///aa wwg////aaaoo aab//ab/// 22002 11055.... 1 6317. (Cited on page 9)
Jiaming Song, Chenlin Meng, and Stefano Ermon. (Cited on page 9)
Denoising diffusion implicit models, 2020a.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310.141/ 2023. (Cited on pages 2 and 9)
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution, http://wwwww....aa//aaao/ wwtpabs////aao.ooooooo........ 907. 05600. (Cited on page 1)
2020. URL https: /tt////axiii arxiv.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. baore-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13345 2020b. (Cited on page 21)
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models, 2023a. URL
https: //taiiivv./////////////////aaavl wwwp:iivoorr//oooooooooo alss 22033. 01469. (Cited on page 18)
Yang Song, Prafulla Dhariwal, Mark Chen, and llya Sutskever. Consistency models. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023b. (Cited on pages 2, 9, and 10)
Esteban G. Tabak and Eric Vanden-Eiinnden. Density estimation by dual ascent of the log-likelihood. Communications Mathematical Sciences, 8(1):217 233, 2010. (Cited page 9)
Zhiwei Tang, Jiasheng Tang, Hao Luo, Fan Wang, and Tsung-Hui Chang. Accelerating parallel sampling of diffusion models. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver. Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 47800-47818 PMLR, 21-27 Jul 2024. URL. https: /ttes://prroce ptoc/edoeeeeeeepeeooooooccceeeee ngs mllr. press/v233 21355/tan224555 html. (Cited page 9)
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet. Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302. 00482, 2023. (Cited on page 3)
14

---

Published as a conference paper at ICLR 2026
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eise- nach, Woody Ahemn Andrew Borst, Robert Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 620(7976;:1089-11000 2023. (Cited on page I)
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arX1v:2204. 13902, 2022. (Cited on page 2)
15

---

Published as a conference paper at ICLR 2026
A SELF-DISTILLATION ON EUCLIDEAN SPACES
In this section, we provide a self-contained review of the three self-distillation based objective func- tions introduced by Boffi et al. (2025). These extend upon the original distillation-baat objectives introduced in Boffi et al. (2024). To this end, we first recall three key properties of the flow map.
Proposition 5 (Flow map characterization). Let d, = (x())) denote a probability flow ODE, and let Xa.t [0, 1]2 H& denote its flow map. Then,
Xa is the unique solution to the Lagrangian equation,
o(X,,((x) (x))
(2,8,t) 10, 1]3.
(14)
ii. Xs is the unique solution the Eulerian equation,
0,X,,((z) + (z),(z)), =
A (Z,S,t)) 10,1]
(15)
iii. Xs, satisfies the semigroup condition,
Xx,(2)) (X,,
(z,,su,,t) E R x [0,
(16)
Proof. We simply prove the forward direction: that the flow map satisfies each property. The proof of uniqueness follows similarly and can be found in Boffi al. (2025).
Each property follows from the defining condition
Xx,( = It
(s, E 10.1
(17)
Taking the derivative with respect to time, we find that
0 (x 0,st: = U, (x()),
(18)
which is the Lagrangian equation. Noting that x, was arbitrary completes the proof.
Taking the total derivative of equation 17 with respect to s, we find
(xX,,((x.) (x,) 0, O,X,, +
(19)
= O.
Again, noting that x, was arbitrary completes the proof.
The semigroup condition follows simply by noting that
Xu.:(Xa.u (xx)) (x,( = -Xaa(xo).
(20)
This completes the proof.
As discussed in the main text, from the Lagrangian condition in equation 14, we may observe that
lim (x),, = lim vi vx(Xx,(<))) = (x)
(21)
8-2
aut
assuming continuity of X and using that (x. for all (x, E Rdd [0, 1). Moreover, using s)ve,(())
the parameterization
Xx,(2)
(22)
find that
lim 0,X..((t) = lim ((t s)0,v, + v22( vx,(T) (x))..
(23)
This observation immediately leads three training the flow (x) on the diagonal s = t via flow matching, and simultaneously distilling into a flow map by minimizing the square residual one of the conditions in Proposition 5.
schemes, each of which proceeds by
16

---

Published as a conference paper at ICLR 2026
Proposition 6 (Euclidean self-distillation (Boffi et al., 2025)), Let (²) + (t s))),(())
denote a candidate flow map, and assume that v continuous in both time arguments. Moreover, let
- ((L((() O.le 2
(24)
denote the flow matching loss on the diagonal = t. Then, the ideal flow map Xs,( - x + s)Us, (z) corresponds to the unique minimizer over v of each of the following objective funiiiinns
The Lagrangian self-distillation loss,
(LSD( = Lo(v") + [. (() -f, (T,,(x dsdt,
(25)
it The Eulerian self-distillation loss,
LESD (0) = L6((") + [1
dsdt,
(26)
iii The progressive self-distillation loss,
CpsD(ü) Ls((²)
fo ( E [. (T.. (1 7),,, (T.))) dsdtdy, where (1 - ?)s + yt with y E [0,
(27)
Proof. For completeness, and for ease of the reader, we sketch the proof from Boffi et al. (2025). We first observe that (b()) 2 LE(D) where (x)) = I = x is the ideal flow. From this, it follows that (ū) Z Lv(,), (0) (L(v), and LPSD 2 Lo(u). because the second term in each loss non-negative. Moreover, by Proposition 5, the ideal flow map satisfies the Lagrangian equation 14 and the Eulerian equation 15, that the second term is zero for both LisD(() (0) and (ū) at the true flow map.
It remains show that the second loss in equation 27 imposes the semigroup condition. To see this, we observe that we may write the semigroup condition as
Xx,(æ)) (X,
z + sz),,(()) (z, + u)u
x (t (z) + (u (x) + (t u)un, (No,u
(t ())),(()) ((1 - -7)s+++ (,)),((x)
(28)
+ + (1-7)s- - (X (z)), s),,,( - s)u. + (1 (1 8) s)Ua, (X,,,()))
U1 7V,, (z) + (1-7)u,,,, (z))
From the sequence of operations in equation 28, we see that the second term in equation 27 pe- nalises the square residual on a rescaled semigroup condition. The minimiser therefore satisfies the semigroup condition, and by Proposition is the ideal flow map.
Stopgradients. In practice, it is beneficial, when optimising over neural networks, to use the stopgrad(-) operator to control the flow of information. This allows to interpret the flow model of, as a "teacher", which is used to provide training signal for the map on the off diagonal A t. If the squared residuals presented Proposition are minimized directly, gradients can align the trained flow model with the untrained flow map, rather than vice-versa. To this end, we provide some concrete recommendations inspired by the setting when frozen pre-trained teacher model available.
For LSD, we recommend:
( (s()) + / / E [axxe, (f,
dsdt, (29)
17

---

Published as a conference paper at ICLR 2026
2. For ESD, recommend:
Los = Ls(0)+ L. %. (10.xg,1115 stopgrad(VX,, (4))) dsdt,
(30)
3. For PSD, recommend:
(ū) = Lo)
[ ( 0 J
E (L)) stopgrad((r (1)) (1 7) (X2, (4.)1)) dsdtdy,
(31)
where 71- 7t with [0,
The recommendations in equation 29 and equation 31 are precisely what would arise given pre- trained teacher. The recommendation in equation 30 is similar, but also places stopgrad(.) on the spatial Jacobian of the model, which often improves optimization stability significantly.
Mean Flows and consistency training. Given the choice of stopgrad(-) in equation 30, the re- sulting parameter gradient will be linear in V.s" We may then replace (1,) by the Monte Carlo estimate 0,I, of the ideal flow b (x) E, = 2|. using the tower property of the conditional expectation E [[1L)
E (f for any function f: IR" + Rd Without this choice of stopgrad, this replacement is not possible due to the quadratic term | V a. (T,)) + XI (L,)v, (T,) [2] because of the nonlinearity in |. 12. With this stopgrad, the quadratic term is a constant from the perspective of the gradient and hence this dis- crepancy can be ignored.
E [f [4]],
This trick is used by both consistency training (Song et al., 2023a; Lu & Song, 2025) and Mean Flows (Geng et al., 2025); in fact, these algorithms can recovered by expanding
0xX",(x) = -o + (t s)0,u.
(32)
Vx2,(x)= 1+ Ve,
Plugging equation 32 into the Fulerian loss equation 30 yields
C())=
= (0)
+ L. //- E 1- + s00, stopgrad (1-
{]
dsdl.
Jo
Replacing vg (1a) by the Monte Carlo estimate of the ideal flow as described above yields
[(v0) = (v))
S/ E (L) s)),0,,,, stopgrad (0 (t- 5)- (t=s)Ve []] dsdt,
(33)
and then choosing stopgrad the 0,vg, term well as the spatial gradient yields the Mean Flow/Consistency Training objective,
((00) = (6)))
+ L / E (lo.,,(1.) stopgrad ((t (1) (t- (1-))4L |] dsdt.
(34)
B PROOFS
B. GENERALISED FLOW MAP CHARACTERIZATIONS
Proposition 1. Let X.,: be parametrised as Xx,(x) exp, (it - s)),,(1)). Then X is the unique GFM for equation if and only if it satisfies any of the following conditions:
18

---

Published as a conference paper at ICLR 2026
(i) Generalised Lagrangian Condition: V(s,t) t) E O0, [0,1],,0 xz, E M,
0x,,, = (x(,s,r(z.))
(3)
(ii) Generalised Eulerian Condition:
(s, 1) E [0, 1]"; 2s E
0gXa,((as) + (v.
(4)
(iii) Generalised Semigroup Condition:
(8,,,uu) u) E [0, 1,,,,,,,,,,, E,,
Xu, (X.,,((())) = Xx,r(2s).
(5)
To show that the the propositions hold, let us first demonstrate the legality of the above claims, that the compared vectors indeed belong to the same tangent planes. For the semigroup condition, observe that xx, (x) E M for any 8, f and æ, and therefore evidently coherent.
Lemma 2. For any M. s,t E [0, 1J?. (,)),)())) E (x..
Proof of Lemma 2. Recall that, for any z, 8, x,,x expx - s),,,(())) and, without loss of generality, suppose that < t. (The case >> is perfectly symmetric.) For 1, / = 1 and the proof is trivial, as Ō E M. Let < 1, and let (z) = Xs, (1-ss) (z). which coincides with (x) for (1 s)u +
> 1- ū E [0, 1]. Remarking that ((1 s)u s) E [0, 1] for all 0 S < 1, we have that X defines a geodesic between Xo(r) and (x) = (x). Now, let < 1, as u
= s from where the proof also trivial. Therefore, it follows that (x) logx,.. (z) (X,.1 (x))
u), and considering the previous expression for u finalises the proof from the definition of the logarithmic map.
The above validates the Lagrangian condition. Let also justify the Bulerian condition: by Lemma 2, (xx)) E Tx..(tr)) and let us now prove that ((X..)) is on the same tangent space.
Lemma 3. For any M, (8, [0, Il. d(Xs.c)alua(x)) E Tx..
Proof of Lemma 3. This is true by definition. To see this, let 7 [0, 1] M, + 7(T) be curve,
such that 7(0) = x and 0,y = Eg. We therefore have that X,, o [0, 1] + M, and therefore (X,, 4 (0) E Mx. -T..a Knowing that (X, (x (0)
concludes the proof.
We can now prove the main proposition.
Proof of Proposition 5. Let first prove the Lagrangian condition. It comes naturally
(x,( =
0, (x..,((ts) 0,x; = vxr),
(35)
which true by definition, since (0t)1Ejo.1) is a solution of Equation (2). As for the Eulerian condition, we note too that, for any x E M, (X,, = æ. (So, it is also invertible.) Taking
the derivative through
0,X ox,, = 00a OX + (dX...)x...aa (0. (en (x))
(36)
(37)
where the last line due the chain rule on manifolds. Moreover, by the probability flow ODE, we have that (x) = es (X,, (z)), and, letting = (x), we have that
0,X,, + (dX,,:) (v>( 0.
(38)
Since X,, invertible for any and and defined on M. is bijective, and thus we can freely rename into to conclude that the Eulerian condition holds. Finally, the semigroup condition is trivially true from the definition: (X,,, o (x.) = It Xx.(xx.).
B.2 PROOF POR LEMMA
19

---

Published as a conference paper at ICLR 2026
Lemma 1. (Generalised Tangent Condition) Let ve( = (x 101k and Ie E M be the drift of equation 2. Then it holds that lim,-st 8, X (xx) = ve(xt).
(0.) 1. xz. for any t
Proof. Proof of Lemma 1. This property follows from the Lagrangian condition. We have that (xx) = vx(Xx,(x..) Taking the limit on both sides, we find that
= Ili or (x,,((2))) lim d X. (xx))
(39)
a-t
x-2
lim 0, (x,) [lim Xs, (xx))))) =x())),
(40)
5-
where the last line comes from the continuity of the limit.
Lemma 4, For any (s, E [0, 1]2 and M, lim Ox (2))
Proof. Following elementary Riemannian geometry, we find that
(x)),((z)) 0 exp, ((t s)Us, (x))
(U) (v) + (t ((--))00 (x),,,,
(41)
Knowing that d(exp=)ū = ld and taking the limit, we find the desired conclusion.
B.3 GENERALISED SELF-DISTILLATIN LOSSES
Proposition 2. (Generalised Lagrangian self-distillation) The GFM is the unique minimiser of the objective equation 7 for LGFM-SD (0) = LG-LSD(O), where
(GILD(()) = ( |1...x2..1 (T.) (x 8.t
(T..(TT))
(8)
Proof. The loss is zero if and only if 0,x (T.) 0. conclude by Proposition 1.
vx (. almost everywhere.
We can
Proposition 3. (Generalised Eulerian self-distillation) The GFM is the unique minimiser of the objective in equation 7 for LGFM-SD (0) = LG-ESD (6), where
LG.ESD (0) = (x...() [0,00 d(x) 1x (4,1111)))
(9)
Proof. The loss zero if and only if X (2,((x,) - d(x (x.)) almost everywhere. We conclude by Proposition 1.
Proposition 4. (Generalised Progressive Self-Distillation) The GFM is the unique minimiser over Ug of equation 7 for LGFM-DD (0) = LG.PSD (0) and (s, 1) ~ U(s (), where
LG.PSD (0) (E,,,, dd (x9,((4) x (x (1.)))
(10)
Proof. The loss is zero if and only if (x clude by Proposition 1.
(. (x)) = x x2,(x) almost everywhere. We can con-
B.4
CONNECTIONS TO EXISTING METHODS
B.4.1 GENERALISED MEAN FLOWS
Generalising Mean Flows directly non-trivial as it requires defining the integral of vector field on manifold properly, which would involve parallel transport and therefore derivatives thereof. Also, Mean Flows operate on the vector field level as opposed to Flow Map Matching which operates on the level of the flow map. It is difficult to go from one level another directly, as it will involve non- trivial curvature terms. Instead, we propose to heuristically follow our derivations in Appendix A, in the "stopgradients" section. Indeed,
that loss involves the instantaneous vector field
20

---

Published as a conference paper at ICLR 2026
of the modelled flow map, vl,e- opposed the ideal flow d.Ig hence the use of the Levi-Civita connection along Ug (4,) instead of the differential evaluated v.
CG (0) (x. (T,,, stopgrad (Ö.I. -t-
[.]
(42)
which indeed recovers the Euclidean case a special case. The Levi-Civita connection, V x(M) X(M) - x(M), (u) + V,X, the unique torsion-free, bilinear, metric com- patible connection on gM. that respects the Leibniz rule in X, and it defines notion of covariant derivative for the vector fields on (M.g).
B.4.2 GENERALISED SHORTCUT MODELS
Shortcut models (Frans al., 2025) are exactly trained to enforce the semigroup property but on discrete time grid. Expressed in our notation, the loss amounts to
(())- (0) + E. (x....)) | ((a.)+2a( - stopgrad (Co+d,,42a(T, + du a (43)
Indeed, d can be seen the time difference s, and, letting d be uniformly distributed, recover exactly Euclidean Progressive self-distillation, as noted in Boffi et al. (2024). The Riemannian case strietly analogous, where the operation is replaced by the exponential map, and the appropriate manifold distance is used.
C ADDITIONAL EXPERIMENTAL DETAILS
C.I THE RELEVANCE OF NLL
The NLL is standard generative modelling benchmark for generative models (Chen & Lipman, 2024; Lipman et al., 2023; Song et al., 2020b). It is true, however, that in our context this metric does not evaluate directly the quality of our learnt GFMs. is still relevant least two important ways:
1. It shows that the instantaneous vector field is still well-learnt, despite the deff-disiillttinn loss, which could have impacted the training dynamics on this part of the loss. We show that it does not. So, overall, our objective did not "sacrifice" the instantaneous vector field, and the sum of the self-distillation and flow matching losses are not inherently "incompatible".
2. The flow map and its learnt instantaneous vector field are approximations of one another. Therefore, while the likelihood of one is not exactly equal that of the other, the difference between the two being typically very low (as measured by the self-distillation losses, down 10 in MSE), the instantaneous vector field allows us to compute a good approximation of the likelihood induced by the flow map.
C.2 EMPIRICAL MMD CALCULATION
For any distributions p and 9 with support (M, g), and with respective independent samples (p.)ssisn and (9:)1515n=
MMD g) = Σ exp (-Rd{(PE,P))) + exp (-Rd q;)) 2exp (-nd 4:) (44) {J==1
We choose n to be equal to the size of the test set.
C.3 ADDITIONAL RESULTS
Figure includes additional (all) qualitative results the Earth datasets.
C.4 THE COMPUTATIONAL COST OF SELF-DISTILLATION
We include, here, Table exhibiting the differences in training run-ting run-times (wall clock time) be- tween our GFM objectives and plain RFM; the numbers reported have been gathered on the Volcano
21

---

Published as a conference paper at ICLR 2026
Figure 7: Plots of densities for the various datasets and all compared methods. Depicted in red the test-set samples. Datasets from left to right: volcano, earthquake, flood, fire.
G-PSD
30
-10
230
-30
D0
110
Letitao
G-ESD
-10
300
-30
300
-00
100
100
100
G-MF
Lattale
Loughtkk
RFM
G-ESD
G-LSD
G-PSD
0.60 + 0.07 epochs/min
0.29 0.01 epochs/min
0.36 + 0.02 epochs/min
0.43 + 0.04 epochs/min
Table 5: Wall clock throughput time of proposed methods the volcano dataset (higher is faster).
dataset, which ran the largest model (and therefore is arguably most relevant; about 45M parame- ters). The gradation is rather intuitive: RFM has simple flow matching loss; G-ESD computes spatial Jacobian large mode; G-LSD only computes time derivative of the same model; and G-PSD requires two additional forward passes without any gradients, compared to RFM.
22


Reviews:
[{"review_id": "i42mYf5SWR", "reviewer": "Reviewer_f4VY", "summary": "This paper introduces Generalised Flow Maps (GFM), a new class of generative models for Riemannian manifolds that enables fast, few-step inference. It generalizes Euclidean flow map principles by proposing three self-distillation training methods: Generalised Lagrangian, Eulerian, and Progressive Flow Maps (G-LSD, G-ESD, G-PSD). These methods are derived from three equivalent theoretical conditions that characterize a manifold-constrained flow map. Experiments on various geometric datasets (tori, spheres, SO(3), hyperbolic) show GFMs achieve state-of-the-art sample quality in single- and few-step evaluations, outperforming traditional geometric generative models that require many simulation steps.", "questions": "See **Weaknesses**.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "**Originality & Significance**: The paper's core contribution is the novel and significant generalization of flow map learning from Euclidean spaces to arbitrary Riemannian manifolds. This directly addresses the critical bottleneck of slow inference in existing geometric generative models.\n\n**Quality & Clarity**: The work is technically sound, providing a rigorous theoretical foundation for the new GFM variants. The empirical validation is comprehensive, testing across a diverse set of non-trivial manifolds and demonstrating clear, state-of-the-art performance in few-step sampling. The paper is exceptionally well-written and easy to follow.", "weaknesses": "**Unexplained Performance Gaps**: The three GFM variants are derived from theoretically equivalent conditions but show vast empirical performance differences (e.g., G-LSD is far superior to G-ESD in Table 2). The paper notes this but lacks an in-depth analysis of why, which is a key practical limitation.\n\n**Missing Training Cost Analysis**: The paper focuses exclusively on inference speed. However, the GFM objectives add a self-distillation loss to the standard RFM loss, implying a more expensive training procedure. This training cost trade-off is not discussed or quantified.\n\n**Validity of NLL Metric**: The NLL is computed using the \"instantaneous velocity\" (the implicit vector field $v_{t,t}$). Given that GFMs are trained to learn the global flow map $X_{s,t}$ for few-step sampling, how valid is the NLL of the instantaneous flow as a primary measure of model quality?\n\n**Missing Protein Backbone Experiments**: The introduction prominently features protein backbone generation ($SE(3)^N$) as a key high-impact application. Why were experiments on this crucial manifold omitted in favor of simpler torsion angle datasets?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "CdfuPsJ5IZ", "reviewer": "Reviewer_bSeb", "summary": "The paper proposes Generalised Flow Maps (GFM), a novel framework for few-step generative modeling on Riemannian manifolds. It addresses the high computational cost of geometric diffusion and flow-matching models by directly learning the flow map (integrator), enabling single or few-step high-quality sampling. The approach generalizes existing Euclidean few-step models and achieves state-of-the-art sample quality across various geometric datasets.", "questions": "Why in Tab. 3 volcano dataset the proposed algorithm worse than several existing algorithms? Is it because lack of data since the volcano has only <1000 data.\n\nThe Euclidean version that requires fewer NFE algorithms, in my view, highly depends on the flat property of the space. Will the curved space itself lead to some difficulty in GFM? Do you have some intuition?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "GFM is tested on TNA torsion, Earth dataset, SO3, and Hyperbolic spaces. The exponential results are well proposed.", "weaknesses": "The experiments are demonstrated on relative low-dim spaces, unlike the Euclidean algorithms, e.g., Meanflow, that can be applied to large datasets.\n\nIn the case with relatively large NFE, e.g., NFE=8, there is almost no improvement. (Fig. 4)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "o1uKiubYNS", "reviewer": "Reviewer_qmRH", "summary": "The paper introduces the Generalized Flow Maps (GFM) framework, which extends the flow map matching framework to Riemannian manifolds. It further generalizes the self-distillation loss from previous works to the Riemannian setting and demonstrates that each variant can be efficiently learned by neural networks. Experiments on RNA backbones, protein side-chains, and various toy datasets validate the empirical performance of the proposed methods.", "questions": "1. Does the proposed method require the geodesics of the Riemannian manifold to have a closed-form expression? In addition, does it rely on an embedding of the manifold into a Euclidean space? I would like to confirm whether my understanding is correct.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper is well-written, and the proposed method is easy to follow and understand.\n2. It extends the Flow Map Matching framework [1] to general Riemannian manifolds.\n\n\n[1]. Boffi, Nicholas Matthew, Michael Samuel Albergo, and Eric Vanden-Eijnden. \"Flow map matching with stochastic interpolants: A mathematical framework for consistency models.\" Transactions on Machine Learning Research (2025).", "weaknesses": "1. **Limited experimental scope.** The experiments are restricted to relatively simple toy datasets. As shown in previous works, Riemannian generative models have broad applications in various domains, such as material generation [1], molecular conformer generation [2], and protein backbone generation [3]. As mentioned in the paper, the inference processes of these models are often inefficient. Therefore, additional experiments on real-world applications would strengthen the paper’s contributions and demonstrate that the proposed methods can indeed improve the inference efficiency of Riemannian diffusion models.\n\n2. **Applicability to general manifolds.** The proposed method appears to be applicable only to well-known manifolds with closed-form geodesics. In addition, it relies on an embedding of the manifold into Euclidean space. Consequently, for more complex manifolds without such prior geometric knowledge, the method would be difficult to apply in practice.\n\n3. The results appear less innovative, as the proposed method can be seen as a direct generalization of existing approaches in Euclidean space [4, 5]. However, this would be acceptable if the method proves to offer meaningful benefits in real-world applications (see Weakness 1).\n\n\n[1]. Miller, Benjamin Kurt, et al. \"Flowmm: Generating materials with riemannian flow matching.\" arXiv preprint arXiv:2406.04713 (2024).\n\n[2]. Jing, Bowen, et al. \"Torsional diffusion for molecular conformer generation.\" Advances in neural information processing systems 35 (2022): 24240-24253.\n\n[3]. Yim, Jason, et al. \"SE (3) diffusion model with application to protein backbone generation.\" arXiv preprint arXiv:2302.02277 (2023).\n\n[4]. Boffi, Nicholas Matthew, Michael Samuel Albergo, and Eric Vanden-Eijnden. \"Flow map matching with stochastic interpolants: A mathematical framework for consistency models.\" Transactions on Machine Learning Research (2025).\n\n[5]. Geng, Zhengyang, et al. \"Mean flows for one-step generative modeling.\" arXiv preprint arXiv:2505.13447 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "gay4TKZkEC", "reviewer": "Reviewer_73fW", "summary": "This paper introduces Generalised Flow Maps (GFM), a unified framework for few-step generative modeling on Riemannian manifolds. It generalizes the notion of flow maps to non-Euclidean geometries such as spheres, tori, and Lie groups.\nThe core idea is to replace the Euclidean linear interpolant with geodesic-based Riemannian interpolation and to define manifold-consistent flow maps that jump between points along the probability flow ODE.\nThe authors derive three generalized self-distillation objectives, Generalised Lagrangian, Eulerian, and Progressive Flow Maps.\nEmpirically, GFM achieves state-of-the-art few-step generation on diverse geometric domains, such as protein torsion angles, geospatial data, and synthetic rotations.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper extends flow-map-based generative modeling from Euclidean to general Riemannian manifolds, encompassing multiple geometric structures (tori, spheres, Lie groups, hyperbolic spaces).\n\n* The paper is theoretically rigor.\n\n* The paper tries to bridges recent few-step generation paradigms (Consistency Models, Mean Flows, Shortcut Models) with the geometric setting via G-LSD, G-ESD, and G-PSD objectives, which is a appealing story.", "weaknesses": "A potential weakness of this work lies in its limited empirical validation on large-scale or high-dimensional settings, where the proposed few-step framework would be most impactful. The observed gains are primarily evident at very few inference steps (as shown in Figures 3 and 4) and on relatively small datasets, making it unclear how well the method scales to more complex manifolds or realistic applications.\nUnlike Euclidean diffusion or flow matching models, where few-step sampling is crucial for accelerating large-scale tasks such as protein generation or real-time image editing, the computational benefits in Riemannian domains remain less compelling.\nIn particular, for manifolds like torsion or angle spaces, existing diffusion or flow models already achieve fast inference, so the advantage of adopting a few-step Riemannian formulation is not yet convincingly demonstrated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)