SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jinhui Hou $^{1}$ , Zhiyu Zhu $^{1}$ , Junhui Hou $^{1*}$ , Hui Liu $^{2}$ , Huanqiang Zeng $^{3}$ , and Hui Yuan $^{4}$
<sup>1</sup>City University of Hong Kong, <sup>2</sup>Caritas Institute of Higher Education
$^{3}$ Huaqiao University,  $^{4}$ Shandong University
This paper studies a diffusion-based framework to address the low-light image enhancement problem. To harness the capabilities of diffusion models, we delve into this intricate process and advocate for the regularization of its inherent ODE-trajectory. To be specific, inspired by the recent research that low curvature ODE-trajectory results in a stable and effective diffusion process, we formulate a curvature regularization term anchored in the intrinsic non-local structures of image data, i.e., global structure-aware regularization, which gradually facilitates the preservation of complicated details and the augmentation of contrast during the diffusion process. This incorporation mitigates the adverse effects of noise and artifacts resulting from the diffusion process, leading to a more precise and flexible enhancement. To additionally promote learning in challenging regions, we introduce an uncertainty-guided regularization technique, which wisely relaxes constraints on the most extreme regions of the image. Experimental evaluations reveal that the proposed diffusion-based framework, complemented by rank-informed regularization, attains distinguished performance in low-light enhancement. The outcomes indicate substantial advancements in image quality, noise suppression, and contrast amplification in comparison with state-of-the-art methods. We believe this innovative approach will stimulate further exploration and advancement in low-light image processing, with potential implications for other applications of diffusion models. The code is publicly available at https://github.com/jinnh/GSAD.
Low-light image enhancement (LLIE) aims to improve the visibility and contrast of the image captured in poor lighting while preserving the natural-looking details, which contributes to many downstream applications, such as object detection [15, 58] and semantic segmentation [62, 48].
Traditional works employ some techniques, such as histogram equalization [28], Retinex theory [6, 23], and gamma correction [29], to correct the image illumination. In recent years, owing to the powerful representational ability and large collections of data, a considerable number of deep learning-based methods [53, 54, 63, 5, 21, 64, 43, 51, 40, 65] have been presented to significantly improve the performance of low-light enhancement by learning the mapping between low-light and normal-light images. Generally, most of the existing works tend to adopt pixel-wise objective functions to optimize a deterministic relationship. Consequently, such regularization frequently produces suboptimal reconstructions for indeterminate regions and poor local structures, resulting in visibly lower reconstruction quality. Although the adversarial loss might mitigate this issue, these methods necessitate careful training adjustments, which might lead to overfitting on specific features or data distributions and even creating new content or artifacts. Recently, the popular diffusion
*Corresponding author (Email: jh.hou@cityu.edu.hk). Jinhui Hou and Zhiyu Zhu contributed to this work equally.
Algorithm 1 Forward process
1:Repeat
2:  $\mathbf{X}_0\sim q(\mathbf{X}_0)$
3:  $t\sim \mathrm{Uniform}(\{1,2,\dots,T\})$
4:  $\epsilon \sim \mathcal{N}(0,\mathbf{I})$
5: Take gradient descent step on
6:  $\nabla_{\theta}\| \epsilon -\epsilon_{\theta}(\sqrt{\overline{\alpha}_{t}}\mathbf{X}_{0} + \sqrt{1 - \overline{\alpha}_{t}}\epsilon ,t)\|^{2}$
7: Until converged
Algorithm 2 Reverse process
1:  $\mathbf{X}_T\sim \mathcal{N}(0,\mathbf{I})$
2: For  $t = T, \dots, 1$  do
3:  $\mathbf{z} \sim \mathcal{N}(0, \mathbf{I})$  if  $t > 1$ , else  $\mathbf{z} = 0$
4:  $\mathbf{X}_{t - 1} = \frac{1}{\sqrt{\alpha_t}} (\mathbf{X}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}}\epsilon_\theta (\mathbf{X}_t,t)) + \sigma_t\mathbf{z}$
5: End for
6: Return  $\mathbf{X}_0$
denoising probabilistic model (DDPM) [7] has garnered notable interest in low-level vision domains [9, 32], credited with its outstanding capacity in modeling the distribution of image pixels. However, a straightforward implementation of the diffusion model for low-light image enhancement is insufficient to address this issue.
In this paper, we present a novel diffusion-based method to boost the performance of low-light enhancement from the perspective of regularizing the ODE-trajectory. Drawing inspiration from the evidence that suggests the potential of a low-curvature trajectory in yielding superior reconstruction performance [12], we endeavor to modulate the ODE-trajectory curvature. Specifically, we introduce a global structure-aware regularization scheme into the diffusion-based framework by gradually exploiting the intrinsic structure of image data. This innovative constraint promotes structural consistency and content coherence across similar regions, contributing to a more natural and aesthetically pleasing image enhancement while preserving the image's fine details and textures. In addition, we devise an uncertainty-guided regularization by integrating an uncertainty map into the diffusion process, facilitating adaptive modulation of the regularization strength. Experimental results on a comprehensive set of benchmark datasets consistently demonstrate the superior performance of our proposed method compared to state-of-the-art ones. A thorough ablation study emphasizes the significance of both the global structure-aware regularization and uncertainty-guided constraint components in our approach, revealing their synergistic effects on the overall quality of enhancement.
Low-light image enhancement. The early works tackle low-light image enhancement by applying some traditional techniques like histogram equalization [28, 36], gamma correction [29], and Retinex theory [6, 23]. Some researchers also seek to improve the visibility using additional sensors[67, 66, 42, 41, 37, 17]. In recent years, with the advancement of low-light data collection [11, 6, 45, 14, 54] a significant number of deep learning-based methods [2, 60, 8, 50, 53, 54, 63, 5, 21, 64, 43, 51, 40] have been proposed, which greatly improved the restoration quality of traditional methods. For example, Retinex-based deep learning methods [45, 60, 59, 47] employed deep learning to decompose low-light images into two smaller subspaces, i.e., illumination and reflectance maps. Wang et al. [43] proposed a normalizing flow-based low-light image enhancement approach that models the distributions across normally exposed images. Xu et al. [51] incorporated the Signal-to-Noise-Ratio (SNR) prior to achieve spatial-varying low-light image enhancement. Wang et al. [40] proposed a transformer-based low-light enhancement method. We refer readers to [13] for a comprehensive review of this field.
Diffusion-based image restoration. Recently, diffusion-based generative models [34] have delivered astounding results with the advancements in denoising diffusion probabilistic models (DDPM) [7, 24], making them increasingly influential in low-level vision tasks, such as image super-resolution [9, 32], image inpainting [30, 19], image deraining [26], and image deblurring [46]. Saharia et al. [32] employed a denoising diffusion probabilistic model to image super-resolution by conditioning the low-resolution images in the diffusion process. Lugmayr et al. [19] proposed a mask-agnostic approach for image inpainting by leveraging the pre-trained unconditional DDPM as the generative prior. Ozan et al. [26] developed a patch-based diffusion model for weather removal that utilized a mean estimated noise-guided sampling update strategy across overlapping patches. These models are based on a diffusion process that transforms a clean image into a noisy one during training time, and then reverses that Markov Chain to generate new images in the testing phase. Consider a clean image  $\mathbf{X}_0\in \mathbb{R}^{h\times w}$  and a noise variable  $\epsilon \in \mathbb{R}^{h\times w}\sim \mathcal{N}(0,\sigma^2 I)$ . The training steps of a diffusion process act as Algorithm 1.
Figure 1: Reverse trajectories of diffusion model with and without our global structure-aware regularization on the testing set of LOLv1 dataset. Our method effectively compactifies the distribution of reverse trajectories of multiple samplings, producing low-curvature reverse trajectories, which makes them stably approach the Ground Truth (GT) values.
Note  $\sqrt{\overline{\alpha}_t}\mathbf{X}_0 + \sqrt{1 - \overline{\alpha}_t\epsilon}$  is a closed-form noisy sample at timestamp  $t$ , and  $\epsilon_{\theta}(\cdot ,\cdot)$  denotes a noise estimation network, which perceives the closed-form samples with timestamps and predicts the inherent noise in such samples. By applying this process in reverse, shown as Algorithm 2, from  $\mathbf{X}_T\sim \mathcal{N}(0,\mathbf{I})$  to  $\mathbf{X}_0$ , we obtain a generated image  $\mathbf{X}_0$ , where  $\mathbf{X}_t\in \mathbb{R}^{h\times w}$  is the intermediate reconstructed sample at timestamp  $t\in [0,1,\dots,T]$ ,  $\alpha_{t}$  denotes a scaling scalar, usually parameterized as a linearly decreasing sequence [7, 32], and  $\bar{\alpha}_{t} = \prod_{i = 1}^{t}\alpha_{i}$ .
3.1 Problem Statement and Overview
Low-light image enhancement is a crucial research area in the field of image processing and computer vision, owing to the increasing demand for high-quality images captured in various low-light environments. Mathematically, the degradation of an image captured under low-light conditions can be modeled as
$$
\mathbf {Y} = \mathbf {X} _ {0} \odot \mathbf {S} + \mathbf {N}, \tag {1}
$$
where  $\mathbf{Y} \in \mathbb{R}^{H \times W \times 3}$  is the observed low-light image of dimensions  $H \times W$ ,  $\mathbf{X}_0 \in \mathbb{R}^{H \times W \times 3}$  is the latent high-quality image under normal-light conditions,  $\mathbf{S} \in \mathbb{R}^{H \times W \times 3}$  represents the spatially varying illumination map that accounts for the uneven lighting, and  $\mathbf{N} \in \mathbb{R}^{H \times W}$  is the noise term introduced due to sensor limitations and other factors. Thus, low-light image enhancement, particularly in reconstructing extreme low-light regions, is of paramount challenge due to the scarcity of available content in such images. The intricacy is further intensified by the content-dependent nature of optimal lighting conditions, necessitating a flexible and adaptive solution. Fortunately, diffusion models, celebrated for their remarkable capacity to synthesize images conditioned on ancillary images or textual cues, emerge as a propitious candidate for tackling the low-light image enhancement conundrum.
In this paper, we advocate capitalizing on the potent generalization capabilities inherent in diffusion models to overcome the obstacles associated with low-light image enhancement. By incorporating low-light images  $(\mathbf{Y})$  as conditioning inputs into the diffusion models, we can obtain a naive implementation of low-light enhancement diffusion model  $(\epsilon_{\theta}(\mathbf{Y},\mathbf{X}_t,\bar{\alpha}_t))$ . Nonetheless, the performance achieved by these preliminary models remains unsatisfactory. In pursuit of procuring high-quality diffusion processes, various strategies have been presented, including alterations to the network structure [31], the incorporation of data augmentation techniques [38], and the introduction of innovative loss terms [22]. Among them, trajectory rectification emerges as a promising strategy, facilitating a seamless and stable diffusion process through the regularization of a low curvature reverse trajectory
Figure 2: Illustration of the training workflow of the proposed method. We begin by obtaining a closed-form sample  $\mathbf{X}_t$  from  $\mathbf{X}_0$  and the noise level  $\bar{\alpha}_t$ , which is subsequently fed into a U-Net  $\epsilon(\cdot)$  conditioned with the corresponding low-light image  $\mathbf{Y}$  and  $\bar{\alpha}_t$  to estimate the employed noise  $\epsilon$ . To comprehensively regularize reconstructed images, we proceed to construct a learnable closed-form sample  $\mathbf{X}_{t-1}$  using the estimated noise  $\epsilon_{\theta}(\mathbf{Y}, \mathbf{X}_t, \bar{\alpha}_t)$  and optimize the model using uncertainty-guided noise estimation and intrinsic image global structure. Meanwhile, we use a pre-trained uncertainty estimation model  $\mu_{\theta_1}$  to obtain the uncertainty map  $\mathbf{P}_t$ .
[12, 18]. Encouraged by current successes, we endeavor to augment this preliminary model from the following two distinct perspectives.
1) Learning low curvature trajectories with global structure-aware regularization. To learn a low curvature reverse trajectory, we propose a simple yet effective way via directly minimizing the gap between a learnable sample on the reverse trajectory and the ground-truth sample. This process is further elucidated through pertinent experimental results presented in Fig. 1. Moreover, it is crucial to acknowledge that images intrinsically encompass analogous textures or patterns distributed across different locations [1, 27, 3, 4, 49, 16, 56]. As a result, cultivating a diffusion process that adeptly captures such global structures is indispensable to the success of low-light image reconstruction. By integrating the matrix rank-based regularization within the diffusion-based framework, we astutely exploit the intrinsic structure of image data, thereby facilitating the preservation of nuanced details and the enhancement of contrast.
Despite the effectiveness of the trajectory regularization, precipitating it during the beginning steps of diffusion models could unintentionally subvert sample diversity and quality, attributable to the pronounced fluctuations in large components. Consequently, in pursuit of enhancing the effectiveness of our model, it becomes imperative to architect a progressively adaptive regularization schedule.
2) Adaptive regularization with uncertainty. Furthermore, most contemporary loss functions treat all pixels equally in the field of low-light enhancement tasks, regardless of whether they exist in conditions of slightly low-light or extremely low-brightness, or whether they are subject to noise-deterioration or are noise-free. Therefore, drawing inspiration from [25], we introduce an uncertainty-guided regularization into the diffusion process to further improve restoration learning in challenging regions.
Those strategic incorporations ensure the precise and resilient recovery of image details in struggling low-light regions while accommodating scene-specific lighting variations. In the following sections, we give technical details of our method.
3.2 Exploring Global Structures via Matrix Rank Modeling
The denoising diffusion probabilistic models exhibit a unique learning scheme, which involves learning the distribution of latent noise during training. This characteristic makes it challenging to regularize network parameters  $\theta$  for learning global structure-aware representations. To address this
issue, we first construct a learnable closed-form sample of  $\mathbf{X}_{t - 1}$ . Subsequently, we implement two further steps to enhance the content-aware regularization in those images.
1) Constructing learnable closed-form samples. The training of diffusion models commences with the procurement of a closed-form  $\mathbf{X}_t$  at an arbitrary timestep  $t$ , as delineated in Line 6 from Algorithm 1. In the ensuing step, a learnable function  $\epsilon_{\theta}(\mathbf{Y},\mathbf{X}_t,\bar{\alpha}_t)$  is employed to estimate the latent noise  $\epsilon$ . Given that the network can successfully learn this noise, it becomes feasible to progressively eliminate the noise, as demonstrated by:
$$
\mathbf {X} _ {t - 1} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {X} _ {t} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \epsilon_ {\theta} \left(\mathbf {Y}, \mathbf {X} _ {t}, \bar {\alpha} _ {t}\right)\right) + \sigma_ {t} \mathbf {Z}. \tag {2}
$$
Consequently, throughout the training phase, the regularization is strategically applied to the learnable sample  $\mathbf{X}_{t-1}$  in lieu of the closed-form sample  $\mathbf{X}_t$ . This approach seeks to optimize the performance of the model by focusing on regularization where it can most effectively influence the learning process.
2) Non-local patch-based matrix rank modeling. To find global structures of a clean image  $\mathbf{X}_0$ , we first patchify it into a set of non-overlapping blocks  $\{\mathbf{x}_0^{(i)}\}_{i=1}^n$ , where  $\mathbf{x}_0^{(i)} \in \mathbb{R}^m$  denotes the vectorization of a small block from  $\mathbf{X}_0$ . To capture the intrinsic structures from the sample  $\mathbf{X}_0$ , we then apply the commonly-used clustering algorithm to get  $k$  clusters of  $\{\mathbf{x}_0^{(i)}\}_{i=1}^n$  as  $\left\{\{\mathbf{x}_0^{(i)}\}_{i=1}^{n_j}\right\}_{j=1}^k$ . It is worth noting that our method benefits from advanced clustering algorithms greatly (See Sec. 4.3 for details). Through stacking each cluster of vectors into a matrix, we could get  $\{\mathbf{M}_0^{n_j}\}_{j=1}^k$ , where the rank of matrix  $\mathbf{M}_0^{n_j} \in \mathbb{R}^{m \times n_j}$  is an effective measurement of the global structures. Consequently, to regularize structure consistency in  $\mathbf{X}_{t-1}$  and  $\mathbf{X}_0$ , we then apply the same aggregation to the learnable sample  $\mathbf{X}_{t-1}$ , resulting in  $\{\mathbf{M}_{t-1}^{n_j}\}_{j=1}^k$ , and regularize
$$
\mathcal {L} _ {s} \left(\mathbf {M} _ {t - 1} ^ {n _ {j}}, \mathbf {M} _ {0} ^ {n _ {j}}\right) = \| \operatorname {d i a g} \left(\mathbf {S} _ {t - 1} ^ {j}\right) - \operatorname {d i a g} \left(\mathbf {S} _ {0} ^ {j}\right) \| _ {1}, \mathbf {M} _ {t} ^ {n _ {j}} = \mathbf {U} _ {t} ^ {j} \mathbf {S} _ {t} ^ {j} \mathbf {V} _ {t} ^ {j}, \tag {3}
$$
where  $\mathcal{L}_s(\mathbf{M}_{t - 1}^{n_j},\mathbf{M}_0^{n_j})$  indicates the regularization function,  $\mathbf{U}_t^j\in \mathbb{R}^{m\times m}$ $\mathbf{S}_t^j\in \mathbb{R}^{m\times n_j}$ , and  $\mathbf{V}_t^j\in \mathbb{R}^{n_j\times n_j}$  are the outputs of the singular value decomposition of matrix  $\mathbf{M}_t^{n_j}$ , and  $\mathrm{diag}(\cdot)$  returns a vector of the matrix main diagonal.
3) Gradually injecting structure-aware regularization. This strategy delicately calibrates the focus on patches throughout different learning steps, thereby ensuring a nuanced, stepwise infusion of structure-aware regularization. This judicious approach not only preserves the heterogeneity of the reconstructed samples but also optimizes the noise suppression process, culminating in a notable improvement in the overall performance of our model. Specifically, the modified regularization function  $\mathcal{L}_s^t$  is written as
$$
\mathcal {L} _ {s} ^ {t} = \kappa_ {t} \mathcal {L} _ {s} (\mathbf {M} _ {t - 1} ^ {n _ {j}}, \mathbf {M} _ {0} ^ {n _ {j}}), \kappa_ {t} \sim \{\bar {\alpha} _ {1} ^ {2}, \bar {\alpha} _ {2} ^ {2}, \dots , \bar {\alpha} _ {t} ^ {2} \}, t \sim \{1, 2, \dots , T \}, \tag {4}
$$
where  $\kappa_{t}$  denotes an adaptive factor for regularization term  $\mathcal{L}_s(\mathbf{M}_{t - 1}^{n_j},\mathbf{M}_0^{n_j})$ . This schedule allows for a more adaptive and flexible approach to regularization, enabling the model to capture better the global structures and details within the low-light images.
Analysis of choosing rank-based global structure-aware regularization. To minimize the curvature of ODE-trajectory, we bridge the deviation between the learnable closed-form trajectory samples  $X_{t-1}$  and the GT  $X_0$  by our rank-based global structure-aware regularization. Moreover, it is essential to acknowledge that while conventional pixel-wise regularization terms, including  $L_1$ ,  $L_2$ , and SSIM, do offer some degree of enhancement to the original diffusion model, their impact is constrained. This limitation arises from their inadequate ability to fully encapsulate the non-local structures and intricate patterns within an image. Although regularization within the feature domain potentially aids in modeling a given image's structure, such structure is usually confined within a local region due to the kernel perceptual regions. This type of regularization, moreover, cannot explicitly characterize the properties of the structure and is lacking in theoretical guidance. Another significant concern is that regularization features could result in significant fluctuations in the back-propagated gradients, thereby impeding network training. Lastly, the empirical evidence documented in Table 4 corroborates the superiority of our rank-based global structure-aware regularization.
3.3 Integrating Uncertainty into Diffusion Process
To integrate uncertainty into the diffusion process, inspired by [25], we introduce an additional convolutional branch after the last layer of U-Net of the diffusion model for generating the uncertainty map. Specifically, we train the uncertainty model  $\mu_{\theta_1}(\cdot)$  with the same step as the denoising diffusion models, like Algorithm 1, except that we optimize the following objective:
$$
\mathcal {L} _ {\mu} = \left\| e ^ {- \mathbf {P} _ {t}} \odot \left(\epsilon - \epsilon_ {\theta} \left(\mathbf {Y}, \mathbf {X} _ {t}, \bar {\alpha} _ {t}\right)\right) \right\| _ {1} + 2 \mathbf {P} _ {t}, \tag {5}
$$
where  $\mathbf{P}_t\coloneqq \mu_{\theta_1}(\mathbf{Y},\mathbf{X}_t,\bar{\alpha}_t)$  denotes a pixel-wise uncertainty map, i.e.,  $\mathbf{P}_t\in \mathbb{R}^{H\times W\times 3}$ ,  $\epsilon_{\theta}(\cdot)$  represents the corresponding function with weights of a denoising diffusion model. Upon completion of the uncertainty estimation model's pretraining phase, we obtain  $\mu_{\hat{\theta}_1}$ , and  $\epsilon_{\hat{\theta}}$ . Subsequently, we fix the parameter of  $\mu_{\hat{\theta}_1}$  to generate the uncertainty map  $\mathbf{P}_t$ . This map serves as a flexible weight, calibrating the importance of each pixel such that pixels with higher uncertainty are assigned greater weight. Further details are provided in Algorithm 3. The integration of uncertainty into the diffusion process in this manner incentivizes the diffusion model to intensify its focus on uncertain or challenging regions. This strategic approach yields visually pleasing results, underscoring the value of incorporating uncertainty-aware regularization in the enhancement of low-light images.
Algorithm 3 Training a diffusion-based low-light enhancement model
Require: Normal and low-light image pairs  $(\mathbf{X}_0,\mathbf{Y})$  , and a pre-trained models  $\mu_{\theta_1}$  and  $\epsilon_{\hat{\theta}}$  1: Initialize parameters of  $\mu_{\theta_1}(\cdot)$  (resp.  $\epsilon_{\theta}(\cdot))$  from pretrained  $\hat{\theta}_{1}$  (resp.  $\hat{\theta}$  ), and freeze  $\hat{\theta}_{1}$    
2: Repeat   
3: Sample an image pair  $(\mathbf{X}_0,\mathbf{Y})$ $\bar{\alpha}_t\sim p(\bar{\alpha})$  and  $\epsilon \sim \mathcal{N}(0,\mathbf{I})$    
4:  $\mathbf{X}_t = \sqrt{\bar{\alpha}_t}\mathbf{X}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon$    
5:  $\mathbf{P}_t\gets \mu_{\theta_1}(\mathbf{Y},\mathbf{X}_t,\bar{\alpha}_t)$    
6: Construct the learnable closed-form sample  $\mathbf{X}_{t - 1}$  via Eq. (2)   
7: Perform a gradient descent step on   
8:  $\nabla_{\theta}\{\lambda \| \mathbf{P}_{t}\odot \epsilon -\mathbf{P}_{t}\odot \epsilon_{\theta}(\mathbf{Y},\mathbf{X}_{t},\bar{\alpha}_{t})\|_{1} + \mathcal{L}_{s}^{t}\}$  , where  $\lambda$  is a balancing factor, and  $\mathcal{L}_s^t$  refers to Eq. (4)   
9: Until converged   
10: Return the resulting diffusion model  $\epsilon_{\tilde{\theta}}$