Title: IR-CM: The Fast and General-purpose Image Restoration Method Based on Consistency Model
Venue: NeurIPS 2024
Accepted: yes
Keywords: Image restoration

Image enhencement

SDE-based model

Consistency model
Abstract: This paper proposes a fast and general-purpose image restoration method. The key idea is to achieve few-step or even one-step inference by conducting consistency distilling or training on a specific mean-reverting stochastic differential equations. Furthermore, based on this, we propose a novel linear-nonlinear decoupling training strategy, significantly enhancing training effectiveness and surpassing consistency distillation on inference performance. This allows our method to be independent of any pre-trained checkpoint, enabling it to serve as an effective standalone image-to-image transformation model. Finally, to avoid trivial solutions and stabilize model training, we introduce a simple origin-guided loss. To validate the effectiveness of our proposed method, we conducted experiments on tasks including image deraining, denoising, deblurring, and low-light image enhancement. The experiments show that our method achieves highly competitive results with only one-step inference. And with just two-step inference, it can achieve state-of-the-art performance in low-light image enhancement. Furthermore, a number of ablation experiments demonstrate the effectiveness of the proposed training strategy. our code is available at https://github.com/XiaoxuanGong/IR-CM.

IR-CM: The Fast and General-purpose Image Restoration Method Based on Consistency Model

Abstract

This paper proposes a fast and general-purpose image restoration method. The key idea is to achieve few-step or even one-step inference by conducting consistency distilling or training on a specific mean-reverting stochastic differential equations. Furthermore, based on this, we propose a novel linear-nonlinear decoupling training strategy, significantly enhancing training effectiveness and surpassing consistency distillation on inference performance. This allows our method to be independent of any pre-trained checkpoint, enabling it to serve as an effective standalone imageto-image transformation model. Finally, to avoid trivial solutions and stabilize model training, we introduce a simple origin-guided loss. To validate the effectiveness of our proposed method, we conducted experiments on tasks including image deraining, denoising, deblurring, and low-light image enhancement. The experiments show that our method achieves highly competitive results with only one-step inference. And with just two-step inference, it can achieve state-of-the-art performance in low-light image enhancement. Furthermore, a number of ablation experiments demonstrate the effectiveness of the proposed training strategy. our code is available at https://github.com/XiaoxuanGong/IR-CM .

1 Introduction

Image restoration is a classic problem in the field of computer vision. It aims to transform low-quality or noisy images into their corresponding high-quality or noise-free counterparts. In many industrial applications (such as autonomous driving), there are various complex types of image degradation, including rain and fog interference, glare interference, low-light conditions, and motion blur etc. This necessitates high generality in image restoration algorithms. Moreover, there is often a high demand for real-time performance in practical application scenarios, posing even greater challenges to the design of algorithms and models.

Common image restoration tasks include image deraining[1, 2, 3, 4, 5, 6], denoising[7, 8, 9, 10, 11], and deblurring[12, 13, 14, 15] etc. Due to the demands of industrial applications, research on low-light image enhancement[16, 17, 18, 19, 20, 21] is also gradually increasing. However, these methods are often heuristic and difficult to apply to general tasks because they typically require domain-specific prior knowledge for training. A more general approach is the recently proposed IR-SDE[22], which

does not rely on any prior knowledge. It only requires pairs of high-quality(HQ)/low-quality(LQ) images for training and achieves good results across multiple image restoration tasks. However, like most diffusion-based methods, it requires multi-step sampling for inference, making it difficult to meet the real-time requirements of practical applications.

In recent years, diffusion models have achieved remarkable results in both unconditional image generation[23, 24, 25, 26] and conditional image generation tasks[27, 28, 20, 29, 22, 30, 21, 31, 32]. In a more general description, diffusion models can be described using stochastic differential equations (SDEs)[33, 34]. It achieves this by converting the original data distribution into a fixed Gaussian distribution, then learning the corresponding distribution gradients (score) through a network, and finally gradually reconstructing the original image using SDE solvers or ODE solvers. Most SDEbased methods can generate high-quality samples, however they often require many steps of sampling to accomplish this. Despite the existence of many methods for accelerating sampling[25, 35, 36], it remains challenging to meet the real-time requirements in industrial applications. Recently, the consistency model[37] has been proposed. It aims to map any point on the ODE trajectory of the SDE-based model to its origin. Once trained, the model can achieve one-step inference with some decrease in model performance. The training of consistency models is divided into two approaches: consistency distillation(CD) and consistency training(CT). The former involves distilling training based on the ODE trajectories of a pre-trained SDE model. The latter, while not requiring a pretrained SDE model and can be considered as an independent generative model, typically exhibits lower performance compared to the former[37]. This means that training consistency models often rely on pre-trained models, which inevitably leads to performance degradation.

The purpose of this paper is to design a universal image restoration model with fast inference. The proposed model, named IR-CM ( I mage R estoration C onsistency M odel ), achieves one-step or fewstep fast inference through consistency model training. Due to its versatility and flexibility in different image transformation tasks, the IR-SDE serves as the foundation for our method. IR-CM can be trained using consistency distillation (CD) on a pre-trained model, but at the cost of a slight decrease in performance. Therefore, we focus more on consistency training (CT), enabling IR-CM to become an independent image restoration model without relying on any pre-trained checkpoint. Furthermore, to improve the effectiveness of CT, we propose a novel linear-nonlinear decoupling training strategy and a novel origin-estimated consistency function, these allows the model's performance to reach or even surpass that of the original model. Finally, we propose a simple origin-guided loss to stabilize the training process. In summary, the main contributions of our method are as follows:

We propose a universal and fast image restoration method that can obtain high-quality images with one step or few-step sampling. For different tasks, training only requires dataset replacement, without the need for any additional prior knowledge.

We have introduced a novel linear-nonlinear decoupling training strategy, enabling our method to achieve even surpass the performance of the original model without relying on any pre-trained checkpoint.

