Multi-stage Neural Networks: Function Approximator of Machine Precision
Deep learning techniques are increasingly applied to scientific problems,
where the precision of networks is crucial. Despite being deemed as universal
function approximators, neural networks, in practice, struggle to reduce the
prediction errors below $O(10^{-5})$ even with large network size and extended
training iterations. To address this issue, we developed the multi-stage neural
networks that divides the training process into different stages, with each
stage using a new network that is optimized to fit the residue from the
previous stage. Across successive stages, the residue magnitudes decreases
substantially and follows an inverse power-law relationship with the residue
frequencies. The multi-stage neural networks effectively mitigate the spectral
biases associated with regular neural networks, enabling them to capture the
high frequency feature of target functions. We demonstrate that the prediction
error from the multi-stage training for both regression problems and
physics-informed neural networks can nearly reach the machine-precision
$O(10^{-16})$ of double-floating point within a finite number of iterations.
Such levels of accuracy are rarely attainable using single neural networks
alone.
Yongji Wang, Ching-Yao Lai
# Multi-stage Neural Networks: Function Approximator of Machine Precision

###### Abstract

Deep learning techniques are increasingly applied to scientific problems, where the precision of networks is crucial. Despite being deemed as universal function approximators, neural networks, in practice, struggle to reduce the prediction errors below \(O(10^{-5})\) even with large network size and extended training iterations. To address this issue, we developed the multi-stage neural networks that divides the training process into different stages, with each stage using a new network that is optimized to fit the residue from the previous stage. Across successive stages, the residue magnitudes decreases substantially and follows an inverse power-law relationship with the residue frequencies. The multi-stage neural networks effectively mitigate the spectral biases associated with regular neural networks, enabling them to capture the high frequency feature of target functions. We demonstrate that the prediction error from the multi-stage training for both regression problems and physics-informed neural networks can nearly reach the machine-precision \(O(10^{-16})\) of double-floating point within a finite number of iterations. Such levels of accuracy are rarely attainable using single neural networks alone.

**keywords**: scientific machine learning; neural networks; physics-informed neural networks, multi-stage training

## 1 Introduction

Deep learning techniques [1] have been well developed in the fields of computer vision [2; 3] and natural language processing [4; 5; 6]. More recently, neural networks have been increasingly applied to the mathematical and physical sciences [7; 8; 9], where the demand for precision is high. In particular, physics-informed neural networks (PINNs) [10; 11] have emerged as a new class of numerical solver for partial differential equations, where computing high-precision solutions becomes an intrinsic requirement of the method.

Neural networks have been proven to be universal function approximators [12; 13]. However, in practice, neural network training often falls into local minima [14; 15], causing the training loss to plateau after a certain number of iterations \(n_{iters}\). This issue cause the failure modes of PINNs [15]. Advanced methods focusing on different aspects, such as activation function selection [16; 17], network configuration [18; 19; 20], optimization techniques [21; 22], trainable weights [23], and loss function [24; 25], have been developed to effectively enhance the convergence rate of the loss function for various problems. However, few of these methods manage to reduce the training error less than \(O(10^{-5})\). In contrast, classical numerical methods (e.g., finite difference) can systematically enhance solution's accuracy by simply reducing the grid size [26]. This is a major shortcoming of neural networks for solving many problems within mathematical and physical sciences.

In this work, we proposed the multi-stage neural networks that effectively addresses this limitation. Our novel method involves dividing the network training into multiple stages, where each stage
incorporates a separate neural network. The setting of each network in a given stage are optimized based on the residues from the preceding stage. By executing training stage by stage, we significantly enhance the convergence rate, ensuring that it remains consistently high throughout the iterations. As a result, the combined neural networks from different stages can approximate the target function with remarkable accuracy, with the error approaching the machine precision \(O(10^{-16})\) for double-floating point numbers.

We begin with the introduction of the multi-stage neural network for regression problems in Section 2. By exploring the limitations of classical neural network training, we highlight the benefits of the multi-stage training in overcoming these constraints. We then propose and substantiate the optimal settings for each training stage. In Section 3, we extend the method to physics-informed neural networks (PINNs) for solving differential equations. Unlike regression problems, the optimal settings for PINNs in each stage are implicitly tied to the equation residues of previous stages. Both theoretical investigation and practical algorithmic solutions are presented to address this challenge. Additional techniques that can expedite the multi-stage training for PINNs are also discussed. In Section 4, we generalize the multi-stage training scheme to solve combined-forward-and-inverse problems, which are of great importance in mathematical and physical sciences. Lastly, we provide further discussions on the challenges and potential development of the MSNN method in Section 5 and conclude the paper in Section 6.

## 2 Multi-stage training scheme for regression problems

We first illustrate the multistage idea with regression problems that involve predicting a continuous output variable \(u\) as a function of the input variable \(x\). We train a neural network that represents \(u(x)\) to fit \(N_{d}\) data points, denoted \((x_{i},u_{i})\). The loss function for a regression problem is typically the mean squared error (MSE), defined as

\[\mathcal{L}=\frac{1}{N_{d}}\sum_{i=1}^{N_{d}}[u(x_{i})-u_{i}]^{2}\,, \tag{2.1}\]

In this study, we consider all training data with no noise.

### Limitation of regular neural network training

To illustrate the limitation of neural network's function approximation capacity, we consider a target function,

\[u_{g}(x)=\sin(2x+1)+0.2e^{1.3x}\,, \tag{2.2}\]

we created training data by sampling 300 data points from it with no noise, uniformly distributed within the domain \(x\in[-1,1]\). To fit the training data, we create a fully-connected neural network made of three hidden layers with 20 units in each layer and use hyperbolic tangent as the activation function for each unit. Using Adam [27] optimizer, figure 1\((a)\) shows that the trained neural network \(u_{0}(x)\) captures the target function \(u_{g}(x)\) well. During the iterations, the training loss \(\mathcal{L}\) based on mean squared error (MSE) between the data and network, decreases significantly at the early stage (figure 1\(b\)). However, after 5000 iterations, it reaches a plateau around \(O(10^{-7})\) with very small convergence rate. The error function \(e_{1}(x)\) between \(u_{g}\) and \(u_{0}\) across the training domain is, thus, trapped around \(10^{-4}\) (inset of figure 1\(a\)). Further experiments, elucidated in A, affirm that this plateau value of the error remains consistent even with larger networks and additional data, and not optimizer-specific.
Neural networks are known for their spectral biases [28], also referred to as the frequency principle [29]. Utilizing the tool of neural tangent kernels [30], prior studies [31; 32] demonstrated that a standard multi-layer neural network struggles to learn the high frequencies of target functions in both theory and practice. The plateau of training loss corresponds to a mismatch between the trained network \(u_{0}(x)\) and target function \(u_{g}(x)\) at high frequencies. Figure 1\((b)\) demonstrates that the error function \(e_{1}(x)=u_{g}(x)-u_{0}(x)\) within the domain is indeed a high-frequency function.

### Key settings of multi-stage training scheme

Since training a single neural network struggles with learning the high frequencies of the target function, an intuitive approach is to train a second neural network to capture the error function \(e_{1}(x)\), or the residue, between the training data and the first trained network [33]. The original training data from (2) is denoted with \((x^{(i)},u_{g}^{(i)})\). The training data for the second neural network would be \((x^{(i)},e_{1}(x^{(i)}))\), where \(e_{1}(x^{(i)})\) denotes the error of network at \(x^{(i)}\). At this point, extra care should be taken when setting up the second neural network, particularly concerning two key aspects.

#### ii.2.1 Magnitude of the second neural network

Considering that the original training data has a magnitude of \(O(1)\), then the training data for the second neural network, which is the residue \(e_{1}(x)\), would be much smaller than 1 (figure 1\(b\)). We observe that a neural network employing regular weight initialization methods, such as Xavier [34], often struggles to capture training data whose magnitude is significantly larger or smaller than 1 (see Appendix B). A straightforward solution to this issue is to normalize the training data by its

Figure 1: **Comparison of single-stage with multi-stage training**. \((a)\) Fitting of a neural network \(u_{0}(x)\) with tanh activation function to the data from (2). \((b)\) Fitting of second-stage neural network to the error \(e_{1}(x)\) between the data from (2) and the first-stage trained network \(u_{0}(x)\) as shown in \((a)\). \((c)\) The error \(e_{2}(x)\) between the data and the sum of two-stage networks, which reaches the machine precision of a single float (32-bit). \((d)\) Comparison of the loss convergence between a single-stage training (pink and red) and a two-stage training (black). For a single-stage training, the convergence rate of loss suddenly reduces (for Adam) after the loss reaches \(O(10^{-6})\) or terminates (for L-BFGS). For a two-stage training, even with less number of weights and biases, the convergence rate is significantly faster than that for the single-stage training.
root mean square value \(\epsilon_{1}\), defined as

\[\epsilon_{1}=\sqrt{\frac{1}{N_{d}}\sum_{i=1}^{N_{d}}[e_{1}(x^{(i)})]^{2}}=\sqrt{ \frac{1}{N_{d}}\sum_{i=1}^{N_{d}}[u_{g}^{(i)}-u_{0}(x^{(i)})]^{2}}\,. \tag{3}\]

Then, the normalized training data for the second neural network becomes \((x^{(i)},e_{1}(x^{(i)})/\epsilon_{1})\). Denoting the second trained network as \(u_{1}(x)\), the combined networks for the original data become

\[u_{c}^{(1)}(x)=u_{0}(x)+\epsilon_{1}u_{1}(x)\,. \tag{4}\]

Subsequently, we can continue training the third or even further neural networks to reach higher accuracy for our model. The training data for the \((n+1)\)-th neural network \(u_{n}\) is the residue \(e_{n}\) between the original training data \(u_{g}\) and the output of the previously combined \(n\) neural networks, \(u_{c}^{(n-1)}(x^{(i)})\), normalized by its own magnitude (root mean square value) \(\epsilon_{n}\), namely \((x^{(i)},e_{n}(x^{(i)})/\epsilon_{n})\). Then, the final model that combines all the \((n+1)\) neural networks reads,

\[u_{c}^{(n)}(x)=\sum_{j=0}^{n}\epsilon_{j}u_{j}(x)\,, \tag{5}\]

where \(\epsilon_{i}\) stands for the magnitude for the \(i\)-th neural network. When the original training data \(u_{g}\) is normalized, \(\epsilon_{0}\) is set to be 1.

#### ii.2.2 Frequency of the second neural network

Even with normalization, the second neural network, if initialized with regular weights, could still struggle to fit the high-frequency data due to the inherent spectral biases of neural networks. To

Figure 2: **Spectral biases of neural networks**. \((a)\) Fitting of neural networks with tanh activation function to the data from (6) for different \(m\). Under regular settings, neural networks have difficulty fitting high-frequency functions. \((b)\) Frequency domain of the function (6) for \(m=30\) with the dominant frequency \(f_{d}=5.5\). \((c)\) Derivative of the function (6) \(du/dx\) for \(m=30\), which scale as \(O(2\pi f_{d})\). \((d)\) Schematic diagram of a single-hidden layer neural network. \((e)\) Comparison between single-neuron outputs for different weights \(w^{(0)}\) within the tanh activation function and the function (6) for \(m=30\). To capture high-frequency functions, it shows that the weight within the activation function needs to increase from \(O(1)\) to \(O(2\pi f_{d})\).
illustrate this, we consider a target function

\[u(x)=\left(1-\frac{x^{2}}{2}\right)\cos\left[m\left(x+0.5x^{3}\right)\right]. \tag{6}\]

with \(m\) the free parameter related to the frequency of the function. Figure 2 shows the function (6) for \(m=3\), \(15\), and \(30\), respectively. For each \(m\), we generate \(300\) sample points \((x_{i},u_{i})\) that satisfy (6) as our training data, with \(x_{i}\) uniformly distributed in the domain \([-1,1]\). Figure 2 shows that the neural network, using regular weight initialization, fits the data well for \(m=3\), partially misses the data for \(m=15\), and completely fails to fit the data for \(m=30\).

To understand the challenge in fitting high-frequency data, let's consider a shallow neural network with a single input, single output, and one hidden layer that uses the hyperbolic tangent as its activation function:

\[u(x)=\sum_{i=1}^{N}w_{i}^{(1)}\,\tanh\left(w_{i}^{(0)}x+b_{i}^{(0)}\right)+b_{ 0}\,, \tag{7}\]

where \(w_{i}^{(0)}\) denote the weights between the input and hidden layers, and \(w_{i}^{(1)}\) are the ones between the hidden and output layers. \(b_{i}^{(0)}\) is the bias for the hidden units and \(b_{0}\) is the bias for the output unit. The magnitude of the output function is determined by \(w_{i}^{(1)}\), while \(w_{i}^{(0)}\), within the activation function, influence the local gradient of the function (figure 2\(e\)). Common practice involves initializing the weights of the network to follow a Gaussian distribution with zero mean and a specified variance. For example, Xavier initialization uses a specified variance \(V_{ar}=\sqrt{2/(N_{l-1}+N_{l})}\), where \(N_{l-1}\) and

