Decoupled Latent Optimization of Diffusion Models for Full Waveform Inversion

Abstract: Full waveform inversion (FWI) recovers subsurface velocity from seismic recordings by solving a severely ill-posed, nonconvex PDE-constrained optimization. Classical regularizers stabilize the inversion but fail to reproduce realistic geological structures; recent diffusion-prior methods improve realism at the cost of a fragile trade-off between data fidelity and prior consistency. We propose Decoupled Latent Optimization (DLO), which relaxes the standard latent-optimization formulation into a quadratic-penalty objective over an auxiliary physical variable and a latent variable. The data-fidelity gradient acts in physical space, the diffusion sampler contributes only through a decoded prior sample, and the standard smoothed-velocity initialization of classical FWI is preserved. On the OpenFWI benchmark, DLO outperforms classical regularizers and existing diffusion-based methods under clean, noisy, and missing-trace acquisitions. The prior, trained on 70*70 OpenFWI models, transfers directly to the Marmousi and Overthrust benchmarks, where DLO recovers intricate fault structures and remains robust to initialization smoothing and measurement noise.

• The paper proposes a quadratic-penalty strategy to decouple physical data fidelity from diffusion-based prior guidance.
• It integrates PDE-constrained inversion with pretrained DDPM priors to achieve improved reconstruction accuracy and effective uncertainty quantification.
• Experiments on OpenFWI and benchmarks like Marmousi and SEG/EAGE Overthrust confirm DLO’s robustness under noisy, incomplete, and out-of-distribution conditions.

Decoupled Latent Optimization of Diffusion Models for Full Waveform Inversion

Introduction and Motivation

This paper introduces Decoupled Latent Optimization (DLO), a framework that leverages pretrained diffusion models as priors for Full Waveform Inversion (FWI), a high-dimensional, ill-posed PDE-constrained inverse problem ubiquitous in seismic imaging. Classical regularizers such as Tikhonov and total variation (TV) offer stabilization but lack the expressive power to recover realistic geological structures, while recent advances in diffusion-prior regularization for PDE inverse problems have shown promise yet suffer from inconsistent data/prior trade-offs and optimization fragility. DLO proposes a principled relaxation of latent optimization: it introduces a quadratic-penalty scheme over an auxiliary physical variable and a latent variable, allowing the explicit decoupling of data-fidelity optimization (operating strictly in physical space) from guidance by the diffusion-generated prior.

Background: FWI and Diffusion Model Regularization

Standard FWI seeks to minimize the discrepancy between observed and simulated seismic data by solving a highly nonconvex optimization problem, with gradients typically obtained via the adjoint-state method. Due to nonuniqueness and sensitivity to initialization, regularization is essential but classical approaches are insufficient for capturing subsurface complexities evident in geological structures. Diffusion models, notably DDPMs, score-based SDEs, and their discretizations (e.g., DDIM), have emerged as high-quality priors with strong generative and generalization capacity, especially in image-based inverse problems. However, prior approaches embedding diffusion priors into FWI either perform constraint injection directly into the noisy sampling trajectory (risking inconsistency between physics and prior) or add diffusion-based regularization terms with only indirect influence on solution fidelity. Approaches optimizing in the diffusion latent space exclusively (e.g., D-Flow, DMPlug) suffer from optimization instability and forgo the advantages of physical-space initialization, both of which are crucial for highly nonlinear, PDE-governed settings such as FWI.

Figure 1: Forward modeling and FWI inversion pipeline illustrating the wave equation with surface sources and simulated seismic surveys.

DLO Formulation

DLO recasts the latent optimization objective as a quadratic-penalty relaxation: rather than enforcing $v = R_\theta(z)$ as a hard constraint (with $R_\theta$ the deterministic DDIM sampler for the latent $z$), the objective is

$$\min_{v,\,z}\;\|\mathcal{F}(v) - \mathbf{d}_\text{obs}\|^2 + \lambda \|v - R_\theta(z)\|^2,$$

where $\mathcal{F}$ is the forward PDE operator and the penalty hyperparameter $\lambda > 0$ balances data fidelity and prior consistency. This decoupling admits alternating minimization: gradients for $v$ are computed via adjoint-state methods (not requiring backpropagation through the sampler chain), and the $z$-updates are performed via gradient descent through the DDIM chain, using automatic differentiation. This approach preserves the advantages of classical smoothed-model initialization and avoids the instability of hard prior projection or backpropagation through deep PDE-sampler compositions.

Figure 2: DLO algorithm overview with the physical and prior velocity fields jointly evolved via alternating minimization.

Convergence Analysis

Theoretical analysis confirms that as $\lambda \to \infty$, minimizers of the relaxed penalty formulation converge to minimizers of the original constrained problem, recovering strict prior consistency in the limit. Moreover, for finite (well-chosen) $\lambda$, DLO achieves a trade-off between numerical stability and prior adherence. The alternating block-coordinate minimization framework is robust to inexact updates, and the decoupling enables efficient application of the adjoint-state method for PDE gradients.

