Riemannian Gaussian Variational Flow Matching

• RG-VFM is a generative modeling framework that extends Variational Flow Matching to data on Riemannian manifolds using Riemannian Gaussian distributions.
• It leverages closed-form geodesic and log/exp maps to perform manifold-aware learning, enabling efficient endpoint prediction along geodesics.
• Empirical evaluations on spherical data demonstrate that RG-VFM better preserves geometric properties and checkerboard contrast compared to Euclidean methods.

Riemannian Gaussian Variational Flow Matching (RG-VFM) is a generative modeling framework that extends the principle of Variational Flow Matching (VFM) to data residing on Riemannian manifolds. It achieves this by utilizing Riemannian Gaussian distributions as variational families and matching endpoint predictions along geodesics in spaces where the geometric structure is non-Euclidean. RG-VFM is designed for manifolds with closed-form geodesic and log/exp map computations—exemplified by spheres, hyperbolic spaces, and flat tori—allowing manifold-aware learning while maintaining computational tractability (Zaghen et al., 18 Feb 2025).

1. Riemannian Gaussian Distributions

Given a complete Riemannian manifold $(\mathcal{M}, g)$ with metric $g$, geodesic distance $\mathrm{dist}_g(\cdot,\cdot)$, and volume form $d\mathcal{M}_z$, the Riemannian Gaussian (RG) centered at $\mu \in \mathcal{M}$ with scale $\sigma > 0$ is defined as

$$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$

where the normalizing constant

$$C(\mu,\sigma) = \int_{\mathcal{M}} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right) d\mathcal{M}_z$$

ensures unit mass. In Euclidean space, this recovers the standard isotropic Gaussian. If $\mathcal{M}$ is homogeneous, $C(\mu,\sigma)$ is independent of $g$0. This property is essential for parameterizing densities and performing variational inference on structured manifolds (Zaghen et al., 18 Feb 2025).

2. RG-VFM: Variational Objective and Training

RG-VFM generalizes Euclidean VFM to manifolds by introducing joint densities $g$1, interpolating between a base $g$2 and a target $g$3. The variational flow matching objective is

$g$4

with the variational posterior

$g$5

where $g$6 is the predicted geodesic endpoint.

Assuming (a) $g$7 is homogeneous and (b) closed-form geodesics are available, the RG-VFM loss reduces (up to an additive constant, absorbing $g$8) to

$g$9

where $\mathrm{dist}_g(\cdot,\cdot)$0 denotes the Riemannian logarithm map at $\mathrm{dist}_g(\cdot,\cdot)$1 and $\mathrm{dist}_g(\cdot,\cdot)$2 is the norm induced by the metric on the tangent space (Zaghen et al., 18 Feb 2025).

RG-VFM recovers the underlying velocity field $\mathrm{dist}_g(\cdot,\cdot)$3 for the Riemannian continuity equation

$\mathrm{dist}_g(\cdot,\cdot)$4

by averaging the conditional geodesic velocities: $\mathrm{dist}_g(\cdot,\cdot)$5 connecting endpoint-based training to continuous-time flows over $\mathrm{dist}_g(\cdot,\cdot)$6.

3. Comparison with Riemannian Flow Matching (RFM)

In contrast to RG-VFM, Riemannian Flow Matching (RFM) learns a vector field directly on $\mathrm{dist}_g(\cdot,\cdot)$7 by minimizing

$\mathrm{dist}_g(\cdot,\cdot)$8

Key differences include:

• Objective: RFM matches instantaneous velocity fields at each $\mathrm{dist}_g(\cdot,\cdot)$9, whereas RG-VFM matches predicted endpoints.
• Parameterization: RG-VFM is variational: it learns the geodesic endpoint $d\mathcal{M}_z$0 (and optionally the scale $d\mathcal{M}_z$1), reconstructing the velocity field indirectly.
• Geometric requirements: Both methods utilize geodesic log maps. RG-VFM requires these only at predicted endpoints, while RFM requires them at all training points.
• Support of base distribution: RG-VFM can leverage a base $d\mathcal{M}_z$2 in $d\mathcal{M}_z$3 (yielding the RG-VFM$d\mathcal{M}_z$4 variant) with linear off-manifold interpolation. RFM requires $d\mathcal{M}_z$5 to be supported on $d\mathcal{M}_z$6 and geodesic-based interpolation throughout [(Zaghen et al., 18 Feb 2025), Chen & Lipman 2024].

