Koopman operator learning using invertible neural networks
In Koopman operator theory, a finite-dimensional nonlinear system is
transformed into an infinite but linear system using a set of observable
functions. However, manually selecting observable functions that span the
invariant subspace of the Koopman operator based on prior knowledge is
inefficient and challenging, particularly when little or no information is
available about the underlying systems. Furthermore, current methodologies tend
to disregard the importance of the invertibility of observable functions, which
leads to inaccurate results. To address these challenges, we propose the
so-called FlowDMD, aka Flow-based Dynamic Mode Decomposition, that utilizes the
Coupling Flow Invertible Neural Network (CF-INN) framework. FlowDMD leverages
the intrinsically invertible characteristics of the CF-INN to learn the
invariant subspaces of the Koopman operator and accurately reconstruct state
variables. Numerical experiments demonstrate the superior performance of our
algorithm compared to state-of-the-art methodologies.
Yuhuang Meng, Jianguo Huang, Yue Qiu
# Physics-informed invertible neural network for the Koopman operator learning 1

###### Abstract

In Koopman operator theory, a finite-dimensional nonlinear system is transformed into an infinite but linear system using a set of observable functions. However, manually selecting observable functions that span the invariant subspace of the Koopman operator based on prior knowledge is inefficient and challenging, particularly when little or no information is available about the underlying systems. Furthermore, current methodologies tend to disregard the importance of the invertibility of observable functions, which leads to inaccurate results. To address these challenges, we propose the so-called FlowDMD, a Flow-based Dynamic Mode Decomposition that utilizes the Coupling Flow Invertible Neural Network (CF-INN) framework. FlowDMD leverages the intrinsically invertible characteristics of the CF-INN to learn the invariant subspaces of the Koopman operator and accurately reconstruct state variables. Numerical experiments demonstrate the superior performance of our algorithm compared to state-of-the-art methodologies.

keywords: Koopman operator, Generative models, Invertible neural networks +
## 1 Introduction

Nonlinear dynamic systems are widely prevalent in both theory and engineering applications. Since the governing equations are generally unknown in many situations, it can be challenging to study the systems directly based on the first principles. Fortunately, the data about the systems of interest could be available by experiments or observations. Instead, one could seek to understand the behavior of the nonlinear system through the data-driven approaches [1; 2; 3; 4; 5].

The Koopman operator [6], which embeds the nonlinear system of interest into an infinite dimensional linear space by observable functions has attracted lots of attention. The Koopman operator acts on the infinite dimensional Hilbert space and aims to capture the full representations of the nonlinear systems. Dynamic mode decomposition (DMD) calculates the spectral decomposition of the Koopman operator numerically by extracting dynamic information from the collected data. Concretely, DMD devises a procedure to extract the spectral information directly from a data sequence without an explicit formulation of the Koopman operator, which is efficient for handling high dimensional data [7]. Variants of DMD are proposed to address challenges in different scenarios [8; 9; 10; 11; 12; 13; 14; 15].

The selection of observable functions plays an essential role in the DMD algorithm. Exact DMD [8] exploits the identity mapping as the observables. This implies that one uses a linear system to approximate a nonlinear system with given data [16]. This would yield inaccurate or even completely mistaken outcomes. Furthermore, the short-term prediction of Exact DMD might be acceptable for some cases, but the long-term prediction is probably unreliable. Typically, prior knowledge is required to select the observable functions that span the invariant subspace of the Koopman operator. However, the invariant subspace is not simply available. In order to overcome the limitations of the Exact DMD algorithm and capture the full feature of the nonlinear system, several data-driven selection strategies for observable functions have been proposed. Extended DMD (EDMD) [17] lifts the state variables from the original space into a higher dimensional space using the dictionary functions. The accuracy and rate of convergence of EDMD depend on the choice of the dictionary functions. Therefore, EDMD needs as many dictionary functions as possible. This implies that the set of dictionary functions (nonlinear transformations) should be sufficiently complex, which results in enormous computational cost. Kernel based DMD (KDMD) [18] differs from EDMD in that it utilizes the kernel trick to exploit the implicit expression of dictionary functions, whereas EDMD uses the explicit expression of dictionary functions. Nonetheless, both EDMD and KDMD are prone to overfitting [19], which leads to large generalization error. How to efficiently choose the observable functions that span the invariant subspace of the Koopman operator
becomes a significant challenge.

In contrast to EDMD and KDMD, observable functions can be represented by neural networks. Dictionary learning [20] couples the EDMD with a set of trainable dictionary functions, where dictionary functions are represented by a fully connected neural network and an untrainable component. Fixing the partial dictionary function facilitates the reconstruction of the state variables, however, this setting implicitly assumes that linear term lies in the invariant subspace of the Koopman operator. Yeung et al. [21] select low-dimensional dictionary functions more efficiently using deep neural networks.

Autoencoder (AE) neural networks have been widely applied to learn the optimal observable functions and reconstruction functions in Koopman embedding [19; 22; 23; 24; 25; 26]. Concretely, the invariant subspace of the Koopman operator and reconstruction functions are represented by the encoder and decoder network in AE, respectively. Lusch et al. [23] utilize neural networks to identify the Koopman eigenfunctions and introduced an auxiliary network to cope with the dynamic systems with continuous spectrum. Azencot et al. [24] propose the Consistent Koopman AE model that combines the forward-backward DMD method [27] with the AE model. This approach extracts the latent representation of high-dimensional non-linear data and eliminates the effect of noise in the data simultaneously. Pan and Duraisamy [25] parameterize the structure of the transition matrix in linear space and construct an AE model to learn the residual of the DMD. Li and Jiang [26] utilize deep learning and the Koopman operator to model the nonlinear multiscale dynamical problems, where coarse-scale data is used to learn the fine-scale information through a set of multiscale basis functions. Wang et al. [28] propose Koopman Neural Forecaster (KNF) combining AE with Koopman operator theory to predict the data with distributional shifts.

