Separable Physics-Informed Neural Networks
Physics-informed neural networks (PINNs) have recently emerged as promising
data-driven PDE solvers showing encouraging results on various PDEs. However,
there is a fundamental limitation of training PINNs to solve multi-dimensional
PDEs and approximate highly complex solution functions. The number of training
points (collocation points) required on these challenging PDEs grows
substantially, but it is severely limited due to the expensive computational
costs and heavy memory overhead. To overcome this issue, we propose a network
architecture and training algorithm for PINNs. The proposed method, separable
PINN (SPINN), operates on a per-axis basis to significantly reduce the number
of network propagations in multi-dimensional PDEs unlike point-wise processing
in conventional PINNs. We also propose using forward-mode automatic
differentiation to reduce the computational cost of computing PDE residuals,
enabling a large number of collocation points (>10^7) on a single commodity
GPU. The experimental results show drastically reduced computational costs (62x
in wall-clock time, 1,394x in FLOPs given the same number of collocation
points) in multi-dimensional PDEs while achieving better accuracy. Furthermore,
we present that SPINN can solve a chaotic (2+1)-d Navier-Stokes equation
significantly faster than the best-performing prior method (9 minutes vs 10
hours in a single GPU), maintaining accuracy. Finally, we showcase that SPINN
can accurately obtain the solution of a highly nonlinear and multi-dimensional
PDE, a (3+1)-d Navier-Stokes equation. For visualized results and code, please
see https://jwcho5576.github.io/spinn.github.io/.
Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, Eunbyung Park
# Separable Physics-Informed Neural Networks

###### Abstract

