Accelerating Diffusion Models for Inverse Problems through Shortcut Sampling

1 Introduction

Inverse problem is a classic problem in the field of machine learning. Given a low-quality (LQ) input measurement image y and a degradation operator H, inverse problems aim to restore the original high-quality (HQ) image x from y = Hx + n. Many image restoration tasks, including super-resolution [Haris et al., 2018; Wang et al., 2018], colorization [Larsson et al., 2016], inpainting [Yeh et al., 2017], deblurring [Guo et al., 2019; Zhang et al., 2018a] and denoising [Wang et al., 2022], can be considered as applications of solving inverse problems. In general, the restored image should exhibit two critical attributes: Realism and Faithfulness. The former indicates that the restored image should be of high-quality and photo-realistic, while the latter denotes that the restored image should be consistent with the input image in the degenerate subspace.

Recently, diffusion models [Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021c] have demonstrated phenomenal performance in generation tasks [Rombach et al., 2022; Dhariwal and Nichol, 2021; Xiao et al., 2022]. Due to their powerful capability in modeling complex distributions, recent methods [Kawar et al., 2021; Kawar et al., 2022; Chung et al., 2022b; Chung et al., 2023; Wang et al., 2023; Lugmayr et al., 2022; Song et al., 2021b] have sought to utilize pre-trained diffusion models for solving inverse problems in an unsupervised manner. These methods leverage the generative priors to enhance realism, and enforce additional consistency constraints to ensure faithfulness. Specifically, diffusion models have predefined a forward process expressed as p(xt|xt-1) and a generation process expressed as p(xt-1|xt). Existing methods are primarily concerned with the generation process, altering the posterior sampling to a conditional sampling p(xt-1|xt, y) based on the LQ image y.

Despite the successful application of existing methods in solving various inverse problems, their relatively slow sampling speed is a major drawback. These methods overlook the importance of modifying the forward process to achieve an improved initial state, simply initiating from random noise xT ∼ N(0, I). However, as depicted in Fig. 1, since the initial state of pure noise xT is considerably distant from the target HQ images x0, prior methods have to travel through a long journey of sampling, typically requiring at least 100-250 neural function evaluations (NFEs), to generate the overall layout, structure, appearance and detailed texture of the restored image, and finally achieve a satisfactory result.

In this work, we propose Shortcut Sampling for Diffusion (SSD), a novel approach for solving inverse problems in a zero-shot manner. The primary concept behind SSD is to find a specific transitional state that bridges the gap between the input image y and the target restoration x. For convenience, we use "E " to refer to this specific transitional state. As depicted in Fig. 1, by employing a shortcut path of "Input - E - Target"(H† y → xt → x0) instead of the previous "Noise-Target"(xT → x0), SSD enables precise and fast restoration.

For the forward process denoted as p(xt|xt-1, y), we observe that the original forward process erodes information from the input image, yielding realistic yet unfaithful results. Conversely, an alternative, DDIM Inversion[Song et al., 2021a], widely adopted for image editing[Hertz et al., 2022; Tumanyan et al., 2023], tends to produce unrealistic outcomes. To address this dilemma, we introduce Distortion Adaptive Inversion (DA Inversion). By incorporating a controllable random disturbance at each forward step, DA Inversion is capable of deriving E that adheres to the predetermined noise distribution while preserving the majority of the input image's information.

For the generation process denoted as p(xt-1|xt, y), we utilize generation priors to produce extra details and texture, and introduce the back projection technique[Tirer and Giryes, 2018; Wang et al., 2023] as additional consistency constraints. In particular, we add a projection step after each denoising step to project the coarse restored image onto the degenerate subspace, obtaining a revised version of the restored image that aligns with the input image in the degenerate subspace. We further propose SSD+ to extend applicability to scenarios with unknown noise and more intricate degradation conditions.