Representing Koopman embedding by dictionary learning or AE networks has several drawbacks. Firstly, the reconstruction in dictionary learning partially fixes the dictionary functions, which leads to a low level of interpretability of the model. Secondly, the encoder and decoder in an AE model are trained simultaneously, but neither of them is invertible, cf. [29] for more details. Moreover, due to the structural noninvertibility of the encoder and decoder, it typically requires a large amount of training data in order to obtain accurate representations, which makes the AE model prone to overfitting. Alford-Lago et al. [29] analyze the property of both the encoder and decoder in AE and proposed the deep learning dynamic mode decomposition (DLDMD). Bevanda et al. [30] constructed a conjugate map between the nonlinear system and its Jacobian linearization, which is learned by a diffeomorphic neural network.
In this paper, we develop a novel architecture that incorporates physical knowledge to learn the Koopman embedding. Specifically, we apply the coupling flow invertible neural networks (CF-INN) to learn the observable functions and reconstruction functions. The invertibility of the learned observable functions makes our method more flexible than dictionary learning or AE learning. Our contributions are three-folds:

1. We utilize an structurally invertible mapping to reconstruct state variables, which increases the interpretability of the neural network and alleviates the overfitting of AE.
2. The difficulty of learning the observable functions and observable functions is reduced by exploiting their structural invertibility of neural networks. Therefore, the reconstruction error in the loss function could be eliminated.
3. As the physical information is embedded into the model, the number of parameters is reduced to achieve comparable accuracy with other methods. Additionally, the parameters to be optimized are reduced dramatically since the learned mappings and their inverse share the same parameters.

This paper is organized as follows. In Section 2, we briefly review the Koopman operator theory and DMD. In Section 3, we present the structure of CF-INN and introduce how to learn the invariant subspace of the Koopman operator and the reconstruction functions. In Section 4, several numerical experiments are performed to demonstrate the performance of our method, and we summarize our work in Section 5.

## 2 Preliminaries

### Koopman operator theory

Consider the nonlinear autonomous system in discrete form,

\[\mathbf{x}_{k+1}=f(\mathbf{x}_{k}),\quad\mathbf{x}_{k}\in\mathcal{M}\subset \mathbb{R}^{m}, \tag{1}\]

where \(\mathcal{M}\) represents the set of state space, \(f:\mathcal{M}\rightarrow\mathcal{M}\) is an unknown nonlinear map, and \(k\) is the time index.

**Definition 1** (Koopman operator [16]).: _For the nonlinear system (1), the Koopman operator \(\mathcal{K}\) is an infinite-dimensional linear operator that acts on all observable functions \(g:\mathcal{M}\rightarrow\mathbb{C}\) such that_

\[\mathcal{K}g(\mathbf{x})=g(f(\mathbf{x})).\]

_Here, \(g(x)\in\mathcal{H}\) and \(\mathcal{H}\) represents the infinite dimensional Hilbert space._
Through the observable functions, the nonlinear system (1) could be transformed into an infinite-dimensional linear system using the Koopman operator,

\[g(\mathbf{x}_{k+1})=g(f(\mathbf{x}_{k}))=\mathcal{K}g(\mathbf{x}_{k}). \tag{2}\]

Note that the Koopman operator is linear, _i.e._, \(\mathcal{K}(\alpha_{1}g_{1}(\mathbf{x})+\alpha_{2}g_{2}(\mathbf{x}))=\alpha_{1} g_{1}(f(\mathbf{x}))+\alpha_{2}g_{2}(f(\mathbf{x}))\), with \(g_{1}(\mathbf{x}),g_{2}(\mathbf{x})\in\mathcal{H}\) and \(\alpha_{1},\alpha_{2}\in\mathbb{R}\). As \(\mathcal{K}\) is an infinite-dimensional operator, we denote its eigenfunctions and eigenvalues by \(\{\lambda_{i},\varphi_{i}(x)\}_{i=0}^{\infty}\) such that \(\mathcal{K}\varphi_{i}(\mathbf{x})=\lambda_{i}\varphi_{i}(\mathbf{x})\), where \(\varphi_{i}(\mathbf{x}):\mathcal{M}\rightarrow\mathbb{R}\), \(\lambda_{i}\in\mathbb{C}\).

The Koopman eigenfunctions define a set of intrinsic measurement coordinates, then a vector-valued observable function \(\mathbf{g}(\mathbf{x})=[g_{1}(\mathbf{x}),\cdots,g_{n}(\mathbf{x})]^{T}\) could be written in terms of the Koopman eigenfunctions,

\[\mathbf{g}(\mathbf{x}_{k})=\begin{bmatrix}g_{1}(\mathbf{x}_{k})\\ \vdots\\ g_{n}(\mathbf{x}_{k})\end{bmatrix}=\sum_{i=1}^{\infty}\varphi_{i}(\mathbf{x}_{ k})\begin{bmatrix}<\varphi_{i},g_{1}>\\ \vdots\\ <\varphi_{i},g_{n}>\end{bmatrix}=\sum_{i=1}^{\infty}\varphi_{i}(\mathbf{x}_{k}) \mathbf{v}_{i}, \tag{3}\]

where \(\mathbf{v}_{i}\) refers to the \(i\)-th Koopman mode with respect to the Koopman eigenfunction \(\varphi_{i}(\mathbf{x})\). Combining (2) and (3), we have the decomposition of a vector-valued observable functions

\[\mathbf{g}(\mathbf{x}_{k+1})=\mathcal{K}\mathbf{g}(\mathbf{x}_{k})=\mathcal{K }\sum_{i=1}^{\infty}\varphi_{i}(\mathbf{x}_{k})\mathbf{v}_{i}=\sum_{i=1}^{ \infty}\lambda_{i}\varphi_{i}(\mathbf{x}_{k})\mathbf{v}_{i}.\]

Furthermore, the decomposition could be rewritten as

\[\mathbf{g}(\mathbf{x}_{k})=\sum_{i=1}^{\infty}\lambda_{i}^{k}\varphi_{i}( \mathbf{x}_{0})\mathbf{v}_{i}.\]

