PaperlayerView exact revision on arXiv

Source-linked AI summary

Separable Physics-Informed Neural Networks

Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, Eunbyung Park

arXiv:2306.15969v4cs.LGcs.AI

TL;DR

PINNs face computational difficulty solving multidimensional PDEs and representing complex solutions as collocation requirements grow. SPINN uses per-axis subnetworks with forward-mode automatic differentiation to reduce computation. At equal collocation counts, it reports substantially lower training time and FLOPs alongside better accuracy.

  • Problem

    PINNs face difficulty solving multidimensional PDEs and representing complex solutions, while required collocation points increase computational demands.

    Source: §1 · p. 2

  • Method

    SPINN uses per-axis MLPs, feature merging into a low-rank solution representation, and forward-mode AD to reduce PDE-residual computation.

    Source: §1 · p. 2, §4.1 · p. 5, §4.2 · p. 5

  • Results

    At equal collocation counts, SPINN achieves up to 62× in wall-clock training time and 1,394× in FLOPs while achieving better accuracy on multidimensional PDEs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Theory and experiments support SPINN's approximation of complex solution functions across challenging PDEs, including diffusion, Helmholtz, Klein-Gordon, and Navier–Stokes.

    Source: §4.4 · p. 7, §7 · p. 10

    SPINN tends to train better when solutions align with variable-separation form, though it remains effective on tested non-low-rank solutions.

    Source: §6 · p. 9

Abstract

from arXiv · show

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/.

1 Introduction

PINNs offer a flexible, mesh-free approach to PDE solving, but pointwise computation and growing collocation requirements make multidimensional, complex problems costly. SPINN separates computation by coordinate axis and uses forward-mode AD to reduce that burden.

1 Introduction

  • For multidimensional PDEs and complex solutions, pointwise coordinate MLPs require multiple propagations per collocation point, while the number of required points grows substantially.The burden worsens as PDE dimensionality or solution complexity increases.Source: §1 · p. 2
  • Larger collocation batches have been associated with improved PINN precision and faster convergence, motivating the use of more training points when resources permit.The cited studies connect convergence speed to batch-size scaling and learning rate.Source: §1 · p. 2
  • SPINN uses one subnetwork per coordinate axis and feature aggregation, reducing network propagations from O(N^d) to O(Nd) while representing the solution as a low-rank tensor approximation.Forward-mode automatic differentiation enables more than 10^7 collocation points on a single GPU.Source: §1 · p. 2
  • Up to 62× in wall-clock training time and 1,394× in FLOPs at equal collocation counts, with better accuracy; turbulent Navier–Stokes takes 9 minutes versus 10 hours.The runtime comparison is against the state-of-the-art PINN method [43].Source: §1 · p. 2
  • Navier–Stokes experiments indicate that SPINN can solve highly nonlinear PDEs and represent complex functions, supported by a theoretical result.The passage also describes an inductive bias in the loss function.Source: §1 · p. 3

2 Related Works

Prior work improves PINN training, studies collocation-point effects, and explores derivative computation and multiple-network designs; this paper distinguishes its approach by separating input dimensions and targeting efficient PDE residuals.

2 Related Works

  • PINNs use coordinate-based neural networks trained against PDE, initial, and boundary-condition losses; prior work targets challenging systems [26] [43] and training speed [20] [37].The paper describes its computational perspective as orthogonal to most of these techniques.Source: §2 · p. 3
  • Prior studies examine collocation sampling [23] [43], including evidence that larger batches improve accuracy and convergence and a statistical-error bound tied to point count [23].The cited bound indicates that larger networks require more collocation points.Source: §2 · p. 3
  • Numerical and hybrid differentiation approaches [33] [35] [42] can avoid network backpropagation but retain O(N^d) complexity and numerical truncation error; causal PINN [43] uses Taylor-mode AD.The authors present forward-mode AD as applicable to time-dependent and independent PDEs without truncation errors.Source: §2 · p. 3
  • Earlier multi-MLP PINNs divide the domain into subdomains [32], whereas SPINN decomposes input dimensions; NAM, Haghighat et al., and CoordX [31] target different purposes.The passage notes that CoordX [31] adds layers after feature merging and is limited by available ground-truth points.Source: §2 · p. 3

