Residual Diffusion Bridge Model for Image Restoration

AI summary

Overview

• Research area: Computer vision, specifically universal image restoration using diffusion bridge models and stochastic differential equations (SDEs).
• Technical level: Advanced. The paper is built around Ornstein–Uhlenbeck SDEs, Doob's h-transform, and stochastic interpolants, and assumes familiarity with diffusion model formulations.
• Scope: The paper proposes the Residual Diffusion Bridge Model (RDBM), a unified analytical framework for diffusion bridge models that uses the residual between clean and degraded images to adaptively modulate noise injection and removal across five image restoration tasks.

What This Paper Is About

Existing diffusion bridge methods for image restoration add noise globally across the whole image and then remove it globally, which means undegraded regions get unnecessarily perturbed and are not perfectly reconstructed. The paper's goal is twofold: build a unified theoretical account of how existing bridge formulations relate to one another, and use the difference (residual) between the clean image and the degraded observation to decide where noise should be applied, so degraded regions are restored while intact regions are left alone.

Key Contributions

1. A scalable, unified diffusion bridge framework. The authors recast a general mean-reverting Ornstein–Uhlenbeck process with Doob's h-transform to derive a generalized diffusion bridge, describing probabilistic transitions between arbitrary paired distributions and characterized as generalized stochastic interpolants.
2. Residual-modulated noise injection and removal. Because the terminal state is deterministic, they set π = x₀ − μ (the residual between clean image x₀ and degraded image μ), so noise perturbation is applied selectively to regions of differing degradation levels rather than uniformly.
3. Unification of existing bridge models. They show analytically and empirically that mainstream bridge formulations — Flow Matching, VE Bridge, VP Bridge, Brownian Bridge, and OU Bridge — are all special cases of RDBM under specific configurations of θ_t, λ, and π.
4. Extensive empirical validation. Experiments across five restoration tasks (deraining, low-light enhancement, desnowing, dehazing, deblurring), plus image translation, inpainting, and zero-shot real-world generalization, are used to substantiate the framework.

Main Findings

• State-of-the-art restoration results: RDBM-L achieves 30.31 PSNR / 0.884 SSIM on deraining, 24.53 / 0.812 on low-light enhancement, 32.59 / 0.961 on desnowing, 33.45 / 0.965 on dehazing, and 29.04 / 0.877 on deblurring, for an average of 31.04 PSNR / 0.917 SSIM. The paper reports this as average gains of 1.55 dB in PSNR and 0.013 in SSIM over comparison methods.
• Efficiency relative to baselines: RDBM-L uses 7.73M parameters and 32.93G FLOPs. The smallest variant, RDBM-T, uses 0.45M parameters and 5.74G FLOPs; RDBM-S uses 1.07M / 8.01G; RDBM-B uses 3.65M / 23.97G. RDBM-B is reported as attaining the best average metrics among lightweight variants.
• Residual modulation is the decisive design choice: In the π ablation, π = 0 (flow-matching-like) averages 28.21 PSNR / 0.872 SSIM, π = 1 (stochastic-interpolant-like) averages 30.15 / 0.903, π = x₀ − x_T averages 31.04 / 0.917, and |x₀ − x_T| averages 30.94 / 0.915. The two residual-based variants perform best.
• Noise schedule matters, and the best one is task-dependent: Comparing linear, cosine, and sigmoid schedules at λ = 10/255, the cosine schedule gives the best overall average (31.04 / 0.917), though the paper notes optimal schedules differ across tasks.
• Stationary variance of λ = 10/255 is optimal: Across λ values of 1/255, 10/255, 20/255, 50/255 and 100/255, the 10/255 setting with a cosine schedule gives the best average (31.04 / 0.917); 1/255 averages 30.36 / 0.915 and 100/255 averages 29.08 / 0.906.
• Performance peaks at 10 sampling steps: Averaged over tasks, 2 NFEs give 22.81 PSNR / 0.815 SSIM, 5 NFEs give 29.61 / 0.905, 10 NFEs give 31.04 / 0.917, and performance declines beyond that (20 NFEs: 30.58 / 0.915; 50: 30.28 / 0.913; 100: 30.19 / 0.912). The authors attribute the decline to the model prioritizing primary degradations before secondary ones in multi-degradation samples.
• Strong zero-shot generalization on unknown tasks: On the under-display-camera datasets POLED and TOLED, RDBM reaches 19.834 PSNR / 0.715 SSIM / 0.012 MSE / 0.351 LPIPS on POLED and 30.809 / 0.870 / 0.001 / 0.202 on TOLED, outperforming all listed baselines including Restomer, AirNet, Prompt-IR, ProRes, IDR, IRSDE, AutoDIR, DA-CLIP, GOUB, ConvIR, DeepSNNet, AWRaCLe, and MaIR.
• Known-task generalization: Using the non-reference metrics MetaIQA and NIQE, the paper states RDBM outperforms other universal models across benchmarks; the numerical table for this comparison is truncated in the provided content and the RDBM row is not shown.
• Smooth residual-to-noise ratio: The paper derives that fixing π = x₀ − μ makes the pixel-wise residual-to-noise ratio independent of pixel position, R(t,i,j) = R(t), preserving the smooth monotonic decay that π = 1 can disrupt.

