Research map / Probabilistic and scientific modeling / Variational and generative sampling
Learning Distributions on Manifolds with Free-Form Flows
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends free-form flows to Riemannian manifolds embedded in Euclidean space, training a single feed-forward encoder/decoder pair with a manifold maximum-likelihood gradient estimator for one-step sampling and density estimation.
Problem
Learning distributions on arbitrary Riemannian manifolds with generative models is expensive at inference because existing methods (continuous-time flows, diffusion, flow matching) require solving ODEs/SDEs with many function evaluations, while specialized single-step flows are tailored to specific
Approach
Use free-form (non-invertible) encoder and decoder networks working in an embedding space with outputs projected onto the manifold, and estimate the gradient of the manifold change-of-variables NLL with a Hutchinson trace estimator using tangent-space-projected noise and the decoder Jacobian in place of the inverse encoder Jacobian.
Key results
- On SO(3) mixtures, M-FFF achieves NLL of -0.87, -0.21, and 0.45 for M=16, 32, 64, outperforming SO(3)-NF and multi-step baselines at M=32 and 64.
- On earth sphere datasets, M-FFF achieves -2.25 on Volcano vs. -7.93 for Riemannian flow matching (best); M-FFF is best single-step on all four datasets, with mixed results vs. multi-step methods.
- M-FFF is typically two orders of magnitude faster at inference than multi-step diffusion and flow matching methods, requiring a single function evaluation plus projection.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use M-FFF instead of the specialized SO(3)-NF | SO(3) synthetic mixtures · SO(3) mixture M=16, M=32, and M=64 | test NLL: M=16: -0.81 ± 0.01 -> -0.87 ± 0.02; M=32: -0.12 ± 0.004 -> -0.21 ± 0.02; M=64: 0.61 ± 0.01 -> 0.45 ± 0.02 | Helped |
| Use M-FFF instead of Riemannian flow matching | Sphere earth datasets · Volcano | test NLL: -7.93 ±1.67 -> -2.25 ± 0.02 | Hurt |
| Use M-FFF instead of the strongest reported multi-step baseline | Sphere earth datasets · Earthquake | test NLL: -0.40 ± 0.05 -> -0.23 ± 0.01 | Helped |
| Use M-FFF instead of the strongest reported multi-step baseline | Sphere earth datasets · Flood | test NLL: 0.42 ± 0.05 -> 0.51 ± 0.01 | Hurt |
| Use M-FFF instead of the strongest reported multi-step baseline | Sphere earth datasets · Fire | test NLL: -1.86 ± 0.11 -> -1.19 ± 0.03 | Hurt |
| Use M-FFF instead of circular spline coupling flows | Torus distributions · Glycine | NLL: 1.91 ± 0.04 -> 1.89 ± 0.05 | Helped |
| Use M-FFF instead of Riemannian flow matching | Torus distributions · RNA | NLL: -5.20 ± 0.067 -> -4.27 ± 0.09 | Hurt |
| Use M-FFF instead of Riemannian flow matching | Stanford bunny distributions · Stanford bunny k=50 | test NLL: 1.48 ± 0.01 -> 1.34 ± 0.01 | Helped |
| Use M-FFF instead of Riemannian flow matching | Stanford bunny distributions · Stanford bunny k=100 | test NLL: 1.53 ± 0.01 -> 1.28 ± 0.01 | Helped |
Surprising findings
- On the Volcano earth dataset, Riemannian flow matching (-7.93) outperforms M-FFF (-2.25) by a wide margin, and multi-step methods generally beat M-FFF when density changes sharply or low-density regions exist.
Limitations
- Likelihoods rely on decoder Jacobian approximation, and become inaccurate when reconstruction loss is high.
- Performs less favorably than multi-step methods when density changes sharply or very low-density regions are present.
- Gradient estimator is stochastic with a single noise sample per batch item, adding variance.
Experimental setup
Datasets: SO(3) synthetic mixture distributions Earth data (volcano, earthquake, flood, fire) Torus torsion angle datasets (General, Glycine, Proline, Pre-Pro, RNA) Stanford bunny synthetic distributions Poincare ball toy distributions
Benchmarks: SO(3) synthetic mixtures Earth sciences sphere datasets Torus protein and RNA torsion angles Stanford bunny synthetic distributions
Metrics: test negative log-likelihood inference speed (function evaluations)
Baselines: Moser flow Riemannian score-based models SO(3)-NF Riemannian CNF Riemannian diffusion Riemannian flow matching Mixture of Kent Circular Spline Coupling Flows Mixture of power spherical
Builds on: free-form flows Riemannian diffusion models Riemannian flow matching Hutchinson trace estimator
Closest papers by meaning
- Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes (ICML 2024 · similarity 0.91)
- Manifold Diffusion Fields (ICLR 2024 · similarity 0.89)
- Categorical Flow Matching on Statistical Manifolds (NeurIPS 2024 · similarity 0.88)
- Flow Matching on General Geometries (ICLR 2024 · similarity 0.88)
- Rectifying Categorical Flows on Statistical Manifolds for One-Step Generation (NeurIPS 2026 · similarity 0.88)
- Residual Deep Gaussian Processes on Manifolds (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.