Implicit Stochastic Gradient Descent for Training Physics-informed
  Neural Networks
Physics-informed neural networks (PINNs) have effectively been demonstrated
in solving forward and inverse differential equation problems, but they are
still trapped in training failures when the target functions to be approximated
exhibit high-frequency or multi-scale features. In this paper, we propose to
employ implicit stochastic gradient descent (ISGD) method to train PINNs for
improving the stability of training process. We heuristically analyze how ISGD
overcome stiffness in the gradient flow dynamics of PINNs, especially for
problems with multi-scale solutions. We theoretically prove that for two-layer
fully connected neural networks with large hidden nodes, randomly initialized
ISGD converges to a globally optimal solution for the quadratic loss function.
Empirical results demonstrate that ISGD works well in practice and compares
favorably to other gradient-based optimization methods such as SGD and Adam,
while can also effectively address the numerical stiffness in training dynamics
via gradient descent.
Ye Li, Song-Can Chen, Sheng-Jun Huang
# Implicit Stochastic Gradient Descent for Training Physics-informed Neural Networks

###### Abstract

Physics-informed neural networks (PINNs) have effectively been demonstrated in solving forward and inverse differential equation problems, but they are still trapped in training failures when the target functions to be approximated exhibit high-frequency or multi-scale features. In this paper, we propose to employ implicit stochastic gradient descent (ISGD) method to train PINNs for improving the stability of training process. We heuristically analyze how ISGD overcome stiffness in the gradient flow dynamics of PINNs, especially for problems with multi-scale solutions. We theoretically prove that for two-layer fully connected neural networks with large hidden nodes, randomly initialized ISGD converges to a globally optimal solution for the quadratic loss function. Empirical results demonstrate that ISGD works well in practice and compares favorably to other gradient-based optimization methods such as SGD and Adam, while can also effectively address the numerical stiffness in training dynamics via gradient descent.

## 1 Introduction

Gradient descent (GD) and practical stochastic gradient descent with mini-batch gradients (SGD) are widely used optimization algorithms, especially in optimizing deep neural networks. Formally, the goal of optimization is to find a weight vector \(\hat{\theta}\) in parameter space \(\mathbb{R}^{m}\) that minimizes a loss \(L(\theta)\). The GD algorithm is the updating procedure of model weights in the direction of the steepest loss gradient:

\[\theta_{n+1}=\theta_{n}-\alpha\cdot\nabla L(\theta_{n}), \tag{1}\]

where \(\alpha\) is the learning rate. The SGD replaces the gradient \(\nabla L(\theta)\) with a mini-batch gradient \(\nabla\hat{L}_{i}(\theta)\), where \(\hat{L}_{i}\) is the loss computed on mini-batch data instead of the whole dataset. The continuous gradient flow is defined as a curvature \(\theta(t)\) that satisfies the following ordinary differential equation (ODE):

\[\frac{d}{dt}\theta(t)=-\nabla_{\theta}L(\theta(t)). \tag{2}\]

It is easy to show that when the learning rate is sufficiently small, the discrete updates \(\{\theta_{n}\}_{n=0}^{\infty}\) computed by Eq.(1) stay close to a function \(\{\theta(t_{n})\}_{n=0}^{\infty}\) where \(t_{n}=n\alpha\). Variants based on GD/SGD, such as AdaGrad [2], RMSDrop [11], and Adam [12], have been developed in recent years.

Despite its numerous successes in practical optimization tasks such as optimizing deep neural networks, GD/SGD may suffer from numerical instability in some key hyper-parameters, such as the learning rate and batch size. For example, if the learning rate is misspecified, GD/SGD may numerically diverge, and the model training fails. The main reason is the _stiffness_ in the gradient flow dynamics. Typically, the gradient flow dynamics is called a _stiff ODE_ when the gap between the maximum and minimum eigenvalues of the Hessian matrix is large [23]. We can simply perform a linearization for the gradient flow (2) and obtain

\[\frac{d}{dt}\tilde{\theta}(t)=-\nabla_{\theta}^{2}L(\tilde{\theta}(t))\cdot \tilde{\theta}(t). \tag{3}\]

