Neural-net-induced Gaussian process regression for function approximation and PDE solution

Introduction

The high expressivity of deep NNs can be combined with the predictive ability of GP regression by taking advantage of the apparent equivalence between GPs and fully-connected, infinitely-wide, deep NNs with random weights and biases. This equivalence was first proved by for shallow NNs. In a subsequent paper, the authors of obtained analytically the covariance functions of the GPs induced by the shallow NNs having two specific activation functions: error and Gaussian functions. More recently, the authors of extended ’s work to the case of deep NNs having generic activation functions, with reference to the mean field theory of signal propagation .

The term NNGP was coined by the authors of , but until now this method despite its potential merits has not been used extensively in applications. In particular, the authors of validated the NNGP on two image recognition datasets: MNIST (handwritten digits data) and CIFAR-10 (color images data), and observed that NNGP outperforms finite width NNs in terms of classification accuracy. The primary objective of our work is to exploit the flexibility of NNGP in order to approximate smooth and discontinuous functions, and furthermore in assessing the effectiveness of NNGP to solve linear and nonlinear PDEs, compared to the GP and deep NN methods formulated recently in . Specifically, so far GP regression has been successfully applied to solution of PDEs including linear and nonlinear forward problems as well as inverse problems . Use of GP regression can bypass discretizing differential operators by properly placing GP priors, and can easily handle noisy data as well as the associated uncertainty propagation in time-dependent problems. Hence, the motivation of the current research is to evaluate if indeed NNGP possesses not only the high expressivity of deep NNs but also the uncertainty quantification property of GPs. To this end, we will perform three-way comparisons between NNGP, GP and deep NN for prototypical problems involving function approximation and PDE solution.

The present paper makes two main contributions.

First, to pursue high expressivity, we increase the number of hyperparameters of NNGP by assuming that the prior variances of weights and biases vary layer by layer and that the weight variances between the input layer and the first hidden layer are different for different neurons in the input layer. We also train the resulting NNGP regression model by using maximum likelihood estimation. We note that the authors of assumed only two hyperparameters: one weight variance and one bias variance, which are kept fixed over different layers; the two hyperparameters are adjusted to achieve the best performance on validation dataset. However, no training was performed in .

Second, we derive the analytical covariance function for the NNGP induced by the deep NN with the error-function nonlinearity. This is motivated by the observation that the analytical covariance function induced by ReLU activated NN, given in , cannot be used to solve the two-dimensional Poisson equation. Denoting by $k(\boldsymbol{x},\boldsymbol{x}^{\prime})$ the analytical covariance function from the ReLU case where $\boldsymbol{x}$ and $\boldsymbol{x}^{\prime}$ are input vectors of length two, we observed that the Laplacian of the covariance function, namely, $\Delta_{\boldsymbol{x}}\Delta_{\boldsymbol{x}^{\prime}}k(\boldsymbol{x},\boldsymbol{x}^{\prime})$ , which has to be computed in order to solve the problem (see Eq.(18)), will tend to infinity as $\boldsymbol{x}\rightarrow\boldsymbol{x}^{\prime}$.

The paper is organized as follows. In Section 2, we review the equivalence between GPs and infinitely wide NNs, which is a cornerstone of NNGP. We change the notations for weight and bias variances compared to the previous work in order to show a different hyperparameter setup. Section 3 introduces the analytically tractable covariance functions induced by the deep NNs with the ReLU and error-function nonlinearities. We employed a version of Gaussian quadrature, proposed in , to validate numerically the two analytical covariance functions. The methodology on GP regression for function approximation and PDE solution is briefly reviewed and summarized in Section 4. Section 5 compares the accuracy of GP, NNGP and NN in function approximation and PDE solution. The uncertainty estimates of GP and NNGP are also presented in the section. Finally, we conclude in Section 5 with a brief summary.

Equivalence between GPs and infinitely wide NNs

It follows from and that for analytical derivation of the covariance function the distribution $\mathcal{D}_{0}$ must be Gaussian. The specific forms of $D_{l},l\geq 1$ do not really matter, since using the Central Limit Theorem one always obtains the Gaussian distribution whatever $\mathcal{D}_{l}$ is.

The $i$-th component of the input of the second hidden layer is computed as

