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

In wave equation modeling, wavefields propagating over time can be regarded as feedforward in a recurrent neural network. Therefore, the seismic inversion problem based on partial differential wave equations can be addressed using automatic differentiation in the state-of-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 2D and 3D 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/GPU memory occupation.