Figure 3: **Neural network settings for high-frequency functions**. (\(a\)) Fitting of a neural network to the data from (6) with \(m=30\), by either changing the activation function for the first hidden layer to \(\sin(x)\), or multiplying the weight \(w_{i}^{(0)}\) before the first hidden layer by a scale factor \(\kappa\). None of them captures the data. (\(b\)) Comparison of training loss for the neural networks with different settings. (\(c\)) Neural network using \(\sin(x)\) activation function and modified scale factor \(\hat{\kappa}=35\) fits well the high-frequency data (6) with \(m=30\), reaching the same accuracy as (\(d\)) Fourier Let network [31].
\(N_{l}\) are the number of units in the preceding and succeeding layers, respectively. This initialization ensures that the variance of the sum of all unit outputs in each hidden layer remains \(O(1)\), which prevents gradient vanishing or explosion during the training. However, a side effect of this approach is that the neural network becomes a slowly varying function with respect to normalized inputs.

For a high-frequency function with normalized input and output, and a dominant frequency \(f_{d}\), the magnitude of its gradient scales as \(O(2\pi f_{d})\) (figure 2\(c\)). To capture these large gradients, considering a one-hidden layer network with single input and output (2.7), the weights \(w_{i}^{(0)}\) within the activation function need to increase from their initialized value of \(O(\sqrt{1/V_{ar}})\) to \(O(2\pi f_{d})\) during training. This large shift in weight values, particularly for large \(f_{d}\), leads to slower convergence during training or an inaccurate approximation of the data.

To address this issue, we multiply weights within the activation function by a large scale factor \(\kappa\)[35] to expedite the convergence of weights towards their optimal high value when fitting high-frequency data. We only multiply the scale factor \(\kappa\) to the weights between the input and the first hidden layer, rather than all weights, to prevent gradient explosion [36] during the training.

Besides large gradient, high-frequency functions also have a large amount of inflection points. In contrast, the hyperbolic tangent function, being a monotonic function, struggles to capture this feature. Periodic functions, such as the sine or cosine function, are more suitable choices for activation functions in this case [16]. In our approach, we use the sine function solely for the first hidden layer while retaining the hyperbolic tangent function for the remaining layers. This combination allows us to capture both low and high-frequency data effectively.

Figure 3 illustrates the impact of the scale factor \(\kappa\) and the choice of activation function on

Figure 4: **Importance of the modified scale factor \(\hat{\kappa}\). (\(a\)) Fitting of neural networks to the data from (2.6) for \(m=30\) using different modified scale factor \(\hat{\kappa}\). The networks start overfitting the data when \(\hat{\kappa}\geq 60\). (\(b\)) Training loss for the neural networks with different modified scale factors \(\hat{\kappa}\). When \(\hat{\kappa}\geq 60\), the training loss decreases significantly fast due to over-fitting. (\(c\)) Relation of the root mean square value \(\epsilon\) of the error \(e(x)\) between the trained network \(u(x)\) and target function \(u_{g}(x)\) with the modified scale factor \(\hat{\kappa}\). The minimal error is reached when \(\pi f_{d}<\hat{\kappa}<N_{d}/6\), where \(f_{d}\) denotes the dominant frequency and \(N_{d}\) the total number of training data points.**
improving the fit for high-frequency data. Using a combination of the scale factor \(\kappa\) and sine function for the first hidden layer yields the best training result. This combination equates to applying a Fourier feature mapping (i.e. Fourier let network) [31] to the input before it is passed through the multi-layer network. Figure 3(_c_&_d_) compares the convergence rate and final error of both methods when fitting the same high-frequency data. The results are consistently good, verifying the efficacy of both methods in fitting high-frequency data.

To expedite the convergence of weights from their initialized value \(O(\sqrt{1/V_{ar}})\) to the high gradient value \(O(2\pi f_{d})\) of a high frequency function, the optimal value of \(\kappa\) is expected to depend on the variance \(V_{ar}\), which is relevant to the weight initialization approaches, and the size of neural network. To isolate the impact of \(V_{ar}\) on determining the optimal value of the scale factor, we introduce a modified scale factor \(\hat{\kappa}\),

\[\hat{\kappa}=\kappa/\sqrt{V_{ar}}\,, \tag{2.8}\]

Figure 4(_c_) shows that the minimal fitting error is achieved when the modified scale factor is

\[\hat{\kappa}>\pi f_{d}\,,\qquad\text{namely}\qquad\kappa>\pi f_{d}\sqrt{V_{ar }}\,, \tag{2.9}\]

where \(f_{d}\) denotes the dominant frequency of the data. This finding is intuitive, as a scale factor \(\hat{\kappa}\) that meets the criterion (2.9) allows the neural network to directly capture the large gradient \(O(2\pi f_{d})\) of the high-frequency data.

However, setting \(\hat{\kappa}\) too high, close to the number of data points \(N_{d}\), results in overfitting of the neural network. Figure 4(_a_) shows that a neural network trained with a scale factor \(\hat{\kappa}=300\) to fit the training data (\(N_{d}=300\)) sampled from the high-frequency function (2.6) with \(m=30\) overfits the data. While the training loss is significantly small (figure 4_b_), the validation error is extremely large. To mitigate overfitting, figure 4(_c_) suggests that the modified scale factor \(\hat{\kappa}\) should be less than one-sixth of the total number of data points \(N_{d}\). As a rule of thumb, for the optimal fitting of high-frequency data, besides satisfying (2.9), the number of training data points \(N_{d}\) should also meet the criterion

\[N_{d}/6>\pi f_{d}\qquad\Longrightarrow\qquad N_{d}>6\pi f_{d}. \tag{2.10}\]

Given that the training domain is often normalized within \([-1,1]\), which contains \(2f_{d}\) dominant periods, the criterion (2.10) essentially requires a minimum of \(3\pi\approx 10\) data points within each dominant period \(1/(2f_{d})\), to ensure optimal fitting of the neural network to the high-frequency data. Without specific clarification, the criterion (2.10) is applied to all the example problems in this section.

### Algorithm of multi-stage training scheme for regression problems

Incorporating these key settings for higher-stage neural network training, we summarize a complete procedure of multi-stage training scheme for regression problems as shown in Algorithm 1.

