Numerical Approximations of the Allen-Cahn-Ohta-Kawasaki (ACOK) Equation with Modified Physics Informed Neural Networks (PINNs)

Ohta-Kawasaki(OK) free energy[22], usually associated with a volume constraint, has been used to simulate the phase separation of diblock copolymers, i.e., polymers consisting of two types of monomers, $A$ and $B$, that are chemically incompatible and connected by a covalent chemical bond. It is formulated as

Here $\mathbb{T}^{d}=\prod_{i}[-X_{i},X_{i})\subset\mathbb{R}^{d},d=1,2,3$ is a periodic domain, and $u=u(x)$ is a phase field labeling function representing the density of $A$ species with interfacial width $\epsilon$. Function $W(u)=18(u^{2}-u)^{2}$ is a double well potential. The nonlinear function $f(u)=6u^{5}-15u^{4}+10u^{3}$ keeps the $0\text{-}1$ structure of $u$ as $f(0)=0$ and $f(1)=1$. More importantly it also has the property that $f^{\prime}(0)=f^{\prime}(1)=0$, which localize the force and avoid possible non-zeros and non-ones near the boundary of the interface [29] when considering $L^{2}$ gradient flow dynamics of (1.1). The first integral is a local surface energy, the second term represents the long-range interaction between the molecules with $\gamma$ controlling its strength and $\omega\in(0,1)$ being the fraction of species $A$, and the third term is a penalty term to fulfill the volume constraint

Allen-Cahn-Ohta-Kawasaki(ACOK) model [29], which is the focus of this paper, takes the $L^{2}$ gradient dynamics of the OK free energy in (1.1) as

Various successful and efficient PDE solvers can be applied to the model, such as Finite Difference, Finite Element [20], Spectral method [10, 15], etc. In recent years, some novel numerical methods have also been studied for Cahn-Hilliard dynamics, i.e., the $H^{-1}$ gradient flow dynamics of OK energy, such as the midpoint spectral method [1], and the Invariant Energy Quadratization (IEQ) method [5]. In [29], a first-order energy stable linear semi-implicit method is introduced and analyzed for the ACOK equation. Our goal in this paper is to study the approximation of the solution of the ACOK equation with neural networks.

Recently, there has been an increasing interest in solving PDEs with machine learning methods, among which the physics-informed neural networks(PINNs) [24] have gained tremendous popularity. PINNs are based on a fully-connected feed-forward deep neural network (DNN)[14, 26], which is often interpreted utilizing universal approximation theorem [6, 18, 12, 11]. The structure of a fully-connected DNN, consisting of an input layer, hidden layers, and an output layer, is shown in Figure 1. Every node in one layer is connected with every node in another by a series of computations. The nodes in the input layer are the specific data studied, and the nodes in the output are the expected outcome. For example, in the case of facial recognition, the input could be the RBG values (features) of each pixel of the sample pictures, while the output could be the assigned numerical values representing the individuals (classes). The nodes in the hidden layers are called neurons, and they are responsible for the performance of the neural network. There can be many hidden layers in a neural network; in the case of Figure 1, the number of hidden layers is three.

A column vector is fed to each node of the input layer. Then in the first hidden layer, we multiply it by another row vector, i.e., weight, to get the product. This is the case that there is only one neuron in the relevant hidden layer. If there are more neurons, the row vector will be a weight matrix instead, where the number of rows of the matrix represents the number of neurons. A bias vector is then added to the product, and an activation function is applied to the new vector, which contributes to the output of the current hidden layer, which is also the input of the next layer. In a nutshell, let $M$ be the number of hidden layers, $p$ be the input of the neural network, then the output $a$ of the neural network will be $a^{M}$, with

where $a^{i-1}$ is the input, $W^{i}$ represents the weights, $b^{i}$ represents the bias, and $f^{i}$ is the activation function in the $i\text{th}$ layer.

To obtain the proper weights and bias, we perform an optimization process, for example, gradient descent, on the weight and bias to update them. To calculate the derivatives involved in the optimization process, the network goes through a process called back-propagation [25], which is essentially a chain rule applied on each layer starting from the final one. We start with a certain loss function:

where $y$ is the object, the network is approximating, and $\Theta$ denotes the weights and bias of the DNN. The loss function, depending on the specific problem, could be a mean square error (MSE), mean absolute error (MAE), log loss function, etc. We then perform back-propagation on the loss function, trace all the way back to get the derivatives of weights and bias for each layer, and then use the optimization tool to update the weights and bias. A gradient descent update is as follows:

where $\eta$ is the learning rate, and $\nabla_{\Theta}\text{Loss}$ is the direction of update. It works essentially like a directing system, telling you which direction to go and how far to go in that direction based on your current location and destination. The optimization process will be repeated several times until the prescribed accuracy is achieved. We call this number of repetitions the number of epochs.

Furthermore, there are a variety of optimizers that can be selected. The previously mentioned gradient descent is one of them. While it has many advantages, it does not work well in our non-convex problem due to its notoriously poor performance on escaping local minimums [7]. Therefore, other optimizers, such as momentum [23], ADAM [16], Adadelta [30], RMSprop [13], etc are developed to better address this issue. The momentum method adds an accelerator in the relevant direction. For example, think about rolling a ball down a ramp. As the ball goes down, it gains momentum and travels faster and faster. Other methods, like Adadelta, use different forms of adaptive learning rates. It is worth mentioning that RMSprop works by dividing “the learning rate for a weight by a running average of the magnitudes of recent gradients for that weight” [13].

In our case, we adopt ADAM as our optimizer for each epoch and L-BFGS-B with a certain stopping condition as the optimizer afterward. Introduced in [16], ADAM can be viewed as a combination of the momentum method and RMSprop. It uses the mean $m_{t}$ and uncentered variance $v_{t}$ of the gradient with regards to time $t$ to adapt the learning rate for each weight and bias of the neural network:

where $g_{t}$ is the gradient of the loss function at time step $t$, and the hyper-parameters $\beta_{1},\beta_{2}\in[0,1)$ direct the decay rates of $m_{t}$ and $v_{t}$. A bias correction is also applied to $m_{t}$ and $v_{t}$, which leads to

and to update the weights and bias, we perform $\Theta_{t}=\Theta_{t-1}-\eta\frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}}+\varepsilon}$, where $\eta$ is the learning rate, and $\varepsilon$ is arbitrary. On the other hand, the L-BFGS-B optimizer [3] uses a “limited memory” Broyden–Fletcher– Goldfarb–Shanno (BFGS) [2, 8, 9, 27] matrix to represent the Hessian matrix of the objective function, making it great for optimization problems with a large number of variables. Stems from the L-BFGS method [19, 21], L-BFGS-B places upper and lower bound constraints on variables. It identifies fixed and free variables for each step with a basic gradient method. It then performs the L-BFGS method on the free variables to achieve better accuracy.

For PINNS, sampled collocation points are fed to the DNN, and the function value is returned to the output layer. Weight and bias in each layer can be learned or updated by optimizing the error function to reach the prescribed accuracy, in which the back-propagation technique is needed to compute the first-order derivatives. Different from traditional machine learning DNN, which requires a very large amount of ground truth data to make an accurate prediction, PINNs take advantage of the physics information and make up for the lack of sufficient ground truth data. It divides the error function into two components, one corresponding to the known ground truth data and the other corresponding to the physics information. Consequently, our goal is to minimize the error between the limited amount of ground truth and the prediction, as well as how much the prediction ”fits” the physics information. Furthermore, as shown in [4], PINNs can be employed to solve inverse problems, in particular, inverse scatter problems in photonic metamaterials and nano-optics technologies, by transforming them into parameter retrieving problems. However, PINNs also have their downsides. As suggested in [17], “existing PINN methodologies can learn good models for relatively trivial problems”, and it can fail when predicting solutions for a more complex model, which is likely due to the fact that the setup of PINNs makes it hard to optimize the less-smooth loss landscape. They have, therefore, proposed a “curriculum regularization”, which is finding a good initialization for weights and training the model gradually. Notably, this has made the optimization process easier. They also mentioned a sequence-to-sequence learning method, which is similar to the Time-Adaptive Method proposed in [28]. This method can make learning easier by learning in small intervals.

