Separable Physics-Informed Neural Networks

Summary

This paper introduces a novel approach to using separable neural networks to represent solutions of partial differential equations (PDEs). In contrast to the conventional form used in physics-informed neural network (PINN) methods, the specialized form of neural networks proposed in this work takes advantage of forward-mode auto-differentiation, leading to faster training times in PINN frameworks. The authors provide a mathematical demonstration of the proposed form's approximation ability and conduct experiments on several PDEs. The results reveal that the proposed method achieves comparable accuracy to previous work, while significantly reducing computational and memory costs.

Strengths

1. The speed-up achieved by the proposed method is highly promising. The time required for training a single network to convergence has been a limitation in the application of PINN methods. This work addresses this issue by introducing a specialized neural network form, and in combination with the forward-mode auto-differentiation method, the training speed is significantly improved. 2. In addition to the improved auto-differentiation computation, this work employs fixed collocation points, enabling further acceleration of the total training time by reusing calculations on each point. By leveraging this approach, the computational efficiency of the method is significantly enhanced.

Weaknesses

Although the proposed method in this paper is novel and interesting, there are several concerns that the reviewer would like to address: 1. The choice of using a manufactured solution that aligns with the variable separation form for evaluating the results on the Helmholtz Equation, Klein-Gordon Equation, and (3+1)-d Navier-Stokes Equation raises some questions. While this choice is suitable for evaluating original PINN methods that do not utilize the variable separation property of the ground truth, it may be unfair to compare SPINN with the original method on these tasks. 2. One of the key advantages of PINN-based methods is their flexibility to be applied to any geometric surface without modifying the grids. However, the variable separation form of the neural network in SPINN may limit this flexibility. It would be challenging for SPINN to fit the boundary conditions of complex geometries due to the constraints of the variable separation form. 3. Throughout the paper, the evaluation metric predominantly used is the $L_2$ evaluation metric. However, the reliability of the $L_2$ metric, especially for non-linear PDEs, can sometimes be questionable. It would be valuable to also report metrics such as the Sobolev norm or the PINN loss to provide a more comprehensive evaluation.

Questions

1. To assess the performance of SPINN on more complex solutions that do not directly align with its form, the authors should conduct more experiments. The experiments could involve using a Multilayer Perceptron (MLP) in the classical PINN as the manufactured solution and training SPINN under the associated PDE. 2. The reviewer is curious about the performance of SPINN on PDEs with sophisticated boundary. For instance, it would be interesting to explore how SPINN performs on a sphere boundary condition, where the variables are still separated into x, y, and z components. 3. Including the Sobolev norm or PINN loss as additional evaluation metrics would provide the reviewer with more information to assess the performance of SPINN comprehensively. 4. The reviewer wonders if it is possible to reduce the dimension of the problem through variable separation. Most PINN methods work on high-dimension manifolds, incurring significant computational costs. While SPINN employs low-dimensional functions to construct the high-dimensional function, its domain is still a high-dimensional manifold. The reviewer believes that there might be a way to directly optimize these low-dimensional functions within the low-dimensional manifold, potentially offering benefits in solving high-dimensional PDEs.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The authors have talked their limitation in the paper

Summary

This paper presents a novel variant of PINN for learning multi-dimensional PDEs. Unlike traditional approaches considering point-wise inputs, this method incorporates each axis information separately during the learning process and leverages tensor decomposition to interpret the predicted solution variables. Additionally, the authors consider a forward-mode auto-differentiation technique that facilitates training with a more significant number of collocation points. The extensive experiments proved the efficiency, scalability, and effectiveness of the proposed SPINN.

Strengths

- This paper contributes to alleviating the scalability and efficiency issues of PINNs. Also, the proposed forward-mode auto diff method overcomes collocation point constraints, leading to more accurate solutions. I think it is of interest to the PINN community and the experimental results are convincing. - The authors provide the theoretical foundation of SPINN. Also, the authors have done extensive experiments on high-dimensional PDEs to prove the scalability of SPINNs. Although the tested PDEs are not that complicated, I don't think that is a big issue in this paper since it is more related to the optimization issue of PINNs. - This paper is well-written and well-organized.

Weaknesses

