Title: Stabilized Neural Differential Equations for Learning Dynamics with Explicit Constraints
Venue: NeurIPS 2023
Accepted: yes
Keywords: neural differential equations

neural ordinary differential equations

constraints

conservation laws

stabilization

dynamical systems

dynamics

scientific machine learning

physics-informed machine learning
Abstract: Many successful methods to learn dynamical systems from data have recently been introduced. However, ensuring that the inferred dynamics preserve known constraints, such as conservation laws or restrictions on the allowed system states, remains challenging. We propose stabilized neural differential equations (SNDEs), a method to enforce arbitrary manifold constraints for neural differential equations. Our approach is based on a stabilization term that, when added to the original dynamics, renders the constraint manifold provably asymptotically stable. Due to its simplicity, our method is compatible with all common neural differential equation (NDE) models and broadly applicable. In extensive empirical evaluations, we demonstrate that SNDEs outperform existing methods while broadening the types of constraints that can be incorporated into NDE training.

Stabilized Neural Differential Equations for Learning Dynamics with Explicit Constraints

Abstract

Many successful methods to learn dynamical systems from data have recently been introduced. However, ensuring that the inferred dynamics preserve known constraints, such as conservation laws or restrictions on the allowed system states, remains challenging. We propose stabilized neural differential equations (SNDEs), a method to enforce arbitrary manifold constraints for neural differential equations. Our approach is based on a stabilization term that, when added to the original dynamics, renders the constraint manifold provably asymptotically stable. Due to its simplicity, our method is compatible with all common neural differential equation (NDE) models and broadly applicable. In extensive empirical evaluations, we demonstrate that SNDEs outperform existing methods while broadening the types of constraints that can be incorporated into NDE training.

1 Introduction

Advances in machine learning have recently spurred hopes of displacing or at least enhancing the process of scientific discovery by inferring natural laws directly from observational data. In particular, there has been a surge of interest in data-driven methods for learning dynamical laws in the form of differential equations directly from data [18, 67, 16, 23, 3, 11]. Assuming there is a ground truth system with dynamics governed by an ordinary differential equation

with u : R → R n and f : R n × R → R n , the question is whether we can learn f from (potentially noisy and irregularly sampled) observations ( t i , u ( t i )) N i =1 .

Neural ordinary differential equations (NODEs) provide a prominent and successful method for this task, which leverages machine learning by directly parameterizing the vector field f of the ODE as a

(a)

(b)

(c)

NDE

Stabilization

SNDE

SNDE Trajectory

neural network [18] (see also Kidger [49] for an overview). A related approach, termed universal differential equations (UDEs) [67], combines mechanistic or process-based model components with universal function approximators, typically also neural networks. In this paper, we will refer to these methods collectively as neural differential equations (NDEs), meaning any ordinary differential equation model in explicit form, where the right-hand side is either partially or entirely parameterized by a neural network. Due to the use of flexible neural networks, NDEs have certain universal approximation properties [72, 75], which are often interpreted as 'in principle an NDE can learn any vector field f ' [30]. While this can be a desirable property in terms of applicability, in typical settings one often has prior knowledge about the dynamical system that should be incorporated.

As in other areas of machine learning - particularly deep learning - inductive biases can substantially aid generalization, learning speed, and stability, especially in the low data regime. Learning dynamics from data is no exception [4]. In scientific applications, physical priors are often not only a natural source of inductive biases, but can even impose hard constraints on the allowed dynamics. For instance, when observing mechanical systems, a popular approach is to directly parameterize either the Lagrangian or Hamiltonian via a neural network [39, 58, 24]. Constraints such as energy conservation can then be 'baked into the model', in the sense that the parameterization of the vector field is designed to only represent functions that satisfy the constraints. Finzi et al. [32] build upon these works and demonstrate how to impose explicit constraints in the Hamiltonian and Lagrangian settings.

In this work, we propose a stabilization technique to enforce arbitrary, even time-dependent manifold constraints for any class of NDEs, not limited to second order mechanical systems and not requiring observations in particular (canonical) coordinates. Our method is compatible with all common differential equation solvers as well as adjoint sensitivity methods. All code is publicly available at https://github.com/white-alistair/Stabilized-Neural-Differential-Equations .

2 Background

A first order neural differential equation is typically given as

where u : R → R n and the vector field f θ : R n × R → R n is at least partially parameterized by a neural network with parameters θ ∈ R d . We restrict our attention to ground truth dynamics f in Equation (1) that are continuous in t and Lipschitz continuous in u such that the existence of a unique (local) solution to the initial value problem is guaranteed by the Picard-Lindelöf theorem. As universal function approximators [44, 25], neural networks f θ can in principle approximate any such f to arbitrary precision, that is to say, the problem of learning f is realizable.

In practice, the parameters θ are typically optimized via stochastic gradient descent by integrating a trajectory ˆ u ( t ) = ODESolve( u 0 , f θ , t ) and taking gradients of the loss L ( u, ˆ u ) with respect to θ . Gradients of trajectories can be computed using adjoint sensitivity analysis ( optimize-then-discretize ) or automatic differentiation of the solver operations ( discretize-then-optimize ) [18, 49]. While these approaches have different (dis)advantages [59, 49], we use the adjoint sensitivity method due to reportedly improved stability [18]. We use the standard squared error loss L ( u, ˆ u ) = ∥ u -ˆ u ∥ 2 2 throughout.

We focus on NODEs as they can handle arbitrary nonlinear vector fields f and outperform traditional ODE parameter estimation techniques. In particular, they do not require a pre-specified parameterization of f in terms of a small set of semantically meaningful parameters. The original NODE model [18] has quickly been extended to augmented neural ODEs (ANODEs), with improved universal approximation properties [30] and performance on second order systems [62]. Further extensions of NODEs include support for irregularly-sampled observations [70] and partial differential equations [35], Bayesian NODEs [26], universal differential equations [67], and more-see, e.g., Kidger [49] for an overview. All such variants fall under our definition of neural differential equations and can in principle be stabilized using our method. We demonstrate this empirically by stabilizing both standard NODEs as well as ANODEs and UDEs in this paper. Finally, we note that all experiments with second order structure are implemented as second order neural ODEs due to reportedly improved generalization compared to first order NODEs [62, 40].

NODEs were originally introduced as the infinite depth limit of residual neural networks (typically for classification), where there is no single true underlying dynamic law, but rather 'some' vector field f is learned that allows subsequent (linear) separation of the inputs into different classes. A number of techniques have been introduced to restrict the number of function evaluations needed during training in order to improve efficiency, which also typically results in relatively simple learned dynamics [31, 38, 47, 63, 50]. These are orthogonal to our method and are mentioned here for completeness, as they could also be viewed as 'regularized' or 'stabilized' NODEs from a different perspective.

3 Related Work

Hamiltonian and Lagrangian Neural Networks. A large body of related work has focused on learning Hamiltonian or Lagrangian dynamics via architectural constraints. By directly parameterizing and learning the underlying Hamiltonian or Lagrangian, rather than the equations of motion, first integrals corresponding to symmetries of the Hamiltonian or Lagrangian may be (approximately) conserved automatically.

Hamiltonian Neural Networks (HNNs) [39] assume second order Hamiltonian dynamics u ( t ) , where u ( t ) = ( q ( t ) , p ( t )) is measured in canonical coordinates, and directly parameterize the Hamiltonian H (instead of f ) from which the vector field can then be derived via f ( q, p ) = ( ∂ H ∂p , -∂ H ∂q ) . Symplectic ODE-Net (SymODEN) [76] takes a similar approach but makes the learned dynamics symplectic by construction, as well as allowing more general coordinate systems such as angles or velocities (instead of conjugate momenta). HNNs have also been made agnostic to the coordinate system altogether by modeling the underlying coordinate-free symplectic two-form directly [19], studied extensively with respect to the importance of symplectic integrators [77], and adapted specifically to robotic systems measured in terms of their SE(3) pose and generalized velocity [29]. Recently, Gruver et al. [40] showed that what makes HNNs effective in practice is not so much their built-in energy conservation or symplectic structure, but rather that they inherently assume that the system is governed by a single second order differential equation ('second order bias'). Chen et al. [20] provide a recent overview of learning Hamiltonian dynamics using neural architectures.

A related line of work is Lagrangian Neural Networks (LNNs), which instead assume second order Lagrangian dynamics and parameterize the inertia matrix and divergence of the potential [58], or indeed any generic Lagrangian function [24], with the dynamics f then uniquely determined via the Euler-Lagrange equations. While our own work is related to HNNs and LNNs in purpose, it is more general, being applicable to any system governed by a differential equation. In addition, where we consider fundamental conservation laws as constraints, e.g. conservation of energy, we assume the conservation law is known in closed-form.