To validate the effectiveness of SSD, we conduct experiments across various inverse problems, including superresolution, colorization, inpainting, and deblurring on CelebA and ImageNet. Experiments demonstrate that SSD achieves competitive results in comparison to state-of-the-art zero-shot methods (with 100 NFEs), even though it employs only 30 NFEs. Additionally, we observe that SSD, when operated with 100 NFEs, can surpass state-of-the-art methods in certain IR tasks.

2 RelatedWorks

2.1 DiffusionModels

Denoising Diffusion Probabilistic Models. Diffusion models [Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021c] are a family of generative models designed to model complex probability distributions of high-dimensional data. Denoising Diffusion Probabilistic Models(DDPM) [Ho et al., 2020] comprise both a forward process and a generation process. In the forward process, an image x0 is transformed into Gaussian noise xT ∼ N(0, 1) by gradually adding random noise over T steps. We can describe each step in the forward process as:

where [xt]Tt=0 is the noisy image at time-step t, [βt]Tt=0 is the predefined variance schedule, and [zt]Tt=0 is the random gauss noise added at time-step t. Using reparameterization tricks [Kingma and Welling, 2013], The resulting noisy image xt can be expressed as:

where αt = ∏ti=1 (1-βi) and ϵ is the reparameterized noise. The generation process transforms gaussian noise xT to image x0, the transition from xt to xt-1 can be expressed as:

and ϵθ(xt, t) is a neural network trained to predict the noise ϵ from noisy image xt at time-step t. The noise approximation model ϵθ(xt, t) can be trained by minimize the following objective:

Denoising Diffusion Implicit Models. Meanwhile, [Song et al., 2021a] generalize DDPM via a non-Markov diffusion process that shares the same training objective, whose generation process is outlined as follows:

where fθ(xt, t) is the prediction of x0 at time-step t:

When σt = 0, DDIM samples images through a deterministic generation process, which allows for high-quality sampling with fewer steps:

2.2 Solving Inverse Problems in a Zero-shot Way

A general inverse problem aims to restore a high-quality image x from a known degradation operator H and the degraded measurement y with random additional noise n:

Some methods investigate leveraging the generative priors of pre-trained generative models to restore degraded images in a zero-shot way. GAN Inversion aims to find the closest latent vector in the GAN space for an input image. Utilizing the GAN Inversion technique, existing methods[Ulyanov et al., 2018; Pan et al., 2021; Menon et al., 2020] have exhibited remarkable effectiveness in tackling inverse problems.

Compared to GAN, diffusion models provide a forward process that enables the direct acquisition of latent vectors within the Gaussian noise space. By performing generation processes and using consistency constraints at each step, diffusion models can be applied to various IR problems. DDRM[Kawar et al., 2022] applies SVD to decompose the degradation operators and perform diffusion in its spectral space for various IR tasks. DDNM[Wang et al., 2023] uses range-null space decomposition to decompose the restored image as a null-space part and a range-space part, they keep the range-space part unchanged to force consistency, and obtain the null-space part through iterative refinement. Meanwhile, inspired by guided diffusion[Dhariwal and Nichol, 2021; Liu et al., 2023], some recently developed methods[Chung et al., 2022b; Chung et al., 2023; Song et al., 2022; Mardani et al., 2023; Pokle et al., 2023] adopted gradient-based guidance to generate faithful restoration. MCG[Chung et al., 2022b] applies a gradient-based measurement consistency step at each denoising step to achieve image restoration. These methods typically perform better in terms of perceptual quality, while requiring additional inference consumes compared to non-gradient methods.

3 Method

3.1 Shortcut Sampling

As discussed above, diffusion models comprise two processes: a forward process denoted as p(xt|xt-1), which progressively adds noise to the image until complete Gaussian noise; and a generation process denoted as p(xt-1|xt), which generates realistic images through iteratively denoising. Previous methods mainly focus on modifying the posterior sampling process to p(xt-1|xt, y) during the generation process, while ignoring the utilization of the forward process. Instead, these methods typically sample xT directly from the Gaussian prior as the initial state.

