SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alistair White
Potsdam Institute for Climate Impact Research
Niki Kilbertus
Hemholtz AI, Munich
Maximilian Gelbrecht
Niklas Boers
{alistair.white, niki.kilbertus, maximilian.gelbrecht, n.boers}@tum.de
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.
Advances in machine learning have recently spurred hopes of displacing or at least enhancing the process of scientific discovery by inferring natural laws directly from observational data. In particular, there has been a surge of interest in data-driven methods for learning dynamical laws in the form of differential equations directly from data [18, 67, 16, 23, 3, 11]. Assuming there is a ground truth system with dynamics governed by an ordinary differential equation
$$
\frac {d u (t)}{d t} = f (u (t), t) \quad (\text {w i t h i n i t i a l c o n d i t i o n} u (0) = u _ {0}) \tag {1}
$$
with  $u:\mathbb{R}\to \mathbb{R}^n$  and  $f:\mathbb{R}^n\times \mathbb{R}\to \mathbb{R}^n$ , the question is whether we can learn  $f$  from (potentially noisy and irregularly sampled) observations  $(t_i,u(t_i))_{i = 1}^N$ .
Neural ordinary differential equations (NODEs) provide a prominent and successful method for this task, which leverages machine learning by directly parameterizing the vector field  $f$  of the ODE as a
(a)
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 a stabilized trajectory (green). The stabilization pushes any trajectory starting away from (but near) the manifold to converge to it at a rate  $\gamma$  (see Section 4).
neural network [18] (see also Kidger [49] for an overview). A related approach, termed universal differential equations (UDEs) [67], combines mechanistic or process-based model components with universal function approximators, typically also neural networks. In this paper, we will refer to these methods collectively as neural differential equations (NDEs), meaning any ordinary differential equation model in explicit form, where the right-hand side is either partially or entirely parameterized by a neural network. Due to the use of flexible neural networks, NDEs have certain universal approximation properties [72, 75], which are often interpreted as "in principle an NDE can learn any vector field  $f$ " [30]. While this can be a desirable property in terms of applicability, in typical settings one often has prior knowledge about the dynamical system that should be incorporated.
As in other areas of machine learning – particularly deep learning – inductive biases can substantially aid generalization, learning speed, and stability, especially in the low data regime. Learning dynamics from data is no exception [4]. In scientific applications, physical priors are often not only a natural source of inductive biases, but can even impose hard constraints on the allowed dynamics. For instance, when observing mechanical systems, a popular approach is to directly parameterize either the Lagrangian or Hamiltonian via a neural network [39, 58, 24]. Constraints such as energy conservation can then be “baked into the model”, in the sense that the parameterization of the vector field is designed to only represent functions that satisfy the constraints. Finzi et al. [32] build upon these works and demonstrate how to impose explicit constraints in the Hamiltonian and Lagrangian settings.
In this work, we propose a stabilization technique to enforce arbitrary, even time-dependent manifold constraints for any class of NDEs, not limited to second order mechanical systems and not requiring observations in particular (canonical) coordinates. Our method is compatible with all common differential equation solvers as well as adjoint sensitivity methods. All code is publicly available at https://github.com/white-alistair/Stabilized-Neural-Differential-Equations.
A first order neural differential equation is typically given as
$$
\frac {d u (t)}{d t} = f _ {\theta} (u (t), t) \quad u (0) = u _ {0}, \tag {2}
$$
where  $u: \mathbb{R} \to \mathbb{R}^n$  and the vector field  $f_\theta: \mathbb{R}^n \times \mathbb{R} \to \mathbb{R}^n$  is at least partially parameterized by a neural network with parameters  $\theta \in \mathbb{R}^d$ . We restrict our attention to ground truth dynamics  $f$  in Equation (1) that are continuous in  $t$  and Lipschitz continuous in  $u$  such that the existence of a unique (local) solution to the initial value problem is guaranteed by the Picard-Lindelöf theorem. As universal function approximators [44, 25], neural networks  $f_\theta$  can in principle approximate any such  $f$  to arbitrary precision, that is to say, the problem of learning  $f$  is realizable.
In practice, the parameters  $\theta$  are typically optimized via stochastic gradient descent by integrating a trajectory  $\hat{u}(t) = \mathrm{ODESolve}(u_0, f_\theta, t)$  and taking gradients of the loss  $\mathcal{L}(u, \hat{u})$  with respect to  $\theta$ . Gradients of trajectories can be computed using adjoint sensitivity analysis (optimize-then-discretize) or automatic differentiation of the solver operations (discretize-then-optimize) [18, 49]. While these approaches have different (dis)advantages [59, 49], we use the adjoint sensitivity method due to reportedly improved stability [18]. We use the standard squared error loss  $\mathcal{L}(u, \hat{u}) = \| u - \hat{u} \|_2^2$  throughout.
We focus on NODEs as they can handle arbitrary nonlinear vector fields  $f$  and outperform traditional ODE parameter estimation techniques. In particular, they do not require a pre-specified parameterization of  $f$  in terms of a small set of semantically meaningful parameters. The original NODE model [18] has quickly been extended to augmented neural ODEs (ANODEs), with improved universal approximation properties [30] and performance on second order systems [62]. Further extensions of NODEs include support for irregularly-sampled observations [70] and partial differential equations [35], Bayesian NODEs [26], universal differential equations [67], and more-see, e.g., Kidger [49] for an overview. All such variants fall under our definition of neural differential equations and can in principle be stabilized using our method. We demonstrate this empirically by stabilizing both standard NODEs as well as ANODEs and UDEs in this paper. Finally, we note that all experiments with second order structure are implemented as second order neural ODEs due to reportedly improved generalization compared to first order NODEs [62, 40].
NODEs were originally introduced as the infinite depth limit of residual neural networks (typically for classification), where there is no single true underlying dynamic law, but rather "some" vector field  $f$  is learned that allows subsequent (linear) separation of the inputs into different classes. A number of techniques have been introduced to restrict the number of function evaluations needed during training in order to improve efficiency, which also typically results in relatively simple learned dynamics [31, 38, 47, 63, 50]. These are orthogonal to our method and are mentioned here for completeness, as they could also be viewed as "regularized" or "stabilized" NODEs from a different perspective.
Hamiltonian and Lagrangian Neural Networks. A large body of related work has focused on learning Hamiltonian or Lagrangian dynamics via architectural constraints. By directly parameterizing and learning the underlying Hamiltonian or Lagrangian, rather than the equations of motion, first integrals corresponding to symmetries of the Hamiltonian or Lagrangian may be (approximately) conserved automatically.
Hamiltonian Neural Networks (HNNs) [39] assume second order Hamiltonian dynamics  $u(t)$ , where  $u(t) = (q(t), p(t))$  is measured in canonical coordinates, and directly parameterize the Hamiltonian  $\mathcal{H}$  (instead of  $f$ ) from which the vector field can then be derived via  $f(q, p) = (\frac{\partial \mathcal{H}}{\partial p}, -\frac{\partial \mathcal{H}}{\partial q})$ . Symplectic ODE-Net (SymODEN) [76] takes a similar approach but makes the learned dynamics symplectic by construction, as well as allowing more general coordinate systems such as angles or velocities (instead of conjugate momenta). HNNs have also been made agnostic to the coordinate system altogether by modeling the underlying coordinate-free symplectic two-form directly [19], studied extensively with respect to the importance of symplectic integrators [77], and adapted specifically to robotic systems measured in terms of their SE(3) pose and generalized velocity [29]. Recently, Gruver et al. [40] showed that what makes HNNs effective in practice is not so much their built-in energy conservation or symplectic structure, but rather that they inherently assume that the system is governed by a single second order differential equation ("second order bias"). Chen et al. [20] provide a recent overview of learning Hamiltonian dynamics using neural architectures.
A related line of work is Lagrangian Neural Networks (LNNs), which instead assume second order Lagrangian dynamics and parameterize the inertia matrix and divergence of the potential [58], or indeed any generic Lagrangian function [24], with the dynamics  $f$  then uniquely determined via the Euler-Lagrange equations. While our own work is related to HNNs and LNNs in purpose, it is more general, being applicable to any system governed by a differential equation. In addition, where we consider fundamental conservation laws as constraints, e.g. conservation of energy, we assume the conservation law is known in closed-form.
The work most closely related to ours is by Finzi et al. [32], who demonstrate how to impose explicit constraints on HNNs and LNNs. The present work differs in a number of ways with the key advances of our approach being that SNDEs (a) are applicable to any type of ODE, allowing us to go beyond second order systems with primarily Hamiltonian or Lagrangian type constraints, (b) are compatible with hybrid models, i.e., the UDE approach where part of the dynamics is assumed to be known and only the remaining unknown part is learned while still constraining the overall dynamics, and (c) can incorporate any type of manifold constraint. Regarding (c), we can for instance also enforce time-dependent first integrals, which do not correspond to constants of motion or conserved quantities arising directly from symmetries in the Lagrangian.
Continuous Normalizing Flows on Riemannian Manifolds. Another large body of related work aims to learn continuous normalizing flows (CNFs) on Riemannian manifolds. Since CNFs transform probability densities via a NODE, many of these methods work in practice by constraining NODE trajectories to a pre-specified Riemannian manifold. For example, Lou et al. [57] propose "Neural Manifold ODE", a method that directly adjusts the forward mode integration and backward mode adjoint gradient computation to ensure that the trajectory is confined to a given manifold. This approach requires a new training procedure and relies on an explicit chart representation of the manifold. The authors limit their attention to the hyperbolic space  $\mathbb{H}^2$  and the sphere  $\mathbb{S}^2$ , both of which are prototypical Riemannian manifolds with easily derived closed-form expressions for the chart. Many similar works have been proposed [37, 14, 68, 60], but all are typically limited in scope to model Riemannian manifolds such as spheres, hyperbola, and tori. Extending these approaches to the manifolds studied by us, which can possess nontrivial geometries arising from arbitrary constraints, is not straightforward. Nonetheless, we consider this an interesting opportunity for further work.
Riemannian Optimization. A vast literature exists for optimization on Riemannian manifolds [2, 53, 61]. In deep learning, in particular, orthogonality constraints have been used to avoid the vanishing and exploding gradients problem in recurrent neural networks [5, 54, 1, 48] and as a form of regularization for convolutional neural networks [9]. Where these methods differ crucially from this paper is that they seek to constrain neural network weights to a given matrix manifold, while our aim is to constrain trajectories of an NDE, the weights of which are not themselves directly constrained.
Constraints via Regularization. Instead of adapting the neural network architecture to satisfy certain properties by design, Lim and Kasim [55] instead manually craft different types of regularization terms that, when added to the loss function, aim to enforce different constraints or conservation laws. While similar to our approach, in that no special architecture is required for different constraints, the key difference is that their approach requires crafting specific loss terms for different types of dynamics. Moreover, tuning the regularization parameter can be rather difficult in practice.
In this work, we vastly broaden the scope of learning constrained dynamics by demonstrating the effectiveness of our approach on both first and second order systems, including chaotic and nonchaotic as well as autonomous and non-autonomous examples. We cover constraints arising from holonomic restrictions on system states, conservation laws, and constraints imposed by controls.
4 Stabilized Neural Differential Equations
General approach. Given  $m < n$  explicit constraints, we require that solution trajectories of the NDE in Equation (2) are confined to an  $(n - m)$ -dimensional submanifold of  $\mathbb{R}^n$  defined by
$$
\mathcal {M} = \{(u, t) \in \mathbb {R} ^ {n} \times \mathbb {R}; g (u, t) = 0 \}, \tag {3}
$$
where  $g\colon \mathbb{R}^n\times \mathbb{R}\to \mathbb{R}^m$  is a smooth function with  $0\in \mathbb{R}^{m}$  being a regular value of  $g$ . In other words, we have an NDE on a manifold
$$
\dot {u} = f _ {\theta} (u, t) \quad \text {w i t h} \quad g (u, t) = 0. \tag {4}
$$
Any non-autonomous system can equivalently be represented as an autonomous system by adding time as an additional coordinate with constant derivative 1 and initial condition  $t_0 = 0$ . For ease of notation, and without loss of generality, we will only consider autonomous systems in the rest of this section. However, we stress that our method applies equally to autonomous and non-autonomous systems.
While methods exist for constraining neural network outputs to lie on a pre-specified manifold, the added difficulty in our setting is that we learn the vector field  $f$  but constrain the solution trajectory  $u$  that solves a given initial value problem for the ODE defined by  $f$ . Inspired by Chin [21], we propose the following stabilization of the vector field in Equation (4)
$$
\boxed {\dot {u} = f _ {\theta} (u) - \gamma F (u) g (u),} \quad [ \text {g e n e r a l S N D E} ] \tag {5}
$$
where  $\gamma \geq 0$  is a scalar parameter of our method and  $F:\mathbb{R}^n\to \mathbb{R}^{n\times m}$  is a so-called stabilization matrix. We call Equation (5) a stabilized neural differential equation (SNDE) and say that it is
<sup>1</sup>The preimage theorem ensures that  $\mathcal{M}$  is indeed an  $(n - m)$ -dimensional submanifold of  $\mathbb{R}^n$ .
stabilized with respect to the invariant manifold  $\mathcal{M}$ . We illustrate the main idea in Figure 1; while even small deviations in the unstabilized vector field (left) can lead to (potentially accumulating) constraint violations, our stabilization (right) adjusts the vector field near the invariant manifold to render it asymptotically stable, all the while leaving the vector field on  $\mathcal{M}$  unaffected.
Our stabilization approach is related to the practice of index reduction of differential algebraic equations (DAEs). We refer the interested reader to Appendix A for a brief overview of these connections.
Theoretical guarantees. First, note that ultimately we still want  $f_{\theta}$  to approximate the assumed ground truth dynamics  $f$ . However, Equation (5) explicitly modifies the right-hand side of the NDE. The following theorem provides necessary and sufficient conditions under which  $f_{\theta}$  can still learn the correct dynamics when using a different right-hand side.
Theorem 1 (adapted from Chin [21]). Consider an NDE
$$
\dot {u} = f _ {\theta} (u) \tag {6}
$$
on an invariant manifold  $\mathcal{M} = \{u\in \mathbb{R}^n;g(u) = 0\}$ . A vector field  $\dot{u} = h_{\theta}(u)$  admits all solutions of Equation (6) on  $\mathcal{M}$  if and only if  $h_\theta |_{\mathcal{M}} = f_\theta |_{\mathcal{M}}$ .
Since  $g(u) = 0$  on  $\mathcal{M}$ , the second term on the right-hand side of Equation (5) vanishes on  $\mathcal{M}$ . Therefore the SNDE Equation (5) admits all solutions of the constrained NDE Equation (4). Next, we will show that, under mild conditions, the additional stabilization term in Equation (5) "nudges" the solution trajectory to lie on the constraint manifold such that  $\mathcal{M}$  is asymptotically stable.
Theorem 2 (adapted from Chin [21]²). Suppose the stabilization matrix  $F(u)$  is chosen such that the matrix  $G(u)F(u)$ , where  $G(u) = g_u$  is the Jacobian of  $g$  at  $u$ , is symmetric positive definite with the smallest eigenvalue  $\lambda(u)$  satisfying  $\lambda(u) > \lambda_0 > 0$  for all  $u$ . Assume further that there is a positive number  $\gamma_0$  such that
$$
\| G (u) f _ {\theta} (u) \| \leq \gamma_ {0} \| g (u) \| \tag {7}
$$
for all  $u$  near  $\mathcal{M}$ . Then the invariant manifold  $\mathcal{M}$  is asymptotically stable in the SNDE Equation (5) if  $\gamma \geq \gamma_0 / \lambda_0$ .
Proof. Consider the Lyapunov function  $V(u) = \frac{1}{2} g^T(u)g(u)$ . Then (omitting arguments)
$$
\frac {d}{d t} V (t) = \frac {1}{2} \frac {d}{d t} \| g (u (t)) \| ^ {2} = g ^ {T} \frac {d g}{d t} = g ^ {T} \frac {d g}{d u} \dot {u} = g ^ {T} G \left(f _ {\theta} - \gamma F g\right), \tag {8}
$$
where we substitute Equation (5) for  $\dot{u}$ . With Equation (7), we have  $g^{T}Gf_{\theta} \leq \gamma_{0}g^{T}g$ , and since the eigenvalues of  $GF$  are assumed to be at least  $\lambda_0 > 0$ , we have  $g^{T}GFg \geq \lambda_{0}g^{T}g$ . Hence
$$
\frac {d}{d t} V \leq \left(\gamma_ {0} - \gamma \lambda_ {0}\right) \| g \| ^ {2}, \tag {9}
$$
so the manifold  $\mathcal{M}$  is asymptotically stable whenever  $\gamma_0 - \gamma \lambda_0\leq 0$
When  $f_{\theta}(u)$  and  $g(u)$  are given,  $\mathcal{M}$  is asymptotically stable in the SNDE Equation (5) as long as
$$
\gamma \geq \frac {\| G (u) f _ {\theta} (u) \|}{\lambda_ {0} \| g (u) \|}. \tag {10}
$$
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  $\mathcal{M}$ .
2.  $\mathcal{M}$  is asymptotically stable in the SNDE for sufficiently large values of  $\gamma$ .
The stabilization parameter  $\gamma$ , with units of inverse time, determines the rate of relaxation to the invariant manifold; intuitively, it is the strength of the "nudge towards  $\mathcal{M}$ " experienced by a trajectory. Here,  $\gamma$  is neither a Lagrangian parameter (corresponding to a constraint on  $\theta$ ), nor a regularization parameter (to overcome an ill-posedness by regularization). Therefore, there is no "correct" value for  $\gamma$ . In particular, Theorem 1 holds for all  $\gamma$ , while Theorem 2 only requires  $\gamma$  to be "sufficiently large".
In the limit  $\gamma \to \infty$ , the SNDE in Equation (5) is equivalent to a Hessenberg index-2 DAE (see Appendix A for more details).
2We note that there is a minor error in the proof of Theorem 2 in Chin [21], which we correct here.
Practical implementation. This leaves us to find a concrete instantiation of the stabilization matrix  $F(u)$  that should (a) satisfy that  $F(u)G(u)$  is symmetric positive definite with the smallest eigenvalue bounded away from zero near  $\mathcal{M}$ , (b) be efficiently computable, and (c) be compatible with gradient-based optimization of  $\theta$  as part of an NDE. In our experiments, we use the Moore-Penrose pseudoinverse of the Jacobian of  $g$  at  $u$  as the stabilization matrix,
$$
F (u) = G ^ {+} (u) = G ^ {T} (u) \left(G (u) G ^ {T} (u)\right) ^ {- 1} \in \mathbb {R} ^ {n \times m}. \tag {11}
$$
Let us analyze the properties of this choice. Regarding the requirements (b) and (c), the pseudoinverse can be computed efficiently via a singular value decomposition with highly optimized implementations in all common numerical linear algebra libraries (including deep learning frameworks) and does not interfere with gradient-based optimization. In particular, the computational cost for the pseudoinverse is  $\mathcal{O}(m^2 n)$ , i.e., it scales well with the problem size. The quadratic scaling in the number of constraints is often tolerable in practice, since the number of constraints is typically small (in Appendix B, we discuss faster alternatives to the pseudoinverse that can be used, for example, when there are many constraints). Moreover, the Jacobian  $G(u)$  of  $g$  can be obtained via automatic differentiation in the respective frameworks.
Regarding requirement (a), the pseudoinverse  $G^{+}(u)$  is an orthogonal projection onto the tangent space  $T_{u}\mathcal{M}$  of the manifold at  $u$ . Hence, locally in a neighborhood of  $u\in \mathcal{M}$ , we consider the stabilization matrix as a projection back onto the invariant manifold  $\mathcal{M}$  (see Figure 1). In particular,  $G(u)$  has full rank for  $u\in \mathcal{M}$  and  $G^{+}G = G^{T}(GG^{T})^{-1}G$  is symmetric and positive definite near  $\mathcal{M}$ . From here on, we thus consider the following specific form for the SNDE in Equation (5),
$$
\boxed {\dot {u} = f _ {\theta} (u) - \gamma G ^ {+} (u) g (u).} \quad [ \text {p r a c t i c a l S N D E} ] \tag {12}
$$
We now demonstrate the effectiveness of SNDEs on examples that cover autonomous first and second order systems with either a conserved first integral of motion or holonomic constraints, a non-autonomous first order system with a conserved quantity, a non-autonomous controlled first order system with a time-dependent constraint stemming from the control, and a chaotic second order system with a conservation law. To demonstrate the flexibility of our approach with respect to the variant of underlying NDE model, for the first three experiments we include both vanilla NODEs and augmented NODEs as baselines and apply stabilization to both (with the stabilized variants labeled SNDROME and SANODE, respectively).
As a metric for the predicted state  $\hat{u}(t)$  versus ground truth  $u(t)$ , we use the relative error  $\| u(t) - \hat{u}(t)\|_2 / \| u(t)\|_2$ , averaged over many trajectories with independent initial conditions. As a metric for the constraints, we compute analogous relative errors for  $g(u)$ .