Intelligent AttractorsD. A. Serino, A. Alvarez Loya, et al.

# Intelligent Attractors for Singularly Perturbed Dynamical Systems††thanks:
Submitted to the editors February 2024. Corresponding author: This work was
partially supported by the U.S. Department of Energy Advanced Scientific
Computing Research (ASCR) under DOE-FOA-2493 “Data-intensive scientific
machine learning”. It was also partially supported by the ASCR program of
Mathematical Multifaceted Integrated Capability Center (MMICC). Alvarez Loya
is supported by the U.S. National Science Foundation under NSF-DMS-2213261.

Daniel A. Serino Los Alamos National Laboratory, Los Alamos, NM (, , ).
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>Allen Alvarez
Loya22footnotemark: 2 J. W. Burby Department of Physics, University of Texas
at Austin, Austin, TX<EMAIL_ADDRESS>Ioannis G. Kevrekidis
Department of Chemical and Biomolecular Engineering and Department of Applied
Mathematics and Statistics, Johns Hopkins University, Baltimore, MD ().
<EMAIL_ADDRESS>Qi Tang22footnotemark: 2

###### Abstract

Singularly perturbed dynamical systems, commonly known as fast-slow systems,
play a crucial role in various applications such as plasma physics. They are
closely related to reduced order modeling, closures, and structure-preserving
numerical algorithms for multiscale modeling. A powerful and well-known tool
to address these systems is the Fenichel normal form, which significantly
simplifies fast dynamics near slow manifolds through a transformation.
However, the Fenichel normal form is difficult to realize in conventional
numerical algorithms. In this work, we explore an alternative way of realizing
it through structure-preserving machine learning. Specifically, a fast-slow
neural network (FSNN) is proposed for learning data-driven models of
singularly perturbed dynamical systems with dissipative fast timescale
dynamics. Our method enforces the existence of a trainable, attracting
invariant slow manifold as a hard constraint. Closed-form representation of
the slow manifold enables efficient integration on the slow time scale and
significantly improves prediction accuracy beyond the training data. We
demonstrate the FSNN on several examples that exhibit multiple timescales,
including the Grad moment system from hydrodynamics, two-scale Lorentz96
equations for modeling atmospheric dynamics, and Abraham-Lorentz dynamics
modeling radiation reaction of electrons in a magnetic field.

###### keywords:

Singularly perturbed dynamical systems, slow manifold, model discovery,
closures

##  1 Introduction

In this article, a new architecture called the Fast-Slow Neural Network (FSNN)
is developed for learning data-driven models of fast-slow dynamical systems
with dissipative fast timescale dynamics. Our method enforces the existence of
a trainable, attracting invariant slow manifold [10] as a hard constraint. The
constraint is imposed by way of a novel neural network architecture with
structured weight matrices. By leveraging the neural ODE framework from [11],
our technique produces deterministic evolution equations that approximately
reproduce trajectory training data. When the training data comes from an
unknown dynamical system, the method achieves data-driven _full-order_ model
discovery with stability guarantees due to provable existence of the slow
manifold. In short, our architecture strongly contains an attractor of
singularly perturbed dynamical systems by design.

Various dynamical systems encountered in nature exhibit a timescale separation
such that the fast timescale is dissipative. Prominent examples include the
kinetic dynamics of strongly-collisional gasses and plasmas, and radiative
transfer in optically-thick media. For ordinary differential equations (ODEs)
of this type, Fenichel theory [21] reveals that fast timescale dissipation
implies existence of an invariant slow manifold that attracts nearby
trajectories and governs the long-term dynamics. For partial differential
equations (PDEs), or more generally integro-differential equations, Fenichel
theory generally only applies formally, but its predictions are often valid
anyway. For PDEs, inertial manifold theory [24] has been studied and shown to
be useful in applications such as model reductions of fluid dynamics.

Nonlinear dimensionality reduction inspired by inertial manifold theory has a
rich history (see [36] and the references therein) originating in the 1980s
and 1990s [22, 24, 23, 33, 56, 57, 31, 28]. Data-driven approaches for
discovering inertial manifolds in dynamical systems and using separations of
scales to evolve a lower-dimensional reduced order model first materialized in
the 90s [34, 59, 37, 38, 55, 53, 54] and recently reemerged in popularity due
to boom in machine-learning-based modeling [4, 12, 48, 3, 39, 42, 61, 43, 60,
16, 44, 41]. The work in this manuscript tackles a special class of systems
with attracting invariant sets, namely fast-slow systems with dissipative fast
timescale dynamics [10]. Unlike previous data-driven approaches, the FSNN is a
_full-order_ model which, through the use of neural network architectures with
special properties, strictly enforces the existence of an attracting slow
manifold. This method may be viewed as a major step along a pathway in
developing data-driven model discovery techniques with embedded invariant
manifolds that do not use explicit timescale separation as a crutch.

The secondary goal of the proposed network is to discover an efficient and
asymptotic-preserving _reduced-order_ model. In the common methodology of
physics-based model reduction, one first decides a latent space (e.g. lower-
order moments of a distribution) and then looks for a closure using tools such
as asymptotic analysis. Many complicated closures have been discovered through
this approach. The classical examples in rarefied gases and plasma physics
include Grad’s moment method [27], Braginskii closure [8] and Hammett-Perkins
closure [29]. The challenge of this methodology is that an exact closure may
not exist. The accuracy of the approximate closure is related to the dimension
of the latent space. In contrast, Fenichel theory guarantees the existence of
an exact closure in the neighborhood of the singular limit of the equations in
the form of a diffeomorphic coordinate transformation. This transformation
exposes an invariant manifold by flattening the fast variable space. Often in
practice, as asymptotic expansion is the best analytical approximation to the
exact closure. In our work, we leverage structured neural network
architectures and optimization to learn an approximation to the exact closure
that is hidden in the dynamics. The data-driven approach described in this
work can discover such a transformation only using short-term trajectory data.
Therefore, this approach may be useful when the underlying dynamical system is
known but the targeting latent space or a formal closure are not known.

The design of the FSNN architecture is motivated by the the first few theorems
discovered by Fenichel [35]. The components of the architecture are
constructed using structured weight matrices. The coordinate transformation
into Fenichel normal form is based on an invertible coupling flow network [58]
in which we introduced the additional structure-preserving property of bi-
Lipschitz control through the use of bi-Lipschitz affine transformation (bLAT)
layers. This network and its inverse can be evaluated using closed-form
expressions with identical computational complexity. We note that “vanilla”
neural networks are generally non-invertible [25, 52, 14]. We also introduce
the Schur form network, which is a novel architecture based on the Schur
decomposition for controlling the locations of eigenvalues in the complex
plane, to parameterize the linear term in the dynamics. Additionally, a low-
rank bilinear form network is included to increase the approximation power in
the fast variable space. Our approach is then integrated using adaptive time
stepping that treats the linear term implicitly, further leveraging back-
substitution on the Schur form. Some of the structures introduced are related
to the idea of trivializing manifold optimization through parameterizations
[40]. Structure-preserving flow maps through machine learning are an active
area of research, and include such as the symplectic neural networks for
Hamiltonian dynamics [32, 9]. Previous studies on flow map learning for multi-
scale dynamical systems include [46] and [20]. In particular, [20]
simultaneously preserves symplecticity and adiabatic invariance in a multi-
scale setting.

Before going into a detailed presentation of our new method, it is worth
highlighting an important shortcoming it suffers from, and how that
shortcoming might be mitigated in practice. The method assumes that the
dimension of the attracting invariant manifold is known in advance. While this
is not a limitation in the surrogate modeling context, where the (expensive to
simulate) dynamical system is known in detail, it does present a moderate
challenge in the model discovery context. As such, when performing data-driven
model discovery using our method, it will be necessary to first determine the
dimension of the slow manifold, either through physical reasoning, or through
application of auxiliary data mining techniques such as diffusion maps [13],
autoencoders [61], or other techniques from reduced order modeling [36].

The remainder of the manuscript is organized as follows. A review of fast-slow
dynamical systems is presented in Section 2. The fast-slow neural network
architecture is described in detail in Section 3. The numerical implementation
of the architecture is presented in Section 4. Results from applying the
architecture to various examples are shown in Section 5. Concluding remarks
are made in Section 6.

##  2 Fast-Slow Dynamics and the Fenichel Normal Form

In this paper, we develop a novel machine learning (ML)-inspired architecture,
the Fast-Slow Neural Network (FSNN). FSNNs are a parameterization for a class
of dynamical systems exhibiting multiscale behavior in time that, to leading
order, are dissipative on the fast time scale. Fast-slow systems, which
exhibit both short-term and long-term behavior, are relevant to many physical
systems observed in nature, including the climate and plasma systems. In this
section, we motivate the FSNN architecture with a review the theory of fast-
slow systems.

A fast-slow system is an ODE of the form [10]

(1a) $\displaystyle\epsilon\frac{d}{d\tau}y$ $\displaystyle=f(x,y,\epsilon),$
(1b) $\displaystyle\frac{d}{d\tau}x$ $\displaystyle=g(x,y,\epsilon),$

where $x=x(\tau)\in\mathbb{R}^{N_{x}}$, $y=y(\tau)\in\mathbb{R}^{N_{y}}$, $f$
and $g$ are smooth functions of their arguments, and $D_{y}f(x,y,0)$ is
invertible for all $(x,y)$ where $f(x,y,0)=0$.

In the above definition, $y$ and $x$ are the fast and slow variables,
respectively, over corresponding timescale variables $t$ and $\tau=\epsilon
t$.

In some instances, a dynamical system may exhibit a fast-slow split, but a
representation for the fast and slow variables is unknown. Consider the
following representation of a dynamical system for $z(\tau)\in\mathbb{R}^{N}$,
$\epsilon>0$,

(2) $\displaystyle\epsilon\frac{d}{d\tau}z$ $\displaystyle=U(z,\epsilon),$

where $U$ is a sufficiently smooth function of $z$ and $\epsilon$.

The dynamical system (2) exhibits a fast-slow split if there is an
$\epsilon$-dependent invertible transformation of the dependent variables into
fast-slow coordinates $(x,y)$, $x\in\mathbb{R}^{N_{x}}$,
$y\in\mathbb{R}^{N_{y}}$, $N=N_{x}+N_{y}$, where the dynamics are described by
a fast-slow system.

In this paper, we restrict our attention to the normally stable fast-slow
system, which is a fast-slow system where the fast directions are stable.

A normally stable fast-slow system is a fast-slow system where the eigenvalues
of the Jacobian, $D_{y}f(x,y,0)\in\mathbb{R}^{N_{y}\times N_{y}}$, are
strictly to the left of the imaginary axis when $f(x,y,0)=0$.

Due to the dissapitve fast-scale dynamics of normally stable fast-slow
systems, the solutions will be attracted to an invariant slow manifold in
forward time. We will now define the notion of the slow manifold for normally
stable fast-slow systems. Consider a fast-slow system in the form (1). For
$\epsilon>0$ sufficiently small, there exists a function, $y^{*}(x,\epsilon)$,
such that the graph

(3) $\displaystyle M_{\epsilon}=\\{(x,y):y=y^{*}(x,\epsilon)\\},$

is locally invariant under (1). We define $M_{\epsilon}$ to be the slow
manifold.

This lemma is proved following Theorem 4 in [35]. The purpose of the FSNN is
to be a universal approximator for normally stable fast-slow systems. The
network is based on the following well known result due to Fenichel.

When $\epsilon>0$ is sufficiently small and $(x,y)$ lies is in a neighborhood
around the slow manifold $M_{\epsilon}$, there exists a change of coordinates
$(x,y)\rightarrow(\tilde{x},\tilde{y})$ where the dynamics of normally stable
fast-slow systems can be written as,

(4a) $\displaystyle\frac{d}{dt}\tilde{y}$
$\displaystyle=A(\tilde{x},\tilde{y},\epsilon)\tilde{y},$ (4b)
$\displaystyle\frac{d}{dt}\tilde{x}$
$\displaystyle=\epsilon{\tilde{g}}(\tilde{x},\tilde{y},\epsilon),$

where $A(\tilde{x},\tilde{y},\epsilon)\in\mathbb{R}^{N_{y}\times N_{y}}$,
${\tilde{g}}(\tilde{x},\tilde{y},\epsilon)\in\mathbb{R}^{N_{x}}$, and
$A(\tilde{x},0,0)$ is a stable matrix, i.e., its eigenvalues are on the left
of the imaginary axis.

