Bayesian Deep Convolutional Encoder-Decoder Networks for Surrogate Modeling and Uncertainty Quantification

Introduction

Uncertainty in complex systems arises from model error and model parametrization, unknown/incomplete material properties, boundary conditions or forcing terms, and other. Uncertainty propagation takes place by reformulating the problem of interest as a system of stochastic partial differential equations (SPDEs). Solution of such problems often needs to rely on the solution of the deterministic problem at a finite number of realizations of the random input using Monte Carlo sampling, or collocation methods. Considering the computational cost of solving complex multiscale/multiphysics deterministic problems, one often relies on Bayesian surrogate models that are trained with only a small number of deterministic solution runs while at the same time they are capable of capturing the epistemic uncertainty introduced from the limited training data .

For realistic problems in science and engineering, we only have access to limited number (e.g. $100$ or so) of deterministic simulation runs. Vanilla Monte Carlo for uncertainty propagation is thus hopeless. A dominant solution is to train a surrogate model using the limited simulation-based training data, and then perform prediction and uncertainty propagation tasks using the surrogate instead of solving the actual PDEs. Unfortunately most existing surrogate models have difficulty scaling to high-dimensional problems, such as the ones based on Gaussian processes (GP) or generalized polynomial chaos expansions (gPC ). High dimensionality often arises from the discretization of properties with small correlation lengths (e.g. permeability in heterogeneous media flows), random distributed sources or force input fields with multiple scales .

To alleviate the curse of stochastic input dimensionality, we usually assume that the given input data lie on an embedded non-linear manifold within the higher dimensional space. This intrinsic dimensionality is captured by dimensionality reduction techniques , such as the Karhunen-Loève expansion (KLE), t-SNE , auto-encoders , probabilistic methods like variational auto-encoders , Gaussian process latent variable models (GP-LVM) , and many more. Most dimensionality reduction models are unsupervised learning problems that do not explicitly take the regression task into account. Thus the classical approach to uncertainty quantification is to first reduce the dimensionality of the input to obtain a low-dimensional latent representation of the input field, then to built a regression model from this latent representation to the output. This approach is certainly not efficient is particular when the map from the latent representation to the physical space of the input data is not available. In , KLE was used for dimensionality reduction of the permeability field and then GP was performed as independent task for Bayesian regression. In , this approach was taken one step further with the probabilistic mappings from input to latent space and from latent space to output being modeled by generalized linear models both trained simultaneously end-to-end using stochastic variational inference instead of performing the unsupervised and supervised/regression tasks separately.

One of the essential upcoming approaches for handling high-dimensional data is to learn the latent input representation automatically by supervision with the output in regression tasks. This is the central idea of deep neural networks , especially convolutional neural networks (CNNs) which stack (deeper) layers of linear convolutions with nonlinear activations to automatically extract the multi-scale features or concepts from high-dimensional input , thus alleviating the hand-craft feature engineering, such as searching for the right set of basis functions, or relying on expert knowledge.

However, the general perspective for using deep neural networks in the context of surrogate modeling is that physical problems in uncertainty quantification (UQ) are not big data problems, thus not suitable for addressing them with deep learning approaches. However, we argue otherwise in the sense that each simulation run generates large amount of data which potentially reveal the essential characteristics about the underlying system. In addition, even for a relatively small dataset, deep neural networks show unique generalization property . These are typically over-parameterized models (hundreds and thousands of times more parameters than training data), but they do not overfit, i.e. the test error does not grow as the network parameters increase. Deep learning has been explored as a competitive methodology across fiels such as fluid mechanics , hydrology , bioinformatics , high energy physics and other.

This unique generalization behavior makes it possible to use deep neural networks for surrogate modeling. They are capable to capture the complex nonlinear mapping between high-dimensional input and output due to their expressiveness , while they only use small number of data from simulation runs. In addition, there has been a resurgence of interest in putting deep neural network under a formal Bayesian framework. Bayesian deep learning enables the network to express its uncertainty on its predictions when using a small number of training data. The Bayesian neural networks can quantify the predictive uncertainty by treating the network parameters as random variables, and perform Bayesian inference on those uncertain parameters conditioned on limited observations.

In this work, we mainly consider surrogate modeling of physical systems governed by stochastic partial differential equations with high-dimensional stochastic input such as flow in random porous media . The spatially discretized stochastic input field and the corresponding output fields are high-dimensional. We adopt an end-to-end image-to-image regression approach for this challenging surrogate modeling problem. More specifically, a fully convolutional encoder-decoder network is designed to capture the complex mapping directly from the high-dimensional input field to the output fields without using any explicit intermediate dimensionality reduction method. To make the model more parameter efficient and compact, we use DenseNet to build the feature extractor within the encoder and decoder paths . Intuitively, the encoder network extracts the multi-scale features from the input data which are used by the decoder network to reconstruct the output fields. In similarity with problems in computer vision, we treat the input-output map as an image-to-image map. To account for the limited training data and endow the network with uncertainty estimates, we further treat the convolutional encoder-decoder network to be Bayesian and scale a recently proposed approximate inference method called Stein Variational Gradient Descent to modern deep convolutional networks. We will show that the methodology can learn a Bayesian surrogate for a problem with an intrinsic dimensionality of $50$, achieving promising results on both predictive accuracy and uncertainty estimates using as few as $32$ training data. More importantly we develop a surrogate for the case of $4225$ dimensionality using $512$ training data. We also show these uncertainty estimates are well-calibrated using a reliability diagram. To this end, we believe that Bayesian neural networks are strong candidates for surrogate modeling and uncertainty propagation in high-dimensional problems with limited training data.

The remaining of the paper is organized as follows: In Section 2, we present the problem setup for surrogate modeling with high-dimensional input and the proposed approach in treating it as an image regression problem. We then introduce the CNNs and the encoder-decoder network used in our model. In Section 3, we present the Bayesian formulation of neural networks and a non-parametric variational method for the underlying challenging approximate inference task. In Section 4, we provide implementation details and show the performed experiments on a porous media flow problem. We finally conclude and discuss the various unexplored research directions in Section 5.

Methodology

With the discretization described above and assuming for simplicity fixed boundary and initial conditions and source terms as appropriate, we consider the computation simulation as a black-box mapping of the form:

In order to tackle the limitations of using the deterministic computationally expensive simulator for uncertainty propagation, a surrogate model $\mathbf{y}=\mathbf{f}(\mathbf{x},\bm{\theta})$ is trained using limited simulation data $\mathcal{D}=\{\mathbf{x}^{i},\mathbf{y}^{i}\}_{i=1}^{N}$, to approximate the ‘ground-truth’ simulation-induced function $\mathbf{y}=\bm{\eta}(\mathbf{x})$, where $\bm{\theta}$ are the model parameters, and $N$ is the number of simulation runs (number of training simulation-based data).

Therefore, we transform the surrogate modeling problem to an image-to-image regression problem, with the regression function as

In distinction from an image classification problem which requires image-wise prediction, the image regression problem is concerned with pixel-wise predictions, e.g. predicting the depth of each pixel in an image, or in our physical problem, predicting the output fields at each grid point. Such problems have been intensively studied within the computer vision community by leveraging the rapid recent progress of convolutional neural networks (CNNs), such as AlexNet , VGG , Inception , ResNet , DenseNet , and many more. A common model design pattern for semantic segmentation or depth regression is the encoder-decoder architecture. The intuition behind regression between two high-dimensional objects is to go through a coarse-refine process, i.e. to reduce the spatial dimension of the input image to high-level coarse features using an encoder, and then recover the spatial dimension by refining the coarse features through a decoder. One of the characteristics shared by those vision tasks is that the input and output images share the underlying structure, or they are different renderings of the same underlying structure . However, for our surrogate modeling tasks, the input and output images appear to be quite different, due to the complex physical influences (defined by PDEs) of the random input field, forcing terms and boundary conditions on the system response. This was after all the reason of pursuing the training of a surrogate model that avoids the repeated solution of the PDEs for different input realizations. Surprisingly, as we will discuss later on in this paper, the encoder-decoder network still works very well.

The training data for the surrogate model of interest here include the realizations of the random input field and the corresponding multi-output obtained from simulation. Of interest is to address problems with limited training data sets considering the high-computational cost of each simulation run. However, note that key UQ tasks include the ability to predict the system response and our confidence on it using input realizations (testing dataset) consistent with the given training data but also the computation of the statistics of the response induced by the random input. Both of these tasks require the availability of a high-number of input data sets (e.g. $500$ data points for testing, and $10^{4}$ input data points for Monte Carlo calculation of the output statistics). The problem of generating more input realizations using only the training dataset is the solution of a generative model problem. There has been significant progress in recent years in the topic, such as the generative adversarial networks (GANs) and its ever exploding variants, variational auto-encoders (VAEs) , autoregressive models like PixelCNN , PixelRNN , and other. However, note that in this work our focus is on the image-to-image mapping and its performance on uncertainty quantification tasks. We will thus assume that enough input samples are provided both for testing and output statistics calculation even though only a small dataset will be used for training. In our examples in Section 4, synthetic log-permeability datasets are generated by sampling a Gaussian random field with an exponential kernel. The output for each permeability sample is generated using a deterministic simulator.

2 Dense Convolutional Encoder-Decoder Networks

In this subsection, we briefly introduce a state-of-the-art CNN architecture called DenseNet and fully convolutional encoder-decoder networks developed in computer vision, and then present how to utilize these advances to build our baseline network for surrogate modeling in uncertainty quantification.

DenseNet is a recently proposed CNN architecture which extends the ideas of ResNet and Highway Networks to create dense connections between all layers, so as to improve the information (gradient) flow through the network for better parameter efficiency.

Let $\mathbf{x}_{l}$ be the output of the $l^{th}$ layer. Traditionally CNNs pass the output of one layer only to the input of the next layer, i.e. $\mathbf{x}_{l}=h_{l}(\mathbf{x}_{l-1})$, where $h_{l}$ denotes the nonlinear function of the $l^{th}$ hidden layer. In current CNNs, $h$ is commonly defined as a composition of Batch Normalization (BatchNorm), Rectified Linear Unit (ReLU) and Convolution (Conv) or transposed convolution (ConvT) . ResNets create an additional identity mapping that bypasses the nonlinear layer, i.e. $\mathbf{x}_{l}=h_{l}(\mathbf{x}_{l-1})+\mathbf{x}_{l-1}$. In this way, the nonlinear layer only needs to learn a residual function which facilitates the training of deeper networks.

DenseNets introduce connections from any layer to all subsequent layers, i.e. $\mathbf{x}_{l}=h_{l}([\mathbf{x}_{l-1},\mathbf{x}_{l-2},\cdots,\mathbf{x}_{0}])$. To put it in another way, the input features of one layer are concatenated to the output features of this layer and this serves as the input features to the next layer. Assume that the input image has $K_{0}$ channels, and each layer outputs $K$ feature maps, then the $l^{th}$ layer would have input with $K_{0}+(l-1)\cdot K$ feature maps, i.e. the number of feature maps in DenseNet grows linearly with the depth. $K$ is here referred to as the growth rate.

For image regression based on encoder-decoder networks, downsampling and upsampling are required to change the size of feature maps, which makes concatenation of feature maps unfeasible. Dense blocks and transition layers are introduced to solve this problem and modularize the network design. A dense block contains multiple densely connected layers whose input and output feature maps are of the same size. It contains two design parameters, namely the number $L$ of layers within and the growth rate $K$ for each layer. An illustration of the dense block is in Fig. 1.

Transition layers are used to change the size of feature maps and reduce their number between dense blocks. More specifically, the encoding layer typically halfs the size of feature maps, while the decoding layer doubles the feature map size. Both of the two layers reduce the number of feature maps. This is illustrated in Fig. 2.

2.2 Fully Convolutional Networks