We propose a novel origin-estimated consistency function, which enables our model to have a more stable initial state and a smaller solution space, and a simple origin-guided loss to stabilize the training process. This makes our method more robust. The ablative experiments demonstrated its effectiveness.

Our method achieves highly competitive performance in multiple tasks (including image deraining, denoising, deblurring, and low-light image enhancement) with one-step inference. With two-step inference, our method achieves state-of-the-art performance in low-light image enhancement task.

2 Preliminaries

2.1 Mean-reverting stochastic differential equation

Our method requires a SDE-based model as the base model for consistency training. Specifically, we choose IR-SDE[22] as the base SDE model due to its excellent generality and applicability. Its forward process involves gradually transforming high-quality images into corresponding low-quality versions with fixed-variance Gaussian noise. It can be represented as follows:

where µ ∈ R d typically represents the low-quality image, and x (0) ∈ R d represents its corresponding high-quality version, θ t , σ t are time-dependent positive parameters, and they satisfy σ 2 t /θ t = 2 λ 2 for all times t with positive constant λ , and dw represents Brownian motion. In [22], it has been proven that at each time t, the marginal probability distribution of x can be represented as follows:

where ¯ θ t = ∫ t 0 θ k dk . As t increases gradually, m t approaches µ and v t approaches λ 2 . Thus, x (0) (the high-quality image) will gradually transform into its corresponding low-quality version µ accompanied by Gaussian noise with variance λ 2 .

With the above conclusion, we can naturally sample x ( t ) by x ( t ) = m t + √ v t ϵ t . Then we can train a network to estimate the noise ϵ t . In the inference phase, the reverse-time process of the IR-SDE can be represented as follows[34]:

where ∇ x log p t ( x ) is called score function and ˆ ϵ ( x, µ, t ) is the noise estimated by network. Then, similar to other SDE-based models, we can adopt a SDE-solver (or ODE-solver) to reverse the process to restore the low-quality image x ( t ) back to the high-quality version x (0) progressively.

2.2 Consistency model

For a solution trajectory { x t } t ∈ [ η,T ] of any PF-ODE such as (3), a consistency function can be defined as f ( x t , t ) ≡ x η , where η is a small positive number. This means that when sampling any pair ( x t , t ) on the trajectory of the PF-ODE, the output of the consistency function is always the initial point x η of the trajectory. This property is referred to as self-consistency[37]. The example of the consistency function in [37] is as follows:

where c skip ( t ) , c out ( t ) are differentiable functions, and they satisfy c skip ( η ) = 1 , c out ( η ) = 0 , F ϕ ( x t , t ) is a trainable network initialized by a pre-trained noise estimation model. Once training is complete, we only need to input x T and apply f ϕ ( x T , T ) = x η to obtain high-quality sample in one step. There are two training methods for consistency models:

Consistency Distillation (CD) For a discrete time sequence t 1 = η < t 2 < ... < t N = T , given an arbitrary point ( x t n +1 , t n +1 ) on PF-ODE trajectory, we can estimate the x t n by following formula:

where Φ( · ) represents the update function of a one step ODE solver applied to the PF-ODE, and φ is the weights of a pre-trained score matching network. Then the CD loss can represent as follow:

where ϕ -represents the exponential moving average (EMA) version of training weights ϕ , it is frozen during backward, and λ ( · ) is a positive weighting function, d ( · ) denotes a distance function, such as the L1 or L2 distance. This approach essentially involves distillation training on the pre-trained score matching network, hence referred to as consistency distillation.

Consistency Training (CT) Unlike CD, the CT does not rely on a pre-trained score matching network and can independently train any SDE-based model. The CT loss is represented as follow:

Here x t n and x t n +1 are both sampled from the forward process of SDE model.

3 Method

The key idea of our method lies in employing our proposed two-stage training strategy and the origin-estimated consistency function to conduct consistency training on the IR-SDE model. We thus refer to it as an I mage R estoration C onsistency M odel (IR-CM).We begin by describing the novel origin-estimated consistency function, followed by an explanation of the proposed two-stage training strategy and the origin-guided loss. The overall architecture is depicted in Fig 1.

3.1 Origin-estimated consistency function

For any SDE-based model, each step of the reverse process is essentially a prediction of the PF-ODE solution[34]. In other words, any point on the PF-ODE trajectory actually contains information about the origin. Based on this, according to equations (2) and (3), the origin prediction function at each point on the PF-ODE trajectory of the IR-SDE is as follows:

During the consistency training phase(both for CD and CT), differ from the formula (4) proposed in [37], we train using the following origin-estimated consistency function (OECF):

where the ˆ x 0 ( x, t ; ϕ ) is obtained by (8). Suppose the pre-trained score model matches the ground truth, i.e., ∀ t ∈ [ η, T ] : ˆ ϵ ϕ ( x, µ, t ) = ϵ ( t ) + o (∆ t ) and c skip ( T ) = 0 , c out ( T ) = 1 . For formulas (4) and (9), c skip ( t ) , x ( t ) , and c out ( t ) are all same and constant at each time t. Thus, for simplicity in analysis, we specifically consider the moment t = T . Then substituting equation (2) into equations (4) and (9) respectively, we obtain:

From equation (10), it can be observed that if equation (4) is chosen as the consistency function, the initial solution will have a significant fixed error. In contrast, using the OECF, the initial solution will only have a time-dependent higher-order infinitesimal error. Therefore, for a pre-trained score matching model, the OECF offers a more stable initial state and hence a smaller solution space for all time t .

Intuitively, OECF can effectively enhance the training performance of CD. In practice, we have found that for CT, OECF also exhibits significant performance improvements compared to equation (4). The relevant ablative experiment results will be presented in Chapter 4. Note that, unlike CD, the pre-trained model in CT is only used to initialize the training weights and is not involved in any training process. In practice, training with randomly initialized weights is also feasible, albeit usually resulting in slightly longer convergence times.

3.2 Origin-guided loss

