1 Introduction↩︎

Data-driven reduced order models (ROMs) of dynamical systems arising from high-resolution simulations or experiments can be used as efficient forward-model surrogates. The learned surrogates can be deployed in the traditionally expensive tasks of optimization and parameter inference. However, conventional data-driven ROMs for dynamical systems generally struggle with accurately approximating stiff dynamical systems [1] and maintaining long-term stability [2], [3]. In this work, we address both of these challenges in the context of neural ordinary differential equations (NODEs) [4], which are a class of models that approximate dynamical systems continuously as ODEs using neural networks. The two ingredients that determine the effectiveness of NODE-based models are the neural network architecture and the numerical integrator. For the former we proposed a structure-preserving NODE, which integrates with a linear and nonlinear split and several small components that preserve key mathematical structures strongly. For the latter we incorporate NODE with advanced time stepping through exponential integrators. The use of an exponential integrator improves the performance of explicit and semi-implicit integrators that are commonly used in packages such as diffrax [5] and torchode [6].

In [2], motivated by the inertial manifold theorem, we developed a special NODE approach that embeds a linear and nonlinear partition of the right-hand side (RHS) being approximated by two neural networks. This partitioning method provides better short-time tracking, prediction of the energy spectrum, and robustness to noisy initial conditions than standard NODEs. The recent work [7], [8] extended the linear and nonlinear partition to a semi-implicit NODE. In the forward pass of their training they solve a partitioned ODE using implicit-explicit Runge-Kutta (IMEX-RK) methods. The linear and nonlinear partition allows them to avoid solving a fully nonlinear system. The backward pass is handled through a discrete adjoint approach, resulting into a linear system to be solved. In general, the matrices from these linear systems are dense, requiring a direct solve which can be computationally expensive and have poor scalability. Our NODE uses an explicit approach based on exponential integrators, which do not require direct solvers but enjoy stability properties comparable to implicit approaches. To the best of our knowledge, exponential integration schemes have first been incorporated into NODEs in recent work [9] The authors show that exponential integrators can outperform implicit schemes for low-dimensional ODEs. However, this work shows the straightforward implementation of exponential integrators for high-dimensional systems to be too computationally expensive for practical purposes. Our work relies on several key components to overcome the challenge of training when applying such an exponential integrator with NODE to high-dimensional problems.

The contributions of our work include both structure-preserving enhancements and improvements in time stepping and training. In our approach, a coordinate transformation is learned via an autoencoder in which the dynamics are described by a NODE with a linear and nonlinear partitioning. Our approach is summarized by the schematic in Figure 1. In related work, augmented Neural ODEs [10] expand the hidden state space to enhance expressivity and improve numerical stability. Likewise, the structure-preserving formulation introduced here implicitly augments the state space through a linear-nonlinear partition, achieving similar advantages for modeling stiff systems.

The linear part of the system is parameterized using a novel Hurwitz matrix decomposition based on [11]. This promotes long-term stability through bounding the real parts of the linear operator eigenvalues from above by a specified shift, \(\mu\in\mathbb{R}\). The nonlinear operator is represented using either a bilinear form or a Lipshitz controlled neural network based on the innovations of our previous work [3]. The bilinear form has the advantage of clearly separating the linear and nonlinear operators while a neural network still contains a nonzero linear contribution in it’s Taylor expansion. In either case, we show that a local Lipschitz bound on the nonlinear operator can provably guarantee stable long-term predictions for certain dynamical systems. Lyapunov-based neural networks [12] offer a related perspective, as they enforce stability through explicit Lyapunov-function parameterizations that guarantee monotonic energy decay along system trajectories. While sharing the goal of stability preservation, these methods are rooted in control-oriented formulations, whereas our approach achieves Lyapunov stability implicitly through Hurwitz spectral constraints on the linear operator and Lipschitz control on the nonlinear term, tailored specifically for stiff dynamical systems. The other major contribution of our work involves using an exponential integrator for the time integration. We propose a matrix-free method for advancing the dynamics based on the algorithm developed in [13]. In our numerical examples, we show that the combination of structure-preserving networks and exponential integrators enables us to robustly learn the Robertson system, which is a stiff system bridging time scales over multiple orders of magnitude. For cases with high-dimensional data we use an autoencoder to significantly reduce the dimensionality and aid in our training. We demonstrate using the Kuromoto-Sivashinsky equations that evolving dynamics on a low-dimensional latent space while preserving the linear and nonlinear partition allows a significant reduction in training time while retaining the stability of the model. We note that latent dynamics discovery has seen success in other scientific machine learning approaches (e.g. [14], [15]).

Training and deploying neural ODEs in practical problems is often complicated by stiffness in the dynamics, which occurs when some or all of the solution components undergo rapid changes relative to transient behavior. Stiff behavior forces explicit solvers to take prohibitively small timesteps. A number of recent works have utilized ML to address the challenges of stiff systems. In [16] the stiffness issue is resolved by introducing a suitable reparameterization in time. The map produces a nonstiff surrogate that can be cheaply solved with an explicit method. In [15] the authors utilize a physics-assisted autoencoder to learn a low-dimensonal latent space of a collisional radiative model. In [1] the authors utilize equation scaling to mitigate the stiffness problem. The method proposed in this work is capable of resolving the stiff dynamics by learning a suitable coordinate transformation into a new system composed of a linear operator capturing the stiffness of the system and a nonlinear operator which can be easily integrated using a standard first-order exponential integrator. This capability is highlighted in Roberton’s chemical reaction problem in Section 5.2, in which conventional exponential integrators fail to eliminate the stiff time step restriction resulting from the original form of the equations.

