Riemannian Conditional Flow Matching

• RCFM extends Conditional Flow Matching to non-Euclidean settings by employing Riemannian metrics to generate geodesic interpolants that respect manifold constraints.
• It uses a dual-stage loss—kinetic energy minimization followed by vector field regression—to reduce uncertainty and produce robust generative models.
• RCFM achieves state-of-the-art performance in domains like robotic control, neuroimaging, and material design while ensuring computational efficiency.

Riemannian Conditional Flow Matching (RCFM) is an advanced generative modeling framework that extends Conditional Flow Matching (CFM) to non-Euclidean geometries, leveraging data-adaptive Riemannian metrics to construct transport paths that more accurately reflect the underlying manifold structure of realistic data. RCFM enables learned flows that respect manifold constraints, intrinsically reducing uncertainty and yielding high-fidelity interpolations. It has been developed and applied in diverse contexts, such as data interpolation, robotic visuomotor policies, material generation, and neuroimaging, by several research groups (Kapuśniak et al., 2024, Ding et al., 2024, Collas et al., 20 May 2025, Miller et al., 2024, Chen et al., 2023).

1. Formalism: From Euclidean CFM to Riemannian CFM

Standard CFM chooses straight-line interpolants in ℝᵈ between source and target densities (such as $x_t = (1-t)x_0 + t x_1$ on Euclidean space), minimizing the mean squared difference between the predicted time-dependent vector field $v_{t,\theta}$ and the analytical velocity of the straight path. The loss is

$$\mathcal{L}_{\mathrm{CFM}}(\theta) = \mathbb{E}_{t,(x_0,x_1)\sim q}\left[\|v_{t,\theta}(x_t) - \dot{x}_t\|^2\right]$$

where $q(x_0,x_1)$ is a coupling (e.g., independent or OT).

RCFM generalizes this by equipping the ambient space or data manifold $\mathcal{M}$ with a data-induced Riemannian metric $g(x)$, so that conditional paths become manifold geodesics (or their approximations), and the loss is adapted to the Riemannian norm. This construction is critical for modeling data with intrinsic geometric constraints (such as robot states in $\mathbb{R}^3 \times S^3$, symmetric positive-definite matrices, crystalline structures, or cell trajectories), where straight-line segments typically lie off the data manifold and yield geometrically implausible interpolations (Chen et al., 2023, Kapuśniak et al., 2024, Collas et al., 20 May 2025, Miller et al., 2024).

2. Construction of Data-Induced Riemannian Metrics and Geodesics

A Riemannian metric takes the form $G(x) \in \mathrm{SPD}(d)$, assigning each point $x \in \mathbb{R}^d$ a positive-definite matrix, endowing the space with a position-dependent (possibly highly curved) inner product. Concrete metric designs include the LAND metric:

$$G_\epsilon(x) = [\mathrm{diag}(h(x)) + \epsilon I]^{-1}$$

where $$h(x)_\alpha = \sum_i (x_i^\alpha - x^\alpha)^2 \exp(-\|x - x_i\|^2 / 2\sigma^2)$$, and the RBF metric, a learned diagonal variant tuned to saturate on the data cloud (Kapuśniak et al., 2024). The geodesic $\gamma^*$ connecting $x_0$ and $x_1$ on $(\mathbb{R}^d, g)$ is the minimizer of the kinetic energy functional

$$E[\gamma] = \frac{1}{2} \int_0^1 \dot{\gamma}(t)^\top G(\gamma(t)) \dot{\gamma}(t) dt$$

which yields paths curving toward high-density regions.

In matrix manifold settings, such as covariance and correlation matrices, global diffeomorphisms (e.g., matrix logarithm, normalized Cholesky factorization) are used to define pullback metrics, allowing the computational machinery of Euclidean CFM to be transferred to manifold-constrained domains (Collas et al., 20 May 2025).

3. Riemannian Conditional Flow Matching Objective

RCFM introduces a dual-stage loss:

1. Kinetic energy pre-training with a neural interpolant $\phi_{t,\eta}(x_0,x_1)$, learns geodesic-like conditional paths:

$$x_{t,\eta} = (1-t)x_0 + t x_1 + t(1-t)\phi_{t,\eta}(x_0,x_1)$$

$\eta$ is fit by minimizing the expected kinetic energy under $G$:

$$\mathcal{L}_g(\eta) = \mathbb{E}_{t,(x_0,x_1)\sim q}\left[\dot{x}_{t,\eta}^\top G(x_{t,\eta}) \dot{x}_{t,\eta}\right]$$

2. Vector field regression: With the optimized $\eta^*$, RCFM minimizes the Riemannian conditional flow matching loss:

$$\mathcal{L}_{\mathrm{RCFM}}(\theta) = \mathbb{E}_{t,(x_0,x_1)\sim q} \Bigl[ \| v_{t,\theta}(x_{t,\eta^*}) - \dot{x}_{t,\eta^*} \|_{g(x_{t,\eta^*})}^2 \Bigr ]$$

and in manifold contexts, the analytic conditional velocity $u_t(x_t|x_1)$ is given by the tangent of the manifold geodesic $\psi_t(x_0|x_1)$, with

$u_t(x_t|x_1) = \frac{d}{dt} \psi_t(x_0|x_1)$