The work most closely related to ours is by Finzi et al. [32], who demonstrate how to impose explicit constraints on HNNs and LNNs. The present work differs in a number of ways with the key advances of our approach being that SNDEs (a) are applicable to any type of ODE, allowing us to go beyond second order systems with primarily Hamiltonian or Lagrangian type constraints, (b) are compatible with hybrid models, i.e., the UDE approach where part of the dynamics is assumed to be known and only the remaining unknown part is learned while still constraining the overall dynamics, and (c) can incorporate any type of manifold constraint. Regarding (c), we can for instance also enforce time-dependent first integrals, which do not correspond to constants of motion or conserved quantities arising directly from symmetries in the Lagrangian.

Continuous Normalizing Flows on Riemannian Manifolds. Another large body of related work aims to learn continuous normalizing flows (CNFs) on Riemannian manifolds. Since CNFs transform probability densities via a NODE, many of these methods work in practice by constraining NODE trajectories to a pre-specified Riemannian manifold. For example, Lou et al. [57] propose 'Neural Manifold ODE', a method that directly adjusts the forward mode integration and backward mode adjoint gradient computation to ensure that the trajectory is confined to a given manifold. This approach requires a new training procedure and relies on an explicit chart representation of the manifold. The authors limit their attention to the hyperbolic space H 2 and the sphere S 2 , both of which are prototypical Riemannian manifolds with easily derived closed-form expressions for the chart. Many similar works have been proposed [37, 14, 68, 60], but all are typically limited in scope to model Riemannian manifolds such as spheres, hyperbola, and tori. Extending these approaches to the manifolds studied by us, which can possess nontrivial geometries arising from arbitrary constraints, is not straightforward. Nonetheless, we consider this an interesting opportunity for further work.

Riemannian Optimization. A vast literature exists for optimization on Riemannian manifolds [2, 53, 61]. In deep learning, in particular, orthogonality constraints have been used to avoid the vanishing and exploding gradients problem in recurrent neural networks [5, 54, 1, 48] and as a form of regularization for convolutional neural networks [9]. Where these methods differ crucially from this paper is that they seek to constrain neural network weights to a given matrix manifold, while our aim is to constrain trajectories of an NDE, the weights of which are not themselves directly constrained.

Constraints via Regularization. Instead of adapting the neural network architecture to satisfy certain properties by design, Lim and Kasim [55] instead manually craft different types of regularization terms that, when added to the loss function, aim to enforce different constraints or conservation laws. While similar to our approach, in that no special architecture is required for different constraints, the key difference is that their approach requires crafting specific loss terms for different types of dynamics. Moreover, tuning the regularization parameter can be rather difficult in practice.

In this work, we vastly broaden the scope of learning constrained dynamics by demonstrating the effectiveness of our approach on both first and second order systems, including chaotic and nonchaotic as well as autonomous and non-autonomous examples. We cover constraints arising from holonomic restrictions on system states, conservation laws, and constraints imposed by controls.

4 Stabilized Neural Differential Equations

General approach. Given m<n explicit constraints, we require that solution trajectories of the NDE in Equation (2) are confined to an ( n -m ) -dimensional submanifold of R n defined by

where g : R n × R → R m is a smooth function with 0 ∈ R m being a regular value of g . 1 In other words, we have an NDE on a manifold

Any non-autonomous system can equivalently be represented as an autonomous system by adding time as an additional coordinate with constant derivative 1 and initial condition t 0 = 0 . For ease of notation, and without loss of generality, we will only consider autonomous systems in the rest of this section. However, we stress that our method applies equally to autonomous and non-autonomous systems.

While methods exist for constraining neural network outputs to lie on a pre-specified manifold, the added difficulty in our setting is that we learn the vector field f but constrain the solution trajectory u that solves a given initial value problem for the ODE defined by f . Inspired by Chin [21], we propose the following stabilization of the vector field in Equation (4)

where γ ≥ 0 is a scalar parameter of our method and F : R n → R n × m is a so-called stabilization matrix . We call Equation (5) a stabilized neural differential equation (SNDE) and say that it is

stabilized with respect to the invariant manifold M . We illustrate the main idea in Figure 1; while even small deviations in the unstabilized vector field (left) can lead to (potentially accumulating) constraint violations, our stabilization (right) adjusts the vector field near the invariant manifold to render it asymptotically stable, all the while leaving the vector field on M unaffected.

Our stabilization approach is related to the practice of index reduction of differential algebraic equations (DAEs). We refer the interested reader to Appendix A for a brief overview of these connections.

Theoretical guarantees. First, note that ultimately we still want f θ to approximate the assumed ground truth dynamics f . However, Equation (5) explicitly modifies the right-hand side of the NDE. The following theorem provides necessary and sufficient conditions under which f θ can still learn the correct dynamics when using a different right-hand side.

Theorem 1 (adapted from Chin [21]) . Consider an NDE

on an invariant manifold M = { u ∈ R n ; g ( u ) = 0 } . A vector field ˙ u = h θ ( u ) admits all solutions of Equation (6) on M if and only if h θ | M = f θ | M .

Since g ( u ) = 0 on M , the second term on the right-hand side of Equation (5) vanishes on M . Therefore the SNDE Equation (5) admits all solutions of the constrained NDE Equation (4). Next, we will show that, under mild conditions, the additional stabilization term in Equation (5) 'nudges' the solution trajectory to lie on the constraint manifold such that M is asymptotically stable.

Theorem 2 (adapted from Chin [21] 2 ) . Suppose the stabilization matrix F ( u ) is chosen such that the matrix G ( u ) F ( u ) , where G ( u ) = g u is the Jacobian of g at u , is symmetric positive definite with the smallest eigenvalue λ ( u ) satisfying λ ( u ) > λ 0 > 0 for all u . Assume further that there is a positive number γ 0 such that

for all u near M . Then the invariant manifold M is asymptotically stable in the SNDE Equation (5) if γ ≥ γ 0 /λ 0 .

Proof. Consider the Lyapunov function V ( u ) = 1 2 g T ( u ) g ( u ) . Then (omitting arguments)

where we substitute Equation (5) for ˙ u . With Equation (7), we have g T Gf θ ≤ γ 0 g T g , and since the eigenvalues of GF are assumed to be at least λ 0 > 0 , we have g T GFg ≥ λ 0 g T g . Hence

so the manifold M is asymptotically stable whenever γ 0 -γλ 0 ≤ 0

.

When f θ ( u ) and g ( u ) are given, M is asymptotically stable in the SNDE Equation (5) as long as

To summarize, the general form of the SNDE in Equation (5) has the following important properties:

The SNDE admits all solutions of the constrained NDE Equation (4) on M .

M is asymptotically stable in the SNDE for sufficiently large values of γ .

The stabilization parameter γ , with units of inverse time, determines the rate of relaxation to the invariant manifold; intuitively, it is the strength of the 'nudge towards M ' experienced by a trajectory. Here, γ is neither a Lagrangian parameter (corresponding to a constraint on θ ), nor a regularization parameter (to overcome an ill-posedness by regularization). Therefore, there is no 'correct' value for γ . In particular, Theorem 1 holds for all γ , while Theorem 2 only requires γ to be 'sufficiently large'.

In the limit γ → ∞ , the SNDE in Equation (5) is equivalent to a Hessenberg index-2 DAE (see Appendix A for more details).

Practical implementation. This leaves us to find a concrete instantiation of the stabilization matrix F ( u ) that should (a) satisfy that F ( u ) G ( u ) is symmetric positive definite with the smallest eigenvalue bounded away from zero near M , (b) be efficiently computable, and (c) be compatible with gradient-based optimization of θ as part of an NDE. In our experiments, we use the Moore-Penrose pseudoinverse of the Jacobian of g at u as the stabilization matrix,

Let us analyze the properties of this choice. Regarding the requirements (b) and (c), the pseudoinverse can be computed efficiently via a singular value decomposition with highly optimized implementations in all common numerical linear algebra libraries (including deep learning frameworks) and does not interfere with gradient-based optimization. In particular, the computational cost for the pseudoinverse is O ( m 2 n ) , i.e., it scales well with the problem size. The quadratic scaling in the number of constraints is often tolerable in practice, since the number of constraints is typically small (in Appendix B, we discuss faster alternatives to the pseudoinverse that can be used, for example, when there are many constraints). Moreover, the Jacobian G ( u ) of g can be obtained via automatic differentiation in the respective frameworks.

Regarding requirement (a), the pseudoinverse G + ( u ) is an orthogonal projection onto the tangent space T u M of the manifold at u . Hence, locally in a neighborhood of u ∈ M , we consider the stabilization matrix as a projection back onto the invariant manifold M (see Figure 1). In particular, G ( u ) has full rank for u ∈ M and G + G = G T ( GG T ) -1 G is symmetric and positive definite near M . From here on, we thus consider the following specific form for the SNDE in Equation (5),

