Separable Physics-Informed Neural Networks
Authors: Junwoo Cho
Seungtae Nam
Hyunmo Yang
Seok-Bae Yun
Youngjoon Hong
Eunbyung Park
Venue: NeurIPS 2023 spotlight (2023)
Keywords: partial differential equations
scientific machine learning
physics-informed neural networks
fluid dynamics
TL;DR: proposed a PINNs (physics-informed neural networks) architecture for fast training

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
Junwoo Cho
Seungtae Nam
Hyunmo Yang'
Scok-Bae Yun
Youngjoon Hong
Eunbyung Park 1.
I
"Department of Artificial Intelligence, Sungkyunkwan University
Depaatment of Mathematics, Sungkyunkwan University DDpartmentt of Mathematical Sciences, KAIST
Department of Electrical and Computer Engineering, Sungkyunkwan University
Abstract
Physics-inforin neural networks (PINNs) have recently emerged as promising data-driven PDE solvers showing encouraging results on various PDEs. However, there is 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 severely limited due to the expensive computational costs and heavy mem- ory overhead. To overcome this issue, we propose a network architecture and training algorithm for PINNs. The proposed method, separable PINN (SPINN), operates per-axis basis to significantly reduce the number of network propa- gations multi-dimensional PDEs unlike point-wise processing in conventional PINNs. We also propose using forward-nen modee automatic differentiation to reduce the computational cost computing PDE residuals, enabling large number col- location points (> 10") single commodity GPU. The experimental results show drastically reduced computational costs (62x x in wall-clock time, 1, 394 in FLOPs given the same number of collocation points) in multi-dimenssonal 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 single GPU), maintaining accuracy. Finally, we showcase that SPINN accurately obtain the solution of highly nonlinear and multi-dimensional PDE, a (3+1)-d Navier-Stoker equation. For visualized results and code, please see https://jxcho555: github. 10/ spinn. github. io/
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 often infeasible in many useful and important PDEs [10]. Hence, numerical approximation methods to solutions have been extensively studied [17] c.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 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 promising methods, chysics-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
24
38
Among
"quual contribution.
Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Causal PINN
PINN
Preducted
SPINN (ours)
Reference vorticity 0.1)
vorticils
0
relative L, error
hours later...
-20
10
Absolute ettce
4.01e-3
5.89e-3
0.0
2.5 3.0
training time (hours)
4.31e-3
Training time-. - Ours at 5 sec
Ours at mins 1.48e-2
Ours at mins
Relative La error 8.50e-2
4.26e-3
Figure 1: Training speed (w/ a single GPU) of our model compared to the causal PINN (2+1)-d Wavier-Stokes equation of time interval [0, 0.1].
45
in
gradient-based optimization algorithms to represent and obtain the solutions, leveraging automatic differentiation to enforce the physical constraints of underlying PDE. has enjoyed great success in various forward and inverse problems thanks to its numerous benefits, such 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-dine multi-diimensonal PDEs and approximate very complex solution functions. It primarily stems from using coordinate-basse 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 grows substantially. The situation deteriorates as the dimensionality of the PDE or the solution's complexity increases.
MLP architectures represent the solution function, which takes input
PDEs and obtain more accurate solutions
Recent studies have presented empirical evidence showing that choosing larger batch size (i.e., 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 the scaling characteristic between the batch size and the learning rate (the larger the batch size, the higher the learning rate) 112
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.
27 40
We propose a novel PINN architecture, separable PINN (SPINN), which utilizes forward-mode? automatic differentiation (AD) to enable large number of collocation points (> 10" in single GPU), reducing the computational cost of solving multi-dimensional PDEs. Instead of feeding every multi dilli-dimensional coordinate into single MLP, we use separated sub-networks, in which each sub-network takes independent one-dimensions coordinates as input. The final output is generated by an aggregation module such simple outer product and element-wise summation where the predicted solution can be interpreted by low-rank tensor approximation (25 (Fig. (Fig.Eb). The suggested architecture obviates the need to query every multi dimei-idimenii coordinate input pair, exponentially reducing the number of network propagations to generate a solution, ((N)) + O(Nd). where N is the resolution of the solution for each dimension, and d 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 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 62x on commodity GPUs and FLOPs up to by 1, 394x while achieving better accuracy. Furthermore, with sragesccale collocation points, SPINN can solve a turbulent Anvier-Stoker equation much faster than the sfate-of-the art PINN method 43 (9 minutes vs 10 hours in single GPU) without bells and whistles, such as causal
aa..a. forward accumulation mode tangent linear mode.

---

inductive bias in the loss function (Fig. Our experimental results of the Navier- Stokks equation show that SPINN can solve highly nonlinear PDEs and 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.
Related Works
Physics-informn neural networks. Physics-Informe Neural Networks (PINNs) [3] have re- ceived great attention as a promising learning-basee PDE solver. Given the underlying PDE and initial, boundary conditions embedded in loss function, 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 computational perspective.
The effect of collocation points. The residual loss of PINNs calculated by Monte Carlo integra- tion, not exact definite integral. Therefore, it inherits 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 27 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 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). with AD for the calculation, since numerical methods such as finite differences do not need to back- proppppa through the network. However, they are still burdened by computational complexity of O(N"), 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 4 to handle high-order PDEs such as Kuramoto-SSvaahinnky Sivashinsky equation (28). To the best of our knowledge, the proposed method is the first approach to leverage forward-mode AD training PINNs, which is fully applicable to both time-dependent and independent PDEs and does not incur any truncation errors.
or hybrid approach [5] 37 46
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 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 (.), Haghighat et al. [14. and CoordX [31. NAM T suggested separated network architectures and inputs, but only for achieving the interpretability of the model's 's prediction in multi-task learning. Haghighat al. [14] used multiple MLPs to individually predict each component of the output vector. CoordX's 's [31 primary purpose 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.
3

---

Merge
xi
X2 0-0-0
( m=
(b) (0)
R
Nd points
No points
N points
N points
(a) Separated
(b) Non-sepanraeed
Figure 2: Simple neural networks with dif- ferent input and output dimensions. To compute Øy
Figure 3: An illustrative example of separated approach non-separated approach when f(() R - R (an example case of N 3, d = 3 is shown above). The number of JVP (forward-modk evaluations (propaga- tions) of computing the Jacobian for separated approach (a) is Nd, while the number of VJP (reverse-mode) eval- uations for non-separated approach (b) is Nd.
85 (a) requires one forward pass using forward-mode AD or three backward passes using reverse-mode AD, (b) re- quires three forward passes using forward- mode AD or one backward pass using reverse-mode AD.
3 Preliminaries: Forward/Reverse-nodd AD
For the completeness of this paper, we start by briefly introducing the two
Table 1: An example of forward and eveess-mooee AD in two-layers tanh MLP. Here to denotes the input variable, uk the primals, ég the tangents, 0g the adjoints, and W,, W2 the weight matrices. Biases are omitted for brevity.
types of AD and how Jacobian ma-
trices are evaluated. For clarity, will follow the notations used in [2]
and 113 Suppose our function R" - IIm is two-layer MLP with 'tanh' activation. The seft-hand side of Tab. demonstrates a single forward trace of /. To obtain bian matrix Jy in forward-mode, we compute the Jacobian-vector product (JVP),
Forward primal trace
Forward langent Irace
Backward adjuint trace
Up
11 W1 :
81 W1 Ag
fig
W1
1/
tanb(er))
④2 tanh (1e) b
tanhi (ex)
Jaco-
W2 02
8g = W2 82
F2
Beg
Dey
W2
8y1 0e
Byl Za
J
(I)
@ Zhr,
azn
Bar
8x.
i E {1, (1,...,m) .n). The forward-mode AD is one-phase process: while tracing primals (intermediate values) Uk. it continues to evaluate and accumulate their tangents dj = dr./ /Ox; (the middle column of Tab. [ This equivalent to decomposing one large JVP into series of JVPs by the chain rule the initial tangents vo as the first column vector of an identity matrix In gives the first column of Jy. Thus, the full Jacobian be obtained in forward passes.
0rz/Or;
and computing them from right to left. A run of JVP with
On the other hand, the reverse-mode AD computes vector-Jacobiar product (VJP):
Oyn
0y
8m
Ar.
y =
Oy dys
oy: Bym
8ym
for j (!....,m). which is the reverse-order operation of JVP. This two-phase process. The first phase corresponds to forward propagation, storing all the primals, Uj. and recording the elementary operations in the computational graph. In the second phase, the derivatives are computed by accumulating the adjoints Tn 0y;/Or, (the right-hand side of Tab. Tb. Since VJP builds one row of a Jacobian at a time, takes m evaluations to obtain the full Jacobian. To sum up, the forward-mode is more efficient for tall Jacobian (m ni), while the eeverse-mode is better suited for a wide Jacobian (n > m). Fig. shows an illustrative example and please refer to Baydin et al. 22] for more details.

---

Fua F.11
Fizay Rea
Fira
8(x)
f
(1)),
Para
Σ
Element-wise Product
R,12
Raa
Fizz
Batchified
Mode-1 (columnn) fibers of FiL
Input
X.1 R
Fan
ausn X RN
B Neraa
Fa1 E
Fina
Pira
X1 E
X2 E
E R
X t8
X
F.1
R
Faa
Fir
(a) Acchttecture (single input)
(b) Low-rank tensor decomposition (batchified input)
Figure 4: (a) SPINN architecture 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 given in batches. Each outer product between the column vectors F. from the feature tensor F constructs a rank- tensor and summing all the tensors gives a rank-r tensor. The output tensor of SPINN can be interpreted as low-rank decomposed representation of solution.
4 Separable PINN
4.1 Forward-Mode AD with Separated Functions
We demonstrate that leveraging forward-ne 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- 3R), we first sample N one-dimensional coordinates on each of 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, feature merging function h is used to construct the solution of the entire Nd discretized points. The amount of computations for AD is known to be 2~ times more expensive than the forward propagation [ B According Fig. 3 and with the scale constants Cf, Ca. E [2, 3] we can approximate the total number of operations to compute the Jacobian matrix of the proposed separated approach (Fig. Bh).
Coep Ndeeops(f) + N ch))sh).
(3)
For non-separated approach (Fig.
3)),
Cnon-sop No cfop((f),
we can make ops(h) sufficiently small, then the ratio Cuep/ CCop becomes Na converging 0 as d and N increases. Our separated approach has linear complexity with respect N in network propagations, implying it can obtain more accurate solutions (high-resolution) efficiently.
Nd < 1, quickly
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 1-dimensional coordinate component as an input. Each body-network y(0.) R R" (parameterized
by 0:) is vector-valued function which transforms the coordinates of i-th axis into r-dimensional feature representation. The final prediction computed by feature merging:
ú(z1, .2,,,,,,
Ia) = EI
(0) (2.)
(5)
1=
where à : d - IR. is the predicted solution function, I; R is coordinate of i-th axis, and f/(0) denotes the j-th element of ((0.), We used 'tanh' activation function throughout the paper. As shown in Eq. 5 the feature merging operation is a simple product (II) and summation (E) which
5

---

corresponds to the merging function h described in Eq. Due to its simplicity, h operations are much cheaper than operations MLP layers (i.e., ops (h)
ops(f) f in B. Note that SPINN can
also approximate any m-dimensional vector functions : IR + R by using a larger output feature size (see section D.4 the appendix for details).
The collocation points of our model and conven- tional PINNs have a distinet difference (Fig. 5 Both are uniformly evaluated hypercube, but collocation points of SPINN form lattice-like structure, which we call as factorizable coordinates. In SPINN, dimensional input points from each axis ran- domly 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 co- ordinates with our separated MLP architecture enable us to evaluate functions on dense (N collocation points with a small number (Nd) of input points.
dddimennional
(a) Factorizable coondnates (b) Non-factorizable coordinates coordiins
Figure 5: An illustrative 2-dimensional example of (a) factorizable and (b) non-factorizable coordi- nates, (a) has lattice-like structure, where SPINN can be evaluated more dense collocation points with fewer input points. Conventional PINNs sam- ple factoorireeee coordinates which do not have any structures.
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 be the same. The input coordinates X RRxx xd is now matrix. The batchified form of feature representation F E IR Xrxd and Eq. 5 now becomes
X x)
Fii.a
(6)
where Ú e RN.xNx....X the discretized solution tensor, X denotes outer product, F.,. R.Nx
an i-th frontal slice matrix of tensor F, and F.J.t R is the j-th column of the matrix F F..... Fig. 4 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 is the rank of the reconstructed tensor. Among many decomposition methods, SPINN corresponds to uP-decompoittion [16), which approximates tensor by finite summation of rank-1 tensors. While traditional methods use iterative methods such as alternating least- least-squares (ALS) or alternating slicewise diagonalization (ASD) 2 directly 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 (Vū(z)) is Nd, where N is the number of coordinates sampled from each axis and d the input dimension. According Eq. the i-th element of Vū(x)) is:
Oú dz,
Σ
5 (0)) (2r))) (x2) (x2)...
... ((00)) (2t)).
dz,
Computing this derivative requires feature representations (())) which be reused from the forward pass results computed beforehand. The entire r components of the feature derivatives (x:))a; R + IR" can be obtained by a single pass thanks to forward-mode AD. To obtain the full gradient Vū(x)= Rdd - Rd, 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 Pi Oxx
can also be obtained in the same way.

---

SPINN
mood
PINN
PINN-mod
runbing (ISEES
to errrr
10 12325
120-35P
co.
A of
rundima (sec)
of col. prs
of col. pts
nancime (Sec)
(a) Diffusion
(b) Helmholtz
Pa. G crar
runsine imsit)
(molt)
Decroryy (GB)
12 # of co. pts
2507
230
06 (6eC)
121
& numtme
col pls
pes
(c) (2+1)-d Mordon
(d) --+ Klein-Gordon
Figure 6: Overall results of (a) diffusion, (b) Helmholtz, (c) (2+ D)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 (54% for PINN with modified MI...) For (3+1 ) Klein-Gordon, the training curve is plotted when each model is trained with the number of collocation points of 16 The error and the training curves are averaged over 7 different runs and 70% confidence intervals provided. Note that the y-axis scale of every plot is a log scale.
4.4 Universal Approximation Property
is widely known that neural networks with sufficiently many hidden units have the expressive power of approximating any continuous functions 18). However, it 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 RR. Choose L2 Y). Then, for arbitrary e > O, we can find a sufficiently large r > and neural networks f, and gj such that
Σ 139 leacer <E.
(8)
By repeatedly applying Theorem I we can show that SPINN can approximate any functions in L in the high-dimensional input space. An approximation property for broader function space is fruitful research area, and we leave it to future work. To support the expressive power of the suggested method, empirically showed that SPINN can accurately approximate solution functions of various challenging PDEs, including diffusion, Helmholtz, Klein-Gordon, and Navier-Stokes equations.
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 MI.P introduced in [44] to both PINN and SPINN. For 3-d systems, the number of collocation points of 643 was the upper limit for the vanilla PINN (543 for modified MLP) when we trained with 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²) to get more accurate solutions. This because SPINN stores much smaller batch of tensors, which are the primals (2% in Tab.[]) while building the computational graph. All reported error metrics are average relative L2 emoor computed by l u1 /|ull2 where is the model prediction, and u is the reference solution. Every experiment is performed seven times (three times for 12- and five times for (3+1) Navier-Stokes, respectively) with different random seeds. More detailed experimental settings are provided in the appendix.
7

