Robust Physics-Guided Diffusion for Full-Waveform Inversion

This paper proposes a robust physics-guided diffusion framework for full-waveform inversion that integrates a score-based generative prior with a transport-based data-consistency potential and a preconditioned guided reverse-diffusion scheme to achieve superior reconstruction quality and stability compared to deterministic baselines and standard diffusion posterior sampling.

The Gist

Imagine you are trying to figure out what's inside a giant, opaque cake without cutting it open. You can only tap the top of the cake with a spoon and listen to the sound the vibrations make as they travel through the layers. This is essentially what Full-Waveform Inversion (FWI) does for the Earth. Seismologists tap the ground (using earthquakes or explosions) and listen to the echoes to build a 3D map of the underground rocks, oil, or gas.

However, this is a notoriously difficult puzzle. The "echoes" are messy, the cake has layers of different thicknesses, and sometimes the sound waves bounce around in confusing ways. If you try to solve this puzzle with standard math, you often get stuck in a "local trap"—you think you've found the right answer, but you're actually looking at a fake version of the cake that just looks similar.

This paper introduces a new, smarter way to solve this puzzle using a combination of AI intuition and physics rules. Here is how it works, broken down into simple concepts:

1. The Problem: The "Wrong Turn" Trap

Imagine you are trying to match two songs playing on different speakers. If one song is slightly delayed (even by a fraction of a second), a standard computer trying to match them point-by-point will scream, "These are totally different!" and get confused. In seismology, this is called cycle-skipping. The computer gets stuck trying to fix tiny timing errors instead of seeing the big picture, leading to a blurry or wrong map of the underground.

Also, the "loud" parts of the sound (the first big bangs) drown out the "quiet" parts (the subtle echoes from deep underground). Standard methods listen only to the loud parts and ignore the quiet ones, missing the deep secrets of the Earth.

2. The Solution: A "Physics-Guided" AI Detective

The authors built a system that acts like a detective who has two superpowers:

1. A Library of Imagination (The AI Prior): The AI has "seen" thousands of realistic underground maps before. It knows what a healthy geological structure should look like. It's like having a chef who knows exactly what a perfect cake layering looks like, even without seeing the specific cake you are trying to solve.
2. A Physics Compass (The Guidance): The AI doesn't just guess; it checks its guesses against the actual sound data using the laws of physics.

3. The Three Magic Tricks

The paper introduces three specific "tricks" to make this detective work better than anyone else:

Trick A: The "Volume Knob" (Amplitude Balancing)

In the past, the computer was too obsessed with the loudest sounds. The authors gave the computer a smart volume knob.

• The Metaphor: Imagine a party where one person is shouting and everyone else is whispering. A normal listener only hears the shouter. This new method turns down the volume of the shouter and turns up the volume of the whisperers.
• The Result: The AI now pays attention to the quiet, deep echoes that reveal hidden faults and layers, not just the loud surface noise.

Trick B: The "Shape Shifter" (Optimal Transport)

Instead of comparing the sound waves point-by-point (which fails if they are slightly out of sync), the AI compares the shape of the waves.

• The Metaphor: Imagine you have two piles of sand. One pile is slightly shifted to the left. A standard ruler would say, "These are completely different!" But if you look at the shape of the pile, you realize it's the same sand, just moved a little bit. This method uses Optimal Transport (a fancy math way of measuring how much "sand" needs to be moved to match the shapes).
• The Result: The AI is no longer confused by tiny timing errors. It understands that a wave is the same wave, even if it arrives a split second late. This stops the "local trap" problem.

Trick C: The "Smart Step" (Preconditioned Guidance)

When the AI tries to fix its guess, it usually takes one giant step. But sometimes, the ground is slippery (noisy data), and a giant step makes you fall. Other times, the ground is solid, and you need to walk faster.