In practice, we need a finite-dimensional representation of the infinite-dimensional Koopman operator. Denote the \(n\)-dimensional invariant subspace of the Koopman operator \(\mathcal{K}\) by \(\mathcal{H}_{g}\), _i.e._, \(\forall g(\mathbf{x})\in\mathcal{H}_{g},\mathcal{K}g(\mathbf{x})\in\mathcal{H }_{g}\). Let \(\{g_{i}(\mathbf{x})\}_{i=1}^{n}\) be one set of basis of \(\mathcal{H}_{g}\), this induces a finite-dimensional linear operator \(\mathbf{K}\)[16], which projects the Koopman operator \(\mathcal{K}\) onto \(\mathcal{H}_{g}\), _i.e._, for the \(n\)-dimensional vector-valued observable functions \(\mathbf{g}(\mathbf{x})=[g_{1}(\mathbf{x}),\cdots,g_{n}(\mathbf{x})]^{T}\), we have

\[\mathbf{g}(x_{k+1})=\begin{bmatrix}g_{1}(x_{k+1})\\ \vdots\\ g_{n}(x_{k+1})\end{bmatrix}=\begin{bmatrix}\mathcal{K}g_{1}(x_{k})\\ \vdots\\ \mathcal{K}g_{n}(x_{k})\end{bmatrix}=\mathbf{K}\begin{bmatrix}g_{1}(x_{k})\\ \vdots\\ g_{n}(x_{k})\end{bmatrix}=\mathbf{K}\mathbf{g}(x_{k}) \tag{4}\]
### Dynamic mode decomposition

DMD approximates the spectral decomposition of the Koopman operator numerically. Given the state variables \(\{\mathbf{x}_{0},\mathbf{x}_{1},\cdots,\mathbf{x}_{p}\}\) and a vector-valued observable function \(\mathbf{g}(\mathbf{x})=[g_{1}(\mathbf{x}),\cdots,g_{n}(\mathbf{x})]^{T}\), then we get the sequence \(\{\mathbf{g}(\mathbf{x}_{0}),\mathbf{g}(\mathbf{x}_{1}),\cdots,\mathbf{g}( \mathbf{x}_{p})\}\), where each \(\mathbf{g}(\mathbf{x}_{k})\in\mathbb{R}^{n}\) is the observable snapshot of the \(k\)-th time step. According to (4), we have

\[\mathbf{g}(\mathbf{x}_{k+1})=\mathbf{K}\mathbf{g}(\mathbf{x}_{k}),\]

where \(\mathbf{K}\in\mathbb{R}^{n\times n}\) is the matrix form of the finite-dimensional operator. For the two data matrices, \(\mathbf{X}=[\mathbf{g}(\mathbf{x}_{0}),\cdots,\mathbf{g}(\mathbf{x}_{p-1})]\) and \(\mathbf{Y}=[\mathbf{g}(\mathbf{x}_{1}),\cdots,\mathbf{g}(\mathbf{x}_{p})]\), where \(\mathbf{X}\) and \(\mathbf{Y}\) are both in \(\mathbb{R}^{n\times p}\), which satisfies \(\mathbf{Y}=\mathbf{K}\mathbf{X}\). Therefore, \(\mathbf{K}\) can be represented by

\[\mathbf{K}=\mathbf{Y}\mathbf{X}^{\dagger},\]

where \(\mathbf{X}^{\dagger}\) denotes the Moore-Penrose inverse of \(\mathbf{X}\).

The Exact DMD algorithm developed by Tu et al. [8] computes dominant eigen-pairs (eigenvalue and eigenvector) of \(\mathbf{K}\) without the explicit formulation of \(\mathbf{K}\). In Algorithm 1, we present the DMD algorithm on the observable space, which is a general form of the Exact DMD algorithm. When using the identical mapping as the observable functions, _i.e._, \(\mathbf{g}(\mathbf{x})=\mathbf{x}\), Algorithm 1 is identical to the Exact DMD algorithm.

```
1. Compute the (reduced) SVD of \(\mathbf{X}\), \(\mathbf{X}=\mathbf{U}_{\mathbf{r}}\mathbf{\Sigma}_{\mathbf{r}}\mathbf{V}_{ \mathbf{r}}^{*}\), where \(\mathbf{U}_{\mathbf{r}}\in\mathbb{C}^{n\times r}\), \(\mathbf{\Sigma}_{\mathbf{r}}\in\mathbb{R}^{r\times r}\), \(\mathbf{V}_{\mathbf{r}}\in\mathbb{C}^{p\times r}\).
2. Compute \(\tilde{\mathbf{K}}=\mathbf{U}_{\mathbf{r}}^{*}\mathbf{Y}\mathbf{V}_{\mathbf{r} }\mathbf{\Sigma}_{\mathbf{r}}^{-1}\).
3. Compute the eigen-pairs of \(\tilde{\mathbf{K}}\): \(\tilde{\mathbf{K}}\mathbf{W}=\mathbf{W}\mathbf{\Lambda}\).
4. Reconstruct the eigen-pairs of \(\mathbf{K}\), where eigenvalues of \(\mathbf{K}\) are diagonal entries of \(\Lambda\), the corresponding eigenvectors of \(\mathbf{K}\)(DMD modes) are columns of \(\mathbf{\Phi}=\mathbf{Y}\mathbf{V}_{\mathbf{r}}\mathbf{\Sigma}_{\mathbf{r}}^{ -1}\mathbf{W}\).
5. Approximate the observation data via DMD, \(\hat{\mathbf{g}}(\mathbf{x}_{k})=\mathbf{\Phi}\mathbf{\Lambda}^{k}\mathbf{b}\), where \(\mathbf{b}=\mathbf{\Phi}^{\dagger}\mathbf{g}(\mathbf{x}_{0})\).
6. Reconstruct the state variables \(\hat{\mathbf{x}}_{k}=\mathbf{g}^{-1}(\hat{\mathbf{g}}(\mathbf{x}_{k}))= \mathbf{g}^{-1}\left(\mathbf{\Phi}\mathbf{\Lambda}^{k}\mathbf{b}\right)\).
```

