[Paper Review] Image Restoration Through Generalized Ornstein-Uhlenbeck Bridge

Conghan Yue, Zhengwei Peng+4 more|arXiv (Cornell University)|Dec 16, 2023

Generative Adversarial Networks and Image Synthesis4 citations

TL;DR

This paper proposes the Generalized Ornstein-Uhlenbeck Bridge (GOUB), a diffusion-based image restoration model that leverages the mean-reverting property of the generalized Ornstein-Uhlenbeck process and Doob’s h-transform to enable efficient, end-to-end learning of point-to-point mappings from low-quality to high-quality images. The method achieves state-of-the-art performance in image restoration tasks such as deraining, super-resolution, and inpainting by enabling optimal diffusion with minimal computational cost.

ABSTRACT

Diffusion models exhibit powerful generative capabilities enabling noise mapping to data via reverse stochastic differential equations. However, in image restoration, the focus is on the mapping relationship from low-quality to high-quality images. Regarding this issue, we introduce the Generalized Ornstein-Uhlenbeck Bridge (GOUB) model. By leveraging the natural mean-reverting property of the generalized OU process and further eliminating the variance of its steady-state distribution through the Doob's h-transform, we achieve diffusion mappings from point to point enabling the recovery of high-quality images from low-quality ones. Moreover, we unravel the fundamental mathematical essence shared by various bridge models, all of which are special instances of GOUB and empirically demonstrate the optimality of our proposed models. Additionally, we present the corresponding Mean-ODE model adept at capturing both pixel-level details and structural perceptions. Experimental outcomes showcase the state-of-the-art performance achieved by both models across diverse tasks, including inpainting, deraining, and super-resolution. Code is available at \url{https://github.com/Hammour-steak/GOUB}.

Motivation & Objective

• To address the challenge of learning efficient, task-agnostic image restoration mappings from low-quality to high-quality images.
• To overcome limitations in existing bridge models that rely on iterative approximations or suboptimal diffusion paths.
• To unify various existing bridge models under a single mathematical framework by revealing their special cases within the GOUB model.
• To develop a parameterization mechanism that enhances pixel-level and structural fidelity in restored images.
• To enable end-to-end training via maximum likelihood estimation on a diffusion process with fixed endpoints.

Proposed method

• The GOUB model uses the generalized Ornstein-Uhlenbeck (GOU) process as the forward diffusion process, naturally exhibiting mean reversion from high-quality to low-quality images.
• It applies Doob’s h-transform to adjust the steady-state variance of the GOU process, enabling precise bridging between high-quality and low-quality image distributions.
• The reverse process is derived via maximum likelihood estimation, allowing end-to-end training without task-specific priors or degradation matrices.
• A novel Mean-ODE model is introduced based on the GOUB’s parameterization, enhancing capture of pixel-level details and structural coherence.
• The model’s SDE formulation is derived from the Fokker-Planck equation, with drift and diffusion coefficients optimized through score matching.
• The framework unifies existing bridge models as special cases, demonstrating mathematical optimality and computational efficiency.

Figure 1: The overview of the proposed image restoration GOUB. It can be observed that the GOU process is capable of transferring High Quality (HQ) into noisy Low Quality (LQ). Additionally, through the application of h -transform, minor adjustments can be made to the entire process, enabling the GO

Experimental results

Research questions

• RQ1 Can a diffusion-based model achieve optimal, end-to-end image restoration without task-specific priors or iterative refinement?
• RQ2 How does the mean-reverting property of the generalized Ornstein-Uhlenbeck process enable efficient point-to-point image mapping?
• RQ3 What is the mathematical relationship between existing bridge models and the proposed GOUB framework?
• RQ4 Can a reparameterized ODE model (Mean-ODE) outperform standard diffusion models in capturing fine-grained image details and structural perception?
• RQ5 Does the GOUB model achieve state-of-the-art performance across diverse image restoration benchmarks?

Key findings

• The GOUB model achieves state-of-the-art performance on multiple image restoration benchmarks, including Rain100H for deraining, DIV2K for super-resolution, and CelebA-HQ for inpainting.
• The Mean-ODE variant demonstrates superior capability in preserving pixel-level details and structural consistency compared to standard diffusion baselines.
• All existing bridge models are shown to be special cases of the GOUB framework, confirming its mathematical generality and unifying power.
• The model enables efficient training via maximum likelihood estimation without requiring iterative sampling or classifier guidance.
• Empirical results confirm the optimality of the proposed diffusion path, with lower cost and better performance than prior methods.
• The method generalizes across tasks without requiring task-specific degradation modeling, enhancing practicality and scalability.

Figure 2: Qualitative comparison of the visual results of different inpainting methods on the CelebA-HQ datasets.