5 Results

We now demonstrate the effectiveness of SNDEs on examples that cover autonomous first and second order systems with either a conserved first integral of motion or holonomic constraints, a non-autonomous first order system with a conserved quantity, a non-autonomous controlled first order system with a time-dependent constraint stemming from the control, and a chaotic second order system with a conservation law. To demonstrate the flexibility of our approach with respect to the variant of underlying NDE model, for the first three experiments we include both vanilla NODEs and augmented NODEs as baselines and apply stabilization to both (with the stabilized variants labeled SNODE and SANODE, respectively).

As a metric for the predicted state ˆ u ( t ) versus ground truth u ( t ) , we use the relative error ∥ u ( t ) -ˆ u ( t ) ∥ 2 / ∥ u ( t ) ∥ 2 , averaged over many trajectories with independent initial conditions. As a metric for the constraints, we compute analogous relative errors for g ( u ) .

In Appendix C, we further demonstrate empirically that SNDEs are insensitive to the specific choice of γ over a large range (beyond a minimum value, consistent with Theorem 2). We also provide more intuition about choosing γ and the computational implications of this choice. In practice, we find that SNDEs are easy to use across a wide variety of settings with minimal tuning and incur only moderate training overhead compared to vanilla NDEs. In some cases, SNDEs are even computationally cheaper than vanilla NDEs at inference time despite the cost of computing the pseudoinverse. Finally, Appendix D provides results of additional experiments not included here due to space constraints.

5.1 Two-Body Problem [second order, autonomous, non-chaotic, conservation law]

The motion of two bodies attracting each other with a force inversely proportional to their squared distance (e.g., gravitational interaction in the non-relativistic limit) can be written as

where one body is fixed at the origin and x, y are the (normalized) Cartesian coordinates of the other body in the plane of its orbit [43]. We stabilize the dynamics with respect to the conserved angular momentum L , yielding

y

y

x

-1.5 -1.0 -0.5

0.0

0.5

-1.0

-0.5

0.0

0.5

1.0

Time (seconds)

0

50

100

Relative Error (State)

10 -5

10 0

10 5

Time (seconds)

0

50

100

Relative Error (Ang. Mom.)

10 -5

10 0

10 5

y

1

y

2

3

Time (seconds)

0

400

800

Relative Error (State)

10 -5

10 0

10 5

Time (seconds)

0

400

800

Relative Error (Casimir)

10 -5

10 0

10 5

Time (seconds)

0

5

10

15

v

1

[

V

]

-2

-1

0

1

2

Time (seconds)

0

40

80

Relative Error (State)

10 -5

10 0

10 5

Time (seconds)

0

40

80

Relative Error (Energy)

10 -5

10 0

10 5

(a)

(b)

(c)

(d)

(e)

(f)

(g)

(h)

(i)

Ground Truth

NODE

ANODE

SNODE

SANODE

-0.2

0.0

0.2

-0.2

0.0

0.2

0.97

0.99

1.01

where L 0 is the initial value of L . We train on 40 trajectories with initial conditions ( x, y, ˙ x, ˙ y ) = (1 -e, 0 , 0 , √ 1 -e / 1+ e ) , where the eccentricity e is sampled uniformly via e ∼ U (0 . 5 , 0 . 7) . Each trajectory consists of a single period of the orbit sampled with a timestep of ∆ t = 0 . 1 .

The top row of Figure 2 shows that SNODEs, ANODEs and SANODEs all achieve stable long-term prediction over multiple orbits, while unstabilized NODEs diverge exponentially from the correct orbit.

5.2 Motion of a Rigid Body [first order, autonomous, non-chaotic, holonomic constraint]

The angular momentum vector y = ( y 1 , y 2 , y 3 ) T of a rigid body with arbitrary shape and mass distribution satisfies Euler's equations of motion,

where the coordinate axes are the principal axes of the body, I 1 , I 2 , I 3 are the principal moments of inertia, and the origin of the coordinate system is fixed at the body's centre of mass [43]. The motion of y conserves the Casimir function C ( y ) = 1 2 ( y 2 1 + y 2 2 + y 2 3 ) , which is equivalent to conservation of angular momentum in the orthogonal body frame and constitutes a holonomic constraint on the allowed states of the system. We therefore have the manifold

We train on 40 trajectories with initial conditions ( y 1 , y 2 , y 3 ) = (cos( ϕ ) , 0 , sin( ϕ )) , where ϕ is drawn from a uniform distribution ϕ ∼ U (0 . 5 , 1 . 5) . Each trajectory consists of a 15 second sample with a timestep of ∆ t = 0 . 1 seconds.

The middle row of Figure 2 demonstrates that, unlike vanilla NODEs and ANODEs, SNODEs and SANODEs are constrained to the sphere and stabilize the predicted dynamics over a long time horizon in this first order system.

5.3 DC-to-DC Converter [first order, non-autonomous, non-chaotic, conservation law]

We now consider an idealized DC-to-DC converter [52, 33], illustrated in Figure 3, with dynamics

where v 1 , v 2 are the state voltages across capacitors C 1 , C 2 , respectively, i 3 is the state current across an inductor L 3 , and u ∈ { 0 , 1 } is a control input (a switch) that can be used to transfer energy between the two capacitors via the inductor. The total energy in the circuit, E = 1 2 ( C 1 v 2 1 + C 2 v 2 2 + L 3 i 2 3 ) , is conserved, yielding the manifold

u

L

3

1

C

2

0

C

1

We train on 40 trajectories integrated over 10 seconds with a timestep of ∆ t = 0 . 1 seconds, where C 1 = 0 . 1 , C 2 = 0 . 2 , L 3 = 0 . 5 , and a switching period of 3 seconds, i.e., the switch is toggled every 1 . 5 seconds. The initial conditions for ( v 1 , v 2 , i 3 ) are each drawn independently from a uniform distribution U (0 , 1) .

The bottom row of Figure 2 shows the voltage across C 1 over multiple switching events (g), with the NODE and ANODE quickly accumulating errors every time the switch is applied, while the SNODE and SANODE remain accurate for longer. Panels (h,i) show the familiar exponentially accumulating errors for vanilla NODE and ANODE, versus stable relative errors for SNODE and SANODE.

5.4 Controlled Robot Arm [first order, non-autonomous, non-chaotic, time-dependent control]

Next, we apply SNDEs to solve a data-driven inverse kinematics problem [64], that is, learning the dynamics of a robot arm that satisfy a prescribed path p ( t ) . We consider an articulated robot arm consisting of three connected segments of fixed length 1, illustrated in Figure 4(a). Assuming one end of the first segment is fixed at the origin and the robot arm is restricted to move in a plane, the endpoint e ( θ ) of the last segment is given by

where θ j is the angle of the j -th segment with respect to the horizontal and θ = ( θ 1 , θ 2 , θ 3 ) . The problem consists of finding the motion of the three segments θ ( t ) such that the endpoint e ( θ ) follows

Time (seconds)

100

102

104

e

(

𝜃

)

0.70

0.75

0.80

Time (seconds)

0

250

500

Relative Error (Endpoint)

10 -5

10 -4

10 -3

10 -2

10 -1

(a)

(b)

(c)

Ground Truth

NODE

SNODE

a prescribed path p ( t ) in the plane, i.e., e ( θ ) = p ( t ) . Minimizing || ˙ θ ( t ) || , it can be shown [42] that the optimal path satisfies

where e ′ is the Jacobian of e . These will be our ground truth equations of motion.

We stabilize the SNODE with respect to the (time-dependent) manifold

In particular, we prescribe the path

where e 0 is the initial position of the endpoint, such that e ( θ ) traces a line back and forth on the x -axis. We train on 40 trajectories of duration 5 seconds, with timestep ∆ t = 0 . 1 and initial conditions ( θ 1 , θ 2 , θ 3 ) = ( θ 0 , -θ 0 , θ 0 ) , where θ 0 is drawn from a uniform distribution θ 0 ∼ U ( π/ 4 , 3 π/ 8) . Additionally, we provide the network with ˙ p , the time derivative of the prescribed control.

Figure 4 shows that the unconstrained NODE drifts substantially from the prescribed path during a long integration, while the SNODE implements the control to a high degree of accuracy and without drift.

5.5 Double Pendulum [second order, autonomous, chaotic, conservation law]

Finally, we apply stabilization to the chaotic dynamics of the frictionless double pendulum system. The total energy E of the system is conserved [7], yielding the manifold,