Denoting by $\mathcal{GP}(\mu(\boldsymbol{x}),k(\boldsymbol{x},\boldsymbol{x}^{\prime}))$ the GP with the mean function $\mu(\cdot)$ and the covariance function $k(\cdot,\cdot)$, we learn that $z^{1}_{i}(\boldsymbol{x})\sim\mathcal{GP}(\mu^{1}(\boldsymbol{x}),k^{1}(\boldsymbol{x},\boldsymbol{x}^{\prime}))$. Also, $\{z^{1}_{1}(\boldsymbol{x}),z^{1}_{2}(\boldsymbol{x}),\cdots,z^{1}_{N_{2}}(\boldsymbol{x})\}$ are independent GPs with same mean and covariance functions. The mean function of GP $z^{1}_{i}(\boldsymbol{x})$ is zero due to the zero mean of the distribution $\mathcal{D}_{1}$, and the covariance function is computed as

Two NN-induced covariance kernels of GP

The arguments of the previous section can be extended to deeper layers by induction. For layer $l$ ($\geq 2$), using the multidimensional Central Limit Theorem, we learn that $z^{l}_{i}(\boldsymbol{x})$ and $z^{l-1}_{i}(\boldsymbol{x})$ are both GPs. Specifically, $z^{l}_{i}(\boldsymbol{x})$ has zero mean function and the covariance function computed as

The expectation in Eq. (4) is taken over the GP governing $z^{l-1}_{i}(\boldsymbol{x})$, and this amounts to integrating over the joint distribution of two Gaussian random variables $z^{l-1}_{i}(\boldsymbol{x})$ and $z^{l-1}_{i}(\boldsymbol{x}^{\prime})$ given $\boldsymbol{x}$ and $\boldsymbol{x}^{\prime}$. The covariance function is rewritten as

where the density function of the joint distribution of $z^{l}_{i}(\boldsymbol{x})$ and $z^{l}_{i}(\boldsymbol{x}^{\prime})$ is

Eq. (5) is actually an iteration formula, given by

where the diagonal matrix $\boldsymbol{\Lambda}$ is defined by $\boldsymbol{\Lambda}=diag(\sigma^{2}_{w,0,1},\sigma^{2}_{w,0,2},\cdots,\sigma^{2}_{w,0,d_{in}})$. Note that $k^{l-1}(\boldsymbol{x},\boldsymbol{x}^{\prime})=k^{l-1}(\boldsymbol{x}^{\prime},\boldsymbol{x})$ due to symmetry of the kernel. The trivariate function $F_{\phi}()$ represents the two-fold integral in Eq. (5), which depends on specific form of activation function $\phi(\cdot)$. For certain $\phi(\cdot)$, the function $F_{\phi}()$ is analytically tractable. Two iteration formulas that are analytically derivable will be given in the next two subsections.

Letting the activation function be a rectifier $\phi(x)=\max(x,0)$, we can derive the analytical expression for $F_{\phi}()$, and the corresponding iteration formula is

Here we have increased the number of undetermined variances compared to the original formula in in which $\sigma^{2}_{w,l}=\sigma^{2}_{w}$, $\sigma^{2}_{b,l}=\sigma^{2}_{b}$, and $\boldsymbol{\Lambda}=\frac{\sigma^{2}_{w}}{d_{in}}\boldsymbol{I}$. There are totally $d_{in}+1+2L$ undetermined parameters $\boldsymbol{\theta}$: $\sigma^{2}_{w,0,j}$ ($j=1,2,\cdots,d_{in}$), $\sigma^{2}_{b,0},\sigma^{2}_{w,l}$ ($l=1,2,\cdots,L$), and $\sigma^{2}_{b,l}$ ($l=1,2,\cdots,L$). These parameters (a.k.a. hyperparameters in GP regression) can be learned by using maximum likelihood estimation.

Since the trivariate function $F_{\phi}()$ is not analytically tractable for a generic nonlinearity, the authors of developed an efficient Gaussian quadrature scheme to approximate the two-fold integral in Eq.(5). Fig. 2 (left) compares the numerical results computed by the Gaussian quadrature scheme with the analytical results computed by Eq. (9), showing good agreement.

2 Kernel from the error-function nonlinearity

The error-function nonlinearity can also yield an analytical iteration formula just as the ReLU nonlinearity does. However, the existing formula only holds for single-hidden layer, which is given by . The present paper extends the single-hidden layer case to multi-layer case using the strategy presented in Section 2.3 of . The resulting iteration formula is

