Neural Manifold Ordinary Differential Equations

TL;DR

This paper introduces Neural Manifold Ordinary Differential Equations, a general framework for continuous normalizing flows on arbitrary manifolds, improving density estimation and downstream task performance by leveraging continuous manifold dynamics.

Contribution

It generalizes Neural ODEs to manifolds, enabling flexible, continuous normalizing flows applicable to any manifold with improved performance.

Findings

Marked improvement in density estimation

Enhanced performance on downstream tasks

General applicability to arbitrary manifolds

Abstract

To better conform to data geometry, recent deep generative modelling techniques adapt Euclidean constructions to non-Euclidean spaces. In this paper, we study normalizing flows on manifolds. Previous work has developed flow models for specific cases; however, these advancements hand craft layers on a manifold-by-manifold basis, restricting generality and inducing cumbersome design constraints. We overcome these issues by introducing Neural Manifold Ordinary Differential Equations, a manifold generalization of Neural ODEs, which enables the construction of Manifold Continuous Normalizing Flows (MCNFs). MCNFs require only local geometry (therefore generalizing to arbitrary manifolds) and compute probabilities with continuous change of variables (allowing for a simple and expressive flow construction). We find that leveraging continuous manifold dynamics produces a marked improvement for both density estimation and downstream tasks.