Learning Distributions on Manifolds with Free-Form Flows

TL;DR

The key innovation is to optimize a neural network via maximum likelihood on the manifold by adapting the free-form flow framework to Riemannian manifolds, possible by adapting the free-form flow framework to Riemannian manifolds.

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.

Figures (6)

• Figure 1: Manifold Free-Form Flows (M-FFF) learn generative models on a variety of manifolds. (Left) The learned distributions (colored surface) accurately match the test points (black dots). (Right) We parameterize M-FFF using a neural network in an embedding space, whose outputs are projected to the manifold. This enables simulation-free training and inference, and naturally respects the corresponding geometry, yielding fast sampling and continuous distributions regardless of the manifold.
• Figure 2: Computation of the volume change in the tangent space of the manifold: The manifold change of variables formula in \ref{['eq:manifold-cov-isometric']} requires to compute the change of a volume element in the tangent spaces under f, which in this example is given by the ratio of lengths of $dt$ and $dt'$. Since $f$ is a map in the embedding space, $f'(x)$ defines a linear map between vectors from the embedding space. To correctly compute the change in volume, we use $Q$ and $R$ to change coordinates to the intrinsic tangent spaces, resulting in the linear map $R^T f'(x) Q: \mathcal{T}_x\mathcal{M} \rightarrow \mathcal{T}_{f(x)}\mathcal{M}$, which maps $dt$ to $dt'$.
• Figure 3: Manifold free-form flows on a synthetic $SO(3)$ mixture distribution with $M=64$ mixture components proposed by de2022riemannian. (Left) 10,000 samples each from the ground truth distribution and (right) our model. This visualization computes three Euler angles, which fully describe a rotation matrix, and then plot the first two angles on the projection of a sphere and the last by color murphy2021implicit. We find that our model nicely samples from the distribution with few outliers between the modes.
• Figure 4: Density estimates of our model on the earth datasets. Blue points show the training dataset, red points the test dataset.
• Figure 5: Log density of M-FFF models in the $(\Phi, \Psi)$-plane of protein backbone dihedral angles (known as Ramachandran plotramachandran1963stereochemistry). The learned density matches the true density indicated by the test dataset (black dots) very well. Note also that the learned distribution obeys the periodic boundary conditions.

Theorems & Definitions (14)

• Theorem 1: Volume change gradient estimator, draxler2024freeform
• Theorem 2: Manifold change of variables
• Theorem 3
• proof
• proof
• proof
• Corollary 4
• proof
• Theorem 5
• proof