Spectral PINNs: Fast Uncertainty Propagation with Physics-Informed Neural Networks
Abstract
Physics-informed neural networks (PINNs) promise to signiﬁcantly speed up partial differential equation (PDE) solvers.
However, most PINNs can only solve deterministic PDEs.
Here, we consider stochastic PDEs that contain partially unknown parameters.
We aim to quickly quantify the impact of uncertain parameters onto the solution of a PDE - that is - we want to perform fast uncertainty propagation.
Classical uncertainty propagation methods such as Monte Carlo sampling, stochastic Galerkin, collocation, or discrete projection methods become computationally too expensive with an increasing number of stochastic parameters.
For example, the well-known spectral or polynomial chaos expansions achieve to separate the spatiotemporal and probabilistic domains and offer theoretical guarantees and fast computation of stochastic summaries (e.g., mean), but can be computationally expensive to form.
Our Spectral PINNs approximate the underlying spectral coefﬁcients with a neural network and reduce the computational cost of the spectral expansion while maintaining guarantees.
We derive the method for partial differential equations, discuss runtime, demonstrate initial results on the convection-diffusion equation, and provide steps towards convergence guarantees.
Introduction
Solving partial differential equations (PDEs) in ﬂuid dynamics, astrophysics or climate science is computationally expensive [1].
For example, solving a climate model at 1km resolution can take 10 days on a 5000-GPU supercomputer [2].
Decades of research in numerical methods [3] and physical approximations, such as parametrizations or closures [4] has achieved drastic reductions in the complexity of solving PDEs.
However, numerical methods remain computationally expensive for high-dimensional problems or uncertainty quantiﬁcation [5] and physical approximations require extensive domain knowledge and often sacriﬁce accuracy [4].
Recently, researchers in the ﬁeld of physics-informed neural networks (PINNs) have proposed to combine numerical methods with deep neural networks (NNs) [6, 5].
The combination builds a perfect team: NNs have fast inference times despite approximating high-dimensional functions, but lack guarantees and constraints [7].
Numerical methods are slow for high-dimensional problems, but have rigorous convergence guarantees and physical constraints [8].
Each balancing the other’s weaknesses, physical knowledge has been integrated into NNs as inputs [9, 10, 11], loss function [6], latent [Figure] state [12, 13, 14, 15], analytic constraints [16, 17], or evaluation function [18, 19].
From the other side, NNs have been embedded into numerical methods by approximating the PDE’s parameters [20, 6], dynamics [21, 12], residual [22, 23], differential operator [24, 25], or solution [6, 26].
Our Spectral PINNs leverage NNs to approximate the solution of a PDE, creating a learning-based “surrogate model” [27, 12] that can be used for fast propagation of uncertainties or sensitivity analyses.
PINNs have achieved to speed-up PDE solvers by 1-3 orders of magnitude [28, 26], but are mostly limited to deterministic PDEs [5].
In practice, many parameters of PDEs are partially unknown and uncertainty quantiﬁcation plays a vital role in the development of real-life safety-critical systems [29], climate risk estimation [30], or parameter inference [27].
A small set of PINNs has targeted stochastic differential equations [31, 32, 33], but has been developed rather independently from the existing tools in uncertainty quantiﬁcation (UQ) [27].
Our work builds a bridge between UQ and deep learning by combining well-known spectral expansions with neural networks.
Uncertainty propagation, the UQ subﬁeld of our interest, aims to calculate statistical summaries of the solution given knowledge about parameter distributions [27].
We can roughly divide methods for uncertainty propagation in intrusive and non-intrusive techniques.
Intrusive techniques, such as Stochastic Galerkin [34], require domain knowledge of the PDEs and modiﬁability of the PDE solver which often does not exist [8].
Non-intrusive techniques, such as Monte Carlo sampling [35], collocation methods [36], or discrete projection [36], are purely data-driven methods that only assume solver-generated datasets.
The runtime complexity of many non-intrusive techniques, however, can scale cubically with the number of stochastic parameters [34].
Purely deep learning-based generative models such as GANs [37] or VAEs [38] can offer fast uncertainty propagation via Monte Carlo sampling, but often come without theoretical guarantees [39, 40].
Our method is non-intrusive, faster than conventional UQ methods, and leverages a widely adopted UQ technique called spectral expansion or polynomial chaos expansion (PCE) offering some guarantees [41, 42, 27].
In particular, we propose to approximate the spectral coefﬁcients with NNs.
The integration of NNs into spectral expansions enables fast calculation of statistical summaries, sensitivity analyses [43], and the potential for theoretical guarantees of optimality or convergence [27].
The closest to our work proposes NNs to estimate the coefﬁcients of an arbitrary polynomial chaos [44], but has only been demonstrated on a 1D stochastic elliptic diffusion equation, not on PDEs, and is intrusive [45].
We contribute Spectral PINNs , speciﬁcally:
• The extension of spectral PINNs to PDEs, spatiotemporal, parabolic, hyperpolic, and evolutionary equations.
• A scalable, mesh-free, and data-driven method for non-intrusive uncertainty propagation in potentially high-dimensional parameter and spatiotemporal spaces.
• A potentially faster method for uncertainty propagation than traditional UQ methods, such as stochastic Galerkin or collocation.
The ﬁrst demonstration of Spectral PINNs on the convection-diffusion equation.
Approach
Quantifying uncertainties in solutions to stochastic PDEs (SPDEs), via computing large ensembles is computationally expensive.
In this section, we approximate the solution using a spectral expansion, and estimate the coefﬁcients of that spectral expansion with NNs to build a fast and accurate surrogate for uncertainty propagation.
2.1 Problem deﬁnition
First, we deﬁne a stochastic PDE with the solution, u , and stochastic parameters, Q , as: [Equation] x [ [Equation]
t
[Equation]
x [Equation] with spatial domain, D x ⊂ R n x , temporal domain, D t ⊂ R n t + , sample space, Ω , domain boundary, ∂ , potentially nonlinear operator, N t,x , and initial and Dirichlet boundary conditions, B t,x .
We aim to learn a surrogate model, i.e, an approximation to the solution, ˆ u :
[Equation]
The domain of ˆ u are points of the form ( t, x, Q ( ξ )) , where Q ( ξ ) is a sample from a ﬁnite dimensional distribution of Q , i.e. Q ( ξ )
=
[ Q ( x 0 ; ξ ) , . . .
, Q ( x n x − 1 ; ξ )] for some n x > 0 and x 0 , . . . , x n x − 1 ∈ D x and the realizations of a new stochastic basis, ξ .
The solution then maps a distribution of parameters to a distribution of solutions.
2.2 The Local Convection-Diffusion Equation
We use the convection-diffusion equation to illustrate our method. This is a well-known evolutionary PDE that combines parabolic and hyperbolic equations and is conceptually similar to incompressible
Navier-Stokes.
The equation models the temperature distribution in a vertical ocean column over time as,
[Equation]
with height, x ∈ D x = [0 , 1] , time, t ∈ D t = [0 , 0 . 01] , source, f = 0 , noise, ω ∈ Ω , univariate ( n u =1 ) temperature, u = T : D t × D x × Q (cid:55)→ R 1 , stochastic diffusivity, Q ( x ; ω ) = κ ( x ; ω ) (detailed in Appendix A.2), partial derivative, δ , constant vertical velocity, w = 100 , initial temperature,
T (0 , x ; ω ) = − 4( x − 0 . 5) 2 + 1 , and Dirichlet boundary conditions, T ( t, 0; ω ) =
T ( t, 1; ω ) = 0 . 2.3 Spectral expansion A core idea in our paper is to approximate the solution, u , as a spectral expansion, also called polynomial chaos expansion [42, 27].
The spectral expansion projects the solution onto a space of stochastic orthogonal polynomials.
The solution then becomes a linear combination of spatiotemporal coefﬁcients and stochastic polynomials.
The projection is optimal in the L2 sense [27].
The spectral expansion approximates the solution as,
[Equation]
where ˆ C α j ( t, x ) are deterministic spectral coefﬁcients of the expansion of the random variable ˆ u ( t, x ; ξ ) . In our approach we replace these coefﬁcients with neural network function approximations that will be learned from data. The unknown stochastic domain ω ∈ Ω is approximated by a transformed domain that can easily be sampled from, for example, a set of iid. Gaussians with ξ = { ξ 0 , . .
. , ξ n ξ − 1 } ∼ N (0 , 1) n ξ − 1 with the truncated stochastic dimension n ξ ≤ n x .
We choose the set of polynomials, Ψ A ( ξ ) , to transform a Gaussian distribution, speciﬁcally a multivariate orthogonal Gaussian-Hermite polynomials (as displayed and displayed and detailed in Appendix A.3.1).
The polynomials and spectral coefﬁcients are indexed by the multi-indices, α j ∈ A , which contain index-pairs into the stochastic and polynomial degree, as detailed in Appendix A.3.2.
2.4 Spectral PINNs
The NN jointly estimates all | A | spectral coefﬁcients based on 2D spatiotemporal input pairs, { t, x } ,
[Equation]
Using single input pairs makes our surrogate model mesh-free as the NN can be queried at any location to interpolate the training samples or collocation points.
The NN is trained to approximate the spectral coefﬁcients while only using a limited number of measurements as target data.
We formulate the mean-squared error (MSE) loss for episode, e , and batch size, B , as:
[Equation]
where the realizations of Gaussian random vectors, ξ b ∼ N (0 , 1) n ξ − 1 are shared between the target and approximated solution.
The batch size is chosen to ﬁt one solution sample, B = n t n x with t-grid size, n t , and x-grid size, n x .
The training and validation/test set are using different realizations of the random vector and the same { t, x } -pairs as grid points.
Preliminary Results
Figure 1 shows that the Spectral PINNs in Fig.
1b can successfully approximate the mean and standard deviation of the target solution Fig.
1a.
Importantly the explicit formulation as spectral expansion allows us to compute the mean and standard deviation without any sampling as a function of the spectral coefﬁcients, e.g., µ u ( t, x ) =
C [0 , 0 , 0] ( t, u ) [27].
We can note that the Spectral PINNsapproximated standard deviation captures the growing trend towards the center location ( x = 0 . 5 ) and increasing time ( t = 1 ).
While the accuracy of estimating the mean and std.
dev.
could be improved towards regions of high uncertainty ( x = 0 .
5 , t = 0 . 05 )
the plots indicate promise and we expect hyperparameter tuning to further increase the accuracy.
Figure 4 shows the approximated spectral coefﬁcients.
We used a 6 -layer 2048 -unit fully-connected neural network with ResNet blocks, in-/output normalization, and ReLu activation.
The network was trained with the ADAM optimizer with learning rate, lr = 0 .
001 , and β = [0 . 9 , 0 . 999] for E = 8 epochs in ≈ 1 hr on one 5-year old Intel i 7 − 7500 U @2 .
70 GHz quad-core CPU.
The target data was generated with n t = 128 temporal and n z = 16 grid points and { n s, train = 800 , n s, val = 200 } samples of the solved PDE.
The maximum polynomial degree was chosen to be, n ξ = 3 , s.t.
the number of spectral coefﬁcients is | A | = 10 .
3.1 Linear vs. cubic runtime complexity
Introducing NNs into uncertainty propagation (UP) methods aims to reduce runtime complexity while maintaining accuracy in conducting statistic analysis of high-dimensional PDEs.
Runtime complexity.
Classical UP methods can scale cubically with the number of spectral coefﬁcients in the worst-case [27].
For example, stochastic Galerkin requires to solve one system of
J ( K + 1) equations [34] and collocation methods require to solve M systems of size J [36].
Here,
J is the number of spectral coefﬁcients, K the number of stochastic polynomials, M the number of training samples and in our case J = K
= | A | =10 , M = n s, train =800 .
In comparison, inference of our FCNN scales linearly with the number of outputs/spectral coefﬁcients, J , assuming the network architecture is independent of J [46].
In practice, the architecture depends on the number of outputs and future work will investigate this relationship.
Further, a full comparison will include a discussion on how training time scales with the number of outputs and experimental plots.
Convergence rate.
The ﬁnal accuracy of collocation methods scales linearly in the number of samples, M [36].
However, the convergence rate towards an accurate solution has poor scaling with the number of stochastic parameters n ξ as O ( M − α regularity /n ξ ) .
In comparison, MC sampling converges with O ( M − 1 / 2 ) which is faster for a high number of stochastic parameters, n ξ [27].
Future work, will include a discussion on how the accuracy of NNs converges with the number of stochastic training samples, M , and stochastic dimension, n ξ [47], and provide examples with increased n ξ > n KL =(see Appendix A.2 for current stochastic dimension).
Memory complexity.
In some cases, leveraging NN-based surrogate models can not only reduce computational complexity but also memory complexity [23].
Our network, however, contains n w = 2 n units + n layers n 2 units + n units | A | = 25 M weights which occupy as ﬂoats n weights 4 B ≈ 96 MB .
As this is quite large, future work, will investigate pruning [48] and mixed-precision [49] methods.
Discussion and Future works
We contributed Spectral PINNs , a spectral method for fast uncertainty propagation of highdimensional stochastic PDEs.
We demonstrated the ﬁrst application of NN-based spectral expansions on the convection-diffusion equation.
There exist various avenues for future work.
First, we would like to extend hyperparameter search to increase accuracy and decrease memory consumption.
Second, our work assumes that the stochastic parameters are mutually independent and Gaussian which can be alleviated by an extension to arbitrary polynomial chaos expansion for arbitrary distributions [50].
Third, future work will investigate efﬁcient sampling strategies to reduce the amount of training samples, such as Latin hypercube sampling [51] or Quasi Monte Carlo sampling [52].
Fourth, our work relies purely on FCNNs, but the approximation of as spatiotemporal coefﬁcients naturally lends itself to RNNs [53], CNNs [54], GNNs [55], Transformers [56], or more to exploit spatiotemporal features and propagate uncertaintes of higher-dimensional systems.
Lastly, we will continue the analysis of convergence guarantees based on recent work in converging PINNs [47].
Acknowledgments and Disclosure of Funding
The authors greatly appreciate the discussions with Chris Hill, Shu-Yu Lin, Nicholas Mehrle, Yanni Yuval, Paul O’Gorman, Jaume Peraire, and Youssef Marzouk. Research was sponsored by the United States Air Force Research Laboratory and the United States Air Force Artiﬁcial Intelligence Accelerator and was accomplished under Cooperative Agreement
Number FA8750-19-2-1000.
The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the ofﬁcial policies, either expressed or implied, of the United States Air Force or the U.S. Government. The U S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
We generated n s = 1000 samples of the stochastic convection-diffusion equation via a 2nd-order ﬁnite difference scheme in Python.
We used the number of discretization points, n t = 128 and n x = 16 .
We used a 2nd order forward/central/backward difference scheme for the left boundary/interior/right boundary of the domain, respectively. The initial condition was then time-stepped with T ( t + 1) =
T ( t ) + ∆ t dT ( t ) dt .
We assume that the distribution over the diffusivity is known, for example, through domain knowledge, data assimilation, or Bayesian parameter estimation.
Speciﬁcally, the diffusivity is assumed to follow an exponential Gaussian process (GP) with κ ( x ; ω ) =
exp( Y κ ( x ; ω )) .
The GP, Y κ ( x ; ω ) , is deﬁned by mean, µ Y κ = 3 , correlation length, L = 0 .
3 , variance, σ Y κ = 1 . 0 , exponent, p GP = 1 . 0 , and a covariance kernel that is similar to the non-smooth Ornstein-Uhlenbeck kernel: [Equation]
During training we sample the GP via a Karhunen-Loève (KL) expansion [27]:
[Equation]
Here, λ j and φ j ( x i ) are the j -th eigenvalues/-vectors of the covariance matrix respectively, s.t., (cid:80) n x i Cov ( x j , x i ) φ j ( x i ) = λ j φ j ( x j ) . The KL-expansion has a maximum of n KL,max = n x modes which have been truncated at n KL = 3 modes.
The GP which is deﬁned by random variable, ω , is then approximated by drawing iid.
Gaussian random samples, ξ b = { ξ 0 , ..., ξ n KL − 1 } b ∼ N (0 , 1) n KL − 1 , and projecting them onto the eigenspace of the covariance matrix.
A.3 Polynomial basis A.3.1 Hermite polynomials
We choose the set of polynomials, Ψ A ( ξ ) , to be a set of multivariate orthogonal Gaussian-Hermite polynomials, as displayed in Fig. [Equation]
[Equation]
(8) with the one-term (monic) polynomials, ψ α ji of polynomial degree, α ji , and maximum degree, n .
We are choosing the random vector of each stochastic dimension, i ∈ { 0 , ..., n − 1 } , to be a Gaussian, ξ i ∼ N (0 , 1) and use the associated probabilists’ Hermite polynomials, He α ji .
We chose Hermite polynomials, because we assumed the stochastic parameter, κ , to follow a Gaussian
Process and Hermite polynomials are orthogonal with respect to a Gaussian measure [8]. For other distributions we choose other polynomials, e.g., uniform → Legendre, γ → Laguerre, as listed in table 2.4 of [8]. Future work, will analyse which space of functions can be approximated by Spectral PINNs that rely on Gaussian-Hermite polynomials.
A.3.2 Multi-indices
We choose the vector of polynomial degrees, α j , that builds the polynomial set, Ψ α j to be all combinations of degrees that sum to a maximum of n − 1 . In other words, the multi-index, α j ∈ A , with j ∈ { 0 , ..., | A |− 1 } is deﬁned by the total-degree multi-index set,
A = { α j ∈ N n 0 : ||
α j ||
= (cid:80) n − 1 i =
α ji ≤ n − 1 } [27].
For example, A = { α 0 , ..., α | A |− 1 } = { [0 , 0 , 0] , [0 , 0 , 1] , [0 , 1 , 0] , [1 , 0 , 0] , [0 , 0 , 2] , [0 , 1 , 1] , [0 , 2 , 0] , [1 , 0 , 1] , [1 , 1 , 0] , [2 , 0 , 0] } and | A | = 10 for n = 3 .
A.4 Learned spectral coefﬁcients
Figure 4 shows the approximated spectral coefﬁcients.