where $\psi_t$ typically interpolates between $x_0$ and $x_1$ using exponential/logarithm maps (Ding et al., 2024, Collas et al., 20 May 2025, Chen et al., 2023). The loss remains simulation-free so long as geodesic evaluation is tractable or can be approximated.

4. Algorithmic Aspects and Practical Implementations

Algorithmic implementation in RCFM proceeds via:

• Sampling $(x_0,x_1)$ from the chosen coupling $q$, and $t\sim \mathrm{Unif}[0,1]$
• Computing interpolant positions and velocities (either via analytic forms or neural approximations)
• Pretraining $\eta$ by gradient descent on $\mathcal{L}_g$, fitting the interpolant to kinetic energy geodesics
• Freezing $\eta^*$ and regressing $v_{t,\theta}$ using Riemannian vector-norm loss
• Efficient inference by Euler or Runge–Kutta integration of $x' = v_{t,\theta}(x)$ with respect to the Riemannian metric, preserving manifold constraints by construction

When the metric is diagonal (as in LAND/RBF), the norm simplifies to a dimension-wise scaling, $$\|u\|_g^2 = \sum_\alpha u_\alpha^2 / (h_\alpha(x)+\epsilon)$$. In matrix domains, the diffeomorphism trick reduces the full manifold operation to standard CFM in Euclidean space with a change of variables (Collas et al., 20 May 2025).

Architectures for $v_{t,\theta}$ typically employ MLPs for tabular or matrix data, 1D-UNets (with FiLM conditioning) for trajectory prediction in visually-conditioned settings, and E(3)-inspired graph networks for crystalline material generation, ensuring symmetry invariance and geometric fidelity (Miller et al., 2024, Ding et al., 2024).

5. Theoretical and Empirical Advantages of Riemannian Structure

RCFM yields several key benefits over Euclidean flow matching:

• Lower uncertainty: Matching and evaluating the flow on paths that hug the data manifold reduces diffusive uncertainty and ensures that the loss penalizes discrepancies only where data is observed, enhancing robustness and sample quality (Kapuśniak et al., 2024).
• Meaningful interpolations: The learned interpolants are approximate geodesics, yielding synthetic trajectories or samples that naturally respect the non-linear geometry present in empirical data, producing semantically and physically plausible outputs (Chen et al., 2023, Miller et al., 2024).
• Consistency between path geometry and vector field regression: The two-stage process (kinetic energy minimization followed by vector field regression) closely mirrors the variational principles underlying geodesic flows, so learned flows align with manifold dynamics.
• Computational efficiency: The simulation-free nature of RCFM enables fast training and inference, often with an order of magnitude fewer integration steps to reach top performance, compared to score-based diffusion models (Miller et al., 2024, Ding et al., 2024, Collas et al., 20 May 2025).

6. Domains of Application and Empirical Validation

RCFM and related algorithms have achieved state-of-the-art results across multiple domains:

Domain	RCFM Variant / Paper	Notable Results (as reported)
Synthetic trajectory inference	OT-MFM (Kapuśniak et al., 2024)	Earth-Mover distance at $t=1/2$: 0.081±0.009 (OT-MFM) vs 0.608±0.023 (OT-CFM)
LiDAR surface navigation (3D)	OT-MFM (Kapuśniak et al., 2024)	Paths track mountain surface; Euclidean CFM cuts through air
Unpaired image translation	OT-MFM (Kapuśniak et al., 2024)	AFHQ cats↔dogs, FID: 37.9 (OT-MFM) vs 41.4 (OT-CFM), LPIPS: 0.50 vs 0.51
Single-cell trajectory (scRNA-seq)	OT-MFM (Kapuśniak et al., 2024)	5D PCA: Wasserstein-1: 0.724±0.070 (OT-MFM) vs 0.882±0.058 (OT-CFM)
Robotic visuomotor control	RFMP/SRFMP (Ding et al., 2024)	Inference speed (NFE=3): 0.021s (RFMP); success rate 92.3% (RFMP/NFE=3)
Brain connectivity (fMRI/EEG)	DiffeoCFM (Collas et al., 20 May 2025)	α-precision/recall/F1, ROC-AUC surpasses baselines, valid SPD/corr matrices
Crystalline material generation	FlowMM (Miller et al., 2024)	CSP match-rate: 61.4% (FlowMM); stable structure search ∼3× faster than SOTA

RCFM’s geometry-aware flows outperform Euclidean and score-based baselines in domains that require manifold-respecting generative models, from cell trajectories and vision-to-action robotics to matrix-valued neuroimaging and atomic crystal synthesis.

7. Extensions, Equivalences, and Stability

Variants of RCFM incorporate stability principles (LaSalle invariance, Lyapunov functions) to further enhance robustness, ensuring flows remain in the support of target distributions for $t \geq 1$ (Ding et al., 2024). In matrix manifold settings, RCFM under a pullback metric is provably equivalent to standard CFM in Euclidean coordinates after suitable transformation, yielding substantial computational benefits (Collas et al., 20 May 2025). On arbitrary geometries or meshes, spectral premetrics can be employed to efficiently construct analytic target fields without numerical ODE backpropagation (Chen et al., 2023).

A plausible implication is that continued methodological advances in metric learning, manifold embedding, and conditional path construction will further broaden the reach of RCFM, enabling high-fidelity generative modeling wherever data reside on nonlinear spaces.