The fully convolutional networks (FCNs) are extensions of CNNs for pixel-wise prediction, e.g. semantic segmentation. FCNs replace the fully connected layers in CNNs with convolution layers, add upsampling layers in the end to recover the input spatial resolution, and introduce the skip connections between feature maps in downsampling and upsampling path to recover finer information lost in the downsampling path. Most of the recent work focuses in improving the upsampling path and increase the connectivity within and between upsampling and downsampling paths. U-nets extend the upsampling path as symmetric to the downsampling path and add skip connections between each size of feature maps in the downsampling and upsampling paths. Within SegNets , the decoder uses pooling indices computed in the max-pooling step of the corresponding encoder to perform non-linear upsampling. Fully convolutional DenseNets extend DenseNets to FCNs, which are closest to our network design but with several differences. We keep all the feature maps of a dense block concatenated so far before passing to the transition layers, while they only keep the output feature maps of the last convolution layer within the dense block. The feature maps explosion problem is addressed by the first convolution layer within the transition layer. Besides that we do not use skip connections between encoding and decoding paths because of the weak correspondence between the input and output images. We also do not use max-pooling for encoding layers, instead we use convolution with stride $2$.

3 Network architecture: DenseED

We follow the fully convolutional networks (FCNs) for image segmentation without using any fully connected layers, and encode-decoder architecture similar to U-net and SegNet but without the concatenation of feature maps between the encoder paths and decoder paths. Furthermore, we adapt the DenseNet structure into the encoder and decoder networks. After extensive hyperparameter and architecture search, we arrived at a baseline dense convolutional encoder-decoder network, called DenseED, similar to the network proposed in but with noticeable differences as stated above and shown in Fig. 3.

In the encoding path, the input field realizations are fed into the first convolution layer with large kernel size $k=7$, stride $s=2$ and zero padding $p=2$. Then the extracted feature maps are passed through an alternative cascade of dense blocks and encoding layers as introduced in Figs. 1 and 2. The dense block after the last encoding layer outputs the high-level coarse feature maps extracted from the input, as shown in purple at the right end of the network in Fig. 3, which are subsequently fed into the decoder path. The decoding network consists of an alternation of dense blocks and decoding layers, with the last decoding layer directly leading to the prediction of the output fields.

4 Network architecture engineering and hyperparameter search

Network architecture engineering and hyperparameter search are among the main challenges and source of innovations in deep learning, mostly problem-specific and empirical. The general network architecture is introduced in Section 2.2, which is based on the recent development of neural network design for image segmentation. For our image regression problem, the main design considerations include the following:

Downsampling layers: convolution or pooling;

Upsampling layers: bilinear upsampling or transposed convolution;

Smallest spatial dimensions of feature maps: this is determined by the number of downsampling layers;

Add or not of skip connections between the encoding and decoding paths;

Kernel of convolution layers, kernel size $k$, stride $s$, zero padding $p$;

Number of layers $L$ and growth rate $K$ within each dense block;

Regularizations: weight decay, batch normalization, dropout, etc;

Optimizer: stochastic gradient descent algorithms and their variants, such as Adam, Adagrad, RMSprop, and others;

Training hyperparameters: batch size, learning rate and its scheduler.

The details of architecture search and hyperparameter selection for the particular problem considered are presented in A where we also report various experiments using DenseED for surrogate modeling with limited training data. No overfitting was observed in our calculations and the obtained results were quite good. This is an intriguing and active research topic in the deep learning community .

In our non-Bayesian calculations, we have considered $\mathcal{L}_{2}$ or $\mathcal{L}_{1}$ regularized MSE training loss function. Given an input image $\mathbf{x}$, and a target image $\mathbf{y}$, the prediction $\mathbf{f}(\mathbf{x},\mathbf{w})$, the regularized MSE loss is

where the penalty function $\Omega(\mathbf{w})=\frac{1}{2}\mathbf{w}^{\top}\mathbf{w}$ for $\mathcal{L}_{2}$ regularization, and $\Omega(\mathbf{w})=\left\lVert\mathbf{w}\right\rVert_{1}=\sum_{i}|w_{i}|$ for $\mathcal{L}_{1}$ regularization, and $n=C_{out}\cdot H_{out}\cdot W_{out}$ is the number of pixels in all channels of one output image. $\mathbf{w}$ denotes all the parameters in the network. For our case, it includes the kernel weights in all the convolution and transposed convolution layers (no bias is used in convolutional kernel), the scale and shift parameters in all the batch normalization layers. See Section 4.3 for an example of the fully convolutional encoder-decoder network used for the Darcy flow problem. Note that $\mathcal{L}_{2}$ regularization is implemented in PyTorch optimizers by specifying weight decay, which is $\alpha$ in Eq. (3).

The network architecture selected for the non-Bayesian model is the same as that used for our Bayesian model introduced next.

In the encoder-decoder network, the batch normalization layer used after each convolutional layer can also be considered as an effective regularizer https://openreview.net/forum?id=BJlrSmbAZ¬eId=BJlrSmbAZ. It is commonly adopted nowadays in deep convolutional networks https://github.com/pytorch/vision/tree/master/torchvision/models replacing dropout https://www.reddit.com/r/MachineLearning/comments/5l3f1c/d_what_happened_to_dropout/.

Bayesian Neural Networks

Consider a deterministic neural net $\mathbf{y}=\mathbf{f}(\mathbf{x},\mathbf{w})$ with input $\mathbf{x}$, output $\mathbf{y}$, and all parameters $\mathbf{w}$ including the weights and biases. http://pytorch.org/docs/master/nn.html?ht=conv2d##torch.nn.functional.conv2d Bayesian neural networks (BNNs) treat the parameters $\mathbf{w}$ as random variables instead of deterministic unknowns to account for epistemic uncertainty induced by lack of training data. Besides that, usually additive noise $\mathbf{n}$ is introduced to model the aleatoric uncertainty which can not be reduced by having more observations, also to make the probabilistic model have an explicit likelihood depending on the noise distribution, i.e.

where $\mathbf{f}(\mathbf{x},\mathbf{w})$ is the output of a neural network with the uncertain $\mathbf{w}$, and $\mathbf{n}$ is the additive noise.

Given the large amount of ‘un-interpretable’ parameters $\mathbf{w}$ in a deep neural net, there are not many choices of priors. But the demand for compression of the neural net for lower memory and computation cost calls for sparsity promoting priors. We assume a fully factorized Gaussian prior with zero mean and Gamma-distributed precision $\alpha$ on parameters $\mathbf{w}$

This results in a student’s t-prior for $\mathbf{w}$, which has heavy tails and more mass close to zero.

2 Additive Noise Model

Additive noise can be considered of the following form:

Output-wise: $\mathbf{n}=\sigma\bm{\epsilon}$, same for all output pixels;

Channel-wise: $\mathbf{n}=[\sigma_{1}\bm{\epsilon}_{1},\cdots,\sigma_{d_{y}}\bm{\epsilon}_{d_{y}}]$, same across each of the $d_{y}$ output channels/fields;

Pixel-wise: $\mathbf{n}=\bm{\sigma}\odot\bm{\epsilon}$, distinct for each output pixel.

We can also model the noise varying with input (pixel-wise noise model), resulting in a heteroscedastic noise model, i.e. $\mathbf{n}(\mathbf{x},\mathbf{w})=\sigma(\mathbf{x},\mathbf{w})\bm{\epsilon}$ or $\mathbf{n}(\mathbf{x},\mathbf{w})=\bm{\sigma}(\mathbf{x},\mathbf{w})\odot\bm{\epsilon}$. Again $\bm{\epsilon}$ can be Gaussian or Laplacian. The heteroscedastic noise can be implemented as extending the output of the neural net as:

3 Stein Variational Gradient Descent (SVGD)

Approximate inference for Bayesian deep neural network is a daunting task because of the large number of uncertain parameters, e.g. tens or hundreds of millions in modern deep networks. In our surrogate problem, the task is to find a high-dimensional posterior distribution over millions of random variables using less than hundreds or thousands of training data. As reviewed in Section 1, most of variational inference methods restrict the approximate posterior within certain parametric variational family, while sampling-based methods are slow and difficult to converge. Here we adopt a recently proposed non-parametric variational inference method called stochastic variational gradient descent (SVGD) that is similar to standard gradient descent while maintaining the efficiency of particle methods.

For a prescribed probabilistic model with likelihood function $p(\mathbf{y}\mid\bm{\theta},\mathbf{x})$ and prior $p_{0}(\bm{\theta})$, we are interested in Bayesian inference of the uncertain parameters $\bm{\theta}$, i.e. to find the posterior distribution $p(\bm{\theta}\mid\mathcal{D})$, where $\mathcal{D}$ denote the i.i.d. observations (training data), i.e. $\mathcal{D}=\{\mathbf{x}^{i},\mathbf{y}^{i}\}_{i=1}^{N}$. For the BNNs with homoescedastic Gaussian noise case, $\bm{\theta}=\{\mathbf{w},\beta\}$. Variational inference aims to approximate the target posterior distribution $p(\bm{\theta}\mid\mathcal{D})$ with a variational distribution $q^{*}(\bm{\theta})$ which lies in a restricted set of distributions $\mathcal{Q}$ by minimizing the KL divergence between the two, i.e.

The variational family considered here is a set of distributions obtained by smooth transforms from an initial tractable distribution (e.g. the prior) represented in terms of particles. The transforms applied to each particle take the following form:

where $\epsilon$ is the step size, $\bm{\phi}(\bm{\theta})\in\mathcal{F}$ is the perturbation direction within a function space $\mathcal{F}$. When $\epsilon$ is small, $\mathbf{T}$ transforms the initial density $q(\bm{\theta})$ to

Instead of using parametric form for the variational posterior, a particle approximation is used, i.e. a set of particles $\{\bm{\theta}^{i}\}_{i=1}^{S}$ with empirical measure $\mu_{S}(d\bm{\theta})=\frac{1}{S}\sum_{i=1}^{S}\bm{\delta}(\bm{\theta}-\bm{\theta}^{i})d\bm{\theta}$. We would like to have $\mu$ to weakly converge to the measure of the true posterior $\nu_{p}(d\bm{\theta})=p(\bm{\theta})d\bm{\theta}$. We apply the transform $\mathbf{T}$ to those particles, and denote the pushforward measure of $\mu$ as $\mathbf{T}\mu$.

The problem is to find out the direction to maximally decrease the KL divergence of the variational approximation and the target distribution, i.e. to solve the following functional optimization problem:

where $\mathcal{T}_{p}$ is called the Stein operator associated to the distribution $p$,

The expectation $\E_{\mu}[\mathcal{T}_{p}\bm{\phi}]$ evaluates the difference between $p$ and $\mu$, and its maximum is defined as the Stein discrepancy,

It has been shown that when the functional space $\mathcal{F}$ is chosen to be the unit ball in a product reproducing kernel Hilbert space $\mathcal{H}$ with the positive kernel $k(\mathbf{x},\mathbf{x}^{\prime})$, the maximal direction to perturb (or the Stein discrepancy) has a closed-form solution,

Thus we have the following algorithm to transform an initial distribution $\mu_{0}$ to the target posterior $\nu_{p}$.

This is an one-line algorithm, where the gradient $\bm{\phi}(\bm{\theta})$ pushes the particles towards the high posterior region by kernel smoothed gradient term $k(\cdot,\cdot)\nabla\log p$, while maintaining a degree of diversity by repulsive force term $\nabla k(\bm{\theta},\bm{\theta}^{\prime})$. When the number of particles becomes $1$, then the algorithm reduces to the MAP estimate of the posterior. This algorithm is implemented in PyTorch with GPU acceleration and scaled to our deep convolutional encoder-decoder network DenseED.

Here we use one toy example to illustrate the idea of SVGD. We start with the $20$ particles from the Normal distribution $\mathcal{N}(-10,1)$, and tranport the particles iteratively to the target Gaussian mixture distribution $0.8\mathcal{N}(-2,1)+0.2\mathcal{N}(2,1)$ with Algorithm 1.

We use $S$ samples of $\bm{\theta}$ to approximate the empirical measure of the posterior. They are initialized and stored in $20$ deterministic DenseED neural networks.

At each step $t$, for each model $j$, the gradient of its joint likelihood (or unnormalized posterior) $\nabla_{\bm{\theta}^{j}_{t}}\log p(\bm{\theta}^{j}_{t})$ is computed by the automatic differentiation tool in PyTorch. We first compute the joint likelihood $\log p(\bm{\theta}^{j}_{t})=\prod_{i=1}^{N}p(\mathbf{y}_{i}\mid\bm{\theta}^{j}_{t},\mathbf{x}_{i})p(\bm{\theta}^{j}_{t})$ by feeding forward the data $\{\mathbf{x}^{i},\mathbf{y}^{i}\}_{i=1}^{N}$, then back propagate to compute its gradient $\nabla_{\bm{\theta}^{j}_{t}}\log p(\bm{\theta}^{j}_{t})$. Note that this gradient is stored in PyTorch module associated to the weights $\bm{\theta}_{t}^{j}$ in each network.