The largest eigenvalue of the Hessian dictates the fastest time-scale of the ODEs. In the language of numerical analysis, to ensure the numerical stability of GD, we need \(\alpha\leq 2/\lambda_{\max}(\nabla_{\theta}^{2}L(\theta))\), where \(\lambda_{\max}(\nabla_{\theta}^{2}L(\theta))\) is the maximum eigenvalue of the Hessian matrix [1].

From the theory of numerical analysis, GD/SGD is not suitable for stiff ODEs, because a very small learning rate and very large number of iterations are required to maintain numerical stability. One of the outstanding first-order solvers with strong stability for stiff ODEs is the implicit (backward) Euler method:

\[\theta_{n+1}=\theta_{n}-\alpha\cdot\nabla L(\theta_{n+1}), \tag{4}\]

where a large learning rate can be used. Eq.(4) is also known as the implicit gradient descent (IGD) or implicit stochastic gradient descent (ISGD) method, as the next iteration \(\theta_{n+1}\) appears implicitly on the right side of Eq.(4), and cannot be computed explicitly.

Physics-informed neural networks (PINNs) are neural networks with outputs constrained to approximately satisfy a system of partial differential equations (PDEs) by using a regularization functional \(\mathcal{R}(u_{\theta}(\mathbf{x}))\) that typically represents the residual of PDEs. A general loss function representation
of PINNs takes the form

\[L(\theta)=\frac{1}{N_{u}}\sum_{i=1}^{N_{u}}|\mathbf{u}_{i}-u_{\theta}(\mathbf{x}_{ i})|^{2}+\mathcal{R}(u_{\theta}(\mathbf{x})), \tag{5}\]

where the given set input output pairs \((\mathbf{x}_{i},\mathbf{u}_{i})\) are corresponding to the initial/boundary conditions of PDEs. The most popular optimizers for training PINNs are gradient descent based Adam optimizer and quasi-Newton based L-BFGS optimizer [13]. However, the additional regularization term \(\mathcal{R}(u_{\theta}(\mathbf{x}))\) has been shown to _increase the stiffness_ of the gradient dynamics [20], causing model training failures especially when the target functions to be approximated exhibit high-frequency or multi-scale features. Typically, for stiff solutions, L-BFGS is more likely to be stuck at a bad local minimum, and Adam may need very small learning rate and very large number of iterations. We claim that IGD/ISGD is more stable than GD/SGD and L-BFGS in the PINNs training when fitting multi-scale solutions. As an example, Figure 1 contrasts these approaches. As the solution of Poisson equation changes from smooth to multi-scale, the maximum eigenvalue of Hessian increases significantly and the gradient flow dynamics of PINN becomes stiff, both Adam and L-BFGS become divergent while our IGD/ISGD is still convergent.

### Contributions

Our main contributions can be summarized in the following points:

* We first propose to employ the IGD/ISGD method to train PINNs. We theoretically and numerically show that IGD/ISGD can overcome the stiffness in the gradient flow dynamics of PINNs, especially for PDEs with multi-scale solutions.
* We used practical L-BFGS and Adam optimizer to deal with the implicit updates in IGD/ISGD, which is effective in practice. The computational cost is comparable to Adam. Furthermore, the method is stable for the learning rate and batch size, making it easier for nonexperts to process neural network training tasks.
* The IGD global convergence property is proven. We theoretically prove that for two-layer fully connected neural networks with large hidden nodes, randomly initialized IGD converges to a globally optimal solution at a linear convergence rate for the quadratic loss function.

### Related Work

**Gradient Descent.** Global convergence of gradient descent based methods have been proved when training deep neural network despite the objective function being non-convex [14, 15, 16, 17]. The dynamics of neural network weights under GD converge to a point that is close to the minimum norm solution under proper conditions [2].

Toulis and his collaborators [18, 19, 20], for the robustness and generalization ability. The IGD/ISGD method was also applied to optimize the k-means clustering problem [21] and the objective matrix factorization loss function that appears in recommendation systems [23], and the convergence time was effectively improved.