In general training, we randomly sample ( x t , t ) and ( x t +1 , t + 1) using equation (2), and then simply apply (6) for CD or (7) for CT, as illustrated in the lower half of Figure 1. However, we have empirically found that its performance is not ideal, and occasionally, pattern collapse occurs. Upon further analysis, we discovered that this often occurs when t is not sufficiently small during the random sampling in the early stages of training. This also leads to the emergence of mediocre solutions. Below we provide a simple theoretical proof.

Theorem 1. Let c out ( t ) be monotonic differentiable and satisfy c out ( η ) = 0 , c out ( T ) = 1 , consider (9) as consistency function. When t > η , for any y ∈ R d , there always exists a ˆ ϵ ϕ ( x t , µ, t ) such that f ϕ ( x t , t ) = y .

Proof. Recall (9), there is

· · ·

HQ image

𝒙𝒙

(

𝟎𝟎

)

𝒙𝒙

(

𝑻𝑻

)

𝒅𝒅

𝒙𝒙

=

𝜽𝜽

𝒕𝒕

𝝁𝝁 - 𝒙𝒙

𝒅𝒅𝒕𝒕

+

𝝈𝝈

𝒕𝒕

𝒅𝒅𝒅𝒅

Image noising process

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝒕𝒕

𝟏𝟏

,

𝒕𝒕

𝟏𝟏

)

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝒕𝒕

𝟐𝟐

,

𝒕𝒕

𝟐𝟐

)

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝑻𝑻

,

𝑻𝑻

)

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝟎𝟎

,

𝟎𝟎

)

· · ·

𝒙𝒙

(

𝟎𝟎

)

𝒙𝒙

(

𝑻𝑻

)

𝒅𝒅

𝒙𝒙

=

𝜽𝜽

𝒕𝒕

𝝁𝝁 - 𝒙𝒙

𝒅𝒅𝒕𝒕

+

𝝈𝝈

𝒕𝒕

𝒅𝒅𝒅𝒅

Image degradation process

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝒕𝒕

𝟐𝟐

,

𝒕𝒕

𝟐𝟐

)

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝒕𝒕

𝟑𝟑

,

𝒕𝒕

𝟑𝟑

)

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝑻𝑻

,

𝑻𝑻

)

𝒇𝒇

𝝓𝝓

(

𝒙𝒙

𝒕𝒕

𝟏𝟏

,

𝒕𝒕

𝟏𝟏

)

Linear Fitting Stage:

𝒙𝒙

𝟎𝟎

: =

𝝁𝝁

Nonlinear Fitting Stage:

𝒙𝒙

𝟎𝟎

: =

HQ image

≈

+

𝝁𝝁

𝓝𝓝

(

𝟎𝟎

,

𝝀𝝀

𝟐𝟐

)

≈

+

𝝁𝝁

𝓝𝓝

(

𝟎𝟎

,

𝝀𝝀

𝟐𝟐

)

LQ image

noise

LQ image

noise

𝒇𝒇

𝝓𝝓

𝒙𝒙

𝒕𝒕

,

𝒕𝒕

=

𝒄𝒄

𝒔𝒔𝒔𝒔𝒔𝒔𝒔𝒔

(

𝒕𝒕

)[

𝝁𝝁

+

𝒗𝒗

𝒕𝒕

𝜺𝜺

] +

𝒄𝒄

𝒙𝒙

𝟎𝟎

(

𝒙𝒙

,

𝒕𝒕

;

𝝓𝝓

)

𝒇𝒇

𝝓𝝓

𝒙𝒙

𝒕𝒕

,

𝒕𝒕

=

𝒄𝒄

𝒔𝒔𝒔𝒔𝒔𝒔𝒔𝒔

𝒕𝒕

[

𝝁𝝁

+(

𝒙𝒙

𝟎𝟎

-𝝁𝝁

)

𝒆𝒆

-� 𝜽𝜽

𝒕𝒕

+

𝒗𝒗

𝒕𝒕

𝜺𝜺

] +

𝒄𝒄

𝒙𝒙

𝟎𝟎

(

𝒙𝒙

,

𝒕𝒕

;

𝝓𝝓

)

then substitute (8) into (12), we obtain:

Note that when t > η , both v t and ˆ θ t are greater than 0, so the right-hand side of equation (13) is nonsingular. This implies that under loss functions (6) and (7), the consistency function admits arbitrary non-singular solutions, leading to pattern collapse. In other hand, (6) and (7) only emphasize the self-consistency between any two points on the PF-ODE trajectory, leading to a lack of determinism in the training process. This uncertainty indirectly contributes to slower convergence of the model.

To address this issue, we additionally introduce the following origin-guided(OG) loss to stabilize the training process.

where λ perc is a positive constant, and Φ( · ) represents a VGG16[38] feature extractor from 2nd and 5th pooling layers. This is equivalent to performing a consistency computation with the origin after each random sampling, hence referred to as the origin-guided loss. This effectively avoids the emergence of mediocre solutions and adds some certain determinism to the training process, resulting in faster convergence. Then the final loss function is represented as follows:

where λ OG is a positive constant. The selection of λ OG will be discussed in the ablation experiment section of Chapter 4.

3.3 Linear-nonlinear decoupling training strategy

With the foundation laid in Sections 3.1 and 3.2, we can naturally train IR-CM by a regular CD or CT process, as illustrated in the lower half of Figure 1. Despite this method offers convenience in training IR-CM, we have empirically found that its performance is not optimal in practice. Let us recall the SDE (2) and OECF (9), we obtain:

Algorithm 1: Linear-fitting stage

Algorithm 2: NonLinear-fitting stage

