SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Junwoo Cho $^{1*}$
Seok-Bae Yun
Seungtae  $\mathbf{N}\mathbf{a}\mathbf{m}^{1*}$
Youngjoon Hong<sup>3</sup>
Hyunmo Yang
Eunbyung Park1,4†
$^{2}$ Department of Mathematics, Sungkyunkwan University
$^{3}$ Department of Mathematical Sciences, KAIST
$^{4}$ Department of Electrical and Computer Engineering, Sungkyunkwan University
Physics-informed neural networks (PINNs) have recently emerged as promising data-driven PDE solvers showing encouraging results on various PDEs. However, there is a fundamental limitation of training PINNs to solve multi-dimensional PDEs and approximate highly complex solution functions. The number of training points (collocation points) required on these challenging PDEs grows substantially, but it is severely limited due to the expensive computational costs and heavy memory overhead. To overcome this issue, we propose a network architecture and training algorithm for PINNs. The proposed method, separable PINN (SPINN), operates on a per-axis basis to significantly reduce the number of network propagations in multi-dimensional PDEs unlike point-wise processing in conventional PINNs. We also propose using forward-mode automatic differentiation to reduce the computational cost of computing PDE residuals, enabling a large number of collocation points  $(>10^{7})$  on a single commodity GPU. The experimental results show drastically reduced computational costs  $(62\times$  in wall-clock time,  $1,394\times$  in FLOPs given the same number of collocation points) in multi-dimensional PDEs while achieving better accuracy. Furthermore, we present that SPINN can solve a chaotic  $(2 + 1)$ -d Navier-Stokes equation significantly faster than the best-performing prior method (9 minutes vs 10 hours in a single GPU), maintaining accuracy. Finally, we showcase that SPINN can accurately obtain the solution of a highly nonlinear and multi-dimensional PDE, a  $(3 + 1)$ -d Navier-Stokes equation. For visualized results and code, please see https://jwcho5576.github.io/spinn.github.io/.
Solving partial differential equations (PDEs) has been a long-standing problem in various science and engineering domains. Finding analytic solutions requires in-depth expertise and is often infeasible in many useful and important PDEs [10]. Hence, numerical approximation methods to solutions have been extensively studied [17], e.g., spectral methods [3], finite volume methods (FVM) [9], finite difference method (FDM) [39], and finite element methods (FEM) [47]. While successful, classical methods have several limitations, such as expensive computational costs, requiring sophisticated techniques to support multi-physics and multi-scale systems, and the curse of dimensionality in high dimensional PDEs.
With the vast increases in computational power and methodological advances in machine learning, researchers have explored data-driven and learning-based methods [4, 24, 30, 38]. Among the promising methods, physics-informed neural networks (PINNs) have recently emerged as new data-driven PDE solvers for both forward and inverse problems [34]. PINNs employ neural networks and
Figure 1: Training speed (w/ a single GPU) of our model compared to the causal PINN [43] in  $(2 + 1)$ -d Navier-Stokes equation of time interval [0, 0.1].
gradient-based optimization algorithms to represent and obtain the solutions, leveraging automatic differentiation to enforce the physical constraints of underlying PDE. It has enjoyed great success in various forward and inverse problems thanks to its numerous benefits, such as flexibility in handling a wide range of forward and inverse problems, mesh-free solutions, and not requiring observational data, hence, unsupervised training.
Despite its advantages and promising results, there is a fundamental limitation of training PINN to solve multi-dimensional PDEs and approximate very complex solution functions. It primarily stems from using coordinate-based MLP architectures to represent the solution function, which takes input coordinates and outputs corresponding solution quantities. For each training point, computing PDE residual loss involves multiple forward and backward propagations, and the number of training points (collocation points) required to solve multi-dimensional PDEs and obtain more accurate solutions grows substantially. The situation deteriorates as the dimensionality of the PDE or the solution's complexity increases.
Recent studies have presented empirical evidence showing that choosing a larger batch size (i.e., a large number of collocation points) in training PINNs leads to enhanced precision [34, 36, 37, 43]. Furthermore, more collocation points also accelerate the convergence speed due to the scaling characteristic between the batch size and the learning rate (the larger the batch size, the higher the learning rate) [12, 27, 40]. As long as the computational resources allow, we have the flexibility to utilize a substantial number of training points in PINNs since they can be continuously sampled from the input domain in an unsupervised training manner.
We propose a novel PINN architecture, separable PINN (SPINN), which utilizes forward-mode* automatic differentiation (AD) to enable a large number of collocation points ( $>10^{7}$  in a single GPU), reducing the computational cost of solving multi-dimensional PDEs. Instead of feeding every multi-dimensional coordinate into a single MLP, we use separated sub-networks, in which each sub-network takes independent one-dimensional coordinates as input. The final output is generated by an aggregation module such as simple outer product and element-wise summation where the predicted solution can be interpreted by low-rank tensor approximation [25] (Fig. 4b). The suggested architecture obviates the need to query every multi-dimensional coordinate input pair, exponentially reducing the number of network propagations to generate a solution,  $\mathcal{O}(N^d)\rightarrow \mathcal{O}(Nd)$ , where  $N$  is the resolution of the solution for each dimension, and  $d$  is the dimension of the system.
We have conducted comprehensive experiments on the representative PDEs to show the effectiveness of the suggested method. The experimental results demonstrate that the training runtime of the conventional PINN increase linearly with the number of collocation points, while the proposed model shows logarithmic growths. This allows SPINN to accommodate orders of magnitude larger number of collocation points in a single batch during training. We also show that given the same number of training points, SPINN improves wall-clock training time up to by  $62 \times$  on commodity GPUs and FLOPs up to by  $1,394 \times$  while achieving better accuracy. Furthermore, with large-scale collocation points, SPINN can solve a turbulent Navier-Stokes equation much faster than the state-of-the-art PINN method [43] (9 minutes vs 10 hours in a single GPU) without bells and whistles, such as causal
$^{*}$  a.k.a. forward accumulation mode or tangent linear mode.
inductive bias in the loss function (Fig. 1). Our experimental results of the Navier-Stokes equation show that SPINN can solve highly nonlinear PDEs and is sufficiently expressive to represent complex functions. This is further supported by the provided theoretical result, potentiating the use of our method for more challenging and various PDEs.
Physics-informed neural networks. Physics-Informed Neural Networks (PINNs) [34] have received great attention as a promising learning-based PDE solver. Given the underlying PDE and initial, boundary conditions embedded in a loss function, a coordinate-based neural network is trained to approximate the desired solution. Since its inception, many techniques have been studied to improve training PINNs for more challenging PDE systems [26, 43, 44], or to accelerate the training speed [20, 37]. Our method is orthogonal to most of the previously suggested techniques above and improves PINNs' training from a computational perspective.
The effect of collocation points. The residual loss of PINNs is calculated by Monte Carlo integration, not an exact definite integral. Therefore, it inherits a core property of Monte Carlo methods [15]: the impact of the number of sampled points. The importance of sampling strategy and the number of collocation points in training PINNs has been highlighted by recent works [7, 23, 36, 43]. Especially, Sankaran et al. [36] empirically found that training with a large number of collocation points is unconditionally favorable for PINNs in terms of accuracy and convergence speed. Another line of research established a theoretical upper bound of PINN's statistical error with respect to the number of collocation points [23]. It showed that a larger number of collocation points are required to use a bigger network size for training. Our work builds on this evidence to bring PINNs out in more practical scenarios and overcome their limitation.
**Derivative computations in scientific machine learning.** Embedding physical constraints in the loss function is widely used in scientific machine learning, and computing the derivatives is an essential process for this formulation. Several works employed numerical differentiation [33, 35, 42], or hybrid approach [5, 37, 46] with AD for the calculation, since numerical methods such as finite differences do not need to back-propagate through the network. However, they are still burdened by a computational complexity of  $\mathcal{O}(N^d)$ , thereby limiting them to handle large-scale collocation points or meshes. Furthermore, numerical differentiation has truncation errors depending on the step size. Employing Taylor-mode AD [13] in training PINNs was introduced by causal PINN [43] to handle high-order PDEs such as Kuramoto-Sivashinsky equation [28]. To the best of our knowledge, the proposed method is the first approach to leverage forward-mode AD in training PINNs, which is fully applicable to both time-dependent and independent PDEs and does not incur any truncation errors.
Multiple MLP networks. Employing multiple MLPs for PINNs has been introduced by several works to utilize parallelized training [19, 21, 32]. They share the same concept of dividing the entire spatio-temporal domain and training multiple individual MLPs on each sub-domain. Although these methods showed promising results, they still suffer from the fundamental problem of heavy computation as the number of collocation points increases. While these methods decompose input domains, and each small MLP is used to cover a particular sub-domain, we decompose input dimensions and solve PDEs over the entire domain cooperated by all separated MLPs. In terms of the model architecture and function representation, our work is also related to NAM [1], Haghighat et al. [14], and CoordX [31]. NAM [1] suggested separated network architectures and inputs, but only for achieving the interpretability of the model's prediction in multi-task learning. Haghighat et al. [14] used multiple MLPs to individually predict each component of the output vector. CoordX's [31] primary purpose is to reconstruct natural signals, such as images or 3D shapes. Hence, they are not motivated to improve the efficiency of computing higher-order gradients. In addition, they had to use additional layers after the feature merging step, which made their model more computationally expensive than ours. Furthermore, in neural fields [45], there is an explicit limitation in the number of ground truth data points (e.g., the number of pixels in an image). Therefore, CoordX cannot fully maximize the advantage of enabling a large number of input coordinates. We focus on solving PDEs and carefully devise the architecture to exploit forward-mode AD to efficiently compute PDE residual losses.
Figure 2: Simple neural networks with different input and output dimensions. To compute  $\frac{\partial y}{\partial x}$ , (a) requires one forward pass using forward-mode AD or three backward passes using reverse-mode AD, (b) requires three forward passes using forward-mode AD or one backward pass using reverse-mode AD.
Figure 3: An illustrative example of separated approach vs non-separated approach when  $f^{(\theta)} : \mathbb{R}^d \to \mathbb{R}$  (an example case of  $N = 3$ ,  $d = 3$  is shown above). The number of JVP (forward-mode) evaluations (propagations) of computing the Jacobian for separated approach (a) is  $Nd$ , while the number of VJP (reverse-mode) evaluations for non-separated approach (b) is  $N^d$ .
3 Preliminaries: Forward/Reverse-mode AD
For the completeness of this paper, we start by briefly introducing the two types of AD and how Jacobian matrices are evaluated. For clarity, we will follow the notations used in [2] and [13]. Suppose our function  $f:\mathbb{R}^n\to \mathbb{R}^m$  is a two-layer MLP with 'tanh' activation. The left-hand side of Tab. 1 demonstrates a single forward trace of  $f$ . To obtain a  $m\times n$  Jacobian matrix  $\mathbb{J}_f$  in forward-mode, we compute the Jacobian-vector product (JVP),
Table 1: An example of forward and reverse-mode AD in a two-layers tanh MLP. Here  $v_{0}$  denotes the input variable,  $v_{k}$  the primals,  $\dot{v}_{k}$  the tangents,  $\bar{v}_{k}$  the adjoints, and  $W_{1}, W_{2}$  the weight matrices. Biases are omitted for brevity.
Forward primal trace	Forward tangent trace	Backward adjoint trace
v0=x	\(\dot{v}_0=\dot{x}\)	\(\bar{x}=\bar{v}_0\)
v1=W1·v0	\(\dot{v}_1=W_1\cdot\dot{v}_0\)	\(\bar{v}_0=\bar{v}_1\cdot\frac{\partial v_1}{\partial v_0}=\bar{v}_1\cdot W_1\)
v2=tanh(v1)	\(\dot{v}_2=\tanh^{\prime}(v_1)\circ\dot{v}_1\)	\(\bar{v}_1=\bar{v}_2\cdot\frac{\partial v_2}{\partial v_1}=\bar{v}_2\circ\tanh^{\prime}(v_1)\)
v3=W2·v2	\(\dot{v}_3=W_2\cdot\dot{v}_2\)	\(\bar{v}_2=\bar{v}_3\cdot\frac{\partial v_3}{\partial v_2}=\bar{v}_3\cdot W_2\)
y=v3	\(\dot{y}=\dot{v}_3\)	\(\bar{v}_3=\bar{y}\)
$$
\mathbb {J} _ {f} r = \left[ \begin{array}{c c c} \frac {\partial y _ {1}}{\partial x _ {1}} & \dots & \frac {\partial y _ {1}}{\partial x _ {n}} \\ \vdots & \ddots & \vdots \\ \frac {\partial y _ {m}}{\partial x _ {1}} & \dots & \frac {\partial y _ {m}}{\partial x _ {n}} \end{array} \right] \left[ \begin{array}{c} \frac {\partial x _ {1}}{\partial x _ {i}} \\ \vdots \\ \frac {\partial x _ {n}}{\partial x _ {i}} \end{array} \right], \tag {1}
$$
for  $i \in \{1, \dots, n\}$ . The forward-mode AD is a one-phase process: while tracing primals (intermediate values)  $v_{k}$ , it continues to evaluate and accumulate their tangents  $\dot{v}_{k} = \partial v_{k} / \partial x_{i}$  (the middle column of Tab. 1). This is equivalent to decomposing one large JVP into a series of JVPs by the chain rule and computing them from right to left. A run of JVP with the initial tangents  $v_{0}$  as the first column vector of an identity matrix  $I_{n}$  gives the first column of  $\mathbb{J}_f$ . Thus, the full Jacobian can be obtained in  $n$  forward passes.
On the other hand, the reverse-mode AD computes vector-Jacobian product (VJP):
$$
r ^ {\top} \mathbb {J} _ {f} = \left[ \begin{array}{l l l} \frac {\partial y _ {j}}{\partial y _ {1}} & \dots & \frac {\partial y _ {j}}{\partial y _ {m}} \end{array} \right] \left[ \begin{array}{c c c} \frac {\partial y _ {1}}{\partial x _ {1}} & \dots & \frac {\partial y _ {1}}{\partial x _ {n}} \\ \vdots & \ddots & \vdots \\ \frac {\partial y _ {m}}{\partial x _ {1}} & \dots & \frac {\partial y _ {m}}{\partial x _ {n}} \end{array} \right], \tag {2}
$$
for  $j \in \{1, \dots, m\}$ , which is the reverse-order operation of JVP. This is a two-phase process. The first phase corresponds to forward propagation, storing all the primals,  $v_{k}$ , and recording the elementary operations in the computational graph. In the second phase, the derivatives are computed by accumulating the adjoints  $\bar{v}_{k} = \partial y_{j} / \partial v_{k}$  (the right-hand side of Tab. 1). Since VJP builds one row of a Jacobian at a time, it takes  $m$  evaluations to obtain the full Jacobian. To sum up, the forward-mode is more efficient for a tall Jacobian ( $m > n$ ), while the reverse-mode is better suited for a wide Jacobian ( $n > m$ ). Fig. 2 shows an illustrative example and please refer to Baydin et al. [2] for more details.
Figure 4: (a) SPINN architecture in a 3-dimensional system. To solve a  $d$ -dimensional PDE, our model requires  $d$  body MLP networks, each of which takes individual scalar coordinate values as input and gives  $r$ -dimensional feature vector. The final output is obtained by element-wise product and summation. (b) Construction process of the entire discretized solution tensor when the inputs are given in batches. Each outer product between the column vectors  $F_{:,j,i}$  from the feature tensor  $F$  constructs a rank-1 tensor and summing all the  $r$  tensors gives a rank- $r$  tensor. The output tensor of SPINN can be interpreted as a low-rank decomposed representation of a solution.
4 Separable PINN
4.1 Forward-Mode AD with Separated Functions
We demonstrate that leveraging forward-mode AD and separating the function into multiple functions with respect to input axes can significantly reduce the cost of computing Jacobian. In the proposed separated approach (Fig. 3a), we first sample  $N$  one-dimensional coordinates on each of  $d$  axes, which makes a total of  $Nd$  batch size. Next, these coordinates are fed into  $d$  individual functions. Let  $f$  be a function which takes coordinate as input to produce feature representation and we denote the number of operations of  $f$  as  $\mathrm{ops}(f)$ . Then, a feature merging function  $h$  is used to construct the solution of the entire  $N^d$  discretized points. The amount of computations for AD is known to be  $2\sim 3$  times more expensive than the forward propagation [2, 13]. According to Fig. 3 and with the scale constants  $c_{f}, c_{h} \in [2,3]$  we can approximate the total number of operations to compute the Jacobian matrix of the proposed separated approach (Fig. 3a).
$$
\mathcal {C} _ {\text {s e p}} = N d c _ {f} \operatorname {o p s} (f) + N ^ {d} c _ {h} \operatorname {o p s} (h). \tag {3}
$$
For a non-separated approach (Fig. 3b),
$$
\mathcal {C} _ {\text {n o n - s e p}} = N ^ {d} c _ {f} \mathsf {o p s} (f), \tag {4}
$$
If we can make  $\mathrm{ops}(h)$  sufficiently small, then the ratio  $\mathcal{C}_{\mathrm{sep}} / \mathcal{C}_{\mathrm{non - sep}}$  becomes  $\frac{Nd}{Nd} \ll 1$ , quickly converging to 0 as  $d$  and  $N$  increases. Our separated approach has linear complexity with respect to  $N$  in network propagations, implying it can obtain more accurate solutions (high-resolution) efficiently.
4.2 Network Architecture
Fig. 4a illustrates the overall SPINN architecture, parameterizing multiple separated functions with neural networks. SPINN consists of  $d$  body-networks (MLPs), each of which takes an individual 1-dimensional coordinate component as an input. Each body-network  $f^{(\theta_i)}: \mathbb{R} \to \mathbb{R}^r$  (parameterized by  $\theta_i$ ) is a vector-valued function which transforms the coordinates of  $i$ -th axis into a  $r$ -dimensional feature representation. The final prediction is computed by feature merging:
$$
\hat {u} \left(x _ {1}, x _ {2}, \dots , x _ {d}\right) = \sum_ {j = 1} ^ {r} \prod_ {i = 1} ^ {d} f _ {j} ^ {\left(\theta_ {i}\right)} \left(x _ {i}\right) \tag {5}
$$
where  $\hat{u}:\mathbb{R}^d\to \mathbb{R}$  is the predicted solution function,  $x_{i}\in \mathbb{R}$  is a coordinate of  $i$  -th axis, and  $f_{j}^{(\theta_{i})}$  denotes the  $j$  -th element of  $f^{(\theta_i)}$ . We used 'tanh' activation function throughout the paper. As shown in Eq. 5, the feature merging operation is a simple product (II) and summation (Sigma) which
corresponds to the merging function  $h$  described in Eq. 3. Due to its simplicity,  $h$  operations are much cheaper than operations in MLP layers (i.e.,  $\mathrm{ops}(h) \ll \mathrm{ops}(f)$  in Eq. 3). Note that SPINN can also approximate any  $m$ -dimensional vector functions  $\hat{u} : \mathbb{R}^d \to \mathbb{R}^m$  by using a larger output feature size (see section D.4 in the appendix for details).
The collocation points of our model and conventional PINNs have a distinct difference (Fig. 5). Both are uniformly evaluated on a  $d$ -dimensional hypercube, but collocation points of SPINN form a lattice-like structure, which we call as factorizable coordinates. In SPINN, 1-dimensional input points from each axis are randomly sampled, and  $d$ -dimensional points are generated via the cartesian product of the point sets from each axis. On the other hand, non-factorizable coordinates are randomly sampled points without any structure. Factorizable coordinates with our separated MLP architecture enable us to evaluate functions on dense  $(N^d)$  collocation points with a small number  $(Nd)$  of input points.
(a) Factorizable coordinates
Figure 5: An illustrative 2-dimensional example of (a) factorizable and (b) non-factorizable coordinates. (a) has a lattice-like structure, where SPINN can be evaluated on more dense collocation points with fewer input points. Conventional PINNs sample non-factorizable coordinates which do not have any structures.
(b) Non-factorizable coordinates
In practice, the input coordinates are given in a batch during training and inference. Assume that  $N$  input coordinates (training points) are sampled from each axis. Note that the sampling resolutions for each axis need not be the same. The input coordinates  $X \in \mathbb{R}^{N \times d}$  is now a matrix. The batchified form of feature representation  $F \in \mathbb{R}^{N \times r \times d}$  and Eq. 5 now becomes
$$
\hat {U} \left(X _ {:, 1}, X _ {:, 2}, \dots , X _ {:, d}\right) = \sum_ {j = 1} ^ {r} \bigotimes_ {i = 1} ^ {d} F _ {:, j, i}, \tag {6}
$$
where  $\hat{U} \in \mathbb{R}^{N \times N \times \ldots \times N}$  is the discretized solution tensor,  $\otimes$  denotes outer product,  $F_{\vdots ,\cdot ,i} \in \mathbb{R}^{N \times r}$  is an  $i$ -th frontal slice matrix of tensor  $F$ , and  $F_{\vdots ,j,i} \in \mathbb{R}^N$  is the  $j$ -th column of the matrix  $F_{\vdots ,\cdot ,i}$ . Fig. 4b shows an illustrative procedure of Eq. 6. Due to its structural input points and outer products between feature vectors, SPINN's solution approximation can be viewed as a low-rank tensor decomposition where the feature size  $r$  is the rank of the reconstructed tensor. Among many decomposition methods, SPINN corresponds to CP-decomposition [16], which approximates a tensor by finite summation of rank-1 tensors. While traditional methods use iterative methods such as alternating least-squares (ALS) or alternating slicewise diagonalization (ASD) [22] to directly fit the decomposed vectors, we train neural networks to learn the decomposed vector representation and approximate the solution functions in continuous input domains. This, in turn, allows for the calculation of derivatives with respect to arbitrary input coordinates.
4.3 Gradient Computation of SPINN
In this section, we show that the number of JVP (forward-mode AD) evaluations for computing the full gradient of SPINN  $(\nabla \hat{u}(x))$  is  $Nd$ , where  $N$  is the number of coordinates sampled from each axis and  $d$  is the input dimension. According to Eq. 5, the  $i$ -th element of  $\nabla \hat{u}(x)$  is:
$$
\frac {\partial \hat {u}}{\partial x _ {i}} = \sum_ {j = 1} ^ {r} f _ {j} ^ {\left(\theta_ {1}\right)} \left(x _ {1}\right) f _ {j} ^ {\left(\theta_ {2}\right)} \left(x _ {2}\right) \dots \frac {\partial f _ {j} ^ {\left(\theta_ {i}\right)} \left(x _ {i}\right)}{\partial x _ {i}} \dots f _ {j} ^ {\left(\theta_ {d}\right)} \left(x _ {d}\right). \tag {7}
$$