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

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. Code is available at https://github.com/Jabruson/UDBM.

Real-world photos are often affected by a mix of unpredictable degradations like rain, blur, haze, or low light. Fixing all these issues with a single model is difficult because treating one degradation often conflicts with treating another. Current methods generally rely on rigid rules that struggle to account for the unique severity of these different types of damage. To solve this, we developed the Uncertainty-Aware Diffusion Bridge Model (UDBM). Our model examines a damaged image and calculates uncertainty—essentially determining how severely damaged and difficult to fix each pixel is. It then builds a "bridge" to smoothly transform the damaged image into a clean one. The model adaptively spends more effort repairing highly uncertain, heavily damaged areas, while moving quickly through the confident, clean parts. This approach allows our model to achieve state-of-the-art image restoration across a wide variety of tasks. It efficiently adapts to the exact damage of each image in just a single computational step, making it fast and reliable for real-world applications.