Then we can proceed to compute the kernel matrix $\Big[k(\bm{\theta}^{j}_{t},\bm{\theta}^{i}_{t})\Big]_{i,j\in\{1,\cdots,S\}}$, and its gradient $\nabla_{\bm{\theta}^{j}_{t}}k(\bm{\theta}^{j}_{t},\bm{\theta}^{i}_{t})$, as well as the kernel weighted gradient of the joint likelihood $k(\bm{\theta}^{j}_{t},\bm{\theta}^{i}_{t})\nabla_{\bm{\theta}^{j}_{t}}\log p(\bm{\theta}^{j}_{t})$. For this to happen, we need to vectorize (extract out) the parameters $\bm{\theta}_{t}^{j}$ and the computed gradient $\nabla_{\bm{\theta}^{j}_{t}}\log p(\bm{\theta}^{j}_{t})$ from each neural network. The optimal perturbation direction $\bm{\phi}$ is further computed by the sum of the two terms as shown in the SVGD algorithm.

After $\{\bm{\phi}(\bm{\theta}_{t}^{j})\}_{j=1}^{S}$ is computed, we send $\bm{\phi}(\bm{\theta}_{t}^{j})$ back to each neural network the gradient w.r.t. its parameters $\bm{\theta}_{t}^{j}$ (overwriting the previously computed $\nabla_{\bm{\theta}^{j}_{t}}\log p(\bm{\theta}^{j}_{t})$), and updating locally in each neural network using PyTorch’s optimization library such as Adam or SGD to compute $\bm{\theta}_{t+1}^{j}$.

With one iteration complete, the algorithm repeats the above steps until convergence in the parameters is achieved.

4 Uncertainty Quantification

Of interest to the classical UQ problem is the computation of the posterior predictive distribution (predict the system response for a test input) as well as the computation of the output response averaged over the input probability distribution. In particular, we are interested in computing the following:

We can use Monte Carlo to approximate the moments of the predictive distribution

with mean (by the law of total expectation)

Note that the SVGD algorithm provides a sample representation of the joint posterior of all parameters $p(\mathbf{w},\beta\mid\mathcal{D})$. To obtain the samples of the marginal posterior $p(\mathbf{w}\mid\mathcal{D})$ as needed above, one simply needs to use the samples corresponding to $\mathbf{w}$.

The predictive covariance can also be easily calculated using the law of total variance. The variance and expectation below are w.r.t. to the posterior of the parameters. We can show the following:

where $\beta^{i}\sim p(\beta\mid\mathcal{D}),\mathbf{w}^{i}\sim p(\mathbf{w}\mid\mathcal{D})$. The predictive variance is the diagonal of the predictive covariance:

where $\bm{1}$ is a vector of ones with the same dimension as $\mathbf{f}$, and the square $(\cdot)^{2}$ is applied element-wise to the vectors.

The above computation is the prediction at a specific input $\mathbf{x}^{*}$. We would also like to compute the average prediction over the distribution of the uncertain input. We first compute the output statistics given the realizations of the uncertain parameters $\bm{\theta}=\{\mathbf{w},\beta\}$, where $\bm{\theta}\sim p(\bm{\theta}\mid\mathcal{D})$. The conditional predictive mean is

and the conditional predictive covariance is

Also the conditional predictive variance (at each spatial location) is

Numerical Implementation and Results

We study the two-dimensional, single phase, steady-state flow through a random permeability field following the case study in Section $3.2$ in . Consider the random permeability field $K$ on a unit square spatial domain $\mathcal{S}=^{2}$, the pressure field $p$ and velocity field $\mathbf{u}$ of the fluid through the porous media are governed by Darcy’s law:

where $\hat{\mathbf{n}}$ denotes the unit normal vector to the boundary and the source term $f$ is used to model an injection well on the left-bottom corner of $\mathcal{S}$ and a production well on the right-top corner. We also enforce no-flux boundary condition, and an integral constraint to ensure the uniqueness of the solution as in . More specifically,

where $r$ is the rate of the wells and $w$ is their size. The input log-permeability field is restricted in this work to be a Gaussian random field, i.e.

where $m$ is the constant mean and covariance function $k$ is specified in the following form using the $L_{2}$ norm in the exponent instead of the $L_{1}$ norm in , i.e.

The Gaussian random field with exponential kernel for the one-dimensional case corresponds to the Ornstein-Uhlenbeck process which is mean-square continuous but not mean-square differentiable. Thus when we discretize the field over a grid, the field value jumps (varies highly) as we move from pixel to pixel. The field does not become smoother when we use a finer grid over a fixed spatial domain. This high variability creates a significant challenge for data-driven models to capture, i.e. the intrinsic dimensionality of the discretized random field is the total number of pixels, e.g. $4,225$ for $65\times 65$ grids (which will be our reference grid for our calculations). However, a common assumption for natural images is that the underlying dimensionality is actually small (few hundreds) despite their complex appearance. To evaluate the generality and effectiveness of the methodology, we use KLE to control the intrinsic dimensionality of the permeability dataset. We evaluated our model using datasets produced with increasing dimensionality of $50$, $500$, $4225$ (called KLE$50$, KLE$500$, and KLE$4225$, respectively). Notice that when the number of KLE terms is $4225$, the permeability field is directly sampled from the exponential Gaussian field without any dimensionality reduction. The intrinsic dimensionality of dataset is hidden from our model, i.e. our model do not built a map from the KLE terms to the system output. Instead, it models an end-to-end mapping from input fields to output fields. In fact, we will show one specific network architecture that works well for all three datasets obtained from the different intrinsic input data dimensions.

We consider solving the Darcy flow Eq. (18) over a unit squared domain $\mathcal{S}=^{2}$ with fixed $65\times 65$ grid, and length scale $l=0.1$, kernel mean $m=0$, rate of source $r=10$, size of source $w=0.125$. The ratio of the cumulative sum of eigenvalues (in decreasing order) over the total sum of them is shown in Fig. 5.