Experimental Results

OpenFWI Benchmarks

DLO is systematically benchmarked on OpenFWI, spanning four geologically distinct data families (FV-B, FF-B, CV-B, CF-B). For each, an independent DDPM prior is trained and reused (rather than finetuned). Comparisons include unregularized FWI, Tikhonov regularized, TV regularized FWI, and state-of-the-art diffusion-based baselines (DiffusionFWI and RED-DiffEq).

Figure 3: Ground truth (left) vs diffusion-generated prior samples (right) across four OpenFWI dataset families.

DLO achieves superior or competitive performance under clean, noisy, and highly incomplete (missing-trace) data acquisition scenarios. Numerical results indicate that DLO consistently outperforms classical regularizers and matches or exceeds advanced diffusion-based regularization methods according to RMSE, MAE, and SSIM. Under extreme noise or acquisition sparsity, DLO maintains structural fidelity and avoids spurious artifacts prevalent in baseline results.

Figure 4: Clean-data OpenFWI reconstruction comparisons showing faithful recovery of fine structure with DLO.

Figure 5: Quantitative comparison of normalized MAE, RMSE, and SSIM for all methods across OpenFWI families.

Robustness and Generalization

DLO demonstrates strong robustness to measurement noise and missing traces.

Figure 6: DLO and baseline reconstructions under increasing Gaussian noise, showing maintenance of high-fidelity with DLO.

Figure 7: Reconstructions with increasing receiver removal; DLO's prior preserves geologically plausible results even with heavily subsampled data.

Uncertainty Quantification

Owing to its stochastic latent initializations, DLO enables natural ensemble-based uncertainty quantification. Pixelwise variability across latent initializations correlates strongly with prediction error, localizing model uncertainty to geologically ambiguous regions (often corresponding to faults or poorly illuminated subsurface zones).

Figure 8: DLO ensemble sensitivity to latent initialization, highlighting uncertainty localization at faults.

Figure 9: Correlation between per-pixel ensemble standard deviation and absolute reconstruction error.

Large-Scale Out-of-Distribution Generalization

DLO is validated on Marmousi and SEG/EAGE Overthrust benchmarks—both out-of-distribution from the OpenFWI training set (notably for larger spatial domains and increased structural complexity). Without retraining, DLO delivers high-fidelity velocity inversions, capturing intricate faulting and stratigraphy, and maintains superiority over all tested baselines across varying initialization smoothness and increasing additive noise levels.

Figure 10: DLO reconstructions on Marmousi, showing high-resolution recovery of complex geological features.

Figure 11: Qualitative Overthrust results; DLO preserves stratigraphic and fault complexity unseen in training.

Figure 12: Robustness sweep for Marmousi: DLO remains superior across varying initialization smoothness and noise.

Figure 13: Overthrust robustness: DLO exhibits minimal performance degradation with increased measurement noise.

Implementation

DDPM priors are trained per-family with a U-Net (as in Hugging Face Diffusers) and a standard VP linear beta schedule. The physical forward operator $R_\theta$0 is a fourth-order finite-difference acoustic wave solver with a realistic acquisition geometry. DLO alternates Adam updates for $R_\theta$1 (velocity field) and $R_\theta$2 (latent) variables over 300 iterations. The computational overhead of DLO is modest: per-iteration runtime is $R_\theta$32.2x standard FWI iteration, with moderate GPU memory consumption. No tuning of priors or re-initialization is performed for "transfer" experiments on Marmousi or Overthrust, demonstrating prior robustness.

Implications and Future Outlook

DLO presents a structured methodology for integrating high-fidelity diffusion-model priors with PDE-constrained inverse problems, offering improved inversion accuracy, robust uncertainty quantification, and substantial resilience to noise, incomplete acquisition, and initialization. The approach naturally admits extension to other high-dimensional inverse problems governed by physical constraints when pretrained generative priors are accessible. Notably, DLO preserves the practical advantages of physical-space initialization, which is essential in exploratory geophysics and other engineering contexts.

Future extensions include scaling to three-dimensional geometries, integration with more expressive or hybrid priors, adaptive penalty annealing, and application to additional scientific imaging modalities. DLO’s decoupling strategy may inform the design of general algorithms for combining generative models with physical inversion beyond FWI settings.

Conclusion

DLO establishes a new standard for physics-guided, generative-model-regularized FWI, consistently outperforming both classical and contemporary diffusion-based methods across challenging scenarios and demonstrating strong generalization without retraining. Its decoupled penalty approach bridges the gap between strict prior-consistency and physically-informed data optimization, achieving high-fidelity subsurface reconstructions critical for seismic imaging. The methodology is broadly applicable in the context of PDE-inverse problems with learnable priors, with clear theoretical justification and demonstrated practical effectiveness.

Figure 14: Diffusion prior pretraining curve illustrating consistent noise-prediction loss convergence across training and validation.