The above lemma is derived in Chapter 3.2 of [35]. The lemma introduces (4),
which is often called the _Fenichel normal form_ due to the pioneer work of
Neil Fenichel [21]. Note that the Fenichel normal form given in [35] is more
general, including fixed points that are either stable or unstable in the
linearization of the fast variable. In this work, we only consider fast
variables that are long-time stable. Moreover, due to the existence of the
coordinate transformation, we consider the general form of the singularly
perturbed systems (2). However, the advantage of the split form (1) is that
the slow manifold dimension is known.

Our FSNN takes advantage of a modified formulation of Fenichel normal form
which will be described in the following discussion. The following theorem
forms the framework for the FSNN architecture.

Suppose the dynamics of $z(t;\epsilon)\in\mathbb{R}^{N}$, $\epsilon>0$, admit
a fast-slow split and can be described by normally stable fast-slow system.
Then there exists a $\epsilon$-dependent invertible transformation

(7) $\displaystyle\left[\begin{array}[]{c}y\\\ x\end{array}\right]$
$\displaystyle=h(z;\epsilon),$

where $x(t)\in\mathbb{R}^{N_{x}}$, $y(t)\in\mathbb{R}^{N_{y}}$, are slow and
fast variables, respectively, and for $|y|\leq\Delta$, where $\Delta>0$, the
dynamics for $(x,y)$ can be described using

(8a) $\displaystyle\frac{d}{dt}y$
$\displaystyle=T(x)y+B(x,y,\epsilon)(y,y)+\epsilon C(x,y,\epsilon)y,$ (8b)
$\displaystyle\frac{d}{dt}x$ $\displaystyle=\epsilon g(x,y,\epsilon).$

$T(x)\in\mathbb{R}^{N_{y}\times N_{y}}$ is a block upper triangular matrix
consisting of either $\mathbb{R}^{2\times 2}$ or $\mathbb{R}$ blocks on the
diagonal and the eigenvalues of these diagonal blocks are strictly to the left
of the imaginary axis for all $x$.
$B(x,y,\epsilon):\mathbb{R}^{N_{y}}\times\mathbb{R}^{N}_{y}\rightarrow\mathbb{R}^{N}_{y}$
is a bilinear map, $C(x,y,\epsilon)\in\mathbb{R}^{N_{y}\times N_{y}}$ is a
matrix, and $g(x,y,\epsilon)\in\mathbb{R}^{N_{x}}$. $T$, $B$, $C,$ and $g$ are
smooth functions of their arguments.

In the following discussion, we will prove Theorem 8 by refining the Fenichel
normal form in (4) and applying the matrix Schur decomposition. First, we
modify (4) by performing a Taylor expansion of the term
$A(\tilde{x},\tilde{y},\epsilon)\tilde{y}$ about $y=0$ and $\epsilon=0$. With
sufficient smoothness on $A$, we have

(9) $\displaystyle
A(\tilde{x},\tilde{y},\epsilon)\tilde{y}=A(\tilde{x},0,0)\tilde{y}+B(\tilde{x},\tilde{y},\epsilon)(\tilde{y},\tilde{y})+\epsilon
C(\tilde{x},\tilde{y},\epsilon)\tilde{y},$

where
$B(\tilde{x},\tilde{y},\epsilon):\mathbb{R}^{N_{y}}\times\mathbb{R}^{N_{y}}\rightarrow\mathbb{R}^{N_{y}}$
is a bilinear map and
$C(\tilde{x},\tilde{y},\epsilon)\in\mathbb{R}^{N_{y}\times N_{y}}$ are
remainder terms. Next, the Schur decomposition of $A(\tilde{x},0,0)$ will be
used to define a new set of fast coordinates.

The Schur form of a matrix, $A\in\mathbb{R}^{M\times M}$, is a matrix
$T\in\mathbb{R}^{M\times M}$ that satisfies $A=QTQ^{T}$, where
$Q\in\mathbb{R}^{M\times M}$ is orthogonal ($Q^{T}Q=I$) and $T$ is block upper
triangular where the diagonal blocks consist of $\mathbb{R}$ and
$\mathbb{R}^{2\times 2}$ blocks.

The Schur decomposition of a matrix $A\in\mathbb{R}^{M\times M}$, is given by

(10) $\displaystyle A=QTQ^{\rm T},$

where $Q,T\in\mathbb{R}^{M\times M}$, $Q$ is an orthogonal matrix ($Q^{\rm
T}Q=I$), and $T$ is the Schur form of $A$. The eigenvalues of $A$ are given by
the eigenvalues of the diagonal blocks of $T$.

The existence of the Schur decomposition is proven in many references,
including [30, 2], though for any given matrix, the decomposition is not
unique. Let the Schur decomposition of $A(\tilde{x},0,0)$ be given by

(11) $\displaystyle A(\tilde{x},0,0)=Q(\tilde{x})T(\tilde{x})Q(\tilde{x})^{\rm
T},$

where $Q(\tilde{x})\in\mathbb{R}^{N_{y}\times N_{y}}$ is an orthogonal matrix
and $T(\tilde{x})\in\mathbb{R}^{N_{y}\times N_{y}}$ is the Schur form of
$A(\tilde{x},0,0)$. Due to our assumptions on $A(\tilde{x},0,0)$, all of the
eigenvalues of $T(\tilde{x})$ are to the left of the imaginary axis. Consider
the coordinate transformation $\hat{y}=Q(\tilde{x})^{\rm T}\tilde{y}$. We form
the relevant ODE for $\hat{y}$ to use in place of the ODE for $\tilde{y}$.

$\displaystyle\frac{d}{dt}\hat{y}=\;$ $\displaystyle Q(\tilde{x})^{\rm
T}\frac{d}{dt}\tilde{y}+\left(\nabla Q(\tilde{x})^{\rm
T}\right)\left(\frac{d}{dt}\tilde{x},\tilde{y}\right)$ $\displaystyle=\;$
$\displaystyle{Q(\tilde{x})^{T}}Q(\tilde{x})T(\tilde{x})Q(\tilde{x})^{T}\tilde{y}+Q(\tilde{x})^{\rm
T}B(\tilde{x},Q(\tilde{x})\hat{y},\epsilon)(Q(\tilde{x})\hat{y},Q(\tilde{x})\hat{y})$
$\displaystyle+\epsilon Q(\tilde{x})^{\rm
T}C(\tilde{x},\epsilon)Q(\tilde{x})\hat{y}+\epsilon\left(\nabla
Q(\tilde{x})^{\rm
T}\right)\left(g_{\epsilon}(\tilde{x},\hat{y}),\hat{y}\right).$
$\displaystyle=\;$ $\displaystyle T(\tilde{x})\hat{y}+Q(\tilde{x})^{\rm
T}B(\tilde{x},Q(\tilde{x})\hat{y},\epsilon)(Q(\tilde{x})\hat{y},Q(\tilde{x})\hat{y})$
$\displaystyle+\epsilon Q(\tilde{x})^{\rm
T}C(\tilde{x},\epsilon)Q(\tilde{x})\hat{y}+\epsilon\left(\nabla
Q(\tilde{x})^{\rm
T}\right)\left(g_{\epsilon}(\tilde{x},\hat{y}),\hat{y}\right).$

By redefining terms, it can be shown that the ODE is equivalent to

$\displaystyle\frac{d}{dt}\hat{y}=T(\tilde{x})\hat{y}+\tilde{B}(\tilde{x},\hat{y},\epsilon)(\hat{y},\hat{y})+\epsilon\tilde{C}(\tilde{x},\hat{y},\epsilon)\hat{y},$

where $\tilde{B}(\tilde{x},\hat{y},\epsilon)$ is a bilinear form operator and
$\tilde{C}(\tilde{x},\hat{y},\epsilon)$ is a matrix. This establishes an
equivalence to (8).

Theorem 8 takes advantage of many useful properties. In the fast-slow
coordinate system of Theorem 8, the definition of the slow-manifold simplifies
significantly. The graph, $M=\\{(x,y):x\in\mathbb{R}^{N_{x}},\,y=0\\}$ is
locally invariant under (8).

This theorem can be easily verified by substituting $y=0$ into (8a) and
observing that $\frac{d}{dt}y=0$ on the slow manifold. Trajectories on the
slow manifold will remain on the slow manifold for all time. In this case, the
dynamics can be reduced to an ODE for the slow variable,

(12) $\displaystyle\frac{d}{d\tau}x=g(x,0,\epsilon),\qquad(x,y)\in M,$

where $\tau$ is the slow timescale. In the case when the coordinates $(x,y)$
are close to, but not on $M$, the following theorem states that trajectories
emanating from $(x,y)$ will approach $M$ exponentially in forward time. If
$\epsilon>0$ is sufficiently small, then, for $|y|\leq\Delta$, where
$\Delta>0$, there are $\kappa>0$ and $\alpha<0$ such that

(13) $\displaystyle d\left((x(t),y(t)),M\right)\leq\kappa\exp(\alpha t),$

where $d$ represents the Euclidean distance.

This theorem is proved following Theorem 5 in [35].

Theorem 8 provides a framework for the FSNN architecture, which expresses each
of the functions $h,T,B,C,$ and $g$ using neural networks with specially
designed architectures described in the following section. There are a few
considerations that went into using the modified Fenichel normal form. By
expanding $A(\tilde{x},\tilde{y},\epsilon)$ in (9), we expose three separate
terms with known structures. We further use the Schur decomposition to
explicitly expose the eigenvalues of $A(\tilde{x},0,0)$, which are required to
be left of the imaginary axis. Finally, we further simplify the linear term by
absorbing the matrix $Q(\tilde{x})$ into the representation. Though we use the
Schur decomposition to represent $A(\tilde{x},0,0)$ in this paper, other
decompositions for representing matrices with eigenvalues to the left of the
imaginary axis have been studied in literature [19, 2].

##  3 Fast-Slow Neural Networks (FSNNs)

In this section, we present a neural network-based approach to approximating
the set of solutions to normally stable fast-slow systems using the framework
of Theorem 8. We refer to this approach as the fast-slow neural network
(FSNN), which is defined below and summarized in Figure 1.

The fast slow neural network (FSNN) is a transformation of an initial
condition, $z_{0}$, and parameter, $\epsilon>0$, to a time series
$\\{z_{n}\\}_{n=0}^{N}$ corresponding to times $t_{0},\dots,t_{N}$, where
$z_{n}\in\mathbb{R}^{N}$ for $n=0,\dots,N$. The initial condition is
transformed into fast-slow coordinates, $(y_{0},x_{0})$, via (7) where $h$ is
represented using an invertible coupling flow network. A numerical integration
is used to approximate the solution to the fast-slow system in (8) to obtain
the time series $\\{(y_{n},x_{n})\\}_{n=0}^{N_{t}}$ for corresponding times
$t_{0},\dots,t_{N}$, where $T$ is represented using a negative Schur form
network, $B$ is represented using a low-rank bilinear form network, and $C$
and $g$ are represented using feed-forward neural networks. The trajectories
are then transformed back into the original coordinate system using

(15) $\displaystyle
z_{n}=h^{-1}\left(\left[\begin{array}[]{cc}y_{n}&x_{n}\end{array}\right]^{\rm
T};\epsilon\right),\qquad n=1,\dots,N.$

The key components of Definition 3, including the invertible coupling flow
network, Schur form network, and low-rank bilinear form network, will be
defined in the following sections. First we will begin by stating the main
result of the paper.

The FSNN is a universal approximator for the set of solutions of normally
stable fast-slow systems in the neighborhood of the slow manifold.

Due to Theorem 8, normally stable fast-slow systems can be represented using
(7)–(8). Therefore, the FSNN inherits the universal approximation properties
of its components, which we will later show are universal approximators for
their corresponding class of functions. The invertible coupling flow network
which represents $h$ is described in Section 3.1 and was proven to be a
universal approximator to invertible maps in [58]. The universal approximation
theorems of the negative Schur form network representing $T$ and the bilinear
form network representing $B$ are proved in Sections 3.2 and 3.3,
respectively. The universal approximation of nonlinear functions using
feedforward neural networks for $C$ and $g$ are well known, e.g. [15].
Finally, we describe a consistent numerical scheme to obtain trajectories in
time, which has an error bound that vanishes as the time steps $\Delta
t_{n}=t_{n}-t_{n-1}\rightarrow 0$.