To validate the above formula, we compare the analytical results from the above formula with the numerical results from the aforementioned Gaussian quadrature scheme, and as we see in Fig.2 (right) the new iteration formula (10) is correct.

We specifically call the GPs with the kernels that are induced iteratively by NNs “NNGP” for short. NNGP is actually a specific case of GP, but we still distinguish it from the standard GP involving an ad hoc kernel (e.g., square exponential and Matern kernels). An exception is the output of the first hidden layer, $x^{1}_{i}(\boldsymbol{x})$. The output is also a GP and its covariance function is

We still regard the GP with this kernel as a standard GP rather than a NNGP, as the kernel $k(\boldsymbol{x},\boldsymbol{x}^{\prime})$ corresponds to an incomplete, shallow NN without output layer. In Section 5.2.2 we will compare the predictive ability of this kernel with our generic kernel $k^{l}(\boldsymbol{x},\boldsymbol{x}^{\prime})$ computed by formula (10).

GP regression

The essence of GP regression is the conditional distribution of quantities of interest, $\boldsymbol{q}$, given the known observations $\boldsymbol{o}$, where $\boldsymbol{q}$ and $\boldsymbol{o}$ are both Gaussian random vectors. Thanks to the analytical tractability of Gaussian distribution, we can derive analytically mean vector and covariance matrix of the conditional distribution. Supposed that $\boldsymbol{q}$ and $\boldsymbol{o}$ have the joint multivariate Gaussian distribution with zero mean vector and block-wise covariance matrix

the conditional distribution $p(\boldsymbol{q}|\boldsymbol{o})$ is also a multivariate Gaussian distribution with the mean vector

(See (A.5) and (A.6) of ). Zero-mean and $\sigma_{noise}^{2}$-variance Gaussian white noise is assumed in observations.

Note that observations and quantities of interest are usually evaluated at specific time-space coordinates. To preserve a multivariate Gaussian distribution (12) at arbitrary time-space coordinates, the observations and the quantities of interest can be seen as samples of GPs. In other words, $\boldsymbol{o}=[X_{o}(\boldsymbol{x}_{1}),\cdots,X_{o}(\boldsymbol{x}_{N_{o}})]^{T}$ is a sample of $X_{o}(\boldsymbol{x})\sim\mathcal{GP}(0,k_{o}(\boldsymbol{x},\boldsymbol{x}^{\prime}))$ and $\boldsymbol{q}=[X_{q}(\boldsymbol{x}_{1}),\cdots,X_{q}(\boldsymbol{x}_{N_{q}})]^{T}$ is a sample of $X_{q}(\boldsymbol{x})\sim\mathcal{GP}(0,k_{q}(\boldsymbol{x},\boldsymbol{x}^{\prime}))$, where the index set $\boldsymbol{x}$ represents arbitrary time and/or space coordinates. Particularly, for solving PDEs, the observations $\boldsymbol{o}$ consist of samples from different GPs. Different covariance function $k_{o}(\cdot,\cdot)$ will capture different trends and features that the observations $\boldsymbol{o}$ exhibit.

Before computing mean and covariance of conditional distribution from formulas (13) and (14), we need to first learn the parameters hidden in covariance function as well as the noise variance $\sigma_{noise}^{2}$. These parameters are also known as hyperparameters and can be learned by using maximum likelihood estimation (MLE) . In practice, we find the optimal hyperparameters (including undetermined parameters $\boldsymbol{\theta}$ in covariance function and the noise variance $\sigma_{noise}^{2}$) that minimize the negative log marginal-likelihood

where $|\cdot|$ denotes the determinant and the covariance matrix for observations (or training data) $\boldsymbol{K}_{oo}$ depends on $\boldsymbol{\theta}$ and $\sigma^{2}_{noise}$.

In what follows, we will briefly review that different formulations of observations $\boldsymbol{o}$ will enable the GP regression to solve different problems including function approximation and PDE solution.

