Geometric Neural Ordinary Differential Equations: From Manifolds to Lie Groups

Abstract

Neural ordinary differential equations (neural ODEs) are a well-established tool for optimizing the parameters of dynamical systems, with applications in image classification, optimal control, and physics learning. Although dynamical systems of interest often evolve on Lie groups and more general differentiable manifolds, theoretical results for neural ODEs are frequently phrased on ${\mathbb{R}}^n$. We collect recent results for neural ODEs on manifolds and present a unifying derivation of various results that serves as a tutorial to extend existing methods to differentiable manifolds. We also extend the results to the recent class of neural ODEs on Lie groups, highlighting a non-trivial extension of manifold neural ODEs that exploits the Lie group structure.

1. Introduction

Ordinary differential equations (ODEs) are ubiquitous in the engineering sciences, from modeling and control of simple physical systems like pendulums and mass–spring–dampers, or more complicated robotic arms and drones, to the description of high-dimensional spatial discretizations of distributed systems, such as fluid flows, chemical reactions, or quantum oscillators. Neural ordinary differential equations (neural ODEs) [1,2] are ODEs parameterized by neural networks. Given a state x, and parameters $\theta$ representing the weights and biases of a neural network, a neural ODE reads as follows:

$$\dot{x}=f_{\theta }(x,t),x\left(0\right)=x_0.$$

(1)

First introduced by [1] as the continuum limit of recurrent neural networks, the number of applications of neural ODEs quickly exploded beyond simple classification tasks: learning highly nonlinear dynamics of multi-physical systems from sparse data [3,4,5], optimal control of nonlinear systems [6], medical imaging [7], and real-time handling of irregular time series [8], to name but a few. Discontinuous state transitions and dynamics [9,10], time-dependent parameters [11], augmented neural ODEs [12], and physics-preserving formulations [13,14] present further extensions that increase the expressivity of neural ODEs.

However, these methods are typically phrased for states $x\in {\mathbb{R}}^n$. For many physical systems of interest, such as robot arms, humanoid robots, and drones, the state lives on differentiable manifolds and Lie groups [15,16]. More generally, the manifold hypothesis in machine learning raises the expectation that many high-dimensional data-sets evolve on intrinsically lower-dimensional, albeit more complicated, manifolds [17]. Neural ODEs on manifolds [18,19] presented significant steps to address this gap, with the first optimization methods for neural ODEs on manifolds. Yet, the general tools and approaches available on ${\mathbb{R}}^n$, such as running costs, augmented states, time-dependent parameters, control inputs, or discontinuous state transitions, are rarely addressed in a manifold context. Similar issues persist in a Lie group context, where neural ODEs on Lie groups [20,21] have been formalized.

Our goal is to extend further architectures and costs for neural ODEs from ${\mathbb{R}}^n$ to arbitrary manifolds (cf.

Table 1. Summary of neural ODEs on manifolds and Lie groups presented in this article.

Name of Neural ODE	Subtype	Trajectory Cost	Subsection	Originally Introduced in
Neural ODEs on manifolds (Section 3)	Extrinsic	Running and final cost	Section 3.1.1	Final cost [19], running cost [21]
	Intrinsic	Running and final cost, intermittent cost	Section 3.1.2 and Section 3.2.1	Final cost [18], running cost [21], intermittent cost (this work)
	Augmented, time-dependent parameters	Final cost	Section 3.2.2	Augmenting $\mathcal{M}$ to $T\mathcal{M}$ [23], Augmenting $\mathcal{M}$ to $\mathcal{M}\times \mathcal{N}$ (this work)
Neural ODEs on Lie groups (Section 4)	Extrinsic	Final cost and intermittent cost	Section 4.1	In [20]
	Intrinsic, dynamics in local charts	Running and final cost	Section 4.2	In [21,24]
	Intrinsic, dynamics on Lie algebra	Running and final cost	Section 4.2	In [21]

), and in particular Lie groups, and to equip the reader with the technical background for their own extensions. Here the main conceptual challenge lies in phrasing chart-independent optimization methods [18,19] in a manner that easily adapts to a variety of neural ODE architectures and cost functions [1,3,12]. To this end we present a systematic approach for deriving geometric versions of the adjoint sensitivity method [1,2], which is a memory-efficient and scalable tool for the optimization of neural ODEs (cf. Section 1.1). Such benefits extend to manifolds and Lie groups [18,19,20,21]. A second challenge, both conceptual and practical, lies in expressing various manifolds in terms of local charts and in expressing neural-net-parameterized functions, dynamics, and tensor fields in local charts. To this end we classify existing methods into extrinsic [19,20] and intrinsic [18,21,22] approaches, a distinction inspired by well-known differential geometric concepts. In our context the distinction suggests different parameterizations, affects numerical integration techniques, and affects scaling to high-dimensional dynamics. Specifically, our contributions are as follows:

• Systematic derivation of adjoint methods for neural ODEs on manifolds and Lie groups, highlighting the differences and equivalence of various approaches—for an overview, see also Table 1;
• Summarizing the state of the art of manifold and Lie group neural ODEs by formalizing the notion of extrinsic and intrinsic neural ODEs;
• A tutorial on neural ODEs on manifolds and Lie groups, with a focus on the derivation of coordinate-agnostic adjoint methods for optimization of various neural ODE architectures. Readers will gain a conceptual understanding of the geometric nature of the underlying variables, a coordinate-free derivation of adjoint methods and learn to incorporate additional geometric and physical structures. On the practical side, this will aid in the derivation and implementation of adjoint methods with non-trivial terms for various architectures, also with regard to coordinate expressions and chart transformations.

The remainder of this article is organized as follows. A brief state of the art on neural ODEs concludes this introduction. Section 2 provides a background on differentiable manifolds, Lie groups, and the coordinate-free adjoint method. Section 3 describes neural ODEs on manifolds and derives parameter updates via the adjoint method for various common architectures and cost functions, including time-dependent parameters, augmented neural ODEs, running costs, and intermediate cost terms. Section 4 describes neural ODEs on matrix Lie groups, explaining the merits of treating Lie groups separately from general differentiable manifolds. Both Section 3 and Section 4 also classify methods into extrinsic and intrinsic approaches. We conclude with a discussion in Section 5, highlighting advantages, disadvantages, challenges, and promise of the presented material. Appendix A includes a background on Hamiltonian systems, which appear when transforming the adjoint method into a form that is unique to Lie groups.

1.1. Literature Review

For a general introduction to neural ODEs, see [25]. Neural ODEs on ${\mathbb{R}}^n$ with fixed parameters were first introduced by [1], and parameter optimization via the adjoint method allowed for intermittent and final cost terms on each trajectory. The generalized adjoint method [2] also allows for running cost terms. Memory-efficient checkpointing is introduced in [26] to address stability issues of adjoint methods. Augmented neural ODEs [12] introduced augmented state spaces to allow neural ODEs to express arbitrary diffeomorphisms. Time-varying parameters were introduced by [11], with similar benefits to augmented neural ODEs. Neural ODEs with discrete transitions were formulated in [9,10], with [9] also learning event-triggered transitions common in engineering applications. Neural controlled differential equations (CDEs) were introduced in [27] for handling irregular time series, and parameter updates reapply the adjoint method [1]. Neural stochastic differential equations (SDEs) were introduced in [28], relying on a stochastic variant of the adjoint method for the parameter update. The previously mentioned literature phrases dynamics of neural ODEs on ${\mathbb{R}}^n$.

Recent trends in research on neural ODEs focus on structure preservation to improve performance and reduce training time by appropriately restricting the class of parameterized vector fields. This includes symmetry preservation by equivariant [23] and approximately equivariant neural ODEs [29], which tackle symmetric and approximately symmetric time series and dynamics, e.g., in N-body dynamics and molecular dynamics. It also includes physics preservation in a physics learning context, where Hamiltonian neural networks [30,31] and (generalized) Lagrangian neural networks [32,33,34] improve performance by guaranteeing energy conservation. In control and model order reduction, port-Hamiltonian neural ODEs [3,35] further allow for learning models that interact with external ports in a power-preserving manner. These methods also phrase dynamics on ${\mathbb{R}}^n$ and frequently apply the adjoint method for parameter updates.

Neural ODEs on manifolds were first introduced by [19], including an adjoint method on manifolds for final cost terms and application to continuous normalizing flows on Riemannian manifolds, but embedding manifolds into ${\mathbb{R}}^n$. Neural ODEs on Riemannian manifolds are expressed in local exponential charts in [18], avoiding embedding into ${\mathbb{R}}^n$ and considering final cost terms in the optimization. Charts for unknown, non-trivial latent manifolds together with dynamics in local charts are learned from high-dimensional data in [22], also including discretized solutions to partial differential equations. Parameterized equivariant neural ODEs on manifolds are constructed in [23], also commenting on state augmentation to express arbitrary (equivariant) flows on manifolds.

Neural ODEs on Lie groups were first introduced in [36] on the Lie group $SE\left(3\right)$ to learn the port-Hamiltonian dynamics of a drone from an experiment, expressing group elements on an embedding ${\mathbb{R}}^{12}$, and the approach was formalized to port-Hamiltonian systems on arbitrary matrix Lie groups in [20], embedding $m\times m$ matrices in ${\mathbb{R}}^{m^2}$.

Neural ODEs on $SE\left(3\right)$ were phrased in local exponential charts in [24] to optimize a controller for a rigid body using a chart-based adjoint method in local exponential charts. As an alternative, a Lie algebra-based adjoint method on general Lie groups was introduced in [21], foregoing Lie group-specific numerical issues of applying the adjoint method in local charts.

The choice of numerical solver in integrating neural ODEs and adjoint sensitivity equations is a nuanced area with much active research, especially for highly stiff [37], highly nonlinear [38,39], and structure-preserving neural ODEs [40]. We point towards the aforementioned sources for the interested reader. Results are expected to carry over into a manifold and Lie group context, where they hold in local charts. Also Lie group integrators [41,42] may be of interest for geometrically exact integration but are not well-investigated in a neural ODE context [20,21].

The optimization of neural ODEs via adjoint sensitivity methods is also referred to as “optimize-then-discretize” [25,43], since the formulation of the continuous adjoint system (called “optimize”) precedes their numerical solution (called “discretize”). This is opposed to “discretize-then-optimize” approaches, in which the neural ODE is first solved numerically (discretize) and gradients are then backpropagated through the numerical solver (optimize) [25,37,43]. Comparing the two, the constant memory efficiency of “optimize-then-discretize” approaches allows them to scale better to high-dimensional systems, giving them an edge for cases with more than 100 parameters and states [43]. Instead, “discretize-then-optimize” boasts higher accuracy and speed for low-dimensional systems, as well as highly stiff systems in which adjoint methods struggle with stability [37]. A popular discrete alternative to neural ODEs for physics-informed dynamics learning is given by variational integrator networks (VINs) [44,45], phrasing Lagrangian and Hamiltonian dynamics as discrete systems that conserve energy and the symplectic structure of the continuum dynamics [46,47]. Recent work [48] on Lie group forced VINs (LieFVINs) also allows inputs to the Lagrangian and Hamiltonian dynamics to be included in the variational formulation, allowing discrete optimal control. Both VINs and LieFVINs are applicable in a Lie group context, where they conserve geometry, symplecticity, and energy. The approach does not use adjoint methods for optimization and outperforms neural ODEs in the investigated conservative, low-dimensional dynamical systems [44,45,48]. Compared to continuous neural ODEs, both VINs and LieFVINs are discrete, which removes overhead from ODE solvers for lightweight applications, but their necessarily energy-based formulation presently restricts their use cases to conservative physical systems. We mention this promising area for completeness but narrow our attention to a geometric “optimize-then-discretize” approach via adjoint methods in the remainder of this article.