**Algorithm 1** DMD on observable space [16; 31]

### State reconstruction

Koopman operator theory utilizes observable functions \(\mathbf{g}\) to transform the nonlinear system (1) into a linear system while preserving the nonlinearity. Evolving the nonlinear system (1) is computationally expensive or even impossible when \(f\) is
unknown, whereas evolving through the Koopman operator (2) offers a promising and computationally efficient approach.

Figure 1 illustrates the relation between the nonlinear evolution \(f\) and the Koopman operator evolution where the system evolves linearly in the observation space \(\mathcal{H}\). By computing the Koopman eigenvalues and modes, we can make predictions of the observable functions \(\mathbf{g}(\mathbf{x})\). We could reconstruct the state \(\mathbf{x}\) by the inverse of the observable functions \(\mathbf{g}^{-1}(\mathbf{x})\) provided that \(\mathbf{g}(\mathbf{x})\) is invertible. The invertibility of observable functions is essential to ensure the reconstruction accuracy and the interpretability of the outcomes.

Typical observable functions \(\mathbf{g}(\mathbf{x})\) selection are performed manually based on prior knowledge. Exact DMD takes the identical mapping, while the EDMD utilizes a set of pre-defined functions such as polynomials, Fourier modes, radial basis functions, and so forth [17]. However, these methods can be inaccurate and inefficient for Koopman embeddings learning. Deep neural networks, as efficient global nonlinear approximators, could be applied to represent the observable function \(\mathbf{g}(\mathbf{x})\) and the reconstruction function \(\mathbf{g}^{-1}(\mathbf{x})\). Several studies have demonstrated that the encoder and decoder networks in AE correspond to \(\mathbf{g}(\mathbf{x})\) and \(\mathbf{g}^{-1}(\mathbf{x})\), respectively [19; 22; 23; 24; 25; 26].

In practical applications, it is not always guaranteed that \(\mathbf{g}(\mathbf{x})\) is invertible. In the learning Koopman embedding via AE, the invertibility of \(\mathbf{g}(\mathbf{x})\) is enforced through numerical constraints, _i.e._, the reconstruction error \(\|\mathbf{x}-\mathbf{g}^{-1}(\mathbf{g}(\mathbf{x}))\|_{2}^{2}\), which tends to result in overfitting and suboptimal performance [29]. Besides, the reconstruction error is trained simultaneously with the prediction error and the linearity error [23]. The weights assigned to each loss term are hyperparameters that can be challenging to tune. In this paper, we propose a structurally invertible mapping learning framework, which eliminates the need for the reconstruction term in the loss function and yields more robust and accurate results. We present the details of our method in Section 3.

Figure 1: Koopman operator and inverse of observable functions
## 3 Learning Koopman embedding by invertible neural networks

In this section, we first briefly review the AE neural network and demonstrate the limitation of this class of neural networks in the Koopman embedding learning. Then, we introduce our method to overcome this limitation.

### Drawback of AE in the Koopman embedding learning

Most of the work use the Autoencoder (AE) neural networks as the backbone to learn the invariant subspace of the Koopman operator and reconstruct the state variables. AE as the frequently-used unsupervised learning structure of neural networks, consists of two parts, _i.e._, the encoder \(\mathcal{E}\) and the decoder \(\mathcal{D}\). AE learns these two mappings (functions) \(\mathcal{E}\) and \(\mathcal{D}\) by optimizing

\[\min_{\mathcal{E},\mathcal{D}}\mathbb{E}_{x\sim m(x)}[\text{loss}(x,\mathcal{ D}\circ\mathcal{E}(x))]. \tag{5}\]

Here \(m(x)\) denotes the distribution of the input data, \(\text{loss}(x,y)\) describes the difference between \(x\) and \(y\), and \(\mathbb{E}(\cdot)\) represents the expectation.

**Definition 2**.: _Let \(f_{1}:S\to S^{\prime}\) be an arbitrary mapping, and it is said to be invertible if there exists a mapping \(f_{2}:S^{\prime}\to S\) such that_

\[f_{1}\circ f_{2}=\mathcal{I},f_{2}\circ f_{1}=\mathcal{I},\]

_where \(\mathcal{I}\) is the identity mapping. Then, \(f_{2}\) is said to be the inverse mapping of \(f_{1}\)._

Let \(\mathcal{E}\) and \(\mathcal{D}\) be two mappings learned by AE such that \(\mathcal{D}\circ\mathcal{E}\approx\mathcal{I}\). However, the reverse order of the mapping \(\mathcal{E}\circ\mathcal{D}\) is not always a good approximation to the identity mapping, moreover, \(\mathcal{E}\) and \(\mathcal{D}\) are generally not invertible [29]. The main reason is that while AE strives to reach \(\mathcal{D}\circ\mathcal{E}\approx\mathcal{I}\), it omits the additional constraint \(\mathcal{E}\circ\mathcal{D}\approx\mathcal{I}\) which requires the latent variable data to train. Unfortunately, the latent variables are not accessible, thus rendering it impossible for AE to satisfy \(\mathcal{E}\circ\mathcal{D}\approx\mathcal{I}\) and \(\mathcal{D}\circ\mathcal{E}\approx\mathcal{I}\) simultaneously.

