Galerkin meets Laplace: Fast uncertainty estimation in neural PDEs

Authors: christian jimenez beltran
antonio  vergari
aretha l teckentrup
konstantinos c zygalakis
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 19:52:00

Abstract: The solution of partial differential equations (PDEs) by deep neural networks
trained to satisfy the differential operator has become increasingly popular. While
these approaches can lead to very accurate approximations, they tend to be over-
confident and fail to capture the uncertainty around the approximation. In this
work, we propose a Bayesian treatment to the deep Galerkin method (Sirignano &
Spiliopoulos, 2018), a popular neural approach for solving parametric PDEs. In
particular, we reinterpret the deep Galerkin method as the maximum a posteriori
estimator corresponding to a likelihood term over a fictitious dataset, leading thus
to a natural definition of a posterior. Then, we propose to model such posterior via
the Laplace approximation, a fast approximation that allows us to capture mean-
ingful uncertainty in out of domain interpolation of the PDE solution and in low
data regimes with little overhead, as shown in our preliminary experiments.

Main paper: 
 introduction: 
 Partial differential equation (PDE) models appear in numerous areas of science and engineering, including geophysics, climate modelling and medical imaging to name a few (De Marsily, 1986; Isaacson et al., 2004; Schneider et al., 2017) . In the last few years, deep learning approaches that use neural networks to approximate the solution of PDEs have shown remarkable success in terms of approximating certain classes of high-dimensional PDEs (Raissi et al., 2019; Kelshaw et al., 2022) , thus providing fast alternatives to classical numerical methods such as finite differences, elements, volumes and spectral methods (Larsson & Thomée, 2003) .One of the first deep learning approaches designed for solving PDEs is the physics-informed neural networks (PINNs) (Raissi et al., 2019) . PINNs define a suitable loss function that takes explicitly into account the PDE in question over specific collocation points, as well as over specific initial and boundary data. To do so, PINNs require fixed collocation points at training time, an aspect that can become problematic for high-dimensional PDEs. In addition, in many applications, one is interested in solving the PDE for a range of parameters, which is something that cannot be addressed by PINNs. The deep Galerkin method (DGM) (Sirignano & Spiliopoulos, 2018) can deal with these two issues by averaging oversampled collocation points and initial and boundary data during training and adding the PDE parameters as additional inputs to the neural network.While these neural approaches can scale PDE solving to high dimensions, they still require large amounts of data to be trained and tend to be highly overconfident even when their approximations are far from the ground truth. This is especially problematic when they are used as surrogate solutions on collocations unseen during training (out-of-sample) and when trying to perform statistical analysis of real-world differential equation models, as it is important to have calibrated estimates of uncertainty with regards to our numerical approximations (Hennig et al., 2015; Conrad et al., 2017) .In this work, we address this issue for DGMs. We start by reinterpreting the DGM as the maximum a posteriori estimator corresponding to a likelihood term over a fictitious data set, leading thus to a natural Bayesian treatment. Then, we propose a simple and effective solution: approximating the Figure 1 : Our DEEPGALA extends the DGM to provide meaningful uncertainty on out-ofsample input locations (bottom) for the PDE solution of (PDE1) for θ = 0.001 and when trained over 1000 (second column) and 3000 (third column) samples respectively. Monte Carlo Dropout (MCD) struggles to approximate the solution even in-sample.corresponding posterior via Laplace's method. This allows us to deliver fast uncertainty estimation over the PDEs that can be meaningful when training data is scarce or we evaluate our surrogate solution on out-of-sample collocations (see Figure 1 and Figure 2 ). 
