Stabilized Neural Differential Equations for Learning Dynamics with Explicit Constraints
Authors: Alistair White
Niki Kilbertus
Maximilian Gelbrecht
Niklas Boers
Venue: NeurIPS 2023 poster (2023)
Keywords: neural differential equations
neural ordinary differential equations
constraints
conservation laws
stabilization
dynamical systems
dynamics
scientific machine learning
physics-informed machine learning
TL;DR: We introduce a simple yet powerful stabilization method that allows us to learn dynamics with arbitrary explicit constraints from observed solution trajectories.

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
Alistair White
Niki Kilbertus
Technical University of Munich Potsdam Institute for Climate Impact Research
Technical University of Munich
Hemholtz AI, Munich
Maximilian Gelbrecht
Technical University of Munich Potsdam Institute for Climate Impact Research
Niklas Boers
Technical University of Munich
Potsdam Institute for Climate Impact Research University of Exeter
(alistair. white, niki kilbertus, maximilian. gelbrecht, aermmmllinngglbrecht,
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 stabilization term that, when added to the original dynamics, renders the constraint manifold provably asymptotically stable. Due its simplicity, method is compatible with all common neural differential equation (NDE) models and broadly applicable. In extensive empirical evaluations, demonstrate that SNDEs outperform existing methods while broadening the types of constraints that can be incorporated into NDE training.
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 surge of interest data-driven methods for learning dynamical laws the form of differential equations directly from data [18 67 [6] 23 3 [1. Assuming there is ground truth system with dynamics governed by an ordinary differential equation
du(t) dt
= t)(
(with initial condition u(0) = 24)
(1)
with u : IR + R" and f R" x R + R", the question is whether we can learn f from (potentially noisy and irregularly sampled) observations (t,, u u(t:));
Neural ordinary differential equations (NODEs) provide a prominent and successful method for this task, which leverages machine learning by directly parameterizing the vector field of the ODE as
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(a)
(b)
(c)
NDE
Constraint Manifold
Stabilization
SNDE
SNDE Trajectory
Figure 1: Sketch of the basic idea behind stabilized neural differential equations (SNDEs). (a) An idealized, unstabilized NDE vector field (blue arrows). (b) A constraint manifold (black circle) and the corresponding stabilization of the vector field (pink arrows), (c) The overall, stabilized vector field of the SNDE (orange arrows) and stabilized trajectory (green). The stabilization pushes any trajectory starting away from (but near) the manifold to converge to it at rate (see Section
neural network [8] (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 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 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 substantially aid generalization, learning speed, and stability, especially the low data regime. Learning dynamics from data is exception [4. In scientific applications, physical priors are often not only 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] (2). 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 stabilization technique enforce arbitrary, even time-dependent manifold constraints for any class 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://giiee http:///ggihuuu...///oeooae com htt.:///ttttt.a.aaaaaaaaaaaallea Neural Diferraaaliffeeentiall -aaaalllluuiinnnn
2 Background
first order neural differential equation is typically given as
du(()
dt (t()))) u(()= uo:
(2)
where R RR and the vector field fq: R" R + R is at least partially parameterized by a neural network with parameters 0 E Rdd We restrict our attention to ground truth dynamics in Equation (T) that are continuous in / and Lipschitz continuous in u such that the existence of unique (local) solution to the initial value problem is guaranteed by the Picard-Lindeldo theorem. As universal function approximators [4 255 neural networks fo can in principle approximate any such if to arbitrary precision, that to say, the problem of learning f is realizable.
In practice, the parameters 0 are typically optimized via stochastic gradient descent by integrating a trajectory â(t) = ODESolve uo. fo. t) and taking gradients of the loss C(u, û) with respect Gradients of trajectories can be computed using adjoint sensitivity analysis (optimize-ther -iseeenreeeeeieee or automatic differentiation of the solver operations ddiscrelize-hhen-opiiizze [8 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 41/ 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-speccifid parameteriza- tion 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 pregularly-samplee
observations (70] and partial differential equations [35. Bayesian NODEs (26). universal differential equations 167.. and more-sce, 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 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 66 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 [7 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 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 parameterizz- ing 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.
39 assume second order Hamiltonian dynamics u(t), where u(t)= = (q(t). p(f)) 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(g p) = ( ( अq ). Symplectic
Hamiltonian Neural Networks (HNNs)
0H
ODE-Net (SymÖDEN) (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-fres 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 SB(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 recent overview of learning Hamiltonian dynamics using neural architectures.
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. 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 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.
3

---

Continuous Normalizing Flows on Riemannian Manifolds. Another large body of related work aims learn continuous normalizing flows (CNFs) on Riemannian manifolds. Since CNFs transform probability densities via 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 manifold. The authors limit their attention to the hyperbolic space IHI and the sphere S both of which are prototypical Riemannian manifolds with easily derived closed-form expressions for the chart. Many similar works have been proposed [7 14 68 60), but all are typically limited in scope to model Riemannian manifolds such 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 interesting opportunity for further work.
explicit chart representation of the
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] 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 cer- tain 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 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 non- chaotic 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.
Stabilized Neural Differential Equations
General approach. Given < n explicit constraints, we require that solution trajectories of the are confined to an (n m) dimensiion submanifold of R. defined by
NDE in Equation 2
M [(u t) E R x R; g(u,t) = 0}.
(3)
where g: R x R + R is a smooth function with O E RT being a regular value of g In other words, we have NDE on manifold
i fo( with g(u,t) 0.
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 to O. 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 that solves a given initial value problem for the ODE defined by f. Inspired by Chin (1), we propose the following stabilization of the vector field in Equation
- ((())))F(w)g(u)))
[general SNDE]
(5)
where y 2 0 is a scalar parameter of our method and F R² + Raxm is a so-called stabilization matrix. We call Equation ( stabilized neural differential equation (SNDE) and say that it is
'The preimage theorem ensures that M is indeed an (n m))dimennionall submanifold of R

---

stabilized with respect to the invariant manifold M.. We illustrate the main idea in Figure I while even small deviations in the unstabilized vector field (left) can lead (potentially accumulating) constraint violations, our stabilization (right) adjusts the vector field near the invariant manifold to render 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 Appendix A for brief overview of these connections.
Theoretical guarantees. First, note that ultimately still want fo 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 fo can still learn the correct dynamics when using a different right-hand side.
Theorem (adapted from Chin (2]). Consider an NDE
ù ύ= [e
(6)
on an invariant manifold M = { ER"; g(u) 0}. A vector field ú = ho(u) admits all solutions of Equation (6) on M if and only if ho folM-
Since (()) = 0 on M, the second term on the right-hand side of Equation 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 the constraint manifold such that M is asymptotically stable.
Theorem 2 (adapted from Chin 121 Suppose the stabilization matrix F()) is chosen such that the matrix F(u), where G(u) Iu the Jacobian of g at u, is symmetric positive definite with the smallest eigenvalue X()) satisfying A(u) > do > O forrr all u. Assume further that there is a positive
number "No such that
(G((u)fo (u)il S 70 (ool)(4)]]
for all u near M. Then the invariant manifold M is asymptotically stable in the SNDE Equation if 2 7o/>o
Proof. Consider the Lyapunov function V(u)= - (u)g(u). Then (omitting arguments)
2 di Ilo(u(t))) -- -d
1 d
dg dg di g" d g G(fo - G('g)
v(1
(8)
where we substitute Equation (5) for û. With Equation 7 we have GG S "70g" and since the eigenvalues of GF are assumed be least Ao 0, we have g GFg 2 Aug g. Hence
d
V (10 >
di
so the manifold M is asymptotically stable whenever To 7Ao S 0.
When so(u) and g(u) are given, M is asymptotically stable in the SNDE Equation as long as
[G(u)
(10)
Au)l
To summarize, the general form of the SNDE in Equation 5 has the following important properties:
1. The SNDE admits all solutions of the constrained NDE Equation 4 on M.
2. .M is asymptotically stable in the SNDE for sufficiently large values of 7.
The stabilization parameter 7, 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, 7 is neither Lagrangian parameter (corresponding to a constraint 0), nor a regularization parameter (to overcome an ill- ill-posedness by regularization). Therefore, there is no "correct" value for y. particular, Theorem holds for all 7, while Theorem Sineem@oon only requires ? to be "sufficiently large".
In the limit ? + co, the SNDE in Equation 5 is equivalent to Hessenberg index-2 DAE (see Appendix A for more details).
We note that there is minor error in the proof of Theorem in Chin 21. which correct here.
5

---

Practical implementation. This leaves us to find 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. our experiments, we use the Moore-Penrose pseudoinverse of the Jacobian of 9 at u as the stabilization matrix,
F(u) = (T) = (u) (G)(G( (w) E RR
(11)
Let us analyze the properties of this choice. Regarding the requirements (b) and (c), the pseudoinverse can be computed efficiently via singular value decomposition with highly optimized implementations in all common numerical linear algebra libraries (including deep learning frameworks) and does not interfere with gradient-basee
optimization. In particular, the computational cost for the pseudoinverse is (m,2)) i.c.. it scales well with the problem size. The quadratic scaling in the number of constraints 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+ (+) is orthogonal projection onto the tangent space T.M of the manifold at u. Hence, locally in a neighborhood of u E .M, we consider the projection back onto the invariant manifold M (see Figure In particular, G(u) has full rank for u E M and G+G = G (GG")) is symmetric and positive definite near M. From here on, we thus consider the following specific form for the SNDE in Equation (5).
stabilization matrix
û =f(u)) - - (w)g(uu))
[practical SNDE]
(12)
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, non-autonomous first order system with a conserved quantity, a non-autonomous controlled first order system with 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 i(() versus ground truth u()). we use the relative error [|x(()= a (()))a/lllu(()) averaged over many trajectories with independent initial conditions. As 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 y over 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, 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 space constraints.
find that
5.1 Two-Body Problem [second order, autonomous, non-chaotic, conservation law]
The motion of two bodies attracting each other with force inversely proportional to their squared distance (e.g., gravitational interaction in the non-relativistic limit) can be written as
y
(13)
2 33/2
323/22
where one body is fixed 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
M M= R2² zÿ + yē Lo 0}.
(14)

---

Ground Truth
NODE
ANODE
SNODE
SANODE
(a)
(b)
Relative Error (Ang. Mom.)
(c)
1.0
Relative Error (State)
10
10
0.5
0.0
10
10
-0.5
10.
10
-1.5
-1.0
-0.5
0.0
0.5
50
100
0
50
100
27
Time (seconds)
Time (seconds)
(e)
Relative Error (Casimir)
Relative Error (State)
10
10$
1.01
0.9.9
10
10
0.97
0.2
0.2
0.0
0.2
-0.
0.0
10
10
1//
0
400
800
0
400
800
Time (seconds)
Time (seconds)
(g)
(h)
Relative Error (Energy)
Relative Error (State)
10
10
2
1
0
10
10
2
10.5
10.5
5
10
15
40
80
0
40
80
Time (seconds)
Time (seconds)
Time (seconds)
Figure Top row: Results for two-body problem experiment, showing a single test trajectory in (a) and averages 100 test trajectories in (b-c). Middle row: Results for the rigid body rotation experiment, showing a single test trajectory in (d) and averages over 100 test trajectories in (c-f). Bottom row: Results for the DC-to-De DC converter experiment, showing the voltage t1 across the first capacitor during single test trajectory in (g), and averages over 100 test trajectories in (h-i). The vanilla NODE (blue) is unstable in all settings, quickly drifting from the constraint manifold and subsequently diverging exponentially, while the vanilla ANODE (green) is unstable for the rigid body and DC-to-DC converter experiments. In contrast, the SNODE (red) and SANODE (purple) constrained to the manifold with accurate predictions over long horizon in all settings. Confidence intervals not shown they diverge along with the unstabilized trajectories.
where Lo is the initial value of L. We train 40 trajectories with initial conditions (z. y, ÿ) (1 e, 0, 0, v1-s/1+e), where the eccentricity is sampled uniformly via U(0.5, 0.7). Each trajectory consists of a single period of the orbit sampled with a timestep of At 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] of rigid body with arbitrary shape and mass
The angular momentum vector y = (y1, V2, ys) distribution satisfies Euler's equations of motion,
ú1 1/2
0
y2
V1 I
1/3
0
y1
12/L2 12
(15)
1/1
0
U3/ 1s
where the coordinate axes are the principal axes of the body. 11. 12. 13 are the principal moments of inertia, and the origin of the coordinate system fixed the body's 's centre of mass 43). The motion of y conserves the Casimir function Cly) = [8ï + v/ + vš), which is equivalent to conservation of angular momentum the orthogonal body frame and constitutes holonomic constraint on the
allowed states of the system. We therefore have the manifold
M {111, U2, Us) E RR² vी + v/ + v3 Co 03.
(16)
We train on 40 trajectories with initial conditions (y1 122 ys) = (cos(Ø), O, si)( where @ is drawn from uniform distribution timestep of A = 0.1 seconds.
(0.5, 1.5). Each trajectory consists of a 15 second sample with
The middle row of Figure 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-autonommous, non-chaotic, conservation law]
0
We now consider an idcalized DC-to-DC converter 52 illustrated in Figure 3 with dynamics
C1 (1 u)1s. C2bı udg,
Lyig = (1 u)ul uV2, (17)
Ci
where U1, U2 are the state voltages across capacitors C₁, C2. respectively, is is the state current across an inductor L3, and u E {0, 1 is control input (a switch) that can be used transfer energy between the two capacitors via the inductor. The total energy in the circuit, E
2000
L3
(Cuff
Figure 3:
Idealized
C2 + Tz3ī),, is conserved, yielding the manifold
schematic of a DCCto-
DC converter.
M [(vi, U2, ia) E R Cfuf + C2v2 + L3g Eo = 0}.
(18)
We train on 40 trajectories integrated over 10 seconds with timestep of At 0.1 seconds, where G = 0.1, C2 = 0.2, L3 = 0.5. and a switching period of 3 seconds, i.e., the switch is toggled every 1.5 seconds. The initial conditions for (v1, U2 is) are each drawn independently from a uniform distribution U(0,1 1).
The bottom row of Figure 2 shows the voltage across C1 over multiple switching events (g), with the NODE and ANODE quickly accumulating errors every time the switch 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 non-autonomous, non-chaotic, time-dependent control]
Next, apply SNDEs to solve 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 aa). 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 c(f) of the last segment is given by
[cos(() (00) + (0)) cos( cos(O2)) sin + sin(Oz) + sin(0z)
e(f)
(19)
where 0y is the angle the j-th segment with respect to the horizontal and (0,, 02, 0z). The problem consists of finding the motion of the three segments 0(t) such that the endpoint e(0) follows

---

(a)
(b)
(c)
Relative Error (Endpoint)
10"1
0.80
VVV
10
10
0.75
origin
10
0.70
10.5
102
104
500
Time (seconds)
Time (seconds)
Ground Truth
NODE
SNODE
Figure 4: Controlled robot arm. (a) Schematic of the robot arm. (b) Snapshot of single test trajectory. After 100 seconds the NODE (blue) has drifted significantly from the prescribed control while the SNODE (red) accurately captures the ground truth dynamics (black). (c) Relative error in the endpoint e(0) averaged over 100 test trajectories. The NODE (blue) accumulates errors and leaves the prescribed path, while the SNODE (red) remains accurate. Shadings in (c) are 95% confidence intervals.
prescribed path p(t) t) in the plane, i.e., e(0) = p(t). Minimizing [[o(t)]. it can be shown 42 that the optimal path satisfies
i e'(0) (())
(20)
where e' is the Jacobian of These will be our ground truth equations of motion.
We stabilize the SNODE with respect to the (time-dependent) manifold
M {(0, ESE R; e(0) p(t) 0).
(21)
In particular, we prescribe the path
p(t) = Eo (
(sin(20t1/22 0
(22)
where eo is the initial position of the endpoint, such that e(f) traces line back and forth the r-axiss
We train on 40 trajectories of duration seconds, with timestep At = 0.1 and initial conditions (0,, 02, 03) (0o, Oo On), where Op is drawn from a uniform distribution Og U(=/4, /4, 3r/8)
Additionally, we provide the network with ý. the time derivative of the prescribed control.
Figure shows that the unconstrained NODE drifts substantially from the prescribed path during long integration, while the SNODE implements the control to high degree of accuracy and without drift.
5.5 Double Pendulum [second order, autonomous, chaotic, conservation law]
Finally, we apply stabilization the chaotic dynamics the frictionless double pendulum system. The total energy E of the system is conserved [77. yielding the manifold,
M {011, 92 (1,44) w2) E 2 R2: E(O1, 02 w1, W2) E. = 0),
(23)
where 0, is the angle of the i-th arm with the vertical and wi = è, We refer the reader to Arnold 77) (or the excellent Wikipedia entry) for the lengthy equations of motion and an expression for the total energy. For simplicity we take my = m2 = 1 kg. 4 = L2 = 1 m, and g = 9.81 ms ms- We train on 40 trajectories, each consisting of 10 seconds equally sampled with At = 0.05, and with initial conditions (0,, 02, W1, w2y = (ò, Φ, 0, 0), where á is drawn randomly from a uniform distribution ~ (TT/44 3T 3T/44 /4). We emphasize that this is 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

---

Relative Error (State)
Relative Error (Energy)
(c)
10:0.
Hellinger Distance
10
0.4
10.10
NODE
10-15
0.3
10
SNODE
10:20
0.2
10
10
0.1
2.5
5.0
250
500
1
Time (seconds)
Time (seconds)
Figure 5: Results for the double pendulum. (a) Relative error in the state 300 short test trials, shown with 95% confidence intervals (shaded). Compared to the SNODE, the NODE diverges rapidly it begins to accumulate errors in the energy. (b) Relative error in the energy averaged over 5 long test trials. (c) Comparison of the double pendulum's 's invariant measure estimated by the (hybrid) UDE, with and without stabilization, versus ground truth, with 95% confidence intervals.
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 short time due to the highly chaotic underlying dynamics. Motivated by hybrid models of the Earth's 's climate 33 7T]. 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 01 are assumed known, while the dynamics of the second arm 02 are learned from data. We train on single trajectory of duration 60 seconds with A/ 0.05. For each trained model, we then integrate ten trajectories of duration one hour longer than the observed data. An invariant measure is estimated from each long trajectory (see Appendix F and compared with the ground truth using the Hellinger distance.
Figure ((c) shows that, as is increased, our ability to accurately learn the double pendulum's invariant measure increases dramatically due 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), method for learning dynamical systems from observational data subject to arbitrary explicit constraints. Our approach is based stabilization term that is cheap to compute and provably renders the invariant manifold asymptotically stable while still admitting all solutions of the vanilla NDE the invariant manifold. Key benefits of our method are simplicity and generality, making it compatible with all common NDE methods without requiring further architectural changes. Crucially, SNDEs allow entirely constraints, such as those arising from known conservation laws and controls, to be incorporated into neural differential equation models. We demonstrate their effectiveness across 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 time-dependent restrictions on the system state. SNDEs are robust with respect to the only tuneable parameter and incur only moderate computational overhead compared vanilla NDEs.
types of
The current key limitations and simultaneously interesting directions for future work include adapting existing methods for NODEs Riemannian manifolds to setting, generalizations to partial dif- ferential 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, non- linear 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 (c.g., human robot interactions), especially when going beyond the training distribution.
10

---

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 [33 and the open- source packages DifferentialEuuationn]]] [66], Flux.jl [46, Zygote.jl [45], ComplexityMaauures, [4] 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
[1] 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.
[2] Pierre-Antoine Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization Algorithms on Matrix Manifolds. Princeton University Press, 2008
[3] Hananeh Aliee, Fabian J. Theis, and Niki Kilbertus. Beyond predictions in neural ODEs: Identification and interventions. arXiv preprint arXiv:2106. 12430, 2021.
[4] Hananeh Aliee, Till Richter, Mikhail Solonin, Ignacio Ibarra, Fabian Theis, and Niki Kilbertus. Sparsity in continuous-deptl neural networks. In Advances in Neural Information Processing Systems, pages 901-914, 2022.
[5] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In Proceedings of The 33rd International Conference on Machine Learning, pages 1 1120-1128. PMLR, 2016.
[6] Ludwig Amold, Christopher K.R.T. Jones, Konstantin Mischaikow, Geneviève Raugel, and Ludwig Amood. Random Dynamical Systems. Springer, 1995.
[7] Vladimir Igorevich Amold. Mathematical Methods of Classical Mechanics. Springer New York, 2nd edition, 1989.
[8] Uri M. Ascher and Linda R. Petzold. Computer Methods for Ordinary Differential Equations and Differential-Algebraie Equations. Society for Industrial and Applied Mathematics, 1998.
[9] Nitin Bansal, Xiaohan Chen, and Zhangyang Wang. Can gain more from orthogonality regularizations in training deep networks? In Advances in Neural Information Processing Systems, pages 4261-4271 2018
[10] Baumgarte. Stabilization of constraints and integrals of motion in dynamical systems. Computer Methods in Applied Mechanics and Engineering, 1(1):1-16, 1972.
[11] Sören Becker, Michal Klein, Alexander Neitz, Giambattista Parascandolo, and Niki Kilber- tus. Predicting ordinary differential equations with transformers. In Proceedings of the 40th International Conference on Machine Learning, pages 1978 27082 PMLR, 2023.
[12] 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 1763. PMLR, 2022.
[13] Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral Shah. Julia: A fresh approach to numerical computing. SIAM Review, 59(1):65 98, 2017.
[14] 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 Machine Learning, pages 1045-1055. PMLR, 2020.
11

---

[15] Kathryn E. Brenan, Stephen L. Campbell, and Linda R. Petzold. Numerical Solution of Initial- Vnitial-Value Problems Differential-Algebene Equations. Society for Industrial and Applied Mathematics, 1995.
[16] 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, 1 113(15):3932-393)
2016.
[17] Mickaël D. Chekroun, Eric Simonnet, and Michael Ghil. Stochastic climate dynamics: Random attractors and deeedeepndent invariant measures. Physica D: Nonlinear Phenomena, 240(21): 1685-1700, 2011.
[18] 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.
[19] Yuhan Chen, Takashi Matsubara, and Takaharu Yaguchi. Neural symplectic form: learning Hamiltonian equations on general coordinate systems. pages 16659-16670, 2021.
[20] Zhijie Chen, Mingquan Feng, Junchi Yan, and Hongyuan Zha. Learning neural Hamiltonian dynamics: A methodological overview. arXiv preprint 2203:0200300//2 2022.
[21] Hong Sheng Chin. Stabilization Methods for Simulations of Constrained Multibody Dynamics. PhD thesis, University of British Columbia, 1995. URI baattpp:///pen..bbaaat ubc. ca/
collections/uiccheeee//////// 000000
[22] Harald Cramer. Mathematical Methods of Statisties. Princeton University Press, 1999.
[23] Miles Cranmer. Interpretable machine learning for science with PySR and SymbolicRegres- sion.jl. arXiv preprini arXiv:2305.01558 2023.
[24] Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley
Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.0466 2020.
[25] George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Control, Signals and Systems, 2(4):303-314 1989.
[26] Raj Dandekar, Karen Chung, Vaibhav Dixit, Mohamed Tarck, Aslan Garcia Valcia-Valadzz Kr- ishna Vishal Vemula, and Christopher Rackauckas. Bayesian neural ordinary differential equations. arXiv preprint arXiv:2012.00244. 2020.
[27] Simon Danisch and Julius Krumbiegel. Makie.jl: Flexible high-performance data visualization for Julia. Journal of Open Source Software, 6(65):3349, 2021.
[28] David Diego, Kristian Agasøster Haaga, and Bjarte Hannisdal. Transfer entropy computation using the Perron-Frooeenus operator. Phys. Rev. E, 99:042212, 2019,,
[29] Thai P. Duong and Nikolay A. Atanasov. Hamiltonian- neural ODE networks the SE(3) manifold for dynamics learning and control. In Proceedings of Robotics: Science and Systems, 2021.
[30] Emilien Dupont, Arnaud Doucet, and Yee whye Teh. Augmented neural ODEs, In Advances in Neural Information Processing Systems, pages 3140-3150, 2019.
[31] Chris Finlay, Joern-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ODE: the world Jacobian and kinetic regularization. In Proceedings of the 37th International Conference Machine Learning, pages 3154-31666 PMLR, 2020.
[32] Mare Finzi, Ke Alexander Wang, and Andrew G Wilson. Simplifying Hamiltonian and La- grangian neural networks via explicit constraints. In Advances in Neural Information Processing Systems, pages 13880-13889, 2020.
[33] Simone Fiori. Manifold calculus in system theory and funddamntass and first-order systems. Symmetry, 13(11), 2021.
12

---

[34] C. Williimm Gear. sefferential-algebrae equation index transformations, SIAM Journal on Scientific and Statistical Computing, 9(1):39-4) 1988.
[35] Maximilian Gelbrecht, Niklas Boers, and Jürgen Kurths. Neural partial differential equations for chaotic systems. New Journal of Physics, 23(4)043005, 2021.
[36] Maximilian Gelbrecht, Alistair White, Sehastian Bathiany, and Niklas Boers. Differentiable programming for earth system modeling. Geoscientific Model Development, 16(11):3123-3135,
2023.
[37] Mevlana C. Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows Riemannian manifolds. arXiv preprint arXiv:1611.02333 1.00344, 2016.
[38] 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-14844 2020.
[39] Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, pages 15379-15389, 2019.
[40] Nate Gruver, Mare Anton Finzi, Samuel Don Stanton, and Andrew Gordon Wilson. Decon- structing the inductive biases of Hamiltonian neural networks. In The Tenth International Conference on Learning Representations, 2022.
[41] Kristian Agasester Haaga and George Datseris. Complexity Zenodo.
[42]
Ernst Hairer. Solving differential equations on manifolds. 2011. URL https: unige ch/ hta/hhairrrppolysddee sdanmmindd .dff
[43] Ernst Hairer, Christian Lubich, and Gerhard Wanner. Springer- Opringer-Verlag, Berlin, 2nd edition, 2006.
Geometric Numerical Integration.
[44] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural Networks, 2(5):359-366) 1989.
[45] Michael Innes. Don't unroll adjoint: Differentiating SSA-form programs. abs/1810.07951, 2018
CoRR,
[46] 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//811.01eeel 1.01557, 2018.
[47] Jacob Kelly, Jesse Bettencourt, Matthew J. Johnson, and David K. Duvenaud. Leaaning differential equations that are easy to solve. In Advances in Neural Information Processing Systems, pages 4370-4388) 2020.
[48] 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.
[49] Patrick Kidger. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
[50] 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.
[51] 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.
[52] Naomi F. 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.
[53] Mario Lezcano Casado. Trivializations for gradient-based optimization on manifolds. In Advances in Neural Information Processing Systems, pages 9157-9168, 2019.
13

---

[54] Mario Lezcano-Casado and David Mtrtinez-Rubble 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.
[55] Yi Heng Lim and Muhammad Firmansyah Kasim. Unifying physical systems' inductive biases in neural ODE using dynamics constraints. arXiv preprint (xXiv:2208.02632 2022.
[56] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, 2019.
[57] 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-17555 2020.
[58] 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.
[59] Yingbo Ma, Vaibhav Dixit, Michael Innes, Xingjian Guo, and Christopher Rackauckas. A comparison of automatic differentiation and continuous sensitivity analysis for derivatives of differential equation solutions. In 202/ IEEE High Performance Extreme Computing Conference (HPEC), pages 1-9. IEEE, 2021.
[60] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. In Advances in Neural Information Processing Systems, pages 2503-2515, 2020.
[61] Y. Nishimori. Learning algorithm for independent component analysis by geodesic flows on orthogonal group. In IJCNN"99. International Joint Conference Neural Networks. Proceedings (Cat. No.99CH36339) volume 2, pages 933-938 vol.2, 1999. doi: 10.1109/UCNN. 1999.831078.
[62] 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 591 15911, 2020.
[63] Avik Pal, Yingbo Ma, Viral Shah, and Christopher Rackauckas. Opening the blackbox: Accel- erating neural differential equations by regularizing internal solver heuristics. In Proceedings of the 38th International Conference on Machine Learning, pages 5325-8335, PMLR, 2021.
[64] Suhan Park, Mathew Schwartz, and Jaeheung Park. NODE IK: Solving inverse kinematics with neural ordinary differential equations for path planning, arXiv preprint arX1v:2209.00498, 2022.
[65] Linda Petzold. Differential/algebraic equations are not ODE's. SIAM Journal on Scientific and Statistical Computing, 3(3):367-384, 1982.
[66] Christopher Rackauckas and Qing Nie. DifferentialEguutionss]] jl-- performant and feature-rich ecosystem for solving differential equations Julia. Journal of Open Research Software, 5(1), 2017.
[67] 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.0433 2020.
[68] 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. PMIR, 2020.
[69] Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence- based generative modeling on manifolds. Advances Neural Information Processing Systems, pages 17680, 2021.
14

---

[70] Yulia Rubanova, Ricky T.Q. Chen, and David K. Duvenaud. Latent ordinary differential equations for -aaple time series. In Advances in Neural Information Processing Systems, pages 5320-5330, 2019.
[71]
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 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.
[72] Takeshi Teshima, Koichi Tojo, Masahiro Ikeda, Isao Ishikawa, and Kenta Oono. Universal ap- proximation property of neural ordinary differential equations. arXiv preprint arXiv:2012.02411, 2020.
[73] Ch. Tsitouras. Kutta pairs of order 5 (4) satisfying only the first column simplifying assumption. Computers & Mathematics with Applications, 62(2):770-775, 2011.
[74] James H. Verner. Numerically optimal Runge-Kutta pairs with interpolants. Numerical Algorithms, 53(2-3):383-396, 2010.
[75] 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 Machine Learning, pages 11066-11000 1095. PMLR, 2020.
[76] Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty, Symplectic ODE-Net: Learn- ing Hamiltonian dynamics with control. In 8th International Conference on Learning Represen- tations, 2020.
[77] Aiqing Zhu, Pengzhan Jin, and Yifa Tang. Deep Hamiltonian networks based on symplectic integrators. arXiv preprint arXiv:2004. 13830, 2020.
15

---

A Differential Algebraic Equations
A differential algebraic equation (DAB) in its most general, implicit form is
F(), = 0,
(24)
where x E __ and F: R x R" RR + R", When aF/Oi is nonsingular, Equation (24) is an implicit ODE and by the implicit function theorem may be written explicit ODE in the form = f(xtt) [3]]. In the more interesting case of singular aF/Oi. an important special case of Equation 24 is given by semi-cxplicit DAEs in Hessenberg form, for example,
ÿ = f(t,y.2)
(25a)
0 g(t y. 2).
(25b)
where = (y. 2). We call the differential variables, since their derivatives appear in the equations. and 2 the algebraic variables, since their derivatives do not. The semi-explicit form of Equation (25)
highlights the connection between certain classes of DABs and ODEs subject to constraints.
generally possible to differentiate the constraints in Equation 25b) a number of times and substitute the result into Equation (25a) to obtain 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 [T5]. Each differentiation of the constraints reduces the index the system by one (ODEs have index O).
Of particular interest in the context of this paper are constrained ODEs of the form
= f(1, 0 g(t, u),
(26a)
(26b)
where u E R" R RM - IR", and y: IR x RR" + Rm Equation (26) can be written as a semi- explicit Hessenberg index-2 DAE, with u as the differential variables and m Lagrange multipliers as the algebraic variables, for example,
à = f(t, u) - D(u) g(l, u):
(27a)
(27b)
where À E I and Duu) is any bounded matrix function such that GD, where G = yn is the
Jacobian of g, is boundedly invertible for all B We can therefore approach the task of solving the constrained ODE Equation (26) from the perspective of solving the Hessenberg index-2 DAE Equation|
DAEs are not ODEs, however, and number of additional complications arise when we seek numerical solutions [65]. Generally, the higher the index, the harder it is solve 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 's stabilization is, in turn, special case of the stabilization procedure later proposed by Chin (20 and adopted by us in this paper. We emphasize, however, that our stabilization procedure addresses 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 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.
16

---

(a)
(b)
Relative Error (Endpoint)
10:2
0.80
el
10:3
0.75
10
0.70
10.8
100
102
104
250
500
Time (seconds)
Time (seconds)
Ground Truth
NODE
SNODE
SNODE (fast)
Figure 6: Comparison of the alternative, faster SNDE formulation in Equation 29 with the standard SNDE formulation in Equation (12). applied the controlled robot experiment. (a) After 100 seconds, both the fast SNODE (green) and the standard SNODE (orange, not visible behind green) continue implement the prescribed path exactly, while the NODE (blue) has drifted significantly. (b) Relative errors in the position of the endpoint, averaged over 10 test trajectories. The fast SNODE (green) admits slightly larger errors than the standard SNODE (orange), although the difference negligible in this setting.
B Faster Alternatives to the Pseudoinverse
The computational cost of the pseudoinverse is O(m²n), where m is the number of constraints and is the dimension of the problem. While cheap compute for the systems studied in this paper, the pseudoinverse may become prohibitively expensive for applications with a large number of constraints. As alternative choice for the stabilization matrix (()) in Equation (5), using the transpose of the Jacobian G(u) Jus i.c.,
we recommend
F(u) G* (u).
(28)
which is instead O1). Since G (u) is symmetric positive definite, the theoretical guarantees of Theorem T and Theorem still apply, such that
fo(u) G (u)g(u)
[alternative SNDE]
(29)
is valid SNDE.
We apply an SNDE in the form of Equation 29
the controlled robot arm experiment, in which the Jacobian has dimensions 2 x 3. In this instance, single evaluation of G" (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 il- lustrates the tradeoff when choosing the stabilization matrix F(u). The pseudoinverse G+ can be expensive to compute but, as orthogonal projection, yields the most "direct" stabilization to the constraint manifold and therefore the most accurate implementation of the constraint. The transpose GT, on the other hand, is always cheap to compute and yields valid stabilization, at the cost of only slightly larger errors in the constraint.
C The Stabilization Parameter Y
C.1 Runtime Implications
We assess the computational cost SNODEs compared to vanilla NODEs. SNODEs require the computation of (and backpropagation through) the pscudoinverse of the Jacobian of the constraint
17

---

Table 1: Training time of NODES vs SNODEs. All experiments are trained for 1,000 epochs on an Intel(R) Xeon(R) CPU E5-2667 v3 @ 3.20GHz. Statistics are calculated over 5 random seeds.
Training Time (seconds)
Two-Body Problem
Rigid Body
DC-to-DC Converter
Model
Mean
Std. Dev.
Mean
Std. Dev.
Mean
Std. Dev.
NODE
10,580
271
9,730
71
14,000
316
0.1
12,060
206
12,000
283
18,060
524
12,180
194
12,500
126
18.020
549
2
12,160
102
12,340
301
18,280
240
SNODE
12,980
147
14,160
280
18,020
354
8
14,000
268
15,320
376
18,200
482
14,260
162
15,680
519
32
15,300
167
17,140
680
function g. Additionally, tolerance, with the SNODE eventually becoming stiff for sufficiently large 7, Naively, one may thus expect noticeable increase in runtime. However, as described in Section the computational cost of training NODEs also depends on the "complexity" the learned dynamics, which in turn determines how many function evaluations between the added computation enforcing constraints and the thereby potentially regularized "simpler" dynamies, which may require fewer function evaluations by the solver.
increased, SNODEs may require more solver steps at given error
required by the solver. This leads to nontrivial interactions
In Table we report comparisons training times between SNODEs and NODEs for different values of y for three settings. SNODEs take roughly 1.2 to 1.8 times longer to train, with smaller values y incurring less overhead. Overall, this manageable increase for most relevant scenarios.
We also show inference times in Table 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 tuning knob that trades off training time versus inference time.
C.2 Constraint Implications
To complement these results, Figure shows that the relative error remains almost unchanged for a large range of y values, that is, beyond certain minimum value, SNODEs not sensitive the specific choice of y. Even a value of y works well in the settings we have considered, indicating that we can typically get away with runtime increases of factor of 1.2. However, larger values of ? only lead to slightly increased training times, while generally enforcing the constraints to higher degree of accuracy.
C,3 A Practical Guide to Choosing ?
As described in the previous sections, choosing the optimal value of requires tradeoff between accuracy and training time; larger values of y 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 y ~ 1 and remains effective as y is increased, until the system becomes stiff around y 100. Within those limits, suggest experimenting with 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 once fo is close to
18

---

10"
Relative Error (Ang. Mom.)
Relalive Error (State)
10
10
10
10
10
10
(c)
(d)
Relative Enor (State)
Relative Error (Casimir)
10
0.1
10
10
10
10
16
10
32
400
800
0
400
800
(e)
(f)
10a
Relative Error (Energy)
Relatie Error (State)
10
10
10
10"
10a
10a
10
10
0
40
80
0
40
80
Time (seconds)
Time (seconds)
Figure 7: Effect of y on relative errors. Top row: Two-body problem (a-b). Middle Rigid body (c-d). Bottom row: DC-to-DC converter (e-f). Beyond certain value, SNDEs are not highly sensitive the choice of 7, although larger values may enforce the constraints more accurately.
Table 2: Inference time and (adaptive) solver statistics SNODEs vs NODEs for the two-body problem experiment. Inference time statistics are calculated using 100 test initial conditions. each of which is integrated for 20 seconds (short enough so that the NODE solution does not diverge). Solver step statistics are reported for single test trial, intended to illustrate the observed trends in inference time. SNODEs are cheaper at inference time due to significantly fewer rejected solver steps.
Model
Inference Time (seconds)
Solver Steps
Type
Median
Mean
Accepted
Rejected
RHS Evaluations
NODE
2.44
2.51 + 0.04
2,379
3,343
34,335
0.1
2.14
2.17 + 0.03
2,040
2,874
29,487
1
2.19
+ 0.03
2,054
2,704
28,551
2
2.22
2.27
2,061
2,541
27,615
SNODE
2.07
2.15
2,180
2,382
27,375
8
2.05
2.7+
2,355
1,877
25,395
16
1.98
2.03 0.05
2,677
1,437
24,687
32
1.96
1.99 0.04
3,219
1,029
25,491
19

---

Table 3: Stable time of NODEs for the same 100 test trials as shown in Figure (not shown) did not diverge during any trial.
SNODE models
NODE Stable Time (seconds)
Experiment
Trial Length (seconds)
Min.
Max.
Median
Mean
Std. Dev.
Two-Body Problem
200.0
52.0
182.8
121.7
117.3
30.3
Rigid Body DC-to-DC Converter
1600.0
341.7
1600.0
1600.0
1349.9
468 4
160.0
19.3
160.0
113.9
110.49
45.6
D Additional Experiments
D.1 Hamiltonian Neural Networks
We train a Hamiltonian neural network (HNN) on the two-body problem, which is Hamiltonian system (Figure 8 The HNN initially conserves angular momentum, suggesting that it has learned reasonable approximation of the true Hamiltonina. However, it becomes unstable after approximately 50 seconds, highlighting the benefits of stabilization when data limited and conservation laws are known.
D.2 Tangent Projection Operator
The invariant manifold of the rigid body experiment is 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 tangent projection operator (TPO),
wu
P(u)
(30)
|
where / the identity matrix. P(u) projects velocities at u to the tangent space of the sphere T., Applying the tangent projection operator of Equation (30) to the standard NDE model fo for rigid body dynamics, we obtain
(())fo(u) E T,S?
(31)
Equation 31) is manifold ODE. In general, specialized numerical methods must be used when inte- grating manifold ODEs 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 ())), 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 characterized by initial drift from the invariant manifold that gives way to a subsequent rapid divergence. practice, however, certain test trials may remain stable for significantly longer than others. In this section, characterize the stable time Taab of an individual test trial as the time elapsed until the relative error E(I) in the predicted system state i()) exceeds given threshold value Estab, i.e.
T Tstab max { E(t) < Estab)
(32)
Taking Estab 10 Table shows Tstab for the the two-body problem, rigid body, and DC-to-DC converter experiments. Across several thousand trials this paper, not one stabilized NDE model diverged.
E Invariant Measure
Given that the double pendulum 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
20

---

NODE
ANODE
SNODE
SANODE
HNN
TPO
(b)
10
10
Relative Error (Ang. Mom.)
Relative Error (State)
10
10
10.
10
50
100
50
100
Time (seconds)
Time (seconds)
(c)
(d)
109
10
Relative Error (Casimir)
Relative Error (State)
10
10
10
10-10
10
10
400
800
400 Time (seconds)
800
Time (seconds)
Figure 8: Top row: Comparison with a Hamiltonian neural network (HNN) for the two-body problem experiment. Bottom row: Comparison with tangent projection operator (TPO) for the rigid body experiment. Relative errors are calculated using the same 100 test initial conditions as in Figure 2 The HNN (a-b, pink) initially conserves angular momentum but is unstable for some initial conditions. The TPO (c-d, navy blue) conserves the Casimir function exactly. However, this does not translate to better prediction of the system state.
of their ability capture the double pendulum's invariant measure. We refer Arnold al. 6 and Chekroun et al. [17] for detailed definitions of invariant measures; in short, measure is said to be invariant under some flow 4 if μ( (()) (A)) = {a(AA for all measurable sets A. Invariant measures are commonly used to characterize the long-tern tergteem dynamical characteristics of chaotic dynamical systems (see, for example, Arnold al. 6 Chekroun al. [17] for a discussion of the invariant measure of the paradigmatic Lorenz-63 63 system). Since the double pendulum is an ergodic system, averages over long times approximate ensemble averages. We can therefore obtain sample of the invariant measure numerically by integrating the system for 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 MeaaureiyyMasuuus..jii 41,, based numerical estimate of the transfer operator. We then use the Hellinger distance 222 to compare the resulting probability distribution with the ground truth value for the double pendulum.
21

---

Table 4: Additional hyperparameters.
Experiment
Two- Two-Body Problem
Rigid Body
DC-to-DC Converter
Double Pendulum
32
8
16
Hidden Layers
2
Hidden Width
128
64
64
128
128
Max LR
10-3
10-4
10-3
10-3
10-2
Min LR
10-5
10-3
10-5 5
10-3
10-4
Augmented Dimension
2
1
F Architecture and Training
Trajectories are generated using the 9()) explicit Runge- Kutta algorithm due to Verner [74]. im- plemented in Differentiallquations, [66] Vern9. To ensure that invariants 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 independent. that is, given trajectory is used exclusively either for training or validation. Each trajectory is split using a multiple-shooting approach into non-overlapping chunks of timesteps each.
Networks are implemented using Flux.jl [66] 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 integrated using the 5(4) explicit Runge- Kutta algorithm due to Tsitouras 73.. implemented in DifferentialEquations]] [66] as Tsit5, with absolute and relative tolerances of
10
The stabilization hyperparameter y 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
22


