PSO-PINN: Physics-Informed Neural Networks Trained with Particle Swarm Optimization

The feed-forward fully-connected neural network is the basic architecture used in deep learning algorithms[28]. A fully-connected neural network with $L$ layers is a function $f_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{k}$ described by

where $\sigma$ is an entry-wise activation function, $W^{[l]}$ and $b^{[l]}$ are respectively the weight matrices and the bias corresponding to each layer $l$, and $\theta$ is the set of weights and biases:

Among popular choices for the activation function are the sigmoid function, the hyperbolic tangent function (tanh), and the rectified linear unit (ReLU) [29]. To fit the neural network $f_{\theta}(\mathbf{x})$ to data $\{(\mathbf{x}_{i},\mathbf{y}_{i})\}_{i=1}^{n}$, one minimizes a suitable loss function. A popular choice in traditional deep learning is the mean square error (MSE):

Consider a non-linear differential equation of the general form:

where $\Omega\subset\mathbb{R}^{d}$, $u:\Omega\rightarrow\mathbb{R}^{k}$, $\mathcal{N}[\cdot]$ is a differential operator, and $g,h:\Omega\rightarrow\mathbb{R}^{k}$ specify forcing and boundary conditions, respectively. We employ a neural network $f_{\theta}(\mathbf{x})$ to approximate the unknown function $u(\mathbf{x})$. The PDE is enforced through the loss function

where $\{\mathbf{x}^{r}_{i}\}_{i=1}^{n_{r}}\subset\Omega$ are collocation points, and $\mathcal{N}[f_{\theta}(\mathbf{x})]$ is computed accurately using automatic differentiation methods [30]. The neural network is fitted to the initial and boundary condition $h(\mathbf{x})$ using a traditional data-fitting loss function:

where $\{\mathbf{x}^{b}_{i}\}^{n_{b}}_{i=1}\subset\partial\Omega$ are initial and boundary points. If there are experimental or simulation data $\{(\mathbf{x}^{d}_{i},\mathbf{y}^{d}_{i})\}^{n_{d}}_{i=1}$ available, they may be included in the usual way through the loss:

The previous multi-objective optimization problem is typically solved by simple linear scalarization, in which case the PINN is trained by minimizing the total loss function $\mathcal{L}$:

This framework is easily extended to more complex initial and boundary conditions involving derivatives of $u(\mathbf{x})$.

The Particle Swarm Optimization (PSO) algorithm [31, 32] is a population-based stochastic optimization algorithm that emulates the swarm behavior of particles distributed in a $n$-dimensional search space [33]. Each individual in this swarm represents a candidate solution. At each iteration, the particles in the swarm exchange information and use it to update their positions. Particle $\theta^{t}$ at iteration $t$ is guided by a velocity determined by three factors: its own velocity inertia $\beta v^{t}$, its best-known position $p_{best}$ in the search-space, as well as the entire swarm’s best-known position $g_{best}$:

where $c_{1}$ and $c_{2}$ are the cognitive and social coefficients, respectively, referred to jointly as the behavioral coefficients, and $r_{1}$ and $r_{2}$ are uniformly distributed random numbers in range $[0,1)$. Then the particle position is updated as

Many variations of the PSO algorithm were proposed over time. PSO can be combined with gradient descent to train neural networks, with the gradients computed efficiently by automatic differentiation (backpropagation), as was done in [27]. In this algorithm, which was called PSO-BP by the authors, the particle velocity has an additional component, namely, the gradient vector of the loss function$\nabla\mathcal{L}(\theta)$:

where $\alpha$ is a learning rate. Therefore, the gradient participates in the velocity magnitude and direction, by an amount that is specified by the learning rate, which helps guide the swarm to a good local optimum. See Figure 1 for an illustration of the PSO-BP update.

We consider the following three variants of the PSO algorithm:

1. 1.

   PSO - The algorithm in its original form, as described in (9) and (10).

2. 2.

   PSO-BP - This version adds a component based on the gradient of the loss function, as described in (11).

3. 3.

   PSO-BP-CD In this variant, we introduce a behavioral coefficient schedule. The behavioral components $c_{1}$ and $c_{2}$ in PSO-BP-CD decay linearly according to the iteration number as follows:

   $$c^{t+1}_{k}=c^{t}_{k}-\frac{2c^{t}_{k}}{n}$$

   (12)

   for $k={1,2}$ and $t={1,2,\ldots,n}$, where $n$ is the total number of iterations.

We propose PSO-BP-CD as a variant of PSO-BP, which puts more weight on the gradient vector near the end of training, when the swarm should already have zeroed in on a good local optimum, and less communication between the particles is needed. This allows PSO-BP-CD to be more effective than the PSO-BP and pure PSO methods in training PINN ensembles, as will be seen in the experimental results section.

