Least-Squares-Embedded Optimization for Accelerated Convergence of PINNs in Acoustic Wavefield Simulations

TL;DR

The paper tackles slow convergence in Physics-Informed Neural Networks (PINNs) for high-frequency acoustic wavefields governed by the Helmholtz equation. It introduces a hybrid least-squares–gradient descent (LS-GD) optimization by embedding a least-squares solver for the linear output layer directly into the training loss, enabling optimal output-layer updates with minimal overhead. The approach accommodates both with and without Perfectly Matched Layers (PML) and leverages forward-mode differentiation and a Cholesky LS solve on a small $P\times P$ system, yielding faster convergence, higher accuracy, and improved stability compared to standard GD PINNs, including challenging Marmousi-like scenarios. The method is scalable and practical for large-scale wavefield simulations, offering a robust path to accelerate PINN-based seismic modeling. A TensorFlow implementation is provided to facilitate adoption and experimentation.

Abstract

Physics-Informed Neural Networks (PINNs) have shown promise in solving partial differential equations (PDEs), including the frequency-domain Helmholtz equation. However, standard training of PINNs using gradient descent (GD) suffers from slow convergence and instability, particularly for high-frequency wavefields. For scattered acoustic wavefield simulation based on Helmholtz equation, we derive a hybrid optimization framework that accelerates training convergence by embedding a least-squares (LS) solver directly into the GD loss function. This formulation enables optimal updates for the linear output layer. Our method is applicable with or without perfectly matched layers (PML), and we provide practical tensor-based implementations for both scenarios. Numerical experiments on benchmark velocity models demonstrate that our approach achieves faster convergence, higher accuracy, and improved stability compared to conventional PINN training. In particular, our results show that the LS-enhanced method converges rapidly even in cases where standard GD-based training fails. The LS solver operates on a small normal matrix, ensuring minimal computational overhead and making the method scalable for large-scale wavefield simulations.

Figures (5)

• Figure 1: Schematic representations of the neural network architectures. $\mathit{w}_{i,j}$ show the elements of the output layer's weights matrix $\mathbf{W}$, calculated by the LS solver. Dashed connections have no associated weight.
• Figure 2: Comparison of a 10Hz scattered wavefield predictions. (a) The velocity model used for simulation. (b) The real part of the finite-difference computed scattered wavefield. (c) Best prediction of a simple PINN within 50,000 epochs. (d) Prediction using the proposed LS-GD PINN, which incorporates the LS solver in the training process while keeping all other parameters the same as the simple PINN.
• Figure 3: Evolution of training losses and validation errors for different training configurations (related to Figure \ref{['fig:comparison']}). Loss evolution for (a) the simple PINN, and (b) the proposed LS-GD PINN. (c) Validation errors compared to the finite difference reference solution. Training was repeated with three different collocation point counts per epoch. The LS solver improved convergence, even with very few points, where the simple PINN failed after 50,000 epochs.
• Figure 4: Evolution of training losses and validation errors for different values of $P$, the numbers of neurons in the penultimate layer (related to Figure \ref{['fig:comparison']}). $P$ directly determines the number of weights estimated by the LS step. In all cases, the LS solver improves convergence and achieves lower errors than the standard PINN, even for smaller values of $P$.
• Figure 5: Comparison of simple PINN and LS-GD PINN results for a 30 Hz scattered wavefield simulation on Marmousi model. (a) The selected portion of the Marmousi velocity model. (b) The FD solution. (c) The simple PINN prediction. (d) The LS-GD PINN prediction. (e) Evolution of training losses. (f) Evolution of validation errors, excluding the PML region.