Number Theoretic Accelerated Learning of Physics-Informed Neural Networks

For simplicity, we consider PDEs defined on an $s$-dimensional unit cube $[0,1]^{s}$. PINNs employ a PDE that describes the target physical phenomena as loss functions. Specifically, first, an appropriate set of collocation points $L^{*}=\{\bm{x}_{j}\mid j=0,\ldots,N-1\}$ is determined, and then the sum of the residuals of the PDE at these points

$$\textstyle\frac{1}{N}\sum_{j=0}^{N-1}\mathcal{P}[\tilde{u}](\bm{x}_{j})=\frac{1}{N}\sum_{{\bm{x}}_{j}\in L^{*}}\mathcal{P}[\tilde{u}](\bm{x}_{j}),$$

(1)

is minimized as a loss function, where $\mathcal{P}$ is a differential operator. The physics-informed loss satisfies $\mathcal{P}[\tilde{u}](\bm{x})=\|\mathcal{N}[\tilde{u}](\bm{x})\|^{2}$. Then, the neural network’s output $\tilde{u}$ becomes an approximate solution of the PDE. However, for $\tilde{u}$ to be the exact solution, the loss function should be 0 for all $\bm{x}\in\Omega$, not just at collocation points. Therefore, the following integral must be minimized as the loss function:

$$\textstyle\int_{[0,1]^{s}}\mathcal{P}[\tilde{u}](\bm{x})\mathrm{d}\bm{x}.$$

(2)

In other words, the practical minimization of (1) essentially minimizes the approximation of (2) with the expectation that (2) will be small enough, and hence $\tilde{u}$ becomes an accurate approximation to the exact solution.

More precisely, we show the following theorem, which is an improvement of an existing error analysis (Mishra2023estimates) in the sense that the approximation error bound of neural networks is considered.

For a proof, see Appendix “Theoretical Background.” $\varepsilon_{1}$ is determined by the architecture of the network and the function space to which the solution belongs. For example, approximation rates of neural networks in Sobolev spaces are given in GUHRING2021107. This explains why increasing the number of collocation points beyond a certain point does not further reduce the error. $\varepsilon_{2}$ depends on the accuracy of the approximation of the integral. In this paper, we investigate a training method that easily gives small $\varepsilon_{2}$.

One standard strategy often used in practice is to set $\bm{x}_{j}$’s to be uniformly distributed random numbers, which can be interpreted as the Monte Carlo approximation of the integral (2). As is widely known, the Monte Carlo method can approximate the integral within an error of $O(1/N^{\frac{1}{2}})$ independently from the number $s$ of dimensions (Sloan1994-cl). However, most PDEs for physical simulations are two to four-dimensional, incorporating a three-dimensional space and a one-dimensional time. Hence, in this paper, we propose a sampling strategy specialized for low-dimensional cases, inspired by number-theoretic numerical analysis.

Note that some variants, such as CPINN (Zeng2023), have proposed alternative objective functions. We hereafter denote any variant of physics-informed loss by (1), without loss of generality, as long as it can be regarded as a finite approximation to an integral over a domain, (2).

In this section, we propose the good lattice training (GLT), in which a number theoretic numerical analysis is used to accelerate the training of PINNs (Niederreiter1992-qb; Sloan1994-cl; Zaremba1972-gj). In the following, we use some tools from this theory.

While our target is a PDE on the unit cube $[0,1]^{s}$, we now treat the loss function $\mathcal{P}[\tilde{u}]$ as periodic on $\mathbb{R}^{s}$ with a period of 1. Then, we define a lattice.

Given a lattice $L$, the set of collocation points is defined as $L^{*}=\{\bm{x}_{j}\mid j=0,\dots,N-1\}\coloneqq\{\mbox{the decimal part of }\bm{x}\mid\bm{x}\in L\}\in[0,1]^{s}$. Considering that the loss function to be minimized is (2), it is desirable to determine the lattice $L$ (and hence the set of collocation points $\bm{x}_{j}$’s) so that the difference $|(\ref{eq:closs})-(\ref{eq:dloss})|$ of the two functions is minimized.

Suppose that $\varepsilon(\bm{x})\coloneqq\mathcal{P}[\tilde{u}](\bm{x})$ is smooth enough, admitting the Fourier series expansion:

$$\textstyle\varepsilon(\bm{x})\coloneqq\mathcal{P}[\tilde{u}](\bm{x})=\sum_{\bm{h}}\hat{\varepsilon}(\bm{h})\exp(2\pi\mathrm{i}\bm{h}\cdot\bm{x}),$$

where $\mathrm{i}$ denotes the imaginary unit and $\bm{h}=(h_{1},h_{2},\ldots,h_{s})$ $\in\mathbb{Z}^{s}$. Substituting this into (1) yields

