Stochastic Adversarial Koopman Model for
Dynamical Systems

Kaushik Balakrishnan
Ford Greenfield Labs, Palo Alto, CA
kbalak18ford. om

Devesh Upadhyay
Ford Research, Dearborn, MI
dupadhyaford. om

Abstract
We present a general model for non-linear dynamical systems that is based on
the Koopman theory and combines an Autoencoder with a Generative Adversarial
Network (GAN). We show that coupling a GAN as an additional loss term during
training increases the accuracy of the model predictions, and demonstrate it on
test problems in fluid dynamics, combustion, and Li-ion battery systems.

1

Introduction

Many engineering and scientific problems are solved using finite difference or finite volume methods, which are time consuming. Thus, recasting the problem in a lower-dimensional state space
and advancing the solution in time in this reduced dimensional space is preferred, and one class of
methods in this setting are based on the Koopman theory [1]. Several research works have renewedinterest in Koopman models in recent years, as evidenced by [2, 3, 4, 5, 6, 7, 8, 9]. We extend thesolution procedure using a stochastic latent embedding space, and the Koopman operator applies inthis space. Specifically, the Koopman operator applies on the probability distribution of the latentembedding and advances it linearly in time, thereby learning a distribution of possible outcomes. If
z is the latent embedding and p(z) is its probability distribution, past studies [5, 6, 8, 9] apply the
Koopman operator directly on z; in this study, we apply the Koopman operator on p(z) instead. In
addition, a GAN [10] discriminator is also used during training as an additional adversarial loss andthis is found to robustify the predictions. A series of test problems are considered for the analysis.

2

The Koopman operator for dynamical systems

We consider dynamical systems of the form xt+1 = F (xt ), where F is a non-linear function. In
the Koopman model [1], the state vector xt is mapped on to a Hilbert space g(xt ) and the evolutionof the system dynamics in time is linear, where the Koopman operator K advances the system as:
Kg (xt ) = g (F (xt )) = g (xt+1 ). The system is then projected back to the physical state vector
space using an inverse function g −1 [11, 12, 13]: g −1 (Kg (xt )) = xt+1 . Recently, autoencoders[14] are used for this task with the encoder learning the mapping g and the decoder g −1 directlyfrom data snapshots [5, 6, 8, 9]. During training, randomly sampled sequences of length nS areconsidered from the data corpus: xt · · · xt+nS , and used to construct the vectors:
i
h
pred
pred
pred
= xpredX = [xt , xt+1 , · · · , xt+nS −1 ] ; X+1 = [xt+1 , xt+2 , · · · , xt+nS ] ; X+1
t+1 , xt+2 , · · · , xt+nS . (1)predHere, X+1
is the vector of the model’s predictions for the states of the system at subsequent time
predsteps. We take xt as input and output the sequence X+1
using the Koopman dynamics recursivelynS times.

When the codomain of the encoding function g is of finite dimension, the Koopman operator K
reduces to a Koopman matrix K [5], and thus the Koopman operation reduces to that of a matrix
Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

multiplication. We will use zt to represent the embedding at time instant t, i.e., zt = g(xt ) andxt = g −1 (zt ). [9] used a residual Koopman approach where the system dynamics instead of being
represented as zt+1 = Kzt = Kzt , was represented as Kzt = zt+1 = zt + Kzt (note that K is the
Koopman operator, whereas K is the corresponding Koopman matrix). Here, the Koopman model
learns the residual change required to advance the system in time in the latent embedding space.
With this change, the recursive prediction of the future states of the system in the embedding space
is:


pred
(2)
Z+1
= K zt , K2 zt , · · · , KnS zt ,where Kj represents the application of the Koopman opertaor K j times, and nS is the desiredsequence length.
In [9], a Generative Adversarial Network (GAN) [10] was also coupled to the Deep Koopman modelby including an additional loss to train the “generator” networks. It has been demonstrated in computer vision applications that coupling a GAN Discriminator with an autoencoder can improve thequality of samples output from the autoencoder [15]. This is because the feature representations
learned by the GAN discriminator also provide additional learning signals for the primary neuralnetworks, and this improves the output quality [15]. In NLP applications too, the use of GAN has
provided regularization and improved the quality of model predictions [16]. These observations also
extend to dynamical system modeling, as we will demonstrate in this paper.
pred
In this setting, GAN losses are then constructed using X, X+1 and X+1. Two concatenated pairs
pred
are used: the “real” (X, X+1 ) and “fake” (X, X+1). These pairs are fed into the GAN discriminator
which outputs a single real value D(·) from which one can construct the GAN objective, following
the Wasserstein GAN [17] approach due to its robustness against mode collapse:

LGAN objective =E[D(x)] −x∈(X,X+1 )

E

[D(e
x)] .

(3)

pred
x
e∈(X,X+1)

To obtain the Koopman matrix K, an auxiliary neural network is used to output K as a function of
xt , similar to [5, 9]. We consider a tridiagonal structure for the Koopman matrix, since the numberof non-zero entries is one order of magnitude fewer than a full Koopman matrix. That is, if M isthe dimension of the latent embedding, then a full Koopman matrix will have M × M non-zero
entries whereas a tridiagonal Koopman matrix will only have 3M − 2 non-zero entries. See [18]for more discussions on the choice of using a tridiagonal form for the Koopman matrix, albeit theirtridiagonal form had additional constraints not considered here.
The latent embedding zt is modeled as a Gaussian random variable; thus, zt involves two components µzt and σtz , each of which ∈ RM . The auxiliary network takes µzt and σtz as input and outputs
two Koopman matrices Kµ and Kσ which are used to construct the distribution for the next time
step:z
µzt+1 = µzt + Kµ µzt ; ln σt+1
= ln σtz + Kσ ln σtz .

(4)z), which
Once the distributions for the future states are obtained, we sample zt+1 = N (µzt+1 , σt+1
is then passed on to the decoder. We call the model as the Stochastic Adversarial Koopman (SAK)model and a schematic is shown in Fig. 1. A total of four neural networks are used: the encoder, decoder, auxiliary network and the discriminator and their architectures are summarized in
the Supplementary Materials. The loss functions used to train the networks are also elaborated in
the Supplementary Materials. The main contributions of this paper are summarized as follows:• Model zt as a stochastic random variable, i.e., a Gaussian; thus, the embedding is more
precisely p(zt ); we have thus defined a Koopman operator in stochastic embedding space• Use Maximum Mean Discrepancy (MMD) loss for p(zt ), similar to a Wasserstein Autoencoder [19]; this we believe is new in a Koopman setting• Use a GAN Discriminator as an additional adversarial loss term; this acts as a regularizer
and improves the Koopman model predictions; this is novel in the Koopman family ofalgorithms
2

Figure 1: Schematic of the Stochastic Adversarial Koopman Model showing the different neuralnetworks and the connections.

3

Experiments

We will now demonstrate the robustness of the Stochastic Adversarial Koopman (SAK) model ona variety of engineering test problems: (1) von Karman vortex shedding behind a cylinder [20];(2) Flame ball-vortex interaction [21]; and (3) Doyle-Fuller-Newman (DFN) model for Lithium ion
batteries [22, 23]. See the Supplementary Materials for more details.
von Karman vortex shedding
The hyperparameter λGAN (see Supplementary Materials) decides the tradeoff between the GAN
loss in comparison with the other supervised loss terms; λGAN = 0.01 was found to be the optimalvalue for training the SAK model on the von Karman vortex shedding data. The test time model
predictions and ground truth of the velocity components u, v and pressure p at a time instant isshown in Fig. 2 and the model predictions are very similar to the ground truth values, with very low
error. To better understand the role of the GAN loss, we repeat the training from scratch with the
GAN loss turned off, i.e., λGAN = 0.0. In Fig. 3 (a), we compare the test time prediction errors withand without the GAN loss term, and λGAN = 0.01 has lower prediction errors for most time steps.
We also investigated the magnitudes of Kµ and Kσ after training (not presented for brevity), andobserved the entries along the diagonal to be preponderant for both Kµ and Kσ . Although a few offdiagonal entries also had relatively large magnitude values, no specific structure was evident otherthan the diagonal preponderance. In [9] a full Koopman matrix (M ×M ) was output by the auxiliary
network, but in the tridiagonal version considered in this study, the Koopman matrices Kµ and Kσhave 3M − 2 non-zero entries each, thus the number of non-zero entries in the Koopman matrices
are now fewer: O(M 2 ) → O(M ). We consider another training case with a full Koopman matrix,
and compare the mean absolute errors at each time step for the predictions using both formulationsin Fig. 2 (b). The errors are nearly similar in magnitude for both formulations despite having oneorder of magnitude difference in the number of non-zero values in Kµ and Kσ . The tridiagonal
Koopman demonstrated here can be seen as an alternative to the Jordan/diagonal block structures
for K used in [5, 7].
Flame ball-vortex interaction
The results are presented in Fig. 4 and as evident the SAK model is able to accurately capture theshape of the flame ball as the vortex distorts it counter-clockwise. At late time the surface area of theflame increases, which further consumes more of the pre-mixed fuel-oxidizer mixture. The errors(not shown) are very small compared to the magnitudes of the variables, demonstrating the efficacyof the SAK model.
Conditional SAK model on the DFN data
For this problem, we consider a conditional Koopman model where the Koopman matrices are conditioned on the applied current density Iapp (in A/m2 ), which is fed as input to the auxiliary neuralnetwork. The training set consists of the battery discharge states (each state comprises of 6 variables)in time for Iapp in the range 10-35 A/m2 . nS = 16 at the beginning of the training and is increasedby 2.5% every 20k iteration steps until nS = 256, after which it is held fixed. We consider a totalof 1 million training iterations, as this problem has more variety in the training data, necessitating
3

Figure 2: Predictions of u, v and p for the von Karman vortex shedding problem behind a cylinder
with the SAK model. SAK model predictions: “pred”; ground truth: “gt”.

0.0

mean abs error

0.007

0.01

0.006

0.005

0.004

0.003

0

50

(a)

100

150

200

time step

250

300

(b)

Figure 3: (a) Prediction error for each time snapshot for the von Karman vortex problem with λGAN
= 0.0 and 0.01; (b) mean absolute error for every time snapshot w.r.t. the ground truth for the fulland tridiagonal Koopman matrices.
longer training. After the conditional SAK model is trained on the training set, we use it to makepredictions for Iapp values in the test set. The model predictions at time instants: 258, 600, 1100 and2000 seconds during the battery discharge for Iapp = 31 A/m2 are presented in Fig. 5 (in blue), alongwith the ground truth values (in orange). The model predictions are very accurate, although we doobserve some slight differences for the variables “C Sol Surf” and “j main.” Thus, the conditionalmodel is able to learn the dynamics of battery discharge.

4

Discussion

A novel model based on the Koopman family of algorithms to train dynamical systems is developed,
using a Gaussian stochastic embedding of an autoencoder in the latent space, which is advanced intime. During training, the model also couples a GAN discriminator for an adversarial loss term thatis found to improve the accuracy, provided an optimal value for the trade-off parameter λGAN isused. This optimal value varies for different problems and requires experimentation, which needs tobe addressed in future studies. We also extend the model to a conditional Koopman setting where
additional input parameters are supplied to the auxiliary network so that the Koopman matrices areconditioned on these inputs. The SAK model predictions are robust and can be used for predicting
the dynamics of many real-world systems. While the Koopman family of models take a few days
to train, the main advantage of such Reduced Order Models (ROM) is in the design of engineeringsystems where such models can be trained once from data, and used to make inferences in a matterof ∼ 100 milliseconds, which will greatly help engineers in the design iteration process. Our experiments demonstrate inference speed-ups on the order of 10-100X for the problems considered here.
The speedup advantages will be more preponderant for 3D problems, which is of interest in a future
study.

4

(a)

(b)

(c)

(d)
Figure 4: Predictions of T and Y for the flame ball-vortex interaction problem with the SAK modelat time step number: (a) 22, (b) 53, (c) 78 and (d) 112. SAK model prediction: “pred"; ground truth:
“gt".

(a)

(b)

(c)

(d)

Figure 5: Predictions of the DFN battery model variables for Iapp = 31 A/m2 with the conditional
stochastic adversarial Koopman model at time instants: (a) 258, (b) 600, (c) 1101 and (d) 2000
seconds during discharge. Ground truth: orange; SAK model predictions: blue.

References[1] B. O. Koopman. Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences USA,
17:315–318, 1931.
[2] I. Mezic and A. Banaszuk. Comparison of systems with complex behavior. Physica D, 197:101–133, 2004.
[3] H. Arbabi and I. Mezić. Study of dynamics in post-transient flows using koopman mode decomposition. Physical Review Fluids,
2:124402, 2017.
[4] E. Yeung, S. Kundu, and N. Hodas. Learning deep neural network representations for koopman operators of nonlinear dynamicalsystems. arXiv:1708.06850 [cs.LG], 2017.

5

[5] B. Lusch, J. N. Kutz, and S. L. Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature Communications,9(1):4950, 2018.
[6] N. Takeishi, Y. Kawahara, and T. Yairi. Learning koopman invariant subspaces for dynamic mode decomposition. Advances in Neural
Information Processing Systems (NIPS), 2017.
[7] A. Salova, J. Emenheiser, A. Rupe, J. P. Crutchfield, and R. M. DSouza. Koopman operator and its approximations for systems withsymmetries. Chaos, 29:093128, 2019.
[8] J. Morton, F. D. Witherden, A. Jameson, and M. J. Kochenderfer. Deep dynamical modeling and control of unsteady fluid flows.
Advances in Neural Information Processing Systems (NIPS), 2018.
[9] K. Balakrishnan and D. Upadhyay. Deep adversarial koopman model for reaction-diffusion systems. arXiv:2006.05547 [cs.CE], 2020.
[10] I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial
nets. Advances in Neural Information Processing Systems (NIPS), 2014.
[11] J. L. Proctor, S. L. Brunton, and J. N. Kutz. Dynamic mode decomposition with control. arXiv:1409.6358 [math.OC], 2014.
[12] J. N. Kutz, S. L. Brunton, B. W. Brunton, and J. L. Proctor. Dynamic mode decomposition: Data-driven modeling of complex systems.
Society for Industrial and Applied Mathematics, 2016.
[13] E. Kaiser, J. N. Kutz, and S. L. Brunton. Data-driven discovery of koopman eigenfunctions for control. arXiv:1707.01146 [math.OC],
2017.
[14] G. E. Hinton and R. R. Salakhutdinov. Reducing the dimensionality of data with neural networks. Science, 313:504–507, 7 2006.
[15] A. B. L. Larsen, S. K. Sonderby, H. Larochelle, and O. Winther. Autoencoding beyond pixels using a learned similarity metric.
arXiv:1512.09300 [cs.LG], 2016.
[16] A. Lamb, A. Goyal, Y. Zhang, S. Zhang, A. Courville, and Y. Bengio. Professor forcing: A new algorithm for training recurrent
networks. arXiv:1610.09038 [stat.ML], 2016.
[17] I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville. Improved training of wasserstein gans. Advances in Neural
Information Processing Systems (NIPS), 2017.
[18] S. Pan and K. Duraisamy. Physics-informed probabilistic learning of linear embeddings of nonlinear dynamics with guaranteed stability.
SIAM Journal on Applied Dynamical Systems, 19(1):480–509, 2020.
[19] I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schoelkopf. Wasserstein auto-encoders. arXiv:1711.01558 [stat.ML], 2018.
[20] T. von Karman. Aerodynamics. McGraw-Hill, 1963.
[21] O. Roussel and K. Schneider. An adaptive multiresolution method for combustion problems: application to flame ball–vortex interaction. Computers and Fluids, 34(7):817–831, 2005.
[22] M. Doyle, T. F. Fuller, and J. Newman. Modeling of galvanostatic charge and discharge of the lithium/polymer/insertion cell. Journalof The Electrochemical Society, 140(6):1526, 1993.
[23] T. F. Fuller, M. Doyle, and J. Newman. Simulation and optimization of the dual lithium ion insertion cell. Journal of The Electrochemical Society, 141(1):1, 1994.
[24] M. Doyle, T. F. Fuller, and J. Newman.
Dualfoil
http://www.cchem.berkeley.edu/jsngrp/fortran_files/fortran.html.

code

v5.1

(universityofcalifornia,

berkeley).

[25] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. IEEE Conference on Computer Vision and Pattern
Recognition (CVPR), 2016.
[26] D. P. Kingma and J. Ba. Adam: a method for stochastic optimization. 3rd International Conference for Learning Representations(ICLR), 2015.

6

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? blue[Yes](b) Did you describe the limitations of your work? blue[Yes] See the Discussion section(c) Did you discuss any potential negative societal impacts of your work? gray[N/A](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? blue[Yes] This work does not involve ethical issues
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? gray[N/A](b) Did you include complete proofs of all theoretical results? gray[N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? blue[Yes] The code
will be posted on github before the conference(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? blue[Yes] Supplemental Materials specifies this information(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? orange[No] Only errors w.r.t. the model parameter λGAN ispresented in Fig. 3(d) Did you include the total amount of compute and the type of resources used (e.g.,
type of GPUs, internal cluster, or cloud provider)? blue[Yes] This is mentioned in the
Supplementary Materials
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? blue[Yes] We cite [8, 5](b) Did you mention the license of the assets? gray[N/A] The model is in public domain(c) Did you include any new assets either in the supplemental material or as a URL?
blue[Yes] The models are discussed in the supplemental material(d) Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? gray[N/A] All the data is generated by us using open-source
codes and we cite them(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? gray[N/A] Data does not contain any offensive
content
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? gray[N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? gray[N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? gray[N/A]

Supplementary Materials
1. von Karman vortex shedding behind a cylinder
The classical von Karman vortex shedding behind a cylinder [20] is widely used for validating many
Computational Fluid Dynamics (CFD) codes and is of relevance to both the aerospace and automotive industries for aerodynamic drag reduction. The Navier-Stokes equations of fluid dynamics are
solved in a 2D domain for flow past a cylinder from left to right at a Reynolds number Re = 150. The
domain size is 22×4 and the velocity of the flow at the input is 1 (all units are non-dimensionalized).
The domain is discretized using a 660×120 grid and the CFD time step ∆t = 0.01. The open source
code: https://github.com/dorchard/navier is used for solving the Navier-Stokes equations and generating the data corpus. The flow-field comprises of three variables: u, v, p, where u and v are the7x and y velocity components and p is the pressure. For training the SAK model, only the region
immediately behind the cylinder where the vortex shedding is preponderant, is considered; specifically, we consider 360 cells in the horizontal immediately behind the cylinder and the entire 120
cells in the vertical directions. Every 4-th CFD solution snapshot is saved, and thus ∆t = 0.04 for the
Koopman analysis. We consider 400 time snapshots for training the SAK model, which corresponds
to approximately 3.5 cycles of vortex shedding. The SAK model training takes approximately 3days on a single GPU.
2. Flame ball-vortex interaction
Flame-vortex interaction is a classical engineering problem in combustion engines, where vortices
can distort a flame, thereby enhancing the mixing of fuel with the oxidizer. We consider a simple
2D model [21] for dimensionless temperature T and mass fraction/concentration of premixed gas Y
and this consists of a system of 2 PDEs. For brevity, these equations are not presented and can befound in [21]. Outflow boundary conditions are used at all four boundaries for the CFD analysis.
The initial conditions are are not presented here and can be found in [21]. We use a 2nd order central
scheme in space and a 3rd order Runge-Kutta scheme in time to solve the system of equations. For
the CFD analysis, the domain size is 20×20 and is discretized using a 512×512 mesh for the finite
difference method, with a time step ∆t = 10−5 for 140k time steps. The solution at every other nodeis saved once every 1000 time steps, and so the data corpus for training the SAK model is of size
256×256×2 (the 2 is for two variables: T, Y ) for 141 snapshots. The SAK model training takes
approximately 3 days on a single GPU.
3. Doyle-Fuller-Newman Li-ion battery model
The Doyle-Fuller-Newman (DFN) model [22, 23] consists of a system of PDEs in 1D to solve for the
Li-ion concentrations and potential in the electrode particles and the electrolyte of a Li-ion battery.
We use the Dualfoil code v5.1 open sourced from the DFN authors [24] for generating the data
corpus. Specifically, we consider a fully charged Li-ion battery and vary the applied current density,
Iapp , in the range 10-35 A/m2 , and make predictions of the battery potentials in the electrode andelectrolyte, the Li-ion concentrations, and the current as the batetry discharges in time. The crosssection of the battery is discretized using 201 nodes and the Dualfoil code outputs 6 variables:
concentration of the Li-ions in the electrolyte (in mol/m3 ), concentration of Li-ions at the surfaceof the solid particles (non-dimensionalized), potentials in the liquid and solid phases (both in Volts),
liquid phase current density and current density “j main" (both in A/m2 ). These variables are saved
every 1 second time interval during the discharge of the battery until the voltage drops to 3 V, and thetotal number of time snapshots varies depending on Iapp ; for instance, the total number of snapshots
is 8454 for Iapp = 10 A/m2 and is 2197 for Iapp = 35 A/m2 , with the other cases in between. Moredetails on the battery modeling system and solution procedures can be found in the Dualfoil manual[24]. We train the SAK model with the Koopman matrix conditioned on Iapp by supplying this as
an input to the auxiliary network. The SAK model is trained on a training set comprising of 18
different Iapp values in the range 10 ≤ Iapp ≤ 35 A/m2 and used to make predictions of the stateof the battery at different time instants for a value of Iapp not used in the training. The SAK modeltraining takes approximately 7 days on a single GPU.

Neural Network Architectures
The neural network architecture used in the analysis is summarized here. Note that we have a total of
4 neural networks: Encoder, Decoder, auxiliary network AUX (to obtain K matrices) and the GAN
discriminator DISC. We will use several different deep learning building blocks: batch normalization(BN), Dropout (Dropout), convolutional (conv) and deconvolutional (dconv) operators, and the Relu
(Relu) activation function. The notation conv(k,f,S,s) is used for a convolutional layer with kernel
size k, f filters, same padding (identified by S) and a stride of s. The notation Dense(n) is used
to refer to a fully connected dense layer with n neurons. We first define a bottleneck layer for
Encoder with Nf filters as input, BottleNecke (Nf ), comprising of the following in the same order:
BottleNecke (Nf ) = BN → Relu → conv(1,Nf /2,S,1) → BN → Relu → conv(3,Nf /2,S,1) → BN
→ Relu → conv(1,Nf ,S,1). Encoder consists of 5 layers of convolutional operations supplemented
with bottleneck layers added residually, similar to Resnet [25]. For ease of notation, we will refer to
Encoder’s residual block as RESe (Nf ) = conv(3,Nf ,S,2) + BottleNecke (Nf ). Specifically, Encoder
consists of 5 residual layers in succession: RESe (64) → RESe (128) → RESe (256) → RESe (512) →RESe (512), followed by a Relu and a flattening operation. This is then fed into two Dense(M ) layers
8to output the µ and σ of the embedding vector, where the dimension M = 64 is used throughout thisstudy.
For the Decoder we define a similar bottleneck layer with Nf filters, albeit this time using deconvolutional operations: BottleNeckd (Nf ) = BN → Relu → dconv(1,Nf /2,S,1) → BN → Relu →dconv(3,Nf /2,S,1) → BN → Relu → dconv(1,Nf ,S,1). For Decoder, we first add the input to thebottleneck layer akin to Resnet [25], which is then passed though a deconvolutional layer, like so:
RESd (Nf ) = input + BottleNeckd (Nf ) → dconv(3,Nf ,S,2). Decoder starts with Dense(·) with thenumber of neurons used being the same as the dimension of the Encoder’s flattened output. This isreshaped as appropriate and is followed by 5 layers of the Decoder’s residual blocks in succession:
RESd (512) → RESd (256) → RESd (128) → RESd (64) → RESd (nout ). Here, nout is the number of
output channels in the data. Note that some problems are 1D whereas others are 2D, and so theappropriate conv and dconv API calls are used.
For the AUX network, we define a fully connected layer with N neurons as FC(N ) = Dense(N ) →Relu → Dropout. For Dropout, we set the probability of keeping the activations to 0.8 at training,and 1.0 at testing. AUX network consists of 4 fully connected layers: FC(128) → FC(256) →FC(512) → Dense(2ndec ) (the 2 factor arises since we consider two Koopman matrices Kµ andKσ ). For the full Koopman matrices, ndec = M 2 , whereas for the tridiagonal Koopman matrices,
ndec = 3M -2.
For DISC, we will use the Leaky Relu activation function, denoted as LRelu, with a slope of 0.2
in the negative side. We define a block BDISC (Nf ) as conv(5,Nf ,S,2) → BN → LRelu. DISC isthen constructed as: conv(5,64,S,2) → LRelu → BDISC (128) → BDISC (256) → BDISC (512). The
output is then reshaped and passed to a Dense(1) without any activation function to represent the
Wasserstein distance.
Adam [26] optimizer is used to train the neural networks with a learning rate of 1×10−5 . The total
number of iterations used for training varies for the different problems, but is usually in the orderof 100k-500k, where at each iteration step one sequence of nS contiguous snapshots are randomly
sampled from the data corpus and used to train the networks.
Loss terms
The encoder and decoder are represented as g(·) and g −1 (·), respectively. We use the mean squared
error (MSE) and the maximum mean discrepancy (MMD) [19] to construct different loss terms: (1)reconstruction loss Lrecon , (2) prediction loss Lpred , (3) code loss Lcode , (4) gradient loss Lgrad , (5)L2 regularization loss Lreg , (6) GAN loss LGAN , and (7) discriminator loss Ldisc . These different
losses are summarized below:

Lpred =

Lrecon =k xt − g −1 g (xt ) kMSE

(5)

nS
1 X
k xt+m − g −1 (Km g (xt )) kMSE
nS m=1(6)

Lcode = MMD[g (xt+m ) , Km g (xt )]

(7)

Lgrad
j

=

nS


1 X
k ∇j xt+m − g −1 (Km g (xt )) kMSE , j=1, 2, 4
nS m=1Lgrad=

λ1 Lgrad+ λ2 Lgrad+ λ4 Lgrad
1
2
4

(8)

X

wi2

(9)

[D(e
x)]

(10)

Lreg = λreg
LGAN =

E

predxe∈(X,X+1)

9

Ldisc =

E

[D(e
x)] −

E

[D(x)](11)x∈(X,X+1 )predxe∈(X,X+1)

Note that the losses Lrecon , Lpred , and Lcode were also considered in [5, 9], but MSE was considered
for all these terms. In this study, since the latent embedding is stochastic, we use the MMD loss [19]:X
X  pred pred 
1
1
2 X 
Lcode =f (zl , zj ) +f zl , zj
f zl , zjpred , (12)− 2nS (nS − 1)nS (nS − 1)nSl,j,l6=jl,j,l6=jl,jwhere zl is sampled from N (g (xt+l )) and zlpred from N (Kl g (xt )). Note that we essentially havesamples from two distributions g (xt+l ) and Kl g (xt ) that need to be synchronized and so the Maximum Mean Discrepancy (MMD) [19] is one such cost function. Note that the function g takes xt
as input and outputs µzt and σtz . And g −1 takes zt (or zt+1 ) as input to output xt (or xt+1 ). f is theinverse multiquadratics kernel f (x, y) = C/(C+ k x − y k22 ) with C being a constant [19]. In theabove equations, Km denotes the application of the Koopman operator m times. Gradient losses arealso used to improve the overall quality of the output.
The encoder, decoder and the auxiliary network—which we can refer to as the “generator” in GAN
parlance—are trained jointly using the loss function:
Ltotal = Lrecon + Lpred + λcode Lcode + λgrad Lgrad + λreg Lreg + λGAN LGAN ,(13)

and the GAN discriminator is trained using Ldisc . We use λcode = 100, λgrad = 1 and λreg = 10−3
for all the cases considered in this paper. λGAN varies for each test case and is obtained fromexperimentation; it is typically in the range 0–0.1 for best results. The discriminator is trained byminimizing Ldisc , along with an additional gradient penalty loss term similar to WGAN-GP [17].
We alternate between training the generator one step followed by the discriminator for one step, ateach training iteration. At test time, the discriminator is not used.

10