Normalizing Flows on Riemannian Manifolds

TL;DR

This paper extends normalizing flows to Riemannian manifolds, enabling scalable density estimation on curved spaces like spheres, with applications in physics and directional statistics.

Contribution

It introduces a novel generalization of normalizing flows for Riemannian manifolds using differential geometry techniques, broadening their applicability beyond Euclidean spaces.

Findings

Successfully applied to the n-sphere $ extbf{S}^n$

Scalable and easy to implement with automatic differentiation

Demonstrated effectiveness in manifold density estimation

Abstract

We consider the problem of density estimation on Riemannian manifolds. Density estimation on manifolds has many applications in fluid-mechanics, optics and plasma physics and it appears often when dealing with angular variables (such as used in protein folding, robot limbs, gene-expression) and in general directional statistics. In spite of the multitude of algorithms available for density estimation in the Euclidean spaces $\mathbf{R}^n$ that scale to large n (e.g. normalizing flows, kernel methods and variational approximations), most of these methods are not immediately suitable for density estimation in more general Riemannian manifolds. We revisit techniques related to homeomorphisms from differential geometry for projecting densities to sub-manifolds and use it to generalize the idea of normalizing flows to more general Riemannian manifolds. The resulting algorithm is scalable, simple to implement and suitable for use with automatic differentiation. We demonstrate concrete examples of this method on the n-sphere $\mathbf{S}^n$.