**PINNs.** With the rapid growth of deep learning, using neural networks to represent PDE solutions has attracted the attention of many researchers. Based on the early studies of [18, 19], Raissi, Perdikaris, and Karniadakis (2019) proposed the pioneering work of PINNs to solve both forward and inverse problems involving nonlinear PDEs. PINNs have demonstrated remarkable power in applications including fluid dynamics [17, 16, 18, 19], biomedical engineering [15], meta-material design [14, 13], software

Figure 1: _1D Poisson equation_: Results for the heuristic example in Section 2.2 obtained by training a conventional PINN (5-layer, 200 hidden units, \(\tanh\) activations) with gradient descent based Adam optimizer, quasi-Newton based L-BFGS optimizer and our ISGD optimizer. All eigenvalues of \(\nabla_{\theta}^{2}L(\theta)\) are computed and arranged in increasing order. (a): smooth solution \(u_{L}(x)=\sin(2\pi x)\), the maximum eigenvalue is 1.1e+04, non-stiff and all three optimizers trained well. (b): solution with multi-scale features \(u_{H}(x)=\sin(2\pi x)+0.1\sin(50\pi x)\), the maximum eigenvalue is 4.6e+08, stiff and both Adam and L-BFGS failed to train, while our ISGD trained well.
packages (Lu et al., 2021), and numerical simulators (Hennigh et al., 2021; Cai et al., 2021). Adaptive activation functions can be applied to accelerate PINN training (Jagtap, Kawaguchi, and Em Karniadakis, 2020; Jagtap, Kawaguchi, and Karniadakis, 2020; Jagtap et al., 2022). However, despite early empirical success, the original formulations of PINNs often struggles to handle problems exhibiting high-frequency and multi-scale behavior.

Recent works by Wang, Teng, and Perdikaris (2021); Wang, Yu, and Perdikaris (2022); Wang, Wang, and Perdikaris (2021) have identified two fundamental weaknesses in conventional PINN formulations. The first is the remarkable discrepancy in the convergence rate between the data-based loss function and the physical-based loss function. The second is related to the spectral bias, which indeed exists in PINN models and is the leading reason that prevents them from accurately approximating high-frequency or multi-scale functions. In fact, they demonstrated that the gradient flow of PINN models becomes increasingly stiff for PDE solutions exhibiting high-frequency or multi-scale behavior. This result motivates us to use robust implicit numerical schemes such as IGD/ISGD for the numerical solution to the gradient flow of PINN models.

### Organization of the paper

In Section 2, we present the methodology of the proposed IGD/ISGD method. The PINNs framework is also introduced briefly for completeness. Two heuristic examples are presented to show the strong stability of the IGD/ISGD method. In Section 3, we analyze the training dynamics of the IGD/ISGD method when applied to neural network training tasks. Some technical proofs are given in the Appendix. In Section 4, we report various computational examples for inferring the solution of ordinary/partial differential equations by PINNs. Additional computational examples for regression and classification problems are given in the Appendix. Finally, we conclude in Section 5 with a summary.

## 2 Methodology

### Physics-informed neural networks

PINNs are neural networks that imbeds differential equations into neural network training. The initial/boundary condition data of the differential equations are treated as the supervised learning component in the objective loss function, while the residual of the differential equations is applied as an unsupervised regularization factor in the objective loss function. We consider a parametrized PDE system given by:

\[\begin{array}{l}\mathcal{F}(\mathbf{x},u,u_{x},\ldots,\lambda)=0,\quad \mathbf{x}\in\Omega,\\ u(\mathbf{x})=g_{0}(\mathbf{x}),\quad\mathbf{x}\in\partial\Omega,\end{array}\]

where \(\mathbf{x}\) are the spatial and time coordinates, \(u=u(\mathbf{x})\) is the solution to the PDE with boundary/initial data \(g_{0}(\mathbf{x})\), \(\mathcal{F}\) denotes the PDE residual, and \(\lambda\) is the PDE parameter. For example, \(\mathcal{F}=-u_{xx}-f(x)=0\) is the simplest 1D Poisson equation for a given function \(f(x)\). The vanilla PINN uses a fully connected feed-forward neural network \(u_{\theta}(x)\) to approximate the solution \(u(x)\) by minimizing the following loss function:

\[L(\theta)=\omega_{d}L_{data}+\omega_{f}L_{PDE}, \tag{6}\]

where

\[L_{data} = \frac{1}{N_{d}}\sum_{j=1}^{N_{d}}|u_{\theta}(\mathbf{x}_{d}^{j} )-g_{0}(\mathbf{x}_{d}^{j})|^{2},\] \[L_{PDE} = \frac{1}{N_{f}}\sum_{i=1}^{N_{f}}|\mathcal{F}(\mathbf{x}_{f}^{i })|^{2}.\]

Here, \(\{\mathbf{x}_{d}^{j}\}_{j=1}^{N_{d}}\) represents the training data points on \(\partial\Omega\) while \(\{\mathbf{x}_{f}^{i}\}_{i=1}^{N_{f}}\) represents the set of residual points in \(\Omega\). \(\omega_{f}\) and \(\omega_{d}\) are the user-specified weighting coefficients for different loss terms. The first term \(L_{data}\) includes the known boundary/initial conditions and experimental data, which is the usual supervised data-driven part of the neural network. To compute the residuals in the loss function, automatic differentiation is applied to compute the derivatives of the solution with respect to the independent variables. This constitutes the physics-informed part of the neural network as given by the second term \(L_{PDE}\).

The resulting optimization problem is to find the minimum of the loss function by optimizing the trainable parameters \(\theta\). Gradient descent based first-order optimizers such as SGD and Adam (Kingma and Ba, 2014), or quasi-Newton based optimizers like L-BFGS (Liu and Nocedal, 1989), are widely used in PINNs training. However, as Wang, Yu, and Perdikaris (2022) claimed, "...PINNs using fully connected architectures often fail to achieve stable training and produce accurate predictions, especially when the underlying PDE solutions contain high-frequencies or multi-scale features". The gradient flow dynamics of PINNs will become stiff as multi-scale phenomena appear, so explicit GD based optimizers may be unstable, and L-BFGS is more likely to be stuck at a bad local minimum. As we mentioned in the previous section, implicit schemes like IGD/ISGD are more stable to overcome the stiffness problems. Two illustrative examples are presented to show the robustness of IGD/ISGD in the next section.

### Heuristic examples with stability

In this section, we present two heuristic examples to show the stability of IGD/ISGD and the instability of GD/IGD.

**Analytical stiff problem.** The first example is to theoretically analysis the learning rate constraint in the gradient flow dynamics of stiffness problems. We denote a fabricated loss function by

\[L(\theta_{1},\theta_{2})=\frac{K_{1}}{2}(\theta_{1}-\theta_{1}^{*})^{2}+\frac{ K_{2}}{2}(\theta_{2}-\theta_{2}^{*})^{2},\]

where \(\theta_{i}\in\mathbb{R},\ i=1,2\) are two parameters to be optimized, \(K_{i}>0,\ i=1,2\) are two constants. The eigenvalues of the Hessian matrix of \(L(\theta_{1},\theta_{2})\) are characterized by \(K_{1}\) and \(K_{2}\). When \(K_{1}\) and \(K_{2}\) differ in scales, for example, \(K_{1}=10^{-4}\) and \(K_{2}=10^{4}\), the gradient flow of the loss function suffers from the stiffness phenomenon.

A direct computation shows that the loss function update procedure of GD has the following relation:

