A novel meta-learning initialization method for physics-informed neural networks

Generally, the parametrized and nonlinear PDEs is given by

where $\mathcal{N}[;\lambda]$ is a nonlinear operator with the parameter $\lambda$, $\Omega$ is the physical domain and $\mathcal{B}(\cdot)$ is the boundary condition. For example, the nonlinear operator of the one-dimensional Burgers equation is given as $\mathcal{N}[u;\lambda]=u^{2}/2-\lambda u_{x}$, where the parameter $\lambda$ is the viscosity coefficient.

In this work, we mainly consider solving the forward problems as well as inverse problems. The forward problem is solutions of PDEs to be inferred with the fixed parameter $\lambda$. For inverse problems, the unknown parameter $\lambda$ is learned from the observed data. In addition, the solutions can be also inferred.

Recently, PINNs have been proven to be promising for solving PDEs or PDE-based systems sirignano2018dgm ; jagtap2020adaptive ; sun2020surrogate . However, training efficiency and prediction accuracy of PINNs are still a tremendous challenge. A few works focused on modifying the structure of NNs to speed up the convergence. D.Jagtap et al. jagtap2020adaptive ; jagtap2019locally employed a scalable hyper-parameter in the activation function, which changes dynamically during the process of optimizing the loss function. The PINN with the adaptive activation function had a more desirable feature for speeding up the convergence rate and improving the solution accuracy. KIM et al. kim2020fast presented a fast and accurate PINN ROM with a nonlinear manifold solution representation. The structure of NNs included two part, namely encoder and decoder. They trained a shallow masked encoder, where training data was from the full order model simulations. Then, the decoder was used as the nonlinear manifold solution representation. Merging prior information in the structure of NNs, Peng et al. peng2020accelerating proposed a Prior Dictionary based PINNs to capture features provided by dictionaries so as to speed up the convergence.

Some works were to decompose the computational domain for speeding up the convergence. D. Jagtap et al. jagtap2020conservative ; jagtap2020extended proposed conservative PINNs and extended PINNs to decompose the computational domain to several discrete sub-domains, where a shallow PINN was employed in each sub-domain. Inspired by works of D. Jagtap, Shukla et al. shukla2021parallel developed a distributed training framework for PINNs, where domain decomposition methods were used in space and time-space. The distributed framework combined the advantage of conservative PINNs and extended PINNs so as to accelerate the convergence. For solving the PDEs with long time integration, the time-space domain might become very large so that the training cost of NNs would become extremely expensive. To this end, Meng et al. meng2020ppinn proposed a parareal PINN to solve the long-time problem. They used a fast coarse-grained solver to decompose the long-time domain into many discrete short-time domains. Training several PINNs with many small-data sets was much faster than a PINN with a large-data set. The parareal PINN could achieve a significant speedup for PDEs with long-time integration. Combining domain decomposition and projection onto space of high-order polynomials, Kharazm et al. kharazmi2021hp introduced hp-variational PINNs to divide the computational space into the trial space and test space. Trial space was the space of NNs and test space represented the space of high order polynomials. The hp-variational PINNs learned from trial space and test space to accelerate the convergence and improve the solution accuracy.

Other works introduced transfer learning to accelerate the training of PINNs. Many physics systems could obtain high-fidelity and low-fidelity data. Generally, high-fidelity data was a large-data and low-fidelity data was a small-data set. To best use of two data sets, Chakraborty et al. chakraborty2020transfer proposed a novel multi-fidelity PINN. The low-fidelity data was first used to train a low-fidelity PINN. Then, they utilized high-fidelity data to obtain high-fidelity PINN by fine-tuning low-fidelity PINN. The multi-fidelity PINN provided a novel method to extract useful information from both low-fidelity and high-fidelity data. To solve brittle fracture problems, Goswami et al. goswami2020transfer introduced transfer learning to re-train the NN partially, instead of re-training the complete network. With this training method, the training efficiency was significantly enhanced. However, these methods ignored the acceleration effect of initialization. In most works, these existing initialization methods, such as Xavier and random initialization, did not use prior information. Therefore, how to use the prior information to endow PINNs with a good initialization is an open issue.

In this part, we first provide a brief overview of PINNs. Fig.1 is the schematic of a PINN. Let $N_{\theta}(\cdot)$ be a NN of depth $D$, where $\theta$ denotes the vector of initialization of the NN. The NN has an input layer, $D-2$ hidden layers and an output layer, which is a mapping from $\mathbb{R}^{d}$ into $\mathbb{R}^{N}$. We employ a Multi-Layer Perception (MLP) with the activation function $\sigma(\cdot)$.

