Model Reduction with Memory and the Machine Learning of Dynamical Systems

The Mori-Zwanzig formalism

The starting point of the M-Z formalism is to divide all the variables into resolved and unresolved ones. The basic idea is to project functions of all variables into the space of functions of only the resolved variables. The original dynamical system then becomes an equation for the resolved variables with memory and noise. This equation is called the Generalized Langevin Equation (GLE). Below we first demonstrate the idea of the M-Z formalism using a simple linear example. We then use the Kuramoto-Sivashinsky equation as an example to show how M-Z formalism help us develop a reduced model.

with initial value $x(0)=x_{0}$, $y(0)=y_{0}$. We take $x$ as the resolved variable, $y$ as the unresolved variable, and want to develop a reduced system for $x$. To achieve this, we can first fix $x$ in (8) and solve for $y$, and then insert the result into (7). In this way we get

There are three terms at the right hand side of (9). The first term $A_{11}x$ is a Markovian term of $x$; the second term $A_{12}\int_{0}^{t}e^{A_{22}(t-s)}A_{21}x(s)ds$ is a memory term. The third term $A_{12}e^{A_{22}t}y(0)$ contains the initial value of $y$, this term will be viewed as a noise term. Hence, by eliminating $y$ from the original linear system, we obtain an equation for $x$ with memory and noise.

Similar deduction can be applied to non-linear ODE systems. Assume we have a system

In (11), $\hat{\phi}(x,t)$ denotes $\hat{\phi}$ at time $t$ with initial value $x$. Although (11) is more complicated than (9), the essential ingredients are similar. The M-Z formalism tells us that model reduction leads to memory effects, and inspired us to pursuing the application of RNNs as a tool for performing efficient yet rigorous model reduction.

2 Application to the K-S equation

Going back to the K-S equation (1), let $\hat{u}$ and $\hat{G}$ be Fourier transforms of solution $u$ and filter $G$, respectively. Then, by (3), we have $\hat{\bar{u}}_{k}=\hat{G}_{k}\hat{u}_{k}$ for all frequency $k$. Here, we take a spectrally sharp filter $G$ which satisfies

for a certain positive integer $K$. Then, the Fourier transform of $\bar{u}$ is a truncation of the Fourier transform of $u$, and our resolved variables are those $\hat{u}_{k}$ with $|k|\leq K$. Writing the K-S equation in Fourier space, and put all terms with unresolved variables to the right hand side of the equation, we can get

Applying Fourier transform to (4), and compare with (12), we have

for $|k|\leq K$. Using the M-Z theory, the sub-grid stress can be expressed as a memory term and a noise term. By Galilean invariance, the sub-grid stress can be expressed as a function of the history of the resolved strain, $\partial\bar{u}/\partial x$, when the noise effect is negligible.

Recurrent Neural Networks and LSTM

In this section, we briefly introduce recurrent neural networks and long short-term memory networks.

RNNs are neural networks with recurrent connections, designed to deal with time series. Usually, a recurrent connection links some units in a neural network to themselves. This connection means that the values of these units depend on their values at the previous time step. For a simple example, consider a one-layer RNN with time series $\{x_{t}\}$ as input. Let $h_{t}$ and $y_{t}$ be the hidden values and output values at time step $t$, respectively. Then, the simplest RNN model is given by

where $U$, $V$, $W$, $b$, and $d$ are matrices or vectors of trainable parameters, and $\sigma(\cdot)$ is a nonlinear activation function. To deal with complicated time series, sometimes the output of an RNN is fed to another RNN to form a multi-layer RNN.

Note that in RNNs the same group of parameters are shared by all time steps. In this sense, RNNs are stationary models. Among other things, this allows RNNs to learn information in long time series without inducing too many parameters. The training of RNNs is similar to the training of regular feed-forward neural networks, except that we have to consider gradient through time direction, such as the gradient of $y_{t}$ with respect to $h_{s}$ for $s<t$. Usually RNNs are trained by the so-called back-propagation through time (BPTT) method.

2 Long short-term memory networks