Let $\boldsymbol{q}=[f(\boldsymbol{x})]$ and $\boldsymbol{o}=[f(\boldsymbol{x}_{1})+\epsilon,\cdots,f(\boldsymbol{x}_{N_{o}})+\epsilon]^{T}$ where $\epsilon$ is Gaussian white noise having variance $\sigma^{2}_{noise}$. The objective is to approximate the unknown scalar-valued function $f(\boldsymbol{x})$ given the observations $\boldsymbol{o}$. After prescribing the covariance function $k_{o}(\cdot,\cdot)$ for computing the covariance matrix $\boldsymbol{K}_{oo}$ and then training the regression model by the MLE, we derive the function approximation by using formula (13). Moreover, the use of the variance computed from (14) yields the confidence interval. Denoting by $\boldsymbol{X}$ ($N_{o}$ by $d_{in}$ matrix) the collection of the evaluation points $\{\boldsymbol{x}_{j}\}$ (usually called training inputs), we formulate the covariance matrix in (12) by:

2 PDE solution

First we consider the following time-independent linear PDE

where $L_{\boldsymbol{x}}$ is a linear differential operator. The source term $f(\cdot)$ and the boundary condition term $g(\cdot)$ are the only information we know, and thus we put them in the observations by letting $\boldsymbol{o}=[\boldsymbol{o}_{u}^{T},\boldsymbol{o}_{f}^{T}]^{T}=[g(\boldsymbol{x}_{u,1})+\epsilon_{u},\cdots,g(\boldsymbol{x}_{u,N_{u}})+\epsilon_{u},f(\boldsymbol{x}_{f,1})+\epsilon_{f},\cdots,f(\boldsymbol{x}_{f,N_{f}})+\epsilon_{f}]^{T}$, where we select $N_{u}$ boundary points as the evaluation points for $\boldsymbol{o}_{u}$ and $N_{f}$ domain points as the evaluation points for $\boldsymbol{o}_{f}$. The variances of the Gaussian white noises $\epsilon_{u}$ and $\epsilon_{f}$ are $\sigma^{2}_{u}$ and $\sigma^{2}_{f}$, respectively. We assume that the solution $\boldsymbol{q}=[u(\boldsymbol{x})]$ is a GP with zero mean and covariance function $k_{u}(\boldsymbol{x},\boldsymbol{x}^{\prime})$, and the source term is another GP that depends on $u$, whose covariance function is written as

Denoting by $\boldsymbol{X}_{u}$ ($N_{u}$ by $d_{in}$ matrix) the collection of the boundary evaluation points $\{\boldsymbol{x}_{u,j}\}$ and by $\boldsymbol{X}_{f}$ ($N_{f}$ by $d_{in}$ matrix) the collection of domain evaluation points $\{\boldsymbol{x}_{f,j}\}$, we formulate the covariance matrix in (12) as

For time-dependent problem, the above GP regression model also works as we can regard the time variable as the $(d_{in}+1)$-th dimension of space-time domain (See Section 4.3.1 of ).

The alternative way to handle the time-dependent problems is numerical GP regression . Consider the following time-dependent PDE

Generally, we let the covariance function be the same for all $n$, namely $k_{u^{n}}(\boldsymbol{x},\boldsymbol{x}^{\prime})=k_{o}(\boldsymbol{x},\boldsymbol{x}^{\prime};\boldsymbol{\theta}_{n})$, except with varied undetermined parameters $\boldsymbol{\theta}_{n}$. Considering the propagation of the uncertainty in time marching, we need to take into account the uncertainty of the predictions in previous time step (i.e., $u^{n-1}(\boldsymbol{x})$ predicted at $\boldsymbol{X}^{n-1}$). We rewrite the predictive or posterior variance in (14) as

Numerical results

We perform a three-way comparison between NNGP, NN, and GP for function approximation and PDE solution for prototypical problems. The hyperparameters of NNGP and GP are trained by minimizing the negative log marginal-likelihood function. The conjugate gradient algorithm (see the function minimize() in GPML Matlab code (version 3.6) ) is employed for the optimization. For NN, network parameters (weights and biases) are trained by the L-BFGS-B algorithm that is provided in the TensorFlow package.

Because of the non-convex nature of the objective function, to avoid trapping into the local minima, we run our optimization algorithm code ten times with different initializations. Among the ten groups of optimized hyperparameters, for GP/NNGP we choose the one yielding the smallest negative log marginal-likelihood, and for NN we select the one producing the lowest loss for NN. The initializations of hyperparameters for GP/NNGP are taken as the first ten entries of the Halton quasi-random sequence and these initializations are kept deterministic. The initializations of network weights are obtained from Xavier’s initialization algorithm , which is provided by the TensorFlow package. Also, for each initialization at most 200 function evaluations are allowed in conjugate gradient search in GP/NNGP. For NNGP, the central finite difference scheme with the step size $10^{-4}$ is used in computing the gradients of covariance function with respect to hyperparameters while for GP the analytical derivation is employed. Additionally, we compute $l_{2}$ relative error $||\boldsymbol{u}_{approximate}-\boldsymbol{u}_{exact}||_{2}/||\boldsymbol{u}_{exact}||_{2}$ in quantifying accuracy of function approximation and PDE solution, where $\boldsymbol{u}$ is a vector formed by the function values or PDE’s solutions evaluated at test points.

