Title: Separable Physics-Informed Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: partial differential equations

scientific machine learning

physics-informed neural networks

fluid dynamics
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 very complex solution functions.
The number of training points (collocation points) required on these challenging PDEs grows substantially, and it is severely limited due to the expensive computational costs and heavy memory overhead.
To overcome this limit, we propose a network architecture and training algorithm for PINNs.
The proposed method, separable PINN (SPINN), operates on a per-axis basis to decrease the number of network propagations in multi-dimensional PDEs instead of 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 significantly 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 much 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/.

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 × in wall-clock time, 1 , 394 × 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/ .

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 [10]. Hence, numerical approximation methods to solutions have been extensively studied [17], e.g., spectral methods [3], finite volume methods (FVM) [9], finite difference method (FDM) [39], and finite element methods (FEM) [47]. 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 [4, 24, 30, 38]. Among the promising methods, physics-informed neural networks (PINNs) have recently emerged as new datadriven PDE solvers for both forward and inverse problems [34]. PINNs employ neural networks and

Reference vorticity

(

𝑡 = 0.1)

Predicted

vorticity

Ours at 5 sec

Ours at 1 mins

Ours at 5 mins

Absolute

error

4.26e-3

1.48e-2

8.50e-2

Relative

𝐿2

error

Training time

4.31e-3

5.89e-3

2 hours later…

…

4.01e-3

20

SPINN (ours)

PINN

10

0

relative L2 error

-10

10-1

-20

6

10-2

4

2

0.0

0.1

2.0 2.5 3.0

training time(hours)

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, 36, 37, 43]. 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) [12, 27, 40]. 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, O ( N d ) →O ( Nd ) , where N is the resolution of the solution for each dimension, and d is the dimension of the system.

We have conducted comprehensive experiments on the representative PDEs to show the effectiveness of the suggested method. The experimental results demonstrate that the training runtime 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 × on commodity GPUs and FLOPs up to by 1 , 394 × 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 [43] (9 minutes vs 10 hours in a single GPU) without bells and whistles, such as causal

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, 43, 44], or to accelerate the training speed [20, 37]. 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 [7, 23, 36, 43]. Especially, Sankaran et al. [36] 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 scientific machine learning. Embedding physical constraints in the loss function is widely used in scientific machine learning, and computing the derivatives is an essential process for this formulation. Several works employed numerical differentiation [33, 35, 42], or hybrid approach [5, 37, 46] with AD for the calculation, since numerical methods such as finite differences do not need to back-propagate through the network. However, they are still burdened by a computational complexity of O ( N d ) , thereby limiting them to handle large-scale collocation points or meshes. Furthermore, numerical differentiation has truncation errors depending on the step size. Employing Taylor-mode AD [13] in training PINNs was introduced by causal PINN [43] 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, 32]. They share the same concept of dividing the entire 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 sub-domain, 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 [31]. 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 [31] 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 made their model more computationally expensive than ours. Furthermore, in neural fields [45], 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 devise the architecture to exploit forward-mode AD to efficiently compute PDE residual losses.

…

𝑥1

𝑥2

𝑥3

𝑦

…

𝑦1

𝑦2

𝑥

(a)

𝑓 (𝜃)

: ℝ → ℝ 3

(b)

𝑓 (𝜃) : ℝ 3 →ℝ

Merge

(a) Separated

𝑥1

𝑥2

𝑥3

𝑁𝑑

points

𝑁 𝑑

points

𝑁 𝑑

points

(b) Non-separated

𝑥1

𝑥2

𝑥3

𝑁 𝑑

points

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 [2] and [13]. Suppose our function f : R n → 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 × n Jacobian matrix J f in forward-mode, we compute the Jacobian-vector product (JVP),

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, ˙ v k the tangents, ¯ v k the adjoints, and W 1 , W 2 the weight matrices. Biases are omitted for brevity.

for i ∈ { 1 , . . . , 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 = ∂v k /∂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 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):

for j ∈ { 1 , . . . , 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 ¯ v k = ∂y j /∂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. [2] for more details.

(a) Architecture (single input)

Element-wise Product

𝑥1 ∈ ℝ

𝑓 (𝜃 1 )

𝑥2 ∈ ℝ

𝑥3 ∈ ℝ

ො 𝑢(𝑥) = ෍

𝑗=1

𝑟

𝑓

𝑗

(𝜃1 )

𝑥1

𝑓

𝑗

(𝜃 2 )

𝑥2

𝑓

𝑗

(𝜃 3 )

𝑥3

𝑓

1

(𝜃 1 )

…

∘

∑

+

+

+

…

=

𝑓 (𝜃 1 )

…

𝑋 ∈ ℝ 𝑁×3

𝐹:,1,1

𝐹:,2,1

Batchified

Input

𝑋:,3

𝐹:,𝑟,1

…

𝐹:,1,2

𝐹:,2,2

𝐹:,𝑟,2

…

𝐹:,1,3

𝐹:,2,3

𝐹:,𝑟,3

𝑓 (𝜃 2 )

𝑓 (𝜃 3 )

𝑓

2

(𝜃 1 )

𝑓 𝑟

(𝜃 1 )

𝑓

1

(𝜃 2 )

…

𝑓

2

(𝜃 2 )

𝑓 𝑟

(𝜃 2 )

𝑓

1

(𝜃 3 )

…

𝑓

2

(𝜃 3 )

𝑓 𝑟

(𝜃 3 )

𝑓 (𝜃 2 )

𝑓 (𝜃 3 )

𝐹:,1,1

𝐹:,1,2

𝐹:,1,3

𝑋:,1

𝑋:,2

𝐹 ∈ ℝ 𝑁×𝑟×3

𝐹:,2,2

𝐹:,𝑟,2

𝐹:,2,1

𝐹:,2,3

𝐹:,𝑟,3

𝐹:,𝑟,1

Mode-1 (column) fibers of

𝐹

4 Separable PINN

4.1 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. 3a), 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 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 ∼ 3 times more expensive than the forward propagation [2, 13]. According to Fig. 3 and with the scale constants c f , c h ∈ [2 , 3] we can approximate the total number of operations to compute the Jacobian matrix of the proposed separated approach (Fig. 3a).

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

