Compute-Optimal Solutions for Acoustic Wave Equation Using Hard-Constraint PINNs

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 19:45:17

Abstract: This paper explores the optimal imposition of hard constraints, strategic sampling of PDEs, and computational domain scaling for solving the acoustic wave equation within a specified computational budget. First, we derive a formula to systematically enforce hard boundary and initial conditions in Physics-Informed Neural Networks (PINNs), employing continuous functions within the PINN ansatz to ensure that these conditions are satisfied. We demonstrate that optimally selecting these functions significantly enhances the convergence of the solution. Secondly, we introduce an importance sampling strategy that optimizes the efficiency of hard-constraint PINNs under a fixed number of sampling points. Leveraging these strategies, we develop an algorithm to determine the optimal computational domain size, given a computational budget. Our approach offers a practical framework for domain decomposition in large-scale implementation of acoustic wave equation systems.

Main paper: 
 introduction: 
 The concept of using artificial neural networks to solve differential equations was first explored in the 1990s by Lagaris et al. [1998] . In the work of Lagaris et al. [1998] , they developed an ansatz solution that inherently satisfies the boundary conditions (BC) and the initial conditions (IC) of differential equations. More recently, the advent of physics-informed neural networks (PINNs) was marked by the influential study of Raissi et al. [2019] . This work leverages modern deep neural networks to solve forward and inverse problems involving nonlinear partial differential equations (PDEs), incorporating BCs and ICs through soft constraints in loss functions.Subsequent research has introduced various modifications to PINNs to enhance their accuracy, efficiency, and scalability [Lu et al., 2021a] . There are a couple of drawbacks for many PINNs with soft constraints for BCs and ICs. The selection of weights and samples for BCs and ICs cannot certainly be determined and requires many trial-and-error tests. Even when the loss function is minimized, the BCs and ICs are not strictly satisfied. To target the scaling problems of general PDEs and take advantage of parallel computing, XPINNs and FBPINNs have been developed based on domain decomposition methods [Jagtap and Karniadakis, 2020 , Shukla et al., 2021 , Moseley et al., 2023] .There are a few key points that these previous reseearches missed. First, how to formulate ansatz solutions satisfying BCs and ICs, specifically the function multiplier of NN. Second, if BC and IC are inheriently satisfied by constructing the ansatz solution, how to optimally sample the PDEs in the training process. Furthermore, for the existing PINNs handling scaling problems, how to decompose the domain to save the overall compute budget.Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.In this paper, we set up a 1D wave equation problem and investigate the optimal sampling and constraint imposing method given a compute budget.The contributions of this paper are as follows.• We systematically derived the implementation of hard BC and IC constraints in PINNs to solve acoustic wave equations. We give a strategy to select basic functions in the PINN ansatz solution that guarantee the satisfaction of BCs and ICs. We find that optimal selection of the basic function in the PINN ansatz can improve the convergence of PINNs.• We developed a Dynamic Amplitude-Focused Sampling (DAFS) algorithm to improve the convergence of hard-constraint PINNs for wave equations given a fixed number of sampling points.• With the hard constraint and importance sampling strategies, we propose an algorithm to find the optimal size of the computational given a compute budget. This domain size optimization algorithm can help the domain decomposition-based PINNs for large-scale problems save computational cost. 