---

5.2
Results
Fig. showws the overall results of forward problems on three 3-d systems (diffussinn Helmholtz, and (oodnG and one 4-d system (Klein-Gordon). SPINN is significantly more computationally efficient than the baseline PINN in wall-clock run run-time. For every PDE. SPINN with the modified MI.P 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 very effective for solving multi-dimensional PDEs. Furthermore, with the help of the method outlined in Griewank and Walther [3,. we estimated the FLOPs for evaluating the derivatives. Compared to the baseline, SPINN requires 1, x949 fewer
operations to compute the forward pass, first and second order derivatives. Further details regarding the FLOPs estimation are provided in section of the appendix. In the following sections, we give more detailed descriptions and analyses of each experiment.
Diffusion Equation When trained with 128 collocation points, SPINN with the modified MLP finds the most accurate solution. Furthermore, when trained with the same number of collocation points (64%), SPINN-mod is 52x 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. 6p. Due to stiffness in the gradient flow, conventional PINNs hinder finding accurate solutions. Therefore, modified MLP and learning rate annealing algorithm is suggested to mitigate such phenomenon [44). We found that even PINNs with modified MLP fail 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 I I-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%), the training speed of our proposed model with modified MLP is 62x faster, and the memory usage is 29 smaller than the baseline with modified MLP. The exact numerical values are provided in the appendix.
KlrinnGordon Equation Fig. 6p shows the results. Again, our method shows the best performance in terms of accuracy, runtime (62x), and memory usage (29x). 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 Kleinn experiment further, extending it to a 4-d system by adding one more spatial axis. As shown in Fig. 61. baseline PINN can only process 23 (18a 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* (160x larger than the baseline) to obtain more precise solution. The exact numerical values are provided in the appendix.
(2+1)-d Navier-Stokes Equation We constructed SPINN pre- dict the velocity field u and applied forward-mode AD to obtain the vorticity w. We used the same PDE setting used in causal PINN (43
error
training speed
relative error
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 e 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 60x faster than causal PINN in terms of training runtime, The results of the Navier-Stokes equation ensure that SPINN can successfully solve chaotic, highly nonlinear PDE. In Fig, 7 we
peacs
0.2 0,4
100
16
32 rank
64
128
256
Figure 7: Relative error and training speed of (2+1)-d
Navier-Stokes experiment vs. different ranks of SPINN.
8

---

Table 2:
The numerical result of (2+1)-d Navier-Stokes equation compared against PINN with modified MLP mooNN+mod) and causal PINN. Ne the number of collocation points. The fourth row shows the training runtime of a single time window in time-marching training. All relative L errors are averaged over 3 runs.
This is the error (the single run result) reported by the causal PINN paper. Since the causal PINN is sensitive the choice of parameter initialization, we also reported average error obtained from different random seeds by running their official code.
model
mod
44
causal PINN
43
SPINN (ours)
N:
2
2
2
2
2
2
2
2
0.0694
0.0581
0.0578
0.0401
0.0780
0.0363
0,0355
relative L2 error
0.03531
1.000911
1.00135
10001117
10000844
10.0049
10.0018
1.000006
runtime (hh:mm)
03:20
07:52
10:09
23:03
00:07
00:09
00:14
memory (MB) 5,198
17,046
5,200
17,132
764
892
1,276
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 4T. models a decaying vor- tex, widely used for testing Navier-Stokes numerical solvers []. Since the vorticity of the (3+1)-d Navier- Stokes equation is a 3-d vector as opposed to the (2+1)- system, the equation has three independent compo- nents resulting in total of 33 derivative terms (see section D.5 in the appendix). Similar to the (2+1)-d ex- periment, the network's '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. 1.9e-3 less than 30 minutes. Visualized velocity and vorticity vector fields are provided in the appendix.
Table 3: The numerical result of (3+1)4 Navier-Stokes equation. N₂ is the number of collocation points on the entire domain. The relative errors are obtained by compar- ing the vorticity vector field and averaged over five different runs.
model
N:
relative runtime (mins)
memory
L2 error
(MB)
81
0.0090
15.33
768
SPINN
16*
0.0041
16.83
1,192
324
0.0019
26.73
2,946
6 Limitations and Future Works
Although we presented comprehensive experimental results 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 H.3 for more details. Another method could be applying an additional operation after the SPINN's feature merging to map rectangular mesh to an arbitrary physical mesh inspired by PhyGeoNet [T] 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 Di 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 tremendous practical impact. Lastly, we are still far from the theoretical speed-up (wall-clock runtime vs. FL.Pss), which is further discussed appendix section C We expect to additionally reduce runtime by optimizing hardware/software, c.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 various three and four-dimensional PDEs. To our knowledge, it is the first attempt to exploit the power of corward-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 (2022R1A4) 142213443333571 Institute of Information and Communication Technology Planning Evaluation (ITTP) grants for the Al Graduate School program 10TP-2019-0-004211,, Korea Institute of Energy Technology Evaluation and Planning (KETEP) and the Ministry of Trade, Industry & Energy (MOTIE) of the Republic of Korea (No. 20224000000366),
References
[1] 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-477 2021.
[2] Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Auto- matic differentiation in machine learning: a survey. Journal of Marchine Learning Research, pages 1-43, 2018.
[3] John Boyd. Chebyshev and Fourier spectral methods. Courier Corporation, 2001.
[4] Steven Brunton and J Nathan Kutz. Data-driven science and engineering: Machine learning, dynamical systems, and control. Cambridge University Press, 2022.
[5] Pao-Hsiung Chiu, Jian Cheng Wong. Chinchun Ooi, My Dao, and Yew-Soon Ong- Can-pinn: A fast physics-informed neural network based on coupled- automatic- differentiation method. Computer Methods in Applied Mechanics and Engineering, 395:1 14909, 2022.
[6] George Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
[7] Arka Daw, Jie Bu, Sifan Wang, Paris Perdikaris, and Anuj Karpatne. Rethinking the importance of sampling in physics-informed neural networks. arXiv preprint arXi::2207.02338, 2022.
[8] 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.
[9] Robert Eymard, Thierry Gallouët, and Raphaele Herbin. Finite volume methods: handbook of numerical analysis. PG Ciarlet and JL Lions (Eds), 2000.
[10] Charles L Fefferman. Existence and smoothness of the navier-stokes equation. The millennium prize problems, 57:67, 2000.
[11] 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:1 10079, 2021.
[12] 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 hour. arXiv preprint arXiv: 706.7006.22222 02677, 2017."
1rXiv::706.002777
[13] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008
[14] Ehsan Haghighat, Maziar Raissi, Adrian Moure, Hector Gomez, and Ruben Juanes. A physics-informed informed deep learning framework for inversion and surrogate modeling in solid mechanics. Computer Methods in Applied Mechanics and Engineering, 379:113741, 2021.
10

---

[15] John Hammersley. Monte carlo methods. Springer Science & Business Media, 2013.
[16] Frank L Hitchcock. The expression of a tensor or polyadic as a sum of products. Journal of Mathematics and Physics, 6(1-4):164-19) 164-189, 1927.
[17] Joe Hoffman and Steven Frankel. Numerical methods for engineers and scientists. CRC press, 2018.
[18] Kurt Hornik, Maxwell Stincheombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 255)339--66, 1989.
[19] Ameya D Jagtap and George E Karniadakis, Extended Dhysics-informe 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: 2020.
[21] Ameya D Jagtap, Ehsan Kharazmi, and George Em Karniadakis. Conservative physics-infone neu- ral networks on discrete domains for conservation laws: Applications forward and inverse problems. Computer Methods in Applied Mechanics and Engineering, 365: 11302,, 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. rate convergence of physies informed neural networks for the linear second order elliptic pdes. Communications in Computational Physics, 31(4):1272, 2022.
[24] George Em Karniadakis, Ioannis Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-ifformed machine learning. Nature Reviews Physics, 3(6) 2021.
[25] Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review, 5153):55-500 2009.
[26] Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Characterizing possible failure modes physics-infrmed neural networks. Advances in Neural Information Processing Systems, 34:26548-2663 265600, 2021.
[27] Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv:144..5999 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, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes general geometries. arXiv preprint arXiv:2207.05209, 2022.
[30] 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.
[31] Ruofan Liang, Hongyi Sun, and Nandita Vijaykumar. Coordx:
Accelerating implicit neural representation with split mip architecture. In International Conference on Learning Representations, 2022.
[32] Ben Moseley, Andrew Markham, and Tarje Nissen-Meyer. Finite basis physics-informet neural networks (fbpinns): a scalable domain decomposition approach for solving differential cquations. arXiv preprint 1xXi:210770077/. 2021.
[33] Guofei Pang. Lu Lu, and George Em Karniadakis. fpinns: Fractional physics-infommed neural networks. SIAM Journal on Scientific Computing, 41(4): 2019.
[34] Maziar Raissi, Paris Perdikaris, and George F. Karniadakis. Physies-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] Rishikesh Ranade, Chris Hill, and Jay Pathak. Discretizationnet: A machine- leamnnn based solver for navier-stokes equations using finite volume discretization. Computer Methods in Applied Mechanics and Engineering, 137:22 2021.
11