The Darcy flow equation is solved using mixed finite element formulation implemented in FEniCS with third-order Raviart-–Thomas elements for the velocity, and fourth-order discontinuous elements for the pressure. The sample input permeability field and computed output pressure and velocity fields for three datasets are shown in Fig. 6.

When the available data is limited, it is common practice to use cross-validation to evaluate the model. Since our dataset is synthetically generated, we have access to any number of training and test data up to computing constraints to solve the Darcy flow equations. Our current data includes four sets: the training set, validation set, test set, and uncertainty propagation set. The training set is sampled using the simplest design of experiment method, Latin hypercube sampling. More specifically, the KLE for the log-permeability field is

where $\lambda_{k}$ and $\phi_{k}(\mathbf{s})$ are the eigenvalues and eigenfunctions of the exponential covariance function of the Gaussian field specified in Eqs. (20) and (21), $z_{k}$’s are i.i.d. standard Normal, and $q$ is the number of KLE coefficients maintained in the expansion. The maximum number that can be used is finite and equal to the number of grid points used in the discetization of the field over the unit square. We first use Latin hypercube design to sample $\xi_{k}$ from the hypercube $^{q}$, then obtain the eigenvalue by $z_{k}=\Phi^{-1}(\xi_{k})$, where $\Phi$ is the cumulative distribution function of the standard normal distribution. The KLE$50$ case contains $32$, $64$, $128$, and $256$ training data; KLE$500$ contains $64$, $128$, $256$, and $512$ training data; and KLE$4225$ contains $128$, $256$, $512$, and $1024$ training data.

The log-permeability fields in the other three sets are reconstructed directly with $z_{k}$, which are sampled from standard normal. The validation and test set each contains $500$ input permeability fields, and the dataset for uncertainty propagation contains $10,000$ realizations. All datasets are organized as images as discussed in Section 2.1.

2 Evaluation Metrics

Several metrics are used to evaluate the trained models on test data $\{\mathbf{x}^{i},\mathbf{y}^{i}\}_{i=1}^{T}$. In particular, we consider the following:

Coefficient of determination ($R^{2}$-score):

where $\hat{\mathbf{y}}^{i}$ is the output mean of the Bayesian surrogate, i.e. $\sum_{i=1}^{S}\mathbf{f}(\mathbf{x},\mathbf{w}^{i})/S$ as in Eq. (12) or predictive output of the non-Bayesian surrogate, i.e. just $\mathbf{f}(\mathbf{x})$, $\mathbf{y}^{i}$ is the test target, $\bar{\mathbf{y}}$ is the mean of test target, and $T$ is the total number of test data. This metric enables the comparison between different datasets since the error is normalized, with the score closer to $1$ corresponding to better regression. This is the only metric used for evaluating non-Bayesian surrogate, the following metrics are additional metrics for evaluating the Bayesian surrogate. Note that this metric is also used for tracking the performance of the training process, thus it is evaluated for both the training and test data sets.

This is a common metric for regression that is used in our experiments for monitoring the convergence of training.

This metric evaluates the likelihood of the observed data. It is is used to assess the quality of the predictive model.

Predictive uncertainty and Propagated Uncertainty: These metrics were introduced in Section 3.4.

Estimated Distributions: They include histograms or kernel density estimates for the output fields at certain locations of the physical domain.

Reliability Diagram: Given a trained Bayesian surrogate and a test data set, we can compute the $p\%$ predictive interval for each test data point based on the Gaussian quantiles using the predictive mean and variance . We then compute the frequency of the test targets that fall within this predictive interval. For a well-calibrated regression model, the observed frequency should be close to $p\%$. The reliability diagram is the plot of the observed frequency with respect to $p$. Thus a well-calibrated model should have a reliability diagram close to the diagonal.

3 Non-Bayesian Surrogate Model

The hyperparameters to search include the parameters that determine the network architecture and the ones that specify training process, which both affect model performance. We use Hyperband algorithm to optimize those hyperparameters with a constraint that the number of model parameters being less than $0.25$ million. The details of these experiments are given in A. The network configuration with the highest $R^{2}$-score that Hyperband finds is shown in Fig. 7 with more details provided in Table 1. This configuration is referred to as DenseED-c16. The $2$nd–$4$th columns of Table 1 show the number $C_{f}$ of output feature maps, the spatial resolution $H_{f}\times W_{f}$ of output feature maps, the number of parameters of each layer in the network.

The network DenseED-c16 is trained with Adam , a variant of stochastic gradient descent, with the loss function being $L_{2}$ regularized MSE which is implemented as weight decay in modern neural net frameworks, such as PyTorch and TensorFlow. Other loss functions may achieve better results, such as smoothed $L_{1}$ loss, or conditional GAN loss . This requires further investigations to be considered in future publication. The initial learning rate is $0.015$, weight decay (regularization on weights) is $0.0005$, the batch size is $16$. We also use a learning rate scheduler which drops $10$ times on plateau of the rooted MSE. The model is trained $200$ epochs. We train the model with the dataset introduced in Section 4.1.

Training the deterministic neural networks with $L_{2}$ regularized MSE is equivalent to finding the maximum a posterior of the uncertain parameters in Bayesian neural networks whose prior is independent normal. The typical training process is shown in Fig. 8.

We train each network with different number of training data of KLE$50$, KLE$500$, and KLE$4225$. The validation $R^{2}$-score is shown in Fig. 9, which shows that, with the same training data, the $R^{2}$-score is closer to $1$ when the intrinsic dimensionality is smaller, and the $R^{2}$-score is higher with more training data of the same dimensionality. Note that the score is more than $0.9$ with reasonably small size training data set for all the three cases which have dimensionality from $50$ to $4225$. This shows the effectiveness of the network DenseED-c16 for both low-dimensional and high-dimensional problems.

The prediction of the output fields can be easily obtained in the test time by feeding the test input permeability field $\mathbf{x}^{*}$ into the trained network, i.e. $\hat{\mathbf{y}}^{*}=\mathbf{f}(\mathbf{x}^{*})$. We show the prediction of the test input shown in Fig. 6 using DenseED-c16, which is trained with three datasets (KLE$50$, KLE$500$, KLE$4225$) in Figs. 10, 11, and 12, respectively. The predictions are quite good even for the KLE$4225$ case, where both the input and output fields vary rapidly in certain regions of the domain.

