Physics-Informed Neural Networks: Minimizing Residual Loss with Wide Networks and Effective Activations

1 Introduction

The success of deep learning in a wide variety of tasks has motivated its application in scientific domains as well [Sirignano and Spiliopoulos, 2018; Reiser et al., 2022; Li et al., 2020]. PINNs [Raissi et al., 2017] in particular are designed to solve differential equations as an alternative to traditional solvers, benefiting from discretization-free construction and the vast availability of machine learning tools and techniques. As a result, PINNs have been deployed in various physics and engineering problems, such as solving inverse scattering problems in photonics [Chen et al., 2020a], flow problems in fluid dynamics [Cai et al., 2021], and computational neuromusculoskeletal models raised in biomedical and rehabilitation sciences [Zhang et al., 2022].

∗Corresponding Author

Consider a general-form PDE with a Dirichlet boundary condition such asD [u] (x) = f(x) x ∈ Ω

where u is the solution of the PDE on a bounded domain Ω ⊂ Rd of d independent variables with boundaries ∂Ω, f and g are known functions, and D is a differential operator. Here, the operator D expresses the physical rules governing u through a differential expression. PINNs are then trained to respect the underlying physical dynamics given in D by minimizing the residual loss

where ˆu is a neural network approximation of u, l is an error function such as squared error, and x is a set of training collocation points in Ω. To guarantee a unique solution, boundary (and initial) conditions are also imposed by adding other supervised loss terms, referred to as boundary loss, trained with boundary data sampled from ∂Ω. The resulting loss function can then be treated as a multi-objective optimization task [Raissi et al., 2017].

While proven effective, training PINNs is often a challenging task. These challenges usually originate from either the discrepancy between the residual loss and the boundary loss [Wang et al., 2020b; Farhani et al., 2022; Wang et al., 2020a] or the nature of the residual loss [Krishnapriyan et al., 2021; Wang et al., 2022]. Notably, as Eq. 2 involves differentiation over a neural network with respect to (w.r.t.) its inputs, the outputs of the network undergo a significant structural transformation. To better understand the aforementioned process and its implications, let us consider a simple differential operator ∂u ∂x of a single independent variable and an L-layer feedforward network with an activation function σ. One can find that the application of this differential operator on the neural network, D[ˆu], is given by

where Wi and Gi are weights and linear outputs of layer i, and ◦ and × denote element-wise (Hadamard) and matrix products. In contrast, the original neural network ˆu can be defined recursively as

Eq. 3 shows how the differentiation transforms the outputs of a neural network. Firstly, note that the simple recursive relation between the layers of a feed-forward network is disrupted by applying D, and additional element-wise products with weights emerge as well. More significantly, we observe that the derivative σ′ of the activation function appears in the outputs. This presence of σ′ in PINNs highlights the importance of an activation function with well-behaved derivatives in the model's expressive power in learning D and likewise in the optimization process as it involves higher-order derivatives of σ.

Altogether, the distinct characteristics of D[ˆu] and the resulting residual loss, contribute to a problem that is quite different from common supervised training tasks. Consequently, the existing theory around loss functions and their characteristics cannot readily be applied to PINNs, and the lack of understanding about PINNs and their optimization dynamics leads to uninformed design choices and suboptimal performance even for seemingly easy PDEs.

In this work, we focus on the residual loss and its landscape. Specifically, we are interested in finding what neural networks and design choices enable PINNs to globally minimize the residual loss. To this end, we study the residual loss at a critical point of the network parameter space and search for distinctive characteristics of a global minimum compared to other critical points. Once these characteristics are identified, our investigation shifts towards determining sufficient conditions within the network design, in particular, width and activation function, that guarantee the existence of global minima in the parameter space. Our findings underscore the importance of the width and activation functions with wellbehaved high-order derivatives in acquiring a high expressive power in learning the differential operator. Finally, we verify our findings by conducting extensive experiments on several PDEs.