AE learns an identity mapping \(\mathcal{I}\) from a training data set \(\mathcal{S}\), _i.e._, for any \(\mathbf{x}\in\mathcal{S},\mathcal{D}\circ\mathcal{E}(\mathbf{x})\approx\mathbf{x}\). For data out of the set \(\mathcal{S}\), the mapping learned by AE may perform badly. In other words, AE may have poor generalization capability. Next, we use a preliminary experiment to demonstrate this limitation. The details of this numerical example are given in Section 4.1. We use the structure of AE defined in [26] and randomly generate 120 trajectories to train the AE, and the results are depicted by Figure 2.
Figure 2 compares the input data points out of the distribution of the training data with the corresponding reconstructed data points using the trained AE model. Figure 2(a) shows the density distribution of training data set \(\mathcal{S}\), which provides a rough illustration of the data space \(\mathcal{S}\). For the reconstruction test of AE, we generate three types of data, _i.e._, the sin-shaped scatters, the S-shaped scatters, and scatters from the standard 2-d normal distribution. We plot the corresponding input points (blue) and reconstructed data points (red) of the AE. The results shown in the next three subfigures illustrate that AE can reconstruct the input data points nearby the training data set \(\mathcal{S}\) very well. But for the data points far away from \(\mathcal{S}\), AE performs badly. The same situation happens in learning the Koopman embedding. Specifically, in the training process of AE, one aims to find the Koopman invariant space by minimizing the error of the Koopman embedding learning and the reconstruction error. However, minimizing the error between latent variables and their corresponding reconstruction denoted by \(\text{loss}(\mathbf{x},\mathcal{E}\circ\mathcal{D}(\mathbf{x}))\) is intractable. This result is in poor stability and generalization capability.

### Structure of Cf-Inn

We have shown that the mapping learned by AE performs poorly, which inspires us that invertibility can greatly reduce computational complexity and yields better

Figure 2: Generalization capability test of AE. (a) the training data distribution. (b) the \(sin(x)\) test function. (c) S-shaped scatters test. (d) random scatters from 2-d standard normal distribution.
generalization capability. Next, we introduce an invertible neural network to overcome the drawback of AE. Let \(\mathbf{g}_{\boldsymbol{\theta}}(\mathbf{x}):\mathbf{X}\rightarrow\mathbf{Y}\) denote the input-output mapping of the invertible neural network, where \(\boldsymbol{\theta}\) represents the parameters of the neural network. Let \(\mathbf{f}_{\boldsymbol{\theta}}\) be the inverse mapping of \(\mathbf{g}_{\boldsymbol{\theta}}\) which shares the same parameters with \(\mathbf{g}_{\boldsymbol{\theta}}\). Then we can reconstruct \(x\) in the backward direction by \(\mathbf{f}_{\boldsymbol{\theta}}(\mathbf{y}):\mathbf{Y}\rightarrow\mathbf{X}\). In generative tasks of machine learning, the forward generating direction is called the flow direction and the backward direction is called the normalizing direction. Next, we introduce the concept of coupling flows, which belongs to the invertible neural networks.

**Definition 3** (Coupling flow [32]).: _Let \(m\in\mathbb{N}\) and \(m\geq 2\), for a vector \(\mathbf{z}\in\mathbb{R}^{m}\) and \(2\leq q\leq m-1\), we define \(\mathbf{z}_{up}\) as the vector \((z_{1},\ldots,z_{q})^{\top}\in\mathbb{R}^{q}\) and \(\mathbf{z}_{low}\) as the vector \((z_{q+1},\ldots,z_{m})^{\top}\in\mathbb{R}^{m-q}\). A coupling flow (CF), denoted by \(h_{q,\tau}\), has the following form,_

\[h_{q,\tau}(\mathbf{z}_{up},\mathbf{z}_{low})=(\mathbf{z}_{up},\tau(\mathbf{z}_ {low},\sigma(\mathbf{z}_{up}))),\]

_where \(\sigma:\mathbb{R}^{q}\rightarrow\mathbb{R}^{l}\), and \(\tau(\cdot,\sigma(\mathbf{y})):\mathbb{R}^{m-q}\times\mathbb{R}^{l}\rightarrow \mathbb{R}^{m-q}\) is a bijection mapping for any \(\mathbf{y}\in\mathbb{R}^{q}\)._

A coupling flow defined in _Definition 3_ is invertible if and only if \(\tau\) is invertible and its inverse \(h_{q,\tau}^{-1}(\mathbf{z}_{up},\mathbf{z}_{low})=(\mathbf{z}_{up},\tau^{-1}( \mathbf{z}_{low},\sigma(\mathbf{z}_{up})))\)[33]. The key point of making the CF invertible is the invertibility of \(\tau\). One of the mostly used CF is the affine coupling function (ACF) [34, 35, 36], where \(\tau\) is an invertible element-wise function.

**Definition 4** (Affine coupling function [33]).: _Define an affine coupling function by the mapping \(\Psi_{q,s,t}\) from \(\mathbb{R}^{q}\times\mathbb{R}^{m-q}\) to \(\mathbb{R}^{m}\) such that_

\[\Psi_{q,s,t}(\mathbf{z}_{up},\mathbf{z}_{low})=(\mathbf{z}_{up},(\mathbf{z}_{ low}+t(\mathbf{z}_{up}))\odot s(\mathbf{z}_{up})), \tag{6}\]

_where \(\odot\) is the Hadamard product, \(s,t:\mathbb{R}^{q}\rightarrow\mathbb{R}^{m-q}\) are two arbitrary vector-valued mappings._

Definition 4 defines the forward direction of computations, and the backward direction of computations is given by \(\Psi_{q,s,t}^{-1}(\mathbf{z}_{up},\mathbf{z}_{low})=(\mathbf{z}_{up},\mathbf{ z}_{low}\oslashs(\mathbf{z}_{up})-t(\mathbf{z}_{up}))\), where \(\oslash\) denotes the element-wise division of vectors. The mappings \(s\) and \(t\) in Definition 4 can be any nonlinear functions, neural networks such as fully-connected neural network (FNN) are typically used to parameterize \(t\) and \(s\).

Let \(\Psi_{1},\ldots,\Psi_{L}\) be a sequence of \(L\) affine coupling functions and define \(\mathbf{g}_{\boldsymbol{\theta}}=\Psi_{L}\circ\Psi_{L-1}\circ\cdots\Psi_{1}\), where \(\boldsymbol{\theta}\) represents the parameters of \(\{\Psi_{i}\}_{i=1}^{L}\). The resulted vector-valued function \(\mathbf{g}_{\boldsymbol{\theta}}\) is an invertible neural network and called by coupling flow invertible neural network (CF-INN) in this paper. Moreover, for any \(\Psi_{i}\), the division
index \(q\) of the input vector \(x\) is user-guided. In this paper, we set \(q=\lceil m/2\rceil\), where \(\lceil\cdot\rceil\) is the rounding function. Furthermore, in order to mix the information sufficiently, we can flip the ACF by using the form \(\bar{\Psi}_{q,s,t}(\mathbf{z}_{up},\mathbf{z}_{low})=((\mathbf{z}_{up}+t( \mathbf{z}_{low}))\odot s(\mathbf{z}_{low}),\mathbf{z}_{low})\). We plot the computation process of an ACF and a flipped ACF in Figure 3, where the network structure diagram left shows the forward direction and the network structure diagram right shows the backward direction. The red area is an ACF block and consists of a standard ACF and a flipped ACF, which is a CF-INN of depth 2.