If we can make ops ( h ) sufficiently small, then the ratio C sep / C non-sep becomes Nd N d ≪ 1 , quickly converging to 0 as d and N increases. Our separated approach has linear complexity with respect to N in network propagations, implying it can obtain more accurate solutions (high-resolution) efficiently.

4.2 Network Architecture

Fig. 4a 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 ( θ i ) : R → R r (parameterized by θ 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:

where ˆ u : R d → R is the predicted solution function, x i ∈ R is a coordinate of i -th axis, and f ( θ i ) j denotes the j -th element of f ( θ i ) . We used 'tanh' activation function throughout the paper. As shown in Eq. 5, the feature merging operation is a simple product ( Π ) and summation ( Σ ) which

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., ops ( h ) ≪ ops ( f ) in Eq. 3). Note that SPINN can also approximate any m -dimensional vector functions ˆ u : R d → 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, 1dimensional 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, nonfactorizable 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.

(a) Factorizable coordinates (b) Non-factorizable coordinates

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 ∈ R N × d is now a matrix. The batchified form of feature representation F ∈ R N × r × d and Eq. 5 now becomes

where ˆ U ∈ R N × N × ... × N is the discretized solution tensor, ⊗ denotes outer product, F : , : ,i ∈ R N × r is an i -th frontal slice matrix of tensor F , and F : ,j,i ∈ 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.

4.3 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 ( ∇ ˆ u ( x ) ) is Nd , where N is the number of coordinates sampled from each axis and d is the input dimension. According to Eq. 5, the i -th element of ∇ ˆ u ( x ) is:

Computing this derivative requires feature representations f ( θ i ) j , which can be reused from the forward pass results computed beforehand. The entire r components of the feature derivatives ∂f ( θ i ) ( x i ) /∂x i : R → R r can be obtained by a single pass thanks to forward-mode AD. To obtain the full gradient ∇ ˆ u ( x ) : R d → 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. Note that any p -th order derivative ∂ p ˆ u/∂x p i can also be obtained in the same way.

(a) Diffusion

(b) Helmholtz

(c) (2+1)-d Klein-Gordon

(d) (3+1)-d Klein-Gordon

SPINN

SPINN-mod

PINN

PINN-mod

rel. L2 error

102

(GB)

101

100

rel. L2 error

100

102

memory(GB)

0

rel. L2 error

0

memory

00

10-2

101

-0T

101

100

-O

***

-0

163 323 64312832563

1633236431282563

1633236431282563

020406080

16332364128256

163323 64312832563

16332364312832563

0

20406080

# of col.pts

# of col.pts

# of col.pts

runtime (sec)

# of col.pts

#of col.pts

# of col.pts

runtime(sec)

rel.L2 error

10~2

memory (GB)

101

00

rel.L2error

10-2

)

102

memory(GB)

101

101

rel.

2-01

101

-0

16332364312832563

# of col.pts

16332364312832563

# of col.pts

100

16332364312832563

# of col.pts

020406080

runtime (sec)

164

#of col.pts

324

644

164

# of col.pts

324

644

164

# of col.pts

324

644

0

runtime (sec)

306090120

4.4 Universal Approximation Property

It is widely known that neural networks with sufficiently many hidden units have the expressive power of approximating any continuous functions [6, 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 R d . Choose u ∈ L 2 ( X × Y ) . Then, for arbitrary ε > 0 , we can find a sufficiently large r > 0 and neural networks f j and g j such that

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

5.1 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 [44] 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 ∥ ˆ u -u ∥ 2 / ∥ u ∥ 2 , where ˆ 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.

5.2 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 × 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 Equation When 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 × faster and 29 × 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 Equation The results of the Helmholtz experiments are shown in Fig. 6b. 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 [44]. 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 × faster, and the memory usage is 29 × smaller than the baseline with modified MLP. The exact numerical values are provided in the appendix.

Klein-Gordon Equation Fig. 6c shows the results. Again, our method shows the best performance in terms of accuracy, runtime ( 62 × ), and memory usage ( 29 × ). Note that both Helmholtz and KleinGordon 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. 6d, 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 × larger than the baseline) to obtain a more precise solution. The exact numerical values are provided in the appendix.

(2+1)-d Navier-Stokes Equation We constructed SPINN to predict the velocity field u and applied forward-mode AD to obtain the vorticity ω . We used the same PDE setting used in causal PINN [43] 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 ϵ in the causal loss function, they had to run 3 ∼ 4 times more training epochs, depending on the stopping criterion. As a result, given the same number of collocation points, SPINN converges 60 × 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

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 Equation We proceeded to explore our model on (3+1)-d Navier-Stokes equation by devising a manufactured solution corresponding to the Taylor-Green vortex [41]. It models a decaying vortex, widely used for testing Navier-Stokes numerical solvers [8]. Since the vorticity of the (3+1)-d NavierStokes 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. Handling any geometric surface is one of the key advantages of PINNs. In the appendix, we presented a simple way to handle arbitrary domains with SPINN by testing on the Poisson equation in an L-shaped domain. Please see section E.3 for more details. Another method could be applying an additional operation after the SPINN's feature merging to map a rectangular mesh to an arbitrary physical mesh, inspired by PhyGeoNet [11] and Geo-FNO [29]. Combining such techniques for handling arbitrary boundary conditions of complex geometries into ours is an exciting research direction, and we leave it to future work.