related work: 
 Hard constraint Hard constraint PINNs can guarantee the satisfaction of BCs, ICs, symmetries, and/or conservation laws. There are comprehensive studies of embedding BCs in PINNs. Lu et al. [2021b] demontrated various ansatz equations to strictly meet Dirichlet and periodic BCs, and proposed the penalty method and the augamented Lagrangian method to impose inequality constraints as hard constraints. Liu et al. [2022] developed a unified ansatz formula to enforce the Dirichlet, Neumann, and Robin boundary conditions for high-dimensional and geometrically complex domains. Moseley et al. [2023] implemented the hard Dirichlet in the subdomain using a tanh 2 (ωx) function as the multiplier function of the neural networks in their FBPINN ansatz solution. However, studies on how to impose both hard BC and IC constraints in PINNs for acoustic wave equations that have a second-order time dirivative term are still limited. Alkhadhr and Almekkawy [2023] compared the accuracy and performance of PINNs with a combination of hard-BC/soft-BC and hard-IC/soft-IC for solving a 1D wave equation with a time-dependent point source function. This implementation of the hard-IC only considers the satisfaction of the wavefield values at the initial time u(x, t = 0), but neglects the hard constraint of the first-order time derivative of the wavefield u(x, t), i.e., ∂ t u(x, t = 0). Brecht et al. [2023] proposed improved physics-informed DeepONets with hard constraints, and presented a numerical example of a 1D standing wave equation with Dirichlet BCs. The DeepONet framework used in the paper has an inherent satisfaction of the initial wavefield, but ∂ t u(x, t = 0) is also neglected. This neglection does not affect the numerical results for the 1D standing wave equation in their paper, since they simply assume ∂ t u(x, t = 0) = 0.Strategic Sampling Many sampling algorithms have been developed to improve the training efficiency, mitigating failure modes of PINNs. [Wu et al., 2023] provided a comprehensive comparison of ten sampling methods, including non-adaptive and residual-based adaptive methods. Daw et al. [2023] proposed a Retain-Resample-Release (R3) Sampling algorithm to mitigate the failure propagation during the training processes of PINNs. [Gao et al., 2023a,b] developed failure informed adamptive sampling for PINNs, with the extentions of combining re-sampling and subset simulation. Yang et al.[2023] introduced a Dynamic Mesh-Based Importance Sampling (DMIS) method to enhance the training of PINNs. Additionally, [Zhang et al., 2024] proposed an annealed adaptive importance sampling method for solving high-dimensional partial differential equations using PINNs. Domain Scaling Computational domain scaling is a key issue to apply PINNs to real-world large spatial-temporal scale applications. [Jagtap and Karniadakis, 2020] proposed a generalized spacetime domain decomposition framework for PINNs, named extended PINNs (XPINNs), which can handle nonlinear PDEs on complex-geometry domains. XPINNs provide large representation and parallelization capacity by deploying multiple neural networks in smaller subdomains, offering both space and time parallelization to reduce training costs effectively. Shukla et al. [2021] developed a distributed framework for PINNs based on two extensions: conservative PINNs (cPINNs) and XPINNs. These methods employ domain decomposition in space and time-space, respectively, enhancing the parallelization capacity, representation capacity, and efficient hyperparameter tuning of PINNs. The framework allows for optimizing all hyperparameters of each neural network separately in each subdomain, providing significant advantages for multi-scale and multi-physics problems. They demonstrated the efficiency of cPINNs and XPINNs through various forward problems, highlighting that cPINNs are more communication-efficient while XPINNs offer greater flexibility for handling complex subdomains. Moseley et al. [2023] addressed the limitations of PINNs in solving large domains and multi-scale solutions by proposing Finite Basis PINNs (FBPINNs). FBPINNs use neural networks to learn basis functions defined over small, overlapping subdomains, inspired by classical finite element methods. This approach mitigates the spectral bias of neural networks and reduces the complexity of the optimization problem by using smaller neural networks in a parallel, divide-andconquer approach. Their experiments showed that FBPINNs outperform standard PINNs in accuracy and computational efficiency for both small and large, multi-scale problems. Chalapathi et al. [2024] introduced a scalable approach to enforce hard physical constraints using Mixture-of-Experts (MoE) in neural network architectures. This method imposes constraints over smaller decomposed domains, with each domain solved by an expert through differentiable optimization. The independence of each expert allows for parallelization across multiple GPUs, improving accuracy, training stability, and computational efficiency for predicting the dynamics of complex nonlinear systems. The optimal decomposition of subdomains is critical to the effectiveness of these scaling methods, given a fixed compute budget. Our work focuses on finding the maximum subdomain size that even a 64x2 small PINN can handle within a compute budget. 