where θ i is the angle of the i -th arm with the vertical and ω i = ˙ θ i . We refer the reader to Arnold [7] (or the excellent Wikipedia entry) for the lengthy equations of motion and an expression for the total energy. For simplicity we take m 1 = m 2 = 1kg , l 1 = l 2 = 1m , and g = 9 . 81 ms -2 . We train on 40 trajectories, each consisting of 10 seconds equally sampled with ∆ t = 0 . 05 , and with initial conditions ( θ 1 , θ 2 , ω 1 , ω 2 ) = ( ϕ, ϕ, 0 , 0) , where ϕ is drawn randomly from a uniform distribution ϕ ∼ U ( π/ 4 , 3 π/ 4) . We emphasize that this is a highly limited amount of data when it comes to describing the chaotic motion of the double pendulum system, intended to highlight the effect of stabilization in the low-data regime.

Figure 5(a-b) shows that, while initially the SNODE only marginally outperforms the vanilla NODE in terms of the relative error of the state, the longer term relative error in energy is substantially

Time (seconds)

0.0

2.5

5.0

Relative Error (State)

10 -2.5

10 -2.0

10 -1.5

10 -1.0

10 -0.5

Time (seconds)

0

250

500

Relative Error (Energy)

10 -4

10 -2

10 0

𝛾

0

1

2

4

8

0.1

0.2

0.3

0.4

(a)

(b)

(c)

NODE

SNODE

larger for NODE than for SNODE. A certain relative error in state is indeed unavoidable for chaotic systems.

In addition to predicting individual trajectories of the double pendulum, we also consider an additional important task: learning the invariant measure of this chaotic system. This can be motivated by analogy with climate predictions, where one also focuses on long-term prediction of the invariant measure of the system, as opposed to predicting individual trajectories in the sense of weather forecasting, which must break down after a short time due to the highly chaotic underlying dynamics. Motivated by hybrid models of the Earth's climate [36, 71], we choose a slightly different training strategy than before, namely a hybrid setup in line with the UDE approach mentioned above. In particular, the dynamics of the first arm ¨ θ 1 are assumed known, while the dynamics of the second arm ¨ θ 2 are learned from data. We train on a single trajectory of duration 60 seconds with ∆ t = 0 . 05 . For each trained model, we then integrate ten trajectories of duration one hour - far longer than the observed data. An invariant measure is estimated from each long trajectory (see Appendix E) and compared with the ground truth using the Hellinger distance.

Figure 5(c) shows that, as γ is increased, our ability to accurately learn the double pendulum's invariant measure increases dramatically due to stabilization, demonstrating that the 'climate' of this system is captured much more accurately by the SNDE than by the NDE baseline.

6 Conclusion

We have introduced stabilized neural differential equations (SNDEs), a method for learning dynamical systems from observational data subject to arbitrary explicit constraints. Our approach is based on a stabilization term that is cheap to compute and provably renders the invariant manifold asymptotically stable while still admitting all solutions of the vanilla NDE on the invariant manifold. Key benefits of our method are its simplicity and generality, making it compatible with all common NDE methods without requiring further architectural changes. Crucially, SNDEs allow entirely new types of constraints, such as those arising from known conservation laws and controls, to be incorporated into neural differential equation models. We demonstrate their effectiveness across a range of settings, including first and second order systems, autonomous and non-autonomous systems, with constraints stemming from holonomic constraints, conserved first integrals of motion, as well as time-dependent restrictions on the system state. SNDEs are robust with respect to the only tuneable parameter and incur only moderate computational overhead compared to vanilla NDEs.

The current key limitations and simultaneously interesting directions for future work include adapting existing methods for NODEs on Riemannian manifolds to our setting, generalizations to partial differential equations, allowing observations and constraints to be provided in different coordinates, and scaling the method to high-dimensional settings such as learning dynamics from pixel observations, for example in fluid dynamics or climate modeling. Finally, we emphasize that high-dimensional, nonlinear dynamics may not be identifiable from just a small number of solution trajectories. Hence, care must be taken when using learned dynamics in high-stakes scenarios (e.g., human robot interactions), especially when going beyond the training distribution.

Acknowledgments

This work was funded by the Volkswagen Foundation. The authors would like to thank Christof Schötz, Philipp Hess and Michael Lindner for many insightful discussions while preparing this work, as well as the anonymous reviewers whose thoughtful feedback significantly improved the paper. The authors also thank the creators and maintainers of the Julia programming language [13] and the opensource packages DifferentialEquations.jl [66], Flux.jl [46], Zygote.jl [45], ComplexityMeasures.jl [41] and Makie.jl [27], all of which were essential in preparing this work. The authors gratefully acknowledge the European Regional Development Fund (ERDF), the German Federal Ministry of Education and Research and the Land Brandenburg for supporting this project by providing resources on the high performance computer system at the Potsdam Institute for Climate Impact Research.

References

Pierre Ablin and Gabriel Peyré. Fast and accurate optimization on the orthogonal manifold without retraction. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , pages 5636-5657. PMLR, 2022.

Pierre-Antoine Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization Algorithms on Matrix Manifolds . Princeton University Press, 2008.

Hananeh Aliee, Fabian J. Theis, and Niki Kilbertus. Beyond predictions in neural ODEs: Identification and interventions. arXiv preprint arXiv:2106.12430 , 2021.

Hananeh Aliee, Till Richter, Mikhail Solonin, Ignacio Ibarra, Fabian J. Theis, and Niki Kilbertus. Sparsity in continuous-depth neural networks. In Advances in Neural Information Processing Systems , pages 901-914, 2022.

Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In Proceedings of The 33rd International Conference on Machine Learning , pages 1120-1128. PMLR, 2016.

Ludwig Arnold, Christopher K.R.T. Jones, Konstantin Mischaikow, Geneviève Raugel, and Ludwig Arnold. Random Dynamical Systems . Springer, 1995.

Vladimir Igorevich Arnold. Mathematical Methods of Classical Mechanics . Springer New York, 2nd edition, 1989.

Uri M. Ascher and Linda R. Petzold. Computer Methods for Ordinary Differential Equations and Differential-Algebraic Equations . Society for Industrial and Applied Mathematics, 1998.

Nitin Bansal, Xiaohan Chen, and Zhangyang Wang. Can we gain more from orthogonality regularizations in training deep networks? In Advances in Neural Information Processing Systems , pages 4261-4271, 2018.

J. Baumgarte. Stabilization of constraints and integrals of motion in dynamical systems. Computer Methods in Applied Mechanics and Engineering , 1(1):1-16, 1972.

Sören Becker, Michal Klein, Alexander Neitz, Giambattista Parascandolo, and Niki Kilbertus. Predicting ordinary differential equations with transformers. In Proceedings of the 40th International Conference on Machine Learning , pages 1978-2002. PMLR, 2023.

Heli Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Maximillian Nickel, Aditya Grover, Ricky T. Q. Chen, and Yaron Lipman. Matching normalizing flows and probability paths on manifolds. In Proceedings of the 39th International Conference on Machine Learning , pages 1749-1763. PMLR, 2022.

Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral B Shah. Julia: A fresh approach to numerical computing. SIAM Review , 59(1):65-98, 2017.

Joey Bose, Ariella Smofsky, Renjie Liao, Prakash Panangaden, and Will Hamilton. Latent variable modelling with hyperbolic normalizing flows. In Proceedings of the 37th International Conference on Machine Learning , pages 1045-1055. PMLR, 2020.

Kathryn E. Brenan, Stephen L. Campbell, and Linda R. Petzold. Numerical Solution of Initial-Value Problems in Differential-Algebraic Equations . Society for Industrial and Applied Mathematics, 1995.

Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences , 113(15):3932-3937, 2016.

Mickaël D. Chekroun, Eric Simonnet, and Michael Ghil. Stochastic climate dynamics: Random attractors and time-dependent invariant measures. Physica D: Nonlinear Phenomena , 240(21): 1685-1700, 2011.

Ricky T.Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems , pages 6571-6583, 2018.

Yuhan Chen, Takashi Matsubara, and Takaharu Yaguchi. Neural symplectic form: learning Hamiltonian equations on general coordinate systems. pages 16659-16670, 2021.

Zhijie Chen, Mingquan Feng, Junchi Yan, and Hongyuan Zha. Learning neural Hamiltonian dynamics: A methodological overview. arXiv preprint arXiv:2203.00128 , 2022.

Hong Sheng Chin. Stabilization Methods for Simulations of Constrained Multibody Dynamics . PhD thesis, University of British Columbia, 1995. URL https://open.library.ubc.ca/ collections/ubctheses/831/items/1.0080031 .

Harald Cramer. Mathematical Methods of Statistics . Princeton University Press, 1999.