```
1:Input: \(\hat{
\(3\times 30^{2}\approx 2700\)). This underscores that a larger neural network is not inherently advantageous; an appropriate training scheme is more vital and efficient for the reduction of training loss.

In fact, the power of multi-stage training scheme lies not only in boosting the convergence of training, but also fundamentally enabling neural networks to approximate a target function with arbitrary accuracy as required. We now convert the weights, biases and training data from single-float precision to double precision, and create the third and fourth-stage neural networks in accordance with Algorithm 1. Figure 5\((d)\) shows that the error \(e_{4}(x)\) between the sum of four-staged networks and the data successfully approaches the machine precision of a 64-bit double float. As long as higher-precision floating-point is used, the error can be further reduced with additional stages of training.

Figure 5\((c)\) shows that the overall convergence rate of the root mean square value \(\epsilon\) of the error \(e(x)\) between the network and data using multi-stage training scheme follows \(\epsilon\sim\exp(-\sqrt{n_{iters}}/25)\), closely approximating exponential decay. In contrast, the regular single-stage training only exhibits a linear decay, \(\epsilon\sim 1/n_{iters}\). That is to say, without considering the risk of being trapped in local minima, it would take at least \(O(10^{10})\) iterations for a single-stage training to reach an error of \(O(10^{-10})\). With the multi-stage training scheme, it only requires \(250^{2}\approx 6\times 10^{5}\) iterations to reach the same error, which is _four orders of magnitude faster_.

Moreover, we note that the number of data points required for higher-stage training also needs to increase following the criterion (2.10). Figure 5\((d)\) shows that the dominant frequency of residue after three-stage training can reach \(f_{d}=150\). This implies that the minimal number of training data to guarantee the success of the fourth-stage training needs to be \(N_{d}>6\pi f_{d}=2830\). Figure 5\((c)\) shows that the relation between the dominant frequency \(f_{d}\) and magnitude \(\epsilon\) of the residue after different stages of training empirically follows a power law

\[f_{d}\approx f_{0}\epsilon^{-\alpha},\quad\text{with the exponent }\alpha=1/6\,, \tag{2.11}\]

where \(f_{0}\) denotes the dominant frequency of the original training data. In practice, we anticipate a gradual increase in the frequency of the residue \(e(x)\) corresponding to the decrease in its magnitude, namely the exponent \(\alpha\) should be close to \(0\). Considering the error between a neural network \(u(x)\) and target function \(u_{g}(x)\) of magnitude \(\epsilon_{0}\) and dominant frequency \(f_{d}\), the error between the \(m\)-th derivative of \(\hat{u}(x)\) and \(u_{g}(x)\) becomes,

\[\frac{d^{m}}{dx^{m}}u(x)-\frac{d^{m}}{dx^{m}}u_{g}(x)=\frac{d^{m}}{dx^{m}}e(x) \sim(2\pi f_{d})^{m}\epsilon_{0}\sim\epsilon_{0}^{1-\alpha m}\,, \tag{2.12}\]

where we derive the last expression using (2.11). We find that when \(1-\alpha m>0\), even if the magnitude of error \(\epsilon_{0}\) is small, the error at high derivatives \(m>1/\alpha\) can still exceed \(1\). This indicates that
the trained neural network with high-frequency error tends to miss the high-derivative (\(m>1/\alpha\)) information of the target function underlying the training data. Hence, our goal is to achieve a smaller \(\alpha\) value during training, which enables the neural network to learn the high-derivative information from the data more accurately.

However, figure 5\((e)\) shows that, for regression problems, the exponent \(\alpha\) from multi-stage training

Figure 5: **Multi-stage neural networks**. \((a)\) Fitting of the first-stage neural network (red dashed curve) to the data from a given target function (blue curve). \((b)\) Training loss \(\mathcal{L}\) over the iterations based on multi-stage training scheme. \((c)\) Evolution of the root mean square value \(\epsilon\) of the error \(e_{n}(x)\) over the iterations, which follows \(\epsilon\sim\exp(-\sqrt{n_{iters}}/25)\), close to an exponential decay. However, for single-stage training (\(c\)-inset), the error convergence only follows a linear decay, \(\epsilon\sim 1/n_{iters}\). \((d)\) Fitting of higher-stage networks to the error of lower-stage training. Frequency domain of the error \(e_{n}(x)\) for different stages are shown in the right column. After four stages of training, the error between the data and combined networks is close to the machine precision of a double float (64-bit) \((e)\) Relation of the dominant frequency \(f_{d}\) and the root mean square value \(\epsilon\) of the error \(e_{n}(x)\) after different stages of training follows a power law (2.11) with an exponent \(\alpha\) independent of (i) target functions, (ii) neural network size, (iii) and the number of data points.
scheme appears to be universal, independent of both target functions and neural network settings. To further reduce \(\alpha\), high-derivative information about the target function would be required for the training. However, this information is often absent in regression problems, while it is readily available for physics-informed neural networks. The methodology of reducing the exponent \(\alpha\) for PINNs will be addressed in a later section (SS3.4).

Figure 6(_a-d_) shows that the multi-stage training scheme is equally applicable for 2D regression problems. The convergence rate of the loss function for the 2D problem roughly follows \(\epsilon\sim\exp(-\sqrt[3]{n_{iters}}/7)\) (figure 6_f_), slower than that for the 1D problem, but still much faster than the linear decay seen with regular single-stage training. Figure 6(_e_) shows that the relation between the dominant frequency \(f_{d}\) and the root mean square value \(\epsilon\) of the 2D residue \(e(x,y)\) follows the same power law (2.11) with the exponent \(\alpha\approx 1/6\).

Figure 6: **Multi-stage neural networks for a 2D target function**. (\(a\)) Fitting of first-stage neural network \(u_{0}(x,y)\) to the data from a 2D target function \(u_{g}(x,y)\). (\(b\)-\(d\)) Fitting of higher-stage networks \(u_{i}(x,y)\) to the error \(e_{i}(x,y)\) of lower-stage training. Frequency domain of the error at each stage is given. (\(e\)) Relation of the dominant frequency \(f_{d}\) with the root mean square value \(\epsilon\) of the error \(e_{n}(x,y)\) after different stages of training follows the same power law with the 1D problem, of which the exponent \(\alpha=1/6\). (\(f\)) Training loss \(\mathcal{L}\) over iterations of the multi-stage neural networks. The inset shows that the evolution of the root mean square error \(\epsilon\) over iterations for the 2D regression problem follows \(\epsilon\sim\exp(-\sqrt[3]{n_{iters}}/7)\), which is slightly slower than that for the 1D problem (see inset of (\(e\)).
Multi-stage training for physics-informed neural network

The multi-stage training scheme is particularly critical when we use neural network to approximate solutions governed by equations, where the demand for precision is high and essential for the usefulness of the solution. Here we apply the multistage idea to the physics-informed neural networks (PINNs) to improve their accuracy to machine precision. Unlike classical numerical method (i.e. finite difference) which can steadily enhance the accuracy of solution by reducing the grid size, PINNs cannot efficiently reduce solution errors merely by adding more collocation points or enlarging the neural network size, similar to the issue seen with regression problems (see A). This has made PINNs a less favored method for many scientific research that demands high-precision prediction. In this section, we show that the multi-stage training scheme can be extended to address this limitation of PINNs.

The general procedure of multi-stage training scheme for physics-informed neural networks (PINNs) mirrors that for regression problems (Algorithm 1). However, _two_ new challenges emerge when applying multi-stage training to PINNs. _First_, for regression problems, we can directly determine the magnitude \(\epsilon\) and dominant frequency \(f_{d}\) of the target function for each stage of training from the residue of lower-stage training. However, for PINNs, these two quantities are not readily obtainable because we lack the exact solution required to estimate the error of lower-stage training.

In addition, the loss function of PINNs involves both data loss and equation loss, defined as

\[\mathcal{L}=(1-\gamma)\mathcal{L}_{d}+\gamma\mathcal{L}_{e}\qquad \text{with} \tag{3.1}\] \[\mathcal{L}_{d}=\frac{1}{N_{d}}\sum_{i=1}^{N_{d}}[u(x_{i})-u_{i }]^{2}\quad\text{and}\quad\mathcal{L}_{e}=\frac{1}{N_{e}}\sum_{j=1}^{N_{e}}[r( x_{j},u(x_{j}))]^{2}, \tag{3.2}\]

where \(N_{d}\) represents the number of data points, commonly employed as the boundary condition, and \(N_{e}\) is the number of collocation points, which are utilized to examine the equation residue \(r(x,u)\) at various positions within the domain. In comparison to regression problems, \(\gamma\), known as the equation weight, is the additional hyper-parameter that balances the significance of the two losses during training. How to determine an appropriate value of \(\gamma\) for higher stages of training becomes the _second_ challenge. Using a simple example, we will demonstrate new algorithms to address these challenges and develop a modified multi-stage training scheme for PINNs.

### First challenge: magnitude and frequency of higher-stage network

As discussed in Section 2.2, the effectiveness of multi-stage training scheme depends largely on the optimal setting of the higher-stage neural networks \(u_{n}\), which is based on the magnitude and frequency of the residue \(e_{n}\) between the combined lower-stage networks and the ground truth \(u_{g}\). However, these pieces of information are not directly accessible for PINNs because we don't have the exact solution \(u_{g}(x)\) to the equation that is required to estimate the error \(e(x)=u_{g}-u_{0}\) of the lower-stage trained networks. Instead, the only information we have is the equation residue \(r(x,u_{0})\) associated with the trained first-stage network \(u_{0}\). Thus, understanding the relation between the equation residue \(r(x,u_{0})\) and the error \(e(x)\) of the lower-stage networks with the exact solution is crucial for determining the settings for the higher-stage training of PINNs.
#### 3.1.1 A simple example

We consider a first-order ordinary differential equation with the boundary condition

\[\frac{du}{dx}=u+x\qquad\text{with}\qquad u(0)=1\,, \tag{3.3}\]

which has the exact solution \(u_{g}(x)=e^{x}-x-1\). Figure 7\((a)\) shows the single-stage trained network \(u_{0}(x)\) to solve the equation (3.3) via PINN, which matches the exact solution \(u_{g}(x)\) well. The equation residue \(r_{1}(x,u_{0})\) associated with the network \(u_{0}(x)\) gives

\[r_{1}(x,u_{0})=\frac{du_{0}}{dx}-(u_{0}+x)\,, \tag{3.4}\]

which has the same dominant frequency with the error \(e_{1}(x)=u_{g}(x)-u_{0}(x)\) between the trained network \(u_{0}(x)\) and the exact solution \(u_{g}(x)\). However, the magnitude of equation residue \(r_{1}(x,u_{0})\) is one-order of magnitude larger than that of the error \(e_{1}(x)\). To elucidate their relations, we introduce the ansatz,

\[u_{g}(x)=u_{0}(x)+\epsilon_{1}u_{1}(x)\qquad\text{with}\quad e_{1}(x)= \epsilon_{1}u_{1}(x)\,, \tag{3.5}\]

where \(\epsilon_{1}\) denotes the magnitude of the error \(e_{1}(x)\), and \(u_{1}(x)\) becomes the normalized function within the domain. Substituting the ansatz (3.5) into (3.3) and re-arranging the equation gives

\[-\epsilon_{1}\left(\frac{du_{1}}{dx}-u_{1}\right)=\frac{du_{0}}{ dx}-(u_{0}+x)\,. \tag{3.6}\]

Recalling (3.4), the right-hand side of (3.6) is the equation residue. Thus, the relation between the prediction error \(e_{1}(x)\) and the equation residue \(r_{1}(x,u_{0})\) gives

\[-\epsilon_{1}\left(\frac{du_{1}}{dx}-u_{1}\right)=r_{1}(x,u_{0})\,, \tag{3.7}\]

which also becomes the governing equation for the second-stage training and \(u_{1}(x)\) is the second-stage neural network. The boundary condition of \(u_{1}\), based on (3.3) and (3.5), is

\[\epsilon_{1}u_{1}(0)=1-u_{0}(0)\qquad\Longrightarrow\qquad u_{1}(0)=\frac{1- u_{0}(0)}{\epsilon_{1}}\,. \tag{3.8}\]

Figure 7: **Comparison of prediction error with equation residue of PINNs.**\((a)\) Exact solution \(u_{g}(x)\) and neural network prediction \(u_{0}(x)\) to equation (3.3). \((b)\) Comparison of the equation residue \(r_{1}(x,u_{0})\) associated with the neural network prediction \(u_{0}(x)\) with the prediction error \(e_{1}(x)\) between \(u_{0}(x)\) and the exact solution \(u_{g}(x)\), which has different magnitude. \((c)\) The frequency domain of the equation residue \(r_{1}(x,u_{0})\) and the prediction error \(e_{1}(x)\), which has the same dominant frequency.
With the appropriate setting of the equation weight \(\gamma\) (as discussed in a later section SS3.2), the data loss of the first-stage training should be much smaller than that of the equation loss. This indicates that the error \(e_{1}(x)\), as well as \(u_{1}(x)\), has much smaller value at the boundary than within the domain. Namely, the boundary condition of \(u_{1}(0)\) can be considered as \(0\).

With zero boundary conditions, the magnitude and frequency of the solution \(u_{1}(x)\) are governed by the source function. For a linear equation, the dominant frequency of \(u_{1}(x)\) must be equal to that of the source function, namely the equation residue \(r_{1}(x,u_{0})\). Otherwise, the equation cannot be balanced in the frequency domain.

From (3.7), the magnitude \(\epsilon_{1}\) of the error \(e_{1}(x)\) also appears to be the same as that of the equation residue \(r_{1}(x,u_{0})\). However, this is only true when the solution \(u_{1}(x)\) is a low-frequency function. For a high-frequency function, its derivative, which represents its local gradient, becomes large and scales as \(O(2\pi f_{d})\), as discussed in Section 2.2, where \(f_{d}\) is the dominant frequency of the function. Given that \(u_{1}(x)\) shares the same dominant frequency with the equation residue \(r_{1}(x,u_{0})\), the magnitude \(\epsilon_{1}\) of the error \(e_{1}(x)\) between the network \(u_{0}(x)\) and exact solution \(u_{g}(x)\) can be determined by equating the magnitudes of the leading-order terms on both sides of the equation (3.7), which gives,

\[2\pi f_{d}\epsilon_{1}\sim\epsilon_{r_{1}}\qquad\Longrightarrow\qquad\epsilon _{1}=\frac{\epsilon_{r_{1}}}{2\pi f_{d}}\quad\text{with}\quad\epsilon_{r_{1}}= \text{RMS}(r_{1}(x,u_{0})) \tag{3.9}\]

where we use the root mean square (RMS) value \(\epsilon_{r_{1}}\) to represent the magnitude of the equation residue \(r_{1}(x,u_{0})\). Figure 7(\(c\)) shows that the dominant frequency for the equation residue \(r_{1}(x,u_{0})\) and prediction error \(e_{1}(x)\) are the same, around \(f_{d}\approx 1.5\). Based on (3.9), the magnitude \(\epsilon_{1}\) of the error should be \(2\pi f_{d}\approx 10\) times less than that of the equation residue, consistent with the result shown in figure 7(\(b\)).

#### 3.1.2 Magnitude and frequency estimation for general differential equations

To extend the relations between the properties of equation residue \(r_{1}(x,u_{0})\) and prediction error \(e_{1}(x)\) for general equations, we now consider a general form of ordinary differential equations

\[\mathcal{N}\left(x,u,u^{(1)},...u^{(m)}\right)=F(x)\qquad\text{with}\ \ u^{(i)}=\frac{d^{i}u}{dx^{i}}\ \ \text{for}\ \ i=1,2,...,m \tag{3.10}\]

where \(\mathcal{N}\) is a nonlinear differential operator that involves \(x\), \(u\) and its derivative \(u^{(i)}\) at different orders. \(m\) represents the highest order of derivative of \(u\) in the equation. \(F(x)\) is a source function with known expression. We denote \(u_{g}(x)\) as the exact solution to the equation and \(u_{0}\) the first-stage neural network prediction. By introducing the ansatz (3.5) and substituting into (3.10), we have

\[\mathcal{N}\left(x,(u_{0}+\epsilon_{1}u_{1}),[u_{0}+\epsilon_{1}u_{1}]^{(1)},...,[u_{0}+\epsilon_{1}u_{1}]^{(m)}\right)=F(x) \tag{3.11}\]

Considering that the first-stage neural network \(u_{0}\) captures the main variation of the exact solution \(u_{g}\), the magnitude \(\epsilon_{1}\) of the error \(e_{1}(x)\) between \(u_{g}\) and \(u_{0}\) would then be much smaller than one, namely \(\epsilon_{1}\ll 1\). In that case, the equation (3.11) can be rewritten in terms of a Taylor expansion of the nonlinear function \(\mathcal{N}\). After re-arrangement, it gives

\[-\epsilon_{1}\left(\left.\frac{\partial\mathcal{N}}{\partial u}\right|_{u=u_{0 }}u_{1}+\left.\frac{\partial\mathcal{N}}{\partial u^{(1)}}\right|_{u=u_{0}}u _{1}^{(1)}+...+\left.\frac{\partial\mathcal{N}}{\partial u^{(m)}}\right|_{u=u_ {0}}u_{1}^{(m)}\right)+O(\epsilon_{1}^{2})=\mathcal{N}\left(x,...,u_{0}^{(m)} \right)-F(x) \tag{3.12}\]

where \(u\) and its derivative \(u^{(i)}\) at different orders are considered as separate independent variables of the function \(\mathcal{N}\). Because \(\epsilon_{1}\ll 1\), all the nonlinear terms of \(u_{1}\) fall into the high-order \(O(\epsilon_{1}^{2})\) term,
and can generally be disregarded. This suggests that regardless of whether the original equation is linear or nonlinear, the governing equations for higher-stage networks essentially become _linear_ equations. This is a key factor that ensures the success of multi-stage training scheme for PINNs.

Since the right-hand side of (3.12) is the equation residue of \(u_{0}\), the final equation for \(u_{1}\) gives

\[\epsilon_{1}\left(\left.\frac{\partial\mathcal{N}}{\partial u}\right|_{u=u_{0} }u_{1}+\left.\frac{\partial\mathcal{N}}{\partial u^{(1)}}\right|_{u^{(1)}=u_{ 0}^{(1)}}u_{1}^{(1)}+...+\left.\frac{\partial\mathcal{N}}{\partial u^{(m)}} \right|_{u^{(m)}=u_{0}^{(m)}}u_{1}^{(m)}\right)=r_{1}(x,u_{0}) \tag{3.13}\]

or, in a short form,

\[-\epsilon_{1}\sum_{k=0}^{m}\beta_{k}u_{1}^{(k)}=r_{1}(x,u_{0}) \qquad\text{with}\quad\beta_{k}=\left.\frac{\partial\mathcal{N}}{\partial u^{ (k)}}\right|_{u=u_{0}}\quad\text{and}\quad u_{1}^{(k)}=\frac{d^{k}u_{1}}{dx^{ k}} \tag{3.14}\]

As mentioned earlier, if \(u_{0}\) is correctly trained, the boundary condition for \(u_{1}\) should be close to \(0\). Given that (3.14) is linear, the magnitude and frequency of \(u_{1}\) should be determined from the equation residue \(r_{1}(x,u_{0})\) by matching the magnitude and frequency of the dominant term (the term with the largest magnitude) on the left-hand side of (3.14) with that of \(r_{1}(x,u_{0})\).

Considering a physical equation with coefficients of similar scale before each term, and assuming \(u_{1}\) to be a high-frequency function with a dominant frequency far exceeding that of \(u_{0}\), the dominant term on the left-hand side of (3.14) is expected to be the one involving the highest-order derivative of \(u_{1}\), namely \(\epsilon_{1}\beta_{m}u_{1}^{(m)}\). We denote the dominant frequency of \(\beta_{m}\) and \(r_{1}(x,u_{0})\) as \(f_{d}^{(\beta)}\) and \(f_{d}^{(r)}\), respectively. Then, the dominant frequency \(f_{d}^{(1)}\) of \(u_{1}\) gives

\[f_{d}^{(\beta)}+f_{d}^{(1)}=f_{d}^{(r)}\qquad\Longrightarrow\qquad f_{d}^{(1 )}=f_{d}^{(r)}-f_{d}^{(\beta)}. \tag{3.15}\]

Given that \(\beta_{m}\) is a function only with respect to the lower-order network \(u_{0}\), and the dominant frequency \(f_{d}^{(1)}\) of the higher-stage network \(u_{1}\) is much larger than that of \(u_{0}\). Thus, even if \(\beta_{m}\) is a highly-nonlinear function, we have \(f_{d}^{(1)}\gg f_{d}^{(\beta)}\). Combined with (3.15), we have

\[f_{d}^{(1)}=f_{d}^{(r)}-f_{d}^{(\beta)}\approx f_{d}^{(r)} \tag{3.16}\]

where the dominant frequency of \(u_{1}\) is mainly governed by that of equation residue. With the dominant frequency determined, the magnitude \(\epsilon_{1}\) of the error between \(u_{g}\) and \(u_{0}\) can be derived by matching the magnitude of the term \(\epsilon_{1}\beta_{m}u_{1}^{(m)}\) and \(r_{1}(x,u_{0})\), which gives

\[\epsilon_{1}\cdot\epsilon_{\beta}\left[2\pi f_{d}^{(r)}\right]^{m}=\epsilon_{ r_{1}}\qquad\Longrightarrow\qquad\epsilon_{1}=\frac{\epsilon_{r_{1}}}{ \left[2\pi f_{d}^{(r)}\right]^{m}\epsilon_{\beta}} \tag{3.17}\]

\[\text{with}\qquad\epsilon_{\beta}=\text{RMS}(\beta_{m})\quad\text{and}\quad \epsilon_{r_{1}}=\text{RMS}(r_{1}(x,u_{0}))\]

where we use the root mean square (RMS) value \(\epsilon_{\beta}\) and \(\epsilon_{r_{1}}\) to represent the magnitude of \(\beta_{m}\) and equation residue \(r_{1}(x,u_{0})\), respectively.

The relations (3.16) and (3.17) can also be generalized to _partial_ differential equations, for which we need to calculate the dominant frequency of \(u_{1}\) with respect to each independent variable \(x_{i}\), namely

\[f_{d}^{(1,x_{i})}\approx f_{d}^{(r,x_{i})}\qquad\text{for}\quad i=1,2,...,N \tag{3.18}\]

where \(N\) is the total number of independent variables of the equation. Then, the magnitude \(\epsilon_{1}\) of the error between the first-stage network \(u_{0}\) and the exact solution \(u_{g}\) would be

\[\epsilon_{1}=\frac{\epsilon_{r_{1}}}{\epsilon_{\beta}\left[2\pi f_{d}^{(r,x_{1} )}\right]^{m_{1}}\cdot\left[2\pi f_{d}^{(r,x_{2})}\right]^{m_{2}}...\left[2 \pi f_{d}^{(r,x_{N})}\right]^{m_{N}}} \tag{3.19}\]
where \(m_{1}+m_{2}+...+m_{N}\) represents the highest order of partial derivative of \(u_{1}\) in the equation. For most equations, the relations (3.16)-(3.19) are sufficiently accurate to estimate the magnitude and frequency of the network for higher-stage PINN training. However, there are two types of nonlinear equations where these relations may not hold exactly. They are discussed in Appendix C.

#### 3.1.3 Importance of magnitude and frequency for higher-stage PINN training

The proper setting of the magnitude and frequency of a neural network, as shown to be essential for regression problems in Section 2.2, is critical for physics-informed neural networks, especially during

Figure 8: **Importance of rescaling PINN magnitude and frequency**. \((a)\) Exact solution \(u_{g}(x,y)\) to equation (3.20) and the general setting of PINNs for 4 cases. \((b)\) Evolution of the data loss and equation loss over the training of solving (3.20) via PINNs for different magnitude prefactor \(\epsilon_{0}\) and modified scale factors \(\hat{\kappa}_{0}\). \((c)\) Trained network \(u_{0}(x,y)\) of the solution to (3.20), the associated equation residue \(r_{1}(x,y,u_{0})\) and prediction error \(e_{1}(x,y)\) under different magnitude prefactor \(\epsilon_{0}\) and modified scale factor \(\hat{\kappa}_{0}\) (Case 1-4). The trained network with \(\epsilon_{0}\) and \(\hat{\kappa}_{0}\) from (3.18) and (3.19) gives the prediction with the lowest relative error \(e_{rr}\).
higher-stage training. To illustrate this, we use a Poisson equation with a high-frequency source function to represent the equation residue, and zero boundary conditions. This setup effectively mimics the governing equation for higher-stage PINN training. The equation reads

\[u_{xx}+u_{yy}=-\sin(6\pi x)\sin(6\pi y)\qquad\text{with}\quad u(x,\pm 1)=u(\pm 1,y)=0 \tag{3.20}\]

At first glance, one might assume that the solution \(u\) has the same order of magnitude with the source function, which is \(O(1)\). However, based on the analysis (3.18) in Section 3.1.2, the solution should have a dominant frequency \(f_{d}=3\) with respect to both independent variables \(x\) and \(y\). Given that the highest order of derivative in (3.20) is \(m=2\), the magnitude of the solution \(u\) can, then, be derived from (3.19), as \(1/[2\pi f_{d}]^{2}\sim O(10^{-3})\), as shown in the exact solution,

\[u_{g}(x,y)=\frac{1}{2(6\pi)^{2}}\sin(6\pi x)\sin(6\pi y)\,. \tag{3.21}\]

Figure 8 shows the neural network predicted solution \(u_{0}\) to (3.20) via PINNs under different setting of magnitude (via magnitude prefactor \(\epsilon_{0}\)) and frequency (via modified scale factor \(\hat{\kappa}_{0}\)). In these cases, we assume that the correct value of the equation weight \(\gamma\) is used. As shown in figure 8(\(b\)), only when both magnitude and frequency are correctly set in accordance with (3.18) and (3.19) does the neural network successfully converge to the exact solution at a rapid convergence rate.

#### 3.1.4 Algorithm for determining the solution magnitude for higher-stage PINN training

Besides the theoretical relations (3.16)-(3.19) derived in Section 3.1.2, we also develop a general algorithm to determine the magnitude of the solution to linear differential equations with high-frequency source functions and zero boundary conditions to mimic the higher-stage training of PINN. The algorithm can subsequently be combined with Algorithm 1 to extend the multi-stage training scheme for PINNs. The specific steps of the algorithm are given in Algorithm 2.

The principle underlying the algorithm is based on the fact that the dominant frequency of the solution \(u(\mathbf{x})\) mirrors that of the source function \(s(\mathbf{x})\). Therefore, the amplification effect of the derivative, attributable to the high-frequency property of the solution, can be well-estimated by taking \(s(\mathbf{x})\) as the guess solution. We define the ratio \(R\) as the magnitude of the differential operator \(\mathcal{N}_{s}\) relative to that of the source function \(s(\mathbf{x})\).

If \(R\) is larger than \(10\), the magnitude of the differential operator \(\mathcal{N}_{s}\) associated with the guess solution (3.23) much larger than the source function \(s(\mathbf{x})\). In that case, the magnitude of the guess solution should be reduced by decreasing \(\alpha\). Conversely, when \(R\) is less than \(0.1\), it suggests that differential operator \(\mathcal{N}_{s}\) associated with the guess solution (3.23) is too small. Hence, we should increase the magnitude of the solution by increasing \(\alpha\). The recursive relation (3.27) in Algorithm 2 is designed to achieve this objective. Here, the learning rate \(\eta\) is a user-defined positive hyper-parameter, which determines the rate at which \(\alpha\) and \(R(\alpha)\) converges to satisfy the criterion (3.26). Finally, the magnitude of the solution \(\epsilon\) can be estimated using (3.28).

Applying Algorithm 2 to equation (3.20), we obtain that \(\epsilon=1.41\times 10^{-3}\), which is very close to the magnitude of the exact solution, \(\epsilon=1/[2(6\pi)^{2}]=1.43\times 10^{-3}\). Here, we note that Algorithm (2) is mainly applicable to linear differential equations with a single dependent variable. For nonlinear equations or a group of differential equations with multiple dependent variables, a more advanced algorithm may be required to determine the magnitude for each variable, which is beyond the scope of this paper.
**Algorithm 2** Determine the magnitude \(\epsilon\) of solution to a linear equation with _high-frequency_ source function and **zero** boundary condition

### Second challenge: equation weight \(\gamma\) for higher-stage network

Equation weight \(\gamma\), as shown in (3.1), is a hyper-parameter to balance the contribution of data loss and equation loss in the loss function for physics-informed neural networks (PINNs). In the context of PINNs as a differential equation solver, boundary conditions are often implemented as data loss and the governing equations constitute the equation loss. Given that boundary conditions determine the uniqueness of the solution, a general rule of thumb is weighting the data loss more than the equation loss in the loss function [23]. This ensures that the boundary conditions are prioritized and satisfied during the training.

The relative contribution of data loss and equation loss in the loss function (3.1) is the ratio of the first term, \(I_{1}=(1-\gamma)\mathcal{L}_{d}\), to the second term, \(I_{2}=\gamma\mathcal{L}_{e}\) in (3.1), i.e. \(I_{1}/I_{2}\). For normalized linear differential equations that have low-frequency solutions, such as (3.3), the equation loss, \(\mathcal{L}_{e}\), remains the same order of magnitude with the data loss \(\mathcal{L}_{d}\), around \(O(1)\). In that case, by setting \(0.1<\gamma<0.5\), we can ensure that the contribution of data loss \(I_{1}>I_{2}\) in the loss function (3.1). However, this setting of \(\gamma\) does not hold for differential equations with high-frequency solutions. A
systematic mathematical justification was provided by a prior study [25], showing that the magnitude of the equation loss increases with the frequency of the solution. When considering the same Poisson equation (3.20) with zero boundary conditions and a high-frequency source function, similar to the equation we solve during higher-stage training, the magnitude of equation loss at the beginning of the training can be estimated as

\[\mathcal{L}_{e}=\frac{1}{N_{e}}\sum_{i=0}^{N_{e}}\left[u_{xx}+u_{yy}+\sin(6\pi x _{i})\sin(6\pi y_{i})\right]^{2}\sim O(1), \tag{3.29}\]

which is determined by the magnitude of the source function \(\sin(6\pi x)\sin(6\pi y)\). The magnitude of data loss at the beginning of training reads

\[\mathcal{L}_{d}=\frac{1}{N_{1}}\sum_{i=0}^{N_{1}}[u(x_{i},\pm 1)]^{2}+\frac{1}{ N_{2}}\sum_{j=0}^{N_{2}}[u(\pm 1,y_{j})]^{2}\sim O(u^{2}) \tag{3.30}\]

which is determined by the initial magnitude of the solution. Considering that the magnitude \(\epsilon_{0}\) of the solution has been estimated from the relation (3.19) or Algorithm 2, which gives \(\epsilon_{0}\approx 1/(6\pi)^{2}\) for the solution to (3.20), the magnitude of the data loss thus becomes

\[\mathcal{L}_{d}\sim O(\epsilon_{0}^{2})\sim O(10^{-6}) \tag{3.31}\]

which is six orders of magnitude smaller than the equation loss. If we still use \(\gamma\sim O(0.1)\), as \(I_{2}\gg I_{1}\), the optimization process will primarily focus on minimizing the \(I_{2}\) during the training, largely neglecting the contribution from the data loss. Utilizing the appropriate values of the magnitude prefactor \(\epsilon_{0}\) and the modified scale factor \(\hat{\kappa}_{0}\) from Section 3, figure 9\((b)\) shows the evolution of the data loss and equation loss over iterations by setting \(\gamma=0.5\). Compared with the equation loss, which was reduced by seven orders of magnitude in total, the data loss decays at a much slower rate, significantly limiting the errors of the trained neural network 9\((c)\) to be round 5%.

#### 3.2.1 Theoretical approach of determining \(\gamma\)

To improve the accuracy of trained network via PINN, we need to minimize the data loss prior to the equation loss. Therefore, the optimal value of \(\gamma\) should yield a larger contribution from the data loss larger than from the equation loss, namely,

\[I_{1}=(1-\gamma)\mathcal{L}_{d}\geq\gamma\mathcal{L}_{e}=I_{2}\qquad\Longrightarrow \qquad\gamma\leq\frac{\mathcal{L}_{d}}{\mathcal{L}_{e}+\mathcal{L}_{d}}, \tag{3.32}\]

which is consistent with the expression proposed in a prior study [25]. For equation (3.20), with the magnitude of the equation loss and data loss determined from (3.29) and (3.30) respectively, equation (3.32) yields \(\gamma=2\times 10^{-6}\). With this \(\gamma\), figure 9\((b)\) shows that both data loss and equation loss rapidly decrease over the training by more than five orders of magnitude. This suggests that both the equation and boundary condition are progressively satisfied by the network throughout the training. Although the reduction in equation loss is slightly less than in the case with \(\gamma=0.5\), the relative error \(e_{rr}\) between the trained neural network and the exact solution is reduced by more than one hundred times (figure 9\(a\)). However, generally without prior knowledge of the solution we do not know its corresponding \(\mathcal{L}_{d}\) and \(\mathcal{L}_{e}\), so estimating \(\gamma\) theoretically is difficult. Therefore below we develop an alternative approach.
#### 3.2.2 Algorithm for determining \(\gamma\) for general equations

Besides the theoretical expression (3.32), we also develop a more general algorithm to determine \(\gamma\) through a pre-training process. This approach provides higher accuracy and adaptability for a broad range of problems. As mentioned previously, the optimal value of \(\gamma\) should result in similar convergence rates for the data loss and equation loss over the course of training. We propose a heuristic approach for determining the optimal \(\gamma\), which is outlined in Algorithm 3.

The principle underlying the algorithm lies in estimating the _convergence rates_ of both data loss \(C_{d}\) and equation loss \(C_{e}\). This estimation involves calculating the ratio of the initial loss \(\mathcal{L}_{d}^{0}\) and \(\mathcal{L}_{e}^{0}\), to the respective losses \(\mathcal{L}_{d}^{m}\) and \(\mathcal{L}_{e}^{m}\) after a short period of pre-training. Here, we use the minimal value during the last \(10\%\) of the pre-training iterations to calculate \(\mathcal{L}_{d}^{m}\) and \(\mathcal{L}_{e}^{m}\) to counteract any potential spikes in the loss evolution.

If the convergence rate of the data loss \(C_{d}\) is substantially lower than that of the equation loss \(C_{e}\), it indicates that the \(\gamma\) used in training is too large and needs to be reduced. Vice versa. The recursive relation (3.37) is designed to reach this goal. \(\eta\) can be considered as the learning rate, a hyper-parameter that determines how fast \(R_{c}(\gamma)\) meets the criterion (3.36). We note that, when \(\gamma\) is

Figure 9: **Importance of the equation weight \(\gamma\) of PINNs. \((a)\) Equation residue \(r_{1}(x,y)\) and prediction error \(e_{1}(x,y)\) of solving (3.20) via PINN for different \(\gamma\). \((b)\) Evolution of data loss and equation loss over the training for different \(\gamma\) shown in \((a)\). The inset shows the loss evolution for different \(\gamma\) that satisfy the criterion (3.37), which are close to each other. \((c)\) The root mean square value \(\epsilon_{1}\) of the prediction error \(e_{1}(x,y)\) (lower panel), and the corresponding ratio \(R_{c}\) of the data loss convergence rate over that of equation loss (upper panel) as a function of \(\gamma\). The optimal range of \(\gamma\) with minimal prediction error corresponds to \(0.1<R_{c}<5\). Error bars show the standard deviation of five repetitive experiments with different random initialization.**
updated, one should re-train the neural network from the beginning to compute the updated \(R_{c}(\gamma)\), instead of continuing the previous training.

Apply Algorithm 3 to the equation (3.20) with \(N_{0}\) set to be 500, one gives \(\gamma\approx 10^{-4}\). Figure 9\((a)\) shows the trained network using this \(\gamma\), which reaches further higher accuracy than that using \(\gamma=2\times 10^{-6}\) from (3.32). Compared with the case of \(\gamma=0.5\) and \(\gamma=2\times 10^{-6}\), the convergence rate of both data loss and equation loss when using \(\gamma\) from Algorithm 3 are maximized (figure 9\(b\)), leading to the smallest errors of the neural network prediction.

The criterion range in (3.36) suggests that the training accuracy is not overly sensitive to the value of \(\gamma\), provided that the convergence rate of data loss \(C_{d}\) and equation loss \(C_{d}\) remains within the same order of magnitude. Figure 9\((c)\) shows the optimal range of \(\gamma\) that yields the minimal root mean square value \(\epsilon_{1}\) of the error \(e_{1}(x)\) between the trained network and the exact solution to (3.20). The corresponding range of \(R_{c}\) is found to be \(0.1<R_{c}<5\), aligning with the range in criterion (3.37). The inset of figure 9\((b)\) shows that the evolution of both data loss \(\mathcal{L}_{d}\) and equation loss \(\mathcal{L}_{e}\) for different \(R_{c}(\gamma)\) with this range (3.36) are closely matched. For \(\gamma=0.5\), the ratio \(R_{c}\), based on Algorithm 3, is found to be \(R_{c}(\gamma=0.5)=10^{-4}\), which largely deviates from the criterion, thus, resulting in large prediction error.

### Additional setting of PINN training for higher-stage network

Besides the most critical settings, i.e. magnitude prefactor \(\epsilon\), scale factor \(\kappa\) of frequency, and the equation weight \(\gamma\), there are other settings and advanced algorithm developed in the literature that can ensure the success of PINN training for the high-stage networks.
#### 3.3.1 Optimization method and re-sampling collocation points

Two other critical settings in the training of high-frequency function includes the selection of optimization method and the number of collocation point. Common choices for PINN training optimizer include Adam and L-BFGS, a second-order quasi Newton method. For general equations with low-frequency solution, L-BFGS is often the preferred optimization method. However, for equations with high-frequency solutions, this is not always the case. Figure 10\((a)\) presents a comparison of the loss evolution and final prediction error of trained network between using Adam and L-BFGS for solving the equation (3.20), where Adam shows a better overall convergence rate. Furthermore,

Figure 10: **Importance of re-sampling collocation points of PINNs**. \((a)\) Comparison of prediction error \(e_{1}(x,y)\) of solving (3.26) via PINNs for different number of collocation points \(N_{c}\) and using different optimizer, including L-BFGS, Adam with gradient descend (GD) (fixed collocation points) and Adam with stochastic gradient descent (re-sample collocation points over the iterations). \((b)\) The evolution of data loss and equation loss over the iterations for different optimizer using the number of collocation points \(N_{c}\) below or above (inset) the critical value \(N_{crit}\). \((c)\) The relation of the root mean square value \(\epsilon_{1}\) of the prediction error \(e_{1}(x,y)\) with the number of collocation points \(N_{c}\) for Adam (GD) and Adam (SGD). When the number of collocation points \(N_{c}\) is less than the critical value \(N_{crit}\approx 3100\), stochastic gradient descend can reach better performance for predicting high-frequency solutions. Error bars show the standard deviation of five repetitive experiments with different random initialization.
Adam has the added advantage of utilizing stochastic gradient descent (SGD) by re-sampling the collocation points every few iterations [37, 38], which shows a even higher convergence rate.

Collocation points in PINN training is as important as data points in regression problems. As discussed in Section 2.2.2, when training the neural network to fit high-frequency data, a sufficient number of data points (\(3\pi\approx 10\) per dominant period) are needed to ensure accurate predictions. This principle remains valid for PINN training. Unlike regression problems, which are limited by the availability of finite data points, PINN could potentially utilize as many collocation points as computationally feasible. For equation (3.20), the dominant frequency \(f_{d}=3\) in each dimension. Given that the domain is defined in \((x,y)\in[-1,1]\), there are 6 dominant periods in each dimension. Thus approximately \(N_{crit}=(3\pi\times 6)^{2}\approx 3100\) collocation points are required. Figure 10\((a)\) compares the accuracy of the trained network for different number of collocation points. For L-BFGS and Adam (GD) with fixed and small number of collocation points, the neural network predictions significantly deviate from the exact solution. When the number of collocation points reach the criterion \(N_{crit}\), the prediction error drops sharply and only improves marginally with the addition of more collocation points.

Compared with using fixed collocation points, predictions using Adam with stochastic gradient descent (SGD) are less sensitive to the number of collocation points. Figure 10\((a)\) shows that the prediction error using Adam (SGD) can attain optimal precision even when the number \(N_{c}\) of collocation points falls below the critical value \(N_{crit}\). Figure 10\((c)\) further compare the root mean square error (RMSE) \(\epsilon\) of the trained network between utilizing Adam (GD) and Adam (SGD) for different number of collocation points \(N_{c}\). It confirms that SGD is an essential tool in PINN training for predicting high-frequency solution.

#### 3.3.2 Advanced methods from the literature: RAR and gPINNs

Having discussed the essential settings, we note that many advanced algorithms developed in the literature can also largely improve the PINN training of higher-stage networks. Two of most useful methods we found are the adaptive residual-based collocation refinement (RAR) method [39, 40] and the gradient-enhanced physics-informed networks (gPINNs) [41].

A usual practice in PINN training is to uniformly distribute the collocation points across the domain. However, this approach proves inadequate for equations whose solution feature steep gradients [42]. As discussed in Section 2.2, high-frequency solutions exhibit large gradients throughout the domain. Despite setting a large scale factor to align with the gradient, there remain regions where the local gradient exceeds the averaged gradient \(O(2\pi f_{d})\) with a dominant frequency \(f_{d}\). It can be challenging to minimize the local residue of equation in these areas. To address this issue, we employ the residual-based refinement (RAR [39]) of collocation point. By continuing adding collocation points in areas of high equation residue throughout the training, the equation residue across the entire domain can be efficiently reduced. This technique thus becomes a vital tool for optimizing PINN training.

An additional method to boost the training performance of PINNs involves incorporating the gradient of the equation residue function \(r(x,u)\) into the loss function \(\mathcal{L}\), known as the gradient-enhanced physics-informed network (gPINN) [41]. Thus, the loss function can be expressed as,

\[\mathcal{L}=(1-\gamma)\mathcal{L}_{d}+\gamma(\mathcal{L}_{e}+\gamma_{g} \mathcal{L}_{g})\qquad\text{with}\quad\mathcal{L}_{g}=\frac{1}{N_{g}}\sum_{j= 1}^{N_{g}}|\nabla r(x_{j},u(x_{j}))|^{2}, \tag{3.38}\]

where \(N_{g}\) denotes the number of collocation points used to examine the gradient of the equation residue \(r(x,u)\) within the domain. \(\gamma_{g}\) is an additional hyper-parameter, akin to \(\gamma\), that control the balance between the equation loss \(\mathcal{L}_{e}\) and gradient loss \(\mathcal{L}_{g}\) during training.
By incorporating the gradient constraint \(\mathcal{L}_{g}\), we obligate the neural networks to learn the high-derivative information of the solution involved in the gradient of the equation. This can significantly improve the convergence rate of the training loss, provided we choose the appropriate value of the weight \(\gamma_{g}\) for the gradient constraint \(\mathcal{L}_{g}\). Analogous to (3.32), the value of \(\gamma_{g}\) can be estimated by

\[\mathcal{L}_{e}\geq\gamma_{g}\mathcal{L}_{g}\qquad\Longrightarrow\qquad\gamma_ {g}\leq\frac{\mathcal{L}_{e}}{\mathcal{L}_{g}}\sim\frac{||r||^{2}}{||\nabla r|| ^{2}} \tag{3.39}\]

where \(||\cdot||^{2}\) represents the \(l_{2}\)-norm. As discussed in Section 3.2.2, for the equation with high-frequency solutions, the equation residue has roughly the same frequency with the solution. Thus, the magnitude ratio of the equation residue \(||r||\) with its gradient should scale as \(||r||/||\nabla r||\sim O(2\pi f_{d})^{-1}\), where \(f_{d}\) is the dominant frequency of the solution. Thus, the optimal value of \(\gamma_{g}\) can be selected as

\[\gamma_{g}=\frac{||r||^{2}}{||\nabla r||^{2}}\sim O(2\pi f_{d})^{-2} \tag{3.40}\]

The effect of the gPINN on the higher-stage training is shown and discussed in Section 3.4.

### Algorithm of multi-stage training for PINNs

Leveraging the multi-stage training algorithm for regression problems and incorporating the results discussed in the previous sections, we have extend the multistage training scheme to physics informed neural networks (PINNs). The details of the algorithm are provided in Algorithm 4.

Here, we note that the primary distinction between the multistage training scheme for PINNs and that for regression problems lies in the fact that we _lack_ training data for the solution itself for PINNs. Contrasting with the multistage framework for regression problems, where the second network is trained directly using the error \(e_{1}=u_{g}-u_{0}\) between the first trained network \(u_{0}\) with the data \(u_{g}\), we don't necessarily have access to the error of the first trained network in the context of PINNs. Thus, the method of training the second network \(u_{1}\) for PINNs involves creating a combined network \(u_{k}^{(c)}\) (3.41) that involves the previously trained network \(u_{k-1}^{(c)}\) and a new network \(u_{k}(x,\kappa_{k})\), with an appropriately-estimated magnitude prefactor \(\epsilon_{k}\) and scale factor \(\kappa_{k}\). A key advantage of this approach is that it circumvents the need to derive a new equation, as shown in (3.13), for each higher-stage network. By fixing the trained weights and biases in the previous networks, the training process for solving the original equation becomes mathematically equivalent to solving the higher-stage governing equation (3.13) with the high-frequency source function from the equation residue for the lower-stage training.

Using Algorithm 4, figure 11 shows the three-staged PINN training for solving the ordinary differential equation (2.2). For the first two stages, we employ a combination of Adam and L-BFGS for training, which maximizes the convergence rate. However, given the high-frequency residue from the second stage of training, it indicates a high-frequency solution for the third stage of training. Thus, we only use Adam with stochastic gradient descent (SGD) to optimize the performance of the third-stage training, in accordance with the suggestions made in Section 3.3.1. By combining all the optimal settings as discussed in the previous sections (SS3.1-SS3.3), the prediction error at each stage can be reduced by 3-5 orders of magnitude within \(10^{5}\) iterations.

Compared with single-stage training, figure 12\((a\&b)\) shows that multi-stage training can reduce both the data loss and equation loss by more than 20 orders of magnitude within the same number of iterations. In this instance, the number of weights in the single-stage network has been selected to be approximately equivalent to the total number across all three-stage networks. These results suggest that employing appropriate network settings and an effective training scheme plays a more essential role in successful training than simply increasing the size of neural network and the number of
collocation points. Additionally, when combined with gPINN, the multistage training demonstrates an accelerated convergence rate. Figure 11(\(f\)) shows that, after the three stages of gradient-enhanced PINN training, the prediction error of the final trained network with the exact solution reaches the machine precision of double floating points The observed oscillation in \(e_{3}\) is primary attributable to round-off error.

The right panel of figure 11(\(d\)-\(f\)) displays the spectrum and dominant frequency of the equation residue after each stage of training. Figure 12(\(c\)) further shows the relation of the dominant frequency \(f_{d}\) with the root mean square value \(\epsilon\) of the prediction error \(e_{n}(x)\) over the stages, which follows a

Figure 11: **Multi-stage gPINNs for 1D equations**. (\(a\)) Comparison of the first-stage trained neural network (red dashed curve) with the exact solution \(u_{g}(x)\) (blue curve) to equation (3.20). (\(b\)) Data loss and (\(c\)) equation loss over iterations of three-stage training. The inset of (\(b\)) shows the evolution of the total loss \(\mathcal{L}\) over iterations. The inset of (\(c\)) shows that the evolution of the root mean square value \(\epsilon_{r}\) of the equation residue \(r(x,u)\) of the multi-stage neural networks follows \(\epsilon_{r}\sim\exp(-\sqrt{n_{iters}})\), which is consistent with that for regression problems (figure 5\(c\)). (\(d\) & \(e\)) Comparison of the higher-stage trained network with the error of lower-stage training is shown in the left column. The equation residue \(r_{n}(x)\) for different stages of training is in the middle. Frequency domain of the equation residue \(r_{n}(x)\) at each stage is shown in the right column. (\(f\)) Prediction error \(e_{3}(x)\) and the equation residue \(r_{3}(x)\) after the third-stage of the training. The zoom-in figure (on the right) shows fluctuations in the prediction error \(e_{3}(x)\), which is caused by the round-off error of the machine-precision of double-floating point.
power law \(f_{d}\sim\epsilon^{-\alpha}\) for both regression problems and PINNs. We recall that the power law exponent \(\alpha\) for regression problems is around \(1/6\). Compared with that, the power law exponent \(\alpha\) for PINNs becomes noticeably smaller, around \(1/7\) for multi-stage training with regular PINNs and reduced further to \(1/8\) when using gradient-enhanced PINNs. As discussed in Section 2.2, this indicates that trained neural networks in PINNs achieve higher accuracy in capturing higher-order derivatives compared to regression problems. This is reasonable as PINNs involve differential equations that contain the derivatives of the solution. By minimizing the equation loss, PINNs constrain both the neural network and its derivatives to approach the exact solution, enhancing the capture of the high derivative information of the solution. The same reasoning applies to the gradient-enhanced PINNs, which result in an even lower exponent \(\alpha\) since the gradient of equation residue involved further higher derivatives of the solution.
### Application to 2D partial differential equations

Multi-stage training scheme for PINNs can also be applied to solve partial different equations (PDE). Figure 13\((a)\) shows the three-stage training to solve the diffusion equation

\[\frac{\partial u}{\partial t}=\frac{\partial^{2}u}{\partial x^{2}}+(1-x^{2}+t) \qquad\text{with}\ \ u(x,0)=u(\pm 1,x)=0\,, \tag{3.44}\]

which has the exact solution,

\[u_{g}(t,x)=t(1-x^{2})\,. \tag{3.45}\]

Consistent with regression problems, the convergence rate of multistage PINN method for solving 2D problem is slightly slower than that for 1D problem (inset of figure 13\(e\)). After three stage of training using RAR method and gradient-enhanced PINN, the prediction error \(e_{3}(t,x)\) of the combined trained networks with the exact solution \(u_{g}(t,x)\) is around \(O(10^{-11})\). The accuracy of multistage training is still seven order-of-magnitude higher than that of the single-stage training. Figure 13\((e)\) shows that, when employing the multi-stage training scheme with gPINNs, the relation of the dominant frequency \(f_{d}\) with the root mean square value \(\epsilon\) of the prediction error \(e(x,y)\) for solving both 1D and 2D equations, follows the same power law (2.11) with an exponent of \(\alpha\approx 1/8\). This observation is consistent with the results observed in regression problems (figure 6\(e\)).

We note that achieving a low prediction error of \(O(10^{-11})\) for solving 2D partial differential equations via classical numerical method, such as finite difference, would require an extensive number of grid points. For instance, considering the central difference method along the \(x\)-direction, to reach \(10^{-11}\), we would need a grid size in the x-direction of \(h_{(x)}\sim O(\sqrt{10^{-11}})\sim O(10^{-5})\), namely \(10^{6}\) grid points for each time step. Even with a 4th-order Runge-Kutta method along the \(t\) direction, the step size in \(t\)-direction would need to be \(h_{(x)}\sim O(\sqrt{(10^{-11})}^{1/4}\sim O(10^{-3})\), requiring \(10^{3}\) time steps. Consequently, the total number of grid points needed to achieve this accuracy across the entire domain would be on the order of \(O(10^{9})\).

In contrast, our approach utilizes fully connected neural network with 4-hidden layers consisting of 30 units for each stage of the training. Thus, the total number of weights and biases used to express the solution is only around \(3\times 4\times 30^{2}\approx 10^{4}\), which is five order of magnitude less than the number of grid points used in a discretized solution. This demonstrates the efficiency and effectiveness of

Figure 12: **Comparison of single-stage with multi-stage PINN training**. \((a\) & \(b)\) Comparison of the data loss \((a)\) and equation loss \((b)\) evolution over iterations between the single-stage training, multistage training of PINN, and multistage training of gPINN. \((c)\) Relation of the dominant frequency \(f_{d}\) with the root mean square value \(\epsilon\) of the error \(e_{n}(x)\) after different stages of training for multi-stage training for regression problems (blue), PINNs (red) and gPINNs (black).
the multi-stage PINN in achieving accurate solution with significantly fewer parameters compared to classical numerical methods.

## IV Generalization of multistage PINN to a combined forward-and-inverse problem

Here we investigate a specific class of problems in mathematics that requires solving equations (forward problem) and simultaneously inferring unknown parameters in the equation (inverse problem) with a high demand for accuracy, for example, finding self-similar blow-up solutions for nonlinear fluid equations [9]. The physical significance of the problem was explained in Eggers (2015) [43], and a prior study [9] has developed the implementation of PINNs to solve it. In these problems, the multistage PINN method can play a critical role in achieving accurate results.

Here we focus on the 1D inviscid Burgers' equation for which we know the exact solutions. In Appendix D, we provide a summary of the background knowledge and PINN implementation. The task for the PINN involves discovering the _smooth_ solution to the nonlinear self-similar Burgers'

Figure 13: **Multi-stage gPINNs for 2D equations**. (\(a\)) Comparison of the first-stage trained neural network with the exact solution \(u_{g}(x,t)\) to the equation (3.44). (\(b\)-\(d\)) The error \(e_{n}(x,t)\) of higher-stage trained network \(u_{n}(x,t)\) with the exact solution \(u_{g}(x,t)\) is shown in the left panel. The equation residue \(r_{n}(x,t)\) for different stages of training and their frequency domains are shown in the right panel. (\(e\)) Relation of the dominant frequency \(f_{d}\) with the root mean square value \(\epsilon\) of the error \(e_{n}(x,t)\) after different stages of training follows the same power law as 1D problems, of which the exponent \(\alpha=1/8\). The inset shows that the number of iterations required for 2D problems to reach the same accuracy \(\epsilon\) is more than that for 1D problems.
equation,

\[-\lambda U+[(1+\lambda)y+U]\partial_{y}U=0,\qquad\text{with}\quad U(-2)=1\,, \tag{4.1}\]

where the solution \(U\) should be an odd function with respect to the independent variable \(y\), and \(\lambda\) is the unknown parameter to be predicted by PINNs. Smooth solutions to (4.1) exist when \(\lambda=1/(2+2i)\) for \(i=1,2,3...\). For other \(\lambda\) values, the solution is non-smooth, exhibiting discontinuity at certain order of its derivatives at the origin \(y=0\). Finding the smooth solution with the correct value of \(\lambda\) numerically is challenging.

To address the issue, prior study [9] leveraged PINNs and introduced an additional _smoothness_ constraint into the loss function, which penalizes the higher-order derivative of the equation residue around the non-smooth position. We note that the minimal order of derivative needed for the smoothness constraint depends on specific problems. In general, it should be larger than the order of smoothness for the non-smooth solution (see Appendix D). Any higher derivative with order larger than the minimal value can be involved in the smoothness constraint as long as it remains computationally feasible. Here we focus on the first smooth solution of the self-similar Burgers' equation (4.1). The loss function can be expressed as

\[\mathcal{L}=(1-\gamma)\mathcal{L}_{d}+\gamma(\mathcal{L}_{e}+ \gamma_{g}\mathcal{L}_{g})+\gamma_{s}\mathcal{L}_{s}\qquad\text{with} \tag{4.2}\] \[\mathcal{L}_{d}=(U(y=-2)-1)^{2}\qquad\text{and}\qquad\mathcal{L} _{e}=\frac{1}{N_{c}}\sum_{i=1}^{N_{c}}|r\left(y_{i},U(y_{i})\right)|^{2}\] (4.3) \[\mathcal{L}_{g}=\frac{1}{N_{c}}\sum_{i=1}^{N_{c}}\left|\frac{ \partial r}{\partial y}\left(y_{i},U(y_{i})\right)\right|^{2}\qquad\text{and} \qquad\mathcal{L}_{s}=\frac{1}{N_{s}}\sum_{j=1}^{N_{s}}\left|\frac{\partial^{ 3}r}{\partial y^{3}}\left(y_{j},U(y_{j})\right)\right|^{2}\] (4.4) \[\text{with}\qquad r(y,U)=-\lambda U+[(1+\lambda)y+U]\partial_{y}U \tag{4.5}\]

where \(\mathcal{L}_{d}\) and \(\mathcal{L}_{e}\) are the data loss and equation loss, respectively. \(\mathcal{L}_{g}\) is the gPINNs implementation, which involves the first-order gradient of the equation residue \(r(y,U)\). \(\mathcal{L}_{s}\) is the smoothness constraint that incorporates the third derivative of the equation residue. While the equation loss \(\mathcal{L}_{e}\) and gradient loss \(\mathcal{L}_{g}\) are examined at \(N_{c}\) random collocation points \(y_{i}\) across the entire domain, the smoothness constraint \(\mathcal{L}_{s}\) is calculated at \(y_{j}\) close to the origin (e.g. \(|y_{j}|<0.1\)) with number \(N_{s}\ll N_{c}\). Although the smoothness constraint depends on the equation residue, it can be viewed as an additional boundary condition for the solution to determine the value of \(\lambda\).

Following Algorithm 4, figure 14(\(a\)-\(d\)) shows the first two stages of training for solving the self-similar Burgers' equation (4.1). We observe that the second-stage training successfully improves both the prediction error \(e_{2}(y)\) of the trained network and the inferred lambda \(\lambda_{2}\) by _four orders of magnitude_. However, in addition to the high-frequency error as previously seen for the higher-stages, we observed that the prediction error \(e_{2}(y)\) from the second-stage training contains a low-frequency profile, which dominates over the high-frequency error. This disparity hinders the further reduction of the error by adding more stages of training based on Algorithm 4.

To understand the issue, we first study the occurrence of the high-frequency error in \(e_{2}(y)\). The middle panel in figure 14(\(e\)) reveals that the equation residue \(r_{2}(y)\) after the second stage of training exhibits a similar dominant frequency \(f_{d}\) to the high-frequency error in the prediction error \(e_{2}(y)\). Using (3.9), we estimate the magnitude \(\epsilon_{2}\) of the prediction error \(e_{2}(y)\) based on the magnitude \(\epsilon_{r_{2}}\) of the equation residue \(r_{2}(y)\) as \(\epsilon_{2}=\epsilon_{r_{2}}/(2\pi f_{d})\sim O(10^{-13})\). This is consistent with the magnitude of the high-frequency error in \(e_{2}(y)\). Then, following Algorithm 4, we create a new network \(U_{2}(y)\) multiplied by the magnitude of \(O(10^{-13})\) for the third-stage training. Figure 14(\(f\)) shows that the high-frequency error in \(e_{2}(y)\) after the third-stage training does vanish in \(e_{3}(y)\). However, the magnitude of the prediction error \(e_{3}(y)\) and the inferred \(\lambda_{3}\) after the third-stage of training (figure 14\(f\)) remains nearly the same as those of the second-stage training (figure 14\(e\)).
The issue appears to be related to the existence of the low-frequency profile in \(e_{2}(y)\). We recall that the prediction error of the training is estimated by comparing the trained networks at the

Figure 14: **Multi-stage gPINNs for a combined forward and inverse problem**. \((a)\) Comparison of the first-stage trained network \(U_{0}(y)\) at the inferred \(\lambda_{1}\) (red dashed curve) with the exact profile of the first smooth solution \(U_{g}(y)\) with \(\lambda_{g}=0.5\) to the self-similar Burgers’ equation (4.1). \((b)\) Data loss and \((c)\) equation loss over iterations of three-stage training. The inset of \((b)\) shows the relation of the error \(\epsilon_{\lambda}\) of inferred \(\lambda\) with the loss of the smoothness constraint \(\mathcal{L}_{s}\) after different stages of training. The dash line indicates the relation \(\epsilon=\sqrt{\mathcal{L}_{f}}\). The inset of \((c)\) shows that the evolution of the root mean square value \(\epsilon_{r}\) of the equation residue \(r(y,u)\) over iterations of the multi-stage neural networks follows \(\epsilon_{r}\sim\exp(-\sqrt{n_{iters}})\), consistent with that of regular forward problem (figure 11c). \((d\ \&\ e)\) The prediction error \(e_{n}(y)\) (left), equation residue \(r_{n}(y)\) (right) and its frequency domain (right) for the first \((d)\) and second \((e)\) stages of training. Comparison of higher-stage trained networks with the lower-stage prediction error is shown in the left panel. \((f)\) Prediction error \(e_{3}(y)\) and equation residue \(r_{3}(y)\) for the third-stage training using only one additional neural network \(U_{3}(y)\). It successfully reduce the high-frequency error from the second stage but fails to reduce its low-frequency error. \((g)\) Prediction error \(\hat{e}_{3}(y)\) and equation residue \(\hat{r}_{3}(y)\) for the third-stage training using two neural networks \(U_{3}(y)\) and \(U_{\lambda}(y)\), which successfully reduce both the high-frequency error associated with lower-stage equation residue \(r_{2}(y)\) and the low-frequency error associated with the error \(\epsilon_{\lambda}\) of inferred \(\lambda_{2}\). The zoom-in figure shows the prediction after three stages of training approaches the machine precision of double floating point.
inferred \(\lambda_{2}\) with the exact smooth solution at \(\lambda_{g}=0.5\). Therefore, the error of the trained network is influenced not only by the equation residue, but also by the error \(\epsilon_{\lambda}\) of the inferred \(\lambda\). To assess the impact of the inference error \(\epsilon_{\lambda}\) on the prediction error \(e_{2}(y)\), we perform a similar analysis as discussed in Section 3.1.1, introducing the ansatz of the exact solution \(U_{g}\) and exact value of \(\lambda_{g}\) as

\[U_{g}(y)=U_{0}(y)+\epsilon U_{e}(y)\qquad\text{and}\quad\lambda_{g}=\lambda_{0 }+\epsilon_{\lambda} \tag{4.6}\]

where \(U_{0}\) represents the lower-stage trained network and \(\lambda_{0}\) is the inferred \(\lambda\) from the lower-stage training. \(\epsilon U_{e}(y)\) represents the prediction error of the trained network and \(\epsilon_{\lambda}\) is the error of inferred \(\lambda\). Both \(\epsilon\) and \(\epsilon_{\lambda}\) are much smaller than \(1\). Substituting (4.6) into (4.1) and removing higher-order small terms \(O(\epsilon^{2})\), we have

\[\epsilon\left\{(\partial_{y}U_{0}-\lambda_{0})U_{e}+[(1+\lambda_{0})y+U_{0}] \partial_{y}U_{e}\right\}=\underbrace{(\lambda U_{0}-[(1+\lambda_{0})y+U_{0}] \partial_{y}U_{0})}_{\text{equation residue: }-r_{0}}+\underbrace{\epsilon_{ \lambda}(U_{0}-y\partial_{y}U_{0})}_{\text{term from }\epsilon_{\lambda}:\,r_{ \lambda}} \tag{4.7}\]

which can be viewed as the governing equation for the higher-stage network. In addition to the equation residue \(r_{0}(y)\) from the lower-stage training, the higher-stage equation (4.7) involves a new source function \(r_{\lambda}(y)\) that is associated with the error \(\epsilon_{\lambda}\) of the inferred \(\lambda\). While the equation residue \(r_{0}(y)\) exhibits high-frequency behavior, the source function \(r_{\lambda}(y)\) is influenced by the profile of the trained network \(U_{0}(y)\), exhibiting the low-frequency profile in the prediction error \(e_{2}(y)\) (figure 14_e_).

Considering the low frequency nature of the source function \(r_{\lambda}(y)\), the magnitude of prediction error \(\epsilon\) in (4.7) associated with \(r_{\lambda}(y)\) is expected to be similar to the error \(\epsilon_{\lambda}\) of the inferred \(\lambda\), which is approximately \(O(10^{-12})\), consistent with our results (figure 14_e_). In contrast, the prediction error associated with the high-frequency equation residue \(r_{0}\), as discussed earlier, is only around \(O(10^{-13})\). This explains why the low-frequency profile dominates the prediction error \(e_{2}(y)\).

Here, we note that the error \(\epsilon_{\lambda}\) of inferred \(\lambda\) is calculated using the known exact value \(\lambda_{g}=0.5\). However, in many other problems, the exact value of \(\lambda_{g}\) is unknown. Thus, an alternative way to quantify the inference error \(\epsilon_{\lambda}\) is from the loss \(\mathcal{L}_{s}\) of the smoothness constraint. The inset of figure 14(_b_) shows that the inference error \(\epsilon_{\lambda}\) after different stages of training is proportional to \(\sqrt{L_{s}}\), i.e. \(\epsilon_{\lambda}=\sqrt{L_{s}}\) (dashed line). This suggests that we can use \(\sqrt{L_{s}}\) to estimate the inference error \(\epsilon_{\lambda}\), as well as the magnitude prefactor for the higher-stage network \(U_{\lambda}\) associated with \(r_{\lambda}(y)\).

Since the prediction error is dominated by the low-frequency source function \(r_{\lambda}(y)\), one might intuitively consider creating a single low-frequency network multiplied by the error \(\epsilon_{\lambda}\) of the inferred \(\lambda\) for the third-stage training. However, this approach is not effective because the smoothness constraint (4.4) depends on the higher-order derivative of the equation residue. By using only a low-frequency network, it would be challenging to reduce the high-frequency equation residue. Therefore, our proposed solution is to create two networks for both source functions in (4.7) at the third-stage training, namely

\[\epsilon_{2}U_{2}(y)+\epsilon_{\lambda}U_{\lambda}(y), \tag{4.8}\]

where the magnitude prefactor \(\epsilon_{2}\) and modified scale factor \(\hat{\kappa}\) (for frequency) for the high-frequency network \(U_{2}(y)\) associated with the equation residue \(r_{2}\) can be determined by the relations (3.17) and (3.16) or Algorithm 2. The low-frequency network \(U_{\lambda}\) associated with the error \(\epsilon_{\lambda}\) of the inferred \(\lambda_{2}\) can be directly multiplied by the inference error \(\epsilon_{\lambda}\). Figure 14(_g_) shows that, using combined two networks for the third-stage training, the prediction error \(\hat{e}_{3}(y)\) is successfully reduced by another three orders of magnitude, eventually approaching the machine precision of double-floating points.

## 5 Discussion

We note that the principle of multi-stage neural networks is similar to that of Fourier series, which combines a series of sine or cosine functions, ranging from low to high frequencies, to approximate
functions. Provided the series converge, the error between the Fourier series expansion of a given order and the target function possesses lower magnitudes but higher frequency than any terms in the series. To further minimize the error, higher-order sine or cosine functions need to be incorporated into the series, leading to additional higher-frequency error.

The introduction of new neural networks in multi-stage neural networks (MSNNs) is analogous to the inclusion of higher-order trigonometric functions in the Fourier series expansions. However, in contrast to sines and cosines, deep neural networks with appropriate settings offer stronger function representation capacity. Our finding indicates that the magnitude \(\epsilon\) of error after each stage of training follows a inverse power law relation with the dominant frequency \(f_{d}\) of the error, i.e. \(\epsilon\sim f_{d}^{-\alpha}\), with the exponent \(\alpha\approx 1/6\) for regression problems, \(\alpha\approx 1/7\) for regular PINNs, and \(\alpha\approx 1/8\) for gradient-enhanced PINNs (gPINNs).

In comparison, the power law exponent \(\alpha\) for Fourier series is roughly around \(\alpha\approx 0.5\)[44], much larger than that of MSNNs. This indicates that, to achieve the same error magnitude, the error frequency generated by MSNNs could be several orders of magnitude lower than that by Fourier series. This observation confirms that MSNNs serve as a superior tool capable of accurately approximating target functions, as well as their high derivative information.

The multi-stage neural networks (MSNNs) developed in this work remains in their _early_ stage, and mainly serve as a proof of concept to demonstrate that neural networks can practically achieve high accuracy. It is crucial to recognize that MSNNs should not be regarded as a substitute for classical numerical methods, but rather as a complementary approach. In fact, there remains several challenges that need to be addressed in the MSNN method. One of the primary challenges pertains to high-dimension problems. As shown in figure 6 and 13, the convergence rate of MSNNs for both 2D regression problems and PINNs are consistently slower than that for 1D problems. It is expected that this challenges will become more pronounced in higher-dimensional problems.

The second major challenge pertains to approximating functions or predict solutions with steep gradients. Near the regions where the target function exhibits steep gradients, neural networks often encounter local peaks in the error or the equation residue during training. The presence of these peaks hinder the reduction of error in successive stages, necessitating their removal before proceeding to the next stage of training. We note that functions with steep gradients are commonly encountered in differential equations, such as stiff equations, nonlinear equations, or singular perturbation equations (see Appendix C). Solving these types of equation via PINNs is beyond the scope of this paper.

There are additional questions to be addressed that could further improve the MSNN method. One of the critical questions concerns the optimal timing for transitioning to the next stage of training. In each stage, the convergence rate of training loss gradually decreases over the iterations. The decision whether to switch to the next stage quickly for higher convergence rates, or to stay in the current stage until the loss plateaus in order to maximize the error reduction at each stage requires careful consideration and further investigation. Moreover, with multiple stages of training, oversized networks are no longer required to achieve high accuracy within a single stage of training. The optimal strategy for selecting the neural network size at different stages that can minimize the numbers of training parameters (weights and biases) and thus computational expense for the entire MSNN training becomes another future direction that is worth investigating.

## 6 Conclusion

We introduced the multi-stage neural networks (MSNNs) for both regression problems and physics-informed neural networks. Inspired by perturbation theory, we sequentially introduced new stages of training with new neural networks to capture the residue from the previous stage of training. This enable MSNNs to reach unprecedentedly high accuracy over stages. We showed that three stages of MSNN training can reach machine precision, making neural networks truly universal func
tion approximators in practice. This new method can be widely applied to many scientific domains, such as mathematical and nonlinear physical science where the precision matters.

The success of MSNNs lies in two aspects. The first is the idea of staged training itself. Deep neural networks often suffer from spectral biases, making it challenging to capture the full spectrum of the target function in a single stage of training, even when employing large-sized networks with an increased number of data or collocation points. As a result, the training loss tends to plateau after a certain number of iterations. However, by employing multi-stage training, the previously plateaued error can be substantially reduced in each successive stage, which enables MSNNs to progressively capture finer details of the target function.

The second aspect for the success of MSNNs is the specific design of each new-stage network based on the training error from the previous stage. The neural-network predictions in successive stages exhibits significantly small magnitude and high frequency compared to the previous stages. We showed that, by employing optimal magnitude prefactor \(\epsilon\) and scale factor \(\kappa\) with sin activation function, accurate predictions of functions with small magnitudes and high frequencies can be achieved. This enables the effective capture of intricate features in each successive stage.

To maximize the performance of each stage of training, we also studied the optimal value of \(\epsilon_{n}\) and \(\kappa_{n}\) for each stage. For regression problems, the \(\epsilon_{n}\) is equal to the magnitude (root mean square value) of the error \(e_{n}(\mathbf{x})\) between the trained networks in the previous stage and the ground truth \(u_{g}(\mathbf{x})\), and \(\kappa_{n}\) is proportional to the dominant frequency \(f_{d}\) of the error \(e_{n}(\mathbf{x})\).

However, for physics-informed neural networks (PINNs), the prediction error \(e_{n}(\mathbf{x})\) is not directly available and needs to be inferred from the equation residue \(r_{n}(\mathbf{x})\) from the previous stage of training. Based on the fact that the governing equations for higher-stage training are essentially linear, we provided the theoretical relations between the magnitude and frequency of the prediction error \(e_{n}(\mathbf{x})\) and the equation residue \(r_{n}(\mathbf{x})\) in Section 3. We also presented an algorithm that can effectively estimate the magnitude of the prediction error from the equation residue.

Moreover, we discussed several other optimal settings that can enhance the efficiency of multi-stage PINN training. These include the equation weight \(\gamma\), number of collocation points \(N_{c}\), choice of optimizer, and advanced PINN techniques in the literature, such as RAR method and gPINNs.

Leveraging all the optimal settings discussed in this work, we showed that multi-stage neural networks (MSNNs) can significantly reduce the prediction error for both regression problems and PINNs, approaching machine precision. Furthermore, MSNNs showcases their capability in solving combined-forward-and-inverse problems to machine precision, a task typically challenging for classical numerical methods, but of great importance in mathematical and physical sciences. However, there remains many questions and challenges to be addressed for further enhancing the MSNN method.

## Acknowledgements

We thank T. Buckmaster and J. Gomez-Serrano for helpful discussions regarding the application of multi-stage neural networks to critical mathematical questions. We acknowledge the Office of the Dean for Research at Princeton University for partial funding support via the Dean for Research Fund for New Ideas in the Natural Sciences. C.-Y.L acknowledge the National Science Foundation for funding via Grant No. DMS-2245228. We also gratefully acknowledges financial support from the Schmidt Data X Fund at Princeton University made possible through a major gift from the Schmidt Futures Foundation.
## Appendix A Neural network error under different settings

Systematic experiments (figure 15) show that the root mean square value (RMS) \(\epsilon\) of the error \(e(x)\) between the trained network \(u_{0}(x)\) and the data from the target function \(u_{g}(x)\) remains unchanged even when the number of either layers (figure 15\(a\)) or units (figure 15\(b\)) is increased. Although the RMS error \(\epsilon\) does slightly decrease with an increase in training data (figure 15\(c\)), this reduction is smaller than the standard deviation of eight repetitive experiments conducted with different random initializations, and thus is negligible. These results suggest that the plateau in training loss is not due to insufficient neural network size or lack of training data, but instead arises from inherent limitations of the training process itself. Figure 15\((d)\) presents the training loss for two different optimization methods. Compared to Adam [27], a first-order gradient descent method, L-BFGS [45], a quasi-Newton method, exhibits a higher convergence rate. However, training with L-BFGS quickly falls into a local minimum after reaching the same plateau as Adam. This suggests that the loss plateau is not optimizer-specific.

## Appendix B Effect of data magnitude on neural network training

Figure 16: Fitting of neural networks to the data with different magnitudes without normalization. It shows that the network is hard to fit data with magnitude either too much larger or smaller than 1.

Figure 15: Root mean square (RMS) value \(\epsilon\) of the error between the target function and trained network using different number of \((a)\) hidden units, \((b)\) layers, and \((c)\) training data, and \((d)\) different types of optimizers, which show no big difference. Error bars show the standard deviation of eight repetitive experiments with different random initialization.
## Appendix C Two extreme types of equations

There are two extreme types of equation where the general settings of networks derived in (3.19) and (3.18) for the high-stage training do not strictly hold. The _first_ case is when the equation involves nonlinear term with high-order of derivatives, for example,

\[\left(\frac{d^{8}u}{dx^{8}}\right)^{2}-u=F(x).\] (C.1)

Substituting the ansatz (3.5) into (C.1) gives

\[-\epsilon\left(2\frac{d^{8}u_{0}}{dx^{8}}\frac{d^{8}u_{1}}{dx^{8}}-u_{1} \right)-\epsilon^{2}\left(\frac{d^{8}u_{1}}{dx^{8}}\right)^{2}=r(x,u_{0})= \left(\frac{d^{8}u}{dx^{8}}\right)^{2}-u_{0}-F(x)\] (C.2)

where \(r(x,u_{0})\) is the equation residue of \(u_{0}\). When \(u_{1}\) is a high-frequency function with dominant frequency \(f_{d}\) satisfying the criterion,

\[(2\pi f_{d})^{8}\epsilon>1\qquad\Longrightarrow\qquad f_{d}>\epsilon^{-1/8},\] (C.3)

the nonlinear term of \(u_{1}\) on the left-hand side of (C.2) is no longer negligible and becomes the dominant term in the equation. The magnitude and frequency of \(u_{1}(x)\), thus, need to be reassessed by balancing the nonlinear term with the equation residue \(r(x,u_{0})\). This results in the dominant frequency \(f_{d}^{(1)}\) of \(u_{1}(x)\) to be \(f_{d}^{(1)}=f_{d}^{(e)}/2\), rather than \(f_{d}^{(1)}=f_{d}^{(e)}\) from (3.16), where \(f_{d}^{(e)}\) represents the dominant frequency of the equation residue \(r(x,u_{0})\).

However, we note that, although the dominant frequency \(f_{d}^{(e)}\) of equation residue is larger, it still capture the order of magnitude of the actual frequency of \(u_{1}(x)\). We recall from figure 4\((c)\) that neural networks with modified scale factor \(\hat{\kappa}\) larger than the criterion \(\hat{\kappa}>\pi f_{d}\) can reach the same high-accuracy of fitting to high-frequency functions. This indicates that setting the scale factor \(\kappa\) based on a larger dominant frequency \(f_{d}^{(e)}\) for the network of \(u_{1}\) remains a good option to solve (C.2).

The _second_ case is when the equation involves singular perturbation term, for example

\[\alpha\frac{d^{4}u}{dx^{4}}+\frac{d^{2}u}{dx^{2}}-u=F(x)\qquad\text{with} \quad\alpha\ll 1.\] (C.4)

where the coefficient before the highest-order derivative term is much smaller than the others. This type of equation is very common in physical sciences, such as boundary-layer problem. Substituting the ansatz (3.5) into (C.4) gives

\[-\epsilon\left(\alpha\frac{d^{4}u_{1}}{dx^{4}}+\frac{d^{2}u_{1}}{dx^{2}}-u_{1 }\right)=r(x,u_{0})=\alpha\frac{d^{4}u_{0}}{dx^{4}}+\frac{d^{2}u_{0}}{dx^{2}} -u_{0}-F(x)\] (C.5)

where \(r(x,u_{0})\) is the equation residue of \(u_{0}\). Based on (C.5), the dominant frequency \(f_{d}^{(1)}\) of \(u_{1}(x)\) remains equal to that of the equation residue \(r(x,u_{0})\). However, when \(\alpha<(2\pi f_{d})^{-2}\), the dominant term on the right-hand side of (C.5) is not the one with the highest-order derivative of \(u_{1}\), but the term with the second-order derivative. The magnitude \(\epsilon\) of the error is, then, determined by

\[\epsilon=\frac{\epsilon_{r}}{\left[2\pi f_{d}^{(1)}\right]^{2}},\qquad\text{ rather than}\quad\frac{\epsilon_{f}}{\left[2\pi f_{d}^{(1)}\right]^{4}\alpha},\] (C.6)

which is based on (3.17). Here, we note that the actual challenges of solving singular perturbation equation via PINNs is more than the violation of the expression (3.17) for setting the higher-stage
neural network. According to asymptotic analysis, the existence of singular perturbation term in the equation indicates that the solution to the equation has a narrow inner region where local gradient is very large. This property makes both the first-stage and higher-stage training of networks difficult. More discussion of the challenge is given in the Discussion (Section 6) of the paper. However, the solution to this challenge is beyond the scope of this paper.

## Appendix D 1D inviscid Burgers' equation

This section summarizes the exact self-similar blow-up solutions to the 1D inviscid Burgers' equation [43] and the PINN implementation developed in Wang _et al._[9] to find it numerically.

Without viscous dissipation, the 1D inviscid Burgers' equation is given as

\[\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}=0.\] (D.1)

which has a shock wave solution where the velocity becomes discontinuous at a finite time, exhibiting a singularity where the solution blow up. However, right before the time when the shock/singularity is formed, the velocity profile remains smooth and continuous, and follow a self-similar structure near the singularity formation. We suppose the singularity occurs at \(t=t_{0}\) and \(x=x_{0}\). The self-similar coordinates can be written as

\[s=-\log(t_{0}-t),\qquad y=\frac{x-x_{0}}{(t_{0}-t)^{1+\lambda}},\] (D.2)

where \(s\) and \(y\) are the local time and spatial coordinates, respectively. When \(s\) goes to infinity, time \(t\) approaches to the blow-up time \(t_{0}\), but can never go beyond that. In the meantime, the new self similar coordinate \(y\) allow us to zoom into and examine the solution profile around the singularity as time approaches \(t_{0}\). The solution \(u\) follows the ansatz [9]

\[u(x,t)=(t-t_{0})^{\beta}U(y,s)\] (D.3)

where \(U(y,s)\) indicates the self-similar profile near the singularity with \(\beta\) to be determined. Substituting the ansatz (D.3) into the equation (D.1) gives \(\beta=\lambda\). Thus, the self-similar form of the Burgers' equation becomes

\[(\partial_{s}-1)\lambda U+[(1+\lambda)y+U]\partial_{y}U=0.\] (D.4)

We assume that when approaching the blow-up time \(t_{0}\), namely \(s\) goes to infinity, time derivative term in (D.4) vanishes, and the self-similar profile \(U(y,s)\) reaches steady state. Then, the steady state profile \(\tilde{U}(y)\) is governed by

\[-\lambda\tilde{U}+[(1+\lambda)y+\tilde{U}]\partial_{y}\tilde{U}=0.\] (D.5)

For simplicity, we consistently use \(U\) to represent the steady state solution in the rest of the section. The parameter \(\lambda\) (D.2), the rate at which singularity forms, remains unknown and is the key parameter to be inferred via the multi-stage neural networks.

To guarantee the equation (D.5) is well-posed globally in the local coordinates, the self-similar solution \(U\) must be an odd function. Theoretically, there exists solutions to (D.5) for each value of \(\lambda\). The analytic solutions to the self-similar Burgers' equation (D.5) are

\[y=\begin{cases}-U-CU^{1+\frac{1}{\lambda}}&\text{for}\quad x\geq 0\\ -U+C(-U)^{1+\frac{1}{\lambda}}&\text{for}\quad x<0\end{cases}\] (D.6)
where \(C\) is a constant determined by the boundary condition. Here, we use \(U(y=2)=-1\), which gives \(C=1\). From the analytic expression (46), we can see that \(\lambda\), in fact, determines the smoothness of the solution. Here, the smoothness indicates the solution is continuous at all its derivative. When \(\lambda=1/(2+2i)\) with \(i=0,1,2,...\), the solution is smooth everywhere in the domain. However, when \(\lambda\neq 1/(2+2i)\), the expression (46) involves fractional power, causing the solution to be non-smooth at the origin. For example, figure 17\((c)\) shows that the fourth derivative of the solution for \(\lambda=0.4\) is discontinuous at the origin. Here, we note that the non-smooth solutions have no physical meaning. Thus, finding the smooth solutions to (41) is the goal.

Prior study by Wang _et al._[9], leveraged the continuous property of neural networks, showing that PINN can discover the smooth solution with associated \(\lambda\) by imposing the high-order derivative constraint at the non-smooth position, known as the _smoothness_ constraint. Additionally, we impose odd symmetry of the solution \(U\) by constructing the function form \(U=y[\text{NN}_{u}(y)+\text{NN}_{u}(-y)]\), where \(\text{NN}_{u}\) indicates a fully-connected neural network created for \(U\). The _data loss_ and _equation loss_ for solving the Burgers' equation (45) are given as

\[\mathcal{L}_{d}=(U(y=-2)-1)^{2}\qquad\text{and}\qquad\mathcal{L}_{e}=\frac{1}{ N_{c}}\sum_{i=1}^{N_{c}}|r\left(y_{i},U(y_{i})\right)|^{2} \tag{47}\]

\[\text{with}\qquad r(y,U)=-\lambda U+((1+\lambda)y+U)\partial_{y}U \tag{48}\]

where \(y_{i}\) indicates the random collocation points in the training domain \(y\in[-2,2]\) and \(N_{c}\) is their total number. Here we focus on finding the first smooth solution with known \(\lambda_{g}=1/2\). Utilizing the fact that the non-smooth solutions in the neighborhood of \(\lambda=0.5\) has unbounded fourth-order derivative, which appears in the third order of derivative of equation residue, the smoothness constraint is given as

\[\mathcal{L}_{s}=\frac{1}{N_{s}}\sum_{i=1}^{N_{s}}\left|\frac{d^{3}r}{dy^{3}} \left(y_{i},\,U(y_{i})\right)\right|^{2}\,. \tag{49}\]

where \(y_{i}\) indicates the random collocation points close to the origin (e.g. \(|y_{i}|<0.1\)) and \(N_{s}\) is their total number. Although the smoothness constraint depends on the equation residue, it can be simply considered as a boundary condition for the solution that help determine the value of \(\lambda\).