The rest of the paper is organized as follows. Section 2 presents the details behind the parameterizations of the linear and nonlinear parts of NODE along with implementation details. Section 3 gives a result about the Lyapunov stability and provides an error bound for the NODE. Section 4 provides a brief introduction to exponential integrators. The details related to training and their results are given in Section 5. Conclusive remarks are given in 7.

2 Structure-Preserving NODE↩︎

This work focuses on dynamical systems that can be transformed into systems of the form \[\frac{du}{dt} = Au + g(u), \label{eq:linearPlusNonlinear}\tag{1}\] where \(u \in \mathbb{R}^n\), \(g:\mathbb{R}^n\rightarrow\mathbb{R}^n\) is a nonlinear function, and \(A \in \mathbb{R}^{n \times n}\) has eigenvalues whose real parts are bounded above by a known estimate \(\mu\). Many semi-discretized time-dependent PDEs, such as the Navier-Stokes equations and Kuramoto-Sivashinsky equation, can be modeled using this framework.

Our proposed structure-preserving NODE takes the form \[\frac{du}{dt} = A_{L}u + g_{NN}(u). \label{eq:linearPlusNonlinearNODE}\tag{2}\] Here the linear trainable matrix \(A_{L}\) is parameterized using \(A_{L} = A_H + I\mu\), where \(A_H\) has eigenvalues that all have negative real part. The estimate for \(\mu\) is problem specific. For example, if the linearized dynamics are known to be dissipative, then \(\mu = 0\) is a suitable choice. In the case of Kuramoto-Sivashinky considered in Section 5.3, we applied a Fourier analysis to the equations to determine an estimate. Section 2.1 describes how \(A_H\) can be parameterized using a Hurwitz matrix decomposition. The nonlinear operator, \(g_{NN}\), is parameterized using either a bilinear form neural network or a Lipschitz controlled neural network, as described in Sections 2.2 and 2.3, respectively.

2.1 Hurwitz Matrix Parameterization↩︎

A matrix where all the eigenvalues have negative real parts is a Hurwitz matrix. We develop a general decomposition for a Hurwitz matrix based on the following theorem.

This theorem and its proof can be found in [11].

The Cholesky factorization, which involves a product of a lower triangular matrix with its transpose, is a parameterization for symmetric positive definite matrices and involves \(\frac{n(n+1)}{2}\) unique parameters. The Cholesky factorization is used to parameterize \(S\) and \(G_s\), \[\begin{align} S &= L^S(L^S)^{T},\\ G_s &= -L^G(L^G)^{T}, \end{align}\] where \(L^S\) and \(L^G\) are lower triangular matrices. The skew-symmetric matrix \(G_u\) is formed using \[\begin{align} G_u = L_s - L_s^T, \end{align}\] where \(L_s\) is a strictly lower triangular matrix, \(L_s\) and therefore requires \(\frac{n(n-1)}{2}\) unique parameters. The matrix \(A_H\) is thus parameterized by a combination of \(L^S\), \(L^G\), and \(L_s\). The total number of parameters needed to describe \(A_H\) is \(\frac{3n^2+n}{2}\).

2.2 Bilinear Form Neural Network↩︎

One option to approximate the nonlinear part, \(g\), is to use a bilinear form. By construction, the bilinear form has zero contribution to the linear operator since it is built with quadratic forms. This property is not guaranteed when using a vanilla feed-forward network for \(g\). This creates a strict separation of the linear and nonlinear dynamics that allows recovery of the linear operator \(A\) in an equation learning setting. We utilize the low rank-bilinear form network as defined in [3]. It is defined by \[g_{NN}(u) = \sum \limits_{r = 1}^{R} (C^{(r)} u) \odot (D^{(r)} u) + b, \label{eq:bilinearForm}\tag{3}\] where \(u \in \mathbb{R}^{n}, \;C^{(r)} \in \mathbb{R}^{n \times n}, \;D^{(r)} \in \mathbb{R}^{n \times n},\) \(b \in \mathbb{R}^{n \times 1}\), \(\odot\) represents the Hadamard product (element-wise multiplication), and \(R, \;1 \leq R \leq n\) represents the rank of approximation which is chosen much smaller than \(n\) during the training.

2.3 Lipschitz Controlled Neural Network↩︎

The second option we explore to approximate \(g\) with is a neural network with a bounded Lipschitz constant. As shown in Section 3.1, having a Lipschitz bound on \(g_{NN}\) has an impact on the long-term stability of the dynamics. We introduce a novel Lipschitz-controlled neural network architecture. The network is a feedforward network given by \[\begin{align} g_{NN}(u) = l_p \odot l_{p-1} \dots \odot l_1(u) \end{align}\] where \[\begin{align} l_i(z) = \sigma(B_i z + b_i), \qquad B_i = L^{1/p} \frac{W_i}{\max(1,\, \sqrt{\|W_i\|_{\infty} \|W_i\|_1})}, \end{align}\] \(W_i\in\mathbb{R}^{n_i\times n_{i+1}}\) are trainable matrices, \(n_1=n_{p+1}=n\), \(n_2=\dots=n_p=d\) is a chosen latent dimension and \(\sigma\) is a \(1\)-Lipschitz activation function. \(L>0\) is the desired Lipschitz constant of the network \(g_{NN}\).