$$\textstyle|(\ref{eq:closs})-(\ref{eq:dloss})|=\left|\frac{1}{N}\sum_{j=0}^{N-1}\sum_{\bm{h}\in\mathbb{Z}^{s},\bm{h}\neq 0}\hat{\varepsilon}(\bm{h})\exp(2\pi\mathrm{i}\bm{h}\cdot\bm{x}_{j})\right|,$$

(3)

because the Fourier mode of $\bm{h}=0$ is equal to the integral $\int_{[0,1]^{s}}\varepsilon(\bm{x})\mathrm{d}\bm{x}$. Before optimizing (3), the dual lattice of lattice $L$ and an insightful lemma are introduced as follows.

Lemma 4 follows directly from the properties of Fourier series. Based on this lemma, we restrict the lattice point $L$ to the form $\{\bm{x}\mid\bm{x}=\frac{j}{N}\bm{z}\mbox{ for }j\in\mathbb{Z}\}$ with a fixed integer vector $\bm{z}$; the set $L^{*}$ of collocation points is $\{\mbox{the decimal part of\ }\frac{j}{N}\bm{z}\mid j=0,\ldots,N-1\}$. Then, instead of searching $\bm{x}_{j}$’s, a vector $\bm{z}$ is searched. By restricting to this form, $\bm{x}_{j}$’s can be obtained automatically from a given $\bm{z}$, and hence the optimal collocation points $\bm{x}_{j}$’s do not need to be stored as a table of numbers, making a significant advantage in implementation. Another advantage is theoretical; the optimization problem of the collocation points can be reformulated in a number theoretic way. In fact, for $L$ as shown above, it is confirmed that $L^{\top}=\{\bm{h}\mid\bm{h}\cdot\bm{z}\equiv 0\pmod{N}\}$. If $\bm{h}\cdot\bm{z}\equiv 0\pmod{N}$ then there exists an $m\in\mathbb{Z}$ such that $\bm{h}\cdot\bm{z}=mN$ and hence $\frac{j}{N}\bm{h}\cdot\bm{z}=mj\in\mathbb{Z}$. Conversely, if $\bm{h}\cdot\bm{z}\not\equiv 0\pmod{N}$, clearly $\frac{1}{N}\bm{h}\cdot\bm{z}\notin\mathbb{Z}$.

From the above lemma,

$$\textstyle(\ref{eq:interror0})\leq\sum_{\bm{h}\in\mathbb{Z}^{s},\bm{h}\neq 0,\bm{h}\cdot\bm{z}\equiv 0\pmod{N}}|\hat{\varepsilon}(\bm{h})|,$$

(4)

and hence the collocation points $\bm{x}_{j}$’s should be determined so that (4) becomes small. This problem is a number theoretic problem in the sense that it is a minimization problem of finding an integer vector $\bm{h}$ subject to the condition $\bm{h}\cdot\bm{z}\equiv 0\pmod{N}$. This problem has been considered in the field of number theoretic numerical analysis. In particular, optimal solutions have been investigated for integrands in the Korobov spaces, which are spaces of functions that satisfy a certain smoothness condition.

It is known that if $\alpha$ is an integer, for a function $f$ to be in $E_{\alpha}$, it is sufficient that $f$ has continuous partial derivatives $\frac{\partial^{q_{1}+q_{2}+\cdots+q_{s}}}{\partial_{1}^{q_{1}}\cdot\partial_{2}^{q_{2}}\cdots\partial_{s}^{q_{s}}}f,0\leq q_{k}\leq\alpha\ (k=1,\ldots,s)$. For example, if a function $f(x,y):\mathbb{R}^{2}\to\mathbb{R}$ has continuous $f_{x},f_{y},f_{xy}$, then $f\in E_{1}$. Hence, if $\mathcal{P}[\tilde{u}]$ and the neural network belong to Koborov space,

$$\textstyle(\ref{eq:interror})\leq\sum_{\bm{h}\in\mathbb{Z}^{s},\bm{h}\neq 0,\bm{h}\cdot\bm{z}\equiv 0\pmod{N}}\frac{c}{(\bar{h}_{1}\bar{h}_{2}\cdots\bar{h}_{s})^{\alpha}}.$$

(5)

Here, we introduce a theorem in the field of number theoretic numerical analysis:

The main result of this paper is the following.