Our ML-based approach naturally comes with a set of hyperparameters. The
standard hyperparmeters are those related to the architectures of $h$, $T$,
$C$ and $g$. The only systemic hyperparameters are the dimensions of the slow
and fast subspaces, $N_{x}$ and $N_{y}$, respectively. In the examples
considered in Section 5, the true dimensions of slow manifolds are known and
used. In applications when the size of the slow manifold dimension is not
known, tools like diffusion maps [13] or proper orthogonal decomposition (POD)
[5] may be used to infer the size of the reduced dimension before applying the
proposed FSNN.

###  3.1 The Invertible Coupling Flow Network

Figure 1: Schematic of the fast-slow neural network (FSNN). The original
coordinates are projected into fast-slow coordinates using the invertible
neural network (INN). In these coordinates, numerical integration is used to
advance the solution to later time steps. The inverse function of the INN is
used to convert the fast-slow coordinates into the original coordinate system.

Now we will describe further details of the FSNN architecture. The first
component is the invertible function, $h(z;\epsilon)$, which represents the
transformation into fast-slow coordinates. The network is based on the
invertible coupling flow network developed in [1], where the authors closely
follow the architecture provided by [18]. These networks have been shown to be
universal approximators to invertible functions in [58]. In this work, we
modify the original network by introducing a regularization property which
controls the Lipschitz constant of both the forward and inverse map. Such a
regularization improves the robustness of the network during training. The
constant is a tunable parameter of the network that, similar to the hidden
dimensions of the network, can be made arbitrarily large to increase the
expressiveness of the network. Therefore, this modification is compatible with
the universal approximation result derived in [58]. As such, we introduce the
key ingredient of Lipschitz control.

The bi-Lipschitz affine transformation (bLAT) layer is given by the
transformation, $f:\mathbb{R}^{N}\rightarrow\mathbb{R}^{M}$, defined by

(16) $\displaystyle f(x)=U\Sigma V^{\rm T}x+b,$

where $U\in\mathbb{R}^{M\times r}$, $V\in\mathbb{R}^{r\times N}$,
$\Sigma\in\mathbb{R}^{r\times r}$, $b\in\mathbb{R}^{M}$, $r=\min(M,N)$, $U$
and $V$ are orthogonal matrices satisfying $U^{\rm T}U=V^{\rm T}V=I$, and

(18) $\displaystyle\Sigma={\rm
diag}\left[\begin{array}[]{ccc}\sigma_{1}&\dots&\sigma_{r}\end{array}\right],$

where $\frac{1}{L}\leq\sigma_{i}\leq L$ for $i=1,\dots,r$ and some $L\geq 1$.
When $M=N$, the mapping is invertible.

(19) $\displaystyle f^{-1}(x)=V\Sigma^{-1}U^{\rm T}(x-b).$

Our new contribution is to include the matrix term $U\Sigma V^{T}$, which is
written as a singular value decomposition (SVD), which exists for every matrix
in $\mathbb{R}^{M\times N}$. One advantage of the bLAT layer is the following
theorem.

The forward and inverse bLAT layers are both $L$-bi-Lipschitz.

This follows from the fact that $\frac{1}{L}\leq\|\nabla f(x)\|_{2}\leq L$ and
$\frac{1}{L}\leq\|\nabla f^{-1}(x)\|_{2}\leq L$. Another advantage of using
the SVD is that, by construction, the evaluation of $f(x)$ and $f^{-1}(x)$
have an equivalent computational cost. In our implementation of the $M=N$
case, the orthogonal matrices of the bLAT layers are parameterized using the
matrix exponential of a skew-symmetric matrix input, i.e. $U={\rm expm}(S)$,
where $S=-S^{T}$. When $M\neq N$, the orthogonal matrices are parameterized
using the Householder factorization.

The second component of the invertible coupling flow network is the affine
coupling flow layer, defined below. We introduce a similar regularization idea
to the above.

The affine coupling flow layer is defined by the transformation $y=f(x)$,
$x,y\in\mathbb{R}^{M}$, which is implicitly defined by

(28) $\displaystyle\left[\begin{array}[]{c}y_{1}\\\
y_{2}\end{array}\right]=\left[\begin{array}[]{cc}{\rm diag}(F(x_{2}))&\\\
&{\rm diag}(G(y_{1}))\end{array}\right]\left[\begin{array}[]{c}x_{1}\\\
x_{2}\end{array}\right]+\left[\begin{array}[]{c}B(x_{2})\\\
C(y_{1})\end{array}\right],$

where $x=\left[\begin{array}[]{cc}x_{1}&x_{2}\end{array}\right]^{\rm T}$
$y=\left[\begin{array}[]{cc}y_{1}&y_{2}\end{array}\right]^{\rm T}$
$x_{1},y_{1}\in\mathbb{R}^{d}$, $x_{2},y_{2}\in\mathbb{R}^{M-d}$,
$F,B:\mathbb{R}^{M-d}\rightarrow\mathbb{R}^{d}$,
$G,C:\mathbb{R}^{d}\rightarrow\mathbb{R}^{M-d}$, $1\leq d<M$ and the range of
the output elements of $F$ and $G$ is constrained to
$\left[\frac{1}{L},\,L\right]$ for some $L\geq 1$. The inverse mapping,
$x=f^{-1}(y)$, is implicitly defined by

(37) $\displaystyle\left[\begin{array}[]{c}x_{1}\\\
x_{2}\end{array}\right]=\left[\begin{array}[]{cc}{\rm diag}(F(x_{2}))^{-1}&\\\
&{\rm
diag}(G(y_{1}))^{-1}\end{array}\right]\left(\left[\begin{array}[]{c}y_{1}\\\
y_{2}\end{array}\right]-\left[\begin{array}[]{c}B(x_{2})\\\
C(y_{1})\end{array}\right]\right).$

The functions $F,G,B,C$ are represented using feed-forward neural networks. In
order to constrain outputs to the range $\left[\frac{1}{L},L\right]$, we apply
the map $\theta_{L}:\mathbb{R}\rightarrow\left[\frac{1}{L},L\right]$ defined
by $\theta_{L}(s)=\exp(\log(L)\tanh(s))$. This constraint ensures
invertibility of the affine coupling layer. The Lipschitz constant is a
function of $L$ and $\nabla F,\nabla G,\nabla B,\nabla C$ and can optionally
be constrained using bLAT layers in the feedforward networks.

The invertible coupling flow network combines bLAT layers and affine coupling
flow layers into a network.

Let $f_{1},\dots f_{K}:\mathbb{R}^{M}\rightarrow\mathbb{R}^{M}$ define a set
of bLAT layers and $g_{1},\dots g_{K}:\mathbb{R}^{M}\rightarrow\mathbb{R}^{M}$
define a set of affine coupling flow layers. The invertible coupling flow
network, $h:\mathbb{R}^{M}\rightarrow\mathbb{R}^{M}$ is defined as the
composition

(38) $\displaystyle h=g_{K}\circ f_{K}\circ\dots\circ g_{1}\circ f_{1}$

and the inverse map is defined by

(39) $\displaystyle h^{-1}=f_{1}^{-1}\circ g_{1}^{-1}\circ\dots\circ
f_{K}^{-1}\circ g_{K}^{-1}.$

The free parameter $L$ in each of the layers adjusts the complexity of the
overall network and can be used as an effective regularizer. Despite the
introduction of this regularization parameter, the network can be made
arbitrarily expressive with the tuning of $L$, and therefore the universal
approximation result from [58] still holds true.

The invertible coupling flow network is a universal approximator for
invertible maps.

In Theorem 8, the invertible transformation is $\epsilon$-dependent.
Therefore, we bake this dependence into the network by including $\epsilon$ as
a functional dependence in $F,G,B,C$.

###  3.2 The Negative Schur Form Network

Next, a parameterization suitable for representing $T$ is discussed. We
present the negative Schur form network, which is a differentiable
parameterization for Schur forms with eigenvalues to the left of the imaginary
axis. The Schur form network is composed of a feedforward neural network in
which the outputs are re-formatted into a negative Schur form matrix, which is
defined below.

The negative Schur form matrix is given by the block matrix
$T\in\mathbb{R}^{N_{y}\times N_{y}}$, where

(44) $\displaystyle T=\left[\begin{array}[]{cccc}T_{11}&T_{12}&\dots&T_{1p}\\\
&T_{22}&\dots&T_{2p}\\\ &&\ddots&\vdots\\\ &&&T_{pp}\end{array}\right].$

When $M$ is even, $p=M/2$ and $T_{ij}\in\mathbb{R}^{2\times 2}$ for $i=1,\dots
p$, $j=i,\dots,p$. When $M$ is odd, $p=(M+1)/2$, $T_{ij}\in\mathbb{R}^{2\times
2}$ for $i=1,\dots p-1$, $j=i,\dots,p-1$, $T_{ip}\in\mathbb{R}^{2}$ for
$i=1,\dots,p-1$, and $T_{pp}\in(-\infty,0)$. Each of the $\mathbb{R}^{2\times
2}$ blocks on the diagonal are further parameterized using the map

(49)
$\displaystyle\mathcal{M}(R,r,\theta,\phi)=R\left[\begin{array}[]{cc}\cos(\theta)&\sin(\theta)\\\
-\sin(\theta)&\cos(\theta)\end{array}\right]+r\left[\begin{array}[]{cc}\cos(\phi)&\sin(\phi)\\\
\sin(\phi)&-\cos(\phi)\end{array}\right],$

where the region of parameters is restricted to

(50) $\displaystyle|R|\geq|r|,\quad
R<0,\quad\theta\in\left(-\frac{\pi}{2},\frac{\pi}{2}\right).$

Due to this parameterization, we have the following theorem.

The negative Schur form network is a parameterization for the set of Schur
forms with eigenvalues to the left of the imaginary axis.

This can be proven by showing that (49)–(50) can smoothly represent both
complex conjugate pairs and real pairs of eigenvalues with negative real part.
Equation (49) represents the $\mathbb{R}^{2\times 2}$ matrix as a 4-parameter
family sum of a symmetric and antisymmetric operator [6]. Let
$m=\frac{1}{2}{\rm Tr}(\mathcal{M})=R\cos(\theta)$ and $p={\rm
det}(\mathcal{M})=R^{2}-r^{2}$. Then the eigenvalues of $\mathcal{M}$ are
given by

(51) $\displaystyle\lambda_{\pm}=m\pm\sqrt{m^{2}-p}.$

The real parts of the eigenvalues are given by

(52) $\displaystyle{\rm Re}(\lambda_{\pm})=m\pm\sqrt{\max(m^{2}-p,0)}.$

Figure 2 shows the regions of the $m{\rm-}p$ plane where both eigenvalues are
less than zero (green region), greater than zero (red region), complex (blue
region), and real (outside of the blue region). In order for ${\rm
Re}(\lambda_{\pm})<0$, we must have $m<0$ and $p>0$. This translates to the
constraints in (50).

Figure 2: Plot showing the eigenvalues encountered in various regions of the
$m{\rm-}p$ plane, where $m=\frac{1}{2}{\rm Tr}(\mathcal{M})=R\cos(\theta)$ and
$p={\rm det}(\mathcal{M})=R^{2}-r^{2}$. The green region corresponds to linear
stability, ${\rm Re}(\lambda_{+},\lambda_{-})<0$, and the red region
corresponds to linear instability, ${\rm Re}(\lambda_{+},\lambda_{-})>0$.
Above the parabola $p=m^{2}$ (in the blue region), eigenvalues are complex.
Below the parabola (outside of the blue region), the eigenvalues are real.

In the implementation of the network, each of the $\mathbb{R}^{2\times 2}$
blocks have four tunable parameters,
$\bar{R},\,\bar{r},\,\bar{\theta},\,\bar{\phi}\in\mathbb{R}$, that are mapped
onto the restricted domain using

(53) $\displaystyle R=-|\bar{r}|e^{|\bar{R}|},\quad
r=\bar{r},\quad\theta=\frac{\pi}{2}\tanh(\bar{\theta}),\quad\phi=\bar{\phi}.$