Theoretically, RNNs is capable of learning long-term memory effects in the time series. However, in practice it is hard for RNN to catch such dependencies, because of the exploding or shrinking gradient effects , . The Long Short-Term Memory (LSTM) network is designed to solve this problem. Proposed by Hochreiter et al. , the LSTM introduces a new group of hidden units called states, and uses gates to control the information flow through the states. Since the updating rule of the states is incremental instead of compositional as in RNN, the gradients are less likely to explode or shrink. The computational rule of an LSTM cell is

where $f_{t}$, $i_{t}$, $o_{t}$ are called the forget gate, input gate, and output gate, respectively, $S_{t}$ is the state, and $h_{t}$ is the hidden value. The LSTMs can also be trained by BPTT.

The Two Training Models

For simplicity, we will focus on learning the memory-dependent terms in the GLE, neglecting for now the noise term. For the K-S equation, as shown in Figure 3, by directly fitting the stress using the history of the strain, we can reduce the error to nearly $1\%$ while maintaining a small gap between the testing and training error. This shows that we are not overfitting, and the history of the strain determines most of the stress.

We will discuss two models for learning the memory effect: a direct training model and a coupled training model. For both models, we generate data using direct numerical simulation (DNS). In the direct training model, after generating the data, we directly train the RNN to fit the sub-grid stress $\tau$ using the time history of the strain. The loss function is defined as the difference between the output of the RNN and the true stress. The neural network model for the stress is then used in the macro-scale equation for the reduced system. In the coupled training model, the loss function is defined as the difference between the solutions of the reduced model (with stress represented by the neural network) and the ground truth solution. Therefore in the coupled model, the RNN is coupled with the solver of the macro-scale equation when training.

In the direct training model, a RNN is used to represent the stress as a function of the time history of the macrocell strain. The loss function is simply the squared difference of the predicted stress and the ground truth stress. This RNN model is then used in the reduced system. Figure 1 shows the flow chart of the direct training model.

2 The coupled training model

Figure 2 shows part of the computational graph of the coupled training model.

In the coupled training model, the trainable parameters are still in the RNN, but to compute the gradient of the loss with respect to the parameters, we have to do back-propagation (BP) through the solver of the macro-scale equation, i.e. through (16). In many applications, this solver is complicated and it is hard to do BP through the solver. To deal with this problem, we take one step back to perform BP directly through the differential equations. This is done by writing down the differential equation satisfied by the gradient, which we call the backward equation. Another solver is used to solve this backward equation. In this way, BP is done by solving the backward equation, and is decoupled from the macro-scale solver.

As an illustration, let $\hat{u}(t)$ be the solution of the K-S equation at time $t$ in the Fourier space, and assume that we want to perform back-propagation through the K-S equation from time $t+\delta$ back to $t$, which means we are going to compute the Jacobian $\partial\hat{u}(t+\delta)/\partial\hat{u}(t)$. Let

then we have $J(0)=I$ and we want to compute $J(\delta)$. In the reduced system, we solve the K-S equation with stress,

where $\hat{\tau}$ is the sub-grid stress in Fourier space, $\ast$ means convolution, and $\mbox{diag}(v)$ represents a diagonal matrix with vector $v$ being the diagonal entries. Taking derivative of $J$ with respect to $s$, and assuming that $\partial\hat{\tau}(t)/\partial\hat{u}(t)=0$, we obtain

Hence, as long as we know the solution $\hat{u}(t+s)$ for $0\leq s\leq\delta$, we can compute the Jacobian $J(\delta)$ by solving the equation (19) from to $\delta$, with initial condition $J(0)=I$.

Numerical Experiments

We now present some numerical results of the proposed models for the K-S equation and the 2-dimensional shear flow problem. Below when we talk about true solution or ground truth, we mean the exact solution after filtering.