When the depth (L) of a CF-INN is large, its training becomes challenging. The main curse is that the dividend term \(s\) is too small in \(\Psi\) in the backward direction computations. This can be solved by replacing the affine coupling functions with residual coupling functions. Similar idea has also been applied in the residual term of ResNet.

**Definition 5** (Residual coupling functions [37]).: _Define a residual affine coupling function (RCF) by the map \(\Psi_{q,s,t}\) from \(\mathbb{R}^{q}\times\mathbb{R}^{m-q}\) to \(\mathbb{R}^{m}\) such that_

\[\Psi_{q,t}(\mathbf{z}_{up},\mathbf{z}_{low})=(\mathbf{z}_{up},\mathbf{z}_{low }+t(\mathbf{z}_{up})),\]

_where \(t:\mathbb{R}^{q}\rightarrow\mathbb{R}^{m-q}\) is a nonlinear mapping._

RCFs are simplifications of ACFs and when we connect a RCF with a flipped RCF, we obtain a RCF block, which is a simplified ACF block in Figure 3.

### Loss function for Koopman embedding

In this paper, we use the CF-INN to learn the Koopman invariant subspace and the reconstructions simultaneously, where the forward direction of CF-INN is represented by \(\mathbf{g}_{\boldsymbol{\theta}}\) and its backward direction is represented by \(\mathbf{f}_{\boldsymbol{\theta}}\). The observable

Figure 3: The illustration of the forward and backward direction in an ACF block.
functions evolve linearly in the Koopman invariant subspace. Hence, the linearity constrained loss function that represents the DMD approximation error is given by

\[\mathcal{L}_{\text{linear}}=\sum_{t=1}^{T}||\mathbf{g}_{\boldsymbol{\theta}}( \mathbf{x}_{t})-\Phi\Lambda^{t}\Phi^{\dagger}\mathbf{g}_{\boldsymbol{\theta}}( \mathbf{x}_{0})||^{2}=\sum_{t=1}^{T}||\mathbf{g}_{\boldsymbol{\theta}}( \mathbf{x}_{t})-\hat{\mathbf{g}}_{\boldsymbol{\theta}}(\mathbf{x}_{t})||^{2},\]

where \(\hat{\mathbf{g}}_{\boldsymbol{\theta}}(\mathbf{x}_{t})=\Phi\Lambda^{t}\Phi^{ \dagger}\mathbf{g}_{\boldsymbol{\theta}}(\mathbf{x}_{0})\) is the DMD approximation of the observable functions \(\{\mathbf{g}(\mathbf{x}_{t})\}_{t=1}^{T}\) by using Algorithm 1. To reconstruct the states \(x_{t}\), the inverse mapping of \(\mathbf{g}\), _i.e_, \(\mathbf{f}_{\theta}\) corresponds to the backward direction of CF-INN. \(\mathbf{f}_{\theta}\) shares the same network structure and parameters with \(\mathbf{g}_{\theta}\). Therefore, the computational cost is greatly reduced, compared with AE that another neural network is required to parameterize the inverse mapping of \(\mathbf{g}_{\theta}\). The reconstruction loss due to the DMD approximation error is given by

\[\mathcal{L}_{\text{rec}}=\sum_{t=1}^{T}||\mathbf{x}_{t}-\mathbf{f}_{ \boldsymbol{\theta}}(\hat{\mathbf{g}}_{\boldsymbol{\theta}}(\mathbf{x}_{t})) ||^{2}.\]

The optimal parameters \(\boldsymbol{\theta}^{*}\) is given by

\[\boldsymbol{\theta}^{*}=\operatorname*{arg\,min}_{\boldsymbol{\theta}} \mathcal{L}_{\text{linear}}+\alpha\mathcal{L}_{\text{rec}},\]

where \(\alpha\) is a user-guard hyperparameter.

Compared with other Koopman embedding learning frameworks, the loss function in our approach is much more simplified. We summarize our CF-INN framework for the Koopman embedding learning in Figure 4 and our method is called FlowDMD since this framework uses a flow model based Dynamic Model Decomposition to compute the finite dimensional Koopman operator approximation and reconstruct system states.

Figure 4: The general framework of FlowDMD.
## 4 Numerical experiments

In this section, we use three numerical examples to demonstrate the efficiency of our method for learning the Koopman embedding and compare its performance with LIR-DMD [26] and Exact DMD. We use the Python library _FEniCS_[38] to compute the numerical solutions of PDEs, the Python library _PyDMD_[39] to complete the calculations of Exact DMD, and the Python library _PyTroch_[40] to train the neural networks. Besides, the Xavier normal initialization scheme [41] is utilized to initialize the weights of all neural networks, while the biases of all nodes are set to zero. All the networks are trained by the Adam optimizer [42] with an initial learning rate of \(10^{-3}\). In order to find the optimal parameters of the network, we use _ReduceLROnPlateau_[43] to adjust the learning rate during the training process for all numerical examples. For fairness, all the methods share the same training strategies. Denote \(x\) as the "true" value of the states and \(\hat{x}\) as its reconstruction. We use three metrics to evaluate different methods synthetically., i.e., the relative \(L_{2}\) error

\[\text{RL2E}(t)=\frac{||\hat{x}_{t}-x_{t}||_{2}}{||x_{t}||_{2}},\]

the mean squared error (MSE),

