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

2026-01-29

AI Summary

This work addresses the challenges of conflicting optimization objectives and limited adaptability in All-in-One image restoration caused by heterogeneous degradations. To this end, we propose an uncertainty-aware diffusion bridge model that formulates the restoration process as a stochastic transport problem guided by pixel-level uncertainty. By introducing relaxed terminal constraints to mitigate drift singularities and designing a dual scheduling mechanism for both noise and transport paths, our method adaptively regulates the dynamics of the transport process. This unified framework effectively handles diverse degradation types and achieves state-of-the-art performance across multiple image restoration tasks under single-step inference.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Diffusion Models for Vision
• Planning, Routing, and Scheduling: Scheduling under Uncertainty

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• All-in-One Image Restoration
• Heterogeneous Degradations
• Conflicting Optimization Objectives
• Diffusion Bridge
• Uncertainty Modeling

Innovation

• Uncertainty-Aware Diffusion Bridge
• All-in-One Image Restoration
• Stochastic Transport
• Relaxed Diffusion Bridge
• Dual Modulation Strategy