WaveDiffusion: Exploring Full Waveform Inversion via Joint Diffusion in the Latent Space

Authors: Hanchen Wang, Yinpeng Chen, Jeeun Kang, Yixuan Wu, Young Jin Kim, Youzuo Lin
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Full waveform inversion, Diffusion model, Partial differential equation

Abstract:
Full Waveform Inversion (FWI) is a vital technique for reconstructing high-resolution subsurface velocity maps from seismic waveform data, governed by partial differential equations (PDEs) that model wave propagation. Traditional machine learning approaches typically map seismic data to velocity maps by encoding seismic waveforms into latent embeddings and decoding them into velocity maps. In this paper, we introduce a novel framework that reframes FWI as a joint diffusion process in a shared latent space, bridging seismic waveform data and velocity maps. Our approach has two key components: first, we merge the bottlenecks of two separate autoencoders—one for seismic data and one for velocity maps—into a unified latent space using vector quantization to establish a shared codebook. Second, we train a diffusion model in this latent space, enabling the simultaneous generation of seismic and velocity map pairs by sampling and denoising the latent representations, followed by decoding each modality with its respective decoder. Remarkably, our jointly generated seismic-velocity pairs approximately satisfy the governing PDE without any additional constraint, offering a new geometric interpretation of FWI. The diffusion process learns to score the latent space according to its deviation from the PDE, with higher scores representing smaller deviations from the true solutions. By following this diffusion process, the model traces a path from random initialization to a valid solution of the governing PDE. Our experiments on the OpenFWI dataset demonstrate that the generated seismic and velocity map pairs not only exhibit high fidelity and diversity but also adhere to the physical constraints imposed by the governing PDE.

Decision: reject
Reviewer scores: 6, 3, 6, 3

Reviews and discussion:
<thread>
creation date:
1729725381192

thread_type:
official_review

<thread_object>
creator: reviewer_rgtj

<content>
summary:
This paper deals with the problem of full waveform inversion.
There are two mechanisms that the paper proposes.
1. The paper uses the same latent space for both the model and the data
2. They train a diffusion model in the latent space. Such a diffusion model can therefore generate a plethora of models and their data.

Results look reasonable even though the models that are being trained on are very simple.

strengths:
The idea of using a joint feature space is good and then using a diffusion model on this space is also a good idea. The results are interesting and it seems that the approach works for the models in the data base.

weaknesses:
Unfortunately, the idea of using an AE with common feature spaces for data and model is not new. See https://paperswithcode.com/paper/paired-autoencoders-for-inverse-problems
This is the main problem that the paper have. I understand that in this fast moving field some papers are missed but in this case, the work that was already done makes much of the paper not relevant.
I would recommend the authors to withdraw the paper, concentrate of the diffusion aspect of the paper and resubmit to a different venue.

questions:
The interesting parts of the paper are actually hiding towards the end.

1. How do you actually do coarse to fine?

2. Given some data $d$ how to you use diffusion to find an appropriate model

I would recommend re-writing the paper with section 3.3 in mind. Since training a dual AE is not very innovative and using diffusion of the latent space is not very innovating, the innovation is exactly what you do in 3.3. You could easily develop it to a full paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730429004519

thread_type:
official_review

<thread_object>
creator: reviewer_sthj

<content>
summary:
Full waveform inversion (FWI) is a seismic imaging technique that traditionally reconstructs the subsurface velocity model by iteratively comparing observed and predicted seismic data. More recently, machine learning-based approaches would solve FWI by treating it as an image-to-image translation problem. Furthermore, generative diffusion models mainly treated FWI as a conditional generation problem where the velocity map is generated from a given seismic data. This paper offers a new perspective on FWI by considering it as a joint generative process. Namely, the paper considers whether the two modalities -- seismic data and velocity map -- can be generated simultaneously. Two key steps are proposed: first, a dual autoencoder encodes the two modalities in a shared latent space that provides a coarse approximation of the wave equation solution. Second, a diffusion process in the latent space refines the coarse latent representations which are later decoded into seismic data and velocity maps. In contrast to seismic-velocity pairs generated by the conditional models which often lack physical consistency, the jointly generated pairs approximately satisfy the governing PDE without any additional constraint. The paper's main goal is to offer a new perspective by extending FWI from a conditional generation problem to a joint generation problem.

strengths:
- The proposed paper is well-organized and the idea is clearly presented.
- The paper offers a new perspective on the FWI generation problem by simultaneously generating two modalities -- seismic data and velocity maps -- from the shared latent space. This is a novel idea in contrast to the existing related work which treats these two modalities separately.
- Treating seismic data and velocity maps separately limits the ability to generate physically consistent seismic-velocity pairs. In contrast, jointly generating these modalities makes them approximately consistent with the governing PDE that describes the relationship between them.
- The extensive experiments confirm the soundness of the proposed method and show that the jointly generated seismic-velocity pairs can be a useful supplement to real training data.

