Separable Physics-Informed Neural Networks (2306.15969v4)

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

• The paper introduces SPINNs to reduce computational costs by employing a separable, per-axis network design.
• It leverages forward-mode automatic differentiation to efficiently process over 10^7 collocation points on a single GPU.
• Experimental results show dramatic speedups, with training times and FLOPs reduced by up to 62× in solving complex PDEs.

Separable Physics-Informed Neural Networks

The paper "Separable Physics-Informed Neural Networks" proposes a novel approach to enhance the computational efficiency and accuracy of physics-informed neural networks (PINNs) in solving multi-dimensional partial differential equations (PDEs). While PINNs hold promise as data-driven solvers for PDEs across various scientific domains, they face challenges in handling high-dimensional PDEs and complex solution functions due to the substantial number of training points required. This paper introduces Separable Physics-Informed Neural Networks (SPINNs) as a solution to these challenges, utilizing a per-axis approach to significantly reduce the computational burden.

Methodology

SPINNs leverage the concept of forward-mode automatic differentiation to compute PDE residuals, facilitating the use of large batches of collocation points—exceeding $10^7$—on a single commodity GPU. This approach is contrasted with traditional PINNs, which process collocation points point-wise, resulting in prohibitive computational costs and memory usage for high-dimensional problems.

The SPINN architecture is characterized by separate networks for each axis, each of which independently processes one-dimensional coordinates. An aggregation module subsequently combines these outputs via simple operations such as outer product and element-wise summation. This feature enables SPINNs to maintain a computational complexity that grows linearly with the dimensionality of the problem and the resolution of the solution, $\mathcal{O}(Nd)$, as opposed to the exponential growth observed in conventional PINNs, $\mathcal{O}(N^d)$, where $N$ represents the resolution per dimension and $d$ is the dimensionality of the system.

Experimental Results

The experimental evaluations demonstrate the efficacy of SPINNs over traditional PINNs across various PDE scenarios. For instance, in solving the chaotic (2+1)-dimensional Navier-Stokes equation, SPINNs achieved comparable accuracy while reducing computational time from 10 hours to 9 minutes on a single GPU—a dramatic improvement. Moreover, SPINNs excelled in solving complex equations like the (3+1)-dimensional Navier-Stokes equation.

SPINNs exhibit logarithmic growth in training runtime concerning the number of collocation points, thus harnessing larger numbers of training points to enhance solution accuracy. The reduction in computational costs was remarkable, with decreases in wall-clock time by $62\times$ and in FLOPs by $1,394\times$ for the same number of collocation points relative to standard PINNs.

Theoretical Implications and Future Directions

The paper establishes the universal approximation capabilities of SPINNs, supported by theoretical exposition and empirical success in various PDE experiments. However, challenges persist in extending SPINNs to handle arbitrary boundary conditions and more complex geometric domains. Addressing these limitations could involve integrating geometric mapping techniques or mesh transformations with SPINNs for enhanced flexibility in diverse practical applications.

As computational methods and hardware continue to evolve, optimizing SPINNs’ efficiency further—potentially through customized software and hardware solutions—remains a pivotal area for future research. The separable design of SPINNs opens novel avenues for efficiently solving high-dimensional and computationally intensive PDEs in scientific machine learning, with significant implications for both theoretical advancements and practical applications.

In conclusion, SPINNs represent a substantial stride toward resolving the computational challenges encountered in training physics-informed neural networks on complex multidimensional PDE systems, offering both numerical efficiency and enhanced model accuracy. This innovative approach has the potential to redefine computational techniques in solving PDEs, showcasing significant promise for future scientific and engineering developments.