Note that the \(1\)-norm and \(\infty\)-norm of a matrix are the maximum absolute column sum and the maximum absolute row sum, respectively. Therefore, the realization of \(B_i\) can be implemented in a manner compatible with automatic differentiation. The following theorem shows \(g_{NN}\) has the desired property. The network \(g_{NN}(u)\) given in Definition [def:lcnn] is \(L\)-Lipschitz.

Our previous work [3] introduced a bi-Lipschitz affine transformation (bLAT) network, achieving Lipschitz control using a parameterization based on the singular value decomposition (SVD). The parameteization presented here is more efficient since the SVD is more expensive to evaluate than the \(1\) and \(\infty\) norms. However, unlike the above, the bLAT network is bi-Lipschitz.

2.4 Autoencoder↩︎

Inspired by the seminal work [17] for complex dynamical systems involving high-dimensional data, an autoencoder produces a low dimensional representation of the dynamics which can be integrated efficiently. For high-dimensional problems like the Kuramoto-Sivashinsky equation, we find it is necessary to include the dimension reduction in structure-preserving NODE to avoid training directly on a high-dimensional space, which is more challenging and expensive. For low-dimensional problems such as the Robertson problem in Section 5.2, the autoencoder is additionally used to learn a transformation into coordinates that yields a system of the same dimension but with stiffness coming from the linear part. Instead of simplifying our model by reducing the dimension, the system is transformed such that our exponential integrator is able to integrate utilizing large timesteps.

We consider some dynamical system for \(U(t)\in \mathbb{R}^{m}\) and assume that it can be transformed into a variable \(u\in \mathbb{R}^n\), \(n < m\) where the dynamics are described by a system of the form 1 . Two feed-forward neural networks are trained \(f_E(U;\alpha)\) and \(f_D(u;\beta)\), which represent the encoder and decoder and have network parameters \(\alpha\) and \(\beta\), respectively. The network \(f_E:\mathbb{R}^m \to \mathbb{R}^n\) takes in the initial condition in physical space and network parameters and outputs the initial condition in the latent space \(\mathbb{R}^n\) with \(n\le m\). Once the initial condition is projected onto the latent space we use a time integrator to evolve the dynamics forward in time. By evolving the dynamics in the low-dimensional latent space \(\mathbb{R}^n\), efficiency is gained in the numerical integration step. In mathematical symbols we have \(f_E(U_0) = u_0\) then we evolve \(u_0\) several timesteps (using the exponential integrator in this work) to obtain \(\lbrace u_1, u_2, \dots, u_t \rbrace\). After this, these are each transformed back into the full space data using the decoder \(U_i = f_D(u_i)\), which are then used to compute the loss term. A schematic of this process is given in Figure 1.

3 Analysis of the Network↩︎

In this section two results regarding the proposed structure-preserving NODE are presented. The results include a Lyapunov stability theorem in Section 3.1 and an error bound for a trained model in Section 3.2.

3.1 Lyapunov Stability of the Network↩︎

The NODE is parameterized such that the spectrum of the linear operator \(A_{L}\) is bounded from above by some \(\mu\in\mathbb{R}\). In the case of \(\mu < 0\), we show that any fixed point is an asymptotically stable attractor for the dynamical system defined by the NODE when a weak constraint on the nonlinear part is satisfied. The following theorem is applicable to the case of learning near any fixed point \(u_e\).

If the structure-preserving NODE 2 is parameterized such that the linear operator \(A_{L}\) satisfies \({\rm Re}(\lambda(A_L)) \le \mu < 0\) and \(g_{NN}\) is an \(L\)-Lipschitz function in the neighborhood of a fixed point \(u=u_e\), then for sufficiently small \(L\) the NODE is stable in the sense of Lyapunov at the fixed point \(u=u_e\).

It is important to note that the above Lyapunov stability result of the proposed network is only applicable to a neighborhood of an equilibrium. It does not extend to other attractors or chaotic systems. Nevertheless, we show later empirically that this predicts a good long-time solution for chaotic dynamics when evaluated using appropriate statistical measures. The above assumption of addressing network stability near equilibrium is motivated by the practical stiff dynamics such as chemical reaction and collisional radiative models [15], where the dynamics of interest is a transition from various initial conditions to the corresponding equilibria.

3.2 Bounding the Error of a Learned Model↩︎

Given an estimate on the approximation error of the learned operators, a bound can be obtained for the error between the ground truth and learned dynamics.

4 Exponential Integrators↩︎