In general training, the model attempts to simultaneously fit the variations of both linear part f 1 and nonlinear part f 2 . Obviously, this is more challenging than fitting f 1 alone. In practice, the model's performance at larger values of t is more crucial during training, because any intermediate state for x ( t ) , ∀ t ∈ [ η, T ) are unknown during inference, thus we can only set t = T for inference. And note that there is lim t → T f 2 ≈ 0 , therefore, the influence of f 2 is negligible when t = T . Based on the analysis above, we can then eliminate f 2 for all t by simply setting x (0) as µ . By doing so, the original image degradation process is transformed into a purely image noising process (as shown in the upper half of Figure 1). Since only the linear part f 1 and noise f 3 are being fitted, the model will achieve better performance over the entire PF-ODE trajectory, leading to improved inference performance as well. We refer to this training as the linear-fitting stage and the corresponding pseudocode is shown in Algorithm 1.

Despite the Linear-fitting stage achieves good performance for one-step inference after training, we can still further improve the performance by employing appropriate multi-step sampling. However, the model trained in the Linear-fitting stage cannot perform multi-step sampling inference because the outputed HQ image is not the origin x (0) = µ of PF-ODE trajectory, making it unable to estimate any intermediate states x t after one inference. Therefore, after the Linear-fitting stage, we set x (0) to be the HQ image instead of µ and fine-tune the model (Algorithm 2). Since in the previous stage, the model has already fitted f 1 and f 3 , this stage mainly focuses on fitting the non-linear part f 2 . Therefore, this stage is naturally referred to as the non-Linear-fitting stage. After training in the non-linear stage, the model's performance at t = T is almost unchanged, meaning there is no change in the performance of one-step inference. But the model will be able to perform multi-step sampling inference (Algorithm 3 in appendix C). In practice, even just two-step sampling inference brings a noticeable performance improvement. Relevant comparative experiments and ablative experiments will be presented in Chapter 4.

4 Experiments

4.1 Comparative Experiment

In this section, we will validate the effectiveness of our proposed method on five tasks: image de-raining, image denoising, image deblurring, low-light image enhancement, and nighttime glare removal. The implementation details are described in Appendix A. Specifically, we will compare our method with some of the state-of-the-art methods on PSNR, SSIM[39], LPIPS[40] metrics and NFE. The NFE (Number of Function Evaluations) refers to the number of function evaluations required to generate an image or data. In other words, it is the number of evaluations needed at each step of the diffusion process. Notably, like other SDE-based models, we prioritize perceptual scores LPIPS over distortion scores PSNR and SSIM. And our metric settings are same as other mentioned baseline methods. For PSNR metric, we perform the calculation in the luminance space (Y channel). For SSIM metric, it refer to [39], and for LPIPS metric, it refer to [40]. All comparison experiments were

input

JORDER

MAXIM

PReNet

Restormer

IR-SDE

Ours

GT

input

Transweather

Ours

GT

MAXIM

conducted using the original resolution of each dataset (for ease of presentation, the image sizes were adjusted in Figures 2-5).

4.1.1 Image deraining

We validate the effectiveness of the proposed IR-CM based on two datasets: R100L and R100H[41]. A total of 2000 images are used for training, while 200 images are reserved for testing. Our method is qualitatively and quantitatively compared with some of the state-of-the-art image deraining methods including IR-SDE[22], JORDER[3], Restormer[5], PReNet[1], and MAXIM[4]. The comparison results are shown in Table 1, Table 2 and Figure 2. More visual results are available in appendix D. To further validate the effectiveness of our method in real-world rainy scenarios, we conducted comparative experiments on the Raindrop[42] dataset containing 1119 pairs of real-world rainy/nonrainy images. The results are shown in Table 3 and Figure 3.

The IR-CM model we proposed surpasses the baseline IR-SDE model and achieves optimal performance on both SSIM and LPIPS metrics in scenarios of either one-step or two-step reasoning. Furthermore, the comparison results with the CNN-baseline demonstrate that our approach improves network performance while ensuring real-time capabilities. Note that our IR-CM is based on IR-SDE for CD or CT, thus IR-CM-CD represents consistency distillation based on pre-trained IR-SDE model as teacher model. And IR-CM-CT only initializes the model using pre-trained IR-SDE checkpoint of faster convergence and does not use any teacher model during the CT process. Of course, random

weight initialization is also an option. In practice, we found that when given enough training time, the performance of both approaches is quite similar.

4.1.2 Image deblurring

We validated the effectiveness of the proposed IR-CM model for image deblurring task based on the GoPro[43] dataset. A total of 2103 images are used for training, while 1111 images are reserved for testing. Our method is qualitatively and quantitatively compared with some of the milestone image deblurring methods including DeepDeblur[43], DeblurGAN-v2[14], DBGAN[15], MAXIM[4], DiffIR[44] and of course IR-SDE[22]. The comparison results are shown in Table 4 and Figure 4. More visual results are available in appendix D.

Our proposed method achieves optimal performance on SSIM and LPIPS metrics, surpassing the IRSDE baseline model as well. This strongly demonstrates the effectiveness of our proposed consistency training approach. Notably, like other SDE-based models, we prioritize perceptual scores LPIPS over distortion scores PSNR and SSIM.

4.1.3 Image denoising

Note that the last term in (1) is a Gaussian process. Hence, we can consider a special case where µ = x (0) . In this case, the IR-SDE degenerates into a pure additive noise process. This implies that any point along the PF-ODE trajectory can serve as a noisy image to be processed. In this scenario, we can only apply the conventional consistency training approach and cannot utilize the proposed linear-nonlinear decoupled training strategy because x (0) must be set as the HQ image and x ( t ) represents the corresponding low-quality (LQ) image.

Similar to [22], we collected approximately 5000 HQ images from the DIV2K[45], Flickr2K[45], and BSD500[46] datasets for training, and subsequently tested on the McMaster[47] dataset. To demonstrate the competitiveness of our method against state-of-the-art approaches, we compare it with DnCNN[7], FFDNet[8], as well as the special denoising methods proposed in [22], namely Denoising-SDE and Denoising-ODE. Our method achieves the optimal result in perceptual scores. Comparative experimental results are shown in Table 5, and visual results are available in appendix D.

input

DeepDeblur

DBGAN

DeblurGANv2

MAXIM

IR-SDE

Ours

GT

input

DiffLL

GlobalDiff