\[\text{MSE}(t)=\frac{||\hat{x}_{t}-x_{t}||_{2}^{2}}{m},\]

and the total relative \(L_{2}\) error

\[\text{TRL2E}=\sqrt{\frac{\sum_{t=1}^{T}||\hat{x}_{t}-x_{t}||_{2}^{2}}{\sum_{i= 1}^{T}||x_{t}||_{2}^{2}}}.\]

### Fixed-point attractor

The fixed-point attractor example [23] is given by

\[\begin{cases}x_{t+1,1}=\lambda x_{t,1},\\ x_{t+1,2}=\mu x_{t,2}+(\lambda^{2}-\mu)x_{t,1}^{2}.\end{cases}\]

The initial state is chosen randomly by \(x_{0,1}\sim U(0.2,4.2)\), \(x_{0,2}\sim U(0.2,4.2)\) and \(\lambda=0.9,\mu=0.5\). We divide the data set into three parts where the ratio of training, validation, and test is \(60\%,20\%\), and \(20\%\), respectively. The number of neurons of each layer for the encoder network in LIR-DMD is \(2,10,10,3\) and the number of neurons of decoder network is \(3,10,10,2\). This results in \(345\) trainable parameters for
LIR-DMD. We use three ACFs for this problem. The mappings \(t\) and \(s\) are parameterized by FNN with three layers and the width of each layer is 1,8,2, respectively. This results in 102 trainable parameters in total.

We randomly choose one example from the test set and plot its results in Figure 5. Both Figure 5(a) and Figure 5(b) show that the reconstruction calculated by LIR-DMD and FlowDMD are better than that by the Exact DMD and the difference of trajectories between LIR-DMD and FlowDMD is very small. Figure 5(c) and Figure 5(d) illustrate that the reconstruction error of FlowDMD is the smallest. In the first 30 time steps, LIR-DMD has a similar error to FlowDMD. The error of FlowDMD increases much more slowly than that of LIR-DMD for the following 30 time steps. We conclude that FlowDMD has better generalization ability than LIR-DMD.

We test FlowDMD, LIR-DMD and Exact DMD using 40 randomly generated examples and the results are depicted by Figure 6. We use the total relative \(L_{2}\) error to evaluate the reconstruction results of trajectories. For FlowDMD, the reconstruction error is the lowest among almost all of the test examples, and the average total relative \(L_{2}\) error is only 0.3%. Compared with LIR-DMD, FlowDMD has better generalization ability and learning ability of the Koopman invariant subspace.

Figure 5: Comparison of three methods for Example 4.1. The total relative \(L_{2}\) error of the Exact DMD, LIR-DMD, and FlowDMD are 0.2448, 0.0111 and 0.0018, respectively.
### Burgers' equation

The 1-D Burgers' equation [44] is given by

\[\begin{cases}\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}=\frac{0.01}{\pi}\frac{\partial^{2}u}{\partial x^{2}}\quad x\in(-1,1),t\in(0,1],\\ u(1,t)=u(-1,t)=0,\\ u(x,0)=-\xi*sin(\pi x),\end{cases} \tag{7}\]

where \(\xi\) is a random variable that satisfies a uniform distribution \(U(0.2,1.2)\). We use the finite element method with 30 equidistant grid points for the spatial discretization and the implicit Euler method with a step size of 0.01 for temporal discretization. We generate 100 samples of \(\xi\) for the initial state and compute the corresponding solutions. The examples are then divided into three parts, with proportions 60% for training, 20% for validation, and 20% for test. We test the performance of the Exact DMD, LIR-DMD, and FlowDMD. The rank of Exact DMD is 3 and the same rank is also used in LIR-DMD and FlowDMD to embed the Koopman linearity. The structure of the encoder network for LIR-DMD is \([30,40,50,40]\), and the decoder network is \([40,50,40,30]\) where the numbers in the brackets represent the width of each layer and we use RCFs to replace ACFs. This results in an invertible neural network of depth of 3 with one RCF block and one RCF. In each RCF, the width of each layer in FNN to parameterize the mapping \(t\) is 15, 40, 15, which results in 7530 parameters in FlowDMD, whereas LIR-DMD has 10650 parameters.

Figure 7 depicts that FlowDMD has the smallest absolute reconstruction error and total relative reconstruction error. Figure 8(a) and Figure 8(b) show that the reconstruction error of Exact DMD and LIR-DMD increase with time, but FlowDMD maintains in a very low level. Figure 9 summarizes the TRL2E of reconstruction on all test examples and depicts that the FlowDMD has the smallest error on almost all test examples, where the average TRL2E of FlowDMD is 1.5%. For some test examples, Exact DMD has the same TRL2E with FlowDMD, but for most test

Figure 6: Total relative \(L_{2}\) error in Example 4.1.
examples, FlowDMD performs better than Exact DMD. The TRL2E of LIR-DMD are bigger than FlowDMD over all the test examples and are slightly better than Exact DMD for some test examples.

### Allen-Cahn equation

The 1-D Allen-Cahn equation [44] is given by

\[\begin{cases}\dfrac{\partial u}{\partial t}-\gamma_{1}\dfrac{ \partial^{2}u}{\partial x^{2}}+\gamma_{2}\left(u^{3}-u\right)=0,x\in(-1,1),t \in(0,1],\\ u(0,x)=\xi*x^{2}\cos(2\pi x),\\ u(t,-1)=u(t,1),\end{cases} \tag{8}\]

where \(\gamma_{1}=0.0001\), \(\gamma_{2}=5\), and \(\xi\sim\mathcal{N}(-0.1,0.04)\). We use the finite element method with 20 equidistant grid points for the spatial discretization and the implicit Euler with a step size of 0.02 for the temporal discretization. Furthermore, we generate 100 samples of \(\xi\) and use _FEniCS_ to compute the numerical solutions. The data set is segmented according to a ratio of 60%, 20%, 20%, respectively to be used as