Intuitively, if $\mathcal{P}[\tilde{u}]$ satisfies certain conditions, we can find a set $L^{*}$ of collocation points with which the objective function (1) approximates the integral (2) only within an error of $O(\frac{(\log N)^{\alpha s}}{N^{\alpha}})$. This rate is much better than that of the uniformly random sampling (i.e., the Monte Carlo method), which is of $O(1/N^{\frac{1}{2}})$ (Sloan1994-cl), if the activation function of $\tilde{u}$ and hence the neural network $\tilde{u}$ itself are sufficiently smooth so that $\mathcal{P}[\tilde{u}]\in E_{\alpha}$ for a large $\alpha$. Hence, in this paper, we call the training method that minimizes (1) for a lattice $L$ satisfying (6) the good lattice training (GLT).

While any set of collocation points that satisfies the above condition will have the same convergence rate, a set constructed by the vector $\bm{z}\in\mathbb{Z}^{s}$ that minimizes (5) leads to better accuracy. When $s=2$, it is known that a good lattice can be constructed by setting $N=F_{k}$ and $\bm{z}=(1,F_{k-1})$, where $F_{k}$ denotes the $k$-th Fibonacci number (Niederreiter1992-qb; Sloan1994-cl). In general, an algorithm exists that can determine the optimal $\bm{z}$ with a computational cost of $O(N^{2})$. See Appendix “Theoretical Background” for more details. Also, we can retrieve the optimal $\bm{z}$ from numerical tables found in references, such as Fang1994; Keng1981.

The integrand $\mathcal{P}[\tilde{u}]$ of the loss function (2) does not always belong to the Korobov space $E_{\alpha}$ with high smoothness $\alpha$. To align the proposed GLT with theoretical expectations, we propose periodization tricks for ensuring periodicity and smoothness.

Given an initial condition at time $t=0$, the periodicity is ensured by extending the lattice twice as much along the time coordinate and folding it. Specifically, instead of $t$, we use $\hat{t}$ as the time coordinate that satisfies $t=2\hat{t}$ if $\hat{t}<0.5$ and $t=2(1-\hat{t})$ otherwise (see the lower right panel of Fig. 1, where the time is put on the horizontal axis). Also, while not mandatory, we combine the initial condition $u_{0}(\bm{x}_{\!/t})$ and the neural network’s output $\tilde{u}(t,\dots)$ as $\exp(-t)u_{0}(\bm{x}_{\!/t})+(1-\exp(-t))\tilde{u}(t,\bm{x}_{\!/t})$, thereby ensuring the initial condition, where $\bm{x}_{\!/t}$ denotes the set of coordinates except for the time coordinate $t$. A similar idea was proposed in Lagaris1998, which however does not ensure the initial condition strictly. If a periodic boundary condition is given to the $k$-th space coordinate, we bypass learning it and instead map the coordinate $x_{k}$ to a unit circle in two-dimensional space. Specifically, we map $x_{k}$ to $(x_{k}^{(1)},x_{k}^{(2)})=(\cos(2\pi x_{k}),\sin(2\pi x_{k}))$, assuring the loss function $\mathcal{P}[\tilde{u}]$ to take the same value at the both edges ($x_{k}=0$ and $x_{k}=1$) and be periodic. Given a Dirichlet boundary condition $u=0$ at $\partial\Omega$ to the $k$-th axis, we multiply the neural network’s output $\tilde{u}(\dots,x_{k},\dots)$ by $x_{k}(1-x_{k})$, and treat the result as the approximated solution. This ensures the Dirichlet boundary condition is met, and the loss function $\mathcal{P}[\tilde{u}]$ takes zero at the boundary $\partial\Omega$, thereby ensuring the periodicity. If a more complicated Dirichlet boundary condition is given, one can fold the lattice along the space coordinate in the same manner as the time coordinate and ensure the periodicity of the loss function $\mathcal{P}[\tilde{u}]$.

These periodization tricks aim to satisfy the periodicity conditions necessary for GLT to exhibit the performance shown in Theorem 7. However, they are also available for other sampling methods and potentially improve the practical performance by liberating them from the effort of learning initial and boundary conditions.

Since the GLT is grounded on the Fourier series foundation, it allows for the periodic shifting of the lattice. Hence, we randomize the collocation points as

$$\textstyle L^{*}=\{\mbox{the decimal part of\ }\frac{j}{N}\bm{z}+{\bm{r}}\mid j=0,\ldots,N-1\},$$

where ${\bm{r}}$ follows the uniform distribution over the unit cube $[0,1]^{s}$. Our preliminary experiments confirmed that, if using the stochastic gradient descent (SGD) algorithm, resampling the random numbers ${\bm{r}}$ at each training iteration prevents the neural network from overfitting and improves training efficiency. We call this approach the randomization trick.