Physics-informed neural networks (PINNs) have recently emerged as promising data-driven PDE solvers showing encouraging results on various PDEs. However, there is a fundamental limitation of training PINNs to solve multi-dimensional PDEs and approximate highly complex solution functions. The number of training points (collocation points) required on these challenging PDEs grows substantially, but it is severely limited due to the expensive computational costs and heavy memory overhead. To overcome this issue, we propose a network architecture and training algorithm for PINNs. The proposed method, _separable PINN (SPINN)_, operates on a per-axis basis to significantly reduce the number of network propagations in multi-dimensional PDEs unlike point-wise processing in conventional PINNs. We also propose using forward-mode automatic differentiation to reduce the computational cost of computing PDE residuals, enabling a large number of collocation points (\(>10^{7}\)) on a single commodity GPU. The experimental results show drastically reduced computational costs (\(62\times\) in wall-clock time, \(1,394\times\) in FLOPs given the same number of collocation points) in multi-dimensional PDEs while achieving better accuracy. Furthermore, we present that SPINN can solve a chaotic (2+1)-d Navier-Stokes equation significantly faster than the best-performing prior method (9 minutes vs 10 hours in a single GPU), maintaining accuracy. Finally, we showcase that SPINN can accurately obtain the solution of a highly nonlinear and multi-dimensional PDE, a (3+1)-d Navier-Stokes equation. For visualized results and code, please see [https://jwcho5576.github.io/spinn.github.io/](https://jwcho5576.github.io/spinn.github.io/).

## 1 Introduction

Solving partial differential equations (PDEs) has been a long-standing problem in various science and engineering domains. Finding analytic solutions requires in-depth expertise and is often infeasible in many useful and important PDEs [11]. Hence, numerical approximation methods to solutions have been extensively studied [17], e.g., spectral methods [4], finite volume methods (FVM) [10], finite difference method (FDM) [38], and finite element methods (FEM) [45]. While successful, classical methods have several limitations, such as expensive computational costs, requiring sophisticated techniques to support multi-physics and multi-scale systems, and the curse of dimensionality in high dimensional PDEs.

With the vast increases in computational power and methodological advances in machine learning, researchers have explored data-driven and learning-based methods [2; 5; 24; 29; 33; 37]. Among the promising methods, physics-informed neural networks (PINNs) have recently emerged as new data-driven PDE solvers for both forward and inverse problems [34]. PINNs employ neural networks and
gradient-based optimization algorithms to represent and obtain the solutions, leveraging automatic differentiation to enforce the physical constraints of underlying PDE. It has enjoyed great success in various forward and inverse problems thanks to its numerous benefits, such as flexibility in handling a wide range of forward and inverse problems, mesh-free solutions, and not requiring observational data, hence, unsupervised training.

Despite its advantages and promising results, there is a fundamental limitation of training PINN to solve multi-dimensional PDEs and approximate very complex solution functions. It primarily stems from using coordinate-based MLP architectures to represent the solution function, which takes input coordinates and outputs corresponding solution quantities. For each training point, computing PDE residual loss involves multiple forward and backward propagations, and the number of training points (collocation points) required to solve multi-dimensional PDEs and obtain more accurate solutions grows substantially. The situation deteriorates as the dimensionality of the PDE or the solution's complexity increases.

Recent studies have presented empirical evidence showing that choosing a larger batch size (i.e., a large number of collocation points) in training PINNs leads to enhanced precision [34, 35, 36, 42]. Furthermore, more collocation points also accelerate the convergence speed due to the scaling characteristic between the batch size and the learning rate (the larger the batch size, the higher the learning rate) [39, 27, 12]. As long as the computational resources allow, we have the flexibility to utilize a substantial number of training points in PINNs since they can be continuously sampled from the input domain in an unsupervised training manner.

We propose a novel PINN architecture, _separable PINN (SPINN)_, which utilizes forward-mode* automatic differentiation (AD) to enable a large number of collocation points (\(>10^{7}\) in a single GPU), reducing the computational cost of solving multi-dimensional PDEs. Instead of feeding every multi-dimensional coordinate into a single MLP, we use separated sub-networks, in which each sub-network takes independent one-dimensional coordinates as input. The final output is generated by an aggregation module such as simple outer product and element-wise summation where the predicted solution can be interpreted by low-rank tensor approximation [25] (Fig. 4b). The suggested architecture obviates the need to query every multi-dimensional coordinate input pair, exponentially reducing the number of network propagations to generate a solution, \(\mathcal{O}(N^{d})\rightarrow\mathcal{O}(Nd)\), where \(N\) is the resolution of the solution for each dimension, and \(d\) is the dimension of the system.

Footnote *: a.k.a. forward accumulation mode or tangent linear mode.

We have conducted comprehensive experiments on the representative PDEs to show the effectiveness of the suggested method. The experimental results demonstrate that the computational costs of the conventional PINN increase linearly with the number of collocation points, while the proposed model shows logarithmic growths. This allows SPINN to accommodate orders of magnitude larger number of collocation points _in a single batch_ during training. We also show that given the same number of training points, SPINN improves wall-clock training time up to by \(62\times\) on commodity GPUs and FLOPs up to by \(1,394\times\) while achieving better accuracy. Furthermore, with large-scale collocation points, SPINN can solve a turbulent Navier-Stokes equation much faster than the state-of-the-art PINN method [42] (9 minutes vs 10 hours in a single GPU) without bells and whistles, such as causal

Figure 1: Training speed (w/ a single GPU) of our model compared to the causal PINN [42] in (2+1)-d Navier-Stokes equation of time interval [0, 0.1].
inductive bias in the loss function (Fig. 1). Our experimental results of the Navier-Stokes equation show that SPINN can solve highly nonlinear PDEs and is sufficiently expressive to represent complex functions. This is further supported by the provided theoretical result, potentiating the use of our method for more challenging and various PDEs.

## 2 Related Works

Physics-informed neural networks.Physics-Informed Neural Networks (PINNs) [34] have received great attention as a promising learning-based PDE solver. Given the underlying PDE and initial, boundary conditions embedded in a loss function, a coordinate-based neural network is trained to approximate the desired solution. Since its inception, many techniques have been studied to improve training PINNs for more challenging PDE systems [26; 42; 43], or to accelerate the training speed [20; 36]. Our method is orthogonal to most of the previously suggested techniques above and improves PINNs' training from a computational perspective.

The effect of collocation points.The residual loss of PINNs is calculated by Monte Carlo integration, not an exact definite integral. Therefore, it inherits a core property of Monte Carlo methods [15]: the impact of the number of sampled points. The importance of sampling strategy and the number of collocation points in training PINNs has been highlighted by recent works [8; 23; 35; 42]. Especially, Sankaran et al. [35] empirically found that training with a large number of collocation points is unconditionally favorable for PINNs in terms of accuracy and convergence speed. Another line of research established a theoretical upper bound of PINN's statistical error with respect to the number of collocation points [23]. It showed that a larger number of collocation points are required to use a bigger network size for training. Our work builds on this evidence to bring PINNs out in more practical scenarios and overcome their limitation.

Derivative computations in PINNs.Several studies have introduced techniques to increase the efficiency of derivative calculations in training PINNs, where obtaining derivatives with respect to input coordinates is an essential process, often achieved through reverse-mode automatic differentiation. DT-PINN [36] utilized the finite difference (FD) method to accelerate the training. They specifically employed the radial basis function-finite differences (RBF-FD) method [41], which can be applied to the irregular domain to maintain the mesh-agnostic characteristic of PINNs. However, the RBF-FD method can only be used for calculating spatial derivative, and DT-PINN still had to employ AD for temporal derivative. Another work, fPINN [32], also used the FD method to solve fractional-order PDEs, which addresses the limitation of automatic differentiation. CAN-PINN [6] adopted both FD and AD to redefine the loss function and accelerate the training. Employing Taylor-mode AD [13] in training PINNs was introduced by causal PINN [42] to handle high-order PDEs such as Kuramoto-Sivashinsky equation [28]. To the best of our knowledge, the proposed method is the first approach to leverage forward-mode AD in training PINNs, which is fully applicable to both time-dependent and independent PDEs and does not incur any truncation errors.

Multiple MLP networks.Employing multiple MLPs for PINNs has been introduced by several works to utilize parallelized training [19; 21; 31]. They share the same concept of decomposing the spatio-temporal domain and training multiple individual MLPs on each sub-domain. Although these methods showed promising results, they still suffer from the fundamental problem of heavy computation as the number of collocation points increases. While these methods decompose input domains, and each small MLP is used to cover a particular subdomain, we decompose input dimensions and solve PDEs over the entire domain cooperated by all separated MLPs. In terms of the model architecture and function representation, our work is also related to NAM [1], Haghighat et al. [14], and CoordX [30]. NAM [1] suggested separated network architectures and inputs, but only for achieving the interpretability of the model's prediction in multi-task learning. Haghighat et al. [14] used multiple MLPs to individually predict each component of the _output_ vector. CoordX's [30] primary purpose is to reconstruct natural signals, such as images or 3D shapes. Hence, they are not motivated to improve the efficiency of computing higher-order gradients. In addition, they had to use additional layers after the feature merging step, which makes their model more computationally expensive than ours. Furthermore, in neural fields [44], there is an explicit limitation in the number of ground truth data points (e.g., the number of pixels in an image). Therefore, CoordX cannot fully maximize the advantage of enabling a large number of input coordinates. We focus on solving
PDEs and carefully devised the architecture to exploit forward-mode AD to efficiently compute PDE residual losses.

## 3 Preliminaries: Forward/Reverse-mode AD

For the completeness of this paper, we start by briefly introducing the two types of AD and how Jacobian matrices are evaluated. For clarity, we will follow the notations used in [3] and [13]. Suppose our function \(f:\mathbb{R}^{n}\to\mathbb{R}^{m}\) is a two-layer MLP with 'tanh' activation. The left-hand side of Tab. 1 demonstrates a single forward trace of \(f\). To obtain a \(m\times n\) Jacobian matrix \(\mathbb{J}_{f}\) in forward-mode, we compute the Jacobian-vector product (JVP),

\[\mathbb{J}_{f}r=\begin{bmatrix}\frac{\partial y_{1}}{\partial x_{1}}&\dots& \frac{\partial y_{1}}{\partial x_{n}}\\ \vdots&\ddots&\vdots\\ \frac{\partial y_{m}}{\partial x_{1}}&\dots&\frac{\partial y_{m}}{\partial x _{n}}\end{bmatrix}\begin{bmatrix}\frac{\partial x_{1}}{\partial x_{1}}\\ \vdots\\ \frac{\partial x_{n}}{\partial x_{1}}\end{bmatrix}, \tag{1}\]

for \(i\in\{1,\dots,n\}\). The forward-mode AD is a one-phase process: while tracing primals (intermediate values) \(v_{k}\), it continues to evaluate and accumulate their tangents \(v_{k}=\partial v_{k}/\partial x_{i}\) (the middle column of Tab. 1). This is equivalent to decomposing one large JVP into a series of JVPs by the chain rule and computing them from right to left. A run of JVP with the initial tangents \(v_{0}\) as the first column vector of an identity matrix \(I_{n}\) gives the first column of \(\mathbb{J}_{f}\). Thus, the full Jacobian can be obtained in \(n\) forward passes.

On the other hand, the reverse-mode AD computes vector-Jacobian product (VJP):

\[r^{\top}\mathbb{J}_{f}=\begin{bmatrix}\frac{\partial y_{j}}{ \partial y_{1}}&\dots&\frac{\partial y_{j}}{\partial y_{m}}\\ \vdots&\ddots&\vdots\\ \frac{\partial y_{m}}{\partial x_{1}}&\dots&\frac{\partial y_{m}}{\partial x _{n}}\end{bmatrix}, \tag{2}\]

for \(j\in\{1,\dots,m\}\), which is the reverse-order operation of JVP. This is a two-phase process. The first phase corresponds to forward propagation, storing all the primals, \(v_{k}\), and recording the elementary operations in the computational graph. In the second phase, the derivatives are computed by accumulating the adjoints \(\bar{v}_{k}=\partial y_{j}/\partial v_{k}\) (the right-hand side of Tab. 1). Since VJP builds one row of a Jacobian at a time, it takes \(m\) evaluations to obtain the full Jacobian. To sum up, the forward-mode is more efficient for a tall Jacobian (\(m>n\)), while the reverse-mode is better suited for a wide Jacobian (\(n>m\)). Fig. 2 shows an illustrative example and please refer to Baydin et al. [3] for more details.

\begin{table}
\begin{tabular}{l l l l} \hline Forward primal trace & Forward tangent trace & Backward adjoint trace \\ \hline \(v_{0}=x\) & \(\frac{\dot{v}_{0}}{v_{1}}=W_{1}\cdot v_{0}\) & \(\frac{\dot{v}_{0}}{\dot{v}_{1}}=W_{1}\cdot\dot{v}_{0}\) & \(\frac{\dot{x}}{\dot{v}_{0}}=\bar{v}_{1}\cdot\frac{\partial v_{1}}{\partial x_ {0}}=\bar{v}_{1}\cdot W_{1}\) \\ \(v_{2}=\text{tanh}(v_{1})\) & \(\dot{v}_{2}=\text{tanh}^{\prime}(v_{1})\circ\dot{v}_{1}\) & \(\bar{v}_{1}=\bar{v}_{2}\cdot\frac{\partial v_{2}}{\partial v_{1}}=\bar{v}_{2} \circ\text{tanh}^{\prime}(v_{1})\) \\ \(v_{3}=W_{2}\cdot v_{2}\) & \(\frac{\dot{v}_{3}}{\dot{v}_{2}}=W_{2}\cdot\dot{v}_{2}\) & \(\frac{\dot{v}_{2}}{\dot{v}_{3}}=\bar{v}_{3}\cdot\frac{\dot{v}_{3}}{\dot{v}_{2}}= \bar{v}_{3}\cdot W_{2}\) \\ \hline \(y=v_{3}\) & \(\bar{y}=v_{3}\) & \(\bar{y}=v_{3}\) & \(\frac{\dot{x}}{\dot{v}_{3}}=\bar{y}\) \\ \hline \end{tabular}
\end{table}
Table 1: An example of forward and reverse-mode AD in a two-layers tanh MLP. Here \(v_{0}\) denotes the input variable, \(v_{k}\) the primals, \(\dot{v}_{k}\) the tangents, \(\bar{v}_{k}\) the adjoints, and \(W_{1},W_{2}\) the weight matrices. Biases are omitted for brevity.

Figure 3: An illustrative example of separated approach vs non-separated approach when \(f^{(\theta)}:\mathbb{R}^{d}\to\mathbb{R}\) (an example case of \(N=3\), \(d=3\) is shown above). The number of JVP (forward-mode) evaluations (propagations) of computing the Jacobian for separated approach (a) is \(Nd\), while the number of VJP (reverse-mode) evaluations for non-separated approach (b) is \(N^{d}\).

Figure 2: Simple neural networks with different input and output dimensions. To compute \(\frac{\partial y}{\partial x}\), (a) requires one forward pass using forward-mode AD or three backward passes using reverse-mode AD, (b) requires three forward passes using forward-mode AD or one backward pass using reverse-mode AD or one backward pass using reverse-mode AD.
## 4 Separable PINN

### Forward-Mode AD with Separated Functions

We demonstrate that leveraging forward-mode AD and separating the function into multiple functions with respect to input axes can significantly reduce the cost of computing Jacobian. In the proposed separated approach (Fig. 2(a)), we first sample \(N\) one-dimensional coordinates on each of \(d\) axes, which makes a total of \(Nd\) batch size. Next, these coordinates are fed into \(d\) individual functions. Let \(f\) be a function which takes coordinate as input to produce feature representation and we denote the number of operations of \(f\) as \(\mathsf{ops}(f)\). Then, a feature merging function \(h\) is used to construct the solution of the entire \(N^{d}\) discretized points. The amount of computations for AD is known to be 2\(\sim\)3 times more expensive than the forward propagation [3; 13]. According to Fig. 3 and with the scale constants \(c_{f},c_{h}\in[2,3]\) we can approximate the total number of operations to compute the Jacobian matrix of the proposed separated approach (Fig. 2(a)).

\[\mathcal{C}_{\text{sep}}=Ndc_{f}\mathsf{ops}(f)+N^{d}c_{h}\mathsf{ops}(h). \tag{3}\]

For a non-separated approach (Fig. 2(b)),

\[\mathcal{C}_{\text{non-sep}}=N^{d}c_{f}\mathsf{ops}(f), \tag{4}\]

If we can make \(\mathsf{ops}(h)\) sufficiently small, then the ratio \(\mathcal{C}_{\text{sep}}/\mathcal{C}_{\text{non-sep}}\) becomes \(\frac{Nd}{N^{d}}\ll 1\), quickly converging to 0 as \(d\) and \(N\) increases. Our separated approach has almost linear complexity with respect to \(N\), implying it can obtain more accurate solutions (high-resolution) efficiently.

### Network Architecture

Fig. 3(a) illustrates the overall SPINN architecture, parameterizing multiple separated functions with neural networks. SPINN consists of \(d\) body-networks (MLPs), each of which takes an individual 1-dimensional coordinate component as an input. Each body-network \(f^{(\theta_{i})}:\mathbb{R}\rightarrow\mathbb{R}^{r}\) (parameterized by \(\theta_{i}\)) is a vector-valued function which transforms the coordinates of \(i\)-th axis into a \(r\)-dimensional feature representation. The final prediction is computed by feature merging:

\[\hat{u}(x_{1},x_{2},\ldots,x_{d})=\sum_{j=1}^{r}\prod_{i=1}^{d}f_{j}^{(\theta _{i})}(x_{i}) \tag{5}\]

where \(\hat{u}:\mathbb{R}^{d}\rightarrow\mathbb{R}\) is the predicted solution function, \(x_{i}\in\mathbb{R}\) is a coordinate of \(i\)-th axis, and \(f_{j}^{(\theta_{i})}\) denotes the \(j\)-th element of \(f^{(\theta_{i})}\). We used 'tanh' activation function throughout the paper. As shown in Eq. 5, the feature merging operation is a simple product (\(\Pi\)) and summation (\(\Sigma\)) which

Figure 4: (a) SPINN architecture in a 3-dimensional system. To solve a \(d\)-dimensional PDE, our model requires \(d\) body MLP networks, each of which takes individual scalar coordinate values as input and gives \(r\)-dimensional feature vector. The final output is obtained by element-wise product and summation. (b) Construction process of the entire discretized solution tensor when the inputs are given in batches. Each outer product between the column vectors \(F_{:,j,i}\) from the feature tensor \(F\) constructs a rank-1 tensor and summing all the \(r\) tensors gives a rank-\(r\) tensor. The output tensor of SPINN can be interpreted as a low-rank decomposed representation of a solution.
corresponds to the merging function \(h\) described in Eq. 3. Due to its simplicity, \(h\) operations are much cheaper than operations in MLP layers (i.e., \(\mathsf{ops}(h)\ll\mathsf{ops}(f)\) in Eq. 3). Note that SPINN can also approximate any \(m\)-dimensional vector functions \(\hat{u}:\mathbb{R}^{d}\to\mathbb{R}^{m}\) by using a larger output feature size (see section D.4 in the appendix for details).

The collocation points of our model and conventional PINNs have a distinct difference (Fig. 5). Both are _uniformly_ evaluated on a \(d\)-dimensional hypercube, but collocation points of SPINN form a lattice-like structure, which we call as _factorizable coordinates_. In SPINN, 1-dimensional input points from each axis are randomly sampled, and \(d\)-dimensional points are generated via the cartesian product of the point sets from each axis. On the other hand, non-factorizable coordinates are randomly sampled points without any structure. Factorizable coordinates with our separated MLP architecture enable us to evaluate functions on dense (\(N^{d}\)) collocation points with a small number (\(Nd\)) of input points.

In practice, the input coordinates are given in a batch during training and inference. Assume that \(N\) input coordinates (training points) are sampled from each axis. Note that the sampling resolutions for each axis need not be the same. The input coordinates \(X\in\mathbb{R}^{N\times d}\) is now a matrix. The batchified form of feature representation \(F\in\mathbb{R}^{N\times r\times d}\) and Eq. 5 now becomes

\[\hat{U}(X_{:,1},X_{:,2},\ldots,X_{:,d})=\sum_{j=1}^{r}\bigotimes_{i=1}^{d}F_{:,j,i}, \tag{6}\]

where \(\hat{U}\in\mathbb{R}^{N\times N\times\ldots\times N}\) is the discretized solution tensor, \(\bigotimes\) denotes outer product, \(F_{:,:,i}\in\mathbb{R}^{N\times r}\) is an \(i\)-th frontal slice matrix of tensor \(F\), and \(F_{:,j,i}\in\mathbb{R}^{N}\) is the \(j\)-th column of the matrix \(F_{:,:,i}\). Fig. 4b shows an illustrative procedure of Eq. 6. Due to its structural input points and outer products between feature vectors, SPINN's solution approximation can be viewed as a low-rank tensor decomposition where the feature size \(r\) is the rank of the reconstructed tensor. Among many decomposition methods, SPINN corresponds to CP-decomposition [16], which approximates a tensor by finite summation of rank-1 tensors. While traditional methods use iterative methods such as alternating least-squares (ALS) or alternating slicewise diagonalization (ASD) [22] to directly fit the decomposed vectors, we train neural networks to learn the decomposed vector representation and approximate the solution functions in continuous input domains. This, in turn, allows for the calculation of derivatives with respect to arbitrary input coordinates.

### Gradient Computation of SPINN

In this section, we show that the number of JVP (forward-mode AD) evaluations for computing the full gradient of SPINN (\(\nabla\hat{u}(x)\)) is \(Nd\), where \(N\) is the number of coordinates for each axis and \(d\) is the input dimension. According to Eq. 5, the \(i\)-th element of \(\nabla\hat{u}(x)\) is:

\[\frac{\partial\hat{u}}{\partial x_{i}}=\sum_{j=1}^{r}f_{j}^{(\theta_{1})}(x_{ 1})f_{j}^{(\theta_{2})}(x_{2})\ldots\frac{\partial f_{j}^{(\theta_{i})}(x_{i} )}{\partial x_{i}}\ldots f_{j}^{(\theta_{d})}(x_{d}). \tag{7}\]

Computing this derivative requires feature representations \(f_{j}^{(\theta_{i})}\), but we can reuse them from the forward pass results computed beforehand. Note that the entire \(r\) components of the feature derivatives \(\partial f^{(\theta_{i})}(x_{i})/\partial x_{i}:\mathbb{R}\to\mathbb{R}^{r}\) can be obtained by a single pass thanks to forward-mode AD. To obtain the full gradient vector \(\nabla\hat{u}(x):\mathbb{R}^{d}\to\mathbb{R}^{d}\), we iterate the calculation of Eq. 7 over \(d\) times, switching the input axis \(i\). Also, since each iteration involves \(N\) training samples, the total number of JVP evaluations becomes \(Nd\), which is consistent with Eq. 3.

Figure 5: An illustrative 2-dimensional example of (a) factorizable and (b) non-factorizable coordinates. (a) has a lattice-like structure, where SPINN can be evaluated on more dense collocation points with fewer input points. Conventional PINNs sample non-factorizable coordinates which do not have any structures.
### Universal Approximation Property

It is widely known that neural networks with sufficiently many hidden units have the expressive power of approximating any continuous functions [7; 18]. However, it is not straightforward that our suggested architecture enjoys the same capability. For the completeness of the paper, we provide a universal approximation property of the proposed method.

**Theorem 1**.: _(Proof in appendix) Let \(X,Y\) be compact subsets of \(\mathbb{R}^{d}\). Choose \(u\in L^{2}(X\times Y)\). Then, for arbitrary \(\varepsilon>0\), we can find a sufficiently large \(r>0\) and neural networks \(f_{j}\) and \(g_{j}\) such that_

\[\left\|u-\sum_{j=1}^{r}f_{j}g_{j}\right\|_{L^{2}(X\times Y)}<\varepsilon. \tag{8}\]

By repeatedly applying Theorem 1, we can show that SPINN can approximate any functions in \(L^{2}\) in the high-dimensional input space. An approximation property for a broader function space is a fruitful research area, and we leave it to future work. To support the expressive power of the suggested method, we empirically showed that SPINN can accurately approximate solution functions of various challenging PDEs, including diffusion, Helmholtz, Klein-Gordon, and Navier-Stokes equations.

## 5 Experiments

### Experimental setups

We compared SPINN against vanilla PINN [34] on 3-d (diffusion, Helmholtz, Klein-Gordon, and Navier-Stokes equation) and 4-d (Klein-Gordon and Navier-Stokes) PDE systems. Every experiment was run on a different number of collocation points, and we also applied the modified MLP introduced in [43] to both PINN and SPINN. For 3-d systems, the number of collocation points of \(64^{3}\) was the upper limit for the vanilla PINN (\(54^{3}\) for modified MLP) when we trained with a single NVIDIA RTX3090 GPU with 24GB of memory. However, the memory usage of our model was significantly smaller, enabling SPINN to use a larger number of collocation points (up to \(256^{3}\)) to get more accurate solutions. This is because SPINN stores a much smaller batch of tensors, which are the primals (\(v_{k}\) in Tab. 1) while building the computational graph. All reported error metrics are average relative \(L_{2}\) errors computed by \(\|\hat{u}-u\|^{2}/\|u\|^{2}\), where \(\hat{u}\) is the model prediction, and \(u\) is the reference solution. Every experiment is performed seven times (three times for (2+1)-d and five times for (3+1)-d Navier-Stokes, respectively) with different random seeds. More detailed experimental settings are provided in the appendix.

Figure 6: Overall results of (a) diffusion, (b) Helmholtz, (c) (2+1)-d Klein-Gordon, and (d) (3+1)-d Klein-Gordon experiments. It shows comparisons among PINN, PINN with modified MLP (PINN-mod), SPINN, and SPINN with modified MLP (SPINN-mod). For each experiment, the first to third columns show the relative error, runtime, and GPU memory versus different numbers of collocation points, respectively. The rightmost column shows the training curves of each model when the number of collocation points is \(64^{3}\) (\(54^{3}\) for PINN with modified MLP.) For (3+1)-d Klein-Gordon, the training curve is plotted when each model is trained with the number of collocation points of \(16^{4}\). The error and the training curves are averaged over 7 different runs and 70% confidence intervals are provided. Note that the y-axis scale of every plot is a log scale.
### Results

Fig. 6 shows the overall results of forward problems on three 3-d systems (diffusion, Helmholtz, and Klein-Gordon) and one 4-d system (Klein-Gordon). SPINN is significantly more computationally efficient than the baseline PINN in wall-clock run-time. For every PDE, SPINN with the modified MLP found the most accurate solution. Furthermore, when the number of collocation points grows exponentially, the memory usage and the actual run-time of SPINN increase almost linearly. We also confirmed that we can get more accurate solutions with more collocation points. This training characteristic of SPINN substantiates that our method is very effective for solving multi-dimensional PDEs. Furthermore, with the help of the method outlined in Griewank and Walther [13], we estimated the FLOPs for evaluating the derivatives. Compared to the baseline, SPINN requires \(1,394\times\) fewer operations to compute the forward pass, first and second order derivatives. Further details regarding the FLOPs estimation are provided in section C of the appendix. In the following sections, we give more detailed descriptions and analyses of each experiment.

Diffusion EquationWhen trained with \(128^{3}\) collocation points, SPINN with the modified MLP finds the most accurate solution. Furthermore, when trained with the same number of collocation points (\(64^{3}\)), SPINN-mod is \(52\times\) faster and \(29\times\) more memory efficient than the baseline with modified MLP. You can find the visualized solution in the appendix. Since we construct the solution relatively simple compared to other PDE experiments, baseline PINNs also find a comparatively good solution. However, our model finds a more accurate solution with a larger number of collocation points with minor computational overhead. The exact numerical values are provided in the appendix.

Helmholtz EquationThe results of the Helmholtz experiments are shown in Fig. 5(b). Due to stiffness in the gradient flow, conventional PINNs hinder finding accurate solutions. Therefore, a modified MLP and learning rate annealing algorithm is suggested to mitigate such phenomenon [43]. We found that even PINNs with modified MLP fail to solve the equation when the solution is complex (contains high-frequency components). However, SPINN obtains one order of magnitude lower relative error without bells and whistles. This is because each body network of the SPINN learns an individual 1-dimensional function which mitigates the burden of representing the entire 3-dimensional function. This makes our model much easier to learn complex functions. Given the same number of collocation points (\(64^{3}\)), the training speed of our proposed model with modified MLP is \(62\times\) faster, and the memory usage is \(29\times\) smaller than the baseline with modified MLP. The exact numerical values are provided in the appendix.

Klein-Gordon EquationFig. 5(c) shows the results. Again, our method shows the best performance in terms of accuracy, runtime (\(62\times\)), and memory usage (\(29\times\)). Note that both Helmholtz and Klein-Gordon equations contain three second-order derivative terms, while the diffusion equation contains only two. Our results showed the largest differences in runtime and memory usage in Helmholtz and Klein-Gordon equations. This is because SPINN significantly reduces the AD computations with forward-mode AD, implying SPINN is very efficient for solving high-order PDEs.

We investigated the Klein-Gordon experiment further, extending it to a 4-d system by adding one more spatial axis. As shown in Fig. 5(d), baseline PINN can only process \(23^{4}\) (\(18^{4}\) for modified MLP) collocation points at once due to a high memory throughput. On the other hand, SPINN can exploit the number of collocation points of \(64^{4}\) (\(160\times\) larger than the baseline) to obtain a more precise solution. The exact numerical values are provided in the appendix.

(2+1)-d Navier-Stokes EquationWe constructed SPINN to predict the velocity field \(u\) and applied forward-mode AD to obtain the vorticity \(\omega\). We used the same PDE setting used in causal PINN [42] and compared SPINN against their model. As shown in Tab. 2, our model finds a comparably accurate solution even without the causal loss function. Furthermore, since causal PINN had to iterate over four different tolerance values \(\epsilon\) in the causal loss function, they had to run 3\(\sim\)4 times more training epochs, depending on the stopping criterion. As a result, given the same number of collocation points, SPINN converges 60\(\times\) faster than causal PINN in terms of training runtime. The results of the Navier-Stokes equation ensure that SPINN can successfully solve a chaotic, highly nonlinear PDE. In Fig. 7, we

Figure 7: Relative error and training speed of (2+1)-d Navier-Stokes experiment vs. different ranks of SPINN.
demonstrate the effect of the rank \(r\) for solving the equation. We observed that the performance almost converges at the rank of 128, and increasing the rank does not slow down the training speed too much. Further details and visualization of the predicted vorticity map is shown in appendix section D.4.

(3+1)-d Navier-Stokes EquationWe proceeded to explore our model on (3+1)-d Navier-Stokes equation by devising a manufactured solution corresponding to the Taylor-Green vortex [40]. It models a decaying vortex, widely used for testing Navier-Stokes numerical solvers [9]. Since the vorticity of the (3+1)-d Navier-Stokes equation is a 3-d vector as opposed to the (2+1)-d system, the equation has three independent components resulting in a total of 33 derivative terms (see section D.5 in the appendix). Similar to the (2+1)-d experiment, the network's output is the velocity field, and we obtained 3-d vorticity by forward-mode AD. Tab. 3 shows the relative error, runtime, and GPU memory. Trained with a single GPU, SPINN achieves a relative error of 1.9e-3 less than 30 minutes. Visualized velocity and vorticity vector fields are provided in the appendix.

## 6 Limitations and Future Works

Although we presented comprehensive experimental results to show the effectiveness of our model, there remain questions on solving more challenging and higher dimensional PDEs. Combining the existing PINNs training techniques, such as learning-rate annealing [43], curriculum learning [26], and causal loss function [42], into ours is an exciting research direction, and we leave it to future work. Applying SPINN to higher dimensional PDEs, such as the BGK equation, would also have a tremendous practical impact. Lastly, we are still far from the theoretical speed-up (wall-clock runtime vs. FLOPs), which is further discussed in appendix section C We expect to additionally reduce runtime by optimizing hardware/software, e.g., using customized CUDA kernels.

## 7 Conclusion

We showed a simple yet powerful method to employ large-scale collocation points for PINNs training, leveraging forward-mode AD with the separated MLPs. Experimental results demonstrated that our method significantly reduces both spatial and computational complexity while achieving better accuracy on various three and four-dimensional PDEs. To our knowledge, it is the first attempt to exploit the power of forward-mode AD. We believe this work opens up a new direction to rethink the architectural design of neural networks in many scientific applications.

\begin{table}
\begin{tabular}{c|c c|c c c|c c c} \hline \hline model & \multicolumn{2}{c|}{PINN+mod [43]} & \multicolumn{3}{c|}{causal PINN [42]} & \multicolumn{3}{c}{SPINN (ours)} \\ \hline \(N_{c}\) & \(2^{12}\) & \(2^{15}\) & \(2^{12}\) & \(2^{15}\) & \(2^{15}\) & \(2^{15}\) & \(2^{18}\) & \(2^{21}\) \\ \hline \multirow{3}{*}{relative \(L_{2}\) error} & 0.0694 & 0.0581 & 0.0578 & 0.0401 & \multirow{3}{*}{0.0353\({}^{\dagger}\)} & 0.0780 & 0.0363 & 0.0355 \\  & \(\pm\)0.0091 & \(\pm\)0.0135 & \(\pm\)0.0117 & \(\pm\)0.0084 & & & & \\ \cline{1-1}  & 03:20 & 07:52 & 10:09 & 23:03 & - & 00:07 & 00:09 & 00:14 \\ \cline{1-1}  & 5,198 & 17,046 & 5,200 & 17,132 & - & 764 & 892 & 1,276 \\ \hline \hline \end{tabular}
\end{table}
Table 2: The numerical result of (2+1)-d Navier-Stokes equation compared against PINN with modified MLP (PINN+mod) and causal PINN. \(N_{c}\) is the number of collocation points. The fourth row shows the training runtime of a single time window in time-marching training. All relative \(L_{2}\) errors are averaged over 3 runs.

\({}^{\dagger}\) This is the error (the single run result) reported by the causal PINN paper. Since the causal PINN is sensitive to the choice of parameter initialization, we also reported average error obtained from different random seeds by running their official code.

\begin{table}
\begin{tabular}{c|c c c c} \hline \hline model & \(N_{c}\) & \(\begin{array}{c}\text{relative}\\ L_{2}\text{ error}\end{array}\) & runtime & memory \\ \hline  & \(8^{4}\) & 0.0090 & 15.33 & 768 \\ SPINN & \(16^{4}\) & 0.0041 & 16.83 & 1,192 \\  & \(32^{4}\) & 0.0019 & 26.73 & 2,946 \\ \hline \hline \end{tabular}
\end{table}
Table 3: The numerical result of (3+1)-d Navier-Stokes equation. \(N_{c}\) is the number of collocation points on the entire domain. The relative errors are obtained by comparing the vorticity vector field and averaged over five different runs.
## Acknowledgments and Disclosure of Funding

This research was supported by the Ministry of Science and ICT (MSIT) of Korea, under the National Research Foundation (NRF) grant (2022R1F1A1064184, 2022R1A4A3033571), Institute of Information and Communication Technology Planning Evaluation (IITP) grants for the AI Graduate School program (IITP-2019-0-00421). The research of Seok-Bae Yun was supported by Samsung Science and Technology Foundation under Project Number SSTF-BA1801-02.

## References

* Agarwal et al. [2021] Rishabh Agarwal, Levi Melnick, Nicholas Frosst, Xuezhou Zhang, Ben Lengerich, Rich Caruana, and Geoffrey E Hinton. Neural additive models: Interpretable machine learning with neural nets. _Advances in Neural Information Processing Systems_, 34:4699-4711, 2021.
* Avrutskiy [2020] Vsevolod I Avrutskiy. Neural networks catching up with finite differences in solving partial differential equations in higher dimensions. _Neural Computing and Applications_, 32(17):13425-13440, 2020.
* Baydin et al. [2018] Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. _Journal of Machine Learning Research_, pages 1-43, 2018.
* Boyd [2001] John P Boyd. _Chebyshev and Fourier spectral methods_. Courier Corporation, 2001.
* Brunton and Kutz [2022] Steven L Brunton and J Nathan Kutz. _Data-driven science and engineering: Machine learning, dynamical systems, and control_. Cambridge University Press, 2022.
* Chiu et al. [2022] Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong. Can-pinn: A fast physics-informed neural network based on coupled-automatic-numerical differentiation method. _Computer Methods in Applied Mechanics and Engineering_, 395:114909, 2022.
* Cybenko [1989] George Cybenko. Approximation by superpositions of a sigmoidal function. _Mathematics of control, signals and systems_, 2(4):303-314, 1989.
* Daw et al. [2022] Arka Daw, Jie Bu, Sifan Wang, Paris Perdikaris, and Anuj Karpatne. Rethinking the importance of sampling in physics-informed neural networks. _arXiv preprint arXiv:2207.02338_, 2022.
* Ethier and Steinman [1994] C Ross Ethier and DA Steinman. Exact fully 3d navier-stokes solutions for benchmarking. _International Journal for Numerical Methods in Fluids_, 19(5):369-375, 1994.
* Eymard et al. [2000] Robert Eymard, Thierry Gallouet, and Raphaele Herbin. Finite volume methods: handbook of numerical analysis. _PG Ciarlet and JL Lions (Eds)_, 2000.
* Fefferman [2000] Charles L Fefferman. Existence and smoothness of the navier-stokes equation. _The millennium prize problems_, 57:67, 2000.
* Goyal et al. [2017] Priya Goyal, Piotr Dollar, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. _arXiv preprint arXiv:1706.02677_, 2017.
* Griewank and Walther [2008] Andreas Griewank and Andrea Walther. _Evaluating derivatives: principles and techniques of algorithmic differentiation_. SIAM, 2008.
* Haghighat et al. [2021] Ehsan Haghighat, Maziar Raissi, Adrian Moure, Hector Gomez, and Ruben Juanes. A physics-informed deep learning framework for inversion and surrogate modeling in solid mechanics. _Computer Methods in Applied Mechanics and Engineering_, 379:113741, 2021.
* Hammersley [2013] John Hammersley. _Monte carlo methods_. Springer Science & Business Media, 2013.
* Hitchcock [1927] Frank L Hitchcock. The expression of a tensor or a polyadic as a sum of products. _Journal of Mathematics and Physics_, 6(1-4):164-189, 1927.
* Hoffman and Frankel [2018] Joe D Hoffman and Steven Frankel. _Numerical methods for engineers and scientists_. CRC press, 2018.
* Hornik et al. [1989] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. _Neural networks_, 2(5):359-366, 1989.
* [19] Ameya D Jagtap and George E Karniadakis. Extended physics-informed neural networks (xpinns): A generalized space-time domain decomposition based deep learning framework for nonlinear partial differential equations. In _AAAI Spring Symposium: MLPS_, 2021.
* [20] Ameya D Jagtap, Kenji Kawaguchi, and George Em Karniadakis. Adaptive activation functions accelerate convergence in deep and physics-informed neural networks. _Journal of Computational Physics_, 404:109136, 2020.
* [21] Ameya D Jagtap, Ehsan Kharazmi, and George Em Karniadakis. Conservative physics-informed neural networks on discrete domains for conservation laws: Applications to forward and inverse problems. _Computer Methods in Applied Mechanics and Engineering_, 365:113028, 2020.
* [22] Jian-hui Jiang, Hai-long Wu, Yang Li, and Ru-qin Yu. Three-way data resolution by alternating slice-wise diagonalization (asd) method. _Journal of Chemometrics: A Journal of the Chemometrics Society_, 14(1):15-36, 2000.
* [23] Yuling Jiao, Yanming Lai, Dingwei Li, Xiliang Lu, Fengru Wang, Yang Wang, and Jerry Zhijian Yang. A rate of convergence of physics informed neural networks for the linear second order elliptic pdes. _Communications in Computational Physics_, 31(4):1272, 2022.
* [24] George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. _Nature Reviews Physics_, 3(6):422-440, 2021.
* [25] Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. _SIAM review_, 51(3):455-500, 2009.
* [26] Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Characterizing possible failure modes in physics-informed neural networks. _Advances in Neural Information Processing Systems_, 34:26548-26560, 2021.
* [27] Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. _arXiv preprint arXiv:1404.5997_, 2014.
* [28] Yoshiki Kuramoto and Toshio Tsuzuki. Persistent propagation of concentration waves in dissipative media far from thermal equilibrium. _Progress of theoretical physics_, 55(2):356-369, 1976.
* [29] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In _International Conference on Learning Representations_, 2021.
* [30] Ruofan Liang, Hongyi Sun, and Nandita Vijaykumar. Coordx: Accelerating implicit neural representation with a split mlp architecture. In _International Conference on Learning Representations_, 2022.
* [31] Ben Moseley, Andrew Markham, and Tarje Nissen-Meyer. Finite basis physics-informed neural networks (fpbirns): a scalable domain decomposition approach for solving differential equations. _arXiv preprint arXiv:2107.07871_, 2021.
* [32] Guofei Pang, Lu Lu, and George Em Karniadakis. fpinns: Fractional physics-informed neural networks. _SIAM Journal on Scientific Computing_, 41(4):A2603-A2626, 2019.
* [33] Maziar Raissi. Deep hidden physics models: Deep learning of nonlinear partial differential equations. _The Journal of Machine Learning Research_, 19(1):932-955, 2018.
* [34] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. _Journal of Computational physics_, 378:686-707, 2019.
* [35] Shyam Sankaran, Hanwen Wang, Leonardo Ferreira Guilhoto, and Paris Perdikaris. On the impact of larger batch size in the training of physics informed neural networks. In _The Symbiosis of Deep Learning and Differential Equations II_, 2022.
* [36] Ramansh Sharma and Varun Shankar. Accelerated training of physics informed neural networks (pinns) using meshless discretizations. _Advances in Neural Information Processing Systems_, 2022.
* [37] Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. _Journal of computational physics_, 375:1339-1364, 2018.
* [38] Gordon Dennis Smith. _Numerical solution of partial differential equations: finite difference methods_. Oxford university press, 1985.
* [39] Samuel L. Smith, Pieter-Jan Kindermans, and Quoc V. Le. Don't decay the learning rate, increase the batch size. In _International Conference on Learning Representations_, 2018.
* [40] Geoffrey Ingram Taylor and Albert Edward Green. Mechanism of the production of small eddies from large ones. _Proceedings of the Royal Society of London. Series A-Mathematical and Physical Sciences_, 158(895):499-521, 1937.
* [41] AI Tolstykh and DA Shirobokov. On using radial basis functions in a "finite difference mode" with applications to elasticity problems. _Computational Mechanics_, 33(1):68-79, 2003.
* [42] Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality is all you need for training physics-informed neural networks. _arXiv preprint arXiv:2203.07404_, 2022.
* [43] Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow pathologies in physics-informed neural networks. _SIAM Journal on Scientific Computing_, 43(5):A3055-A3081, 2021.
* [44] Yiheng Xie, Towaki Takikawa, Shunsuke Saito, Or Litany, Shiqin Yan, Numair Khan, Federico Tombari, James Tompkin, Vincent Sitzmann, and Srinath Sridhar. Neural fields in visual computing and beyond. In _Computer Graphics Forum_, pages 641-676. Wiley Online Library, 2022.
* [45] Olek C Zienkiewicz, Robert Leroy Taylor, and Jian Z Zhu. _The finite element method: its basis and fundamentals_. Elsevier, 2005.
## Appendix A Proof of Theorem 1

Here we show the preliminary lemmas and proofs for Theorem 1. in the main paper. We start by defining a general tensor product between two Hilbert spaces.

**Definition 1**.: Let \(\{v_{\beta}\}\) be an orthonormal basis for \(\mathcal{H}_{2}\). **Tensor product** between Hilbert spaces \(\mathcal{H}_{1}\) and \(\mathcal{H}_{2}\), denoted by \(\mathcal{H}_{1}\otimes\mathcal{H}_{2}\), is a set of all antilinear mappings \(A:\mathcal{H}_{2}\rightarrow\mathcal{H}_{1}\) such that \(\sum_{\beta}\lVert Av_{\beta}\rVert^{2}<\infty\) for every orthonormal basis for \(\mathcal{H}_{2}\).

Then by Theorem 7.12 in Folland [3], \(\mathcal{H}_{1}\otimes\mathcal{H}_{2}\) is also a Hilbert space with respect to norm \(\left|\kern-1.075pt\left|\kern-1.075pt\left|\cdot\right|\right|\kern-1.075pt\right|\) and associated inner product \(\langle\cdot,\cdot\rangle\):

\[\left|\kern-1.075pt\left|\kern-1.075pt\left|A\right|\kern-1.075pt\right| \kern-1.075pt\right|^{2} \equiv\sum_{\beta}\lVert Av_{\beta}\rVert^{2}, \tag{1}\] \[\langle A,B\rangle \equiv\sum_{\beta}\langle Av_{\beta},Bv_{\beta}\rangle, \tag{2}\]

where \(A,B\in\mathcal{H}_{1}\otimes\mathcal{H}_{2}\), and \(\{v_{\beta}\}\) is any orthonormal basis of \(\mathcal{H}_{2}\).

**Lemma 1**.: _Let \(x\in\mathcal{H}_{1}\) and \(y,y^{\prime}\in\mathcal{H}_{2}\). Then, \((x\otimes y)y^{\prime}=\langle y,y^{\prime}\rangle x\)._

Proof.: Let \(A\) be a mapping \(A:y^{\prime}\rightarrow\langle y,y^{\prime}\rangle x\), where \(\lVert y\rVert=1\). We expand \(y^{\prime}\) with orthonormal basis \(\{y,z_{\beta}\}\). i.e., \(\{y,z_{\beta}\}\) is a basis for \(\mathcal{H}_{2}\). Then,

\[\lVert Ay\rVert^{2}+\sum_{\beta}\lVert Az_{\beta}\rVert^{2} =\lVert Ay\rVert^{2}+\sum_{\beta}\lVert\langle y,z_{\beta}\rangle x \rVert^{2} \tag{3}\] \[=\lVert Ay\rVert^{2}<\infty \tag{4}\]

It is obvious that \(A\) is antilinear. Then by Definition 1, \(A\in\mathcal{H}_{1}\otimes\mathcal{H}_{2}\) which is, \(Ay^{\prime}=(x\otimes y)y^{\prime}\). Therefore, \((x\otimes y)y^{\prime}=\langle y,y^{\prime}\rangle x\).

**Lemma 2**.: \(\{u_{\alpha}\otimes v_{\beta}\}\) _is an orthonormal basis for \(\mathcal{H}_{1}\otimes\mathcal{H}_{2}\)._

Proof.: Let \(A,B\in\mathcal{H}_{1}\otimes\mathcal{H}_{2}\), where \(B=u_{\alpha}\otimes v_{\beta}\). Then by the definition of inner product in Eq. 2,

\[\langle A,B\rangle =\sum_{i}\langle Av_{i},Bv_{i}\rangle\] (5) \[=\sum_{i}\langle Av_{i},(u_{\alpha}\otimes v_{\beta})v_{i}\rangle\] (6) \[=\sum_{i}\langle Av_{i},\langle v_{\beta},v_{i}\rangle u_{\alpha} \rangle\qquad(\because\text{Lemma 1})\] (7) \[=\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\langle\!\langle\!\langle\!\langle\langle\! \langle\!\langle\!\langle\!\langle\!\langle\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\! \langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\! \langle\!\langle\!\langle\!\langle\!\langle\!\!\langle\!\langle\!\langle\!\langle\!\! \langle\!\langle\!\langle\!\langle\!\langle\!\!\langle\!\langle\!\!\langle\!\langle\!\langle \!\langle\!\langle\!\langle\!\!\langle\!\langle\!\!\langle\!\!\langle\!\langle\!\langle\!\! \langle\!\!\langle\!\langle\!\langle\!\langle\!\langle\!\langle\!\!\langle\!\langle\!\! \langle\!\langle\!\langle\!\!\langle\!\langle\!\langle\!\!\langle\!\langle\!\langle\!\! \langle\!\langle\!\!\langle\!\langle\!\!\langle\!\langle\!\langle\!\langle\!\langle\!\!\langle\!\langle \!\langle\!\!\langle\!\langle\!\!\langle\!\!\langle\!\langle\!\langle\!\!\langle\!\langle\!\!\langle\!\! \langle\!\langle\!\langle\!\!\langle\!\!\langle\!\!\langle\!\!\langle\!\!\langle\!\!\langle\!\!\langle\!\langle\!\! \langle\!\!\langle\!\langle\!\langle\!\langle\!\!\langle\!\!\langle\!\langle\!\!\langle\!\!\langle\!\!\langle\!\langle\!\! \langle\!\langle\!\langle\!\!\langle\!\langle\!\!\langle\!\langle\!\!\langle\
\(\therefore\{u_{\alpha}\otimes v_{\beta}\}\) is a basis. 

Now we begin the proof of Theorem 1. in the main paper.

**Theorem 1**.: _Let \(X,Y\) be compact subsets of \(\mathbb{R}^{d}\). Choose \(u\in L^{2}(X\times Y)\). Then, for arbitrary \(\varepsilon>0\), we can find a sufficiently large \(r>0\) and neural networks \(f_{j}\) and \(g_{j}\) such that_

\[\left\|u-\sum_{j=1}^{r}f_{j}g_{j}\right\|_{L^{2}(X\times Y)}<\varepsilon. \tag{13}\]

Proof.: Let \(\{\phi_{i}\}\) and \(\{\psi_{j}\}\) be orthonormal basis for \(L^{2}(X)\) and \(L^{2}(Y)\) respectively. Then \(\{\phi_{i}\psi_{j}\}\) forms an orthonormal basis for \(L^{2}(X\times Y)\) (\(\therefore\) Lemma 2). Therefore, we can find a sufficiently large \(r\) such that

\[\left\|u-\sum_{i,j}^{r}a_{ij}\phi_{i}\psi_{j}\right\|_{L^{2}(X\times Y)}<\frac {\varepsilon}{2}, \tag{14}\]

where \(a_{ij}\) denotes

\[a_{ij}=\int_{X\times Y}u(x,y)\phi_{i}(x)\psi_{j}(y)dxdy.\]

On the other hand, by the universal approximation theorem [1], we can find neural networks \(f_{j}\) and \(g_{j}\) such that

\[\|\phi_{i}-f_{j}\|_{L^{2}(X)}\leq\frac{\varepsilon}{3^{j}\|u\|_{L^{2}(X \times Y)}}\text{ and }\|\psi_{j}-g_{j}\|_{L^{2}(Y)}\leq\frac{\varepsilon}{3^{j}\|u\|_{L^{2}(X \times Y)}}. \tag{15}\]

We first consider the difference between \(u\) and \(\sum_{i,j}^{r}a_{ij}f_{i}g_{j}\):

\[\left\|u-\sum_{i,j}^{r}a_{ij}f_{i}g_{j}\right\|_{L^{2}(X\times Y)} \tag{16}\] \[\leq\left\|u-\sum_{i,j}^{r}a_{ij}\phi_{i}\psi_{j}\right\|_{L^{2} (X\times Y)}+\left\|\sum_{i,j}^{r}a_{ij}\phi_{i}\psi_{j}-\sum_{i,j}^{r}a_{ij} f_{i}g_{j}\right\|_{L^{2}(X\times Y)}\] (17) \[\equiv I+II \tag{18}\]

Since \(|I|<\varepsilon/2\) from (14), it is enough to estimate \(II\). For this, we consider

\[\sum_{i,j}^{r}a_{ij}\phi_{i}\psi_{j}-\sum_{i,j}^{r}a_{ij}f_{i}g_{j} =\sum_{i,j}^{r}a_{ij}\phi_{i}\big{(}\psi_{j}-g_{j}\big{)}+\sum_{i,j}^{r}a_{ij}\big{(}\phi_{i}-f_{i}\big{)}g_{j} \tag{19}\] \[\equiv II_{1}+II_{2}. \tag{20}\]

We first compute \(II_{1}\):

\[\|II_{1}\|_{L^{2}(X\times Y)}^{2} =\int_{X\times Y}\left\{\sum_{i,j}^{r}a_{ij}\phi_{i}\big{(}\psi_ {j}-f_{j}\big{)}\right\}^{2}dxdy \tag{21}\] \[=\int_{X\times Y}\left\{\sum_{j}^{r}\left(\sum_{i}^{r}a_{ij} \phi_{i}\right)(\psi_{j}-f_{j})\right\}^{2}dxdy \tag{22}\]

We set

\[A_{j}(x)=\sum_{i}^{r}a_{ij}\phi_{i}(x),\qquad B_{j}(y)=\psi_{j}(y)-g_{j}(y) \tag{23}\]
to write \(II_{1}\) as

\[\|II_{1}\|_{L^{2}(X\times Y)}^{2}=\int_{X\times Y}\left\{\sum_{j}^{r}A_{j}(x)B_{j} (y)\right\}^{2}dxdy. \tag{24}\]

We then apply Cauchy-Scharwz inequality to get

\[\|II_{1}\|_{L^{2}(X\times Y)}^{2} \leq\int_{X\times Y}\Big{(}\sum_{j}^{r}|A_{j}(x)|^{2}\Big{)}\Big{(} \sum_{j}^{r}|B_{j}(y)|^{2}\Big{)}dxdy \tag{25}\] \[=\left(\int_{X}\sum_{j}^{r}|A_{j}(x)|^{2}dx\right)\left(\int_{Y} \sum_{j}^{r}|B_{j}(y)|^{2}dy\right). \tag{26}\]

Now

\[\int_{X}\sum_{j}^{r}|A_{j}(x)|^{2}dx =\int_{X}\sum_{j}^{r}\left(\sum_{i}^{r}a_{ij}\phi_{i}\right)^{2}dx \tag{27}\] \[=\sum_{j}^{r}\Big{\|}\sum_{i}^{r}a_{ij}\phi_{i}\Big{\|}_{L(X)}^{ 2}. \tag{28}\]

Since \(\{\phi_{i}\}\) is an orthonormal basis, we see that

\[\Big{\|}\sum_{i}^{r}a_{ij}\phi_{i}\Big{\|}_{L(X)}\leq\sum_{i}^{r}|a_{ij}| \big{\|}\phi_{i}\big{\|}_{L(X)}=\sum_{i}^{r}|a_{ij}|. \tag{29}\]

Therefore,

\[\int_{X}\sum_{j}^{r}|A_{j}(x)|^{2}dx =\int_{X}\sum_{j}^{r}\left(\sum_{i}^{r}a_{ij}\phi_{i}\right)^{2}dx \tag{30}\] \[=\sum_{j}^{r}\int_{X}\left(\sum_{i}^{r}a_{ij}\phi_{i}\right)^{2}dx\] (31) \[=\sum_{j}^{r}\Big{\|}\sum_{i}^{r}a_{ij}\phi_{i}\Big{\|}_{L(X)}^{ 2}\] (32) \[=\sum_{j}^{r}\Big{(}\sum_{i}^{r}|a_{ij}|\Big{)}^{2}. \tag{33}\]

Finally, we recall

\[\Big{(}\sum_{i}^{r}|a_{ij}|\Big{)}^{2}\leq 2\sum_{i}^{r}|a_{ij}|^{2} \tag{34}\]

to conclude

\[\int_{X}\sum_{j}^{r}|A_{j}(x)|^{2}dx\leq 2\sum_{j}^{r}\sum_{i}^{r}|a_{ij}|^{2 }<2\sum_{i,j}^{\infty}|a_{ij}|^{2}=2\|u\|_{L^{2}(X\times Y)}^{2}. \tag{35}\]
On the other hand, we have from (15)

\[\int_{Y}\sum_{j}^{r}|B_{j}(y)|^{2}dy =\sum_{j}^{r}\int_{Y}|\psi_{j}(y)-g_{j}(y)|^{2}dy \tag{36}\] \[=\sum_{j}^{r}\left\|\psi_{j}(y)-g_{j}(y)\right\|_{L^{2}(Y)}^{2}\] (37) \[\leq\sum_{j}^{r}\frac{\varepsilon^{2}}{\mathfrak{g}^{j}\|u\|_{L^ {2}(X\times Y)}^{2}}\] (38) \[<\frac{\varepsilon^{2}}{8\|u\|_{L^{2}(X\times Y)}^{2}}. \tag{39}\]

Hence we have

\[\|II_{1}\|_{L^{2}(X\times Y)}^{2}<\frac{\varepsilon^{2}}{8\|u\|_{L^{2}(X \times Y)}^{2}}2\|u\|_{L^{2}(X\times Y)}^{2}=\frac{\varepsilon^{2}}{4}. \tag{40}\]

Likewise,

\[\|II_{2}\|_{L^{2}(X\times Y)}^{2}<\frac{\varepsilon^{2}}{4}. \tag{41}\]

Therefore,

\[\|II\|_{L^{2}(X\times Y)}^{2}<\frac{\varepsilon^{2}}{2}. \tag{42}\]

We go back to (16) with the estimates (14) and (42) to derive

\[\left\|u-\sum_{i,j}^{r}a_{ij}f_{i}g_{j}\right\|_{L^{2}(X\times Y)} <\frac{\varepsilon}{2}+\frac{\varepsilon}{2}=\varepsilon. \tag{43}\]

Finally, we reorder the index to rewrite

\[\sum_{i,j}^{r}a_{ij}f_{i}g_{j} =\sum_{i}^{\tilde{r}}b_{i}\tilde{f}_{i}\tilde{g}_{i} \tag{44}\] \[=\sum_{i}^{\tilde{r}}\left\{b_{i}\tilde{f}_{i}\right\}\tilde{g}_ {i}\] (45) \[=\sum_{i}^{\tilde{r}}k_{i}\tilde{g}_{i} \tag{46}\]

Without loss of generality, we rewrite \(\tilde{r},k_{i}\), \(\tilde{g}_{i}\) into \(r\), \(f_{i}\), \(g_{i}\) respectively, to complete the proof. 

## Appendix B Training with Physics-Informed Loss

After SPINN predicts an output function with the methods described above, the rest of the training procedure follows the same process used in conventional PINN training [11], except we use forward-mode AD to compute PDE residuals (standard back-propagation, a.k.a. reverse-mode AD for parameter updates). With the slight abuse of notation, our predicted solution function is denoted as \(\hat{u}^{(\theta)}(x,t)\) from onwards, explicitly expressing time coordinates. Given an underlying PDE (or ODE), the initial, and the boundary conditions, SPINN is trained with a 'physics-informed' loss function:

\[\min_{\theta}\mathcal{L}(\hat{u}^{(\theta)}(x,t))=\min_{\theta} \lambda_{\text{pde}}\mathcal{L}_{\text{pde}}+\lambda_{\text{ic}}\mathcal{L}_{ \text{ic}}+\lambda_{\text{bc}}\mathcal{L}_{\text{bc}}, \tag{47}\] \[\mathcal{L}_{\text{pde}}=\int_{\Gamma}\int_{\Omega}\|\mathcal{N} [\hat{u}^{(\theta)}](x,t)\|^{2}dxdt,\] (48) \[\mathcal{L}_{\text{ic}}=\int_{\Omega}\|\hat{u}^{(\theta)}(x,0)-u _{\text{ic}}(x)\|^{2}dx,\] (49) \[\mathcal{L}_{\text{bc}}=\int_{\Gamma}\int_{\partial\Omega}\| \mathcal{B}[\hat{u}^{(\theta)}](x,t)-u_{\text{bc}}(x,t)\|^{2}dxdt, \tag{50}\]
where \(\Omega\) is an input domain, \(\mathcal{N},\mathcal{B}\) are generic differential operators and \(u_{\text{ic}}\), \(u_{\text{bc}}\) are initial, boundary conditions, respectively. \(\lambda\) are weighting factors for each loss term. When calculating the PDE loss (\(\mathcal{L}_{\text{pole}}\)) with Monte-Carlo integral approximation, we sampled collocation points from factorized coordinates and used forward-mode AD. The remaining \(\mathcal{L}_{\text{ic}}\) and \(\mathcal{L}_{\text{bc}}\) are then computed with initial and boundary coordinates to regress the given conditions. By minimizing the objective loss in Eq. 47, the model output is enforced to satisfy the given equation, the initial, and the boundary conditions.

## Appendix C FLOPs Estimation

The FLOPs for evaluating the derivatives can be systematically calculated by disassembling the computational graph into elementary operations such as additions and multiplications. Given a computational graph of forward pass for computing the primals, AD augments each elementary operation into other elementary operations. The FLOPs in the forward pass can be precisely calculated since it consists of a series of matrix multiplications and additions. We used the method described in [4] to estimate FLOPs for evaluating the derivatives. Table 1 shows the number of additions (ADDS) and multiplications (MULTS) in each evaluation process. Note that FLOPs is a summation of ADDS and MULTS by definition.

One thing to note here is that this is a theoretical estimation. Theoretically, the number of JVP evaluations for computing the gradient with respect to the input coordinates is \(Nd\), when \(N\) is the number of coordinates for each axis and \(d\) is the input dimension (see section 4.3 in the main paper). However, our actual implementation of gradient calculation involves re-computing the feature representations \(f^{(\theta_{i})}\), which makes the complexity of network propagations from \(\mathcal{O}(Nd)\) to \(\mathcal{O}(Nd^{2})\). Ideally, these feature vectors can be computed only once and stored to be used later for gradient computations. Although it is still significantly more efficient than the conventional PINN's complexity (\(Nd^{2}\ll N^{d}\)), there is a room to bridge the gap between theoretical FLOPs and actual training runtime by further software optimization.

## Appendix D Experimental Details and Results

In this section, we provide experimental details, numerical results, and visualizations for each experiment in the main paper.

### Diffusion Equation

The diffusion equation is one of the most representative parabolic PDEs, often used for modeling the heat diffusion process. We especially choose a nonlinear diffusion equation where it can be written as:

\[\partial_{t}u=\alpha\left(\|\nabla u\|^{2}+u\Delta u\right), x\in\Omega,t\in\Gamma, \tag{51}\] \[u(x,0)=u_{\text{ic}}(x), x\in\Omega,\] (52) \[u(x,t)=0, x\in\partial\Omega,t\in\Gamma. \tag{53}\]

We used diffusivity \(\alpha=0.05\), spatial domain \(\Omega=[-1,1]^{2}\), temporal domain \(\Gamma=[0,1]\) and used superposition of three Gaussian functions for the initial condition \(u_{\text{ic}}\). We obtained the reference

\begin{table}
\begin{tabular}{c|c|c|c|c} \hline \hline  & \multicolumn{2}{c|}{SPINN (ours)} & \multicolumn{2}{c}{PINN (baseline)} \\ \hline  & ADDS (\(\times 10^{6}\)) & MULTS (\(\times 10^{6}\)) & ADDS (\(\times 10^{6}\)) & MULTS (\(\times 10^{6}\)) \\ \hline forward pass & 20 & 20 & 21,609 & 21,609 \\
1st-order derivative & 40 & 40 & 86,638 & 43,419 \\
2nd-order derivative & 80 & 80 & 130,057 & 87,040 \\ \hline MFLOPs (total) & \multicolumn{2}{c|}{280} & \multicolumn{2}{c}{390,370} \\ \hline \hline \end{tabular}
\end{table}
Table 1: The number of elementary operations for evaluating forward pass, first and second-order derivatives. The calculation is based on \(64^{3}\) collocation points in a 3-d system and the vanilla MLP settings used for diffusion, Helmholtz, and Klein-Gordon equations. We assumed that each derivative is evaluated on every coordinate axis.
solution (\(101\times 101\times 101\) resolution) through a widely-used PDE solver platform FEniCS [9]. Note that FEniCS is a FEM-based solver. We particularly set the initial condition to be a superposition of three gaussian functions:

\[u_{\text{ic}}(x,y)=0.25\exp\left[-10\{(x-0.2)^{2}+(y-0.3)^{2}\} \right]+0.4\exp\left[-15\{(x+0.1)^{2}+(y+0.5)^{2}\}\right]\] \[+0.3\exp\left[-20\{(x+0.5)^{2}+y^{2}\}\right]. \tag{54}\]

For our model, we used three body networks of 4 hidden layers with 64/32 hidden feature/output size each. For the baseline model, we used a single MLP of 5 hidden layers with 128 hidden feature sizes. We used Adam [6] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors \(\lambda\) in the loss function in Eq. 47 are set to 1. The final reported errors are extracted where the total loss was minimum across the entire training iteration. We also resampled the input points every 100 epochs. Tab. 3 shows the numerical results, and the visualized solutions are provided in Fig. 1.

### Helmholtz Equation

The Helmholtz equation is a time-independent wave equation that takes the form:

\[\Delta u+k^{2}u=q, x\in\Omega, \tag{55}\] \[u(x)=0, x\in\partial\Omega, \tag{56}\]

where the spatial domain is \(\Omega=[-1,1]^{3}\). For a given source term \(q=-(a_{1}\pi)^{2}u-(a_{2}\pi)^{2}u-(a_{3}\pi)^{2}u+k^{2}u\), we devised a manufactured solution \(u=\sin(a_{1}\pi x_{1})\sin(a_{2}\pi x_{2})\sin(a_{3}\pi x_{3})\), where we take \(k=1,a_{1}=4,a_{2}=4,a_{3}=3\).

For our model, we used three body networks of 4 hidden layers with 64/32 hidden feature/output size each. For the baseline model, we used a single MLP of 5 hidden layers with 128 hidden feature sizes. We used Adam [6] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors \(\lambda\) in the loss function in Eq. 47 are set to 1. The final reported errors are extracted where the total loss was minimum across the entire training iteration. We also resampled the input points every 100 epochs. Tab. 4 shows the numerical results and the visualized solutions are provided in Fig. 2.

### Klein-Gordon Equation

The Klein-Gordon equation is a nonlinear hyperbolic PDE, which arises in diverse applied physics for modeling relativistic wave propagation. The inhomogeneous Klein-Gordon equation is given by

\[\partial_{tt}u-\Delta u+u^{2}=f, x\in\Omega,t\in\Gamma, \tag{57}\] \[u(x,0)=x_{1}+x_{2}, x\in\Omega,\] (58) \[u(x,t)=u_{\text{bc}}(x), x\in\partial\Omega,t\in\Gamma, \tag{59}\]

where we chose the spatial/temporal domain to be \(\Omega=[-1,1]^{2}\) and \(\Gamma=[0,10]\), respectively. For error measurement, we used a manufactured solution \(u=(x_{1}+x_{2})\cos(2t)+x_{1}x_{2}\sin(2t)\) and \(f\), \(u_{bc}\) are extracted from this exact solution.

For our model, we used three body networks of 4 hidden layers with 64/32 hidden feature/output size each. For the baseline model, we used a single MLP of 5 hidden layers with 128 hidden feature sizes. We used Adam [6] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors \(\lambda\) in the loss function in Eq. 47 are set to 1. The final reported errors are extracted where the total loss was minimum across the entire training iteration. We also resampled the input points every 100 epochs. Tab. 5 shows the numerical results.

We used the same settings used in (2+1)-d Klein-Gordon experiment for the (3+1)-d experiment except the manufactured solution was chosen as:

\[u=(x_{1}+x_{2}+x_{3})\cos{(t)}+x_{1}x_{2}x_{3}\sin{(t)}, \tag{60}\]

where \(f\), \(u_{bc}\) are extracted from this exact solution. Tab. 6 shows the numerical results. The collocation sizes of \(23^{4},18^{4}\) were the maximum value for PINN and PINN with modified MLP, respectively.
### (2+1)-d Navier-Stokes Equation

Navier-Stokes equation is a nonlinear time-dependent PDE that describes the motion of a viscous fluid. Various engineering fields rely on this equation, such as modeling the weather, airflow, or ocean currents. The vorticity form for incompressible fluid can be written as below:

\[\partial_{t}\omega+u\cdot\nabla\omega=\nu\Delta\omega, x\in\Omega,t\in\Gamma, \tag{61}\] \[\nabla\cdot u=0, x\in\Omega,t\in\Gamma,\] (62) \[\omega(x,0)=\omega_{0}(x), x\in\Omega, \tag{63}\]

where \(u\in\mathbb{R}^{2}\) is the velocity field, \(\omega=\nabla\times u\) is the vorticity, \(\omega_{0}\) is the initial vorticity, and \(\nu\) is the viscosity. We used the viscosity 0.01 and made the spatial/temporal domain \(\Omega=[0,2\pi]^{2}\) and \(\Gamma=[0,1]\), respectively. Note that Eq. 61 models decaying turbulence since there is no forcing term and Eq. 62 is the incompressible fluid condition. The reference solution is generated by JAX-CFD solver [7] which specifically used the pseudo-spectral method. The initial condition was generated using the gaussian random field with a maximum velocity of 5. The resolution of the obtained solution is \(100\times 128\times 128\left(N_{t}\times N_{x}\times N_{y}\right)\), and we tested our model on this data.

For our model, we used three body networks (modified MLP) of 3 hidden layers with 128/256 hidden feature/output sizes each. We divided the temporal domain into ten time windows to adopt the time marching method [8; 13]. We used Adam [6] optimizer with a learning rate of 0.002 and each time window is trained for 100,000 iterations. Followed by causal PINN, the PDE (residual) loss and initial condition loss function are written as follows.

\[\mathcal{L}_{\text{pde}} =\frac{\lambda_{w}}{N_{c}}\sum_{i=1}^{N_{c}}|\partial_{t}w+u_{x} \partial_{x}w+u_{y}\partial_{y}w-\nu(\partial_{xx}w+\partial_{yy}w)|^{2}+\frac {\lambda_{c}}{N_{c}}\sum_{i=1}^{N_{c}}|\partial_{x}u_{x}+\partial_{y}u_{y}|^{2}, \tag{64}\] \[\mathcal{L}_{\text{ic}} =\frac{\lambda_{ic}}{N_{ic}}\sum_{i=1}^{N_{ic}}\left(|u_{x}-u_{x 0}|^{2}+|u_{y}-u_{y0}|^{2}+|w-w_{0}|^{2}\right), \tag{65}\]

where \(u_{x},u_{y}\) are \(x,y\) components of predicted velocity, \(w=\partial_{x}u_{y}-\partial_{y}u_{x}\), \(N_{c}\) is the collocation size, \(N_{ic}\) is the number of coordinates for initial condition and \(u_{x0},u_{y0},u_{w0}\) are the initial conditions. We chose the weighting factors \(\lambda_{w}=1\), \(\lambda_{c}=5,000\), and \(\lambda_{ic}=10,000\). We also resampled the input points every 100 epochs. The periodic boundary condition can be explicitly enforced by positional encoding [2], and we specifically used the following encoding function only for the spatial input coordinates.

\[\gamma(x)=[1,\sin(x),\sin(2x),\sin(3x),\sin(4x),\sin(5x),\cos(x), \cos(2x),\cos(3x),\cos(4x),\cos(5x)]^{\top}. \tag{66}\]

Unlike other experiments, the solution of the Navier-Stokes equation is a 2-dimensional vector-valued function \(u:\mathbb{R}^{3}\rightarrow\mathbb{R}^{2}\). We can rewrite the feature merging equation Eq. 5 in the main paper to construct SPINN into a 2-dimensional vector function:

\[u_{1} =\sum_{j=1}^{r}\prod_{i=1}^{d}f_{j}^{(\theta_{i})}(x_{i}), \tag{67}\] \[u_{2} =\sum_{j=r+1}^{2r}\prod_{i=1}^{d}f_{j}^{(\theta_{i})}(x_{i}). \tag{68}\]

For example in our (2+1)-d Navier-Stokes equation setting, \(r=128\) since the network output feature size is 256. This can be applied to any \(m\)-dimensional vector function if we use a larger output feature size. More formally, if we want to construct an \(m\)-dimensional vector function with SPINN of rank \(r\), the \(k\)-th element of the function output can be written as

\[u_{k}=\sum_{j=(k-1)r+1}^{kr}\prod_{i=1}^{d}f_{j}^{(\theta_{i})}(x_{i}), \tag{69}\]

where the output feature size of each body network is \(mr\).
### (3+1)-d Navier-Stokes Equation

The vorticity form of (3+1)-d Navier-Stokes equation is given as:

\[\partial_{t}\omega+(u\cdot\nabla)\omega=(\omega\cdot\nabla)u+\nu \Delta\omega+F, x\in\Omega,t\in\Gamma, \tag{70}\] \[\nabla\cdot u=0, x\in\Omega,t\in\Gamma,\] (71) \[\omega(x,0)=\omega_{0}(x), x\in\Omega. \tag{72}\]

We constructed the spatial/temporal domain to be \(\Omega=[0,2\pi]^{3}\) and \(\Gamma=[0,5]\), respectively. We constructed the analytic solution for (3+1)-d Navier-Stokes equation introduced in Taylor et al. [12]. The manufactured velocity and vorticity are

\[u_{x} =2e^{-9\nu t}\cos(2x)\sin(2y)\sin(z), \tag{73}\] \[u_{y} =-e^{-9\nu t}\sin(2x)\cos(2y)\sin(z),\] (74) \[u_{z} =-2e^{-9\nu t}\sin(2x)\sin(2y)\cos(z),\] (75) \[\omega_{x} =-3e^{-9\nu t}\sin(2x)\cos(2y)\cos(z),\] (76) \[\omega_{y} =6e^{-9\nu t}\cos(2x)\sin(2y)\cos(z),\] (77) \[\omega_{z} =-6e^{-9\nu t}\cos(2x)\cos(2y)\sin(z), \tag{78}\]

where we chose the viscosity to be \(\nu=0.05\). Each forcing term \(F\) in the Eq. 70 is then given as

\[F_{x} =-6e^{-18\nu t}\sin(4y)\sin(2z), \tag{79}\] \[F_{y} =-6e^{-18\nu t}\sin(4x)\sin(2z),\] (80) \[F_{z} =6e^{-18\nu t}\sin(4x)\sin(4y). \tag{81}\]

We constructed SPINN to be four body networks (modified MLP) of 5 hidden layers with 64/384 hidden feature/output sizes each. We used Adam [6] optimizer with a learning rate of 0.001 and trained for 50,000 iterations. The weight factors in the loss function in Eq. 48 are chosen as \(\lambda_{\text{pde}}=1\), \(\lambda_{\text{ic}}=10\), and \(\lambda_{\text{bc}}=1\). We also weighted the incompressibility loss (Eq. 71) with 100. The visualized solution vector field is shown in Fig. 4.

## Appendix E Additional Experiments

### (5+1)-d Heat Equation

We tested our model on (5+1)-d heat equation to verify the effectiveness of our model on higher dimensional PDE:

\[\frac{\partial u(t,x)}{\partial t}=\Delta u(t,x),\qquad\qquad x\in[-1,1]^{5}, t\in[0,1], \tag{82}\]

where the manufactured solution is chosen to be \(\|x\|^{2}+10t\). When trained with \(8^{6}\) collocation points, SPINN achieved a relative \(L_{2}\) error of 0.0074 within 2 minutes.

### Fine Tuning with L-BFGS

We also conducted some experiments to explore the use of L-BFGS when training SPINN. We found that training with Adam first and then fine-tuning with L-BFGS showed a slight increase in accuracy. Note that this training strategy is used by other works [5, 10] and is known to be effective in some cases. Tab. 2 shows the numerical results on three 3-d PDEs. Understanding the effect of the optimization algorithm is still an open question in PINNs, we believe that investigating this issue in the context of SPINN would be a valuable direction for future study.
\begin{table}
\begin{tabular}{c|c c c} \hline \hline  & Diffusion & Helmholtz & Klein-Gordon \\ \hline Adam & 0.0041 & 0.0360 & 0.0013 \\ Adam + L-BFGS & 0.0041 & 0.0308 & 0.0010 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Numerical result of 3-d PDEs with L-BFGS fine tuning. The number of training collocation points is \(64^{3}\) and a single outer loop L-BFGS is applied.
\begin{table}
\begin{tabular}{c|c c c c} \hline \hline \multirow{2}{*}{model} & \multirow{2}{*}{\(N_{c}\)} & relative & runtime & memory \\  & & \(L_{2}\) error & (ms/iter.) & (MB) \\ \hline \multirow{3}{*}{PINN} & \(16^{3}\) & 0.0193 & 4.70 & 2,810 \\  & \(32^{3}\) & 0.0281 & 14.95 & 2,938 \\  & \(64^{3}\) & 0.0299 & 112.00 & 18,118 \\ \hline \multirow{3}{*}{PINN +} & \(16^{3}\) & 0.0158 & 17.87 & 7,036 \\  & \(32^{3}\) & 0.0185 & 34.61 & 9,082 \\  & \(54^{3}\) & 0.0163 & 159.20 & 22,246 \\ \hline \multirow{3}{*}{SPINN} & \(16^{3}\) & 0.0193 & 1.55 & 762 \\  & \(32^{3}\) & 0.0060 & 1.71 & 762 \\  & \(64^{3}\) & 0.0045 & 1.82 & 762 \\  & \(128^{3}\) & 0.0040 & 1.85 & 890 \\  & \(256^{3}\) & 0.0039 & 3.98 & 1,658 \\ \hline \multirow{3}{*}{SPINN +} & \(16^{3}\) & 0.0062 & 2.20 & 764 \\  & \(32^{3}\) & 0.0020 & 2.41 & 764 \\  & \(64^{3}\) & 0.0013 & 2.57 & 764 \\ \cline{1-1}  & \(128^{3}\) & **0.0008** & 2.79 & 892 \\ \cline{1-1}  & \(256^{3}\) & 0.0009 & 5.61 & 1,660 \\ \hline \hline \end{tabular}
\end{table}
Table 6: Full results of the **(3+1)-d Klein-Gordon equation**. \(N_{c}\) is the number of collocation points.

\begin{table}
\begin{tabular}{c|c c c c} \hline \hline \multirow{2}{*}{model} & \multirow{2}{*}{\(N_{c}\)} & relative & runtime & memory \\  & & \(L_{2}\) error & (ms/iter.) & (MB) \\ \hline \multirow{3}{*}{PINN} & \(16^{3}\) & 0.0919 & 4.84 & 2,810 \\  & \(32^{3}\) & 0.9757 & 14.84 & 2,938 \\  & \(64^{3}\) & 0.9723 & 110.23 & 18,118 \\ \hline \multirow{3}{*}{\begin{tabular}{c} PINN + \\ modified MLP \\ \end{tabular} } & \(16^{3}\) & 0.4770 & 18.32 & 7,034 \\  & \(32^{3}\) & 0.5176 & 35.02 & 9,082 \\  & \(54^{3}\) & 0.4770 & 159.90 & 22,244 \\ \hline \multirow{3}{*}{SPINN} & \(16^{3}\) & 0.1177 & 1.54 & 762 \\  & \(32^{3}\) & 0.0809 & 1.71 & 762 \\  & \(64^{3}\) & 0.0592 & 1.85 & 762 \\  & \(128^{3}\) & 0.0449 & 1.89 & 762 \\  & \(256^{3}\) & 0.0435 & 3.84 & 1,146 \\ \hline \multirow{3}{*}{
\begin{tabular}{c} SPINN \\ + \\ modified MLP \\ \end{tabular} } & \(16^{3}\) & 0.1161 & 2.24 & 764 \\  & \(32^{3}\) & 0.0595 & 2.50 & 764 \\  & \(64^{3}\) & 0.0360 & 2.57 & 764 \\ \cline{1-1}  & \(128^{3}\) & **0.0300** & 2.76 & 764 \\ \cline{1-1}  & \(256^{3}\) & 0.0311 & 5.50 & 1,148 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Full results of the **Helmholtz equation**. \(N_{c}\) is the number of collocation points.

\begin{table}
\begin{tabular}{c|c c c c} \hline \hline \multirow{2}{*}{model} & \multirow{2}{*}{\(N_{c}\)} & relative & runtime & memory \\  & & \(L_{2}\) error & (ms/iter.) & (MB) \\ \hline \multirow{3}{*}{PINN} & \(16^{3}\) & 0.0095 & 3.98 & 1,022 \\  & \(32^{3}\) & 0.0082 & 12.82 & 2,942 \\  & \(64^{3}\) & 0.0081 & 95.22 & 18,122 \\ \hline \multirow{3}{*}{\begin{tabular}{c} PINN \\ + \\ modified MLP \\ \end{tabular} } & \(16^{3}\) & 0.0048 & 14.94 & 1,918 \\  & \(52^{3}\) & 0.0043 & 29.91 & 4,990 \\  & \(54^{3}\) & 0.0041 & 134.64 & 22,248 \\ \hline \multirow{3}{*}{SPINN} & \(16^{3}\) & 0.0447 & 1.45 & 766 \\  & \(32^{3}\) & 0.0115 & 1.76 & 766 \\  & \(64^{3}\) & 0.0075 & 1.90 & 766 \\  & \(128^{3}\) & 0.0061 & 2.09 & 894 \\  & \(256^{3}\) & 0.0061 & 10.54 & 2,174 \\ \hline \multirow{3}{*}{
\begin{tabular}{c} SPINN \\ + \\ modified MLP \\ \end{tabular} } & \(16^{3}\) & 0.0390 & 2.17 & 766 \\  & \(32^{3}\) & 0.0067 & 2.44 & 768 \\  & \(64^{3}\) & 0.0041 & 2.59 & 768 \\ \cline{1-1}  & \(128^{3}\) & **0.0036** & 3.06 & 896 \\ \cline{1-1}  & \(256^{3}\) & **0.0036** & 12.13 & 2,176 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Full results of **diffusion equation**. \(N_{c}\) is the number of collocation points.

\begin{table}
\begin{tabular}{c|c c c c} \hline \hline \multirow{2}{*}{model} & \multirow{2}{*}{\(N_{c}\)} & relative & runtime & memory \\  & & \(L_{2}\) error & (ms/iter.) & (MB) \\ \hline \multirow{3}{*}{PINN} & \(16^{3}\) & 0.0343 & 4.70 & 2,810 \\  & \(32^{3}\) & 0.0281 & 14.95 & 2,938 \\  & \(64^{3}\) & 0.0299 & 112.00 & 18,118 \\ \hline \multirow{3}{*}{\begin{tabular}{c} PINN \\ + \\ modified MLP \\ \end{tabular} } & \(16^{3}\) & 0.0158 & 17.87 & 7,036 \\  & \(54^{3}\) & 0.0185 & 34.61 & 9,082 \\  & \(54^{3}\) & 0.0163 & 159.20 & 22,246 \\ \hline \multirow{3}{*}{\begin{tabular}{c} SPINN \\ + \\ modified MLP \\ \end{tabular} } & \(16^{3}\) & 0.0193 & 1.55 & 762 \\  & \(32^{3}\) & 0.0060 & 1.71 & 762 \\  & \(64^{3}\) & 0.0045 & 1.82 & 762 \\  & \(128^{3}\) & 0.0040 & 1.85 & 890 \\  & \(256^{3}\) & 0.0039 & 3.98 & 1,658 \\ \hline \multirow{3}{*}{
\begin{tabular}{c} SPINN \\ + \\ modified MLP \\ \end{tabular} } & \(16^{3}\) & 0.0062 & 2.20 & 764 \\  & \(32^{3}\) & 0.0020 & 2.41 & 764 \\  & \(64^{3}\) & 0.0013 & 2.57 & 764 \\ \cline{1-1}  & \(128^{3}\) & **0.0008** & 2.79 & 892 \\ \cline{1-
Figure 1: Visualized solution of **nonlinear diffusion equation** obtained by the baseline PINN and SPINN, both trained on \(64^{3}\) collocation points.
Figure 3: Visualized vorticity maps of **(2+1)-d Navier-Stokes equation** experiment predicted by SPINN. Three snapshots at timestamps \(t=0.0\), \(0.5\), \(1.0\) are presented.

Figure 2: Visualized solution of **Helmholtz equation** obtained by the baseline PINN and SPINN, both trained on \(64^{3}\) collocation points.
Figure 4: Visualized solution of **(3+1)-d Navier-Stokes equation** obtained by SPINN, trained on \(32^{4}\) collocation points. Each arrow is colored by its zenith angle.