Figure 7: Comparison of three methods in Example 4.2. The total relative \(L_{2}\) errors for exact DMD, LIR-DMD, and FlowDMD are 0.08, 0.119, and 0.017, respectively.
the training set, the validation set, and the test set. The structure of the encoder network for LIR-DMD is [20, 30, 40, 30] and the decoder network is [30, 40, 30, 20], where the numbers in the bracket indicate the width of each layer. This results in 6190 parameters for LIR-DMD. For FlowDMD, we also use RCFs to replace the ACFs. The neural network for FlowDMD consists of one RCF block and one RCF, which results in a network with depth \(L=3\). In each RCF, the width of each layer of the FNN to parameterize \(t\) is 10, 20, 10. Finally, we obtain 2580 parameters for FlowDMD. The rank of Exact DMD is 3, and the same rank is also used in LIR-DMD and FlowDMD to embed the Koopman linearity.

Figure 10 clearly shows that FlowDMD can reconstruct the original state most accurately. It reveals that the absolute error of both exact DMD and LIR-DMD increase over time, but FlowDMD can maintain the error in a low level all the time. Numerical results show that FlowDMD is more robust and generalizes better than Exact DMD and LIR-DMD. The error of the state reconstruction for three methods are given in Figure 11. At the beginning time, FlowDMD has the biggest relative error because the norm of the true state variables is too small, which leads to a

Figure 8: Relative error of three methods for Example 4.2.

Figure 9: Total relative \(L_{2}\) error in Example 4.2.
large relative error. As time evolves, the error of FlowDMD reaches the lowest level among all three methods. In Figure 12, we use the test data set to evaluate the generalization ability. The FlowDMD has almost the smallest total relative \(L_{2}\) error in most examples and the average of the total relative \(L_{2}\) error is 9%. It also shows that the fluctuation of error for FlowDMD is smaller than that of LIR-DMD, which demonstrates that FlowDMD has a better generalization ability and is more robust than LIR-DMD.

### Sensitivity study

Here, we study the sensitivity of FlowDMD systematically with respect to the following four aspects:

1. The neural network initialization.
2. The hyperparameter \(\alpha\) in the loss function.
3. The structure of neural networks.
4. The rank \(r\) used by DMD in Algorithm 1.

We study the sensitivity of FlowDMD using the Allen-Cahn equation in Section 4.3.

Figure 10: Comparison of three methods in Example 4.3. The total relative \(L_{2}\) error for exact DMD, LIR-DMD, and FlowDMD are 0.6129, 0.4038, and 0.0725, respectively.
#### 4.4.1 Sensitivity with respect to the neural network initialization

In order to quantify the sensitivity of FlowDMD with respect to the initialization, we consider the same data set with Section 4.3. Simultaneously, we fix the structure for FlowDMD to include only one RCF block and one RCF. Each RCF has a FNN to parameterize \(t\) where the width of each layer is \(10,20,10\). Moreover, all FNNs use the rectified linear unit (ReLU) as activation functions. We use \(15\) random seeds to initialize models and train all the models with the same setting. In Figure 13, we report the total relative \(L_{2}\) error between the reconstructed states and the"true" states. Evidently, the TRL2E remains stable for different initializations of neural networks, as demonstrated by the consistent results obtained within the following interval,

\[[\mu_{TRL2E}-\sigma_{TRL2E},\mu_{TRL2E}+\sigma_{TRL2E}]=[6.5\times 10^{-2}-1.6 \times 10^{-2},6.5\times 10^{-2}+1.6\times 10^{-2}]\]

#### 4.4.2 Sensitivity with respect to \(\alpha\)

We utilize the same training set with Section 4.3 and select \(\alpha\) from the list \([0.01,0.1,1,10,100]\). As shown in Table 1, the different weights \(\alpha\) in the loss function

Figure 11: Relative error of three methods for Example 4.3.

Figure 12: Total relative \(L_{2}\) error in Example 4.3.
have little influence on the final results. We observe that the error is minimized when \(\alpha=10\), which suggests the use of an adaptive weight selection algorithm. The gradient flow provided by the neural tangent kernel [45] can be employed to adjust the weight \(\alpha\) and accelerate the training process, and we leave this for our future work.

#### 4.4.3 Sensitivity with respect to the structure of neural networks

We study the impact of the number of RCFs and the number of neurons in the FNN to parameterize the mapping \(t\) on the performance of the FlowDMD. Specifically, the sensitivity of FlowDMD is being quantified with respect to two parameters: the number of RCFs (\(N_{f}\)) and the number of neurons (\(N_{n}\)) in the middle layer of the FNN. Here, the FNN used to parameterize \(t\) is restricted to a three layer structure of \([10,N_{n},10]\). The results are summarized in Table 2, which indicate that the reconstruction of FlowDMD has little to do with its structure while adding more neurons or more RCFs will not improve the final results to a big extent.

#### 4.4.4 Sensitivity with respect to the rank of DMD

As we increase the rank \(r\) used for the DMD computations in Algorithm 1, we include more physical information, but the computation time also increases. In this study, we investigate how the DMD rank affects the model and its reconstruction. The results in Table 3 show that as we increase the rank \(r\), the corresponding error decreases rapidly.

\begin{table}
\begin{tabular}{c c c c c c} \hline \(\alpha\) & 0.01 & 0.1 & 1 & 10 & 100 \\ \hline TRL2E & 6.2e-02 & 6.8e-02 & 8.2e-02 & 3.2e-02 & 6.9e-02 \\ \hline \end{tabular}
\end{table}
Table 1: Total relative \(L_{2}\) error for different \(\alpha\).

Figure 13: Total relative \(L_{2}\) error for different neural network initializations.
## 5 Conclusion

In this paper, we propose a coupling flow invertible neural network approach to learn both the observable functions and reconstruction functions for the Koopman operator learning. Our method generate more accurate Koopman embedding model and better approximations of the Koopman operator than state-of-the-art methods. Our FlowDMD is structurally invertible, which simplifies the loss function and improves the accuracy of the state reconstruction. Numerical experiments show that our approach is more accurate, efficient, and interpretable than the state-of-the-art methods.

## Acknowledgments

The authors would like to thank Mengnan Li and Lijian Jiang for sharing their code.
