Efficient Diffusion Model for Image Restoration by Residual Shifting

Most of the conventional IR methods can be cast into the Maximum a Posteriori (MAP) framework, a Bayesian paradigm encompassing a likelihood (loss) term and a prior (regularization) term. The likelihood reflects the underlying noise distribution of the LQ image. The commonly used $L_{2}$ or $L_{1}$ loss indeed corresponds to a Gaussian or Laplacian assumption on image noise. To more accurately depict the noise distribution, some robust alternatives were introduced, such as Poissonian-Gaussian [42], MoG [43], MoEP [44], Dirichlet MoG [45, 46] and so on. Simultaneously, there has been an increased focus on employing image priors to address the inherent ill-posedness of IR over recent decades. Typical image priors encompass total variation [47], wavelet coring [48], non-local similarity [49, 3], sparsity [50, 51], low-rankness [1, 52], dark channel [53, 54], among others. These conventional methods are mainly limited by the model capacity and the subjectivity inherited from the manually designed assumptions on image noise and prior.

In recent years, the landscape of IR has been dominated by deep learning (DL)-based methodologies. The seminal works [4, 55, 2] proposed to solve the IR problem using a convolution neural network, outperforming traditional model-based methods on the tasks of image denoising, super-resolution, and deblurring. Then, many studies [56, 57, 58, 59, 60, 61, 62, 63, 64, 6] have emerged, mainly concentrating on designing more delicate network architectures to further improve the restoration performance. Besides, there have been some discernible investigations that seek to combine current DL tools and classical IR modeling ideas. Notable works include the plug-and play or unfolding paradigm [65, 66, 67], learnable image priors [68, 69, 70, 71], and the loss-oriented methods [72, 73, 74]. The infusion of deep neural networks, owing to their large model capacity, has substantively extended the frontiers of IR tasks.

Inspired by principles from non-equilibrium statistical physics, Sohl-Dickstein et al. [7] proposed the diffusion model to fit complex distributions. Subsequent advancements by Ho et al. [8] and Song et al. [20] further improve its theoretical foundation by integrating denoising score matching and stochastic differential equation, thereby achieving impressive success in image generation [25]. Owing to its powerful generative capability, diffusion models have also found successful applications in the field of IR. Next, we provide a comprehensive overview of recent developments in diffusion-based IR methods.

The most straightforward solution to solve the IR problem using diffusion models is to introduce the LQ image as an addition condition in each timestep. Pioneering this approach, Saharia et al. [23] have successfully trained a diffusion model for image super-resolution. Subsequent studies [21, 14, 27] further expanded upon this approach, exploring its applicability in image deblurring, colorization, and denoising. To circumvent the resource-intensive process of training from scratch, an alternative strategy involves harnessing a pre-trained diffusion model to facilitate IR tasks. Numerous investigations, such as [18, 24, 28, 17, 32, 33, 15], reformulated the reverse sampling procedure of a pre-trained diffusion model into an optimization problem by incorporating the degradation model, enabling solving the IR problem in a zero-shot manner. Most of these methods, however, cannot deal with the blind IR problem, as they rely on a pre-defined degradation model. In contrast, some other works [29, 75, 30] directly introduced a trainable module that takes the LQ image as input. This module modulates the feature maps of the pre-trained diffusion model, steering it toward the direction of generating a desirable HQ image. Such a paradigm eliminates the reliance on a degradation model in the test phase, rendering it capable of handling the blind IR tasks.

The methodologies mentioned above are grounded in the foundational diffusion model initially crafted for image generation, necessitating a large number of sampling steps. This inefficiency presents a constraint on their application in real scenarios. The primary goal of our investigation is to devise a new diffusion model customized for IR, which facilitates a swift transition between the LQ/HQ image pair, thereby enhancing efficiency during inference.

The iterative sampling paradigm of diffusion models has proven highly effective in generating intricate and vivid image details, inspiring us to design an iterative approach to address the IR problem as well. Our proposed method builds up a Markov chain to facilitate a transition from the HQ image to its LQ counterpart as shown in Fig. 2. In this way, the restoration of the desirable HQ image can be achieved by sampling along the reverse trajectory of this Markov chain that starts at the given LQ image. Next, we will detail how to construct such a Markov chain specifically for IR.

