---
title: Riemannian Gaussian Variational Flow Matching
url: https://www.emergentmind.com/topics/riemannian-gaussian-variational-flow-matching-rg-vfm
type: topic
---

# Riemannian Gaussian Variational Flow Matching

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

## 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 $\mu$. This property is essential for parameterizing densities and performing variational inference on structured manifolds [2502.12981].

## 2. RG-VFM: Variational Objective and Training

RG-VFM generalizes Euclidean VFM to manifolds by introducing joint densities $p_t(x_1, x), 0 \leq t \leq 1$, interpolating between a base $p_0$ and a target $p_1$. The variational flow matching objective is
\[
\mathbb{E}_{t}\left[\mathrm{KL}\left(p_t(x_1,x)\,\|\,q_t^\theta(x_1,x)\right)\right]
= -\,\mathbb{E}_{t, x_1, x}\bigl[\log q_t^\theta(x_1\mid x)\bigr] + \text{const}
\]
with the variational posterior
\[
q_t^\theta(x_1\mid x) = \mathcal{N}_{\rm Riem}\left(x_1\mid \mu_t^\theta(x),\,\sigma_t(x)\right)
\]
where $\mu_t^\theta(x) \in \mathcal{M}$ is the predicted geodesic endpoint.

Assuming (a) $\mathcal{M}$ is homogeneous and (b) closed-form geodesics are available, the RG-VFM loss reduces (up to an additive constant, absorbing $1/\sigma^2$) to
\[
\mathcal{L}_{\rm RG\text{-}VFM}(\theta)
=
\mathbb{E}_{t, x_1, x} \left\|\log_{x_1}\left(\mu_t^\theta(x)\right)\right\|_g^{2}
\]
where $\log_{x_1}$ denotes the Riemannian logarithm map at $x_1$ and $\| \cdot \|_g$ is the norm induced by the metric on the tangent space [2502.12981].

RG-VFM recovers the underlying velocity field $u_t$ for the Riemannian continuity equation
\[
\partial_t\,p_t(z) + \mathrm{div}_g\left(p_t(z)\,u_t(z)\right) = 0
\]
by averaging the conditional geodesic velocities:
\[
u_t^\theta(x) = \mathbb{E}_{x_1 \sim q_t^\theta(\cdot\mid x)} \left[ \frac{\log_x(x_1)}{1-t} \right]
\]
connecting endpoint-based training to continuous-time flows over $\mathcal{M}$.

## 3. Comparison with Riemannian Flow Matching (RFM)

In contrast to RG-VFM, Riemannian Flow Matching (RFM) learns a vector field directly on $\mathcal{M}$ by minimizing
\[
\mathcal{L}_{\rm RFM}(\theta) =
\mathbb{E}_{t, x_1, x}
\left\| v_t^\theta(x) - \frac{\log_x(x_1)}{1-t} \right\|_g^2
\]
Key differences include:

- **Objective:** RFM matches instantaneous velocity fields at each $x$, whereas RG-VFM matches predicted endpoints.
- **Parameterization:** RG-VFM is variational: it learns the geodesic endpoint $\mu_t^\theta(x)$ (and optionally the scale $\sigma_t(x)$), 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 $p_0$ in $\mathbb{R}^D \supset \mathcal{M}$ (yielding the RG-VFM$_{\mathbb{R}^3}$ variant) with linear off-manifold interpolation. RFM requires $p_0$ to be supported on $\mathcal{M}$ and geodesic-based interpolation throughout [2502.12981, 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 $S^d \subset \mathbb{R}^{d+1}$ | $\operatorname{arccos}(\langle x, y \rangle)$ | $\log_x(y) = \frac{\theta}{\sin\theta}\,(y - \cos\theta\,x),\ \theta = \arccos\langle x, y\rangle$ |
| Hyperbolic space $\mathbb{H}^d$ (Poincaré ball) | $\operatorname{arccosh}\left(1+2\,\frac{\|x-y\|^2}{(1-\|x\|^2)(1-\|y\|^2)}\right)$ | Closed form via Möbius transforms |
| Flat torus $T^d$          | Modulo Euclidean norm         | Standard log map mod $2\pi$               |

The explicit availability of exponential and logarithm maps in these geometries underpins both training and generative sampling [2502.12981].

## 5. Training and Sampling Procedures

Training follows a stochastic gradient approach:

- **Sample** $t \sim \operatorname{Uniform}[0,1]$.
- **Sample** $x_1 \sim p_1$ (the target).
- **Interpolate**:
    - If $p_0 \subset \mathcal{M}$: $x$ via uniform geodesic interpolation, $x \leftarrow \mathrm{Exp}_{x_1}((t – 1)\,\log_{x_1}(x_0))$, with $x_0 \sim p_0$.
    - If $p_0 \subset \mathbb{R}^D$: $x = (1-t)x_0 + t x_1$.
- **Compute Loss**: $\mathcal{L} = \|\log_{x_1}(\mu_t^\theta(x))\|_g^2$.
- **Update**: $\theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L}$.

Sampling from the learned model proceeds by solving the ODE:
\[
\frac{dx}{dt} = v_t^\theta\left(x(t)\right) =
\begin{cases}
\frac{1}{1-t}\,\log_{x(t)}\left(\mu_t^\theta(x(t))\right), & x(t)\in \mathcal{M} \\
\frac{\mu_t^\theta(x(t)) - x(t)}{1-t}, & x(t)\in \mathbb{R}^D
\end{cases}
\]
with any suitable ODE solver (e.g., RK4), yielding samples $x(1) \sim p_1$ [2502.12981].

## 6. Empirical Evaluation: Spherical Checkerboard

In the checkerboard experiment, data are embedded on $S^2$ via:
\[
p_1(\theta, \phi) \propto 1[\lfloor k\theta/\pi \rfloor + \lfloor k\phi/2\pi \rfloor\ \text{is even}]\cdot\sin\theta
\]
The following methods are compared:

- CFM (Conditional Flow Matching)
- VFM (Euclidean)
- RFM (Riemannian)
- RG-VFM$_{\mathbb{R}^3}$
- RG-VFM$_{\mathcal{M}}$

Evaluation metrics:

- **Norm deviation**: $\|x\|_2$ deviation from 1 (ideal on $S^2$).
- **Visual quality**: Checker patch concentration.

Key results:

- Euclidean methods yield points off $S^2$ (standard deviation of $\|x\| \approx 0.09$).
- Geometric methods maintain $\|x\| \approx 1$ (std $<0.01$).
- Variational methods (VFM and RG-VFM) better preserve checker contrast than CFM/RFM.
- RG-VFM$_{\mathcal{M}}$ and RG-VFM$_{\mathbb{R}^3}$ achieve the closest qualitative match to the geometric checkerboard target [2502.12981].

## 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 $\sigma$ are assumed; extending to nonhomogeneous or anisotropic Gaussians is nontrivial.
- Off-manifold sampling in RG-VFM$_{\mathbb{R}^3}$ 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 [2502.12981].

Source: https://www.emergentmind.com/topics/riemannian-gaussian-variational-flow-matching-rg-vfm