We summarize our contributions as follows. (1) We theoretically show that the residual loss of PINNs can be globally minimized, given a two-layer neural network with a width equal to or greater than the number of collocation points. (2) Through our analysis, we establish that the residual loss for a k-th order differential operator is optimally minimized when using an activation function with a bijective k-th order derivative. We leverage this theoretical foundation as a guideline for selecting activation functions, justifying the choice of sinusoidal activations, and subsequently validating their effectiveness through empirical demonstrations and experiments.

2 RelatedWorks

2.1 Wide Neural Networks

Wide neural networks have historically been of significant interest in machine learning. With classical results such as Universal Approximation and Gaussian processes, and more recently, NTK theory [Jacot et al., 2018], wide networks have been studied to understand neural networks in certain regimes

[Chen et al., 2020b; Lee et al., 2019]. The optimal width of a neural network is also studied for convergence guarantees [Oymak and Soltanolkotabi, 2020; Du et al., 2019; Allen-Zhu et al., 2019; Nguyen and Mondelli, 2020] and loss geometry [Safran and Shamir, 2016; Nguyen and Hein, 2017] with certain classes of neural networks and optimizers. The convergence guarantees are often provided for a width polynomial in the number of training samples and the number of layers [Allen-Zhu et al., 2019]. Safran and Shamir [2016] studied the basins of the loss function for wide twolayer ReLU networks, showing that wider networks are initialized at a good basin with higher probability. [Nguyen and Hein, 2017] also showed that most of the critical points in a wide neural network are also global minima. The developed theory in most of the aforementioned works cannot be directly applied to PINNs, as they either rely on specific neural network formulations [Nguyen and Hein, 2017; Nguyen and Mondelli, 2020], which are disrupted by differentiation, or certain hyper-parameters that are not effective for PINNs, such as ReLU activation function [Du and Hu, 2019; Safran and Shamir, 2016; Allen-Zhu et al., 2019].

2.2 Periodic Activation Functions

Sitzmann et al. [2020] proposed using sinusoidal activation functions in neural networks with low-dimensional inputs for learning differentiable signals. Notably, they also showed the capability of Sine networks in solving Wave and Helmholtz PDEs with PINNs. Since then, few works have explored the behaviour of neural networks with periodic activation functions at initialization [Belbute-Peres and Kolter, 2022] and their expressive power as function approximators [Meronen et al., 2021]. Meronen et al. [Meronen et al., 2021] studied the inductive bias introduced by periodic activation functions on the neural network functional space, and showed that such networks are less sensitive to input shifts.

2.3 Physics-Informed Neural Networks

Besides the applications of PINNs in solving various PDEs, there has been a surge in analyzing the behaviour and pitfalls of PINNs in recent years, especially from the optimization perspective [Wang et al., 2020a; Liu et al., 2020; Farhani et al., 2022]. Using Neural Tangent Kernel (NTK) theory from infinitely wide neural networks, [Wang et al., 2020b] showed that high-frequency terms in a PDE result in discrepancies in the convergence rate between the loss objectives when trained with Gradient Descent, leading the model to exhibit behaviours similar to spectral bias [Rahaman et al., 2018]. Wang et al. [2020a] also showed similar results by studying the magnitude of the loss gradients at different layers. Later on, the momentum term was shown to address the discrepancy in optimization in the infinite-width regime [Farhani et al., 2022].

Many recent works alleviate the optimization challenges in PINNs and improve their performance by assigning weights to each loss term [Wang et al., 2020b; McClenny and BragaNeto, 2020; Wight and Zhao, 2020], designing new architectures and embeddings [Wong et al., 2022; Wang et al., 2020a; Wang et al., 2021; Dong and Ni, 2021], and using sophisticated training strategies such as curriculum learning [Krishnapriyan et al., 2021; Wang et al., 2022]. Among them, [Wang et al., 2020b] and [Wong et al., 2022], suggested mapping the inputs to random or trainable Fourier features and the use of sinusoidal activation functions to overcome the spectral bias and the convergence discrepancy.