The remaining off-diagonal terms are unrestricted in parameter space. In order
to achieve dependence with respect to $x$, (i.e. $T(x)$), we introduce a
feedforward network which takes $x$ as an input and outputs both the tunable
parameters $\bar{R},\,\bar{r},\,\bar{\theta},\,\bar{\phi}$ for each diagonal
block and the elements of the remaining blocks. The resulting negative Schur
form matrix is returned as the output of $T(x)$.

###  3.3 The Low-Rank Bilinear Form Network

Finally, we introduce an approximation for the bilinear form term, $B$.

The low-rank bilinear form network
$B:\mathbb{R}^{N_{x}}\times\mathbb{R}^{N_{y}}\rightarrow\mathbb{R}^{N_{z}}$ is
defined by

(54) $\displaystyle B(x,y)=\sum_{r=1}^{R}(C^{(r)}x)\odot(D^{(r)}y),$

where $x\in\mathbb{R}^{N_{x}}$, $y\in\mathbb{R}^{N_{y}}$,
$C^{(r)}\in\mathbb{R}^{N_{z}\times N_{x}}$, $D^{(r)}\in\mathbb{R}^{N_{z}\times
N_{y}}$, $\odot$ represents the Hadamard product (element-wise
multiplication), and $R$, $1\leq R\leq\min(N_{x},N_{y})$, represents the rank
of approximation.

Here, $R$ is a tunable parameter. Increasing $R$ gives more expressive power
at the expense of computational cost. This network has the following property:

The low-rank bilinear form network is a universal approximator to the set of
bilinear forms.

Consider the bilinear form,
$\tilde{B}:\mathbb{R}^{N_{x}}\times\mathbb{R}^{N_{y}}\rightarrow\mathbb{R}^{N_{z}}$.
We can represent the operator of this function as $\tilde{B}_{ijk}$ where
$B_{i}\in\mathbb{R}^{N_{x}\times N_{y}}$ for $i=1,\dots,N_{z}$. Let
$p=\min(N_{x},N_{y})$, the we can use the singular value decomposition to
represent the matrix $\tilde{B}_{i}$ as

(55)
$\displaystyle\tilde{B}_{i}=\sum_{r=1}^{p}\sigma_{i}^{(r)}u_{i}^{(r)}(v_{i}^{(r)})^{T},$

where $\sigma_{i}^{(r)}$, $u^{(r)}$, $v^{(r)}$, $r=1,\dots,p$, are the
singular values, left singular vectors, and right singular vectors of
$\tilde{B}_{i}$, respectively, and $\sigma_{1}\geq\dots\geq\sigma_{p}\geq 0$.
The $i^{\rm th}$ term of $\tilde{B}(x,y)$ is given by

(56) $\displaystyle x^{T}\tilde{B}_{i}y$
$\displaystyle=\sum_{r=1}^{p}\sigma_{i}^{(r)}x^{T}u_{i}^{(r)}(v_{i}^{(r)})^{T}y,=\sum_{r=1}^{p}\sigma_{i}^{(r)}\left(\sum_{j=1}^{N_{x}}u_{ij}^{(r)}x_{j}\right)\left(\sum_{k=1}^{N_{y}}v_{ik}^{(r)}y_{k}\right),$

We can find some $C\in\mathbb{R}^{N_{z}\times N_{x}}$ and
$D\in\mathbb{R}^{N_{z}\times N_{y}}$ such that

(57) $\displaystyle x^{T}\tilde{B}_{i}y$
$\displaystyle=\sum_{r=1}^{p}\left(\sum_{j=1}^{N_{x}}C_{ij}^{(r)}x_{j}\right)\left(\sum_{k=1}^{N_{y}}D_{ik}^{(r)}y_{k}\right).$

This is equivalent to

(58) $\displaystyle\tilde{B}(x,y)$
$\displaystyle=\sum_{r=1}^{p}(C^{(r)}x)\odot(D^{(r)}y).$

The network uses a truncated version of this sum to approximate the action of
the bilinear operator. The approximation error of the low-rank bilinear form
network is $\|B_{i}-\tilde{B}_{i}\|_{2}=\sigma_{i}^{(R+1)}$. This error can be
decreased by increasing $R$.

In our formulation, the operator $B$ has coefficients which are functions of
$x,y,$ and $\epsilon$. We inject this functional dependence by making the
elements of $C^{(r)}$ and $D^{(r)}$ outputs of a feed-forward network with
inputs $(x,y,\epsilon)$.

###  3.4 Numerical Integration

The Schur form, $T$ will have eigenvalues that can be of arbitrarily large
magnitude. Due to this, the dynamics of (8) will be stiff. Therefore, an
Implicit-Explicit (IMEX) time integrator can be used to solve (8) efficiently.
The IMEX scheme has been applied to a neural ODE framework in [62]. As shown
below, an additional benefit of the proposed FSNN is to eliminiate the need
for iterative or direct linear solvers, thereby simplifying the training
process.

Consider an IMEX Runge-Kutta method consisting of a diagonally implicit for
treating the stiff linear term and an explicit method for treating the
nonlinear term. The method can be represented by the following Butcher
tableau.

$a_{11}$ |  |  |  |  |  |  |
---|---|---|---|---|---|---|---
$a_{21}$ | $a_{22}$ |  |  | $\alpha_{21}$ |  |  |
$\vdots$ | $\vdots$ | $\ddots$ |  | $\vdots$ | $\ddots$ |  |
$a_{s1}$ | $\dots$ | $\dots$ | $a_{ss}$ | $\alpha_{s1}$ | $\dots$ | $\alpha_{s,{s-1}}$ |
$b_{1}$ | $\dots$ | $\dots$ | $b_{s}$ | $\beta_{1}$ | $\dots$ | $\beta_{s-1}$ | $\beta_{s}$

Let $x_{n}\approx x(t_{n})$, $y_{n}\approx y(t_{n})$, and $\Delta
t=t_{n}-t_{n-1}$. Let $f(x,y,\epsilon)=B(x,y,\epsilon)(y,y)+\epsilon
C(x,y,\epsilon)y$. The intermediate stages are implicitly defined by

(59a) $\displaystyle X^{i}$ $\displaystyle=x_{n-1}+\epsilon\Delta
t\sum_{j=1}^{i-1}\alpha_{ij}{g(X^{j},Y^{j},\epsilon)},$ (59b) $\displaystyle
Y^{i}$ $\displaystyle=y_{n-1}+\Delta
t\sum_{j=1}^{i}a_{ij}{T(X^{j})Y^{j}}+\Delta
t\sum_{j=1}^{i-1}\alpha_{ij}{f(X^{j},Y^{j},\epsilon)}$

for $i=1,\dots,s$. The above equation can be written to expose the linear
system that needs to be solve.

(60) $\displaystyle Y^{i}$ $\displaystyle=(I-\Delta
ta_{ii}T(X^{i}))^{-1}\left(y_{n-1}+\Delta
t\sum_{j=1}^{i-1}\left(a_{ij}{T(X^{j})Y^{j}}+\alpha_{ij}{f(X^{j},Y^{j},\epsilon)}\right)\right)$

Note that the matrix given by $I-\Delta ta_{ii}T(X^{i})$ is a block 2-by-2
upper triangular matrix. Due to the proposed architecture, this inverse can be
handled efficiently using a backward substitution solve. Since $T$
additionally has functional dependence, we treat this in a vectorization-
friendly matrix-free fashion, which avoids storing multiple matrices
$T(X^{i})$ for each data pair in the batch set during training. The solution
at the next time step is obtained through the update

(61a) $\displaystyle x_{n}$ $\displaystyle=x_{n-1}+\epsilon\Delta
t\sum_{i=1}^{s}\beta_{i}{g(X^{i},Y^{i},\epsilon)},$ (61b) $\displaystyle
y_{n}$ $\displaystyle=y_{n-1}+\Delta
t\sum_{i=1}^{s}(b_{i}{T(X^{i})Y^{i}}+\beta_{i}{f(X^{i},Y^{i},\epsilon)}).$

In our implementation, we use the two-stage second-order accurate IMEX strong-
stability-preserving L-stable scheme defined in [49].

When the solution is initially on the slow manifold ($y=0$), the stages
$Y^{i}$ for $i=1,\dots,s$ evaluate exactly to zero and therefore the numerical
scheme can be simplified to an integration of the slow variables given by

(62a) $\displaystyle x_{n}$ $\displaystyle=x_{n-1}+\epsilon\Delta
t\sum_{i=1}^{s}\beta_{i}{g(X^{i},0,\epsilon)},$ (62b) $\displaystyle X^{i}$
$\displaystyle=x_{n-1}+\epsilon\Delta
t\sum_{j=1}^{i-1}\alpha_{ij}{g(X^{j},0,\epsilon)},\qquad i=1,\dots,s.$

In this case, one can take larger time steps (e.g. $\Delta
t\sim\frac{1}{\epsilon}$) and step over the fast dynamics formally. A hybrid
scheme can also be used where the full system in (61) can be integrated until
$\|y\|$ is sufficiently small, after which slow integration in (62) is pursued
using larger time steps. We note that some of the first works to use Runge-
Kutta based residual networks to learn equations from time series was pursued
in [51, 50].

#### Data-driven closure

A closure discovery procedure is implied through the above integration with
large time steps. For a given full state vector $z_{0}$, the transformation
$h$ first maps to the separate state of $(x_{0},y_{0})$. The state vector is
then projected onto the closure of $y=0$; the system for $x$ is thus closed
and evolved using (62). Finally, the evolved state of $x$ is lifted back to
the full state $z$ using $h^{-1}$.

##  4 Numerical Implementation

Figure 3: Schematic of training the fast-slow neural network (FSNN). The
network is trained on dynamics data, $\mathcal{D}$, which contains examples of
trajectories on both the fast and slow timescales and manifold data,
$\mathcal{M}$, which contains examples of solution vectors that lie on the
slow manifold. The overall loss is a combination of the dynamics loss, which
includes $L_{2}$ losses in the original coordinates ($\mathcal{L}_{\rm
system}$), fast coordinates ($\mathcal{L}_{\rm fast}$), and slow coordinates
($\mathcal{L}_{\rm slow}$), and the manifold loss ($\mathcal{L}_{\rm
manifold}$), which is an $L_{2}$ loss of the fast coordinate, which should map
to $y=0$ on the slow manifold.

In this section, we describe the numerical implementation of the fast-slow
neural network and its corresponding training algorithm. The computation of
the loss function and its gradient is summarized in Figure 3. The training set
consists of both trajectory data, which describes the behavior of the
dynamics, and manifold data, which describes the limiting behavior of the
system. The trajectory dataset, $\mathcal{D}$, contains time series,
$\\{(t_{n}^{(k)},z_{n}^{(k)})\\}_{n=1}^{N_{t}}$, and parameters
$\epsilon^{(k)}$, for $k=1,\dots,N_{d}$, where $N_{d}$ is the number of
trajectory data samples. The manifold dataset, $\mathcal{M}$, contains
multiple examples of the solution on the slow manifold, $z^{(k)}$, with
corresponding parameters $\epsilon^{(k)}$, for $k=1,\dots,N_{m}$. Examples of
manifold data can be obtained either experimentally by integrating a model for
a sufficiently long time with a small $\epsilon$ or analytically if the
equations permit an $\epsilon=0$ reduction or asymptotic expansion.

The loss function is given by

(63) $\displaystyle\mathcal{L}=\mathcal{L}_{\rm system}+\mathcal{L}_{\rm
fast}+\mathcal{L}_{\rm slow}+\mathcal{L}_{\rm manifold}.$

The first three terms are evaluated using the trajectory data set,
$\mathcal{D}$, and the last term is evaluated using the manifold data set,
$\mathcal{M}$. For the dynamics-oriented loss terms, the fast-slow neural
network is applied to initial conditions in $\mathcal{D}$ to obtain the model
trajectories,
$\\{(t_{n}^{(k)},\tilde{x}_{n}^{(k)},\tilde{y}_{n}^{(k)},\tilde{z}_{n}^{(k)})\\}$.
Additionally, the invertible coupling flow network is applied to the
trajectory in $\mathcal{D}$ to augment the data
$\\{(t_{n}^{(k)},{x}_{n}^{(k)},{y}_{n}^{(k)},{z}_{n}^{(k)})\\}$. Let $L$ be an
operator representing the approximation to the time integral of squared error,
given by

(64) $\displaystyle
L(t,z,\tilde{z})=\sum_{n=0}^{N_{t}-1}\frac{1}{2}\left((z_{n+1}-\tilde{z}_{n+1})^{2}+(z_{n}-\tilde{z}_{n})^{2}\right)(t_{n+1}-t_{n}).$

