Kronecker neural networks for the win: Overcoming spectral bias for PINN-based wavefield computation

Wavefield computation constitutes the majority of the computational cost for seismic applications, including reverse-time migration and full waveform inversion (FWI). One of the popular approaches is to solve for the wavefields in the frequency domain by using the Helmholtz equation. However, Helmholtz solvers require inversion of a large stiffness matrix that can become computationally intractable for large 3D models or in the case of modeling high frequencies. Recently, researchers have explored the potential of physics-informed neural networks (PINNs) in solving the Helmholtz equation with limited success. While a number of attractive features have been demonstrated for the PINN-based Helmholtz solvers, their large training time has been the main impediment in their widespread adoption for wavefield computations. This is mainly caused by the spectral bias of neural networks, which poses difficulty in training the PINN model for high-frequency wavefields. We employ Kronecker neural networks (KNNs) that form a general framework for neural networks with adaptive activation functions. We implement it using a standard feed-forward neural network employing a composite activation formed by using the inverse tangent (atan), exponential linear unit (elu), locally adaptive sine (l-sin), and locally adaptive cosine (l-cos) activation functions. Thanks to the oscillatory noise added by the sine and cosine terms, the network is able to explore more and learn faster. This allows the network to get rid of saturation regions from the output of each layer and overcome slow convergence. Through numerical tests, we show the efficacy of the proposed approach in fast and accurate wavefield modeling in the frequency domain.