• The Metaphor: Imagine walking through a foggy forest. In the beginning, when you can't see well (the early guesses are rough), you take tiny, cautious steps. As the fog clears and you see the path better, you take bigger, more confident strides. Also, if you are walking on a steep hill (a hard-to-see area), you slow down; if you are on flat ground, you speed up.
• The Result: The AI adapts its speed and direction dynamically. It doesn't crash into errors early on, and it doesn't miss details later on.

4. The Result: A Clearer Picture

When the authors tested this new method on standard geological datasets, the results were impressive:

• Sharper Images: The underground maps were much clearer, with distinct layers and faults visible.
• Fewer Mistakes: The AI didn't get stuck in "local traps" as often as older methods.
• Robustness: Even when the data was noisy (like trying to hear a whisper in a storm), the method still worked well.
• Generalization: The AI could even look at completely different types of geological maps (ones it had never seen before) and still do a great job, proving it learned the rules of geology, not just memorized specific pictures.

Summary

Think of this paper as upgrading a GPS system. Old GPS systems would get confused if you took a slightly wrong turn, leading you to the wrong destination. This new system uses a smart map of the world (the AI prior) combined with real-time traffic rules (the physics guidance) and adaptive driving (the smart steps) to get you to the correct underground map, no matter how messy the road conditions are. It's a more robust, smarter, and clearer way to see inside the Earth.

Technical Summary

1. Problem Statement

Full-Waveform Inversion (FWI) is a high-resolution seismic imaging technique used to reconstruct subsurface velocity models by minimizing the misfit between observed and simulated seismic wavefields. However, FWI faces three fundamental challenges:

1. Non-convexity and Cycle-Skipping: Standard pointwise $\ell_2$ misfit functions are highly sensitive to small time/phase shifts between observed and synthetic traces. This leads to a non-convex objective landscape with numerous local minima, causing iterative solvers to converge to incorrect phase alignments (cycle-skipping).
2. Amplitude Imbalance: Seismic data often exhibit large dynamic ranges where early-arriving high-amplitude phases dominate the gradient, causing the inversion to underfit weaker, later arrivals that contain crucial information about deeper structures.
3. Ill-posedness and Heterogeneous Sensitivity: Limited acquisition geometry and band-limited data make the problem ill-posed. Furthermore, sensitivity to velocity variations is highly spatially heterogeneous (e.g., poor illumination at depth), which standard scalar guidance steps in diffusion models fail to address effectively.

While Diffusion Posterior Sampling (DPS) has emerged as a promising Bayesian approach to incorporate learned priors, standard DPS implementations using $\ell_2$ potentials and scalar guidance steps struggle with the specific structural deficiencies of FWI.

2. Methodology

The authors propose a Robust Physics-Guided Diffusion Framework that integrates a score-based generative prior with a modified reverse-diffusion process tailored for FWI. The method consists of two primary innovations:

A. OT-Based Data-Consistency Potential

Instead of using a standard least-squares ($\ell_2$) misfit, the authors design a robust likelihood potential $\Phi(v)$ based on Optimal Transport (OT) and Wavefield Enhancement:

• Bounded Amplitude-Adaptive Weighting: A weight function $\omega(s,r,t)$ is applied to seismic traces based on their instantaneous amplitude relative to the maximum observed amplitude. This suppresses the influence of high-amplitude early arrivals while preserving low-amplitude later phases, preventing gradient dominance by strong events.
• Wasserstein-2 Discrepancy: The misfit is computed using the 1D Wasserstein-2 distance ($W_2$) between probability density functions (PDFs) derived from the weighted traces. By transforming waveforms into PDFs (via shift-rescale normalization) and measuring discrepancy via quantile functions, the method becomes robust to small time/phase shifts, significantly mitigating cycle-skipping.
• Observation-Dependent Normalization: The total misfit is normalized by a characteristic magnitude derived solely from the observed data. This stabilizes the numerical scale of the guidance term, making the inversion less sensitive to hyperparameter tuning.

B. Preconditioned Variable-Metric Guidance

