PFNN-2: A Domain Decomposed Penalty-Free Neural Network Method for Solving Partial Differential Equations

In recent years, neural network methods are becoming an attractive alternative for solving partial differential equations (PDEs) arising from applications such as fluid dynamics [2, 3, 4], quantum mechanics [5, 6, 7], molecular dynamics [8], material sciences [9] and geophysics [10, 11]. In contrast to most traditional numerical approaches, methods based on neural networks are naturally meshfree and intrinsically nonlinear therefore can be applied without going through the cumbersome step of mesh generation and could be more potentially applicable to complicated nonlinear problems. These advantages have enabled neural network methods to draw increasingly more attention with both early studies using shallow neural networks [12, 13, 14, 15, 16, 17, 18] and recent works with the advent of deep learning technology [19, 20, 21, 5, 22, 23, 24, 25].

Despite of the tremendous efforts to improve the performance of neural network methods for solving PDEs, there are still issues that require further study. The first issue is related to the accuracy of neural network methods. It is still not fully understood how well a neural network can approximate the solution of a PDE in either theory or practice. Recent investigations were carried out to reveal some preliminary approximation properties of neural networks for simplified problems [26, 27, 28, 29], or to improve the accuracy of neural network methods from various aspects such as by introducing weak form loss functions to relax the smoothness constraints [21, 30, 31, 32, 33, 34], by modifying the solution structure to automatically satisfy the initial-boundary conditions [15, 17, 18, 35, 36], by revising the network structure to enhance its representative capability [37, 38, 39, 40], and by introducing advanced sampling strategies to reduce the statistical error [41, 42]. These improvements so far are usually effective in respective situations, but are not general enough to adapt with different types of PDEs defined on complex geometries.

Another difficulty for solving PDEs with neural network methods is related to the efficiency. It is widely known that training a neural network is usually much more costly than solving a linear system resulted from a traditional discretization of the PDE [43, 44, 45]. To improve the efficiency, it is natural and has been extensively considered [4, 46] to utilize distributed training techniques [47, 48] based on data or model parallelism, by which the training task is split into a number of small sub-tasks according to the partitioned datasets or model parameters so that multiple processors can be exploited. Although this distributed training is general and successful in handling many machine learning tasks [49, 50, 51, 52], it is not the most effective choice for solving PDEs because no specific knowledge of the original problem is adopted throughout the training process.

Inspired by the idea of classical domain decomposition methods [53], it was proposed to introduce domain decomposition strategies into neural network based PDE solvers [54, 55, 56, 57, 58, 59, 60, 61, 62] by dividing the learning task into training a series of sub-networks related to solutions on sub-domains. This tends to be more natural than the plain distributed training approaches because by introducing domain decomposition the training of each sub-network only requires a small part of dataset related to the corresponding sub-domain, therefore can significantly decrease the computational cost. However, these methods may still suffer from issues related to low convergence speed and poor parallel scalability due in large part to the straightforward treatment of artificial sub-domain boundaries.

In this paper, we present PFNN-2, a new penalty-free neural network method that is a subsequent improvement of our previously proposed PFNN method [1]. Inheriting all advantages of PFNN in handling the smoothness constraints and essential boundary conditions of self-adjoint PDEs, PFNN-2 extends the application to a broader range of non-self-adjoint time-dependent problems and introduces a domain decomposition strategy to improve the training efficiency. In particular, PFNN-2 employs a Galerkin variational principle to transform the non-self-adjoint time-dependent equation into a weak form, and adopts compactly supported test functions that are convenient to build on complex geometries. To further reduce the difficulty of the training process and improve the accuracy, PFNN-2 separates the penalty terms related to the initial-boundary constraints from the loss function, with a dedicated neural network learning the initial and essential boundary conditions, and the true solution on the rest of the domain approximated by another network that is totally disentangled from the former one. On top of the above, PFNN-2 introduces an overlapping domain decomposition strategy to speedup multi-processor training without sacrificing accuracy, thanks to the penalty-free treatment of initial-boundary conditions of the decomposed sub-problems and the weak form loss function with reduced smoothness requirement. We carry out a series of numerical experiments to demonstrate that PFNN-2 is advantageous over existing state-of-the-art algorithms in terms of both numerical accuracy and computational efficiency.