The design of the transition distribution presented in Eq. (1) is guided by two primary principles. The first principle concerns the standard deviation, i.e., $\kappa\sqrt{\alpha_{t}}$, which aims to facilitate a smooth transition between $\bm{x}_{t}$ and $\bm{x}_{t-1}$. This is achieved by bounding the expected distance between $\bm{x}_{t}$ and $\bm{x}_{t-1}$ with $\sqrt{\alpha_{t}}$, given that the image data falls within the range of $[0,1]$. Mathematically, this is expressed as:

where $\text{max}[\cdot]$ represents the pixel-wise maximizing operation. The hyper-parameter $\kappa$ is introduced to increase the flexibility of this design. The second principle pertains to the mean parameter, i.e., $\bm{x}_{0}+\alpha_{t}\bm{e}_{0}$. Combining with the definition of $\alpha_{t}$, namely $\alpha_{t}=\eta_{t}-\eta_{t-1}$, it induces the marginal distribution in Eq. (2). Furthermore, the marginal distributions of $\bm{x}_{1}$ and $\bm{x}_{T}$ converges to $\delta_{\bm{x}_{0}}(\bm{x})$¹¹1$\delta_{\bm{\mu}}(\bm{x})$ denotes the Dirac distribution centered at $\bm{\mu}$. and $\mathcal{N}(\bm{x};\bm{y}_{0},\kappa^{2}\bm{I})$, respectively, serving as approximations for the HQ image and the LQ image. By constructing the Markov chain in such a thoughtful way, it is possible to handle the IR task by inversely sampling from it given the LQ image $\bm{y}_{0}$.

By combining Eq. (1) and Eq. (2), the target distribution $q(\bm{x}_{t-1}|\bm{x}_{t},\bm{x}_{0},\bm{y}_{0})$ in Eq. (6) can be rendered tractable and expressed in an explicit form given below:

The detailed calculation of this derivation can be found in Appendix -A. Considering that the variance parameter is independent of $\bm{x}_{t}$ and $\bm{y}_{0}$, we thus set it to be the true variance, i.e.,

The mean parameter $\bm{\mu}_{\bm{\theta}}(\bm{x}_{t},\bm{y}_{0},t)$ is reparameterized as:

where $f_{\bm{\theta}}(\cdot)$ is a deep neural network with parameter $\bm{\theta}$, aiming to predict $\bm{x}_{0}$. We explored different parameterization forms on $\bm{\mu}_{\theta}$ and found that Eq. (9) exhibits superior stability and performance.

Based on Eq. (9), the objective function in Eq. (6) is then simplified as:

where $w_{t}=\frac{\alpha_{t}}{2\kappa^{2}\eta_{t}\eta_{t-1}}$. In practice, we empirically find that the omission of weight $w_{t}$ results in a notable performance improvement, aligning with the conclusion in [8].

Unfortunately, attempts at further acceleration, particularly with fewer than 5 sampling steps, tend to produce over-smooth results. This phenomenon is primarily attributed to that the $L_{2}$-based loss in Eq. (10) favors the prediction of an average over plausible solutions. To overcome this limitation, we introduce the perceptual regularization [77] to impose an additional constraint on the solution space, namely,

where $l_{p}(\cdot,\cdot)$ denotes the pre-trained LPIPS metric, $\lambda$ is a hyper-parameter controlling the relative importance of these two constraints. This simple solution effectively curtails the sampling trajectory to fewer steps, e.g., 4 steps in this study, as shown in Sec. IV-B, while concurrently maintaining superior performance.

The proposed method employs a hyper-parameter $\kappa$ and a shifting sequence $\{\eta_{t}\}_{t=1}^{T}$ to determine the noise schedule of the diffusion process. In particular, the hyper-parameter $\kappa$ regulates the overall noise intensity during the transition, and its influence on performance is empirically discussed in Sec. IV-B. The subsequent exposition mainly revolves around the construction of the shifting sequence $\{\eta_{t}\}_{t=1}^{T}$.