Then we choose

(65a) $\displaystyle\mathcal{L}_{\rm system}$
$\displaystyle=\frac{1}{N_{d}}\sum_{k=1}^{N_{d}}L(t,z^{(k)},\tilde{z}^{(k)}),$
(65b) $\displaystyle\mathcal{L}_{\rm fast}$
$\displaystyle=\frac{1}{N_{d}}\sum_{k=1}^{N_{d}}L(t,y^{(k)},\tilde{y}^{(k)}),$
(65c) $\displaystyle\mathcal{L}_{\rm slow}$
$\displaystyle=\frac{1}{\epsilon^{(k)}N_{d}}\sum_{k=1}^{N_{d}}L(t,x^{(k)},\tilde{x}^{(k)}).$

Including the three terms individually encourages the network to learn the
dynamics for the original and fast-slow coordinates to an adequate level.

In practice, we split the trajectory dataset, $\mathcal{D}$, into a fast-scale
subset, with time steps chosen to smoothly capture the fast dynamics
converging to the slow manifold, and a slow-scale subset, with larger time
steps chosen to capture the variation of the slow variables after the dynamics
sufficiently close to the slow manifold. The trajectories in fast-scale subset
are treated using the full integration in (61) while the trajectories from the
slow-scale subset are treated using the slow-manifold integration in (62).

For the manifold loss term, the invertible coupling flow network is applied to
examples in $\mathcal{M}$ to obtain fast-slow coordinates $(x^{(k)},y^{(k)})$.
If the data is on the slow manifold, then $y^{(k)}$ should be zero. We
therefore include the following as a penalty term.

(66) $\displaystyle\mathcal{L}_{\rm
manifold}=\frac{1}{N_{m}}\sum_{k=1}^{N_{m}}(y^{(k)})^{2}$

Our implementation of the fast-slow neural network uses Jax [7] as a backend.
Auto-differentiation is used to compute gradients of the loss function with
respect to the tunable parameters and the Adam optimizer from the Optax [17]
package is used adjust the weights towards an locally optimal solution for the
dataset.

##  5 Examples

In this section, we consider multiple example problems to demonstrate the
applicability of the FSNN to various problems, including multi-scale problems
arising in hydrodynamics, climate dynamics, and plasma physics.

###  5.1 Visualization of the Attracting Slow Manifold Property

A significant advantage of the FSNN is that, for any choice of model
parameters, the network will have an attracting slow manifold that, due to
Theorem 8, will attract solutions for sufficiently small values of $\epsilon$.
The slow manifold is represented by the graph, $y=0$. In Figure 4, we
demonstrate this property on multiple random initializations of FSNNs for
$N_{y}=N_{x}=1$. Initial conditions are drawn randomly from a standard normal
distribution (left plot) and are integrated to a later time $t=4$. The
trajectory points at $t=4$ (blue dots) and the resulting slow manifold (orange
curves) are plotted at this final time, demonstrating convergence of solutions
to the slow manifold.

Figure 4: Demonstration of attracting slow manifold property of the FSNN. For
three different random initializations of the FSNN, initial conditions are
drawn randomly from a standard normal distribution and are evolved in time to
a later time ($t=4$). The dynamics are shown to limit to the curve
representing the slow manifold.

###  5.2 Learning a Simple Attracting Manifold

We demonstrate the ability of the FSNN to learn dynamics for a simple test
problem where the dynamics can be described by the following system of ODEs
for $z_{1},z_{2}\in\mathbb{R}$.

(67a) $\displaystyle\frac{dz_{1}}{dt}$
$\displaystyle=-\epsilon(\sin(z_{1})+z_{2})$ (67b)
$\displaystyle\frac{dz_{2}}{dt}$
$\displaystyle=\lambda(z_{1})(z_{2}-\theta(z_{1}))$

In this example, $z_{1}$ represents the slow variable and $z_{2}$ represents
the fast variable. When $\epsilon=0$, the slow manifold is given by
$\theta(z_{1})$. In this example, we choose $\lambda$ and $\theta$ to be

(68)
$\displaystyle\lambda(x)=-1-\frac{1}{10}\cos(2x),\qquad\theta(x)=2\tanh(x).$

The trajectory dataset, $\mathcal{D}$, consists of $10^{3}$ fast-scale
trajectories, corresponding to initial conditions for $z_{1},z_{2}$ randomly
generated uniformly in $[-1,1]$, $\epsilon$ randomly generated log-uniformly
in $[10^{-5},10^{-2}]$, and integrated forward 5 timesteps using an adaptive
RK45 solver with an error tolerance of $10^{-8}$. The manifold dataset,
$\mathcal{M}$, consists of $10^{3}$ examples of data pairs
$((z_{1},z_{2}),\epsilon=0)$ on the graph $z_{2}=\theta(z_{1})$, where $z_{1}$
is randomly generated uniformly in $[-1,1]$.

A FSNN is trained on this dataset using the Adam optimizer from Optax. The INN
consists of one outer layer, where the inner coupling flow layers consists of
feedforward networks with one layer and a hidden dimension of 10. East of the
feedforward neural networks on the RHS of the neural ODE system consists of 10
layers and the bilinear form has a rank of 2.

Figure 5 shows a comparison of the trained model and the reference solution in
phase space for various choices of $\epsilon$. The model shows great agreement
to the reference solution. In all of the plots, it is evident that the
trajectories eventually get trapped in a slow manifold. In the rightmost plot,
a comparison of the learned eigenvalues with that of the ground truth,
$\lambda(z_{1})$, is shown. The results demonstrate that the FSNN is able to
successfully recover details of the original operator.

Figure 5: Results of learning the dynamics of the simple attracting manifold
example described by (67). Left three plots: phase portraits comparing the
ground truth (g.t.) reference solution (bold curves) and that of the trained
model (dots) for $\epsilon=0,0.001,0.01$. Rightmost plot: comparison of the
eigenvalues learned by the FSNN with the ground truth eigenvalues.

###  5.3 Grad’s Moment System

In this example, we show that the FSNN can be applied towards learning
hydrodynamic slow manifolds. We consider a linear Grad moment system [26, 27],
defined by

(69a) $\displaystyle\partial_{t}\rho$
$\displaystyle=-\epsilon\,\nabla\cdot\mathbf{u},$ (69b)
$\displaystyle\partial_{t}\mathbf{u}$
$\displaystyle=-\epsilon\left(\nabla\rho+\nabla T+\nabla\cdot\sigma\right),$
(69c) $\displaystyle\partial_{t}T$
$\displaystyle=-\frac{2\epsilon}{3}\left(\nabla\cdot\mathbf{u}+\nabla\cdot\mathbf{q}\right),$
(69d) $\displaystyle\partial_{t}\mathbf{\sigma}$
$\displaystyle=-\mathbf{\sigma}-\epsilon\left(2\overline{\nabla\mathbf{u}}+\frac{4}{5}\overline{\nabla\mathbf{q}}\right),$
(69e) $\displaystyle\partial_{t}\mathbf{q}$
$\displaystyle=-\frac{2}{3}\mathbf{q}-\epsilon\left(\frac{5}{2}\nabla
T+\mathbf{\nabla}\cdot\mathbf{\sigma}\right),$

where $\rho(\mathbf{x},t)$ is the density, $\mathbf{u}(\mathbf{x},t)$ is
velocity, $T(\mathbf{x},t)$ is temperature, $\sigma(\mathbf{x},t)$ is the
stress tensor, $\mathbf{q}(\mathbf{x},t)$ is the heat flux, and
$\mathbf{x}=(x,y,z)^{T}\in\mathbb{R}^{3}$ is the spatial coordinate. The
overline represents the symmetric, traceless part of a tensor, e.g.

(70) $\displaystyle\overline{\mathbf{a}}$
$\displaystyle=\frac{1}{2}(\mathbf{a}+\mathbf{a}^{T})-\frac{1}{3}\mathbf{I}\,\text{Tr}(\mathbf{a}).$

We consider solutions that are one-dimensional ($\partial_{y}=\partial_{z}=0$)
and $2\pi$-periodic in space, e.g.

(71)
$\displaystyle\rho(\mathbf{x},t)=\sum_{k=-\infty}^{\infty}\rho_{k}(t)e^{ikx}.$

The Fourier coefficients are governed by the following system of ODEs.

(87) $\displaystyle\frac{d}{dt}\left[\begin{array}[]{c}\rho_{k}\\\ u_{k}\\\
T_{k}\\\ \sigma_{k}\\\
q_{k}\end{array}\right]=\left[\begin{array}[]{ccccc}0&-ik\epsilon&0&0&0\\\
-ik\epsilon&0&-ik\epsilon&-ik\epsilon&0\\\
0&-\frac{2}{3}ik\epsilon&0&0&-\frac{2}{3}ik\epsilon\\\
0&-\frac{4}{3}ik\epsilon&0&-1&-\frac{8}{15}ik\epsilon\\\
0&0&-\frac{5}{2}ik\epsilon&-ik\epsilon&-\frac{2}{3}\end{array}\right]\left[\begin{array}[]{c}\rho_{k}\\\
u_{k}\\\ T_{k}\\\ \sigma_{k}\\\ q_{k}\end{array}\right],\quad k=0,\pm 1,\dots$

For this example, we truncate the Fourier series to consider only the $k=0,\pm
1,\pm 2$ models. Initial conditions are randomly generated for each Fourier
coefficient uniformly in the complex region $[-1,1]\times[-i,i]$. The solution
is computed for 10 time steps both using $\Delta t_{\rm fast}=\frac{1}{4}$ and
$\Delta t_{\rm slow}=\frac{1}{4\epsilon}$. The slow-scale trajectories are
offset $\Delta t_{\rm slow}$ from $t=0$ to ensure the solution is sufficiently
close to the slow manifold. A FSNN is trained on this dataset using the Adam
optimizer from Optax. The INN consists of one outer layer, where the inner
coupling flow layers consists of feedforward networks with one layer and a
hidden dimension of 40. The feedforward neural networks on the RHS of the
neural ODE system each consist of 50 layers and the bilinear form has a rank
of 1. Considering each coefficient and its real and imaginary parts, the
dimension of the learned model is $\mathbb{R}^{30}$, with $N_{x}=18$ and
$N_{y}=12$.

Table 1: Comparison of ground truth eigenvalues of the Jacobian of the fast
equations with that of the trained model for the Grad Moment System problem.
Grad Eigenvalues
---
g.t. | model | error | g.t. | model | error
-1 | -0.9971 | 2.9e-03 | -$\frac{2}{3}$ | -0.6659 | 8.2e-04
-1 | -0.9969 | 3.1e-03 | -$\frac{2}{3}$ | -0.6658 | 8.3e-04
-1 | -0.9969 | 3.1e-03 | -$\frac{2}{3}$ | -0.6658 | 9.1e-04
-1 | -0.9968 | 3.2e-03 | -$\frac{2}{3}$ | -0.6657 | 9.3e-04
-1 | -0.9968 | 3.2e-03 | -$\frac{2}{3}$ | -0.6657 | 9.4e-04
-1 | -0.9968 | 3.2e-03 | -$\frac{2}{3}$ | -0.6657 | 9.4e-04

Figure 6 shows comparisons of the trained model and reference solutions on the
fast time scale and slow time scale. The results demonstrate that the model is
accurate on both time scales. Table 1 shows a comparison of the eigenvalues of
the trained model with the analytical eigenvalues. Great agreement is
observed, which implies that the model was successful in mapping the slow
manifold approximately to the true slow manifold.

It is critical to stress that the slow dynamics comparison in Figure 6 is done
through integration on the slow manifold using (62). Specifically, the initial
condition state vector $z_{0}$ is first transformed into the fast-slow
coordinates $(x_{0},y_{0})$ and projected onto the learned slow manifold
$(x_{0},0)$ before integrating the slow variable $x$ on the slow time scale
only and then transforming back into original coordinates using $h^{-1}$. This
approach is significantly more efficient than integrating the full system. The
excellent accuracy of the predicted slow variable solution indicates a data-
driven  _closure_ is learned through this procedure. The closure of (69)
formally exists when $\epsilon$ approaches 0. It is readily seen that the
reduced system is a linearized Euler system.