To measure the difference between the NN and the physics constraints (governing equation and boundary condition), the left-hand-side of Eq.(1) is defined as $f(t,x)$, which is given by

The loss function includes partial differentiable structure loss (PDE loss), boundary value condition loss (BC loss), initial value condition loss (IC loss), and true value condition (Data loss) which are defined as

where

and $w_{f}$, $w_{i}$, $w_{b}$ as well as $w_{d}$ are predefined hyper-parameters of weights. In this manuscript, we choose equal weights. $\mathcal{T}_{f}$, $\mathcal{T}_{i}$, $\mathcal{T}_{b}$, and $\mathcal{T}_{data}$ denote the sets of residual points from PDE, initial value, boundary value, and real value, respectively. The $\hat{u}$ denotes the output of NN and $|\cdot|$ denotes the size of $\mathcal{T}$.

In the last step, the training procedure is to search for the best parameter $\theta^{*}$ of NN by minimizing the loss function $\mathcal{L}\left(\theta\right)$, where automatic differentiation is used to minimize the highly nonlinear and non-convex loss function $\mathcal{L}\left(\theta\right)$.

MAML and Reptile are important algorithms for meta-learning initialization. Compared to MAML, the Reptile algorithm nichol2018first only performs stochastic gradient descent (SGD) bottou1991stochastic or Adam kingma2014adam (a gradient descent optimization algorithm) on each task once without twice like MAML, making less computation and memory resources, but the Reptile algorithm can achieve similar performance to MAML. The key of the Reptile algorithm is to obtain the initialization parameters so that the model can get the maximum performance of accuracy in the new task after updating the parameters through one or more gradient steps.

A detailed description of the Reptile algorithm is shown in Algorithm 1. The optimization problem of the Reptile is to find a set of initialization parameters $\phi$. The learner constructs the corresponding loss $L_{\tau}$ through a sampled task $\tau$. Each sampled task contains labeled data for classification or regression. For instance, a task might be a classification task from a labeled dataset with cats and dogs. Next, the learner finds parameters $\tilde{\phi}$ to minimize the loss $L_{\tau}$ after $k$ gradient steps. Then, the parameters are updated in the direction $\phi+\epsilon(\widetilde{\phi}-\phi)$. In the last step, the parameters for Reptile initialization can be obtained by gradient updates under different tasks. The NN with initialization parameters for a new task can achieve the maximum performance of accuracy after a few gradient updates. The Reptile algorithm requires a large number of labeled data for supervised learning. However, when the amount of labeled data is limited, the initialization effect for accelerating convergence is restricted.

PINNs add PDEs as a penalty term into the loss function and then can be learned with less labeled data or even without any labeled data. Inspired by this idea, we propose the new Reptile initialization through modifying the task sampling process and the penalty term of the loss. The new Reptile initialization is extended from supervised learning to unsupervised and semi-supervised learning. To this end, we propose a new Reptile initialization based physics-informed neural network (NRPINN), which can achieve fast adapting and improve accuracy for the PINN. Our implementation is all done in Torch (version 1.4.0) using the Python application programming interface. According to the learning tasks in the new Reptile initialization, the NRPINN is divided into NRPINN with supervised learning (NRPINN-s), NRPINN with unsupervised learning (NRPINN-un) and NRPINN with semi-supervised learning (NRPINN-semi). The detailed description of the NRPINN is shown in Algorithm 2. We consider a MLP represented by $\hat{u}(x;\theta)$ with parameters $\theta$. The PDE to be solved is in the general form

where $\mathcal{N}[\cdot;\lambda]$ is a nonlinear operator parametrized by the known parameter $\lambda$, and $\Omega$ is a subset of $\mathbb{R}^{D}$. We summarize the prior information as two different information, namely high-order information and zero-order information. The two types of information can be represented as follows.

1. 1)

   The high-order information contains high-order derivative information of the solution $u$, which is described by

   $$u_{t}+\mathcal{N}[u;\lambda]=0,\lambda\in P_{h}(\tau),x\in\Omega,t\in[0,T],$$

   (6)

   where $\lambda$ is unknown and $P_{h}(\tau)$ denotes the set of $\lambda$ value. In other words, we consider the known parameter $\lambda$ in Eq.(5) as variables to obtain the Eq.(6), which is called the high-order information of solution to be solved. Under high-order information, the governing equation with different $\lambda$ is considered as a task.