Retinexformer

LLFlow

IR-SDE

Ours

GT

4.1.4 Low-light image enhancement

We compare our method with existing researches for low-light image enhancement tasks on the LOLv2[48] dataset consisted of 689 pairs training images and 100 pairs testing images. We compare our proposed method with some recent approaches for low-light image enhancement task, including: DiffLL[21], GlobalDiff[20], Retinexformer[16], LLFlow[17]. Note that IR-SDE is trained following the experimental settings described in [22], while IR-CM-CD is based on consistency distillation using the trained weights of IR-SDE. Our method outperforms the state-of-the-art diffusion-based methods, DiffLL and GlobalDiff, and requires only two sampling steps. The corresponding quantitative and qualitative comparison results are shown in Table 6 and Figure 5 respectively. More visual results are available in appendix D

4.1.5 Runtime Comparation

To validate the superiority of our method in inference real-time performance, we compared it with several other state-of-the-art (SOTA) methods across three common resolution sizes. The results are shown in Table8. Compared to the baseline method IR-SDE, our method significantly reduces inference time by introducing a consistency training process. In some real-time demanding applications, one-step sampling inference allows for fast predictions while maintaining competitive performance. On the other hand, two-step sampling inference can substantially improve model

performance with only a slight trade-off in inference speed. Users can choose the approach based on their specific requirements. Additionally, a discussion on the model complexity will be provided in Appendix E.

4.2 Ablation Experiments

4.2.1 Components ablation experiments

The main innovations of this paper lie in the OECF, Origin-guided loss, and the linear-nonlinear decoupling training strategy proposed in the previous section. To verify the effectiveness of each component, we conducted a series of ablation experiments using low-light image enhancement as an example. To ensure fairness, each control group used the same checkpoint for model initialization, and the total number of epochs was kept consistent during training. And only the performance of one-step sampling inference was considered. The corresponding results are shown in Table 7. The effectiveness of the proposed OECF can be demonstrated by group pairs (1, 3), (2, 4), (5, 7), (6, 8). The significant improvement in perceptual scores due to the origin-guided loss is evident from the comparisons between group pairs (1, 2), (3, 4), (5, 6), (7, 8). Finally, the notable performance enhancement of the model due to the linear-nonlinear decoupling training strategy can be seen from the comparisons between groups (1, 5), (2, 6), (3, 7), (4, 8). Note that for groups 5 and 7, we only use the origin-guided loss during the Linear-fitting stage, and not during the nonlinear-fitting stage.

4.2.2 Selection of origin-guided loss weight

The setting of λ OG also affects the performance of IR-CM. If set too low, the model may become unstable and perform poorly, as discussed in Section 3.2. If set too high, it will affect self-consistency property, leading to a degradation into a CNN-baseline method. We experimented with multiple values and empirically found that the best performance is achieved around 0.8. The related experimental results are discussed in Appendix B.

4.2.3 Linear-fitting stage only & Multiple sampling

In fact, after completing the linear-fitting stage, the model can already perform one-step sampling inference, however, it cannot perform multi-step sampling as discussed in Section 3.3. To further evaluate the model's performance, we tested it under the conditions of linear-fitting stage only, and with 1-step, 2-step, 4-step, and 6-step sampling. Empirically, we consider 2-step sampling to be the most cost-effective choice. The related experimental results are discussed in Appendix C.

5 Conclusion

This paper proposes a multi-task image restoration and enhancement method, IR-CM, based on a consistency training approach, enabling few-step or even one-step sampling inference. Specifically, we proposed the Origin-estimated Consistency Function (OECF), which provides a more stable initial state and a smaller solution space for the consistency training process. Furthermore, to make the training process more robust and prevent trivial solutions, we introduced the Origin-guided Loss (OE Loss). Based on these, we developed a Linear-Nonlinear Decoupling Training Strategy, which not only accelerates the training process but, more importantly, enables the model to perform multi-step sampling inference, further enhancing its performance. Finally, a series of comparative experiments and ablation studies demonstrated the effectiveness of the proposed method.

6 Acknowledgment

Our work is supported by the weapons and equipment pre-research fund, Grant No. 50911020603.

References

Dongwei Ren, Wangmeng Zuo, Qinghua Hu, Pengfei Zhu, and Deyu Meng. Progressive image deraining networks: A better and simpler baseline. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 3937-3946, 2019.

Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, MingHsuan Yang, and Ling Shao. Multi-stage progressive image restoration. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 14821-14831, 2021.

Wenhan Yang, Robby T Tan, Jiashi Feng, Zongming Guo, Shuicheng Yan, and Jiaying Liu. Joint rain detection and removal from a single image with contextualized deep networks. IEEE transactions on pattern analysis and machine intelligence , 42(6):1377-1393, 2019.

Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, and Yinxiao Li. Maxim: Multi-axis mlp for image processing. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 5769-5780, 2022.

Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang. Restormer: Efficient transformer for high-resolution image restoration. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 5728-5739, 2022.

Yeachan Park, Myeongho Jeon, Junho Lee, and Myungjoo Kang. Mcw-net: Single image deraining with multi-level connections and wide regional non-local blocks. Signal Processing: Image Communication , 105:116701, 2022.

Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising. IEEE transactions on image processing , 26(7):3142-3155, 2017.

Kai Zhang, Wangmeng Zuo, and Lei Zhang. Ffdnet: Toward a fast and flexible solution for cnn-based image denoising. IEEE Transactions on Image Processing , 27(9):4608-4622, 2018.

Rongkai Zhang, Jiang Zhu, Zhiyuan Zha, Justin Dauwels, and Bihan Wen. R3l: Connecting deep reinforcement learning to recurrent neural networks for image denoising via residual recovery. In 2021 IEEE International Conference on Image Processing (ICIP) , pages 1624-1628. IEEE, 2021.

Shen Cheng, Yuzhi Wang, Haibin Huang, Donghao Liu, Haoqiang Fan, and Shuaicheng Liu. Nbnet: Noise basis learning for image denoising with subspace projection. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4896-4906, 2021.