We also found that due to its architectural characteristic, SPINN tends to be better trained when the solutions align with the variable separation form (discretized into low-rank tensor). However, note that SPINN is still effective in solving equations where the solution is not a low-rank tensor, such as our examples on diffusion and (2+1)-d Navier-Stokes equation. Please see section D in the appendix for additional experiments and information of which example aligns with the variable separation form or not. 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 in PINNs. We believe this work opens up a new direction to rethink the architectural design of neural networks in many scientific applications.

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 (2022R1A4A3033571), Institute of Information and Communication Technology Planning Evaluation (IITP) grants for the AI Graduate School program (IITP-2019-0-00421), Korea Institute of Energy Technology Evaluation and Planning (KETEP) and the Ministry of Trade, Industry & Energy (MOTIE) of the Republic of Korea (No. 20224000000360).

References

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.

Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of Marchine Learning Research , pages 1-43, 2018.

John P Boyd. Chebyshev and Fourier spectral methods . Courier Corporation, 2001.

Steven L Brunton and J Nathan Kutz. Data-driven science and engineering: Machine learning, dynamical systems, and control . Cambridge University Press, 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.

George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems , 2(4):303-314, 1989.

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.

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.

Robert Eymard, Thierry Gallouët, and Raphaele Herbin. Finite volume methods: handbook of numerical analysis. PG Ciarlet and JL Lions (Eds) , 2000.

Charles L Fefferman. Existence and smoothness of the navier-stokes equation. The millennium prize problems , 57:67, 2000.

Han Gao, Luning Sun, and Jian-Xun Wang. Phygeonet: Physics-informed geometry-adaptive convolutional neural networks for solving parameterized steady-state pdes on irregular domain. Journal of Computational Physics , 428:110079, 2021.

Priya Goyal, Piotr Dollár, 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.

Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation . SIAM, 2008.

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.

John Hammersley. Monte carlo methods . Springer Science & Business Media, 2013.

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.

Joe D Hoffman and Steven Frankel. Numerical methods for engineers and scientists . CRC press, 2018.

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks , 2(5):359-366, 1989.

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.

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.

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.

Jian-hui Jiang, Hai-long Wu, Yang Li, and Ru-qin Yu. Three-way data resolution by alternating slicewise diagonalization (asd) method. Journal of Chemometrics: A Journal of the Chemometrics Society , 14(1):15-36, 2000.

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.

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.

Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review , 51(3):455-500, 2009.

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.

Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv:1404.5997 , 2014.

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.

Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint arXiv:2207.05209 , 2022.

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.

Ruofan Liang, Hongyi Sun, and Nandita Vijaykumar. Coordx: Accelerating implicit neural representation with a split mlp architecture. In International Conference on Learning Representations , 2022.

Ben Moseley, Andrew Markham, and Tarje Nissen-Meyer. Finite basis physics-informed neural networks (fbpinns): a scalable domain decomposition approach for solving differential equations. arXiv preprint arXiv:2107.07871 , 2021.

Guofei Pang, Lu Lu, and George Em Karniadakis. fpinns: Fractional physics-informed neural networks. SIAM Journal on Scientific Computing , 41(4):A2603-A2626, 2019.

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.

Rishikesh Ranade, Chris Hill, and Jay Pathak. Discretizationnet: A machine-learning based solver for navier-stokes equations using finite volume discretization. Computer Methods in Applied Mechanics and Engineering , 378:113722, 2021.

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.

Ramansh Sharma and Varun Shankar. Accelerated training of physics informed neural networks (pinns) using meshless discretizations. Advances in Neural Information Processing Systems , 2022.

Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics , 375:1339-1364, 2018.

Gordon Dennis Smith. Numerical solution of partial differential equations: finite difference methods . Oxford university press, 1985.

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.

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.

Nils Wandel, Michael Weinmann, Michael Neidlin, and Reinhard Klein. Spline-pinn: Approaching pdes without data using fast, physics-informed hermite-spline cnns. In Proceedings of the AAAI Conference on Artificial Intelligence , pages 8529-8538, 2022.

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.

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.

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.

Yinhao Zhu, Nicholas Zabaras, Phaedon-Stelios Koutsourelakis, and Paris Perdikaris. Physics-constrained deep learning for high-dimensional surrogate modeling and uncertainty quantification without labeled data. Journal of Computational Physics , 394:56-81, 2019.

Olek C Zienkiewicz, Robert Leroy Taylor, and Jian Z Zhu. The finite element method: its basis and fundamentals . Elsevier, 2005.

Separable PINN - Supplementary Materials

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 β } be an orthonormal basis for H 2 . Tensor product between Hilbert spaces H 1 and H 2 , denoted by H 1 ⊗ H 2 , is a set of all antilinear mappings A : H 2 → H 1 such that ∑ β ∥ Av β ∥ 2 < ∞ for every orthonormal basis for H 2 .

Then by Theorem 7.12 in Folland [4], H 1 ⊗H 2 is also a Hilbert space with respect to norm | | |·| | | and associated inner product ⟨· , ·⟩ :

where A,B ∈ H 1 ⊗H 2 , and { v β } is any orthonormal basis of H 2 .

Lemma 1. Let x ∈ H 1 and y, y ′ ∈ H 2 . Then, ( x ⊗ y ) y ′ = ⟨ y, y ′ ⟩ x .

Proof. Let A be a mapping A : y ′ →⟨ y, y ′ ⟩ x , where ∥ y ∥ = 1 . We expand y ′ with orthonormal basis { y, z β } . i.e., { y, z β } is a basis for H 2 . Then,