weaknesses:
I think there are three main problems in the experiments:
- The method wasn't compared to any existing conditional generative methods. There is even a section 4.2.4. that compares separate vs. joint diffusion but there the separate diffusion was the same model as for the joint diffusion but with a single branch kept active and the latent space no longer shared. I think it would be useful to see how the proposed method compares to the existing methods (e.g., [1]) both in terms of the diversity of the generated data and the performance of the reconstruction methods when trained on the generated data.
- The results might also differ based on a different reconstruction method other than InversionNet (e.g., [2] and/or [3]). I think it would be beneficial to add at least one additional data-driven solver.
- Some of the experiments in the results section do not seem to be realistic. (see more in the questions section)

---

[1] F. Wang, X. Huang, and T. A. Alkhalifah. "A prior regularized full waveform inversion using generative diffusion models." IEEE Transactions on Geoscience and Remote Sensing, 61:1-11, 2023.

[2] P. Jin, X. Zhang, Y. Chen, S. Huang, Z. Liu, and Y. Lin. "Unsupervised learning of full-waveform inversion: Connecting CNN and partial differential equation in a loop." ICLR, 2022.

[3] Z. Zhang, Y. Wu, Z. Zhou, and Y. Lin. "VelocityGAN: Subsurface velocity image estimation using conditional adversarial networks." In 2019 IEEE Winter Conference on Applications of Computer Vision (WACV), 2019, pp. 705-714.

questions:
1. Could you comment on the comparison with the existing conditional generative models? Why didn't you compare to any of the existing methods at least in 4.2.4 section?
2. Could you comment on the choice of the reconstruction method? I think it would be beneficial to add at least one additional data-driven solver. It would be interesting to see how the reconstruction methods work with data generated by different generative models.
3. The generative model was trained on the same OpenFWI dataset on which InversionNet was later evaluated. What is the amount of data your generative model should be trained with and how does it compare to the size of a dataset reconstruction methods (e.g., InversionNet) should be trained with? If the size of a dataset for reconstruction methods is satisfying what is the rationale of doing this? I think you should address the limitations of such a setup.
4. In continuation to the previous question, how realistic is the Gen+1\% case? In this case, you trained your generative model on the same data distribution as in the 1\% of the original dataset. If a real dataset is small, wouldn't it be more realistic to train your generative model with real data that differ from the distribution in the small dataset? Maybe a more realistic case would be to train the generative model on the two subsets and add 1\% of the third subset of OpenFWI. Could you comment on this? What are the implications of the existing setup for real-world applications of the method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730541526227

thread_type:
official_review

<thread_object>
creator: reviewer_a1pp

<content>
summary:
The paper introduces a new framework for Full Waveform Inversion (FWI) that uses a joint diffusion process in a shared latent space. This approach merges the bottlenecks of two separate autoencoders (one for seismic data and one for velocity maps) into a unified latent space.

strengths:
The paper is well written and the diffusion approach in the latent space is an interesting extention to dual autoencoder approaches. With convincing results to support this research.