1.2. Notation

For a complete introduction to differential geometry see, e.g, [49], and for Lie group theory see [50].

Calligraphic letters $\mathcal{M},\mathcal{N},\dots$ denote smooth manifolds. For conceptual clarity, the reader may think of these manifolds as embedded in a high-dimensional ${\mathbb{R}}^N$, e.g., $\mathcal{M}\subset {\mathbb{R}}^N$. The set $C^{\infty }(\mathcal{M},\mathcal{N})$ contains smooth functions between $\mathcal{M}$ and $\mathcal{N}$, and we define $C^{\infty }\left(\mathcal{M}\right):=C^{\infty }(\mathcal{M},\mathbb{R})$.

The tangent space at $x\in \mathcal{M}$ is $T_x\mathcal{M}$ and the cotangent space is $T_x^{*}\mathcal{M}$. The tangent bundle of $\mathcal{M}$ is $T\mathcal{M}$, and the cotangent bundle of $\mathcal{M}$ is $T^{*}\mathcal{M}$. Then $\mathfrak{X}\left(\mathcal{M}\right)$ denotes the set of vector fields over $\mathcal{M}$, and ${\mathrm{\Omega }}^k\left(\mathcal{M}\right)$ denotes the set of k forms, where ${\mathrm{\Omega }}^1\left(\mathcal{M}\right)$ are co-vector fields and ${\mathrm{\Omega }}^0\left(\mathcal{M}\right)=C^{\infty }\left(\mathcal{M}\right)$ are smooth functions $V:\mathcal{M}\to \mathbb{R}$. The exterior derivative is denoted as $\mathrm{d}:{\mathrm{\Omega }}^k\left(\mathcal{M}\right)\to {\mathrm{\Omega }}^{k+1}\left(\mathcal{M}\right)$. For functions $V\in C^{\infty }(\mathcal{M}\times \mathcal{N},\mathbb{R})$, with $x\in \mathcal{M}$, $y\in \mathcal{N}$, we denote by $\left({\mathrm{d}}_{x}V\right)\left(y\right)\in T_x^{*}\mathcal{M}$ the partial differential at $x\in \mathcal{M}$. Curves $x:\mathbb{R}\to \mathcal{M}$ are denoted as $x\left(t\right)$, and their tangent vectors are denoted as $\dot{x}\in T_{x\left(t\right)}\mathcal{M}$.

A Lie group is denoted by G and its elements by $g,h$. The group identity is $e\in G$, and I denotes the identity matrix. The Lie algebra of G is $\mathfrak{g}$, and its dual is ${\mathfrak{g}}^{*}$. Letters $\tilde{A},\tilde{B}$ denote vectors in the Lie algebra, while letters $A,B$ denote vectors in ${\mathbb{R}}^n$.

In coordinate expressions, lower indices are covariant and upper indices are contravariant components of tensors. For example for a $(0,2)$-tensor M the components $M_{ij}$ are covariant, and for non-degenerate M the components of its inverse $M^{-1}$ are $M^{ij}$, which are contravariant. We use the Einstein summation convention $a_i{b}^i:={\sum }_i{a}_i{b}^i$; i.e., the product of variables with repeated lower and upper indices implies a sum.

Denoting W as a topological space, D the Borel $\sigma$-algebra, and $\mathbb{P}:D\to [0,1]$ a probability measure, the tuple $(W,D,\mathbb{P})$ denotes a probability space. Given a vector space L and a random variable $C:X\to L$, the expectation of C with respect to $\mathbb{P}$ is ${\mathbb{E}}_{w\sim \mathbb{P}}\left(C\right):={\int }_{W}C\left(w\right)\mathrm{d}\mathbb{P}\left(w\right)$.

2. Background

2.1. Smooth Manifolds

Given an n-dimensional manifold $\mathcal{M}$, with $U\subset \mathcal{M}$ being an open set and $Q:U\to {\mathbb{R}}^n$ a homeomorphism, we call $(U,Q)$ a chart and we denote the coordinates of $x\in U$ as

$$(q^1,\dots ,q^n):=Q\left(x\right),x\in U\subset \mathcal{M}.$$

(2)

Smooth manifolds admit charts $(U_1,Q_1)$ and $(U_2,Q_2)$ with smooth transition maps $Q_{21}=Q_2\circ Q_1^{-1}$ defined on the intersection $U_1\bigcap U_2$, and a collection $\mathcal{A}$ of charts $(U,Q)$ with smooth transition maps is called a smooth atlas. For examples of local charts for particular manifolds, see [49], Example 1.4, Example 1.5. A vector field $f\in \mathfrak{X}\left(\mathcal{M}\right)$ assigns a vector $f\left(x\right)\in T_x\mathcal{M}$ at any point $x\in \mathcal{M}$. This defines a dynamic system, also shown in a local chart $(U,Q)$ with components $f^i\left(q\right),{\dot{q}}^i\in \mathbb{R}$:

$$\begin{array}{cccc}\hfill \dot{x}& =f\left(x\right)=f^i\left(q\right){\displaystyle \frac{\partial }{\partial Q^i}};\hfill & \hfill & x\left(0\right)=x_0,\hfill \end{array}$$

(3)

$$\begin{array}{cccc}\hfill {\dot{q}}^i& =f^i\left(q\right);\hfill & \hfill & q\left(0\right)=Q\left(x_0\right).\hfill \end{array}$$

(4)

Solutions of (3) are then found by numerical integration of (4), applying chart transitions (e.g., $q_2\left(t\right)=Q_{21}\left(q_1\left(t\right)\right)$ from $q_1\left(t\right)=Q_1\left(x\left(t\right)\right)$ to $q_2\left(t\right)=Q_2\left(x\left(t\right)\right)$) during integration to avoid coordinate singularities (cf. Section 3.1.2). Denote the solution of (3) by the flow operator

$${\mathrm{\Psi }}_f^t:\mathcal{M}\to \mathcal{M};{\mathrm{\Psi }}_f^t\left(x_0\right):=x\left(t\right).$$

(5)

For a real-valued function $V\in C^{\infty }\left(\mathcal{M}\right)$, its differential is the covector field

$$\mathrm{d}V\in {\mathrm{\Omega }}^1\left(\mathcal{M}\right);\mathrm{d}V={\displaystyle \frac{\partial V}{\partial q^i}}\mathrm{d}{Q}^i.$$

(6)

Additionally, given a smooth manifold $\mathcal{N}$ and a smooth map $\phi :\mathcal{N}\to \mathcal{M}$, with $(U,Q)$ and $(\overline{U},\overline{Q})$ appropriate charts of $\mathcal{M}$ and $\mathcal{N}$, respectively, the pullback of $\mathrm{d}V$ via $\phi$ is

$${\phi }^{*}\mathrm{d}V\in {\mathrm{\Omega }}^1\left(\mathcal{N}\right);{\phi }^{*}\mathrm{d}V:=\mathrm{d}(V\circ \phi )={\displaystyle \frac{\partial {\phi }^j}{\partial {\overline{q}}^i}}{\displaystyle \frac{\partial V}{\partial q^j}}\mathrm{d}{\overline{Q}}^i.$$

(7)

With a Riemannian metric M (i.e., a symmetric, non-degenerate (0,2) tensor field) on $\mathcal{M}$, the gradient of V is a uniquely defined vector field $\nabla V\in \mathfrak{X}\left(\mathcal{M}\right)$ given by

$$\nabla V:=M^{-1}\mathrm{d}V=M^{ij}{\displaystyle \frac{\partial V}{\partial q^j}}{\displaystyle \frac{\partial }{\partial q^i}}.$$

(8)

When $\mathcal{M}={\mathbb{R}}^n$, we assume that M is the Euclidean metric and pick coordinates such that the components of the gradient and differential are the same. Finally, we define the Lie derivative of 1-forms, which differentiates $\omega \in {\mathrm{\Omega }}^1\left(\mathcal{M}\right)$ along a vector field $f\in \mathfrak{X}\left(\mathcal{M}\right)$ and returns ${\mathcal{L}}_f\omega \in {\mathrm{\Omega }}^1\left(\mathcal{M}\right)$:

$${\mathcal{L}}_f\omega :={\displaystyle \frac{\mathrm{d}}{\mathrm{d}t}}{\left({{\mathrm{\Psi }}_f^t}^{*}\omega \right)}_{t=0}={\omega }_j\left({\displaystyle \frac{\partial }{\partial q^i}}{f}^j\right)\mathrm{d}{Q}^i+\left({\displaystyle \frac{\partial }{\partial q^j}}{\omega }_i\right)f^j\mathrm{d}{Q}^i.$$

(9)

2.2. Lie Groups

Lie groups are smooth manifolds with a compatible group structure. We consider real matrix Lie groups $G\subseteq GL(m,\mathbb{R})$, i.e., subgroups of the general linear group

$$GL(m,\mathbb{R}):=\{g\in {\mathbb{R}}^{m\times m}|\det \left(g\right)\ne 0\}.$$

(10)

For $g,h\in G$ the left and right translations by h are, respectively, the matrix multiplications

$$\begin{array}{c}\hfill L_h\left(g\right):=hg,\end{array}$$

(11)

$$\begin{array}{c}\hfill R_h\left(g\right):=gh.\end{array}$$

(12)

The Lie algebra of G is the vector space $\mathfrak{g}\subseteq \mathfrak{gl}(m,\mathbb{R})$, with $\mathfrak{gl}(m,\mathbb{R})={\mathbb{R}}^{m\times m}$ being the Lie algebra of $GL(m,\mathbb{R})$.

Define a basis $E:=\{{\tilde{E}}_1,\dots ,{\tilde{E}}_n\}$ with ${\tilde{E}}_i\in \mathfrak{g}\subset {\mathbb{R}}^{m\times m}$, and define the (invertible linear) map $\mathrm{\Lambda }:{\mathbb{R}}^n\to \mathfrak{g}$ as (equivalently (e.g, [51]), $\mathrm{\Lambda }$ and ${\mathrm{\Lambda }}^{-1}$ are often denoted as the operators “hat” $\wedge :{\mathbb{R}}^n\to {\mathbb{R}}^{m\times m}$ and “vee” $\vee :{\mathbb{R}}^{m\times m}\to {\mathbb{R}}^n$, respectively)

$$\mathrm{\Lambda }:{\mathbb{R}}^n\to \mathfrak{g};(A^1,\dots ,A^n)\mapsto \sum _i{A}^i{\tilde{E}}_i.$$

(13)

The dual of $\mathfrak{g}$ is denoted ${\mathfrak{g}}^{*}$, and given the map $\mathrm{\Lambda }$ we call ${\mathrm{\Lambda }}^{*}:{\mathfrak{g}}^{*}\to {\mathbb{R}}^n$ its dual. For $\tilde{A},\tilde{B}\in \mathfrak{g}$ the small adjoint ${\mathrm{ad}}_{\tilde{A}}\left(\tilde{B}\right)$ is a bilinear map, and the large adjoint ${\mathrm{Ad}}_g\left(\tilde{A}\right)$ is a linear map