Miles Cranmer. Interpretable machine learning for science with PySR and SymbolicRegression.jl. arXiv preprint arXiv:2305.01582 , 2023.

Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.04630 , 2020.

George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Control, Signals and Systems , 2(4):303-314, 1989.

Raj Dandekar, Karen Chung, Vaibhav Dixit, Mohamed Tarek, Aslan Garcia-Valadez, Krishna Vishal Vemula, and Christopher Rackauckas. Bayesian neural ordinary differential equations. arXiv preprint arXiv:2012.07244 , 2020.

Simon Danisch and Julius Krumbiegel. Makie.jl: Flexible high-performance data visualization for Julia. Journal of Open Source Software , 6(65):3349, 2021.

David Diego, Kristian Agasøster Haaga, and Bjarte Hannisdal. Transfer entropy computation using the Perron-Frobenius operator. Phys. Rev. E , 99:042212, 2019.

Thai P. Duong and Nikolay A. Atanasov. Hamiltonian-based neural ODE networks on the SE(3) manifold for dynamics learning and control. In Proceedings of Robotics: Science and Systems , 2021.

Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural ODEs. In Advances in Neural Information Processing Systems , pages 3140-3150, 2019.

Chris Finlay, Joern-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ODE: the world of Jacobian and kinetic regularization. In Proceedings of the 37th International Conference on Machine Learning , pages 3154-3164. PMLR, 2020.

Marc Finzi, Ke Alexander Wang, and Andrew G Wilson. Simplifying Hamiltonian and Lagrangian neural networks via explicit constraints. In Advances in Neural Information Processing Systems , pages 13880-13889, 2020.

Simone Fiori. Manifold calculus in system theory and control-fundamentals and first-order systems. Symmetry , 13(11), 2021.

C. William Gear. Differential-algebraic equation index transformations. SIAM Journal on Scientific and Statistical Computing , 9(1):39-47, 1988.

Maximilian Gelbrecht, Niklas Boers, and Jürgen Kurths. Neural partial differential equations for chaotic systems. New Journal of Physics , 23(4):043005, 2021.

Maximilian Gelbrecht, Alistair White, Sebastian Bathiany, and Niklas Boers. Differentiable programming for earth system modeling. Geoscientific Model Development , 16(11):3123-3135, 2023.

Mevlana C. Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on Riemannian manifolds. arXiv preprint arXiv:1611.02304 , 2016.

Arnab Ghosh, Harkirat Behl, Emilien Dupont, Philip Torr, and Vinay Namboodiri. STEER: Simple temporal regularization for neural ODEs. In Advances in Neural Information Processing Systems , pages 14831-14843, 2020.

Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems , pages 15379-15389, 2019.

Nate Gruver, Marc Anton Finzi, Samuel Don Stanton, and Andrew Gordon Wilson. Deconstructing the inductive biases of Hamiltonian neural networks. In The Tenth International Conference on Learning Representations , 2022.

Kristian Agasøster Haaga and George Datseris. ComplexityMeasures.jl. Zenodo .

Ernst Hairer. Solving differential equations on manifolds. 2011. URL https://www.unige. ch/~hairer/poly-sde-mani.pdf .

Ernst Hairer, Christian Lubich, and Gerhard Wanner. Geometric Numerical Integration . Springer-Verlag, Berlin, 2nd edition, 2006.

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural Networks , 2(5):359-366, 1989.

Michael Innes. Don't unroll adjoint: Differentiating SSA-form programs. CoRR , abs/1810.07951, 2018.

Michael Innes, Elliot Saba, Keno Fischer, Dhairya Gandhi, Marco Concetto Rudilosso, Neethu Mariya Joy, Tejan Karmali, Avik Pal, and Viral Shah. Fashionable modelling with Flux. CoRR , abs/1811.01457, 2018.

Jacob Kelly, Jesse Bettencourt, Matthew J. Johnson, and David K. Duvenaud. Learning differential equations that are easy to solve. In Advances in Neural Information Processing Systems , pages 4370-4380, 2020.

Bobak Kiani, Randall Balestriero, Yann LeCun, and Seth Lloyd. projUNN: efficient method for training deep networks with unitary matrices. In Advances in Neural Information Processing Systems , pages 14448-14463, 2022.

Patrick Kidger. On Neural Differential Equations . PhD thesis, University of Oxford, 2021.

Patrick Kidger, Ricky T.Q. Chen, and Terry J Lyons. "Hey, that's not an ODE": Faster ODE adjoints via seminorms. In Proceedings of the 38th International Conference on Machine Learning , pages 5443-5452. PMLR, 2021.

Suyong Kim, Weiqi Ji, Sili Deng, Yingbo Ma, and Christopher Rackauckas. Stiff neural ordinary differential equations. Chaos: An Interdisciplinary Journal of Nonlinear Science , 31(9), 2021.

Naomi E. Leonard and P.S. Krishnaprasad. Control of switched electrical networks using averaging on Lie groups. In Proceedings of 1994 33rd IEEE Conference on Decision and Control , pages 1919-1924, 1994.

Mario Lezcano Casado. Trivializations for gradient-based optimization on manifolds. In Advances in Neural Information Processing Systems , pages 9157-9168, 2019.

Mario Lezcano-Casado and David Martínez-Rubio. Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In Proceedings of the 36th International Conference on Machine Learning , pages 3794-3803. PMLR, 2019.

Yi Heng Lim and Muhammad Firmansyah Kasim. Unifying physical systems' inductive biases in neural ODE using dynamics constraints. arXiv preprint arXiv:2208.02632 , 2022.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations , 2019.

Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser Nam Lim, and Christopher M. De Sa. Neural manifold ordinary differential equations. In Advances in Neural Information Processing Systems , pages 17548-17558, 2020.

Michael Lutter, Christian Ritter, and Jan Peters. Deep Lagrangian networks: Using physics as model prior for deep learning. In 7th International Conference on Learning Representations , 2019.

Yingbo Ma, Vaibhav Dixit, Michael J Innes, Xingjian Guo, and Christopher Rackauckas. A comparison of automatic differentiation and continuous sensitivity analysis for derivatives of differential equation solutions. In 2021 IEEE High Performance Extreme Computing Conference (HPEC) , pages 1-9. IEEE, 2021.

Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. In Advances in Neural Information Processing Systems , pages 2503-2515, 2020.

Y. Nishimori. Learning algorithm for independent component analysis by geodesic flows on orthogonal group. In IJCNN'99. International Joint Conference on Neural Networks. Proceedings (Cat. No.99CH36339) , volume 2, pages 933-938 vol.2, 1999. doi: 10.1109/IJCNN. 1999.831078.

Alexander Norcliffe, Cristian Bodnar, Ben Day, Nikola Simidjievski, and Pietro Liò. On second order behaviour in augmented neural ODEs. In Advances in Neural Information Processing Systems , pages 5911-5921, 2020.

Avik Pal, Yingbo Ma, Viral Shah, and Christopher Rackauckas. Opening the blackbox: Accelerating neural differential equations by regularizing internal solver heuristics. In Proceedings of the 38th International Conference on Machine Learning , pages 8325-8335. PMLR, 2021.

Suhan Park, Mathew Schwartz, and Jaeheung Park. NODE IK: Solving inverse kinematics with neural ordinary differential equations for path planning. arXiv preprint arXiv:2209.00498 , 2022.

Linda Petzold. Differential/algebraic equations are not ODE's. SIAM Journal on Scientific and Statistical Computing , 3(3):367-384, 1982.

Christopher Rackauckas and Qing Nie. DifferentialEquations.jl-a performant and feature-rich ecosystem for solving differential equations in Julia. Journal of Open Research Software , 5(1), 2017.

Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Dominic Skinner, Ali Ramadhan, and Alan Edelman. Universal differential equations for scientific machine learning. arXiv preprint arXiv:2001.04385 , 2020.

Danilo Jimenez Rezende, George Papamakarios, Sebastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. In Proceedings of the 37th International Conference on Machine Learning , pages 8083-8092. PMLR, 2020.

Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergencebased generative modeling on manifolds. In Advances in Neural Information Processing Systems , pages 17669-17680, 2021.

Yulia Rubanova, Ricky T.Q. Chen, and David K. Duvenaud. Latent ordinary differential equations for irregularly-sampled time series. In Advances in Neural Information Processing Systems , pages 5320-5330, 2019.