The proposed NODE form \(\eqref{eq:linearPlusNonlinear}\) assumes a linear and nonlinear split where the stiffness of the system shall come primarily from the linear matrix \(A\). For such a system, the exponential integrator is an ideal technique to overcome its stiffness, and therefore it is leveraged in this work in the NODE training. In this section a brief review exponential integrators is presented. For an in depth introduction see [19]. Consider a system that may be written in the form \[\frac{du}{dt} = Au + g(u), \quad u(t_0) = u_0, \quad t \geq t_0, \label{eq:linearPlusNonlinear2}\tag{4}\] where \(u \in \mathbb{R}^n, \;A \in \mathbb{R}^{n \times n}\), and \(g:\mathbb{R}^n\rightarrow\mathbb{R}^n\) is a nonlinear function. The solution of 4 satisfies the nonlinear integral equation \[u(t) = e^{(t-t_0)A}u_0 + \int \limits_{t_0}^t e^{(t - \tau)A} g(\tau,u(\tau)) \;d\tau.\] A \(p\)-th order Taylor series expansion of \(g\) about \(t_0\) gives \[u(t) = e^{(t - t_0)A}u_0 + \sum \limits_{k = 1}^{p} \varphi_k((t - t_0)A)(t - t_0)^k g_k + O((t-t_0)^{p+1}), \label{eq:expIntPthOrder}\tag{5}\] where \[g_k = \left. \frac{d^{k-1}}{dt^{k-1}}g(t,u(t)) \right|_{t = t_0}, \quad \varphi_k(z) = \frac{1}{(k-1)!} \int \limits_{0}^{1} e^{(1-\theta)z}\theta^{k-1} \;d\theta, \quad k \geq 1.\]

The functions \(\varphi_l(z)\) satisfy the recurrence relation \[\varphi_{l+1}(z) = \frac{\varphi_{l}(z) - \frac{1}{l!}}{z}, \quad \varphi_0(z) = e^z,\] and have the Taylor series expansion \[\varphi_l(z) = \sum \limits_{k = 0}^{\infty} \frac{z^k}{(k+l)!}.\]

The integrator used in this study is exponential time differencing (ETD) of order 1 [20]. The ETD methods are based on approximating the nonlinear term by an algebraic polynomial. Note that there are many more exotic versions of exponential integrators such as exponential Rosenbrock methods and Lawson methods [19].

Evaluating the \(\varphi\) functions is a nontrivial task. The computation of \(\varphi_1\) is a well known problem in numerical analysis [21]–[23]. Ref. [13] shows that it is possible to avoid computing the \(\varphi\) functions 5 by solving a system with a single matrix exponential of size \((n + p) \times (n + p)\), where \(n\) is the dimension of the system and \(p\) is the order of the approximation. The system 1 is approximated using \[\hat{u}(t) = \begin{bmatrix} I_n & 0 \end{bmatrix} \exp \left((t-t_0) \begin{bmatrix} A & W\\ 0 & J \end{bmatrix} \right) \begin{bmatrix} u_0 \\ e_p \end{bmatrix}, \label{eq:expSystem}\tag{6}\] where \(e_p\) is a vector with a \(1\) in the \(p^{\rm th}\) element and \(0\)’s elsewhere and \(W\in\mathbb{R}^{n\times p}\) and \(J\in\mathbb{R}^{p\times p}\) are defined to be \[\begin{align} J = \begin{bmatrix} 0 & I_{p-1} \\ 0 & 0 \end{bmatrix},\qquad W_{ij} = \begin{cases} g_k, & j=p-k+1 \\ 0, & \text{otherwise} \end{cases}. \end{align}\] \(I_n\) and \(I_{p-1}\) are identity matrices with sizes \(\mathbb{R}^{n\times n}\) and \(\mathbb{R}^{p-1\times p-1}\), respectively. When \(p = 1\), then \(W = u_0\) and \(J = 0\). Thus, for \(p = 1\), 6 reduces to \[\hat{u}(t) = \begin{bmatrix} I_n & 0 \end{bmatrix} \exp \left((t-t_0) \begin{bmatrix} A & g_0\\ 0 & 0 \end{bmatrix} \right) \begin{bmatrix} u_0 \\ 1 \end{bmatrix}. \label{eq:expSystemP1}\tag{7}\]

A matrix-free algorithm developed in [13] is used to compute 7 . This algorithm adapts the scaling and squaring method [24] by computing \(e^{t A}b \approx(T_m(s^{-1} t A))^s b\), where \(T_m\) is a truncated Taylor series. Algorithmically, the computation is performed using \[\begin{align} e^{t A} b \approx \texttt{expmv}(A, b, t) = f_{s}, \end{align}\] where \(f_s\) is defined by the following recurrence relations \[\begin{alignat}{2} b_{1, 0} &= b, \\ b_{i,k} &= \frac{t}{s k} A b_{i, k-1}, \qquad &&i=1,\dots,s, \quad k=1,\dots,m \\ b_{i+1,0} = f_i &= \sum_{k=0}^m b_{i,k}, \qquad &&i=1,\dots,s \label{eq:recursesum} \end{alignat}\tag{8}\] The choice of \(s\) and \(m\) impact both accuracy and efficiency. The number of matrix vector multiplications is \(sm\) and an analysis for the optimal choice of \(s\) and \(m\) may be found in [13]. In practice, we choose \(s\) to be proportional to the maximum time step of the training data and the norm of \(A\), \(s \sim \Delta t_{\rm max} \|A\|_2\), and we choose \(m\) such that 8 converges to some tolerance, \(\tau\). Utilizing the matrix-free computation eliminates the need to compute the matrix exponential in equation 6 , and can provide a faster alternative when the dimension of \(A\) is large. Moreover, the computation of a matrix exponential is required for each trajectory due to the time dependency. Batching of data is improved since a matrix-free algorithm allows for vectorization. This form is used to evolve the dynamics for the structure-preserving NODE in Section 5. Note that higher-order versions of this scheme require the derivatives of the solution.

