FastVPINNs: A fast, versatile and robust Variational PINNs framework for forward and inverse problems in science

Authors: divij  ghose
thivin  anandh
sashikumaar  ganesan
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 12:50:33

Abstract: Variational physics-informed neural networks (VPINNs), with h- and p-refinement, show promise over conventional PINNs. But existing frameworks are computationally inefficient and unable to deal with complex meshes. As such, VPINNs have had limited application when it comes to practical problems in science and engineering. In the present work, we propose a novel VPINNs framework, that achieves up to a 100x speed-up over SOTA codes. We demonstrate the flexibility of this framework by solving different forward and inverse problems on complex geometries, and by applying VPINNs to vector-valued partial differential equations.

Main paper: 
 introduction: 
 In recent years, PINNs have been a popular technique to solve forward and inverse problems related to differential equations (Haghighat et al., 2021; Mao et al., 2020) . PINNs approximate the solution of a PDE using a deep neural network, which is optimized my minimizing the residual of the PDE (Raissi et al., 2019) . Since their introduction, several variants of PINNs have been proposed, which exploit different architectures or numerics to obtain better solutions (Cuomo et al., 2022) . One such extension uses the variational form of the differential equations and numerical integration to compute the loss (Kharazmi et al., 2019; Khodayi-Mehr & Zavlanos, 2020) . While several works have demonstrated than variational PINNs (VPINNs) to be more accurate than vanilla PINNs, particularly with h and p-refinement (hp-VPINNS), the vast majority of VPINN applications to date have only been limited to scalar problems in small domains with limited flexibility in the type of geometry that can be used (Kharazmi et al., 2021; Yang & Foster, 2021; Radin et al., 2023) . One reason for this is the poor computational efficiency of VPINNs, where most implementations calculate the total loss by looping through each cell. Moreover, the simple reference transformation used in existing implementations means that they can only be applied to geometries that can be decomposed into rectangular cells. As a result, solutions of coupled equations on complex meshes, for example the one shown in Fig. 1 , are infeasible in state-of-the-art hp-VPINNs codes.Figure 1 : Mesh for a spur gear with 14,000 quad cells.In this work, we propose a novel implementation of hp-VPINNs based on mapped finite elements (Wilbrandt et al., 2017) , that vectorizes cell-based loops and replaces them with tensor calculations. This leads to considerable speed-ups by leveraging GPU computing, and lends itself to solutions on virtually any geometry that traditional techniques like the finite element method can handle. 
methodology: 
 We present the mathematical preliminaries of both PINNs and hp-VPINNs in A. Assuming our domain is decomposed into N cells cells, each having N quad quadrature points and N test test functions, we compute the variational loss in a manner shown in Fig. 2a . We calculate the derivatives of the test functions, Jacobians and the quadrature weights once, since these do no change while training the network, and assemble them into a pre-multiplier matrix. However, when dealing with skewed cells in complex geometries, the Jacobian may vary at different quadrature points for each test function. By implementing a bilinear transformation to map the gradients of the test function from the reference cell to the actual cell, as shown in Fig. 2b , we can handle any complex geometry that can be decomposed into quadrilateral cells, without the need for a uniform mesh. Hence, by stacking the pre-multiplier matrices for each cell, we obtain a third-order tensor with dimensions N cells × N test × N quad . We collect all the quadrature points from each cell and pass it through the neural network go, thereby obtaining the solutions at all points in a single forward pass, and the gradients of the solutions in a single backward pass. The solution gradients are assembled into a two-dimensional matrix of dimensions N quad × N cells . This approach guarantees that the gradient computations for every cell across the domain are performed in a single backward pass through the neural network, instead of being repeated N cells times as in existing implementation of hp-VPINNs. A detailed description of the bilinear transformation can be found in B  