\[\frac{L(\theta_{1}^{n+1},\theta_{2}^{n+1})}{L(\theta_{1}^{n},\theta_{2}^{n})} \leq\max\{(1-\alpha K_{1})^{2},(1-\alpha K_{2})^{2}\}. \tag{7}\]
Typically, we need \(D=\max\{(1-\alpha K_{1})^{2},(1-\alpha K_{2})^{2}\}\leq 1\) to guarantee loss decay, which implies \(\alpha\leq\frac{1}{\max\{K_{1},K_{2}\}}\). When \(K_{1}=10^{-4}\) and \(K_{2}=10^{4}\), we have \(\alpha\leq 10^{-4}\) and \(D\leq 1-10^{-8}\), meaning that the loss decays very slowly, and very large number of iterations (at least \(\mathcal{O}(10^{8})\)) are needed to converge. For a large learning rate \(\alpha\), the loss decay rate \(D\) may be greater than 1, and the loss may increase as the iterations increase, causing numerical instability in the gradient flow dynamics computation.

For IGD method, the loss function update procedure has the following relation:

\[\frac{L(\theta_{1}^{n+1},\theta_{2}^{n+1})}{L(\theta_{1}^{n},\theta_{2}^{n})} \leq\max\{\frac{1}{(1+\alpha K_{1})^{2}},\frac{1}{(1+\alpha K_{2})^{2}}\}. \tag{8}\]

The loss decay rate \(D=\max\{\frac{1}{(1+\alpha K_{1})^{2}},\frac{1}{(1+\alpha K_{2})^{2}}\}\) satisfies \(D<1\) automatically for all learning rates \(\alpha>0\) and regardless of the scales of \(K_{1},K_{2}\), and \(D\) is even smaller for larger \(\alpha\). This shows the strong stability of IGD to deal with stiffness phenomena.

**1D Poisson equation with multi-scale solution.** This heuristic example is to show the advantage of IGD/ISGD when the gradient flow dynamics of PINN is stiff. We consider a simple 1D Poisson equation

\[-\Delta u(x)=f(x),\quad x\in(0,1) \tag{9}\]

subject to the boundary condition

\[u(0)=u(1)=0.\]

We consider two fabricated solutions: one is \(u_{L}(x)=\sin(2\pi x)\) exhibiting low frequency on the whole domain, and another is \(u_{H}(x)=\sin(2\pi x)+0.1\sin(50\pi x)\) exhibiting low frequency in the macro-scale and high frequency in the micro-scale. Though this example is simple and pedagogical, it resembles many practical scenarios with multi-scale phenomenons.

We represent the unknown solution \(u(x)\) by a 5-layer fully-connected neural network \(u_{\theta}(x)\) with 200 units per hidden layer. \(N_{r}=1000\) training points \(\{x_{i},f(x_{i})\}\) are uniformly sampled in the interval \((0,1)\). Figure 1 shows the results obtained by training PINN with gradient descent based Adam optimizer (Kingma and Ba 2014) with default settings for a maximum \(10^{7}\) epochs, quasi-Newton based L-BFGS optimizer (Liu and Nocedal 1989) with default settings, and our ISGD method with learning rate 0.1 for a maximum \(10^{4}\) epochs. We observe that all three optimizers can train PINN well for smooth solution \(u_{L}(x)\) when there is non-stiff. As multi-scale solution \(u_{H}(x)\) appears, the maximum eigenvalue of Hessian has a significant rise from 1.1e+04 to 4.6e+08. The gradient flow dynamics of PINN becomes stiff, and the popular Adam optimizer is incapable of training PINN to the correct solution even after a million training epochs. The L-BFGS optimizer is also failed to train. As a comparison, our ISGD method can train PINN well both for smooth \(u_{L}(x)\) as well as multi-scale \(u_{H}(x)\) with larger learning rate and smaller iterations.

### Loss decay of GD/IGD

Wang, Teng, and Perdikaris (2021) shows that the loss decay of GD is

\[L(\theta_{n+1})-L(\theta_{n}) \tag{10}\] \[= \alpha\left\|\nabla_{\theta}L(\theta_{n})\right\|_{2}^{2}\left(-1 +\frac{1}{2}\alpha\sum_{i=1}^{N}\lambda_{i}y_{i}^{2}\right),\]

