Riemannian Continuous Normalizing Flows

Emile Mathieu; Maximilian Nickel

Source-linked AI summary

Riemannian Continuous Normalizing Flows

Emile Mathieu, Maximilian Nickel

arXiv:2006.10605v2Jun 18, 2020stat.MLcs.LG

TL;DR

Standard normalizing flows assume flat Euclidean geometry, limiting their suitability for manifold-valued data. The paper introduces Riemannian continuous normalizing flows based on manifold vector fields and ODEs; experiments show substantial improvements over conventional and projected flows, while related projection methods have scope and numerical limitations.

Problem
Conventional normalizing flows assume Euclidean geometry, making them ill-suited or misspecified for data on Riemannian manifolds.
Method
Riemannian continuous normalizing flows define vector fields on manifolds and compute trajectories by solving geometry-aware ordinary differential equations.
Results
The method yields substantial improvements on constant-curvature manifolds compared with conventional or projected flows and better generalization with faster convergence on earth-science data.
Takeaways & Limitations
Geometry-aware continuous flows provide a principled way to parametrize expressive probability distributions directly on Riemannian manifolds.
Takeaways & Limitations
Projected-flow approaches are limited to manifolds homeomorphic to Euclidean space and can face numerical instability near singularities.

Abstract

from arXiv · show

Normalizing flows have shown great promise for modelling flexible probability distributions in a computationally tractable way. However, whilst data is often naturally described on Riemannian manifolds such as spheres, torii, and hyperbolic spaces, most normalizing flows implicitly assume a flat geometry, making them either misspecified or ill-suited in these situations. To overcome this problem, we introduce Riemannian continuous normalizing flows, a model which admits the parametrization of flexible probability measures on smooth manifolds by defining flows as the solution to ordinary differential equations. We show that this approach can lead to substantial improvements on both synthetic and real-world data when compared to standard flows or previously introduced projected flows.

1 Introduction

Normalizing flows model flexible distributions but are poorly suited to data on Riemannian manifolds because they assume Euclidean geometry. The paper introduces geometry-aware continuous flows that remain on the manifold and reports benefits on synthetic and earth-science data.

Riemannian manifolds naturally describe data in applications including storm trajectories, protein shapes, cell development, and human actions.
Conventional normalizing flows assume Euclidean geometry, so their distributions may place mass outside the manifold and become ill-suited or misspecified.
Riemannian continuous normalizing flows define manifold vector fields through ODEs, using geometry-aware neural networks and numerical solvers to keep trajectories on the manifold.
The approach combines expressive, unconstrained continuous flows with a Riemannian formulation that avoids ad-hoc projection maps and reduces numerical artefacts.
Experiments on the Poincaré disk, sphere, and earth-science datasets compare the method with non-Riemannian and projected approaches, showing better generalization and faster convergence.

2 Continuous Normalizing Flows on Riemannian Manifolds

The paper extends continuous normalizing flows to Riemannian manifolds by defining manifold-valued particle dynamics through vector fields and ODEs. It derives geometry-aware likelihood computation and specifies practical choices for manifolds, solvers, divergence estimation, and base distributions.

Vector flows: The method pushes samples from a simple base distribution through time-evolving manifold-valued particles, producing a complex multimodal distribution without mapping through Euclidean space.A manifold-aware numerical solver integrates the particle evolution, while the approach permits almost free-form neural networks.
Likelihood: For constant-curvature manifolds, closed-form geometric quantities support practical likelihood computation, while the flow integral is approximated by discretizing the solver trajectory.The discussion focuses on these manifolds because their relevant geometrical quantities are available in closed form, although the proposed approach is described as applicable to broader classes such as tori and Grassmannians.
Vector flows: RCNFs define manifold flows as ODE solutions generated by vector fields fθ: M × R → TM, allowing flexible neural-network parameterizations on general Riemannian manifolds.Particles evolve from an initial position z(0)=z0 according to a manifold-valued vector field, and the resulting flow is bijective and differentiable under stated regularity conditions.
Practical choices: The implementation uses an adaptive fourth-order Runge–Kutta integrator, a uniform base distribution on the sphere, and a standard wrapped Gaussian base distribution on the Poincaré ball.The base-distribution choice reflects the stated practical settings: uniformity for the sphere and convenience for the Poincaré ball.
Likelihood: Likelihood evaluation replaces direct Jacobian-determinant computation with an integral of instantaneous log-density change expressed using the Riemannian metric G(z).The metric contributes a geometry-dependent term, while Hutchinson’s estimator approximates the divergence and backward autodifferentiation computes vector-Jacobian products with linear complexity.
Likelihood: The stochastic divergence estimator can have high variance because its variance scales with the inverse determinant term √|G(z(t))|.This is identified as a limitation of the estimator rather than a limitation of the manifold-flow construction itself.

3 Related work

