Normalizing Flows on Riemannian Manifolds

TL;DR

This work revisits techniques related to homeomorphisms from differential geometry for projecting densities to sub-manifolds and uses it to generalize the idea of normalizing flows to more general Riemannian manifolds.

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$.

Figures (1)

• Figure 1: Left: Construction of a complex density on $\mathbf{S}^n$ by first projecting the manifold to $\mathbf{R}^{n}$, transforming the density and projecting it back to $\mathbf{S}^n$. Right: Illustration of transformed ($\mathbf{S}^2 \rightarrow \mathbf{R}^2$) densities corresponding to an uniform density on the sphere. Blue: empirical density (obtained by Monte Carlo); Red: Analytical density from equation \ref{['eq.sphere']}; Green: Density computed ignoring the intrinsic dimensionality of $\mathbf{S}^n$.