Meng Chang, Qi Li, Huajun Feng, and Zhihai Xu. Spatial-adaptive network for single image denoising. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXX 16 , pages 171-187. Springer, 2020.

Seungjun Nah, Tae Hyun Kim, and Kyoung Mu Lee. Deep multi-scale convolutional neural network for dynamic scene deblurring. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 3883-3891, 2017.

Orest Kupyn, Volodymyr Budzan, Mykola Mykhailych, Dmytro Mishkin, and Jiˇ rí Matas. Deblurgan: Blind motion deblurring using conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 8183-8192, 2018.

Orest Kupyn, Tetiana Martyniuk, Junru Wu, and Zhangyang Wang. Deblurgan-v2: Deblurring (orders-of-magnitude) faster and better. In Proceedings of the IEEE/CVF international conference on computer vision , pages 8878-8887, 2019.

Kaihao Zhang, Wenhan Luo, Yiran Zhong, Lin Ma, Bjorn Stenger, Wei Liu, and Hongdong Li. Deblurring by realistic blurring. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 2737-2746, 2020.

Yuanhao Cai, Hao Bian, Jing Lin, Haoqian Wang, Radu Timofte, and Yulun Zhang. Retinexformer: One-stage retinex-based transformer for low-light image enhancement. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , pages 12504-12513, October 2023.

Yufei Wang, Renjie Wan, Wenhan Yang, Haoliang Li, Lap-Pui Chau, and Alex C Kot. Low-light image enhancement with normalizing flow. arXiv preprint arXiv:2109.05923 , 2021.

Alexandru Brateanu, Raul Balmez, Adrian Avram, and Ciprian Orhei. Lyt-net: Lightweight yuv transformer-based network for low-light image enhancement. arXiv preprint arXiv:2401.15204 , 2024.

Wu Yuhui, Pan Chen, Wang Guoqing, Yang Yang, Wei Jiwei, Li Chongyi, and Heng Tao Shen. Learning semantic-aware knowledge guidance for low-light image enhancement. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2023.

Hou Jinhui, Zhu Zhiyu, Liu Hui, Zeng Huanqiang, and Yuan Hui. Global structure-aware diffusion process for low-light image enhancement. Advances in Neural Information Processing Systems , 2023.

Hai Jiang, Ao Luo, Haoqiang Fan, Songchen Han, and Shuaicheng Liu. Low-light image enhancement with wavelet-based diffusion models. ACM Transactions on Graphics (TOG) , 42(6):1-14, 2023.

Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Image restoration with mean-reverting stochastic differential equations. International Conference on Machine Learning , 2023.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems , 34:8780-8794, 2021.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv:2010.02502 , October 2020.

Alexia Jolicoeur-Martineau, Ke Li, Rémi Piché-Taillefer, Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2105.14080 , 2021.

Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE transactions on pattern analysis and machine intelligence , 45(4):4713-4726, 2022.

Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 11461-11471, 2022.

Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research , 23(47):1-33, 2022.

Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Refusion: Enabling large-size realistic image restoration with latent-space diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops , pages 1680-1691, 2023.

Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Controlling vision-language models for universal image restoration. arXiv preprint arXiv:2310.01018 , 2023.

Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Photorealistic image restoration in the wild with controlled vision-language models. arXiv preprint arXiv:2404.09732 , 2024.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems , 35:5775-5787, 2022.

Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Dpm-solver-v3: Improved diffusion ode solver with empirical model statistics. Advances in Neural Information Processing Systems , 36, 2024.

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 , 2023.

Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556 , 2014.

Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing , 13(4):600612, 2004.

Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 586-595, 2018.

Wenhan Yang, Robby T Tan, Jiashi Feng, Jiaying Liu, Zongming Guo, and Shuicheng Yan. Deep joint rain detection and removal from a single image. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1357-1366, 2017.

Rui Qian, Robby T Tan, Wenhan Yang, Jiajun Su, and Jiaying Liu. Attentive generative adversarial network for raindrop removal from a single image. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2482-2491, 2018.

Tae Hyun Kim, Seungjun Nah, and Kyoung Mu Lee. Deep multi-scale convolutional neural network for dynamic scene deblurring. In Conference on Computer Vision and Pattern Recognition , pages 1-21. IEEE, 2017.

Bin Xia, Yulun Zhang, Shiyin Wang, Yitong Wang, Xinglong Wu, Yapeng Tian, Wenming Yang, and Luc Van Gool. Diffir: Efficient diffusion model for image restoration. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 13095-13105, 2023.

Eirikur Agustsson and Radu Timofte. Ntire 2017 challenge on single image super-resolution: Dataset and study. In Proceedings of the IEEE conference on computer vision and pattern recognition workshops , pages 126-135, 2017.

Pablo Arbelaez, Michael Maire, Charless Fowlkes, and Jitendra Malik. Contour detection and hierarchical image segmentation. IEEE transactions on pattern analysis and machine intelligence , 33(5):898-916, 2010.

Lei Zhang, Xiaolin Wu, Antoni Buades, and Xin Li. Color demosaicking by local directional interpolation and nonlocal adaptive thresholding. Journal of Electronic imaging , 20(2):023016023016, 2011.

Wenhan Yang, Wenjing Wang, Haofeng Huang, Shiqi Wang, and Jiaying Liu. Sparse gradient regularized deep retinex network for robust low-light image enhancement. IEEE Transactions on Image Processing , 30:2072-2086, 2021.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

A Implementation details

We chose the same U-Net backbone network as used in IR-SDE[22]. Note that the CNN-baseline uses the same U-Net backbone, with its time embedding set to a constant, and it directly outputs HQ images from the input LQ images. For all tasks, we set the training patch-size to be 128 × 128 and use a batch size of 4. We used the AdamW[49] optimizer with a weight decay set to 0.01 and set learning rate to 10 -5 . Our models are trained on two 2080Ti GPUs for 200 epochs for each task, with the linear-fitting stage and the nonlinear-fitting stage each accounting for 100 epochs. The training hyperparameters are set as follows: λ = 10 , λ OG = 0 . 8 , λ perc = 0 . 125 .