The remainder of the paper is organized as follows. In Section 2 we present the basic algorithms employed in PFNN-2 to handle the smoothness constraints and initial-boundary conditions. Following that the domain decomposition strategy adopted by PFNN-2 for parallel computing is introduced in Section 3. We then discuss and compare PFNN-2 with several related state-of-the-art neural network methods in Section 4. Numerical experiment results on a series of linear and nonlinear, time-independent and time-dependent, initial-boundary value problems on simple and complex geometries are reported in Section 5. The paper is concluded in Section 6.

For the ease of discussion, we first consider the case of a single sub-domain $\Omega\subset\mathbb{R}^{d_{\bm{x}}}$ ($d_{\bm{x}}\in\mathbb{N}_{+}$) on which we try to solve the following initial-boundary value problem:

where $\mathcal{A}$, $\mathcal{B}$ and $\mathcal{C}$ are dependent on the solution $u$ as well as the spatial and temporal coordinates $\bm{x}$ and $t$, and $\bm{n}$ is the outward unit normal. Both Dirichlet and Neumann boundary conditions are considered, with the corresponding boundaries $\Gamma_{D}\cup\Gamma_{N}=\partial\Omega$ and $\Gamma_{D}\cap\Gamma_{N}=\varnothing$.

To reduce the smoothness requirement, we introduce the following weak form

with the solution $u$ belonging to a trial (hypothesis) space $\mathcal{H}$ and satisfying initial-boundary conditions

where $\mathcal{V}$ is a certain test space, in which any function $v$ satisfies $v=0$ on $\Gamma_{D}\times(0,T]$.

In PFNN-2, the approximate solution of problem (2) is found within the hypothesis space $\mathcal{H}$ that is formed by neural networks. Unlike most of the existing methods [19, 20, 21, 30, 31, 32, 34] that merely apply a single network to construct $\mathcal{H}$, PFNN-2 adopts two neural networks, with one network $g_{\bm{\theta}_{1}}$ learning the solution with the initial condition and essential boundary condition more accurately and swiftly, and another network $f_{\bm{\theta}_{2}}$ approximating the solution on the rest part of the domain, where $\bm{\theta}_{1}$, $\bm{\theta}_{2}$ donate the sets of weights and biases forming the two networks, respectively. To eliminate the influence of $f_{\bm{\theta}_{2}}$ on the initial condition and essential boundary condition, a length factor function is introduced to impose appropriate restriction, which satisfies

With the help of the length factor function, any function in the hypothesis space $\mathcal{H}$ is constructed as follows:

where $\bm{\theta}=\{\bm{\theta}_{1},\bm{\theta}_{2}\}$.

The length factor function $\ell$ is established by utilizing spline functions, which are more flexible to adapt to various complicated geometries than the analytical functions proposed in previous works [15, 35, 36]. To be specific, we divide the essential boundary $\Gamma_{D}$ into several segments $\{\gamma_{j}\}_{j=1}^{n_{D}}$, $n_{D}\in\mathbb{N}_{+}$. For each segment $\gamma_{j}$, select another segment of the boundary as its companion $\gamma_{j_{o}}$ (note that $\gamma_{j_{o}}$ should not be adjacent to $\gamma_{j}$). Then, for each $\gamma_{j}$, a spline function $l_{j}$ is built, which satisfies

The specific details to construct $l_{j}$ can be found in our previous work of PFNN [1]. Utilizing all the $l_{j}$ functions, the length factor function $\ell$, which is dependent on both space and time, is established by:

where $\mu>0$ is adopted for adjusting the shape of $\ell$. In order to avoid that the value of $\ell$ drops dramatically with the growth of $n_{D}$, we suggest that $\mu=n_{D}$.