The PSO-PINN swarm is an ensemble of PINN candidates. Each particle in the swarm is a vector $\theta$ containing the weights and bias of the corresponding PINN. The population $\Theta$ of all particles is randomly initialized, as will be detailed later in Section 4.2. The objective function $f_{\theta}(\cdot)$ is the total loss function in (8). The $p_{best}$ vector should be initialized to the values in the original population $\Theta$, as the particles know nothing but their initial point at this time. The $g_{best}$ parameter is equal to the value of $\theta\in\Theta$ that achieves the minimum value of $f(\theta)$. After the initialization step, the updates to the swarm attempt to minimize the expected value of the objective function. Each particle velocity is updated following (11). Despite the use of the gradient update $\alpha\nabla\mathcal{L}(\theta)$ in the previous equation, in practice, the update is made using ADAM optimization [6]. After the velocity update, the particle locations (i.e., the network weights) are updated according to (10). The values of $p_{best}$ and $g_{best}$ are updated accordingly. The training loop runs until a maximum number of iterations is reached. The procedure is summarized in Algorithm 1.

Besides the fact that swarms are a well-established optimization method, they also provide an ensemble of solutions. Thus, we can take advantage of the various properties provided by ensembles. This category of learning methods combines several individual models to create a “collective” solution, which is expected to display improved performance and stability with respect to any of the individual models. There are different methods to seek the consensus decision in an ensemble [25]. The most straightforward one perhaps would be simply choosing the best-fitted model according to the loss function. Although this may yield a good model, it could fail due to overfitting. To avoid this, a better approach exploits the ensemble diversity through model averaging:

Uncertainty of the prediction at each point $\mathbf{x}$ of the domain can be quantified naturally through the sample variance:

Fist, we use the 1D Poisson equation to illustrate the PSO-PINN algorithm. Due to its simplicity, this example allows us to visualize the evolution of training and the results. The 1D Poisson equation considered here is:

where $g(x)=-(2\pi)^{2}\sin(2\pi x)$ is manufactured so that the solution is $u(x)=\sin(2\pi x)$. For this simple problem, we use a relatively shallow architecture consisting of 3 layers of 10 neurons. The PSO-PINN ensemble was composed by 50 particles and used the proposed PSO-BP-CD optimization method, with hyperparameter values $\alpha=0.005$, $\beta=0.9$, $c_{1}=0.08$, and $c_{2}=0.5$.

We can see in Figure 2 the progression of training, and how the PSO-PINN swarm zeros in on the solution, rather quickly, after only 90 iterations. We can observe that, initially, the mean is off and the variance is large, which is the desired behavior, as the variance is supposed to quantify the uncertainty associated with the approximation. As training progresses, the swarm approaches a consensus, the mean converges to the solution, and the variance simultaneously shrinks to zero, indicating more confidence in the final result. The PSO-PINN swarm is thus accurate, with well-calibrated uncertainty quantification.

In this section, we use classical PDE benchmarks to study the performance of PSO-PINN using the three PSO variants described in 3.1. The number of training iterations was fixed at 2,000, while the architecture of the neural networks was fixed at 5 layers and 8 neurons per layer. The ensemble size was kept fixed at 100 neural networks. We kept the same number of initial condition points, boundary condition points and residual points for all experiments in this section. In all experiments, we used 1024 evenly spaced points for the initial condition, 512 evenly spaced points for the boundary conditions, and 1000 residual points randomly distributed over the solution domain using latin hypercube sampling. For the basic PSO algorithm, the hyperparameters were set to $\beta=0.9$, $c_{1}=0.8$ and $c_{2}=0.5$ in all experiments.

In this section we show that PSO-PINN ensembles produce better results than traditional ensembles of PINNs trained individually with ADAM gradient descent. For this, we selected two particularly hard examples, in order to appreciate the improvement over traditional ensembles afforded by PSO-PINN over traditional ensembles. In both benchmarks, the PSO-BP-CD variant is used with the PSO-PINN, with $\alpha=0.005$, $\beta=0.99$, $c_{1}=0.08$ and $c_{2}=0.5$, while the ADAM parameters were set to $\alpha=0.001$, $\beta_{1}=0.99$, and $\beta_{2}=0.999$.

The experimental results displayed promising results for PINN training using swarm intelligence and ensemble models. It was seen in Section 4.2 that both PSO-BP and PSO-BP-CD show superior results to using a basic PSO algorithm. This is no doubt due to the extra information contained in the gradient of the function. We also observed that PSO-BP-CD, the variant we proposed in this paper, performed better than PSO-BP, especially in the case of the nonlinear Burgers benchmark.

PSO-PINN provides the advantages of an ensemble approach, allowing the use of the mean of the ensemble as the predictor and its variance as a measure of the corresponding prediction uncertainty. However, PSO-PINN has clear advantages over a standard ensemble of PINNs trained with ADAM, as seen in the results of Section 4.3. The PSO-PINN ensembles are more precise and converge faster, with well-calibrated variance, leading to a greater confidence in the solution. This no doubt results from the fact that there is communication between the PINNs in a PSO-PINN swarm, while the traditional ensemble the PINNs do not communicate.

Setting the hyperparameters of the PSO-PINN ensemble, i.e., $\alpha$, $\beta$, $c_{1}$, and $c_{2}$, is a model selection problem comparable to setting the correct learning rates and initialization in training a deep neural neural network. Indeed, the hyper parameter $\alpha$ is directly comparable to the learning rate used in traditional gradient descent. We observed in our experimental results that suitable hyperparameter values were fairly consistent across experiments. Also, there was no need to set a schedule for the learning rate hyperparameter $\alpha$, as is usually needed in traditional gradient descent algorithms