The K-S equation is considered in the Fourier space. To generate the training data, we solve the K-S equation with $N=256$ Fourier modes to approximate the accurate solution, using a $3^{rd}$ order integral factor Runge-Kutta method . We set $L=2\pi/\sqrt{0.085}$ and the time step $dt=0.001$. The settings are similar to that in . The micro-scale equation is solved for $1.2\times 10^{5}$ time units and filterd outputs are saved for every $0.1$ time units. We drop the results of the first $10^{4}$ time units, and take the results from the following $10^{5}$ time units as the training data, and the results of the last $10^{4}$ time units as the testing data.

For the reduced system, we solve (4) in Fourier space. We take $K=16$ Fourier modes and the time step $dt_{r}=0.1$. The macro-scale solver is still a $3^{rd}$ order integral factor Runge-Kutta method. The solver works in the Fourier space, and takes the output of the RNN as the sub-grid stress.

On the machine learning side, we use an LSTM to predict the stress. The LSTM has two layers and $64$ hidden units in each layer. An output layer with linear activation is applied to ensure that the dimension of the outputs is $16$. The LSTM works in the physical space: it takes strains in the physical space as inputs, and outputs predicted stresses in the physical space. A Fourier transform is applied to the output of the LSTM before feeding it into the macro solver.

Direct training

For the direct training model, we train the LSTM to fit the stress as a (memory-dependent) function of the strains for the last $20$ time steps. The network is trained by the Adam algorithm for $2\times 10^{5}$ iterations, with batch size being $64$. Figure 3 shows the relative training and testing error during the training process. We can see that the two curves go down simultaneously, finally reaching a relative error of about $1\%$. There is almost no gap between the training and the testing error, hence little overfitting. This also suggests that most contribution to the sub-grid stress can be explained by the time history of the strain.

First we consider some a priori results. Figure 4 presents the relative error and the correlation coefficients between the predicted stress and the true stress at different locations in the physical space. We can see that our prediction of the stress has relative errors of about $1\%$ and correlation coefficients very close to $1$.

We next examine some a posteriori results. After training the model, we pick some times in the testing set, and solve the reduced system initialized from the true solution at these times. Then, we compare the reduced solution with the true solution at later times. Figure 5 shows two examples. In these figures, the $x$-axis measures the number of time units from the initial solution. From these figures we can see that the reduced solution given by the direct training model produces satisfactory prediction for $150$-$200$ time units ($1500$-$2000$ time steps of macro-scale solver).

As for the eventual deviation between the two solutions, it is not clear at this point what is more responsible, the intrinsic unstable behavior in the model or the model error. We will conduct careful studies to resolve this issue in future work.

Next we compare the long-term statistical properties of the solution to the reduced model with the true solution. We solve the reduced system for a long time ($10^{4}$ time units in our case). We then compute the distributions and the autocorrelation curves of Fourier modes of the reduced solution, and compare with that of the true solution. Figure 6 shows some results. From these figures we see that the reduced model can satisfactorily recover statistical properties of the true solution.

Coupled training

For the coupled training model, we train the LSTM together with the equation solver. A detailed description of this training model is given in Section 3.2. Here we choose $L=20$ during training. The size of the LSTM is the same as that in the direct training model. A simple one-step forward scheme is used to solve the backward equation for the Jacobian (19). $2\times 10^{4}$ iterations were performed using an Adam optimizer with batch size $16$. During the training process, we measured the relative error of predicted solutions $l$ steps later from the initial true solution. The results for different $l$’s are given in Figure 7. From the figure we can see that the relative error for different $l$ goes down to $6-7\%$ after training, and there is no gap between the different curves, which means that solutions with different steps from the initial time are fitted nearly equally well.

Short-term prediction and long-term statistical performance are also considered for the coupled training model. Results are shown on Figure 8 and 6. From the figures we see that, the reduced system trained by the coupled training model gives satisfactory prediction for $50-100$ time units, which is shorter than the direct training model. However, the auto-correlation of Fourier modes match better with the true solution, while the distribution is as good as the direct training model. This suggests that, compared to the direct training model, the coupled model is less competitive for short-term prediction, but performs better for long-term statistical properties.

2 The 222-dimensional shear flow