Figure 6: Comparison of fast scale (top) and slow scale (bottom) dynamics
between the ground truth and trained model for the Grad moment system example.

###  5.4 Multi-Scale Lorenz96 Equations

In this section we consider the two-scale Lorenz96 equations, which was a
model proposed in [47] as a test bed model for multi-scale behavior in
atmospheric dynamics. The original system is formulated as

(88a) $\displaystyle\frac{d}{d\tau}X_{k}$
$\displaystyle=-X_{k-1}(X_{k-2}-X_{k+1})-X_{k}+F-(hc/b)\sum_{j=1}^{J}Y_{j,k},$
(88b) $\displaystyle\frac{d}{d\tau}Y_{j,k}$
$\displaystyle=-cbY_{j+1,k}(Y_{j+2,k}-Y_{j-1,k})-cY_{j,k}+(hc/b)X_{k},$

for $j=1,\dots,J$ and $k=1,\dots,K$, where the solutions satisfy periodicity
($X_{k+K}=X_{k}$, $Y_{j,k+K}=Y_{j,k}$, $Y_{j+J,k}=Y_{j,k}$). Let
$c=b=1/\epsilon$, $h=1$, $Y=\epsilon y$, $\tau=\epsilon t$, and $F=0$. The
system becomes

(89a) $\displaystyle\frac{d}{dt}x_{k}$ $\displaystyle=-\epsilon
x_{k-1}(x_{k-2}-x_{k+1})-\epsilon x_{k}-\epsilon^{2}\sum_{j=1}^{J}y_{j,k},$
(89b) $\displaystyle\frac{d}{dt}y_{j,k}$
$\displaystyle=-y_{j+1,k}(y_{j+2,k}-y_{j-1,k})-y_{j,k}+x_{k}.$

Since the FSNN is limited to representing normally stable fast-slow systems,
we explore the region of applicability where (89) is a normally fast-slow
system. The normally stable region of Lorenz96 is identified by the following
theorem. The system of equations defined by $\eqref{eq:lorenz96}$ represent a
normally stable fast-slow system when $-\frac{1}{2}<x_{k}<\frac{8}{9}$, for
$k=1,\dots,K$.

Using Definition 2, we need to show that the Jacobian of the right-hand-side
of (89b) has eigenvalues strictly to the left of the imaginary axis. Consider
the following nonlinear equation $f(y;\xi)\in\mathbb{R}^{J}$ for
$y\in\mathbb{R}^{J}$ and $\chi\in\mathbb{R}$.

(90) $\displaystyle f(y;\chi)=-y_{j+1}(y_{j+2}-y_{j-1})-y_{j}+\chi.$

The Jacobian of this equation, $D_{y}f$ is given by

(91) $\displaystyle(D_{y}f)_{jk}=\frac{\partial f_{j}}{\partial
y_{k}}=-\delta_{jk}-\delta_{{j+1},k}(y_{j+2}-y_{j-1})-y_{j+1}(\delta_{j+2,k}-\delta_{j-1,k}),$

where $\delta_{jk}$ represents the Kronecker delta. We are interested in the
eigenvalues of $D_{y}$ when $f(y;\chi)=0$. Note that one solution to
$f(y;\chi)=0$ is given by $y_{j}=\chi$ for $j=1,\dots,J$. At this fixed point,
the Jacobian becomes

(92)
$\displaystyle(D_{y}f)_{jk}|_{y=\xi}=-\delta_{jk}-\chi(\delta_{j+2,k}-\delta_{j-1,k})$

This represents a circulant matrix, of which the eigenvalues are given by

(93) $\displaystyle\lambda_{j}=-1-\chi(e^{ij4\pi/J}-e^{-ij2\pi/J}),\qquad
j=0,1,\dots,J-1.$

Note that

(94) $\displaystyle{\rm
Re}(\lambda_{j})=-1-\chi(\cos(4\xi_{j})-\cos(2\xi_{j})),$

where $\xi_{j}=4\pi j/J$. We are interested in the region for $\chi$ where
${\rm Re}(\lambda_{j})<0$ for all $j$. As $J\rightarrow\infty$, we have the
bound $-\frac{9}{8}\leq\cos(4\xi_{j})-\cos(2\xi_{j})\leq 2$. This implies that
the region is given by $-\frac{1}{2}<x<\frac{8}{9}$. For a finite $J$, the
bounds of $\cos(4\xi_{j})-\cos(2\xi_{j})$ are determined by the resolution of
$\xi_{j}$, which is a discrete function of $j$. This affords us to take a
slightly wider region of $x$ where ${\rm Re}(\lambda_{j})<0$ for all $j$. For
example, when $J=4$, $-\frac{1}{2}<x<1$.

Homotopy continuation can be used on the polynomial system of equations
defined by $f$ to reveal a high-order polynomial containing all possible
roots. Maple was used to evaluate all the roots for multiple values of $\chi$
and for $J=4,5,6$. Within the region $\chi\in(-\frac{1}{2},\frac{8}{9})$, the
only real-valued solution is given by $y_{j}=\chi$ for $j=1,\dots,J$. We
conjecture that there is only a single real-valued solution to (90) when
$\chi\in(-\frac{1}{2},\frac{8}{9})$ for $J>6$.

We have shown that the multi-scale Lorenz96 equations are normally-stable for
a restricted range of $x_{k}$. However, for any $x_{k}$, $y_{j,l}$,
$\epsilon>0$, the system is energy stable, which implies that, for any initial
condition, the energy will decay in time and the dynamics will eventually be
able to be described by a normally fast-slow system for some time $t>0$.

The equations defined by (89) are energy-stable for $\epsilon>0$.

Consider again the original system. We multiply the equations by $x_{k}$ and
$y_{j,k}$, respectively. Let $(\cdot)^{\prime}=\frac{d}{dt}(\cdot)$, then

(95a) $\displaystyle x_{k}x_{k}^{\prime}$ $\displaystyle=-\epsilon
x_{k}x_{k-1}(x_{k-2}-x_{k+1})-\epsilon
x_{k}^{2}-\epsilon^{2}\sum_{j=1}^{J}x_{k}y_{j,k}$ (95b) $\displaystyle
y_{j,k}y_{j,k}^{\prime}$
$\displaystyle=-y_{j,k}y_{j+1,k}(y_{j+2,k}-y_{j-1,k})-y_{j,k}^{2}+x_{k}y_{j,k}.$

We collect the terms and take the sum over the discrete variables.

$\displaystyle\frac{1}{2}\sum_{k=1}^{K}(x_{k}^{2})^{\prime}$
$\displaystyle=-\epsilon\sum_{k=1}^{K}x_{k}x_{k-1}x_{k-2}+\epsilon\sum_{k=1}^{K}x_{k+1}x_{k}x_{k-1}-\epsilon\sum_{k=1}^{K}x_{k}^{2}-\epsilon^{2}\sum_{k=1}^{K}\sum_{j=1}^{J}x_{k}y_{j,k}$
$\displaystyle\frac{1}{2}\sum_{j=1}^{J}(y_{j,k}^{2})^{\prime}$
$\displaystyle=-\sum_{j=1}^{J}y_{j,k}y_{j+1,k}y_{j+2,k}+\sum_{j=1}^{J}y_{j-1,k}y_{j,k}y_{j+1,k}-\sum_{j=1}^{J}y_{j,k}^{2}+\sum_{j=1}^{J}x_{k}y_{j,k}$

Due to periodicity, the cubic terms evaluate to zero. The remaining terms give
us the following energy estimate.

(96)
$\displaystyle\frac{1}{2}\sum_{k=1}^{K}(x_{k}^{2})^{\prime}+\frac{1}{2}\epsilon^{2}\sum_{k=1}^{K}\sum_{j=1}^{J}(y_{j,k}^{2})^{\prime}=-\epsilon\sum_{k=1}^{K}x_{k}^{2}-\epsilon^{2}\sum_{k=1}^{K}\sum_{j=1}^{J}y_{j,k}^{2}<0$

We train a FSNN to learn the multi-scale Lorenz96 equations from data. The
trajectory dataset, $\mathcal{D}$, consists of $10^{4}$ fast-scale
trajectories, corresponding to initial conditions for $x_{k}$ randomly
generated uniformly in $[-\frac{1}{2},1]$ and $y_{j,k}$ randomly generated
from a normal distribution, $\epsilon$ randomly generated log-uniformly in
$[10^{-5},10^{-2}]$, and integrated forward 10 timesteps using an adaptive
RK45 solver with error tolerance $10^{-8}$. The manifold dataset,
$\mathcal{M}$ consists of $10^{4}$ examples of data pairs
$((x_{k},y_{j,k}),\epsilon=0)$ on the graph $y_{j,k}=x_{k}$, where $x_{k}$ is
randomly generated uniformly in $[-\frac{1}{2},1]$.

A FSNN is trained on this dataset using the Adam optimizer from Optax. The INN
consists of one outer layer, where the inner coupling flow layers consists of
feedforward networks with one layer and a hidden dimension of 40. The
feedforward neural networks on the RHS of the neural ODE system each consist
of 40 layers and the bilinear form has a rank of 4.

Figure 7: The top five plots are trajectory plots comparing the reference
solution computed using RK45 (bold curves) and that of the trained model
(dotted markers) for an initial condition with $\epsilon=1.1\cdot 10^{-3}$ for
the multi-scale Lorenz96 example described by the dynamics in (89). The bottom
five plots show the corresponding errors as a function of time. The plots in
the left column show the solution and error on the slow time scale
($t\in[0,3000]$) while the plots in the right column zoom in on the fast time
scale ($t\in[0,3]$).

Figure 7 shows a comparison of the trained model and the reference solution
for an example initial condition with $\epsilon=1.1\cdot 10^{-3}$. The model
shows quantitative agreement to the reference solution. Figure 8 shows a
comparison of the eigenvalues of the trained model with the analytical
eigenvalues. Agreement is observed, implying that the model was successful in
mapping the slow manifold approximately to the true slow manifold. In this
example, the time interval of all the training data is much shorter than
$t=3$. Nevertheless, the prediction can be done up to $T=3000$ or longer with
excellent accuracy.

Figure 8: Comparison of the true eigenvalues (bold lines) with those learned
from the trained model (dots) for the multi-scale Lorenz96 example described
by the dynamics in (89)

###  5.5 Abraham-Lorentz Equations

In this section, we consider the Abraham-Lorentz equations [10], which model
the radiation reaction in the case of a single electron with charge $e$ and
mass $m$ moving through a static external magnetic field
$\mathbf{B}(\mathbf{x})\in\mathbb{R}^{3}$. The equations are given by

(97a) $\displaystyle\frac{2}{3}\frac{e^{2}}{c^{3}}\frac{d\mathbf{a}}{d\tau}$
$\displaystyle=m\mathbf{a}-\frac{e}{c}\mathbf{v}\times\mathbf{B}(\mathbf{x}),$
(97b) $\displaystyle\frac{d\mathbf{v}}{d\tau}$ $\displaystyle=\mathbf{a},$
(97c) $\displaystyle\frac{d\mathbf{x}}{d\tau}$ $\displaystyle=\mathbf{v},$

where $\mathbf{x}(\tau),\mathbf{v}(\tau),\mathbf{a}(\tau)\in\mathbb{R}^{3}$
are the position, velocity, and acceleration, respectively, and $c$ is the
speed of light. We rescale these equations following the approach in [10].
Time is scaled by the observer time scale $T$ as $\tau=T\overline{\tau}$,
space by the observer length scale $\mathbf{x}=L\overline{\mathbf{x}}$,
velocity as $\mathbf{v}=(L/T)\overline{\mathbf{v}}$, the magnetic field as
$\mathbf{B}(\mathbf{x})=B_{0}\overline{\mathbf{B}}(\mathbf{x})$, and
acceleration as $a=(L/T)(|e|B_{0})(mc)^{-1}\overline{\mathbf{a}}.$ We
reformulate the rescaled Abraham-Lorentz equations in terms of the
nondimensional variables and drop the overline for notational convenience.

(98a)
$\displaystyle\left(\frac{r_{0}}{cT}\right)\frac{2}{3}\frac{d\mathbf{a}}{d\tau}$
$\displaystyle=\mathbf{a}-\zeta\mathbf{v}\times\mathbf{B}(\mathbf{x}),$ (98b)
$\displaystyle\frac{d\mathbf{v}}{d\tau}$
$\displaystyle=(\omega_{c}T)\mathbf{a},$ (98c)
$\displaystyle\frac{d\mathbf{x}}{d\tau}$ $\displaystyle=\mathbf{v},$