While many successful and efficient traditional numerical methods exist, our motivation to study the ACOK equation using machine learning methods is that they can learn the unknown operators in a family of generalized equations from a set of data samples, i.e., the inverse problem. This would benefit future studies focusing on nonlocal PDEs with general nonlocal operators. It can also potentially help tackle real-world challenges in fields such as physics, medicine, fluid dynamics, aerospace, etc. Our contribution to this work is to tackle the challenges that the complexity of the ACOK model has imposed on PINNs. As shown in 1.1, the long-range interaction term, the volume constraint, and the zero-mean condition make it necessary to apply computational and structural modifications to the basic PINNs. Our solution for the long-range interaction term is to apply the Laplace operator on a second neural network output and ensure its accuracy. We introduce a separate neural network into the model for the volume constraint and the zero-mean condition, which takes a separate set of uniform sampling points as its input.

The rest of this paper is organized as follows. In Section 2, we present the numerical methods. Specifically, we briefly revisit the baseline PINNs in Section 2.1. Then we elaborate on the modified PINNs, including the application of the periodic boundary condition, the approximation of the long-range interaction term, the volume constraint, and the zero-mean condition, in Section 2.2. We will then present and discuss some 1-dimensional results in Section 3. In the end, we give a detailed discussion and explain our future work in Section 4.

Our goal in this paper is to take the ground truth $u(0,x)$ data as input, and predict the solution $u(t,x)$ of ACOK equation with neural networks.

Now we have all the pieces regarding how to apply PINNs on ACOK. We can combine them and implement them in Python. The code is based on the work of Maziar Raissi et al. [24]. We now show the results of the modified PINNs on 1-dimensional ACOK with different time intervals, ranging from $1\times 10^{-3}$s to $1\times 10^{-2}$s. Fine-tuning on hyper-parameter is needed to reach the best results, but it is out of our current research scope.

Figure 6 are run for $495$ epochs, with weights $W_{u_{in}}=1\times 10^{5}$, $W_{((-\Delta)^{-1}u)_{in}}=5\times 10^{6}$, $W_{u_{b}}=1$, $W_{((-\Delta)^{-1}u)_{b}}=30$, $W_{F}=1$, $W_{Lap}=500$, $W_{int}=1$, $W_{zm}=30$. There are $10$ hidden layers in $\text{Net}_{u}$, each containing 20 neurons, while only $3$ hidden layers with $10$ neurons for each layer were needed for $\text{Net}_{v}$, since it is a much simpler neural network.

We can see that the modified PINNs perform remarkably well in smaller time periods. However, as the time interval increases, the prediction of $(-\Delta)^{-1}u$ at the right tail is getting notably worse. There are several ways to address this problem. One approach is the Time-Adaptive Method [28], or similarly, the sequence-to-sequence learning [17], which learns only one small-time period at a time and then uses the prediction as to the initial data for the next period. Since our model approximates well with short time intervals, we can compile all the sub-results together to get the entire solution. For example, we could use the data of the last time column in Figure 6(a) as the initial data for the next 100 columns and repeat the process in 100 column increments to achieve good results.

Furthermore, the problem can be improved by rescaling the sampling points. As shown in Figure 7(a), changes in phase separation occur most rapidly during the earlier time period. Therefore, we propose rescaling the sampling points to be denser in earlier time columns. Recall Figure 4, we continue to use the LHS sampling, but rescale the points polynomially to achieve the denser in the beginning effect, as in Figure 7(b).

The results are shown as follows, Figure 8(a) is the result with original LHS sampling, with time interval being $3\times 10^{-3}$s. Notice that the right tail on $(-\Delta)^{-1}u$ fits quite badly in this case due to the long time interval. Figure 8(b) is the result with a $x^{2}$ rescaling on the points. Although it does not fix the problem entirely, we can see that the tail on $(-\Delta)^{-1}u$ fits much better with a $x^{2}$ rescaling. In some cases, $x^{2}$ rescaling does not improve the result much. Therefore, $x^{3}$ rescaling could potentially be used to improve the fitting further.

