[Paper Review] Riemannian Continuous Normalizing Flows

Émile Mathieu, Maximilian Nickel|arXiv (Cornell University)|Jun 18, 2020

Mathematical Dynamics and FractalsMathematics54 references30 citations

TL;DR

This paper introduces Riemannian continuous normalizing flows (RCNFs) that define flows as ODE-evolved vector fields on manifolds, enabling expressive probability models that respect underlying geometry. They demonstrate advantages over standard and projected flows on spherical and hyperbolic manifolds and in earth-science density estimation.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for probabilistic modeling that respects Riemannian geometry of manifold-valued data.
• Propose a principled framework to define continuous normalizing flows on smooth manifolds via vector fields.
• Develop geometry-aware methods for flow evaluation, likelihood computation, and base distribution choices.
• Empirically compare RCNFs against naive, wrapped, and projected methods on constant-curvature manifolds (sphere and Poincaré disk) and real-world earth-science datasets.
• Demonstrate improvements in density estimation and training efficiency over baselines.

Proposed method

• Define manifold flows via the time evolution dz(t)/dt = fθ(z(t), t) on a smooth complete manifold M.
• Prove vector flows yield a C1-diffeomorphic flow φ(·, t) on M under mild conditions.
• Compute likelihood via Liouville-style change in density using the Riemannian metric G(z) and a Hutchinson trace estimator for divergence.
• Use a geometry-aware Runge-Kutta + projection-based solver to integrate flows while keeping samples on the manifold.
• Parameterize the vector field fθ with a neural network that receives geodesic-distance features and outputs tangent vectors, with outputs projected to TM.
• Train with maximum likelihood or reverse KL objectives, backpropagating through the solver with memory-efficient autodiff.

Experimental results

Research questions

• RQ1 How can continuous normalizing flows be formulated to operate directly on Riemannian manifolds rather than Euclidean space?
• RQ2 Do RCNFs provide improved expressiveness and numerical stability compared to naive, wrapped, or projected manifold-flow approaches across constant-curvature manifolds?
• RQ3 What are the effects of geometry-aware likelihood computation and divergence estimation on training stability and performance?
• RQ4 How do RCNFs perform in density estimation tasks on real-world spherical data, such as Earth-science datasets, relative to baselines?

Key findings

• RCNFs consistently outperform naive and wrapped projected models on synthetic hyperbolic and spherical tasks as the target distribution approaches challenging regions near manifold boundaries.
• On spherical data, RCNFs yield better log-likelihood and reverse KL performance than stereographic projection methods, especially as target mass nears singularities.
• RCNFs converge faster and require fewer iterations than projected methods in density estimation on Earth-science datasets (volcanoes, earthquakes, floods, fires).
• Geodesic-distance input layers and proper scaling of vector-field outputs are crucial for stable training and reduced function evaluations.
• Density training on real-world datasets shows that RCNFs provide tighter fits and improved generalization compared to projection-based approaches.