Probing optimisation in physics-informed neural networks
A novel comparison is presented of the effect of optimiser choice on the
accuracy of physics-informed neural networks (PINNs). To give insight into why
some optimisers are better, a new approach is proposed that tracks the training
trajectory curvature and can be evaluated on the fly at a low computational
cost. The linear advection equation is studied for several advective
velocities, and we show that the optimiser choice substantially impacts PINNs
model performance and accuracy. Furthermore, using the curvature measure, we
found a negative correlation between the convergence error and the curvature in
the optimiser local reference frame. It is concluded that, in this case, larger
local curvature values result in better solutions. Consequently, optimisation
of PINNs is made more difficult as minima are in highly curved regions.
Nayara Fonseca, Veronica Guidetti, Will Trojak
# Probing optimisation in physics-informed neural networks

###### Abstract

A novel comparison is presented of the effect of optimiser choice on the accuracy of physics-informed neural networks (PINNs). To give insight into why some optimisers are better, a new approach is proposed that tracks the training trajectory curvature and can be evaluated on the fly at a low computational cost. The linear advection equation is studied for several advective velocities, and we show that the optimiser choice substantially impacts PINNs model performance and accuracy. Furthermore, using the curvature measure, we found a negative correlation between the convergence error and the curvature in the optimiser local reference frame. It is concluded that, in this case, larger local curvature values result in better solutions. Consequently, optimisation of PINNs is made more difficult as minima are in highly curved regions.

## 1 Introduction

The idea of solving PDE problems using neural networks (NNs) was put forward by Lagaris et al. (1997; 1998); Lagaris et al. (2000) in the second half of the '90s and then revised in 2017 by Raissi et al. (2017a;b) who named the methodology Physics-Informed Neural Networks (PINNs). Relying on the universal approximation theorem of Cybenko (1989); Hornik (1991); Pinkus (1999), PINNs aim to deliver a universal regressor that can represent any bounded continuous function and solve any PDE/ODE problems, having input and output shape as the only limitation. Although the goal of PINNs was to produce a unifying method of solving PDE/ODEs, satisfactory results could not be achieved in a multitude of cases. The recent works of Karniadakis et al. (2021); Hao et al. (2022) provide an overview of the state-of-the-art; furthermore, Cuomo et al. (2022) focus on algorithms and applications and Beck et al. (2020) on theoretical results.

Several studies have analysed the effects of choosing different architectures, loss function formulations, and treatment of domain and collocation points. However, the effect of the optimiser choice on PINN performance needs more attention. Recently, new PINN-specific optimisation methods were developed or applied to improve poor convergence performance. For example, De Luca & Silverstein (2022) propose a _relativistic_ optimisation algorithm that introduces chaotic jumps and Davi & Braga-Neto (2022) use a metaheuristic optimisation method, namely, particle swarm optimisation, to eliminate gradient-related problems.

This work aimed to improve the understanding of how PINNs performance is affected by optimiser choice. Specifically, we considered the following algorithms spanning different optimiser categories: gradient descent (GD) without momentum, LBFGS (Dennis & More, 1973; Goodfellow et al., 2016) (a second order quasi-Newton method), ADAM (Kingma & Ba, 2014) (an adaptive stochastic GD algorithm), and bouncing Born-Infeld (BBI) (De Luca & Silverstein, 2022). The first three optimisers are widely used in ML optimisation problems, whereas BBI is a recent realisation of a frictionless energy-conserving optimiser. See Appendix A for details.
Moreover, we introduce a new method to study optimisation in neural networks, which provides training trajectory curvature data at low computational cost. Using this approach, we studied the evolution of different optimisers through the network parameter space during training. Further discussion is given in Section 2.1.

Specifically, we apply PINNs to solve the linear advection equation as described in Section 2.2. Linear advection is a simple PDE with a single parameter, i.e., the wave speed \(\beta\), whose variation can tune the complexity of the PINN landscape (Krishnapriyan et al., 2021). Moreover, to understand how a network adjusts to different depths and widths, we considered two configurations of multi-layer perceptron architectures with different numbers of hidden layers and nodes per layer.

The main contributions of this work are:

* We found that the choice of optimisation algorithm significantly impacts the convergence of PINNs models. Specifically, second-order optimisers that do not directly follow gradient directions, such as LBFGS, performed better, as shown in Fig. 1a.
* We introduce a new low-cost method for studying the dynamics of optimisers. This relies on defining a local reference frame and evaluating the local curvature of the training trajectory in the NN parameter space.
* We found a negative correlation between PINNs convergence error and the newly introduced training trajectory curvature which is shown in Fig. 1b. This implies that good PINNs solutions lie in highly curved regions of the optimiser reference frame. Therefore, making PINNs converge is hard for those techniques designed to explore shallow landscapes that generalise well under perturbation.

## 2 Background

### Tracking local curvature

Solely observing the convergence performance of an optimiser will show which optimiser is the best for a given task but will give limited information as to why. To move a step towards explaining why one optimiser is better than another, we analyze the trajectory of every optimiser in their local reference frame.

