Image Restoration with Mean-Reverting Stochastic Differential Equations

Abstract: This paper presents a stochastic differential equation (SDE) approach for general-purpose image restoration. The key construction consists in a mean-reverting SDE that transforms a high-quality image into a degraded counterpart as a mean state with fixed Gaussian noise. Then, by simulating the corresponding reverse-time SDE, we are able to restore the origin of the low-quality image without relying on any task-specific prior knowledge. Crucially, the proposed mean-reverting SDE has a closed-form solution, allowing us to compute the ground truth time-dependent score and learn it with a neural network. Moreover, we propose a maximum likelihood objective to learn an optimal reverse trajectory that stabilizes the training and improves the restoration results. The experiments show that our proposed method achieves highly competitive performance in quantitative comparisons on image deraining, deblurring, and denoising, setting a new state-of-the-art on two deraining datasets. Finally, the general applicability of our approach is further demonstrated via qualitative results on image super-resolution, inpainting, and dehazing. Code is available at https://github.com/Algolzw/image-restoration-sde.

• The paper presents a novel SDE-based approach that degrades images with fixed Gaussian noise and reverses the process to restore them.
• It employs a maximum likelihood training objective that stabilizes network learning and enhances restoration quality over traditional methods.
• Results across tasks like deraining, deblurring, and denoising demonstrate state-of-the-art performance measured by PSNR, SSIM, and other metrics.

Image Restoration with Mean-Reverting Stochastic Differential Equations

The paper presents a novel approach to image restoration utilizing mean-reverting stochastic differential equations (SDEs), known as the Image Restoration SDE (IR-SDE). This approach is characterized by a mean-reverting SDE that degrades high-quality images by adding fixed Gaussian noise and reverses this process to restore images without requiring task-specific prior knowledge. The technique leverages a closed-form solution, enabling the computation of the true time-dependent score for training a neural network, thus allowing effective restoration across diverse applications.

Key Contributions

1. Mean-Reverting SDE for Image Degradation and Restoration: The authors present an SDE-based approach where images undergo degradation using a mean-reverting process that adds fixed Gaussian noise. The reversal of the SDE restores high-quality images by reverse-simulating the degradation process, sidestepping the need for detailed degradation information.
2. Maximum Likelihood Objective: A maximum likelihood-based objective is proposed for training the neural network, which stabilizes training and enhances restoration quality compared to conventional score matching techniques.
3. Versatility Across Restoration Tasks: The paper demonstrates the method's broad applicability by deploying it on tasks such as image deraining, deblurring, denoising, super-resolution, inpainting, and dehazing, achieving competitive quantitative performance.

Results and Evaluation

Experiments show that IR-SDE achieves state-of-the-art results in quantitative metrics such as PSNR, SSIM, LPIPS, and FID across several image restoration tasks:

• Image Deraining: The IR-SDE outperformed other deraining methods, setting new benchmarks on the Rain100H and Rain100L datasets, offering superior perceptual quality as measured by LPIPS and FID scores.
• Image Deblurring: Achievements in deblurring tasks on the GoPro dataset reveal IR-SDE's capability to produce sharp and perceptually pleasing images, surpassing GAN-based alternatives in fidelity and perceptual quality.
• Image Denoising: A special case of IR-SDE was tested for denoising and successfully improved efficiency using fewer steps. A deterministic ODE variant showed even better fidelity scores, overcoming challenges faced by stochastic processes in Gaussian denoising tasks.

Discussion

IR-SDE provides an innovative framework in image restoration, proposing a unique method for modeling degradations via SDEs and reversing them efficiently. The introduction of a maximum likelihood training objective highlights a robust means of stabilizing network training while achieving superior restoration quality. Moreover, the accommodation of various time-varying schedules for $\theta$ in their SDE model indicates the adaptability and potential improvements available with this approach.

Future Directions

While IR-SDE demonstrates impressive performance, challenges remain, particularly concerning the smooth variance change in the last steps, which may hinder learning efficiency. Further exploration into $\theta$ scheduling and sampling techniques is necessary to enhance computational efficiency during inference. Additionally, broadening the scope of applicable SDE choices could provide even more comprehensive solutions for complex image restoration tasks.

This method's flexibility suggests promising advances in developing new stochastic models for wide-ranging restoration challenges, pointing toward significant future contributions in AI-driven image processing.