A Unified Framework for Forward and Inverse Problems in Subsurface Imaging using Latent Space Translations

Comment 1: This paper resembles a survey and might be more suitable for journal submission. Aside from extensive experiments on factor influences, the work offers limited novelty and lacks theoretical analysis.

One of the key innovations of our work is the novel Generalized Forward-Inverse (GFI) framework, which unifies forward and inverse problems in subsurface imaging. Within this framework, we propose two new architectures—Latent U-Net and Invertible X-Net—that achieve state-of-the-art performance on multiple datasets, including OpenFWI benchmarks. By systematically addressing forward and inverse problems, GFI bridges existing methods and provides a systematic approach to enhancing performance in scientific machine learning tasks.

Similar to recent works such as "A Unified and General Framework for Continual Learning" (ICLR 2024) [1] and "Mirror Learning: A Unifying Framework of Policy Optimisation" (ICML 2022) [2], which highlight the importance of unifying diverse methodologies under a single framework, our GFI framework provides a comprehensive lens to interpret, compare, and extend approaches to subsurface imaging. These frameworks demonstrate the value of unifying state-of-the-art approaches while introducing novel components, just as we do with Latent U-Net and Invertible X-Net. Moreover, our work complements practical advancements with theoretical insights into invertibility and latent space translation. This dual focus aligns with the trends seen in leading ML venues, reinforcing the relevance and novelty of our contributions.

Along with introducing GFI, we also propose a novel Invertible X-Net architecture that is the first state-of-the-art model designed to jointly solve forward and inverse problems using a single-step unified approach. Unlike prior methods such as InversionNet or AutoLinear, which treat these tasks independently, X-Net leverages a shared invertible latent space that allows for consistent and efficient bidirectional mapping between the spatial (velocity maps) and spatio-temporal (seismic waveforms) domains. This joint optimization is a key innovation, as it ensures that the learning for one task benefits the other, leading to more robust and accurate results across both tasks.

Regarding the lack of theoretical analysis, we have detailed the theoretical foundations of GFI in Appendix Section C, where we explain how invertibility and latent space translation enable this unified approach.

[1] Wang, Z., Li, Y., Shen, L., & Huang, H. A Unified and General Framework for Continual Learning. In The Twelfth International Conference on Learning Representations, ICLR 2024.

[2] Grudzien, Jakub, Christian A. Schroeder De Witt, and Jakob Foerster. "Mirror learning: A unifying framework of policy optimisation." International Conference on Machine Learning. PMLR, 2022.

Comment 2: Additional factors should be investigated, such as the robustness of data noise across different network architectures.

We thank the reviewer for this valuable suggestion. In response, we conducted comprehensive noise evaluation experiments to assess the robustness of Latent U-Net and Invertible X-Net under varying levels of Gaussian noise in the input data. For these experiments, Gaussian noise with varying variances was added to the input velocity or waveform, and the models’ predictions were evaluated based on these noisy inputs.

Detailed results for the CurveFault-B and FlatFault-B datasets are provided in Appendix Tables 11–14, covering both forward and inverse tasks. These results show that Latent U-Net and Invertible X-Net maintain strong performance even with low PSNR of input signals, exhibiting minimal degradation in accuracy compared to contemporary methods. The performance of Latent U-Net and Invertible X-Net for both forward and inverse problems at the noisiest signal level is better than the Auto-Linear predictions at the zero noise level.

For the inverse problem, the noise levels correspond to PSNR values of 84.07 dB, 77.08 dB, 74.07 dB, and 67.08 dB. Both Latent U-Net and Invertible X-Net demonstrate robust performance under varying noisy conditions. For instance, on CurveFault-B, the Latent U-Net and Invertible X-Net show degradation in SSIM of about 1.24% and 1.13% respectively at the noisiest signal level (PSNR = 67.08 dB), outperforming Auto-Linear with significant margin.

For the forward problem, the noise levels correspond to PSNR values of 56.02 dB, 49.03 dB, 46.02 dB, and 39.03 dB. Both Latent U-Net and Invertible X-Net demonstrate strong robustness at all noise levels. For example, on CurveFault-B, both Latent U-Net and Invertible X-Net show similar degradation in SSIM of about 0.82% and 0.73% respectively at highest noise level (PSNR = 39.03 dB). The performance of these models across all noise levels exceeds that of the Auto-Linear at the zero noise level for all metrics, indicating significant improvement for forward problem as well.

