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

arXiv: https://arxiv.org/abs/2601.21592 · HTML: https://mlsanity.com/paper/2601.21592/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 63 references that are in the map.
- First posted on arXiv 29 Jan 2026.

## Builds on (most cited first, 30 of 40)

- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Restormer: Efficient Transformer for High-Resolution Image Restoration](https://mlsanity.com/paper/2111.09881/) (2021)
- [Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising](https://mlsanity.com/paper/1608.03981/) (2016)
- [MUSIQ: Multi-scale Image Quality Transformer](https://mlsanity.com/paper/2108.05997/) (2021)
- [Uformer: A General U-Shaped Transformer for Image Restoration](https://mlsanity.com/paper/2106.03106/) (2021)
- [Deep Multi-scale Convolutional Neural Network for Dynamic Scene Deblurring](https://mlsanity.com/paper/1612.02177/) (2016)
- [Deep Retinex Decomposition for Low-Light Enhancement](https://mlsanity.com/paper/1808.04560/) (2018)
- [MANIQA: Multi-dimension Attention Network for No-Reference Image Quality Assessment](https://mlsanity.com/paper/2204.08958/) (2022)
- [FFDNet: Toward a Fast and Flexible Solution for CNN based Image Denoising](https://mlsanity.com/paper/1710.04026/) (2017)
- [Benchmarking Single Image Dehazing and Beyond](https://mlsanity.com/paper/1712.04143/) (2017)
- [Learning Deep CNN Denoiser Prior for Image Restoration](https://mlsanity.com/paper/1704.03264/) (2017)
- [Deep Joint Rain Detection and Removal from a Single Image](https://mlsanity.com/paper/1609.07769/) (2016)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Blind Image Quality Assessment via Vision-Language Correspondence: A Multitask Learning Perspective](https://mlsanity.com/paper/2303.14968/) (2023)
- [Multi-Scale Progressive Fusion Network for Single Image Deraining](https://mlsanity.com/paper/2003.10985/) (2020)
- [BBDM: Image-to-image Translation with Brownian Bridge Diffusion Models](https://mlsanity.com/paper/2205.07680/) (2022)
- [DesnowNet: Context-Aware Deep Network for Snow Removal](https://mlsanity.com/paper/1708.04512/) (2017)
- [Image Restoration with Mean-Reverting Stochastic Differential Equations](https://mlsanity.com/paper/2301.11699/) (2023)
- [Denoising Diffusion Bridge Models](https://mlsanity.com/paper/2309.16948/) (2023)
- [Human-Aware Motion Deblurring](https://mlsanity.com/paper/2001.06816/) (2020)
- [Learning Enriched Features for Fast Image Restoration and Enhancement](https://mlsanity.com/paper/2205.01649/) (2022)
- [AdaIR: Adaptive All-in-One Image Restoration via Frequency Mining and Modulation](https://mlsanity.com/paper/2403.14614/) (2024)
- [AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion](https://mlsanity.com/paper/2310.10123/) (2023)
- [OneRestore: A Universal Restoration Framework for Composite Degradation](https://mlsanity.com/paper/2407.04621/) (2024)
- [Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models](https://mlsanity.com/paper/2304.08291/) (2023)
- [RestoreFormer: High-Quality Blind Face Restoration from Undegraded Key-Value Pairs](https://mlsanity.com/paper/2201.06374/) (2022)
- [ProRes: Exploring Degradation-aware Visual Prompt for Universal Image Restoration](https://mlsanity.com/paper/2306.13653/) (2023)
- [Complexity Experts are Task-Discriminative Learners for Any Image Restoration](https://mlsanity.com/paper/2411.18466/) (2024)

## Built on by (most cited first, 1 of 1)

- [DocPure: Prompt-Free Unified Document Restoration via Degradation-Aware Structure-Guided Wavelet Modulation](https://mlsanity.com/paper/2608.09536/) (2026)

## Abstract (arXiv)

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.