Equivariant Manifold Neural ODEs and Differential Invariants

TL;DR

A manifestly geometric framework for equivariant manifold neural ordinary differential equations (NODEs) is developed and it is shown that universality persists in the equivariant case and that the augmented equivariant manifold NODEs can be incorporated into the geometric framework using higher-order differential invariants.

Abstract

In this paper, we develop a manifestly geometric framework for equivariant manifold neural ordinary differential equations (NODEs) and use it to analyse their modelling capabilities for symmetric data. First, we consider the action of a Lie group $G$ on a smooth manifold $M$ and establish the equivalence between equivariance of vector fields, symmetries of the corresponding Cauchy problems, and equivariance of the associated NODEs. We also propose a novel formulation, based on Lie theory for symmetries of differential equations, of the equivariant manifold NODEs in terms of the differential invariants of the action of $G$ on $M$, which provides an efficient parameterisation of the space of equivariant vector fields in a way that is agnostic to both the manifold $M$ and the symmetry group $G$. Second, we construct augmented manifold NODEs, through embeddings into flows on the tangent bundle $TM$, and show that they are universal approximators of diffeomorphisms on any connected $M$. Furthermore, we show that universality persists in the equivariant case and that the augmented equivariant manifold NODEs can be incorporated into the geometric framework using higher-order differential invariants. Finally, we consider the induced action of $G$ on different fields on $M$ and show how it can be used to generalise previous work, on, e.g., continuous normalizing flows, to equivariant models in any geometry.

Figures (8)

• Figure 1: The figure shows $A(r)$ (blue line) and $B(r)$ (orange line) in \ref{['example:rot_eq_NODE']} as functions of $r$ after $300$ epochs. The dashed line in the top figure illustrates the function $a(r)=1/r$.
• Figure 2: $A$ and $B$ in \ref{['example:so2_sphere']} as functions of $\theta$ after 400 epochs of training. The dashed lines represent the functions $a(\theta)=\theta$ and $b(\theta)=0.05$, respectively.
• Figure 3: Intersection of paths $\gamma_p$ and $\gamma_{p'}$.
• Figure 4: The lift $\Gamma_p$ of $\gamma_p$ and the vector $\frac{d}{dt}\Gamma_p$ tangent to $TM$.
• Figure 5: In \ref{['example:augmentation']}, we train a non-augmented NODE model and an augmented NODE model to approximate the equivariant diffeomorphism $h(r,\theta) = (1/r,\theta)$. The orange dots represent the resulting outputs by the models, while the blue dots represent the target values. In (a), the neural ODE has not been augmented. We see that, as a result of the solution curves not being able to cross, the outputs freeze on a circle minimising the distance to the targets after around 1000 epochs. In (b), the NODE model has been augmented to the tangent bundle. As a result, the outputs are able to better approximate the targets.

Theorems & Definitions (32)

• Definition 2.1
• Definition 2.2
• Definition 2.3
• Theorem 2.4: Theorem 2.34 Olver1995
• Theorem 2.5: Proposition 2.56 Olver1993
• Remark 2.6
• Definition 3.1: Manifold neural ODE
• Remark 3.2
• Definition 3.3: Equivariant manifold neural ODE
• Theorem 3.4