weaknesses:
My major problem with this manuscript is that the main approach to generating two joint autoencoders is not novel. A similar approach including similar experiments has been proposed and published the approach on dual autoencoder before this submission (https://arxiv.org/pdf/2305.13314) and another publication on dual autoencoder can be found at (https://arxiv.org/pdf/2405.13220). These contributions are neither acknowledged nor cited.  The remaining novelty is the diffusion process within the latent spaces which is by itself an interesting idea and should have been stated as the contribution of this manuscript.

questions:
Please address the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730666321790

thread_type:
official_review

<thread_object>
creator: reviewer_1caw

<content>
summary:
The manuscript introduces a new approach to invert acoustic wave equation data based on a joint generative process. Although there were earlier papers on the use of generative models for data inversion, the presented approach looks fairly original. The authors study the famous geophysical problem known as full waveform inversion (FWI). The approach was tested on 2D spatial data from public dataset OpenFWI.

strengths:
Generative AI is transforming different industries in our days and its use for data inversion looks like a promising research direction. Both theoretical and experimental parts are well-present and easy-to-follow. An important original feature of the work is joint generation of acoustic data and velocity models.

weaknesses:
1.	It is not clear (at least none of the experiments show this) how to use the presented algorithm to invert actual data. It is shown how to generate acoustic data and velocities. But what is typically expected by the reader is the answer on what to do when we are given with some specific seismic data.
2.	Section 4.2.3 Comparison with Inversionnet is not sufficiently complete and convincing. See Questions below.
3.	The geophysical terminology is mixed in the manuscript. Notice the used wave equation models $acoustic$ data. This is a significant simplification of seismic phenomena. In other words, the terms $acoustic$ and $seismic$ are not interchangeable.

questions:
1.	How acoustic data and velocity models were preprocessed before training?
2.	The authors trained the model for 1000 epochs. How long was it in terms CPU/GPU time (depending on the dataset)?
3.	The discussion in the manuscript covers only generation and inversion of 2D spatial data. While 3D models/data are of much higher interest.  Could the proposed algorithm be used in the 3D case? What will the implication on computational complexity?
4.	An important test for an inversion code is to check that symmetric data with respect to some plane produces a symmetric velocity model. Would the presented generation model obey this principle?
5.	It is not clear from the experiments how the presented algorithm compares to baselines. Section 4.2.3 Comparison with Inversionnet does not give the answer on the obvious question: which of the two algorithms is better.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732325585176

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Submission of Revised Manuscript and Summary of Revisions

comment:
Dear Reviewers,

We hope this message finds you well. We are writing to inform you that we have uploaded the revised version of our manuscript based on your insightful feedback and suggestions. We greatly appreciate the time and effort you invested in reviewing our work and providing constructive comments, which have significantly improved the quality and clarity of our paper. Below is a summary of the main revisions addressing your feedback:

**Section 4.6 (page 9) FWI with Seismic Data Only:**

We added an experiment utilizing the “one-in-two-out” WaveDiffusion model to perform conventional FWI tasks when only seismic data is available. This experiment highlights the model’s capability to invert seismic data into velocity maps effectively.
(***Addressing comments from reviewers 1caw, rGTJ, and STHJ.***)

**Line 687 (page 13) Training Time Details:**
Details about the training times for the autoencoders (1000 GPU hours) and the joint diffusion model (2000 GPU hours) have been added.
(***Addressing reviewer 1caw.***)

**Appendix A.6 (page 16) Comparison with VelocityGAN and UPFWI:**
We conducted additional experiments comparing WaveDiffusion-generated samples with VelocityGAN and UPFWI. This comparison evaluates the effectiveness of our generated samples under different setups.
(***Addressing comments from reviewers 1caw and STHJ.***)

**Section 3.2 (page 4) Paper Novelty Clarification:**
We have cited the suggested references and clarified that the dual autoencoder serves as a preliminary step for implementing the joint diffusion process, rather than being the main contribution of the paper. The primary contribution of our work lies in Section 3.3 (page 4) and is further elaborated upon through the subsequent analysis and experiments. This revision emphasizes the novel aspects of our joint diffusion framework and provides appropriate context for the dual autoencoder's role within the overall methodology.
(***Addressing comments from reviewers rGTJ and A1pp.***)

**Line 105 (page 2) Clarification on Terminology:**
We included a statement clarifying that acoustic seismic data, representing a simplified wave phenomenon, is used in this study.
(***Addressing reviewer 1caw.***)

**Line 688 (page 13) Data Preprocessing Details:**
Details about the data preprocessing steps, including resizing and normalization, have been added for clarity.
(***Addressing reviewer 1caw.***)

**Appendix A.4 (page 14) Symmetry in FVB:**
Joint generation examples across multiple datasets, including FVB, were added to demonstrate the model’s ability to preserve the symmetry of flat-layer structures.
(***Addressing reviewer 1caw.***)

**Appendix A.5 (page 15) Handling Noisy Input Data:**
We included experiments demonstrating the model’s ability to reconstruct clean, noise-free data from noisy inputs, showcasing WaveDiffusion’s robustness under challenging data conditions.
(***Addressing reviewer STHJ.***)

We have carefully revised the manuscript to incorporate your feedback and ensure the improvements are well-documented. If you have any further comments or questions, we would be grateful for your insights. Thank you again for your valuable input and for considering our work.

We look forward to your thoughts on the revised manuscript.

Best regards,

The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732403026011

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional Experiment: Handling Noisy Seismic Data

comment:
Dear reviewers,

We would like to share a new discovery by our new experiement that emerged during the course of our revisions: the ***robustness*** of *WaveDiffusion* in ***handling noisy input seismic data in FWI***. While this was not explicitly raised in the initial review, we believe this experiment adds significant value to the paper and further demonstrates the versatility of our method.

In this experiment, we tested *WaveDiffusion*, InversionNet, and VelocityGAN on seismic data corrupted with Gaussian noise (mean=0, std=0.05), a realistic scenario often encountered in field data acquisition. Our results reveal that *WaveDiffusion* robustly refines the noisy seismic data and accurately inverts the corresponding velocity maps, achieving the lowest MAE (0.2227) and RMSE (0.3776) and the highest SSIM (0.6142) among the three methods. Neither InversionNet nor VelocityGAN, which rely on direct image-to-image mapping, were able to handle the noisy inputs effectively.

The visual comparisons are shown in the figure at https://bit.ly/3V5k0Uc, and the quantitative results are detailed in the table at https://bit.ly/4f0rAqB. The new experiment and results are added to the latest revised manuscript on **Page 19 A. 8** section.

Despite being trained only on clean seismic data, *WaveDiffusion* demonstrates its ability to refine noisy latent representations, providing both noise-free seismic data and high-fidelity velocity maps. This capability further emphasizes the strength of the joint diffusion process in real-world scenarios.

We believe this additional evidence underscores the potential of *WaveDiffusion* as a robust and practical tool for FWI tasks, extending its applicability to challenging noisy data conditions.

Best regards,

The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732517212683

thread_type:
official_comment

<thread_object>
creator: reviewer_rgtj

<content>
title:
I stay with my score

comment:
I find the claim that the joint diffusion being the main idea of this paper a bit ridiculous. The idea is discussed from lines 192 to 211, barely a third of a page. 
Again, I believe that there maybe novelty in joint diffusion but clearly, the authors did not intend this to be the main contributions originally and the paper was written accordingly. 
I understand the pressure to publish at these main conferences, but to be honest, the authors will do more justice to the paper and their own work if they resubmit the paper to the next conference in line. They can revise the paper, be serious about their latest findings (rather than write the referees that they just discovered something major) and write a winning paper that will last.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper addresses full waveform inversion by training autoencoders on models and data, with shared latent spaces, followed by training a diffusion model in the latent space to generate samples. While the reviewers acknowledged the idea as reasonable and potentially effective, the contribution and presentation fall significantly below ICLR standards.

Two reviewers noted that paired autoencoders have been applied to inverse problems before, and multimodal generative models have long explored similar ideas (some of which the authors acknowledge). Unfortunately, these reviewers did not substantively engage in the discussion (please see below), which rightly frustrated the authors, but the core critique remains valid.

The authors emphasize joint diffusion as the primary contribution, but this is effectively just latent diffusion applied to an highly stylized problem (layered phantoms with few degrees of freedom). The two channels are generated by two decoders whose semantics are entirely determined by the first stage. Whether or not paired autoencoders are considered part of the contributio, the method relies on them—or equivalent architectures—in such a way that the remainder is not novel enough for publication. Once a joint latent space is learned, various generative models could be employed. 

Another key issue is the lack of discussion on the tradeoffs between conditional and joint generation. There are numerous involved tradeoffs which are not mentioned, let alone investigated through experiments. The paper’s claim of tackling the harder problem of simultaneous generation across modalities is valid, but it is important to know why? What are the tradeoffs? How do they manifest? What are the right experiments or possible theoretical insights? Joint generative modeling is generally harder than discriminative, where observations localize or guide the generator, especially if there is a cheap approximate inverses.

The problem of stability, key for inverse problems, is unaddressed. It would be informative to train a model with noisy measurements (at varying noise levels), which is different from inputting noisy data into a trained model. The current joint distribution is degenerate—measurements are functions of unknowns and we can sample joint dist by generating models and training an FNO.The authors report thousands of GPU hours to train their models on simple phantoms, with massive network architectures, whereas much smaller nets have achieved strong results in similar settings. The justification for this computational expense is unclear.

The paper’s presentation is another issue. Design choices are not clearly motivated or lack explanation. The VQ is a crucial part of the model. The authors say it ensures a "compact and discrete representation," but what does this mean? Why does VQ enable the model to "capture structured patterns and long-range dependencies"? Also, what exactly is being vector-quantized? The supplementary mentions 8192 32-dimensional vectors—is that for the entire 16×16 latent? If it is then in Section 4.6 the reconstruction consists in picking one out of 8192 templates. What is really the case remains unclear.

Same is the case for the latent-space diffusion ("joint diffusion"). It is described so tersely and handwavily that it is unclear how the model is trained; the supplementary contains information about hyperparameters and learning rates but not about inputs, outputs, losses, ... . The only equation in the first block is z_t = z + eps_t (which is missing a _t), but in the second block t seems to be the terminal time, L is never defined so we don't know what it is, ..., there is little chance to reproduce results.

Experiments are also problematic. The method is tested only on simple layered phantoms (very few DoFs) with an acoustic wave equation, comparing to a couple of deep learning baselines (Table 5). Other experiments analyze the proposed method, without establishing its utility. The authors claim they are not focused on achieving strong reconstruction results, but no clear alternative insights emerge. The experiments with separate AEs show that learning the joint distribution of synthetic data is simpler than solving the ill-posed full waveform inversion. It would be interesting to navigate this spectrum. The paper's repeated references to "seismic data" are misleading, as it uses simulated 2D acoustic data from 64x64 layered toy models. Real-world geophysical models are far large, 3D, and waves are elastic. While academic compute limits are understood (although this does not seem to be a block here), the extensive resources for such stylized setups are hard to justify.