---

[36] 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 2022.
[37] Ramansh Sharma and Varun Shankar. Accelerated training of physics informed neural networks (pinns) using meshless discretizations. Advances in Neural Information Processing Systems, 2022.
[38] Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339 1364, 2018.
[39] Gordon Dennis Smith. Numerical solution partial differential equations: finite difference methods. Oxford university press, 1985.
[40] Samuel L.. Smith, Pieter-Jan Kindermans, and Quoe V. Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations. 2018.
[41] 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-Meahematical and Physical Sciences, 158(895):499-21, 1937.
[42] Nils Wandel, Michael Weinmann, Michael Neidlin, and Reinhard Klein. Spline-pimi Approaching pdes without data using fast, physics-informed hermite-spline cans. In Proceedings of the AAAI Conference on Artificiad Intelligence, pages
2022.
[43] 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.
[44] 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-A3308 2021.
[45] Yiheng Xie, Towaki Takikawa, Shunsuke Saito, Or Litany, Shiqin Yan, Numair Khan, Federico Tomhari, 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.
[46] Yinhao Zhu, Nicholas Zabaras, Sinedon-Sttllon Koutsourelakis, and Paris Perdikaris. Physics-constranned deep learning for high dineediinnnional surrogate modeling and uncertainty quantification without labeled data. Journal of Computational Physics, 394:56 81, 2019.
[47] Olek C Zienkiewicz, Robert Leroy Taylor, and Jian Z Zhu. The finite element method: its basis and fundamentals. Elsevier, 2005.
12

---