In this work, we propose Shortcut Sampling for Diffusion (SSD), a novel pipeline for solving inverse problems in a zero-shot manner. Different from previous methods that initiate from pure noise, SSD enhances the forward process to obtain an intermediate state E , which serves as a bridge between the measurement image y and the restored image x. Throughout the shortcut sample path of "input-E -output", SSD can achieve efficient and satisfactory restoration results.

For convenience, we denote the transition from the measurement image y to E as the "inversion process"; and the transition from E to the restored image x0 as the "generation process". Given a LQ image y and corresponding degraded operator H, we start from H† y and apply Distortion Adaptive Inversion (DA Inversion) to derive E in the inversion process(Sec. 3.2). Subsequently, during the generation process, we iteratively perform the denoising step and the back projection step to generate both faithful and realistic results(Sec. 3.3). Further, due to SSD relies on an accurate estimation of degraded operators to exhibit high performance, we proposed an enhanced version called SSD+ that makes SSD suitable for noisy situations or inaccurate estimation of H.(Sec. 3.4). The overall pipeline is illustrated in Fig. 2.

3.2 Distortion Adaptive Inversion.

We expect to obtain the transitional state by enhancing the forward process. As previously discussed, the transitional state E should satisfy the following criterias:

Criteria (i): The transitional state should contain information from the input image;

Criteria (ii): The transitional state should retain the capacity for generating a high-quality image.

WhyDDIMInversionCannotWorkWell. To satisfy Criteria (i), a naive approach is to apply DDIM Inversion, capable of providing a deterministic mapping from the input image y to a noisy transitional state[Hertz et al., 2022; Tumanyan et al., 2023; Zhang et al., 2024]. Given the deterministic nature of the DDIM generation process, we can establish the DDIM Inversion process by reversing Eq. (7) in the following manner:

The transitional state obtained through DDIM Inversion preserves most information of the input images since we can reconstruct it by iteratively executing Eq. (7). However, as depicted in Fig. 3 (a), the application of E in the generation process produces faithful but unrealistic results, thus violating Criteria (ii).

We attribute the failure to the observation that, the obtained E deviates from the predefined noise distribution. More specifically, given a low-quality input image y, the predicted noise {ϵθ(xt, t)} during DDIM Inversion process deviates from the standard normal distribution, resulting in the deviation of E from the predefined noise distribution. During the generation process, the pre-trained model receives outof-domain input distributions, thereby generating unrealistic results. We summarize this observation as follows, more details are available in Appendix B.

Assumption 1. Diffusion Models rely on in-domain noise distribution to generate high-quality images. When facing low-quality input images, the distribution of predicted noise ϵθ(xt, t) during the DDIM Inversion process exhibits a greater deviation from the standard normal distribution.

Why Original Forward Process Cannot WorkWell. Another extreme scenario is the original forward process, which can be regarded as a special inversion technique termed DDPM Inversion. In DDPM Inversion, the predicted noise is replaced with randomly sampled noise from Gaussian distribution. We can redefine the forward process in Eq. 1 in a similar form of DDIM Inversion in Eq. 9:

As shown in Fig. 3(b), DDPM Inversion converts the input image y into pure noise, thereby violating Criterion (i) and producing results that are realistic yet lack faithfulness.

Distortion Adaptive Inversion. Since a deterministic inversion process like DDIM Inversion produces unrealistic results, while a stochastic inversion process like DDPM Inversion yields unfaithful results. To resolve this dilemma, we propose a novel inversion approach called Distortion Adaptive Inversion(DA Inversion). The definition of DA Inversion is stated as follows:

Definition 1 (Distortion Adaptive Inversion). We define the iterative process of DA Inversion as:

where η control the proportion of random disturbances and 0 < η < 1.

For ease of exposition, the predicted noise in DA Inversion at each timestep can be rephrased as