where we have introduced the cyclotron frequency $\omega_{c}=|e|B_{0}/(mc)$,
the classical electron radius $r_{0}=e^{2}/(mc)^{2}$, and the sign of the
charge $\zeta\in\\{-1,+1\\}$, where $\zeta=-1$ corresponds to electrons and
$\zeta=+1$ to positrons. The dimensionless parameters

(99)
$\displaystyle\gamma_{R}=\frac{r_{0}}{cT},\qquad\gamma_{B}=\frac{1}{\omega_{c}T},$

represent the ratio of the electron size to the distance light travels during
a time interval $T$, and ratio of the cyclotron period to $T$, respectively.
The system can be transformed into a two-scale fast-slow system with the
choice $\gamma_{B}=1$ and $\gamma_{R}~{}=~{}\epsilon~{}\ll~{}1$. We scale time
through the transformation $\tau=\frac{2}{3}\epsilon t$ to reveal a normally-
unstable fast-slow system given by

(100a) $\displaystyle\frac{d\mathbf{a}}{dt}$
$\displaystyle=\mathbf{a}-\zeta\mathbf{v}\times\mathbf{B}(\mathbf{x}),$ (100b)
$\displaystyle\frac{d\mathbf{v}}{dt}$
$\displaystyle=\epsilon\frac{2}{3}\mathbf{a},$ (100c)
$\displaystyle\frac{d\mathbf{x}}{dt}$
$\displaystyle=\epsilon\frac{2}{3}\mathbf{v},$

We will refer to (100) as the _forward-time Abraham-Lorentz system_. A
normally-stable fast-slow system can be obtained by reversing time using
$t_{-}=-t$,

(101a) $\displaystyle\frac{d\mathbf{a}}{dt_{-}}$
$\displaystyle=-\mathbf{a}+\zeta\mathbf{v}\times\mathbf{B}(\mathbf{x}),$
(101b) $\displaystyle\frac{d\mathbf{v}}{dt_{-}}$
$\displaystyle=-\epsilon\frac{2}{3}\mathbf{a},$ (101c)
$\displaystyle\frac{d\mathbf{x}}{dt_{-}}$
$\displaystyle=-\epsilon\frac{2}{3}\mathbf{v}.$

We refer to (101) as the _reverse-time Abraham-Lorentz system_. In both
systems, $\mathbf{a}$ is the fast variable and $\mathbf{v},\mathbf{x}$ are the
slow variables. In reverse time, the acceleration, $\mathbf{a}$ will be
attracted to the slow manifold, given by

(102)
$\displaystyle\mathbf{a}^{*}(\mathbf{x},\mathbf{v})=\zeta\mathbf{v}\times\mathbf{B}(\mathbf{x})+\epsilon\frac{3}{2}\mathbf{v}\times\mathbf{B}(\mathbf{x})\times\mathbf{B}(\mathbf{x})+\mathcal{O}(\epsilon^{2}).$

Though the reverse-time system is stable on the fast time scale, the system
exhibits an instability tangential to the slow manifold for which solutions
blow up on the slow time scale. This is a more challenging test than the
previous ones.

It is of interest to obtain forward-time trajectories of electrons along the
slow manifold. A standard numerical integration of (100) will result in
unstable trajectories on the fast time scale since any numerical error will
push trajectories off of the slow-manifold, resulting in exponential growth of
the acceleration. We instead apply FSNNs to learn the Fenichel normal form of
the system in _reverse time_. We then obtain forward-time trajectories along
the learned slow manifold by performing an integration of the slow variables
along $y=0$ using (62). This test is therefore a more practical application of
a data-driven closure.

We train a FSNN to learn the Abraham-Lorentz equations from data. The
trajectory dataset, $\mathcal{D}$, consists of $10^{4}$ trajectories,
corresponding to initial conditions randomly generated in the region,
$\mathcal{R}$, specified by

$\displaystyle\mathcal{R}=\left\\{\mathbf{x},\mathbf{v},\mathbf{a}:\quad 2\leq
r\leq 4,\quad|z|\leq 1,\quad\|\mathbf{v}\|\leq
1,\quad\|\mathbf{a}-\mathbf{a}^{*}(\mathbf{x},\mathbf{v})\|\leq 1\right\\},$

where $r=\sqrt{x_{1}^{2}+x_{2}^{2}}$ and $z=x_{3}$. $\epsilon$ is randomly
generated log-uniformly in the range $[10^{-5},10^{-2}]$. The trajectories are
integrated in reverse time using equations (101) using an adaptive RK45 solver
set to an error tolerance of $10^{-8}$. The first 40 time steps batched in
sequences of length 10 are used as fast-scale training data. These
trajectories are also interpolated over a slower time scale with $\Delta
t=0.025/\epsilon$ and $10^{4}$ sequences of length 10 are chosen at random for
slow-scale training data. The manifold dataset, $\mathcal{M}$ consists of
$10^{4}$ examples of data pairs
$((\mathbf{x},\mathbf{v},\mathbf{a}),\epsilon=0)$ on the graph
$\mathbf{a}=\mathbf{v}\times\mathbf{B}(\mathbf{x})$, where
$\mathbf{x},\mathbf{v}$ are randomly generated in the region $\mathcal{R}$.

The magnetic field is chosen to be
$\mathbf{B}=B_{r}\mathbf{e}_{r}+B_{\phi}\mathbf{e}_{\phi}+B_{z}\mathbf{e}_{z}$,
where $\mathbf{e}_{r},\mathbf{e}_{\phi},\mathbf{e}_{z}$ represent basis
vectors for cylindrical coordinates, and

(103) $\displaystyle B_{r}$ $\displaystyle=-\frac{z}{q(r,z)r}B_{0},\qquad
B_{\phi}=\frac{R_{0}}{r}B_{0},\qquad B_{z}=\frac{(r-R_{0})}{q(r,z)r}B_{0},$

where $q(r,z)=q_{0}+q_{2}((r-R_{0})(r-R_{0})+z^{2})/a^{2}$ and $q_{0}=1.2$,
$q_{2}=2.8$, $a=1.5$, $R_{0}=3.0$, $Z_{0}=3.0$, and $B_{0}=1.0$. This field is
an idealized equilibrium of a tokamak, where the field is physical inside a
torus with major radius of $R_{0}$ and minor radius of $a$. Here $q$ is called
a safety factor [45] and the poloidal magnetic flux function is
$\psi=\frac{B_{0}a^{2}}{2q_{2}}\ln(q).$ Under this definition, the given field
can be rewritten as

(104)
$\displaystyle\mathbf{B}=\frac{1}{R}\nabla\psi\times\mathbf{e}_{\phi}+\frac{F_{0}}{R}\mathbf{e}_{\phi}.$

where the poloidal field is a constant of $F_{0}=R_{0}B_{0}$. This test thus
models a charged particle moving through a tokamak with radiation reaction.

A FSNN is trained on this dataset using the Adam optimizer from Optax. The INN
consists of one outer layer, where the inner coupling flow layers consists of
feedforward networks with one layer and a hidden dimension of 50. The function
$g$ on the right-hand-side of the neural ODE uses a single-layer feedforward
neural network with hidden dimension 100 and is added to a bilinear form
network with rank 1. The remaining feedforward neural networks on the right-
hand-side of the neural ODE system each consist of 50 layers and the bilinear
form has a rank of 4.

Figure 9 shows comparisons of the solutions in $r{\rm-}z$ space between the
trained network and trajectories computed using RK45 on the reverse-time
system. Using the reverse-time trajectories, endpoints were supplied as
initial conditions to the trained network to predict forward in time. For both
the reverse-time predictions (red arrows) and forward-time predictions (green
arrows), the the trained model agrees well with the solution computed
numerically. The corresponding $L_{2}$ errors are shown for position,
velocity, and acceleration in Figure 10, demonstrating that errors are low
when integrating on the slow time scale.

Table 2: Comparison of true eigenvalues to the eigenvalues of the trained
model for the Abraham-Lorentz problem.  Reverse-Time Abraham-Lorentz System
Eigenvalues
---
g.t. |  model |  error
-1 |  -0.99962 |  3.8e-04
-1 |  -0.99958 |  4.2e-04
-1 |  -0.99958 |  4.2e-04

Table 2 shows the comparison of the true eigenvalues of the linear part of the
fast dynamics of the reverse-time system. The eigenvalues were recovered with
a great deal of accuracy, demonstrating that the fast-scale dynamics are
learned accurately.

Figure 11 shows comparisons of orbits between the trained model and orbits
obtained when integrating the system by substituting the first two terms of
the asymptotic expansion in (102). Note again that when integrating along the
slow manifold, only the slow variable is evolved using the transformation $h$
and the closure $y=0$. The trained model is able to qualitatively capture the
gyro-motion behavior while also accurately capturing the guiding center
dynamics over the slow time scale. Despite training using a limited number of
trajectories over a relatively short time frame (2 or less gyration periods)
due to the slow time scale instability of the reverse-time system, the trained
network is highly capable of predicting long-time dynamics in forward-time for
many more gyration cycles. We remind the reader that standard numerical
methods cannot integrate the forward-time Abraham-Lorentz system along the
slow manifold without experiencing an exponential growth instability. Direct
integration of the asymptotic expansion removes the instability, but
introduces an error dependent on $\epsilon$ due to the asymptotic
approximation. This result highlights the unique capability of the proposed
network.

Figure 9: Displayed are the trajectories for three values of $\epsilon$
evolved backwards in time along the slow manifold (top three plots, red
arrows) and three values of $\epsilon$ evolved forwards in time along the slow
manifold (bottom three plots, green arrows). The solid lines represent the
ground truth trajectories and the dotted lines represent the model.

Figure 10: Error trajectories for the reverse time (top three plots) and
forward time simulations are displayed. For each simulation, the $L_{2}$ norm
of position, velocity and acceleration are plotted.

Figure 11: Comparison of orbits of the forward-time Abraham-Lorentz system
(100) along the slow manifold for the first two terms of the asymptotic
expansion in (102) (left) and the trained model (right) for two different
initial conditions and choices of $\epsilon$.

##  6 Conclusions

In this paper, we introduced the fast-slow neural network (FSNN), which is a
data-driven approach suitable for learning singularly perturbed dynamical
systems where the fast-scale dynamics are dissipative. Our method enforces the
existence of a trainable, attracting invariant slow manifold as a hard
constraint. Invertible neural networks and neural ODEs are key components of
the FSNN. We introduce bi-Lipschitz affine transformation (bLAT) layers for
enforcing regularity of the invertible neural network. We also develop a
negative Schur form network which parameterizes matrices with eigenvalues to
the left of the imaginary axis. This in addition to the use of an additive
L-stable diagonally-implicit Runge Kutta scheme for integration are key pieces
in ensuring that the fast-scale dynamics are stable. The Schur form network
also enables the use of back-substitution to efficiently handle the implicit
solve. Analytical representation of the slow manifold enables efficient
integration on the slow time scale. We demonstrate the FSNN on many examples
that exhibit two timescales, including the Grad moment system from
hydrodynamics, two-scale Lorentz96 equations for modeling atmospheric
dynamics, and Abraham-Lorentz dynamics modeling radiation reaction of
electrons in a magnetic field.

The proposed network can be used to discover a full-order model with long-term
accuracy near a slow manifold when only short duration trajectories are
available. The proposed network is also applicable to discover a reduced-order
model as a unique closure discovery procedure.

## Acknowledgment

This research used resources provided by the Los Alamos National Laboratory
Institutional Computing Program, which is supported by the U.S. Department of
Energy National Nuclear Security Administration under Contract No.
89233218CNA000001, and 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.