$$\begin{array}{cc}\hfill \mathrm{ad}& :\mathfrak{g}\times \mathfrak{g}\to \mathfrak{g};{\mathrm{ad}}_{\tilde{A}}\left(\tilde{B}\right)=\tilde{A}\tilde{B}-\tilde{B}\tilde{A},\hfill \end{array}$$

(14)

$$\begin{array}{cc}\hfill \mathrm{Ad}& :G\times \mathfrak{g}\to \mathfrak{g};{\mathrm{Ad}}_g\left(\tilde{A}\right)=g\tilde{A}{g}^{-1}.\hfill \end{array}$$

(15)

In the remainder of this article, we exclusively use the adjoint representation ${\mathrm{ad}}_A:{\mathbb{R}}^n\to {\mathbb{R}}^n$, written without a tilde in the subscript A, and adjoint representation ${\mathrm{Ad}}_g:{\mathbb{R}}^n\times {\mathbb{R}}^n$, which are obtained as

$$\begin{array}{cc}\hfill {\mathrm{ad}}_A& :={\mathrm{\Lambda }}^{-1}\left({\mathrm{ad}}_{\mathrm{\Lambda }\left(A\right)}\mathrm{\Lambda }\left(·\right)\right),\hfill \end{array}$$

(16)

$$\begin{array}{cc}\hfill {Ad}_g& :={\mathrm{\Lambda }}^{-1}\left({\mathrm{Ad}}_g\mathrm{\Lambda }\left(·\right)\right).\hfill \end{array}$$

(17)

The exponential map $exp:\mathfrak{g}\to G$ is a local diffeomorphism given by the matrix exponential ([50], Chapter 3.7)

$$exp\left(\tilde{A}\right):=\sum _{n=0}^{\infty }{\displaystyle \frac{1}{n!}}{\tilde{A}}^n.$$

(18)

Its inverse $log:U_{log}\to \mathfrak{g}$ is given by the matrix logarithm, and it is well-defined on a subset $U_{log}\subset G$ ([50], Chapter 2.3):

$$log\left(g\right)=\sum _{n=1}^{\infty }{(-1)}^{n+1}{\displaystyle \frac{{(g-I)}^n}{n}}.$$

(19)

Often, these infinite sums in (18) and (19) can be further reduced to a finite sums in m terms by use of the Cayley–Hamilton theorem [52]. A chart $(U_h,Q_h)$ on G that assigns zero coordinates to $h\in G$ can be defined using (19) and (13):

$$\begin{array}{cc}\hfill U_h& =\left\{hg\right|g\in U_{log}\},\hfill \end{array}$$

(20)

$$\begin{array}{cc}\hfill Q_h& :U_h\to {\mathbb{R}}^n;g\mapsto {\mathrm{\Lambda }}^{-1}log\left(h^{-1}g\right),\hfill \end{array}$$

(21)

$$\begin{array}{cc}\hfill Q_h^{-1}& :{\mathbb{R}}^n\to G;q\mapsto hexp\left(\mathrm{\Lambda }\left(q\right)\right).\hfill \end{array}$$

(22)

The chart $(U_h,Q_h)$ is called an exponential chart, and a collection $\mathcal{A}$ of exponential charts $(U_h,Q_h)$ that cover the manifold is called an exponential atlas.

The differential of a function $V\in C^{\infty }(G,\mathbb{R})$ is the co-vector field $\mathrm{d}V\in {\mathrm{\Omega }}^1\left(G\right)$ (see also Equation (6)). For any given $g\in G$ we further transform the co-vector $\mathrm{d}V\left(g\right)\in T_g^{*}G$ to a left-trivialized differential, which collects the components of the gradient expressed in ${\mathfrak{g}}^{*}$:

$$\begin{array}{c}\hfill {\mathrm{d}}_g^{L}V:={\mathrm{\Lambda }}^{*}{L}_g^{*}\mathrm{d}V\left(g\right)={\displaystyle \frac{\partial }{\partial q}}V{\left(g\left(I+\mathrm{\Lambda }\left(q\right)\right)\right)}_{|q=0}\in {\mathbb{R}}^n.\end{array}$$

(23)

For a derivation of this coordinate expression, see ([21], Section 3).

2.3. Gradient over a Flow

We are interested in computing the gradient of functions with respect to the initial state of a flow. The adjoint sensitivity equations are a set of differential equations that achieve this. In the following, we show a derivation of the adjoint sensitivity on manifolds ([21], App. A2). Given a function $\mathrm{C}:\mathcal{M}\to \mathbb{R}$, a vector field $f\in \mathfrak{X}\left(\mathcal{M}\right)$, the associated flow ${\mathrm{\Psi }}_f^t:\mathcal{M}\to \mathcal{M}$, and a final time $T\in \mathbb{R}$, the goal of the adjoint sensitivity method on manifolds is to compute the gradient

$$\mathrm{d}\left(C\circ {\mathrm{\Psi }}_f^T\right)\left(x_0\right).$$

In the adjoint method we define a co-state $\lambda \left(t\right)=\mathrm{d}(C\circ {\mathrm{\Psi }}^{T-t})\left(x\left(t\right)\right)\in T_{x\left(t\right)}^{*}\mathcal{M}$, which represents the differential of $C\left(x\left(T\right)\right)$ with respect to $x\left(t\right)$. The adjoint sensitivity method describes its dynamics, which are integrated backwards in time from the known final condition $\lambda \left(T\right)=\mathrm{d}C\left(x\left(T\right)\right)$, see also Figure 1. The adjoint sensitivity method is stated in Theorem 1.

Theorem 1

(Adjoint sensitivity on manifolds). The gradient of a function $C\circ {\mathrm{\Psi }}_f^T$ is

$$d\left(C\circ {\mathrm{\Psi }}_f^T\right)\left(x_0\right)=\lambda \left(0\right),$$

(24)

where $\lambda \left(t\right)\in T_{x\left(t\right)}^{*}\mathcal{M}$ is the co-state. In a local chart $(U,Q)$ of $\mathcal{M}$ with coordinates $q\left(t\right)=Q\left(x\right(t\left)\right)$, $\lambda \left(t\right)={\lambda }_i\left(t\right)\mathrm{d}{Q}^i$, the state and co-state satisfy the dynamics

$$\begin{array}{cc}\hfill {\dot{q}}^j& =f^j\left(q\right),q\left(0\right)=Q\left(x_0\right),\hfill \end{array}$$

(25)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_i& =-{\lambda }_j{\displaystyle \frac{\partial }{\partial q^i}}{f}^j\left(q\right),{\lambda }_i\left(T\right)={\displaystyle \frac{\partial C}{\partial q^i}}\left(x\left(T\right)\right).\hfill \end{array}$$

(26)

Proof. 

Define the co-state $\lambda \left(t\right)\in T_{x\left(t\right)}^{*}\mathcal{M}$ as

$$\begin{array}{cc}\hfill \lambda \left(t\right)& :={\left({\mathrm{\Psi }}_f^{T-t}\right)}^{*}\mathrm{d}C\left(x\left(T\right)\right).\hfill \end{array}$$

(27)

Then Equation (24) is recovered by application of Equation (7):

$$\lambda \left(0\right)={\left({\mathrm{\Psi }}_f^T\right)}^{*}\mathrm{d}C\left(x\left(T\right)\right)=(\mathrm{d}C\circ {\mathrm{\Psi }}_f^T)\left(x_0\right),$$

(28)

A derivation of the dynamics governing $\lambda \left(t\right)$ constitutes the remainder of this proof. By definition of $\lambda \left(t\right)$ and the Lie derivative (9), we have that ${\mathcal{L}}_f\lambda \left(t\right)=0$:

$$\begin{array}{cc}\hfill {\mathcal{L}}_f\lambda \left(t\right)& ={\displaystyle \frac{d}{ds}}{\left({\left({\mathrm{\Psi }}_f^s\right)}^{*}\lambda (t+s)\right)}_{s=0}\hfill \\ \hfill & ={\displaystyle \frac{d}{ds}}\lambda \left(t\right)=0.\hfill \end{array}$$

(29)

If we further treat $\lambda$ as a 1-form $\lambda \in {\mathrm{\Omega }}^1\left(\mathcal{M}\right)$ (denoted as $\lambda$ by an abuse of notation), we obtain

$$\begin{array}{cc}\hfill {\mathcal{L}}_f\lambda =& {\lambda }_j\left({\displaystyle \frac{\partial }{\partial q^i}}{f}^j\right)\mathrm{d}{Q}^i+\left({\displaystyle \frac{\partial }{\partial q^j}}{\lambda }_i\right)f^j\mathrm{d}{Q}^i=0.\hfill \end{array}$$

The components satisfy the partial differential equation

$${\lambda }_j{\displaystyle \frac{\partial }{\partial q^i}}{f}^j+f^j{\displaystyle \frac{\partial }{\partial q^j}}{\lambda }_i=0.$$

(30)

Impose that $\lambda \left(t\right)=\lambda \left({\mathrm{\Psi }}_f^t\left(x_0\right)\right)$ (this defines the 1-form $\lambda$ along $x\left(t\right)$); then

$${\dot{\lambda }}_i={\displaystyle \frac{\partial {\lambda }_i}{\partial q^j}}{\dot{q}}^j={\displaystyle \frac{\partial {\lambda }_i}{\partial q^j}}{f}^j.$$

(31)

Combining Equations (30) and (31) leads to Equation (26):

$${\dot{\lambda }}_i=-{\lambda }_j{\displaystyle \frac{\partial }{\partial q^i}}{f}^j.$$

(32)

Expanding the final condition $\lambda \left(T\right)=\mathrm{d}C\left(x\left(T\right)\right)$ in local coordinates (see Equation (6)) gives

$$\lambda \left(T\right)={\displaystyle \frac{\partial C}{\partial q^i}}\left(x\left(T\right)\right)\mathrm{d}{Q}^i={\lambda }_i\left(T\right)\mathrm{d}{Q}^i\iff {\lambda }_i\left(T\right)={\displaystyle \frac{\partial C}{\partial q^i}}\left(x\left(T\right)\right).$$

(33)

□

Given a chart transition from a chart $(U_1,Q_1)$ to a chart $(U_2,Q_2)$, e.g., during numerical integration of (26), the respective co-state components ${\lambda }_{1,i}$ and ${\lambda }_{2,i}$ are related by a transformation $A_i^j={\partial }_i\left(Q_1^j\circ Q_2^{-1}\right)$ as follows:

$$\begin{array}{cc}\hfill {\lambda }_{i,2}& =A_i^j{\lambda }_{j,1}.\hfill \end{array}$$

(34)

A fact that will become useful in Section 4 is that Equations (25) and (26) have a Hamiltonian form. Define the control Hamiltonian $H_c:T^{*}\mathcal{M}\to \mathbb{R}$ as

$$H_c(x,\lambda )=\lambda \left(f(x,t)\right)={\lambda }_i\left(f^i(q,t)\right).$$

(35)

Then Equation (25) and Equation (26), respectively, of Theorem 1 follow as the Hamiltonian equations on $T^{*}\mathcal{M}$:

$$\begin{array}{cc}\hfill {\dot{q}}^j& ={\displaystyle \frac{\partial H_c}{\partial {\lambda }_j}}=f^j(q,t),\hfill \end{array}$$

(36)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_i& =-{\displaystyle \frac{\partial H_c}{\partial q^i}}=-{\lambda }_j{\displaystyle \frac{\partial }{\partial q^i}}{f}^j(q,t).\hfill \end{array}$$

(37)