Chaopeng Shen, Alison P. Appling, Pierre Gentine, Toshiyuki Bandai, Hoshin Gupta, Alexandre Tartakovsky, Marco Baity-Jesi, Fabrizio Fenicia, Daniel Kifer, Li Li, Xiaofeng Liu, Wei Ren, Yi Zheng, Ciaran J. Harman, Martyn Clark, Matthew Farthing, Dapeng Feng, Praveen Kumar, Doaa Aboelyazeed, Farshid Rahmani, Yalan Song, Hylke E. Beck, Tadd Bindas, Dipankar Dwivedi, Kuai Fang, Marvin Höge, Chris Rackauckas, Binayak Mohanty, Tirthankar Roy, Chonggang Xu, and Kathryn Lawson. Differentiable modelling to unify machine learning and physical models for geosciences. Nature Reviews Earth & Environment , 4(8):552-567, 2023.

Takeshi Teshima, Koichi Tojo, Masahiro Ikeda, Isao Ishikawa, and Kenta Oono. Universal approximation property of neural ordinary differential equations. arXiv preprint arXiv:2012.02414 , 2020.

Ch. Tsitouras. Runge-Kutta pairs of order 5 (4) satisfying only the first column simplifying assumption. Computers & Mathematics with Applications , 62(2):770-775, 2011.

James H. Verner. Numerically optimal Runge-Kutta pairs with interpolants. Numerical Algorithms , 53(2-3):383-396, 2010.

Han Zhang, Xi Gao, Jacob Unterman, and Tom Arodz. Approximation capabilities of neural ODEs and invertible residual networks. In Proceedings of the 37th International Conference on Machine Learning , pages 11086-11095. PMLR, 2020.

Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ODE-Net: Learning Hamiltonian dynamics with control. In 8th International Conference on Learning Representations , 2020.

Aiqing Zhu, Pengzhan Jin, and Yifa Tang. Deep Hamiltonian networks based on symplectic integrators. arXiv preprint arXiv:2004.13830 , 2020.

A Differential Algebraic Equations

A differential algebraic equation (DAE) in its most general, implicit form is

where x ∈ R n and F : R × R n × R n → R n . When ∂F/∂ ˙ x is nonsingular, Equation (24) is an implicit ODE and by the implicit function theorem may be written as an explicit ODE in the form ˙ x = f ( x, t ) [34]. In the more interesting case of singular ∂F/∂ ˙ x , an important special case of Equation (24) is given by semi-explicit DAEs in Hessenberg form, for example,

where x = ( y, z ) . We call y the differential variables , since their derivatives appear in the equations, and z the algebraic variables , since their derivatives do not. The semi-explicit form of Equation (25) highlights the connection between certain classes of DAEs and ODEs subject to constraints.

It is generally possible to differentiate the constraints in Equation (25b) a number of times and substitute the result into Equation (25a) to obtain a mathematically equivalent ODE. The number of differentiations required to do so is the differential index of the DAE, and corresponds loosely to the 'distance' of the DAE from an equivalent ODE. The differential index - and related measures of index not directly based on differentiation - is used extensively to classify DAEs, especially in the context of numerical methods for their solution [15]. Each differentiation of the constraints reduces the index of the system by one (ODEs have index 0).

Of particular interest in the context of this paper are constrained ODEs of the form

where u ∈ R n , f : R × R n → R n , and g : R × R n → R m . Equation (26) can be written as a semiexplicit Hessenberg index-2 DAE, with u as the differential variables and m Lagrange multipliers as the algebraic variables, for example,

where λ ∈ R m and D ( u ) is any bounded matrix function such that GD , where G = g u is the Jacobian of g , is boundedly invertible for all t [8]. We can therefore approach the task of solving the constrained ODE Equation (26) from the perspective of solving the Hessenberg index-2 DAE Equation (27).

DAEs are not ODEs, however, and a number of additional complications arise when we seek numerical solutions [65]. Generally, the higher the index, the harder it is to solve a given DAE. For this reason, it is common to first perform an index reduction (i.e. differentiate the constraints) before applying numerical methods. However, the numerical solution of the resulting index-reduced system may exhibit drift off from the invariant manifold defined by the original constraints. For this reason, Baumgarte [10] proposed a stabilization procedure for index-reduced DAEs that renders the invariant manifold asymptotically stable. Baumgarte's stabilization is, in turn, a special case of the stabilization procedure later proposed by Chin [21] and adopted by us in this paper. We emphasize, however, that our stabilization procedure addresses a different application and problem than these related methods; while Baumgarte and Chin sought to stabilize drift off from the invariant manifold due to discretization error in the numerical solution of an index-reduced DAE, we seek to constrain some learned dynamics imperfectly approximated by a neural network.

Finally, one may ask why a neural network could not be incorporated directly into Equation (27) and the resulting index-2 DAE solved directly. While possible in principle, DAEs require implicit numerical methods, with the result that the computational complexity of computing gradients of solutions - whether via automatic differentiation or adjoint sensitivity analysis - scales with the cube of the system size [51]. In contrast, backpropagating through explicit solvers has linear complexity.

Time (seconds)

100

102

104

e

(

𝜃

)

0.70

0.75

0.80

Time (seconds)

0

250

500

Relative Error (Endpoint)

10 -5

10 -4

10 -3

10 -2

(a)

(b)

Ground Truth

NODE

SNODE

SNODE (fast)

B Faster Alternatives to the Pseudoinverse

The computational cost of the pseudoinverse is O ( m 2 n ) , where m is the number of constraints and n is the dimension of the problem. While cheap to compute for the systems studied in this paper, the pseudoinverse may become prohibitively expensive for applications with a large number of constraints. As an alternative choice for the stabilization matrix F ( u ) in Equation (5), we recommend using the transpose of the Jacobian G ( u ) = g u , i.e.,

which is instead O (1) . Since G T ( u ) G ( u ) is symmetric positive definite, the theoretical guarantees of Theorem 1 and Theorem 2 still apply, such that

is a valid SNDE.

We apply an SNDE in the form of Equation (29) to the controlled robot arm experiment, in which the Jacobian has dimensions 2 × 3 . In this instance, a single evaluation of G T ( u ) is approximately 10 times faster than the pseudoinverse G + ( u ) , which translates to a 10% speedup when evaluating the entire right-hand side of the SNDE (the cost of evaluating of the neural network still dominates in this setting).

Figure 6 shows that the alternative SNDE formulation implements the control effectively and illustrates the tradeoff when choosing the stabilization matrix F ( u ) . The pseudoinverse G + can be expensive to compute but, as an orthogonal projection, yields the most 'direct' stabilization to the constraint manifold and therefore the most accurate implementation of the constraint. The transpose G T , on the other hand, is always cheap to compute and yields a valid stabilization, at the cost of only slightly larger errors in the constraint.

C The Stabilization Parameter γ

C.1 Runtime Implications

We assess the computational cost of SNODEs compared to vanilla NODEs. SNODEs require the computation of (and backpropagation through) the pseudoinverse of the Jacobian of the constraint

function g . Additionally, as γ is increased, SNODEs may require more solver steps at a given error tolerance, with the SNODE eventually becoming stiff for sufficiently large γ . Naively, one may thus expect a noticeable increase in runtime. However, as described in Section 2, the computational cost of training NODEs also depends on the 'complexity' of the learned dynamics, which in turn determines how many function evaluations are required by the solver. This leads to nontrivial interactions between the added computation of enforcing constraints and the thereby potentially regularized 'simpler' dynamics, which may require fewer function evaluations by the solver.

In Table 1, we report comparisons of training times between SNODEs and NODEs for different values of γ for three settings. SNODEs take roughly 1.2 to 1.8 times longer to train, with smaller values of γ incurring less overhead. Overall, this is a manageable increase for most relevant scenarios.

We also show inference times in Table 2. Here, the trend reverses and larger values of γ lead to lower inference times. This is because the solver requires fewer steps (has higher acceptance rates of proposed step sizes) for stronger stabilization. Hence, while predictive performance is largely unaffected, one can use the specific choice of γ as a tuning knob that trades off training time versus inference time.

C.2 Constraint Implications

To complement these results, Figure 7 shows that the relative error remains almost unchanged for a large range of γ values, that is, beyond a certain minimum value, SNODEs are not sensitive to the specific choice of γ . Even a value of γ = 1 works well in the settings we have considered, indicating that we can typically get away with runtime increases of a factor of 1.2. However, larger values of γ only lead to slightly increased training times, while generally enforcing the constraints to a higher degree of accuracy.

C.3 A Practical Guide to Choosing γ

As described in the previous sections, choosing the optimal value of γ requires a tradeoff between accuracy and training time; larger values of γ will enforce the constraints more accurately, at the cost of (mildly) additional training time. Our experience with the systems in this paper is that stabilization begins to work around γ ∼ 1 and remains effective as γ is increased, until the system becomes stiff around γ ∼ 100 . Within those limits, we suggest experimenting with a range of values, for example, powers of two.

Finally, we note that the stabilization term can be switched 'on and off' at any time during training. For example, if computing the pseudoinverse is expensive, it may be beneficial to start training without stabilization, before switching the stabilization on once f θ is close to f .

