Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Riemannian Consistency Model

Chaoran Cheng, Yusong Wang, Yuxin Chen, Xiangxin Zhou, Nanning Zheng, Ge Liu

TL;DR

The paper introduces the Riemannian Consistency Model (RCM) to enable few-step generation on non-Euclidean manifolds by enforcing intrinsic geometry through an exponential-map parameterization and the covariant derivative. It derives closed-form discrete- and continuous-time objectives, proves the equivalence between Riemannian Consistency Distillation (RCD) and Riemannian Consistency Training (RCT) via marginalization, and presents a simplified objective (sRCM) that avoids differential maps while preserving performance. A kinematics-based interpretation provides intuition for geometry-aware motion along PF-ODEs on manifolds, with concrete instantiations on Euclidean space, spheres, and SO(3). Empirical results across flat tori, spheres, and SO(3) demonstrate improved few-step generation quality and scalability compared to naïve Euclidean adaptations and prior Riemannian flow methods, highlighting practical potential for fast manifold-aware generative modeling in domains like protein design and robotics.

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

Riemannian Consistency Model

TL;DR

The paper introduces the Riemannian Consistency Model (RCM) to enable few-step generation on non-Euclidean manifolds by enforcing intrinsic geometry through an exponential-map parameterization and the covariant derivative. It derives closed-form discrete- and continuous-time objectives, proves the equivalence between Riemannian Consistency Distillation (RCD) and Riemannian Consistency Training (RCT) via marginalization, and presents a simplified objective (sRCM) that avoids differential maps while preserving performance. A kinematics-based interpretation provides intuition for geometry-aware motion along PF-ODEs on manifolds, with concrete instantiations on Euclidean space, spheres, and SO(3). Empirical results across flat tori, spheres, and SO(3) demonstrate improved few-step generation quality and scalability compared to naïve Euclidean adaptations and prior Riemannian flow methods, highlighting practical potential for fast manifold-aware generative modeling in domains like protein design and robotics.

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

Paper Structure

This paper contains 43 sections, 5 theorems, 46 equations, 6 figures, 5 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Riemannian Flow Matching
  4. Consistency Model
  5. Riemannian Consistency Model
  6. Consistency Model on Riemannian Manifolds
  7. Simplified Riemannian Consistency Model
  8. Kinematics Perspective of Riemannian Consistency Model
  9. Euclidean Space.
  10. Spherical Manifold.
  11. 3D Rotation Group.
  12. Experiments
  13. Spherical Manifold
  14. Flat Torus
  15. 3D Rotation Group
  16. ...and 28 more sections

Key Result

Theorem 3.1

When $N\to\infty,\Delta t\to0$, the continuous-time RCM loss is where $\mathrm{d}(\exp_x)_u,\mathrm{d}(\exp u)_{x}:T_x\mathcal{M}\to T_{f(x)}\mathcal{M}$ are the differentials of the exponential map with respect to the tangent vector $u=\kappa v$ and the base point $x$. $\nabla_{\dot x}$ denotes the covariant derivative with respect to the Levi-Civita connection

Figures (6)

  • Figure 1: Riemannian Consistency Model (RCM). The denoiser $f_\theta$ along the Riemannian PF-ODE should be consistent.
  • Figure 2: Three components of the variations in the consistency objective.
  • 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.
  • Figure 4: KLD vs NFEs for different models on the Earthquake dataset. RCD variants are colored in green, and RCT variants are colored in purple.
  • Figure 5: Generations on the 2D flat torus. The FM, CD, and CT models are colored in blue, green, and orange.
  • ...and 1 more figures

Theorems & Definitions (9)

  • Theorem 3.1
  • proof
  • Theorem 3.2
  • Lemma 3.1
  • proof
  • Lemma 3.2
  • proof : Proof for Theorem \ref{['thm:rct']}
  • Proposition 3.1
  • proof