2. 2)

   The zero-order information represents the solution $u_{\lambda}$ under different $\lambda$ can be obtained, where the solution $u_{\lambda}$ is composed of labeled data. The zero-order information is given by

   $$u_{\lambda}\in P_{z}(\tau),x\in\Omega,t\in[0,T],$$

   (7)

   where $P_{z}(\tau)$ represents the set of the labeled data under different parameter $\lambda$. For the zero-order information, the solution $u_{\lambda}$ under different parameter $\lambda$ is considered as a task.

The two types of information can be regarded as tasks of the new Reptile initialization. There are three ways to acquire parameters: (i) when only the zero-order information can be obtained, the new Reptile initialization acquires parameters through the tasks with labeled data. (ii) when only the high-order information can be obtained, the new Reptile initialization acquires parameters through the tasks with unlabeled data. (iii) when the two types of information is obtained, the new Reptile initialization acquires parameters through the tasks with labeled and unlabeled data.

The procedure of the NRPINN algorithm is mainly divided into two parts, the new Reptile initialization and solving PDEs by PINNs. In the process of the new Reptile initialization, $L_{z}$ and $L_{h}$ represent the number of tasks for supervised and unsupervised learning, respectively. $L$ and $N$ represent the number of all tasks and updating all tasks. The new Reptile initialization is used for supervised, unsupervised, and semi-supervised learning. According to zero-order or high-order information, different types of tasks are obtained. Under each task, the difference between the three learning types is whether labeled data is used in the task sampling process. The three ways to learn from the tasks are as follows.

1. 1)

   For supervised learning when $L_{h}=0$, we sample task $\tau_{j}$ from $P_{z}(\tau)$ to obtain labeled data. Based on exact and predicted values, the loss function $\mathcal{L}_{z}(\theta)$ is confirmed.

2. 2)

   For unsupervised learning when $L_{z}=0$, we sample task $\tau_{j}$ from $P_{h}(\tau)$ and specify training points set $\mathcal{T}_{f}$ and $\mathcal{T}_{b}$ from PDE and boundary conditions. Based on the training data, the loss function $\mathcal{L}_{h}(\theta)$ about PDE and boundary condition is specified.

3. 3)

   For semi-supervised learning, task $\tau_{j}$ can be obtained from $P_{z}(\tau)$ and $P_{h}(\tau)$, respectively. We consider the number of $L_{z}$ tasks for supervised learning and the remaining number $L_{h}=L-L_{z}$ of tasks is used for unsupervised learning.

Then, for the same task, the parameter $\theta$ is updated $k$ times by minimizing the loss function through SGD or Adam, which is described by $\widetilde{\theta}=U_{\tau}^{k}(\theta)$. For different tasks, the parameter $\theta$ is updated in the direction $\widetilde{\theta}-\theta$ to find the best parameters $\theta^{*}$, and the update step size $\epsilon$ is a linear policy, $\epsilon=\epsilon_{0}(1-j/N)$.

The second part is the solving PDEs by PINNs. Through the new Reptile initialization, initialization parameters can be obtained. After loading the parameter $\theta^{*}$, the NN is restricted to satisfy the physics constraints. Then the training sets $\mathcal{T}_{f}$, $\mathcal{T}_{i}$, and $\mathcal{T}_{b}$ are specified for PDE, initial, boundary, and real data conditions. The loss $\mathcal{L}_{PINN}(\theta)$ by Eq.(3) is constructed to measure the discrepancy between the NN and the physics constraints. In the last step, the training procedure is to search for the best parameter $\theta^{**}$ by minimizing the loss function $\mathcal{L}_{PINN}\left(\theta\right)$, where automatic differentiation is used to minimize the highly nonlinear and non-convex loss function $\mathcal{L}_{PINN}\left(\theta\right)$. It is worth noting that the new Reptile initialization is used for initialization part, instead of changing the structure of the NN. Therefore, the new Reptile initialization can also be flexibly used for the variants of PINNs, such as PINNs with adaptive activation jagtap2020adaptive , parareal PINNs meng2020ppinn and conservative PINNs jagtap2020conservative , etc.

Poisson equation is represented by

where $\Delta$ denotes Laplace operator and $f$ is penalty term.

Before solving the Burgers equation, the high-order information is given by