Separable PINN Supplementary Materials
A Proof of Theorem 1
Here we show the preliminary lemmas and proof' for Theorem 1. in the main paper. We start by defining general lensor product between Hilbert spaces.
Definition 1. Let {es} be orthonormal basis for H2. Tensor product between Hilbert spaces Hi and H2. denoted by Hi & H2, is a set of all antilinear mappings A : H2 -> Mi such that Aus1 co for every orthonormal basis for H2.
Then by Theorem 7.12 in Folland [4. Hi @ H2 is also a Hilbert space with respect to norm ||| and associated inner product (;
|||
>
(1)
8
(A, Σ (Aus, Beg),
(2)
where A, H, H2, and {ug} is any orthonormal basis of H2.
Lemma 1. Let x E Hi and 3/ E H2. Then, (x @ v)y' (v.y))z v²)z.
Proof. Let A be mapping A : 3/ + "'), where ||y/| 1. We expand y with orthonormal basis {u, 20). i.e. {y, 2a} is basis for H2. Then,
| 2
Σ
Azæll2 Al
Σ || (y,, 2g)=²|
|| < 00
(4)
obvious that A is antilinear. Then by Definition [ Therefore, (x & y)v' /...
A H₂ @ H2 which is, Ay' (ar v)y/
Lemma 2. {wa @ Ug} is an orthonormal basis for Hi ? H2-
Proof. Let A, Hi H2, where ua 2 ug. Then by the definition of inner product in Eq.
(A, B) = Σ (Aus, Bui)
(5)
Σ (Aus, (ua 8 vs)us)
(6)
Σ (Aus, (vg, vijual
(:. Lemma] T
(Au vitual +
+... + (Ave, (13, (oy,og)ua +...
(8)
(Ang, ud)
(9)
Now we check the Parseval identity:
Σ (A, ua & ug)1² = Σ ((Agg, ua)l
(: Ua Ø 1 BB B)
(10)
0,3
o.,8
Antell
(11)
B
= ?
(12)

---

fua @ vo is a basis.
Now we begin the proof of Theorem in the main paper.
Theorem Let Y be compact subsets of Choose (xxx)), Then, for arbitrary we can find sufficiently large r > 0 and neural networks fg and 9, such that
|w- Σ f39 Lexxn Se.
(13)
Proof. Let { and {1y} be orthonormal basis for (² and (Y respectively, Then {dmb forms orthonormal basis for (2 x Y) (: Lemma 2 Therefore, can find sufficiently large r such that
angont (x(xxV)
< 2
(14)
where di denotes
u(x, (1)
Ozj
XxY On the other hand, by the universal approximation theorem we can find neural networks fy and Bj such that
| falle(x) 30]|u|1 and || 9,llea(r) 311lal1a(xx))
(15)
We first consider the difference between u and Ei, aigSigsi
Σ ai fisllllexx
(16)
lexcer + +12 az aayoriy -a
11
(17)
i.
1.j
+
IT
(18)
Since ||| < e/2 from
[4], is enough to estimate II. For this, we consider
Σ
anjobby
asySegi Σ anjos (1b, - 995) + Σ
ais ais (1)))
(19)
1.j
4
111 + 112.
(20)
We first compute l11:
2 aiyos 1.)) }
1h 1 ((X = J
dxdy
(21)
Σ aijide (w/ - 15) } drdy
=JxWy
Σ
(1s)
(22)
We set
Az(x) = Σ (1)))(x))
By(x) +,) (2(4)
(23)
14

---

to write 1[1 as
)
2
| [ull = SKx XXY {E
dody.
(24)
We then apply Cauchy-Scharwz inequality to get
| 1Xxxx) - JJOV (Σ ((A(())))) (E1s(())P) )xddy
drdy
(25)
= (1 (1 ) (1.5)
By (u))) (1.513(01120
(26)
Now
1 (ESSOO)P
/
r
Σ
Σ aiyli
dr
(27)
Σ
Σ
atj
Φs l(x)
2
(28)
Since { is an orthonormal basis, we see that
Σ 0:j LLx,)
S >Σ
|ays lasll lxxx =
Ibaagl layl-
(29)
1(X)
Therefore,
J Σ
1
Σ aij
dx
(30)
Σ/ x
/
Σ aij
aijoi
dx
(31)
Σ
Σ
Dij Oi
(32)
L(X)
Σ
lass ))
(33)
Finally, we recall
Σ
) 2
< Σ
lass1
(34)
to conclude
r " Σ
r
xc >
Σ 1xxx)) 3d $2
lais" 2
< 2
laiz1 22
(35)
1
i
15

---

On the other hand, we have from 15
1 Σ B dy Σ / (15(()) (v)
(36)
Σ (10s(()) - WW
(37)
e²
(38)
9 1| 11
(39)
1 8111122(Xx))
Hence we have
c²
8²
|
<
2 lu|l (lll22(xxx))
(40)
8||u1
Likewise,
8²
<
(41)
Therefore,
8²
|
(42)
2
We go back to (16) with the estimates
14 and
42
to derive
Σ anfig (xXXY)
2 2 E.
(43)
Finally, we reorder the index rewrite
Σ
Cayfig Σ
big
(44)
Σ (b gt
(45)
Σ kgg
(46)
Without loss of generality, we rewrite i, k,, Oj into r, fú. gi 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 forward- mode AD to compute PDE residuals (standard back proppapapag,,nnn a.k.a. roverse-mode AD for parameter updates). With the slight abuse of notation, our predicted solution function is denoted as (( (I, t) from onwards, explicitly expressing time coordinates. Given an underlying PDE (or ODE), the initial, and the boundary conditions, SPINN is trained with "physics-ifformed' loss function:
min C(u (0) (I, t)) = min Xpdo Cpde + AicCio + Abc 0
(47)
Cpde 1 INTT (1,t))]dddtt
(48)
J
Lre
99 (10100) (x. 0) "u<(t))l2
(49)
Lbe
[|BTu(0) (x, t) unc drdt,
(50)
18
16

---

where Ω is an imput domain, N, B are generic differential operators and die, Uhe are initial, boundary conditions, respectively. À are weighting factors for each loss term. When calculating the PDE loss pae)) with Monte-Carlo integral approximation, we sampled collocation points from factorized coordinates and used forward-modk AD. The remaining Lie and Lie 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 clementary operations such as additions and multiplications. Given 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 series of matrix multiplications and additions. We used the method described in [6] to estimate FLOPs for evaluating the derivatives. Table | shows the number of additions (ADDS) and multiplications (MULTS) in each evaluation process. Note that FLOPs is summation of ADDS and MULTS by definition.
One thing to note here is that this is theoretical estimation. Theoretically, the number of JVP evaluations for computing the gradient with respect to the input coordinates is Nd, when 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())), which makes the complexity network propagations from O (NV)) to ((N)). Ideally, these feature vectors can be computed only once and stored to be used later for gradient computations. Although
still significantly more efficient than the conventional PINN's complexity (Ndf < N, there is room to bridge the gap between theoretical FLOPs and actual training runtime by further software optimization.
Table 1: The number of elementary operations for evaluating forward pass, first and second-order derivatives. The calculation is based on 643 collocation points in 3-d system and the vanilla MLP settings used for diffusion, Helmholtz, and evaluated on every coordinate axis.
equations. We assumed that each derivative
SPINN (ours)
PINN (baseline)
ADDS (x10))
MULTS (x 109)
ADDS (x (x10%)
MULTS (x 10%)
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
MFLOPs (total)
280
390,370
D Experimental Details and Results
In this section, we provide experimental details, numerical results, and visualizations for each experiment in the main paper. Below, 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 Navier- Stokes, (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 Diffusson Equation
The diffusion equation is one of the most representative parabolic PDEs, often used for modeling the heat diffusion process. We especially choose nonlinear diffusion equation where it can be written
17

---

as:
O.u a (1/u/) 2 + uAu).
x. E Q I 9
(51)
u(z, 0) (x(
(52)
u(2, 1) = 0,
E 00,101
(53)
We used diffusivity a = 0.05, spatial domain 52 = 1=-1,11 temporal domain = [0. 1] and used superposition of three Gaussian functions for the initial condition vic- We obtained the reference solution (101 x 101 x 101 resolution) through a widely-used PDE solver platform FEniCS [2]. Note
that FEniCS is FEM-based solver. We particularly set the initial condition be a superposition of three gaussian functions:
Ugo (z. 9) 0.25 exp [-10 - 0.2) (y 1 0.3))]] + 0.4 exp (-15((1+ 0.1)² (y+
+ 0.3 0x3 -0 ( +2]] (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 single MLP 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. shows the numerical results, and the visualized solutions are provided in Fig.
D.2 Helmholtz Equation
The Helmholtz equation a time-independent wave equation that takes the form:
Au + k²u q, u( (xx) 0,
x E Ω
(55)
x E an,
(56)
where the spatial domain is Ω 1- 1ز For a given source term q - (a (a2 )2 (azT)²u + u2u, we devised a manufactured solution u = sin(ayrzz) (T)) sin (Q2TTz2 sin(aa)rz)) where we take k 1, a1 4, a2 4, 03 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, used 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.
D.3 Klein-Gordon Equation
The Elein-Gordon equation is a nonlinear hyperbolic PDE, which arises in diverse applied physics for modeling relativistic wave propagation. The inhomogeneous Glrion equation is given by
duuu- Au + 2
x E e 2 E Ω I E 00, / E E,
(57)
u(x, 0) = I1 x2,
(58)
u(x, t) 14c (2:).
(59)
where we chose the spatial/temporal domain to be Ω [-1, 1]2 and I [0, 10], respectively. For error measurement, we used a manufactured solution u = (œ1 + x2) cos (2t) + 21 (2t) and f. Uhe 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 single MLP of 5 hidden layers with 128 hidden feature sizes. We used Adam (8 optimizer with 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.
18

---

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;
(x1 22 + 23) cos (t) + 21 sin (t),
(60)
where ʃ. lhe are extracted from this exact solution. Tab. shows the numerical results. The number of collocation points of 23", 184 were the maximum value for PINN and PINN with modified MLP, respectively.
D.4 (2+1)-d Navier-Stokes Equation
Ravier-Stokes equation is a nonlinear Lime-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 written as below:
Orw + Vw vdw,
E VERLEET. x E EENIET,
(61)
V v = 0, wxx, 0) (x)
(62)
E
(63)
where u E R is the velocity field, w x u is the vorticity, wp is the initial vorticity, and is the viscosity. We used the viscosity 0.01 and made the spatial/temprral domain f2 = [0. 2T and [0, 11, respectively. Note that Eq. 61 models decaying turbulence since there 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-spectre method. The initial condition was generated using the gaussian random field with maximum velocity of 5. The resolution of the obtained solution is 100 X 128 x 128 (N₂ X N₂ x Ny), 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] [8]. We used Adam [8] optimizer with 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.
N. de
Lpde
Au Ne
10:w + U.w.w + uyyOyw - v(Orzw w + )yy No
Σ
10 u, + +,uuyl? (64)
Nac
Aic
Lio
(lu-
Uno unol2+ + My uyo + - wol²).
(65)
Nic
where Ur, My are x. components of predicted velocity, collocation points, Nic is the number of coordinates for initial condition and uzo, Uyo, Mho are the initial conditions. We chose the weighting factors A = 1, A 5, 000, and Aio = 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.
0_ 8y N₂ is the number of
(x) = [1, sin(z), sin()x), sin(3x), sin(4x), (ine),,in((()) (5x), cos(x), cos(2x), cos(3x), cos(4x), cos (osr)))))
(66)
Unlike other experiments, the solution of the Havier-Stoker equation is a 2-dimensional vector-valued function : IR3 + 2 We can rewrite the feature merging equation Eq. 5 in the main paper to construct SPINN into 2-dimensional vector function:
"1
ΣI
(8.) (.))
(zi),
(67)
#=1
2r
U2 Σ
(x:))
(z:).
(68)
For example in our (2+1)-d -d Savier-Stoker equation setting, r r= 128 since the network output feature size is 256. This can be applied to any m-dimensional vector function if we use larger output feature
19

