Separable PINN: Mitigating the Curse of Dimensionality in Physics-Informed Neural Networks
Abstract
Physics-informed neural networks (PINNs) have emerged as new data-driven PDE solvers for both forward and inverse problems.
While promising, the expensive computational costs to obtain solutions often restrict their broader applicability.
We demonstrate that the computations in automatic differentiation (AD) can be significantly reduced by leveraging forward-mode AD when training PINN.
However, a naive application of forward-mode AD to conventional PINNs results in higher computation, losing its practical benefit.
Therefore, we propose a network architecture, called separable PINN (SPINN), which can facilitate forward-mode
AD for more efficient computation.
SPINN operates on a per-axis basis instead of point-wise processing in conventional PINNs, decreasing the number of network forward passes.
Besides, while the computation and memory costs of standard
PINNs grow exponentially along with the grid resolution, that of our model is remarkably less susceptible, mitigating the curse of dimensionality.
We demonstrate the effectiveness of our model in various PDE systems by significantly reducing the training run-time while achieving comparable accuracy. Project page: https:/ jwcho5576.github.io/spinn/
Introduction
With the vast increases in computational power and breakthroughs in machine learning, researchers have explored data-driven and learning-based methods for solving partial differential equations [1, 3, 4, 5, 9, 10, 12, 13].
Among the promising methods, physics-informed neural networks (PINNs) have recently emerged as new data-driven PDE solvers for both forward and inverse problems [11]
PINNs employ neural networks and gradient-based optimization algorithms to represent and obtain the solutions, leveraging automatic differentiation to enforce physical constraints of underlying PDEs.
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.
PINNs, using neural networks and gradient-based iterative optimization algorithms, generally require a large number of iterations to converge.
Furthermore, the exponential growth of the number of forward and backward propagations due to the curse of dimensionality restricts their capability in solving high-dimensional PDEs.
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 required to obtain more accurate solutions grows substantially.
This paper aims to find a solution architecture to this fundamental limitation of training PINNs.
[Figure]
We propose a novel PINN architecture, separable PINN (SPINN) , which utilizes forwardmode automatic differentiation (AD) * to significantly reduce the computational cost for solving multi-dimensional PDEs.
Instead of feeding every multi-dimensional coordinate into a single MLP, e.g. f ( x 1 , x 2 , x 3 ) , we use factorized coordinates and separated subnetworks, where each sub-network takes independent one-dimensional coordinates as input, e.g. f 1 ( x 1 ) , f 2 ( x 2 ) , f 3 ( x 3 ) .
The final output is generated by a simple outer product and elementwise summation (Fig 3).
The suggested architecture obviates the need to query every multidimensional coordinate input pair, exponentially reducing the number of network propagations to generate a solution, O ( n d ) → O ( nd ) , where n is the resolution of the solution for each dimension, and d is the input dimension.
Experimental results demonstrate that given the same number of training points, SPINN significantly reduces the computational costs both in terms of FLOPs and wall-clock training time on commodity GPUs while achieving better accuracy.
Figure 1 shows the run-time of training conventional PINN and our proposed method trained on a (2+1)-dimensional Klein-Gordon equation.
Forward-Mode AD with Separated Functions
[Figure]
We demonstrate that leveraging forward mode
AD and separating the function into multiple functions can significantly reduce the computational cost of computing Jacobian matrix J .
In the separated approach (Fig 2 (b)), we first sample n one-dimensional coordinates on each of d axes, which makes a total of nd points.
Next, these coordinates are fed into d individual functions f 1 , . . .
, f d to generate the features.
We denote the number of operations of f as ops ( f ) .
Then, a feature merging function g is used to construct the solution of the entire n d discretized domain.
The amount of computations for AD is known to be 2 or 3 times more expensive than the forward propagation [2, 6].
With the scale constants c f , c g ∈ [2 , 3] we can approximate the total number of operations to compute the Jacobian matrix of a separated approach,
[Equation]
For a non-separated approach (Fig 2 (a)),
[Equation]
If we can make ops ( g ) sufficiently small, then the ratio C sep / C non-sep becomes ndn d ≪ 1 , quickly converging to 0 as d and n increases.
This mitigates the curse of dimensionality for solving highdimensional PDEs.
It also has linear complexity with respect to n , implying it can obtain more accurate solutions (high-resolution) efficiently.
(a) Architecture
Element-wise Product 𝑥𝑓 𝜃
𝑥𝑓
𝜃
𝑥𝑓 𝜃 ො𝑢(𝑥) = ෍ 𝑗= 𝑟 𝑓 𝜃 1 ,𝑗 𝑥 1 𝑓
𝜃 2 ,𝑗 𝑥 2 𝑓 𝜃 3 ,𝑗 𝑥𝑓 𝜃 1 ,1 𝑓 𝜃 1 ,2 𝑓 𝜃 1 ,𝑟 … ∘ … …
∑ 𝑓 𝜃 2 ,1 𝑓 𝜃 2 ,2 𝑓 𝜃 2 ,𝑟 𝑓 𝜃 3 ,1 𝑓 𝜃 3 ,2 𝑓 𝜃 3 ,𝑟 (b) Batchfied processing (Low-rank tensor decomposition)
+ + + … =
𝑓 𝜃 𝑓
𝜃 𝑓 𝜃 … …
𝑓 𝜃 1 , 𝑓 𝜃 2 , 𝑓 𝜃 3 , 𝑓 𝜃 1 , 𝑓 𝜃 2 , 𝑓 𝜃 3 , 𝑓 𝜃 1 ,𝑟 𝑓 𝜃 2 ,𝑟 𝑓 𝜃 3 ,𝑟 𝑥𝑥𝑥𝑓 𝜃 1 , 𝑓 𝜃 2 , 𝑓 𝜃 3 , 𝑓 𝜃 1 , 𝑓 𝜃 2 , 𝑓 𝜃 3 , 𝑓 𝜃 1 ,𝑟 𝑓 𝜃 2 ,𝑟 𝑓 𝜃 3 ,𝑟 …
Batchified Inputs
Separable PINN
Figure 3 (a) illustrates the overall SPINN architecture, parameterizing multiple separated functions with neural networks.
SPINN consists of d body-networks (MLPs), each of which takes an individual 1-dimensional coordinate component as an input.
Each body-network f θ i : R → R r (parameterized by θ i ) transforms an i -th input coordinate into a feature representation.
We define our predicted solution function as below:
[Equation]
where x ∈ R
d is an input coordinate, x i denotes i -th input coordinate.
f θ i is a vector-valued function and f θ i ,j denotes the j -th component of the output vector.
We used a tanh activation function throughout the paper.
In Eq.
3, the feature merging operation is a simple product and summation which corresponds to the merging function g described in Eq.
Due to its simplicity, g operations are cheaper than operations in MLP layers and thus, neglectable.
This makes the assumption C sep / C non-sep
≈ ndn d in Eq.
1 and 2 valid if we choose a such feature merging scheme.
The evaluated coordinate points of our model and conventional PINN have a distinct difference.
Both are uniformly evaluated on a d -dimensional hypercube, but points of SPINN form a lattice-like structure, which we call as factorizable coordinates .
This is because SPINN samples 1-dimensional coordinates from each axis and merges the features with outer product.
On the other hand, nonfactorizable coordinates are random points without any structure.
Factorizable coordinates with our separated MLP architecture enable us to evaluate functions on a much more dense ( n d ) grid, with a small number ( nd ) of inputs.
Due to its structural points and outer products between feature vectors,
SPINN’s solution approximation can also be viewed as a low-rank tensor decomposition (Figure(b)).
Among many decomposition methods, SPINN corresponds to CP-decomposition [7] 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) [8] to directly fit the decomposed vectors, we train neural networks which learn a continuous representation, enabling us to take derivative w.r.t.
input coordinates.
Experiments
We compared SPINN against vanilla PINN [11] on 3-d systems.
For every experiment, the PDE loss is calculated on fixed 90 3 collocation points, which means we sampled 90 · 3 factorized coordinates for SPINN, while 90 3 points are directly fed into PINN.
The number of collocation points of 90 3 was the upper limit for the baseline model when we trained with a single NVIDIA A100 GPU with 40GB of memory.
However, the memory usage of our model was significantly smaller, enabling SPINN to use more collocation points ( 180 3 and 360 3 ) to get more accurate solutions.
All experiments are trained with full batch for both models and reported error metrics are average relative L 2 errors computed by ∥
ˆ u − u
∥ 2 / ∥ u ∥ 2 , where ˆ u is the model prediction and u is the reference solution.
Table 1 shows our overall results.
SPINN is significantly more computationally efficient than the baseline PINN in terms of wall-clock run-time.
Klein-Gordon equation showed the largest gap, where our model achieved 57 × speed-up given the same input size.
With the help of the method outlined in [6], we also estimated the FLOPs for evaluating the derivatives.
In comparison to the baseline, SPINN requires 1 , 195 × less computations to evaluate the forward pass, first and second order derivatives.
See section D in the appendix for details.
In all cases, our model finds the most accurate solution.
Furthermore, while the number of collocation points grows exponentially, the memory usage and actual run-time of our model show logarithmic growth (Figure 4).
SPINN is also memory-efficient because it stores a much smaller batch of tensors, which are the intermediate values (primals) while building the computational graph.
We also identified from the results that we get more accurate solutions with more collocation points.
This training characteristic of PINN substantiates that our method is very effective for solving multi-dimensional
PDEs.
Figure 5 shows the error curves with run-time and visualized solutions of Helmholtz and
Klein-Gordon equation.
Additional experimental details and analyses are provided in the appendix.
[Table]
[Figure]
[Table]
[Figure]
I NN ( O u r s )
P I NN
[Figure]
(d) Klein-
[Figure]
A Related Works on Separated Architectures
CoordX [6] and NAM [1] have recently explored using separated MLP architectures.
Unlike ours,
CoordX had to use additional layers after the feature merging step, which makes their model more computationally expensive than ours.
In addition, their primary purpose is to reconstruct natural signals, such as images or 3D shapes.
We focus on solving PDEs and carefully devised the architecture to exploit forward-mode differentiation to efficiently compute PDE residual losses.
NAM [1] suggested separated network architectures and inputs, but only for achieving interpretability of the model’s prediction.
TensoRF [3] also shares similarities with ours.
It used CP-decomposition to factorize the output with low-rank tensors (Figure 3 (b)).
However, TensoRF directly optimizes low-rank tensors and uses tri-linear interpolation to implement continuous functions, or neural fields [11].
On the other hand, we use coordinated-based MLPs to predict low-rank tensors, which is a critical difference since solving PDEs mandates computing partial derivatives with respect to input coordinates.
FBPINNs [8] is another approach to use multiple MLPs to solve PDEs.
They decompose input domains , and each small MLP is used to cover a particular subdomain.
In contrast, we decompose input dimensions and solve PDEs over the entire domain cooperated by all separated
MLPs.
To the best of our knowledge, this is the first work to accelerate PINN training with separated
MLPs leveraging forward-mode AD.
B Forward/Reverse-Mode AD
[Table]
∂v 2 =
¯ v 3 · W 2 y =
v˙ y = ˙
v¯ v 3 = ¯ y
In this section, we provide some background on the two modes of AD and how Jacobian matrices are evaluated.
For clarity, we will follow the notations used in [2] and [4].
Suppose our function f : R n → R m is a two-layers MLP with tanh activation.
First column of Table 2 demonstrates a single forward trace of f .
To obtain a m × n Jacobian matrix J f in forward-mode, we compute the
Jacobian-vector product (JVP),
[Equation]
for i ∈ { 1 , . . . , n } .
The forward-mode AD is a one-phase process: while tracing primals (intermediate values), v k , it continues to evaluate and accumulate their tangents ˙ v
k = ∂v k
/∂x i .
This is equivalent to decomposing one large JVP into a series of JVPs by the chain rule and computing them from right to left.
A run of JVP with the initial tangents ˙ v 0 as the first column vector of an identity matrix I n gives the first column of J f .
Thus, the full Jacobian can be obtained in n evaluations.
On the other hand, the reverse-mode AD computes vector-Jacobian product (VJP): r ⊤ J f = [ ∂y j /∂y. . [Equation] ∂y
[Equation]
[Equation]
]  
[Equation]
[Equation] [Equation] [Equation] (5) for j ∈ { 1 , . . .
, m } , which is the reverse-order operation of JVP.
This is a two-phase process.
The first phase corresponds to forward propagation, storing all the primals, v k , and recording the elementary operations in the computational graph.
In the second phase, the derivatives are computed by accumulating the adjoints ¯ v k =
∂y j /∂v k (third of Table 2).
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 ).
Please refer to [2] for a more comprehensive review.
C Training with Physics-Informed Loss
After SPINN predicts an output function with the methods described in section 3, the rest of the training procedure follows the same process used in conventional PINN training [9], except we use forward mode AD (standard back-propagation, a.k.a.
reverse mode AD, used for parameter updates).
With the slight abuse of notation, our predicted solution function is denoted as ˆ u ( x, t ) , explicitly expressing time coordinates.
Given an underlying PDE (or ODE) and initial, boundary conditions,
SPINN is trained with a “physics-informed” loss function:
[Equation]
where Ω is an input domain, N , B are generic differential operators and u ic , u bc are initial, boundary conditions, respectively.
λ are weighting factors for each loss term.
When calculating the PDE loss ( L pde ) we sampled collocation points from factorized coordinates and used forward-mode AD.
The remaining L ic and L bc are then computed with initial and boundary coordinates to regress the given conditions.
By minimizing the objective loss in Eq. 6, the model output is enforced to satisfy the given equation and the initial, boundary conditions.
D FLOPs Estimation
The FLOPs for evaluating the derivatives can be systematically calculated by disassembling the computational graph into elementary operations such as additions and multiplications.
Given the 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 calculated precisely since it consists of a series of matrix multiplications and additions.
We used the method described in [4] to estimate FLOPs for evaluating the derivatives.
Table 3 shows the number of additions ( ADDS ), and multiplications ( MULTS ) in each evaluation process.
Note that FLOPs is a summation of ADDS and MULTS by definition.
[Table]
E Experimental Details and Analyses
We matched the total number of learnable parameters between each model for a fair comparison.
For our model, we used 3 body networks of 5 hidden layers with 50 hidden/output feature sizes each.
For the baseline model, we used a single MLP of 5 hidden layers with 100 hidden feature sizes.
We used Adam [5] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment.
All experiments were performed five times with different random seeds.
All weight factors λ in the loss function in Eq.
6 are set to 1.
In the following sections, we give more detailed descriptions and analyses of each experiment.
E.1 Linear/Non-linear Diffusion Equation
Diffusion equation is one of the most representative parabolic PDEs, often used for modeling the heat diffusion process.
It can be written as:
[Equation]
where Eq.
10 is a linear form and Eq.
11 is a non-linear form.
We used diffusivity α = 0 .
05 , spatial domain Ω ∈ [ − 1 , 1] 2 , temporal domain Γ ∈ [0 , 1] and used superposition of three Gaussian functions for initial condition u ic .
For error measurement, we compared against the reference solution obtained by a widely-used PDE solver platform FEniCS [7].
Note that FEniCS is a FEM-based solver.
The first and second row of Table 1 shows the numerical results of the linear and non-linear diffusion equation, respectively.
With 40 × less memory usage, our method finds more accurate solutions, 49 × ( 47 × for nonlinear case) faster than the baseline where 180 3 collocation points showed the lowest error.
E.2 Helmholtz Equation
The Helmholtz equation is a time-independent wave equation that takes the form:
[Equation]
[Equation]
(15) where the spatial domain is Ω ∈ [ − 1 , 1] 3 .
For a given source term q = − ( a 1 π ) 2 u − ( a 2 π ) 2 u − ( a 3 π ) 2 u + k 2 u , we devised a manufactured solution u = sin( a 1 πx 1 ) sin( a 2 πx 2 ) sin( a 3 πx 3 ) , where we take k = 1 , a 1 = 3 , a 2 = 3 , a 3 = 2 .
The results of our experiments are shown in the third row of
Table 1.
Due to stiffness in the gradient flow, it hinders conventional PINNs to find accurate solutions.
Therefore, a learning rate annealing algorithm is suggested to mitigate such phenomenon [10].
However, without bells and whistles, SPINN successfully finds accurate solutions.
On the same number of collocation points, the memory usage of our model is 41 × smaller, and the training speed is 56 × faster than the baseline.
Visualized solutions are illustrated in Figure 5 (e) left.
Helmholtz equation contains three second-order derivative terms, which makes the baseline model most time-consuming (more than two hours for full training).
Meanwhile, the linear diffusion equation contains one first-order derivative and two second-order derivative terms, which has fewer second-order derivatives than Helmholtz.
Comparing Helmholtz’s experiment against the linear diffusion experiment, the training time increased 28% ( 129 .
→ 166 . 06 ) for the baseline model, but SPINN has only increased 12% ( 2 . 64 → 2 . 95 ).
This is due to SPINN’s significant reduction of the forward-mode AD computations, which makes our model also very efficient at solving high-order
PDEs.
Figure 4 also shows the result of training Helmholtz equation to identify the error and run-time dynamics along the number of collocation points.
We increased the resolution of each axis fromto 90 (maximum size for baseline) and recorded the relative errors and training run-times for each resolution.
As expected, the relative errors of both models tend to decrease along the resolution.
However, the errors of the baseline model starts to converge at about 40 3 , while our model keeps finding a more accurate solution with larger points.
More importantly, the training run-times of our model barely increase ( 2 . 43 → 2 . 95 ) while that of the baseline model is very sensitive to the resolution ( 2 . 13 → 166 . 06 ).
E.3 Klein-Gordon Equation
The Klein-Gordon equation is a non-linear hyperbolic PDE, which arises in diverse applied physics for modeling relativistic wave propagation.
The inhomogeneous Klein-Gordon equation is given by
[Equation]
+ x 2 , x
[Equation]
Ω (17) u ( x, t ) =
u bc ( x ) , x
[Equation] t [Equation] where we chose the spatial/temporal domain to be Ω ∈ [ − 1 , 1] 2 and Γ ∈ [0 , 10] , respectively.
For error measurement, We used a manufactured solution u
= ( x 1 + x 2 )
cos( t ) + x 1 x 2 sin( t ) and f , u bc are extracted from this exact solution.
Fourth row of Table 1 shows the quantitative results.
Again, our method finds the best solution in every resolution with much efficient memory usage ( 40 × ) and training run-time ( 57 × ).
Note that this equation showed the largest difference in terms of run-time.
Analogous to Helmholtz equation, Klein-Gordon equation contains the same number of second-order derivatives and thus, shows a similar aspect on training run-time.
Visualized results are shown in
Figure 5 (e) right.