Generalised Flow Maps for Few-Step Generative Modelling on Riemannian Manifolds

TL;DR

The paper addresses the challenge of costly inference in geometric generative models by introducing Generalised Flow Maps (GFM) that extend flow-map-based sampling to arbitrary Riemannian manifolds. GFM defines manifold-constrained jump maps X_{s,t} via geodesic interpolation and the exponential map, and presents three self-distillation losses—Generalised Lagrangian Self-Distillation (G-LSD), Generalised Eulerian Self-Distillation (G-ESD), and Generalised Progressive Self-Distillation (G-PSD)—to learn these maps from scratch. The authors prove that GFM subsumes and improves upon Euclidean few-step models (consistency models, shortcut models, mean flows) in the Riemannian setting and validate the approach across torsion-angle data on tori, geospatial data on the sphere, SO(3) rotations, and hyperbolic manifolds, achieving state-of-the-art or competitive sample quality and log-likelihoods in low steps. They also establish that the Lagrangian variant frequently provides the strongest performance, while the implicit flow within the flow map can yield better NLL, highlighting practical and theoretical avenues for further exploration in geometric generative modeling.

Abstract

Geometric data and purpose-built generative models on them have become ubiquitous in high-impact deep learning application domains, ranging from protein backbone generation and computational chemistry to geospatial data. Current geometric generative models remain computationally expensive at inference -- requiring many steps of complex numerical simulation -- as they are derived from dynamical measure transport frameworks such as diffusion and flow-matching on Riemannian manifolds. In this paper, we propose Generalised Flow Maps (GFM), a new class of few-step generative models that generalises the Flow Map framework in Euclidean spaces to arbitrary Riemannian manifolds. We instantiate GFMs with three self-distillation-based training methods: Generalised Lagrangian Flow Maps, Generalised Eulerian Flow Maps, and Generalised Progressive Flow Maps. We theoretically show that GFMs, under specific design decisions, unify and elevate existing Euclidean few-step generative models, such as consistency models, shortcut models, and meanflows, to the Riemannian setting. We benchmark GFMs against other geometric generative models on a suite of geometric datasets, including geospatial data, RNA torsion angles, and hyperbolic manifolds, and achieve state-of-the-art sample quality for single- and few-step evaluations, and superior or competitive log-likelihoods using the implicit probability flow.

Figures (7)

• Figure 1: Left: The instantaneous vector field $v_t$ of the Riemannian probability flow ODE. Right: The Generalised flow map $X_{s,t}$ that jumps along the trajectory of the Riemannian probability flow ODE by steps of size $t-s$. Both models are depicted on $\mathbb{S}^2$.
• Figure 2: Ramachandran plots on the General protein dataset. Test-set samples depicted in red.
• Figure 3: MMD on protein datasets against the NFE.
• Figure 4: MMD on Earth datasets against the NFE.
• Figure 5: Plots of densities for the various datasets and all compared methods. Depicted in red are the test-set samples. Datasets from left to right: volcano, earthquake, flood, fire.

Theorems & Definitions (26)

• Definition 1
• Proposition 1
• Lemma 1
• Proposition 2
• Proposition 3
• Proposition 4
• Proposition 5: Flow map characterization
• proof
• Proposition 6: Euclidean self-distillation flowmaps
• proof