RED-DiffEq: Diffusion-Based PDE Inversion

• RED-DiffEq is a computational framework for PDE-governed inverse problems that couples a physics-based data misfit with a diffusion-denoising learned prior.
• It employs a RED-style regularizer that integrates a pretrained DDPM to predict noise and enforce structural realism in high-resolution seismic velocity models.
• The framework achieves a balance between data fidelity and regularization, outperforming conventional methods under noisy and incomplete measurements.

Searching arXiv for the primary RED-DiffEq paper and closely related context. RED-DiffEq is a computational framework for solving PDE-governed inverse problems by combining a physics-based data misfit with a learned prior implemented via a denoising diffusion model, using the Regularization by Denoising (RED) paradigm. It is developed and instantiated for full waveform inversion (FWI) in seismic imaging, where the objective is to reconstruct high-resolution subsurface velocity models from seismic measurement data. The framework sits at the intersection of physics-driven inversion and data-driven learning: it uses a pretrained DDPM as a denoiser or score model and embeds it as a RED-style regularizer inside an explicit PDE-based inversion procedure (Shan et al., 25 Sep 2025).

1. Problem class and position within inverse-problem regularization

RED-DiffEq targets inverse problems in which the forward map is defined by a PDE, the unknowns are spatial fields, and the data are indirect, noisy, and possibly incomplete. The canonical example in the paper is acoustic FWI, where one seeks a subsurface velocity field $\mathbf{x}(\mathbf{r})$ from seismic recordings $\mathbf{u}_{\text{data}}$, with the wavefield $\mathbf{u}(\mathbf{r},t)$ governed by

$$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$

Such problems are difficult because the dependence of data on $\mathbf{x}$ is nonlinear, the inverse map is ill-posed, and the model parameters are high-dimensional spatial fields (Shan et al., 25 Sep 2025).

Within the regularization landscape, RED-DiffEq is positioned against both classical and modern alternatives. The paper contrasts it with Tikhonov regularization, which smoothes but oversmooths geologic boundaries, and total variation regularization, which preserves sharp interfaces but causes staircase artifacts. It also distinguishes RED-DiffEq from plug-and-play methods, in which the prior is implicit, and from end-to-end neural inversion. The framework is closer to RED than to end-to-end neural inversion, but uses a richer prior derived from a pretrained DDPM.

A central point of positioning is that existing diffusion-based inverse methods mostly treat the forward operator as a linear image operator, or work in latent space, whereas RED-DiffEq explicitly couples a diffusion prior to a PDE solver. This coupling is the defining feature of the framework: the inversion remains physics-based, but the prior is learned from a dataset of plausible model parameters.

2. Core construction: RED with a diffusion denoiser

The prior is represented by a DDPM trained on clean velocity models. During inversion, for a current iterate $\mathbf{x}_k$, the method first corrupts the model according to a variance-preserving diffusion schedule,

$$\mathbf{x}_{k,t}=\sqrt{\gamma(t)}\,\mathbf{x}_k+\sqrt{1-\gamma(t)}\,\boldsymbol{\epsilon}, \qquad \boldsymbol{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I}),$$

then evaluates the DDPM to predict the added noise $$\hat{\boldsymbol{\epsilon}}_\theta(\mathbf{x}_{k,t},t)$$, and uses Tweedie’s formula to construct a denoising operator $\mathcal{D}_\theta$. Intuitively, if $\mathbf{x}_k$ lies in a high-density region of the learned prior, the DDPM’s predicted noise aligns with the sampled noise and the regularization penalty is small; if $\mathbf{u}_{\text{data}}$0 is implausible under the prior, the mismatch is large (Shan et al., 25 Sep 2025).

The RED-style regularizer is written as

$\mathbf{u}_{\text{data}}$1

Using the DDPM score approximation

$\mathbf{u}_{\text{data}}$2

the practical denoiser becomes

$\mathbf{u}_{\text{data}}$3

This leads to the expanded RED regularizer

$\mathbf{u}_{\text{data}}$4

In experiments, the time-dependent weight $\mathbf{u}_{\text{data}}$5 is dropped for stability, and the scaling is absorbed into $\mathbf{u}_{\text{data}}$6. The resulting empirical inversion objective is

