Learning Distributions on Manifolds with Free-Form Flows

TL;DR

The paper introduces Manifold Free-Form Flows (M-FFF), a fast and efficient generative model for data on manifolds that simplifies sampling by using a single neural network evaluation, outperforming previous methods.

Contribution

It adapts the free-form flow framework to Riemannian manifolds, enabling fast, single-step sampling and maximum likelihood training on manifolds with known projections.

Findings

M-FFF matches or outperforms specialized methods.

It is two orders of magnitude faster than diffusion-based methods.

Achieves better likelihoods in experiments.

Abstract

We propose Manifold Free-Form Flows (M-FFF), a simple new generative model for data on manifolds. The existing approaches to learning a distribution on arbitrary manifolds are expensive at inference time, since sampling requires solving a differential equation. Our method overcomes this limitation by sampling in a single function evaluation. The key innovation is to optimize a neural network via maximum likelihood on the manifold, possible by adapting the free-form flow framework to Riemannian manifolds. M-FFF is straightforwardly adapted to any manifold with a known projection. It consistently matches or outperforms previous single-step methods specialized to specific manifolds. It is typically two orders of magnitude faster than multi-step methods based on diffusion or flow matching, achieving better likelihoods in several experiments. We provide our code at https://github.com/vislearn/FFF.