---

size. More formally, if we want to construct an m-dimensional vector function with SPINN of rank the k-thh element of the function output can be written
kr Σ
W
II
(69)
J=
1)++li=l
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
Drw + (u V) = (w V)u + DAw +.F V 0,
x ee,ter I E R,tte,, x JE Ω
(70)
(71)
(xx, 0) (2),
(72)
We constructed the spatial/temporal domain to be Ω = [0,2*] and r [0. respectively. We constructed the analytic solution for (3+1)-d Navier-Stokes equation introduced Taylor et The manufactured velocity and vorticity are
Ug 2e 29-9vt cos (2x) sin (2y) sin(z),
(73)
Bl gut sin sin( cos (2y) sin (z),
(74)
u, -2 sin (2x) sin (2y) cos(2).
(75)
We = 3e sin (2x)) cos (2 cos(2);
(76)
Wy = be 9ut cos (2x) sin (2 cos(2);
(77)
W3 = - 6e got c2s cos (2 sin
(78)
where we chose the viscosity to be 0.05. Each forcing term F in the Eq. 70 is then given
Fz 6e
18vt sin 4y) sin(2:),
(79)
Fy 6c 18st sin (z) sin (22),
(80)
F= 6e 18ut sin(4x) sin (yy)
(81)
We constructed SPINN to be four body networks (modified MLP) of hidden layers with 64/384
hidden feature/output sizes each. We used Adam g 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 Apde Djc = 10, and Abc = 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:
@u(t, x)
Au(t, x)).
x E 1-1,11 [0,
(82)
Ot
where the manufactured solution is chosen to be || + 10t. When trained with 86 collocation points, SPINN achieved relative L2 error of 0.0074 within 2 minutes.
E.2 -d Flow Mixing Problem
This is time dependent PDE which describes the behavior of two fluids being mixed at the interface in 2-d environment. Following the settings in CAN-PINNs [. the equation is
@u(t,x,V) x. v) + dx
Ou(t, x. y) +
ou(t,,,,))
E e (- (-4,4,, (83)
at
Oy
20

---

where
V/
a(x, y)
(84)
max
b(z, y)
(85)
Ut mt,max = sech (r) (r))na((r))
(86)
r 2² y²,
(87)
mt,max 0.385.
(88)
It has an analytic solution 16
utt, v, y) - tanh cos(wt) 2 sin(wt)
(89)
where w = Memas The initial and boundary conditions are extracted from the exact solution Eq. 89
When trained with 256 collocation points, SPINN achieved relative Lg error 0.0029 in minutes. Figure shows the visualized result.
t-
t-2
Exact solution
Predicted solution
Figure 1: Visualized solution of flow mixing problem obtained by SPINN.
E.3 2-d Poisson equation L-shaped domain
Besides the coordinate or mesh transformations (5] ITI 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- monnfactorizable coordinates). For collocation points inside the domain, we train separable way (factorizable coordinates) and ignore the points outside the domain. This can be easily achieved by finding a tight bounding box the input domain and masking out the PDE loss when 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
-ul(x, y) 1, (x. y) E Ω u(x.v) (x. v) E an,
(90)
(91)
where ! - (-1,) 1]2. The reference solution was obtained by the spectral method. SPINN achieved relative L2 error of 0.0322, while PINNs with modified MLP achieved 0.0392. Figure 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--haped domain we showed just example case, and we believe that this method can be applied to any arbitrary input domain.
21

---

Reference solution
Prediction (SPINN)
Prediction (PINN)
L error: 3.22e-2
Rel. L error: 3.92e-2
100
1.00
675
012
010
008
000
000
0.0
006
025
334
004
006
002
002
082
-180
800
180
-.0
-05
60
1.00
-05
Figure 2: Visualized solution of Poisson equation obtained by SPINN and PINNs with modified MLP.
E.4 Fine Tuning with L-BFGS
We also conducted some experiments to explore the 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 [4] and is known to be effective some cases. Tab. shows the numerical results on three 3-d PDEs. Understanding the effect the optimization algorithm is still an open question in PINNs, believe that investigating this issue in the context of SPINN would be a valuable direction for future study.
Table 2: Numerical result of 3-d PDEs with L-BFGS fine-tuning. The number of training collocation points is 64 and a single outer loop L-BFGS applied.
Diffusion
Helmholtz
Klein-Gordon
Adam Adam L-BFGS
0.0041 0.0041
0.0360 0.0308
0.0013
0.0010
Table Full results diffusion equation. N₂ is the number collocation points.
relative L2 error
runtime (ms/iter.)
memory (MB)
model
N₂
RMSE
163
0.0095
0.00082
3.98
1,022
PINN
323
0.0082
0.00071
12.82
2,942
643
0.0081
0.00070
95.22
18,122
163
0.0048
0.00042
14.94
1,918
PINN modified MLP
323
0.0043
0.00037
29.91
4,990
543
0.0041
0.00036
134.64
22,248
163
0.0447
0.00387
1.45
766
323
0.0115
0.00100
1.76
766
SPINN
643
0.0075
0.00065
1.90
766
1283
0.0061
0.00053
2.09
894
2563
0.0061
0.00053
10.54
2,174
163
0.0390
0.00338
2.17
766
323
0.0067
0.00058
2.44
768
SPINN + modified MLP
643
0.0041
0.00036
2.59
768
1283
0.0036
0.00031
3.06
896
256
0.0036
0.00031
12.13
2,176
22

---