Equation (2) indicates that the stochastic perturbation in state $\bm{x}_{t}$ is proportional to $\sqrt{\eta_{t}}$, incorporating a scaling factor $\kappa$. This observation motivates us to focus on the design of $\sqrt{\eta_{t}}$ rather than $\eta_{t}$. Previous work by Song et al. [78] has suggested that maintaining a sufficiently small value for $\kappa\sqrt{\eta_{1}}$, such as 0.04 in LDM [25], is imperative to ensure $q(\bm{x}_{1}|\bm{x}_{0},\bm{y}_{0})\approx q(\bm{x}_{0})$. Further considering $\eta_{1}\to 0$, we set $\eta_{1}$ to be the minimum value between $(\nicefrac{{0.04}}{{\kappa}})^{2}$ and $0.001$. For the terminal step $T$, we set $\eta_{T}$ as 0.999, guaranteeing $\eta_{T}\to 1$. For the intermediate timesteps $t\in[2,T-1]$, we propose a non-uniform geometric schedule for $\sqrt{\eta_{t}}$ as follows:

where

It should be noted that the choice of $\beta_{t}$ and $b_{0}$ is grounded in the assumption of $\beta_{1}=0$, $\beta_{T}=T-1$, and $\sqrt{\eta_{T}}=\sqrt{\eta_{1}}\times b_{0}^{T-1}$. The hyper-parameter $p$ controls the growth rate of $\sqrt{\eta_{t}}$, as depicted in Fig. 3(h).

The proposed noise schedule exhibits a high degree of flexibility in three key aspects. First, in the case of small values of $\kappa$, the final state $\bm{x}_{T}$ converges towards a perturbation around the LR image, as illustrated in Fig. 3(c)-(d). Compared to the diffusion process ended at Gaussian noise, this design significantly shortens the length of the Markov chain, thereby improving the inference efficiency. Second, the hyper-parameter $p$ provides precise control over the shifting speed, enabling a fidelity-realism trade-off in the SR results as analyzed in Sec. IV-B. Third, by setting $\kappa=40$ and $p=0.8$, our method achieves a diffusion process that remarkably resembles to LDM [25]. This is clearly demonstrated by the visual results of the diffusion process presented in Fig. 3(e)-(f), and further supported by the comparisons on the relative noise strength as shown in Fig. 3(g).

@articleliang2024control, title=Control Color: Multimodal Diffusion-based Interactive Image Colorization, author=Liang, Zhexin and Li, Zhaochen and Zhou, Shangchen and Li, Chongyi and Loy, Chen Change, journal=arXiv preprint arXiv:2402.10855, year=2024

It is widely acknowledged that the self-attention layer [79], a pivotal component in recent diffusion architectures, plays a crucial role in capturing global information in image generation. In the field of IR, however, it causes a blurring issue in handling the test images with arbitrary resolutions, particularly when the test resolution largely diverges from the training resolution. One typical example is provided in Figure 4, considering two LQ images with different resolutions. The baseline model with multiple self-attention layers, which is trained on a resolution of $64\times 64$, performs well when the LQ image aligns with the training resolution but yields blurred results when confronted with a mismatched resolution of $128\times 128$.

To analyze the underlying reason, we visualize the attention maps extracted from the first attention layer in this baseline network, as shown in Fig. 4(c1) and (c2). Note that these two attention maps are both interpolated to the resolution of $256\times 256$ for ease of comparison. Evidently, the example with a larger resolution tends to generate a more uniformly distributed attention map, i.e., Fig. 4(c2), being consistent with the entropy values annotated on the left-upper corner²²2The principle of maximum entropy posits that it achieves the maximum entropy when the attention map conforms to a uniform distribution. Consequently, a uniformly distributed attention map often leads to an over-smooth outcome, introducing undesirable distortions in performance.

To address this issue, some recent studies [14, 12] have chosen to discard the self-attention layers, a strategy that typically results in a noticeable decline in performance. Inspired by Liang et al. [40], we propose a solution by substituting the self-attention layer with that in Swin Transformer [80]. This straightforward replacement not only alleviates the blurring problem but also maintains the promised performance, as shown in Fig. 4(d1) and (d2). This is because the Swin Transformer computes the attention map in a local window, thus being independent of the image resolution.