MultiAuto-DeepONet: A Multi-resolution Autoencoder DeepONet for Nonlinear Dimension Reduction, Uncertainty Quantification and Operator Learning of Forward and Inverse Stochastic Problems¹¹footnotemark: 1

As we known, autoencoders were first introduced by Hinton and the PDP group [49] in the 1980s. Together with Hebbian learning rules [19, 39], autoencoders lay the foundation of unsupervised learning. An autoencoder is a special type of feedforward neural network where the output is set to be the same as the input and it often consists of two main parts: an encoder and an decoder. The encoder is used to process the inputs and the decoder takes the output of encoder and uses it to reconstruct the inputs. The outputs from the encoder can be understood as a latent representation of the inputs. Autoencoder can be trained by minimizing the reconstruction error which measures the differences between the original input and output from the reconstruction process of decoder. Traditional application for autoencoder includes dimensionality reduction, feature learning and etc. We apply autoencoder in our model to find the hidden features in inputs from high-dimensional stochastic process.

Suppose our target problem is a simple forward problem of stochastic DE, i.e. to find the solution $u(x;\omega)$. The randomness comes from the coefficient $k(x;\omega)$, where $\omega\in\Omega$ and $\Omega$ is the random space. Figure 1 presents the unsupervised part of our MultiAuto-DeepONet model. We can see it is an autoencoder as stated above except that we use a convolutional encoder and the decoder part is in the form of a DeepONet. The encoder consists of several convolutional layers followed by a dense layer. The filter size, strides in convolutional layers and the activation functions in each layer depend on the specific problem and will be tuned to achieve the best performance. The hidden features $z(x;\omega)$ discovered by the convolutional encoder are fed into the branch net of the decoder. Moreover, $z(x;\omega)$ can be used to generate samples of input $k(x;\omega)$. This makes uncertainty quantification of our model much easier. Assume the number of training data is $N$, the network loss function denoted by $L_{k}$ for unsupervised part is exactly the reconstruction error of this autoencoder,

where $k_{i}$ are training data and $\tilde{k}_{i}$ are network reconstruction of $k(x;\omega)$ corresponding to each $k_{i}$.

We also present the second supervised part of our MultiAuto-DeepONet model here in Figure 2. It has similar structure as unsupervised part in Figure 1 and they share one common branch net. The difference is that the output of supervised trunk net is used to approximate the SDE solution $u(x;\omega)$. So assume the number of training data is $M$, loss function for this supervised part denoted by $L_{u}$ is,

where $u_{i}$ are training data and $\tilde{u}_{i}$ are network prediction of $u(x;\omega)$ corresponding to each $u_{i}$.

The authors in [31] introduce DeepONet, a universal operator approximator which has shown remarkable approximation and generalization capabilities. Furthermore, they demonstrate that DeepONet can learn both explicit and implicit operators including those representing SDEs. We use DeepONet as the building block in this paper. As stated in the section of learning stochastic operators in [31], the input of the branch net is a random process and its dimension can be very high. As shown in Figure 3, a convolutional autoencoder is used to reduce the dimensionality thus discovering the hidden features $z(x;\omega)$ from the inputs $k(x;\omega)$ first. Then $z(x;\omega)$ is treated as the input to the common branch net of two DeepONets. The upper DeepONet in Figure 3 is exactly the unsupervised part in Figure 1 and the bottom DeepONet is the supervised part shown in Figure 2. Figure 4 presents the architecture of our convolutional encoder. The details of architectures of the two DeepONets is shown in Figures 5 and 6 respectively. Note that the unsupervised part of our MultiAuto-DeepONet model reconstructs $\tilde{k}(x;\omega)$ from the inputs and the supervised part aims at learning the solution $u(x;\omega)$. The following is the loss function of our model:

Note that in MultiAuto-DeepONet model, DeepONet is used as the decoder to find a nonlinear basis for both $k(x;\omega)$ and $u(x;\omega)$. This improves the interpretation of our model. We also add $L_{1}$ regularization in the two DeepONets to introduce the sparse structures. Thus a relative small number of basis is required to express the solution of a complex high-dimensional problem. Other regularization can be added similarly for different purposes. The network performance and uncertainty quantification of our model will be presented in next section.