Table 4: Full results of the Helmholtz equation. No is the number of collocation points.
model
N
relative L2 error
RMSE
runtime (ms/iter.)
memory
(MB)
163
0.9819
0.3420
4.84
2,810
PINN
323
0.9757
0.3398
14.84
2,938
643
0.9723
0.3386
110.23
18,118
163
0,4770
0.1661
18.32
7,034
PINN modified MLP
323
0.5176
0.1803
35.02
9,082
543
0.4770
0.1661
159.90
22,244
163
0.1177
0.0410
1.54
762
323
0.0809
0.0282
1.71
762
SPINN
64
0.0592
0.0206
1.85
762
1283
0.0449
0.0156
1,89
762
256
0.0435
0.0151
3.84
1,146
163
0.1161
0.0404
2.24
764
323
0.0595
0.0207
2.50
764
SPINN + modified MLP
643
0.0360
0.0125
2.57
764
1283
0.0300
0.0104
2.76
764
256
0.0311
0.0108
5.50
1,148
Table Full results of the (2+1)-d Klein-Gordon equation. N₂ is the number collocation points.
model
N₂
relative
RMSE
runtime (ms/iter.)
memory
L2 enror
(MB)
163
0.0343
0.0218
4.70
2,810
PINN
323
0.0281
0.0178
14.95
2,938
643
0.0299
0.0190
112.00
18,118
163
0.0158
0.0100
17.87
7,036
PINN modified MLP
323
0.0185
0,0118
34.61
9,082
543
0.0163
0.0104
159.20
22.246
163
0.0193
0.0123
1.55
762
323
0.0060
0.0038
1.71
762
SPINN
643
0.0045
0.0029
1.82
762
1283
0.0040
0.0025
1.85
890
256
0.0039
0.0025
3.98
1,658
163
0.0062
0.0039
2.20
764
323
0.0020
0.0013
2.41
764
SPINN + modified MLP
644
0.0013
0.0008
2.57
764
1283
0.0008
0.0005
2.79
892
250
0.0009
0.0006
5.61
1,660
23

---

Table Full results of the (3+1)-d Klein-Gordon equation. N₂ is the number of collocation points.
model
N₂
relative
RMSE
runtime (ms/iter))
memory (MB)
L2 error
164
0.0129
0.0096
43.51
5,246
PINN
234
0.0121
0.0090
154.24
22,244
164
0.0061
0.0045
100.06
17,534
PINN modified MLP
18a
0.0059
0.0044
174.00
22,246
164
0.0122
0.0091
2.45
890
SPINN
32
0.0095
0.0071
2.98
892
64
0.0093
0.0069
9.22
2,172
164
0.0064
0.0048
3.48
892
SPINN + modified MLP
324
0.0022
0.0016
3.66
892
64
0.0012
0.0009
10.96
2,172
24

---

0.225
0.320
0323
00000
0.20
5.00
2000
(a) Training time: 5 sec
PINN
0.000
0.000
0.20
0.00
(b) Training time: 30 sec
MNN
(SanO) SNIIS
0.5
0.06
(c) Training time: 1 min
Referenec
0222
033
003
0.0
t= 0.25
t.5
t= t=1.0
Figure 3: Visualized solution of nonlinear diffusion equation obtained by the baseline PINN and SPINN, both trained 6 collocation points.
25

---

SPINN (0mms)
boolaat error (SPINN)
SPINN (Ours)
booluta erfor (SPINN)
SPINN (Ours)
bsooluto error (SPINN)
Exact solution
PINN
Absollte (PINN)
PINN
Absolute (PINN)
Absolute (PINN))
(a) Training time: 5 sec
(b) Training time: 30 sec
(c) Training time: 1 min
Figure 4: Visualized solution of Helmholtz equation obtained by the baseline PINN and SPINN, both trained on 64 collocation points.
Reference w(0.0, x, y)
Predicted w(0.0, x, y)
Absolute Error
20
20
0.10
10
10
0.08
0.06
10
10
0.04
-20
-20
0.02
0
0
0.00
1
2
3
4
5
1
2
3
4
5
2
3
4
5
Reference (0.5, y)
Predicted w0..5,
Absolute Error
20
3.5
3.0
10
10
2.5
>
0
0
2.0
1.5
10
10
0.5
-20
20
Reference (1.0, x, y)
Predicted w11.0, x, y)
Absolute Error
6
15
15
10
10
5
5
0
0
3
-5
2
2
2
2
10
1
1
0
20
20
5
Figure 5: Visualized vorticity maps (2+1)-d Navier-Stokes equation experiment predicted by
SPINN. Three snapshots at timestamps 0.0, 0.5, 1.0 presented.
26

---

(a) Exact solution (velocity)
(b) Prediction (velocity)
(c) solution (vorticity)
(d) Prediction (vorticity)
Figure Visualized solution of (3+1)-d Navier-Stokes equation obtained by SPINN, trained on 324 collocation points. Each arrow is colored by its zenith angle.
27

---

References
[1] Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong. Can-pinn: differentiation
A fast physics-informed neural network based on satplee-automatic-nnnnen method. Computer Methods in Applied Mechanics and Engineering. 14909, 2022.
[2] George Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314 1989.
[3] 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: 10242, 2021.
[4] Gerald B Folland. course in abstract harmonic analysis, volume 29. CRC press, 2016.
[5] Han Gao, Luning Sun, and Jian-Xun Wang. Phygeonet: Physics-informed ccometry-adapteee convolutional neural networks for solving parameterized steady-state pdes on irregular domain. Journal of Computational Physics, 428: 10079, 2021.
[6] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.
[7] 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.
[8] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[9] Dmitrii Kochkov, Jamie Smith, Ayya Alieva, Qing Wang, Michael Brenner, and Stephan -eaningg-cccelerated
Hoyer. Machine learning accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118021))e21007741188 2021.
[10] Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Characterizing possible failure modes in physics-infommn neural networks. Advances in Neural Information Processing Systems, 34:26548-266 -25560, 2021.
[11] Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural oper- ator with learned deformations for pdes on general geometries. arXiv preprint +xXix:2207.05209 2022.
[12] 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.
[13] 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.
[14] Xuhui Meng, Zhen Li, Dongkun Zhang, and George Em Karniadakis. Ppinn: Parareal physics- informed neural network for time-dependent pdes. Computer Methods in Applied Mechanics and Engineering. 370: 13250, 2020.
[15) Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physies-informed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707,2011
[16] Panos Tamamidis and Dennis N Assanis. Evaluation of various high-order-acccuracy schemes with and without flux limiters. International Journal for Numerical Methods in Fluids, 16(10):931-94,, 1993.
[17] 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):4995522 1937.
[18) Colby I. Wight and Jia Zhao. Solving allen and cahn-hilliard equations using the adaptive physics informed neural networks. arXiv preprint arXiv:2007.04542. 2020.
28


