Generalised Flow Maps on Riemannian Manifolds

NeurIPS 2025 Workshop FPI Submission73 Authors

Track: Main Track

Keywords: generative modelling, Riemannian geometry, few-step models

TL;DR: Generalisation of few-step models on Riemannian manifolds

Abstract: Recent advances in generative modelling on Euclidean spaces have shown how to train models that achieve state-of-the-art quality from scratch, while requiring only few function evaluations. In the meanwhile, data supported on Riemannian manifolds, such as protein backbones or geological data, has lagged behind: inference and training on such manifolds remain computationally challenging and numerically unstable due to the need for manifold-specific operations, and require numerous evaluations of a potentially expensive model to obtain samples of high quality. In this paper, we propose Generalised Flow Maps, a new class of few-step generative models that generalise the flow map framework~\citep{boffi2024flow} to arbitrary Riemannian manifolds. We port and design three self-distillation-based training methods: Generalised Lagrange Flow Maps, Generalised Eulerian Flow Maps, and Generalised Semi-Group Flow Maps, all of which recover their Euclidean counterparts. Empirically, we test Generalised Flow Maps on a host of standard datasets, and achieve, amongst Riemannian generative models, state-of-the-art sample quality for single- and few-step evaluations, and log-likelihoods.

Submission Number: 73