Let us first consider a pedagogical example, i.e. one dimensional stochastic ODE representing the population growth model:

Here, $\Omega$ represents the random space and the initial condition is set to be $u(0;\omega)=1$. As in [31], the coefficient $k(t;\omega)$ is modeled as the following Gaussian random process,

where $k_{0}(t)=0$ is the mean funtion and the covariance function is in the squared exponential form $\text{Cov}(t_{1},t_{2})=\sigma^{2}\text{exp}(-\frac{\left\lVert t_{1}-t_{2}\right\rVert^{2}}{2l^{2}})$. Here, $\sigma$ is set to be $1$ and the correlation length $l$ is in $[1,2]$.

The input in this problem comes from a random process and the authors in [31] employ the Karhunen-Loeve(KL) expansion as in Equation 1 to handle it. In their example, $N$ is chosen to be $5$ which can be viewed as a truncation process. Although they did not assume to know the covariance function, the KL expansion must be known in advance. In our MultiAuto-DeepONet model, we can directly treat the observations of $k(t;\omega)$ as the inputs. The model automatically performs the nonlinear dimension reduction and discovers the hidden features first. The output of our model consists of both the solution of the equation and the reconstructed random process $k(t;\omega)$. Note that this reconstruction process of $k(t;\omega)$ can be used to generate new samples once the model is well calibrated.

In our first experiment, we train MultiAuto-DeepONet model with a data set of $25,000$ different $k(t;\omega)$. The training data set consists of $1000$ samples with $25$ input sensors for each trajectories. The test MSE is approximately $0.026\%$ and the average $L^{2}$ relative error is around $0.52\%$. In Figure 7a, the blue solid line represents the training loss and the red dashed line represents the validation loss. The difference between the two losses can be ignored after hundreds of epoches. Part $(b)$ shows the predictions for three different random samples from $k(t;\omega)$. The reference solutions are solid lines with different colors and the model predictions are triangle up marker lines. We can see the predictions from MultiAuto-DeepONet model matches the reference pretty well. Table 1 shows the MSE and average relative $l^{2}$ error of MultiAuto-DeepONet model for different the number of input $k$ sensors. Both errors decrease as the number of sensors increases.

Next, we conduct a experiment to compare our model with the other two different models. The first one is the original DeepONet model. The second is PCA based DeepONet model which we call it PCA-DeepONet. This model performs PCA to the inputs from $k(t;\omega)$ instead of reducing the spatial dimension using autoencoder. Figure 8 presents the predicted mean and variance of the solution $u$ for the reference and different methods. Table 2 lists the relative $l^{2}$ error of the predicted mean and variance when the number of input sensors is $20$. From Figure 8 and Table 2, we can see that our MultiAuto-DeepONet model achieves the best accuracy for predicted mean and variance among all three models when the training data set is relatively small.

As we have mentioned, we add $L_{1}$ regularization in the two trunk nets of MultiAuto-DeepONet model in order to introduce the sparsity. Figure 9 shows the coefficients and basis learned from our model for one specific samples at the corresponding output sensor locations. Part $(a)$ shows the learned coefficients from the unsupervised trunk net and part $(b)$ is the learned coefficients from the supervised trunk net. The learned basis from the branch net is presented in part $(c)$. The $x$ axes in the figures are the output sensor locations and the $y$ axes are corresponding values. We can see that only few numbers of them are not zero. This reduces the number of parameters in the network thus model training and calibration become much easier and faster.

One major difference between our model and original DeepONet model is that we can generate samples of $k(x;\omega)$ and $u(x;\omega)$ simultaneously after the network is trained. This is achieved by using the hidden features $z(x;\omega)$ learned from the MultiAuto-DeepONet model and estimate the probability density function of $z$ to construct a generator. The method we apply in this paper is the kernel density estimation(KDE) and $3000$ samples of $u(x;\omega)$ are generated. We present the mean and variance of those samples in Figure 10. We can see the accuracy of the mean and variance of generated samples are relatively good considering the size of training data set.

In this example, we consider a stochastic inverse problem, i.e. the following one-dimensional stochastic Poisson equation:

with homogeneous boundary conditions,

