DiLO: Decoupling Generative Priors and Neural Operators via Diffusion Latent Optimization for Inverse Problems

Abstract

Diffusion models have emerged as powerful generative priors for solving PDE-constrained inverse problems. Compared with end-to-end approaches that learn inverse mappings from large paired observation--parameter datasets, decoupling the parameter prior from the forward physical model reduces reliance on paired data and improves flexibility. Fast neural operators can accelerate inference in such decoupled frameworks, but standard diffusion guidance evaluates them on partially denoised, non-physical intermediate states, introducing out-of-distribution (OOD) inputs. We therefore formulate the Manifold Consistency Requirement, under which the physical surrogate is evaluated only on fully denoised parameters. Based on this principle, we present Diffusion Latent Optimization (DiLO), which reformulates inference as deterministic optimization of the initial latent state and backpropagates measurement gradients through the full denoising trajectory. This avoids surrogate evaluations on partially denoised intermediate states while enabling accurate reconstruction. We establish asymptotic stationarity under bounded surrogate-gradient error. Experiments on Electrical Impedance Tomography, Inverse Scattering, and Inverse Navier-Stokes demonstrate DiLO's accuracy and robustness to noise. DiLO's code is available at \url{https://github.com/haibo-research/DiLO}.