methodology: 
 In this section, we outline our approach to effectively implement hard constraints, strategically sampling partial differential equations (PDEs), and optimizing the scaling of computational domains.These methods are utilized to solve the acoustic wave equation within a specified computational budget.We focus on an acoustic wave equation defined by:EQUATIONwhere:• D represents the differential operator. For a simplified one-dimensional acoustic wave equation, D = ∂ tt -c 2 (x)∇ 2 , indicating the second temporal derivative minus the spatial derivative scaled by the square of the local speed of sound, c(x).• B i denotes the boundary condition operator applied at x ∈ ∂Ω i .• I j signifies the initial condition operator, defining the state of the system at t = t 0 across the domain Ω.A prevalent ansatz employed in prior studies on hard-constraint PINNs for 1D wave equations is expressed as:EQUATION) where ũ(x, t) represents the neural network output with inputs x and t, and τ (t) is a function that satisfies τ (0) = 0. This design ensures that the initial condition u(x, 0) is met precisely when t = 0.To accommodate boundary conditions (BCs) at x = 0 and x = L, the ansatz is often modified to:EQUATIONcan ensure both Dirichlet BCs and the initial condition u(x, t)| t=0 = u(x, 0). However, this ansatz does not account for ∂ t u(x, t)| t=0 , unless it is assumed to be zero.We propose a more general hard constraint imposition formula:EQUATIONwhich guarantees satisfaction of the conditions:EQUATIONwhere U i (x, t), V j (x), W j (x) are the specified functions in BCs and ICs, and τ (t) is an arbitraryfunction satisfying τ (0) = d t τ (0) = 0.It is straightforward to demonstrate that the proposed ansatz correctly imposes all BCs and ICs as required:EQUATIONIn Section 4.2, we will explore numerical tests to optimize the selection of τ (t) by evaluating convergence rates and mean absolute errors (MAE).The primary advantage of employing hard constraints in our model is the elimination of the need to fine-tune the weights of PDE, BC, and IC loss terms typically required in soft-constraint PINNs.Sampling is crucial for efficient training of PINNs, ensuring rapid convergence and mitigating potential failure modes. To enhance the computational efficiency of our hard-constraint PINNs, we introduce the Dynamic Amplitude-Focused Sampling (DAFS) method. This strategy optimally selects the number of points, N pde , used in the training.Initially, we segmented the computational domain to identify regions with high-amplitude acoustic wave fields, based on low-resolution finite difference (FD) simulations. These high-amplitude regions are defined by a threshold δ, which determines the intensity level above which areas are considered to be of high amplitude. Within these identified regions, we uniformly sampled αN pde points. This was supplemented by uniformly sampling (1 -α)N pde points in the remaining areas of the domain.Both and α are parameters crucial to the sampling process and are optimally chosen to balance the computational budget and the accuracy of the simulations. By adjusting these parameters, we can tailor the distribution of sample points to areas that are most influential in the wave dynamics, thereby improving the efficiency of our PINN training.The pseudocode for the DAFS algorithm is provided in Algorithm 1. This sampling strategy, characterized by its focus on dynamically identified regions of interest based on wave amplitude, significantly optimizes the efficiency of the computation during the PINN training phase. The numerical tests for DAFS are in Section 4.3. 
