Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Riemannian Variational Flow Matching for Material and Protein Design

Olga Zaghen, Floor Eijkelboom, Alison Pouplin, Cong Liu, Max Welling, Jan-Willem van de Meent, Erik J. Bekkers

TL;DR

This work introduces Riemannian Gaussian Variational Flow Matching (RG-VFM), a geometric extension of variational flow matching for data on manifolds that uses a Riemannian Gaussian posterior to model endpoint distributions. By analyzing Jacobi fields, the authors show that RG-VFM captures curvature information via the endpoint distance $\|\log_{x_1}(\mu_t^{\theta}(x))\|^2$ and differs from Riemannian Flow Matching (RFM) by a curvature-dependent term, with Euclidean space recovering the standard VFM/CFM equivalence. The paper proves that RG-VFM reduces to the Fréchet mean under suitable assumptions and demonstrates, through synthetic curved spaces and real-world MOF and protein design tasks, that endpoint-based variational training yields sharper distributions and better geometry alignment than velocity-based methods. Empirical results on MOFFlow and protein backbone generation show consistent improvements in structure prediction, designability, and RMSD metrics, validating the practical value of incorporating curvature-aware variational objectives into manifold-valued generative modeling. Overall, RG-VFM offers a principled, geometry-respecting, and computationally efficient approach for learning probability paths on complex geometries with broad applicability in materials and biomolecular design.

Abstract

We present Riemannian Gaussian Variational Flow Matching (RG-VFM), a geometric extension of Variational Flow Matching (VFM) for generative modeling on manifolds. In Euclidean space, predicting endpoints (VFM), velocities (FM), or noise (diffusion) are largely equivalent due to affine interpolations. On curved manifolds this equivalence breaks down, and we hypothesize that endpoint prediction provides a stronger learning signal by directly minimizing geodesic distances. Building on this insight, we derive a variational flow matching objective based on Riemannian Gaussian distributions, applicable to manifolds with closed-form geodesics. We formally analyze its relationship to Riemannian Flow Matching (RFM), exposing that the RFM objective lacks a curvature-dependent penalty - encoded via Jacobi fields - that is naturally present in RG-VFM. Experiments on synthetic spherical and hyperbolic benchmarks, as well as real-world tasks in material and protein generation, demonstrate that RG-VFM more effectively captures manifold structure and improves downstream performance over Euclidean and velocity-based baselines.

Riemannian Variational Flow Matching for Material and Protein Design

TL;DR

This work introduces Riemannian Gaussian Variational Flow Matching (RG-VFM), a geometric extension of variational flow matching for data on manifolds that uses a Riemannian Gaussian posterior to model endpoint distributions. By analyzing Jacobi fields, the authors show that RG-VFM captures curvature information via the endpoint distance and differs from Riemannian Flow Matching (RFM) by a curvature-dependent term, with Euclidean space recovering the standard VFM/CFM equivalence. The paper proves that RG-VFM reduces to the Fréchet mean under suitable assumptions and demonstrates, through synthetic curved spaces and real-world MOF and protein design tasks, that endpoint-based variational training yields sharper distributions and better geometry alignment than velocity-based methods. Empirical results on MOFFlow and protein backbone generation show consistent improvements in structure prediction, designability, and RMSD metrics, validating the practical value of incorporating curvature-aware variational objectives into manifold-valued generative modeling. Overall, RG-VFM offers a principled, geometry-respecting, and computationally efficient approach for learning probability paths on complex geometries with broad applicability in materials and biomolecular design.

Abstract

We present Riemannian Gaussian Variational Flow Matching (RG-VFM), a geometric extension of Variational Flow Matching (VFM) for generative modeling on manifolds. In Euclidean space, predicting endpoints (VFM), velocities (FM), or noise (diffusion) are largely equivalent due to affine interpolations. On curved manifolds this equivalence breaks down, and we hypothesize that endpoint prediction provides a stronger learning signal by directly minimizing geodesic distances. Building on this insight, we derive a variational flow matching objective based on Riemannian Gaussian distributions, applicable to manifolds with closed-form geodesics. We formally analyze its relationship to Riemannian Flow Matching (RFM), exposing that the RFM objective lacks a curvature-dependent penalty - encoded via Jacobi fields - that is naturally present in RG-VFM. Experiments on synthetic spherical and hyperbolic benchmarks, as well as real-world tasks in material and protein generation, demonstrate that RG-VFM more effectively captures manifold structure and improves downstream performance over Euclidean and velocity-based baselines.

Paper Structure

This paper contains 65 sections, 10 theorems, 70 equations, 17 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Flow Matching.
  4. Riemannian Flow Matching.
  5. Variational Flow Matching.
  6. Riemannian Gaussian Variational Flow Matching
  7. The Riemannian Gaussian VFM Objective
  8. RG-VFM vs RFM: a comparison based on Jacobi fields
  9. Jacobi Field formulation of the Flow Matching Objectives
  10. Relation between RG-VFM and RFM Objectives
  11. Euclidean case.
  12. Geometric interpretation.
  13. Experiments
  14. Goal of the experiments.
  15. Motivation for material and protein generation and common pattern.
  16. ...and 50 more sections

Key Result

Proposition 3.1

Let $\mathcal{M}$ be a homogeneous manifold with closed-form geodesics. Then, the RG-VFM objective reduces to where $\log$ denotes the logarithmic map on the manifold and $\mathrm{dist}_{\mathbf{g}}$ is the geodesic distance.

Figures (17)

  • Figure 1: Overview of the models relevant for our framework. The square represents Euclidean space, while the sphere represents a general $\mathcal{M}$.
  • Figure 2: a: Representation of a shooting family of geodesics on $\mathbb{S}^2$ with corresponding Jacobi field. b: Visualization of predictions (color-coded to match the name colors) of each model relevant to our framework, for a target distribution $p_1$ supported on the sphere $\mathbb{S}^2$.
  • Figure 3: Comparison of the spherical checkerboard distribution generated with CFM, VFM, RFM and our methods RG-VFM-$\mathbb{R}^3$ and RG-VFM-$\mathcal{M}$.
  • Figure 4: Flow trajectories of 10,000 samples, initially drawn from the noisy distribution $p_0$ at $t=0$, evolving to reach their final configuration by $t=1$. In all variational cases, the posterior distribution is Normal, and ${p_1}$ is the checkerboard distribution on $\mathbb{S}^2$.
  • Figure 5: Flow trajectories of 10,000 samples, initially drawn from the noisy distribution $p_0$ at $t=0$, evolving to reach their final configuration by $t=1$. In all variational cases, the posterior distribution is Laplace, and ${p_1}$ is the checkerboard distribution on $\mathbb{S}^2$.
  • ...and 12 more figures

Theorems & Definitions (21)

  • Proposition 3.1
  • Definition 4.1: Jacobi field at a vanishing starting point
  • Proposition 4.1: Loss functions as evaluation of Jacobi fields
  • Proposition 4.2
  • Proposition 4.3: Difference of loss functions as a curvature term
  • Definition C.1: Normal law pennec2006intrinsic
  • Definition C.2: Riemannian Gaussian
  • Proposition D.1
  • proof
  • Proposition D.1
  • ...and 11 more