Comment 1: The experiment 5.1.3 addresses the question ‘should we train latent spaces solely for reconstruction?’, which I assume corresponds to the question 3) formulated in line 083 ‘what is the role of manifold learning?’. 1.1 For increased clarity I’d suggest not changing the phrasing of the questions

We apologize for the inconsistency in our section titles. We have now revised the phrasing of questions to ensure consistency throughout the paper. Section 5.1.3 is now renamed as “What is the role of manifold learning in DL4SI?”

Comment 2: What are training fractions? A fraction of the training set or a fraction of the training time? And why is it instructive to look at 5% and 10% training fraction?

Training fraction refers to the subset of the dataset used for training, where the remaining data is masked and not utilized. In our experiments, we followed the same precedent for training with lesser data as followed in AutoLinear, which evaluated model performance using 5%, 10%, and 100% training fractions on the same dataset. Evaluating model performance across different training fractions provides insights into model robustness and efficiency, particularly in data-constrained scenarios that are often encountered in real-world applications.

Comment 3: Importance of manifold learning in the GFI framework and inclusion of reconstruction loss for evaluation.

We appreciate the reviewer’s suggestion to explore a third learning objective combining reconstruction and translation losses. Following this, we trained all three models—Latent U-Net (Large), Latent U-Net (Small), and Invertible X-Net—on this combined objective across various datasets. These results have been included in the updated Figure 6 in the main paper and Figure 22 in the Appendix. Our findings indicate that directly learning to translate yields better performance for translation tasks compared to a two-step approach of first learning a latent space optimized for reconstruction followed by translation. Directly learning to translate (without using reconstruction loss) is also better than simultaneously optimizing both reconstruction and translation loss terms. This suggests that translation and reconstruction can have competing objectives, where a manifold optimized for reconstruction may not be best-tailored for translation. Hence, manifold learning, while beneficial for reconstruction, may lead to suboptimal performance when translation is the primary goal. We have now included this discussion in the main paper.

Comment 4: The idea of the Latent U-net seems to be very similar to the InversionNet following the comparison in Table 1. So to my understanding the difference between the two boils solely down to the architectural design and size of the network but not to any conceptual differences. If that is the case, it should be stated more directly in the paper and the exact architectural differences should be explained in more detail.

We agree with the reviewer that the differences between Latent U-Net and InversionNet primarily stem from architectural design. As suggested by the reviewer, InversionNet can be viewed as a special case of Latent U-Net, lacking skip connections and incorporating a smaller bottleneck layer. Both InversionNet and Latent U-Net (Large) have similar numbers of parameters. We also agree that some of the middle layers of InversionNet are performing latent space translation. However, the Latent U-Net (Large) has more layers in its translation model (U-Net) while having lighter encoder-decoder models compared to InversionNet. Architectural differences, such as skip connections, encoder/decoder size, latent dimension size, and the choice of translation architecture (e.g., U-Net vs. iUNet) play significant roles on model performance, as explored through the four salient questions of our paper. Note that GFI represents a spectrum of methods where InversionNet is just one point in the spectrum. By drawing connections between InversionNet and Latent U-net, we are expanding the space of model choices available for DL4SI.

Thank you for your responses.

I went through the revised versions and find that the clarity of certain statements is improved as asked for in my review. Further, I think the additional experiments with the joint reconstruction and translation loss are interesting and improve the discussion of the question on training latent spaces for reconstruction.

I adjusted my score accordingy. In summary I would recommend the paper for acceptance as I think the unifying framework and the advancement of the state of the art is interesting for the audience. The reason for not further raising the score is that I find the novelty of the best performing method (latent U-net) to be somewhat limited to equipping the idea behind InversionNet with a better network architecture (which I still think should be communicated and compared more clearly in the paper).

Comment 3: While the authors aim to model the forward problem using neural networks, it is uncertain if the learned mapping can capture complex physical dynamics, as seen in the governing equations of seismic data. Extensive out-of-domain experiments are recommended to assess the model’s robustness and generalizability to real-world data.

