Memory Optimization in RNN-Based Full Waveform Inversion Using Boundary Saving Wavefield Reconstruction

• Published in: IEEE transactions on geoscience and remote sensing Vol. 61; pp. 1 - 12
• Main Authors: Wang, Shaowen, Jiang, Yong, Song, Peng, Tan, Jun, Liu, Zhaolun, He, Bingshou
• Format: Journal Article
• Language: English
• Published: New York IEEE 2023; The Institute of Electrical and Electronics Engineers, Inc. (IEEE)
• Subjects: Automatic differentiation; Back propagation; Back propagation networks; Backpropagation; boundary saving; Computation; Computational modeling; Computer applications; Computer memory; Deep learning; Differential equations; Differentiation; full waveform inversion; Graphics; Graphics processing units; Image reconstruction; Linear programming; Mathematical models; Modelling; Neural networks; Optimization; Propagation; Recurrent neural networks; Seismic surveys; Wave equations; Wave propagation; Waveforms
• ISSN: 0196-2892; 1558-0644
• DOI: 10.1109/TGRS.2023.3317529

In wave equation modeling, wavefields propagating over time can be regarded as feedforward in a recurrent neural network (RNN). Therefore, the seismic inversion problem based on partial differential wave equations can be addressed using automatic differentiation in the state-of-the-art deep learning frameworks, eliminating the need for explicit backpropagating the residual wavefield. However, one challenge that arises in the context of automatic differentiation is the significant memory usage due to the necessity of storing the hidden states of the RNN (i.e., wavefields in seismic modeling) during forward computation for constructing the computational graph and computing the derivatives during backpropagation. This memory overhead can become a bottleneck, particularly when dealing with large-scale inversion problems. To mitigate this issue, we propose an effective boundary-saving strategy that allows for the reconstruction of the computational graph during the backpropagation process. Instead of storing all the intermediate wavefields at each time step, we selectively save the necessary information at the boundaries, thereby significantly reducing the memory footprint. This approach enables us to maintain the convenience and efficiency of automatic differentiation computations while minimizing the memory requirements. Both 2-D and 3-D numerical experiments validate the accurate reconstruction of wavefields with minimal loss in precision, while the computational graph is simultaneously reconstructed. Consequently, the gradients can also be calculated correctly by automatic differentiation with minimal CPU/graphics processing unit (GPU) memory occupation.