The covariance matrices of GP/NNGP are formulated according to (16).

We approximate the step function $f(x)=1,x\geq 0;f(x)=0$, otherwise. Ten evenly distributed training inputs and 100 equispaced test points are chosen. Fig. 3 shows the approximation by the GPs with two commonly used kernels: squared exponential (SE) and Matern . The first kernel can describe well the data exhibiting smooth trend while the second one is suitable for finite regularity. However for the step function, a non-smooth function, the performance of these kernels is poor. The NNGP predictions shown in Figs. 4 and 5 are much better. In the domain away from the discontinuity, the NNGP approximation is good with small uncertainty. Particularly, the ReLU-induced kernel even succeeds in capturing the discontinuity. An important observation is that the ReLU induced kernel is more suitable for non-smooth functions than the error-function induced kernel. Another observation regarding NNGP is that deepening the inducing NN does not change the approximation accuracy for the step function. Figs. 6 and 7 show the predictions of NN, which are less accurate than NNGP’s. Unlike GP/NNGP, NN does not quantify any uncertainty for approximation.

1.2 Hartmann 3D function

The Hartmann function is frequently used for testing global optimization algorithms. Here we consider the trivariate case. The function is much smoother than the step function and thus we can expect GP to perform well. To test the approximation accuracy, we first generate $N$ ($N=100,200,500$, and $1000$) points by choosing the first $N$ entries of the Halton sequence, and then randomly permutate the points, followed by selecting the first 70% points as training points and the remaining 30% points as the test points.

Fig. 8 compares NN, NNGP, and GP in terms of approximation accuracy. Since the Hartmann 3D function is smooth, the NNGP with error-function induced kernel and the GP with squared exponential kernel are both very accurate. ReLU is not suitable for smooth function in contrast to approximating the non-smooth step function. Additionally, the GP and NNGP both outperform NN. It is observed again that NN does not give any uncertainty estimate, whereas GP and NNGP do. Another observation is that increasing the depth of the inducing NN in NNGP does not change the accuracy.

2 PDE solution

We use the proposed covariance function from error-function nonlinearity to solve the following two PDEs. We replace the covariance function $k_{u}(\boldsymbol{x},\boldsymbol{x}^{\prime})$ in Eq. (19) as well as $k_{u^{n}}(\boldsymbol{x},\boldsymbol{x}^{\prime})$ in Eq. (21) with the new kernel (10). The derivation of kernel’s derivatives for the case of Poisson equation is given in Appendix.

with two fabricated solutions (1) $\sin(\pi x)(y^{2}+\exp(-y))$ and (2) $\sin(\pi x)\cos(2\pi(y^{2}+x))$. Dirichlet boundary conditions are assumed. The second solution is more complex than the first one and thus we use more training points in the second solution case. The training inputs $\boldsymbol{X}_{f}$ are chosen from the first $N_{f}$ entries of the Halton sequence and $\boldsymbol{X}_{u}$ is equispaced on the boundary, which are shown in Fig. 9.

The covariance matrices of GP/NNGP are obtained by (19). Figs. 10 and 11 give the error plots of NN, NNGP, and GP for fabricated solution 1 and 2, respectively. Totally 441 evenly distributed test points are taken to evaluate the relative error. The NN results are obtained according to the approach proposed in . To keep a fair comparison, the same training inputs are selected in the NN case. It is observed from the figures that for different fabricated solutions, the approximation accuracy of the GP and NNGP is comparable. Also, the NN accuracy is nearly one order of magnitude lower than that of GP/NNGP.

Fig. 12 shows the uncertainty estimates of GP and one-layer NNGP evaluated on the cut line $y=x$ of the square domain, for the case of fabricated solution 2. We see that for both methods uncertainty is reduced with increasing number of training points. Moreover, the uncertainty is strongly correlated with the prediction error as smaller uncertainty corresponds to lower error. Note that the uncertainty at the endpoints of the cut line, namely, $(x,y)=(0,0)$ and $(x,y)=(1,1)$, is exactly zero, due to the boundary condition.