Not only is this true for the proposed GLT, but most strategies to determine collocation points are not directly applicable to non-rectangular or non-flat domain $\Omega$ (Shankar2018). To achieve the best performance, the PDEs should be transformed to such a domain by an appropriate coordinate transformation. See Knupp2020-ix; Thompson1985-tx for examples.

Previous studies on numerical analysis addressed the periodicity and smoothness conditions on the integrand by variable transformations (Sloan1994-cl) (see also Appendix “Theoretical Background”). However, our preliminary experiments confirmed that it did not perform optimally in typical problem settings. Intuitively, these variable transformations reduce the weights of regions that are difficult to integrate, suppressing the discretization error. This implies that, when used for training, the regions with small weights remain unlearned. As a viable alternative, we introduced the periodization tricks to ensure periodicity.

The performance depends on the smoothness of the physics-informed loss, and hence on the smoothness of the neural network and the true solution. See Appendix “Theoretical Background” for details.

We modified the code from the official repository¹¹1https://github.com/maziarraissi/PINNs (MIT license) of Raissi2019, the original paper on PINNs. We obtained the datasets of the nonlinear Schrödinger (NLS) equation, Korteweg–De Vries (KdV) equation, and Allen-Cahn (AC) equation from the repository. The NLS equation governs wave functions in quantum mechanics, while the KdV equation models shallow water waves, and the AC equation characterizes phase separation in co-polymer melts. These datasets provide numerical solutions to initial value problems with periodic boundary conditions. Although they contain numerical errors, we treated them as the true solutions $u$. These equations are nonlinear versions of hyperbolic or parabolic PDEs. Additionally, as a nonlinear version of an elliptic PDE, we created a dataset for $s$-dimensional Poisson’s equation, which produces analytically solvable solutions with $2^{s}$ modes with the Dirichlet boundary condition. We examined the cases where $s\in\{2,4\}$. See Appendix “Experimental Settings” for further information.

Unless otherwise stated, we followed the repository’s experimental settings for the NLS equation. The physics-informed loss was defined as $\mathcal{P}[\tilde{u}]=\frac{1}{N}\sum_{j=0}^{N-1}\|\mathcal{N}[\tilde{u}](\bm{x}_{j})\|^{2}$ given $N$ collocation points $\{\bm{x}_{j}\}_{j=0}^{N-1}$. This can be regarded as a finite approximation to the squared 2-norm $|\mathcal{N}[\tilde{u}]|_{2}^{2}=\int_{\Omega}\|\mathcal{N}[\tilde{u}](\bm{x})\|^{2}\mathrm{d}\bm{x}$. The state of the NLS equation is complex; we simply treated it as a 2D real vector for training and used its absolute value for evaluation and visualization. Following Raissi2019, we evaluated the performance using the relative error, which is the normalized squared error $\mathcal{L}(\tilde{u},u;\bm{x}_{j})=(\sum_{j=0}^{N_{e}-1}\|\tilde{u}(\bm{x}_{j})-u(\bm{x}_{j})\|^{2})^{\frac{1}{2}}/(\sum_{j=0}^{N_{e}-1}\|u(\bm{x}_{j})\|^{2})^{\frac{1}{2}}$ at predefined $N_{e}$ collocation points $\{\bm{x}_{j}\}_{j=0}^{N_{e}-1}$. This is also a finite approximation to $|\tilde{u}-u|_{2}/|u|_{2}$.

We applied the above periodization tricks to each test and model for a fair comparison. For Poisson’s equation with $s=2$, which gives the exact solutions, we followed the original learning strategy using the L-BFGS-B method preceded by the Adam optimizer (Kingma2014b) for 50,000 iterations to ensure precise convergence. For other datasets, which contain the numerical solutions, we trained PINNs using the Adam optimizer with cosine decay of a single cycle to zero (Loshchilov2017) for 200,000 iterations and sampled a different set of collocation points at each iteration. See Appendix “Experimental Settings” for details.

We determined the collocation points using uniformly random sampling, uniformly spaced sampling, LHS, the Sobol sequence, and the proposed GLT. For the GLT, we took the number $N$ of collocation points and the corresponding integer vector $\bm{z}$ from numerical tables in (Fang1994; Keng1981). We used the same values for uniformly random sampling and LHS to maintain consistency. For uniformly spaced sampling, we selected numbers $N=m^{s}$ for $m\in\mathbb{N}$ that were closest to a number used in the GLT, creating a unit cube of $m$ points on each side. We additionally applied the randomization trick. For the Sobol sequence, we used $N=2^{m}$ for $m\in\mathbb{N}$ due to its theoretical background. We conducted five trials for each number $N$ and each method. All experiments were conducted using Python v3.7.16 and tensorflow v1.15.5 (tensorflow) on servers with Intel Xeon Platinum 8368.