4 Bayesian Surrogate Model

For all the experiments we only consider the homoscedastic noise model for Bayesian neural networks, i.e. output-wise Gaussian noise with Gamma prior on its precision $\beta$, and Student’s t-prior on $\mathbf{w}$.

The algorithm is implemented in PyTorch and runs on a single NVIDIA GeForce GTX $1080$ Ti X GPU which requires about $2000-7000$ seconds for training $300$ epochs, when the training data size varies from $32$ to $512$. The training time depends heavily on the training mini-batch size, which is $16$ for all cases. Potential ways to speed up significantly the training process include increasing the mini-batch size, or implementing the SVGD in parallel using multi-GPUs. The python source code will become available upon publication at https://github.com/bmmi/bayesnn.

We report next the $R^{2}$-score computed similar to the non-Bayesian case, except the predicted output mean is used to compare with the test target. The scores are shown in Fig. 13. We can see that the Bayesian surrogate improves the $R^{2}$-score significantly over the non-Bayesian version.

We also report the MNLP for test data in Fig. 14, which is a proper scoring rule and usually used to access the quality of predictive uncertainty .

The $R^{2}$-score gives us a general estimate of how well the regression performs. For a given input permeability field, the Bayesian neural network can predict the mean of corresponding output fields, and also gives uncertainty estimate represented as predictive variance at each spatial location, which is unavailable for deterministic models, and desirable when the training data is small. In Figs. 15, 16, and 17, we show predictions for the test input shown in Fig. 6 with training data from KLE$50$, KLE$500$, and KLE$4225$, respectively. We can see that the predictive accuracy improves as the size of the training dataset increases, while the predictive uncertainty drops.

We also performed uncertainty propagation by feeding the trained Bayesian surrogate with $10,000$ input realizations sampled from the Gaussian field, and calculating the output statistics as in Section 3.4. In Figs. 18, 19 and 20, we show the uncertainty propagation results and compare with Monte Carlo using the $10,000$ UP data for the Bayesian surrogate trained with the datasets KLE$50$, KLE$500$, and KLE$4225$.

We show the estimate of pressure $p$, velocity components $u_{x}$, $u_{y}$ at locations $(0.85,0.88)$, and $(0.68,0.05)$ on the unit square for $128$, and $512$ training data of KLE$4225$ in Fig. 21, and 22, respectively. The PDF is obtained by kernel density estimation using the predictive mean. We can see that the density estimate is close to the Monte Carlo result even when the training dataset is small, and becomes closer as the training dataset increases. From Fig. 22, we observe that the predictions for the velocity fields are better than the pressure field especially in locations away from the diagonal of the unit square domain, and this is in general the case for our current network architecture, where the three output fields are treated the same.

In order to access the quality of the computed uncertainty, we adopt the reliability diagram which expresses the discrepancy between the predictive probability and the frequency of falling in the predictive interval for the test data. The diagram is shown in Fig. 23. Overall our models are well-calibrated since they are quite close to the ideal diagonal, especially the case when the training dataset size is $128$ as shown in Fig. 23(d). In general the model turns to be over-confident (small predictive uncertainty) when the training data is small, and gradually becomes prudent (larger predictive uncertainty) when the training data increases. The main reason for this observation is that the predictive uncertainty is dominated by the variation seen in the training data, which is small when small data is observed. The initial learning rates and their scheduling scheme for network parameters $\mathbf{w}$ and noise precision $\beta$ may also play roles here since the uncertainty is determined by the optimum that that stochastic optimization obtained.

The training processes with different datasets are shown in Fig. 24. We can see that the training and test RMSE converges around $50\sim 75$ epochs of training for KLE$50$, and KLE$500$, but the convergence for KLE$4225$ seems to take longer time. The training dataset size is $256$ for all three sets.

To empirically show that using $20$ samples of the Bayesian neural nets for the SVGD algorithm is sufficient in our problem, we show in Fig. 25 the convergence of the test and training RMSE when we vary the number of samples.

Conclusions

In this work, we explore to use an alternative Bayesian surrogate, i.e. a Bayesian neural network, to predict the output fields of a systems governed by stochastic partial differential equations with high-dimensional stochastic input. The approach is built on the highly expressive deep convolutional encoder-decoder networks. The end-to-end image-to-image regression avoids the usual linear dimensionality reduction step, and achieves promising results in terms of predictive performance and uncertainty modeling, even with limited training data.

We show in experiments that one network (DenseED-c16) works well across problems with different intrinsic dimensionality. We believe the performance at the small dataset domain of deep neural networks is due to its unique generalization property which effectively says the over-parameterized deep neural nets lack over-fitting.

The Bayesian model provides uncertainty estimates for our predictions thus accounting for epistemic uncertainty when training with small datasets. We show that Bayesian inference based on SVGD works well for training the Bayesian neural network, and presented results on uncertainty propagation tasks. The uncertainty of the trained model is well-calibrated by investigating the reliability diagrams.

There are three potential directions to improve the regression performance of non-Bayesian surrogate:

Add skip connections between the encoder and decoder path in each level of feature map size. The reason is that the input permiability field directly affects the output velocity field pixel-wise through the equation $\mathbf{u}(\mathbf{s})=-K(\mathbf{s})\nabla p(\mathbf{s})$. Concatenate the features extracted in the encoder path to the decoder path may help to recover the discontinuity of the output velocity fields which are caused by the discontinuity in the input.

Conditional GANs loss. The idea is that we should use a stronger loss function such as an adversarial classifier to enforce the model output to obey the discontinuity in the data. Similar situation has been explored in jet maps in high-energy particle physics .

Using group convolution in the last decoding layer to separate the influences between the pressure field and the velocity fields in the end. This is to address the observation that the prediction performance for the velocity fields is better than for the pressure field using our current network architecture DenseED-c16.

For the Bayesian neural network, the inference task is genuinely difficult since it is asked to find the posterior of millions of random variables based on hundreds of training data. Exploring other priors for the network weights, and using recent advances in variational inference for Bayesian neural networks are definitely valuable directions to pursue.

