Zero-Shot Adaptation for Approximate Posterior Sampling of Diffusion Models in Inverse Problems

Area of Science:

• Computational imaging
• Machine learning for inverse problems
• Generative models

Background:

• Diffusion models are powerful for inverse problems but suffer from slow inference.
• Sophisticated noise schedules improve diffusion model performance but are challenging for inverse problems.

Purpose of the Study:

• To develop a method that accelerates diffusion model inference for inverse problems.
• To improve the robustness and reconstruction quality of diffusion models in inverse problems.

Main Methods:

• Proposed Zero-Shot Approximate Posterior Sampling (ZAPS), a physics-driven deep learning approach.
• ZAPS uses zero-shot training with a physics-guided loss to learn log-likelihood weights for irregular timesteps.
• Approximated the Hessian of the prior logarithm using a learnable diagonalization approach for efficiency.

Main Results:

• ZAPS significantly reduces inference time compared to baseline methods.
• Demonstrated improved reconstruction quality across various inverse problems like deblurring, inpainting, and super-resolution.
• Showcased robustness to irregular noise schedules inherent in advanced diffusion models.

Conclusions:

• ZAPS offers a computationally efficient and effective solution for accelerating diffusion models in inverse problem solving.
• The method enhances the practical applicability of diffusion models in imaging and other scientific domains.
• ZAPS provides a flexible framework adaptable to various posterior sampling diffusion models.