For a background on Hamilton’s equations, see also Appendix A.

3. Neural ODEs on Manifolds

A neural ODE on a manifold is an NN-parameterized vector field in $\mathfrak{X}\left(\mathcal{M}\right)$—or including time dependence, it is an NN-parameterized vector field in $\mathfrak{X}(\mathcal{M}\times \mathbb{R})$, with t in the $\mathbb{R}$ slot and $\dot{t}=1$. Given parameters $\theta \in {\mathbb{R}}^{n_{\theta }}$, we denote this parameterized vector field as $f_{\theta }(x,t):=f(x,t,\theta )$. This results in the dynamic system

$$\dot{x}=f_{\theta }(x,t),x\left(0\right)=x_0.$$

(38)

The key idea of neural ODEs is to tackle various flow approximation tasks by optimizing the parameters with respect to a to-be-specified optimization problem. Denote a finite time horizon T and intermittent times $T_1,T_2,\dots <T$. Denote a general trajectory cost by

$$C_{f_{\theta }}^T(x_0,\theta )=F(\theta ,{\mathrm{\Psi }}_{f_{\theta }}^{T_0}\left(x_0\right),{\mathrm{\Psi }}_{f_{\theta }}^{T_1}\left(x_0\right),\dots ,{\mathrm{\Psi }}_{f_{\theta }}^T\left(x_0\right))+{\int }_0^{T}r({\mathrm{\Psi }}_{f_{\theta }}^s\left(x_0\right),s)\mathrm{d}s,$$

(39)

with an intermittent and final cost term F and running cost r. Indicating a probability space $(\mathcal{M},D,\mathbb{P})$, we define the total cost as

$$J\left(\theta \right):={\mathbb{E}}_{x_0\sim \mathbb{P}}{C}_{f_{\theta }}^T(x_0,\theta ).$$

(40)

The minimization problem takes the form

$$\begin{array}{c}\hfill \underset{\theta }{min}J\left(\theta \right).\end{array}$$

(41)

Note that (41) is not subject to any dynamic constraint—the flow already appears explicitly in the cost $C_{f_{\theta }}^T$.

Normally, the optimization problem is solved by means of a stochastic gradient descent optimization algorithm [53]. In this, a batch of N initial conditions $x_i$ is sampled from the probability distribution corresponding to the probability measure $\mathbb{P}$. Writing $C_i=C_{f_{\theta }}^T(x_i,\theta )$, the parameter gradient ${\displaystyle \frac{\partial }{\partial \theta }}J\left(\theta \right)$ is approximated as