0

50

100

Relative Error (State)

10 -3

10 -2

10 -1

10 0

0

50

100

Relative Error (Ang. Mom.)

10 -4

10 -3

10 -2

0

400

800

Relative Error (State)

10 -4

10 -2

10 0

0

400

800

Relative Error (Casimir)

10 -5

10 -4

10 -3

10 -2

Time (seconds)

0

40

80

Relative Error (State)

10 -3

10 -2

10 -1

10 0

Time (seconds)

0

40

80

Relative Error (Energy)

10 -4

10 -3

10 -2

10 -1

10 0

10 1

10 2

(a)

(b)

(c)

(d)

(e)

(f)

𝛾

0.1

1

2

4

8

16

32

D Additional Experiments

D.1 Hamiltonian Neural Networks

We train a Hamiltonian neural network (HNN) on the two-body problem, which is a Hamiltonian system (Figure 8). The HNN initially conserves angular momentum, suggesting that it has learned a reasonable approximation of the true Hamiltonina. However, it becomes unstable after approximately 50 seconds, highlighting the benefits of stabilization when data is limited and conservation laws are known.

D.2 Tangent Projection Operator

The invariant manifold of the rigid body experiment is a sphere. To constrain trajectories to the sphere, we can follow an approach similar to Rozen et al. [69] and Ben-Hamu et al. [12] and define a tangent projection operator (TPO),

where I is the identity matrix. P ( u ) projects velocities at u to the tangent space of the sphere T u S 2 . Applying the tangent projection operator of Equation (30) to the standard NDE model f θ for rigid body dynamics, we obtain

Equation (31) is a manifold ODE. In general, specialized numerical methods must be used when integrating manifold ODEs to guarantee the numerical solution remains on the manifold [43]. However, for the sake of this comparison, we find regular ODE solvers to be sufficient in practice.

Figure 8(d) shows that the TPO satisfies the constraint exactly, with relative errors of ∼ 10 -10 , compared to ∼ 10 -4 for the SNODE and SANODE. However, this does not translate to improved prediction of the system state (Figure 8(c)), suggesting that, in this setting, there is little practical benefit to the additional accuracy in the constraint.

D.3 Stable Time

The unstabilized, vanilla NODEs of Figure 2 are characterized by an initial drift from the invariant manifold that gives way to a subsequent rapid divergence. In practice, however, certain test trials may remain stable for significantly longer than others. In this section, we characterize the stable time T stab of an individual test trial as the time elapsed until the relative error E ( t ) in the predicted system state ˆ u ( t ) exceeds a given threshold value E stab , i.e.

Taking E stab = 10 3 , Table 3 shows T stab for the the two-body problem, rigid body, and DC-to-DC converter experiments. Across several thousand test trials in this paper, not one stabilized NDE model diverged.

E Invariant Measure

Given that the double pendulum is a chaotic system, predictions of individual trajectories will break down after short times. We therefore also quantify the performance of NODEs and SNDEs in terms

Time (seconds)

0

50

100

Relative Error (State)

10 -5

10 0

10 5

Time (seconds)

0

50

100

Relative Error (Ang. Mom.)

10 -5

10 0

10 5

Time (seconds)

0

400

800

Relative Error (State)

10 -5

10 0

10 5

Time (seconds)

0

400

800

Relative Error (Casimir)

10 -15

10 -10

10 -5

10 0

10 5

(a)

(b)

(c)

(d)

NODE

ANODE

SNODE

SANODE

HNN

TPO

of their ability to capture the double pendulum's invariant measure. We refer to Arnold et al. [6] and Chekroun et al. [17] for detailed definitions of invariant measures; in short, a measure µ is said to be invariant under some flow Φ if µ (Φ -1 ( t )( A )) = µ ( A ) for all measurable sets A . Invariant measures are commonly used to characterize the long-term dynamical characteristics of chaotic dynamical systems (see, for example, Arnold et al. [6] or Chekroun et al. [17] for a discussion of the invariant measure of the paradigmatic Lorenz-63 system). Since the double pendulum is an ergodic system, averages over long times approximate ensemble averages. We can therefore obtain a sample of the invariant measure numerically by integrating the system for a very long time. Concretely, we estimate the invariant measure from a single long trajectory using an algorithm due to Diego et al. [28], implemented in ComplexityMeasures.jl [41], based on a numerical estimate of the transfer operator. We then use the Hellinger distance [22] to compare the resulting probability distribution with the ground truth value for the double pendulum.

F Architecture and Training

Trajectories are generated using the 9(8) explicit Runge-Kutta algorithm due to Verner [74], implemented in DifferentialEquations.jl [66] as Vern9 . To ensure that invariants are satisfied exactly in the training data, we use absolute and relative solver tolerances of 10 -24 in conjunction with Julia's BigFloat number type. Trajectories for training and validation are independent, that is, a given trajectory is used exclusively either for training or validation. Each trajectory is split using a multiple-shooting approach into non-overlapping chunks of 3 timesteps each.

Networks are implemented using Flux.jl [46] and consist of fully-connected dense layers with ReLU activation functions. All experiments are trained for 1,000 epochs using the AdamW optimizer [56] with weight decay of 10 -6 and an exponentially decaying learning rate schedule. During training, trajectories are integrated using the 5(4) explicit Runge-Kutta algorithm due to Tsitouras [73], implemented in DifferentialEquations.jl [66] as Tsit5 , with absolute and relative tolerances of 10 -6

The stabilization hyperparameter γ as well as network sizes, learning rates, and the number of additional dimensions for ANODEs are optimized for each experiment and are summarized in Table 4.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nNeural ODEs which learn a particularly dynamics can be unstable during evolution of the learned dynamics. This papers studies instabilities which take solutions outside of the constraint space of a given evolution. The authors show that a simple added control term to the learned evolution can improve stability noticeably. This added term, intuitively, projects part of the flow onto the tangent space of the constraint manifold. This stability term has the benefit that it is not complicated to implement, has some theoretical guarantees of stability, and works well in the models they consider.\n\nSTRENGTHS:\nAs mentioned in the summary, I see a few benefits to the authors' proposed techniques:\n- The method is rather simple and easy to follow. The reason why it enhances stability is readily apparent and experiments back this up.\n- The authors raise an important point about neural ODEs in the context of constraints. As someone not completely familiar with the field, I appreciated how they setup the problem and discussed it.\n- There are some simple theoretical guarantees of the stability of the approach. The experiments back up these theoretical guarantees, though the experiments appear to be focused on mostly toy models.\n\nWEAKNESSES:\nMain comments:\n- There are a host of tools in the Riemannian optimization literature that are relevant here and not discussed. For example, there exists many algorithms for learning flows or vector fields on manifolds (e.g. [1], [2]) and performing transport on those manifolds. There are also simpler methods such as projected gradient descent or riemannian descent via retractions of the exponential map which can easily stay on a manifold and could be adapted to this setting. I suppose the authors would argue that their idea is simpler, but I would not be convinced unless these are compared to at least empirically or theoretically.\n- The authors' main contribution is addressing a control problem and not really a learning problem. The paper includes an additional term to the ODE evolution to better control the trajectory, but this parameter is not learned in any way. Thus, the problem they address is really one of control and not learning. This problem of stable control is well studied in control theory as far as I understand. I raise this point for two reasons. First, I worry the problem the paper addresses is a bit outside of the scope of this conference. Second, as I am not an expert here, it is not clear to me how effective their particular controller is. Perhaps the authors can expand on this point and address comparisons to methods in the control literature. \n- The theoretical guarantees and experiments, I think, could be expanded. Related to my previous comments, the theorems are two straightforward statements about the added stability of the control term. There are no theorems relating to how much better this is for specifically learning tasks (e.g., some proofs about generalization error). On the experiments side, all experiments seem to be on toy models of small dimensional systems. It is difficult to see how much better this will work in higher dimensional or more complex systems. \n\n\n\nSmaller comments:\n- Eq. 3 missing $u \\in \\mathbb{R}^n$\n- A number of references have no journal or indication of where it was published (even an arXiv indicator)\n- The pseudo-inverse can be expensive to calculate in high dimensions with large numbers of constraints. Is this an issue in practice? Have you considered faster approximations to this?\n- Some portions of the text use $f_\\theta(u)$ and others $f_\\theta(u,t)$. I would stick with one notation and express whether or not the setting the authors study depends on $t$ explicitly throughout the work.\n- All of the experiments seem to be on low dimensional systems. Is there a reason why the authors did not go to higher dimensional examples?\n- Many of the problems people care about are on continuous spaces (e.g., fluid flow in Navier stokes). Does the method here extend to that setting in an obvious way?\n\n\nReferences:\n\n[1] Lou, Aaron, et al. \"Neural manifold ordinary differential equations.\" Advances in Neural Information Processing Systems33 (2020): 17548-17558.\n\n[2] Ben-Hamu, Heli, et al. \"Matching normalizing flows and probability paths on manifolds.\" arXiv preprint arXiv:2207.04711 (2022).\n\nQUESTIONS:\nPlease see weaknesses section for questions.\n\nLIMITATIONS:\nThere is a brief discussion of limitations of the work in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-12"}

