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

TL;DR

This paper introduces an Uncertainty-Aware Diffusion Bridge Model for All-in-One Image Restoration, reformulating the problem as a stochastic transport task guided by pixel-wise uncertainty, enabling effective handling of diverse degradations in a single step.

Contribution

It proposes a novel diffusion bridge formulation with relaxed constraints and a dual modulation strategy to unify heterogeneous image restoration tasks.

Findings

Achieves state-of-the-art results across multiple restoration tasks

Effectively models degradation uncertainty with pixel-wise precision

Performs restoration in a single inference step

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.