Source-linked AI summary
Moser Flow: Divergence-based Generative Modeling on Manifolds
Noam Rozen, Aditya Grover, Maximilian Nickel, Yaron Lipman
TL;DR
Generative models for manifold-valued data are limited by geometry-specific methods and costly CNF training. Moser Flow uses a divergence-based density parameterization to avoid ODE solving during training, proves universality under suitable assumptions, and improves empirical performance while extending flow modeling to general curved surfaces.
Problem
Existing generative models are limited on manifold data because extensions are often geometry-specific, computationally costly, or struggle with general curved surfaces.
Method
Moser Flow parameterizes the target density as the source density minus the divergence of a neural network, yielding a CNF without ODE solving during training.
Results
Moser Flow improves density estimation and generation on challenging benchmarks, including a 49% average improvement over Riemannian CNFs on four earth and climate datasets, and supports general curved surfaces.
Takeaways & Limitations
Moser Flow provides a universal generative model under suitable assumptions and broadens flow-based sampling to complex manifold geometries.
Takeaways & Limitations
Scaling Moser Flow to higher dimensions is limited by probability underflow, increasing integral-estimation variance, and costly exact divergence computation.
Abstract
from arXiv · showhide
We are interested in learning generative models for complex geometries described via manifolds, such as spheres, tori, and other implicit surfaces. Current extensions of existing (Euclidean) generative models are restricted to specific geometries and typically suffer from high computational costs. We introduce Moser Flow (MF), a new class of generative models within the family of continuous normalizing flows (CNF). MF also produces a CNF via a solution to the change-of-variable formula, however differently from other CNF methods, its model (learned) density is parameterized as the source (prior) density minus the divergence of a neural network (NN). The divergence is a local, linear differential operator, easy to approximate and calculate on manifolds. Therefore, unlike other CNFs, MF does not require invoking or backpropagating through an ODE solver during training. Furthermore, representing the model density explicitly as the divergence of a NN rather than as a solution of an ODE facilitates learning high fidelity densities. Theoretically, we prove that MF constitutes a universal density approximator under suitable assumptions. Empirically, we demonstrate for the first time the use of flow models for sampling from general curved surfaces and achieve significant improvements in density estimation, sample quality, and training complexity over existing CNFs on challenging synthetic geometries and real-world benchmarks from the earth and climate sciences.
1 Introduction
Generative modeling must extend beyond Euclidean data to scientific and geometric datasets on manifolds. Moser Flow addresses limitations of existing manifold CNFs with a divergence-based parameterization and shows strong empirical performance across synthetic, spherical, and curved-surface settings.
- Manifold-valued scientific data include earth, climate, life-science, and graphics applications where Euclidean models can assign probability outside the natural geometry.
- Existing manifold flow methods target specific geometries, while Riemannian CNFs can converge slowly, produce lower-fidelity samples, and struggle with general curved surfaces.
- Moser Flow models the target density as the source density minus a neural-network divergence, avoiding explicit ODE solving during maximum-likelihood training.
- 49% average improvement over Riemannian CNFs was reported for density estimation across four earth and climate datasets, alongside high-fidelity generation on spherical geometries.
- Moser Flow also demonstrates generative modeling on general curved surfaces, extending flow-based sampling beyond spherical geometries.
2 Preliminaries
Manifold CNFs transform a source density into a target density through an orientation-preserving diffeomorphism and the change-of-variables formula. Their continuous formulation uses an ODE-generated flow whose endpoint maps source samples to the target distribution.
- The manifold framework represents probabilities as positive densities integrated against the Riemannian volume form, which defines geometric volume.
- A manifold CNF is an orientation-preserving diffeomorphism that maps a simple source distribution on a manifold to a complex target distribution.
- The change-of-variables construction pulls back the target volume form through the map and requires that pullback to equal the source density.
- When the normalizing equation holds, a source sample z distributed by ν maps through x = Φ(z) to a target-distributed sample, making Φ the generator.
- Continuous normalizing flows construct the map by solving an ODE for a time-dependent vector field, initialized at the identity and evaluated at its terminal time.
- In Euclidean space, the manifold formulation reduces to the standard normalizing-flow equation using the Jacobian matrix of partial derivatives.
3 Moser Flow
Moser Flow constructs a continuous normalizing flow by interpolating source and target densities, solving a divergence-based PDE on the manifold, and deriving a vector field whose flow solves the normalization equation. Its generative model parameterizes density with a neural vector field, trains with likelihood and positivity constraints, and is theoretically guaranteed to recover the target under stated conditions.
- 3.1 Solving the normalizing equation: Moser’s construction interpolates source ν and target µ through α_t, then defines a vector field whose flow solves the continuous normalization equation.The interpolant satisfies α_0 = ν and α_1 = µ; the resulting flow at t = 1 solves the normalizing equation.
- 3.1 Solving the normalizing equation: The method first solves a manifold PDE involving divergence, then constructs v_t from its solution; on boundaryless manifolds, divergence uses the Riemannian covariant derivative.The vector field must also be tangent to the boundary when the manifold has one.
- 3.1 Solving the normalizing equation: Theorem 1 guarantees that the time-one diffeomorphism generated by the ODE solves the normalization equation.A simple linear interpolant makes dα_t/dt = µ − ν independent of time, allowing a time-constant auxiliary vector field.
- 3.2 Generative model utilizing Moser Flow: Moser Flow defines its learned density from a neural vector field through the divergence-based parameterization and trains it with likelihood plus penalties enforcing positivity.The model uses a neural network for the vector-field degree of freedom and replaces the open constraint μ̄ > 0 with μ̄ ≥ ϵ.
4 Generative modeling over Euclidean submanifolds
The section adapts Moser Flows to compact Euclidean submanifolds by constructing tangent vector fields from ambient neural networks and exploiting their normal-direction behavior. This yields an implementable divergence calculation and a universal density model under the stated manifold assumptions.
- Geometric scope: The construction applies to orientable, compact, boundaryless, connected submanifolds M ⊂ R^d, including implicit surfaces, triangulated surfaces, and manifold simplicial complexes.The manifold has dimension n < d and uses the metric induced from the ambient Euclidean space.
- Model construction: MFs on Euclidean submanifolds use an ambient MLP-derived vector field whose restriction to the manifold is tangent.The construction defines u from an ambient vector field vθ, with Softplus activation and learnable parameters; for x ∈ M, u(x) ∈ T_xM.
- Normal-direction construction: The vector field is constant along manifold-normal directions, a property called infinitesimal constancy in the normal direction.This follows because the closest-point projection π(x) remains constant along normal directions.
- Divergence computation: Infinitesimal normal constancy makes the manifold divergence equal to the ambient Euclidean divergence, simplifying MF implementation.The paper states that div(u(x)) = divE(u(x)) for x ∈ M, so the Euclidean divergence can be computed in practice.
- Universality: Under the stated differentiability and compactness assumptions, MF is universal and can approximate any continuous target density on a Euclidean submanifold arbitrarily well.The theorem guarantees an MLP and parameter choice achieving arbitrary ε-accuracy for the density defined by equations 12 and 18.
5 Experiments
Moser Flow is evaluated on flat-torus toy data, implicit surfaces, Euclidean CNF comparisons, and earth-science datasets. The experiments examine density consistency, fidelity to complex geometries, and computational efficiency.
- Implicit surfaces: Moser Flow is applied to implicit surfaces including the sphere and Stanford Bunny, with experiments using increasing-frequency target distributions on the curved surface.Surfaces are represented as zero level sets of signed distance functions, and vector fields are modeled on the resulting submanifolds.
- Synthetic data: Toy experiments on flat-torus datasets visualize input samples, generated samples, and learned densities across distributions with differing complexity.More complex bottom-row datasets use four hidden layers, while the remaining datasets use three.
- Choice of hyper-parameter λ: Higher λ makes the learned positive density ¯µ+ closer to the generated density Φ∗ν, while even λ = 1 yields rather consistent densities.The effect is attributed to lower ¯µ− values, meaning ¯µ better represents a probability density.
- Flat torus comparison: Moser Flow captures high-frequency torus details better than FFJORD, with faster NLL decrease and iteration times roughly 1–2 orders of magnitude lower.MF iteration time does not increase during training, whereas FFJORD iterations take longer as training progresses.
- Earth and climate science data: 49% average improvement over Riemannian CNFs is reported for density estimation across four earth and climate science datasets.The datasets concern global locations of volcano eruptions, earthquakes, floods, and related earth-science phenomena.
6 Related Work
Related manifold flow methods include parametric approaches that map Euclidean flows onto manifolds and Riemannian approaches that operate directly on manifolds. Their scope and numerical behavior differ, with examples spanning wrapped distributions, Lie groups, hyperbolic space, spheres, and tori.
- Parametric methods: Parametric manifold flows push forward Euclidean normalizing flows through invertible maps, but global representation requires topological equivalence between the Euclidean space and the manifold.This requirement limits the scope of the approach.
- Curved-surface example: Moser Flow experiments on Stanford Bunny provide a curved-surface example with target distributions whose frequencies increase across conditions.Each condition displays the target density, data samples, generated samples, and learned density.
- Parametric methods: Exponential-map-based parametric methods produce wrapped distributions for Lie groups and hyperbolic space but can introduce numerical and computational challenges on compact manifolds.The cited examples include methods by Falorsi et al. (2019) and Bose et al. (2020).
- Riemannian methods: Riemannian methods operate directly on the manifold, avoiding numerical instabilities from mapping onto the manifold.Examples include geodesic transformations on hyperspheres and discrete flows based on Möbius transformations for hyperspheres and tori.
7 Discussion and limitations
Moser Flow combines a divergence-based CNF parameterization with theoretical universality and empirical gains on manifold data. Its main current limitation is scaling to higher dimensions, where probability, integration, and divergence computations become difficult.
- Discussion: MF represents target densities through a divergence-based neural parameterization, avoids ODE solving during training, and is theoretically universal under suitable assumptions.The model can approximately generate arbitrary positive target densities from arbitrary positive prior densities.
- Limitations: Scaling MF to higher dimensions remains a current limitation and important future-work direction.The limitation involves multiple computational components that become harder to manage as dimension increases.
- Limitations: In high dimensions, model probabilities should be computed in log-scale, integral-approximation variance increases, and exact divergence evaluation becomes too costly.The divergence term therefore needs approximation similarly to other CNF approaches.
- Limitations: MF, like other generative models, could be used to generate fake data and amplify harmful dataset biases.The paper identifies mitigation of these harms as an active area of ongoing research.
A Proof of Moser’s Theorem.
The proof constructs a time-dependent vector field whose flow transports the source density to the target density. Differential forms, Lie derivatives, Hodge decomposition, Stokes’ theorem, and vector-field correspondence establish existence and the required divergence relation.
- A Proof of Moser’s Theorem: Lie derivatives and Cartan’s formula reduce the evolution of the pulled-back volume form to a differential equation involving contraction and exterior differentiation.For top-degree forms, the exterior derivative term vanishes, simplifying the expression.
- A Proof of Moser’s Theorem: Using the correspondence between vector fields and differential forms, the exact component defines a vector field whose divergence satisfies the required equation.The proof then defines the final time-dependent vector field and verifies equation 23.
- A Proof of Moser’s Theorem: Hodge decomposition expresses the relevant form as exact plus harmonic components, while Stokes’ theorem and the boundaryless assumption eliminate the harmonic constant.This yields an exact form suitable for constructing the needed vector field.
- A Proof of Moser’s Theorem: Practically finding the vector field requires solving a Hodge-decomposition equation equivalent to a PDE on the manifold.This identifies the computational step needed to implement the theoretical construction.
B Other proofs
The supplementary proofs establish loss properties, universal approximation, and practical manifold implementations. They also describe local projection approximations, unnormalized-density scaling, and the setup used for Figure 7.
- Other proofs: The loss is minimized when the modeled density equals the target density, with both the generalized KL term and negative-density penalty minimized.The argument uses the assumption that the target density is bounded below by a positive constant.
- Other proofs: For submanifolds with induced metrics, the Riemannian covariant derivative is obtained by projecting the ambient derivative onto the tangent space.This supplies the geometric identity used in the manifold divergence analysis.
- Other proofs: An MLP can approximate the required vector field and first derivatives, enabling approximation of the divergence-based density on a compact manifold neighborhood.The proof relies on smooth activations and bounded derivatives of the projection map and its differential.
- Implementation: Mesh implementations use a cotangent Laplacian and mass matrix to compute Laplace–Beltrami eigenfunctions.The mass matrix contains Voronoi-cell areas for mesh vertices.
- Implementation: A first-order projection approximation is sufficient because the exact projection and its local approximation share value and first derivatives at each sampled manifold point.This avoids requiring the full projection operator.
- Implementation: Unnormalized densities can replace normalized ones on general manifolds because the relevant vector-field definition is invariant to volume scaling.Simple manifolds such as spheres and flat tori have closed-form volume, while general cases need not.
F Additional Experimental Details
The experiments used one NVIDIA Quadro RTX 6000 GPU per run and seed configuration, with remaining details provided in the main paper.
- F Additional Experimental Details: Each run and seed configuration used one NVIDIA Quadro RTX 6000 GPU, with code implemented in PyTorch.The authors state that the code would be open-sourced after review.