Reviews:
[{"review_id": "shJWddpn9H", "reviewer": "Reviewer_pbb8", "summary": "The 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. ", "questions": "1. 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? ", "limitations": "Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The 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. ", "weaknesses": "The 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> 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\nWe thank the referee for motivating us to compare our method more generally but are unsure what they mean by “another variation of vanilla NODE”.\nHowever, we have now added two additional baselines: 1. a HNN applied to the two-body problem experiment, and 2. a manifold ODE, based on a spherical tangent projection operator, applied to the rigid body rotation experiment.\nWe hope these additional baselines make the paper stronger in the eyes of the reviewer and we emphasize again that we are not aware of any existing method for constraining neural ODEs that is applicable across all of our examples.\n\n> 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\nWe apologize for not having explained this sufficiently clearly in our original submission.\nWe have carefully revised our paper to ensure that each step is motivated well and that the logical connections between the different parts is clear.\nIn particular, the loss is the usual MSE loss and is not adapted in any way for our application.\nGradients of the ODE solution are computed using adjoint sensitivity analysis and the vector-Jacobian product of the parameterized vector field, *including the stabilization term*, is calculated using automatic differentiation.\n\n> Could the authors discuss further if there are potentially other candidates for the matrix F besides the ones that the author's used?\n\nIn practice we have only used the pseudoinverse $F = G^+ = G^T(GG^T)^{-1}$ because it has the desirable property that it nudges trajectories towards the manifold along the normal direction to the manifold, i.e. the shortest path, and it has proved computationally tractable.\nHowever, the formulation is quite general and other options are certainly possible.\nFor example, Chin [2] even reports competitive results by simply using the transpose of the Jacobian, i.e. $F = G^T$, which may be an attractive alternative for applications where the pseudoinverse becomes expensive, e.g. when there are many constraints.\n\n> 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\nWe thank the reviewer for bringing this to our attention.\nThis was an oversight on our part and we have fixed it in the revised manuscript.\n\n> In Figure 2, what does (SO) NODE represent?\n\nSO(NODE) indicates that the baseline is either a vanilla neural ODE (NODE) or a second-order neural ODE (SONODE), depending on whether the problem itself is first- or second-order.\nWe adopt the term SONODE from [1] where the authors find that the most impactful inductive bias in the HNN and LNN architectures is in fact the second-order structure inherent in the architecture rather than Hamiltonian or Lagrangian structure.\nIn other words, the majority of the benefit of HNNs and LLNs over vanilla first-order neural ODEs is obtained by simply modelling the dynamics as a second-order ODE.\nFor this reason we use SONODE as a more challenging baseline than vanilla NODE for the two-body problem and double-pendulum experiments.\n\n> 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\nYes.\nFor each experiment, we chose the length of the trajectory to provide a reasonable “snapshot” of the system's qualitative dynamics. The overall integration length is in this sense chosen in view of the characteristic time scale of the system. \nFor the double pendulum experiment, we chose a shorter timestep because it leads to better convergence for this case of a chaotic system. We note that Finzi et al. also chose a timestep specific to the system.\n\n> 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\nThis is certainly an interesting question that we will investigate and report on the results.\n\n**References**\n\n[1] Gruver, Nate, et al. \"Deconstructing the inductive biases of hamiltonian neural networks.\" arXiv preprint arXiv:2202.04836 (2022).\n\n[2] Chin, H. S. (1995). Stabilization methods for simulations of constrained multibody dynamics (T). University of British Columbia. Retrieved from https://open.library.ubc.ca/collections/ubctheses/831/items/1.0080031\n\n"}, {"review_id": "saShqGG4GA", "reviewer": "Reviewer_cZV4", "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.", "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", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper has good theoretical analysis and the practical implementation is based on the theories.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> 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 -dimensional, the network can only predict dimensions and use the conservation equation to compute the last dimension directly.\n\nThis is an interesting idea and we thank the reviewer for the insightful suggestion.\nIn practice, this procedure would have to be carried out after every single step of the adaptive solver.\nSolving for the constraint would then be straightforward for linear constraint equations but would quickly become more complicated for nonlinear constraints or when there are multiple constraints, possibly requiring an iterative root-finding procedure at each step (through which gradients must also be computed, which can be problematic, see e.g. [1]).\nWe would also be concerned about the effect of this procedure on the solver's internal error estimates. Since the constrained degree of freedom is not evolved by the solver, its error will not contribute to the overall error.\nFurthermore, if one degree of freedom is calculated based on the network's estimates of the remaining (m-1) degrees of freedom, we would expect that variable’s error to scale as sqrt(m-1).\nWe would also like to draw the reviewer's attention to Beucler et al.'s ACNet [2], which constrains the output of a neural network (as opposed to trajectories of a neural ODE) using a similar approach to the one suggested by the reviewer.\n\n> 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?\n\nWe have formulated the problem in such a way that the constraint is a function of the system's state only (and potentially time).\nIn practice, computing the pseudoinverse has then turned out to be computationally tractable in all cases.\nWe are not entirely sure when the situation described by the reviewer might arise, however it is also possible that we have not understood the question properly.\nIf the reviewer has some specific applications in mind, or would like to expand further, we would be happy to clarify.\n\n[1] Jean Charles Gilbert (1992) Automatic differentiation and iterative processes, Optimization Methods and Software, 1:1, 13-21\n\n[2] Beucler, Tom, et al. \"Enforcing analytic constraints in neural networks emulating physical systems.\" Physical Review Letters 126.9 (2021): 098302.\n"}, {"review_id": "8QmAZLfkdr", "reviewer": "Reviewer_kVz8", "summary": "The 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. ", "questions": "The 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\n\n\n", "limitations": "The authors have addressed limitations, in particular, as a key limitation it is stated that at the current implementation the method is limited to ODEs. ", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "Originality: \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\n", "weaknesses": "Quality:\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.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Despite 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. As 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\nWe thank the reviewer for the suggestion to include more comparisons.\nThe reason we initially only compared against (SO)NODEs is that, despite the various related work, we are not aware of an existing method for imposing arbitrary constraints on NODEs that would be generally applicable to our setting.\n\nIn particular, Finzi et al. incorporate holonomic constraints to enable easier learning of Hamiltonians (and Lagrangians).\nThe applicability of this work to our setting is limited by the fact that we are not solely interested in learning Hamiltonian dynamics, and often we wish to explicitly enforce known conservation laws as inductive biases.\nDespite this, we tried to adapt their method to our two-body problem experiment by imposing conservation of angular momentum as an explicit constraint \"on top of\" the Hamiltonian.\nHowever, this did not converge well.\n\nAs an alternative, **we applied a vanilla HNN to the two-body problem experiment (Figure 2 in the attached PDF)**.\nThe HNN is competitive initially but then diverges for some initial conditions, highlighting the value of our method for enforcing conservation laws as inductive biases.\n\n**We also added another baseline for the rigid body rotation experiment, based on a manifold ODE with a spherical tangent projection operator.**\nWe hope that these additional comparisons have improved the paper in the opinion of the reviewer.\n\nFinally, we also added a discussion of SymODEN as a means of including a Hamiltonian constraint in the related work section.\n\n> Eq (3). I assume there is a missing $\\in \\mathbb{R}^n$\n\nWe thank the reviewer for bringing this to our attention and we have fixed it in the revised manuscript.\n\n> The 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\nThe proof is due to Chin but presented with a minor correction due to us, as well as notation adapted to the present setting.\n\n> For 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\nWe are not sure if we understand this question correctly and apologize in advance if we have misunderstood.\nEq 9 is part of the standard procedure for demonstrating stability of manifolds via suitably defined Lyapunov functions.\nThe constraint manifold sketched in Figure 1 is indeed asymptotically stable since the vector field points towards it from all directions.\n\n> 4.1., 4.2., 4.3., 4.4., 4.5 what are the test settings? what is the value used in these experiments? \n\nWe thank the reviewer for the question and we agree that this was not clear in the manuscript.\nAll test statistics are calculated over 100 trials with initial conditions drawn from the same distribution as the training data (except the double pendulum, where 300 test trajectories are evaluated instead of 100).\nThe gamma values used are summarized in Table 4 of the Appendix and are 8, 32, 8, 16, 16 for the two-body problem, rigid body, DC-to-DC converter, robot arm, and double pendulum, respectively.\nWe have clarified this information in the revised manuscript.\n\n> In section 3 you mention that it might be beneficial to start training without stabilization, do you test this in practice?\n\nWe have not tested this yet as all of the examples we studied converged easily with stabilization enabled throughout.\nHowever, this may become relevant for systems where the constraint is computationally more expensive to compute, e.g. due to a large number of constraints.\n\n> In 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 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 value, or a general advise what is a good starting value etc, which would be informative to the reader as a starting point?\n\nWe thank the reviewer for pointing this out and we have removed the reference to “no tuning”.\nThe reviewer is correct that, like any hyperparameter, gamma requires some tuning, although we emphasize again that the tuning required is minimal.\nThe x-axis on Figure 5(c) is on a log scale, demonstrating that the method is not sensitive to the choice of gamma within a large range.\nMore concretely, we draw the reviewer’s attention to Section B and Figure 6 of the Appendix, where we explicitly test the effect of gamma on the first three experiments.\nThe results demonstrate are consistent with our theoretical guarantees and demonstrate that, beyond a minimum threshold value, SNDEs perform well across a wide range of values for gamma.\nWe agree with the reviewer that the paper would be improved by adding some practical advice on how to choose gamma.\nWe have added this to the Appendix."}, {"review_id": "8wxhXuV8yo", "reviewer": "Reviewer_gVTB", "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.", "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?\n\n- 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- This paper proposes a new stabilization term that makes it possible to learn vector fields that satisfy various types of conditions.\n- The proposed method has been theoretically and experimentally evaluated, and the results are valid.\n- The paper is well written.", "weaknesses": "- Lack of clarity on how the problem setup differs from HNN and LNN.\n- Insufficient discussion of applicability to PDEs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Lack of clarity on how the problem setup differs from HNN and LNN.\n\nWe thank the reviewer for drawing our attention to this shortcoming and we have updated the related work section of the manuscript to further clarify how our method and the problem setup relates to and differs from HNNs and LNNs.\nWe outline these differences below, in response to the reviewer's first question.\n\n> Insufficient discussion of applicability to PDEs.\n\nWe thank the reviewer for raising this point.\nOur method is also applicable to PDEs when solved via a method of lines approach and when the constraints are formulated in terms of the discretized variables.\nIn this context, we would like to emphasize that our method scales linearly in the number of dimensions, implying that our method should remain competitive when extended to PDEs. \nHowever, while we actively work on the PDE case, we do consider a rigorous treatment of PDEs to be beyond the scope of this present paper and would like to outline it as a subject for a follow-up study.\n\n> 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?\n\nYes, when considering conservation laws as constraints, our method requires a closed-form expression for the conservation law in terms of the system state (from which, e.g. the total energy, can then be computed using the initial conditions).\nOur method can thus be used to tackle the rather typical situation where one wants to learn the dynamics of some system for which some fundamental conservation laws are known, but not the precise physical equations of motion.\n\nThis is indeed the key point of difference between our method and HNNs/LNNs, where conservation laws are not necessarily known in advance but can be approximately satisfied if the Hamiltonian/Lagrangian is sufficiently well learned - however, this comes at the cost that these methods are restricted to Hamiltonian / Lagrangian dynamics.\nIn contrast, we often wish to enforce known conservation laws as physics-based inductive biases.\n\nIt is for this reason that we do not consider our method to be a direct alternative to HNNs and LNNs.\nWe once again emphasize the flexibility and generality of our method, which is applicable to systems with first-order dynamics, systems with arbitrary and potentially time-dependent constraints not necessarily arising from conservation laws, as well as hybrid (UDE) systems.\nNone of these problems can be solved using HNNs/LNNs.\n\nAs for the input data, this also highlights key differences between our approach and HNNs/LNNs.\nHNNs require the input data to be expressed in canonical coordinates, while LNNs require generalized coordinates.\nIn contrast, our method places no restriction on the form of the input data.\n\n> 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?\n\nWe are also very interested in applications of our method to PDEs and we consider this an exciting opportunity for follow-up work.\nIndeed, given that the computational complexity of the pseudoinverse scales linearly with the dimension of the system, we expect our approach to generalize well and efficiently to PDEs when suitably discretized.\n\nWith respect to boundary conditions, this is certainly an interesting question that we hadn’t considered before and we thank the reviewer for bringing it to our attention.\nIt should be possible to satisfy PDE boundary conditions in an approximate sense, in that the boundary conditions may be satisfied on average but small violations of the boundary conditions may be possible at any given timestep.\nIt would of course be possible to enforce the boundary conditions more strictly by increasing the value of gamma, however excessively large values of gamma will introduce stiffness to the system.\nOne must therefore consider a trade-off between how accurately the boundary conditions should be satisfied and the number of steps the solver must take."}, {"review_id": "41ucBFhrbx", "reviewer": "Reviewer_RTPu", "summary": "Neural 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.", "questions": "Please see weaknesses section for questions.", "limitations": "There is a brief discussion of limitations of the work in the conclusion section. ", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "As 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.", "weaknesses": "Main 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">There are a host of tools in the Riemannian optimization literature...For example, there exists many algorithms for learning flows or vector fields on manifolds (e.g. [1], [2]) and performing transport on those manifolds.\n\nWe thank the reviewer for bringing our attention to these very interesting papers.\nWhile there are similarities with our work, there are also significant differences.\n\nIn particular, [1] solves manifold ODEs by using a local chart to map between the manifold and its tangent space.\n*This projection depends critically on the availability of the chart, in this case via the exponential map.*\nWe emphasize, however, that the authors of [1] restrict their attention to model Riemannian manifolds with easily derived closed form expressions for the exponential map.\nExtending this technique to the settings studied by us would be cumbersome.\nIn general, it requires a case-by-case construction of the chart, which becomes tedious for nontrivial submanifolds arising from arbitrary algebraic constraints, such as those in our work.\n\nHence, not only is our method “simpler”, but also more general in terms of its applicability.\nWe argue that [1] is much more limited, which is reflected in the systems studied in their paper.\nThis also demonstrates that the systems we study, while not overly complex, are still difficult problems in the context of existing work.\n\n[2] is also a very interesting approach, although again we note that the authors restrict their attention to spheres and products of Euclidean space with spheres, for which they derive closed form expressions for the tangent projection operator.\nExtending this to the arbitrary constraints studied by us is again cumbersome.\nHowever, for our rigid body rotation experiment, which possesses spherical geometry, **we have implemented an additional baseline based on their tangent projection operator for spherical manifolds**.\n\nWe have significantly expanded the discussion of [1], [2] and similar works in the related work section of our revised manuscript to clarify how they relate to and differ from our work.\nIf we have overlooked any aspect of these methods that would enable us to adapt them further to our setting, we would be happy to discuss this with the reviewer.\n\n> 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.\n\nWith the analogy to projected gradient descent (PGD), we assume the reviewer suggests projecting the predicted state back to the constraint manifold after every solver step.\nHowever, a number of complications arise.\nFirstly, if an iterative procedure is used for the projection, this may become expensive (particularly when gradients are also calculated via AD).\nSecondly, care must be taken when differentiating iterative processes via AD, see e.g. [3].\nTo get around these problems, we previously carried out an experiment where we performed a single Newton iteration to project the predicted state *towards* the manifold.\nThe intuition is that, for small errors, one iteration may prove sufficient for near convergence.\nHowever, training did not converge, perhaps due to a flat loss surface (all points in the ambient space that are projected to a given small region of the manifold will have similar loss).\nIt is therefore our experience that PGD is not easily adapted to the current setting.\nAs for Riemannian descent via retractions of the exponential map, this suffers from the same limitations discussed earlier in the context of [1].\n\n[3] Jean Charles Gilbert (1992) Automatic differentiation and iterative processes, Optimization Methods and Software, 1:1, 13-21\n\n> The authors' main contribution is addressing a control problem and not really a learning problem...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.\n\nWe respectfully disagree that our work is outside the scope of this conference.\n*The primary goal of our work is to learn dynamical laws from observational data.*\nThis is a timely and important problem for the core machine learning community, as acknowledged by several reviewers.\nWe also refer to the extensive literature from this conference, e.g. on LNNs and HNNs, that is strongly related to our work in terms of purpose, suggesting our work is well within the scope of NeurIPS.\n\nMore specifically, we assume no controllability of the system in general.\nFor example, in our experiment on the two-body problem, we do not seek to control the orbit of the observed bodies.\nRather, our stabilization incorporates a physics-based inductive bias that helps us to learn the dynamics of the system.\n\nFurthermore, we do not think one can meaningfully discuss the effectiveness of our \"controller\" from a control perspective; in the language of optimal control, there is no \"cost\" to our stabilization.\nThe primary measure of effectiveness we are interested in is generalization error, and from this perspective we have demonstrated empirically that our method is highly effective.\n\n> The theoretical guarantees and experiments, I think, could be expanded...There are no theorems relating to how much better this is for specifically learning tasks (e.g., some proofs about generalization error). \n\nWe have provided theoretical guarantees that our method enforces arbitrary constraints on neural ODEs.\nIn particular, Theorem 1 guarantees that our method can learn the correct dynamics on the manifold, while Theorem 2 guarantees that the learned dynamics are asymptotically stable with respect to the manifold.\nIn addition, we have provided extensive empirical evidence that both theorems are satisfied in practice and demonstrated that this has significant benefits for generalization and stability.\n\n*We will respond to the smaller comments in the discussion phase.*"}]