with viscous parameter $v\in[0,0.1/\pi]$. Each task obtained from the high-order information represents a governing equation in Eq.(16). The zero-order information is the set of numerical solutions under different $v$. Each task obtained from the zero-order information represents the numerical solution under one value $v\in[0.005/\pi,0.1/\pi]$. In this case, the Burgers equation to be solved is when $v=0.01/\pi$. The analytical solution can be obtained using the Hopf–Cole transformation and its graph is shown in Fig.8, see Basdevant et al. basdevant1986spectral for more details.

In this case, the number of hidden layers is eight and in each layer 20 neurons are used. In the new Reptile initialization part, we consider three types of new Reptile initialization: (i) for supervised learning, we assume that 20 tasks can be obtained from the zero-order information. Then, we sample 1,000 training points under each task, confirm loss function $\mathcal{L}_{z}(\theta)$, and Adam is used to update the parameter $\theta$ 10,000 times with the learning rate of 0.001. (ii) for unsupervised learning, we randomly select 50 tasks, sample 1,000 residual points and 1,000 training points on the boundary, and update the parameters of the NN 10,000 times. (iii) for semi-supervised learning, we select 75 tasks, of which half is used for supervised learning and the other half for unsupervised learning. In the part of solving PDEs by PINNs, the PINN with initialization parameters is used to solve the PDEs. The number of residual training points is 10,000, the number of training data points on the boundary is 5,000, and the number of true data is 0.

Fig.9 shows the predicted solution (top row) as well as the point-wise error (bottom row) by PINNs with Xavier initialization, NRPINN-s, NRPINN-un, and NRPINN-semi. NRPINN-s, NRPINN-un, and NRPINN-semi accurately capture all the dispersive waves in the solution and the point-wise errors in the entire domain are very low. In contrast, the PINN with Xavier initialization fails to predict the solution accuracy and its point-wise error is very high compared to NRPINN-s, NRPINN-un, and NRPINN-semi. Fig.10 gives the comparison of solutions by the PINN with Xavier initialization, NRPINN-s, NRPINN-un, and NRPINN-semi at various t locations. In the left column, Xavier initialization is used for the PINN and the right three columns present the results of NRPINN-s, NRPINN-un, and NRPINN-semi, where one can see that the accuracy of NRPINN-s, NRPINN-un, and NRPINN-semi over the PINN is significantly improved. Fig.11 shows the loss history and relative $L_{2}$ error over the number of iterations with different initializations methods. The loss and relative $L_{2}$ error of the NRPINN over the PINN with other initialization methods are visibly decreasing faster. The PINN with Xavier initialization outperforms other initializations methods but its loss and relative error reach $1e-02$, far less than the NRPINN-un.

Table 4 lists the MAEs between prediction solutions and exact solutions for different initialization methods after 2,000 iterations. The MAEs of NRPINN-s, NRPINN-un, and NRPINN-semi are far less than the PINN with other initialization methods. The performance of NRPINN-un is the best and the MAE of NRPINN-un tends to be $8.582e-05$. The MAE of the PINN with uniform distribution initialization tends to be $3.774e-01$, which is worse than NRPINN-s, NRPINN-un and NRPINN-semi.

Before solving the Schrödinger equation, the high-order information is given by

where $x\in[-5,5]$, $t\in[0,\pi/2]$ and $\lambda\sim\text{uniform}[0,1]$. Each task obtained from the high-order information represents a governing equation in Eq.(17). In this case, the Schrödinger equation to be solved is the PDE when $\lambda=0.5$ and its graph is shown in Fig.12. Considering the Eq.17, we consider unsupervised learning for the new Reptile initialization. We use four hidden layers with 100 neurons in each layer. In the part of new Reptile initialization, we randomly select 100 tasks for unsupervised learning, sample 1,000 residual points and 1,000 training points on the boundary, and update the parameters $\theta$ 10,000 times. In the part of solving PDEs by PINNs, we consider 20,000 residual training points, 1,000 training data points on the boundary, and 0 real training data.

Fig.13 shows the predicted solution at various iterations, and the point-wise error at 30,000 iterations by the PINN (top row) with Xavier initialization and NRPINN-un (bottom row). Compared to the PINN, NRPINN-un accurately captures the feature of the solution faster and the point-wise error at 30,000 iterations in the entire domain is lower. Fig.14 gives the comparison of solution by the PINN and the NRPINN-un at $t=0.25,0.5,0.75$. The top figure shows results of the PINN with Xavier initialization and the bottom figure presents results of NRPINN-un, where the accuracy of NRPINN-un over PINNs is visibly improved. Especially, the predicted solution at $t=0.79$ of PINNs is worst than NRPINN-un, which accords with points-wise error in Fig.14. Fig.15 shows the loss history and the relative $L_{2}$ error over the number of iterations for different initialization methods, where the loss and the relative $L_{2}$ error of NRPINN-un is decreasing faster over other initialization methods.