During implementation, we have discovered that our modified PINNs are challenging to train due to their many hyper-parameters. Some hyper-parameters, such as learning rate, number of hidden layers, and number of neurons of each layer, are fixed during our training. However, the weights, number of epochs, and number of sampling points need fine-tuning to achieve the best results.

The most challenging part of our fine-tuning is finding a well-balanced set of weights. Each weight controls a different part of the estimation. For instance, $W_{u_{in}}$, $W_{((-\Delta)^{-1}u)_{in}}$ controls the fitting of initial data, which is undoubtedly an essential part of our loss function since the initial data is the only ground truth data we have. But they differ from each other, as we later find out that the scale of $((-\Delta)^{-1}u)_{in}$ is 100 times smaller than that of $u_{in}$. Therefore, we also consider that by increasing $W_{((-\Delta)^{-1}u)_{in}}$ to counter the scale issue. A similar scaling technique is applied to $W_{u_{b}}$ and $W_{((-\Delta)^{-1}u)_{b}}$, which controls the periodic boundary conditions.

$W_{F}$ dominates the physics information in our model, while $W_{Lap}$ controls the accuracy of the approximation of the long-range interaction term. We have learned that the accuracy of the long-range interaction term is one of the key players in our approximation. $W_{int}$ directs the volume constraint and, therefore, can be increased if the predicted result has an incorrect volume. $W_{zm}$ applies only to the $((-\Delta)^{-1}u)_{in}$, which means the scaling problem also needs to be taken into consideration. We have also discovered that the interval of a number of epochs is quite important. The model will not train well with too few epochs. However, if we boost the number of epochs to the scale of the thousands, the model is over-trained and does not perform well.

Moreover, we tune the number of the sampling points using the principle that we want to keep the randomness while at the same time preserving as much of the known information as possible. We also remain attentive to the efficiency of the program. Again, the initial data, being the only piece of ground truth information available, need to have the most sampling points proportionally. We do not take all the initial data because we want to keep the randomness. We also sample a relatively large percentage of the boundary points for the same reason. As the total number of the grid points is quite large for the internal points, we cannot take as many as the sampling points while keeping the program’s efficiency. Randomness also plays a crucial part here. Therefore, we only select less than half of the points. The number of randomly selected $t$ columns in the uniform mesh, $N_{t_{uni}}$, follows the same principle.

To increase the randomness, we also implemented a mini-batch method. The idea is that instead of considering all the points simultaneously, we consider a small subset of the points at a time. In theory, it should need fewer epochs and converge faster. However, as the computation in our case is quite complex, we have discovered that the mini-batch method takes significantly longer time than the regular approach.

We are currently working on the 2-dimensional case of this problem. As the number of points has drastically increased for the same time period, we have run into technical issues, in particular, insufficient memory capacity. To address this issue, we have to dramatically decrease the sampling points, even on the initial data. Also, we are keeping the time interval relatively small for the same reason. Additionally, for the boundary condition, we now have the upper and lower bound and the front and back bound. A more complicated back-propagation further lengthens the training time due to the second-order derivative. These are the difficulties we have encountered in the 2-dimensional case.

Possible future work includes investigating the inverse problem for nonlocal phase-field models with unknown nonlocal operators. We could train our model to learn the dynamics using the existing data and obtain a set of trained weights and bias representing the unknown operator to compare with a number of known possible operators. Furthermore, we could also explore other methods to improve the efficiency of the training process. The current one-dimensional results run within an hour. However, as the structure of the DNN becomes more complicated, training time increases drastically, likely caused by the more complex computations and the increased number of epochs. One possible solution is to start the currently randomized initial weights and bias closer to the solution, which will significantly lower the number of epochs, and by extension, decrease the training time overall. Additionally, various other machine learning methods, such as FNO, conventional neural networks, and transformer-based methods, can be applied to PDEs. These are also some exciting directions to explore.