To begin, let \(\mathbf{\omega}\) be the neural network parameters, the weight update rule for update \(k\) is defined as \(\Delta\mathbf{\omega}_{k}\equiv\mathbf{\omega}_{k+1}-\mathbf{\omega}_{k}=\mathbf{V}(\mathbf{ \omega}_{k},\mathbf{\eta})\), where \(\mathbf{\eta}\) are the model hyperparameters. Assuming the continuum time limit, the trajectory in the parameter space during training can be described by:

\[\hat{\mathbf{\omega}}=\mathbf{V}(\mathbf{\omega},\mathbf{\eta}), \tag{1}\]

where the overdot corresponds to a time derivative, i.e. \(\hat{\mathbf{\omega}}=\mathrm{d}\mathbf{\omega}/\mathrm{d}t\). When using gradient descent, Eq. (1) can be interpreted as the equation of motion of a classical particle in a friction-dominated setup and \(\mathbf{V}(\mathbf{\omega},\lambda)=-\lambda\frac{1}{n}\sum_{i}\nabla_{\mathbf{\omega} }\ \mathcal{L}(\mathbf{x}^{(i)},\mathbf{\omega})\). Here, \(\lambda\) is the learning rate, \(n\) is the number of samples in the training set \(\{\mathbf{x}^{(1)},\mathbf{x}^{(2)},\cdots,\mathbf{x}^{(n)}\}\), and \(\mathcal{L}\) is the loss function. For other optimisation algorithms, the definition of \(\mathbf{V}\) changes and the ODE in Eq. (1) is no longer obtained via the Euler-Lagrange equations. Therefore, it does not have a dynamic interpretation. However, Eq. (1) can be efficiently used to understand the kinematics of the optimiser trajectories. This aspect is discussed further in Section 3.

To examine the motion in parameter space described by Eq. (1), we introduce the unit vector tangential to the training trajectory, defined as

\[\hat{\mathbf{T}}=\frac{\hat{\mathbf{\omega}}}{\|\hat{\mathbf{\omega}}\|_{2}},\quad \text{with}\quad\|\hat{\mathbf{\omega}}\|_{2}=\sqrt{\langle\hat{\mathbf{\omega}},\hat {\mathbf{\omega}}\rangle}. \tag{2}\]

Here, \(\hat{\mathbf{T}}\) is a time-dependent vector in a \(N_{\mathbf{w}}\)-dimensional space, for a \(N_{\mathbf{w}}\) parameter NN. To track the local training trajectory curvature in this time-dependent reference frame, we can define the local (time-dependent) curvature, \(\kappa_{t}\). This is the rate at which the tangent unit-vector changes with respect to time, and is given by:

\[\kappa_{t}=\left\|\frac{\mathrm{d}\hat{\mathbf{T}}}{\mathrm{d}t}\right\|_{2}= \frac{1}{\|\hat{\mathbf{\omega}}\|_{2}}\left[\mathbf{\tilde{\omega}}\cdot\mathbf{\tilde{ \omega}}-\left(\frac{\mathrm{d}\|\hat{\mathbf{\omega}}\|_{2}}{\mathrm{d}t}\right)^{ 2}\right]^{1/2}. \tag{3}\]
Another parametrisation with a clearer geometric meaning is \(\kappa_{\omega}=\left\|\frac{\mathrm{d}\mathbf{\hat{\Sigma}}}{\mathrm{d}\omega} \right\|_{2}=\kappa_{t}/\|\boldsymbol{\dot{\omega}}\|_{2}\). This represents a local geometric quantity, i.e., the local curvature in the parameter space, removing the effects of time and trajectory speed1. To calculate the curvature at each training step a first-order approximation is used, details of which are given in Appendix C.

Footnote 1: This is the local curvature of the training trajectory, not the curvature of the space of weights, which is flat.

In comparison with a Hessian-based curvature calculation, the method presented here has a significantly lower computational cost as it does not require the Hessian of the loss function. Furthermore, in most cases, an optimiser does not solely follow the gradient loss and it is generally non-trivial to find the function relating the loss, its derivatives, and the trajectory followed by the optimiser. Therefore, although every optimiser is constructed with the aim of sharing its minima with the loss function, analyzing the loss Hessian may not characterise the optimiser trajectories and may not explain why a certain minimum is found. Moreover, different optimisers are based on different assumptions--such as energy conservation, friction-dominated motion, or others--further modifying the loss seen by the optimiser. Fig. 2 in Appendix C gives a pictorial representation of the loss function (or its distorted version seen by the optimiser) and its relationship to the trajectory in the parameter space.

### Linear Advection

The PDE used for the tests in this work was the one-dimensional linear advection equation on the periodic spatial domain \(\Omega=[0,2\pi)\), given by

\[\frac{\partial u}{\partial t}+\beta\frac{\partial u}{\partial x} =0,\quad\text{for}\quad\mathrm{u}:\mathrm{T}\times\Omega\mapsto \mathbb{R},\quad\Omega=[0,2\pi),\mathrm{T}\in[0,1], \tag{4a}\] \[u(x,0) =u_{0}(x),\] (4b) \[u(0,t) =u(2\pi,t), \tag{4c}\]

