Beyond First-Order Tweedie: Solving Inverse Problems using Latent Diffusion

Litu Rout, Yujia Chen, Abhishek Kumar, Constantine Caramanis, Sanjay Shakkottai, Wen-Sheng Chu

Abstract: Sampling from the posterior distribution in latent diffusion models for inverse problems is computationally challenging. Existing methods often rely on Tweedie's first-order moments that tend to induce biased results [32]. Second-order approximations are computationally prohibitive, making standard reverse diffusion processes in-tractable for posterior sampling. We present Second-order Tweedie sampler from Surrogate Loss (STSL), a novel sampler offering efficiency comparable to first-order Tweedie while enabling tractable reverse processes using second-order approximation. Theoretical results reveal that our approach establishes a lower bound through a surrogate loss and enables a tractable reverse process using the trace of the Hessian with only $\mathcal{O}(1)$ compute. We show STSL out-performs SoTA solvers PSLD [43] and P2L [10] by reducing neural function evaluations by 4X and 8X, respectively, while enhancing sampling quality on FFHQ, ImageNet, and COCO benchmarks. Moreover, STSL extends to text-guided image editing, effectively mitigating residual distortions in corrupted images. To our best knowledge, this is the first work to offer an efficient second-order approximation for solving inverse problems using latent diffusion, which further enables editing real-world images with corruptions.