problem setup: 
 We applied our method to three numerical examples for three different types of 1D acoustic wave equations -standing waves, string waves, and traveling waves. The ground truth wavefields are shown in Figure 1 . Standing waves for Dirichlet BCs Our first numerical example is a standing wave solution for the following 1D wave equation with Dirichlet BCs:EQUATIONThe analytical solution u(x, t) for Equation 8isEQUATIONA standing wave solutionu(x, t) = sin kπx L cos kπct L , k ∈ Z + (10)can be achieved if we assume U (x) = sin kπx L and V (x) = 0. We show the solutions for k = 1, 2, 3in Figure 1(a) .String waves for time-dependent BCs Our third example is a string wave solution for timedependent BCs shown in Equation 11. The ground truth solutions in Figuer 1(b) are achieved by finite different simulation.∂ 2 u(x, t) ∂t 2 -c 2 ∂ 2 u ∂x 2 = 0, x ∈ (0, L) B.C.: u(0, t) = u(L, t) = sin(2πt), I.C.: u(x, 0) = 0, ∂u ∂t (x, 0) = 2π cos 2kπx L (11)Traveling waves for Gaussian source time functions Our third example is a traveling wave solution for initial conditions of Gaussian source time functions shown in Equation 12. The ground truth solutions in Figuer 1(c) are computed by finite different simulation.∂ 2 u(x, t) ∂t 2 -c 2 ∂ 2 u ∂x 2 = 0, x ∈ (0, L) B.C.: u(0, t) = u(L, t) = 0, I.C.: u(x, 0) = 1 σ √ 2π exp - (x -µ) 2 2σ 2 , ∂u ∂t (x, 0) = 0 (12)We selected six candidate functions for τ (t) to construct PINNs with a network configuration of only 64x2 neurons. Figures 2 through 4 illustrate the L 2 loss and L 1 error as functions of training epochs.Our findings suggest that τ (t) significantly influences both the convergence rate and the emergence of failure modes. In general, t 2 , 2t 2 1+t 2 performs better in general, especially for higher modes k = 2, 3.We show a few training dynmaics in Appendix C.Our analysis indicates that the frequency characteristics of τ (t) and the corresponding wavefields may be critical for selecting an appropriate τ (t). Matching these characteristics can potentially enhance the model's efficiency by aligning τ (t)'s influence on the neural network's learning dynamics with the physical properties of the wave phenomena being modeled. We demonstrate the efficacy of our proposed Dynamic Amplitude-Focused Sampling (DAFS) in enhancing both the convergence and accuracy of Physics-Informed Neural Networks (PINNs).Experiments varying α from 0 to 0.5 to 1 indicate that optimal results are typically achieved when α is around 0.5.This suggests a balanced sampling strategy, where a significant portion of the samples is concentrated in regions of higher amplitude. However, exclusively focusing on these high-amplitude areas can hinder information transfer from boundary conditions to the interior of the domain, potentially leading to failure modes. Figures 5 and 6 illustrate these dynamics, showing the L 2 loss and L 1 error across different values of α, and the impact on the predicted wavefield and its accuracy. We then propose an optimal subdomain selection method shown in a flow chart in Figure 7 . This method will automatically determine the optimal k our 64x2 small PINNs can handle, given a compute budget. 
limitations and training dynamics: 
 While our proposed methods significantly enhance the functionality and efficiency of PINNs, the determination of the optimal function τ (t) presents certain limitations. The choice of τ (t) is crucial as it directly affects the model's ability to satisfy boundary and initial conditions rigidly. However, finding an ideal τ (t) that adapts across different problems and boundary conditions without extensive trial and error remains challenging. The training dynamics are also sensitive to the form of τ (t), where inappropriate selections can lead to slower convergence or even divergence in some cases. These issues underscore the need for a more automated, perhaps adaptive, approach to selecting τ (t) that can dynamically adjust based on the evolving training characteristics and the specific requirements of the PDE being solved.  
conclusion: 
 This work presented a comprehensive approach to improving the effectiveness and efficiency of Physics-Informed Neural Networks (PINNs) for solving acoustic wave equations. By integrating a well-formulated hard constraint imposition strategy and the novel Dynamic Amplitude-Focused Sampling (DAFS) method, we have significantly enhanced both the accuracy and convergence of PINNs.Our methodological innovations include:• A systematic derivation of hard boundary and initial conditions in PINNs that ensures these constraints are inherently satisfied, leading to better convergence and stability of the solution.• The introduction of DAFS, which optimally allocates computational resources by focusing sampling in regions of high amplitude while ensuring adequate coverage across the computational domain to prevent information isolation.• Development of a domain size optimization algorithm that assists in domain decomposition, enabling efficient scaling of PINNs for large-scale applications while managing computational costs.These contributions mark a significant step forward in the practical deployment of PINNs, especially in fields requiring the simulation of complex physical phenomena over large scales. Future work will focus on extending these strategies to other types of partial differential equations and exploring the integration of our methods with other deep learning frameworks to further enhance the adaptability and efficiency of PINNs in diverse applications, for example, we will explore the integration of our methods with existing PINNs frameworks that employ domain decomposition techniques, such as XPINNs and FBPINNs, to further enhance their scalability and adaptability. We aim to make PINNs more adaptable and efficient for a broader range of applications, particularly in complex systems where traditional numerical methods struggle. By advancing these strategies, we can significantly contribute to the deployment of PINNs in real-world scenarios, tackling large-scale and multi-scale challenges effectively.A Phase diagrams of loss weights  
 
 Appendix: The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:• You should answer [Yes] , [No] , or [NA] .• [NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.• Please provide a short (1-2 sentence) justification right after your answer (even for NA).The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions)with the final version of your paper, and its final version will be published with the paper.The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer[Yes] to a question, in the justification please point to the section(s) where related material for the question can be found. IMPORTANT, please:• Delete this instruction block, but keep the section heading "NeurIPS paper checklist",• Keep the checklist subsection headings, questions/answers and guidelines below.• Do not modify the questions and only use the provided macros for your answers.  