Reviews:
[{"review_id": "tMDWtcuK8g", "reviewer": "Reviewer_3ENg", "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. ", "questions": "- What is the rank in Figure 7?\n- 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...\n- Where are the comparisons with other methods like deepOnets, FNOs, etc. \n- is 1.9e-3 sufficient loss for Navier-Stokes? What is the typical acceptable error in real applications? and can SPINN achieve it?\n\n", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ Significant amount of speedup in performance.\n+ 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.\n+ They can even solve (3+1) d problems.\n", "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)\ne.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\n- lack of rigorous comparisons\n- lack of ablation studies", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**R: Expressions when stating our contribution**  \nA: In the related work section, we mentioned that we utilized forward-mode AD *in training PINNs* (lines 108-109), but regrettably, this detail was inadvertently omitted in the conclusion section (line 348). Thank you for pointing this out. We will carefully examine the entire text for any additional inaccuracies and make the necessary corrections.  \n***  \n**R: Forward-mode AD is not very new**  \nA: We want to clarify that our contribution lies in the demonstration of novel PINNs architecture that *utilizes* the forward-mode AD *for fast and accurate PINNs training*. Although forward-mode AD itself is not a new concept, we have rediscovered its potential within the context of PINNs research.\n***\n**R: PINN-type papers that use numerical derivatives**  \nA: We appreciate the provided references. We will cite them in the related works section. It is true that PINN models that use numerical differentiation do not require any backward pass for obtaining the input derivatives. However, these models are still burdened by a computational complexity of $O(N^d)$ during the network propagation, thereby limiting them to handle large-scale collocation points. Furthermore, numerical differentiation has truncation errors depending on the step size.\n***\n**R: Lack of rigorous comparisons**  \nA: We are afraid that we do not know which comparison you would like to see, so we tried our best to give more information in this context. But, please let us know the details of the rigorous comparisons. In the experiment section, we focused on demonstrating 1) training speed and 2) the effectiveness of using large-scale collocation points.\n* Training speed: For a fair comparison, we matched all factors which can affect the training runtime between our model and the baselines. Each experiment was run on a single RTX 3090 GPU and the training source code for both our model and the baselines was implemented with JAX. We matched the total number of learnable parameters, and the training speed of each model is compared under the same number of collocation points. \n* The effectiveness of using large-scale collocation points: We wanted to show that using large-scale collocation points during training can enhance the model’s ability to predict more accurate solutions. We numerically compared the accuracy using the relative L2 error metric. However, as the reviewers ‘KaJa’ and ‘dosJ’ pointed out, we acknowledged the potential limitations of the relative error in universally explaining the results. To address this, we have introduced two additional metrics – RMSE and PINN loss value – as detailed in the attached PDF file. We will provide the complete results in the revised version of the paper if permitted.\n***\n**R: Lack of rigorous ablation studies**  \nA: Again, please let us know which particular ablation study you would like to see. The most essential factor that determines the performance of SPINN is the rank of the reconstructed tensor. We carried out experiments with varying ranks on the (2+1)-dimensional Navier-Stokes equation to showcase its impact on both the prediction accuracy and training speed (Figure 7).\n***\n**R: Rank in Figure 7**  \nA: This is the rank of the reconstructed solution tensor. SPINN constructs multiple rank-1 tensors via the outer product, and the rank is the number of these rank-1 tensors to be added to represent the final solution tensor. Please see Figure 4 in the main paper for details. Increasing the rank gives more expressive power to the model supported by Theorem 1 in the main paper. Figure 7 is intended to demonstrate whether increasing the rank actually helps the model find more accurate solutions.\n***\n**R: No comparisons against operator methods**  \nA: We did not take the neural operator methods into account because they require a pre-training stage to train a neural network over a large set of ground truth data. Thus, the resulting neural network is not capable of solving other PDEs (not shown in the pre-training stage). For example, an FNO trained on Burgers’ equations cannot be used to solve the Helmholtz equations. On the other hand, SPINN (or PINN in general) can be applied to solve any PDEs without any data. To summarize, PINNs and operator learnings (deepOnets, FNO) differ in many aspects, and we thought it is hard to perform an apples-to-apples comparison. We will add more discussions about the differences in the final version of the paper.  \n***  \n**R: Typical error in Navier-Stokes equation**  \nA: For the (2+1)-d case, we’ve shown that SPINN achieved higher accuracy compared to the causal PINNs, which is the best-performing prior method. We have found a PINNs paper [1] that also tried to solve a (3+1)-d Navier-Stokes equation. Although there is a discrepancy in the equation settings, we leave their result below for a rough reference.  \nEquation: (3+1) dimensional Beltrami flow  \nManufactured solution (velocity vectors):  \n$u_x = -a[e^{ax}\\sin(ay+dz)+e^{az}\\cos(ax+dy)]e^{-d^2t}$  \n$u_y = -a[e^{ay}\\sin(az+dx)+e^{ax}\\cos(ay+dz)]e^{-d^2t}$  \n$u_z = -a[e^{az}\\sin(ax+dy)+e^{ay}\\cos(az+dx)]e^{-d^2t}$  \nRelative L2 error: 0.023886  \nWe want to note that the error encountered when solving the (3+1)-d Navier-Stokes equation using a numerical solver can vary significantly based on the specific solver, the chosen discretization methods, the grid resolution, and the nature of the problem to being solved. It is hard to pin down a single “typical” error value as it depends on various factors as mentioned above. Moreover, since there’s no analytical solution available for such a general PDE setting, engineers and researchers often assess the accuracy of their numerical solutions by studying the convergence behavior of the numerical methods, akin to our proposed theorem.  \n  \n[1] Jin et al., NSFnets (Navier-Stokes flow nets): Physics-informed neural networks for the incompressible Navier-Stokes equations, Journal of Computational Physics (2021)"}, {"review_id": "G8pQ7Ahqsv", "reviewer": "Reviewer_5Jte", "summary": "The 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.", "questions": "1. 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", "limitations": "The limitations have been adequately addressed in the paper.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**R: Computational stability and gradient vanishing/exploding problems**  \nA: We have tested on the (5+1)-d system to check whether our model can handle higher dimension PDEs (please see section G.1 in the supplementary material). Additionally, we acknowledge the potential numerical instability introduced by the outer product operation as the dimensionality ($d$) increases. To address this, a promising approach for solving higher-dimensional PDEs involves integrating multiple body networks into one, thereby reducing the multiplications during feature merging:  \n$U(x_1, x_2, …, x_d) = \\sum\\prod f(x_1, x_2)f(x_3, x_4), …, f(x_{d-1}, x_d)$  \nWe view the exploration of SPINN's performance on high-dimensional PDEs as an intriguing avenue for future research. Thank you for your suggestion.\n***\n**R: Size of Figure 6**  \nA: You can see the magnified version of Figure 6 in our attached PDF file (Figure 1). We will also update the figure in our revised version of the paper. Thank you for your advice.\n***  \n**R: Choice of the PDEs**  \nA: The Helmholtz and Klein-Gordon equations pose notable challenges for conventional PINNs, even in 2-dimensional cases [1]. Moreover, we intentionally introduced high-frequency components into the reference solution of the Helmholtz equation to examine the strength of SPINN’s capability to approximate highly complex solutions. Due to neural networks' inherent spectral bias [2], both the standard PINN and the PINN with a modified MLP failed in this case, regardless of the number of collocation points used. During the rebuttal period, we tested our model on the chaotic (1+1)-d Kuramoto-Sivashinky equation following the causal PINNs’ experiment [3]. Within the temporal domain [0, 0.4], our model achieved a relative L2 error of 3.81e-2 without using a causal loss, and the training speed was 15 times faster than the causal PINNs. You can see the visualized result in our attached PDF, Figure 2.\n***\n**R: Does SPINN architecture has any optimization problems?**  \nA: Since our work addresses the scalability issues of PINNs in priority, SPINN inherits the difficulty of optimization in PINNs’ context. However, our experimental results, when compared with prior methods that focused on PINNs optimization (such as modified MLP, and causal PINNs), suggest that the adoption of large-scale collocation points can serve as an effective means of optimizing PINNs. Standard PINNs were infeasible to employ such large-scale collocation points due to the computational issue. Regarding optimization, we have conducted some experiments to explore the use of L-BFGS when training SPINN. Please see section G.2 in the supplementary materials. 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.\n***\n**R: Why PINNs perform better than SPINN for a small number of collocation points for the diffusion equation?**  \nA: Among the PDEs in our experiments, the diffusion equation was the simplest case where the solution is a superposition of three Gaussians. It seems that even standard PINNs trained with a small amount of collocation points can accurately predict the relatively simple solution function. Because of the structural collocation points, SPINN with small training inputs seems to suffer in some cases. However, we want to note that SPINN eventually finds a more accurate solution with more collocation points ($>10^6$). Training standard PINNs in this scale would require much more training time and memory space.\n\n  \n[1] Wang et al., \"Understanding and mitigating gradient pathologies in physics-informed neural networks\", SIAM Journal on Scientific Computing (2021)  \n[2] Rahaman et al., \"On the spectral bias of neural networks\", ICML (2019)  \n[3] Wang et al., \"Respecting causality is all you need for training physics-informed neural networks\", arXiv (2022)\n\n\n\n\n"}, {"review_id": "ejBN6r7KlG", "reviewer": "Reviewer_kaJa", "summary": "The 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.", "questions": "See above.", "limitations": "Yes.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The 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", "weaknesses": "While 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**R: Discussion of why SPINN delivers more improvements**  \nA: Our performance enhancement mainly stems from the adoption of large-scale collocation points. This is particularly effective when dealing with PDEs where the solution is complex, such as a solution that contains high-frequency components (as demonstrated in our Helmholtz equation) or a solution that shows turbulent and chaotic behaviors (Navier-Stokes equation). Moreover, as the input dimension increases, our model can derive substantial benefits from the adoption of large-scale collocation points, effectively spanning the high-dimensional input space.\n***\n**R: Higher-order derivatives for section 4.3**  \nA: Thank you for your suggestion. We could easily derive the formula for higher-order derivates of Equation 7 and found that as the order increases, the overall pattern in the equation does not change a lot. You can just replace the derivative $df/dx$ with its higher-order counterpart $d^nf/dx^n$. As you mentioned, since computing the derivatives dominates the PINNs training, we will provide a detailed explanation of the higher-order case in this section.\n***\n**R: Other metrics**  \nA: Based on your and reviewer dosJ’s suggestion, we have introduced two additional metrics – RMSE and PINN loss value – as detailed in the attached PDF file. We will provide the complete results in the revised version of the paper if permitted. Thank you for your valuable comment.\n***\n**R: Clarification in the expressions: “complexity is linear in network propagations”**  \nA: Thank you for pointing this out. We will carefully examine the entire text for issues and make the necessary corrections.\n"}, {"review_id": "SfQ9HrhPMe", "reviewer": "Reviewer_Zd5Z", "summary": "This 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.", "questions": "- 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?", "limitations": "The discussion of the limitations in this paper is comprehensive. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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. ", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**R: The primary concern of PINNs is the optimization issue. Therefore, the contribution is moderate.**  \nA: We believe that scalability is also an important issue to tackle in PINNs’ literature. In our Navier-Stokes experiment, where the solution shows complex and turbulent behavior, previous studies [1, 2] have demonstrated the necessity for a large number of collocation points. Particularly in the causal PINNs [1], they had to adopt parallelized training in multi-GPUs and sophisticated software optimization to employ large-scale collocation points. However, without bells and whistles, our proposed method effectively mitigates the scalability issue. Moreover, our experimental results, when compared with prior methods that focused on PINNs' optimization (such as modified MLP, and causal PINNs), suggest that the adoption of large-scale collocation points can serve as an effective means of optimizing PINNs. Standard PINNs were infeasible to employ such large-scale collocation points due to the computational issue. Also, since the seq2seq (time-marching) and the modified MLP are orthogonal to our proposed method, we additionally adopted these techniques to SPINN and observed their positive impact. Other optimization methods like adaptive activation function [3], and learning rate annealing [4] are also applicable to SPINN and we believe that investigating this would be a valuable direction for future study.\n***\n**R: Related works on scalability for PINNs**  \nA: Thank you for the references. We will add an additional paragraph in the related works section. We will also try to run the provided baselines on our PDE settings and compare them with SPINN’s result for future work.\n***\n**R: Higher-order PDEs**  \nA: During the rebuttal period, we tested our model on the chaotic (1+1)-d Kuramoto-Sivashinky equation following the causal PINNs’ experiment [1]. Within the temporal domain [0, 0.4], our model achieved a relative L2 error of 3.81e-2 without using a causal loss, and the training speed was 15 times faster than the causal PINNs. You can see the visualized result in our attached PDF, Figure 2. We also want to notify the reviewer that the PDE we used in (2+1)-d Navier-Stokes equation is a third-order. The vorticity form of the equation is a second-order and our model's prediction is a velocity, in which additional curl operation needs to be taken to obtain the vorticity.\n***\n**R: In the theorem, rank should be sufficiently large. What about in practice?**  \nA: Empirically we showed that SPINN w/ 128 output units (meaning the rank, not too large) can more accurately approximate the (2+1)-d Navier-Stokes equations than any other existing PINNs. As you can see in Figure 7, using a larger rank does not show a meaningful performance increase. Knowing the sufficient number of output units (the rank in SPINN) beforehand is also challenging since we do not know the ranks of the solution functions in general. It is a similar question as 'how many layers and units are sufficient to solve a certain PDE in PINN?'. We believe it is an important question and very challenging, and we plan to further study as a future work.\n***\n**R: Random sampling may not be an optimal choice. Any other ways to improve in the context of SPINN?**  \nA: One straightforward way is to sample more points where the loss value is large. When solving time-dependent PDE, we can sample more temporal coordinates near the initial point (t=0) at the beginning of the training and progressively extend the sampling to later times as the training proceeds. This would guide the model to learn the causal nature of the solution.\n***\n**R: Hyper-parameter tuning**  \nA: Our hyper-parameters are tuned by grid search. We mentioned the learning-rate annealing because it helps the model to search for the best weight balances of the loss terms.  \n  \n[1] Wang et al., \"Respecting causality is all you need for training physics-informed neural networks\", arXiv (2022)  \n[2] Sankaran et al., \"On the impact of larger batch size in the training of Physics Informed Neural Networks\", The Symbiosis of Deep Learning and Differential Equations II (2022)  \n[3] Jagtap et al., \"Adaptive activation functions accelerate convergence in deep and physics-informed neural networks\", Journal of Computational Physics (2020)  \n[4] Wang et al., \"Understanding and mitigating gradient pathologies in physics-informed neural networks\", SIAM Journal on Scientific Computing (2021)\n"}, {"review_id": "urLvtt6G1N", "reviewer": "Reviewer_dosJ", "summary": "This 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.", "questions": "1. 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.", "limitations": "The authors have talked their limitation in the paper", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "Although 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**R: It is unfair to test SPINN on solutions that align with the variable separation form.**  \nA: We tested on manufactured solutions because this is the most conventional way of making an exact reference solution. As you can see in the diffusion and (2+1)-d Navier-Stokes equation, SPINN is still effective in solving such equations where the solution does not align with the variable separation. During the rebuttal period, we conducted additional experiments on the flow mixing problem. This is a (2+1)-d PDE of two fluids being mixed at the interface. We followed the problem settings used in CAN-PINNs [1]. It has an analytical solution that does not explicitly decompose into each input variable:  \n$u(t, x, y)=-\\tanh(\\frac{y}{2}\\cos(\\omega t)-\\frac{x}{2}\\sin(\\omega t))$, where  \n$\\omega=\\frac{1}{r}\\frac{\\tanh(r)}{0.385\\cosh^2(r)}, r=\\sqrt{x^2+y^2}.$  \nSince CAN-PINN reported mean-squared errors on a box plot without specifying the exact numerical value, we compared against their rough result. You can see the visualization in Figure 3 of the attached PDF.  \nStandard PINNs: 1.08e-2  \nCAN-PINN: 2e-5~3e-5  \nSPINN: 2.95e-5  \nWe also added a result of (1+1)-d chaotic Kuramoto-Sivashinsky equation following the causal PINNs’ experiment [2]. Within the temporal domain [0, 0.4], our model achieved a relative L2 error of 3.81e-2 without using a causal loss, and the training speed was 15 times faster than the causal PINNs. You can see the visualized result in our attached PDF, Figure 2. We then have total of 8 PDEs:  \n*separable form*: Helmholtz, (2+1)-d Klein-Gordon, (3+1)-d Klein-Gordon, (3+1)-d Navier-Stokes  \n*non-separable form*: diffusion, (2+1)-d Navier-Stokes, flow mixing, Kuramoto-Sivashinsky  \nMoreover, we will add a paragraph in the experiment section to indicate to the readers that some examples align with the variable separation form or not.\n***\n**R: SPINN on complex boundary conditions**  \nA: Although our current work showed examples of a rectangular domain, it is premature to conclude that SPINN *cannot* be applied to arbitrary geometries. Inspired by PhyGeoNet [3] and Geo-FNO [4], we can apply an additional operation after the SPINN’s feature merging to map a rectangular mesh to an arbitrary physical mesh. We are eager to explore combining these works with SPINN as a future study. We also believe that current SPINN can be applied to circular or spherical boundaries. The exploration of SPINN across a range of coordinate systems greatly intrigues us; however, we could not find a suitable example involving a spherical coordinate-defined PDE with a non-decomposable solution along each axis for testing purposes during the rebuttal phase. We will conduct experiments whenever we find such examples. Thank you for your suggestion.\n***\n**R: Other metrics**  \nA: Based on your and reviewer KaJa’s suggestion, we have introduced two additional metrics – RMSE and PINN loss value – as detailed in the attached PDF file. We will provide the complete results in the revised version of the paper if permitted. Thank you for your valuable comment.\n***\n**R: Dimensionality reduction through variable separation**  \nA: This is a very good idea. We also believe that optimizing the each separated function would greatly reduce the computational cost since it does not require any operations on a high dimensional space. At first thought, if we were able to decompose the original high-dim PDE into a set of ODEs, we could directly optimize each function using the ODE of corresponding dimension. There still exists some challenges that needs to be solved such as 'how can we actually decompose a PDE into ODEs?' or 'can such decomposition applied universally?'. There may be a completely different approach we have not thought of. We find this avenue of research very interesting and are enthusiastic about exchanging further insights with you.\n  \n[1] Chiu et al., \"CAN-PINN: A fast physics-informed neural network based on coupled-automatic–numerical differentiation method\", Computer Methods in Applied Mechanics and Engineering 395 (2022)  \n[2] Wang et al., \"Respecting causality is all you need for training physics-informed neural networks\", arXiv (2022)  \n[3] Gao et al., \"PhyGeoNet: Physics-Informed Geometry-Adaptive Convolutional Neural Networks for Solving Parameterized Steady-State PDEs on Irregular Domain\", Journal of Computational Physics (2021)  \n[4] Li et al., \"Fourier Neural Operator with Learned Deformations for PDEs on General Geometries\", arXiv (2022)"}]

Author rebuttal:
## For all reviewers  

We thank the reviewers for taking their time to suggest all the invaluable comments and constructive feedback.
**We attached a PDF file** to show tables and visualizations of additional experiments conducted during the rebuttal period.
Please see our responses below.

Meta-review: In this paper, a new class of architecture of PINN (Physics-informed Neural Network) called Separable PINN (SPINN) is proposed as an alternative to the standard coordinate-based MLP architecture. This significantly reduces the computational complexity of PINNs, enables the use of numerous collocation points, and allows for improved solution accuracy. The proposed architecture also provides a universal approximation theorem. Through experiments, it is demonstrated that SPINN is substantially faster during training than the original PINN while maintaining accuracy. During the rebuttal phase, the authors appropriately addressed reviewers' questions and comments, including additional experiments. While the practicality of PINNs, particularly in high-dimensional problems, is currently limited, this approach achieves remarkable efficiency improvements, making a significant contribution to the advancement of PINNs. This can be regarded as an excellent research achievement.
Decision: Accept (spotlight)