3 Preliminaries: Forward/Reverse-mode AD

Forward-mode AD builds Jacobian columns through JVPs, while reverse-mode AD builds rows through VJPs; which is more efficient depends on the Jacobian's shape.

3 Preliminaries: Forward/Reverse-mode AD

  • Forward-mode AD propagates primals and tangents together; one JVP gives a Jacobian column, so n forward passes recover the full Jacobian.Table 1 illustrates the procedure in a two-layer tanh MLP.Source: §3 · p. 4, Table 1 · p. 4
  • Reverse-mode AD records primals and the computation graph, then accumulates adjoints; m VJPs build the Jacobian row by row, favoring wide Jacobians over forward-mode's tall-Jacobian regime.The table caption identifies v0 as input, vk as primals, tangent and adjoint quantities, and W1 and W2 as weight matrices.Source: §3 · p. 4, Table 1 · p. 4
  • The SPINN schematic uses d coordinate-specific MLPs and merges their features by element-wise product and summation into a low-rank solution tensor.Each feature-vector outer product contributes a rank-1 tensor to the sum.Source: Figure 4 · p. 5

4 Separable PINN

SPINN separates coordinate processing by axis and uses forward-mode automatic differentiation to reduce the cost of computing PDE derivatives. Its architecture supports dense evaluation from structured inputs, and the paper establishes an L2 approximation property while leaving broader function spaces for future work.

4.1 Forward-Mode AD with Separated Functions

  • When feature merging is sufficiently cheap, the separated-to-nonseparated Jacobian-cost ratio tends toward zero as dimension and samples increase, while network propagations scale linearly with N.The analysis compares separated functions with a non-separated approach and highlights the cost of the feature-merging operation.Source: §4.1 · p. 5

4.2 Network Architecture

  • Each axis-specific MLP maps a one-dimensional coordinate to features, which product-and-sum merging combines into a solution; Cartesian-product inputs form a lattice.This structure lets SPINN evaluate dense collocation points using fewer input points.Source: §4.2 · pp. 5–6, Figure 5 · p. 6
  • The discretized solution can be viewed as a CP tensor decomposition, with neural networks learning its decomposed vectors over continuous input domains.This representation allows derivatives with respect to arbitrary input coordinates.Source: §4.2 · p. 6

4.3 Gradient Computation of SPINN

  • The full gradient requires Nd forward-mode JVP evaluations, reusing feature representations from the forward pass; the same approach obtains higher-order derivatives.Figure 6 compares relative error, runtime, memory, and training curves for diffusion, Helmholtz, and 3-d and 4-d Klein-Gordon experiments.Source: §4.3 · p. 6, Figure 6 · p. 7

4.4 Universal Approximation Property

  • For compact subsets of R^d, the theorem states that sufficiently large feature size allows neural networks to approximate an L2 function within arbitrary ε.The authors say repeated application extends the property to L2 functions in high-dimensional input spaces; broader function spaces remain future work, and they report accurate approximations for several challenging PDEs.Source: §4.4 · p. 7

5 Experiments

SPINN was evaluated on 3-d and 4-d forward PDEs and Navier–Stokes systems, with results reporting improved efficiency and accuracy against PINN baselines. The experiments also examine memory limits, collocation-point scaling, and rank sensitivity.

5.1 Experimental setups

  • Experiments compared SPINN with vanilla PINN on 3-d diffusion, Helmholtz, Klein-Gordon, and Navier-Stokes systems and 4-d Klein-Gordon and Navier-Stokes systems.The 3-d vanilla PINN limit was 64^3 collocation points (54^3 with modified MLP) on a 24-GB RTX3090; SPINN used up to 256^3 points.Source: §5.1 · p. 7