Author rebuttal:
We thank the reviewers for their thoughtful reviews and insightful feedback. We are glad they found the problem we study important (**RTPu**) and relevant (**kVz8**) while appreciating the novelty (**kVz8**) and simplicity (**RTPu**, **pbb8**) of our approach. We are encouraged that they found our work technically sound (**kVz8**) with theoretical guarantees backed up by experiments (**RTPu**, **gVTB**, **cZV4**). We are also pleased that they found the paper well-written and easy to follow (**gVTB**, **kVz8**, **pbb8**). Finally, we are glad they recognise the potential for our method to be extended to other variations of neural ODE (**pbb8**).

### Additional Comparisons
Several reviewers (**pbb8**, **cZV4**, **kVz8**, **RTPu**) expressed a desire for further empirical evaluations versus existing work. Owing to the generality of our method, we believe there is no existing constrained method that can be applied to all of our examples. However, where certain existing methods can be applied to individual examples, we have added comparisons. In particular,

1. For the two-body problem experiment, we have added a Hamiltonian Neural Network (HNN) [1] baseline. While this does not directly impose conservation of angular momentum as in our example, the dynamics are nonetheless constrained to be Hamiltonian and should conserve angular momentum to a high degree of accuracy.
2. For the rigid body dynamics experiment, where the constraint manifold is a sphere, we have added a baseline based on the spherical tangent projection operator (TPO) used by Ben-Hamu et al. [2], which was cited by **RTPu**.