2.2 1D Burgers’ equation

with $u(-1,t)=u(1,t)=0$. The initial condition is $u(x,0)=-\sin(\pi x)$. After the linearization through replacing the nonlinear term $u^{n}\frac{\partial u^{n}}{\partial x}$ by $\mu^{n-1}\frac{\partial u^{n}}{\partial x}$, where $\mu^{n}$ is the mean vector computed for the previous time step by (13), we derive the differential operator $L_{x}=\frac{0.01}{\pi}\frac{\partial^{2}}{\partial x^{2}}-\mu^{n-1}\frac{\partial}{\partial x}$. The covariance matrices of GP/NNGP are formulated by (21).

The time step size is fixed to be $\Delta t=0.01$. To evaluate the relative error at each time step, we place 400 equispaced test points in the space domain $$. We also solve the same equation using the NN approach proposed in . One main difference between the NN approach and the numerical GP/NNGP regression is the sampling strategy for training inputs. For NN, we sample the training inputs in the time-space domain$(x,t)\in^{2}$, whereas for GP/NNGP, we sample the training inputs merely in the space domain and then perform time-marching. To make the comparison as fair as possible, we sample 10000 training points for NN, as in GP/NNGP we have at most 100 (number of time step)$\times$ 100 (number of training inputs in space) = 10000 (number of time-space sampling points). The Latin hypercube sampling strategy is adopted in sampling the training inputs in NN, GP, and NNGP. To accelerate training of numerical GP/NNGP, the initial guess of hyperparameters for current time step is taken as the optimized hyperparameters attained at previous time step. This strategy is reasonable since for a small time step two successive GPs are highly correlated and therefore the respective hyperparameters should be close.

The kernel (11), rather than the SE considered in solving Poisson equation, is taken in GP. We choose a non-stationary kernel instead of a stationary one, because the solution exhibits discontinuity that cannot be well captured by stationary kernel. Additionally, the NN with four hidden layers of 40-unit width and the hyperbolic tangent nonlinearity is employed in the NN approach.

We first consider the case where the initial condition is noise-free. Figs. 13, 14, and 15 show the numerical solutions computed by the GP, the NNGP with one-hidden-layer ($L=1$), and the NNGP with three-hidden-layer ($L=3$), respectively. We can see that the NNGP has higher solution accuracy than the GP. Importantly, different from previous examples, increasing the depth of inducing NN in NNGP improves the results in the current example. The high accuracy of NNGP is attributed to a larger number of hyperparameters compared to the GP case. More hyperparameters presumably implies higher expressivity for function approximation, which means that we can use the kernel to approximate a wider spectrum of functions. For simple function or solution (as in the Hartmann 3D function and Poisson equation examples), the advantage of higher expressivity is not fully realized, but for complex ones, such as the step function and the Burgers’ equation here, the advantage can be clearly seen.

The NN results are shown in Fig.16. Although NN can derive sufficiently accurate solutions, it does not give any uncertainty estimate. Fig. 17 collects the errors from GP, NNGP, and NN for noise-free case. The NNGP and NN achieve errors of the same order of magnitude, while the GP performs worst. This is because NNGP and NN both have a large set of parameters to be tuned and thus possess higher expressivity.

Next we consider the case where the initial condition is contaminated by Gaussian white noise of zero mean and variance $0.15^{2}$. The NN approach for the noise case is beyond the scope of the present paper, since without proper regularization methods (such as dropout , early stopping, and weighted L1/L2 penalty), NN will easily encounter overfitting. Unlike NN, GP and NNGP inherently include the model complexity penalty in the negative log-marginal likelihood and have less risk in overfitting (see the discussion in Section 5.4.1 of the book ).

Numerical solutions computed by the GP and NNGP are ploted in Figs. 18, 19, and 20. NNGP still has higher solution accuracy than the GP due to the higher expressivity. Data noise amplifies uncertainty represented by the orange shadowed region. GP and NNGP can handle noise well, because noise variance can be directly learned from the training data and the corresponding uncertainty is quantified by the conditional (or posterior) distribution. Fig. 21 compares the solution accuracy of GP and NNGP. For long-term simulations, the accuracy of the NNGP is roughly one order higher than that of GP.

