Riemannian Consistency Model
About
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).
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Generation | Torus Protein General (test) | MMD0.01 | 7 | |
| Generation | Torus Protein Glycine (test) | MMD0.04 | 7 | |
| Generation | Torus Protein Proline (test) | MMD0.05 | 7 | |
| Generation | Torus Protein PrePro (test) | MMD0.06 | 7 | |
| Generation | Torus RNA (test) | MMD0.11 | 7 | |
| One-step generation | Earth Volcano (test) | MMD0.155 | 7 | |
| One-step generation | Earth Earthquake (test) | MMD0.053 | 7 | |
| One-step generation | Earth Flood (test) | MMD0.086 | 7 | |
| One-step generation | Earth Fire (test) | MMD0.08 | 7 | |
| One-step manifold generation | SO(3) Fisher (test) | MMD0.071 | 5 |