3 Global Minima of the Residual Loss

In this section, we study the residual loss at its critical points to obtain sufficient conditions for the existence of global minima. We present the lemmas and theorems for a simple differential operator with a single independent variable and generalize in the Supplementary Material 1 . First, we introduce the notation and the setup used throughout this section.

3.1 Notation and Setup

We use ˆuW : Rd ×R|W| → RnL to denote an L-layer feedforward neural network parameterized byW = {Wi, bi | 1 ≤ i ≤ L,Wi ∈ Rni-1×ni , bi ∈ Rni }, where ni is the number of neurons in layer i, n0 = d, and nL = 1. We drop W from ˆuW for simplicity if there is no ambiguity. The neural network ˆu for an input x = (x1, . . . , xd) is formulated as

where σ is an activation function, and F0(x) = x. We further define F (k) i (x) as

where σ(k) is the k-th derivative of σ. In the case of k = 1, we simply use F ′ i and σ′ instead. For a batch x of N samples, Fi(x), Gi(x), and F (k) i (x) are N × ni matrices. Also, the matrix power W k represents an element-wise power.

In a PINN, the neural network ˆu is trained to approximate the solution u of a differential equation denoted as in Eq. 1. In this work, we assume that D is a linear differential operator, i.e., the PDE is linear in the derivatives of u and u itself. We reformulate the residual loss in Eq. 2 to be a function of weights W and choose l(r) = r2 .

Throughout the rest of this section, we consider a two-layer neural network and a simple k-th order differential operator D [u] = ∂ k u ∂xk for a single independent variable x (i.e., d = 1).

3.2 Residual Loss of a Two-layer PINN

To study the residual loss and its critical points, we first need to derive the analytic formula for the residual loss and its gradients. The next two lemmas, provide us with these tools by finding the differentiation D [ˆu] and then deriving the gradients of the resulting residual loss.

1 Supplementary Material containing the generalized analysis, proofs, and additional experiments can be found at https://arxiv.org/ abs/2405.01680.

Lemma 1. For a two-layer neural network ˆu defined in Eq. 4, and a k-th order differential operator D [u] = ∂ k u ∂xk of a single independent variable x, D[ˆu] is

With the analytic formula for D[ˆu] in hand, it is easy to plug it into Eq. 5 to get the residual loss. The next lemma derives the gradients ∇W2 ϕr(x;W) of the residual loss w.r.t. the weights of the last layer.

Lemma 2. For ˆu and D[ˆu] given in Lemma 1, gradients of the residual loss w.r.t. the weights of the second layer over the training collocation data x of N samples are given by

Remark 1. Lemmas 1 and 2 generalize the appearance of derivatives in the outputs of the neural network as in Eq. 3, showing that a k-th order differential term similarly contains the k-th derivative of the activation function. Thus, activation functions with vanishing high-order derivatives, such as ReLU, significantly reduce the network representation power in approximating the residuals. Note that the gradients w.r.t. W1 contain the (k+1)-th derivative of the activation function, further highlighting the importance of well-behaved derivatives in optimization.

In the following section, the gradients given in Lemma 2 are studied at a critical point to find the characteristics of global minima of the residual loss. Note that global minimum in this context refers to the parameters that make the loss zero.

3.3 Critical Points of Wide PINNs

We are eventually interested in finding sufficient conditions for the existence of a global minimum of the residual loss, i.e., ϕr(x;W) = 0. The following theorem takes the first step by providing a necessary condition for globally minimizing the residual loss. We then turn this requirement into a sufficient condition by establishing a set of assumptions. Note that the squared error l(r) is a non-negative convex function of the residuals r, and l′ (r) = 0 results in l(r) = 0. Thus, a critical point W of ϕr(x;W) in the parameter space globally minimizes the residual loss if l′ (D[ˆuW ](x)-f(x)) = 0 for every training sample in x.