By adding controllable random perturbations in each inverse step, DA Inversion has the capability to generate highquality images while preserving the essential information including layout and structure, which is shown in Fig. 3 (c).

For Criterion (i), since the random perturbation only replaces a portion of the predicted noise, the E obtained through DA Inversion actually preserves a substantial amount of information from the input image. For Criterion (ii), we have verified that incorporating random disturbances can bring the predicted noise closer to N(0, 1). Proofs are available in Appendix A.

Theorem 1. Assuming ϵθ(xt, t) ∼N(µ, σ2 ), We have:

thus:

which indicates that after adding random disturbance, ϵDA becomes closer to N(0, 1).

In practice, inspired by [Meng et al., 2021; Kim et al., 2022; Chung et al., 2022a], rather than performing the inversion process until the last time-step T , we find that we can achieve acceleration by performing until time-step t0 < T .

3.3 Back Projection

Although E obtained from DA Inversion carries information about the input image, and the generation process started from which can produce images with high quality, the result may not entirely align with the input LQ image in the degenerate subspace. Following [Tirer and Giryes, 2018; Kawar et al., 2022; Wang et al., 2023], we introduce the back projection technique as consistency constraints to address this problem.

For details, given a measurement image y and corresponding degraded operator H, we can project the coarse restored image x onto the affine subspace {HRn = y} by the following operation to force consistency:

where H† is the pseudo-inverse of H, details about calculating H and H† is available in Appendix E.

The refined restored image x′ consists of two parts: the former part, denoted as (I - H†H)x, represents the residual between x and the image obtained after projection-in and projection-back, which can be interpreted as the enhancement details of x. The latter part, denoted as H† y, can be regarded as the preservation of input image y. After the backprojection step, we have:

which indicates that x′ entirely aligns with the input measurement image y in the degenerate subspace.

At each timestep, we initiate the process by performing a denoising step to obtain the predicted x0 in Eq. 6. Subsequently, we process with a back projection step to refine the results of x0 and derive xt-1 through Eq. 5. The complete transformations from xt to xt-1 can be expressed as follows:

3.4 Expand SSD to Noisy IR Tasks

Although SSD is effective in addressing various noiseless inverse problems, it tends to exhibit poor performance when faced with noisy tasks. This limitation can be primarily attributed to the utilization of back projection. The success of back projection hinges on a precise estimation of the degraded operator H. When applied to blind image restoration or noisy IR tasks, back-projection tends to result in disappointing restorations because of the inability to satisfy Eq. 18.

To solve this problem, we proposed SSD+, an enhanced version that makes SSD suitable for noisy situations or inaccurate estimation of H. Earlier studies[Meng et al., 2021; Hertz et al., 2022; Tumanyan et al., 2023] have indicated that diffusion models typically generate the overall layout and color in the early stage, generate the structure and appearance in the middle stage, and generate the texture details in the final stage. We notice that SSD employs shortcut sampling to skip the early stage and ensure the preservation of the overall layout. However, during the middle final stage, the utilization of back projection with an inaccurate H has the potential to deteriorate the fine-textured details, leading to suboptimal outcomes. In SSD+, rather than performing back projection throughout the generation process, we restrict its use to the middle stage, where it still plays a crucial role in maintaining structure consistency. During the final stage of generation, we rely exclusively on diffusion priors to ensure texture details without compromising the integrity of the original structure.

4 Experiments

4.1 Experimental Setup

Pretrained Models and Datasets. To evaluate the performance of SSD, we conduct experiments on two datasets with different distribution characters: CelebA 256×256 [Karras et al., 2017] for face images and ImageNet 256×256 [Deng et al., 2009] for natural images, both containing 1k validation images independent of the training dataset. For CelebA 256×256, we use the denoising network VP-SDE[Song et al., 2021c; Meng et al., 2021] 1 . For ImageNet 256×256, we use the denoising network guided-diffusion [Dhariwal and Nichol, 2021] 2 .