the deep galerkin method: 
 We consider a general parametric PDE (see e.g. Deveney et al. (2023) ), given byA(x, u(x; θ); θ A ) = h(x; θ h ), u(∂x; θ) = b(∂x; θ b ), x ∈ Ω, ∂x ∈ ∂Ω, θ ∈ Θ (1)where the domain Ω ⊂ R d is the input space with boundary ∂Ω, and Θ ⊂ R p is the parameter space for θ = (θ A , θ h , θ b ). The differential operator A is parameterized by θ A , whereas the functions h and b are parameterized by θ h and θ b , respectively. We denote by u(x; θ) the solution of the parametric PDE. Note that A could be non-linear, and that time could be a component of the input x. We want to use a neural network (NN) to approximate u(x; θ). To this end, consider a NN f W with L layers and parameterized by W = {W 1 , . . . , W L }. 1 The NN takes as inputs a point x ∈ Ω in the closure of the input domain and a parameter value θ ∈ Θ, and it is trained to approximate the solution of the PDE, f W (x; θ) ≈ u(x; θ).The deep Galerkin method (DGM) (Sirignano & Spiliopoulos, 2018 ) uses a mesh-free strategy: At each iteration of the training step, K collocations for x i , ∂x i , and θ i are sampled from densities π p , π b , and π θ , respectively. Then, the NN f W is trained to minimize the following loss:EQUATIONwhereD = {D i } K i=1 = {D h,i } K i=1 ∪ {D b,i } K i=1 , and D h,i = {x i , θ i } and D b,i = {∂x i , θ i }. Here, ℓ h (D h,i ; W) = (A(x i , f W (x i , θ i ); θ A,i ) -h(x i ; θ h,i ))2 measures the error in the approximation of the differential operator andℓ b (D b,i ; W) = (f W (∂x i , θ i ) -b(∂x i ; θ b,i ))2 measures the error in the boundary conditions. Optimization is carried out by gradient-based approaches, e.g., SGD, possibly increasing the number of samples K and the number of iterations. In this work, we followed a two-step training as suggested by He et al. (2020) . Although conceptually simple, there are no convergence guarantees for DGM, and the error in predictions made by the NN f W could be significant. In applications where the NN is then used in a computational pipeline, such as using NNs to approximate the likelihood in an inverse problem (Deveney et al., 2023) , it becomes crucial to quantify the uncertainty in the predictions of the NN to avoid overconfident and biased inference. 
deepgala: laplace approximation for bayesian dgm: 
 In this section, we explain how to extend the DGM to allow for uncertainty quantification in its predictions, using Bayesian theory and the Laplace approximation. Uncertainty quantification has been thoroughly investigated for NNs as function approximators (Wilson & Izmailov, 2020; Wilson, 2020) and more recently also for neural PDE solvers such as PINNS and DeepOnets (Psaros et al., 2023) . To the best of our knowledge, this has not been done for neural parametric PDE solvers such as DGM yet. Moreover, our approach emphasises more precisely assessing the uncertainty resulting from limited sample numbers, training methods, and NN hyperparameters.In a Bayesian version of DGM, we are interested not in a single parameter configuration W, but in its posterior distribution p(W | D). To this end, we first interpret the data D to be the fictious dataset built by K collocations points sampled from π p , π b , and π θ . Then, we interpret the loss in equation 2 as a negative log-likelihood by considering Gaussian noise around the fictitious data observed at the collocation points. By doing this, the likelihood p(D|W) characterises how well a choice of weights W approximates the PDE solution u(x; θ), and is given byEQUATION)where we assume that the likelihoods of the collocations D h and D b are conditionally independent given W. We further introduce a Gaussian prior distribution p(W ). By Bayes' rule, we then have the posterior distribution p(W |D) ∝ p(D|W)p(W ) on the weights W given the fictitious training data D. Note that although not explicitly indicated, the posterior p(W |D) is also conditioned on A, h and b from equation 1, since these are used to define the likelihood. While intractable in general, p(W |D) is amenable to a fast and and effective approximation, as discussed next.We propose a Laplace Approximation (LA) of the posterior of our Bayesian DGM, which we call DEEPGALA-Deep Galerkin via Laplace. The LA is a well-known method for approximating intractable posterior distributions as a Gaussian distribution centred on the maximum a posteriori (MAP) solution W MAP , i.e.,EQUATIONwhere the negative log-prior r(W ) = -log p(W ) is an L2 regularizer (weight decay) in our experiments. Despite its simplicity (assuming unimodality of the posterior), the LA has been demonstrated to be a solid alternative to more complex approximations such as MCMC and VI on a number of uncertainty quantification tasks for Bayesian deep learning (Daxberger* et al., 2021) . To compute the LA, the first step is to find the minimizer W MAP of equation 4, which can be achieved by training f W via gradient-based optimization. The second step is to fit the local Gaussian distribution, namely N (W ; W MAP , Λ) where Λ -1 is the Hessian of the negative log-posterior evaluated at W MAP . In our case, r(W ) is a weight decay regularizer, which corresponds to a Gaussian prior distribution p(W ) = N (W ; 0, γ2 I) (Daxberger* et al., 2021) , Λ -1 taking the formEQUATIONExactly computing Λ -1 above can be computationally demanding, as W in modern NNs can be very large (Nilsen et al., 2019) . To this end, we employ a fast approximation of the covariance in DEEP-GALA, by leveraging a number of heuristics from the modern Bayesian deep learning literature. First, we consider only the weights W L of the last layer of f , which is generally enough to deliver good uncertainty estimates as noted in Sharma et al. (2023) . 2 Second, we subsample a subset K ′ of all the training points seen by f during training. Third, we either compute the full Hessian or approximate the Hessian by its diagonal (Pearlmutter, 1994) . Appendix A details the whole process.  
