SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zheng Chen $^{1}$ , Yulun Zhang $^{2*}$ , Ding Liu $^{3}$ , Bin Xia $^{4}$ , Jinjin Gu $^{5,6}$ , Linghe Kong $^{1*}$ , Xin Yuan $^{7}$
$^{1}$ Shanghai Jiao Tong University,  $^{2}$ ETH Zürich,  $^{3}$ Bytedance Inc,  $^{4}$ Tsinghua University,
$^{5}$ Shanghai AI Laboratory,  $^{6}$ The University of Sydney,  $^{7}$ Westlake University
Diffusion models (DMs) have recently been introduced in image deblurring and exhibited promising performance, particularly in terms of details reconstruction. However, the diffusion model requires a large number of inference iterations to recover the clean image from pure Gaussian noise, which consumes massive computational resources. Moreover, the distribution synthesized by the diffusion model is often misaligned with the target results, leading to restrictions in distortion-based metrics. To address the above issues, we propose the Hierarchical Integration Diffusion Model (HI-Diff), for realistic image deblurring. Specifically, we perform the DM in a highly compacted latent space to generate the prior feature for the deblurring process. The deblurring process is implemented by a regression-based method to obtain better distortion accuracy. Meanwhile, the highly compact latent space ensures the efficiency of the DM. Furthermore, we design the hierarchical integration module to fuse the prior into the regression-based model from multiple scales, enabling better generalization in complex blurry scenarios. Comprehensive experiments on synthetic and real-world blur datasets demonstrate that our HI-Diff outperforms state-of-the-art methods. Code and trained models are available at https://github.com/zhengchen1999/HI-Diff.
Image deblurring is a long-standing computer vision task, aiming to recover a sharp image from a blurred observation. Blurred artifacts can arise from various factors, e.g., camera shake and fast-moving objects. To address this challenge, traditional methods typically formulate the task as an optimization problem and incorporate natural priors [12, 38, 24, 29] to regularize the solution space. However, since blur in real scenarios is complex and non-uniform, which is hard to be modeled by the specific priors, these algorithms suffer from poor generalization in complex situations.
With the development of deep learning, convolutional neural networks (CNNs) have been applied in image deblurring [37, 43, 13, 55, 1, 11]. Among them, the regression-based methods have shown remarkable success, especially in terms of distortion-based metrics (e.g., PSNR). Moreover, Transformer-based approaches [48, 53, 44, 21], which can capture long-distance dependencies, are introduced as an alternative to CNNs. These methods further boost the deblurring performance. However, regression-based methods are prone to recovering images with fewer details, since the regression losses are conservative with high-frequency details [36].
Apart from regression-based methods, deep generative models, like generative adversarial networks (GANs) [14] and normalizing flows [8], provide other solutions for generating complex details. Recently, Diffusion Models (DMs) [40, 15] have exhibited impressive performance in image synthesis [41, 35] and restoration tasks (including image deblurring) [50, 33, 36, 51]. DMs generate high-fidelity images through a stochastic iterative denoising process from a pure white Gaussian noise. Compared to other generative models, such as GANs, DMs generate a more accurate target distribution without encountering optimization instability or mode collapse.
*Corresponding authors: Yulun Zhang, yulun100@gmail.com; Linghe Kong, linghe.kong@sjtu.edu.cn
Nonetheless, DM-based methods face several challenges. First, DMs are limited by the high computational cost of generating samples. It requires a high number of inference steps. Some methods reduce the number of iterations by predicting the residual distribution [50] or applying advanced sampling strategies [41, 3]. However, the overall computational complexity is still high, especially for high-resolution images. Furthermore, although some methods perform DMs on latent space [35], the compression ratio is small (e.g., 8 times), due to limitations in generation quality. Second, generative models, including DMs, tend to produce undesired artifacts not present in the original clean image. Besides, the generated details may also be misaligned with real targets. These lead to the poor performance of DM regarding some distortion-based metrics (e.g., PSNR).
The aforementioned issues promote us to apply DMs from another perspective. The motivation for our work is threefold. First, due to the high computational overhead in the image space, we also consider applying DMs on the low-dimensional latent space. Meanwhile, we increase the compression ratio of latent to effectively reduce the computational complexity. Second, since the advantages of regression-based methods in distortion accuracy, we integrate DMs and regression-based methods to improve the performance of DMs regarding distortion. Third, considering the non-uniform blur in real scenarios, we apply a hierarchical approach to enhance the generalization of the method.
Specifically, we apply DMs to generate priors in latent space, inspired by the application of priors in traditional deblurring algorithms. The priors are integrated into the regression-based model in a hierarchical manner to improve the details of restored images. The design contains three benefits: (1) The dimension of latent space can be very low, since the regression-based model restores most of the distribution. (2) The issue of distortion caused by misaligned details generated by DMs can be avoided. (3) The hierarchical integration enables better generalization in complex blurry scenarios.
Based on the above analysis, we propose a novel approach called the Hierarchical Integration Diffusion Model (HI-Diff) for realistic image deblurring. Following previous practice [10, 35, 51], we perform a two-stage training to realize latent compression and the training of the DM, respectively. Without loss of generality, for the regression-based methods, we choose the encoder-decoder Transformer architecture. In the first stage, we compress the ground-truth image into a highly compact latent representation as the prior feature through a latent encoder (LE). We propose the hierarchical integration module (HIM) that can fuse the prior and intermediate features of Transformer at multiple levels. We jointly train the LE and Transformer to effectively construct the prior. In the second stage, we train a latent diffusion model to generate the prior feature in the latent space from the Gaussian noise and guide Transformer through the HIM. Similar to the first stage, we jointly train the DM and Transformer. In summary, our main contributions are three-fold as follows:
- We propose the Hierarchical Integration Diffusion Model (HI-Diff), for realistic image deblurring. The HI-Diff leverages the power of diffusion models to generate informative priors, that are integrated into the deblurring process hierarchically for better results.
- We apply the diffusion model in highly compact latent space to generate the prior. Meanwhile, we propose the hierarchical integration module to fuse the prior into the regression-based model from multiple scales, enabling the generalization in complex blurry scenarios.
- Extensive experiments conducted on synthetic and real-world blur datasets demonstrate the superior performance of the HI-Diff in comparison to state-of-the-art deblurring methods.
2.1 Image Deblurring
Traditional Methods. Traditional deblurring methods generally formulate the problem as an optimization problem [12, 38, 24, 6, 29]. They utilize various natural image priors for sharp images and blur kernels. Common methods include local smoothness prior [38], sparse image prior [24],  $L_0$ -norm gradient prior [52], and dark channel prior [29]. However, these methods rely on manually crafted priors, which results in poor generalization ability and limited performance in complex situations.
Deep CNN-based Methods. With the rapid development of deep learning, significant progress has been made in image deblurring using CNN-based methods [37, 28, 43, 55, 1, 16]. For instance, MSCNN [28] designs a multi-scale CNN to restore sharp images. SRN [43] proposes a coarse-to-fine scale-recurrent network for more efficient multi-scale image deblurring. DMPHN [55] devises a deep multi-patch hierarchical deblurring network to exploit the deblurring information at different scales. XYDeblur [16], on the other hand, divides the original deblurring problem into two sub-problems.
(a) Transformer
(c) Multi-Scale
(d) HIM
Figure 1: The overall framework of our HI-Diff. (a) Transformer, adopts hierarchical encoder-decoder architecture, equipped with HIM, for the deblurring process. (b) Diffusion Model, is performed in a highly compact latent space for computational efficiency. (c) The multi-scale prior feature  $\{\mathbf{z}_1,\mathbf{z}_2,\mathbf{z}_3\}$  is obtained by downsampling the prior feature multiple times. In stage one,  $\mathbf{z}_1 = \mathbf{z}$ ; in stage two,  $\mathbf{z}_1 = \hat{\mathbf{z}}$ . (d) The hierarchical integration module (HIM), calculates cross-attention between the intermediate feature of Transformer and the multi-scale prior feature. (e) The latent encoder (LE), where the size of the input feature  $(in)$  is  $H\times W\times 6$  for LE, and  $H\times W\times 3$  for  $\mathrm{LE}_{\mathrm{DM}}$ .
(e) LE
Transformer-based Methods. Recently, Transformer [46, 9] proposed in natural language processing (NLP), which utilizes the self-attention (SA) mechanism to model long-range dependence, has shown remarkable performance in image deblurring [48, 53, 44]. For example, Restormer [53] utilizes "transposed" attention, which calculates the SA across channel dimension instead of spatial dimension, to model global information and keep computational efficiency to large images. Stripformer [44] constructs horizontal and vertical strip tokens to catch region-specific blurred patterns with different orientations in dynamic scenes. These methods further improved the deblurring performance compared with CNN-based Methods. However, they are limited in recovering image details, since the regression-based methods are conservative with high-frequency details [36].
Diffusion Models (DMs) [40, 15] are probabilistic generative models that enable the construction of desired data samples from the Gaussian noise via a stochastic iterative denoising process. DMs have demonstrated excellent performance in various image restoration tasks [18, 17, 47, 7], such as super-resolution [36, 25], inpainting [35], and deblurring [50]. A survey [26] summarizes more diffusion-based restoration methods. For example, DiffIR [51] utilizes the diffusion model to generate a prior representation for image restoration, and applies a two-stage training approach. DvSR [50] introduces conditional DMs into image deblurring with a residual model, showing more realistic results than regression-based methods. InDI [7] directly produces high-quality results by restoring low-quality images. However, one of the restrictions of DMs is the large number of iteration steps required in inference. Some methods attempt to mitigate this limitation by applying well-designed noise schedules [3] and sampling strategies [41] or performing DM on the latent space [35]. Despite these efforts, the overall complexity remains high, especially for high-resolution images common in image deblurring. Additionally, these methods are susceptible to misaligned details distribution and undesired artifacts, resulting in poor performance in distortion-based metrics, e.g., PSNR.
Our goal is to integrate the diffusion model (DM) and the regression-based model, thus overcoming the shortcomings of the DM and realizing better deblurring. We propose the Hierarchical Integration Diffusion Model (HI-Diff). The HI-Diff performs the DM to generate the prior feature that is integrated hierarchically into Transformer (the regression-based model).
The overall framework of the proposed HI-Diff is depicted in Fig. 1. The HI-Diff consists of two parts: Transformer and the latent diffusion model. We adopt Restormer [53], a hierarchical encoder-decoder Transformer architecture, in our method. Compared with other Transformer networks designed specifically for image deblurring, Restormer is a general restoration model. Applying this model in our method can better illustrate the effectiveness of our proposed method. Meanwhile, following previous practice [10, 35, 51], we train our HI-Diff with a two-stage training strategy, to realize latent compression and the training of the DM. In this section, we first elaborate on the two-stage training framework and then illustrate the whole deblurring inference process.
3.1 Stage One: Latent Compression
In stage one, our purpose is to compress the ground truth image into the highly compact latent space, and utilize it to guide Transformer in the deblurring process. As shown in Fig. 1(a, b), we compress the ground truth images through a latent encoder (LE) to obtain a compact representation as the prior feature. Then we integrate the prior feature into Transformer through the hierarchical integration module (HIM). The prior feature can provide explicit guidance for Transformer, thus increasing the details of the reconstructed image. Next, we describe these parts in detail.
Latent Encoder. As shown in Fig. 1(b), given the blurry input image  $\mathbf{I}_{Blur} \in \mathbb{R}^{H \times W \times 3}$  and its corresponding ground truth counterpart  $\mathbf{I}_{GT} \in \mathbb{R}^{H \times W \times 3}$ , we first concatenate them along the channel dimension and feed them into the latent encoder (LE) to generate the prior feature  $\mathbf{z} \in \mathbb{R}^{N \times C'}$ . Here  $H$  and  $W$  represent the image height and width, while  $N$  and  $C'$  are the token number and channel dimensions of  $\mathbf{z}$ . Importantly, the token number  $N$  is a constant much smaller than  $H \times W$ . The compression ratio  $(\frac{H \times W}{N})$  is much higher than that applied in previous latent diffusion [35] (e.g., 8 times). Therefore, the computational burden of the subsequent latent diffusion model is effectively reduced. The details of the latent encoder are depicted in Fig. 1(e), which contains  $L$  residual blocks.
Hierarchical Integration Module. To effectively integrate the prior feature and intermediate feature of Transformer, we propose the hierarchical integration module (HIM). As illustrated in Fig. 1(a), the HIM is placed in front of each encoder and decoder. For each HIM, cross-attention is computed between the prior and intermediate features for feature fusion. This module allows the information in the prior feature to be aggregated into features of Transformer.
Specifically, as shown in Fig. 1(d), given the intermediate feature  $\mathbf{X}_{in} \in \mathbb{R}^{\hat{H} \times \hat{W} \times \hat{C}}$ , we reshaped it as tokens  $\mathbf{X}_r \in \mathbb{R}^{\hat{H}\hat{W} \times \hat{C}}$ ; where  $\hat{H} \times \hat{W}$  is spatial resolution, and  $\hat{C}$  denotes channel dimension. Then we linearly project  $\mathbf{X}_r$  into  $\mathbf{Q} \in \mathbb{R}^{\hat{H}\hat{W} \times \hat{C}}$  (query). Similarly, we project the prior feature  $\mathbf{z}_i \in \mathbb{R}^{\hat{N} \times C'}$  as  $\mathbf{K} \in \mathbb{R}^{\hat{N} \times \hat{C}}$  (key) and  $\mathbf{V} \in \mathbb{R}^{\hat{N} \times \hat{C}}$  (value). The cross-attention is formulated as:
$$
\mathbf {Q} = \mathbf {W} _ {Q} \mathbf {X} _ {r}, \mathbf {K} = \mathbf {W} _ {K} \mathbf {z} _ {i}, \mathbf {V} = \mathbf {W} _ {V} \mathbf {z} _ {i},
$$
$$
\operatorname {A t t e n t i o n} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \operatorname {S o f t M a x} \left(\mathbf {Q} \mathbf {K} ^ {T} / \sqrt {\hat {C}}\right) \cdot \mathbf {V}, \tag {1}
$$
where  $\mathbf{W}_Q\in \mathbb{R}^{\hat{C}\times \hat{C}}$ $\mathbf{W}_K\in \mathbb{R}^{C'\times \hat{C}}$  , and  $\mathbf{W}_V\in \mathbb{R}^{C'\times \hat{C}}$  represent learnable parameters of linear projections without bias. As vanilla multi-head self-attention [46, 9], we separate channels into multiple "heads" and calculate the attention operations. Note that Fig. 1(d) depicts the situation with a single head and omits some details for simplification. Finally, we reshape and project the output of cross-attention, and add it with  $\mathbf{X}_{in}$  to derive the output feature  $\mathbf{X}_{out}\in \mathbb{R}^{\hat{H}\times \hat{W}\times \hat{C}}$
Moreover, since the non-uniform blur in real scenarios, the single-scale prior feature cannot adapt well to complex blurry situations. Therefore, we generate the multiple-scale prior feature  $\{\mathbf{z}_1,\mathbf{z}_2,\mathbf{z}_3\}$  (where  $\mathbf{z}_1 = \mathbf{z}$ ), by downsampling the prior feature  $\mathbf{z}$ , as shown in Fig. 1(c). The multiple-scale prior feature adapts to different scale intermediate features for better fusion. The effectiveness of the hierarchical integration with the multiple-scale prior feature is demonstrated in Sec. 4.2.
Training Strategy. To ensure the effectiveness of the latent encoder (LE) in constructing the prior feature, we optimize it jointly with Transformer using the  $L_{1}$  loss function, defined as:
$$
\mathcal {L} _ {\mathrm {d e b l u r}} = \left\| \mathbf {I} _ {D B} - \mathbf {I} _ {G T} \right\| _ {1}, \tag {2}
$$
where  $\mathbf{I}_{DB}$  is the deblurred image, and  $\mathbf{I}_{GT}$  represents its corresponding ground truth.
3.2 Stage Two: Latent Diffusion Model
In stage two, a latent diffusion model (DM) is trained to learn to generate the prior feature, that enhances the deblurring process of Transformer through HIM.
Diffusion Model. Specifically, our latent diffusion model is based on conditional denoising diffusion probabilistic models [15, 3, 36, 35]. The diffusion model involves a forward diffusion process and a reverse denoising process, as illustrated in Fig. 1(b).
In the diffusion process, given a ground truth image, we first adopt the latent encoder (LE) trained in stage one to generate the corresponding prior feature  $\mathbf{z} \in \mathbb{R}^{N \times C'}$ . We take  $\mathbf{z}$  as the starting point of the forward Markov process, and gradually add Gaussian noise to it over  $T$  iterations as follows:
$$
q \left(\mathbf {z} _ {1: T} \mid \mathbf {z} _ {0}\right) = \prod_ {t = 1} ^ {T} q \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {t - 1}\right), \quad q \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {t - 1}\right) = \mathcal {N} \left(\mathbf {z} _ {t}; \sqrt {1 - \beta_ {t}} \mathbf {z} _ {t - 1}, \beta_ {t} \mathbf {I}\right), \tag {3}
$$
where  $t = 1, \dots, T$ ;  $\mathbf{z}_t$  represents the noisy features at the  $t$ -th step;  $\mathbf{z}_0 = \mathbf{z}$  for unification;  $\beta_{1:T} \in (0,1)$  are hyperparameters that control the variance of the noise;  $\mathcal{N}$  denotes the Gaussian distribution. Through iterative derivation with reparameterization [20], Eq. (3) can be written as:
$$
q \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {0}\right) = \mathcal {N} \left(\mathbf {z} _ {t}; \sqrt {\bar {\alpha} _ {t}} \mathbf {z} _ {0}, (1 - \bar {\alpha} _ {t}) \mathbf {I}\right), \quad \alpha = 1 - \beta_ {t}, \quad \bar {\alpha} _ {t} = \prod_ {i = 1} ^ {t} \alpha_ {i}. \tag {4}
$$
In the reverse process, we aim to generate the prior feature from a pure Gaussian distribution. The reverse process is a  $T$ -step Markov chain that runs backwards from  $\mathbf{z}_T$  to  $\mathbf{z}_0$ . Specifically, for the reverse step from  $\mathbf{z}_t$  to  $\mathbf{z}_{t-1}$ , we use the posterior distribution as:
$$
q \left(\mathbf {z} _ {t - 1} \mid \mathbf {z} _ {t}, \mathbf {z} _ {0}\right) = \mathcal {N} \left(\mathbf {z} _ {t - 1}; \boldsymbol {\mu} _ {t} \left(\mathbf {z} _ {t}, \mathbf {z} _ {0}\right), \frac {1 - \bar {\alpha} _ {t - 1}}{1 - \bar {\alpha} _ {t}} \beta_ {t} \mathbf {I}\right), \quad \boldsymbol {\mu} _ {t} \left(\mathbf {z} _ {t}, \mathbf {z} _ {0}\right) = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {z} _ {t} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon}\right), \tag {5}
$$
where  $\epsilon$  represents the noise in  $\mathbf{z}_t$ , and is the only uncertain variable. Following previous work [15, 36, 35, 51], we adopt a neural network (denoted as denoising network,  $\epsilon_{\theta}$ ) to estimate the noise  $\epsilon$  for each step. Since DM operates in the latent space, we utilize another latent encoder, denoted as  $\mathrm{LE}_{\mathrm{DM}}$ , with the same structure as LE.  $\mathrm{LE}_{\mathrm{DM}}$  compresses the blurry image  $\mathbf{I}_{Blur}$  into latent space to get the condition latent  $\mathbf{c} \in \mathbb{R}^{N \times C'}$ . The denoising network predicts the noise conditioned on the  $\mathbf{z}_t$  and  $\mathbf{c}$ , i.e.,  $\epsilon_{\theta}(\mathbf{z}_t, \mathbf{c}, t)$ . With the substitution of  $\epsilon_{\theta}$  in Eq. (5) and set the variance to  $(1 - \alpha_t)$ , we get:
$$
\boldsymbol {z} _ {t - 1} = \frac {1}{\sqrt {\alpha_ {t}}} (\boldsymbol {y} _ {t} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon} _ {\theta} (\boldsymbol {z} _ {t}, \mathbf {c}, t)) + \sqrt {1 - \alpha_ {t}} \boldsymbol {\epsilon} _ {t}, \tag {6}
$$
where  $\epsilon_t \sim \mathcal{N}(0, \mathbf{I})$ . By iteratively sampling  $\mathbf{z}_t$  using Eq. (6)  $T$  times, we can generate the predicted prior feature  $\hat{\mathbf{z}} \in \mathbb{R}^{N \times C'}$ , as shown in Fig. 1(b). The predicted prior feature is then used to guide Transformer, that is,  $\mathbf{z}_1 = \hat{\mathbf{z}}$  in Fig. 1(c). Notably, since the distribution of the latent space  $(\mathbb{R}^{N \times C'})$  is much simpler than that of images  $(\mathbb{R}^{H \times W \times C'})$  [50, 51], the prior feature can be generated with a small number of iterations. We further explore the iteration numbers  $T$  in Sec. 4.2.
Training Strategy. Training DM means training denoising network  $\epsilon_{\theta}$ . Previous works [15, 41] train the model by optimizing the weighted variational bound. The training objective is:
$$
\nabla_ {\boldsymbol {\theta}} \| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} _ {\theta} (\sqrt {\bar {\alpha} _ {t}} \mathbf {z} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon}, \mathbf {c}, t) \| _ {2} ^ {2}, \tag {7}
$$
where  $\mathbf{z}$  and  $\mathbf{c}$  are prior feature and condition latent defined above;  $t\in [1,T]$  is a random time-step;  $\epsilon \sim \mathcal{N}(0,\mathbf{I})$  denotes sampled noise. However, the objective in Eq. (7) only trains DM. Since the slight deviation between the predicted prior feature and the actual prior  $\mathbf{z}$ , directly combining DM with Transformer could cause a mismatch, which restricts the deblurring performance.
To overcome this issue, we jointly train the diffusion model and Transformer, inspired by previous work [51]. In this approach, for each training iteration, we use the prior feature  $\mathbf{z}$  to generate the noise sample  $\mathbf{z}_T$  through Eq. (4). As the time-step  $T$  is small in latent space, we then run the complete  $T$  iteration reverse processes (Eq. (6)) to generate the predicted prior feature  $\hat{\mathbf{z}}$ . The  $\hat{\mathbf{z}}$  is used to guide Transformer through HIM ( $\mathbf{z}_1 = \hat{\mathbf{z}}$ ). Finally, the training loss is given by the sum of the deblurring loss  $\mathcal{L}_{\mathrm{deblur}}$  and the diffusion loss  $\mathcal{L}_{\mathrm{diffusion}}$ , where the diffusion loss  $\mathcal{L}_{\mathrm{diffusion}} = \| \hat{\mathbf{z}} - \mathbf{z} \|_1$ .