Tensorial Implementation for Robust Variational Physics-Informed Neural Networks

Abstract

Variational Physics-Informed Neural Networks (VPINN) train the parameters of neural networks (NN) to solve partial differential equations (PDEs). They perform unsupervised training based on the physical laws described by the weak-form residuals of the PDE over an underlying discretized variational setting; thus defining a loss function in the form of a weighted sum of multiple definite integrals representing a testing scheme. However, this classical VPINN loss function is not robust. To overcome this, we employ Robust Variational Physics-Informed Neural Networks (RVPINN), which modifies the original VPINN loss into a robust counterpart that produces both lower and upper bounds of the true error. The robust loss modifies the original VPINN loss by using the inverse of the Gram matrix computed with the inner product of the energy norm. The drawback of this robust loss is the computational cost related to the need to compute several integrals of residuals, one for each test function, multiplied by the inverse of the proper Gram matrix. In this work, we show how to perform efficient generation of the loss and training of RVPINN method on GPGPU using a sequence of einsum tensor operations. As a result, we can solve our 2D model problem within 350 s on A100 GPGPU card from Google Colab Pro. We advocate using the RVPINN with proper tensor operations to solve PDEs efficiently and robustly. Our tensorial implementation allows for 18 times speed up in comparison to for-loop type implementation on the A100 GPGPU card