PaperlayerView exact revision on arXiv

Source-linked AI summary

Image Restoration with Mean-Reverting Stochastic Differential Equations

Ziwei Luo, Fredrik K. Gustafsson, Zheng Zhao, Jens Sjölund, Thomas B. Schön

arXiv:2301.11699v3cs.LGcs.CV

TL;DR

Image restoration methods based on standard diffusion processes can start reverse generation from high-variance noise, making recovery of the original high-quality image difficult. This paper instead uses a mean-reverting SDE that models degradation directly, learns its reverse dynamics with score and maximum-likelihood objectives, and achieves highly competitive results across six restoration tasks, including state-of-the-art performance on two deraining datasets.

  • Problem

    Standard diffusion restoration methods can initialize reverse processes with high-variance noise, which may produce poor restoration of the ground-truth high-quality image.

  • Method

    The method uses a mean-reverting SDE to model image degradation directly, computes a ground-truth time-dependent score from its closed-form solution, and trains with score-based and maximum-likelihood objectives.

  • Results

    The approach achieves highly competitive performance on deraining, deblurring, and denoising, setting a new state-of-the-art on two deraining datasets.

  • Takeaways & Limitations

    Experiments across six diverse restoration tasks demonstrate the method’s wide applicability and highly competitive restoration performance.

  • Takeaways & Limitations

    The exponential variance schedule can make neighboring late-stage states too similar for effective learning, especially with maximum-likelihood training.

Abstract

from arXiv · show

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.

1. Introduction

The paper adapts SDE diffusion modeling to image restoration by directly modeling degradation with a mean-reverting process and reversing it to recover high-quality images. The approach combines a closed-form score-learning formulation with maximum-likelihood training and performs competitively across diverse restoration tasks.

  • Motivation: Standard diffusion restoration methods can initialize reverse processes with high-variance noise, which may yield poor restoration of the ground-truth high-quality image.These methods typically diffuse images to pure noise and may require known degradation models or parameters at test time.
  • Method: The proposed mean-reverting SDE models degradation from a high-quality image toward a noisy low-quality counterpart, then restores the image by simulating the reverse-time SDE.The forward process moves x(0) toward µ + ϵ, where µ is the low-quality image.
  • Method: The formulation has a closed-form solution that provides the ground-truth time-dependent score for training a neural network.This adapts score-based SDE methods while directly representing the degradation process.
  • Training: A maximum-likelihood reverse-trajectory objective stabilizes training and consistently improves restoration performance over the common score-matching objective.The objective is presented as an alternative loss for neural-network training.
  • Experiments: The approach is evaluated on deraining, deblurring, denoising, super-resolution, inpainting, and dehazing, with new state-of-the-art results on two deraining datasets.The quantitative comparisons cover deraining, deblurring, and denoising; the remaining tasks are demonstrated qualitatively.

2. Background

SDE-based diffusion models transform data into a tractable terminal noise distribution and generate samples by simulating a learned reverse-time process. Their reverse dynamics depend on estimating the generally intractable time-dependent score function.

  • SDE formulation: An SDE diffusion process uses drift and dispersion functions to evolve an initial data state x(0) over continuous time.The Wiener process supplies the stochastic component of the process.
  • Forward process: The terminal state is typically designed to follow a Gaussian distribution with fixed mean and variance, representing the data-to-noise transformation.The SDE gradually transforms the data distribution into fixed Gaussian noise.
  • Reverse process: Sampling reverses the diffusion process by simulating the corresponding SDE backward in time from the terminal distribution.The reverse-time representation uses a reverse-time Wiener process.
  • Score estimation: Because the score ∇x log p_t(x) is generally intractable, diffusion models train a time-dependent neural network sθ(x,t) to approximate it.The approximation is learned under a score-matching objective.

3. Method

IR-SDE combines a mean-reverting forward process that models image degradation with reverse-time sampling for restoration. Its analytically tractable transitions support score learning, while a maximum-likelihood objective targets an optimal reverse trajectory and improves training stability.

  • Method: IR-SDE combines a mean-reverting SDE with a maximum-likelihood objective to model degradation and learn restoration trajectories.The forward process is paired with reverse-time simulation, while training optimizes the reverse trajectory likelihood.
  • 3.1. Forward SDE for Image Degradation: The forward SDE transfers a high-quality image toward its paired low-quality counterpart through noisy interpolation.The state mean converges to the degraded image and the variance converges to a stationary Gaussian variance.
  • 3.2. Reverse-Time SDE for Image Restoration: The trained noise network generates high-quality images by numerically solving the reverse-time IR-SDE from a noisy terminal state.Euler–Maruyama or Milstein’s method can be used for iterative reverse simulation.
  • 3.2. Reverse-Time SDE for Image Restoration: During training, the available high-quality image enables computation of the conditional score, which a neural network estimates from the degraded state and noise.The score can be reparameterized using standard Gaussian noise, and a conditional noise network receives the state, low-quality image, and time.
  • 3.3. Maximum Likelihood Training: Maximum-likelihood training addresses instability observed with score matching by optimizing the reverse trajectory toward the likelihood-maximizing path.The reverse transition is derived using Gaussian distributions and Bayes’ rule, with an analytically defined optimal reversed state.