The image-to-image regression approach can be used to handle the prediction of systems with different source terms, by adding the source field as another channel in the input besides the material property field. For implementation, one only needs to increase the number of input channels by $1$, and leave everything else unchanged.

Our current network does not utilize any information from physics, such as the governing equations or constraints between the three output fields. Incorporating physics information into deep neural networks is a more principled ways to improve the model performance.

Acknowledgments

This work was supported from the University of Notre Dame, the Center for Research Computing (CRC) and the Center for Informatics and Computational Science (CICS). Early developments were supported by the Computer Science and Mathematics Division of ORNL under the DARPA EQUiPS program. N.Z. thanks the Technische Universität München, Institute for Advanced Study for support through a Hans Fisher Senior Fellowship, funded by the German Excellence Initiative and the European Union Seventh Framework Programme under Grant Agreement No. $291763$. The authors acknowledge Dr. Steven Atkinson from the CICS for generating and providing us the Darcy flow datasets.

Appendix A Experiments on Darcy Flow Problem

We optimize the network hyperparameters following the general architecture as shown in Fig. 3. The encoding layer in Fig. 2(a) down-samples its input feature maps by $2$ using convolution of stride $2$ as in the first convolution layer. The decoding layer in Fig. 2(b) up-samples the feature maps by $2$ using transposed convolution instead of pooling layers (as commonly seen in image classification networks) since the location information is critical for regression. The number of down-sampling and up-sampling layers are the same. The datasets are described in Section 4.1. We aim to find one general network architecture that works well across different datasets from both Bayesian and non-Bayesian network models.

Experiment 1 - Spatial dimension of feature maps at the coarsest scale: This is determined by how many down-sampling layers are used. Shrinking the spatial dimension of feature maps can extract high-level or coarse information of the input permeability field, which is subsequently used to predict the output pressure and velocity fields. This design choice is related to a central concept in CNNs called receptive field of an unit within a certain layer, which is the region in the input image that affects this unit feedforward, or the region in the output image that affects this unit when back-propagating gradients. For dense prediction tasks such as our surrogate modeling problems, it is important for each pixel in the code feature maps to have the suitable receptive field in the input and output images. We observe that for the dataset generated with fewer KLE terms, both the input and output velocity fields are smoother, or the output has stronger correlation across pixels, such as KLE$50$ in Fig. 6(a). But for KLE$500$ or KLE$4225$ in Figs. 6(b) and 6(c), the fields vary more rapidly from pixel to pixel, but still there is weak long-range correlation between pixel values.

In our experimental setup we have used no skip connections, the convolution kernels were fixed to be $k7s2p2$ for the first Conv layer, $k3s1p1$ in the layers within dense block, $k1s1p0$ and $k3s2p1$ in the encoding and decoding layers, and $k5s2p2$ in the last ConvtT layer. No dropout was used and the growth rate of dense blocks was taken as $16$.

The loss function is taken as the regularized MSE in Eq. (3), and the validation metric as the $R^{2}$-score in Eq. (23). Adam optimizer is used for training $200$ epochs, with learning rate $0.001$ and a plateau scheduler on the test RMSE. Batch size is always smaller than the number of training data (e.g. $16,32,64$ for dataset with $32,64,128$ data, respectively). Weight decay is set to $0.0005$.

We vary the dense block configurations, i.e. a list of integers specifying the number of layers within each dense block. This list contains odd number of integers because of the symmetry between encoding and decoding paths. For example, blocks $(3,6,3)$ specify the network architecture in Fig. 26(a). $K16L6$ in the second dense block means there are $6$ layers within the dense blocks, and the growth rate for each layer is $16$. In this case, the code dimension (purple feature maps) is $16\times 16$ since there are two down-sampling layers in the encoding path.

The second block configuration is $(12,)$, i.e. only one dense block sits after the first Conv layer, as shown in Fig. 26(b) or $(2,2,4,2,2)$ as shown in Fig. 27.

The validation $R^{2}$-scores of the above three networks DenseED-c16, DenseED-c8, DenseED-c32 on test set ($500$ Monte Carlo samples from each KLE case) are summarized in Tables 2 , 3, 4. From these tables, it can be seen that the $DenseED-c16$ network works well for all three datasets. The network DenseED-c8 does not work well for the KLE$4225$ dataset, since its code dimension is too small. This enforces the output field to be too smooth, which is still favorable to the KLE$50$ dataset. We can also clearly note that the network DenseED-c32 does not work well for KLE$50$, even though it performs well for the KLE$4225$ dataset.

Predictions and learning curve for the selected network design DenseED-c16 are presented in Section 4.3.

Experiment 2 - Hyperparameter optimization with Hyperband: The hyperparameters to select include the ones specifying the network architecture and the training process. We separately optimize those two sets of hyperparameters with one of them fixed. Hyperband is a bandit random search algorithm which has several rounds of successive halving. There are two input for this algorithm, i.e. $R=243$, the maximum iterations for each hyperparameter configuration (here each iteration corresponding to one epoch of training) and $\eta=3$ which specifies to keep $1/3$ of the best configurations after running through all the candidate configurations.

The search space for network architecture with the search space specified in Table 5 with the constraint that the number of parameters being less than $0.25$ million, gives the network architecture presented in Table 1.

The search for the training process hyperparameters as in Table 6 gives approximately the following hyperparameters: initial learning rate $0.002$, weight decay $0.0005$, batch size $16$, and Adam optimizer.

Note that the hyperparameters found by Hyperband are only sub-optimal, but indicate the potential range. The actual hyperparameters used for training are presented in Section 4.3.

Appendix B Generalization behavior

Several numerical experiments were conducted to show the generalization behavior of the network DenseED-c16. The network was trained with different block configurations while keeping the number of training data to be $256$. The number of parameters ranged from $37,892$ to $805,204$. The training and test RMSE are shown in Fig. 28. Clearly, the model is over-parameterized but still shows no overfitting behavior, i.e. the training error does not become smaller and the test (generalization) error does not go higher as we increase the number of parameters. When the number of parameters is less than $200,000$ (still in the over-parameterized regime), there is plenty of room to improve in both training loss and test loss. We configure the number of parameters of the baseline network DenseED-c16 as in Fig. 7 to have $241,164$ parameters, which is favorable according to the generalization curve in Fig. 28.

References