Riemannian Consistency Model
Authors: Chaoran Cheng
Yusong Wang
Yuxin Chen
Xiangxin Zhou
Nanning Zheng
Ge Liu
Venue: NeurIPS 2025 poster ()
Keywords: Consistency Model
Flow Matching
Generative Model
Riemannian Manifold
Primary area: deep_learning
TL;DR: We propose Riemannian Consistency Model (RCM) as an extension of the original consistency model for non-Euclidean domains.

Abstract: Consistency models are a class of generative models that enable few-step generation for diffusion and flow matching models. While consistency models have achieved promising results on Euclidean domains like images, their applications to Riemannian manifolds remain challenging due to the curved geometry. In this work, we propose the Riemannian Consistency Model (RCM), which, for the first time, enables few-step consistency modeling while respecting the intrinsic manifold constraint imposed by the Riemannian geometry. Leveraging the covariant derivative and exponential-map-based parameterization, we derive the closed-form solutions for both discrete- and continuous-time training objectives for RCM. We then demonstrate theoretical equivalence between the two variants of RCM: Riemannian consistency distillation (RCD) that relies on a teacher model to approximate the marginal vector field, and Riemannian consistency training (RCT) that utilizes the conditional vector field for training. We further propose a simplified training objective that eliminates the need for the complicated differential calculation. Finally, we provide a unique kinematics perspective for interpreting the RCM objective, offering new theoretical angles. Through extensive experiments, we manifest the superior generative quality of RCM in few-step generation on various non-Euclidean manifolds, including flat-tori, spheres, and the 3D rotation group SO(3), spanning a variety of crucial real-world applications such as RNA and protein generation.

Riemannian Consistency Model
Chaoran Cheng", 1 Yusong Wang"?, Yuxin Chen', Xiangxin Zhou², Nanning Zheng', Ge Liu
Uniiersitty of Illinois Urbana- Chraana-hammainn, 2Xi'an Jiaotong University, UUniversity of Chinese Academy of Sciences
Abstract
Consistency models are a class of generative models that enable few-step gener- ation for diffusion and flow matching models. While consistency models have achieved promising results on Euclidean domains like images, their applications to Riemannian manifolds remain challenging due to the curved geometry. In this work, we propose the Riemannian Consistency Model (RCM), which, for the first time, enables few-step consistency modeling while respecting the intrin- sic manifold constraint imposed by the Riemannian geometry. Leveraging the covariant derivative and exponential-map--asse closed-form solutions for both discrete- and continuous-time training objectives for RCM. We then demonstrate theoretical equivalence between the variants of RCM: Riemannian consistency distillation (RCD) that relies on a teacher model to approximate the marginal vector field, and Riemannian consistency training (RCT) that utilizes the conditional vector field for training. We further propose a sim- plified training objective that eliminates the need for the complicated differential calculation. Finally, we provide a unique kinematics perspective for interpreting the RCM objective, offering theoretical angles. Through extensive experiments, we manifest the superior generative quality of RCM in few-step generation on various -uclldcidnnn manifolds, including flat-tori, spheres, and the 3D rotation group SO(3), spanning a variety of crucial real-world applications such as RNA and protein generation.
parameterization, we derive the
Introduction
Diffusion [47, 48, 19] and flow matching [27, 7] models have achieved remarkable success generative modeling in various domains including image generation [42, 11], protein design [55, 3]. and text generation [2, 14, 9]. As an intrinsically iterative process that gradually transforms the data from random noises into meaningful samples, the inference procedure of diffusion and flow matching usually requires hundreds up to thousand steps for decent generation. To mitigate such high computational cost, a new family of generative models known as the Consistency Model (CM) [50] was proposed. By "shortcutting" the probability flow and enforcing consistent model outputs, consistency models are able to generate high-quality samples using one or a few steps. In the image generation task, consistency models have surpassed the existing distillation approaches [44] and rectified flows [29, 28] in one-step generation.
Besides Euclidean domains like images, generative models various Riemannian manifolds have a potentially profound impact in scientific domains, including protein generation [55, 3]. peptide design [26, 25], robotics [5], and geoscience [41]. For example, the generative modeling of protein requires deseriptors its position, orientation, and torsion angles for each amino acid. While the position is Euclidean, the orientation lies in the 3D rotation group SO(3) and the torsion angle lies in
'Equal contribution.
"Coresponding author to chaoran7@illinoia edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

the flat torus where a translation by 2n will result in the same angle. Existing works for protein design typical require 200 to 1000 sampling steps, limiting the overall throughput for virtual screening. fast and effective few-step generative model on Riemannian manifolds, therefore, can significantly accelerate the drug discovery and enzyme design processes, further facilitating crucial real-world pharmaceutical applications.
While the current Euclidean consistency model has achieved superior performance on images where the distance between two predictions is measured with the standard Euclidean norm, its counterpart in Riemannian manifolds, the Riemannian consistenyy model (RCM), poses additional challenges and remains largely unexplored. Specifically, the intrinsically curved manifold requires the consistency parameterizations to lie on the manifold. In this way, simple linear interpolation like Lu and Song [33], Yang et al. [54] is not always feasible. Furthermore, such a manifold constraint will impose an additional constraint on the vector field when the consistency loss enforced, that is, the vector field at different points should also lie in their corresponding tangent spaces, necessitating additional corrections in the time derivative.
In this work, we address the challenge of Riemannian Consistency Modeling with a novel consistency parameterization based on the exponential map to ensure the manifold constraint and use the covariant derivative to account for the intrinsic curved geometry when calculating the time derivative. We provide the flosed-form solutions for both discrete- and continuous- time formulations of the RCM objective. Similar to the Euclidean CM, we theoretically prove that the Riemannian consistency distillation (RCD), which relies on a teacher model to approximate the marginal vector field, can be extended to Riemannian consistency training (RCT), which directly utilizes the conditional vector field with marginalization techniques for training. We further propose a simplified training objective that empirically improves model performance and eliminates the need for potentially complicated calculations of the differentials of the exponential map. Intriguingly, we provide an intuitive interpretation of the RCM objective from the perspective of kinematics curved geometries, offering new theoretical angles. We carry out extensive experiments on non-Euclidean manifolds, including flat tori, spheres, and the 3D rotation group SO(3). Compared to the vanilla flow-matching and the naive Euclidean adaptation of the consistency model, our RCM recipe demonstrates higher- quality generations with better distributional fitness in the step generation setting.
2 Preliminary
2.1 Riemannian Flow Matching
Conditional flow matching (CFM) [27] learns time-dependent vector field that pushes the prior noise distribution to any target data distribution. Such flow-based model can be viewed as the continuous generalization of the score matching (diffusion) model [47, 48, 19] while allowing for a more flexible design of the denoising process. Riemannian flow matching (RFM) [7] further extends CFM to general manifolds on which a well-defined distance metric can be computed.
Mathematically, consider a smooth Riemannian manifold .M with the Riemannian metric a probability path Pt : [0, 1] + P(M) is a curve of probability densities over M. A flow thz : [0, 1] x
M + M is a time-dependent diffeomorphism defined by a time-dependent vector field u. [0, 1] X M + TM via the probability flow ordinary differential equation (PF-ODE): ă (x() = u. (x).
The flow matching objective directly regresses the conditional vector field u, (Ia/200.2 21) a with a time-dependent neural net vo c0 t) where H, := (x))... The Riemannian flow matching objective can be formulated as:
PREM- En po0.) [1l 0) ue x1) 11]
where q is the data distribution, Po is the prior distribution, and 2, == (/ Ir, 21) denotes the conditional flow. Chen and Lipman [7] further demonstrated that if the exponential map and logarithm map can be evaluated in closed-form, the condition flow can be defined as the geodesic interpolation J, = expa, (Ns loga, 2o), where K; is monotonically decreasing schedule satisfying Ki = 1, 0. In this way, the corresponding vector field can be calculated as u, (x, Io, 21) = K₁ log,, 21 /k,k, In
this work, we use ă, u, etc, to denote the time derivative of xr, v and follow the Einstein summation notation. If necessary, the time will be noted in the additional subscript.

---