The matrix-free exponential integrator used in this work is currently implemented as a prototype in high-level Python, and its performance therefore reflects an unoptimized version of the algorithm. Its computational overhead depends on factors such as the number of matrix-vector multiplications, the system dimension, and the convergence properties of the chosen polynomial or rational approximation, which in turn depend on the spectrum of the matrix. In practice, the matrix-free approach scales favorably for large or sparse systems where direct matrix exponentiation or repeated linear solves are computationally expensive. Future compiled or GPU-accelerated implementations have the potential to substantially reduce the overhead, and we expect that further development of this algorithm by the numerical analysis community could yield significant performance gains.

Exponential integrators are explicit methods, but enjoy stability properties comparable to implicit methods. These methods are used to solve two types of stiff problems. The first are systems with a Jacobian whose eigenvalues has large negative real parts. These systems usually arise in the discretization of parabolic partial differential equations. The second type of problems are systems that are highly oscillatory in nature these systems have purely imaginary eigenvalues. The focus of this paper is on the former.

5 Numerical Examples↩︎

In this section the structure-preserving NODE framework is applied to three examples, including a weakly nonlinear system with transient growth, the Robertson’s chemical reaction problem, and the Kuramoto-Sivashinsky equation. The first example highlights the benefits of using an exponential integrator compared to an implicit-explicit (IMEX) integration scheme, which includes improved accuracy when learning dynamics using large time steps. Moreover, we demonstrate that the properties of the exponential integrator allow us to deploy our NODE using different time steps than those used for training. The Robertson example highlights the ability to deal with a problem which has stiff dynamics with time scale separations. Finally, the Kuramoto-Sivashinsky equation demonstrates learning chaotic dynamics using our NODE.

5.0.0.1 Training details

The input to the NODE is a set of initial conditions and the output a the set of trajectories evolved from the initial conditions. For the linear part of the proposed NODE, the proposed parameterization is used for the Hurwitz matrix learning. For the nonlinear part, either a bilinear form or a feed-forward neural network with Lipschitz controlled layers is used. The networks have 2 layers with 100 to 200 hidden dimensions. The activation function is \(\tanh\). The autoencoder is made up of two simple feed-forward neural networks, one for the encoder and one for the decoder. The loss is computed using an approximation to the time integral of the mean squared error of the solution, given by \[\begin{align} L(t, U, \tilde{U}) = \sum_{n=0}^{N_t-1} \frac{1}{2}\left( (U_{n+1} - f_D(\tilde{u}_{n+1}))^2 + (U_{n} - f_D(\tilde{u}_{n}))^2 \right)(t_{n+1} - t_{n}), \label{eq:LossPrediction2} \end{align}\tag{9}\] where the initial latent variable is projected through \(\tilde{u}_0=f_E(U_0)\). Note that the loss is computed in the original coordinates; thus, the autoencoder is optimized along with the NODE during optimization step. The inclusion of this autoencoder network is optional in our framework. Our implementation of the NODE uses Jax [25] as a backend. Automatic differentiation is used to compute gradients of the loss function with respect to the trainable parameters, and the Adam optimizer from the Optax [26] package is used for optimizations.

5.1 Taming Transient Growth↩︎

In this first example we apply our structure-preserving NODE to learn transient growth for a simple system. We focus on a weakly nonlinear dynamical system in which the right-hand-side can be decomposed into the sum of linear and nonlinear operators, \[\frac{dx}{dt} = Ax + \epsilon \sin(x), \quad x(0) = x_0, \label{eq:simpleLinear}\tag{10}\] where \(\sin(x)\) is interpreted to be a component-wise application of the function. The linear operator, \(A\), is chosen such that in the limit of \(\epsilon\to 0\) there is initial transient growth followed by long-term decay. This example is influenced by work done in [27] where the author analyzes how the norm of the matrix exponential contributes to the runtime of quantum algorithms for linear ordinary differential equations (i.e., \(\epsilon=0\)). Some relevant quantities used in their analysis are \[\begin{align} \alpha(A) &= \max \{\text{Re}(\lambda) \;\vert\; \lambda \in \sigma(A) \} \tag{11} \\ \gamma(A) &= \max \left\{\lambda \;\vert\; \lambda \in \sigma\left((A + A^{T})/2\right) \right\}, \tag{12} \end{align}\] where \(\sigma(\cdot)\) represents the the set of eigenvalues of the argument. The values of \((\alpha, \gamma)\) describe the long term behavior of the norm of \(e^{At}\). We focus on the case when \(\alpha < 0\), which implies \(\lim \limits_{t \to \infty}||e^{At}|| \to 0\). While this gives the long term behavior, the quantity \(\alpha\) does not give any insight into the short term behavior around \(t = 0\). The sign of \(\gamma\) tells us whether \(||e^{At}||\) has an initial grown or decay around \(t = 0\), \(\gamma > 0\) corresponding to growth and \(\gamma < 0\) corresponding to decay. To illustrate how \(\gamma\) impacts the short timescale we consider two matrices in \(\mathbb{R}^{2\times 2}\), \[A_{\text{decay}} = \begin{bmatrix} -2 & 1 \\ 0 & -2 \end{bmatrix}, \quad A_{\text{bump}} = \begin{bmatrix} -2 & 10 \\ 0 & -2 \end{bmatrix}. \label{eq:decayAndBumpMatrices}\tag{13}\] As both matrices have eigenvalues equal to \(-2\), both systems have \(\alpha < 0\), but \(\gamma(A_{\text{decay}}) < 0 <\gamma(A_{\text{bump}})\). This shows that \(e^{tA_{\text{decay}}}\) has norm that decays for all \(t > 0\) and \(e^{tA_{\text{bump}}}\) has an initial growth in the norm before the decay as is shown in Figure 2.