For NNGP, increasing the depth, $L$, of inducing NN does not guarantee the increase of accuracy. For example, one-layer NNGP with 101 training points outperforms three-layer NNGP with 101 training points, as is shown in Fig. 21. A larger $L$ amounts to a larger number of hyperparameters, which merely increases the possibility in fitting complex functions better. Actually, $L$ can also be seen as another hyperparameter of NNGP.

Summary

A larger number of hyperparameters enables NNGP to achieve higher or comparable accuracy for both smooth and non-smooth functions in comparison with GP, which can be seen from Table 1. The deep NN that induces NNGP contributes its prior variances of network parameters (weights and biases) as well as its depth to the hyperparameter list of NNGP, which endows NNGP with high expressivity. On the other hand, NNGP is able to estimate uncertainty of predictions, which is crucial to noisy-data handling and active learning . Unlike NN, Bayesian NN can provide uncertainty estimate. However, the conventional Bayesian NN could be time-consuming to train due to approximation of a high-dimensional integral over weight space.

Due to the need for inverting the covariance matrix, NNGP has cubic time complexity for training (see Table 2). In Table 2, for GP and NNGP $N_{training}$ is the size of training set, and $m_{GP}$ and $m_{NNGP}$ are numbers of evaluations of negative log marginal-likelihood in conjugate gradient algorithm for GP and NNGP, respectively. For NN, the accurate estimate of time complexity for training is still an open question . Generally, the training of a fully-connected NN is faster than that of GP, because one does not need to invert a matrix. For each training point, the forward and backward propagation only need linear cost, namely $O(N_{weight})$, where $N_{weight}$ is the total number of weights in the network. $N_{training}$ for NN means batch size; in this paper we fed the whole training set to optimization algorithm and thus the batch size is exactly the size of training set. $m_{NN}$ is number of iterations in optimization algorithm. In numerical examples of this paper, training set size does not exceed 1000. However, for very large datasets, GP and NNGP will be less attractive compared to NN. In future work, we intend to leverage scalable GPs recently developed in to tackle large dataset problems.

Acknowledgements

We thank Mr. Dongkun Zhang and Dr. Maziar Raissi for useful discussions. This work was supported by AFOSR (FA9550-17-1-0013) and NSF (DMS-1736088). The first author was also supported by the National Natural Science Foundation of China (11701025). The second author was also supported by a gift from Swiss company BH-Robotics.

References

Appendix: Derivatives of NNGP kernel from the error-function nonlinearity

Absorbing the constant coefficients $\frac{2}{\pi}$ and $2$ into the variances, the kernel derived from the error-function nonlinearity, namely (10), can be further simplified to

To solve PDEs, we need to compute the derivatives of the kernel $k^{l}(\boldsymbol{x},\boldsymbol{x}^{\prime})$. We take the 2D Poisson equation as an example. We need to know the explicit forms of $-\left(\frac{\partial^{2}}{\partial x^{2}}+\frac{\partial^{2}}{\partial y^{2}}\right)k^{l}(\boldsymbol{x},\boldsymbol{x}^{\prime})$ and $\left(\frac{\partial^{2}}{\partial x^{2}}+\frac{\partial^{2}}{\partial y^{2}}\right)\left(\frac{\partial^{2}}{\partial^{2}x^{\prime}}+\frac{\partial^{2}}{\partial^{2}y^{\prime}}\right)k^{l}(\boldsymbol{x},\boldsymbol{x}^{\prime})$ where $\boldsymbol{x}=[x,y]^{T}$ and $\boldsymbol{x}^{\prime}=[x^{\prime},y^{\prime}]^{T}$. The partial derivatives up to fourth order needs to be derived. Denoting by the trivariate function $F_{\phi}()$ the $\arcsin()$ term in the above iteration formula, the use of chain rule yields

Generally, the iteration formulas for kernel and its derivatives can be written as

where $i,j\in{0,1,2}$ and $F_{0,0}=F_{\phi}$. The initial values for the iteration are the known kernel $k^{0}$ and its derivatives. It should be noted that the explicit form of the function $F_{i,j}$ will become rather lengthy and ugly for higher order derivatives, but fortunately, we have derived for the readers these lengthy formulas using Maple. The Matlab code computing the covariance functions and its partial derivatives up to fourth order can be downloaded at \urlhttps://github.com/Pang1987/nngp_kernel_derivative.