Fast and Stable Diffusion Inverse Solver with History Gradient Update

TL;DR

This paper first proves that, in previous work, using the gradient descent method to optimize the data fidelity term is convergent, and introduces the incorporation of historical gradients into this optimization process, termed History Gradient Update (HGU).

Abstract

Diffusion models have recently been recognised as efficient inverse problem solvers due to their ability to produce high-quality reconstruction results without relying on pairwise data training. Existing diffusion-based solvers utilize Gradient Descent strategy to get a optimal sample solution. However, these solvers only calculate the current gradient and have not utilized any history information of sampling process, thus resulting in unstable optimization progresses and suboptimal solutions. To address this issue, we propose to utilize the history information of the diffusion-based inverse solvers. In this paper, we first prove that, in previous work, using the gradient descent method to optimize the data fidelity term is convergent. Building on this, we introduce the incorporation of historical gradients into this optimization process, termed History Gradient Update (HGU). We also provide theoretical evidence that HGU ensures the convergence of the entire algorithm. It's worth noting that HGU is applicable to both pixel-based and latent-based diffusion model solvers. Experimental results demonstrate that, compared to previous sampling algorithms, sampling algorithms with HGU achieves state-of-the-art results in medical image reconstruction, surpassing even supervised learning methods. Additionally, it achieves competitive results on natural images.

Figures (8)

• Figure 1: Our method is capable of reconstructing both nature images and medical data in a zero-shot manner with stable and fast optimization processes. In this paper, we demonstrate the reconstruction ability of our method on various datasets and measurements. FBP is filtered backprojection.
• Figure 2: Left: the proposed framework for inverse problem solving.. We extend the previous work in the pixel-based diffusion models ( chung2022diffusion) to the latent-based diffusion models, which are more efficient. We design a new approach to update the latent after every data consistency step, making the optimization process of latent more stable and the results more accuracy. Right: an illustration of the latent updating process and the proposed History Gradient Updating. The history gradients are essential to stabilize optimization processing and improve the quality of samples. We collect the history information from previous step, and estimate the optimal gradient factors $\{{\boldsymbol m}_t, {\boldsymbol v}_t\}$ based on the current gradient ${\boldsymbol g}_t$. Finally, we compute the optimal gradients based on the estimated momentum through our Momentum-variant HGU (GDM) and Improved-Momentum-variant HGU (iGDM) algorithms and obtain the next latent ${\boldsymbol z}_{t-1}$.
• Figure 3: Qualitative results of medical sparse data reconstruction. (a) represents the visual results of 18 sparse-view CT reconstruction. (b) represents the visual results of 32 sparse-view CT reconstruction. (c) represents the visual results of 45 degree limited-angle CT reconstruction. (d) represents the visual results of 90 degree limited-angle CT reconstruction. The resolution of CT images is $256\times256$ The display window of CT images is set to $\left[-150, 256\right]$ HU. The standard measurement of CT is 512 views around 180 degrees.
• Figure 4: Qualitative results of $99\%$ inpaint nature image reconstruction with Gaussian noise ($\sigma=0.05$).
• Figure 5: Qualitative results of $8\times$ super-resolution nature image reconstruction with Gaussian noise ($\sigma=0.05$).

Theorems & Definitions (12)

• Proposition 1
• Proposition 2
• Theorem 1: Latent Diffusion Solver
• Definition 1: Data fidelity term
• Definition 2: Convergence Criteria
• Theorem 2
• Theorem 3
• Remark 1
• Theorem 2
• proof