comments: SUMMARY:
This paper proposes a general framework to stabilize NODE models. By introducing a new stabilization term, it can learn vector fields that satisfy various conditions. The authors evaluate the effectiveness of the proposed method on several physical systems represented by ODEs.

STRENGTHS:
- This paper proposes a new stabilization term that makes it possible to learn vector fields that satisfy various types of conditions.
- The proposed method has been theoretically and experimentally evaluated, and the results are valid.
- The paper is well written.

WEAKNESSES:
- Lack of clarity on how the problem setup differs from HNN and LNN.
- Insufficient discussion of applicability to PDEs.

QUESTIONS:
- When considering conservation laws as constraints, does the proposed method require the energy function to be given explicitly, and does it also require the total energy to be known? HNN can learn a conserved system even when the energy function is unknown by incorporating a symplectic structure into the NN. Can you clarify the difference in problem setup (especially in input data) between HNN or LNN and the proposed method?

- In section 5, application to PDEs is mentioned. Are there any difficulties other than the issue of computational complexity? Also, is it possible to consider boundary conditions in PDEs as stabilization terms?

LIMITATIONS:
As written in Questions, it is better to be clear about the problem setting and add a discussion about the scope of application of the proposed method.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a new method called stabilized neural differential equations (SNDEs) that enforce arbitrary manifold constraints for neural differential equations. The proposed method is applicable to any type of ODE, hybrid models (UDEs), and can incorporate any type of manifold constraints. The authors' showcase their work on first- and second-order systems including non-chaotic and chaotic, non-autonomous and autonomous examples.\n\nSTRENGTHS:\nOriginality: \nThe related work is nicely covered with a specific focus on Hamiltonian/Lagrangian and other constrained neural differential equation methods. It has been made clear out of the related work, which previous approach is most similar to the work presented and how the current work differs/builds upon it. \nThe proposed work is a novel perspective/approach on the existing work on constrained neural DE methods. It is largely based on earlier PhD thesis by Chin (1995) on stabilization methods for simulations of constrained multi-body dynamics. \n\nQuality: \nThe submission is technically sound. I have not analyzed Proof 2 in depth, but based on the information that the authors' have provided, it seems supported. \n\nClarity: \nThe submission is very clearly written making it enjoyable to read as it is clear and understandable for the reader. The submission is well organized, the notation is clear, the function input/output dimensions are clarified, as well as matrices are defined.\n\nWEAKNESSES:\nQuality:\nDespite providing a thorough overview of related work that also incorporate constraints, experimentally the authors only compare to NODE and SONODE which are known not to extrapolate well beyond the training regime due to the limited constrains/bias for the differential function/vector field. If the authors would have included comparisons to some of the other works mentioned in the related works section, such as Finzi et al., or other method that include a Hamiltonian constraint (SymODEN etc), it would put the present work in a better comparison to the related work in the field, as it provides an alternative to it. \nAs such an explicit comparison has not been made, it is hard to evaluate the present method in context of other, constraint based Neural ODE type, methods. \n\nClarity:\nEq (3). I assume there is a missing $\\in \\mathbb{R}^n$\n\nSignificance: \nThe proposed method is a relevant contribution to the field, however, I miss a better empirical comparison to other recent constrained neural ODE methods, this would make the submission a lot stronger and a better contribution to the existing work in the field.\n\nQUESTIONS:\nThe proof of Theorem 2, is this the author contribution or as the Theorem 2 comes from the work of Chin, 1995, was the proof already established then? \n\nFor Eq (9), this states that the derivative is either  0 or negative, meaning that the vector field is decreasing or stable (0), how does this then relate to the example given in Fig. 1 where at around coordinates (0,0) the vector field is 'expanding'? \n\n4.1., 4.2., 4.3., 4.4., 4.5 what are the test settings? what is the $\\gamma$ value used in these experiments? \n\nIn section 3 you mention that it might be beneficial to start training without stabilization, do you test this in practice?\n\nIn general, you mention that 'are easy to use in practice across a wide variety of settings with minimal to no tuning', however based on results on double pendulum it could be suggested that as any hyperparameter also $\\gamma$ requires tuning, meaning that the performance is sensitive to it. Have you done additional experiments to comment on whether there are certain experimental set-ups that require a certain $\\gamma$ value, or a general advise what is a good starting value etc, which would be informative to the reader as a starting point?\n\nAre there additional experiments available that also compare the present method to existing constrained neural ODE models? \nAt the moment, this is a major limitation of the presented work, which negatively affects the quality of the paper, even if the proposed method is sound.\n\nLIMITATIONS:\nThe authors have addressed limitations, in particular, as a key limitation it is stated that at the current implementation the method is limited to ODEs.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

comments: SUMMARY:
This paper aims to learn ODE dynamics where the dynamics has some extra equation constraint. For example, a double pendulum with energy conservation. In the problem setting, a neural network is used to learn the time derivative of the system, the dynamics is then rolled out with a classical solver. To ensure the rollout respects the conservation laws, the authors propose to add a stabilizing term to the time derivative, which can be proven to make the rollout asymptotically converge to the manifold defined by the conservation law. The stabilizing term is implemented with peudo-inverse of the Jacobin of the conservation function. Experiments show that the proposed method can stabilize the rollout w.r.t. the conservation law, compared to the method without stabilization.

STRENGTHS:
- The paper has good theoretical analysis and the practical implementation is based on the theories.
- The experimental results are improved.

WEAKNESSES:
- The baseline method involves no stabilization. Is there other simple method that can ensure the stabilization? For example, since the conservation equation is known explicitly, can one use one degree of freedom only to satisfy the constraints? I.e., if the output is $m$-dimensional, the network can only predict $m-1$ dimensions and use the conservation equation to compute the last dimension directly.

QUESTIONS:
- The method involve computing the pseudo-inverse of the Jacobin matrix of $g$. If $g$ itself is complicated, e.g., contains higher order derivatives, would the computational cost become very significant?

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-19

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors of this paper proposed to augment neural ODEs with an additional stabilization term that is derived from the pseudo-inverse of a Jacobian matrix to learn dynamics of a system from data. The authors provided theoretical grantees to show that their augmented version is still capable of learning the original dynamics despite having an additional term and that the additional term actually stabilizes the learned dynamics. Experimental results were shown and compared with a vanilla NODE across a variety of dynamical systems with different characteristics.\n\nSTRENGTHS:\nThe overall paper, description and explanations of the work is very well written and relatively accessible even to a reader who is not very familiar to the topic. I also appreciate the simplicity of the idea, in that the proposed stabilization term could potentially be easily combined with other variations of neural ODE methods to achieve potentially more accurate predictions. The experimental results were presented in a wide variety of scenarios and the benefits of the proposed method does seem significant when compared to Vanilla ODEs.\n\nWEAKNESSES:\nThe following are the two improvements that I believe could improve the paper.\n\n1. The idea behind this paper seems simple and the results seems promising, which given it's simplicity, could really be an impactful contribution. Hence, if the authors could also demonstrate the gains of this method on another variation of vanilla NODE and also perhaps an actual application, I believe this paper will be significantly stronger. \n\n2. The overall paper is relatively accessible to readers not familiar with the field from a conceptual point of view, however, it'll be more informative if the authors could more candidly explain the connection between certain parts of the paper. For example, it is not very clear from an initial pass how does the manifold relates to the loss function and training of the NDE and the additional stabilization term.\n\nQUESTIONS:\n1. Could the authors discuss further if there are potentially other candidates for the matrix F besides the ones that the author's used?\n\n2. From reading the main paper alone, it wasn't very clear to me if the results shown were from the training trajectories or a set of testing trajectories with unseen initial conditions. I later realized that the information was provided in the supplementary material but I believe it at least be briefly stated in the main paper.\n\n3. In Figure 2, what does (SO) NODE represent?\n\n4. I noticed that the length of the trajectories and the resolution of each time step is sometimes different across the different experiments. Is there specific reason for this?\n\n5. From what I understand, each trajectory is split into 3 time steps of non-overlapping chunks and then divided into training and validation. I am curious how does this method perform if the training and validation is not randomized, but is only trained on the initial segment of the trajectory and tested on the later segment of a trajectory?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-20"}