Prior manifold normalizing-flow methods either project Euclidean flows onto manifolds or define transformations directly on the manifold. Projected approaches face topological, numerical, or computational limitations, whereas Riemannian methods avoid projection artifacts and continuous flows reduce structural constraints.

Projected methods: Projected methods push a Euclidean normalizing flow through an invertible map ψ: R^d → M, limiting applicability because such a map requires M to be homeomorphic to R^d.
Projected methods: Exponential-map flows define wrapped distributions by pushing a Euclidean probability measure through exp_μ, but compact manifolds require truncated infinite density summations.
Projected methods: Stereographic projection is non-injective and maps the projection point to infinity, making distributions near that point numerically challenging.
Riemannian methods: Riemannian methods operate directly on the manifold rather than mapping it to Euclidean space, avoiding projection-related artifacts and numerical instability.
Riemannian methods: The paper introduces continuous flows on general Riemannian manifolds, whose ODE-based parametrization alleviates structural constraints associated with discrete flows.

4 Experimental results

Experiments on hyperbolic and spherical synthetic targets, as well as earth-sciences data, compare Riemannian continuous normalizing flows with naive, wrapped, stereographic, and mixture-based alternatives. The Riemannian model performs better near geometric pathologies and across spherical datasets, while some components introduce computational or estimator-variance costs.

Experimental setup: The experiments compare Riemannian, naive, wrapped, stereographic, and mixture-based models on synthetic manifold targets and four spherical earth-sciences datasets.All models were designed with approximately the same number of parameters, and confidence intervals were computed over 12 runs.
Hyperbolic geometry: As targets move toward the Poincaré-disk boundary, the Riemannian model outperforms naive and wrapped models, which fail to fit targets far from the origin.The target location is controlled by α, and performance is evaluated by maximum likelihood.
Spherical geometry: Near the spherical stereographic singularity, the Riemannian model perfectly fits the target, whereas the stereographic model fails to push sufficient mass toward −µ0.The Riemannian model outperforms the stereographic model on both negative log-likelihood and reverse KL objectives, with the gap shrinking at lower concentration.
Real-world spherical data: On all four earth-sciences datasets, the Riemannian model outperforms the stereographic counterpart and mixture of vMF distributions by a large margin and converges faster.The learned stereographic distributions are generally more diffuse, allocate some mass outside the target support, and cover fewer data points.
Limitations and costs: The stochastic divergence estimator has high variance near the spherical north pole, requiring the exact estimator, while this variance behavior does not occur on the Poincaré ball.A different vector-field basis could alleviate the spherical issue on large-scale datasets.
Limitations and costs: The geodesic distance layer adds computational cost, improves hyperbolic performance empirically, and has less effect in the spherical setting.It is therefore an optional component that may improve model quality at additional computational cost.

5 Discussion

The paper proposes Riemannian continuous normalizing flows for expressive probability distributions on Riemannian manifolds. Experiments show substantial improvements over conventional and projected flows on constant-curvature manifolds.

Riemannian continuous normalizing flows parametrize expressive probability distributions directly on Riemannian manifolds.
The method defines flows through manifold vector fields and computes them as solutions to associated ordinary differential equations.
The approach yields substantial improvements for modelling data on constant-curvature manifolds compared with conventional or projected flows.

Broader impact

The work targets well-specified probabilistic modelling for manifold-valued data, with potential relevance to scientific fields including earth and climate science, computational biology, and computer vision. It motivates intrinsic manifold modelling because conventional models assume flat geometry and can be ill-suited or misspecified.

The method focuses on learning well-specified probabilistic models for manifold-valued data.
The work may support scientific understanding in earth and climate science, computational biology, and computer vision, while inheriting broader ethical aspects of machine learning.
Conventional probabilistic models assume flat geometry, making them ill-suited or misspecified for manifold distributions.
When a manifold is embedded in a higher-dimensional ambient space, the naive model is defined on a different probability space from the target and is misspecified.

Proof of Proposition 2

The proof derives density evolution for ODE-driven random variables on Riemannian manifolds. It expresses the change through the divergence of the vector field, whose computation incorporates the manifold’s metric tensor.

The density evolution follows a Liouville equation involving the divergence of the vector field.
On a Riemannian manifold, vector-field divergence is computed in local coordinates using the manifold’s metric tensor.
Equation 12 is used for practical numerical experiments.

D.1 l2-norm

The section connects continuous normalizing-flow regularization to dynamical optimal transport and extends the relevant formulation to complete Riemannian manifolds. It also motivates Jacobian regularization through variance reduction for divergence estimation.