Methodology in Plain English

The starting point is an Ornstein–Uhlenbeck process, a kind of noise process that drifts back toward a mean value instead of wandering freely. The authors let that mean be the degraded image, so the forward process moves a clean image toward its degraded counterpart rather than toward pure random noise. They then apply Doob's h-transform, a mathematical tool that forces a random process to arrive at a specific endpoint at the final time step, which turns the process into a diffusion bridge connecting the clean-image distribution and the degraded-image distribution.

The key modification is a scaling factor π on the noise term. Previous work sets π = 1, applying noise everywhere. Instead, the authors set π to the residual, the per-pixel difference between the clean image and the degraded image. Regions that are already fine have a small residual and therefore receive almost no noise; heavily degraded regions have a large residual and receive more. Because the residual is known in closed form, the noise scale becomes spatially adaptive while still following a clean mathematical rule.

The network does not have to predict the residual and the noise separately. From the reverse-process equation, one unknown network output can predict the product of the residual and the noise at once, and training minimizes the L1 distance between that predicted product and the true product. Sampling reverses the process step by step from the degraded image back to the restored one.

The residual-to-noise ratio (RNR) is the quantity the authors use to describe how the balance between signal and noise shifts over time. It has two parts: one depending on the residual and the fixed ratio λ, and one depending entirely on the θ_t series that decreases monotonically. Setting π equal to the residual eliminates the discontinuous per-pixel jumps that π = 1 can produce.

Why This Matters

Impact on research. The paper supplies a single theoretical framework in which Flow Matching, VE Bridge, VP Bridge, Brownian Bridge, and OU Bridge all appear as special configurations of three parameters (θ_t, λ, π). That lets researchers compare apparently different methods on common ground rather than treating them as unrelated variants, and it pins down which assumption in prior bridge models causes redundant reconstruction of undamaged image regions.

Real-world applications.

• Restoring photos taken with under-display cameras on phones, where the sensor is behind the screen — the paper tests this on POLED and TOLED.
• Enhancing images captured in rain, snow, haze, low light, or motion blur, as a preprocessing step for downstream vision systems.
• Zero-shot cleanup of real-world imagery where no ground truth exists, evaluated using non-reference metrics.
• Image-to-image translation and inpainting, which the paper also demonstrates.

Industry relevance. By making restoration spatially adaptive, the method avoids degrading already-good parts of an image, which matters for pipelines where over-processing is as harmful as under-processing. The parameter range from 0.45M (RDBM-T) to 7.73M (RDBM-L) parameters supports deployment from constrained mobile hardware up to more capable systems, and the ability to act on multiple degradation types with one unified model reduces the need for task-specific models.

Future Directions

• Resolving the multi-degradation ordering problem. The paper observes that beyond 10 NFEs performance degrades because the model removes a primary degradation before secondary ones, causing output to deviate from the available reference. How to handle compound degradations in a principled order is left open.
• Task-adaptive noise schedules. The ablation shows the optimal schedule differs by task (cosine is best on average), suggesting a per-task or learned schedule could do better than a single global choice.
• Extending the framework to more than two paired distributions. The framework is presented as connecting any paired distributions; whether it scales to chained or multi-stage transitions is not demonstrated.
• Standardizing the theoretical mapping. The paper claims all mainstream bridge models are special cases; a full exploration of the space of (θ_t, λ, π) configurations beyond those tabulated could reveal further useful variants.

Target Audience

Graduate students and researchers working on diffusion models, generative models, or image restoration who already understand SDE-based diffusion; practitioners building image enhancement or restoration pipelines who want to know whether adaptive noise modulation improves results in practice; and readers interested in unifying theoretical accounts of diffusion bridge methods. Beginners will find the mathematical development in Section 4 demanding but may still follow the empirical comparisons and the spatial-adaptivity intuition.

Authors’ abstract

Diffusion bridge models establish probabilistic paths between arbitrary paired distributions and exhibit great potential for universal image restoration. Most existing methods merely treat them as simple variants of stochastic interpolants, lacking a unified analytical perspective. Besides, they indiscriminately reconstruct images through global noise injection and removal, inevitably distorting undegraded regions due to imperfect reconstruction. To address these challenges, we propose the Residual Diffusion Bridge Model (RDBM). Specifically, we theoretically reformulate the stochastic differential equations of generalized diffusion bridge and derive the analytical formulas of its forward and reverse processes. Crucially, we leverage the residuals from given distributions to modulate the noise injection and removal, enabling adaptive restoration of degraded regions while preserving intact others. Moreover, we unravel the fundamental mathematical essence of existing bridge models, all of which are special cases of RDBM and empirically demonstrate the optimality of our proposed models. Extensive experiments are conducted to demonstrate the state-of-the-art performance of our method both qualitatively and quantitatively across diverse image restoration tasks. Code is publicly available at https://github.com/MiliLab/RDBM.