Learning Distributions on Manifolds with Free-Form Flows
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Density Estimation | Volcano (test) | NLL-2.25 | 14 | |
| Density Estimation | Earthquakes NGDC/WDS, 2022a (test) | Negative Log-Likelihood-0.23 | 8 | |
| Density Estimation | Floods (test) | NLL0.51 | 8 | |
| Density Estimation | Wildfires EOSDIS, 2020 (test) | NLL-1.19 | 8 | |
| Density Estimation | Synthetic SO(3) M=32 (test) | NLL-0.21 | 5 | |
| Density Estimation | Synthetic SO(3) M=64 (test) | NLL0.45 | 5 | |
| Density Estimation | Protein T2 Glycine 500 proteins (test) | NLL1.89 | 5 | |
| Density Estimation | Protein T2 Pre-Pro 500 proteins (test) | NLL1.23 | 5 | |
| Density Estimation | RNA T7 (test) | NLL-4.27 | 5 | |
| Density Estimation | SO(3) M=16 synthetic (test) | NLL-0.87 | 5 |