2.2 Consistency Model
The marginal vector field learned by the (Euclidean) flow matching model is not necessarily straight, which requires solving the PF-ODE with hundreds of iterative steps. The consistency model (CM) [50], by design, can generate high-quality samples by directly mapping noise to data. CM achieves one-few generation by "short-cutting" the PF-ODE such that the denoiser fo t) along the PF-ODE should output consistency predictions:
CCM = Er, [wellfo (2 1) fo- (II+ t + At)
(2)
where At 1/N is the discretization step, f is the nop-gradient operation, and we is a weighing function. Ie+a: is defined by following the marginal vector field at 2; along the PF-ODE. To avoid trivial solutions, an additional consistency constraint fo : 1) Ij needs be enforced. Song et al. [50] utilized parameterization of fo(x,, 1) Gn(()); + (L) (2,, where Fo(x,, () is the unconstrained denoiser and Cin (4), Cour are schedulers such that Gin (in(1 = 1, Coun (1) = Yang et al. [54], Lu and Song [33] further related the denoiser parameterization with the vector field (score) parameterization. In practice, a pre-trained flow matching model is used to approximate the marginal vector field (consistency distillation, CD). Additionally, Song al. [50] also demonstrated that, with the coop-gradien operation, training on the conditional vector field (consistency training, CT) has the same marginalization effect. Existing CMs have achieved high-quality one-step few-step generations on Euclidean domains like image generation. Their Riemannian counterpart, however, remains unexplored.
3 Riemannian Consistency Model
3.1 Consistency Model on Riemannian Manifolds
As discussed above, the extension of the consistency model to Riemannian manifolds poses additional challenges, as the model needs to be fully aware of the intrinsic geometry to shortcut the Riemannian PF- ODE. To address the above challenges, we choose to learn the vector field vy (3t, () and adopt the consis- tency parameterization utilizing the exponential map to ensure the manifold geometry (see Figure l):
T&M
x
v
T
U
M
Jo (Zx,, 1) == exp2, (IU !)-
(3)
(t)
fo(æ,, t) = expx, (x t)
It is easy to verify that for any bounded vector field uo (1,,, t). the above parameterization satisfies the con- sistency constraint (x 1) I1 K1 O.
Figure 1: Riemannian Consistency Model (RCM). The denoiser fe along the Rieman- nian PF-ODE should be consistent.
The geometric property of Riemannian manifolds allows us to define geodesics as the "straight lines" on the manifold, whose lengths are geodesic distance, the shortest distance between two points on the manifold. Therefore, following the core idea that the predictions along the same PF-ODE should be consistent, we use the geodesic distance to measure the consistency:
C E2 [wrd (fo t), fo (2440 + At))
Similar to Euclidean cases, the next data point Tr+ is defined along the marginal PF-ODE on the manifold. When N + co, we have the following continuous-time limit of RCM:
Theorem 3.1. When N + 00, At + O. the continuous-time RCM loss
RRC = lim RCM Be,or [w (Id(exp))a (RU + RVavv) + d(expu), (+)) N
(5)
where d((xpx) u) T,M + T are the differentials of the exponential map with respect the tangent vector u = ww and the base point x. V denotes the covariant derivative with respect to the Levi- Civita connection of the Riemannian manifold (M. g) along the PF-ODE.
3

---

Proof. Note that (z, y) = log,r vll holds for any z, E M. As At - 0, the two target points in Eq.4 also approach each other. Therefore, the first-order approximation log, y = y - x holds. Now consider (fo (Fe+Artt + At) fo (z, t)) + f. Applying the chain rule, we get:
f d u (Rv + xVav) + d(exp u), (#).
(6)
Marginalization over L and 2t, concludes the proof.
In the above formulation, we assume that the marginal vector field å is available via the pre-trained RFM model, leading to the Riemannian consistency distillation (RCD) where the marginal is learned by pre-trained Riemannian flow matching model. Similar to [50], in the following theorem, we demonstrate that the same objective can utilized with the conditional vector field 2|z1 leading to Riemannian consistency training (RCT).
Theorem 3.2. With the stop-g pra-gradient operator 0- in Eq.4, the marginal vector field i in the RCD loss can be substituted with the conditional vector field 2|21
The following lemma is the key result for extending RCD RCT:
Lemma 3.1. f is linear in i.
Proof. By the definition covariant derivative, V linear with respect to ø. Also note that d(exp,)u )u ((exp u)z both linear mappings. Therefore, the final result is linear in ø.
Li
Lemma 3.2. 3 (d(()a |2s,), or using the Riemannian integral J.M u(x) ( dvol
Lemma 3.2 generalizes the Euclidean case and can be verified using the Bayes' rule Riemannian manifolds [6]. It draws the connection between the conditional and marginal vector fields. The key to the proof of Theorem 3.2 is that want move the expectation outside use the conditional vector fields.
Proof for Theorem 3.2. Let denote the denoising result in Eq.3 at time step +++ At and A = f-f. Note that dg (S. fa) |og follz Taking the gradient on both sides, we have
ood (fo, (fo. fo-) 2
2 v Vollog,, fo-,log,, fo-) )g-- = (log for Vo log fo-)g
(7)
(Afo. volfe- fo))g fo. Vofolg
Using Lemma 3.1, Afe Ar f is linear in à and the second argument Vefo is now independent of a. This indicates Vod? (S. fo-) also linear in å. Therefore, when using Lemma 3.2 to marginalize over it, we can simply move the expectation over outside the gradient operation, leaving the conditional vector field á , inside the expectation.
The proof above inspires us to use an alternative loss as
CRCM = Er,,, w (fo- fo + fo-, >,}
d(epp)) (ke + <V + d(exp u)z (+)
(8)
where the differentials are calculated along conditional vector fields. It is easy to verify that the loss Eq.8 has the same value Eq.5 and the same gradient as demonstrated in the proof above.
3.2 Simplified Riemannian Consistency Model
The loss in Eq.8 involves the calculation of the differentials of the exponential map d(expa)a, d(exp u).. which may require additional complex symbolic calculation for efficient implementation (see Appendix B). Instead, we proposed an alternative loss on the vector fields that eliminates the need compute these differentials:
C30 RRC = -E,, l ko + K V&vll +V&ll]"
Er,zi [w (ug- uo + ug- u
(9)

---

where dgg := 2 Kug + AVV Here, we make one key approximation that d(exp=)u = d(exp u)x such that the second term can be combined with the first one. Also note that d(exp,), linear mapping that does not change the optimality of the loss. Therefore, we simply ignore such transform and optimize the norm of dgg at st. Indeed, for the flat-torus and any manifold whose exponential map is symmetric, the identity d(expg) d(exp u), always holds. For general Riemannian manifolds, the following result holds: Proposition 3.1. For u, U E M, if u,, are parallel, then d(exp=)u(v) = d(exp u) (v).
Proof. As both differentials linear, it suffices verify ((ex))))(()) d(exp u) (u). Consider the consistency parameterization f(r, f) == expx, (1 E)u,, where 14 = Izz.. (u) is the parallel- transported tangent vector along the geodesic y from zy to Z,. Therefore, f(z,, t) constant dynamics along the geodesics defined by point in constant speed, we have f(z,, t) exp2 u, V E [0, 1]. Taking the derivative with respect to t, we obtain:
and vector field th = u E M.M. As travels
f= d(epp n (- u + - -)V + d(exp u)., (4) (6)=0) =0
(10)
As = u, the covariant derivative V = V vanishes as traces a geodesic. Therefore, we have d(exp,))u()) d(exp (t) d(exp (u), which concludes the proof.
In the original RCM loss, the two tangent vectors follow the predicted and the marginal vector field directions, respectively. This indicates that, if the pre-trained model approximates the marginal vector field well, the approximation shall be more accurate. The adaptation from RCD to RCT follows a similar procedure of marginalization in the proof for Theorem 3.2. We summarize the simplified RCD and RCT training procedure in Algorithm and 2, where the key differences are highlighted in red. The original RCD and RCT follow similar training algorithms except for optimizing the loss in Eq.8. Sampling from RCM follows the same setup as the Euclidean CM, which is described in Appendix C.2.
Algorithm 1 Simplified Riemannian Consistency Distillation (sRCD)
Algorithm 2 Simplified Riemannian Consistency Training (sRCT)
1: Input: Pre-trained RFM 8.-
Input: None.
2: while not converged do
2: while not converged do
3: Sample data 21, noise wi, and I.
3:
Sample data #1. noise 2o, and
4: Calculate x; expx, (Kr log,, (zo)).
4:
Calculate x; = expa, (Ke log,,, (zo)). Calculate ã; = k₃ log., 21 Kip.
5:
Calculate (x(1,,
5:
6:
Calculate úig 5 + Amg + KV 28.
Calculate üg + KoD + K Vo.
7:
Optimize the loss with Vow(pg- Ug + ugg- -gg)) )--
7:
Optimize the loss with Vow(ng- uo + thg- úg- )9-
8: end while
8: end while
3.3 Kinematics Perspective of Riemannian Consistency Model
T.M
T.M
T.M
T&M
T.
M
M
M
7(t)
((t)
((t)
(a) Vector field difference è (b) Vector field derivative ò
(c) Covariant derivative V
Figure 2: Three components of the variations in the consistency objective.
We now provide a perspective from kinematics that intuitively explains the terms the RCM objective. Consider a point moving on the Riemannian manifold, the RCM objective essentially enforces the infinitesimal equilibrium on the target [o(,,, f) along the PF-ODE of z. Locally, the change of the final target can be decomposed into three components:
5

---

The difference in the predicted and marginal vector fields. Figure 2(a).
The intrinsic change of the vector field, Figure 2(b).
. The extrinsic change of the vector field due to the geometric constraint, Figure 2(c).
The first and second terms more intuitive, represented by the difference and derivative of the vector field. Also appearing in the Euclidean CM in [54]'s parameterization, these two terms can be directly generalized to any Riemannian manifold with the standard calculation in the tangent space M2 a vector space where all vector calculations are valid with the additional Riemannian metric.
The third term, however, is unique for general Riemannian manifolds, as it involves the geometric properties of the manifold. Intuitively, the curved geometry leads to different tangent spaces at different points, where tangent vectors not directly comparable. In order to differentiate the vector field at adjacent points, the covariant derivative is introduced as a generalization of the directional derivative in the Fuclidean case. More concretely, the covariant derivative Vae describes how the vector field v changes along the curve defined by j. In other words, even if the vector field u is "constant" along the curve, its time derivative is not necessarily zero; is only when we consider such additional geometric properties that we can arrive at the results that the covariant derivative is indeed zero for a constant vector field. In this case, the vector field is called garallel-transporrs along the curve i.
In this way, we have included all three components that affect the target position in the kinematics of the Riemannian manifold, with the last extrinsic change capturing the manifold's s geometric properties. We now discuss some concrete examples of Riemannian manifolds demonstrate how different RCMs can explicitly or implicitly learn the consistent kinematics on the manifold. Specifically, we focus the covariant derivative Vi. which appears in all continuous-time RCM losses. The explicit mathematical formulae can be found in Appendix B.
Euclidean Space. The Euclidean space is flat manifold with the canonical Euclidean metric Pi = ō.g. The covariant derivative reduces to the common time derivative as V&U = è. This intuitively makes sense as the vector field can be trivially transported everywhere on the flat manifold. With the linear scheduler Ae = - L, we obtain the original vector field consistency model loss + (1 t)ui| in [54] as expected. A similar analysis holds for the flat torus T (8*)), where S is the ID spherical manifold (see below).
||
Spherical Manifold. The n-sphere S [||=||| = = E R'+!} is n-dimensional manifold which inherits the Euclidean metric of RR++1 The covariant derivative on sphere reads V + #).: = 0. If v á such that the first variation is perfectly optimized, the acceleration becomes à = - su, v) with a magnitude of (v, u) = ||>|| and points in the inverse direction of x towards the origin. This result coincides with the acceleration formula for uniform circular motion, which reads ||a|| = ||u|l//r, with the direction also pointing to the center. In this way, with the non-zero curvature, the covariant derivative provides additional geometry-awart information for the consistency objective.
3D Rotation Group. The 3D rotation group SO(3) of all 3D rotation matrices is a Lie group with Riemannian structure. With the group property, the tangent space every point is isomorphic to the tangent space = T.G at the identity element e = T, also known as the Lie algebra. The Lie algebra of SO(3) (3). the vector space of skew-symmetric matrices. Using the 3-vector representations, SO(3) has natural bi-invariant Riemannian metric g(u, v) = (u, v). The covariant derivative can be calculated as VAV- = i + _ e). where [u. v] = u X v is the cross-product for 3-vectors. This formulation has a close relation to the rotating frame of reference in kinematics, where the additional cross product represents the Coriolis term induced by the non-inertial rotating reference frame. Furthermore, as [p, 2) vanishes, we have V.u d. Essentially, this condition indicates that u should be a invarinvtr vector field to trace a geodesic on SO(3).
4 Experiments
To demonstrate the effectiveness of the RCM framework on Riemannian manifolds, we carry out extensive experiments on various Bunluulid settings. In addition to Riemannian consistency distillation and training (RCD, RCT), we also test with the simplified objective in Eq.9 (sRCD, sRCT). Furthermore, to demonstrate the advantage of continuous-time RCM, we include the discrete- time RCD (dRCD) with discretization step At 10-2 In principle, as all Riemannian manifolds

---

used in this work have natural inclusion mapping to the ambient Euclidean space, the Euclidean CM can be directly applied. We call such naive adaptation (on distillation) CD Ca naive to distinguish from our RCD. We use 2-step generation setup for all the models above during inference and set the fixed intermediate time step to = 0.8 for all manifolds and datasets. As reference, the results for the 100-step (RFM-100) and 2-step (RFM-2) Riemannian flow matching model are also provided.
Across all experiments on all manifolds, we employ the same network architecture, with the sole difference being the input and output dimensions that vary according to the manifold dimension. We use the linear scheduler K2 = 1 - with the weighing function w/ ( = //(1-t) t)² following [33]. Given the critical importance of stable Jacobian-vector products (JVP) é = 0,c + 20," for training consistency models, we followed EDM2 [23] and utilized magnitude-presene connected (MP-FC) layers with force weight normalization. Additionally, we incorporated time information into the model using pagnitude-preservvee detailed network architecture and hyperparameters are provided in Appendix C.
fully-
Fourier features and concatenation. The
To provide model-agnostic evaluation metric for the generation quality, we use nonparametric clustering approaches of kernel density estimation (KDE) and maximum mean discrepancy (MMD) [18]. Specifically, for the 2-sphere datasets, we use KDE with the von Mises-Fisher kernel with a bandwidth of 0.02 and the haversine distance (geodesic distance on sphere) estimate the spherical densities for the ground truth data and the generated samples of the same size. We then mesh-grid the sphere and calculate the Kullback-Leibie flat-tori and S3)3), the analogue of the standard Gaussian kernel is not well-defined, so we switch to MMD to measure distributional fitness. MMD is kernel-based, distribution- free two-sample test, with lower value indicating better distributional fitness. On the two sample sets, the MMD can be calculated as:
divergence (KLD) between the two kernel densities. For
MMD (X,Y) n(n
uz) 2k. k(s, y) exp(- (2. 4)),
{J=1
(11)
where the symmetric Gaussian-lin kernel function l: is defined using the geodesic distance dg with the bandwidth parameter y
4.1 Spherical Manifold
Table 1: KL. divergence between the estimated kernel densities on the 2-sphere datasets. The dataset size is noted. Except for RFM-100, the best results are bold and the second best are underlined.
Flow Matching
Consistency Distillation
Consistency Training
KLD1
RFM-100
RFM-2
sRCD
RCD
dRCD
CD
sRCT
RCT
Farthquake 6,224
1.51
10.99
2.13
2.22
6,20
3.66
2.38
2.38
Volcano Fire
1.77
35.40
3.36
3.84
17.19
5.44
4.47
4.78
0.53
9.79
1.65
1.71
8.01
3.39
1.74
1.72
Flood 12,310
1.33
8.17
2.27
2.41
6.21
2.81
2.39
2.23
For our experiments on spherical manifolds, we utilize the real-world data comprising four distinct earth location datasets: volcanic eruptions [40], earthquakes [39], floods [4], and wildfires [46], collected by [36]. Following [7], we assume the Earth's surface to a perfect sphere. For the
calculation of KLD between estimated kernel densities, we sampled the same number of points as the
dataset for each model.
As shown in Table and Figure 3 on the Flood data, the 100-step RFM achieves the lowest KLD as expected. However, its performance degrades significantly with 2-step sampling. In contrast, all consistency methods markedly improved sample quality in the 2-step setting. Directly using a naive Euclidean consistency model loss yielded the poorest results, highlighting the necessity of RCM's covariant derivative formula to accurately characterize the manifold's geometry. Similar to the results in [50], the results from consistency distillation were superior to those from consistency training. Furthermore, the simplified loss we proposed streamlines the calculation of the differentials of the exponential map without sacrificing generation quality. As a result, the simplified loss yields even better results empirically.
7

---

Data
RFM-100
RFM-2
sRCD
RCD
dRCD
Ci
SRCT
RCT
Figure 3: Generations on the Flood dataset on the 2-sphere. Except for RFM- 100, all models use 2-step generation. The FM, CD, and CT models are colored in blue, green, and orange.
4.2 Flat Torus
We evaluate our RCM framework on flat tori using synthetic checkerboard dataset! well as pre-processed protein [32] and RNA [38] datasets, whose torsion angles can be represented on the 2D and 7D tori, respectively. Specifically, for the checkerboard data, randomly generated 100k sampling points for training. The protein dataset contains 166,305 samples, and the RNA dataset contains 9,473 samples.
The MMD results are pre- sented in Table 2, with additional visualizations in Appendix D.3. RFM with 100 sampling steps still achieves the best performance, while its performance is poor with only two sampling steps. Since the torus is flat, cl(exp. d(exp u)z are identity mappings such that the simplified loss is precisely equivalent to the original loss. Surprisingly. the discrete version of RCD achieves decent results. We argue that this might be because the flat torus is more similar Fuclidean space and less prone to numerical issues compared to other manifolds. Finally, the trivial Fuclidean consistency model CDnaive ignores the periodic nature of the torus, leading boundary issues and significantly underperforming RCMs.
Table Maximum mean discrepancy (MMD) scores on the 2D flat torus. Except for RFM-100, the best results in bold and the second best are underlined.
Flow Matching
Consistency Training
Consistency Distillation
MMDJ/ 10-2
RFM-100
RFM-2
RCD
dRCD
CD
RCT
Board
0.47
14.04
0.61
0.37
3.49
1.89
Protein
0.35
15.19
3.84
3.79
8.40
5.54
RNA
0.97
14.03
7.82
3.39
10.44
5.60
4.3
3D Rotation Group
For the generative modeling on the 3D Rotation Group SO(3). we follow [30] to use synthetic rotation datasets with three different modes (Cone, Fisher, Line) on SO(3). We additionally include more challenging dataset that projects the 2D Swiss roll onto the SO(3) manifold. For each dataset, 100k samples are generated for training the Riemannian flow matching and consistency models. During evaluation, 10k rotations are sampled for each model for MMD calculation.
We refer Flow Matching Guide and Code implement synthetic dataset.

---

Table Maximum mean discrepancy (MMD) scores on the SO(3) datasets. Except for RFM-100, the best results are in bold and the second best are underlined.
Flow Matching
Consistency Distillation
Consistency Training
MMD1/10
RFM-100
RFM-2
sRCD
RCD
dRCD
CDuave
sRCT
RCT
Swiss Roll
1.35
19.64
1.51
1.47
8.69
2.75
4.17
8.23
Cone
7.38
19.96
5.47
6.30
20.39
21.46
7.53
3.78
Fisher
4.02
17.41
5.81
5.71
16.41
6.87
8.59
7.00
Line
4.87
15.50
3.06
2.39
14.93
9.36
3.75
3.32
The quantitative results of MMD on SO(3) are summarized in Table Again, the simplified version of RCD and RCT achieves similar performance to the exact one. Noticeably, the discrete-time RCD performs worse on SO(3). probably due to larger discretization errors on the curve geometry. The naive Euclidean CD also falls short in terms of MMD, even though the 3-vector representation allows for arbitrary rotation vectors as the output. Interestingly, the RCT model on the Cone dataset achieved the best MMD score, even better than the RCD counterpart. This is probably because the pre- pre-trained RFM this case is not as good as the other datasets, as demonstrated by a relatively high MMD for the RFM- 100 model. In this way, the learned marginal vector field may not be accurate enough, leading to worse performance on the distillation approaches. The RCT model, on the other hand, does not suffer from this limitation, relies the conditional vector fields instead.
4.4 Ablation Study
Scalability to sigher-dimenssioal Riemannian manifolds and sampling efficiency are two crucial aspects for the generalization of RCM and other baseline models. To demonstrate scalability, we follow the setup in [27] on high-dimensional tori. The target distribution is a wrapped standard Gaussian distribution centered at the original, with the torus defined on [0, ) Therefore, the high-density regions cover every corner of the hypercube, making it necessary to enforce the manifold constraint (periodic boundary condition). We estimate the Gaussian parameters based on maximum likelihood estimation, and calculate the Fréchet distance to the ground truth. 5k points are sampled as the training dataset, and the same number of samples generated for parameter estimation. The results are shown in Table 4, where RCT consistently outperforms RFM and the naive Euclidean CM in the few-step setup for manifold dimension up to 128. Specifically, we noted that the performance manifolds. We hypothesize is because the
of CT drastically degrades increasing manifold dimension makes it exponentially harder for the Euclidean model to comply with the manifold constraint, the number of constraints grows exponentially. Our results demonstrate the scalability of RCM, further necessitating the need to respect manifold properties in consistency models.
higher-dimensiee
Table 4: Fréchet distance (lower is better) between the ground truth Gaussian parameters and those estimated from the generations on high-dimensional tori. All samples are generated using steps.
Torus dim
2
4
8
16
32
64
128
RFM
0.52
0.70
1.01
1.41
1.95
1.47
1.83
RCT
0.22
0.31
0.54
0.81
0.46
0.58
0.62
CT naive
0.73
2.69
1.58
2.16
2.41
24.80
35.16
We also provide a comprehensive study on the impact of the number of sampling steps on the generation quality for different models in Appendix D.I, where RCM variants consistently outperform the naive Euclidean baseline across all NFEs. The empirical benchmark on the sampling stage speedup provided in Appendix D.2.
Related Work
Riemannian Generative Models. Modeling data distributions beyond Fuclidean space is essen-
tial for applications like protein modeling [56] and geosciences [22]. [10, 20] successfully extend

---

diffusion models [19, 49] to Riemannian manifolds that learn to reverse the noising process. [31] in- troduced practical improvements for Riemannian diffusion models for certain Riemannian symmetric spaces to enhance scalability. [21] tackled scalability by employing bridge processes for generative diffusion on Riemannian manifolds. Another approach involves extending continuous normalizing flows [8, 17, 27] to manifolds. Early attempts leverage mapping between manifolds and Euclidean spaces [15] or focused on simple manifolds using maximum likelihood training [36]. [43] proposed simulation-free training for continuous flows, though they scaled mainly to high dimensions. [6] expanded flow matching [27, 28] and achieved simulation-free training for simple geometries. [51] introduced one-step form flow generator for manifolds. Despite these advances, the develop- ment of a one-step Riemannian Generator remains largely unexplored, motivating our research in this area.
One-Step Generator. Diffusion models, also known as score-based generative models [19, 47, 49], along with flow models [27], have seen substantial success in various fields. However, these models typically necessitate hundreds of functional evaluations to ensure high- ligh-quality outputs. To address this, recent research has focused on reducing the number of function evaluations (NFEs) by adopting advanced solvers [34, 59], introducing better prior-data couplings [29, 52], and leveraging progressive distillation [37, 44]. There is growing interest in developing one-step generators capable of producing samples with single function evaluation. Notably, some studies have successfully
distilled multi-step diffusion models into single-step student generators [35, 57, 53, 45], often drawing from pre-trained teacher models. Another emerging approach involves consistency models [50, 54, 33], which learn consistency function that maps noisy data along ODE trajectory to the corresponding clean data. These models can be trained either by distillation or from scratch. The idea of "consistency" inspires subsequent works such shortcut models [12].
Conclusion
In this work, we propose the Riemannian consistency model (RCM) as an extension to the existing consistency model on Euclidean domains. We have provided an explicit closed-form continuous- time RCM objective with manifold-spere speciffcc operations. Utilizing the marginalization trick on the Riemannian manifold, we rigorously prove that Riemannian consistency training (RCT) is mathematically equivalent to Riemannian consistency distillation (RCD). We further propose simpler objective without the need for the explicit form of the differentials of the exponential map while still achieving similar and even better empirical performance. We provide novel perspective from kinematics that offers an intuitive interpretation of the RCM objective as geometry- motions.
Our RCM, together with its CD and CT variants, achieved significantly better generation quality than the vanilla flow matching and Euclidean consistency model in the few-step generation setup, demonstrating the necessity of incorporating intrinsic geometric information in the model design. We will continue to explore the possibility of scaling up Riemannian flow matching and hope to inspire more efficient flow model architectures in various downstream domains like protein design.
Acknowledgement
This work funded by the DOE Center for Advanced Bioenergy and Bioproducts Innovation (U.S. Department of Energy, Office of Science, Biological and Environmental Research Program under Award Number DE-SC0018420).
References
[1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08779, 2023.
[2] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances neurul information processing systems, 1498 17779991199991.1 2021.
10

---

[3] Avishek Joey Bose, Tara Akhound-Sadeghe Kilian Fatras, Guillaume Huguet, Jarrid Rector- Brooks, Cheng- Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Se (to-staca flow matching for protein backbone generation. arXiv preprint arXiv:2310.0239) 2023.
[4] G Robert Brakenridge. Global active archive of large flood events, 1985-present. Available from DFO Flood Observatory, University of Colorado.(Accents) on October, 2021 through personal communication), 2021.
[5] Max Braun, Noémie Jaquier, Leonel Rozo, and Tamim Asfour. Riemannian flow matching policy for robot motion learning. In 2024 JEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pages 5144-5151. IEEE, 2024.
[6] Ricky Т. Q. Chen and Yaron Lipman. Flow matching on general geometries. In The Twelfth International Conference on Learning Representations, 2024. URL fetps://openree net/forum? net//oorum22d=//obDI11
[7] Ricky TQ Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03366 03660, 2023.
[8] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[9] Chaoran Cheng, Jiahan Li, Jian Peng, and Ge Liu. Categorical flow matching on statistical manifolds. arXiv preprint 1rX1v:2405.16441 2024.
[10] Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. Advances in neural information processing systems, 35:2406-244 24222 2022.
[11] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow trans- formers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024.
[12] Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models, arXiv preprint arXiv:2410. 12557, 2024.
[13] S. Gallot, D. Hulin, and Lafontaine. Riemannian Geometry, Universitext. Springer Berlin Heidelberg, 2004. ISBN 9783540204930. URL ttttps://books google. com/books id 6F4UmpwS_gUC.
[14] Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. Advances Neural Information Processing Systems, 17:133345-133388 2024.
[15] Mevlana Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on riemannian manifolds. arXiv preprint arXiv: 0026611.020044 2016.
[16] Zhengyang Geng. Ashwini Pokle, William Luo, Justin Lin, and Zico Kolter. Consistency models made easy. arXiv preprint 2rXxv:24006.4448. 2024.
[17] Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint 1ixiv 2018.
[18] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learning Research, 13(1):233-73, 2012.
[19] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851,, 2020.

---

[20] Chin- Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Riemannian diffusion models. Advances in Neural Information Processing Systems, 35:2750- 2761, 2022.
[21] Jachyeong Jo and Sung Ju Hwang. Generative modeling on manifolds through mixture of riemannian diffusion processes. In International Conference on Machine Learning, 2024.
[22] Anuj Karpatne, Imme AibertUUphfft Sai Ravela, Hassan Ali Babaie, and Vipin Kumar. Machine learning for the geosciences: Challenges and opportunities. IEEE Transactions Knowledge and Data Engineering, 114)):544-1554, 2018.
[23] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 24174-24184, 2024.
[24] Dongjun Kim, Chich Heiin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.0227,, 2023.
[25] Xiangzhe Kong, Yinjun Jia, Wenbing Huang, and Yang Liu. Full-atom peptide design with geometric latent diffusion. arXiv preprint arXiv:2402.13555, 2024.
[26] Jiahan Li, Chaoran Cheng, Zuofan Wu, Ruihan Guo, Shitong Luo. Zhizhou Ren, Jian Peng, and Jianzhu Ma. Full-atom peptide design based on multi-modal flow matching. arXiv preprint 27X1v:2406.00735, 2024.
[27] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[28] Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.1455 2022.
[29] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference Learning Representations, 2023. URL https: /
[30] 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 on computer vision and pattern recognition, pages 21264-21273, 2023.
[31] Aaron Lou, Minkai Xu, Adam Farris, and Stefano Ermon. Scaling riemannian diffusion models. Advances in Neural Information Processing Systems, 00588091--8000 2023.
[32] Simon C Lovell, Ian W Davis, W Bryan Arendall III, Paul IW De Bakker, Michael Word, Michael G Prisant, Jane S Richardson, and David C Richardson. Structure validation by ca geometry: Φ, se and cB deviation. Proteins: Structure, Function, and Bioinformatics, 50(3): 437-450, 2003.
[33] Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous time consistency models, arXiv preprint wrXiv:2410.11088 2024.
[34] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
[35] Weijian Luo, Tianyang Hu, Shifeng Zhang. Jiacheng Sun, Zhenguo Li, and Zhihua Zhang. Diff- instruct: universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36:76525-76546, 2023.
[36] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503-2515, 2020.
12

---

[37] Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 14297-14306, 2023.
[38] Laura JW Murray, W Bryan Arendall III David C Richardson, and Jane S Richardson. Rna backbone is rotameric. Proceedings of the National Academy of Sciences, 139024))1300-11999 2003.
[39] National Geophysical Data Center/World Data Service (NGDC/WDS). Ncei/wds global signifi- cant earthquake database, 2019.
[40] National Geophysical Data Center/World Data Service (NGDC/WDS). Ncei/wds global signifi- cant volcanic eruptions database, 2019.
[41] Ilan Price, Alvaro Sanchez-Gonzalez,, Ferran Alet, Tom Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, et al. Probabilistic weather forecasting with machine learning. Nature, 637(80444:84 90, 2025.
[42] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[43] Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence- based generative modeling on manifolds. Advances in neural information processing systems, 34:17669- 17680, 2021.
[44] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprini arXi::2202.00512. 2022.
[45] Axel Sauer, Dominík Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. In European Conference on Computer Vision, pages 87-103. Springer, 2024.
[46] Deborah K Smith, Sherry Harrison, Hong Lin, Shannon Flynn, Mary Nair, Helen Conover, and Sara J Graves. Land, atmosphere near real-time capability for eos (lance) amsr2 data system. In AGU Fall Meeting Abstracts, volume 2016, pages IN33B-1809 2016.
[47] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[48] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 12:12438-12444 2020.
[49] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. URL httpps://openreviev.
[50] Yang Song. Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023.
[51] Peter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich, and Ullrich Köthe. Learning distributions on manifolds with free-form flows. Advances in Neural Information Processing Systems, 37:66961-66667
2024.
[52] 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.
[53] Yilun Xu, Weili Nie, and Arash Vahdat. One-step diffusion models with f-divergence distribu- tion matching. arXiv preprint arXiv:2502. 15681, 2025.
[54] Ling Yang, Zixiang Zhang, Zhilong Zhang, Xingchao Liu, Minkai Xu, Wentao Zhang, Chenlin Meng, Stefano Ermon, and Bin Cui. Consistency flow matching: Defining straight flows with velocity consistency. arXiv preprint arXiv:2407.023398 2024.
13

---

[55] Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Gareia Satorras. Bastiaan S Veeling, Regina Barzilay, Tommi Jaakkola, et al. Fast protein backbone generation with se (3) flow matching. arXiv preprini 10.05223 2023.
[56] Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application protein backbone generation. arXiv preprint arXiv:2302.02277 2023.
[57] Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6613-6623, 2024.
[58] Angxiao Yue. Zichong Wang, and Hongteng Xu. Reqflow: Rectified quaternion flow for efficient and high-qualig, protein backbone generation. arXiv preprint arXiv:2502.1460 2025,
[59] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2023, URL https://openrevice www.eeeeooouuee www.oomiiiooke//ttaaaaaoooo
14

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We have clearly and accurately summarized our main theoretical contributions and practical evaluation results in the abstract and introduction.
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 the 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 can be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have mentioned in Section that one limitation of our proposed work is the additional experiments larger datasets. As theoretical framework, have demonstrated the theoretical stability of RCM on larger datasets and models, and also provided concrete evidence of the superior performance of RCM smaller datasets. In this way, we will leave the scaling experiments for future work.
Guidelines:
The answer NA means that the paper has 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 separate "Limitations" section in 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-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.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a 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 or 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 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 fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, 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 not penalize honesty concerning limitations.
3. Theory assumptions and proofs
15

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We have provided detailed proof in the main text. References and attributes to existing results are also provided.
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 or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be 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 it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We have provided detailed implementation algorithms, dataset information, model architecture, and training hyperparameters in the main text and Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution 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 it possible for others to replicate the model with the same dataset, 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 replicate the results, access 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 NeurIPS 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) If the contribution is primarily new algorithm, the paper should make it 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) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
16

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: All datasets used in paper publicly available. We will also publicize our code once our paper gets accepted.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/ public/gun Soadde/s/CooeSSmmssiiiiiooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / /nips. wttpcc/plc.cccceeeeeeeeeoee gwwp//uuie////ee// /CodeSubmic wwwiiisssiiiiiieieeeeoe
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If 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 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?
Answer: [Yes]
Justification: We have provided detailed dataset information, model architecture, and training hyperparameters in the main text and Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary 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
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: To minimize the impact of stochasticity of the generative model, we draw large number of samples from each model and evaluate the distributional metrics to reflect the significance of our experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

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 error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of 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 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: We have provided the compute resource information in Appendix C.
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, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics attpps://net /tttu::///nurips... wttpeeewww...ppppppeeeeeeeeeeeee
Answer:
[Yes]
Justification:
We conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS 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: [Yes]
Justification: Our work focuses on developing a general generative framework on Rieman- nian manifolds. Our proposed RCM and its variants offer a novel approach to extend the consistency models to Riemannian manifolds. It applies to diverse downstream applications, including image generation, geoscience, and protein design. In this way, our framework can potentially positively impact various domains, and we are dedicated to ensuring responsible usage of our framework.
18

---

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., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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: Our proposed approach is general theoretical framework that poses few risks for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-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 are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All the original creators owners of the external assets used in this work properly cited and attributed.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
19

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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 package should be provided. For popular datasets, paperswi thcode. cow com/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 available online, the authors are 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:
This work does not introduce new assets.
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:
This work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in 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 in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to 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: This work does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
20

---

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 we 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 as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development this work does not involve L.MSs as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https homps://nnuuee //nnruurpp.... /ccc//ooeeeeecc//225///... for what should or should not be described.
21

---

Supplementary Material
A Riemannian Geometry
In this section, we give more detailed mathematical background on Riemannian geometry. A more comprehensive background on the Riemannian manifold can be found in standard mathematical textbooks like [13].
A Riemannian manifold M is real, smooth manifold equipped with a positive definite inner product g on the tangent space T.- at each point I E M. Let TM =ZEM((z) TM be the tangent bundle of the manifold ,M, time-dependent vector field M is mapping u, [0, 1] x M + TM where We (2) E M. A geodesic is locally distance miniaiiinin curve on the manifold. The existence and the uniqueness of the geodesic state that for any point
.M and for any tangent vector u E T,M, there exists a unique geodesic 7 [0, 1] + M such that 7(0) = x and (0) = u. The exponential map exp .M T.M M uniquely defined to be expx (u) 7(1). The logarithm map log M x M + TM is defined as the inverse mapping of the exponential map such that exp, (log, y) y, Vx, y E M.
Let MTTM) denote the space of vector fields on M, a covariant derivative or affine connection
V P(TM) I(TM) + I(TM),
(12)
(u. + V.v
a bilinear map on M that satisfies;
1. V is tensorial in the first argument:
Vartuz 2 Vu, + U. Yel, U2 w E I(TM),
(13)
V fuu fV.V Af U
(14)
2. V is R-llneer in the second argument:
Vu(ii + v2) = Vavi + v2 Vu, 02 ET(TM),
(15)
and it satisfies the Leibniz rule:
Vu(fu) = =(f AS Vf (M)
(16)
The concept of the affine connection generalizes the idea of directional derivatives in the Euclidean case and allows for the differentiation of vector fields at different points. Consider a smooth curve Ya. vector field is said to be purallel-transpottel along the curve 7 if V O. Intuitively, such a condition indicates the vector field u remains "constant" along the curve 7y. We call the mapping Ur M - M, ug + Ue the parallel transport of the tangent vector un along the curve y. Furthermore, a curve 7 is said to be autoparallel if V = 0. Intuitively, a point along such a curve travels at "constant" speed evaluated by the affine connection.
Levi-Civita connection is the unique affine connection that is torsion-fiee and metric compatible Vg = 0. The covariant derivative can be explicitly expanded as
VaV 8h + +1 o'u',
(17)
where 1' are the Christoffel symbols that can be calculated explicitly using the Riemannian metric Piy and the inverse metric g"d as:
1 km (d:gm. + 0.gms - @m.ti)) r 29 9
(18)
In this way, the Christoffel symbols serve as the geometry-aware terms that correct the time derivative of the vector field in the Euclidean case. One important property of the Levi-Civita connection is that every geodesic y is autoparallel:
V
(19)
In this way, the geodesic equation in Eq. 19 can be written locally as:
n + +
0, k 1,..,, n.
(20)
22

---

The differentials of the exponential map d(exp,)u. d(exp T&M + TyM. exp, are linear mappings between tangent spaces. For d(exp.. we z and evaluate the differential with respect to the tangent vector at u; for (l(epp u).. we fix u and evaluate the differential with respect to the base point at x. Intuitively, these differentials describe how a small change change at expx Note that, from the parallel transport equation in 17.17, we can obtain the approximation infinitesimal parallel transport along the direction dæ as
will affect the
du = - u² (dz)²
(21)
Such result indicates that, for non-flat Riemannian manifolds with non-zero Christoffel symbols, there exists an acceleration for enforeing the manifold constraint even if the vector field is "constant" (more precisely, parallel transported along the curve). This is especially important for the calculation of d(exp u) a) where an additional term involving d(exp,)) )u needs to be added to account for the extrinsic change in the vector field because of the base point change.
B Geometry on Specific Riemannian Manifold
In this section, we further provide additional information on the geometric properties of the Rieman- nian manifolds used in this work.
B.1 Euclidean Space and Flat Torus
The Euclidean space with the canonical Euclidean inner product forms Riemannian structure. As expected, the geodesic distance is simply the Euclidean distance. As the manifold is flat, the exponential and logarithm maps simply read:
exp3
u,
(22)
log,
x.
(23)
All Christoffel symbols vanish on the Euclidean space, so the covariant derivative V&V = è reduces to the ordinary time derivative. Direct calculation gives d((xp,)).(t d(exp =))(v))) identity map. This makes sense as the tangent vectors can be parallel transported freely anywhere in the flat space.
the
The flat torus T = (st)a can be viewed as the quotient manifold of the Euclidean space by identifying + 2KT ~ for & E Z along each dimension. Therefore, it is also flat manifold that inherits most properties of the Euclidean space. The only difference that, when calculating the exponential and logarithm maps, shall follow the minimum-image convention along each dimension and wrap the results back to the interval of [0, 2n].
B.2 Spherical Manifold
The n-sphere S is an n-dimennional Riemannian manifold that inherits the canonical inner product from 1=+1 as
(u, u) = (u
U. E ST
(24)
The tangent space Tn = (u, x) = 0} is a n-diiensional hyperplane perpendicular to the vector J. The geodesic the sphere follows the great circle between two points, and the geodesic distance can be calculated as
ds (x. y) arccos v)- The exponential and logarithm maps can be calculated as:
(25)
exp,, u = x cos |ul| + ||z/|| sin ||u]]]
(26)
log, y arccos( (zz, v)) (y x (x, - x)x)).
(27)
(I., y)
The embedded Christoffel symbols read 13 = -² 8.z. and the covariant derivative can be calculated as Vau = i + (u, ..... This indicates that small perturbation dr on the point will lead to a change
of (v, dx)æ on the tangent vector
23

---

Direct differentiation on Eq.26 gives:
sin |/4]
d(exp=)u(v) Un cos ||
+
(U1
(2, u)2),
(28)
where 21| (u, ///// is the parallel component of with respect to and ty ln is the
orthogonal component. Similarly, for d(exp u) we have:
d(cxp (u) = u cos |||]| sin ||u]]
(u, 2).x
(29)
||/]|
The first term is the direct differentiation result. The second term is exactly d((xp..)u( (- v)x), which arises because, as mentioned above, a change in x will lead an additional change in u in the covariant derivative. One can verify that, for any v E TyS", we have (d(exp,)a(x)) exp, u)
(d(exp (v), expa u) = O.
3D Rotation Group B.3
The 3D rotation group SO(3) is a 3-dimensional Riemannian manifold with the Lie group structure. Usually, can be considered as the group of all 3D rotation matrices SO(3) = {R E R3x3 R R I, det (R) 1} [55, 3], with it Lie algebra so(3) consisting of all 3D skew-symmetric matrics so(3) {A E 3333 A -. However, it is worth noting that other formalisms can also be adopted, e.g., the quaternion representation [58]. In this work, we use the 3-vector representation
(i.c., rotation vectors axis-angle representations) for easy derivative calculation. Mathematically, define the vee-hat isomorphism as
a1
03
02
() 3 s0(3),
a2
4
0
(30)
43
02
@1
03
02
d1
( s3 +
R
d3
0
01
H
(31)
02
01
0
03
Then, the 3-vector representations related to rotation matrices via the following Lie exponential map Exp : s0) - 50(3) and Lie logarithm map Log SO(3) + s3):
R = Exp(0) = I + sin ||@||
cos ||0|| go
θ +
(32)
||@||
|lll|
|0||
(R RT), |0|| = arccos Tr(R)= -
0 = Log(R) 2sin [0|
(33)
2
For simplicity, we will omit the vee-hat isomorphism notation and allow the Lie exponential map to apply to a 3-vector R Exp 8 as if the hat has been applied, and also allow the Lie logarithm map to return a 3-vector 0 = Log R if the vee has been applied. The Riemannian exponential and
logarithm maps exp, log : 3 x IR - IR3 can be expressed using the Lie exponential and logarithm maps
exp, Log(Expaa x Exp
(34)
log, exp(-x) Log(Exp(-x)) Exp
(35)
Here, we follow the body-frame convention for the rotation vectors such that the tangent vector can be interpreted as the local angular velocity relative to the global frame. Using the vector,, the canonical bi invariaan Riemannian metric is given by g(u, v) = (u, v), which is equivalent to g(A, B) = TrAA B), A, E so)).. The geodesic distance then reads
d(x, v) log, yl ||Log(Exp(-x) Exp
(36)
The Christoffel symbols r Ekij the Levi-Civits symbols, and the covariant derivative can be calculated as VEU + [ w). where for 3- vector representations, the Lie bracket becomes the cross product [u, v]
24

---

For the calculation of the differentials, the left- and right-Jacobiang JR(0), relate small pertur- bations in the 3-vectors. They are defined as:
1z Σ
1
1)! (0) JR(0) E
(n 1)! (-0)"
(37)
n=0
n=
Specifically, we have (0) = J (0) R (t), where Exp(O). The Baker-
Campbell-Hausdorff (BCH) formula on SO(3) reads:
Exp((0 + d0) 2 Exp / Exp( J (J)(0) d0),
(38)
Exp(O + d0) Exp( (0) d0) Exp
(39)
(ox Exp d0) + JR (0) d0.
(40)
~ + - (0) d8. dd Exp Log(Exp
(41)
With the BCH formula, we can obtain the following approximations of the Riemannian exponential and logarithm maps:
log, (y + dy) ~ JR(Y)dy,
(42)
exp= + du) J (v)
exp2 u, exp2+de (w) (v) J ( = expx u.
(43)
(44)
Combining the above approximations, have
d(exp,)) Jn(u)u.
(45)
Similarly, we can calculate deexp u)x using the above approximations with the additional covariant derivative term as:
((eppu)= = Rv JR(=)o - 2 J u].
(46)
where R., denotes the rotation by the rotation vector u, and the second term again arises from the covariant derivative. The time derivative of the model using the Jacobian-vector product (JVP) should also be modified as:
de dr J (x)
dv ōt
(47)
dt Or where the additional left inverse Jacobian comes from the expansion of the reference vector field using the right expansion BCH formula in Eq.40.
For any u R²3 the Jacobians and their inverse can be calculated in closed form as:
Jzlu)e -J(n-uue -u)e - cos ||u|||
|||||| sin ||u||| [ur, [u,
[u, v]
wI],
(48)
||u|²'
|
1 + cos ||<|||
(u)u " (-u)e ju, u]
[u, [u, ul]-
(49)
|u/12 2||uu||
sin |||]]]
For any v E IR3 the application of rotation vector can calculated using the vector-form Rodrigues' formula:
RAv = v sin |||||| [|0]]
1 1 cos [||u|]
[u, v|
[u, [u, v]]-
(50)
|||
We elaborate why the formulae for SO(3) will incur additional terms and transformations. Specifically, one may notice a seemingly contradiction between the Taylor expansion we used in our proof of Theorem 3.1 of log,, y 2 y - and the expansion for SO(3) in Eq.42. This is because use non- canonical representation for the rotations (rotation vectors instead of rotation matrices), effectively adding an additional coordinate transformation that leads the Jacobians. By definition, smooth Riemannian manifold resembles Euclidean space everywhere locally; therefore, the first-order Taylor expansion the exponential map exp, u + u always holds. However, by considering smooth coordinate mapping & .M- + R",, the Jacobian of such map Jpp TM - IR" needs to be composed during the derivative calculation. Furthermore, for a non-Abelian Lie group, one must define both the left- and raght-Jacobians. Indeed, the above procedure of calculating the Jacobians can be generalized to arbitrary smooth mappings yp, and the other canonical coordinates we used for the sphere and torus can be thought of as adopting the identity mapping with the identity Jacobian Jrp id. Under the assumption of canonical coordinate, our theoretical result Theorem 3.1 safely holds for all smooth Riemannian manifolds.
25

---

C
Experimental Detail
In this section, we provide comprehensive details on the model parameterization, training setup, and additional techniques used to improve the performance of RCM.
C.1 Model Parameterization
We found that the magnitude-pressrrinng design principle of EDM2 [23] is very effective in stabilizing the JVP computation. However, since our task and datasets are relatively small (compared to the image generation task), rebuilt simple presservng multilayer perceptron (MLP) using magnitude preseervingg EDM2 modules.
Magnitude--reservvinn fully-connected layer fully-connected layer with input activation x and output activation is calculated as follows:
MP-FC(x)
(51)
|
Bagnitude-preservivn Fourier features function as follows:
The Fourier features are scaled by v2 using cosine
c2 91)) (ft (2))
MP-Fourier(t)-
(52)
V (fx +NN)
Magnitude-pressrvinn SiLU The SiLU nonlinear function should also be scaled to maintain the magnitude as follows:
MP-SiLU(z)) 0.596 (1 ee))
x
(53)
Magnitude-preserving Concatenation and Sum The concatenation of two input activations x and y, scaled by constants w, and wy. with blend factor and is calculated as follows:
[0, 1] to adjust the balance between
N₂ N₂ MP-Cat(Z, V (1 a)² a²
1-a VN=
y
(54)
VNy"
where, N₂ and Ny is the size of and y concatenation dimension. A similar formula is used for the sum operation follows:
MP-Sum a) (1 a)x + ay
(55)
v(1 - a)2
Magnitude-preserviie Block By using the above operations, we design a simple multilayer perceptron and residual connections fuse noisy data and time embedding. Concretely, for inputs and emb, we perform the following computation:
h MP-FC(MMP-SLLUNNomm22
(56)
Remb MP-FC(emb)
(57)
MP-Sum(Norm(z), (PM hemb)) 0.3)
(58)
Model Architectural We use MP-Blocks to build our model. First, we stack the blocks in the encoder gradually increase the hidden dimension. Subsequently, the decoder, we use the same blocks and concatenate the output of the corresponding layer of the encoder, similar to what UNet does. The time encoded using MP-Fourier and passed through an MP-FC, which is subsequently fed into each block.
26

