Generative modeling of probability distributions on curved spaces
When Riemann flows with Wasserstein: Generative Modeling of Probability Distributions on Manifolds
Machine Learning
Summary
Many real-world datasets have complex shapes that can't be described using simple flat geometry. This paper introduces a new method that respects the curved nature of these spaces to more accurately generate new data. The authors developed a technique that learns paths between probability distributions on curved surfaces, making it possible to create realistic samples in settings like molecule shapes or cell data. Their approach works efficiently with different curved shapes and even irregular surfaces.
What this means in practice
- •For bioinformatics teams: Generate realistic molecular conformations and protein ensembles by modeling distributions on curved geometric spaces derived from molecular data.
- •For single-cell data scientists: Create full single-cell samples within hyperspherical latent spaces to better capture cellular heterogeneity respecting intrinsic data geometry.
Authors
Doron Haviv, Edward De Brouwer, Rishabh Anand, Rex Ying, Aïcha Bentaieb, Gabriele Scalia, Hector Corrada Bravo
Abstract
Many scientific datasets, such as molecular conformational ensembles or single-cell tissue measurements, are naturally modeled as meta-distributions: distributions over probability measures on non-Euclidean domains. Existing generative methods largely assume Euclidean geometry and fail to capture this structure. We introduce Riemannian Wasserstein Entropic Flow Matching (RWEFM), a generative framework on the Wasserstein space $\mathcal{P}_2(\mathcal{M})$ of a Riemannian manifold $(\mathcal{M},g)$. RWEFM is trained by regressing a neural vector field onto Riemannian optimal transport velocities, using McCann displacement interpolations as conditional paths. We confirm theoretically that this construction leads to a valid flow matching approach on $\mathcal{P}_2(\mathcal{M})$ and introduce the Riemannian Entropic Map, a GPU-efficient approximation of the optimal transport map on manifolds. Our experiments show that by respecting the intrinsic geometry of the data, RWEFM can generate whole single-cell samples in hyperspherical latent spaces and protein conformational ensembles on the torus. As RWEFM requires only a geodesic distance and a projection operator, it is not restricted to manifolds with closed-form geometry, which we demonstrate by generating distributions on a general triangulated mesh.