experiments and discussion: 
 We aim to answer the following research questions: RQ1) does DEEPGALA capture meaningful uncertainty over out-of-sample inputs x? RQ2) does increasing training size properly reduce the epistemic uncertainty highlighted by DEEPGALA?PDEs. To this end, we consider two PDEs defined over x = (y, t): A heat equation with external heat source:EQUATIONwith conditions u(y, 0) = 4 sin(3πy)+9 sin(7πy), and u(0, t) = u(1, t) = 0; and Burgers EquationEQUATIONwith conditions u(y, 0) = -sin(πy), and u(-1, t) = u(1, t) = 0. PDE1 admits analytic solutions, allowing us to directly compare the performance of DEEPGALA both in terms of accuracy and uncertainty. For PDE2, we use the results obtained by the work of Raissi et al. (2019) for θ = 0.01/π as a baseline.The experiments were carried out in a machine with 12th Gen Intel Core i7-1265Ux12, 32 GB of RAM and Ubuntu 22.04 operating system.NNs models. For both PDEs we use a feed-forward NN with three layers and 40 neurons using hyperbolic tangent non-linearities, trained with Adam for 1,200 epochs, a learning rate of 0.01 and regularizer, γ, of 0.0015 for PDE1 and 0.0001 for PDE2, followed by the L-BFGS optimizer in a two-step optimisation, as recommended by He et al. (2020) . We show results for the diagonal approximation in Appendix B. The mean time to fit and call the full Hessian LA for PDE1 and PDE2 was 0.081 and 0.003, and 0.077 and 0.0044 seconds respectively.Alternative UQ baselines. We test our approach against Deep Ensemble (DE) (Lakshminarayanan et al., 2017b) and Monte Carlo Dropout (MCD) (Lakshminarayanan et al., 2017a) . In order to calculate the mean and variance for the MCD, 1,000 samples were used, with the dropout layer set to p = 0.05 for PDE1 and p = 0.1 for PDE2, respectively. For the PDE2, the NN architecture stays the same, however for the PDE1, five hidden layers comprising forty neurons were employed. The reason for the size modification was that smaller NN were not producing satisfactory results for this specific PDE. Five NN's of three layers and forty neurons were trained to approximate the PDE2 using the DE method in accordance with Lakshminarayanan et al. (2017a) . Nevertheless, obtaining satisfactory results for the PDE1 was surprisingly challenging; this difficulty may have arisen from the PDE's initial condition. That's the reason why we don't present the results utilizing this baselines for the PDE1.In-vs out-of-sample. For both PDEs, we provide DEEPGALA only a portion of the input at training time (in-sample), expecting it to extrapolate poorly -but delivering good uncertainty -on out-of-sample inputs. In particular, we draw our in-sample collocations from the following intervals:y i ∼ U[0, 1], t i ∼ U[0, 1], θ i ∼ U[0.0001, 0.05] for PDE1 and y i ∼ U[-1, 1], t i ∼ U[0, 0.5], θ i ∼ U[0.0001, 0.05] for PDE2, where U denotes the uniform distribution. The top (insample) and bottom (out-of-sample) sides of Figure 1 and Figure 2 show two uncertainty surfaces representing 95% and 68% confidence intervals around the MAP estimate for DEEPGALA for PDE1 and PDE2 respectively. Reassuringly, these uncertainties increase when the model makes more mistakes, highlighting how DEEPGALA can be used to indicate that the NN are "aware" that they may be forecasting incorrectly. In contrast, the uncertainty exhibits no specific behaviour when examined both in-and out-of-sample for the methods MCD and DE, nor do they provide good approximations at the mean.) training size effect. We train DEEPGALA over sample sizes of 1,000, and 3,000 for PDE1 and 500, and 1,500 for PDE2, plotting the results in the top and bottom halves of Figure 1 and Figure 2 . All models trained with more samples outperform the other (both for in-sample and out-sample values of t). Nonetheless, the exact solution is within the uncertainty range of the models for this particular value of t. Again, DEEPGALA is shown to be an effective way to estimate the epistemic uncertainty of NNs for PDE solving, which can be useful in quantifying how much data can improve performance. As such, we plan to use the uncertainty DEEPGALA provides in an active learning loop to select what are the most promising collocation to improve accuracy of the PDE solution. 