In the attached PDF we have uploaded a revised version of Figure 2 with the additional comparisons. The HNN is initially competitive with the SNDE but diverges for certain initial conditions (panels b-c), highlighting the value of using our method to directly enforce known conservation laws as physical priors. The TPO implements the spherical constraint exactly, with a relative error in the Casimir function of $\sim 10^{-10}$ versus $\sim 10^{-4}$ for the SNDE (f). However, in practical terms the difference between these relative errors is negligible, highlighted by the fact that the SNDE predicts the state of the system at least as well as the TPO (e).

*These results highlight that, where comparisons can be made to existing work, our method is at least as good or better, while being vastly more general and flexible.*

Finally, we note that we did attempt to adapt the constrained HNN approach of Finzi et al. to impose conservation of angular momentum "on top of" the learned Hamiltonian in the two-body problem example. However, this approach did not converge.
We also note that our double pendulum example is a hybrid / UDE model, intended to highlight the applicability of our method to this setting, and we do not believe the HNN/LNN approaches are applicable here.

### On Our Choice of Experiments
We would once again like to highlight that our five experiments were carefully chosen to showcase the effectiveness of our method across a wide range of settings:
- First- and second-order systems
- Hamiltonian and non-Hamiltonian systems
- Autonomous and non-autonomous systems
- Chaotic and non-chaotic systems
- Purely data-driven and hybrid systems
- Time-dependent and time-independent constraint manifolds