Here, $x\in[-1,1]$ and $\omega\in\Omega$ which is the random space and the randomness comes from the forcing term $f(x;\omega)$. We model $f(x;\omega)$ as a Gaussian process,

with mean function $f_{0}(x)=sin(\pi x)$ and covariance function $\text{Cov}(x_{1},x_{2})=\sigma^{2}\text{exp}(-\frac{\left\lVert x_{1}-x_{2}\right\rVert^{2}}{2l^{2}})$. The standard deviation $\sigma$ is chosen to be $1$ and the correlation length $l=1.5$.

First, we train the model with a data set of $40,000$ different $u(x;\omega)$. The number of sample trajectories is $1000$ and the number of input $u$ sensors is $40$. The test MSE is approximately $0.12\%$ and the average $L^{2}$ relative error is around $0.016$. In Figure 11a , the blue solid line represents the training loss and the red dashed line represents the validation loss. Part $(b)$ shows the predictions for three different random samples from $f(x;\omega)$. We can see the predictions from MultiAuto-DeepONet model is accurate. Table 3 shows the MSE and average relative $l^{2}$ errors of MultiAuto-DeepONet model for different number of input $u(x;\omega)$ sensors. We see they both decrease as the number of input $u$ sensors increases.

Now, a comparison between our model and PCA based DeepONet model is performed. Figure 12 presents the predicted mean and variance of unknown forcing term $f$ and the reference of different methods for the inverse stochastic Poisson equation. Table 4 lists the relative $l^{2}$ error of the predicted mean and variance when the number of input sensor is $40$. We can see that our MultiAuto-DeepONet model achieves the best performance for predicted mean and variance.

Figure 13 shows the coefficients learned from our model for one specific sample of $u$ at the corresponding output sensor locations. Part $(a)$ shows the learned coefficients from the unsupervised trunk net and part $(b)$ is the learned coefficients from the supervised trunk net. The learned basis from the branch net is presented in part $(c)$. The $x$ axes in the figures are the output sensor locations and the $y$ axes are corresponding values. We can see that most of those value are zero in this problem. This results from the $L_{1}$ regularization in the trunk nets.

In this example, we consider the forward problem of a two-dimensional stochastic Poisson equation:

with homogeneous boundary conditions,

Here, $\Omega$ is the random space and the randomness comes from the forcing term $f(\vec{t};\omega)$ ($\vec{t}=(x,y)$) as in the last example except that the spatial dimension is $2$. Again as the previous example, $f(\vec{t};\omega)$ is modelled as a Gaussian process:

with mean function $f_{0}(\vec{t})=20sin(\pi(x+y))$ and covariance function $\text{Cov}(\vec{t}_{1},\vec{t}_{2})=\sigma^{2}\text{exp}(-\frac{\left\lVert\vec{t}_{1}-\vec{t}_{2}\right\rVert^{2}}{2l^{2}})$. The hyper-parameters $\sigma$ and $l$ are set to be $0.5$ and $1.5$ respectively.

In this two dimensional problem, the latent dimension of the encoder is chosen to be $15$. There are two convolutional layers and two dense layers with relu activation function. For convolution layers, the stride length is $1$ and the filter size is $5$. We add $L_{1}$ regularization in the two trunk nets of the decoder to introduce the spasity. This step is crucial because the spatial dimension is two in this problem. The Adam optimizer is applied to train the network with batch size $256$. For our training data set, the number of sample trajectories of $f(x,y;\omega)$ is $2000$ and the number of input $f(x,y;\omega)$ sensors is $400$, i.e. $20\times 20$ grid.

In Figure 14a, the blue solid line and red dashed line represent training and validation loss respectively. We can see the generalization error is very small. Figure 14b presents the model prediction of one particular sample path for the solution $u(x,y;\omega)$. Figure 14c is the ground true plot and Figure 14d shows the difference plot between MultiAuto-DeepONet approximation and reference solution. We can see they match pretty well. In Table 5, we list the test MSE and average relative $l^{2}$ error of MultiAuto-DeepONet model for different number of input sensors. The number of input sample path is fixed to be $2000$ in this case. Table 6 shows the test MSE and average $l^{2}$ error for different number of input sample path. The number of input sensors is fixed to be $361$. The results shown in the two tables illustrate the convergence of the model.