4. Experiments

Experiments evaluate IR-SDE across deraining, deblurring, denoising, and qualitative extensions to super-resolution, inpainting, and dehazing. The method achieves strong perceptual restoration, while deterministic denoising variants improve fidelity and reduce sampling steps.

  • Image Deraining: IR-SDE achieves the best performance on all metrics across the Rain100H and Rain100L deraining datasets.Its LPIPS and FID scores are markedly better than competing approaches, and the visual results are described as realistic and high-fidelity.
  • Gaussian Image Denoising: Denoising-ODE and Denoising-SDE require 22 steps to recover clean images, compared with 100 steps for IR-SDE.The denoising visual and tabular comparisons use noise level σ = 25.
  • Image Deblurring: IR-SDE surpasses DeblurGAN-v2 by 1.15 dB in PSNR on GoPro and achieves the best perceptual performance overall.The reported results indicate realistic sharp images that remain consistent with ground truths.
  • Gaussian Image Denoising: The stochastic IR-SDE and Denoising-SDE have worse PSNR and SSIM than CNN-based methods, whereas Denoising-ODE achieves significantly better PSNR on all test datasets.IR-SDE has high perceptual performance, while Denoising-ODE produces less noisy realistic results and outperforms DDRM across all reported metrics on cropped images.
  • Qualitative Experiments: Qualitative experiments apply IR-SDE to super-resolution, face inpainting, and dehazing, producing detailed realistic images, semantically reasonable faces, and haze-free indoor scenes.Super-resolution comparisons use L2-trained EDSR, inpainting uses CelebA-HQ, and dehazing uses RESIDE ITS and SOTS.

5. Discussion and Analysis

The analysis examines how IR-SDE reverses degradation and noise, how maximum likelihood training affects stability, and how theta schedules shape restoration. It also identifies overly smooth late-step variance changes as a limitation and considers broader SDE choices.

  • Reverse-Time Restoration Process: IR-SDE prioritizes removing original degradation during reverse restoration and performs Gaussian denoising mainly in the final steps.For deraining and deblurring, most degradation is removed by the middle timesteps.
  • Reverse-Time Restoration Process: After 20 reverse steps, deblurring PSNR and LPIPS increase and then converge during the final steps.The performance curves are computed for a single-image deblurring case.
  • Maximum Likelihood Objective: The maximum likelihood objective stabilizes training and improves restoration compared with the noise-matching objective.Noise-matching PSNR fluctuates and can deteriorate during training in deraining and denoising.
  • Theta Schedule: IR-SDE varies theta while constraining sigma_t through the stationary variance, with constant, linear, and cosine schedules explored.A constant theta reduces the forward process to an Ornstein–Uhlenbeck process.
  • Limitations and Future Work: The exponential variance term produces overly smooth changes in the last few steps, making neighboring states difficult to distinguish during learning.The authors plan to explore alternative theta schedules, especially when maximum likelihood training optimizes state differences.
  • Generalized SDE Choices: The conditional score can be generalized by selecting alternative SDEs, including geometric Brownian motion, whose score corresponds to an exponential distribution.This extension is described through Tweedie’s formula.

6. Related Work

Related work spans supervised deep restoration, mean-reverting SDEs for speech and JPEG artifact removal, and diffusion-based restoration. IR-SDE differs through its stationary-variance design, cosine schedule, and maximum likelihood loss.

  • Image Restoration: Supervised deep learning remains a common approach to image restoration, with CNN and transformer architectures widely explored.Image restoration includes tasks such as deraining, deblurring, denoising, and super-resolution.
  • Mean-Reverting SDEs: Prior mean-reverting SDE work addressed speech enhancement, dereverberation, and JPEG artifact removal using OU-like or linear-theta processes.These approaches used different diffusion coefficients and did not impose the stationary variance condition.
  • Schedule Analysis: The cosine schedule yields a variance that changes overly smoothly in the final steps.This behavior is shown in the variance analysis associated with Figure 12.
  • Mean-Reverting SDEs: For JPEG artifact removal, the concurrent mean-reverting SDE baselines are outperformed by IR-SDE’s cosine theta scheduler.Those prior methods also use the standard score-matching objective rather than the proposed maximum likelihood-based loss.