Theorem 1. For ˆu and D[ˆu] as in Lemma 1, a critical point W of the residual loss ϕr(x,W) is a global minimum if the following conditions are satisfied:

1. Weights W1 of the first layer are strictly non-zero,

Theorem 1 distinguishes the global minima from other critical points of the residual loss. However, there is no guarantee that an arbitrary neural network can satisfy the conditions in this theorem. In other words, a critical point that makes F (k) 1 (x) full row rank may not exist in the parameter space of a neural network at all. Still, this theorem does give out a necessary condition for such a neural network. Since F (k) 1 (x) is an N ×n1 matrix, the width n1 of the first layer must be at least N for it to be full row rank. In fact, given other assumptions, the next theorem shows that n1 ≥ N is also a sufficient condition for the existence of a global minimum. Note that the first condition on W1 is satisfied with a high probability in a continuous high-dimensional parameter space.

We first define the non-degenerate critical points used in the next theorem and establish a set of assumptions that connect the two theorems together.

Definition 1 (Non-degenerate Critical Point [Nguyen and Hein, 2017]). For a function f ∈ C 2 : U ⊂ Rn → R (i.e., f has continuous second-order derivatives), a critical point x = (x1, . . . , xn) ∈ Uof f is non-degenerate if its Hessian matrix at x is non-singular. Furthermore, x is non-degenerate on a subset of variables s ⊂ {x1, . . . , xn} if the Hessian w.r.t. only the variables in s is non-singular at x.

Assumptions 1. For the collocation training data x of N points, the activation function σ in ˆu, and the k-th order differential operator defined in Lemma 1, we assume that

1. samples in x are distinct,
2. σ(k) is a continuous and strictly monotonically increasing function, and
3. σ(k) is a bounded function with an infimum of zero.

Theorem 2. With Assumptions 1 holding and for D [ˆu] as in Lemma 1, if n1 ≥ N, then every critical point W of ϕr(x; W) that is non-degenerate on {W2, b2} is a global minimum of ϕr .

The following remark allows us to make our final conclusion from Theorem 2.

Remark 2. As explained in [Milnor et al., 1965] and [Nguyen and Hein, 2017], for a function f in C 2 that maps an open subset U ⊂ Rn to R, the degenerate critical points in U are rare as the set of all degenerate points has Lebesgue measure zero.

Theorem 2 provides sufficient conditions for a global minimum of the residual loss with a wide network of width N or higher. Based on this theorem, if a PINN with a width of at least N has a non-degenerate critical point, then it also has a global minimum for the residual loss. Since the residual loss ϕr(x; ·) is a function from R|W| to R and has continuous second derivatives, the degenerate critical points are rare, and the wide PINN in Theorem 2 has a global minimum.

Note that the residual loss is a strong regularizer that results in a data-efficient training process. Thus, PINNs are often trained with O(1000) collocation points and even fewer boundary data [Raissi et al., 2017; Krishnapriyan et al., 2021]. Therefore, the constraint on the width is well within the practical settings of neural networks. Furthermore, as we observe in the experiments in Section 5, while satisfying the constraint on the width improves the performance, one can expect relatively good results with smaller width as long as the other conditions in Assumptions 1 are almost satisfied.

4 On the Choice of Activation Function

The conditions outlined in Theorem 2 and Assumptions 1 collectively establish an important set of necessities for achieving global minimization of the residual loss. Notably, the requirement of strictly monotonically increasing σ(k) implies that it should be a bijection, providing an important guideline in choosing effective activation functions for PINNs. It is noteworthy that bijective activation functions are widely prevalent in deep learning, and extending this characteristic to their derivatives for improved expressiveness in representing differential operators is a plausible goal.