We agree with the reviewer that extensive out-of-domain experiments are required to validate our performance on forward problems. Toward this goal, we have devoted an entire Section 5.2 to evaluate zero-shot generalization capabilities of our models on both forward and inverse problems. First, we demonstrate out-of-distribution generalization across all OpenFWI datasets, where a model trained on a certain dataset is tested on all remaining datasets (see Appendix E.4 to E.6 for additional visualizations). Second, we perform zero-shot evaluation on two real-world benchmarks widely used in seismic imaging studies, Marmousi and Overthrust. Quantitative results for these datasets are presented in Appendix Section E.2, with visualizations included both in the main paper and the Appendix. Similar out-of-domain evaluations have been performed in previous state-of-the-art works as well [1].

To further assess the robustness of our models, we have conducted new experiments analyzing the effects of adding noise to the inputs of forward and inverse problems. In these experiments, Gaussian noise with varying variances is added to the input velocity or waveform, and the models' predictions are evaluated based on the noisy inputs. We observe that for the inverse problem, both Latent U-Net and Invertible X-Net demonstrate robust performance even under noisy conditions. In the forward case, Invertible X-Net has the highest SSIM and the lowest MAE even at the lowest PSNR signal. Detailed results are provided in Appendix Tables 11–14 for CurveFault-B and FlatFault-B datasets, covering both forward and inverse problems.

[1] Y. Feng, Y. Chen, P. Jin, S. Feng, Y. Lin, "Auto-Linear Phenomenon in Subsurface Imaging ." ICML, PMLR:235, 2024, pp. 13153-13174.

Comment 4: The paper's structure could be streamlined to highlight the proposed methodology. Currently, a substantial background section precedes a brief description of the methods. Specific sections of the introduction and related work and preliminary knowledge could be condensed.

We appreciate the reviewer’s suggestion and will try to streamline the introduction and related works sections to reduce redundancy and increase focus on our contributions. However, we would like to emphasize that the main contribution of our work is introducing the unified framework of GFI, which subsumes previous works in the field that can be viewed as special cases of GFI. As a result, we begin Section 3 on “Proposed Methodologies” by first introducing GFI and then describing how previous works are different instantiations of GFI (which is also a novel contribution), before discussing the proposed architectures of Latent U-Net and Invertible X-Net. We will revise Section 3 to bring more prominence to our contributions and emphasize the connections of prior works to GFI.

Comment 5: The model specifications lack detail; adaptations made to standard UNet for subsurface imaging and the mechanisms enabling X-Net’s invertibility are unclear. Clarifying these aspects and specifying when each model might be preferred would improve understanding.

To provide more details of our model specifications, we have added two new Sections in the Appendix (D.3.1 and D.3.2) that explain how U-Net adaptations are tailored for subsurface imaging and summarize the architecture of U-Nets and iUNets used as backbones in our proposed models, including the type and number of blocks and parameter counts (also described in Appendix Tables 3 to 6). Additionally, the code has been made available to ensure full transparency and reproducibility. Our approach leverages the strengths of U-Nets, widely successful in computer vision, to be applied in the domain of subsurface imaging.

Regarding the comment on when we should choose one model over another (e.g., Invertible X-Net over Latent U-Net), we would like to refer to our global response on this comment. In summary, we recommend using Invertible X-Net if the goal is to solve both forward and inverse problems, or just the forward problem. We recommend Latent U-Net (Large) if the goal is to only solve the inverse problem. We have included this explicit recommendation in the main paper.

Comment 2: Comparisons to prior work are outdated, with main competing methods from 2019, like InversionNet and VelocityGAN. Experimental comparisons with recent baselines (published within the last 2-3 years) would add rigor to the evaluation, as several recent methods could serve as relevant benchmarks. To name a few, “Physics-Informed Robust and Implicit Full Waveform Inversion Without Prior and Low-Frequency Information 2024”; “Full-waveform inversion using a learned regularization 2023”; “Wasserstein distance-based full-waveform inversion with a regularizer powered by learned gradient 2023” and more.