For our experiments, we consider \(A=A_{\text{bump}}\). and \(\epsilon\in\{0, 1.0\}\). The initial bump in the norm \(||e^{tA_{\text{bump}}}||\) makes the dynamics more challenging to learn, as this implies initial transient growth before the long term decay. Data is generated using an adaptive Runge-Kutta-Fehlberg scheme. There are 1000 initial conditions evolved using 50 adaptive timesteps and a relative tolerance of \(10^{-8}\). For training, only the data at a few points in time are utilized. For \(\epsilon = 0.0\) only the initial and final timesteps are utilized and for \(\epsilon = 1.0\) only timesteps \(0, 10, 20, 30, 40\), and \(50\) are utilized. Note that the exponential integrator can learn the dynamics for \(\epsilon = 0.0\) with a single timestep, but the IMEX scheme is unable to capture the dynamics. In our network, the nonlinear term is treated using a bilinear form network when \(\epsilon =0.0\) and Lipschitz controlled layers when \(\epsilon = 1.0\). In both cases 2 layers and 100 hidden dimensions are used. The Adam optimizer from Optax is used for training with a learning rate of 0.01 and batch size of 1000. The training is run for 1000 iterations for each case. For comparison, the integration scheme is replaced with the implicit-explicit (IMEX) SSP2(2,2,2) L-Stable Scheme defined in Table II of [28]. The predictions of the models for \(\epsilon=0.0\) are shown in Figure 3 and for \(\epsilon=1.0\) are shown in Figure 4. Learning with the exponential integrator yields more accurate results than the IMEX integrator for both values of \(\epsilon\), the difference is more pronounced for the \(\epsilon = 1.0\).

5.2 Robertson’s Chemical Reaction↩︎

The Robertson problem is a stiff three-dimensional system given by \[\begin{align} \frac{dy_1}{dt} &= -0.04y_1 + 10^4 y_2y_3,\\ \frac{dy_2}{dt} &= 0.04y_1 -10^4y_2y_3 - (3\times 10^7)y_2^2,\\ \frac{dy_3}{dt} &= (3 \times10^7)y_2^2. \end{align}\] This system describes the kinetics of an autocatalytic chemical reaction and is a benchmark problem for stiff ODE solvers.

The training data consists of 1000 trajectories generated using MATLAB’s ODE15s solver. Each trajectory has initial condition \((1,0,0)\) plus a small perturbation. The timesteps are logarithmically spaced in the interval \([4\times10^{-6},4\times10^6]\) with 50 timesteps. The minimum and maximum timesteps are \(3.03\times 10^{-6}\) and \(1.72 \times 10^6\), respectively. The model is trained for 10,000 epochs with a learing rate of 0.001 decaying by a factor of 0.99 after each epoch. In [1] it is shown that training the imbalance in the magnitude of the components can lead to gradient pathologies. The trajectory for \(y_2\) is rescaled to \(\hat{y}_2 = 10^4 y_2\) so that inputs and outputs of the network are \(O(1)\). To recapture the original system all that is required is scaling \(\hat{y_2}\) by \(10^{-4}\).

Two NODEs are created to approximate the Robertson problem: one that utilizes the autoencoder and one that does not. All other model and training parameters remain constant. The nonlinear networks in the latent dynamics consist of 2 layers and 100 hidden dimensions. The batch size is chosen to be 100. The learning rate is initially 0.01 and reduced by a factor of 0.99 at each step. The training is run for 10,000 epochs in each model. Lipschitz controlled layers are used to represent the nonlinear operator. In Figure 5 (a), predictions for both models with initial condition \((1,0,0)\) are displayed. It is clear from this figure that the application of the autoencoder produces an accurate model for all components. Here the autoencoder acts like a coordinate transformation. The latent space dynamics are presented in Figure 5 (b), illustrating how the non-trivial transformation induced by the autoencoder shifts the stiffness of the original system from the nonlinear operator into the linear operator. A non-trivial, stiff, latent dynamics is learned under a large temporal variation. We stress that both versions of the model are able to robustly integrate using timesteps ranging from \(3.03\times 10^{-6}\) to \(1.72 \times 10^6\) without rescaling time. In both cases, the model is learning to extract a stiff linear operator that enables efficient integration on all scales.

Many existing approaches for learning stiff dynamics such as the Robertson problem rely on reparameterizing time and transforming the system into a non-autonomous system. In [16] the authors resolve the stiffness by learning a reparameterization in time. The resulting learned dynamical system is a nonstiff surrogate and may be cheaply evaluated using an explicit integrator. In data driven approaches for stiff dynamics [15], [29] a fixed time transformation is used, such as \(\Delta \tilde{t} = -\frac{1}{\log{\Delta t}}\). Our proposed approach avoids the need to know the time scale change a prior or learning the time reparameterization. The proposed approach can therefore bridge large variations in time scales–a desirable property for multi-scale problems that are often found in plasma physics and fluid dynamics. We observe that our approach performs better than [16] with a much smaller error for the same test, particularly in the initial stage of the dynamics. [16] obtains a worse result for this particular case due to its working on a rescaled dynamics, which does not prioritize the error in the true dynamics. The only other approach that we are aware of learning directly in the original time domain is [30], in which a multiscale structure is explicitly built in neural ODE.

