Accelerated Training of Physics Informed Neural Networks (PINNs) using Meshless Discretizations
Abstract
Physics-informed neural networks (PINNs) are neural networks trained by usingphysical laws in the form of partial differential equations (PDEs) as soft constraints.
We present a new technique for the accelerated training of PINNs that combinesmodern scientific computing techniques with machine learning: discretely-trained
PINNs (DT-PINNs).
The repeated computation of the partial derivative terms inthe PINN loss functions via automatic differentiation during training is known tobe computationally expensive, especially for higher-order derivatives.
DT-PINNsare trained by replacing these exact spatial derivatives with high-order accuratenumerical discretizations computed using meshless radial basis function-finitedifferences (RBF-FD) and applied via sparse-matrix vector multiplication.
Whilein principle any high-order discretization may be used, the use of RBF-FD allowsfor DT-PINNs to be trained even on point cloud samples placed on irregulardomain geometries.
Additionally, though traditional PINNs (vanilla-PINNs) aretypically stored and trained in 32-bit floating-point (fp32) on the GPU, we showthat for DT-PINNs, using fp64 on the GPU leads to significantly faster trainingtimes than fp32 vanilla-PINNs with comparable accuracy.
We demonstrate theefficiency and accuracy of DT-PINNs via a series of experiments.
First, we explorethe effect of network depth on both numerical and automatic differentiation of aneural network with random weights and show that RBF-FD approximations ofthird-order accuracy and above are more efficient while being sufficiently accurate.
We then compare the DT-PINNs to vanilla-PINNs on both linear and nonlinear
Poisson equations and show that DT-PINNs achieve similar losses with 2-4x fastertraining times on a consumer GPU.
Finally, we also demonstrate that similarresults can be obtained for the PINN solution to the heat equation (a space-timeproblem) by discretizing the spatial derivatives using RBF-FD and using automaticdifferentiation for the temporal derivative.
Our results show that fp
DT-PINNsoffer a superior cost-accuracy profile to fp32 vanilla-PINNs, opening the door to anew paradigm of leveraging scientific computing techniques to support machinelearning.
Introduction
Partial differential equations (PDEs) provide a convenient framework to model a large number ofphenomena across science and engineering.
In real-world scenarios, PDEs are typically challengingor impossible to solve using analytical techniques, and must instead be approximately solved using anumerical method.
A variety of numerical methods to solve these PDEs have been developed includingbut not limited to finite difference (FD) methods (19) (which work primarily on rectangular domainspartitioned into Cartesian grids) and finite element (FE) methods (36) (which work on domains withcurved boundaries but require partitioning the domain into multidimensional simplices).
A modernclass of numerical methods called meshless or meshfree methods generalizes finite difference methods in such a way as to remove the dependence on Cartesian grids, thereby allowing for the numericalsolution of PDEs on point clouds.
Of these, radial basis function-finite differences (RBF-FD) areamong the most popular and widely-used (3; 5; 37; 8; 9; 1; 11; 12; 13; 10; 25; 26; 14; 33; 18), thougha host of other such methods also exist.
Much like FD or FE methods, these meshless methods canalso approximate solutions to a desired order of accuracy.
More recently, PDE solvers based on machine learning (ML) have begun to gain in popularity due tothe inherent ability of ML techniques such as neural networks (NNs) to recover highly complicatedfunctions from data specified at arbitrary locations (15; 20).
We focus on a popular class of ML-basedmeshless methods called physics-informed neural networks (PINNs) (27).
PINNs can be used bothto discover/infer PDEs that govern a given data set, and as direct PDE solvers.
Our focus in thiswork is on the latter problem, though our techniques extend straightforwardly to inferring PDEs aswell.
PINNs are typically multilayer feedforward deep NNs (DNNs) that are trained using PDEsand boundary conditions as soft constraints, leveraging automatic differentiation (autograd) forcomputing derivatives appearing in the PDE terms.
The original PINNs, often referred to as vanillaPINNs, are challenging to train, at least partly because PDE-based constraints lead to complicatedloss landscapes (17).
These issues are somewhat ameliorated by using domain decomposition (XPINNs) (16) or gradient-enhanced training (G-PINNs) (38).
Other approaches for ameliorating theseissues involve curriculum training or sequence-to-sequence learning (17).
Many of these extensionscan also help improve training and test accuracy.
Much like other DNNs, PINNs are typically trainedin 32-bit floating-point ( i.e. , fp32 or single precision).
In this work, we introduce a new technique for accelerating the training of vanilla-PINNs.
Ourtechnique relies on two key features: (a) using RBF-FD to compute highly accurate (neverthelessapproximate) spatial derivatives in place of autograd, and (b) training the DNN in fp64 rather thanfp32.
These new discretely-trained PINNs (DT-PINNs) can be trained significantly faster than fp32vanilla-PINNs on consumer desktop GPUs with no loss in accuracy or change in DNN architecture.
The use of RBF-FD allows DT-PINNs to retain the meshless nature of vanilla-PINNs, therebyallowing for the solution of PDEs on domains with curved boundaries.
As RBF-FD uses sparsematrix vector multiplication (SpMV) to approximate the derivatives, DT-PINNs are also parallelizableon modern GPU architectures.
It is important to note that DT-PINNs use autograd for the actualoptimization of the PINN weights; only PDE derivatives are discretized using RBF-FD.
The NN literature does contain efforts to replace automatic differentiation with numerical differentiation.
For instance, recent work showed that FD approximations can be efficient for learning generativemodels via score matching (23).
Another example is an NN architecture that involves learning FD-likefilters for faster prediction of PDEs (35).
In the PINN literature, fractional-PINNs (F-PINNs) usenumerical differentiation as autograd cannot compute fractional derivatives (22).
Neverthless, to thebest of our knowledge, ours is the first work on using meshless high-order accurate FD-like methodsin conjunction with PINNs, allowing them to be trained without any loss in accuracy on domainswith curved boundaries (just as autograd does).
An alternative would involve eliminating autogradinefficiencies via Taylor-mode differentiation (4).
However, we show that at least part of the speedupsobserved in DT-PINNs is because numerical differentiation results in training completing in fewerepochs than if autograd were to be used.
To alleviate concerns about replacing autograd with RBF-FD, we first compare fp64 RBF-FDapproximation of different orders of accuracy against fp32 autograd for DNNs and show the costbenefits of using higher-order accurate RBF-FD.
Then, to illustrate the features of DT-PINNs, wefocus for brevity on two purely spatial PDEs (the nonlinear and linear Poisson equations) and onespace-time PDE (the heat equation).
We use these settings to compare DT-PINNs and vanilla-PINNsfor relative errors, timings, and speedups on a simple desktop GPU.
We demonstrate through ourexperiments that DT-PINNs offer a superior cost-accuracy profile over vanilla-PINNs.
The remainder of this paper is organized as follows.
In Section 2, we review both vanilla-PINNs and
RBF-FD.
Next, in Section 3, we discuss how to train DT-PINNs to solve both the Poisson and heatequations.
Then, in Section 4, we present experimental results comparing RBF-FD and autograd,and comparing DT-PINNs against vanilla-PINN on the Poisson and heat equations.
We summarizeour results and discuss possible future work in Section 5.
Finally, the appendix contains additionalresults, code snippets, and key implementation details.
Notation : We use x to refer to spatial coordinates in d dimensions.
On the other hand, a boldedquantity such as c or u indicates a vector with more than d elements (an array).
Finally, the ∼ symbolon top of a quantity indicates that the quantity is an approximation.
Review
We now provide a brief mathematical review of both vanilla-PINNs and RBF-FD discretizations.
Unless we note otherwise, all derivatives in this section are spatial or temporal.
We focus on threeprototypical PDEs: the nonlinear Poisson equation, the linear Poisson equation, and the heat equation.
2.1 Physics-informed neural networks
Let Ω ⊂ R d be a domain with boundary given by ∂ Ω ; here, d is the spatial dimension.
We will focuson the solution of the nonlinear Poisson equation on Ω using PINNs.
Let x ∈ R d , and let u : R d → Rbe the solution to
[Equation]
where ∆ is the Laplacian in R d , ∇ is the R d gradient, n ( x ) is the unit outward normal vector on theboundary ∂ Ω , f ( x ) and g ( x ) are known functions, and α, β ∈ R are known coefficients.
If the e u ( x )term is dropped from (1), we obtain the simpler linear Poisson equation:
[Equation]
The vanilla-PINN technique for solving either Poisson problem involves approximating the unknown 107 solution u ( x ) by a DNN ˜ u ( x, w ) (where w is a vector of unknown NN weights), so that ∥ ˜ u ( x, w ) −u ( x ) ∥ ≤ ϵ for some norm ∥ .
∥ and some tolerance ϵ .
In the absence of existing solution data, thisis accomplished by enforcing (1) and (2) as soft constraints on ˜ u
( x ) to find the weights w duringtraining.
Denote by X = { x k } Nk =1 the set of training points at which these constraints are enforced; 111 in the context of PDEs, these are also called collocation points .
For convenience, we divide X intotwo sets: N i interior points in the set X i and N b boundary points in the set X b ; then, X
=
X i ∪ X b , 113 and N = N i + N b .
Further, let B =
αn
( x ) · ∇ + β .
The vanilla-PINN training loss e ( x, w ) can then 114 be written as:
[Equation]
where ∆ and the ∇ term in B are both applied through autograd.
The tanh activation function istypically used, L-BFGS is used as the optimizer for finding the weights w , and training is typicallydone in fp32 (17).
For the linear Poisson equation, one simply omits the e ˜ u term from the loss above.
For time-dependent PDEs, the PINN becomes a function of space and time ˜ u ( x, t ) .
We focus on theforced heat equation, given by
[Equation]
where (7) is an initial condition and u 0 ( x ) is some known function.
While the ∆ term is handledvia autograd, there are two options to handle temporal derivatives: in a continuous fashion ora time-discrete fashion.
In the former, one samples the full space-time interval Ω × [0 , T ] withcollocation/training points, and then uses autograd to compute all spatial and temporal derivatives.
The loss terms are also augmented with the initial condition (7), which is enforced on the full spacetime solution.
In the time-discrete approach, one typically discretizes the time derivative using anappropriate scheme (such as a Runge-Kutta method), and then proceeds in a step by step fashion.
Wefocus on the continuous approach in this work.
2.2 Radial basis function-finite differences (RBF-FD)
We now briefly review RBF-FD methods.
Given some function f : R d → R , the goal of any FDformula is to approximate the action of a linear operator L on that function ( i.e. , to approximate L f ) at some location x 1 .
This is typically accomplished by using a weighted linear combination of f atx 1 and its n − 1 nearest neighbors.
Mathematically, this can be written as:
[Equation]
where the real numbers c k are called FD weights, and the set of points x 1 , . .
, x n is called an FDstencil.
In general, given a set of samples X = { x j } Nj =1 , one can repeat the above procedure to find
FD weights at every single point.
These weights can be assembled into an N × N differentiation 136 matrix L so that L f ( x ) |
X ≈ L f ( x ) | X .
If n << N , L will be a sparse matrix with at most n 137 non-zero elements per row.
If X lies on a Cartesian grid, the entries of L ( i.e. , the FD weights c k )are known in advance.
However, if X is a more general point cloud, standard FD cannot be used to 139 generate the entries of L (see Mairhuber-Curtis theorem (7)).
The RBF-FD method involves using aninterpolatory combination of RBFs and polynomials instead.
Without loss of generality, we describe 141 the RBF-FD procedure for x 1 and its n − 1 nearest neighbors.
Let ϕ ( r ) =
r m , where m is odd, bea radial kernel (a polyharmonic spline), and q j ( x ) , j = 1 , . . .
, (cid:0) ℓ + dd (cid:1) be a basis for polynomials oftotal degree ℓ in d dimensions; we use tensor-product Legendre polynomials.
The RBF-FD weights 144 for the operator L at the point x 1 are computed by solving the following dense (block) linear systemon this stencil:
[Equation]
where
[Equation]
= ϕ ( ∥ x i − x j ∥ ) , i, j = 1 , . . . , n, P ij = q j
( x i ) , i = 1 , . .
. , n, j = 1 , . .
, (cid:32) ℓ + d d (cid:33) , (10)
L a = L ϕ ( ∥ x − x j ∥ ) | x = x 1 , L q =
L q j ( x ) | x = x 1 , j = 1 , . .
, (cid:32) ℓ + d d (cid:33) , (11) where c is the (column) vector of n RBF-FD weights.
The vector λ is a set of Lagrange multipliersenforcing the condition P T c =
L q , thereby ensuring that (a) the RBF-FD weights c can exactlydifferentiate all polynomials up to total degree ℓ ; and that (b) the error in the RBF-FD approximationto L when applied to all other functions is O ( h ℓ +1 − θ ) , where 0 ≤ h ≤ 1 is a measure of samplespacing in the stencil, and θ is the number of derivatives in the differential operator L (6).
We setℓ =
p + θ − 1 based on the desired order of convergence p so that the error is O ( h p ) .
We then setthe stencil size to n = 2 (cid:0) ℓ + dd (cid:1) + 1 as this ensures that (9) has a solution (2), and also set m = ℓif ℓ is odd, and m = ℓ − 1 if ℓ is even (29).
L becomes more dense for higher values of p anddimension d , as n =
O ( p d ) .
When this procedure is repeated for each point in the set X , the costscales as O ( N ) for fixed n , with large speedups possible by computing multiple sets of weightsusing each stencil (28; 29; 31; 34; 32).
For domains with fixed boundaries, the RBF-FD weights canbe precomputed and reused during simulation.
However, domains with moving boundaries requirerecomputation of RBF-FD weights proximal to the boundary every time-step; fortunately, this can bedone quite efficiently (32).
Ghost points When tackling boundary conditions involving derivatives (such as in (2)) using RBFFD, it is common to include a set of N b ghost points outside the domain boundary ∂ Ω into the set ofsamples to ensure that RBF-FD stencils at the boundary are less one-sided; this aids in numericalstability and accuracy.
Ghost points allow us to also enforce the PDE at both the interior and boundarypoints.
We therefore define and use the extended set ˜ X
=
X i ∪ X
b
∪ X g , where X g is the set ofghost points.
For the remainder of this article, let the RBF-FD differentiation matrix for ∆ be L(dimensions ( N i + N b ) ×
( N i + 2 N b ) ), and for B be B (dimensions N b × ( N i + 2 N b ) ).
Discretely-Trained PINNs (DT-PINNs)
Having described both vanilla-PINNs and RBF-FD, we are now ready to describe DT-PINNs.
Inshort, DT-PINNs are PINNs that are trained using the sparse differentiation matrices L and B inplace of the autograd operations used to compute the Laplacian and boundary operators in the lossfunction (4) (and its heat equation equivalent).
All operations are carried out in fp64.
Poisson Equation Focusing first on the nonlinear Poisson equation (1), recall that ˜ u ( x, w ) is the
PINN approximation to the true solution u ( x ) .
Let the evaluation of ˜ u
( x, w ) on the set ˜ X be ˜ u , i.e. ,
˜ u is obtained by evaluating ˜ u ( x, w ) at interior, boundary, and ghost points.
Further define the vectore , which is the loss function evaluated at only the interior and boundary points, i.e. , e = e ( x, w ) | X .
Then, the DT-PINN loss function can be written as:
[Equation]
where L and B were defined previously, exp(˜ u ) is the element-wise exponential of the vector ˜ u , andf =
f ( x ) | X , and g = g ( x ) | X b ; here, f has dimension ( N i + N b ) × 1 , and g has dimension N b × 1 .
For efficiency, L and B can be precomputed using RBF-FD before the training process begins, andthen simply multiplied with the vector ˜ u to obtain its numerical derivatives.
The loss function (12)is then minimized over w as usual using autograd in conjunction with a suitable optimizer.
For thelinear Poisson equation (3), we simply drop the exp(˜ u ) term.
Heat Equation When using DT-PINNs for the heat equation, we demonstrate the flexibility of our 185 method by using a mixed training technique where the time derivative is handled with autogradand the spatial derivatives are discretized with RBF-FD; this also allows us to bypass the Courant- 187 Friedrichs-Lewy (CFL) constraint on the time-step.
We carefully order the evaluations of the networkso that L and B multiply the right quantities.
Let ˜ u
( x, t, w ) be the PINN, and recall that we have 189 N t time steps over the interval [0 , T ] ; in addition, we also have the initial condition at time t = 0 , 190 making for a total of N t + 1 steps.
Define ˜ u k = ˜
u
| x = ˜ X,t = k △ t , where △ t is the timestep.
This vectoris the evaluation of ˜ u on all spatial locations (including ghost nodes) for the k -th time slice.
This 192 definition in turn allows us to define two vectors, ˜ u ∆ and ˜ u B as follows:
[Equation]
The vector ˜ u ∆ has dimensions ( N t + 1)( N i + N b ) × 1 , and ˜ u B has dimensions
N t N b × 1 .
Next, 194 we define the data vectors f and g as follows:
[Equation]
where f k = f
( x, t )
| x
= X,t = k △ t , and g k
= g ( x, t ) | x = X b ,t = k △ t .
Finally, we define two more vectors:u 0 = u 0 ( x ) | X , the vector evaluating the initial condition on the set X (interior and boundary points);and ˜
u t , the vector of evaluations of ∂ ˜ u∂t at spatial locations (interior and boundary) for each time slice:
[Equation] [Equation]
where (cid:0) ∂ ˜ u ∂t (cid:1) k
= ∂ ˜ u∂t
(cid:12)(cid:12) x
= X,t = k △ t .
This vector is computed using autograd.
With these differentvectors defined, we can finally write the DT-PINN loss vector e for the heat equation ase
[Equation]
Results
We now present experimental results comparing DT-PINN and vanilla-PINN performance on thelinear Poisson equation (3), the nonlinear Poisson equation (1), and the forced heat equation (5).
Setup All experiments were run for 5000 epochs on an NVIDIA GeForce RTX 2070.
All resultsare reproducible with the seeds we used in the experiments.
We used the L-BFGS optimizer with manually fine-tuned learning rates for both vanilla-PINNs and DT-PINNs.
Both DT-PINNs andvanilla-PINNs used a constant NN depth of s = 4 layers with 50 nodes each across all runs.
We usequasi-uniform collocation points generated using a node generator (30).
For the Poisson experiments,we report errors on a test set of N test = 21748 points.
For the heat equation, we report results directlyat the collocation points for convenience.
For all experiments, the spatial domain Ω is set to the unitdisk
[Equation]
In the 2D heat equation experiment, the space-time domain is chosen to be Ω × [0 , 1] .
The time 212 interval [0 , 1] is evenly divided into 24 time steps so that N t = 24 (excluding t = 0 ), and thetime-step was set to △ t = 124 .
We measure all errors against a manufactured (specified) solution u ,and specify f so that the solution holds true.
The boundary condition term g is computed by applying 215 the operator B to u ; we use α = β = 1 for all tests.
To compare DT-PINNs and vanilla-PINNs to themanufactured solutions u , we report the relative ℓ 2 error
[Equation]
[Equation]
where u is the true solution vector, and ˜ u is either the DT-PINN or vanilla-PINN solution vector.
4.1 Effect of neural network depth
[Figure]
We first study the effect of PINN depth (fixing the number of nodes per layer)
s on computing the
Laplacian ∆ of the output with respect to the spatial variable x using either autograd or RBF-FD.
Wecompute errors against fp64 autograd for fp32 autograd and for RBF-FD with p = 2 , 3 , 4 , and 5 .
Allerrors were computed on N = 19638 quasi-uniform collocation points.
The results are in shown in
Figure 1a.
We see fp32 autograd is the most accurate, and that increasing p increases the accuracy of
RBF-FD by about two orders of magnitude.
Only the p = 5 case appears to match fp32 autograd inaccuracy, but errors are reasonably low for p = 3 and p = 4 also.
In Figure 1b, we report the timetaken for the same test.
It is immediately clear that fp64 autograd is significantly more expensivethan the fp32 variant, though both costs scale slowly with the network depth s .
More importantly, thetime taken for fp64 RBF-FD (for all orders) is both lower than both fp32 and fp64 autograd and isindependent of the network depth s , primarily since the RBF-FD weights can be precomputed andrepeatedly reused during training.
4.2 Linear Poisson equation
Next, we study the performance of fp64 DT-PINNs and fp32 vanilla-PINNs on the linear Poissonequation (3) on the domain (17).
Letting x = [ x 1 , x 2 ] , we specify the true solution u to be
[Equation]
[Figure]
and enforce this by setting f = ∆ u .
We then solve for ˜ u as described in Section 3.
The results of thisexperiment are shown in Figure 2.
We present relative errors (Figure 2a), wall clock time (Figure2b), and speedup (Figure 2c).
We also present results for fp64 vanilla PINNs.
It is important tonote that fp64 DT-PINNs were completely stored and trained in fp64, a format widely known to besignificantly slower on the GPU than fp32.
Figure 2a shows the relative errors for DT-PINNs as a function of the number of collocation points
N .
DT-PINNs for p = 3 , 4 , 5 produce similar relative errors to both fp32 and fp64 vanilla-PINNsfor the same value of N .
In contrast, the DT-PINN using p = 2 is generally less accurate, showingthat higher-order accuracy is needed to reach the same relative errors as vanilla-PINNs.
Examining
Figures 2b and 2c, we also see that all fp64 DT-PINNs can be trained much more rapidly than bothfp32 and fp64 vanilla-PINNs.
In fact, Figure 2c shows a maximum training speedup of 4x for
DT-PINNs
even if p = 2 is ignored .
In general, fp64 DT-PINNs for p > 2 are trained much morequickly than vanilla-PINNs without a significant loss in accuracy .
We also note that using fp
DT-PINNs did not lead to greater speedups over the fp64 DT-PINNs, with a loss in accuracy.
Theseresults are shown in Appendix A.1.2.
The superior performance of fp64 DT-PINNs becomes clearerwhen we examine the number of training epochs as a function of the number of collocation points N(Figure 3).
Figures 3a and 3b both illustrate that both fp32 and fp64 DT-PINNs reach their lowestrelative errors in fewer epochs than vanilla-PINNs.
These results provide evidence that DT-PINNshave simpler loss function landscapes than their vanilla-PINN counterparts, also implying that lossfunctions involving linear combinations of NN values are easier to minimize than loss functionsinvolving derivatives of NNs.
Figure 3b also shows that only fp64 DT-PINNs take fewer epochs totrain as N is increased.
We also see that moving to fp64 does not appear to significantly speed upvanilla-PINNs.
It is therefore the combination of discrete training and fp64 that results in speedupsfor increasing N . {+
+}
4.3 Nonlinear Poisson equation
Next, to understand the influence of nonlinearities in terms not including the differential operator,we test the performance of DT-PINNs on the nonlinear Poisson equation (1).
To measure errors,we use the manufactured solution given by (19), and set f = ∆ u − e u .
The results are shown in
Figure 4; for simplicity, we omit p = 2 and fp64 vanilla-PINNs as both these have poor cost-accuracytradeoffs.
First, Figure 4a shows that despite some outliers, fp64 DT-PINNs achieve comparablerelative errors to fp32 vanilla-PINNs.
Further, Figure 4b shows that DT-PINNs are still trained faster
[Figure] [Figure]
than vanilla-PINNs.
However, when comparing Figure 4c to Figure 2c (linear Poisson equation), wesee that the average speedup is higher for the linear Poisson equation.
This highlights one of thepotential limitations of DT-PINNs: they may not offer speedups over vanilla-PINNs if nonlinearterms not involving differential operators dominate training times .
4.4 Heat equation
Next, we compare fp64 DT-PINNs and fp32 vanilla-PINNs on the 2D heat equation.
In order todemonstrate the flexibility of our method, we adopt a mixed training approach where only spatialderivatives are discretized with RBF-FD.
We specify the true solution u to be
[Equation]
and specify f = ∂u∂t − ∆ u so that the solution u satisfies the heat equation for all space-time.
Wecompute the initial condition as u 0 ( x, 0) =
u ( x 1 , x 2 , 0)
= 1 .
We trained on N = 828 spatial
[Figure]
collocation points over 25 time slices (including time t = 0 ) for a total of 20 , 700 spacetimecollocation points; we express all results as a function of p .
These results are shown in Figure 5.
First, Figure 5a shows similar results to the 2D Poisson equation, with p > 2 achieving relativeerrors similar to fp32 DT-PINNs.
Figure 5b shows that we achieve 2-4x speedups over vanilla-PINNs.
We observed in our experiments that the speedup appears to increase as a function of the numberof time-steps N t (results not shown).
It is likely that one could achieve further speedups by alsodiscretizing the temporal derivatives, but we leave this exploration for future work.
Summary and future work
We presented a novel technique, DT-PINNs, that involves training PINNs by using RBF-FD forspatial derivatives, and using fp64 weights and training instead of fp32.
This involved replacingall autograd operations (dense matrix-matrix multiplies) related to PDE loss terms with an SpMVoperation.
We showed that using an RBF-FD approximation order of p > 2 resulted in DT-PINNsthat were comparable in accuracy to vanilla-PINNs while offering 2-4x speedups in training times forboth the linear and nonlinear Poisson equations.
We also showed that DT-PINNs trained in a mixedfashion (autograd for time, RBF-FD for space) also achieved comparable accuracy and speedup onthe heat equation.
DT-PINNs therefore constitute a new paradigm for scientific machine learning thatallow practitioners to leverage existing sophisticated scientific computing techniques to accelerate
ML training times.
There are several possible extensions to our current work.
It is likely that using DT-PINNs inconjunction with X-PINNs and G-PINNs will yield even greater speedups in training times.
Further,
DT-PINNs open the door to leveraging compute more efficiently.
For instance, the SpMV operationscould be parallelized using distributed memory systems in conjunction with GPUs, thereby allowingscaling to very large training sets; alternatively, the SpMV operation could be parallelized on manycore CPUs while other operations are conducted on the GPU.
It may also be profitable to exploremixed-precision training of DT-PINNs.
Finally, DT-PINNs can be viewed as vanilla-PINNs withpartially linearized constraints; it may be profitable to explore other types of constraint linearizationto accelerate training and simplify loss function landscapes.
AppendixA.1 Additional results
In this section, we present additional results that help clarify details of our method.
A.1.1 Higher order derivatives
We also study the effect of PINN depth on computing the biharmonic operator ∆ 2 of the output withrespect to the spatial variable x using either autograd or RBF-FD.
We compute errors against fp64autograd for fp32 autograd and for RBF-FD with p = 2 , 3 , 4 , and 5 .
All errors were computed on
N = 4977 quasi-uniform collocation points.
The results are in shown in Figure 6a.
Much like for the
Laplacian ∆ , we see fp32 autograd is the most accurate, and that increasing p increases the accuracyof RBF-FD by about two orders of magnitude.
In Figure 6b, we report the time taken for the sametest.
It is clear that RBF-FD offers even greater speedups for approximating high-order differentialoperators.
A.1.2 fp32 DT-PINN
Figure 7 shows results for fp32 DT-PINNs on the linear Poisson equation.
While speedups overvanilla-PINNs are similar to fp64 DT-PINNs, they suffer from a degradation in accuracy for all
[Figure]
[Figure]
values of p compared to both fp64 DT-PINNs and vanilla-PINNs.
These results demonstrate thatthe combination of fp64 and discrete training is key to achieving training speedups without a loss inaccuracy, but that fp32 DT-PINNs are also viable alternatives to vanilla-PINNs.
A.2 Implementation
From a practical implementation standpoint, DT-PINNs differ from vanilla-PINNs in several waysbeyond the replacement of autograd with differentiation matrices.
We discuss some of the details ofour implementation within the PyTorch framework (24).
Efficient sparse matrix storage and operations on the GPU Since DT-PINNs replace all autogradcomputations in the loss function with an SpMV, we use the compressed sparse row (CSR) format forstoring L and B ; this ensures both ease of parallelization and compact storage.
As of this writing,
PyTorch support for SpMV operations in the CSR format on the GPU is limited and somewhat inefficient.
To overcome this limitation, we used the CuPy library (21) for all SpMV operations.
Wealso use the DLPack library to enable memory and context sharing between CuPy and PyTorch.
Custom autograd implementation As of this writing, PyTorch is in general unable to apply autogradwith respect to the weight vector w to loss terms involving CuPy CSR matrices.
To overcome thisissue, we wrote custom autograd classes in PyTorch.
Consider the second term in (12), ∥ L ˜ u − f ∥ 22 .
According to the PyTorch API, a custom autograd class for handling this term must supply twomethods: a forward method that tells PyTorch how to evaluate it, and a backward method that tells
PyTorch how to compute the product ( ∇ ˜ u
( L ˜ u − f ))
( ∇ w ˜ u ) .
As ∇ ˜ u ( L ˜ u − f )
=
L T , our custom
PyTorch backward method can be expressed as a CuPy-based SpMV of L T and ∇ w ˜ u , the latter ofwhich PyTorch automatically supplies; loss terms involving B are handled similarly.
The forwardmethod to evaluate the loss term is also an SpMV between L and ˜ u .
The overall procedure is similarfor the heat equation as well.
Code listings We first show the CuPy custom autograd methods for the SpMV operations in theautograd term.
The required imports are:476 1 import
The following Python class defines the forward and backward methods to connect the CuPy SpMVwith the native PyTorch tensors.484 1 class Cupy_L(torch.autograd.
Function): @staticmethod def
): return
cupy.from_dlpack(to_dlpack(pinn_pred)) ).toDlpack
())7 491@staticmethod def backward(ctx , grad_output): return from_dlpack(L_t.dot( cupy.from_dlpack(to_dlpack(grad_output)) ).toDlpack
()), None13 497 14 # class Cupy_B can be similarly definedwhere L_t is the transpose of the L matrix:501 1 L_t = csr_matrix(self.
L.transpose(), dtype=np.float64)
Each DT-PINN training step in linear Poisson can then be written as follows:505
# we use the L-BFGS optimizer provided by PyTorch2 def closure (): self.optimizer.zero_grad
() u_tilde = self.pinn_weights.forward(self.
X_tilde)5 510pde_residual
= L_mul(u_tilde , self.
L) - self.f
bdry_residual = B_mul(u_tilde , self.
B) - self.g8 513pde_loss
= torch.mean(torch.square(torch.flatten(pde_residual )))
bdry_loss = torch.mean(torch.square(torch.flatten(bdry_residual )))11 516train_loss = interior_loss + boundary_loss train_loss.backward(retain_graph
=
True) return train_loss.item()15 520 16 loss_value =
self.optimizer.step(closure)where self.
X_tilde is ˜ X and L_mul and B_mul are:
524 1 L_mul = Cupy_L.apply apply