conclusion and discussion: 
 We proposed DEEPGALA as a fast and effective way to estimate uncertainty in deep neural PDE solvers, and tested it on preliminary benchmarks with encouraging results. We plan to evaluate DEEPGALA to more complex PDEs, extending it with recent advancements in Bayesian deep learning (Wilson, 2020) and comparing to other Bayesian PDE solvers such as the ones presented in Psaros et al. (2023) , which however are limited to non-parametric PDEs. We hope that scaling uncertainty estimation can lead to a renewed interest for efficient probabilistic numerics for complex neural PDEs in the real-world. 
 
 Appendix: Following Sharma et al. (2023) , we consider partially stochastic networks and introduce uncertainty only in W L , the parameters in the last layer of f W . The last layer, involving only an affine transformation, can be represented asf W (x, θ) = W L • [1, ϕ(y L-1 (x, θ))], where W L is the vector of weights and biases, and ϕ(y L-1 (x, θ)) is the vector of outputs of the penultimate layer for input (x, θ). Let's just focus on calculating the Hessian for one input z i = (x i , θ i ) and compute ∇ 2 W L ℓ(z i ; W), where we will use the abbreviation l i and f i,W to denote ℓ(z i ; W ) and f W (z i ). Therefore,∇ 2 W L l i = ∇ W        ∂li ∂f i,W ∂f i,W ∂w 0 L ∂li ∂f i,W ∂f i,W ∂w 1 L . . . ∂li ∂f i,W ∂f i,W ∂w K L        =      ∂ 2 li ∂f 2 i,W ∂f i,W ∂w 0 L ∂f i,W ∂w 0 L . . . ∂ 2 li ∂f 2 i,W ∂f i,W ∂w 0 L ∂f i,W ∂w K L . . . . . . . . . ∂ 2 li ∂f 2 i,W ∂f i,W ∂w K L ∂f i,W ∂w 0 L . . . ∂ 2 li ∂f 2 i,W ∂f i,W ∂w K L ∂f i,W ∂w K L      +      ∂li ∂f i,W ∂ 2 f i,W ∂w 02 L . . . ∂li ∂f i,W ∂ 2 f i,W ∂w 0 L ∂w K L . . . . . . . . . ∂li ∂f i,W ∂ 2 f i,W ∂w K L ∂w 0 L . . . ∂li ∂f i,W ∂ 2 f i,W ∂w K2 L      . (6)In the special case of a last layer with only an affine transformation, the second term in the previous equation becomes zero. Additionally, the terms∂f i,W ∂w k L = ϕ k (y L-1 (z i )) and ∂f i,W ∂w 0 L = 1.Consequently, we can reformulate the previous equation as:EQUATIONAnother observation is that the last equation is akin to the Gauss-Newton Matrix (GNN), up to a constant G(W ) = 1 N N n=1 J W (f W (z n )) T ∇ 2 f l(f W (z n ))J W (f W (z n )).The GNN is a first-order Taylor approximation of the loss function. We can compute the equation 5 by using the next equation:EQUATIONwhereJ W L (f W (z i )) = [ ∂f W (zi) ∂w 0 L , ..., ∂f W (zi) ∂w c L ]T . Note that equation 8 represents a full Hessian; for a big enough NN, we can approximate even further to obtain a diagonal approximation by using the Hadamard product. Thus, the Diagonal approximation of the Hessian is computed in the following way:EQUATIONA.1 POSTERIOR PREDICTIVE Finally, in order to perform predictions, one needs to compute the posterior predictive distribution. This is done in the following way. Given our focus on the last layer of f W , linearity in the weights W is observed. Therefore, for a test input ẑi of the neural network, we obtain p(f * |f W (ẑ i ), D) = N (f * ; f W M AP (ẑ i ), ϕ(y zi,L-1 ) T Λϕ(y zi,L-1 )), where f WMAP is the output of the neural network at the MAP. In the case of a diagonal Hessian, Λ is easily obtained by computing the inverse of each value on the diagonal. For a full Hessian, the process involves first computing the lower triangular decomposition, where Λ -1 = L T L, then finding the inverse of L, and finally obtaining Λ = L -1 (L -1 ) T .B DIAGONAL HESSIAN UNCERTAINTY Figures 3 and 4 depict the results of DEEPGALA when applying a diagonal approximation to the Hessian to quantify uncertainty for the solutions of PDE1 and PDE2. The mean time to fit and call the diagonal LA for the PDE1 and PDE was 0.07 and 0.004, and 0.067 and 0.06 seconds respectively. As we can see, the uncertainty estimation has been reduced; yet, our method continues to be a good tool to predict how much performance may be improved. However, we still need to understand and quantify how much uncertainty estimation changes when utilising full or diagonal Hessian. We consider the biases to be included in each W l . Note that other non-sequential architectures are possible.This also allows us to compute the posterior predictive in closed-form, see Appendix A.1 for details. 