Figure 1 is an example to demonstrate how to construct the function $\ell$, in which the essential boundary $\Gamma_{D}$ is divided into two parts, $\gamma_{1}$ and $\gamma_{2}$. Correspondingly, another two segments, $\gamma_{3}$ and $\gamma_{4}$, are taken from the Neumann boundary $\Gamma_{N}$ as the companions $\gamma_{1_{o}}$ and $\gamma_{2_{o}}$, respectively. Then two spline functions $l_{1}$ and $l_{2}$ satisfying (5) can be built according to the method in PFNN [1] and the length factor function is established according to (6).

To improve the flexibility to complex geometries and reduce the cost to evaluation residual (2), rather than the commonly utilized orthogonal polynomials [32, 33], a set of compactly supported test functions

are employed to span the test space $\mathcal{V}$, where $(\widehat{\bm{x}}_{s},\widehat{t}_{s})$ and $h_{s}$ are the center and radius of the support $G_{s}=\cap_{j=1}^{d_{\bm{x}}}\{(\bm{x},t)||x_{j}-\widehat{x}_{s,j}|<h_{s}\}\cap\{(\bm{x},t)||t-\widehat{t}_{s}|<h_{s}\}$ of $v_{s}$, respectively, and $x_{j}$ and $\widehat{x}_{s,j}$ are the $j$th element of $\bm{x}$ and $\widehat{\bm{x}}_{s}$, respectively. The center $(\widehat{\bm{x}}_{s},\widehat{t}_{s})$ is sampled on $(\Omega\cup\Gamma_{N})\times[0,T]$ while the radius $h_{s}$ is small to ensure that the residual can be just evaluated on a tiny region instead of the whole domain. To achieve this, the radius $h_{s}$ is set to:

where $\delta(\widehat{\bm{x}}_{s},\Gamma_{D})$ is the minimum distance from $\widehat{\bm{x}}_{s}$ to $\Gamma_{D}$ so that $v_{s}=0$ on $\Gamma_{D}$, and $\overline{h}$ is the upper bound of $h_{s}$, which is defined as

Such configuration keeps $\sum_{s=1}^{n_{v}}|G_{s}|\approx 2^{d_{\bm{x}}+1}|\Omega|T$ and guarantees that the supports of all the test functions could cover most part of the computing domain.

Since the two neural networks $g_{\bm{\theta}_{1}}$ and $f_{\bm{\theta}_{2}}$ are decoupled by the length factor function $\ell$, they can be trained separately. First, $g_{\bm{\theta}_{1}}$ is trained to learn the initial and essential boundary condition via minimizing the loss function

where $S(\Box)$ represents a set of points sampled on $\Box$. Then, let $g_{\bm{\theta}_{1}^{\ast}}$ be the obtained network after minimizing (9) and $\widetilde{w}_{\bm{\theta}_{2}}=g_{\bm{\theta}_{1}^{\ast}}+\ell f_{\bm{\theta}_{2}}$. An approximate solution $u^{\ast}$ of problem (2) can be obtained by minimizing $L$ as follows:

where

is the estimate of the residual of the governing equation (2), $\#S(\Box)$ donates the size of $S(\Box)$. Eq. (10) is taken as the loss function for training $f_{\bm{\theta}_{2}}$. It can be seen that neither of the loss functions (9) and (10) involves any penalty term, therefore is easier to minimize than those with penalty terms frequently found in various neural network methods for solving partial differential equations. Besides, the training of each network only needs data samples on part of the computing domain, which has less computational cost than training a single network with data samples from the whole domain.

To enable parallel computing capability in PFNN-2, we first partition the computing domain into a group of non-overlapping sub-domains $\{\widetilde{\Omega}_{i}\}_{i=1}^{m}$, and then extend each sub-domain to obtain a series of overlapping sub-domains $\{\Omega_{i}\}_{i=1}^{m}$, where $m\in\mathbb{N}_{+}$ is the number of sub-domains. Figure 2 shows an example of the overlapping domain decomposition partition for the case of $m=9$. Starting with a suitable initial guess $u^{0}$, problem (1) can be solved by constructing a solution sequence $u^{k}$ that converges to the true solution, where $k=1,2,\cdots$ represents the count of iteration. And the approximate solution for each iteration is formed by solving a series of sub-domain problems concurrently on parallel computers. In particular, at the $k$-th iteration the sub-domain problem defined on $\Omega_{i}$ is