Consider the $2$-dimensional shear flow in channel, whose governing equation is given by (6). We take $(x,y)\in\times$ and use periodic boundary condition in $x$ direction and zero boundary condition at $y=\pm 1$. We choose $Re=10000$, and $f=2/Re$ as a constant driving force. To numerically solve the equation, we employ the spectral method used in . We take $256$ Fourier modes in $x$ direction and $33$ Legendre modes in $y$ direction. The time step $\Delta t$ are chosen to be $0.005$.

For ease of implementation, we only do model reduction for Fourier modes in $x$ direction, and keep all Legendre modes in $y$ direction. The macro solution has $64$ Fourier modes in $x$ direction. Still, we generate accurate solution by DNS and compute the macro-scale strain and stress, and use an LSTM to fit the stress as a function of the history of the strain. In this experiment, the neural network we use is a $4$-layer LSTM with $256$ hidden units in each layer. As for the K-S equation, the input and output of the LSTM are in the physical space.

In this problem, the stress has $3$ components ($\tau_{11}=\mkern 1.5mu\overline{\mkern-1.5muuu\mkern-1.5mu}\mkern 1.5mu-\bar{u}\bar{u}$, $\tau_{12}=\mkern 1.5mu\overline{\mkern-1.5muuv\mkern-1.5mu}\mkern 1.5mu-\bar{u}\bar{v}$, $\tau_{22}=\mkern 1.5mu\overline{\mkern-1.5muvv\mkern-1.5mu}\mkern 1.5mu-\bar{v}\bar{v}$). Since each component has $33\times 64=2112$ modes in the spectral space, we need at least $2112$ variables in the physical space to represent the stress. If directly fit the stress, our LSTM will have about $6$ thousand outputs. This can make the training very difficult. Here, noting that both the stress and the strain are periodic in the $x$ direction, we choose to fit the stress by column. We train an LSTM to predict one column of the stress (the stress at the same $x$ in the physical space), using strains at this column and the neighboring columns. Figure 9 shows this idea. In practice, when predicting the $k$-th column of the stress, we use strains from the $(k-2)$-th to the $(k+2)$-th column.

For the $2$-dimensional shear flow problem, we only show results from the direct training model. The LSTM is trained by an Adam optimizer for $10^{5}$ steps, with batch size being 64. Still, the stress is predicted using the strains from the last $20$ time steps.

Numerical results

Figure 10 shows the relative training and testing error during the training process. We can see that the training error goes down to $3-4\%$, while the testing error goes down to about $5\%$. Figure 11 shows the correlation coefficients of the predicted stress and the true stress at different $y$. The three curves represent three components of the stress, respectively. we can see that the correlation coefficients are close to $1$ except in the area close to the boundary. Considering the boundary condition, the stress near the boundary is close to . Hence, its reasonable for the prediction to have a low correlation with the true stress.

Next, we solve the reduced system initialized from a true solution, and compare the quality of the reduced solution at later times with the solution given by the Smagorinsky model . The Smagorinsky model is a classical and widely used model for large eddy simulation (LES). In the two dimensional case, the Smagorinsky model for the sub-grid stress can be written as

where $\tau_{kk}=\tau_{11}+\tau_{22}$, $\nu_{t}$ is called the turbulent eddy viscosity, and

The turbulent eddy viscosity can be expressed as

with $\Delta$ being the grid size and $C_{s}$ being a constant. In Figure 12, we compare the deviation of our reduced solution and Smagorinsky solution from the true solution. We see that the prediction given by our reduced system is much better than that by the Smagorinsky model.

Conclusion

Much work needs to be done to develop the methods proposed here into a systematic and practical approach for model reduction for a wide variety of problems. For example, in the coupled training model, it is crucial to find an accurate and efficient way to compute the Jacobian. How to make the method scalable for larger systems is a problem that should be studied. For reduced systems where the noise effects cannot be neglected, how to model the noise term in the GLE is a problem that should be dealt with. Finally, we also need theoretical understanding to answer questions such as how to choose the size of the neural network, and how large the dataset should be in order to obtain a good reduced system.

References