where \(\lambda_{1}\leq\lambda_{2}\leq\cdots\leq\lambda_{N}\) are eigenvalues of the Hessian matrix \(\nabla_{\theta}^{2}L(\xi)\), and \(\mathbf{y}=(y_{1},...,y_{N})\) is a normalized vector. When \(\{\theta_{n}\}_{n=0}^{\infty}\) reaches a local or global minimum, the Hessian matrix \(\nabla_{\theta}^{2}L(\xi)\) is semi-positive definite and all \(\lambda_{i}\geq 0\) for all \(i=1,...,N\). Moreover, for the multi-scale solution \(u_{H}(x)\), computational results show that many eigenvalues of \(\nabla_{\theta}^{2}L(\xi)\) every large (see Figure 1), i.e., stiff during gradient flow dynamics. As a result, it is very possible that \(L(\theta_{n+1})-L(\theta_{n})>0\), which implies that the GD method fails to decrease the loss. A similar computation approach (see the Appendix) shows that the loss decay of IGD is

\[L(\theta_{n+1})-L(\theta_{n}) \tag{11}\] \[= \alpha\left\|\nabla_{\theta}L(\theta_{n+1})\right\|_{2}^{2}\left( -1-\frac{1}{2}\alpha\sum_{i=1}^{N}\lambda_{i}y_{i}^{2}\right),\]

means that the loss will always decay regardless of the stiffness of the gradient flow dynamics of PINNs. In addition, the linear convergence rate of IGD is strictly proven in Section 3.

### Implementation of the IGD/ISGD method

Although the IGD/ISGD method Eq.4 looks simple and theoretically stable, one difficulty that can not be ignored is the implicity of the nonlinear Eq.4. It can also be expressed as the celebrated proximal point algorithm (Yin et al. 2018; Rockafellar 1976):

\[\theta_{n+1}=arg\min_{\theta}\left\{\frac{1}{2}\left\|\theta-\theta_{n}\right\| ^{2}+\alpha\cdot L(\theta)\right\}. \tag{12}\]

Hence, when \(\alpha\) is sufficiently small, \(\theta_{n+1}\) is approximately close to its previous updates \(\theta_{n}\) with the original loss as a regularizer. This sub-optimization task requires additional computation and brings difficulties for the whole optimization process.

To reduce the computational burden, we take a practical "ISGD,L-BFGS" (or "ISGD,Adam") optimizer for PINNs training with multi-scale solutions. Here "ISGD,L-BFGS" means that we first use ISGD with large learning rate for a certain number of iterations, and then switch to L-BFGS with default settings. In the sub-optimization problem (12), we also apply L-BFGS to compute \(\theta_{n+1}\). The optimizer L-BFGS does not require learning rate, and the neural network is trained until convergence, so the number of iterations is also ignored for L-BFGS (Liu and Nocedal 1989). Here, the successful application of L-BFGS in "ISGD,L-BFGS" optimizer is that both the sub-optimization problem and the subsequent optimization problem have good initial point \(\theta_{n}\), thus are easier for L-BFGS to achieve good convergence properties. The "ISGD,Adam" optimizer is to replace L-BFGS by Adam optimizer with default settings in the "ISGD,L-BFGS" optimizer when the parameters of PINNs are too large for the quasi-Hessian matrix computation. The details are illustrated in Algorithm 1.
**Input**: initial \(\theta_{0}\); ISGD learning rate \(\alpha\) and maximum iterations \(K_{0}\); the inner Adam learning rate \(\gamma\) and maximum iterations \(K_{1}\) ; the outer Adam learning rate \(\eta\) and maximum iterations \(K_{2}\)

**Output**: the optimized \(\theta^{*}\)