## References

  * [1] L. Ardizzone, J. Kruse, S. Wirkert, D. Rahner, E. W. Pellegrini, R. S. Klessen, L. Maier-Hein, C. Rother, and U. Köthe, Analyzing inverse problems with invertible neural networks, arXiv preprint arXiv:1808.04730, (2018).
  * [2] Z.-Z. Bai and J.-Y. Pan, Matrix Analysis and Computations, 01 2021.
  * [3] Y. Bar-Sinai, S. Hoyer, J. Hickey, and M. P. Brenner, Learning data-driven discretizations for partial differential equations, Proceedings of the National Academy of Sciences, 116 (2019), pp. 15344–15349.
  * [4] P. Benner, S. Gugercin, and K. Willcox, A survey of projection-based model reduction methods for parametric dynamical systems, SIAM Review, 57 (2015), pp. 483–531.
  * [5] G. Berkooz, P. Holmes, and J. L. Lumley, The proper orthogonal decomposition in the analysis of turbulent flows, Annual review of fluid mechanics, 25 (1993), pp. 539–575.
  * [6] J. F. Blinn, Consider the lowly 2 x 2 matrix, IEEE Computer Graphics and Applications, 16 (1996), pp. 82–88.
  * [7] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, JAX: composable transformations of Python+NumPy programs, 2018.
  * [8] S. Braginskii, Transport processes in a plasma, Reviews of plasma physics, 1 (1965), p. 205.
  * [9] J. Burby, Q. Tang, and R. Maulik, Fast neural Poincaré maps for toroidal magnetic fields, Plasma Physics and Controlled Fusion, 63 (2020).
  * [10] J. W. Burby and T. Klotz, Slow manifold reduction for plasma science, Commun. Nonlinear Sci. Numer. Simul., 89 (2020), p. 105289.
  * [11] R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud, Neural ordinary differential equations, Advances in neural information processing systems, 31 (2018).
  * [12] A. J. Chorin and F. Lu, Discrete approach to stochastic parametrization and dimension reduction in nonlinear dynamics, Proceedings of the National Academy of Sciences, 112 (2015), pp. 9804–9809.
  * [13] R. R. Coifman and S. Lafon, Diffusion maps, Applied and computational harmonic analysis, 21 (2006), pp. 5–30.
  * [14] T. Cui, T. Bertalan, G. J. Pappas, M. Morari, I. G. Kevrekidis, and M. Fazlyab, Certified invertibility in neural networks via mixed-integer programming, 2023, https://arxiv.org/abs/2301.11783.
  * [15] G. V. Cybenko, Approximation by superpositions of a sigmoidal function, Mathematics of Control, Signals and Systems, 2 (1989), pp. 303–314.
  * [16] C. E. P. De Jesús and M. D. Graham, Data-driven low-dimensional dynamic model of kolmogorov flow, Phys. Rev. Fluids, 8 (2023), p. 044402.
  * [17] DeepMind, I. Babuschkin, K. Baumli, A. Bell, S. Bhupatiraju, J. Bruce, P. Buchlovsky, D. Budden, T. Cai, A. Clark, I. Danihelka, A. Dedieu, C. Fantacci, J. Godwin, C. Jones, R. Hemsley, T. Hennigan, M. Hessel, S. Hou, S. Kapturowski, T. Keck, I. Kemaev, M. King, M. Kunesch, L. Martens, H. Merzic, V. Mikulik, T. Norman, G. Papamakarios, J. Quan, R. Ring, F. Ruiz, A. Sanchez, L. Sartran, R. Schneider, E. Sezener, S. Spencer, S. Srinivasan, M. Stanojević, W. Stokowiec, L. Wang, G. Zhou, and F. Viola, The DeepMind JAX Ecosystem, 2020, http://github.com/google-deepmind.
  * [18] L. Dinh, J. Sohl-Dickstein, and S. Bengio, Density estimation using real NVP, arXiv preprint arXiv:1605.08803, (2016).
  * [19] G.-R. Duan and R. J. Patton, A note on hurwitz stability of matrices, Automatica, 34 (1998), pp. 509–511.
  * [20] V. Duruisseaux, J. W. Burby, and Q. Tang, Approximation of nearly-periodic symplectic maps via structure-preserving neural networks, Scientific Reports, 13 (2023), p. 8351.
  * [21] N. Fenichel, Geometric singular perturbation theory for ordinary differential equations, Journal of Differential Equations, 31 (1979), pp. 53–98.
  * [22] C. Foias, M. Jolly, I. Kevrekidis, G. Sell, and E. Titi, On the computation of inertial manifolds, Physics Letters A, 131 (1988), pp. 433–436.
  * [23] C. Foias, G. Sell, and E. Titi, Exponential tracking and approximation of inertial manifolds for dissipative nonlinear equations, Journal of Dynamics and Differential Equations, 1 (1989), pp. 199–244.
  * [24] C. Foias, G. R. Sell, and R. Temam, Inertial manifolds for nonlinear evolutionary equations, Journal of differential equations, 73 (1988), pp. 309–353.
  * [25] N. Gicquel, J. Anderson, and I. Kevrekidis, Noninvertibility and resonance in discrete-time neural networks for time-series processing, Physics Letters A, 238 (1998), pp. 8–18.
  * [26] A. Gorban and I. Karlin, Hilbert’s 6th problem: Exact and approximate hydrodynamic manifolds for kinetic equations, Bulletin of the American Mathematical Society, 51 (2013), pp. 186–246.
  * [27] H. Grad, On the kinetic theory of rarefied gases, Communications on Pure and Applied Mathematics, 2 (1949), pp. 331–407.
  * [28] J.-L. Guermond and S. Prudhomme, A fully discrete nonlinear galerkin method for the 3d navier-stokes equations, Numerical Methods for Partial Differential Equations, 24 (2008), pp. 759 – 775.
  * [29] G. W. Hammett and F. W. Perkins, Fluid moment models for Landau damping with application to the ion-temperature-gradient instability, Physical review letters, 64 (1990), p. 3019.
  * [30] R. A. Horn and C. R. Johnson, Matrix Analysis, Cambridge University Press, 1985.
  * [31] F. Jauberteau, C. Rosier, and R. Temam, A nonlinear galerkin method for the navier-stokes equations, Computer Methods in Applied Mechanics and Engineering, 80 (1990), pp. 245–260.
  * [32] P. Jin, Z. Zhang, A. Zhu, Y. Tang, and G. E. Karniadakis, SympNets: Intrinsic structure-preserving symplectic networks for identifying hamiltonian systems, Neural Networks, 132 (2020), pp. 166–179.
  * [33] M. S. Jolly, Explicit construction of an inertial manifold for a reaction diffusion equation, Journal of Differential Equations, 78 (1989), pp. 220–261.
  * [34] M. S. Jolly, I. Kevrekidis, and E. S. Titi, Approximate inertial manifolds for the Kuramoto-Sivashinsky equation: analysis and computations, Physica D: Nonlinear Phenomena, 44 (1990), pp. 38–60.
  * [35] C. Jones, Geometric singular perturbation theory, 1995.
  * [36] E. D. Koronaki, N. Evangelou, C. P. Martin-Linares, E. S. Titi, and I. G. Kevrekidis, Nonlinear dimensionality reduction then and now: AIMs for dissipative PDEs in the ML era, 2023, https://arxiv.org/abs/2310.15816.
  * [37] M. A. Kramer, Nonlinear principal component analysis using autoassociative neural networks, AIChE Journal, 37 (1991), pp. 233–243.
  * [38] K. Krischer, R. Rico-Martínez, I. G. Kevrekidis, H. H. Rotermund, G. Ertl, and J. L. Hudson, Model identification of a spatiotemporally varying catalytic reaction, AIChE Journal, 39 (1993), pp. 89–98.
  * [39] K. Lee and K. T. Carlberg, Model reduction of dynamical systems on nonlinear manifolds using deep convolutional autoencoders, Journal of Computational Physics, 404 (2020), p. 108973.
  * [40] M. Lezcano-Casado, Trivializations for gradient-based optimization on manifolds, in Advances in Neural Information Processing Systems, NeurIPS, 2019, pp. 9154–9164.
  * [41] A. J. Linot, J. W. Burby, Q. Tang, P. Balaprakash, M. D. Graham, and R. Maulik, Stabilized neural ordinary differential equations for long-time forecasting of dynamical systems, Journal of Computational Physics, 474 (2023), p. 111838.
  * [42] A. J. Linot and M. D. Graham, Deep learning to discover and predict dynamics on an inertial manifold, Phys. Rev. E, 101 (2020), p. 062209.
  * [43] A. J. Linot and M. D. Graham, Data-driven reduced-order modeling of spatiotemporal chaos with neural ordinary differential equations, Chaos: An Interdisciplinary Journal of Nonlinear Science, 32 (2022), p. 073110.
  * [44] A. J. Linot, K. Zeng, and M. D. Graham, Turbulence control in plane couette flow using low-dimensional neural ode-based models and deep reinforcement learning, International Journal of Heat and Fluid Flow, 101 (2023), p. 109139.
  * [45] S. Liu, Q. Tang, and X.-z. Tang, A parallel cut-cell algorithm for the free-boundary Grad–Shafranov problem, SIAM Journal on Scientific Computing, 43 (2021), pp. B1198–B1225.
  * [46] Y. Liu, J. N. Kutz, and S. L. Brunton, Hierarchical deep learning of multiscale differential equation time-steppers, Philosophical Transactions of the Royal Society A, 380 (2022), p. 20210200.
  * [47] E. Lorenz, Predictability: a problem partly solved, Seminar on Predictability, 4-8 September 1995, 1 (1995), pp. 1–18.
  * [48] F. Lu, K. K. Lin, and A. J. Chorin, Data-based stochastic model reduction for the kuramoto–sivashinsky equation, Physica D: Nonlinear Phenomena, 340 (2017), pp. 46–57.
  * [49] L. Pareschi and G. Russo, Implicit–explicit Runge–Kutta schemes and applications to hyperbolic systems with relaxation, Journal of Scientific Computing, 25 (2005), pp. 129–155.
  * [50] R. Rico-Martínez, I. Kevrekidis, and K. Krischer, Nonlinear system identification using neural networks: dynamics and instabilities, Neural networks for chemical engineers, (1995), pp. 409–442.
  * [51] R. Rico-Martínez, K. Krischer, I. G. Kevrekidis, M. C. Kube, and J. L. Hudson, Discrete- vs. continuous-time nonlinear signal processing of Cu electrodissolution data, Chemical Engineering Communications, 118 (1992), pp. 25–48.
  * [52] R. Rico-Martínez, R. Adomaitis, and I. Kevrekidis, Noninvertibility in neural networks, Computers & Chemical Engineering, 24 (2000), pp. 2417–2433.
  * [53] S. Shvartsman, C. Theodoropoulos, R. Rico-Martínez, I. Kevrekidis, E. Titi, and T. Mountziaris, Order reduction for nonlinear dynamic models of distributed reacting systems, Journal of Process Control, 10 (2000), pp. 177–184.
  * [54] S. Shvartsman, C. Theodoropoulos, R. Rico-Martínez, I. Kevrekidis, E. Titi, and T. Mountziaris, Order reduction for nonlinear dynamic models of distributed reacting systems, Journal of Process Control, 10 (2000), pp. 177–184.
  * [55] S. Y. Shvartsman and I. G. Kevrekidis, Nonlinear model reduction for control of distributed systems: A computer-assisted study, AIChE Journal, 44 (1998), pp. 1579–1595.
  * [56] R. Temam, Do inertial manifolds apply to turbulence?, Physica D: Nonlinear Phenomena, 37 (1989), pp. 146–152.
  * [57] R. Temam, Induced trajectories and approximate inertial manifolds, ESAIM: M2AN, 23 (1989), pp. 541–561.
  * [58] T. Teshima, I. Ishikawa, K. Tojo, K. Oono, M. Ikeda, and M. Sugiyama, Coupling-based invertible neural networks are universal diffeomorphism approximators, in Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS’20, Red Hook, NY, USA, 2020, Curran Associates Inc.
  * [59] E. S. Titi, On approximate inertial manifolds to the navier-stokes equations, Journal of Mathematical Analysis and Applications, 149 (1990), pp. 540–557.
  * [60] K. Zeng, C. E. P. D. Jesús, A. J. Fox, and M. D. Graham, Autoencoders for discovering manifold dimension and coordinates in data from complex dynamical systems, 2023, https://arxiv.org/abs/2305.01090.
  * [61] K. Zeng, A. J. Linot, and M. D. Graham, Data-driven control of spatiotemporal chaos with reduced-order neural ode-based models and reinforcement learning, Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 478 (2022), p. 20220297.
  * [62] H. Zhang, Y. Liu, and R. Maulik, Semi-implicit neural ordinary differential equations for learning chaotic systems, in NeurIPS 2023 Workshop Heavy Tails in Machine Learning, 2023.