We agree with the reviewers **gVTB** and **RTPu** that an extension to PDEs would make for an interesting follow-up study.

### Revisions
We once again thank the reviewers for their comments on the manuscript. We have made the following changes and clarifications in the revised manuscript, which we believe is significantly improved.
- Significantly expanded the discussion of related methods for Riemannian manifolds and added a comparisons (**RTPu**).
- Significantly expanded the discussion of HNNs and LNNs, along with their constrained versions, and added a comparisons (**kVz8**, **gVTB**).
- Clarified the connection between different parts of the paper, for example, how the stabilization term is related to the loss function and the training procedure (**pbb8**).
- Added a discussion of SymODEN to the related work section (**kVz8**).
- Clarified the test settings in the captions of Figures 2, 4, and 5 (**pbb8**, **kVz8**).
- Added a "how-to guide" to the appendix for choosing the stabilization parameter $\gamma$ (**kVz8**).
- Expanded the discussion of applicability to PDEs (**gVTB**, **RTPu**).
- Fixed incomplete references in the bibliography (**RTPu**).
- Clarified the use of $f_{\theta}(u)$ versus $f_{\theta}(u,t)$ throughout (**RTPu**).
- Added missing $\in \mathbb{R}^n$ in Eq. (3) (**kVz8**, **RTPu**).
- Removed reference to "no tuning" of $\gamma$ in the results section (**kVz8**).