The θ schedule is defined same as [22]:

where s = 0 . 008 . And the consistency weights c skip and c out are set as follow:

where T is set to 100.

B Selection of origin-guided loss weight

We take the low-light image enhancement task as an example to consider the impact of different λ OG values on the final performance of the model considering only the one-step sampling scenario. The quantitative experimental results are shown in Table 9.

The results show that the best performance is achieved when λ OG is set to 0.8. Although a higher PSNR score was obtained with λ OG = 1 . 0 , we prioritize perceptual scores. Therefore, we chose λ OG = 0 . 8 as the experimental setting in Section 4.1.

C linear-fitting stage only & Multiple sampling

As shown in Algorithm 3, our method, like most SDE-based methods, can improve model performance through multi-step sampling. Specifically, if set M = 1 , one-step sampling can be achieved. To further evaluate the model's performance, we tested it under 1-step, 2-step, 3-step, and 4-step sampling. Additionally, we evaluated the model's performance after completing only the linear-fitting stage (LFS). Similarly, using low-light enhancement as an example, the test results are shown in Table 10 and Figure 6.

It can be observed that the optimal perceptual score is achieved when the sampling step is 2. Further increasing the sampling steps leads to only slight improvements in model performance but implies longer inference times. In the case of one-step sampling, the performance of LFS only is not significantly different from that of the two-stage trained model. Considering all factors, we believe that 2-step sampling inference is the most cost-effective.

Algorithm 3:

input

1-step

GT

2-step

3-step

4-step

D More visual results

input

JORDER

MAXIM

PReNet

Restormer

IR-SDE

Ours

GT

input

JORDER

MAXIM

PReNet

Restormer

IR-SDE

Ours

GT

input

DeepDeblur

DBGAN

DeblurGANv2

MAXIM

IR-SDE

Ours

GT

input

IR-SDE

ours

GT

input

IR-SDE

Ours

GT

input

DiffLL

GlobalDiff

Retinexformer

LLFlow

IR-SDE

Ours

GT

090000

中员享受

会员享受

公员享受

会员享受

会员享受

会员享受

快员享受

台拆优惠

扬优惠

E Model complexity

Our backbone network is a conditional Unet, and its computational complexity is proportional to the number of convolutional layers and the size of the image. The complexity can be calculated as follows:

where L is number of convolutional layers, H,W are the height and width of the image, C in,i , C out,i represent the input and output sizes of the convolutional layer respectively, K is the size of the convolutional kernel, and S denotes the sample steps. From the above formula, it can be seen that our model has a linear relationship with the height and width of the input image. And it is also linearly related to the number of sampling steps. This is advantageous for applying our model to larger-sized images. Additionally, by adjusting the number of sampling steps, a trade-off between real-time performance and model performance can be achieved.

F Potential societal impact

Our work on conditional image generation tasks, such as image restoration and enhancement, has several potential societal impacts. These improvements can benefit fields like medical imaging, where enhanced image quality can aid in better diagnosis and treatment. Additionally, improved image restoration techniques can be valuable in preserving and restoring historical photographs and artworks.

However, there are also potential negative impacts to consider. Enhanced image generation techniques could be misused for creating deceptive content, such as deepfakes, which can have serious ethical and social implications. Therefore, it is crucial to implement safeguards and ethical guidelines to prevent misuse and ensure that the technology is used for beneficial purposes.

We are committed to promoting the positive applications of our research while being aware of and mitigating potential risks.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We accurately reflect the contributions and scope of the paper in the abstract and introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We emphasized in the introduction and abstract that our work is applicable to conditional image generation tasks, such as image restoration and image enhancement. It is not suitable for unconditional generation tasks.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: In the Chapter 3, we introduced and demonstrated the effectiveness of the three main contributions of this paper.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In Appendix A, we elaborate on the experimental details. Additionally, our code will be open-sourced after the review process.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We are pleased to open-source our code, which will be made available on https://github.com .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: In Chapter 4 and Appendix A, we provide detailed explanations of our experimental procedures, including dataset splits and experimental settings.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Our experiments primarily compared the performance of various existing methods on specific public datasets. Once the model is trained, these metrics (PSNR, SSIM, LPIPS) are typically deterministic and do not change due to other random factors. Therefore, it is unnecessary to report error bars in this paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In Appendix A, we provide detailed information about the GPU models and memory sizes used in our experiments, as well as the number of training iterations required.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have carefully read the NeurIPS Code of Ethics and ensure strict adherence to it.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the potential negative societal impacts of our method in Appendix E.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our proposed method poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have provided clear citations for the public datasets used in our comparative experiments, as well as for foundational works in the field.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: Due to submission size constraints, we only provide our code in the supplementary material. The corresponding datasets, test results, and pretrained models will be open-sourced in the future.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourced experiments or research involving human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourced experiments or research involving human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
This paper presents IR-CM, a fast and universal image restoration method leveraging consistency models. The key innovations include a novel linear-nonlinear decoupling training strategy and an origin-estimated consistency function to enhance training effectiveness and inference performance. The proposed method is evaluated across several image restoration tasks, including draining, deblurring, denoising, and low-light image enhancement, showing competitive results with minimal inference steps.

STRENGTHS:
1. The proposed original-estimated consistency function (OECF) provides a more stable initial state and reduces the solution space.

2. The origin-guided loss stabilizes training and prevents pattern collapse.

3. Extensive experiments on four restoration tasks are conducted to validate the effectiveness of the proposed method.

WEAKNESSES:
1. There is an error in Eq. (8), $v_t$ should be $\sqrt{v_t}$. Please have a check.

2. There is an inconsistency in Theorem 1, ``$c_{out}(\eta)=0, c_{out}(\eta)=1$''. Please have a check. I suggest the author carefully check the writing of this paper and avoid such errors or typos.