5.2 Results

  • Across four forward PDE systems, SPINN was more computationally efficient, and modified-MLP SPINN found the most accurate solution for every PDE.SPINN required 1,394× fewer operations than the baseline for the forward pass and first- and second-order derivatives.Source: §5.2 · p. 8
  • At 64^3 collocation points, modified-MLP SPINN was 52× faster and 29× more memory efficient than the modified-MLP baseline on diffusion.At 128^3 points, modified-MLP SPINN found the most accurate solution.Source: §5.2 · p. 8
  • On Helmholtz, SPINN obtained one order of magnitude lower relative error than PINNs with modified MLP, and was 62× faster with 29× lower memory usage at 64^3 points.The authors attribute the accuracy result to each SPINN body network learning an individual one-dimensional function, reducing the burden of representing the full 3-d function.Source: §5.2 · p. 8
  • On Klein-Gordon, SPINN had the best accuracy, runtime, and memory usage, with runtime 62× faster and memory usage 29× lower.The largest runtime and memory differences occurred on Helmholtz and Klein-Gordon, which have three second-order derivative terms, versus two for diffusion.Source: §5.2 · p. 8
  • For 4-d Klein-Gordon, SPINN handled 64^4 collocation points, 160× the baseline count, to obtain a more precise solution.The baseline was limited by memory throughput to 23^4 points (18^4 with modified MLP).Source: §5.2 · p. 8
  • On chaotic (2+1)-d Navier-Stokes, SPINN was comparably accurate without causal loss and converged 60× faster than causal PINN at equal collocation counts.Performance nearly converged at rank 128, and increasing rank did not slow training much; Figure 7 compares relative error and training speed across ranks.Source: §5.2 · pp. 8–9, Figure 7 · p. 8, Table 2 · p. 9
  • On manufactured (3+1)-d Navier-Stokes, SPINN achieved relative error 1.9e-3 in under 30 minutes on one GPU.The experiment used a Taylor-Green vortex solution and evaluated vorticity-vector-field relative errors averaged over five runs.Source: §5.2 · p. 9

6 Limitations and Future Works

The paper identifies open challenges in applying SPINN to more demanding PDEs and complex geometries. It also notes that training tends to be better when solutions align with variable separation, while non-low-rank examples remain solvable.

6 Limitations and Future Works

  • Combining SPINN with techniques for arbitrary boundary conditions on complex geometries is left to future work; the demonstrated arbitrary-domain approach is a simple L-shaped Poisson example.The paper also identifies more challenging and higher-dimensional PDEs as unresolved questions.Source: §6 · p. 9
  • SPINN tends to train better when solutions align with variable separation, though diffusion and (2+1)-d Navier-Stokes show it can still solve non-low-rank cases.The authors also note a gap between theoretical speed-up and wall-clock runtime, with hardware/software optimization such as customized CUDA kernels left as a direction.Source: §6 · p. 9

7 Conclusion

The conclusion presents separated MLPs with forward-mode automatic differentiation as a way to train PINNs with large collocation sets. It reports reduced complexity and better accuracy on three- and four-dimensional PDEs.

7 Conclusion

  • SPINN combines separated MLPs with forward-mode automatic differentiation to use large-scale collocation points in PINN training.The authors describe this as their first use of forward-mode automatic differentiation in PINNs.Source: §7 · p. 10
  • The method reduces spatial and computational complexity while achieving better accuracy on various three- and four-dimensional PDEs.Source: §7 · p. 10
  • The authors conclude that this approach opens a direction for rethinking neural-network architecture in scientific applications.Source: §7 · p. 10

A Proof of Theorem 1

The proof develops an approximation argument using tensor products of Hilbert-space basis functions and neural networks. It also establishes a basic identity for the action of a tensor product on a vector.

A Proof of Theorem 1

  • For x in H1 and y, y′ in H2, the tensor product operator satisfies (x ⊗ y)y′ = ⟨y, y′⟩x.Source: A Proof of Theorem 1 · p. 13
  • For u in L2(X × Y), Theorem 1 states that neural networks on X and Y can be chosen to approximate it to arbitrary ε.The theorem assumes compact X and Y and permits a sufficiently large rank r.Source: A Proof of Theorem 1 · pp. 14, 16
  • The proof uses product basis functions on X × Y, then approximates their component functions with neural networks.It bounds the approximation error for the basis components and combines these estimates to complete the proof.Source: A Proof of Theorem 1 · pp. 14, 16

B Training with Physics-Informed Loss

SPINN uses a physics-informed loss for PDE residuals and initial and boundary conditions, with forward-mode automatic differentiation for the PDE loss. The section also estimates derivative-operation costs and notes a gap between theoretical and implemented propagation complexity.