To address the spatial heterogeneity of FWI sensitivities, the authors replace the standard scalar guidance step in DPS with a diagonal preconditioned guidance scheme:

• Adaptive Scalar Schedule ($\rho_i$): The guidance strength is modulated by a factor $\rho_i$ that depends on the "roughness" (Total Variation) of the current denoised estimate. Early in the reverse process (when estimates are noisy), guidance is conservative to prevent instability; as the estimate smooths, guidance strength increases.
• Spatial Diagonal Scaling ($D_i$): A diagonal matrix $D_i$ is constructed based on the magnitude of the current misfit gradient. This rescales updates spatially, balancing the influence of well-illuminated regions against poorly illuminated ones, effectively acting as a variable-metric preconditioner.

The final update rule for the reverse diffusion step $i$ is:
$$v_{i-1} = v'_{i-1} - \rho_i D_i \nabla_{v_i} \Phi(\hat{v}^{(i)}_0)$$
where $\hat{v}^{(i)}_0$ is the denoised estimate of the clean velocity model.

3. Key Contributions

1. Robust OT Potential: Introduction of a novel data-consistency potential combining bounded amplitude weighting and Wasserstein-2 distance, which effectively alleviates cycle-skipping and amplitude imbalance in FWI.
2. Stabilized Variable-Metric Guidance: Development of an adaptive guidance mechanism that dynamically adjusts both the temporal strength (based on estimate roughness) and spatial scaling (based on sensitivity) during the reverse diffusion process.
3. Decoupled Prior Training: The score-based prior is trained solely on velocity model samples without requiring paired (wavefield, velocity) data or embedding a forward solver in the training loop, enhancing computational efficiency and flexibility.
4. Comprehensive Validation: Extensive numerical experiments demonstrating superior performance over deterministic baselines and standard DPS.

4. Experimental Results

The method was evaluated on the OpenFWI benchmark datasets (CurveVel, FlatFault, CurveFault families) and the Marmousi2 dataset.

• Quantitative Performance: The proposed method (OT-WE-PDPS) achieved significantly lower relative $\ell_2$ errors, higher PSNR, and higher SSIM compared to:
  • Deterministic baselines (W2 + TV, OT-WE + TV).
  • Standard DPS with $\ell_2$ guidance.
  • Example: On the CurveVel-B dataset, the proposed method achieved an $\ell_2$ error of 2.04% and SSIM of 0.95, compared to 15.70% and 0.43 for standard DPS.
• Qualitative Improvements: Reconstructions preserved sharp interfaces and fault geometries with fewer artifacts and less over-smoothing than deterministic methods.
• Robustness:
  • Noise: The method remained stable under additive Gaussian noise levels up to $\sigma = 0.05$.
  • Generalization: The model trained on OpenFWI data successfully inverted the Marmousi2 dataset (which was not in the training set) and handled variations in source frequency, receiver depth, and source count without retraining the prior.
• Ablation Study: Confirmed that both the OT-based wavefield enhancement and the preconditioned guidance are necessary; combining them yielded the best results, outperforming either component used in isolation.

5. Significance

This work represents a significant advancement in data-driven seismic inversion. By bridging the gap between generative modeling and physics-based inversion, the proposed framework:

• Solves the Cycle-Skipping Problem: The OT-based metric provides a convex-like landscape for the guidance step, allowing the diffusion process to escape local minima that trap traditional gradient-based methods.
• Enhances Stability: The variable-metric guidance prevents the instability often associated with applying strong data-consistency constraints to noisy intermediate states in diffusion models.
• Offers Practical Flexibility: The ability to use a single, unified prior for diverse geological structures (via mixed-model training) and different acquisition geometries makes the approach highly scalable for real-world exploration scenarios where training data is limited or heterogeneous.

In summary, the paper establishes a new paradigm for Physics-Guided Diffusion, demonstrating that carefully designed data-consistency potentials and adaptive guidance mechanisms can overcome the inherent non-convexity and ill-posedness of Full-Waveform Inversion.