Another interesting finding from this example is that the proposed approach enables ETD for a system that might otherwise be considered unsuitable for it. Specifically, one may not be able to use ETD to integrate the Robertson problem in its present form due to the stiffness being in the nonlinear part. Although an exponential integrator allows for larger timesteps than many standard explicit methods, a larger step can be only achieved if the variables are transformed so that the stiffness is transferred to the linear terms. This transformation is learned through data when we use an autoencoder to learn a favorable coordinate system. That is, two neural networks are applied, one to transform into the favorable coordinates for evolution and one to transform back into the original coordinates to compute the loss. These transformations along with restricting the linear operator to be Hurwitz stable are ingredients that enable robust training and, more importantly, efficient forward predictions for a challenging system such as Robertson.

5.2.1 Network Ablation Studies↩︎

To evaluate the robustness of the proposed network, we perform two ablation studies: (1) assessing the model’s sensitivity to noisy inputs and (2) analyzing the effect of the Lipschitz constraint on the nonlinear operator. In the first study, Gaussian noise is added to each trajectory used for training, with noise levels scaled by factors of 0.01 and 0.1. As shown in Figure 6, the model remains robust across all noise amplitudes, demonstrating strong resilience to perturbations in the training data. Table 1 presents the results of the second study, which investigates the impact of varying the Lipschitz constant on model performance. Three configurations are compared: no constraint, \(L=1.0\), and \(L=2.0\). Additionally, the scaling factor applied to the initial weights of the nonlinear function is varied from \(10^{-3}\) to \(10^0\). The results indicate that imposing a Lipschitz constraint significantly improves accuracy and stability across all initialization scales. In contrast, models trained without the constraint exhibit less consistent and occasionally unstable behavior. These findings also highlight that the performance remains sensitive to the magnitude of the initial weights, underscoring the importance of proper initialization when applying Lipschitz regularization. In each case the learning rate is initialized at 0.01 and has a decay rate of 0.9. The network has 100 hidden dimensions and 2 layers. A total of 20000 iterations are done and the final loss is reported in Table 1.

5.3 Kuramoto-Sivashinsky Equation↩︎

In this section we consider the one-dimensional Kuramoto-Sivashinsky (KS) equation \[\frac{du}{dt} = - u\frac{du}{dx} - \frac{d^2u}{dx^2} - \frac{d^4u}{dx^4}, \label{eq:ksEQN}\tag{14}\] with periodic boundary conditions on a domain of length \(L = 22\). The KS equation exhibits chaotic behavior; therefore, we do not expect pointwise accuracy in the long term. To check the validity of our solution we rely on statistical measures, namely, the probability distribution function of the first and the second spatial derivatives, \(u_x\) and \(u_{xx}\).

We consider the dataset used in [2]. Solutions are found by performing a Galerkin projection onto Fourier modes and using exponential time differencing to evolve the ODE forward in time [31]. After solving in Fourier space, we then transform the data back to physical space for training the neural ODEs. We sample the data at every 1.0 time units. The training data is constructed by separating the long trajectory into trajectories with 8 steps. Our batch size is 2000 and our learning rate is 0.001. The learning rate is reduced by a factor of 0.99 at each step. For this problem we use a Lipschitz controlled feed-forward neural network for the nonlinear part using 2 layers and inner dimension 200. We observe that Lipschitz control simultaneously improves training performance and the long-term stability of the dynamic model.

There are a few additional steps taken to train this model effectively. First, a shift of the spectrum to the Hurwitz matrix is applied during training. This is required since the KS system has both positive and negative eigenvalues and the Hurwitz parametrization requires the real part of the eigenvalues to be negative. Second, an autoencoder that projects the dynamics down to a latent space. There are many attractive options for model reduction. A popular method is utilizing projection based reduced-order modeling, which attempt to find a low-dimensional trial subspace that can represent the state space. Most of these trial subspaces are linear and can not capture the complicated dynamics of the KS equation. Moreover, the KS equation has a slow decaying Kolomogorov \(n\)-width [32], which implies linear trial subspaces may not accurately represent the system. The power behind autoencoders is that they attempt to find an identity mapping with fewer dimensions. If the data can be represented in a lower dimensional nonlinear manifold, then optimization of the encoder and decoder finds the optimal low-order representation. For this example 64 grid points are projected into a latent space of dimension 20 using an autoencoder. For \(L=22\), the KS equation has 3 positive Lyapunov exponents, resulting in the Kaplan-York dimension of its attractor to be 5.198 [33]. Therefore, it needs an inertial manifold of at least 6 dimensions. Our second model utilizes a latent space of 7 dimensions. This finding is consistent with the previous work [34], where the authors found that for the exact same KS test, the minimal latent dimension for their model to train reasonably is 7. We note that relative to our baseline architecture without an autoencoder, we observe factors of 12 and 24 speedup for the models with dimension 20 and 7, respectively.

In Figure 7 the evolution of the true solution (\(U_{\text{true}}\)), latent space dimension 20 model and latent space dimension 7 model (\(U_{20},U_7\)), and the errors (\(\text{error}_{20},\text{error}_7\)) are displayed until time \(t = 200\). The dynamics agree pointwise until around time \(t = 50\) for latent space dimension 20 and until time \(t = 25\) for latent space dimension 7. After the time of pointwise agreement we must rely on other measures to assess the validity of the model.