$\mathbf{u}_{\text{data}}$7

The second term is described as an unbiased Monte Carlo estimator of the RED regularizer.

3. Mathematical formulation and optimization procedure

The forward operator is defined by numerically solving the acoustic wave equation and sampling the simulated wavefield at receiver locations and times,

$\mathbf{u}_{\text{data}}$8

The inverse problem is cast in the standard regularized form

$\mathbf{u}_{\text{data}}$9

For RED-DiffEq, the regularizer $\mathbf{u}(\mathbf{r},t)$0 is the diffusion-based RED term constructed from the frozen DDPM (Shan et al., 25 Sep 2025).

Training of the diffusion model is separate from inversion. Clean velocity models $\mathbf{u}(\mathbf{r},t)$1 are corrupted via

$\mathbf{u}(\mathbf{r},t)$2

and the DDPM is trained with the $\mathbf{u}(\mathbf{r},t)$3-prediction objective

$\mathbf{u}(\mathbf{r},t)$4

During inversion, the DDPM is frozen. At each iteration, the method samples $\mathbf{u}(\mathbf{r},t)$5 and $\mathbf{u}(\mathbf{r},t)$6, forms $\mathbf{u}(\mathbf{r},t)$7, evaluates $\mathbf{u}(\mathbf{r},t)$8, and inserts the scalar penalty into the loss. Crucially, gradients are stopped through the DDPM prediction, denoted $\mathbf{u}(\mathbf{r},t)$9 in the paper, so the optimization objective used in practice is

$$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$0

The update is

$$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$1

implemented with Adam and typically cosine learning-rate scheduling. The gradient of the misfit is computed via the adjoint-state method implicit in the PDE solver, while the gradient of the regularizer is straightforward because $$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$2 is treated as a constant.

The paper emphasizes that the framework relies on established RED and score-matching theory, but focuses more on empirical performance than on formal convergence proofs. It also notes that stopping gradients through $$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$3 is justified empirically, though it breaks exact variational interpretation.

4. Instantiation for full waveform inversion

The experimental setting is 2D acoustic FWI on synthetic benchmarks. The acquisition geometry uses five sources, labeled A–E in the examples, together with receiver lines. The data conditions include clean seismic waveforms, waveforms with Gaussian noise with standard deviation $$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$4–$$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$5, and waveforms with randomly missing traces with $$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$6–$$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$7 traces removed (Shan et al., 25 Sep 2025).

The FWI workflow follows a conventional physics-based structure augmented by the diffusion regularizer. Initialization uses a smoothed version of the ground truth generated by a Gaussian filter, with examples including $$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$8 on OpenFWI and $$\frac{1}{\mathbf{x}^2(\mathbf{r})}\frac{\partial^2 \mathbf{u}(\mathbf{r},t)}{\partial t^2}-\nabla^2 \mathbf{u}(\mathbf{r},t)=q(\mathbf{r},t).$$9–$\mathbf{x}$0 for Marmousi and Overthrust. For each iterate, the acoustic PDE is solved numerically to generate synthetic data, the physics gradient is computed through standard FWI gradient computation or adjoint-state machinery, the regularization gradient is added, and the model is updated by gradient descent.

The regularization gradient is especially simple in implementation because the DDPM output is frozen during each optimization step. Since $\mathbf{x}$1 is treated as constant, the derivative of

$\mathbf{x}$2

reduces to $\mathbf{x}$3 projected appropriately. The paper also mentions an optional post-processing refinement step using the pretrained diffusion model to further refine the recovered velocity maps after inversion.

Conceptually, the FWI instantiation illustrates the intended role of the prior. The learned DDPM captures plausible geologic structures, while the PDE misfit enforces consistency with the observed seismic waveforms. This suggests that RED-DiffEq is designed to balance structural realism and data fidelity without replacing the physical forward model.

5. Training data, empirical performance, and generalization

The diffusion model is trained on OpenFWI velocity families: CF-B, FV-B, FF-B, and CV-B. These are described as curved faults, flat velocity layers, flat faults, and curved velocity variations. The DDPM architecture is a U-Net with residual blocks and attention, using $\mathbf{x}$4 diffusion steps, a sigmoid noise schedule, base channels $\mathbf{x}$5 with multipliers $\mathbf{x}$6, and $\mathbf{x}$7 attention heads. Training runs for $\mathbf{x}$8 iterations with batch size $\mathbf{x}$9 and learning rate $\mathbf{x}_k$0 (Shan et al., 25 Sep 2025).