results: 
 In the following section, we examine the performance of FastVPINNs in terms of speed and accuracy. We use an NVIDIA RTX A6000 GPU with 48GB of device memory for training. The NVIDIA-Modulus library(NVIDIA Modulus) was used for PINNs results, while the code available on (Kharazmi, 2023) was used for obtaining a baseline performance for hp-VPINNs. For our test functions, we use Jacobi polynomials, denoted as P n for degree n, such that P n = P n+1 -P n-1 .We first demonstrate the performance of our code, with h and p-refinement, on forward problems by solving the two-dimensional Poisson equation on the unit square, with the given forcing function,-∆u(x, y) = -2ω 2 sin (ωx) sin (ωy) (x, y) ∈ Ω = [0, 1] 2 , u(0, •) = u(•, 0) = 0 (1)This problem has the exact solution u(x, y) = -sin(ωx) sin(ωy) For a test case with ω = 4π, Fig. 3 shows that our framework exhibits improvement in the accuracy of the solution using h and p-refinement, as expected from any hp-VPINNs solver. We then illustrate the efficiency of our framework with that of PINNs (utilizing the Nvidia-Modulus Library, as shown in Modulus (2023) and hp-VPINNs (as shown in Kharazmi ( 2023)) in Fig. 4 . Fig. 4a compares the median time per epoch required for solving Eq. 1 as the number of cells increases, with each cell having 25 quadrature points in total and 5 test functions in each direction. For a fair comparison, the PINNs solution is obtained using the same number of residual points, and we show our results for both FP32 and FP64 precision. We can see that the existing hp-VPINNs code scales linearly as the number of residual points increases, whereas the time required by our framework remains largely constant, offering up to a 100x speed-up. Further, in Fig. 4b , we show that our code consistently converges faster than PINNs, with the improvements being more noticeable as the solution frequency increases. We observe that our code is faster than PINNs, which could be because we skip the backward pass to compute second derivatives for calculating the loss.To demonstrate the application of our framework to complex geometries, we solve the convectiondiffusion equation given in Eq. 3 on the mesh shown in Fig. 1 .-ε∆u + b • ∇u = f, x ∈ Ω (3)where,f = 50 sin x + cos x; ϵ = 1; b = [0.1, 0] T ;Fig. 5 shows that our code manages to achieve reasonable accuracy even on a complex mesh with a large number of cells, which was thus far infeasible for hp-VPINNs codes.Since existing codes do not scale very well, as is evident from Fig. 4a , hp-VPINNs have not been applied to more complex problems in science that are of practical interest, like the incompressible Navier-Stokes equation. We present a novel application of hp-VPINNs to this problem, using our fastVPINNs framework, for the solution of the Kovasznay flow and lid-driven cavity flow. Fig. 6 shows that our framework can solve such coupled flow problems with relatively good accuracy.To complete the discussion of our hp-VPINNs framework, we demonstrate that it is equally capable in solving inverse problems as shown in Section 3.4 To complete the discussion of our framework, we show that it is equally capable in solving inverse problems on complex meshes. We add a sensor loss to our loss function by generating data from a finite cell solver for the problem shown in in Eq. 4. We then predict the unknown spatially varying diffusion coefficient and the solution, as shown in Fig. 7 . Here, v p represents a polynomial function of degree p. This selection of test and solution spaces results in a Petrov-Galerkin finite element method. Specifically, u h is estimated by u NN (x; W, b), which is the neural network solution, whereas the test function v h is a predetermined polynomial function. By utilizing these functions, we establish the element-wise residual of the variational form equation 7 with u NN (x; W, b) asEQUATIONFurther, define the variational loss byL v (W, b) = 1 N cells Ncells k=1 |W k (x; W, b)| 2 L b (W, b) = 1 N D N D d=1 |B(x; W, b)| 2 ,and the cost function of the neural network in hp-VPINN as .EQUATIONwhere, x c0 = (x 0 + x 1 + x 2 + x 3 ) 4 , x c1 = (-x 0 + x 1 + x 2 -x 3 ) 4 , x c2 = (-x 0 -x 1 + x 2 + x 3 ) 4 , x c3 = (x 0 -x 1 + x 2 -x 3 ) 4 , 
 
 Appendix:  


Decision: The paper proposes a framework that speeds up training of hp-VPINNs for physics simulations by vectorizing, achieving considerable speedups over standard PINNs. While the motivation to improve efficiency is sound, the reviewers note limitations like insufficient mathematical details on the methods and some unclear implementation aspects that need to be addressed for the camera-ready version.

Reviews: review: title: title: Review of fast VPINNs paper
main_review: main_review: This paper improves the efficiency of training hp-VPINNs. This is achieved by vectorizing computations over the hp-VPINN’s cells when computing the hp-VPINN loss function, allowing tensorized GPU operations to be used. The authors claim that this approach is orders of magnitude faster to train than PINNs and standard hp-VPINNs.

In general, reducing the training time of hp-VPINNs is important for them to be useful in realistic applications, and so the subject of this paper is well motivated. It does appear that the authors achieve considerable speedups compared to PINNs and standard hp-VPINNs.

There are some weaknesses that should be addressed:

-	This paper lacks a mathematical description of hp-VPINNs, their trial/test functions, and their loss function - such a description is essential to include. Furthermore, it is unclear whether the proposed approach uses a global trial function (like the original hp-VPINN paper) or local trial functions as well as local test functions – this should be clarified -  and the authors do not mentioned which test functions are used – this is essential to include.

-	It would be useful to have a mathematical explanation on how the bilinear transform allows vectorization, in the context of the hp-VPINN loss function.

-	The authors claim “the existing hp-VPINNs code scales exponentially as the number of residual points increases, whereas the time required by our framework remains largely constant”. This sounds unrealistic to me – what is the actual computational complexity of evaluating the hp-VPINN loss function and how does it scale with the number of residual points? This would be more convincing if stated mathematically (big O notation). Furthermore, I am unsure why the proposed approach is faster in terms of time/epoch than PINNs, because PINNs already use tensorized operations and are evaluated using the same number of residual points – more explanation would be useful.

-	In places the paper needs more proof-reading; please check for typos (e.g. Figure 4 I am not sure the subplot labels should be h-refinement and p-refinement, and in the introduction “where most implementations loop calculate the”)

In summary, the motivation is sound, but the paper lacks some essential implementation descriptions which limit its reliability in my opinion.
title: title: Authors propose Fast Variational physics-informed neural networks framework using tensor-based computation of variational loss and bilinear transformations for quadrilateral cell.
main_review: main_review: Authors propose FASTVPINN framework using tensor-based computation of variational loss and bilinear transformations for quadrilateral cell.