To verify the long-time behavior of predicting this chaotic system, Figure 8 presents the joint probability distribution function of \(u_x\) and \(u_{xx}\). Qualitatively they match well demonstrating that the statistics of the dynamics agree between both models and the ground truth.

6 Discussions↩︎

A critical component of solving high-dimensional problems is dimensionality reduction. We pursue this through an autoencoder that compresses the state variables into a latent representation whose evolution satisfies an ordinary differential equation. Our approach can be viewed as an extension of latent neural ODEs [35], but with additional structure-preserving mechanisms built into the latent dynamics. Latent neural ODEs learn a low-dimensional latent space in which continuous-time dynamics evolve, enabling the modeling of irregularly sampled or partially observed trajectories. Our framework extends this idea by introducing Hurwitz-stable linear operators, Lipschitz-controlled nonlinear components, and exponential integration, ensuring stable and efficient learning even for stiff or multiscale systems. This structure makes the method particularly suitable for physical systems where long-term accuracy and numerical robustness are essential. Beyond latent ODEs, there are other extensions of the neural ODE framework designed to enhance representational capacity. Augmented neural ODEs [10] expand the state space with additional coordinates, allowing the model to represent more complex trajectories and overcome topological limitations of standard NODEs. In our examples, which involve first-order systems solved directly for the full state variable, augmentation is not physically motivated. However, the principles of exponential integration, Hurwitz stability, and Lipschitz control can naturally extend to augmented formulations. For PDEs in multiple spatial dimensions, the number of degrees of freedom in the discretization grows rapidly with the spatial dimension. To address such large-scale problems, it may be necessary to employ sparse or localized operators, such as convolutional, patch-based, or graph-based embeddings, to perform efficient coarse-graining before applying the learned structure-preserving neural ODE. Extending these ideas to high-dimensional PDEs represents an important step toward modeling realistic multi-scale physical phenomena, including fluid turbulence, plasma dynamics, and climate processes.

The proposed framework shows potential for extension to more complex dynamical systems beyond the examples considered in this work. In particular, it can naturally be applied to multi-scale systems, where components evolve on widely separated time scales. In our recent work on fast-slow dynamical systems [3], we demonstrated that a structure-preserving formulation can accurately capture the fast, dissipative dynamics through a linear-nonlinear split while maintaining stability and efficiency. The present approach builds on these ideas, and can be directly applied to the fast subsystem to achieve stable integration across multiple scales, while future work will target systems that exhibit strongly oscillatory or resonant behavior. For non-smooth systems, such as those involving shocks, contact discontinuities, or material interfaces, the performance of the method will depend primarily on the spatial discretization employed. Nevertheless, the exponential time-stepper introduced in this paper can be readily combined with shock-capturing or discontinuity-aware spatial schemes, enabling robust integration even in the presence of non-smooth features. We also note that hybrid dynamical systems, which involve discrete-continuous interactions or state-dependent switching, fall outside the continuous ODE framework considered here. Extending the current formulation to handle event-driven updates or discontinuous vector fields would require additional mechanisms, such as switching-aware integration or piecewise structure-preserving embeddings. Addressing these challenges represents a compelling direction for future research toward applying structure-preserving neural ODEs to a broader class of real-world, multi-physics problems.

7 Conclusions↩︎

This work presents a novel approach to addressing the long-term stability issues of neural ordinary differential equations (NODEs), a crucial aspect for efficient and robust modeling of dynamical systems. By combining a structure-preserving NODE with a linear and nonlinear split, an exponential integrator, constraints on the linear operator through Hurwitz matrix decomposition, and a Lipschitz-controlled neural network for the nonlinear operator, the proposed approach demonstrates significant advantages in both learning and deployment over standard explicit and implicit NODE methods. This approach enables the efficient modeling of complex stiff systems and provides a stable foundation for deploying NODE models in real-world applications. We demonstrate our approach in various examples. For the weakly nonlinear ODE example, we show that using an exponential integrator provides significant accuracy over a semi-implicit scheme. We also demonstrate the structure-preserving NODE approach using stiff ODEs of the Robertson chemical reaction problem and PDEs of the Kuramoto-Sivashinksy system, showcasing its ability to learn both multi-scale and chaotic systems.

While our experiments focus on canonical stiff systems, the proposed framework is general and can be applied to other domains. In particular, neural ODEs trained on image, speech, or sequential data may exhibit effective stiffness during optimization, making our structure-preserving approach potentially beneficial for broader deep learning tasks. We plan to investigate this connection in future work by evaluating the method on standard benchmarks where stiffness may arise implicitly during training. Beyond deep learning, the framework can also be extended to real-world scientific and engineering problems such as collisional-radiative modeling, climate and geophysical forecasting, and computational fluid dynamics, where stiffness and multi-scale interactions are prevalent. These applications represent natural testbeds for assessing the scalability, stability, and interpretability advantages of structure-preserving neural ODEs in realistic, high-dimensional environments.

Acknowledgments↩︎

This research used resources provided by the National Energy Research Scientific Computing Center (NERSC), a U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory, operated under Contract No. DE-AC02-05CH11231 using NERSC award ASCR-ERCAP0023112. We thank Prof. A. J. Roberts from University of Adelaide on the suggestion related to the last example training.

Data Availability↩︎

The code and datasets analyzed during the current study are available from the corresponding author on reasonable request.

References↩︎