Riemannian Variational Flow Matching for Material and Protein Design

TL;DR

This paper introduces Riemannian Gaussian Variational Flow Matching (RG-VFM), a geometric extension of VFM for generative modeling on curved manifolds, improving material and protein design tasks by respecting manifold structures.

Contribution

The paper develops RG-VFM, incorporating Riemannian geometry into flow matching, and demonstrates its advantages over Euclidean methods in manifold-aware generative modeling.

Findings

RG-VFM outperforms Euclidean baselines on synthetic and real-world benchmarks.

Incorporating curvature via Jacobi fields improves geodesic distance approximation.

Endpoint prediction enhances learning signals on curved manifolds.

Abstract

We present Riemannian Gaussian Variational Flow Matching (RG-VFM), a geometric extension of Variational Flow Matching (VFM) for generative modeling on manifolds. Motivated by the benefits of VFM, we derive a variational flow matching objective for manifolds with closed-form geodesics based on Riemannian Gaussian distributions. Crucially, in Euclidean space, predicting endpoints (VFM), velocities (FM), or noise (diffusion) is largely equivalent due to affine interpolations. However, on curved manifolds this equivalence breaks down. We formally analyze the relationship between our model and Riemannian Flow Matching (RFM), revealing that the RFM objective lacks a curvature-dependent penalty -- encoded via Jacobi fields -- that is naturally present in RG-VFM. Based on this relationship, we hypothesize that endpoint prediction provides a stronger learning signal by directly minimizing geodesic distances. 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. Code available at https://github.com/olgatticus/rg-vfm.