- This paper only alleviates the scalability issue, but the primary concern of PINNs lies in the challenges of optimization. Therefore, its contribution to the scientific machine-learning community is moderate. - It would be good to have a paragraph discussing the related work on scalability for PINNs. There are many existing techniques to handle the scalability issue, such as domain decomposition [1], seq2seq [2], and adaptive sampling method [3]. Using adaptive sampling can also help reduce the collocation points. Moreover, in the baseline comparison, it may strengthen the paper to compare SPINN with one of these techniques though my impression is that SPINN will be better based on the results of this paper. [1] Jagtap, A. D., & Karniadakis, G. E. (2021, March). Extended Physics-informed Neural Networks (XPINNs): A Generalized Space-Time Domain Decomposition based Deep Learning Framework for Nonlinear Partial Differential Equations. In AAAI spring symposium: MLPS (Vol. 10). [2] Krishnapriyan, A., Gholami, A., Zhe, S., Kirby, R., & Mahoney, M. W. (2021). Characterizing possible failure modes in physics-informed neural networks. Advances in Neural Information Processing Systems, 34, 26548-26560. [3] Subramanian, S., Kirby, R. M., Mahoney, M. W., & Gholami, A. (2022). Adaptive self-supervision algorithms for physics-informed neural networks. arXiv preprint arXiv:2207.04084.

Questions

- What about the performance of forward-mode auto diff for higher-order derivatives (e.g., fourth-order in the Kuramoto-Sivashinsky equation)? - How do the authors consider the rank *r* (i.e., *r*-dimensional feature representation in Eq. (5))? In Theorem 1, the authors claim *r* should be sufficiently large. What about in practice? - In lines 193-195, for each axis, this paper considers random sampling for 1D input points. Random sampling may not be an optimal choice. Is there any way to improve it in the context of SPINN? Can the authors add some discussions on this part? - How do the authors consider the optimization in PINNs? In the section on Limitations (lines 335-342), the authors claim that the learning-rate annealing will be considered in the future. I assume the authors were just doing the hyper-parameter tuning by using a grid search. Is that correct?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The discussion of the limitations in this paper is comprehensive.

Summary

The paper proposes a methodology for significant savings in computation and memory in learning physics informed networks for approximately solving PDEs. This enables learning with a much larger number of points, resulting in better accuracy. The idea is to use a specific function class, where features are constructed from each input dimension, and the output is simply the product of these features. The function class is shown to be universal. Since the same feature encoding can be used for all points that share the same value in each dimension, feature embeddings are calculated only once in a linear time for an exponential number of points, resulting in substantial savings. Using forward mode automatic differentiation, derivatives also benefit from the same savings. Experimental results demonstrate the effectiveness of this approach in approximating the solution of diffusion, Helmholtz, Klein-Gordon, and Navier-Stokes PDEs.

Strengths

The paper is nicely written, the background material is quite useful and many figures help demonstrate the main idea of the paper. In particular, the paper briefly discusses connections to similar ideas in implicit neural representation, and tensor factorization methods, which I found quite relevant. The main idea is simple, and a priori it is not obvious that the chosen function class can perform so well in practice. Therefore the primary strengths of the paper is its rather surprising experimental results, where for some important PDEs orders of magnitude in speed-up results in accurate estimates compared to vanilla PINN and some other variants. For the same reason, the main theorem of the paper that states the universality of the proposed architecture is also quite useful. Overall, I enjoyed reading the paper, and I believe the proposed methodology can have a high practical impact.

Weaknesses

While I did not identify a major weakness, I believe the paper can benefit from the following changes: - While experimental results suggest that the proposed architecture can deliver more improvements in some settings than others (e.g., results for diffusion), there is no discussion of this. In general, while the universality result is reassuring, it is unclear in what kind of PDEs it is a useful inductive bias. - I found Section 4.3 on gradient computation for SPINN is useful; the paper can benefit from doing the same exercise for higher-order derivatives since the complexity of calculating these derivatives in terms of the dimensionality of the embedding space will dominate the complexity of SPINN. - While I’m aware that reporting relative error is a common exercise, it is quite sensitive and unstable (due to division by the true value of the PDE), and, therefore, could become misleading. I suggest also reporting the RMSE for all experiments. - For clarity, the paper should state that the overall complexity still remains exponential in resolution (N), despite reducing it to linear in “network propagations”.

Questions

See above.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Yes.

Summary

The paper proposes a new class of architecture called Separable PINN (SPINN) that operates on a per-dimension basis instead of the standard coordinate based MLP architecture. This architecture significantly reduces the computational complexity and allows the use of a large number of collocation points to enforce the PDE residuals. The paper also provides an universal approximation theorem for the proposed architecture. The paper empirically demonstrates that the proposed SPINN is significantly faster than the original PINN architecture during training while maintaining accuracy.

Strengths

1. The idea of using a different sub-network for every dimension combined with an outer-product operation to evaluate on a grid is both technically sound and interesting. 2. The paper provides theoretical foundations of the SPINN architecture in the form of the universal approximation theorem. 3. The gradient computation using forward mode AD and the network architecture together leads to significant speed-ups in the training time of PINNs (which is quite impressive).

Weaknesses

