Riemannian Consistency Model

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.

Q1 Motivation

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

Q2 Theoretical Guarantees of Convergence

The theoretical guarantees of convergence can be established using Theorem 3.2 in our paper. We provide a brief sketch of proof as follows.

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

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

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

Q3 Alternative Loss in Eq.8

We apologize for the confusion and provide a detailed mathematical deduction below. Specifically, for the loss value, we have $\langle f_{\theta^-}-f_\theta+\dot f_{\theta^-},\dot f_{\theta^-}\rangle_g=\\|\dot f_{\theta^-}\\|\_g^2.$ For the loss gradient, we have

$$\nabla_\theta\langle f_{\theta^-}-f_\theta+\dot f_{\theta^-},\dot f_{\theta^-}\rangle_g =\nabla_\theta\langle f_{\theta^-}-f_\theta,\dot f_{\theta^-}\rangle_g =\langle \nabla_\theta (f_{\theta^-}-f_\theta),\dot f_{\theta^-}\rangle_g\\ =-\langle \nabla_\theta f_\theta,\dot f_{\theta^-}\rangle_g\\ \approx\frac{\Delta t}{2}\nabla_\theta d_g^2(f_\theta,\tilde{f}_{\theta^-})$$

where the last approximation follows the deduction in Eq.7. This is also the technique used in the sCM paper.

Q4 Conceptual Innovation

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

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

Q5 Scalability to Higher Dimension

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

Q6 Additional Baselines

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

Nonetheless, we follow your suggestion to conduct additional experiments using the FFF model on our datasets. The results are summarized in the table below:

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

Q7 NLL Calculation

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

We 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 $\nabla\_\theta\log|f'\_\theta(x)|=\mathrm{tr}(\nabla\_\theta f'\_\theta(x)f^{'-1}\_\theta(z)),$ where $z=f\_\theta(x)$ (Theorem 1 in the FFF paper). Therefore,

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

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

Q8 Error bars

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

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.

Q1 Empirical Improvement over Existing Work

We first thank your recognition of our theoretical contributions and applications. Here, we provide a more comprehensive discussion on the empirical improvements.

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

Q2 Computational Cost & Advantage of Riemannian Constraint

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

Q3 Computational Infrastructure

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

For code, we will publicly release our code in an effort to inspire future work once our paper is published.

Q4 Theoretical Guarantees on Convergence

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

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

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

Q5 Impact of KDE Bandwidth

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

[1] Mathieu, Emile, and Maximilian Nickel. "Riemannian continuous normalizing flows." Advances in neural information processing systems 33 (2020): 2503-2515.

Q6 Training Instability

It is already known that training CMs will generally encounter more instability (e.g., in the sCM paper). Similar 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.

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

Q7 Miscellaneous Questions

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

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.

Q1. Tangent Warmup

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

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

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

Q2 Practical Applications to Different Manifolds

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

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

Q3 Quality vs NFE

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

[2] Kim, Dongjun, et al. "Consistency trajectory models: Learning probability flow ode trajectory of diffusion." arXiv preprint arXiv:2310.02279 (2023).

Q4 Scalability

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

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

Q5 Runtime Complexity Analysis

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

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

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

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.

Q1 Application to High-Dimensional Manifolds

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

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

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

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

[1] De Bortoli, Valentin, et al. "Riemannian score-based generative modelling." Advances in neural information processing systems 35 (2022): 2406-2422.

[2] Mathieu, Emile, and Maximilian Nickel. "Riemannian continuous normalizing flows." Advances in neural information processing systems 33 (2020): 2503-2515.

Q2 Baselines Question

Comprehensive 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:

• 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.
• 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.
• The Riemannian diffusion and score models fall short of the RFM model (in the RFM paper results).

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

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

Q3 Performance vs NFEs

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

We have the following interesting observations:

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

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

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

[3] Kim, Dongjun, et al. "Consistency trajectory models: Learning probability flow ode trajectory of diffusion." arXiv preprint arXiv:2310.02279 (2023).

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.

• Scalability to higher-dimension manifolds with a dimension up to 256, in Reviewer aiD6, Q1. RCM consistently achieves better generation quality.
• NFEs vs performance evaluation, in Reviewer aiD6, Q3. RCM has a stable performance compared to RFM.
• Additional FFF baseline results (a Riemannian normalizing flow), in Reviewer J7YW, Q6. RCM consistently outperforms FFF.

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