PIXEL: Physics-Informed Cell Representations for Fast and Accurate PDE Solvers
Abstract
Physics-informed neural networks (PINNs) have recently emerged and succeeded in various PDEs problems with their mesh-free properties, and unsupervised training.
However, their slower convergence speed and relatively inaccurate solutions often limit their broader applicability.
This paper proposes a new kind of data-driven
PDEs solver, physics-informed cell representations (PIXEL), elegantly combining classical numerical methods and learning-based approaches.
We adopt a grid structure from the numerical methods to improve accuracy and convergence speed and overcome the spectral bias presented in PINNs.
Moreover, the proposed method enjoys the same benefits in PINNs, e.g., using the same optimization frameworks to solve both forward and inverse PDE problems and readily enforcing
PDE constraints with modern automatic differentiation techniques.
Project URL: https://namgyukang.github.io/PIXEL/
Introduction
Physics-informed neural networks (PINNs) have recently received significant attention as new data-driven Partial Differential Equation (PDE) solvers [11].
However, PINNs suffer from slow convergence rates and they often fall short of the desired accuracy [3, 16, 17, 19].
In addition, Multi-layer perceptron (MLP) architecture is known to have spectral bias, which prioritizes learning low-frequency components.
Recent studies have shown that spectral bias [10] indeed exists in PINN models [8, 18] and this tendency towards smooth function approximation often leads to failure to accurately capture high-frequency components or singular behaviors in solution functions.
In this paper we propose physics-informed cell representations (coined as PIXEL ), which is jointly trained with shallow neural networks to improve convergence rates and accuracy.
Inspired by classical numerical solvers that use grid points, we divide solution space into many subspaces and allocate trainable parameters for each cell.
The key motivation of PIXEL is to disentangle the trainable parameters with respect to the input coordinates.
In neural network-only approaches, such as PINNs, all network parameters are affected by the entire input domain space.
Therefore, parameter updates for specific input coordinates influence the outputs of other input subspaces.
On the other hand, each input coordinate has dedicated trainable parameters updated only for certain input coordinates in PIXEL.
This parameter separation technique has been explored in visual computing domains [2, 5, 12, 13, 9, 1] and has shown remarkable success in terms of convergence speed of the training procedure.
Furthermore, PIXEL is immuned to spectral bias presented in PINNs.
A root cause of the bias is the shared parameters of neural networks for the entire input space.
However, PIXEL, each cell is only responsible for a small sub-region of the input domain.
Therefore, a large difference between neighboring cell values can better approximate high-frequency components or singular behaviors.
Our proposed model PIXEL uses a differentiable interpolation scheme to implement virtually infinite resolution grids, and the resulting representations are differentiable with respect to input coordinates.
PIXEL 2.1 Physics-informed neural networks
We briefly review physics-informed neural networks (PINNs).
Let us begin with the initial-boundary value problems with PDEs.
A general formulation can be written as follows:
[Equation]
where N x,t [ · ] is a linear or non-linear differential operator, B x,t [ · ] is also a differential operator for boundary conditions, and the initial conditions are denoted as u ( x, 0)
=
g ( x ) .
u ( x, t ) represents the unknown solution function and PINNs use neural networks, u θ ( x, t ) , parameterized by the trainable model parameters θ , to approximate the solution.
Then, neural networks are trained by minimizing the following loss function.
[Equation]
where, L res , L ic , L bc , L data are PDE residual, initial condition, boundary condition, and observational data loss functions, respectively.
λ are weighting factors for each loss term.
Each loss term is usually defined as mean square loss functions over sampled points.
For example, a PDE residual loss L res over N res collocation points can be written as,
[Equation]
2.2 Neural networks and grid representations
The proposed architecture consists of a small neural network and a feature extractor of input coordinates.
We approximate solution functions by a neural network f parameterized by θ ,
[Equation]
where C is a grid representation and ϕ is a feature extractor given input coordinates and the grid representation using an interpolation scheme.
Note that both C and θ are model parameters and updated during the training procedure.
The dimension of C is determined by the dimension of the spatial input domain.
For example, if x ∈ Ω ⊂ R then C ∈ R c × H × W is a three dimensional tensor, where the channel size c , and H and W are spatial and temporal grid sizes, respectively.
ˆ x ∈ [1 , H ] and ˆ t ∈ [1 , W ] are normalized input coordinates assuming input domain Ω ⊂ R and [0 , T ] are tightly bounded by a rectangular grid.
If x ∈ Ω ⊂ R 2 then C is a four dimensional tensor, and if x ∈ Ω ⊂ Rthen C is a five dimensional tensor 2 .
2.3 Mesh-agnostic representations through interpolation
In two dimensional grid cases, x ∈ Ω ⊂ R and C ∈ R c × H × W , the following is a feature extractor.
[Equation]
where C ij ∈ R c denotes cell representations at ( i, j ) , and k : [0 , 1] → [0 , 1] represents a monotonically increasing smooth function.
Given normalized coordinates (ˆ x, ˆ t ) , it looks up neighboring points ( 2 d +1 points) and computes the weighted sum of the representations according to a predefined kernel function.
It is differentiable w.r.t input coordinates so that we can easily compute partial derivatives for PDE differential operator N [ · ] by using automatic differentiation.
To support higher-order gradients, we use a cosine interpolation kernel which is multiple differentiable.
We use, k ( x ) :
= 12 (1 − cos( πx )) because it is everywhere continuous and infinitely differentiable.
2.4 Multigrid representations
The more fine-grained grids, the less chance a grid cell would see the points.
It would result in highly overfitted solutions.
Inspired by recent hierarchical grid representations [14, 9], we suggest to use multigrid representations.
We stack up multiple coarse-grained grid representations and the representations at each collocation point are computed by summing up the representations from all grids.
With a slight abuse of notation, a multigrid representation is defined as four dimensional tensors in two dimensional grids C ∈ R M × c
× H × W .
Then we can reformulate an interpolation function as,
[Equation]
where C i ∈
R
c
× H × W denotes a grid representation.
Multigrid approach was very critical to overall
PIXEL performance.
Without this, we observed that PIXEL suffers from serious overfitting issues.
Experiments
We experimented on 6 various 2D and 3D PDEs by using L-BFGS optimizer.
We implemented a customized CUDA kernel of the triple backward grid sampler that supports high-order gradients [15], the runtime and memory capacity were efficiently reduced.
Navier-Stokes Eq , [11] shows the result of the inverse problem, which is the multivariate coefficient simultaneous prediction.
Convection Eq , [3] shows that without sequential training, the original PINNs has failed to find accurate solutions.
Reaction-diffusion Eq , which is also a PDE that the original PINNs have worked poorly [3].
About Helmholtz Eq , [17] has reported that the original PINN has struggled to find an accurate solution.
We tested both low and high frequency parameters setting.
Allen-Cahn Eq is a non-linear second-order PDE that is known to be challenging to solve using conventional PINNs [17], and a few techniques, including adaptive re-sampling [20] and weighting algorithms [7, 4, 19], have been proposed.
Finally Burgers Eq , which is non-linear and known to have a singular behavior.
[Figure]
[Table]
3.1 Results
For the forward problem , In Convection equation , the resulting solution image of PINN was not correctly updated for the later time domain, t > 0 .
However, PIXEL converged to a high accurate solution.
For Reaction-diffusion equation , PINN showed a constant curve shape after 285 iterations in the averaged loss curve.
Whereas PIXEL showed an exponential decay shape until 10,000 iterations.
In Helmholtz which has the high-frequency components, as we expected, PINN has failed to converge to an accurate solution due to the spectral bias.
In contrast, PIXEL obtained high-accuracy solutions quickly in 3-dimension as well as 2-dimension.
For Allen-Cahn , which is known to be notoriously difficult, the previous studies have demonstrated that PINNs perform very poorly without additional training techniques, such as time marching techniques [6] or causal training [16].
However, our method can obtain accurate solutions without any additional methods.
For the inverse problem , PINN showed fluctuation in the prediction curve due to the random seed.
In contrast, PIXEL showed robustness in predicting regardless of random seed in 3-dimension
Navier-Stokes equation as well as 2-dimension equations.
Except for Helmholtz equation which
PINN failed to train, PIXEL showed convergence in a few iterations for PDEs.
It was found that
PIXEL achieves better performance than PINN at multivariate coefficient prediction of the inverse problem in 3-dimension Navier-Stokes equation . Acknowledgments
We are thankful to Junwoo Cho for helpful discussion and contributions.
This research was supported by the Ministry of Science and ICT (MSIT) of Korea, under the National Research Foundation (NRF) grant (2022R1F1A1064184, 2022R1A4A3033571), Institute of Information and Communication
Technology Planning Evaluation (IITP) grants for the AI Graduate School program (IITP-20190-00421).
The research of Seok-Bae Yun was supported by Samsung Science and Technology
Foundation under Project Number SSTF-BA1801-02. A Experimental setup and details [Table] u t + βu x =u ( x, 0) =
sin x u (0 ,t ) = u (2 π,t ) β x ∈ [0 , 2 π ] , t ∈ [0 ,T ]
[Table]
u t − νu xx − ρu (1 − u ) =
u ( x, 0)
=
h ( x ) u (0 ,t ) = u (2 π,t ) ν
Diffusion x ∈ [0 , 2 π ] ,t ∈ [0 ,T ] ∆ u ( x,y,z ) + k 2 u ( x,y )
=
q ( x,y,z ) u
( x,y,z )
= 0 , q ( x,y,z ) =
k [Table] in our experiments of the forward and the inverse problem. For all experiments, we used Limited-memory BFGS (L-BFGS) second-order optimization algorithms.
In many cases of training PINNs, it outperforms other first-order optimization algorithms, such as
ADAM or SGD.
We set the learning rate to 1.0 and used the strong-wolfe line search algorithm.
In every L-BFGS iteration, we randomly sample collocation points to make the model robust to the entire input and time domains for training PIXELs.
We found that PINNs often have struggled to converge in this setting, so we initially sampled the collocation points and fixed them, which has been a common practice in PINNs literature.
To compute the accuracy of the approximated solutions, we used the relative L 2 error, defined as || u − ˆ u
||
|| u
|| 2 , where ˆ u is a predicted solution and u is a reference solution.
We used NVIDIA RTX3090 GPUs and A100 GPUs with 40 GB of memory.
For all experiments, we used 2 hidden layers and 16 hidden dimensions for shallow MLP architecture, and a hyperbolic tangent activation function (tanh) was used.
For coefficients of the loss function, we used λ ic = λ bc = 1 , λ data = 0 .
Convection equation.
A shallow MLP of 2 layers with 16 hidden units was used.
The baseline PINN model was trained with the same number of data points from PIXEL.
For the PINN model, we usedhidden layers and 50 hidden dimensions following the architecture in [1].
Reaction-diffusion equation.
We used the same formulation in [1], and conducted experiments with the same PDE parameters ( ρ = 5 , ν = 3 ).
For training PINNs, we used 3 hidden layers andhidden dimensions following the architecture in [1].
2D Helmholtz equation.
We used the same formulation in [4].
The source term is given as q ( x, y )
= − ( a 1 π ) 2 u ( x, y ) − ( a 2 π ) 2 u ( x, y )
+ k 2 u ( x, y ) .
The analytic solution of this formulation is known as u ( x, y )
=
sin ( a 1 πx ) sin ( a 2 πy ) .
We tested the PDE parameters k = 1 , a 1 = 1 ,and a 2 = 4 .
For a more complex setting, we also tested k = 1 , a 1 = 10 and a 2 = 10 .
For the baseline
PINN model, we used 7 hidden layers and 100 hidden dimensions following the architecture in [4].
3D Helmholtz equation.
For 3D Helmholtz equation, the source term with q ( x, y, z ) = − ( a 1 π ) 2 u ( x, y, z ) − ( a 2 π ) 2 u ( x, y, z ) − ( a 3 π ) 2 u ( x, y, z ) + k 2 u ( x, y, z ) is used.
We tested the complex PDE parameters setting k = 1 , a 1 = 7 , a 2 = 7 , and a 3 = 7 .
MLP architecture is same as 2D problems.
Allen-Cahn equation.
For the baseline PINN model, we used 6 hidden layers and 128 hidden dimensions following the architecture in [2] In the case of Allen-Cahn, there was a problem that NaN occurs in PINN when the seed is 400.
Unlike PIXEL, PINN excludes seed 400 only in the case of
Allen-Cahn.
Burgers equation.
We used the same PDE parameter, ν = 0 .
01 /π in [3].
For the baseline PINN model, we adopted the same architecture in [3], using 8 hidden layers and 40 hidden dimensions.
3D Navier-Stokes equation In this paper, the same Navier-Stokes equation with the original PINN paper [11] is used.
We used 9 hidden layers and 20 hidden dimensions for the baseline PINN model to predict the multivariate coefficient in the inverse problem [Table]
A.1 Hyperparameter of experiments
[Table] [Table]
Table 3 and Table 4 shows the experimental details for the forward and inverse problems respectively.
Note that all other hyperparameters of the forward problems and the inverse problems are same explained in main text, including the architecture of PINN, confirming that the proposed method is not sensitive to hyperparameters.
For the inverse problems, The number of data points used for ground truth data points is 1,000,000 for Navier-Stokes equation, 25,600 for convection equation,
Burger equation, and Reaction-diffusion equation.
The Helmholtz equation uses 490,000 and the
Allen-Cahn equation uses 102,912.
B Grid size and the number of data points
This section studies the relationship between the amounts of training data points and the grid sizes.
We introduced multigrid representations that inject a smooth prior into the whole framework, which would reduce the required data points for each training iteration.
We demonstrate this with the convection equation example by varying the number of training data points (collocation and initial condition) and the number of multigrid representations.
We fixed the grid size 16 and the channel size [Table] and varied the number of multigrids from 4 to 64.
We reported the L 2 relative errors after
L-BFGS iterations.
As shown in Table 5, our method is robust to the amounts of training data points.
Although we can achieve higher accuracy with more training points, it performs comparably with a few data regimes
[Table]
C The visualization of multigrid representations
64) and (4,4,16,16) multigrid representations for Burgers and Helmholtz, respectively.
We demonstrate the intermediate results in Figure 5.
In Burgers example (the first and the second rows), we used (4,1,64,64) configuration and two layers of MLP with 16 hidden units.
As we expect, each grid show foggy images since the final solution will be the sum of all multigrid representations.
Also, we shifted each grid, which resulted in slightly different images from each other.
The final solution is completed in the last stage by filling up the remaining contents using an MLP.
Importantly, we note that the singular behavior (shock, a thin line in the middle of the solution image) is already well captured by the grid representations.
The role of MLP was to represent the smooth global component in solution function.
Therefore, the proposed grid representations and an MLP combine each other’s strengths to obtain better final solutions.
In Helmholtz example, we used small size grids (4,4,16,16).
Thus, we can observe notable differences after cosine interpolation (grid-like pattern before the interpolation).
We also note that the grid representations already represent complex patterns, and the last MLP stage also refined the solution by darkening the colors in the solution image.
D Visualizations