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 usinga numerical method.
A variety of numerical methods to solve these PDEs have been developedincluding but not limited to finite difference (FD) methods (19) (which work primarily on regulardomains partitioned into Cartesian grids) and finite element (FE) methods (36) (which work onirregular domains but require partitioning the domain into multidimensional simplices).
A modernclass of numerical methods called meshless or meshfree methods generalize 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 replacements for meshless numericalmethods such as RBF-FD for directly approximating solutions to PDEs.
Our focus in this work ison the latter problem, though our techniques extend straightforwardly to inferring PDEs as well.
PINNs are typically multilayer feedforward deep NNs (DNNs) that are trained using PDEs andboundary conditions as soft constraints, leveraging automatic differentiation (autograd) for computingspatial and/or temporal derivatives appearing in the PDE terms.
The original PINNs, often referredto as vanilla-PINNs, are challenging to train, at least partly because PDE-based constraints leadto complicated loss landscapes (17).
These issues are somewhat ameliorated by using domaindecomposition (X-PINNs) (16) or gradient-enhanced training (G-PINNs) (38).
Other approaches forameliorating these issues involve curriculum training or sequence-to-sequence learning (17).
Manyof these extensions can help improve training and test accuracy in addition to reducing trainingtime.
Much like other DNNs, PINNs are typically trained in 32-bit floating-point ( i.e. , fp32 or singleprecision).
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 irregular domains.
As RBF-FD uses sparse-matrix vectormultiplication (SpMV) to approximate the derivatives, DT-PINNs are also parallelizable on modern GPU architectures.
It is important to note that DT-PINNs use autograd for the actual optimization ofthe 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 FD-Net, an NN architecture that involves learning FD-like filters for faster prediction of PDEs (35).
In the PINN literature, fractional-PINNs(F-PINNs) use numerical differentiation as autograd cannot compute fractional derivatives (22).
Neverthless, to the best of our knowledge, ours is the first work on using meshless high-order accurate FD-like methods in conjunction with PINNs, allowing them to be trained without any loss inaccuracy on irregular domains (just as autograd does).
An alternative to replacing autograd wouldinvolve eliminating inefficiencies via Taylor-mode differentiation (4).
However, we show that at leastpart of the speedups observed in DT-PINNs is because numerical differentiation results in trainingcompleting in fewer epochs 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 with different depthsand show the cost benefits of using higher-order accurate RBF-FD.
Then, to illustrate the featuresof DT-PINNs, we focus for brevity on two purely spatial PDEs (the nonlinear and linear Poissonequations) and one space-time PDE (the heat equation).
We use these settings to compare DTPINNs and vanilla-PINNs for relative errors, timings, and speedups on a simple desktop GPU.
Wedemonstrate through our experiments that DT-PINNs offer a superior cost-accuracy profile overvanilla-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, andcomparing DT-PINNs against vanilla-PINN on the Poisson and heat equations.
We summarize our [Equation] results and discuss possible future work in Section 5.
Finally, the appendix contains code snippetsand 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 unknownsolution 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 (3) 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;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 ,and N = N i + N b .
Further, let B =
αn
( x ) · ∇ + β .
The vanilla-PINN training loss e ( x, w ) can thenbe 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:
L f ( x )
[Equation]
n (cid:88) [Equation] [Equation] ) , (8) 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 differentiationmatrix L so that L f ( x ) | X ≈ L f ( x ) | X .
If n << N , L will be a sparse matrix with at most nnon-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, the weights c k are unknown,and standard FD (generated from polynomial interpolation) cannot be used to generate the entriesof L ; this follows from the Mairhuber-Curtis theorem (7).
The RBF-FD method involves using aninterpolatory combination of RBFs and polynomials instead.
Without loss of generality, we describethe 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 weightsfor the operator L at the point x 1 are computed by solving the following dense (block) linear systemon this stencil:
[Equation]
where
A
[Equation]
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 set thestencil size to n = 2 (cid:0) ℓ + dd (cid:1) + 1 as this ensures that the above linear system has a solution (2), andalso set m = ℓ if ℓ is odd, and m = ℓ − 1 if ℓ is even (29).
L becomes more dense for higher valuesof p and dimension d , as n =
O ( p d ) .
When this procedure is repeated for each point in the set X ,the cost scales as O ( N ) for fixed n , with further speedups possible by computing multiple sets ofweights using each stencil (28; 29; 31; 34; 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 use a mixed training technique sothat all temporal derivatives are discretized with autograd, and the spatial derivatives are discretizedwith RBF-FD.
We carefully order the evaluations of the network so that L and B multiply the rightquantities.
Let ˜ u
( x, t, w ) be the PINN, and recall that we have N t time steps over the interval [0 , T ] ;in addition, we also have the initial condition at time t = 0 , making for a total of N t + 1 steps.
Define ˜ u k = ˜
u
| x = ˜ X,t = k △ t , where △ t is the timestep.
This vector is the evaluation of ˜ u on all spatiallocations (including ghost nodes) for the k -th time slice.
This definition in turn allows us to definetwo 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,we define the data vectors f and g as follows:
[Equation]
[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 different vectorsdefined, 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 results arereproducible with the consistent random seeds we used in the experiments.
We used the L-BFGSoptimizer with manually fine-tuned learning rates for both vanilla-PINNs and DT-PINNs.
Both DT-PINNs and vanilla-PINNs used a constant NN depth of s = 4 layers with 50 nodes each acrossall runs.
We use quasi-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 directly at the collocation points for convenience.
For all experiments, the spatialdomain Ω is set to
[Equation]
This corresponds to the unit disk in 2D, whose boundary is the unit circle ∥ x ∥ 22 = 1 . In the timedependent 2D heat equation experiment, the space-time domain is chosen to be Ω × [0 , 1] , i.e. , thetime domain runs from time t = 0 to time t = 1 .
The interval [0 , 1] is evenly divided into 24 timesteps so that N t = 24 (excluding t = 0 ), and the time-step was set to △ t = 124 .
We measure allerrors against a manufactured (specified) solution u , and specify f so that the solution holds true.
The boundary condition term g is computed by applying the operator B to u ; we use α = β = 1 forall tests.
To compare DT-PINNs and vanilla-PINNs to the manufactured solutions u , we report therelative ℓ 2 error
[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]
(b) time taken for one application of autograd on fp32 and fp64, compared to the time taken for SpMV using RBF-FD.
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 . [Figure]
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]
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 do not present results for fp64 vanilla PINNs as they were almostidentical to the fp32 vanilla-PINNs in terms of relative errors while being far slower.
It is importantto note 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 vanilla-PINN for the same value of N .
In contrast, the DT-PINN using p = 2 is generally less accurate, showing that higher-order accuracyis needed to reach the same relative errors as fp32 vanilla-PINNs.
Examining Figures 2b and 2c,we also see that all fp64 DT-PINNs (using RBF-FD) can be trained much more rapidly than fp32vanilla-PINNs.
In fact, Figure 2c shows a maximum training speedup of 4x for DT-PINNs even ifp = 2 is ignored .
In general, fp64 DT-PINNs for p > 2 are trained much more quickly thanvanilla-PINNs without a significant loss in accuracy .
We also note that using [-fp
-]fp32 DT-PINNs didnot lead to greater speedups over the fp64 DT-PINNs (results not shown); in addition, fp32 DT-PINNswere not as accurate as their fp64 counterparts.
The reason for the superior performance of {+fp
+} DT-PINNs becomes clearer when we examine the number of training epochs as a function of thenumber of collocation points N (Figure 3).
Figures 3a and 3b both illustrate that both fp32 and {+fp
+} DT-PINNs reach their lowest relative errors in fewer epochs than both fp32 and fp64 vanilla-PINNs.
This implies that the discrete losses are more easily minimized with respect to the NN weight vectorsw .
However, Figure 3b shows that fp64 DT-PINNs take fewer epochs to train as N is increased (whilealso achieving lower errors); in contrast, fp32 DT-PINNs exhibit no such behavior. 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 . 
-]attemptedto train fp32 vanilla-PINNs using ghost points, but using ghost points offered no improvement (resultsnot shown).
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, weuse the manufactured solution given by (19), and set f = ∆ u − e u .
The results are shown in Figure4; we omit p = 2 to make the trends in the higher-order methods clearer. First, Figure 4a shows [Figure]
[Figure] that despite some outliers, fp64 DT-PINNs achieve comparable relative errors to fp32 vanilla-PINNs.
Further, Figure 4b shows that DT-PINNs are still trained faster than vanilla-PINNs.
However, whencomparing Figure 4c to Figure 2c (linear Poisson equation), we see that the average speedup is higherfor the linear Poisson equation than the nonlinear one, though we still obtain reasonable speedups inthe latter case.
This highlights one of the potential limitations of DT-PINNs: they may not offerspeedups over vanilla-PINNs if nonlinear terms not involving differential operators dominatetraining times .
4.4 Heat equation
Next, we compare fp64 DT-PINNs and fp32 vanilla-PINNs on the 2D heat equation.
We specify thetrue solution u to be
[Equation]
[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 spatialcollocation points over 25 time slices (including time t = 0 ) for a total of 20 , 700 spacetime
[Figure]
collocation 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 relative errorssimilar to fp32 DT-PINNs.
Figure 5b shows that we achieve 2-2.5x speedups over vanilla-PINNs.
Weobserved in our experiments that the speedup appears to increase as a function of the number of timesteps N t (results not shown).
We conclude that DT-PINNs for p > 2 offer a superior cost-accuracyprofile over vanilla-PINNs even in a time-dependent problem with mixed training.
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 timesfor both the linear and nonlinear Poisson equations.
For p = 2 (a second-order discretization), weachieved much higher speedups (20x or above for small training sets), but with a loss of accuracy forsmall N .
We also showed that DT-PINNs trained in a mixed fashion (autograd for time, RBF-FD forspace) also achieved comparable accuracy with a 2-2.5x speedup on the heat equation.
DT-PINNstherefore constitute a new paradigm for scientific machine learning that allow practitioners to leverageexisting 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 onmany-core CPUs while other operations are conducted on the GPU.
It may also be profitable toexplore mixed-precision training of DT-PINNs using combinations of fp64, fp32, and even fp16in conjunction with RBF-FD.
Finally, DT-PINNs can be viewed as vanilla-PINNs with partiallylinearized constraints; it may be profitable to explore other types of constraint linearization toaccelerate training and simplify loss function landscapes.
A Appendix
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 somewhatinefficient.
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:
[Figure]
The following Python class defines the forward and backward methods to connect the CuPy SpMVwith the native PyTorch tensors.473 1 class Cupy_L(torch.autograd.
Function): @staticmethod def forward(ctx , pinn_pred , sparse_mat): return from_dlpack( ( cupy _dlpack(to_dlpack(pinn_pred)) ).toDlpack
())7 480@staticmethod def backward(ctx , grad_output): return from_dlpack( .dot ( cupy.from_dlpack to_dlpack grad_output)) ).toDlpack
, None13 486 14 # class Cupy_B can be similarly definedwhere L_t is the transpose of the L matrix:490 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:494 1 # 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 499pde_residual =
L_mul(u_tilde , self.
L) - self.f
bdry_residual = B_mul(u_tilde , self.
B) - [-self.g
 522pde_loss-]self.g8 502pde_loss
= torch.mean(torch.square(torch.flatten(pde_residual )))
bdry_loss = torch.mean(torch.square(torch.flatten(bdry_residual )))11 505train = interior + boundary loss
) return
)15 509 16 loss
)where self.
X_tilde is ˜ X and L_mul and B_mul are:513 1 L_mul = Cupy_L.apply
B_mul = Cupy_B.apply