where \(\beta\) is the wave speed and \(u_{0}(x)\) is the initial condition. The exact solution of this system can be straightforwardly found using the superposition of Bloch waves. To find approximate solutions to this system we applied PINNs with loss functions similar to those used by Krishnapriyan et al. (2021), defined as

\[\mathcal{L}(\theta)=\frac{1}{N_{u}}\sum_{i=1}^{N_{u}}\left(\hat{u}-u_{0}^{i} \right)^{2}+\frac{1}{N_{f}}\sum_{i=1}^{N_{f}}\lambda_{i}\left(\frac{\partial \hat{u}}{\partial t}+\beta\frac{\partial\hat{u}}{\partial x}\right)^{2}+\frac {1}{N_{b}}\sum_{i=1}^{N_{b}}\left(\hat{u}(\theta,0,t)-\hat{u}(\theta,2\pi,t) \right)^{2}.\] (5a) Here \[\hat{u}=\hat{u}(\theta,x,t)\] is the neural network output, parameterised by \[\theta\]. The initial condition used through out this work was \[u(x,0)=\sin(x)\] and \[u(0,t)=u(2\pi,t)\]. For all the experiments we fix \[\lambda_{i}=1\].

## 3 Results and Discussion

Fig. 0(a) shows the median over 10 samples of the mean squared error (MSE) evaluated on the entire domain between the \(\hat{u}\) predicted by the PINN and the analytical solution. Here the samples were formed by using 10 different random network initialisation. The corresponding training and test losses are shown in Appendix D.1. During testing, it was observed that the learning rate significantly impacts the final performances; therefore, a grid search was performed over a range of learning rates to estimate the best configuration for each optimiser. See Appendix B for the details.

Two multi-layer perception architectures were considered, with layer structures \(S=[2,25,25,1]\) and \(L=[2,50,50,50,50,1]\) resulting in \(751\) and \(7851\) trainable parameters respectively2. The number of training points was around 2000, as described in Appendix B; therefore, the larger network is in the over-parameterised regime, i.e. it has more parameters than training data. Fig. 0(a) shows that the large network
generally has a lower error, in agreement with common knowledge in ML3. This effect is most noticeable for BBI 4. A significant observation is that for all configurations of NN and optimiser, PINNs failed to produce good approximations for systems with large values of \(\beta\).

Footnote 3: From classical statistical learning, one expects that over-parameterised models over-fit. However, there is overwhelming evidence that large models generalise well in several cases of interest, e.g. Szegedy et al. (2015); Huang et al. (2019).

Footnote 4: The total phase space volume for BBI can be analytically estimated and is parametrically large as the objective function is close to zero with an exponential dependence on the number of dimensions (see Appendix A.3 in De Luca and Silverstein (2022)).

In Fig. (b)b, we show the inverse relation between the final values of MSE and the curvature \(\kappa_{\omega}\). Here, only \(\beta=1,5\) were considered as they converged to reasonable errors for most of the optimisers, see Appendix D.2. Significantly, LBFGS achieves the highest values of \(\kappa_{\omega}\) and the lowest error. Moreover, LBFGS has significant memory overhead but it typically converges in fewer epochs compared to the other optimisers. Note that such a negative correlation in Fig. (b)b links generalisation (measured by the MSE on the entire domain) with the curvature \(\kappa_{\omega}\). Deriving their exact relation is non-trivial and we leave further investigations for future work. Nevertheless, we would like to stress that the high curvature values is not solely associated with orbiting the final local or global minima. In fact, this is due to \(\kappa_{\omega}\) and MSE being negatively correlated throughout the whole trajectory, including its initial stages. We show examples of this behavior in Appendix D.2.

Outlook. In contrast to traditional ML tasks such as image recognition and NLP, ML methods applied to science require more accurate models which, in general, are trained with high-quality data. This suggests that the training process on data from physical phenomena can be fundamentally different from common ML applications. A promising future research direction would be exploring the connection between model generalisation and the flatness of minima for problems requiring high accuracy. For example, the work of Dinh et al. (2017); Huang et al. (2020) discuss this in the context of traditional ML tasks. Finally, in the future, it would be insightful to compare Hessian-based methods such as that of Michaud et al. (2023), with our approach linking accuracy and local trajectory curvature.

Figure 1: (a) Median MSE over 10 samples for various optimisers and \(\beta\) values. _Solid/dashed_ lines refer to small/large (S/L) network architectures. (b) Relation between final values of convergence (MSE) and curvature (\(\kappa_{\omega}\)) for \(\beta=1\) (left) and \(5\) (right). Dots and crosses refer to small (S) and large (L) networks, respectively.
#### Acknowledgments

We thank Eloisa Bentivegna and Imran Nasim for discussions. NF and WT acknowledge the UKRI support through the grant MR/T041862/1. The authors acknowledge the IBM Research Cognitive Computing Cluster service for providing resources that have contributed to the research results reported within this work.