Decision: This paper proposes to improve the training efficiency of physics-informed neural networks, first by extending ansatz to also take the first derivative into account, second by a sampling method that focuses on high-amplitude regions, and third by a framework for domain decomposition. The reviewers generally think that, in the current state, the scope of the paper is quite narrow as it only tackles the 1D wave equation, the paper lacks comparisons with existing approaches, and the presentation needs to be significantly improved to avoid confusions.

Reviews: review: paper_summary: paper_summary: This paper improves the training efficiency of original physics-informed neural networks to solve the 1D wave equation threefold: first by extending ansatz to also take the first derivative into account, second by a sampling method that focuses on high-amplitude regions, and third by a framework for domain decomposition.
strength_weakness: strength_weakness: [Originality] While considering the first derivative for the ansatz is a good addition, the contribution is only minor. 
Sampling more collocation points in the regions that might be more difficult to solve is a practical approach however the comparison and distinction to other sampling methods is missing.
Lastly if I understand the domain decomposition framework correctly, the contribution is to wrap the entire training into a loop and, based on the training process's results, increase or decrease the subdomain size. 

Evaluation results are only provided for the 1D wave equation. Further results for other differential equations are necessary to demonstrate the benefits of the proposed method.

[Clarity] 
The framework for domain decomposition is not presented clearly. While the flow chart in Figure 7 provides an overview of the method additional textual explanations in Section 4.4 are needed.
There were few to no remarks about the training regime (#training points, optimizer, learning rate…, etc.), making it more difficult to reproduce results.
Minor remarks:
-            N_pde is not introduced. It is probably the number of collocation points?
-            Most of the Figures (e.g. Fig. 1, Fig 6.) are hard to read.
-            Line 46: (…) optimal size of the computational [domain?] given (…)
-            Line 149: Both [N_pde?] and alpha (…) 
-            Line 178: (…) In general, (...) performs better in general
questions: questions: Q1: In the abstract the DAFS method is described as a method "that optimizes the efficiency (…) under a FIXED number of sampling points. However in Section 3.2: "This strategy optimally selects the number of points, N_pde, used in training." I assumed that DAFS only distributes the collection points into high and low-amplitude regions. Is that correct?
Q2: How does DAFS compare to other state-of-the-art collocation point sampling methods? When should one use DAFS, and in which cases a different method?
limitations: limitations: While the authors clearly state that they are interested in the 1D wave equation it would have been interesting to see their proposed methods applied to the 2D wave equation of any other differential equations what are typically used in PINN benchmarks.
paper_summary: paper_summary: This paper explores to solve the acoustic wave equation in the context of PINNs. Hard boundary and initial conditions are enforced by employing continuous functions within the PINN ansatz to ensure that these conditions are satisfied. A Dynamic Amplitude-Focused Sampling (DAFS) method is introduced to improve the efficiency of hard-constraint PINNs under a fixed number of sampling points.
strength_weakness: strength_weakness: 1. Only the wave equation is discussed.
2. The proposed Dynamic Amplitude-Focused Sampling (DAFS) method is trivial.
3. There are no comparisons with other methods in the experiments.
4. In the experiments, the relative errors between exact solutions and predictions are not given.
5. In the context of PINNs, it is better to give explicitly the formulation of training loss. Training details are also lacking. 
6. Instead of tuning \tau (t) manually, it is better to train \tilde{u}(x,t) and \tau (t) simutanuously.
7. Many typos and grammar errors, such as "both and \alpha" in line 149, "x \in {\partial \Omega}_i" in line 125, "computational" in line 46.
8. The quality of Fig.7 should be improved.
questions: questions: In line 41, what does the "basic function" mean? Is it the function \tau (t)?
limitations: limitations: Only the wave equation is discussed. There are no comparisons with other methods in the experiments.
paper_summary: paper_summary: The manuscript treats the one dimensional wave equation with a PINN approach and discusses the imposition of boundary and initial conditions directly into the network, as common practice in PINNs. The authors then propose a quadrature scheme based on a coarse finite difference discretization of the wave equation.
strength_weakness: strength_weakness: The main weakness of the manuscript is the focus on the very special and simple toy problem of the one dimensional wave equation. Solving the one-dimensional wave equation with PINNs is only of academic interest and insights obtained from it for the training of PINNs might not generalize. More specifically:
- The exact imposition of the time derivative should also work for general time dependent equations. The authors should comment on this.
- The sampling strategy employing a finite difference simulation to determine regions of high sampling density is not a generalizable approach. If a finite difference solver for the equation at hand is available, a PINN solver is typically not required.
- The authors determine an optimal function $\tau$ via considering six concrete examples. There is no guarantee that this approach will generalize to different equation types and is therefore of limited practical use.
- The authors might want to discuss the theoretical literature that proves the theoretical advantage of exactly imposed boundary conditions [1, 2, 3] and more elaborate constructions of distance functions.

[1] https://proceedings.mlr.press/v190/muller22b/muller22b.pdf

[2] https://arxiv.org/abs/2311.00529

[3] https://www.sciencedirect.com/science/article/abs/pii/S0045782521006186
questions: questions: See weaknesses.
limitations: limitations: The scope of the paper is too narrow.
paper_summary: paper_summary: The paper attempts to train PINNs which solves acoustic wave equations. They do so by using hard-constrained PINNs which can enforce IC and BCs, and propose a collocation point sampling method (DAFS) based on the amplitude of the solution at different regions.
strength_weakness: strength_weakness: The paper itself feels less coherent, and seems like just an application of many existing PINN training techniques (e.g., hard constraint PINNs, collocation point sampling) into solving a certain problem, rather than providing a novel method or a coherent framework into solving a domain-specific problem.

The experimental section feels incomplete. Different point selection algorithms have not been extensively compared with, e.g., from that in Wu et. al. (2023). Furthermore, it would be interesting to see how the method can scale to more realistic acoustic problems (i.e., outside of 1D settings).

The paper itself also seems incomplete. The Appendix and the NeurIPS checklist are partially filled and have half-finished sentences.

The labels within the graphs can also be enlarged slightly to make them more readable.
questions: questions: 1. Can the optimal \alpha be selected without having to test out different values of \alpha? Is there some recommended value that can be used for different acoustics problems?

2. Are there any relation between selecting points in high-amplitude regions with selecting points in high-residual regions, such as those methods benchmarked in Wu et. al. (2023)? In the sense that these two are indirectly doing the same thing, or they end up selecting very similar points anyway.

3. In Figure 5, are there any intuition to why L1 loss peaks at larger \alpha but L2 peaks at smaller \alpha?

4. How is the computation time of the methods proposed?
limitations: limitations: The authors have provided limitations with selection of \tau.

Comments: comment: I thank the author for their response. I still believe that the paper requires some amount of revision and therefore will keep the current score.
Thank you for your answer.  Since some concerns have not been solved, I would like to keep my score.
Thank you for your answer. I think the pre-print still needs a decent amount of work,  I am still critical about relying on a finite difference solver within a PINN workflow. I would like to keep my score.