It is obvious that A is antilinear. Then by Definition 1, A ∈ H 1 ⊗H 2 which is, Ay ′ = ( x ⊗ y ) y ′ . Therefore, ( x ⊗ y ) y ′ = ⟨ y, y ′ ⟩ x .

Lemma 2. { u α ⊗ v β } is an orthonormal basis for H 1 ⊗H 2 .

Proof. Let A,B ∈ H 1 ⊗H 2 , where B = u α ⊗ v β . Then by the definition of inner product in Eq. 2,

=

∑

i

Now we check the Parseval identity:

⟨

Av

,

(

u

⊗

v

)

v

⟩

(6)

i

α

β

i

∴ { u α ⊗ v β } is a basis.

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

Theorem 1. Let X,Y be compact subsets of R d . Choose u ∈ L 2 ( X × Y ) . Then, for arbitrary ε > 0 , we can find a sufficiently large r > 0 and neural networks f j and g j such that

Proof. Let { ϕ i } and { ψ j } be orthonormal basis for L 2 ( X ) and L 2 ( Y ) respectively. Then { ϕ i ψ j } forms an orthonormal basis for L 2 ( X × Y ) ( ∵ Lemma 2). Therefore, we can find a sufficiently large r such that

On the other hand, by the universal approximation theorem [2], we can find neural networks f j and g j such that

We first consider the difference between u and ∑ r i,j a ij f i g j :

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

We first compute II 1 :

We set where a ij denotes

to write II 1 as

Now

Since { ϕ i } is an orthonormal basis, we see that

Therefore,

Finally, we recall

to conclude

On the other hand, we have from (15)

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

Finally, we reorder the index to rewrite

Without loss of generality, we rewrite ˜ r, k i , ˜ g i into r, f i , g i respectively, to complete the proof.

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 [15], except we use forwardmode 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 ˆ u ( θ ) ( 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:

Hence we have

Likewise,

Therefore,

where Ω is an input domain, N , B are generic differential operators and u ic , u bc are initial, boundary conditions, respectively. λ are weighting factors for each loss term. When calculating the PDE loss ( L pde) with Monte-Carlo integral approximation, we sampled collocation points from factorized coordinates and used forward-mode AD. The remaining L ic and L 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.

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 [6] 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 ( θ i ) , which makes the complexity of network propagations from O ( Nd ) to 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 ≪ N d ), there is a room to bridge the gap between theoretical FLOPs and actual training runtime by further software optimization.

D Experimental Details and Results

In this section, we provide experimental details, numerical results, and visualizations for each experiment in the main paper. Below, we show the list of examples where the solution aligns with the variable separation form or not.

separable form : 3-d Helmholtz, (2+1)-d Klein-Gordon, (3+1)-d Klein-Gordon, (3+1)-d NavierStokes, (5+1)-d diffusion non-separable form : (2+1)-d diffusion, (2+1)-d Navier-Stokes, (2+1)-d flow mixing, 2-d L-shaped Poisson

D.1 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:

We used diffusivity α = 0 . 05 , spatial domain Ω = [ -1 , 1] 2 , temporal domain Γ = [0 , 1] and used superposition of three Gaussian functions for the initial condition u ic. We obtained the reference solution ( 101 × 101 × 101 resolution) through a widely-used PDE solver platform FEniCS [12]. Note that FEniCS is a FEM-based solver. We particularly set the initial condition to be a superposition of three gaussian functions:

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 [8] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors λ 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. 3.

D.2 Helmholtz Equation

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