1. Although the SPINN architecture is quite efficient, I am concerned about the scalability of the architecture for high dimensional PDEs (>4). (Please see question 1) 2. The size of the plots in Figure 6 is excessively small, making it difficult to compare the performance between the baselines. 3. The PDEs chosen for the baselines are relatively simple (other than the Navier Stokes). It would be interesting to compare the architectures for equations where standard PINNs fail such as Convection Equation (with larger values of $\beta$), Kuramoto Sivashinsky Equation (in chaotic regimes).

Questions

1. One of the main claims of the paper is that the computational complexity of the SPINN architecture linearly scales with the number of collocation points $N$ and the dimensions $d$, in contrast to the standard PINN which scales as $N^d$. However, for the SPINN architecture if we have a different branch for each dimension $d$, the outer product operation of multiplying $d$ different quantities can be numerically unstable. For example, can it lead to gradient vanishing or gradient exploding problems? Is some form of normalization required to ensure that the network is stable? 2. The universal approximation theorem guarantees that the approximation error can be as low as possible. However, does the SPINN architecture have any optimization problems? 3. In Figure 6, for a small number of collocation points(16, 32, 64) for Diffusion Equation, the PINNs perform much better than SPINN. Is there a specific reason why SPINNs did not perform well?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The limitations have been adequately addressed in the paper.

Summary

In this paper, the authors propose two main contributions: (i) forward mode AD for PDEs (ii) separating the contributions from each dimension (just like in separable conv) and performing computation using tensor multiplication. The results naturally show significant improvement over PINNs.

Strengths

+ Significant amount of speedup in performance. + To be honest, this is probably the only PINN variant I have seen showing results for 3D with as many collocation points as 64^3 etc. + They can even solve (3+1) d problems.

Weaknesses

- the language needs to be worked upon. It's not that forward AD is a very new area. In the broad community of Neurips, there are a lot of works that use similar ideas. Statements like "To our knowledge, it is the first attempt to exploit the power of forward-mode AD" needs to be made with caution. Perhaps make it clear that in the purview of PINNs or something. In fact, even in PINN-type papers, there are works that use Finite difference and Finite element based approaches (which don't need a backward pass) e.g. (i) https://arxiv.org/pdf/2005.08357.pdf (ii) https://arxiv.org/abs/2112.04960 (iii) https://arxiv.org/abs/2211.03241 (iv) https://arxiv.org/abs/1901.06314 (v) https://arxiv.org/pdf/2109.07143.pdf - lack of rigorous comparisons - lack of ablation studies

Questions

- What is the rank in Figure 7? - It seems that your accuracy with even higher number of collocation points is higher than PINNs. Yes, you are solving faster. But is it better? I don't seem to see it clearly... - Where are the comparisons with other methods like deepOnets, FNOs, etc. - is 1.9e-3 sufficient loss for Navier-Stokes? What is the typical acceptable error in real applications? and can SPINN achieve it?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Handling arbitrary domains

We thank the reviewers for taking their time to suggest all the invaluable comments and constructive feedback. To address the training of SPINN on arbitrary domains, we provide a description and experimental result obtained during the discussion phase. Besides the coordinate or mesh transformations [1, 2] mentioned in the rebuttal, we would like to present yet another simpler way to handle arbitrary domains with SPINNs. The following is our suggestion. For boundary points, we train in a non-separable way, just as normal pinn training (non-factorizable coordinates). For collocation points inside the domain, we train in a separable way (factorizable coordinates) and ignore the points outside the domain. We can easily achieve this by finding a tight bounding box of the input domain and masking out the PDE loss when the coordinates are outside. To show the effectiveness, we trained SPINN on an L-shaped domain and tested on a 2D Poisson equation following the settings used in DeepXDE [3]. The reference solution was obtained by a spectral method (not a manufactured solution), and SPINN achieved a relative L2 error of 0.0322 (PINN got 0.0392). You can see the visualization via an anonymous link we gave to the ACs ('poisson_result.png'). For practical reasons, we split the L-shaped domain into two rectangular domains instead of masking out the top-right part of the domain. Note that the L-shaped domain we showed is just an example case, and we believe that this method can be applied to any arbitrary input domains. We hope this will resolve the concerns raised by the reviewer 'dosJ', and we are eager to explore further to enable SPINN to handle more complicated domains in future works. **Reference** [1] Gao et al., "PhyGeoNet: Physics-Informed Geometry-Adaptive Convolutional Neural Networks for Solving Parameterized Steady-State PDEs on Irregular Domain", Journal of Computational Physics (2021) [2] Li et al., "Fourier Neural Operator with Learned Deformations for PDEs on General Geometries", arXiv (2022) [3] Lu et al., "DeepXDE: A deep learning library for solving differential equations", SIAM Review (2021)