Table 5 lists the MAEs between prediction solutions and exact solutions for different initialization methods under 30,000 iterations. The MAE of the NRPINN-un tends to be $1.3406e-02$, which is better than $1.0742e-01$ by the PINN with Xavier initialization. The PINN with uniform distribution initialization tends to be $1.0060$, which is far worse than $1.3406e-02$ by Xavier initialization method. In summary, the NRPINN endows the PINN with a good start to achieve higher accuracy.

In response to the lack of data, in this section, we consider quickly identifying the unknown parameters in the PDEs with a small amount of real data by NRPINN. The parameterized viscous Burgers equation is given by

where $u(t,x)$ is hidden solution and $\mathcal{N}[u;v]=u^{2}/2-vu_{x}$ is a parameterized nonlinear term. In this case, we consider solving the PDE when viscosity coefficient $v$is equal to $0.01/\pi$. Eq.18 can be considered as the high-order information for the new Reptile initialization. The zero-order information is the set of the solutions for different $v$. The detailed experimental parameters setting for supervised, unsupervised, and semi-supervised learning are the same as section 4.2.

Fig.16 (left to right) shows the $v$ history, $L_{2}$ error, and loss versus number of iterations, respectively. In the left figure, the predicted $v$ by NRPINN-s, NRPINN-un, and NRPINN-semi visibly converge faster towards its exact value over the PINN. After 80,000 iterations, the predicted $v$ by NRPINN-s, NRPINN-un, and NRPINN-semi are $0.0034965$, $0.0033251$, and $0.0034558$, respectively. The relative errors by NRPINN-s, NRPINN-un, and NRPINN-semi are $9.85\%$, $4.46\%$, and $8.57\%$, whereas the predicted $v$ by the PINN with Xavier initialization is $0.0036330$ and the relative error is $14.13\%$. To analyze the dependence by the PINN with Xavier initialization and NRPINN on real training data ($N_{train}$), we use different number of real data. Fig.17 shows the $v$ history, $L_{2}$ error, and loss versus number of iterations with different real training data, where NRPINN-s, NRPINN-un, and NRPINN-semi always predict unknown parameters $v$ faster and achieve higher accuracy than the PINN with Xavier initialization. For $N_{train}$ is equal to $200$, the predicted $v$ by the PINN is $0.004990$ and the relative error is $56.94\%$, whereas the relative error by NRPINN-s, NRPINN-un, and NRPINN-semi reaches $36.41\%$, $14.80\%$, and $23.32\%$, respectively, where one can see that NRPINN maintains good predictive performance when the amount of real training data is small.

Fig.18 shows the history and relative loss with different initializations methods. The NRPINN has a better performance, which outperforms other initialization methods. Especially, the historical loss and relative error by NRPINN-un reach $1e-06$. In this case, the variation trend of Xavier initialization has similar to random distribution initialization. Fig.19 shows the $v$ history with clean data, $1\%$, and $2\%$ noise data under 80,000 iterations. Whatever $1\%$ or $2\%$ noise data, NRPINN-s, NRPINN-un, and NRPINN-semi outperform the PINN with clean data, whose predicted $v$ is $0.0034554$, $0.0033584$ and $0.003395$. Faced $1\%$ and $2\%$ noise data, the predicted $v$ by PINNs is $0.003828$ and $0.0038560$, far from the exact value $0.0031831$. Fig.19 shows the historical loss with clean and noisy data, where NRPINN-s, NRPINN-un, NRPINN-semi achieve lower and faster historical losses than the PINN with Xavier initialization. Table 6 shows the relative error of predicted $v$ under the different number of noise data, where NRPINN-s, NRPINN-un, and NRPINN-semi achieve smaller errors than the PINN with Xavier initialization when a small number of real data is available. The new Reptile initialization is also used for the variants of PINNs. To demonstrate the new Reptile can be used for the variants of PINNs, we use new Reptile initialization for PINNs with adaptive activation function proposed by D.Jagtap et al. jagtap2020adaptive . Fig.21 shows the predicted $v$ and historical loss, where the PINN with variable $a$ outperform the PINN with $a=1$. Apparently, the new Reptile initialization can also be used for PINNs with adaptive activation functions to accelerate training.