D.1 l2-norm: The l2 norm regularization is motivated by a connection between continuous normalizing flows and the dynamical formulation of optimal transport.
D.1 l2-norm: The dynamical formulation represents optimal transport as a flow generated by an optimal vector field satisfying the continuity equation.
D.1 l2-norm: The Euclidean optimal-transport result is extended to complete Riemannian manifolds under stated regularity, completeness, boundary, and moment assumptions.
D.1 l2-norm: The optimal flow has geodesic, constant-velocity integral paths, motivating regularization toward smoother flows and potentially lower solver NFE at fixed tolerance.
D.1 l2-norm: Regularizing the vector-field Jacobian can reduce the variance of Hutchinson’s divergence estimator and thereby improve training.

E Vector flows and neural architecture

The paper establishes conditions for existence, uniqueness, and globality of manifold flows, then uses geometry-aware neural layers to parameterize expressive vector fields on manifolds and product manifolds.

Vector flows: C1 time-dependent vector fields generate unique smooth local flows whose time maps are local C1-diffeomorphisms.The local flow satisfies the vector-field ODE and starts from the identity map.
Vector flows: Local existence and uniqueness do not by themselves guarantee that flows are defined for all times or across the whole manifold.The paper distinguishes these local guarantees from global-flow requirements.
Global flows: On complete manifolds, bounded vector fields provide a sufficient condition for global flows; compact manifolds such as spheres and tori satisfy globality directly.The proof uses bounded speed to establish Lipschitz continuity and geodesic completeness to extend trajectories.
Neural architecture: A geodesic distance layer builds manifold-aware features from geodesic-distance neurons, after which standard feed-forward networks can be stacked.For hyperspheres, the construction uses great-circle geometry and hyperplanes represented by vectors.
Product manifolds: Product-manifold densities can be modeled autoregressively by composing conditional flows, yielding a lower-triangular Jacobian whose determinant is efficiently computed from diagonal terms.Each conditional density is defined on one component manifold and its vector-field parameters depend on preceding components.

G Experimental details

Experiments use comparable neural architectures and Monte Carlo objectives, with adaptive ODE integration and manifold projection for training the Riemannian models.

Architecture: The vector field uses a three-hidden-layer, 64-unit tanh MLP, while RCNF replaces its input layer with a geodesic distance layer.Projected and naive models use the standard MLP input architecture.
Objectives: The experiments optimize Monte Carlo estimators of negative log-likelihood and reverse KL divergence.These are the two stated training objectives.
Optimization: Optimization uses Adam with β1 = 0.9, β2 = 0.999, batch size 400, and learning rate 1e−3.These settings are reported for all models in the experimental setup.
Training: Training solves the ODE with an adaptive Dormand–Prince Runge–Kutta 4(5) method at absolute and relative tolerance 1e−5, projecting each solver step onto the manifold.The computations use GeForce RTX 2080 Ti GPUs.

G.1 Hyperbolic geometry and limits of conventional and wrapped methods

The hyperbolic experiment evaluates RCNF against conventional Euclidean and wrapped-normal baselines on a wrapped-normal target in the Poincaré disk.

Hyperbolic experiment: The target is a wrapped normal on B2 with covariance Σ = diag(0.3, 1.0), while α controls its distance from the disk origin.The target is centered through exp0(α ∂x), allowing different radial locations.
Hyperbolic experiment: The comparison includes a conventional CNF on R2, the proposed RCNF on B2, a standard Gaussian, and a standard wrapped normal.The named flows are φR2 and φB2, with Gaussian and wrapped-normal base distributions.
Evaluation protocol: All three models train for 1500 iterations by minimizing negative log-likelihood, with results averaged over 12 runs.Divergence during training is approximated with the Hutchinson stochastic estimator.

G.2 Spherical geometry

Spherical experiments compare RCNF with stereographic-projection models on synthetic and earth-science data, using tangent-vector constraints and dataset-specific training protocols.

Spherical models: The spherical comparison uses an inverse stereographic projection model and RCNF, with Gaussian and uniform-on-S2 base distributions.The models are denoted by ρ−1 ◦ φR2 and φS2.
Spherical models: RCNF projects each vector-field output onto the tangent space of S2, enforcing tangent output vectors.The spherical experiments use the exact divergence estimator.
Projection limits: The stereographic-projection limit experiment places a vMF target at μ = −μ0 and trains both models with negative log-likelihood and reverse KL objectives.Training lasts 3000 iterations and reported results average four runs.
Earth-science data: Earth-science density estimation covers volcano eruptions, earthquakes, floods, and wildfires, with results averaged over 12 runs.The datasets represent locations on the sphere.
Earth-science data: CNF training uses 1000 epochs except for volcano eruptions, which requires 3000 epochs; training tolerances are relaxed relative to evaluation because solver tolerance 1e−5 is computationally intensive.Training tolerances are 1e−3 for volcanoes, earthquakes, and wildfires, and 1e−4 for floods; evaluation retains 1e−5.
Earth-science data: Mixture-of-von-Mises–Fisher baselines use Riemannian Adam, with 10000 epochs generally and 30000 for volcano eruptions.Learning rates and mixture counts are selected by grid search.