where $\Gamma_{i,D}=\partial\Omega_{i}\cap\Gamma_{D}$, $\Gamma_{i,N}=\partial\Omega_{i}\cap\Gamma_{N}$, and $\Gamma_{i}=\partial\Omega_{i}\backslash(\Gamma_{D}\cup\Gamma_{N})$. After solving the sub-domain problems for iteration $k$, the approximate solution defined on the whole domain is composed as $u^{k}\big{|}_{\widetilde{\Omega}_{i}}=u_{i}^{k}$.

To construct the solution sequence, instead of solving the sub-domain problem (12) directly, PFNN-2 finds an approximate solution to the corresponding sub-domain problem in the weak form

where $u_{i}^{k}$ belongs to a hypothesis space $\mathcal{H}_{i}$ and satisfies the initial-boundary conditions

and $\mathcal{V}_{i}$ is a certain test space.

Since the weak form (13) is similar to (2), we can utilize the method presented in the previous section to establish the hypothesis space $\mathcal{H}_{i}$ and test space $\mathcal{V}_{i}$, for all $i=1,2,\cdots,m$. Specifically on sub-domain $\Omega_{i}$, $\mathcal{V}_{i}$ is spanned by a set of compactly supported functions $\{v_{i,s}\}_{s=1}^{n_{v}^{i}}$ ($n_{v}^{i}\in\mathbb{N}_{+}$), and any function $w_{i}\in\mathcal{H}_{i}$ is constructed according to

where $g_{i}$ and $f_{i}$ are two sub-networks that play roles similar to $g_{\bm{\theta}_{1}}$ and $f_{\bm{\theta}_{2}}$, and $\ell_{i}$ is a length factor function that decouples the two sub-networks. Note that here and hereafter, we drop the subscripts representing the learnable parameters for brevity.

Starting from a proper initial state $w_{i}^{0}=g_{i}^{0}+\ell_{i}f_{i}^{0}$, each local approximate solution $w_{i}$ is updated by a dedicated processor to produce a sequence $w_{i}^{k}$ converging to the corresponding local true solution, where $g_{i}^{0}$ and $f_{i}^{0}$ are the initial states of sub-networks $g_{i}$ and $f_{i}$, $k=1,2,\cdots$ is the count of iteration. In the $k$-th iteration, $w_{i}^{k}$ is acquired by solving the corresponding sub-problem (13). And following that, all the $w_{i}^{k}$ are merged together to form the approximate solution $w^{k}$ on the whole domain according to $w^{k}\big{|}_{\widetilde{\Omega}_{i}}=w_{i}^{k}$. To illustrate the solving process accurately and intuitively, Algorithm 1 and Figure 2 are provided, which show that each iteration is comprised of two stages:

• •

  Online stage (line 4-13 of Algorithm 1): As demonstrated in step ① of Figure 2, each processor conducts communication to the immediate neighbors to acquire the current interior boundary condition, where $w_{i}$ satisfies $w_{i}^{k}=w^{k-1}=w_{j}^{k-1}$, on $(\Gamma_{i}\cap\Omega_{j})\times(0,T]$ $(j\neq i)$.