3. Even though this work contains rich experiments on multiple tasks, more recent diffusion-based methods should be considered as a comparison, such as DiffIR.

4. All the experiments for draining and denoising are based on the synthetic training and testing data. To better validate the practical value in real applications, more real-world datasets should be considered.

5. It would be more comprehensive to provide a comparison on model complexity and runtime.

QUESTIONS:
See weakness.

LIMITATIONS:
It would be better to analyze the limitations or some failed cases of the proposed methods.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper modifies the consistency model from image generation to image restoration with three modifications. Each module contains robust theoretical proof and shows great performance improvement.  Linear-nonlinear decoupling training strategy will motivate future works in using consistency model in image restoration. The additional results show the power of the paper both in inference speed and performance.\n\nSTRENGTHS:\n1. The modification from Consistency Model in image generation to image restoration is terrific. Each module shows robust theoretical proof 2. The extensive experiments show the great power of IR-CM in inference speed and performance.\n\nWEAKNESSES:\n1. The comparison with other methods seems to be unfair. For example, the result of retinexformer in low-light enhancement is tested on full resolution with y-channel calculation. It seems that all of the results of this paper are tested in 512*512 resolution, which is observed by the unexpectedly high SSIM metrics (Table 4 with comparable SSIM but significantly lower PSNR, Table 5 with over 30 PSNR and 0.95 SSIM, etc.). In my view, it is impossible. $\\textbf{Please show the result in full resolution}$ (i.e., results in 512 resolution will change the result largely and it is meaningless at all) with uniform metrics (like y-channel calculation) but not directly report the result of other papers.\n2. The phenomenon of first-stage training is interesting, as to why setting the $x_0$ equal to $\\mu$ will benefit learning image restoration tasks as it is just a denoising task. I understand that this will ease training, but that's still not enough to explain the phenomenon. \n3. Mistaking the meaning of the setting. Universal image restoration means using one model to train different degradation types at once, not meaning a framework that can be used to train different tasks (the author can freely search in the arxiv).\n4. Typos. Line 143 should change to \"the training process\" as in the inference process neither strategy uses pre-trained model. Line 152 the $c_{out}(T)=1$, not $\\eta$\n\nQUESTIONS:\nSee the weaknesses.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a method called IR-CM (Image Restoration Consistency Model) for fast and general-purpose image restoration.\nThe key idea is to achieve few-step or one-step inference by employing consistency training on specific mean-reverting stochastic differential equations (SDEs).\nA novel linear-nonlinear decoupling training strategy is proposed to enhance training effectiveness and inference performance without relying on pre-trained models.\nAn origin-guided loss is introduced to avoid trivial solutions and stabilize model training.\nExperiments conducted on tasks such as image deraining, denoising, deblurring, and low-light image enhancement demonstrate highly competitive results with one-step inference and state-of-the-art performance in low-light image enhancement with two-step inference.\n\nSTRENGTHS:\n- The paper is well-structured and easy to follow.\n- The method is designed for fast inference, which is crucial for real-time applications.\n- It is a general-purpose model that can be applied to various image restoration tasks without domain-specific prior knowledge.\n- The proposed linear-nonlinear decoupling strategy and origin-estimated consistency function enhance training and inference performance.\n- IR-CM does not depend on any pre-trained checkpoints, making it an independent model.\n- The method achieves competitive or even state-of-the-art results in various image restoration tasks with minimal inference steps.\n\nWEAKNESSES:\n1. While the method is designed to be universal, its performance across different types of image degradations and real-world scenarios might vary.\n2. The introduction of a novel training strategy and loss function may complicate the training process compared to more straightforward models.\n3. The method's performance may be sensitive to the choice of hyperparameters, such as the weight of the origin-guided loss.\n4. Without sufficient data or proper regularization, the model might overfit to the training data, which could reduce its effectiveness on unseen data.\n\nQUESTIONS:\nBesides the above Weaknesses, I still have some questions:\n\nThe one-step inference model cannot perform multi-step sampling, which might limit its application in certain scenarios where more refinement is needed.\n\nHow does the method handle different levels of image degradation beyond the tested conditions?\n\nWhat is the computational complexity of the method, and how does it scale with the size of the image or video?\n\nCan the method be extended to handle video restoration tasks, and if so, what would be the challenges?\n\nHow does the method compare to other state-of-the-art methods in terms of computational resources and efficiency?\n\nAre there any specific use cases or applications where the method excels or falls short?\n\nLIMITATIONS:\nThe societal impacts are discussed in Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper proposes a SDE-based two-stage network to tackle the image restoration tasks, including deraining, deblurring denoising, and low-light image enhancement. Based on the existing stochastic differential equation works, the proposed method focuses on the training efficiency and inference speed. It proposes a linear and non-linear decoupling training strategy to enhance the training effectiveness and to surpass consistency distillation on inference. Extensive experiments on multiple different image degradation data show that 
the proposed method is able to outperforms baseline methods.

STRENGTHS:
Strength
1. the proposed method is a fast universal method that is able to solve multiple image degradation problems including deraining, denoising, low-light enhancement etc.  
2. Extensive experiments demonstrate the proposed method outperforms the baseline dedicated methods on multiple benchmarking datasets
3. The proposed training schema may bring impact and benefits to the development of SDE study.

WEAKNESSES:
Weakness
1. What does NFE in the quantitative comparison tables (Table 1,2,3,4) stand for? 
2. As the paper is claiming the proposed method is a fast and universal method at inference time, it is suggested to add the inferencing time comparison with respect to baseline methods. 
3. There is no conclusion in the submitted paper.
4. All the experiments on deraining topic use synthetic rain images. It is more useful and applicable to compare methods on real rain images. 

Minors
1. All the functions are lack of punctuations.

QUESTIONS:
The authors are suggested to address the issues raised in the weakness section during rebuttal period.

LIMITATIONS:
There is no limitation mentioned in the paper or the supplementary material.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-14