### References
[1] Greydanus, Samuel, Misko Dzamba, and Jason Yosinski. "Hamiltonian neural networks." Advances in neural information processing systems 32 (2019).

[2] Ben-Hamu, Heli, et al. "Matching Normalizing Flows and Probability Paths on Manifolds" Proceedings of the 39th International Conference on Machine Learning, PMLR 162:1749-1763 

Meta-review: The paper focuses on improving the stability of Neural ODE methods by introducing a regularization term which is related to the (pseudo)-inverse of the Jacobian matrix of a conservation function that describes the conservation laws in the system. Extensive experiments on a variety of settings confirms the applicability and relevance. All but one reviewer agree on both the simplicity, elegance, and relevance of the method.

One reviewer finds that the paper lacks sufficient comparison with mathematical literature on Riemannian optimization. In my opinion the authors have sufficiently demonstrated why their method is different from existing *dynamics* literature and how existing literature specializes on particular geometries abstractly, in the context of which the current method focuses on a machine learning perspective that is generalizable. That said, reviewer RTPu has a point that not all comparisons are properly made, specifically with algorithms from the constrained optimization literature.

Given that the focus, however, is on the dynamics, and all other reviewers actually like the paper, I ask the authors to include the suggested literature below, as recommended by reviewer RTPu in their related work, as well as their limitations. Specifically, in the limitations there should be a clarification, that better constrained optimization algorithms exist and that there is ample room for further improvements.

