Tweedie Moment Projected Diffusions For Inverse Problems

TL;DR

This work introduces Tweedie Moment Projected Diffusions (TMPD) for solving linear inverse problems with diffusion priors by deriving a Gaussian posterior approximation via Tweedie's formula and moment projection. The approach yields a principled, parameter-free data term that augments the reverse diffusion, with theoretical guarantees in the Gaussian case and a bound for general data. Empirically, TMPD and its fast DTMPD variant demonstrate robust performance across VP and VE forward processes, noise levels, and inverse tasks like inpainting and super-resolution, while avoiding time-varying step-size hyperparameters. The results highlight a scalable, theoretically grounded alternative to existing conditional diffusion methods for inverse problems, with potential for further efficiency and generalization improvements.

Abstract

Diffusion generative models unlock new possibilities for inverse problems as they allow for the incorporation of strong empirical priors in scientific inference. Recently, diffusion models are repurposed for solving inverse problems using Gaussian approximations to conditional densities of the reverse process via Tweedie's formula to parameterise the mean, complemented with various heuristics. To address various challenges arising from these approximations, we leverage higher order information using Tweedie's formula and obtain a statistically principled approximation. We further provide a theoretical guarantee specifically for posterior sampling which can lead to a better theoretical understanding of diffusion-based conditional sampling. Finally, we illustrate the empirical effectiveness of our approach for general linear inverse problems on toy synthetic examples as well as image restoration. We show that our method (i) removes any time-dependent step-size hyperparameters required by earlier methods, (ii) brings stability and better sample quality across multiple noise levels, (iii) is the only method that works in a stable way with variance exploding (VE) forward processes as opposed to earlier works.

Figures (17)

• Figure 1: Error to target posterior for a Gaussian random field. (Top row) visualisation of the empirical mean and variance of the 1500 samples that were used to compute this error against the analytical moments. (Bottom) Wasserstein distances of different methods w.r.t. sample size. For details, see Appendix \ref{['appendix:gaussian']}.
• Figure 2: We display the first two dimensions of the GMM inverse problem for one of the measurement models tested (${\mathbf{H}}, \sigma_{{\textnormal{y}}}=0.1, (d_{{\textnormal{x}}}, d_{{\textnormal{y}}}) = (80, 1)$). The blue dots represent samples from the target posterior, while the red dots correspond to samples generated by each of the algorithms used (the names of the algorithms are given at the bottom of each column).
• Figure 3: 4$\times$ bicubic super-resolution FID vs LPIPS (left) and SSIM (right) using the VP-SDE on FFHQ-1k validation dataset for increasing observation noise.
• Figure 4: 'box' mask inpainting FID vs LPIPS (left) and SSIM (right) using the VP-SDE on FFHQ-1k validation dataset for increasing observation noise.
• Figure 5: 'random' mask inpainting FID vs LPIPS (left) and SSIM (right) using the VP-SDE on FFHQ-1k validation dataset for increasing observation noise.

Theorems & Definitions (5)

• Proposition 1: Tweedie's formula
• Proposition 2: Moment projection
• Proposition 3: Gaussian data distribution
• Theorem 1: General data distribution
• proof : Proof of Theorem \ref{['thm:bound']}