Now, we compare our model with PCA based DeepONet for the predicted mean and variance. The number of sample trajectories of is $1800$ and the number of input sensors is $400$ for the test set. Figure 15a shows the predicted mean of our model and Figure 15b shows the difference between our model and the reference. Figure 15c presents the PCA based DeepONet predicted mean and Figure 15d is the corresponding difference. The mean of reference is shown in Figure 15e. In Figure 16, we present the predicted variance of our model, PCA based DeepONet model and their difference with the reference. Table 7 shows the relative $l^{2}$ errors of predicted mean and variance for the two models. We can see MultiAuto-DeepONet outperforms PCA based DeepONet in this case for both mean and variance. Figure 17 presents the coefficients and basis learned from MultiAuto-DeepONet model for one specific sample at corresponding sensor locations. We can see most of those values are zeros indicating a sparse structure.

We perform another experiment in this section to show the effectiveness of prediction basis $\phi(\vec{t})$ from the branch net in MultiAuto-DeepONet model. First we replace the branch net part, i.e. the basis $\phi(\vec{t})$ with the basis learned using polynomial chaos expansion(PCE) and fix this part during the training process of network. Then the MSE of this method is compared with our original model. Note that we choose the order of Legendre polynomials in PCE method to be $3$ considering both the accuracy and computational efficiency in this case. During the training process of MultiAuto-DeepONet model, we adjust the number of basis and coefficients learned from the branch net and trunk nets to be the same with the number of basis in PCE method. The number of input sensors is $361$ and the number of input sample path is $2000$. The test MSE of PCE method and our model are approximately $0.102\%$ and $0.047\%$ respectively. This means MultiAuto-DeepONet model can capture the basis of the SDE solution better for this forward stochastic problem.

In last example, we consider the forward problem of the Korteweg–De Vries (KdV) equation with time-dependent additive noise. It is a more complicated and nonlinear equation modeling waves on shallow water surfaces,

with initial condition,

Here, $\omega\in\Omega$ is the random space in this example. We know the analytical solution is,

where $W(t;\omega)$ is defined to be,

Here, we use KL expansion to model $f(t;\omega)$ as a Gaussian random field,

where $\sigma$ is a constant and we set it to be $0.1$, $\{\lambda_{i}\}_{i=1}^{d}$ and $\{\phi_{i}(t)\}_{i=1}^{d}$ are $d$ largest eigenvalues and corresponding eigenfunctions of the covariance kernel:

In this example, $d$ is set to be $3$ and $l_{c}=0.25$. We also limit the spatial $x\in[0,4]$ and the final time is set to be $t=0.1$s. The input of our encoder is similar to the previous examples. The unsupervised trunk net is used to process the temporal input $t$ and the supervised trunk net is used to process the input of both spatial and temporal inputs $(x,t)$. This also illustrates the capability of our model to deal with inputs of different dimensions.

We first train the model with $4000$ sample trajectories of $f(t;\omega)$. The number of input sensor for $f(t;\omega)$ is $400$, i.e. $40\times 10$ grid for $(x,t)$. The latent dimension for the encoder is fixed to be $4$. Other settings are the same as in two dimensional Poisson example. Figure 18a shows the training and validation loss. The red and blue lines match well meaning generalization error is very small. In Figure 18b, we present the model prediction and the reference solution at $t=0.1$s of $u(x,t;\omega)$ for one sample trajectory of $f(t;\omega)$. The test MSE is about $6.34E-05$ and average relative $l^{2}$ error is around $0.097\%$.

Now, we compare our model with PCA based DeepONet for the predicted mean and variance. The number of sample trajectories of is $2000$ for the test set and the number of input sensor is the same as training data set. Figure 19a shows the predicted mean of our model and Figure 19c shows the predicted mean of PCA based DeepONet model at final time $t=0.1$s. In Figures 19b and 19d, the predicted variance of our model and PCA based DeepONet model are presented. Table 8 shows the relative $l^{2}$ errors of predicted mean and variance for the two models. We can conclude that our model achieves better accuracy in terms of predicted mean and variance in this example. Again, Figure 20 presents the sparse coefficients and basis learned from MultiAuto-DeepONet model for one specific sample at corresponding sensor locations.