We thank the reviewer for highlighting these impactful works [2,3,4], which we are happy to include in our related works discussion. While these methods focus on solving inverse problems using physics-informed priors, they can be categorized as per-sample methods according to the taxonomy introduced in [1]. In particular, a different model is optimized in these methods to solve a single sample of the inverse problem (e.g., a single waveform) independently of the other samples. In contrast, our GFI framework belongs to a completely different category of per-domain approaches, where the trained model can solve forward and inverse problems for any input sample. This offers generalizability across the entire domain of input and output samples, while enabling fast “online” predictions without per-instance optimization.

We did not compare our results with these prior works as per-sample methods are outside the scope of GFI, where we are interested in analyzing shared latent spaces for jointly solving forward and inverse problems. Direct comparisons were also not feasible as their code and model checkpoints are not publicly available.

Regarding the comment on using recent baselines, we would like to emphasize that we compare our results with AutoLinear [5], a 2024 ICML publication that is the most recent state-of-the-art baseline in the field.

[1] Lin, Youzuo, et al. "Physics and Deep Learning in Computational Wave Imaging." arXiv preprint arXiv:2410.08329 (2024).

[2] P. Sun, F. Yang, H. Liang and J. Ma, "Full-Waveform Inversion Using a Learned Regularization," in IEEE Transactions on Geoscience and Remote Sensing, vol. 61, pp. 1-15, 2023, Art no. 5920715, doi: 10.1109/TGRS.2023.3322964.

[3] B. Du, J. Sun, A. Jia, N. Wang and H. Liu, "Physics-Informed Robust and Implicit Full Waveform Inversion Without Prior and Low-Frequency Information," in IEEE Transactions on Geoscience and Remote Sensing, vol. 62, pp. 1-12, 2024, Art no. 5918712, doi: 10.1109/TGRS.2024.3416547.

[4] F. Yang and J. Ma, "Wasserstein Distance-Based Full-Waveform Inversion With a Regularizer Powered by Learned Gradient," in IEEE Transactions on Geoscience and Remote Sensing, vol. 61, pp. 1-13, 2023, Art no. 5904813, doi: 10.1109/TGRS.2023.3241723.

[5] Y. Feng, Y. Chen, P. Jin, S. Feng, Y. Lin, "Auto-Linear Phenomenon in Subsurface Imaging ." ICML, PMLR:235, 2024, pp. 13153-13174.

Comment 1: The paper’s narrow focus on subsurface imaging may limit its relevance for a broader audience at ICLR. If there are analogous problems in other fields where their framework could be applied?

Thank you for the suggestion to discuss our work’s broader applicability outside of the domain of subsurface imaging. Our proposed GFI framework is generic enough to be applied in any domain-to-domain translation problem, especially where the input and output domains are from two different spaces with varying dimensionalities. This includes applications in medical imaging, climate modeling, and material property predictions. For example, full-waveform inversion has been successfully applied to neuroimaging, enabling non-invasive 3D imaging of the human brain with sub-millimeter resolution [1]. Our work provides a generalized perspective of solving forward and inverse problems for domain translation, using encoder-decoder blocks coupled with a translation model. While our work analyzes different variations of GFI for applications in seismic imaging, similar analyses can be performed for other domain translation problems. We will include this broader motivation in the revised manuscript.

We would also like to highlight that research in deep learning for subsurface imaging (DL4SI) is gaining growing attention at top-tier venues in machine learning, with influential papers such as AutoLinear [2], OpenFWI [3], UPFWI [4], and InvLINT [5] being published at ICML, NeurIPS, and ICLR. This reflects the broader scope of research in DL4SI beyond the application of seismic imaging.

[1] Guasch, L., Calderón Agudo, O., Tang, M. X., Nachev, P., & Warner, M. (2020). Full-waveform inversion imaging of the human brain. NPJ digital medicine, 3(1), 28.

[2] Y. Feng, Y. Chen, P. Jin, S. Feng, Y. Lin, "Auto-Linear Phenomenon in Subsurface Imaging ." ICML, PMLR:235, 2024, pp. 13153-13174.

[3] C. Deng, S. Feng, H. Wang, X. Zhang, P. Jin, Y. Feng, Q. Zeng, Y. Chen, Y. Lin, "OpenFWI: Large-scale Multi-structural Benchmark Datasets for Full Waveform Inversion," Thirty-sixth Conference on Neural Information Processing Systems (NeurIPS), Datasets and Benchmarks Track, 2022.