Decision: This paper proposes to add uncertainty quantification, which is very important to the SciML field. to the Deep Galerkin method.  It is nice that the authors consider a shock problem in Burgers' equation and that the uncertainty bounds contain the shock. The authors propose to use a Bayesian approach with the Laplace approximation of the posterior. I would like to see this UQ method compared to other UQ baselines, e.g., simple ensembling, MC-Dropout and the answers to the reviewer's questions addressed. In particular, I agree with (2) to clarify what is meant by "fictious dataset" in the final version and also motivate why the Laplace approximation via the MAP is taken. I vote for acceptance with the above revisions clarified.

Reviews: review: title: title: the method makes sense, but the experiments are toy examples
main_review: main_review: Neural PDE solvers need accurate methods for quantifying the approximation uncertainty. This work focuses on the Galerkin method as a neural PDE solver. It then interprets the Galerkin as maximum a posteriori estimator, and then the posterior is approximated via Laplace approximation to capture the uncertainty. Experiments with heat equations and burgers PDE demonstrate the effectiveness of the method.

The proposed method based on Laplace approximation makes sense. The experiments however are for toy PDEs, and more extensive experiments for realistic scenarios are needed to prove the method. The inversion in laplace approximation also impedes scaling the method.
title: title: review of Galerkin meets Laplace: Fast uncertainty estimation in neural PDEs
main_review: main_review: This is a nice paper discussing about uncertainty quantification in physics informed machine learning, by using a Bayesian perspective for the uncertainty quantification of DGM. Here are afew questions:
1. What is the advantage of the propsoed method compared to "Yang, Yibo, and Paris Perdikaris. "Adversarial uncertainty quantification in physics-informed neural networks." Journal of Computational Physics 394 (2019): 136-152."? Why not use the same approach there?
2. Where does the uncertainty comes from? Please be more specific. The authors claimed that "To this end, we first interpret the data D to be the fictious dataset", where does the uncertainty of the collocation points comes from in reality? 

Thank you!