B Training with Physics-Informed Loss

  • The PDE loss uses Monte Carlo integration over factorized collocation coordinates and forward-mode automatic differentiation, while initial and boundary losses regress their given conditions.The total objective combines these terms with weighting factors.Source: B Training with Physics-Informed Loss · pp. 16–17

C FLOPs Estimation

  • Actual feature recomputation raises network-propagation complexity from theoretical O(Nd) to O(Nd²), though it remains below the conventional PINN complexity.The authors identify further software optimization as a way to narrow the gap between theoretical FLOPs and actual runtime.Source: C FLOPs Estimation · p. 17
  • Table 1 reports additions and multiplications for forward-pass and derivative evaluations using 64^3 collocation points in a 3D system.The estimate assumes derivatives are evaluated on every coordinate axis and uses vanilla MLP settings.Source: Table 1 · p. 17, C FLOPs Estimation · p. 17

D Experimental Details and Results

The experiments cover equations with separable and non-separable solution forms, including Navier–Stokes and higher-dimensional diffusion. The supplied results include accurate solutions for the (5+1)-d diffusion and flow-mixing problems.

D Experimental Details and Results

  • The appendix lists separable-form and non-separable-form solution examples across diffusion, Helmholtz, Klein-Gordon, Navier–Stokes, flow mixing, and L-shaped Poisson problems.The listed forms include both (3+1)-d Navier–Stokes and (2+1)-d Navier–Stokes.Source: D Experimental Details and Results · p. 17

D.1 Diffusion Equation

  • The diffusion experiment uses diffusivity α = 0.05, a two-dimensional spatial domain, and a FEniCS reference solution at 101 × 101 × 101 resolution.The initial condition is a superposition of three Gaussian functions.Source: §D.1 · p. 18

D.2 Helmholtz Equation

  • The Helmholtz experiment uses a manufactured three-dimensional sine-product solution with k = 1 and coefficients a1 = 4, a2 = 4, and a3 = 3.Source: §D.2 · p. 18

D.4 (2+1)-d Navier-Stokes Equation

  • The (2+1)-d Navier–Stokes test uses a JAX-CFD reference solution at 100 × 128 × 128 resolution for decaying turbulence with viscosity 0.01.The initial condition is a Gaussian random field with maximum velocity 5; training divides time into ten windows.Source: §D.4 · p. 19

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

  • For the (3+1)-d Navier–Stokes experiment, the authors constructed an analytic solution on [0, 2π]^3 over time [0, 5] and visualized the solution vector field in Figure 6.The experiment uses viscosity ν = 0.05 and a forcing term.Source: §D.5 · p. 20

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

  • With 86 collocation points, SPINN achieved a relative L2 error of 0.0074 on the (5+1)-d diffusion equation within 2 minutes.The test used the manufactured solution ∥x∥2 + 10t.Source: §E.1 · p. 20

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

  • With 2563 collocation points, SPINN achieved a relative L2 error of 0.0029 in 6 minutes on the flow-mixing problem.Figure 1 shows the visualized result.Source: §E.2 · p. 21, Figure 1 · p. 21

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

SPINN handles an L-shaped Poisson domain by combining separable interior training with conventional boundary-point training. The section also reports a slight accuracy increase from Adam followed by L-BFGS fine-tuning on three 3-d PDEs.

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

  • SPINN trains boundary points non-separably and interior collocation points separably, masking PDE loss outside the domain’s tight bounding box.The experiment split the L-shaped domain into two rectangular domains rather than masking its top-right region.Source: §E.3 · p. 21
  • The test follows DeepXDE settings for −∆u(x, y) = 1 inside Ω = [−1, 1]2 \ [0, 1]2, with zero boundary values.The reference solution was obtained by the spectral method.Source: §E.3 · p. 21
  • 0.0322 relative L2 error for SPINN versus 0.0392 for PINNs with modified MLP on the L-shaped Poisson problem.The reference solution used the spectral method; Figure 2 visualizes the solutions.Source: §E.3 · p. 21, Figure 2 · p. 22

E.4 Fine Tuning with L-BFGS

  • Adam training followed by L-BFGS fine-tuning showed a slight accuracy increase across three 3-d PDEs, while optimizer effects remain an open question in PINNs.The reported setup used 643 collocation points and a single outer-loop L-BFGS application.Source: §E.4 · p. 22

Related papers

Loading 2306.15969v4…