All in all, I find the paper is convincing and I recommend acceptance given the above changes.

Lezcano-Casado, Mario, and David Martınez-Rubio. "Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group." International Conference on Machine Learning. PMLR, 2019.

Kiani, Bobak, et al. "projUNN: Efficient method for training deep networks with unitary matrices." Advances in Neural Information Processing Systems 35 (2022): 14448-14463.

Ablin, Pierre, and Gabriel Peyré. "Fast and accurate optimization on the orthogonal manifold without retraction." International Conference on Artificial Intelligence and Statistics. PMLR, 2022.

M. Arjovsky, A. Shah, and Y. Bengio. Unitary evolution recurrent neural networks. In International Conference on Machine Learning, pages 1120–1128, 2016.

N. Bansal, X. Chen, and Z. Wang. Can we gain more from orthogonality regularizations in training deep networks? 31:4261–4271, 201

M. Lezcano Casado. Trivializations for gradient- based optimization on manifolds. Advances in Neural Information Processing Systems, 32:9157– 9168, 2019

Y. Nishimori. Learning algorithm for independent component analysis by geodesic flows on orthog- onal group. In IJCNN’99. International Joint Conference on Neural Networks. Proceedings (Cat. No. 99CH36339), volume 2, pages 933–938. IEEE, 1999.
Decision: Accept (poster)