---

Hyperparameter In all our experiments, we use 4 blocks, each with dimensions [256, 512, 512, 256]. and the dimension of time embedding 256. For flow matching, we use learning rate of 10-9, while for the consistency model, we use a learning rate of 10-1 with the Adam optimizer. The batch size varies depending on the dataset size and typically 512 or 1024. We not use dropout, nor do we employ other tricks such as learning rate decay. We train our model using total of 50 million data samples.
All experiments were carried out on a single A100. The maximum GPU memory consumed is only around 5GB, which can be easily fit into GPUs with smaller memory. As we fixed the number of iterations, the entire training of each variant of CM models took approximately 4 hours, regardless of the dataset size.
C.2 Model Sampling
We described the Riemannian flow matching sampling and our Riemannian consistency model sampling Algorithm 3 and respectively. The sampling algorithm for RCM consistent with [50] on Euclidean cases. In each sampling step, the model first makes a prediction of the denoised results with the learned vector field (consistency parameterization in Eq.3). A geodesic interpolation step then follows to bring the prediction back to intermediate noisy data if there is more than one sampling step.
Algorithm 3 RFM Sampling
1: Input: Trained RFM uy. number of steps N.
2: Sample noise Tp-
3: for in 1/N, 2/N, ,..., - 1/N do
4: T++1/N exPx, (-Krvo(Z,, t)/N).
>> Euler step
5: end for
6: Return: T1.
Algorithm 4 RCM Sampling
1: Input: Trained RFM vo, number of steps N.
2: Sample initial noise Tn-
3: for in 0, 1/N,2/N,.....
1/N do
4: Ii = exP,, (I),
Consistency parameterization
5: Sample noise Z0.
6: T++1/N = expx, (K441/N log,, zo).
Interpolation to the next time step
7: end for
8: Return: 21-
C.3 Training Technique
Tangent Warmup In line with the findings of Lu and Song [33], we identified the covariant derivative term in our loss likely contributor to JVP instability. Consequently, we implemented strategy of incrementally adding this term throughout the training process, detailed below:
f d(exp,) (Ra) + d(exp u),, (+) + d(exx.)) (N 2
(59)
where r linearly increases from 0 to over the first 10k training iterations. We made an interesting observation, not previously noted in and Song [33], regarding the parameter r. Specifically, when O, the loss function degenerates into a flow matching loss. As progressively increases, the model gradually transitions from a flow matching objective towards that of a consistency model. This progression is analogous to the methodology in ECT [16], which involves a reduction of dt in discrete time steps. This insight offers unified explanation for these two approaches, despite their different conceptual starting points.
Tangent Clipping To further stabilize the gradient variance consistency models, we build upon the concept of tangent normalization proposed in Lu and Song [33]. We observed that when the
27

---

tangent norm is small, standard normalization preserves only directional information, which can impede further optimization. Therefore, drawing inspiration from gradient clipping, we introduce 'tangent clipping' mechanism, detailed below.
clip(/) = max((M/l/iIl)
(60)
where represents the maximum tangent norm, which we set to our experiments. This approach designed to maintain low gradient variance without inappropriately rescaling the tangent when its norm is small.
C.4 Evaluation Details
For the estimation of kernel density on the 2-sphere, we use the off-the-shelf implementation from Scikit-Leam with the haversine distance and the von Mises- kernel to match the spherical manifold. We follow [36] choose the bandwidth of 0.02 and use meshgrid of 90 longitudes and 90 latitudes for calculating the empirical KLD, For MMD calculation, we always use a bandwidth of 1 and use the geodesic distance the exponential kernel match the corresponding manifold.
For the scalability experiment on high-dimensional tori [0, 2nr] D we first estimate the (diagonal) wrapped Gaussian using the maximum likelihood estimation (MLE) formula as
arctan
Σ sin Zi
cOs Tj Σ
(61)
2
2
N
N
1
8² log
Σ
sin F2
COS I
(62)
N
N -
where the index is over the data points. The above formula is applied each manifold dimension (assuming a diagonal wrapped Gaussian). We then calculate the Fréchet distance (FD) to the ground truth standard Gaussian as:
FD Σns+ 1)
(63)
where the index is over the manifold dimension.
D Additional Results
In this section, provide additional experimental results to further support the effectiveness, efficiency, and scalability of RCM.
D.1 Impact of NFEs
We further provide comprehensive study on the impact of the number of function evaluations (NFEs) on the final generation quality. We use the same model checkpoints trained on the earthquake dataset and alter only the sampling steps in {2, 10, 20, 50, 100) as the NFEs. The results are summarized in Figure 4. We have the following interesting observations:
For the RFM model, increasing the NFEs has significant impact on the generation quality, as the learned marginal vector fields are not necessarily straight, leading to significant errors in the few-step generation setup.
RCD and RCT generations more stable with respect NFEs, consistently achieving decent generation quality. This indicates RCM can indeed shorteut the probability path on the Riemannian manifold.
RCM variants consistently outperform the baselines, especially the naive Euclidean CM approach, demonstrating the necessity of Riemannian constraints. For RFM. the two lines intersect at around 20 NFEs, below which the RFM generation quality drastically drops.
hatps://scikitteeamm.rr//
28

---

10.0
RFM
sRCD
RCD
dRCD
CD
6.0
RCT
SRCT
KLD
4.0
3.0
2.0
1.5
2
10
50
NFEs
Figure 4: KLD vs NFEs for different models on the Earthquake dataset. RCD variants are colored in green, and RCT variants colored in purple.
shall be noted that the speed-quality tradeoff for the standard flow matching model (as demonstrated in the blue RFM curve) does not generally apply to CMs. Our NFE vs performance results generally echo the findings Kim al. [24], in which high NFEs may instead lead poorer generations. In practice, the improvement in terms of KLD is also less significant for all CMs. It is also possible to adapt the techniques in Kim al. [24] for Riemannian manifolds, resulting finer-grained control over the sampling stage, which we leave for future work.
D.2 Practical Time Complexity and Speedup
Intuitively, as the training and sampling procedure of RCM follows the Fuclidean consistency model, the training time should be approximately the same while enjoying speedup of the ratio of the sampling steps (NFEs) needed. In this way, the 2-step setup will incur 50x speedup during inference. To provide more concrete quantitative results, we benchmarked the training stage time RCM and sRCM compared to the Euclidean version (CD and the sampling stage speedup of two-step RCM compared to the 100-step RFM. The results are summarized in Table 5.
Table 5: Training time overhead compared CD and sampling time speedup for RCM-2 compared to RFM-100 on three different datasets.
Dataset
Earth
RNA
SO(3)
RCM train overhead
+5.9%
+<0.1%
+10.7%
sRCM train overhead
+4.1%
+<0.1 %
+7.8%
RCM-2 sample speedup
x46.5
x43.2
For training, it clear that even with fairly small model, the additional Riemannian operations have minimal overhead in general. It is also expected that the overhead for SO(3) is larger than the sphere, as the former requires more operations, like the calculation of left right Jacobians. Similarly, the simplified loss leads to fewer overheads as it does not require the calculation of the differentials of exponential map. Generally, the overheads for RCm are still within 10% range and will be even smaller for larger models, as the Riemannian operators only scale linearly with the data dimension but are model-indeeennce For sampling, the speedups are slightly below 50, probably due to the additional Riemannian operators and noising sampling in Algorithm 4. Despite this, the speedup
29

---

is drastic, demonstrating the effectiveness of RCM variants in few-step generation scenarios and offering perfect solution for sampling efficiency.
D.3
Additional Visualization
Data
RFM-100
RFM-2
RCD
dRCD
CD
RCT
Board
Protein
Figure 5: Generations on the 2D flat torus. The FM, CD, and CT models are colored in blue, green, and orange.
For the 2D torus datasets, we provide visualizations of the generations the heat maps in Figure 5. Similar trends can be observed comparing different baselines. Interestingly, although the discrete-time RCD model achieved better MMD score on the protein dataset, its visual results are not good as RCD, as there are noticeable artifacts on the minor mode on the right in the dRCD generation.
For the spherical manifold, we provide generated points on all four datasets across all models in Figure The number of generated samples is exactly the same the corresponding ground truth data for fair comparison. It can be clearly seen that, while RFM-2 could still capture some coarse-grain features of the overall distribution, failed capture the finer-grain modes accurately. Similar phenomena can be observed on the CD baseline, where the Euclidean assumption does not necessarily capture the curved Riemannian geometry.
30

---

Earthquake
Volcano
Fire
Flood
Data
RFM-100
RFM-2
SRCD
RCD
dRCD
CD
naive
SRCT
RCT
Figure 6: Generations on the four datasets on the 2-sphere. The FM, CD, and CT models are colored in blue, green, and orange.
31


Reviews:
[{"review_id": "iKNtN04d2q", "reviewer": "Reviewer_Ezj8", "summary": "This paper presents an extension of Consistency Models to the case of non-Euclidean Riemannian manifolds. The ability to accelerate diffusion/flow matching on such manifolds is impactful for a number of scientific applications ranging from protein generation to astrophysics and had until this paper not received sufficient attention. The paper develops the key results necessary to derive both consistency distillation and consistency training on Riemannian manifolds, introduces simplified objectives that bypass some of complexities of geometric calculus, and performs adequate evaluations of the different variants of the proposed models on established and appropriate benchmark problems.", "questions": "- I noticed a missing stop gradient in Eq 2", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n- The paper is very well written, provides appropriate background and survey of the related literature, and includes useful insights to help the reader understand the different aspects of the consistency loss on Riemanian manifolds. \n- The derivations are mathematically sound (as far as I can tell), providing the necessary key results to build consistency models on Riemannian manifolds.\n- The evaluations are appropriate and demonstrate 1. the gain of using a proper riemannian distillation strategy over a naive implementation of a Euclidean consistency model in the ambient Euclidean space, 2. the speed up made possible by consistency modeling compared to standard flow matching.\n\nThis is a well rounded paper, I do not have significant weaknesses to point out.", "was_revised": true, "final_justification": "This is a sound paper that develops the required results for building consistency models on Riemannian manifolds and I am satisfied with the practical evaluations of the proposed method.", "rebuttal": "We sincerely appreciate your high recognition of our work's clear structure, solid mathematical background, and comprehensive evaluation. As per request by the other reviewers, we also provide the following additional experiments to make our work more comprehensive.\n- Scalability to higher-dimension manifolds with a dimension up to 256, in Reviewer aiD6, Q1. RCM consistently achieves better generation quality.\n- NFEs vs performance evaluation, in Reviewer aiD6, Q3. RCM has a stable performance compared to RFM.\n- Additional FFF baseline results (a Riemannian normalizing flow), in Reviewer J7YW, Q6. RCM consistently outperforms FFF.\n\nWe will continue to refine our work and explore the important applications of our proposed RCM in various AI4Science domains, where Riemannian manifolds are frequently encountered. We hope that our framework, as one of the earliest attempts to introduce a mathematically sound approach to Riemannian few-step generation, can also inspire future work."}, {"review_id": "grM0AnObRt", "reviewer": "Reviewer_nYt4", "summary": "This paper introduces the Riemannian Consistency Model (RCM), a framework to extend consistency models to general Riemannian manifolds for fast, few-step generation. The authors present a principled approach using the exponential map and covariant derivative to correctly enforce the consistency objective while respecting the manifold's geometry. The work derives closed-form objectives for both distillation-based (RCD) and training-from-scratch (RCT) versions of the model and proves their theoretical equivalence. A key practical contribution is a simplified objective that is easier to implement without sacrificing performance. Through extensive experiments on spheres, tori, and the SO(3) rotation group, the paper demonstrates that RCM significantly outperforms naive baselines and few-step flow matching, highlighting the importance of a geometry-aware approach for generative modeling on non-Euclidean data.", "questions": "- Your appendix mentions the use of \"Tangent Warmup\" and \"Tangent Clipping\" for stable training. Could you reflect on the practical difficulty of working with the proposed objectives? What is the primary source of the instability that necessitates these heuristics, and does the simplified sRCM objective require them to the same extent as the full RCM?\n- The paper successfully applies RCM to tori, spheres, and SO(3). For which other classes of manifolds are the required geometric quantities (exponential/logarithm maps, Christoffel symbols) readily available or straightforward to derive? A brief discussion on the scope of manifolds where RCM is practically applicable would be insightful.\n- The experiments focus almost exclusively on a 2-step generation scenario. How does the model's performance and the quality of the generated samples evolve as the number of inference steps increases to 4, 8, or 16? An analysis of this quality-vs-computation trade-off is crucial for understanding the model's behavior.\n- Could you provide an analysis of the runtime complexity introduced by the covariant derivative calculation during training? How does this computational overhead compare to the Euclidean consistency models, and how is it expected to scale with the dimensionality of the manifold?", "limitations": "Yes.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## **Strengths**\n\n- As far as I am aware, this work appears to be the first to develop a comprehensive and theoretically sound framework for consistency models on general Riemannian manifolds. This is a significant step forward, as it enables fast, high-quality generative modeling for a wide range of important scientific data (e.g., protein structures, robotics, geoscience) that naturally reside on non-Euclidean domains. \n- The paper is technically very strong; both the the continuous-time RCM objective which seems a highly non-trivial result, but also the equivalence between the distillation (RCD) and training (RCT) variants  to further solidify the theoretical foundation. Also the kinematic perspective is really great, and the author are able to present rather intricate formalisms in a physical-intuitive manner making the paper quite accessible.  \n-  The authors demonstrate strong practical utility by not only presenting the full, complex objective but also proposing a simplified version (sRCM). This simplified loss eliminates the need to compute the differentials of the exponential map, which can be a significant implementation issue, yet it achieves excellent empirical results. \n\n\n## **Weaknesses** \n- The appendix shows that the method requires non-standard techniques like \"Tangent Warmup\" and \"Tangent Clipping\" to ensure stable training. The need for these specialized heuristics suggests that the proposed loss functions, while theoretically derived, may be inherently difficult to optimize and sensitive to the magnitude of the geometric correction terms. Could the authors reflect on how hard it is to work with the method in practice? \n-  The practical application of RCM is challenging. First, applying the model to a new manifold requires deriving the specific forms of the exponential map, logarithm map, and Christoffel symbols for the covariant derivative. Can the author illustrate more for what manifolds this is doable?  \n- The empirical evaluation, while covering several manifolds, is limited in two key ways. First, the experiments are almost exclusively focused on a 2-step generation scenario, with no analysis of the quality vs. computation trade-off at other few-step counts (e.g., 4, 8, 16 steps), which would give significant insight in the general trend in comparison. Second, the datasets used are of a relatively small scale and low dimensionality. This makes it difficult to assess the method's scalability and performance characteristics on larger, more complex problems.\n- The paper seems to not analyze the runtime complexity of the proposed method. The calculation of the covariant derivative term at each training step is computationally more expensive than the simple time derivative used in Euclidean models. This introduces an additional, unquantified computational overhead that may become significant for high-dimensional manifolds.", "was_revised": true, "final_justification": "This work is very solid - I had some initial doubts, but after the rebuttal in which my main issues were addressed, I recommend accepting this work.", "rebuttal": "We sincerely appreciate your recognition of the solid theoretical contributions and innovations in our work, which extend the existing consistency models on Euclidean data. We will address your questions as follows.\n\n## Q1. Tangent Warmup\n\nWe thank your insightful suggestions on the stability of the training stage of consistency models, and we are more than happy to share our findings on the practical training stage of our proposed RCM model. \n\nWe first note that training stability has been recognized as a common challenge in many continuous-time consistency models, and not particularly for Riemannian manifolds. For example, the original CM paper did not run continuous-time CM on image generation, although it was also proposed; the ECT paper addressed such \"curse of consistency\" by gradually annealing the $\\Delta t$; the sCM paper proposed many essential techniques for scaling up continuous-time CM models. \n\nIn fact, the special techniques (e.g., tangent warmup and clipping) were directly inspired by the Euclidean sCM paper to address the issue that the initial JVP of the model, even when loaded from pre-trained weights, is noisy for the consistency constraint. Empirically, we did not find the additional geometric correction terms add special challenges to the training, and therefore, the standard tangent warmup technique is effective enough to stabilize the training process, allowing us to obtain decent and stable results across all continuous-time variants. As a concrete example on the Earthquake dataset with RCD, RCD was able to achieve 2.22 with both techniques, while no-warmup achieved 3.81, no-clip achieved 2.11, and no-warmup-no-clip achieved 2.70. It can be seen that the tangent warmup step is more important for the performance. While clipping has some negative impact on the performance, we empirically found it helps prevent gradient exploding in the early stage.\n\n## Q2 Practical Applications to Different Manifolds\n\nWe have proposed an alternative **simplified loss** in Section 3.2 that **only requires the additional information of the covariant derivative** on the manifold. The simplified loss eliminates the need for potentially complex symbolic calculations of the differentials of the exponential map for each different Riemannian manifold, achieving almost the same performance. In this way, the RCM framework can be easily applied to a wide range of manifolds. For example, in addition to the manifold we have experimented with, the standard hyperbolic manifold has known Christoffel symbol expressions, and all matrix Lie groups (e.g., SE(3), SU(2)) share a similar covariant derivative expression [1]. We will follow your suggestion to add a discussion in our revised manuscript to inspire potentially new applications.\n\n[1] Guigui, Nicolas, and Xavier Pennec. \"A reduced parallel transport equation on Lie groups with a left-invariant metric.\" International Conference on Geometric Science of Information. Cham: Springer International Publishing, 2021.\n\n## Q3 Quality vs NFE\n\nWe provide additional ablation results on the model performance vs NFEs. Please kindly refer to our rebuttal to Reviewer aiD6, Q3. To summarize, in contrast to RFM, whose performance drastically drops when decreasing the NFEs, our RCM variants exhibit relatively stable performance. Such results demonstrate RCM can indeed shortcut the probability path on the Riemannian manifold. Our NFE vs performance results generally echo the findings in the consistency trajectory model [2], in which high NFEs may instead lead to poorer generations.\n\n[2] Kim, Dongjun, et al. \"Consistency trajectory models: Learning probability flow ode trajectory of diffusion.\" arXiv preprint arXiv:2310.02279 (2023).\n\n## Q4 Scalability\n\nScalability has always been a key target for RCM. However, due to the limited availability of large-scale datasets in this domain compared to traditional image generation tasks in previous CM papers, we follow previous Riemannian generative models (e.g., RFM, Riemannian diffusion, Riemannian score matching) to focus on a wide range of different manifolds instead. Additionally, to demonstrate RCM's scalability on higher-dimensional manifolds, we have conducted experiments with a manifold dimension up to 256 following the RFM paper. Please refer to Reviewer aiD6, Q1, for the detailed results. To summarize, RCM **consistently outperforms RFM in few-step generation** even on higher-dimensional manifolds.\n\nIn the runtime complexity analysis below, we further demonstrate that the overhead from the additional geometric components is both theoretically and practically negligible. Therefore, we believe RCM can be easily adapted for larger data and models, which we will continue to explore in future work.\n\n\n## Q5 Runtime Complexity Analysis\n\nWe provide a comprehensive analysis of the Riemannian operations. In short, we demonstrate **both asymptotically and practically** that the overheads of the additional Riemannian operations (e.g., covariative derivative, differentials of the exponential map) are **negligible**.\n\nAsymptotically, as the Riemannian operations used in the RCM algorithm are linear operators, the time complexity is $O(LD^2)$ where $L$ is the data dimension and $D$ is the number of manifold dimensions where the Jacobian is non-diagonal. (Therefore, the n-torus is actually the data dimension of n because it is the direct product of n 1-tori.) In practice, $L$ can be large (e.g., number of pixels in the Euclidean case), but $D$ is limited, so the behavior **scales linearly as the data dimension and is model-independent**. Practically, we observed that the model's forward and backward passes were the bottleneck. Specifically, the relative overheads during training are provided below.\n\n| overhead/% | RCM  | sRCM |\n| ---------- | :--: | :--: |\n| Earth      | 5.9  | 4.1  |\n| RNA        | <0.1 |  -   |\n| SO(3)      | 10.7 | 7.8  |\n\nIt can be clearly seen that, even with a fairly small model, the additional Riemannian operations have little overhead in general. It is also expected that the overhead for SO(3) is larger than the sphere, as the former requires more operations, like the calculation of left/right Jacobians. However, it is still within 10% and will be even smaller for larger models."}, {"review_id": "j9ISlI0mtQ", "reviewer": "Reviewer_aiD6", "summary": "This paper proposes the Riemannian Consistency Model (RCM), allowing few-step generative modeling on Riemannian manifolds (e.g., spheres, tori, SO(3)). It adapts consistency models to respect manifold geometry using exponential maps and covariant derivatives, and includes both theoretical and practical improvements. RCM outperforms previous Euclidean and flow matching baselines in experiments.", "questions": "Please see the weaknesses.", "limitations": "Yes, although authors mention experiments on larger datasets are their limitations, however, baseline [1] do report performance on these datasets. I would advise authors to include at least one such dataset to showcase the scalability and others can fall under this limitation.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- Extends consistency models to non-Euclidean manifolds and provides the simplified training objective\n- Strong theoretical analysis and proofs to support the proposed approach\n- Results on multiple manifolds and datasets with respect to the selected baselines\n\nWeaknesses:\n- The main experimental limitations of the RCM paper, compared to RFM, are a narrower focus on simple manifolds, lack of high-dimensional and general (mesh/bounded) manifold experiments after the literature [1].\n- Missing baselines: Diffusion Model, Riemannian Diffusion Model, Riemannian Score-Model, etc.\n- Missing ablation: (NFE vs. Performance) plot could give a better overview of the effectiveness of the proposed approach. \n\n[1] Chen, Ricky TQ, and Yaron Lipman. \"Riemannian flow matching on general geometries.\" arXiv e-prints (2023): arXiv-2302.", "was_revised": true, "final_justification": "Authors have addressed my concerns.", "rebuttal": "We sincerely thank you for recognizing our work's strong analysis and proof to extend consistency models to multiple Riemannian manifolds. We are happy to answer your questions as follows.\n\n## Q1 Application to High-Dimensional Manifolds\n\nWe thank you for your suggestions on the additional experiments to demonstrate the scalability of our RCM to high-dimensional manifolds. To achieve this, we follow the RFM and [1] to conduct additional experiments on high-dimensional synthetic data on the $n$-torus. Following [1], we used the truncated standard Gaussian on the domain $[0,2\\pi]^n$ as the target and sampled 10k points as the training set. Each model was trained for 2000 epochs, and we estimated the Gaussian parameters using maximum likelihood estimation with truncated Gaussians on 10k generated samples. We then calculated the Fréchet distance (Wasserstein-2 distance) to the ground truth standard normal parameters as the metric. For a manifold dimension of $2^k$ where $k=1,2,...,8$, the results are summarized in the following table.\n\n| dim | 2 | 4 | 8 | 16 | 32 | 64 | 128 | 256 |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| RFM-2 | 0.52 | 0.70 | 1.01 | 1.41 | 1.95 | 1.47 | 1.83 | 2.41 |\n| RCT-2 | **0.22** | **0.31** | **0.54** | **0.81** | **0.46** | **0.58** | **0.62** | **0.96** |\n| CT naive-2 | 0.73 | 2.69 | 1.58 | 2.16 | 2.41 | 24.80 | 35.16 | 49.24 |\n\nIt can be clearly seen that our proposed RCM model **consistently outperforms the RFM model in a few-step generation setup**. Interestingly, the naive Euclidean CT baseline quickly fails as the manifold growing dimension adds to the difficulty in enforcing the periodic constraint in every dimension.\n\nWe also provide an asymptotical analysis on the additional cost of Riemannian operations (e.g., covariant derivative). Please kindly refer to our rebuttal to Reviewer nYt4, Q5. In short, the overheads regarding the additional Riemannian operators are **small compared to the naive Euclidean CM setup and are also model-independent**. Therefore, we believe our proposed RCM can scale to large data and models.\n\nRegarding the extension to general manifolds like meshes, we first noted that, in principle, like Riemannian flow matching, RCM can be used for general geometries but requires expensive simulation (of the geodesic and covariant derivatives). Therefore, we acknowledge that this is one of the limitations of our model so far, which we are actively exploring. We noted that most existing work, like [1] and [2], operates on the manifolds with known exponential and logarithm maps. We emphasize our theoretical contributions in such manifolds, which already exhibit profound practical applications as indicated in previous work.\n\n[1] De Bortoli, Valentin, et al. \"Riemannian score-based generative modelling.\" Advances in neural information processing systems 35 (2022): 2406-2422.\n\n[2] Mathieu, Emile, and Maximilian Nickel. \"Riemannian continuous normalizing flows.\" Advances in neural information processing systems 33 (2020): 2503-2515.\n\n\n## Q2 Baselines Question\n\nComprehensive baselines have always been a key target for us. However, we noted that the Riemannian diffusion and Riemannian score model are not directly comparable baselines in our setup, based on the following reasons:\n\n- Our proposed RCM framework **targets the few-step generation setup** on the Riemannian manifold, whereas the above two baselines require 100-1000 steps for decent results. They are baselines to models like RFM but **not directly comparable** to our few-step generation results.\n- Even for simple geometries like spheres, the diffusion-based approach may require **expensive simulation** of the model Jacobian trace (as detailed in the RFM paper, Appendix D), due to its reliance on the Gaussian-like noise.\n- The Riemannian diffusion and score models fall short of the RFM model (in the RFM paper results).\n\nInstead, our baselines comprise comparable algorithms in the few-step generation setup on the Riemannian manifold, e.g., the original CM model, designed for Euclidean data (referred to as CD/CT-naive in our paper), and the RFM with its few-step generation results. In this way, our baselines already covered the comparable algorithms and demonstrated the effectiveness of our proposed RCM. \n\nIn addition, we provide additional baseline results of FFF, a Riemannian normalizing flow that enables one-step generation, as per request by Reviewer J7YW, Q6. RCM consistently outperforms this new baseline, further demonstrating its effectiveness. \n\n## Q3 Performance vs NFEs\n\nWe thank your insightful suggestions on further ablations of the performance vs NFEs. Below, we provide comprehensive results on the Earthquake dataset, examining the impact of NFEs on generation quality (measured by KDE-KLD in this case). The best results, except for RFM, are highlighted in bold, while the second-best result is italicized.\n\n| NFE | RFM | sRCD | RCD | dRCD | CD-naive | sRCT | RCT |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| 2 | 10.31 | **1.91** | _2.14_ | 6.52 | 3.14 | 2.26 | 2.40 |\n| 5 | 4.74 | _1.83_ | **1.76** | 4.10 | 2.78 | 1.84 | 2.64 |\n| 10 | 3.24 | _2.09_ | **1.82** | 3.87 | 3.02 | 2.31 | 2.72 |\n| 20 | 1.90 | _2.32_ | **2.10** | 4.06 | 2.96 | 2.30 | 2.66 |\n| 50 | 1.53 | _2.41_ | **2.04** | 4.18 | 3.12 | 2.23 | 2.78 |\n| 100 | 1.47 | 2.64 | **2.02** | 4.72 | 3.16 | _2.15_ | 2.75 |\n\nWe have the following interesting observations:\n\n- For the RFM model, increasing the NFEs has a significant impact on the generation quality, as its learned marginal vector fields are not necessarily straight for few-step generation.\n- On the other hand, RCD & RCT generations are more stable with respect to the NFEs, achieving consistently good generation quality. This indicates RCM can indeed shortcut the probability path on the Riemannian manifold.\n- RCM variants **consistently outperform the baselines**, especially the naive Euclidean CM approach, demonstrating the necessity of Riemannian constraints. For RFM, the two lines intersect at around 20 NFEs, below which the RFM generation quality drastically drops.\n\nIt shall be noted that the speed-quality tradeoff for the standard flow matching model does not apply to CMs. Our NFE vs performance results generally echo the findings in the consistency trajectory model [3], in which high NFEs may instead lead to poorer generations.\n\nWe also observe similar trends in other manifolds where RCD & RCT models have stable generation quality compared to RFM. Due to the limited space, we will add these results to our revised manuscript.\n\n\n[3] Kim, Dongjun, et al. \"Consistency trajectory models: Learning probability flow ode trajectory of diffusion.\" arXiv preprint arXiv:2310.02279 (2023)."}, {"review_id": "l50qVfmrbf", "reviewer": "Reviewer_iNgD", "summary": "The paper introduces Riemannian Consistency Models (RCM) by extending the traditional consistency model to Riemannian manifolds. In the process, they derive a closed-form continuous-time objective tailored to manifold-specific operations. By applying the marginalization trick, the authors prove the mathematical equivalence of Riemannian consistency training and distillation. Additionally, a simpler objective is proposed that avoids computing differentials of the exponential map, yet maintains performance, with a kinematic interpretation offering intuitive insight into the method.", "questions": "- _[Line 86-87]_ Kindly define the weights $w_t$. Do they follow a similar specification to that in [1] or [3]?\n- _[Line 127]_ As pertaining will always include an inherent additional error due to generalization to domains, what should be a sufficient condition for $\\dot{x}$ such that RCM converges to the target distribution? \n- _[Line 245]_ Is there any rationale behind sticking to the linear scheduler? \n- _[Line 252]_ The choice of KDE to obtain density estimates is questionable, since in higher dimensions (even as low as 10-15), choosing suitable bandwidths becomes difficult. At best, carefully chosen kernels can recover error bounds adapted to intrinsic dimensions under the manifold hypothesis (see [4]), theoretically. In practice, there is hardly any improvement.   \n- There remains the issue of training instability that can occur when consistency models are trained from scratch in a naïve manner—specifically, by directly optimizing the objective as in (4) while closely adhering to $\\Delta t \\approx 0$. The authors of [3] call it the Curse of Consistency. In particular, the model tends to converge slowly after a certain point, perhaps due to accumulated consistency errors from each interval. Does the proposed model experience anything similar, since it follows a simple training regime akin to [2] (following _Line 87_,  $\\Delta t$ is quite small for large $N$)? \n- Heuristically speaking, deep generative models often can be observed to work well enough, being agnostic to the fact that underlying data distributions are supported on intrinsic structures embedded into the ambient space. They silently adopt the intrinsic dimensions in their error rates (e.g., GANs [5], WAE [6], etc. However, to my knowledge, error bounds on diffusion model-based generations are yet to be shown to follow the same.). In this context, does RCM really generalize the Euclidean case _[Line 139]_ or rather extend the framework onto manifolds? Given the additional costs due to hyperbolic/Riemannian operations, is it really meaningful to propose a new model altogether? \n- Also, kindly respond to the points raised as 'Weakness'.\n\n> [3] Zhengyang Geng, Ashwini Pokle, Weijian Luo, Justin Lin, and J Zico Kolter. 2025. Consistency Models Made Easy. In The Thirteenth International Conference on Learning Representations. \n> \n> [4] Jisu Kim, Jaehyeok Shin, Alessandro Rinaldo, and Larry Wasserman. Uniform convergence rate of the kernel density estimator adaptive to intrinsic volume dimension. In International Conference on Machine Learning, pp. 3398-3407. PMLR, 2019.\n>\n> [5] Jian Huang, Yuling Jiao, Zhen Li, Shiao Liu, Yang Wang, and Yunfei Yang. An error analysis of generative adversarial networks for learning distributions. Journal of machine learning research 23, no. 116 (2022): 1-43.\n>\n> [6] Saptarshi Chakraborty and Peter Bartlett. A Statistical Analysis of Wasserstein Autoencoders for Intrinsically Low-dimensional Data. In The Twelfth International Conference on Learning Representations. 2024.", "limitations": "While the article does not have a dedicated section explaining limitations, the apparent ones mostly relate to empirical validation.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strength**\n> The paper is very well-written, easy to follow, with sufficient discussion on similar approaches extending diffusion/flow models in non-Euclidean data supports. Supplementary notes on rudimentary notions of geometry are appreciable.\n\n**Weakness**\n>- While the motivation to adopt consistency models, amenable to data supported on a Riemannian manifold, is clear, the paper lacks empirical validation. For example, what is the extent of improvement in inference time (or convergence) over both Riemannian diffusion models [1] and/or vanilla consistency models [2].  \n>- Given that the calculations in a non-Euclidean space are far more computationally demanding compared to their Euclidean counterparts, the paper must show that the advantages of adopting this approach outweigh the added processing cost. This is missing in the current article.\n>- Computational infrastructure under which experiments were carried out, code repositories, and specific execution details seem to be missing in the current article. \n\n\n> [1] Huang, Chin-Wei, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C. Courville. Riemannian diffusion models. Advances in Neural Information Processing Systems 35 (2022): 2750-2761\n>\n> [2] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. 2023. Consistency models. In Proceedings of the 40th International Conference on Machine Learning (ICML'23), Vol. 202. JMLR.org, Article 1335, 32211–32252.", "was_revised": true, "final_justification": "The authors' responses have been prompt and insightful. Considering all the reviews, the subsequent discussions, and my initial assessment, I believe it is appropriate to revise my score upward.", "rebuttal": "We sincerely appreciate your comprehensive review and your recognition of the clear structure of our work. We are happy to address any remaining questions you may have.\n\n## Q1 Empirical Improvement over Existing Work\n\nWe first thank your recognition of our theoretical contributions and applications. Here, we provide a more comprehensive discussion on the empirical improvements. \n\nThe inference time improvement of RCM, as is the case for regular CMs for Euclidean data, is straightforward. This is because the inference code is identical to that in the RFM except for a smaller number of Euler steps. As a concrete example, in most of our experimental setups, the RFM model requires approximately 100 inference steps. In contrast, our RCM runs with just two steps, resulting in a **50x speedup during inference**. Furthermore, compared to RFM or Riemannian diffusion under a few-step setup, our results also demonstrate **better generation quality** (as shown in the RFM-2 column), highlighting the effectiveness of RCM in this context on Riemannian manifolds. Lastly, we provided the CD-naive baseline, which is **equivalent to the classic consistency model on Euclidean data**. On most datasets, even after the final projection to the manifold, the Euclidean assumptions often fail to capture the true data distribution, resulting in lower generation quality and necessitating manifold constraints. In short, our experimental results demonstrate the effectiveness of RCM in terms of both **inference efficiency and few-step generation quality**, as well as the necessity of the Riemannian constraint, which motivates our framework.\n\n## Q2 Computational Cost & Advantage of Riemannian Constraint\n\nWe first point out what might be a **misconception** that the additional Riemannian operation is \"computationally demanding\". In fact, we provide a comprehensive analysis of the overhead, both asymptotic and empirical, compared to the Euclidean CM model. Please kindly refer to Reviewer nYt4, Q5, for details. To summarize, such an overhead is **asymptotically negligible** as model size increases (as it is model-independent) and **practically small** (within 10% even for a complex manifold of SO(3) and a small model). We also provide additional scalability experiments in our rebuttal to Reviewer aiD6, Q1, in which our proposed **RCM consistently outperformed the few-step RFM and the Euclidean CM** (CD-naive) on a high-dimensional manifold (up to 256 dim). Such experimental evidence validates the need for our proposed Riemannian corrections, as Euclidean CM often fails to capture the manifold constraint, especially when the dimension is large (e.g., periodicity in our scalability experiments).\n\n## Q3 Computational Infrastructure\n\nWe have provided the detailed experimental setup in Appendix C (including the techniques we used to stabilize the training). Here, we provide additional infrastructure specifications. All experiments were carried out on a single A100. The maximum GPU memory consumed is only around 5GB, which can be easily fit into GPUs with smaller memory. As we fixed the number of iterations, the whole training of each variant of CM models took about 4 hours regardless of the dataset size.\n\nFor code, we will publicly release our code in an effort to inspire future work once our paper is published.\n\n## Q4 Theoretical Guarantees on Convergence\nWe first note that RCT formulation, as guaranteed by our Theorem 3.2, does not rely on pre-trained models and will converge to the data distribution assuming a perfectly trained model.\n\nFor RCD that does rely on a pre-trained model, the asymptotic analysis in the CM paper, Theorem 1, can be adapted for RCM under some regularity conditions. Informally speaking, such a theorem guarantees that if a perfect CD model is trained (with the pre-trained flow-based model), it can perfectly \"shortcut\" the marginal probability path. Therefore, theoretically, the errors can be fully bounded by the errors of the pre-trained flow-based model.\n\nIn practice, we observe that the CM model is generally more challenging to train than the flow matching model. See our discussion on the training instability in Q6. In this way, extending CM to Riemannian manifolds is a non-trivial task.\n\n## Q5 Impact of KDE Bandwidth\n\nWe appreciate your insightful suggestions regarding the impact of KDE bandwidth. We chose the bandwidth parameters following [1]. Additionally, we provide additional ablations on the impact of the KDE bandwidth in the table below. It can be seen that, although the choice of bandwidth did impact the numbers of KLD, their relative rankings between different models remain largely unchanged. Therefore, we believe it remains an informative metric for measuring generation quality. As a more concrete example, we run additional evaluation on the Earthquake dataset with a KDE bandwidth of 0.04 (instead of 0.02). The results are as follows. Although the specific KLD values do change a lot, the relative ordering is perfectly preserved.\n\n| Model | RFM-100 | RFM-2 | sRCD | RCD | dRCD | CD-naive | sRCT | RCT |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| KLD | 0.40 | 6.13 | **0.66** | 0.75 | 3.77 | 1.48 | 1.04 | 0.94 |\n\n[1] Mathieu, Emile, and Maximilian Nickel. \"Riemannian continuous normalizing flows.\" Advances in neural information processing systems 33 (2020): 2503-2515.\n\n## Q6 Training Instability\nIt is already known that training CMs will generally encounter more instability (e.g., in the sCM paper).\nSimilar to the original CM, we also encountered instabilities during training, both for discrete-time RCM that uses Δ𝑡 and for continuous-time RCM that relies on the model JVP. For the discrete-time RCM (referred to as dRCD in the results), we follow the original CM paper and use $\\Delta t=0.01$. We observed NaN loss sometimes for $\\Delta t=1𝑒-3$ and almost surely NaN for $\\Delta t=1𝑒-4$. The results echo with the CM and  ECT paper, where the former also uses a $\\Delta t$ around 0.01.\n\nFor the continuous-time scenario, we use multiple techniques inspired by sCM to mitigate the instability caused by the initially noisy JVP signals. The details were provided in Appendix C.2, and additional ablations are provided in our rebuttal to Reviewer nYt4, Q1.\n\n## Q7 Miscellaneous Questions\n\n- **Weight** $w_t$. Yes, we follow the CM and sCM paper in choosing the weight as $t^2/(1-t)^2$. In practice, we do not experience a significant impact from the choice of scheduler, as they perform equally well.\n- **Scheduler**. We simply follow RFM to use the linear scheduler. Despite fixing this so far, our theoretical framework allows for the flexibility of other schedulers.\n- **Comparison to other generative models (GAN, VAE)**. We interpret your question as the advantage of enforcing the manifold constraint over the pure data-driven approaches (with GAN/VAE) that ignore such a constraint. While we agree that the bounds for VAE/GAN are true, they only hold asymptotically with an infinite amount of data with perfect learning assumptions. In practice, they do not necessarily hold, as demonstrated in our comparisons with the CD-naive baseline and discussions in Q2. Therefore, RCM provides a **more data-efficient approach** in which the **manifold constraint holds by design**. Regarding the computational cost, we have clarified in Q2 that the additional Riemannian operations add negligible overhead."}, {"review_id": "kDzJjcCYPn", "reviewer": "Reviewer_J7YW", "summary": "This work proposes a new way to train few-step generative models on Riemannian manifolds, taking inspiration from and generalising existing literature on consistency models. They show a simplified training objective that allows for simpler and faster computation, and, for small ($N = 2$), equal amounts of steps, they show that their method outperforms typical Riemannian flow matching.", "questions": "1. Do you have any theoretical results in terms of the learnt distribution? It seems that the authors have not delved too deep in Probability Theory over Riemannian manifolds—which is fine, but I think it would consolidate the paper theoretically rather substantially.\n2. Have you tried any other datasets, such as those in [4]?\n3. Have you tried to run [1] on the datasets you use?\n4. To compare the different methods, it would have been particularly convenient to also obtain the NLL on the test sets. Is it possible to evaluate the NLL using your method?\n5. I see no errors on your numbers. Have you re-run your experiments to obtain confidence intervals? They should be relatively cheap, considering the dimensionality of the considered datasets.", "limitations": "It is unclear whether the method could be scaled up to higher dimensions.\n\nReferences:\n---\n\n[1] Peter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich, Ullrich Köthe (2023). Learning distributions on manifolds with free-form flows.\n\n[2] Linqi Zhou, Stefano Ermon, Jiaming Song (2025). Inductive Moment Matching\n\n[3] Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, Stefano Ermon (2023). Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion.\n\n[4] Ricky T. Q. Chen, Yaron Lipman (2023). Flow matching on General Geometries.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. The paper presents a relatively simple and intuitive method to train a few-step generative model on a Riemannian manifold, which is a relatively unexplored topic.\n2. The proposed method seems to improve significantly on its usual flow matching counterpart.\n3. The authors illustrate with understandable examples their approach in this work, namely in section 3.3, and they present differential geometry in rather clear terms.\n\nWeaknesses:\n\n1. I believe that the need for few-step generative models on Riemannian geometries could have been better motivated in the introduction.\n2. There are no theoretical guarantees relating to the learnt distribution of the model, although I do not believe it is of primary importance.\n3. I find the argument at the end of section 3.1 (“The proof above immediately…”) particularly hand-wavy, and do not get how equation 8 is a natural consequence.\n4. There is not much new conceptual content introduced in this work—which would have been fine if the experiments were more extensive.\n5. The experiments are not extensive enough, as they show only low-dimensional cases (sphere, flat torus and $\\mathrm{SO}(3)$).\n6. The experiments include no comparable baselines, such as [1] (cited in section 5, which is perhaps the most important one to include), or Euclidean few-step generators [2, 3]. It is also quite unclear which consistency method has been used as a baseline for “Euclidean consistency models”: is it the authors’ framework on an Euclidean geometry, or another method, which only works on Euclidean geometry?\n\nI believe that I would be happy to reconsider my assessment if given more empirical evidence: more datasets (in higher dimensions), more comparable baselines.", "was_revised": true, "final_justification": "The authors have demonstrated through the rebuttals the soundness and novelty of their method, which does indeed set a new high in the field of few-step generative models over Riemannian manifolds. The discussions incurred will enhance the paper to match the standards required.", "rebuttal": "We thank you for your comprehensive review and your high recognition of our proposed RCM's experimental improvement. We are happy to answer your questions as follows.\n\n## Q1 Motivation\n\nThanks for your suggestions. We will further elaborate on the motivations as follows. While diffusion or flow-based generative models that respect the intrinsic geometry have been proposed, one of their limitations is that they often **require hundreds to thousands of iterative sampling steps** to obtain decent results, which is a time-consuming process. While consistency models have been introduced to accelerate the Euclidean diffusion/flow models, **their generalization to Riemannian manifolds is non-trivial**. In this regard, we proposed RCM as a mathematically sound extension of CM to Riemannian manifolds, aiming to **enhance the inference time efficiency** of existing diffusion/flow models. In practice, such a few-step generation setup can accelerate important real-world applications, such as protein design (e.g., FrameFlow/FoldFlow, which operates on SE(3)), with significantly reduced inference time. Therefore, we hope RCM can have a profound impact on such AI4Science domains where Riemannian manifolds are encountered.\n\n## Q2 Theoretical Guarantees of Convergence\n\nThe theoretical guarantees of convergence can be established using Theorem 3.2 in our paper. We provide a brief sketch of proof as follows.\n\n- According to Theorem 3.2 in our paper, RCD and RCT are equivalent, assuming a perfectly learned pre-trained RFM model that approximates the marginal vector field.\n- According to Theorem 1 in the CM paper, a perfect CD loss indicates a perfect recovery of the marginal probability paths. Such a result can also be extended to Riemannian cases with some regularity conditions.\n- According to the RFM paper, Theorem 3.1, a perfect conditional Riemannian flow matching model can perfectly approximate the marginal vector field, thus giving the true marginal probability path.\n\nCombining the three results above, we arrive at the conclusion that, given a perfect pre-trained RFM, both RCD and RCT can recover the ground truth data distribution if perfectly optimized.\n\nDespite theoretical guarantees, the training procedure in practice cannot be perfect due to limited data, resources, and numerical issues. Still, our empirical evaluation demonstrates the effectiveness of RCM in practice in such a few-step generation setup.\n\n## Q3 Alternative Loss in Eq.8\n\nWe apologize for the confusion and provide a detailed mathematical deduction below. Specifically, for the loss value, we have \n$$\\langle f_{\\theta^-}-f_\\theta+\\dot f_{\\theta^-},\\dot f_{\\theta^-}\\rangle_g=\\\\|\\dot f_{\\theta^-}\\\\|\\_g^2.$$\nFor the loss gradient, we have\n$$\n\\nabla_\\theta\\langle f_{\\theta^-}-f_\\theta+\\dot f_{\\theta^-},\\dot f_{\\theta^-}\\rangle_g\n=\\nabla_\\theta\\langle f_{\\theta^-}-f_\\theta,\\dot f_{\\theta^-}\\rangle_g\n=\\langle \\nabla_\\theta (f_{\\theta^-}-f_\\theta),\\dot f_{\\theta^-}\\rangle_g\\\\\n=-\\langle \\nabla_\\theta f_\\theta,\\dot f_{\\theta^-}\\rangle_g\\\\\n\\approx\\frac{\\Delta t}{2}\\nabla_\\theta d_g^2(f_\\theta,\\tilde{f}_{\\theta^-})\n$$\nwhere the last approximation follows the deduction in Eq.7. This is also the technique used in the sCM paper.\n\n## Q4 Conceptual Innovation\n\nWe respectfully disagree with the reviewer's argument that our work has \"not much new conceptual content\". Our theoretical contributions and innovations on the few-step generation setup on the Riemannian manifold have been unanimously acknowledged by other reviewers. We further elaborate on our new contributions as follows.\n\n- Regarding the task, our work is the **first to systematically explore few-step generative modeling** for Riemanian diffusion/flow-based models. Although generative frameworks on Riemannian manifolds exist, they predominantly require a large number of steps to achieve decent results.\n- To the best of our knowledge, our RCM is the **first to introduce covariant derivatives and differentials of the exponential map** in the context of generative modeling. We provide a **novel kinematics perspective** in Section 3.3 that links these important mathematical concepts with efficient generative models on a Riemannian manifold. RCM not only fills this gap but also provides solid theoretical grounding and substantial empirical evidence, further supporting our claim that such Riemannian information is crucial for generating good results.\n- Our framework is **theoretically grounded**, with Theorem 3.2 guaranteeing the equivalence between RCD and RCT. Due to the complexity of Riemannian geometries, such a result is a **non-trivial extension** of the Euclidean counterparts.\n\n## Q5 Scalability to Higher Dimension\n\nWe provide additional scalability experiments on the manifold with up to 256 dimensions. Please kindly refer to our rebuttal to Reviewer aiD6, Q1. To summarize, our proposed **RCM consistently outperformed the few-step RFM and the Euclidean CM** (CD-naive). Such experimental results also necessitate the need for Riemannian constraints, as Euclidean CM often fails to capture the manifold constraint, especially when the dimension is large (e.g., periodicity in our scalability experiments).\n\n## Q6 Additional Baselines\nWe first note that our experiments have already included the CD-naive baseline that corresponds to the standard CM model for Euclidean data. Such a CM model has proven powerful and effective in few-step image generation and thus offers a strong baseline for comparison.\n\nNonetheless, we follow your suggestion to conduct additional experiments using the FFF model on our datasets. The results are summarized in the table below:\n\n| Dataset | Earthquake | Volcano | Fire | Flood | Protein | RNA | Fisher |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Metric | 2.76 KLD | 22.57 KLD | 2.02 KLD | 2.34 KLD | 4.96×10⁻² MMD  | 8.52×10⁻² MMD  | 6.67×10⁻² MMD  |\n\nIt can be seen that FFF provides decent generation results for larger datasets, but **still falls short compared to the best RCM variant** on all datasets. Interestingly, for smaller datasets like Volcano and RNA, FFF's performance is considerably worse, although it still achieves better performance compared to the few-step RFM setup. Such additional experimental results further demonstrate the effectiveness of RCM.\n\n## Q7 NLL Calculation\n\nWe fully understand that NLL is one of the important metrics for generative models. However, we noted that, due to the few-step generation setup, NLL cannot be calculated as the original RFM model. We further elaborate on the reasons as follows.\n\nWe first note that we can consider diffusion/flow models as a pushforward transformation of the prior probability distribution. The change-of-measure (change-of-probability) terms can be calculated as\n$$\\nabla\\_\\theta\\log|f'\\_\\theta(x)|=\\mathrm{tr}(\\nabla\\_\\theta f'\\_\\theta(x)f^{'-1}\\_\\theta(z)),$$\nwhere $z=f\\_\\theta(x)$ (Theorem 1 in the FFF paper). Therefore,\n\n- Normalizing flow can naturally compute this quantity as its model is designed such that the trace of the Jacobian of each layer in the model can be calculated in closed form. \n- FFF approaches this by learning an inverse model $g_\\phi$ such that $g_\\phi(f_\\theta(x))\\equiv x$. In this way, the inverse Jacobian can be approximated using the learned $g$. \n- The standard flow matching model, as a special case of continuous normalizing flow, also enjoys exact likelihood calculation by integrating the model Jacobian. Here, we note that the continuous-time is essential, as the flow is always invertible locally. In this way, the inverse is simply the negative vector field. However, for CM models that shortcut the probability path, such an invertibility no longer holds. In other words, we do not know the inverse function of $f_\\theta$ for NLL calculation. It is, though, possible to combine FFF to CM models to learn an \"inverse\" shortcut flow for NLL estimation.\n\nTherefore, we resorted to alternative metrics for a fair comparison between baselines. The KDE with KLD approach follows the Riemannian normalizing flow paper. MMD is also widely used in point cloud matching. In this regard, our existing metrics already provide comprehensive evaluations that demonstrate the effectiveness of RCM over the baselines.\n\n## Q8 Error bars\n\nWe thank you for your suggestions on adding error bars to the evaluation metrics. In the table below, we provide the new results with three repeats on the Earthquake dataset as an example. It can be seen that most original metrics fall inside the new variance range. Interestingly, all RCT models have higher variance than the CD models, probably due to the noisier training procedure with the teacher model. Still, given the current results, **our improvements over the baselines are indeed significant**. Due to the limits of time and rebuttal space, we will add other results in the revised manuscript.\n\n| Model | RFM-100 | RFM-2 | sRCD | RCD | dRCD | CD-naive | sRCT | RCT |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| KLD | 1.57 ± 0.26 | 10.57 ± 0.50 | **1.97 ± 0.06** | 2.12 ± 0.04 | 6.38 ± 0.06 | 3.14 ± 0.02 | 2.27 ± 0.40 | 2.22 ± 0.23 |"}]

Author rebuttal:

Meta-review: This paper proposes to extend consistency models to Riemannian manifolds. All reviewers agree that the method is technically sound and leads to a new class of few-step generative models for geometric domains. Despite this, the reviewers and I agree that the work in terms of novelty is incremental as it combines known ideas in consistency models and Riemannian generative models in a straightforward---but well executed---manner. This means that the theoretical results are used to substantiate the algorithm design, but by themselves don't offer profound implications on Riemannian generative models. In addition, the experimental evaluation is limited and higher-dimensional datasets are not considered (e.g. protein backbone generation)---despite the authors' implication in the rebuttal. Having said all of this, I still believe this paper fills a gap in the literature and is an appropriate contribution to the community to deserve publication at NeurIPS.
Decision: Accept (poster)