4. Manifolds, Metrics, and Geodesics in Practice

RG-VFM is applicable to homogeneous manifolds with computable closed-form geodesics and log/exp maps, such as:

Manifold	Geodesic Distance	Logarithm Map
Sphere $d\mathcal{M}_z$7	$d\mathcal{M}_z$8	$d\mathcal{M}_z$9
Hyperbolic space $\mu \in \mathcal{M}$0 (Poincaré ball)	$\mu \in \mathcal{M}$1	Closed form via Möbius transforms
Flat torus $\mu \in \mathcal{M}$2	Modulo Euclidean norm	Standard log map mod $\mu \in \mathcal{M}$3

The explicit availability of exponential and logarithm maps in these geometries underpins both training and generative sampling (Zaghen et al., 18 Feb 2025).

5. Training and Sampling Procedures

Training follows a stochastic gradient approach:

• Sample $\mu \in \mathcal{M}$4.
• Sample $\mu \in \mathcal{M}$5 (the target).
• Interpolate:
  • If $\mu \in \mathcal{M}$6: $\mu \in \mathcal{M}$7 via uniform geodesic interpolation, $\mu \in \mathcal{M}$8, with $\mu \in \mathcal{M}$9.
  • If $\sigma > 0$0: $\sigma > 0$1.
• Compute Loss: $\sigma > 0$2.
• Update: $\sigma > 0$3.

Sampling from the learned model proceeds by solving the ODE: $\sigma > 0$4 with any suitable ODE solver (e.g., RK4), yielding samples $\sigma > 0$5 (Zaghen et al., 18 Feb 2025).

6. Empirical Evaluation: Spherical Checkerboard

In the checkerboard experiment, data are embedded on $\sigma > 0$6 via: $\sigma > 0$7 The following methods are compared:

• CFM (Conditional Flow Matching)
• VFM (Euclidean)
• RFM (Riemannian)
• RG-VFM$\sigma > 0$8
• RG-VFM$\sigma > 0$9

Evaluation metrics:

• Norm deviation: $$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$0 deviation from 1 (ideal on $$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$1).
• Visual quality: Checker patch concentration.

Key results:

• Euclidean methods yield points off $$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$2 (standard deviation of $$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$3).
• Geometric methods maintain $$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$4 (std $$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$5).
• Variational methods (VFM and RG-VFM) better preserve checker contrast than CFM/RFM.
• RG-VFM$$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$6 and RG-VFM$$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$7 achieve the closest qualitative match to the geometric checkerboard target (Zaghen et al., 18 Feb 2025).

7. Limitations and Open Problems

Several constraints shape the current scope and challenge potential extensions:

• Applicability is restricted to homogeneous manifolds with closed-form log/exp maps; general learned manifolds are excluded.
• Constant normalizer and fixed scale parameter $$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$8 are assumed; extending to nonhomogeneous or anisotropic Gaussians is nontrivial.
• Off-manifold sampling in RG-VFM$$\mathcal{N}_{\rm Riem}(z\mid \mu,\sigma) = \frac{1}{C(\mu,\sigma)} \exp\!\left(-\tfrac{1}{2\sigma^2}\,\mathrm{dist}_g^2(z,\mu)\right)$$9 via Euclidean linear interpolation may behave unpredictably on high-curvature manifolds.
• Extension to manifolds with boundaries, singularities, or non-global geodesics is an open question.
• Non-Gaussian variational families (e.g., von Mises–Fisher on spheres) may further improve modeling on compact manifolds.

Further developments may involve addressing these limitations, especially integration with flexible non-Gaussian families and arbitrary manifold topologies (Zaghen et al., 18 Feb 2025).