Evaluation is performed on $\mathbf{x}_k$1 unseen models per OpenFWI family, for a total of $\mathbf{x}_k$2 held-out test cases, and on Marmousi and Overthrust as out-of-distribution benchmarks. On OpenFWI, RED-DiffEq is reported to reconstruct fault planes and smooth gradients simultaneously, yielding the lowest RMSE and MAE and the highest SSIM, with narrow error distributions centered at better medians. On Marmousi and Overthrust, despite being trained only on simpler OpenFWI families, the method is reported to recover complex faults and sharp discontinuities, maintain smooth gradients where appropriate, and outperform the baselines both quantitatively and qualitatively.

The robustness study addresses both measurement noise and incomplete acquisition. Under Gaussian noise, RED-DiffEq is described as maintaining structural coherence even at high noise levels, whereas other methods either oversmooth or introduce artifacts. Under missing traces, the metrics are said to be relatively stable across methods, but RED-DiffEq better preserves deeper structures visually. Performance degrades mildly with more aggressive smoothing of the initial model, but remains superior in the reported experiments.

A plausible implication is that the learned prior is not merely memorizing the training families. The combination of a diffusion prior and a PDE misfit appears to enable transfer to more complex velocity models than those seen during DDPM training, although the paper presents this as empirical evidence rather than a formal generalization theorem.

6. Comparisons, limitations, and broader applicability

The reported baselines are standard FWI without regularization, FWI with Tikhonov regularization, FWI with TV regularization, and a diffusion-based prior method by Wang et al. (2023) that uses generative diffusion models as prior regularizers but not the RED-style formulation. The comparison metrics are RMSE, MAE, and SSIM (Shan et al., 25 Sep 2025).

On OpenFWI, standard FWI is described as having large errors and failing to recover key structures; Tikhonov produces smoother results but loses sharp boundaries and fine structural detail; TV preserves edges but exhibits staircase artifacts and distorted geology; and the Wang et al. method performs adequately on simple structures but struggles with fine faults and complex velocity patterns. RED-DiffEq is reported as achieving the best RMSE, MAE, and SSIM, with visually accurate faults and layers, stable convergence curves, and tight error distributions. On Marmousi and Overthrust, the paper reports that the standard, Tikhonov, and TV baselines show significant artifacts and difficulty with complex structures, especially under poor initial models, while the Wang et al. method shows limited generalization and can be worse than physics-only baselines. RED-DiffEq is reported as consistently superior, with vertical profiles that closely track the ground truth across depths.

The paper also delineates several limitations and sensitivities. The learned prior depends on synthetic OpenFWI velocity families, so realism depends on how representative these are of actual geology. The regularization weight $\mathbf{x}_k$3 is chosen empirically as $\mathbf{x}_k$4, and over- or under-regularization can degrade results. The noise schedule $\mathbf{x}_k$5 uses a sigmoid schedule and affects both DDPM behavior and the RED term, but is not extensively ablated. FWI remains sensitive to initialization, and very poor initial models still challenge convergence. Each iteration requires at least one full wave equation solve, forward plus adjoint, so the computational cost is similar to conventional FWI per iteration; each iteration also requires one DDPM forward pass, which the paper describes as modest compared to PDE solve cost. The reported experiments use $\mathbf{x}_k$6 iterations.

The framework is formulated abstractly enough to extend beyond seismic inversion. The paper states that the forward operator may be any PDE solver $\mathbf{x}_k$7, the unknown may be any field such as conductivity, permeability, or stiffness, and the prior may be any dataset of plausible fields on which a DDPM can be trained. The generic formulation is

$\mathbf{x}_k$8

The paper lists medical imaging, electromagnetic inversion, flow or diffusion-reaction systems, and climate or fluid dynamics as candidate application areas. This suggests that RED-DiffEq is intended as a general recipe for PDE inversion with a diffusion prior and RED-style regularization, rather than as a method restricted to FWI.