Degradation Operators. We conduct experiments on several typical IR tasks, including Super-Resolution(× 4, × 8), Colorization, Inpainting, and Deblurring. Details of degradation operators can be found in Appendix E.

Evaluation. We use PSNR, FID[Heusel et al., 2017], and LPIPS[Zhang et al., 2018b] as the main metrics to quantitatively evaluate the performance of image restoration. Especially due to the inability of PSNR to capture colorization performance, we use FID and LPIPS for colorization. Additionally, we adopt Neural Function Evaluations(NFEs) as the metrics of sampling speed, which is a commonly employed benchmark in diffusion model-based methods. For most methods, NFEs are equivalent to the sampling steps of the generation process. Since SSD introduces an additional inversion process, the NFEs of SSD are calculated by summing the steps involved in both the inversion process and the generation process. All of our experiments are conducted on a single NVIDIA RTX 2080Ti GPU.

ComparisonMethods. We compare the restoration performance of the proposed method with recent State-Of-The-Art zero-shot image restoration methods using pre-trained diffusion models: DDRM[Kawar et al., 2022], DPS[Chung et al., 2023] and DDNM[Wang et al., 2023]. For a fair comparison, all methods above use the same pre-trained denoising networks and degradation operator.

1 https://github.com/ermongroup/SDEdit

2 https://github.com/openai/guided-diffusion

4.2 Noiseless Image Restoration Results

We compare SSD (with 30 and 100 steps) with previous methods mentioned in Sec 4.1. The quantitative evaluation results shown in Table 1 illustrate that the proposed method achieves competitive results compared to state-of-the-art methods. When setting NFEs to 30, SSD-30 outperforms other methods known for fast sampling such as DDRM30 and achieves better perception-oriented metris (i.e., FID, LPIPS) than SOTA methods (DDNM with 100 NFEs). When setting NFEs to 100, SSD-100 achieves SOTA performance in many IR tasks, including SR ×4 and colorization. As shown in Fig. 4, 5, SSD generates high-quality restoration results in all tested datasets and tasks.

4.3 Noisy Image Restoration Results

To illustrate the robustness of SSD+ in the face of noisy situations and complex degradation, we evaluate SSD and SSD+ on diverse inverse problems with gaussian noise and JPEG compression[Shin and Song, 2017]. For gaussian noise, we add gaussian noise z ∼ N(0, σ2 ) to the degraded image y, where σ represents the intensity of noise and is randomly distributed in [0.0, 0.2]. For JPEG compression, we perform JPEG compression[Wang et al., 2021] with a quality factor of 60 after the degraded operator H is applied. The quantitative result is shown in Tab. 2. Qualitative results are available in Fig. 6. Compared to SSD, SSD+ exhibits enhanced robustness when encountering more intricate degradation, producing satisfactory restoration results.

4.4 Sampling Speed

We further investigate the performance of various methods concerning FID in relation to the change in NFEs, as illustrated in Fig. 7. We conduct experiments with SR ×4 task on the CelebA dataset. DPS [Chung et al., 2023] is excluded from the comparison as it fails to generate reasonable results at low NFEs(< 100). The results indicate that SSD surpasses all other methods at both high and low NFEs, with a greater advantage when the NFEs is relatively small.

5 Conclusion

In this paper, we propose SSD, a novel framework for solving inverse problems in a zero-shot manner. We have departed from the conventional "Noise-Target" paradigm and instead proposed a shortcut sampling pathway of "Input-EmbryoTarget". This novel approach enables us to achieve satisfactory results with reduced steps. We further propose SSD+, an enhanced version of SSD tailored to excel in scenarios where degradation estimation is less accurate or in the presence of noise. We hope the proposed pipeline may inspire future work on inverse problems to solve them in a more efficient manner.

Acknowledgements

This work was partly supported by the National Natural Science Foundation of China (Grant No. 61991451)

and the Shenzhen Science and Technology Program (JSGG20220831093004008).