[4] P. Jin, X. Zhang, Y. Chen, S. X. Huang, Z. Liu, Y. Lin, "Unsupervised Learning of Full-Waveform Inversion: Connecting CNN and Partial Differential Equation in a Loop," ICLR, 2022

[5] Y. Feng, Y. Chen, S. Feng, P. Jin, Z. Liu, Y. Lin, "An Intriguing Property of Geophysics Inversion," Proceedings of the 39th International Conference on Machine Learning (ICML), PMLR:162, 2022, pp. 6434–6446

Comment 1: Section 5.1.1 states that the ideal size of the latent space should be decided based on the “complexity” of the problem being solved. It is not clear how to determine the complexity in real-world applications.

We apologize for the lack of clarity in our wording. We meant to say that the latent space size is a hyper-parameter that needs be determined based on the *needs” of the problem. Intuitively, problems with simpler geological structures and velocity variations (e.g., A-family of datasets in OpenFWI) can be solved with smaller latent dimensions, since the amount of information that we need to encode is also small. On the other hand, problems with more complex structures (e.g., B-family of datasets) require large latent space sizes to encode more sophistical patterns in the waveform and velocity fields. We will revise Section 5.1.1 to clarify this point.

Comment 2: What is the relation between the number of parameters in encoders and decoders and the U-Net and how does it influence performance?

Appendix Table 6 provides a detailed comparison of the number of parameters in the encoder-decoder and translation models across different baselines. While Latent U-Net (Large), Latent U-Net (Small), and Invertible X-Net maintain a similar ratio between these components, AutoLinear represents a contrasting case with a more complex encoder-decoder and a lightweight translation model. Despite having a similar total number of parameters to Latent U-Net (Large), AutoLinear demonstrates different performance characteristics, indicating that the parameter ratio is not the sole factor influencing performance. Architectural differences, such as skip connections, latent dimension size, and the choice of translation architecture (e.g., U-Net vs. iUNet), also play significant roles, as evaluated in Figure 5.

We also agree that conducting more ablations of Latent U-net with different ratios of encoder/decoder to translation model parameters is an important direction of research, which we could not complete in the time constraints of the discussion period but are happy to explore in future works.

Comment 3: Invertible X-Net was not included in the experiments on varying model size and only one size was used throughout the paper. It would be useful to see how its size influences the performance.

To address this, we have added a new plot in Figure 5(c) of the main paper that shows how Invertible X-Net's performance changes with varying latent dimensions across 8 OpenFWI datasets, similar to the analysis done for Latent U-Net (Large) in Figure 5(a). We observe a similar trend that the performance of Invertible X-Net does not mostly change with latent size for A-family of datasets, while it shows a decreasing trend for B-family.

To answer why we only chose one size for Invertible X-Net while for Latent U-Net we had two versions (small and large), we would like to remind that we wanted to keep the number of parameters in our proposed models (Latent U-net and Invertible X-Net) similar to the range of parameters used in other baselines like InversionNet. Since the iUNet used in Invertible X-Net requires two coupling blocks for every convolution in a standard U-Net, adjusting the overall model size of Invertible X-Net required a significant change in total parameters. For instance, doubling the number of invertible blocks in the iUNet increases parameters from 25M to 50M, while halving it reduces the parameters to 13M. As a result, a larger version of Invertible X-Net would have significantly more parameters than other models, while a smaller version would have too few parameters leading to underfitting.

Comment: Summarize Key Findings from Appendix Tables and Figures in the Main Paper (vWh3, 74KD, nmvk)

Thank you for this suggestion. To summarize the key findings across all 10 OpenFWI datasets used in our paper, we have now added a new column of average rank for every model in the Appendix Tables 7 and 9. We will use these average ranks in our discussions of results throughout the main paper. To improve the readability of the Appendix section from the main paper, we have now also added cross-references in the main paper to specific Appendix sections wherever appropriate.

We hope that our responses to the individual review comments provided below address the main concerns of the reviewers. If we missed any detail, we will be happy to provide more clarifications during the discussion period. If our responses have adequately addressed your concerns, we kindly request that you consider revising your scores. Thank you very much for your time and effort.