$${\displaystyle \frac{\partial }{\partial \theta }}J\left(\theta \right)={\mathbb{E}}_{x_0\sim \mathbb{P}}{\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left(x_0\right)\approx {\displaystyle \frac{1}{N}}\sum _{i=0}^N{\displaystyle \frac{\partial }{\partial \theta }}{C}_i.$$

(42)

In this section, we show how to optimize the parameters $\theta$ for various choices of neural ODEs and cost functions, with (39) being the most general case of a cost, and highlight similarities in the various derivations. In the following, the gradient ${\displaystyle \frac{\partial }{\partial \theta }}{C}_i$ is computed via the adjoint method on manifolds for various scenarios. The advantage of the adjoint method over, e.g., automatic differentiation of $C_i$/backpropagation through an ODE solver is that it has a constant memory efficiency with respect to the network depth T.

3.1. Constant Parameters and Running and Final Cost

Here we consider neural ODEs of the form (38), with constant parameters $\theta$ and cost functions of the form

$$C_{f_{\theta }}^T(x_0,\theta )=F({\mathrm{\Psi }}_{f_{\theta }}^T\left(x_0\right),\theta )+{\int }_0^{T}r({\mathrm{\Psi }}_{f_{\theta }}^s\left(x_0\right),\theta ,s)\mathrm{d}s,$$

(43)

with a final cost term F and a running cost term r. This generalizes [2] to manifolds. Compared to existing manifold methods for neural ODEs [18,54], the running cost is new.

The parameter gradient’s components ${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left((x_0,t_0),\theta \right)\in {\mathbb{R}}^{n_{\theta }}$ are then computed by Theorem 2 (see also [21]):

Theorem 2

(Generalized Adjoint Method on Manifolds). Given the dynamics (38) and the cost (43), the parameter gradient’s components ${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left((x_0,t_0),\theta \right)\in {\mathbb{R}}^{n_{\theta }}$ are computed by

$${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left((x_0,t_0),\theta \right)=\left({\displaystyle \frac{\partial F}{\partial \theta }}\right)(x\left(T\right),\theta )+{\int }_0^T{\displaystyle \frac{\partial }{\partial \theta }}\left({\lambda }_j{f}_{\theta }^j\left(q\left(s\right)\right)+r(q\left(s\right),\theta ,s)\right)ds.$$

(44)

where the state $x\left(s\right)\in \mathcal{M}$ and co-state $\lambda \left(s\right)\in T_{x\left(s\right)}^{*}\mathcal{M}$ satisfy, in a local chart $(U,Q)$ with $q\left(t\right)=Q\left(x\right(t\left)\right)$, $\lambda \left(t\right)={\lambda }_i\left(t\right)\mathrm{d}{Q}^i$,

$$\begin{array}{cc}\hfill {\dot{q}}^j& =f_{\theta }^j(q,t),q\left(0\right)=Q\left(x_0\right),t\left(0\right)=t_0,\hfill \end{array}$$

(45)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_i& =-{\lambda }_j{\displaystyle \frac{\partial }{\partial q^i}}{f}_{\theta }^j(q,t)-{\displaystyle \frac{\partial r}{\partial q^i}},{\lambda }_i\left(T\right)={\displaystyle \frac{\partial F}{\partial q^i}}\left(x\left(T\right),\theta \right).\hfill \end{array}$$

(46)

Proof. 

Define the augmented state space as ${\mathcal{M}}^{\prime }=\mathcal{M}\times {\mathbb{R}}^{n_{\theta }}\times \mathbb{R}\times \mathbb{R}$ to include the original state $x\in \mathcal{M}$, parameters $\theta \in {\mathbb{R}}^{n_{\theta }}$, accumulated running cost $L\in \mathbb{R}$, and time $t\in \mathbb{R}$ in the augmented state $x^{\prime }:=(x,\theta ,L,t)\in {\mathcal{M}}^{\prime }$. In addition, define the augmented dynamics $f_{\mathrm{aug}}\in \mathfrak{X}\left({\mathcal{M}}^{\prime }\right)$ as

$${\dot{x}}^{\prime }=f_{\mathrm{aug}}\left(x^{\prime }\right)=\left(\begin{array}{c}{f}_{\theta }(x,t)\\ 0\\ r(x,\theta ,t)\\ 1\end{array}\right),x^{\prime }\left(0\right)=x_0^{\prime }:=\left(\begin{array}{c}{x}_0\\ \theta \\ 0\\ t_0\end{array}\right).$$

(47)

This is an autonomous system with final state $x^{\prime }\left(T\right)=(x\left(T\right),\theta ,{\int }_0^{T}r(x,\theta ,s)\mathrm{d}s,T)$. Next, define the cost $C_{\mathrm{aug}}:{\mathcal{M}}^{\prime }\to \mathbb{R}$ on the augmented space:

$$C_{\mathrm{aug}}\left(x^{\prime }\right)=F(x,\theta )+L.$$

(48)

Then Equation (43) can be rewritten as the evaluation of a terminal cost $C_{\mathrm{aug}}\left(x^{\prime }\left(T\right)\right)$:

$$C_{f_{\theta }}^T\left(x_0\right)=(C_{\mathrm{aug}}\circ {\mathrm{\Psi }}_{f_{\mathrm{aug}}}^T)\left(x_0^{\prime }\right).$$

(49)

By Theorem 1, the gradient $\mathrm{d}\left(C_{\mathrm{aug}}\circ {\mathrm{\Psi }}_{f_{\mathrm{aug}}}^T\right)$ is given by

$$\mathrm{d}(C_{\mathrm{aug}}\circ {\mathrm{\Psi }}_{f_{\mathrm{aug}}}^T)\left(x_0^{\prime }\right)=\lambda \left(0\right),$$

(50)

and by Equation (26), the components of $\lambda \left(s\right)$ satisfy

$$\begin{array}{c}\hfill {\dot{\lambda }}_i=-{\lambda }_j{\displaystyle \frac{\partial }{\partial {q^i}^{\prime }}}{f}_{\mathrm{aug}}^j,{\lambda }_i\left(T\right)={\displaystyle \frac{\partial }{\partial {q^i}^{\prime }}}{C}_{\mathrm{aug}}\left(x^{\prime }\left(T\right)\right)\end{array}$$

(51)

Split the co-state into ${\lambda }_q,{\lambda }_{\theta },{\lambda }_L,{\lambda }_t$; then their components’ dynamics are as follows:

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{q,i}& =-{\displaystyle \frac{\partial }{\partial q^i}}\left({\lambda }_{q,j}{f}_{\theta }^j(q,t)+{\lambda }_{L}r(q,\theta ,t)\right),{\lambda }_q\left(T\right)={\displaystyle \frac{\partial F}{\partial q}}(x\left(T\right),\theta ),\hfill \end{array}$$

(52)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{\theta ,i}& =-{\displaystyle \frac{\partial }{\partial {\theta }^i}}\left({\lambda }_{q,j}{f}_{\theta }^j(q,t)+{\lambda }_{L}r(q,\theta ,t)\right),{\lambda }_{\theta }\left(T\right)={\displaystyle \frac{\partial F}{\partial \theta }}(x\left(T\right),\theta ),\hfill \end{array}$$

(53)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_L& =0,{\lambda }_L\left(T\right)={\displaystyle \frac{\partial }{\partial L}}{C}_{\mathrm{aug}}(x\left(T\right),\theta )=1,\hfill \end{array}$$

(54)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_t& =-{\displaystyle \frac{\partial }{\partial t}}\left({\lambda }_{q,j}{f}_{\theta }^j(q,t)+{\lambda }_{L}r(q,\theta ,t)\right),{\lambda }_t\left(T\right)={\displaystyle \frac{\partial }{\partial t}}{C}_{\mathrm{aug}}(x\left(T\right),\theta )=0.\hfill \end{array}$$

(55)

The component ${\lambda }_L=1$ is constant, so Equation (52) coincides with (46). Integrating (53) from $s=0$ to $s=T$ recovers Equation (44). ${\lambda }_t$ does not appear in any of the other equations, so Equation (55) may be ignored. □

In summary, the above proof depends on identifying a suitable augmented manifold ${\mathcal{M}}^{\prime }$, with the goal that augmented dynamics $f_{\mathrm{aug}}\in \mathfrak{X}\left({\mathcal{M}}^{\prime }\right)$ are autonomous such that the cost function $C_{\mathrm{aug}}:{\mathcal{M}}^{\prime }\to \mathbb{R}$ on the augmented manifold rephrases the cost (43) as a final cost $C_{\mathrm{aug}}\left(x\left(T\right)\right)$. This allows Theorem 1 to be applied to derive the corresponding adjoint method. In later sections (Section 3.2), this process will be the main technical tool for generalizations of Theorem 2. The next sections describe common special cases of (38) and Theorem 2.

3.1.1. Vanilla Neural ODEs and Extrinsic Neural ODEs on Manifolds

The case of neural ODEs on ${\mathbb{R}}^n$ (e.g., [2]) is obtained by setting $\mathcal{M}={\mathbb{R}}^n$. Scalar functions, vector fields, and tensor fields are readily expressed, see

Table 2. Parameterization of functions in extrinsic neural ODEs.

Function	Vanilla Neural ODE	Extrinsic Neural ODE
Scalar fields $V_{\theta }\left(x\right)\in \mathbb{R}$	$V_{\theta }:{\mathbb{R}}^n\to \mathbb{R}$	$V_{\theta }:{\mathbb{R}}^N\to \mathbb{R}$
Vector fields $f_{\theta }(x,t)\in T_x\mathcal{M}$	$f_{\theta }:{\mathbb{R}}^n\times \mathbb{R}\to {\mathbb{R}}^n$	$f_{\theta }^{\uparrow }:{\mathbb{R}}^N\times \mathbb{R}\to {\mathbb{R}}^N$ with tangency constraint [19], optional stabilization [55]
Components of $(p,q)$-tensor fields $M_{j_1,\dots ,j_q}^{i_1,\dots ,i_p}\left(x\right)\in \mathbb{R}$	$M_{j_1,\dots ,j_q}^{i_1,\dots ,i_p}:{\mathbb{R}}^n\to \mathbb{R}$	$M_{j_1,\dots ,j_q}^{i_1,\dots ,i_p}:{\mathbb{R}}^N\to \mathbb{R}$, see footnote 1

¹ A tangency constraint on contravariant components of $(p,q)$-tensors is not necessarily required for the vector field $f_{\theta }^{\uparrow }$ to remain tangent to $\iota \left(\mathcal{M}\right)$ and depends on the vector field under investigation.

.

There is an overlap with extrinsic neural ODEs on manifolds (described, for instance, in [19]), which optimize the neural ODE on an embedding space ${\mathbb{R}}^N$, see also Figure 2.

We denote the embedding as $\iota :\mathcal{M}\to {\mathbb{R}}^N$, where $x\in \mathcal{M}$ and $y\in {\mathbb{R}}^N$. Optimizing the neural ODE on ${\mathbb{R}}^N$ requires extending the dynamics $f_{\theta }(·,t)\in \mathfrak{X}\left(\mathcal{M}\right)$ to a vector field $f_{\theta }^{\uparrow }(·,t)\in \mathfrak{X}\left({\mathbb{R}}^N\right)$ such that

$${\iota }_{*}{f}_{\theta }(x,t)=f_{\theta }^{\uparrow }(\iota \left(x\right),t).$$

(56)

The dynamics $f_{\theta }^{\uparrow }(y,t)$ are then used in Theorem 2, and also the co-state lives in $T^{*}{\mathbb{R}}^N$.

As shown in [19], the resulting parameter gradients are equivalent to those resulting from an application in local charts, as long as it can be guaranteed that the integral curves of $f^{\uparrow }(y,t)$ remain within $\iota \left(\mathcal{M}\right)\subset {\mathbb{R}}^N$, i.e., are geometrically exact. Geometrically exact integration has to be guaranteed separately, either by integration in local charts [18] or stabilization techniques [55].

A strong upside of an extrinsic formulation is that existing neural ODE packages (e.g., [56]) can be applied directly. A downside to extrinsic neural ODEs is that finding $f^{\uparrow }(y,t)$ may not be immediate, since tangency to $\iota \left(\mathcal{M}\right)$ is required, see also Table 2. Finally, the extrinsic dimension N can be much larger than the intrinsic dimension $n=dim\mathcal{M}$, leading to computational overhead that does not fully exploit the manifold hypothesis. Extrinsic methods for neural ODEs are the preferred choice when the intrinsic dimension n is small and there is a known embedding $\iota \left(\mathcal{M}\right)\subset {\mathbb{R}}^N$ with low extrinsic dimension N. Then the computational overhead due to $N>n$ is negligible, and stabilization techniques [55] can be applied to guarantee geometrically exact integration.

3.1.2. Intrinsic Neural ODEs on Manifolds

The intrinsic case of neural ODEs on manifolds [18] is described by integrating the dynamics in local charts, see also Figure 3.

The advantage of intrinsic over extrinsic neural ODEs on manifolds is that the dimension of the resulting equations is as low as possible in the intrinsic case for a given manifold. The flexibility of chart representations gives intrinsic neural ODEs on manifolds the power to represent high-dimensional data distributions at their latent dimension, see especially [22] for learning charts from data and [18,21] for chart-switching methods during numerical integration. Numerical integration in local charts is also geometrically exact by default. However parameterization of scalar functions, vector fields, and tensor fields with neural networks in local charts, as well as their differentiation with respect to parameters, presents a source of complexity. There are three common methods to parameterize scalar-valued functions $V\in C^{\infty }\left(\mathcal{M}\right)$ in local charts (vector fields and tensor-valued functions directly follow by parameterizing their scalar component functions in an analogous way):

• A partition of unity ${\sigma }_i:{\mathbb{R}}^n\to \mathbb{R}$ with respect to a collection of charts $(U_i,Q_i)$ can be used to sum over chart -components $V_i:{\mathbb{R}}^n\to \mathbb{R}$ as $V\left(x\right):={\sum }_i{\sigma }_i\left(Q_i\left(x\right)\right)V_i\left(Q_i\left(x\right)\right)$, see examples in [49], Chapter 2, and [24].
• The function V can be directly defined by chart representatives $V_i:=V\circ Q_i^{-1}$, enforcing compatibility between overlapping charts $(U_i,Q_i),(U_j,Q_j)$ by soft constraints, which are implemented as additional cost terms $\parallel V_i\left(q_i\right)-V_j\circ Q_j\circ Q_i^{-1}\left(q_i\right)\parallel$ that are minimized on chart overlaps $U_i\cap U_j$, see [18,22].
• Given an embedding $\iota :\mathcal{M}\to {\mathbb{R}}^N$ and $\overline{V}\in C^{\infty }\left({\mathbb{R}}^N\right)$, an extrinsic representation $V_i:=\overline{V}\circ \iota \circ Q_i^{-1}$ is possible, see [19,21].

Advantages and disadvantages are summarized in

Table 3. Parameterization of scalar functions and tensor components in intrinsic neural ODEs.

Partition of Unity [24,49]	Soft Constraint [18,22]	Pullback [19,21]
Components from all local charts are summed, weighted by a partition of unity.	Function is directly represented in local charts.	Function is pulled back to local chart.
Allows representation of arbitrary smooth functions.	Functions are smooth where charts do not overlap, but are not well-defined at chart transitions.	Allows representation of arbitrary smooth functions.
Differentiating functions generally requires differentiating through chart transition maps, creating computational overload [24].	Chart transition maps do not have to be differentiated.	Chart representations of the embedding $\iota \left(\mathcal{M}\right)$ are differentiated, possibly creating computational overload.

.

In available state-of-the-art packages for neural ODEs, the chart dynamics are phrased as discontinuous dynamics with state transitions $Q_1\circ Q_2^{-1}:{\mathbb{R}}^n\to {\mathbb{R}}^n$, but implementation is not yet streamlined for local charts, chart transitions of the co-state (cf. Section 2.3), and custom adjoint sensitivity equations.

3.1.3. Structure Preservation

Structure-preserving architectures narrow down the class of learnable neural ODEs from arbitrary vector fields $\mathfrak{X}\left(\mathcal{M}\right)$ to subsets of $\mathfrak{X}\left(\mathcal{M}\right)$ with particular properties, improving training speed and performance (cf. Section 1.1). Examples of such subsets are (symmetry-preserving) equivariant dynamical systems [23] and (physics-preserving) Hamiltonian, Lagrangian, and port-Hamiltonian dynamical systems [3]. Given that a structure-preserving parameterization of the neural ODE is known in closed form, these are readily implemented in the above formalism.

For example, reusing results from Table 2 and Table 3, Hamiltonian and Lagrangian neural ODEs [30,32] are fully determined by scalar functions $H_{\theta }\in C^{\infty }\left(T^{*}\mathcal{Q}\right),L_{\theta }\in C^{\infty }\left(T\mathcal{Q}\right)$, respectively, and their gradients. Hamiltonian neural ODEs are advantageous for joint learning of the dynamics and energy of conservative physical systems, where the learned Hamiltonian vector fields ${\mathfrak{X}}_{H_{\theta }}\left(T^{*}\mathcal{Q}\right)$ are guaranteed to conserve the Hamiltonian $H_{\theta }$ representing the total energy. Lagrangian neural ODEs likewise enable learning the dynamics of conservative physical systems and enable incorporation of dissipative terms [34] but do not directly represent the total energy.

Port-Hamiltonian neural ODEs [3,35] offer further expressiveness: besides a scalar Hamiltonian $H_{\theta }\in C^{\infty }\left(\mathcal{M}\right)$, they offer degrees of freedom in a skew-symmetric $(2,0)$-tensor $J_{\theta }$ (called a Poisson tensor), a positive-definite $(2,0)$-tensor $R_{\theta }$ (called a dissipation tensor), and a linear input map $B_{\theta }\left(x\right):{\mathbb{R}}^k\to T_x\mathcal{M}$. This allows learning the dynamics of non-conservative dynamical systems that can be coupled with known physical systems and control inputs through the input map, while jointly learning the total energy $H_{\theta }$, rate of energy dissipation $R_{\theta }(\mathrm{d}{H}_{\theta },\mathrm{d}{H}_{\theta })$, and externally supplied power (see [57] for an introductory overview). Most physical systems can be represented in a port-Hamiltonian form [58], giving this parametrization a high degree of expressiveness that has been used in dynamics learning [3], control [20,21], and model order reduction [35]. Albeit not investigated in practice, this expressiveness may also be a disadvantage compared to Lagrangian or Hamiltonian neural networks, resulting in overfitting when, e.g., small dissipation terms are learned where there is no dissipation. Generally speaking, choosing the most specific structure-preserving neural network is advised.

3.2. Extensions

The proof of Theorem 2 depended on identifying a suitable augmented manifold ${\mathcal{M}}^{\prime }$, autonomous augmented dynamics $f_{\mathrm{aug}}\in \mathfrak{X}\left({\mathcal{M}}^{\prime }\right)$, and an augmented cost function $C_{\mathrm{aug}}:{\mathcal{M}}^{\prime }\to \mathbb{R}$ that rephrases the cost (43) as a final cost $C_{\mathrm{aug}}\left(x\left(T\right)\right)$ to apply Theorem 1. This approach generalizes to various other scenarios, including different cost terms, augmented neural ODEs on manifolds, and time-dependent parameters, presented in the following.

3.2.1. Nonlinear and Intermittent Cost Terms

We consider here the case of neural ODEs on manifolds of the form (38) with cost (39). This is a generalization of [1], in which intermittent cost terms appear for neural ODEs on ${\mathbb{R}}^n$. For the final and intermittent cost term $F_{\theta }:\mathcal{M}\times \mathcal{M}\times \dots \times \mathcal{M}\to \mathbb{R}$, we denote by ${\mathrm{d}}_k{F}_{\theta }\in T_x^{*}\mathcal{M}$ the differential with respect to the k-th slot and denote $\theta$ as a subscript to avoid confusion. The components of ${\mathrm{d}}_{k}F$ will be denoted ${\displaystyle \frac{\partial F}{{\partial }^k{q}^i}}$. In this case, the parameter gradient is determined by repeated application of Theorem 2:

Theorem 3

(Generalized Adjoint Method on Manifolds). Given the dynamics (38) and the cost (39), the parameter gradient’s components ${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left((x_0,t_0),\theta \right)\in {\mathbb{R}}^{n_{\theta }}$ are computed by

$$\begin{array}{cc}\hfill {\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left((x_0,t_0),\theta \right)=& \left({\displaystyle \frac{\partial F}{\partial \theta }}\right)(\theta ,x\left(T_1\right),x\left(T_2\right),\dots ,x\left(T\right))\hfill \\ \hfill & +{\int }_0^T{\displaystyle \frac{\partial }{\partial \theta }}\left({\lambda }_j{f}_{\theta }^j\left(q\left(s\right)\right)+r(q\left(s\right),\theta ,s)\right)ds.\hfill \end{array}$$

(57)

where the state $x\left(s\right)\in \mathcal{M}$ satisfies (45) and the co-state $\lambda \left(s\right)\in T_{x\left(s\right)}^{*}\mathcal{M}$ satisfies dynamics with discrete updates at times $T_1,\dots ,T_{N-1}$ given by

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{q,i}& ={\displaystyle \frac{\partial }{\partial q^i}}\left({\lambda }_{q,j}{f}_{\theta }^j(q,t)+r(q,\theta ,t)\right);{\lambda }_{q,i}\left(T\right)={\displaystyle \frac{\partial F_{\theta }}{{\partial }^N{q}^i}}(x\left(T_1\right),\dots ,x\left(T\right))\hfill \end{array}$$

(58)

$$\begin{array}{cc}\hfill {\lambda }_i\left(T_{k,-}\right)& ={\lambda }_i\left(T_{k,+}\right)+{\displaystyle \frac{\partial F_{\theta }}{{\partial }^k{q}^i}}\left(x\left(T_1\right),\dots ,x\left(T\right)\right),\hfill \end{array}$$

(59)

with $T_{k,-}$ being the instance after a discrete update at time $T_k$ (recall that co-state dynamics are integrated backwards, so $T_{k,-}<T_k<T_{k,+}$) and $T_{k,+}$ the instance before.

Proof. 

We introduce an augmented manifold ${\mathcal{M}}^{\prime }=\mathcal{M}\times \dots \times \mathcal{M}\times {\mathbb{R}}^{n_{\theta }}\times \mathbb{R}\times \mathbb{R}$ to include N copies of the original state $x\in \mathcal{M}$, parameters $\theta \in {\mathbb{R}}^{n_{\theta }}$, accumulated running cost $L\in \mathbb{R}$, and time $t\in \mathbb{R}$ in the augmented state $x^{\prime }:=(x_1,\dots ,x_N,\theta ,L,t)\in {\mathcal{M}}^{\prime }$. Let

$${\varrho }_{T_i}\left(t\right)=\left\{\begin{array}{c}1t\le T_i\hfill \\ 0t>T_i\hfill \end{array}\right.,$$

(60)

and define the augmented dynamics $f_{\mathrm{aug}}\in \mathfrak{X}\left({\mathcal{M}}^{\prime }\right)$ as

$${\dot{x}}^{\prime }=f_{\mathrm{aug}}\left(x^{\prime }\right)=\left(\begin{array}{c}{\varrho }_{T_1}\left(t\right)f_{\theta }(x_1,t)\\ ⋮\\ {\varrho }_{T_{N-1}}\left(t\right)f_{\theta }(x_{N-1},t)\\ f_{\theta }(x_N,t)\\ 0\\ r(x_N,\theta ,t)\\ 1\end{array}\right),x^{\prime }\left(0\right)=x_0^{\prime }:=\left(\begin{array}{c}{x}_0\\ ⋮\\ x_0\\ x_0\\ \theta \\ 0\\ t_0\end{array}\right).$$

(61)

This is an autonomous system with final state

$$x^{\prime }\left(T\right)=(x\left(T_1\right),\dots ,x\left(T_{N-1}\right),x\left(T\right),\theta ,{\int }_0^{T}r(x,\theta ,s)\mathrm{d}s,T).$$

(62)

Next, define the cost $C_{\mathrm{aug}}:{\mathcal{M}}^{\prime }\to \mathbb{R}$ on the augmented space:

$$C_{\mathrm{aug}}\left(x^{\prime }\right)=F_{\theta }(x_1,\dots ,x_N)+L.$$

(63)

Then Equation (39) can be rewritten as the evaluation of a terminal cost $C_{\mathrm{aug}}\left(x^{\prime }\left(T\right)\right)$:

$$C_{f_{\theta }}^T\left(x_0\right)=(C_{\mathrm{aug}}\circ {\mathrm{\Psi }}_{f_{\mathrm{aug}}}^T)\left(x_0^{\prime }\right).$$

(64)

Apply Equation (26), and split the co-state into ${\lambda }_{q_1},\dots ,{\lambda }_{q_N},{\lambda }_{\theta },{\lambda }_L,{\lambda }_t$; then their components’ dynamics are as follows:

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{q_1,i}& =-{\displaystyle \frac{\partial }{\partial q^i}}\left({\lambda }_{q_1,j}{\varrho }_{T_1}\left(t\right)f_{\theta }^j(q_1,t)\right),{\lambda }_{q_1}\left(T\right)={\displaystyle \frac{\partial F_{\theta }}{{\partial }^{1}q}}(x\left(T_1\right),\dots ,x\left(T\right)),\hfill \\ \hfill ⋮\end{array}$$

(65)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{q_N,i}& =-{\displaystyle \frac{\partial }{{\partial }^N{q}^i}}\left({\lambda }_{q_N,j}{f}_{\theta }^j(q_N,t)+{\lambda }_{L}r(q_N,\theta ,t)\right),{\lambda }_{q_N}\left(T\right)={\displaystyle \frac{\partial F_{\theta }}{{\partial }^{N}q}}(x\left(T_1\right),\dots ,x\left(T\right)),\hfill \end{array}$$

(66)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{\theta ,i}& =-{\displaystyle \frac{\partial }{\partial {\theta }^i}}\left({\lambda }_{q_1,j}{\varrho }_{T_1}\left(t\right)f_{\theta }^j(q_1,t)+\dots +{\lambda }_{q_N,j}{f}_{\theta }^j(q_N,t)+{\lambda }_{L}r(q,\theta ,t)\right),\hfill \\ \hfill {\lambda }_{\theta }\left(T\right)& ={\displaystyle \frac{\partial F_{\theta }}{\partial \theta }}(x\left(T_1\right),\dots ,x\left(T\right)).\hfill \end{array}$$

(67)

We excluded the dynamics of ${\lambda }_t$, which does not appear in any of the other equations, and the constant ${\lambda }_L=1$. Finally, define the cumulative co-state

$${\lambda }_q={\varrho }_{T_1}\left(t\right){\lambda }_{q_1}+\dots +{\lambda }_{q_N}.$$

(68)

Its dynamics at $t\in [0,T]T_1,·,T_{N-1}$ are given by the sum of (65) to (66), letting $q=q_N$:

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{q,i}& ={\dot{\lambda }}_{q_1,i}+\dots +{\dot{\lambda }}_{q_N,i}\hfill \end{array}$$

(69)

$$\begin{array}{cc}\hfill & ={\displaystyle \frac{\partial }{\partial q^i}}\left({\lambda }_{q,j}{f}_{\theta }^j(q,t)+r(q,\theta ,t)\right)\hfill \end{array}$$

(70)

$$\begin{array}{cc}\hfill {\lambda }_q\left(T\right)& ={\displaystyle \frac{\partial F_{\theta }}{{\partial }^{N}q}}(x\left(T_1\right),\dots ,x\left(T\right)),\hfill \end{array}$$

(71)

with discrete jumps (58) accounting for the final conditions of ${\lambda }_{q_1},\dots ,{\lambda }_{q_N}$, and the dynamics (67) can be rewritten as

$${\dot{\lambda }}_{\theta ,i}={\displaystyle \frac{\partial }{\partial {\theta }^i}}\left({\lambda }_{q,j}{f}_{\theta }^j(q,t)+r(q,\theta ,t)\right);{\lambda }_{\theta }\left(T\right)={\displaystyle \frac{\partial F_{\theta }}{\partial \theta }}(x\left(T_1\right),\dots ,x\left(T\right)).$$

(72)

Integrating this from $s=0$ to $s=T$ recovers Equation (57). □

Cost terms of this form are interesting for optimization of, e.g., periodic orbits [59] or trajectories on manifolds, where conditions at multiple checkpoints ${\mathrm{\Psi }}_{f_{\theta }}^{T_i}\left(x_0\right)$ may appear in the cost.

3.2.2. Augmented Neural ODEs on Manifolds and Time-Dependent Parameters

With state $x\in \mathcal{M}$, augmented state $\alpha \in \mathcal{N}$ (not to be confused with $x^{\prime }\in {\mathcal{M}}^{\prime }$), and parameterized ${\phi }_{\theta }:\mathcal{M}\to \mathcal{N}$, augmented neural ODEs on manifolds are neural ODEs on the manifold $\mathcal{M}\times \mathcal{N}$ of the form

$$\left(\begin{array}{c}\dot{x}\\ \dot{\alpha }\end{array}\right)=\left(\begin{array}{c}{f}_{\theta }(x,\alpha )\\ g_{\theta }(x,\alpha )\end{array}\right);\left(\begin{array}{c}x\left(0\right)\\ \alpha \left(0\right)\end{array}\right)=\left(\begin{array}{c}{x}_0\\ {\phi }_{\theta }\left(x_0\right)\end{array}\right).$$

(73)

Time t is not included explicitly in these dynamics, since it can be included in $\alpha$. This case also includes the scenario of time-dependent parameters $\overline{\theta }\left(t\right)$ as part of $\alpha$. As the trajectory cost, we take a final cost

$$C_{f_{\theta },g_{\theta }}^T(x_0,\theta )=F({\mathrm{\Psi }}_{f_{\theta },g_{\theta }}^T\left(x_0,{\phi }_{\theta }\left(x_0\right)\right),\theta ).$$

(74)

This is a generalization of [11,12].

Theorem 4

(Adjoint Method for Augmented Neural ODEs on Manifolds). Given the dynamics (73) and the cost (74), the parameter gradient’s components ${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left((x_0,t_0),\theta \right)\in {\mathbb{R}}^{n_{\theta }}$ are computed by

$$\begin{array}{cc}\hfill {\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta },g_{\theta }}^T\left((x_0,\phi \left(x_0\right)),\theta \right)=& \left({\displaystyle \frac{\partial F}{\partial \theta }}\right)(x\left(T\right),\alpha \left(T\right),\theta )+{\displaystyle \frac{\partial {\phi }^j}{\partial \theta }}{\lambda }_{\alpha ,j}\left(0\right)\hfill \\ \hfill & +{\int }_0^T{\displaystyle \frac{\partial }{\partial \theta }}\left({\lambda }_{x,j}{f}_{\theta }^j\left(q\left(s\right)\right)+{\lambda }_{\alpha ,j}{g}_{\theta }^j\left(q\left(s\right)\right)\right)ds.\hfill \end{array}$$

(75)

where the states $x\left(s\right)\in \mathcal{M},\alpha \left(s\right)\in \mathcal{N}$ satisfy (73) and co-states ${\lambda }_x\left(s\right)\in T_{x\left(s\right)}^{*}\mathcal{M},{\lambda }_{\alpha }\left(s\right)\in T_{\alpha \left(s\right)}^{*}\mathcal{N},$ satisfy, in a local chart $(U,Q)$ on $\mathcal{M}$ and $\overline{U},\overline{Q}$ on $\mathcal{N}$,

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{x,i}& =-{\displaystyle \frac{\partial }{\partial q^i}}\left({\lambda }_{x,j}{f}_{\theta }^j(q,\overline{q},t)+{\lambda }_{\alpha ,j}{g}_{\theta }^j(q,\overline{q},t)\right),{\lambda }_{x,i}\left(T\right)={\displaystyle \frac{\partial F}{\partial q^i}}\left(x\left(T\right),\alpha \left(T\right),\theta \right),\hfill \end{array}$$

(76)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{\alpha ,i}& =-{\displaystyle \frac{\partial }{\partial {\overline{q}}^i}}\left({\lambda }_{x,j}{f}_{\theta }^j(q,\overline{q},t)+{\lambda }_{\alpha ,j}{g}_{\theta }^j(q,\overline{q},t)\right),{\lambda }_{\alpha ,i}\left(T\right)={\displaystyle \frac{\partial F}{\partial {\overline{q}}^i}}\left(x\left(T\right),\alpha \left(T\right),\theta \right).\hfill \end{array}$$

(77)

Proof. 

Define the augmented state space as ${\mathcal{M}}^{\prime }=\mathcal{M}\times \mathcal{N}\times {\mathbb{R}}^{n_{\theta }}$ to include the states $x\in \mathcal{M},\alpha \left(s\right)\in \mathcal{N}$ and parameters $\theta \in {\mathbb{R}}^{n_{\theta }}$ in the augmented state $x^{\prime }:=(x,\alpha ,\theta )\in {\mathcal{M}}^{\prime }$. In addition, define the augmented dynamics $f_{\mathrm{aug}}\in \mathfrak{X}\left({\mathcal{M}}^{\prime }\right)$ as

$${\dot{x}}^{\prime }=f_{\mathrm{aug}}\left(x^{\prime }\right)=\left(\begin{array}{c}{f}_{\theta }(x,\alpha )\\ g_{\theta }(x,\alpha )\\ 0\end{array}\right),x^{\prime }\left(0\right)=x_0^{\prime }:=\left(\begin{array}{c}{x}_0\\ {\phi }_{\theta }\left(x_0\right)\\ \theta \end{array}\right).$$

(78)

This is an autonomous system with final state $x^{\prime }\left(T\right)=(x\left(T\right),\alpha \left(T\right),\theta )$. Next, define the cost $C_{\mathrm{aug}}:{\mathcal{M}}^{\prime }\to \mathbb{R}$ on the augmented space:

$$C_{\mathrm{aug}}\left(x^{\prime }\right)=F(x,\alpha ,\theta ).$$

(79)

Then Equation (43) can be rewritten as the evaluation of a terminal cost $C_{\mathrm{aug}}\left(x^{\prime }\left(T\right)\right)$. The gradient $\mathrm{d}\left(C_{\mathrm{aug}}\circ {\mathrm{\Psi }}_{f_{\mathrm{aug}}}^T\right)$ is given by an application of Equation (26). Split the co-state into ${\lambda }_x,{\lambda }_{\alpha },{\lambda }_{\theta }$; then their components’ dynamics are as follows:

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{x,i}& =-{\displaystyle \frac{\partial }{\partial q^i}}\left({\lambda }_{x,j}{f}_{\theta }^j(q,\overline{q},t)+{\lambda }_{\alpha ,j}{g}_{\theta }^j(q,\overline{q},t)\right),{\lambda }_x\left(T\right)={\displaystyle \frac{\partial F}{\partial q}}(x\left(T\right),\alpha \left(T\right),\theta ),\hfill \end{array}$$

(80)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{\alpha ,i}& =-{\displaystyle \frac{\partial }{\partial {\overline{q}}^i}}\left({\lambda }_{x,j}{f}_{\theta }^j(q,\overline{q},t)+{\lambda }_{\alpha ,j}{g}_{\theta }^j(q,\overline{q},t)\right),{\lambda }_{\alpha ,i}\left(T\right)={\displaystyle \frac{\partial F}{\partial {\overline{q}}^i}}\left(x\left(T\right),\alpha \left(T\right),\theta \right),\hfill \end{array}$$

(81)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_{\theta ,i}& =-{\displaystyle \frac{\partial }{\partial {\theta }^i}}({\lambda }_{x,j}{f}_{\theta }^j(q,\overline{q},t)+{\lambda }_{\alpha ,j}{g}_{\theta }^j(q,\overline{q},t)),{\lambda }_{\theta }\left(T\right)={\displaystyle \frac{\partial F}{\partial \theta }}(x\left(T\right),\alpha \left(T\right),\theta ).\hfill \end{array}$$

(82)

Since $\alpha \left(0\right)={\phi }_{\theta }\left(x_0\right)$ also depends on $\theta$, the total gradient of the cost with respect to $\theta$ is given by

$${\displaystyle \frac{\partial }{\partial {\theta }^i}}{C}_{f_{\theta }}^T\left((x_0,{\phi }_{\theta }\left(x_0\right)),\theta \right)={\lambda }_{\theta ,i}\left(0\right)+{\displaystyle \frac{\partial {\phi }^j}{\partial {\theta }^i}}{\lambda }_{\alpha ,j}\left(0\right).$$

(83)

Integrate (82) to find ${\lambda }_{\theta ,i}\left(0\right)$; then Equation (75) is recovered. □

Augmented neural ODEs are universal function approximators ([25], Chapter 2). Potential applications of augmented neural ODEs on manifolds include, e.g., the optimization of guiding vector fields for path-following of closed or self-intersecting paths [60], where state augmentation sits at the core of formulating singularity-free guiding vector fields for self-intersecting paths. In the same context, discontinuous initializations $g_{\alpha }\left(x_0\right)$ allow globally stabilizing controllers to be represented for topologically non-trivial manifolds (e.g., the sphere $S^2$), where smooth controllers are necessarily not globally stable. A further degenerate application of Theorem 4 is obtained by removing x, i.e., fixing $x=0$ and $f_{\theta }(x,\alpha )=0$ in Equation (73). Then both the dynamics $g_{\theta }\left(\alpha \right)$ and initial condition $\alpha \left(0\right)={\phi }_{\theta }\left(0\right)$ are parameterized by $\theta$, allowing joint optimization of the parameters and initial condition. This is interesting for joint optimization and numerical continuation, e.g., [59].

4. Neural ODEs on Lie Groups

Just as a neural ODE on a manifold is an NN-parameterized vector field in $\mathfrak{X}\left(\mathcal{M}\right)$ (or, including time, $\mathfrak{X}(\mathcal{M}\times \mathbb{R})$), a neural ODE on a Lie group can be seen as a parameterized vector field in $\mathfrak{X}\left(G\right)$ (or $\mathfrak{X}(G\times \mathbb{R})$). Similarly to Equation (38), this results in a dynamic system

$$\dot{g}=f_{\theta }(g,t),g\left(0\right)=g_0.$$

(84)

Yet, Lie groups offer more structure than manifolds: the Lie algebra $\mathfrak{g}$ provides a canonical space to represent tangent vectors, and its dual ${\mathfrak{g}}^{*}$ provides a canonical space to represent the co-state. Similarly, canonical (exponential) charts offer a structure for integrating dynamic systems [41]. Frequently, dynamics on a Lie group induce dynamics on a manifold $\mathcal{M}$: by means of an action

$$\mathrm{\Phi }:G\times \mathcal{M}\to \mathcal{M};(g,x)\mapsto \mathrm{\Phi }(g,x),$$

(85)

evolutions $g\left(t\right)$ induce evolutions $x\left(t\right)=\mathrm{\Phi }(g\left(t\right),x_0)$ on $\mathcal{M}$. This makes neural ODEs on Lie groups interesting in their own right.

In this section, we describe optimizing (41) for the cost

$$C_{f_{\theta }}^T(g_0,\theta )=F({\mathrm{\Psi }}_{f_{\theta }}^T\left(g_0\right),\theta )+{\int }_0^{T}r({\mathrm{\Psi }}_{f_{\theta }}^s\left(g_0\right),\theta ,s)\mathrm{d}s,$$

(86)

with a final cost term F and a running cost term r. We highlight the extrinsic approach and two intrinsic approaches, where one of the latter is particular to Lie groups.

4.1. Extrinsic Neural ODEs on Lie Groups

The extrinsic formulation of neural ODEs on Lie groups was first introduced by [20] and applies ideas of [54] (see also Section 3.1.1). Given $G\subset GL(m,\mathbb{R})$, this formulation treats the dynamic system (84) as a dynamic system on ${\mathbb{R}}^{m^2}$. Denote $\mathrm{vec}:{\mathbb{R}}^{m\times m}\to {\mathbb{R}}^{m^2}$ as an invertible map that stacks the components of an input matrix into a component vector (in canonical coordinates on ${\mathbb{R}}^{m\times m}$ and ${\mathbb{R}}^{m^2}$, though this choice is not required.) and let ${\mathrm{proj}}_G:{\mathbb{R}}^{m\times m}\to G$ be a projection onto $G\subset {\mathbb{R}}^{m\times m}$. Further denote $A_y={\mathrm{vec}}^{-1}\left(y\right)$ and $g_y={\mathrm{proj}}_G\left(A_y\right)$. A lift $f_{\theta }^{\uparrow }(y,t)$ can then be defined as

$$f_{\theta }^{\uparrow }(y,t)=\mathrm{vec}\left(A_y{g}_y^{-1}f(g_y,\theta ,t)\right).$$

(87)

As was the case for extrinsic neural ODEs on manifolds, the cost gradient resulting from this optimization is well-defined and equivalent to any intrinsically defined procedure. However, the dimension $m^2$ of the vectorization can be significantly larger than the intrinsic dimension of the Lie group.

4.2. Intrinsic Neural ODEs on Lie Groups

Theorem 2 directly applies to optimization of neural ODEs on Lie groups, given the local exponential charts (20) and (21) on G. This does not make full use of the available structure on Lie groups. Frequently, dynamical systems are of a left-invariant form (88) or a right-invariant form (89)

$$\begin{array}{cc}\hfill \dot{g}& =g\mathrm{\Lambda }\left({\rho }_{\theta }^L(g,t)\right),\hfill \end{array}$$

(88)

$$\begin{array}{cc}\hfill \dot{g}& =\mathrm{\Lambda }\left({\rho }_{\theta }^R(g,t)\right)g.\hfill \end{array}$$

(89)

Denote $K\left(q\right):T_q{\mathbb{R}}^n\to {\mathbb{R}}^n$ as the derivative of the exponential map (see [21] for details). Then the chart representatives $f_{\theta }^i$ in a local exponential chart $(U_h,Q_h)$ are

$$\begin{array}{cc}\hfill f_{\theta }^{L,i}(q,t)& ={\left(K^{-1}\right)}_j^i\left(q\right){\rho }^{L,j}\left(Q_h^{-1}\left(q\right)\right),\hfill \end{array}$$

(90)

$$\begin{array}{cc}\hfill f_{\theta }^{R,i}(q,t)& ={\left(K^{-1}\right)}_j^i\left(q\right){\mathrm{Ad}}_{Q_h^{-1}\left(q\right)}{\rho }^{R,j}\left(Q_h^{-1}\left(q\right)\right).\hfill \end{array}$$

(91)

Application of Theorem 2 then requires computing ${\displaystyle \frac{\partial }{\partial q^j}}{f}_{\theta }^{L,i}(q,t)$ or ${\displaystyle \frac{\partial }{\partial q^j}}{f}_{\theta }^{R,i}(q,t)$. But this leads to significant computational overhead due to differentiation of the terms ${\left(K^{-1}\right)}_j^i\left(q\right)$ (see [21]). Instead of applying Theorem 2, i.e., expressing dynamics in local charts, the dynamics can also be expressed in the Lie algebra $\mathfrak{g}$. Theorem 1 has a Hamiltonian form, which can be directly transformed into Hamiltonian equations on a Lie group (see also Appendix A). Applying this reasoning to Theorem 2, we arrive at the following form, which foregoes differentiating ${\left(K^{-1}\right)}_j^i\left(x\right)$:

Theorem 5

(Left Generalized Adjoint Method on Matrix Lie Groups). Given the dynamics (88) and the cost (86), or the dynamics (89) with ${\rho }_{\theta }^L(g,t)={\mathit{Ad}}_{g^{-1}}{\rho }_{\theta }^R(g,t)$, the parameter gradient ${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left(g_0\right)$ of the cost is given by the integral equation

$${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left(g_0\right)={\displaystyle \frac{\partial F}{\partial \theta }}(g\left(T\right),\theta )+{\int }_0^T{\displaystyle \frac{\partial }{\partial \theta }}\left({\lambda }_g^{\top }{\rho }_{\theta }^L(g,s)+r(g,\theta ,s)\right)ds,$$

(92)

where the state $g\left(t\right)\in G$ and co-state ${\lambda }_g\left(t\right)\in {\mathbb{R}}^n$ are the solutions of the system of equations

$$\begin{array}{cc}\hfill \dot{g}& =f_{\theta }(g,t),g\left(0\right)=g_0,\hfill \end{array}$$

(93)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_g& =-d_g^L\left({\lambda }_g^{\top }{\rho }_{\theta }^L(g,s)+r(g,\theta ,s)\right)+{\mathit{ad}}_{{\rho }_{\theta }^L(g,t)}^{\top }{\lambda }_g,{\lambda }_g\left(T\right)=d_g^{L}F(g\left(T\right),\theta ).\hfill \end{array}$$

(94)

Proof. 

This is proven in two steps. First, define the time- and parameter-dependent control Hamiltonian $H_c:T^{*}\mathcal{M}\times {\mathbb{R}}^{n_{\theta }}\times \mathbb{R}\to \mathbb{R}$ as

$$H_c(x,\lambda ,\theta ,t)=\lambda \left(f_{\theta }(x,t)\right)+r(x,\theta ,t)={\lambda }_i\left(f_{\theta }^i(q,t)\right)+r(q,\theta ,t).$$

(95)

The equations for the state and co-state dynamics (45) and (46), respectively, of Theorem 2 follow as the Hamiltonian equations on $T^{*}\mathcal{M}$:

$$\begin{array}{cc}\hfill {\dot{q}}^j& ={\displaystyle \frac{\partial H_c}{\partial {\lambda }_j}}=f_{\theta }^j(q,t),\hfill \end{array}$$

(96)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_i& =-{\displaystyle \frac{\partial H_c}{\partial q^i}}=-{\lambda }_j{\displaystyle \frac{\partial }{\partial q^i}}{f}_{\theta }^j(q,t)-{\displaystyle \frac{\partial r}{\partial q^i}}.\hfill \end{array}$$

(97)

And the integral Equation (44) reads

$${\displaystyle \frac{\partial }{\partial \theta }}{C}_{f_{\theta }}^T\left((x_0,t_0),\theta \right)={\displaystyle \frac{\partial F}{\partial \theta }}(x\left(T\right),\theta )+{\int }_0^T{\displaystyle \frac{\partial H_c}{\partial \theta }}\mathrm{d}t.$$

(98)

Second, rewrite the control Hamiltonian (95) on a Lie group G, i.e., $H_c:T^{*}G\times \times {\mathbb{R}}^{n_{\theta }}\times \mathbb{R}\to \mathbb{R}$. By substituting ${\lambda }_g\left(t\right)={\mathrm{\Lambda }}^{*}{L}_g^{*}\lambda \left(t\right)$ (see also Equation (A6)), this induces ${\mathcal{H}}_c:G\times {\mathfrak{g}}^{*}\times {\mathbb{R}}^{n_{\theta }}\times \mathbb{R}\to \mathbb{R}$,

$${\mathcal{H}}_c(g,{\lambda }_g,\theta ,t)={\lambda }_g^{\top }{\rho }_{\theta }^L(g,t)+r(g,\theta ,t).$$

(99)

Finally Hamilton’s equations (96) and (97) are rewritten in their form on a matrix Lie group by means of (A7) and (A8), which recovers Equations (93) and (94):

$$\begin{array}{cc}\hfill \dot{g}& =g\mathrm{\Lambda }\left({\displaystyle \frac{\partial {\mathcal{H}}_c}{\partial {\lambda }_g}}\right),\hfill \end{array}$$

(100)

$$\begin{array}{cc}\hfill {\dot{\lambda }}_g& =-{\mathrm{d}}_g^L{\mathcal{H}}_c+{\mathrm{ad}}_{{\displaystyle \frac{\partial {\mathcal{H}}_c}{\partial {\lambda }_g}}}^{\top }{\lambda }_g.\hfill \end{array}$$

(101)

To find the final condition for ${\lambda }_g$, use that ${\lambda }_g\left(t\right)={\mathrm{\Lambda }}^{*}{L}_g^{*}\lambda \left(t\right)$:

$${\lambda }_g\left(T\right)={\mathrm{\Lambda }}^{*}{L}_g^{*}\lambda \left(T\right)={\mathrm{\Lambda }}^{*}{L}_g^{*}\mathrm{d}F(g\left(T\right),\theta )=d_g^{L}F(g\left(T\right),\theta ).$$

(102)

□

Similar equations also hold on abstract (non-matrix) Lie groups, see [21]. Compared to the extrinsic method of Section 4.1, Theorem 5 has the advantage that the dimension of the co-state ${\lambda }_g$ is as low as possible. Compared to the chart-based approach on Lie groups, Theorem 5 foregoes differentiating through the terms $K_j^i\left(q\right)$, avoiding overhead. Compared to a chart-based approach on manifolds, the choice of charts is also canonical on Lie groups. Although the Lie group approach foregoes many of the pitfalls of intrinsic neural ODEs on manifolds, implementation in existing neural ODE packages is currently cumbersome: the adjoint sensitivity equations (94) have a non-standard form, requiring an adapted dynamics of the co-state $\lambda$, but these equations are rarely intended for modification in existing packages. Packages for geometry-preserving integrators on Lie groups, such as [41], are also not readily available for arbitrary Lie groups.

4.3. Extensions

The proof of Theorem 5 relied on finding a control Hamiltonian formulation for Theorem 2. This approach generalizes to methods in Section 3.2, which rely on the use of Theorem 1. This is because Theorem 1 itself has a Hamiltonian form ([21,54]).

A further straightforward extension of the methods presented in this Section are port-Hamiltonian neural ODEs on Lie groups [20]. In [20], these are systems with a configuration on a Lie group G and momentum on ${\mathfrak{g}}^{*}$. In terms of the theory presented above, such port-Hamiltonian dynamics can be phrased as a dynamic system on a product Lie group $G\times {\mathfrak{g}}^{*}$ (taking vector addition as the group operation on ${\mathfrak{g}}^{*}$), recovering both extrinsic [20] and intrinsic [21] port-Hamiltonian neural ODEs on Lie groups.

5. Discussion

We discuss advantages and disadvantages of the main flavors of the presented formulations for manifold neural ODEs, expanding on the previous sections. We focus on extrinsic (embedding dynamics in ${\mathbb{R}}^N$) and intrinsic (integrating in local charts) formulations. The prior comments can be summarized as follows:

• The extrinsic formulation is readily implemented if the low-dimensional manifold $\mathcal{M}$ and an embedding into ${\mathbb{R}}^N$ are known. This comes at the possible cost of geometric inexactness and a higher dimension of the co-state and sensitivity equations.
• The co-state in the intrinsic formulation has a generally lower dimension, which reduces the dimension of the sensitivity equations. The chart-based formulation also guarantees geometrically exact integration of dynamics. This comes at the mild cost of having to define local charts and chart-transitions.

This dimensionality reduction is unlikely to have a high impact when the manifold $\mathcal{M}$ is known and low-dimensional, e.g., for the sphere $\mathcal{M}=S^2$ or similar manifolds. However, when applying the manifold hypothesis to high-dimensional data, there might be non-trivial latent manifolds for which the embedding is not immediate and where the latent manifold is of a much lower dimension than the embedding data manifold. Then the intrinsic method becomes difficult to avoid. If geometric exactness of the integration is desired, local charts also need to be defined for the extrinsic approach, in which case the intrinsic approach may offer further advantages.

In order to derive neural ODEs on Lie groups, three approaches are possible: the extrinsic and intrinsic formulations on manifolds directly carry over to matrix Lie groups, embedding $G\subset GL(m,\mathbb{R})$ in ${\mathbb{R}}^{m^2}$ or using local exponential charts, respectively. A third option is a novel intrinsic method for neural ODEs on matrix Lie groups, which makes full use of the Lie group structure by phrasing dynamics on $\mathfrak{g}$ (as is more common on Lie groups) and the co-state on ${\mathfrak{g}}^{*}$, avoiding difficulties of the chart-based formalism in differentiating extra terms.

Prior comments on advantages and disadvantages of these flavors can be summarized as follows:

• The extrinsic formulation on matrix Lie groups can come at much higher cost than that on manifolds, since the intrinsic dimension of G can be much lower than $m^2$ and a higher dimension of the co-state and sensitivity equations can be obtained. Geometrically exact integration procedures are more readily available for matrix Lie groups, integrating $\dot{g}$ in local exponential charts.
• The chart-based formulation on matrix Lie groups struggles when dynamics are not naturally phrased in local charts. This is common; dynamics are often more naturally phrased on $\mathfrak{g}$. This was alleviated by an algebra-based formulation on matrix Lie groups. Both are intrinsic approaches that feature co-state dynamics that are as low as possible. However, the algebra-based approach still lacks readily available software implementation.

The authors believe that the algebra-based formulation is more convenient in principle and consider software implementations of the algebra-based approach as possible future work.

In summary, we presented a unified, geometric approach to extend various methods for neural ODEs on ${\mathbb{R}}^N$ to neural ODEs on manifolds and Lie groups. Optimization of neural ODEs on manifolds was based on the adjoint method on manifolds. Given a novel cost function C and neural ODE architecture f, the strategy to present the results in a unified fashion was to identify a suitable augmented manifold ${\mathcal{M}}_{\mathrm{aug}}$, augmented dynamics $f_{\mathrm{aug}}\in \mathfrak{X}\left({\mathcal{M}}_{\mathrm{aug}}\right)$, and cost $C_{\mathrm{aug}}:{\mathcal{M}}_{\mathrm{aug}}\to \mathbb{R}$ such that the original cost function can be rephrased as $C=C_{\mathrm{aug}}\circ {\mathrm{\Psi }}_{f_{\mathrm{aug}}}^T$. To further derive optimization of intrinsic neural ODEs on Lie groups, we found a Hamiltonian formulation of the adjoint method on manifolds and subsequently transformed it into Hamiltonian equations on a matrix Lie group.