However, the activation functions frequently used in deep learning do not satisfy the bijection property even for the first-order derivatives. As depicted in Figure 1, only Softplus has a bijective first-order derivative, and as we show in Section 5, it indeed improves the performance of the firstorder Transport PINN significantly. Meanwhile, there has been an increasing interest in the use of sinusoidal functions either as feature embeddings [Wong et al., 2022; Wang et al., 2021] or activation functions [Sitzmann et al., 2020; Belbute-Peres and Kolter, 2022] for PINNs.

As shown in Figure 2, we observe that the linear outputs of the layers in a neural network with the Sine activation function at initialization are centred at zero with low variance when initialized with normal Xavier initialization. Sitzmann et al. [2020] also proposed a uniform initialization scheme for Sine networks that produces normal linear outputs at all layers with a desired variance. Consequently, in both cases, most of the linear outputs of the layers lie in the [-π/2, π/2] interval where Sine is bijective. Furthermore, as we train the PINNs with the Sine activation function, we observe that layers still exhibit the same behaviour, i.e., most of the linear outputs of the layers are between -π/2 and π/2 after convergence, especially as the width grows larger. Figure 3 illustrates the output distributions for each layer of the trained Wave and Klein-Gordon PINNs (We later define these equations in Section 5).

The observations above suggest that the sinusoidal functions can be utilized to almost satisfy the bijective condition of the activation function. Specifically, we use Cosine and Sine activation functions to train PINNs with first- and second-order terms, respectively. As a result, as long as the width is adequately large to produce low-variance preactivations within [-π/2, π/2], the first-order terms in Cosine networks and the second-order terms in Sine networks are determined with the bijective interval of Sine. The same approach can be taken when solving PDEs with higher odd or even terms.

The experiments in the next section show that sinusoidal non-linearity greatly improves the performance of PINNs compared to the common Tanh activation, and the gains are often greater as the width grows. We note that while the Assumption 1.3 facilitates the proof of the Theorem 2, the crucial property is the bijective σ(k) , and we relax the assumption on the infimum of the derivatives.

5 Experiments

In this section, we provide numerical results for several PDEs, revealing the impact of the activation functions and the width. We first experiment with the first-order Transport equation, comparing the Softplus and Cosine activation functions with Tanh. Next, we study three second-order PDEs using Sine and Tanh activation functions. We empirically show that Sine significantly improves the performance of PDEs with secondorder terms with a noticeable decrease in error as the width exceeds the number of training samples.

In all of the experiments, we use a three-layer feed-forward network with a width varying from 64 neurons up to 1024 neurons and initialized with Normal Xavier initialization. All the models are trained with normalized inputs for 80K epochs using the Adam optimizer and an exponential learning rate decay scheme. The only exception is the Wave equation, for which the models are trained for 120K epochs for better con- Neuron Neuron vergence. We repeat each experiment three times with a different random initialization and report the average and the best results.

Figure 4: Transport Equation. Top row: Exact solution, Middle row: Predicted solution, Bottom row: Absolute error

5.1 Transport Equation

The transport equation is a first-order linear PDE that describes a quantity as it moves through time and space. We experiment with the one-dimensional equation with the following formulation:

We also impose a periodic boundary condition u(t, 0) = u(t, 2π) and a Dirichlet initial condition u(0, x) consistent with the solution in [Krishnapriyan et al., 2021]. The PINN is then trained with 256 collocation training points and 200 boundary samples. To verify our results in Section 3.3, we choose Softplus as the activation function as it is a smooth version of ReLU with a bijective first derivative equal to the Sigmoid function. We also use Cosine, as it results in a Sine network for the first-order terms.

Table 1 reports the mean absolute errors for the transport equation. In all cases, Softplus and Cosine perform significantly better than Tanh, as shown in Figure 4. Furthermore, as the width becomes equal to the number of collocation samples (256), we observe a noticeable decrease in the absolute error. The same improvement is also evident in the training curve of the residual loss as shown in Figure 5, where the wide models follow a steep curve. For the PINNs with a width of 256 or wider with both Cosine and Softplus activation functions, the absolute errors are between 10-2 and 10-3 , outperforming the reported 1.1 × 10-2 absolute error achieved with curriculum learning in [Krishnapriyan et al., 2021].