where the spatial domain is Ω = [ -1 , 1] 3 . For a given source term q = -( a 1 π ) 2 u -( a 2 π ) 2 u -( a 3 π ) 2 u + k 2 u , we devised a manufactured solution u = sin( a 1 πx 1 ) sin( a 2 πx 2 ) sin( a 3 π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 [8] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors λ 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. 4.

D.3 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

where we chose the spatial/temporal domain to be Ω = [ -1 , 1] 2 and Γ = [0 , 10] , respectively. For error measurement, we used a manufactured solution u = ( x 1 + x 2 ) cos(2 t ) + x 1 x 2 sin(2 t ) 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 [8] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors λ 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:

where f , u bc are extracted from this exact solution. Tab. 6 shows the numerical results. The number of collocation points of 23 4 , 18 4 were the maximum value for PINN and PINN with modified MLP, respectively.

D.4 (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:

where u ∈ R 2 is the velocity field, ω = ∇× u is the vorticity, ω 0 is the initial vorticity, and ν is the viscosity. We used the viscosity 0.01 and made the spatial/temporal domain Ω = [0 , 2 π ] 2 and Γ = [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 [9] 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 × 128 × 128 ( N t × N x × N y ) , 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 [10, 18]. We used Adam [8] 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.

where u x , u y are x, y components of predicted velocity, w = ∂ x u y -∂ y u x , N c is the number of collocation points, N ic is the number of coordinates for initial condition and u x 0 , u y 0 , u w 0 are the initial conditions. We chose the weighting factors λ w = 1 , λ c = 5 , 000 , and λ ic = 10 , 000 . We also resampled the input points every 100 epochs. The periodic boundary condition can be explicitly enforced by positional encoding [3], and we specifically used the following encoding function only for the spatial input coordinates.

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

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

where the output feature size of each body network is mr .

D.5 (3+1)-d Navier-Stokes Equation

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

We constructed the spatial/temporal domain to be Ω = [0 , 2 π ] 3 and Γ = [0 , 5] , respectively. We constructed the analytic solution for (3+1)-d Navier-Stokes equation introduced in Taylor et al. [17]. The manufactured velocity and vorticity are

where we chose the viscosity to be ν = 0 . 05 . Each forcing term F in the Eq. 70 is then given as

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 [8] 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 λ pde = 1 , λ ic = 10 , and λ bc = 1 . We also weighted the incompressibility loss (Eq. 71) with 100. The visualized solution vector field is shown in Fig. 6.

E Additional Experiments

E.1 (5+1)-d diffusion Equation

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

where the manufactured solution is chosen to be ∥ x ∥ 2 +10 t . When trained with 8 6 collocation points, SPINN achieved a relative L 2 error of 0.0074 within 2 minutes.

E.2 (2+1)-d Flow Mixing Problem

This is a time dependent PDE which describes the behavior of two fluids being mixed at the interface in a 2-d environment. Following the settings in CAN-PINNs [1], the equation is

where

It has an analytic solution [16]:

where ω = 1 r v t v t, max . The initial and boundary conditions are extracted from the exact solution Eq. 89. When trained with 256 3 collocation points, SPINN achieved a relative L 2 error of 0.0029 in 6 minutes. Figure 1 shows the visualized result.

t=0

t=2

t=4

Exact solution

Predicted solution

E.3 2-d Poisson equation on an L-shaped domain

Besides the coordinate or mesh transformations [5, 11] mentioned in the limitations section 6, we would like to present yet another simpler way to handle arbitrary domain shapes with SPINN. For boundary points, we train in a non-separable way, just as normal PINN training (non-factorizable coordinates). For collocation points inside the domain, we train in a separable way (factorizable coordinates) and ignore the points outside the domain. This can be easily achieved by finding a tight bounding box of the input domain and masking out the PDE loss when the coordinates are outside.

To show the effectiveness, we tested SPINN on the Poisson equation on an L-shaped domain following the settings used in DeepXDE [13]:

where Ω = [ -1 , 1] 2 \ [0 , 1] 2 . The reference solution was obtained by the spectral method. SPINN achieved a relative L 2 error of 0.0322, while PINNs with modified MLP achieved 0.0392. Figure 2 shows the visualized result. For practical reasons, we split the L-shaped domain into two rectangular domains instead of masking out the top-right part of the domain. The L-shaped domain we showed is just an example case, and we believe that this method can be applied to any arbitrary input domain.

Referencesolution

Prediction(SPINN)

Prediction(PINN)

1.00

1.00

1.00

0.75

0.12

0.75

0.12

0.75

0.12

0.50

0.10

0.50

0.10

0.50

0.10

0.25

0.08

0.25

0.08

0.25

0.08

0.00

0.06

0.00

0.06

0.00

0.06

0.25

0.25

0.04

0.25

0.04

0.04

0.50

0.50

0.50

0.02

0.02

0.75

0.02

0.75

0.75

0.00

0.00

1.00

0.00

1.00

1.00

-1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

0.5

1.0

E.4 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 [7, 14] 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.

0s

PINN

SPINN (Ours)

(c) Training time: 1 min

(a) Training time: 5 sec

PINN

SPINN (Ours)

PINN

SPINN (Ours)

(b) Training time: 30 sec

Reference

t

= 0.0

t

= 0.25

t

= 0.5

t

= 1.0

0.0275

0.30

0.35

0.0250

0.0225

025

0.0200

0.0175

0.0150

0.0

0.0125

1.0

0.5

00

0.5

1.0

0.5

00

1.0

0.5

00

0.5

1.0

1.0

0.5

00

0.5

1.0

1.0

0.5

0.0

0.5

1.0

1.0

0.5

00

05

10

1.0

0.5

0.0

05

1.0

0.5

00

05

10

0.0100

0.0075

0.20

0.15

0.05

0.10

1.0

1.0

0.00

1.0

10

0.05

0.5

0.5

0.5

0.5

1.0

0.5

0.0

0.5

0.5

0.0

1.0

0.5

0.0

0.5

10 -1.0

0.0

1.0

-0.5

05

10-1.0

0.5

0.0

1.0

0.5

0.0

0.5

101.0

-0.5

00

0.00

0.150

0.125

0.100

0.075

0.050

0.00

0.025

0.5

10

0.5

1.0

1.0

0.5

10

0.000

-1.0

=0.5

00

0.5

1.0

-1.0

0.5

0.0

-1.0

0.5

00

0.5

1.0

1.0

0.5

0.0

1.0

0.5

00

0.5

1.0

0.5

0.0

1.0

0.5

0.0

05

10

1.0

0.5

00

0.025

#

0.20

0.15

0.10

10

10

10

10

0.05

0.5

0.5

0.5

05

1.0

0.5

0.0

0.5

10-1.0

0.5

0.0

1.0

0.5

00

0.5

10 1.0

0.5

00

1.0

-0.5

0.0

05

10-1.0

0.5

00

1.0

0.5

0.0

05

101.0

-0.5

0.0

0.00

0.35

0.150

0.125

0.100

0.075

0.00

0.050

-1.0

0.5

00

05

10

1.0

0.5

0.0

0.5

1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

Q.5

1.0

1.0

0.5

0.0

0.5

10

1.0

0.5

0.0

0.5

10

0.025

0.000

0.20

0.15

0.10

10

10

1.0

10

0.05

1.0

0.5

0.0

0.5

10 -1.0

0.5

0.0

1.0

0.5

0.0

0.5

10 1.0

0.5

0.5

1.0

0.5

0.0

0.5

10 1.0

-0.5

0.0

0.5

1.0

0.5

0.0

0.5

10 1.0

-0.5

0.0

0.5

0.00

0.20

0.15

0.10

1.0

0.5

10

0.5

10

05

10

0.05

1.0

0.5

00

0.5

10 -1.0

0.5

0.0

1.0

0.5

00

0.5

0.5

0.0

1.0

-0.5

0.0

0.5

0.5

0.0

-1.0

0.5

0.0

0.5

101.0

-0.5

0.0

0.00

SPINN (Ours)

PINN

Absolute error (SPINN)

Absolute error (PINN)

SPINN (Ours)

PINN

Absolute error (SPINN)

Absolute error (PINN)

(a) Training time: 5 sec

(b) Training time: 30 sec

Exact solution

(c) Training time: 1 min

SPINN (Ours)

PINN

Absolute error (SPINN)

Absolute error (PINN)

-4.75

4.0#

-43#

-4.40

18

-448

4.75

-431

12

-.5

15858

-45

-4.99

-65

-.5

-625

-450

-4.99

-1.8

4.万

100

-1.90

0

-1.90

-475

4万

-1.%

100

-1.90

-47

-4.50

42

-1.00

1.99

858585

-0.25

-4.7%

4.9#

1/5

##

4.5

4.48

4.40

4.25

S8585

-25

-0.75

-4.99

1#

-.n

-35

459

196

1.90

-1.99

45

459

999999

-425

45

-1.8

-010

4 02

-1.00

-4

-1.90

-15e

-1.%

-47

-436

42

100

12i

159

-1.90

-0.75

-4.50

-4.90

-4.59

-421

125

950

6

Reference w(o.o, x, y)

AbsoluteError

6

6

6

20

20

0.10

5

5

10

10

0.08

4

4

4

>3

0>3

0>31

0.06

2

-102

-102

0.04

1

-20 1

-201

0.02

0

0

0

10.00

1

2

3

4

5

6

O

1

2

3

4

5

6

2

3

4

5

6

X

Reference w(o.5,x, y)

Predictedw(0.5,x,y)

AbsoluteError

6

20

6

20

3.5

5

5

5

3.0

10

10

2.5

4

4

4

0>3

0>3

2.0

>3

1.5

2

-10 2

-102

1.0

1

1

1

0.5

-20

20

0

0

0

0

1

2

3

4

5

9

0

1

2

3

4

5

9

0

1

3

4

5

X

X

Predicted w(1.0, x, y)

Absolute Error

6

6

6

15

15

5

10

5

10

5

4

4

5

4

5

4

3

>3

0>3

>3

-5

-5

2

2

2

-10

-10

1

-15

-151

0

-200

-200

0

1

2

3

4

5

6

0

1

2

3

4

5

6

0

1

2

3

4

5

6

X

X

X

2

(a) Exact solution (velocity)

(c) Exact solution (vorticity)

(d) Prediction (vorticity)

(b) Prediction (velocity)

=V

[=1

=Z

=3

t=5

=1

t=2

t=3

=5

t=3

=0

=2

t=3

t=5

References

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.

George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems , 2(4):303-314, 1989.

Suchuan Dong and Naxian Ni. A method for representing periodic functions and enforcing exactly periodic boundary conditions with deep neural networks. Journal of Computational Physics , 435:110242, 2021.

Gerald B Folland. A course in abstract harmonic analysis , volume 29. CRC press, 2016.

Han Gao, Luning Sun, and Jian-Xun Wang. Phygeonet: Physics-informed geometry-adaptive convolutional neural networks for solving parameterized steady-state pdes on irregular domain. Journal of Computational Physics , 428:110079, 2021.

Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation . SIAM, 2008.

Xiaowei Jin, Shengze Cai, Hui Li, and George Em Karniadakis. Nsfnets (navier-stokes flow nets): Physics-informed neural networks for the incompressible navier-stokes equations. Journal of Computational Physics , 426:109951, 2021.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning-accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences , 118(21):e2101784118, 2021.

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.

Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint arXiv:2207.05209 , 2022.

Anders Logg, Kent-Andre Mardal, and Garth Wells. Automated solution of differential equations by the finite element method: The FEniCS book , volume 84. Springer Science & Business Media, 2012.

Lu Lu, Xuhui Meng, Zhiping Mao, and George Em Karniadakis. Deepxde: A deep learning library for solving differential equations. SIAM review , 63(1):208-228, 2021.

Xuhui Meng, Zhen Li, Dongkun Zhang, and George Em Karniadakis. Ppinn: Parareal physicsinformed neural network for time-dependent pdes. Computer Methods in Applied Mechanics and Engineering , 370:113250, 2020.

Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: Adeep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics , 378:686-707, 2019.

Panos Tamamidis and Dennis N Assanis. Evaluation of various high-order-accuracy schemes with and without flux limiters. International Journal for Numerical Methods in Fluids , 16(10):931-948, 1993.

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.

Colby L Wight and Jia Zhao. Solving allen-cahn and cahn-hilliard equations using the adaptive physics informed neural networks. arXiv preprint arXiv:2007.04542 , 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel approach to using separable neural networks to represent solutions of partial differential equations (PDEs). In contrast to the conventional form used in physics-informed neural network (PINN) methods, the specialized form of neural networks proposed in this work takes advantage of forward-mode auto-differentiation, leading to faster training times in PINN frameworks. The authors provide a mathematical demonstration of the proposed form's approximation ability and conduct experiments on several PDEs. The results reveal that the proposed method achieves comparable accuracy to previous work, while significantly reducing computational and memory costs.\n\nSTRENGTHS:\n1. The speed-up achieved by the proposed method is highly promising. The time required for training a single network to convergence has been a limitation in the application of PINN methods. This work addresses this issue by introducing a specialized neural network form, and in combination with the forward-mode auto-differentiation method, the training speed is significantly improved.\n2. In addition to the improved auto-differentiation computation, this work employs fixed collocation points, enabling further acceleration of the total training time by reusing calculations on each point. By leveraging this approach, the computational efficiency of the method is significantly enhanced.\n\nWEAKNESSES:\nAlthough the proposed method in this paper is novel and interesting, there are several concerns that the reviewer would like to address:\n\n1. The choice of using a manufactured solution that aligns with the variable separation form for evaluating the results on the Helmholtz Equation, Klein-Gordon Equation, and (3+1)-d Navier-Stokes Equation raises some questions. While this choice is suitable for evaluating original PINN methods that do not utilize the variable separation property of the ground truth, it may be unfair to compare SPINN with the original method on these tasks.\n\n2. One of the key advantages of PINN-based methods is their flexibility to be applied to any geometric surface without modifying the grids. However, the variable separation form of the neural network in SPINN may limit this flexibility. It would be challenging for SPINN to fit the boundary conditions of complex geometries due to the constraints of the variable separation form.\n\n3. Throughout the paper, the evaluation metric predominantly used is the $L_2$ evaluation metric. However, the reliability of the $L_2$ metric, especially for non-linear PDEs, can sometimes be questionable. It would be valuable to also report metrics such as the Sobolev norm or the PINN loss to provide a more comprehensive evaluation.\n\nQUESTIONS:\n1. To assess the performance of SPINN on more complex solutions that do not directly align with its form, the authors should conduct more experiments. The experiments could involve using a Multilayer Perceptron (MLP) in the classical PINN as the manufactured solution and training SPINN under the associated PDE.\n2. The reviewer is curious about the performance of SPINN on PDEs with sophisticated boundary. For instance, it would be interesting to explore how SPINN performs on a sphere boundary condition, where the variables are still separated into x, y, and z components.\n3. Including the Sobolev norm or PINN loss as additional evaluation metrics would provide the reviewer with more information to assess the performance of SPINN comprehensively.\n4. The reviewer wonders if it is possible to reduce the dimension of the problem through variable separation. Most PINN methods work on high-dimension manifolds, incurring significant computational costs. While SPINN employs low-dimensional functions to construct the high-dimensional function, its domain is still a high-dimensional manifold. The reviewer believes that there might be a way to directly optimize these low-dimensional functions within the low-dimensional manifold, potentially offering benefits in solving high-dimensional PDEs.\n\nLIMITATIONS:\nThe authors have talked their limitation in the paper", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel variant of PINN for learning multi-dimensional PDEs. Unlike traditional approaches considering point-wise inputs, this method incorporates each axis information separately during the learning process and leverages tensor decomposition to interpret the predicted solution variables. Additionally, the authors consider a forward-mode auto-differentiation technique that facilitates training with a more significant number of collocation points. The extensive experiments proved the efficiency, scalability, and effectiveness of the proposed SPINN.\n\nSTRENGTHS:\n- This paper contributes to alleviating the scalability and efficiency issues of PINNs. Also, the proposed forward-mode auto diff method overcomes collocation point constraints, leading to more accurate solutions. I think it is of interest to the PINN community and the experimental results are convincing. \n\n- The authors provide the theoretical foundation of SPINN. Also, the authors have done extensive experiments on high-dimensional PDEs to prove the scalability of SPINNs. Although the tested PDEs are not that complicated, I don't think that is a big issue in this paper since it is more related to the optimization issue of PINNs. \n\n- This paper is well-written and well-organized.\n\nWEAKNESSES:\n- This paper only alleviates the scalability issue, but the primary concern of PINNs lies in the challenges of optimization. Therefore, its contribution to the scientific machine-learning community is moderate. \n\n- It would be good to have a paragraph discussing the related work on scalability for PINNs. There are many existing techniques to handle the scalability issue, such as domain decomposition [1], seq2seq [2], and adaptive sampling method [3]. Using adaptive sampling can also help reduce the collocation points. Moreover, in the baseline comparison, it may strengthen the paper to compare SPINN with one of these techniques though my impression is that SPINN will be better based on the results of this paper.\n\n[1] Jagtap, A. D., & Karniadakis, G. E. (2021, March). 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 (Vol. 10).\n\n[2] Krishnapriyan, A., Gholami, A., Zhe, S., Kirby, R., & Mahoney, M. W. (2021). Characterizing possible failure modes in physics-informed neural networks. Advances in Neural Information Processing Systems, 34, 26548-26560.\n\n[3] Subramanian, S., Kirby, R. M., Mahoney, M. W., & Gholami, A. (2022). Adaptive self-supervision algorithms for physics-informed neural networks. arXiv preprint arXiv:2207.04084.\n\nQUESTIONS:\n- What about the performance of forward-mode auto diff for higher-order derivatives (e.g., fourth-order in the Kuramoto-Sivashinsky equation)?\n\n- How do the authors consider the rank *r* (i.e., *r*-dimensional feature representation in Eq. (5))? In Theorem 1, the authors claim *r* should be sufficiently large. What about in practice?\n\n- In lines 193-195, for each axis, this paper considers random sampling for 1D input points. Random sampling may not be an optimal choice. Is there any way to improve it in the context of SPINN? Can the authors add some discussions on this part?\n\n- How do the authors consider the optimization in PINNs? In the section on Limitations (lines 335-342), the authors claim that the learning-rate annealing will be considered in the future. I assume the authors were just doing the hyper-parameter tuning by using a grid search. Is that correct?\n\nLIMITATIONS:\nThe discussion of the limitations in this paper is comprehensive.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a methodology for significant savings in computation and memory in learning physics informed networks for approximately solving PDEs. This enables learning with a much larger number of points, resulting in better accuracy. The idea is to use a specific function class, where features are constructed from each input dimension, and the output is simply the product of these features. The function class is shown to be universal. Since the same feature encoding can be used for all points that share the same value in each dimension, feature embeddings are calculated only once in a linear time for an exponential number of points, resulting in substantial savings. Using forward mode automatic differentiation, derivatives also benefit from the same savings. Experimental results demonstrate the effectiveness of this approach in approximating the solution of diffusion, Helmholtz, Klein-Gordon, and Navier-Stokes PDEs.\n\nSTRENGTHS:\nThe paper is nicely written, the background material is quite useful and many figures help demonstrate the main idea of the paper. In particular, the paper briefly discusses connections to similar ideas in implicit neural representation, and tensor factorization methods, which I found quite relevant.  \n\nThe main idea is simple, and a priori it is not obvious that the chosen function class can perform so well in practice. Therefore the primary strengths of the paper is its rather surprising experimental results, where for some important PDEs orders of magnitude in speed-up results in accurate estimates compared to vanilla PINN and some other variants. For the same reason, the main theorem of the paper that states the universality of the proposed architecture is also quite useful.\n\nOverall, I enjoyed reading the paper, and I believe the proposed methodology can have a high practical impact.\n\nWEAKNESSES:\nWhile I did not identify a major weakness, I believe the paper can benefit from the following changes:\n\n- While experimental results suggest that the proposed architecture can deliver more improvements in some settings than others (e.g., results for diffusion), there is no discussion of this. In general, while the universality result is reassuring, it is unclear in what kind of PDEs it is a useful inductive bias.\n\n- I found Section 4.3 on gradient computation for SPINN is useful; the paper can benefit from doing the same exercise for higher-order derivatives since the complexity of calculating these derivatives in terms of the dimensionality of the embedding space will dominate the complexity of SPINN.\n\n- While I’m aware that reporting relative error is a common exercise, it is quite sensitive and unstable (due to division by the true value of the PDE), and, therefore, could become misleading. I suggest also reporting the RMSE for all experiments.\n\n- For clarity, the paper should state that the overall complexity still remains exponential in resolution (N), despite reducing it to linear in “network propagations”.\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new class of architecture called Separable PINN (SPINN) that operates on a per-dimension basis instead of the standard coordinate based MLP architecture. This architecture significantly reduces the computational complexity and allows the use of a large number of collocation points to enforce the PDE residuals. The paper also provides an universal approximation theorem for the proposed architecture. The paper empirically demonstrates that the proposed SPINN is significantly faster than the original PINN architecture during training while maintaining accuracy.\n\nSTRENGTHS:\n1. The idea of using a different sub-network for every dimension combined with an outer-product operation to evaluate on a grid is both technically sound and interesting.\n2. The paper provides theoretical foundations of the SPINN architecture in the form of the universal approximation theorem.\n3. The gradient computation using forward mode AD and the network architecture together leads to significant speed-ups in the training time of PINNs (which is quite impressive).\n\nWEAKNESSES:\n1. Although the SPINN architecture is quite efficient, I am concerned about the scalability of the architecture for high dimensional PDEs (>4). (Please see question 1)\n2. The size of the plots in Figure 6 is excessively small, making it difficult to compare the performance between the baselines. \n3. The PDEs chosen for the baselines are relatively simple (other than the Navier Stokes). It would be interesting to compare the architectures for equations where standard PINNs fail such as Convection Equation (with larger values of $\\beta$), Kuramoto Sivashinsky Equation (in chaotic regimes).\n\nQUESTIONS:\n1. One of the main claims of the paper is that the computational complexity of the SPINN architecture linearly scales with the number of collocation points $N$ and the dimensions $d$, in contrast to the standard PINN which scales as $N^d$. However, for the SPINN architecture if we have a different branch for each dimension $d$, the outer product operation of multiplying $d$ different quantities can be numerically unstable. For example, can it lead to gradient vanishing or gradient exploding problems? Is some form of normalization required to ensure that the network is stable?\n\n2. The universal approximation theorem guarantees that the approximation error can be as low as possible. However, does the SPINN architecture have any optimization problems?\n\n3. In Figure 6, for a small number of collocation points(16, 32, 64) for Diffusion Equation, the PINNs perform much better than SPINN. Is there a specific reason why SPINNs did not perform well?\n\nLIMITATIONS:\nThe limitations have been adequately addressed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
In this paper, the authors propose two main contributions: (i) forward mode AD for PDEs (ii) separating the contributions from each dimension (just like in separable conv) and performing computation using tensor multiplication. The results naturally show significant improvement over PINNs.

STRENGTHS:
+ Significant amount of speedup in performance.
+ To be honest, this is probably the only PINN variant I have seen showing results for 3D with as many collocation points as 64^3 etc.
+ They can even solve (3+1) d problems.

WEAKNESSES:
- the language needs to be worked upon. It's not that forward AD is a very new area. In the broad community of Neurips, there are a lot of works that use similar ideas. Statements like "To our knowledge, it is the first attempt to exploit the power of forward-mode AD" needs to be made with caution. Perhaps make it clear that in the purview of PINNs or something. In fact, even in PINN-type papers, there are works that use Finite difference and Finite element based approaches (which don't need a backward pass)
e.g. (i) https://arxiv.org/pdf/2005.08357.pdf (ii) https://arxiv.org/abs/2112.04960 (iii) https://arxiv.org/abs/2211.03241 (iv) https://arxiv.org/abs/1901.06314 (v) https://arxiv.org/pdf/2109.07143.pdf
- lack of rigorous comparisons
- lack of ablation studies

QUESTIONS:
- What is the rank in Figure 7?
- It seems that your accuracy with even higher number of collocation points is higher than PINNs. Yes, you are solving faster. But is it better? I don't seem to see it clearly...
- Where are the comparisons with other methods like deepOnets, FNOs, etc. 
- is 1.9e-3 sufficient loss for Navier-Stokes? What is the typical acceptable error in real applications? and can SPINN achieve it?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06