7. Conclusion

The paper concludes that mean-reverting SDEs provide a broadly applicable restoration framework, while closed-form transitions and maximum likelihood training support score learning and stable performance. It reports applicability across six restoration tasks and identifies schedule optimization and cheaper sampling as future directions.

  • Conclusion: IR-SDE applies to a wide class of image restoration tasks and achieves highly competitive performance across six diverse tasks.The method combines a closed-form mean-reverting SDE with a maximum likelihood-based training objective.
  • Conclusion: The closed-form solution enables computation of the ground-truth time-dependent score and neural-network score estimation.This is a central role of the proposed SDE formulation.
  • Conclusion: The maximum likelihood objective significantly stabilizes neural-network training and consistently improves restoration performance.The conclusion presents this as a distinct contribution alongside the mean-reverting SDE.
  • Conclusion: Future work will optimize the theta schedule and sampling procedures to reduce computational cost at test time.

A. Proofs

The proofs establish tractable Gaussian transitions and an analytically solvable mean-reverting SDE, then derive the optimal reverse transition through negative log likelihood. The denoising specialization supports both stochastic and deterministic recovery and permits starting from intermediate noise levels.

  • SDE Solution: The mean-reverting SDE has analytically computable integrals, yielding a closed-form transition with Gaussian mean and variance.The proof uses a differentiable surrogate and Itô’s formula before identifying the transition kernel.
  • Optimal Reverse Path: The reverse transition is optimized by minimizing negative log likelihood, with Bayes’ rule and tractable transitions producing the optimal reversing solution.The second derivative is positive, confirming the derived point is optimal.
  • Denoising SDE/ODE: The denoising SDE is obtained by specializing the IR-SDE transition and reverse-time process to a clean-image mean state.The appendix also gives the corresponding sampling strategy and probability-flow ODE.
  • Denoising SDE/ODE: The probability-flow ODE provides a deterministic trajectory with the same marginal probability densities as the IR-SDE.This deterministic recovery becomes available once the score is known.
  • Denoising SDE/ODE: A known real noise level determines a timestep whose transition variance matches that noise level, allowing denoising from an intermediate state.This supports arbitrary noise levels and improves sample efficiency.

C. Relationship between Maximum Likelihood Objective and DDPM

The paper connects its maximum likelihood objective to DDPM by showing that minimizing reverse-process negative log-likelihood recovers DDPM’s reverse mean.

  • The maximum likelihood objective is applied to DDPM to mathematically establish its connection with diffusion models.
  • Minimizing the negative log-likelihood yields an optimal reverse state x*_{t−1}.
  • The resulting optimum matches DDPM’s reverse mean and guarantees learning for the reverse process.

D. Additional Implementation Details

The experiments use a shared efficient U-Net setup, fixed optimization settings, and a cosine-based schedule for the mean-reverting SDE.

  • All experiments use a DDPM-like U-Net without group normalization or self-attention layers, with padding for size-preserving outputs.The stationary variance is set to 10 over 255, and inference uses 100 steps.
  • Training generally uses 128 × 128 patches, batch size 16, Adam optimization, and 500 thousand total steps.The initial learning rate is 10^-4 and decays by half every 200 thousand iterations.
  • The θ schedule is defined as the flipped version of the cosine noise schedule.
  • The diffusion coefficient σ_t is computed from the stationary condition σ_t^2 θ_t = λ^2, using a discrete approximation for accumulated θ.
  • The terminal exponential term is set to δ = 0.005 rather than zero to alleviate over-smoothing.The time step is then computed from δ.

E. Additional Experimental Results

Additional experiments report loss, dehazing, denoising, complexity, and qualitative results across restoration tasks, with generally favorable outcomes for the proposed methods.

  • Loss comparison: Maximum likelihood significantly improves deraining performance over noise matching across all reported criteria.
  • Additional quantitative results: Denoising-ODE achieves the best perceptual performance across all scenes in the reported denoising comparisons.The results cover three test sets and different noise levels, including CBSD68 comparisons with KBNet.
  • Additional quantitative results: Dehazing results compare the proposed method with the CNN baseline because DDRM requires degradation parameters known at test time.
  • Model complexities: The proposed model only slightly increases CNN-baseline parameters and FLOPs, while diffusion-style reverse sampling increases inference cost.The reverse process requires repeated network evaluations; training generally converges quickly after sampling noises.
  • Additional qualitative results: Qualitative results span denoising, deraining, deblurring, super-resolution, inpainting, and dehazing.In most tasks, the proposed outputs are described as sharper and more realistic.
  • Additional qualitative results: The additional visual results include dehazing comparisons on the reported dehazing task.
Loading 2301.11699v3…