• •

  Offline stage (line 15-23 of Algorithm 1): First, as shown in steps ②-③ of Figure 2, each sub-network $g_{i}$ is trained on the corresponding processor to learn the local solution by minimizing the loss function

  $$\small\begin{array}[]{l}\Psi_{i}^{k}[g_{i}]:=\displaystyle\sum\limits_{(\bm{x},t)\in S(\Gamma_{i}\times(0,T])}\big{(}g_{i}(\bm{x},t)-w^{k-1}(\bm{x},t)\big{)}^{2}+\displaystyle\sum\limits_{(\bm{x},t)\in S(\Gamma_{i,D}\times(0,T])}\big{(}g_{i}(\bm{x},t)-r_{D}(\bm{x},t)\big{)}^{2}\\[17.07164pt] \qquad\quad+\displaystyle\sum\limits_{(\bm{x},t)\in S(\overline{\Omega}_{i}\times\{0\})}\big{(}g_{i}(\bm{x},t)-r_{0}(\bm{x},t)\big{)}^{2}.\end{array}$$

  (15)

  Let $g_{i}^{k}$ be the acquired sub-network after minimizing (15) and $\widetilde{w}_{i}^{k}=g_{i}^{k}+\ell_{i}f_{i}$. Then, as shown in steps ④-⑤ of Figure 2, each sub-network $f_{i}$ is trained on the corresponding processor to approximate the local solution on the remainder of the sub-domain via minimizing

  $$\small\begin{array}[]{l}L_{i}^{k}[f_{i}]:=\displaystyle\sum_{s=1}^{n_{v}^{i}}\Bigg{[}\displaystyle\frac{|\Omega_{i}|T}{\#S(\Omega_{i}\times(0,T])}\sum\limits_{(\bm{x},t)\in S(\Omega_{i}\times[0,T])}\Big{[}\mathcal{A}\nabla\widetilde{w}_{i}^{k}\cdot\nabla v_{i,s}+\Big{(}\dfrac{\partial\widetilde{w}_{i}^{k}}{\partial t}+\nabla\cdot\mathcal{B}+\mathcal{C}\Big{)}v_{i,s}\Big{]}\Big{(}\bm{x},t\Big{)}\\[5.69054pt] \qquad\qquad\qquad\ -\displaystyle\frac{|\Gamma_{i,N}|T}{\#S(\Gamma_{i,N}\times(0,T])}\sum\limits_{(\bm{x},t)\in S(\Gamma_{i,N}\times[0,T])}\big{[}r_{N}v_{i,s}\big{]}\big{(}\bm{x},t\big{)}\Bigg{]}^{2}.\end{array}$$

  (16)

  Afterwards, the obtained sub-network $f_{i}^{k}$ is utilized to form the local approximate solution $w_{i}^{k}=g_{i}^{k}+\ell_{i}f_{i}^{k}$. And finally, the approximate solution on the whole domain is composed as $w^{k}\big{|}_{\widetilde{\Omega}_{i}}=w_{i}^{k}$.

In this section, we review several state-of-the-art neural network methods for solving PDEs and compare them with PFNN-2. For simplicity we consider a model boundary value problem

The most commonly applied neural network algorithms are based on the idea of least squares; examples can be found in e.g., Refs. [14, 15, 17, 19, 2, 63]. In a least-squares based method, an approximate solution of (17) is found by minimizing the loss function in the least-squares form as follows:

where $\beta>0$ is a penalty coefficient. Despite of being straightforward to implement, this approach may encounter issues of low convergence speed and poor accuracy, primarily because the complicated loss function usually requires simultaneously approximating the high-order derivatives of the solution and the additional initial and boundary conditions, thus imposing extreme difficulty for optimization. In addition, the training of the highly nonlinear neural network may also demand large computational cost and lead to low computational efficiency.

In this section, we carry out a series of numerical experiments to investigate the performance of the proposed PFNN-2 method, and compare it with several existing state-of-the-art approaches. Four test cases are utilized to examine the accuracy and efficiency of different numerical methods. In all tests involving neural networks, the ResNet model [67] with sinusoid activation function is utilized for constructing the network, in which each ResNet block is comprised of two fully connected layers and a residual connection. Note that PFNN-2 utilizes two neural networks instead of one. To make fair comparison between various methods, we try to adjust the number of ResNet blocks and the width of hidden layers so that the total degree of freedoms (i.e., parameters including weights and biases for neural network methods, and unknowns for traditional methods) is kept to a similar level. The L-BFGS optimizer [68] is applied for training the neural networks. All the experiments are conducted on a computer equipped with Intel Xeon Gold 6130 CPUs. And all calculations are done with 32-bit floating point computations. The relative $\ell_{2}$-error is taken to measure the accuracy of the approximate solutions. All the results obtained with neural network methods are reported based on ten times of independent experiments.