```
1: Let \(n=0\).
2:while\(n<K_{0}\)do
3: Let \(\tilde{\theta}_{0}=\theta_{n}\) and \(k=0\).
4:while\(k\leq K_{1}\)do
5: Update \(\tilde{\theta}_{k+1}\) \(\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad
### PINN for ordinary differential equations

Singularly perturbed ordinary differential equations have been successfully applied to many fields including gas dynamics, chemical reaction, fluid mechanics, elasticity, etc. To find the solution is a hot and difficult problem because it contains a very small parameter \(\epsilon\). We consider the second-order linear singularly perturbed boundary value differential equation

\[\left\{\begin{array}{l}-\epsilon y^{{}^{\prime\prime}}(x)+y^{{}^{\prime}}(x)= f(x),\quad x\in(0,1),\\ y(0)=0,\quad y(1)=0.\end{array}\right. \tag{19}\]

The true solution is chosen as \(y(x)=\frac{1-e^{\frac{x}{\epsilon}}}{e^{\frac{1}{\epsilon}}-1}+\sin(\frac{\pi }{2}x)\), and \(f(x)\) is given according to Eq.(19). \(\epsilon>0\) is a constant; when \(\epsilon\) is very small, a boundary layer exists near the boundary \(x=1\). Let \(y_{\theta}(x)\) be the neural network approximation of \(y(x)\), then the PINN loss function can be defined as

\[L(\theta) = \frac{1}{2}\left[|y_{\theta}(0)-y(0)|^{2}+|y_{\theta}(1)-y(1)|^{ 2}\right]\] \[+\frac{1}{N}\sum_{i=1}^{N}\left|-\epsilon y_{\theta}^{{}^{\prime \prime}}(x_{i})+y_{\theta}^{{}^{\prime}}(x_{i})-f(x_{i})\right|^{2}.\]

We choose \(N=400\) randomly sampled points to compute the loss function, a batch size of 40 for a small learning rate \(\alpha=0.001\), and a full batch size for a large learning rate \(\alpha=0.5\). A neural network with 4 hidden layers, every 50 units with _tanh_ activations, is applied in all the computations. The results are shown in Figure 2. For the case \(\epsilon=2\), the true solution is smooth. As shown in Fig. 2(a)(b), we find that the ISGD optimizer can significantly improve training convergence and remain stable for different learning rates. For the case \(\epsilon=0.01\), as shown in Fig. 2(f), the true solution has a boundary layer near \(x=1\), and the large gradient creates difficulties for the optimizers. As shown in Fig. 2(d)(e), more epochs and a smaller learning rate are required to be convergent for this singularity phenomenon. While the SGD and Adam optimizers are not convergent for large learning rates, the ISGD can still have stable convergent results, demonstrating the robustness of the proposed method.

### PINN for Poisson equation

Poisson equation is an elliptic partial differential equation of broad utility in theoretical physics. We consider the Poisson equation on the domain \(\Omega=[0,1]\times[0,1]\)

\[\left\{\begin{array}{l}-\frac{\partial^{2}u}{\partial x^{2}}-\frac{\partial ^{2}u}{\partial y^{2}}=f(x,y),\quad(x,y)\in\Omega,\\ u(x,y)=0,\quad(x,y)\in\partial\Omega.\end{array}\right. \tag{20}\]

\begin{table}
\begin{tabular}{c|l c c} \hline Example & Optimizer & Learning rate & \#Iterations \\ \hline
4.1 & SGD(Adam) & 0.001 & 120,000 \\ \cline{2-4} (\(\epsilon=2\)) & ISGD, Adam & 0.5, 0.001 & 102,000 \\ \hline
4.1 & SGD(Adam) & 0.001 & 400,000 \\ \cline{2-4} (\(\epsilon=0.01\)) & ISGD, Adam & 0.5, 0.001 & 360,000 \\ \hline  & SGD(Adam) & 0.0005 & 2,000,000 \\ \cline{2-4}
4.2 & ISGD, Adam & 0.5, 0.0005 & 1,100,000 \\ \hline  & SGD(Adam) & 0.0005 & 1,000,000 \\ \cline{2-4}
4.3 & ISGD, Adam & 0.5, 0.0005 & 550,000 \\ \hline \end{tabular}
\end{table}
Table 1: Hyper-parameters used in the three optimizers for the following 3 examples. “SGD(Adam)” represents SGD shares the same hyper-parameters with Adam. “ISGD, Adam” is referred in Algorithm 1.

Figure 2: The optimization training results for ODE Eq.(19). Top row: \(\epsilon=2.0\). Down row: \(\epsilon=0.01\).
The true solution is chosen as \(u(x,y)=\sin(\pi x)\sin(\pi y)+0.1\sin(10\pi x)\sin(10\pi y)\) with multi-scale features. The PINN loss function is defined as

\[L(\theta)=\frac{1}{N_{b}}\sum_{i=1}^{N_{b}}\left|u_{\theta}(x_{i},y_{i})-u(x_{i},y_{i})\right|^{2}+\]

\[\frac{1}{N_{f}}\sum_{j=1}^{N_{f}}\left|\frac{\partial^{2}u_{\theta}(x_{j},y_{j })}{\partial x^{2}}+\frac{\partial^{2}u_{\theta}(x_{j},y_{j})}{\partial y^{2} }-f(x_{j},y_{j})\right|^{2}.\]

We choose \(N_{b}=400\) randomly sampled points on \(\partial\Omega\), and \(N_{f}=4,000\) randomly sampled points in \(\Omega\) to compute the loss function. A neural network with 6 hidden layers, every 100 units with _tanh_ activations, is applied in all the computations. The three optimizer training results for \(\alpha=0.0005\) and \(0.5\) are shown in Fig. 3(a) and Fig. 3(b), respectively. We see that neither SGD nor Adam can train well as learning rate increases, but our ISGD trains well for different values of \(\alpha\). The PINN prediction is plotted in Fig. 3(c), and the absolute error is shown in Fig. 3(d), with an absolute error less than \(0.2\%\). We see that the PINN trained by the ISGD optimizer can obtain stable and accurate results for the Poisson equation (20).

### PINN for Helmholtz equation

The Helmholtz equation is one of the fundamental equations of mathematical physics arising in many physical problems, such as vibrating membranes, acoustics, and electromagnetism equations. We solve the two-dimensional Helmholtz equation given by

\[\left\{\begin{array}{l}\frac{\partial^{2}u}{\partial x^{2}}+\frac{\partial^ {2}u}{\partial y^{2}}+k^{2}u(x,y)=f(x,y),\quad(x,y)\in\Omega,\\ u(x,y)=0,\quad(x,y)\in\partial\Omega.\end{array}\right. \tag{21}\]

The exact solution for \(k=4\) is \(u(x,y)=\sin(\pi x)\sin(4\pi y)\), and the force term \(f(x,y)\) is given by the Eq.(21). We choose \(N_{b}=400\) randomly sampled points on \(\partial\Omega\), and \(N_{f}=4,000\) randomly sampled points in \(\Omega\) to compute the loss function. A neural network with 6 hidden layers, every 100 units with _tanh_ activations, is applied in all the computations. The three optimizer training results for \(\alpha=0.0005\) and \(0.5\) are shown in Fig. 4(a) and Fig. 4(b), respectively. The PINN solution is plotted in Fig. 4(c), and the absolute error is shown in Fig. 4(d), with an absolute error less than \(0.7\%\). We see that the PINN trained by the ISGD optimizer can obtain stable and accurate results for the Helmholtz equation (21).

## 5 Conclusion

To overcome the numerical instability of traditional gradient descent methods to some key hyper-parameters, a stable IGD/ISGD method was proposed, analyzed and tested in this paper. The IGD/ISGD method includes implicit updates, and the L-BFGS or Adam optimizer can be combined to forward the updates. The global convergence of IGD/ISGD are theoretically analyzed and proven. We apply the IGD/ISGD method to train deep as well as physics-informed neural networks, showing that the IGD/ISGD method can effectively deal with stiffness phenomenon in the training dynamics via gradient descent. The techniques proposed in this paper stabilize the training of neural network models. This may result in making it easier for non-experts to train such models for beneficial applications, such as solving PDEs.

## Acknowledgments

The first author is supported by the National Natural Science Foundation of China (No.62106103), Fundamental Research Funds for the Central Universities (No.ILA22023) and 173 Program Technical Field Fund (No.2021-JCJQ-JJ-0018).

Figure 4: PINN training for Helmholtz equation (21).

Figure 3: PINN training for the Poisson equation (20).