Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unifying Heterogeneous Degradations: Uncertainty-Aware Diffusion Bridge Model for All-in-One Image Restoration

Luwei Tu, Jiawei Wu, Xing Luo, Zhi Jin

TL;DR

An Uncertainty-Aware Diffusion Bridge Model (UDBM) is proposed, which innovatively reformulates AiOIR as a stochastic transport problem steered by pixel-wise uncertainty, and introduces a relaxed diffusion bridge formulation which replaces the strict terminal constraint with a relaxed constraint.

Abstract

All-in-One Image Restoration (AiOIR) faces the fundamental challenge in reconciling conflicting optimization objectives across heterogeneous degradations. Existing methods are often constrained by coarse-grained control mechanisms or fixed mapping schedules, yielding suboptimal adaptation. To address this, we propose an Uncertainty-Aware Diffusion Bridge Model (UDBM), which innovatively reformulates AiOIR as a stochastic transport problem steered by pixel-wise uncertainty. By introducing a relaxed diffusion bridge formulation which replaces the strict terminal constraint with a relaxed constraint, we model the uncertainty of degradations while theoretically resolving the drift singularity inherent in standard diffusion bridges. Furthermore, we devise a dual modulation strategy: the noise schedule aligns diverse degradations into a shared high-entropy latent space, while the path schedule adaptively regulates the transport trajectory motivated by the viscous dynamics of entropy regularization. By effectively rectifying the transport geometry and dynamics, UDBM achieves state-of-the-art performance across diverse restoration tasks within a single inference step.

Unifying Heterogeneous Degradations: Uncertainty-Aware Diffusion Bridge Model for All-in-One Image Restoration

TL;DR

An Uncertainty-Aware Diffusion Bridge Model (UDBM) is proposed, which innovatively reformulates AiOIR as a stochastic transport problem steered by pixel-wise uncertainty, and introduces a relaxed diffusion bridge formulation which replaces the strict terminal constraint with a relaxed constraint.

Abstract

All-in-One Image Restoration (AiOIR) faces the fundamental challenge in reconciling conflicting optimization objectives across heterogeneous degradations. Existing methods are often constrained by coarse-grained control mechanisms or fixed mapping schedules, yielding suboptimal adaptation. To address this, we propose an Uncertainty-Aware Diffusion Bridge Model (UDBM), which innovatively reformulates AiOIR as a stochastic transport problem steered by pixel-wise uncertainty. By introducing a relaxed diffusion bridge formulation which replaces the strict terminal constraint with a relaxed constraint, we model the uncertainty of degradations while theoretically resolving the drift singularity inherent in standard diffusion bridges. Furthermore, we devise a dual modulation strategy: the noise schedule aligns diverse degradations into a shared high-entropy latent space, while the path schedule adaptively regulates the transport trajectory motivated by the viscous dynamics of entropy regularization. By effectively rectifying the transport geometry and dynamics, UDBM achieves state-of-the-art performance across diverse restoration tasks within a single inference step.
Paper Structure (67 sections, 4 theorems, 111 equations, 17 figures, 11 tables, 2 algorithms)

This paper contains 67 sections, 4 theorems, 111 equations, 17 figures, 11 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Diffusion Bridge for Image Restoration
  4. All-in-One Image Restoration
  5. Methodology
  6. Problem Formulation
  7. Relaxed Diffusion Bridge
  8. Uncertainty-Aware Entropic Regularization
  9. The UDBM Framework
  10. Experiments
  11. Setting (a): All-in-One
  12. Setting (b): Task-Specific
  13. Setting (c): Generalization
  14. Ablation Study
  15. Conclusion
  16. ...and 52 more sections

Key Result

Proposition 3.2

For a strict terminal constraint $p(\mathbf{x}_1)=\delta(\mathbf{x}_1-\mathbf{x}_{lq})$ in standard diffusion bridge, the harmonic function $h(\mathbf{x}_t, t)$ in Eq. eq:bridge_sde_labeled_full is instantiated as $h_{\text{strict}}(\mathbf{x}_t, t) \triangleq p(\mathbf{x}_1 = \mathbf{x}_{lq} \mid \

Figures (17)

  • Figure 1: Motivation and Analysis. (a) Schematic comparison of different paradigms. Fixed schedule methods employ rigid mappings regardless of degradation severity, while control modulation methods suffer from erroneous flows caused by imprecise conditional. In contrast, our method simultaneously accounts for degradation heterogeneity and shared information. (b) Imposing incorrect prompts in AutoDIR jiang2024autodir induces artifacts. (c) Uncertainty Visualization. Uncertainty maps zhang2025uncertainty naturally capture the severity of diverse degradations.
  • Figure 2: Overview of the UDBM. (a) The schematic diagram of marginal distribution of UDBM. (b) The dual uncertainty-guided path and noise schedules. (c) t-SNE visualization of $\mathbf{x}_t$ of different degradations, showing the gradually alignment of diverse degradations in the shared high-entropy latent space.
  • Figure 3: Visual comparison of AiOIR. Our UDBM exhibits better artifacts removal and details restoration. Bottom-right: error maps.
  • Figure 4: Visual comparison of unseen scenarios on the CDD11 dataset. The proposed UDBM shows better degradation removal compared to other methods.
  • Figure 5: Ablation studies on hyper-parameter $\pi_\mathrm{SB}$ (a) and $\lambda_b$ (b).
  • ...and 12 more figures

Theorems & Definitions (6)

  • Proposition 3.2
  • Theorem 3.3
  • Proposition 3.4
  • Remark 3.5
  • Theorem 6.1: Zero Geometric Curvature
  • proof