5.2 Wave Equation

The wave equation describes mechanical and electromagnetic waves and has the following form in 1-D:

Here c is the velocity of the wave. For c = 1 and the solution

we train the PINN with 512 collocation points for t, x ∈ [0, 1] and impose the initial and boundary conditions below with 256 boundary data points:

The residual loss training curves in Figure 7 highlight the impact of the width in training PINNs, with wider models consistently achieving smaller loss values, and the 512- and 1024-neuron models following almost the same path. This behaviour is similar to the curves in Figure 5 for the transport equation, where all the models with a width of at least N achieve very close loss values. As evident in Table 2, PINNs with both Tanh and Sine activation functions perform notably better than narrow PINNs, with a mean absolute error of 3.11 × 10-2 and 5.62 × 10-2 for Tanh and Sine respectively. Figure 6 illustrates the inability of the narrow Tanh network compared to Sine and wide Tanh models in representing the solution. Also, while narrow Sine PINNs are able to find good solutions, the training is more unstable and the performance is worse on average compared to wider models.

5.3 Helmholtz Equation

We consider a 2D Helmholtz equation of the following form with t, x ∈ [-1, 1] as in [Wong et al., 2022]:

With zero boundary conditions, the solution is given by

The PINN is then trained with 512 collocation points and 256 boundary data, using Sine and Tanh as the activation function. As reported in Table 2, Sine performs remarkably better than Tanh across all widths. Similar to the Wave and Transport equations, there is also a slight decrease in the errors as the width exceeds the number of collocation points. As in the Wave equation, PINNs with the Tanh activation function start to perform better as the width grows larger. However, as illustrated in Figure 8, even the 1024 neurons-wide Tanh network is still unable to capture the solution, while the Sine network finds an acceptable solution even with a width of 128 neurons.

5.4 Klein-Gordon Equation

We conduct the same experiments on the non-linear onedimensional Klein-Gordon equation of the following form:

We adopt the solution provided in [Wang et al., 2020a] and derive the source term f(t, x) in Eq. 6 to be consistent with the solution below:

The model is then trained with a zero initial condition and a Dirichlet boundary condition corresponding to the solution u, using 256 collocation training points and 200 boundary data.

The resulting absolute errors are reported in Table 2. Similar to other equations, the Sine activation function performs consistently better for all widths, achieving the best mean absolute error of 3.1 × 10-3 with the 256 neurons-wide model.

(a) Exact Solution (b) Tanh 128-Neuron (c) Sine 256-Neuron

Figure 9: Absolute error of the Klein-Gordon PINNs.

Tanh also performs reasonably well with the 128 neuronswide model, although it still performs slightly worse than the Sine models, as shown in Figure 9.

6 Conclusion

The differentiation process in the residual loss of PINNs transforms the structure of the neural networks and their outputs, rendering the existing theory around loss functions and common supervised tasks ineffective in analyzing PINNs. In this work, we aim to fill the gap in our understanding of the residual loss and derive the requirements in network design that lead to achieving global minimization of this loss function. To this end, we study the residual loss at a critical point in the parameter space of the neural network and look for distinct characteristics of a global minimum that sets it apart from other critical points. We then use those characteristics to derive the requirements in the neural network design that ensure the existence of a global minimum. In particular, we show that under certain conditions, wide networks globally minimize the residual loss. Additionally, we reveal that activation functions with well-behaved high-order derivatives are crucial in the optimal minimization of the residual loss. We then use the established theory and empirical observations to choose activation functions and verify their effectiveness by conducting a set of experiments. The theory developed in this work paves the way for further development of better activation functions and provides a guideline for designing effective PINNs.

Acknowledgements

This research is supported by Natural Sciences and Engineering Research Council of Canada (NSERC), Discovery Grants program, and the Vector Scholarship in Artificial Intelligence, provided through the Vector Institute.